diff --git a/.agents/skills/triton-kernel-writing/SKILL.md b/.agents/skills/triton-kernel-writing/SKILL.md
index 6dd6a8520be0..a41de016b78d 100644
--- a/.agents/skills/triton-kernel-writing/SKILL.md
+++ b/.agents/skills/triton-kernel-writing/SKILL.md
@@ -22,10 +22,14 @@ description: Write or review Triton kernels for vLLM, with practical guidance fo
   a small, legible heuristic when workloads need different choices. Use
   `triton.autotune` only when tuning is critical to performance, such as for a
   matrix multiplication. Otherwise prioritize simple code and fast startup.
-- Be careful to avoid unintended runtime JIT compilation. For example, put
-  unimportant runtime integer scalars in `do_not_specialize`, especially those
-  that may alternate between values such as 0 and 1, which can produce
-  different specialization keys.
+- Avoid unintended runtime JIT compilation. Put unimportant runtime integer
+  scalars in `do_not_specialize`, especially values that alternate between 0
+  and 1. Pointer alignment also affects specialization: an offset tensor view
+  can trigger a new variant after warmup used an aligned allocation. For
+  scalar-only pointers whose alignment provides no useful optimization, use
+  `@triton.jit(do_not_specialize_on_alignment=["end_ptr"])`. Retain alignment
+  specialization when vectorized accesses benefit from it. See the
+  [Triton JIT options](https://triton-lang.org/main/python-api/generated/triton.jit.html).
 - The Triton compiler does not guarantee safe ordering when a kernel writes to
   a pointer and subsequently reads from the same pointer. This pattern must
   have a `tl.debug_barrier()` between the write and read. The barrier
@@ -49,7 +53,9 @@ description: Write or review Triton kernels for vLLM, with practical guidance fo
 ## Validation
 
 - Check correctness at boundary shapes and at sizes that exercise masks and
-  large offsets.
+  large offsets. Include offset tensor views, not only fresh allocations.
+  Warm aligned buffers, then exercise offset views with vLLM's
+  `--jit-monitor-mode error --jit-monitor-verbose` to catch missing variants.
 - Choose accumulation and intermediate dtypes explicitly. Test numerically
   difficult inputs, not only random, well-scaled tensors.
 - Use `$kernel-microbenchmark` for benchmark construction, measurement, and
diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml
index 769e3d4adc3a..b50a082727f5 100644
--- a/.buildkite/hardware_tests/cpu.yaml
+++ b/.buildkite/hardware_tests/cpu.yaml
@@ -52,6 +52,52 @@ steps:
       tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
   parallelism: 2
 
+# Speculative decoding under a CPU platform. The unit tests already pass on
+# CPU but no CPU job runs them; the e2e tests gate the engine path itself,
+# which unit tests cannot reach.
+- label: CPU-Spec Decode Tests
+  key: cpu-spec-decode-tests
+  depends_on: []
+  soft_fail: true
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  - vllm/v1/spec_decode/
+  - vllm/v1/worker/cpu/
+  - vllm/v1/worker/cpu_model_runner.py
+  - vllm/v1/worker/gpu/spec_decode/
+  - vllm/v1/attention/backends/cpu_attn.py
+  - vllm/config/speculative.py
+  - vllm/model_executor/models/qwen3_5.py
+  - vllm/model_executor/models/qwen3_5_mtp.py
+  - vllm/model_executor/models/llama_eagle3.py
+  - vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+  - vllm/model_executor/layers/mamba/mamba_mixer2.py
+  - vllm/model_executor/layers/mamba/ops/cpu/
+  - tests/v1/e2e/test_cpu_spec_decode.py
+  - tests/v1/spec_decode/
+  - tests/v1/worker/
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      pytest -x -v -s \
+      tests/v1/worker/test_gpu_autoregressive_speculator.py \
+      tests/v1/worker/test_gpu_extract_hidden_states_speculator.py \
+      tests/v1/worker/test_gpu_batch_ordering.py \
+      tests/v1/worker/test_gpu_rejection_sampler_chunking.py \
+      tests/v1/worker/test_eagle3_aux_hidden_states_pp.py \
+      tests/v1/worker/test_spec_decode_embed_sharing_pp.py \
+      tests/v1/spec_decode/test_adaptive_verification.py \
+      tests/v1/spec_decode/test_dflash2.py \
+      tests/v1/spec_decode/test_dflash_causality.py \
+      tests/v1/spec_decode/test_eagle_draft_attn_metadata.py \
+      tests/v1/spec_decode/test_draft_attention_backend_override.py \
+      tests/v1/spec_decode/test_draft_moe_backend_override.py \
+      tests/v1/spec_decode/test_ngram.py \
+      tests/v1/spec_decode/test_llm_base_proposer_sampling.py
+      unset VLLM_CPU_CI_ENV
+      pytest -x -v -s tests/v1/e2e/test_cpu_spec_decode.py"
+
 # Note: SDE can't be downloaded from CI host because of AWS WAF
 # - label: CPU-Compatibility Tests
 #   depends_on: []
diff --git a/.buildkite/image_build/image_build.sh b/.buildkite/image_build/image_build.sh
index b81cd4ec0cbe..6d22acd5b16c 100755
--- a/.buildkite/image_build/image_build.sh
+++ b/.buildkite/image_build/image_build.sh
@@ -79,6 +79,54 @@ setup_buildx_builder() {
     docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
 }
 
+download_ci_hcl() {
+    # Returns non-zero instead of exiting: the optional export path must be
+    # able to give up without failing a build whose image already exists.
+    echo "--- :arrow_down: Downloading ci.hcl"
+    curl -sSfL -o "${CI_HCL_PATH}" "${CI_HCL_URL}" || return $?
+    echo "Downloaded to ${CI_HCL_PATH}"
+
+    if [[ ! -f "${CI_HCL_PATH}" ]]; then
+        echo "Error: ci.hcl not found at ${CI_HCL_PATH}"
+        return 1
+    fi
+}
+
+export_kernel_symbol_map_from_cache() {
+    # The existing-image path: everything here is best effort. Called as
+    # `... || echo`, so errexit is suspended inside and any failure just
+    # returns; the build keeps the image it already has.
+    local rc=0
+    download_ci_hcl \
+        && setup_buildx_builder \
+        && resolve_parent_commit \
+        && export PARENT_COMMIT \
+        && BUILD_TMP_DIR="$(mktemp -d)" \
+        && export_kernel_symbol_map || rc=$?
+    rm -rf -- "${BUILD_TMP_DIR:-}"
+    return "${rc}"
+}
+
+export_kernel_symbol_map() {
+    # Only when the build asked for it (nightly/post-merge). The map was
+    # produced inside the csrc-build stage during the main bake; this second
+    # bake of the scratch stage is a cache hit that just writes the file out.
+    if [[ "${VLLM_KERNEL_SYMBOL_MAP:-0}" != "1" ]]; then
+        return 0
+    fi
+    echo "--- :world_map: Exporting kernel symbol map"
+    local out_dir="${BUILD_TMP_DIR}/kernel-symbol-map"
+    mkdir -p "${out_dir}"
+    if docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --progress plain \
+        --set "kernel-symbol-map.output=type=local,dest=${out_dir}" kernel-symbol-map \
+        && [[ -s "${out_dir}/kernel_symbol_map.json.gz" ]]; then
+        ls -la "${out_dir}"
+        (cd "${out_dir}" && buildkite-agent artifact upload "kernel_symbol_map.json.gz")
+    else
+        echo "kernel symbol map export failed; continuing without it" >&2
+    fi
+}
+
 annotate_image_tags() {
     .buildkite/scripts/annotate-image-build.sh \
         "${IMAGE_TAG:-}" "${IMAGE_TAG_LATEST:-}"
@@ -90,6 +138,15 @@ check_and_skip_if_image_exists() {
         if docker manifest inspect "${IMAGE_TAG}" >/dev/null 2>&1; then
             echo "Image already exists: ${IMAGE_TAG}"
             echo "Skipping build"
+            if [[ "${VLLM_KERNEL_SYMBOL_MAP:-0}" == "1" ]]; then
+                # The image is reused, but the symbol map is an artifact of
+                # this build. Bake just the export target: its layers come
+                # from the registry cache, and the map step itself reruns
+                # if the cached image was built without the arg. Never
+                # fails the build: the image is already there.
+                export_kernel_symbol_map_from_cache \
+                    || echo "kernel symbol map: export from the cached image failed; continuing" >&2
+            fi
             annotate_image_tags
             exit 0
         fi
@@ -154,7 +211,7 @@ print_bake_config() {
     local bake_tmp
     bake_tmp="$(mktemp -d)"
     BAKE_CONFIG_FILE="${bake_tmp}/bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
-    docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true
+    docker buildx bake "${BAKE_FILES[@]}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true
     echo "Saved bake config to ${BAKE_CONFIG_FILE}"
     echo "--- :arrow_down: Uploading bake config to Buildkite"
     (cd "$(dirname "${BAKE_CONFIG_FILE}")" && buildkite-agent artifact upload "$(basename "${BAKE_CONFIG_FILE}")")
@@ -230,8 +287,10 @@ TARGET="test-ci"
 VLLM_BAKE_FILE_PATH="${VLLM_BAKE_FILE_PATH:-docker/docker-bake.hcl}"
 BUILDER_NAME="${BUILDER_NAME:-vllm-builder}"
 CI_HCL_URL="${CI_HCL_URL:-https://raw.githubusercontent.com/vllm-project/ci-infra/main/docker/ci.hcl}"
-CI_HCL_PATH="/tmp/ci.hcl"
+CI_HCL_PATH="${CI_HCL_PATH:-/tmp/ci.hcl}"
+ZSTD_HCL_PATH="${ZSTD_HCL_PATH:-.buildkite/image_build/zstd.hcl}"
 BUILDKIT_SOCKET="/run/buildkit/buildkitd.sock"
+BAKE_FILES=(-f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" -f "${ZSTD_HCL_PATH}")
 
 prepare_cache_tags
 ecr_login
@@ -288,14 +347,7 @@ if [[ ! -f "${VLLM_BAKE_FILE_PATH}" ]]; then
     exit 1
 fi
 
-echo "--- :arrow_down: Downloading ci.hcl"
-curl -sSfL -o "${CI_HCL_PATH}" "${CI_HCL_URL}"
-echo "Downloaded to ${CI_HCL_PATH}"
-
-if [[ ! -f "${CI_HCL_PATH}" ]]; then
-    echo "Error: ci.hcl not found at ${CI_HCL_PATH}"
-    exit 1
-fi
+download_ci_hcl || exit 1
 
 setup_buildx_builder
 
@@ -309,7 +361,7 @@ BUILD_TMP_DIR="$(mktemp -d)"
 trap 'rm -rf -- "${BUILD_TMP_DIR}"' EXIT
 BUILD_METADATA_FILE="${BUILD_TMP_DIR}/build-metadata.json"
 BUILD_STATUS=0
-docker --debug buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" \
+docker --debug buildx bake "${BAKE_FILES[@]}" \
     --progress plain --metadata-file "${BUILD_METADATA_FILE}" "${TARGET}" || BUILD_STATUS=$?
 
 record_buildkit_trace "${BUILD_METADATA_FILE}" || true
@@ -319,5 +371,6 @@ if [[ "${BUILD_STATUS}" -ne 0 ]]; then
 fi
 
 echo "--- :white_check_mark: Build complete"
+export_kernel_symbol_map || true
 
 annotate_image_tags
diff --git a/.buildkite/image_build/image_build_torch_nightly.sh b/.buildkite/image_build/image_build_torch_nightly.sh
index 8aa313d0e6d4..56c7116e14ac 100755
--- a/.buildkite/image_build/image_build_torch_nightly.sh
+++ b/.buildkite/image_build/image_build_torch_nightly.sh
@@ -49,7 +49,16 @@ echo "Image not found, proceeding with build..."
 # Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
 NIGHTLY_CUDA_VERSION="13.0.2"
 NIGHTLY_BUILD_BASE_IMAGE="pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359"
-NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
+# Must track ARG UBUNTU_VERSION in docker/Dockerfile. The nightly lane exists to
+# run the normal pipeline against nightly torch, so it has to differ from the
+# regular image in torch only -- not in the OS underneath it. This was pinned to
+# 22.04 while everything else moved to 24.04, and the runtime dependency list in
+# tools/install_snapshot_runtime.sh is 24.04-only (libgnutls30t64 is the 64-bit
+# time_t rename and does not exist on jammy), so the nightly build failed with
+# "E: Unable to locate package libgnutls30t64" while regular CI stayed green on
+# a cached layer.
+NIGHTLY_UBUNTU_VERSION="24.04"
+NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu${NIGHTLY_UBUNTU_VERSION}"
 
 echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})"
 docker buildx build --file docker/Dockerfile \
diff --git a/.buildkite/image_build/zstd.hcl b/.buildkite/image_build/zstd.hcl
new file mode 100644
index 000000000000..1f795c5f9150
--- /dev/null
+++ b/.buildkite/image_build/zstd.hcl
@@ -0,0 +1,3 @@
+target "test-ci" {
+  output = ["type=registry,compression=zstd,compression-level=3,force-compression=true,oci-mediatypes=true"]
+}
diff --git a/.buildkite/intel_jobs/basic_correctness_intel.yaml b/.buildkite/intel_jobs/basic_correctness_intel.yaml
index 292920c47ea1..173baf40f62c 100644
--- a/.buildkite/intel_jobs/basic_correctness_intel.yaml
+++ b/.buildkite/intel_jobs/basic_correctness_intel.yaml
@@ -3,6 +3,7 @@ depends_on:
   - image-build-xpu
 steps:
 - label: XPU Sleep Mode
+  key: xpu-sleep-mode
   timeout_in_minutes: 45
   device: intel_gpu
   agent_tags:
diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml
index d0aabd6d3299..98218f23b785 100644
--- a/.buildkite/intel_jobs/engine_intel.yaml
+++ b/.buildkite/intel_jobs/engine_intel.yaml
@@ -41,6 +41,7 @@ steps:
       pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py'
 
 - label: Engine (1 GPU)
+  key: engine-1-gpu
   timeout_in_minutes: 30
   device: intel_gpu
   agent_tags:
@@ -62,9 +63,10 @@ steps:
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'cd tests &&
       pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py &&
-      VLLM_XPU_ENABLE_XPU_GRAPH=1 pytest -v -s test_config.py'
+      pytest -v -s test_config.py'
 
 - label: V1 e2e (2 GPUs)
+  key: v1-e2e-2-gpus
   timeout_in_minutes: 30
   device: intel_gpu
   agent_tags:
diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml
index efed04c1b383..e8006bac8f66 100644
--- a/.buildkite/intel_jobs/entrypoints_intel.yaml
+++ b/.buildkite/intel_jobs/entrypoints_intel.yaml
@@ -134,7 +134,8 @@ steps:
     'cd tests &&
     pytest -v -s tool_use
     --deselect="tests/tool_use/mistral/test_mistral_tool_calls.py::test_tool_call_none_with_tools[ministral-3b]"
-    --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]" &&
+    --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]"
+    --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[llama3.2]" &&
     pytest -v -s entrypoints/tool_parsers &&
     pytest -v -s entrypoints/generate &&
     pytest -v -s entrypoints/anthropic &&
diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml
index e914c10d8bfe..da6086c1dc43 100644
--- a/.buildkite/intel_jobs/kernels_intel.yaml
+++ b/.buildkite/intel_jobs/kernels_intel.yaml
@@ -3,6 +3,7 @@ depends_on:
   - image-build-xpu
 steps:
 - label: vLLM IR Tests
+  key: vllm-ir-tests
   timeout_in_minutes: 30
   device: intel_gpu
   agent_tags:
diff --git a/.buildkite/intel_jobs/lm_eval_intel.yaml b/.buildkite/intel_jobs/lm_eval_intel.yaml
index 966ced8156d9..e1f649ad91c7 100644
--- a/.buildkite/intel_jobs/lm_eval_intel.yaml
+++ b/.buildkite/intel_jobs/lm_eval_intel.yaml
@@ -23,6 +23,30 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install "gpt-oss[eval]==0.0.5" &&
+      'uv pip install "gpt-oss[eval]==0.0.5" &&
       cd tests &&
       pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt'
+- label: LM Eval Large Models(2 GPUs)
+  key: lm-eval-large-models-2-gpus
+  timeout_in_minutes: 45
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
+  no_plugin: true
+  optional: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/
+    - .buildkite/lm-eval-harness/
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'cd .buildkite/lm-eval-harness &&
+       export VLLM_USE_DEEP_GEMM=0 &&
+      pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-xpu-fp8.txt --tp-size=2'
diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml
index cab0cb66aa45..2cecbf18f773 100644
--- a/.buildkite/intel_jobs/lora_intel.yaml
+++ b/.buildkite/intel_jobs/lora_intel.yaml
@@ -3,6 +3,7 @@ depends_on:
   - image-build-xpu
 steps:
 - label: LoRA Runtime + Utils
+  key: lora-runtime-utils
   timeout_in_minutes: 45
   device: intel_gpu
   agent_tags:
@@ -36,6 +37,7 @@ steps:
       pytest -v -s lora/test_worker.py'
 
 - label: LoRA Fused/MoE Kernels
+  key: lora-fused-moe-kernels
   timeout_in_minutes: 45
   device: intel_gpu
   agent_tags:
@@ -60,6 +62,7 @@ steps:
       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
+  key: lora-punica-kernels
   timeout_in_minutes: 45
   device: intel_gpu
   agent_tags:
@@ -86,6 +89,7 @@ steps:
       pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit'
 
 - label: LoRA Punica FP8/XPU Ops
+  key: lora-punica-fp8-xpu-ops
   timeout_in_minutes: 60
   device: intel_gpu
   agent_tags:
@@ -110,6 +114,7 @@ steps:
       pytest -v -s lora/test_punica_xpu_ops.py'
 
 - label: LoRA Models
+  key: lora-models
   timeout_in_minutes: 45
   device: intel_gpu
   agent_tags:
@@ -137,6 +142,7 @@ steps:
       pytest -s -v lora/test_minicpmv_tp.py'
 
 - label: LoRA Multimodal
+  key: lora-multimodal
   timeout_in_minutes: 45
   device: intel_gpu
   agent_tags:
diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml
index 752870abe36b..5fa4d833d7f0 100644
--- a/.buildkite/intel_jobs/misc_intel.yaml
+++ b/.buildkite/intel_jobs/misc_intel.yaml
@@ -2,7 +2,36 @@ group: Miscellaneous Intel
 depends_on:
   - image-build-xpu
 steps:
+- label: XPU Batch Invariance
+  key: xpu-batch-invariance
+  timeout_in_minutes: 60
+  num_devices: 2
+  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/
+    - tests/v1/determinism/test_xpu_batch_invariant_ut.py
+    - .buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh
+    - .buildkite/intel_jobs/misc_intel.yaml
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'export VLLM_BATCH_INVARIANT=1 &&
+      cd tests &&
+      pytest -v -s v1/determinism/test_xpu_batch_invariant_ut.py &&
+      bash ../.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh'
+
 - label: V1 Core + KV + Metrics
+  key: v1-core-kv-metrics-intel
   timeout_in_minutes: 45
   device: intel_gpu
   agent_tags:
@@ -32,9 +61,10 @@ steps:
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
       cd tests &&
       pytest -v -s v1/executor &&
-      VLLM_BATCH_INVARIANT=1 pytest -v -s -m 'not cpu_test' v1/ec_connector/unit'
+      VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit'
 
 - label: V1 Sample + Logits
+  key: v1-sample-logits-intel
   timeout_in_minutes: 90
   device: intel_gpu
   agent_tags:
@@ -67,8 +97,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install lm_eval[api]>=0.4.12 &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
+      'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
       cd tests &&
       pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
       pytest -v -s v1/test_oracle.py v1/test_request.py v1/test_outputs.py &&
@@ -80,6 +109,7 @@ steps:
   parallelism: 4
 
 - label: Basic Models Tests (Initialization)
+  key: basic-models-tests-initialization
   timeout_in_minutes: 60
   device: intel_gpu
   agent_tags:
@@ -104,6 +134,7 @@ steps:
       pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]'
 
 - label: XPU CPU Offload
+  key: xpu-cpu-offload
   timeout_in_minutes: 60
   device: intel_gpu
   agent_tags:
@@ -131,6 +162,7 @@ steps:
       pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
 
 - label: NixlConnector PD accuracy (4 GPUs)
+  key: nixlconnector-pd-accuracy-4-gpus
   timeout_in_minutes: 60
   num_devices: 4
   device: intel_gpu
@@ -188,8 +220,7 @@ steps:
   commands:
   - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install modelscope\<1.38 &&
-      cd tests &&
+      'cd tests &&
       pytest -v -s test_regression.py'
 
 - label: Metrics, Tracing (2 GPUs)
@@ -224,8 +255,7 @@ steps:
   commands:
   - >-
     bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-    'pip install opentelemetry-sdk\>=1.26.0 opentelemetry-api\>=1.26.0 opentelemetry-exporter-otlp\>=1.26.0 opentelemetry-semantic-conventions-ai\>=0.4.1 &&
-    cd tests &&
+    'cd tests &&
     pytest -v -s v1/tracing'
 
 - label: Async Engine, Inputs, Utils, Worker
@@ -263,12 +293,12 @@ steps:
   - >-
     bash .buildkite/scripts/hardware_ci/run-intel-test.sh
     'cd tests &&
-    pip install av &&
     pytest -v -s detokenizer &&
     pytest -v -s -m "not cpu_test" ./multimodal &&
     pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
 
 - label: Fusion Unit Tests
+  key: fusion-unit-tests
   timeout_in_minutes: 30
   device: intel_gpu
   agent_tags:
@@ -288,4 +318,4 @@ steps:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'cd tests &&
-      pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
\ No newline at end of file
+      pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml
index 14a853544eca..d7c683f8cea3 100644
--- a/.buildkite/intel_jobs/model_executor_intel.yaml
+++ b/.buildkite/intel_jobs/model_executor_intel.yaml
@@ -25,8 +25,6 @@ steps:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
       'apt-get update && apt-get install -y curl libsodium23 &&
-      pip3 install tensorizer==2.10.1 &&
-      pip3 install runai-model-streamer[s3,gcs,azure]\>=0.15.7 &&
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
       export PYTHONFAULTHANDLER=1 &&
       cd tests &&
diff --git a/.buildkite/intel_jobs/model_runner_v2_intel.yaml b/.buildkite/intel_jobs/model_runner_v2_intel.yaml
deleted file mode 100644
index 1f68333e976f..000000000000
--- a/.buildkite/intel_jobs/model_runner_v2_intel.yaml
+++ /dev/null
@@ -1,116 +0,0 @@
-group: Model Runner V2 Intel
-depends_on:
-  - image-build-xpu
-steps:
-- label: Model Runner V2 Core Tests (Intel)
-  timeout_in_minutes: 45
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 2+
-    mem: 24+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/v1/worker/gpu/
-    - vllm/v1/worker/gpu_worker.py
-    - vllm/v1/core/sched/
-    - vllm/v1/attention/
-    - tests/v1/engine/test_llm_engine.py
-    - tests/v1/e2e/
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
-      cd tests &&
-      pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" &&
-      pytest -v -s v1/e2e/general/test_context_length.py &&
-      ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" &&
-      pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" &&
-      pytest -v -s v1/e2e/general/test_min_tokens.py'
-
-- label: Model Runner V2 Examples (Intel)
-  timeout_in_minutes: 45
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 1+
-    mem: 24+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/v1/worker/gpu/
-    - vllm/v1/core/sched/
-    - vllm/v1/worker/gpu_worker.py
-    - examples/basic/offline_inference/
-    - examples/generate/multimodal/
-    - examples/features/
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
-      cd examples &&
-      python3 basic/offline_inference/chat.py &&
-      python3 basic/offline_inference/generate.py --model facebook/opt-125m &&
-      python3 generate/multimodal/vision_language_offline.py --seed 0 &&
-      python3 features/automatic_prefix_caching/prefix_caching_offline.py'
-
-- label: Model Runner V2 Distributed (2 GPUs)
-  timeout_in_minutes: 50
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 2+
-    mem: 16+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/v1/worker/gpu/
-    - vllm/v1/worker/gpu_worker.py
-    - tests/basic_correctness/test_basic_correctness.py
-    - tests/v1/distributed/test_async_llm_dp.py
-    - tests/v1/distributed/test_eagle_dp.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
-      cd tests &&
-      TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"'
-
-- label: Model Runner V2 Spec Decode
-  timeout_in_minutes: 50
-  device: intel_gpu
-  agent_tags:
-    label: production
-    gpu: 1+
-    mem: 24+
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - tests/v1/spec_decode/test_max_len.py
-  - tests/v1/spec_decode/test_rejection_sampler_utils.py
-  - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
-      cd tests &&
-      pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py'
diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml
index d99a7debb0b5..58ecb2d28209 100644
--- a/.buildkite/intel_jobs/models_multimodal_intel.yaml
+++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml
@@ -22,8 +22,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av &&
-      cd tests &&
+      'cd tests &&
       pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" &&
       pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model'
 
@@ -96,8 +95,7 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av &&
-      cd tests &&
+      'cd tests &&
       pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing'
 
 - label: Multi-Modal Processor # 44min
@@ -121,10 +119,9 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av matplotlib ftfy &&
-      pip install open-clip-torch --no-deps &&
-      cd tests &&
+      'cd tests &&
       pytest -v -s models/multimodal/processing/test_tensor_schema.py
+      --deselect="tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[deepseek-ai/DeepSeek-V4.1-Flash]"
       --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB'
   parallelism: 4
 
diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml
index 0bc25261a2b3..e15b20ea405e 100644
--- a/.buildkite/intel_jobs/test-intel.yaml
+++ b/.buildkite/intel_jobs/test-intel.yaml
@@ -15,9 +15,10 @@ steps:
         - exit_status: -10  # Agent was lost
           limit: 2
   - label: "XPU example Test"
+    key: xpu-example-test
     depends_on:
       - image-build-xpu
-    timeout_in_minutes: 50
+    timeout_in_minutes: 70
     device: intel_gpu
     agent_tags:
       label: production
@@ -33,8 +34,7 @@ steps:
     commands:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-        'pip install tblib==3.1.0 &&
-        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager &&
+        'python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --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 mp &&
         python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN &&
@@ -47,9 +47,10 @@ steps:
         python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
         python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 &&
         VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
-        python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
+        python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC -tp 2 --max-model-len 8192
         '
   - label: "XPU W8A8 FP8 Linear Examples"
+    key: xpu-w8a8-fp8-linear-examples
     depends_on:
       - image-build-xpu
     timeout_in_minutes: 60
@@ -77,6 +78,7 @@ steps:
         python3 examples/basic/offline_inference/generate.py --linear-backend torch --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8 --enforce-eager --max-model-len 4096
         '
   - label: "XPU V1 test"
+    key: xpu-v1-test
     depends_on:
       - image-build-xpu
     timeout_in_minutes: 70
@@ -95,7 +97,7 @@ steps:
     commands:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-        'cd tests &&
+        'cd tests && pip install dpctl &&
         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 --ignore=v1/engine/test_output_processor.py &&
         pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
@@ -103,10 +105,11 @@ steps:
         pytest -v -s v1/structured_output &&
         pytest -v -s v1/test_serial_utils.py &&
         pytest -v -s v1/e2e/general/test_correctness_sliding_window.py &&
-        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
-        pytest -v -s v1/kv_connector/unit --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 --ignore=v1/kv_connector/unit/test_offloading_connector.py'
+        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head &&
+        pytest -v -s v1/kv_connector/unit --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 --ignore=v1/kv_connector/unit/test_offloading_connector.py --deselect "v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]" --ignore=v1/kv_connector/unit/test_nixl_connector_hma.py'
     parallelism: 2
   - label: "XPU server test"
+    key: xpu-server-test
     depends_on:
       - image-build-xpu
     timeout_in_minutes: 45
@@ -125,11 +128,11 @@ steps:
     commands:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-        'pip install av &&
-        cd tests &&
+        'cd tests &&
         pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py &&
         pytest -v -s benchmarks/test_serve_cli.py'
   - label: "XPU quantization test"
+    key: xpu-quantization-test
     depends_on:
       - image-build-xpu
     timeout_in_minutes: 30
@@ -145,19 +148,14 @@ steps:
     source_file_dependencies:
       - vllm/
       - .buildkite/intel_jobs/test-intel.yaml 
-    # TODO: XPU fused MoE + model runner V2 currently hangs on some
-    # pre-quantized models, e.g.
-    # quantization/test_online.py::test_online_quantization_loads_real_weights[nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable]
-    # Remove the VLLM_USE_V2_MODEL_RUNNER=0 override below once the hang is
-    # root-caused and fixed.
     commands:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests &&
-        export VLLM_USE_V2_MODEL_RUNNER=0 &&
         pytest -v -s quantization/test_auto_round.py &&
         pytest -v -s quantization/test_online.py'
   - label: "XPU compressed tensors FP8 test"
+    key: xpu-compressed-tensors-fp8-test
     depends_on:
       - image-build-xpu
     timeout_in_minutes: 60
@@ -182,6 +180,7 @@ steps:
         pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8 &&
         pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8_block_enabled'
   - label: "XPU Graph"
+    key: xpu-graph
     depends_on:
       - image-build-xpu
     timeout_in_minutes: 60
@@ -197,16 +196,20 @@ steps:
       VLLM_TEST_DEVICE: "xpu"
     source_file_dependencies:
       - vllm/compilation/
+      - vllm/config/compilation.py
       - vllm/platforms/xpu.py
+      - vllm/v1/cudagraph_dispatcher.py
+      - vllm/v1/worker/encoder_cudagraph.py
+      - vllm/v1/worker/encoder_cudagraph_defs.py
       - vllm/v1/worker/xpu_model_runner.py
       - vllm/v1/worker/xpu_worker.py
+      - tests/v1/cudagraph
       - .buildkite/intel_jobs/test-intel.yaml
     commands:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-        'export VLLM_XPU_ENABLE_XPU_GRAPH=1 &&
-        python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-0.6B &&
-        python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-0.6B --kv-cache-dtype fp8 &&
-        python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-CT-AutoRound --max-model-len 4096'
-      # Temporarily disabled (multi-GPU/MXFP8 not yet supported with XPU Graph):
-      # python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP8-CT-AutoRound -tp 2 --max-model-len 4096
+        'cd tests &&
+        pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py &&
+        pytest -v -s v1/cudagraph/test_cudagraph_mode.py &&
+        pytest -v -s v1/cudagraph/test_breakable_cudagraph.py &&
+        pytest -v -s v1/cudagraph/test_encoder_cudagraph.py'
diff --git a/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8-xpu.yaml b/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8-xpu.yaml
new file mode 100644
index 000000000000..72ba8611953c
--- /dev/null
+++ b/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8-xpu.yaml
@@ -0,0 +1,16 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
+tasks:
+- name: "gsm8k"
+  metrics:
+  - name: "exact_match,strict-match"
+    value: 0.7142
+  - name: "exact_match,flexible-extract"
+    value: 0.4579
+limit: 256
+num_fewshot: 5
+max_model_len: 262144
+kv_cache_dtype: fp8
+enforce_eager: true
+apply_chat_template: true
+fewshot_as_multiturn: true
+trust_remote_code: true
diff --git a/.buildkite/lm-eval-harness/configs/models-large-xpu-fp8.txt b/.buildkite/lm-eval-harness/configs/models-large-xpu-fp8.txt
new file mode 100644
index 000000000000..73b28f062c3e
--- /dev/null
+++ b/.buildkite/lm-eval-harness/configs/models-large-xpu-fp8.txt
@@ -0,0 +1 @@
+NVIDIA-Nemotron-3-Nano-30B-A3B-FP8-xpu.yaml
\ No newline at end of file
diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
index fe9e332bdfa7..5576c84b05cb 100644
--- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
+++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-LM eval harness on model to compare vs HF baseline computed offline.
+"""LM eval harness on model to compare vs HF baseline computed offline.
 Configs are found in configs/$MODEL.yaml
 
 pytest -s -v test_lm_eval_correctness.py \
diff --git a/.buildkite/performance-benchmarks/scripts/compare-json-results.py b/.buildkite/performance-benchmarks/scripts/compare-json-results.py
index c9f8139fe62f..823c7488ec59 100644
--- a/.buildkite/performance-benchmarks/scripts/compare-json-results.py
+++ b/.buildkite/performance-benchmarks/scripts/compare-json-results.py
@@ -83,8 +83,7 @@ def compare_data_columns(
     drop_column: str,
     debug: bool = False,
 ):
-    """
-    Align concatenation by keys derived from info_cols instead of row order.
+    """Align concatenation by keys derived from info_cols instead of row order.
     - Pick one canonical key list: subset of info_cols present in ALL files.
     - For each file: set index to those keys, aggregate duplicates
       (mean for metric, first for names).
@@ -363,8 +362,7 @@ def _apply_two_decimals(
 # Export helpers (Excel + CSV)
 # -----------------------------
 def _sanitize_sheet_name(name: str) -> str:
-    """
-    Excel sheet constraints:
+    """Excel sheet constraints:
       - max 31 chars
       - cannot contain: : \ / ? * [ ]
       - cannot be empty
diff --git a/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py b/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py
index e694c38780cb..deb14a65e4a8 100644
--- a/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py
+++ b/.buildkite/performance-benchmarks/scripts/convert-results-json-to-markdown.py
@@ -93,8 +93,7 @@ def results_to_json(latency, throughput, serving):
 
 
 def get_size_with_unit(bytes, suffix="B"):
-    """
-    Scale bytes to its proper format
+    """Scale bytes to its proper format
     e.g:
         1253656 => '1.20MB'
         1253656678 => '1.17GB'
diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml
index 4af4eaf56ebc..2b90fef8201c 100644
--- a/.buildkite/release-pipeline.yaml
+++ b/.buildkite/release-pipeline.yaml
@@ -316,6 +316,66 @@ steps:
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
           - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" release-images'
 
+      - label: "Build release image - x86_64 - Rubin CUDA 13.4"
+        depends_on: ~
+        id: build-release-image-x86-rubin-cuda-13-4
+        if: build.env("NIGHTLY") == "1"
+        agents:
+          queue: cpu_queue_release
+        commands:
+          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+          - |
+            DOCKER_BUILDKIT=1 docker buildx build --load \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134) \
+              --platform linux/amd64 \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=13.4 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg INSTALL_RUBIN_PRERELEASE=true \
+              --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \
+              --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
+              --build-arg RUN_WHEEL_CHECK=false \
+              --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.4 \
+              --build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" release-images'
+
+      - label: "Build release image - aarch64 - Rubin CUDA 13.4"
+        depends_on: ~
+        id: build-release-image-arm64-rubin-cuda-13-4
+        if: build.env("NIGHTLY") == "1"
+        agents:
+          queue: arm64_cpu_queue_release
+        commands:
+          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+          - |
+            DOCKER_BUILDKIT=1 docker buildx build --load \
+              $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134) \
+              --platform linux/arm64 \
+              --build-arg max_jobs=16 \
+              --build-arg USE_SCCACHE=1 \
+              --build-arg GIT_REPO_CHECK=1 \
+              --build-arg CUDA_VERSION=13.4 \
+              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
+              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg INSTALL_RUBIN_PRERELEASE=true \
+              --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \
+              --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \
+              --build-arg RUN_WHEEL_CHECK=false \
+              --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.4 \
+              --build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \
+              --target vllm-openai \
+              --progress plain \
+              -f docker/Dockerfile .
+          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" release-images'
+
       - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
         depends_on: ~
         id: build-release-image-x86-ubuntu2404
@@ -513,6 +573,20 @@ steps:
           - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
           - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" release-manifests'
 
+      - label: "Create multi-arch manifest - Rubin CUDA 13.4"
+        depends_on:
+          - build-release-image-x86-rubin-cuda-13-4
+          - build-release-image-arm64-rubin-cuda-13-4
+        id: create-multi-arch-manifest-rubin-cuda-13-4
+        if: build.env("NIGHTLY") == "1"
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu134 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu134 --amend"
+          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: Rubin CUDA 13.4" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134" release-manifests'
+
       - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
         depends_on:
           - build-release-image-x86-ubuntu2404
@@ -585,6 +659,24 @@ steps:
           DOCKER_BUILDKIT: "1"
           DOCKERHUB_USERNAME: "vllmbot"
 
+      - label: "Publish nightly multi-arch Rubin image to DockerHub - CUDA 13.4"
+        depends_on:
+          - create-multi-arch-manifest-rubin-cuda-13-4
+        if: build.env("NIGHTLY") == "1"
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/push-nightly-builds.sh cu134"
+          # Clean up old nightly builds (keep only last 14)
+          - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu134-nightly-"
+        plugins:
+          - docker-login#v3.0.0:
+              username: vllmbot
+              password-env: DOCKERHUB_TOKEN
+        env:
+          DOCKER_BUILDKIT: "1"
+          DOCKERHUB_USERNAME: "vllmbot"
+
       - label: "Publish nightly CPU multi-arch image to DockerHub"
         depends_on:
           - build-cpu-release-image-x86
diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh
index 15f8274fc62b..0a9e6e7589fb 100644
--- a/.buildkite/scripts/ci-bake-rocm.sh
+++ b/.buildkite/scripts/ci-bake-rocm.sh
@@ -418,7 +418,7 @@ validate_ci_build_context_source() {
 
 describe_ci_revision() {
     git -C "$1" describe --tags --long --abbrev=10 \
-        --match '*[0-9]*' "$2" 2>/dev/null
+        --match 'v[0-9]*' "$2" 2>/dev/null
 }
 
 write_ci_git_archival_metadata() {
@@ -540,6 +540,26 @@ prepare_ci_build_context() {
     echo "Using canonical CI Docker context: ${ROCM_BUILD_CONTEXT_ROOT}"
 }
 
+configure_custom_rocm_stages() {
+    using_custom_rocm_dockerfiles || return 0
+    local context_root="${ROCM_BUILD_CONTEXT_ROOT:-.}"
+    local dockerfile="${context_root}/${CI_BASE_DOCKERFILE}"
+    local stages=""
+
+    (cd "${context_root}" && \
+        validate_rocm_dockerfile "${CI_BASE_DOCKERFILE}" \
+            ci_base test export_vllm export_test_smoke csrc-build rust-build) \
+        || return $?
+    # Hash all stages for custom layouts, including additional helper stages.
+    stages=$(rocm_dockerfile_stages "${dockerfile}" | paste -sd ' ')
+    CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${stages}}"
+    ROCM_CSRC_DOCKERFILE_STAGES="${ROCM_CSRC_DOCKERFILE_STAGES:-${stages}}"
+    ROCM_RUST_DOCKERFILE_STAGES="${ROCM_RUST_DOCKERFILE_STAGES:-${stages}}"
+    if [[ " ${stages} " != *" build_rocshmem "* ]]; then
+        unset ROCSHMEM_BRANCH ROCSHMEM_CACHE_KEY DEEPEP_CACHE_KEY
+    fi
+}
+
 compose_dependency_cache_key() {
     local prefix="$1"
     local material="$2"
@@ -569,11 +589,11 @@ hash_dockerfile_stages() {
             }
             emit = 0
         }
-        $1 == "FROM" {
+        toupper($1) == "FROM" {
             stage = ""
             for (idx = 1; idx <= NF; idx++) {
                 if (tolower($idx) == "as" && idx < NF) {
-                    stage = $(idx + 1)
+                    stage = tolower($(idx + 1))
                 }
             }
             emit = (stage in wanted)
@@ -613,11 +633,11 @@ discover_dockerfile_stage_args() {
         }
         {
             line = $0
-            if ($1 == "FROM") {
+            if (toupper($1) == "FROM") {
                 stage = ""
                 for (idx = 1; idx <= NF; idx++) {
                     if (tolower($idx) == "as" && idx < NF) {
-                        stage = $(idx + 1)
+                        stage = tolower($(idx + 1))
                     }
                 }
                 emit = (stage in wanted)
@@ -630,7 +650,7 @@ discover_dockerfile_stage_args() {
             for (idx = 1; idx <= line_count; idx++) {
                 line = lines[idx]
                 arg_name = line
-                sub(/^[[:space:]]*ARG[[:space:]]+/, "", arg_name)
+                sub(/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+/, "", arg_name)
                 if (arg_name != line) {
                     sub(/[=[:space:]].*/, "", arg_name)
                     if (arg_name ~ /^[A-Za-z_][A-Za-z0-9_]*$/) {
@@ -716,7 +736,7 @@ extract_dockerfile_arg_default() {
     if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then
         physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}"
     fi
-    sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
+    sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
         "${physical_dockerfile}" | head -1
 }
 
@@ -1000,7 +1020,21 @@ create_and_bootstrap_builder() {
 }
 
 init_config() {
+    # shellcheck source=.buildkite/scripts/rocm/build-config.sh
+    source "$(dirname "${BASH_SOURCE[0]}")/rocm/build-config.sh"
+    configure_rocm_build
+
     TARGET="${1:-test-ci}"
+    if using_custom_rocm_dockerfiles; then
+        if [[ -z "${BASE_IMAGE:-}" ]]; then
+            echo "Custom ROCm builds require BASE_IMAGE from the selected base handoff" >&2
+            return 1
+        fi
+        if ! is_ci_base_target && [[ -z "${CI_BASE_IMAGE:-}" ]]; then
+            echo "Custom ROCm builds require CI_BASE_IMAGE from the selected ci_base handoff" >&2
+            return 1
+        fi
+    fi
     BAKE_TARGETS=("${TARGET}")
     DEPENDENCY_CACHE_TARGETS=()
     CI_HCL_SOURCE="${CI_HCL_SOURCE:-${CI_HCL_FILE:-${DEFAULT_CI_HCL_SOURCE}}}"
@@ -1010,10 +1044,12 @@ init_config() {
     PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950}"
     CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}"
     CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}"
-    CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}"
+    if ! using_custom_rocm_dockerfiles; then
+        CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}"
+    fi
     CI_BASE_METADATA_VERSION="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}"
     CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}"
-    export PYTORCH_ROCM_ARCH
+    export PYTORCH_ROCM_ARCH CI_BASE_DOCKERFILE
 
     SCRIPT_TMP_DIR=$(mktemp -d -t ci-bake-rocm.XXXXXX)
     CI_HCL_PATH="${SCRIPT_TMP_DIR}/ci.hcl"
@@ -1042,6 +1078,11 @@ print_header() {
 }
 
 validate_inputs() {
+    if using_custom_rocm_dockerfiles; then
+        validate_rocm_dockerfile "${ROCM_BASE_DOCKERFILE}" || return $?
+        validate_rocm_dockerfile "${CI_BASE_DOCKERFILE}" \
+            ci_base test export_vllm export_test_smoke csrc-build rust-build || return $?
+    fi
     if [[ ! -f "${VLLM_BAKE_FILE}" ]]; then
         echo "Error: vLLM bake file not found at ${VLLM_BAKE_FILE}"
         echo "Make sure you're running from the vLLM repository root"
@@ -1069,6 +1110,11 @@ load_ci_hcl() {
 
 init_bake_files() {
     BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}")
+    if using_custom_rocm_dockerfiles; then
+        # Custom stacks reuse local cache without replacing the
+        # standard ROCm registry caches for the same commit or branch.
+        BAKE_FILES+=(--set '*.cache-to=')
+    fi
 }
 
 compute_ci_base_hash_if_needed() {
@@ -1804,7 +1850,7 @@ compute_rocm_csrc_content_hash() {
     local -a content_args=()
 
     bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
+    dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}"
     read -r -a content_paths <<< "${content_files}"
     mapfile -t content_args < <(
         get_content_arg_names "${dockerfile_rocm}" "${stages}" "${ROCM_CSRC_CONTENT_ARGS:-}"
@@ -1853,7 +1899,7 @@ compute_rocm_rust_content_hash() {
     local -a content_args=()
 
     bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
+    dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}"
     read -r -a content_paths <<< "${content_files}"
     mapfile -t content_args < <(
         get_content_arg_names \
@@ -1943,7 +1989,7 @@ write_rocm_build_arg_override() {
     local arg_value=""
 
     bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
+    dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}"
     mapfile -t arg_names < <(
         {
             get_content_arg_names \
@@ -2259,7 +2305,7 @@ extract_dependency_pins() {
     local val=""
 
     bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
+    dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}"
     physical_dockerfile="${dockerfile_rocm}"
     if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile_rocm}" != /* ]]; then
         physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile_rocm}"
@@ -2275,12 +2321,12 @@ extract_dependency_pins() {
         fi
 
         val=$(
-            sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
+            sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
                 "${physical_dockerfile}" | head -1
         )
         if [[ -n "${val}" ]]; then
             export "${var}=${val}"
-            echo "Extracted ${var}=${val} from Dockerfile.rocm"
+            echo "Extracted ${var}=${val} from ${dockerfile_rocm}"
         fi
     done
 }
@@ -2297,7 +2343,7 @@ compute_dependency_cache_keys() {
     local deepep_material=""
 
     bake_dir=$(dirname "${VLLM_BAKE_FILE}")
-    dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
+    dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}"
     nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
     ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
     rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
@@ -2687,6 +2733,7 @@ main() {
     init_bake_files
     if is_ci_base_target; then
         prepare_ci_build_context
+        configure_custom_rocm_stages
     fi
     compute_ci_base_hash_if_needed
     configure_ci_base_image_refs
@@ -2697,6 +2744,7 @@ main() {
     # version metadata only after that lookup sees the available tag history.
     if ! is_ci_base_target; then
         prepare_ci_build_context
+        configure_custom_rocm_stages
     fi
     extract_dependency_pins
     write_rocm_build_arg_override
diff --git a/.buildkite/scripts/ci-clean-log.sh b/.buildkite/scripts/ci-clean-log.sh
index e2e21483d546..2fe45f3b7745 100644
--- a/.buildkite/scripts/ci-clean-log.sh
+++ b/.buildkite/scripts/ci-clean-log.sh
@@ -10,11 +10,20 @@ fi
 
 INPUT_FILE="$1"
 
+# GNU sed uses "-i [-r]" while BSD/macOS sed needs "-i '' [-E]"; detect once.
+if sed --version >/dev/null 2>&1; then
+    sedi() { sed -i "$@"; }
+    SED_ERE=-r
+else
+    sedi() { sed -i '' "$@"; }
+    SED_ERE=-E
+fi
+
 # Strip timestamps
-sed -i 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE"
+sedi 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE"
 
 # Strip Buildkite inline timestamp markers (ESC _bk;t=<ms> BEL)
-sed -i 's/\x1B_bk;t=[0-9]*\x07//g' "$INPUT_FILE"
+sedi 's/\x1B_bk;t=[0-9]*\x07//g' "$INPUT_FILE"
 
 # Strip colorization
-sed -i -r 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE"
+sedi "$SED_ERE" 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE"
diff --git a/.buildkite/scripts/ci-otel/README.md b/.buildkite/scripts/ci-otel/README.md
new file mode 100644
index 000000000000..318cec5f2169
--- /dev/null
+++ b/.buildkite/scripts/ci-otel/README.md
@@ -0,0 +1,51 @@
+# CI timing and GPU samples
+
+The pipeline generator activates these helpers for trusted, traced CI commands.
+`ci_otel.sh` records command spans and `ci_pytest.sh` injects the pytest timing
+plugin. `ci_gpu.py` runs once per command, outside pytest workers, when
+`nvidia-smi` is available. It queries up to 16 container-visible NVIDIA devices
+once per second without importing PyTorch or creating a CUDA context.
+Set `CI_INFRA_GPU_SAMPLING=0` in the job environment to disable sampling.
+
+Samples contain device UUID, index, name, utilization percent, and used/total
+device memory in bytes. They are `ci.gpu.sample` OTLP events inside
+`ci.gpu.samples` spans, parented to the command span and authenticated using the
+existing job-scoped Buildkite OIDC identity. Batches upload every 30 seconds;
+the final batch is spooled for the shell's normal job-end upload. No host agent
+or long-lived dashboard token is needed. The dashboard companion change adds
+GPU controls alongside jobs, commands, and pytest rows in the Builds timeline.
+
+Device queries and periodic uploads each have a two-second timeout. Three
+consecutive query failures stop sampling. Failed periodic uploads are dropped;
+memory stays bounded, and the dashboard shows the missing intervals as gaps.
+Uploads can delay the next poll. The shell stops and joins the sampler before
+flushing traces, allows at most three seconds for shutdown, and preserves the
+original command exit status. Samples already uploaded survive an interrupted
+job; the last unflushed batch may be lost on forced termination.
+
+These are device readings during each test's time interval, not process-level
+kernel attribution. Other processes and overlapping tests can contribute to
+the readings. Subsecond tests may have no sample. Missing readings never become
+zero utilization. MIG parent-device memory and utilization are omitted because
+they cannot be attributed to the job's partition. When the environment identifies
+an assigned MIG UUID, the collector queries that instance's used/total memory
+through vLLM's standalone bundled NVML binding in a bounded subprocess.
+The binding is loaded directly to avoid vLLM/Torch import side effects. Numeric
+`CUDA_VISIBLE_DEVICES` indices are resolved against explicit MIG UUIDs in
+`NVIDIA_VISIBLE_DEVICES`. For CDI assignments without an explicit UUID, a bounded
+helper initializes the CUDA driver and reads only CUDA-visible device UUIDs,
+without creating CUDA contexts or importing Torch. NVML verifies each MIG handle.
+Discovery or permission failures remain unknown; NVML sibling instances are never
+enumerated. MIG utilization stays absent.
+CPU jobs and AMD mirrors do not collect GPU samples.
+
+Run the helper tests without loading vLLM or requiring a GPU:
+
+```bash
+PATH="$PWD/.venv/bin:$PATH" .venv/bin/python -m pytest \
+  .buildkite/scripts/ci-otel/tests/test_ci_otel.py -q
+```
+
+The suite uses a fake `nvidia-smi` to exercise the shell lifecycle, device-query
+failure handling, disable switch, event encoding, and bounded uploads. Real
+hardware sampling overhead still needs a canary run before broad deployment.
diff --git a/.buildkite/scripts/ci-otel/ci_gpu.py b/.buildkite/scripts/ci-otel/ci_gpu.py
new file mode 100644
index 000000000000..0c77f5170b19
--- /dev/null
+++ b/.buildkite/scripts/ci-otel/ci_gpu.py
@@ -0,0 +1,303 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Best-effort, container-visible NVIDIA device samples for command timelines."""
+
+from __future__ import annotations
+
+import csv
+import ctypes
+import importlib.util
+import io
+import json
+import math
+import os
+import secrets
+import signal
+import subprocess
+import sys
+import threading
+import time
+import uuid as uuid_module
+from pathlib import Path
+
+from ci_otel import Span, export_spans, record_spans
+
+INTERVAL_SECONDS = 1.0
+BATCH_SECONDS = 30
+MAX_DEVICES = 16
+QUERY_TIMEOUT_SECONDS = 2
+QUERY = "index,uuid,name,utilization.gpu,memory.used,memory.total,mig.mode.current"
+
+
+def _number(value: str, scale: int = 1, maximum: int | None = None) -> int | None:
+    try:
+        number = float(value)
+        if not math.isfinite(number) or number < 0:
+            return None
+        if maximum is not None and number > maximum:
+            return None
+        return round(number * scale)
+    except ValueError:
+        return None
+
+
+def parse_samples(output: str, timestamp_ns: int) -> list[dict]:
+    events = []
+    for row in csv.reader(io.StringIO(output)):
+        if len(row) != 7 or len(events) >= MAX_DEVICES:
+            continue
+        index, uuid, name, util, used, total, mig = (value.strip() for value in row)
+        if not index.isdigit() or not uuid.startswith("GPU-"):
+            continue
+        attributes: dict[str, str | int | bool] = {
+            "gpu.uuid": uuid,
+            "gpu.index": int(index),
+            "gpu.name": name,
+        }
+        # Parent-device memory on a MIG GPU is not the job's allocated memory.
+        if mig.lower() == "enabled":
+            attributes["gpu.status"] = "unsupported_mig"
+        else:
+            for key, value in (
+                ("gpu.utilization", _number(util, maximum=100)),
+                ("gpu.memory.used", _number(used, 1024**2)),
+                ("gpu.memory.total", _number(total, 1024**2)),
+            ):
+                if value is not None:
+                    attributes[key] = value
+        events.append(
+            {"time_ns": timestamp_ns, "name": "ci.gpu.sample", "attributes": attributes}
+        )
+    return events
+
+
+def visible_mig_devices() -> list[str]:
+    """Resolve explicit MIG assignments without enumerating sibling instances."""
+    nvidia = os.environ.get("NVIDIA_VISIBLE_DEVICES", "").split(",")
+    nvidia = [value.strip() for value in nvidia]
+    cuda = os.environ.get("CUDA_VISIBLE_DEVICES")
+    devices = nvidia if cuda is None else [value.strip() for value in cuda.split(",")]
+    if devices and all(value.isdigit() for value in devices):
+        if not all(value.startswith("MIG-") for value in nvidia):
+            return []
+        try:
+            devices = [nvidia[int(value)] for value in devices]
+        except IndexError:
+            return []
+    if not devices or not all(value.startswith("MIG-") for value in devices):
+        return []
+    return list(dict.fromkeys(devices))[:MAX_DEVICES]
+
+
+def load_nvml():
+    """Load the standalone bundled binding without importing vLLM or Torch."""
+    # Test images install vLLM as a wheel and keep only tests/helpers in the
+    # checkout. Finding the top-level package does not execute its __init__.
+    package = importlib.util.find_spec("vllm")
+    if package is None or package.origin is None:
+        raise ImportError("Installed vLLM package unavailable")
+    path = Path(package.origin).parent / "third_party/pynvml.py"
+    spec = importlib.util.spec_from_file_location("ci_nvml", path)
+    if spec is None or spec.loader is None:
+        raise ImportError("Bundled NVML binding unavailable")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def cuda_visible_mig_devices() -> list[str]:
+    """Discover CDI assignments in a subprocess without creating CUDA contexts."""
+    pynvml = load_nvml()
+
+    driver = ctypes.CDLL("libcuda.so.1")
+    count = ctypes.c_int()
+    if driver.cuInit(0) or driver.cuDeviceGetCount(ctypes.byref(count)):
+        return []
+    devices = []
+    pynvml.nvmlInit()
+    try:
+        for index in range(min(count.value, MAX_DEVICES)):
+            device = ctypes.c_int()
+            identifier = (ctypes.c_ubyte * 16)()
+            if driver.cuDeviceGet(ctypes.byref(device), index):
+                continue
+            # The v2 entry point returns the compute instance UUID under MIG.
+            if driver.cuDeviceGetUuid_v2(ctypes.byref(identifier), device):
+                continue
+            uuid = "MIG-" + str(uuid_module.UUID(bytes=bytes(identifier)))
+            try:
+                handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
+                if pynvml.nvmlDeviceIsMigDeviceHandle(handle):
+                    devices.append(uuid)
+            except pynvml.NVMLError:
+                continue
+    finally:
+        pynvml.nvmlShutdown()
+    return devices
+
+
+def query_mig_samples(devices: list[str], timestamp_ns: int) -> list[dict]:
+    # Use the binding already shipped with vLLM; never initialize CUDA/Torch.
+    pynvml = load_nvml()
+
+    events = []
+    pynvml.nvmlInit()
+    try:
+        for index, uuid in enumerate(devices[:MAX_DEVICES]):
+            if not uuid.startswith("MIG-"):
+                continue
+            attributes: dict[str, str | int | bool] = {
+                "gpu.uuid": uuid,
+                "gpu.index": index,
+                "gpu.name": "NVIDIA MIG",
+                "gpu.status": "unavailable_mig_memory",
+            }
+            try:
+                handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
+                if not pynvml.nvmlDeviceIsMigDeviceHandle(handle):
+                    continue
+                memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                used = _number(str(memory.used))
+                total = _number(str(memory.total))
+                if used is not None and total and used <= total:
+                    attributes.update(
+                        {
+                            "gpu.memory.used": used,
+                            "gpu.memory.total": total,
+                            "gpu.status": "mig_memory_only",
+                        }
+                    )
+                # NVML does not support per-MIG utilization on these devices.
+            except pynvml.NVMLError:
+                pass
+            events.append(
+                {
+                    "time_ns": timestamp_ns,
+                    "name": "ci.gpu.sample",
+                    "attributes": attributes,
+                }
+            )
+    finally:
+        pynvml.nvmlShutdown()
+    return events
+
+
+def sample_span(events: list[dict]) -> Span:
+    return Span(
+        trace_id=os.environ["CI_INFRA_TRACE_ID"],
+        span_id=secrets.token_hex(8),
+        parent_span_id=os.environ["CI_INFRA_COMMAND_SPAN_ID"],
+        name="ci.gpu.samples",
+        start_ns=events[0]["time_ns"],
+        end_ns=events[-1]["time_ns"],
+        attributes={
+            "ci.span.kind": "gpu-samples",
+            "gpu.sample.interval_ms": int(INTERVAL_SECONDS * 1000),
+            "gpu.scope": "container-visible-devices",
+        },
+        events=events,
+    )
+
+
+def collect(parent_pid: int, stop: threading.Event) -> None:
+    events: list[dict] = []
+    batch_started = time.monotonic()
+    failures = 0
+    discovery_attempted = False
+    mig_devices = visible_mig_devices()
+    query = (
+        [sys.executable, __file__, "--query-mig", *mig_devices]
+        if mig_devices
+        else [
+            "nvidia-smi",
+            f"--query-gpu={QUERY}",
+            "--format=csv,noheader,nounits",
+        ]
+    )
+    try:
+        while not stop.is_set() and os.getppid() == parent_pid:
+            started = time.monotonic()
+            try:
+                result = subprocess.run(
+                    query,
+                    capture_output=True,
+                    text=True,
+                    timeout=QUERY_TIMEOUT_SECONDS,
+                    check=True,
+                )
+                if stop.is_set():
+                    break
+                samples = (
+                    json.loads(result.stdout)
+                    if mig_devices
+                    else parse_samples(result.stdout, time.time_ns())
+                )
+                if (
+                    not mig_devices
+                    and not discovery_attempted
+                    and any(
+                        sample["attributes"].get("gpu.status") == "unsupported_mig"
+                        for sample in samples
+                    )
+                ):
+                    discovery_attempted = True
+                    # CDI can assign a MIG device without exporting its UUID.
+                    # Only CUDA-visible devices are inspected, never NVML siblings.
+                    try:
+                        discovered = subprocess.run(
+                            [sys.executable, __file__, "--discover-mig"],
+                            capture_output=True,
+                            text=True,
+                            timeout=QUERY_TIMEOUT_SECONDS,
+                            check=True,
+                        )
+                        mig_devices = json.loads(discovered.stdout)
+                    except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
+                        mig_devices = []
+                    if mig_devices:
+                        query = [sys.executable, __file__, "--query-mig", *mig_devices]
+                        continue  # Do not create an extra parent-device series.
+                if not samples:
+                    break
+                events.extend(samples)
+                failures = 0
+            except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
+                failures += 1
+                if failures >= 3:
+                    print(
+                        "CI GPU sampling stopped: device query unavailable",
+                        file=sys.stderr,
+                    )
+                    break
+            if events and time.monotonic() - batch_started >= BATCH_SECONDS:
+                # Bounded memory and upload time even if the receiver is down.
+                # Dropped batches appear as gaps, never as zero utilization.
+                export_spans([sample_span(events)], timeout_seconds=2)
+                events = []
+                batch_started = time.monotonic()
+            stop.wait(max(0, INTERVAL_SECONDS - (time.monotonic() - started)))
+    finally:
+        # The shell joins this process before the ordinary job-end spool flush.
+        if events:
+            record_spans([sample_span(events)])
+
+
+def main() -> None:
+    stop = threading.Event()
+    signal.signal(signal.SIGTERM, lambda *_: stop.set())
+    signal.signal(signal.SIGINT, lambda *_: stop.set())
+    try:
+        collect(int(sys.argv[1]), stop)
+    except Exception:
+        print("CI GPU sampling skipped", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1 and sys.argv[1] == "--discover-mig":
+        print(json.dumps(cuda_visible_mig_devices()))
+    elif len(sys.argv) > 1 and sys.argv[1] == "--query-mig":
+        print(json.dumps(query_mig_samples(sys.argv[2:], time.time_ns())))
+    else:
+        main()
diff --git a/.buildkite/scripts/ci-otel/ci_otel.py b/.buildkite/scripts/ci-otel/ci_otel.py
index 0d25594bec19..cec3f4130986 100644
--- a/.buildkite/scripts/ci-otel/ci_otel.py
+++ b/.buildkite/scripts/ci-otel/ci_otel.py
@@ -14,12 +14,13 @@
 import time
 import urllib.request
 from collections.abc import Iterable
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 from pathlib import Path
 
 ENDPOINT = os.getenv("CI_INFRA_OTEL_ENDPOINT", "https://ci.vllm.ai/api/otel/v1/traces")
 AUDIENCE = os.getenv("CI_INFRA_OTEL_AUDIENCE", "https://ci.vllm.ai/api/otel")
 MAX_BATCH_SIZE = 2_000
+MAX_BATCH_BYTES = 1024 * 1024
 TEST_SPOOL_BATCH_SIZE = 100
 
 
@@ -56,6 +57,7 @@ class Span:
     end_ns: int
     attributes: dict[str, str | int | bool]
     status_code: int = 1
+    events: list[dict] = field(default_factory=list)
 
 
 def _varint(value: int) -> bytes:
@@ -137,6 +139,12 @@ def _encode_span(span: Span) -> bytes:
         )
     for name, value in attributes.items():
         encoded += _bytes_field(9, _attribute(name, value))
+    for event in span.events:
+        encoded_event = _fixed64_field(1, int(event["time_ns"]))
+        encoded_event += _string_field(2, event["name"])
+        for name, value in event["attributes"].items():
+            encoded_event += _bytes_field(3, _attribute(name, value))
+        encoded += _bytes_field(11, encoded_event)
     return encoded + _bytes_field(15, _varint_field(3, span.status_code))
 
 
@@ -482,10 +490,24 @@ def export_spans(spans: list[Span], timeout_seconds: float = 3.0) -> bool:
     try:
         deadline = time.monotonic() + max(timeout_seconds, 0.1)
         token = _oidc_token(deadline)
-        for offset in range(0, len(spans), MAX_BATCH_SIZE):
+        batches: list[list[Span]] = []
+        batch: list[Span] = []
+        size = 0
+        for span in spans:
+            span_size = len(_encode_span(span))
+            if batch and (
+                size + span_size > MAX_BATCH_BYTES or len(batch) >= MAX_BATCH_SIZE
+            ):
+                batches.append(batch)
+                batch, size = [], 0
+            batch.append(span)
+            size += span_size
+        if batch:
+            batches.append(batch)
+        for batch in batches:
             request = urllib.request.Request(
                 ENDPOINT,
-                data=encode_request(spans[offset : offset + MAX_BATCH_SIZE]),
+                data=encode_request(batch),
                 method="POST",
                 headers={
                     "Authorization": f"Bearer {token}",
diff --git a/.buildkite/scripts/ci-otel/ci_otel.sh b/.buildkite/scripts/ci-otel/ci_otel.sh
index feec16fa6528..7f1b700fe6bd 100755
--- a/.buildkite/scripts/ci-otel/ci_otel.sh
+++ b/.buildkite/scripts/ci-otel/ci_otel.sh
@@ -26,7 +26,8 @@ if [ -z "${CI_INFRA_OTEL_RUNTIME_DIR}" ] ||
   [ -z "${_CI_INFRA_OTEL_PYTHON}" ] ||
   [ ! -f "${CI_INFRA_OTEL_DIR}/ci_otel.py" ] ||
   ! mkdir -p "${CI_INFRA_OTEL_SPOOL_DIR}" ||
-  ! PYTHONPATH="${CI_INFRA_OTEL_DIR}${PYTHONPATH:+:${PYTHONPATH}}" \
+  ! PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONPATH="${CI_INFRA_OTEL_DIR}${PYTHONPATH:+:${PYTHONPATH}}" \
     "${_CI_INFRA_OTEL_PYTHON}" -c "import ci_otel" >/dev/null 2>&1; then
   echo "vLLM CI OTel: tracing unavailable; test command will run normally" >&2 || :
   if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then
@@ -58,12 +59,45 @@ fi
 
 _ci_otel_python() {
   if command -v timeout >/dev/null 2>&1; then
-    timeout 2s "${_CI_INFRA_OTEL_PYTHON}" "$@"
+    PYTHONDONTWRITEBYTECODE=1 timeout 2s "${_CI_INFRA_OTEL_PYTHON}" "$@"
   else
-    "${_CI_INFRA_OTEL_PYTHON}" "$@"
+    PYTHONDONTWRITEBYTECODE=1 "${_CI_INFRA_OTEL_PYTHON}" "$@"
   fi
 }
 
+# The Docker test commands run as root, while /workdir is owned by the
+# Buildkite checkout user. Run the GPU sampler as that unprivileged identity so
+# any future project-code import cannot leave root-owned files in the checkout.
+_ci_otel_prepare_gpu_user() {
+  _CI_INFRA_GPU_UID=""
+  _CI_INFRA_GPU_GID=""
+  [ "$(id -u 2>/dev/null || :)" = "0" ] || return 0
+  command -v setpriv >/dev/null 2>&1 || return 1
+
+  _CI_INFRA_OTEL_WORKSPACE_DIR="${CI_INFRA_OTEL_WORKSPACE_DIR:-/workdir}"
+  [ -d "${_CI_INFRA_OTEL_WORKSPACE_DIR}" ] || return 1
+  # The owner fields are numeric and whitespace-free on GNU stat.
+  # shellcheck disable=SC2046
+  set -- $(stat -c '%u %g' -- "${_CI_INFRA_OTEL_WORKSPACE_DIR}" 2>/dev/null || :)
+  [ "$#" -eq 2 ] || return 1
+  case "$1:$2" in
+    *[!0-9:]* | 0:* | *:0) return 1 ;;
+  esac
+
+  _CI_INFRA_GPU_UID="$1"
+  _CI_INFRA_GPU_GID="$2"
+  chown "${_CI_INFRA_GPU_UID}:${_CI_INFRA_GPU_GID}" \
+    "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1
+  chmod 0700 "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1
+  if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then
+    # Permit traversal to the private child spool, but not directory listing.
+    chmod 0711 "${CI_INFRA_OTEL_RUNTIME_DIR}" || return 1
+  fi
+  setpriv --reuid="${_CI_INFRA_GPU_UID}" \
+    --regid="${_CI_INFRA_GPU_GID}" --clear-groups --no-new-privs \
+    -- test -w "${CI_INFRA_OTEL_SPOOL_DIR}"
+}
+
 ci_otel_start() {
   _CI_INFRA_OTEL_COMMAND_INDEX="$1"
   _CI_INFRA_OTEL_COMMAND_LABEL="${2:-command ${_CI_INFRA_OTEL_COMMAND_INDEX}}"
@@ -82,6 +116,28 @@ ci_otel_start() {
   CI_INFRA_COMMAND_SPAN_ID="${_CI_INFRA_OTEL_COMMAND_SPAN_ID}"
   export CI_INFRA_TRACE_ID CI_INFRA_COMMAND_SPAN_ID
   _CI_INFRA_OTEL_ACTIVE=1
+  _CI_INFRA_GPU_PID=""
+  if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] &&
+    [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] &&
+    command -v nvidia-smi >/dev/null 2>&1; then
+    if _ci_otel_prepare_gpu_user; then
+      if [ -n "${_CI_INFRA_GPU_UID}" ]; then
+        PYTHONDONTWRITEBYTECODE=1 setpriv \
+          --reuid="${_CI_INFRA_GPU_UID}" --regid="${_CI_INFRA_GPU_GID}" \
+          --clear-groups --no-new-privs -- \
+          "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" \
+          </dev/null &
+      else
+        PYTHONDONTWRITEBYTECODE=1 \
+          "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" \
+          </dev/null &
+      fi
+      _CI_INFRA_GPU_PID=$!
+    else
+      echo "vLLM CI OTel: non-root GPU sampler unavailable; sampling disabled" \
+        >&2 || :
+    fi
+  fi
 }
 
 ci_otel_finish() {
@@ -93,6 +149,18 @@ ci_otel_finish() {
     "${_CI_INFRA_OTEL_PARENT_SPAN_ID}" "${_CI_INFRA_OTEL_START_NS}" \
     "${_CI_INFRA_OTEL_COMMAND_INDEX}" "${_CI_INFRA_OTEL_COMMAND_STATUS}" \
     "${_CI_INFRA_OTEL_COMMAND_LABEL}" || :
+  if [ -n "${_CI_INFRA_GPU_PID:-}" ]; then
+    kill -TERM "${_CI_INFRA_GPU_PID}" 2>/dev/null || :
+    _CI_INFRA_GPU_WAIT=0
+    while kill -0 "${_CI_INFRA_GPU_PID}" 2>/dev/null &&
+      [ "${_CI_INFRA_GPU_WAIT}" -lt 30 ]; do
+      sleep 0.1
+      _CI_INFRA_GPU_WAIT=$((_CI_INFRA_GPU_WAIT + 1))
+    done
+    kill -KILL "${_CI_INFRA_GPU_PID}" 2>/dev/null || :
+    wait "${_CI_INFRA_GPU_PID}" 2>/dev/null || :
+    _CI_INFRA_GPU_PID=""
+  fi
   CI_INFRA_TRACE_ID=""
   CI_INFRA_COMMAND_SPAN_ID=""
   export CI_INFRA_TRACE_ID CI_INFRA_COMMAND_SPAN_ID
@@ -118,9 +186,11 @@ _ci_otel_on_exit() {
   trap - 0
   ci_otel_finish "${_CI_INFRA_OTEL_EXIT_STATUS}" || :
   if command -v timeout >/dev/null 2>&1; then
-    timeout 4s "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || :
+    PYTHONDONTWRITEBYTECODE=1 timeout 4s \
+      "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || :
   else
-    "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || :
+    PYTHONDONTWRITEBYTECODE=1 \
+      "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || :
   fi
   if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then
     rm -rf -- "${CI_INFRA_OTEL_RUNTIME_DIR}" || :
diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py
index ff279de953a5..c66530989f97 100644
--- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py
+++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py
@@ -6,16 +6,20 @@
 import json
 import os
 import shlex
+import shutil
 import subprocess
 import sys
+import threading
 import time
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 
 SCRIPTS_DIR = Path(__file__).resolve().parents[1]
 sys.path.insert(0, str(SCRIPTS_DIR))
 
+import ci_gpu  # noqa: E402
 import ci_otel  # noqa: E402
 from ci_otel import Span  # noqa: E402
 
@@ -251,9 +255,15 @@ def test_export_failure_is_soft_and_bounded(monkeypatch):
     assert time.monotonic() - started < 0.5
 
 
-def test_export_batches_with_one_oidc_token(monkeypatch):
+@pytest.mark.parametrize("byte_limited", [False, True])
+def test_export_batches_with_one_oidc_token(monkeypatch, byte_limited):
     monkeypatch.setenv("BUILDKITE", "true")
-    monkeypatch.setattr(ci_otel, "MAX_BATCH_SIZE", 2)
+    if byte_limited:
+        monkeypatch.setattr(
+            ci_otel, "MAX_BATCH_BYTES", len(ci_otel._encode_span(_span())) * 2
+        )
+    else:
+        monkeypatch.setattr(ci_otel, "MAX_BATCH_SIZE", 2)
     token_calls = []
     requests = []
 
@@ -387,6 +397,35 @@ def test_shell_wrapper_preserves_failure_status(tmp_path):
     assert result.returncode == 1
 
 
+def test_shell_wrapper_does_not_write_bytecode_to_helper_tree(tmp_path):
+    helper_dir = tmp_path / "ci-otel"
+    shutil.copytree(
+        SCRIPTS_DIR,
+        helper_dir,
+        ignore=shutil.ignore_patterns("__pycache__"),
+    )
+    shell = (
+        f'. "{helper_dir / "ci_otel.sh"}"; '
+        f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0"
+    )
+
+    result = subprocess.run(
+        ["/bin/sh", "-c", shell],
+        check=False,
+        capture_output=True,
+        text=True,
+        env={
+            **os.environ,
+            "CI_INFRA_OTEL_DIR": str(helper_dir),
+            "CI_INFRA_GPU_SAMPLING": "0",
+            "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"),
+        },
+    )
+
+    assert result.returncode == 0, result.stderr
+    assert list(helper_dir.rglob("*.pyc")) == []
+
+
 def test_ci_otel_run_records_command_and_preserves_status(tmp_path):
     shell = (
         f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; '
@@ -786,3 +825,398 @@ def test_pytest_plugin_failure_cannot_fail_pytest(tmp_path):
 
     assert result.returncode == 0, result.stderr
     assert "1 passed" in result.stdout
+
+
+def test_gpu_samples_preserve_zero_and_omit_unsupported_metrics():
+    events = ci_gpu.parse_samples(
+        "0, GPU-a, H200, 0, 1024, 2048, Disabled\n"
+        "1, GPU-b, H200, [N/A], [N/A], 2048, Disabled\n"
+        "2, GPU-c, H200, 50, 1024, 2048, Enabled\n"
+        "3, GPU-d, H200, NaN, -1, inf, Disabled\n",
+        1234,
+    )
+    assert len(events) == 4
+    assert events[0]["attributes"]["gpu.utilization"] == 0
+    assert events[0]["attributes"]["gpu.memory.used"] == 1024**3
+    for event in events[1:]:
+        assert "gpu.utilization" not in event["attributes"]
+        assert "gpu.memory.used" not in event["attributes"]
+    assert events[2]["attributes"]["gpu.status"] == "unsupported_mig"
+    assert "gpu.memory.total" not in events[2]["attributes"]
+
+
+@pytest.mark.parametrize(
+    "cuda,nvidia,expected",
+    [
+        ("MIG-a", "GPU-parent", ["MIG-a"]),
+        ("0", "MIG-a", ["MIG-a"]),
+        ("1,0", "MIG-a,MIG-b", ["MIG-b", "MIG-a"]),
+        (None, "MIG-a", ["MIG-a"]),
+        ("", "MIG-a", []),
+        ("-1", "MIG-a", []),
+        ("2", "MIG-a", []),
+        ("GPU-parent", "MIG-a", []),
+        ("0", "all", []),
+    ],
+)
+def test_mig_assignment_obeys_cuda_visibility(monkeypatch, cuda, nvidia, expected):
+    monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", nvidia)
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    if cuda is not None:
+        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda)
+    assert ci_gpu.visible_mig_devices() == expected
+
+
+@pytest.mark.parametrize("used", [0, 3 * 1024**3])
+def test_mig_memory_uses_assigned_handle_without_parent_metrics(monkeypatch, used):
+    handles = []
+    closed = []
+
+    def get_handle(uuid):
+        handles.append(uuid)
+        return "slice-handle"
+
+    def get_memory(handle):
+        assert handle == "slice-handle"
+        return SimpleNamespace(used=used, total=16 * 1024**3)
+
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: closed.append(True),
+        nvmlDeviceGetHandleByUUID=get_handle,
+        nvmlDeviceIsMigDeviceHandle=lambda handle: handle == "slice-handle",
+        nvmlDeviceGetMemoryInfo=get_memory,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    samples = ci_gpu.query_mig_samples(["MIG-a"], 1234)
+    assert handles == ["MIG-a"]
+    assert closed == [True]
+    assert samples[0]["time_ns"] == 1234
+    attributes = samples[0]["attributes"]
+    assert attributes["gpu.uuid"] == "MIG-a"
+    assert attributes["gpu.memory.used"] == used
+    assert attributes["gpu.memory.total"] == 16 * 1024**3
+    assert "gpu.utilization" not in attributes
+
+
+@pytest.mark.parametrize("is_mig", [True, False])
+def test_mig_memory_never_substitutes_parent_or_permission_failure(monkeypatch, is_mig):
+    def unavailable(handle):
+        if not is_mig:
+            pytest.fail("queried parent memory")
+        raise RuntimeError("Insufficient Permissions")
+
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: None,
+        nvmlDeviceGetHandleByUUID=lambda uuid: "handle",
+        nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig,
+        nvmlDeviceGetMemoryInfo=unavailable,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    samples = ci_gpu.query_mig_samples(["MIG-a"], 1234)
+    if is_mig:
+        assert samples[0]["attributes"]["gpu.status"] == "unavailable_mig_memory"
+        assert "gpu.memory.used" not in samples[0]["attributes"]
+    else:
+        assert samples == []
+
+
+def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip(
+    monkeypatch, tmp_path
+):
+    monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16)
+    monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8)
+    monkeypatch.setenv("CI_INFRA_OTEL_SPOOL_DIR", str(tmp_path))
+    span = ci_gpu.sample_span(
+        ci_gpu.parse_samples("0, GPU-a, H200, 90, 1, 2, Disabled", 1234)
+    )
+    assert span.parent_span_id == "02" * 8
+    assert span.trace_id == "01" * 16
+    assert ci_otel.record_spans([span])
+    assert ci_otel.load_spans() == [span]
+    assert b"ci.gpu.sample" in ci_otel.encode_request([span])
+
+
+@pytest.mark.parametrize("mig", [False, True])
+def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(
+    monkeypatch, mig
+):
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "MIG-a" if mig else "0")
+    calls = []
+
+    def query(*args, **kwargs):
+        calls.append(kwargs["timeout"])
+        raise subprocess.TimeoutExpired("nvidia-smi", kwargs["timeout"])
+
+    monkeypatch.setattr(ci_gpu.subprocess, "run", query)
+    monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0)
+    monkeypatch.setattr(
+        ci_gpu, "record_spans", lambda _: pytest.fail("fabricated sample")
+    )
+    ci_gpu.collect(os.getppid(), threading.Event())
+    assert calls == [2, 2, 2]
+
+
+def test_gpu_periodic_batches_are_bounded_even_when_upload_fails(monkeypatch):
+    monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16)
+    monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8)
+    monkeypatch.setattr(ci_gpu, "BATCH_SECONDS", 0)
+    monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0)
+    monkeypatch.setattr(
+        ci_gpu.subprocess,
+        "run",
+        lambda *a, **k: subprocess.CompletedProcess(
+            a, 0, "0, GPU-a, H200, 50, 1, 2, Disabled"
+        ),
+    )
+    stop = threading.Event()
+    batches = []
+
+    def upload(spans, timeout_seconds):
+        batches.append(spans)
+        assert timeout_seconds == 2
+        if len(batches) == 3:
+            stop.set()
+        return False
+
+    monkeypatch.setattr(ci_gpu, "export_spans", upload)
+    ci_gpu.collect(os.getppid(), stop)
+    assert [len(batch[0].events) for batch in batches] == [1, 1, 1]
+
+
+def _write_executable(path: Path, body: str) -> None:
+    path.write_text(f"#!/bin/sh\n{body}")
+    path.chmod(0o755)
+
+
+@pytest.mark.parametrize(
+    "workspace_owner,expect_sampler",
+    [("997 996", True), ("0 0", False)],
+)
+def test_root_shell_runs_gpu_sampler_as_checkout_owner(
+    tmp_path, workspace_owner, expect_sampler
+):
+    bin_dir = tmp_path / "bin"
+    workspace = tmp_path / "workdir"
+    runtime = tmp_path / "runtime"
+    log = tmp_path / "calls.log"
+    bin_dir.mkdir()
+    workspace.mkdir()
+
+    _write_executable(
+        bin_dir / "id",
+        '[ "$1" = "-u" ] && { echo 0; exit 0; }; exec /usr/bin/id "$@"\n',
+    )
+    _write_executable(
+        bin_dir / "stat",
+        f"printf '%s\\n' '{workspace_owner}'\n",
+    )
+    for command in ("chown", "chmod"):
+        _write_executable(bin_dir / command, "exit 0\n")
+    _write_executable(
+        bin_dir / "setpriv",
+        'printf "setpriv %s\\n" "$*" >> "$CI_OTEL_TEST_LOG"\n'
+        'while [ "$#" -gt 0 ] && [ "$1" != "--" ]; do shift; done\n'
+        '[ "$#" -gt 0 ] && shift\n'
+        'exec "$@"\n',
+    )
+    _write_executable(
+        bin_dir / "python3",
+        '[ "$1" = "-c" ] && exit 0\n'
+        '[ "$2" = "new-context" ] && {\n'
+        '  echo "01010101010101010101010101010101 0202020202020202 - 1"\n'
+        "  exit 0\n"
+        "}\n"
+        'case "$1" in *ci_gpu.py) echo gpu >> "$CI_OTEL_TEST_LOG";; esac\n'
+        "exit 0\n",
+    )
+    _write_executable(bin_dir / "nvidia-smi", "exit 0\n")
+    shell = (
+        f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; '
+        f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0"
+    )
+
+    result = subprocess.run(
+        ["/bin/sh", "-c", shell],
+        check=False,
+        capture_output=True,
+        text=True,
+        env={
+            **os.environ,
+            "PATH": f"{bin_dir}:{os.environ['PATH']}",
+            "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR),
+            "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime),
+            "CI_INFRA_OTEL_WORKSPACE_DIR": str(workspace),
+            "CI_OTEL_TEST_LOG": str(log),
+        },
+    )
+
+    assert result.returncode == 0, result.stderr
+    calls = log.read_text() if log.exists() else ""
+    if expect_sampler:
+        expected = "--reuid=997 --regid=996 --clear-groups --no-new-privs --"
+        assert calls.count("setpriv") == 2
+        assert calls.count(expected) == 2
+        assert "gpu" in calls
+    else:
+        assert calls == ""
+        assert "non-root GPU sampler unavailable" in result.stderr
+
+
+@pytest.mark.parametrize("sampling", ["1", "0"])
+def test_gpu_sampler_stops_with_command_and_preserves_failure_status(
+    tmp_path, sampling
+):
+    bin_dir = tmp_path / "bin"
+    bin_dir.mkdir()
+    smi = bin_dir / "nvidia-smi"
+    smi.write_text('#!/bin/sh\necho "0, GPU-a, H200, 70, 1024, 2048, Disabled"\n')
+    smi.chmod(0o755)
+    # Keep the spool for inspection, and avoid any network upload.
+    shell = (
+        f'. "{SCRIPTS_DIR / "ci_otel.sh"}"\n'
+        "ci_otel_run 1 workload sh -c 'sleep 1.3; exit 7'\n"
+        "exit $?"
+    )
+    result = subprocess.run(
+        ["/bin/sh", "-c", shell],
+        capture_output=True,
+        text=True,
+        timeout=8,
+        env={
+            **os.environ,
+            "PATH": f"{bin_dir}:{os.environ['PATH']}",
+            "BUILDKITE": "false",
+            "CI_INFRA_GPU_SAMPLING": sampling,
+            "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR),
+            "CI_INFRA_OTEL_RUNTIME_DIR": str(tmp_path / "runtime"),
+            "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"),
+        },
+    )
+    assert result.returncode == 7, result.stderr
+    records = [
+        json.loads(line)
+        for path in (tmp_path / "spans").glob("spans-*.jsonl")
+        for line in path.read_text().splitlines()
+    ]
+    command = next(record for record in records if record["name"] == "ci.command")
+    batches = [record for record in records if record["name"] == "ci.gpu.samples"]
+    assert bool(batches) == (sampling == "1")
+    if batches:
+        assert batches[0]["parent_span_id"] == command["span_id"]
+        assert all(
+            command["start_ns"] <= event["time_ns"] <= command["end_ns"]
+            for event in batches[0]["events"]
+        )
+
+
+@pytest.mark.parametrize("is_mig", [True, False])
+def test_cdi_discovery_uses_only_cuda_visible_uuids(monkeypatch, is_mig):
+    identifiers = []
+
+    def get_count(output):
+        output._obj.value = 1
+        return 0
+
+    def get_device(output, index):
+        assert index == 0
+        output._obj.value = 7
+        return 0
+
+    def get_uuid(output, device):
+        assert device.value == 7
+        output._obj[:] = bytes(range(16))
+        return 0
+
+    driver = SimpleNamespace(
+        cuInit=lambda flags: 0,
+        cuDeviceGetCount=get_count,
+        cuDeviceGet=get_device,
+        cuDeviceGetUuid_v2=get_uuid,
+    )
+    monkeypatch.setattr(ci_gpu.ctypes, "CDLL", lambda _: driver)
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: None,
+        nvmlDeviceGetHandleByUUID=lambda value: identifiers.append(value) or "handle",
+        nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    expected = "MIG-00010203-0405-0607-0809-0a0b0c0d0e0f"
+    assert ci_gpu.cuda_visible_mig_devices() == ([expected] if is_mig else [])
+    assert identifiers == [expected]
+
+
+@pytest.mark.parametrize("discovery_fails", [False, True])
+def test_cdi_sampling_switches_to_slice_or_keeps_unknown(monkeypatch, discovery_fails):
+    monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16)
+    monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8)
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", "void")
+    monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0)
+    stop = threading.Event()
+    commands = []
+    records = []
+
+    def query(command, **kwargs):
+        assert kwargs["timeout"] == ci_gpu.QUERY_TIMEOUT_SECONDS
+        commands.append(command)
+        if "--discover-mig" in command:
+            if discovery_fails:
+                raise subprocess.TimeoutExpired(command, 2)
+            return SimpleNamespace(stdout='["MIG-assigned"]')
+        if len(commands) == 4:
+            stop.set()
+        if "--query-mig" in command:
+            assert command[-1] == "MIG-assigned"
+            return SimpleNamespace(
+                stdout=json.dumps(
+                    [
+                        {
+                            "time_ns": 2,
+                            "name": "ci.gpu.sample",
+                            "attributes": {
+                                "gpu.uuid": "MIG-assigned",
+                                "gpu.memory.used": 123,
+                            },
+                        }
+                    ]
+                )
+            )
+        return SimpleNamespace(stdout="0, GPU-parent, H200, 99, 100, 200, Enabled")
+
+    monkeypatch.setattr(ci_gpu.subprocess, "run", query)
+    monkeypatch.setattr(ci_gpu, "record_spans", lambda spans: records.extend(spans))
+    ci_gpu.collect(os.getppid(), stop)
+    assert sum("--discover-mig" in command for command in commands) == 1
+    attributes = [event["attributes"] for span in records for event in span.events]
+    if discovery_fails:
+        assert all("gpu.memory.used" not in event for event in attributes)
+    else:
+        assert any(event.get("gpu.memory.used") == 123 for event in attributes)
+        assert all(event["gpu.uuid"] == "MIG-assigned" for event in attributes)
+
+
+def test_bundled_nvml_load_does_not_import_vllm_or_torch(tmp_path):
+    # Match test images: the collector is outside the installed vLLM package.
+    helper = tmp_path / "ci_gpu_relocated.py"
+    helper.write_text((SCRIPTS_DIR / "ci_gpu.py").read_text())
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys; import ci_gpu_relocated as ci_gpu; "
+            "binding = ci_gpu.load_nvml(); "
+            "assert callable(binding.nvmlDeviceGetMemoryInfo); "
+            "assert 'vllm' not in sys.modules; assert 'torch' not in sys.modules",
+        ],
+        env={**os.environ, "PYTHONPATH": f"{tmp_path}{os.pathsep}{SCRIPTS_DIR}"},
+        capture_output=True,
+        text=True,
+        timeout=5,
+    )
+    assert result.returncode == 0, result.stderr
diff --git a/.buildkite/scripts/cleanup-nightly-builds.sh b/.buildkite/scripts/cleanup-nightly-builds.sh
index b324778b1f79..41b04996e455 100755
--- a/.buildkite/scripts/cleanup-nightly-builds.sh
+++ b/.buildkite/scripts/cleanup-nightly-builds.sh
@@ -56,8 +56,8 @@ get_all_tags() {
         set -x
         
         # Get both last_updated timestamp and tag name, separated by |
-        # Exclude nightly-dev tags from cleanup
-        local tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" '.results[] | select(.name | startswith($prefix)) | select(.name | startswith("nightly-dev") | not) | "\(.last_updated)|\(.name)"')
+        # Exclude mutable architecture tags and nightly-dev tags from cleanup
+        local tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" '.results[] | select(.name | startswith($prefix)) | select(.name != ($prefix + "x86_64") and .name != ($prefix + "aarch64")) | select(.name | startswith("nightly-dev") | not) | "\(.last_updated)|\(.name)"')
         
         if [ -z "$tags" ]; then
             break
diff --git a/.buildkite/scripts/crcr-report.sh b/.buildkite/scripts/crcr-report.sh
index e19146fddbed..3b12d451cdb8 100644
--- a/.buildkite/scripts/crcr-report.sh
+++ b/.buildkite/scripts/crcr-report.sh
@@ -40,8 +40,7 @@ fi
 TOKEN_SECRET_KEY="${CRCR_BUILDKITE_TOKEN_SECRET_KEY:-CRCR_BUILDKITE_API_TOKEN}"
 BK_TOKEN="${BUILDKITE_API_TOKEN:-}"
 if [[ -z "${BK_TOKEN}" ]]; then
-    # Not in the job environment, so read it from a Buildkite secret. The agent
-    # redacts values fetched this way from the log.
+    # Not in the job environment, so read it from a Buildkite secret.
     #
     # Report why a lookup failed. Swallowing stderr made a missing secret, a
     # denied policy and an unusable agent indistinguishable, all surfacing as the
@@ -50,14 +49,9 @@ if [[ -z "${BK_TOKEN}" ]]; then
         echo "buildkite-agent is not on PATH; cannot read secret '${TOKEN_SECRET_KEY}'"
     else
         secret_err="$(mktemp)"
-        # Deliberately not passing --skip-redaction. On the deployed agent
-        # (v3.73.1) the flag cannot help: secret_get.go creates the Job API
-        # client unconditionally and only checks SkipRedaction afterwards, so
-        # under the docker plugin -- which does not expose the Job API socket to
-        # the container -- it fails before the flag is read. That ordering was
-        # only fixed in v3.107.0. Skipping redaction would also stop the token
-        # being registered with the log redactor, for no gain here.
-        if BK_TOKEN="$(buildkite-agent secret get "${TOKEN_SECRET_KEY}" 2>"${secret_err}")"; then
+        # Requires agent >= 3.107.0: the Docker plugin does not mount the Job API
+        # socket needed for redaction. Capture the value without logging it.
+        if BK_TOKEN="$(buildkite-agent secret get --skip-redaction "${TOKEN_SECRET_KEY}" 2>"${secret_err}")"; then
             if [[ -z "${BK_TOKEN}" ]]; then
                 echo "secret '${TOKEN_SECRET_KEY}' resolved but is empty"
             fi
@@ -77,7 +71,8 @@ if [[ -z "${BK_TOKEN}" ]]; then
 fi
 
 AUDIENCE="pytorch-cross-repo-ci-relay"
-OIDC_TOKEN="$(buildkite-agent oidc request-token --audience "${AUDIENCE}" 2>/dev/null)"
+# OIDC redaction also requires the unavailable Job API socket.
+OIDC_TOKEN="$(buildkite-agent oidc request-token --skip-redaction --audience "${AUDIENCE}" 2>/dev/null)"
 if [[ -z "${OIDC_TOKEN}" ]]; then
     echo "could not mint a Buildkite OIDC token -- skipping report"
     exit 0
diff --git a/.buildkite/scripts/generate-nightly-index.py b/.buildkite/scripts/generate-nightly-index.py
index 67adf9c92b80..8f92996ff9ab 100644
--- a/.buildkite/scripts/generate-nightly-index.py
+++ b/.buildkite/scripts/generate-nightly-index.py
@@ -18,8 +18,7 @@
 
 
 def normalize_package_name(name: str) -> str:
-    """
-    Normalize package name according to PEP 503.
+    """Normalize package name according to PEP 503.
     https://peps.python.org/pep-0503/#normalized-names
 
     Replace runs of underscores, hyphens, and periods with a single hyphen,
@@ -55,17 +54,18 @@ class WheelFileInfo:
 
 
 def parse_from_filename(file: str) -> WheelFileInfo:
-    """
-    Parse wheel file name to extract metadata.
+    """Parse wheel file name to extract metadata.
 
     The format of wheel names:
         {package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
     All versions could contain a variant like '+cu129' or '.cpu' or `.rocm` (or not).
+
     Example:
         vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl
         vllm-0.10.2rc2+cu129-cp38-abi3-manylinux2014_aarch64.whl
         vllm-0.11.1rc8.dev14+gaa384b3c0-cp38-abi3-manylinux2014_aarch64.whl
         vllm-0.11.1rc8.dev14+gaa384b3c0.cu130-cp38-abi3-manylinux1_x86_64.whl
+
     """
     wheel_file_re = re.compile(
         r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
@@ -111,9 +111,7 @@ def parse_from_filename(file: str) -> WheelFileInfo:
 
 
 def generate_project_list(subdir_names: list[str], comment: str = "") -> str:
-    """
-    Generate project list HTML content linking to each project & variant subdirectory.
-    """
+    """Generate project list HTML content linking to each project & variant subdirectory."""
     href_tags = []
     for name in sorted(subdir_names):
         name = name.strip("/").strip(".")
@@ -127,9 +125,7 @@ def generate_package_index_and_metadata(
     index_base_dir: Path,
     comment: str = "",
 ) -> tuple[str, str]:
-    """
-    Generate package index HTML content for a specific package, linking to actual wheel files.
-    """
+    """Generate package index HTML content for a specific package, linking to actual wheel files."""
     href_tags = []
     metadata = []
     for file in sorted(wheel_files, key=lambda x: x.filename):
@@ -156,8 +152,7 @@ def generate_index_and_metadata(
     alias_to_default: str | None = None,
     comment: str = "",
 ):
-    """
-    Generate index for all wheel files.
+    """Generate index for all wheel files.
 
     Args:
         whl_files (list[str]): List of wheel files (must be directly under `wheel_base_dir`).
@@ -216,8 +211,8 @@ def generate_index_and_metadata(
             },
             ...
         ]
-    """
 
+    """
     parsed_files = [parse_from_filename(f) for f in whl_files]
 
     if not parsed_files:
diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh
index 62fda27bf0be..2705c3b87e7b 100755
--- a/.buildkite/scripts/hardware_ci/run-amd-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh
@@ -34,6 +34,9 @@
 ###############################################################################
 set -o pipefail
 
+# shellcheck source=.buildkite/scripts/rocm/build-config.sh
+source "$(dirname "${BASH_SOURCE[0]}")/../rocm/build-config.sh" || exit $?
+
 : "${BUILDKIT_PROGRESS:=plain}"
 : "${TERM:=xterm-256color}"
 : "${FORCE_COLOR:=1}"
@@ -106,6 +109,10 @@ clear_ci_orchestration_env() {
     VLLM_CI_REQUIRE_WORKSPACE_MOUNT \
     VLLM_TEST_COMMANDS \
     VLLM_CI_BRANCH \
+    CI_ROCM_DOCKERFILE_BASE \
+    CI_ROCM_DOCKERFILE \
+    ROCM_BASE_DOCKERFILE \
+    CI_BASE_DOCKERFILE \
     VLLM_CI_BASE_IMAGE \
     VLLM_CI_FALLBACK_IMAGE \
     VLLM_CI_DOCKER_DISABLED \
@@ -225,6 +232,9 @@ prepare_artifact_image() {
   metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1)
   if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then
     base_image=$(tr -d '[:space:]' < "${metadata_file}")
+  elif using_custom_rocm_dockerfiles; then
+    echo "Custom ROCm ci_base metadata is missing; using the full CI image"
+    return 1
   fi
 
   echo "--- Preparing local ROCm test image"
@@ -1602,6 +1612,13 @@ fi
 
 echo "Final commands: $commands"
 
+# Match native CPU jobs even when the container can see AMD devices.
+cpu_platform_env=()
+if [[ "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" == "0" \
+  && "$commands" != *python_only_compile.sh* ]]; then
+  cpu_platform_env=(-e "VLLM_TARGET_DEVICE=cpu")
+fi
+
 standalone_merge_base_env=()
 if [[ "$commands" == *python_only_compile.sh* ]]; then
   # The ROCm test image often ships /vllm-workspace without .git. Resolve the
@@ -1759,6 +1776,7 @@ else
     -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
     -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
     -e "PYTORCH_ROCM_ARCH=" \
+    "${cpu_platform_env[@]}" \
     "${standalone_merge_base_env[@]}" \
     --name "${container_name}" \
     "${image_name}" \
diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
index 052063f2f69e..689ca92d6a05 100755
--- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
@@ -30,7 +30,7 @@ remove_docker_container
 docker pull "$IMAGE"
 
 # Run the image
-docker run -itd --cpuset-cpus="$CORE_RANGE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name "$CONTAINER_NAME" "$IMAGE"
+docker run -itd --cpuset-cpus="$CORE_RANGE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" -e TERM=xterm-256color -e PY_COLORS=1 -e FORCE_COLOR=1 -e CLICOLOR_FORCE=1 --shm-size=4g --name "$CONTAINER_NAME" "$IMAGE"
 
 print_packages() {
   docker exec "$CONTAINER_NAME" bash -c "
@@ -112,7 +112,9 @@ print_packages
 export CONTAINER_NAME
 export -f kernel_tests model_tests serving_tests
 case "$SHARD_ID" in
-  0) timeout 30m bash -c kernel_tests ;;
+  # The kernel shard contains several large parametrized suites and can exceed
+  # 30 minutes on otherwise healthy Arm runners.
+  0) timeout 60m bash -c kernel_tests ;;
   1) timeout 30m bash -c model_tests ;;
   2) timeout 30m bash -c serving_tests ;;
 esac
diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh
index e47cc22c8fea..eecf70ea2bde 100644
--- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh
@@ -59,7 +59,7 @@ docker pull ubuntu:25.04 || true
 
 # building the docker image
 echo "--- :docker: Building Docker image"
-BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:'
+BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:|The read operation timed out|BrokenPipeError:.*Broken pipe|HTTP/[0-9.]+ stream [0-9]+ was not closed cleanly|fetch-pack: unexpected disconnect|fatal: early EOF|invalid index-pack output'
 BUILD_MAX_ATTEMPTS=4          # 1 initial + 3 retries
 BUILD_RETRY_WAITS=(10 20 40)  # seconds to wait before retry 1/2/3
 build_log="$(mktemp)"
@@ -83,11 +83,23 @@ rm -f "$build_log"
 
 # Run the image, setting --shm-size=4g for tensor parallel. Default to
 # HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network;
-# retry once online if the cache is missing something.
-OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)'
+# retry once online if the cache is missing something. vllm wraps offline
+# cache-miss errors in ways that don't preserve the raw huggingface_hub
+# exception name: get_config() re-raises as a generic ValueError
+# (transformers_utils/config.py), and default_loader.py raises its own
+# RuntimeError when a snapshot dir has a config but no weight files. Match
+# those messages too or the fallback never triggers for those cache misses.
+#
+# Also mount ~/.cache/vllm: multimodal test fixtures (e.g. VideoAsset, via
+# vllm/assets/video.py) download into VLLM_ASSETS_CACHE (~/.cache/vllm/assets
+# by default), a directory distinct from the HF hub cache above. Without a
+# persistent mount there, those assets never survive past one container's
+# lifetime, so every run re-fetches them from the network on the "offline"
+# attempt and unconditionally falls back to the online retry.
+OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified|Cannot find any model weights with'
 run_test() {
     local hf_offline=$1
-    docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \
+    docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface -v ~/.cache/vllm:/root/.cache/vllm --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" -e TERM=xterm-256color -e PY_COLORS=1 -e FORCE_COLOR=1 -e CLICOLOR_FORCE=1 --shm-size=4g "$IMAGE_NAME" \
         timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
 }
 
diff --git a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh
index 58c0403b17a5..905ce0a96373 100644
--- a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh
@@ -43,7 +43,7 @@ case "${test_suite}" in
     pytest -v -s v1/structured_output
     pytest -v -s v1/test_serial_utils.py
     pytest -v -s v1/e2e/general/test_correctness_sliding_window.py
-    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py
+    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head
     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 --ignore=v1/kv_connector/unit/test_offloading_connector.py
     ;;
   server)
@@ -65,8 +65,6 @@ case "${test_suite}" in
     pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8
     ;;
   graph)
-    export VLLM_XPU_ENABLE_XPU_GRAPH=1
-
     python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-0.6B
     python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-0.6B --kv-cache-dtype fp8
     python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-CT-AutoRound --max-model-len 4096
diff --git a/.buildkite/scripts/hardware_ci/run-npu-test.sh b/.buildkite/scripts/hardware_ci/run-npu-test.sh
index f7227095ac2e..d0c973f6daff 100644
--- a/.buildkite/scripts/hardware_ci/run-npu-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-npu-test.sh
@@ -76,21 +76,12 @@ FROM ${BASE_IMAGE_NAME}
 ENV DEBIAN_FRONTEND=noninteractive
 ENV SOC_VERSION="ascend910b1"
 
-# Install uv and configure uv to use internal PyPI mirror
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh
-ENV PATH="/root/.local/bin:$PATH"
-ENV UV_SYSTEM_PYTHON=true
-ENV UV_INDEX_URL="http://cache-service-vllm.nginx-pypi-cache.svc.cluster.local:${PYPI_CACHE_PORT}/pypi/simple"
-ENV UV_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi"
-
-RUN apt-get update -y && \
-    apt-get install -y python3-pip git vim wget net-tools gcc g++ cmake libnuma-dev && \
-    rm -rf /var/cache/apt/* && \
-    rm -rf /var/lib/apt/lists/*
+RUN pip config set global.index-url http://cache-service-vllm.nginx-pypi-cache.svc.cluster.local:${PYPI_CACHE_PORT}/pypi/simple && \
+    pip config set global.trusted-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local
 
 # Install for pytest to make the docker build cache layer always valid
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install pytest>=6.0  'modelscope<1.38' --quiet
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install pytest>=6.0
 
 WORKDIR /workspace/vllm
 
@@ -167,5 +158,5 @@ docker run \
     "${image_name}" \
     bash -c '
     set -e
-    pytest -v -s tests/e2e/vllm_interface/test_vllm_pr_interface_compatibility.py
+    pytest --noconftest -v -s tests/e2e/vllm_interface/test_vllm_pr_interface_compatibility.py
 '
diff --git a/.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh b/.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh
new file mode 100755
index 000000000000..4a8e41d580b7
--- /dev/null
+++ b/.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh
@@ -0,0 +1,243 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../../.." && pwd)"
+
+model="ibm-research/PowerMoE-3b"
+port=8000
+num_decode_tokens=128
+server_pid=
+server_log=
+
+# XPU MoE kernels may not safely handle the -1 routes emitted for padding.
+export VLLM_MOE_SKIP_PADDING=0
+
+cd "$repo_root"
+
+cleanup_server() {
+    local attempt
+
+    if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then
+        kill "$server_pid" 2>/dev/null || true
+
+        for ((attempt = 0; attempt < 30; attempt++)); do
+            if ! kill -0 "$server_pid" 2>/dev/null; then
+                break
+            fi
+            sleep 1
+        done
+
+        if kill -0 "$server_pid" 2>/dev/null; then
+            echo "vLLM did not stop after 30 seconds; sending SIGKILL" >&2
+            kill -KILL "$server_pid" 2>/dev/null || true
+        fi
+
+        wait "$server_pid" 2>/dev/null || true
+    fi
+
+    server_pid=
+}
+
+on_exit() {
+    local status=$?
+
+    cleanup_server
+    if [[ "$status" -ne 0 && -n "$server_log" && -f "$server_log" ]]; then
+        if command -v buildkite-agent >/dev/null 2>&1; then
+            buildkite-agent artifact upload "$server_log" ||
+                echo "Failed to upload $server_log as a Buildkite artifact" >&2
+        fi
+    fi
+    [[ -z "$server_log" ]] || rm -f "$server_log"
+    exit "$status"
+}
+trap on_exit EXIT
+
+wait_for_server() {
+    for _ in $(seq 1 600); do
+        if curl -sf "http://127.0.0.1:$port/health" >/dev/null; then
+            return
+        fi
+        if ! kill -0 "$server_pid" 2>/dev/null; then
+            echo "vLLM exited before becoming healthy; see the server log above." >&2
+            return 1
+        fi
+        sleep 1
+    done
+    echo "Timed out waiting for vLLM server health check" >&2
+    return 1
+}
+
+request_outputs() {
+    local prompts_json=$1
+    local expected_choices=$2
+    local response
+
+    response=$(printf \
+        '{"model":"%s","prompt":%s,"temperature":0.7,"top_p":0.95,"seed":0,"max_tokens":%s,"logprobs":64,"ignore_eos":true,"return_token_ids":true}' \
+        "$model" "$prompts_json" "$num_decode_tokens" |
+        curl --fail --silent --show-error \
+            --max-time 1200 \
+            "http://127.0.0.1:$port/v1/completions" \
+            -H "Content-Type: application/json" \
+            --data-binary @-)
+
+    local -a token_id_arrays logprob_arrays
+    mapfile -t token_id_arrays < <(
+        grep -o '"token_ids":\[[^]]*\]' <<<"$response" |
+            sed 's/^"token_ids":\[//; s/\]$//'
+    )
+    mapfile -t logprob_arrays < <(
+        grep -o '"token_logprobs":\[[^]]*\]' <<<"$response" |
+            sed 's/^"token_logprobs":\[//; s/\]$//'
+    )
+    if [[ "${#token_id_arrays[@]}" -ne "$expected_choices" ]] ||
+        [[ "${#logprob_arrays[@]}" -ne "$expected_choices" ]]; then
+        echo "Expected $expected_choices choices with token IDs and logits" >&2
+        return 1
+    fi
+    local index
+    local array
+    local value
+    local -a values
+    for index in "${!logprob_arrays[@]}"; do
+        array="${logprob_arrays[index]}"
+        IFS=, read -r -a values <<<"$array"
+        if [[ "${#values[@]}" -ne "$num_decode_tokens" ]]; then
+            echo "Completion returned an unexpected number of decode steps" >&2
+            return 1
+        fi
+        IFS=, read -r -a values <<<"${token_id_arrays[index]}"
+        if [[ "${#values[@]}" -ne "$num_decode_tokens" ]]; then
+            echo "Completion returned an unexpected number of generated tokens" >&2
+            return 1
+        fi
+        for value in "${values[@]}"; do
+            if [[ ! "$value" =~ ^[0-9]+$ ]]; then
+                echo "Completion returned an invalid token ID" >&2
+                return 1
+            fi
+        done
+        IFS=, read -r -a values <<<"$array"
+        for value in "${values[@]}"; do
+            if [[ ! "$value" =~ ^-?([0-9]+([.][0-9]*)?|[.][0-9]+)([eE][+-]?[0-9]+)?$ ]]; then
+                echo "Completion returned non-finite raw logits" >&2
+                return 1
+            fi
+        done
+        printf '%s|%s\n' "${token_id_arrays[index]}" "${logprob_arrays[index]}"
+    done
+}
+
+make_prompt() {
+    local length=$1
+    local offset=${2:-0}
+    local i
+    local prompt="Use the following system-design context to answer the question. Scenario $offset:"
+    for ((i = 0; i < length; i++)); do
+        prompt+=" context"
+    done
+    prompt+="\\n\\nQuestion: How should batch-invariant inference behave when this same request is decoded alongside unrelated requests?\\nAnswer:"
+    printf '"%s"' "$prompt"
+}
+
+first_difference() {
+    local actual=$1
+    local expected=$2
+    local position=$3
+    local actual_tokens=${actual%%|*}
+    local expected_tokens=${expected%%|*}
+    local actual_logits=${actual#*|}
+    local expected_logits=${expected#*|}
+    local -a actual_token_values
+    local -a expected_token_values
+    local -a actual_logit_values
+    local -a expected_logit_values
+    local step
+
+    IFS=, read -r -a actual_token_values <<<"$actual_tokens"
+    IFS=, read -r -a expected_token_values <<<"$expected_tokens"
+    IFS=, read -r -a actual_logit_values <<<"$actual_logits"
+    IFS=, read -r -a expected_logit_values <<<"$expected_logits"
+    for ((step = 0; step < num_decode_tokens; step++)); do
+        if [[ "${actual_token_values[step]}" != "${expected_token_values[step]}" ]]; then
+            printf 'batch position %s: step=%s, baseline token=%s, actual token=%s' \
+                "$position" \
+                "$step" \
+                "${expected_token_values[step]}" \
+                "${actual_token_values[step]}"
+            return
+        fi
+        if [[ "${actual_logit_values[step]}" != "${expected_logit_values[step]}" ]]; then
+            printf 'batch position %s: step=%s, token=%s, baseline logit=%s, actual logit=%s' \
+                "$position" \
+                "$step" \
+                "${expected_token_values[step]}" \
+                "${expected_logit_values[step]}" \
+                "${actual_logit_values[step]}"
+            return
+        fi
+    done
+}
+
+run_e2e() {
+    local needle
+    local batch
+    local baseline_output
+    local batch_response
+    local difference
+    local position
+    local -a fillers
+    local -a batch_outputs
+    local -a differences=()
+
+    echo "Testing PowerMoE-3b: TP=2, VLLM_BATCH_INVARIANT=1, VLLM_MOE_SKIP_PADDING=0"
+    server_log=$(mktemp "${TMPDIR:-/tmp}/xpu-batch-invariance.XXXXXX.log")
+    VLLM_BATCH_INVARIANT=1 vllm serve "$model" \
+        --tensor-parallel-size 2 \
+        --max-model-len 1024 \
+        --max-num-seqs 16 \
+        --max-num-batched-tokens 512 \
+        --enable-chunked-prefill \
+        --no-enable-prefix-caching \
+        --enforce-eager \
+        --seed 0 \
+        --max-logprobs 64 \
+        --logprobs-mode raw_logits \
+        --port "$port" > >(tee "$server_log") 2>&1 &
+    server_pid=$!
+    wait_for_server
+
+    needle=$(make_prompt 257)
+    local filler_lengths=(7 31 63 95 127 159 191 223 255 287 319 383 447)
+    for position in "${!filler_lengths[@]}"; do
+        fillers+=( "$(make_prompt "${filler_lengths[position]}" "$((position + 1))")" )
+    done
+    batch="[$needle,${fillers[0]},${fillers[1]},${fillers[2]},${fillers[3]},${fillers[4]},${fillers[5]},$needle,${fillers[6]},${fillers[7]},${fillers[8]},${fillers[9]},${fillers[10]},${fillers[11]},${fillers[12]},$needle]"
+
+    echo "Requesting BS=1 baseline ($num_decode_tokens decode tokens)..."
+    baseline_output=$(request_outputs "$needle" 1)
+    echo "Received BS=1 baseline."
+    echo "Requesting BS=16 batch ($((16 * num_decode_tokens)) decode tokens)..."
+    batch_response=$(request_outputs "$batch" 16)
+    mapfile -t batch_outputs <<<"$batch_response"
+    echo "Received BS=16 batch; comparing generated tokens and logits..."
+
+    for position in 0 7 15; do
+        difference=$(first_difference "${batch_outputs[position]}" "$baseline_output" "$position")
+        if [[ -n "$difference" ]]; then
+            differences+=("$difference")
+        fi
+    done
+
+    if [[ "${#differences[@]}" -ne 0 ]]; then
+        echo "Batch-invariant PowerMoE-3b output changed: ${differences[*]}" >&2
+        return 1
+    fi
+    echo "PowerMoE-3b batch invariance check passed"
+    cleanup_server
+    rm -f "$server_log"
+    server_log=
+}
+
+run_e2e
diff --git a/.buildkite/scripts/initialized-snapshot-e2e.sh b/.buildkite/scripts/initialized-snapshot-e2e.sh
new file mode 100755
index 000000000000..cd0511ffdd9a
--- /dev/null
+++ b/.buildkite/scripts/initialized-snapshot-e2e.sh
@@ -0,0 +1,311 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+set -Eeuo pipefail
+readonly MODEL="Qwen/Qwen3-0.6B"
+readonly MODEL_REVISION="c1899de289a04d12100db370d81485cdf75e47ca"
+readonly EXPECTED_TOKEN_ID="12095"
+readonly EXPECTED_TEXT=" Paris"
+readonly CLEANUP_RESERVE_S=300
+RUN_LIMIT_S="${SNAPSHOT_E2E_RUN_LIMIT_S:-3300}"
+[[ "$RUN_LIMIT_S" =~ ^[1-9][0-9]{2,3}$ ]] \
+    || { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be a bounded decimal integer" >&2; exit 1; }
+(( RUN_LIMIT_S >= CLEANUP_RESERVE_S + 60 && RUN_LIMIT_S <= 3300 )) \
+    || { echo "SNAPSHOT_E2E_RUN_LIMIT_S must be between 360 and 3300" >&2; exit 1; }
+readonly RUN_LIMIT_S
+readonly HOST_HF_CACHE="${SNAPSHOT_E2E_HOST_HF_CACHE:-}"
+HF_CACHE_MOUNT=()
+if [[ -n "$HOST_HF_CACHE" ]]; then
+    [[ "$HOST_HF_CACHE" == /* && -d "$HOST_HF_CACHE" ]] \
+        || { echo "SNAPSHOT_E2E_HOST_HF_CACHE must be an existing absolute directory" >&2; exit 1; }
+    HF_CACHE_MOUNT=(--mount "type=bind,source=$HOST_HF_CACHE,target=/e2e/hf,readonly")
+    PREPARE_ARTIFACT_ROOT='install -d -m 0700 /e2e'
+else
+    PREPARE_ARTIFACT_ROOT='mkdir -m 0700 /e2e && mkdir -m 0700 /e2e/hf'
+fi
+readonly PREPARE_ARTIFACT_ROOT
+if (( $# != 2 )); then
+    echo "usage: $0 IMAGE_TAG BUILDKITE_COMMIT" >&2
+    exit 2
+fi
+readonly IMAGE_TAG="$1"
+readonly EXPECTED_COMMIT="$2"
+readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))"
+RUN_ID="" CONTAINER_NAME="" CONTAINER_ID=""
+GPU_UUID="" PORT_ONE="" PORT_TWO=""
+LINK_REMAP_BASELINE=""
+SNAPSHOT_PIDS=()
+die() {
+    echo "initialized snapshot E2E: $*" >&2
+    exit 1
+}
+
+run() {
+    local label="$1" maximum_s="$2" available limit
+    shift 2
+    available="$((DEADLINE_EPOCH - $(date +%s) - CLEANUP_RESERVE_S))"
+    (( available > 0 )) || die "no time remains before the cleanup reserve"
+    (( available < maximum_s )) && limit="$available" || limit="$maximum_s"
+    echo "--- ${label} (timeout ${limit}s)" >&2
+    timeout --signal=TERM --kill-after=30s "${limit}s" "$@"
+}
+
+link_remaps() {
+    find /dev/shm -maxdepth 1 -name 'link_remap.*' \
+        -printf '%p|%D|%i|%m|%s\n' | sort
+}
+
+gpu_pids() {
+    nvidia-smi --id="$GPU_UUID" --query-compute-apps=pid \
+        --format=csv,noheader,nounits | sed '/^[[:space:]]*$/d'
+}
+
+port_listeners() {
+    ss -H -ltn "sport = :$1"
+}
+
+state_clean() {
+    local pid port output
+    for pid in "${SNAPSHOT_PIDS[@]}"; do
+        [[ ! -e "/proc/$pid" ]] || return 1
+    done
+    output="$(gpu_pids)" || return 1
+    [[ -z "$output" ]] || return 1
+    output="$(link_remaps)" || return 1
+    [[ "$output" == "$LINK_REMAP_BASELINE" ]] || return 1
+    for port in "$PORT_ONE" "$PORT_TWO"; do
+        [[ -n "$port" ]] || continue
+        output="$(port_listeners "$port")" || return 1
+        [[ -z "$output" ]] || return 1
+    done
+}
+
+wait_clean() {
+    local _attempt
+    for ((_attempt = 0; _attempt < 60; _attempt++)); do
+        state_clean && return 0
+        sleep 1
+    done
+    return 1
+}
+
+container_names() {
+    timeout 15s docker container ls --all --format '{{.Names}}'
+}
+
+cleanup() {
+    local failed=0 names="" actual="" expected=""
+    if [[ -n "$CONTAINER_NAME" ]]; then
+        names="$(container_names 2>/dev/null)" || return 1
+        if grep -Fxq "$CONTAINER_NAME" <<< "$names"; then
+            actual="$(docker container inspect \
+                --format '{{.Id}}|{{index .Config.Labels "ai.vllm.snapshot.e2e.run"}}' \
+                "$CONTAINER_NAME" 2>/dev/null)" || return 1
+            expected="${CONTAINER_ID:-${actual%%|*}}|$RUN_ID"
+            if [[ "$actual" != "$expected" ]]; then
+                echo "refusing to remove a container with mismatched identity" >&2
+                failed=1
+            else
+                timeout 60s docker rm --force "$CONTAINER_NAME" >/dev/null \
+                    || failed=1
+            fi
+        fi
+        names="$(container_names 2>/dev/null)" || return 1
+        grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1
+    fi
+    if [[ -n "$PORT_TWO" ]] && ! wait_clean; then
+        echo "process, GPU, port, or link-remap residue remained" >&2
+        failed=1
+    fi
+    return "$failed"
+}
+
+on_exit() {
+    local status=$?
+    trap - EXIT INT TERM
+    set +e
+    cleanup || status=1
+    exit "$status"
+}
+
+trap on_exit EXIT
+trap 'exit 130' INT
+trap 'exit 143' TERM
+
+parse_inspect() {
+    timeout 30s python3 -c '
+import json, math, sys
+data = json.load(sys.stdin)
+model, revision, gpu_uuid, token_id, text = sys.argv[1:]
+expected = {
+    "boundary": "post-engine-init-reloadable-state-released",
+    "model": model, "model_revision": revision,
+    "gpu_uuid": gpu_uuid, "oracle_token_ids": [int(token_id)],
+    "oracle_text": text,
+}
+bad = next((field for field, value in expected.items() if data.get(field) != value), None)
+if bad:
+    raise SystemExit(f"inspect field {bad} mismatch: {data.get(bad)!r}")
+logprob = data.get("oracle_sampled_token_logprob")
+if type(logprob) is not float or not math.isfinite(logprob):
+    raise SystemExit("inspect logprob is not finite")
+pids = data.get("process_tree")
+if not pids:
+    raise SystemExit("inspect process_tree is empty")
+for pid in pids:
+    print(f"pid={pid}")
+print(f"oracle_logprob={logprob}")
+' "$MODEL" "$MODEL_REVISION" "$GPU_UUID" "$EXPECTED_TOKEN_ID" \
+        "$EXPECTED_TEXT"
+}
+
+assert_oracle() {
+    local port="$1" ordinal="$2" response
+    response="$(run "restore ${ordinal} public completion" 120 curl \
+        --fail --silent --show-error --max-time 120 \
+        --header 'Content-Type: application/json' \
+        --data "{\"model\":\"${MODEL}\",\"prompt\":\"The capital of France is\",\"temperature\":0,\"seed\":0,\"min_tokens\":1,\"max_tokens\":1,\"logprobs\":0,\"return_token_ids\":true}" \
+        "http://127.0.0.1:${port}/v1/completions")"
+    printf '%s' "$response" | timeout 30s python3 -c '
+import json, math, sys
+choice = json.load(sys.stdin)["choices"][0]
+token_id, text, expected, ordinal = sys.argv[1:]
+values = choice["logprobs"]["token_logprobs"]
+actual = float(values[0]) if len(values) == 1 else math.nan
+if (choice["token_ids"] != [int(token_id)] or choice["text"] != text
+        or not math.isfinite(actual) or abs(actual - float(expected)) > 1e-3):
+    raise SystemExit(f"public oracle mismatch: {choice!r}")
+print(f"restore={ordinal} token_id={token_id} text={text!r} logprob={actual}")
+' "$EXPECTED_TOKEN_ID" "$EXPECTED_TEXT" "$INSPECT_LOGPROB" "$ordinal"
+}
+
+[[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]] \
+    || die "host must be Linux x86_64"
+[[ "$EXPECTED_COMMIT" =~ ^[0-9a-f]{40}$ ]] \
+    || die "BUILDKITE_COMMIT must be a full lowercase commit hash"
+[[ "$(< /proc/sys/kernel/io_uring_disabled)" == "2" ]] \
+    || die "kernel.io_uring_disabled must already equal 2"
+run "Docker preflight" 60 docker info >/dev/null
+
+VISIBLE_GPU="${NVIDIA_VISIBLE_DEVICES:-}"
+if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then
+    VISIBLE_GPU="${CUDA_VISIBLE_DEVICES:-}"
+fi
+if [[ -z "$VISIBLE_GPU" || "$VISIBLE_GPU" == "all" || "$VISIBLE_GPU" == "void" ]]; then
+    CDI_GPU="${BUILDKITE_PLUGIN_DOCKER_DEVICES_0:-}"
+    if [[ "$CDI_GPU" == nvidia.com/gpu=* ]]; then
+        VISIBLE_GPU="${CDI_GPU#nvidia.com/gpu=}"
+    else
+        VISIBLE_GPU=""
+    fi
+fi
+[[ -n "$VISIBLE_GPU" && "$VISIBLE_GPU" != all && "$VISIBLE_GPU" != void \
+    && "$VISIBLE_GPU" != *","* ]] \
+    || die "the runner must expose exactly one selected GPU"
+GPU_ROW="$(nvidia-smi --id="$VISIBLE_GPU" \
+    --query-gpu=name,uuid,mig.mode.current,memory.total \
+    --format=csv,noheader,nounits)"
+IFS=',' read -r GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB <<< "$GPU_ROW"
+for variable in GPU_NAME GPU_UUID MIG_MODE GPU_MEMORY_MIB; do
+    printf -v "$variable" '%s' "${!variable#"${!variable%%[![:space:]]*}"}"
+done
+[[ "$GPU_NAME" == *H200* && "$GPU_UUID" == GPU-* && "$MIG_MODE" == Disabled ]] \
+    || die "selected device is not a full non-MIG H200: $GPU_ROW"
+(( GPU_MEMORY_MIB >= 130000 )) || die "selected H200 does not expose full memory"
+GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
+[[ -z "$GPU_PIDS" ]] || die "selected H200 is not idle"
+
+RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM"
+CONTAINER_NAME="vllm-$RUN_ID"
+read -r PORT_ONE PORT_TWO < <(python3 -c 'import socket; s=[socket.socket(),socket.socket()]; [x.bind(("127.0.0.1",0)) for x in s]; print(*(x.getsockname()[1] for x in s))')
+LINK_REMAP_BASELINE="$(link_remaps)"
+
+run "pull exact candidate image" 900 docker pull --platform linux/amd64 \
+    "$IMAGE_TAG" >/dev/null
+IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <(
+    docker image inspect \
+        --format '{{.Id}}|{{.Os}}/{{.Architecture}}|{{index .Config.Labels "org.opencontainers.image.revision"}}|{{index .Config.Labels "ai.vllm.build.commit"}}' \
+        "$IMAGE_TAG"
+)
+[[ "$IMAGE_ID" =~ ^sha256:[0-9a-f]{64}$ && "$IMAGE_PLATFORM" == linux/amd64 ]] \
+    || die "candidate image identity is invalid"
+[[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \
+    || die "candidate image commit labels do not match"
+CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \
+    --name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \
+    --gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \
+    --network=host --env "CUDA_VISIBLE_DEVICES=$GPU_UUID" --env HF_HUB_DISABLE_TELEMETRY=1 \
+    --env VLLM_NO_USAGE_STATS=1 --env VLLM_USE_V2_MODEL_RUNNER=1 \
+    ${HF_CACHE_MOUNT[@]+"${HF_CACHE_MOUNT[@]}"} \
+    --entrypoint sleep "$IMAGE_ID" infinity)"
+[[ "$CONTAINER_ID" =~ ^[0-9a-f]{64}$ ]] || die "invalid container ID"
+run "prepare private artifact root" 60 docker exec "$CONTAINER_NAME" sh -c \
+    "$PREPARE_ARTIFACT_ROOT && df -Pk / | awk 'NR == 2 && \$4 >= 10485760 {ok=1} END {exit !ok}'"
+run "verify snapshot runtime" 60 docker exec "$CONTAINER_NAME" sh -c \
+    'criu --version; test -x /usr/local/sbin/cuda-checkpoint; test -f /usr/local/lib/criu/cuda_plugin.so'
+run "verify selected CUDA device" 60 docker exec "$CONTAINER_NAME" python3 -c '
+import sys, torch
+if torch.cuda.device_count() != 1:
+    raise SystemExit("expected exactly one visible CUDA device")
+actual = str(torch.cuda.get_device_properties(0).uuid).removeprefix("GPU-")
+if actual != sys.argv[1].removeprefix("GPU-"):
+    raise SystemExit(f"CUDA device {actual} does not match {sys.argv[1]}")
+print(f"selected_cuda_uuid=GPU-{actual}")
+' "$GPU_UUID"
+OFFLINE_EXEC=(docker exec --env HF_HOME=/e2e/hf --env HF_HUB_OFFLINE=1 \
+    --env TRANSFORMERS_OFFLINE=1 --env VLLM_SNAPSHOT_TIMEOUT_S=900 \
+    "$CONTAINER_NAME")
+if [[ -n "$HOST_HF_CACHE" ]]; then
+    run "validate pinned model in read-only host cache" 120 \
+        "${OFFLINE_EXEC[@]}" python3 -c \
+        'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2], local_files_only=True); assert Path(p).name == sys.argv[2]' \
+        "$MODEL" "$MODEL_REVISION"
+else
+    run "prefetch pinned public model" 900 docker exec --env HF_HOME=/e2e/hf \
+        "$CONTAINER_NAME" python3 -c \
+        'from pathlib import Path; from huggingface_hub import snapshot_download; import sys; p=snapshot_download(repo_id=sys.argv[1], revision=sys.argv[2]); assert Path(p).name == sys.argv[2]' \
+        "$MODEL" "$MODEL_REVISION"
+fi
+
+run "create compact initialized snapshot" 1200 "${OFFLINE_EXEC[@]}" \
+    vllm snapshot create "$MODEL" --revision "$MODEL_REVISION" \
+    --tokenizer-revision "$MODEL_REVISION" --snapshot-dir /e2e/artifact \
+    --dtype float16 --max-model-len 512 --gpu-memory-utilization 0.50
+wait_clean || die "create left process, GPU, port, or link-remap residue"
+
+INSPECT_JSON="$(run "inspect compact initialized snapshot" 120 \
+    "${OFFLINE_EXEC[@]}" vllm snapshot inspect /e2e/artifact)"
+INSPECT_RECEIPT="$(printf '%s' "$INSPECT_JSON" | parse_inspect)"
+mapfile -t SNAPSHOT_PIDS < <(sed -n 's/^pid=//p' <<< "$INSPECT_RECEIPT")
+INSPECT_LOGPROB="$(sed -n 's/^oracle_logprob=//p' <<< "$INSPECT_RECEIPT")"
+(( ${#SNAPSHOT_PIDS[@]} > 0 )) || die "inspect returned no PIDs"
+state_clean || die "create left recorded state alive"
+echo "$INSPECT_RECEIPT"
+
+run "restore compact snapshot 1" 900 "${OFFLINE_EXEC[@]}" \
+    vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_ONE"
+PORT_LISTENERS="$(port_listeners "$PORT_ONE")" || die "failed to inspect restore 1 port"
+[[ -n "$PORT_LISTENERS" ]] || die "restore 1 did not bind its port"
+GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
+[[ -n "$GPU_PIDS" ]] || die "restore 1 did not engage the selected GPU"
+assert_oracle "$PORT_ONE" 1
+run "stop container after restore 1" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null
+[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \
+    || die "docker stop replaced the container"
+wait_clean || die "restore 1 teardown left residue"
+
+run "start the same container" 60 docker start "$CONTAINER_NAME" >/dev/null
+[[ "$(docker container inspect --format '{{.Id}}' "$CONTAINER_NAME")" == "$CONTAINER_ID" ]] \
+    || die "docker start replaced the container"
+run "restore compact snapshot 2" 900 "${OFFLINE_EXEC[@]}" \
+    vllm snapshot restore /e2e/artifact --host 127.0.0.1 --port "$PORT_TWO"
+PORT_LISTENERS="$(port_listeners "$PORT_TWO")" || die "failed to inspect restore 2 port"
+[[ -n "$PORT_LISTENERS" ]] || die "restore 2 did not bind its port"
+GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU"
+[[ -n "$GPU_PIDS" ]] || die "restore 2 did not engage the selected GPU"
+assert_oracle "$PORT_TWO" 2
+run "stop container after restore 2" 60 docker stop --time 20 "$CONTAINER_NAME" >/dev/null
+wait_clean || die "restore 2 teardown left residue"
+cleanup || die "exact cleanup failed"
+echo "container_id=$CONTAINER_ID restore_1_port=$PORT_ONE restore_2_port=$PORT_TWO"
+echo "result=PASS"
diff --git a/.buildkite/scripts/install-kv-connectors.sh b/.buildkite/scripts/install-kv-connectors.sh
index ef47c3625f44..8733a21347d3 100755
--- a/.buildkite/scripts/install-kv-connectors.sh
+++ b/.buildkite/scripts/install-kv-connectors.sh
@@ -9,9 +9,20 @@ if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None
     exit 0
 fi
 
-REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
+# Select requirements before optional packages can change the installed Torch.
+CUDA_MAJOR=$(python3 -c 'import torch; print(torch.version.cuda.split(".", 1)[0])')
+DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors.txt
+if [ "${CUDA_MAJOR}" = "12" ]; then
+    DEFAULT_REQUIREMENTS=/vllm-workspace/requirements/kv_connectors_cu12.txt
+fi
+REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-${DEFAULT_REQUIREMENTS}}"
 
-uv pip install --system -r "${REQUIREMENTS_FILE}"
+CONSTRAINTS_DIR=$(mktemp -d)
+trap 'rm -rf "${CONSTRAINTS_DIR}"' EXIT
+uv pip freeze --system > "${CONSTRAINTS_DIR}/installed.txt"
+grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \
+    "${CONSTRAINTS_DIR}/installed.txt" > "${CONSTRAINTS_DIR}/constraints.txt"
+uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" -r "${REQUIREMENTS_FILE}"
 
 KV_METADATA=$(python3 - <<'PY'
 import importlib.metadata as metadata
@@ -56,5 +67,5 @@ PY
 # build first to avoid a clash.
 if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then
     uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true
-    uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
+    uv pip install --system -c "${CONSTRAINTS_DIR}/constraints.txt" "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
 fi
diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh
new file mode 100644
index 000000000000..46d756b621b1
--- /dev/null
+++ b/.buildkite/scripts/install-kv-offload.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+set -euo pipefail
+
+if [ -z "${CUDA_VERSION:-}" ]; then
+    exit 0
+fi
+
+# TODO: Use the PyPI wheel once it is published.
+uv pip install --system \
+    "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0"
+
+# Keep only the NIXL wheel matching the CI image's CUDA runtime.
+NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p')
+uv pip uninstall --system nixl-cu12 nixl-cu13
+uv pip install --system --no-deps "nixl-cu${CUDA_VERSION%%.*}==${NIXL_VERSION}"
diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh
index 2fa334241fce..dec366390b00 100755
--- a/.buildkite/scripts/publish-release-images.sh
+++ b/.buildkite/scripts/publish-release-images.sh
@@ -61,6 +61,13 @@ if target_enabled cuda-13-0; then
   docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64"
   docker manifest push vllm/vllm-openai:latest
   docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}"
+
+  ZSTD_DIGEST=$(docker buildx imagetools inspect \
+    "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" \
+    --format '{{json .Manifest.Digest}}' | tr -d '"')
+  .buildkite/scripts/publish-zstd-image.sh \
+    "public.ecr.aws/q9t5s3a7/vllm-release-repo@${ZSTD_DIGEST}" \
+    "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-zstd"
 fi
 
 # ---- CUDA 12.9 ----
diff --git a/.buildkite/scripts/publish-zstd-image.sh b/.buildkite/scripts/publish-zstd-image.sh
new file mode 100755
index 000000000000..bd3d2f99d65c
--- /dev/null
+++ b/.buildkite/scripts/publish-zstd-image.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+set -euo pipefail
+
+if [ "$#" -ne 2 ]; then
+  echo "Usage: $0 SOURCE@sha256:DIGEST DESTINATION-zstd" >&2
+  exit 2
+fi
+
+SOURCE=$1
+DESTINATION=$2
+if [[ ! "$SOURCE" =~ ^[a-zA-Z0-9][a-zA-Z0-9._:/-]*@sha256:[a-f0-9]{64}$ ]] ||
+  [[ ! "$DESTINATION" =~ ^[a-zA-Z0-9][a-zA-Z0-9._:/-]*:[a-zA-Z0-9_][a-zA-Z0-9_.-]*-zstd$ ]]; then
+  echo "ERROR: source must use an immutable sha256 digest and destination must end in -zstd" >&2
+  exit 2
+fi
+
+BUILDER=
+cleanup() {
+  status=$?
+  trap - EXIT INT TERM
+  if [ -n "$BUILDER" ] && ! docker buildx rm --force "$BUILDER"; then
+    echo "ERROR: failed to remove Buildx builder $BUILDER" >&2
+    [ "$status" -ne 0 ] || status=1
+  fi
+  exit "$status"
+}
+trap cleanup EXIT
+trap 'exit 130' INT
+trap 'exit 143' TERM
+
+BUILDER=$(docker buildx create \
+  --driver docker-container \
+  --driver-opt network=host \
+  --driver-opt 'image=moby/buildkit:v0.32.2@sha256:28a898719c18a33f4e8000685287fa36fd0dd9560c6440227d3a732d79bb41d8')
+
+printf '# check=error=true\n\nFROM %s\n' "$SOURCE" | docker buildx build \
+  --builder "$BUILDER" \
+  --platform linux/amd64 \
+  --progress plain \
+  --provenance=mode=min \
+  --output "type=image,name=${DESTINATION},push=true,compression=zstd,compression-level=3,force-compression=true,oci-mediatypes=true" \
+  -
diff --git a/.buildkite/scripts/rocm/build-config.sh b/.buildkite/scripts/rocm/build-config.sh
new file mode 100644
index 000000000000..ad0179b16359
--- /dev/null
+++ b/.buildkite/scripts/rocm/build-config.sh
@@ -0,0 +1,85 @@
+#!/usr/bin/env bash
+# Shared Dockerfile selection for the normal AMD CI build chain.
+# CI_ROCM_DOCKERFILE_BASE and CI_ROCM_DOCKERFILE accept repository-relative paths.
+# For Rock, set CI_ROCM_DOCKERFILE_BASE=docker/Dockerfile.rock_base and
+# CI_ROCM_DOCKERFILE=docker/Dockerfile.rock.
+# Custom recipes must preserve the CI targets and artifact layout. Additional
+# COPY inputs also need the corresponding *_CONTENT_FILES overrides.
+
+rocm_base_dockerfile() {
+    local default="docker/Dockerfile.rocm_base"
+    printf '%s\n' "${CI_ROCM_DOCKERFILE_BASE:-${ROCM_BASE_DOCKERFILE:-${default}}}"
+}
+
+rocm_ci_dockerfile() {
+    local default="docker/Dockerfile.rocm"
+    printf '%s\n' "${CI_ROCM_DOCKERFILE:-${CI_BASE_DOCKERFILE:-${default}}}"
+}
+
+using_custom_rocm_dockerfiles() {
+    [[ "$(rocm_base_dockerfile)" != "docker/Dockerfile.rocm_base" \
+        || "$(rocm_ci_dockerfile)" != "docker/Dockerfile.rocm" ]]
+}
+
+rocm_dockerfile_stages() {
+    awk '
+        toupper($1) == "FROM" {
+            for (idx = 2; idx < NF; idx++) {
+                if (tolower($idx) == "as") print tolower($(idx + 1))
+            }
+        }
+    ' "$1"
+}
+
+validate_rocm_dockerfile() {
+    local dockerfile="$1"
+    local stage=""
+    local stages=""
+    shift
+
+    case "${dockerfile}" in
+        /*|../*|*/../*|*/..|*[[:space:]]*)
+            echo "ROCm Dockerfile must be a repository-relative path without '..' or whitespace: ${dockerfile}" >&2
+            return 2
+            ;;
+    esac
+    if [[ ! -f "${dockerfile}" ]]; then
+        echo "ROCm Dockerfile not found: ${dockerfile}" >&2
+        return 2
+    fi
+    if (($#)) && ! awk '
+        toupper($1) == "FROM" {
+            named = 0
+            for (idx = 2; idx < NF; idx++) {
+                if (tolower($idx) == "as") named = 1
+            }
+            if (!named) exit 1
+        }
+    ' "${dockerfile}"; then
+        echo "ROCm Dockerfile must name every build stage: ${dockerfile}" >&2
+        return 2
+    fi
+    stages=$(rocm_dockerfile_stages "${dockerfile}")
+    for stage in "$@"; do
+        if ! grep -Fxq -- "${stage}" <<< "${stages}"; then
+            echo "ROCm Dockerfile ${dockerfile} is missing required stage: ${stage}" >&2
+            return 2
+        fi
+    done
+}
+
+configure_rocm_build() {
+    ROCM_BASE_DOCKERFILE=$(rocm_base_dockerfile)
+    CI_BASE_DOCKERFILE=$(rocm_ci_dockerfile)
+    export ROCM_BASE_DOCKERFILE CI_BASE_DOCKERFILE
+    using_custom_rocm_dockerfiles || return 0
+
+    # Custom stacks may supply dependencies without separate build stages.
+    export ROCM_DEP_CACHE_EXPORT_MODE=never
+
+    # Preserve the standard runtime tags, which are unique to each build,
+    # while preventing experimental builds from promoting stable images.
+    export ROCM_BASE_PUSH_STABLE_TAG=0
+    export CI_BASE_PUSH_STABLE_TAG=0
+    echo "Custom ROCm Dockerfiles: ${ROCM_BASE_DOCKERFILE} and ${CI_BASE_DOCKERFILE}"
+}
diff --git a/.buildkite/scripts/rocm/refresh-base-image.sh b/.buildkite/scripts/rocm/refresh-base-image.sh
index 23d10492a99c..10e594577a87 100755
--- a/.buildkite/scripts/rocm/refresh-base-image.sh
+++ b/.buildkite/scripts/rocm/refresh-base-image.sh
@@ -4,6 +4,10 @@
 
 set -euo pipefail
 
+# shellcheck source=.buildkite/scripts/rocm/build-config.sh
+source "$(dirname "${BASH_SOURCE[0]}")/build-config.sh"
+configure_rocm_build
+
 DOCKERFILE="${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}"
 BASE_REPO="${ROCM_BASE_IMAGE_REPO:-rocm/vllm-dev}"
 CACHE_REPO="${ROCM_BASE_CACHE_REPO:-${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}}"
@@ -12,7 +16,12 @@ DEFAULT_ROCM_BASE_METADATA_VERSION="2"
 DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}"
 
 ROCM_BASE_LAYER_CACHE_REF=""
-ROCM_BASE_TRUSTED_LAYER_CACHE_REF="${CACHE_REPO}:rocm-base-main"
+BASE_CACHE_PREFIX="rocm-base"
+if using_custom_rocm_dockerfiles; then
+    BASE_CACHE_PREFIX="rocm-base-$(printf '%s\n%s\n' \
+        "${ROCM_BASE_DOCKERFILE}" "${CI_BASE_DOCKERFILE}" | sha256sum | cut -c1-12)"
+fi
+ROCM_BASE_TRUSTED_LAYER_CACHE_REF="${CACHE_REPO}:${BASE_CACHE_PREFIX}-main"
 ROCM_BASE_STABLE_TAG_UPDATED=0
 declare -a ROCM_BASE_CACHE_ARGS=()
 
@@ -120,7 +129,7 @@ configure_rocm_base_layer_cache() {
     fi
 
     scope=$(rocm_base_layer_cache_scope)
-    ROCM_BASE_LAYER_CACHE_REF="${CACHE_REPO}:rocm-base-${scope}"
+    ROCM_BASE_LAYER_CACHE_REF="${CACHE_REPO}:${BASE_CACHE_PREFIX}-${scope}"
     ROCM_BASE_CACHE_ARGS+=(
         --cache-from "type=registry,ref=${ROCM_BASE_LAYER_CACHE_REF}"
     )
@@ -139,7 +148,7 @@ configure_rocm_base_layer_cache() {
 extract_arg_default() {
     local arg_name="$1"
 
-    sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
+    sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
         "${DOCKERFILE}" | head -1
 }
 
@@ -384,6 +393,7 @@ build_base_image() {
     local stable_tag="${BASE_REPO}:base"
     local trusted_content_tag=""
     local scoped_content_tag=""
+    local preview_content_tag=""
     local writable_content_tag=""
     local build_ref=""
     local immutable_ref=""
@@ -418,11 +428,16 @@ build_base_image() {
     read -r -a content_paths <<< "${content_files}"
     content_files_hash="$(compute_content_hash "${content_paths[@]}")"
     base_hash=$(compute_base_content_hash "${use_sccache}" "${base_image_digest}")
-    rocm_version="$(rocm_version_from_base_image "${base_image_arg}")"
+    rocm_version="$(extract_arg_default ROCM_SDK_VERSION)"
+    rocm_version="${rocm_version:-$(rocm_version_from_base_image "${base_image_arg}")}"
     triton_arg="$(extract_arg_default TRITON_BRANCH)"
     pytorch_arg="$(extract_arg_default PYTORCH_BRANCH)"
     pytorch_vision_arg="$(extract_arg_default PYTORCH_VISION_BRANCH)"
     pytorch_audio_arg="$(extract_arg_default PYTORCH_AUDIO_BRANCH)"
+    triton_arg="${triton_arg:-$(extract_arg_default TRITON_VERSION)}"
+    pytorch_arg="${pytorch_arg:-$(extract_arg_default TORCH_VERSION)}"
+    pytorch_vision_arg="${pytorch_vision_arg:-$(extract_arg_default TORCHVISION_VERSION)}"
+    pytorch_audio_arg="${pytorch_audio_arg:-$(extract_arg_default TORCHAUDIO_VERSION)}"
     fa_arg="$(extract_arg_default FA_BRANCH)"
     aiter_arg="$(extract_arg_default AITER_BRANCH)"
     mori_arg="$(extract_arg_default MORI_BRANCH)"
@@ -431,6 +446,7 @@ build_base_image() {
     dependency_summary="base=${base_image_arg},rocm=${rocm_version},python=${python_version_arg},pytorch=${pytorch_arg},torchvision=${pytorch_vision_arg},torchaudio=${pytorch_audio_arg},triton=${triton_arg},flash-attn=${fa_arg},aiter=${aiter_arg},mori=${mori_arg},pytorch-rocm-arch=${pytorch_rocm_arch_arg}"
     trusted_content_tag=$(trusted_base_content_ref "${base_hash}" "${metadata_version}")
     scoped_content_tag=$(scoped_base_content_ref "${base_hash}" "${metadata_version}")
+    preview_content_tag="${scoped_content_tag}"
     # Preview writes share an exact-content namespace so stacked PRs can reuse
     # identical images. They may import the trusted ref, but never overwrite it.
     if is_trusted_main_build; then
@@ -471,7 +487,8 @@ build_base_image() {
         if [[ "${trusted_content_tag}" == "${scoped_content_tag}" ]]; then
             immutable_ref=$(find_matching_base_content_ref \
                 "${base_hash}" "${metadata_version}" \
-                "${trusted_content_tag}" "${stable_tag}") || reuse_status=$?
+                "${trusted_content_tag}" "${preview_content_tag}" \
+                "${stable_tag}") || reuse_status=$?
         else
             immutable_ref=$(find_matching_base_content_ref \
                 "${base_hash}" "${metadata_version}" \
@@ -490,8 +507,9 @@ build_base_image() {
 
     if [[ ${reuse_status} -eq 0 && -n "${immutable_ref}" ]]; then
         echo "Reusing ROCm base image with matching content: ${immutable_ref}"
-        if [[ "${immutable_ref%@*}" == "${stable_tag}" \
-            && "${writable_content_tag}" != "${stable_tag}" ]]; then
+        if [[ "${immutable_ref%@*}" != "${writable_content_tag}" \
+            && ( "${immutable_ref%@*}" == "${stable_tag}" \
+                || "${immutable_ref%@*}" == "${preview_content_tag}" ) ]]; then
             echo "Backfilling ROCm base content tag: ${writable_content_tag}"
             docker buildx imagetools create --prefer-index=false \
                 -t "${writable_content_tag}" "${immutable_ref}"
@@ -565,6 +583,13 @@ build_base_image() {
 }
 
 main() {
+    if using_custom_rocm_dockerfiles; then
+        if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then
+            echo "Custom ROCm Dockerfiles require ROCM_BASE_REFRESH_SKIP=0" >&2
+            return 2
+        fi
+        validate_rocm_dockerfile "${DOCKERFILE}" || return $?
+    fi
     metadata_set "rocm-base-refresh" "0"
 
     if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then
diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml
index 9561fec5e04e..326fa3a2cee2 100644
--- a/.buildkite/test-amd.yaml
+++ b/.buildkite/test-amd.yaml
@@ -114,6 +114,7 @@ steps:
 - label: ":computer: (CPU) Basic Models Other" # TBD
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   no_gpu: true
   optional: true
@@ -124,12 +125,15 @@ steps:
   - tests/models/test_vision.py
   - tests/models/test_adapters.py
   - tests/models/transformers/fusers/
+  - tests/model_executor/layers/test_activation.py
+  - tests/models/test_qwen3_5_mtp_config.py
   commands:
-  - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/
+  - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py
 
 - label: ":computer: (CPU) Multimodal Processor Shard %N" # TBD
   timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   no_gpu: true
   parallelism: 6
@@ -145,6 +149,7 @@ steps:
 - label: ":computer: (CPU) V1 Others" # TBD
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   no_gpu: true
   optional: true
@@ -165,6 +170,7 @@ steps:
 - label: ":computer: (CPU) Async Engine, Inputs, Utils, Worker, Config" # TBD
   timeout_in_minutes: 80
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   no_gpu: true
   optional: true
@@ -204,6 +210,7 @@ steps:
 - label: ":computer: (CPU) Rust Frontend Cargo Style + Clippy" # TBD
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   no_gpu: true
   optional: true
@@ -214,11 +221,12 @@ steps:
   - .buildkite/test_areas/rust_frontend_cargo.yaml
   - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
   commands:
-  - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
+  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
 
 - label: ":computer: (CPU) Rust Frontend Cargo" # TBD
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   no_gpu: true
   optional: true
@@ -229,11 +237,12 @@ steps:
   - .buildkite/test_areas/rust_frontend_cargo.yaml
   - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
   commands:
-  - bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
+  - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
 
 - label: ":computer: (CPU) Docker Build Metadata" # TBD
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   no_gpu: true
   optional: true
@@ -264,6 +273,7 @@ steps:
 - label: ":amd: (MI250) PyTorch Fullgraph Smoke" # TBD
   timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -281,8 +291,9 @@ steps:
 #--------------------------------------------------------  mi250 · distributed  --------------------------------------------------------#
 
 - label: ":amd: (MI250) Pipeline + Context Parallelism" # TBD
-  timeout_in_minutes: 80
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_4
   num_gpus: 4
   working_dir: "/vllm-workspace/tests"
@@ -297,15 +308,20 @@ steps:
   - tests/distributed/
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - tests/v1/distributed/test_pp_dp_v2.py
+  optional: true
   commands:
   - pytest -v -s distributed/test_pp_cudagraph.py
   - pytest -v -s distributed/test_pipeline_parallel.py
+  - pytest -v -s v1/distributed/test_pp_dp_v2.py
 
 #----------------------------------------------------------  mi250 · kernels  ----------------------------------------------------------#
 
-- label: ":amd: (MI250) Helion Kernels" # TBD
-  timeout_in_minutes: 40
+- label: ":amd: (MI250) Helion Kernels Shard %N"
+  timeout_in_minutes: 45
+  parallelism: 2
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -316,26 +332,46 @@ steps:
   - vllm/platforms/rocm.py
   commands:
   - pip install helion==1.4.0
-  - pytest -v -s kernels/helion/
+  - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 #-----------------------------------------------------  mi250 · models / language  -----------------------------------------------------#
 
 - label: ":amd: (MI250) Language Models (PPL)" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
   - tests/models/language/generation_ppl_test
+  - vllm/model_executor/models/qwen3_5.py
+  - vllm/model_executor/models/qwen3_5_mtp.py
+  - vllm/transformers_utils/configs/qwen3_5.py
+  - vllm/transformers_utils/configs/qwen3_5_moe.py
+  - vllm/model_executor/models/qwen2.py
+  - vllm/model_executor/models/qwen3.py
+  - vllm/model_executor/models/qwen3_next.py
+  - vllm/model_executor/models/qwen3_next_mtp.py
+  - vllm/model_executor/layers/fla/ops/
+  - vllm/_aiter_ops.py
+  - vllm/v1/attention/backends/triton_attn.py
+  - vllm/v1/attention/backends/rocm_attn.py
+  - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+  - vllm/v1/attention/backends/rocm_aiter_fa.py
+  - vllm/v1/attention/backends/flex_attention.py
+  - vllm/v1/attention/ops/
+  - vllm/platforms/rocm.py
+  optional: true
   commands:
   - pytest -v -s models/language/generation_ppl_test
 
 #----------------------------------------------------  mi250 · models / multimodal  ----------------------------------------------------#
 
 - label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -351,6 +387,7 @@ steps:
 - label: ":amd: (MI250) Batch Invariance" # TBD
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -367,8 +404,9 @@ steps:
   - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
 
 - label: ":amd: (MI250) CUDAGraph" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -378,13 +416,18 @@ steps:
   - vllm/config/compilation.py
   - vllm/compilation
   - vllm/platforms/rocm.py
+  - vllm/v1/worker/encoder_cudagraph.py
+  - vllm/v1/worker/encoder_cudagraph_defs.py
   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_breakable_cudagraph.py
+  - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
 
 - label: ":amd: (MI250) E2E Core" # TBD
   timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -393,11 +436,12 @@ steps:
   - tests/v1/e2e/general/
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
+  - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
 
 - label: ":amd: (MI250) E2E Scheduling" # TBD
   timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -411,6 +455,7 @@ steps:
 - label: ":amd: (MI250) V1 Engine" # TBD
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -423,6 +468,7 @@ steps:
   - vllm/engine/arg_utils.py
   - vllm/transformers_utils/config.py
   - vllm/platforms/rocm.py
+  - vllm/v1/engine/
   commands:
   - pytest -v -s v1/engine/test_preprocess_error_handling.py
   - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
@@ -432,6 +478,7 @@ steps:
 - label: ":amd: (MI250) Spec Decode Draft Model" # TBD
   timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -447,8 +494,9 @@ steps:
   - pytest -v -s v1/e2e/spec_decode/draft_model/
 
 - label: ":amd: (MI250) Spec Decode Speculators + MTP" # TBD
-  timeout_in_minutes: 65
+  timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  dind: false
   agent_pool: mi250_1
   optional: true
   working_dir: "/vllm-workspace/tests"
@@ -461,6 +509,7 @@ steps:
   - vllm/transformers_utils/configs/speculators/
   - tests/v1/e2e/spec_decode/
   - vllm/platforms/rocm.py
+  - vllm/v1/attention/backends/
   commands:
   - pytest -v -s v1/e2e/spec_decode/speculators/
   - pytest -v -s v1/e2e/spec_decode/mtp/
@@ -468,7 +517,7 @@ steps:
 #-------------------------------------------------------------  mi250 · misc  ------------------------------------------------------------#
 
 - label: ":amd: (MI300) Python-only Installation" # TBD
-  timeout_in_minutes: 25
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -501,14 +550,16 @@ steps:
   - tests/basic_correctness/test_basic_correctness
   - tests/basic_correctness/test_cpu_offload
   - tests/basic_correctness/test_mem.py
+  optional: true
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s basic_correctness/test_mem.py
   - VLLM_TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py
   - pytest -v -s basic_correctness/test_cpu_offload.py
 
-- label: ":amd: (MI300) Distributed Models" # TBD
+- label: ":amd: (MI300) Distributed Models Shard %N"
   timeout_in_minutes: 80
+  parallelism: 3
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_2
@@ -516,6 +567,11 @@ steps:
   optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
+  - vllm/models/hy_v4/
+  - vllm/config/
+  - vllm/v1/worker/gpu/
+  - vllm/v1/attention/ops/
+  - vllm/v1/attention/backend.py
   - vllm/model_executor/model_loader/sharded_state_loader.py
   - vllm/model_executor/models/
   - vllm/model_executor/layers/
@@ -526,20 +582,17 @@ steps:
   - tests/basic_correctness/
   - tests/model_executor/model_loader/test_sharded_state_loader.py
   - tests/models/
+  - vllm/v1/worker/
+  - tests/conftest.py
+  - tests/utils.py
+  - tests/models/test_vision.py
   commands:
-  - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
-  - HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
-  - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/language -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
-  - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
-  - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=2)'
+  - case "$$BUILDKITE_PARALLEL_JOB" in 0) TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' && HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' ;; 1) pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' && pytest models/language -v -s -m 'distributed(num_gpus=2)' && VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' ;; 2) pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py && pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' && pytest -v -s models/test_vision.py -m 'distributed(num_gpus=2)' ;; *) exit 2 ;; esac
 
 #--------------------------------------------------------  mi300 · benchmarks  ---------------------------------------------------------#
 
 - label: ":amd: (MI300) Benchmarks CLI" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -548,13 +601,14 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/benchmarks/
+  - rust/src/bench/tests/python_serve_flags.txt
   commands:
   - pytest -v -s benchmarks/
 
 #----------------------------------------------------------  mi300 · compile  ----------------------------------------------------------#
 
 - label: ":amd: (MI300) PyTorch Compilation" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -570,6 +624,29 @@ steps:
   - csrc/
   - tests/compile
   - vllm/platforms/rocm.py
+  - vllm/__init__.py
+  - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/env_override.py
+  - vllm/envs.py
+  - vllm/forward_context.py
+  - vllm/inputs/
+  - vllm/ir/
+  - vllm/kernels/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/transformers_utils/
+  - vllm/triton_utils/
+  - vllm/utils/
+  - vllm/v1/
   commands:
   - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
 
@@ -590,6 +667,7 @@ steps:
   - tests/compile/fusions_e2e/
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  optional: true
   commands:
   - rocm-smi
   - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
@@ -609,6 +687,7 @@ steps:
   - tests/compile/fusions_e2e/
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  optional: true
   commands:
   - rocm-smi
   # Run all models and attn backends but only Inductor partition and native custom ops
@@ -625,11 +704,37 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/compile/passes
+  - vllm/platforms/rocm.py
+  - vllm/__init__.py
+  - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
+  - vllm/compilation/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/env_override.py
+  - vllm/envs.py
+  - vllm/forward_context.py
+  - vllm/inputs/
+  - vllm/ir/
+  - vllm/kernels/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/transformers_utils/
+  - vllm/triton_utils/
+  - vllm/utils/
+  - vllm/v1/
+  optional: true
   commands:
   - pytest -s -v compile/passes --ignore compile/passes/distributed
 
 - label: ":amd: (MI300) PyTorch Fullgraph" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -643,11 +748,33 @@ steps:
   - csrc/
   - tests/compile
   - vllm/platforms/rocm.py
+  - vllm/__init__.py
+  - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/env_override.py
+  - vllm/envs.py
+  - vllm/forward_context.py
+  - vllm/inputs/
+  - vllm/ir/
+  - vllm/kernels/
+  - vllm/logger.py
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/transformers_utils/
+  - vllm/triton_utils/
+  - vllm/utils/
+  - vllm/v1/
   commands:
-  - pytest -v -s compile/fullgraph/test_full_graph.py
+  - find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'
 
 - label: ":amd: (MI300) PyTorch Nightly Dependency Override Check" # TBD
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -698,7 +825,7 @@ steps:
   - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "dynamo_partition"
 
 - label: ":amd: (MI300) Distributed Compile + RPC" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_2
@@ -716,9 +843,9 @@ steps:
   - tests/compile/test_wrapper.py
   - tests/entrypoints/llm/test_collective_rpc.py
   - vllm/platforms/rocm.py
+  optional: true
   commands:
   - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
   - pytest -v -s ./compile/test_wrapper.py
 
 #-----------------------------------------------------------  mi300 · cuda  ------------------------------------------------------------#
@@ -732,9 +859,16 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/cuda
+  - vllm/envs.py
+  - vllm/logger.py
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/utils/
+  optional: true
   commands:
   - pytest -v -s cuda/test_cuda_context.py
   - pytest -v -s cuda/test_platform_no_cuda_init.py
+  - pytest -v -s cuda/test_cuda_compatibility_path.py
 
 #--------------------------------------------------------  mi300 · detokenizer  --------------------------------------------------------#
 
@@ -750,6 +884,19 @@ steps:
   - tests/detokenizer
   - tests/multimodal
   - tests/utils_
+  - vllm/assets/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
   commands:
   - pytest -v -s detokenizer
   - pytest -v -s -m 'not cpu_test' multimodal
@@ -776,7 +923,7 @@ steps:
   - pytest -v -s distributed/test_shm_storage.py
 
 - label: ":amd: (MI300) EPLB Algorithm" # TBD
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -808,7 +955,7 @@ steps:
   - pytest -v -s distributed/test_eplb_spec_decode.py
 
 - label: ":amd: (MI300) Distributed Features" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_2
@@ -825,6 +972,26 @@ steps:
   - examples/features/data_parallel/data_parallel_offline.py
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - csrc/custom_quickreduce.cu
+  - csrc/ops.h
+  - csrc/torch_bindings.cpp
+  - vllm/model_executor/layers/
+  - vllm/entrypoints/llm.py
+  - vllm/config/parallel.py
+  - vllm/v1/engine/
+  - vllm/v1/executor/
+  - vllm/v1/worker/
+  - vllm/_custom_ops.py
+  - vllm/envs.py
+  - tests/distributed/test_rocm_aiter_custom_ar.py
+  - tests/distributed/test_rocm_quick_reduce.py
+  - tests/distributed/test_quick_all_reduce.py
+  - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
+  - tests/utils.py
+  - examples/rl/rlhf_async_new_apis.py
+  - tests/distributed/test_weight_transfer.py
+  - tests/distributed/test_packed_tensor.py
+  optional: true
   commands:
   - pytest -v -s tests/distributed/test_context_parallel.py
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
@@ -833,6 +1000,10 @@ steps:
   - pytest -v -s tests/v1/distributed/test_dbo.py
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
   - pytest -v -s tests/distributed/test_packed_tensor.py
+  - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
+  - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
+  - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
+  - pytest -v -s tests/distributed/test_quick_all_reduce.py
 
 - label: ":amd: (MI300) Distributed" # TBD
   timeout_in_minutes: 40
@@ -987,7 +1158,7 @@ steps:
 #----------------------------------------------------------  mi300 · engine  -----------------------------------------------------------#
 
 - label: ":amd: (MI300) Engine Core" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1002,6 +1173,17 @@ steps:
   - tests/test_vllm_port
   - tests/jit_monitor/test_hooks.py
   - tests/jit_monitor/test_hooks_gpu.py
+  - vllm/compilation/
+  - vllm/config/
+  - vllm/engine/
+  - vllm/entrypoints/logger.py
+  - vllm/envs.py
+  - vllm/logger.py
+  - vllm/logging_utils/
+  - vllm/platforms/
+  - vllm/sequence.py
+  - vllm/triton_utils/
+  - vllm/utils/
   commands:
   - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py
 
@@ -1043,8 +1225,9 @@ steps:
   - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
   - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
 
-- label: ":amd: (MI300) Entrypoints Integration (API Server)" # TBD
+- label: ":amd: (MI300) Entrypoints Integration (API Server) %N"
   timeout_in_minutes: 65
+  parallelism: 4
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1057,12 +1240,12 @@ steps:
   - tests/entrypoints/scale_out
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
-  - pytest -v -s entrypoints/scale_out
+  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc; fi
+  - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
-- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD
-  timeout_in_minutes: 40
+- label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)" # TBD
+  timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1074,10 +1257,11 @@ steps:
   - tests/entrypoints/openai
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
+  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
 
-- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD
-  timeout_in_minutes: 55
+- label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)" # TBD
+  timeout_in_minutes: 70
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1090,7 +1274,6 @@ steps:
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/openai/chat_completion
-  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
 - label: ":amd: (MI300) Entrypoints Integration (API Server Generate)" # TBD
   timeout_in_minutes: 65
@@ -1108,7 +1291,6 @@ steps:
   - tests/entrypoints/cohere
   - tests/entrypoints/generate
   commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s tool_use
   - pytest -v -s entrypoints/tool_parsers
   - pytest -v -s entrypoints/generate
@@ -1116,7 +1298,7 @@ steps:
   - pytest -v -s entrypoints/cohere
 
 - label: ":amd: (MI300) Entrypoints Integration (Responses API)" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1127,7 +1309,6 @@ steps:
   - vllm/
   - tests/entrypoints/openai/responses
   commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/openai/responses
 
 - label: ":amd: (MI300) Entrypoints Integration (Speech to Text)" # TBD
@@ -1140,6 +1321,7 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/entrypoints/speech_to_text
+  optional: true
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/speech_to_text
@@ -1159,23 +1341,8 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/multimodal
 
-- label: ":amd: (MI300) Entrypoints Integration (Pooling)" # TBD
-  timeout_in_minutes: 65
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  optional: true
-  fast_check: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/pooling
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/pooling
-
 - label: ":amd: (MI300) OpenAI API Correctness" # TBD
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1194,6 +1361,23 @@ steps:
   - bash ../tools/install_torchcodec_rocm.sh || exit 1
   - pytest -s entrypoints/openai/correctness/
 
+- label: ":amd: (MI300) Scale-out EC Connector E2E"
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/ec_transfer/
+  - vllm/entrypoints/scale_out/
+  - tests/entrypoints/scale_out/ec_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
+
 #-----------------------------------------------------------  mi300 · evals  -----------------------------------------------------------#
 
 - label: ":amd: (MI300) DeepSeek V2-Lite Prefetch Offload Accuracy" # TBD
@@ -1214,11 +1398,12 @@ steps:
   - .buildkite/scripts/scheduled_integration_test/
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  optional: true
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030
 
 - label: ":amd: (MI300) LM Eval Small Models" # TBD
-  timeout_in_minutes: 40
+  timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1232,11 +1417,12 @@ steps:
   - vllm/v1/attention/selector.py
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  optional: true
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
 
 - label: ":amd: (MI300) MRCR Eval Small Models" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1269,11 +1455,12 @@ steps:
   - vllm/v1/attention/selector.py
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - .buildkite/lm-eval-harness/
   commands:
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt
 
 - label: ":amd: (MI300) Multimodal Accuracy Eval (Small Models)" # TBD
-  timeout_in_minutes: 30
+  timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1289,7 +1476,7 @@ steps:
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
 
 - label: ":amd: (MI300) GPQA Eval (GPT-OSS)" # TBD
-  timeout_in_minutes: 30
+  timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_2
@@ -1311,7 +1498,7 @@ steps:
     - uv pip install --system 'gpt-oss[eval]==0.0.5'
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt
 
-- label: ":amd: (MI300) LM Eval Small Models FP8 + Mixed" # TBD
+- label: ":amd: (MI300) LM Eval Small Models Distributed"
   timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
@@ -1327,6 +1514,8 @@ steps:
   - vllm/v1/attention/selector.py
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt
+  optional: true
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
 
@@ -1511,7 +1700,7 @@ steps:
 #---------------------------------------------------------  mi300 · examples  ----------------------------------------------------------#
 
 - label: ":amd: (MI300) Examples" # TBD
-  timeout_in_minutes: 50
+  timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1523,27 +1712,27 @@ steps:
   - examples/
   - vllm/platforms/rocm.py
   commands:
-    - pip install tensorizer
     # Basic
-    - python3 basic/offline_inference/chat.py
-    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-    - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
-    - python3 basic/offline_inference/classify.py
-    - python3 basic/offline_inference/embed.py
-    - python3 basic/offline_inference/score.py
     # Multi-modal models
-    - python3 generate/multimodal/audio_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
     # Pooling models
-    - python3 pooling/embed/vision_embedding_offline.py --seed 0
     # Features demo
-    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 deployment/llm_engine_example.py
-    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
-    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
-    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
+  - pip install --no-deps tensorizer
+  - python3 basic/offline_inference/chat.py
+  - python3 basic/offline_inference/generate.py --model facebook/opt-125m
+  - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
+  - python3 basic/offline_inference/classify.py
+  - python3 basic/offline_inference/embed.py
+  - python3 basic/offline_inference/score.py
+  - python3 generate/multimodal/audio_language_offline.py --seed 0
+  - python3 generate/multimodal/vision_language_offline.py --seed 0
+  - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
+  - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+  - python3 pooling/embed/vision_embedding_offline.py --seed 0
+  - python3 features/automatic_prefix_caching/prefix_caching_offline.py
+  - python3 deployment/llm_engine_example.py
+  - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
+  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
 
 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#
 
@@ -1559,34 +1748,14 @@ steps:
   - vllm/kernels
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - tests/ir/
+  - tests/kernels/ir/
   commands:
   - pytest -v -s tests/ir
   - pytest -v -s tests/kernels/ir
 
-- label: ":amd: (MI300) MLA Kernels" # TBD
-  timeout_in_minutes: 25
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  working_dir: "/vllm-workspace/"
-  source_file_dependencies:
-  - vllm/v1/attention/backends/mla/rocm_aiter_mla.py
-  - vllm/v1/attention/selector.py
-  - vllm/platforms/rocm.py
-  - vllm/_aiter_ops.py
-  - tests/kernels/attention/test_rocm_aiter_mla_decode.py
-  - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
-  - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
-  - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
-  commands:
-  - rocm-smi
-  - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
-  - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
-  - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
-  - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
-
 - label: ":amd: (MI300) Attention Kernels Shard %N" # TBD
-  timeout_in_minutes: 70
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1601,26 +1770,9 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/envs.py
   - vllm/platforms/rocm.py
+  - vllm/utils/flashinfer.py
   commands:
-  - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-
-- label: ":amd: (MI300) Core Operation Kernels Shard %N" # TBD
-  timeout_in_minutes: 40
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  parallelism: 3
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - csrc/
-  - tests/kernels/core
-  - tests/kernels/test_top_k_per_row.py
-  - tests/kernels/test_concat_mla_q.py
-  - vllm/model_executor/layers/rotary_embedding/
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":amd: (MI300) MiniMax Reduce RMS Kernels" # TBD
   timeout_in_minutes: 45
@@ -1637,6 +1789,12 @@ steps:
   - vllm/platforms/rocm.py
   - tests/kernels/core/test_minimax_reduce_rms.py
   - tests/utils.py
+  - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu
+  - csrc/libtorch_stable/minimax_reduce_rms_kernel.h
+  - csrc/libtorch_stable/ops.h
+  - csrc/libtorch_stable/torch_bindings.cpp
+  - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py
+  optional: true
   commands:
   - pytest -v -s kernels/core/test_minimax_reduce_rms.py
 
@@ -1660,8 +1818,50 @@ steps:
   commands:
   - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py
 
+- label: ":amd: (MI355) Kimi K3" # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/models/kimi_k3/
+  - csrc/libtorch_stable/kimi_k3/
+  - tests/models/kimi_k3/
+  - vllm/platforms/rocm.py
+  commands:
+  # Mirrors the ":amd: (MI355) Kimi K3" step generated from
+  # .buildkite/test_areas/models_basic.yaml. The ROCm KDA/AttnRes kernels are
+  # gated on gfx950, so this has to be MI355. The two kernels/ files the NVIDIA
+  # step runs are CUDA-only (they import cuda.bindings.driver), so they are
+  # deliberately left out here.
+  - pytest -v -s models/kimi_k3
+
+- label: ":amd: (MI355) Kimi K3 Multi-GPU" # TBD
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/models/kimi_k3/
+  - tests/models/kimi_k3/
+  - vllm/platforms/rocm.py
+  commands:
+  # The single-GPU step above skips every @multi_gpu_test case. `-m "distributed"`
+  # selects exactly those: multi_gpu_test() applies pytest.mark.distributed
+  # alongside its skipif (tests/utils.py:2080), so this picks up a new
+  # multi-GPU test anywhere in the directory without editing this list.
+  # Four GPUs run the two tp4 cases in test_amd_latent_moe_runner.py; the tp8
+  # case and the 8/16-GPU cases in test_latent_moe_tail.py still skip.
+  - pytest -v -s -m "distributed" models/kimi_k3
+
 - label: ":amd: (MI300) Mamba Kernels" # TBD
-  timeout_in_minutes: 50
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1670,13 +1870,15 @@ steps:
   source_file_dependencies:
   - csrc/mamba/
   - tests/kernels/mamba
+  - vllm/model_executor/layers/mamba/gdn
   - vllm/model_executor/layers/mamba/ops
+  - vllm/third_party/flash_linear_attention/
   - vllm/platforms/rocm.py
   commands:
   - pytest -v -s kernels/mamba
 
 - label: ":amd: (MI300) MoE Kernels Shard %N" # TBD
-  timeout_in_minutes: 50
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1698,11 +1900,13 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
-- label: ":amd: (MI300) Quantization Kernels" # TBD
-  timeout_in_minutes: 55
+- label: ":amd: (MI300) Quantization Kernels Shard %N"
+  timeout_in_minutes: 120
+  parallelism: 6
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1722,8 +1926,10 @@ steps:
   - vllm/platforms/rocm.py
   - vllm/model_executor/kernels/
   - vllm/v1/attention/backends/rocm_aiter_fa.py
+  - vllm/config/
+  - tests/kernels/quantization/test_rocm_skinny_gemms.py
   commands:
-  - pytest -v -s kernels/quantization
+  - pytest -v -s kernels/quantization -k "not test_nvfp4_dense_emulation" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":amd: (MI300) DeepEP FP8 MoE Kernels" # TBD
   timeout_in_minutes: 55
@@ -1747,7 +1953,7 @@ steps:
 #-----------------------------------------------------------  mi300 · lora  ------------------------------------------------------------#
 
 - label: ":amd: (MI300) LoRA Shard %N" # TBD
-  timeout_in_minutes: 40
+  timeout_in_minutes: 85
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1761,8 +1967,9 @@ steps:
   commands:
   - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
 
-- label: ":amd: (MI300) LoRA TP (Distributed)" # TBD
+- label: ":amd: (MI300) LoRA TP (Distributed) %N"
   timeout_in_minutes: 65
+  parallelism: 4
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_4
@@ -1773,19 +1980,22 @@ steps:
   - vllm/model_executor/layers/fused_moe/
   - tests/lora
   - vllm/platforms/rocm.py
+  optional: true
   commands:
-  - pytest -v -s -x lora/test_chatglm3_tp.py
-  - pytest -v -s -x lora/test_llama_tp.py
-  - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
-  - pytest -v -s -x lora/test_olmoe_tp.py
-  - pytest -v -s -x lora/test_gptoss_tp.py
-  - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
-  - pytest -v -s -x lora/test_gemma4_tp.py
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - case "$$BUILDKITE_PARALLEL_JOB" in 0|1|2|3) ;; *) echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" && exit 1;; esac
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_chatglm3_tp.py; fi
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_llama_tp.py; fi
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_qwen3_with_multi_loras.py; fi
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_olmoe_tp.py; fi
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s -x lora/test_gptoss_tp.py; fi
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_qwen35_densemodel_lora.py; fi
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_gemma4_tp.py; fi
 
 #------------------------------------------------------  mi300 · model_executor  -------------------------------------------------------#
 
 - label: ":amd: (MI300) Model Executor" # TBD
-  timeout_in_minutes: 60
+  timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1799,160 +2009,19 @@ steps:
   - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - vllm/model_executor/warmup
+  - tests/model_executor/test_jit_warmup.py
+  - tests/model_executor/test_jit_warmup_cutedsl_launcher.py
+  - tests/model_executor/test_jit_warmup_triton_launcher.py
   commands:
   - apt-get update && apt-get install -y curl libsodium23
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s model_executor -m '(not slow_test)'
-  - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
-
-#----------------------------------------------------  mi300 · model_runner_v2  -------------------------------------------------------#
-
-- label: ":amd: (MI300) Model Runner V2 Core" # TBD
-  timeout_in_minutes: 55
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - vllm/v1/core/sched/
-  - vllm/v1/attention/
-  - tests/v1/engine/test_llm_engine.py
-  - tests/v1/e2e/
-  - tests/entrypoints/llm/test_struct_output_generate.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
-  - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
-  - pytest -v -s v1/e2e/general/test_context_length.py
-  - pytest -v -s v1/e2e/general/test_min_tokens.py
-  - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
-
-- label: ":amd: (MI300) Model Runner V2 Examples" # TBD
-  timeout_in_minutes: 40
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/examples"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/core/sched/
-  - vllm/v1/worker/gpu_worker.py
-  - examples/basic/offline_inference/
-  - examples/generate/multimodal/
-  - examples/features/
-  - examples/pooling/embed/vision_embedding_offline.py
-  - examples/features/tensorize_vllm_model.py
-  - examples/deployment/llm_engine_example.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pip install tensorizer
-  - python3 basic/offline_inference/chat.py
-  - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-  - python3 generate/multimodal/audio_language_offline.py --seed 0
-  - python3 generate/multimodal/vision_language_offline.py --seed 0
-  - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-  - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
-  - python3 pooling/embed/vision_embedding_offline.py --seed 0
-  - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-  - python3 deployment/llm_engine_example.py
-  - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
-  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
-  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
-
-- label: ":amd: (MI300) Model Runner V2 Distributed" # TBD
-  timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - tests/basic_correctness/test_basic_correctness.py
-  - tests/v1/distributed/test_async_llm_dp.py
-  - tests/v1/distributed/test_eagle_dp.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-
-- label: ":amd: (MI300) Model Runner V2 Pipeline Parallelism" # TBD
-  timeout_in_minutes: 65
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - tests/distributed/test_pipeline_parallel.py
-  - tests/distributed/test_pp_cudagraph.py
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray"
-  - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
-
-- label: ":amd: (MI300) Model Runner V2 Spec Decode" # TBD
-  timeout_in_minutes: 90
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/worker/gpu/
-  - vllm/v1/worker/gpu_worker.py
-  - tests/v1/spec_decode/test_max_len.py
-  - tests/v1/spec_decode/test_rejection_sampler_utils.py
-  - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
-  - tests/v1/e2e/spec_decode/
-  - vllm/platforms/rocm.py
-  commands:
-  - set -x
-  - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
-  - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
-  - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
-  - pytest -v -s v1/e2e/spec_decode/eagle/
-  - pytest -v -s v1/e2e/spec_decode/speculators/
-  - pytest -v -s v1/e2e/spec_decode/mtp/
+  - export PYTHONFAULTHANDLER=1
+  - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
+  - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
 
 #------------------------------------------------------  mi300 · models / basic  -------------------------------------------------------#
 
-- label: ":amd: (MI300) Basic Models (Extra Initialization) Shard %N" # TBD
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  parallelism: 6
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - vllm/model_executor/layers/
-  - tests/models/test_initialization.py
-  - tests/models/registry.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
 - label: ":amd: (MI300) Basic Models (Initialization)" # TBD
   timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1963,11 +2032,12 @@ steps:
   - vllm/
   - tests/models/test_initialization.py
   - tests/models/registry.py
+  optional: true
   commands:
   - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
 
 - label: ":amd: (MI300) Basic Models (Other)" # TBD
-  timeout_in_minutes: 55
+  timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -1978,9 +2048,12 @@ steps:
   - tests/models/test_terratorch.py
   - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
+  - tests/models/test_hyv4_rocm.py
   - tests/models/test_vision.py
+  - tests/models/test_deepseek_v4_vl_rocm.py
   commands:
-  - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
+  - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py
+  - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
   - pytest -v -s models/test_vision.py::test_simple_mrope_vision_model_spatial_merge
 
 #-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#
@@ -1998,8 +2071,9 @@ steps:
   commands:
   - pytest -v -s models/language/pooling_mteb_test
 
-- label: ":amd: (MI300) Language Models (Extended Pooling)" # TBD
-  timeout_in_minutes: 80
+- label: ":amd: (MI300) Language Models (Extended Pooling) Shard %N"
+  timeout_in_minutes: 95
+  parallelism: 4
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2009,10 +2083,10 @@ steps:
   - vllm/
   - tests/models/language/pooling
   commands:
-  - pytest -v -s models/language/pooling -m 'not core_model'
+  - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
 - label: ":amd: (MI300) Language Models (Standard)" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2063,22 +2137,6 @@ steps:
   - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
   - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
-- label: ":amd: (MI300) Language Models (Hybrid) Shard %N" # TBD
-  timeout_in_minutes: 65
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  parallelism: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-  - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-  - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-  - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
 #----------------------------------------------------  mi300 · models / multimodal  ----------------------------------------------------#
 
 - label: ":amd: (MI300) Multimodal Models (Extended Generation 1)" # TBD
@@ -2096,8 +2154,9 @@ steps:
   - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
   - pytest -v -s models/multimodal/test_mapping.py
 
-- label: ":amd: (MI300) Multimodal Models (Extended Generation 2)" # TBD
-  timeout_in_minutes: 65
+- label: ":amd: (MI300) Multimodal Models (Extended Generation 2) Shard %N"
+  timeout_in_minutes: 100
+  parallelism: 4
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2107,11 +2166,11 @@ steps:
   - vllm/
   - tests/models/multimodal/generation
   commands:
-  - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
+  - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
 
 - label: ":amd: (MI300) Multimodal Models (Extended Generation 3)" # TBD
-  timeout_in_minutes: 55
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2124,7 +2183,7 @@ steps:
   - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
 
 - label: ":amd: (MI300) Multimodal Models (Standard) 1: qwen2" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2138,7 +2197,7 @@ steps:
   - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
 
 - label: ":amd: (MI300) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD
-  timeout_in_minutes: 40
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2148,38 +2207,25 @@ steps:
   - vllm/
   - tests/models/multimodal/generation
   - tests/models/multimodal/test_mapping.py
+  - tests/models/multimodal
   commands:
   - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"
   - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
 
-- label: ":amd: (MI300) Multimodal Models (Standard) 4: other + whisper" # TBD
-  timeout_in_minutes: 70
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-  - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
-  - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
-  - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
-  - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
-
-- label: ":amd: (MI300) Multimodal Processor" # 1h 42m
-  timeout_in_minutes: 110
+- label: ":amd: (MI300) Multimodal Processor Shard %N" # TBD
+  timeout_in_minutes: 115
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
   optional: true
+  parallelism: 4
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal
   - tests/models/registry.py
   commands:
-  - pytest -v -s models/multimodal/processing/test_tensor_schema.py
+  - pytest -v -s models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
 - label: ":amd: (MI300) Multimodal Models (Extended Pooling)" # TBD
   timeout_in_minutes: 75
@@ -2195,7 +2241,7 @@ steps:
   - pytest -v -s models/multimodal/pooling -m 'not core_model'
 
 - label: ":amd: (MI300) Multimodal Models (Standard) 2: qwen3 + gemma" # TBD
-  timeout_in_minutes: 50
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2223,7 +2269,9 @@ steps:
   - vllm/platforms/rocm.py
   - tests/models/quantization
   - vllm/model_executor/model_loader/
+  optional: true
   commands:
+  - unset VLLM_USE_V2_MODEL_RUNNER
   - pytest -v -s models/quantization
 
 #--------------------------------------------------  mi300 · models / transformers  ---------------------------------------------------#
@@ -2301,7 +2349,7 @@ steps:
 #----------------------------------------------------------  mi300 · plugins  ----------------------------------------------------------#
 
 - label: ":amd: (MI300) Plugin Integration" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_2
@@ -2314,43 +2362,43 @@ steps:
   - vllm/platforms/rocm.py
   commands:
   # BEGIN: platform plugin and general plugin tests, all the code in-between runs on dummy platform
+  # END: platform plugin tests
+  # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin
+  # END: `io_processor` plugins test
+  # BEGIN: `bge_m3_sparse io_processor` test
+  # END: `bge_m3_sparse io_processor` test
+  # BEGIN: `colbert_query io_processor` test
+  # END: `colbert_query io_processor` test
+  # BEGIN: `stat_logger` plugins test
+  # END: `stat_logger` plugins test
+  # BEGIN: `endpoint` plugins test
+  # END: `endpoint` plugins test
+  # BEGIN: other tests
   - pip install -e ./plugins/vllm_add_dummy_platform
   - pytest -v -s plugins_tests/test_platform_plugins.py
   - pip uninstall vllm_add_dummy_platform -y
-  # END: platform plugin tests
-  # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin
+  - pytest -v -s ./plugins_tests/test_io_processor_plugins.py
   - pip install -e ./plugins/prithvi_io_processor_plugin
-  - pytest -v -s plugins_tests/test_io_processor_plugins.py
   - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py
   - pip uninstall prithvi_io_processor_plugin -y
-  # END: `io_processor` plugins test
-  # BEGIN: `bge_m3_sparse io_processor` test
   - pip install -e ./plugins/bge_m3_sparse_plugin
   - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
   - pip uninstall bge_m3_sparse_plugin -y
-  # END: `bge_m3_sparse io_processor` test
-  # BEGIN: `colbert_query io_processor` test
   - pip install -e ./plugins/colbert_query_plugin
   - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
   - pip uninstall colbert_query_plugin -y
-  # END: `colbert_query io_processor` test
-  # BEGIN: `stat_logger` plugins test
   - pip install -e ./plugins/vllm_add_dummy_stat_logger
   - pytest -v -s plugins_tests/test_stats_logger_plugins.py
   - pip uninstall dummy_stat_logger -y
-  # END: `stat_logger` plugins test
-  # BEGIN: `endpoint` plugins test
   - pip install -e ./plugins/vllm_add_dummy_endpoint_plugin
   - pytest -v -s plugins_tests/test_endpoint_plugins.py
   - pip uninstall vllm_add_dummy_endpoint_plugin -y
-  # END: `endpoint` plugins test
-  # BEGIN: other tests
   - pytest -v -s plugins_tests/test_scheduler_plugins.py
   - pip install -e ./plugins/vllm_add_dummy_model
   - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
-  - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
-  - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
+  - pytest -v -s plugins_tests/test_oot_registration_online.py
+  - pytest -v -s plugins_tests/test_oot_registration_offline.py
+  - pytest -v -s plugins_tests/lora_resolvers
 
 - label: ":amd: (MI300) GGUF Plugin" # TBD
   timeout_in_minutes: 45
@@ -2392,18 +2440,24 @@ steps:
   - tests/entrypoints/serve/instrumentator/test_uds.py
   - tests/v1/sample/test_logprobs_e2e.py
   - vllm/platforms/rocm.py
+  - tests/entrypoints/openai/chat_completion/test_include_reasoning.py
+  - tests/entrypoints/openai/chat_completion/test_serving_chat.py
+  - tests/entrypoints/openai/completion/test_shutdown.py
+  - tests/entrypoints/openai/test_uds.py
   commands:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
+  - pytest -v -s entrypoints/openai/chat_completion/test_include_reasoning.py
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
+  - pytest -v -s entrypoints/openai/chat_completion/test_serving_chat.py -k "test_chat_per_request_metrics_follow_server_flag or test_streaming_reasoning_usage_counts_across_deltas or test_completion_tokens_details"
   - pytest -v -s entrypoints/launchers/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
   - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"
   - pytest -v -s entrypoints/serve/instrumentator/test_uds.py
   - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
 
-- label: ":amd: (MI300) Rust Frontend Serve + Admin Coverage" # TBD
+- label: ":amd: (MI300) Rust Frontend Serve-Admin Coverage"
   timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
@@ -2499,8 +2553,9 @@ steps:
 
 #-------------------------------------------------------  mi300 · quantization  --------------------------------------------------------#
 
-- label: ":amd: (MI300) Quantization" # TBD
+- label: ":amd: (MI300) Quantization Shard %N"
   timeout_in_minutes: 115
+  parallelism: 4
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2514,19 +2569,19 @@ steps:
   - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
   - tests/quantization
   - tests/rocm/test_moe_weight_replay.py
+  optional: true
   commands:
-
   # temporary install here since we need nightly, will move to requirements/test.in
   # after torchao 0.12 release, and pin a working version of torchao nightly here
-
   # since torchao nightly is only compatible with torch nightly currently
   # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now
   # we can only upgrade after this is resolved
   # TODO(jerryzh168): resolve the above comment
-  - uv pip install --system torchao==0.17.0
+  - unset VLLM_USE_V2_MODEL_RUNNER
+  - uv pip install --system torchao==0.18.0
   - uv pip install --system conch-triton-kernels
-  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
-  - pytest -v -s rocm/test_moe_weight_replay.py
+  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi
 
 #---------------------------------------------------------  mi300 · samplers  ----------------------------------------------------------#
 
@@ -2546,8 +2601,9 @@ steps:
   - tests/conftest.py
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - vllm/entrypoints/generate/beam_search
   commands:
-  - pytest -v -s samplers
+  - pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
 
 #------------------------------------------------------------  mi300 · misc  ------------------------------------------------------------#
 
@@ -2560,6 +2616,18 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/test_regression
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  optional: true
   commands:
   - pip install 'modelscope<1.38'
   - pytest -v -s test_regression.py
@@ -2576,6 +2644,8 @@ steps:
   - requirements/
   - setup.py
   - vllm/platforms/rocm.py
+  - .buildkite/scripts/check-ray-compatibility.sh
+  optional: true
   commands:
   - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh
 
@@ -2598,7 +2668,7 @@ steps:
   - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
 
 - label: ":amd: (MI300) E2E Scheduling" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2612,7 +2682,7 @@ steps:
   - pytest -v -s v1/e2e/general/test_async_scheduling.py
 
 - label: ":amd: (MI300) V1 Engine" # TBD
-  timeout_in_minutes: 40
+  timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2622,12 +2692,14 @@ steps:
   - vllm/v1/engine/
   - tests/v1/engine/
   - vllm/platforms/rocm.py
+  - tests/v1/test_tensor_ipc_queue.py
   commands:
   - pytest -v -s v1/engine/test_preprocess_error_handling.py
   - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
+  - pytest -v -s v1/test_tensor_ipc_queue.py
 
 - label: ":amd: (MI300) Spec Decode Draft Model" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2644,7 +2716,7 @@ steps:
   commands:
   - pytest -v -s v1/e2e/spec_decode/draft_model/
 
-- label: ":amd: (MI300) Spec Decode Eagle" # TBD
+- label: ":amd: (MI300) Spec Decode Eagle 1: DeepSeek + Qwen"
   timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
@@ -2660,15 +2732,15 @@ steps:
   - tests/v1/e2e/spec_decode/
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s v1/e2e/spec_decode/eagle/
+  - pytest -v -s v1/e2e/spec_decode/eagle/ -k "deepseek_eagle or qwen3_eagle3"
 
-- label: ":amd: (MI300) Spec Decode N-Gram + Suffix" # TBD
-  timeout_in_minutes: 35
+- label: ":amd: (MI300) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other"
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
   optional: true
-  working_dir: "/vllm-workspace/tests"
+  working_dir: /vllm-workspace/tests
   source_file_dependencies:
   - vllm/v1/spec_decode/
   - vllm/v1/worker/gpu/spec_decode/
@@ -2676,14 +2748,33 @@ steps:
   - vllm/v1/sample/
   - vllm/model_executor/layers/
   - tests/v1/e2e/spec_decode/
-  - tests/spec_decode/
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s v1/e2e/spec_decode/ngram_suffix/
+  - pytest -v -s v1/e2e/spec_decode/eagle/ -k "not deepseek_eagle and not qwen3_eagle3"
+
+
+- label: ":amd: (MI300) Spec Decode N-Gram + Suffix" # TBD
+  timeout_in_minutes: 35
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/spec_decode/
+  - vllm/v1/worker/gpu/spec_decode/
+  - vllm/model_executor/model_loader/
+  - vllm/v1/sample/
+  - vllm/model_executor/layers/
+  - tests/v1/e2e/spec_decode/
+  - tests/spec_decode/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s v1/e2e/spec_decode/ngram_suffix/
   - python3 spec_decode/test_custom_proposer.py
 
 - label: ":amd: (MI300) Spec Decode Speculators + MTP" # TBD
-  timeout_in_minutes: 50
+  timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2698,12 +2789,13 @@ steps:
   - vllm/transformers_utils/configs/speculators/
   - tests/v1/e2e/spec_decode/
   - vllm/platforms/rocm.py
+  - vllm/v1/attention/backends/
   commands:
     - pytest -v -s v1/e2e/spec_decode/speculators/
     - pytest -v -s v1/e2e/spec_decode/mtp/
 
-- label: ":amd: (MI300) Speculators Correctness" # TBD
-  timeout_in_minutes: 35
+- label: ":amd: (MI300) Speculators Correctness Nightly"
+  timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2730,6 +2822,7 @@ steps:
   - tests/evals/gsm8k/
   - tests/v1/spec_decode/test_speculators_correctness.py
   - vllm/platforms/rocm.py
+  - vllm/v1/worker/gpu/spec_decode/dflash/
   commands:
   - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
   - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
@@ -2743,6 +2836,17 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/v1/spec_decode
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/models/qwen4_exp/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  optional: true
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s -m 'not slow_test' v1/spec_decode
@@ -2753,7 +2857,7 @@ steps:
   dind: false
   agent_pool: mi300_2
   num_gpus: 2
-  working_dir: "/vllm-workspace/"
+  working_dir: /vllm-workspace
   source_file_dependencies:
   - vllm/config/speculative.py
   - vllm/distributed/kv_transfer/kv_connector/
@@ -2775,12 +2879,15 @@ steps:
   - vllm/_aiter_ops.py
   - tests/v1/kv_connector/extract_hidden_states_integration/
   - vllm/platforms/rocm.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+  - tests/v1/kv_connector/extract_hidden_states_integration
+  optional: true
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration
+  - pytest -v -s -m 'distributed' tests/v1/kv_connector/extract_hidden_states_integration
 
 - label: ":amd: (MI300) V1 Attention Shard %N" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 125
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_1
@@ -2798,7 +2905,7 @@ steps:
   commands:
   - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
-- label: ":amd: (MI300) V1 Core + KV + Metrics" # TBD
+- label: ":amd: (MI300) V1 Core"
   timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
@@ -2816,23 +2923,204 @@ steps:
   - tests/v1/ec_connector/unit
   - tests/v1/metrics
   - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  optional: true
   commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  # - export HSA_NO_SCRATCH_RECLAIM=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s -m 'not cpu_test' v1/core
+
+- label: ":amd: (MI300) V1 Executor + Worker"
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/simple_kv_offload
+  - tests/v1/worker
+  - tests/v1/streaming_input
+  - tests/v1/kv_connector/unit
+  - tests/v1/ec_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s v1/executor
-  - pytest -v -s v1/kv_offload
-  - pytest -v -s v1/simple_kv_offload
   - pytest -v -s v1/worker
   - pytest -v -s v1/streaming_input
-  - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
-  - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
+  optional: true
+
+- label: ":amd: (MI300) V1 KV Offload"
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/simple_kv_offload
+  - tests/v1/worker
+  - tests/v1/streaming_input
+  - tests/v1/kv_connector/unit
+  - tests/v1/ec_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh
+  - pytest -v -s v1/kv_offload
+  - pytest -v -s v1/simple_kv_offload
+  optional: true
+
+- label: ":amd: (MI300) V1 KV Connectors Shard %N"
+  timeout_in_minutes: 60
+  parallelism: 4
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/simple_kv_offload
+  - tests/v1/worker
+  - tests/v1/streaming_input
+  - tests/v1/kv_connector/unit
+  - tests/v1/ec_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  commands:
+  - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  optional: true
+
+- label: ":amd: (MI300) V1 Metrics + LM Eval"
+  timeout_in_minutes: 65
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/simple_kv_offload
+  - tests/v1/worker
+  - tests/v1/streaming_input
+  - tests/v1/kv_connector/unit
+  - tests/v1/ec_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  commands:
+  - export GIT_CONFIG_COUNT=1
+  - export GIT_CONFIG_KEY_0=http.version
+  - export GIT_CONFIG_VALUE_0=HTTP/1.1
+  - export GIT_TERMINAL_PROMPT=0
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s -m 'not cpu_test' v1/metrics
   - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
-  # - export HSA_NO_SCRATCH_RECLAIM=1
   - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+  optional: true
 
-- label: ":amd: (MI300) V1 Sample + Logits" # TBD
+
+- label: ":amd: (MI300) V1 Sample"
   timeout_in_minutes: 70
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
@@ -2846,41 +3134,56 @@ steps:
   - tests/v1/test_oracle.py
   - tests/v1/test_request.py
   - tests/v1/test_outputs.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s v1/sample
-  - pytest -v -s v1/logits_processors
-  - pytest -v -s v1/test_oracle.py
-  - pytest -v -s v1/test_request.py
-  - pytest -v -s v1/test_outputs.py
 
-- label: ":amd: (MI300) Distributed DP Basic" # TBD
-  timeout_in_minutes: 45
+- label: ":amd: (MI300) V1 Logits + Oracle"
+  timeout_in_minutes: 70
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
-  agent_pool: mi300_2
+  agent_pool: mi300_1
   optional: true
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
+  working_dir: /vllm-workspace/tests
   source_file_dependencies:
+  - vllm/
+  - tests/v1/sample
+  - tests/v1/logits_processors
+  - tests/v1/test_oracle.py
+  - tests/v1/test_request.py
+  - tests/v1/test_outputs.py
+  - vllm/config/
   - vllm/distributed/
   - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/v1/distributed
-  - tests/entrypoints/launchers/api_server/test_multi_api_servers.py
-  - vllm/platforms/rocm.py
+  - vllm/inputs/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
   commands:
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - DP_SIZE=2 pytest -v -s entrypoints/launchers/api_server/test_multi_api_servers.py
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s v1/logits_processors
+  - pytest -v -s v1/test_oracle.py
+  - pytest -v -s v1/test_request.py
+  - pytest -v -s v1/test_outputs.py
+
 
 - label: ":amd: (MI300) Metrics, Tracing" # TBD
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_2
@@ -2891,6 +3194,18 @@ steps:
   - vllm/
   - tests/v1/tracing
   - tests/tracing/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/tracing/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
   commands:
   - "pip install \
       'opentelemetry-sdk>=1.26.0' \
@@ -2901,7 +3216,7 @@ steps:
   - pytest -v -s tracing
 
 - label: ":amd: (MI300) V1 E2E" # TBD
-  timeout_in_minutes: 25
+  timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_2
@@ -2911,31 +3226,84 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/v1/e2e
+  - vllm/compilation/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/envs.py
+  - vllm/forward_context.py
+  - vllm/inputs/
+  - vllm/logger.py
+  - vllm/logging_utils/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/triton_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/v1/e2e/spec_decode/
   commands:
     - >-
       pytest -v -s
       v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism
       v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism
 
-- label: ":amd: (MI300) NixlConnector PD + Spec Decode acceptance" # TBD
-  timeout_in_minutes: 45
+- label: ":amd: (MI300) Mooncake EC TCP E2E"
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_2
   num_gpus: 2
   optional: true
-  working_dir: "/vllm-workspace/tests"
+  working_dir: "/vllm-workspace"
   source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
-  - vllm/v1/worker/kv_connector_model_runner_mixin.py
-  - tests/v1/kv_connector/nixl_integration/
+  - vllm/distributed/ec_transfer/
+  - vllm/config/ec_transfer.py
+  - vllm/config/multimodal.py
+  - vllm/config/vllm.py
+  - vllm/multimodal/
+  - vllm/v1/core/encoder_cache_manager.py
+  - vllm/v1/core/sched/
+  - vllm/v1/engine/
+  - vllm/v1/worker/ec_connector_model_runner_mixin.py
+  - vllm/v1/worker/gpu_model_runner.py
+  - vllm/v1/worker/gpu_worker.py
+  - vllm/v1/worker/gpu/ec_connector.py
+  - vllm/v1/worker/gpu/model_runner.py
+  - vllm/v1/worker/mm_encoder_model_runner.py
+  - vllm/v1/worker/encoder_cudagraph.py
+  - vllm/v1/worker/encoder_cudagraph_defs.py
+  - vllm/v1/attention/ops/vit_attn_wrappers.py
+  - vllm/v1/worker/gpu/model_states/
+  - vllm/v1/worker/gpu/mm/
+  - tests/v1/cudagraph/test_encoder_cudagraph.py
+  - tests/models/multimodal/generation/test_vit_cudagraph.py
+  - examples/disaggregated/disaggregated_encoder/
+  - tests/v1/ec_connector/
+  - requirements/kv_connectors.txt
   - vllm/platforms/rocm.py
+  - requirements/rocm.txt
   commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - export PYTHON_BIN='python'
+  - export PYTHONUNBUFFERED='1'
+  - export VLLM_HOST_IP='127.0.0.1'
+  - export VLLM_USE_V2_MODEL_RUNNER='1'
+  - export E_CUDAGRAPH_MM_ENCODER='1'
+  - export E_MM_ENCODER_ATTN_BACKEND='TRITON_ATTN'
+  - export MOONCAKE_EC_PROTOCOL='tcp'
+  - export USE_MM_PROMPTS='1'
+  - export SKIP_BASELINE='0'
+  - export CONCURRENCY='3'
+  - export REPEAT='2'
+  - export LOG_PATH='/tmp/mooncake-ec-e2e'
+  - export BASELINE_FILE='/tmp/mooncake-ec-e2e/baseline.json'
+  - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
 
 - label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_4
@@ -2951,7 +3319,7 @@ steps:
   - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
 - label: ":amd: (MI300) Distributed DP Extended" # TBD
-  timeout_in_minutes: 60
+  timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_4
@@ -2974,38 +3342,6 @@ steps:
   - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
   - pytest -v -s distributed/test_utils.py
 
-- label: ":amd: (MI300) Distributed NixlConnector PD accuracy" # TBD
-  timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
-- label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" # TBD
-  timeout_in_minutes: 30
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
 - label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD
   timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -3079,48 +3415,21 @@ steps:
   - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
 
 - label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill" # TBD
-  timeout_in_minutes: 30
+  timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   dind: false
   agent_pool: mi300_4
   num_gpus: 4
   optional: true
   source_file_dependencies:
-    - vllm/v1/attention/backends/utils.py
-    - vllm/v1/worker/gpu_model_runner.py
-    - tests/v1/e2e/test_hybrid_chunked_prefill.py
+  - vllm/v1/attention/backends/utils.py
+  - vllm/v1/worker/gpu_model_runner.py
+  - tests/v1/e2e/test_hybrid_chunked_prefill.py
+  - vllm/v1/core/
+  working_dir: /vllm-workspace/tests
   commands:
     - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py
 
-#------------------------------------------------------  mi300 · weight_loading  -------------------------------------------------------#
-
-- label: ":amd: (MI300) Weight Loading Multi-GPU" # TBD
-  timeout_in_minutes: 70
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/weight_loading
-  commands:
-  - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
-
-- label: ":amd: (MI300) Weight Loading Multi-GPU (Large Models)" # TBD
-  timeout_in_minutes: 35
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  dind: false
-  agent_pool: mi300_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/weight_loading
-  commands:
-  - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt
-
 #########################################################################################################################################
 #                                                                                                                                       #
 #                                                         MI355 (gfx950) tests                                                          #
@@ -3144,9 +3453,49 @@ steps:
   commands:
   - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k"
 
-#--------------------------------------------------------  mi355 · distributed  --------------------------------------------------------#
+#----------------------------------------------------------  mi355 · compile  ----------------------------------------------------------#
 
-- label: ":amd: (MI355) Distributed Features" # TBD
+- label: ":amd: (MI355) PyTorch Compilation H100 Cases"
+  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/__init__.py
+  - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
+  - vllm/compilation/
+  - vllm/config/
+  - vllm/distributed/
+  - "!vllm/distributed/kv_transfer/"
+  - vllm/engine/
+  - vllm/env_override.py
+  - vllm/envs.py
+  - vllm/forward_context.py
+  - vllm/inputs/
+  - vllm/ir/
+  - vllm/kernels/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/transformers_utils/
+  - vllm/triton_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/compile/h100/
+  - tests/utils.py
+  - csrc/
+  commands:
+  - pytest -v -s compile/h100/test_startup.py
+
+- label: ":amd: (MI355) Fusion and Compile FP8"
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
@@ -3155,162 +3504,444 @@ steps:
   optional: true
   working_dir: "/vllm-workspace/"
   source_file_dependencies:
-  - csrc/custom_quickreduce.cu
-  - csrc/ops.h
-  - csrc/torch_bindings.cpp
-  - vllm/distributed/
+  - vllm/compilation/
   - vllm/model_executor/layers/
-  - vllm/entrypoints/llm.py
-  - vllm/config/parallel.py
-  - vllm/model_executor/layers/fused_moe/
-  - vllm/v1/engine/
-  - vllm/v1/executor/
-  - vllm/v1/worker/
-  - vllm/v1/distributed/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
+  - vllm/v1/attention/
+  - vllm/distributed/
   - vllm/_aiter_ops.py
-  - vllm/_custom_ops.py
-  - vllm/platforms/rocm.py
-  - vllm/envs.py
-  - examples/offline_inference/data_parallel.py
-  - tests/distributed/test_context_parallel.py
-  - tests/distributed/test_rocm_aiter_custom_ar.py
-  - tests/distributed/test_rocm_quick_reduce.py
-  - tests/distributed/test_quick_all_reduce.py
-  - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
-  - tests/v1/distributed/test_dbo.py
-  - tests/utils.py
+  - tests/compile/passes/
+  - tests/compile/correctness_e2e/test_attn_quant.py
+  - vllm/model_executor/kernels/linear/
+  - vllm/model_executor/models/llama.py
   commands:
-  - pytest -v -s tests/distributed/test_context_parallel.py
-  - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
-  - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
-  - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
-  - pytest -v -s tests/distributed/test_quick_all_reduce.py
-  - pytest -v -s tests/v1/distributed/test_dbo.py
+  - amd-smi
+  - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py
+  - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py
 
-#--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#
+#--------------------------------------------------------  mi355 · distributed  --------------------------------------------------------#
 
-- label: ":amd: (MI355) Entrypoints Integration (API Server)" # TBD
-  timeout_in_minutes: 55
+- label: ":amd: (MI355) Sharded RDT Weight Transfer"
+  timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
+  num_gpus: 1
   optional: true
-  fast_check: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/serve
-  - tests/entrypoints/scale_out
+  - vllm/distributed/weight_transfer/
+  - tests/distributed/test_sharded_rdt_plan.py
+  - tests/distributed/test_sharded_rdt_producer.py
+  - tests/distributed/test_sharded_rdt_trainer.py
+  - vllm/platforms/rocm.py
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
-  - pytest -v -s entrypoints/scale_out
+  - uv pip install --system --no-deps 'ray==2.56.1'
+  - pytest -v -s distributed/test_sharded_rdt_plan.py
+  - pytest -v -s distributed/test_sharded_rdt_producer.py
+  - pytest -v -s distributed/test_sharded_rdt_trainer.py
 
-- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD
-  timeout_in_minutes: 35
+- label: ":amd: (MI355) Distributed DP Basic" # TBD
+  timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_1
-  fast_check: true
+  agent_pool: mi355_2
   optional: true
+  num_gpus: 2
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/openai
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/executor/
+  - vllm/worker/worker_base.py
+  - vllm/v1/engine/
+  - vllm/v1/worker/
+  - tests/v1/distributed
+  - tests/entrypoints/launchers/api_server/test_multi_api_servers.py
+  - vllm/platforms/rocm.py
+  - tests/v1/e2e/general/test_sharded_sampling.py
+  - tests/v1/e2e/spec_decode/test_sharded_sampling.py
+  - tests/v1/e2e/spec_decode/utils.py
   commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
+  - DP_SIZE=2 pytest -v -s entrypoints/launchers/api_server/test_multi_api_servers.py
+  - pytest -v -s v1/e2e/general/test_sharded_sampling.py
+  - pytest -v -s v1/e2e/spec_decode/test_sharded_sampling.py
 
-- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD
-  timeout_in_minutes: 50
+- label: ":amd: (MI355) MoRI EP Healthy Baseline"
+  timeout_in_minutes: 10
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_1
-  fast_check: true
-  optional: true
+  agent_pool: mi355_2
+  num_gpus: 2
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/openai
+  - csrc/moe/
+  - csrc/rocm/
+  - tests/kernels/moe/
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/
+  - vllm/config/
+  - vllm/forward_context.py
+  - vllm/v1/worker/workspace.py
+  - vllm/utils/import_utils.py
+  - vllm/utils/math_utils.py
+  - vllm/utils/torch_utils.py
+  - vllm/platforms/
+  - vllm/_aiter_ops.py
+  optional: true
+  # TODO: Add fault injection, detection and recovery once MoRI supports them.
+  # This job currently checks healthy EP execution only.
   commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion
-  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
+  - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]'
 
-- label: ":amd: (MI355) Entrypoints Integration (API Server Generate)" # TBD
-  timeout_in_minutes: 40
+- label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD
+  timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_1
-  fast_check: true
+  agent_pool: mi355_2
+  num_gpus: 2
   optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/
-  - tests/tool_use
-  - tests/entrypoints/tool_parsers
-  - tests/entrypoints/anthropic
-  - tests/entrypoints/generate
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - vllm/v1/worker/kv_connector_model_runner_mixin.py
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
   commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s tool_use
-  - pytest -v -s entrypoints/tool_parsers
-  - pytest -v -s entrypoints/generate
-  - pytest -v -s entrypoints/anthropic
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
 
-- label: ":amd: (MI355) Entrypoints Integration (Speech to Text)" # TBD
-  timeout_in_minutes: 50
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
+- label: ":amd: (MI355) Distributed NixlConnector PD accuracy" # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_1
-  fast_check: true
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/speech_to_text
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
   commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/speech_to_text
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
-- label: ":amd: (MI355) Entrypoints Integration (Multimodal)"
-  timeout_in_minutes: 40
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
+- label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy"
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_1
-  fast_check: true
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/multimodal
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - tests/v1/kv_connector/nixl_integration/
+  - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
+  - tests/utils.py
+  - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+  - vllm/platforms/rocm.py
+  - requirements/common.txt
+  - requirements/kv_connectors_rocm.txt
   commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/multimodal
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py
 
-- label: ":amd: (MI355) Entrypoints Integration (Pooling)" # TBD
-  timeout_in_minutes: 55
+- label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
+  timeout_in_minutes: 30
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_1
-  fast_check: true
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/pooling
+  - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
+  - tests/v1/kv_connector/mooncake_integration/
+  - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
+  - tests/utils.py
+  - examples/disaggregated/mooncake_connector/
+  - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+  - vllm/platforms/rocm.py
+  - requirements/common.txt
+  - requirements/kv_connectors_rocm.txt
   commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/pooling
-
-#-----------------------------------------------------------  mi355 · evals  -----------------------------------------------------------#
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
 
-- label: ":amd: (MI355) GPQA Eval (GPT-OSS)" # TBD
+- label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" # TBD
   timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_2
-  num_gpus: 2
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+
+- label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP"
+  timeout_in_minutes: 105
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/platforms/rocm.py
+  - vllm/models/kimi_k3/
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+  - tests/v1/kv_connector/nixl_integration/
+  - requirements/kv_connectors_rocm.txt
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - export ENABLE_HMA_FLAG='1'
+  - export DP_EP='1'
+  - export GPU_MEMORY_UTILIZATION='0.9'
+  - export VLLM_SSM_CONV_STATE_LAYOUT='DS'
+  - export PREFILLER_TP_SIZE='2'
+  - export DECODER_TP_SIZE='2'
+  - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct'
+  - export PREFILL_BLOCK_SIZE='2048'
+  - export DECODE_BLOCK_SIZE='2048'
+  - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768'
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
+
+- label: ":amd: (MI355) Distributed Features" # TBD
+  timeout_in_minutes: 90
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - csrc/custom_quickreduce.cu
+  - csrc/ops.h
+  - csrc/torch_bindings.cpp
+  - vllm/distributed/
+  - vllm/model_executor/layers/
+  - vllm/entrypoints/llm.py
+  - vllm/config/parallel.py
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/v1/engine/
+  - vllm/v1/executor/
+  - vllm/v1/worker/
+  - vllm/v1/distributed/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
+  - vllm/platforms/rocm.py
+  - vllm/envs.py
+  - examples/offline_inference/data_parallel.py
+  - tests/distributed/test_context_parallel.py
+  - tests/distributed/test_rocm_aiter_custom_ar.py
+  - tests/distributed/test_rocm_quick_reduce.py
+  - tests/distributed/test_quick_all_reduce.py
+  - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
+  - tests/v1/distributed/test_dbo.py
+  - tests/utils.py
+  - examples/features/data_parallel/data_parallel_offline.py
+  - examples/rl/rlhf_async_new_apis.py
+  - tests/distributed/test_weight_transfer.py
+  - tests/distributed/test_packed_tensor.py
+  commands:
+  - pytest -v -s tests/distributed/test_context_parallel.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
+  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
+  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
+  - pytest -v -s tests/v1/distributed/test_dbo.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
+  - pytest -v -s tests/distributed/test_packed_tensor.py
+  - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
+  - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
+  - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
+  - pytest -v -s tests/distributed/test_quick_all_reduce.py
+
+- label: ":amd: (MI355) Fusion E2E TP2 Quick" # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/"
+  source_file_dependencies:
+  - csrc/quantization/
+  - vllm/compilation/
+  - vllm/distributed/
+  - vllm/model_executor/
+  - vllm/v1/attention/
+  - tests/compile/fusions_e2e/
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+  - >-
+    VLLM_ROCM_USE_AITER=1
+    VLLM_ROCM_USE_AITER_CUSTOM_AR=1
+    pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
+    -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
+    not +quant_fp8 and not +rms_norm and Llama-3 and
+    ROCM_AITER_UNIFIED_ATTN'
+  - >-
+    VLLM_ROCM_USE_AITER=1
+    VLLM_ROCM_USE_AITER_CUSTOM_AR=1
+    pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
+    -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
+    +quant_fp8 and +rms_norm and qwen3 and
+    ROCM_AITER_UNIFIED_ATTN'
+  - >-
+    VLLM_ROCM_USE_AITER=1
+    VLLM_ROCM_USE_AITER_CUSTOM_AR=1
+    pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
+    -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
+    not +quant_fp8 and not +rms_norm and DeepSeek-Coder and
+    ROCM_AITER_MLA'
+
+#--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#
+
+- label: ":amd: (MI355) Entrypoints Integration (API Server) %N"
+  timeout_in_minutes: 65
+  parallelism: 4
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  fast_check: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/serve
+  - tests/entrypoints/scale_out
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc; fi
+  - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+- label: ":amd: (MI355) Entrypoints Integration (OpenAI API completion)" # TBD
+  timeout_in_minutes: 65
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  fast_check: true
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/openai
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
+
+- label: ":amd: (MI355) Entrypoints Integration (OpenAI API chat_completion)" # TBD
+  timeout_in_minutes: 70
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  fast_check: true
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/openai
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/openai/chat_completion
+
+- label: ":amd: (MI355) Entrypoints Integration (API Server Generate)" # TBD
+  timeout_in_minutes: 65
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  fast_check: true
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/tool_use
+  - tests/entrypoints/tool_parsers
+  - tests/entrypoints/anthropic
+  - tests/entrypoints/generate
+  - tests/entrypoints/cohere
+  commands:
+  - pytest -v -s tool_use
+  - pytest -v -s entrypoints/tool_parsers
+  - pytest -v -s entrypoints/generate
+  - pytest -v -s entrypoints/anthropic
+  - pytest -v -s entrypoints/cohere
+
+- label: ":amd: (MI355) Entrypoints Integration (Speech to Text)" # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  fast_check: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/speech_to_text
+  optional: true
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/speech_to_text
+
+- label: ":amd: (MI355) Entrypoints Integration (Multimodal)"
+  timeout_in_minutes: 55
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  fast_check: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/multimodal
+  optional: true
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/multimodal
+
+- label: ":amd: (MI355) Entrypoints Integration (Pooling) Shard %N" # TBD
+  timeout_in_minutes: 65
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  fast_check: true
+  parallelism: 4
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/pooling
+  optional: true
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/pooling --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+#-----------------------------------------------------------  mi355 · evals  -----------------------------------------------------------#
+
+- label: ":amd: (MI355) GPQA Eval (GPT-OSS)" # TBD
+  timeout_in_minutes: 40
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_2
+  num_gpus: 2
   optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -3328,6 +3959,51 @@ steps:
     - uv pip install --system 'gpt-oss[eval]==0.0.5'
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt
 
+- label: ":amd: (MI355) LM Eval Watermarking"
+  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - .buildkite/test_areas/lm_eval.yaml
+  - tests/evals/gsm8k/
+  - vllm/config/watermarking.py
+  - vllm/sampling_params.py
+  - vllm/v1/watermarking/
+  - vllm/v1/worker/gpu/
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
+
+- label: ":amd: (MI355) LM Eval Watermark Feature Combination"
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_2
+  num_gpus: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - .buildkite/test_areas/lm_eval.yaml
+  - tests/evals/gsm8k/
+  - vllm/config/speculative.py
+  - vllm/config/vllm.py
+  - vllm/config/watermarking.py
+  - vllm/sampling_params.py
+  - vllm/model_executor/models/qwen3_dspark.py
+  - vllm/model_executor/models/qwen3_dflash.py
+  - vllm/model_executor/warmup/
+  - vllm/v1/watermarking/
+  - vllm/v1/worker/gpu/
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
+
 - label: ":amd: (MI355) LM Eval Qwen3-5 Models" # TBD
   timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3349,10 +4025,136 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
   commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt
+
+- label: ":amd: (MI355) Qwen3-8-Flash-Next-FP8 Accuracy Eval"
+  timeout_in_minutes: 105
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - csrc/libtorch_stable/gdn/
+  - tests/evals/qwen4_exp/
+  - vllm/model_executor/layers/mamba/
+  - vllm/models/qwen4_exp/
+  - vllm/transformers_utils/configs/qwen4_exp.py
+  - vllm/v1/attention/backends/short_conv_attn.py
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - uv pip install --system 'evalscope==1.10.0'
+  - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt
+
+- label: ":amd: (MI355) LM Eval TurboQuant k8v4"
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - tests/evals/gsm8k/
+  - tests/utils.py
+  - tests/conftest.py
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/backends/triton_attn.py
+  - vllm/v1/attention/ops/triton_unified_attention.py
+  - vllm/v1/attention/selector.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  - vllm/v1/attention/ops/flydsl_turboquant_decode.py
+  - vllm/v1/attention/ops/flydsl_kernels/
+  - vllm/v1/attention/ops/turboquant_soa/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-k8v4-rocm.txt
+
+- label: ":amd: (MI355) LM Eval TurboQuant t4nc"
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - tests/evals/gsm8k/
+  - tests/utils.py
+  - tests/conftest.py
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/backends/triton_attn.py
+  - vllm/v1/attention/ops/triton_unified_attention.py
+  - vllm/v1/attention/selector.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  - vllm/v1/attention/ops/flydsl_turboquant_decode.py
+  - vllm/v1/attention/ops/flydsl_kernels/
+  - vllm/v1/attention/ops/turboquant_soa/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t4nc-rocm.txt
+
+- label: ":amd: (MI355) LM Eval TurboQuant k3v4nc"
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - tests/evals/gsm8k/
+  - tests/utils.py
+  - tests/conftest.py
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/backends/triton_attn.py
+  - vllm/v1/attention/ops/triton_unified_attention.py
+  - vllm/v1/attention/selector.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  - vllm/v1/attention/ops/flydsl_turboquant_decode.py
+  - vllm/v1/attention/ops/flydsl_kernels/
+  - vllm/v1/attention/ops/turboquant_soa/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-k3v4nc-rocm.txt
+
+- label: ":amd: (MI355) LM Eval TurboQuant t3nc"
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - .buildkite/test-amd.yaml
+  - tests/evals/gsm8k/
+  - tests/utils.py
+  - tests/conftest.py
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/backends/triton_attn.py
+  - vllm/v1/attention/ops/triton_unified_attention.py
+  - vllm/v1/attention/selector.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  - vllm/v1/attention/ops/flydsl_turboquant_decode.py
+  - vllm/v1/attention/ops/flydsl_kernels/
+  - vllm/v1/attention/ops/turboquant_soa/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt
 
-- label: ":amd: (MI355) LM Eval Small Models" # TBD
-  timeout_in_minutes: 45
+- label: ":amd: (MI355) LM Eval Small Models Distributed"
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_2
@@ -3368,9 +4170,29 @@ steps:
   - vllm/v1/attention/selector.py
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
 
+- label: ":amd: (MI355) MoE Refactor Integration Shard %N"
+  timeout_in_minutes: 105
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_2
+  num_gpus: 2
+  parallelism: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/platforms/rocm.py
+  - tests/evals/gsm8k/
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/_aiter_ops.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-rocm-shard-$$BUILDKITE_PARALLEL_JOB.txt
+
 - label: ":amd: (MI355) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" # TBD
   timeout_in_minutes: 25
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3414,7 +4236,7 @@ steps:
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-tp4.txt --tp-size=4
 
 - label: ":amd: (MI355) LM Eval Large Models" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_8
@@ -3439,7 +4261,7 @@ steps:
 #---------------------------------------------------------  mi355 · examples  ----------------------------------------------------------#
 
 - label: ":amd: (MI355) Examples" # TBD
-  timeout_in_minutes: 40
+  timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3449,23 +4271,24 @@ steps:
   - vllm/multimodal
   - examples/
   - vllm/platforms/rocm.py
+  optional: true
   commands:
-  - pip install tensorizer
   # Basic
+  # Multi-modal models
+  # Pooling models
+  # Features demo
+  - pip install --no-deps tensorizer
   - python3 basic/offline_inference/chat.py
   - python3 basic/offline_inference/generate.py --model facebook/opt-125m
   - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
   - python3 basic/offline_inference/classify.py
   - python3 basic/offline_inference/embed.py
   - python3 basic/offline_inference/score.py
-  # Multi-modal models
   - python3 generate/multimodal/audio_language_offline.py --seed 0
   - python3 generate/multimodal/vision_language_offline.py --seed 0
   - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
   - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
-  # Pooling models
   - python3 pooling/embed/vision_embedding_offline.py --seed 0
-  # Features demo
   - python3 features/automatic_prefix_caching/prefix_caching_offline.py
   - python3 deployment/llm_engine_example.py
   - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
@@ -3474,6 +4297,23 @@ steps:
 
 #----------------------------------------------------------  mi355 · kernels  ----------------------------------------------------------#
 
+- label: ":amd: (MI355) Attention DiffKV Kernels"
+  timeout_in_minutes: 35
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/attention/ops/triton_unified_attention_diffkv.py
+  - vllm/v1/attention/backends/triton_attn_diffkv.py
+  - tests/kernels/attention/test_triton_unified_attention.py
+  - tests/kernels/attention/test_triton_unified_attention_diffkv.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
+
 - label: ":amd: (MI355) Core Operation Kernels Shard %N" # TBD
   timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3492,8 +4332,10 @@ steps:
   - vllm/model_executor/layers/rotary_embedding/
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - tests/kernels/test_fused_qk_norm_rope_gate.py
+  optional: true
   commands:
-  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_fused_qk_norm_rope_gate.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":amd: (MI355) MiniMax Reduce RMS Kernels" # TBD
   timeout_in_minutes: 45
@@ -3510,6 +4352,12 @@ steps:
   - vllm/platforms/rocm.py
   - tests/kernels/core/test_minimax_reduce_rms.py
   - tests/utils.py
+  - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu
+  - csrc/libtorch_stable/minimax_reduce_rms_kernel.h
+  - csrc/libtorch_stable/ops.h
+  - csrc/libtorch_stable/torch_bindings.cpp
+  - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py
+  optional: true
   commands:
   - pytest -v -s kernels/core/test_minimax_reduce_rms.py
 
@@ -3533,17 +4381,21 @@ steps:
   - vllm/v1/attention/selector.py
   - vllm/platforms/rocm.py
   - vllm/_aiter_ops.py
+  - examples/basic/offline_inference/chat.py
+  - tests/kernels/attention/test_attention_selector.py
+  optional: true
   commands:
-  - rocm-smi
+  - amd-smi
   - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
   - pytest -v -s tests/kernels/attention/test_attention_selector.py
-  - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
 
-- label: ":amd: (MI355) MLA Kernels" # TBD
-  timeout_in_minutes: 25
+- label: ":amd: (MI355) MLA Kernels"
+  timeout_in_minutes: 40
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
   working_dir: "/vllm-workspace/"
   source_file_dependencies:
   - vllm/v1/attention/backends/mla/rocm_aiter_mla.py
@@ -3555,14 +4407,30 @@ steps:
   - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
   - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
   commands:
-  - rocm-smi
+  - amd-smi
   - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
   - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
   - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
   - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
 
+- label: ":amd: (MI355) DeepSeek V4 Fused Kernels"
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+  - vllm/models/deepseek_v4/common/ops/
+  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+
 - label: ":amd: (MI355) Attention Kernels Shard %N" # TBD
-  timeout_in_minutes: 40
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3576,11 +4444,13 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/envs.py
   - vllm/platforms/rocm.py
+  - vllm/utils/flashinfer.py
+  optional: true
   commands:
-  - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":amd: (MI355) MoE Kernels Shard %N" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3600,6 +4470,7 @@ steps:
   - vllm/config
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  optional: true
   commands:
   - pytest -v -s kernels/moe
       --ignore=kernels/moe/test_modular_oai_triton_moe.py
@@ -3636,11 +4507,45 @@ steps:
   - vllm/utils/torch_utils.py
   - vllm/platforms/
   - vllm/_aiter_ops.py
+  - csrc/quantization/cutlass_w8a8/moe/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
-  - pytest -v -s kernels/moe/test_moe_layer.py
+  # MoRI graph cases run in the Fault Tolerance group's healthy EP baseline.
+  - pytest -v -s kernels/moe/test_moe_layer.py -k "not test_moe_layer_mori_graph"
 
-- label: ":amd: (MI355) Quantization Kernels" # TBD
-  timeout_in_minutes: 50
+- label: ":amd: (MI355) MoRI EP Numerics"
+  timeout_in_minutes: 20
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_2
+  num_gpus: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - tests/kernels/moe/test_mori_moe.py
+  - tests/kernels/moe/utils.py
+  - tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+  - tests/kernels/utils.py
+  - tests/utils.py
+  - tests/conftest.py
+  - vllm/utils/torch_utils.py
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/
+  - vllm/config/
+  - vllm/forward_context.py
+  - vllm/v1/worker/workspace.py
+  - vllm/platforms/
+  - vllm/_aiter_ops.py
+  - requirements/rocm-test.txt
+  optional: true
+  commands:
+  - amd-smi
+  - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 pytest -v -s kernels/moe/test_mori_moe.py
+
+- label: ":amd: (MI355) Quantization Kernels Shard %N"
+  timeout_in_minutes: 120
+  parallelism: 6
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3659,8 +4564,34 @@ steps:
   - vllm/platforms/rocm.py
   - vllm/model_executor/kernels/
   - vllm/v1/attention/backends/rocm_aiter_fa.py
+  - vllm/config/
+  - tests/kernels/quantization/test_rocm_skinny_gemms.py
+  optional: true
+  commands:
+  - pytest -v -s kernels/quantization -k "not test_nvfp4_dense_emulation" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+
+- label: ":amd: (MI355) NVFP4 Dense Emulation Shard %N"
+  timeout_in_minutes: 10
+  parallelism: 3
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - tests/kernels/quantization/test_nvfp4_emulation.py
+  - tests/kernels/quantization/nvfp4_utils.py
+  - tests/conftest.py
+  - vllm/model_executor/kernels/linear/
+  - vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+  - vllm/config/kernel.py
+  - vllm/platforms/
+  - vllm/envs.py
+  - vllm/utils/torch_utils.py
+  - requirements/rocm-test.txt
+  optional: true
   commands:
-  - pytest -v -s kernels/quantization
+  - amd-smi
+  - pytest -v -s kernels/quantization/test_nvfp4_emulation.py -k test_nvfp4_dense_emulation --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":amd: (MI355) FP8 MoE Kernels" # TBD
   timeout_in_minutes: 40
@@ -3675,6 +4606,13 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
   - vllm/envs.py
+  - tests/kernels/moe/test_gpt_oss_triton_kernels.py
+  - tests/kernels/moe/test_modular_oai_triton_moe.py
+  - tests/kernels/moe/test_moe.py
+  - tests/kernels/moe/test_block_int8.py
+  - tests/kernels/moe/test_triton_moe_no_act_mul.py
+  - tests/kernels/moe/test_triton_moe_ptpc_fp8.py
+  optional: true
   commands:
     - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py
     - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py
@@ -3684,7 +4622,7 @@ steps:
     - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
 
 - label: ":amd: (MI355) DeepEP FP8 MoE Kernels" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_2
@@ -3702,24 +4640,100 @@ steps:
   commands:
     - pytest -v -s kernels/moe/test_deepep_moe.py
 
+#------------------------------------------------------  mi355 · models / basic  -------------------------------------------------------#
+
+- label: ":amd: (MI355) GLM5Next"
+  timeout_in_minutes: 35
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/models/glm5next/
+  - tests/models/glm5next/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s models/glm5next
+
+- label: ":amd: (MI355) Basic Models (Extra Initialization) Shard %N" # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  parallelism: 6
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/model_executor/models/
+  - vllm/model_executor/layers/
+  - tests/models/test_initialization.py
+  - tests/models/registry.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  optional: true
+  commands:
+  - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
 #-----------------------------------------------------  mi355 · models / language  -----------------------------------------------------#
 
-- label: ":amd: (MI355) Language Models (Extended Generation)" # TBD
+- label: ":amd: (MI355) Language Models (Hybrid) Shard %N" # TBD
   timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
+  parallelism: 2
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/language/generation
+  commands:
+  - export GIT_CONFIG_COUNT=1
+  - export GIT_CONFIG_KEY_0=http.version
+  - export GIT_CONFIG_VALUE_0=HTTP/1.1
+  - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+  - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+  - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+- label: ":amd: (MI355) Granite Language Model Compatibility"
+  timeout_in_minutes: 80
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - "!vllm/distributed/kv_transfer/"
+  - tests/models/language/generation
+  commands:
+  - export GIT_CONFIG_COUNT=1
+  - export GIT_CONFIG_KEY_0=http.version
+  - export GIT_CONFIG_VALUE_0=HTTP/1.1
+  - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+  - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+  - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
+
+- label: ":amd: (MI355) Language Models (Extended Generation)" # TBD
+  timeout_in_minutes: 95
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
   - tests/models/language/generation
+  optional: true
   commands:
   - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'
   - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
   - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
-- label: ":amd: (MI355) Language Models (Extended Pooling)" # TBD
-  timeout_in_minutes: 65
+- label: ":amd: (MI355) Language Models (Extended Pooling) Shard %N"
+  timeout_in_minutes: 95
+  parallelism: 4
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3729,10 +4743,10 @@ steps:
   - vllm/
   - tests/models/language/pooling
   commands:
-  - pytest -v -s models/language/pooling -m 'not core_model'
+  - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
 - label: ":amd: (MI355) Language Models (PPL)" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3757,6 +4771,8 @@ steps:
   - vllm/v1/attention/ops/
   - vllm/platforms/rocm.py
   - tests/models/language/generation_ppl_test
+  - vllm/model_executor/layers/fla/ops/
+  - vllm/
   commands:
   - pytest -v -s models/language/generation_ppl_test
 
@@ -3769,6 +4785,7 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/models/language
+  optional: true
   commands:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and (not slow_test)'
@@ -3776,7 +4793,7 @@ steps:
 #----------------------------------------------------  mi355 · models / multimodal  ----------------------------------------------------#
 
 - label: ":amd: (MI355) Multimodal Models (Extended Generation 1)" # TBD
-  timeout_in_minutes: 60
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3791,7 +4808,7 @@ steps:
   - pytest -v -s models/multimodal/test_mapping.py
 
 - label: ":amd: (MI355) Multimodal Models (Extended Generation 3)" # TBD
-  timeout_in_minutes: 40
+  timeout_in_minutes: 90
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3804,7 +4821,7 @@ steps:
   - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
 
 - label: ":amd: (MI355) Multimodal Models (Extended Pooling)" # TBD
-  timeout_in_minutes: 40
+  timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3817,7 +4834,7 @@ steps:
   - pytest -v -s models/multimodal/pooling -m 'not core_model'
 
 - label: ":amd: (MI355) Multimodal Models (Standard) 1: qwen2" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 65
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3831,7 +4848,7 @@ steps:
   - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
 
 - label: ":amd: (MI355) Multimodal Models (Standard) 4: other + whisper" # TBD
-  timeout_in_minutes: 55
+  timeout_in_minutes: 70
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3843,13 +4860,13 @@ steps:
   commands:
   - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
   - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
-  - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
+  - pytest models/multimodal/generation/test_memory_leak.py -m core_model
   - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
 
 #-----------------------------------------------------  mi355 · models / quantized  -----------------------------------------------------#
 
 - label: ":amd: (MI355) Quantized Models" # TBD
-  timeout_in_minutes: 50
+  timeout_in_minutes: 80
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3860,13 +4877,16 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
   - vllm/model_executor/model_loader/
+  optional: true
   commands:
+  - unset VLLM_USE_V2_MODEL_RUNNER
   - pytest -v -s models/quantization
 
 #-------------------------------------------------------  mi355 · quantization  --------------------------------------------------------#
 
-- label: ":amd: (MI355) Quantization" # TBD
-  timeout_in_minutes: 80
+- label: ":amd: (MI355) Quantization Shard %N"
+  timeout_in_minutes: 115
+  parallelism: 4
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3877,10 +4897,16 @@ steps:
   - tests/quantization
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
+  - vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+  - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+  - tests/rocm/test_moe_weight_replay.py
+  optional: true
   commands:
+  - unset VLLM_USE_V2_MODEL_RUNNER
   - uv pip install --system torchao==0.17.0
   - uv pip install --system conch-triton-kernels
-  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
+  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi
 
 # - label: Quantized MoE Test (B200-MI355) # TBD
 #   timeout_in_minutes: 180
@@ -3909,10 +4935,67 @@ steps:
 #   commands:
 #   - pytest -s -v tests/quantization/test_gfx950_moe.py
 
+#---------------------------------------------------------  mi355 · samplers  ----------------------------------------------------------#
+
+- label: ":amd: (MI355) Samplers Multimodal Beam Search"
+  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/model_executor/layers
+  - vllm/sampling_metadata.py
+  - vllm/v1/sample/
+  - vllm/entrypoints/generate/beam_search/
+  - tests/samplers
+  - tests/conftest.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
+
 #------------------------------------------------------------  mi355 · v1  -------------------------------------------------------------#
 
+- label: ":amd: (MI355) E2E Core Large Memory"
+  timeout_in_minutes: 65
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/v1/
+  - tests/v1/e2e/general/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
+  - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async
+
+- label: ":amd: (MI355) Extract Hidden States Integration Single GPU"
+  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace"
+  source_file_dependencies:
+  - vllm/v1/spec_decode/extract_hidden_states.py
+  - vllm/model_executor/models/extract_hidden_states.py
+  - vllm/transformers_utils/configs/extract_hidden_states.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+  - tests/v1/kv_connector/extract_hidden_states_integration
+  - vllm/platforms/rocm.py
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s -m 'not distributed' tests/v1/kv_connector/extract_hidden_states_integration
+
 - label: ":amd: (MI355) V1 Attention Shard %N" # TBD
-  timeout_in_minutes: 35
+  timeout_in_minutes: 125
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3926,11 +5009,12 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/envs.py
   - vllm/platforms/rocm.py
+  optional: true
   commands:
   - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
-- label: ":amd: (MI355) V1 Core + KV + Metrics" # TBD
-  timeout_in_minutes: 55
+- label: ":amd: (MI355) V1 Core"
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3946,24 +5030,243 @@ steps:
   - tests/v1/kv_connector/unit
   - tests/v1/metrics
   - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/v1/simple_kv_offload
+  - tests/v1/streaming_input
+  - tests/v1/ec_connector/unit
   commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s -m 'not cpu_test' v1/core
+
+- label: ":amd: (MI355) V1 Executor + Worker"
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/worker
+  - tests/v1/cudagraph
+  - tests/v1/kv_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/v1/simple_kv_offload
+  - tests/v1/streaming_input
+  - tests/v1/ec_connector/unit
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s v1/executor
-  - pytest -v -s v1/kv_offload
   - pytest -v -s v1/worker
-  - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
-  - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
+  - pytest -v -s v1/streaming_input
+
+- label: ":amd: (MI355) V1 KV Offload"
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/worker
+  - tests/v1/cudagraph
+  - tests/v1/kv_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/v1/simple_kv_offload
+  - tests/v1/streaming_input
+  - tests/v1/ec_connector/unit
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh
+  - pytest -v -s v1/kv_offload
+  - pytest -v -s v1/simple_kv_offload
+
+- label: ":amd: (MI355) V1 KV Connectors Shard %N"
+  timeout_in_minutes: 60
+  parallelism: 4
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/worker
+  - tests/v1/cudagraph
+  - tests/v1/kv_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/v1/simple_kv_offload
+  - tests/v1/streaming_input
+  - tests/v1/ec_connector/unit
+  commands:
+  - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+  - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+- label: ":amd: (MI355) V1 Metrics + LM Eval"
+  timeout_in_minutes: 65
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/worker
+  - tests/v1/cudagraph
+  - tests/v1/kv_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/entrypoints/pooling/
+  - vllm/inputs/
+  - vllm/lora/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/profiler/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - tests/v1/simple_kv_offload
+  - tests/v1/streaming_input
+  - tests/v1/ec_connector/unit
+  commands:
+  - export GIT_CONFIG_COUNT=1
+  - export GIT_CONFIG_KEY_0=http.version
+  - export GIT_CONFIG_VALUE_0=HTTP/1.1
+  - export GIT_TERMINAL_PROMPT=0
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s -m 'not cpu_test' v1/metrics
   - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
   - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
-- label: ":amd: (MI355) V1 Sample + Logits" # TBD
+- label: ":amd: (MI355) CUDAGraph"
   timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
   optional: true
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/core
+  - tests/v1/executor
+  - tests/v1/kv_offload
+  - tests/v1/worker
+  - tests/v1/cudagraph
+  - tests/v1/kv_connector/unit
+  - tests/v1/metrics
+  - tests/entrypoints/openai/correctness/test_lmeval.py
+  - vllm/v1/cudagraph_dispatcher.py
+  - vllm/config/compilation.py
+  - vllm/compilation
+  - vllm/platforms/rocm.py
+  - vllm/v1/worker/encoder_cudagraph.py
+  - vllm/v1/worker/encoder_cudagraph_defs.py
+  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_breakable_cudagraph.py
+  - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
+
+
+- label: ":amd: (MI355) V1 Sample"
+  timeout_in_minutes: 70
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -3972,15 +5275,56 @@ steps:
   - tests/v1/test_oracle.py
   - tests/v1/test_request.py
   - tests/v1/test_outputs.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
   commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s v1/sample
+
+- label: ":amd: (MI355) V1 Logits + Oracle"
+  timeout_in_minutes: 70
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: /vllm-workspace/tests
+  source_file_dependencies:
+  - vllm/
+  - tests/v1/sample
+  - tests/v1/logits_processors
+  - tests/v1/test_oracle.py
+  - tests/v1/test_request.py
+  - tests/v1/test_outputs.py
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s v1/logits_processors
   - pytest -v -s v1/test_oracle.py
   - pytest -v -s v1/test_request.py
   - pytest -v -s v1/test_outputs.py
 
+
 - label: ":amd: (MI355) V1 Spec Decode" # TBD
-  timeout_in_minutes: 45
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -3988,11 +5332,24 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/v1/spec_decode
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/models/qwen4_exp/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  optional: true
   commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s -m 'not slow_test' v1/spec_decode
 
-- label: ":amd: (MI355) Spec Decode AL MTP + Other Acceptance" # TBD
-  timeout_in_minutes: 180
+- label: ":amd: (MI355) Spec Decode AL MTP + Other Acceptance Nightly"
+  timeout_in_minutes: 55
+  parallelism: 3
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
@@ -4016,50 +5373,90 @@ steps:
   - tests/utils.py
   - tests/v1/e2e/spec_decode/
   - vllm/platforms/rocm.py
+  - vllm/v1/worker/gpu/spec_decode/
+  - tests/v1/e2e/spec_decode/conftest.py
   commands:
-  - pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi
+  - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi
 
-#------------------------------------------------------  mi355 · weight_loading  -------------------------------------------------------#
+#-----------------------------------------------------------  mi355 · misc  ------------------------------------------------------------#
 
-- label: ":amd: (MI355) Weight Loading Multi-GPU" # TBD
-  timeout_in_minutes: 30
+- label: ":amd: (MI355) Regression" # TBD
+  timeout_in_minutes: 25
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
-  agent_pool: mi355_2
-  num_gpus: 2
+  agent_pool: mi355_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
-  - tests/weight_loading
+  - tests/test_regression
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
   commands:
-  - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
+  - pip install 'modelscope<1.38'
+  - pytest -v -s test_regression.py
 
-- label: ":amd: (MI355) Weight Loading Multi-GPU (Large Models)" # TBD
+- label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep'
   timeout_in_minutes: 30
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  working_dir: /vllm-workspace/
+  source_file_dependencies:
+  - vllm/compilation/
+  - vllm/model_executor/layers/linear.py
+  - vllm/model_executor/layers/utils.py
+  - vllm/distributed/
+  - vllm/v1/attention/backends/triton_attn.py
+  - tests/compile/fusions_e2e/
+  - tests/compile/passes/distributed/test_async_tp.py
+  - tests/compile/rocm/test_async_tp.py
+  - tests/compile/correctness_e2e/test_async_tp.py
+  - vllm/model_executor/kernels/linear/
+  - vllm/model_executor/models/qwen3.py
+  commands:
+  - rocm-smi
+  - pytest -v -s tests/compile/rocm/test_async_tp.py
+  - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py
+  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness
   dind: false
   agent_pool: mi355_2
-  working_dir: "/vllm-workspace/tests"
   num_gpus: 2
+  mirror_hardwares:
+  - amdexperimental
+  - amdproduction
+  - amdgfx950nightly
+  - amdmi355
   optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/weight_loading
-  commands:
-  - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt
-
-#-----------------------------------------------------------  mi355 · misc  ------------------------------------------------------------#
 
-- label: ":amd: (MI355) Regression" # TBD
-  timeout_in_minutes: 25
+- label: ":amd: (MI355) MRV2 Sampler JIT Warmup"
+  timeout_in_minutes: 10
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   dind: false
   agent_pool: mi355_1
+  num_gpus: 1
   optional: true
-  working_dir: "/vllm-workspace/tests"
+  working_dir: "/vllm-workspace/"
   source_file_dependencies:
-  - vllm/
-  - tests/test_regression
+  - tests/jit_monitor/
+  - tests/models/utils.py
+  - tests/models/registry.py
+  - tests/utils.py
+  - vllm/v1/worker/gpu/
+  - vllm/v1/worker/gpu_worker.py
+  - vllm/model_executor/warmup/
+  - vllm/v1/sample/ops/topk_topp_sampler.py
+  - vllm/utils/jit_monitor.py
+  - vllm/config/observability.py
+  - vllm/platforms/rocm.py
+  - vllm/envs.py
   commands:
-  - pip install 'modelscope<1.38'
-  - pytest -v -s test_regression.py
+  - pytest -v -s tests/jit_monitor/test_no_runtime_jit_rocm.py --timeout=300
diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml
index 45a9e02fc840..bfa86e5597a4 100644
--- a/.buildkite/test_areas/benchmarks.yaml
+++ b/.buildkite/test_areas/benchmarks.yaml
@@ -34,3 +34,18 @@ steps:
   - vllm/v1/attention/
   commands:
   - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k"
+  mirror:
+    amd:
+      label: ":amd: (MI355) Attention Benchmark Smoke"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 25
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - benchmarks/attention_benchmarks/
+      - vllm/v1/attention/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k"
diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml
index 8eea80068646..2f314aa13893 100644
--- a/.buildkite/test_areas/compile.yaml
+++ b/.buildkite/test_areas/compile.yaml
@@ -67,6 +67,31 @@ steps:
     - pytest -v -s tests/compile/passes/test_silu_mul_quant_manual_fusion.py
     # this runner has 2 GPUs available even though num_devices=2 is not set
     - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Fusion and Compile FP8"
+      dind: false
+      device: mi355_2
+      num_devices: 2
+      working_dir: "/vllm-workspace/"
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/compilation/
+      - vllm/model_executor/layers/
+      - vllm/v1/attention/
+      - vllm/distributed/
+      - vllm/_aiter_ops.py
+      - tests/compile/passes/
+      - tests/compile/correctness_e2e/test_attn_quant.py
+      - vllm/model_executor/kernels/linear/
+      - vllm/model_executor/models/llama.py
+      commands:
+      - amd-smi
+      # MLA static FP8 and distributed AR/RMS are owned by the existing ROCm jobs.
+      - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py
+      - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py
 
 - label: ":nvidia: (H100) Fusion E2E Quick"
   key: fusion-e2e-quick-h100
@@ -92,7 +117,7 @@ steps:
       dind: false
       device: mi355_1
       working_dir: "/vllm-workspace/"
-      timeout_in_minutes: 25
+      timeout_in_minutes: 30
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -133,7 +158,7 @@ steps:
       dind: false
       device: mi355_1
       working_dir: "/vllm-workspace/"
-      timeout_in_minutes: 30
+      timeout_in_minutes: 40
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -181,6 +206,47 @@ steps:
     # Run all models and attn backends but only Inductor partition and native custom ops
     - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
     - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
+  mirror:
+    amd:
+      label: ":amd: (MI355) Fusion E2E TP2 Quick"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd
+      working_dir: /vllm-workspace/
+      source_file_dependencies:
+      - csrc/quantization/
+      - vllm/compilation/
+      - vllm/distributed/
+      - vllm/model_executor/
+      - vllm/v1/attention/
+      - tests/compile/fusions_e2e/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+      - >-
+        VLLM_ROCM_USE_AITER=1
+        VLLM_ROCM_USE_AITER_CUSTOM_AR=1
+        pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
+        -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
+        not +quant_fp8 and not +rms_norm and Llama-3 and
+        ROCM_AITER_UNIFIED_ATTN'
+      - >-
+        VLLM_ROCM_USE_AITER=1
+        VLLM_ROCM_USE_AITER_CUSTOM_AR=1
+        pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
+        -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
+        +quant_fp8 and +rms_norm and qwen3 and
+        ROCM_AITER_UNIFIED_ATTN'
+      - >-
+        VLLM_ROCM_USE_AITER=1
+        VLLM_ROCM_USE_AITER_CUSTOM_AR=1
+        pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
+        -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
+        not +quant_fp8 and not +rms_norm and DeepSeek-Coder and
+        ROCM_AITER_MLA'
 
 - label: ":nvidia: (H100) Fusion E2E TP2 AR-RMS Config Sweep"
   key: fusion-e2e-tp2-ar-rms-config-sweep-h100
@@ -221,6 +287,33 @@ steps:
     - nvidia-smi
     # Run just llama3 (fp8 & bf16) for all config combinations
     - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3"
+  mirror:
+    amd:
+      label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep'
+      dind: false
+      device: mi355_2
+      num_devices: 2
+      working_dir: /vllm-workspace/
+      timeout_in_minutes: 30
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/compilation/
+      - vllm/model_executor/layers/linear.py
+      - vllm/model_executor/layers/utils.py
+      - vllm/distributed/
+      - vllm/v1/attention/backends/triton_attn.py
+      - tests/compile/fusions_e2e/
+      - tests/compile/passes/distributed/test_async_tp.py
+      - tests/compile/rocm/test_async_tp.py
+      - tests/compile/correctness_e2e/test_async_tp.py
+      - vllm/model_executor/kernels/linear/
+      - vllm/model_executor/models/qwen3.py
+      commands:
+      - rocm-smi
+      - pytest -v -s tests/compile/rocm/test_async_tp.py
+      - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py
+      - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness
 
 - label: ":nvidia: (B200) Fusion E2E TP2"
   key: fusion-e2e-tp2-b200
diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml
index f0b1f4af2b80..8ac8e1960ade 100644
--- a/.buildkite/test_areas/cuda.yaml
+++ b/.buildkite/test_areas/cuda.yaml
@@ -23,15 +23,15 @@ steps:
       dind: false
       device: mi355_1
       working_dir: "/vllm-workspace/tests"
-      timeout_in_minutes: 20
+      timeout_in_minutes: 25
       depends_on:
       - image-build-amd
       source_file_dependencies:
       - vllm/
       - tests/cuda
 
-- label: ":nvidia: (H200 MIG 35GB) CUDAGraph"
-  device: h200_35gb
+- label: ":nvidia: (H200 MIG 18GB) CUDAGraph"
+  device: h200_18gb
   key: cudagraph
   timeout_in_minutes: 30
   source_file_dependencies:
@@ -49,8 +49,9 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) CUDAGraph"
+      dind: false
       device: mi250_1
-      timeout_in_minutes: 45
+      timeout_in_minutes: 55
       depends_on:
       - image-build-amd
       source_file_dependencies:
diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml
index 0100b99993ef..1258336009d7 100644
--- a/.buildkite/test_areas/disaggregated.yaml
+++ b/.buildkite/test_areas/disaggregated.yaml
@@ -16,9 +16,9 @@ steps:
     - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
   mirror:
     amd:
-      label: ":amd: (MI300) Distributed NixlConnector PD accuracy"
+      label: ":amd: (MI355) Distributed NixlConnector PD accuracy"
       dind: false
-      device: mi300_4
+      device: mi355_4
       timeout_in_minutes: 60
       depends_on:
         - image-build-amd
@@ -42,6 +42,29 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  mirror:
+    amd:
+      label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy"
+      dind: false
+      device: mi355_4
+      num_devices: 4
+      timeout_in_minutes: 60
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+        - image-build-amd
+      source_file_dependencies:
+        - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+        - tests/v1/kv_connector/nixl_integration/
+        - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
+        - tests/utils.py
+        - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+        - vllm/platforms/rocm.py
+        - requirements/common.txt
+        - requirements/kv_connectors_rocm.txt
+      commands:
+        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+        # Preserve all six model/TP pairs with native AITER attention.
+        - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py
 
 - label: ":nvidia: (L4) Push NixlConnector PP prefill PD accuracy"
   key: push-nixlconnector-pp-prefill-pd-accuracy-4-gpus
@@ -71,9 +94,9 @@ steps:
     - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
   mirror:
     amd:
-      label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy"
+      label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy"
       dind: false
-      device: mi300_4
+      device: mi355_4
       timeout_in_minutes: 40
       depends_on:
         - image-build-amd
@@ -130,7 +153,7 @@ steps:
       label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy"
       dind: false
       device: mi300_4
-      timeout_in_minutes: 55
+      timeout_in_minutes: 60
       depends_on:
         - image-build-amd
       source_file_dependencies:
@@ -172,6 +195,23 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
+  mirror:
+    amd:
+      label: ":amd: (MI355) Hybrid SSM NixlConnector PD prefix cache"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 25
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+      - vllm/v1/core/sched/
+      - vllm/v1/core/kv_cache_coordinator.py
+      - tests/v1/kv_connector/nixl_integration/
+      - vllm/platforms/rocm.py
+      commands:
+      - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+      - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
 
 - label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD accuracy"
   key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
@@ -188,6 +228,25 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
+  mirror:
+    amd:
+      label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD accuracy"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 40
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+      - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+      - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+      - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
+      - tests/v1/kv_connector/nixl_integration/
+      - requirements/kv_connectors_rocm.txt
+      - vllm/platforms/rocm.py
+      commands:
+      - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+      - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
 
 - label: ":nvidia: (A100) NixlConnector PD + Spec Decode acceptance"
   key: nixlconnector-pd-spec-decode-acceptance-2-gpus
@@ -204,9 +263,9 @@ steps:
     - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
   mirror:
     amd:
-      label: ":amd: (MI300) NixlConnector PD + Spec Decode acceptance"
+      label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance"
       dind: false
-      device: mi300_2
+      device: mi355_2
       timeout_in_minutes: 45
       depends_on:
         - image-build-amd
@@ -234,6 +293,24 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+  mirror:
+    amd:
+      label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD edge cases"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 25
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+      - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+      - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+      - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
+      - tests/v1/kv_connector/nixl_integration/
+      - vllm/platforms/rocm.py
+      commands:
+      - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+      - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
 
 # P TP 4 - D DPEP 4 test case for DSv4-Flash
 - label: ":nvidia: (H200) DSv4-Flash Disaggregated DP EP"
@@ -252,6 +329,7 @@ steps:
     PREFILL_BLOCK_SIZE: "256"
     DECODE_BLOCK_SIZE: "256"
     MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
+    VLLM_ENGINE_READY_TIMEOUT_S: "1800"
     VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
     # The V4 indexer only supports block-outermost layouts; let it pick.
     PREFILLER_KV_LAYOUT: ""
@@ -275,7 +353,31 @@ steps:
     PREFILLER_TP_SIZE: "2"
     DECODER_TP_SIZE: "2"
     MODEL_NAMES: "moonshotai/Kimi-Linear-48B-A3B-Instruct"
+    # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima.
+    PREFILL_BLOCK_SIZE: "2048"
+    DECODE_BLOCK_SIZE: "2048"
     VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768"
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
+  mirror:
+    amd:
+      label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP"
+      dind: false
+      device: mi355_4
+      num_devices: 4
+      timeout_in_minutes: 105
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
+      - vllm/models/kimi_k3/
+      - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+      - tests/v1/kv_connector/nixl_integration/
+      - requirements/kv_connectors_rocm.txt
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+      commands:
+      - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+      - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
diff --git a/.buildkite/test_areas/disaggregated_ec.yaml b/.buildkite/test_areas/disaggregated_ec.yaml
index 9100ad0a431e..67bfebaf8cac 100644
--- a/.buildkite/test_areas/disaggregated_ec.yaml
+++ b/.buildkite/test_areas/disaggregated_ec.yaml
@@ -14,3 +14,22 @@ steps:
     - tests/entrypoints/scale_out/ec_integration/
   commands:
     - bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
+  mirror:
+    amd:
+      label: ":amd: (MI300) Scale-out EC Connector E2E"
+      dind: false
+      device: mi300_2
+      num_devices: 2
+      timeout_in_minutes: 60
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/ec_transfer/
+      - vllm/entrypoints/scale_out/
+      - tests/entrypoints/scale_out/ec_integration/
+      - vllm/platforms/rocm.py
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+      commands:
+      - bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
diff --git a/.buildkite/test_areas/disaggregated_mooncake.yaml b/.buildkite/test_areas/disaggregated_mooncake.yaml
index 8c061a95ad8f..030351ba5763 100644
--- a/.buildkite/test_areas/disaggregated_mooncake.yaml
+++ b/.buildkite/test_areas/disaggregated_mooncake.yaml
@@ -19,6 +19,30 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
+  mirror:
+    amd:
+      label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
+      dind: false
+      device: mi355_4
+      num_devices: 4
+      timeout_in_minutes: 30
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+        - image-build-amd
+      source_file_dependencies:
+        - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
+        - tests/v1/kv_connector/mooncake_integration/
+        - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
+        - tests/utils.py
+        - examples/disaggregated/mooncake_connector/
+        - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+        - vllm/platforms/rocm.py
+        - requirements/common.txt
+        - requirements/kv_connectors_rocm.txt
+      commands:
+        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+        # Same full accuracy/TP sweep, using same-host HIP GPU transport.
+        - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
 
 - label: ":nvidia: (L4) Mooncake EC TCP E2E"
   key: mooncake-ec-tcp-e2e-2-gpus
@@ -40,7 +64,14 @@ steps:
     - vllm/v1/worker/gpu_worker.py
     - vllm/v1/worker/gpu/ec_connector.py
     - vllm/v1/worker/gpu/model_runner.py
+    - vllm/v1/worker/mm_encoder_model_runner.py
+    - vllm/v1/worker/encoder_cudagraph.py
+    - vllm/v1/worker/encoder_cudagraph_defs.py
+    - vllm/v1/attention/ops/vit_attn_wrappers.py
+    - vllm/v1/worker/gpu/model_states/
     - vllm/v1/worker/gpu/mm/
+    - tests/v1/cudagraph/test_encoder_cudagraph.py
+    - tests/models/multimodal/generation/test_vit_cudagraph.py
     - examples/disaggregated/disaggregated_encoder/
     - tests/v1/ec_connector/
     - requirements/kv_connectors.txt
@@ -49,6 +80,9 @@ steps:
     PYTHONUNBUFFERED: "1"
     VLLM_HOST_IP: "127.0.0.1"
     VLLM_USE_V2_MODEL_RUNNER: "1"
+    E_CUDAGRAPH_MM_ENCODER: "1"
+    # FA padding NaNs are tracked separately in #57136.
+    E_MM_ENCODER_ATTN_BACKEND: "TRITON_ATTN"
     MOONCAKE_EC_PROTOCOL: "tcp"
     USE_MM_PROMPTS: "1"
     SKIP_BASELINE: "0"
@@ -73,3 +107,45 @@ steps:
       PY
       )"
     - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
+  mirror:
+    amd:
+      label: ":amd: (MI300) Mooncake EC TCP E2E"
+      dind: false
+      device: mi300_2
+      num_devices: 2
+      timeout_in_minutes: 60
+      working_dir: /vllm-workspace
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/ec_transfer/
+      - vllm/config/ec_transfer.py
+      - vllm/config/multimodal.py
+      - vllm/config/vllm.py
+      - vllm/multimodal/
+      - vllm/v1/core/encoder_cache_manager.py
+      - vllm/v1/core/sched/
+      - vllm/v1/engine/
+      - vllm/v1/worker/ec_connector_model_runner_mixin.py
+      - vllm/v1/worker/gpu_model_runner.py
+      - vllm/v1/worker/gpu_worker.py
+      - vllm/v1/worker/gpu/ec_connector.py
+      - vllm/v1/worker/gpu/model_runner.py
+      - vllm/v1/worker/mm_encoder_model_runner.py
+      - vllm/v1/worker/encoder_cudagraph.py
+      - vllm/v1/worker/encoder_cudagraph_defs.py
+      - vllm/v1/attention/ops/vit_attn_wrappers.py
+      - vllm/v1/worker/gpu/model_states/
+      - vllm/v1/worker/gpu/mm/
+      - tests/v1/cudagraph/test_encoder_cudagraph.py
+      - tests/models/multimodal/generation/test_vit_cudagraph.py
+      - examples/disaggregated/disaggregated_encoder/
+      - tests/v1/ec_connector/
+      - requirements/kv_connectors.txt
+      - vllm/platforms/rocm.py
+      - requirements/rocm.txt
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+        PYTHON_BIN: "python"
+      commands:
+      - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml
index 44e96ce63264..15966f8bd1fc 100644
--- a/.buildkite/test_areas/distributed.yaml
+++ b/.buildkite/test_areas/distributed.yaml
@@ -33,6 +33,29 @@ steps:
   - pytest -v -s distributed/test_sharded_rdt_plan.py
   - pytest -v -s distributed/test_sharded_rdt_producer.py
   - pytest -v -s distributed/test_sharded_rdt_trainer.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Sharded RDT Weight Transfer"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 45
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/weight_transfer/
+      - tests/distributed/test_sharded_rdt_plan.py
+      - tests/distributed/test_sharded_rdt_producer.py
+      - tests/distributed/test_sharded_rdt_trainer.py
+      - vllm/platforms/rocm.py
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+      commands:
+      - uv pip install --system --no-deps 'ray==2.56.1'
+      - pytest -v -s distributed/test_sharded_rdt_plan.py
+      - pytest -v -s distributed/test_sharded_rdt_producer.py
+      - pytest -v -s distributed/test_sharded_rdt_trainer.py
 
 - label: ":nvidia: (L4) Distributed DP Basic"
   key: distributed-dp-tests-2-gpus
@@ -65,9 +88,9 @@ steps:
   - pytest -v -s v1/e2e/spec_decode/test_sharded_sampling.py
   mirror:
     amd:
-      label: ":amd: (MI300) Distributed DP Basic"
+      label: ":amd: (MI355) Distributed DP Basic"
       dind: false
-      device: mi300_2
+      device: mi355_2
       timeout_in_minutes: 45
       depends_on:
       - image-build-amd
@@ -105,6 +128,29 @@ steps:
   - export NCCL_CUMEM_HOST_ENABLE=0
   - pytest -v -s entrypoints/llm/test_collective_rpc.py
   - pytest -v -s ./compile/test_wrapper.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Distributed Compile + RPC"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 65
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/compilation/
+      - vllm/distributed/
+      - "!vllm/distributed/kv_transfer/"
+      - vllm/engine/
+      - vllm/executor/
+      - vllm/worker/worker_base.py
+      - vllm/v1/engine/
+      - vllm/v1/worker/
+      - tests/compile/test_wrapper.py
+      - tests/entrypoints/llm/test_collective_rpc.py
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -v -s entrypoints/llm/test_collective_rpc.py
+      - pytest -v -s ./compile/test_wrapper.py
 
 - label: ":nvidia: (L4) Distributed Torchrun + Shutdown"
   key: distributed-torchrun-shutdown-tests-2-gpus
@@ -173,6 +219,41 @@ steps:
   # default does not fit an L4.
   - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
   - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Distributed Torchrun + Examples"
+      dind: false
+      device: mi355_4
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/
+      - "!vllm/distributed/kv_transfer/"
+      - tests/distributed/test_torchrun_example.py
+      - tests/distributed/test_torchrun_example_moe.py
+      - examples/rl/
+      - examples/features/data_parallel/data_parallel_offline.py
+      - .buildkite/scripts/install-kv-connectors.sh
+      - requirements/kv_connectors_rocm.txt
+      - requirements/test/rocm.in
+      - requirements/test/rocm.txt
+      - vllm/platforms/rocm.py
+      commands:
+      # Sharded RDT needs Ray >= 2.56; keep the ROCm image's dependencies.
+      - uv pip install --system --no-deps 'ray==2.56.1'
+      - export NCCL_CUMEM_HOST_ENABLE=0
+      - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
+      - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
+      - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+      - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+      - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+      - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+      - python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
+      - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py
+      - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py
+      - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+      - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py
 
 - label: ":nvidia: (L4) Distributed DP Extended"
   key: distributed-dp-tests-4-gpus
@@ -186,6 +267,9 @@ steps:
   - tests/v1/distributed
   - tests/v1/engine/test_engine_core_client.py
   - tests/distributed/test_utils
+  - tests/distributed/test_engram_dp_shard.py
+  - vllm/config/engram.py
+  - vllm/models/deepseek_v41/
   commands:
   # https://github.com/NVIDIA/nccl/issues/1838
   - export NCCL_CUMEM_HOST_ENABLE=0
@@ -197,6 +281,7 @@ steps:
   - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
   - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
   - pytest -v -s distributed/test_utils.py
+  - pytest -v -s distributed/test_engram_dp_shard.py
 
 - label: ":nvidia: (L4) Distributed Compile + Comm"
   key: distributed-compile-comm-4-gpus
@@ -241,6 +326,23 @@ steps:
   - export NCCL_CUMEM_HOST_ENABLE=0
   # test with torchrun tp=2 and dp=4 with ep
   - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
+  mirror:
+    amd:
+      label: ":amd: (MI355) Distributed DP + EP"
+      dind: false
+      device: mi355_8
+      timeout_in_minutes: 50
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - examples/features/torchrun/torchrun_dp_example_offline.py
+      - vllm/config/parallel.py
+      - vllm/distributed/
+      - "!vllm/distributed/kv_transfer/"
+      - vllm/v1/engine/llm_engine.py
+      - vllm/v1/executor/uniproc_executor.py
+      - vllm/v1/worker/gpu_worker.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (A100) Distributed"
   key: distributed-tests-4xa100
@@ -272,6 +374,57 @@ steps:
     - pytest -v -s tests/v1/distributed/test_dbo.py
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
     - pytest -v -s tests/distributed/test_packed_tensor.py
+  mirror:
+    amd:
+      label: ':amd: (MI355) Distributed Features'
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 90
+      working_dir: /vllm-workspace/
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/
+      - vllm/v1/distributed/
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - tests/v1/distributed/test_dbo.py
+      - tests/distributed/test_context_parallel.py
+      - examples/features/data_parallel/data_parallel_offline.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      - csrc/custom_quickreduce.cu
+      - csrc/ops.h
+      - csrc/torch_bindings.cpp
+      - vllm/model_executor/layers/
+      - vllm/entrypoints/llm.py
+      - vllm/config/parallel.py
+      - vllm/v1/engine/
+      - vllm/v1/executor/
+      - vllm/v1/worker/
+      - vllm/_custom_ops.py
+      - vllm/envs.py
+      - tests/distributed/test_rocm_aiter_custom_ar.py
+      - tests/distributed/test_rocm_quick_reduce.py
+      - tests/distributed/test_quick_all_reduce.py
+      - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
+      - tests/utils.py
+      - examples/rl/rlhf_async_new_apis.py
+      - tests/distributed/test_weight_transfer.py
+      - tests/distributed/test_packed_tensor.py
+      commands:
+      - pytest -v -s tests/distributed/test_context_parallel.py
+      - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
+      - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
+      - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
+      - pytest -v -s tests/v1/distributed/test_dbo.py
+      - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
+      - pytest -v -s tests/distributed/test_packed_tensor.py
+      - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
+      - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
+      - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
+      - pytest -v -s tests/distributed/test_quick_all_reduce.py
 
 - label: ":nvidia: (B200) Distributed"
   key: distributed-tests-2xb200
@@ -286,17 +439,18 @@ steps:
     - pytest -v -s tests/v1/distributed/test_dbo.py
     - pytest -v -s tests/distributed/test_mnnvl_alltoall.py
 
-- label: ":nvidia: (B200) Kimi-K3 GEMM-RS/AR"
-  key: kimi-k3-gemm-rs-ar-2xb200
+- label: ":nvidia: (B200) GEMM-RS/AR"
+  key: gemm-rs-ar-2xb200
   timeout_in_minutes: 20
   device: b200-k8s
   working_dir: "/vllm-workspace/"
   num_devices: 2
   source_file_dependencies:
-    - vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py
-    - tests/kernels/test_kimi_k3_gemm_rs_ar.py
+    - vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
+    - vllm/cute_utils/
+    - tests/kernels/test_gemm_rs_ar.py
   commands:
-    - pytest -v -s tests/kernels/test_kimi_k3_gemm_rs_ar.py
+    - pytest -v -s tests/kernels/test_gemm_rs_ar.py
 
 - label: ":nvidia: (L4) Distributed 2-Node"
   key: 2-node-test-4-gpus
@@ -335,6 +489,26 @@ steps:
   - pytest -v -s distributed/test_pp_cudagraph.py
   - pytest -v -s distributed/test_pipeline_parallel.py
   - pytest -v -s v1/distributed/test_pp_dp_v2.py
+  mirror:
+    amd:
+      label: ":amd: (MI250) Pipeline + Context Parallelism"
+      dind: false
+      device: mi250_4
+      timeout_in_minutes: 90
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/
+      - "!vllm/distributed/kv_transfer/"
+      - vllm/engine/
+      - vllm/executor/
+      - vllm/model_executor/models/
+      - vllm/model_executor/layers/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - tests/distributed/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (L4) RayExecutorV2"
   key: rayexecutorv2-4-gpus
diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml
index e88f19ad9793..4c80f26ca52e 100644
--- a/.buildkite/test_areas/docker.yaml
+++ b/.buildkite/test_areas/docker.yaml
@@ -3,6 +3,7 @@ depends_on:
   - image-build-cpu
 steps:
 - label: ":computer: (CPU) Docker Build Metadata"
+  key: docker-build-metadata
   timeout_in_minutes: 20
   device: cpu-small
   source_file_dependencies:
@@ -17,6 +18,7 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) Docker Build Metadata"
+      dind: false
       device: mi250_1
       no_gpu: true
       timeout_in_minutes: 35
diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml
index 02ba409e3433..12ebee893011 100644
--- a/.buildkite/test_areas/e2e_integration.yaml
+++ b/.buildkite/test_areas/e2e_integration.yaml
@@ -11,6 +11,28 @@ steps:
   working_dir: "/vllm-workspace"
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
+  mirror:
+    amd:
+      label: ":amd: (MI355) DeepSeek V2-Lite Sync EPLB Accuracy"
+      dind: false
+      device: mi355_4
+      timeout_in_minutes: 40
+      env:
+        VLLM_ENGINE_READY_TIMEOUT_S: "1800"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/models/
+      - vllm/model_executor/model_loader/
+      - vllm/distributed/eplb
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/model_executor/layers/quantization/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/backends/mla/
+      - vllm/v1/attention/selector.py
+      - .buildkite/scripts/scheduled_integration_test/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H100) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy"
   key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-4xh100
@@ -21,6 +43,25 @@ steps:
   working_dir: "/vllm-workspace"
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
+  mirror:
+    amd:
+      label: ":amd: (MI355) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy TP2"
+      dind: false
+      device: mi355_4
+      timeout_in_minutes: 30
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/models/
+      - vllm/model_executor/model_loader/
+      - vllm/model_executor/layers/quantization/
+      - vllm/distributed/eplb
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - .buildkite/scripts/scheduled_integration_test/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (B200) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy"
   key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-2xb200
@@ -31,6 +72,25 @@ steps:
   working_dir: "/vllm-workspace"
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
+  mirror:
+    amd:
+      label: ":amd: (MI355) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy TP1"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 25
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/models/
+      - vllm/model_executor/model_loader/
+      - vllm/model_executor/layers/quantization/
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/distributed/eplb
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - .buildkite/scripts/scheduled_integration_test/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H100) Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy"
   key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy
@@ -41,6 +101,25 @@ steps:
   working_dir: "/vllm-workspace"
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
+  mirror:
+    amd:
+      label: ":amd: (MI355) Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy"
+      dind: false
+      device: mi355_4
+      timeout_in_minutes: 30
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/models/
+      - vllm/model_executor/model_loader/
+      - vllm/model_executor/layers/quantization/
+      - vllm/distributed/eplb
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - .buildkite/scripts/scheduled_integration_test/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H100) DeepSeek V2-Lite Prefetch Offload Accuracy"
   key: deepseek-v2-lite-prefetch-offload-accuracy-h100
@@ -57,7 +136,7 @@ steps:
       dind: false
       device: mi355_1
       working_dir: "/vllm-workspace"
-      timeout_in_minutes: 20
+      timeout_in_minutes: 30
       depends_on:
       - image-build-amd
       source_file_dependencies:
diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml
index 4384484c0954..e38cd925732d 100644
--- a/.buildkite/test_areas/engine.yaml
+++ b/.buildkite/test_areas/engine.yaml
@@ -52,10 +52,13 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) V1 Engine"
+      dind: false
       device: mi250_1
       timeout_in_minutes: 45
       depends_on:
       - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 18GB) E2E Scheduling"
   key: e2e-scheduling-1-gpu
@@ -69,23 +72,29 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) E2E Scheduling"
+      dind: false
       device: mi250_1
       timeout_in_minutes: 55
       depends_on:
       - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
 
-- label: ":nvidia: (H200 MIG 35GB) E2E Core"
-  device: h200_35gb
+- label: ":nvidia: (H200 MIG 18GB) E2E Core"
+  device: h200_18gb
   key: e2e-core-1-gpu
-  timeout_in_minutes: 40
+  timeout_in_minutes: 65
   source_file_dependencies:
     - vllm/v1/
     - tests/v1/e2e/general/
   commands:
-    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
+    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
   mirror:
     amd:
       label: ":amd: (MI250) E2E Core"
+      dind: false
+      commands:
+      - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
       device: mi250_1
       timeout_in_minutes: 50
       depends_on:
@@ -95,6 +104,35 @@ steps:
       - tests/v1/e2e/general/
       - vllm/platforms/rocm.py
 
+# Keep the large Gemma-3n and Mamba MRV2 warmup allocations on 35GB slices.
+- label: ":nvidia: (H200 MIG 35GB) E2E Core Large Memory"
+  device: h200_35gb
+  key: e2e-core-large-memory
+  timeout_in_minutes: 50
+  source_file_dependencies:
+    - vllm/v1/
+    - tests/v1/e2e/general/
+  commands:
+    - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
+    - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async
+  mirror:
+    amd:
+      label: ":amd: (MI355) E2E Core Large Memory"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 65
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/
+      - tests/v1/e2e/general/
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
+      - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async
+
 - label: ":nvidia: (L4) V1 E2E"
   key: v1-e2e-2-gpus
   timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability
diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml
index 6ad579f82ad3..075b3733caf8 100644
--- a/.buildkite/test_areas/entrypoints.yaml
+++ b/.buildkite/test_areas/entrypoints.yaml
@@ -2,6 +2,27 @@ group: Entrypoints
 depends_on: 
   - image-build
 steps:
+- label: ":nvidia: (H200) Initialized Snapshot E2E"
+  key: initialized-snapshot-e2e
+  device: h200
+  num_devices: 1
+  no_plugin: true
+  optional: true
+  timeout_in_minutes: 60
+  working_dir: "."
+  env:
+    SNAPSHOT_E2E_RUN_LIMIT_S: "1200"
+    NVIDIA_VISIBLE_DEVICES: "0"
+  source_file_dependencies:
+  - .buildkite/scripts/initialized-snapshot-e2e.sh
+  - .buildkite/test_areas/entrypoints.yaml
+  - docker/Dockerfile
+  - tools/install_snapshot_runtime.sh
+  - vllm/entrypoints/cli/
+  - vllm/snapshot/
+  commands:
+  - bash .buildkite/scripts/initialized-snapshot-e2e.sh "$IMAGE_TAG" "$BUILDKITE_COMMIT"
+
 - label: ":nvidia: (H200 MIG 35GB) Entrypoints Unit"
   device: h200_35gb
   key: entrypoints-unit-tests
@@ -21,9 +42,11 @@ steps:
       label: ":amd: (MI355) Entrypoints Unit"
       dind: false
       device: mi355_1
-      timeout_in_minutes: 35
+      timeout_in_minutes: 50
       depends_on:
       - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (LLM)"
   device: h200_35gb
@@ -44,7 +67,7 @@ steps:
       label: ":amd: (MI300) Entrypoints Integration (LLM)"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 55
+      timeout_in_minutes: 90
       depends_on:
       - image-build-amd
 
@@ -73,9 +96,9 @@ steps:
       depends_on:
       - image-build-amd
 
-- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server OpenAI - Part 1)"
+- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (OpenAI API completion)"
   device: h200_35gb
-  key: entrypoints-integration-api-server-openai-part-1
+  key: entrypoints-integration-api-server-openai-completion
   timeout_in_minutes: 68
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -84,19 +107,20 @@ steps:
   - tests/entrypoints/openai
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
   - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
   mirror:
     amd:
-      label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)"
+      label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)"
       dind: false
       device: mi300_1
       timeout_in_minutes: 65
       depends_on:
       - image-build-amd
 
-- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server OpenAI - Part 2)"
+- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (OpenAI API chat_completion)"
   device: h200_35gb
-  key: entrypoints-integration-api-server-openai-part-2
+  key: entrypoints-integration-api-server-openai-chat_completion
   timeout_in_minutes: 83
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -106,10 +130,9 @@ steps:
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/openai/chat_completion
-  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
   mirror:
     amd:
-      label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)"
+      label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)"
       dind: false
       device: mi300_1
       timeout_in_minutes: 70
@@ -176,6 +199,15 @@ steps:
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/speech_to_text
+  mirror:
+    amd:
+      label: ":amd: (MI355) Entrypoints Integration (Speech to Text)"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 60
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
 
 - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (Multimodal)"
   device: h200_35gb
@@ -210,6 +242,15 @@ steps:
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/pooling
+  mirror:
+    amd:
+      label: ":amd: (MI355) Entrypoints Integration (Pooling)"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 65
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
 
 - label: ":nvidia: (H200 MIG 18GB) OpenAI API Correctness"
   key: openai-api-correctness
diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml
index f74b1adf2e2a..b9ce4f59240c 100644
--- a/.buildkite/test_areas/expert_parallelism.yaml
+++ b/.buildkite/test_areas/expert_parallelism.yaml
@@ -40,10 +40,23 @@ steps:
   commands:
   - pytest -v -s distributed/test_eplb_execute.py
   - pytest -v -s distributed/test_eplb_spec_decode.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) EPLB Execution"
+      dind: false
+      device: mi355_4
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/distributed/eplb
+      - tests/distributed/test_eplb_execute.py
+      - tests/distributed/test_eplb_spec_decode.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H100) Elastic EP Scaling"
   key: elastic-ep-scaling-test
-  timeout_in_minutes: 30
+  timeout_in_minutes: 40
   device: h100
   working_dir: "/vllm-workspace/tests"
   num_devices: 4
diff --git a/.buildkite/test_areas/fault_tolerance.yaml b/.buildkite/test_areas/fault_tolerance.yaml
index 2e85a1c3f856..a578be14eded 100644
--- a/.buildkite/test_areas/fault_tolerance.yaml
+++ b/.buildkite/test_areas/fault_tolerance.yaml
@@ -24,3 +24,33 @@ steps:
   # https://github.com/NVIDIA/nccl/issues/1838
   - export NCCL_CUMEM_HOST_ENABLE=0
   - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) MoRI EP Healthy Baseline"
+      dind: false
+      device: mi355_2
+      num_devices: 2
+      timeout_in_minutes: 10
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/moe/
+      - csrc/rocm/
+      - tests/kernels/moe/
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/model_executor/layers/quantization/
+      - vllm/distributed/
+      - vllm/config/
+      - vllm/forward_context.py
+      - vllm/v1/worker/workspace.py
+      - vllm/utils/import_utils.py
+      - vllm/utils/math_utils.py
+      - vllm/utils/torch_utils.py
+      - vllm/platforms/
+      - vllm/_aiter_ops.py
+      # TODO: Add fault injection, detection and recovery once MoRI supports them.
+      # This mirror currently checks healthy EP execution only.
+      commands:
+      - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
+      - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]'
diff --git a/.buildkite/test_areas/jit_monitor.yaml b/.buildkite/test_areas/jit_monitor.yaml
index 3355470c1712..27de713fa2f0 100644
--- a/.buildkite/test_areas/jit_monitor.yaml
+++ b/.buildkite/test_areas/jit_monitor.yaml
@@ -20,3 +20,24 @@ steps:
     # instead of running until the build timeout.
     - export PYTHONFAULTHANDLER=1
     - pytest -v -s jit_monitor/test_no_runtime_jit.py --timeout=900 --timeout-method=thread
+  mirror:
+    amd:
+      # Focused MRV2 regression; the broader model/backend matrix is pending.
+      label: ':amd: (MI355) MRV2 Sampler JIT Warmup'
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 10
+      working_dir: /vllm-workspace/
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - tests/jit_monitor/
+      - tests/models/utils.py
+      - tests/utils.py
+      - vllm/v1/worker/gpu/
+      - vllm/v1/sample/ops/topk_topp_sampler.py
+      - vllm/platforms/rocm.py
+      - vllm/envs.py
+      commands:
+      - pytest -v -s tests/jit_monitor/test_no_runtime_jit_rocm.py --timeout=300
diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml
index 2b43938060e9..ba683d58adf3 100644
--- a/.buildkite/test_areas/kernels.yaml
+++ b/.buildkite/test_areas/kernels.yaml
@@ -13,6 +13,22 @@ steps:
   commands:
     - pytest -v -s tests/ir
     - pytest -v -s tests/kernels/ir
+  mirror:
+    amd:
+      label: ":amd: (MI355) vLLM IR"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 50
+      working_dir: "/vllm-workspace/"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - tests/ir/
+      - tests/kernels/ir/
+      - vllm/ir
+      - vllm/kernels
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 35GB) Core Operation Kernels Shard %N"
   device: h200_35gb
@@ -26,6 +42,29 @@ steps:
   commands:
     - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 3
+  mirror:
+    amd:
+      label: ":amd: (MI355) Core Operation Kernels Shard %N"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 40
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/
+      - tests/kernels/core
+      - tests/kernels/test_top_k_per_row.py
+      - tests/kernels/test_concat_mla_q.py
+      - tests/kernels/test_rocm_fp32_router_gemm.py
+      - vllm/model_executor/layers/fused_moe/router/gate_linear.py
+      - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
+      - vllm/model_executor/layers/rotary_embedding/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      - tests/kernels/test_fused_qk_norm_rope_gate.py
+      commands:
+      - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_fused_qk_norm_rope_gate.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":nvidia: (H100) MiniMax Reduce RMS Kernels"
   key: kernels-minimax-reduce-rms-test-2-gpus
@@ -72,6 +111,23 @@ steps:
   commands:
     - pytest -v -s kernels/test_fused_deepseek_v4_*.py
     - pytest -v -s kernels/test_top_k_per_row.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) DeepSeek V4 Fused Kernels"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 60
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+      - vllm/models/deepseek_v4/common/ops/
+      - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
 
 - label: ":nvidia: (B200) DeepSeek V4 Kernels"
   key: deepseek-v4-kernel-test-b200
@@ -81,6 +137,7 @@ steps:
   - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
   - vllm/models/deepseek_v4/common/ops/
   - vllm/models/deepseek_v4/nvidia/
+  - vllm/utils/deep_gemm.py
   - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
   - tests/models/test_deepseek_v4_mega_moe.py
   commands:
@@ -102,6 +159,8 @@ steps:
   - "!vllm/distributed/kv_transfer/"
   - tests/kernels/
   commands:
+    # Isolate mHC from FlashInfer's import-time ptxas override.
+    - pytest -v -s kernels/test_mhc_kernels.py
     - pytest -v -s kernels/
         --ignore=kernels/attention
         --ignore=kernels/core
@@ -117,6 +176,7 @@ steps:
         --ignore=kernels/test_kda.py
         --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py
         --ignore=kernels/test_ll_bf16_gemm.py
+        --ignore=kernels/test_mhc_kernels.py
         --ignore=kernels/test_shuffle_rows.py
     # BROKEN on main, pending kernel fixes (B200):
     #   test_shuffle_rows.py (1: test_shuffle_rows_edge_cases)
@@ -142,6 +202,7 @@ steps:
       dind: false
       device: mi300_1
       timeout_in_minutes: 90
+      parallelism: 6
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -152,6 +213,8 @@ steps:
       - vllm/_aiter_ops.py
       - vllm/envs.py
       - vllm/platforms/rocm.py
+      commands:
+      - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":nvidia: (H100) Attention DiffKV Kernels"
   key: kernels-attention-diffkv-test-h100
@@ -165,6 +228,24 @@ steps:
   - tests/kernels/attention/test_triton_unified_attention_diffkv.py
   commands:
     - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Attention DiffKV Kernels"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 35
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/attention/ops/triton_unified_attention_diffkv.py
+      - vllm/v1/attention/backends/triton_attn_diffkv.py
+      - tests/kernels/attention/test_triton_unified_attention.py
+      - tests/kernels/attention/test_triton_unified_attention_diffkv.py
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
 
 - label: ":nvidia: (H100) MLA Kernel Test"
   key: kernels-flashmla-test-h100
@@ -187,6 +268,31 @@ steps:
     - pytest -v -s kernels/attention/test_flashmla_sparse.py
     - pytest -v -s kernels/attention/test_flashinfer_mla_decode.py
     - pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) MLA Kernels"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 40
+      working_dir: /vllm-workspace/
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+      - vllm/v1/attention/selector.py
+      - vllm/platforms/rocm.py
+      - vllm/_aiter_ops.py
+      - tests/kernels/attention/test_rocm_aiter_mla_decode.py
+      - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
+      - tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
+      - tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
+      commands:
+      - amd-smi
+      - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
+      - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
+      - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
+      - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
 
 - label: ":nvidia: (L4) Quantization Kernels Shard %N"
   key: kernels-quantization-test
@@ -205,9 +311,6 @@ steps:
       label: ":amd: (MI300) Quantization Kernels Shard %N"
       dind: false
       device: mi300_1
-      # Pin the AMD mirror at its current 2 shards so raising NVIDIA
-      # parallelism does not multiply AMD copies (needs ci-infra #473).
-      parallelism: 2
       timeout_in_minutes: 120
       source_file_dependencies:
       - csrc/quantization/
@@ -219,8 +322,16 @@ steps:
       - vllm/kernels/aiter_ops.py
       - vllm/platforms/rocm.py
       - vllm/model_executor/kernels/
+      - csrc/rocm/
+      - vllm/v1/attention/backends/rocm_aiter_fa.py
+      - tests/kernels/quant_utils.py
+      - tests/kernels/utils.py
+      - vllm/_custom_ops.py
+      - vllm/envs.py
       depends_on:
       - image-build-amd
+      commands:
+      - pytest -v -s kernels/quantization -k "not test_nvfp4_dense_emulation" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":nvidia: (L4) MoE Kernels Shard %N"
   key: kernels-moe-test
@@ -255,6 +366,10 @@ steps:
       - vllm/config
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
+      - csrc/rocm/
+      - tests/kernels/utils.py
+      - vllm/model_executor/layers/quantization/mxfp4.py
+      - vllm/model_executor/layers/quantization/utils/quant_utils.py
       depends_on:
       - image-build-amd
 
@@ -265,16 +380,21 @@ steps:
   source_file_dependencies:
   - csrc/mamba/
   - tests/kernels/mamba
+  - vllm/model_executor/layers/mamba/gdn
   - vllm/model_executor/layers/mamba/ops
   commands:
     - pytest -v -s kernels/mamba
   mirror:
     amd:
       label: ":amd: (MI250) Mamba Kernels"
+      dind: false
       device: mi250_1
       timeout_in_minutes: 55
       depends_on:
       - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
+      - vllm/third_party/flash_linear_attention/
 
 - label: ":nvidia: (H100) DeepGEMM Kernels"
   key: kernels-deepgemm-test-h100
@@ -314,6 +434,7 @@ steps:
   - csrc/quantization/fp4/
   - csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu
   - csrc/libtorch_stable/quantization/fp4/
+  - csrc/libtorch_stable/quantization/w8a8/cutlass/
   - csrc/attention/mla/
   - csrc/quantization/cutlass_w8a8/moe/
   - vllm/model_executor/layers/fusion/
@@ -355,6 +476,7 @@ steps:
   - tests/kernels/test_ll_bf16_gemm.py
   - tests/kernels/test_top_k_per_row.py
   - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
+  - tests/kernels/quantization/test_block_fp8.py
   - tests/models/quantization/test_nvfp4.py
   commands:
     - nvidia-smi
@@ -389,6 +511,7 @@ steps:
           pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
           pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
           pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
+          pytest -v -s tests/kernels/quantization/test_block_fp8.py::test_w8a8_block_fp8_cutlass_matmul
           pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
           pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
           pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
@@ -402,6 +525,32 @@ steps:
           exit 1
           ;;
       esac
+  mirror:
+    amd:
+      label: ":amd: (MI355) NVFP4 Dense Emulation Shard %N"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 10
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - tests/kernels/quantization/test_nvfp4_emulation.py
+      - tests/kernels/quantization/nvfp4_utils.py
+      - tests/conftest.py
+      - vllm/model_executor/kernels/linear/
+      - vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+      - vllm/config/kernel.py
+      - vllm/platforms/
+      - vllm/envs.py
+      - vllm/utils/torch_utils.py
+      - requirements/rocm-test.txt
+      # W4A4 format numerics and replay via emulation. TODO: native packed
+      # activation/KV-cache layouts and CUDA-only kernels need AMD backends.
+      commands:
+      - amd-smi
+      - pytest -v -s kernels/quantization/test_nvfp4_emulation.py -k test_nvfp4_dense_emulation --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly"
   key: b12x-linear-kernels-dgx-spark-nightly
@@ -456,10 +605,14 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) Helion Kernels Shard %N"
+      dind: false
       device: mi250_1
       timeout_in_minutes: 45
       depends_on:
       - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
+      - scripts/benchmark_helion_kernels.py
 
  
 - label: ":nvidia: (H100) FP8 MoE Kernels"
@@ -488,6 +641,25 @@ steps:
   commands:
     - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py
     - pytest -v -s kernels/moe/test_deepep_moe.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) DeepEP FP8 MoE Kernels"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 55
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/moe/
+      - csrc/quantization/w8a8/cutlass/moe/
+      - vllm/model_executor/layers/fused_moe/
+      - tests/kernels/moe/test_deepep_moe.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      - vllm/envs.py
+      commands:
+      - pytest -v -s kernels/moe/test_deepep_moe.py
 
 - label: ":nvidia: (B200) FP4 MoE Kernels"
   key: kernels-fp4-moe-test-b200
@@ -517,7 +689,33 @@ steps:
   - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py
-
+  mirror:
+    amd:
+      label: ":amd: (MI355) FusedMoE Layer Kernels"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 35
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/moe/
+      - csrc/rocm/
+      - tests/kernels/moe
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/model_executor/layers/quantization/
+      - vllm/distributed/
+      - vllm/config/
+      - vllm/forward_context.py
+      - vllm/v1/worker/workspace.py
+      - vllm/utils/import_utils.py
+      - vllm/utils/math_utils.py
+      - vllm/utils/torch_utils.py
+      - vllm/platforms/
+      - vllm/_aiter_ops.py
+      commands:
+      # MoRI graph cases run in the Fault Tolerance group's healthy EP baseline.
+      - pytest -v -s kernels/moe/test_moe_layer.py -k "not test_moe_layer_mori_graph"
 
 - label: ":nvidia: (B200) FusedMoE Layer Kernels"
   key: kernels-fusedmoe-layer-test-2-b200s
@@ -535,3 +733,95 @@ steps:
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py
     - pytest -v -s kernels/moe/test_deepep_v2_moe.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) MoRI EP Numerics"
+      dind: false
+      device: mi355_2
+      num_devices: 2
+      timeout_in_minutes: 20
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - tests/kernels/moe/test_mori_moe.py
+      - tests/kernels/moe/utils.py
+      - tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+      - tests/kernels/utils.py
+      - tests/utils.py
+      - tests/conftest.py
+      - vllm/utils/torch_utils.py
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/model_executor/layers/quantization/
+      - vllm/distributed/
+      - vllm/config/
+      - vllm/forward_context.py
+      - vllm/v1/worker/workspace.py
+      - vllm/platforms/
+      - vllm/_aiter_ops.py
+      - requirements/rocm-test.txt
+      # Same-node MoRI/AITER BF16 and per-token/channel FP8. TODO: the
+      # DeepEP v2 FlashInfer backend matrix has no direct AMD implementation.
+      commands:
+      - amd-smi
+      - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 pytest -v -s kernels/moe/test_mori_moe.py
+
+- label: ":amd: (MI355) FP8 MoE Kernels"
+  key: amd-fp8-moe-kernels-mi355
+  dind: false
+  device: mi355_1
+  num_devices: 1
+  timeout_in_minutes: 40
+  working_dir: /vllm-workspace/tests
+  depends_on:
+  - image-build-amd
+  source_file_dependencies:
+  - csrc/moe/
+  - vllm/model_executor/layers/fused_moe/
+  - tests/kernels/moe/test_deepep_moe.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  - vllm/envs.py
+  - tests/kernels/moe/test_gpt_oss_triton_kernels.py
+  - tests/kernels/moe/test_modular_oai_triton_moe.py
+  - tests/kernels/moe/test_moe.py
+  - tests/kernels/moe/test_block_int8.py
+  - tests/kernels/moe/test_triton_moe_no_act_mul.py
+  - tests/kernels/moe/test_triton_moe_ptpc_fp8.py
+  commands:
+  - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py
+  - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py
+  - pytest -v -s kernels/moe/test_moe.py
+  - pytest -v -s kernels/moe/test_block_int8.py
+  - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py
+  - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
+
+- label: ":amd: (MI355) Kernels"
+  key: amd-kernels-mi355
+  dind: false
+  device: mi355_1
+  num_devices: 1
+  timeout_in_minutes: 25
+  working_dir: /vllm-workspace/
+  depends_on:
+  - image-build-amd
+  source_file_dependencies:
+  - csrc/quantization/fp4/
+  - csrc/attention/mla/
+  - csrc/quantization/cutlass_w8a8/moe/
+  - vllm/model_executor/layers/fused_moe/cutlass_moe.py
+  - vllm/v1/attention/backends/triton_attn.py
+  - vllm/v1/attention/backends/rocm_attn.py
+  - vllm/v1/attention/backends/rocm_aiter_fa.py
+  - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+  - vllm/v1/attention/backends/mla/aiter_triton_mla.py
+  - vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+  - vllm/v1/attention/selector.py
+  - vllm/platforms/rocm.py
+  - vllm/_aiter_ops.py
+  - examples/basic/offline_inference/chat.py
+  - tests/kernels/attention/test_attention_selector.py
+  commands:
+  - amd-smi
+  - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
+  - pytest -v -s tests/kernels/attention/test_attention_selector.py
diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml
index 0863dfd2767e..dccf57f2e2e5 100644
--- a/.buildkite/test_areas/lm_eval.yaml
+++ b/.buildkite/test_areas/lm_eval.yaml
@@ -30,6 +30,65 @@ steps:
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
 
+- label: ":amd: (MI300) LM Eval Small Models Harness"
+  key: amd-lm-eval-small-models-harness
+  dind: false
+  device: mi300_1
+  timeout_in_minutes: 35
+  depends_on:
+  - image-build-amd
+  working_dir: /vllm-workspace/.buildkite/lm-eval-harness
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  - .buildkite/lm-eval-harness/
+  commands:
+  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt
+
+- label: ":nvidia: (H200 MIG 35GB) LM Eval Watermarking"
+  device: h200_35gb
+  key: lm-eval-watermarking
+  # Recorded runs took 80–140 seconds
+  timeout_in_minutes: 10
+  source_file_dependencies:
+  - .buildkite/test_areas/lm_eval.yaml
+  - tests/evals/gsm8k/
+  - vllm/config/watermarking.py
+  - vllm/sampling_params.py
+  - vllm/v1/watermarking/
+  - vllm/v1/worker/gpu/
+  autorun_on_main: true
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) LM Eval Watermarking"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 45
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - .buildkite/test_areas/lm_eval.yaml
+      - tests/evals/gsm8k/
+      - vllm/config/watermarking.py
+      - vllm/sampling_params.py
+      - vllm/v1/watermarking/
+      - vllm/v1/worker/gpu/
+      - vllm/platforms/rocm.py
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+      commands:
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
+
 # - label: LM Eval Large Models (4xA100)
 #   key: lm-eval-large-models-4xa100
 #   device: a100
@@ -55,6 +114,27 @@ steps:
   commands:
     - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
     - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
+  mirror:
+    amd:
+      label: ":amd: (MI355) LM Eval Large Models FP8"
+      dind: false
+      device: mi355_4
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/
+      - vllm/model_executor/layers/quantization
+      - vllm/model_executor/models/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - .buildkite/lm-eval-harness/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - export VLLM_USE_DEEP_GEMM=0
+      - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-tp4.txt --tp-size=4
 
 - label: ":nvidia: (B200) LM Eval Small Models"
   key: lm-eval-small-models-1xb200
@@ -79,6 +159,44 @@ steps:
   autorun_on_main: true
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) LM Eval Small Models Distributed"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/
+      - vllm/model_executor/layers/quantization
+      - vllm/model_executor/models/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
+
+- label: ":nvidia: (B200) LM Eval HiSparse Nightly"
+  key: lm-eval-hisparse-nightly-4xb200
+  timeout_in_minutes: 60
+  device: b200-k8s
+  num_devices: 4
+  optional: true
+  source_file_dependencies:
+  - .buildkite/test_areas/lm_eval.yaml
+  - tests/evals/gsm8k/
+  - vllm/v1/hisparse/
+  - vllm/distributed/kv_transfer/kv_connector/v1/hisparse/
+  - vllm/v1/attention/backends/mla/
+  - vllm/model_executor/layers/attention/sparse_mla_attention.py
+  - vllm/v1/worker/gpu/
+  - csrc/libtorch_stable/hisparse_kernels.cu
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt
 
 - label: ":nvidia: (B200) LM Eval PCP"
   key: lm-eval-pcp-4xb200
@@ -90,18 +208,72 @@ steps:
   - csrc/
   - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
   - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml
   - tests/evals/gsm8k/configs/models-pcp.txt
   - vllm/model_executor/layers/quantization
   - vllm/config/parallel.py
   - vllm/distributed/parallel_state.py
   - vllm/model_executor/layers/attention/mla_attention.py
   - vllm/model_executor/layers/attention/pcp.py
+  - vllm/model_executor/layers/attention/sparse_mla_attention.py
+  - vllm/model_executor/layers/sparse_attn_indexer.py
+  - vllm/v1/attention/backends/mla/flashmla_sparse.py
+  - vllm/v1/attention/backends/mla/indexer.py
   - vllm/v1/worker/gpu/model_runner.py
   - vllm/v1/worker/gpu/pcp_manager.py
   autorun_on_main: true
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
 
+- label: ":nvidia: (H100) LM Eval Watermark Feature Combination"
+  key: lm-eval-dspark-watermark-2xh100
+  device: h100
+  num_devices: 2
+  # Local TP2 runs took 84–379 seconds, including startup.
+  timeout_in_minutes: 15
+  source_file_dependencies:
+  - .buildkite/test_areas/lm_eval.yaml
+  - tests/evals/gsm8k/
+  - vllm/config/speculative.py
+  - vllm/config/vllm.py
+  - vllm/config/watermarking.py
+  - vllm/sampling_params.py
+  - vllm/model_executor/models/qwen3_dspark.py
+  - vllm/model_executor/models/qwen3_dflash.py
+  - vllm/model_executor/warmup/
+  - vllm/v1/watermarking/
+  - vllm/v1/worker/gpu/
+  autorun_on_main: true
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) LM Eval Watermark Feature Combination"
+      dind: false
+      device: mi355_2
+      num_devices: 2
+      timeout_in_minutes: 60
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - .buildkite/test_areas/lm_eval.yaml
+      - tests/evals/gsm8k/
+      - vllm/config/speculative.py
+      - vllm/config/vllm.py
+      - vllm/config/watermarking.py
+      - vllm/sampling_params.py
+      - vllm/model_executor/models/qwen3_dspark.py
+      - vllm/model_executor/models/qwen3_dflash.py
+      - vllm/model_executor/warmup/
+      - vllm/v1/watermarking/
+      - vllm/v1/worker/gpu/
+      - vllm/platforms/rocm.py
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+      commands:
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
+
 - label: ":nvidia: (B200) LM Eval Spec Decode"
   key: lm-eval-spec-decode-4xb200
   timeout_in_minutes: 120
@@ -191,27 +363,94 @@ steps:
   - vllm/third_party/flash_linear_attention/ops/
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
-
-- label: ":nvidia: (H200) LM Eval Large Models"
-  key: lm-eval-large-models-8xh200
-  timeout_in_minutes: 50
-  device: h200
-  optional: true
-  num_devices: 8
-  commands:
-    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
   mirror:
     amd:
-      label: ":amd: (MI300) LM Eval Large Models"
+      label: ":amd: (MI355) LM Eval Qwen3.5 Models"
       dind: false
-      device: mi300_8
+      device: mi355_2
       timeout_in_minutes: 40
       depends_on:
       - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/models/qwen3_5.py
+      - vllm/model_executor/models/qwen3_5_mtp.py
+      - vllm/transformers_utils/configs/qwen3_5.py
+      - vllm/transformers_utils/configs/qwen3_5_moe.py
+      - vllm/model_executor/models/qwen2.py
+      - vllm/model_executor/models/qwen3.py
+      - vllm/model_executor/models/qwen3_next.py
+      - vllm/model_executor/models/qwen3_next_mtp.py
+      - vllm/third_party/flash_linear_attention/ops/
+      - tests/evals/gsm8k/configs/models-qwen35-mi355.txt
+      - tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-DEP2.yaml
+      - tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml
+      - tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
+      - tests/evals/gsm8k/test_gsm8k_correctness.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
       commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
-      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt
+
+- label: ":nvidia: (H200) LM Eval DeepSeek R1 TP"
+  key: lm-eval-deepseek-r1-tp-8xh200
+  timeout_in_minutes: 35
+  device: h200
+  optional: true
+  num_devices: 8
+  commands:
+    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-r1-tp.txt
+
+- label: ":nvidia: (H200) LM Eval DeepSeek R1 DP"
+  key: lm-eval-deepseek-r1-dp-8xh200
+  timeout_in_minutes: 35
+  device: h200
+  optional: true
+  num_devices: 8
+  commands:
+    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-r1-dp.txt
+
+- label: ":nvidia: (H200) LM Eval DeepSeek V3.2 TP"
+  key: lm-eval-deepseek-v32-tp-8xh200
+  timeout_in_minutes: 35
+  device: h200
+  optional: true
+  num_devices: 8
+  commands:
+    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-v32-tp.txt
+
+- label: ":nvidia: (H200) LM Eval DeepSeek V3.2 DP"
+  key: lm-eval-deepseek-v32-dp-8xh200
+  timeout_in_minutes: 35
+  device: h200
+  optional: true
+  num_devices: 8
+  commands:
+    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-deepseek-v32-dp.txt
+
+- label: ":nvidia: (H200) LM Eval Nemotron 3 Super"
+  key: lm-eval-nemotron-3-super-8xh200
+  timeout_in_minutes: 35
+  device: h200
+  optional: true
+  num_devices: 8
+  commands:
+    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-nemotron-3-super.txt
+
+- label: ":amd: (MI300) LM Eval Large Models"
+  key: amd-lm-eval-large-models
+  dind: false
+  device: mi300_8
+  optional: true
+  timeout_in_minutes: 65
+  depends_on:
+  - image-build-amd
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
+  source_file_dependencies:
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H100) MoE Refactor Integration TEMPORARY"
   key: moe-refactor-integration-test-h100-temporary
@@ -224,7 +463,27 @@ steps:
     # build 83851): shards ~12.4/13.6/12.4/14.5m of eval time. New configs must
     # be added to exactly one config-h100-shard-*.txt.
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt
-  
+  mirror:
+    amd:
+      label: ":amd: (MI355) MoE Refactor Integration Shard %N"
+      dind: false
+      device: mi355_2
+      num_devices: 2
+      timeout_in_minutes: 105
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
+      - tests/evals/gsm8k/
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/model_executor/layers/quantization/
+      - vllm/_aiter_ops.py
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+      commands:
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-rocm-shard-$$BUILDKITE_PARALLEL_JOB.txt
+
 - label: ":nvidia: (B200) MoE Refactor Integration TEMPORARY Shard %N"
   key: moe-refactor-integration-test-b200-temporary
   device: b200-k8s
@@ -242,47 +501,67 @@ steps:
   commands:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
-- label: ":nvidia: (A100) LM Eval Humming FP16 TEMPORARY Shard %N"
-  key: lm-eval-humming-f16-a100
+- label: ":nvidia: (L4) LM Eval Humming TEMPORARY"
+  key: lm-eval-humming-l4
+  optional: true
+  timeout_in_minutes: 90
+  device: l4
+  num_devices: 1
+  source_file_dependencies:
+  - .buildkite/test_areas/lm_eval.yaml
+  - tests/evals/gsm8k/configs/humming/
+  - requirements/cuda.txt
+  - vllm/utils/humming.py
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming/
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/
+  - vllm/model_executor/kernels/linear/
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-l4.txt
+
+- label: ":nvidia: (A100) LM Eval Humming Default TEMPORARY Shard %N"
+  key: lm-eval-humming-a100-default
+  optional: true
   timeout_in_minutes: 75
   device: a100
-  optional: true
   num_devices: 1
   parallelism: 3
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/humming.py
-  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/quantization/utils/humming/
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
   - vllm/model_executor/layers/fused_moe/oracle/
   - vllm/model_executor/kernels/linear/
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt
 
-- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY"
-  key: lm-eval-humming-act-a100
+- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY Shard %N"
+  key: lm-eval-humming-a100-act
+  optional: true
   timeout_in_minutes: 45
   device: a100
-  optional: true
   num_devices: 1
+  parallelism: 2
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/humming.py
-  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/quantization/utils/humming/
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
   - vllm/model_executor/layers/fused_moe/oracle/
   - vllm/model_executor/kernels/linear/
   commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-act-int8-shard-$$BUILDKITE_PARALLEL_JOB.txt
 
-- label: ":nvidia: (H100) LM Eval Humming FP16 TEMPORARY Shard %N"
-  key: lm-eval-humming-f16-h100
+- label: ":nvidia: (H100) LM Eval Humming Default TEMPORARY Shard %N"
+  key: lm-eval-humming-h100-default
+  optional: true
   timeout_in_minutes: 70
   device: h100
-  optional: true
   num_devices: 1
   parallelism: 3
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/humming.py
-  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/quantization/utils/humming/
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
   - vllm/model_executor/layers/fused_moe/oracle/
   - vllm/model_executor/kernels/linear/
@@ -290,30 +569,50 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt
 
 - label: ":nvidia: (H100) LM Eval Humming Activation FP8/INT8 TEMPORARY"
-  key: lm-eval-humming-act-h100
-  timeout_in_minutes: 70
-  device: h100
+  key: lm-eval-humming-h100-act
   optional: true
+  timeout_in_minutes: 90
+  device: h100
   num_devices: 1
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/humming.py
-  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/quantization/utils/humming/
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
   - vllm/model_executor/layers/fused_moe/oracle/
   - vllm/model_executor/kernels/linear/
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
-  # - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
+
+- label: ":nvidia: (H100) LM Eval Humming DeepEP TEMPORARY"
+  key: lm-eval-humming-h100-deepep
+  optional: true
+  timeout_in_minutes: 60
+  device: h100
+  num_devices: 2
+  source_file_dependencies:
+  - .buildkite/test_areas/lm_eval.yaml
+  - tests/evals/gsm8k/configs/humming/
+  - requirements/cuda.txt
+  - vllm/utils/humming.py
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming/
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/kernels/linear/
+  - vllm/distributed/
+  - csrc/moe/
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-deepep.txt
 
-- label: ":nvidia: (B200) LM Eval Humming FP16 TEMPORARY"
-  key: lm-eval-humming-f16-b200
+- label: ":nvidia: (B200) LM Eval Humming Default TEMPORARY"
+  key: lm-eval-humming-b200-default
+  optional: true
   timeout_in_minutes: 50
   device: b200-k8s
-  optional: true
   num_devices: 1
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/humming.py
-  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/quantization/utils/humming/
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
   - vllm/model_executor/layers/fused_moe/oracle/
   - vllm/model_executor/kernels/linear/
@@ -321,14 +620,14 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
 
 - label: ":nvidia: (B200) LM Eval Humming Activation FP8/INT8 TEMPORARY"
-  key: lm-eval-humming-act-b200
+  key: lm-eval-humming-b200-act
+  optional: true
   timeout_in_minutes: 50
   device: b200-k8s
-  optional: true
   num_devices: 1
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/humming.py
-  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/quantization/utils/humming/
   - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
   - vllm/model_executor/layers/fused_moe/oracle/
   - vllm/model_executor/kernels/linear/
@@ -336,9 +635,9 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
 
-- label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant KV Cache"
-  key: lm-eval-turboquant-kv-cache
-  timeout_in_minutes: 55
+- label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant k8v4"
+  key: lm-eval-turboquant-k8v4
+  timeout_in_minutes: 30
   device: h200_18gb
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/turboquant/
@@ -346,7 +645,155 @@ steps:
   - vllm/v1/attention/ops/triton_turboquant_decode.py
   - vllm/v1/attention/ops/triton_turboquant_store.py
   commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-k8v4.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) LM Eval TurboQuant k8v4"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 30
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - .buildkite/test_areas/lm_eval.yaml
+      - tests/evals/gsm8k/
+      - tests/utils.py
+      - tests/conftest.py
+      - vllm/model_executor/layers/quantization/turboquant/
+      - vllm/v1/attention/backends/turboquant_attn.py
+      - vllm/v1/attention/backends/triton_attn.py
+      - vllm/v1/attention/ops/triton_unified_attention.py
+      - vllm/v1/attention/selector.py
+      - vllm/v1/attention/ops/triton_turboquant_decode.py
+      - vllm/v1/attention/ops/triton_turboquant_store.py
+      - vllm/v1/attention/ops/flydsl_turboquant_decode.py
+      - vllm/v1/attention/ops/flydsl_kernels/
+      - vllm/v1/attention/ops/turboquant_soa/
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-k8v4-rocm.txt
+
+- label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant t4nc"
+  key: lm-eval-turboquant-t4nc
+  timeout_in_minutes: 30
+  device: h200_18gb
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t4nc.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) LM Eval TurboQuant t4nc"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 30
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - .buildkite/test_areas/lm_eval.yaml
+      - tests/evals/gsm8k/
+      - tests/utils.py
+      - tests/conftest.py
+      - vllm/model_executor/layers/quantization/turboquant/
+      - vllm/v1/attention/backends/turboquant_attn.py
+      - vllm/v1/attention/backends/triton_attn.py
+      - vllm/v1/attention/ops/triton_unified_attention.py
+      - vllm/v1/attention/selector.py
+      - vllm/v1/attention/ops/triton_turboquant_decode.py
+      - vllm/v1/attention/ops/triton_turboquant_store.py
+      - vllm/v1/attention/ops/flydsl_turboquant_decode.py
+      - vllm/v1/attention/ops/flydsl_kernels/
+      - vllm/v1/attention/ops/turboquant_soa/
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t4nc-rocm.txt
+
+- label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant k3v4nc"
+  key: lm-eval-turboquant-k3v4nc
+  timeout_in_minutes: 30
+  device: h200_18gb
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-k3v4nc.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) LM Eval TurboQuant k3v4nc"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 30
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - .buildkite/test_areas/lm_eval.yaml
+      - tests/evals/gsm8k/
+      - tests/utils.py
+      - tests/conftest.py
+      - vllm/model_executor/layers/quantization/turboquant/
+      - vllm/v1/attention/backends/turboquant_attn.py
+      - vllm/v1/attention/backends/triton_attn.py
+      - vllm/v1/attention/ops/triton_unified_attention.py
+      - vllm/v1/attention/selector.py
+      - vllm/v1/attention/ops/triton_turboquant_decode.py
+      - vllm/v1/attention/ops/triton_turboquant_store.py
+      - vllm/v1/attention/ops/flydsl_turboquant_decode.py
+      - vllm/v1/attention/ops/flydsl_kernels/
+      - vllm/v1/attention/ops/turboquant_soa/
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-k3v4nc-rocm.txt
+
+- label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant t3nc"
+  key: lm-eval-turboquant-t3nc
+  timeout_in_minutes: 30
+  device: h200_18gb
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) LM Eval TurboQuant t3nc"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 30
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - .buildkite/test_areas/lm_eval.yaml
+      - tests/evals/gsm8k/
+      - tests/utils.py
+      - tests/conftest.py
+      - vllm/model_executor/layers/quantization/turboquant/
+      - vllm/v1/attention/backends/turboquant_attn.py
+      - vllm/v1/attention/backends/triton_attn.py
+      - vllm/v1/attention/ops/triton_unified_attention.py
+      - vllm/v1/attention/selector.py
+      - vllm/v1/attention/ops/triton_turboquant_decode.py
+      - vllm/v1/attention/ops/triton_turboquant_store.py
+      - vllm/v1/attention/ops/flydsl_turboquant_decode.py
+      - vllm/v1/attention/ops/flydsl_kernels/
+      - vllm/v1/attention/ops/turboquant_soa/
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt
 
 - label: ":nvidia: (H100) GPQA Eval (GPT-OSS)"
   key: gpqa-eval-gpt-oss-2xh100
@@ -375,6 +822,28 @@ steps:
   commands:
     - uv pip install --system 'gpt-oss[eval]==0.0.5'
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) GPQA Eval (GPT-OSS)"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 40
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/
+      - vllm/model_executor/layers/quantization
+      - vllm/model_executor/models/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - vllm/model_executor/layers/fused_moe/
+      - tests/evals/gpt_oss/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - uv pip install --system 'gpt-oss[eval]==0.0.5'
+      - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt
 
 - label: ":nvidia: (B200) Qwen3.8-Flash-Next-FP8 Accuracy Eval"
   key: accuracy-eval-qwen3-8-flash-next-fp8-b200
@@ -411,6 +880,29 @@ steps:
   commands:
     - uv pip install --system 'evalscope==1.10.0'
     - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-h200.txt
+  mirror:
+    amd:
+      label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval"
+      dind: false
+      device: mi355_4
+      num_devices: 4
+      timeout_in_minutes: 105
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/libtorch_stable/gdn/
+      - tests/evals/qwen4_exp/
+      - vllm/model_executor/layers/mamba/
+      - vllm/models/qwen4_exp/
+      - vllm/transformers_utils/configs/qwen4_exp.py
+      - vllm/v1/attention/backends/short_conv_attn.py
+      - vllm/platforms/rocm.py
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+      commands:
+      - uv pip install --system 'evalscope==1.10.0'
+      - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt
 
 - label: ":nvidia: (DGX) Spark GPQA Eval (GPT-OSS)"
   key: gpqa-eval-gpt-oss-spark
@@ -470,6 +962,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash"
 
 - label: ":nvidia: (H200 MIG 35GB) MRCR Eval Small Models"
+  key: mrcr-eval-small-models
   device: h200_35gb
   timeout_in_minutes: 25
   source_file_dependencies:
diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml
index d63fb567eaf1..28f695f391a9 100644
--- a/.buildkite/test_areas/lora.yaml
+++ b/.buildkite/test_areas/lora.yaml
@@ -54,13 +54,14 @@ steps:
     - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py
     - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py
     - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py
+    - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp
   mirror:
     amd:
       label: ":amd: (MI300) LoRA TP (Distributed) %N"
       dind: false
       device: mi300_4
       working_dir: "/vllm-workspace/tests"
-      timeout_in_minutes: 25
+      timeout_in_minutes: 65
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -72,10 +73,11 @@ steps:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       # expandable_segments breaks custom all-reduce IPC on ROCm.
       - case "$$BUILDKITE_PARALLEL_JOB" in 0|1|2|3) ;; *) echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" && exit 1;; esac
-      - test "$$BUILDKITE_PARALLEL_JOB" != "3" || pytest -v -s -x lora/test_chatglm3_tp.py
-      - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_llama_tp.py
-      - test "$$BUILDKITE_PARALLEL_JOB" != "3" || pytest -v -s -x lora/test_qwen3_with_multi_loras.py
-      - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_olmoe_tp.py
-      - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py
-      - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py
-      - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_chatglm3_tp.py; fi
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_llama_tp.py; fi
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_qwen3_with_multi_loras.py; fi
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_olmoe_tp.py; fi
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s -x lora/test_gptoss_tp.py; fi
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_qwen35_densemodel_lora.py; fi
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_gemma4_tp.py; fi
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp; fi
diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml
index 5ae3bcf73ecb..ad1c58f08717 100644
--- a/.buildkite/test_areas/misc.yaml
+++ b/.buildkite/test_areas/misc.yaml
@@ -99,12 +99,12 @@ steps:
       label: ":amd: (MI300) V1 Core"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 45
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
 
-- label: ":nvidia: (H200 MIG 35GB) V1 Executor + Worker"
-  device: h200_35gb
+- label: ":nvidia: (H200 MIG 18GB) V1 Executor + Worker"
+  device: h200_18gb
   key: v1-executor-worker
   timeout_in_minutes: 45
   source_file_dependencies: *v1-core-kv-metrics-deps
@@ -119,7 +119,7 @@ steps:
       label: ":amd: (MI300) V1 Executor + Worker"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 45
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
 
@@ -130,6 +130,7 @@ steps:
   source_file_dependencies: *v1-core-kv-metrics-deps
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh
     - pytest -v -s v1/kv_offload
     - pytest -v -s v1/simple_kv_offload
   mirror:
@@ -137,26 +138,27 @@ steps:
       label: ":amd: (MI300) V1 KV Offload"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 45
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
 
-- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors"
+- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors Shard %N"
   device: h200_35gb
   key: v1-kv-connectors
   timeout_in_minutes: 45
+  parallelism: 4
   source_file_dependencies: *v1-core-kv-metrics-deps
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
-    - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
+    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+    - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   mirror:
     amd:
-      label: ":amd: (MI300) V1 KV Connectors"
+      label: ":amd: (MI300) V1 KV Connectors Shard %N"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 45
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
 
@@ -235,6 +237,27 @@ steps:
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration
+  mirror:
+    amd:
+      label: ":amd: (MI355) Extract Hidden States Integration Single GPU"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 45
+      working_dir: /vllm-workspace
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/spec_decode/extract_hidden_states.py
+      - vllm/model_executor/models/extract_hidden_states.py
+      - vllm/transformers_utils/configs/extract_hidden_states.py
+      - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+      - tests/v1/kv_connector/extract_hidden_states_integration
+      - vllm/platforms/rocm.py
+      env:
+        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
+      commands:
+      - pytest -v -s -m 'not distributed' tests/v1/kv_connector/extract_hidden_states_integration
 
 - label: ":nvidia: (L4) Extract Hidden States Integration"
   key: extract-hidden-states-integration-2-gpus
@@ -251,6 +274,35 @@ steps:
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s -m 'distributed' tests/v1/kv_connector/extract_hidden_states_integration
+  mirror:
+    amd:
+      label: ":amd: (MI355) Extract Hidden States Integration"
+      dind: false
+      device: mi355_2
+      timeout_in_minutes: 30
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/config/speculative.py
+      - vllm/distributed/kv_transfer/kv_connector/
+      - vllm/model_executor/layers/attention/
+      - vllm/model_executor/layers/mamba/
+      - vllm/model_executor/model_loader/
+      - vllm/model_executor/models/extract_hidden_states.py
+      - vllm/model_executor/models/llama.py
+      - vllm/model_executor/models/qwen3_5.py
+      - vllm/model_executor/models/qwen3_next.py
+      - vllm/model_executor/models/registry.py
+      - vllm/transformers_utils/configs/extract_hidden_states.py
+      - vllm/transformers_utils/configs/qwen3_5.py
+      - vllm/v1/attention/backends/
+      - vllm/v1/attention/selector.py
+      - vllm/v1/kv_cache_interface.py
+      - vllm/v1/spec_decode/extract_hidden_states.py
+      - vllm/v1/worker/gpu_model_runner.py
+      - vllm/_aiter_ops.py
+      - tests/v1/kv_connector/extract_hidden_states_integration/
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 18GB) Regression"
   key: regression
@@ -274,6 +326,15 @@ steps:
   - pip install 'modelscope<1.38'
   - pytest -v -s test_regression.py
   working_dir: "/vllm-workspace/tests" # optional
+  mirror:
+    amd:
+      label: ":amd: (MI355) Regression"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 25
+      depends_on:
+      - image-build-amd
+      working_dir: /vllm-workspace/tests
 
 - label: ":nvidia: (H200 MIG 35GB) Examples"
   device: h200_35gb
@@ -376,7 +437,6 @@ steps:
       dind: false
       device: mi300_1
       timeout_in_minutes: 55
-      soft_fail: true
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -415,9 +475,15 @@ steps:
       label: ":amd: (MI355) Async Engine, Inputs, Utils, Worker"
       dind: false
       device: mi355_1
-      timeout_in_minutes: 40
+      timeout_in_minutes: 50
       depends_on:
       - image-build-amd
+      commands:
+      - pytest -v -s detokenizer
+      - pytest -v -s -m 'not cpu_test' multimodal
+      # test_cache.py marks every test cpu_test, including this GPU regression.
+      - pytest -v -s multimodal/test_cache.py::test_sleep_wake_preserves_mm_cache_consistency
+      - pytest -v -s utils_
 
 - label: ":computer: (CPU) Params, Env, Tokenizers, Parser"
   key: cpu-params-env-tokenizers-parser
@@ -452,6 +518,7 @@ steps:
   - vllm/v1/
   - tests/test_envs.py
   - tests/test_outputs.py
+  - tests/test_pcp_dp.py
   - tests/test_pooling_params.py
   - tests/test_ray_env.py
   - tests/test_sampling_params.py
@@ -459,6 +526,9 @@ steps:
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
   - tests/reasoning
+  - tools/pre_commit/check_test_tethering.py
+  - tools/pre_commit/test_tethering_allowlist.txt
+  - tests/tools/test_check_test_tethering.py
   - tests/tool_parsers
   - tests/tokenizers_
   - tests/parser
@@ -469,6 +539,7 @@ steps:
   - python3 standalone_tests/lazy_imports.py
   - pytest -v -s test_envs.py
   - pytest -v -s test_outputs.py
+  - pytest -v -s test_pcp_dp.py
   - pytest -v -s test_pooling_params.py
   - pytest -v -s test_ray_env.py
   - pytest -v -s test_sampling_params.py
@@ -486,6 +557,7 @@ steps:
   commands:
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s config
+  - pytest -v -s tools/test_check_test_tethering.py
 
 - label: ":computer: (CPU) Reasoning + Renderers"
   key: cpu-reasoning-renderers
@@ -537,6 +609,25 @@ steps:
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
+  mirror:
+    amd:
+      label: ":amd: (MI250) Batch Invariance"
+      dind: false
+      device: mi250_1
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/attention
+      - vllm/model_executor/layers
+      - tests/v1/determinism/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pip install pytest-timeout pytest-forked
+      - pytest -v -s v1/determinism/test_batch_invariance.py
+      - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
 
 - label: ":nvidia: (B200) Batch Invariance"
   key: batch-invariance-b200
@@ -574,3 +665,17 @@ steps:
   commands:
     - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
     - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
+  mirror:
+    amd:
+      label: ":amd: (MI355) Acceptance Length (Large Models)"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      working_dir: /vllm-workspace/tests
+      source_file_dependencies:
+      - vllm/v1/spec_decode/
+      - vllm/model_executor/models/mlp_speculator.py
+      - tests/v1/spec_decode/test_acceptance_length.py
+      - vllm/platforms/rocm.py
diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml
index 17818cc15260..f6905f98ef35 100644
--- a/.buildkite/test_areas/model_executor.yaml
+++ b/.buildkite/test_areas/model_executor.yaml
@@ -33,7 +33,7 @@ steps:
       label: ":amd: (MI300) Model Executor"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 60
+      timeout_in_minutes: 75
       depends_on:
       - image-build-amd
       source_file_dependencies:
diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml
index 93d6962868a4..6f570b3459ef 100644
--- a/.buildkite/test_areas/models_basic.yaml
+++ b/.buildkite/test_areas/models_basic.yaml
@@ -19,7 +19,7 @@ steps:
       label: ":amd: (MI355) Basic Models (Initialization)"
       dind: false
       device: mi355_1
-      timeout_in_minutes: 25
+      timeout_in_minutes: 30
       depends_on:
       - image-build-amd
 
@@ -37,27 +37,51 @@ steps:
     # test.) Also run if model initialization test file is modified
     - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 14
+  mirror:
+    amd:
+      label: ":amd: (MI355) Basic Models (Extra Initialization) Shard %N"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 60
+      parallelism: 10
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/models/
+      - vllm/model_executor/layers/
+      - tests/models/test_initialization.py
+      - tests/models/registry.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 35GB) Basic Models (Other)"
   device: h200_35gb
   key: basic-models-tests-other
-  timeout_in_minutes: 35
+  timeout_in_minutes: 38
   source_file_dependencies:
   - vllm/
   - "!vllm/distributed/kv_transfer/"
   - tests/models/test_terratorch.py
   - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
+  - tests/models/test_deepseek_v4_vl_rocm.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py
   mirror:
     amd:
       label: ":amd: (MI300) Basic Models (Other)"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 50
+      timeout_in_minutes: 65
       depends_on:
       - image-build-amd
+      source_file_dependencies:
+      - tests/models/test_deepseek_v4_vl_rocm.py
+      - tests/models/test_hyv4_rocm.py
+      commands:
+      - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py
+      - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
 
 - label: ":nvidia: (B200) Inkling"
   key: inkling-unit-tests-b200
@@ -85,6 +109,23 @@ steps:
   commands:
     # The native NVIDIA Kimi K3 kernels require the SM100 family.
     - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Kimi K3"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 60
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
+      commands:
+      # The ROCm KDA/AttnRes kernels are gated on gfx950, so this has to be MI355.
+      # The two kernels/ files the NVIDIA job adds are CUDA-only (they import
+      # cuda.bindings.driver), so they are deliberately left out here.
+      - pytest -v -s models/kimi_k3
 
 - label: ":nvidia: (H200) GLM5Next"
   key: glm5next-unit-tests
@@ -95,6 +136,17 @@ steps:
   - tests/models/glm5next/
   commands:
     - pytest -v -s models/glm5next
+  mirror:
+    amd:
+      label: ":amd: (MI355) GLM5Next"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 35
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200) Qwen4 Exp"
   key: qwen4-exp-unit-tests
@@ -146,6 +198,17 @@ steps:
   - tests/models/test_adapters.py
   - tests/models/test_qwen3_5_mtp_config.py
   - tests/models/transformers/fusers/
+  - tests/model_executor/layers/test_activation.py
   device: cpu-small
   commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/
+    - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py
+  mirror:
+    amd:
+      label: ":amd: (MI250) Basic Models Other"
+      dind: false
+      device: mi250_1
+      no_gpu: true
+      working_dir: "/vllm-workspace/tests"
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml
index cf7623ce3ec8..a054c00db860 100644
--- a/.buildkite/test_areas/models_distributed.yaml
+++ b/.buildkite/test_areas/models_distributed.yaml
@@ -38,10 +38,17 @@ steps:
       label: ":amd: (MI300) Distributed Models Shard %N"
       dind: false
       device: mi300_2
-      timeout_in_minutes: 35
+      timeout_in_minutes: 80
       depends_on:
       - image-build-amd
       source_file_dependencies:
+      - vllm/models/hy_v4/
+      - vllm/config/
+      - vllm/v1/worker/
+      - vllm/v1/attention/ops/
+      - vllm/v1/attention/backend.py
+      - tests/conftest.py
+      - tests/utils.py
       - vllm/model_executor/model_loader/sharded_state_loader.py
       - vllm/model_executor/models/
       - vllm/model_executor/layers/
@@ -56,16 +63,17 @@ steps:
       - |
         case "$$BUILDKITE_PARALLEL_JOB" in
           0)
-            TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
+            TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' &&
             HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
             ;;
           1)
-            pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
-            pytest models/language -v -s -m 'distributed(num_gpus=2)'
+            pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' &&
+            pytest models/language -v -s -m 'distributed(num_gpus=2)' &&
+            VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)'
             ;;
           2)
-            pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
-            pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
+            pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py &&
+            pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' &&
             VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
             ;;
           *) exit 2 ;;
diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml
index 68f1cdd429ce..c8473e7f0b60 100644
--- a/.buildkite/test_areas/models_language.yaml
+++ b/.buildkite/test_areas/models_language.yaml
@@ -19,7 +19,7 @@ steps:
       label: ":amd: (MI300) Language Models (Standard)"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 45
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
 
@@ -82,16 +82,16 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      label: ":amd: (MI300) Language Models (Hybrid) Shard %N"
+      label: ":amd: (MI355) Language Models (Hybrid) Shard %N"
       dind: false
-      device: mi300_1
-      timeout_in_minutes: 60
+      device: mi355_1
+      timeout_in_minutes: 65
       depends_on:
       - image-build-amd
       commands:
       - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
       - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-      - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+      - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
 # Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
 # autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
@@ -111,6 +111,27 @@ steps:
     - rm -rf /tmp/mamba-src && git clone --depth 1 --branch v2.3.0 https://github.com/state-spaces/mamba /tmp/mamba-src && sed -i 's/-std=c++17/-std=c++20/g' /tmp/mamba-src/setup.py && MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation /tmp/mamba-src
     - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
     - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
+  mirror:
+    amd:
+      label: ":amd: (MI355) Granite Language Model Compatibility"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 80
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/
+      - "!vllm/distributed/kv_transfer/"
+      - tests/models/language/generation
+      commands:
+      - export GIT_CONFIG_COUNT=1
+      - export GIT_CONFIG_KEY_0=http.version
+      - export GIT_CONFIG_VALUE_0=HTTP/1.1
+      - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+      - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+      - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
 
 - label: ":nvidia: (H200 MIG 35GB) Language Models (Extended Generation)"
   device: h200_35gb
@@ -133,7 +154,7 @@ steps:
       label: ":amd: (MI355) Language Models (Extended Generation)"
       dind: false
       device: mi355_1
-      timeout_in_minutes: 70
+      timeout_in_minutes: 95
       depends_on:
       - image-build-amd
       commands:
@@ -155,6 +176,7 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) Language Models (PPL)"
+      dind: false
       device: mi250_1
       timeout_in_minutes: 65
       depends_on:
@@ -197,11 +219,8 @@ steps:
       dind: false
       device: mi300_1
       timeout_in_minutes: 95
-      parallelism: 2
       depends_on:
       - image-build-amd
-      commands:
-      - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
 
 - label: ":nvidia: (H200 MIG 18GB) Language Models (MTEB)"
   key: language-models-test-mteb
@@ -217,6 +236,7 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) Language Models (MTEB)"
+      dind: false
       device: mi250_1
       timeout_in_minutes: 60
       depends_on:
diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml
index 9893383bff1b..779801a2ea2b 100644
--- a/.buildkite/test_areas/models_multimodal.yaml
+++ b/.buildkite/test_areas/models_multimodal.yaml
@@ -57,6 +57,7 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl"
+      dind: false
       device: mi250_1
       timeout_in_minutes: 55
       depends_on:
@@ -77,9 +78,9 @@ steps:
     - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
   mirror:
     amd:
-      label: ":amd: (MI300) Multimodal Models (Standard) 4: other + whisper"
+      label: ":amd: (MI355) Multimodal Models (Standard) 4: other + whisper"
       dind: false
-      device: mi300_1
+      device: mi355_1
       timeout_in_minutes: 70
       depends_on:
       - image-build-amd
@@ -98,23 +99,36 @@ steps:
   commands:
     - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 4
+  mirror:
+    amd:
+      label: ":amd: (MI250) Multimodal Processor Shard %N"
+      dind: false
+      device: mi250_1
+      no_gpu: true
+      working_dir: "/vllm-workspace/tests"
+      timeout_in_minutes: 90
+      parallelism: 6
+      depends_on:
+      - image-build-amd
 
-- label: ":nvidia: (H200 MIG 18GB) Multimodal Processor"
+- label: ":nvidia: (H200 MIG 18GB) Multimodal Processor Shard %N"
   key: multi-modal-processor
   timeout_in_minutes: 98
   device: h200_18gb
+  parallelism: 4
   source_file_dependencies:
   - vllm/
   - "!vllm/distributed/kv_transfer/"
   - tests/models/multimodal
   - tests/models/registry.py
   commands:
-    - pytest -v -s models/multimodal/processing/test_tensor_schema.py
+    - pytest -v -s models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   mirror:
     amd:
-      label: ":amd: (MI355) Multimodal Processor"
+      label: ":amd: (MI355) Multimodal Processor Shard %N"
       dind: false
       device: mi355_1
+      parallelism: 4
       timeout_in_minutes: 115
       depends_on:
       - image-build-amd
diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml
index 65e7a0b258a0..6b901d8f490b 100644
--- a/.buildkite/test_areas/plugins.yaml
+++ b/.buildkite/test_areas/plugins.yaml
@@ -53,6 +53,7 @@ steps:
   mirror:
     amd:
       label: ":amd: (MI250) Plugin Integration"
+      dind: false
       device: mi250_2
       timeout_in_minutes: 50
       depends_on:
@@ -75,6 +76,20 @@ steps:
   commands:
     - pip install "vllm-gguf-plugin >= 0.0.2"
     - pytest -v -s plugins_tests/gguf
+  mirror:
+    amd:
+      label: ":amd: (MI355) GGUF Plugin"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 45
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/layers/quantization
+      - tests/plugins_tests/test_gguf_plugin.py
+      - tests/plugins_tests/gguf
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 18GB) BitsAndBytes Plugin"
   key: bitsandbytes-plugin
diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml
index eb5b8407033c..8204fb9670a2 100644
--- a/.buildkite/test_areas/pytorch.yaml
+++ b/.buildkite/test_areas/pytorch.yaml
@@ -87,6 +87,20 @@ steps:
     - tests/compile/h100/
   commands:
   - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
+  mirror:
+    amd:
+      label: ":amd: (MI355) PyTorch Compilation H100 Cases"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 45
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - tests/utils.py
+      - csrc/
+      commands:
+      - pytest -v -s compile/h100/test_startup.py
 
 - label: ":nvidia: (L4) PyTorch Compilation Passes"
   key: pytorch-compilation-passes-unit-tests
@@ -122,6 +136,21 @@ steps:
     - tests/compile/passes
   commands:
   - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      label: ':amd: (MI300) PyTorch Compilation Passes'
+      dind: false
+      device: mi300_1
+      timeout_in_minutes: 80
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/
+      - tests/compile/passes
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -s -v compile/passes --ignore compile/passes/distributed
 
 - label: ":nvidia: (H200 MIG 35GB) PyTorch Fullgraph"
   device: h200_35gb
@@ -154,12 +183,21 @@ steps:
   - vllm/utils/
   - vllm/v1/
   - tests/compile
+  - csrc/
   commands:
   # Run smoke tests under fullgraph directory, except test_full_cudagraph.py
   # as it is a heavy test that is covered in other steps.
   # Use `find` to launch multiple instances of pytest so that
   # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
   - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
+  mirror:
+    amd:
+      label: ":amd: (MI355) PyTorch Fullgraph"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 90
+      depends_on:
+      - image-build-amd
 
 # Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model
 # instances and cannot fit a 35GB MIG slice. L4 retains the original coverage:
diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml
index 1d895e1f2f90..702bbe702452 100644
--- a/.buildkite/test_areas/quantization.yaml
+++ b/.buildkite/test_areas/quantization.yaml
@@ -20,6 +20,31 @@ steps:
   # The SM90-only checkpoint currently contains a removed weight_chan_scale
   # parameter. It was not exercised by the previous L4 job.
   - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  mirror:
+    amd:
+      label: ":amd: (MI355) Quantization Shard %N"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 115
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/
+      - vllm/model_executor/layers/quantization
+      - tests/quantization
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      - vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+      - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+      - tests/rocm/test_moe_weight_replay.py
+      commands:
+      # Use the ROCm default instead of the NVIDIA runner override.
+      - unset VLLM_USE_V2_MODEL_RUNNER
+      - uv pip install --system torchao==0.17.0
+      - uv pip install --system conch-triton-kernels
+      - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+      - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi
 
 - label: ":nvidia: (H200 MIG 35GB) Quantized Fusions"
   device: h200_35gb
@@ -63,3 +88,22 @@ steps:
   - tests/models/quantization
   commands:
     - pytest -v -s models/quantization
+  mirror:
+    amd:
+      label: ":amd: (MI355) Quantized Models"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 80
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/layers/quantization
+      - tests/models/quantization
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      - vllm/model_executor/model_loader/
+      commands:
+      # Use the ROCm default instead of the NVIDIA runner override.
+      - unset VLLM_USE_V2_MODEL_RUNNER
+      - pytest -v -s models/quantization
diff --git a/.buildkite/test_areas/ray_compat.yaml b/.buildkite/test_areas/ray_compat.yaml
index d24c38430b8d..dc8571b26045 100644
--- a/.buildkite/test_areas/ray_compat.yaml
+++ b/.buildkite/test_areas/ray_compat.yaml
@@ -16,3 +16,19 @@ steps:
   - setup.py
   commands:
   - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh
+  mirror:
+    amd:
+      label: ":amd: (MI355) Ray Dependency Compatibility Check"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 25
+      working_dir: "/"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - .buildkite/scripts/check-ray-compatibility.sh
+      - requirements/
+      - setup.py
+      - vllm/platforms/rocm.py
+      commands:
+      - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh
diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml
index 8ccc0c50e926..b8d0e58f5810 100644
--- a/.buildkite/test_areas/rust_frontend.yaml
+++ b/.buildkite/test_areas/rust_frontend.yaml
@@ -3,6 +3,7 @@ depends_on:
   - image-build
 steps:
 - label: ":nvidia: (H200 MIG 18GB) Rust Frontend OpenAI Coverage"
+  key: rust-frontend-openai-coverage
   timeout_in_minutes: 30
   device: h200_18gb
   working_dir: "/vllm-workspace/tests"
@@ -67,6 +68,7 @@ steps:
       - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 18GB) Rust Frontend Serve/Admin Coverage"
+  key: rust-frontend-serve-admin-coverage
   timeout_in_minutes: 25
   device: h200_18gb
   working_dir: "/vllm-workspace/tests"
@@ -116,6 +118,7 @@ steps:
       - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 18GB) Rust Frontend Core Correctness"
+  key: rust-frontend-core-correctness
   timeout_in_minutes: 20
   device: h200_18gb
   working_dir: "/vllm-workspace/tests"
@@ -128,8 +131,23 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+  mirror:
+    amd:
+      label: ":amd: (MI355) Rust Frontend Core Correctness"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 25
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - rust/
+      - vllm/entrypoints/openai/
+      - tests/utils.py
+      - tests/entrypoints/openai/correctness/test_lmeval.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (H200 MIG 35GB) Rust Frontend Tool Use"
+  key: rust-frontend-tool-use
   device: h200_35gb
   timeout_in_minutes: 25
   working_dir: "/vllm-workspace/tests"
@@ -143,8 +161,24 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
+  mirror:
+    amd:
+      label: ":amd: (MI355) Rust Frontend Tool Use"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 30
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - rust/
+      - vllm/entrypoints/openai/
+      - vllm/tool_parsers/
+      - tests/utils.py
+      - tests/tool_use/
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (L4) Rust Frontend Distributed"
+  key: rust-frontend-distributed
   timeout_in_minutes: 25
   device: l4
   num_devices: 4
@@ -170,3 +204,27 @@ steps:
   - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
   - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
   - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
+  mirror:
+    amd:
+      label: ":amd: (MI355) Rust Frontend Distributed"
+      dind: false
+      device: mi355_4
+      timeout_in_minutes: 45
+      env:
+        VLLM_ENGINE_READY_TIMEOUT_S: "1800"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - rust/
+      - vllm/distributed/
+      - "!vllm/distributed/kv_transfer/"
+      - vllm/engine/
+      - vllm/v1/executor/
+      - vllm/v1/engine/
+      - vllm/v1/worker/
+      - tests/utils.py
+      - tests/v1/distributed/test_dense_dp_world_size.py
+      - tests/v1/distributed/test_external_lb_dp.py
+      - tests/v1/distributed/test_hybrid_lb_dp.py
+      - tests/v1/distributed/test_internal_lb_dp.py
+      - vllm/platforms/rocm.py
diff --git a/.buildkite/test_areas/rust_frontend_cargo.yaml b/.buildkite/test_areas/rust_frontend_cargo.yaml
index 11bca6252667..53d1da7edf41 100644
--- a/.buildkite/test_areas/rust_frontend_cargo.yaml
+++ b/.buildkite/test_areas/rust_frontend_cargo.yaml
@@ -14,6 +14,16 @@ steps:
   - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
   commands:
   - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
+  mirror:
+    amd:
+      label: ":amd: (MI250) Rust Frontend Cargo Style + Clippy"
+      dind: false
+      device: mi250_1
+      no_gpu: true
+      working_dir: "/vllm-workspace"
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
 
 - label: ":computer: (CPU) Rust Frontend Cargo"
   key: rust-frontend-cargo-tests
@@ -28,3 +38,13 @@ steps:
   - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
   commands:
   - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
+  mirror:
+    amd:
+      label: ":amd: (MI250) Rust Frontend Cargo"
+      dind: false
+      device: mi250_1
+      no_gpu: true
+      working_dir: "/vllm-workspace"
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml
index cfde3717c860..084583c8a732 100644
--- a/.buildkite/test_areas/samplers.yaml
+++ b/.buildkite/test_areas/samplers.yaml
@@ -2,10 +2,10 @@ group: Samplers
 depends_on: 
   - image-build
 steps:
-- label: ":nvidia: (H200 MIG 35GB) Samplers"
-  device: h200_35gb
+- label: ":nvidia: (H200 MIG 18GB) Samplers"
+  device: h200_18gb
   key: samplers-test
-  timeout_in_minutes: 40
+  timeout_in_minutes: 60
   source_file_dependencies:
   - vllm/model_executor/layers
   - vllm/sampling_metadata.py
@@ -15,11 +15,12 @@ steps:
   commands:
     # VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
     # values explicitly to still cover the PyTorch-native (Triton) path.
-    - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers
-    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
+    - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
+    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
   mirror:
     amd:
       label: ":amd: (MI250) Samplers"
+      dind: false
       device: mi250_1
       timeout_in_minutes: 40
       depends_on:
@@ -34,4 +35,40 @@ steps:
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
       commands:
-      - pytest -v -s samplers
+      - pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
+
+# Qwen2-Audio-7B leaves no room for KV cache on a 16 GiB slice.
+- label: ":nvidia: (H200 MIG 35GB) Samplers Multimodal Beam Search"
+  device: h200_35gb
+  key: samplers-multimodal-beam-search
+  timeout_in_minutes: 30
+  source_file_dependencies:
+  - vllm/model_executor/layers
+  - vllm/sampling_metadata.py
+  - tests/samplers
+  - tests/conftest.py
+  - vllm/entrypoints/generate/beam_search
+  commands:
+    - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
+    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
+  mirror:
+    amd:
+      label: ":amd: (MI355) Samplers Multimodal Beam Search"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 45
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/model_executor/layers
+      - vllm/sampling_metadata.py
+      - vllm/v1/sample/
+      - vllm/entrypoints/generate/beam_search/
+      - tests/samplers
+      - tests/conftest.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      commands:
+      - pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml
index 838727035980..fb0130600058 100644
--- a/.buildkite/test_areas/spec_decode.yaml
+++ b/.buildkite/test_areas/spec_decode.yaml
@@ -47,7 +47,7 @@ steps:
       label: ":amd: (MI300) Spec Decode Eagle 1: DeepSeek + Qwen"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 35
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -74,7 +74,7 @@ steps:
       label: ":amd: (MI300) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 25
+      timeout_in_minutes: 60
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -172,22 +172,24 @@ steps:
       - tests/v1/e2e/spec_decode/
       - vllm/platforms/rocm.py
 
-- label: ":nvidia: (H200 MIG 18GB) Spec Decode Draft Model"
+- label: ":nvidia: (H200 MIG 18GB) Spec Decode Draft Model Shard %N"
   key: spec-decode-draft-model
-  timeout_in_minutes: 45
+  timeout_in_minutes: 60
   device: h200_18gb
+  parallelism: 4
   source_file_dependencies:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
     - tests/v1/e2e/spec_decode/
   commands:
-    - pytest -v -s v1/e2e/spec_decode/draft_model/
+    - pytest -v -s v1/e2e/spec_decode/draft_model/ --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   mirror:
     amd:
-      label: ":amd: (MI300) Spec Decode Draft Model"
+      label: ":amd: (MI300) Spec Decode Draft Model Shard %N"
       dind: false
       device: mi300_1
-      timeout_in_minutes: 55
+      timeout_in_minutes: 60
+      parallelism: 2
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -233,6 +235,9 @@ steps:
       timeout_in_minutes: 40
       depends_on:
       - image-build-amd
+      source_file_dependencies:
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
 
 - label: ":nvidia: (B200) Spec Decode DeepSeek MTP Parallel Load"
   key: spec-decode-deepseek-mtp-parallel-load-2xb200-2xmi300
@@ -280,6 +285,8 @@ steps:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
     - vllm/model_executor/models/qwen3_dflash.py
+    - vllm/model_executor/models/qwen3_dflash2.py
+    - vllm/model_executor/models/qwen3_5.py
     - vllm/model_executor/models/laguna_dflash.py
     - tests/v1/e2e/spec_decode/
   commands:
@@ -290,6 +297,8 @@ steps:
       dind: false
       device: mi300_1
       timeout_in_minutes: 120
+      commands:
+      - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2"
       depends_on:
       - image-build-amd
       source_file_dependencies:
@@ -362,6 +371,33 @@ steps:
     # Each Gemma4 parametrization is an indivisible ~23m floor. Keep them in
     # separate shards and put the remaining tests in shard 0, where new tests
     # land by default.
-    - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"
-    - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"
-    - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"
+    - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi
+    - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi
+    - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi
+  mirror:
+    amd:
+      label: ":amd: (MI355) Spec Decode AL MTP + Other Acceptance Nightly"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      working_dir: /vllm-workspace/tests
+      source_file_dependencies:
+      - vllm/v1/spec_decode/
+      - vllm/v1/sample/
+      - vllm/v1/worker/gpu/
+      - vllm/v1/worker/gpu_model_runner.py
+      - vllm/v1/worker/gpu_worker.py
+      - vllm/model_executor/model_loader/
+      - vllm/model_executor/models/gemma4.py
+      - vllm/model_executor/models/gemma4_mm.py
+      - vllm/model_executor/models/gemma4_mtp.py
+      - vllm/model_executor/models/llama.py
+      - vllm/model_executor/models/llama_eagle3.py
+      - vllm/model_executor/models/medusa.py
+      - vllm/model_executor/models/registry.py
+      - tests/evals/gsm8k/
+      - tests/utils.py
+      - tests/v1/e2e/spec_decode/
+      - vllm/platforms/rocm.py
diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml
index 9e8375587adb..201f9b6093a3 100644
--- a/.buildkite/test_areas/torch_abi.yaml
+++ b/.buildkite/test_areas/torch_abi.yaml
@@ -4,7 +4,7 @@ depends_on:
 steps:
 - label: ":nvidia: (L4) Torch Stable ABI Audit"
   key: torch-stable-abi-audit
-  timeout_in_minutes: 5
+  timeout_in_minutes: 10
   source_file_dependencies:
   - .buildkite/check-torch-abi.py
   - csrc/
diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml
deleted file mode 100644
index 154b355778b2..000000000000
--- a/.buildkite/test_areas/weight_loading.yaml
+++ /dev/null
@@ -1,38 +0,0 @@
-group: Weight Loading
-depends_on: 
-  - image-build
-steps:
-- label: ":nvidia: (L4) Weight Loading Multi-GPU"
-  key: weight-loading-multiple-gpu
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  device: l4
-  num_devices: 2
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - "!vllm/distributed/kv_transfer/"
-  - tests/weight_loading
-  commands:
-    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
-  mirror:
-    amd:
-      label: ":amd: (MI355) Weight Loading Multi-GPU"
-      dind: false
-      device: mi355_2
-      timeout_in_minutes: 40
-      depends_on:
-      - image-build-amd
-      commands:
-      - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
-
-# - label: Weight Loading Multiple GPU - Large Models # optional
-#   working_dir: "/vllm-workspace/tests"
-#   num_devices: 2
-#   device: a100
-#   optional: true
-#   source_file_dependencies:
-#   - vllm/
-#   - tests/weight_loading
-#   commands:
-#     - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
diff --git a/.gemini/config.yaml b/.gemini/config.yaml
deleted file mode 100644
index 2499d3f09510..000000000000
--- a/.gemini/config.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github
-have_fun: false  # Just review the code
-code_review:
-  comment_severity_threshold: HIGH  # Reduce quantity of comments
-  pull_request_opened:
-    summary: false  # Don't summarize the PR in a separate comment
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 1f605db84a7e..f64cf5f56107 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -34,9 +34,9 @@
 
 # Entrypoints
 /vllm/entrypoints/anthropic @mgoin @DarkLight1337
-/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
+/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337
 /vllm/entrypoints/mcp @heheda12345
-/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
+/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337
 /vllm/entrypoints/speech_to_text/realtime @njhill
 /vllm/entrypoints/speech_to_text @NickLucche
 /vllm/entrypoints/pooling @noooop
@@ -73,7 +73,7 @@
 /vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium
 /vllm/v1/sample @22quinn @houseroad @njhill
 /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni
-/vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett
+/vllm/v1/structured_output @mgoin @aarnphm @benchislett @arpera
 /vllm/v1/kv_cache_interface.py @heheda12345 @ivanium
 /vllm/v1/kv_offload @ApostaC @orozery
 /vllm/v1/simple_kv_offload @ivanium
@@ -107,10 +107,9 @@
 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan
 /tests/multimodal/media @Isotr0py
 /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas
-/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
-/tests/v1/structured_output @mgoin @russellb @aarnphm
+/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera
+/tests/v1/structured_output @mgoin @aarnphm @arpera
 /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium
-/tests/weight_loading @mgoin @youkaichao @yewentao256
 /tests/lora @jeejeelee
 /tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91
 /tests/v1/kv_connector/nixl_integration @NickLucche
@@ -232,6 +231,6 @@ mkdocs.yaml @hmellor
 /docs/training/weight_transfer @aoshen02
 
 # Security guide and policies
-/docs/usage/security.md @russellb @jperezdealgaba
-/SECURITY.md @russellb @jperezdealgaba
-/docs/contributing/vulnerability_management.md @russellb @jperezdealgaba
+/docs/usage/security.md @jperezdealgaba
+/SECURITY.md @jperezdealgaba
+/docs/contributing/vulnerability_management.md @jperezdealgaba
diff --git a/.github/mergify.yml b/.github/mergify.yml
index 04a623a499d1..1267f0264dfe 100644
--- a/.github/mergify.yml
+++ b/.github/mergify.yml
@@ -330,12 +330,31 @@ pull_request_rules:
     - label != stale
     - or:
       - files~=^vllm/models/deepseek_v4/
-      - title~=(?i)(?:\bDSv4\b|deepseek[-\s]?v4)
+      # Trailing (?![._]?1) keeps V4.1 out: DeepseekV41ForCausalLM is a
+      # distinct architecture with its own tree and label (DSv4.1).
+      - title~=(?i)(?:\bDSv4\b|deepseek[-\s_]?v4)(?![._]?1)
   actions:
     label:
       add:
         - DSv4
 
+- name: label-dsv4-1
+  description: Automatically apply DSv4.1 label
+  conditions:
+    - label != stale
+    - or:
+      # The 4.1 sources use deepseek_v41 and deepseekv41 (tool parsers).
+      - files~=^vllm/models/deepseek_v41/
+      - files~=^vllm/(?:parser|tokenizers|reasoning|transformers_utils/configs)/deepseek_v41.*\.py
+      - files~=^vllm/tool_parsers/deepseekv41.*\.py
+      - files~=^rust/src/(?:chat|parser)/src/.*deepseek_v41
+      - files~=^tests/.*deepseek_?v41
+      - title~=(?i)(?:\bDSv4[._]?1\b|deepseek[-\s_]?v4[._]?1)
+  actions:
+    label:
+      add:
+        - DSv4.1
+
 - name: label-nvidia
   description: Automatically apply nvidia label
   conditions:
diff --git a/.github/workflows/buf.yml b/.github/workflows/buf.yml
index fda4add1d726..dd695d28d529 100644
--- a/.github/workflows/buf.yml
+++ b/.github/workflows/buf.yml
@@ -5,12 +5,6 @@ on:
     paths:
       - "rust/proto/**"
       - ".github/workflows/buf.yml"
-  push:
-    tags:
-      - "v*"
-  schedule:
-    - cron: "0 8 * * *"
-  workflow_dispatch:
 
 permissions:
   contents: read
@@ -21,37 +15,14 @@ concurrency:
 
 jobs:
   proto:
-    if: github.repository_owner == 'vllm-project'
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
-        with:
-          fetch-depth: 0
       - uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0
         with:
           version: 1.72.0
           input: rust/proto
-          token: ${{ github.repository == 'vllm-project/vllm' && github.event_name != 'pull_request' && secrets.BUF_TOKEN || '' }}
           format: false
           breaking: false
           pr_comment: false
           push: false
-      - name: Publish nightly schema
-        if: github.repository == 'vllm-project/vllm' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')
-        run: >-
-          buf push rust/proto
-          --create
-          --create-default-label main
-          --create-visibility public
-          --label nightly
-          --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}"
-      - name: Publish release schema
-        if: github.repository == 'vllm-project/vllm' && github.event_name == 'push'
-        run: >-
-          buf push rust/proto
-          --create
-          --create-default-label main
-          --create-visibility public
-          --label main
-          --label "${GITHUB_REF_NAME}"
-          --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}"
diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml
index b2392e87ce5c..fc6e4921ef83 100644
--- a/.github/workflows/issue_autolabel.yml
+++ b/.github/workflows/issue_autolabel.yml
@@ -197,9 +197,12 @@ jobs:
               deepseek: {
                 keywords: [{ term: "DeepSeek", searchIn: "title" }],
               },
+              // DSv4 and DSv4.1 are distinct architectures; these keep them exclusive.
               DSv4: {
-                keywords: [{ term: "DSv4", searchIn: "title" }],
-                substrings: [{ term: "deepseek-v4", searchIn: "title" }],
+                regexPatterns: [{ pattern: "(?:\\bDSv4\\b|deepseek[-\\s_]?v4)(?![._]?1)", flags: "gi", searchIn: "title" }],
+              },
+              "DSv4.1": {
+                regexPatterns: [{ pattern: "(?:\\bDSv4[._]?1\\b|deepseek[-\\s_]?v4[._]?1)", flags: "gi", searchIn: "title" }],
               },
               llama: {
                 keywords: [{ term: "Llama", searchIn: "title" }],
diff --git a/.github/workflows/proto-crate.yml b/.github/workflows/proto-crate.yml
new file mode 100644
index 000000000000..43be5c8c9214
--- /dev/null
+++ b/.github/workflows/proto-crate.yml
@@ -0,0 +1,70 @@
+name: Rust proto crate
+
+on:
+  pull_request:
+    paths:
+      - "rust/proto/**"
+      - "rust/Cargo.toml"
+      - "rust/Cargo.lock"
+      - "rust-toolchain.toml"
+      - ".github/workflows/proto-crate.yml"
+  push:
+    tags:
+      - "proto-v*"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+jobs:
+  semver:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - uses: obi1kenobi/cargo-semver-checks-action@6b69fcf40e9b5fb17adeb57e4b6ecd020649a239 # v2
+        with:
+          manifest-path: rust/proto/Cargo.toml
+          package: vllm-proto
+
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - run: rustup show active-toolchain
+      - name: Check release version
+        if: github.event_name == 'push'
+        run: |
+          version=$(cargo metadata --manifest-path rust/Cargo.toml --no-deps --format-version 1 |
+            jq -er '.packages[] | select(.name == "vllm-proto") | .version')
+          if [[ "$GITHUB_REF_NAME" != "proto-v$version" ]]; then
+            echo "Tag $GITHUB_REF_NAME does not match vllm-proto version $version" >&2
+            exit 1
+          fi
+      - name: Verify package
+        run: cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run
+
+  publish:
+    needs: [check, semver]
+    if: github.repository == 'vllm-project/vllm' && github.event_name == 'push'
+    runs-on: ubuntu-latest
+    environment: crates-io
+    permissions:
+      contents: read
+      id-token: write
+    steps:
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          ref: ${{ github.sha }}
+      - run: rustup show active-toolchain
+      - name: Verify package before authentication
+        run: cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run
+      - uses: rust-lang/crates-io-auth-action@c6f97d42243bad5fab37ca0427f495c86d5b1a18 # v1.0.5
+        id: auth
+      - name: Publish
+        env:
+          CARGO_REGISTRY_TOKEN: ${{ steps.auth.outputs.token }}
+        run: cargo publish --manifest-path rust/proto/Cargo.toml --locked
diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml
index 747831c7d4c8..3435148ea4e3 100644
--- a/.github/workflows/run-ci-command.yml
+++ b/.github/workflows/run-ci-command.yml
@@ -21,11 +21,17 @@ jobs:
       (github.event.comment.body == '/ci run' ||
       github.event.comment.body == '/ci run all' ||
       github.event.comment.body == '/ci run nightly' ||
+      github.event.comment.body == '/ci run --allow-stale' ||
+      github.event.comment.body == '/ci run all --allow-stale' ||
+      github.event.comment.body == '/ci run nightly --allow-stale' ||
       github.event.comment.body == '/ci retry' ||
       github.event.comment.body == '/ci cancel' ||
       github.event.comment.body == '/amd-ci run' ||
       github.event.comment.body == '/amd-ci run all' ||
       github.event.comment.body == '/amd-ci run nightly' ||
+      github.event.comment.body == '/amd-ci run --allow-stale' ||
+      github.event.comment.body == '/amd-ci run all --allow-stale' ||
+      github.event.comment.body == '/amd-ci run nightly --allow-stale' ||
       github.event.comment.body == '/amd-ci retry' ||
       github.event.comment.body == '/amd-ci cancel')
     runs-on: [self-hosted, linux, x64, vllm-runners]
diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py
index 6454f041ff06..530bd697d008 100644
--- a/.github/workflows/scripts/run_ci_command.py
+++ b/.github/workflows/scripts/run_ci_command.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import http.client
 import json
 import os
 import random
@@ -22,6 +23,7 @@
 COMMAND_RUN_AMD_CI_NIGHTLY = "/amd-ci run nightly"
 COMMAND_RETRY_AMD_FAILED = "/amd-ci retry"
 COMMAND_CANCEL_AMD_CI = "/amd-ci cancel"
+ALLOW_STALE_SUFFIX = " --allow-stale"
 RUN_CI_COMMAND_ENV = {
     COMMAND_RUN_CI: {},
     COMMAND_RUN_CI_ALL: {"RUN_ALL": "1"},
@@ -30,20 +32,25 @@
     COMMAND_RUN_AMD_CI_ALL: {"RUN_ALL": "1"},
     COMMAND_RUN_AMD_CI_NIGHTLY: {"RUN_ALL": "1", "NIGHTLY": "1"},
 }
+ALLOW_STALE_COMMANDS = frozenset(
+    f"{command}{ALLOW_STALE_SUFFIX}" for command in RUN_CI_COMMAND_ENV
+)
+RUN_CI_COMMAND_ENV.update(
+    {
+        command: RUN_CI_COMMAND_ENV[command.removesuffix(ALLOW_STALE_SUFFIX)]
+        for command in ALLOW_STALE_COMMANDS
+    }
+)
 UPSTREAM_CI_COMMANDS = frozenset(
     {
-        COMMAND_RUN_CI,
-        COMMAND_RUN_CI_ALL,
-        COMMAND_RUN_CI_NIGHTLY,
+        *(command for command in RUN_CI_COMMAND_ENV if command.startswith("/ci ")),
         COMMAND_RETRY_FAILED,
         COMMAND_CANCEL_CI,
     }
 )
 AMD_CI_COMMANDS = frozenset(
     {
-        COMMAND_RUN_AMD_CI,
-        COMMAND_RUN_AMD_CI_ALL,
-        COMMAND_RUN_AMD_CI_NIGHTLY,
+        *(command for command in RUN_CI_COMMAND_ENV if command.startswith("/amd-ci ")),
         COMMAND_RETRY_AMD_FAILED,
         COMMAND_CANCEL_AMD_CI,
     }
@@ -79,6 +86,10 @@ def __init__(self, status: int | None, message: str) -> None:
         self.status = status
 
 
+class CiPreparationError(RuntimeError):
+    pass
+
+
 def rate_limit_jitter() -> float:
     return random.uniform(1, 5)
 
@@ -137,6 +148,8 @@ def request(
                     None,
                     f"API request failed: {error.reason}",
                 ) from error
+            except (OSError, http.client.HTTPException) as error:
+                raise ApiError(None, f"API request failed: {error}") from error
 
         if not response_body:
             return None
@@ -234,6 +247,17 @@ def _paginate(self, path: str) -> list[dict[str, Any]]:
     def get_pr(self, number: int) -> dict[str, Any]:
         return self._request(self._repo_path(f"/pulls/{number}"))
 
+    def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int:
+        base_ref = urllib.parse.quote(f"refs/heads/{base_ref}", safe="")
+        head_sha = urllib.parse.quote(head_sha, safe="")
+        response = self._request(
+            self._repo_path(f"/compare/{base_ref}...{head_sha}?per_page=1")
+        )
+        behind = response.get("behind_by") if isinstance(response, dict) else None
+        if type(behind) is not int or behind < 0:
+            raise ApiError(None, "GitHub returned an invalid commits-behind count.")
+        return behind
+
     def list_pulls_for_commit(self, commit: str) -> list[dict[str, Any]]:
         commit = urllib.parse.quote(commit, safe="")
         response = self._request(self._repo_path(f"/commits/{commit}/pulls"))
@@ -747,6 +771,10 @@ def notify_authorized(
         (
             f"✅ @{author}, CI is now available for this PR.\n\n"
             "- `/ci run` starts upstream CI; `/amd-ci run` starts AMD CI only.\n"
+            "- Your branch must contain every commit currently on its upstream "
+            "target branch. Merge or rebase onto the latest target branch, then "
+            "rerun the command. Append `--allow-stale` to a run command to test "
+            "an outdated branch at your own risk.\n"
             "- `/ci retry` retries failed jobs in the CI build for the current "
             "PR head. If the current head has no CI build, it starts a new CI "
             "build for the current head containing only jobs that failed in "
@@ -794,6 +822,56 @@ def find_matching_pr(
     return find_matching_pr(github.list_pulls_for_commit(head_sha))
 
 
+def prepare_pr_for_ci(
+    github: GitHubClient,
+    pr: Mapping[str, Any],
+    command: str,
+) -> tuple[dict[str, Any], str]:
+    base_ref = pr["base"]["ref"]
+    no_action = "No new CI build was started."
+    try:
+        behind = github.get_commits_behind_base(base_ref, pr["head"]["sha"])
+        current_pr = github.get_pr(pr["number"])
+    except ApiError as error:
+        raise CiPreparationError(
+            f"Could not check the PR against upstream `{base_ref}`. "
+            f"{no_action} Comment `{command}` again. {error}"
+        ) from error
+    if (
+        current_pr["state"] != "open"
+        or current_pr["head"]["sha"] != pr["head"]["sha"]
+        or current_pr["base"]["ref"] != pr["base"]["ref"]
+    ):
+        raise CiPreparationError(
+            f"The PR changed while checking its branch. {no_action} "
+            f"Comment `{command}` again."
+        )
+    commits = "commit" if behind == 1 else "commits"
+    lag = f"This PR is {behind} {commits} behind upstream `{base_ref}`."
+    if command in ALLOW_STALE_COMMANDS:
+        warning = ""
+        if behind:
+            run_command = command.removesuffix(ALLOW_STALE_SUFFIX)
+            warning = (
+                f"⚠️ {lag} Running CI at your own risk because "
+                "`--allow-stale` was requested; outdated CI configuration may "
+                "cause failures. Before merging, merge or rebase onto the latest "
+                f"`{base_ref}`, then rerun `{run_command}` on the latest PR commit."
+            )
+            print(warning)
+        return current_pr, warning
+
+    if behind:
+        raise CiPreparationError(
+            f"{lag} Your branch must contain every commit currently on upstream "
+            f"`{base_ref}`. {no_action} "
+            f"Merge or rebase onto the latest `{base_ref}`, then rerun `{command}`."
+            " To test this branch at your own risk, "
+            f"use `{command}{ALLOW_STALE_SUFFIX}`."
+        )
+    return current_pr, ""
+
+
 def handle_run_ci(
     *,
     actor: str,
@@ -828,12 +906,7 @@ def handle_run_ci(
             f"{ci_name} is already running for this commit: {active_build['web_url']}"
         )
 
-    current_pr = github.get_pr(pr["number"])
-    if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]:
-        return (
-            "The PR head changed while processing the command. "
-            f"Comment `{command}` again."
-        )
+    current_pr, warning = prepare_pr_for_ci(github, pr, command)
 
     build = buildkite.create_build(
         create_build_payload(
@@ -846,6 +919,7 @@ def handle_run_ci(
     return (
         f"Triggered [Buildkite {ci_name} #{build['number']}]({build['web_url']}) "
         f"for commit `{current_pr['head']['sha'][:12]}`."
+        + (f"\n\n{warning}" if warning else "")
     )
 
 
@@ -1099,6 +1173,12 @@ def run(
             raise ValueError(f"Unsupported CI command: {command}")
         add_reaction_safely(github, comment_id, "rocket")
         github.add_comment(issue_number, f"✅ {message}")
+    except CiPreparationError as error:
+        add_reaction_safely(github, comment_id, "-1")
+        github.add_comment(
+            issue_number,
+            f"❌ {error}\n\n{command_comment_marker(comment_id)}",
+        )
     except Exception:
         add_reaction_safely(github, comment_id, "confused")
         raise
diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py
index ef93ea8559c8..9750cefaa394 100644
--- a/.github/workflows/scripts/test_run_ci_command.py
+++ b/.github/workflows/scripts/test_run_ci_command.py
@@ -22,8 +22,10 @@
     COMMAND_RUN_CI_ALL,
     COMMAND_RUN_CI_NIGHTLY,
     RETRY_STATES,
+    RUN_CI_COMMAND_ENV,
     ApiError,
     BuildkiteClient,
+    GitHubClient,
     HttpTransport,
     authorize,
     create_build_payload,
@@ -85,6 +87,7 @@ def __init__(
         comments: list[str] | None = None,
         pulls_for_commit: list[dict[str, Any]] | None = None,
         prs: dict[int, dict[str, Any]] | None = None,
+        behind: int = 0,
     ) -> None:
         self.comments = comments or []
         self.permission = permission
@@ -95,6 +98,12 @@ def __init__(
         self.reviews = reviews or []
         self.pulls_for_commit = pulls_for_commit or []
         self.prs = prs or {self.pr["number"]: self.pr}
+        self.behind = behind
+        self.lag_queries: list[tuple[str, str]] = []
+
+    def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int:
+        self.lag_queries.append((base_ref, head_sha))
+        return self.behind
 
     def get_pr(self, number: int) -> dict[str, Any]:
         try:
@@ -322,6 +331,11 @@ def test_only_exact_ci_commands_are_accepted(self) -> None:
         for command in commands:
             with self.subTest(command=command):
                 self.assertEqual(parse_command(command), command)
+                override = f"{command} --allow-stale"
+                if " run" in command:
+                    self.assertEqual(parse_command(override), override)
+                else:
+                    self.assertIsNone(parse_command(override))
 
         self.assertIsNone(parse_command("/ci run please"))
         self.assertIsNone(parse_command("/ci run all please"))
@@ -331,6 +345,8 @@ def test_only_exact_ci_commands_are_accepted(self) -> None:
         self.assertIsNone(parse_command("/amd-ci retry "))
         self.assertIsNone(parse_command("/AMD-CI run"))
         self.assertIsNone(parse_command("/amdci run"))
+        self.assertIsNone(parse_command("/ci run --allow-stale all"))
+        self.assertIsNone(parse_command("/amd-ci run --allow-stale please"))
 
     def test_commands_select_only_their_configured_pipeline(self) -> None:
         cases = (
@@ -530,6 +546,177 @@ def test_write_reviewer_runs_ci_without_delegation(self) -> None:
         self.assertTrue(github.comments[0].startswith("✅ "))
         self.assertIn("Buildkite CI #123", github.comments[0])
 
+    def test_run_variants_require_current_base_without_changing_commit(self) -> None:
+        for command in RUN_CI_COMMAND_ENV:
+            if command.endswith(" --allow-stale"):
+                continue
+            for behind in (0, 1, 5):
+                with self.subTest(command=command, behind=behind):
+                    github = FakeGitHub(
+                        behind=behind, pr=make_pr(base={"ref": "release/1.0"})
+                    )
+                    buildkite = FakeBuildkite([[], []])
+                    run(make_event(command), github, buildkite)
+
+                    self.assertEqual(
+                        github.lag_queries,
+                        [("release/1.0", github.pr["head"]["sha"])],
+                    )
+                    if behind == 0:
+                        self.assertEqual(
+                            buildkite.created_builds[0]["commit"],
+                            github.pr["head"]["sha"],
+                        )
+                        self.assertIn(
+                            "](https://buildkite.example/builds/123)",
+                            github.comments[0],
+                        )
+                    else:
+                        self.assertEqual(buildkite.created_builds, [])
+                        self.assertTrue(github.comments[0].startswith("❌ "))
+                        self.assertIn(str(behind), github.comments[0])
+                        self.assertIn("upstream `release/1.0`", github.comments[0])
+                        self.assertIn("every commit", github.comments[0])
+                        self.assertIn(
+                            "No new CI build was started.", github.comments[0]
+                        )
+                        self.assertIn(command, github.comments[0])
+
+    def test_allow_stale_preserves_author_permissions_and_run_configuration(
+        self,
+    ) -> None:
+        cases = (
+            (COMMAND_RUN_CI, "ci", {}),
+            (COMMAND_RUN_CI_ALL, "ci", {"RUN_ALL": "1"}),
+            (COMMAND_RUN_CI_NIGHTLY, "ci", {"RUN_ALL": "1", "NIGHTLY": "1"}),
+            (COMMAND_RUN_AMD_CI, "amd-ci", {}),
+            (COMMAND_RUN_AMD_CI_ALL, "amd-ci", {"RUN_ALL": "1"}),
+            (
+                COMMAND_RUN_AMD_CI_NIGHTLY,
+                "amd-ci",
+                {"RUN_ALL": "1", "NIGHTLY": "1"},
+            ),
+        )
+        for original_command, pipeline, expected_env in cases:
+            command = f"{original_command} --allow-stale"
+            for ready, behind in ((False, 100), (True, 0), (True, 100)):
+                with self.subTest(command=command, ready=ready, behind=behind):
+                    github = FakeGitHub(
+                        permission="read",
+                        behind=behind,
+                        pr=make_pr(
+                            base={"ref": "feature/stack"},
+                            labels=[{"name": "ready"}] if ready else [],
+                        ),
+                    )
+                    buildkite = FakeBuildkite([[], []])
+
+                    run(make_event(command, "author"), github, buildkite)
+
+                    self.assertEqual(pipeline_for_command(command), pipeline)
+                    if not ready:
+                        self.assertEqual(buildkite.created_builds, [])
+                        self.assertEqual(buildkite.list_calls, [])
+                        self.assertEqual(github.lag_queries, [])
+                        self.assertTrue(github.comments[0].startswith("❌ "))
+                        continue
+                    self.assertEqual(len(buildkite.created_builds), 1)
+                    payload = buildkite.created_builds[0]
+                    self.assertEqual(payload["commit"], github.pr["head"]["sha"])
+                    self.assertEqual(
+                        payload["pull_request_base_branch"], "feature/stack"
+                    )
+                    self.assertEqual(payload["message"], f"PR #42 {command} by @author")
+                    self.assertEqual(
+                        payload["env"],
+                        {
+                            "VLLM_CI_GITHUB_COMMENT_ID": "99",
+                            "VLLM_CI_TRIGGERED_BY": "author",
+                            **expected_env,
+                        },
+                    )
+                    self.assertEqual(
+                        github.lag_queries,
+                        [("feature/stack", github.pr["head"]["sha"])],
+                    )
+                    self.assertIn(
+                        "Buildkite AMD CI #123"
+                        if pipeline == "amd-ci"
+                        else "Buildkite CI #123",
+                        github.comments[0],
+                    )
+                    if behind:
+                        self.assertIn("100 commits behind", github.comments[0])
+                        self.assertIn("upstream `feature/stack`", github.comments[0])
+                        self.assertIn("Before merging", github.comments[0])
+                        self.assertIn("latest PR commit", github.comments[0])
+                        self.assertIn(original_command, github.comments[0])
+
+    def test_base_lag_query_qualifies_branch_refs_and_validates_count(self) -> None:
+        for base_ref, encoded_ref in (
+            ("main", "main"),
+            ("release/1.0", "release%2F1.0"),
+            ("feature/stack", "feature%2Fstack"),
+        ):
+            for count in (0, 5, None, -1, True, "5"):
+                with self.subTest(base_ref=base_ref, count=count):
+                    transport = FakeTransport(
+                        {} if count is None else {"behind_by": count}
+                    )
+                    github = GitHubClient("token", "vllm-project/vllm", transport)
+                    if type(count) is int and count >= 0:
+                        self.assertEqual(
+                            github.get_commits_behind_base(base_ref, "head-sha"), count
+                        )
+                    else:
+                        with self.assertRaises(ApiError):
+                            github.get_commits_behind_base(base_ref, "head-sha")
+                    self.assertEqual(
+                        transport.calls[0]["url"],
+                        "https://api.github.com/repos/vllm-project/vllm/compare/"
+                        f"refs%2Fheads%2F{encoded_ref}...head-sha?per_page=1",
+                    )
+                    self.assertEqual(transport.calls[0]["method"], "GET")
+
+    def test_unknown_lag_or_changed_pr_prevents_runs(self) -> None:
+        for command in (
+            COMMAND_RUN_CI,
+            COMMAND_RUN_AMD_CI,
+            f"{COMMAND_RUN_CI} --allow-stale",
+            f"{COMMAND_RUN_AMD_CI} --allow-stale",
+        ):
+            for fault in ("api", "head", "closed", "base"):
+                with self.subTest(command=command, fault=fault):
+                    github = FakeGitHub()
+                    buildkite = FakeBuildkite([[], []])
+
+                    def check(
+                        base_ref: str,
+                        sha: str,
+                        fault: str = fault,
+                        github: FakeGitHub = github,
+                    ) -> int:
+                        if fault == "api":
+                            raise ApiError(503, "Unavailable")
+                        changed = make_pr()
+                        if fault == "head":
+                            changed["head"]["sha"] = "another-head"
+                        elif fault == "closed":
+                            changed["state"] = "closed"
+                        else:
+                            changed["base"]["ref"] = "another-base"
+                        github.prs[42] = changed
+                        return 0
+
+                    with patch.object(
+                        github, "get_commits_behind_base", side_effect=check
+                    ):
+                        run(make_event(command), github, buildkite)
+                    self.assertEqual(buildkite.created_builds, [])
+                    self.assertEqual(buildkite.retry_calls, [])
+                    self.assertTrue(github.comments[0].startswith("❌ "))
+                    self.assertNotIn("Triggered", github.comments[0])
+
     def test_amd_run_ignores_blocked_builds(self) -> None:
         for metadata in ({}, {"github-comment-id": "98"}):
             with self.subTest(metadata=metadata):
@@ -627,6 +814,7 @@ def test_unapproved_authors_are_denied_without_buildkite(self) -> None:
         run(make_event(COMMAND_RUN_CI, "author"), github, buildkite)
 
         self.assertEqual(buildkite.list_calls, [])
+        self.assertEqual(github.lag_queries, [])
         self.assertEqual(github.reactions, ["eyes"])
         self.assertTrue(github.comments[0].startswith("❌ "))
         self.assertIn("approve the PR", github.comments[0])
@@ -848,6 +1036,7 @@ def test_notification_skips_draft_prs(self) -> None:
 
     def test_ci_retry_retries_failed_jobs_while_build_is_running(self) -> None:
         github = FakeGitHub(
+            behind=100,
             permission="read",
             pr=make_pr(labels=[{"name": "ready"}]),
         )
@@ -867,10 +1056,12 @@ def test_ci_retry_retries_failed_jobs_while_build_is_running(self) -> None:
         run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite)
 
         self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)])
+        self.assertEqual(github.lag_queries, [])
         self.assertIn("Queued 3 failed job", github.comments[0])
 
     def test_amd_ci_retry_retries_only_the_current_head_build(self) -> None:
         github = FakeGitHub(
+            behind=100,
             permission="read",
             pr=make_pr(labels=[{"name": "ready"}]),
         )
@@ -891,6 +1082,7 @@ def test_amd_ci_retry_retries_only_the_current_head_build(self) -> None:
         run(make_event(COMMAND_RETRY_AMD_FAILED, "author"), github, buildkite)
 
         self.assertEqual(buildkite.retry_calls, [(321, RETRY_STATES)])
+        self.assertEqual(github.lag_queries, [])
         self.assertIn("Buildkite AMD CI #321", github.comments[0])
 
     def test_amd_ci_retry_requires_a_build_for_the_current_head(self) -> None:
@@ -912,6 +1104,7 @@ def test_amd_ci_retry_requires_a_build_for_the_current_head(self) -> None:
 
     def test_ci_retry_creates_filtered_build_for_new_head(self) -> None:
         github = FakeGitHub(
+            behind=100,
             permission="read",
             pr=make_pr(labels=[{"name": "ready"}]),
         )
@@ -953,6 +1146,7 @@ def test_ci_retry_creates_filtered_build_for_new_head(self) -> None:
         self.assertEqual(len(buildkite.created_builds), 1)
         payload = buildkite.created_builds[0]
         self.assertEqual(payload["commit"], "0123456789abcdef")
+        self.assertEqual(github.lag_queries, [])
         self.assertEqual(payload["message"], "PR #42 /ci retry by @author")
         self.assertEqual(
             json.loads(payload["env"]["VLLM_CI_ONLY_STEP_KEYS"]),
@@ -1104,6 +1298,7 @@ def test_ci_cancel_cancels_active_builds_for_pr_branch(self) -> None:
         run(make_event(COMMAND_CANCEL_CI), github, buildkite)
 
         self.assertEqual(buildkite.cancel_calls, [123, 124, 126])
+        self.assertEqual(github.lag_queries, [])
         self.assertEqual(
             [request["branch"] for request in buildkite.list_requests],
             ["contributor:feature", "feature"],
@@ -1152,6 +1347,7 @@ def test_amd_ci_cancel_handles_command_and_fork_webhook_branches(self) -> None:
         run(make_event(COMMAND_CANCEL_AMD_CI), github, buildkite)
 
         self.assertEqual(buildkite.cancel_calls, [322, 321])
+        self.assertEqual(github.lag_queries, [])
         self.assertEqual(
             [request["branch"] for request in buildkite.list_requests],
             ["contributor:feature", "feature"],
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 0954c53cb290..4823f788dc3a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -177,6 +177,21 @@ repos:
     files: ^\.github/mergify\.yml$
     pass_filenames: false
     additional_dependencies: [PyYAML, regex]
+  - id: check-buildkite-step-keys
+    name: Check every Buildkite step declares an explicit key
+    language: python
+    entry: python tools/pre_commit/check_buildkite_step_keys.py
+    files: ^\.buildkite/.*\.ya?ml$
+    pass_filenames: false
+    additional_dependencies: [PyYAML]
+  - id: check-test-tethering
+    name: Check new tests are wired into Buildkite CI
+    language: python
+    entry: python tools/pre_commit/check_test_tethering.py
+    files: ^(tests/.*\.py|\.buildkite/(test_areas|hardware_tests|intel_jobs)/.*\.yaml|\.buildkite/(test-amd|ci_config[^/]*)\.yaml|tools/pre_commit/test_tethering_allowlist\.txt)$
+    additional_dependencies: [PyYAML, regex]
+    verbose: true # full-scan `note:` lines are advisory, so the hook passes
+    require_serial: true # one full scan, not one per pre-commit file batch
   - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.10
     entry: python tools/pre_commit/mypy.py "3.10"
@@ -285,7 +300,7 @@ repos:
     pass_filenames: false
     require_serial: true
     stages: [pre-commit] # Only run locally as Buildkite will cover this
-    files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$
+    files: ^rust/(Cargo\.toml|proto/Cargo\.toml|src/.*/Cargo\.toml)$
   - id: rust-cargo-sort
     name: Rust - Sort Cargo manifest sections
     entry: tools/pre_commit/rust-check.sh sort --workspace
@@ -293,7 +308,7 @@ repos:
     pass_filenames: false
     require_serial: true
     stages: [pre-commit] # Only run locally as Buildkite will cover this
-    files: ^rust/(Cargo\.toml|src/.*/Cargo\.toml)$
+    files: ^rust/(Cargo\.toml|proto/Cargo\.toml|src/.*/Cargo\.toml)$
   - id: rust-cargo-fmt
     name: Rust - Format code
     entry: tools/pre_commit/rust-check.sh fmt
diff --git a/CMakeLists.txt b/CMakeLists.txt
index a7d40ae58603..58921b93690d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -460,6 +460,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      DEFINED CMAKE_CUDA_COMPILER_VERSION AND
      CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0)
 
+    list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/all_reduce_mhc.cu")
+
     if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
       cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
         "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
@@ -476,6 +478,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
+      "csrc/libtorch_stable/hisparse_kernels.cu"
       "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu"
       "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu"
       "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu")
@@ -774,22 +777,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
                    "(requires SM90+ and CUDA >= 12.0).")
   endif()
 
-  # Only build AllSpark kernels if we are building for at least some compatible archs.
-  cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
-  if (ALLSPARK_ARCHS)
-    set(ALLSPARK_SRCS
-       "csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu"
-       "csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${ALLSPARK_SRCS}"
-      CUDA_ARCHS "${ALLSPARK_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${ALLSPARK_SRCS}")
-    message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
-  else()
-    message(STATUS "Not building AllSpark kernels as no compatible archs found"
-                   " in CUDA target architectures")
-  endif()
-
   #
   # CUTLASS scaled_mm kernels (moved from _C to _C_stable_libtorch)
   #
@@ -1571,6 +1558,7 @@ endif()
 # For CUDA we also build and ship some external projects.
 if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
+    include(cmake/external_projects/deepselect.cmake)
     include(cmake/external_projects/fmha_sm100.cmake)
     include(cmake/external_projects/flashmla.cmake)
     include(cmake/external_projects/flashkda.cmake)
diff --git a/benchmarks/attention_benchmarks/batch_spec.py b/benchmarks/attention_benchmarks/batch_spec.py
index 9f15f1d8096e..85c689b48944 100644
--- a/benchmarks/attention_benchmarks/batch_spec.py
+++ b/benchmarks/attention_benchmarks/batch_spec.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Simplified batch specification grammar for attention benchmarks.
+"""Simplified batch specification grammar for attention benchmarks.
 
 Grammar (underscore-separated segments):
   Format: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
@@ -24,6 +23,7 @@
   q4s1k            -> [(4, 1024)]              # 4-token extend (spec decode)
   2q1k_32q1s1k     -> [(1024, 1024)] * 2 + [(1, 1024)] * 32  # Mixed batch
   16q4s1k          -> [(4, 1024)] * 16         # 16 spec decode requests
+
 """
 
 from collections import Counter
@@ -71,8 +71,7 @@ def _parse_size(size_str: str, k_suffix: str) -> int:
 
 
 def parse_batch_spec(spec: str) -> list[BatchRequest]:
-    """
-    Parse batch specification string into list of BatchRequest objects.
+    """Parse batch specification string into list of BatchRequest objects.
 
     Grammar: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
 
@@ -84,6 +83,7 @@ def parse_batch_spec(spec: str) -> list[BatchRequest]:
 
     Raises:
         ValueError: If spec format is invalid
+
     """
     requests = []
 
@@ -103,8 +103,7 @@ def parse_batch_spec(spec: str) -> list[BatchRequest]:
 
 
 def format_batch_spec(requests: list[BatchRequest]) -> str:
-    """
-    Format list of BatchRequest into human-readable string.
+    """Format list of BatchRequest into human-readable string.
 
     Groups requests by type and provides counts and sizes.
 
@@ -113,6 +112,7 @@ def format_batch_spec(requests: list[BatchRequest]) -> str:
 
     Returns:
         Formatted string describing the batch
+
     """
     kinds = {
         "prefill": [],
@@ -157,8 +157,7 @@ def format_batch_spec(requests: list[BatchRequest]) -> str:
 
 
 def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:
-    """
-    Reorder requests for FlashInfer: decode first, then prefill.
+    """Reorder requests for FlashInfer: decode first, then prefill.
 
     FlashInfer expects decode requests before prefill requests for
     optimal performance.
@@ -168,6 +167,7 @@ def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:
 
     Returns:
         Reordered list with decode requests first
+
     """
     decodes = [r for r in requests if r.is_decode]
     non_decodes = [r for r in requests if not r.is_decode]
@@ -177,14 +177,14 @@ def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:
 def split_by_type(
     requests: list[BatchRequest],
 ) -> dict[str, list[BatchRequest]]:
-    """
-    Split requests by type for analysis.
+    """Split requests by type for analysis.
 
     Args:
         requests: List of BatchRequest
 
     Returns:
         Dict with keys: 'decode', 'prefill', 'extend'
+
     """
     result = {
         "decode": [],
@@ -204,14 +204,14 @@ def split_by_type(
 
 
 def get_batch_stats(requests: list[BatchRequest]) -> dict:
-    """
-    Compute statistics about a batch.
+    """Compute statistics about a batch.
 
     Args:
         requests: List of BatchRequest
 
     Returns:
         Dict with batch statistics
+
     """
     by_type = split_by_type(requests)
 
@@ -232,8 +232,7 @@ def get_batch_stats(requests: list[BatchRequest]) -> dict:
 
 
 def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
-    """
-    Classify a batch spec into a type string.
+    """Classify a batch spec into a type string.
 
     Args:
         batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k")
@@ -241,6 +240,7 @@ def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
 
     Returns:
         Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)"
+
     """
     requests = parse_batch_spec(batch_spec)
 
diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py
index 9c1ea1add70d..943133252be0 100644
--- a/benchmarks/attention_benchmarks/benchmark.py
+++ b/benchmarks/attention_benchmarks/benchmark.py
@@ -2,8 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Universal vLLM Attention Benchmark
+"""Universal vLLM Attention Benchmark.
 
 Benchmark any attention backend with the extended grammar.
 Supports standard attention (Flash/Triton/FlashInfer) and MLA backends.
@@ -23,6 +22,7 @@
 
     # Parameter sweep (YAML config - recommended)
     python benchmark.py --config configs/cutlass_numsplits.yaml
+
 """
 
 import argparse
@@ -84,8 +84,7 @@ def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
 
 
 def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
-    """
-    Run a single benchmark with proper backend selection.
+    """Run a single benchmark with proper backend selection.
 
     Args:
         config: BenchmarkConfig with backend, batch_spec, and model params
@@ -93,6 +92,7 @@ def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
 
     Returns:
         BenchmarkResult (may have error field set on failure)
+
     """
     try:
         if is_mla_backend(config.backend):
@@ -119,8 +119,7 @@ def run_model_parameter_sweep(
     sweep: ModelParameterSweep,
     console: Console,
 ) -> list[BenchmarkResult]:
-    """
-    Run model parameter sweep for given backends and batch specs.
+    """Run model parameter sweep for given backends and batch specs.
 
     Args:
         backends: List of backend names
@@ -131,6 +130,7 @@ def run_model_parameter_sweep(
 
     Returns:
         List of BenchmarkResult objects
+
     """
     all_results = []
 
@@ -304,8 +304,7 @@ def run_parameter_sweep(
     sweep: ParameterSweep,
     console: Console,
 ) -> list[BenchmarkResult]:
-    """
-    Run parameter sweep for given backends and batch specs.
+    """Run parameter sweep for given backends and batch specs.
 
     Args:
         backends: List of backend names
@@ -316,6 +315,7 @@ def run_parameter_sweep(
 
     Returns:
         List of BenchmarkResult objects
+
     """
     all_results = []
 
@@ -395,8 +395,7 @@ def load_config_from_yaml(config_path: str) -> dict:
 
 
 def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]:
-    """
-    Generate batch specs from range specifications.
+    """Generate batch specs from range specifications.
 
     Args:
         ranges: List of range specifications, each containing:
@@ -420,6 +419,7 @@ def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]:
             }
         ]
         Returns: ["q1kv1k", "q2kv1k", ..., "q16kv1k"]
+
     """
     all_specs = []
 
diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py
index 45858f9148b3..b6fd26c80344 100644
--- a/benchmarks/attention_benchmarks/common.py
+++ b/benchmarks/attention_benchmarks/common.py
@@ -21,8 +21,7 @@
 
 
 def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
-    """
-    Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len).
+    """Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len).
 
     This ensures results are sorted by batch size first, then query length,
     then sequence length, rather than alphabetically.
@@ -354,13 +353,13 @@ def print_table(
         backends: list[str],
         compare_to_fastest: bool = True,
     ):
-        """
-        Print results as a rich table.
+        """Print results as a rich table.
 
         Args:
             results: List of BenchmarkResult
             backends: List of backend names being compared
             compare_to_fastest: Show percentage comparison to fastest
+
         """
         # Group by batch spec, preserving first-occurrence order
         by_spec = {}
@@ -489,14 +488,14 @@ def save_json(self, results: list[BenchmarkResult], path: str):
 
 
 def setup_mla_dims(model_name: str = "deepseek-v3") -> dict:
-    """
-    Get MLA dimensions for known models.
+    """Get MLA dimensions for known models.
 
     Args:
         model_name: Model identifier
 
     Returns:
         Dict with MLA dimension configuration
+
     """
     configs = {
         "deepseek-v2": {
@@ -542,8 +541,7 @@ def get_attention_scale(head_dim: int) -> float:
 
 
 def is_mla_backend(backend: str) -> bool:
-    """
-    Check if backend is an MLA backend using the AttentionBackendEnum.
+    """Check if backend is an MLA backend using the AttentionBackendEnum.
 
     Args:
         backend: Backend name matching AttentionBackendEnum exactly
@@ -551,6 +549,7 @@ def is_mla_backend(backend: str) -> bool:
 
     Returns:
         True if the backend is an MLA backend, False otherwise
+
     """
     from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py
index 790b2f89442b..04dcd73719a3 100644
--- a/benchmarks/attention_benchmarks/mla_runner.py
+++ b/benchmarks/attention_benchmarks/mla_runner.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-MLA benchmark runner - shared utilities for MLA benchmarks.
+"""MLA benchmark runner - shared utilities for MLA benchmarks.
 
 This module provides helpers for running MLA backends without
 needing full VllmConfig integration.
@@ -93,8 +92,7 @@ def _ensure_single_rank_model_parallel() -> None:
 
 
 def _add_mock_methods_to_model_config(model_config: ModelConfig) -> None:
-    """
-    Add mock methods for layer-specific queries to ModelConfig.
+    """Add mock methods for layer-specific queries to ModelConfig.
 
     These methods are needed by metadata builders but aren't normally
     present on ModelConfig when used in benchmark contexts.
@@ -125,14 +123,15 @@ def create_minimal_vllm_config(
     kv_cache_dtype: str = "auto",
     sparse_mla_force_mqa: bool = False,
 ) -> VllmConfig:
-    """
-    Create minimal VllmConfig for MLA benchmarks.
+    """Create minimal VllmConfig for MLA benchmarks.
 
     Args:
         model_name: Model name (deepseek-v2, deepseek-v3, etc.) - used if mla_dims not
                     provided
         block_size: KV cache block size
         max_num_seqs: Maximum number of sequences
+        max_num_batched_tokens: Maximum number of batched tokens
+        max_model_len: Maximum model context length
         mla_dims: Optional custom MLA dimensions dict. If not provided, uses
                   setup_mla_dims(model_name)
         index_topk: Optional topk value for sparse MLA backends. If provided,
@@ -140,11 +139,13 @@ def create_minimal_vllm_config(
         prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",
                         "trtllm"). Configures the attention config to force
                         the specified prefill backend.
+        kv_cache_dtype: KV cache dtype, e.g. "auto" or "fp8"
         sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
                     forward_mqa (even prefill tokens).
 
     Returns:
         VllmConfig for benchmarking
+
     """
     # Get MLA dimensions - use provided or load from model name
     if mla_dims is None:
@@ -320,8 +321,7 @@ def get_prefill_backend_config(prefill_backend: str) -> dict:
 
 
 def _get_backend_config(backend: str) -> dict:
-    """
-    Get backend configuration from AttentionBackendEnum.
+    """Get backend configuration from AttentionBackendEnum.
 
     Uses the registry to get the backend class and extract configuration
     from its methods (get_impl_cls, get_builder_cls, is_sparse, etc.).
@@ -332,6 +332,7 @@ def _get_backend_config(backend: str) -> dict:
 
     Returns:
         Dict with backend configuration
+
     """
     from vllm.v1.attention.backend import MultipleOf
     from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -381,8 +382,7 @@ def _build_attention_metadata(
     device: torch.device,
     builder_instance,
 ) -> tuple:
-    """
-    Build attention metadata from batch requests.
+    """Build attention metadata from batch requests.
 
     Args:
         requests: List of BatchRequest objects
@@ -392,6 +392,7 @@ def _build_attention_metadata(
 
     Returns:
         Tuple of (metadata, kv_cache_num_blocks)
+
     """
     q_lens = [r.q_len for r in requests]
     kv_lens = [r.kv_len for r in requests]
@@ -409,10 +410,6 @@ def _build_attention_metadata(
     seq_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32)
     seq_lens_gpu = seq_lens_cpu.to(device)
 
-    # Build num_computed_tokens (context length for each request)
-    context_lens = [kv_len - q_len for q_len, kv_len in zip(q_lens, kv_lens)]
-    num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
-
     # Build block table
     num_blocks_per_req = [(kv + block_size - 1) // block_size for kv in kv_lens]
     max_num_blocks = max(num_blocks_per_req)
@@ -454,8 +451,6 @@ def _build_attention_metadata(
         query_start_loc_cpu=q_start_cpu,
         seq_lens=seq_lens_gpu,
         seq_lens_cpu_upper_bound=seq_lens_cpu,
-        _seq_lens_cpu=seq_lens_cpu,
-        _num_computed_tokens_cpu=num_computed_tokens_cpu,
         slot_mapping=slot_mapping,
         block_table_tensor=block_table_gpu,
         dcp_local_seq_lens=None,
@@ -478,8 +473,7 @@ def _create_input_tensors(
     device: torch.device,
     dtype: torch.dtype,
 ):
-    """
-    Create input tensors for both decode and prefill modes.
+    """Create input tensors for both decode and prefill modes.
 
     MLA requires different tensor formats for decode vs prefill:
     - Decode: Uses kv_lora_rank (512) dimension
@@ -496,6 +490,7 @@ def _create_input_tensors(
         Tuple of (decode_inputs, prefill_inputs)
         - decode_inputs: Query tensor(s) for decode mode
         - prefill_inputs: Dict with 'q', 'k_c_normed', 'k_pe', 'k_scale' for prefill
+
     """
     if query_format == "tuple":
         # Decode mode format: (q_nope, q_pe) where q_nope has kv_lora_rank dim
@@ -589,8 +584,7 @@ def _create_backend_impl(
     index_topk: int | None = None,
     kv_cache_dtype: str = "auto",
 ):
-    """
-    Create backend implementation instance.
+    """Create backend implementation instance.
 
     Args:
         backend_cfg: Backend configuration dict from _get_backend_config()
@@ -599,9 +593,11 @@ def _create_backend_impl(
         device: Target device
         max_num_tokens: Maximum number of tokens for sparse indexer buffer
         index_topk: Topk value for sparse MLA backends
+        kv_cache_dtype: KV cache dtype, e.g. "auto" or "fp8"
 
     Returns:
         Tuple of (impl, layer, builder_instance, indexer)
+
     """
     # Get classes from backend config (already resolved by _get_backend_config)
     impl_class = backend_cfg["impl_class"]
@@ -707,14 +703,14 @@ def _create_backend_impl(
 
 
 def _extract_mla_dims_from_config(config) -> dict | None:
-    """
-    Extract MLA dimensions from BenchmarkConfig if all required fields are present.
+    """Extract MLA dimensions from BenchmarkConfig if all required fields are present.
 
     Args:
         config: BenchmarkConfig instance
 
     Returns:
         Dict with MLA dimensions if all fields are provided, None otherwise
+
     """
     # Check if all MLA-specific fields are provided
     if all(
@@ -768,8 +764,7 @@ def _run_single_benchmark(
     output_scale: float | None = None,
     fuse_quant_op: bool = False,
 ) -> BenchmarkResult:
-    """
-    Run a single benchmark iteration.
+    """Run a single benchmark iteration.
 
     Args:
         config: BenchmarkConfig instance
@@ -780,6 +775,7 @@ def _run_single_benchmark(
         mla_dims: MLA dimension configuration
         device: Target device
         indexer: Optional MockIndexer for sparse backends
+        kv_cache_dtype: KV cache dtype. None keeps the backend default.
         output_scale: Static per-tensor FP8 scale for prefill output. None
             keeps the plain bf16 output (no quantization).
         fuse_quant_op: With output_scale set, True lets the prefill kernel write
@@ -788,6 +784,7 @@ def _run_single_benchmark(
 
     Returns:
         BenchmarkResult with timing statistics
+
     """
     # Parse batch spec
     requests = parse_batch_spec(config.batch_spec)
@@ -1050,8 +1047,7 @@ def _run_mla_benchmark_batched(
     output_scale: float | None = None,
     fuse_quant_op: bool = False,
 ) -> list[BenchmarkResult]:
-    """
-    Unified batched MLA benchmark runner for all backends.
+    """Unified batched MLA benchmark runner for all backends.
 
     Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
                flashinfer_mla_sparse, flashmla_sparse
@@ -1067,11 +1063,16 @@ def _run_mla_benchmark_batched(
         index_topk: Topk value for sparse MLA backends (default 2048)
         prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
             When set, forces the specified FlashAttention version for prefill.
+        output_scale: Static per-tensor FP8 scale for prefill output. None
+            keeps the plain bf16 output (no quantization).
+        fuse_quant_op: With output_scale set, True lets the prefill kernel write
+            FP8 directly; False runs bf16 attention then a standalone quant.
         sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
             forward_mqa (even prefill tokens).
 
     Returns:
         List of BenchmarkResult objects
+
     """
     if not configs_with_params:
         return []
@@ -1283,8 +1284,7 @@ def run_mla_benchmark(
     output_scale: float | None = None,
     fuse_quant_op: bool = False,
 ) -> BenchmarkResult | list[BenchmarkResult]:
-    """
-    Unified MLA benchmark runner for all backends.
+    """Unified MLA benchmark runner for all backends.
 
     Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
                flashinfer_mla_sparse, flashmla_sparse
@@ -1309,6 +1309,7 @@ def run_mla_benchmark(
 
     Returns:
         BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
+
     """
     # Normalize to batched mode: (config, threshold, num_splits)
     if isinstance(config, list):
diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py
index 01fd558043b7..a1b7d045050e 100644
--- a/benchmarks/attention_benchmarks/runner.py
+++ b/benchmarks/attention_benchmarks/runner.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Standard attention benchmark runner - shared utilities for non-MLA benchmarks.
+"""Standard attention benchmark runner - shared utilities for non-MLA benchmarks.
 
 This module provides helpers for running standard attention backends
 (FlashAttention, Triton, FlashInfer) with real vLLM integration.
@@ -57,8 +56,7 @@
 
 
 def _get_backend_config(backend: str) -> dict:
-    """
-    Get backend configuration from AttentionBackendEnum.
+    """Get backend configuration from AttentionBackendEnum.
 
     Args:
         backend: Backend name matching AttentionBackendEnum exactly
@@ -66,6 +64,7 @@ def _get_backend_config(backend: str) -> dict:
 
     Returns:
         Dict with backend_class
+
     """
     from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
@@ -409,6 +408,7 @@ def _run_single_benchmark(
     Returns:
         (timing_stats, mem_stats) where timing_stats is a dict with
         mean/std/min/max in seconds per layer.
+
     """
     total_q = q_list[0].shape[0]
     v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
@@ -458,8 +458,7 @@ def benchmark_fn():
 
 
 def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
-    """
-    Run standard attention benchmark with real kernels.
+    """Run standard attention benchmark with real kernels.
 
     Supports: FLASH_ATTN, TRITON_ATTN, FLASHINFER
 
@@ -468,6 +467,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
 
     Returns:
         BenchmarkResult with timing and memory statistics
+
     """
     device = torch.device(config.device)
     torch.accelerator.set_device_index(device)
diff --git a/benchmarks/backend_request_func.py b/benchmarks/backend_request_func.py
index 6349095ad729..b8e5f5be4c13 100644
--- a/benchmarks/backend_request_func.py
+++ b/benchmarks/backend_request_func.py
@@ -105,7 +105,7 @@ async def async_request_tgi(
                         timestamp = time.perf_counter()
                         # First token
                         if ttft == 0.0:
-                            ttft = time.perf_counter() - st
+                            ttft = timestamp - st
                             output.ttft = ttft
 
                         # Decoding phase
@@ -330,8 +330,7 @@ async def async_request_openai_completions(
                                 # First token
                                 if not first_chunk_received:
                                     first_chunk_received = True
-                                    ttft = time.perf_counter() - st
-                                    output.ttft = ttft
+                                    output.ttft = timestamp - st
 
                                 # Decoding phase
                                 else:
@@ -416,7 +415,7 @@ async def async_request_openai_chat_completions(
         output.prompt_len = request_func_input.prompt_len
 
         generated_text = ""
-        ttft = 0.0
+        first_chunk_received = False
         st = time.perf_counter()
         most_recent_timestamp = st
         try:
@@ -443,22 +442,30 @@ async def async_request_openai_chat_completions(
                             if choices := data.get("choices"):
                                 content = choices[0]["delta"].get("content")
                                 # First token
-                                if ttft == 0.0:
-                                    ttft = timestamp - st
-                                    output.ttft = ttft
+                                if not first_chunk_received:
+                                    first_chunk_received = True
+                                    output.ttft = timestamp - st
 
                                 # Decoding phase
                                 else:
                                     output.itl.append(timestamp - most_recent_timestamp)
 
                                 generated_text += content or ""
+                                # Only token chunks advance the request end;
+                                # the trailing usage chunk carries no token.
+                                most_recent_timestamp = timestamp
                             elif usage := data.get("usage"):
                                 output.output_tokens = usage.get("completion_tokens")
 
-                            most_recent_timestamp = timestamp
-
                     output.generated_text = generated_text
-                    output.success = True
+                    if first_chunk_received:
+                        output.success = True
+                    else:
+                        output.success = False
+                        output.error = (
+                            "Never received a valid chunk to calculate TTFT."
+                            "This response will be marked as failed!"
+                        )
                     output.latency = most_recent_timestamp - st
                 else:
                     output.error = response.reason or ""
@@ -530,7 +537,7 @@ def to_bytes(y, sr):
             output.prompt_len = request_func_input.prompt_len
 
             generated_text = ""
-            ttft = 0.0
+            first_chunk_received = False
             st = time.perf_counter()
             most_recent_timestamp = st
             try:
@@ -551,9 +558,9 @@ def to_bytes(y, sr):
                                 if choices := data.get("choices"):
                                     content = choices[0]["delta"].get("content")
                                     # First token
-                                    if ttft == 0.0:
-                                        ttft = timestamp - st
-                                        output.ttft = ttft
+                                    if not first_chunk_received:
+                                        first_chunk_received = True
+                                        output.ttft = timestamp - st
 
                                     # Decoding phase
                                     else:
@@ -562,15 +569,24 @@ def to_bytes(y, sr):
                                         )
 
                                     generated_text += content or ""
+                                    # Only token chunks advance the request
+                                    # end; the trailing usage chunk carries no
+                                    # token.
+                                    most_recent_timestamp = timestamp
                                 elif usage := data.get("usage"):
                                     output.output_tokens = usage.get(
                                         "completion_tokens"
                                     )
 
-                                most_recent_timestamp = timestamp
-
                         output.generated_text = generated_text
-                        output.success = True
+                        if first_chunk_received:
+                            output.success = True
+                        else:
+                            output.success = False
+                            output.error = (
+                                "Never received a valid chunk to calculate TTFT."
+                                "This response will be marked as failed!"
+                            )
                         output.latency = most_recent_timestamp - st
                     else:
                         output.error = response.reason or ""
@@ -614,10 +630,6 @@ def get_tokenizer(
         pretrained_model_name_or_path
     ):
         pretrained_model_name_or_path = get_model(pretrained_model_name_or_path)
-    if tokenizer_mode == "slow":
-        if kwargs.get("use_fast", False):
-            raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
-        kwargs["use_fast"] = False
     if tokenizer_mode == "mistral":
         try:
             from vllm.tokenizers.mistral import MistralTokenizer
diff --git a/benchmarks/benchmark_batch_invariance.py b/benchmarks/benchmark_batch_invariance.py
index 7473a41e5140..dd22523ce201 100755
--- a/benchmarks/benchmark_batch_invariance.py
+++ b/benchmarks/benchmark_batch_invariance.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark to measure the performance overhead of VLLM_BATCH_INVARIANT mode.
+"""Benchmark to measure the performance overhead of VLLM_BATCH_INVARIANT mode.
 
 This benchmark runs the same workload twice:
 1. With VLLM_BATCH_INVARIANT=0 (baseline)
@@ -96,8 +95,7 @@ def run_benchmark_with_batch_invariant(
     batch_invariant: bool,
     seed: int = 12345,
 ) -> dict:
-    """
-    Run the benchmark with the specified configuration.
+    """Run the benchmark with the specified configuration.
 
     Returns a dict with timing and throughput metrics.
     """
diff --git a/benchmarks/benchmark_hash.py b/benchmarks/benchmark_hash.py
index 08cdc012d652..c5af4da9998b 100644
--- a/benchmarks/benchmark_hash.py
+++ b/benchmarks/benchmark_hash.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Micro benchmark comparing built-in hash(), SHA-256, and xxHash.
+"""Micro benchmark comparing built-in hash(), SHA-256, and xxHash.
 
 This focuses on a single test payload shaped like the prefix-cache hash input:
     (32-byte bytes object, 32-int tuple)
diff --git a/benchmarks/benchmark_hidden_state_extraction.py b/benchmarks/benchmark_hidden_state_extraction.py
index f0a35a0cf150..878be2c29347 100644
--- a/benchmarks/benchmark_hidden_state_extraction.py
+++ b/benchmarks/benchmark_hidden_state_extraction.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark hidden state extraction throughput.
+"""Benchmark hidden state extraction throughput.
 
 Measures two modes:
   1. Baseline: bulk inference with max_tokens=1, no extraction.
diff --git a/benchmarks/benchmark_long_document_qa_throughput.py b/benchmarks/benchmark_long_document_qa_throughput.py
index 6106b9014952..1f5b5f5504b0 100644
--- a/benchmarks/benchmark_long_document_qa_throughput.py
+++ b/benchmarks/benchmark_long_document_qa_throughput.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Offline benchmark to test the long document QA throughput.
+"""Offline benchmark to test the long document QA throughput.
 
 Example usage:
     # This workload samples 8 different prompts with a default input
@@ -49,14 +48,14 @@
 
 
 def test_long_document_qa(llm=None, sampling_params=None, prompts=None):
-    """
-    Test long document QA with the given prompts and sampling parameters.
+    """Test long document QA with the given prompts and sampling parameters.
     Print the time spent in processing all the prompts.
 
     Args:
         llm: The language model used for generating responses.
         sampling_params: Sampling parameter used to generate the response.
         prompts: A list of prompt strings to be processed by the LLM.
+
     """
     start_time = time.time()
     llm.generate(prompts, sampling_params=sampling_params)
@@ -65,8 +64,7 @@ def test_long_document_qa(llm=None, sampling_params=None, prompts=None):
 
 
 def repeat_prompts(prompts, repeat_count, mode: str):
-    """
-    Repeat each prompt in the list for a specified number of times.
+    """Repeat each prompt in the list for a specified number of times.
     The order of prompts in the output list depends on the mode.
 
     Args:
@@ -84,6 +82,7 @@ def repeat_prompts(prompts, repeat_count, mode: str):
 
     Raises:
         ValueError: If an invalid mode is provided.
+
     """
     print("Repeat mode: ", mode)
     if mode == "random":
diff --git a/benchmarks/benchmark_prefix_block_hash.py b/benchmarks/benchmark_prefix_block_hash.py
index 8bcd8af0d310..522bf9d1ab31 100644
--- a/benchmarks/benchmark_prefix_block_hash.py
+++ b/benchmarks/benchmark_prefix_block_hash.py
@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Simple benchmark to compare prefix-cache block hashing algorithms.
+"""Simple benchmark to compare prefix-cache block hashing algorithms.
 
 Example:
     python benchmark_prefix_block_hash.py --num-blocks 20000 --block-size 32
+
 """
 
 from __future__ import annotations
diff --git a/benchmarks/benchmark_prefix_caching.py b/benchmarks/benchmark_prefix_caching.py
index 0145f6b7c8e5..5abdd65b2aaf 100644
--- a/benchmarks/benchmark_prefix_caching.py
+++ b/benchmarks/benchmark_prefix_caching.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the efficiency of prefix caching.
+"""Benchmark the efficiency of prefix caching.
 
 This script allows you to benchmark the performance of
 a model with and without prefix caching using either fixed prompts
diff --git a/benchmarks/benchmark_serving_structured_output.py b/benchmarks/benchmark_serving_structured_output.py
index 664fa58dd49f..dc8e30268e7f 100644
--- a/benchmarks/benchmark_serving_structured_output.py
+++ b/benchmarks/benchmark_serving_structured_output.py
@@ -102,6 +102,7 @@ class SampleRequest:
             images).
         prompt_len: The length of the prompt in tokens.
         expected_output_len: The expected length of the output in tokens.
+
     """
 
     prompt: str
@@ -305,8 +306,7 @@ async def get_request(
     request_rate: float,
     burstiness: float = 1.0,
 ) -> AsyncGenerator[tuple[int, SampleRequest], None]:
-    """
-    Asynchronously generates requests at a specified rate
+    """Asynchronously generates requests at a specified rate
     with OPTIONAL burstiness.
 
     Args:
@@ -322,6 +322,7 @@ async def get_request(
             A lower burstiness value (0 < burstiness < 1) results
             in more bursty requests, while a higher burstiness value
             (burstiness > 1) results in a more uniform arrival of requests.
+
     """
     input_requests = iter(input_requests)
 
diff --git a/benchmarks/benchmark_topk_topp.py b/benchmarks/benchmark_topk_topp.py
index f727f16ea29c..14111f733b9a 100644
--- a/benchmarks/benchmark_topk_topp.py
+++ b/benchmarks/benchmark_topk_topp.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations.
+"""Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations.
 
 Compares:
 - apply_top_k_top_p_triton (Triton binary search)
@@ -47,8 +46,7 @@ def calculate_ops_pct(
     vocab_size: int,
     batch_size: int,
 ) -> float:
-    """
-    Calculate the percentage of active top-k and top-p operations.
+    """Calculate the percentage of active top-k and top-p operations.
 
     Returns percentage where 100% = batch_size ops.
     E.g., if all rows have both top-k and top-p active, returns 200%.
@@ -117,8 +115,7 @@ def benchmark_function(
     warmup_iters: int = 5,
     benchmark_iters: int = 20,
 ) -> tuple[float, int]:
-    """
-    Benchmark a function and return (avg_time_ms, peak_memory_bytes).
+    """Benchmark a function and return (avg_time_ms, peak_memory_bytes).
 
     Returns average time in milliseconds and peak memory usage.
     """
diff --git a/benchmarks/cutlass_benchmarks/utils.py b/benchmarks/cutlass_benchmarks/utils.py
deleted file mode 100644
index 659c68bb11d7..000000000000
--- a/benchmarks/cutlass_benchmarks/utils.py
+++ /dev/null
@@ -1,39 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Cutlass bench utils
-
-import torch
-
-
-def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
-    finfo = torch.finfo(torch.float8_e4m3fn)
-    return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to(
-        dtype=torch.float8_e4m3fn
-    )
-
-
-def to_int8(tensor: torch.Tensor) -> torch.Tensor:
-    return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)
-
-
-def to_bf16(tensor: torch.Tensor) -> torch.Tensor:
-    return tensor.to(dtype=torch.bfloat16)
-
-
-def to_fp16(tensor: torch.Tensor) -> torch.Tensor:
-    return tensor.to(dtype=torch.float16)
-
-
-def make_rand_tensors(
-    dtype: torch.dtype, m: int, n: int, k: int
-) -> tuple[torch.Tensor, torch.Tensor]:
-    a = torch.randn((m, k), device="cuda") * 5
-    b = torch.randn((n, k), device="cuda").t() * 5
-
-    if dtype == torch.int8:
-        return to_int8(a), to_int8(b)
-    if dtype == torch.float8_e4m3fn:
-        return to_fp8(a), to_fp8(b)
-
-    raise ValueError("unsupported dtype")
diff --git a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py b/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
deleted file mode 100644
index f7325ddd2cbb..000000000000
--- a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
+++ /dev/null
@@ -1,372 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import copy
-import itertools
-import pickle as pkl
-import time
-from collections.abc import Callable, Iterable
-
-import torch
-import torch.utils.benchmark as TBenchmark
-from torch.utils.benchmark import Measurement as TMeasurement
-from utils import make_rand_tensors
-from weight_shapes import WEIGHT_SHAPES
-
-from vllm import _custom_ops as ops
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-    w8a8_triton_block_scaled_mm,
-)
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm.utils.math_utils import cdiv
-
-DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
-DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
-DEFAULT_TP_SIZES = [1]
-
-
-# bench
-def bench_fn(
-    label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs
-) -> TMeasurement:
-    min_run_time = 1
-
-    globals = {
-        "args": args,
-        "kwargs": kwargs,
-        "fn": fn,
-    }
-    return TBenchmark.Timer(
-        stmt="fn(*args, **kwargs)",
-        globals=globals,
-        label=label,
-        sub_label=sub_label,
-        description=description,
-    ).blocked_autorange(min_run_time=min_run_time)
-
-
-def bench_int8(
-    dtype: torch.dtype,
-    m: int,
-    k: int,
-    n: int,
-    label: str,
-    sub_label: str,
-    bench_kernels: list[str] | None = None,
-) -> Iterable[TMeasurement]:
-    """Benchmark INT8-based kernels."""
-    assert dtype == torch.int8
-    a, b = make_rand_tensors(torch.int8, m, n, k)
-    scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
-    scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
-    bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)
-    azp = torch.zeros((m,), device="cuda", dtype=torch.int32)
-    azp_adj = torch.zeros((n,), device="cuda", dtype=torch.int32)
-
-    bench_fns = {
-        "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm(
-            a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16)
-        ),
-        "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm(
-            a.to(dtype=torch.float16), b.to(dtype=torch.float16)
-        ),
-        "cutlass_i8_i8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm(
-            a, b, scale_a, scale_b, torch.bfloat16
-        ),
-        "cutlass_i8_i8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(
-            a, b, scale_a, scale_b, torch.bfloat16, bias
-        ),
-        "cutlass_i8_i8_bf16_scaled_mm_azp": lambda: ops.cutlass_scaled_mm_azp(
-            a, b, scale_a, scale_b, torch.bfloat16, azp_adj
-        ),
-        "cutlass_i8_i8_bf16_scaled_mm_azp_bias": lambda: ops.cutlass_scaled_mm_azp(
-            a, b, scale_a, scale_b, torch.bfloat16, azp_adj, None, bias
-        ),
-        "cutlass_i8_i8_bf16_scaled_mm_azp_pt": lambda: ops.cutlass_scaled_mm_azp(
-            a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp
-        ),
-        "cutlass_i8_i8_bf16_scaled_mm_azp_pt_bias": lambda: ops.cutlass_scaled_mm_azp(
-            a, b, scale_a, scale_b, torch.bfloat16, azp_adj, azp, bias
-        ),
-    }
-
-    timers = []
-    for name, fn in bench_fns.items():
-        # If bench_kernels is None, run all. Otherwise, run only exact matches.
-        if bench_kernels is None or name in bench_kernels:
-            print(f"Running {name}")
-            timers.append(bench_fn(label, sub_label, name, fn))
-
-    return timers
-
-
-def bench_fp8(
-    dtype: torch.dtype,
-    m: int,
-    k: int,
-    n: int,
-    label: str,
-    sub_label: str,
-    bench_kernels: list[str] | None = None,
-) -> Iterable[TMeasurement]:
-    """Benchmark FP8-based kernels."""
-    assert dtype == torch.float8_e4m3fn
-    a, b = make_rand_tensors(torch.float8_e4m3fn, m, n, k)
-    a_cont = a.contiguous()
-    scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
-    scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
-
-    block_scale_a = torch.rand((m, cdiv(k, 128)), device="cuda", dtype=torch.float32)
-    block_scale_b = torch.rand(
-        cdiv(k, 128), cdiv(n, 128), device="cuda", dtype=torch.float32
-    )
-    block_scale_a_M_major = block_scale_a.t().contiguous().t()
-    block_scale_b_K_major = block_scale_b.t().contiguous().t()
-    bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)
-
-    print(m, k, n)
-
-    bench_fns = {
-        "pytorch_bf16_bf16_bf16_matmul-no-scales": lambda: torch.mm(
-            a.to(dtype=torch.bfloat16), b.to(dtype=torch.bfloat16)
-        ),
-        "pytorch_fp16_fp16_fp16_matmul-no-scales": lambda: torch.mm(
-            a.to(dtype=torch.float16), b.to(dtype=torch.float16)
-        ),
-        "pytorch_fp8_fp8_fp16_scaled_mm": lambda: torch._scaled_mm(
-            a, b, scale_a, scale_b, out_dtype=torch.float16
-        ),
-        "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum": lambda: torch._scaled_mm(
-            a, b, scale_a, scale_b, out_dtype=torch.float16, use_fast_accum=True
-        ),
-        "pytorch_fp8_fp8_bf16_scaled_mm": lambda: torch._scaled_mm(
-            a, b, scale_a, scale_b, out_dtype=torch.bfloat16
-        ),
-        "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum": lambda: torch._scaled_mm(
-            a, b, scale_a, scale_b, out_dtype=torch.bfloat16, use_fast_accum=True
-        ),
-        "cutlass_fp8_fp8_bf16_scaled_mm": lambda: ops.cutlass_scaled_mm(
-            a, b, scale_a, scale_b, torch.bfloat16
-        ),
-        "cutlass_fp8_fp8_fp16_scaled_mm": lambda: ops.cutlass_scaled_mm(
-            a, b, scale_a, scale_b, torch.float16
-        ),
-        "cutlass_fp8_fp8_bf16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(
-            a, b, scale_a, scale_b, torch.bfloat16, bias
-        ),
-        "cutlass_fp8_fp8_fp16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(
-            a, b, scale_a, scale_b, torch.float16, bias.to(dtype=torch.float16)
-        ),
-        "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_triton_block_scaled_mm(
-            a_cont, b.t(), block_scale_a, block_scale_b.t(), (128, 128)
-        ),
-        "cutlass_fp8_fp8_fp16_scaled_mm_blockwise": lambda: ops.cutlass_scaled_mm(
-            a, b, block_scale_a_M_major, block_scale_b_K_major, torch.float16
-        ),
-    }
-
-    timers = []
-    for name, fn in bench_fns.items():
-        # If bench_kernels is None, run all. Otherwise, run only exact matches.
-        if bench_kernels is None or name in bench_kernels:
-            print(f"Running {name}")
-            timers.append(bench_fn(label, sub_label, name, fn))
-
-    return timers
-
-
-def bench(
-    dtype: torch.dtype,
-    m: int,
-    k: int,
-    n: int,
-    label: str,
-    sub_label: str,
-    bench_kernels: list[str] | None = None,
-) -> Iterable[TMeasurement]:
-    if dtype == torch.int8:
-        return bench_int8(dtype, m, k, n, label, sub_label, bench_kernels)
-    if dtype == torch.float8_e4m3fn:
-        return bench_fp8(dtype, m, k, n, label, sub_label, bench_kernels)
-    raise ValueError("unsupported type")
-
-
-# runner
-def print_timers(timers: Iterable[TMeasurement]):
-    compare = TBenchmark.Compare(timers)
-    compare.print()
-
-
-def run(
-    dtype: torch.dtype,
-    MKNs: Iterable[tuple[int, int, int]],
-    bench_kernels: list[str] | None = None,
-) -> Iterable[TMeasurement]:
-    results = []
-    for m, k, n in MKNs:
-        timers = bench(
-            dtype,
-            m,
-            k,
-            n,
-            f"scaled-{dtype}-gemm",
-            f"MKN=({m}x{k}x{n})",
-            bench_kernels=bench_kernels,
-        )
-        print_timers(timers)
-        results.extend(timers)
-    return results
-
-
-def make_output(
-    data: Iterable[TMeasurement],
-    MKNs: Iterable[tuple[int, int, int]],
-    base_description: str,
-    timestamp=None,
-):
-    print(f"== All Results {base_description} ====")
-    print_timers(data)
-
-    # pickle all the results
-    timestamp = int(time.time()) if timestamp is None else timestamp
-    with open(f"{base_description}-{timestamp}.pkl", "wb") as f:
-        pkl.dump(data, f)
-
-
-def run_square_bench(args):
-    dim_sizes = list(range(args.dim_start, args.dim_end + 1, args.dim_increment))
-    MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))
-    data = run(args.dtype, MKNs, bench_kernels=args.kernels)
-    make_output(data, MKNs, f"square_bench-{args.dtype}")
-
-
-def run_range_bench(args):
-    dim_sizes = list(range(args.dim_start, args.dim_end, args.dim_increment))
-    n = len(dim_sizes)
-    Ms = [args.m_constant] * n if args.m_constant is not None else dim_sizes
-    Ks = [args.k_constant] * n if args.k_constant is not None else dim_sizes
-    Ns = [args.n_constant] * n if args.n_constant is not None else dim_sizes
-    MKNs = list(zip(Ms, Ks, Ns))
-    data = run(args.dtype, MKNs, bench_kernels=args.kernels)
-    make_output(data, MKNs, f"range_bench-{args.dtype}")
-
-
-def run_model_bench(args):
-    print("Benchmarking models:")
-    for i, model in enumerate(args.models):
-        print(f"[{i}]  {model}")
-
-    def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]:
-        KNs = []
-        for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]):
-            KN[tp_split_dim] = KN[tp_split_dim] // tp_size
-            KNs.append(KN)
-        return KNs
-
-    model_bench_data = []
-    models_tps = list(itertools.product(args.models, args.tp_sizes))
-    for model, tp_size in models_tps:
-        Ms = args.batch_sizes
-        KNs = model_shapes(model, tp_size)
-        MKNs = []
-        for m in Ms:
-            for k, n in KNs:
-                MKNs.append((m, k, n))
-
-        data = run(args.dtype, MKNs, bench_kernels=args.kernels)
-        model_bench_data.append(data)
-
-    # Print all results
-    for data, model_tp in zip(model_bench_data, models_tps):
-        model, tp_size = model_tp
-        print(f"== Results {args.dtype} {model}-TP{tp_size} ====")
-        print_timers(data)
-
-    timestamp = int(time.time())
-
-    all_data = []
-    for d in model_bench_data:
-        all_data.extend(d)
-    # pickle all data
-    with open(f"model_bench-{args.dtype}-{timestamp}.pkl", "wb") as f:
-        pkl.dump(all_data, f)
-
-
-if __name__ == "__main__":
-
-    def to_torch_dtype(dt):
-        if dt == "int8":
-            return torch.int8
-        if dt == "fp8":
-            return torch.float8_e4m3fn
-        raise ValueError("unsupported dtype")
-
-    parser = FlexibleArgumentParser(
-        description="""
-Benchmark Cutlass GEMM.
-
-    To run square GEMMs:
-        python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 square_bench --dim-start 128 --dim-end 512 --dim-increment 64
-    
-    To run constant N and K and sweep M:
-        python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384
-    
-    To run dimensions from a model:
-        python3 ./benchmarks/cutlass_benchmarks/w8a8_benchmarks.py --dtype fp8 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1
-    
-    Output:
-        - a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs.
-            """,  # noqa: E501
-        formatter_class=argparse.RawTextHelpFormatter,
-    )
-
-    parser.add_argument(
-        "--dtype",
-        type=to_torch_dtype,
-        required=True,
-        help="Available options are ['int8', 'fp8']",
-    )
-    parser.add_argument(
-        "--kernels",
-        nargs="+",
-        type=str,
-        default=None,
-        help="Exact names of the kernels to benchmark. If not set, runs all kernels.",
-    )
-
-    subparsers = parser.add_subparsers(dest="cmd")
-
-    square_parser = subparsers.add_parser("square_bench")
-    square_parser.add_argument("--dim-start", type=int, required=True)
-    square_parser.add_argument("--dim-end", type=int, required=True)
-    square_parser.add_argument("--dim-increment", type=int, required=True)
-    square_parser.set_defaults(func=run_square_bench)
-
-    range_parser = subparsers.add_parser("range_bench")
-    range_parser.add_argument("--dim-start", type=int, required=True)
-    range_parser.add_argument("--dim-end", type=int, required=True)
-    range_parser.add_argument("--dim-increment", type=int, required=True)
-    range_parser.add_argument("--m-constant", type=int, default=None)
-    range_parser.add_argument("--n-constant", type=int, default=None)
-    range_parser.add_argument("--k-constant", type=int, default=None)
-    range_parser.set_defaults(func=run_range_bench)
-
-    model_parser = subparsers.add_parser("model_bench")
-    model_parser.add_argument(
-        "--models",
-        nargs="+",
-        type=str,
-        default=DEFAULT_MODELS,
-        choices=WEIGHT_SHAPES.keys(),
-    )
-    model_parser.add_argument(
-        "--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES
-    )
-    model_parser.add_argument(
-        "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES
-    )
-    model_parser.set_defaults(func=run_model_bench)
-
-    args = parser.parse_args()
-    args.func(args)
diff --git a/benchmarks/cutlass_benchmarks/weight_shapes.py b/benchmarks/cutlass_benchmarks/weight_shapes.py
deleted file mode 100644
index 25b96ef56620..000000000000
--- a/benchmarks/cutlass_benchmarks/weight_shapes.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Weight Shapes are in the format
-# ([K, N], TP_SPLIT_DIM)
-# Example:
-#  A shape of ([14336, 4096], 0) indicates the following GEMM shape,
-#   - TP1 : K = 14336, N = 4096
-#   - TP2 : K = 7168, N = 4096
-#  A shape of ([4096, 6144], 1) indicates the following GEMM shape,
-#   - TP1 : K = 4096, N = 6144
-#   - TP4 : K = 4096, N = 1536
-
-# TP1 shapes
-WEIGHT_SHAPES = {
-    "mistralai/Mistral-7B-v0.1": [
-        ([4096, 6144], 1),
-        ([4096, 4096], 0),
-        ([4096, 28672], 1),
-        ([14336, 4096], 0),
-    ],
-    "meta-llama/Llama-2-7b-hf": [
-        ([4096, 12288], 1),
-        ([4096, 4096], 0),
-        ([4096, 22016], 1),
-        ([11008, 4096], 0),
-    ],
-    "meta-llama/Llama-3-8b": [
-        ([4096, 6144], 1),
-        ([4096, 4096], 0),
-        ([4096, 28672], 1),
-        ([14336, 4096], 0),
-    ],
-    "meta-llama/Llama-2-13b-hf": [
-        ([5120, 15360], 1),
-        ([5120, 5120], 0),
-        ([5120, 27648], 1),
-        ([13824, 5120], 0),
-    ],
-    "meta-llama/Llama-2-70b-hf": [
-        ([8192, 10240], 1),
-        ([8192, 8192], 0),
-        ([8192, 57344], 1),
-        ([28672, 8192], 0),
-    ],
-}
diff --git a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py
index 26b04299b353..61dfd7ea127c 100644
--- a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py
+++ b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark: Fused FP8 output quantization in merge_attn_states
+"""Benchmark: Fused FP8 output quantization in merge_attn_states.
 
 Compares fused vs unfused approaches for producing FP8-quantized merged
 attention output:
diff --git a/benchmarks/kernels/bench_concat_mla_q.py b/benchmarks/kernels/bench_concat_mla_q.py
index 8d940484d6b3..c528ab86b295 100644
--- a/benchmarks/kernels/bench_concat_mla_q.py
+++ b/benchmarks/kernels/bench_concat_mla_q.py
@@ -24,9 +24,9 @@ def make_inputs(num_tokens, dtype):
     """Create inputs matching the real code path.
 
     Args:
-        contiguous_nope: If False, simulate the transposed BMM output
-                         (non-contiguous nope with stride pattern from
-                         [N,B,L].transpose(0,1)).
+        num_tokens: Number of tokens to generate inputs for.
+        dtype: dtype of the generated tensors.
+
     """
     # Simulate: bmm output [N, B, L].transpose(0, 1) -> [B, N, L]
     raw = torch.randn(NUM_HEADS, num_tokens, NOPE_DIM, dtype=dtype, device="cuda")
diff --git a/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py b/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py
index 0dd5c6d84882..d0880b3f03cd 100644
--- a/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py
+++ b/benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py
@@ -106,8 +106,7 @@ def make_impl_kwargs(self, impl_type: ImplType) -> dict[str, Any]:
 
 
 def reference_quant(x: torch.Tensor, quant_out: torch.Tensor, use_ue8m0: bool):
-    """
-    Reference triton quant kernel from,
+    """Reference triton quant kernel from,
     vllm.model_executor.layers.quantization.utils.fp8_utils
     """
     assert quant_out.size() == x.size()
diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py
index 03d7fb386f74..bff4f73e2c0f 100644
--- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py
+++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the performance of the cutlass_moe_fp8 kernel vs the triton_moe
+"""Benchmark the performance of the cutlass_moe_fp8 kernel vs the triton_moe
 kernel. Both kernels take in fp8 quantized weights and 16-bit activations,
 but use different quantization strategies and backends.
 """
@@ -190,7 +189,7 @@ def bench_run(
     torch.accelerator.synchronize()
 
     def bench_cuda_graph(graph, num_warmup=5, num_iters=100):
-        """Benchmark CUDA graph using events like benchmark_moe.py"""
+        """Benchmark CUDA graph using events like benchmark_moe.py."""
         # Warmup
         for _ in range(num_warmup):
             graph.replay()
diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py
index 7379bf858889..a14242732a48 100644
--- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py
+++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the performance of the cutlass_moe_fp4 kernel vs the triton_moe
+"""Benchmark the performance of the cutlass_moe_fp4 kernel vs the triton_moe
 kernel. The cutlass_moe_fp4 kernel takes in fp4 quantized weights and 16-bit
 activations. The triton_moe kernel takes in fp8 weights(tensor scaled to fp8)
 and 16-bit activations.
diff --git a/benchmarks/kernels/benchmark_device_communicators.py b/benchmarks/kernels/benchmark_device_communicators.py
index 24e22023b91d..9f9a150471af 100644
--- a/benchmarks/kernels/benchmark_device_communicators.py
+++ b/benchmarks/kernels/benchmark_device_communicators.py
@@ -2,8 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Benchmark script for device communicators:
+"""Benchmark script for device communicators:
 CustomAllreduce (oneshot, twoshot), PyNcclCommunicator,
 and SymmMemCommunicator (multimem, two-shot).
 
@@ -17,6 +16,7 @@
 Example:
     torchrun --nproc_per_node=2 benchmark_device_communicators.py
     --sequence-lengths 512 1024 2048 --num-warmup 10 --num-trials 100
+
 """
 
 import json
@@ -29,6 +29,10 @@
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.distributed.device_communicators.aiter_custom_all_reduce import (
+    AiterCustomAllreduce,
+)
 from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
 from vllm.distributed.device_communicators.flashinfer_all_reduce import (
     FlashInferAllReduce,
@@ -42,6 +46,7 @@
 )
 from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 logger = init_logger(__name__)
@@ -56,6 +61,30 @@
 # CUDA graph settings
 CUDA_GRAPH_CAPTURE_CYCLES = 10
 
+# All-gather and reduce-scatter are distinct collectives from all-reduce, so
+# results are compared within a collective (each against its own NCCL/RCCL
+# baseline) rather than across collectives.
+COMM_COLLECTIVE = {
+    "ca_1stage": "all-reduce",
+    "ca_2stage": "all-reduce",
+    "pynccl": "all-reduce",
+    "pynccl-symm": "all-reduce",
+    "symm_mem_multimem": "all-reduce",
+    "symm_mem_two_shot": "all-reduce",
+    "flashinfer_trtllm": "all-reduce",
+    "flashinfer_mnnvl": "all-reduce",
+    "aiter_ag": "all-gather",
+    "pynccl_ag": "all-gather",
+    "aiter_rs": "reduce-scatter",
+    "pynccl_rs": "reduce-scatter",
+}
+COLLECTIVE_ORDER = ["all-reduce", "all-gather", "reduce-scatter"]
+COLLECTIVE_BASELINE = {
+    "all-reduce": "pynccl",
+    "all-gather": "pynccl_ag",
+    "reduce-scatter": "pynccl_rs",
+}
+
 
 class CommunicatorBenchmark:
     """Benchmark class for testing device communicators."""
@@ -77,6 +106,11 @@ def __init__(
         max_seq_len = max(sequence_lengths)
         max_tensor_elements = max_seq_len * HIDDEN_SIZE
         self.max_size_override = max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1
+        # AITER has effective max size of half the max size, so we double it, and
+        # account for AG output size (which is world_size larger)
+        self.aiter_max_size = (
+            2 * world_size * max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1
+        )
 
         # Initialize communicators
         self.custom_allreduce = None
@@ -85,6 +119,7 @@ def __init__(
         self.symm_mem_comm_multimem = None
         self.symm_mem_comm_two_shot = None
         self.fi_ar_comm = None
+        self.aiter_comm = None
 
         self._init_communicators()
 
@@ -181,11 +216,31 @@ def _init_communicators(self):
             )
             self.fi_ar_comm = None
 
+        if current_platform.is_rocm() and rocm_aiter_ops.is_custom_all_reduce_enabled():
+            try:
+                # Also used for AG/RS for DP
+                self.aiter_comm = AiterCustomAllreduce(
+                    group=self.cpu_group,
+                    device=self.device,
+                    max_size=self.aiter_max_size,
+                )
+                if not self.aiter_comm.disabled:
+                    logger.info("Rank %s: AITER custom AG/RS initialized", self.rank)
+                else:
+                    logger.info("Rank %s: AITER custom AG/RS disabled", self.rank)
+                    self.aiter_comm = None
+            except Exception as e:
+                logger.warning(
+                    "Rank %s: Failed to initialize AITER custom AG/RS: %s",
+                    self.rank,
+                    e,
+                )
+                self.aiter_comm = None
+
     def benchmark_allreduce(
         self, sequence_length: int, num_warmup: int, num_trials: int
     ) -> dict[str, float]:
         """Benchmark allreduce operations for all available communicators."""
-
         results = {}
 
         # Define communicators with their benchmark functions
@@ -385,29 +440,116 @@ def benchmark_allreduce_single(
                 f"CUDA graph benchmark failed for communicator: {e}"
             ) from e
 
+    def benchmark_ag_rs(
+        self, sequence_length: int, num_warmup: int, num_trials: int
+    ) -> dict[str, float]:
+        """Benchmark all-gather and reduce-scatter: custom (AITER) vs NCCL/RCCL.
+
+        Each collective is measured for both the AITER custom kernel and the
+        PyNccl baseline so they compare like-for-like (print_results groups
+        columns by collective). All-gather maps ``(seq, hidden)`` ->
+        ``(world_size * seq, hidden)``; reduce-scatter maps ``(seq, hidden)`` ->
+        ``(seq // world_size, hidden)`` and so only runs when ``seq`` divides
+        evenly across ranks. Both reuse ``benchmark_allreduce_single`` timed on a
+        ``(seq, hidden)`` input; baselines write into a pre-allocated buffer.
+        """
+        results: dict[str, float] = {}
+        entries = []
+
+        # ---- all-gather: (seq, hidden) -> (world_size * seq, hidden) ----
+        if self.aiter_comm is not None:
+            c = self.aiter_comm
+            entries.append(
+                (
+                    "aiter_ag",
+                    lambda t, c=c: c.custom_all_gather(t, dim=0),
+                    lambda t, c=c: c.should_custom_ag(t),
+                    c.capture(),
+                )
+            )
+        if self.pynccl_comm is not None:
+            c = self.pynccl_comm
+            ag_out = torch.empty(
+                sequence_length * self.world_size,
+                HIDDEN_SIZE,
+                dtype=BENCHMARK_DTYPE,
+                device=self.device,
+            )
+            entries.append(
+                (
+                    "pynccl_ag",
+                    lambda t, c=c, o=ag_out: c.all_gather(o, t),
+                    lambda t: True,
+                    nullcontext(),
+                )
+            )
+
+        # ---- reduce-scatter: (seq, hidden) -> (seq // world_size, hidden) ----
+        if sequence_length % self.world_size == 0:
+            rs_shape = (sequence_length // self.world_size, HIDDEN_SIZE)
+            if self.aiter_comm is not None:
+                c = self.aiter_comm
+                rs_out = torch.empty(
+                    rs_shape, dtype=BENCHMARK_DTYPE, device=self.device
+                )
+                entries.append(
+                    (
+                        "aiter_rs",
+                        lambda t, c=c, o=rs_out: c.custom_reduce_scatter(t, o, dim=0),
+                        lambda t, c=c: c.should_custom_rs(t, dim=0),
+                        c.capture(),
+                    )
+                )
+            if self.pynccl_comm is not None:
+                c = self.pynccl_comm
+                rs_out = torch.empty(
+                    rs_shape, dtype=BENCHMARK_DTYPE, device=self.device
+                )
+                entries.append(
+                    (
+                        "pynccl_rs",
+                        lambda t, c=c, o=rs_out: c.reduce_scatter(o, t),
+                        lambda t: True,
+                        nullcontext(),
+                    )
+                )
+
+        for name, fn, should_use_fn, context in entries:
+            latency = self.benchmark_allreduce_single(
+                sequence_length, fn, should_use_fn, context, num_warmup, num_trials
+            )
+            if latency is not None:
+                results[name] = latency
+
+        return results
+
 
-def _calculate_speedup_info(comm_results: dict[str, float]) -> str:
-    """Calculate speedup information for a single tensor size."""
+def _calculate_speedup_info(comm_results: dict[str, float], baseline: str) -> str:
+    """Fastest comm in a collective group and its speedup vs ``baseline``."""
     if not comm_results:
         return "N/A"
 
-    # Find the fastest communicator
     fastest_comm = min(comm_results.keys(), key=lambda k: comm_results[k])
     fastest_time = comm_results[fastest_comm]
 
-    # Calculate speedup vs PyNccl if available
-    if "pynccl" in comm_results:
-        pynccl_time = comm_results["pynccl"]
-        speedup = pynccl_time / fastest_time
+    if baseline in comm_results and fastest_time > 0:
+        speedup = comm_results[baseline] / fastest_time
         return f"{fastest_comm} ({speedup:.2f}x)"
-    else:
-        return f"{fastest_comm} (N/A)"
+    return f"{fastest_comm} (N/A)"
 
 
 def print_results(
-    results: dict[str, dict[str, float]], sequence_lengths: list[int], world_size: int
+    results: dict[int, dict[str, float]], sequence_lengths: list[int], world_size: int
 ):
-    """Print benchmark results in a formatted table."""
+    """Print benchmark results as one table per collective.
+
+    All-reduce, all-gather and reduce-scatter are different collectives, so each
+    gets its own table and its speedup is computed against that collective's
+    NCCL/RCCL baseline rather than across collectives.
+    """
+    present: set[str] = set()
+    for size_results in results.values():
+        present.update(size_results.keys())
 
     print(f"\n{'=' * 130}")
     print("Device Communicator Benchmark Results")
@@ -415,53 +557,45 @@ def print_results(
         f"World Size: {world_size}, Data Type: {BENCHMARK_DTYPE}, "
         f"Hidden Size: {HIDDEN_SIZE}"
     )
+    print("Tensor Shape is the per-rank input; times are ms per operation.")
     print(f"{'=' * 130}")
 
-    # Get all communicator names
-    all_comms = set()
-    for size_results in results.values():
-        all_comms.update(size_results.keys())
-
-    all_comms = sorted(list(all_comms))
-
-    # Print header
-    header = f"{'Tensor Shape':<20}{'Tensor Size':<15}"
-    for comm in all_comms:
-        header += f"{comm:<20}"
-    header += f"{'Best (Speedup vs PyNccl)':<30}"
-    print(header)
-    print("-" * len(header))
-
-    # Print results for each sequence length
-    for seq_len in sequence_lengths:
-        if seq_len in results:
-            # Calculate tensor size in elements and bytes
-            tensor_elements = seq_len * HIDDEN_SIZE
-            tensor_bytes = tensor_elements * BENCHMARK_DTYPE.itemsize
-
-            # Format tensor size (MB)
-            tensor_size_mb = tensor_bytes / (1024 * 1024)
-            tensor_size_str = f"{tensor_size_mb:.2f} MB"
-
-            # Format tensor shape
+    for collective in COLLECTIVE_ORDER:
+        comms = sorted(c for c in present if COMM_COLLECTIVE.get(c) == collective)
+        if not comms:
+            continue
+        baseline = COLLECTIVE_BASELINE[collective]
+
+        print(f"\n{collective} (speedup vs {baseline})")
+        header = f"{'Tensor Shape':<20}{'Tensor Size':<15}"
+        for comm in comms:
+            header += f"{comm:<20}"
+        header += f"{'Best':<30}"
+        print(header)
+        print("-" * len(header))
+
+        for seq_len in sequence_lengths:
+            if seq_len not in results:
+                continue
+            row_results = {
+                c: results[seq_len][c] for c in comms if c in results[seq_len]
+            }
+
+            tensor_bytes = seq_len * HIDDEN_SIZE * BENCHMARK_DTYPE.itemsize
+            tensor_size_str = f"{tensor_bytes / (1024 * 1024):.2f} MB"
             tensor_shape = f"({seq_len}, {HIDDEN_SIZE})"
 
             row = f"{tensor_shape:<20}{tensor_size_str:<15}"
-            for comm in all_comms:
-                if comm in results[seq_len]:
-                    row += f"{results[seq_len][comm]:<20.3f}"
+            for comm in comms:
+                if comm in row_results:
+                    row += f"{row_results[comm]:<20.3f}"
                 else:
                     row += f"{'N/A':<20}"
-
-            # Calculate speedup information
-            speedup_info = _calculate_speedup_info(results[seq_len])
-            row += f"{speedup_info:<30}"
-
+            row += f"{_calculate_speedup_info(row_results, baseline):<30}"
             print(row)
 
-    print(f"{'=' * 130}")
-    print("All times are in milliseconds (ms) per allreduce operation")
-    print("Speedup column shows: fastest_algorithm (speedup_vs_pynccl)")
+    print(f"\n{'=' * 130}")
+    print("All times are in milliseconds (ms) per operation")
 
 
 def main():
@@ -526,6 +660,13 @@ def main():
             num_warmup=args.num_warmup,
             num_trials=args.num_trials,
         )
+        results.update(
+            benchmark.benchmark_ag_rs(
+                sequence_length=seq_len,
+                num_warmup=args.num_warmup,
+                num_trials=args.num_trials,
+            )
+        )
 
         all_results[seq_len] = results
 
@@ -538,12 +679,23 @@ def main():
 
         # Save to JSON if requested
         if args.output_json:
-            # Add speedup information to results
+            # Add per-collective speedup information to results.
             enhanced_results = {}
             for seq_len, comm_results in all_results.items():
+                speedup_info = {}
+                for collective in COLLECTIVE_ORDER:
+                    group = {
+                        k: v
+                        for k, v in comm_results.items()
+                        if COMM_COLLECTIVE.get(k) == collective
+                    }
+                    if group:
+                        speedup_info[collective] = _calculate_speedup_info(
+                            group, COLLECTIVE_BASELINE[collective]
+                        )
                 enhanced_results[seq_len] = {
                     "timings": comm_results,
-                    "speedup_info": _calculate_speedup_info(comm_results),
+                    "speedup_info": speedup_info,
                 }
 
             output_data = {
diff --git a/benchmarks/kernels/benchmark_dsv41_mega_attn.py b/benchmarks/kernels/benchmark_dsv41_mega_attn.py
new file mode 100644
index 000000000000..1b3b5d9c3848
--- /dev/null
+++ b/benchmarks/kernels/benchmark_dsv41_mega_attn.py
@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FlashMLA mega attention vs the split-KV decode pipeline (DeepSeek V4.1).
+
+Times the full set of per-layer ops each path runs in a decode step, captured
+in a CUDA graph -- eager timing is dominated by launch overhead and flatters
+whichever path launches fewer kernels.
+
+    mega:     fused_qnorm_rope_kv_rope_quant_insert (Q pad + KV insert)
+              + mega kernel
+    split-KV: fused_qnorm_rope_kv_rope_quant_insert (Q RoPE + pad + KV insert)
+              + flash_mla_with_kvcache + fused_inv_rope_fp8_quant
+
+Both arms fuse their Q preparation into the KV insert, so that op is charged
+to both; the mega arm's Q side is only a zero-pad (nothing at all when the
+shard is already at the kernel head count) because its kernel does the Q RoPE,
+the inverse RoPE and the FP8 cast.
+
+``--local-heads`` is the variable that decides the outcome: the mega kernel
+absorbs work proportional to the live head count while paying for the padded
+one, so it wins at 64 live heads and is a wash at 16 (TP4 on a 64-head model).
+
+Both arms use the same KV records so the comparison is kernel-for-kernel:
+a V4.1 fp8 (528 B) sliding-window cache and a compressed cache that is either
+V4.1 fp8 or V4.1 NVFP4 (288 B) -- FlashMLA's SM100 sparse decode reads the
+NVFP4 compressed record in both the fused and unfused kernels.
+
+Run: .venv/bin/python benchmarks/kernels/benchmark_dsv41_mega_attn.py
+"""
+
+import argparse
+
+import torch
+
+import vllm.v1.attention.ops.flashmla as fm
+from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
+    fused_inv_rope_fp8_quant,
+)
+from vllm.models.deepseek_v41.common.ops import quantize_and_insert_k_cache
+from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
+    rope_quant_insert,
+)
+from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused
+from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
+    alloc_mega_attn_output,
+    is_flashmla_mega_attn_supported,
+)
+from vllm.utils.math_utils import round_up
+
+HEAD_DIM, ROPE_DIM = 512, 64
+V41_BYTES, V41_FP4_BYTES = 528, 288
+SWA_BLOCK, COMPRESSED_BLOCK = 32, 128
+
+
+def make_cos_sin_cache(max_pos, device):
+    inv_freq = 1.0 / (
+        10000 ** (torch.arange(0, ROPE_DIM, 2, device=device).float() / ROPE_DIM)
+    )
+    freqs = torch.outer(torch.arange(max_pos, device=device).float(), inv_freq)
+    return torch.cat([freqs.cos(), freqs.sin()], -1)
+
+
+def empty_paged_cache(num_rows, block_size, bytes_per_token, device):
+    """A zeroed paged cache view [num_blocks, block_size, bytes_per_token]."""
+    num_blocks = (num_rows + block_size - 1) // block_size + 1
+    page = round_up(block_size * bytes_per_token, 512)
+    backing = torch.zeros(num_blocks, page, dtype=torch.uint8, device=device)
+    return backing, backing.as_strided(
+        (num_blocks, block_size, bytes_per_token), (page, bytes_per_token, 1)
+    )
+
+
+def fill_cache(view, bytes_per_token, cos_sin, device):
+    n = view.shape[0] * view.shape[1]
+    rows = torch.randn(n, HEAD_DIM, device=device, dtype=torch.bfloat16)
+    slots = torch.arange(n, dtype=torch.int64, device=device)
+    if bytes_per_token == V41_FP4_BYTES:
+        # A page-aligned cache has more slots than the cos_sin table has rows;
+        # wrap so the insert never indexes past it.
+        positions = slots % cos_sin.shape[0]
+        rope_quant_insert(rows, positions, cos_sin, view, slots, 1)
+    else:
+        quantize_and_insert_k_cache(
+            rows,
+            view.reshape(view.shape[0], -1),
+            slots,
+            block_size=view.shape[1],
+            bytes_per_token=bytes_per_token,
+        )
+
+
+def time_graph(fn, iters=100, warmup=5):
+    """Capture ``fn`` in a CUDA graph and time its replay."""
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(3):
+            fn()
+    torch.cuda.current_stream().wait_stream(s)
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(warmup):
+        g.replay()
+    torch.accelerator.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        g.replay()
+    end.record()
+    torch.accelerator.synchronize()
+    return start.elapsed_time(end) * 1000 / iters
+
+
+def bench_decode(s_q, device, topk_extra, extra_bytes, local_heads):
+    padded_heads, scale = 64, HEAD_DIM**-0.5
+    n_wv_group = padded_heads // 8
+    cos_sin = make_cos_sin_cache(65536, device)
+    positions = torch.randint(0, 65536, (s_q,), device=device, dtype=torch.int64)
+    pos32 = positions.to(torch.int32)
+
+    # Per-step inputs: the wq_b output (local heads) and this step's KV row.
+    q_local = torch.randn(
+        s_q, local_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
+    )
+    q_fused = permute_q_to_fused(q_local)
+    kv = torch.randn(s_q, HEAD_DIM, device=device, dtype=torch.bfloat16)
+
+    n_swa = s_q * 128 + 128
+    swa_backing, swa = empty_paged_cache(n_swa, SWA_BLOCK, V41_BYTES, device)
+    fill_cache(swa, V41_BYTES, cos_sin, device)
+    _, extra = empty_paged_cache(65536, COMPRESSED_BLOCK, extra_bytes, device)
+    fill_cache(extra, extra_bytes, cos_sin, device)
+
+    slot_mapping = torch.arange(s_q, dtype=torch.int64, device=device)
+    swa_idx = torch.randint(0, n_swa, (s_q, 128), device=device, dtype=torch.int32)
+    swa_len = torch.full((s_q,), 128, device=device, dtype=torch.int32)
+    ex_idx = torch.randint(
+        0, 65536, (s_q, topk_extra), device=device, dtype=torch.int32
+    )
+    ex_len = torch.full((s_q,), topk_extra, device=device, dtype=torch.int32)
+    sink = torch.zeros(padded_heads, device=device)
+    sched = fm.FlashMLASchedMeta()
+    out = alloc_mega_attn_output(s_q, n_wv_group, device)
+    swa4 = swa.unsqueeze(-2)
+    extra4 = extra.unsqueeze(-2)
+    swa_2d = swa_backing
+
+    def mega():
+        # One launch zero-pads the fused-layout Q and inserts the SWA KV.
+        q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+            q_fused,
+            kv,
+            swa_2d,
+            slot_mapping,
+            positions,
+            cos_sin,
+            0 if padded_heads == local_heads else padded_heads,
+            1e-6,
+            SWA_BLOCK,
+            False,  # apply_q_norm
+            True,  # kv_mxfp8
+            False,  # apply_q_rope: the mega kernel rotates Q itself
+            True,  # is_q_interleaved
+        )
+        torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode(
+            q_fused if padded_heads == local_heads else q_pad,
+            swa4,
+            swa_idx,
+            scale,
+            HEAD_DIM,
+            sink,
+            swa_len,
+            extra4,
+            ex_idx,
+            ex_len,
+            False,
+            0.0,
+            pos32,
+            False,
+            64,
+            cos_sin,
+            n_wv_group,
+            32,
+            True,
+            True,
+            True,
+            out.data,
+            out.scale,
+        )
+
+    def split_kv():
+        q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+            q_local,
+            kv,
+            swa_2d,
+            slot_mapping,
+            positions,
+            cos_sin,
+            padded_heads,
+            1e-6,
+            SWA_BLOCK,
+            False,
+            True,
+        )
+        o = fm.flash_mla_with_kvcache(
+            q=q_pad.unsqueeze(1),
+            k_cache=swa4,
+            block_table=None,
+            head_dim_v=HEAD_DIM,
+            tile_scheduler_metadata=sched,
+            cache_seqlens=None,
+            is_fp8_kvcache=True,
+            indices=swa_idx.view(s_q, 1, -1),
+            topk_length=swa_len,
+            softmax_scale=scale,
+            attn_sink=sink,
+            extra_k_cache=extra4,
+            extra_indices_in_kvcache=ex_idx.view(s_q, 1, -1),
+            extra_topk_length=ex_len,
+        )[0]
+        fused_inv_rope_fp8_quant(
+            o.squeeze(1)[:, :local_heads],
+            positions,
+            cos_sin,
+            n_groups=local_heads // 8,
+            heads_per_group=8,
+            quant_group_size=32,
+            tma_aligned_scales=True,
+        )
+
+    return time_graph(mega), time_graph(split_kv)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--decode-s-q",
+        nargs="+",
+        type=int,
+        default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512],
+    )
+    parser.add_argument("--topk-extra", type=int, default=512)
+    parser.add_argument("--local-heads", type=int, default=16)
+    parser.add_argument(
+        "--extra-bytes",
+        type=int,
+        choices=[V41_BYTES, V41_FP4_BYTES],
+        default=V41_BYTES,
+        help="compressed-cache record: 528 (V4.1 fp8) or 288 (V4.1 NVFP4)",
+    )
+    args = parser.parse_args()
+    ok, reason = is_flashmla_mega_attn_supported()
+    if not ok:
+        raise SystemExit(reason)
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    print(
+        f"decode, cudagraph (topk_swa=128, topk_extra={args.topk_extra}, "
+        f"local_heads={args.local_heads}, extra={args.extra_bytes}B)"
+    )
+    print(f"{'s_q':>6} {'mega_us':>9} {'split_kv_us':>12} {'speedup':>8}")
+    for s_q in args.decode_s_q:
+        m, s = bench_decode(
+            s_q, device, args.topk_extra, args.extra_bytes, args.local_heads
+        )
+        print(f"{s_q:>6d} {m:9.1f} {s:12.1f} {s / m:7.2f}x")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/kernels/benchmark_fused_collective.py b/benchmarks/kernels/benchmark_fused_collective.py
index c999c16021bf..aaa39fd9fe06 100644
--- a/benchmarks/kernels/benchmark_fused_collective.py
+++ b/benchmarks/kernels/benchmark_fused_collective.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Benchmark for FlashInfer fused collective operations vs standard operations.
+"""Benchmark for FlashInfer fused collective operations vs standard operations.
 
 This benchmark compares:
 1. FlashInfer's allreduce_fusion with trtllm backend
@@ -446,9 +445,15 @@ def run_benchmarks(
     """Run all benchmarks for given configuration.
 
     Args:
+        num_tokens: Number of tokens in the benchmark input.
+        hidden_dim: Hidden dimension of the benchmark input.
+        dtype: dtype of the benchmark input.
+        use_residual: Whether to include a residual add in the fused op.
+        no_oneshot: Skip the one-shot allreduce variants.
         allreduce_params: Shared parameters for FlashInfer fused allreduce.
         workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace.
         quant_modes: Set of quantization modes: "none", "fp8", "fp4".
+
     """
     (
         input_tensor,
diff --git a/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py b/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py
new file mode 100644
index 000000000000..9f8d97612161
--- /dev/null
+++ b/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import itertools
+
+import torch
+
+import vllm.kernels  # noqa: F401
+from vllm import ir
+from vllm.benchmarks.lib.utils import default_vllm_config
+from vllm.triton_utils import triton
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
+
+NUM_TOKENS = [1, 2, 8, 32, 128, 512, 2048]
+INTERMEDIATE_SIZES = [2048, 4096, 8192, 13824, 16384, 32768]
+CONFIGS = list(itertools.product(NUM_TOKENS, INTERMEDIATE_SIZES))
+STD_MULTIPLIER = 1.6448536269514722
+
+
+@default_vllm_config()
+def benchmark_gelu_and_mul_sparse(
+    num_tokens: int,
+    intermediate_size: int,
+    provider: str,
+    dtype: torch.dtype,
+):
+    set_random_seed(0)
+    torch.set_default_device("cuda")
+    x = torch.randn(num_tokens, 2 * intermediate_size, dtype=dtype)
+    native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn
+    triton_impl = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn
+    args = (x, STD_MULTIPLIER, "tanh")
+
+    expected = native(*args)
+    actual = triton_impl(*args)
+    torch.testing.assert_close(
+        actual,
+        expected,
+        **ir.ops.gelu_and_mul_sparse.get_tolerance(actual.dtype),
+    )
+
+    if provider == "triton":
+        fn = lambda: triton_impl(*args)
+    elif provider == "native_compiled":
+        compiled = torch.compile(native, fullgraph=True)
+        compiled(*args)
+        fn = lambda: compiled(*args)
+    else:
+        fn = lambda: native(*args)
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+        fn, quantiles=[0.5, 0.2, 0.8]
+    )
+    return ms, max_ms, min_ms
+
+
+if __name__ == "__main__":
+    parser = FlexibleArgumentParser(
+        description="Benchmark the Gemma3n sparse GELU activation."
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        choices=["half", "bfloat16", "float"],
+        default="bfloat16",
+    )
+    args = parser.parse_args()
+    dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
+
+    perf_report = triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["num_tokens", "intermediate_size"],
+            x_vals=CONFIGS,
+            line_arg="provider",
+            line_vals=["triton", "native_compiled", "native_eager"],
+            line_names=["Triton", "Native (torch.compile)", "Native (eager)"],
+            styles=[("blue", "-"), ("green", "-"), ("red", "-")],
+            ylabel="ms",
+            plot_name="gelu-and-mul-sparse-performance",
+            args={},
+        )
+    )
+    perf_report(
+        lambda num_tokens, intermediate_size, provider: (
+            benchmark_gelu_and_mul_sparse(
+                num_tokens, intermediate_size, provider, dtype
+            )
+        )
+    ).run(print_data=True)
diff --git a/benchmarks/kernels/benchmark_gemm_rs_ar.py b/benchmarks/kernels/benchmark_gemm_rs_ar.py
new file mode 100644
index 000000000000..3e47f61c74f0
--- /dev/null
+++ b/benchmarks/kernels/benchmark_gemm_rs_ar.py
@@ -0,0 +1,407 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark the SM100 GEMM-RS/AR kernel against the model's unfused path.
+
+The baseline is what a model runs without the fusion: the projection's own
+linear kernel followed by vLLM's collective (``sp_reduce_scatter`` for RS,
+``tensor_model_parallel_all_reduce`` for AR), with the collective backend
+chosen by vLLM's dispatch exactly as in serving. The same two kernels are also
+timed bare (the GEMM alone, the collective alone on a precomputed partial);
+their max is the floor a perfectly overlapped fused kernel could reach.
+
+All ranks must belong to one NVLink domain. Run a TP8 sweep over the Kimi-K3
+projections with:
+
+    torchrun --nproc-per-node=8 -- benchmarks/kernels/benchmark_gemm_rs_ar.py
+
+or the DeepSeek-V4.1 MXFP8 ``wo_b`` projection at TP4 with:
+
+    torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \
+        --model deepseek_v41 --backend flashinfer_cutedsl
+"""
+
+import argparse
+import json
+import os
+import statistics
+from collections.abc import Callable
+
+import torch
+import torch.distributed as dist
+
+from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+from vllm.distributed import (
+    cleanup_dist_env_and_memory,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.distributed.parallel_state import (
+    get_tp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
+from vllm.model_executor.layers.linear import RowParallelLinear
+from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter
+
+# Per model: (hidden size N, global input widths K). Kimi-K3 lists the
+# shared-expert down-proj and attention O-proj; DeepSeek-V4.1 lists the
+# ``wo_b`` output projection (o_groups * o_lora_rank).
+_MODEL_PROJECTIONS = {
+    "kimi_k3": (7168, (6144, 12288)),
+    "deepseek_v41": (5120, (8192,)),
+}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--mode",
+        choices=("rs", "ar"),
+        default="rs",
+        help="Collective mode to benchmark.",
+    )
+    parser.add_argument(
+        "--m",
+        type=int,
+        nargs="+",
+        default=[128, 512, 2048, 8192, 32768],
+        help="Global token counts to benchmark.",
+    )
+    parser.add_argument(
+        "--model",
+        choices=tuple(_MODEL_PROJECTIONS),
+        default="kimi_k3",
+        help="Model whose projection shapes set the default --n and --k.",
+    )
+    parser.add_argument(
+        "--backend",
+        choices=("bf16", "flashinfer_cutlass", "flashinfer_cutedsl"),
+        default="bf16",
+        help="Weight dtype/layout: BF16 or an online-MXFP8 FlashInfer kernel.",
+    )
+    parser.add_argument(
+        "--k",
+        type=int,
+        nargs="+",
+        help=(
+            "Per-rank input dimensions. By default, derive the selected "
+            "model's projection dimensions from the TP world size."
+        ),
+    )
+    parser.add_argument(
+        "--n",
+        type=int,
+        help="Output width. Defaults to the selected model's hidden size.",
+    )
+    parser.add_argument(
+        "--num-workspaces",
+        type=int,
+        default=10,
+        help="Pointer-distinct inputs and CUDA graphs to rotate.",
+    )
+    parser.add_argument("--warmup-replays", type=int, default=5)
+    parser.add_argument("--samples", type=int, default=20)
+    parser.add_argument("--json", type=str, help="Also write the results here.")
+    return parser.parse_args()
+
+
+def make_projection(
+    gemm_rs_ar: GemmRsAr,
+    weight: torch.Tensor,
+    world_size: int,
+    backend: str,
+) -> RowParallelLinear:
+    """Wrap a BF16 [N, K] weight the way the model would hand it to GEMM-RS.
+
+    For the MXFP8 backends this quantizes online through the FlashInfer
+    kernel's post-load step, so the fused kernel and the unfused baseline see
+    the production weight and scale layouts.
+    """
+    N, K = weight.shape
+    quant_config = None
+    if backend != "bf16":
+        from vllm.config.quantization import QuantizationConfigArgs
+        from vllm.model_executor.layers.quantization.online.base import (
+            OnlineQuantizationConfig,
+        )
+
+        quant_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8"))
+    with torch.device(weight.device):
+        linear = RowParallelLinear(
+            K * world_size,
+            N,
+            bias=False,
+            params_dtype=weight.dtype,
+            quant_config=quant_config,
+            # The baseline reduces the local partial itself, like the models.
+            reduce_results=False,
+            return_bias=False,
+        )
+    # Select the fused path before online quantization replaces the weights.
+    assert gemm_rs_ar.can_run(linear)
+    linear.weight = torch.nn.Parameter(weight, requires_grad=False)
+    linear.quant_method.process_weights_after_loading(linear)
+    return linear
+
+
+def capture_graph(
+    op: Callable[[], torch.Tensor],
+    stream: torch.cuda.Stream,
+    cpu_group: dist.ProcessGroup,
+) -> torch.cuda.CUDAGraph:
+    stream.wait_stream(torch.cuda.current_stream())
+    dist.barrier(group=cpu_group)
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            op()
+    stream.synchronize()
+    dist.barrier(group=cpu_group)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        op()
+    torch.cuda.current_stream().wait_stream(stream)
+    dist.barrier(group=cpu_group)
+    return graph
+
+
+def benchmark_graphs(
+    candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]],
+    warmup_replays: int,
+    samples: int,
+    device_group: dist.ProcessGroup,
+    device_barrier: Callable[[], None],
+) -> dict[str, float]:
+    """Median over samples of the slowest rank's replay time, in us."""
+    names = list(candidate_graphs)
+    for round_index in range(warmup_replays):
+        for offset in range(len(names)):
+            name = names[(round_index + offset) % len(names)]
+            graphs = candidate_graphs[name]
+            device_barrier()
+            graphs[round_index % len(graphs)].replay()
+    torch.accelerator.synchronize()
+
+    timings: dict[str, list[float]] = {name: [] for name in names}
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    for sample_index in range(samples):
+        for offset in range(len(names)):
+            name = names[(sample_index + offset) % len(names)]
+            graphs = candidate_graphs[name]
+            device_barrier()
+            start.record()
+            graphs[sample_index % len(graphs)].replay()
+            end.record()
+            end.synchronize()
+
+            elapsed = torch.tensor(
+                start.elapsed_time(end) * 1000,
+                dtype=torch.float64,
+                device=torch.accelerator.current_device_index(),
+            )
+            dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
+            timings[name].append(elapsed.item())
+    return {name: statistics.median(values) for name, values in timings.items()}
+
+
+def benchmark_shape(
+    gemm_rs_ar: GemmRsAr,
+    mode: str,
+    backend: str,
+    M: int,
+    N: int,
+    K: int,
+    num_workspaces: int,
+    warmup_replays: int,
+    samples: int,
+    device_group: dist.ProcessGroup,
+    cpu_group: dist.ProcessGroup,
+    device_barrier: Callable[[], None],
+) -> dict[str, float | int | str]:
+    all_reduce = mode == "ar"
+    world_size = dist.get_world_size(device_group)
+    rank = dist.get_rank(device_group)
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    local_M = (M + world_size - 1) // world_size
+
+    rng = torch.Generator(device=device)
+    rng.manual_seed(1000 + rank * 10 + M + K)
+    inputs = [
+        torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng)
+        for _ in range(num_workspaces)
+    ]
+    projections = [
+        make_projection(
+            gemm_rs_ar,
+            torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng),
+            world_size,
+            backend,
+        )
+        for _ in range(num_workspaces)
+    ]
+
+    def collective(partial: torch.Tensor) -> torch.Tensor:
+        if all_reduce:
+            return tensor_model_parallel_all_reduce(partial)
+        return sp_reduce_scatter(partial)
+
+    def gemm(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
+        return linear(x)
+
+    def unfused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
+        return collective(linear(x))
+
+    def fused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor:
+        return gemm_rs_ar.apply(x, linear)
+
+    # The bare collective runs on the GEMM's own output so it sees the same
+    # message size and layout the unfused path hands to the collective.
+    partials = [linear(x) for x, linear in zip(inputs, projections)]
+
+    expected = unfused(inputs[0], projections[0])
+    actual = fused(inputs[0], projections[0])
+    torch.accelerator.synchronize(device)
+    rows = M if all_reduce else min(max(M - rank * local_M, 0), local_M)
+    torch.testing.assert_close(
+        actual[:rows], expected[:rows], rtol=5e-2, atol=4.0, msg=f"M={M}, K={K}"
+    )
+
+    candidates: dict[str, list[Callable[[], torch.Tensor]]] = {
+        "gemm": [lambda x=x, w=w: gemm(x, w) for x, w in zip(inputs, projections)],
+        "collective": [lambda p=p: collective(p) for p in partials],
+        "vllm": [lambda x=x, w=w: unfused(x, w) for x, w in zip(inputs, projections)],
+        "gemm_rs_ar": [
+            lambda x=x, w=w: fused(x, w) for x, w in zip(inputs, projections)
+        ],
+    }
+
+    candidate_graphs = {}
+    keepalive: list[object] = []
+    for name, ops in candidates.items():
+        stream = torch.cuda.Stream()
+        candidate_graphs[name] = [capture_graph(op, stream, cpu_group) for op in ops]
+        keepalive.append(stream)
+
+    times = benchmark_graphs(
+        candidate_graphs, warmup_replays, samples, device_group, device_barrier
+    )
+    floor = max(times["gemm"], times["collective"])
+    return {
+        "mode": mode.upper(),
+        "backend": backend,
+        "M": M,
+        "N": N,
+        "K": K,
+        "gemm_us": times["gemm"],
+        "collective_us": times["collective"],
+        "floor_us": floor,
+        "vllm_us": times["vllm"],
+        "gemm_rs_ar_us": times["gemm_rs_ar"],
+        "latency_change": (times["gemm_rs_ar"] - times["vllm"]) / times["vllm"],
+        "floor_gap": (times["gemm_rs_ar"] - floor) / floor,
+    }
+
+
+def print_results(results: list[dict[str, float | int | str]]) -> None:
+    collective = results[0]["mode"]
+    backend = results[0]["backend"]
+    print(f"### GEMM-{collective} vs vLLM unfused path ({backend} weights)")
+    print(
+        f"| M | N | K | GEMM (us) | {collective} (us) | floor (us) | vLLM (us) | "
+        f"GEMM-{collective} (us) | vs vLLM | vs floor |"
+    )
+    print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|")
+    for r in results:
+        print(
+            f"| {r['M']} | {r['N']} | {r['K']} | {r['gemm_us']:.1f} | "
+            f"{r['collective_us']:.1f} | {r['floor_us']:.1f} | {r['vllm_us']:.1f} | "
+            f"{r['gemm_rs_ar_us']:.1f} | {r['latency_change']:+.1%} | "
+            f"{r['floor_gap']:+.1%} |"
+        )
+
+
+def main() -> None:
+    args = parse_args()
+    default_n, default_k = _MODEL_PROJECTIONS[args.model]
+    if args.n is None:
+        args.n = default_n
+    assert args.m and min(args.m) > 0
+    assert args.n % 128 == 0
+    assert args.num_workspaces > 0
+    assert args.warmup_replays >= 0
+    assert args.samples > 0
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.accelerator.set_device_index(local_rank)
+    init_distributed_environment()
+    world_size = dist.get_world_size()
+    if args.k is None:
+        assert all(K % world_size == 0 for K in default_k)
+        K_values = [K // world_size for K in default_k]
+    else:
+        K_values = args.k
+    k_alignment = 64 if args.backend == "bf16" else 128
+    assert all(K % k_alignment == 0 for K in K_values)
+
+    config = VllmConfig()
+    config.model_config = ModelConfig(dtype="bfloat16")
+    # Pin the online-MXFP8 kernel the projections quantize through.
+    config.kernel_config.linear_backend = (
+        "auto" if args.backend == "bf16" else args.backend
+    )
+    with set_current_vllm_config(config):
+        # Builds the TP communicator as in serving, so the baseline's
+        # collective is whatever vLLM's dispatch selects (FlashInfer all-reduce
+        # first, then NCCL symmetric memory, custom all-reduce, pynccl).
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+
+    tp_group = get_tp_group()
+    pynccl_comm = tp_group.device_communicator.pynccl_comm
+    assert pynccl_comm is not None
+    sync_input = torch.zeros(1, device=torch.accelerator.current_device_index())
+    sync_output = torch.empty_like(sync_input)
+
+    def device_barrier() -> None:
+        # Order the timed launch after a device-side rank rendezvous without
+        # including the rendezvous itself in the measured event interval.
+        pynccl_comm.all_reduce(sync_input, sync_output)
+
+    gemm_rs_ar = GemmRsAr(
+        max_M=max(args.m),
+        N=args.n,
+        all_reduce=args.mode == "ar",
+    )
+    with set_current_vllm_config(config):
+        results = [
+            benchmark_shape(
+                gemm_rs_ar,
+                args.mode,
+                args.backend,
+                M,
+                args.n,
+                K,
+                args.num_workspaces,
+                args.warmup_replays,
+                args.samples,
+                tp_group.device_group,
+                tp_group.cpu_group,
+                device_barrier,
+            )
+            for K in K_values
+            for M in args.m
+        ]
+    del gemm_rs_ar
+
+    if tp_group.rank_in_group == 0:
+        print_results(results)
+        if args.json:
+            with open(args.json, "w") as f:
+                json.dump(results, f, indent=2)
+
+    dist.barrier(group=tp_group.cpu_group)
+    cleanup_dist_env_and_memory()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/kernels/benchmark_kimi_k3_attn_res.py b/benchmarks/kernels/benchmark_kimi_k3_attn_res.py
index c8cf6f9acc77..dba82a872226 100644
--- a/benchmarks/kernels/benchmark_kimi_k3_attn_res.py
+++ b/benchmarks/kernels/benchmark_kimi_k3_attn_res.py
@@ -10,6 +10,7 @@
 
     .venv/bin/python benchmarks/kernels/benchmark_kimi_k3_attn_res.py \
       --output /tmp/attn_res.json
+
 """
 
 from __future__ import annotations
diff --git a/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py b/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py
deleted file mode 100644
index e4daea26bbce..000000000000
--- a/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py
+++ /dev/null
@@ -1,469 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark the SM100 Kimi-K3 GEMM-RS/AR kernel.
-
-All ranks must belong to one NVLink domain. For example, run a TP8 sweep with:
-
-    torchrun --nproc-per-node=8 \
-        benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py
-"""
-
-import argparse
-import os
-import statistics
-from collections.abc import Callable
-from dataclasses import dataclass
-
-import pandas as pd
-import torch
-import torch.distributed as dist
-import torch.distributed._symmetric_memory as symm_mem
-
-from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.distributed import cleanup_dist_env_and_memory
-from vllm.distributed.parallel_state import (
-    get_tp_group,
-    init_distributed_environment,
-    initialize_model_parallel,
-)
-from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr
-
-# Shared-expert down-proj and attention O-proj.
-_KIMI_K3_PROJECTION_K = (6144, 12288)
-
-
-@dataclass
-class Candidate:
-    name: str
-    runs: list[Callable[[], torch.Tensor]]
-    check_correctness: bool = True
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--mode",
-        choices=("rs", "ar"),
-        default="rs",
-        help="Collective mode to benchmark.",
-    )
-    parser.add_argument(
-        "--m",
-        type=int,
-        nargs="+",
-        default=[128, 512, 2048, 8192, 32768],
-        help="Global token counts to benchmark.",
-    )
-    parser.add_argument(
-        "--k",
-        type=int,
-        nargs="+",
-        help=(
-            "Per-rank input dimensions. By default, derive the Kimi-K3 "
-            "shared-expert down-proj and O-proj dimensions from the TP "
-            "world size."
-        ),
-    )
-    parser.add_argument("--n", type=int, default=7168)
-    parser.add_argument(
-        "--num-workspaces",
-        type=int,
-        default=10,
-        help="Pointer-distinct inputs and CUDA graphs to rotate.",
-    )
-    parser.add_argument("--warmup-replays", type=int, default=5)
-    parser.add_argument("--samples", type=int, default=20)
-    return parser.parse_args()
-
-
-def capture_graph(
-    op: Callable[[], torch.Tensor],
-    stream: torch.cuda.Stream,
-    cpu_group: dist.ProcessGroup,
-) -> tuple[torch.cuda.CUDAGraph, list[torch.Tensor | None]]:
-    result: list[torch.Tensor | None] = [None]
-    stream.wait_stream(torch.cuda.current_stream())
-    dist.barrier(group=cpu_group)
-    with torch.cuda.stream(stream):
-        for _ in range(3):
-            result[0] = op()
-    stream.synchronize()
-    dist.barrier(group=cpu_group)
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        result[0] = op()
-    torch.cuda.current_stream().wait_stream(stream)
-    dist.barrier(group=cpu_group)
-    return graph, result
-
-
-def benchmark_graphs(
-    candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]],
-    warmup_replays: int,
-    samples: int,
-    device_group: dist.ProcessGroup,
-    device_barrier: Callable[[], None],
-) -> dict[str, float]:
-    candidate_names = list(candidate_graphs)
-    for round_index in range(warmup_replays):
-        for candidate_index in range(len(candidate_names)):
-            candidate_id = (round_index + candidate_index) % len(candidate_names)
-            name = candidate_names[candidate_id]
-            graphs = candidate_graphs[name]
-            device_barrier()
-            graphs[round_index % len(graphs)].replay()
-    torch.accelerator.synchronize()
-
-    timings: dict[str, list[float]] = {name: [] for name in candidate_names}
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    for sample_index in range(samples):
-        for candidate_index in range(len(candidate_names)):
-            candidate_id = (sample_index + candidate_index) % len(candidate_names)
-            name = candidate_names[candidate_id]
-            graphs = candidate_graphs[name]
-            device_barrier()
-            start.record()
-            graphs[sample_index % len(graphs)].replay()
-            end.record()
-            end.synchronize()
-
-            elapsed = torch.tensor(
-                start.elapsed_time(end) * 1000,
-                dtype=torch.float64,
-                device=torch.accelerator.current_device_index(),
-            )
-            dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
-            timings[name].append(elapsed.item())
-    return {name: statistics.median(values) for name, values in timings.items()}
-
-
-def valid_rows(M: int, local_M: int, rank: int) -> int:
-    return min(max(M - rank * local_M, 0), local_M)
-
-
-def benchmark_shape(
-    gemm_rs_ar: GemmRsAr,
-    mode: str,
-    M: int,
-    N: int,
-    K: int,
-    num_workspaces: int,
-    warmup_replays: int,
-    samples: int,
-    device_group: dist.ProcessGroup,
-    cpu_group: dist.ProcessGroup,
-    device_barrier: Callable[[], None],
-) -> dict[str, float | int | str]:
-    all_reduce = mode == "ar"
-    world_size = dist.get_world_size(device_group)
-    rank = dist.get_rank(device_group)
-    device = torch.device("cuda", torch.accelerator.current_device_index())
-    padded_M = (M + world_size - 1) // world_size * world_size
-    local_M = padded_M // world_size
-
-    rng = torch.Generator(device=device)
-    rng.manual_seed(1000 + rank * 10 + M + K)
-    inputs = [
-        torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng)
-        for _ in range(num_workspaces)
-    ]
-    weights = [
-        torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng)
-        for _ in range(num_workspaces)
-    ]
-
-    partial = torch.empty((padded_M, N), dtype=torch.bfloat16, device=device)
-    symm_partial = symm_mem.empty((padded_M, N), dtype=torch.bfloat16, device=device)
-    symm_partial_handle = symm_mem.rendezvous(symm_partial, device_group)
-    collective_inputs = [torch.empty_like(partial) for _ in range(num_workspaces)]
-    symm_collective_inputs = []
-    symm_collective_handles = []
-    for _ in range(num_workspaces):
-        collective_input = symm_mem.empty(
-            (padded_M, N),
-            dtype=torch.bfloat16,
-            device=device,
-        )
-        symm_collective_inputs.append(collective_input)
-        symm_collective_handles.append(
-            symm_mem.rendezvous(collective_input, device_group)
-        )
-    torch_output = torch.empty((local_M, N), dtype=torch.bfloat16, device=device)
-    symm_output = torch.empty_like(torch_output)
-    gemm_output = torch.empty((M, N), dtype=torch.bfloat16, device=device)
-
-    if padded_M > M:
-        partial[M:].zero_()
-        symm_partial[M:].zero_()
-
-    def make_torch_ring_ll_gemm_collective(
-        x: torch.Tensor, weight: torch.Tensor
-    ) -> Callable[[], torch.Tensor]:
-        def run() -> torch.Tensor:
-            torch.mm(x, weight.T, out=partial[:M])
-            if all_reduce:
-                dist.all_reduce(partial, group=device_group)
-                return partial[:M]
-            dist.reduce_scatter_single(torch_output, partial, group=device_group)
-            return torch_output
-
-        return run
-
-    def make_torch_ldmc_gemm_collective(
-        x: torch.Tensor, weight: torch.Tensor
-    ) -> Callable[[], torch.Tensor]:
-        def run() -> torch.Tensor:
-            torch.mm(x, weight.T, out=symm_partial[:M])
-            if all_reduce:
-                dist.all_reduce(symm_partial, group=device_group)
-                return symm_partial[:M]
-            dist.reduce_scatter_single(
-                symm_output,
-                symm_partial,
-                group=device_group,
-            )
-            return symm_output
-
-        return run
-
-    def make_fused_gemm_collective(
-        x: torch.Tensor, weight: torch.Tensor
-    ) -> Callable[[], torch.Tensor]:
-        def run() -> torch.Tensor:
-            return gemm_rs_ar(x, weight)
-
-        return run
-
-    def make_torch_gemm(
-        x: torch.Tensor,
-        weight: torch.Tensor,
-    ) -> Callable[[], torch.Tensor]:
-        def run() -> torch.Tensor:
-            return torch.mm(x, weight.T, out=gemm_output)
-
-        return run
-
-    def make_ring_ll_collective(
-        collective_input: torch.Tensor,
-    ) -> Callable[[], torch.Tensor]:
-        def run() -> torch.Tensor:
-            if all_reduce:
-                dist.all_reduce(collective_input, group=device_group)
-                return collective_input
-            dist.reduce_scatter_single(
-                torch_output, collective_input, group=device_group
-            )
-            return torch_output
-
-        return run
-
-    def make_ldmc_collective(
-        collective_input: torch.Tensor,
-    ) -> Callable[[], torch.Tensor]:
-        def run() -> torch.Tensor:
-            if all_reduce:
-                dist.all_reduce(collective_input, group=device_group)
-                return collective_input
-            dist.reduce_scatter_single(
-                symm_output, collective_input, group=device_group
-            )
-            return symm_output
-
-        return run
-
-    candidates = (
-        Candidate(
-            "ring_ll_us",
-            [make_torch_ring_ll_gemm_collective(x, w) for x, w in zip(inputs, weights)],
-        ),
-        Candidate(
-            "ldmc_us",
-            [make_torch_ldmc_gemm_collective(x, w) for x, w in zip(inputs, weights)],
-        ),
-        Candidate(
-            "gemm_rs_ar_us",
-            [make_fused_gemm_collective(x, w) for x, w in zip(inputs, weights)],
-        ),
-        Candidate(
-            "torch_gemm_us",
-            [make_torch_gemm(x, w) for x, w in zip(inputs, weights)],
-            check_correctness=False,
-        ),
-        Candidate(
-            "ring_ll_collective_us",
-            [make_ring_ll_collective(x) for x in collective_inputs],
-            check_correctness=False,
-        ),
-        Candidate(
-            "ldmc_collective_us",
-            [make_ldmc_collective(x) for x in symm_collective_inputs],
-            check_correctness=False,
-        ),
-    )
-
-    expected = candidates[0].runs[0]()
-    rows = M if all_reduce else valid_rows(M, local_M, rank)
-    for candidate in candidates[1:]:
-        if not candidate.check_correctness:
-            continue
-        actual = candidate.runs[0]()
-        torch.accelerator.synchronize(device)
-        torch.testing.assert_close(
-            actual[:rows],
-            expected[:rows],
-            rtol=5e-2,
-            atol=4.0,
-        )
-
-    candidate_graphs = {}
-    graph_keepalive: list[object] = [symm_partial_handle, *symm_collective_handles]
-    for candidate in candidates:
-        stream = torch.cuda.Stream()
-        bundles = [capture_graph(run, stream, cpu_group) for run in candidate.runs]
-        candidate_graphs[candidate.name] = [graph for graph, _ in bundles]
-        graph_keepalive.extend(bundles)
-        graph_keepalive.append(stream)
-
-    times = benchmark_graphs(
-        candidate_graphs,
-        warmup_replays,
-        samples,
-        device_group,
-        device_barrier,
-    )
-
-    best_nccl_us = min(times["ring_ll_collective_us"], times["ldmc_collective_us"])
-    return {
-        "mode": mode.upper(),
-        "M": M,
-        "N": N,
-        "K": K,
-        **times,
-        "best_nccl_us": best_nccl_us,
-        "speedup_vs_ring_ll": times["ring_ll_us"] / times["gemm_rs_ar_us"],
-        "speedup_vs_ldmc": times["ldmc_us"] / times["gemm_rs_ar_us"],
-    }
-
-
-def print_results(results: list[dict[str, float | int | str]]) -> None:
-    results_df = pd.DataFrame(results)
-    collective = str(results_df["mode"].iloc[0])
-    end_to_end = (
-        results_df[
-            [
-                "M",
-                "N",
-                "K",
-                "ring_ll_us",
-                "ldmc_us",
-                "gemm_rs_ar_us",
-                "speedup_vs_ring_ll",
-                "speedup_vs_ldmc",
-            ]
-        ]
-        .rename(
-            columns={
-                "ring_ll_us": f"Torch GEMM + NCCL {collective} (RING_LL) (us)",
-                "ldmc_us": f"Torch GEMM + NCCL {collective} (LDMC) (us)",
-                "gemm_rs_ar_us": f"GEMM-{collective} (us)",
-                "speedup_vs_ring_ll": "Speedup vs RING_LL",
-                "speedup_vs_ldmc": "Speedup vs LDMC",
-            }
-        )
-        .round(3)
-    )
-    components = (
-        results_df[["M", "N", "K", "torch_gemm_us", "best_nccl_us", "gemm_rs_ar_us"]]
-        .rename(
-            columns={
-                "torch_gemm_us": "Torch GEMM (us)",
-                "best_nccl_us": f"NCCL {collective} (best) (us)",
-                "gemm_rs_ar_us": f"GEMM-{collective} (us)",
-            }
-        )
-        .round(2)
-    )
-
-    print(f"### {collective} end-to-end latency")
-    print(end_to_end.to_markdown(index=False))
-    print(f"\n### {collective} component latency")
-    print(components.to_markdown(index=False))
-    print(
-        f"\nNCCL {collective} (best) is the faster of RING_LL and LDMC "
-        "for each shape.\n"
-    )
-
-
-def main() -> None:
-    args = parse_args()
-    assert args.m and min(args.m) > 0
-    assert args.n % 256 == 0
-    assert args.num_workspaces > 0
-    assert args.warmup_replays >= 0
-    assert args.samples > 0
-
-    local_rank = int(os.environ["LOCAL_RANK"])
-    torch.accelerator.set_device_index(local_rank)
-    init_distributed_environment()
-    world_size = dist.get_world_size()
-    if args.k is None:
-        assert all(K % world_size == 0 for K in _KIMI_K3_PROJECTION_K)
-        K_values = [K // world_size for K in _KIMI_K3_PROJECTION_K]
-    else:
-        K_values = args.k
-    assert all(K % 64 == 0 for K in K_values)
-    # Reserve symmetric memory for the NCCL-managed benchmark allocations.
-    os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
-    # NCCL-managed symmetric allocations select the NVLS/LDMC collective path.
-    symm_mem.set_backend("NCCL")
-    with set_current_vllm_config(VllmConfig()):
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
-
-    tp_group = get_tp_group()
-    group_warmup = torch.zeros(1, device=torch.accelerator.current_device_index())
-    dist.all_reduce(group_warmup, group=tp_group.device_group)
-    pynccl_comm = tp_group.device_communicator.pynccl_comm
-    assert pynccl_comm is not None
-    sync_input = torch.zeros(1, device=torch.accelerator.current_device_index())
-    sync_output = torch.empty_like(sync_input)
-
-    def device_barrier() -> None:
-        # Order the timed launch after a device-side rank rendezvous without
-        # including the rendezvous itself in the measured event interval.
-        pynccl_comm.all_reduce(sync_input, sync_output)
-
-    gemm_rs_ar = GemmRsAr(
-        max_M=max(args.m),
-        N=args.n,
-        all_reduce=args.mode == "ar",
-    )
-    results = [
-        benchmark_shape(
-            gemm_rs_ar,
-            args.mode,
-            M,
-            args.n,
-            K,
-            args.num_workspaces,
-            args.warmup_replays,
-            args.samples,
-            tp_group.device_group,
-            tp_group.cpu_group,
-            device_barrier,
-        )
-        for K in K_values
-        for M in args.m
-    ]
-    del gemm_rs_ar
-
-    if tp_group.rank_in_group == 0:
-        print_results(results)
-
-    dist.barrier(group=tp_group.cpu_group)
-    cleanup_dist_env_and_memory()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py b/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py
index fcce64015256..6b6d944455e9 100644
--- a/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py
+++ b/benchmarks/kernels/benchmark_kimi_k3_kda_decode.py
@@ -16,6 +16,7 @@
 Example:
     python benchmarks/kernels/benchmark_kimi_k3_kda_decode.py \
         --tokens 1 8 32 64 128 --heads 12
+
 """
 
 import argparse
diff --git a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py
index 2669dadb348d..855769263d84 100644
--- a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py
+++ b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py
@@ -27,6 +27,7 @@
 
 For multi-node runs, launch one ``torchrun`` agent per node and use a shared
 rendezvous endpoint.
+
 """
 
 from __future__ import annotations
diff --git a/benchmarks/kernels/benchmark_lora.py b/benchmarks/kernels/benchmark_lora.py
index ab930c59d219..606741d0474e 100644
--- a/benchmarks/kernels/benchmark_lora.py
+++ b/benchmarks/kernels/benchmark_lora.py
@@ -102,9 +102,7 @@ def make_rand_tensors(
     num_slices: int,
     device: str = "cuda",
 ) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]:
-    """
-    Make LoRA input/output matrices.
-    """
+    """Make LoRA input/output matrices."""
     A = torch.rand(a_shape, dtype=a_dtype).to(device)
 
     # LoRA weights column major
@@ -117,8 +115,7 @@ def make_rand_tensors(
 def make_prompt_lora_mapping(
     num_prompts: int, num_active_loras: int, sort_by_lora_id: bool, device: str
 ) -> torch.Tensor:
-    """
-    All prompts are mapped to a LoRA ID in range [0, num_active_loras).
+    """All prompts are mapped to a LoRA ID in range [0, num_active_loras).
     where 0 refers to first lora, 1 refers to second lora and so on.
     """
     assert num_active_loras > 0
@@ -147,9 +144,7 @@ def make_token_lora_mapping(
     seq_len_tensor: torch.Tensor,
     device: str,
 ):
-    """
-    Make token_lora_mapping from prompt_lora_mapping and seq_lens_tensor
-    """
+    """Make token_lora_mapping from prompt_lora_mapping and seq_lens_tensor."""
     assert prompt_lora_mapping.shape[0] == num_prompts
 
     # token to lora index mapping
@@ -174,8 +169,7 @@ def ref_group_gemm(
     scaling: float,
     add_inputs: bool | None,
 ):
-    """
-    Torch group gemm reference implementation to test correctness of
+    """Torch group gemm reference implementation to test correctness of
     benchmarking operations.
     """
     batches = seq_lens_cpu.size(0)
@@ -198,9 +192,7 @@ def ref_group_gemm(
 
 
 class OpType(Enum):
-    """
-    LoRA Ops to benchmark and its properties.
-    """
+    """LoRA Ops to benchmark and its properties."""
 
     LORA_SHRINK = auto()
     LORA_EXPAND = auto()
@@ -291,9 +283,7 @@ def mkn(
     def matmul_dtypes(
         self, op_dtype: torch.dtype
     ) -> tuple[torch.dtype, torch.dtype, torch.dtype]:
-        """
-        return a type, b type and c type for A x B = C
-        """
+        """Return a type, b type and c type for A x B = C."""
         if self.is_shrink_fn():
             return op_dtype, op_dtype, torch.float32
         elif self.is_expand_fn():
@@ -338,8 +328,7 @@ def matmul_shapes(
         top_k_num: int | None = None,
         num_experts: int | None = None,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        """
-        Given num_slices, return the shapes of the A, B, and C matrices
+        """Given num_slices, return the shapes of the A, B, and C matrices
         in A x B = C, for the op_type
         """
         m, k, n = self.mkn(batch_size, seq_length, hidden_size, lora_rank)
@@ -419,9 +408,7 @@ def run_ref_group_gemm(
 
 @dataclass
 class BenchmarkContext:
-    """
-    LoRA benchmark context
-    """
+    """LoRA benchmark context."""
 
     batch_size: int
     hidden_size: int
@@ -467,9 +454,7 @@ def bench_sublabel(self, op_type: OpType) -> str:
 
 @dataclass
 class BenchmarkTensors:
-    """
-    Input/Output tensors used for benchmarks
-    """
+    """Input/Output tensors used for benchmarks."""
 
     # matmul tensors
     input: torch.Tensor
@@ -555,9 +540,7 @@ def make(
         )
 
     def sanity_check(self, ctx: BenchmarkContext, op_type: OpType) -> None:
-        """
-        Fails asserts when non-conformality is detected.
-        """
+        """Fails asserts when non-conformality is detected."""
         num_tokens = (
             self.input.shape[1]
             if op_type.is_fused_moe_lora_expand_fn()
@@ -577,9 +560,7 @@ def sanity_check(self, ctx: BenchmarkContext, op_type: OpType) -> None:
         )
 
     def to_device(self, device: str):
-        """
-        Transfer tensors to device if the tensors aren't already on the device
-        """
+        """Transfer tensors to device if the tensors aren't already on the device."""
 
         def to_device(tensor: torch.Tensor):
             if tensor.device != device:
@@ -604,9 +585,7 @@ def to_device(tensor: torch.Tensor):
             )
 
     def metadata(self, ctx: BenchmarkContext, op_type: OpType) -> tuple[int, int, int]:
-        """
-        Return num_seqs, num_tokens and max_seq_len
-        """
+        """Return num_seqs, num_tokens and max_seq_len."""
         num_seqs = self.seq_lens.shape[0]
         num_tokens = self.get_num_tokens(
             self.lora_kernel_meta.token_lora_mapping.shape[0], ctx.top_k_num, op_type
@@ -630,8 +609,7 @@ def moe_lora_align_block_size(
             expert_map: torch.Tensor | None = None,
             pad_sorted_ids: bool = False,
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            """
-            Aligns tokens and experts into block-sized chunks for LoRA-based
+            """Aligns tokens and experts into block-sized chunks for LoRA-based
             mixture-of-experts (MoE) execution.
             """
             max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
@@ -959,8 +937,7 @@ def bench_fn_kwargs(
     def test_correctness(
         self, op_type: OpType, expand_fn_add_inputs: bool | None
     ) -> bool:
-        """
-        Test correctness of op_type implementation against a grouped gemm
+        """Test correctness of op_type implementation against a grouped gemm
         reference implementation.
         """
         seq_lens_cpu = self.seq_lens.to(device="cpu")
@@ -1070,15 +1047,13 @@ def bench_torch_mm(
     op_type: OpType,
     cuda_graph_nops: int | None = None,
 ) -> TMeasurement:
-    """
-    Benchmark basic torch.mm as a roofline.
+    """Benchmark basic torch.mm as a roofline.
 
     When all the input tokens have the same LoRA ID, the LoRA kernels are just
     a matmul. This torch.mm benchmark serves as a roofline for that case.
 
     input op_type is used in determining the m, k, n dimensions for the matmul.
     """
-
     batch_size, hidden_size, lora_rank, seq_length, dtype = (
         ctx.batch_size,
         ctx.hidden_size,
diff --git a/benchmarks/kernels/benchmark_marlin.py b/benchmarks/kernels/benchmark_marlin.py
index 751211993ad3..f2330a4ff0f0 100644
--- a/benchmarks/kernels/benchmark_marlin.py
+++ b/benchmarks/kernels/benchmark_marlin.py
@@ -6,10 +6,6 @@
 from benchmark_shapes import WEIGHT_SHAPES
 
 from vllm import _custom_ops as ops
-from vllm.model_executor.layers.quantization.utils.allspark_utils import (
-    ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
-    ALLSPARK_SUPPORTED_QUANT_TYPES,
-)
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     GPTQ_MARLIN_MAX_PARALLEL,
     GPTQ_MARLIN_MIN_THREAD_N,
@@ -31,7 +27,6 @@
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     gptq_pack,
     gptq_quantize_weights,
-    quantize_weights,
 )
 from vllm.scalar_type import ScalarType, scalar_types
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -62,9 +57,6 @@ def bench_run(
         return
 
     repack_supported = group_size in MARLIN_SUPPORTED_GROUP_SIZES
-    allspark_supported = (
-        quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES and group_size == -1
-    )
 
     def gen_marlin_params():
         # Marlin quant
@@ -104,35 +96,6 @@ def gen_repack_params():
             q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
         return q_w_gptq
 
-    def gen_allspark_params():
-        qw_reorder = s_reorder = zp_reorder = sm_count = sm_version = (
-            CUBLAS_M_THRESHOLD
-        ) = None
-        nonlocal allspark_supported
-        if allspark_supported:
-            properties = torch.cuda.get_device_properties(b.device.index)
-            sm_count = properties.multi_processor_count
-            sm_version = properties.major * 10 + properties.minor
-
-            supported_arch = sm_version >= 80 and sm_version < 90
-            allspark_supported = allspark_supported and supported_arch
-            if supported_arch:
-                w_ref, qw, s, zp = quantize_weights(b, quant_type, group_size, has_zp)
-                qw = qw.to(torch.uint8)
-
-                qw_reorder, s_reorder, zp_reorder = ops.allspark_repack_weight(
-                    qw, s, zp, has_zp
-                )
-                CUBLAS_M_THRESHOLD = ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD
-        return (
-            qw_reorder,
-            s_reorder,
-            zp_reorder,
-            sm_count,
-            sm_version,
-            CUBLAS_M_THRESHOLD,
-        )
-
     (
         marlin_w_ref,
         marlin_q_w,
@@ -141,9 +104,6 @@ def gen_allspark_params():
         marlin_zp,
     ) = gen_marlin_params()
     q_w_gptq = gen_repack_params()
-    qw_reorder, s_reorder, zp_reorder, sm_count, sm_version, CUBLAS_M_THRESHOLD = (
-        gen_allspark_params()
-    )
 
     # Prepare
     marlin_workspace = MarlinWorkspace(
@@ -167,17 +127,9 @@ def gen_allspark_params():
         "marlin_workspace": marlin_workspace,
         # GPTQ params
         "q_w_gptq": q_w_gptq,
-        # AllSpark W8A16 params
-        "qw_reorder": qw_reorder,
-        "s_reorder": s_reorder,
-        "zp_reorder": zp_reorder,
-        "sm_count": sm_count,
-        "sm_version": sm_version,
-        "CUBLAS_M_THRESHOLD": CUBLAS_M_THRESHOLD,
         # Kernels
         "marlin_gemm": ops.marlin_gemm,
         "gptq_marlin_repack": ops.gptq_marlin_repack,
-        "allspark_w8a16_gemm": ops.allspark_w8a16_gemm,
     }
 
     min_run_time = 1
@@ -227,17 +179,6 @@ def gen_allspark_params():
             ).blocked_autorange(min_run_time=min_run_time)
         )
 
-    if allspark_supported:
-        results.append(
-            benchmark.Timer(
-                stmt="output = allspark_w8a16_gemm(a, qw_reorder, s_reorder, zp_reorder, size_n, group_size, sm_count, sm_version, CUBLAS_M_THRESHOLD, False, True)",  # noqa: E501
-                globals=globals,
-                label=label,
-                sub_label=sub_label,
-                description="allspark_w8a16_gemm_fp32",
-            ).blocked_autorange(min_run_time=min_run_time)
-        )
-
 
 def main(args):
     print("Benchmarking models:")
diff --git a/benchmarks/kernels/benchmark_mla_k_concat.py b/benchmarks/kernels/benchmark_mla_k_concat.py
index 7debf3634804..3cae613e82c7 100644
--- a/benchmarks/kernels/benchmark_mla_k_concat.py
+++ b/benchmarks/kernels/benchmark_mla_k_concat.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation
+"""Benchmark script comparing torch.cat vs direct copy for k_nope/k_pe concatenation
 in MLA (Multi-head Latent Attention) prefill.
 
 This validates that the optimization from commit 8d4142bd is beneficial across
diff --git a/benchmarks/kernels/benchmark_moe_align_block_size.py b/benchmarks/kernels/benchmark_moe_align_block_size.py
index a340500379a0..7ddf5b3ecf7c 100644
--- a/benchmarks/kernels/benchmark_moe_align_block_size.py
+++ b/benchmarks/kernels/benchmark_moe_align_block_size.py
@@ -5,6 +5,7 @@
 
 import torch
 
+import vllm._custom_ops as ops
 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
     moe_align_block_size,
 )
@@ -68,8 +69,87 @@ def benchmark(num_tokens, num_experts, topk, ep_size, provider):
     return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
+batched_configs = list(
+    itertools.product(
+        [8, 32, 64, 256],
+        [16, 256, 512, 2048],
+        [8, 32, 48, 64],
+        ["uniform", "skewed", "mostly_empty"],
+    )
+)
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=[
+            "num_batches",
+            "max_tokens_per_batch",
+            "block_size",
+            "distribution",
+        ],
+        x_vals=batched_configs,
+        line_arg="provider",
+        line_vals=["vllm"],
+        line_names=["vLLM"],
+        plot_name="batched-moe-align-block-size-performance",
+        args={},
+    )
+)
+def benchmark_batched(
+    num_batches,
+    max_tokens_per_batch,
+    block_size,
+    distribution,
+    provider,
+):
+    set_random_seed(0)
+    if distribution == "uniform":
+        expert_num_tokens = torch.full(
+            (num_batches,),
+            max_tokens_per_batch,
+            dtype=torch.int32,
+            device="cuda",
+        )
+    elif distribution == "skewed":
+        expert_num_tokens = torch.ones((num_batches,), dtype=torch.int32, device="cuda")
+        expert_num_tokens[0] = max_tokens_per_batch
+    else:
+        expert_num_tokens = torch.zeros(
+            (num_batches,), dtype=torch.int32, device="cuda"
+        )
+        expert_num_tokens[::8] = max_tokens_per_batch
+
+    max_num_tokens_padded = (
+        num_batches
+        * ((max_tokens_per_batch + block_size - 1) // block_size)
+        * block_size
+    )
+    sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda")
+    expert_ids = torch.empty(
+        (max_num_tokens_padded // block_size,), dtype=torch.int32, device="cuda"
+    )
+    num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device="cuda")
+
+    quantiles = [0.5, 0.2, 0.8]
+    if provider == "vllm":
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: ops.batched_moe_align_block_size(
+                max_tokens_per_batch,
+                block_size,
+                expert_num_tokens,
+                sorted_ids,
+                expert_ids,
+                num_tokens_post_pad,
+            ),
+            quantiles=quantiles,
+        )
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
+    parser.add_argument("--batched", action="store_true")
     parser.add_argument(
         "--num_experts",
         type=int,
@@ -85,4 +165,5 @@ def benchmark(num_tokens, num_experts, topk, ep_size, provider):
     )
     args = parser.parse_args()
 
-    benchmark.run(print_data=True, show_plots=True)
+    selected_benchmark = benchmark_batched if args.batched else benchmark
+    selected_benchmark.run(print_data=True, show_plots=True)
diff --git a/benchmarks/kernels/benchmark_moe_defaults.py b/benchmarks/kernels/benchmark_moe_defaults.py
index f6ad59366dca..00c6606db009 100644
--- a/benchmarks/kernels/benchmark_moe_defaults.py
+++ b/benchmarks/kernels/benchmark_moe_defaults.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark comparing old vs new default fused MoE configs.
+"""Benchmark comparing old vs new default fused MoE configs.
 
 Runs the triton fused_moe kernel with three configurations for each scenario:
   1. Tuned config (from JSON file, if available) — the target to match
diff --git a/benchmarks/kernels/benchmark_per_token_group_quant.py b/benchmarks/kernels/benchmark_per_token_group_quant.py
index f2195a6d780b..fda3d2a652b4 100644
--- a/benchmarks/kernels/benchmark_per_token_group_quant.py
+++ b/benchmarks/kernels/benchmark_per_token_group_quant.py
@@ -15,7 +15,7 @@
 
 @contextmanager
 def _triton_mode():
-    """Temporarily force the Triton fallback path"""
+    """Temporarily force the Triton fallback path."""
     with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
         yield
 
diff --git a/benchmarks/kernels/benchmark_per_token_quant_fp8.py b/benchmarks/kernels/benchmark_per_token_quant_fp8.py
index 6ce97e30368b..0069b06c6456 100644
--- a/benchmarks/kernels/benchmark_per_token_quant_fp8.py
+++ b/benchmarks/kernels/benchmark_per_token_quant_fp8.py
@@ -16,7 +16,7 @@
 
 
 def with_triton_mode(fn):
-    """Temporarily force the Triton fallback path"""
+    """Temporarily force the Triton fallback path."""
 
     def wrapped(*args, **kwargs):
         with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
@@ -124,8 +124,7 @@ def compute_geomean_speedups(
     speedup_cols: list[str],
     groupby_cols: list[str] | None = None,
 ) -> pd.DataFrame:
-    """
-    Compute geometric mean speedups over a baseline column.
+    """Compute geometric mean speedups over a baseline column.
 
     Args:
         df: Input dataframe
@@ -135,6 +134,7 @@ def compute_geomean_speedups(
 
     Returns:
         pd.DataFrame with geometric mean speedups
+
     """
     from scipy.stats import gmean
 
diff --git a/benchmarks/kernels/benchmark_persistent_topk.py b/benchmarks/kernels/benchmark_persistent_topk.py
new file mode 100644
index 000000000000..a77a04cfd625
--- /dev/null
+++ b/benchmarks/kernels/benchmark_persistent_topk.py
@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark sparse-indexer decode top-k on variable-length FP32 rows.
+
+Run on a reserved GPU, for example:
+    chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py
+
+Use --capacity to emulate graphs captured at a larger maximum context length.
+Run the same command on the base and candidate builds to compare JSON results.
+"""
+
+import argparse
+import itertools
+import json
+import statistics
+from pathlib import Path
+
+import torch
+from flashinfer.testing import bench_gpu_time_with_cupti
+
+import vllm._custom_ops  # noqa: F401
+from vllm.platforms import current_platform
+
+
+def benchmark(rows: int, width: int, k: int, args) -> dict:
+    capacity = args.capacity or width
+    if capacity < width:
+        raise ValueError("Capacity must be at least the valid-length upper bound")
+    torch.manual_seed(args.seed)
+    logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32)
+    if args.full_length:
+        lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32)
+    else:
+        lengths = torch.randint(
+            int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32
+        )
+    invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None]
+    logits.masked_fill_(invalid, float("nan"))
+    out = torch.empty((rows, k), device="cuda", dtype=torch.int32)
+    workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8)
+    cooperative = (
+        args.backend == "auto"
+        and rows <= 64
+        and capacity % 4 == 0
+        and current_platform.has_device_capability(90)
+        and not current_platform.is_device_capability_family(120)
+    )
+    op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk
+
+    def run():
+        op(logits, lengths, out, workspace, k, capacity)
+
+    run()
+    valid = torch.arange(k, device="cuda")[None] < lengths[:, None]
+    assert torch.all(out[~valid] == -1)
+    assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid)
+    sorted_indices = out.sort(dim=1).values
+    assert torch.all(
+        (sorted_indices[:, 1:] != sorted_indices[:, :-1])
+        | (sorted_indices[:, 1:] == -1)
+    )
+    selected = logits.gather(1, out.clamp_min(0).long()).masked_fill(
+        ~valid, -float("inf")
+    )
+    reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values
+    torch.testing.assert_close(
+        selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0
+    )
+    trials = [
+        statistics.median(
+            bench_gpu_time_with_cupti(
+                run,
+                use_cuda_graph=True,
+                cold_l2_cache=not args.warm,
+                dry_run_iters=5,
+                repeat_iters=args.repeat,
+            )
+        )
+        * 1000
+        for _ in range(args.trials)
+    ]
+    us = statistics.median(trials)
+    return dict(
+        rows=rows,
+        width=width,
+        capacity=capacity,
+        k=k,
+        backend="cooperative" if cooperative else "persistent",
+        us=us,
+        trials_us=trials,
+        # Minimum useful traffic: one valid-score read and one index write.
+        effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024])
+    parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840])
+    parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048])
+    parser.add_argument("--capacity", type=int)
+    parser.add_argument(
+        "--full-length", action="store_true", help="Use exactly --widths valid scores"
+    )
+    parser.add_argument("--backend", choices=["auto", "persistent"], default="auto")
+    parser.add_argument(
+        "--warm", action="store_true", help="Keep L2 warm between replays"
+    )
+    parser.add_argument("--repeat", type=int, default=60)
+    parser.add_argument("--trials", type=int, default=3)
+    parser.add_argument("--seed", type=int, default=123)
+    parser.add_argument("--output", type=Path)
+    args = parser.parse_args()
+    metadata = dict(
+        gpu=str(torch.cuda.get_device_properties(0)),
+        torch=torch.__version__,
+        cuda=torch.version.cuda,
+        args=vars(args),
+    )
+    print(json.dumps(metadata, default=str))
+    results = []
+    for rows, width, k in itertools.product(args.rows, args.widths, args.top_k):
+        result = benchmark(rows, width, k, args)
+        results.append(result)
+        print(json.dumps(result), flush=True)
+    if args.output:
+        args.output.write_text(
+            json.dumps(dict(metadata=metadata, results=results), indent=2, default=str)
+            + "\n"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py
index 4ecd917f31ac..2d5ec80c6aae 100644
--- a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py
+++ b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes.
+"""Benchmark the RDNAHybridW4A16LinearKernel across decode and prefill shapes.
 
 Usage:
     python benchmark_int4_gemm.py
diff --git a/benchmarks/kernels/benchmark_reshape_and_cache.py b/benchmarks/kernels/benchmark_reshape_and_cache.py
index 97af4ac976ee..f49bf61036d0 100644
--- a/benchmarks/kernels/benchmark_reshape_and_cache.py
+++ b/benchmarks/kernels/benchmark_reshape_and_cache.py
@@ -32,7 +32,6 @@ def run_benchmark(
     device: str = "cuda",
 ) -> float:
     """Return latency (seconds) for given num_tokens."""
-
     if kv_cache_dtype == "fp8" and head_size % 16:
         raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.")
 
diff --git a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py
index 55c203725186..84d9cc707c40 100644
--- a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py
+++ b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py
@@ -37,7 +37,6 @@ def run_benchmark(
     device: str = "cuda",
 ) -> float:
     """Return latency (seconds) for given num_tokens."""
-
     if kv_cache_dtype == "fp8" and head_size % 16:
         raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.")
 
diff --git a/benchmarks/kernels/benchmark_rmsnorm.py b/benchmarks/kernels/benchmark_rmsnorm.py
deleted file mode 100644
index d8d7f5bcf9da..000000000000
--- a/benchmarks/kernels/benchmark_rmsnorm.py
+++ /dev/null
@@ -1,255 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import itertools
-
-import torch
-from flashinfer.norm import fused_add_rmsnorm, rmsnorm
-from torch import nn
-
-from vllm import _custom_ops as vllm_ops
-from vllm.triton_utils import triton
-
-
-class HuggingFaceRMSNorm(nn.Module):
-    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor | None = None,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        orig_dtype = x.dtype
-        x = x.to(torch.float32)
-        if residual is not None:
-            x = x + residual.to(torch.float32)
-            residual = x.to(orig_dtype)
-
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
-        x = x * torch.rsqrt(variance + self.variance_epsilon)
-        x = x.to(orig_dtype) * self.weight
-        if residual is None:
-            return x
-        else:
-            return x, residual
-
-
-def rmsnorm_naive(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    residual: torch.Tensor | None = None,
-    eps: float = 1e-6,
-):
-    naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps)
-    naive_norm.weight = nn.Parameter(weight)
-    naive_norm = naive_norm.to(x.device)
-
-    orig_shape = x.shape
-    x = x.view(-1, x.shape[-1])
-    if residual is not None:
-        residual = residual.view(-1, residual.shape[-1])
-
-    output = naive_norm(x, residual)
-
-    if isinstance(output, tuple):
-        output = (output[0].view(orig_shape), output[1].view(orig_shape))
-    else:
-        output = output.view(orig_shape)
-    return output
-
-
-def rmsnorm_flashinfer(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    residual: torch.Tensor | None = None,
-    eps: float = 1e-6,
-):
-    orig_shape = x.shape
-    x = x.view(-1, x.shape[-1])
-    if residual is not None:
-        residual = residual.view(-1, residual.shape[-1])
-
-    if residual is not None:
-        fused_add_rmsnorm(x, residual, weight, eps)
-        output = (x, residual)
-    else:
-        output = rmsnorm(x, weight, eps)
-
-    if isinstance(output, tuple):
-        output = (output[0].view(orig_shape), output[1].view(orig_shape))
-    else:
-        output = output.view(orig_shape)
-    return output
-
-
-def rmsnorm_vllm(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    residual: torch.Tensor | None = None,
-    eps: float = 1e-6,
-):
-    orig_shape = x.shape
-    x = x.view(-1, x.shape[-1])
-    if residual is not None:
-        residual = residual.view(-1, residual.shape[-1])
-
-    if residual is not None:
-        vllm_ops.fused_add_rms_norm(x, residual, weight, eps)
-        output = (x, residual)
-    else:
-        out = torch.empty_like(x)
-        vllm_ops.rms_norm(out, x, weight, eps)
-        output = out
-
-    if isinstance(output, tuple):
-        output = (output[0].view(orig_shape), output[1].view(orig_shape))
-    else:
-        output = output.view(orig_shape)
-    return output
-
-
-def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True):
-    dtype = torch.bfloat16
-    x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda")
-    weight = torch.ones(hidden_size, dtype=dtype, device="cuda")
-    residual = torch.randn_like(x) if use_residual else None
-
-    output_naive = rmsnorm_naive(
-        x.clone(), weight, residual.clone() if residual is not None else None
-    )
-    output_flashinfer = rmsnorm_flashinfer(
-        x.clone(), weight, residual.clone() if residual is not None else None
-    )
-    output_vllm = rmsnorm_vllm(
-        x.clone(), weight, residual.clone() if residual is not None else None
-    )
-
-    if use_residual:
-        output_naive = output_naive[0]
-        output_flashinfer = output_flashinfer[0]
-        output_vllm = output_vllm[0]
-
-    print(f"Naive output={output_naive}")
-    print(f"FlashInfer output={output_flashinfer}")
-    print(f"vLLM output={output_vllm}")
-
-    if torch.allclose(
-        output_naive, output_flashinfer, atol=1e-2, rtol=1e-2
-    ) and torch.allclose(output_naive, output_vllm, atol=1e-2, rtol=1e-2):
-        print("✅ All implementations match")
-    else:
-        print("❌ Implementations differ")
-
-
-batch_size_range = [2**i for i in range(0, 7, 2)]
-seq_length_range = [2**i for i in range(6, 11, 1)]
-head_num_range = [32, 48]
-configs = list(itertools.product(head_num_range, batch_size_range, seq_length_range))
-
-
-def get_benchmark(use_residual):
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["head_num", "batch_size", "seq_len"],
-            x_vals=[list(_) for _ in configs],
-            line_arg="provider",
-            line_vals=["huggingface", "flashinfer", "vllm"],
-            line_names=["HuggingFace", "FlashInfer", "vLLM"],
-            styles=[("blue", "-"), ("green", "-"), ("red", "-")],
-            ylabel="us",
-            plot_name=f"rmsnorm-perf-{'with' if use_residual else 'without'}-residual",
-            args={},
-        )
-    )
-    def benchmark(head_num, batch_size, seq_len, provider):
-        dtype = torch.bfloat16
-        hidden_size = head_num * 128  # assuming head_dim = 128
-
-        x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda")
-        weight = torch.ones(hidden_size, dtype=dtype, device="cuda")
-        residual = torch.randn_like(x) if use_residual else None
-
-        quantiles = [0.5, 0.2, 0.8]
-
-        if provider == "huggingface":
-            ms, min_ms, max_ms = triton.testing.do_bench(
-                lambda: rmsnorm_naive(
-                    x.clone(),
-                    weight,
-                    residual.clone() if residual is not None else None,
-                ),
-                quantiles=quantiles,
-            )
-        elif provider == "flashinfer":
-            ms, min_ms, max_ms = triton.testing.do_bench(
-                lambda: rmsnorm_flashinfer(
-                    x.clone(),
-                    weight,
-                    residual.clone() if residual is not None else None,
-                ),
-                quantiles=quantiles,
-            )
-        else:
-            ms, min_ms, max_ms = triton.testing.do_bench(
-                lambda: rmsnorm_vllm(
-                    x.clone(),
-                    weight,
-                    residual.clone() if residual is not None else None,
-                ),
-                quantiles=quantiles,
-            )
-
-        return 1000 * ms, 1000 * max_ms, 1000 * min_ms
-
-    return benchmark
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=4,
-        help="Batch size",
-    )
-    parser.add_argument(
-        "--seq-len",
-        type=int,
-        default=128,
-        help="Sequence length",
-    )
-    parser.add_argument(
-        "--hidden-size",
-        type=int,
-        default=4096,
-        help="Hidden size (2nd dimension) of the sequence",
-    )
-    parser.add_argument(
-        "--use-residual", action="store_true", help="Whether to use residual connection"
-    )
-    parser.add_argument(
-        "--save-path",
-        type=str,
-        default="./configs/rmsnorm/",
-        help="Path to save rmsnorm benchmark results",
-    )
-
-    args = parser.parse_args()
-
-    # Run correctness test
-    calculate_diff(
-        batch_size=args.batch_size,
-        seq_len=args.seq_len,
-        hidden_size=args.hidden_size,
-        use_residual=args.use_residual,
-    )
-
-    # Get the benchmark function with proper use_residual setting
-    benchmark = get_benchmark(args.use_residual)
-    # Run performance benchmark
-    benchmark.run(print_data=True, save_path=args.save_path)
diff --git a/benchmarks/kernels/benchmark_selective_state_update.py b/benchmarks/kernels/benchmark_selective_state_update.py
index ddf1efd018c8..476dce7dff8d 100644
--- a/benchmarks/kernels/benchmark_selective_state_update.py
+++ b/benchmarks/kernels/benchmark_selective_state_update.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark and tuning script for the Mamba selective_state_update kernel.
+"""Benchmark and tuning script for the Mamba selective_state_update kernel.
 
 Mirrors the fused MoE tuning workflow: sweeps (BLOCK_SIZE_M, num_warps) across
 an effective_batch grid for a given (headdim, dstate, ngroups, cache_dtype) and
@@ -124,8 +123,7 @@ def benchmark_config(
     num_warmup: int = 20,
     graph_batch_size: int = 10,
 ) -> float | None:
-    """
-    Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
+    """Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
     Returns elapsed time in microseconds, or None on error.
 
     Uses accelerator graph capture-and-replay to isolate kernel time from
@@ -356,8 +354,7 @@ def validate_configs(
     rtol: float = 1e-2,
     state_dtype: torch.dtype | None = None,
 ) -> dict[int, bool]:
-    """
-    For every (effective_batch, batch, nheads) in *active* that has a tuned
+    """For every (effective_batch, batch, nheads) in *active* that has a tuned
     config, run the kernel with that config and compare against the reference.
     Returns {effective_batch: passed}.
     """
diff --git a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
index 9fcf278f2ef3..22ef378b0763 100644
--- a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
+++ b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Comprehensive 3-way SiLU Benchmark Suite
+"""Comprehensive 3-way SiLU Benchmark Suite.
 
 This benchmark compares three SiLU implementations:
 1. SiLU V2 (CUDA) - Optimized CUDA kernel implementation
@@ -110,7 +109,7 @@ def silu_mul_fp8_quant_deep_gemm_triton(
     eps: float = 1e-10,
     expert_offsets: torch.Tensor = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
+    """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales.
 
     y has shape (E, T, 2*H). The first half of the last dimension is
     silu-activated, multiplied by the second half, then quantized into FP8.
@@ -206,7 +205,7 @@ def benchmark(
     iterations_per_run: int = 20,
 ):
     def generate_data(seed_offset=0):
-        """Generate input data with given seed offset"""
+        """Generate input data with given seed offset."""
         set_random_seed(42 + seed_offset)
         y = torch.rand((E, T, 2 * H), dtype=torch.bfloat16, device="cuda").contiguous()
 
diff --git a/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py b/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py
index 6bb504b7f090..29ff156b6035 100644
--- a/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py
+++ b/benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py
@@ -10,6 +10,7 @@
 Example:
     python benchmarks/kernels/benchmark_vit_aiter_fp8_attn.py \
         --seq-lens 2304 4096 8192 16384 --head-dim 72
+
 """
 
 from types import SimpleNamespace
diff --git a/benchmarks/kernels/benchmark_w8a8_block_fp8.py b/benchmarks/kernels/benchmark_w8a8_block_fp8.py
index 590d4cfdc6d9..3ad7ec888728 100644
--- a/benchmarks/kernels/benchmark_w8a8_block_fp8.py
+++ b/benchmarks/kernels/benchmark_w8a8_block_fp8.py
@@ -57,10 +57,12 @@ def w8a8_block_matmul(
         Bs: The per-block quantization scale for `B`.
         block_size: The block size for per-block quantization.
                     It should be 2-dim, e.g., [128, 128].
+        config: Triton kernel launch config, e.g. tile sizes and num_warps.
         output_dtype: The dtype of the returned tensor.
 
     Returns:
         torch.Tensor: The result of matmul.
+
     """
     assert len(block_size) == 2
     block_n, block_k = block_size[0], block_size[1]
diff --git a/benchmarks/kernels/ir/bench_ir_ops.py b/benchmarks/kernels/ir/bench_ir_ops.py
index b23c4e8ae327..80fdd87acb21 100644
--- a/benchmarks/kernels/ir/bench_ir_ops.py
+++ b/benchmarks/kernels/ir/bench_ir_ops.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Generic benchmark harness for vLLM IR ops.
+"""Generic benchmark harness for vLLM IR ops.
 
 Usage:
     python benchmarks/kernels/ir/bench_ir_ops.py
diff --git a/benchmarks/kernels/ir/shapes.py b/benchmarks/kernels/ir/shapes.py
index 6cc44cf6cec1..40f5f8947621 100644
--- a/benchmarks/kernels/ir/shapes.py
+++ b/benchmarks/kernels/ir/shapes.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shape configurations for IR op benchmarks.
-"""
+"""Shape configurations for IR op benchmarks."""
 
 import torch
 
@@ -20,10 +18,11 @@
 # Each entry maps an op name to a list of kwarg dicts that will be passed
 # to that op's registered input generator via op.generate_inputs(**kwargs).
 SHAPE_CONFIGS: dict[str, list[dict]] = {
-    "rms_norm": [
+    op_name: [
         {"num_tokens": n, "hidden_size": d, "dtype": dtype}
         for dtype in [torch.float16, torch.bfloat16, torch.float32]
         for d in COMMON_HIDDEN_SIZES
         for n in NUM_TOKENS
-    ],
+    ]
+    for op_name in ("rms_norm", "fused_add_rms_norm")
 }
diff --git a/benchmarks/kernels/utils.py b/benchmarks/kernels/utils.py
index a9af811bbe9c..64d23691bdbb 100644
--- a/benchmarks/kernels/utils.py
+++ b/benchmarks/kernels/utils.py
@@ -17,8 +17,7 @@ class CudaGraphBenchParams:
 
 @dataclasses.dataclass
 class ArgPool:
-    """
-    When some argument of the benchmarking function is annotated with this type,
+    """When some argument of the benchmarking function is annotated with this type,
     the benchmarking class (BenchMM) will collapse the argument to a pick a
     single value from the given list of values, during function invocation.
     For every invocation during a benchmarking run, it will choose a
diff --git a/benchmarks/multi_turn/bench_dataset.py b/benchmarks/multi_turn/bench_dataset.py
index 8cb8a2f386a9..9b159a9ee999 100644
--- a/benchmarks/multi_turn/bench_dataset.py
+++ b/benchmarks/multi_turn/bench_dataset.py
@@ -148,8 +148,7 @@ def __init__(
     def _generate_lognormal_by_median(
         target_average: int, median_ratio: float
     ) -> tuple[float, float]:
-        """
-        Compute (mu, sigma) for a lognormal distribution given:
+        """Compute (mu, sigma) for a lognormal distribution given:
         - a target average (mean of the distribution)
         - a ratio of median / mean (controls skewness), assume mean > median
 
diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py
index 5a60d9c66883..f084a4d84c5f 100644
--- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py
+++ b/benchmarks/multi_turn/benchmark_serving_multi_turn.py
@@ -32,8 +32,11 @@
 from bench_utils import TEXT_SEPARATOR, Color, logger
 from transformers import AutoTokenizer  # type: ignore
 
+from vllm.benchmarks.lib.utils import redact_sensitive_namespace
+
 NUM_TOKENS_FROM_DATASET = 0
 TERM_SIGNAL = None
+_SENSITIVE_ARG_FIELDS = ("api_key", "header")
 
 
 class ConversationSampling(str, Enum):
@@ -1548,7 +1551,7 @@ async def main() -> None:
 
     args = parser.parse_args()
 
-    logger.info(args)
+    logger.info(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS))
 
     logger.info(f"{Color.GREEN}Input parameters:{Color.RESET}")
     logger.info(f"url={args.url}")
diff --git a/benchmarks/multi_turn/convert_sharegpt_to_openai.py b/benchmarks/multi_turn/convert_sharegpt_to_openai.py
index fccab4d0ce21..b415744873b8 100644
--- a/benchmarks/multi_turn/convert_sharegpt_to_openai.py
+++ b/benchmarks/multi_turn/convert_sharegpt_to_openai.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Download dataset from:
+"""Download dataset from:
 https://huggingface.co/datasets/philschmid/sharegpt-raw/blob/main/sharegpt_20230401_clean_lang_split.json
 
 Convert to OpenAI API:
diff --git a/benchmarks/overheads/benchmark_hashing.py b/benchmarks/overheads/benchmark_hashing.py
deleted file mode 100644
index 178599952d5c..000000000000
--- a/benchmarks/overheads/benchmark_hashing.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import cProfile
-import pstats
-
-from vllm import LLM, SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-# A very long prompt, total number of tokens is about 15k.
-LONG_PROMPT = ["You are an expert in large language models, aren't you?"] * 1000
-LONG_PROMPT = " ".join(LONG_PROMPT)
-
-
-def main(args):
-    llm = LLM(
-        model=args.model,
-        enforce_eager=True,
-        enable_prefix_caching=True,
-        tensor_parallel_size=args.tensor_parallel_size,
-    )
-
-    sampling_params = SamplingParams(temperature=0, max_tokens=args.output_len)
-    profiler = cProfile.Profile()
-
-    print("------warm up------")
-    for i in range(3):
-        output = llm.generate(LONG_PROMPT, sampling_params)
-        print(output[0].outputs[0].text)
-
-    print("------start generating------")
-    for i in range(3):
-        profiler.runctx(
-            "llm.generate(LONG_PROMPT, sampling_params)", globals(), locals()
-        )
-
-    # analyze the runtime of hashing function
-    stats = pstats.Stats(profiler)
-    stats.sort_stats("cumulative")
-    total_time = 0
-    total_calls = 0
-    for func in stats.stats:
-        if "hash_of_block" in func[2]:
-            total_time = stats.stats[func][3]
-            total_calls = stats.stats[func][0]
-    percentage = (total_time / stats.total_tt) * 100
-    print(
-        f"Hashing took {total_time:.2f} seconds,{percentage:.2f}% of the total runtime."
-    )
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser(
-        description="Benchmark the performance of hashing function in"
-        "automatic prefix caching."
-    )
-    parser.add_argument("--model", type=str, default="lmsys/longchat-7b-16k")
-    parser.add_argument("--tensor-parallel-size", "-tp", type=int, default=1)
-    parser.add_argument("--output-len", type=int, default=10)
-    parser.add_argument(
-        "--enable-prefix-caching", action="store_true", help="enable prefix caching"
-    )
-    args = parser.parse_args()
-    main(args)
diff --git a/benchmarks/replayssm/e2e_decode_speedup.py b/benchmarks/replayssm/e2e_decode_speedup.py
index b7084a21315b..ff7c30669a12 100644
--- a/benchmarks/replayssm/e2e_decode_speedup.py
+++ b/benchmarks/replayssm/e2e_decode_speedup.py
@@ -15,6 +15,7 @@
     python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
     python e2e_decode_speedup.py --dtype auto --buffer-len 16 \
         --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4   # B300 NVFP4
+
 """
 
 import argparse
diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake
index ff2911dc11c8..414fbc5e238c 100644
--- a/cmake/cpu_extension.cmake
+++ b/cmake/cpu_extension.cmake
@@ -173,6 +173,18 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)
         set(AMX_FP8_SUPPORTED FALSE)
         message(STATUS "AMX-FP8 disabled by ENABLE_AMX_FP8=OFF")
     endif()
+
+    # AVX10.2 arrived in GCC 15, so the target attribute, the fp8 convert
+    # intrinsics and __builtin_cpu_supports("avx10.2") are all hard errors on
+    # the 12.3+ the backend otherwise supports. Probe instead of assuming.
+    check_cxx_compiler_flag("-mavx10.2" COMPILER_SUPPORTS_AVX10_2_FLAG)
+    if (COMPILER_SUPPORTS_AVX10_2_FLAG)
+        set(AVX10_2_SUPPORTED TRUE)
+    else()
+        set(AVX10_2_SUPPORTED FALSE)
+        message(STATUS "AVX10.2 disabled: compiler does not support -mavx10.2 (compiler: ${CMAKE_CXX_COMPILER_ID} ${CMAKE_CXX_COMPILER_VERSION})")
+    endif()
+
     list(APPEND CXX_COMPILE_FLAGS_AVX2
         "-mavx2")
 elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
@@ -488,6 +500,7 @@ set(VLLM_EXT_SRC
     "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
     "csrc/cpu/cpu_fused_moe.cpp"
     "csrc/cpu/cpu_attn.cpp"
+    "csrc/cpu/cpu_isa.cpp"
     "csrc/cpu/torch_bindings.cpp")
 
 if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
@@ -522,6 +535,12 @@ if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
         ${VLLM_EXT_SRC})
 endif()
 
+if (POWER10_FOUND OR POWER11_FOUND)
+    set(VLLM_EXT_SRC
+        "csrc/cpu/cpu_fused_moe_int8.cpp"
+        ${VLLM_EXT_SRC})
+endif()
+
 if(USE_ONEDNN)
     set(VLLM_EXT_SRC
         "csrc/cpu/dnnl_kernels.cpp"
@@ -540,11 +559,13 @@ if (ENABLE_X86_ISA)
         "csrc/cpu/sgl-kernels/gemm.cpp"
         "csrc/cpu/sgl-kernels/gemm_int8.cpp"
         "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
+        "csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp"
         "csrc/cpu/sgl-kernels/gemm_int4.cpp"
         "csrc/cpu/sgl-kernels/moe.cpp"
         "csrc/cpu/sgl-kernels/moe_int8.cpp"
         "csrc/cpu/sgl-kernels/moe_int4.cpp"
         "csrc/cpu/sgl-kernels/moe_fp8.cpp"
+        "csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp"
         "csrc/cpu/sgl-kernels/bmm.cpp"
         "csrc/cpu/sgl-kernels/decode.cpp"
         "csrc/cpu/sgl-kernels/extend.cpp"
@@ -565,6 +586,7 @@ if (ENABLE_X86_ISA)
         "csrc/cpu/utils.cpp"
         "csrc/cpu/spec_decode_utils.cpp"
         "csrc/cpu/cpu_attn.cpp"
+        "csrc/cpu/cpu_isa.cpp"
         "csrc/cpu/dnnl_kernels.cpp"
         "csrc/cpu/mamba_cpu.cpp"
         "csrc/cpu/torch_bindings.cpp"
@@ -581,6 +603,7 @@ if (ENABLE_X86_ISA)
         "csrc/cpu/utils.cpp"
         "csrc/cpu/spec_decode_utils.cpp"
         "csrc/cpu/cpu_attn.cpp"
+        "csrc/cpu/cpu_isa.cpp"
         "csrc/cpu/mamba_cpu.cpp"
         "csrc/cpu/dnnl_kernels.cpp"
         "csrc/cpu/torch_bindings.cpp"
@@ -618,6 +641,12 @@ if (ENABLE_X86_ISA)
         message(STATUS "AMX-FP8 (Diamond Rapids) enabled")
     endif()
 
+    # For the sgl-kernels AVX10.2 fp8 quantize paths
+    if (AVX10_2_SUPPORTED)
+        target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AVX10_2")
+        message(STATUS "AVX10.2 enabled")
+    endif()
+
     # AVX512F 
     define_extension_target(
         _C_AVX512
diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake
index 72f7f5eb66a8..a17031cc5e0f 100644
--- a/cmake/external_projects/deepgemm.cmake
+++ b/cmake/external_projects/deepgemm.cmake
@@ -28,9 +28,12 @@ if(DEEPGEMM_SRC_DIR)
   message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
 else()
   # Keep in sync with tools/install_deepgemm.sh
-  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
-  # Pinned to the tip of the nv_dev branch (SM120 support).
-  set(_DEEPGEMM_UPSTREAM_TAG "8b1392b978f5a03c828dd1711090d7fb50958b8a")
+  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
+  # Pinned to the fork's dev branch: upstream 2.8.0 plus the SM120
+  # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU
+  # alpha/beta, and SiTU for FP8/FP4 Mega MoE. Also includes the CUDA 12.x
+  # layout header fix from vllm-project/DeepGEMM#12.
+  set(_DEEPGEMM_UPSTREAM_TAG "e1f418c2a4f20818221f6b0e578b4c2f634d4c3f")
 
   set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _deepgemm_fc_root)
@@ -51,7 +54,7 @@ else()
       BINARY_DIR "${_deepgemm_bin}"
       GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}"
       GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}"
-      GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
+      GIT_SUBMODULES "third-party/cutlass" "third-party/deep_jit"
       GIT_PROGRESS TRUE
     )
   endif()
@@ -119,13 +122,15 @@ if(DEEPGEMM_ARCHS)
   message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
 
   # add_custom_command does no implicit header scanning; glob explicitly so
-  # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild.
+  # header-only edits in DeepGEMM/cutlass/deep_jit re-trigger the rebuild.
   file(GLOB_RECURSE _dg_headers
     "${deepgemm_SOURCE_DIR}/csrc/*.h"
     "${deepgemm_SOURCE_DIR}/csrc/*.hpp"
     "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
     "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
-    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")
+    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh"
+    "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.h"
+    "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.hpp")
 
   set(_dg_markers)
   set(_dg_seen_soabis)
diff --git a/cmake/external_projects/deepselect.cmake b/cmake/external_projects/deepselect.cmake
new file mode 100644
index 000000000000..e2a5592e6167
--- /dev/null
+++ b/cmake/external_projects/deepselect.cmake
@@ -0,0 +1,97 @@
+include(FetchContent)
+
+# If DEEPSELECT_SRC_DIR is set, DeepSelect is built from that directory
+# instead of downloading.
+# It can be set as an environment variable or passed as a cmake argument.
+# The environment variable takes precedence.
+if (DEFINED ENV{DEEPSELECT_SRC_DIR})
+  set(DEEPSELECT_SRC_DIR $ENV{DEEPSELECT_SRC_DIR})
+endif()
+
+if(DEEPSELECT_SRC_DIR)
+  FetchContent_Declare(
+        deepselect
+        SOURCE_DIR ${DEEPSELECT_SRC_DIR}
+        CONFIGURE_COMMAND ""
+        BUILD_COMMAND ""
+  )
+else()
+  FetchContent_Declare(
+        deepselect
+        GIT_REPOSITORY https://github.com/vllm-project/DeepSelect.git
+        GIT_TAG d96d33afe1fab0d6066da49cdc91e64c2bee65ea
+        GIT_PROGRESS TRUE
+        CONFIGURE_COMMAND ""
+        BUILD_COMMAND ""
+  )
+endif()
+
+FetchContent_MakeAvailable(deepselect)
+message(STATUS "DeepSelect is available at ${deepselect_SOURCE_DIR}")
+
+# DeepSelect kernels require CUDA 12.9+; 10.0f (family) covers SM100 and SM103.
+set(DEEPSELECT_SUPPORT_ARCHS)
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
+    list(APPEND DEEPSELECT_SUPPORT_ARCHS "10.0f")
+endif()
+
+cuda_archs_loose_intersection(DEEPSELECT_ARCHS "${DEEPSELECT_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
+if(DEEPSELECT_ARCHS)
+    message(STATUS "DeepSelect CUDA architectures: ${DEEPSELECT_ARCHS}")
+
+    # All explicit template instantiations must be compiled: api.cpp
+    # references every dispatch combination.
+    file(GLOB DeepSelect_SOURCES CONFIGURE_DEPENDS
+        ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3/instantiations/*.cu
+        ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3_fp32/instantiations/*.cu
+        ${deepselect_SOURCE_DIR}/csrc/cuda_kernels/v3_cluster/instantiations/*.cu
+    )
+    list(PREPEND DeepSelect_SOURCES ${deepselect_SOURCE_DIR}/csrc/api.cpp)
+
+    set(DeepSelect_INCLUDES
+        ${deepselect_SOURCE_DIR}/csrc
+        ${deepselect_SOURCE_DIR}/csrc/3rdparty/cutlass/include
+        ${deepselect_SOURCE_DIR}/csrc/3rdparty/kerutils/include
+    )
+
+    set_gencode_flags_for_srcs(
+        SRCS "${DeepSelect_SOURCES}"
+        CUDA_ARCHS "${DEEPSELECT_ARCHS}")
+
+    set(VLLM_DEEPSELECT_GPU_FLAGS ${VLLM_GPU_FLAGS})
+    list(APPEND VLLM_DEEPSELECT_GPU_FLAGS
+        "-O3" "--expt-relaxed-constexpr" "--expt-extended-lambda"
+        "--use_fast_math" "--ftz=false"
+        "-U__CUDA_NO_HALF_OPERATORS__" "-U__CUDA_NO_HALF_CONVERSIONS__"
+        "-U__CUDA_NO_HALF2_OPERATORS__" "-U__CUDA_NO_BFLOAT16_CONVERSIONS__")
+
+    define_extension_target(
+        _deepselect_C
+        DESTINATION vllm
+        LANGUAGE ${VLLM_GPU_LANG}
+        SOURCES ${DeepSelect_SOURCES}
+        COMPILE_FLAGS ${VLLM_DEEPSELECT_GPU_FLAGS}
+        ARCHITECTURES ${VLLM_GPU_ARCHES}
+        INCLUDE_DIRECTORIES ${DeepSelect_INCLUDES}
+        USE_SABI 3
+        WITH_SOABI)
+
+    # Only use C-shim APIs available in PyTorch 2.10.
+    # _deepselect_C is abi compatible with PyTorch >= TORCH_TARGET_VERSION.
+    target_compile_definitions(_deepselect_C PRIVATE
+        TORCH_TARGET_VERSION=0x020A000000000000ULL)
+
+    # Needed to use cuda APIs from C-shim
+    if(VLLM_GPU_LANG STREQUAL "CUDA")
+        target_compile_definitions(_deepselect_C PRIVATE USE_CUDA)
+    endif()
+
+    # DeepSelect requires C++20
+    target_compile_options(_deepselect_C PRIVATE
+        $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
+        $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
+else()
+    message(STATUS "DeepSelect will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
+    # Create an empty target for setup.py on unsupported systems
+    add_custom_target(_deepselect_C)
+endif()
diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake
index 6ffb46e6a793..5ac4f4b415ea 100644
--- a/cmake/external_projects/flashkda.cmake
+++ b/cmake/external_projects/flashkda.cmake
@@ -13,7 +13,7 @@ else()
   FetchContent_Declare(
     flashkda
     GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
-    GIT_TAG 3b225bf26bb8e218928a1fe14751cb48cf31d11b
+    GIT_TAG b59532f1f464fbd536272780e30df5bf6a2ccc02
     GIT_PROGRESS TRUE
     GIT_SUBMODULES cutlass
   )
diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake
index 4cf1e421f0a5..cfa8243f726c 100644
--- a/cmake/external_projects/flashmla.cmake
+++ b/cmake/external_projects/flashmla.cmake
@@ -19,7 +19,7 @@ else()
   FetchContent_Declare(
         flashmla
         GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
-        GIT_TAG 6bc49418c5ead572ff0339191ddf3b155749e183
+        GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf
         GIT_PROGRESS TRUE
         CONFIGURE_COMMAND ""
         BUILD_COMMAND ""
@@ -76,43 +76,77 @@ if(FLASH_MLA_ARCHS)
 
     set(FlashMLA_SOURCES
         ${flashmla_SOURCE_DIR}/csrc/api/api.cpp
+        ${flashmla_SOURCE_DIR}/csrc/api/sparse_prefill.cpp
+        ${flashmla_SOURCE_DIR}/csrc/api/sparse_decode.cpp
+        ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp
+        ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp
 
         # Misc kernels for decoding
-        ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
-        ${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu
 
         # sm90 dense decode
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu
 
         # sm90 sparse decode
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu
 
         # sm90 sparse prefill
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
 
-        # sm100 dense prefill & backward
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
+        # sm100 dense prefill (inference-only: dense backward stays out)
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
 
         # sm100 sparse prefill
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu
+
+        # sm100 fused norm + rope + attn + rope + cast
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu
 
         # sm100 sparse decode
-        ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu
-        ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu
+        ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu
     )
 
     set(FlashMLA_Extension_SOURCES
@@ -125,7 +159,6 @@ if(FLASH_MLA_ARCHS)
     set(FlashMLA_INCLUDES
         ${flashmla_SOURCE_DIR}/csrc
         ${flashmla_SOURCE_DIR}/csrc/kerutils/include
-        ${flashmla_SOURCE_DIR}/csrc/sm90
         ${flashmla_SOURCE_DIR}/csrc/cutlass/include
         ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
     )
@@ -186,6 +219,15 @@ if(FLASH_MLA_ARCHS)
     if(VLLM_GPU_LANG STREQUAL "CUDA")
         target_compile_definitions(_flashmla_extension_C PRIVATE USE_CUDA)
     endif()
+
+    # FlashMLA's sources use M_LOG2E from <cmath>. MSVC's UCRT only defines the
+    # M_* constants when _USE_MATH_DEFINES is set before <cmath> is first
+    # included, so define it for both targets on MSVC. Scoped to MSVC so the
+    # compile flags (and build caches) on other platforms stay unchanged.
+    if(MSVC)
+        target_compile_definitions(_flashmla_C PRIVATE _USE_MATH_DEFINES)
+        target_compile_definitions(_flashmla_extension_C PRIVATE _USE_MATH_DEFINES)
+    endif()
 else()
     message(STATUS "FlashMLA will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
     # Create empty targets for setup.py on unsupported systems
diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake
index 442097dec7d3..ae5bfec7b47d 100644
--- a/cmake/external_projects/fmha_sm100.cmake
+++ b/cmake/external_projects/fmha_sm100.cmake
@@ -17,7 +17,7 @@ else()
   FetchContent_Declare(
     fmha_sm100
     GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 087c161814d4d9c735b46c21212a09e5f8eb92fa
+    GIT_TAG f355c37eb4e1413f21ee2ad8bbad25079e6bef9d
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND ""
diff --git a/cmake/external_projects/qutlass.cmake b/cmake/external_projects/qutlass.cmake
index cd3665c889e9..b6df33c170b8 100644
--- a/cmake/external_projects/qutlass.cmake
+++ b/cmake/external_projects/qutlass.cmake
@@ -142,9 +142,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
 
   target_compile_definitions(_qutlass_C PRIVATE
     QUTLASS_MINIMAL_BUILD=1
+    QUTLASS_DISABLE_PYBIND=1
     TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
     CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
-    TORCH_TARGET_VERSION=0x020B000000000000ULL
     USE_CUDA)
 
   set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
diff --git a/cmake/external_projects/tml_fa4.cmake b/cmake/external_projects/tml_fa4.cmake
index a965295e2bbd..c6ddff82c339 100644
--- a/cmake/external_projects/tml_fa4.cmake
+++ b/cmake/external_projects/tml_fa4.cmake
@@ -14,7 +14,7 @@ else()
   FetchContent_Declare(
     tml_fa4
     GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
-    GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d
+    GIT_TAG 75765e76a9c2c012c1f6ecd64577eb646eb4d303
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND "")
diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake
index 1b3be244de54..2a26023f740b 100644
--- a/cmake/external_projects/vllm_flash_attn.cmake
+++ b/cmake/external_projects/vllm_flash_attn.cmake
@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 506341a143fcabd4bb79052a7605ada727d6b3f5
+          GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp
index b31a6faa0104..13e430fb4192 100644
--- a/csrc/cpu/cpu_attn.cpp
+++ b/csrc/cpu/cpu_attn.cpp
@@ -1,16 +1,7 @@
 #include "cpu_attn_dispatch_generated.h"
 
-#if defined(__x86_64__) || defined(_M_X64)
-  #include <cpuid.h>
-// Runtime check for AMX-FP8 (Diamond Rapids): CPUID leaf 7, subleaf 1, EAX[21].
-static bool runtime_has_amx_fp8() {
-  unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;
-  if (__get_cpuid_count(7, 1, &eax, &ebx, &ecx, &edx)) {
-    return (eax >> 21) & 1u;
-  }
-  return false;
-}
-#endif  // x86_64
+// Runtime check for AMX-FP8, implemented in cpu_isa.cpp.
+extern bool runtime_has_amx_fp8();
 
 // Maps kv_cache_dtype string to Fp8KVCacheDataType enum.
 // "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.
@@ -252,7 +243,19 @@ void cpu_attention_with_kv_cache(
   input.block_table = block_table.data_ptr<int32_t>();
   input.alibi_slopes =
       alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
-  input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
+  // Attention sinks may be bf16 (native path) or fp32. Anything that is not
+  // bf16 must be provided as fp32 and is executed in full float precision.
+  if (s_aux.has_value()) {
+    TORCH_CHECK(s_aux->scalar_type() == at::ScalarType::BFloat16 ||
+                    s_aux->scalar_type() == at::ScalarType::Float,
+                "cpu_attention_with_kv_cache: s_aux (attention sinks) dtype ",
+                s_aux->scalar_type(), " must be bfloat16 or float32");
+    input.s_aux = s_aux->data_ptr();
+    input.s_aux_is_bf16 = s_aux->scalar_type() == at::ScalarType::BFloat16;
+  } else {
+    input.s_aux = nullptr;
+    input.s_aux_is_bf16 = false;
+  }
   input.dynamic_causal =
       dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
   input.scale = scale;
diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp
index e9ddaacea573..d914e150636c 100644
--- a/csrc/cpu/cpu_attn_impl.hpp
+++ b/csrc/cpu/cpu_attn_impl.hpp
@@ -879,7 +879,11 @@ struct AttentionInput {
   int32_t* seq_lens;
   int32_t* block_table;
   float* alibi_slopes;
-  c10::BFloat16* s_aux;
+  // Attention sinks pointer. May reference bf16 or fp32 data depending on
+  // `s_aux_is_bf16`. bf16 sinks keep the native bf16 path; anything else is
+  // provided as fp32 and executed in full float precision.
+  const void* s_aux;
+  bool s_aux_is_bf16;
   bool* dynamic_causal;
   float scale;
   bool causal;
@@ -956,6 +960,28 @@ void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
   std::printf("%s", ss.str().c_str());
 }
 
+// This allows us to store probabilities in a packed format
+// for GEMM implementations that need packed A.
+template <typename attention_impl_t, typename = void>
+struct ProbabilityTokenStore {
+  using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
+  static constexpr int32_t TokenStride = 1;
+
+  FORCE_INLINE static void store_probabilities(
+      prob_buffer_t* __restrict__ probability, const vec_op::FP32Vec16& values,
+      const int32_t, const int64_t) {
+    using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
+    prob_buffer_vec_t output_vec(values);
+    output_vec.save(probability);
+  }
+};
+
+template <typename attention_impl_t>
+struct ProbabilityTokenStore<
+    attention_impl_t,
+    std::void_t<typename attention_impl_t::ProbabilityTokenStore>>
+    : attention_impl_t::ProbabilityTokenStore {};
+
 template <typename attention_impl_t>
 class AttentionMainLoop {
  public:
@@ -966,6 +992,7 @@ class AttentionMainLoop {
   using partial_output_buffer_t =
       typename attention_impl_t::partial_output_buffer_t;
   using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
+  using probability_token_store_t = ProbabilityTokenStore<attention_impl_t>;
 
   static constexpr int64_t max_q_head_num_per_iter =
       attention_impl_t::MaxQHeadNumPerIteration;
@@ -974,6 +1001,8 @@ class AttentionMainLoop {
   static constexpr int64_t headdim_alignment =
       attention_impl_t::HeadDimAlignment;
   static constexpr int64_t head_dim = attention_impl_t::HeadDim;
+  static constexpr int32_t probability_token_stride =
+      probability_token_store_t::TokenStride;
   static constexpr ISA ISAType = attention_impl_t::ISAType;
   static constexpr bool scale_on_logits =
       attention_impl_t::scale_on_logits;  // apply scale on logits, otherwise
@@ -1152,8 +1181,6 @@ class AttentionMainLoop {
                 ? kv_tile_token_num
                 : kv_tile_token_num *
                       (sizeof(logits_buffer_t) / sizeof(prob_buffer_t));
-        constexpr int64_t prob_buffer_elem_size =
-            prequantize_probabilities ? sizeof(uint8_t) : sizeof(prob_buffer_t);
         partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer;
         bool accum_c = !is_first_iter;
         for (int32_t block_idx = start_block_idx; block_idx < end_block_idx;
@@ -1187,9 +1214,7 @@ class AttentionMainLoop {
           remaining_group_num -= curr_group_num_in_block;
           curr_group_offset = 0;
           curr_group_num_in_block = token_group_num_per_block;
-          curr_prob_buffer = reinterpret_cast<pv_prob_buffer_t*>(
-              reinterpret_cast<uint8_t*>(curr_prob_buffer) +
-              curr_token_num * prob_buffer_elem_size);
+          curr_prob_buffer += curr_token_num * probability_token_stride;
           curr_partial_q_buffer = partial_q_buffer;
           accum_c = true;
         }
@@ -1262,7 +1287,6 @@ class AttentionMainLoop {
           std::is_same_v<query_t, c10::Half>;
 #endif
 
-      using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
       constexpr bool prequantize_probabilities = []() {
         if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) {
           return tile_gemm_t::prequantize_probabilities;
@@ -1331,12 +1355,6 @@ class AttentionMainLoop {
               vec = fast_exp(vec);
             }
 
-            if constexpr (prequantize_probabilities) {
-              vec.save(curr_logits_buffer_iter);
-            } else {
-              prob_buffer_vec_t output_vec(vec);
-              output_vec.save(curr_prob_buffer_iter);
-            }
 #else
             vec.save(curr_logits_buffer_iter);
             for (int32_t k = 0; k < 16; ++k) {
@@ -1345,10 +1363,19 @@ class AttentionMainLoop {
             vec = vec_op::FP32Vec16(curr_logits_buffer_iter);
 #endif
 
+            if constexpr (prequantize_probabilities) {
+              vec.save(curr_logits_buffer_iter);
+            } else {
+              probability_token_store_t::store_probabilities(
+                  curr_prob_buffer_iter, vec, i,
+                  logits_buffer_stride * sizeof(logits_buffer_t) /
+                      sizeof(prob_buffer_t));
+            }
+
             sum_vec = sum_vec + vec;
 
             curr_logits_buffer_iter += 16;
-            curr_prob_buffer_iter += 16;
+            curr_prob_buffer_iter += 16 * probability_token_stride;
           }
           if constexpr (prequantize_probabilities) {
             tile_gemm_t::quantize_probability_row(
@@ -1542,7 +1569,8 @@ class AttentionMainLoop {
       const float scale = input->scale;
       const float softcap_scale = input->softcap;
       const float* alibi_slopes = input->alibi_slopes;
-      const c10::BFloat16* s_aux = input->s_aux;
+      const void* s_aux = input->s_aux;
+      const bool s_aux_is_bf16 = input->s_aux_is_bf16;
       const bool* dynamic_causal = input->dynamic_causal;
       const bool is_dynamic_causal = dynamic_causal != nullptr;
 
@@ -1750,9 +1778,6 @@ class AttentionMainLoop {
               const float* curr_alibi_slopes =
                   (alibi_slopes != nullptr ? alibi_slopes + q_head_start_idx
                                            : nullptr);
-              const c10::BFloat16* curr_s_aux =
-                  (s_aux != nullptr ? s_aux + q_head_start_idx : nullptr);
-
               // copy the Q tile to q_buffer, the logical layout of q_buffer is
               // [actual_q_token_num, curr_q_heads_per_kv, head_dim]
               {
@@ -1764,10 +1789,23 @@ class AttentionMainLoop {
 
               if (use_sink) {
                 alignas(64) float s_aux_fp32[16];
-                // All other platforms have BF16Vec16 available
-                vec_op::BF16Vec16 vec_bf16(curr_s_aux);
-                vec_op::FP32Vec16 vec_fp32(vec_bf16);
-                vec_fp32.save(s_aux_fp32);
+                // Sinks may be provided as bf16 or fp32. Either way they are
+                // loaded into fp32: bf16 sinks are up-converted (native
+                // BF16Vec16 path, available on all platforms), while fp32
+                // sinks are loaded directly and executed in full precision.
+                if (s_aux_is_bf16) {
+                  const c10::BFloat16* curr_s_aux =
+                      static_cast<const c10::BFloat16*>(s_aux) +
+                      q_head_start_idx;
+                  vec_op::BF16Vec16 vec_bf16(curr_s_aux);
+                  vec_op::FP32Vec16 vec_fp32(vec_bf16);
+                  vec_fp32.save(s_aux_fp32);
+                } else {
+                  const float* curr_s_aux =
+                      static_cast<const float*>(s_aux) + q_head_start_idx;
+                  vec_op::FP32Vec16 vec_fp32(curr_s_aux);
+                  vec_fp32.save(s_aux_fp32);
+                }
 
                 float* __restrict__ curr_sum_buffer = sum_buffer;
                 float* __restrict__ curr_max_buffer = max_buffer;
diff --git a/csrc/cpu/cpu_attn_neon.hpp b/csrc/cpu/cpu_attn_neon.hpp
index db4c5df2e88d..eed749944496 100644
--- a/csrc/cpu/cpu_attn_neon.hpp
+++ b/csrc/cpu/cpu_attn_neon.hpp
@@ -11,7 +11,7 @@ namespace cpu_attention {
 
 namespace {
 
-#define BLOCK_SIZE_ALIGNMENT 32
+#define BLOCK_SIZE_ALIGNMENT 16
 #define HEAD_SIZE_ALIGNMENT 32
 #define MAX_Q_HEAD_NUM_PER_ITER 16
 
@@ -387,7 +387,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
 };
 
 #ifdef ARM_BF16_SUPPORT
-// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
+// For BF16 on Arm, reuse the BFMMLA kernels with the same token alignment.
 template <int64_t head_dim>
 class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
     : public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
diff --git a/csrc/cpu/cpu_attn_neon_bfmmla.hpp b/csrc/cpu/cpu_attn_neon_bfmmla.hpp
index 4e4578a74f5b..33bc0d5e029d 100644
--- a/csrc/cpu/cpu_attn_neon_bfmmla.hpp
+++ b/csrc/cpu/cpu_attn_neon_bfmmla.hpp
@@ -5,432 +5,275 @@
 
 #include "cpu_attn_impl.hpp"
 
+#include <arm_bf16.h>
 #include <arm_neon.h>
+#include <c10/util/BFloat16.h>
 
+#include <algorithm>
 #include <cstdint>
-#include <vector>
 
 namespace cpu_attention {
 
-namespace {
-
-// BFMMLA tile dimensions
-constexpr int32_t TILE_ROWS = 2;  // M dimension
-constexpr int32_t TILE_K = 4;     // K reduction
-constexpr int32_t TILE_COLS = 2;  // N dimension (column-pair)
-
-// Derived constants
-constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8;   // 4 column-pairs
-constexpr int32_t K_TOKENS_PER_GROUP = 8;      // Tokens grouped in K cache
-constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4;  // Tokens per V cache row block
-constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
-constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
-constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K;  // A packing
-
-// Matrix Packing and Accumulator
-// Reshape two rows of Q into BFMMLA-friendly interleaved
-// Input:  row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3]
-// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7]
-// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk
-FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0,
-                                           const c10::BFloat16* __restrict r1,
-                                           c10::BFloat16* __restrict dst,
-                                           int32_t K) {
-  const uint16_t* s0 = reinterpret_cast<const uint16_t*>(r0);
-  const uint16_t* s1 = reinterpret_cast<const uint16_t*>(r1);
-  uint16_t* d = reinterpret_cast<uint16_t*>(dst);
-
-  // Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output)
-  int32_t k = 0;
-  for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) {
-    vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k)));
-  }
-
-  // Handle K tail: pack remaining elements with zero-padding
-  const int32_t tail = K - k;
-  if (tail > 0) {
-    // Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1],
-    // pad]
-    for (int32_t t = 0; t < tail; ++t) {
-      d[t] = s0[k + t];
-      d[t + TILE_K] = s1[k + t];
-    }
-    // Zero-pad the rest
-    for (int32_t t = tail; t < TILE_K; ++t) {
-      d[t] = 0;
-      d[t + TILE_K] = 0;
+class BfmmlaGemm {
+ public:
+  static constexpr int32_t KTile = 4;
+  static constexpr int32_t NTile = 8;
+  static constexpr int32_t MaxRows = 8;
+
+  FORCE_INLINE static void gemm(const c10::BFloat16* __restrict__ a,
+                                const c10::BFloat16* __restrict__ b,
+                                float* __restrict__ c, const int32_t m,
+                                const int32_t n, const int32_t k,
+                                const int64_t a_pair_stride,
+                                const int64_t b_n_group_stride,
+                                const int64_t b_k_group_stride,
+                                const int64_t ldc, const bool accumulate) {
+    const auto* a_ptr = reinterpret_cast<const bfloat16_t*>(a);
+    const auto* b_ptr = reinterpret_cast<const bfloat16_t*>(b);
+
+    for (int32_t n_idx = 0; n_idx < n; n_idx += 16) {
+      const auto* b_panel = b_ptr + (n_idx / NTile) * b_n_group_stride;
+      float* c_panel = c + n_idx;
+
+      // Preserve this range so the inactive row pair is optimized away.
+      if (m <= 2) {
+        gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride,
+                  b_n_group_stride, b_k_group_stride, ldc, accumulate);
+      } else if (m <= 4) {
+        gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride,
+                  b_n_group_stride, b_k_group_stride, ldc, accumulate);
+      } else {
+        gemm_8x8(a_ptr, b_panel, c_panel, m, k, a_pair_stride, b_k_group_stride,
+                 ldc, accumulate);
+        gemm_8x8(a_ptr, b_panel + b_n_group_stride, c_panel + NTile, m, k,
+                 a_pair_stride, b_k_group_stride, ldc, accumulate);
+      }
     }
   }
-}
-
-// 2x2 accumulator load/store with compile-time row count
-template <int32_t m_rows>
-FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) {
-  static_assert(m_rows == 1 || m_rows == 2);
-  float32x2_t row0 = vld1_f32(base + col_off);
-  float32x2_t row1 =
-      (m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f);
-  return vcombine_f32(row0, row1);
-}
-
-template <int32_t m_rows>
-FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc,
-                                int col_off) {
-  static_assert(m_rows == 1 || m_rows == 2);
-  vst1_f32(base + col_off, vget_low_f32(acc));
-  if constexpr (m_rows == 2) {
-    vst1_f32(base + ldc + col_off, vget_high_f32(acc));
-  }
-}
-
-// Initialize 4 column-pair accumulators for 2 rows (8 columns total)
-#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \
-  do {                                                               \
-    if (accum) {                                                     \
-      if (m_rows == 2) {                                             \
-        a0 = load_acc_2x2<2>(Crow, ldc, 0);                          \
-        a1 = load_acc_2x2<2>(Crow, ldc, 2);                          \
-        a2 = load_acc_2x2<2>(Crow, ldc, 4);                          \
-        a3 = load_acc_2x2<2>(Crow, ldc, 6);                          \
-      } else {                                                       \
-        a0 = load_acc_2x2<1>(Crow, ldc, 0);                          \
-        a1 = load_acc_2x2<1>(Crow, ldc, 2);                          \
-        a2 = load_acc_2x2<1>(Crow, ldc, 4);                          \
-        a3 = load_acc_2x2<1>(Crow, ldc, 6);                          \
-      }                                                              \
-    } else {                                                         \
-      a0 = a1 = a2 = a3 = vdupq_n_f32(0.f);                          \
-    }                                                                \
-  } while (0)
-
-// Store 4 column-pair accumulators back to C matrix
-#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \
-  do {                                                         \
-    if (m_rows == 2) {                                         \
-      store_acc_2x2<2>(a0, Crow, ldc, 0);                      \
-      store_acc_2x2<2>(a1, Crow, ldc, 2);                      \
-      store_acc_2x2<2>(a2, Crow, ldc, 4);                      \
-      store_acc_2x2<2>(a3, Crow, ldc, 6);                      \
-    } else {                                                   \
-      store_acc_2x2<1>(a0, Crow, ldc, 0);                      \
-      store_acc_2x2<1>(a1, Crow, ldc, 2);                      \
-      store_acc_2x2<1>(a2, Crow, ldc, 4);                      \
-      store_acc_2x2<1>(a3, Crow, ldc, 6);                      \
-    }                                                          \
-  } while (0)
-
-// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs
-#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \
-  do {                                                      \
-    r0 = vbfmmlaq_f32(r0, a, b0);                           \
-    r1 = vbfmmlaq_f32(r1, a, b1);                           \
-    r2 = vbfmmlaq_f32(r2, a, b2);                           \
-    r3 = vbfmmlaq_f32(r3, a, b3);                           \
-  } while (0)
-
-// Micro-kernel: updates a small fixed tile using BFMMLA.
-// RP = number of row-pairs (1,2,4)
-// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B.
-// A_packed interleaves RP row-pairs; B layout is driven by the attention phase:
-// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T)
-// - AttentionGemmPhase::PV -> token-row layout (P @ V)
-// K_static < 0 enables runtime K (PV only)
-template <int32_t RP, int32_t K_static, AttentionGemmPhase phase>
-FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon(
-    const bfloat16_t* const* __restrict A_packed_rp,
-    const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk,
-    float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride,
-    int32_t K_runtime = 0) {
-  static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4");
-  static_assert(K_static < 0 || K_static % TILE_K == 0,
-                "K must be divisible by TILE_K");
-  static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
-                "Runtime K only supported for PV");
-
-  constexpr bool runtime_k = (K_static < 0);
-  const int32_t K_iters =
-      runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K);
-  const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0;
-
-  if (!runtime_k) {
-    // Help the compiler fold away unused K_runtime when K is compile-time
-    (void)K_runtime;
-  }
-
-  auto* C_al = C;
-  const auto* B_al = B_blk;
 
-  // Setup A pointers
-  const bfloat16_t* a_ptr[4] = {
-      A_packed_rp[0],
-      (RP >= 2) ? A_packed_rp[1] : nullptr,
-      (RP >= 4) ? A_packed_rp[2] : nullptr,
-      (RP >= 4) ? A_packed_rp[3] : nullptr,
-  };
-
-  // Setup B pointers based on layout
-  const bfloat16_t* b_ptr[4];
-  if constexpr (phase == AttentionGemmPhase::PV) {
-    b_ptr[0] = B_blk + 0 * b_stride;
-    b_ptr[1] = B_blk + 1 * b_stride;
-    b_ptr[2] = B_blk + 2 * b_stride;
-    b_ptr[3] = B_blk + 3 * b_stride;
+ private:
+  FORCE_INLINE static float32x4_t zip_low_pairs(const float32x4_t a,
+                                                const float32x4_t b) {
+    return vreinterpretq_f32_f64(
+        vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
   }
 
-  float32x4_t acc[4][4];
-
-// Initialize accumulators
-#define INIT_RP(rp)                                                            \
-  if constexpr (RP > rp) {                                                     \
-    INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3],         \
-                       C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \
+  FORCE_INLINE static float32x4_t zip_high_pairs(const float32x4_t a,
+                                                 const float32x4_t b) {
+    return vreinterpretq_f32_f64(
+        vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
   }
-  INIT_RP(0);
-  INIT_RP(1);
-  INIT_RP(2);
-  INIT_RP(3);
-#undef INIT_RP
-
-  // Main compute loop
-  for (int32_t ki = 0; ki < K_iters; ++ki) {
-    bfloat16x8_t b0, b1, b2, b3;
-    if constexpr (phase == AttentionGemmPhase::PV) {
-      b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE);
-      b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE);
-      b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE);
-      b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE);
-    } else {
-      const bfloat16_t* b_base = B_al + ki * b_stride;
-      b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE);
-      b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE);
-      b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE);
-      b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE);
+
+  FORCE_INLINE static void init_accumulators(
+      float32x4_t& acc01, float32x4_t& acc23, float32x4_t& acc45,
+      float32x4_t& acc67, const float* __restrict__ c, const int64_t ldc,
+      const int32_t rows, const bool accumulate) {
+    if (!accumulate || rows == 0) {
+      acc01 = vdupq_n_f32(0.0f);
+      acc23 = vdupq_n_f32(0.0f);
+      acc45 = vdupq_n_f32(0.0f);
+      acc67 = vdupq_n_f32(0.0f);
+      return;
     }
 
-#define COMPUTE_RP(rp)                                                       \
-  if constexpr (RP > rp) {                                                   \
-    bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \
-    BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0,  \
-                     b1, b2, b3);                                            \
-  }
-    COMPUTE_RP(0);
-    COMPUTE_RP(1);
-    COMPUTE_RP(2);
-    COMPUTE_RP(3);
-#undef COMPUTE_RP
+    const float32x4_t row0_0123 = vld1q_f32(c);
+    const float32x4_t row0_4567 = vld1q_f32(c + 4);
+    const float32x4_t row1_0123 =
+        (rows == 2) ? vld1q_f32(c + ldc) : vdupq_n_f32(0.0f);
+    const float32x4_t row1_4567 =
+        (rows == 2) ? vld1q_f32(c + ldc + 4) : vdupq_n_f32(0.0f);
+
+    acc01 = zip_low_pairs(row0_0123, row1_0123);
+    acc23 = zip_high_pairs(row0_0123, row1_0123);
+    acc45 = zip_low_pairs(row0_4567, row1_4567);
+    acc67 = zip_high_pairs(row0_4567, row1_4567);
   }
 
-  // K tail for runtime PV: fallback path
-  if constexpr (runtime_k) {
-    if (K_tail > 0) {
-      const int32_t tail_offset = K_iters * V_INNER_STRIDE;
-      const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK;
-      for (int32_t kt = 0; kt < K_tail; ++kt) {
-        float32x4_t b_vecs[4];
-        for (int32_t p = 0; p < 4; ++p) {
-          const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS;
-          const float b0 = vcvtah_f32_bf16(bp[0]);
-          const float b1 = vcvtah_f32_bf16(bp[1]);
-          const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1);
-          b_vecs[p] = vcombine_f32(b_pair, b_pair);
-        }
+  FORCE_INLINE static void store_accumulators(
+      const float32x4_t acc01, const float32x4_t acc23, const float32x4_t acc45,
+      const float32x4_t acc67, float* __restrict__ c, const int64_t ldc,
+      const int32_t rows) {
+    if (rows == 0) {
+      return;
+    }
 
-#define TAIL_RP(rp)                                                     \
-  if constexpr (RP > rp) {                                              \
-    const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset;             \
-    float a_row0 = vcvtah_f32_bf16(ap[kt]);                             \
-    float a_row1 =                                                      \
-        (m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \
-    const float32x4_t a_vec =                                           \
-        vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1));           \
-    for (int32_t p = 0; p < 4; ++p) {                                   \
-      acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]);             \
-    }                                                                   \
-  }
-        TAIL_RP(0);
-        TAIL_RP(1);
-        TAIL_RP(2);
-        TAIL_RP(3);
-#undef TAIL_RP
-      }
+    vst1q_f32(c, zip_low_pairs(acc01, acc23));
+    vst1q_f32(c + 4, zip_low_pairs(acc45, acc67));
+    if (rows == 2) {
+      vst1q_f32(c + ldc, zip_high_pairs(acc01, acc23));
+      vst1q_f32(c + ldc + 4, zip_high_pairs(acc45, acc67));
     }
   }
 
-  // Store results
-#define STORE_RP(rp)                                                    \
-  if constexpr (RP > rp) {                                              \
-    STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
-                        C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]);     \
-  }
-  STORE_RP(0);
-  STORE_RP(1);
-  STORE_RP(2);
-  STORE_RP(3);
-#undef STORE_RP
-}
-
-// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the
-// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables
-// runtime K (PV only).
-template <int32_t MB, int32_t N, int32_t K_static, AttentionGemmPhase phase>
-FORCE_INLINE void gemm_packA_compute_MB_xN(
-    const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
-    float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc,
-    int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
-  static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]");
-  static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
-                "N must be a multiple of OUTPUT_COLS_PER_BLOCK");
-  static_assert(K_static < 0 || K_static % TILE_K == 0,
-                "K must be divisible by TILE_K");
-  static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
-                "Runtime K only supported for PV");
-
-  constexpr bool runtime_k = (K_static < 0);
-  const int32_t K_val = runtime_k ? K_runtime : K_static;
-
-  // Keep small packs on-stack to avoid heap churn
-  constexpr int32_t STACK_PACK_STRIDE =
-      (1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
-
-  constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS;
-  const int32_t pack_stride =
-      runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK
-                : (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
-
-  alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE];
-  std::vector<c10::BFloat16> A_packed_heap;
-  c10::BFloat16* A_packed =
-      (pack_stride <= STACK_PACK_STRIDE)
-          ? A_packed_stack
-          : (A_packed_heap.resize(ROW_PAIRS * pack_stride),
-             A_packed_heap.data());
-
-  for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) {
-    const int32_t m = rp * TILE_ROWS;
-    const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1;
-    const c10::BFloat16* A0 = A + m * lda;
-    const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0;
-    reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val);
+  FORCE_INLINE static bfloat16x8_t load_a_pair(const bfloat16_t* __restrict__ a,
+                                               const int32_t rows) {
+    if (rows == 0) {
+      return vdupq_n_bf16(bfloat16_t{});
+    }
+    // Packed A reserves both rows for an M tail.
+    return vld1q_bf16(a);
   }
 
-  for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) {
-    const c10::BFloat16* B_blk_c10 =
-        (phase == AttentionGemmPhase::PV)
-            ? (B + (n / TILE_COLS) * b_layout_stride)
-            : (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride);
-    const bfloat16_t* B_blk = reinterpret_cast<const bfloat16_t*>(B_blk_c10);
-
-    // Process row-pairs in groups of 4, 2, then 1
-    int32_t row_pair_idx = 0;
-
-#define PROCESS_RP_GROUP(group_size)                                       \
-  for (; row_pair_idx + (group_size - 1) < ROW_PAIRS;                      \
-       row_pair_idx += group_size) {                                       \
-    const bfloat16_t* Ap[group_size];                                      \
-    int32_t mr[group_size];                                                \
-    for (int32_t i = 0; i < group_size; ++i) {                             \
-      Ap[i] = reinterpret_cast<const bfloat16_t*>(                         \
-          A_packed + (row_pair_idx + i) * pack_stride);                    \
-      mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \
-    }                                                                      \
-    float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n;               \
-    if constexpr (runtime_k) {                                             \
-      gemm_rowpairs_x8_bfmmla_neon<group_size, -1, phase>(                 \
-          Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val);  \
-    } else {                                                               \
-      gemm_rowpairs_x8_bfmmla_neon<group_size, K_static, phase>(           \
-          Ap, mr, B_blk, C_blk, ldc, accumulate,                           \
-          (phase == AttentionGemmPhase::PV) ? b_layout_stride              \
-                                            : b_reduction_stride);         \
-    }                                                                      \
-  }
+  FORCE_INLINE static void gemm_4x16(const bfloat16_t* __restrict__ a,
+                                     const bfloat16_t* __restrict__ b,
+                                     float* __restrict__ c, const int32_t m,
+                                     const int32_t k,
+                                     const int64_t a_pair_stride,
+                                     const int64_t b_n_group_stride,
+                                     const int64_t b_k_group_stride,
+                                     const int64_t ldc, const bool accumulate) {
+    const int32_t rows01 = std::min(2, std::max(0, m));
+    const int32_t rows23 = std::min(2, std::max(0, m - 2));
+    float32x4_t acc0101, acc0123, acc0145, acc0167;
+    float32x4_t acc2301, acc2323, acc2345, acc2367;
+    float32x4_t acc0189, acc011011, acc011213, acc011415;
+    float32x4_t acc2389, acc231011, acc231213, acc231415;
+    init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01,
+                      accumulate);
+    init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
+                      rows23, accumulate);
+    init_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc,
+                      rows01, accumulate);
+    init_accumulators(acc2389, acc231011, acc231213, acc231415, c + 2 * ldc + 8,
+                      ldc, rows23, accumulate);
+
+    const bfloat16_t* a01 = a;
+    const bfloat16_t* a23 = a + a_pair_stride;
+    const bfloat16_t* b0 = b;
+    const bfloat16_t* b1 = b + b_n_group_stride;
+
+#pragma GCC unroll 4
+    for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) {
+      const bfloat16x8_t av01 = load_a_pair(a01, rows01);
+      const bfloat16x8_t av23 = load_a_pair(a23, rows23);
+      const bfloat16x8_t b01 = vld1q_bf16(b0);
+      const bfloat16x8_t b23 = vld1q_bf16(b0 + NTile);
+      const bfloat16x8_t b45 = vld1q_bf16(b0 + 2 * NTile);
+      const bfloat16x8_t b67 = vld1q_bf16(b0 + 3 * NTile);
+      const bfloat16x8_t b89 = vld1q_bf16(b1);
+      const bfloat16x8_t b1011 = vld1q_bf16(b1 + NTile);
+      const bfloat16x8_t b1213 = vld1q_bf16(b1 + 2 * NTile);
+      const bfloat16x8_t b1415 = vld1q_bf16(b1 + 3 * NTile);
+
+      acc0101 = vbfmmlaq_f32(acc0101, av01, b01);
+      acc2301 = vbfmmlaq_f32(acc2301, av23, b01);
+      acc0123 = vbfmmlaq_f32(acc0123, av01, b23);
+      acc2323 = vbfmmlaq_f32(acc2323, av23, b23);
+      acc0145 = vbfmmlaq_f32(acc0145, av01, b45);
+      acc2345 = vbfmmlaq_f32(acc2345, av23, b45);
+      acc0167 = vbfmmlaq_f32(acc0167, av01, b67);
+      acc2367 = vbfmmlaq_f32(acc2367, av23, b67);
+      acc0189 = vbfmmlaq_f32(acc0189, av01, b89);
+      acc2389 = vbfmmlaq_f32(acc2389, av23, b89);
+      acc011011 = vbfmmlaq_f32(acc011011, av01, b1011);
+      acc231011 = vbfmmlaq_f32(acc231011, av23, b1011);
+      acc011213 = vbfmmlaq_f32(acc011213, av01, b1213);
+      acc231213 = vbfmmlaq_f32(acc231213, av23, b1213);
+      acc011415 = vbfmmlaq_f32(acc011415, av01, b1415);
+      acc231415 = vbfmmlaq_f32(acc231415, av23, b1415);
+
+      a01 += 2 * KTile;
+      a23 += 2 * KTile;
+      b0 += b_k_group_stride;
+      b1 += b_k_group_stride;
+    }
 
-    PROCESS_RP_GROUP(4);
-    PROCESS_RP_GROUP(2);
-    PROCESS_RP_GROUP(1);
-#undef PROCESS_RP_GROUP
+    store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01);
+    store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
+                       rows23);
+    store_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc,
+                       rows01);
+    store_accumulators(acc2389, acc231011, acc231213, acc231415,
+                       c + 2 * ldc + 8, ldc, rows23);
   }
-}
-
-// Macro-kernel: iterates over M in MB={8,4,2,1} chunks.
-// Supports compile-time K specialization when K >= 0; otherwise uses runtime K
-// (runtime K path is only supported for PV).
-template <AttentionGemmPhase phase, int32_t N, int32_t K = -1>
-FORCE_INLINE void gemm_macro_neon_bfmmla(
-    const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
-    float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc,
-    int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
-  static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
-                "N must be a multiple of OUTPUT_COLS_PER_BLOCK");
-
-  if constexpr (K >= 0) {
-    static_assert(K % TILE_K == 0, "K must be divisible by TILE_K");
-    for (int32_t m = 0; m < M;) {
-      const int32_t rem = M - m;
-      const c10::BFloat16* A_blk = A + m * lda;
-      float* C_blk = C + m * ldc;
-
-#define DISPATCH_MB(mb)                                                   \
-  gemm_packA_compute_MB_xN<mb, N, K, phase>(A_blk, B, C_blk, 0, lda, ldc, \
-                                            b_layout_stride,              \
-                                            b_reduction_stride, accumulate)
-
-      if (rem >= 8) {
-        DISPATCH_MB(8);
-        m += 8;
-      } else if (rem >= 4) {
-        DISPATCH_MB(4);
-        m += 4;
-      } else if (rem >= 2) {
-        DISPATCH_MB(2);
-        m += 2;
-      } else {
-        DISPATCH_MB(1);
-        m += 1;
-      }
-#undef DISPATCH_MB
-    }
-  } else {
-    static_assert(phase == AttentionGemmPhase::PV,
-                  "Runtime K specialization only supported for PV.");
-    const int32_t K_val = K_runtime;
-
-    for (int32_t m = 0; m < M;) {
-      const int32_t rem = M - m;
-      const c10::BFloat16* A_blk = A + m * lda;
-      float* C_blk = C + m * ldc;
-
-#define DISPATCH_MB_RUNTIME(mb)                                                \
-  gemm_packA_compute_MB_xN<mb, N, -1, phase>(A_blk, B, C_blk, K_val, lda, ldc, \
-                                             b_layout_stride,                  \
-                                             b_reduction_stride, accumulate)
-
-      if (rem >= 8) {
-        DISPATCH_MB_RUNTIME(8);
-        m += 8;
-      } else if (rem >= 4) {
-        DISPATCH_MB_RUNTIME(4);
-        m += 4;
-      } else if (rem >= 2) {
-        DISPATCH_MB_RUNTIME(2);
-        m += 2;
-      } else {
-        DISPATCH_MB_RUNTIME(1);
-        m += 1;
-      }
-#undef DISPATCH_MB_RUNTIME
+
+  FORCE_INLINE static void gemm_8x8(const bfloat16_t* __restrict__ a,
+                                    const bfloat16_t* __restrict__ b,
+                                    float* __restrict__ c, const int32_t m,
+                                    const int32_t k,
+                                    const int64_t a_pair_stride,
+                                    const int64_t b_k_group_stride,
+                                    const int64_t ldc, const bool accumulate) {
+    const int32_t rows01 = std::min(2, std::max(0, m));
+    const int32_t rows23 = std::min(2, std::max(0, m - 2));
+    const int32_t rows45 = std::min(2, std::max(0, m - 4));
+    const int32_t rows67 = std::min(2, std::max(0, m - 6));
+    float32x4_t acc0101, acc0123, acc0145, acc0167;
+    float32x4_t acc2301, acc2323, acc2345, acc2367;
+    float32x4_t acc4501, acc4523, acc4545, acc4567;
+    float32x4_t acc6701, acc6723, acc6745, acc6767;
+    init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01,
+                      accumulate);
+    init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
+                      rows23, accumulate);
+    init_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc,
+                      rows45, accumulate);
+    init_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc,
+                      rows67, accumulate);
+
+    const bfloat16_t* a01 = a;
+    const bfloat16_t* a23 = a + a_pair_stride;
+    const bfloat16_t* a45 = a + 2 * a_pair_stride;
+    const bfloat16_t* a67 = a + 3 * a_pair_stride;
+    const bfloat16_t* b_ptr = b;
+
+#pragma GCC unroll 4
+    for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) {
+      const bfloat16x8_t av01 = load_a_pair(a01, rows01);
+      const bfloat16x8_t av23 = load_a_pair(a23, rows23);
+      const bfloat16x8_t av45 = load_a_pair(a45, rows45);
+      const bfloat16x8_t av67 = load_a_pair(a67, rows67);
+      const bfloat16x8_t b01 = vld1q_bf16(b_ptr);
+      const bfloat16x8_t b23 = vld1q_bf16(b_ptr + NTile);
+      const bfloat16x8_t b45 = vld1q_bf16(b_ptr + 2 * NTile);
+      const bfloat16x8_t b67 = vld1q_bf16(b_ptr + 3 * NTile);
+
+      acc0101 = vbfmmlaq_f32(acc0101, av01, b01);
+      acc2301 = vbfmmlaq_f32(acc2301, av23, b01);
+      acc4501 = vbfmmlaq_f32(acc4501, av45, b01);
+      acc6701 = vbfmmlaq_f32(acc6701, av67, b01);
+      acc0123 = vbfmmlaq_f32(acc0123, av01, b23);
+      acc2323 = vbfmmlaq_f32(acc2323, av23, b23);
+      acc4523 = vbfmmlaq_f32(acc4523, av45, b23);
+      acc6723 = vbfmmlaq_f32(acc6723, av67, b23);
+      acc0145 = vbfmmlaq_f32(acc0145, av01, b45);
+      acc2345 = vbfmmlaq_f32(acc2345, av23, b45);
+      acc4545 = vbfmmlaq_f32(acc4545, av45, b45);
+      acc6745 = vbfmmlaq_f32(acc6745, av67, b45);
+      acc0167 = vbfmmlaq_f32(acc0167, av01, b67);
+      acc2367 = vbfmmlaq_f32(acc2367, av23, b67);
+      acc4567 = vbfmmlaq_f32(acc4567, av45, b67);
+      acc6767 = vbfmmlaq_f32(acc6767, av67, b67);
+
+      a01 += 2 * KTile;
+      a23 += 2 * KTile;
+      a45 += 2 * KTile;
+      a67 += 2 * KTile;
+      b_ptr += b_k_group_stride;
     }
+
+    store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01);
+    store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
+                       rows23);
+    store_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc,
+                       rows45);
+    store_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc,
+                       rows67);
   }
-}
+};
+
+namespace {
 
-#undef INIT_ACC_ROWPAIR_4
-#undef STORE_ACC_ROWPAIR_4
-#undef BFMMLA_COMPUTE_4
+constexpr int32_t TILE_K = BfmmlaGemm::KTile;
+constexpr int32_t TILE_COLS = 2;
+constexpr int32_t OUTPUT_COLS_PER_BLOCK = BfmmlaGemm::NTile;
+constexpr int32_t K_TOKENS_PER_GROUP = 8;
+constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4;
+constexpr int32_t K_CACHE_K_GROUP_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
+constexpr int32_t B_COL_PAIR_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
 
 }  // namespace
 
-// TileGemm Adapter for Attention
-
 template <typename kv_cache_t, int32_t BlockTokens, int32_t HeadDim>
 class TileGemmNEONBFMMLA {
  public:
@@ -443,49 +286,31 @@ class TileGemmNEONBFMMLA {
                                 [[maybe_unused]] const int32_t block_size,
                                 [[maybe_unused]] const int32_t dynamic_k_size,
                                 const bool accum_c) {
-    static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0);
-    // BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for
-    // API parity with other tile_gemm implementations.
+    static_assert(BlockTokens % 16 == 0);
     if constexpr (head_dim_ct >= 0) {
-      static_assert(head_dim_ct == HeadDim,
-                    "BFMMLA expects head_dim_ct to match HeadDim; PV passes "
-                    "-1 for API parity.");
+      static_assert(head_dim_ct == HeadDim);
     }
 
+    const auto* a = reinterpret_cast<const c10::BFloat16*>(a_tile);
     if constexpr (phase == AttentionGemmPhase::QK) {
-      const int64_t b_reduction_stride = K_INNER_STRIDE;
-      const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE;
-
-      gemm_macro_neon_bfmmla<AttentionGemmPhase::QK, BlockTokens, HeadDim>(
-          reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-          m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride,
-          accum_c);
+      constexpr int64_t b_n_group_stride =
+          (HeadDim / BfmmlaGemm::KTile) * K_CACHE_K_GROUP_STRIDE;
+
+      for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) {
+        const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row);
+        BfmmlaGemm::gemm(a + row * HeadDim, b_tile, c_tile + row * ldc, panel_m,
+                         BlockTokens, HeadDim, 2 * HeadDim, b_n_group_stride,
+                         K_CACHE_K_GROUP_STRIDE, ldc, accum_c);
+      }
     } else {
-      const int64_t b_pair_stride =
-          (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
-
-      // PV gemm with runtime K specialization
-      switch (dynamic_k_size) {
-        case 32:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 32>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 32, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        case 128:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 128>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 128, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        case 256:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 256>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, 256, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
-        default:
-          gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim>(
-              reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
-              m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c);
-          break;
+      const int64_t b_n_group_stride =
+          (block_size / V_TOKENS_PER_ROW_BLOCK) * K_CACHE_K_GROUP_STRIDE;
+
+      for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) {
+        const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row);
+        BfmmlaGemm::gemm(a + row * lda, b_tile, c_tile + row * ldc, panel_m,
+                         HeadDim, dynamic_k_size, 2 * lda, b_n_group_stride,
+                         K_CACHE_K_GROUP_STRIDE, ldc, accum_c);
       }
     }
   }
@@ -512,9 +337,11 @@ class AttentionImplNEONBFMMLA {
   static constexpr int64_t HeadDim = head_dim;
   static constexpr ISA ISAType = isa_type;
   static constexpr bool scale_on_logits = false;
+  static constexpr int64_t VCacheNGroup = OUTPUT_COLS_PER_BLOCK;
+  static constexpr int64_t VCacheKGroupStride = VCacheNGroup * TILE_K;
 
-  static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0);
-  static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0);
+  static_assert(HeadDim % (2 * OUTPUT_COLS_PER_BLOCK) == 0);
+  static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
   static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K");
 
  public:
@@ -527,19 +354,36 @@ class AttentionImplNEONBFMMLA {
     attention_iteration(CPU_ATTENTION_PARAMS);
   }
 
+  struct ProbabilityTokenStore {
+    static constexpr int32_t TokenStride = 2;
+
+    FORCE_INLINE static void store_probabilities(
+        c10::BFloat16* __restrict__ probability,
+        const vec_op::FP32Vec16& values, const int32_t row,
+        const int64_t row_stride) {
+      const int32_t row_in_pair = row & 1;
+      auto* dst = reinterpret_cast<bfloat16_t*>(
+          probability + row_in_pair * (BfmmlaGemm::KTile - row_stride));
+      vst1_bf16(dst, vcvt_bf16_f32(values.reg.val[0]));
+      vst1_bf16(dst + 2 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[1]));
+      vst1_bf16(dst + 4 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[2]));
+      vst1_bf16(dst + 6 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[3]));
+    }
+  };
+
   // Key cache stride per token group (TokenColumn layout; QK)
   static constexpr int64_t k_cache_token_group_stride(
       [[maybe_unused]] const int32_t block_size) {
     static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
     return (BlockSizeAlignment / K_TOKENS_PER_GROUP) *
-           ((head_dim / TILE_K) * K_INNER_STRIDE);
+           ((head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE);
   }
 
   // Value cache stride per token group (TokenRow layout; PV)
   static constexpr int64_t v_cache_token_group_stride(
       [[maybe_unused]] const int32_t block_size) {
     static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0);
-    return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
+    return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride;
   }
 
   // The stride to move to the "next" head_dim group
@@ -550,7 +394,7 @@ class AttentionImplNEONBFMMLA {
     return head_dim * block_size;
   }
 
-  // Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics
+  // Scale Q and write row pairs in BFMMLA reduction order.
   static void copy_q_heads_tile(c10::BFloat16* __restrict__ src,
                                 c10::BFloat16* __restrict__ q_buffer,
                                 const int32_t q_num,
@@ -559,25 +403,42 @@ class AttentionImplNEONBFMMLA {
                                 const int64_t q_head_stride, float scale) {
     constexpr int32_t dim = static_cast<int32_t>(head_dim);
     const float32x4_t scale_vec = vdupq_n_f32(scale);
-
-    for (int32_t qi = 0; qi < q_num; ++qi) {
-      for (int32_t hi = 0; hi < q_heads_per_kv; ++hi) {
-        c10::BFloat16* __restrict__ curr_q =
-            src + qi * q_num_stride + hi * q_head_stride;
-        c10::BFloat16* __restrict__ dst =
-            q_buffer + qi * q_heads_per_kv * head_dim + hi * head_dim;
-
-        for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) {
-          bfloat16x8_t in8 =
-              vld1q_bf16(reinterpret_cast<const bfloat16_t*>(curr_q + i));
-          float32x4_t lo = vmulq_f32(vcvtq_low_f32_bf16(in8), scale_vec);
-          float32x4_t hi = vmulq_f32(vcvtq_high_f32_bf16(in8), scale_vec);
-
-          bfloat16x4_t lo_b = vcvt_bf16_f32(lo);
-          bfloat16x4_t hi_b = vcvt_bf16_f32(hi);
-          bfloat16x8_t out = vcombine_bf16(lo_b, hi_b);
-          vst1q_bf16(reinterpret_cast<bfloat16_t*>(dst + i), out);
+    const bfloat16x4_t zero = vdup_n_bf16(bfloat16_t{});
+    const int32_t row_num = q_num * q_heads_per_kv;
+
+    for (int32_t row = 0; row < row_num; row += 2) {
+      const int32_t q0 = row / q_heads_per_kv;
+      const int32_t h0 = row % q_heads_per_kv;
+      const auto* row0 = reinterpret_cast<const bfloat16_t*>(
+          src + q0 * q_num_stride + h0 * q_head_stride);
+      const bool has_row1 = row + 1 < row_num;
+      const int32_t q1 = (row + 1) / q_heads_per_kv;
+      const int32_t h1 = (row + 1) % q_heads_per_kv;
+      const auto* row1 = has_row1
+                             ? reinterpret_cast<const bfloat16_t*>(
+                                   src + q1 * q_num_stride + h1 * q_head_stride)
+                             : nullptr;
+      auto* dst = reinterpret_cast<bfloat16_t*>(q_buffer + row * head_dim);
+
+      for (int32_t k = 0; k < dim; k += OUTPUT_COLS_PER_BLOCK) {
+        const bfloat16x8_t in0 = vld1q_bf16(row0 + k);
+        const bfloat16x4_t out0_lo =
+            vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in0), scale_vec));
+        const bfloat16x4_t out0_hi =
+            vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in0), scale_vec));
+
+        bfloat16x4_t out1_lo = zero;
+        bfloat16x4_t out1_hi = zero;
+        if (has_row1) {
+          const bfloat16x8_t in1 = vld1q_bf16(row1 + k);
+          out1_lo =
+              vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in1), scale_vec));
+          out1_hi =
+              vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in1), scale_vec));
         }
+
+        vst1q_bf16(dst + 2 * k, vcombine_bf16(out0_lo, out1_lo));
+        vst1q_bf16(dst + 2 * k + 8, vcombine_bf16(out0_hi, out1_hi));
       }
     }
   }
@@ -585,11 +446,12 @@ class AttentionImplNEONBFMMLA {
  public:
   // Reshape and cache K/V into BFMMLA-optimized layouts
   // K cache:
-  // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE]
+  // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K]
+  // [K_CACHE_K_GROUP_STRIDE]
   // - TokenColumn
   // V cache:
-  // [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE]
-  // - TokenRows
+  // [head_dim/VCacheNGroup][block_size/V_TOKENS_PER_ROW_BLOCK]
+  // [VCacheKGroupStride]
   static void reshape_and_cache(
       const c10::BFloat16* __restrict__ key,
       const c10::BFloat16* __restrict__ value,
@@ -604,13 +466,13 @@ class AttentionImplNEONBFMMLA {
       const int64_t block_size,
       [[maybe_unused]] const int64_t block_size_stride,
       const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
-    const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
-    const int64_t v_pair_stride =
-        (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
+    const int64_t k_block_stride = (head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE;
+    const int64_t v_n_group_stride =
+        (block_size / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride;
 
-#pragma omp parallel for
-    for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
-      for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
+#pragma omp parallel for collapse(2)
+    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
+      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
         const int64_t pos = slot_mapping[token_idx];
         if (pos < 0) continue;
 
@@ -636,12 +498,12 @@ class AttentionImplNEONBFMMLA {
               key_base + block_in_block * k_block_stride;
 
           for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) {
-            uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(
-                block_base + hd4 * K_INNER_STRIDE +
-                pair_in_block * V_INNER_STRIDE + lane_base);
-            const uint16_t* src_u16 =
-                reinterpret_cast<const uint16_t*>(key_src + hd4 * TILE_K);
-            vst1_u16(dst_u16, vld1_u16(src_u16));
+            bfloat16_t* dst = reinterpret_cast<bfloat16_t*>(
+                block_base + hd4 * K_CACHE_K_GROUP_STRIDE +
+                pair_in_block * B_COL_PAIR_STRIDE + lane_base);
+            const bfloat16_t* src =
+                reinterpret_cast<const bfloat16_t*>(key_src + hd4 * TILE_K);
+            vst1_bf16(dst, vld1_bf16(src));
           }
         }
 
@@ -655,22 +517,33 @@ class AttentionImplNEONBFMMLA {
               value_cache + block_idx * num_blocks_stride +
               head_idx * cache_head_num_stride;
 
-          const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK;
+          const int64_t token_group = block_offset / V_TOKENS_PER_ROW_BLOCK;
           const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1);
 
-          c10::BFloat16* __restrict row_block_base =
-              value_base + row_block * V_INNER_STRIDE;
-
-          for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) {
-            c10::BFloat16* __restrict dst_val =
-                row_block_base + hd2 * v_pair_stride;
-
-            const uint16_t* src_u16 =
-                reinterpret_cast<const uint16_t*>(value_src);
-            uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(dst_val);
-            dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0];
-            dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] =
-                src_u16[hd2 * TILE_COLS + 1];
+          const auto* src = reinterpret_cast<const bfloat16_t*>(value_src);
+          auto* dst = reinterpret_cast<bfloat16_t*>(value_base);
+          for (int64_t hd8 = 0; hd8 < head_dim / OUTPUT_COLS_PER_BLOCK; ++hd8) {
+            const bfloat16x8_t values =
+                vld1q_bf16(src + hd8 * OUTPUT_COLS_PER_BLOCK);
+            const bfloat16x4_t low = vget_low_bf16(values);
+            const bfloat16x4_t high = vget_high_bf16(values);
+            bfloat16_t* group =
+                dst + hd8 * v_n_group_stride + token_group * VCacheKGroupStride;
+
+            vst1_lane_bf16(group + lane, low, 0);
+            vst1_lane_bf16(group + V_TOKENS_PER_ROW_BLOCK + lane, low, 1);
+            vst1_lane_bf16(group + B_COL_PAIR_STRIDE + lane, low, 2);
+            vst1_lane_bf16(
+                group + B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, low,
+                3);
+            vst1_lane_bf16(group + 2 * B_COL_PAIR_STRIDE + lane, high, 0);
+            vst1_lane_bf16(
+                group + 2 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane,
+                high, 1);
+            vst1_lane_bf16(group + 3 * B_COL_PAIR_STRIDE + lane, high, 2);
+            vst1_lane_bf16(
+                group + 3 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane,
+                high, 3);
           }
         }
       }
@@ -680,4 +553,4 @@ class AttentionImplNEONBFMMLA {
 
 }  // namespace cpu_attention
 
-#endif  // CPU_ATTN_ASIMD_BFMMLA_HPP
+#endif  // CPU_ATTN_NEON_BFMMLA_HPP
diff --git a/csrc/cpu/cpu_fused_moe_int8.cpp b/csrc/cpu/cpu_fused_moe_int8.cpp
index 649317b4786a..27d0f6a7d539 100644
--- a/csrc/cpu/cpu_fused_moe_int8.cpp
+++ b/csrc/cpu/cpu_fused_moe_int8.cpp
@@ -26,10 +26,23 @@
   #define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON:
 #endif
 
+#if defined(__powerpc64__)
+  #include "cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp"
+  #define VSX_DISPATCH(SCALAR_TYPE, ...)                                   \
+    case cpu_utils::ISA::VSX: {                                            \
+      using gemm_t =                                                       \
+          cpu_micro_gemm::MicroGemmINT8<cpu_utils::ISA::VSX, SCALAR_TYPE>; \
+      return __VA_ARGS__();                                                \
+    }
+#else
+  #define VSX_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::VSX:
+#endif
+
 #define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \
   [&] {                                                        \
     switch (ISA_TYPE) {                                        \
       NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__)                  \
+      VSX_DISPATCH(SCALAR_TYPE, __VA_ARGS__)                   \
       default: {                                               \
         TORCH_CHECK(false, "Invalid CPU ISA type.");           \
       }                                                        \
@@ -81,9 +94,9 @@ void fused_moe_int8_impl(
 
   const int32_t thread_num = cpu_utils::get_max_threads();
   const int32_t w13_input_buffer_size = cpu_utils::round_up<64>(
-      gemm_m_tile_size * input_size_13 * sizeof(int8_t));
-  const int32_t w2_input_buffer_size =
-      cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t));
+      (gemm_m_tile_size * input_size_13 + 64) * sizeof(int8_t));
+  const int32_t w2_input_buffer_size = cpu_utils::round_up<64>(
+      (gemm_m_tile_size * input_size_2 + 64) * sizeof(int8_t));
 
   const int32_t w13_n_tile_size = [&]() {
     const int64_t cache_size = cpu_utils::get_available_l2_size();
diff --git a/csrc/cpu/cpu_isa.cpp b/csrc/cpu/cpu_isa.cpp
new file mode 100644
index 000000000000..1c4a8faf1918
--- /dev/null
+++ b/csrc/cpu/cpu_isa.cpp
@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+// AMX-FP8 is reported at CPUID leaf 0x1E, subleaf 1, EAX[4].
+#if defined(__x86_64__) || defined(_M_X64)
+  #include <cpuid.h>
+
+bool runtime_has_amx_fp8() {
+  unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;
+  if (__get_cpuid_max(0, nullptr) < 0x1e) return false;
+  if (__get_cpuid_count(0x1e, 1, &eax, &ebx, &ecx, &edx)) {
+    return (eax >> 4) & 1u;
+  }
+  return false;
+}
+
+bool cpu_has_amx_fp8() { return runtime_has_amx_fp8(); }
+#else
+bool cpu_has_amx_fp8() { return false; }
+#endif
diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp
index d46906130e92..eaea921d98b8 100644
--- a/csrc/cpu/cpu_types_vsx.hpp
+++ b/csrc/cpu/cpu_types_vsx.hpp
@@ -13,10 +13,10 @@ namespace vec_op {
 struct fp8_e4m3_tag {};
 struct fp8_e5m2_tag {};
 
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)            \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+// Note: FP16 (Half) is not supported on POWER VSX — no FP16Vec16 type.
 
 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
   AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -416,6 +416,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
       tmp.val[0] = two_x;
       tmp.val[1] = two_x;
       FP32Vec8 temp_vec(tmp);
+      FP32Vec8 exp_vec = temp_vec.exp();
       vector float e = temp_vec.exp().reg.val[0];
 
       vector float num = vec_sub(e, one);
@@ -634,6 +635,16 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     reg.val[3] = vec_ctf(v.reg.val[3], 0);
   }
 
+  explicit FP32Vec16(const float* __restrict__ base, const INT32Vec16& index) {
+    INT32Vec16::AliasReg idx;
+    idx.reg = index.reg;
+    AliasReg ar;
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      ar.values[i] = base[idx.values[i]];
+    }
+    reg = ar.reg;
+  }
+
   FP32Vec16 operator*(const FP32Vec16& b) const {
     return FP32Vec16(f32x4x4_t({vec_mul(reg.val[0], b.reg.val[0]),
                                 vec_mul(reg.val[1], b.reg.val[1]),
@@ -648,6 +659,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
                                 vec_add(reg.val[3], b.reg.val[3])}));
   }
 
+  FP32Vec16 operator-() const {
+    const __vector float zero = vec_splats(0.0f);
+    return FP32Vec16(
+        f32x4x4_t({vec_sub(zero, reg.val[0]), vec_sub(zero, reg.val[1]),
+                   vec_sub(zero, reg.val[2]), vec_sub(zero, reg.val[3])}));
+  }
+
   FP32Vec16 operator-(const FP32Vec16& b) const {
     return FP32Vec16(f32x4x4_t({vec_sub(reg.val[0], b.reg.val[0]),
                                 vec_sub(reg.val[1], b.reg.val[1]),
@@ -655,11 +673,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
                                 vec_sub(reg.val[3], b.reg.val[3])}));
   }
 
-  FP32Vec16 operator-() const {
-    return FP32Vec16(
-        f32x4x4_t({-reg.val[0], -reg.val[1], -reg.val[2], -reg.val[3]}));
-  }
-
   FP32Vec16 operator/(const FP32Vec16& b) const {
     return FP32Vec16(f32x4x4_t({vec_div(reg.val[0], b.reg.val[0]),
                                 vec_div(reg.val[1], b.reg.val[1]),
diff --git a/csrc/cpu/dnnl_helper.cpp b/csrc/cpu/dnnl_helper.cpp
index 14c136dcbbf0..1ac0401f540d 100644
--- a/csrc/cpu/dnnl_helper.cpp
+++ b/csrc/cpu/dnnl_helper.cpp
@@ -172,7 +172,8 @@ struct hash<MatMulPrimitiveHandler::ClassMatmulCacheKey> {
   size_t operator()(
       const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const {
     return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^
-           hash<int>()(static_cast<int>(val.b_type));
+           hash<int>()(static_cast<int>(val.b_type)) ^
+           hash<int>()(static_cast<int>(val.c_type));
   }
 };
 
@@ -202,7 +203,7 @@ bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l,
 bool operator==(const MatMulPrimitiveHandler::ClassMatmulCacheKey& l,
                 const MatMulPrimitiveHandler::ClassMatmulCacheKey& r) {
   return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size &&
-         l.b_type == r.b_type;
+         l.b_type == r.b_type && l.c_type == r.c_type;
 }
 
 bool operator==(const MatMulPrimitiveHandler::MSizeCacheKey& l,
@@ -482,8 +483,10 @@ void MatMulPrimitiveHandler::execute(ExecArgs& args) {
 dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache(
     const MSizeCacheKey& key) {
   if (m_size_cache_.get() == nullptr) {
-    ClassMatmulCacheKey class_key = {
-        .b_n_size = b_n_size_, .b_k_size = b_k_size_, .b_type = b_type_};
+    ClassMatmulCacheKey class_key = {.b_n_size = b_n_size_,
+                                     .b_k_size = b_k_size_,
+                                     .b_type = b_type_,
+                                     .c_type = c_type_};
     m_size_cache_ =
         get_matul_class_primitive_cache(class_key, primitive_cache_size_);
   }
diff --git a/csrc/cpu/dnnl_helper.h b/csrc/cpu/dnnl_helper.h
index b841ed73058e..05fb116ecd1e 100644
--- a/csrc/cpu/dnnl_helper.h
+++ b/csrc/cpu/dnnl_helper.h
@@ -176,6 +176,7 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
     dnnl_dim_t b_n_size;
     dnnl_dim_t b_k_size;
     dnnl::memory::data_type b_type;
+    dnnl::memory::data_type c_type;
 
     friend bool operator==(const ClassMatmulCacheKey& l,
                            const ClassMatmulCacheKey& r);
@@ -203,6 +204,8 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
  public:
   MatMulPrimitiveHandler(const Args& args);
 
+  dnnl::memory::data_type get_output_type() const { return c_type_; }
+
   void execute(ExecArgs& args);
 
  private:
diff --git a/csrc/cpu/dnnl_kernels.cpp b/csrc/cpu/dnnl_kernels.cpp
index 6dda09296160..e81d365628c0 100644
--- a/csrc/cpu/dnnl_kernels.cpp
+++ b/csrc/cpu/dnnl_kernels.cpp
@@ -513,6 +513,17 @@ int64_t create_onednn_mm_handler(const torch::Tensor& b,
                                  args.ab_type = get_dnnl_type<scalar_t>();
                                });
 
+#ifdef VLLM_USE_ACL
+  // TODO(fadara01): this is a hot-fix, remove once problem is addressed in the
+  // AArch64 stack. hybrid kernels in ACL which are the fastest for LM head
+  // layers are only enabled for BF16 x BF16 -> FP32 problem desc as they can't
+  // convert FP32->BF16 internally. In this case do: oneDNN matmul with
+  // bf16:bf16:fp32 then cast output down to bf16.
+  if (b.scalar_type() == at::kBFloat16 && b.size(1) >= 64000) {
+    args.c_type = dnnl::memory::data_type::f32;
+  }
+#endif
+
   return reinterpret_cast<int64_t>(new MatMulPrimitiveHandler(args));
 }
 
@@ -527,12 +538,22 @@ void onednn_mm(torch::Tensor& c,        // [M, OC], row-major
   MatMulPrimitiveHandler* ptr =
       reinterpret_cast<MatMulPrimitiveHandler*>(handler_tensor.item<int64_t>());
 
+  torch::Tensor dnnl_output = c;
+  // AArch64 specific case where we do a bf16 x bf16 -> fp32 matmul
+  const bool convert_output =
+      ptr->get_output_type() == dnnl::memory::data_type::f32 &&
+      c.scalar_type() != at::kFloat;
+  if (convert_output) {
+    dnnl_output = torch::empty(c.sizes(), c.options().dtype(at::kFloat));
+  }
+
 // ACL matmuls expect contiguous source tensors
 #ifdef VLLM_USE_ACL
   torch::Tensor a_contig = a.contiguous();
 #endif
 
   MatMulPrimitiveHandler::ExecArgs exec_args;
+  exec_args.c_ptr = dnnl_output.data_ptr();
 
 #ifdef VLLM_USE_ACL
   exec_args.a_m_size = a_contig.size(0);
@@ -548,7 +569,7 @@ void onednn_mm(torch::Tensor& c,        // [M, OC], row-major
 #ifdef VLLM_USE_ACL
       // ACL matmuls in oneDNN do not support a bias.
       // We handle a matmul with bias by doing: c = bias; c += matmul(a, b)
-      c.copy_(bias.value());
+      dnnl_output.copy_(bias.value());
 #else
       exec_args.bias_ptr = bias->data_ptr<scalar_t>();
 #endif
@@ -563,8 +584,11 @@ void onednn_mm(torch::Tensor& c,        // [M, OC], row-major
     exec_args.a_ptr = a.data_ptr<scalar_t>();
 
 #endif
-    exec_args.c_ptr = c.data_ptr<scalar_t>();
 
     ptr->execute(exec_args);
   });
+
+  if (convert_output) {
+    c.copy_(dnnl_output);
+  }
 }
diff --git a/csrc/cpu/generate_cpu_attn_dispatch.py b/csrc/cpu/generate_cpu_attn_dispatch.py
index a82a27b844a4..12a14ef58774 100644
--- a/csrc/cpu/generate_cpu_attn_dispatch.py
+++ b/csrc/cpu/generate_cpu_attn_dispatch.py
@@ -1,9 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Generate CPU attention dispatch switch cases and kernel instantiations.
-"""
+"""Generate CPU attention dispatch switch cases and kernel instantiations."""
 
 import os
 
diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp
new file mode 100644
index 000000000000..60333537281d
--- /dev/null
+++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp
@@ -0,0 +1,287 @@
+#ifndef CPU_MICRO_GEMM_INT8_VSX_HPP
+#define CPU_MICRO_GEMM_INT8_VSX_HPP
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <type_traits>
+#include <utility>
+
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
+
+#include <altivec.h>
+#include <c10/util/BFloat16.h>
+#include <c10/util/Exception.h>
+
+namespace cpu_micro_gemm {
+
+namespace vsx_int8 {
+
+constexpr int32_t K = 4;
+constexpr int32_t Mr = 4;
+constexpr int32_t Nr = 32;
+
+// ---------------------------------------------------------------------------
+// 8-Accumulator MMA 4x32 micro-GEMM (POWER10 xvi8ger4pp).
+//
+// Computes 4 rows x 32 columns (128 outputs) using all 8 hardware accumulators.
+// Operates on pre-biased uint8 weights (B + 128) and subtracts 128 * a_sums.
+// Works for all m in [1, 4] with zero scalar loop overhead.
+// ---------------------------------------------------------------------------
+__attribute__((target("cpu=power10"))) FORCE_INLINE void gemm_mma_4x32(
+    const int8_t* __restrict__ a_packed, const uint8_t* __restrict__ b_panel,
+    const int32_t* __restrict__ a_sums, int32_t* __restrict__ c_ptr,
+    const int32_t m, const int32_t k_size, const int64_t ldc) {
+  TORCH_CHECK(m > 0 && m <= Mr);
+  TORCH_CHECK(k_size % K == 0);
+
+  __vector_quad acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7;
+  __builtin_mma_xxsetaccz(&acc0);
+  __builtin_mma_xxsetaccz(&acc1);
+  __builtin_mma_xxsetaccz(&acc2);
+  __builtin_mma_xxsetaccz(&acc3);
+  __builtin_mma_xxsetaccz(&acc4);
+  __builtin_mma_xxsetaccz(&acc5);
+  __builtin_mma_xxsetaccz(&acc6);
+  __builtin_mma_xxsetaccz(&acc7);
+
+  const int8_t* __restrict__ ap = a_packed;
+  const uint8_t* __restrict__ bp = b_panel;
+
+#pragma GCC unroll 4
+  for (int32_t k = 0; k < k_size; k += 4) {
+    const __vector unsigned char avec =
+        vec_xl(0, reinterpret_cast<const unsigned char*>(ap));
+    __builtin_mma_xvi8ger4pp(&acc0, avec, vec_xl(0, bp));
+    __builtin_mma_xvi8ger4pp(&acc1, avec, vec_xl(16, bp));
+    __builtin_mma_xvi8ger4pp(&acc2, avec, vec_xl(32, bp));
+    __builtin_mma_xvi8ger4pp(&acc3, avec, vec_xl(48, bp));
+    __builtin_mma_xvi8ger4pp(&acc4, avec, vec_xl(64, bp));
+    __builtin_mma_xvi8ger4pp(&acc5, avec, vec_xl(80, bp));
+    __builtin_mma_xvi8ger4pp(&acc6, avec, vec_xl(96, bp));
+    __builtin_mma_xvi8ger4pp(&acc7, avec, vec_xl(112, bp));
+    ap += 16;
+    bp += 128;
+  }
+
+  __vector signed int res[4];
+#define STORE_ACC(ACC, CG_IDX)                                                \
+  {                                                                           \
+    __builtin_mma_disassemble_acc(res, &ACC);                                 \
+    for (int r = 0; r < m; ++r) {                                             \
+      const __vector signed int vcorr = vec_splats(128 * a_sums[r]);          \
+      const __vector signed int val = vec_sub(res[r], vcorr);                 \
+      vec_xst(val, 0,                                                         \
+              reinterpret_cast<signed int*>(c_ptr + r * ldc + (CG_IDX) * 4)); \
+    }                                                                         \
+  }
+  STORE_ACC(acc0, 0);
+  STORE_ACC(acc1, 1);
+  STORE_ACC(acc2, 2);
+  STORE_ACC(acc3, 3);
+  STORE_ACC(acc4, 4);
+  STORE_ACC(acc5, 5);
+  STORE_ACC(acc6, 6);
+  STORE_ACC(acc7, 7);
+#undef STORE_ACC
+}
+
+}  // namespace vsx_int8
+
+// =============================================================================
+// MicroGemmINT8 specialisation for VSX (POWER10+)
+// =============================================================================
+template <typename scalar_t>
+class MicroGemmINT8<cpu_utils::ISA::VSX, scalar_t> {
+ public:
+  static constexpr int32_t K = vsx_int8::K;    // 4
+  static constexpr int32_t Mr = vsx_int8::Mr;  // 4
+  static constexpr int32_t Nr = vsx_int8::Nr;  // 32
+  static constexpr int32_t NrGemv = Nr;
+  static constexpr int32_t MaxMSize = 8;
+  static constexpr int32_t NSize = Nr;
+  static constexpr int32_t WeightOCGroupSize = Nr;
+  static_assert(MaxMSize % Mr == 0);
+
+  // -------------------------------------------------------------------------
+  // quantize_row: scalar_t -> INT8 + per-token scale
+  // Single-pass 16-element vector conversion with aligned 16-byte stores.
+  // -------------------------------------------------------------------------
+  static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output,
+                                        float& scale, const int32_t size) {
+    TORCH_CHECK_EQ(size % 16, 0);
+
+    const __vector signed short z = vec_splats((signed short)0);
+
+    auto load_lo_hi =
+        [&](int32_t i) -> std::pair<__vector float, __vector float> {
+      if constexpr (std::is_same_v<scalar_t, float>) {
+        const float* p = reinterpret_cast<const float*>(input) + i;
+        return {vec_xl(0, p), vec_xl(16, p)};
+      } else {
+        const signed short* p =
+            reinterpret_cast<const signed short*>(input + i);
+        const __vector signed short raw = (__vector signed short)vec_xl(0, p);
+        return {(__vector float)vec_mergeh(z, raw),
+                (__vector float)vec_mergel(z, raw)};
+      }
+    };
+
+    __vector float max_v = vec_splats(0.0f);
+    for (int32_t i = 0; i < size; i += 8) {
+      auto [lo, hi] = load_lo_hi(i);
+      max_v = vec_max(max_v, vec_abs(lo));
+      max_v = vec_max(max_v, vec_abs(hi));
+    }
+    __vector float t = vec_max(max_v, vec_sld(max_v, max_v, 8));
+    t = vec_max(t, vec_sld(t, t, 4));
+    const float abs_max = std::max(vec_extract(t, 0), 1.0e-7f);
+    scale = abs_max / 127.0f;
+    const __vector float inv_scale = vec_splats(127.0f / abs_max);
+
+    for (int32_t i = 0; i < size; i += 16) {
+      auto [lo0, hi0] = load_lo_hi(i);
+      auto [lo1, hi1] = load_lo_hi(i + 8);
+
+      const __vector signed int i0 =
+          vec_cts(vec_round(vec_mul(lo0, inv_scale)), 0);
+      const __vector signed int i1 =
+          vec_cts(vec_round(vec_mul(hi0, inv_scale)), 0);
+      const __vector signed int i2 =
+          vec_cts(vec_round(vec_mul(lo1, inv_scale)), 0);
+      const __vector signed int i3 =
+          vec_cts(vec_round(vec_mul(hi1, inv_scale)), 0);
+
+      const __vector signed short s0 = vec_packs(i0, i1);
+      const __vector signed short s1 = vec_packs(i2, i3);
+      const __vector signed char s8 = vec_packs(s0, s1);
+
+      vec_xst(s8, 0, reinterpret_cast<signed char*>(output + i));
+    }
+  }
+
+  // -------------------------------------------------------------------------
+  // dequantize_tile: INT32 accumulator -> FP32 scaled output
+  // -------------------------------------------------------------------------
+  static FORCE_INLINE void dequantize_tile(
+      int32_t* input, float* output, const float* __restrict__ input_scales,
+      const float* __restrict__ weight_scales, const int32_t m, const int32_t n,
+      const int32_t stride) {
+    TORCH_CHECK_EQ(n % 4, 0);
+    for (int32_t mi = 0; mi < m; ++mi) {
+      const __vector float a_scale = vec_splats(input_scales[mi]);
+      for (int32_t ni = 0; ni < n; ni += 4) {
+        const __vector signed int i32 = vec_xl(0, input + mi * stride + ni);
+        const __vector float wscale = vec_xl(0, weight_scales + ni);
+        vec_xst(vec_mul(vec_ctf(i32, 0), vec_mul(a_scale, wscale)), 0,
+                output + mi * stride + ni);
+      }
+    }
+  }
+
+  // -------------------------------------------------------------------------
+  // pack_input_from_rows: pack rows into 4-row MMA tiles and compute row sums
+  // Header: 64 bytes (contains int32 a_sums[0..7])
+  // Data:   a_packed + 64 (contains 4x4 tiles of A)
+  // -------------------------------------------------------------------------
+  static void pack_input_from_rows(const int8_t* const* __restrict__ rows,
+                                   int8_t* __restrict__ a_packed,
+                                   const int32_t m, const int32_t k) {
+    TORCH_CHECK(m > 0 && m <= MaxMSize);
+    TORCH_CHECK_EQ(k % K, 0);
+
+    int32_t* a_sums = reinterpret_cast<int32_t*>(a_packed);
+    for (int32_t r = 0; r < 8; ++r) a_sums[r] = 0;
+
+    int8_t* a_data = a_packed + 64;
+
+    for (int32_t row_base = 0; row_base < m; row_base += Mr) {
+      const int32_t panel_m = std::min(Mr, m - row_base);
+      const int8_t* const* panel_rows = rows + row_base;
+      int32_t* panel_sums = a_sums + row_base;
+      int8_t* dst_panel = a_data + (row_base / Mr) * (k * Mr);
+
+      for (int32_t ki = 0; ki < k; ki += 4) {
+        int8_t* dst = dst_panel + (ki / 4) * 16;
+        for (int32_t r = 0; r < 4; ++r) {
+          if (r < panel_m) {
+            const int8_t* src = panel_rows[r] + ki;
+            dst[r * 4 + 0] = src[0];
+            dst[r * 4 + 1] = src[1];
+            dst[r * 4 + 2] = src[2];
+            dst[r * 4 + 3] = src[3];
+            panel_sums[r] +=
+                static_cast<int32_t>(src[0]) + src[1] + src[2] + src[3];
+          } else {
+            dst[r * 4 + 0] = 0;
+            dst[r * 4 + 1] = 0;
+            dst[r * 4 + 2] = 0;
+            dst[r * 4 + 3] = 0;
+          }
+        }
+      }
+    }
+  }
+
+  // -------------------------------------------------------------------------
+  // pack_weight: [output_size, input_size] INT8 -> packed layout
+  //   Pre-biases with +128 in [N/32][K/4][8 col-groups][4 cols x 4 k] layout
+  // -------------------------------------------------------------------------
+  static void pack_weight(const int8_t* __restrict__ weight,
+                          int8_t* __restrict__ packed_weight,
+                          const int32_t output_size, const int32_t input_size) {
+    TORCH_CHECK_EQ(output_size % NSize, 0);
+    TORCH_CHECK_EQ(input_size % K, 0);
+
+    uint8_t* dst_base = reinterpret_cast<uint8_t*>(packed_weight);
+
+    for (int32_t o = 0; o < output_size; o += Nr) {
+      uint8_t* dst_panel = dst_base + o * input_size;
+      for (int32_t k = 0; k < input_size; k += 4) {
+        for (int32_t cg = 0; cg < 8; ++cg) {
+          uint8_t* dst = dst_panel + (k / 4) * 128 + cg * 16;
+          for (int32_t col = 0; col < 4; ++col) {
+            for (int32_t ki = 0; ki < 4; ++ki) {
+              const int8_t w = weight[(o + cg * 4 + col) * input_size + k + ki];
+              dst[col * 4 + ki] =
+                  static_cast<uint8_t>(static_cast<int>(w) + 128);
+            }
+          }
+        }
+      }
+    }
+  }
+
+  // -------------------------------------------------------------------------
+  // gemm: 8-accumulator MMA GEMM over 32-column tiles
+  // -------------------------------------------------------------------------
+  void gemm(const int8_t* __restrict__ a_packed,
+            const int8_t* __restrict__ b_packed, int32_t* __restrict__ c,
+            const int32_t m, const int32_t k, const int64_t b_n_group_stride,
+            const int64_t ldc) const {
+    TORCH_CHECK(m > 0 && m <= MaxMSize);
+    TORCH_CHECK_EQ(k % K, 0);
+    (void)b_n_group_stride;
+
+    const int32_t* a_sums = reinterpret_cast<const int32_t*>(a_packed);
+    const int8_t* a_data = a_packed + 64;
+    const uint8_t* b_u8 = reinterpret_cast<const uint8_t*>(b_packed);
+
+    for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr) {
+      const uint8_t* b_panel = b_u8 + n_idx * k;
+
+      for (int32_t row_base = 0; row_base < m; row_base += Mr) {
+        const int32_t panel_m = std::min(Mr, m - row_base);
+        const int8_t* a_panel = a_data + (row_base / Mr) * (k * Mr);
+        const int32_t* panel_sums = a_sums + row_base;
+
+        vsx_int8::gemm_mma_4x32(a_panel, b_panel, panel_sums,
+                                c + row_base * ldc + n_idx, panel_m, k, ldc);
+      }
+    }
+  }
+};
+
+}  // namespace cpu_micro_gemm
+
+#endif  // CPU_MICRO_GEMM_INT8_VSX_HPP
diff --git a/csrc/cpu/sgl-kernels/common.h b/csrc/cpu/sgl-kernels/common.h
index 6f27ba2d5c10..acf75912819f 100644
--- a/csrc/cpu/sgl-kernels/common.h
+++ b/csrc/cpu/sgl-kernels/common.h
@@ -12,11 +12,45 @@
 #include <ATen/ATen.h>
 #include <ATen/Dispatch.h>
 #include <ATen/Parallel.h>
+#include <cstring>
 
 #if defined(_OPENMP)
 #include <omp.h>
 #endif
 
+// Adapted from sglang: dispatch utilities for FP8 W8A8 kernels
+#define AT_DISPATCH_OUT_TYPES(TYPE, NAME, ...)                                          \
+  AT_DISPATCH_SWITCH(                                                                   \
+      TYPE,                                                                             \
+      NAME,                                                                             \
+      AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::Float, out_t, __VA_ARGS__)        \
+          AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::BFloat16, out_t, __VA_ARGS__) \
+              AT_PRIVATE_CASE_TYPE_USING_HINT(at::ScalarType::Half, out_t, __VA_ARGS__))
+
+#define AT_DISPATCH_CASE_ENUM_NO_RETURN(VALUE, TYPE, HINT, ...) \
+  case VALUE: {                                                 \
+    constexpr TYPE HINT = VALUE;                                \
+    __VA_ARGS__;                                                \
+    break;                                                      \
+  }
+
+#define AT_DISPATCH_BOOL_NO_RETURN(VALUE, NAME, HINT, ...)            \
+  [&]() {                                                             \
+    switch (VALUE) {                                                  \
+      AT_DISPATCH_CASE_ENUM_NO_RETURN(true, bool, HINT, __VA_ARGS__)  \
+      AT_DISPATCH_CASE_ENUM_NO_RETURN(false, bool, HINT, __VA_ARGS__) \
+    }                                                                 \
+  }()
+
+#define AT_DISPATCH_QUANT_MODE_NO_RETURN(MODE, NAME, HINT, ...)           \
+  [&]() {                                                                 \
+    switch (MODE) {                                                       \
+      AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_TENSOR, int, HINT, __VA_ARGS__) \
+      AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_ROW, int, HINT, __VA_ARGS__)    \
+      AT_DISPATCH_CASE_ENUM_NO_RETURN(PER_GROUP, int, HINT, __VA_ARGS__)  \
+    }                                                                     \
+  }()
+
 namespace {
 
 // dispatch bool
@@ -246,26 +280,10 @@ inline int get_thread_num() {
 // balance payload across each thread
 template <typename T>
 inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
-#if 0
-    // onednn partition pattern
-    T& n_my = n_end;
-    if (nth <= 1 || n == 0) {
-        n_start = 0;
-        n_my = n;
-    } else {
-        T n1 = div_up(n, nth);
-        T n2 = n1 - 1;
-        T T1 = n - n2 * nth;
-        n_my = ith < T1 ? n1 : n2;
-        n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2;
-    }
-    n_end += n_start;
-#else
   // pytorch aten partition pattern
   T n_my = div_up(n, nth);
   n_start = ith * n_my;
   n_end = std::min(n_start + n_my, n);
-#endif
 }
 
 template <typename func_t>
@@ -472,4 +490,60 @@ inline T* conditional_data_ptr(const std::optional<at::Tensor>& opt) {
   return opt.has_value() ? opt.value().data_ptr<T>() : nullptr;
 }
 
+// Adapted from sglang: M blocking utilities for FP8 W8A8 kernels
+inline int64_t get_m_block(int64_t M) {
+  if (M <= 48) {
+    return M;
+  } else if (M < 64) {
+    return 32;
+  } else if (M < 96) {
+    return 64;
+  } else {
+    return 128;
+  }
+}
+
+inline std::tuple<bool, int64_t, int64_t, int64_t> get_m_blocking(int64_t M) {
+  bool parallel_on_M = M > 128;
+  int64_t block_m = get_m_block(M);
+  int64_t Mc = (M + block_m - 1) / block_m;
+  int64_t Mc_parallel = parallel_on_M ? Mc : 1;
+  return std::make_tuple(parallel_on_M, block_m, Mc, Mc_parallel);
+}
+
+#if defined(CPU_CAPABILITY_AVX512)
+template <typename T>
+void zero_buffer(T* data, int64_t size) {
+  const int32_t vec_size = at::vec::Vectorized<T>::size();
+  auto zero_vec = at::vec::Vectorized<T>(0);
+  int64_t d = 0;
+  for (; d < size - (size % vec_size); d += vec_size) {
+    zero_vec.store(data + d);
+  }
+  if (d < size) {
+    zero_vec.store(data + d, size - d);
+  }
+}
+#else
+template <typename T>
+void zero_buffer(T* data, int64_t size) {
+  memset(data, 0, sizeof(T) * size);
+}
+#endif
+
+#if defined(__x86_64__)
+// CPU_CAPABILITY_AVX10_2 is defined by cmake only when the compiler accepts
+// -mavx10.2; without it the builtin below is not merely absent but a hard
+// error, so the AVX512 paths have to stand in.
+bool avx10_2_available() {
+#ifdef CPU_CAPABILITY_AVX10_2
+  // __builtin_cpu_supports returns the masked feature bit rather than 0/1.
+  static const bool available = __builtin_cpu_supports("avx10.2");
+  return available;
+#else
+  return false;
+#endif
+}
+#endif
+
 }  // anonymous namespace
diff --git a/csrc/cpu/sgl-kernels/gemm.h b/csrc/cpu/sgl-kernels/gemm.h
index db7ad884d09a..dc5f9c585d5d 100644
--- a/csrc/cpu/sgl-kernels/gemm.h
+++ b/csrc/cpu/sgl-kernels/gemm.h
@@ -94,7 +94,7 @@ enum class CPUActMethod : int {
   gelu_and_mul = 2,
 };
 
-enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };
+enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4, FP8_W8A8 = 5 };
 
 constexpr bool operator==(CPUQuantMethod a, int64_t b) {
   return static_cast<int64_t>(a) == b;
@@ -395,3 +395,55 @@ void tinygemm_kernel(
     bool brg,
     int64_t block_size_K,
     bool do_unpack = true);
+
+// FP8 W8A8 tinygemm kernel (defined in gemm_fp8_w8a8.cpp)
+void tinygemm_kernel(
+    float* C,
+    const at::Float8_e4m3fn* A,
+    const float* scales_a,
+    const at::Float8_e4m3fn* B,
+    const float* scales_b,
+    int64_t M,
+    int64_t K,
+    int64_t lda,
+    int64_t ldc,
+    int64_t ldsa,
+    float* ukernel_buf,
+    at::BFloat16* dqA_buf,
+    at::BFloat16* dqB_buf);
+
+// FP8 W8A8 quantize activation (defined in gemm_fp8_w8a8.cpp)
+std::tuple<at::Tensor, at::Tensor> quantize_fp8e4m3_vec(
+    const at::Tensor& t,
+    bool channelwise,
+    c10::optional<at::Tensor> scale_opt);
+
+// FP8 W8A8 fused MoE kernel (defined in moe_fp8_w8a8.cpp)
+template <typename scalar_t>
+void fused_experts_fp8_a8_kernel_impl(
+    scalar_t* __restrict__ output,
+    scalar_t* __restrict__ ic0,
+    scalar_t* __restrict__ ic1,
+    scalar_t* __restrict__ ic2,
+    at::Float8_e4m3fn* __restrict__ A_tmp,
+    scalar_t* __restrict__ B_tmp,
+    float* __restrict__ C_tmp,
+    float* __restrict__ Ukernel_tmp,
+    const at::Float8_e4m3fn* __restrict__ input,
+    const at::Float8_e4m3fn* __restrict__ packed_w1,
+    const at::Float8_e4m3fn* __restrict__ packed_w2,
+    const float* __restrict__ As,
+    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);
diff --git a/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp b/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp
new file mode 100644
index 000000000000..2014a29fa295
--- /dev/null
+++ b/csrc/cpu/sgl-kernels/gemm_fp8_w8a8.cpp
@@ -0,0 +1,1099 @@
+// Adapted from
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
+
+// clang-format off
+
+#include <ATen/native/CPUBlas.h>
+#include <c10/util/Unroll.h>
+#include <torch/all.h>
+
+#include "common.h"
+#include "vec.h"
+
+// Runtime check for AMX-FP8, implemented in cpu_isa.cpp.
+extern bool cpu_has_amx_fp8();
+
+namespace {
+
+#define BLOCK_N 32
+
+#define PER_TENSOR 1
+#define PER_ROW 2
+#define PER_GROUP 3
+
+static bool cpublas_checked = false;
+static bool cpublas_can_pack = false;
+
+bool cpublas_could_pack() {
+  // the could_pack check requires AMX support implicitly
+  if (cpublas_checked) {
+    return cpublas_can_pack;
+  }
+  cpublas_can_pack = cpu_has_amx_fp8() && at::native::cpublas::could_pack(at::kFloat8_e4m3fn);
+  cpublas_checked = true;
+  return cpublas_can_pack;
+}
+
+#if defined(CPU_CAPABILITY_AVX512)
+static void cvt_f8e4m3_to_bf16(
+    const at::Float8_e4m3fn* __restrict__ in, at::BFloat16* out, int64_t rows, int64_t cols, int64_t stride) {
+  if (stride == cols) {
+    // A contiguous buffer
+    size_t len = rows * cols;
+    size_t i = 0;
+    for (; i < len; i += 32) {
+      __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[i]);
+      __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec);
+      _mm512_storeu_si512((__m512i*)(out + i), (__m512i)bf16_vec);
+    }
+    for (; i < len; ++i) {
+      out[i] = (at::BFloat16)in[i];
+    }
+  } else {
+    // Non-contiguous. Access each row with stride
+    TORCH_CHECK(stride > cols);
+    for (int r = 0; r < rows; ++r) {
+      size_t i = 0;
+      size_t vec_len = cols / 32 * 32;
+      for (; i < vec_len; i += 32) {
+        __m256i fp8_vec = _mm256_loadu_si256((__m256i*)&in[r * stride + i]);
+        __m512bh bf16_vec = cvt_e4m3_bf16_intrinsic_no_nan(fp8_vec);
+        _mm512_storeu_si512((__m512i*)(out + r * cols + i), (__m512i)bf16_vec);
+      }
+      for (; i < cols; ++i) {
+        out[r * cols + i] = (at::BFloat16)in[r * stride + i];
+      }
+    }
+  }
+}
+
+// accumulate and store result to buffer
+// if act/wei are per_group quantized, apply scales
+template <bool accum, int64_t N, int act_quant_mode, int wei_quant_mode>
+static void _accumulate_result(
+    float* __restrict__ output,
+    const float* __restrict__ input,
+    const float* __restrict__ scale_a,
+    const float* __restrict__ scale_b,
+    int M,
+    int ldi,
+    int ldo,
+    int ldsa = 1) {
+  float a_scale, b_scale;
+  __m512 va_scale;
+  __m512 vb_scale;
+  for (int m = 0; m < M; ++m) {
+    if constexpr (act_quant_mode == PER_GROUP) {
+      a_scale = *(scale_a + m * ldsa);
+      va_scale = _mm512_set1_ps(a_scale);
+    }
+    constexpr int N_UNROLL = N / 16;
+    c10::ForcedUnroll<N_UNROLL>{}([&](auto i) {
+      constexpr int n = i * 16;
+      __m512 vc_f = _mm512_loadu_ps(input + m * ldi + n);
+      if constexpr (act_quant_mode == PER_GROUP) {
+        vc_f = _mm512_mul_ps(vc_f, va_scale);
+      }
+      if constexpr (wei_quant_mode == PER_GROUP) {
+        vb_scale = _mm512_loadu_ps(scale_b + n);
+        vc_f = _mm512_mul_ps(vc_f, vb_scale);
+      }
+      if constexpr (accum) {
+        __m512 vo = _mm512_loadu_ps(output + m * ldo + n);
+        _mm512_storeu_ps(output + m * ldo + n, _mm512_add_ps(vo, vc_f));
+      } else {
+        _mm512_storeu_ps(output + m * ldo + n, vc_f);
+      }
+    });
+    constexpr int tail_start = N / 16 * 16;
+    for (int n = tail_start; n < N; ++n) {
+      float dq_val = input[m * ldi + n];
+      if constexpr (act_quant_mode == PER_GROUP) {
+        dq_val = dq_val * a_scale;
+      }
+      if constexpr (wei_quant_mode == PER_GROUP) {
+        b_scale = scale_b[n];
+        dq_val = dq_val * b_scale;
+      }
+      if constexpr (accum) {
+        output[m * ldo + n] += dq_val;
+      } else {
+        output[m * ldo + n] = dq_val;
+      }
+    }
+  }
+}
+
+// Store result to output buffer with dtype conversion
+// If act/wei are per_row or per_tensor quantized, apply scales
+// If bias is not null, add bias
+template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
+inline void store_out(
+    const float* y_buf,
+    out_dtype* c_ptr,
+    int64_t M,
+    int64_t lda,
+    const float* scales_a,
+    const float* scales_b,
+    const float* bias) {
+  float a_scale = 1.0, b_scale = 1.0;
+  __m512 va_scale, vb_scale;
+  if constexpr (act_quant_mode == PER_TENSOR) {
+    a_scale = *scales_a;
+  }
+  if constexpr (wei_quant_mode == PER_TENSOR) {
+    b_scale = *scales_b;
+    vb_scale = _mm512_set1_ps(b_scale);
+  }
+  for (int i = 0; i < M; ++i) {
+    if constexpr (act_quant_mode == PER_ROW) {
+      a_scale = *(scales_a + i);
+    }
+    if constexpr (act_quant_mode != PER_GROUP) {
+      va_scale = _mm512_set1_ps(a_scale);
+    }
+    constexpr int N_UNROLL = N / 16;
+    c10::ForcedUnroll<N_UNROLL>{}([&](auto idx) {
+      constexpr int j = idx * 16;
+      __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
+      __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps();
+      if constexpr (act_quant_mode != PER_GROUP) {
+        y_vec = _mm512_mul_ps(y_vec, va_scale);
+      }
+      if constexpr (wei_quant_mode == PER_ROW) {
+        vb_scale = _mm512_loadu_ps(scales_b + j);
+      }
+      if constexpr (wei_quant_mode != PER_GROUP) {
+        y_vec = _mm512_mul_ps(y_vec, vb_scale);
+      }
+      y_vec = _mm512_add_ps(y_vec, bias_vec);
+      if constexpr (std::is_same<out_dtype, float>::value) {
+        _mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
+      } else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) {
+        __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec);
+        _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec);
+      } else if constexpr (std::is_same<out_dtype, at::Half>::value) {
+        __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec);
+        _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec);
+      } else {
+        TORCH_CHECK(false, "Unsupported output dtype");
+      }
+    });
+    constexpr int tail_start = N / 16 * 16;
+    for (int j = tail_start; j < N; ++j) {
+      if constexpr (wei_quant_mode == PER_ROW) {
+        b_scale = scales_b[j];
+      }
+      c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
+    }
+  }  // for M
+}
+
+#else  // no AVX512
+
+static void cvt_f8e4m3_to_bf16(
+    const at::Float8_e4m3fn* __restrict__ in, at::BFloat16* out, int64_t rows, int64_t cols, int64_t stride) {
+  for (int r = 0; r < rows; ++r) {
+    for (int c = 0; c < cols; ++c) {
+      out[r * cols + c] = (at::BFloat16)in[r * stride + c];
+    }
+  }
+}
+
+// Store result to output buffer with dtype conversion
+// If act/wei are per_row or per_tensor quantized, apply scales
+// If bias is not null, add bias
+template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
+inline void store_out(
+    const float* y_buf,
+    out_dtype* c_ptr,
+    int64_t M,
+    int64_t lda,
+    const float* scales_a,
+    const float* scales_b,
+    const float* bias) {
+  float a_scale = 1.0, b_scale = 1.0;
+  if constexpr (act_quant_mode == PER_TENSOR) {
+    a_scale = *scales_a;
+  }
+  if constexpr (wei_quant_mode == PER_TENSOR) {
+    b_scale = *scales_b;
+  }
+  for (int i = 0; i < M; ++i) {
+    if constexpr (act_quant_mode == PER_ROW) {
+      a_scale = *(scales_a + i);
+    }
+    for (int j = 0; j < N; ++j) {
+      if constexpr (wei_quant_mode == PER_ROW) {
+        b_scale = scales_b[j];
+      }
+      c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
+    }
+  }  // for M
+}
+
+#endif  // CPU_CAPABILITY_AVX512
+
+template <bool cpublas_can_pack, int64_t N, int act_quant_mode, int wei_quant_mode>
+void _micro_gemm(
+    float* C,
+    const at::Float8_e4m3fn* A,
+    const float* scales_a,
+    const at::Float8_e4m3fn* B,
+    const float* scales_b,
+    int64_t M,
+    int64_t K,
+    int64_t lda,
+    int64_t ldc,
+    int64_t ldsa,
+    float* ukernel_buf,
+    at::BFloat16* dqA_buf,
+    at::BFloat16* dqB_buf) {
+  // If FP8 brgemm is not available, convert A/B to bf16 for computation
+  // Compute GEMM fp8 * fp8 -> fp32 (or bf16 * bf16 -> fp32)
+  // If per_group quant, apply scales. Otherwise, don't apply scales here
+  // Finally accumulate and store results
+#if defined(CPU_CAPABILITY_AVX512)
+  if constexpr (cpublas_can_pack) {
+#ifdef CPUBLAS_BRGEMM_F8F8F32
+    at::native::cpublas::brgemm(
+        M, N, K, lda /*lda*/, N /*ldb*/, N /*ldc*/, false /* add_C */, A, B, ukernel_buf, true /* is_vnni */);
+#else
+    cvt_f8e4m3_to_bf16(A, dqA_buf, M, K, lda);
+    cvt_f8e4m3_to_bf16(B, dqB_buf, K, N, N);
+    at::native::cpublas::brgemm(
+        M, N, K, K /*lda*/, N /*ldb*/, N /*ldc*/, false /* add_C */, dqA_buf, dqB_buf, ukernel_buf, true /* is_vnni */);
+#endif
+    _mm_prefetch(B + N * (K + 128), _MM_HINT_T0);
+    _mm_prefetch(A + K + 128, _MM_HINT_T0);
+    _accumulate_result<true, N, act_quant_mode, wei_quant_mode>(
+        C, ukernel_buf, scales_a, scales_b, M, N /*ldi*/, ldc, ldsa);
+  } else
+#endif
+  {
+    for (int64_t i = 0; i < M; ++i) {
+      for (int64_t j = 0; j < N; ++j) {
+        float sum = 0;
+        for (int64_t k = 0; k < K; ++k) {
+          sum += ((float)A[i * lda + k] * (float)B[k * N + j]);
+        }
+        if constexpr (act_quant_mode == PER_GROUP) {
+          sum *= scales_a[i * ldsa];
+        }
+        if constexpr (wei_quant_mode == PER_GROUP) {
+          sum *= scales_b[j];
+        }
+        C[i * ldc + j] += sum;
+      }
+    }
+  }
+}
+
+template <typename out_dtype, bool cpublas_can_pack, int act_quant_mode, int wei_quant_mode>
+void _float8_linear_impl(
+    const at::Tensor& input,
+    const at::Tensor& input_scales,
+    const at::Tensor& weight,
+    const at::Tensor& weight_scales,
+    const std::optional<at::Tensor>& bias,
+    at::Tensor& output) {
+  // input shape = [..., K]
+  // input is per token quantized
+  int64_t K = input.size(-1);
+  auto input_view = input.view({-1, K});
+  int64_t M = input_view.size(0);
+
+  // weight shape = [Nc, Kc, block_k, block_n]
+  // scales shape = [Nc, G, block_n]
+  int64_t Nc = weight.size(0);
+  int64_t Kc = wei_quant_mode != PER_GROUP ? 1 : weight.size(1);
+  int64_t block_k = wei_quant_mode != PER_GROUP ? weight.size(1) * weight.size(2) : weight.size(2);
+  constexpr int64_t block_n = BLOCK_N;
+  TORCH_CHECK(weight.size(3) == block_n, "Float8 linear: unexpected weight shape");
+  int64_t N = Nc * block_n;
+  TORCH_CHECK(K == Kc * block_k, "Float8 linear: weight and input shapes mismatch");
+  auto [parallel_on_M, block_m, Mc, Mc_parallel] = get_m_blocking(M);
+
+  // scales shape = [Nc, G, block_n]
+  int64_t num_groups = wei_quant_mode == PER_TENSOR ? 1 : weight_scales.size(1);
+  TORCH_CHECK(K % num_groups == 0, "K should be divisible by num_groups");
+  int64_t group_size = K / num_groups;
+  TORCH_CHECK(group_size % block_k == 0, "Float8 linear: group_size should be divisible by block_k");
+  int64_t block_per_group = group_size / block_k;
+  TORCH_CHECK(
+      input_scales.numel() == 1 || input_scales.numel() == M || input_scales.numel() == M * num_groups,
+      "Float8 linear: unexpected input scales shape");
+  auto ldsa = act_quant_mode == PER_TENSOR ? 0 : act_quant_mode == PER_ROW ? 1 : num_groups;
+
+  const at::Float8_e4m3fn* a_ptr = input_view.data_ptr<at::Float8_e4m3fn>();
+  const float* a_scales_ptr = input_scales.data_ptr<float>();
+  const at::Float8_e4m3fn* b_ptr = weight.data_ptr<at::Float8_e4m3fn>();
+  const float* b_scales_ptr = weight_scales.data_ptr<float>();
+  out_dtype* c_ptr = output.data_ptr<out_dtype>();
+  const float* bias_ptr = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
+
+  int64_t block_size = block_m * block_n;
+  int64_t num_thread = at::get_num_threads();
+  at::Tensor y_buffer = at::empty({num_thread, block_size}, output.options().dtype(at::kFloat));
+  // Create buffer for brgemm output and dqA/dqB (optional)
+#if defined(CPU_CAPABILITY_AVX512)
+  // buffer for brgemm output in float32
+  int64_t buffer_size = block_size * 2;  // float32 = bfloat16 * 2
+#ifndef CPUBLAS_BRGEMM_F8F8F32
+  // buffers for dqA & dqB in bf16
+  buffer_size += (block_k * block_n + block_m * block_k);
+#endif
+  at::Tensor micro_gemm_buffer = at::empty({num_thread, buffer_size}, output.options().dtype(at::kBFloat16));
+#endif
+
+  parallel_2d(Mc, Nc, [&](int64_t mc0, int64_t mc1, int64_t nc0, int64_t nc1) {
+    int tid = get_thread_num();
+    float* y_buf = y_buffer.data_ptr<float>() + tid * block_size;
+    at::BFloat16 *dqA_buffer = nullptr, *dqB_buffer = nullptr;
+    float* ukernel_buf = nullptr;
+#if defined(CPU_CAPABILITY_AVX512)
+    at::BFloat16* micro_gemm_buf = micro_gemm_buffer.data_ptr<at::BFloat16>() + tid * buffer_size;
+    ukernel_buf = reinterpret_cast<float*>(micro_gemm_buf);
+#ifndef CPUBLAS_BRGEMM_F8F8F32
+    dqA_buffer = micro_gemm_buf;
+    dqB_buffer = micro_gemm_buf + block_m * block_k;
+    ukernel_buf = reinterpret_cast<float*>(micro_gemm_buf + block_m * block_k + block_k * block_n);
+#endif
+#endif
+
+    loop_2d<at::Float8_e4m3fn>(mc0, mc1, nc0, nc1, block_n * K, [&](int64_t mci, int64_t nc, int64_t) {
+      int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m;
+      zero_buffer(y_buf, m_size * block_n);
+      for (int kci = 0; kci < Kc; ++kci) {
+        auto scales_a = a_scales_ptr + mci * block_m * num_groups + kci / block_per_group;
+        auto scales_b = b_scales_ptr + nc * block_n * num_groups + kci / block_per_group * block_n;
+        _micro_gemm<cpublas_can_pack, block_n, act_quant_mode, wei_quant_mode>(
+            /* C */ y_buf,
+            /* A */ a_ptr + mci * block_m * K + kci * block_k,
+            /* scales_a */ scales_a,
+            /* B */ b_ptr + (nc * Kc + kci) * block_n * block_k,
+            /* scales_b */ scales_b,
+            /* M */ m_size,
+            /* K */ block_k,
+            /* lda */ K,
+            /* ldc */ block_n,
+            /* ldsa */ ldsa,
+            /* ukernel_buf */ ukernel_buf,
+            /* dqA_buf */ dqA_buffer,
+            /* dqB_buf */ dqB_buffer);
+      }
+      auto scales_a = act_quant_mode == PER_TENSOR ? a_scales_ptr
+                      : act_quant_mode == PER_ROW  ? a_scales_ptr + mci * block_m
+                                                   : nullptr;
+      auto scales_b = wei_quant_mode == PER_TENSOR ? b_scales_ptr
+                      : wei_quant_mode == PER_ROW  ? b_scales_ptr + nc * block_n
+                                                   : nullptr;
+      auto bias_data = bias_ptr ? bias_ptr + nc * block_n : nullptr;
+      store_out<out_dtype, block_n, act_quant_mode, wei_quant_mode>(
+          y_buf, c_ptr + mci * block_m * N + nc * block_n, m_size, N /*lda*/, scales_a, scales_b, bias_data);
+    });
+
+    if constexpr (cpublas_can_pack) {
+      at::native::cpublas::brgemm_release();
+    }
+  });
+}
+
+}  // anonymous namespace
+
+void tinygemm_kernel(
+    float* C,
+    const at::Float8_e4m3fn* A,
+    const float* scales_a,
+    const at::Float8_e4m3fn* B,
+    const float* scales_b,
+    int64_t M,
+    int64_t K,
+    int64_t lda,
+    int64_t ldc,
+    int64_t ldsa,
+    float* ukernel_buf,
+    at::BFloat16* dqA_buf,
+    at::BFloat16* dqB_buf) {
+  // act_quant_mode = per row (2), wei_quant_mode = per group (3)
+  if (cpublas_could_pack()) {
+    _micro_gemm<true, BLOCK_N, 2, 3>(C, A, scales_a, B, scales_b, M, K, lda, ldc, ldsa, ukernel_buf, dqA_buf, dqB_buf);
+  } else {
+    _micro_gemm<false, BLOCK_N, 2, 3>(C, A, scales_a, B, scales_b, M, K, lda, ldc, ldsa, ukernel_buf, dqA_buf, dqB_buf);
+  }
+}
+
+#define INSTANTIATE_TINYGEMM_TEMPLATE() \
+  void tinygemm_kernel(                 \
+      float* C,                         \
+      const at::Float8_e4m3fn* A,       \
+      const float* scales_a,            \
+      const at::Float8_e4m3fn* B,       \
+      const float* scales_b,            \
+      int64_t M,                        \
+      int64_t K,                        \
+      int64_t lda,                      \
+      int64_t ldc,                      \
+      int64_t ldsa,                     \
+      float* ukernel_buf,               \
+      at::BFloat16* dqA_buf,            \
+      at::BFloat16* dqB_buf)
+
+INSTANTIATE_TINYGEMM_TEMPLATE();
+/*
+return: packed_weight, packed_scales
+*/
+std::tuple<at::Tensor, at::Tensor> float8_linear_prepack_impl(const at::Tensor& weight, const at::Tensor& scales) {
+  // weight shape = [N, K]
+  // scales shape = [N, G]
+  TORCH_CHECK(weight.dim() == 2, "Float8 linear CPU: Weight should be a 2D tensor for packing");
+  TORCH_CHECK(weight.size(1) % 2 == 0, "Float8 linear CPU: Weight should have even number of columns for packing");
+
+  auto new_scales = scales;
+  if (new_scales.dim() == 1) {
+    new_scales.unsqueeze_(1);
+  }
+  new_scales = new_scales.to(at::kFloat);
+  int N = weight.size(0);
+  int K = weight.size(1);
+  int G = scales.size(1);
+  int group_size = K / G;
+  int block_k = group_size > 128 ? 128 : group_size;
+  while (K % block_k != 0) {
+    block_k /= 2;
+  }
+  TORCH_CHECK(
+      block_k > 0 && block_k <= group_size, "Float8 linear CPU: Invalid block_k size, should be in (0, group_size]");
+  constexpr int block_n = BLOCK_N;
+  int Nc = N / block_n;
+  int Kc = K / block_k;
+
+  // Reorder weight to [N/block_n, K/block_k, block_k, block_n]
+  // Reorder scales to [N/block_n, G, block_n]
+  auto weight_view = weight.view({Nc, block_n, Kc, block_k});
+  at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous();
+  at::Tensor blocked_weight;
+  at::Tensor blocked_scales = new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
+
+#if defined(CPU_CAPABILITY_AVX512)
+  if (cpublas_could_pack()) {
+#ifdef CPUBLAS_BRGEMM_F8F8F32
+    constexpr int vnni_size = 4;  // for fp8
+#else
+    constexpr int vnni_size = 2;  // for float16
+#endif
+    blocked_weight = at::empty({Nc, Kc, block_k, block_n}, weight.options());
+    auto weight_ptr = reinterpret_cast<uint8_t*>(weight_reordered.data_ptr());
+    auto blocked_weight_ptr = reinterpret_cast<uint8_t*>(blocked_weight.data_ptr());
+    int64_t num_blocks = Nc * Kc;
+    at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) {
+      for (const auto i : c10::irange(begin, end)) {
+        auto in_ptr = weight_ptr + i * block_k * block_n;
+        auto out_ptr = blocked_weight_ptr + i * block_k * block_n;
+
+        // Reorder weight block to VNNI
+        // plain shape = [block_k, block_n]
+        // packed shape = [block_k / VNNI_SIZE, block_n, VNNI_SIZE] viewed as [block_k, block_n]
+        constexpr int n_group_size = 8;
+        constexpr int n_group = block_n / n_group_size;  // 4
+        for (int nb = 0; nb < n_group; ++nb) {
+          for (int k = 0; k < block_k; k += vnni_size) {
+            for (int ni = 0; ni < n_group_size; ++ni) {
+              for (int ki = 0; ki < vnni_size; ++ki) {
+                int src_idx = nb * n_group_size + ni + (k + ki) * block_n;
+                int dst_idx = (nb * n_group_size + ni) * vnni_size + k * block_n + ki;
+                *(out_ptr + dst_idx) = *(in_ptr + src_idx);
+              }
+            }
+          }
+        }
+      }
+    });
+  } else
+#endif
+  {
+    blocked_weight = weight_reordered;
+  }
+
+  return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales));
+}
+
+// AVX512 optimized channel-wise max computation
+inline void
+compute_channel_max_avx512(const float* data, float* max_vals, int64_t num_channels, int64_t elements_per_channel) {
+  at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
+    for (int64_t c = start; c < end; ++c) {
+      const float* channel_data = data + c * elements_per_channel;
+      float max_val = 0.0f;
+
+      int64_t i = 0;
+      // Process 16 elements at a time
+      __m512 max_vec = _mm512_setzero_ps();
+
+      for (; i <= elements_per_channel - 16; i += 16) {
+        __m512 data_vec = _mm512_loadu_ps(&channel_data[i]);
+        __m512 abs_vec = _mm512_abs_ps(data_vec);
+        max_vec = _mm512_max_ps(max_vec, abs_vec);
+      }
+
+      // Horizontal max of the vector
+      float hmax = _mm512_reduce_max_ps(max_vec);
+      max_val = std::max(max_val, hmax);
+
+      // Handle remaining elements
+      for (; i < elements_per_channel; ++i) {
+        max_val = std::max(max_val, std::abs(channel_data[i]));
+      }
+
+      max_vals[c] = max_val;
+    }
+  });
+}
+
+// AVX512 optimized scaling and clamping for channel-wise quantization
+inline void scale_clamp_channelwise_avx512(
+    const float* src,
+    float* dst,
+    const float* scales,
+    int64_t num_channels,
+    int64_t elements_per_channel,
+    float quant_max,
+    float neg_quant_max) {
+  at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
+    for (int64_t c = start; c < end; ++c) {
+      float scale_val = scales[c];
+      float scale_reciprocal = 1.0f / scale_val;
+
+      const float* channel_src = src + c * elements_per_channel;
+      float* channel_dst = dst + c * elements_per_channel;
+
+      int64_t i = 0;
+      const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
+      const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
+      const __m512 neg_quant_max_vec = _mm512_set1_ps(neg_quant_max);
+
+      for (; i <= elements_per_channel - 16; i += 16) {
+        __m512 src_vec = _mm512_loadu_ps(&channel_src[i]);
+        __m512 scaled_vec = _mm512_mul_ps(src_vec, scale_recip_vec);
+        __m512 clamped_vec = _mm512_min_ps(_mm512_max_ps(scaled_vec, neg_quant_max_vec), quant_max_vec);
+        _mm512_storeu_ps(&channel_dst[i], clamped_vec);
+      }
+
+      // Handle remaining elements
+      for (; i < elements_per_channel; ++i) {
+        float scaled = channel_src[i] * scale_reciprocal;
+        channel_dst[i] = std::clamp(scaled, neg_quant_max, quant_max);
+      }
+    }
+  });
+}
+
+// AVX512 optimized scaling and clamping for global quantization
+inline void scale_clamp_global_avx512(
+    const float* src, float* dst, int64_t size, float scale_reciprocal, float quant_max, float neg_quant_max) {
+  if (size <= 0) return;
+
+  const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
+  const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
+  const __m512 neg_quant_max_vec = _mm512_set1_ps(neg_quant_max);
+
+  at::parallel_for(0, size, 4096, [&](int64_t start, int64_t end) {
+    int64_t block_size = end - start;
+    const float* block_src = src + start;
+    float* block_dst = dst + start;
+
+    int64_t i = 0;
+    for (; i <= block_size - 16; i += 16) {
+      __m512 src_vec = _mm512_loadu_ps(&block_src[i]);
+      __m512 scaled_vec = _mm512_mul_ps(src_vec, scale_recip_vec);
+      __m512 clamped_vec = _mm512_min_ps(_mm512_max_ps(scaled_vec, neg_quant_max_vec), quant_max_vec);
+      _mm512_storeu_ps(&block_dst[i], clamped_vec);
+    }
+
+    // Handle remaining elements in block
+    for (; i < block_size; ++i) {
+      float scaled = block_src[i] * scale_reciprocal;
+      block_dst[i] = std::clamp(scaled, neg_quant_max, quant_max);
+    }
+  });
+}
+
+std::tuple<at::Tensor, at::Tensor>
+_quantize_fp8e4m3(const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt = c10::nullopt) {
+  constexpr float quant_max = 448.0f;  // torch.finfo(torch.float8_e4m3fn).max
+  constexpr float eps = std::numeric_limits<float>::epsilon();
+
+  // Ensure input is contiguous and in float32
+  auto t_float = t.to(at::ScalarType::Float).contiguous();
+  at::Tensor qt;
+  at::Tensor scale_tensor;
+
+  if (channelwise) {
+    // Channel-wise quantization with AVX512 optimization
+    int64_t num_channels = t_float.size(0);
+    int64_t elements_per_channel = t_float.numel() / num_channels;
+
+    if (elements_per_channel * num_channels != t_float.numel()) {
+      throw std::runtime_error("Tensor must be divisible by number of channels for channel-wise quantization");
+    }
+
+    // Allocate scale tensor
+    scale_tensor = at::empty({num_channels}, t_float.options());
+    float* scale_data = scale_tensor.data_ptr<float>();
+
+    // Compute channel-wise max using AVX512
+    compute_channel_max_avx512(t_float.data_ptr<float>(), scale_data, num_channels, elements_per_channel);
+
+    // Apply quant_max and EPS
+    at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
+      for (int64_t c = start; c < end; ++c) {
+        scale_data[c] = std::max(scale_data[c] / quant_max, eps);
+      }
+    });
+
+    // Create output tensor for quantized values
+    qt = at::empty_like(t_float);
+
+    // Scale and clamp using AVX512
+    scale_clamp_channelwise_avx512(
+        t_float.data_ptr<float>(),
+        qt.data_ptr<float>(),
+        scale_data,
+        num_channels,
+        elements_per_channel,
+        quant_max,
+        -quant_max);
+  } else {
+    // Global quantization with AVX512 optimization
+    if (!scale_opt.has_value()) {
+      throw std::runtime_error("Scale must be provided for non-channelwise quantization in AVX512 version");
+    }
+
+    scale_tensor = scale_opt.value().to(at::ScalarType::Float).contiguous();
+
+    // Handle scalar scale case
+    float scale_val;
+    if (scale_tensor.numel() == 1) {
+      scale_val = std::max(scale_tensor.item<float>(), eps);
+    } else {
+      // Compute max of scale tensor if it's not scalar
+      float* scale_data = scale_tensor.data_ptr<float>();
+      __m512 max_vec = _mm512_set1_ps(eps);
+
+      int64_t i = 0;
+      for (; i <= scale_tensor.numel() - 16; i += 16) {
+        __m512 scale_vec = _mm512_loadu_ps(&scale_data[i]);
+        max_vec = _mm512_max_ps(max_vec, scale_vec);
+      }
+
+      float hmax = _mm512_reduce_max_ps(max_vec);
+      for (; i < scale_tensor.numel(); ++i) {
+        hmax = std::max(hmax, scale_data[i]);
+      }
+      scale_val = std::max(hmax, eps);
+    }
+
+    scale_tensor = at::tensor({scale_val}, scale_tensor.options());
+    float scale_reciprocal = 1.0f / scale_val;
+
+    // Create output tensor
+    qt = at::empty_like(t_float);
+
+    // Scale and clamp using AVX512
+    scale_clamp_global_avx512(
+        t_float.data_ptr<float>(), qt.data_ptr<float>(), t_float.numel(), scale_reciprocal, quant_max, -quant_max);
+  }
+
+  // Final conversion to FP8 E4M3FN using ATen's native conversion
+  qt = at::_to_copy(qt, at::ScalarType::Float8_e4m3fn);
+
+  return std::make_tuple(qt, scale_tensor);
+}
+
+inline __m128i cvtfp32_fp8e4m3(__m512& src) {
+  // cvt 16x32 from fp32 to fp8 e4m3
+  const __m512i sign_mask = _mm512_set1_epi32(0x80000000);
+  const __m512i fp8_max = _mm512_set1_epi32(UINT32_C(1087) << 20);
+  const __m512i denorm_thresh = _mm512_set1_epi32(UINT32_C(121) << 23);
+  const __m512i denorm_mask = _mm512_set1_epi32(UINT32_C(141) << 23);
+  const __m512i bias_part1 = _mm512_set1_epi32((uint32_t)(7 - 127) << 23);
+  const __m512i rounding_bias = _mm512_set1_epi32(0x7FFFF);
+  __m512i f_bits = _mm512_castps_si512(src);
+  // Extract and save sign
+  __m512i sign = _mm512_and_epi32(f_bits, sign_mask);
+  f_bits = _mm512_xor_epi32(f_bits, sign);
+
+  // Prepare result containers
+  __m512i result = _mm512_setzero_si512();
+
+  // Step 1: Handle case of overflow
+  // (f_bits >= fp8_max): set result = 0x7f
+  __mmask16 overflow_mask = _mm512_cmpge_epu32_mask(f_bits, fp8_max);
+  if (overflow_mask) {
+    result = _mm512_mask_set1_epi32(result, overflow_mask, 0x7f);
+  }
+
+  // Step 2: Handle small numbers (denormals)
+  // Small numbers (f_bits < denorm_thresh)
+  __mmask16 denorm_thresh_mask = _mm512_cmplt_epu32_mask(f_bits, denorm_thresh);
+
+  if (denorm_thresh_mask) {
+    __m512 small_input = _mm512_castsi512_ps(f_bits);
+    __m512 small_denorm = _mm512_add_ps(small_input, _mm512_castsi512_ps(denorm_mask));
+    __m512i small_denorm_bits = _mm512_castps_si512(small_denorm);
+    __m512i small_result = _mm512_sub_epi32(small_denorm_bits, denorm_mask);
+    result = _mm512_mask_mov_epi32(result, denorm_thresh_mask, small_result);
+  }
+
+  // Step 3: Handle normal numbers
+  __mmask16 normal_mask = ~(overflow_mask | denorm_thresh_mask);
+
+  if (normal_mask) {
+    // mant_odd = (f_bits >> 20) & 1
+    __m512i mant_odd = _mm512_and_epi32(_mm512_srli_epi32(f_bits, 20), _mm512_set1_epi32(1));
+    // f_bits += bias_part1 + rounding_bias
+    __m512i rounded = _mm512_add_epi32(f_bits, bias_part1);
+    rounded = _mm512_add_epi32(rounded, rounding_bias);
+    // Add mant_odd
+    rounded = _mm512_add_epi32(rounded, mant_odd);
+    // Shift right by 20 bits
+    __m512i normal_result = _mm512_srli_epi32(rounded, 20);
+    result = _mm512_mask_mov_epi32(result, normal_mask, normal_result);
+  }
+
+  // Merge back the sign
+  __m512i sign_shifted = _mm512_srli_epi32(sign, 24);
+  result = _mm512_or_epi32(result, sign_shifted);
+
+  // Now result is 16 x 32-bit integers, but we only need 8-bit for each
+  __m512i packed = _mm512_and_si512(result, _mm512_set1_epi32(0xFF));
+
+  // Narrow 32-bit integers to 8-bit
+  return _mm512_cvtepi32_epi8(packed);
+}
+
+#ifdef CPU_CAPABILITY_AVX10_2
+__attribute__((target("avx10.2"))) inline __m128i cvtfp32_fp8e4m3_avx10_2(__m512& src) {
+  __m256i f16_vec = _mm512_cvt_roundps_ph(src, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+  return _mm256_cvtph_hf8(_mm256_castsi256_ph(f16_vec));
+}
+
+// Carries the whole loop rather than just the conversion, as GCC cannot inline an avx10.2
+// function into a caller lacking that target, so branching per vector is costly.
+__attribute__((target("avx10.2"))) void quantize_bf16_block_avx10_2(
+    const at::BFloat16* __restrict__ src, at::Float8_e4m3fn* __restrict__ dst, int64_t count, float scale_reciprocal) {
+  constexpr float quant_max = 448.0f;  // torch.finfo(torch.float8_e4m3fn).max
+  const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
+  const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
+  const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max);
+
+  for (int64_t i = 0; i <= count - 32; i += 32) {
+    __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src[i]));
+    __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src[i + 16]));
+
+    __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16));
+    __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16));
+
+    __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec);
+    __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec);
+
+    __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec);
+    __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec);
+
+    __m128i fp8_vec1 = cvtfp32_fp8e4m3_avx10_2(clamped_vec1);
+    __m128i fp8_vec2 = cvtfp32_fp8e4m3_avx10_2(clamped_vec2);
+    _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), fp8_vec1);
+    _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), fp8_vec2);
+  }
+}
+#endif  // CPU_CAPABILITY_AVX10_2
+
+std::tuple<at::Tensor, at::Tensor> _quantize_fp8e4m3_bf16_per_tensor_no_scale(const at::Tensor& t) {
+  constexpr float quant_max = 448.0f;  // torch.finfo(torch.float8_e4m3fn).max
+  constexpr float eps = std::numeric_limits<float>::epsilon();
+
+  // Input validation and preparation
+  assert(t.scalar_type() == at::ScalarType::BFloat16);
+  auto t_bf16 = t.contiguous();
+  int64_t num_channels = t_bf16.size(0);
+  int64_t elements_per_channel = t_bf16.numel() / num_channels;
+  assert(elements_per_channel % 32 == 0);  // do not consider tile currently
+  at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn));
+
+  // Allocate output tensors
+  at::Tensor scale_tensor = at::empty({num_channels}, t_bf16.options().dtype(at::ScalarType::Float));
+  float* scale_data = scale_tensor.data_ptr<float>();
+
+  // Unified processing: compute max, apply scale, and quantize in single pass
+  at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
+    for (int64_t c = start; c < end; ++c) {
+      const at::BFloat16* channel_src = t_bf16.data_ptr<at::BFloat16>() + c * elements_per_channel;
+      at::Float8_e4m3fn* quant_dst = quant_t.data_ptr<at::Float8_e4m3fn>() + c * elements_per_channel;
+
+      // Step 1: Compute channel-wise max using AVX512
+      float channel_max = 0.0f;
+      int64_t i = 0;
+
+      // Process 32 elements at a time for max computation
+      for (; i <= elements_per_channel - 32; i += 32) {
+        // Load 32 BF16 values (2x 256-bit vectors)
+        __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i]));
+        __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i + 16]));
+
+        // Convert BF16 to FP32
+        __m512i fp32_int1 = _mm512_cvtepu16_epi32(src_vec1);
+        __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int1, 16));
+
+        __m512i fp32_int2 = _mm512_cvtepu16_epi32(src_vec2);
+        __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(fp32_int2, 16));
+
+        // Compute absolute values
+        __m512 abs_vec1 = _mm512_abs_ps(fp32_vec1);
+        __m512 abs_vec2 = _mm512_abs_ps(fp32_vec2);
+
+        // Find max in each vector
+        float max1 = _mm512_reduce_max_ps(abs_vec1);
+        float max2 = _mm512_reduce_max_ps(abs_vec2);
+        channel_max = std::max(channel_max, std::max(max1, max2));
+      }
+
+      // Step 2: Apply quant_max and EPS to compute scale
+      float scale_val = std::max(channel_max / quant_max, eps);
+      float scale_reciprocal = 1.0f / scale_val;
+      scale_data[c] = scale_val;
+
+      // Step 3: Scale and clamp using AVX512 (reuse the same loop structure)
+#ifdef CPU_CAPABILITY_AVX10_2
+      if (avx10_2_available()) {
+        quantize_bf16_block_avx10_2(channel_src, quant_dst, elements_per_channel, scale_reciprocal);
+        continue;
+      }
+#endif
+      i = 0;
+      const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
+      const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
+      const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max);
+
+      for (; i <= elements_per_channel - 32; i += 32) {
+        // Load 32 BF16 values
+        __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i]));
+        __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&channel_src[i + 16]));
+
+        // Convert BF16 to FP32
+        __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16));
+        __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16));
+
+        // Scale and clamp
+        __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec);
+        __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec);
+
+        __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec);
+        __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec);
+
+        __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1);
+        __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2);
+        _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i), fp8_vec1);
+        _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_dst + i + 16), fp8_vec2);
+      }
+    }
+  });
+
+  return std::make_tuple(quant_t, scale_tensor);
+}
+
+std::tuple<at::Tensor, at::Tensor>
+_quantize_fp8e4m3_bf16_per_tensor_with_scale(const at::Tensor& t, at::Tensor& scale_tensor) {
+  constexpr float quant_max = 448.0f;  // torch.finfo(torch.float8_e4m3fn).max
+  constexpr float eps = std::numeric_limits<float>::epsilon();
+
+  // Input validation
+  assert(t.scalar_type() == at::ScalarType::BFloat16);
+
+  auto t_bf16 = t.contiguous();
+  auto scale_tensor_contig = scale_tensor.contiguous().to(at::ScalarType::Float);
+
+  // Get scale value (handle scalar or tensor)
+  float scale_val;
+  if (scale_tensor_contig.numel() == 1) {
+    scale_val = scale_tensor_contig.item<float>();
+  } else {
+    // Take max if scale is a tensor (though should be scalar for global quantization)
+    scale_val = scale_tensor_contig.max().item<float>();
+  }
+
+  // Apply EPS to scale
+  scale_val = std::max(scale_val, eps);
+  float scale_reciprocal = 1.0f / scale_val;
+
+  at::Tensor scale_output = at::tensor({scale_val}, scale_tensor_contig.options());
+
+  // Apply scale and clamp using AVX512
+  const at::BFloat16* src_data = t_bf16.data_ptr<at::BFloat16>();
+  at::Tensor quant_t = at::empty(t_bf16.sizes(), t_bf16.options().dtype(at::kFloat8_e4m3fn));
+  int64_t total_elements = t_bf16.numel();
+  const __m512 scale_recip_vec = _mm512_set1_ps(scale_reciprocal);
+  const __m512 quant_max_vec = _mm512_set1_ps(quant_max);
+  const __m512 neg_quant_max_vec = _mm512_set1_ps(-quant_max);
+  int64_t num_channels = t_bf16.size(0);
+  int64_t elements_per_channel = total_elements / num_channels;
+  assert(elements_per_channel % 32 == 0);  // do not consider tile currently
+  // Process in parallel blocks
+  at::parallel_for(0, num_channels, 1, [&](int64_t start, int64_t end) {
+    for (int64_t c = start; c < end; ++c) {
+      const at::BFloat16* src_data = t_bf16.data_ptr<at::BFloat16>() + c * elements_per_channel;
+      at::Float8_e4m3fn* quant_t_data = quant_t.data_ptr<at::Float8_e4m3fn>() + c * elements_per_channel;
+
+#ifdef CPU_CAPABILITY_AVX10_2
+      if (avx10_2_available()) {
+        quantize_bf16_block_avx10_2(src_data, quant_t_data, elements_per_channel, scale_reciprocal);
+        continue;
+      }
+#endif
+
+      int64_t i = 0;
+      // Process 32 elements at a time using AVX512
+      for (; i <= elements_per_channel - 32; i += 32) {
+        // Load 32 BF16 values (2x 256-bit vectors)
+        __m256i src_vec1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src_data[i]));
+        __m256i src_vec2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&src_data[i + 16]));
+
+        // Convert BF16 to FP32
+        __m512 fp32_vec1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec1), 16));
+        __m512 fp32_vec2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(src_vec2), 16));
+
+        // Scale and clamp
+        __m512 scaled_vec1 = _mm512_mul_ps(fp32_vec1, scale_recip_vec);
+        __m512 clamped_vec1 = _mm512_min_ps(_mm512_max_ps(scaled_vec1, neg_quant_max_vec), quant_max_vec);
+
+        __m512 scaled_vec2 = _mm512_mul_ps(fp32_vec2, scale_recip_vec);
+        __m512 clamped_vec2 = _mm512_min_ps(_mm512_max_ps(scaled_vec2, neg_quant_max_vec), quant_max_vec);
+
+        __m128i fp8_vec1 = cvtfp32_fp8e4m3(clamped_vec1);
+        __m128i fp8_vec2 = cvtfp32_fp8e4m3(clamped_vec2);
+        _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i), fp8_vec1);
+        _mm_storeu_si128(reinterpret_cast<__m128i*>(quant_t_data + i + 16), fp8_vec2);
+      }
+    }
+  });
+
+  return std::make_tuple(quant_t, scale_output);
+}
+
+std::tuple<at::Tensor, at::Tensor>
+_quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt) {
+  if (channelwise) {
+    return _quantize_fp8e4m3_bf16_per_tensor_no_scale(t);
+  } else {
+    assert(scale_opt.has_value());
+    return _quantize_fp8e4m3_bf16_per_tensor_with_scale(t, scale_opt.value());
+  }
+}
+
+// Public wrapper for torch op registration
+std::tuple<at::Tensor, at::Tensor>
+quantize_fp8e4m3_vec(const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt) {
+  return _quantize_fp8e4m3_vec(t, channelwise, scale_opt);
+}
+
+at::Tensor fp8_scaled_mm_with_quant(
+    const at::Tensor& act,
+    const std::optional<at::Tensor>& act_scales,
+    bool channelwise,
+    const at::Tensor& weight,
+    const at::Tensor& weight_scales,
+    const std::optional<at::Tensor>& bias,
+    at::ScalarType output_dtype) {
+  // TODO: refine here to use tensor ptr and in below dispatch api.
+  std::tuple<at::Tensor, at::Tensor> quant_act = act.scalar_type() == at::ScalarType::BFloat16
+                                                     ? _quantize_fp8e4m3_vec(act, channelwise, act_scales)
+                                                     : _quantize_fp8e4m3(act, channelwise, act_scales);
+  auto input = std::get<0>(quant_act);
+  auto input_scales = std::get<1>(quant_act);
+  int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0);
+  int act_quant_mode = input_scales.numel() == 1                                ? PER_TENSOR
+                       : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW
+                                                                                : PER_GROUP;
+  int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP;
+  // Case to fall back
+  if (weight.dim() == 2) {
+    TORCH_CHECK(
+        act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP,
+        "FP8 linear: Per-group quantization is not supported in the fallback path");
+    auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias);
+    return y_fp32.to(output_dtype);
+  }
+
+  static bool cpublas_can_pack = cpublas_could_pack();
+  auto out_sizes = input.sizes().vec();
+  out_sizes.back() = N;
+  auto output = at::empty(out_sizes, input.options().dtype(output_dtype));
+
+#define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \
+  AT_DISPATCH_BOOL_NO_RETURN(                                                               \
+      CAN_PACK,                                                                             \
+      "cpublas_can_pack",                                                                   \
+      can_pack,                                                                             \
+      AT_DISPATCH_QUANT_MODE_NO_RETURN(                                                     \
+          A_QUANT_MODE,                                                                     \
+          "act_quant_mode",                                                                 \
+          a_quant_mode,                                                                     \
+          AT_DISPATCH_QUANT_MODE_NO_RETURN(                                                 \
+              B_QUANT_MODE,                                                                 \
+              "wei_quant_mode",                                                             \
+              b_quant_mode,                                                                 \
+              AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__))))
+
+  AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() {
+    _float8_linear_impl<out_t, can_pack, a_quant_mode, b_quant_mode>(
+        input, input_scales, weight, weight_scales, bias, output);
+  });
+  return output;
+}
+
+at::Tensor float8_linear_impl(
+    const at::Tensor& input,
+    const at::Tensor& input_scales,
+    const at::Tensor& weight,
+    const at::Tensor& weight_scales,
+    const std::optional<at::Tensor>& bias,
+    at::ScalarType output_dtype) {
+  int64_t N = weight.dim() == 4 ? weight.size(0) * weight.size(-1) : weight.size(0);
+  int act_quant_mode = input_scales.numel() == 1                                ? PER_TENSOR
+                       : input_scales.numel() == input.numel() / input.size(-1) ? PER_ROW
+                                                                                : PER_GROUP;
+  int wei_quant_mode = weight_scales.numel() == 1 ? PER_TENSOR : weight_scales.numel() == N ? PER_ROW : PER_GROUP;
+  // Case to fall back
+  if (weight.dim() == 2) {
+    TORCH_CHECK(
+        act_quant_mode != PER_GROUP && wei_quant_mode != PER_GROUP,
+        "FP8 linear: Per-group quantization is not supported in the fallback path");
+    auto y_fp32 = at::linear(input.to(at::kFloat).mul_(input_scales), weight.to(at::kFloat).mul_(weight_scales), bias);
+    return y_fp32.to(output_dtype);
+  }
+
+  static bool cpublas_can_pack = cpublas_could_pack();
+  auto out_sizes = input.sizes().vec();
+  out_sizes.back() = N;
+  auto output = at::empty(out_sizes, input.options().dtype(output_dtype));
+
+#define AT_DISPATCH_FP8_LINEAR_KERNEL(OUT_DTYPE, CAN_PACK, A_QUANT_MODE, B_QUANT_MODE, ...) \
+  AT_DISPATCH_BOOL_NO_RETURN(                                                               \
+      CAN_PACK,                                                                             \
+      "cpublas_can_pack",                                                                   \
+      can_pack,                                                                             \
+      AT_DISPATCH_QUANT_MODE_NO_RETURN(                                                     \
+          A_QUANT_MODE,                                                                     \
+          "act_quant_mode",                                                                 \
+          a_quant_mode,                                                                     \
+          AT_DISPATCH_QUANT_MODE_NO_RETURN(                                                 \
+              B_QUANT_MODE,                                                                 \
+              "wei_quant_mode",                                                             \
+              b_quant_mode,                                                                 \
+              AT_DISPATCH_OUT_TYPES(OUT_DTYPE, "out_dtype", __VA_ARGS__))))
+
+  AT_DISPATCH_FP8_LINEAR_KERNEL(output_dtype, cpublas_can_pack, act_quant_mode, wei_quant_mode, [&]() {
+    _float8_linear_impl<out_t, can_pack, a_quant_mode, b_quant_mode>(
+        input, input_scales, weight, weight_scales, bias, output);
+  });
+  return output;
+}
\ No newline at end of file
diff --git a/csrc/cpu/sgl-kernels/moe.cpp b/csrc/cpu/sgl-kernels/moe.cpp
index 3afdc1bd6791..09a170eb6d4c 100644
--- a/csrc/cpu/sgl-kernels/moe.cpp
+++ b/csrc/cpu/sgl-kernels/moe.cpp
@@ -836,6 +836,11 @@ static inline void check_moe_scales(
     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.");
+  } else if constexpr (quant == CPUQuantMethod::FP8_W8A8) {
+    TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a8.");
+    TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for fp8 w8a8.");
+    TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a8.");
+    TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2.");
   }
 }
 
@@ -850,9 +855,14 @@ static inline void check_moe_scales(
     check_moe_scales<CPUQuantMethod::FP8_W8A16>(w1_scale, w2_scale, block_size);
   } else if (moe_comp_method == CPUQuantMethod::MXFP4) {
     check_moe_scales<CPUQuantMethod::MXFP4>(w1_scale, w2_scale, block_size);
+  } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
+    check_moe_scales<CPUQuantMethod::FP8_W8A8>(w1_scale, w2_scale, block_size);
   }
 }
 
+// CHECK_MOE_SCALES_FP8: validate scale shapes for FP8_W8A16.
+// w1s shape [E, 2N/block_N, K/block_K, block_N]: DIM0=1, DIM1=2
+// w2s shape [E, K/block_N, N/block_K]: DIM0=1, DIM1=2
 #define CHECK_MOE_SCALES_FP8(DIM0, DIM1)                      \
   auto w1s = w1_scale.value();                                \
   auto w2s = w2_scale.value();                                \
@@ -864,6 +874,17 @@ static inline void check_moe_scales(
   TORCH_CHECK(w2s.size(DIM0) == div_up(K, block_size_N));     \
   TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K))
 
+// CHECK_MOE_SCALES_FP8_W8A8: extract scale info for FP8_W8A8 without shape
+// validation, because w1s is pre-packed by float8_linear_prepack_cpu with
+// shape [E, 2N/BLOCK_N, G, BLOCK_N] where BLOCK_N=32 != block_size_N(128).
+// This matches the sglang upstream behavior (checks are commented out there).
+#define CHECK_MOE_SCALES_FP8_W8A8(DIM0, DIM1)                \
+  auto w1s = w1_scale.value();                                \
+  auto w2s = w2_scale.value();                                \
+  auto block_size_val = block_size.value();                   \
+  int64_t block_size_N = block_size_val[0];                   \
+  int64_t block_size_K = block_size_val[1]
+
 // hidden_states: [M, K]
 // w1: [E, 2N, K] or [E, 2N, K / 2] for uint8
 // w2: [E, K, N] or [E, K, N / 2] for uint8
@@ -883,6 +904,7 @@ void 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<at::Tensor>& a1_scale,
     const std::optional<std::vector<int64_t>> block_size,
     const std::optional<at::Tensor>& w1_bias,
     const std::optional<at::Tensor>& w2_bias,
@@ -895,7 +917,10 @@ void fused_experts_cpu(
   constexpr int64_t BLOCK_M = block_size_m();
   constexpr int64_t BLOCK_N = block_size_n();
 
-  const auto st = hidden_states.scalar_type();
+  // FP8_W8A8 hidden_states input is FP8; scalar_t should be BF16 for output
+  const auto st = (moe_comp_method == CPUQuantMethod::FP8_W8A8)
+                      ? at::kBFloat16
+                      : hidden_states.scalar_type();
   // TODO: fused_topk_torch_native (CPU fallback for models like MiniMax)
   // returns int64 topk_ids; fused_experts_cpu requires int32. Remove the typecast after topk kernel is provided
   auto topk_ids_ = topk_ids.scalar_type() == at::kInt ? topk_ids : topk_ids.to(at::kInt);
@@ -908,6 +933,9 @@ void fused_experts_cpu(
   if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) {
     CHECK_DIM(4, w1);
     CHECK_DIM(4, w2);
+  } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
+    // w1 is prepackaged by float8_linear_prepack_cpu (5D), skip w1 dim check
+    CHECK_DIM(3, w2);
   } else {
     CHECK_DIM(3, w1);
     CHECK_DIM(3, w2);
@@ -924,8 +952,9 @@ void fused_experts_cpu(
 
   int64_t M = hidden_states.size(0);
   int64_t K = hidden_states.size(1);
-  int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2
-                                                           : w1.size(1) / 2;
+  int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8
+                  ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2
+                  : (moe_comp_method == CPUQuantMethod::FP8_W8A8 ? w2.size(2) : w1.size(1) / 2);
   int64_t E = w1.size(0);
   int64_t topk = topk_weights_.size(1);
 
@@ -935,7 +964,10 @@ void fused_experts_cpu(
 
   // check weight shapes
   CHECK_EQ(w2.size(0), E);
-  if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) {
+  if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
+    CHECK_EQ(w2.size(1), K);
+    CHECK_EQ(packed_w2.size(2), packed_N);
+  } else if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) {
     CHECK_EQ(w2.size(1), K);
     CHECK_EQ(packed_w1.size(2), packed_K / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
     CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
@@ -993,22 +1025,35 @@ void fused_experts_cpu(
   //   5. Aq_tmp : [M, K] or [M * topk, N]
   //   6. As_tmp : [M * topk]
   //
-  // for fp8 w8a16 and mxfp4:
+  // for fp8 w8a16, mxfp4, and fp8 w8a8:
   //   7. intermediate_cache0 : [M * topk, 2N]
   //   8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
   //
+  // for fp8 w8a8 additionally:
+  //   9. Ukernel_tmp : [T, 2 * BLOCK_M * BLOCK_N]  (float32 for FP8×FP8 brgemm)
+  //
   int64_t buffer_size_nbytes =
       M * topk * N * 2 + M * topk * K * 2 +
       num_threads * BLOCK_M * K *
-          (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 1 : 2) +
+          (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8
+               | moe_comp_method == CPUQuantMethod::FP8_W8A8 ? 1 : 2) +
       num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
 
   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 || moe_comp_method == CPUQuantMethod::MXFP4) {
+  if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4
+      || moe_comp_method == CPUQuantMethod::FP8_W8A8) {
     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::FP8_W8A8) {
+    // Ukernel buffer: float32 workspace per thread for FP8×FP8 brgemm
+    buffer_size_nbytes += num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
+#ifndef CPUBLAS_BRGEMM_F8F8F32
+    // BF16 fallback: also need dqA [T, BLOCK_M, BLOCK_K] + dqB [T, BLOCK_K, BLOCK_N]
+    buffer_size_nbytes += num_threads * (BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N) * 2;
+#endif
+  }
   if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
     buffer_size_nbytes += M * topk * 2 * N * 2 + std::max(M * K, M * topk * N) + M * topk * sizeof(float) +
                           num_threads * 2 * get_4bit_block_k_size(K / w1_scale.value().size(2)) * BLOCK_N;
@@ -1181,6 +1226,47 @@ void fused_experts_cpu(
           E,
           topk,
           num_tokens_post_pad);
+    } else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
+      // Stage 1: FP8×FP8 GEMM (hidden_states is FP8 input)
+      // Stage 2: BF16×FP8 GEMM (W8A16 path for down-projection)
+      at::Float8_e4m3fn* __restrict__ A_tmp =
+          (at::Float8_e4m3fn*)((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));
+      float* __restrict__ Ukernel_tmp =
+          (float*)((void*)(B_tmp + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N)));
+
+      CHECK_MOE_SCALES_FP8_W8A8(1, 2);
+      fused_experts_fp8_a8_kernel_impl(
+          out_hidden_states.data_ptr<scalar_t>(),
+          intermediate_cache0,
+          intermediate_cache1,
+          intermediate_cache2,
+          A_tmp,
+          B_tmp,
+          C_tmp,
+          Ukernel_tmp,
+          hidden_states.data_ptr<at::Float8_e4m3fn>(),
+          packed_w1.data_ptr<at::Float8_e4m3fn>(),
+          packed_w2.data_ptr<at::Float8_e4m3fn>(),
+          a1_scale.has_value() ? a1_scale.value().data_ptr<float>() : nullptr,
+          w1s.data_ptr<float>(),
+          w2s.data_ptr<float>(),
+          block_size_N,
+          block_size_K,
+          topk_weights_.data_ptr<float>(),
+          sorted_ids,
+          expert_ids,
+          offsets,
+          M,
+          N,
+          K,
+          E,
+          topk,
+          num_tokens_post_pad);
     } else {
       scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K;
       float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
diff --git a/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp b/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp
new file mode 100644
index 000000000000..0f393a34a853
--- /dev/null
+++ b/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp
@@ -0,0 +1,432 @@
+// Adapted from
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
+
+// clang-format off
+
+#include <c10/util/Unroll.h>
+
+#include "common.h"
+#include "gemm.h"
+#include "moe.h"
+
+namespace {
+
+// Byte-wise copy for FP8 elements (no remainder assumption)
+inline void copy_stub2(
+    at::Float8_e4m3fn* __restrict__ out,
+    const at::Float8_e4m3fn* __restrict__ input,
+    int64_t size) {
+  for (int64_t d = 0; d < size; d++) {
+    out[d] = input[d];
+  }
+}
+
+#define PER_TENSOR 1
+#define PER_ROW 2
+#define PER_GROUP 3
+
+// Store float32 GEMM result to output buffer with dequantization and dtype
+// conversion. Activations are PER_ROW quantized; weights are PER_GROUP quantized.
+// Uses AVX512 on capable platforms.
+#if defined(CPU_CAPABILITY_AVX512)
+
+template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
+inline void store_out(
+    const float* y_buf,
+    out_dtype* c_ptr,
+    int64_t M,
+    int64_t lda,
+    const float* scales_a,
+    const float* scales_b,
+    const float* bias) {
+  float a_scale = 1.0, b_scale = 1.0;
+  __m512 va_scale, vb_scale;
+  if constexpr (act_quant_mode == PER_TENSOR) {
+    a_scale = *scales_a;
+  }
+  if constexpr (wei_quant_mode == PER_TENSOR) {
+    b_scale = *scales_b;
+    vb_scale = _mm512_set1_ps(b_scale);
+  }
+  for (int i = 0; i < M; ++i) {
+    if constexpr (act_quant_mode == PER_ROW) {
+      a_scale = *(scales_a + i);
+    }
+    if constexpr (act_quant_mode != PER_GROUP) {
+      va_scale = _mm512_set1_ps(a_scale);
+    }
+    constexpr int N_UNROLL = N / 16;
+    c10::ForcedUnroll<N_UNROLL>{}([&](auto idx) {
+      constexpr int j = idx * 16;
+      __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
+      __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps();
+      if constexpr (act_quant_mode != PER_GROUP) {
+        y_vec = _mm512_mul_ps(y_vec, va_scale);
+      }
+      if constexpr (wei_quant_mode == PER_ROW) {
+        vb_scale = _mm512_loadu_ps(scales_b + j);
+      }
+      if constexpr (wei_quant_mode != PER_GROUP) {
+        y_vec = _mm512_mul_ps(y_vec, vb_scale);
+      }
+      y_vec = _mm512_add_ps(y_vec, bias_vec);
+      if constexpr (std::is_same<out_dtype, float>::value) {
+        _mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
+      } else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) {
+        __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec);
+        _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec);
+      } else if constexpr (std::is_same<out_dtype, at::Half>::value) {
+        __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec);
+        _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec);
+      } else {
+        TORCH_CHECK(false, "Unsupported output dtype");
+      }
+    });
+    constexpr int tail_start = N / 16 * 16;
+    for (int j = tail_start; j < N; ++j) {
+      if constexpr (wei_quant_mode == PER_ROW) {
+        b_scale = scales_b[j];
+      }
+      c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale);
+    }
+  }  // for M
+}
+
+#else  // no AVX512
+
+template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode>
+inline void store_out(
+    const float* y_buf,
+    out_dtype* c_ptr,
+    int64_t M,
+    int64_t lda,
+    const float* scales_a,
+    const float* scales_b,
+    const float* bias) {
+  float a_scale = 1.0, b_scale = 1.0;
+  if constexpr (act_quant_mode == PER_TENSOR) {
+    a_scale = *scales_a;
+  }
+  if constexpr (wei_quant_mode == PER_TENSOR) {
+    b_scale = *scales_b;
+  }
+  for (int i = 0; i < M; ++i) {
+    if constexpr (act_quant_mode == PER_ROW) {
+      a_scale = *(scales_a + i);
+    }
+    for (int j = 0; j < N; ++j) {
+      if constexpr (wei_quant_mode == PER_ROW) {
+        b_scale = scales_b[j];
+      }
+      float val = y_buf[i * N + j] * a_scale * b_scale;
+      if (bias) val += bias[j];
+      c_ptr[i * lda + j] = static_cast<out_dtype>(val);
+    }
+  }
+}
+
+#endif  // CPU_CAPABILITY_AVX512
+
+}  // anonymous namespace
+
+// FP8 W8A8 fused MoE kernel.
+//
+//   hidden_states: [M, K]   FP8 E4M3, pre-quantized activation
+//   packed_w1:  [E, 2*Nc, Kc, BLOCK_K, BLOCK_N]  FP8 W8A8 prepack (float8_linear_prepack_cpu)
+//   packed_w2:  [E, K, N_packed]  FP8 W8A16 VNNI pack (convert_weight_packed)
+//   As:  [M]   float32 per-token activation scales
+//   w1s: [E, num_groups, 2N]  float32 per-group weight scales for w1
+//   w2s: [E, ...]             float32 per-group weight scales for w2
+//
+// Stage 1: hidden_states(FP8) @ w1(FP8) → ic0(BF16)  [M*topk, 2N]
+//          Uses FP8×FP8 GEMM (tinygemm_kernel with FP8 inputs)
+// Stage 1.5: ic1 = silu(ic0[:,:N]) * ic0[:,N:]      [M*topk, N]
+// Stage 2: ic1(BF16) @ w2(FP8) → ic2(BF16)          [M*topk, K]
+//          Uses BF16×FP8 GEMM (W8A16 path, tinygemm_kernel<scalar_t>)
+// Stage 3: out = weighted_sum(ic2, topk_weights)     [M, K]
+
+template <typename scalar_t>
+void fused_experts_fp8_a8_kernel_impl(
+    scalar_t* __restrict__ output,
+    scalar_t* __restrict__ ic0,
+    scalar_t* __restrict__ ic1,
+    scalar_t* __restrict__ ic2,
+    at::Float8_e4m3fn* __restrict__ A_tmp,
+    scalar_t* __restrict__ B_tmp,
+    float* __restrict__ C_tmp,
+    float* __restrict__ Ukernel_tmp,
+    const at::Float8_e4m3fn* __restrict__ input,
+    const at::Float8_e4m3fn* __restrict__ packed_w1,
+    const at::Float8_e4m3fn* __restrict__ packed_w2,
+    const float* __restrict__ As,
+    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) {
+  //   1. intermediate_cache1 : [M * topk, N]
+  //   2. intermediate_cache2 : [M * topk, K]
+  //   3. A_tmp : [T, BLOCK_M * K]  FP8
+  //   4. C_tmp : [T, 2 * BLOCK_M * BLOCK_N]  float32
+  //   5. intermediate_cache0 : [M * topk, 2N]
+  //   6. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, max(K, N)]  scalar_t
+  //   7. Ukernel_tmp: [T, 2 * BLOCK_M * BLOCK_N]  float32
+  //   8. (non-FP8-brgemm only) dqA [T, BLOCK_M * BLOCK_K] + dqB [T, BLOCK_K * BLOCK_N] in BF16
+  //      appended after Ukernel_tmp; see moe.cpp FP8_W8A8 buffer allocation
+
+  constexpr int64_t BLOCK_M = block_size_m();
+  constexpr int64_t BLOCK_N_VAL = block_size_n();  // = 32
+  const int64_t KB = K / BLOCK_K;  // K blocks
+
+  int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N_VAL * std::max(K, N);
+
+  // stage 1: intermediate_cache0 = hidden_states @ w1
+  const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
+  const int64_t NB = div_up(N, BLOCK_N_VAL);
+
+  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_VAL;
+  int64_t num_groups = div_up(K, block_size_K);   // G
+  int64_t blocks_k_per_group = block_size_K / BLOCK_K;
+  const int64_t stride_e = 2 * N * K;  // per-expert stride for w1 (total elements)
+  bool use_brgemm = true;
+  int64_t num_thread = at::get_num_threads();
+
+  // For the BF16 fallback path (no native FP8 brgemm), tinygemm_kernel needs temporary
+  // BF16 buffers to dequantize FP8 inputs before running AMX-BF16 brgemm.
+  // These are placed after the Ukernel_tmp region (see moe.cpp FP8_W8A8 buffer allocation).
+#ifndef CPUBLAS_BRGEMM_F8F8F32
+  // dqA: [T, BLOCK_M, BLOCK_K] BF16
+  // dqB: [T, BLOCK_K, BLOCK_N] BF16
+  // Layout: immediately after Ukernel_tmp[num_thread * 2 * BLOCK_M * BLOCK_N_VAL]
+  at::BFloat16* __restrict__ dq_base = reinterpret_cast<at::BFloat16*>(
+      Ukernel_tmp + num_thread * 2 * BLOCK_M * BLOCK_N_VAL);
+  const int64_t dq_per_thread = BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N_VAL;
+#endif
+
+  at::parallel_for(0, MB * NB, 1, [&](int64_t begin, int64_t end) {
+    int tid = get_thread_num();
+    // A_tmp: FP8 activation buffer for this thread [BLOCK_M, K]
+    at::Float8_e4m3fn* __restrict__ A = A_tmp + tid * BLOCK_M * K;
+    float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL;
+    float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N_VAL;
+    alignas(64) float As_[BLOCK_M];
+    float* __restrict__ ukernel_buf_1 = Ukernel_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL;
+#ifndef CPUBLAS_BRGEMM_F8F8F32
+    at::BFloat16* __restrict__ dqA_buf = dq_base + tid * dq_per_thread;
+    at::BFloat16* __restrict__ dqB_buf = dqA_buf + BLOCK_M * BLOCK_K;
+#else
+    at::BFloat16* __restrict__ dqA_buf = nullptr;
+    at::BFloat16* __restrict__ dqB_buf = nullptr;
+#endif
+    float* __restrict__ ukernel_buf_2 = ukernel_buf_1 + BLOCK_M * BLOCK_N_VAL;
+    auto ldsa = 1;  // activation per-row quantized
+
+    for (int64_t i = begin; i < end; ++i) {
+      int64_t mb = i / NB;
+      int64_t nb = i % NB;
+      int64_t nb1 = nb + NB;  // gate is first N columns, up is second N columns
+
+      int64_t n_size = std::min(N - nb * BLOCK_N_VAL, BLOCK_N_VAL);
+
+      // B (weight) for expert expert_id: gate+up fused [2N, K] prepackaged as FP8
+      int32_t expert_id = expert_ids[mb];
+      const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e;
+      const float* __restrict__ Bs = w1s + expert_id * num_groups * (2 * N);
+
+      // Load activation tile for this M-block
+      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
+      int64_t m_size = offsets[mb + 1] - offsets[mb];
+
+      for (int64_t m = 0; m < m_size; ++m) {
+        int32_t index = A_ids[m] / topk;
+        copy_stub2(A + m * K, input + index * K, K);
+        As_[m] = As[index];
+      }
+
+      const int64_t offset = offsets[mb];
+
+      // Gate path (nb): w1[:N, :]
+      zero_buffer(C0, BLOCK_M * BLOCK_N_VAL);
+      for (int kci = 0; kci < KB; ++kci) {
+        // FP8×FP8 tinygemm: act_quant=PER_ROW(2), wei_quant=PER_GROUP(3)
+        tinygemm_kernel(
+            /* C       */ C0,
+            /* A       */ A + kci * BLOCK_K,
+            /* scales_a*/ As_ + kci / blocks_k_per_group,
+            /* B       */ B + (nb * KB + kci) * BLOCK_K * BLOCK_N_VAL,
+            /* scales_b*/ Bs + nb * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL,
+            /* M       */ m_size,
+            /* K       */ BLOCK_K,
+            /* lda     */ K,
+            /* ldc     */ BLOCK_N_VAL,
+            /* ldsa    */ ldsa,
+            /* ukernel */ ukernel_buf_1,
+            /* dqA_buf */ dqA_buf,
+            /* dqB_buf */ dqB_buf);
+      }
+      store_out<scalar_t, BLOCK_N_VAL, PER_ROW, PER_GROUP>(
+          C0,
+          ic0 + offset * 2 * N + nb * BLOCK_N_VAL,
+          m_size,
+          2 * N,  // lda
+          As_,
+          nullptr,
+          nullptr);
+
+      // Up path (nb1): w1[N:, :]
+      zero_buffer(C1, BLOCK_M * BLOCK_N_VAL);
+      for (int kci = 0; kci < KB; ++kci) {
+        tinygemm_kernel(
+            /* C       */ C1,
+            /* A       */ A + kci * BLOCK_K,
+            /* scales_a*/ As_ + kci / blocks_k_per_group,
+            /* B       */ B + (nb1 * KB + kci) * BLOCK_K * BLOCK_N_VAL,
+            /* scales_b*/ Bs + nb1 * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL,
+            /* M       */ m_size,
+            /* K       */ BLOCK_K,
+            /* lda     */ K,
+            /* ldc     */ BLOCK_N_VAL,
+            /* ldsa    */ ldsa,
+            /* ukernel */ ukernel_buf_2,
+            /* dqA_buf */ dqA_buf,
+            /* dqB_buf */ dqB_buf);
+      }
+      store_out<scalar_t, BLOCK_N_VAL, PER_ROW, PER_GROUP>(
+          C1,
+          ic0 + offset * 2 * N + nb1 * BLOCK_N_VAL,
+          m_size,
+          2 * N,  // lda
+          As_,
+          nullptr,
+          nullptr);
+    }
+    if (use_brgemm) {
+      at::native::cpublas::brgemm_release();
+    }
+  });
+
+  // stage 1.5: intermediate_cache1 = silu(ic0[:, :N]) * ic0[:, N:]
+  at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t m = begin; m < end; ++m) {
+      silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
+    }
+  });
+
+  // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
+  //   w2: [E, K, N]  FP8, VNNI packed; treated as W8A16 (BF16 act × FP8 weight)
+  const int64_t OC = K;    // output channels (hidden dim)
+  const int64_t IC = N;    // input channels (expert intermediate dim)
+  const int64_t MB2 = MB;
+  const int64_t NB2 = div_up(OC, BLOCK_N_VAL);
+  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;   // per-expert stride for w2
+  const int64_t stride_oc = IC;
+
+  int64_t avg_M = std::max(int64_t(1), M * topk / E);
+  bool use_brgemm2 = can_use_brgemm<at::Float8_e4m3fn>(avg_M);
+
+  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_VAL * 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_VAL, BLOCK_N_VAL);
+
+      // A from ic1: [M*topk, N] in sorted order (contiguous after stage 1.5)
+      const scalar_t* __restrict__ A2 = ic1 + offsets[mb] * N;
+      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
+
+      // B from w2: [E, K, N] VNNI packed, offset to expert + output block
+      int32_t expert_id = expert_ids[mb];
+      const at::Float8_e4m3fn* __restrict__ B2 =
+          packed_w2 + expert_id * stride_e2 + nb * BLOCK_N_VAL * stride_oc;
+      const float* __restrict__ Bs2 =
+          w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
+
+      // Only unpack B on first visit or new expert
+      int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
+      bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
+
+      tinygemm_kernel<scalar_t>(
+          /* A          */ A2,
+          /* B          */ B2,
+          /* C          */ C,
+          /* Btmp       */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N_VAL * IC,
+          /* Ctmp       */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL,
+          /* Bbias      */ nullptr,
+          /* scale      */ Bs2,
+          /* M          */ m_size,
+          /* N          */ n_size,
+          /* K          */ IC,
+          /* lda        */ IC,
+          /* ldb        */ n_size,
+          /* ldc        */ BLOCK_N_VAL,
+          /* brg        */ use_brgemm2,
+          /* block_sz_K */ block_size_K,
+          /* do_unpack  */ do_unpack);
+
+      // Scatter ic2 in original token order, weighting by topk_weights
+      for (int64_t m = 0; m < m_size; ++m) {
+        int32_t index = A_ids[m];
+        float weight = topk_weights[index];
+        copy_mul_stub(ic2 + index * K + nb * BLOCK_N_VAL, C + m * BLOCK_N_VAL, weight, n_size);
+      }
+    });
+
+    if (use_brgemm2) {
+      at::native::cpublas::brgemm_release();
+    }
+  });
+
+  // stage 3: output = sum(ic2, dim=topk)  [M, K]
+  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t m = begin; m < end; ++m) {
+      sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
+    }
+  });
+}
+
+#define INSTANTIATE_MOE_FP8_A8_TEMPLATE(TYPE)           \
+  template void fused_experts_fp8_a8_kernel_impl<TYPE>( \
+      TYPE* __restrict__ output,                        \
+      TYPE* __restrict__ ic0,                           \
+      TYPE* __restrict__ ic1,                           \
+      TYPE* __restrict__ ic2,                           \
+      at::Float8_e4m3fn* __restrict__ A_tmp,            \
+      TYPE* __restrict__ B_tmp,                         \
+      float* __restrict__ C_tmp,                        \
+      float* __restrict__ Ukernel_tmp,                  \
+      const at::Float8_e4m3fn* __restrict__ input,      \
+      const at::Float8_e4m3fn* __restrict__ packed_w1,  \
+      const at::Float8_e4m3fn* __restrict__ packed_w2,  \
+      const float* __restrict__ As,                     \
+      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)
+
+INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::BFloat16);
+INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::Half);
diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp
index 8f663a93c36d..accd5fcec354 100644
--- a/csrc/cpu/torch_bindings.cpp
+++ b/csrc/cpu/torch_bindings.cpp
@@ -73,6 +73,7 @@ void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states,
                        const std::optional<at::Tensor>& w2_scale,
                        const std::optional<at::Tensor>& w1_zero,
                        const std::optional<at::Tensor>& w2_zero,
+                       const std::optional<at::Tensor>& a1_scale,
                        const std::optional<std::vector<int64_t>> block_size,
                        const std::optional<at::Tensor>& w1_bias,
                        const std::optional<at::Tensor>& w2_bias,
@@ -91,6 +92,18 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
                              const std::optional<at::Tensor>& bias,
                              at::ScalarType out_dtype, bool is_vnni);
 
+// Adapted from sglang: FP8 W8A8 kernels
+std::tuple<at::Tensor, at::Tensor> float8_linear_prepack_impl(
+    const at::Tensor& weight, const at::Tensor& scales);
+at::Tensor fp8_scaled_mm_with_quant(const at::Tensor& act,
+                                    const std::optional<at::Tensor>& act_scales,
+                                    bool channelwise, const at::Tensor& weight,
+                                    const at::Tensor& weight_scales,
+                                    const std::optional<at::Tensor>& bias,
+                                    at::ScalarType output_dtype);
+std::tuple<at::Tensor, at::Tensor> quantize_fp8e4m3_vec(
+    const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt);
+
 // Adapted from sglang: MLA CPU kernels (AMX-only)
 void decode_attention_cpu(at::Tensor& query, at::Tensor& k_buffer,
                           at::Tensor& v_buffer, at::Tensor& output,
@@ -281,6 +294,7 @@ void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
                          const std::string& activation);
 
 bool cpu_attn_has_isa(const std::string& isa);
+bool cpu_has_amx_fp8();
 
 torch::Tensor get_scheduler_metadata(
     const int64_t num_req, const int64_t num_heads_q,
@@ -611,7 +625,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "fused_experts_cpu(Tensor(a0!) out, Tensor hidden_states, Tensor w1, "
       "Tensor w2, Tensor topk_weights, Tensor topk_ids, "
       "int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
-      "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, "
+      "Tensor? w1_zero, Tensor? w2_zero, Tensor? a1_scale, int[]? block_size, "
       "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, "
       "bool is_vnni) -> "
       "()");
@@ -629,7 +643,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "bool is_vnni) -> Tensor");
   ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
 
-  // Adapted from sglang: casual_conv1d kernels
+  // Adapted from sglang: FP8 W8A8 kernels
+  ops.def(
+      "float8_linear_prepack_cpu(Tensor weight, Tensor scales) -> (Tensor, "
+      "Tensor)");
+  ops.impl("float8_linear_prepack_cpu", torch::kCPU,
+           &float8_linear_prepack_impl);
+  ops.def(
+      "fp8_scaled_mm_with_quant(Tensor act, Tensor? act_scales, bool "
+      "channelwise, "
+      "Tensor weight, Tensor weight_scales, Tensor? bias, ScalarType "
+      "output_dtype) -> Tensor");
+  ops.impl("fp8_scaled_mm_with_quant", torch::kCPU, &fp8_scaled_mm_with_quant);
+  ops.def(
+      "quantize_fp8e4m3_vec(Tensor input, bool channelwise, Tensor? scale_opt) "
+      "-> (Tensor, Tensor)");
+  ops.impl("quantize_fp8e4m3_vec", torch::kCPU, &quantize_fp8e4m3_vec);
+
   ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
   ops.impl("causal_conv1d_weight_pack", torch::kCPU,
            &causal_conv1d_weight_pack);
@@ -832,6 +862,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "-> (Tensor, Tensor)");
   ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
 
+  // Pure runtime CPU ISA capability check.
+  ops.def("cpu_has_amx_fp8() -> bool", &cpu_has_amx_fp8);
+
   // CPU attention kernels
   ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa);
   ops.def(
@@ -883,8 +916,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "bool skip_weighted, "
       "str act, str isa) -> ()");
   ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
-#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
-    !defined(__APPLE__)
+
+#if (defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
+     !defined(__APPLE__)) ||                                   \
+    defined(__powerpc64__)
   ops.def(
       "prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, "
       "str isa) -> ()");
@@ -895,8 +930,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool "
       "skip_weighted, str act, str isa) -> ()");
   ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8);
-#endif  // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) &&
-        // !defined(__APPLE__)
+#endif  // ARM_I8MM+BF16 or __powerpc64__
   ops.def(
       "mla_decode_kvcache("
       "   Tensor! out, Tensor query, Tensor kv_cache,"
diff --git a/csrc/libtorch_stable/all_reduce_mhc.cu b/csrc/libtorch_stable/all_reduce_mhc.cu
new file mode 100644
index 000000000000..dc58371156a2
--- /dev/null
+++ b/csrc/libtorch_stable/all_reduce_mhc.cu
@@ -0,0 +1,212 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#include "torch_utils.h"
+#include "custom_all_gather_reduce_scatter.cuh"
+#include <cooperative_groups.h>
+#include <torch/csrc/stable/library.h>
+
+namespace {
+using Tensor = torch::stable::Tensor;
+using Pack = vllm::packed_t<nv_bfloat16>::P;
+constexpr int kHidden = 5120;
+constexpr int kRanks = 4;
+constexpr int kThreads = 128;
+constexpr int kCluster = 5;
+constexpr int kPacks = kHidden / Pack::size;
+
+__global__ __launch_bounds__(kThreads)
+    __cluster_dims__(1, kCluster, 1) void all_reduce_mhc_kernel(
+        const Pack* input, const Pack* residual, const float* post,
+        const float* comb, const float* pre, const Pack* weight, Pack* output,
+        Pack* normalized, Pack* local, Pack* multicast, uint32_t* epochs,
+        int rank, int stage_size, float eps) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  cudaGridDependencySynchronize();
+  const int token = blockIdx.x;
+  const int column = blockIdx.y * kThreads + threadIdx.x;
+  const int idx = token * kPacks + column;
+  const int size = gridDim.x * kPacks;
+  const int stage = epochs[0] % 3;
+  const int dirty = vllm::mnnvl_lamport_dirty_stage(stage);
+  const int dirty_size = epochs[2 + dirty];
+  Pack value = input[idx];
+  vllm::store_multimem_lamport_payload(
+      multicast + stage * stage_size + rank * size + idx,
+      vllm::sanitize_lamport_payload(value));
+  // Multi-token forwards benefit from launching dependents before mHC finishes.
+  if (gridDim.x > 1) cudaTriggerProgrammaticLaunchCompletion();
+  vllm::lamport_cta_arrive(epochs + 1);
+  for (int i = idx; i < dirty_size; i += size) {
+    local[dirty * stage_size + i] = vllm::lamport_sentinel<Pack>();
+  }
+  Pack peers[kRanks];
+  bool pending;
+  do {
+    vllm::LamportPack<Pack> packets[kRanks];
+  #pragma unroll
+    for (int peer = 0; peer < kRanks; ++peer) {
+      auto ptr = local + stage * stage_size + idx + peer * size;
+      asm volatile("ld.relaxed.sys.global.v4.u32 {%0,%1,%2,%3}, [%4];"
+                   : "=r"(packets[peer].words[0]), "=r"(packets[peer].words[1]),
+                     "=r"(packets[peer].words[2]), "=r"(packets[peer].words[3])
+                   : "l"(ptr)
+                   : "memory");
+    }
+    pending = false;
+  #pragma unroll
+    for (int peer = 0; peer < kRanks; ++peer) {
+      pending |= vllm::is_lamport_dirty(packets[peer]);
+      peers[peer] = packets[peer].packed;
+    }
+  } while (pending);
+  auto sum = vllm::upcast(peers[0]);
+  #pragma unroll
+  for (int r = 1; r < kRanks; ++r) {
+    vllm::packed_assign_add(sum, vllm::upcast(peers[r]));
+  }
+  // Keep the BF16 boundary between the collective and mHC post.
+  auto reduced = vllm::upcast(vllm::downcast<Pack>(sum));
+  decltype(sum) streams[4];
+  #pragma unroll
+  for (int r = 0; r < 4; ++r) {
+    streams[r] = vllm::upcast(residual[(token * 4 + r) * kPacks + column]);
+  }
+  decltype(sum) collapsed{};
+  #pragma unroll
+  for (int c = 0; c < 4; ++c) {
+    decltype(sum) mixed;
+  #pragma unroll
+    for (int e = 0; e < Pack::size; ++e) {
+      // Match post*x followed by the four residual FMAs.
+      float acc = __fmul_rn(post[token * 4 + c], reduced.data[e]);
+  #pragma unroll
+      for (int r = 0; r < 4; ++r) {
+        acc = fmaf(comb[token * 16 + r * 4 + c], streams[r].data[e], acc);
+      }
+      mixed.data[e] = acc;
+    }
+    Pack rounded = vllm::downcast<Pack>(mixed);
+    output[(token * 4 + c) * kPacks + column] = rounded;
+    auto as_float = vllm::upcast(rounded);
+  #pragma unroll
+    for (int e = 0; e < Pack::size; ++e) {
+      collapsed.data[e] =
+          fmaf(pre[token * 4 + c], as_float.data[e], collapsed.data[e]);
+    }
+  }
+  // RMSNorm consumes the rounded BF16 collapse in the unfused path.
+  collapsed = vllm::upcast(vllm::downcast<Pack>(collapsed));
+  float sqr = 0;
+  #pragma unroll
+  for (int e = 0; e < Pack::size; ++e) {
+    sqr = fmaf(collapsed.data[e], collapsed.data[e], sqr);
+  }
+  #pragma unroll
+  for (int delta = 16; delta; delta /= 2) {
+    sqr += __shfl_xor_sync(0xffffffff, sqr, delta);
+  }
+  __shared__ float warp_sums[kCluster][kThreads / 32];
+  auto cluster = cooperative_groups::this_cluster();
+  const int lane = threadIdx.x % 32;
+  if (lane < kCluster) {
+    *cluster.map_shared_rank(&warp_sums[blockIdx.y][threadIdx.x / 32], lane) =
+        sqr;
+  }
+  cluster.sync();
+  float total = 0;
+  #pragma unroll
+  for (int c = 0; c < kCluster; ++c) {
+  #pragma unroll
+    for (int w = 0; w < kThreads / 32; ++w) total += warp_sums[c][w];
+  }
+  const float scale = rsqrtf(__fadd_rn(total / kHidden, eps));
+  auto weights = vllm::upcast(weight[column]);
+  #pragma unroll
+  for (int e = 0; e < Pack::size; ++e) {
+    collapsed.data[e] = collapsed.data[e] * scale * weights.data[e];
+  }
+  normalized[idx] = vllm::downcast<Pack>(collapsed);
+  if (idx == 0) {
+    while (*reinterpret_cast<volatile uint32_t*>(epochs + 1) <
+           gridDim.x * gridDim.y) {
+    }
+    epochs[2 + stage] = kRanks * size;
+    epochs[0] = vllm::mnnvl_lamport_next_stage(stage);
+    epochs[1] = 0;
+  }
+  if (gridDim.x == 1) cudaTriggerProgrammaticLaunchCompletion();
+#else
+  asm volatile("trap;");
+#endif
+}
+
+void all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb,
+                    Tensor pre, Tensor weight, Tensor output, Tensor normalized,
+                    int64_t local, int64_t multicast, Tensor epochs,
+                    int64_t rank, int64_t stage_bytes, double eps) {
+  using torch::headeronly::ScalarType;
+  const torch::stable::accelerator::DeviceGuard guard(input.get_device_index());
+  const int64_t n = input.size(0);
+  STD_TORCH_CHECK(input.dim() == 2 && input.size(1) == kHidden && n > 0 &&
+                  n <= 16);
+  for (const auto& t : {input, residual, weight, output, normalized}) {
+    STD_TORCH_CHECK(t.is_contiguous() &&
+                    t.scalar_type() == ScalarType::BFloat16);
+    STD_TORCH_CHECK(t.get_device_index() == input.get_device_index());
+    STD_TORCH_CHECK(reinterpret_cast<uintptr_t>(t.const_data_ptr()) % 16 == 0);
+  }
+  for (const auto& t : {post, comb, pre}) {
+    STD_TORCH_CHECK(t.is_contiguous() && t.scalar_type() == ScalarType::Float);
+    STD_TORCH_CHECK(t.get_device_index() == input.get_device_index());
+  }
+  STD_TORCH_CHECK(residual.numel() == n * 4 * kHidden &&
+                  output.numel() == residual.numel() &&
+                  normalized.numel() == input.numel() &&
+                  weight.numel() == kHidden);
+  STD_TORCH_CHECK(post.numel() == n * 4 && pre.numel() == n * 4 &&
+                  comb.numel() == n * 16);
+  STD_TORCH_CHECK(epochs.is_contiguous() &&
+                  epochs.scalar_type() == ScalarType::Int &&
+                  epochs.numel() >= 5 &&
+                  epochs.get_device_index() == input.get_device_index());
+  STD_TORCH_CHECK(local && multicast && rank >= 0 && rank < kRanks &&
+                  stage_bytes % 16 == 0 &&
+                  stage_bytes >= n * kHidden * 2 * kRanks);
+  STD_TORCH_CHECK(get_device_prop()->major == 10);
+  cudaLaunchAttribute attr{};
+  attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attr.val.programmaticStreamSerializationAllowed = 1;
+  cudaLaunchConfig_t config{
+      .gridDim = dim3(n, kCluster),
+      .blockDim = dim3(kThreads),
+      .dynamicSmemBytes = 0,
+      .stream = get_current_cuda_stream(input.get_device_index()),
+      .attrs = &attr,
+      .numAttrs = 1};
+  STD_CUDA_CHECK(cudaLaunchKernelEx(
+      &config, all_reduce_mhc_kernel,
+      static_cast<const Pack*>(input.const_data_ptr()),
+      static_cast<const Pack*>(residual.const_data_ptr()),
+      static_cast<const float*>(post.const_data_ptr()),
+      static_cast<const float*>(comb.const_data_ptr()),
+      static_cast<const float*>(pre.const_data_ptr()),
+      static_cast<const Pack*>(weight.const_data_ptr()),
+      static_cast<Pack*>(output.mutable_data_ptr()),
+      static_cast<Pack*>(normalized.mutable_data_ptr()),
+      reinterpret_cast<Pack*>(local), reinterpret_cast<Pack*>(multicast),
+      static_cast<uint32_t*>(epochs.mutable_data_ptr()), static_cast<int>(rank),
+      static_cast<int>(stage_bytes / sizeof(Pack)), static_cast<float>(eps)));
+}
+}  // namespace
+
+STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, mhc) {
+  mhc.def(
+      "all_reduce_mhc(Tensor input, Tensor residual, Tensor post, Tensor comb, "
+      "Tensor pre, Tensor weight, Tensor! output, Tensor! normalized, "
+      "int local, int multicast, Tensor! epochs, int rank, int stage_bytes, "
+      "float eps) -> ()");
+}
+STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, mhc) {
+  mhc.impl("all_reduce_mhc", TORCH_BOX(&all_reduce_mhc));
+}
diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu
index 79aa63a5b83a..9130adf87d69 100644
--- a/csrc/libtorch_stable/cache_kernels.cu
+++ b/csrc/libtorch_stable/cache_kernels.cu
@@ -19,7 +19,6 @@
 
 #include <algorithm>
 #include <cassert>
-#include <cfloat>
 #include <cstdlib>
 
 #ifdef USE_ROCM
@@ -84,7 +83,9 @@ void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
 namespace {
 // ROCm hipMemcpyBatchAsync faults for count > 8192 (MI350X/gfx950), so chunk
 // at that ceiling.
+#if defined(USE_ROCM)
 constexpr int64_t kRocmDefaultMaxBatchDescriptors = 8192;
+#endif
 
 constexpr const char* kMaxBatchDescriptorsEnv =
     "VLLM_KV_OFFLOAD_MAX_BATCH_DESCRIPTORS";
@@ -456,14 +457,14 @@ __global__ void concat_and_cache_mla_kernel(
 // Grouped variant of concat_and_cache_mla: inserts the context K/V for every
 // draft layer in a single launch. Grid is (num_tokens, num_layers); each layer
 // reads its own cache base pointer from kv_cache_ptrs (same pointer-array
-// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no
-// scaling or quantization; scalar_t is uint16_t for portability.
-template <typename scalar_t>
+// pattern as copy_blocks_kernel).
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __global__ void concat_and_cache_mla_grouped_kernel(
     const scalar_t* __restrict__ kv_c,  // [num_layers, num_tokens,
                                         // kv_lora_rank]
     const scalar_t* __restrict__ k_pe,  // [num_layers, num_tokens, pe_dim]
     const int64_t* __restrict__ kv_cache_ptrs,  // [num_layers]
+    const float* __restrict__ kv_scales,        // [num_layers] or nullptr
     const int64_t* __restrict__ slot_mapping,   // [num_layers, num_tokens]
     const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride,
     const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride,
@@ -481,12 +482,16 @@ __global__ void concat_and_cache_mla_grouped_kernel(
   const int64_t block_idx = slot_idx / block_size;
   const int64_t block_offset = slot_idx % block_size;
 
-  scalar_t* __restrict__ kv_cache =
-      reinterpret_cast<scalar_t*>(kv_cache_ptrs[layer_idx]);
+  cache_t* __restrict__ kv_cache =
+      reinterpret_cast<cache_t*>(kv_cache_ptrs[layer_idx]);
   const scalar_t* __restrict__ kv_c_layer =
       kv_c + layer_idx * kv_c_layer_stride;
   const scalar_t* __restrict__ k_pe_layer =
       k_pe + layer_idx * k_pe_layer_stride;
+  float scale = 0.0f;
+  if constexpr (kv_dt != Fp8KVCacheDataType::kAuto) {
+    scale = kv_scales[layer_idx];
+  }
 
   auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride,
                   int size, int offset) {
@@ -494,7 +499,12 @@ __global__ void concat_and_cache_mla_grouped_kernel(
       const int64_t src_idx = token_idx * src_token_stride + i;
       const int64_t dst_idx =
           block_idx * block_stride + block_offset * entry_stride + i + offset;
-      kv_cache[dst_idx] = src[src_idx];
+      if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
+        kv_cache[dst_idx] = src[src_idx];
+      } else {
+        kv_cache[dst_idx] =
+            fp8::scaled_convert<cache_t, scalar_t, kv_dt>(src[src_idx], scale);
+      }
     }
   };
 
@@ -539,16 +549,20 @@ __global__ void concat_and_cache_ds_mla_kernel(
   scalar_t* kv_cache_16bit =
       reinterpret_cast<scalar_t*>(&kv_cache[dst_idx_start]);
 
-  // The last warp handles the RoPE part
+  // Zero the reserved RoPE tail for NoPE rows.
   if (threadIdx.x >= 64) {
     // Each thread handles two elements of RoPE
     const int8_t pe_idx_start = (threadIdx.x - 64) * 2;
-    const int64_t src_idx = token_idx * k_pe_stride + pe_idx_start;
-    // Vectorized load of two 16-bit values, performed as one 32-bit load
-    const int32_t vals = *reinterpret_cast<const int32_t*>(&k_pe[src_idx]);
     // RoPE values start after the packed 8-bit NoPE values and the
     // 32-bit scales
     const int64_t dst_idx = kv_lora_rank / 2 + 8 + pe_idx_start;
+    if (pe_dim == 0) {
+      *reinterpret_cast<int32_t*>(&kv_cache_16bit[dst_idx]) = 0;
+      return;
+    }
+    const int64_t src_idx = token_idx * k_pe_stride + pe_idx_start;
+    // Vectorized load of two 16-bit values, performed as one 32-bit load
+    const int32_t vals = *reinterpret_cast<const int32_t*>(&k_pe[src_idx]);
     // Vectorized store of two 16-bit values, performed as one 32-bit store
     *reinterpret_cast<int32_t*>(&kv_cache_16bit[dst_idx]) = vals;
     return;
@@ -578,8 +592,9 @@ __global__ void concat_and_cache_ds_mla_kernel(
     max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 16));
   }
 
-  // Compute the scale for the tile
-  float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN);
+  // Both SM90 and SM100 readers preserve power-of-two fp32 scales exactly.
+  float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f);
+  tile_scale = exp2f(ceilf(log2f(tile_scale)));
 
   // The first lane of each half-warp writes the scale to kv_cache
   if ((lane_idx == 0) || (lane_idx == 16)) {
@@ -589,7 +604,7 @@ __global__ void concat_and_cache_ds_mla_kernel(
   }
 
   // Now all threads in the block scale and write their elements
-  // NoPE data is packed in the first kv_lora_rank/2 bytes (first 256 bytes)
+  // NoPE data occupies the first kv_lora_rank bytes (512 bytes).
   const int64_t dst_idx_base = dst_idx_start + (threadIdx.x * 8);
 
   uint8_t result[8];
@@ -925,7 +940,8 @@ void concat_and_cache_mla(
   if (kv_cache_dtype == "fp8_ds_mla") {
     STD_TORCH_CHECK(kv_lora_rank == 512,
                     "kv_lora_rank must be 512 for fp8_ds_mla");
-    STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla");
+    STD_TORCH_CHECK(pe_dim == 64 || pe_dim == 0,
+                    "pe_dim must be 64 or 0 for fp8_ds_mla");
     STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(),
                     "kv_cache.size(2) must be 656 bytes for fp8_ds_mla");
     STD_TORCH_CHECK(kv_c.element_size() == 2,
@@ -1001,20 +1017,55 @@ void concat_and_cache_mla_grouped(
     torch::stable::Tensor& k_pe,  // [num_layers, num_tokens, pe_dim]
     torch::stable::Tensor& kv_cache_ptrs,  // [num_layers] int64, on device
     torch::stable::Tensor& slot_mapping,   // [num_layers, num_tokens] int64
-    int64_t block_size, int64_t block_stride, int64_t entry_stride) {
-  int num_layers = kv_c.size(0);
-  int num_tokens = kv_c.size(1);
-  int kv_lora_rank = kv_c.size(2);
-  int pe_dim = k_pe.size(2);
+    int64_t block_size, int64_t block_stride, int64_t entry_stride,
+    std::optional<torch::stable::Tensor> kv_scales,  // [num_layers] or None
+    const std::string& kv_cache_dtype) {
+  const bool use_fp8 = kv_cache_dtype == "fp8" ||
+                       kv_cache_dtype == "fp8_e4m3" ||
+                       kv_cache_dtype == "fp8_e5m2";
+#ifdef USE_ROCM
+  STD_TORCH_CHECK(kv_cache_dtype != "fp8_e5m2",
+                  "concat_and_cache_mla_grouped does not support fp8_e5m2 "
+                  "KV cache on ROCm");
+#endif
+  STD_TORCH_CHECK(
+      use_fp8 || kv_cache_dtype == "auto" || kv_cache_dtype == "bfloat16",
+      "concat_and_cache_mla_grouped only supports BF16 and plain "
+      "FP8 KV cache; got ",
+      kv_cache_dtype);
 
   STD_TORCH_CHECK(
       kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
           k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16,
-      "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=",
+      "concat_and_cache_mla_grouped requires BF16 inputs; got kv_c=",
       kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type());
   STD_TORCH_CHECK(
-      kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
-      "kv_cache_ptrs must be int64");
+      kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long &&
+          slot_mapping.scalar_type() == torch::headeronly::ScalarType::Long,
+      "cache pointers and slot mapping must be int64");
+
+  const int num_layers = kv_c.size(0);
+  const int num_tokens = kv_c.size(1);
+  const int kv_lora_rank = kv_c.size(2);
+  const int pe_dim = k_pe.size(2);
+  const float* kv_scales_ptr = nullptr;
+  if (use_fp8) {
+    STD_TORCH_CHECK(kv_scales.has_value(),
+                    "FP8 grouped cache insert requires kv_scales");
+    STD_TORCH_CHECK(
+        kv_scales->scalar_type() == torch::headeronly::ScalarType::Float,
+        "kv_scales must be float32");
+    STD_TORCH_CHECK(kv_scales->numel() == num_layers,
+                    "kv_scales must contain one scale per layer");
+    STD_TORCH_CHECK(kv_scales->is_cuda() && kv_scales->get_device_index() ==
+                                                kv_c.get_device_index(),
+                    "kv_scales must be on the same CUDA device as kv_c");
+    STD_TORCH_CHECK(kv_scales->is_contiguous(), "kv_scales must be contiguous");
+    kv_scales_ptr = kv_scales->const_data_ptr<float>();
+  } else {
+    STD_TORCH_CHECK(!kv_scales.has_value(),
+                    "BF16 grouped cache insert does not use kv_scales");
+  }
 
   if (num_tokens == 0 || num_layers == 0) {
     return;
@@ -1030,17 +1081,40 @@ void concat_and_cache_mla_grouped(
       kv_c.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream();
 
-  dim3 grid(num_tokens, num_layers);
-  dim3 block(std::min(kv_lora_rank, 512));
-  vllm::concat_and_cache_mla_grouped_kernel<uint16_t>
-      <<<grid, block, 0, stream>>>(
-          reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
-          reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
-          kv_cache_ptrs.const_data_ptr<int64_t>(),
-          slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,
-          kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,
-          slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim,
-          block_size);
+  const dim3 grid(num_tokens, num_layers);
+  const dim3 block(std::min(kv_lora_rank, 512));
+
+  if (!use_fp8) {
+    vllm::concat_and_cache_mla_grouped_kernel<uint16_t, uint16_t,
+                                              vllm::Fp8KVCacheDataType::kAuto>
+        <<<grid, block, 0, stream>>>(
+            reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
+            reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
+            kv_cache_ptrs.const_data_ptr<int64_t>(), nullptr,
+            slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,
+            kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,
+            slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim,
+            block_size);
+    return;
+  }
+
+#define LAUNCH_GROUPED_FP8(KV_DTYPE)                                           \
+  vllm::concat_and_cache_mla_grouped_kernel<__nv_bfloat16, uint8_t, KV_DTYPE>  \
+      <<<grid, block, 0, stream>>>(                                            \
+          reinterpret_cast<const __nv_bfloat16*>(kv_c.data_ptr()),             \
+          reinterpret_cast<const __nv_bfloat16*>(k_pe.data_ptr()),             \
+          kv_cache_ptrs.const_data_ptr<int64_t>(), kv_scales_ptr,              \
+          slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,           \
+          kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,             \
+          slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, \
+          block_size)
+
+  if (kv_cache_dtype == "fp8_e5m2") {
+    LAUNCH_GROUPED_FP8(vllm::Fp8KVCacheDataType::kFp8E5M2);
+  } else {
+    LAUNCH_GROUPED_FP8(vllm::Fp8KVCacheDataType::kFp8E4M3);
+  }
+#undef LAUNCH_GROUPED_FP8
 }
 
 namespace vllm {
@@ -1391,7 +1465,11 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page(
     const int32_t block_size, const int32_t total_tokens,
     const int64_t block_table_stride, const int64_t cache_block_stride,
     const int64_t cache_entry_stride, const int64_t dst_entry_stride,
-    const int32_t* __restrict__ seq_starts) {
+    const int32_t* __restrict__ seq_starts,
+    const uint8_t* __restrict__ host_cache,
+    const int32_t* __restrict__ host_row_ids,
+    const int32_t* __restrict__ device_row_ids,
+    const int64_t host_entry_stride) {
   constexpr int32_t warps_per_cta = 16;
   __shared__ GatherPageTask page;
   __shared__ int32_t physical_block;
@@ -1417,8 +1495,22 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache_page(
        page_token < page.page_token_end; page_token += warps_per_cta) {
     const int32_t output_token =
         page.output_token_begin + page_token - page.page_token_begin;
-    const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride +
-                               page_token * cache_entry_stride;
+    const int32_t source_row = physical_block * block_size + page_token;
+    const uint8_t* token_ptr;
+    if (device_row_ids == nullptr) {
+      token_ptr = src_cache + physical_block * cache_block_stride +
+                  page_token * cache_entry_stride;
+    } else {
+      const int32_t device_row = device_row_ids[source_row];
+      if (device_row >= 0) {
+        token_ptr = src_cache + (device_row / block_size) * cache_block_stride +
+                    (device_row % block_size) * cache_entry_stride;
+      } else {
+        token_ptr =
+            host_cache +
+            static_cast<int64_t>(host_row_ids[source_row]) * host_entry_stride;
+      }
+    }
     __nv_bfloat16* dst_ptr = dst + output_token * dst_entry_stride;
     gather_and_upconvert_fp8_token(token_ptr, dst_ptr, lane_id);
   }
@@ -1627,7 +1719,10 @@ void cp_gather_and_upconvert_fp8_kv_cache(
     torch::stable::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
     torch::stable::Tensor const& workspace_starts,  // [BATCH]
     int64_t batch_size,
-    std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt,
+    std::optional<torch::stable::Tensor> host_cache = std::nullopt,
+    std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
+    std::optional<torch::stable::Tensor> device_row_ids = std::nullopt) {
   torch::stable::accelerator::DeviceGuard device_guard(
       src_cache.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream();
@@ -1647,6 +1742,12 @@ void cp_gather_and_upconvert_fp8_kv_cache(
         "seq_starts must be int32");
   }
 
+  const bool has_host_rows = host_cache.has_value();
+  STD_TORCH_CHECK(
+      has_host_rows == host_row_ids.has_value() &&
+          has_host_rows == device_row_ids.has_value(),
+      "host_cache, host_row_ids, and device_row_ids must be provided together");
+
   STD_TORCH_CHECK(src_cache.device() == dst.device(),
                   "src_cache and dst must be on the same device");
   STD_TORCH_CHECK(src_cache.device() == block_table.device(),
@@ -1657,6 +1758,43 @@ void cp_gather_and_upconvert_fp8_kv_cache(
     STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
                     "src_cache and seq_starts must be on the same device");
   }
+  const uint8_t* host_cache_ptr = nullptr;
+  const int32_t* host_row_ids_ptr = nullptr;
+  const int32_t* device_row_ids_ptr = nullptr;
+  int64_t host_entry_stride = 0;
+  if (has_host_rows) {
+    auto const& host = host_cache.value();
+    auto const& host_rows = host_row_ids.value();
+    auto const& device_rows = device_row_ids.value();
+    STD_TORCH_CHECK(
+        host.device().is_cpu() &&
+            host.scalar_type() == torch::headeronly::ScalarType::Byte &&
+            host.dim() == 2 && host.is_contiguous() &&
+            host.size(1) == src_cache.size(2),
+        "host_cache must be a contiguous uint8 CPU row matrix matching "
+        "src_cache");
+    cudaPointerAttributes attributes{};
+    const cudaError_t pointer_status =
+        cudaPointerGetAttributes(&attributes, host.const_data_ptr());
+    if (pointer_status != cudaSuccess) {
+      cudaGetLastError();
+    }
+    STD_TORCH_CHECK(
+        pointer_status == cudaSuccess && attributes.type == cudaMemoryTypeHost,
+        "host_cache must be CUDA-accessible pinned memory");
+    STD_TORCH_CHECK(
+        host_rows.is_cuda() && device_rows.is_cuda() &&
+            host_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
+            device_rows.scalar_type() == torch::headeronly::ScalarType::Int &&
+            host_rows.is_contiguous() && device_rows.is_contiguous() &&
+            host_rows.numel() == device_rows.numel() && host_rows.dim() <= 2 &&
+            device_rows.dim() <= 2,
+        "host/device row maps must be matching contiguous CUDA int32 tensors");
+    host_cache_ptr = reinterpret_cast<const uint8_t*>(host.const_data_ptr());
+    host_row_ids_ptr = host_rows.const_data_ptr<int32_t>();
+    device_row_ids_ptr = device_rows.const_data_ptr<int32_t>();
+    host_entry_stride = host.stride(0);
+  }
   auto dtype = src_cache.scalar_type();
   STD_TORCH_CHECK(
       dtype == torch::headeronly::ScalarType::Byte ||               // uint8
@@ -1699,7 +1837,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
       workspace_starts.const_data_ptr<int32_t>(),
       static_cast<int32_t>(batch_size), block_size, total_tokens,
       block_table_stride, cache_block_stride, cache_entry_stride,
-      dst_entry_stride, seq_starts_ptr);
+      dst_entry_stride, seq_starts_ptr, host_cache_ptr, host_row_ids_ptr,
+      device_row_ids_ptr, host_entry_stride);
 }
 
 void cp_gather_and_upconvert_nvfp4_kv_cache(
diff --git a/csrc/libtorch_stable/cuda_vec_utils.cuh b/csrc/libtorch_stable/cuda_vec_utils.cuh
index ec6e60724e65..c962c99ec255 100644
--- a/csrc/libtorch_stable/cuda_vec_utils.cuh
+++ b/csrc/libtorch_stable/cuda_vec_utils.cuh
@@ -320,10 +320,6 @@ __host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
   return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
 }
 
-__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
-  return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
-}
-
 // ============================================================
 // Packed type conversion and arithmetic
 // ============================================================
diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
index 06fe7108bd61..5a5caacce7a5 100644
--- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
@@ -3,7 +3,10 @@
  * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
  *
  * Horizontally-fused DeepseekV4-MLA kernel:
- *   - Q side:  optional per-head RMSNorm + GPT-J RoPE on last ROPE_DIM
+ *   - Q side:  optional per-head RMSNorm (apply_q_norm) and GPT-J RoPE on the
+ *              last ROPE_DIM (apply_q_rope), plus the zero-pad to
+ *              q_head_padded, in the head-major layout or (is_q_interleaved)
+ *              FlashMLA's mega-attention chunk-interleaved one
  *   - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
  *              cache insert
  *
@@ -23,9 +26,15 @@
  *   cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
  *     second along last dim; each half is rope_dim/2 = 32 values)
  *
- * Cache layout per paged-cache block (block_size tokens):
- *   [0,            bs*576):          token data, 448 fp8 + 128 bf16 each
- *   [bs*576,       bs*576 + bs*8):   UE8M0 scales, 7 real + 1 pad per token
+ * Cache layout per paged-cache block (block_size tokens), selected by the
+ * `kv_mxfp8` argument:
+ *   V4 (584 B/token, kv_mxfp8=false)
+ *     [0,      bs*576):          token data, 448 fp8 + 128 bf16 each
+ *     [bs*576, bs*576 + bs*8):   UE8M0 scales, 7 real + 1 pad per token
+ *   V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the
+ *   whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims)
+ *     [0,      bs*512):          token data, 512 fp8 each
+ *     [bs*512, bs*512 + bs*16):  UE8M0 scales, 16 per token
  */
 
 #include "torch_utils.h"
@@ -91,6 +100,10 @@ constexpr int kQuantBlock = 64;
 constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock;   // 7
 constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1;  // 8 (7 real + 1 pad)
 constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2;  // 448 + 128 = 576
+// V4.1 MXFP8 record: every dim quantized, one UE8M0 scale per 32 dims.
+constexpr int kMxfp8QuantBlock = 32;
+constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock;  // 16
+constexpr int kMxfp8TokenDataBytes = kHeadDim;                     // 512
 // FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw
 // 240.0) to match the rest of vLLM's FNUZ pipeline.
 #if defined(USE_ROCM) && defined(__gfx942__)
@@ -141,6 +154,11 @@ __device__ __forceinline__ uint4 packFp8E4M3x16(float const* values,
 // ────────────────────────────────────────────────────────────────────────────
 // Small inline helpers
 // ────────────────────────────────────────────────────────────────────────────
+__device__ __forceinline__ float warp2MaxAbs(float val) {
+  // Reduce absolute max across 2 consecutive lanes (32 dims = one MXFP8 block).
+  return fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 1));
+}
+
 __device__ __forceinline__ float warp4MaxAbs(float val) {
   // Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
   float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
@@ -159,6 +177,30 @@ __device__ __forceinline__ float warpSum(float val) {
   return val;
 }
 
+// Offset of the 16 elements a lane owns inside a Q tensor of `num_heads`
+// heads.  The default layout is head-major: a head's 512 dims are contiguous.
+// `Q_INTERLEAVED` is FlashMLA's mega-attention layout, 32 chunks of
+// [num_heads, 16] per token, which changes the address only -- lane L still
+// owns dims [16L, 16L+16) of its head either way, so `dim_base` stays the
+// true dim offset and the norm and RoPE below are equally valid in both.
+// What it does change is where the padding heads land: at the tail of every
+// chunk rather than after the live heads.
+template <bool Q_INTERLEAVED>
+__device__ __forceinline__ int64_t qElemOffset(int const tokenIdx,
+                                               int const slotIdx,
+                                               int const laneId,
+                                               int const dim_base,
+                                               int const num_heads) {
+  if constexpr (Q_INTERLEAVED) {
+    return static_cast<int64_t>(tokenIdx) * num_heads * kHeadDim +
+           static_cast<int64_t>(laneId) * num_heads * kElemsPerLane +
+           static_cast<int64_t>(slotIdx) * kElemsPerLane;
+  } else {
+    return (static_cast<int64_t>(tokenIdx) * num_heads + slotIdx) * kHeadDim +
+           dim_base;
+  }
+}
+
 // ────────────────────────────────────────────────────────────────────────────
 // Per-slot inner pipeline
 // ────────────────────────────────────────────────────────────────────────────
@@ -176,7 +218,17 @@ __device__ __forceinline__ float warpSum(float val) {
 //   slot == kNumHeadsQPadded                    → KV       (RoPE + UE8M0 quant
 //                                                           + paged-cache
 //                                                           insert)
-template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
+//
+// `kv_mxfp8` and `apply_q_rope` are grid-uniform, so they stay runtime
+// arguments rather than multiplying every (head count, q-norm, q-layout)
+// instantiation of a kernel whose hot path is the Q branch.  `Q_INTERLEAVED`
+// cannot: it picks the address arithmetic, which has to fold.
+//
+// The three Q knobs are independent.  The mega-attention layer uses
+// Q_INTERLEAVED with both norm and RoPE off, because its attention kernel
+// applies those itself -- not because the layout forbids them.
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM,
+          bool Q_INTERLEAVED = false>
 __device__ __forceinline__ void processDeepseekV4Slot(
     uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
     int const dim_base, int const laneId, int const num_heads_q,
@@ -184,7 +236,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
     uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
     int64_t const* __restrict__ position_ids,
     float const* __restrict__ cos_sin_cache, int const cache_block_size,
-    int const kv_block_stride) {
+    int const kv_block_stride, bool const kv_mxfp8, bool const apply_q_rope) {
   using Converter = vllm::_typeConvert<scalar_t_in>;
   bool const isKV = (slotIdx == kNumHeadsQPadded);
   bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
@@ -194,10 +246,8 @@ __device__ __forceinline__ void processDeepseekV4Slot(
   // zero literal is correct.  Matches the live-Q branch's vectorized store.
   if (isPadQ) {
     scalar_t_in* dst =
-        q_out +
-        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
-            kHeadDim +
-        dim_base;
+        q_out + qElemOffset<Q_INTERLEAVED>(tokenIdx, slotIdx, laneId, dim_base,
+                                           kNumHeadsQPadded);
     uint4 const zero4 = {0u, 0u, 0u, 0u};
     *reinterpret_cast<uint4*>(dst) = zero4;
     *reinterpret_cast<uint4*>(dst + 8) = zero4;
@@ -245,7 +295,9 @@ __device__ __forceinline__ void processDeepseekV4Slot(
 
   // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
   // All math in fp32.  cos_sin_cache is loaded as fp32 (its native storage).
-  bool const is_rope_lane = dim_base >= kNopeDim;
+  // KV is always rotated; Q only when the caller has not already had it
+  // rotated elsewhere (the mega-attention kernel does its own).
+  bool const is_rope_lane = dim_base >= kNopeDim && (isKV || apply_q_rope);
   if (is_rope_lane) {
     int64_t const pos = position_ids[tokenIdx];
     constexpr int kHalfRope = kRopeDim / 2;
@@ -295,26 +347,29 @@ __device__ __forceinline__ void processDeepseekV4Slot(
           make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
     }
     scalar_t_in* dst =
-        q_out +
-        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
-            kHeadDim +
-        dim_base;
+        q_out + qElemOffset<Q_INTERLEAVED>(tokenIdx, slotIdx, laneId, dim_base,
+                                           kNumHeadsQPadded);
     *reinterpret_cast<uint4*>(dst) = out0;
     *reinterpret_cast<uint4*>(dst + 8) = out1;
   } else {
-    // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert.
+    // ── KV: FP8 quant + cache insert.  V4 keeps the RoPE dims in bf16;
+    // V4.1 (kv_mxfp8) quantizes them too.
     int64_t const slot_id = slot_mapping[tokenIdx];
     if (slot_id >= 0) {
+      int const token_data_bytes =
+          kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes;
+      int const scale_bytes =
+          kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken;
       int64_t const block_idx = slot_id / cache_block_size;
       int64_t const pos_in_block = slot_id % cache_block_size;
       uint8_t* block_base =
           k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
-      uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
+      uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes;
       uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
       uint8_t* token_scale_ptr =
           block_base +
-          static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
-          pos_in_block * kScaleBytesPerToken;
+          static_cast<int64_t>(cache_block_size) * token_data_bytes +
+          pos_in_block * scale_bytes;
 
 #pragma unroll
       for (int i = 0; i < kElemsPerLane; i++) {
@@ -326,11 +381,13 @@ __device__ __forceinline__ void processDeepseekV4Slot(
       for (int i = 0; i < kElemsPerLane; i++) {
         local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
       }
-      float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
+      float const absmax = fmaxf(
+          kv_mxfp8 ? warp2MaxAbs(local_absmax) : warp4MaxAbs(local_absmax),
+          1e-4f);
       float const exponent = ceilf(log2f(absmax / kFp8Max));
       float const inv_scale = exp2f(-exponent);
 
-      if (!is_rope_lane) {
+      if (kv_mxfp8 || !is_rope_lane) {
         uint8_t out_bytes[kElemsPerLane];
 #pragma unroll
         for (int i = 0; i < kElemsPerLane; i++) {
@@ -347,13 +404,19 @@ __device__ __forceinline__ void processDeepseekV4Slot(
         *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
             *reinterpret_cast<uint4 const*>(out_bytes);
 
-        if ((laneId & 3) == 0) {
-          int const q_block_idx = laneId >> 2;
-          float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
-          token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
-        }
-        if (laneId == 0) {
-          token_scale_ptr[kNumQuantBlocks] = 0;
+        float const encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
+        if (kv_mxfp8) {
+          // One scale per 32 dims = one per lane pair.
+          if ((laneId & 1) == 0) {
+            token_scale_ptr[laneId >> 1] = static_cast<uint8_t>(encoded);
+          }
+        } else {
+          if ((laneId & 3) == 0) {
+            token_scale_ptr[laneId >> 2] = static_cast<uint8_t>(encoded);
+          }
+          if (laneId == 0) {
+            token_scale_ptr[kNumQuantBlocks] = 0;
+          }
         }
       } else {
         uint4 out0, out1;
@@ -400,14 +463,17 @@ __device__ __forceinline__ void processDeepseekV4Slot(
 // `kNumHeadsQPadded` is a template parameter (compile-time constant) so the
 // divisions in the grid math and the KV-sentinel comparison fold to fast
 // constant operations.  The launch wrapper dispatches the runtime value to
-// the matching instantiation.
+// the matching instantiation.  `kNumHeadsQPadded == 0` leaves one slot per
+// token and is the KV-insert-only instantiation, for a Q that is already in
+// the shape its attention kernel wants.
 //
 // With DP padding, q/kv/position_ids can have more rows than slot_mapping.
 // The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream
 // attention uses them).  The KV branch only inserts the first
 // `num_tokens_insert` tokens (= slot_mapping length) into the paged cache.
 //
-template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM,
+          bool Q_INTERLEAVED>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     scalar_t_in const* __restrict__ q_in,      // [N, num_heads_q,      512]
     scalar_t_in* __restrict__ q_out,           // [N, kNumHeadsQPadded, 512]
@@ -421,7 +487,9 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     int const num_tokens_insert,  // = slot_mapping.size(0), ≤ num_tokens_full
     int const num_heads_q,        // live Q heads (input layout)
     int const cache_block_size,   // tokens per paged-cache block
-    int const kv_block_stride) {  // bytes per paged-cache block
+    int const kv_block_stride,    // bytes per paged-cache block
+    bool const kv_mxfp8,          // V4.1 all-dims MXFP8 record
+    bool const apply_q_rope) {    // rotate Q too, not just KV
 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
   // BF16 _typeConvert specialization is unavailable on pre-Ampere.  The
   // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
@@ -464,20 +532,18 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
       if (isKV) {
         src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
       } else {
-        int64_t const q_row_offset =
-            (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
-                kHeadDim +
-            dim_base;
-        src_ptr = q_in + q_row_offset;
+        src_ptr = q_in + qElemOffset<Q_INTERLEAVED>(tokenIdx, slotIdx, laneId,
+                                                    dim_base, num_heads_q);
       }
       v0 = *reinterpret_cast<uint4 const*>(src_ptr);
       v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
     }
 
-    processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
+    processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM,
+                          Q_INTERLEAVED>(
         v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
         k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
-        kv_block_stride);
+        kv_block_stride, kv_mxfp8, apply_q_rope);
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
     cudaTriggerProgrammaticLaunchCompletion();
@@ -497,7 +563,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
 // Q branch (optional RMSNorm + RoPE, in place) head_slot == num_heads_q
 // KV branch (RoPE + UE8M0 quant + insert)
 //
-template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM,
+          bool Q_INTERLEAVED>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
     scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
     scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
@@ -506,7 +573,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
     float const* __restrict__ cos_sin_cache, float const eps,
     int const num_tokens_full, int const num_tokens_insert,
     int const num_heads_q, int const cache_block_size,
-    int const kv_block_stride) {
+    int const kv_block_stride, bool const kv_mxfp8, bool const apply_q_rope) {
 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
   if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
     return;
@@ -536,11 +603,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
       if (s == kNumHeadsQPadded) {
         src = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
       } else {
-        src = q_in +
-              (static_cast<int64_t>(tokenIdx) * num_heads_q +
-               static_cast<int64_t>(s)) *
-                  kHeadDim +
-              dim_base;
+        src = q_in + qElemOffset<Q_INTERLEAVED>(tokenIdx, s, laneId, dim_base,
+                                                num_heads_q);
       }
       va = *reinterpret_cast<uint4 const*>(src);
       vb = *reinterpret_cast<uint4 const*>(src + 8);
@@ -561,10 +625,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
           load_slot(next_slot, v0_next, v1_next);
         }
 
-        processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
+        processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM,
+                              Q_INTERLEAVED>(
             v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
             num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
-            cos_sin_cache, cache_block_size, kv_block_stride);
+            cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8,
+            apply_q_rope);
 
         // ── Buffer rotation: hand the prefetched LDGs to the next iter.
         v0_curr = v0_next;
@@ -584,14 +650,15 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 // ────────────────────────────────────────────────────────────────────────────
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
+template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM,
+          bool Q_INTERLEAVED>
 static void launchFusedDeepseekV4Templated(
     scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
     uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
     float const* cos_sin_cache, float const eps, int const num_tokens_full,
     int const num_tokens_insert, int const num_heads_q,
-    int const cache_block_size, int const kv_block_stride,
-    cudaStream_t stream) {
+    int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8,
+    bool const apply_q_rope, cudaStream_t stream) {
   constexpr int kBlockSize = 256;
   constexpr int kWarpsPerBlock = kBlockSize / 32;
   int64_t const total_warps =
@@ -624,37 +691,46 @@ static void launchFusedDeepseekV4Templated(
   config.attrs = attrs;
   config.numAttrs = (sm_version >= 90) ? 1 : 0;
 
-  if (num_tokens_full < NUM_TOKEN_CUTOFF) {
-    cudaLaunchKernelEx(&config,
-                       fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
-                           scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
-                       q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
-                       cos_sin_cache, eps, num_tokens_full, num_tokens_insert,
-                       num_heads_q, cache_block_size, kv_block_stride);
+  // The reduced grid pays off by iterating a token's Q heads inside one CTA.
+  // With zero query heads there is nothing to iterate, so one warp per block
+  // would do all the work and the other seven idle; keep the warp-per-slot
+  // grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single
+  // insert, where it is worth ~2x.
+  if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) {
+    cudaLaunchKernelEx(
+        &config,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
+            scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        kv_block_stride, kv_mxfp8, apply_q_rope);
   } else {
     config.gridDim = dim3(num_tokens_full);
     cudaLaunchKernelEx(
         &config,
         fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<
-            scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
+            scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>,
         q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
         eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
-        kv_block_stride);
+        kv_block_stride, kv_mxfp8, apply_q_rope);
   }
 #else
   // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
   // clang-format off
   fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
-      scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>
+      scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>
       <<<grid, kBlockSize, 0, stream>>>(
           q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
           cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
-          cache_block_size, kv_block_stride);
+          cache_block_size, kv_block_stride, kv_mxfp8, apply_q_rope);
 #endif
 }
 
 // Runtime dispatch into one of the precompiled `kNumHeadsQPadded`
-// instantiations.  Supported padded head counts: 8, 16, 32, 64, 128.
+// instantiations.  Supported padded head counts: 8, 16, 32, 64, 128, plus 0
+// for a KV-only launch that writes no Q at all.  `apply_q_norm` and
+// `is_q_interleaved` are template parameters, `apply_q_rope` is not (see
+// processDeepseekV4Slot), so each head count carries four instantiations.
 template <typename scalar_t_in>
 void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
     scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
@@ -663,22 +739,32 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
     int const num_tokens_full, int const num_tokens_insert,
     int const num_heads_q, int const num_heads_q_padded,
     int const cache_block_size, int const kv_block_stride,
-    bool const apply_q_norm, cudaStream_t stream) {
-#define DISPATCH(N)                                                         \
-  case N:                                                                   \
-    if (apply_q_norm) {                                                     \
-      launchFusedDeepseekV4Templated<scalar_t_in, N, true>(                 \
-          q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,          \
-          cos_sin_cache, eps, num_tokens_full, num_tokens_insert,            \
-          num_heads_q, cache_block_size, kv_block_stride, stream);           \
-    } else {                                                                \
-      launchFusedDeepseekV4Templated<scalar_t_in, N, false>(                \
-          q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,          \
-          cos_sin_cache, eps, num_tokens_full, num_tokens_insert,            \
-          num_heads_q, cache_block_size, kv_block_stride, stream);           \
-    }                                                                       \
+    bool const apply_q_norm, bool const kv_mxfp8, bool const apply_q_rope,
+    bool const is_q_interleaved, cudaStream_t stream) {
+#define LAUNCH(N, NORM, INTERLEAVED)                                         \
+  launchFusedDeepseekV4Templated<scalar_t_in, N, NORM, INTERLEAVED>(         \
+      q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, \
+      eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, \
+      kv_block_stride, kv_mxfp8, apply_q_rope, stream)
+#define DISPATCH(N)                  \
+  case N:                            \
+    if (is_q_interleaved) {          \
+      if (apply_q_norm) {            \
+        LAUNCH(N, true, true);       \
+      } else {                       \
+        LAUNCH(N, false, true);      \
+      }                              \
+    } else if (apply_q_norm) {       \
+      LAUNCH(N, true, false);        \
+    } else {                         \
+      LAUNCH(N, false, false);       \
+    }                                \
     return;
 
+  if (num_heads_q_padded == 0) {
+    LAUNCH(0, false, false);
+    return;
+  }
   switch (num_heads_q_padded) {
     DISPATCH(8)
     DISPATCH(16)
@@ -690,9 +776,10 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
                   "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
                   "unsupported num_heads_q_padded=",
                   num_heads_q_padded,
-                  " (compiled instantiations: 8, 16, 32, 64, 128).");
+                  " (compiled instantiations: 0, 8, 16, 32, 64, 128).");
   }
 #undef DISPATCH
+#undef LAUNCH
 }
 
 // ────────────────────────────────────────────────────────────────────────────
@@ -959,6 +1046,91 @@ static void launchFullCacheKernel(
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
+void fused_deepseek_v4_kv_rope_insert(
+    torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping,
+    torch::stable::Tensor const& position_ids,
+    torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> fp8_scale, bool kv_mxfp8) {
+  using torch::headeronly::ScalarType;
+  STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous() &&
+                      kv.dim() == 2 && kv.size(1) == 512 &&
+                      kv.scalar_type() == ScalarType::BFloat16,
+                  "kv must be contiguous CUDA bfloat16 [tokens, 512]");
+  auto const device = kv.get_device_index();
+  STD_TORCH_CHECK(k_cache.device().is_cuda() &&
+                      k_cache.get_device_index() == device &&
+                      k_cache.dim() == 3 && cache_block_size > 0 &&
+                      k_cache.size(1) == cache_block_size &&
+                      k_cache.stride(2) == 1,
+                  "k_cache must be a paged cache on the KV device");
+  for (auto const* tensor : {&slot_mapping, &position_ids}) {
+    STD_TORCH_CHECK(tensor->device().is_cuda() &&
+                        tensor->get_device_index() == device &&
+                        tensor->scalar_type() == ScalarType::Long &&
+                        tensor->dim() == 1 && tensor->is_contiguous(),
+                    "slots and positions must be contiguous int64 on the KV device");
+  }
+  STD_TORCH_CHECK(position_ids.size(0) == kv.size(0) &&
+                      slot_mapping.size(0) <= kv.size(0),
+                  "positions must match KV rows; slots may omit padded rows");
+  STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
+                      cos_sin_cache.get_device_index() == device &&
+                      cos_sin_cache.scalar_type() == ScalarType::Float &&
+                      cos_sin_cache.dim() == 2 &&
+                      cos_sin_cache.size(1) == 64 && cos_sin_cache.is_contiguous(),
+                  "cos_sin_cache must be contiguous float32 [max_pos, 64]");
+  bool const packed = k_cache.scalar_type() == ScalarType::Byte;
+  bool const fp8 = k_cache.scalar_type() == ScalarType::Float8_e4m3fn;
+  STD_TORCH_CHECK(packed || fp8 ||
+                      k_cache.scalar_type() == ScalarType::BFloat16,
+                  "cache must be uint8, bfloat16 or float8_e4m3fn");
+  STD_TORCH_CHECK(k_cache.size(2) == (packed ? (kv_mxfp8 ? 528 : 584) : 512),
+                  "cache row width does not match its dtype");
+  if (fp8) {
+    STD_TORCH_CHECK(fp8_scale.has_value() && fp8_scale->device().is_cuda() &&
+                        fp8_scale->get_device_index() == device &&
+                        fp8_scale->scalar_type() == ScalarType::Float &&
+                        fp8_scale->numel() == 1,
+                    "fp8 cache requires one float32 KV scale on the KV device");
+  }
+  int const num_tokens = static_cast<int>(slot_mapping.size(0));
+  if (num_tokens == 0) return;
+  const torch::stable::accelerator::DeviceGuard device_guard(device);
+  const cudaStream_t stream = get_current_cuda_stream(device);
+  using scalar_t = c10::BFloat16;
+  auto const* input = reinterpret_cast<scalar_t const*>(kv.const_data_ptr());
+  auto* cache = reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr());
+  auto const* slots = slot_mapping.const_data_ptr<int64_t>();
+  auto const* positions = position_ids.const_data_ptr<int64_t>();
+  auto const* cos_sin = cos_sin_cache.const_data_ptr<float>();
+  // Zero query heads schedules only the existing KV branch, preserving its
+  // RoPE rounding and quantization without allocating any query tensors.
+  if (packed) {
+    // With zero query heads no Q slot is scheduled, so neither Q knob is
+    // reachable: Q_INTERLEAVED picks addressing nothing uses, and apply_q_rope
+    // gates a branch only a Q slot takes (KV is rotated unconditionally).
+    vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated<scalar_t, 0,
+                                                                false, false>(
+        nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f,
+        num_tokens, num_tokens, 0, static_cast<int>(cache_block_size),
+        static_cast<int>(k_cache.stride(0)), kv_mxfp8, /*apply_q_rope=*/true,
+        stream);
+  } else if (fp8) {
+    vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false, true>(
+        nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin,
+        fp8_scale->const_data_ptr<float>(), nullptr, 0.0f, num_tokens, num_tokens,
+        0, static_cast<int>(cache_block_size), k_cache.stride(0),
+        k_cache.stride(1), false, "fused_deepseek_v4_kv_rope_insert", stream);
+  } else {
+    vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false, false>(
+        nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, nullptr,
+        nullptr, 0.0f, num_tokens, num_tokens, 0, static_cast<int>(cache_block_size),
+        k_cache.stride(0) * 2, k_cache.stride(1) * 2, false,
+        "fused_deepseek_v4_kv_rope_insert", stream);
+  }
+}
+
 torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::stable::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
     torch::stable::Tensor const& kv,             // [N, 512] bf16 (read-only)
@@ -966,8 +1138,9 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::stable::Tensor const& slot_mapping,   // [N] int64
     torch::stable::Tensor const& position_ids,   // [N] int64
     torch::stable::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
-    int64_t q_head_padded,                       // padded Q head count for output
-    double eps, int64_t cache_block_size, bool apply_q_norm) {
+    int64_t q_head_padded,  // padded Q head count for output, 0 for no Q
+    double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8,
+    bool apply_q_rope, bool is_q_interleaved) {
   STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
                   "q_in must be contiguous CUDA");
   STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
@@ -987,8 +1160,11 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
   STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
                   "q_in and kv dtype must match");
-  STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
-                  "q_head_padded must be >= q_in.size(1) (num_heads_q)");
+  // q_head_padded == 0 asks for the KV insert alone: Q is left untouched and
+  // an empty tensor comes back, for a caller whose Q is already in the shape
+  // its attention kernel reads.
+  STD_TORCH_CHECK(q_head_padded == 0 || q_head_padded >= q_in.size(1),
+                  "q_head_padded must be 0 or >= q_in.size(1) (num_heads_q)");
   STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
                   "k_cache must be uint8");
   STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
@@ -1019,8 +1195,12 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   // Allocate the padded q output.  The kernel writes every element (live
   // region gets optional RMSNorm+RoPE; pad region gets zeros), so `empty` is
   // safe.
-  auto q_out = torch::stable::new_empty(
-      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
+  auto q_out =
+      q_head_padded == 0
+          ? torch::stable::new_empty(q_in, {int64_t{0}}, q_in.scalar_type())
+          : torch::stable::new_empty(
+                q_in, {q_in.size(0), q_head_padded, q_in.size(2)},
+                q_in.scalar_type());
 
   VLLM_STABLE_DISPATCH_HALF_TYPES(
       q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
@@ -1036,7 +1216,8 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                 cos_sin_cache.const_data_ptr<float>(), static_cast<float>(eps),
                 num_tokens_full, num_tokens_insert, num_heads_q,
                 num_heads_q_padded, cache_block_size_i, kv_block_stride,
-                apply_q_norm, stream);
+                apply_q_norm, kv_mxfp8, apply_q_rope, is_q_interleaved,
+                stream);
       });
   return q_out;
 }
diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
index 05c7bfc6d13e..3892790c4040 100644
--- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
+++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
@@ -35,8 +35,9 @@
  *
  *   fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert):
  *     - full key concat (bf16), and cache insert in DeepSeek's 656-byte
- *       block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic
- *       scales, RoPE bf16), bit-compatible with concat_and_cache_ds_mla_kernel.
+ *       block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile
+ *       power-of-two scales, RoPE bf16), bit-compatible with
+ *       concat_and_cache_ds_mla_kernel.
  *
  * Both use Programmatic Dependent Launch (PDL) to overlap the tail of the
  * producing GEMMs on sm_90+, and are structured after
@@ -68,7 +69,6 @@
   #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
 #endif
 #include <cuda_runtime.h>
-#include <cfloat>
 #include <type_traits>
 
 #ifdef USE_ROCM
@@ -111,9 +111,9 @@ constexpr float kFp8Max = 224.0f;
 #else
 constexpr float kFp8Max = 448.0f;
 #endif
-// Divisor for fp8_ds_mla per-tile dynamic scales (matches cache_kernels.cu).
-// fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), [512,528) 4 fp32
-// tile scales, [528,656) RoPE 64 bf16.
+// Divisor for fp8_ds_mla per-tile power-of-two scales (matches
+// cache_kernels.cu). fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128),
+// [512,528) 4 fp32 tile scales, [528,656) RoPE 64 bf16.
 constexpr float kFp8ScaleDivisor = kFp8Max;
 
 // Copy 8 source elements (one uint4 of bf16/fp16) to `dst`. FP8=false stores a
@@ -325,7 +325,7 @@ __device__ __forceinline__ void writeLatent576(
 }
 
 // Write [kv_c | k_pe] into the fp8_ds_mla 656B entry using one warp: NoPE 512
-// as fp8 in 4 tiles of 128 (per-tile dynamic absmax scale, 4 fp32 scales at
+// as fp8 in 4 tiles of 128 (per-tile power-of-two scale, 4 fp32 scales at
 // [512,528)), RoPE 64 as bf16 at [528,656). Bit-compatible with
 // concat_and_cache_ds_mla_kernel.
 template <typename scalar_t, bool APPLY_ROPE = false>
@@ -349,7 +349,8 @@ __device__ __forceinline__ void writeDsMlaCache(
   for (int offset = 4; offset > 0; offset /= 2) {
     max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 8));
   }
-  float const tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN);
+  float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f);
+  tile_scale = exp2f(ceilf(log2f(tile_scale)));
   if ((laneId & 7) == 0) {
     reinterpret_cast<float*>(row)[kKvLoraRank / 4 + tile] = tile_scale;
   }
@@ -598,7 +599,7 @@ __global__ void fusedKimiK3MLAKVConcatPackKernel(
 // ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert
 //
 // Cache entry (656 bytes), matching concat_and_cache_ds_mla_kernel:
-//   [0, 512)   NoPE 512 vals as fp8, 4 tiles of 128, each dynamically scaled
+//   [0, 512)   NoPE 512 vals as fp8, 4 tiles of 128, power-of-two scaled
 //   [512, 528) 4 fp32 per-tile scales
 //   [528, 656) RoPE 64 vals as bf16 (unquantized)
 // The cache slot uses one warp: lane L quantizes NoPE elems [L*16, L*16+16)
diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/libtorch_stable/hisparse_kernels.cu
new file mode 100644
index 000000000000..97ce9246cc75
--- /dev/null
+++ b/csrc/libtorch_stable/hisparse_kernels.cu
@@ -0,0 +1,1244 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// HiSparse decode hot-buffer kernels for sparse MLA (GLM-5 / DeepSeek-V3.2).
+//
+// The swap-in algorithm is a port of SGLang's hisparse
+// load_cache_to_device_buffer kernel (sgl jit_kernel/csrc/hisparse.cuh),
+// adapted to vLLM addressing:
+//  - tokens are keyed by their global KV slot id (block_table-converted
+//    indexer output) instead of in-request positions, so no per-request
+//    host-location table is needed: host pool row i is global slot i.
+//  - each batch row owns a fixed region of `region_stride` hot rows;
+//    slots [0, hot_size) are LRU-managed, slot `hot_size` holds the row's
+//    newest token (written directly by the KV-cache update).
+//
+// The kernels only move bytes; they are dtype agnostic.
+
+#include "torch_utils.h"
+#include "ops.h"
+#include "../cuda_utils.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <optional>
+
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/headeronly/version.h>
+
+namespace {
+
+constexpr int kWarpSize = 32;
+// Sentinel in the shared top-k scratch: entry already resolved (hit /
+// newest / invalid), no miss handling needed.
+constexpr int32_t kTokenDone = -1;
+constexpr int32_t kHashEmpty = -1;
+
+bool is_pinned_cpu_tensor(const torch::stable::Tensor& tensor) {
+  cudaPointerAttributes attributes{};
+  const auto status =
+      cudaPointerGetAttributes(&attributes, tensor.const_data_ptr());
+  if (status != cudaSuccess) {
+    cudaGetLastError();
+    return false;
+  }
+  return attributes.type == cudaMemoryTypeHost;
+}
+
+__device__ __forceinline__ int32_t hash_slot(int32_t key, int32_t hash_size) {
+  // Knuth multiplicative hash for the open-addressing table.
+  return static_cast<int32_t>((static_cast<uint32_t>(key) * 2654435761u) %
+                              static_cast<uint32_t>(hash_size));
+}
+
+// Copy one row of `row_bytes` bytes with a single warp. MLA rows take the
+// 16-byte vectorized path; packed indexer rows use the scalar tail-safe path.
+__device__ __forceinline__ void copy_row_warp(int lane_id, const char* src,
+                                              char* dst, int64_t row_bytes) {
+  const auto alignment = reinterpret_cast<uintptr_t>(src) |
+                         reinterpret_cast<uintptr_t>(dst) |
+                         static_cast<uintptr_t>(row_bytes);
+  if ((alignment & 15) == 0) {
+    const int64_t num_vec = row_bytes / 16;
+    const uint4* src4 = reinterpret_cast<const uint4*>(src);
+    uint4* dst4 = reinterpret_cast<uint4*>(dst);
+    for (int64_t j = lane_id; j < num_vec; j += kWarpSize) {
+      __stcg(dst4 + j, __ldcg(src4 + j));
+    }
+    return;
+  }
+
+  if ((alignment & 3) == 0) {
+    const int64_t num_words = row_bytes / 4;
+    const unsigned int* src_words = reinterpret_cast<const unsigned int*>(src);
+    unsigned int* dst_words = reinterpret_cast<unsigned int*>(dst);
+    for (int64_t j = lane_id; j < num_words; j += kWarpSize) {
+      __stcg(dst_words + j, __ldcg(src_words + j));
+    }
+    return;
+  }
+
+  for (int64_t j = lane_id; j < row_bytes; j += kWarpSize) {
+    __stcg(dst + j, __ldcg(src + j));
+  }
+}
+
+// Zero one row with a single warp (16B vectorized, L2-only stores), same
+// addressing contract as copy_row_warp.
+__device__ __forceinline__ void zero_row_warp(int lane_id, char* dst,
+                                              int64_t row_bytes) {
+  const auto alignment =
+      reinterpret_cast<uintptr_t>(dst) | static_cast<uintptr_t>(row_bytes);
+  if ((alignment & 15) == 0) {
+    const int64_t num_vec = row_bytes / 16;
+    uint64_t* dst8 = reinterpret_cast<uint64_t*>(dst);
+    for (int64_t j = lane_id; j < num_vec; j += kWarpSize) {
+      uint64_t* d = dst8 + j * 2;
+      asm volatile("st.global.cg.v2.b64 [%0],{%1,%2};" ::"l"(d), "l"(0ULL),
+                   "l"(0ULL)
+                   : "memory");
+    }
+    return;
+  }
+
+  if ((alignment & 3) == 0) {
+    const int64_t num_words = row_bytes / 4;
+    unsigned int* dst_words = reinterpret_cast<unsigned int*>(dst);
+    for (int64_t j = lane_id; j < num_words; j += kWarpSize) {
+      __stcg(dst_words + j, 0u);
+    }
+    return;
+  }
+
+  for (int64_t j = lane_id; j < row_bytes; j += kWarpSize) {
+    __stcg(dst + j, static_cast<char>(0));
+  }
+}
+
+__device__ __forceinline__ char* cache_row_ptr(char* cache, int64_t row,
+                                               int32_t block_size,
+                                               int64_t block_stride,
+                                               int64_t row_bytes) {
+  return cache + (row / block_size) * block_stride +
+         (row % block_size) * row_bytes;
+}
+
+__device__ __forceinline__ void zero_cache_row_warp(int lane_id, char* cache,
+                                                    int64_t row,
+                                                    int32_t block_size,
+                                                    int64_t block_stride,
+                                                    int64_t row_bytes) {
+  zero_row_warp(lane_id,
+                cache_row_ptr(cache, row, block_size, block_stride, row_bytes),
+                row_bytes);
+}
+
+// In-place inclusive scan over s_data[offset, count) performed by warp 0,
+// carrying `accumulator` across calls. Returns the running total.
+__device__ __forceinline__ int warp_inclusive_scan(int32_t* s_data, int lane_id,
+                                                   int offset, int count,
+                                                   int accumulator) {
+  int idx = lane_id + offset;
+  int val = (idx < count) ? s_data[idx] : 0;
+#pragma unroll
+  for (int i = 1; i < 32; i *= 2) {
+    int n = __shfl_up_sync(0xffffffff, val, i);
+    if (lane_id >= i) val += n;
+  }
+  val += accumulator;
+  if (idx < count) {
+    s_data[idx] = val;
+  }
+  return __shfl_sync(0xffffffff, val, 31);
+}
+
+__device__ __forceinline__ int64_t
+get_physical_hot_row(const int32_t* hot_block_table, int32_t row,
+                     int64_t table_stride, int32_t block_size, int32_t slot) {
+  const int32_t block =
+      hot_block_table[static_cast<int64_t>(row) * table_stride +
+                      slot / block_size];
+  return static_cast<int64_t>(block) * block_size + slot % block_size;
+}
+
+__device__ __forceinline__ void store_hot_index(
+    int32_t* hot_indices, int32_t* attention_indices, int32_t index,
+    int32_t physical_row, int32_t hot_block_size,
+    int64_t attention_block_stride) {
+  hot_indices[index] = physical_row;
+  if (attention_indices != nullptr) {
+    attention_indices[index] =
+        physical_row < 0
+            ? -1
+            : (physical_row / hot_block_size) * attention_block_stride +
+                  physical_row % hot_block_size;
+  }
+}
+
+// One block per batch row.
+//
+// Shared memory layout (int32 region followed by int16 region):
+//   s_topk[top_k]            top-k global ids; reused as miss scratch
+//   s_chunk_off[nbc + 1]     prefix sums for hit (then miss) compaction
+//   s_evict_off[nbc + 1]     prefix sums for evictable compaction
+//   s_hash_keys[hash_size]   open addressing: global id -> top-k index
+//   s_counters[3]            hits, phase-1 resolved, valid count
+//   s_lru_out[hot_size]      int16, compacted slots: [hits fwd | evict bwd]
+//   s_hash_vals[hash_size]   int16 hash values (top-k index)
+// Valid global ids must be unique within each row.
+__global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
+    const int32_t* __restrict__ hot_block_table,  // [max_rows, hot_blocks]
+    const int32_t* __restrict__ global_indices,   // global or request-relative
+    const int32_t* __restrict__ request_ids,      // [num_rows] or nullptr
+    const int32_t* __restrict__ source_block_table,  // [num_reqs, max_blocks]
+    const int32_t* __restrict__ resident_block_table,
+    int32_t* __restrict__ resolved_global_indices,    // [num_rows, top_k]
+    int32_t* __restrict__ valid_counts,               // [num_rows] or nullptr
+    int32_t* __restrict__ swap_host_physical_rows,    // [num_rows, top_k]
+    int32_t* __restrict__ swap_device_physical_rows,  // [num_rows, top_k]
+    int32_t* __restrict__ swap_counts,                // [num_rows]
+    int32_t* __restrict__ hot_indices,                // [num_rows, top_k]
+    int32_t* __restrict__ attention_indices,          // [num_rows, top_k]
+    int32_t* __restrict__ miss_mask,  // [num_rows, top_k] or nullptr
+    int32_t* __restrict__ device_global_indices,  // [max_rows, region_stride]
+    int16_t* __restrict__ lru_slots,              // [max_rows, hot_size]
+    unsigned long long* __restrict__ stats,       // [2] hits,misses or nullptr
+    const int32_t* __restrict__ request_state_indices,  // [num_requests] or
+                                                        // nullptr
+    const int32_t request_state_count, const int64_t host_rows,
+    const int64_t hot_table_stride, const int32_t hot_block_size,
+    const int32_t top_k, const int32_t hot_size, const int32_t hash_size,
+    const int64_t region_stride, const int64_t attention_block_stride,
+    const int64_t source_bt_stride, const int32_t source_num_reqs,
+    const int32_t source_num_blocks, const int32_t source_block_size,
+    const int64_t resident_bt_stride, const int32_t resident_num_reqs,
+    const int32_t resident_num_blocks, const int32_t resident_block_size,
+    const int32_t resident_null_block, const int64_t input_row_stride,
+    const int64_t attention_row_stride, const int64_t valid_count_stride) {
+  const int NUM_WARPS = blockDim.x / kWarpSize;
+  const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
+  const int num_token_chunks = (top_k + kWarpSize - 1) / kWarpSize;
+
+  const int batch_row = blockIdx.x;
+  const int request_row =
+      request_ids != nullptr ? request_ids[batch_row] : batch_row;
+  const int state_row =
+      request_state_indices != nullptr && request_row >= 0 &&
+              request_row < request_state_count
+          ? request_state_indices[request_row]
+          : (request_state_indices == nullptr ? request_row : -1);
+  // V2 publishes -1 for CUDA-graph padding rows.
+  if (state_row < 0) {
+    for (int i = threadIdx.x; i < top_k; i += blockDim.x) {
+      const int64_t index = static_cast<int64_t>(batch_row) * top_k + i;
+      hot_indices[index] = -1;
+      if (attention_indices != nullptr) {
+        attention_indices[static_cast<int64_t>(batch_row) *
+                              attention_row_stride +
+                          i] = -1;
+      }
+      if (resolved_global_indices != nullptr) {
+        resolved_global_indices[index] = -1;
+      }
+    }
+    if (valid_counts != nullptr && threadIdx.x == 0) {
+      valid_counts[static_cast<int64_t>(batch_row) * valid_count_stride] = 0;
+    }
+    if (swap_counts != nullptr && threadIdx.x == 0) {
+      swap_counts[batch_row] = 0;
+    }
+    return;
+  }
+  const int tid = threadIdx.x;
+  const int warp_id = tid / kWarpSize;
+  const int lane_id = tid % kWarpSize;
+  const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1;
+
+  const int32_t* row_topk =
+      global_indices + static_cast<int64_t>(batch_row) * input_row_stride;
+  int32_t* row_out = hot_indices + static_cast<int64_t>(batch_row) * top_k;
+  int32_t* row_attention =
+      attention_indices != nullptr
+          ? attention_indices +
+                static_cast<int64_t>(batch_row) * attention_row_stride
+          : nullptr;
+  int32_t* row_miss = (miss_mask != nullptr)
+                          ? miss_mask + static_cast<int64_t>(batch_row) * top_k
+                          : nullptr;
+  int32_t* row_dgi =
+      device_global_indices + static_cast<int64_t>(state_row) * region_stride;
+  int16_t* row_lru = lru_slots + static_cast<int64_t>(state_row) * hot_size;
+
+  extern __shared__ char smem_raw[];
+  int32_t* s_topk = reinterpret_cast<int32_t*>(smem_raw);
+  int32_t* s_chunk_off = s_topk + top_k;
+  int32_t* s_evict_off = s_chunk_off + (num_buffer_chunks + 1);
+  int32_t* s_hash_keys = s_evict_off + (num_buffer_chunks + 1);
+  int32_t* s_counters = s_hash_keys + hash_size;
+  int16_t* s_lru_out = reinterpret_cast<int16_t*>(s_counters + 3);
+  int16_t* s_hash_vals = s_lru_out + hot_size;
+
+  if (tid < 3) {
+    s_counters[tid] = 0;
+  }
+  __syncthreads();
+
+  // Phase 1: translate request-relative positions and resolve resident rows.
+  for (int i = tid; i < top_k; i += blockDim.x) {
+    const int32_t token_index = row_topk[i];
+    int32_t g = token_index;
+    int32_t resident_row = -1;
+    if (source_block_table != nullptr) {
+      const int32_t request_id = request_ids[batch_row];
+      const int32_t source_block =
+          token_index >= 0 ? token_index / source_block_size : -1;
+      if (request_id >= 0 && request_id < source_num_reqs &&
+          source_block >= 0 && source_block < source_num_blocks) {
+        const int32_t physical_block =
+            source_block_table[static_cast<int64_t>(request_id) *
+                                   source_bt_stride +
+                               source_block];
+        g = physical_block > 0 ? physical_block * source_block_size +
+                                     token_index % source_block_size
+                               : -1;
+      } else {
+        g = -1;
+      }
+      if (resident_block_table != nullptr) {
+        const int32_t resident_block =
+            token_index >= 0 ? token_index / resident_block_size : -1;
+        if (request_id >= 0 && request_id < resident_num_reqs &&
+            resident_block >= 0 && resident_block < resident_num_blocks) {
+          const int32_t physical_block =
+              resident_block_table[static_cast<int64_t>(request_id) *
+                                       resident_bt_stride +
+                                   resident_block];
+          if (physical_block != resident_null_block && physical_block >= 0) {
+            resident_row = physical_block * resident_block_size +
+                           token_index % resident_block_size;
+          }
+        }
+      }
+    }
+    if (g >= host_rows) {
+      g = -1;
+    }
+    if (resolved_global_indices != nullptr) {
+      resolved_global_indices[static_cast<int64_t>(batch_row) * top_k + i] = g;
+    }
+    if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1);
+    if (row_miss != nullptr) row_miss[i] = 0;
+    if (resident_row >= 0) {
+      store_hot_index(row_out, row_attention, i, resident_row, hot_block_size,
+                      attention_block_stride);
+      s_topk[i] = kTokenDone;
+      atomicAdd(&s_counters[1], 1);
+    } else if (g < 0) {
+      store_hot_index(row_out, row_attention, i, -1, hot_block_size,
+                      attention_block_stride);
+      s_topk[i] = kTokenDone;
+      atomicAdd(&s_counters[1], 1);
+    } else {
+      s_topk[i] = g;
+    }
+  }
+  __syncthreads();
+  if (valid_counts != nullptr && tid == 0) {
+    valid_counts[static_cast<int64_t>(batch_row) * valid_count_stride] =
+        s_counters[2];
+  }
+  // Fully resident rows need only request-relative page translation. Avoid
+  // scanning or rewriting the hot LRU when no selected row can consult it.
+  if (s_counters[1] == top_k) {
+    if (tid == 0 && swap_counts != nullptr) {
+      swap_counts[batch_row] = 0;
+    }
+    return;
+  }
+
+  for (int i = tid; i < hash_size; i += blockDim.x) {
+    s_hash_keys[i] = kHashEmpty;
+  }
+  for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) {
+    s_chunk_off[i] = 0;
+    s_evict_off[i] = 0;
+  }
+  __syncthreads();
+
+  for (int i = tid; i < top_k; i += blockDim.x) {
+    const int32_t g = s_topk[i];
+    if (g == kTokenDone) {
+      continue;
+    }
+    int slot = hash_slot(g, hash_size);
+    while (true) {
+      const int32_t old = atomicCAS(&s_hash_keys[slot], kHashEmpty, g);
+      if (old == kHashEmpty || old == g) {
+        s_hash_vals[slot] = static_cast<int16_t>(i);
+        break;
+      }
+      slot = (slot + 1) % hash_size;
+    }
+  }
+  __syncthreads();
+
+  // Phase 2: walk hot slots in LRU order, classify hit / evictable, and
+  // compact them (hits forward, evictables backward) into s_lru_out.
+  const int iters_buffer = (num_buffer_chunks + NUM_WARPS - 1) / NUM_WARPS;
+  int total_hit_count = 0;
+  int total_evict_count = 0;
+  for (int iter = 0; iter < iters_buffer; iter++) {
+    const int chunk_idx = warp_id + iter * NUM_WARPS;
+    const bool has_valid_chunk = chunk_idx < num_buffer_chunks;
+
+    const int pos = chunk_idx * kWarpSize + lane_id;
+    const bool has_valid_pos = has_valid_chunk && (pos < hot_size);
+    const int16_t slot = has_valid_pos ? row_lru[pos] : int16_t(-1);
+    // Corruption tripwire: lru/dgi are long-lived device state; if some
+    // external writer (e.g. stray RDMA into reused VRAM) corrupts a slot
+    // out of [0, region_stride), treat it as no-hit so it degrades to a re-miss
+    // instead of an unbounded dgi read (phases 3/5 bound the write side).
+    const int32_t cached_g =
+        (slot >= 0 && slot < region_stride) ? row_dgi[slot] : -1;
+
+    int found_topk_idx = -1;
+    if (cached_g >= 0) {
+      int h = hash_slot(cached_g, hash_size);
+      while (true) {
+        const int32_t k = s_hash_keys[h];
+        if (k == cached_g) {
+          found_topk_idx = static_cast<int32_t>(s_hash_vals[h]);
+          break;
+        }
+        if (k == kHashEmpty) break;
+        h = (h + 1) % hash_size;
+      }
+    }
+    const bool is_hit = found_topk_idx >= 0;
+    const bool is_evictable = has_valid_pos && !is_hit;
+
+    if (is_hit) {
+      s_topk[found_topk_idx] = kTokenDone;
+      store_hot_index(row_out, row_attention, found_topk_idx,
+                      static_cast<int32_t>(get_physical_hot_row(
+                          hot_block_table, request_row, hot_table_stride,
+                          hot_block_size, slot)),
+                      hot_block_size, attention_block_stride);
+    }
+
+    int local_hit_off = 0;
+    int local_evict_off = 0;
+    if (has_valid_chunk) {
+      const unsigned int hit_mask = __ballot_sync(0xFFFFFFFF, is_hit);
+      const unsigned int evict_mask = __ballot_sync(0xFFFFFFFF, is_evictable);
+      local_hit_off = __popc(hit_mask & lanes_before);
+      local_evict_off = __popc(evict_mask & lanes_before);
+      if (lane_id == 0) {
+        s_chunk_off[chunk_idx + 1] = __popc(hit_mask);
+        s_evict_off[chunk_idx + 1] = __popc(evict_mask);
+      }
+    }
+    __syncthreads();
+
+    if (warp_id == 0) {
+      total_hit_count =
+          warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1,
+                              num_buffer_chunks + 1, total_hit_count);
+      total_evict_count =
+          warp_inclusive_scan(s_evict_off, lane_id, chunk_idx + 1,
+                              num_buffer_chunks + 1, total_evict_count);
+      if (tid == 0) {
+        s_counters[0] = total_hit_count;
+      }
+    }
+    __syncthreads();
+
+    if (is_hit) {
+      const int off = s_chunk_off[chunk_idx] + local_hit_off;
+      s_lru_out[off] = slot;
+    }
+    if (is_evictable) {
+      const int off = s_evict_off[chunk_idx] + local_evict_off;
+      s_lru_out[hot_size - 1 - off] = slot;
+    }
+  }
+  __syncthreads();
+
+  // Reset prefix sums for the miss compaction (token chunks <= buffer
+  // chunks because hot_size >= top_k).
+  for (int i = tid; i < num_token_chunks + 1; i += blockDim.x) {
+    s_chunk_off[i] = 0;
+  }
+  __syncthreads();
+
+  // Phase 3: compact misses, assign them eviction slots (oldest first) and
+  // record the new ownership in device_global_indices.
+  const int iters_token = (num_token_chunks + NUM_WARPS - 1) / NUM_WARPS;
+  int miss_running_total = 0;
+  for (int iter = 0; iter < iters_token; iter++) {
+    const int chunk_idx = warp_id + iter * NUM_WARPS;
+    const bool has_valid_chunk = chunk_idx < num_token_chunks;
+
+    const int i = chunk_idx * kWarpSize + lane_id;
+    const bool has_valid_token = has_valid_chunk && (i < top_k);
+
+    int32_t g = 0;
+    bool is_miss = false;
+    if (has_valid_token) {
+      is_miss = s_topk[i] != kTokenDone;
+      if (is_miss) {
+        g = s_topk[i];
+      }
+    }
+
+    int local_miss_off = 0;
+    if (has_valid_chunk) {
+      const unsigned int miss_mask = __ballot_sync(0xFFFFFFFF, is_miss);
+      local_miss_off = __popc(miss_mask & lanes_before);
+      if (lane_id == 0) {
+        s_chunk_off[chunk_idx + 1] = __popc(miss_mask);
+      }
+    }
+    __syncthreads();
+
+    if (warp_id == 0) {
+      miss_running_total =
+          warp_inclusive_scan(s_chunk_off, lane_id, chunk_idx + 1,
+                              num_token_chunks + 1, miss_running_total);
+    }
+    __syncthreads();
+
+    if (is_miss) {
+      const int m = s_chunk_off[chunk_idx] + local_miss_off;
+      const int16_t evict_slot = s_lru_out[hot_size - 1 - m];
+      if (evict_slot < 0 || evict_slot >= region_stride) {
+        // Corruption tripwire: an out-of-range slot (corrupted lru state,
+        // see phase 2) must not become a hot-cache/dgi write. Resolve the
+        // entry as invalid (-1, masked by attention, not in the miss set);
+        // it re-misses on a later step. Phase 5 re-checks the same
+        // s_lru_out value, so its copy is skipped consistently.
+        store_hot_index(row_out, row_attention, i, -1, hot_block_size,
+                        attention_block_stride);
+        if (swap_host_physical_rows != nullptr) {
+          const int64_t compact_index =
+              static_cast<int64_t>(batch_row) * top_k + m;
+          swap_host_physical_rows[compact_index] = g;
+          swap_device_physical_rows[compact_index] = -1;
+        }
+      } else {
+        // Reuse s_topk as compacted miss scratch: m < i always holds (done
+        // entries are skipped), so writes never overrun pending reads.
+        s_topk[m] = g;
+        const int32_t physical_row = static_cast<int32_t>(
+            get_physical_hot_row(hot_block_table, request_row, hot_table_stride,
+                                 hot_block_size, evict_slot));
+        store_hot_index(row_out, row_attention, i, physical_row, hot_block_size,
+                        attention_block_stride);
+        if (swap_host_physical_rows != nullptr) {
+          const int64_t compact_index =
+              static_cast<int64_t>(batch_row) * top_k + m;
+          swap_host_physical_rows[compact_index] = g;
+          swap_device_physical_rows[compact_index] = physical_row;
+        }
+        if (row_miss != nullptr) row_miss[i] = 1;
+        row_dgi[evict_slot] = g;
+      }
+    }
+  }
+  __syncthreads();
+
+  const int total_hits = s_counters[0];
+  const int total_misses = top_k - total_hits - s_counters[1];
+  if (swap_counts != nullptr && tid == 0) {
+    swap_counts[batch_row] = total_misses;
+  }
+  if (stats != nullptr && tid == 0) {
+    atomicAdd(&stats[0], static_cast<unsigned long long>(total_hits));
+    atomicAdd(&stats[1], static_cast<unsigned long long>(total_misses));
+  }
+
+  // Phase 4: write back the LRU order: stale evictables at the front,
+  // freshly loaded misses next, then hits at MRU.
+  const int total_evictable = hot_size - total_hits;
+  const int remaining_evictable = total_evictable - total_misses;
+  for (int i = tid; i < hot_size; i += blockDim.x) {
+    if (i < remaining_evictable) {
+      row_lru[i] = s_lru_out[hot_size - 1 - total_misses - i];
+    } else if (i < remaining_evictable + total_misses) {
+      row_lru[i] = s_lru_out[hot_size - 1 - (i - remaining_evictable)];
+    } else {
+      row_lru[i] = s_lru_out[i - remaining_evictable - total_misses];
+    }
+  }
+}
+
+// Shared-layer plan replay. Given a plan (hot_indices + miss_mask) already
+// computed by the group's index-producing layer, gather THIS
+// layer's own missed KV rows into the planned hot slots. No LRU resolution:
+// index-sharing shared layers see identical global_indices/hot_indices, so the
+// slot assignment is identical -- only the per-layer bytes differ. Fixed shape
+// (num_rows x top_k), so it is CUDA-graph-capture safe.
+__global__ void hisparse_gather_plan_kernel(
+    const char* __restrict__ host_cache,         // [host_rows, row_bytes]
+    char* __restrict__ hot_cache,                // [hot_rows, row_bytes]
+    const int32_t* __restrict__ global_indices,  // [num_rows, top_k]
+    const int32_t* __restrict__ hot_indices,  // [num_rows, top_k] abs hot rows
+    const int32_t* __restrict__ miss_mask,    // [num_rows, top_k]
+    int32_t* __restrict__ attention_indices,  // [num_rows, top_k]
+    const int32_t* __restrict__ request_state_indices,  // [num_rows] or nullptr
+    const int64_t host_rows, const int64_t hot_rows, const int64_t row_bytes,
+    const int64_t hot_block_stride, const int32_t hot_block_size,
+    const int32_t top_k, const int64_t attention_block_stride) {
+  const int NUM_WARPS = blockDim.x / kWarpSize;
+  // Columns are interleaved across gridDim.y blocks per row so few-row
+  // launches (local-prefill staging's single-row layout, small decode
+  // batches) still fill the device with outstanding host reads instead of
+  // starving on one block per row.
+  const int row = blockIdx.x;
+  const bool is_padding =
+      request_state_indices != nullptr && request_state_indices[row] < 0;
+  const int warp_id = threadIdx.x / kWarpSize;
+  const int lane_id = threadIdx.x % kWarpSize;
+  const int col_start = blockIdx.y * NUM_WARPS + warp_id;
+  const int col_stride = gridDim.y * NUM_WARPS;
+  const int64_t base = static_cast<int64_t>(row) * top_k;
+  for (int col = col_start; col < top_k; col += col_stride) {
+    const int32_t dst = hot_indices[base + col];
+    if (lane_id == 0 && attention_indices != nullptr) {
+      attention_indices[base + col] =
+          is_padding
+              ? -1
+              : (dst < 0 ? -1
+                         : (dst / hot_block_size) * attention_block_stride +
+                               dst % hot_block_size);
+    }
+    if (is_padding) {
+      continue;
+    }
+    if (miss_mask[base + col] == 0) {
+      continue;
+    }
+    const int32_t g = global_indices[base + col];
+    if (g < 0 || dst < 0 || dst >= hot_rows) {
+      continue;
+    }
+    if (g < host_rows) {
+      copy_row_warp(lane_id, host_cache + static_cast<int64_t>(g) * row_bytes,
+                    cache_row_ptr(hot_cache, dst, hot_block_size,
+                                  hot_block_stride, row_bytes),
+                    row_bytes);
+    } else {
+      // No source row for g: zero the planned slot rather than serving
+      // whatever bytes it held (see the swap-in kernel's phase 5).
+      zero_cache_row_warp(lane_id, hot_cache, dst, hot_block_size,
+                          hot_block_stride, row_bytes);
+    }
+  }
+}
+
+// Copy the compact swap rows produced by the residency resolver.
+__global__ void hisparse_gather_compact_kernel(
+    const char* __restrict__ host_cache, char* __restrict__ hot_cache,
+    const int32_t* __restrict__ miss_global_indices,
+    const int32_t* __restrict__ miss_hot_indices,
+    const int32_t* __restrict__ miss_counts, const int64_t host_rows,
+    const int64_t hot_rows, const int64_t row_bytes,
+    const int64_t hot_block_stride, const int32_t hot_block_size,
+    const int32_t top_k) {
+  const int NUM_WARPS = blockDim.x / kWarpSize;
+  const int row = blockIdx.x;
+  const int warp_id = threadIdx.x / kWarpSize;
+  const int lane_id = threadIdx.x % kWarpSize;
+  const int miss_count = min(max(miss_counts[row], 0), top_k);
+  const int col_start = blockIdx.y * NUM_WARPS + warp_id;
+  const int col_stride = gridDim.y * NUM_WARPS;
+  const int64_t base = static_cast<int64_t>(row) * top_k;
+  for (int col = col_start; col < miss_count; col += col_stride) {
+    const int32_t g = miss_global_indices[base + col];
+    const int32_t dst = miss_hot_indices[base + col];
+    if (g < 0 || dst < 0 || dst >= hot_rows) {
+      continue;
+    }
+    if (g < host_rows) {
+      copy_row_warp(lane_id, host_cache + static_cast<int64_t>(g) * row_bytes,
+                    cache_row_ptr(hot_cache, dst, hot_block_size,
+                                  hot_block_stride, row_bytes),
+                    row_bytes);
+    } else {
+      zero_cache_row_warp(lane_id, hot_cache, dst, hot_block_size,
+                          hot_block_stride, row_bytes);
+    }
+  }
+}
+
+__global__ void hisparse_invalidate_written_slots_kernel(
+    int32_t* __restrict__ device_global_indices,
+    const int32_t* __restrict__ request_state_indices,
+    const int32_t* __restrict__ req_id_per_token,
+    const int64_t* __restrict__ written_slots, const int64_t num_tokens,
+    const int64_t num_request_ids, const int64_t num_state_rows,
+    const int64_t region_stride) {
+  const int64_t token_idx = blockIdx.x;
+  if (token_idx >= num_tokens) {
+    return;
+  }
+  const int32_t req_idx = req_id_per_token[token_idx];
+  if (req_idx < 0 || req_idx >= num_request_ids) {
+    return;
+  }
+  const int32_t state_idx = request_state_indices[req_idx];
+  const int64_t written_slot = written_slots[token_idx];
+  if (state_idx < 0 || state_idx >= num_state_rows || written_slot < 0) {
+    return;
+  }
+  int32_t* row = device_global_indices + state_idx * region_stride;
+  for (int64_t offset = threadIdx.x; offset < region_stride;
+       offset += blockDim.x) {
+    if (row[offset] == written_slot) {
+      row[offset] = -1;
+    }
+  }
+}
+
+int64_t check_2d_rows(const torch::stable::Tensor& t, const char* name,
+                      int64_t row_bytes) {
+  STD_TORCH_CHECK(t.dim() == 2, name, " must be 2D");
+  STD_TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+  STD_TORCH_CHECK(t.size(1) * t.element_size() == row_bytes, name,
+                  " row width mismatch");
+  return t.size(0);
+}
+
+}  // namespace
+
+void hisparse_invalidate_written_slots(
+    torch::stable::Tensor& device_global_indices,
+    torch::stable::Tensor const& request_state_indices,
+    torch::stable::Tensor const& req_id_per_token,
+    torch::stable::Tensor const& written_slots) {
+  STD_TORCH_CHECK(device_global_indices.is_cuda() &&
+                      request_state_indices.is_cuda() &&
+                      req_id_per_token.is_cuda() && written_slots.is_cuda(),
+                  "HiSparse invalidation tensors must be on CUDA");
+  STD_TORCH_CHECK(
+      device_global_indices.scalar_type() ==
+              torch::headeronly::ScalarType::Int &&
+          device_global_indices.dim() == 2 &&
+          device_global_indices.is_contiguous(),
+      "device_global_indices must be a contiguous 2D int32 CUDA tensor");
+  STD_TORCH_CHECK(
+      request_state_indices.scalar_type() ==
+              torch::headeronly::ScalarType::Int &&
+          request_state_indices.dim() == 1 &&
+          request_state_indices.is_contiguous(),
+      "request_state_indices must be a contiguous 1D int32 CUDA tensor");
+  STD_TORCH_CHECK(
+      req_id_per_token.scalar_type() == torch::headeronly::ScalarType::Int &&
+          req_id_per_token.dim() == 1 && req_id_per_token.is_contiguous(),
+      "req_id_per_token must be a contiguous 1D int32 CUDA tensor");
+  STD_TORCH_CHECK(
+      written_slots.scalar_type() == torch::headeronly::ScalarType::Long &&
+          written_slots.dim() == 1 && written_slots.is_contiguous(),
+      "written_slots must be a contiguous 1D int64 CUDA tensor");
+  STD_TORCH_CHECK(req_id_per_token.numel() == written_slots.numel(),
+                  "req_id_per_token and written_slots must have equal length");
+  const int device_index = device_global_indices.get_device_index();
+  STD_TORCH_CHECK(
+      request_state_indices.get_device_index() == device_index &&
+          req_id_per_token.get_device_index() == device_index &&
+          written_slots.get_device_index() == device_index,
+      "HiSparse invalidation tensors must be on the same CUDA device");
+
+  const int64_t num_tokens = written_slots.numel();
+  const int64_t num_state_rows = device_global_indices.size(0);
+  const int64_t region_stride = device_global_indices.size(1);
+  if (num_tokens == 0 || num_state_rows == 0 || region_stride == 0) {
+    return;
+  }
+  STD_TORCH_CHECK(num_tokens <= INT32_MAX,
+                  "HiSparse invalidation grid exceeds CUDA limits");
+
+  constexpr int kBlockSize = 256;
+  const torch::stable::accelerator::DeviceGuard device_guard(device_index);
+  const cudaStream_t stream = get_current_cuda_stream();
+  hisparse_invalidate_written_slots_kernel<<<static_cast<int>(num_tokens),
+                                             kBlockSize, 0, stream>>>(
+      device_global_indices.mutable_data_ptr<int32_t>(),
+      request_state_indices.const_data_ptr<int32_t>(),
+      req_id_per_token.const_data_ptr<int32_t>(),
+      written_slots.const_data_ptr<int64_t>(), num_tokens,
+      request_state_indices.numel(), num_state_rows, region_stride);
+  const cudaError_t launch_error = cudaGetLastError();
+  STD_TORCH_CHECK(launch_error == cudaSuccess,
+                  "HiSparse invalidation kernel launch failed: ",
+                  cudaGetErrorString(launch_error));
+}
+
+void hisparse_resolve_residency(
+    torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
+    torch::stable::Tensor const& hot_block_table,
+    torch::stable::Tensor const& global_indices,
+    torch::stable::Tensor& hot_indices,
+    torch::stable::Tensor& device_global_indices,
+    torch::stable::Tensor& lru_slots,
+    std::optional<torch::stable::Tensor> const& request_state_indices,
+    int64_t region_stride,
+    std::optional<torch::stable::Tensor> const& miss_mask,
+    std::optional<torch::stable::Tensor> const& stats,
+    std::optional<torch::stable::Tensor> const& attention_indices,
+    int64_t attention_block_stride,
+    std::optional<torch::stable::Tensor> const& request_ids,
+    std::optional<torch::stable::Tensor> const& source_block_table,
+    int64_t source_block_size,
+    std::optional<torch::stable::Tensor> const& resolved_global_indices,
+    std::optional<torch::stable::Tensor> const& valid_counts,
+    std::optional<torch::stable::Tensor> const& swap_host_physical_rows,
+    std::optional<torch::stable::Tensor> const& swap_device_physical_rows,
+    std::optional<torch::stable::Tensor> const& swap_counts,
+    std::optional<torch::stable::Tensor> const& resident_block_table,
+    int64_t resident_block_size, int64_t resident_null_block) {
+  STD_TORCH_CHECK(
+      host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache),
+      "host_cache must be pinned CPU memory");
+  STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA");
+  STD_TORCH_CHECK(
+      hot_block_table.is_cuda() &&
+          hot_block_table.scalar_type() == torch::headeronly::ScalarType::Int &&
+          hot_block_table.dim() == 2,
+      "hot_block_table must be a 2D int32 CUDA tensor");
+  STD_TORCH_CHECK(global_indices.is_cuda() && hot_indices.is_cuda() &&
+                      device_global_indices.is_cuda() && lru_slots.is_cuda(),
+                  "index tensors must be on CUDA");
+  STD_TORCH_CHECK(
+      global_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
+          hot_indices.scalar_type() == torch::headeronly::ScalarType::Int,
+      "global_indices/hot_indices must be int32");
+  STD_TORCH_CHECK(
+      device_global_indices.scalar_type() == torch::headeronly::ScalarType::Int,
+      "device_global_indices must be int32");
+  STD_TORCH_CHECK(
+      lru_slots.scalar_type() == torch::headeronly::ScalarType::Short,
+      "lru_slots must be int16");
+  STD_TORCH_CHECK(global_indices.dim() == 2 && global_indices.stride(1) == 1,
+                  "global_indices must be row-major 2D");
+  const int64_t num_rows = global_indices.size(0);
+  STD_TORCH_CHECK(num_rows >= 0 && num_rows <= INT32_MAX,
+                  "num_rows must fit the CUDA grid");
+  STD_TORCH_CHECK(hot_indices.size(0) == num_rows &&
+                      hot_indices.size(1) == global_indices.size(1) &&
+                      hot_indices.is_contiguous(),
+                  "hot_indices must have one contiguous row per launch row");
+  STD_TORCH_CHECK(
+      device_global_indices.dim() == 2 && device_global_indices.is_contiguous(),
+      "device_global_indices must be contiguous 2D");
+  STD_TORCH_CHECK(
+      lru_slots.dim() == 2 &&
+          lru_slots.size(0) == device_global_indices.size(0) &&
+          lru_slots.is_contiguous(),
+      "lru_slots must be contiguous with one row per request state");
+  STD_TORCH_CHECK(hot_cache.dim() == 3,
+                  "hot_cache must be [num_blocks, block_size, row_width]");
+  const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size();
+  STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned");
+  STD_TORCH_CHECK(hot_cache.stride(1) * hot_cache.element_size() == row_bytes,
+                  "hot-cache rows must be contiguous");
+  const int64_t hot_block_size = hot_cache.size(1);
+  const int64_t hot_rows = hot_cache.size(0) * hot_block_size;
+  const int64_t hot_block_stride =
+      hot_cache.stride(0) * hot_cache.element_size();
+  const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes);
+  const auto launch_rows = static_cast<int32_t>(num_rows);
+  const auto top_k = static_cast<int32_t>(global_indices.size(1));
+  const auto hot_size = static_cast<int32_t>(lru_slots.size(1));
+  STD_TORCH_CHECK(hot_size >= top_k, "hot buffer size must be >= top_k");
+  STD_TORCH_CHECK(hot_size <= 32768, "hot buffer size must fit int16 slots");
+  STD_TORCH_CHECK(region_stride == hot_size,
+                  "region_stride must match the LRU size");
+  STD_TORCH_CHECK(device_global_indices.size(1) == region_stride,
+                  "device_global_indices must cover the full hot region");
+  STD_TORCH_CHECK(hot_block_table.size(1) >=
+                      (region_stride + hot_block_size - 1) / hot_block_size,
+                  "hot_block_table has too few columns");
+  STD_TORCH_CHECK(hot_rows < INT32_MAX, "hot indices must fit int32");
+
+  const int32_t* request_ids_ptr = nullptr;
+  const int32_t* source_block_table_ptr = nullptr;
+  int64_t source_bt_stride = 0;
+  int32_t source_num_reqs = 0;
+  int32_t source_num_blocks = 0;
+  STD_TORCH_CHECK(
+      request_ids.has_value() == source_block_table.has_value(),
+      "request_ids and source_block_table must be provided together");
+  if (source_block_table.has_value()) {
+    auto const& req = request_ids.value();
+    auto const& table = source_block_table.value();
+    STD_TORCH_CHECK(
+        req.is_cuda() &&
+            req.scalar_type() == torch::headeronly::ScalarType::Int &&
+            req.numel() >= launch_rows && req.is_contiguous(),
+        "request_ids must be contiguous int32 on CUDA with one entry per row");
+    STD_TORCH_CHECK(
+        table.is_cuda() &&
+            table.scalar_type() == torch::headeronly::ScalarType::Int &&
+            table.dim() == 2 && table.stride(1) == 1,
+        "source_block_table must be a row-major 2D int32 CUDA tensor");
+    STD_TORCH_CHECK(source_block_size > 0,
+                    "source_block_size must be positive");
+    request_ids_ptr = req.const_data_ptr<int32_t>();
+    source_block_table_ptr = table.const_data_ptr<int32_t>();
+    source_bt_stride = table.stride(0);
+    source_num_reqs = static_cast<int32_t>(table.size(0));
+    source_num_blocks = static_cast<int32_t>(table.size(1));
+  }
+  const int32_t required_hot_rows =
+      request_ids_ptr != nullptr ? source_num_reqs : launch_rows;
+  STD_TORCH_CHECK(hot_block_table.size(0) >= required_hot_rows,
+                  "hot_block_table has too few rows");
+
+  const int32_t* resident_block_table_ptr = nullptr;
+  int64_t resident_bt_stride = 0;
+  int32_t resident_num_reqs = 0;
+  int32_t resident_num_blocks = 0;
+  if (resident_block_table.has_value()) {
+    auto const& table = resident_block_table.value();
+    STD_TORCH_CHECK(
+        source_block_table.has_value() && request_ids.has_value(),
+        "resident lookup requires request_ids and source_block_table");
+    STD_TORCH_CHECK(
+        table.is_cuda() &&
+            table.scalar_type() == torch::headeronly::ScalarType::Int &&
+            table.dim() == 2 && table.stride(1) == 1,
+        "resident_block_table must be a row-major 2D int32 CUDA tensor");
+    STD_TORCH_CHECK(resident_block_size == hot_block_size,
+                    "resident and hot block sizes must match");
+    STD_TORCH_CHECK(resident_null_block >= 0,
+                    "resident null block must be non-negative");
+    resident_block_table_ptr = table.const_data_ptr<int32_t>();
+    resident_bt_stride = table.stride(0);
+    resident_num_reqs = static_cast<int32_t>(table.size(0));
+    resident_num_blocks = static_cast<int32_t>(table.size(1));
+  }
+
+  int32_t* resolved_global_indices_ptr = nullptr;
+  if (resolved_global_indices.has_value()) {
+    auto const& resolved = resolved_global_indices.value();
+    STD_TORCH_CHECK(
+        resolved.is_cuda() &&
+            resolved.scalar_type() == torch::headeronly::ScalarType::Int &&
+            resolved.dim() == 2 && resolved.size(0) == launch_rows &&
+            resolved.size(1) == global_indices.size(1) &&
+            resolved.is_contiguous(),
+        "resolved_global_indices must be contiguous int32 matching indices");
+    resolved_global_indices_ptr = resolved.mutable_data_ptr<int32_t>();
+  }
+
+  int32_t* valid_counts_ptr = nullptr;
+  if (valid_counts.has_value()) {
+    auto const& counts = valid_counts.value();
+    STD_TORCH_CHECK(
+        counts.is_cuda() &&
+            counts.scalar_type() == torch::headeronly::ScalarType::Int &&
+            counts.dim() == 1 && counts.size(0) == launch_rows,
+        "valid_counts must be int32 on CUDA with one entry per row");
+    valid_counts_ptr = counts.mutable_data_ptr<int32_t>();
+  }
+
+  int32_t* swap_host_physical_rows_ptr = nullptr;
+  int32_t* swap_device_physical_rows_ptr = nullptr;
+  int32_t* swap_counts_ptr = nullptr;
+  const bool has_compact_swaps = swap_host_physical_rows.has_value();
+  STD_TORCH_CHECK(has_compact_swaps == swap_device_physical_rows.has_value() &&
+                      has_compact_swaps == swap_counts.has_value(),
+                  "compact swap tensors must be provided together");
+  if (has_compact_swaps) {
+    auto const& globals = swap_host_physical_rows.value();
+    auto const& hots = swap_device_physical_rows.value();
+    auto const& counts = swap_counts.value();
+    STD_TORCH_CHECK(
+        globals.is_cuda() && hots.is_cuda() && counts.is_cuda() &&
+            globals.scalar_type() == torch::headeronly::ScalarType::Int &&
+            hots.scalar_type() == torch::headeronly::ScalarType::Int &&
+            counts.scalar_type() == torch::headeronly::ScalarType::Int &&
+            globals.dim() == 2 && hots.dim() == 2 &&
+            globals.size(0) == launch_rows &&
+            globals.size(1) == global_indices.size(1) &&
+            hots.size(0) == launch_rows &&
+            hots.size(1) == global_indices.size(1) &&
+            counts.numel() >= launch_rows && globals.is_contiguous() &&
+            hots.is_contiguous() && counts.is_contiguous(),
+        "compact swaps must be contiguous int32 matching indices");
+    swap_host_physical_rows_ptr = globals.mutable_data_ptr<int32_t>();
+    swap_device_physical_rows_ptr = hots.mutable_data_ptr<int32_t>();
+    swap_counts_ptr = counts.mutable_data_ptr<int32_t>();
+  }
+
+  const int32_t* request_state_ptr = nullptr;
+  int32_t request_state_count = 0;
+  if (request_state_indices.has_value()) {
+    auto const& state_indices = request_state_indices.value();
+    STD_TORCH_CHECK(
+        state_indices.is_cuda() &&
+            state_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
+            state_indices.dim() == 1 && state_indices.is_contiguous() &&
+            state_indices.numel() <= INT32_MAX,
+        "request_state_indices must be contiguous 1D int32 on CUDA");
+    request_state_ptr = state_indices.const_data_ptr<int32_t>();
+    request_state_count = static_cast<int32_t>(state_indices.numel());
+  } else {
+    STD_TORCH_CHECK(device_global_indices.size(0) >= launch_rows,
+                    "device_global_indices has too few rows");
+  }
+
+  // Optional output: 1 at columns requiring a host-to-device swap, 0 elsewhere.
+  int32_t* miss_mask_ptr = nullptr;
+  if (miss_mask.has_value()) {
+    auto const& mm = miss_mask.value();
+    STD_TORCH_CHECK(
+        mm.is_cuda() &&
+            mm.scalar_type() == torch::headeronly::ScalarType::Int &&
+            mm.dim() == 2 && mm.is_contiguous() && mm.size(0) == launch_rows &&
+            mm.size(1) == global_indices.size(1),
+        "miss_mask must be a contiguous int32 CUDA tensor matching "
+        "global_indices");
+    miss_mask_ptr = mm.mutable_data_ptr<int32_t>();
+  }
+
+  unsigned long long* stats_ptr = nullptr;
+  if (stats.has_value()) {
+    auto const& st = stats.value();
+    STD_TORCH_CHECK(
+        st.is_cuda() && st.is_contiguous() && st.dim() == 1 &&
+            st.scalar_type() == torch::headeronly::ScalarType::UInt64 &&
+            st.numel() == 2,
+        "stats must be a contiguous two-element uint64 CUDA tensor");
+    stats_ptr = static_cast<unsigned long long*>(st.mutable_data_ptr());
+  }
+
+  int32_t* attention_indices_ptr = nullptr;
+  if (attention_indices.has_value()) {
+    auto const& indices = attention_indices.value();
+    STD_TORCH_CHECK(
+        indices.is_cuda() &&
+            indices.scalar_type() == torch::headeronly::ScalarType::Int &&
+            indices.dim() == 2 && indices.size(0) == launch_rows &&
+            indices.stride(1) == 1 && indices.size(1) == global_indices.size(1),
+        "attention_indices must be row-major int32 matching global_indices");
+    STD_TORCH_CHECK(attention_block_stride >= hot_block_size,
+                    "attention block stride must cover one hot block");
+    attention_indices_ptr = indices.mutable_data_ptr<int32_t>();
+  }
+
+  if (launch_rows == 0 || top_k == 0) {
+    return;
+  }
+
+  constexpr int kBlockSize = 1024;
+  const int hash_size = 2 * top_k;
+  const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
+  const size_t smem_bytes =
+      sizeof(int32_t) * (top_k + 2 * (num_buffer_chunks + 1) + hash_size + 3) +
+      sizeof(int16_t) * (hot_size + hash_size);
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      hot_cache.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  const int64_t attention_row_stride =
+      attention_indices.has_value() ? attention_indices.value().stride(0) : 0;
+  const int64_t valid_count_stride =
+      valid_counts.has_value() ? valid_counts.value().stride(0) : 0;
+  auto kernel = hisparse_resolve_residency_kernel;
+  if (smem_bytes > 48 * 1024) {
+    const cudaError_t attribute_error = cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+    STD_TORCH_CHECK(attribute_error == cudaSuccess,
+                    "failed to configure HiSparse swap-in shared memory: ",
+                    cudaGetErrorString(attribute_error));
+  }
+  kernel<<<launch_rows, kBlockSize, smem_bytes, stream>>>(
+      hot_block_table.const_data_ptr<int32_t>(),
+      global_indices.const_data_ptr<int32_t>(), request_ids_ptr,
+      source_block_table_ptr, resident_block_table_ptr,
+      resolved_global_indices_ptr, valid_counts_ptr,
+      swap_host_physical_rows_ptr, swap_device_physical_rows_ptr,
+      swap_counts_ptr, hot_indices.mutable_data_ptr<int32_t>(),
+      attention_indices_ptr, miss_mask_ptr,
+      device_global_indices.mutable_data_ptr<int32_t>(),
+      lru_slots.mutable_data_ptr<int16_t>(), stats_ptr, request_state_ptr,
+      request_state_count, host_rows, hot_block_table.stride(0), hot_block_size,
+      top_k, hot_size, hash_size, region_stride, attention_block_stride,
+      source_bt_stride, source_num_reqs, source_num_blocks,
+      static_cast<int32_t>(source_block_size), resident_bt_stride,
+      resident_num_reqs, resident_num_blocks,
+      static_cast<int32_t>(resident_block_size),
+      static_cast<int32_t>(resident_null_block), global_indices.stride(0),
+      attention_row_stride, valid_count_stride);
+  const cudaError_t launch_error = cudaGetLastError();
+  STD_TORCH_CHECK(launch_error == cudaSuccess,
+                  "HiSparse residency kernel launch failed: ",
+                  cudaGetErrorString(launch_error));
+}
+
+void hisparse_gather_plan(
+    torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
+    torch::stable::Tensor const& global_indices,
+    torch::stable::Tensor const& hot_indices,
+    torch::stable::Tensor const& miss_mask,
+    std::optional<torch::stable::Tensor> const& request_state_indices,
+    std::optional<torch::stable::Tensor> const& attention_indices,
+    int64_t attention_block_stride) {
+  STD_TORCH_CHECK(
+      host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache),
+      "host_cache must be pinned CPU memory");
+  STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA");
+  STD_TORCH_CHECK(
+      global_indices.is_cuda() && hot_indices.is_cuda() && miss_mask.is_cuda(),
+      "plan tensors must be on CUDA");
+  STD_TORCH_CHECK(
+      global_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
+          hot_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
+          miss_mask.scalar_type() == torch::headeronly::ScalarType::Int,
+      "plan tensors must be int32");
+  STD_TORCH_CHECK(global_indices.dim() == 2 && global_indices.is_contiguous(),
+                  "global_indices must be contiguous 2D");
+  STD_TORCH_CHECK(
+      hot_indices.size(0) == global_indices.size(0) &&
+          hot_indices.size(1) == global_indices.size(1) &&
+          miss_mask.size(0) == global_indices.size(0) &&
+          miss_mask.size(1) == global_indices.size(1) &&
+          hot_indices.is_contiguous() && miss_mask.is_contiguous(),
+      "hot_indices/miss_mask must match contiguous 2D global_indices");
+
+  STD_TORCH_CHECK(hot_cache.dim() == 2 || hot_cache.dim() == 3,
+                  "hot_cache must be a 2D staging buffer or paged 3D cache");
+  const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size();
+  STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned");
+  STD_TORCH_CHECK(
+      hot_cache.stride(hot_cache.dim() - 2) * hot_cache.element_size() ==
+          row_bytes,
+      "hot-cache rows must be contiguous");
+  const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes);
+
+  const auto num_rows = static_cast<int32_t>(global_indices.size(0));
+  const auto top_k = static_cast<int32_t>(global_indices.size(1));
+
+  const int32_t* request_state_ptr = nullptr;
+  if (request_state_indices.has_value()) {
+    auto const& state_indices = request_state_indices.value();
+    STD_TORCH_CHECK(
+        state_indices.is_cuda() &&
+            state_indices.scalar_type() == torch::headeronly::ScalarType::Int &&
+            state_indices.numel() >= num_rows,
+        "request_state_indices must be int32 on CUDA with one entry per row");
+    request_state_ptr = state_indices.const_data_ptr<int32_t>();
+  }
+
+  int32_t* attention_indices_ptr = nullptr;
+  if (attention_indices.has_value()) {
+    auto const& indices = attention_indices.value();
+    STD_TORCH_CHECK(
+        indices.is_cuda() &&
+            indices.scalar_type() == torch::headeronly::ScalarType::Int &&
+            indices.dim() == 2 && indices.is_contiguous() &&
+            indices.size(0) == global_indices.size(0) &&
+            indices.size(1) == global_indices.size(1),
+        "attention_indices must be a contiguous int32 CUDA tensor matching "
+        "global_indices");
+    attention_indices_ptr = indices.mutable_data_ptr<int32_t>();
+  }
+
+  if (num_rows == 0 || top_k == 0) {
+    return;
+  }
+
+  // Match the swap-in kernel's block size: the gather serves 3 of every 4
+  // layers' misses (index-sharing replay), so per-row copy parallelism is
+  // the throughput limiter on cold rows.
+  constexpr int kBlockSize = 1024;
+  constexpr int kNumWarps = kBlockSize / kWarpSize;
+  // Interleave columns over enough blocks per row to cover the device even
+  // for few-row launches (local-prefill staging's single-row layout, small
+  // decode batches), keeping >= 1 column per warp.
+  constexpr int kTargetBlocks = 256;
+  const int max_chunks = std::max(1, (top_k + kNumWarps - 1) / kNumWarps);
+  const int num_chunks =
+      std::min(max_chunks, std::max(1, kTargetBlocks / num_rows));
+  const dim3 grid(num_rows, num_chunks);
+  const int32_t hot_block_size =
+      hot_cache.dim() == 3 ? static_cast<int32_t>(hot_cache.size(1)) : 1;
+  if (attention_indices.has_value()) {
+    STD_TORCH_CHECK(attention_block_stride >= hot_block_size,
+                    "attention block stride must cover one hot block");
+  }
+  const int64_t hot_rows = hot_cache.size(0) * hot_block_size;
+  const int64_t hot_block_stride =
+      hot_cache.stride(0) * hot_cache.element_size();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      hot_cache.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  hisparse_gather_plan_kernel<<<grid, kBlockSize, 0, stream>>>(
+      static_cast<const char*>(host_cache.const_data_ptr()),
+      static_cast<char*>(hot_cache.mutable_data_ptr()),
+      global_indices.const_data_ptr<int32_t>(),
+      hot_indices.const_data_ptr<int32_t>(),
+      miss_mask.const_data_ptr<int32_t>(), attention_indices_ptr,
+      request_state_ptr, host_rows, hot_rows, row_bytes, hot_block_stride,
+      hot_block_size, top_k, attention_block_stride);
+}
+
+void hisparse_gather_compact(torch::stable::Tensor const& host_cache,
+                             torch::stable::Tensor& hot_cache,
+                             torch::stable::Tensor const& miss_global_indices,
+                             torch::stable::Tensor const& miss_hot_indices,
+                             torch::stable::Tensor const& miss_counts) {
+  STD_TORCH_CHECK(
+      host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache),
+      "host_cache must be pinned CPU memory");
+  STD_TORCH_CHECK(hot_cache.is_cuda(), "hot_cache must be on CUDA");
+  STD_TORCH_CHECK(
+      miss_global_indices.is_cuda() && miss_hot_indices.is_cuda() &&
+          miss_counts.is_cuda() &&
+          miss_global_indices.scalar_type() ==
+              torch::headeronly::ScalarType::Int &&
+          miss_hot_indices.scalar_type() ==
+              torch::headeronly::ScalarType::Int &&
+          miss_counts.scalar_type() == torch::headeronly::ScalarType::Int &&
+          miss_global_indices.dim() == 2 && miss_hot_indices.dim() == 2 &&
+          miss_global_indices.size(0) == miss_hot_indices.size(0) &&
+          miss_global_indices.size(1) == miss_hot_indices.size(1) &&
+          miss_counts.numel() >= miss_global_indices.size(0) &&
+          miss_global_indices.is_contiguous() &&
+          miss_hot_indices.is_contiguous() && miss_counts.is_contiguous(),
+      "compact miss plan must be matching contiguous int32 CUDA tensors");
+  STD_TORCH_CHECK(hot_cache.dim() == 3,
+                  "hot_cache must be [num_blocks, block_size, row_width]");
+  const int64_t row_bytes = hot_cache.size(-1) * hot_cache.element_size();
+  STD_TORCH_CHECK(row_bytes % 16 == 0, "KV rows must be 16-byte aligned");
+  STD_TORCH_CHECK(hot_cache.stride(1) * hot_cache.element_size() == row_bytes,
+                  "hot-cache rows must be contiguous");
+  const int64_t host_rows = check_2d_rows(host_cache, "host_cache", row_bytes);
+  const auto num_rows = static_cast<int32_t>(miss_global_indices.size(0));
+  const auto top_k = static_cast<int32_t>(miss_global_indices.size(1));
+  if (num_rows == 0 || top_k == 0) {
+    return;
+  }
+
+  constexpr int kBlockSize = 512;
+  constexpr int kTargetBlocks = 64;
+  const int num_chunks = std::min(8, std::max(1, kTargetBlocks / num_rows));
+  const dim3 grid(num_rows, num_chunks);
+  const int32_t hot_block_size = static_cast<int32_t>(hot_cache.size(1));
+  const int64_t hot_rows = hot_cache.size(0) * hot_block_size;
+  const int64_t hot_block_stride =
+      hot_cache.stride(0) * hot_cache.element_size();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      hot_cache.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  hisparse_gather_compact_kernel<<<grid, kBlockSize, 0, stream>>>(
+      static_cast<const char*>(host_cache.const_data_ptr()),
+      static_cast<char*>(hot_cache.mutable_data_ptr()),
+      miss_global_indices.const_data_ptr<int32_t>(),
+      miss_hot_indices.const_data_ptr<int32_t>(),
+      miss_counts.const_data_ptr<int32_t>(), host_rows, hot_rows, row_bytes,
+      hot_block_stride, hot_block_size, top_k);
+}
diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
index 98a04de783d4..5c9a522feb7b 100644
--- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
@@ -23,20 +23,18 @@ namespace batched_moe_align_block_size {
 // Note num_threads needs to be 1024 for BlockScan Reduction in the kernel.
 static constexpr int32_t num_threads = 1024;
 static constexpr int32_t num_blocks = 1;
+template <bool cooperative_writes>
 __global__ void batched_moe_align_block_size_kernel(
     int32_t const num_batches, int32_t const max_tokens_per_batch,
     int32_t const block_size, int32_t const* __restrict__ batch_num_tokens,
     int32_t* __restrict__ sorted_ids, int32_t* __restrict__ block_ids,
     int32_t* __restrict__ num_tokens_post_pad) {
-  // TODO(varun): This is a naive implementation. Could be optimized.
-
-  size_t const batch_id = threadIdx.x;
   size_t const stride = blockDim.x * gridDim.x;
   int32_t const num_blocks_per_batch =
       CEILDIV(max_tokens_per_batch, block_size);
-  int32_t const sorted_ids_size =
-      num_blocks_per_batch * num_batches * block_size;
-  int32_t const block_ids_size = sorted_ids_size / block_size;
+  size_t const sorted_ids_size =
+      static_cast<size_t>(num_blocks_per_batch) * num_batches * block_size;
+  size_t const block_ids_size = sorted_ids_size / block_size;
   int32_t const SENTINEL =
       num_batches * max_tokens_per_batch;  // To denote invalid entries.
   // Initialize sorted_ids
@@ -49,8 +47,8 @@ __global__ void batched_moe_align_block_size_kernel(
   }
 
   int32_t b_num_tokens = 0;
-  if (batch_id < num_batches) {
-    b_num_tokens = batch_num_tokens[batch_id];
+  if (threadIdx.x < num_batches) {
+    b_num_tokens = batch_num_tokens[threadIdx.x];
   }
   int32_t const ceil_b_num_tokens =
       CEILDIV(b_num_tokens, block_size) * block_size;
@@ -62,20 +60,56 @@ __global__ void batched_moe_align_block_size_kernel(
   BlockScan(temp_storage).ExclusiveSum(ceil_b_num_tokens, cumsum_val);
   __syncthreads();
 
-  bool const is_last_batch = batch_id == (num_batches - 1);
-  if (is_last_batch) {
+  if (threadIdx.x == num_batches - 1) {
     *num_tokens_post_pad = cumsum_val + ceil_b_num_tokens;
   }
 
-  if (batch_id < num_batches) {
-    int32_t const batch_offset = batch_id * max_tokens_per_batch;
+  if constexpr (cooperative_writes) {
+    __shared__ int32_t batch_cumsum[num_threads];
+    __shared__ int32_t valid_tokens[num_threads];
+    __shared__ int32_t batch_num_blocks[num_threads];
+    if (threadIdx.x < num_batches) {
+      batch_cumsum[threadIdx.x] = cumsum_val;
+      valid_tokens[threadIdx.x] = b_num_tokens;
+      batch_num_blocks[threadIdx.x] = ceil_b_num_tokens / block_size;
+    }
+    __syncthreads();
+
+    int32_t const max_num_groups = blockDim.x / WARP_SIZE;
+    int32_t num_groups = 1;
+    while (num_groups < num_batches && num_groups < max_num_groups) {
+      num_groups *= 2;
+    }
+    int32_t const threads_per_batch = blockDim.x / num_groups;
+    int32_t const group_id = threadIdx.x / threads_per_batch;
+    int32_t const group_offset = threadIdx.x % threads_per_batch;
+
+    // Assign at least one warp to each batch when possible.
+    for (int32_t batch_id = group_id; batch_id < num_batches;
+         batch_id += num_groups) {
+      size_t const batch_offset =
+          static_cast<size_t>(batch_id) * max_tokens_per_batch;
+      size_t const cumsum = batch_cumsum[batch_id];
+      for (size_t i = group_offset; i < valid_tokens[batch_id];
+           i += threads_per_batch) {
+        sorted_ids[cumsum + i] = static_cast<int32_t>(batch_offset + i);
+      }
+
+      size_t const block_start = cumsum / block_size;
+      for (size_t i = group_offset; i < batch_num_blocks[batch_id];
+           i += threads_per_batch) {
+        block_ids[block_start + i] = batch_id;
+      }
+    }
+  } else if (threadIdx.x < num_batches) {
+    size_t const batch_id = threadIdx.x;
+    size_t const batch_offset = batch_id * max_tokens_per_batch;
     for (size_t i = 0; i < b_num_tokens; ++i) {
-      sorted_ids[cumsum_val + i] = batch_offset + i;
+      sorted_ids[cumsum_val + i] = static_cast<int32_t>(batch_offset + i);
     }
 
-    int32_t const block_start = cumsum_val / block_size;
-    int32_t const num_blocks = ceil_b_num_tokens / block_size;
-    for (size_t i = 0; i < num_blocks; ++i) {
+    size_t const block_start = cumsum_val / block_size;
+    for (size_t i = 0; i < ceil_b_num_tokens / block_size; ++i) {
       block_ids[block_start + i] = batch_id;
     }
   }
@@ -202,7 +236,7 @@ __device__ void _moe_align_block_size_small_batch_expert(
     int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
     size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks,
     int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num,
-    int32_t* token_mask, bool has_expert_map) {
+    int32_t* token_mask, bool has_expert_map, int32_t* scatter_idx = nullptr) {
   // Compute input buffer offsets. Typically these will all be 0, except when
   // using Multi LoRA.
   int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
@@ -284,9 +318,14 @@ __device__ void _moe_align_block_size_small_batch_expert(
   }
 
   for (size_t i = tid; i < numel; i += stride) {
-    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
-                                            num_experts, has_expert_map);
-        expert_id != -1) {
+    int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts,
+                                        has_expert_map);
+    bool valid =
+        expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]);
+    if (scatter_idx != nullptr) {
+      scatter_idx[i] = valid ? static_cast<int32_t>(i) : -1;
+    }
+    if (expert_id != -1) {
       int32_t rank_post_pad =
           tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
 
@@ -304,14 +343,20 @@ __device__ void _count_and_sort_expert_tokens(
     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* __restrict__ token_mask,
-    int32_t model_offset, int32_t topk_num, bool has_expert_map) {
+    int32_t model_offset, int32_t topk_num, bool has_expert_map,
+    int32_t* scatter_idx = nullptr) {
   const size_t tid = blockIdx.y * blockDim.x + threadIdx.x;
   const size_t stride = blockDim.x * gridDim.y;
 
   for (size_t i = tid; i < numel; i += stride) {
-    if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
-                                            num_experts, has_expert_map);
-        expert_id != -1) {
+    int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts,
+                                        has_expert_map);
+    bool valid =
+        expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]);
+    if (scatter_idx != nullptr) {
+      scatter_idx[i] = valid ? static_cast<int32_t>(i) : -1;
+    }
+    if (expert_id != -1) {
       if (token_mask == nullptr || token_mask[i / topk_num]) {
         int32_t rank_post_pad = atomicAdd(
             &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
@@ -343,10 +388,11 @@ __global__ void count_and_sort_expert_tokens_kernel(
     const scalar_t* __restrict__ topk_ids,
     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, bool has_expert_map) {
+    int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map,
+    int32_t* scatter_idx) {
   _count_and_sort_expert_tokens(
       topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts,
-      max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map);
+      max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map, scatter_idx);
 }
 
 // Reduce the topk expert outputs per token (summed in fp32). The output is
@@ -499,12 +545,12 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
     int32_t* __restrict__ total_tokens_post_pad,
     int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
     size_t numel, int32_t max_num_tokens_padded, int32_t topk_num,
-    bool has_expert_map) {
+    bool has_expert_map, int32_t* scatter_idx) {
   _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
       topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
       num_experts, block_size, numel, max_num_tokens_padded,
       CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr,
-      has_expert_map);
+      has_expert_map, scatter_idx);
 }
 
 template <typename scalar_t>
@@ -624,16 +670,31 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel(
 
 // taken from
 // https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc
-void moe_align_block_size(
-    torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
-    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
-    torch::stable::Tensor num_tokens_post_pad,
-    std::optional<torch::stable::Tensor> maybe_expert_map) {
+void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts,
+                          int64_t block_size,
+                          torch::stable::Tensor sorted_token_ids,
+                          torch::stable::Tensor experts_ids,
+                          torch::stable::Tensor num_tokens_post_pad,
+                          std::optional<torch::stable::Tensor> maybe_expert_map,
+                          std::optional<torch::stable::Tensor> scatter_idx) {
   const torch::stable::accelerator::DeviceGuard device_guard(
       topk_ids.get_device_index());
   const cudaStream_t stream =
       get_current_cuda_stream(topk_ids.get_device_index());
 
+  int32_t* scatter_ptr = nullptr;
+  if (scatter_idx.has_value()) {
+    STD_TORCH_CHECK(
+        scatter_idx->scalar_type() == torch::headeronly::ScalarType::Int,
+        "scatter_idx must be int32");
+    STD_TORCH_CHECK(scatter_idx->device() == topk_ids.device() &&
+                        scatter_idx->is_contiguous() &&
+                        scatter_idx->numel() == topk_ids.numel(),
+                    "scatter_idx must be contiguous, on the same device, and "
+                    "match topk_ids.numel()");
+    scatter_ptr = reinterpret_cast<int32_t*>(scatter_idx->mutable_data_ptr());
+  }
+
   int64_t padded_num_experts =
       ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
   int experts_per_warp = WARP_SIZE;
@@ -679,7 +740,8 @@ void moe_align_block_size(
                   num_tokens_post_pad.mutable_data_ptr()),
               reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
               num_experts, block_size, topk_ids.numel(),
-              sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
+              sorted_token_ids.size(0), topk_ids.size(1), has_expert_map,
+              scatter_ptr);
         } else {
           torch::stable::Tensor cumsum_buffer = torch::stable::new_empty(
               topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
@@ -719,7 +781,7 @@ void moe_align_block_size(
               reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
               reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
               topk_ids.numel(), num_experts, sorted_token_ids.size(0),
-              topk_ids.size(1), has_expert_map);
+              topk_ids.size(1), has_expert_map, scatter_ptr);
         }
       });
 }
@@ -747,8 +809,16 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
   STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
   STD_TORCH_CHECK(B <= batched_kernel::num_threads);
 
-  batched_kernel::batched_moe_align_block_size_kernel<<<
-      batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>(
+  // Avoid coordination overhead for small capacities or many batches.
+  int64_t const cooperative_threshold = std::max<int64_t>(256, 8 * B);
+  bool const use_cooperative_writes =
+      max_tokens_per_batch >= cooperative_threshold;
+  auto kernel =
+      use_cooperative_writes
+          ? batched_kernel::batched_moe_align_block_size_kernel<true>
+          : batched_kernel::batched_moe_align_block_size_kernel<false>;
+  kernel<<<batched_kernel::num_blocks, batched_kernel::num_threads, 0,
+           stream>>>(
       B, max_tokens_per_batch, block_size,
       reinterpret_cast<const int32_t*>(batch_num_tokens.const_data_ptr()),
       reinterpret_cast<int32_t*>(sorted_ids.mutable_data_ptr()),
diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h
index 03057c3b1421..549a19e98148 100644
--- a/csrc/libtorch_stable/moe/moe_ops.h
+++ b/csrc/libtorch_stable/moe/moe_ops.h
@@ -36,11 +36,13 @@ void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
              std::optional<torch::stable::Tensor> topk_ids,
              std::optional<torch::stable::Tensor> expert_map);
 
-void moe_align_block_size(
-    torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
-    torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
-    torch::stable::Tensor num_tokens_post_pad,
-    std::optional<torch::stable::Tensor> maybe_expert_map);
+void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts,
+                          int64_t block_size,
+                          torch::stable::Tensor sorted_token_ids,
+                          torch::stable::Tensor experts_ids,
+                          torch::stable::Tensor num_tokens_post_pad,
+                          std::optional<torch::stable::Tensor> maybe_expert_map,
+                          std::optional<torch::stable::Tensor> scatter_idx);
 
 void batched_moe_align_block_size(
     int64_t max_tokens_per_batch, int64_t block_size,
diff --git a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu
index 251097456f9e..957cd60e49de 100644
--- a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu
+++ b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu
@@ -56,20 +56,18 @@ int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
       CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
 }
 
-void moe_permute_impl(
-    const torch::stable::Tensor& input,                 // [n_token, hidden]
-    const torch::stable::Tensor& topk_ids,              // [n_token, topk]
-    const torch::stable::Tensor& token_expert_indices,  // [n_token, topk]
+torch::stable::Tensor moe_sort_routing(
+    const torch::stable::Tensor& topk_ids,                   // [n_token, topk]
+    const torch::stable::Tensor& token_expert_indices,       // [n_token, topk]
     const std::optional<torch::stable::Tensor>& expert_map,  // [n_expert]
     int64_t n_expert, int64_t n_local_expert, int64_t topk,
-    torch::stable::Tensor& permuted_input,  // [permuted_size, hidden]
     torch::stable::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
     torch::stable::Tensor& inv_permuted_idx,           // [n_token, topk]
-    torch::stable::Tensor& permuted_idx,               // [permute_size]
     const std::optional<torch::stable::Tensor>& maybe_sort_workspace,
     const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
     const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
-    const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
+    const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort,
+    bool build_inverse) {
   STD_TORCH_CHECK(expert_first_token_offset.scalar_type() ==
                       torch::headeronly::ScalarType::Long,
                   "expert_first_token_offset must be int64");
@@ -88,12 +86,11 @@ void moe_permute_impl(
       "token_expert_indices shape must be same as inv_permuted_idx");
 
   const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  auto device = input.device();
-  auto n_token = input.sizes()[0];
-  auto n_hidden = input.sizes()[1];
+      topk_ids.get_device_index());
+  auto device = topk_ids.device();
+  auto n_token = topk_ids.sizes()[0];
   auto expanded_rows = n_token * topk;
-  auto stream = get_current_cuda_stream(input.get_device_index());
+  auto stream = get_current_cuda_stream(topk_ids.get_device_index());
 
   auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
   auto sort_workspace = maybe_allocate_tensor(
@@ -107,13 +104,11 @@ void moe_permute_impl(
       torch::headeronly::ScalarType::Int, device, "sorted_row_idx");
 
   CubKeyValueSorter sorter{};
-  int64_t* valid_num_ptr = nullptr;
   torch::stable::Tensor topk_ids_for_sort = topk_ids;
 
-  if (expert_map.has_value()) {
-    const int* expert_map_ptr = get_ptr<int>(expert_map.value());
-    valid_num_ptr =
-        get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
+  if (expert_map.has_value() || build_inverse) {
+    const int* expert_map_ptr =
+        expert_map.has_value() ? get_ptr<int>(expert_map.value()) : nullptr;
     topk_ids_for_sort = maybe_allocate_tensor(
         maybe_topk_ids_for_sort, topk_ids.sizes(),
         torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort");
@@ -126,7 +121,39 @@ void moe_permute_impl(
       get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
       get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
       get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
-      n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
+      n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream,
+      build_inverse ? get_ptr<int>(inv_permuted_idx) : nullptr);
+  return sorted_row_idx;
+}
+
+void moe_permute_impl(
+    const torch::stable::Tensor& input,                 // [n_token, hidden]
+    const torch::stable::Tensor& topk_ids,              // [n_token, topk]
+    const torch::stable::Tensor& token_expert_indices,  // [n_token, topk]
+    const std::optional<torch::stable::Tensor>& expert_map,  // [n_expert]
+    int64_t n_expert, int64_t n_local_expert, int64_t topk,
+    torch::stable::Tensor& permuted_input,  // [permuted_size, hidden]
+    torch::stable::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
+    torch::stable::Tensor& inv_permuted_idx,           // [n_token, topk]
+    torch::stable::Tensor& permuted_idx,               // [permute_size]
+    const std::optional<torch::stable::Tensor>& maybe_sort_workspace,
+    const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
+    const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
+    const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  auto sorted_row_idx = moe_sort_routing(
+      topk_ids, token_expert_indices, expert_map, n_expert, n_local_expert,
+      topk, expert_first_token_offset, inv_permuted_idx, maybe_sort_workspace,
+      maybe_permuted_experts_id, maybe_sorted_row_idx, maybe_topk_ids_for_sort,
+      false);
+  auto stream = get_current_cuda_stream(input.get_device_index());
+  auto n_token = input.size(0);
+  auto n_hidden = input.size(1);
+  auto valid_num_ptr =
+      expert_map.has_value()
+          ? get_ptr<int64_t>(expert_first_token_offset) + n_local_expert
+          : nullptr;
 
   MOE_DISPATCH(input.scalar_type(), [&] {
     expandInputRowsKernelLauncher<scalar_t>(
@@ -171,6 +198,22 @@ void moe_permute_with_scratch(
                    topk_ids_for_sort);
 }
 
+void moe_prepare_scatter(const torch::stable::Tensor& topk_ids,
+                         const torch::stable::Tensor& token_expert_indices,
+                         const std::optional<torch::stable::Tensor>& expert_map,
+                         int64_t n_expert, int64_t n_local_expert,
+                         torch::stable::Tensor& expert_first_token_offset,
+                         torch::stable::Tensor& scatter_idx,
+                         torch::stable::Tensor& sort_workspace,
+                         torch::stable::Tensor& sorted_experts,
+                         torch::stable::Tensor& sorted_rows,
+                         torch::stable::Tensor& topk_ids_for_sort) {
+  moe_sort_routing(topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk_ids.size(1), expert_first_token_offset,
+                   scatter_idx, sort_workspace, sorted_experts, sorted_rows,
+                   topk_ids_for_sort, true);
+}
+
 void moe_unpermute(
     const torch::stable::Tensor&
         permuted_hidden_states,                     // [n_token * topk, hidden]
@@ -315,6 +358,19 @@ void moe_unpermute(
   STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");
 }
 
+void moe_prepare_scatter(const torch::stable::Tensor& topk_ids,
+                         const torch::stable::Tensor& token_expert_indices,
+                         const std::optional<torch::stable::Tensor>& expert_map,
+                         int64_t n_expert, int64_t n_local_expert,
+                         torch::stable::Tensor& expert_first_token_offset,
+                         torch::stable::Tensor& scatter_idx,
+                         torch::stable::Tensor& sort_workspace,
+                         torch::stable::Tensor& sorted_experts,
+                         torch::stable::Tensor& sorted_rows,
+                         torch::stable::Tensor& topk_ids_for_sort) {
+  STD_TORCH_CHECK(false, "moe_prepare_scatter is not supported on CUDA < 12.0");
+}
+
 #endif
 
 bool moe_permute_unpermute_supported() {
@@ -328,5 +384,6 @@ bool moe_permute_unpermute_supported() {
 STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
   m.impl("moe_permute", TORCH_BOX(&moe_permute));
   m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch));
+  m.impl("moe_prepare_scatter", TORCH_BOX(&moe_prepare_scatter));
   m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute));
 }
diff --git a/csrc/libtorch_stable/moe/moe_wna16_utils.h b/csrc/libtorch_stable/moe/moe_wna16_utils.h
index 8ef03f0e6052..6ce52ee5b1d8 100644
--- a/csrc/libtorch_stable/moe/moe_wna16_utils.h
+++ b/csrc/libtorch_stable/moe/moe_wna16_utils.h
@@ -34,10 +34,6 @@ class ScalarType<half> {
   static __host__ __device__ float2 inline num22float2(const half2 x) {
     return __half22float2(x);
   }
-
-  static __host__ __device__ half2 inline float22num2(const float2 x) {
-    return __float22half2_rn(x);
-  }
 };
 
 template <>
@@ -71,10 +67,6 @@ class ScalarType<nv_bfloat16> {
   static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) {
     return __bfloat1622float2(x);
   }
-
-  static __host__ __device__ nv_bfloat162 inline float22num2(const float2 x) {
-    return __float22bfloat162_rn(x);
-  }
 #endif
 };
 
diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
index ec7851da3d3d..8d8bdf8baaf6 100644
--- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
@@ -82,34 +82,40 @@ __device__ inline int64_t findTotalEltsLessThanTarget(T const* sorted_indices,
   return target_location + 1;
 }
 
-// Calculates the start offset of the tokens for a given expert. The last
-// element is the total number of valid tokens
-__global__ void computeExpertFirstTokenOffsetKernel(
+// Computes expert offsets (ending with the valid row count) and optional
+// inverse indices.
+__global__ void computeExpertOffsetsAndInverseKernel(
     int const* sorted_experts, int64_t const sorted_experts_len,
-    int const num_experts, int64_t* expert_first_token_offset) {
-  // First, compute the global tid. We only need 1 thread per expert.
-  int const expert = blockIdx.x * blockDim.x + threadIdx.x;
+    int const num_experts, int64_t* expert_first_token_offset,
+    int const* sorted_rows, int* inverse) {
+  int const index = blockIdx.x * blockDim.x + threadIdx.x;
 
   // Note that expert goes [0, num_experts] (inclusive) because we want a count
   // for the total number of active tokens at the end of the scan.
-  if (expert >= num_experts + 1) {
-    return;
+  if (index < num_experts + 1) {
+    expert_first_token_offset[index] =
+        findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, index);
+  }
+  if (inverse != nullptr && index < sorted_experts_len) {
+    inverse[sorted_rows[index]] = index;
   }
-  expert_first_token_offset[expert] =
-      findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, expert);
 }
 
-void computeExpertFirstTokenOffset(int const* sorted_indices,
-                                   int const total_indices,
-                                   int const num_experts,
-                                   int64_t* expert_first_token_offset,
-                                   cudaStream_t stream) {
-  int const num_entries = num_experts + 1;
-  int const threads = std::min(1024, num_entries);
+void computeExpertOffsetsAndInverse(int const* sorted_indices,
+                                    int const total_indices,
+                                    int const num_experts,
+                                    int64_t* expert_first_token_offset,
+                                    int const* sorted_rows, int* inverse,
+                                    cudaStream_t stream) {
+  int const num_entries = inverse != nullptr
+                              ? std::max(num_experts + 1, total_indices)
+                              : num_experts + 1;
+  int const threads = inverse != nullptr ? 256 : std::min(1024, num_entries);
   int const blocks = (num_entries + threads - 1) / threads;
 
-  computeExpertFirstTokenOffsetKernel<<<blocks, threads, 0, stream>>>(
-      sorted_indices, total_indices, num_experts, expert_first_token_offset);
+  computeExpertOffsetsAndInverseKernel<<<blocks, threads, 0, stream>>>(
+      sorted_indices, total_indices, num_experts, expert_first_token_offset,
+      sorted_rows, inverse);
 }
 
 void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
@@ -117,7 +123,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
                        int64_t* expert_first_token_offset, int num_rows,
                        int num_experts, int num_experts_per_node, int k,
                        CubKeyValueSorter& sorter, void* sorter_ws,
-                       cudaStream_t stream) {
+                       cudaStream_t stream, int* inverse) {
   int64_t const expanded_num_rows = static_cast<int64_t>(k) * num_rows;
   // We need to use the full num_experts because that is the sentinel value used
   // by topk for disabled experts
@@ -127,9 +133,9 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
   sorter.run((void*)sorter_ws, sorter_ws_size_bytes, expert_for_source_row,
              permuted_experts, source_rows, permuted_rows, expanded_num_rows,
              stream);
-  computeExpertFirstTokenOffset(permuted_experts, expanded_num_rows,
-                                num_experts_per_node, expert_first_token_offset,
-                                stream);
+  computeExpertOffsetsAndInverse(
+      permuted_experts, expanded_num_rows, num_experts_per_node,
+      expert_first_token_offset, permuted_rows, inverse, stream);
 }
 
 __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
@@ -142,7 +148,7 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
   extern __shared__ int smem_expert_map[];
   // store expert_map in smem
   for (int i = tidx; i < num_experts; i += blockDim.x) {
-    smem_expert_map[i] = expert_map_ptr[i];
+    smem_expert_map[i] = expert_map_ptr ? expert_map_ptr[i] : i;
   }
   __syncthreads();
 
diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
index 89c278a4ed4b..53afad2f5d41 100644
--- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
@@ -43,18 +43,19 @@ class CubKeyValueSorter {
   int num_bits_;
 };
 
-void computeExpertFirstTokenOffset(int const* sorted_indices,
-                                   int const total_indices,
-                                   int const num_experts,
-                                   int64_t* expert_first_token_offset,
-                                   cudaStream_t stream);
+void computeExpertOffsetsAndInverse(int const* sorted_indices,
+                                    int const total_indices,
+                                    int const num_experts,
+                                    int64_t* expert_first_token_offset,
+                                    int const* sorted_rows, int* inverse,
+                                    cudaStream_t stream);
 
 void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
                        int* permuted_experts, int* permuted_rows,
                        int64_t* expert_first_token_offset, int num_rows,
                        int num_experts, int num_experts_per_node, int k,
                        CubKeyValueSorter& sorter, void* sorter_ws,
-                       cudaStream_t stream);
+                       cudaStream_t stream, int* inverse = nullptr);
 
 template <typename T>
 void expandInputRowsKernelLauncher(
diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp
index c0e966b2be9c..2feec3838993 100644
--- a/csrc/libtorch_stable/moe/torch_bindings.cpp
+++ b/csrc/libtorch_stable/moe/torch_bindings.cpp
@@ -39,7 +39,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
       "                     int block_size, Tensor! sorted_token_ids,"
       "                     Tensor! experts_ids,"
       "                     Tensor! num_tokens_post_pad,"
-      "                     Tensor? maybe_expert_map) -> ()");
+      "                     Tensor? maybe_expert_map,"
+      "                     Tensor(a!)? scatter_idx=None) -> ()");
 
   // Aligning the number of tokens to be processed by each expert such
   // that it is divisible by the block size, but for the batched case.
@@ -105,6 +106,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
       "permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, "
       "Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()");
 
+  m.def(
+      "moe_prepare_scatter(Tensor topk_ids, Tensor token_expert_indices, "
+      "Tensor? expert_map, int n_expert, int n_local_expert, "
+      "Tensor! expert_first_token_offset, Tensor! scatter_idx, "
+      "Tensor! sort_workspace, Tensor! sorted_experts, Tensor! sorted_rows, "
+      "Tensor! topk_ids_for_sort)->()");
+
   m.def(
       "moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
       "Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu b/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu
index c2faf7b76023..b1f369e5ab00 100644
--- a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu
+++ b/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu
@@ -109,9 +109,6 @@ __device__ __forceinline__ void nvfp4_unpack16_e2m1(uint64_t raw,
 #endif
 }
 
-// Rounded UP to the next e4m3 value so amax / sf never exceeds the payload
-// range: round-to-nearest can land up to 33% low near e4m3's 2^-9 floor, which
-// would saturate the tile's largest values.
 __device__ __forceinline__ float nvfp4_tile_scale(const float* vals,
                                                   float max_val,
                                                   uint8_t* sf_out) {
@@ -121,15 +118,8 @@ __device__ __forceinline__ float nvfp4_tile_scale(const float* vals,
     amax = fmaxf(amax, fabsf(vals[i]));
   }
   const float sf_f = fmaxf(amax / max_val, 0.001953125f);  // 2^-9
-  __nv_fp8_e4m3 sf8(sf_f);
-  uint8_t sf_bits = *reinterpret_cast<const uint8_t*>(&sf8);
-  if (float(sf8) < sf_f && sf_bits < 0x7E) {
-    // Bump to the next e4m3 value (bit patterns of positive e4m3 values are
-    // monotonic; 0x7E = 448 is the max finite value).
-    sf_bits += 1;
-    sf8 = *reinterpret_cast<const __nv_fp8_e4m3*>(&sf_bits);
-  }
-  *sf_out = sf_bits;
+  const __nv_fp8_e4m3 sf8(sf_f);
+  *sf_out = *reinterpret_cast<const uint8_t*>(&sf8);
   return float(sf8);
 }
 
diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h
index a8beef5cda3d..331cd1fbb87f 100644
--- a/csrc/libtorch_stable/ops.h
+++ b/csrc/libtorch_stable/ops.h
@@ -175,9 +175,6 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 // DSV3 fused A GEMM: conditionally compiled so declaration and impl
 // registration are in the source file (dsv3_fused_a_gemm.cu)
 
-// AllSpark ops: declarations are in the source files
-// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
-
 #endif
 
 // CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
@@ -259,12 +256,20 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
                         torch::stable::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
+void fused_deepseek_v4_kv_rope_insert(
+    torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache,
+    torch::stable::Tensor const& slot_mapping,
+    torch::stable::Tensor const& position_ids,
+    torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size,
+    std::optional<torch::stable::Tensor> fp8_scale, bool kv_mxfp8);
+
 torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
     torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
     torch::stable::Tensor const& position_ids,
     torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
-    double eps, int64_t cache_block_size, bool apply_q_norm);
+    double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8,
+    bool apply_q_rope, bool is_q_interleaved);
 
 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
     torch::stable::Tensor& q, torch::stable::Tensor const& kv,
@@ -654,12 +659,61 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c,
                           const std::string& kv_cache_dtype,
                           torch::stable::Tensor& scale);
 
-void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c,
-                                  torch::stable::Tensor& k_pe,
-                                  torch::stable::Tensor& kv_cache_ptrs,
-                                  torch::stable::Tensor& slot_mapping,
-                                  int64_t block_size, int64_t block_stride,
-                                  int64_t entry_stride);
+void concat_and_cache_mla_grouped(
+    torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe,
+    torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping,
+    int64_t block_size, int64_t block_stride, int64_t entry_stride,
+    std::optional<torch::stable::Tensor> kv_scales,
+    const std::string& kv_cache_dtype);
+
+#ifndef USE_ROCM
+// HiSparse kernels use raw PTX in the row copy; CUDA-only.
+void hisparse_resolve_residency(
+    torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
+    torch::stable::Tensor const& hot_block_table,
+    torch::stable::Tensor const& global_indices,
+    torch::stable::Tensor& hot_indices,
+    torch::stable::Tensor& device_global_indices,
+    torch::stable::Tensor& lru_slots,
+    std::optional<torch::stable::Tensor> const& request_state_indices,
+    int64_t region_stride,
+    std::optional<torch::stable::Tensor> const& miss_mask,
+    std::optional<torch::stable::Tensor> const& stats,
+    std::optional<torch::stable::Tensor> const& attention_indices,
+    int64_t attention_block_stride,
+    std::optional<torch::stable::Tensor> const& request_ids,
+    std::optional<torch::stable::Tensor> const& source_block_table,
+    int64_t source_block_size,
+    std::optional<torch::stable::Tensor> const& resolved_global_indices,
+    std::optional<torch::stable::Tensor> const& valid_counts,
+    std::optional<torch::stable::Tensor> const& swap_host_physical_rows,
+    std::optional<torch::stable::Tensor> const& swap_device_physical_rows,
+    std::optional<torch::stable::Tensor> const& swap_counts,
+    std::optional<torch::stable::Tensor> const& resident_block_table,
+    int64_t resident_block_size, int64_t resident_null_block);
+
+void hisparse_invalidate_written_slots(
+    torch::stable::Tensor& device_global_indices,
+    torch::stable::Tensor const& request_state_indices,
+    torch::stable::Tensor const& req_id_per_token,
+    torch::stable::Tensor const& written_slots);
+
+void hisparse_gather_plan(
+    torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
+    torch::stable::Tensor const& global_indices,
+    torch::stable::Tensor const& hot_indices,
+    torch::stable::Tensor const& miss_mask,
+    std::optional<torch::stable::Tensor> const& request_state_indices,
+    std::optional<torch::stable::Tensor> const& attention_indices,
+    int64_t attention_block_stride);
+
+void hisparse_gather_compact(torch::stable::Tensor const& host_cache,
+                             torch::stable::Tensor& hot_cache,
+                             torch::stable::Tensor const& miss_global_indices,
+                             torch::stable::Tensor const& miss_hot_indices,
+                             torch::stable::Tensor const& miss_counts);
+
+#endif  // !USE_ROCM
 
 // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
 void concat_and_cache_mla_rope_fused(
@@ -704,7 +758,10 @@ void cp_gather_and_upconvert_fp8_kv_cache(
     torch::stable::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
     torch::stable::Tensor const& workspace_starts,  // [BATCH]
     int64_t batch_size,
-    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt,
+    std::optional<torch::stable::Tensor> host_cache = std::nullopt,
+    std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
+    std::optional<torch::stable::Tensor> device_row_ids = std::nullopt);
 
 // Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace
 void cp_gather_and_upconvert_nvfp4_kv_cache(
diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh
index 9ae0fd4bcbaa..bf5b46ab1875 100644
--- a/csrc/libtorch_stable/persistent_topk.cuh
+++ b/csrc/libtorch_stable/persistent_topk.cuh
@@ -1030,41 +1030,34 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
 constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
     sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE;  // 128KB
 
-/*!
- * \brief Filtered Top-K kernel for ragged sequences.
- *
- * \tparam DType Data type (float, half, nv_bfloat16)
- * \tparam IdType Index type (int32_t)
- * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
- */
-template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
-          bool UsePredicatedShortLoads = false>
-__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
-    FilteredTopKUnifiedKernel(const DType* __restrict__ input,
-                              IdType* __restrict__ output,
-                              const IdType* __restrict__ lengths,
-                              uint32_t num_rows, uint32_t top_k,
-                              uint32_t max_len) {
+template <uint32_t MAX_K>
+struct FilteredTopKStorage {
+  alignas(128) int histogram[2][256 + 128];
+  alignas(128) int counter;
+  alignas(128) int threshold_bin;
+  alignas(128) int num_input[2];
+  alignas(128) int indices[MAX_K];
+  int last_remain;
+};
+
+// With CheckOverflow, return false before using a truncated stash so the caller
+// can retry with exact full-row selection.
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K,
+          bool UsePredicatedShortLoads, bool CheckOverflow = false>
+__device__ __forceinline__ bool filtered_topk_row(
+    const DType* score, IdType* dst, int length, uint32_t top_k,
+    FilteredTopKStorage<MAX_K>& storage) {
   constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS;
   constexpr int RADIX = 256;
   constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE;
-
-  const uint32_t bid = blockIdx.x;
   const int tx = threadIdx.x;
 
-  if (bid >= num_rows) return;
-
-  const int length =
-      (lengths != nullptr) ? lengths[bid] : static_cast<int>(max_len);
-  const DType* score = input + bid * max_len;
-  IdType* dst = output + bid * top_k;
-
   // Trivial case: length <= top_k
   if (length <= static_cast<int>(top_k)) {
     for (int i = tx; i < static_cast<int>(top_k); i += BLOCK_SIZE) {
       dst[i] = (i < length) ? static_cast<IdType>(i) : static_cast<IdType>(-1);
     }
-    return;
+    return true;
   }
 
   // Short path
@@ -1077,16 +1070,14 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
       hist4096::histogram_4096_topk<MAX_K, 12, 8>(score, dst, length,
                                                   _smem_reg);
     }
-    return;
+    return true;
   }
 
-  // Static shared memory
-  alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
-  alignas(128) __shared__ int s_counter;
-  alignas(128) __shared__ int s_threshold_bin_id;
-  alignas(128) __shared__ int s_num_input[2];
-  alignas(128) __shared__ int s_indices[MAX_K];
-
+  auto& s_histogram_buf = storage.histogram;
+  auto& s_counter = storage.counter;
+  auto& s_threshold_bin_id = storage.threshold_bin;
+  auto& s_num_input = storage.num_input;
+  auto& s_indices = storage.indices;
   auto& s_histogram = s_histogram_buf[0];
 
   // Dynamic shared memory for input double buffer
@@ -1213,10 +1204,13 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     // Stage 2: refine with 8bit radix passes
 #pragma unroll
     for (int round = 0; round < NUM_ROUNDS; ++round) {
-      __shared__ int s_last_remain;
+      auto& s_last_remain = storage.last_remain;
       const auto r_idx = round % 2;
 
       const auto _raw_num_input = s_num_input[r_idx];
+      if constexpr (CheckOverflow) {
+        if (_raw_num_input > SMEM_INPUT_SIZE) return false;
+      }
       const auto num_input =
           (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE;
 
@@ -1284,6 +1278,30 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
     const int idx = s_indices[base];
     dst[base] = static_cast<IdType>(idx);
   }
+  return true;
+}
+
+/*!
+ * \brief Filtered Top-K kernel for ragged sequences.
+ *
+ * \tparam DType Data type (float, half, nv_bfloat16)
+ * \tparam IdType Index type (int32_t)
+ * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
+ */
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
+          bool UsePredicatedShortLoads = false>
+__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
+    FilteredTopKUnifiedKernel(const DType* __restrict__ input,
+                              IdType* __restrict__ output,
+                              const IdType* __restrict__ lengths,
+                              uint32_t num_rows, uint32_t top_k,
+                              uint32_t max_len) {
+  const uint32_t bid = blockIdx.x;
+  if (bid >= num_rows) return;
+  const int length = lengths ? lengths[bid] : static_cast<int>(max_len);
+  __shared__ FilteredTopKStorage<MAX_K> storage;
+  filtered_topk_row<DType, IdType, VEC_SIZE, MAX_K, UsePredicatedShortLoads>(
+      input + bid * max_len, output + bid * top_k, length, top_k, storage);
 }
 
 // Helper to compute GCD for VEC_SIZE selection
diff --git a/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh b/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh
index 3d82e4393166..abe3368cb16f 100644
--- a/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh
+++ b/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh
@@ -74,18 +74,9 @@ class MatrixView_half_rw {
                                                 const int width)
       : data(data), height(height), width(width) {}
 
-  __device__ __forceinline__ half item(int row, int column) const {
-    return data[row * width + column];
-  }
   __device__ __forceinline__ half* item_ptr(int row, int column) {
     return &data[row * width + column];
   }
-  __device__ __forceinline__ void set(int row, int column, half value) {
-    data[row * width + column] = value;
-  }
-  __device__ __forceinline__ void set_half2(int row, int column, half2 value) {
-    ((half2*)data)[(row * width + column) / 2] = value;
-  }
 
   __device__ __forceinline__ void set4(int row, int column, half v0, half v1,
                                        half v2, half v3) {
@@ -113,14 +104,6 @@ class MatrixView_q4_row {
     return (data[row * width / 8 + column / 8] >> shift) & 0x0f;
   }
 
-  __device__ __forceinline__ void item2(int (&items)[2], int row,
-                                        int column) const {
-    int shift = (column & 0x07) * 4;
-    uint32_t d = data[row * width / 8 + column / 8] >> shift;
-    items[0] = d & 0x0f;
-    items[1] = (d >> 4) & 0x0f;
-  }
-
   __device__ __forceinline__ void item4(int (&items)[4], int row,
                                         int column) const {
     int shift = (column & 0x07) * 4;
@@ -148,14 +131,6 @@ class MatrixView_q2_row {
     return (data[row * width / 16 + column / 16] >> shift) & 0x03;
   }
 
-  __device__ __forceinline__ void item2(int (&items)[2], int row,
-                                        int column) const {
-    int shift = (column & 0x0f) * 2;
-    uint32_t d = data[row * width / 16 + column / 16] >> shift;
-    items[0] = d & 0x03;
-    items[1] = (d >> 2) & 0x03;
-  }
-
   __device__ __forceinline__ void item4(int (&items)[4], int row,
                                         int column) const {
     int shift = (column & 0x0f) * 2;
@@ -237,14 +212,6 @@ class MatrixView_q8_row {
     return (data[row * width / 4 + column / 4] >> shift) & 0xff;
   }
 
-  __device__ __forceinline__ void item2(int (&items)[2], int row,
-                                        int column) const {
-    int shift = (column & 0x03) * 8;
-    uint32_t d = data[row * width / 4 + column / 4] >> shift;
-    items[0] = d & 0xff;
-    items[1] = (d >> 8) & 0xff;
-  }
-
   __device__ __forceinline__ void item4(int (&items)[4], int row,
                                         int column) const {
     int shift = (column & 0x03) * 2;
diff --git a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu
deleted file mode 100644
index 96dc3ecfc860..000000000000
--- a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu
+++ /dev/null
@@ -1,1024 +0,0 @@
-#include "allspark_utils.cuh"
-
-#include <torch/csrc/stable/c/shim.h>
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/core/ScalarType.h>
-
-#include <cublas_v2.h>
-
-#include "core/registration.h"
-#include "libtorch_stable/torch_utils.h"
-
-torch::stable::Tensor as_g_workspace;
-
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
-
-torch::stable::Tensor allspark_w8a16_gemm(
-    torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
-    torch::stable::Tensor const& b_scales,
-    std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
-    int64_t group_size, int64_t sm_count, int64_t sm_version,
-    int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
-  STD_TORCH_CHECK_NOT_IMPLEMENTED(
-      false, "allspark_w8a16_gemm(..) requires CUDA_ARCH >= 8.0");
-  return torch::stable::empty({1, 1});
-}
-
-#else
-namespace allspark {
-/*
- * GemmTile manage data movement from Global Memory to Shared Memory
- * requiring N % 8 == 0， K % 16 == 0 by loading uint
- * BN is obtained by padding the original N to a multiple of 32
- * weight B is rearranged as N32K16 order,
- * i.e. a initial data block of size 32(n)x16(k) is reordered as n8k4n4k4，
- * in order to put data loaded by the same thread of 32x16 data block together
- * continuously (see
- * https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type)
- */
-template <typename FType, typename QType, int Mtile, int Ntile, int NStage,
-          int BLOCK>
-struct GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
-  // element num loaded by a LDG inst.
-  static constexpr int LDG_ELEMENT_CNT_A = 8;
-  static constexpr int LDG_ELEMENT_CNT_B = 16;
-  static constexpr int WARP_SIZE = 32;
-  static constexpr int M_SIZE_ONE_LOAD = (BLOCK * LDG_ELEMENT_CNT_A) / 32;
-  static constexpr int N_SIZE_ONE_LOAD = (BLOCK * LDG_ELEMENT_CNT_B) / 32;
-
-  __device__ GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK(
-      const SM8x_GEMM_W8A16_Splitk_Params<FType, QType>& k_params,
-      const uint32_t& A_smem_addr, const uint32_t& BQ_smem_addr,
-      const uint32_t& A_stage_stride, const uint32_t& BQ_stage_stride)
-      : params(k_params),
-        A_smem_base_addr(A_smem_addr),
-        BQ_smem_base_addr(BQ_smem_addr),
-        A_smem_stage_stride(A_stage_stride),
-        BQ_smem_stage_stride(BQ_stage_stride) {
-    this_block_A_base_ptr = params.A_ptr + blockIdx.x * Mtile * params.K +
-                            blockIdx.z * params.SplitK;
-    // here B is rearranged as N32K16 order, i.e. 4 continuous N-direction
-    // 8(N)x16(K) size data blocks are packed together
-    this_block_B_base_ptr = params.B_ptr + blockIdx.y * Ntile * params.K +
-                            blockIdx.z * params.SplitK * 4;
-
-    const auto lane_id = threadIdx.x % WARP_SIZE;
-
-    // For matrix A, a block load/store Mtile(row) x 32(col) elements in
-    // multiple iters, 8x4 warp load/store 8(row) x 32(col) elements per iter
-    const auto Aldg_row_base_idx = threadIdx.x / 4;
-    Aldg_col_idx = (threadIdx.x % 4) * LDG_ELEMENT_CNT_A;
-    const int Aldg_base_offset = Aldg_row_base_idx * params.K + Aldg_col_idx;
-
-    // For matrix B, a block load/store elements of (Ntile / 4) row x 128 col
-    // elements of N32K16 packing in multiple iters, 4x8 warp load/store 4(row)
-    // * 128(col) per iter
-    Bldg_col_idx = (threadIdx.x % 8) * LDG_ELEMENT_CNT_B;
-    const auto Bldg_row_base_idx = threadIdx.x / 8;
-    const int Bldg_base_offset =
-        Bldg_row_base_idx * params.K * 4 + Bldg_col_idx;
-
-    this_block_A_base_ptr += Aldg_base_offset;
-    this_block_B_base_ptr += Bldg_base_offset;
-
-    const int sts_a_base_offset =
-        (threadIdx.x / 4) * 32 +
-        ((lane_id % 4) ^ ((lane_id / 4) % 4) ^ ((lane_id / 4) / 4)) *
-            LDG_ELEMENT_CNT_A;
-    const int sts_bq_base_offset =
-        Bldg_row_base_idx * 32 * 4 +
-        ((threadIdx.x % 8) ^ (((threadIdx.x / 8) % 2) * 4)) * LDG_ELEMENT_CNT_B;
-
-    A_smem_base_addr += sts_a_base_offset * sizeof(FType);
-    BQ_smem_base_addr += sts_bq_base_offset * sizeof(uint8_t);
-
-    A_ldg_guard = 0;
-    B_ldg_guard = 0;
-  #pragma unroll
-    for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD; ++i) {
-      auto m_idx = blockIdx.x * Mtile + Aldg_row_base_idx + i * M_SIZE_ONE_LOAD;
-      if (m_idx < params.M) {
-        A_ldg_guard |= (1u << i);
-      }
-    }
-
-    const int N_padded = (params.N + 31) / 32 * 32;
-  #pragma unroll
-    for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD; ++i) {
-      auto n_idx = blockIdx.y * Ntile + (Bldg_row_base_idx / 8) * 32 +
-                   i * N_SIZE_ONE_LOAD;
-      if (n_idx < N_padded) {
-        B_ldg_guard |= (1u << i);
-      }
-    }
-  }
-
-  __device__ void ldgsts_first_ktiles(const int& first_k_tile,
-                                      const int& k_tiles) {
-    // load first k_tile
-    // load A
-    const int A_src_size = Aldg_col_idx < first_k_tile ? 16 : 0;
-  #pragma unroll
-    for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD; ++i) {
-      cp_async<16>(
-          A_smem_base_addr + (i * M_SIZE_ONE_LOAD * 32) * sizeof(FType),
-          this_block_A_base_ptr + i * M_SIZE_ONE_LOAD * params.K, A_src_size,
-          (A_ldg_guard & (1u << i)) != 0);
-    }
-
-    // load B
-    const int B_src_size = (Bldg_col_idx / 4) < first_k_tile ? 16 : 0;
-  #pragma unroll
-    for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD; ++i) {
-      cp_async<16>(
-          BQ_smem_base_addr + (i * N_SIZE_ONE_LOAD * 32) * sizeof(uint8_t),
-          this_block_B_base_ptr + i * N_SIZE_ONE_LOAD * params.K, B_src_size,
-          (B_ldg_guard & (1u << i)) != 0);
-    }
-
-    cp_async_commit_group();
-    this_block_A_base_ptr += first_k_tile;
-    this_block_B_base_ptr += (first_k_tile * 4);
-
-    // load second to (N-stage - 1) k_tiles
-    for (int stage_idx = 1; stage_idx < NStage - 1; ++stage_idx) {
-      if (stage_idx < k_tiles) {
-  #pragma unroll
-        for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD;
-             ++i) {
-          cp_async<16>(A_smem_base_addr + stage_idx * A_smem_stage_stride +
-                           (i * M_SIZE_ONE_LOAD * 32) * sizeof(FType),
-                       this_block_A_base_ptr + i * M_SIZE_ONE_LOAD * params.K,
-                       16, (A_ldg_guard & (1u << i)) != 0);
-        }
-
-  #pragma unroll
-        for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD;
-             ++i) {
-          cp_async<16>(BQ_smem_base_addr + stage_idx * BQ_smem_stage_stride +
-                           (i * N_SIZE_ONE_LOAD * 32) * sizeof(uint8_t),
-                       this_block_B_base_ptr + i * N_SIZE_ONE_LOAD * params.K,
-                       16, (B_ldg_guard & (1u << i)) != 0);
-        }
-
-        this_block_A_base_ptr += 32;
-        this_block_B_base_ptr += (32 * 4);
-      }
-      cp_async_commit_group();
-    }
-  }
-
-  __device__ void ldgsts(const int& sts_stage_idx) {
-    const int a_stage_offset = sts_stage_idx * A_smem_stage_stride;
-    const int bq_stage_offset = sts_stage_idx * BQ_smem_stage_stride;
-  #pragma unroll
-    for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD; ++i) {
-      cp_async<16>(A_smem_base_addr + a_stage_offset +
-                       (i * M_SIZE_ONE_LOAD * 32) * sizeof(FType),
-                   this_block_A_base_ptr + i * M_SIZE_ONE_LOAD * params.K, 16,
-                   (A_ldg_guard & (1u << i)) != 0);
-    }
-
-  #pragma unroll
-    for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD; ++i) {
-      cp_async<16>(BQ_smem_base_addr + bq_stage_offset +
-                       (i * N_SIZE_ONE_LOAD * 32) * sizeof(uint8_t),
-                   this_block_B_base_ptr + i * N_SIZE_ONE_LOAD * params.K, 16,
-                   (B_ldg_guard & (1u << i)) != 0);
-    }
-
-    cp_async_commit_group();
-    this_block_A_base_ptr += 32;
-    this_block_B_base_ptr += (32 * 4);
-  }
-
-  const FType* this_block_A_base_ptr = nullptr;
-  const QType* this_block_B_base_ptr = nullptr;
-
-  int Aldg_col_idx;
-  int Bldg_col_idx;
-
-  uint32_t A_ldg_guard;
-  uint32_t B_ldg_guard;
-
-  uint32_t A_smem_base_addr, BQ_smem_base_addr;
-  const uint32_t A_smem_stage_stride, BQ_smem_stage_stride;
-
-  const SM8x_GEMM_W8A16_Splitk_Params<FType, QType>& params;
-};
-
-/*
- * requiring N % 8 == 0
- */
-template <typename FType, typename QType, int Mtile, int Ntile, int BLOCK,
-          bool EnableFuse, bool has_zp>
-struct ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
-  static constexpr int WARP_SIZE = 32;
-  static constexpr int WARP_CNT = BLOCK / WARP_SIZE;
-  static constexpr int WARP_NTILE = Ntile / WARP_CNT;
-  static constexpr int WARP_NITER = WARP_NTILE / 8;  // hmma16816
-  static_assert(WARP_NTILE == 32 or WARP_NTILE == 64,
-                "now only support WARP_NTILE = 32 or 64!");
-
-  __device__ ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK(
-      const SM8x_GEMM_W8A16_Splitk_Params<FType, QType>& k_params,
-      const uint32_t& A_smem_addr, const uint32_t& BQ_smem_addr,
-      const uint32_t& A_stage_stride, const uint32_t& BQ_stage_stride)
-      : params(k_params),
-        A_smem_base_addr(A_smem_addr),
-        BQ_smem_base_addr(BQ_smem_addr),
-        A_smem_stage_stride(A_stage_stride),
-        BQ_smem_stage_stride(BQ_stage_stride) {
-    warp_id = threadIdx.x / WARP_SIZE;
-    lane_id = threadIdx.x % WARP_SIZE;
-
-    load_a_base_offset[0] =
-        (lane_id % 16) * 32 +
-        ((lane_id / 16) ^ (lane_id % 4) ^ ((lane_id / 4) % 2)) * 8;
-    load_a_base_offset[1] =
-        (lane_id % 16) * 32 +
-        ((lane_id / 16 + 2) ^ (lane_id % 4) ^ ((lane_id / 4) % 2)) * 8;
-
-    load_b_base_offset[0] =
-        (lane_id / 4 + warp_id * (WARP_NTILE / 4)) * 32 * 4 +
-        (lane_id % 4) * 16 + ((lane_id / 4) % 2) * 16 * 4;
-    load_b_base_offset[1] =
-        (lane_id / 4 + warp_id * (WARP_NTILE / 4)) * 32 * 4 +
-        (lane_id % 4) * 16 + (((lane_id / 4) % 2) ^ 1) * 16 * 4;
-
-    sts_c_base_offset = warp_id * Mtile * WARP_NTILE +
-                        (lane_id / 4) * WARP_NTILE + (lane_id % 4) * 2;
-
-    if (EnableFuse) {
-      this_block_C_base_ptr =
-          params.C_ptr + blockIdx.x * Mtile * params.N + blockIdx.y * Ntile;
-    } else {
-      this_block_C_base_ptr =
-          params.C_split_ptr + blockIdx.z * params.M * params.N +
-          blockIdx.x * Mtile * params.N + blockIdx.y * Ntile;
-    }
-    int store_thds_in_row = WARP_NTILE / 8;
-    store_c_row_base_idx = lane_id / store_thds_in_row;
-    store_c_col_idx = warp_id * WARP_NTILE + (lane_id % store_thds_in_row) * 8;
-    store_c_base_offset = store_c_row_base_idx * params.N + store_c_col_idx;
-
-  #pragma unroll
-    for (int i = 0; i < Mtile / 16; ++i) {
-  #pragma unroll
-      for (int j = 0; j < WARP_NITER; ++j) {
-  #pragma unroll
-        for (int k = 0; k < 4; ++k) {
-          C_frag[i][j][k] = 0.f;
-        }
-      }
-    }
-    params_n_idx =
-        blockIdx.y * Ntile + warp_id * WARP_NTILE + (lane_id / 4) * 4;
-  }
-
-  __device__ void lds(const int& smem_stage_idx, const int& reg_buf_idx,
-                      const int& k_phase_idx) {
-    uint32_t A_smem_addr =
-        A_smem_base_addr + A_smem_stage_stride * smem_stage_idx;
-    uint32_t B_smem_addr =
-        BQ_smem_base_addr + BQ_smem_stage_stride * smem_stage_idx;
-
-  #pragma unroll
-    for (int i = 0; i < Mtile / 16; ++i) {
-      ldsm_4(A_frag[reg_buf_idx][i][0], A_frag[reg_buf_idx][i][1],
-             A_frag[reg_buf_idx][i][2], A_frag[reg_buf_idx][i][3],
-             A_smem_addr + (load_a_base_offset[k_phase_idx] + i * 16 * 32) *
-                               sizeof(FType));
-    }
-  #pragma unroll
-    for (int i = 0; i < WARP_NTILE / 32; ++i) {
-      lds128(BQ_frag[reg_buf_idx][4 * i + 0], BQ_frag[reg_buf_idx][4 * i + 1],
-             BQ_frag[reg_buf_idx][4 * i + 2], BQ_frag[reg_buf_idx][4 * i + 3],
-             B_smem_addr + (load_b_base_offset[k_phase_idx] + i * 32 * 32) *
-                               sizeof(uint8_t));
-    }
-
-  // dequant B
-  #pragma unroll
-    for (int i = 0; i < WARP_NITER / 2; ++i) {
-      cvt_8bx4_to_16bx4_bias128(BQ_frag[reg_buf_idx][2 * i],
-                                BF_frag[reg_buf_idx][2 * i]);
-      if (has_zp) {
-        BF_frag[reg_buf_idx][2 * i][0] =
-            __hsub2(BF_frag[reg_buf_idx][2 * i][0], num2num2(B_zero[i].x));
-        BF_frag[reg_buf_idx][2 * i][1] =
-            __hsub2(BF_frag[reg_buf_idx][2 * i][1], num2num2(B_zero[i].x));
-      }
-
-      BF_frag[reg_buf_idx][2 * i][0] =
-          __hmul2(BF_frag[reg_buf_idx][2 * i][0], num2num2(B_scale[i].x));
-      BF_frag[reg_buf_idx][2 * i][1] =
-          __hmul2(BF_frag[reg_buf_idx][2 * i][1], num2num2(B_scale[i].x));
-
-      cvt_8bx4_to_16bx4_bias128(BQ_frag[reg_buf_idx][2 * i + 1],
-                                BF_frag[reg_buf_idx][2 * i + 1]);
-      if (has_zp) {
-        BF_frag[reg_buf_idx][2 * i + 1][0] =
-            __hsub2(BF_frag[reg_buf_idx][2 * i + 1][0], num2num2(B_zero[i].y));
-        BF_frag[reg_buf_idx][2 * i + 1][1] =
-            __hsub2(BF_frag[reg_buf_idx][2 * i + 1][1], num2num2(B_zero[i].y));
-      }
-
-      BF_frag[reg_buf_idx][2 * i + 1][0] =
-          __hmul2(BF_frag[reg_buf_idx][2 * i + 1][0], num2num2(B_scale[i].y));
-      BF_frag[reg_buf_idx][2 * i + 1][1] =
-          __hmul2(BF_frag[reg_buf_idx][2 * i + 1][1], num2num2(B_scale[i].y));
-    }
-  }
-
-  __device__ void ldg_params() {
-    const int N_padded = (params.N + 31) / 32 * 32;
-    // load B scale and zero_point
-  #pragma unroll
-    for (int i = 0; i < WARP_NTILE / 32; ++i) {
-      ldg64_ca(B_scale[2 * i + 0], B_scale[2 * i + 1],
-               params.B_scale_ptr + params_n_idx + i * 32,
-               (params_n_idx + i * 32) < N_padded);
-      if (has_zp) {
-        ldg64_ca(B_zero[2 * i + 0], B_zero[2 * i + 1],
-                 params.B_zero_ptr + params_n_idx + i * 32,
-                 (params_n_idx + i * 32) < N_padded);
-      }
-    }
-  }
-
-  __device__ void mma(const int& reg_buf_idx) {
-  #pragma unroll
-    for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
-  #pragma unroll
-      for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
-        hmma16816_f32<FType>(
-            C_frag[m_idx][n_idx], A_frag[reg_buf_idx][m_idx],
-            reinterpret_cast<uint32_t (&)[2]>(BF_frag[reg_buf_idx][n_idx]));
-      }
-    }
-  }
-
-  __device__ void fused_splitk_reduce() {
-    // need splitk-reduce if enable splitk
-    if (gridDim.z > 1) {
-      auto blk_red_idx = blockIdx.x * gridDim.y + blockIdx.y;
-      // Wait for all previous blocks in the splitk direction to accumulate the
-      // results into C_tmp
-      if (threadIdx.x == 0) {
-        uint32_t* red_count_ptr = params.red_count_ptr + blk_red_idx;
-        uint32_t count;
-        do {
-          // make sure the ld.cg inside the do-wile loop
-          __threadfence_block();
-          asm volatile("ld.global.cg.b32 %0, [%1];"
-                       : "=r"(count)
-                       : "l"(red_count_ptr));
-        } while (count != blockIdx.z);
-      }
-      __syncthreads();
-
-      auto C_tmp_base_offset = blk_red_idx * Mtile * Ntile + threadIdx.x * 4;
-      if (blockIdx.z != 0) {
-        // expecting that temporary register here reuses the previous A&B frag
-        // register
-        float temp_frag[Mtile / 16][WARP_NITER][4];
-  #pragma unroll
-        for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
-  #pragma unroll
-          for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
-            int offset =
-                C_tmp_base_offset + (m_idx * WARP_NITER + n_idx) * BLOCK * 4;
-            *reinterpret_cast<int4*>(temp_frag[m_idx][n_idx]) =
-                *reinterpret_cast<int4*>(params.C_tmp_ptr + offset);
-          }
-        }
-  #pragma unroll
-        for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
-  #pragma unroll
-          for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
-  #pragma unroll
-            for (int idx = 0; idx < 4; ++idx) {
-              C_frag[m_idx][n_idx][idx] += temp_frag[m_idx][n_idx][idx];
-            }
-          }
-        }
-      }
-
-      // first splitk - 1 blocks need to write partial results into C_tmp
-      if (blockIdx.z != gridDim.z - 1) {
-  #pragma unroll
-        for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
-  #pragma unroll
-          for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
-            int offset =
-                C_tmp_base_offset + (m_idx * WARP_NITER + n_idx) * BLOCK * 4;
-            asm volatile(
-                "{st.global.cg.v4.b32 [%0], {%1, %2, %3, %4};}\n"
-                :
-                : "l"(params.C_tmp_ptr + offset), "f"(C_frag[m_idx][n_idx][0]),
-                  "f"(C_frag[m_idx][n_idx][1]), "f"(C_frag[m_idx][n_idx][2]),
-                  "f"(C_frag[m_idx][n_idx][3]));
-          }
-        }
-        __threadfence();
-        __syncthreads();
-        if (threadIdx.x == 0) {
-          uint32_t* red_count_ptr = params.red_count_ptr + blk_red_idx;
-          atomicInc(red_count_ptr, gridDim.z);
-        }
-      }
-    }
-  }
-
-  __device__ void stg(char* smem) {
-    if (EnableFuse) {
-      if (blockIdx.z != gridDim.z - 1) return;
-    }
-    uint32_t* C_sts_ptr =
-        reinterpret_cast<uint32_t*>(smem + sts_c_base_offset * sizeof(FType));
-    // C_tile sts
-  #pragma unroll
-    for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
-  #pragma unroll
-      for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
-  #pragma unroll
-        for (int k_idx = 0; k_idx < 2; ++k_idx) {
-          FType low16 = MarlinScalarType2<FType>::float2num(
-              C_frag[m_idx][n_idx][k_idx * 2]);
-          FType high16 = MarlinScalarType2<FType>::float2num(
-              C_frag[m_idx][n_idx][k_idx * 2 + 1]);
-          uint32_t tmp = (reinterpret_cast<uint32_t&>(low16) & 0xffff) |
-                         (reinterpret_cast<uint32_t&>(high16) << 16);
-          int sts_offset =
-              m_idx * 16 * (WARP_NTILE / 2) +
-              (((lane_id / (32 / WARP_NITER)) + n_idx) % WARP_NITER) * (8 / 2) +
-              k_idx * 8 * (WARP_NTILE / 2);
-          C_sts_ptr[sts_offset] = tmp;
-        }
-      }
-    }
-
-    __syncthreads();
-
-    FType* C_base_ptr = this_block_C_base_ptr + store_c_base_offset;
-    // C_tile lds and stg
-    auto m_base_idx = store_c_row_base_idx + blockIdx.x * Mtile;
-    bool n_guard = (store_c_col_idx + blockIdx.y * Ntile) < params.N;
-    if (WARP_NTILE == 32) {
-      int lds_c_base_offset = warp_id * Mtile * WARP_NTILE +
-                              (lane_id / 4) * WARP_NTILE +
-                              ((lane_id % 4 + lane_id / 8) % 4) * 8;
-      uint4* C_lds_ptr =
-          reinterpret_cast<uint4*>(smem + lds_c_base_offset * sizeof(FType));
-  #pragma unroll
-      for (int i = 0; i < (Mtile / 16) * (WARP_NITER / 2); ++i) {
-        uint4 stg_reg = C_lds_ptr[i * 8 * 4];
-        stg128(stg_reg.x, stg_reg.y, stg_reg.z, stg_reg.w,
-               C_base_ptr + i * 8 * params.N,
-               (m_base_idx + i * 8) < params.M && n_guard);
-      }
-    } else if (WARP_NTILE == 64) {
-      int lds_c_base_offset =
-          warp_id * Mtile * WARP_NTILE + (lane_id / 8) * WARP_NTILE;
-  #pragma unroll
-      for (int i = 0; i < (Mtile / 16) * (WARP_NITER / 2); ++i) {
-        int lds_c_offset = lds_c_base_offset + i * 4 * WARP_NTILE +
-                           ((lane_id % 8 + lane_id / 8 + (i % 2) * 4) % 8) * 8;
-        uint4 stg_reg =
-            *reinterpret_cast<uint4*>(smem + lds_c_offset * sizeof(FType));
-        stg128(stg_reg.x, stg_reg.y, stg_reg.z, stg_reg.w,
-               C_base_ptr + i * 4 * params.N,
-               (m_base_idx + i * 4) < params.M && n_guard);
-      }
-    }
-  }
-
-  const SM8x_GEMM_W8A16_Splitk_Params<FType, QType>& params;
-
-  int load_a_base_offset[2];
-  int load_b_base_offset[2];
-  int sts_c_base_offset;
-
-  int store_c_base_offset;
-
-  int store_c_row_base_idx, store_c_col_idx;
-  FType* this_block_C_base_ptr = nullptr;
-
-  int params_n_idx;
-  const uint32_t A_smem_base_addr, BQ_smem_base_addr;
-  const uint32_t A_smem_stage_stride, BQ_smem_stage_stride;
-
-  int lane_id;
-  int warp_id;
-  // first 2 denotes double buffer, second dim denotes M direction
-  uint32_t A_frag[2][Mtile / 16][4];
-
-  typename HalfType<FType>::T2 B_scale[WARP_NITER / 2];
-  typename HalfType<FType>::T2 B_zero[WARP_NITER / 2];
-  uint32_t BQ_frag[2][WARP_NITER];
-  // first 2 denotes double buffer, second dim denotes N direction, last 2
-  // denotes K direction
-  typename HalfType<FType>::T2 BF_frag[2][WARP_NITER][2];
-  // first dim denotes M direction, second dim denotes N direction
-  float C_frag[Mtile / 16][WARP_NITER][4];
-};
-
-/*
- *  @brief W8A16 Perchannel Quantization GEMM,
- *         requires N % 8 == 0, K % 16 == 0
- *         accumulator precision: FP32
- *  @tparam FType: DataType for A, B_scale, B_zero, and C, supports half or
- * nv_bfloat16
- *  @tparam QType: DataType for B, support uint8(bias128)
- *  @tparam Mtile: M-dimensional size of the gemm block tile, supports 16, 32,
- * 48 or 64
- *  @tparam Ntile: N-dimensional size of the gemm block tile, supports 128 or
- * 256
- *  @tparam NStage: Num of stages for async copy
- *  @tparam BLOCK: BLOCK size
- *  @tparam EnableFuse: If true, use fused splitk-reduce, otherwise use
- * non-fused splitk-reduce
- *  @tparam has_zp: whether to use zero_point
- *
- *  @fparam params struct consists of following parameters:
- *      @param A_ptr: Matrix A value ptr, A = (M, K)
- *      @param B_ptr: Matrix B value ptr, B = (N32_align, K) (N32K16 special
- * format), N32_align = (N + 32 - 1) / 32 * 32
- *      @param B_scale_ptr: B_scale value ptr, B_scale = (N32_align,) (N32K16
- * special format)
- *      @param B_zero_ptr: B_zero value ptr, B_zero = (N32_align,) (N32K16
- * special format)
- *      @param C_ptr: Matrix C value ptr, C = (M, N)
- *      @param M: dimnesion m
- *      @param N: dimnesion n
- *      @param K: dimnesion k
- *      @param SplitK: split size along K-dimension
- *      @param C_split_ptr: Matrix C_split value ptr, used only in non-fused
- * splitk-reduce
- *      @param C_tmp_ptr: Matrix C_tmp value ptr, used only in fused
- * splitk-reduce
- *      @param red_count_ptr: 1-D red_count value ptr, used only in fused
- * splitk-reduce
- */
-template <typename FType, typename QType, int Mtile, int Ntile, int NStage,
-          int BLOCK, bool EnableFuse, bool has_zp>
-__global__ void __launch_bounds__(BLOCK)
-    ampere_hgemm_W8A16_perc_f16_f16_MtilexNtilex32_hmma16816_multistage_AN_BTN32K16_CN_splitk_kernel(
-        const SM8x_GEMM_W8A16_Splitk_Params<FType, QType> params) {
-  // A smem size = 64 * 32 * 2B/elem * 4(stage) = 16KB
-  // B smem size = 128 * 32 * 1B/elem * 4(stage) = 16KB
-  constexpr int smem_size_one_stage = Mtile * 32 * 2 + Ntile * 32;
-  __shared__ char smem[NStage * smem_size_one_stage];
-  char* A_smem = smem;
-  char* BQ_smem = smem + Mtile * 32 * 2 * NStage;
-
-  uint32_t A_smem_addr = smem_u32addr(A_smem);
-  uint32_t BQ_smem_addr = smem_u32addr(BQ_smem);
-  uint32_t A_smem_stage_stride = Mtile * 32 * 2;
-  uint32_t BQ_smem_stage_stride = Ntile * 32;
-
-  // initialize the data move process from GM to SMEM for this block
-  GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK<
-      FType, QType, Mtile, Ntile, NStage, BLOCK>
-      gmem_tile(params, A_smem_addr, BQ_smem_addr, A_smem_stage_stride,
-                BQ_smem_stage_stride);
-
-  int sts_stage_idx = 0;
-  int lds_stage_idx = 0;
-
-  auto tb_k_slice = blockIdx.z * params.SplitK + params.SplitK <= params.K
-                        ? params.SplitK
-                        : params.K - blockIdx.z * params.SplitK;
-  int k_tiles = (tb_k_slice + 31) / 32;
-  int first_k_tile = tb_k_slice - (k_tiles - 1) * 32;
-
-  // load first three tiles to shared memory
-  gmem_tile.ldgsts_first_ktiles(first_k_tile, k_tiles);
-  sts_stage_idx += (NStage - 2);
-  ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK<
-      FType, QType, Mtile, Ntile, BLOCK, EnableFuse, has_zp>
-      compute_tile(params, A_smem_addr, BQ_smem_addr, A_smem_stage_stride,
-                   BQ_smem_stage_stride);
-  compute_tile.ldg_params();
-  cp_asyc_wait_group<NStage - 2>();
-  __syncthreads();
-
-  compute_tile.lds(lds_stage_idx, 0, 0);
-  int reg_buf_idx = 1;
-
-  // main loop
-  for (; k_tiles > NStage - 1; --k_tiles) {
-    // load next A&B tile
-    sts_stage_idx = sts_stage_idx < NStage - 1 ? sts_stage_idx + 1 : 0;
-    gmem_tile.ldgsts(sts_stage_idx);
-
-  #pragma unroll
-    for (int k_phase_idx = 0; k_phase_idx < 2; k_phase_idx++) {
-      // dequantize next B tile
-      if (k_phase_idx == 1) {
-        cp_asyc_wait_group<NStage - 2>();
-        __syncthreads();
-        lds_stage_idx = lds_stage_idx < NStage - 1 ? lds_stage_idx + 1 : 0;
-      }
-
-      compute_tile.lds(lds_stage_idx, reg_buf_idx, (k_phase_idx + 1) % 2);
-
-      compute_tile.mma(reg_buf_idx ^ 1);
-      reg_buf_idx ^= 1;
-    }
-  }
-
-  // last NStage-1 tiles
-  for (; k_tiles > 0; --k_tiles) {
-    cp_async_commit_group();
-  #pragma unroll
-    for (int k_phase_idx = 0; k_phase_idx < 2; k_phase_idx++) {
-      // dequantize next B tile
-      if (k_phase_idx == 1) {
-        cp_asyc_wait_group<NStage - 2>();
-        __syncthreads();
-        lds_stage_idx = lds_stage_idx < NStage - 1 ? lds_stage_idx + 1 : 0;
-      }
-
-      compute_tile.lds(lds_stage_idx, reg_buf_idx, (k_phase_idx + 1) % 2);
-
-      compute_tile.mma(reg_buf_idx ^ 1);
-      reg_buf_idx ^= 1;
-    }
-  }
-
-  if (EnableFuse) {
-    compute_tile.fused_splitk_reduce();
-  }
-  compute_tile.stg(smem);
-}
-
-  #define __CALL_IF(MTILE, NTILE, NUM_THREADS, ENABLE_FUSE, HAS_ZP)                                     \
-    else if (Mtile == MTILE && Ntile == NTILE && BLOCK == NUM_THREADS &&                                \
-             enable_fuse == ENABLE_FUSE && has_zp == HAS_ZP) {                                          \
-      ampere_hgemm_W8A16_perc_f16_f16_MtilexNtilex32_hmma16816_multistage_AN_BTN32K16_CN_splitk_kernel< \
-          FType, QType, MTILE, NTILE, 4, NUM_THREADS, ENABLE_FUSE, HAS_ZP>                              \
-          <<<grid, block, 0, stream>>>(params);                                                         \
-    }
-
-template <typename FType, typename QType>
-void ampere_hgemm_W8A16_perc_f16_f16_MtilexNtilex32_mma16816_multistage_AN_BTN32K16_CN_splitk(
-    const FType* A, const QType* B, const FType* B_scale, const FType* B_zero,
-    FType* C, const int M, const int N, const int K, void* workspace,
-    const int sm_version, const BlockTileSplitkParams& fused_gemm_params,
-    cudaStream_t stream) {
-  int Mtile = fused_gemm_params.Mtile;
-  int grid_x = (M + Mtile - 1) / Mtile;
-  int Ntile = fused_gemm_params.Ntile;
-  int grid_y = (N + Ntile - 1) / Ntile;
-  int SplitK = fused_gemm_params.SplitK;
-  int grid_z = (K + SplitK - 1) / SplitK;
-
-  int BLOCK = (Ntile == 256) ? 256 : 128;
-
-  dim3 grid(grid_x, grid_y, grid_z);
-  dim3 block(BLOCK);
-
-  bool enable_fuse = fused_gemm_params.EnableFuse;
-  bool has_zp = B_zero != nullptr;
-  if (enable_fuse) {
-    float* C_tmp = reinterpret_cast<float*>(workspace);
-    uint32_t* red_count = reinterpret_cast<uint32_t*>(
-        (char*)workspace + grid_x * Mtile * grid_y * Ntile * sizeof(float));
-    CHECK_CUDA(cudaMemsetAsync(red_count, 0, grid_x * grid_y * sizeof(uint32_t),
-                               stream));
-    SM8x_GEMM_W8A16_Splitk_Params<FType, QType> params{
-        A, B,      B_scale, B_zero, C,       M,     N,
-        K, SplitK, 0,       -1,     nullptr, C_tmp, red_count};
-
-    if (false) {
-    }
-    // Select the template parameters for kernel launch
-    // according to the above settings. Tuning is not supported.
-    __CALL_IF(16, 256, 256, true, false)
-    __CALL_IF(32, 256, 256, true, false)
-    __CALL_IF(48, 256, 256, true, false)
-    __CALL_IF(64, 128, 128, true, false)
-    __CALL_IF(64, 256, 256, true, false)
-    __CALL_IF(16, 256, 256, true, true)
-    __CALL_IF(32, 256, 256, true, true)
-    __CALL_IF(48, 256, 256, true, true)
-    __CALL_IF(64, 128, 128, true, true)
-    __CALL_IF(64, 256, 256, true, true)
-  } else {
-    FType* C_split = reinterpret_cast<FType*>(workspace);
-    SM8x_GEMM_W8A16_Splitk_Params<FType, QType> params{
-        A, B,      B_scale, B_zero, C,       M,       N,
-        K, SplitK, 0,       -1,     C_split, nullptr, nullptr};
-
-    if (false) {
-    }
-    // Select the template parameters for kernel launch
-    // according to the above settings. Tuning is not supported.
-    __CALL_IF(16, 256, 256, false, false)
-    __CALL_IF(32, 256, 256, false, false)
-    __CALL_IF(48, 256, 256, false, false)
-    __CALL_IF(64, 128, 128, false, false)
-    __CALL_IF(64, 256, 256, false, false)
-    __CALL_IF(16, 256, 256, false, true)
-    __CALL_IF(32, 256, 256, false, true)
-    __CALL_IF(48, 256, 256, false, true)
-    __CALL_IF(64, 128, 128, false, true)
-    __CALL_IF(64, 256, 256, false, true)
-
-    // SplitK reduce
-    f16_gemm_splitk_reduce(C_split, C, M, N, grid_z, stream);
-  }
-}
-
-size_t allspark_qgemm_w8a16_perc_n32k16_ampere_workspace_size(
-    int m, int n, int k, int sm_count,
-    BlockTileSplitkParams& fused_gemm_params) {
-  // Determine the block tile and splitk strategy
-  int m16_times = (m + 16 - 1) / 16;
-  int Mtile = m16_times <= 4 ? m16_times * 16 : 64;
-  int grid_x = (m + Mtile - 1) / Mtile;
-  int Ntile =
-      (float(grid_x * ((n + 127) / 128)) / sm_count > 10) || (Mtile < 64) ? 256
-                                                                          : 128;
-  int grid_y = (n + Ntile - 1) / Ntile;
-  int grid_z;
-
-  // split-k
-  const float SPLIT_THRESHOLD = 0.8;
-  int n_slice;
-  for (n_slice = 1; n_slice < k / 256; ++n_slice) {
-    int n_block = grid_x * grid_y * n_slice;
-    if (n_block >= sm_count * SPLIT_THRESHOLD &&
-        (n_block % sm_count == 0 || n_block % sm_count >= sm_count * 0.5)) {
-      break;
-    }
-  }
-
-  int k_slice =
-      (k / n_slice) % 32 == 0 ? k / n_slice : k / n_slice / 32 * 32 + 32;
-  grid_z = (k + k_slice - 1) / k_slice;
-  bool enable_fuse = float(grid_x * grid_y) / sm_count >= 0.5 ? 1 : 0;
-
-  size_t ws_size;
-  if (enable_fuse) {
-    ws_size = grid_x * Mtile * grid_y * Ntile * sizeof(float)  // For C_tmp
-              + grid_x * grid_y * sizeof(uint32_t);            // For red_count
-  } else {
-    ws_size = grid_z * m * n * sizeof(__half);
-  }
-
-  fused_gemm_params.Mtile = Mtile;
-  fused_gemm_params.Ntile = Ntile;
-  fused_gemm_params.SplitK = k_slice;
-  fused_gemm_params.EnableFuse = enable_fuse;
-  return ws_size;
-}
-
-// restore from N32K16 order to original N-major order
-// K % 16 == 0, N % 8 == 0
-// each block process 64(k) * 32(n) result elements
-template <typename FT, typename QT>
-__global__ void restore_N32_K16_dequantize_rhs_w8a16_perc_kernel(
-    const QT* qdata, const FT* scales, const FT* zeros, FT* fdata,
-    const int N_32align, const int N, const int K) {
-  __shared__ FT smem[64 * 32];
-  auto warp_id = threadIdx.x / 32;
-  auto lane_id = threadIdx.x % 32;
-  const auto src_row_idx = blockIdx.x * 8 + lane_id / 4;
-  const int src_col_idx =
-      blockIdx.y * 64 * 4 + warp_id * 16 * 4 + (lane_id % 4) * 16;
-  const int src_offset = src_row_idx * K * 4 + src_col_idx;
-  auto params_nidx = blockIdx.x * 32 + (lane_id / 4) * 4;
-
-  QT qval_reg[16];
-  const QT* pdata = qdata + src_offset;
-  if (src_col_idx < (K * 4)) {
-    *(reinterpret_cast<uint4*>(qval_reg)) =
-        *(reinterpret_cast<const uint4*>(qdata + src_offset));
-  }
-  FT scale_reg[4];
-  *(reinterpret_cast<uint2*>(scale_reg)) =
-      *(reinterpret_cast<const uint2*>(scales + params_nidx));
-  FT zero_reg[4];
-  if (zeros != nullptr) {
-    *(reinterpret_cast<uint2*>(zero_reg)) =
-        *(reinterpret_cast<const uint2*>(zeros + params_nidx));
-  }
-  FT fval_reg[16];
-
-  const int sts_base_offset =
-      (warp_id * 16 + (lane_id % 4) * 2) * 32 + lane_id / 4;
-  #pragma unroll
-  for (int ni = 0; ni < 4; ++ni) {
-    cvt_8bx4_to_16bx4_bias128(
-        *reinterpret_cast<uint32_t*>(&qval_reg[ni * 4]),
-        reinterpret_cast<typename HalfType<FT>::T2*>(&(fval_reg[ni * 4])));
-  #pragma unroll
-    for (int ki = 0; ki < 4; ++ki) {
-      if (zeros != nullptr) {
-        fval_reg[ni * 4 + ki] = __hsub(fval_reg[ni * 4 + ki], zero_reg[ni]);
-      }
-      fval_reg[ni * 4 + ki] = __hmul(fval_reg[ni * 4 + ki], scale_reg[ni]);
-      int sts_offset = sts_base_offset + ((ki / 2) * 8 + (ki % 2)) * 32 +
-                       ((ni + lane_id % 4) % 4) * 8;
-      smem[sts_offset] = fval_reg[ni * 4 + ki];
-    }
-  }
-  __syncthreads();
-
-  const int lds_base_offset =
-      (threadIdx.x / 4) * 32 + ((threadIdx.x % 4 + threadIdx.x / 8) % 4) * 8;
-  #pragma unroll
-  for (int i = 0; i < 2; ++i) {
-    *reinterpret_cast<uint4*>(fval_reg + i * 8) =
-        *reinterpret_cast<uint4*>(smem + lds_base_offset + i * 32 * 32);
-  }
-
-  const auto dst_row_base_kidx = blockIdx.y * 64 + threadIdx.x / 4;
-  const auto dst_col_nidx = blockIdx.x * 32 + (threadIdx.x % 4) * 8;
-  #pragma unroll
-  for (int i = 0; i < 2; ++i) {
-    int dst_row_kidx = dst_row_base_kidx + i * 32;
-    int dst_offset = dst_row_kidx * N + dst_col_nidx;
-    if (dst_row_kidx < K && dst_col_nidx < N) {
-      *reinterpret_cast<uint4*>(fdata + dst_offset) =
-          *reinterpret_cast<uint4*>(fval_reg + i * 8);
-    }
-  }
-}
-
-template <typename FT, typename QT>
-void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
-                                          const FT* zeros, FT* fdata,
-                                          const int N_32align, const int N,
-                                          const int K, const int GroupSize,
-                                          cudaStream_t stream) {
-  STD_TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
-                  "Unsupported shape");
-  if (GroupSize == -1) {
-    const int BLOCK = 128;
-    dim3 grid(N_32align / 32, ((K / 16) + 3) / 4);
-    restore_N32_K16_dequantize_rhs_w8a16_perc_kernel<FT, QT>
-        <<<grid, BLOCK, 0, stream>>>(qdata, scales, zeros, fdata, N_32align, N,
-                                     K);
-  }
-  // TODO: Support SubChannel
-  else {
-    STD_TORCH_CHECK(false, "Now only support PerChannel");
-  }
-}
-
-template <typename FT, typename QT>
-void w8a16_gemm_dq_cublas(const FT* in, const QT* rhs_qdata_ptr,
-                          const FT* rhs_scales_ptr, const FT* rhs_zeros_ptr,
-                          FT* out, void* workspace, const int M,
-                          const int N_32align, const int N, const int K,
-                          const int group_size, cudaStream_t stream,
-                          cublasHandle_t handle) {
-  static_assert(
-      std::is_same<FT, half>::value || std::is_same<FT, nv_bfloat16>::value,
-      "only float16 and bfloat16 is supported");
-  // Dequant
-  FT* rhs_fdata_ptr = static_cast<FT*>(workspace);
-  restore_N32_K16_dequantize_rhs_w8a16(rhs_qdata_ptr, rhs_scales_ptr,
-                                       rhs_zeros_ptr, rhs_fdata_ptr, N_32align,
-                                       N, K, group_size, stream);
-  // cuBLAS GEMM
-  int lda = K;
-  int ldb = N;
-  int ldc = N;
-  const float alpha = 1.0f;
-  const float beta = 0.0f;
-  cudaDataType_t cuda_type;
-  if (std::is_same<FT, __half>::value) {
-    cuda_type = CUDA_R_16F;
-  } else {
-    cuda_type = CUDA_R_16BF;
-  }
-  CHECK_CUBLAS(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha,
-                            rhs_fdata_ptr, cuda_type, ldb, in, cuda_type, lda,
-                            &beta, out, cuda_type, ldc, CUDA_R_32F,
-                            CUBLAS_GEMM_DEFAULT_TENSOR_OP));
-}
-
-template <typename FType, typename QType>
-void allspark_qgemm_w8a16_perc_ampere(
-    const FType* A, const QType* B, const FType* B_scale, const FType* B_zero,
-    FType* C, const int M, const int N_32align, const int N, const int K,
-    void* workspace, const BlockTileSplitkParams& fused_gemm_params,
-    const int group_size, int CUBLAS_M_THRESHOLD, const int sm_version,
-    cudaStream_t stream, cublasHandle_t handle) {
-  if (M > CUBLAS_M_THRESHOLD) {
-    w8a16_gemm_dq_cublas<FType, QType>(A, B, B_scale, B_zero, C, workspace, M,
-                                       N_32align, N, K, group_size, stream,
-                                       handle);
-  } else {
-    ampere_hgemm_W8A16_perc_f16_f16_MtilexNtilex32_mma16816_multistage_AN_BTN32K16_CN_splitk<
-        FType, QType>(A, B, B_scale, B_zero, C, M, N, K, workspace, sm_version,
-                      fused_gemm_params, stream);
-  }
-}
-
-}  // namespace allspark
-
-torch::stable::Tensor allspark_w8a16_gemm(
-    torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
-    torch::stable::Tensor const& b_scales,
-    std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
-    int64_t group_size, int64_t sm_count, int64_t sm_version,
-    int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
-  // Verify device and strides
-  STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
-  STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
-
-  STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
-  STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
-
-  STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
-  STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
-
-  if (has_zp) {
-    STD_TORCH_CHECK(b_qzeros.value().device().is_cuda(),
-                    "b_qzeros is not on GPU");
-    STD_TORCH_CHECK(b_qzeros.value().is_contiguous(),
-                    "b_qzeros is not contiguous");
-  }
-
-  int m = a.size(0);
-  int n_32align = (n + 32 - 1) / 32 * 32;
-  int k = a.size(1);
-
-  // Verify shape
-  STD_TORCH_CHECK(b_qweight.size(0) == n_32align,
-                  "Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
-                  ", n_32align = ", n_32align);
-  STD_TORCH_CHECK(b_qweight.size(1) == k,
-                  "Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
-                  ", k = ", k);
-
-  STD_TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      a.get_device_index());
-  const void* a_ptr = reinterpret_cast<const void*>(a.data_ptr());
-  const uint8_t* b_ptr = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
-  const void* b_scale_ptr = reinterpret_cast<const void*>(b_scales.data_ptr());
-  const void* b_zero_ptr = nullptr;
-  if (b_qzeros.has_value()) {
-    b_zero_ptr = reinterpret_cast<const void*>(b_qzeros.value().data_ptr());
-  }
-
-  auto c =
-      torch::stable::empty({m, n}, a.scalar_type(), std::nullopt, a.device());
-  void* c_ptr = reinterpret_cast<void*>(c.mutable_data_ptr());
-
-  cudaStream_t stream = get_current_cuda_stream();
-  cublasHandle_t handle = get_current_cuda_blas_handle();
-
-  allspark::BlockTileSplitkParams fused_gemm_params;
-
-  size_t ws_size = 0;
-  if (m > CUBLAS_M_THRESHOLD) {
-    ws_size = k * n * 2;  // sizeof(f16)==2
-  } else {
-    ws_size = allspark::allspark_qgemm_w8a16_perc_n32k16_ampere_workspace_size(
-        m, n, k, sm_count, fused_gemm_params);
-  }
-
-  if (as_g_workspace.numel() <
-      ws_size) {  // ws_options: kChar, so numel() is bytes
-    as_g_workspace = torch::stable::empty({static_cast<int64_t>(ws_size)},
-                                          torch::headeronly::ScalarType::Char,
-                                          std::nullopt, a.device());
-  }
-  void* ws = reinterpret_cast<void*>(as_g_workspace.data_ptr());
-
-  if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
-    allspark::allspark_qgemm_w8a16_perc_ampere<__half, uint8_t>(
-        reinterpret_cast<const __half*>(a_ptr), b_ptr,
-        reinterpret_cast<const __half*>(b_scale_ptr),
-        reinterpret_cast<const __half*>(b_zero_ptr),
-        reinterpret_cast<__half*>(c_ptr), m, n_32align, n, k, ws,
-        fused_gemm_params, group_size, CUBLAS_M_THRESHOLD, sm_version, stream,
-        handle);
-  } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    allspark::allspark_qgemm_w8a16_perc_ampere<__nv_bfloat16, uint8_t>(
-        reinterpret_cast<const __nv_bfloat16*>(a_ptr), b_ptr,
-        reinterpret_cast<const __nv_bfloat16*>(b_scale_ptr),
-        reinterpret_cast<const __nv_bfloat16*>(b_zero_ptr),
-        reinterpret_cast<__nv_bfloat16*>(c_ptr), m, n_32align, n, k, ws,
-        fused_gemm_params, group_size, CUBLAS_M_THRESHOLD, sm_version, stream,
-        handle);
-  }
-
-  return c;
-}
-
-#endif
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("allspark_w8a16_gemm", TORCH_BOX(&allspark_w8a16_gemm));
-}
diff --git a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu
deleted file mode 100644
index 9a00f66241b5..000000000000
--- a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu
+++ /dev/null
@@ -1,176 +0,0 @@
-#include "allspark_utils.cuh"
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/core/ScalarType.h>
-
-#include "core/registration.h"
-#include "libtorch_stable/torch_utils.h"
-
-namespace allspark {
-
-// Rearrange B to facilitate Ampere Tensor Core load data
-// reorder B from (K, N) to (N_32align / 4, K * 4)
-// K % 16 == 0, N % 16 == 0, N_32align % 32 == 0
-template <typename FType>
-__global__ void __launch_bounds__(128)
-    rearrange_kn_weight_as_n32k16_order_ldg16_kernel(
-        const uint8_t* B, const FType* B_scale, const FType* B_zero,
-        uint8_t* B_result, FType* B_scale_result, FType* B_zero_result,
-        const int K, const int N, const int N_32align) {
-  const auto lane_id = threadIdx.x % 32;
-  const auto warp_id = threadIdx.x / 32;
-
-  if (blockIdx.x != gridDim.x - 1) {
-    // Load B
-    // per block process 64(k) * 128(n) B elements
-    // per warp process 16(k) * 128 B elements
-    const int src_row_base_idx =
-        blockIdx.x * 64 + warp_id * 16 + ((lane_id % 8) / 2) * 2;
-    const int src_col_idx =
-        blockIdx.y * 128 + (lane_id / 8) * 32 + (lane_id % 2) * 16;
-    uint8_t B_frag[4][16];
-#pragma unroll
-    for (int i = 0; i < 4; ++i) {
-      int src_row_idx = src_row_base_idx + (i / 2) * 8 + (i % 2);
-      int src_offset = src_row_idx * N + src_col_idx;
-      bool guard = src_row_idx < K && src_col_idx < N;
-      ldg128_cg_0(*reinterpret_cast<uint32_t*>(B_frag[i]),
-                  *(reinterpret_cast<uint32_t*>(B_frag[i]) + 1),
-                  *(reinterpret_cast<uint32_t*>(B_frag[i]) + 2),
-                  *(reinterpret_cast<uint32_t*>(B_frag[i]) + 3), B + src_offset,
-                  guard);
-    }
-
-    // reorder B
-    uint8_t B_reorder_frag[8][8];
-#pragma unroll
-    for (int i = 0; i < 4; ++i) {
-#pragma unroll
-      for (int j = 0; j < 16; ++j) {
-        int dst_i = j % 8;
-        int dst_j = i + (j / 8) * 4;
-        B_reorder_frag[dst_i][dst_j] = B_frag[i][j];
-      }
-    }
-
-    // Store B
-    const auto dst_row_base_idx = blockIdx.y * (128 / 4) + (lane_id / 8) * 8;
-    const int dst_col_idx =
-        blockIdx.x * (64 * 4) + warp_id * 64 + (lane_id % 8) * 8;
-    for (int i = 0; i < 8; ++i) {
-      int dst_row_idx = dst_row_base_idx + i;
-      int dst_offset = dst_row_idx * K * 4 + dst_col_idx;
-      bool guard = (dst_row_base_idx < N_32align / 4) && (dst_col_idx < K * 4);
-      if (guard) {
-        *reinterpret_cast<int2*>(B_result + dst_offset) =
-            *reinterpret_cast<int2*>(B_reorder_frag[i]);
-      }
-    }
-  } else {
-    // Load B_scale and B_zero
-    FType b_scale_reg, b_zero_reg;
-    auto src_offset = blockIdx.y * 128 + threadIdx.x;
-    ldg16_cg_0(b_scale_reg, B_scale + src_offset, src_offset < N);
-    if (B_zero != nullptr)
-      ldg16_cg_0(b_zero_reg, B_zero + src_offset, src_offset < N);
-    int dst_offset =
-        blockIdx.y * 128 + warp_id * 32 + (lane_id % 8) * 4 + lane_id / 8;
-    if (dst_offset < N_32align) {
-      B_scale_result[dst_offset] = b_scale_reg;
-      if (B_zero != nullptr) B_zero_result[dst_offset] = b_zero_reg;
-    }
-  }
-}
-
-template <typename FType>
-void rearrange_kn_weight_as_n32k16_order_ldg16(
-    const uint8_t* B, const FType* B_scale, const FType* B_zero,
-    uint8_t* B_result, FType* B_scale_result, FType* B_zero_result,
-    const int64_t K, const int64_t N, const int64_t N_32align,
-    cudaStream_t stream) {
-  if (N % 16 != 0 || K % 16 != 0) {
-    std::cerr << "Now only support N and K is multiples of 16" << std::endl;
-  }
-  const int BLOCK = 128;
-  int grid_x = (K + 64 - 1) / 64 + 1;
-  int grid_y = (N + 128 - 1) / 128;
-  dim3 grid(grid_x, grid_y);
-
-  rearrange_kn_weight_as_n32k16_order_ldg16_kernel<FType>
-      <<<grid, BLOCK, 0, stream>>>(B, B_scale, B_zero, B_result, B_scale_result,
-                                   B_zero_result, K, N, N_32align);
-}
-}  // namespace allspark
-
-void rearrange_kn_weight_as_n32k16_order(
-    torch::stable::Tensor const& b_qweight,
-    torch::stable::Tensor const& b_scales,
-    std::optional<torch::stable::Tensor> const& b_zeros, bool has_zp,
-    torch::stable::Tensor& b_qweight_reorder,
-    torch::stable::Tensor& b_scales_reorder,
-    std::optional<torch::stable::Tensor> const& b_zeros_reorder,
-    const int64_t K, const int64_t N, const int64_t N_32align) {
-  // Verify device and strides
-  STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
-  STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
-
-  STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
-  STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
-
-  STD_TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
-                  "b_qweight_reorder is not on GPU");
-  STD_TORCH_CHECK(b_qweight_reorder.is_contiguous(),
-                  "b_qweight_reorder is not contiguous");
-
-  STD_TORCH_CHECK(b_scales_reorder.device().is_cuda(),
-                  "b_scales_reorder is not on GPU");
-  STD_TORCH_CHECK(b_scales_reorder.is_contiguous(),
-                  "b_scales_reorder is not contiguous");
-
-  if (has_zp) {
-    STD_TORCH_CHECK(b_zeros.value().device().is_cuda(),
-                    "b_zeros is not on GPU");
-    STD_TORCH_CHECK(b_zeros.value().is_contiguous(),
-                    "b_zeros is not contiguous");
-
-    STD_TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
-                    "b_zeros_reorder is not on GPU");
-    STD_TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
-                    "b_zeros_reorder is not contiguous");
-  }
-
-  const uint8_t* matB = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
-  const void* b_scale = b_scales.data_ptr();
-  const void* b_zero = has_zp ? b_zeros.value().data_ptr() : nullptr;
-
-  uint8_t* matB_reorder =
-      reinterpret_cast<uint8_t*>(b_qweight_reorder.mutable_data_ptr());
-  void* b_scale_reorder = b_scales_reorder.mutable_data_ptr();
-  void* b_zero_reorder =
-      has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr;
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      b_qweight.get_device_index());
-  cudaStream_t stream = get_current_cuda_stream();
-  if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
-    allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(
-        matB, reinterpret_cast<const __half*>(b_scale),
-        reinterpret_cast<const __half*>(b_zero), matB_reorder,
-        reinterpret_cast<__half*>(b_scale_reorder),
-        reinterpret_cast<__half*>(b_zero_reorder), K, N, N_32align, stream);
-  } else if (b_scales.scalar_type() ==
-             torch::headeronly::ScalarType::BFloat16) {
-    allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__nv_bfloat16>(
-        matB, reinterpret_cast<const __nv_bfloat16*>(b_scale),
-        reinterpret_cast<const __nv_bfloat16*>(b_zero), matB_reorder,
-        reinterpret_cast<__nv_bfloat16*>(b_scale_reorder),
-        reinterpret_cast<__nv_bfloat16*>(b_zero_reorder), K, N, N_32align,
-        stream);
-  }
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("rearrange_kn_weight_as_n32k16_order",
-         TORCH_BOX(&rearrange_kn_weight_as_n32k16_order));
-}
diff --git a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh
deleted file mode 100644
index 628f6551ea88..000000000000
--- a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh
+++ /dev/null
@@ -1,387 +0,0 @@
-#pragma once
-
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-
-#include <iostream>
-
-#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
-using marlin::MarlinScalarType2;
-
-namespace allspark {
-
-#define CHECK_CUDA(cmd)                                             \
-  do {                                                              \
-    cudaError_t cuda_status = cmd;                                  \
-    if (cuda_status != cudaSuccess) {                               \
-      std::string err_str = cudaGetErrorString(cuda_status);        \
-      std::cerr << "Failed: " << __FILE__ << ":" << __LINE__ << " " \
-                << err_str;                                         \
-      exit(-1);                                                     \
-    }                                                               \
-  } while (0)
-
-#define CHECK_CUBLAS(cmd)                                            \
-  do {                                                               \
-    cublasStatus_t cublas_status = cmd;                              \
-    if (cublas_status != CUBLAS_STATUS_SUCCESS) {                    \
-      std::cerr << "Failed:  " << __FILE__ << ":" << __LINE__ << " " \
-                << cublas_status << std::endl;                       \
-      exit(-1);                                                      \
-    }                                                                \
-  } while (0)
-
-template <typename FType, typename QType>
-struct SM8x_GEMM_W8A16_Splitk_Params {
-  const FType* A_ptr;
-  const QType* B_ptr;
-  const FType* B_scale_ptr;
-  const FType* B_zero_ptr;
-  FType* C_ptr;
-  int M;
-  int N;
-  int K;
-  int SplitK;
-  int GroupCnt;
-  int GroupSize;
-  FType* C_split_ptr;       // for non-fused splitk reduce
-  float* C_tmp_ptr;         // for fused splitk reduce
-  uint32_t* red_count_ptr;  // for fused splitk reduce
-};
-
-struct alignas(16) BlockTileSplitkParams {
-  int Mtile;
-  int Ntile;
-  int SplitK;
-  bool EnableFuse;
-};
-
-template <typename FType, int BLOCK, int N_MATRIX>
-__global__ void f16_gemm_splitk_reduce_kernel(const FType* C_split, FType* C,
-                                              uint32_t n, uint32_t n_matrix,
-                                              uint32_t matrix_size) {
-  auto idx = blockIdx.x * BLOCK + threadIdx.x;
-
-  if (idx >= matrix_size) {
-    return;
-  }
-
-  float sum = 0.f;
-
-  int n_mat = N_MATRIX > 0 ? N_MATRIX : (int)n_matrix;
-  for (int i = 0; i < n_mat; ++i) {
-    sum += MarlinScalarType2<FType>::num2float(C_split[idx + i * matrix_size]);
-  }
-
-  C[idx] = MarlinScalarType2<FType>::float2num(sum);
-}
-
-template <typename FType>
-void f16_gemm_splitk_reduce(const FType* C_split, FType* C, const uint32_t m,
-                            const uint32_t n, const uint32_t n_matrix,
-                            cudaStream_t stream) {
-  const int BLOCK = 128;
-  uint32_t matrix_size = m * n;
-  int grid = (matrix_size + BLOCK - 1) / BLOCK;
-
-  void (*kernel)(const FType*, FType*, uint32_t, uint32_t, uint32_t) = nullptr;
-
-  switch (n_matrix) {
-    case 4:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 4>;
-      break;
-    case 5:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 5>;
-      break;
-    case 6:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 6>;
-      break;
-    case 7:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 7>;
-      break;
-    case 8:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 8>;
-      break;
-    case 9:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 9>;
-      break;
-    case 10:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 10>;
-      break;
-    case 11:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 11>;
-      break;
-    case 12:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 12>;
-      break;
-    default:
-      kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, -1>;
-      break;
-  }
-
-  kernel<<<grid, BLOCK, 0, stream>>>(C_split, C, n, n_matrix, matrix_size);
-}
-
-template <typename T>
-struct HalfType;
-template <>
-struct HalfType<half> {
-  using T1 = __half;
-  using T2 = __half2;
-};
-template <>
-struct HalfType<__nv_bfloat16> {
-  using T1 = __nv_bfloat16;
-  using T2 = __nv_bfloat162;
-};
-
-// convert 64-bit pointer to 32-bit smem addr
-__device__ __forceinline__ uint32_t smem_u32addr(const void* smem_ptr) {
-  uint32_t addr;
-  asm("{.reg .u64 u64addr;\n"
-      " cvta.to.shared.u64 u64addr, %1;\n"
-      " cvt.u32.u64 %0, u64addr;}\n"
-      : "=r"(addr)
-      : "l"(smem_ptr));
-
-  return addr;
-}
-
-template <typename T>
-__device__ __forceinline__ void ldg16_cg_0(T& r0, const void* ptr, bool guard) {
-  static_assert(sizeof(T) == 2, "ldg16_cg_0: invalid T");
-
-  asm volatile(
-      "{.reg .pred p;\n"
-      " setp.ne.b32 p, %2, 0;\n"
-      " @!p mov.b16 %0, 0;\n"
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDACC_VER_MINOR__ >= 4 && \
-    __CUDA_ARCH__ >= 750
-      " @p ld.global.cg.L2::128B.b16 {%0}, [%1];}\n"
-#else
-      " @p ld.global.ca.b16 {%0}, [%1];}\n"
-#endif
-      : "=h"(reinterpret_cast<uint16_t&>(r0))
-      : "l"(ptr), "r"((int)guard));
-}
-
-template <typename T>
-__device__ __forceinline__ void ldg64_ca(T& r0, T& r1, const void* ptr,
-                                         bool guard) {
-  static_assert(sizeof(T) == 4, "ldg64_ca: invalid T");
-
-  asm volatile(
-      "{.reg .pred p;\n"
-      " setp.ne.b32 p, %3, 0;\n"
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDACC_VER_MINOR__ >= 4 && \
-    __CUDA_ARCH__ >= 750
-      " @p ld.global.ca.L2::128B.v2.b32 {%0, %1}, [%2];}\n"
-#else
-      " @p ld.global.ca.v2.b32 {%0, %1}, [%2];}\n"
-#endif
-      : "=r"(reinterpret_cast<uint32_t&>(r0)),
-        "=r"(reinterpret_cast<uint32_t&>(r1))
-      : "l"(ptr), "r"((int)guard));
-}
-
-template <typename T>
-__device__ __forceinline__ void ldg128_cg_0(T& r0, T& r1, T& r2, T& r3,
-                                            const void* ptr, bool guard) {
-  static_assert(sizeof(T) == 4, "ldg128_cg_0: invalid T");
-
-  asm volatile(
-      "{.reg .pred p;\n"
-      " setp.ne.b32 p, %5, 0;\n"
-      " @!p mov.b32 %0, 0;\n"
-      " @!p mov.b32 %1, 0;\n"
-      " @!p mov.b32 %2, 0;\n"
-      " @!p mov.b32 %3, 0;\n"
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDACC_VER_MINOR__ >= 4 && \
-    __CUDA_ARCH__ >= 750
-      " @p ld.global.cg.L2::128B.v4.b32 {%0, %1, %2, %3}, [%4];}\n"
-#else
-      " @p ld.global.cg.v4.b32 {%0, %1, %2, %3}, [%4];}\n"
-#endif
-      : "=r"(reinterpret_cast<uint32_t&>(r0)),
-        "=r"(reinterpret_cast<uint32_t&>(r1)),
-        "=r"(reinterpret_cast<uint32_t&>(r2)),
-        "=r"(reinterpret_cast<uint32_t&>(r3))
-      : "l"(ptr), "r"((int)guard));
-}
-
-template <typename T>
-__device__ __forceinline__ void lds128(T& reg0, T& reg1, T& reg2, T& reg3,
-                                       const uint32_t addr) {
-  static_assert(sizeof(T) == 4, "lds128: invalid T");
-
-  asm volatile("ld.shared.v4.b32 {%0, %1, %2, %3}, [%4];\n"
-               : "=r"(reinterpret_cast<uint32_t&>(reg0)),
-                 "=r"(reinterpret_cast<uint32_t&>(reg1)),
-                 "=r"(reinterpret_cast<uint32_t&>(reg2)),
-                 "=r"(reinterpret_cast<uint32_t&>(reg3))
-               : "r"(addr));
-}
-
-template <typename T>
-__device__ __forceinline__ void stg128(const T& r0, const T& r1, const T& r2,
-                                       const T& r3, const void* ptr,
-                                       bool guard) {
-  static_assert(sizeof(T) == 4, "stg128: invalid T");
-
-  asm volatile(
-      "{.reg .pred p;\n"
-      " setp.ne.b32 p, %1, 0;\n"
-      " @p st.global.v4.b32 [%0], {%2, %3, %4, %5};}\n"
-      :
-      : "l"(ptr), "r"((int)guard), "r"(reinterpret_cast<const uint32_t&>(r0)),
-        "r"(reinterpret_cast<const uint32_t&>(r1)),
-        "r"(reinterpret_cast<const uint32_t&>(r2)),
-        "r"(reinterpret_cast<const uint32_t&>(r3)));
-}
-
-template <typename T>
-__device__ __forceinline__ void ldsm_4(T& r0, T& r1, T& r2, T& r3,
-                                       const uint32_t& addr) {
-  static_assert(sizeof(T) == 4, "ldsm_4: invalid T");
-#if (__CUDA_ARCH__ >= 750) && (__CUDACC_VER_MAJOR__ >= 11)
-  asm volatile(
-      "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
-      : "=r"(reinterpret_cast<uint32_t&>(r0)),
-        "=r"(reinterpret_cast<uint32_t&>(r1)),
-        "=r"(reinterpret_cast<uint32_t&>(r2)),
-        "=r"(reinterpret_cast<uint32_t&>(r3))
-      : "r"(addr));
-#endif
-}
-
-template <typename FType>
-__device__ __forceinline__ void hmma16816_f32(float (&d)[4],
-                                              const uint32_t (&a)[4],
-                                              const uint32_t (&b)[2]);
-
-template <>
-__device__ __forceinline__ void hmma16816_f32<__half>(float (&d)[4],
-                                                      const uint32_t (&a)[4],
-                                                      const uint32_t (&b)[2]) {
-#if (__CUDA_ARCH__ >= 800) && (__CUDACC_VER_MAJOR__ >= 11)
-  asm volatile(
-      "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, "
-      "{%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
-      : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
-      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
-#endif
-}
-
-template <>
-__device__ __forceinline__ void hmma16816_f32<__nv_bfloat16>(
-    float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) {
-#if (__CUDA_ARCH__ >= 800) && (__CUDACC_VER_MAJOR__ >= 11)
-  asm volatile(
-      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, "
-      "{%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
-      : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
-      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
-#endif
-}
-
-template <int SIZE_IN_BYTES>
-__device__ __forceinline__ void cp_async(const uint32_t smem_addr,
-                                         const void* gmem_ptr,
-                                         const int src_in_bytes, bool guard) {
-  static_assert(
-      (SIZE_IN_BYTES == 4 || SIZE_IN_BYTES == 8 || SIZE_IN_BYTES == 16),
-      "Size is not supported");
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
-  asm volatile(
-      "{.reg.pred p;\n"
-      " setp.ne.b32 p, %4, 0;\n"
-  #if __CUDACC_VER_MINOR__ >= 4
-      " @p cp.async.cg.shared.global.L2::256B [%0], [%1], %2, %3;}\n"
-  #else
-      " @p cp.async.cg.shared.global [%0], [%1], %2, %3;}\n"
-  #endif
-      ::"r"(smem_addr),
-      "l"(gmem_ptr), "n"(SIZE_IN_BYTES), "r"(src_in_bytes), "r"((int)guard));
-#endif
-}
-
-__device__ __forceinline__ void cp_async_commit_group() {
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
-  asm volatile("cp.async.commit_group;\n");
-#endif
-}
-
-template <int N>
-__device__ __forceinline__ void cp_asyc_wait_group() {
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
-  asm volatile("cp.async.wait_group %0;\n" : : "n"(N));
-#endif
-}
-
-template <typename T>
-__device__ __forceinline__ void cvt_8bx4_to_16bx4_bias128(const uint32_t& idata,
-                                                          T* fdata);
-
-template <>
-// fast conversion: 4xuint8 to 4xhalf, subtracting bias = 128
-__device__ __forceinline__ void cvt_8bx4_to_16bx4_bias128<__half2>(
-    const uint32_t& idata, __half2* fdata) {
-  uint32_t i10, i32;
-  asm volatile(
-      "prmt.b32 %0, %2, 0x64, 0x4140;"
-      "prmt.b32 %1, %2, 0x64, 0x4342;"
-      : "=r"(i10), "=r"(i32)
-      : "r"(idata));
-
-  static constexpr uint32_t MAGIC_NUM = 0x64806480;
-  fdata[0] = __hsub2(reinterpret_cast<const __half2&>(i10),
-                     reinterpret_cast<const __half2&>(MAGIC_NUM));
-  fdata[1] = __hsub2(reinterpret_cast<const __half2&>(i32),
-                     reinterpret_cast<const __half2&>(MAGIC_NUM));
-}
-
-template <>
-// fast conversion: 4xuint8 to 4xbfloat16, subtracting bias = 128
-// reference from marlin fast implementation
-__device__ __forceinline__ void cvt_8bx4_to_16bx4_bias128<__nv_bfloat162>(
-    const uint32_t& idata, __nv_bfloat162* fdata) {
-  float fp32_imd[4];
-  uint32_t* fp32_imd_casted = reinterpret_cast<uint32_t*>(fp32_imd);
-  asm volatile(
-      "prmt.b32 %0, %4, 0x4B000000, 0x7650;"
-      "prmt.b32 %1, %4, 0x4B000000, 0x7651;"
-      "prmt.b32 %2, %4, 0x4B000000, 0x7652;"
-      "prmt.b32 %3, %4, 0x4B000000, 0x7653;"
-      : "=r"(fp32_imd_casted[0]), "=r"(fp32_imd_casted[1]),
-        "=r"(fp32_imd_casted[2]), "=r"(fp32_imd_casted[3])
-      : "r"(idata));
-
-  fp32_imd[0] -= 8388736.f;
-  fp32_imd[1] -= 8388736.f;
-  fp32_imd[2] -= 8388736.f;
-  fp32_imd[3] -= 8388736.f;
-
-  uint32_t* bf16_res = reinterpret_cast<uint32_t*>(fdata);
-  asm volatile(
-      "prmt.b32 %0, %2, %3, 0x7632;"
-      "prmt.b32 %1, %4, %5, 0x7632;"
-      : "=r"(bf16_res[0]), "=r"(bf16_res[1])
-      : "r"(fp32_imd_casted[0]), "r"(fp32_imd_casted[1]),
-        "r"(fp32_imd_casted[2]), "r"(fp32_imd_casted[3]));
-}
-
-static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
-  assert(false);
-#else
-  return __bfloat162bfloat162(x);
-#endif
-  __builtin_unreachable();  // Suppress missing return statement warning
-}
-
-static __device__ half2 inline num2num2(const half x) {
-  return __half2half2(x);
-}
-
-}  // namespace allspark
diff --git a/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu b/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu
index 665585caa46c..a5d199db2737 100644
--- a/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu
+++ b/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu
@@ -796,7 +796,7 @@ torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, bool inplace)
 
     VLLM_STABLE_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
       auto constexpr SCALAR_TYPE = torch::headeronly::CppTypeToScalarType<scalar_t>::value;
-      hadacore::run_fht<SCALAR_TYPE>(x.data_ptr(), x.data_ptr(), x.numel(), had_size, stream);
+      hadacore::run_fht<SCALAR_TYPE>(x.data_ptr(), out.data_ptr(), x.numel(), had_size, stream);
     });
 
     if (numel % 256 != 0) {
diff --git a/csrc/libtorch_stable/quantization/marlin/marlin.cuh b/csrc/libtorch_stable/quantization/marlin/marlin.cuh
index bfb65e874b3d..af4d4730ba3a 100644
--- a/csrc/libtorch_stable/quantization/marlin/marlin.cuh
+++ b/csrc/libtorch_stable/quantization/marlin/marlin.cuh
@@ -67,14 +67,6 @@ __device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
   }
 }
 
-__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
-                                         bool pred = true) {
-  if (pred) {
-    reinterpret_cast<int4*>(smem_ptr)[0] =
-        reinterpret_cast<const int4*>(glob_ptr)[0];
-  }
-}
-
 __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
                                       bool pred = true) {
   if (pred) {
@@ -121,19 +113,6 @@ __device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
       "r"(smem), "l"(glob_ptr), "n"(BYTES));
 }
 
-__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
-                                         bool pred = true) {
-  const int BYTES = 16;
-  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
-  asm volatile(
-      "{\n"
-      "   .reg .pred p;\n"
-      "   setp.ne.b32 p, %0, 0;\n"
-      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
-      "}\n" ::"r"((int)pred),
-      "r"(smem), "l"(glob_ptr), "n"(BYTES));
-}
-
 __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
                                       bool pred = true) {
   const int BYTES = 16;
@@ -170,4 +149,4 @@ __device__ inline void cp_async_wait() {
 
 }  // namespace MARLIN_NAMESPACE_NAME
 
-#endif
\ No newline at end of file
+#endif
diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh
index 8f74f2991469..fa2add918213 100644
--- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh
+++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh
@@ -152,6 +152,29 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
 
   int32_t m = a.size(0), n = b.size(1), k = a.size(1);
 
+  // The SM100 blockwise SF copy requires the SFA M extent to be a multiple
+  // of 4 (4-row copy atom). For misaligned M on the non-swapAB path, pad only
+  // the per-token activation scales (column-major (m, k/128), 1/64 the size
+  // of the bf16 activation) and build layout_SFA from the padded extent.
+  int32_t m_sf = m;
+  std::optional<torch::stable::Tensor> a_scales_pad;
+  void const* a_scales_data = a_scales.data_ptr();
+  if constexpr (!swap_ab) {
+    if (m % 4 != 0) {
+      m_sf = (m + 3) & ~3;
+      int64_t num_groups = a_scales.numel() / m;
+      a_scales_pad = torch::stable::new_empty(
+          a_scales, {num_groups * m_sf},
+          torch::headeronly::ScalarType::Float);
+      cudaMemcpy2DAsync(a_scales_pad->data_ptr(), m_sf * sizeof(float),
+                        a_scales.data_ptr(), m * sizeof(float),
+                        m * sizeof(float), num_groups,
+                        cudaMemcpyDeviceToDevice,
+                        get_current_cuda_stream(a.get_device()));
+      a_scales_data = a_scales_pad->data_ptr();
+    }
+  }
+
   StrideA a_stride;
   StrideB b_stride;
   StrideC c_stride;
@@ -162,16 +185,16 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
   c_stride =
       cutlass::make_cute_packed_stride(StrideC{}, swap_ab ? cute::make_shape(n, m, 1) : cute::make_shape(m, n, 1));
 
-  LayoutSFA layout_SFA = swap_ab ? 
+  LayoutSFA layout_SFA = swap_ab ?
       ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)) :
-      ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
+      ScaleConfig::tile_atom_to_shape_SFA(make_shape(m_sf, n, k, 1));
   LayoutSFB layout_SFB = swap_ab ?
       ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) :
       ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
 
   auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
   auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
-  auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales.data_ptr());
+  auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales_data);
   auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
 
   typename GemmKernel::MainloopArguments mainloop_args{};
@@ -212,7 +235,9 @@ void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::stable::Tensor& out,
 
   constexpr int TILE_K = 128;
   // TODO: better heuristics
-  bool swap_ab = (m < 16) || (m % 4 != 0);
+  // swapAB only wins for small M; misaligned M >= 65 is faster on the
+  // non-swapAB path with padded activation scales (measured on GB300).
+  bool swap_ab = (m <= 64);
   bool use_tma_epilogue = (m * n) % 4 == 0;
   if (!swap_ab) {
     constexpr int TILE_N = 128;
diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
index 06fa29eaba29..cfd1fab1ad00 100644
--- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+++ b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
@@ -96,7 +96,8 @@ __device__ __forceinline__ void QuantizeGroup(
 }
 
 template <typename T, typename DST_DTYPE, bool IS_COLUMN_MAJOR = false,
-          bool SCALE_UE8M0 = false, typename scale_packed_t = float>
+          bool SCALE_UE8M0 = false, typename scale_packed_t = float,
+          bool CHECK_BOUNDS = false>
 __global__ void per_token_group_quant_8bit_kernel(
     const T* __restrict__ input, void* __restrict__ output_q,
     scale_packed_t* __restrict__ output_s, const int group_size,
@@ -114,49 +115,54 @@ __global__ void per_token_group_quant_8bit_kernel(
   using scale_element_t = float;
   static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
 
-  const T* group_input = input + block_group_offset;
-  DST_DTYPE* group_output =
-      static_cast<DST_DTYPE*>(output_q) + block_group_offset;
-  scale_element_t* scale_output;
-
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
   cudaGridDependencySynchronize();
 #endif
 
-  if constexpr (IS_COLUMN_MAJOR) {
-    const int num_elems_per_pack =
-        static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
-    const int scale_num_rows_element = scale_num_rows * num_elems_per_pack;
-    const int row_idx = global_group_id / scale_num_rows_element;
-    const int col_idx_raw = global_group_id % scale_num_rows_element;
-    const int col_idx = col_idx_raw / num_elems_per_pack;
-    const int pack_idx = col_idx_raw % num_elems_per_pack;
-    scale_output = reinterpret_cast<scale_element_t*>(output_s) +
-                   (col_idx * scale_stride * num_elems_per_pack +
-                    row_idx * num_elems_per_pack + pack_idx);
-  } else {
-    scale_output = output_s + global_group_id;
-  }
-
   // shared memory to cache each group's data to avoid double DRAM reads.
   extern __shared__ __align__(16) char smem_raw[];
   T* smem = reinterpret_cast<T*>(smem_raw);
   T* smem_group = smem + local_group_id * group_size;
+  const bool valid_group = !CHECK_BOUNDS || global_group_id < num_groups;
+  float y_s = 1.0f;
+  if (valid_group) {
+    const T* group_input = input + block_group_offset;
+    scale_element_t* scale_output;
+
+    if constexpr (IS_COLUMN_MAJOR) {
+      const int num_elems_per_pack =
+          static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
+      const int scale_num_rows_element = scale_num_rows * num_elems_per_pack;
+      const int row_idx = global_group_id / scale_num_rows_element;
+      const int col_idx_raw = global_group_id % scale_num_rows_element;
+      const int col_idx = col_idx_raw / num_elems_per_pack;
+      const int pack_idx = col_idx_raw % num_elems_per_pack;
+      scale_output = reinterpret_cast<scale_element_t*>(output_s) +
+                     (col_idx * scale_stride * num_elems_per_pack +
+                      row_idx * num_elems_per_pack + pack_idx);
+    } else {
+      scale_output = output_s + global_group_id;
+    }
 
-  const float y_s = ComputeGroupScale<T, SCALE_UE8M0>(
-      group_input, smem_group, group_size, lane_id, threads_per_group, eps,
-      max_8bit);
+    y_s = ComputeGroupScale<T, SCALE_UE8M0>(group_input, smem_group, group_size,
+                                            lane_id, threads_per_group, eps,
+                                            max_8bit);
 
-  scale_element_t y_s_quant = y_s;
+    scale_element_t y_s_quant = y_s;
 
-  if (lane_id == 0) {
-    *scale_output = y_s_quant;
+    if (lane_id == 0) {
+      *scale_output = y_s_quant;
+    }
   }
 
   __syncthreads();
 
-  QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
-                              threads_per_group, y_s, min_8bit, max_8bit);
+  if (valid_group) {
+    DST_DTYPE* group_output =
+        static_cast<DST_DTYPE*>(output_q) + block_group_offset;
+    QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
+                                threads_per_group, y_s, min_8bit, max_8bit);
+  }
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
   cudaTriggerProgrammaticLaunchCompletion();
@@ -209,10 +215,20 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
 
   constexpr int THREADS_PER_GROUP = 16;
 
-  const int groups_per_block = GetGroupsPerBlock(num_groups);
+  if (num_groups == 0) {
+    return;
+  }
+  STD_TORCH_CHECK(input.size(-1) % group_size == 0);
+  const int64_t groups_per_row = input.size(-1) / group_size;
+  const int64_t num_rows = num_groups / groups_per_row;
+  // Odd row counts inherit the next even row count's launch geometry.
+  const int64_t policy_num_groups =
+      (num_rows + (num_rows & 1)) * groups_per_row;
+  const int groups_per_block = GetGroupsPerBlock(policy_num_groups);
+  const bool needs_bounds_check = num_groups % groups_per_block != 0;
 
   auto dst_type = output_q.scalar_type();
-  const int num_blocks = num_groups / groups_per_block;
+  const int num_blocks = (num_groups + groups_per_block - 1) / groups_per_block;
   const int num_threads = groups_per_block * THREADS_PER_GROUP;
 
   const bool is_column_major = output_s.stride(0) < output_s.stride(1);
@@ -220,38 +236,51 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
   const int scale_stride = output_s.stride(1);
 
 #ifndef USE_ROCM
-  #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES)     \
-    do {                                                                     \
-      cudaLaunchConfig_t config = {};                                        \
-      config.gridDim = dim3(num_blocks);                                     \
-      config.blockDim = dim3(num_threads);                                   \
-      config.dynamicSmemBytes = (SMEM_BYTES);                                \
-      config.stream = stream;                                                \
-      cudaLaunchAttribute attrs[1];                                          \
-      attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;      \
-      attrs[0].val.programmaticStreamSerializationAllowed = 1;               \
-      config.numAttrs = 1;                                                   \
-      config.attrs = attrs;                                                  \
-      cudaLaunchKernelEx(                                                    \
-          &config,                                                           \
-          per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0>, \
-          static_cast<T*>(input.data_ptr()), output_q.data_ptr(),            \
-          static_cast<float*>(output_s.data_ptr()), group_size, num_groups,  \
-          groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit,    \
-          scale_num_rows, scale_stride);                                     \
+  #define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS,  \
+                             SMEM_BYTES)                                    \
+    do {                                                                    \
+      cudaLaunchConfig_t config = {};                                       \
+      config.gridDim = dim3(num_blocks);                                    \
+      config.blockDim = dim3(num_threads);                                  \
+      config.dynamicSmemBytes = (SMEM_BYTES);                               \
+      config.stream = stream;                                               \
+      cudaLaunchAttribute attrs[1];                                         \
+      attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;     \
+      attrs[0].val.programmaticStreamSerializationAllowed = 1;              \
+      config.numAttrs = 1;                                                  \
+      config.attrs = attrs;                                                 \
+      cudaLaunchKernelEx(                                                   \
+          &config,                                                          \
+          per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0, \
+                                            float, CHECK_BOUNDS>,           \
+          static_cast<T*>(input.data_ptr()), output_q.data_ptr(),           \
+          static_cast<float*>(output_s.data_ptr()), group_size, num_groups, \
+          groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit,   \
+          scale_num_rows, scale_stride);                                    \
     } while (0)
 #else
-  #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES)   \
-    do {                                                                   \
-      per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0>    \
-          <<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>( \
-              static_cast<T*>(input.data_ptr()), output_q.data_ptr(),      \
-              static_cast<float*>(output_s.data_ptr()), group_size,        \
-              num_groups, groups_per_block, (float)eps, (float)min_8bit,   \
-              (float)max_8bit, scale_num_rows, scale_stride);              \
+  #define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS,     \
+                             SMEM_BYTES)                                       \
+    do {                                                                       \
+      per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0, float, \
+                                        CHECK_BOUNDS>                          \
+          <<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>(     \
+              static_cast<T*>(input.data_ptr()), output_q.data_ptr(),          \
+              static_cast<float*>(output_s.data_ptr()), group_size,            \
+              num_groups, groups_per_block, (float)eps, (float)min_8bit,       \
+              (float)max_8bit, scale_num_rows, scale_stride);                  \
     } while (0)
 #endif
 
+#define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES)       \
+  do {                                                                       \
+    if (needs_bounds_check) {                                                \
+      LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, true, SMEM_BYTES);  \
+    } else {                                                                 \
+      LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, false, SMEM_BYTES); \
+    }                                                                        \
+  } while (0)
+
 #define LAUNCH_KERNEL(T, DST_DTYPE)                                     \
   do {                                                                  \
     size_t smem_bytes =                                                 \
@@ -284,6 +313,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
 
 #undef LAUNCH_KERNEL
 #undef LAUNCH_KERNEL_INST
+#undef LAUNCH_KERNEL_IMPL
 }
 
 // Register-resident fast path for group_size==128.
diff --git a/csrc/libtorch_stable/sampled_topk.cuh b/csrc/libtorch_stable/sampled_topk.cuh
new file mode 100644
index 000000000000..867abbd63542
--- /dev/null
+++ b/csrc/libtorch_stable/sampled_topk.cuh
@@ -0,0 +1,253 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+// Sampled top-k for long FP32 sparse-indexer rows, inspired by DeepSelect:
+// https://github.com/deepseek-ai/DeepSelect
+// A sample estimates a cutoff, then survivors are compacted in shared memory
+// for exact FP32 selection. Too few or too many candidates trigger an exact
+// full-row fallback.
+
+#pragma once
+
+#include <cub/cub.cuh>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+#include "persistent_topk.cuh"
+
+namespace vllm::sampled_topk {
+
+constexpr int kThreads = 1024;
+// Conservative crossover bounds from the B300 batch/length sweep.
+template <int K>
+constexpr int kMinSampledLength = K == 512 ? 98304 : 65536;
+// Target half the buffer to leave room for sampling error.
+constexpr int kSample = 4096;
+constexpr int kCapacity = 8192;
+
+struct Storage {
+  uint2 candidates[kCapacity];
+  int histogram[2048];
+  typename cub::BlockScan<int, kThreads>::TempStorage scan;
+  int count;
+  int remaining;
+  int bin;
+  uint32_t prefix;
+};
+
+__device__ __forceinline__ uint32_t ordered(float value) {
+  const uint32_t bits = __float_as_uint(value);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+__device__ __forceinline__ uint32_t coarse(uint32_t key) {
+  const uint32_t bits = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key;
+  return topk_histogram_4096::extract_coarse_bin_N<11>(__uint_as_float(bits));
+}
+
+// Select the bin containing the rank and update the rank within that bin.
+template <int Bins>
+__device__ void select_bin(Storage& s) {
+  constexpr int kItems = (Bins + kThreads - 1) / kThreads;
+  int sums[kItems];
+#pragma unroll
+  for (int j = 0; j < kItems; ++j) {
+    const int bin = Bins - 1 - (threadIdx.x * kItems + j);
+    sums[j] = bin >= 0 ? s.histogram[bin] : 0;
+  }
+  cub::BlockScan<int, kThreads>(s.scan).InclusiveSum(sums, sums);
+  const int remaining = s.remaining;
+  __syncthreads();
+#pragma unroll
+  for (int j = 0; j < kItems; ++j) {
+    const int bin = Bins - 1 - (threadIdx.x * kItems + j);
+    if (bin >= 0 && sums[j] >= remaining &&
+        sums[j] - s.histogram[bin] < remaining) {
+      s.bin = bin;
+      s.remaining = remaining - (sums[j] - s.histogram[bin]);
+    }
+  }
+  __syncthreads();
+}
+
+// Find a coarse bin, then refine its full FP32 keys one byte at a time.
+// Read shared candidates normally, or the full row for exact fallback.
+template <bool Buffered>
+__device__ uint32_t threshold(Storage& s, const float* row, int count,
+                              int rank) {
+  const int tid = threadIdx.x;
+  if (tid == 0) {
+    s.remaining = rank;
+    s.prefix = 0;
+  }
+  for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0;
+  __syncthreads();
+  for (int i = tid; i < count; i += kThreads) {
+    const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]);
+    atomicAdd(&s.histogram[coarse(key)], 1);
+  }
+  __syncthreads();
+  select_bin<2048>(s);
+  const int coarse_bin = s.bin;
+  uint32_t mask = 0;
+  for (int shift = 24; shift >= 0; shift -= 8) {
+    if (tid < 256) s.histogram[tid] = 0;
+    __syncthreads();
+    const uint32_t prefix = s.prefix;
+    for (int i = tid; i < count; i += kThreads) {
+      const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]);
+      if (coarse(key) == coarse_bin && (key & mask) == prefix) {
+        atomicAdd(&s.histogram[(key >> shift) & 255], 1);
+      }
+    }
+    __syncthreads();
+    select_bin<256>(s);
+    if (tid == 0) s.prefix = prefix | (static_cast<uint32_t>(s.bin) << shift);
+    mask |= 255u << shift;
+    __syncthreads();
+  }
+  return s.prefix;
+}
+
+__device__ __forceinline__ int reserve(bool hit, int* counter) {
+  const uint32_t mask = __ballot_sync(0xffffffffu, hit);
+  if (mask == 0) return 0;
+  const int lane = threadIdx.x & 31;
+  int base = 0;
+  if (lane == 0) base = atomicAdd(counter, __popc(mask));
+  base = __shfl_sync(0xffffffffu, base, 0);
+  return base + __popc(mask & ((1u << lane) - 1));
+}
+
+template <bool Buffered>
+__device__ void emit(Storage& s, const float* row, int32_t* dst, int count,
+                     uint32_t cutoff) {
+  const int tid = threadIdx.x;
+  if (tid == 0) s.count = 0;
+  __syncthreads();
+  for (int base = 0; base < count; base += kThreads) {
+    const int i = base + tid;
+    const uint2 pair = i < count ? (Buffered ? s.candidates[i]
+                                             : make_uint2(i, ordered(row[i])))
+                                 : make_uint2(0, 0);
+    bool keep = i < count && pair.y > cutoff;
+    if (i < count && pair.y == cutoff) {
+      keep = atomicSub(&s.remaining, 1) > 0;
+    }
+    const int offset = reserve(keep, &s.count);
+    if (keep) dst[offset] = pair.x;
+  }
+}
+
+template <int K>
+__global__ void __launch_bounds__(kThreads)
+    sampled_topk_kernel(const float* __restrict__ input,
+                        const int32_t* __restrict__ lengths,
+                        int32_t* __restrict__ output, int64_t stride,
+                        int max_length) {
+  extern __shared__ __align__(16) unsigned char smem[];
+  auto& s = *reinterpret_cast<Storage*>(smem);
+  const int tid = threadIdx.x;
+  const int length = max(0, min(lengths[blockIdx.x], max_length));
+  const float* row = input + blockIdx.x * stride;
+  int32_t* dst = output + blockIdx.x * K;
+  // Graph allocations can hold rows much shorter than the host length bound.
+  if (length < kMinSampledLength<K>) {
+    __shared__ filtered_topk::FilteredTopKStorage<K> fallback_storage;
+    bool selected;
+    if ((stride & 3) == 0 && (reinterpret_cast<uintptr_t>(row) & 15) == 0) {
+      selected =
+          filtered_topk::filtered_topk_row<float, int32_t, 4, K, false, true>(
+              row, dst, length, K, fallback_storage);
+    } else {
+      selected =
+          filtered_topk::filtered_topk_row<float, int32_t, 1, K, true, true>(
+              row, dst, length, K, fallback_storage);
+    }
+    if (!selected) {
+      const uint32_t exact_cutoff = threshold<false>(s, row, length, K);
+      emit<false>(s, row, dst, length, exact_cutoff);
+    }
+    return;
+  }
+  if (tid == 0) s.remaining = max(1, (kCapacity / 2) * kSample / length);
+  for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0;
+  __syncthreads();
+  // Sample contiguous warps spread across the row to keep loads coalesced.
+  for (int i = tid; i < kSample; i += kThreads) {
+    const int index = (i / 32) * (length / (kSample / 32)) + i % 32;
+    atomicAdd(&s.histogram[coarse(ordered(row[index]))], 1);
+  }
+  __syncthreads();
+  select_bin<2048>(s);
+  const uint16_t half_key = s.bin << 5;
+  const uint16_t half_bits =
+      (half_key & 0x8000u) ? (half_key ^ 0x8000u) : ~half_key;
+  const uint32_t cutoff = ordered(__half2float(__ushort_as_half(half_bits)));
+  if (tid == 0) s.count = 0;
+  __syncthreads();
+  constexpr int kItems = 16;
+  // Compact survivors with one reservation per warp per tile. Keep counting
+  // beyond capacity so overflow triggers fallback instead of using a subset.
+  for (int base = 0; base < length; base += kItems * kThreads) {
+    uint32_t keys[kItems];
+    uint32_t hit_mask = 0;
+    if ((reinterpret_cast<uintptr_t>(row) & 15) == 0 &&
+        base + kItems * kThreads <= length) {
+#pragma unroll
+      for (int j = 0; j < kItems / 4; ++j) {
+        const float4 values =
+            reinterpret_cast<const float4*>(row + base)[tid + j * kThreads];
+        keys[4 * j] = ordered(values.x);
+        keys[4 * j + 1] = ordered(values.y);
+        keys[4 * j + 2] = ordered(values.z);
+        keys[4 * j + 3] = ordered(values.w);
+      }
+    } else {
+#pragma unroll
+      for (int j = 0; j < kItems; ++j) {
+        const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4;
+        keys[j] = i < length ? ordered(row[i]) : 0;
+      }
+    }
+#pragma unroll
+    for (int j = 0; j < kItems; ++j) {
+      const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4;
+      if (i < length && keys[j] >= cutoff) hit_mask |= 1u << j;
+    }
+    const int lane = tid & 31;
+    const int hits = __popc(hit_mask);
+    int inclusive = hits;
+#pragma unroll
+    for (int delta = 1; delta < 32; delta *= 2) {
+      const int previous = __shfl_up_sync(0xffffffffu, inclusive, delta);
+      if (lane >= delta) inclusive += previous;
+    }
+    const int total = __shfl_sync(0xffffffffu, inclusive, 31);
+    int warp_base = 0;
+    if (lane == 31 && total != 0) warp_base = atomicAdd(&s.count, total);
+    warp_base = __shfl_sync(0xffffffffu, warp_base, 31);
+    int offset = warp_base + inclusive - hits;
+#pragma unroll
+    for (int j = 0; j < kItems; ++j) {
+      if (hit_mask & (1u << j)) {
+        if (offset < kCapacity) {
+          s.candidates[offset] = make_uint2(
+              base + tid * 4 + (j / 4) * kThreads * 4 + j % 4, keys[j]);
+        }
+        ++offset;
+      }
+    }
+  }
+  __syncthreads();
+  const int count = s.count;
+  if (count >= K && count <= kCapacity) {
+    const uint32_t exact_cutoff = threshold<true>(s, row, count, K);
+    emit<true>(s, row, dst, count, exact_cutoff);
+  } else {
+    const uint32_t exact_cutoff = threshold<false>(s, row, length, K);
+    emit<false>(s, row, dst, length, exact_cutoff);
+  }
+}
+
+}  // namespace vllm::sampled_topk
diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu
index ada1c4ba39c0..703a748ce686 100644
--- a/csrc/libtorch_stable/sampler.cu
+++ b/csrc/libtorch_stable/sampler.cu
@@ -45,11 +45,6 @@ __global__ void apply_repetition_penalties_kernel(
   }
 }
 
-__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
-  uint32_t bits = __float_as_uint(x);
-  return (bits & 0x80000000) ? bits : ~bits & 0x7fffffff;
-}
-
 template <int step>
 static inline __device__ uint32_t extractBinIdx(float x) {
   if constexpr (step == 0) {
@@ -154,7 +149,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads,
 }
 
 template <int step, int kNumThreadsPerBlock, int kNumBins, int kNumFinalItems,
-          bool multipleBlocksPerRow, bool mergeBlocks, typename SmemFinalType,
+          bool multipleBlocksPerRow, bool mergeBlocks,
+          bool useTopK512Optimization, typename SmemFinalType,
           typename SmemOutputType>
 __device__ bool processHistogramStep(
     const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern,
@@ -180,10 +176,15 @@ __device__ bool processHistogramStep(
                     << patternShift;
   }
 
+  int negativeInfinityCount = 0;
   auto distributeToBins = [&](float logit, int /* idx */ = 0) {
     if (isPartialMatch<patternShift>(logit, logitPattern)) {
       uint32_t binIdx = extractBinIdx<step>(logit);
-      atomicAdd(&smemFinal.histo.data[binIdx], 1);
+      if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) {
+        ++negativeInfinityCount;
+      } else {
+        atomicAdd(&smemFinal.histo.data[binIdx], 1);
+      }
     }
   };
 
@@ -198,6 +199,12 @@ __device__ bool processHistogramStep(
       distributeToBins(logit, idx);
     }
   }
+  if constexpr (useTopK512Optimization) {
+    if (negativeInfinityCount > 0) {
+      const auto bin = extractBinIdx<step>(-INFINITY);
+      atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount);
+    }
+  }
   // Make sure the histogram is ready.
   __syncthreads();
 
@@ -254,6 +261,12 @@ __device__ bool processHistogramStep(
 
   // The threshold bin.
   thresholdBinIdx = smemThresholdBinIdx[0];
+  // Resolve a -inf cutoff with exact tie emission, keeping sort padding out.
+  const bool finalBinFits =
+      smemFinalBinSize[0] <= kNumFinalItems &&
+      !(useTopK512Optimization && step < 3 &&
+        isPartialMatch<patternShift>(-INFINITY, logitPattern) &&
+        thresholdBinIdx == extractBinIdx<step>(-INFINITY));
 
   auto processBins = [&](float logit, int idx) {
     if (isPartialMatch<patternShift>(logit, logitPattern)) {
@@ -262,8 +275,7 @@ __device__ bool processHistogramStep(
       // 1. This is step 0 and the threshold bin is small enough (no step 1)
       // 2. This is step >= 1 (where pattern matching filters correctly)
       // This prevents duplicates when step 0 and step 1 both run.
-      bool shouldWriteDirectly =
-          (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
+      bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1);
       if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
         // The element is part of the top-k selection
         int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
@@ -279,8 +291,7 @@ __device__ bool processHistogramStep(
       }
       if constexpr (step < 3) {
         // Only fill the final items for sorting if the threshold bin fits
-        if (binIdx == thresholdBinIdx &&
-            smemFinalBinSize[0] <= kNumFinalItems) {
+        if (binIdx == thresholdBinIdx && finalBinFits) {
           int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1);
           smemFinal.items.logits[dstIdx] = logit;
           if constexpr (mergeBlocks) {
@@ -325,20 +336,21 @@ __device__ bool processHistogramStep(
   __syncthreads();
 
   // Check if we should continue to next step
-  return smemFinalBinSize[0] > kNumFinalItems;
+  return !finalBinFits;
 }
 
 // Follows half - 11 - 11 - 10 bit iterations
 // Keep the adaptive kernel's device instantiation separate from legacy calls.
 template <int kNumThreadsPerBlock, int kNumBins, bool useRadixSort,
           bool multipleBlocksPerRow = false, bool mergeBlocks = false,
-          bool deviceLengthAware = false>
+          bool deviceLengthAware = false, bool useTopK512Optimization = false>
 static __device__ void topKPerRowJob(const int* indices, const float* logits,
                                      int rowStart, int rowEnd, int* outIndices,
                                      float* outLogits, int stride1, int topK) {
   static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks);
+  static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024);
   // The number of slots for the final pass.
-  static constexpr int kNumFinalItems = 2048;
+  static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 2048;
   // The number of elements per thread for the final sort.
   static constexpr int kNumFinalItemsPerThread =
       kNumFinalItems / kNumThreadsPerBlock;
@@ -404,7 +416,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
          rowIt += kNumThreadsPerBlock) {
       outIndices[rowIt] = -1;
       if constexpr (multipleBlocksPerRow) {
-        outLogits[rowIt] = -FLT_MAX;
+        outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX;
       }
     }
 
@@ -422,7 +434,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
   // Step 0: Process first 11 bits of half representation
   bool continueToNextStep =
       processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
-                           multipleBlocksPerRow, mergeBlocks>(
+                           multipleBlocksPerRow, mergeBlocks,
+                           useTopK512Optimization>(
           indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
           smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
           smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -431,7 +444,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
     // Step 1: Process next 11 bits
     continueToNextStep =
         processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
-                             multipleBlocksPerRow, mergeBlocks>(
+                             multipleBlocksPerRow, mergeBlocks,
+                             useTopK512Optimization>(
             indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
             smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
             smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -441,7 +455,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
     // Step 2: Process next 11 bits
     continueToNextStep =
         processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
-                             multipleBlocksPerRow, mergeBlocks>(
+                             multipleBlocksPerRow, mergeBlocks,
+                             useTopK512Optimization>(
             indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
             smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
             smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -450,7 +465,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
   if (continueToNextStep) {
     // Step 3: Process last 10 bits
     processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
-                         multipleBlocksPerRow, mergeBlocks>(
+                         multipleBlocksPerRow, mergeBlocks,
+                         useTopK512Optimization>(
         indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
         smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
         smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -460,50 +476,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
     // The histogram did not proceed to the final 10 bits, therefore we need to
     // sort the final items The logits of the elements to be sorted in the final
     // pass.
-    if constexpr (useRadixSort) {
-      // Sorting with radix sort
-      float finalLogits[kNumFinalItemsPerThread];
-      // The indices of the elements to be sorted in the final pass.
-      int finalIndices[kNumFinalItemsPerThread];
-
-#pragma unroll
-      for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
-        finalLogits[ii] = -FLT_MAX;
-      }
-
-      // Read the elements from SMEM.
-#pragma unroll
-      for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
-        int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
-        if (srcIdx < smemFinalDstIdx[0]) {
-          finalLogits[ii] = smemFinal.items.logits[srcIdx];
-          finalIndices[ii] = smemFinal.items.indices[srcIdx];
-        }
-      }
-      // Make sure the shared memory has been read.
-      __syncthreads();
-
-      // Sort the elements.
-      FinalSort(smemFinal.finalSort)
-          .SortDescendingBlockedToStriped(finalLogits, finalIndices);
-
-      // Copy the data back to the shared memory storage.
-      int baseIdx = smemFoundTopKValues[0];
-
-#pragma unroll
-      for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
-        int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
-        int dstIdx = baseIdx + srcIdx;
-
-        if (dstIdx < topK) {
-          smemOutput[dstIdx] = finalIndices[ii];
-          if constexpr (multipleBlocksPerRow) {
-            reinterpret_cast<float*>(smemOutput + topK)[dstIdx] =
-                finalLogits[ii];
-          }
-        }
-      }
-    } else {
+    auto insertionSort = [&]() {
       // Sorting with insertion sort
       auto baseIdx = smemFoundTopKValues[0];
       for (int i = threadIdx.x; i < smemFinalDstIdx[0];
@@ -525,6 +498,57 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
           }
         }
       }
+    };
+    if constexpr (useRadixSort) {
+      if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) {
+        insertionSort();
+      } else {
+        // Sorting with radix sort
+        float finalLogits[kNumFinalItemsPerThread];
+        // The indices of the elements to be sorted in the final pass.
+        int finalIndices[kNumFinalItemsPerThread];
+
+#pragma unroll
+        for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
+          finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX;
+          if constexpr (useTopK512Optimization) finalIndices[ii] = -1;
+        }
+
+        // Read the elements from SMEM.
+#pragma unroll
+        for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
+          int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
+          if (srcIdx < smemFinalDstIdx[0]) {
+            finalLogits[ii] = smemFinal.items.logits[srcIdx];
+            finalIndices[ii] = smemFinal.items.indices[srcIdx];
+          }
+        }
+        // Make sure the shared memory has been read.
+        __syncthreads();
+
+        // Sort the elements.
+        FinalSort(smemFinal.finalSort)
+            .SortDescendingBlockedToStriped(finalLogits, finalIndices);
+
+        // Copy the data back to the shared memory storage.
+        int baseIdx = smemFoundTopKValues[0];
+
+#pragma unroll
+        for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
+          int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
+          int dstIdx = baseIdx + srcIdx;
+
+          if (dstIdx < topK) {
+            smemOutput[dstIdx] = finalIndices[ii];
+            if constexpr (multipleBlocksPerRow) {
+              reinterpret_cast<float*>(smemOutput + topK)[dstIdx] =
+                  finalLogits[ii];
+            }
+          }
+        }
+      }
+    } else {
+      insertionSort();
     }
     __syncthreads();
   }
@@ -669,6 +693,70 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware(
                       stride1, topK);
 }
 
+#ifdef USE_ROCM
+// The grid and workspace stay fixed while graph replays change device lengths.
+__device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows,
+                                         int maxBlocks) {
+  int activeBlocks = 1;
+  if (numRows <= 32 && rowLength > 16 * 1024) {
+    activeBlocks = rowLength <= 128 * 1024   ? 8
+                   : rowLength <= 256 * 1024 ? 10
+                                             : 16;
+  } else if (numRows <= 128 && rowLength > 64 * 1024) {
+    activeBlocks = rowLength <= 512 * 1024 ? 4 : 10;
+  } else if (numRows > 128 && numRows <= 256 && rowLength > 256 * 1024) {
+    activeBlocks = numRows <= 160 ? 3 : 2;
+  }
+  // Every block in a merged row must contain at least k real indices.
+  return min(min(activeBlocks, maxBlocks), max(1, rowLength / 512));
+}
+
+template <int kNumThreadsPerBlock, bool mergeBlocks = false>
+static __global__
+__launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware(
+    const float* logits, const int* seqLens, int* outIndices,
+    int* outIndicesAux, float* outLogitsAux, int stride0, int next_n,
+    int seqLensIs2D, int maxBlocks) {
+  #if defined(__gfx950__)
+  constexpr int topK = 512;
+  constexpr int kNumBins = 2048;
+  const int rowIdx = blockIdx.x;
+  const int rowEnd =
+      seqLensIs2D
+          ? max(0, seqLens[rowIdx])
+          : max(0, seqLens[rowIdx / next_n] - next_n + rowIdx % next_n + 1);
+  const int activeBlocks =
+      gfx950TopK512ActiveBlocks(rowEnd, gridDim.x, maxBlocks);
+  outIndices += static_cast<int64_t>(rowIdx) * topK;
+  if constexpr (mergeBlocks) {
+    if (activeBlocks == 1) return;
+    const int64_t offset = static_cast<int64_t>(rowIdx) * maxBlocks * topK;
+    topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, false, true, false,
+                  true>(outIndicesAux + offset, outLogitsAux + offset, 0,
+                        activeBlocks * topK, outIndices, nullptr, 1, topK);
+  } else {
+    if (blockIdx.y >= activeBlocks) return;
+    logits += static_cast<int64_t>(rowIdx) * stride0;
+    if (activeBlocks == 1) {
+      topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, false, false, false,
+                    true>(nullptr, logits, 0, rowEnd, outIndices, nullptr, 1,
+                          topK);
+      return;
+    }
+    const int blockSize = rowEnd / activeBlocks;
+    const int rowStart = blockSize * blockIdx.y;
+    const int blockRowEnd =
+        blockIdx.y + 1 == activeBlocks ? rowEnd : rowStart + blockSize;
+    const int64_t offset =
+        (static_cast<int64_t>(rowIdx) * maxBlocks + blockIdx.y) * topK;
+    topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, true, false, false,
+                  true>(nullptr, logits, rowStart, blockRowEnd,
+                        outIndicesAux + offset, outLogitsAux + offset, 1, topK);
+  }
+  #endif
+}
+#endif
+
 }  // namespace vllm
 
 void apply_repetition_penalties_(
@@ -731,6 +819,60 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
   // the same seq_len and the kernel computes the per-row offset itself.
   int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0;
 
+#ifdef USE_ROCM
+  const bool useGfx950DeviceLengthAwareTopK512 =
+      topK == 512 && stride1 == 1 && numRows > 0 && numRows <= 384 &&
+      numColumns <= 1024 * 1024 &&
+      std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0;
+  if (useGfx950DeviceLengthAwareTopK512) {
+    int multipleBlocksPerRowConfig = 1;
+    if (numRows <= 32 && numColumns > 16 * 1024) {
+      multipleBlocksPerRowConfig = 16;
+    } else if (numRows <= 128 && numColumns > 64 * 1024) {
+      multipleBlocksPerRowConfig = 10;
+    } else if (numRows <= 160 && numColumns > 256 * 1024) {
+      multipleBlocksPerRowConfig = 3;
+    } else if (numRows <= 256 && numColumns > 256 * 1024) {
+      multipleBlocksPerRowConfig = 2;
+    }
+
+    constexpr int kNumThreadsPerSplitBlock = 1024;
+    if (multipleBlocksPerRowConfig == 1) {
+      vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerSplitBlock>
+          <<<numRows, kNumThreadsPerSplitBlock, 2 * topK * sizeof(int32_t),
+             stream>>>(logits.const_data_ptr<float>(),
+                       seqLens.const_data_ptr<int>(),
+                       indices.mutable_data_ptr<int>(), nullptr, nullptr,
+                       static_cast<int>(stride0), static_cast<int>(next_n),
+                       seqLensIs2D, multipleBlocksPerRowConfig);
+      return;
+    }
+    const auto outIndicesAux = torch::stable::empty(
+        {numRows, multipleBlocksPerRowConfig, topK},
+        torch::headeronly::ScalarType::Int, std::nullopt, logits.device());
+    const auto outLogitsAux = torch::stable::empty(
+        {numRows, multipleBlocksPerRowConfig, topK},
+        torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
+    vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerSplitBlock>
+        <<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerSplitBlock,
+           2 * topK * sizeof(int32_t), stream>>>(
+            logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            indices.mutable_data_ptr<int>(),
+            outIndicesAux.mutable_data_ptr<int>(),
+            outLogitsAux.mutable_data_ptr<float>(), static_cast<int>(stride0),
+            static_cast<int>(next_n), seqLensIs2D, multipleBlocksPerRowConfig);
+    constexpr int kNumThreadsPerBlockMerge = 1024;
+    vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerBlockMerge, true>
+        <<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
+            nullptr, seqLens.const_data_ptr<int>(),
+            indices.mutable_data_ptr<int>(),
+            outIndicesAux.mutable_data_ptr<int>(),
+            outLogitsAux.mutable_data_ptr<float>(), 0, static_cast<int>(next_n),
+            seqLensIs2D, multipleBlocksPerRowConfig);
+    return;
+  }
+#endif
+
   if (numColumns < kSortingAlgorithmThreshold) {
     // Use insertion sort
     vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu
index b3e1a8623f3e..fb3d946b96b0 100644
--- a/csrc/libtorch_stable/topk.cu
+++ b/csrc/libtorch_stable/topk.cu
@@ -9,6 +9,7 @@
 
 #ifndef USE_ROCM
   #include "persistent_topk.cuh"
+  #include "sampled_topk.cuh"
 #endif
 
 namespace {
@@ -36,7 +37,22 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
     max_smem_per_block = device_prop->sharedMemPerBlockOptin;
   }
 
-  if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
+  // Allow static fallback storage in addition to the 128 KiB dynamic buffer.
+  if (num_rows > 64 &&
+      max_seq_len >= vllm::sampled_topk::kMinSampledLength<TopK> &&
+      max_smem_per_block >= 144 * 1024) {
+    auto kernel = vllm::sampled_topk::sampled_topk_kernel<TopK>;
+    constexpr size_t smem_size =
+        vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC;
+    cudaError_t status = cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+    STD_TORCH_CHECK(status == cudaSuccess,
+                    "sampled_topk smem failed: ", cudaGetErrorString(status));
+    kernel<<<num_rows, vllm::sampled_topk::kThreads, smem_size, stream>>>(
+        logits.const_data_ptr<float>(), lengths.const_data_ptr<int32_t>(),
+        output.mutable_data_ptr<int32_t>(), stride,
+        static_cast<int>(std::min(max_seq_len, logits.size(1))));
+  } else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
     cudaError_t status =
         vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
             logits.const_data_ptr<float>(), output.mutable_data_ptr<int32_t>(),
diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp
index 3781eeaebf26..3833c95fde00 100644
--- a/csrc/libtorch_stable/torch_bindings.cpp
+++ b/csrc/libtorch_stable/torch_bindings.cpp
@@ -327,21 +327,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
   // conditionally compiled so impl registration is in source file
   ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
-
-  // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
-  ops.def(
-      "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
-      "Tensor? b_zeros, "
-      "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
-      "Tensor!? b_zeros_reorder, "
-      "int K, int N, int N_32align) -> ()");
-
-  // AllSpark quantization ops
-  ops.def(
-      "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
-      "Tensor? b_qzeros, "
-      "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
-      "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
 #endif
 
   // Merge attn states
@@ -425,12 +410,25 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "bool is_neox, Tensor position_ids, "
       "int forced_token_heads_per_warp=-1) -> ()");
 
+  // q_head_padded is the padded Q head count of the returned tensor, or 0 to
+  // do the KV insert alone and return an empty tensor.  The Q knobs are
+  // independent: apply_q_norm and apply_q_rope each drop that step for Q
+  // alone (KV is always rotated), and is_q_interleaved reads and writes Q in
+  // FlashMLA's mega-attention chunk-interleaved layout, which moves the
+  // padding heads to the tail of every head-dim chunk.
+  ops.def(
+      "fused_deepseek_v4_kv_rope_insert("
+      "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, "
+      "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, "
+      "bool kv_mxfp8=False) -> "
+      "()");
   ops.def(
       "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
       "Tensor q_in, Tensor kv, Tensor! k_cache, "
       "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
       "int q_head_padded, float eps, int cache_block_size, "
-      "bool apply_q_norm=True) -> Tensor");
+      "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, "
+      "bool is_q_interleaved=False) -> Tensor");
 
   // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
   // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
@@ -781,9 +779,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
 
   // DSV3 fused A GEMM: conditionally compiled so impl registration is in
   // source file (dsv3_fused_a_gemm.cu)
-
-  // AllSpark ops: conditionally compiled so impl registrations are in source
-  // files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
 #endif
 
   ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states));
@@ -807,6 +802,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // Positional encoding kernels (shared CUDA/ROCm)
   ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
   ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
+  ops.impl("fused_deepseek_v4_kv_rope_insert",
+           TORCH_BOX(&fused_deepseek_v4_kv_rope_insert));
   ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
            TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
   ops.impl(
@@ -976,14 +973,68 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
       "                     str kv_cache_dtype,"
       "                     Tensor scale) -> ()");
 
-  // Grouped concat_and_cache_mla across all layers (bf16 only). Each
-  // layer's cache base pointer is read from kv_cache_ptrs.
+  // Grouped concat_and_cache_mla across all layers. Each layer's cache base
+  // pointer and optional plain-FP8 scale are read from device tensors.
   ops.def(
       "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe,"
       "                             Tensor kv_cache_ptrs,"
       "                             Tensor slot_mapping,"
       "                             int block_size, int block_stride,"
-      "                             int entry_stride) -> ()");
+      "                             int entry_stride,"
+      "                             Tensor? kv_scales=None,"
+      "                             str kv_cache_dtype='auto') -> ()");
+
+#ifndef USE_ROCM
+  ops.def(
+      "hisparse_resolve_residency(Tensor host_cache,"
+      "                 Tensor! hot_cache,"
+      "                 Tensor hot_block_table,"
+      "                 Tensor global_indices,"
+      "                 Tensor! hot_indices,"
+      "                 Tensor! device_global_indices,"
+      "                 Tensor! lru_slots,"
+      "                 Tensor? request_state_indices,"
+      "                 int region_stride,"
+      "                 Tensor(a!)? miss_mask=None,"
+      "                 Tensor(b!)? stats=None,"
+      "                 Tensor(c!)? attention_indices=None,"
+      "                 int attention_block_stride=0,"
+      "                 Tensor? request_ids=None,"
+      "                 Tensor? source_block_table=None,"
+      "                 int source_block_size=0,"
+      "                 Tensor(d!)? resolved_global_indices=None,"
+      "                 Tensor(e!)? valid_counts=None,"
+      "                 Tensor(f!)? swap_host_physical_rows=None,"
+      "                 Tensor(g!)? swap_device_physical_rows=None,"
+      "                 Tensor(h!)? swap_counts=None,"
+      "                 Tensor? resident_block_table=None,"
+      "                 int resident_block_size=0,"
+      "                 int resident_null_block=0) -> ()");
+
+  ops.def(
+      "hisparse_invalidate_written_slots(Tensor! device_global_indices,"
+      "                                   Tensor request_state_indices,"
+      "                                   Tensor req_id_per_token,"
+      "                                   Tensor written_slots) -> ()");
+
+  ops.def(
+      "hisparse_gather_plan(Tensor host_cache,"
+      "                     Tensor! hot_cache,"
+      "                     Tensor global_indices,"
+      "                     Tensor hot_indices,"
+      "                     Tensor miss_mask,"
+      "                     Tensor? request_state_indices,"
+      "                     Tensor(a!)? attention_indices=None,"
+      "                     int attention_block_stride=0) -> ()");
+
+  ops.def(
+      "hisparse_gather_compact(Tensor host_cache,"
+      "                        Tensor! hot_cache,"
+      "                        Tensor miss_global_indices,"
+      "                        Tensor miss_hot_indices,"
+      "                        Tensor miss_counts) -> ()");
+
+#endif  // !USE_ROCM
 
   // Rotate Q and K, then write to kv cache for MLA
   ops.def(
@@ -1021,7 +1072,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
   ops.def(
       "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
       "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
-      "seq_starts) -> ()");
+      "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? "
+      "device_row_ids=None) -> ()");
 
   ops.def(
       "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, "
@@ -1089,6 +1141,15 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
   ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
   ops.impl("concat_and_cache_mla_grouped",
            TORCH_BOX(&concat_and_cache_mla_grouped));
+
+#ifndef USE_ROCM
+  ops.impl("hisparse_resolve_residency",
+           TORCH_BOX(&hisparse_resolve_residency));
+  ops.impl("hisparse_invalidate_written_slots",
+           TORCH_BOX(&hisparse_invalidate_written_slots));
+  ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan));
+  ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact));
+#endif  // !USE_ROCM
   ops.impl("concat_and_cache_mla_rope_fused",
            TORCH_BOX(&concat_and_cache_mla_rope_fused));
   ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh
index 58d19de73494..32925c0c3c95 100644
--- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh
+++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh
@@ -67,104 +67,6 @@ using fp8_type = __hip_fp8_e4m3_fnuz;
 using fp8x2_type = __hip_fp8x2_e4m3_fnuz;
     #endif
 
-// fp8 -> half
-template <>
-__inline__ __device__ uint16_t
-vec_conversion<uint16_t, uint8_t>(const uint8_t& a) {
-  return __hip_cvt_fp8_to_halfraw(a, fp8_type::__default_interpret).x;
-}
-
-// fp8x2 -> half2
-template <>
-__inline__ __device__ uint32_t
-vec_conversion<uint32_t, uint16_t>(const uint16_t& a) {
-  union {
-    __half2_raw h2r;
-    uint32_t ui32;
-  } tmp;
-  tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret);
-  return tmp.ui32;
-}
-
-// fp8x4 -> half2x2
-template <>
-__inline__ __device__ uint2 vec_conversion<uint2, uint32_t>(const uint32_t& a) {
-  union {
-    uint2 u32x2;
-    uint32_t u32[2];
-  } tmp;
-  tmp.u32[0] = vec_conversion<uint32_t, uint16_t>((uint16_t)a);
-  tmp.u32[1] = vec_conversion<uint32_t, uint16_t>((uint16_t)(a >> 16U));
-  return tmp.u32x2;
-}
-
-// fp8x8 -> half2x4
-template <>
-__inline__ __device__ uint4 vec_conversion<uint4, uint2>(const uint2& a) {
-  union {
-    uint4 u64x2;
-    uint2 u64[2];
-  } tmp;
-  tmp.u64[0] = vec_conversion<uint2, uint32_t>(a.x);
-  tmp.u64[1] = vec_conversion<uint2, uint32_t>(a.y);
-  return tmp.u64x2;
-}
-
-using __nv_bfloat16 = __hip_bfloat16;
-
-// fp8 -> __nv_bfloat16
-template <>
-__inline__ __device__ __nv_bfloat16
-vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a) {
-  fp8_type f8;
-  f8.__x = a;
-  return __float2bfloat16(static_cast<float>(f8));
-}
-
-using __nv_bfloat162 = __hip_bfloat162;
-
-// fp8x2 -> __nv_bfloat162
-template <>
-__inline__ __device__ __nv_bfloat162
-vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a) {
-  __nv_bfloat162 res;
-  res.x = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a);
-  res.y = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U));
-  return res;
-}
-
-// fp8x4 -> bf16_4_t
-template <>
-__inline__ __device__ bf16_4_t
-vec_conversion<bf16_4_t, uint32_t>(const uint32_t& a) {
-  bf16_4_t res;
-  res.x = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a);
-  res.y = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U));
-  return res;
-}
-
-// fp8x8 -> bf16_8_t
-template <>
-__inline__ __device__ bf16_8_t vec_conversion<bf16_8_t, uint2>(const uint2& a) {
-  bf16_4_t tmp1, tmp2;
-  tmp1 = vec_conversion<bf16_4_t, uint32_t>(a.x);
-  tmp2 = vec_conversion<bf16_4_t, uint32_t>(a.y);
-  bf16_8_t res;
-  res.x = tmp1.x;
-  res.y = tmp1.y;
-  res.z = tmp2.x;
-  res.w = tmp2.y;
-  return res;
-}
-
-// fp8 -> float
-template <>
-__inline__ __device__ float vec_conversion<float, uint8_t>(const uint8_t& a) {
-  fp8_type f8;
-  f8.__x = a;
-  return static_cast<float>(f8);
-}
-
 // fp8x2 -> float2
 template <>
 __inline__ __device__ float2
@@ -193,138 +95,6 @@ vec_conversion<float4, uint32_t>(const uint32_t& a) {
   return res;
 }
 
-// fp8x8 -> float8
-template <>
-__inline__ __device__ Float8_ vec_conversion<Float8_, uint2>(const uint2& a) {
-  Float4_ tmp1, tmp2;
-  tmp1 = vec_conversion<Float4_, uint32_t>(a.x);
-  tmp2 = vec_conversion<Float4_, uint32_t>(a.y);
-  Float8_ res;
-  res.x = tmp1.x;
-  res.y = tmp1.y;
-  res.z = tmp2.x;
-  res.w = tmp2.y;
-  return res;
-}
-
-// half -> fp8
-template <>
-__inline__ __device__ uint8_t
-vec_conversion<uint8_t, uint16_t>(const uint16_t& a) {
-  __half_raw tmp;
-  tmp.x = a;
-  return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation,
-                                  fp8_type::__default_interpret);
-}
-
-template <>
-__inline__ __device__ uint16_t
-vec_conversion<uint16_t, uint32_t>(const uint32_t& a) {
-  union {
-    uint32_t ui32;
-    __half2_raw h2r;
-  } tmp;
-  tmp.ui32 = a;
-  return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation,
-                                     fp8_type::__default_interpret);
-}
-
-// bf16 -> fp8
-template <>
-__inline__ __device__ uint8_t
-vec_conversion<uint8_t, __nv_bfloat16>(const __nv_bfloat16& a) {
-  return __hip_cvt_float_to_fp8(__bfloat162float(a),
-                                fp8_type::__default_saturation,
-                                fp8_type::__default_interpret);
-}
-
-// float -> fp8
-template <>
-__inline__ __device__ uint8_t vec_conversion<uint8_t, float>(const float& a) {
-  return __hip_cvt_float_to_fp8(a, fp8_type::__default_saturation,
-                                fp8_type::__default_interpret);
-}
-
-// float2 -> half2
-template <>
-__inline__ __device__ uint32_t
-vec_conversion<uint32_t, float2>(const float2& a) {
-  union {
-    half2 float16;
-    uint32_t uint32;
-  };
-
-  float16 = __float22half2_rn(a);
-  return uint32;
-}
-
-// Float4 -> half2x2
-template <>
-__inline__ __device__ uint2 vec_conversion<uint2, Float4_>(const Float4_& a) {
-  uint2 b;
-  float2 val;
-  val.x = a.x.x;
-  val.y = a.x.y;
-  b.x = vec_conversion<uint32_t, float2>(val);
-
-  val.x = a.y.x;
-  val.y = a.y.y;
-  b.y = vec_conversion<uint32_t, float2>(val);
-  return b;
-}
-
-// Float4 -> float4
-template <>
-__inline__ __device__ float4 vec_conversion<float4, Float4_>(const Float4_& a) {
-  float4 b;
-  b.x = a.x.x;
-  b.y = a.x.y;
-  b.z = a.y.x;
-  b.w = a.y.y;
-  return b;
-}
-
-// Float8 -> half2x4
-template <>
-__inline__ __device__ uint4 vec_conversion<uint4, Float8_>(const Float8_& a) {
-  uint4 b;
-  b.x = vec_conversion<uint32_t, float2>(a.x);
-  b.y = vec_conversion<uint32_t, float2>(a.y);
-  b.z = vec_conversion<uint32_t, float2>(a.z);
-  b.w = vec_conversion<uint32_t, float2>(a.w);
-  return b;
-}
-
-// float2 -> bfloat162
-template <>
-__inline__ __device__ __nv_bfloat162
-vec_conversion<__nv_bfloat162, float2>(const float2& a) {
-  __nv_bfloat162 b = __float22bfloat162_rn(a);
-  return b;
-}
-
-// Float4 -> bfloat162x2
-template <>
-__inline__ __device__ bf16_4_t
-vec_conversion<bf16_4_t, Float4_>(const Float4_& a) {
-  bf16_4_t b;
-  b.x = __float22bfloat162_rn(a.x);
-  b.y = __float22bfloat162_rn(a.y);
-  return b;
-}
-
-// Float8 -> bfloat162x4
-template <>
-__inline__ __device__ bf16_8_t
-vec_conversion<bf16_8_t, Float8_>(const Float8_& a) {
-  bf16_8_t b;
-  b.x = __float22bfloat162_rn(a.x);
-  b.y = __float22bfloat162_rn(a.y);
-  b.z = __float22bfloat162_rn(a.z);
-  b.w = __float22bfloat162_rn(a.w);
-  return b;
-}
-
 /* Scaled and vectorized conversions, for data exchange between high and low
    precision domains
 
@@ -335,6 +105,7 @@ vec_conversion<bf16_8_t, Float8_>(const Float8_& a) {
  */
 
 using __nv_bfloat16 = __hip_bfloat16;
+using __nv_bfloat162 = __hip_bfloat162;
 
 // fp8 -> __nv_bfloat16
 template <>
@@ -613,17 +384,6 @@ scaled_vec_conversion<uint32_t, float4>(const float4& a, float scale) {
 }
   #endif  // ENABLE_FP8
 
-template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
-__inline__ __device__ Tout convert(const Tin& x) {
-  #ifdef ENABLE_FP8
-  if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) {
-    return vec_conversion<Tout, Tin>(x);
-  }
-  #endif
-  assert(false);
-  return {};  // Squash missing return statement warning
-}
-
 template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
 __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
   #ifdef ENABLE_FP8
diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh
index 11a9b0818dc6..7fcf541fd64d 100644
--- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh
+++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh
@@ -282,19 +282,6 @@ __inline__ __device__ float4 scaled_vec_conversion<float4, uint32_t>(
 }
   #endif  // ENABLE_FP8
 
-template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
-__inline__ __device__ Tout convert(const Tin& x) {
-  #if 0  // Disable the following code to reduce the binary size.
-  if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) {
-    return vec_conversion<Tout, Tin>(x, __NV_E4M3);
-  } else if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E5M2) {
-    return vec_conversion<Tout, Tin>(x, __NV_E5M2);
-  }
-  #endif
-  assert(false);
-  __builtin_unreachable();  // Suppress missing return statement warning
-}
-
 template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
 __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
   #ifdef ENABLE_FP8
diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu
index b747b870e23f..2a11653a58d7 100644
--- a/csrc/rocm/attention.cu
+++ b/csrc/rocm/attention.cu
@@ -342,22 +342,6 @@ typedef union u64_cvt {
   int64_t i64;
 } _T8x8;
 
-__device__ __forceinline__ _B8x8 convert_b16x8(const _B16x8& input,
-                                               _T8x8& Mtemp) {
-  _T8x8 Qtmp8x8;
-
-  for (int i = 0; i < 2; i++) {
-    floatx4 q_out = {0, 0, 0, 0};
-    q_out = gcn_mfma16x16x16_instr<_Float16, 0, 0, 0>(Mtemp.b64, input.xy[i],
-                                                      q_out);
-    Qtmp8x8.b16x4[i * 2] =
-        __builtin_amdgcn_cvt_pk_fp8_f32(q_out[0], q_out[1], 0, false);
-    Qtmp8x8.b16x4[i * 2 + 1] =
-        __builtin_amdgcn_cvt_pk_fp8_f32(q_out[2], q_out[3], 0, false);
-  }
-  return Qtmp8x8.b8x8;
-}
-
 __device__ float warpReduceMax(float val) {
   for (int offset = warpSize / 2; offset > 0; offset /= 2) {
     val = max(
@@ -2583,23 +2567,6 @@ __device__ __forceinline__ float to_float(const T& inp) {
   }
 }
 
-template <typename T>
-__device__ __forceinline__ float to_float_b16(const bit16_t& inp) {
-  union tmpcvt {
-    bit16_t u;
-    _Float16 f;
-    __hip_bfloat16 b;
-  } t16;
-  t16.u = inp;
-  if constexpr (std::is_same<T, _Float16>::value) {
-    return (float)t16.f;
-  } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
-    return __bfloat162float(t16.b);
-  } else {
-    static_assert(false, "unsupported 16b dtype");
-  }
-}
-
 template <typename T>
 __device__ __forceinline__ T from_float(const float& inp) {
   if constexpr (std::is_same<T, _Float16>::value) {
diff --git a/csrc/rocm/moe_q_gemm_rdna3.cu b/csrc/rocm/moe_q_gemm_rdna3.cu
index 6c25ed7e4bc0..61ffbb9987e9 100644
--- a/csrc/rocm/moe_q_gemm_rdna3.cu
+++ b/csrc/rocm/moe_q_gemm_rdna3.cu
@@ -67,21 +67,6 @@ __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
   return result;
 }
 
-__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
-                                           const bf16_t* a_ptr) {
-  float result = 0.0f;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    uint32_t aw;
-    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
-    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
-    float a_y = __uint_as_float(aw & 0xFFFF0000u);
-    result = __fmaf_rn(dq[2 * i + 0], a_x, result);
-    result = __fmaf_rn(dq[2 * i + 1], a_y, result);
-  }
-  return result;
-}
-
 __forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
                                                    half2 v23) {
   unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu
index 104af66d1cc5..bf2368269964 100644
--- a/csrc/rocm/q_gemm_rdna3.cu
+++ b/csrc/rocm/q_gemm_rdna3.cu
@@ -115,61 +115,6 @@ __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
   return result;
 }
 
-__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4],
-                                           const bf16_t* a_ptr) {
-  // RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in
-  // the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc
-  // therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised
-  // fallback (single-element FMAs or fp32 round-trips), which empirically
-  // runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16
-  // decode path was paying that tax in full, scaling linearly with M (the
-  // fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate
-  // and the kernel becomes memory-bound).
-  //
-  // Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs)
-  // and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA
-  // count, but each FMA is fast. Bonus: the accumulator is now fp32
-  // throughout instead of bf16, which is also numerically more accurate
-  // (no compounding bf16-rounding inside the dot loop).
-  float result = 0.0f;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    uint32_t aw, dw;
-    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
-    __builtin_memcpy(&dw, &dq[i], sizeof(uint32_t));
-    // bf16 in low 16 bits  → fp32 by left-shifting into the upper half.
-    // bf16 in high 16 bits → already aligned with fp32's upper half.
-    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
-    float a_y = __uint_as_float(aw & 0xFFFF0000u);
-    float d_x = __uint_as_float((dw & 0xFFFFu) << 16);
-    float d_y = __uint_as_float(dw & 0xFFFF0000u);
-    result = __fmaf_rn(d_x, a_x, result);
-    result = __fmaf_rn(d_y, a_y, result);
-  }
-  return result;
-}
-
-// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already
-// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t
-// overload above does for dq (still need to widen A from bf16). Wins more
-// at high N: the bf162_t version's per-call widening cost scales with the
-// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only
-// for A widening (M_COUNT × 4 × 4 widens, half as many).
-__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
-                                           const bf16_t* a_ptr) {
-  float result = 0.0f;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    uint32_t aw;
-    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
-    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
-    float a_y = __uint_as_float(aw & 0xFFFF0000u);
-    result = __fmaf_rn(dq[2 * i + 0], a_x, result);
-    result = __fmaf_rn(dq[2 * i + 1], a_y, result);
-  }
-  return result;
-}
-
 // ---------------------------------------------------------------------------
 // Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS).
 // RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those
@@ -327,9 +272,8 @@ __global__ void gemm_q4_kernel_rdna3(const T* __restrict__ a,
 
   // Per-column dequant constants. We hold one set of (z, y) pairs per column.
   // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper-
-  // nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant
-  // produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA
-  // bypass for the missing v_pk_fma_bf16 on gfx11.
+  // nibble-*16 trick. bf16 uses fp32 scalars (z, y) for the inline v_dot2
+  // bias correction.
   half2 z1z16_h[4][2], y1y16_h[4][2];
   float z_b_f[4], y_b_f[4];
 
diff --git a/csrc/rocm/qdq_4_rdna3.cuh b/csrc/rocm/qdq_4_rdna3.cuh
index f4be668c1f08..43a7b33ba153 100644
--- a/csrc/rocm/qdq_4_rdna3.cuh
+++ b/csrc/rocm/qdq_4_rdna3.cuh
@@ -113,63 +113,6 @@ __forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4],
 // bf16 path
 // ---------------------------------------------------------------------------
 
-// Bit-trick magic for bf16:
-//   bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0).
-//   For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because
-//   bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So
-//   ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi).
-//
-// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble
-// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16]
-// via a single 4/8/12-bit right-shift before the OR. That costs one extra
-// shift per pair vs fp16, but keeps the FMA structure identical.
-__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero,
-                                                     bf16_t scale,
-                                                     bf162_t& z_prep,
-                                                     bf162_t& y_prep) {
-  // z = scale * -(128 + zero); y = scale.
-  float scale_f = __bfloat162float(scale);
-  float zf = -(128.0f + (float)zero) * scale_f;
-  bf16_t zb = __float2bfloat16(zf);
-  z_prep = __bfloat162bfloat162(zb);
-  y_prep = __bfloat162bfloat162(scale);
-}
-
-__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa,
-                                                    bf162_t (&dq)[4],
-                                                    bf162_t z_prep,
-                                                    bf162_t y_prep) {
-  const uint32_t c0 = 0x43004300;
-
-  union {
-    uint32_t u;
-    bf162_t b2;
-  } q0, q1, q2, q3;
-  q0.u = ((qa >> 0) & 0x000F000F) | c0;   // bf162(128+q[0], 128+q[1])
-  q1.u = ((qa >> 4) & 0x000F000F) | c0;   // bf162(128+q[2], 128+q[3])
-  q2.u = ((qa >> 8) & 0x000F000F) | c0;   // bf162(128+q[4], 128+q[5])
-  q3.u = ((qa >> 12) & 0x000F000F) | c0;  // bf162(128+q[6], 128+q[7])
-
-  // dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale
-  dq[0] = __hfma2(q0.b2, y_prep, z_prep);
-  dq[1] = __hfma2(q1.b2, y_prep, z_prep);
-  dq[2] = __hfma2(q2.b2, y_prep, z_prep);
-  dq[3] = __hfma2(q3.b2, y_prep, z_prep);
-}
-
-// ---------------------------------------------------------------------------
-// bf16-input → fp32-output dequant (RDNA3 scalar path).
-//
-// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow
-// fallback. Rather than computing dq in bf16 and widening at FMA time in
-// the dot product, we widen to fp32 here once (a free left-shift by 16) and
-// emit the (q - zero) * scale FMA directly in fp32. This:
-//   * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32.
-//   * Eliminates 4× bf16→fp32 widens that the dot product would do.
-//   * Keeps the dot product accumulator in fp32 without a roundtrip.
-//
-// Output: fp32 dq[8], one element per K position (consumed by the
-// fp32-overload of dot22_8_f in q_gemm_rdna3.cu).
 __forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
                                                          bf16_t scale,
                                                          float& z_prep,
@@ -179,60 +122,6 @@ __forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
   y_prep = scale_f;
 }
 
-// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32
-// values 128+nibble, without folding scale/zero. The caller folds scale/zb
-// into the accumulator outside the inner loop using a precomputed sum_a,
-// which saves ~27% of the FMA count vs the per-col-dequant approach above
-// (only beneficial at M_COUNT=1; break-even at M_COUNT=2).
-//
-// Cost: 0 FMAs (pure bit-trick + as_float reinterprets).
-__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa,
-                                                           float (&q_f32)[8]) {
-  const uint32_t c0 = 0x43004300;
-  const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
-  const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
-  const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
-  const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
-  q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16);
-  q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u);
-  q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16);
-  q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u);
-  q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16);
-  q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u);
-  q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16);
-  q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u);
-}
-
-__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa,
-                                                        float (&dq)[8],
-                                                        float z_prep,
-                                                        float y_prep) {
-  const uint32_t c0 = 0x43004300;
-  const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
-  const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
-  const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
-  const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
-  // bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16
-  // (just zero-extends the mantissa from 7 to 23 bits; exponent preserved).
-  const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16);
-  const float q0y = __uint_as_float(q0 & 0xFFFF0000u);
-  const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16);
-  const float q1y = __uint_as_float(q1 & 0xFFFF0000u);
-  const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16);
-  const float q2y = __uint_as_float(q2 & 0xFFFF0000u);
-  const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16);
-  const float q3y = __uint_as_float(q3 & 0xFFFF0000u);
-  // dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale
-  dq[0] = __fmaf_rn(q0x, y_prep, z_prep);
-  dq[1] = __fmaf_rn(q0y, y_prep, z_prep);
-  dq[2] = __fmaf_rn(q1x, y_prep, z_prep);
-  dq[3] = __fmaf_rn(q1y, y_prep, z_prep);
-  dq[4] = __fmaf_rn(q2x, y_prep, z_prep);
-  dq[5] = __fmaf_rn(q2y, y_prep, z_prep);
-  dq[6] = __fmaf_rn(q3x, y_prep, z_prep);
-  dq[7] = __fmaf_rn(q3y, y_prep, z_prep);
-}
-
 }  // namespace gptq_rdna3
 }  // namespace vllm
 
diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu
index 615cdabed585..79af0814550d 100644
--- a/csrc/rocm/skinny_gemms.cu
+++ b/csrc/rocm/skinny_gemms.cu
@@ -8,12 +8,20 @@
 
 #include <stdexcept>
 #include <algorithm>
+#include <map>
+#include <mutex>
+#include <utility>
 
 #include "../cuda_compat.h"
 #include "dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"
 #include "core/batch_invariant.hpp"
 
+// Number of streams the pre-allocated split-K pool covers. ~7.5 MiB per slot.
+static constexpr int64_t kWvSlots = 8;
+// Slots held back for capturing streams, which cannot allocate one themselves.
+static constexpr int64_t kWvCaptureSlots = 8;
+
 // TODO(rasmith): The kernels in this file are susceptible to integer overflow
 // issues, do not take strides, and are unable to handle PyTorch tensors that
 // return is_contiguous() as False (the tensors may actually be contiguous
@@ -1364,9 +1372,18 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   };
   using big4 = __attribute__((__vector_size__(4 * sizeof(bigType)))) __bf16;
 
-  __shared__ scalar_t stg[WvPrGrp * WVLDS / GrpsShrB];
+  // The main-loop staging areas are separated from the split-K readback by a
+  // __syncthreads(), so the readback may reuse all of LDS.
+  __shared__ union {
+    struct {
+      scalar_t stg[WvPrGrp * WVLDS / GrpsShrB];
+      scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB];
+    } mainloop;
+    scalar_t all[max_lds_len];
+  } lds;
+  auto& stg = lds.mainloop.stg;
+  auto& s = lds.mainloop.s;
   unsigned int* myStg = (unsigned int*)(&stg[WVLDS * (threadIdx.y / GrpsShrB)]);
-  __shared__ scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB];
 
   #ifndef WVSPLITKRC_1KPASS
   constexpr int TUC_ = (THRDS * UNRL * A_CHUNK);
@@ -1682,33 +1699,69 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     if (my_cntr + 1 == k_rnd) {
       cntr[adr_] = 0;  // clear for next round
       if constexpr (DTRMNSTC) {
+        constexpr uint32_t NT = N / NTILE / GrpsShrB;
+        // The readback stages k-shard partials through LDS, which bounds how
+        // many shards may be in flight at once.
+        constexpr uint32_t LDS_F4 =
+            max_lds_len * sizeof(scalar_t) / sizeof(float4);
+        constexpr uint32_t KSBMAX = LDS_F4 / (THRDS * 4 * NT);
+        static_assert(KSBMAX >= 1, "LDS cannot stage even one k-shard");
+        const uint32_t KSB = (k_rnd < KSBMAX) ? k_rnd : KSBMAX;
+
+        // First batch is peeled; it is the only batch for most shapes.
   #pragma unroll
-        for (int ks = 0; ks < k_rnd; ks++) {
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
+        for (uint32_t kb = 0; kb < KSB; kb++) {
+          for (uint32_t nt = 0; nt < NT; nt++) {
             int g_nindx =
                 (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
             int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
             __builtin_amdgcn_global_load_lds(
-                (float4*)(&glbl[g_adr + M * N * ks]),
-                &(((float4*)s)[(threadIdx.y * THRDS) + ks * THRDS * 4 +
-                               nt * THRDS * 4 * k_rnd]),
+                (float4*)(&glbl[g_adr + M * N * kb]),
+                &(((float4*)lds.all)[(threadIdx.y * THRDS) + kb * THRDS * 4 +
+                                     nt * THRDS * 4 * KSB]),
                 16, 0, 0);
           }
         }
-        if (BIAS)
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
+        if (BIAS)  // overlaps the DMAs above
+          for (uint32_t nt = 0; nt < NT; nt++) {
             for (uint32_t j = 0; j < 4; j++) {
               int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
                           (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
               biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
             }
           }
-        asm volatile("s_waitcnt 0");
-        for (int ks = 0; ks < k_rnd; ks++) {
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            float4 eval = ((float4*)s)[(threadIdx.x + threadIdx.y * THRDS) +
-                                       ks * THRDS * 4 + nt * THRDS * 4 * k_rnd];
-            vals[nt].f4 += eval;
+        asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
+        for (uint32_t kb = 0; kb < KSB; kb++) {
+          for (uint32_t nt = 0; nt < NT; nt++) {
+            vals[nt].f4 +=
+                ((float4*)lds.all)[(threadIdx.x + threadIdx.y * THRDS) +
+                                   kb * THRDS * 4 + nt * THRDS * 4 * KSB];
+          }
+        }
+
+        for (uint32_t ks0 = KSB; ks0 < k_rnd; ks0 += KSB) {
+          const uint32_t ksn = (k_rnd - ks0 < KSB) ? (k_rnd - ks0) : KSB;
+          // prior reads must retire before the staging is overwritten
+          asm volatile("s_waitcnt lgkmcnt(0)" ::: "memory");
+          for (uint32_t kb = 0; kb < ksn; kb++) {
+            for (uint32_t nt = 0; nt < NT; nt++) {
+              int g_nindx =
+                  (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
+              int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
+              __builtin_amdgcn_global_load_lds(
+                  (float4*)(&glbl[g_adr + M * N * (ks0 + kb)]),
+                  &(((float4*)lds.all)[(threadIdx.y * THRDS) + kb * THRDS * 4 +
+                                       nt * THRDS * 4 * KSB]),
+                  16, 0, 0);
+            }
+          }
+          asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
+          for (uint32_t kb = 0; kb < ksn; kb++) {
+            for (uint32_t nt = 0; nt < NT; nt++) {
+              vals[nt].f4 +=
+                  ((float4*)lds.all)[(threadIdx.x + threadIdx.y * THRDS) +
+                                     kb * THRDS * 4 + nt * THRDS * 4 * KSB];
+            }
           }
         }
       } else {
@@ -1771,7 +1824,7 @@ __global__ void wvSplitKrc_(const int actlN, const int K, const int Kap,
 torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
                          const std::optional<at::Tensor>& in_bias,
                          const int64_t CuCount) {
-  int _DTRMNSTC = 1;  // vllm::vllm_is_batch_invariant();
+  static constexpr bool _DTRMNSTC = true;
 
   auto M_in = in_b.size(0);
   auto N_in = in_a.size(0);
@@ -1819,26 +1872,157 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
 
   if (CuNeeded > CuCount) throw std::runtime_error("Invalid wvSplitKrc size");
 
-  // Can we increase SplitK by shrinking the K-shared to 256?
-  int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;
-
-  static torch::Tensor axl_glbl =
-      torch::zeros(
-          128 * 1024 * (_DTRMNSTC ? 12 : 1),
-          torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device()))
-          .detach();
-  static torch::Tensor axl_cntr =
-      torch::zeros(
-          128 * 1024 * (_DTRMNSTC ? 12 : 1) / 4,
-          torch::TensorOptions().dtype(torch::kInt).device(in_a.device()))
-          .detach();
-  auto glbl = axl_glbl.data_ptr<float>();
-  auto cntr = axl_cntr.data_ptr<int>();
+  // Can we increase SplitK by shrinking the K-shard to 256? That doubles both
+  // the CUs needed and the k-shards to read back, and the readback is serial in
+  // the shard count, so it pays only up to a crossover. The crossover is not a
+  // constant -- it falls as M grows -- so one cap cannot sit on it for every
+  // shape; this value minimizes total time over the admitted shapes. Retune it
+  // if the readback cost changes, and note the directions are not symmetric:
+  // capping too high costs about twice what capping too low does.
+  // Mirrored by rocm_unquantized_gemm_impl() in
+  // vllm/model_executor/layers/utils.py; both must pick the same chunkk or the
+  // host-side fit check bounds the wrong k_rnd.
+  constexpr int64_t CHUNKK2_MAX_SHARDS = 11;
+  const int64_t shards_chunkk2 = (K_in + 256 - 1) / 256;  // 256 == 512 / 2
+  const int chunkk = (N_p2 != 16 && CuNeeded * 2 <= CuCount &&
+                      shards_chunkk2 <= CHUNKK2_MAX_SHARDS)
+                         ? 2
+                         : 1;
+
+  // One fp32 partial per (M, N, k-shard) must fit the split-K workspace.
+  const int64_t k_rnd = (K_in + 512 / chunkk - 1) / (512 / chunkk);
+  constexpr int64_t glbl_cap = 128 * 1024 * (_DTRMNSTC ? 12 : 1);
+  const int64_t glbl_needed =
+      _DTRMNSTC ? (int64_t)N_p2 * M_in * k_rnd : (int64_t)N_p2 * M_in;
+  TORCH_CHECK(glbl_needed <= glbl_cap,
+              "wvSplitKrc: split-K workspace too small", " (need ", glbl_needed,
+              " > ", glbl_cap, " floats) for N=", N_in, " K=", K_in,
+              " M=", M_in);
+
+  // Split-K workspaces: one pre-partitioned pool per device.
+  // These must be zeroed before first use, and the kernel re-zeros them after
+  // each use.
+  // To avoid unnecessary overhead, we do not allocate them zeroed
+  // per-invocation. However, a simple static allocation does not work with
+  // multi-streams, because concurrent streams can alias into the same
+  // workspace. Solution: pre-allocate kWvSlots = 8 slots per device, plus
+  // kWvCaptureSlots reserved for capturing streams, and zero the whole pool
+  // once at creation. On first use, a stream takes one of the slots if any are
+  // still available. On subsequent uses, it re-uses the slot. If no slots are
+  // available, fall back to allocating a zeroed one per call, which is safe but
+  // adds overhead.
+  // A capturing stream must neither allocate nor enqueue anything -- either
+  // would be captured into the graph -- so it is only handed a reserved slot.
+  struct WvSplitKrcPool {
+    torch::Tensor glbl, cntr;
+    int64_t glbl_stride = 0, cntr_stride = 0;
+    int next_slot = 0;
+    int next_capture_slot = 0;
+    std::map<cudaStream_t, int> slot_of;
+    std::map<cudaStream_t, int> capture_slot_of;
+  };
+  // mutex not strictly necessary if caller is always single-threaded (Python)
+  // but this preserves Torch thread_local semantics for the C++ API
+  static std::mutex wv_pool_mu;
+  static std::map<int, WvSplitKrcPool> wv_pools;
+
+  float* glbl;
+  int* cntr;
+  // Fallback space
+  torch::Tensor spill_glbl, spill_cntr;
+  {
+    std::lock_guard<std::mutex> wv_lk(wv_pool_mu);
+    const int wv_dev = static_cast<int>(in_a.device().index());
+    WvSplitKrcPool& P = wv_pools[wv_dev];
+    constexpr int64_t wv_gf = 128 * 1024 * (_DTRMNSTC ? 12 : 1);
+    constexpr int64_t wv_ci = wv_gf / 4;
+    constexpr int64_t wv_all_slots = kWvSlots + kWvCaptureSlots;
+
+    // A stream whose capture was already invalidated reports an error; treat
+    // that as not capturing, since the capture is lost either way.
+    cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
+    unsigned long long capture_id = 0;
+    if (cudaStreamGetCaptureInfo(stream, &capture_status, &capture_id) !=
+        cudaSuccess) {
+      capture_status = cudaStreamCaptureStatusNone;
+    }
+    const bool wv_capturing = capture_status != cudaStreamCaptureStatusNone;
+
+    if (!P.glbl.defined()) {
+      TORCH_CHECK(!wv_capturing,
+                  "wvSplitKrc: the split-K workspace pool cannot be allocated "
+                  "during graph capture. Call "
+                  "warmup_rocm_skinny_gemm_workspaces() "
+                  "(vllm/model_executor/layers/utils.py) eagerly first.");
+      P.glbl_stride = wv_gf;
+      P.cntr_stride = wv_ci;
+      P.glbl = torch::zeros(wv_gf * wv_all_slots, torch::TensorOptions()
+                                                      .dtype(torch::kFloat32)
+                                                      .device(in_a.device()))
+                   .detach();
+      P.cntr =
+          torch::zeros(
+              wv_ci * wv_all_slots,
+              torch::TensorOptions().dtype(torch::kInt).device(in_a.device()))
+              .detach();
+      // The mutex orders host access to the map, not device execution, so it
+      // alone would let another stream be handed a slot the fill above is
+      // still clearing. Synchronizing here, inside the lock, is what makes
+      // every slot safe to use without any further per-slot zeroing.
+      const cudaError_t wv_err = cudaStreamSynchronize(stream);
+      TORCH_CHECK(wv_err == cudaSuccess,
+                  "wvSplitKrc: failed to zero the split-K workspace pool: ",
+                  cudaGetErrorString(wv_err));
+    }
+
+    int wv_slot = -1;
+    if (wv_capturing) {
+      // One workspace per capturing stream, so graphs captured on the same
+      // stream share it -- safe because at most one of them replays at a time
+      // (vLLM captures one graph per padded batch size, and runs one per step).
+      auto cap_it = P.capture_slot_of.find(stream);
+      if (cap_it == P.capture_slot_of.end()) {
+        TORCH_CHECK(P.next_capture_slot < kWvCaptureSlots,
+                    "wvSplitKrc: no free capture workspace slot; more than ",
+                    kWvCaptureSlots,
+                    " distinct streams have captured a graph containing "
+                    "wvSplitKrc. Raise kWvCaptureSlots in "
+                    "csrc/rocm/skinny_gemms.cu.");
+        cap_it = P.capture_slot_of
+                     .emplace(stream, static_cast<int>(kWvSlots) +
+                                          P.next_capture_slot++)
+                     .first;
+      }
+      wv_slot = cap_it->second;
+    } else {
+      auto slot_it = P.slot_of.find(stream);
+      if (slot_it != P.slot_of.end()) {
+        wv_slot = slot_it->second;
+      } else if (P.next_slot < kWvSlots) {
+        wv_slot = P.next_slot++;
+        P.slot_of.emplace(stream, wv_slot);
+      }
+    }
+
+    if (wv_slot >= 0) {
+      glbl = P.glbl.data_ptr<float>() + wv_slot * P.glbl_stride;
+      cntr = P.cntr.data_ptr<int>() + wv_slot * P.cntr_stride;
+    } else {
+      spill_glbl = torch::zeros(
+          wv_gf,
+          torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device()));
+      spill_cntr = torch::zeros(
+          wv_ci,
+          torch::TensorOptions().dtype(torch::kInt).device(in_a.device()));
+      glbl = spill_glbl.data_ptr<float>();
+      cntr = spill_cntr.data_ptr<int>();
+    }
+  }
 
 #define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK)                                     \
   {                                                                            \
     dim3 block(64, 4);                                                         \
-    if (_DTRMNSTC)                                                             \
+    if constexpr (_DTRMNSTC)                                                   \
       wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 1>          \
           <<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
                                        af4, bf4, biasf4, glbl, cntr, c,        \
diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu
index 8b7077d14ed0..a20ca41d413e 100644
--- a/csrc/rocm/skinny_gemms_int4.cu
+++ b/csrc/rocm/skinny_gemms_int4.cu
@@ -102,6 +102,18 @@ __device__ inline unsigned int min__(uint32_t a, uint32_t b) {
   return min(a, b);
 }
 
+// Reads row `row` of packed zero points [M/8, num_groups] uint32.
+//
+// The parameters must be signed and the indexing must use shift/mask: with
+// unsigned parameters and row/8, row%8, the (A_CHUNK=32, UNRL=8)
+// instantiation needs 256 VGPRs and 68 bytes of scratch instead of 239.
+__device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed,
+                                              const int row, const int group,
+                                              const int num_groups) {
+  const uint32_t w = zp_packed[(row >> 3) * num_groups + group];
+  return (w >> (4 * (row & 7))) & 0xFu;
+}
+
 // W4A16 skinny GEMM kernel: packed int4 weights, fp16/bf16 activations
 // Targets the "sml" case where activations fit in LDS.
 // A_CHUNK: number of K-elements processed per thread per step.
@@ -116,7 +128,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By,
                           const uint8_t* B_packed,
                           const scalar_t* __restrict__ A, const scalar_t* scale,
-                          const scalar_t* zero_points,
+                          const uint32_t* zero_points,
                           const scalar_t* __restrict__ BIAS, scalar_t* C,
                           const int _WvPrGrp, const int CuCount) {
   constexpr int max_lds_len = LDS_SIZE / 2;
@@ -257,7 +269,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
             if constexpr (!std::is_same_v<scalar_t, __hip_bfloat16>) {
               if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {
                 uint32_t group_idx = k_ / GROUP_SIZE;
-                scalar_t zp = zero_points[(m + y) * num_groups + group_idx];
+                scalar_t zp = __float2s<scalar_t>(static_cast<float>(
+                    zp_nibble(zero_points, m + y, group_idx, num_groups)));
   #pragma unroll
                 for (uint32_t b = 0; b < A_CHUNK; b++) {
                   cvtB.h[b] = cvtB.h[b] - zp;
@@ -280,8 +293,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
                 }
                 if constexpr (HAS_ZERO_POINTS) {
                   uint32_t group_idx_zp = k_ / GROUP_SIZE;
-                  float zp_f = __s2float(
-                      zero_points[(m + y) * num_groups + group_idx_zp]);
+                  float zp_f = static_cast<float>(
+                      zp_nibble(zero_points, m + y, group_idx_zp, num_groups));
                   partial -= (128.0f + zp_f) * act_sum;
                 } else {
                   partial -= 136.0f * act_sum;
@@ -335,7 +348,7 @@ __global__ void wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx,
                                       const int By, const uint8_t* B_packed,
                                       const scalar_t* __restrict__ A,
                                       const scalar_t* scale,
-                                      const scalar_t* zero_points,
+                                      const uint32_t* zero_points,
                                       const scalar_t* __restrict__ BIAS,
                                       scalar_t* C, const int _WvPrGrp,
                                       const int CuCount) {
@@ -352,7 +365,7 @@ template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
     wvSplitK_int4_hf_(const int K, const int M, const int Bx, const int By,
                       const uint8_t* B_packed, const scalar_t* __restrict__ A,
-                      const scalar_t* scale, const scalar_t* zero_points,
+                      const scalar_t* scale, const uint32_t* zero_points,
                       const scalar_t* __restrict__ BIAS, scalar_t* C,
                       const int _WvPrGrp, const int CuCount) {
   constexpr int max_lds_len = LDS_SIZE / 2;
@@ -508,7 +521,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
             if constexpr (!std::is_same_v<scalar_t, __hip_bfloat16>) {
               if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {
                 uint32_t group_idx = k_ / GROUP_SIZE;
-                scalar_t zp = zero_points[(m + y) * num_groups + group_idx];
+                scalar_t zp = __float2s<scalar_t>(static_cast<float>(
+                    zp_nibble(zero_points, m + y, group_idx, num_groups)));
   #pragma unroll
                 for (uint32_t b = 0; b < A_CHUNK; b++) {
                   cvtB.h[b] = cvtB.h[b] - zp;
@@ -531,8 +545,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
                 }
                 if constexpr (HAS_ZERO_POINTS) {
                   uint32_t group_idx_zp = k_ / GROUP_SIZE;
-                  float zp_f = __s2float(
-                      zero_points[(m + y) * num_groups + group_idx_zp]);
+                  float zp_f = static_cast<float>(
+                      zp_nibble(zero_points, m + y, group_idx_zp, num_groups));
                   partial -= (128.0f + zp_f) * act_sum;
                 } else {
                   partial -= 136.0f * act_sum;
@@ -596,7 +610,7 @@ __global__ void wvSplitK_int4_hf_(const int K, const int M, const int Bx,
                                   const int By, const uint8_t* B_packed,
                                   const scalar_t* __restrict__ A,
                                   const scalar_t* scale,
-                                  const scalar_t* zero_points,
+                                  const uint32_t* zero_points,
                                   const scalar_t* __restrict__ BIAS,
                                   scalar_t* C, const int _WvPrGrp,
                                   const int CuCount) {
@@ -621,7 +635,8 @@ static int mindiv_int4(int N, int div1, int div2) {
 // in_a: packed int4 weights [M, K/2] (int8) or [M, K/8] (int32)
 // in_b: activations [N, K] (fp16/bf16)
 // in_scale: group scales [M, K/group_size] (fp16/bf16)
-// in_zero_points: optional raw zero points [M, K/group_size] (fp16/bf16)
+// in_zero_points: optional raw zero points [M/8, K/group_size] (int32,
+//   8 uint4 per word along dim 0)
 //   If provided, kernel dequants as (nibble - zp_raw) * scale (asymmetric).
 //   If absent, kernel dequants as (nibble - 8) * scale (symmetric uint4b8).
 // group_size: 32, 64, or 128
@@ -664,14 +679,22 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
               "Scale must be [M, K/group_size] = [", M_in, ", ", num_groups,
               "] but got [", in_scale.size(0), ", ", in_scale.size(1), "]");
   if (in_zero_points.has_value()) {
-    TORCH_CHECK(in_zero_points->dtype() == in_b.dtype(),
-                "Zero points dtype must match activation dtype");
+    // The kernel reads the words as uint32, so either signedness is accepted.
+    TORCH_CHECK(in_zero_points->dtype() == at::kInt ||
+                    in_zero_points->dtype() == at::kUInt32,
+                "Zero points must be int32 or uint32 (packed 8x uint4 along "
+                "dim 0), got ",
+                in_zero_points->dtype());
     TORCH_CHECK(in_zero_points->dim() == 2,
-                "Zero points must be 2D [M, K/group_size], got shape ",
+                "Zero points must be 2D [M/8, K/group_size], got shape ",
                 in_zero_points->sizes());
-    TORCH_CHECK(in_zero_points->size(0) == M_in &&
+    TORCH_CHECK(M_in % 8 == 0,
+                "M must be divisible by 8 for packed zero points, got ", M_in);
+    TORCH_CHECK(in_zero_points->is_contiguous(),
+                "Zero points must be contiguous");
+    TORCH_CHECK(in_zero_points->size(0) == M_in / 8 &&
                     in_zero_points->size(1) == num_groups,
-                "Zero points must be [M, K/group_size] = [", M_in, ", ",
+                "Zero points must be [M/8, K/group_size] = [", M_in / 8, ", ",
                 num_groups, "] but got [", in_zero_points->size(0), ", ",
                 in_zero_points->size(1), "]");
   }
@@ -769,9 +792,9 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
         const fptype* aptr = reinterpret_cast<const fptype*>(in_b.data_ptr());
         const fptype* sptr =
             reinterpret_cast<const fptype*>(in_scale.data_ptr());
-        const fptype* zpptr =
+        const uint32_t* zpptr =
             in_zero_points.has_value()
-                ? reinterpret_cast<const fptype*>(in_zero_points->data_ptr())
+                ? reinterpret_cast<const uint32_t*>(in_zero_points->data_ptr())
                 : nullptr;
         const fptype* biasptr =
             (in_bias.has_value() && in_bias->numel() > 0)
diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp
index a8114ce1969a..14d1822707ea 100644
--- a/csrc/rocm/torch_bindings.cpp
+++ b/csrc/rocm/torch_bindings.cpp
@@ -31,7 +31,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
   rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
 
   // W4A16 grouped skinny GEMM: packed int4 weights, per-group scales,
-  // optional zero points for asymmetric quantization
+  // optional zero points [M/8, K/group_size] int32 for asymmetric
+  // quantization
   rocm_ops.def(
       "wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, "
       "Tensor? in_zero_points, Tensor? in_bias, int CuCount, "
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 9b344552c5ae..26b615688100 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -131,6 +131,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \
         git \
         curl \
         sudo \
+        elfutils-devel \
         rdma-core-devel \
     && dnf clean all \
     && rm -rf /var/cache/dnf
@@ -230,10 +231,12 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         echo "Installing Rubin prerelease requirements..." \
         && /opt/venv/bin/python3 use_existing_torch.py --prefix \
         && uv pip install --python /opt/venv/bin/python3 \
+            --prerelease=allow --index-strategy unsafe-first-match \
             -r /tmp/vllm-requirements/rubin-prerelease.txt \
         && uv pip install --python /opt/venv/bin/python3 \
             -r requirements/cuda.txt \
         && uv pip install --python /opt/venv/bin/python3 \
+            --prerelease=allow --index-strategy unsafe-first-match \
             -r /tmp/vllm-requirements/rubin-prerelease.txt; \
     elif [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing torch nightly..." \
@@ -474,6 +477,24 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 #################### EXTENSIONS BUILD IMAGE ####################
 # Build DeepEP - runs in PARALLEL with csrc-build
 # This stage is independent and doesn't affect csrc cache
+# Kernel symbol map for CI test selection: which source file produced each GPU
+# kernel symbol, read from the objects this stage just compiled. Only when
+# VLLM_KERNEL_SYMBOL_MAP=1 (nightly), and placed after the compile layers so
+# their cache is untouched. Never fails the build; an empty map is a signal.
+ARG VLLM_KERNEL_SYMBOL_MAP=""
+ARG buildkite_commit=""
+COPY tools/ci/kernel_symbol_map.py tools/ci/kernel_symbol_map.py
+RUN if [ "${VLLM_KERNEL_SYMBOL_MAP}" = "1" ]; then \
+        mkdir -p dist && python3 tools/ci/kernel_symbol_map.py \
+            --build-root build --source-root . --commit "${buildkite_commit}" \
+            --out dist/kernel_symbol_map.json.gz; \
+    else echo "Skipping kernel symbol map (VLLM_KERNEL_SYMBOL_MAP != 1)"; fi
+
+# Exported on its own by the CI bake target of the same name (ci-infra
+# docker/ci.hcl) with output=type=local, so the map becomes a build artifact.
+FROM scratch AS kernel-symbol-map
+COPY --from=csrc-build /workspace/dist/kernel_symbol_map.json.gz /kernel_symbol_map.json.gz
+
 FROM base AS extensions-build
 ARG TARGETARCH
 ARG CUDA_VERSION
@@ -640,7 +661,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         python3 use_existing_torch.py --prefix; \
     fi \
     && if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
-        sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\])[[:space:]]*==/Id' \
+        sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\]|quack-kernels)[[:space:]]*==/Id' \
             requirements/cuda.txt; \
     fi
 
@@ -732,6 +753,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 FROM ${FINAL_BASE_IMAGE} AS vllm-runtime-base
 
 ARG INSTALL_RUBIN_PRERELEASE
+ARG TARGETPLATFORM
+ARG max_jobs=2
 ARG CUDA_VERSION
 ARG PYTHON_VERSION
 ARG NCCL_VERSION
@@ -851,6 +874,20 @@ RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
         && ldconfig; \
     fi
 
+# Install the initialized-snapshot runtime on the architecture supported by
+# `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned
+# source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's
+# Apache-2.0 source against the public CUDA Driver checkpoint APIs instead of
+# redistributing NVIDIA's proprietary cuda-checkpoint binary.
+COPY tools/install_snapshot_runtime.sh /usr/local/bin/install_snapshot_runtime.sh
+RUN if [ "${TARGETPLATFORM}" = "linux/amd64" ]; then \
+        MAX_JOBS="${max_jobs}" /usr/local/bin/install_snapshot_runtime.sh; \
+    else \
+        echo "Initialized snapshots are not installed for ${TARGETPLATFORM}"; \
+    fi
+
+ENV CRIU_CUDA_PLUGIN_DIR=/usr/local/lib/criu
+
 # Install uv for faster pip installs
 RUN python3 -m pip install uv
 
@@ -940,11 +977,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     fi && \
     if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
         uv pip install --system \
+            --prerelease=allow --index-strategy unsafe-first-match \
             -r /tmp/vllm-requirements/rubin-prerelease.txt \
         && sed -Ei '/^[[:space:]]*(torch|torchvision|torchaudio)[[:space:]]*==/Id' \
             /tmp/requirements-cuda.txt \
         && uv pip install --system -r /tmp/requirements-cuda.txt \
         && uv pip install --system \
+            --prerelease=allow --index-strategy unsafe-first-match \
             -r /tmp/vllm-requirements/rubin-prerelease.txt \
         && printf '%s\n' \
             "/usr/local/lib/python${PYTHON_VERSION}/dist-packages/nvpl/lib" \
@@ -960,7 +999,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.18.post1
+ARG FLASHINFER_VERSION=0.7.0
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "${INSTALL_RUBIN_PRERELEASE}" != "true" ]; then \
         uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
@@ -1042,6 +1081,10 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
     fi
 
+# Expose the Rust frontend bundled in the wheel on the container's PATH.
+RUN ln -s "$(python3 -c 'import sysconfig; print(sysconfig.get_path("platlib"))')/vllm/vllm-rs" /usr/local/bin/vllm-rs \
+    && vllm-rs --version
+
 RUN --mount=type=cache,target=/opt/uv/cache \
 . /etc/environment && \
 uv pip list
@@ -1221,7 +1264,8 @@ ENV UV_HTTP_TIMEOUT=500
 ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 RUN --mount=type=cache,target=/opt/uv/cache \
-    --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
+    --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \
+    set -eu; \
     CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
     CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-'); \
     CUDA_HOME=/usr/local/cuda; \
@@ -1230,11 +1274,19 @@ RUN --mount=type=cache,target=/opt/uv/cache \
                 libcublas-dev-${CUDA_VERSION_DASH} \
                 libcusolver-dev-${CUDA_VERSION_DASH}"; \
     if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
-        uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \
+        KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors.txt; \
+        if [ "${CUDA_MAJOR}" = "12" ]; then \
+            KV_REQUIREMENTS=/tmp/vllm-requirements/kv_connectors_cu12.txt; \
+        fi; \
+        # Preserve the stack used to compile vLLM when resolving optional dependencies. \
+        uv pip freeze --system > /tmp/kv-installed.txt; \
+        grep -E '^(torch|torchaudio|torchvision|triton|cuda-bindings|cuda-python|cuda-toolkit)==' \
+            /tmp/kv-installed.txt > /tmp/kv-constraints.txt; \
+        uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build || ( \
             # if the above fails, install from source
             apt-get update -y && \
             apt-get install -y --no-install-recommends --allow-change-held-packages ${BUILD_PKGS} && \
-            uv pip install --system -r /tmp/kv_connectors.txt --no-build-isolation && \
+            uv pip install --system -c /tmp/kv-constraints.txt -r "${KV_REQUIREMENTS}" --no-build-isolation && \
             apt-get purge -y ${BUILD_PKGS} && \
             # clean up -dev packages, keep runtime libraries
             rm -rf /var/lib/apt/lists/* \
@@ -1245,8 +1297,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         MOONCAKE_VERSION=$(python3 -c "import importlib.metadata as m; print(m.version('mooncake-transfer-engine'))" 2>/dev/null || true) && \
         if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then \
             uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true; \
-            uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
+            uv pip install --system -c /tmp/kv-constraints.txt "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"; \
         fi; \
+        rm -f /tmp/kv-installed.txt /tmp/kv-constraints.txt; \
     fi
 
 # Mooncake registers GPU memory for RDMA via nvidia-peermem by default. Hosts
@@ -1261,6 +1314,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
         echo "Installing Rubin prerelease requirements..." \
         && uv pip install --system \
+            --prerelease=allow --index-strategy unsafe-first-match \
             -r /tmp/vllm-requirements/rubin-prerelease.txt; \
     fi
 
diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu
index f81ae880966c..05d337fa1361 100644
--- a/docker/Dockerfile.cpu
+++ b/docker/Dockerfile.cpu
@@ -107,55 +107,6 @@ ENV UV_LINK_MODE="copy"
 
 ENV TARGETARCH=${TARGETARCH}
 
-######################### TRITON-CPU BUILD IMAGE #########################
-# Branches off base-common directly (before requirements/cpu.txt is installed
-# below) so a torch/etc. version bump doesn't invalidate this stage's cache --
-# it builds triton-cpu in an isolated `uv build` environment and never touches
-# these packages.
-FROM base-common AS vllm-triton-cpu-build
-
-# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
-# Re-declared here because this stage is `FROM base-common` (not `vllm-build`),
-# so it does not inherit the ARG/ENV defined there. Without it, the guard
-# below would see an empty value and build triton-cpu on non-x86 targets
-# (e.g. arm64).
-ARG VLLM_CPU_X86=0
-
-ARG USE_SCCACHE
-ARG SCCACHE_ENDPOINT
-ARG SCCACHE_LOCAL_ONLY=0
-
-WORKDIR /vllm-workspace
-
-RUN mkdir dist
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=cache,target=/root/.cache/ccache \
-    --mount=type=cache,target=/root/.cache/sccache-local,sharing=locked \
-    --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
-    --mount=type=cache,target=/root/.triton \
-    --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
-    set -ex; \
-    if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
-        git clone --recurse-submodules --shallow-submodules --filter=blob:none "https://github.com/triton-lang/triton-cpu.git"; \
-        cd triton-cpu; \
-        git checkout "270e696d"; \
-        if [ "$USE_SCCACHE" = "1" ]; then \
-            if [ "$SCCACHE_LOCAL_ONLY" = "1" ]; then \
-                unset SCCACHE_BUCKET SCCACHE_REGION; \
-                export SCCACHE_DIR=/root/.cache/sccache-local SCCACHE_CACHE_SIZE=20G; \
-            elif [ -n "${SCCACHE_ENDPOINT}" ]; then \
-                export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; \
-            fi; \
-            export TRITON_APPEND_CMAKE_ARGS="-DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache"; \
-            sccache --show-stats; \
-        fi; \
-        uv build --wheel --out-dir=../dist; \
-        if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi; \
-        cd ..; \
-        rm -rf triton-cpu; \
-    fi
-
 ######################### COMMON BASE IMAGE + PYTHON DEPS #########################
 FROM base-common AS base-common-pydeps
 
@@ -379,9 +330,12 @@ WORKDIR /vllm-workspace
 COPY requirements/test/cpu.txt requirements/test/cpu.txt
 
 # cpu.txt is compiled for x86_64, so platform markers are resolved away. Drop
-# packages unavailable on aarch64 (decord, terratorch) for arm builds.
+# packages unavailable on aarch64 (decord, terratorch) for arm builds. Also
+# drop the generic PyPI triton pulled in transitively (via xgrammar) on x86
+# builds, so it doesn't clobber the triton-cpu wheel from requirements/cpu.txt.
 RUN case "$(uname -m)" in \
     aarch64|arm64) sed -i '/^decord==/d; /^terratorch==/d' requirements/test/cpu.txt ;; \
+    x86_64) sed -i '/^triton==/d' requirements/test/cpu.txt ;; \
     esac
 
 # pytrec-eval-terrier bundles pre-C23 trec_eval sources. GCC 15 defaults to C23;
@@ -434,14 +388,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     fi && \
     pre-commit install --hook-type pre-commit --hook-type commit-msg
 
-# Installed after requirements/test/cpu.txt (which pins triton==3.6.0) so the
-# triton-cpu wheel isn't overwritten by that pin.
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
-    if ls dist/*.whl >/dev/null 2>&1; then \
-        uv pip install dist/*.whl; \
-    fi
-
 ENTRYPOINT ["bash"]
 
 ######################### TEST IMAGE #########################
@@ -453,18 +399,16 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
     uv pip install dist/*.whl
 
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
-    if ls dist/*.whl >/dev/null 2>&1; then \
-        uv pip install dist/*.whl; \
-    fi
-
 ADD ./tests/ ./tests/
 ADD ./examples/ ./examples/
 ADD ./benchmarks/ ./benchmarks/
 ADD ./vllm/collect_env.py .
 ADD ./docker/ ./docker/
 ADD ./.buildkite/ ./.buildkite/
+# tools/ carries repo-tooling code that a few tests exercise directly
+# (e.g. tests/tools/test_check_test_tethering.py imports
+# tools.pre_commit.check_test_tethering).
+ADD ./tools/ ./tools/
 
 # install development dependencies (for testing)
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -484,13 +428,7 @@ WORKDIR /vllm-workspace
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
-    uv pip install "$(realpath dist/*.whl)[audio]"
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
-    if ls dist/*.whl >/dev/null 2>&1; then \
-        uv pip install dist/*.whl; \
-    fi
+    uv pip install "$(realpath dist/*.whl)[audio,bench]"
 
 # Add labels to document build configuration
 LABEL org.opencontainers.image.title="vLLM CPU"
diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock
index c3ce785c679d..16e01ad723b9 100644
--- a/docker/Dockerfile.rock
+++ b/docker/Dockerfile.rock
@@ -5,7 +5,7 @@ ARG VLLM_BRANCH="main"
 ARG COMMON_WORKDIR=/app
 ARG BASE_IMAGE=rocm/vllm-dev:base
 ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
-ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7
+ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557
 ARG INSTALL_LMCACHE=false
 # NIC backend for MoRI RDMA support.
 # By default (all), drivers and userspace libraries for all supported NIC types
@@ -17,8 +17,9 @@ ARG INSTALL_LMCACHE=false
 ARG NIC_BACKEND=all
 # AMD AINIC apt repo settings
 # Users can specify a custom version compatible with their host drivers.
-# The default version has been tested with ioinic-dkms=25.11.1.001
-ARG AINIC_VERSION=1.117.3-hydra
+# The default version has been tested with ionic-dkms=26.03.3.001
+# (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC.
+ARG AINIC_VERSION=1.117.5-a-77
 ARG UBUNTU_CODENAME=jammy
 
 # Sccache configuration. Release builds use this today; CI can opt in when a
@@ -322,8 +323,9 @@ ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build"
 
 RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \
     --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
-    grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
+    { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
         cmake/external_projects/triton_kernels.cmake \
+        || echo "WARNING: Dockerfile.rock ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \
     && export CCACHE_BASEDIR="$PWD" \
     && echo "=== ccache stats before ROCm native build ===" \
     && (ccache --show-stats || true) \
@@ -669,8 +671,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
 
 # Install the packaging helpers before modifying requirements/rocm.txt.
 RUN --mount=type=cache,target=/root/.cache/uv \
-    grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
+    { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
         vllm/cmake/external_projects/triton_kernels.cmake \
+        || echo "WARNING: Dockerfile.rock ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \
     && cd vllm \
     && uv pip install --system setuptools_scm regex
 
@@ -875,6 +878,15 @@ RUN mkdir src && mv vllm src/vllm
 # Catch GPU<->CPU syncs in execute_model/sample_tokens during tests.
 ENV VLLM_GPU_SYNC_CHECK=error
 
+FROM test AS test_smoke
+ARG ROCM_SMOKE_ID
+RUN --network=none \
+    bash /vllm-workspace/.buildkite/scripts/rocm/smoke-test-image.sh --inside \
+    && printf '%s\n' "${ROCM_SMOKE_ID:-local}" > /vllm-smoke-ok
+
+FROM scratch AS export_test_smoke
+COPY --from=test_smoke /vllm-smoke-ok /
+
 # -----------------------
 # Final vLLM image
 FROM mori_base AS final_common
@@ -902,7 +914,7 @@ RUN case "$(which python3)" in \
         *) ;; esac
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --upgrade huggingface-hub[cli]
+    uv pip install --system --upgrade huggingface-hub
 
 # Install vLLM using uv (inherited from base stage)
 # Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base
index 1b1d08988123..79896933e7eb 100644
--- a/docker/Dockerfile.rock_base
+++ b/docker/Dockerfile.rock_base
@@ -3,10 +3,12 @@ ARG PYTHON_VERSION=3.12
 ARG TORCH_VERSION=2.12.0+rocm10.0.0
 ARG TORCHVISION_VERSION=0.27.0+rocm10.0.0
 ARG TORCHAUDIO_VERSION=2.11.0+rocm10.0.0
-ARG TRITON_VERSION=3.8.0+git4cff872c.rocm10.0.0
+ARG TRITON_BRANCH="c0a142ffcbb42f9a67a56bf791cad472be1d6425" # release/internal/3.8.x as of 9/21/2026
+ARG TRITON_VERSION=3.8.0+gitc0a142ff
+ARG TRITON_REPO="https://github.com/ROCm/triton.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="v0.1.21.post1"
+ARG AITER_BRANCH="v0.1.22.post1"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="f95cd5e0ca7b730f92f9ffdfb358484d6965d9f7" # TheRock compatible v1.1.0
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -90,19 +92,15 @@ RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
     && rocm-sdk init
 
 # Install PyTorch wheels (nightly for latest torch, release as fallback)
-# Triton is pinned here rather than left to torch's Requires-Dist, so a wheel bump that
-# moves the triton torch wants fails resolution instead of quietly changing compilers.
 ARG TORCH_VERSION
 ARG TORCHAUDIO_VERSION
 ARG TORCHVISION_VERSION
-ARG TRITON_VERSION
 RUN pip install \
     --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \
     --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \
     "torch==${TORCH_VERSION}" \
     "torchaudio==${TORCHAUDIO_VERSION}" \
-    "torchvision==${TORCHVISION_VERSION}" \
-    "triton==${TRITON_VERSION}"
+    "torchvision==${TORCHVISION_VERSION}"
 
 # Install the AMD-compiled device kernels corresponding to PYTORCH_ROCM_ARCH:
 
@@ -146,6 +144,15 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \
             ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \
         fi; \
     done \
+    && for f in librocprofiler-sdk.so.1 librocprofiler-register.so.0 \
+              librocprofiler-sdk-roctx.so.1; do \
+        if [ -e "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}" ] \
+           && [ -e "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" ]; then \
+            rm -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \
+            ln -sfn "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" \
+                "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \
+        fi; \
+    done \
     && ldconfig
 
 # Install sccache if USE_SCCACHE is enabled (for release builds)
@@ -248,6 +255,20 @@ RUN pip install pyyaml && cd aiter \
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
 
 
+###
+### Triton Build
+###
+FROM base AS build_triton
+ARG TRITON_BRANCH
+ARG TRITON_REPO
+RUN git clone ${TRITON_REPO}
+RUN cd triton \
+    && git checkout ${TRITON_BRANCH} \
+    && if [ ! -f setup.py ]; then cd python; fi \
+    && python3 setup.py bdist_wheel --dist-dir=dist \
+    && mkdir -p /app/install && cp dist/*.whl /app/install
+
+
 ###
 ### ROCr + CLR build
 ###
@@ -311,6 +332,8 @@ RUN mkdir -p /staging/lib \
 # only includes dependencies used by wheel release pipeline
 FROM base AS debs_wheel_release
 RUN mkdir /app/debs
+RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
+    cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
@@ -318,6 +341,8 @@ RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
 
 FROM base AS debs
 RUN mkdir /app/debs
+RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
+    cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
@@ -326,6 +351,7 @@ RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
 
 FROM base AS final
+ARG TRITON_VERSION
 
 # torch 2.12 implicitly attaches rocprofiler-sdk; its queue interposition hangs at HSA
 # teardown, so disable it. TODO: drop once ROCm's stock rocprofiler carries the fix.
@@ -350,7 +376,15 @@ RUN set -eux; \
 
 RUN ldconfig
 
-ARG TRITON_VERSION
+# This is to prevent resolvers from checking METADATA and RECORD and clobbering
+# the Triton source build we installed. Constraints files are unable to help
+# with this. However, these entries are optional and can be blanked which
+# solves the problem.
+RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \
+    && grep -q '^Requires-Dist: triton' "$TORCH_DIST_INFO/METADATA" \
+    && sed -i '/^Requires-Dist: triton/d' "$TORCH_DIST_INFO/METADATA" \
+    && sed -i "s|^\(${TORCH_DIST_INFO##*/}/METADATA\),.*|\1,,|" "$TORCH_DIST_INFO/RECORD"
+
 RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
     pip install /install/*.whl \
     && ! pip check 2>&1 | grep -iE '^(torch|triton) ' \
@@ -374,6 +408,8 @@ ARG TORCHAUDIO_VERSION
 ARG TORCHVISION_VERSION
 ARG PYTORCH_ROCM_ARCH
 ARG AITER_ROCM_ARCH
+ARG TRITON_BRANCH
+ARG TRITON_REPO
 ARG FA_BRANCH
 ARG FA_REPO
 ARG AITER_BRANCH
@@ -382,6 +418,7 @@ ARG MORI_BRANCH
 ARG MORI_REPO
 ARG ROCM_SYSTEMS_REPO
 ARG ROCM_RUNTIME_COMMIT
+ARG TRITON_VERSION
 # Record hashes for Triton and pybind11 so we can verify later.
 RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
     && echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \
@@ -391,6 +428,8 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
     && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \
     && echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}" >> /app/versions.txt \
     && echo "AITER_ROCM_ARCH: ${AITER_ROCM_ARCH}" >> /app/versions.txt \
+    && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
+    && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
     && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
     && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
     && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
index 8d5f4f8f408a..cfd0b4557cd0 100644
--- a/docker/Dockerfile.rocm
+++ b/docker/Dockerfile.rocm
@@ -19,7 +19,7 @@ ARG NIC_BACKEND=all
 # Users can specify a custom version compatible with their host drivers.
 # The default version has been tested with ionic-dkms=26.03.3.001
 # (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC.
-ARG AINIC_VERSION=1.117.5
+ARG AINIC_VERSION=1.117.5-a-77
 ARG UBUNTU_CODENAME=jammy
 
 # Sccache configuration. Release builds use this today; CI can opt in when a
@@ -367,7 +367,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/scripts/benchmark_helion_kernels.py /scripts/benchmark_helion_kernels.py
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools /tools
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
@@ -750,7 +750,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmar
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/scripts/benchmark_helion_kernels.py /scripts/benchmark_helion_kernels.py
-COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools /tools
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
 COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
@@ -938,7 +938,7 @@ RUN case "$(which python3)" in \
         *) ;; esac
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --upgrade huggingface-hub[cli]
+    uv pip install --system --upgrade huggingface-hub
 
 # Install vLLM using uv (inherited from base stage)
 # Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base
index d842a7b07e68..eb4bb5cc9750 100644
--- a/docker/Dockerfile.rocm_base
+++ b/docker/Dockerfile.rocm_base
@@ -9,13 +9,13 @@ ARG PYTORCH_AUDIO_BRANCH="v2.11.0"
 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="v0.1.21.post2"
+ARG AITER_BRANCH="v0.1.22.post1"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="v1.1.0"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
 ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git"
 ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350"
-ARG ROCM_RUNTIME_COMMIT="b539bf7eebfd99ad0a69668caa1f4037034d501f"
+ARG ROCM_RUNTIME_COMMIT="af70edc91f1dc5d228fd4c076df99e65e49c4521"
 
 # Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE
diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250
index 95dad539136d..318bf702720b 100644
--- a/docker/Dockerfile.rocm_gfx1250
+++ b/docker/Dockerfile.rocm_gfx1250
@@ -15,7 +15,6 @@ ARG NIC_BACKEND=all
 # Users can specify a custom version compatible with their host drivers.
 # The default version has been tested with ionic-dkms=26.03.3.001
 # (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC.
-# TODO (jmathias) : Check if we can revert this Pinned AINIC version
 ARG AINIC_VERSION=1.117.5-a-77
 ARG UBUNTU_CODENAME=jammy
 
@@ -672,7 +671,7 @@ RUN case "$(which python3)" in \
         *) ;; esac
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --upgrade huggingface-hub[cli]
+    uv pip install --system --upgrade huggingface-hub
 
 # Install vLLM using uv (inherited from base stage)
 # Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x
index 9517ff371850..6faba540a92b 100644
--- a/docker/Dockerfile.s390x
+++ b/docker/Dockerfile.s390x
@@ -169,8 +169,6 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/
 ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
 ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
-# Force pure Python protobuf to avoid s390x C++ extension crashes
-ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
 COPY . /workspace/vllm
 WORKDIR /workspace/vllm
 
@@ -214,10 +212,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
     uv pip install "$(echo dist/*.whl)[tensorizer]"
 
-# Remove protobuf C++ extension that crashes on s390x
-RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \
-           /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true
-
 # setup non-root user for vllm
 RUN umask 002 && \
     /usr/sbin/useradd --uid 2000 --gid 0 vllm && \
diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu
index ee4c2a0d3b5a..3b0df0082643 100644
--- a/docker/Dockerfile.xpu
+++ b/docker/Dockerfile.xpu
@@ -87,11 +87,14 @@ RUN mkdir neo && \
 
 # Install XPU-Manager payloads without registering conflicting libze package
 # dependencies in dpkg.
+# Note that ocl-icd-opencl-dev is required by DLE upgrade but xpu-smi, should
+# can be removed in future.
 RUN apt-get update -y && \
     apt-get install -y --no-install-recommends software-properties-common && \
     add-apt-repository -y ppa:kobuk-team/intel-graphics && \
     apt-get update -y && \
     apt-get install -y --no-install-recommends libhwloc15 libpciaccess0 libigsc1 intel-metrics-library && \
+    apt-get install -y ocl-icd-opencl-dev && \
     mkdir xpu-m && \
     cd xpu-m && \
     wget https://github.com/intel/xpumanager/releases/download/v2.1.0/xpu-smi_2.1.0+26.33.6468cec-1.24.04_amd64.deb && \
@@ -264,4 +267,19 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -e tests/vllm_test_utils
 
+# Replace 2026.1.2 DLE packages to 2026.1.3
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=.git,target=.git \
+    uv pip install \
+      "intel-cmplr-lib-rt==2026.1.1" \
+      "intel-cmplr-lib-ur==2026.1.1" \
+      "intel-cmplr-lic-rt==2026.1.1" \
+      "intel-sycl-rt==2026.1.1" \
+      "oneccl-devel==2022.1.2; platform_system == 'Linux' and platform_machine == 'x86_64'" \
+      "oneccl==2022.1.2; platform_system == 'Linux' and platform_machine == 'x86_64'" \
+      "dpcpp-cpp-rt==2026.1.1" \
+      "intel-opencl-rt==2026.1.1" \
+      "intel-openmp==2026.1.1" \
+      "intel-pti==1.1.0"
+
 ENTRYPOINT ["vllm", "serve"]
diff --git a/docker/docker-bake-rocm.hcl b/docker/docker-bake-rocm.hcl
index 9d73a6505eab..ab523a296a93 100644
--- a/docker/docker-bake-rocm.hcl
+++ b/docker/docker-bake-rocm.hcl
@@ -33,6 +33,10 @@ variable "CI_BASE_CONTENT_HASH" {
   default = ""
 }
 
+variable "CI_BASE_DOCKERFILE" {
+  default = "docker/Dockerfile.rocm"
+}
+
 # REMOTE_VLLM=0: use local source via Docker build context (ONBUILD COPY ./ vllm/)
 # REMOTE_VLLM=1: clone from GitHub at VLLM_BRANCH (standalone builds without local source)
 variable "REMOTE_VLLM" {
@@ -74,7 +78,7 @@ group "default" {
 }
 
 target "_common-rocm" {
-  dockerfile = "docker/Dockerfile.rocm"
+  dockerfile = CI_BASE_DOCKERFILE
   context    = "."
   args = {
     max_jobs                        = MAX_JOBS
diff --git a/docker/versions.json b/docker/versions.json
index 7a5b9f44abd7..ec23b95186ac 100644
--- a/docker/versions.json
+++ b/docker/versions.json
@@ -71,7 +71,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.18.post1"
+      "default": "0.7.0"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"
diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png
index 9eca3ae91357..8ab91038016e 100644
Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ
diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md
index d46ab04b0e2a..e7ee794da6f1 100644
--- a/docs/benchmarking/cli.md
+++ b/docs/benchmarking/cli.md
@@ -147,7 +147,7 @@ With standard decoding, each streamed output usually contains one token, so ITL
 and TPOT are typically similar.
 
 With speculative decoding, one streamed output can contain multiple tokens,
-such as several accepted draft tokens within a single engine tstep. ITL records
+such as several accepted draft tokens within a single engine step. ITL records
 only the gaps between streamed outputs; it does not add zero-duration gaps for
 tokens in the same output. TPOT instead amortizes the request's decoding time
 over every output token.
@@ -451,7 +451,7 @@ vllm bench serve \
     --num-prompts -1
 ```
 
-Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
+Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructured data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
 
 #### BFCL (Tool-Calling) Benchmark
 
diff --git a/docs/benchmarking/dashboard.md b/docs/benchmarking/dashboard.md
index 44effc078e35..4c1132c17fd1 100644
--- a/docs/benchmarking/dashboard.md
+++ b/docs/benchmarking/dashboard.md
@@ -36,7 +36,7 @@ When run, benchmark script generates results under **benchmark/results** folder,
 - `ON_CPU`: set the value to '1' on Intel® Xeon® and Arm® Neoverse™ Processors. Default value is 0.
 - `SERVING_JSON`: JSON file to use for the serving tests. Default value is empty string (use default file).
 - `LATENCY_JSON`: JSON file to use for the latency tests. Default value is empty string (use default file).
-- `THROUGHPUT_JSON`: JSON file to use for the throughout tests. Default value is empty string (use default file).
+- `THROUGHPUT_JSON`: JSON file to use for the throughput tests. Default value is empty string (use default file).
 - `REMOTE_HOST`: IP for the remote vLLM service to benchmark. Default value is empty string.
 - `REMOTE_PORT`: Port for the remote vLLM service to benchmark. Default value is empty string.
 - `PROMPTS_PER_CONCURRENCY`: Multiplier to compute `num_prompts` for serving tests (`num_prompts = max_concurrency × value`). Overrides JSON `num_prompts`. Default is NULL.
diff --git a/docs/contributing/README.md b/docs/contributing/README.md
index 55966f2854fd..e4b86f99aed5 100644
--- a/docs/contributing/README.md
+++ b/docs/contributing/README.md
@@ -50,7 +50,7 @@ To rebuild only the Rust frontend binary:
 ./build_rust.sh --debug  # faster build for development
 ```
 
-If you are developing vLLM's Python and CUDA/C++ code, install Pytorch first:
+If you are developing vLLM's Python and CUDA/C++ code, install PyTorch first:
 
 ```bash
 uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129
@@ -315,6 +315,34 @@ review process:
   the `ready` label, the PR author can use `/ci run`, `/ci retry`, `/ci cancel`,
   or the corresponding `/amd-ci` variants. New commits do not start upstream
   CI automatically.
+- Before creating a Buildkite build, `/ci run` and `/amd-ci run`, including
+  `all` and `nightly`, check the requested PR commit against its target branch
+  in the upstream repository. Your branch must contain every commit currently
+  on that branch: **zero commits behind**. This also applies to PRs targeting
+  release branches or another PR's branch.
+- If the PR is behind, the bot reports the count without creating a new build.
+  Merge or rebase onto the latest target branch, then rerun the command. If the
+  target branch advances in the meantime, update your branch and try again.
+- To run CI on an outdated branch at your own risk, append `--allow-stale` to
+  any run command, for example `/ci run --allow-stale` or
+  `/amd-ci run all --allow-stale`. The same authorization requirements apply.
+  The bot reports the lag and warns that outdated CI configuration may cause
+  failures. Before merging, merge or rebase onto the latest target branch and
+  rerun CI without `--allow-stale` on the latest PR commit.
+- Retrying or cancelling builds does not check branch freshness.
+- These commands do not modify your branch or enforce merge requirements.
+  Direct Buildkite launches are outside this workflow. No additional token is
+  needed.
+
+Maintainers must configure required Buildkite status checks on the latest PR
+commit and enable GitHub's
+[Require branches to be up to date before merging](https://docs.github.com/en/repositories/configuring-branches-and-merges-in-your-repository/managing-protected-branches/about-protected-branches#require-status-checks-before-merging)
+protection on the relevant target branches to enforce the merge policy. Select
+statuses that confirm the required tests completed, using Buildkite as their
+expected source. An aggregate status that reports "passed and blocked" does not
+by itself guarantee those tests ran; verify the pipeline's required and optional
+blocked steps before choosing statuses or changing blocked-build reporting.
+The CI command workflow does not configure these repository or pipeline settings.
 
 ### Pull Request Limits and Escalation
 
diff --git a/docs/contributing/model/multimodal.md b/docs/contributing/model/multimodal.md
index 951d25587ad9..d374a5ecbc27 100644
--- a/docs/contributing/model/multimodal.md
+++ b/docs/contributing/model/multimodal.md
@@ -454,17 +454,17 @@ return a schema of the tensors outputted by the HF processor that are related to
     ??? code
 
         ```python
-        def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+        def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
             # Mistral3Processor requires text corresponding to the images
             return self.dummy_inputs.get_dummy_text(mm_counts)
 
         def _postprocess_hf_mm_data(
             self,
-            mm_data: Mapping[str, object],
-            hf_processor_mm_kwargs: Mapping[str, object],
+            hf_data: Mapping[str, object],
+            hf_kwargs: Mapping[str, object],
             processed_data: BatchFeature,
         ) -> BatchFeature:
-            if not mm_data:
+            if not hf_data:
                 return processed_data
 
             pixel_values = processed_data.get("pixel_values")
@@ -486,25 +486,36 @@ return a schema of the tensors outputted by the HF processor that are related to
     calls the HF processor on the multi-modal data without passing any text.
     If the HF processor instead requires text corresponding to the multi-modal items,
     you should override
-    [_get_hf_processor_text][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_processor_text]
+    [_get_hf_mm_text][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_text]
     to return the dummy text from
     [BaseDummyInputsBuilder.get_dummy_text][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_text]
     like in the example above. If the HF processor expects the multi-modal data
-    under different keys than those provided by the multi-modal items
-    (e.g. `audio` instead of `audios`), or requires additional keyword arguments
+    under different keys than those provided by the multi-modal items,
+    or requires additional keyword arguments
     (e.g. `sampling_rate`), you should override
-    [_preprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._preprocess_hf_mm_data].
+    [_get_hf_mm_inputs][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs].
     If you need to modify the output of the HF processor,
     you should override
     [_postprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._postprocess_hf_mm_data].
+    Passthrough fields have already been merged into `processed_data` when this
+    hook is called.
+    If you need to change how the normalized inputs are passed to the HF
+    processor (for example, to call a method other than `__call__`, or to use
+    different kwargs to construct and invoke it), override
+    [_call_hf_processor][vllm.multimodal.processing.BaseMultiModalProcessor._call_hf_processor].
     If the HF processor applies additional transformations to the prompt tokens
     that are not reflected in its multi-modal outputs (e.g. prepending a BOS token),
     you should override
     [_postprocess_prompt][vllm.multimodal.processing.BaseMultiModalProcessor._postprocess_prompt]
     to replicate them.
-    For even more control over how the HF processor is called, you can override
+    Only processors that require multiple HF calls or manually construct the
+    complete output should override
     [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main]
-    directly.
+    directly. Such implementations must obtain their normalized inputs through
+    [_get_hf_mm_inputs][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs]
+    and return through
+    [_finalize_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._finalize_hf_mm_data]
+    so passthrough fields and model-specific postprocessing are applied consistently.
 
     Since `pixel_values` is now a list with one tensor per image, we can override
     [_get_mm_fields_config][vllm.multimodal.processing.BaseMultiModalProcessor._get_mm_fields_config] as follows:
@@ -724,7 +735,7 @@ Examples:
 
 ### Custom HF processor
 
-Some models don't define an HF processor class on HF Hub. In that case, you can define a custom HF processor that has the same call signature as HF processors and return it from [BaseProcessingInfo.get_hf_processor][vllm.multimodal.processing.BaseProcessingInfo.get_hf_processor]. It is then applied to the multi-modal data via [InputProcessingContext.call_hf_processor][vllm.multimodal.processing.InputProcessingContext.call_hf_processor] inside [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main].
+Some models don't define an HF processor class on HF Hub. In that case, you can define a custom HF processor that has the same call signature as HF processors and return it from [BaseProcessingInfo.get_hf_processor][vllm.multimodal.processing.BaseProcessingInfo.get_hf_processor]. The default [_call_hf_processor][vllm.multimodal.processing.BaseMultiModalProcessor._call_hf_processor] applies it to the normalized multi-modal data via [InputProcessingContext.call_hf_processor][vllm.multimodal.processing.InputProcessingContext.call_hf_processor]. Override `_call_hf_processor` if the custom processor instead exposes a different callable interface.
 
 Examples:
 
diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md
index e420a3aa68e1..86bd68acef0c 100644
--- a/docs/contributing/profiling.md
+++ b/docs/contributing/profiling.md
@@ -22,6 +22,10 @@ To use the `torch.profiler` module, set the `profiler` entry to `'torch'` and `t
 - `torch_profiler_with_flops` to enable recording FLOPs, off by default
 - `torch_profiler_use_gzip` to control gzip-compressing profiling files, on by default
 - `torch_profiler_dump_cuda_time_total` to control dumping and printing the aggregated CUDA self time table, on by default
+- `torch_profiler_activities` to select worker activities. Defaults are
+  platform-specific: `["CPU"]` on CPU, `["CPU", "CUDA"]` on NVIDIA GPUs, and
+  `["CPU", "XPU"]` on Intel GPUs. Selecting `["CUDA"]` omits CPU annotations
+  and the AsyncLLM CPU trace to reduce profiling overhead and trace size.
 
 When using `vllm bench serve`, you can enable profiling by passing the `--profile` flag.
 
@@ -90,18 +94,18 @@ $ curl -X POST http://localhost:8000/stop_profile
 is Triton's GPU profiler. It can collect a low-overhead aggregate tree or a
 Chrome trace and works through the same vLLM profiling controls as the PyTorch
 and CUDA profilers. Proton currently supports NVIDIA GPUs through CUPTI and
-requires eager execution.
+supports CUDA graph attribution.
 
-Start a server with a local output directory:
+Start a server with a local output directory and graph attribution:
 
 ```bash
 vllm serve meta-llama/Llama-3.1-8B-Instruct \
-    --enforce-eager \
     --profiler-config '{
         "profiler": "proton",
         "proton_profiler_dir": "./proton_profile",
         "proton_output_format": "hatchet",
-        "proton_hook": "triton"
+        "proton_hook": "triton",
+        "proton_graph_attribution": true
     }'
 ```
 
@@ -111,7 +115,7 @@ rank-qualified output name, such as
 `proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet`,
 so distributed workers, restarted servers, and repeated profiling runs do not
 overwrite one another. A `profile_prefix` is included when supplied. Each
-profile is finalized by `/stop_profile` and is ready to inspect immediately.
+profile is written by `/stop_profile` and is ready to inspect immediately.
 
 The Proton-specific options are:
 
@@ -121,21 +125,39 @@ The Proton-specific options are:
 - `proton_mode`: an optional backend mode string
 - `proton_hook`: `triton` to record Triton launch metadata, or unset
 - `proton_output_format`: `hatchet`, `hatchet_msgpack`, `chrome_trace`, or unset
+- `proton_graph_attribution`: observe CUDA graph capture for replay attribution;
+  disabled by default and requires `proton_data: "tree"`
 
 `hatchet` and `hatchet_msgpack` require `proton_data: "tree"`, while
-`chrome_trace` requires `proton_data: "trace"`. CUDA graph profiling is not yet
-supported, so `--enforce-eager` is required for every Proton mode.
+`chrome_trace` requires `proton_data: "trace"`.
 
 Automatic backend selection is recommended. vLLM currently supports Proton's
 `cupti` backend on NVIDIA GPUs. ROCm support is not yet available. vLLM does not
 expose Proton's experimental instrumentation backend because current upstream
-Triton builds can produce profiles without timing metrics. Backend-specific
-modes, including `pcsampling`, can be selected with `proton_mode`.
-
-Triton 3.6 supports explicit `hatchet` and `chrome_trace` output. The
-`hatchet_msgpack` format and `periodic_flushing` mode require Triton 3.7 or
-newer; vLLM rejects these options with the detected version before starting a
-session.
+can produce profiles without timing metrics. When `proton_graph_attribution` is
+enabled, Proton observes vLLM's CUDA graph capture with the configured profiling
+session active, then deactivates that same session until profiling starts. This
+lets later profiles attribute replayed kernels without retaining model-startup
+activity. Backend-specific modes can be selected with `proton_mode`;
+`pcsampling` synchronizes the CUDA context and therefore requires
+`--enforce-eager`. When CUDA graphs are enabled (including encoder graphs),
+Proton requires `proton_graph_attribution: true` to collect replayed kernels.
+For Chrome traces, disable CUDA graphs with `--enforce-eager`.
+
+CUDA graph-attributed profiles support repeated `start_profile`/`stop_profile`
+runs. Each stop flushes and writes one tree-data phase while preserving the
+graph-aware session. Without `proton_graph_attribution`, each `stop_profile`
+instead finalizes and writes an independent Proton session.
+
+CUDA graph attribution requires Triton 3.7 or newer. It uses the phase data API
+to discard graph-capture activity and separate profiling runs. The `hatchet_msgpack`
+output format and `periodic_flushing` mode also require Triton 3.7 or newer.
+`periodic_flushing` cannot be combined with graph attribution because both
+manage the session's data phases. Ordinary Proton profiling remains available
+with Triton 3.6.
+
+Graph attribution retains capture metadata for the worker lifetime. With eager
+execution or no graphs to capture, profiling uses ordinary independent sessions.
 
 Inspect tree profiles with:
 
diff --git a/docs/contributing/vulnerability_management.md b/docs/contributing/vulnerability_management.md
index 8dd908a7ce5d..a73d451a6217 100644
--- a/docs/contributing/vulnerability_management.md
+++ b/docs/contributing/vulnerability_management.md
@@ -33,9 +33,7 @@ on GitHub. However, if you need to contact the VMT directly for an urgent issue,
 you may contact the following individuals:
 
 - Simon Mo - <simon.mo@hey.com>
-- Russell Bryant - <rbryant@redhat.com>
 - Juan Pérez de Algaba - <jperezde@redhat.com>
-- Huzaifa Sidhpurwala - <huzaifas@redhat.com>
 
 ## Slack Discussion
 
diff --git a/docs/deployment/frameworks/nebius.md b/docs/deployment/frameworks/nebius.md
new file mode 100644
index 000000000000..94d66c1de8b7
--- /dev/null
+++ b/docs/deployment/frameworks/nebius.md
@@ -0,0 +1,148 @@
+# Nebius Serverless AI
+
+[Nebius Serverless AI](https://docs.nebius.com/serverless/overview) runs GPU containers as persistent HTTP endpoints or finite jobs. This guide deploys the vLLM OpenAI-compatible server as a single-GPU Endpoint using the Nebius CLI.
+
+The Endpoint keeps running until you stop or delete it. This recipe uses a regular GPU instance; it does not configure autoscaling or automatic scale-to-zero.
+
+## Prerequisites
+
+- A Nebius project with permission to create and delete Serverless Endpoints, and sufficient GPU, VM, disk, and networking quotas.
+- The [Nebius CLI](https://docs.nebius.com/cli/install), installed and authenticated for that project.
+- A subnet in `eu-north1` with outbound access to Docker Hub and Hugging Face. The example uses the `gpu-l40s-a` platform and `1gpu-8vcpu-32gb` preset; check availability in your project before starting.
+- `curl`, `jq`, and `openssl` on your local machine.
+
+The example model, `Qwen/Qwen3-0.6B`, does not require a Hugging Face token. Other models may require accepting their license and supplying a token through Nebius [secret environment variables](https://docs.nebius.com/serverless/endpoints/manage).
+
+## Create the Endpoint
+
+The commands below were tested with Nebius CLI `0.12.265`. Set your project and subnet IDs explicitly. Use the same CLI profile throughout the guide.
+
+```bash
+export PROJECT_ID="<project-id>"
+export SUBNET_ID="<subnet-id>"
+export ENDPOINT_NAME="vllm-qwen-$(openssl rand -hex 6)"
+export AUTH_TOKEN="$(openssl rand -hex 32)"
+export MODEL_ID="Qwen/Qwen3-0.6B"
+export MODEL_REVISION="c1899de289a04d12100db370d81485cdf75e47ca"
+# vllm/vllm-openai:v0.19.1, Linux/amd64 image digest.
+export VLLM_IMAGE="vllm/vllm-openai@sha256:89c1d0629d377daa3f7f369cbea6167a7b48ea89aaacd12555e2b0b2f7f740d3"
+
+nebius ai endpoint create \
+    --parent-id "$PROJECT_ID" \
+    --subnet-id "$SUBNET_ID" \
+    --name "$ENDPOINT_NAME" \
+    --image "$VLLM_IMAGE" \
+    --container-command vllm \
+    --args "serve $MODEL_ID --revision $MODEL_REVISION --tokenizer-revision $MODEL_REVISION --host 0.0.0.0 --port 8000 --tensor-parallel-size 1 --max-model-len 4096 --gpu-memory-utilization 0.8" \
+    --platform gpu-l40s-a \
+    --preset 1gpu-8vcpu-32gb \
+    --container-port 8000/http \
+    --auth token --token "$AUTH_TOKEN" \
+    --disk-size 250Gi --shm-size 16Gi \
+    --public=false --preemptible=false \
+    --retries 1
+```
+
+The image and model revisions are pinned for reproducibility. When changing the image, check its CUDA/driver requirements against the selected Nebius platform. The model downloads into the container disk on startup; this example does not configure persistent model storage.
+
+The Endpoint token is separate from your Nebius CLI credentials. Keep it private. Nebius token authentication requires exactly one HTTP port. This example relies on authentication at the managed HTTPS URL; it does not set a second vLLM API key. Access from inside the subnet must be restricted to trusted clients.
+
+A public VM IP is not required to use the managed HTTPS URL. The subnet still needs outbound connectivity for image and model downloads.
+
+Save the ID of the Endpoint you created:
+
+```bash
+export ENDPOINT_ID="$(nebius ai endpoint get-by-name \
+    --parent-id "$PROJECT_ID" --name "$ENDPOINT_NAME" \
+    --format jsonpath='{.metadata.id}')"
+```
+
+If creation times out or your terminal disconnects, use the same project and name to find the Endpoint before attempting another creation. A local timeout does not mean that provisioning was cancelled.
+
+## Check readiness
+
+Inspect the state and recent logs:
+
+```bash
+nebius ai endpoint get "$ENDPOINT_ID"
+nebius ai endpoint logs "$ENDPOINT_ID" --tail 100 --timestamps
+```
+
+Wait for `RUNNING`, then select the managed HTTPS URL. The URL can appear while the Endpoint is still `STARTING`, when ingress may return HTTP 404. The `jq` expression requires exactly one HTTPS URL and preserves its scheme:
+
+```bash
+ENDPOINT_URL="$(nebius ai endpoint get "$ENDPOINT_ID" --format json \
+    | jq -er '[.status.public_endpoints[]? | select(startswith("https://"))]
+        | if length == 1 then .[0] else error("Expected one HTTPS URL") end')"
+export ENDPOINT_URL="${ENDPOINT_URL%/}"
+
+curl --fail-with-body --silent --show-error --max-time 10 \
+    "$ENDPOINT_URL/health" -H "Authorization: Bearer $AUTH_TOKEN"
+curl --fail-with-body --silent --show-error --max-time 10 \
+    "$ENDPOINT_URL/v1/models" -H "Authorization: Bearer $AUTH_TOKEN" | jq
+```
+
+`RUNNING` alone does not establish that the model has finished loading. Wait for `/health` to succeed and `/v1/models` to list `Qwen/Qwen3-0.6B` before sending a chat request. While waiting, inspect logs instead of creating another Endpoint. If startup does not succeed within your chosen time or spending limit, delete the Endpoint and investigate the error.
+
+## Send a chat request
+
+```bash
+curl --fail-with-body --silent --show-error --max-time 120 \
+    "$ENDPOINT_URL/v1/chat/completions" \
+    -H "Authorization: Bearer $AUTH_TOKEN" \
+    -H 'Content-Type: application/json' \
+    -d '{
+        "model": "Qwen/Qwen3-0.6B",
+        "messages": [{"role": "user", "content": "Say hello in one short sentence."}],
+        "max_tokens": 128,
+        "temperature": 0,
+        "chat_template_kwargs": {"enable_thinking": false}
+    }' | jq
+```
+
+Expect a chat completion containing an assistant message. The Qwen-specific template option disables thinking for this short example; it is not a universal option for other models.
+
+For streaming, add `"stream": true` to the JSON and use `curl -N` without piping into `jq`. A successful stream contains completion chunks followed by `data: [DONE]`. For an OpenAI-compatible client, use `${ENDPOINT_URL}/v1` as the base URL and the Endpoint token as the API key.
+
+Verify that the same inference URL rejects requests without the token:
+
+```bash
+curl --silent --show-error --max-time 10 --output /dev/null \
+    --write-out '%{http_code}\n' "$ENDPOINT_URL/v1/models"
+```
+
+Expect HTTP `401` or `403`. Do not share an Endpoint that accepts unauthenticated inference unexpectedly.
+
+## Stop or delete the Endpoint
+
+To stop serving while preserving the Endpoint configuration:
+
+```bash
+nebius ai endpoint stop "$ENDPOINT_ID"
+nebius ai endpoint get "$ENDPOINT_ID"
+```
+
+Let the synchronous stop command complete and confirm `STOPPED` before starting again. To start it again, run `nebius ai endpoint start "$ENDPOINT_ID"`, retrieve its current URL, and repeat the readiness checks. Starting again can require fresh capacity and new image/model downloads; do not rely on the container disk as persistent storage.
+
+When finished, delete the Endpoint and confirm that getting the same ID returns `NotFound`:
+
+```bash
+nebius ai endpoint delete "$ENDPOINT_ID"
+nebius ai endpoint get "$ENDPOINT_ID"
+unset AUTH_TOKEN
+```
+
+Nebius deletes the managed VM and container disk with the Endpoint. Separately mounted storage has its own lifecycle and billing. Closing your terminal or interrupting a request does not stop the Endpoint. If deletion fails, keep the ID and retry cleanup; a generic connection or authentication error is not confirmation of deletion.
+
+## Troubleshooting
+
+| Symptom | Check |
+| --- | --- |
+| `PROVISIONING` or `NotEnoughResources` | Inspect quota, region, platform and preset. Capacity availability can differ from quota. |
+| `RUNNING`, but HTTP 502/503 | Check model-download and engine-start logs, subnet egress, host binding and port 8000. |
+| HTTP 401/403 | Use the Endpoint token, not the Nebius control-plane token. |
+| CUDA or out-of-memory error | Check image/driver compatibility, the model's memory requirements and maximum context length. |
+| Model not found | Verify the requested model ID and pinned revision. |
+| Interrupted stream | Inspect ingress and server logs. A partial response should not be automatically replayed as a new generation. |
+
+See the [Nebius vLLM cookbook example](https://github.com/nebius/serverless-ai-cookbook/tree/main/inference/vllm-endpoint) and [Endpoint management documentation](https://docs.nebius.com/serverless/endpoints/manage) for further details.
diff --git a/docs/deployment/integrations/kthena.md b/docs/deployment/integrations/kthena.md
index 0285b733c6b8..73dcb3dafd72 100644
--- a/docs/deployment/integrations/kthena.md
+++ b/docs/deployment/integrations/kthena.md
@@ -2,7 +2,9 @@
 
 [**Kthena**](https://github.com/volcano-sh/kthena) is a Kubernetes-native LLM inference platform that transforms how organizations deploy and manage Large Language Models in production. Built with declarative model lifecycle management and intelligent request routing, it provides high performance and enterprise-grade scalability for LLM inference workloads.
 
-This guide shows how to deploy a production-grade, **multi-node vLLM** service on Kubernetes.
+Kthena manages inference workloads through `ModelServing` and provides model-aware request routing through `ModelServer` and `ModelRoute`.
+
+This guide shows how to deploy a **multi-node vLLM** service on Kubernetes using `ModelServing` and expose it through a Kubernetes Service.
 
 We’ll:
 
@@ -18,7 +20,7 @@ You need:
 
 - A Kubernetes cluster with **GPU nodes**.
 - `kubectl` access with cluster-admin or equivalent permissions.
-- **Volcano** installed for gang scheduling.
+- **Volcano 1.14 or later** for the role-level gang scheduling used in this example.
 - **Kthena** installed with the `ModelServing` CRD available.
 - A valid **Hugging Face token** if loading models from Hugging Face Hub.
 
@@ -30,21 +32,22 @@ helm repo update
 helm install volcano volcano-sh/volcano -n volcano-system --create-namespace
 ```
 
-This provides the gang-scheduling and network topology features used by Kthena.
+Volcano provides the gang scheduling used in this example. It is optional for Kthena deployments that do not use Volcano scheduling.
 
 ### 1.2 Install Kthena
 
 ```bash
-helm install kthena oci://ghcr.io/volcano-sh/charts/kthena --version v0.1.0 --namespace kthena-system --create-namespace
+helm install kthena oci://ghcr.io/volcano-sh/charts/kthena --version v1.0.0 --namespace kthena-system --create-namespace
 ```
 
 - The `kthena-system` namespace is created.
-- Kthena controllers and CRDs, including `ModelServing`, are installed and healthy.
+- Kthena controllers and CRDs, including `ModelServing`, are installed.
 
 Validate:
 
 ```bash
-kubectl get crd | grep modelserving
+kubectl get pods -n kthena-system
+kubectl get crd modelservings.workload.serving.volcano.sh
 ```
 
 You should see:
@@ -57,7 +60,7 @@ modelservings.workload.serving.volcano.sh   ...
 
 ## 2. The Multi-Node vLLM `ModelServing` Example
 
-Kthena provides an example manifest to deploy a **multi-node vLLM cluster running Llama**. Conceptually this is equivalent to the vLLM production stack Helm deployment, but expressed with `ModelServing`.
+Kthena provides an example manifest to deploy a **multi-node vLLM cluster running Llama**. Each role replica contains an entry pod and its worker pods.
 
 A simplified version of the example (`llama-multinode`) looks like:
 
@@ -116,10 +119,9 @@ Commands:
           - -c
           - >
             bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh
-            leader --ray_cluster_size=2; python3 -m
-            vllm.entrypoints.launchers.api_server.entry --port 8080 --model
-            meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8
-            --pipeline-parallel-size 2
+            leader --ray_cluster_size=2;
+            vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8
+            --pipeline-parallel-size 2 --distributed-executor-backend ray
         ```
 
         Worker:
@@ -166,9 +168,9 @@ Save one of the following manifests to `modelserving.yaml`:
             restartGracePeriodSeconds: 60
             gangPolicy:
               minRoleReplicas:
-                405b: 1
+                llama-405b: 1
             roles:
-              - name: 405b
+              - name: llama-405b
                 replicas: 2
                 entryTemplate:
                   spec:
@@ -256,9 +258,9 @@ Save one of the following manifests to `modelserving.yaml`:
             restartGracePeriodSeconds: 60
             gangPolicy:
               minRoleReplicas:
-                405b: 1
+                llama-405b: 1
             roles:
-              - name: 405b
+              - name: llama-405b
                 replicas: 2
                 entryTemplate:
                   spec:
@@ -275,7 +277,7 @@ Save one of the following manifests to `modelserving.yaml`:
                           - sh
                           - -c
                           - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=2;
-                            vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2"
+                            vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2 --distributed-executor-backend ray"
                         resources:
                           limits:
                             nvidia.com/gpu: "8"
@@ -383,24 +385,24 @@ kubectl get pod -owide -l modelserving.volcano.sh/name=llama-multinode
 Example output (from docs):
 
 ```text
-NAMESPACE   NAME                          READY   STATUS    RESTARTS   AGE   IP            NODE           ...
-default     llama-multinode-0-405b-0-0    1/1     Running   0          15m   10.244.0.56   192.168.5.12   ...
-default     llama-multinode-0-405b-0-1    1/1     Running   0          15m   10.244.0.58   192.168.5.43   ...
-default     llama-multinode-0-405b-1-0    1/1     Running   0          15m   10.244.0.57   192.168.5.58   ...
-default     llama-multinode-0-405b-1-1    1/1     Running   0          15m   10.244.0.53   192.168.5.36   ...
+NAMESPACE   NAME                                   READY   STATUS    RESTARTS   AGE   IP            NODE           ...
+default     llama-multinode-0-llama-405b-0-0    1/1     Running   0          15m   10.244.0.56   192.168.5.12   ...
+default     llama-multinode-0-llama-405b-0-1    1/1     Running   0          15m   10.244.0.58   192.168.5.43   ...
+default     llama-multinode-0-llama-405b-1-0    1/1     Running   0          15m   10.244.0.57   192.168.5.58   ...
+default     llama-multinode-0-llama-405b-1-1    1/1     Running   0          15m   10.244.0.53   192.168.5.36   ...
 ```
 
 Pod name pattern:
 
 - `llama-multinode-<group-idx>-<role-name>-<replica-idx>-<ordinal>`.
 
-The first number indicates `ServingGroup`. The second (`405b`) is the `Role`. The remaining indices identify the pod within the role.
+The first index identifies the `ServingGroup`, followed by the role name (`llama-405b`), role replica index, and pod index.
 
 ---
 
-## 6. Accessing the vLLM OpenAI-Compatible API
+## 5. Accessing the vLLM OpenAI-Compatible API
 
-Expose the entry via a Service:
+Save the following Service as `service.yaml` to expose the entry pods:
 
 ```yaml
 apiVersion: v1
@@ -412,7 +414,6 @@ spec:
   selector:
     modelserving.volcano.sh/name: llama-multinode
     modelserving.volcano.sh/entry: "true"
-    # optionally further narrow to leader role if you label it
   ports:
     - name: http
       port: 80
@@ -420,9 +421,10 @@ spec:
   type: ClusterIP
 ```
 
-Port-forward from your local machine:
+Apply the Service, then port-forward from your local machine:
 
 ```bash
+kubectl apply -f service.yaml
 kubectl port-forward svc/llama-multinode-openai 30080:80 -n default
 ```
 
@@ -450,11 +452,12 @@ You should see an OpenAI-style response from vLLM.
 
 ---
 
-## 7. Clean Up
+## 6. Clean Up
 
 To remove the deployment and its resources:
 
 ```bash
+kubectl delete -f service.yaml
 kubectl delete modelserving llama-multinode -n default
 ```
 
@@ -464,3 +467,5 @@ If you’re done with the entire stack:
 helm uninstall kthena -n kthena-system   # or your Kthena release name
 helm uninstall volcano -n volcano-system
 ```
+
+For model-aware routing and prefill-decode disaggregation, see the [Kthena documentation](https://kthena.volcano.sh/docs/user-guide/router-routing).
diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md
index 3fdc23ff8422..d175add40b84 100644
--- a/docs/design/attention_backends.md
+++ b/docs/design/attention_backends.md
@@ -56,6 +56,55 @@ llm = LLM(
 
 ## Backend Selection Behavior
 
+### Triton/FlashAttention Composite
+
+On Hopper, `TRITON_FLASH_ATTN` is preferred for compatible multimodal-prefix
+configurations. It uses Triton when the current queries need bidirectional
+image attention and FlashAttention for causal text prefills and decode, with a
+shared KV cache. The causal child must resolve to FA4. This happens
+automatically for FA4-only shapes such as head size 512 and models whose
+version policy promotes all layers to FA4, including Gemma 4. A standalone
+head-size-256 configuration whose version policy selects FA3 falls back to
+Triton for the whole layer.
+
+No attention override is needed for supported multimodal models:
+
+```bash
+vllm serve google/gemma-4-31B-it
+```
+
+### Triton/FlashInfer Composite
+
+On Blackwell, `TRITON_FLASHINFER` is preferred for compatible
+multimodal-prefix configurations, including Gemma 4 with BF16 or FP8 KV cache. It uses
+Triton when a batch's current queries require bidirectional image attention,
+and FlashInfer for causal text prefills and decode. Historical image tokens
+alone do not select Triton.
+
+This instantiates `create_composite_attention_backend` from
+`vllm/v1/attention/backends/composite.py` with Triton,
+FlashInfer, and `MMPrefixAttentionRouting`. The reusable factory owns child
+implementations, metadata dispatch, compatible cache requirements, and workspace
+sharing. The routing policy selects the child and defines graph-capture safety.
+Other combinations can reuse the same machinery.
+
+The backend supports head dimensions 256/512, FP16/BF16 and FP8 KV cache,
+and 64-token kernel pages with a head-major cache layout. TRTLLM handles causal
+attention at both head dimensions; its hdim512 kernels do not support 128-token
+pages.
+For Gemma 4, full CUDA graphs cover single-token batches; multi-token batches
+use the non-full-graph execution path. Models whose image masks extend beyond
+the sliding window cannot use full attention graphs with this composite.
+Context parallelism, R-SWA, attention sinks,
+and adaptive verification are not supported by this composite.
+
+It can also be selected explicitly:
+
+```bash
+vllm serve google/gemma-4-31B-it \
+    --attention-backend TRITON_FLASHINFER
+```
+
 ### Manual Selection
 
 When you explicitly set a backend via `--attention-backend` or `AttentionConfig`:
diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md
new file mode 100644
index 000000000000..53a309371783
--- /dev/null
+++ b/docs/design/hisparse.md
@@ -0,0 +1,254 @@
+# HiSparse local KV offload architecture
+
+Status: experimental
+
+## The short version
+
+There are three different jobs:
+
+1. The normal KV cache system manages GPU block pools and tables.
+2. `HiSparseCoordinator` manages logical host blocks, source-prefix identity, and
+   host/GPU residency transitions.
+3. `HiSparseConnector` carries residency work between scheduler and worker;
+   `HiSparseWorker` coordinates transfers, while per-cache
+   `HiSparseRuntime` objects own host/hot views and GPU replacement state.
+
+Neither worker-side object allocates or frees logical blocks. HMA provides the
+GPU allocation shared by the resident and hot groups; it does not manage CPU
+memory or KV identity.
+
+```text
+request ────► HiSparseCoordinator ── source blocks + residency policy
+                    │
+                    ├──► KV cache manager ── resident/hot GPU leases (HMA)
+                    │
+                    └──► HiSparseConnector ── host bytes + copies + GPU LRU
+```
+
+This is a local KV connector. When P/D or another offload connector is also
+configured, `MultiConnector` composes it with `HiSparseConnector`.
+
+`host_pool_gib` is configured on `HiSparseConnector` and is the usable
+host-cache capacity per data-parallel replica, not a node-wide memory budget.
+Tensor-parallel ranks hold replicated views of that logical cache. Those views
+may use private per-rank backing or one shared
+physical allocation without changing the configured capacity. Physical host
+memory consumption is therefore topology- and implementation-dependent. The
+realized capacity may be slightly smaller because the budget is rounded down
+to complete host blocks.
+
+## Ownership
+
+| Thing | Owner | What “owner” means |
+| --- | --- | --- |
+| HiSparse source and prefix identity | `HiSparseCoordinator` | maps tokens to logical host blocks |
+| Resident GPU block leases | normal KV cache manager | allocates and frees HMA blocks |
+| Resident block tables | normal KV cache manager | tells attention where resident pages are |
+| Residency transitions | `HiSparseCoordinator` | plans spill-before-free transactions |
+| Logical host block allocation | `HiSparseCoordinator` | owns the separate CPU block pool and its lifecycle |
+| Pinned host-pool lifecycle | `HiSparseWorker` | worker-wide backing and teardown |
+| Per-cache host view and hot contents | `HiSparseRuntime` | binds host/hot storage and fills cache-manager-provided hot leases |
+| Hot row map and LRU | `HiSparseRuntime` | resolves hits and chooses victims on GPU |
+| Resident-cache route | `HiSparseCacheHandle` | exposes resident or host/hot resolution to attention |
+| Sparse attention | attention backend | consumes a device cache and physical row IDs |
+| HMA | allocator | provides GPU capacity; owns no KV meaning |
+
+The key distinction is logical allocation versus contents. `HiSparseCoordinator`
+owns host block IDs and their request/prefix associations. `HiSparseWorker` and
+its per-cache runtimes own the corresponding bytes. The normal cache manager
+sees only device pools.
+The source group has `block_pool_id=None`; device-pool consumers must narrow it
+before indexing, so host ownership cannot masquerade as a numeric GPU pool.
+
+For single-node MP tensor parallelism, every TP worker maps the same pinned host
+pool and uses the same block and layer offsets. MLA source KV is replicated
+across TP ranks, so this stores one physical copy instead of one copy per rank.
+TP rank 0 writes the shared host pool; peers wait on its IPC events before
+reading it. Other executor and parallel layouts retain private per-rank pools.
+
+The shared layout backs the per-replica logical capacity with one physical pool;
+the private layout allocates one physical pool per rank. Physical pool size
+includes block-stride alignment.
+
+## Code boundary
+
+```text
+scheduler process                           worker process
+
+HiSparseConnector                          HiSparseConnector
+  └─ HiSparseCoordinator                         └─ HiSparseWorker
+       │                                          │
+       │ connector metadata                       ├─ host bytes
+       │ - page transfers                         ├─ copy scheduling
+       │ - block-table replacements               └─ per-layer hot state
+       └───────────────────────────────────────────────►│
+       ◄──────── connector worker metadata ─────────────┘
+                    enqueued and completed transfer IDs
+```
+
+The command travels in `kv_connector_metadata`; transfer updates return in
+`KVConnectorOutput.kv_connector_worker_meta`. The model runner does not
+interpret page transfers. Enqueue acknowledgements let the scheduler release
+source leases in stream order; completion acknowledgements publish the copied
+host pages.
+
+## Resident device pages
+
+Resident pages are intentionally outside `HiSparseRuntime`.
+
+KV-cache initialization binds cache-manager allocations to the attention-facing
+`HiSparseCacheHandle` before constructing `HiSparseWorker`. That same
+handle's runtime retains the resident source index needed by a transfer plan.
+There is no second resident object or registration wrapper.
+
+```text
+KV cache setup
+   │
+   ├─ bind resident allocation ──► HiSparseCacheHandle
+   │                               cache + block table + slot mapping
+   │
+   ├─ bind host/hot allocation ──► HiSparseRuntime
+   │                               host + hot + GPU LRU
+   │
+   └─ register cache handles ────► HiSparseWorker
+                                   step-level transfers
+
+HiSparseWorker registers the same HiSparseCacheHandle objects directly
+```
+
+Attention construction links each layer to the most recent layer that actually
+owns an indexer. This releases a follower's duplicate LRU tensors before GPU
+memory profiling. Cache binding only attaches storage; it does not infer
+semantic groups from the physical packed-tensor order. The construction cursor
+is discarded with the worker's pinned state.
+
+Every HiSparse decode batch uses the same fused resolver. It checks resident
+pages first, then hot rows, then pinned host memory. A resident hit exits inside
+the kernel before hot-LRU lookup or host copying; there is no framework-level
+residency route or separate CUDA graph. No CPU decision is added to the decode
+path. The resolver consumes the existing graph-stable request mapping from
+attention metadata; neither the worker nor individual cache handles keep a
+duplicate mapping.
+
+Speculative decoding resolves and consumes each verification step in order.
+Each step receives distinct replayable plan rows while sharing the request's
+hot-cache state, so a later step cannot reuse a hot row before an earlier step
+has consumed it.
+
+## P/D import target
+
+The decoder chooses the landing target once per request from the normal cache
+admission calculation. If the complete imported prefix fits the device pools,
+NIXL transfers it directly into resident GPU pages. Otherwise, if the fixed
+host-backed GPU footprint and host source blocks fit, the request imports into
+the host tier. There is no context-length threshold or other heuristic, and a
+request waiting for capacity retains its choice across admission retries.
+
+A host import reads through a bounded decoder-GPU staging pool before copying
+into registered host memory. Pages needed immediately are mirrored into their
+resident destinations during that copy. Both landing targets then use the same
+fused decode resolver described above.
+
+## Indexer KV offloading
+
+HiSparse does not keep a private CPU copy of indexer KV. The indexer remains a
+normal prefix-cacheable GPU cache group. If `OffloadingConnector` is configured
+with HiSparse, it stores and restores that group through the generic KV
+offloading path; HiSparse continues to own only the sparse MLA host tier.
+
+The two prefix sources can have different hit lengths. When the HiSparse host
+prefix extends beyond the GPU-resident indexer prefix, the scheduler asks
+`OffloadingConnector` to restore only the missing indexer suffix, capped at the
+host prefix boundary. If that suffix is unavailable, all groups fall back to
+the shorter prefix they share. NIXL P/D transfers continue to place indexer KV
+directly in its GPU group.
+
+## Spill transaction
+
+A resident block cannot be reused until its contents have been handed to the
+worker.
+
+```text
+HiSparseCoordinator                            HiSparseWorker
+          │                                     │
+          │ pin source and destination leases   │
+          │── SparseKVPageTransfer ─────────────►│
+          │                                     │ enqueue GPU-to-host copy
+          │◄── enqueued transfer ID ────────────│
+          │ replace resident table entry        │
+          │ release resident lease to HMA       │
+          │                                     │ copy reaches its event
+          │◄── completed transfer ID ───────────│
+          │ mark host page valid                │
+          │ release destination host lease      │
+```
+
+“Enqueued” means the copy has entered the worker stream. Stream ordering makes
+it safe to reuse the resident GPU block for later work, but the host page is
+not yet published. “Completed” means the worker has observed the copy's event;
+only then does the coordinator publish the host page for prefix reuse and
+release its destination lease. A host-write event separately protects direct
+CPU readers from writes already queued on the accelerator.
+
+The worker transfer contains only its transfer ID and physical copy
+coordinates. Request identity and logical page state remain in the scheduler.
+
+## Hot lookup and LRU
+
+The NVIDIA CUDA path keeps replacement entirely on the accelerator:
+
+```text
+top-K logical positions
+        │
+        ▼
+resident page? ── yes ──► resident physical row
+        │ no
+        ▼
+hot row? ──────── yes ──► existing hot physical row + update GPU LRU
+        │ no
+        ▼
+choose GPU LRU victim ──► copy pinned host row ──► hot physical row
+```
+
+ROCm is not currently supported because the fused HiSparse cache operations are
+implemented only by CUDA kernels. A future platform-specific worker may provide
+the same command, output, and cache-resolution boundaries.
+
+## Main classes
+
+| Class | Inherits / implements | Responsibility |
+| --- | --- | --- |
+| `HiSparseCoordinator` | plain scheduler component | host allocation, source prefixes, resident leases, and spill state machine |
+| `HiSparseConnector` | `KVConnectorBase_V1`, `SupportsHMA` | scheduler/worker metadata and lifecycle boundary |
+| `HiSparseResidentManager` | `SingleTypeKVCacheManager` | normal block-pool bookkeeping with host-backed holes |
+| `PagedCacheView` | immutable data object | shared resident/hot HMA tensor binding |
+| `HiSparseWorker` | connector-owned worker component | worker-wide transfer scheduling and host-pool lifecycle |
+| `HiSparseRuntime` | plain worker-owned component | per-cache host/hot tensors, GPU LRU, and fused resolution |
+| `HiSparseCacheHandle` | plain attention component | resident view and fused cache resolution |
+| `SparseKVOffloadCommand` | dataclass | opaque scheduler-to-worker work |
+
+## Performance invariants
+
+- Resident hits bypass hot-LRU lookup and host copies inside the fused resolver.
+- Hot lookup, victim selection, and LRU updates stay on the GPU.
+- A hot miss still copies directly from registered pinned host memory.
+- Top-K resolution stays inside the attention invocation and remains graph
+  capturable.
+- Compatible layers still share one miss plan.
+- Index-sharing followers release their private LRU state before memory sizing.
+- Indexer KV is untouched by HiSparse unless a generic KV offloader is configured.
+- Resident and hot leases can still share one packed HMA allocation.
+- No device scalar readback or CPU/device metadata round trip is added.
+- The abstraction wraps the fused kernel; it does not add another kernel
+  launch.
+- When HiSparse is disabled, the scheduler does not construct an offload
+  command or empty update table.
+
+## What remains platform-specific
+
+The command/result and attention-layer boundaries can be shared. The host
+allocator, copy implementation, hot layout, and replacement policy should stay
+platform-specific. NVIDIA uses the current accelerator LRU and fused host/hot
+kernel. ROCm is not currently supported; AMD or other accelerator backends can
+implement their own worker without forcing NVIDIA's policy into the shared
+boundary.
diff --git a/docs/design/metrics.md b/docs/design/metrics.md
index c7bc85756e0d..3e070f588890 100644
--- a/docs/design/metrics.md
+++ b/docs/design/metrics.md
@@ -33,7 +33,8 @@ In v1, an extensive set of metrics are exposed via a Prometheus-compatible `/met
 - `vllm:request_prompt_tokens` (Histogram) - Histogram of input prompt token counts.
 - `vllm:request_generation_tokens` (Histogram) - Histogram of generation token counts.
 - `vllm:time_to_first_token_seconds` (Histogram) - Time to first token (TTFT).
-- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency.
+- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency (time between consecutive streamed outputs).
+- `vllm:request_time_per_output_token_seconds` (Histogram) - Per-request Time Per Output Token (TPOT).
 - `vllm:e2e_request_latency_seconds` (Histogram) - End-to-end request latency.
 - `vllm:request_prefill_time_seconds` (Histogram) - Request prefill time.
 - `vllm:request_decode_time_seconds` (Histogram) - Request decode time.
@@ -49,7 +50,8 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of
 - `vllm:e2e_request_latency_seconds_bucket` - End to end request latency measured in seconds.
 - `vllm:prompt_tokens` - Prompt tokens.
 - `vllm:generation_tokens` - Generation tokens.
-- `vllm:inter_token_latency_seconds` - Inter-token latency (Time Per Output Token, TPOT) in seconds.
+- `vllm:inter_token_latency_seconds` - Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains exactly one token, but differs from request-level TPOT when an output contains multiple tokens or when aggregation weights differ.
+- `vllm:request_time_per_output_token_seconds` - Per-request Time Per Output Token (TPOT) in seconds, computed as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. It is recorded as zero for requests that generate no more than one token.
 - `vllm:time_to_first_token_seconds` - Time to First Token (TTFT) latency in seconds.
 - `vllm:num_requests_running` (also, `_swapped` and `_waiting`) - Number of requests in the RUNNING, WAITING, and SWAPPED states.
 - `vllm:kv_cache_usage_perc` - Percentage of used cache blocks by vLLM.
@@ -63,6 +65,20 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of
 
 See [the PR which added this Dashboard](https://github.com/vllm-project/vllm/pull/2316) for interesting and useful background on the choices made here.
 
+`vllm:inter_token_latency_seconds` records one sample per streamed output event
+(the wall-clock gap between successive outputs), whereas
+`vllm:request_time_per_output_token_seconds` is recorded once per finished request
+as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. These two
+metrics differ when an output bundles multiple tokens (e.g. speculative
+decoding) or when mean ITL and mean per-request TPOT are aggregated over
+differing weights. See [Understanding the latency metrics](../benchmarking/cli.md#understanding-the-latency-metrics)
+for the benchmark definitions. Requests that generate no more than one token
+are recorded with a TPOT of zero, while `vllm bench serve` excludes them from
+its TPOT statistics. Therefore, the Prometheus histogram and benchmark TPOT
+statistics can differ. Use `vllm:request_time_per_output_token_seconds` for
+request-level TPOT and `vllm:inter_token_latency_seconds` when you specifically
+want the inter-output latency distribution.
+
 ### Prometheus Client Library
 
 Prometheus support was initially added [using the aioprometheus library](https://github.com/vllm-project/vllm/pull/1890), but a switch was made quickly to [prometheus_client](https://github.com/vllm-project/vllm/pull/2730). The rationale is discussed in both linked PRs.
@@ -238,8 +254,7 @@ statistics relating to that iteration:
   iteration.
 - The prefill intervals for any requests that completed prefill in
   this iteration.
-- The inter-token intervals (Time Per Output Token, TPOT), for all
-  requests included in this iteration.
+- The inter-token intervals, for all requests included in this iteration.
 - The Time-To-First-Token (TTFT) for any requests that completed
   prefill in this iteration. However, we calculate this interval
   relative to when the request was first received by the frontend
@@ -655,7 +670,7 @@ fall under the more general heading of "Observability".
 vLLM has support for OpenTelemetry tracing:
 
 - Added by <https://github.com/vllm-project/vllm/pull/4687> and reinstated by <https://github.com/vllm-project/vllm/pull/20372>
-- Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces`
+- Configured with `--otlp-traces-endpoint` and `--collect-detailed-traces`
 - [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/)
 - [User-facing docs](../../examples/observability/opentelemetry/README.md)
 - [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f)
diff --git a/docs/design/mm_processing.md b/docs/design/mm_processing.md
index a2112a2922ed..202d8d04c166 100644
--- a/docs/design/mm_processing.md
+++ b/docs/design/mm_processing.md
@@ -8,9 +8,16 @@ In vLLM's rendering pipeline (see [BaseRenderer][vllm.renderers.base.BaseRendere
 
 Since Transformers 5.10, `ProcessorMixin` now allows multi-modal inputs to be passed by themselves. However, certain subclasses (such as `ChameleonProcessor`) and older out-of-tree implementations may still define their own `__call__` method that assumes the presence of text with corresponding placeholder tokens. This causes a problem as we don't have the original text anymore to pass to these processors.
 
-To work around this, each model defines how to generate dummy text based on the number of multi-modal inputs, via [get_dummy_text][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_text], which its override of [_get_hf_processor_text][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_processor_text] returns so that [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main] passes it to the HF processor together with the multi-modal inputs to obtain the processed multi-modal data.
+To work around this, each model defines how to generate dummy text based on the number of multi-modal inputs, via [get_dummy_text][vllm.multimodal.processing.BaseDummyInputsBuilder.get_dummy_text], which its override of [_get_hf_mm_text][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_text] returns. [_get_hf_mm_inputs][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs] adds this text and `truncation=False` to the normalized HF inputs. The truncation setting is attached only to calls that use dummy text, so unrelated processor calls do not receive a text-specific option.
 
-Similarly, since the multi-modal data extracted by vLLM may not match what a specific HF processor expects, [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main] allows each model to adapt the inputs via [_preprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._preprocess_hf_mm_data] (e.g. renaming keys like `audios` to `audio` or injecting extra keyword arguments such as `sampling_rate`) and the outputs via [_postprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._postprocess_hf_mm_data], without having to reimplement the entire method.
+HF processing is split into four overridable stages:
+
+1. [_get_hf_mm_inputs][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs] extracts and normalizes processor data, kwargs, and passthrough data. It centrally renames vLLM's `audios` field to the standard HF `audio` keyword; processors that require another spelling can adapt it in an override.
+2. [_call_hf_processor][vllm.multimodal.processing.BaseMultiModalProcessor._call_hf_processor] invokes the HF processor. Models can override it when they use a different callable or need separate construction and invocation kwargs.
+3. [_finalize_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._finalize_hf_mm_data] merges passthrough fields into the processed output. It lazily creates an empty output when there is no processor data.
+4. [_postprocess_hf_mm_data][vllm.multimodal.processing.BaseMultiModalProcessor._postprocess_hf_mm_data] performs model-specific output transformations after passthrough fields have been merged.
+
+This division lets most models customize one stage without reimplementing [_apply_hf_processor_main][vllm.multimodal.processing.BaseMultiModalProcessor._apply_hf_processor_main]. Implementations that genuinely require multiple processor calls or manually construct their complete output may still override that method, but must use `_get_hf_mm_inputs` and `_finalize_hf_mm_data` to retain the common input and output behavior.
 
 ## Prompt Update Detection
 
diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md
index 6f66be5eb9f5..a6debd7dba81 100644
--- a/docs/design/nixl_kv_push_connector.md
+++ b/docs/design/nixl_kv_push_connector.md
@@ -181,6 +181,50 @@ The completion notif sent from P to D after a WRITE is the existing
 is D's own request id, taken from the registration), so the D-side
 accounting code is unchanged.
 
+## Pipeline parallelism and hybrid KV caches
+
+The pull connector requires the local and remote workers to expose a
+congruent list of KV regions — region *i* here corresponds to region
+*i* there. That assumption breaks under pipeline parallelism (PP)
+combined with a hybrid (HMA) KV layout:
+
+* a PP-sharded prefiller (**P**) holds only a slice of the model's
+  layers while the `PP=1` decoder (**D**) holds them all, so region
+  counts differ;
+* with HMA, several layer names are pooled into one region and the layer
+  that represents a pooled region can differ between P and D.
+
+`NixlPushConnector` handles this for PP-sharded producers by routing
+**by layer-name (member) identity** instead of by region index. Each
+worker advertises which layer names back each of its NIXL regions in the
+handshake metadata (`NixlAgentMetadata.region_members`). A producer that
+needs member routing derives its member-major layout once, when it
+registers its KV caches, and every transfer it issues uses that order.
+`add_remote_agent` then selects exactly the remote regions this stage
+owns and reorders them to match, so both sides stay paired regardless of
+how each remote rank happens to order its metadata.
+
+Addresses, block lengths, strides, and per-region capacities follow the same
+member order. Descriptor offsets use each member's region capacity, so P and
+D need not allocate the same number of blocks. Physical allocations are still
+registered once, even when multiple members share them.
+
+Invariants enforced when the remote regions are aligned:
+
+* every locally owned member must be advertised exactly once by the
+  remote; a missing member fails the handshake rather than silently
+  leaving that layer's KV stale, and remote-only members (owned by other
+  PP stages) are ignored;
+* a remote that omits member metadata while the local layout requires
+  member routing fails the handshake instead of falling back to
+  region-index routing;
+* member order is a property of the local layout alone, so the same local
+  source descriptors serve every remote engine and TP rank; only the
+  remote descriptor list is rebuilt per rank.
+
+Decode-side PP is unsupported because completions are counted per
+consumer rank. Mamba/SSM hybrids are unsupported under PP.
+
 ## Scheduler-side responsibilities
 
 `NixlPushConnectorScheduler` extends the base scheduler with:
diff --git a/docs/design/torch_compile.md b/docs/design/torch_compile.md
index 4dc0da0c7d65..1ffc0ac53a40 100644
--- a/docs/design/torch_compile.md
+++ b/docs/design/torch_compile.md
@@ -27,6 +27,12 @@ With all these factors taken into consideration, usually we can guarantee that t
 
 A unique aspect of vLLM's `torch.compile` integration, is that we guarantee all the compilation finishes before we serve any requests. No requests will trigger new compilations. Otherwise, the engine would be blocked on that request, and the response time will have unexpected spikes.
 
+An explicit `TRITON_CACHE_DIR` is preserved, allowing processes and model
+configurations to reuse Triton artifacts. The directory can start empty: the
+first compilation populates it and later matching specializations reuse it.
+When copying compilation caches, also copy this directory if set. Cold-cache
+tests should set it to a fresh directory.
+
 By default, the cache saves compiled artifacts as binary files. If you would like to interact with the generated code for debugging purposes, set the field `compile_cache_save_format=unpacked` in the compilation config, or omit this and set the env variable `VLLM_COMPILE_CACHE_SAVE_FORMAT=unpacked`.
 
 ## Dynamic shapes and vllm guard dropping
diff --git a/docs/features/README.md b/docs/features/README.md
index 28362f401477..1df8960ee354 100644
--- a/docs/features/README.md
+++ b/docs/features/README.md
@@ -64,7 +64,7 @@ th:not(:first-child) {
 | [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ |
+| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
 | [pooling](../models/pooling_models/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
diff --git a/docs/features/automatic_prefix_caching.md b/docs/features/automatic_prefix_caching.md
index 03cca3b38d6c..cfc8c19acb2b 100644
--- a/docs/features/automatic_prefix_caching.md
+++ b/docs/features/automatic_prefix_caching.md
@@ -15,7 +15,7 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl
 
 ## Hybrid Mamba models
 
-Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-fine-grained-prefix-cache` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it.
+Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it.
 
 This helps when many requests share a long system prompt and then diverge. It is off by default, and takes effect only when all of the following hold:
 
@@ -28,7 +28,7 @@ This helps when many requests share a long system prompt and then diverge. It is
 vllm serve <hybrid-model> \
     --mamba-cache-mode align \
     --prefix-match-unit 64 \
-    --enable-mamba-fine-grained-prefix-cache
+    --enable-mamba-shared-prefix-checkpoint
 ```
 
 `--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. Under `align` that is the block size itself, so no sub-block boundary exists and the flag has no effect.
diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md
index 789878f48a96..7353bd359462 100644
--- a/docs/features/batch_invariance.md
+++ b/docs/features/batch_invariance.md
@@ -145,6 +145,9 @@ When batch invariance is enabled, vLLM:
 1. Uses deterministic kernel implementations for attention and other operations
 2. Ensures consistent numerical behavior across different batch sizes
 3. Disables certain optimizations that may introduce non-determinism (such as custom all-reduce operations in tensor parallel mode)
+4. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell),
+   runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out
+   (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes).
 
 !!! note
     Enabling batch invariance may impact performance compared to the default non-deterministic mode. This trade-off is intentional to guarantee reproducibility.
diff --git a/docs/features/cross_encoder_cache.md b/docs/features/cross_encoder_cache.md
new file mode 100644
index 000000000000..02fdb00f180b
--- /dev/null
+++ b/docs/features/cross_encoder_cache.md
@@ -0,0 +1,109 @@
+# Cross-Encoder Output Reuse
+
+Cross-Encoder output reuse lets an Encoder load image embeddings computed by
+another Encoder from a shared Mooncake Store. Local cache misses check the Store
+before encoding; newly computed outputs are published asynchronously. Both paths
+use the existing `ECMooncakeConnector` P2P delivery to Prefill.
+
+## Requirements
+
+- A working [disaggregated Encoder](disagg_encoder.md) deployment using
+  `ECMooncakeConnector`. See the [Mooncake integration example](../../tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh)
+  for the base E/PD configuration.
+- Model Runner V2 (`VLLM_USE_V2_MODEL_RUNNER=1`), Encoder TP=1, and no dynamic LoRA.
+- `--mm-processor-cache-gb` greater than zero to preserve content identifiers.
+- Mooncake 0.3.12 or later with a running RAM Store. Disk offloading
+  (`enable_offload: true`) is not supported.
+
+Shared reuse supports images with contiguous 2D FP16, BF16 or FP32 encoder outputs.
+Other modalities use normal encoding; non-contiguous outputs are not published.
+
+## Usage
+
+On each Encoder, add `cross_encoder_cache: true` to the existing producer
+`--ec-transfer-config`, preserving its P2P settings:
+
+```json
+{
+  "ec_connector": "ECMooncakeConnector",
+  "ec_role": "ec_producer",
+  "ec_connector_extra_config": {
+    "cross_encoder_cache": true
+  }
+}
+```
+
+Prefill keeps its existing `ECMooncakeConnector` consumer configuration and does
+not need a Store client.
+
+Create a JSON file pointing to an existing Store. For an independently managed
+RAM pool, use `standalone-store` with `global_segment_size` set to zero:
+
+```json
+{
+  "metadata_server": "http://STORE_HOST:2379/metadata",
+  "master_server_address": "STORE_HOST:50051",
+  "protocol": "tcp",
+  "device_name": "",
+  "mode": "standalone-store",
+  "global_segment_size": 0,
+  "local_buffer_size": "4GB"
+}
+```
+
+Replace the addresses with your Store endpoints and set the path on each Encoder
+before starting vLLM:
+
+```bash
+export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
+```
+
+This connects to an existing Store; it does not start the storage services.
+For Store setup and tenant configuration, see the
+[Mooncake Store guide](mooncake_store_connector_usage.md#prerequisites).
+
+## Configuration
+
+These options belong in the producer's `ec_connector_extra_config`:
+
+| Option | Default | Description |
+| --- | --- | --- |
+| `cross_encoder_cache` | `false` | Enable shared output reuse. |
+| `embedding_cache_prefix` | `""` | Namespace prefix for shared embeddings. |
+| `embedding_model_identity` | Configured model path | Override the model field in Store keys. Matching multimodal identifiers are still required. |
+| `store_max_pending_items` | `32` | Maximum pending publications per Encoder. Must be positive. |
+| `store_max_pending_bytes` | `2147483648` (2 GiB) | Maximum retained tensor storage for pending publications per Encoder. Must be positive. |
+| `store_read_buffer_bytes` | `134217728` (128 MiB) | Maximum reusable CPU staging buffer per Encoder; pinned for CUDA outputs. Must be positive. |
+
+Encoders sharing outputs must use the same immutable weights, compatible
+preprocessing and matching multimodal identifiers. Use the same configured model
+path for automatic identifiers: `embedding_model_identity` alone does not enable
+reuse across different paths. Change `embedding_cache_prefix` when replacing
+weights in place or changing output-affecting settings. Caller-provided UUIDs
+must consistently identify the same input; see [cached inputs](multimodal_inputs.md#cached-inputs).
+
+## Limitations
+
+- Store reads are synchronous. Reuse skips Encoder computation, but not
+  preprocessing, P2P delivery or the scheduler's Encoder budget reservation.
+- Hits are read in batches through one lazily allocated, registered CPU staging
+  buffer. Each chunk is copied to independent output tensors before the buffer
+  is reused. Reads are split by `store_read_buffer_bytes`; an individual object
+  larger than this capacity (including its 24-byte header) falls back to encoding.
+  Staging memory is additional to the Store's `local_buffer_size` and is released
+  on healthy shutdown. The first load includes allocation and registration costs.
+- Publication is best-effort. A request can finish before its outputs reach the
+  Store, so concurrent cold requests may still encode the same image.
+- Publication budgets count each retained backing storage once across pending
+  views. Its charge is released after the last view is safely reclaimed.
+  Exceeding a budget skips the write. Store client buffers and the P2P pool
+  consume additional memory.
+- Cache misses and recoverable read errors fall back to encoding. Recoverable
+  publication errors skip the write. Unexpected native errors or unconfirmed
+  I/O completion or buffer release can fail the worker.
+- Incompatible objects are rejected without replacement; they can trigger
+  repeated fallback until evicted.
+- An independent Store can retain embeddings across Encoder restarts. Capacity,
+  eviction and resilience to Store failures depend on the Store deployment.
+- The `protocol:v3` namespace isolates the compact embedding format from older
+  cache objects. Encoders using different format versions do not share hits.
diff --git a/docs/features/custom_logitsprocs.md b/docs/features/custom_logitsprocs.md
index 232f4363efee..8c1c25939928 100644
--- a/docs/features/custom_logitsprocs.md
+++ b/docs/features/custom_logitsprocs.md
@@ -8,84 +8,40 @@ A "custom" logits processor is written by a user of vLLM and is loaded into vLLM
 
 This document shows how to write, load and use a custom logits processor.
 
+!!! note
+    This document covers the Model Runner V2 (MRV2) interface. The legacy V1
+    model runner (`VLLM_USE_V2_MODEL_RUNNER=0`) uses a different interface,
+    `vllm.v1.sample.logits_processor.LogitsProcessor`, which is not covered
+    here.
+
 ## Logits Processors Background
 
 A logits processor adjusts the next-token probability distribution, usually with the intention of steering the model towards a desired type of behavior.
 
-In vLLM, logits processors operate at batch granularity. During a given engine step, the logits processor consumes a `(num_requests) x (vocab_size)` tensor of raw logits output by the model. For all requests which enable the logits processor, the logits processor applies a transformation to the corresponding row of the logits tensor, while leaving other rows unmodified. The transformed logits tensor is then passed to softmax.  
+In vLLM, logits processors operate at batch granularity. During a given engine step, the logits processor consumes a `(num_logits_rows) x (vocab_size)` tensor of raw logits output by the model. Note that a logits row is not a request: rows are reordered every step, and under speculative decoding a request owns one row per draft token. The logits processor applies a transformation to the rows of the logits tensor, while leaving other rows unmodified. The transformed logits tensor is then used for sampling.
 
 ## Creating a Custom Logits Processor
 
-Custom logits processors must subclass `vllm.v1.sample.logits_processor.LogitsProcessor` and define (at minimum) the following methods:
+Custom logits processors must subclass `vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor` and define (at minimum) the following methods:
 
-* `validate_params(cls, sampling_params: SamplingParams)`:
-    * Raise `ValueError` if `SamplingParams` has invalid arguments (especially custom arguments) used by logits processor.
-    * When request is sent to entrypoint, `validate_params()` will validate `SamplingParams` and refuse request with invalid arguments.
-    * **Note:** it's important to implement `validate_params()` to prevent invalid parameters for custom logits processor. Otherwise requests with invalid parameters can cause unexpected behaviour in custom logits processor.
-
-* `__init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool)`
+* `__init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState)`:
     * `vllm_config`: engine configuration data structure
-    * `device`: hardware accelerator device info
-    * `is_pin_memory`: flag indicating whether pin memory is available to support logits processor implementation
-
-* `apply(self, logits: torch.Tensor) -> torch.Tensor`:
-    * Consume a `(num_requests) x (vocab_size)` logits tensor (`logits`)
-    * Apply logits processor transformation at batch granularity
-    * Return a transformed `(num_requests) x (vocab_size)` logits tensor
-    * You can modify the input logits processors in-place or out-of-place; in-place is more memory-efficient
-
-* `is_argmax_invariant(self) -> bool`:
-    * Return `True` if the logits processor is argmax invariant (never changes what is the highest-logit-value token ID for a given request), `False` if the logits processor may modify argmax
-    * `is_argmax_invariant()` is evaluated once at startup; if `True`, vLLM will skip applying this logits processor in a given step when all requests use greedy sampling
-
-* `update_state(self, batch_update: Optional["BatchUpdate"]) -> None`:
-    * Consume a `BatchUpdate` data structure representing persistent batch state changes at the beginning of the current engine step
-    * Use the `BatchUpdate` members to update logits processor internal state
-    * **Note:** batch update data structure may be `None`, signaling no change to the batch constituents. In this case, the LogitsProcessor might still want to update its state based on the updated `output_token_ids` lists that it could have retained when they were added.
-
-### How the vLLM engine builds the `BatchUpdate` data structure
-
-!!! important
-    Some logits processors design changes are still in progress. We expect
-    that in the future you will not need to account for batch state changes
-    when implementing a logits processor, and the information in this section
-    will become irrelevant.
-
-Logits processor `update_state()` implementations should assume the following model for how the model runner updates persistent batch state (expressed here in terms of the `BatchUpdate` abstraction):
-
-1. Identify indices of requests which finished in the current engine step
-
-2. Identify new requests introduced in the current step
-
-3. Use Add operations to replace as many finished requests with new requests, in order of increasing index of the replaced request starting with the lowest index
-
-4. Based on the relative number of new and finished requests:
-
-    1. If the numbers of new and finished requests are the same, proceed to next step
+    * `req_states`: a narrow, read-only view of the persistent batch, exposing the on-device token history (`all_token_ids`, `prompt_len`, `prefill_len`, `total_len`) plus `device`, `max_num_reqs` and `vocab_size`
 
-    2. *If there are more new requests than finished requests:* apply Add operations to extend the batch with the remaining new requests which did not replace finished requests. Assign consecutive indices to these new requests, starting with `current_max_batch_index + 1`
+* `add_request(self, req_idx, sampling_params) -> bool`:
+    * Initialize per-slot state for a request entering the batch. Slots are recycled through a free list, so per-slot state must be fully (re)initialized here; there is no removal hook, since freed slots are never read
+    * Return whether this processor modifies logits for the request. The sampler skips the logits-processing pipeline for batches in which no request needs it
 
-    3. *If there are fewer new requests than finished requests:*
+* `apply_staged_writes(self) -> None` (optional):
+    * Flush host-side writes staged by `add_request()` to the device; called once per step before the forward pass
 
-        * Apply Remove operations to finished requests which were not replaced with new requests. These removed request indices will necessarily be greater than the greatest index of the finished requests which were replaced in the previous step. The Removes may leave the batch in a non-contiguous state
+* `validate_params(cls, sampling_params) -> None` (optional classmethod):
+    * Raise `ValueError` for invalid per-request arguments (especially custom arguments); runs at request admission, so invalid arguments fail the request instead of reaching the sampler
 
-        * **"Condense" the batch to be contiguous:** starting with the lowest-index empty slot (which was caused by a Remove), apply a Unidirectional Move from the current highest non-empty slot in the batch to fill the empty slot. Proceed with additional Unidirectional Move operations in order of increasing empty slot destination index and decreasing non-empty slot source index until the batch is contiguous
-
-        * **Shrink the batch:** a side effect of condensing the batch is that empty slots resulting from Remove operations are grouped in a contiguous block at the end of the batch array. Thus, after condensing, update `BatchUpdate.batch_size` to reflect the number of non-empty slots
-
-5. Reorder the batch for improved efficiency. Depending on the attention backend implementation and the current characteristics of the batch, zero or more Swap Move operations may be applied to reorder the batch
-
-Notes:
-
-* A logits processor `update_state()` method must process batch update operations in the following order: removes, adds, moves
-
-* The index argument for Add operations refers to the index *at the time the Add occurred*, i.e. before any Move operations
-    * Example: if a request is Added at index 5 and then swapped with index 3, the Add operation in `BatchUpdate.added` will be associated with index 5 not 3
-    * In other words Move operations can be assumed to be applied after Adds and Removes
-
-* Move operations can be assumed to be applied in the order in which they appear in `BatchUpdate.moved`
-
-* If there are no new/finished requests and there is no batch reordering, then the batch update for the logits processors will be `None`
+* `apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor`:
+    * Consume a `(num_logits_rows) x (vocab_size)` logits tensor and the step's batch layout (`ctx`: row-to-slot mappings, `input_ids`, positions)
+    * Apply logits processor transformation at batch granularity
+    * Return a transformed logits tensor. You can modify the input logits in-place or out-of-place; in-place is more memory-efficient
 
 ### Passing Custom Argument to a Custom Logits Processor
 
@@ -93,7 +49,7 @@ Unlike built-in logits processors, custom logits processors may require configur
 
 ### Example Custom Logits Processor Implementation
 
-The contrived example below implements a custom logits processor which consumes a `(num\_requests) \times (vocab\_size)` logits tensor and masks out all tokens except for one (`target_token`) with `float(-inf)`. The logits processor is disabled for any request that does not specify `target_token`. To determine whether the logits processor is enabled and which token to leave unmasked, the logits processor checks `SamplingParams.extra_args` for a `target_token` custom argument associated with each request:
+The contrived example below implements a custom logits processor which masks out all tokens except for one (`target_token`) for the requests that enable it. The processor is disabled for any request that does not specify `target_token`; `add_request()` reports this via its return value, so the sampler can skip the logits-processing pipeline entirely when no request enables the processor. To determine whether the logits processor is enabled and which token to leave unmasked, the processor checks `SamplingParams.extra_args` for a `target_token` custom argument associated with each request:
 
 ??? code "Example custom logits processor definition"
 
@@ -101,193 +57,55 @@ The contrived example below implements a custom logits processor which consumes
     import torch
     from vllm.config import VllmConfig
     from vllm.sampling_params import SamplingParams
-    from vllm.v1.sample.logits_processor import (BatchUpdate,
-                                                LogitsProcessor,
-                                                MoveDirectionality)
+    from vllm.v1.worker.gpu.sample.logits_processor import (
+        LogitsProcRequestState,
+        LogitsContext,
+        LogitsProcessor,
+    )
 
-    class DummyLogitsProcessor(LogitsProcessor):
-        """Fake logit processor to support unit testing and examples"""
 
-        @classmethod
-        def validate_params(cls, params: SamplingParams):
-            target_token: int | None = params.extra_args and params.extra_args.get(
-                "target_token"
-            )
-            if target_token is not None and not isinstance(target_token, int):
-                raise ValueError(f"target_token value {target_token} is not int")
-
-        def __init__(self, vllm_config: "VllmConfig", device: torch.device,
-                    is_pin_memory: bool):
-            self.req_info: dict[int, int] = {}
-
-        def is_argmax_invariant(self) -> bool:
-            """Never impacts greedy sampling"""
-            return False
-
-        def update_state(self, batch_update: BatchUpdate | None):
-            if not batch_update:
-                return
-
-            # Process added requests.
-            for index, params, _, _ in batch_update.added:
-                assert params is not None
-                self.validate_params(params)
-                if params.extra_args and (target_token :=
-                                        params.extra_args.get("target_token")):
-                    self.req_info[index] = target_token
-                else: 
-                    self.req_info.pop(index, None)
-
-            if self.req_info:
-                # Process removed requests.
-                for index in batch_update.removed:
-                    self.req_info.pop(index, None)
-
-                # Process moved requests, unidirectional move (a->b) and swap
-                # (a<->b)
-                for adx, bdx, direct in batch_update.moved:
-                    a_val = self.req_info.pop(adx, None)
-                    b_val = self.req_info.pop(bdx, None)
-                    if a_val is not None:
-                        self.req_info[bdx] = a_val
-                    if direct == MoveDirectionality.SWAP and b_val is not None:
-                        self.req_info[adx] = b_val
-
-        def apply(self, logits: torch.Tensor) -> torch.Tensor:
-            if not self.req_info:
-                return logits
+    class TargetTokenLogitsProcessor(LogitsProcessor):
+        """Masks out all tokens except `target_token` (a per-request custom
+        argument); requests without it are left alone."""
 
-            # Save target values before modification
-            cols = torch.tensor(
-                list(self.req_info.values()), dtype=torch.long, device=logits.device
+        def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
+            # Per-slot target; -1 means disabled. Staged on the host and
+            # flushed to the device once per step in apply_staged_writes().
+            self.target_token = torch.full(
+                (req_states.max_num_reqs,), -1, dtype=torch.int64
             )
-            rows = torch.tensor(
-                list(self.req_info.keys()), dtype=torch.long, device=logits.device
+            self.target_token_dev = torch.full(
+                (req_states.max_num_reqs,), -1, dtype=torch.int64, device=req_states.device
             )
-            values_to_keep = logits[rows, cols].clone()
-
-            # Mask all but target tokens
-            logits[rows] = float('-inf')
-            logits[rows, cols] = values_to_keep
-
-            return logits
-
-    ```
 
-In the rest of this document, we will use `DummyLogitsProcessor` as an example of a custom logits processor.
+        def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+            target = (sampling_params.extra_args or {}).get("target_token")
+            self.target_token[req_idx] = target if target is not None else -1
+            return target is not None
 
-The `DummyLogitsProcessor.update_state()` implementation maintains a "sparse" representation of the batched requests in the `self.req_info` dictionary: only those requests which specify a `target_token` value have a key in the dictionary. `update_state()` adjusts the stored request indices and `target_token` values (keys and values respectively in `self.req_info`) in response to Add, Remove and Move operations against the persistent batch.
+        def apply_staged_writes(self) -> None:
+            self.target_token_dev.copy_(self.target_token, non_blocking=True)
 
-### Wrapping an Existing Request-Level Logits Processor
-
-Although the vLLM engine applies logits processors at batch granularity, some users may want to use vLLM with a "request-level" logits processor implementation - an implementation which operates on individual requests. This will be especially true if your logits processor was developed for vLLM version 0, which required it to be a `Callable` (as described [here][vllm.logits_process]) conforming to the following type annotation:
-
-``` python
-RequestLogitsProcessor = Union[
-
-    # (output token ids, logits tensor) -> logits tensor
-    Callable[[list[int], Tensor], Tensor],
-
-    # (prompt token ids, output token ids, logits tensor) -> logits tensor
-    Callable[[list[int], list[int], Tensor], Tensor],
-]
-```
-
-While request-level logits processors are explicitly *not* supported in the vLLM engine, vLLM *does* provide a convenient process to wrap an existing `Callable` request-level logits processor and create a batch-level logits processor that is compatible with vLLM. The `Callable` must conform to the type annotation above; if your request-level logits processor has a different interface, then in order to wrap it, you may need to modify it or implement an additional wrapper layer to comply with the interface specification above.
-
-You can wrap the request-level logits processor by subclassing `AdapterLogitsProcessor` as shown in the example below (in this example, `DummyPerReqLogitsProcessor` is a stand-in for your request-level logits processor which needs to be wrapped.):
-
-* Override `AdapterLogitsProcessor.validate_params(cls,params)` to validate request's sampling parameters.
-
-* Override `AdapterLogitsProcessor.is_argmax_invariant(self)` to accurately reflect whether your request-level logits processor may impact which token has the highest-value logit.
-
-* Override `AdapterLogitsProcessor.new_req_logits_processor(self,params)` to create a new request-level logits processor instance from a `SamplingParams` instance:
-
-??? code "Example of Wrapping a Request-Level Logits Processor"
-
-    ``` python
-    ...
-
-    from vllm.v1.sample.logits_processor import (
-        AdapterLogitsProcessor, # Wrapper base-class
-        RequestLogitsProcessor, # Request-level logitsproc type annotation
-    )
-
-    ...
-
-    # Stand-in for your request-level logits processor:
-    class DummyPerReqLogitsProcessor:
-        """The request-level logits processor masks out all logits except the
-        token id identified by `target_token`"""
-
-        def __init__(self, target_token: int) -> None:
-            """Specify `target_token`"""
-            self.target_token = target_token
-
-        def __call__(
-            self,
-            output_ids: list[int],
-            logits: torch.Tensor,
-        ) -> torch.Tensor:
-            val_to_keep = logits[self.target_token].item()
-            logits[:] = float("-inf")
-            logits[self.target_token] = val_to_keep
+        def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+            # ctx.expanded_idx_mapping maps each logits row to its request slot.
+            cols = self.target_token_dev[ctx.expanded_idx_mapping.long()]
+            rows = torch.nonzero(cols >= 0).squeeze(1)
+            if rows.numel() == 0:
+                return logits
+            kept = logits[rows, cols[rows]].clone()
+            logits[rows] = float("-inf")
+            logits[rows, cols[rows]] = kept
             return logits
-
-    ...
-
-    # Example of wrapping the request-level logits processor:
-    class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
-        """Example of wrapping a fake request-level logit processor to create a
-        batch-level logits processor"""
-
-        @classmethod
-        def validate_params(cls, params: SamplingParams):
-            target_token: Any | None = params.extra_args and params.extra_args.get(
-                "target_token"
-            )
-            if target_token is not None and not isinstance(target_token, int):
-                raise ValueError(
-                    f"target_token value {target_token} is not int"
-                )
-
-        def is_argmax_invariant(self) -> bool:
-            return False
-
-        def new_req_logits_processor(
-            self,
-            params: SamplingParams,
-        ) -> Optional[RequestLogitsProcessor]:
-            """This method returns a new request-level logits processor, customized
-            to the `target_token` value associated with a particular request.
-
-            Returns None if the logits processor should not be applied to the
-            particular request. To use the logits processor the request must have
-            a "target_token" custom argument with an integer value.
-
-            Args:
-            params: per-request sampling params
-
-            Returns:
-            `Callable` request logits processor, or None
-            """
-            target_token: Any | None = params.extra_args and params.extra_args.get(
-                "target_token"
-            )
-            if target_token is None:
-                return None
-            return DummyPerReqLogitsProcessor(target_token)
     ```
 
-!!! note
-    Your `new_req_logits_processor()` override can return `None` to signal that the wrapped logits processor should not be applied to the request in question.
-
-Once you have created a custom subclass (like `WrappedPerReqLogitsProcessor`) which wraps your request level logits processor, you can pass the custom subclass to vLLM via any of the methods described in the following section.
+Per-request state is keyed by the request slot index, and slots are recycled through a free list. The example keeps a per-slot state where `-1` means disabled; note that `add_request()` overwrites the slot's entry unconditionally, which is what keeps recycled slots from leaking state between requests.
 
 ## Ways to Load Your Custom Logits Processor in vLLM
 
 Logits processors are loaded at initialization. Critically, the set of loaded logits processors cannot be modified after the vLLM engine finishes loading, and new logits processors cannot be loaded on-demand for individual requests.
 
+Loaded classes are validated against the Model Runner V2 interface; a class that does not subclass `vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor` is rejected with a clear error.
+
 This section details different ways of making your logits processor visible to vLLM and triggering vLLM to load your logits processor.
 
 ### Method 1: Pass the Custom Logits Processor Fully-Qualified Class Name (FQCN) to vLLM at Initialization Time
@@ -314,7 +132,7 @@ See examples below:
     # Pass in FQCN
     llm = LLM(
         model="facebook/opt-125m",
-        logits_processors=["your.module.path:DummyLogitsProcessor"],
+        logits_processors=["your.module.path:TargetTokenLogitsProcessor"],
     )
     ```
 
@@ -323,14 +141,14 @@ See examples below:
     ``` python
     # Pass in FQCN
     engine_args = AsyncEngineArgs(model="facebook/opt-125m",
-                                  logits_processors=["your.module.path:DummyLogitsProcessor"])
+                                  logits_processors=["your.module.path:TargetTokenLogitsProcessor"])
     async_llm = AsyncLLM.from_engine_args(engine_args)
     ```
 
 ??? code "Passing custom logits processor FQCN to vLLM server via CLI"
 
     ```bash
-    vllm serve facebook/opt-125m --logits_processors your.module.path:DummyLogitsProcessor
+    vllm serve facebook/opt-125m --logits_processors your.module.path:TargetTokenLogitsProcessor
     ```
 
 ### Method 2: Automatically Detect Custom Logits Processors Installed in Your Python Environment As Entry Points
@@ -345,7 +163,7 @@ Suppose that you have developed a Python package that holds your custom logits p
 
     ``` toml
     [project.entry-points."vllm.logits_processors"]
-    dummy_logits_processor = "your.module.path:DummyLogitsProcessor"
+    target_token_logits_processor = "your.module.path:TargetTokenLogitsProcessor"
     ```
 
 Once your package is installed, your custom logits processor will be loaded automatically whenever vLLM is initialized. You do *not* need to pass the custom logits processor to the `LLM` or `AsyncLLM` constructors or to the vLLM server explicitly at initialization time if your logits processor is exposed as an entry point.
@@ -361,26 +179,26 @@ You can pass one or more custom logits processor class objects to the `LLM` and
 
     ``` python
     # Import custom logits processor
-    from some.module import DummyLogitsProcessor
+    from some.module import TargetTokenLogitsProcessor
 
     # ...or...
 
     # Define custom logits processor locally
-    from vllm.v1.sample.logits_processor import LogitsProcessor
+    from vllm.v1.worker.gpu.sample.logits_processor import LogitsProcessor
 
-    class DummyLogitsProcessor(LogitsProcessor):
-        # See DummyLogitsProcessor implementation above
+    class TargetTokenLogitsProcessor(LogitsProcessor):
+        # See TargetTokenLogitsProcessor implementation above
         ...
 
     # Pass class object to LLM constructor
     llm = LLM(
         model="facebook/opt-125m",
-        logits_processors=[DummyLogitsProcessor],
+        logits_processors=[TargetTokenLogitsProcessor],
     )
 
     # Pass class object to AsyncLLM constructor
     engine_args = AsyncEngineArgs(model="facebook/opt-125m",
-                                  logits_processors=[DummyLogitsProcessor])
+                                  logits_processors=[TargetTokenLogitsProcessor])
     async_llm = AsyncLLM.from_engine_args(engine_args)
     ```
 
@@ -388,7 +206,7 @@ You can pass one or more custom logits processor class objects to the `LLM` and
 
 The design of the custom logits processor determines whether the logits processor must be enabled/disabled for a given request, and what arguments must be provided to configure the logits processor.
 
-The examples below show how a user would pass a custom argument (`target_token`) to `DummyLogitsProcessor` in order to (1) enable the logits processor for that particular request and (2) control the logits processor's behavior.
+The examples below show how a user would pass a custom argument (`target_token`) to `TargetTokenLogitsProcessor` in order to (1) enable the logits processor for that particular request and (2) control the logits processor's behavior.
 
 ??? code "vLLM REST API: configure custom logits processor for a request"
 
@@ -419,7 +237,7 @@ The examples below show how a user would pass a custom argument (`target_token`)
 ??? code "Offline: configure custom logits processor for an `LLM` request"
 
     ``` python
-    outputs_logitproc = llm.generate("your prompt", 
+    outputs_logitproc = llm.generate("your prompt",
                                      SamplingParams(...,
                                         extra_args={"target_token": 67}))
     ```
@@ -438,31 +256,14 @@ The examples below show how a user would pass a custom argument (`target_token`)
 
 ## Best Practices for Writing Custom Logits Processors
 
-Once vLLM loads a logits processor during initialization, then vLLM will invoke `update_state()` and `apply()` against that logits processor in every engine step. Both methods operate on all requests which currently reside in the vLLM persistent batch. Thus, it is important to implement these methods efficiently.
-
-* Write efficient `apply()` and `update_state()` implementations in light of the fact that logits processors operate at batch granularity
-    * For example, you may be able to use efficient vectorized operations to implement `apply()` or update internal state vectors in `update_state()`
-    * However, if you think that a logits processor may be used infrequently, it may be appropriate to use a "sparse" representation of request state i.e. the class can represent request configuration using a dictionary which only stores metadata about requests that enable the logits processor
-    * **Note:** wrapped request-level logits processors do not need to implement `apply()` and `update_state()`; the default `AdapterLogitsProcessor.update_state()` implementation maintains a sparse representation of request state, wherein requests for which `new_req_logits_processor()` returns `None` are not represented in the base-class state dictionary. The default implementation of `AdapterLogitsProcessor.apply()` applies the request-level logits processor to each row of input logits sequentially and assembles the output logits tensor. If the performance of this `AdapterLogitsProcessor` default implementation is insufficient, then avoid wrapping your request-level logits processor and instead re-implement it as a `LogitsProcessor` subclass with optimized `apply()` and `update_state()` implementations that operate at batch granularity
+* Write an efficient `apply()` implementation in light of the fact that logits processors operate at batch granularity. For example, you may be able to use efficient vectorized operations to implement `apply()`, and to stage per-request state into contiguous tensors in `add_request()` instead of per-row Python loops
 
 * It is up to the logits processor author to determine:
 
-    1. **The per-request attributes which configure the logits processor's behavior against that request.** Your custom logits processor's `update_state()` override determines how `SamplingParams` fields are mapped into logits processor state
-
-        * **Note:** for wrapped request-level logits processors, `new_req_logits_processor()` determines how `SamplingParams` fields are used to initialize a request-level logits processor instance.
-
-    2. **The conditions under which the logits processor is or is not enabled on a per-request basis.** Unless your intention is for the custom logits processor to act on all requests all the time, you should write your logits processor in such a way that it is possible to disable the logits processor for a given request, i.e. by defaulting an argument to `None` or by passing in a specific do-nothing argument value i.e. `0.0`. Try to save compute and memory for requests which disable the logits processor
-
-        * **Note:** for wrapped per-request logits processors, the default `AdapterLogitsProcessor.update_state()` implementation ensures that the request-level logits processor is disabled when `new_req_logits_processor()` returns `None` for that request
-
-    3. **The conditions under which the logits processor is short-circuited at the batch level.** Even if you have defined a way to disable the custom logits processor at the request level, it may be difficult to translate this into compute savings i.e. if your `update_state()` and `apply()` implementations use efficient vectorized implementations that operate on the whole persistent batch in a single command. For example, you cannot skip an entire vectorized operation in `apply()` just because one request disabled the logits processor. To save compute in the edge-case where no running requests utilize the custom logits processor, we recommend designing `apply()` to return the unmodified input tensor if all requests have the logits processor disabled. Similarly, consider whether steps can be skipped in `update_state()` if no requests enable the logits processor
-
-        * Additionally, an easy way to save compute in `update_state()` is to exit early when the `batch_update` is `None`
-
-        * **Note:** for wrapped per-request logits processors, the `AdapterLogitsProcessor` base-class implements the above optimizations by default
+    1. **The per-request attributes which configure the logits processor's behavior against that request.** Your custom logits processor's `add_request()` override determines how `SamplingParams` fields are mapped into logits processor state
 
-* Ensure that the logits processor `update_state` method discards information about finished requests (i.e. requests which are replaced by an Add or which are subject to a Remove)
+    2. **The conditions under which the logits processor is or is not enabled on a per-request basis.** Unless your intention is for the custom logits processor to act on all requests all the time, you should write your logits processor in such a way that it is possible to disable the logits processor for a given request, i.e. by defaulting an argument to `None` or by passing in a specific do-nothing argument value i.e. `0.0`. Return `False` from `add_request()` for such requests so that the sampler can skip the logits-processing pipeline for batches in which no request needs it
 
-    * **Note:** for wrapped per-request logits processors, the `AdapterLogitsProcessor` base-class handles this by default
+* Stage host-side writes in `add_request()` and flush them once per step in `apply_staged_writes()`, instead of writing device tensors directly; this is the discipline every built-in sampler state follows
 
-* `is_argmax_invariant()` can be hard-coded to `True` or `False` if the logits processor has consistent behavior. However, the argmax invariance may also be determined programmatically (i.e. if your logits processor is user-customizable in some way that impacts whether the logits processor is argmax invariant). For this reason, `is_argmax_invariant()` is not a class method
+* Since there is no removal hook, `add_request()` must fully (re)initialize the slot's state; do not rely on state left behind by the slot's previous occupant
diff --git a/docs/features/disagg_encoder.md b/docs/features/disagg_encoder.md
index c27d6b277284..ac304345fd1e 100644
--- a/docs/features/disagg_encoder.md
+++ b/docs/features/disagg_encoder.md
@@ -32,8 +32,9 @@ Design doc: <https://docs.google.com/document/d/1aed8KtC6XkXtdoV87pWT0a8OJlZ-Cpn
 
 ## 2  Usage Example
 
-The current reference pathway is **ExampleConnector**.  
-Below ready-to-run scripts shows the workflow:
+### ExampleConnector
+
+The following scripts demonstrate the workflow with **ExampleConnector**:
 
 1 Encoder instance + 1 PD instance:
 `examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh`
@@ -41,10 +42,45 @@ Below ready-to-run scripts shows the workflow:
 1 Encoder instance + 1 Prefill instance + 1 Decode instance:
 `examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh`
 
+### ECMooncakeConnector
+
+**ECMooncakeConnector** transfers encoder outputs using the Mooncake TransferEngine.
+See the [Mooncake integration example](../../tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh)
+for a complete 1 Encoder instance + 1 PD instance setup, including the producer and
+consumer `--ec-transfer-config` settings and proxy configuration.
+
+With vLLM and Mooncake installed, run the example from the repository root:
+
+```bash
+GPU_E=0 GPU_PD=1 MOONCAKE_EC_PROTOCOL=tcp \
+    bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
+```
+
+The script uses `Qwen/Qwen2.5-VL-3B-Instruct` by default, runs a single-GPU baseline,
+and compares the disaggregated outputs against it. Set `MODEL` to use another model
+or `MOONCAKE_EC_PROTOCOL=rdma` to use RDMA on supported hardware. See the script for
+additional configuration options.
+
+### Audio inputs
+
+The Python EPD proxy also rewrites audio in `/v1/chat/completions` into
+metadata-only references for Qwen2-Audio, AudioFlamingo3, Ultravox,
+Qwen2.5-Omni, and Qwen3-Omni. The encoder publishes `audio_num_tokens`
+(the placeholder token count of each audio) together with any feature
+lengths the model needs; the consumer uses them to reconstruct the audio
+placeholders and receives the embeddings through its EC connector without
+repeating audio preprocessing.
+
+This covers pure audio inputs, not video with an embedded audio track. It does
+not add `/v1/audio/transcriptions` or realtime routes to the example proxy.
+
 ---
 
 ## 3  Test Script
 
+For optional cross-Encoder output reuse while retaining Mooncake P2P delivery,
+see [Cross-encoder output reuse](cross_encoder_cache.md).
+
 Please refer to the directories `tests/v1/ec_connector`
 
 ## 4  Development
diff --git a/docs/features/engram.md b/docs/features/engram.md
new file mode 100644
index 000000000000..73dc05e9f144
--- /dev/null
+++ b/docs/features/engram.md
@@ -0,0 +1,148 @@
+# Engram: conditional memory via n-gram lookups
+
+Engram is a *conditional memory* module attached to specific layers of a
+transformer backbone. At each position, the preceding few tokens are hashed
+into large static embedding tables, and the retrieved rows are fused into the
+hidden state, gated by how well they match it. Where Mixture-of-Experts scales
+capacity through conditional computation, Engram scales it through conditional
+memory: parameters that are never computed over, only looked up, so capacity
+grows at O(1) per-token cost. The mechanism and its design rationale are
+described in [Conditional Memory via Scalable Lookup: A New Axis of Sparsity
+for Large Language Models](https://arxiv.org/abs/2601.07372).
+
+vLLM enables Engram automatically whenever the checkpoint declares Engram
+layers, so basic serving needs no extra flags. This page explains what Engram
+does and how to control its memory footprint and multi-GPU topology.
+
+## Supported models
+
+Engram layers are declared by the checkpoint, not chosen at serve time:
+
+| Architecture | HF config field | Notes |
+| --- | --- | --- |
+| `DeepseekV41ForCausalLM` | `engram_layer_ids` | e.g. `deepseek-ai/DeepSeek-V4.1-Flash` |
+| `Qwen4ExpForCausalLM` / `Qwen4ExpForConditionalGeneration` | `ple_layer_ids` | same mechanism, called PLE |
+
+Passing `--engram-config` for a model without n-gram embedding layers fails at
+startup, as does running on a non-CUDA-alike platform.
+
+## How it works
+
+Engram processes each position in two phases: *retrieval* and *fusion*.
+
+**Retrieval** maps the local token context to static table rows:
+
+1. **Tokenizer compression.** At startup, vLLM builds a many-to-one map from
+   the tokenizer that collapses tokens which normalize alike (NFKC, accent
+   stripping, lowercasing, whitespace collapsing), so `" The"`, `"the"` and
+   `"THE"` hash identically. This concentrates semantic density into a much
+   smaller id space; its size must match `engram_compressed_vocab_size` from
+   the checkpoint, or startup fails — every hash multiplier derives from it.
+2. **Multi-head hashing.** For each n-gram order `n` (up to
+   `engram_max_ngram_size - 1`), several hash heads each apply an independent
+   multiplicative-XOR hash over the compressed suffix n-gram. Every (n-gram
+   order, head) pair indexes its own prime-sized bucket range in the layer's
+   table — disjoint across pairs, so collisions cannot collide twice — and
+   the retrieved rows are concatenated into one memory vector per position.
+
+**Fusion** turns the retrieved rows into a context-aware contribution to the
+hidden state. The rows are context-independent priors and may be noisy (hash
+collisions, polysemy), so they are gated: a `wkv` projection produces one key
+per hyper-connection branch plus a shared value, and each branch computes a
+sigmoid gate over the normalized dot product between its hidden state —
+which has already aggregated global context through preceding attention —
+and its key. If the retrieved memory contradicts the current context, the
+gate tends toward zero and suppresses it. The gated value is added
+residually into the branch's stream (`hidden + gate * value`); positions that
+take no part in an n-gram (e.g. image spans) are masked so they pass through
+untouched.
+
+Engram sits only on the layers listed in the checkpoint. Placement is a
+deliberate hardware-algorithm trade-off: early enough that the backbone is
+relieved from reconstructing static patterns in its first layers, deep enough
+that vLLM can hide table lookups behind the compute of the preceding layers
+(see below).
+
+In vLLM, hashing, lookup, and gating run as fused Triton kernels over FP8
+rows with ue8m0 block scales. The layer layout — table sizes, n-gram size,
+head count — is fixed by the checkpoint and is not user-configurable. Only the
+storage and sharding described below are configurable, via
+`--engram-config`, which takes a JSON object and can also be set field by
+field as `--engram-config.<field>`
+([CLI reference](../cli/serve.md#-engram-config)).
+
+## CPU offload
+
+Engram tables dwarf typical embedding weights: [DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash)
+carries two Engram layers of roughly 384M table rows each, with 256-byte FP8
+entries plus their ue8m0 block scales — about 200 GB of table weights alone.
+Because lookups are deterministic — the indices are known from the token ids
+alone, before the forward pass — the tables need not sit in GPU memory at
+all. By default vLLM stores them in **pinned host memory** and serves lookups
+directly from the Triton kernel through unified virtual addressing (UVA),
+prefetching on a side CUDA stream so transfers overlap with the compute of
+preceding layers. The freed GPU memory goes to the KV cache.
+
+`cpu_offload` defaults to `true` (on): it follows the
+`VLLM_PLE_CPU_OFFLOAD` environment variable, which also defaults to `true`.
+An explicit `--engram-config` value takes precedence over the environment
+variable.
+
+To keep the tables resident on GPU instead (e.g. when host memory is
+scarce):
+
+```bash
+vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
+  --engram-config.cpu_offload false
+```
+
+or equivalently `VLLM_PLE_CPU_OFFLOAD=0`. CPU offload requires a GPU with UVA
+support; vLLM fails fast when it is unavailable.
+
+## Data-parallel topologies (DeepSeek V4.1)
+
+By default (`embedding_across_dp: false`), each DP replica keeps its own
+TP-sharded copy of the tables — no cross-replica communication, but one host
+copy per replica.
+
+With `embedding_across_dp: true`, the hash heads are sharded across all
+TP × DP ranks into a single table copy. Each step gathers the hash ids of
+every co-located DP replica, each rank looks up the heads it owns, and the
+rows are exchanged back. This trades per-step DP collectives for a much
+smaller table footprint:
+
+```bash
+vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
+  --tensor-parallel-size 2 --data-parallel-size 4 \
+  --engram-config '{"embedding_across_dp": true}'
+```
+
+`embedding_across_dp` is not supported with elastic expert parallelism yet.
+Qwen4Exp PLE tables are ETP-sharded instead and honor `cpu_offload` only.
+
+## Sharing host tables across DP replicas
+
+When the tables are CPU-offloaded and DP replicas are co-located on one node,
+`dp_shared_memory` stores **one copy of each TP shard in `/dev/shm`**, shared
+by all co-located replicas through `cudaHostRegister`. Each replica then
+prefetches only its own tokens — no per-step Engram DP collectives — and host
+memory drops from one table copy per replica to one per node.
+
+`dp_shared_memory` defaults to enabled whenever the other settings allow it,
+and falls back to per-replica (or DP-sharded) tables with a warning when:
+
+- the DP replicas are not co-located on a single node,
+- `/dev/shm` cannot hold the full tables — in containers, raise the limit
+  with `--shm-size` or `--ipc=host`,
+- expert parallelism is elastic, which is unsupported.
+
+Requirements: `cpu_offload` enabled, `--data-parallel-size > 1`, and a shared
+IPC namespace across the co-located replicas.
+
+## Limitations
+
+- CUDA-alike platforms only.
+- DeepSeek V4.1 Engram does not support DBO or microbatching; disable
+  `--enable-dbo` and set `--ubatch-size 0`.
+- `dp_shared_memory` requires `cpu_offload`, `data_parallel_size > 1`, and is
+  unsupported with elastic expert parallelism.
diff --git a/docs/features/initialized_snapshots.md b/docs/features/initialized_snapshots.md
new file mode 100644
index 000000000000..538695d0a8be
--- /dev/null
+++ b/docs/features/initialized_snapshots.md
@@ -0,0 +1,152 @@
+# Initialized engine snapshots
+
+Initialized engine snapshots are an experimental way to trade local disk space
+and host privileges for a faster vLLM activation. Snapshot creation initializes
+the engine and records deterministic generation output before CRIU captures the
+process tree and CUDA state. Restore validates the saved environment, restores
+the engine, binds the HTTP server, and reproduces the recorded token and sampled
+token log probability before returning.
+
+This path is intended for repeatedly activating the same model and engine
+configuration on the same machine. It is not a portable model artifact.
+
+## Requirements
+
+Snapshots currently require:
+
+- Linux on x86-64 with one NVIDIA GPU.
+- TP1 with one unauthenticated plaintext HTTP server. Other parallel sizes,
+  TLS, middleware, Unix sockets, and speculative decoding are unsupported.
+- Because snapshot mode requires an unauthenticated plaintext HTTP server, use
+  a trusted host or network boundary, or external controls. See the vLLM
+  [Security guide](../usage/security.md).
+- [CRIU](https://github.com/checkpoint-restore/criu), its CUDA plugin, a
+  `cuda-checkpoint`-compatible helper, and `nvidia-smi` on `PATH`.
+- Root or passwordless `sudo` for CRIU.
+- `io_uring` disabled before launch because CRIU cannot dump it. Use
+  `kernel.io_uring_disabled=1` for an unprivileged process. A process running
+  as root, including the `docker exec` flow below, bypasses `=1`, so use `=2`
+  host-wide.
+- No established TCP connection to a peer outside the captured process tree.
+  Current Hugging Face hub clients hold their connections for the process
+  lifetime, so download the model in a separate step and run create with
+  `HF_HUB_OFFLINE=1`, as the quickstart below does.
+- A remote model ID and an immutable 40-character `--revision`. Local model
+  directories and mutable revisions are not supported.
+- Enough disk for the artifact, with the same installed vLLM package, model
+  files, container filesystem, and generated-cache paths available at restore.
+  Generated-cache files are not copied into the artifact. The manifest
+  fingerprints the ones the captured tree holds open, so restore fails early
+  and names the file when one is removed or replaced. Only open descriptors
+  are recorded. A library the tree mapped and then closed is not fingerprinted,
+  and CRIU still reopens it by path, so replacing that file permanently
+  invalidates the artifact without an early error.
+
+The official `vllm/vllm-openai` Linux x86-64 image includes the snapshot
+runtime. It still requires a compatible host driver, kernel, and privileges.
+Arm64 images omit it. Source installs must set `CRIU_CUDA_PLUGIN_DIR` to the
+directory containing `cuda_plugin.so`.
+
+Run snapshot commands with `docker exec` inside a long-lived container. Restore
+hands the API server off as a detached process, so a one-shot container would
+stop that server when its PID 1 exits. Snapshot preflight requires every
+component of the artifact path to be owned by the invoking user or root, with
+the directory itself at mode 0700, and the commands below run as root inside
+the official image, so the bind-mounted host directory is created with `sudo`
+and root ownership. The model downloads in its own step so the captured tree
+holds no hub connection, and create runs offline. This example also keeps the
+container filesystem and `/dev/shm` namespace stable for the lifetime of the
+artifact:
+
+```bash
+sudo sysctl kernel.io_uring_disabled=2
+
+snapshot_root="$(pwd)/vllm-snapshots"
+sudo install -d -m 0700 -o root -g root "${snapshot_root}"
+
+docker run --detach --name vllm-snapshot \
+  --gpus all \
+  --privileged \
+  --pid=host \
+  --ipc=host \
+  --network=host \
+  --mount "type=bind,source=${snapshot_root},target=/snapshots" \
+  --entrypoint sleep \
+  vllm/vllm-openai:latest infinity
+
+docker exec vllm-snapshot hf download Qwen/Qwen3-0.6B \
+  --revision c1899de289a04d12100db370d81485cdf75e47ca
+
+docker exec -e HF_HUB_OFFLINE=1 vllm-snapshot vllm snapshot create Qwen/Qwen3-0.6B \
+  --snapshot-dir /snapshots/qwen3-0.6b \
+  --revision c1899de289a04d12100db370d81485cdf75e47ca \
+  --dtype float16 \
+  --max-model-len 512
+
+docker exec vllm-snapshot vllm snapshot inspect /snapshots/qwen3-0.6b
+
+docker exec vllm-snapshot vllm snapshot restore \
+  /snapshots/qwen3-0.6b --host 0.0.0.0 --port 8000
+```
+
+Keep that container and its mounts available while the snapshot is in use.
+Stop and remove it only after the restored API server is no longer needed.
+
+Create initializes the engine, records a one-token canary, releases and reloads
+weights and KV cache to rehearse restore, then releases them again for capture.
+The manifest is published only after CRIU completes and the source tree stops.
+Creation is offline preparation and is not part of restore latency.
+
+The private `0700` artifact contains process memory, CUDA state, engine
+arguments, compatibility identity, and canary output; its manifest is `0600`.
+Literal API keys and Hugging Face tokens are redacted from the manifest's engine
+arguments. The manifest records selected environment names plus deterministic
+name-bound fingerprints, not their values. The protected CRIU artifact can
+still contain process secrets, so treat it as sensitive data. Restore reloads
+model files and KV cache before binding HTTP, then reproduces the canary or
+tears down the restored tree. The inspect command prints that identity and
+canary without executing the saved process.
+
+## Restore behavior
+
+Restore fails before CRIU runs if the saved identity does not match the current
+host. It does not silently fall back to ordinary startup. After CRIU restores
+the process tree, vLLM releases the saved engine to bind the requested HTTP
+address and checks the first generated token and sampled-token log probability
+against the snapshot canary. The command returns only after that check passes.
+The restored API server continues to run as a detached process.
+
+The pre-release port probe is best-effort only. It neither reserves the port
+nor authenticates the listener that appears afterward.
+
+Rollback terminates and waits for the restored tree after process identity
+verification succeeds. If verification fails, vLLM writes an abort marker for
+the snapshot server instead of signaling unverified PIDs. If that server has
+already exited, surviving engine processes can require operator cleanup.
+Verify that the previous tree has stopped before retrying; do not kill a process
+solely because its PID appears in the manifest, since PIDs can be reused.
+
+The artifact is reusable after its previous restored tree stops. Only one
+snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time.
+
+## Tradeoffs and limitations
+
+- Creation has its own latency and briefly requires the full engine. Artifact
+  size can approach the captured process and GPU memory.
+- Artifacts are not guaranteed to survive a power loss. The manifest is
+  fsynced, but publication does not explicitly flush all captured payload files
+  and directories. Recreate the artifact after an unclean host shutdown.
+- Restore currently requires the same host, GPU, driver, kernel, Python,
+  PyTorch, installed vLLM version, model revision, engine arguments,
+  selected environment variables, and CRIU plugin binaries.
+- Only dense float16 TP1 has been validated. Other model formats depend on their
+  existing sleep level 2 reload support; NCCL state is not restored.
+- CRIU support varies by kernel and driver. Preserve package, library, model,
+  and generated-cache paths for the artifact lifetime.
+- A snapshot can include application secrets or request state present in the
+  process. Build it before traffic and protect it like process memory.
+- Any feature that opens an external connection must close it before capture.
+
+For a lower-complexity option that retains a live process, see
+[Sleep mode](sleep_mode.md). Sleep mode and initialized snapshots retain
+different amounts of state and have different idle resource costs.
diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md
index 54a54d89887d..22a2350d8095 100644
--- a/docs/features/kv_offloading_usage.md
+++ b/docs/features/kv_offloading_usage.md
@@ -22,6 +22,35 @@ flowchart LR
     CPU <--> SN["..."]
 ```
 
+## Per-request load control
+
+Individual requests can cap how many tokens are loaded from offloaded storage
+by setting `max_load_tokens` in `kv_transfer_params`. The cap applies to tokens
+beyond those already available in the GPU prefix cache. Set it to `0` to
+disable external loading for the request:
+
+```json
+{
+  "kv_transfer_params": {
+    "max_load_tokens": 0
+  }
+}
+```
+
+GPU prefix-cache reuse remains enabled, and tokens beyond the aligned load cap
+are recomputed. Omitting the field leaves loading uncapped. The value must be a
+non-negative integer; invalid values are ignored. Positive caps are rounded
+down to a boundary supported by the configured KV cache groups. Offloading
+remains enabled unless it is controlled separately with `max_offload_tokens`.
+
+!!! warning
+    `max_load_tokens` is experimental and subject to change.
+
+`kv_load_tiers` continues to select secondary tiers. CPU is always included:
+it can satisfy a resident hit directly and is the required staging tier for
+secondary-to-GPU loads. Consequently, an empty tier list allows CPU hits but
+does not query any secondary tier.
+
 ## Terminology: Chunks
 
 The unit of operation is a **chunk** — a fixed-size piece of KV data covering a group of tokens. By default, a chunk maps to a single accelerator block. A configurable `blocks_per_chunk` parameter allows larger chunks, yielding larger I/Os to the host and secondary tiers.
@@ -203,9 +232,31 @@ Block content hashes must match across instances for peers to exchange blocks (s
 | `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. |
 | `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. |
 | `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. |
+| `unbound_store_timeout_s` | no | `60` | Seconds a producer holds stored blocks for a consumer that has not fetched them yet. Raise it for deployments whose prefills outlast the default; the blocks keep primary-tier CPU slots pinned for the whole window. Once it expires a late fetch is rejected in a single round trip, so the consumer falls back to local prefill immediately. |
 
 The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process.
 
+A producer parks a request's blocks until the consumer's `FetchMsg` arrives; if none arrives
+within `unbound_store_timeout_s` the blocks are released so they stop pinning primary-tier
+slots. Raise it when prefills legitimately take longer than the default:
+
+```bash
+vllm serve <model> \
+  --kv-transfer-config '{
+    "kv_connector": "OffloadingConnector",
+    "kv_role": "kv_both",
+    "kv_connector_extra_config": {
+      "spec_name": "TieringOffloadingSpec",
+      "secondary_tiers": [
+        {
+          "type": "p2p",
+          "unbound_store_timeout_s": 180
+        }
+      ]
+    }
+  }'
+```
+
 #### Environment Variables
 
 Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence.
diff --git a/docs/features/lora.md b/docs/features/lora.md
index 839624542ec5..78004722a935 100644
--- a/docs/features/lora.md
+++ b/docs/features/lora.md
@@ -422,6 +422,18 @@ vllm serve ibm-granite/granite-speech-3.3-2b \
 
 Note: Default multimodal LoRAs are currently only available for `.generate` and chat completions.
 
+## Sequence-Classification LoRA Adapters
+
+vLLM supports PEFT sequence-classification adapters that save a complete, single-layer linear classification head through `modules_to_save`. The saved module must be named `score` or `classifier`.
+
+See [classification_with_lora_offline.py](../../examples/pooling/classify/classification_with_lora_offline.py) for an offline classification example using a LoRA adapter.
+
+This support has the following limitations:
+
+- All adapters in one engine must have the same `num_labels` and hidden size as the base classification head. An incompatible adapter is rejected when it is loaded.
+- A classification head stored as float32 is converted to the runtime head dtype when it is loaded.
+- Token-classification adapters are not supported by this feature.
+
 ## Using Tips
 
 ### Configuring `max_lora_rank`
diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md
index 0a5bda79a552..6c906e7db18a 100644
--- a/docs/features/moriio_connector_usage.md
+++ b/docs/features/moriio_connector_usage.md
@@ -253,7 +253,7 @@ To run MoRI with RDMA, your environment must have the necessary RDMA userspace l
 
 The official image `vllm/vllm-openai-rocm:nightly` comes pre-installed with userspace libraries for the following NICs and kernel module versions:
 
-- AINIC (AMD Pensando Pollara): version `1.117.5`, ships `libionic1=54.0-187-1`, tested with `ionic-dkms=26.03.3.001`
+- AINIC (AMD Pensando Pollara): version `1.117.5-a-77`, ships `libionic1=54.0-187-1`, tested with `ionic-dkms=26.03.3.001`
 - Thor2 (Broadcom): version `235.2.86.0`, tested with `bnxt-en-dkms=1.10.3.235.2.86.0`, `bnxt-re-dkms=235.2.86.0`
 
 Refer to [Dockerfile.rocm](../../docker/Dockerfile.rocm) for more details. For users with NICs, kernel modules, and/or FW other than those stated above we refer to
diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md
index 031fba5b2f04..a41d5868b169 100644
--- a/docs/features/nixl_connector_compatibility.md
+++ b/docs/features/nixl_connector_compatibility.md
@@ -104,6 +104,24 @@ By default, a **compatibility hash** is checked during handshake. P and D instan
 - `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting.
 - Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA.
 
+### Pipeline parallelism
+
+The default (pull) NixlConnector does **not** support
+`pipeline-parallel-size > 1` together with a hybrid KV cache (HMA)
+layout: region indices are not a stable identity across the
+prefill/decode layer split, so the connector raises at startup. Use the
+push connector (`NixlPushConnector`) for pipeline parallelism with
+hybrid KV caches — it routes transfers by layer-name (member) identity,
+so a PP-sharded prefiller can write into a `PP=1` decoder. See
+[NIXL push-mode KV transfer](../design/nixl_kv_push_connector.md).
+
+Current push PP + HMA limitations:
+
+- Only the prefiller (producer) may be PP-sharded; decode-side PP is not supported.
+- Hybrid SSM/Mamba layouts are not supported under PP.
+- HMA requires the same block size on P and D.
+- Attention-HMA member routing requires decode TP to be no greater than prefill TP.
+
 ### Quantized KV cache
 
 [Quantized KV cache](quantization/quantized_kvcache.md) (e.g., FP8) requires both P and D instances to use the **same** `cache_dtype`. Mismatched cache dtypes will fail the compatibility hash check during handshake.
diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md
index 9799eb9cf1b8..0d1f16e2aa45 100644
--- a/docs/features/nixl_connector_usage.md
+++ b/docs/features/nixl_connector_usage.md
@@ -418,12 +418,13 @@ activity for the last reporting interval. Example output:
 ```text
 KV Transfer metrics: Num successful transfers=4, Avg xfer time (ms)=1.381,
 P90 xfer time (ms)=2.601, Avg post time (ms)=0.672, P90 post time (ms)=0.801,
-Avg MB per transfer=2.25, Throughput (MB/s)=1629.549, Avg number of descriptors=72.0
+Avg MB per transfer=2.25, Throughput (MB/s)=1629.549, Avg number of descriptors=72.0,
+Num failed transfers=0, Num KV expired reqs=0
 ```
 
 The table below describes each field. All timing values cover only the
-successful transfers recorded in the current interval; failed transfers are
-counted separately via Prometheus (see
+successful transfers recorded in the current interval; failures are reported
+as interval counts in the same line and via Prometheus (see
 [Prometheus metrics](#prometheus-metrics) below).
 
 | Metric | Unit | Description |
@@ -436,6 +437,8 @@ counted separately via Prometheus (see
 | `Avg MB per transfer` | MB | Mean payload size per transfer, computed as `total bytes transferred / number of transfers`. Reflects the average KV cache footprint of a single request (sequence length × layers × head dimension × dtype bytes). |
 | `Throughput (MB/s)` | MB/s | Effective bandwidth over the interval: `total MB transferred / total xfer time (s)` across all successful transfers. This is aggregate throughput, not per-request bandwidth. |
 | `Avg number of descriptors` | count | Mean number of NIXL memory descriptors (scatter-gather segments) submitted per transfer. More descriptors indicate more fragmented or larger KV cache allocations; very high counts can increase descriptor-registration overhead. |
+| `Num failed transfers` | count | Failed NIXL transfers, handshakes, and completion notifications (`send_notif`) in the interval, grouped because all are sporadic lower-transport-layer events. |
+| `Num KV expired reqs` | count | Requests whose KV blocks expired before being read. Reported separately from the failure counts above because it is an autoscaler/lease-tuning signal rather than a transport issue. |
 
 ### Prometheus metrics
 
@@ -448,9 +451,9 @@ exported when NixlConnector is active:
 | `vllm:nixl_post_time_seconds` | Histogram | Time to submit the transfer request to the RDMA backend (seconds). |
 | `vllm:nixl_bytes_transferred` | Histogram | Bytes moved per transfer. |
 | `vllm:nixl_num_descriptors` | Histogram | Descriptor count per transfer. |
-| `vllm:nixl_num_failed_transfers` | Counter | Cumulative count of failed NIXL KV-block transfers. |
-| `vllm:nixl_num_failed_notifications` | Counter | Cumulative count of failed completion notifications (`send_notif`). |
-| `vllm:nixl_num_kv_expired_reqs` | Counter | Requests whose KV blocks expired on the prefiller before the decoder read them (tracked on the P instance). |
+| `vllm:nixl_num_failed_transfers` | Counter | Cumulative count of failed NIXL KV-block transfers, including handshake and notification (`send_notif`) failures. These are grouped because all are sporadic lower-transport-layer events. |
+| `vllm:nixl_num_failed_notifications` | Counter | Cumulative count of failed notifications (`send_notif`), retained for compatibility. These failures are also included in `vllm:nixl_num_failed_transfers`; do not sum the two counters. |
+| `vllm:nixl_num_kv_expired_reqs` | Counter | Requests whose KV blocks expired on the prefiller before the decoder read them (tracked on the P instance). Kept separate from the failure counters above: KV expiry is an autoscaler/lease-tuning signal, not a transport issue. |
 
 !!! tip
     High `vllm:nixl_num_kv_expired_reqs` indicates that the prefiller's lease
diff --git a/docs/features/per_request_metrics.md b/docs/features/per_request_metrics.md
index 25211ecfc39b..da1685f92925 100644
--- a/docs/features/per_request_metrics.md
+++ b/docs/features/per_request_metrics.md
@@ -120,6 +120,18 @@ the same `metrics` response field. As with `n > 1`, metrics are omitted for
 requests with multiple prompts, because the timing data cannot be attributed to
 a single prompt's generation.
 
+## Responses API
+
+Per-request metrics are available on the `/v1/responses` endpoint using the
+common `metrics` object described above. Non-streaming responses include it at
+the top level. Streaming responses include it in the final response carried by
+the `response.completed` event; intermediate events do not include metrics.
+
+Metrics are omitted for Responses requests that perform multiple
+model-generation turns, such as built-in tool-call workflows, because the
+response retains timing data for only one generation turn while token usage is
+accumulated across all turns.
+
 ## Relationship to Prometheus Metrics
 
 The `metrics` response field provides per-request values for a single request.
diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md
index 69ece3607610..ba1a5b86fe45 100644
--- a/docs/features/quantization/README.md
+++ b/docs/features/quantization/README.md
@@ -23,6 +23,23 @@ The following are the supported quantization formats for vLLM:
 - [TorchAO](torchao.md)
 - [FP8 ViT Encoder Attention](fp8_vit_attn.md)
 
+## Selecting Linear Backends per Quantization
+
+`--linear-backend` selects one backend for all quantized linear layers. For
+mixed-precision models that use more than one linear quantization scheme, use
+`linear_backend_per_quant` to override the backend for individual schemes:
+
+```bash
+vllm serve <model> \
+  --linear-backend cutlass \
+  --kernel-config '{"linear_backend_per_quant":{"nvfp4_w4a16":"humming"}}'
+```
+
+Here, NVFP4 W4A16 linear layers use Humming, while all other quantized linear
+layers use CUTLASS. Per-quantization overrides take precedence over
+`--linear-backend`; schemes without an override continue to use the global
+setting, including automatic selection when it is `auto`.
+
 ## Supported Hardware
 
 The table below shows the compatibility of various quantization implementations with different hardware platforms in vLLM:
diff --git a/docs/features/quantization/quark.md b/docs/features/quantization/quark.md
index 8445288ff97f..e5cfb8653bd6 100644
--- a/docs/features/quantization/quark.md
+++ b/docs/features/quantization/quark.md
@@ -100,7 +100,7 @@ kv-cache and the quantization algorithm is AutoSmoothQuant.
 
 !!! note
     Note the quantization algorithm needs a JSON config file and the config file is located in
-    [Quark Pytorch examples](https://quark.docs.amd.com/latest/pytorch/pytorch_examples.html),
+    [Quark PyTorch examples](https://quark.docs.amd.com/latest/pytorch/pytorch_examples.html),
     under the directory `examples/torch/language_modeling/llm_ptq/models`. For example,
     AutoSmoothQuant config file for Llama is
     `examples/torch/language_modeling/llm_ptq/models/llama/autosmoothquant_config.json`.
diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md
index fbd23fcdd581..a58dc6033e13 100644
--- a/docs/features/reasoning_outputs.md
+++ b/docs/features/reasoning_outputs.md
@@ -15,6 +15,7 @@ vLLM currently supports the following reasoning models:
 | Model Series | Parser Name | Structured Output Support | Tool Calling |
 | ------------ | ----------- | ---------------- | ----------- |
 | [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
+| [Cohere Command A Plus](https://huggingface.co/CohereLabs/command-a-plus-05-2026-bf16) | `cohere_command4` | `json`, `regex` | ✅ |
 | [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
 | [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ |
 | [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md
index dcb6ad8c940c..3e419707f219 100644
--- a/docs/features/sleep_mode.md
+++ b/docs/features/sleep_mode.md
@@ -74,6 +74,33 @@ llm.wake_up(tags=["weights"])
 llm.wake_up(tags=["kv_cache"])
 ```
 
+#### Retaining frozen weights during RLHF updates
+
+Set `sleep_preserve_parameter_names` (CLI: `--sleep-preserve-parameter-names`)
+to runtime parameter-name glob
+patterns for weights that stay frozen during training. Each pattern must match
+`model.named_parameters()`; checkpoint names and `requires_grad` are not used.
+
+Level-2 sleep backs up selected GPU parameters to pageable CPU memory; weights
+wake-up restores them in place and releases the backups. CPU parameters need no
+copy. Allow enough host memory for backups. Level-1 behavior is unchanged.
+
+The trainer must omit these parameters from updates, and reload/post-processing
+must preserve their values and storage (including avoiding replacement with meta
+tensors). This option does not filter updates and applies only to the target model.
+
+#### Release only KV cache memory
+
+`LLM.release_kv_cache_memory()` discards KV cache physical memory while keeping model weights resident. It requires a completed pause and all executor memory to be resident: full sleep, partial wake-up, and repeated release without restoring memory are rejected. Requests retained with `mode="keep"` are recomputed after wake-up.
+
+Use the default `cumem` backend with managed allocations: `enable_cumem_allocator=True` on CUDA/ROCm (also enabled by `enable_sleep_mode=True`), or `enable_sleep_mode=True` on XPU. Other backends must implement selective discard; CPU and unmanaged allocations are unsupported.
+
+```python
+llm.sleep(level=0, mode="keep")  # Wait for the pause to complete.
+llm.release_kv_cache_memory()
+llm.wake_up(tags=["kv_cache"])  # Reallocate KV cache and resume scheduling.
+```
+
 ### Online Serving
 
 To enable sleep mode in a vLLM server you need to initialize it with the flag `VLLM_SERVER_DEV_MODE=1` and pass `--enable-sleep-mode` to the vLLM server.
@@ -107,9 +134,18 @@ curl -X POST 'http://localhost:8000/collective_rpc' -H 'Content-Type: applicatio
 curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache'
 ```
 
+To release only KV cache memory, wait for the level 0 pause to complete before calling the release endpoint. The same resident-memory and backend requirements as the Python API apply.
+
+```bash
+curl -X POST 'http://localhost:8000/sleep?level=0&mode=keep'
+curl -X POST 'http://localhost:8000/release_kv_cache_memory'
+curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache'
+```
+
 #### HTTP endpoints
 
 - `POST /sleep?level=1` — Put the model to sleep (`level=1`).
+- `POST /release_kv_cache_memory` — Discard KV cache memory after a completed pause, while all executor memory is resident.
 - `POST /wake_up` — Wake up the model. Supports optional `tags` query parameters for partial wake-up (e.g., `?tags=weights`).
 - `POST /collective_rpc` — Perform a collective remote procedure call (RPC).
 - `GET /is_sleeping` — Check if the model is sleeping.
diff --git a/docs/features/speculative_decoding/draft_model.md b/docs/features/speculative_decoding/draft_model.md
index 5ab4629409f4..880ea1df5989 100644
--- a/docs/features/speculative_decoding/draft_model.md
+++ b/docs/features/speculative_decoding/draft_model.md
@@ -96,7 +96,7 @@ The code used to request completions as a client remains unchanged:
       },
       gpu_memory_utilization=0.5,
   )
-outputs = llm.generate(prompts，sampling_params)
+outputs = llm.generate(prompts, sampling_params)
 
 for output in outputs:
       prompt = output.prompt
diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md
index ae65231919a3..fb70c5024835 100644
--- a/docs/features/tool_calling.md
+++ b/docs/features/tool_calling.md
@@ -109,18 +109,22 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t
 
 ## Constrained Decoding Behavior
 
-Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode and the per-tool `strict` field:
+Structural-tag parsers resolve call obligation, grammar activation, and argument-schema enforcement separately. With structural-tag enforcement enabled:
 
-| `tool_choice` value | Schema-constrained decoding | Behavior |
+| `tool_choice` value | Call obligation | Structural-tag activation |
 | --- | --- | --- |
-| Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. |
-| `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. |
-| `"auto"` | Only when `strict: true` is set on at least one tool | Structural-tag parsers constrain tool-call arguments when a tool opts in with `strict: true`. Without it, the model generates freely and tool calls are extracted from raw text. |
-| `"none"` | N/A | No tool calls are produced. |
+| Named function | Call the selected function | Always |
+| `"required"` | Produce at least one tool call | Always |
+| `"auto"` | Tool calls are optional | When at least one tool sets `strict: true`, or `--tool-strict-level` is `function` or `parameter` |
+| `"none"` | Tool calling is disabled | Disabled |
+
+When a structural tag applies, each tool's declared parameter schema is enforced only when that tool sets `strict: true` or the server uses `--tool-strict-level parameter`. Tools with omitted or false `strict` receive broad argument-syntax constraints at levels `auto` and `function`, including for required and named calls.
+
+For parsers using schema-derived JSON constraints, required and named calls continue to enforce the declared parameter schemas.
 
 ### Strict Mode
 
-For `tool_choice="required"` or named function calling, structural-tag constraints are always applied regardless of the `strict` field. For `tool_choice="auto"`, setting `strict: true` on at least one tool opts in to structural-tag constraints; without it, the model generates freely and tool calls are extracted from raw text. The `strict` field is supported across all three API surfaces: Chat Completion, Responses, and Anthropic Messages.
+Structural-tag constraints have two layers. The tool-call *envelope* (markup and function name) is constrained whenever a structural tag applies: always for `tool_choice="required"` and named function calling, and for `tool_choice="auto"` once at least one tool sets `strict: true` or the server raises the floor. The *argument schema* of an individual tool is pinned only when that tool sets `strict: true` (or under `--tool-strict-level parameter`); a tool that omits `strict` keeps its arguments unconstrained, even when another tool in the same request is strict. The `strict` field is supported across all three API surfaces: Chat Completion, Responses, and Anthropic Messages.
 
 For best compatibility with strict schema enforcement, define tool parameter schemas in the OpenAI strict-schema style:
 
@@ -134,6 +138,22 @@ vLLM also provides a global toggle via the `VLLM_ENFORCE_STRICT_TOOL_CALLING` en
 VLLM_ENFORCE_STRICT_TOOL_CALLING=false vllm serve ...
 ```
 
+### Server-Side Strictness Floor
+
+Most OpenAI-compatible clients and agent frameworks never set `strict` on their tools, so with `tool_choice="auto"` the model generates tool calls without any grammar and malformed markup can leak into the response. The `--tool-strict-level` option lets the server operator raise the floor for every request that carries tools, independently of what the client declares:
+
+| Value | Behavior |
+| --- | --- |
+| `auto` (default) | Follow the request's tool choice and per-tool strictness. Required/named choices activate structural tags; `tool_choice="auto"` activates them when at least one tool sets `strict: true`. |
+| `function` | Constrain the tool-call envelope (markup and the function name) for every request with tools, leaving argument contents free unless the client marked the tool `strict: true`. |
+| `parameter` | Additionally pin argument schemas for every tool, as if every tool had `strict: true`. |
+
+```bash
+vllm serve ... --tool-strict-level function
+```
+
+The floor never relaxes a constraint the request would already receive: tools the client marked `strict: true` keep their schemas at every level. With `tool_choice="auto"`, the grammar does not force a tool call; a plain text response stays valid. `VLLM_ENFORCE_STRICT_TOOL_CALLING=false` disables structural tags entirely and takes precedence over this option.
+
 ## Automatic Function Calling
 
 To enable this feature, you should set the following flags:
@@ -152,7 +172,7 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso
 If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template!
 
 !!! note
-    With `tool_choice="auto"`, schema-level constraint requires both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`. When these conditions are met and the selected parser supports structural tags, vLLM constrains tool-call arguments. Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema.
+    With `tool_choice="auto"`, structural-tag constraints require both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`, or a server-side floor set via `--tool-strict-level`. When these conditions are met and the selected parser supports structural tags, vLLM constrains the tool-call envelope and pins the argument schema of each tool that sets `strict: true` (or of every tool under `--tool-strict-level parameter`). Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema.
 
 ### Hermes Models (`hermes`)
 
@@ -162,8 +182,8 @@ All Nous Research Hermes-series models newer than Hermes 2 Pro should be support
 * `NousResearch/Hermes-2-Theta-*`
 * `NousResearch/Hermes-3-*`
 
-_Note that the Hermes 2 **Theta** models are known to have degraded tool call quality and capabilities due to the merge
-step in their creation_.
+*Note that the Hermes 2 **Theta** models are known to have degraded tool call quality and capabilities due to the merge
+step in their creation*.
 
 Flags: `--tool-call-parser hermes`
 
@@ -366,15 +386,16 @@ Flags:
 * For non-reasoning: `--tool-call-parser hunyuan_a13b`
 * For reasoning: `--tool-call-parser hunyuan_a13b --reasoning-parser hunyuan_a13b`
 
-### Cohere Command A Reasoning (`cohere_command3`)
+### Cohere Command (`cohere_command3` / `cohere_command4`)
 
 Supported models:
 
-* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025)
+* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) (`cohere_command3`)
+* [`CohereLabs/command-a-plus-05-2026`](https://huggingface.co/CohereLabs/command-a-plus-05-2026-bf16) and [`CohereLabs/North-Mini-Code-1.0`](https://huggingface.co/CohereLabs/North-Mini-Code-1.0) (`cohere_command4`)
 
-Flags: `--tool-call-parser cohere_command3 --reasoning-parser cohere_command3`
+Flags: `--tool-call-parser cohere_command4 --reasoning-parser cohere_command4`.
 
-Note: the Cohere tool parser requires the `cohere_melody` package, which is not installed by default. Before using this parser please install the [cohere_melody](https://pypi.org/project/cohere-melody/) package.
+Note: the Cohere parser requires the `cohere_melody` package, which is not installed by default. Before using this parser please install the [cohere_melody](https://pypi.org/project/cohere-melody/) package.
 
 ### LongCat-Flash-Chat Models (`longcat`)
 
diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md
index 7c601e13aaa2..54c7e58ce7d9 100644
--- a/docs/features/watermarking.md
+++ b/docs/features/watermarking.md
@@ -29,16 +29,27 @@ from vllm import SamplingParams
 sampling_params = SamplingParams(watermarking=False)
 ```
 
-The OpenAI-compatible APIs accept the same `watermarking: false` request field.
+The Python OpenAI-compatible APIs and the Rust chat/completions APIs accept the
+same `watermarking: false` request field. The Rust token API accepts it in
+`sampling_params`, and the Rust gRPC `GenerateRequest` accepts
+the optional `watermarking` field. Omitting the field keeps watermarking enabled
+when the engine has a watermark configuration.
+
 Deployments that require watermarking must restrict this field to trusted
 callers, or strip and validate it at the ingress boundary, so untrusted clients
 cannot opt out.
 
-`context_width` controls how many prior output tokens seed each watermark
-decision and defaults to 4. Larger values make the watermark less robust to
+`context_width` controls how many prior tokens seed each watermark decision
+and defaults to 4. Larger values make the watermark less robust to
 edits because an insertion, deletion, or substitution changes more subsequent
 contexts. Values above 16 are allowed but emit a warning.
 
+`allow_target_only_watermarking` defaults to false and only has an effect when
+speculative decoding is enabled. It permits speculative decoding with a
+watermarking algorithm that does not support it natively, at the cost of
+weaker detectability. See
+[Speculative decoding](#speculative-decoding).
+
 ## Architecture
 
 `WatermarkConfig` selects an algorithm and PRF. Model Runner V2 constructs the
@@ -52,21 +63,108 @@ the reference algorithms. `WatermarkDetector` consumes token IDs, so callers
 remain responsible for using the tokenizer and watermark profile that match
 generation.
 
+## Speculative decoding
+
+Watermarking requires speculative decoding to use probabilistic draft sampling,
+standard rejection sampling, and an autoregressive model-based method (`dspark`,
+`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dspark`.
+
+A watermarking algorithm without native speculative-decoding support is
+rejected before model loading. Set
+`"allow_target_only_watermarking": true` to allow it: accepted draft tokens are
+not watermarked, while target-side rejection recovery and bonus sampling remain
+watermarked. The watermark signal is diluted in proportion to the share of
+output tokens supplied by accepted drafts; rejected drafts do not dilute it
+because their recovery tokens are watermarked.
+
+Speculative-decoding token paths do not currently support generation-side
+context deduplication. The configured `deduplicate_contexts` policy is not
+applied to accepted drafts, rejection-recovery tokens, or bonus tokens.
+
+For `dual_key_gumbel`, `alpha` has no effect under speculative decoding. The
+speculative protocol selects the key for each token instead.
+
 ## Algorithms
 
 ### Gumbel-max
 
-Gumbel-max derives a deterministic pseudorandom value from the key, prior
-generated-token context, and every candidate token, then uses the resulting
-Gumbel noise for categorical sampling. See
+Gumbel-max derives a deterministic pseudorandom value from the key, prior token
+context, and every candidate token, then uses the resulting Gumbel noise for
+categorical sampling. See
 [Aaronson's original presentation](https://simons.berkeley.edu/sites/default/files/2024-10/LLM24-2%20Slides%20-%20Scott%20Aaronson.pdf).
 
 Gumbel-max requires stochastic sampling. Greedy requests (`temperature=0`)
 bypass watermarking and emit a warning once per worker.
 
-Single-key Gumbel-max watermarking may increase degenerate generations. In
-particular, reusing the same keyed random vector when a context repeats can
-reinforce repetition loops that continue until the maximum token limit.
+When a token context is repeated, generation can use ordinary sampling for that
+occurrence. Reusing the repeated context leads to a bias over the sequence, as
+certain token choices would be correlated. Ordinary sampling at these positions
+allows for single-sequence non-distortion (see section G.3 of the
+[SynthID-Text supplementary materials](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf)).
+The detector independently deduplicates contexts so repeated keyed random
+vectors are not treated as independent evidence, meaning that context
+deduplication at generation time does not reduce the watermarking signal, unless
+a user tampers with the output and removes the first occurrence while leaving an
+unwatermarked repetition.
+
+`deduplicate_contexts` controls which history is searched for a repeated
+context:
+
+- `"none"` disables context deduplication.
+- `"single_turn"`, the default, searches the tokens generated for the current
+  request. This gives single-turn non-distortion.
+- `"all"` also searches the prompt, which extends non-distortion across the
+  turns of a conversation. The first `context_width` generated tokens use
+  ordinary sampling, and subsequent tokens whose context already occurs in the
+  prompt are not watermarked. Previous turns of a session, or a prompt that
+  already contains the structure of the answer such as a tool result the model
+  extends, may therefore leave little of the answer watermarked. Given the
+  impact of `"all"` on both performance (sampling steps must scan over a
+  potentially large amount of previous context) and watermarking detectability
+  (we may significantly reduce the amount of watermarked tokens), you should
+  only use `"all"` when non-distortion across multiple turns is required.
+
+`deduplicate_contexts_max_history` limits the search to the most recent
+positions and defaults to 8,192. Each position is compared over the
+`context_width` tokens before it, so the window of tokens read is that much
+longer. A smaller value reduces scanning cost but only provides the guarantee
+within that window. Set it to `null` to search back to the start of the
+generation for `"single_turn"`, or of the request for `"all"`; an unbounded
+search costs more as the sequence grows. The setting has no effect when
+`deduplicate_contexts` is `"none"`. Values below 1,024 emit a warning because a
+short window can miss repetition loops whose contexts recur farther apart.
+
+For example, this checks prompt and completion history within the default
+8,192-position window:
+
+```bash
+vllm serve MODEL \
+  --watermark-config \
+  '{"algorithm":"gumbel","key":42,"deduplicate_contexts":"all"}'
+```
+
+### Dual-key Gumbel-max
+
+Dual-key Gumbel-max derives independent keys A and B from one configured master
+key. During ordinary generation, each token uses key A with probability
+`1 - alpha` and key B with probability `alpha`. The algorithm-specific `alpha`
+parameter defaults to 0.1. Detection scores every token against both keys.
+
+The same two key streams support speculative decoding without changing its
+acceptance rate. In this mode, the speculative protocol selects the key instead
+of `alpha`: draft tokens use key A, while rejection recovery and bonus tokens
+use key B. The ordinary target-to-draft probability-ratio test remains unchanged,
+implementing [SynthID-Text Supplementary Algorithm
+6](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf).
+
+Select `dual_key_gumbel` together with probabilistic drafting:
+
+```bash
+vllm serve MODEL \
+  --speculative-config \
+  '{"method":"mtp","num_speculative_tokens":3,"draft_sample_method":"probabilistic"}' \
+  --watermark-config '{"algorithm":"dual_key_gumbel","key":42,"alpha":0.1}'
+```
 
 ### SynthID-Text
 
@@ -106,7 +204,14 @@ result = GumbelWatermarkDetector(key=42, prf="philox").detect(token_ids)
 print(result.p_value, result.is_watermarked)
 ```
 
-The tokenizer, algorithm, PRF, key, and context width must match generation.
+The detection configuration must match the generation configuration, including
+the tokenizer, PRF, watermarking algorithm, algorithm-specific watermarking
+configuration, and key. In practice, this information is often unavailable
+when checking a piece of text. Deployments should therefore retain the set of
+candidate configurations they have served, test the text against each
+candidate, and correct for multiple testing, for example with a Bonferroni
+correction to the resulting p-values.
+
 Gumbel-max detection scores repeated contexts once by default so identical PRF
 random vectors are not treated as independent evidence. Keep
 `deduplicate_contexts=True` unless the detector's calibration has been adjusted
@@ -140,7 +245,7 @@ watermarked output or to modify watermarked text so it is no longer detected.
 ## Limitations
 
 - Watermarking is currently available only with Model Runner V2.
-- Gumbel-max cannot be configured with speculative decoding.
+- Not all watermarking algorithms have native speculative-decoding support.
 - Beam search expands candidates from model log probabilities and does not apply
   Gumbel-max watermarking.
 - Models that replace the vLLM sampler with a custom sampler cannot use
diff --git a/docs/getting_started/installation/cpu.arm.inc.md b/docs/getting_started/installation/cpu.arm.inc.md
index e719d65dd95e..db49e96386e0 100644
--- a/docs/getting_started/installation/cpu.arm.inc.md
+++ b/docs/getting_started/installation/cpu.arm.inc.md
@@ -20,12 +20,12 @@ Pre-built vLLM wheels for Arm are available since version 0.11.2. These wheels c
 
 ```bash
 export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//')
-uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --torch-backend cpu
+uv pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl" --torch-backend cpu
 ```
 
 ??? console "pip"
     ```bash
-    pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu
+    pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl" --extra-index-url https://download.pytorch.org/whl/cpu
     ```
 
 The `uv` approach works for vLLM `v0.6.6` and later. A unique feature of `uv` is that packages in `--extra-index-url` have [higher priority than the default index](https://docs.astral.sh/uv/pip/compatibility/#packages-that-exist-on-multiple-indexes). If the latest public release is `v0.6.6.post1`, `uv`'s behavior allows installing a commit before `v0.6.6.post1` by specifying the `--extra-index-url`. In contrast, `pip` combines packages from `--extra-index-url` and the default index, choosing only the latest version, which makes it difficult to install a development version prior to the released version.
diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md
index 539c28d7097b..53b8eb596d28 100644
--- a/docs/getting_started/installation/cpu.md
+++ b/docs/getting_started/installation/cpu.md
@@ -150,7 +150,7 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu
 
 - `VLLM_CPU_KVCACHE_SPACE`: specify the KV Cache size (e.g, `VLLM_CPU_KVCACHE_SPACE=40` means 40 GiB space for KV cache), larger setting will allow vLLM to run more requests in parallel. This parameter should be set based on the hardware configuration and memory management pattern of users. Default value is `0`.
 - `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads, can be set as CPU id lists, `auto` (by default), or `nobind` (to disable binding to individual CPU cores and to inherit user-defined OpenMP variables). For example, `VLLM_CPU_OMP_THREADS_BIND=0-31` means there will be 32 OpenMP threads bound on 0-31 CPU cores. `VLLM_CPU_OMP_THREADS_BIND=0-31|32-63` means there will be 2 tensor parallel processes, 32 OpenMP threads of rank0 are bound on 0-31 CPU cores, and the OpenMP threads of rank1 are bound on 32-63 CPU cores. By setting to `auto`, the OpenMP threads of each rank are bound to the CPU cores in each NUMA node respectively. If set to `nobind`, the number of OpenMP threads is determined by the standard `OMP_NUM_THREADS` environment variable.
-- `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `None`. If the value is not set and use `auto` thread binding, no CPU will be reserved for `world_size == 1`, 1 CPU per rank will be reserved for `world_size > 1`.
+- `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the total number of CPU cores which are not dedicated to OpenMP threads. The variable only takes effect when `VLLM_CPU_OMP_THREADS_BIND` is set to `auto`. By default, x86, ARM, and RISC-V reserve `local_world_size` CPUs in total, while PowerPC and S390X retain a one-CPU default. When KV transfer is enabled, one additional CPU is reserved per local rank. Reserved CPUs are distributed across the local ranks.
 - `CPU_VISIBLE_MEMORY_NODES`: specify visible NUMA memory nodes for vLLM CPU workers, similar to ```CUDA_VISIBLE_DEVICES```. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. The variable provides more control for the auto thread-binding feature, such as masking nodes and changing nodes binding sequence.
 - `VLLM_ZENTORCH_WEIGHT_PREPACK` (AMD Zen only): when `ZenCpuPlatform` is active, eagerly prepack linear weights into ZenDNN's blocked layout at model load time, eliminating per-inference layout conversion overhead. Default is `1` (enabled). See [AMD Zen optimizations](#amd-zen-optimizations).
 
@@ -253,7 +253,7 @@ See [AMD Zen optimizations](#amd-zen-optimizations) for detection rules, support
 
 ### How to decide `VLLM_CPU_OMP_THREADS_BIND`?
 
-- Default `auto` thread-binding is recommended for most cases. Ideally, each OpenMP thread will be bound to a dedicated physical core respectively, threads of each rank will be bound to the same NUMA node respectively, and 1 CPU per rank will be reserved for other vLLM components when `world_size > 1`. If you have any performance problems or unexpected binding behaviours, please try to bind threads as following.
+- Default `auto` thread-binding is recommended for most cases. Ideally, each OpenMP thread will be bound to a dedicated physical core respectively, threads of each rank will be bound to the same NUMA node respectively, and the default reserved CPUs will be kept for other vLLM components. On x86, ARM, and RISC-V, `local_world_size` CPUs are reserved in total, plus one additional CPU per local rank when KV transfer is enabled. If you have any performance problems or unexpected binding behaviours, please try to bind threads as following.
 
 - On a hyper-threading enabled platform with 16 logical CPU cores / 8 physical CPU cores:
 
diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md
index 88767e4c6ed2..e1ff4ee3cb0d 100644
--- a/docs/getting_started/installation/cpu.s390x.inc.md
+++ b/docs/getting_started/installation/cpu.s390x.inc.md
@@ -138,19 +138,6 @@ VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
     The Docker image (`Dockerfile.s390x`) already includes TCMalloc and sets
     `LD_PRELOAD` automatically.
 
-!!! warning "Protobuf workaround for s390x"
-    The C++ protobuf extension crashes on s390x. After installation, set the
-    following environment variable and remove the C++ extensions:
-
-    ```bash
-    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
-
-    # Remove C++ protobuf extensions that crash on s390x
-    SITE_PKGS=$(python -c "import site; print(site.getsitepackages()[0])")
-    rm -rf "$SITE_PKGS/google/_upb/"*.so \
-           "$SITE_PKGS/google/protobuf/pyext/"*.so 2>/dev/null || true
-    ```
-
 --8<-- [end:build-wheel-from-source]
 --8<-- [start:pre-built-images]
 
diff --git a/docs/getting_started/installation/cpu.x86.inc.md b/docs/getting_started/installation/cpu.x86.inc.md
index b2ef384bd93b..260ffe286362 100644
--- a/docs/getting_started/installation/cpu.x86.inc.md
+++ b/docs/getting_started/installation/cpu.x86.inc.md
@@ -24,16 +24,16 @@ Pre-built vLLM wheels for x86 with AVX512/AVX2 are available since version 0.17.
 export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//')
 
 # use uv
-uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --torch-backend cpu
+uv pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl" --torch-backend cpu
 ```
 
 ??? console "pip"
     ```bash
     # use pip
-    pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cpu
+    pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl" --extra-index-url https://download.pytorch.org/whl/cpu
     ```
 !!! warning "set `LD_PRELOAD`"
-    Before use vLLM CPU installed via wheels, make Intel OpenMP is added to `LD_PRELOAD`:
+    Before use vLLM CPU installed via wheels, make sure Intel OpenMP is added to `LD_PRELOAD`:
     ```bash
     # manually find the path
     sudo find / -iname *libiomp5.so
diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md
index cd6f2892c390..d110e377da8d 100644
--- a/docs/getting_started/installation/gpu.cuda.inc.md
+++ b/docs/getting_started/installation/gpu.cuda.inc.md
@@ -43,7 +43,7 @@ As of now, vLLM's binaries are compiled with CUDA 12.9 and public PyTorch releas
 export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//')
 export CUDA_VERSION=130 # or other
 export CPU_ARCH=$(uname -m) # x86_64 or aarch64
-uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
+uv pip install "https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl" --extra-index-url "https://download.pytorch.org/whl/cu${CUDA_VERSION}"
 ```
 
 #### Install the latest code
diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md
index ad5af7994054..0169d0154608 100644
--- a/docs/getting_started/installation/gpu.rocm.inc.md
+++ b/docs/getting_started/installation/gpu.rocm.inc.md
@@ -324,7 +324,7 @@ You can find more information about vLLM's wheels in
         # Install dependencies
         pip install --upgrade numba \
             scipy \
-            huggingface-hub[cli] \
+            huggingface-hub \
             setuptools_scm
         pip install -r requirements/rocm.txt
 
diff --git a/docs/getting_started/installation/gpu.xpu.inc.md b/docs/getting_started/installation/gpu.xpu.inc.md
index 84dc69ca84c5..893150148aa2 100644
--- a/docs/getting_started/installation/gpu.xpu.inc.md
+++ b/docs/getting_started/installation/gpu.xpu.inc.md
@@ -32,6 +32,22 @@ To install the wheel built from the latest main branch:
 uv pip install vllm --extra-index-url https://wheels.vllm.ai/nightly/xpu --extra-index-url https://download.pytorch.org/whl/xpu --index-strategy unsafe-best-match
 ```
 
+Note that after upgrading to PyTorch 2.14, xpu graph support requires specific oneAPI packages.
+
+```bash
+uv pip install \
+       "intel-cmplr-lib-rt==2026.1.1" \
+       "intel-cmplr-lib-ur==2026.1.1" \
+       "intel-cmplr-lic-rt==2026.1.1" \
+       "intel-sycl-rt==2026.1.1" \
+       "oneccl-devel==2022.1.2; platform_system == 'Linux' and platform_machine == 'x86_64'" \
+       "oneccl==2022.1.2; platform_system == 'Linux' and platform_machine == 'x86_64'" \
+       "dpcpp-cpp-rt==2026.1.1" \
+       "intel-opencl-rt==2026.1.1" \
+       "intel-openmp==2026.1.1" \
+       "intel-pti==1.1.0"
+```
+
 #### Install specific revisions
 
 If you want to access the wheels for previous commits (e.g. to bisect the behavior change, performance regression), you can specify the commit hash in the URL:
diff --git a/docs/governance/committers.md b/docs/governance/committers.md
index 0ca6e24ded0a..48ed4b26073f 100644
--- a/docs/governance/committers.md
+++ b/docs/governance/committers.md
@@ -2,7 +2,7 @@
 
 This document lists the current committers of the vLLM project and the core areas they maintain.
 Committers have write access to the vLLM repository and are responsible for reviewing and merging PRs.
-You can also refer to the [CODEOWNERS](https://github.com/vllm-project/vllm/blob/main/.github/CODEOWNERS) file for concrete file-level ownership and reviewers. Both this documents and the CODEOWNERS file are living documents and they complement each other.
+You can also refer to the [CODEOWNERS](https://github.com/vllm-project/vllm/blob/main/.github/CODEOWNERS) file for concrete file-level ownership and reviewers. Both these documents and the CODEOWNERS file are living documents and they complement each other.
 
 ## Active Committers
 
@@ -32,6 +32,7 @@ Sorted alphabetically by GitHub handle:
 - [@ivanium](https://github.com/ivanium): KV Connector and offload
 - [@jeejeelee](https://github.com/jeejeelee): LoRA, new model support
 - [@jikunshang](https://github.com/jikunshang): Intel CPU/XPU integration
+- [@jperezdealgaba](https://github.com/jperezdealgaba): Security
 - [@khluu](https://github.com/khluu): CI infrastructure
 - [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance
 - [@markmc](https://github.com/markmc): Observability
@@ -44,7 +45,6 @@ Sorted alphabetically by GitHub handle:
 - [@pavanimajety](https://github.com/pavanimajety): NVIDIA GPU integration
 - [@ProExpertProg](https://github.com/ProExpertProg): Compilation, startup UX
 - [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg
-- [@russellb](https://github.com/russellb): Structured output, engine core, security
 - [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser
 - [@shen-shanshan](https://github.com/shen-shanshan): AMD GPU / ROCm integration, multimodality, ViT CUDA graph
 - [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community
@@ -82,6 +82,7 @@ Committers who have contributed to vLLM significantly in the past (thank you!) b
 - [@pcmoritz](https://github.com/pcmoritz): MoE
 - [@rkooo567](https://github.com/rkooo567): Chunked prefill
 - [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support
+- [@russellb](https://github.com/russellb): Structured output, engine core, security
 - [@sighingnow](https://github.com/sighingnow): Qwen models, new model support
 - [@sroy745](https://github.com/sroy745): Speculative decoding
 - [@Yard1](https://github.com/Yard1): kernels and performance
@@ -100,7 +101,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - KV Cache Manager: memory management layer within scheduler maintaining KV cache logical block data
     - @heheda12345, @WoosukKwon
 - AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints
-    - @robertgshaw2-redhat, @njhill, @russellb
+    - @robertgshaw2-redhat, @njhill
 - ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation
     - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni
 - KV Connector: Connector interface and implementation for KV cache offload and transfer
@@ -153,7 +154,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP.
     - @WoosukKwon, @benchislett, @MatthewBonanni
 - Structured Output: The structured output implementation
-    - @russellb, @aarnphm
+    - @aarnphm
 - RL: The RL related features such as collective rpc, sleep mode, etc.
     - @youkaichao, @zhuohan123
 - LoRA: @jeejeelee
@@ -167,7 +168,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - Documentation: @hmellor, @DarkLight1337, @simon-mo
 - Benchmarks: @ywang96, @simon-mo
 - CI, Build, Release Process: @khluu, @njhill, @simon-mo, @Harry-Chen, @vadiklyutiy
-- Security: @russellb
+- Security: @jperezdealgaba
 
 ### External Kernels Integration
 
diff --git a/docs/mkdocs/gen_files/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py
index 1c7025366f4b..ecdce7761fb4 100644
--- a/docs/mkdocs/gen_files/generate_argparse.py
+++ b/docs/mkdocs/gen_files/generate_argparse.py
@@ -98,8 +98,7 @@ class MockModule:
 # Mock torch.library.infer_schema for vllm.ir.ops.IrOpInplaceOverload.__init__
 # We need to return the corresponding number of inputs, as IR infra will assert it
 def get_outputs(native_fn: Callable) -> str:
-    """
-    Extract output schema from function's return type annotation,
+    """Extract output schema from function's return type annotation,
     e.g. 'Tensor' or 'Tensor, Tensor'.
     """
     import typing
@@ -249,11 +248,12 @@ def create_parser(add_cli_args, **kwargs) -> FlexibleArgumentParser:
     """Create a parser for the given class with markdown formatting.
 
     Args:
-        cls: The class to create a parser for
+        add_cli_args: The `add_cli_args` classmethod to build the parser from.
         **kwargs: Additional keyword arguments to pass to `cls.add_cli_args`.
 
     Returns:
         FlexibleArgumentParser: A parser with markdown formatting for the class.
+
     """
     try:
         parser = FlexibleArgumentParser(add_json_tip=False)
diff --git a/docs/mkdocs/gen_files/generate_attention_backends.py b/docs/mkdocs/gen_files/generate_attention_backends.py
index 942165eaa940..725d7072ca92 100644
--- a/docs/mkdocs/gen_files/generate_attention_backends.py
+++ b/docs/mkdocs/gen_files/generate_attention_backends.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Generates documentation table for attention backends showing feature support.
+"""Generates documentation table for attention backends showing feature support.
 
 This script parses all registered attention backends using AST (no imports needed)
 and generates a markdown table showing what features each backend supports,
@@ -304,6 +303,7 @@ def parse_mla_prefill_registry() -> dict[str, str]:
 
     Returns:
         A dict mapping backend names to their class paths.
+
     """
     if not MLA_PREFILL_REGISTRY_FILE.exists():
         return {}
@@ -325,6 +325,7 @@ def parse_mla_prefill_priorities() -> dict[str, list[str]]:
     Returns:
         A dict with keys like 'blackwell' and 'default' containing
         lists of backend enum names in priority order.
+
     """
     if not MLA_PREFILL_SELECTOR_FILE.exists():
         return {}
@@ -549,6 +550,7 @@ def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:
 
     Returns:
         A dict with backend properties, or None if parsing fails.
+
     """
     file_path = get_file_from_class_path(class_path)
     if file_path is None:
diff --git a/docs/mkdocs/gen_files/generate_examples.py b/docs/mkdocs/gen_files/generate_examples.py
index c9c10018ecfb..a514dcf879ff 100644
--- a/docs/mkdocs/gen_files/generate_examples.py
+++ b/docs/mkdocs/gen_files/generate_examples.py
@@ -49,8 +49,7 @@ def title(text: str) -> str:
 
 @dataclass
 class Example:
-    """
-    Example class for generating documentation content from a given path.
+    """Example class for generating documentation content from a given path.
 
     Attributes:
         path (Path): The path to the main directory or file.
@@ -64,6 +63,7 @@ class Example:
 
     Methods:
         generate() -> str: Generates the documentation content.
+
     """
 
     path: Path
@@ -116,8 +116,7 @@ def title(self) -> str:
         raise ValueError(f"Title not found in {self.main_file}")
 
     def fix_relative_links(self, content: str) -> str:
-        """
-        Fix relative links in markdown content by converting them to gh-file
+        """Fix relative links in markdown content by converting them to gh-file
         format.
 
         Args:
@@ -125,6 +124,7 @@ def fix_relative_links(self, content: str) -> str:
 
         Returns:
             str: Content with relative links converted to gh-file format
+
         """
         # Regex to match markdown links [text](relative_path)
         # This matches links that don't start with http, https, ftp, or #
diff --git a/docs/mkdocs/gen_files/generated_content.py b/docs/mkdocs/gen_files/generated_content.py
index fa3f31b294bf..d6691944246d 100644
--- a/docs/mkdocs/gen_files/generated_content.py
+++ b/docs/mkdocs/gen_files/generated_content.py
@@ -35,6 +35,7 @@ def fill_markers(doc_path: str, blocks: dict[str, str]) -> None:
         FileNotFoundError: If the source page does not exist.
         ValueError: If an expected marker is missing, or any `gen:` marker is
             left unsubstituted after filling.
+
     """
     source = DOCS_DIR / doc_path
     if not source.exists():
diff --git a/docs/mkdocs/hooks/autoref_code.py b/docs/mkdocs/hooks/autoref_code.py
index 647f74f202d8..58cec14d45b3 100644
--- a/docs/mkdocs/hooks/autoref_code.py
+++ b/docs/mkdocs/hooks/autoref_code.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-MkDocs hook to automatically convert inline code references to API doc links.
+"""MkDocs hook to automatically convert inline code references to API doc links.
 
 For example, `WeightTransferConfig` becomes
 [`WeightTransferConfig`][vllm.config.WeightTransferConfig]
diff --git a/docs/mkdocs/hooks/url_schemes.py b/docs/mkdocs/hooks/url_schemes.py
index d208b9d6779e..cee418936f47 100644
--- a/docs/mkdocs/hooks/url_schemes.py
+++ b/docs/mkdocs/hooks/url_schemes.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-MkDocs hook + markdown extension to enable the following links to render correctly,
+"""MkDocs hook + markdown extension to enable the following links to render correctly,
 including inside content included via pymdownx.snippets:
 
 - Relative file links outside of the `docs/` directory, e.g.:
@@ -71,9 +70,7 @@ def run(self, lines):
             return lines
 
         def replace_relative_link(match: re.Match) -> str:
-            """
-            Replace relative file links with URLs if they point outside the docs dir.
-            """
+            """Replace relative file links pointing outside the docs dir with URLs."""
             title = match.group("title")
             path = match.group("path")
             path = ((DOC_DIR / page.file.src_uri).parent / path).resolve()
@@ -91,9 +88,7 @@ def replace_relative_link(match: re.Match) -> str:
             return f"[{gh_icon} {title}]({url})"
 
         def replace_github_link(match: re.Match) -> str:
-            """
-            Replace GitHub issue, PR, and project links with enhanced Markdown links.
-            """
+            """Replace GitHub issue, PR, and project links with rich Markdown links."""
             repo = match.group("repo")
             type = match.group("type")
             number = match.group("number")
diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md
index a60bd4307a1e..13a2e0d38147 100644
--- a/docs/models/pooling_models/embed.md
+++ b/docs/models/pooling_models/embed.md
@@ -47,7 +47,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th
 | `JinaEmbeddingsV5Model`<sup>C</sup> | Qwen3-decoder or EuroBERT-encoder backbone with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small`, `jinaai/jina-embeddings-v5-text-nano` (see note) | ✅︎ | ✅︎ |
 | `LlamaBidirectionalModel`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ |
 | `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ |
-| `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | |
+| `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | ✅︎ | |
 | `NomicBertModel` | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | |
 | `Qwen2Model`<sup>C</sup>, `Qwen2ForCausalLM`<sup>C</sup> | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ |
 | `Qwen3Model`<sup>C</sup>, `Qwen3ForCausalLM`<sup>C</sup> | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, etc. | ✅︎ | ✅︎ |
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index d06154fb5bd5..df44a3e8b0e5 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -164,9 +164,9 @@ For your model to be compatible with vLLM's tensor parallel and/or pipeline para
 
 ```python
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
-class MyConfig(PretrainedConfig):
+class MyConfig(PreTrainedConfig):
     base_model_tp_plan = {
         "layers.*.self_attn.k_proj": "colwise",
         "layers.*.self_attn.v_proj": "colwise",
@@ -240,7 +240,7 @@ Alternatively, you can [open an issue on GitHub](https://github.com/vllm-project
 
 #### Download a model
 
-If you prefer, you can use the Hugging Face CLI to [download a model](https://huggingface.co/docs/huggingface_hub/guides/cli#huggingface-cli-download) or specific files from a model repository:
+If you prefer, you can use the Hugging Face CLI to [download a model](https://huggingface.co/docs/huggingface_hub/guides/cli#hf-download) or specific files from a model repository:
 
 ```bash
 # Download a model
@@ -490,6 +490,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 | `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
+| `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ |
 | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |
 | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
 | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |
@@ -547,7 +548,8 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
 | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
 | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
-| `DeepseekV4ForConditionalGeneration` | DeepSeek-V4 (vision) | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-V4-Flash-Vision-Exp` | | ✅︎ |
+| `DeepseekV4ForConditionalGeneration` | DeepSeek-V4 (vision) | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-V4-Flash-Vision-Exp` | ✅︎ | ✅︎ |
+| `DiffusionGemmaForBlockDiffusion` | DiffusionGemma | T + I<sup>+</sup> + V<sup>+</sup> | `google/diffusiongemma-26B-A4B-it`, etc. | | |
 | `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + I<sup>E+</sup> | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ |
 | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I<sup>+</sup>/ V<sup>+</sup> | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ |
 | `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + I<sup>E+</sup> | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ |
@@ -601,6 +603,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ |
 | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
 | `MuseGlimmerForCausalLM`, `MuseGlimmerForConditionalGeneration` | Muse Glimmer | T + I<sup>+</sup> + V<sup>+</sup> | `meta-models/Muse-Glimmer-30B` | ✅︎ | ✅︎ |
+| `NemotronH_Nano_Omni_Reasoning_V3`, `NemotronH_Nano_VL_V2` | Nemotron Nano 3 Omni, Nemotron Nano V2 VL | T + I<sup>E+</sup> + V<sup>+</sup> + A<sup>*</sup> | `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16` | ✅︎ | |
 | `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
 | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
 | `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
@@ -700,6 +703,12 @@ Some models are supported only via the [Transformers modeling backend](#transfor
     `meta-models/Muse-Glimmer-30B-assistant` checkpoint, which vLLM serves through
     its [DFlash](../features/speculative_decoding/README.md) path.
 
+!!! note
+    Both Nemotron Nano architecture names map to the same vLLM implementation.
+    Audio inputs are only supported by the Omni variants; `NemotronH_Nano_VL_V2`
+    checkpoints such as `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` accept text,
+    image and video only.
+
 !!! note
     The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now.
     For more details, please see: <https://github.com/vllm-project/vllm/pull/4087#issuecomment-2250397630>
diff --git a/docs/serving/context_parallel_deployment.md b/docs/serving/context_parallel_deployment.md
index 3df2b06635f5..cb9a1ece71a4 100644
--- a/docs/serving/context_parallel_deployment.md
+++ b/docs/serving/context_parallel_deployment.md
@@ -20,6 +20,15 @@ Both approaches are under active development.
 
 Due to the auto-regressive nature of decoding, every decoding step needs to compute a small amount of query tokens w.r.t. a large number of key/value tokens stored in the paged KV cache. The core of decode context parallel is how to shard the KV cache across GPUs.
 
+`engine_client.vllm_config.cache_config.effective_attention_block_size` reports the
+initialized full-attention block size in tokens, including DCP: a physical block of
+16 tokens with DCP=4 represents 64 tokens. The same value is available through
+gRPC `Control.GetServerInfo.effective_attention_block_size`. Physical block-size
+fields retain their existing meaning. The new value is unavailable (`None` in
+Python, absent in gRPC) if there is no common full-attention size, the scheduler
+does not expose it, or an older engine omits it. gRPC also omits the value when
+engine replicas disagree. Partial cache events carry their own actual size.
+
 For a model with `H` kv-heads, a request with `T` tokens in the context needs to store `H * T` key/value tensors in the KV cache.
 
 1. If one GPU can hold them all, and the performance is good enough, then no parallelization is needed.
diff --git a/docs/serving/data_parallel_deployment.md b/docs/serving/data_parallel_deployment.md
index 3b92fa06a5ee..d1f9b4a4624d 100644
--- a/docs/serving/data_parallel_deployment.md
+++ b/docs/serving/data_parallel_deployment.md
@@ -126,6 +126,14 @@ vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 1 \
                   --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
 ```
 
+When using `--data-parallel-external-lb` with `--nnodes`, the data parallel rank
+can be inferred from `--node-rank` only when `--nnodes` is evenly divisible by
+`--data-parallel-size`. If it is not evenly divisible, automatic rank inference
+is rejected and `--data-parallel-rank` must be set explicitly. When multiple
+external-LB ranks are co-located on a node, assign each process a disjoint
+device set with `CUDA_VISIBLE_DEVICES` or `--device-ids`; devices are not
+partitioned automatically between processes.
+
 The coordinator process also runs in this scenario, co-located with the DP rank 0 engine.
 
 <figure markdown="1">
diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md
index 5c93ca17877c..a292a4a3a436 100644
--- a/docs/serving/online_serving/README.md
+++ b/docs/serving/online_serving/README.md
@@ -121,12 +121,7 @@ For further details on profiling vLLM, please refer to [this page](../../contrib
 
 ## Scale-Out APIs
 
-Scale-out APIs are disabled by default on `vllm serve`. The environment
-variable accepts only `0` or `1`; set `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1` to
-register the endpoints below. The
-dedicated `vllm launch render` and `vllm serve --tokens-only` modes are explicit
-opt-ins and enable their required endpoints when the variable is unset; an
-explicit value of `0` is rejected for those modes.
+Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out` to register the endpoints below. The dedicated `vllm launch render` and `vllm serve --tokens-only` modes always register their required endpoints regardless of `--enable-scale-out`.
 
 ### Tokens IN <> Tokens OUT APIs
 
@@ -135,7 +130,12 @@ explicit value of `0` is rejected for those modes.
 
 ### Renderer APIs
 
-For further details on renderer APIs, please refer to [this page](renderer.md).
+Renderer APIs preprocess completion, chat, and Responses requests without running
+inference. They handle tokenization, model-specific prompt formatting, and
+multimodal preprocessing, returning prompt token IDs, sampling parameters, and
+any processed multimodal inputs for generation.
+
+See the [renderer guide](renderer.md) for setup instructions and examples.
 
 - [Completions Render API](renderer.md) (`/v1/completions/render`)
     - Render completion requests
diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md
index d536498f5c60..0e865796f81d 100644
--- a/docs/serving/online_serving/derenderer.md
+++ b/docs/serving/online_serving/derenderer.md
@@ -6,7 +6,7 @@ This closes the loop for a token-in / token-out engine in disaggregated serving:
 
 - **GPU less post processing**: Detokenization, reasoning parsing, and tool call parsing run on the same GPU less frontend that hosts `/render`
 - **Parser parity**: The derenderer reuses vLLM's tool and reasoning parsers, so a disaggregated deployment produces the same `content`/`reasoning`/ `tool_calls` split as a standard `vllm serve` server
-- **Non-streaming**: The endpoints expect a complete `GenerateResponse` with all token IDs present and perform one-shot parsing. Streaming derender would require a separate endpoint design and is not currently supported but is in the pipeline
+- **Streaming and one-shot parsing**: Non-streaming calls send a complete `GenerateResponse` with all token IDs present and get one-shot parsing. Both endpoints also accept `stream: true`, taking one `GenerateStreamResponse` delta plus a client carried `stream_state` and returning `{chunk, stream_state}`. The chat endpoint's streaming path supports reasoning and tool call parsing while emitting the same `reasoning`/`content`/`tool_calls` deltas the generate streaming path would for the same token IDs (see [Streaming](#streaming))
 
 Both endpoints are hosted by the GPU less rendering server started with [`vllm launch render`](../../cli/launch/render.md), alongside the `/render`
 endpoints.
@@ -21,14 +21,16 @@ endpoints.
                         └─────────────── request + prompt_tokens ──┘
 ```
 
-The derender step needs more than the engine's `token_ids`. It also consumes the original `chat_request`/`completion_request` and `prompt_tokens` carried over from the render step (see [Request format](#request-format)) so the tool and reasoning parsers have the context they need.
+The derender step needs more than the engine's `token_ids`. It also consumes the original `chat_request`/`completion_request` and `prompt_tokens` carried over from the render step (see [Request format](#request-format)) so the tool and reasoning parsers have the context they need. The chat streaming path additionally needs `prompt_token_ids` (see [Streaming cost](#streaming-cost)) when a parser is configured.
 
 ## API Reference
 
 - Chat Completions Derender API (`/v1/chat/completions/derender`)
     - Post process a single `GenerateResponse` into a `ChatCompletionResponse`
+    - With `stream: true`, post process one `GenerateStreamResponse` chunk into a `ChatCompletionStreamResponse` chunk
 - Completions Derender API (`/v1/completions/derender`)
     - Post process a list of `GenerateResponse` objects (one per prompt) into a `CompletionResponse`
+    - With `stream: true`, post process one `GenerateStreamResponse` chunk into a `CompletionStreamResponse` chunk
 
 ## Request format
 
@@ -50,8 +52,69 @@ Each request wraps the engine's `GenerateResponse`(s) together with the caller m
     --8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-request"
     ```
 
+Streaming requests set `stream: true` and carry one generate chunk (`generate_chunk`) plus `stream_state` instead of a complete response.
+
+`/v1/chat/completions/derender` with `stream: true`:
+
+??? code
+
+    ```python
+    --8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-chat-stream-request"
+    ```
+
+`/v1/completions/derender` with `stream: true`:
+
+??? code
+
+    ```python
+    --8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-stream-request"
+    ```
+
+Both return `{"chunk": ..., "stream_state": ...}`. `chunk` is a `ChatCompletionStreamResponse` or `CompletionStreamResponse` and `stream_state` goes with the next call.
+
 Oversized payloads are rejected with a `400` before any `tokenizer.decode()` or parser runs.
 
+## Parser configuration
+
+The derenderer builds its tool and reasoning parsers from its own server flags plus the `chat_request` sent with each call. It can't see how the request was rendered or served anywhere else, so for its output to match a `vllm serve` server:
+
+- Start the render server with the same `--tool-call-parser`, `--reasoning-parser`, `--enable-auto-tool-choice`, `--chat-template` and `--default-chat-template-kwargs` you'd give `vllm serve` for this model. If `/render` and `/derender` run on different servers, give both the same values.
+- Send the full `chat_request` that went to `/render`, not just `messages` and `tools`. Fields like `chat_template_kwargs`, `reasoning_effort`, `tool_choice` and `include_reasoning` change how the output is parsed.
+
+A mismatch doesn't fail. The parser just splits `reasoning`, `content` and `tool_calls` differently from what `vllm serve` would return.
+
+`/v1/completions/derender` only detokenizes. It never runs tool or reasoning parsers, the same as `/v1/completions` on `vllm serve`. It only reads `skip_special_tokens` from `completion_request`.
+
+## Streaming
+
+Streaming derender mirrors `/inference/v1/generate`. Set `stream: true` in the body on the same path and the endpoint takes one generate stream chunk instead of a complete response. It answers with JSON, not SSE. Each call turns one generate chunk into one derendered chunk and the client forwards that to its own caller.
+
+The server keeps no state between calls. Everything the next call needs is in `stream_state` which the client carries:
+
+- **Echo `stream_state` back.** Leave it out (or send `null`) on the first call, then send the `stream_state` from each response with the next request, unchanged. State that doesn't add up (e.g. `output_chunk_lens` that don't sum to the number of output tokens) is rejected with a `400`.
+- **One choice per chunk.** `/inference/v1/generate` sends one choice per SSE event and each derender call accepts at most one. More than one is rejected with a `400`.
+- **`n > 1` is N streams.** Generate interleaves chunks for different choices and tags each with its `index`. Keep one `stream_state` per index and send each chunk with the state for its index. Each index gets its own `role` delta on its first chunk.
+- **Send every chunk through, including the last ones.** The finish reason usually arrives with the final tokens. A chunk with a `finish_reason` and no tokens is accepted too and flushes any buffered tool call arguments. With `stream_options: {"include_usage": true}`, generate ends with a usage only chunk (`choices: []`). Send it through derender like any other chunk to get the usage chunk for the response. `usage.prompt_tokens` is the request's `prompt_tokens` if set, otherwise the generate chunk's.
+- **Send the same context on every call.** `chat_request` and `prompt_token_ids` aren't kept between calls, so they go with every chunk including the usage chunk. Both are required when a tool or reasoning parser is configured. `prompt_token_ids` is the `token_ids` of the `GenerateRequest` returned by `/render`.
+- **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk.
+
+Streaming chunks don't carry logprobs yet, so `logprobs` on a generate chunk are dropped. The non-streaming endpoints do resolve them, including `token_id:N` placeholders.
+
+## Streaming cost
+
+Streaming derender threads a client carried `stream_state` across per-chunk calls instead of keeping session state on the server.
+
+Plain detokenization (no parser configured) carries only a small, bounded incremental decode window in `stream_state` independent of generation length.
+
+When a tool or reasoning parser is configured, parser internal state (buffered markup, reasoning/tool phase) can't be serialized. `stream_state` therefore instead carries the full `output_token_ids` seen so far plus `output_chunk_lens`, the token count of each chunk they arrived in. Each chunk rebuilds a fresh parser and replays that history through `parse_delta`, one call per original chunk, before processing the new tokens for real. For the parser path only, this means:
+
+- **Transport**: `output_token_ids` and `output_chunk_lens` round-trip in full in both directions on every call. `output_chunk_lens` has one entry per chunk, which is one per token without speculative decoding. This means O(n) bytes per chunk, O(n²) bytes over a full generation. Bounded by `max_model_len`. `prompt_token_ids` is sent in full on every call too and it isn't trimmed as `output_token_ids` grows. This means that for most of a stream it dominates the per chunk payload. A 100k token prompt with 1k tokens of output means `prompt_token_ids` is ~99% of the request body on every chunk.
+- **Compute**: replay is O(n) `parse_delta` calls per chunk (O(n²) per generation). `parse_delta` itself is O(n) for parsers that re-scan accumulated text (e.g. Hermes tool-call JSON, DeepSeek-R1 reasoning). The per-generation cost is O(n³) character work, not O(n²). This is a deliberately minimal first implementation with no caching layer.
+- The parser path also requires `prompt_token_ids` so `parse_delta` can settle whether the prompt left reasoning open or not. Since parser state can't be carried across calls, it re-scans the full prompt once per chunk.
+- Replay runs off the event loop on the renderer's executor (`renderer_num_workers`, default `1`). Size it for the expected number of concurrent parser configured streams.
+
+`output_token_ids` and `prompt_token_ids` are both bounded by `max_model_len` but callers streaming long reasoning traces through a parser configured model should expect materially more state transport and CPU cost than the plain detokenization path.
+
 ## Example
 
 The example below drives the full `render → generate → derender` round trip for a chat request against a GPU less render server (`/render`, `/derender`) and a token-in / token-out engine (`/inference/v1/generate`).
@@ -60,8 +123,7 @@ Launch the two servers first:
 
 ```bash
 vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100
-VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve \
-    meta-llama/Llama-3.2-1B-Instruct --tokens-only --port 8200
+vllm serve meta-llama/Llama-3.2-1B-Instruct --tokens-only --port 8200
 ```
 
 ```python
@@ -103,4 +165,65 @@ with httpx.Client(timeout=60.0) as client:
 print(response["choices"][0]["message"]["content"])
 ```
 
-Passing `chat_request` lets the derenderer run the configured tool and reasoning parsers. This means `response["choices"][0]["message"]` carries the same `content` / `reasoning` / `tool_calls` split a `vllm serve` server would produce. Omit `chat_request` for plain detokenization only.
+Passing `chat_request` lets the derenderer run the configured tool and reasoning parsers. This means `response["choices"][0]["message"]` carries the same `content` / `reasoning` / `tool_calls` split a `vllm serve` server would produce. `chat_request` can only be omitted for a model with no tool or reasoning parser configured. A  parser configured model rejects a missing `chat_request` with a 400 rather than silently falling back to plain detokenization.
+
+### Streaming example
+
+The same round trip, streamed. `/render` keeps `stream` and `stream_options` on the `GenerateRequest` it returns, so the generate call streams too. Each generate chunk goes through derender with the `stream_state` from the previous call (see [Streaming](#streaming)).
+
+```python
+import json
+
+import httpx
+
+MODEL = "meta-llama/Llama-3.2-1B-Instruct"
+RENDER = "http://localhost:8100"  # vllm launch render ...
+ENGINE = "http://localhost:8200"  # token-in / token-out engine
+
+chat_request = {
+    "model": MODEL,
+    "messages": [{"role": "user", "content": "What is 2+2?"}],
+    "max_tokens": 32,
+    "stream": True,
+    "stream_options": {"include_usage": True},
+}
+
+with httpx.Client(timeout=60.0) as client:
+    # 1. Render: request -> token IDs (GPU less)
+    generate_request = client.post(
+        f"{RENDER}/v1/chat/completions/render", json=chat_request
+    ).json()
+    prompt_token_ids = generate_request["token_ids"]
+
+    # 2. Generate: stream token IDs (token-in / token-out engine)
+    stream_state = None  # one per choice index when n > 1
+    with client.stream(
+        "POST", f"{ENGINE}/inference/v1/generate", json=generate_request
+    ) as generate_stream:
+        for line in generate_stream.iter_lines():
+            if not line.startswith("data: ") or line == "data: [DONE]":
+                continue
+
+            # 3. Derender: one generate chunk -> one chat completion chunk
+            derendered = client.post(
+                f"{RENDER}/v1/chat/completions/derender",
+                json={
+                    "stream": True,
+                    "model": MODEL,
+                    "generate_chunk": json.loads(line[len("data: ") :]),
+                    "stream_state": stream_state,
+                    "prompt_tokens": len(prompt_token_ids),
+                    "prompt_token_ids": prompt_token_ids,
+                    "chat_request": chat_request,
+                },
+            ).json()
+            stream_state = derendered["stream_state"]
+
+            chunk = derendered["chunk"]
+            for choice in chunk["choices"]:
+                print(choice["delta"].get("content") or "", end="", flush=True)
+            if chunk.get("usage"):
+                print(f"\n{chunk['usage']}")
+```
+
+With a tool or reasoning parser configured, `delta` also carries `reasoning` and `tool_calls`, the same deltas `/v1/chat/completions` streams for those tokens.
diff --git a/docs/serving/online_serving/openai_compatible_server.md b/docs/serving/online_serving/openai_compatible_server.md
index 22e2e9a60844..c50910040d32 100644
--- a/docs/serving/online_serving/openai_compatible_server.md
+++ b/docs/serving/online_serving/openai_compatible_server.md
@@ -199,6 +199,11 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
 
 Code example: [examples/tool_calling/openai_responses_client_with_tools.py](../../../examples/tool_calling/openai_responses_client_with_tools.py)
 
+To get prompt token IDs before generation, use
+[`/v1/responses/render`](renderer.md#get-responses-prompt-token-ids). It preprocesses
+a self-contained Responses request without inference and returns the prompt
+token IDs, so a separate `/tokenize` call is not needed.
+
 #### Extra parameters
 
 The following extra parameters in the request object are supported:
diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md
index 1f5e893570ef..9d494df22efe 100644
--- a/docs/serving/online_serving/renderer.md
+++ b/docs/serving/online_serving/renderer.md
@@ -7,16 +7,14 @@ Our renderer API is designed to disaggregate the render phase(preprocessing) and
 - Tokens-in / tokens-out engine: Make the engine a pure token-in / token-out service, decoupled from request preprocessing.
 
 The dedicated `vllm launch render` server always exposes the `/render` and
-`/derender` endpoints when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS` is unset or set to
-`1`. An explicit value of `0` conflicts with the renderer command and is
-rejected at startup.
+`/derender` endpoints.
 
 Scale-out endpoints, including `/render`, `/derender`, and
 `/inference/v1/generate`, are disabled by default on a standard inference
 server. To expose them with `vllm serve`, opt in explicitly:
 
 ```bash
-VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve <model>
+vllm serve <model> --enable-scale-out
 ```
 
 ## API Reference
@@ -28,27 +26,58 @@ VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve <model>
 - [Responses Render API](renderer.md) (`/v1/responses/render`)
     - Render a self-contained Responses request
 
+## Get Responses prompt token IDs
+
+Use `/v1/responses/render` to get prompt token IDs before choosing a model
+replica. Rendering applies prompt construction and preprocessing without running
+inference.
+
 The Responses render endpoint uses the same prompt construction as
 `/v1/responses` and returns one token-in `GenerateRequest`. It is stateless:
 inline history is supported, but `previous_response_id` is not. Callers must
 resolve stored response state and include the resulting history in the request
 before rendering.
 
+Configure the renderer and generation workers with the same model, tokenizer,
+chat template, and preprocessing options. Render the full request, including
+instructions, history, tools, and any template or truncation options, so the
+returned IDs reflect the prompt the model will receive.
+
 For multimodal requests, the `GenerateRequest` contains the model-processed
 multimodal payload, which can be substantially larger than the source image or
 video. The caller must forward that payload unchanged to the generation service
 and provision transport limits and memory accordingly.
 
+For example, start a standard inference server with scale-out endpoints enabled:
+
+```bash
+VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \
+    vllm serve meta-llama/Llama-3.1-8B-Instruct
+```
+
+Send a Responses request and use `jq` to extract the `token_ids` field:
+
 ```bash
-curl http://localhost:8000/v1/responses/render \
+curl --fail --silent --show-error http://localhost:8000/v1/responses/render \
     -H "Content-Type: application/json" \
     -d '{
         "model": "meta-llama/Llama-3.1-8B-Instruct",
         "input": "Explain prefix caching in one sentence.",
         "max_output_tokens": 32
-    }'
+    }' | jq '.token_ids'
 ```
 
+To count the rendered prompt tokens for this text request, replace the `jq`
+filter with `'.token_ids | length'`. The endpoint returns the full
+`GenerateRequest`; `jq` filters the response on the client. Keep the complete
+response if you will forward it to `/inference/v1/generate`, including any
+multimodal features.
+
+If the server has `--api-key` or `VLLM_API_KEY` configured, add
+`-H "Authorization: Bearer <api-key>"` to the request. See
+[API key authentication limitations](../../usage/security.md#api-key-authentication-limitations)
+for the existing authentication boundaries.
+
 For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md).
 
 ## Multimodal Render Features
diff --git a/docs/training/weight_transfer/base.md b/docs/training/weight_transfer/base.md
index d73c88436ae0..759c55ef6fe9 100644
--- a/docs/training/weight_transfer/base.md
+++ b/docs/training/weight_transfer/base.md
@@ -233,7 +233,7 @@ Two implementations ship with vLLM:
 | `RayVLLMWeightSyncClient(handle)` | One or more `AsyncLLM`/`LLM` Ray actors. Accepts a list and fans each call out to every handle, blocking on all of them, so a multi-actor (e.g. multi-DP) deployment is driven as one unit |
 | `HTTPVLLMWeightSyncClient(base_url, timeout=300)` | A vLLM server over the RLHF HTTP routes |
 
-Custom weight sync clients can be implement like so:
+Custom weight sync clients can be implemented like so:
 
 ```python
 class MyFrameworkWeightSyncClient:
diff --git a/docs/usage/security.md b/docs/usage/security.md
index ba730bcb2056..5d4d429c6fed 100644
--- a/docs/usage/security.md
+++ b/docs/usage/security.md
@@ -147,38 +147,38 @@ 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**, and other similar `/v2`, `/inference` 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 endpoints under the `/v1`, `/v2`, `/inference`, and `/cohere` path prefixes**. 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.
 
 ### Protected Endpoints (Require API Key)
 
-When `--api-key` is configured, the following `/v1` endpoints require Bearer token authentication:
+When `--api-key` is configured, the following endpoints require Bearer token authentication:
 
 - `/v1/models` - List available models
 - `/v1/chat/completions` - Chat completions
 - `/v1/chat/completions/batch` - Batch chat completions
-- `/v1/chat/completions/render` - Render chat completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
-- `/v1/chat/completions/derender` - Derender chat completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
+- `/v1/chat/completions/render` - Render chat completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
+- `/v1/chat/completions/derender` - Derender chat completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
 - `/v1/completions` - Text completions
-- `/v1/completions/render` - Render completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
-- `/v1/completions/derender` - Derender completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
+- `/v1/completions/render` - Render completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
+- `/v1/completions/derender` - Derender completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
 - `/v1/embeddings` - Generate embeddings
 - `/v1/audio/transcriptions` - Audio transcription
 - `/v1/audio/translations` - Audio translation
 - `/v1/messages` - Anthropic-compatible messages API
-- `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
+- `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
 - `/v1/messages/count_tokens` - Count tokens for Anthropic messages
 - `/v1/responses` - Create a response
-- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
+- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render` unless explicitly disabled)
 - `/v1/responses/{response_id}` - Retrieve a response
 - `/v1/responses/{response_id}/cancel` - Cancel a response
 - `/v1/score` - Scoring API
 - `/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 (available when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or with `--tokens-only` unless explicitly disabled)
-- `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires both `VLLM_ENABLE_COHERE_API=1` and `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`)
+- `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`)
+- `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires `VLLM_ENABLE_COHERE_API=1` and either `--enable-scale-out` or `--tokens-only`)
 - `/v2/embed` - Cohere Embed API
 - `/v2/rerank` - Cohere Rerank API
 
@@ -206,11 +206,11 @@ The following endpoints **do not require authentication** even when `--api-key`
 - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
 - `/update_weights` - Update model weights (can alter model behavior)
 - `/get_world_size` - Get distributed world size
-- `/abort_requests` - Abort in-flight requests (available with `--tokens-only` unless `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0`)
+- `/abort_requests` - Abort in-flight requests (available with `--tokens-only`)
 
 **Utility endpoints:**
 
-- `/tokenize` - Tokenize text (not disabled by `VLLM_ENABLE_SCALE_OUT_ENDPOINTS`)
+- `/tokenize` - Tokenize text (not gated by `--enable-scale-out`)
 - `/detokenize` - Detokenize tokens
 - `/health` - Health check
 - `/ping` - SageMaker health check
@@ -249,7 +249,7 @@ These endpoints are only available when profiling is enabled and should only be
 
 An attacker who can reach the vLLM HTTP server can:
 
-1. **Bypass authentication** by using non-`/v1` endpoints like `/invocations`, `/inference/v1/generate`, `/generative_scoring`, `/pooling`, `/classify`, `/score`, or `/rerank` to run arbitrary inference without credentials
+1. **Bypass authentication** by using endpoints outside the protected path prefixes, such as `/invocations`, `/generative_scoring`, `/pooling`, `/classify`, `/score`, or `/rerank`, to run arbitrary inference without credentials
 2. **Cause denial of service** by calling `/pause`, `/scale_elastic_ep`, or `/abort_requests` without a token
 3. **Access operational controls** to manipulate server state (e.g., pausing generation, updating model weights via `/update_weights`)
 4. **If `--enable-tokenizer-info-endpoint` is set:** Access sensitive tokenizer configuration including chat templates, which may reveal prompt engineering strategies or other implementation details
@@ -346,7 +346,7 @@ vLLM supports loading out-of-tree HTTP routes via the `vllm.endpoint_plugins` en
 
 1. **Only allowlist plugins you trust.** Set `VLLM_PLUGINS` to the exact plugin names you intend to run and never wildcard or copy an allowlist between deployments without reviewing what each named plugin does.
 2. **Audit routes before deploying.** A plugin's `attach_router` can add routes under any path, including ones that duplicate existing `/v1/*` paths. There is currently no route conflict enforcement (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)), so a malicious or buggy plugin can **shadow a core route** and silently replace its behavior. Prefer plugins that namespace their routes under a distinct prefix (e.g. `/plugins/<plugin-name>/...`) instead of reusing `/v1/...` and review `app.routes` after startup if you need certainty about what is actually being served.
-3. **Treat plugin routes like any other unauthenticated by default surface.** `--api-key` only protects the `/v1`, `/v2`, and `/inference` path prefixes (see [API Key Authentication Limitations](#api-key-authentication-limitations)). A plugin route outside those prefixes is unauthenticated unless the plugin implements its own authentication. Deploy behind a reverse proxy that allowlists only the plugin routes you intend to expose externally.
+3. **Treat plugin routes like any other unauthenticated by default surface.** `--api-key` only protects the `/v1`, `/v2`, `/inference`, and `/cohere` path prefixes (see [API Key Authentication Limitations](#api-key-authentication-limitations)). A plugin route outside those prefixes is unauthenticated unless the plugin implements its own authentication. Deploy behind a reverse proxy that allowlists only the plugin routes you intend to expose externally.
 4. **Remember the `vllm.general_plugins` pairing.** A plugin that also needs new engine side behavior ships that half separately via `vllm.general_plugins` which loads in every worker process under the default (load all unless restricted) posture. Allowlisting the endpoint plugin does not by itself restrict its paired engine side plugin. Need to review both.
 
 ## gRPC Interface
diff --git a/docs/usage/troubleshooting.md b/docs/usage/troubleshooting.md
index dc1cd89f8209..bbd4109a4196 100644
--- a/docs/usage/troubleshooting.md
+++ b/docs/usage/troubleshooting.md
@@ -8,7 +8,7 @@ This document outlines some troubleshooting strategies you can consider. If you
 ## Hangs downloading a model
 
 If the model isn't already downloaded to disk, vLLM will download it from the internet which can take time and depend on your internet connection.
-It's recommended to download the model first using the [huggingface-cli](https://huggingface.co/docs/huggingface_hub/en/guides/cli) and passing the local path to the model to vLLM. This way, you can isolate the issue.
+It's recommended to download the model first using the [`hf` CLI](https://huggingface.co/docs/huggingface_hub/en/guides/cli) and passing the local path to the model to vLLM. This way, you can isolate the issue.
 
 ## Hangs loading a model from disk
 
diff --git a/examples/applications/api_server/server.py b/examples/applications/api_server/server.py
index e56a7931f666..36474e210184 100644
--- a/examples/applications/api_server/server.py
+++ b/examples/applications/api_server/server.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-NOTE: This API server is used only for demonstrating usage of AsyncEngine
+"""NOTE: This API server is used only for demonstrating usage of AsyncEngine
 and simple performance benchmarks. It is not intended for production use.
 For production use, we recommend using our OpenAI compatible server.
 We are also not going to accept PRs modifying this file.
diff --git a/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py b/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py
index d60dbf4d753a..eb18b41a5ad2 100644
--- a/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py
+++ b/examples/applications/chatbot/streamlit_openai_chatbot_webserver.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-vLLM Chat Assistant - A Streamlit Web Interface
+"""vLLM Chat Assistant - A Streamlit Web Interface.
 
 A streamlined chat interface that quickly integrates
 with vLLM API server.
@@ -78,6 +77,7 @@ def create_new_chat_session():
         - current_session: Sets to new session ID
         - active_session: Sets to new session ID
         - messages: Resets to empty list
+
     """
     session_id = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     st.session_state.sessions[session_id] = []
@@ -101,6 +101,7 @@ def switch_to_chat_session(session_id):
         - current_session: Updated to specified session_id
         - active_session: Updated to specified session_id
         - messages: Loaded from sessions[session_id]
+
     """
     st.session_state.current_session = session_id
     st.session_state.active_session = session_id
@@ -134,6 +135,7 @@ def get_llm_response(messages, model, reason, content_ph=None, reasoning_ph=None
     Note:
         The function uses streamlit placeholders for live updates.
         When reason=True, the reasoning process appears above the response.
+
     """
     full_text = ""
     think_text = ""
@@ -256,6 +258,7 @@ def server_supports_reasoning():
 
     Returns:
         bool: True if the model supports reasoning, False otherwise
+
     """
     resp = client.chat.completions.create(
         model=model,
diff --git a/examples/applications/rag/retrieval_augmented_generation_with_langchain.py b/examples/applications/rag/retrieval_augmented_generation_with_langchain.py
index d9a4cadb036e..1681ea146665 100644
--- a/examples/applications/rag/retrieval_augmented_generation_with_langchain.py
+++ b/examples/applications/rag/retrieval_augmented_generation_with_langchain.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Retrieval Augmented Generation (RAG) Implementation with Langchain
+"""Retrieval Augmented Generation (RAG) Implementation with Langchain
 ==================================================================
 
 This script demonstrates a RAG implementation using LangChain, Milvus
@@ -35,6 +34,7 @@
     - Ensure both vLLM services are running before executing
     - Default ports: 8000 (embedding), 8001 (chat)
     - First run may take time to download models
+
 """
 
 import argparse
@@ -52,9 +52,7 @@
 
 
 def load_and_split_documents(config: dict[str, Any]):
-    """
-    Load and split documents from web URL
-    """
+    """Load and split documents from web URL."""
     try:
         loader = WebBaseLoader(web_paths=(config["url"],))
         docs = loader.load()
@@ -70,9 +68,7 @@ def load_and_split_documents(config: dict[str, Any]):
 
 
 def init_vectorstore(config: dict[str, Any], documents: list[Document]):
-    """
-    Initialize vector store with documents
-    """
+    """Initialize vector store with documents."""
     return Milvus.from_documents(
         documents=documents,
         embedding=OpenAIEmbeddings(
@@ -86,9 +82,7 @@ def init_vectorstore(config: dict[str, Any], documents: list[Document]):
 
 
 def init_llm(config: dict[str, Any]):
-    """
-    Initialize llm
-    """
+    """Initialize llm."""
     return ChatOpenAI(
         model=config["chat_model"],
         openai_api_key=config["vllm_api_key"],
@@ -97,9 +91,7 @@ def init_llm(config: dict[str, Any]):
 
 
 def get_qa_prompt():
-    """
-    Get question answering prompt template
-    """
+    """Get question answering prompt template."""
     template = """You are an assistant for question-answering tasks.
 Use the following pieces of retrieved context to answer the question.
 If you don't know the answer, just say that you don't know.
@@ -112,16 +104,12 @@ def get_qa_prompt():
 
 
 def format_docs(docs: list[Document]):
-    """
-    Format documents for prompt
-    """
+    """Format documents for prompt."""
     return "\n\n".join(doc.page_content for doc in docs)
 
 
 def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate):
-    """
-    Set up question answering chain
-    """
+    """Set up question answering chain."""
     return (
         {
             "context": retriever | format_docs,
@@ -134,9 +122,7 @@ def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate):
 
 
 def get_parser() -> argparse.ArgumentParser:
-    """
-    Parse command line arguments
-    """
+    """Parse command line arguments."""
     parser = argparse.ArgumentParser(description="RAG with vLLM and langchain")
 
     # Add command line arguments
@@ -192,10 +178,7 @@ def get_parser() -> argparse.ArgumentParser:
 
 
 def init_config(args: Namespace):
-    """
-    Initialize configuration settings from command line arguments
-    """
-
+    """Initialize configuration settings from command line arguments."""
     return {
         "vllm_api_key": args.vllm_api_key,
         "vllm_embedding_endpoint": args.vllm_embedding_endpoint,
diff --git a/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py b/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py
index be4796acd1b6..0df32206bf07 100644
--- a/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py
+++ b/examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-RAG (Retrieval Augmented Generation) Implementation with LlamaIndex
+"""RAG (Retrieval Augmented Generation) Implementation with LlamaIndex
 ================================================================
 
 This script demonstrates a RAG system using:
@@ -35,6 +34,7 @@
     - Ensure both vLLM services are running before executing
     - Default ports: 8000 (embedding), 8001 (chat)
     - First run may take time to download models
+
 """
 
 import argparse
@@ -50,7 +50,7 @@
 
 
 def init_config(args: Namespace):
-    """Initialize configuration with command line arguments"""
+    """Initialize configuration with command line arguments."""
     return {
         "url": args.url,
         "embedding_model": args.embedding_model,
@@ -66,12 +66,12 @@ def init_config(args: Namespace):
 
 
 def load_documents(url: str) -> list:
-    """Load and process web documents"""
+    """Load and process web documents."""
     return SimpleWebPageReader(html_to_text=True).load_data([url])
 
 
 def setup_models(config: dict[str, Any]):
-    """Configure embedding and chat models"""
+    """Configure embedding and chat models."""
     Settings.embed_model = OpenAILikeEmbedding(
         api_base=config["embedding_endpoint"],
         api_key=config["vllm_api_key"],
@@ -96,14 +96,14 @@ def setup_models(config: dict[str, Any]):
 
 
 def setup_vector_store(db_path: str) -> MilvusVectorStore:
-    """Initialize vector store"""
+    """Initialize vector store."""
     sample_emb = Settings.embed_model.get_text_embedding("test")
     print(f"Embedding dimension: {len(sample_emb)}")
     return MilvusVectorStore(uri=db_path, dim=len(sample_emb), overwrite=True)
 
 
 def create_index(documents: list, vector_store: MilvusVectorStore):
-    """Create document index"""
+    """Create document index."""
     storage_context = StorageContext.from_defaults(vector_store=vector_store)
     return VectorStoreIndex.from_documents(
         documents,
@@ -112,13 +112,13 @@ def create_index(documents: list, vector_store: MilvusVectorStore):
 
 
 def query_document(index: VectorStoreIndex, question: str, top_k: int):
-    """Query document with given question"""
+    """Query document with given question."""
     query_engine = index.as_query_engine(similarity_top_k=top_k)
     return query_engine.query(question)
 
 
 def get_parser() -> argparse.ArgumentParser:
-    """Parse command line arguments"""
+    """Parse command line arguments."""
     parser = argparse.ArgumentParser(description="RAG with vLLM and LlamaIndex")
 
     # Add command line arguments
diff --git a/examples/deployment/async_llm_streaming.py b/examples/deployment/async_llm_streaming.py
index ef69089a0460..bf9b36d05cb1 100644
--- a/examples/deployment/async_llm_streaming.py
+++ b/examples/deployment/async_llm_streaming.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Simple example demonstrating streaming offline inference with AsyncLLM (V1 engine).
+"""Simple example demonstrating streaming offline inference with AsyncLLM (V1 engine).
 
 This script shows the core functionality of vLLM's AsyncLLM engine for streaming
 token-by-token output in offline inference scenarios. It demonstrates DELTA mode
@@ -20,8 +19,7 @@
 
 
 async def stream_response(engine: AsyncLLM, prompt: str, request_id: str) -> None:
-    """
-    Stream response from AsyncLLM and display tokens as they arrive.
+    """Stream response from AsyncLLM and display tokens as they arrive.
 
     This function demonstrates the core streaming pattern:
     1. Create SamplingParams with DELTA output kind
diff --git a/examples/deployment/llm_engine_example.py b/examples/deployment/llm_engine_example.py
index d9215255a808..0d25a157b416 100644
--- a/examples/deployment/llm_engine_example.py
+++ b/examples/deployment/llm_engine_example.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file demonstrates using the `LLMEngine`
+"""This file demonstrates using the `LLMEngine`
 for processing prompts with various sampling parameters.
 """
 
diff --git a/examples/disaggregated/disaggregated_encoder/README.md b/examples/disaggregated/disaggregated_encoder/README.md
index bd2a46abbe52..ec10f23fc074 100644
--- a/examples/disaggregated/disaggregated_encoder/README.md
+++ b/examples/disaggregated/disaggregated_encoder/README.md
@@ -106,6 +106,98 @@ If you enable prefill instance (`--prefill-servers-urls` not disabled), you will
 | `--decode-servers-urls` | Comma-separated list of decode endpoints. Non-stream and stream paths both round-robin over this list. |
 | `--host`, `--port` | Bind address for the proxy itself (defaults: `0.0.0.0:8000`). |
 
+### Dynamic registration
+
+Alternatively, let the external launcher register ready instances over HTTP.
+No vLLM configuration changes or worker registration threads are required.
+Set `ADMIN_API_KEY` on the proxy and supply it as `X-API-Key` for registration
+and removal. These are trusted control-plane APIs; do not expose them publicly.
+
+```bash
+export ADMIN_API_KEY="your-admin-key"
+python disagg_epd_proxy.py --port 8000 --dynamic-registration
+```
+
+Omit the static server URL flags and register every stage. The roles determine
+the topology: `encode` + `prefill_decode` gives E+PD; `encode` + `prefill` +
+`decode` gives E+P+D. Standalone `decode` always requires an available `prefill`,
+even when D registers first or all P instances go offline; otherwise requests
+return `503`. The proxy rejects mixing combined PD with standalone P/D,
+including unhealthy instances still in the registry. Explicitly remove the
+old topology's P/D or PD registrations before switching topologies.
+
+After each instance is ready, the launcher registers its reachable HTTP URL:
+
+```bash
+curl --fail-with-body http://proxy-host:8000/instances \
+    -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
+    -d '{"role":"encode","url":"http://e-host:8001"}'
+curl --fail-with-body http://proxy-host:8000/instances \
+    -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
+    -d '{"role":"prefill_decode","url":"http://pd-host:8002"}'
+```
+
+For E+P+D, register P with `role: "prefill"` and D with `role: "decode"`.
+Example and NIXL EC connectors need no additional registration fields.
+For Mooncake, register `ec_zmq_addrs` on the **EC consumer** (`prefill_decode`
+or `prefill`), using its configured `ec_ip` and `ec_port`. Keep the fixed-port
+layout: supply the TP-rank-0 address for each DP replica, in DP-rank order;
+replica `r` uses `ec_port + r * tensor_parallel_size`. The connector discovers
+the remaining TP ranks itself. For example, DP=2, TP=2, `ec_port=19019`:
+
+```json
+{
+  "role": "prefill_decode",
+  "url": "http://pd-host:8002",
+  "dp_size": 2,
+  "ec_zmq_addrs": ["tcp://pd-host:19019", "tcp://pd-host:19021"]
+}
+```
+
+The proxy selects one consumer replica and uses it for both the encoder push
+and the consumer HTTP request. Standalone D does not need EC control addresses.
+Port allocation and avoiding collisions remain the launcher's responsibility.
+
+Inspect or remove instances without restarting the proxy:
+
+```bash
+curl http://proxy-host:8000/instances
+curl --fail-with-body -X DELETE \
+    'http://proxy-host:8000/instances?url=http://e-host:8001' \
+    -H "X-API-Key: $ADMIN_API_KEY"
+```
+
+Registration is idempotent. The proxy probes registered instances every
+`--probe-interval` seconds (default 5), with `--probe-timeout` seconds per probe
+(default 2). After `--fail-threshold` consecutive failures (default 3), it stops
+sending new requests to that instance. Healthy instances rejoin automatically;
+unreachable ones are forgotten after `--evicted-ttl` seconds (default 900;
+`0` retains them indefinitely). Removal stops new routing; already routed
+requests retain their selected endpoints, so drain requests before stopping
+the instance. The launcher must re-register instances after a proxy restart.
+
+The example launch scripts also support this flow: export `ADMIN_API_KEY` and
+set `DYNAMIC_REGISTRATION=1` when running `disagg_1e1pd_example.sh` or
+`disagg_1e1p1d_example.sh`. Their default remains static routing.
+
+### Static configuration
+
+The proxy batches images from the same user request assigned to the same encoder.
+Set the proxy environment variable `ENCODER_MAX_BATCH_SIZE` to limit the number
+of images per encoder subrequest. It defaults to `0` (unlimited); `1` sends each
+image separately. Audio and video remain separate subrequests.
+
+For encoders with a smaller `--limit-mm-per-prompt` image limit than P/PD, set
+`ENCODER_MAX_BATCH_SIZE` no higher than the smallest encoder image limit.
+For example, for encoders configured with `--limit-mm-per-prompt '{"image": 2}'`:
+
+```bash
+ENCODER_MAX_BATCH_SIZE=2 python disagg_epd_proxy.py \
+    --encode-servers-urls "http://e1:8001,http://e2:8002" \
+    --prefill-servers-urls disable \
+    --decode-servers-urls "http://pd1:8003"
+```
+
 Example usage:
 For E + PD setup:
 
@@ -121,6 +213,6 @@ For E + P + D setup:
 ```bash
 $ python disagg_encoder_proxy.py \
       --encode-servers-urls "http://e1:8001,http://e2:8001" \
-      --prefill-servers-urls "http://p1:8003,http://p2:8004" \ 
+      --prefill-servers-urls "http://p1:8003,http://p2:8004" \
       --decode-servers-urls "http://d1:8005,http://d2:8006"
 ```
diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh
index cb51333e908d..57392a0408c8 100644
--- a/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh
+++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh
@@ -189,17 +189,34 @@ wait_for_server "$DECODE_PORT"
 ###############################################################################
 # Proxy
 ###############################################################################
+PROXY_ARGS=()
+if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
+    : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}"
+    export ADMIN_API_KEY
+    PROXY_ARGS=(--dynamic-registration)
+else
+    PROXY_ARGS=(
+        --encode-servers-urls "http://localhost:$ENCODE_PORT"
+        --prefill-servers-urls "http://localhost:$PREFILL_PORT"
+        --decode-servers-urls "http://localhost:$DECODE_PORT"
+    )
+fi
 python disagg_epd_proxy.py \
     --host "0.0.0.0" \
     --port "$PROXY_PORT" \
-    --encode-servers-urls "http://localhost:$ENCODE_PORT" \
-    --prefill-servers-urls "http://localhost:$PREFILL_PORT" \
-    --decode-servers-urls "http://localhost:$DECODE_PORT" \
+    "${PROXY_ARGS[@]}" \
     >"${PROXY_LOG}" 2>&1 &
 
 PIDS+=($!)
 
 wait_for_server "$PROXY_PORT"
+if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
+    for endpoint in "encode:$ENCODE_PORT" "prefill:$PREFILL_PORT" "decode:$DECODE_PORT"; do
+        curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \
+            -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
+            -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}"
+    done
+fi
 echo "All services are up!"
 
 ###############################################################################
@@ -238,4 +255,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
 
 # cleanup
 echo "cleanup..."
-cleanup
\ No newline at end of file
+cleanup
diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh
index 4bf60ea66d61..44e43b8d7f6e 100644
--- a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh
+++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh
@@ -152,17 +152,33 @@ wait_for_server "$PREFILL_DECODE_PORT"
 ###############################################################################
 # Proxy
 ###############################################################################
+PROXY_ARGS=()
+if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
+    : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}"
+    export ADMIN_API_KEY
+    PROXY_ARGS=(--dynamic-registration)
+else
+    PROXY_ARGS=(
+        --encode-servers-urls "http://localhost:$ENCODE_PORT"
+        --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT"
+    )
+fi
 python disagg_epd_proxy.py \
     --host "0.0.0.0" \
     --port "$PROXY_PORT" \
-    --encode-servers-urls "http://localhost:$ENCODE_PORT" \
-    --prefill-servers-urls "disable" \
-    --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
+    "${PROXY_ARGS[@]}" \
     >"${PROXY_LOG}" 2>&1 &
 
 PIDS+=($!)
 
 wait_for_server "$PROXY_PORT"
+if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
+    for endpoint in "encode:$ENCODE_PORT" "prefill_decode:$PREFILL_DECODE_PORT"; do
+        curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \
+            -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
+            -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}"
+    done
+fi
 echo "All services are up!"
 
 ###############################################################################
@@ -201,4 +217,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
 
 # cleanup
 echo "cleanup..."
-cleanup
\ No newline at end of file
+cleanup
diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
index ace44558287e..4715a15f8679 100644
--- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
@@ -1,18 +1,18 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-disagg_encoder_proxy.py
+"""OpenAI-compatible proxy for E+PD and E+P+D disaggregation.
 
-Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two
-clusters:
-  • encode  (multimodal feature extraction)
-  • decode  (language-model inference)
+Use static server URLs, or --dynamic-registration for launcher-managed HTTP
+registration of E, P/PD and D instances. Both modes route OpenAI-compatible
+``/v1/chat/completions`` requests through encoder and inference clusters and
+share encoder fan-out, metadata-only rewriting, prefill, and retry/streaming
+forwarding.
 
 For MM input we:
     1. Extract *every* image/audio/video item.
-    2. Fire N concurrent requests to the encoder cluster
-       (one request per item, with **all text removed**).
+    2. Send concurrent encoder requests with all text removed,
+       grouping images assigned to the same encoder.
     3. Wait for all of them to succeed.
     4. Forward the *original* request to a decode server.
 """
@@ -21,21 +21,27 @@
 
 import argparse
 import asyncio
+import enum
 import hashlib
 import itertools
 import json
 import logging
 import os
 import random
+import secrets
 import time
 import uuid
 from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass, field
 from typing import Any
 
 import aiohttp
+import msgspec
 import uvicorn
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi import Depends, FastAPI, Header, HTTPException, Request
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from pydantic import AnyHttpUrl, BaseModel, Field, model_validator
 
 ###############################################################################
 # FastAPI app & global state
@@ -44,6 +50,340 @@
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
 logger = logging.getLogger("proxy")
 
+DEFAULT_PROBE_INTERVAL = 5.0
+DEFAULT_PROBE_TIMEOUT = 2.0
+DEFAULT_FAIL_THRESHOLD = 3
+# Stop probing an instance that has been down this long. 0 probes forever,
+# which is what a cluster that restarts instances in place wants.
+DEFAULT_EVICTED_TTL = 900.0
+
+
+class InstanceRole(str, enum.Enum):
+    ENCODE = "encode"
+    PREFILL = "prefill"
+    DECODE = "decode"
+    PREFILL_DECODE = "prefill_decode"
+
+
+@dataclass
+class InstanceRecord:
+    """One registered instance.
+
+    Attributes:
+        role: Which stage this instance serves.
+        url: Base OpenAI-compatible URL, e.g. ``http://host:8000``.
+        ec_zmq_addrs: Mooncake TP-rank-0 control addresses, one per DP replica,
+            supplied by the launcher from the consumer's fixed port config.
+        dp_size: Data-parallel replicas behind `url`, so the proxy can pick
+            a replica and name the same one to both halves of a request.
+
+    """
+
+    role: InstanceRole
+    url: str
+    ec_zmq_addrs: list[str] = field(default_factory=list)
+    dp_size: int = 1
+
+
+class InstanceRegistry:
+    def __init__(
+        self,
+        probe_interval: float = DEFAULT_PROBE_INTERVAL,
+        probe_timeout: float = DEFAULT_PROBE_TIMEOUT,
+        fail_threshold: int = DEFAULT_FAIL_THRESHOLD,
+        evicted_ttl: float = DEFAULT_EVICTED_TTL,
+    ):
+        self._probe_interval = probe_interval
+        self._probe_timeout = probe_timeout
+        self._fail_threshold = fail_threshold
+        self._evicted_ttl = evicted_ttl
+
+        self._live: dict[str, InstanceRecord] = {}
+        self._evicted: dict[str, InstanceRecord] = {}
+        self._evicted_since: dict[str, float] = {}
+        self._fail_counts: dict[str, int] = {}
+        # One cursor per role, only ever incremented. Rebuilding it whenever
+        # the roster changes -- what an `itertools.cycle` over a mutable list
+        # forces -- restarts every fan-out at the first instance and hot-spots
+        # it after each registration.
+        self._cursors: dict[InstanceRole, int] = {role: 0 for role in InstanceRole}
+        self._replica_cursors: dict[str, int] = {}
+        self._probe_task: asyncio.Task | None = None
+
+    def register(self, record: InstanceRecord) -> bool:
+        """Add or refresh an instance without overriding its health status."""
+        roles = {
+            other.role
+            for other in itertools.chain(self._live.values(), self._evicted.values())
+        }
+        split_roles = {InstanceRole.PREFILL, InstanceRole.DECODE}
+        if (record.role is InstanceRole.PREFILL_DECODE and roles & split_roles) or (
+            record.role in split_roles and InstanceRole.PREFILL_DECODE in roles
+        ):
+            raise ValueError("Cannot mix prefill_decode with standalone prefill/decode")
+        key = record.url
+        previous = self._live.get(key) or self._evicted.get(key)
+        if previous is not None:
+            if previous == record:
+                return False
+            if previous.role is not record.role:
+                raise ValueError("Unregister the instance before changing its role")
+            target = self._live if key in self._live else self._evicted
+            target[key] = record
+            return False
+        self._live[key] = record
+        logger.info("Registered instance %s: %s", record.role.value, record.url)
+        return True
+
+    def unregister(self, url: str) -> bool:
+        """Drop an instance for good, so a probe cannot bring it back."""
+        key = url
+        if key not in self._live and key not in self._evicted:
+            return False
+        self._live.pop(key, None)
+        self._evicted.pop(key, None)
+        self._evicted_since.pop(key, None)
+        self._fail_counts.pop(key, None)
+        self._replica_cursors.pop(key, None)
+        logger.info("Unregistered instance: %s", url)
+        return True
+
+    def instances(self, role: InstanceRole) -> list[InstanceRecord]:
+        return [record for record in self._live.values() if record.role is role]
+
+    def urls(self, role: InstanceRole) -> list[str]:
+        return [record.url for record in self.instances(role)]
+
+    def pick(self, role: InstanceRole) -> InstanceRecord | None:
+        """Take the next instance of `role` in round-robin order."""
+        picked = self.pick_many(role, 1)
+        return picked[0] if picked else None
+
+    def pick_many(self, role: InstanceRole, count: int) -> list[InstanceRecord]:
+        """Take `count` instances, continuing the rotation across calls.
+
+        A multimodal request fans out one encoder request per item, so the
+        assignment has to be contiguous with the previous request's rather
+        than restart at the first instance every time.
+        """
+        alive = self.instances(role)
+        if not alive or count <= 0:
+            return []
+        start = self._cursors[role]
+        self._cursors[role] = start + count
+        return [alive[(start + offset) % len(alive)] for offset in range(count)]
+
+    def next_replica(self, record: InstanceRecord) -> int:
+        """Take the next data-parallel replica of `record`, round-robin.
+
+        The encoder pushes to one replica's receive channel, so the request
+        has to run on that same replica; the caller names the rank to both
+        halves.
+        """
+        if record.dp_size <= 1:
+            return 0
+        cursor = self._replica_cursors.get(record.url, 0)
+        self._replica_cursors[record.url] = cursor + 1
+        return cursor % record.dp_size
+
+    def status(self) -> dict[str, Any]:
+        return {
+            role.value: {
+                "live": [record.url for record in self.instances(role)],
+                "evicted": [
+                    record.url
+                    for record in self._evicted.values()
+                    if record.role is role
+                ],
+            }
+            for role in InstanceRole
+        }
+
+    def start_probing(self) -> None:
+        if self._probe_task is None and self._probe_interval > 0:
+            self._probe_task = asyncio.create_task(self._probe_loop())
+
+    async def stop_probing(self) -> None:
+        if self._probe_task is None:
+            return
+        self._probe_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await self._probe_task
+        self._probe_task = None
+
+    async def _probe_loop(self) -> None:
+        timeout = aiohttp.ClientTimeout(total=self._probe_timeout)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            while True:
+                await asyncio.sleep(self._probe_interval)
+                try:
+                    await self._probe_once(session)
+                except Exception:
+                    logger.exception("EPD registry probe round failed")
+
+    async def _probe_once(self, session: aiohttp.ClientSession) -> None:
+        targets = list(self._live.values()) + list(self._evicted.values())
+        if not targets:
+            return
+        results = await asyncio.gather(
+            *(self._probe(session, record.url) for record in targets),
+            return_exceptions=True,
+        )
+        now = time.monotonic()
+        for record, healthy in zip(targets, results):
+            current = self._live.get(record.url) or self._evicted.get(record.url)
+            if current is not record:
+                continue
+            if healthy is True:
+                self._on_probe_success(record)
+            else:
+                self._on_probe_failure(record, now)
+        self._drop_expired(now)
+
+    async def _probe(self, session: aiohttp.ClientSession, url: str) -> bool:
+        async with session.get(f"{url}/health") as resp:
+            return resp.status == 200
+
+    def _on_probe_success(self, record: InstanceRecord) -> None:
+        key = record.url
+        self._fail_counts.pop(key, None)
+        if key in self._evicted:
+            self._evicted.pop(key, None)
+            self._evicted_since.pop(key, None)
+            self._live[key] = record
+            logger.info(
+                "Instance %s (%s) is healthy again; routing resumed",
+                record.url,
+                record.role.value,
+            )
+
+    def _on_probe_failure(self, record: InstanceRecord, now: float) -> None:
+        key = record.url
+        if key not in self._live:
+            # Either already evicted, or unregistered while this probe was in
+            # flight. A round snapshots its targets and then awaits, so a
+            # removal inside that window would otherwise be undone by a result
+            # describing a registry that no longer exists.
+            return
+        failures = self._fail_counts.get(key, 0) + 1
+        self._fail_counts[key] = failures
+        if failures < self._fail_threshold:
+            return
+        self._live.pop(key, None)
+        self._evicted[key] = record
+        self._evicted_since[key] = now
+        logger.warning(
+            "Instance %s (%s) failed %d consecutive probes; stopped routing "
+            "to it. It rejoins on its own once it responds again.",
+            record.url,
+            record.role.value,
+            failures,
+        )
+
+    def _drop_expired(self, now: float) -> None:
+        if self._evicted_ttl <= 0:
+            return
+        for key, since in list(self._evicted_since.items()):
+            if now - since < self._evicted_ttl:
+                continue
+            record = self._evicted.pop(key, None)
+            self._evicted_since.pop(key, None)
+            self._fail_counts.pop(key, None)
+            self._replica_cursors.pop(key, None)
+            if record is not None:
+                logger.warning("Instance %s stayed down; forgetting it", record.url)
+
+
+class InstanceRegistration(BaseModel):
+    role: InstanceRole
+    url: AnyHttpUrl
+    ec_zmq_addrs: list[str] = Field(default_factory=list)
+    dp_size: int = Field(default=1, ge=1)
+
+    @model_validator(mode="after")
+    def validate_consumer_addresses(self):
+        if self.ec_zmq_addrs and len(self.ec_zmq_addrs) != self.dp_size:
+            raise ValueError("Provide one Mooncake control address per DP replica")
+        return self
+
+
+def require_admin_key(x_api_key: str = Header(default="")) -> None:
+    expected = os.getenv("ADMIN_API_KEY", "")
+    if not expected or not secrets.compare_digest(x_api_key, expected):
+        raise HTTPException(403, "Invalid admin API key")
+
+
+@dataclass
+class EPDProxyConfig:
+    probe_interval: float = 5.0
+    probe_timeout: float = 2.0
+    fail_threshold: int = 3
+    evicted_ttl: float = 900.0
+
+
+@dataclass
+class _Route:
+    """The instances one request was assigned to.
+
+    Attributes:
+        consumer: The stage that receives the embedding, if any stage does.
+        consumer_zmq: The receive address named to the encoders.
+        dp_rank: Which replica of `consumer` was named, so the request can be
+            pinned to it.
+
+    """
+
+    encoder_urls: list[str]
+    prefill: InstanceRecord | None
+    decode: InstanceRecord
+    consumer: InstanceRecord | None = None
+    consumer_zmq: str | None = None
+    dp_rank: int | None = None
+
+
+class EPDProxy:
+    def __init__(self, registry: InstanceRegistry):
+        self.registry = registry
+
+    # ---------------------------------------------------------------- #
+    # Routing                                                          #
+    # ---------------------------------------------------------------- #
+    def route(self, num_items: int) -> _Route:
+        decode = self.registry.pick(InstanceRole.PREFILL_DECODE) or self.registry.pick(
+            InstanceRole.DECODE
+        )
+        if decode is None:
+            raise HTTPException(
+                status_code=503, detail="No decode instance is registered"
+            )
+        prefill = (
+            self.registry.pick(InstanceRole.PREFILL)
+            if decode.role is InstanceRole.DECODE
+            else None
+        )
+        if decode.role is InstanceRole.DECODE and prefill is None:
+            raise HTTPException(
+                status_code=503, detail="No prefill instance is registered"
+            )
+        encoders = self.registry.urls(InstanceRole.ENCODE)
+        if num_items and not encoders:
+            raise HTTPException(
+                status_code=503, detail="No encode instance is registered"
+            )
+        route = _Route(encoder_urls=encoders, prefill=prefill, decode=decode)
+        self._name_consumer(route)
+        return route
+
+    def _name_consumer(self, route: _Route) -> None:
+        """Pin the EC consumer replica and, for push connectors, its endpoint."""
+        candidate = route.prefill or route.decode
+        route.consumer = candidate
+        rank = self.registry.next_replica(candidate)
+        route.dp_rank = rank if candidate.dp_size > 1 else None
+        if candidate.ec_zmq_addrs:
+            route.consumer_zmq = candidate.ec_zmq_addrs[rank]
+
+
 app = FastAPI()
 encode_session: aiohttp.ClientSession | None = None
 prefill_session: aiohttp.ClientSession | None = None
@@ -99,6 +439,11 @@ def validate_ec_consumer_routing(
 # Diagnostic switch: forward the original request to the decoder so the
 # only difference from the rewrite path is the rewrite itself.
 NO_REWRITE = False
+# Maximum images per encoder subrequest; 0 leaves batches unlimited.
+ENCODER_MAX_BATCH_SIZE = int(os.getenv("ENCODER_MAX_BATCH_SIZE", "0"))
+if ENCODER_MAX_BATCH_SIZE < 0:
+    raise ValueError("ENCODER_MAX_BATCH_SIZE must be non-negative")
+
 # Decode-side retries for a retryable internal error (`finish_reason="error"`,
 # e.g. an encoder embedding the connector could not deliver). Re-issuing runs
 # the encode again, which produces a fresh transfer.
@@ -135,6 +480,17 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
     the grid its processor actually produced), so the grid is never re-derived
     here -- a second derivation could disagree with the encoder's.
     """
+    mm_items = extract_mm_items(req_data)
+    # Audio processors may consume placeholders from other modalities too.
+    # Keep the whole request raw if any item cannot be rewritten safely.
+    control_fields = {"mm_hash", "ec_mm_hash", "transfer_id"}
+    raw_audio_request = any(
+        item["type"] in {"audio_url", "input_audio"} for item in mm_items
+    ) and any(
+        not item_meta.get(i, {}).get("mm_hash")
+        or not (item_meta.get(i, {}).keys() - control_fields)
+        for i in range(len(mm_items))
+    )
     rewritten = 0
     transfer_items = []
     idx = 0
@@ -154,6 +510,18 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
             item_uuid = meta.pop("mm_hash", None)
             ec_mm_hash = meta.pop("ec_mm_hash", None) or item_uuid
             transfer_id = meta.pop("transfer_id", None)
+            if raw_audio_request:
+                new_content.append(
+                    {
+                        **item,
+                        "uuid": item_uuid or item.get("uuid") or content_uuid(item),
+                    }
+                )
+                if transfer_id is not None:
+                    transfer_items.append(
+                        {"mm_hash": ec_mm_hash, "transfer_id": transfer_id}
+                    )
+                continue
             # Whatever keys the encoder reported are the metadata its model
             # declared as needed to size the placeholder range; the proxy does
             # not need to know their names.
@@ -187,9 +555,10 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
             rewritten += 1
         new_messages.append({**msg, "content": new_content})
 
-    if not rewritten:
+    if not rewritten and not raw_audio_request:
         return req_data
-    logger.info("Rewrote %d media item(s) as metadata references", rewritten)
+    if rewritten:
+        logger.info("Rewrote %d media item(s) as metadata references", rewritten)
     rewritten_request = {**req_data, "messages": new_messages}
     if transfer_items:
         ec_transfer_params = dict(req_data.get("ec_transfer_params") or {})
@@ -199,8 +568,7 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
 
 
 def extract_mm_items(request_data: dict) -> list[dict]:
-    """
-    Return *all* image/audio/video items that appear anywhere in `messages`.
+    """Return *all* image/audio/video items that appear anywhere in `messages`.
 
     Each returned dict looks like:
         { "type": "image_url", "image_url": {...} }
@@ -223,8 +591,7 @@ async def fanout_encoder_primer(
     req_id: str,
     consumer_zmq: str | None = None,
 ) -> tuple[dict[int, dict], dict[str, Any]]:
-    """
-    1. Build one request *per MM item* with all text removed.
+    """1. Group images by encoder, retaining per-item round-robin assignment.
     2. Send them concurrently to the encode cluster.
     3. Raise if any of them fails.
 
@@ -250,6 +617,7 @@ async def fanout_encoder_primer(
     item_uuids: dict[int, str] = {}
     item_transfer_ids: dict[int, str] = {}
     item_meta: dict[int, dict] = {}
+    transfer_items: dict[int, dict[str, str]] = {}
     ec_params: dict[str, Any] = {}
 
     # Round-robin over encode servers to distribute load a bit. The cursor
@@ -261,53 +629,72 @@ async def fanout_encoder_primer(
             e_urls, encoder_rr_idx, len(mm_items)
         )
 
+    groups: list[tuple[str, list[int]]] = []
+    image_groups: dict[str, list[int]] = {}
     for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)):
+        if item["type"] == "image_url":
+            indices = image_groups.get(target_url)
+            if indices is None or (
+                ENCODER_MAX_BATCH_SIZE and len(indices) >= ENCODER_MAX_BATCH_SIZE
+            ):
+                indices = []
+                image_groups[target_url] = indices
+                groups.append((target_url, indices))
+            indices.append(idx)
+        else:
+            groups.append((target_url, [idx]))
+
+    for target_url, indices in groups:
         # Derive a *child* request id:  <parent>:<index>:<random-short>
-        child_req_id = f"{req_id}:{idx}:{uuid.uuid4().hex[:6]}"
-        headers = {"x-request-id": child_req_id}
+        child_req_id = f"{req_id}:{indices[0]}:{uuid.uuid4().hex[:6]}"
+        headers = {"x-request-id": child_req_id, "Content-Type": "application/json"}
 
         # With --no-rewrite the decoder still receives the raw image and derives
         # the cache key by hashing it, so the encoder must do the same -- passing
         # a uuid here would make the two disagree and silently defeat the EC
         # transfer, leaving the decoder to encode the image itself.
-        item_uuid = None if NO_REWRITE else content_uuid(item)
-        if item_uuid is not None:
-            item_uuids[idx] = item_uuid
-        transfer_id = uuid.uuid4().hex
-        item_transfer_ids[idx] = transfer_id
+        content = []
+        for idx in indices:
+            item = mm_items[idx]
+            item_uuid = None if NO_REWRITE else (item.get("uuid") or content_uuid(item))
+            if item_uuid is not None:
+                item_uuids[idx] = item_uuid
+            item_transfer_ids[idx] = uuid.uuid4().hex
+            content.append(item if item_uuid is None else {**item, "uuid": item_uuid})
 
         encoder_req = {
-            # You *may* need to keep additional fields
             "model": orig_request.get("model"),
             "messages": [
                 {
                     "role": "user",
-                    "content": [
-                        item if item_uuid is None else {**item, "uuid": item_uuid}
-                    ],
+                    "content": content,
                 },
             ],
             # No max_tokens cap: the encoder instance never samples, it finishes
             # once the prompt is encoded and its embeddings are published.
             "stream": False,
         }
+        for key in (
+            "mm_processor_kwargs",
+            "media_io_kwargs",
+            "priority",
+            "session_id",
+        ):
+            if key in orig_request:
+                encoder_req[key] = orig_request[key]
         if consumer_zmq is not None:
-            # No mm_hash here on purpose. The encoder's own
-            # `mm_features[i].identifier` is derived from the uuid *and* the
-            # engine's media_io_kwargs / mm_processor_kwargs, so this proxy
-            # cannot know it before the encoder runs. Sending the bare uuid
-            # would fail the connector's hash match and make the producer
-            # invent its own transfer id, which the consumer could then never
-            # cancel. Omitting it lets the connector match by position, which
-            # is exact: one encoder request carries exactly one item.
+            # The engine may rehash UUIDs with processing options. Match by
+            # position instead; batches contain only images in input order.
             encoder_req["ec_transfer_params"] = {
                 "consumer_zmq": consumer_zmq,
-                "ec_items": [{"transfer_id": transfer_id}],
+                "ec_items": [
+                    {"transfer_id": item_transfer_ids[idx]} for idx in indices
+                ],
             }
         tasks.append(
             encode_session.post(
                 f"{target_url}/v1/chat/completions",
-                json=encoder_req,
+                data=msgspec.json.encode(encoder_req),
                 headers=headers,
             )
         )
@@ -315,7 +702,8 @@ async def fanout_encoder_primer(
     results = await asyncio.gather(*tasks, return_exceptions=True)
 
     # Fail fast if any sub-request failed
-    for idx, r in enumerate(results):
+    for (_, indices), r in zip(groups, results):
+        idx = indices[0]
         if isinstance(r, Exception):
             logger.error(
                 "[%s] Encoder request #%d raised exception: %s",
@@ -347,44 +735,74 @@ async def fanout_encoder_primer(
         # The encoder reports each mm_hash's metadata (e.g. the grid) here,
         # keyed by the same uuid this proxy assigned above.
         try:
-            params = (await r.json()).get("ec_transfer_params") or {}
+            params = msgspec.json.decode(await r.read()).get("ec_transfer_params") or {}
         except Exception:
             logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx)
             params = {}
         if params:
-            # One encoder request carries exactly one item, so there is a
-            # single reported entry. Do not key it by this proxy's uuid: when
-            # media_io_kwargs or mm_processor_kwargs are set the engine
-            # re-hashes the uuid together with them, so the encoder's own
-            # `mm_features[i].identifier` is a derived value this proxy cannot
-            # predict. Fall back to the sole entry, and carry the key the
-            # encoder actually used through as `ec_mm_hash`.
-            ec_mm_hash = item_uuids.get(idx)
-            reported = params.get(ec_mm_hash)
-            if reported is None and len(params) == 1:
-                ((ec_mm_hash, reported),) = params.items()
-            if reported:
+            by_index = {}
+            for mm_hash, reported in params.items():
+                if len(indices) == 1:
+                    break
+                for local_index in reported.get("item_indices", []):
+                    if type(local_index) is not int or not 0 <= local_index < len(
+                        indices
+                    ):
+                        raise HTTPException(502, "Invalid encoder item index")
+                    if local_index in by_index:
+                        raise HTTPException(502, "Duplicate encoder item index")
+                    by_index[local_index] = (mm_hash, reported)
+            for local_index, idx in enumerate(indices):
+                matched = by_index.get(local_index)
+                if matched is None:
+                    # Compatibility with encoders without position metadata.
+                    mm_hash = item_uuids.get(idx)
+                    if mm_hash in params:
+                        matched = (mm_hash, params[mm_hash])
+                    elif len(indices) == 1 and len(params) == 1:
+                        matched = next(iter(params.items()))
+                    elif (
+                        mm_items[idx]["type"] == "video_url"
+                        and (orig_request.get("mm_processor_kwargs") or {}).get(
+                            "use_audio_in_video"
+                        )
+                        and len(params) == 2
+                        and consumer_zmq is None
+                        and all(
+                            entry.keys() <= {"metadata", "item_indices"}
+                            for entry in params.values()
+                        )
+                    ):
+                        # Keep raw video for its audio/video features when no transfer
+                        # handles need to be forwarded.
+                        continue
+                    else:
+                        raise HTTPException(502, "Encoder metadata cannot be matched")
+                ec_mm_hash, reported = matched
                 metadata = reported.get("metadata") or {}
-                if metadata:
-                    item_meta[idx] = {
-                        **metadata,
-                        "mm_hash": item_uuids.get(idx, ec_mm_hash),
-                        "ec_mm_hash": ec_mm_hash,
-                    }
-                    if idx in item_transfer_ids:
-                        item_meta[idx]["transfer_id"] = item_transfer_ids[idx]
+                item_meta[idx] = {
+                    **metadata,
+                    "mm_hash": item_uuids.get(idx, ec_mm_hash),
+                    "ec_mm_hash": ec_mm_hash,
+                }
+                if idx in item_transfer_ids:
+                    item_meta[idx]["transfer_id"] = item_transfer_ids[idx]
                 # Whatever the encoder reported alongside `metadata` is the
                 # connector's own handle on the published embedding (for NIXL,
                 # peer_host/peer_port/size_bytes). The decoder's connector
                 # looks it up by mm_hash on the request, so carry it through.
-                ec_params[item_uuids.get(idx, ec_mm_hash)] = reported
+                ec_params[ec_mm_hash] = reported
                 if NO_REWRITE and consumer_zmq is not None:
-                    ec_params.setdefault("ec_items", []).append(
-                        {"mm_hash": ec_mm_hash, "transfer_id": item_transfer_ids[idx]}
-                    )
+                    transfer_items[idx] = {
+                        "mm_hash": ec_mm_hash,
+                        "transfer_id": item_transfer_ids[idx],
+                    }
+
+    if transfer_items:
+        ec_params["ec_items"] = [transfer_items[idx] for idx in sorted(transfer_items)]
 
     logger.info(
-        "[%s] All %d encoder requests completed successfully", req_id, len(mm_items)
+        "[%s] All %d encoder requests completed successfully", req_id, len(groups)
     )
     return item_meta, ec_params
 
@@ -393,9 +811,9 @@ async def maybe_prefill(
     req_data: dict,
     p_url: str,
     req_id: str,
+    dp_rank: int | None = None,
 ) -> dict:
-    """
-    - Do prefill-only task if p_url exist;
+    """- Do prefill-only task if p_url exist;
     - Return a new body carrying kv transfer params (for nixl connector)
     - Else, skip and return the original request data for decode
 
@@ -405,9 +823,9 @@ async def maybe_prefill(
     if p_url:
         logger.info("[%s] Processing through prefill: %s", req_id, p_url)
 
-        prefill_response = await process_prefill_stage(req_data, p_url, req_id)
+        prefill_response = await process_prefill_stage(req_data, p_url, req_id, dp_rank)
         # for nixl connector to facilitate kv transfer...
-        prefill_response_json = await prefill_response.json()
+        prefill_response_json = msgspec.json.decode(await prefill_response.read())
         kv_transfer_params = prefill_response_json.get("kv_transfer_params", {})
         if kv_transfer_params:
             return {**req_data, "kv_transfer_params": kv_transfer_params}
@@ -419,6 +837,7 @@ async def process_prefill_stage(
     req_data: dict,
     p_url: str,
     req_id: str,
+    dp_rank: int | None = None,
 ) -> dict:
     """Process request through Prefill stage and return kv_transfer_params"""
     logger.info("[%s] Sending prefill request to: %s", req_id, p_url)
@@ -439,10 +858,14 @@ async def process_prefill_stage(
     if "stream_options" in prefill_request:
         del prefill_request["stream_options"]
 
-    headers = {"x-request-id": req_id}
+    headers = {"x-request-id": req_id, "Content-Type": "application/json"}
+    if dp_rank is not None:
+        headers["X-data-parallel-rank"] = str(dp_rank)
     try:
         prefill_response = await prefill_session.post(
-            f"{p_url}/v1/chat/completions", json=prefill_request, headers=headers
+            f"{p_url}/v1/chat/completions",
+            data=msgspec.json.encode(prefill_request),
+            headers=headers,
         )
         prefill_response.raise_for_status()
 
@@ -538,9 +961,7 @@ async def on_startup() -> None:
         ),
     )
     encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
-    if app.state.p_urls:
-        # only setup if prefill instance(s) exist
-        prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
+    prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
     decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
 
 
@@ -566,6 +987,7 @@ async def prepare_for_decode(
     e_urls: list[str],
     p_url: str,
     consumer_zmq: str | None,
+    prefill_dp_rank: int | None = None,
 ) -> tuple[dict, float, float]:
     """Encode, rewrite and prefill, returning the body to send to decode.
 
@@ -586,7 +1008,7 @@ async def prepare_for_decode(
         handles.update(ec_params)
         prepared = {**prepared, "ec_transfer_params": handles}
     _t2 = time.perf_counter()
-    prepared = await maybe_prefill(prepared, p_url, req_id)
+    prepared = await maybe_prefill(prepared, p_url, req_id, prefill_dp_rank)
     return prepared, _t1 - _t0, _t2 - _t1
 
 
@@ -598,22 +1020,25 @@ async def forward_non_stream(
     d_url: str,
     consumer_zmq: str | None,
     dp_rank: int | None = None,
-) -> dict:
+    prefill_dp_rank: int | None = None,
+) -> Response:
     try:
         for attempt in range(DECODE_RETRIES + 1):
             _t0 = time.perf_counter()
             prepared, encode_s, rewrite_s = await prepare_for_decode(
-                req_data, req_id, e_urls, p_url, consumer_zmq
+                req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank
             )
             _t2 = time.perf_counter()
 
             logger.info("[%s] Forwarding to decode: %s", req_id, d_url)
-            headers = {"x-request-id": req_id}
+            headers = {"x-request-id": req_id, "Content-Type": "application/json"}
             if dp_rank is not None:
                 headers["X-data-parallel-rank"] = str(dp_rank)
 
             async with decode_session.post(
-                f"{d_url}/v1/chat/completions", json=prepared, headers=headers
+                f"{d_url}/v1/chat/completions",
+                data=msgspec.json.encode(prepared),
+                headers=headers,
             ) as resp:
                 if resp.status >= 400:
                     detail = await resp.text()
@@ -637,7 +1062,7 @@ async def forward_non_stream(
                         detail,
                     )
                     raise HTTPException(status_code=resp.status, detail=detail)
-                out = await resp.json()
+                out = await resp.read()
                 _t3 = time.perf_counter()
                 logger.info(
                     "STAGE %s encode=%.1f rewrite=%.1f decode=%.1f total=%.1f "
@@ -649,7 +1074,15 @@ async def forward_non_stream(
                     (_t3 - _t0) * 1e3,
                     attempt,
                 )
-                return out
+                return Response(
+                    content=out,
+                    status_code=resp.status,
+                    headers={
+                        "Content-Type": resp.headers.get(
+                            "Content-Type", "application/json"
+                        )
+                    },
+                )
         raise HTTPException(status_code=500, detail="Decode failed after re-encoding")
 
     except HTTPException:
@@ -667,24 +1100,25 @@ async def forward_stream(
     d_url: str,
     consumer_zmq: str | None,
     dp_rank: int | None = None,
-) -> AsyncIterator[str]:
+    prefill_dp_rank: int | None = None,
+) -> AsyncIterator[bytes]:
     try:
         for attempt in range(DECODE_RETRIES + 1):
             _t0 = time.perf_counter()
             prepared, encode_s, rewrite_s = await prepare_for_decode(
-                req_data, req_id, e_urls, p_url, consumer_zmq
+                req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank
             )
             _t2 = time.perf_counter()
 
             logger.info("[%s] Starting streaming from decode: %s", req_id, d_url)
-            headers = {"x-request-id": req_id}
+            headers = {"x-request-id": req_id, "Content-Type": "application/json"}
             if dp_rank is not None:
                 headers["X-data-parallel-rank"] = str(dp_rank)
 
             _first = None
             async with decode_session.post(
                 f"{d_url}/v1/chat/completions",
-                json=prepared,
+                data=msgspec.json.encode(prepared),
                 headers=headers,
             ) as resp:
                 # Retry only before the first chunk: once anything reached the
@@ -701,11 +1135,11 @@ async def forward_stream(
                     )
                     continue
                 resp.raise_for_status()
-                async for chunk in resp.content.iter_chunked(1024):
+                async for chunk in resp.content.iter_any():
                     if chunk:
                         if _first is None:
                             _first = time.perf_counter()
-                        yield chunk.decode("utf-8", errors="ignore")
+                        yield chunk
             _t3 = time.perf_counter()
 
             logger.info(
@@ -739,7 +1173,7 @@ async def forward_stream(
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request):
     try:
-        req_data = await request.json()
+        req_data = msgspec.json.decode(await request.body())
         req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
 
         e_urls = app.state.e_urls  # we want the full list for fan-out
@@ -764,10 +1198,9 @@ async def chat_completions(request: Request):
                 ),
                 media_type="text/event-stream",
             )
-        result = await forward_non_stream(
+        return await forward_non_stream(
             req_data, req_id, e_urls, p_url, d_url, consumer_zmq, dp_rank
         )
-        return JSONResponse(content=result)
 
     except HTTPException:
         raise
@@ -830,14 +1263,14 @@ async def _post_if_available(
     payload: dict,
     headers: dict,
 ) -> dict | None:
-    """
-    POST `payload` to `url`.
+    """POST `payload` to `url`.
 
     Returns
     -------
     • The decoded JSON body on success (2xx)
     • None if the endpoint does not exist (404)
     • Raises for anything else.
+
     """
     try:
         resp = await session.post(url, json=payload, headers=headers)
@@ -858,9 +1291,7 @@ async def _post_if_available(
 
 
 async def _profile_cmd(cmd: str, payload: dict, e_url: str, p_url: str, d_url: str):
-    """
-    Fire & forget to both clusters, tolerate 404.
-    """
+    """Fire & forget to both clusters, tolerate 404."""
     headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"}
 
     encode_task = _post_if_available(
@@ -913,10 +1344,147 @@ async def stop_profile(request: Request):
     return await _profile_cmd("stop", body, e_url, p_url, d_url)
 
 
+def build_app(config: EPDProxyConfig | None = None) -> FastAPI:
+    config = config or EPDProxyConfig()
+    registry = InstanceRegistry(
+        probe_interval=config.probe_interval,
+        probe_timeout=config.probe_timeout,
+        fail_threshold=config.fail_threshold,
+        evicted_ttl=config.evicted_ttl,
+    )
+    proxy = EPDProxy(registry)
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        await on_startup()
+        try:
+            registry.start_probing()
+            yield
+        finally:
+            await registry.stop_probing()
+            await on_shutdown()
+
+    app = FastAPI(lifespan=lifespan)
+    app.state.proxy = proxy
+    app.state.registry = registry
+
+    @app.get("/instances")
+    async def list_instances():
+        return registry.status()
+
+    @app.post("/instances", dependencies=[Depends(require_admin_key)])
+    async def register_instance(body: InstanceRegistration):
+        if body.ec_zmq_addrs and body.role not in (
+            InstanceRole.PREFILL,
+            InstanceRole.PREFILL_DECODE,
+        ):
+            raise HTTPException(400, "Only EC consumers accept Mooncake addresses")
+        try:
+            created = registry.register(
+                InstanceRecord(
+                    body.role,
+                    str(body.url).rstrip("/"),
+                    body.ec_zmq_addrs,
+                    body.dp_size,
+                )
+            )
+        except ValueError as error:
+            raise HTTPException(409, str(error)) from error
+        return {"registered": created}
+
+    @app.delete("/instances", dependencies=[Depends(require_admin_key)])
+    async def unregister_instance(url: AnyHttpUrl):
+        return {"removed": registry.unregister(str(url).rstrip("/"))}
+
+    @app.post("/v1/chat/completions")
+    async def chat_completions(request: Request):
+        req_data = msgspec.json.decode(await request.body())
+        req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
+        route = proxy.route(len(extract_mm_items(req_data)))
+        args = (
+            req_data,
+            req_id,
+            route.encoder_urls,
+            route.prefill.url if route.prefill else None,
+            route.decode.url,
+            route.consumer_zmq,
+        )
+        ranks = {
+            "dp_rank": route.dp_rank if route.prefill is None else None,
+            "prefill_dp_rank": route.dp_rank if route.prefill is not None else None,
+        }
+        if req_data.get("stream", False):
+            return StreamingResponse(
+                forward_stream(*args, **ranks), media_type="text/event-stream"
+            )
+        return await forward_non_stream(*args, **ranks)
+
+    @app.get("/v1/models")
+    async def list_models():
+        decode = registry.pick(InstanceRole.PREFILL_DECODE) or registry.pick(
+            InstanceRole.DECODE
+        )
+        if decode is None:
+            raise HTTPException(
+                status_code=503, detail="No decode instance is registered"
+            )
+        async with decode_session.get(f"{decode.url}/v1/models") as resp:
+            resp.raise_for_status()
+            return await resp.json()
+
+    @app.get("/health")
+    async def health():
+        status = registry.status()
+        # An empty roster is not unhealthy: the proxy is meant to come up
+        # before anything registers with it.
+        return JSONResponse({"proxy": "healthy", "instances": status})
+
+    @app.post("/start_profile")
+    async def start_profile(request: Request):
+        return await _profile(registry, "start", await request.json())
+
+    @app.post("/stop_profile")
+    async def stop_profile(request: Request):
+        return await _profile(registry, "stop", await request.json())
+
+    return app
+
+
+async def _profile(registry: InstanceRegistry, cmd: str, payload: dict) -> dict:
+    headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"}
+    targets = {role.value: registry.pick(role) for role in InstanceRole}
+    results = await asyncio.gather(
+        *(
+            _post_if_available(
+                decode_session, f"{record.url}/{cmd}_profile", payload, headers
+            )
+            for record in targets.values()
+            if record is not None
+        )
+    )
+    reachable = [result for result in results if result is not None]
+    if not reachable:
+        raise HTTPException(
+            status_code=503,
+            detail="Profiling endpoints are disabled on every instance",
+        )
+    live = [name for name, record in targets.items() if record is not None]
+    return dict(zip(live, results))
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", default="0.0.0.0")
     parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument(
+        "--dynamic-registration",
+        action="store_true",
+        help="Enable launcher-managed HTTP registration for E+PD or E+P+D.",
+    )
+    parser.add_argument("--probe-interval", type=float, default=DEFAULT_PROBE_INTERVAL)
+    parser.add_argument("--probe-timeout", type=float, default=DEFAULT_PROBE_TIMEOUT)
+    parser.add_argument("--fail-threshold", type=int, default=DEFAULT_FAIL_THRESHOLD)
+    parser.add_argument("--evicted-ttl", type=float, default=DEFAULT_EVICTED_TTL)
     parser.add_argument(
         "--log-requests",
         action="store_true",
@@ -932,12 +1500,12 @@ async def stop_profile(request: Request):
     )
     parser.add_argument(
         "--encode-servers-urls",
-        required=True,
+        default="",
         help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")',
     )
     parser.add_argument(
         "--prefill-servers-urls",
-        required=True,
+        default="none",
         help=(
             'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") '
             'to enable E->P->D, set "disable" or "none" to enable E->PD'
@@ -945,7 +1513,7 @@ async def stop_profile(request: Request):
     )
     parser.add_argument(
         "--decode-servers-urls",
-        required=True,
+        default="",
         help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")',
     )
     parser.add_argument(
@@ -980,58 +1548,87 @@ async def stop_profile(request: Request):
     )
 
     args = parser.parse_args()
+    if args.fail_threshold < 1:
+        parser.error("--fail-threshold must be at least 1")
     if args.log_requests:
         logging.getLogger().setLevel(logging.DEBUG)
         app.middleware("http")(log_requests)
     NO_REWRITE = args.no_rewrite
     DECODE_RETRIES = max(0, args.decode_retries)
-    app.state.e_urls = [
-        u.strip() for u in args.encode_servers_urls.split(",") if u.strip()
-    ]
-    app.state.d_urls = [
-        u.strip() for u in args.decode_servers_urls.split(",") if u.strip()
-    ]
-    app.state.d_ec_urls = [
-        u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip()
-    ]
-    if args.ec_consumer_dp_size < 1:
-        parser.error("--ec-consumer-dp-size must be at least 1")
-    app.state.ec_consumer_dp_size = args.ec_consumer_dp_size
-    app.state.replica_counter = itertools.count()
-    expected = len(app.state.d_urls) * args.ec_consumer_dp_size
-    if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected:
-        parser.error(
-            "--ec-consumer-zmq-addrs must contain one address per consumer "
-            f"replica: expected {expected} "
-            f"({len(app.state.d_urls)} servers x {args.ec_consumer_dp_size} replicas), "
-            f"got {len(app.state.d_ec_urls)}"
-        )
-    # handle prefill instances
-    if args.prefill_servers_urls.lower() in ("disable", "none", ""):
-        app.state.p_urls = []
-        logger.info(
-            "Disaggregated prefill phase explicitly disabled by user. Running E + PD..."
+    if args.dynamic_registration:
+        if not os.getenv("ADMIN_API_KEY"):
+            parser.error("--dynamic-registration requires ADMIN_API_KEY")
+        if (
+            args.encode_servers_urls
+            or args.prefill_servers_urls.lower() not in ("disable", "none", "")
+            or args.decode_servers_urls
+            or args.ec_consumer_zmq_addrs
+        ):
+            parser.error(
+                "With --dynamic-registration, instances register through /instances"
+            )
+        app = build_app(
+            EPDProxyConfig(
+                probe_interval=args.probe_interval,
+                probe_timeout=args.probe_timeout,
+                fail_threshold=args.fail_threshold,
+                evicted_ttl=args.evicted_ttl,
+            )
         )
+        if args.log_requests:
+            app.middleware("http")(log_requests)
     else:
-        app.state.p_urls = [
-            u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()
+        if not args.encode_servers_urls or not args.decode_servers_urls:
+            parser.error(
+                "Static routing requires --encode-servers-urls "
+                "and --decode-servers-urls"
+            )
+        app.state.e_urls = [
+            u.strip() for u in args.encode_servers_urls.split(",") if u.strip()
         ]
-        logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
-    try:
-        validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls)
-    except ValueError as exc:
-        parser.error(str(exc))
-
-    logger.info("Proxy listening on %s:%s", args.host, args.port)
-    logger.info("Encode servers: %s", app.state.e_urls)
-    logger.info("Prefill instances %s", app.state.p_urls)
-    logger.info("Decode servers: %s", app.state.d_urls)
-    if app.state.ec_consumer_dp_size > 1:
-        logger.info(
-            "EC consumer replicas per server: %d (control addresses: %s)",
-            app.state.ec_consumer_dp_size,
-            app.state.d_ec_urls,
-        )
+        app.state.d_urls = [
+            u.strip() for u in args.decode_servers_urls.split(",") if u.strip()
+        ]
+        app.state.d_ec_urls = [
+            u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip()
+        ]
+        if args.ec_consumer_dp_size < 1:
+            parser.error("--ec-consumer-dp-size must be at least 1")
+        app.state.ec_consumer_dp_size = args.ec_consumer_dp_size
+        app.state.replica_counter = itertools.count()
+        expected = len(app.state.d_urls) * args.ec_consumer_dp_size
+        if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected:
+            parser.error(
+                "--ec-consumer-zmq-addrs must contain one address per consumer "
+                f"replica: expected {expected} "
+                f"({len(app.state.d_urls)} servers x "
+                f"{args.ec_consumer_dp_size} replicas), "
+                f"got {len(app.state.d_ec_urls)}"
+            )
+        # handle prefill instances
+        if args.prefill_servers_urls.lower() in ("disable", "none", ""):
+            app.state.p_urls = []
+            logger.info("Disaggregated prefill disabled. Running E + PD...")
+        else:
+            app.state.p_urls = [
+                u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()
+            ]
+            logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
+        try:
+            validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls)
+        except ValueError as exc:
+            parser.error(str(exc))
+
+        logger.info("Proxy listening on %s:%s", args.host, args.port)
+        logger.info("Encode servers: %s", app.state.e_urls)
+        logger.info("Prefill instances %s", app.state.p_urls)
+        logger.info("Decode servers: %s", app.state.d_urls)
+        if app.state.ec_consumer_dp_size > 1:
+            logger.info(
+                "EC consumer replicas per server: %d (control addresses: %s)",
+                app.state.ec_consumer_dp_size,
+                app.state.d_ec_urls,
+            )
 
     uvicorn.run(
         app,
diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py
index 57deef6a15d2..ae35ec6d7695 100644
--- a/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py
+++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file provides a disaggregated prefilling proxy demo to demonstrate an
+"""This file provides a disaggregated prefilling proxy demo to demonstrate an
 example usage of XpYd disaggregated prefilling.
 We can launch multiple vllm instances (2 for prefill and 2 for decode), and
 launch this proxy demo through:
diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py
index cc1cc402d29e..3102f344b027 100644
--- a/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py
+++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer
+"""Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer.
 
 This proxy sits between clients and a vLLM Prefill/Decode (P/D) deployment,
 routing multi-turn chat requests so that each turn reuses KV cache blocks
@@ -54,7 +53,7 @@
     it the benchmark sends OpenAI-schema-compliant payloads and every
     turn lands as a cache MISS in this proxy.
 
-    Example:
+Example:
         python benchmarks/multi_turn/benchmark_serving_multi_turn.py \\
             --model <MODEL> --served-model-name <NAME> \\
             --url http://<proxy_host>:8000 \\
@@ -67,6 +66,7 @@
 
 Dependencies:
     pip install fastapi uvicorn httpx
+
 """
 
 from __future__ import annotations
diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py
index 672162f8ba34..53d8b210d677 100644
--- a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py
+++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Push-mode disaggregated prefilling proxy demo.
+"""Push-mode disaggregated prefilling proxy demo.
 
 Companion to ``disagg_proxy_demo.py`` (pull mode). The client-facing API is
 the same; the difference is in how P and D coordinate the KV transfer:
@@ -31,6 +30,7 @@
 """
 
 import argparse
+import asyncio
 import contextlib
 import ipaddress
 import itertools
@@ -259,17 +259,24 @@ async def _push_completion(self, raw_request: Request, path: str):
         decode_instance = self.schedule(self.decode_cycler)
         headers = self._common_headers(request_id)
 
-        # Fire prefill; we don't read its body but must drain the
-        # connection so the upstream server can free its slot.
-        async for _ in self.forward_request(
-            f"http://{prefill_instance}{path}", prefill_request, headers
-        ):
-            continue
+        async def drain_prefill():
+            async for _ in self.forward_request(
+                f"http://{prefill_instance}{path}", prefill_request, headers
+            ):
+                continue
 
-        generator = self.forward_request(
-            f"http://{decode_instance}{path}", decode_request, headers
-        )
-        return StreamingResponse(generator, media_type="application/json")
+        prefill_task = asyncio.create_task(drain_prefill())
+
+        async def stream_decode():
+            try:
+                async for chunk in self.forward_request(
+                    f"http://{decode_instance}{path}", decode_request, headers
+                ):
+                    yield chunk
+            finally:
+                await prefill_task
+
+        return StreamingResponse(stream_decode(), media_type="application/json")
 
     async def create_completion(self, raw_request: Request):
         try:
diff --git a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
index 8b2c46999efb..ce4d558fb919 100644
--- a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+++ b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
@@ -367,6 +367,8 @@ async def handle_request(api: str, request: Request):
                 "remote_block_ids"
             ]
             req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"]
+            if prefill_kv.get("tp_size") is not None:
+                req_data["kv_transfer_params"]["tp_size"] = prefill_kv["tp_size"]
 
         req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
             "dp_size"
diff --git a/examples/disaggregated/example_connector/decode_example.py b/examples/disaggregated/example_connector/decode_example.py
index 2d575840e6a7..9b76e2f00178 100644
--- a/examples/disaggregated/example_connector/decode_example.py
+++ b/examples/disaggregated/example_connector/decode_example.py
@@ -6,7 +6,7 @@
 
 
 def read_prompts():
-    """Read prompts from output.txt"""
+    """Read prompts from output.txt."""
     prompts = []
     try:
         with open("output.txt") as f:
diff --git a/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py b/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py
index b67c2459319f..55f4fe52007d 100644
--- a/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py
+++ b/examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use FlexKV with vLLM for prefix caching.
+"""This example shows how to use FlexKV with vLLM for prefix caching.
 
 FlexKV is a distributed KV Store and multi-level cache management system for
 ultra-large-scale LLM inference.
diff --git a/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py b/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py
index db9c5a85f7f0..3c3d5dbff985 100644
--- a/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py
+++ b/examples/disaggregated/kv_load_failure_recovery_offline/decode_example.py
@@ -7,7 +7,7 @@
 
 
 def read_prompts():
-    """Read prompts from prefill_output.txt"""
+    """Read prompts from prefill_output.txt."""
     prompts = []
     try:
         with open("prefill_output.txt") as f:
diff --git a/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py b/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py
index 5bae04426746..2b2e94aa746b 100644
--- a/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py
+++ b/examples/disaggregated/kv_load_failure_recovery_offline/load_recovery_example_connector.py
@@ -119,8 +119,7 @@ def get_num_new_matched_tokens(
     def update_state_after_alloc(
         self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
     ):
-        """
-        Update KVConnector state after block allocation.
+        """Update KVConnector state after block allocation.
 
         If blocks were allocated, add to _requests_need_load,
         such that we load the KVs in the next forward pass.
diff --git a/examples/disaggregated/lmcache/cpu_offload_lmcache.py b/examples/disaggregated/lmcache/cpu_offload_lmcache.py
index b67a929e5d96..8a84675952ab 100644
--- a/examples/disaggregated/lmcache/cpu_offload_lmcache.py
+++ b/examples/disaggregated/lmcache/cpu_offload_lmcache.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file demonstrates the example usage of CPU offloading
+"""This file demonstrates the example usage of CPU offloading
 with LMCache in vLLM v1.
 
 Note that `lmcache` is needed to run this example.
diff --git a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py
index e048aecff9f7..3a64139c637a 100644
--- a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py
+++ b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py
@@ -14,9 +14,7 @@
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """
-    Lifespan context manager to handle startup and shutdown events.
-    """
+    """Lifespan context manager to handle startup and shutdown events."""
     # Startup: Initialize clients
     prefiller_base_url = (
         f"http://{global_args.prefiller_host}:{global_args.prefiller_port}/v1"
@@ -106,9 +104,7 @@ def parse_args():
 async def send_request_to_service(
     client: httpx.AsyncClient, endpoint: str, req_data: dict
 ):
-    """
-    Send a request to a service using a persistent client.
-    """
+    """Send a request to a service using a persistent client."""
     req_data = req_data.copy()
     req_data["max_tokens"] = 1
     if "max_completion_tokens" in req_data:
@@ -128,9 +124,7 @@ async def send_request_to_service(
 async def stream_service_response(
     client: httpx.AsyncClient, endpoint: str, req_data: dict
 ):
-    """
-    Asynchronously stream the response from a service using a persistent client.
-    """
+    """Asynchronously stream the response from a service using a persistent client."""
     headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
     async with client.stream(
         "POST", endpoint, json=req_data, headers=headers
diff --git a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py
index 489ff1321229..a50d6cbbc30d 100644
--- a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py
+++ b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file demonstrates the example usage of remote KV cache sharing
+"""This file demonstrates the example usage of remote KV cache sharing
 with LMCache.
 We will launch 2 vllm instances, and launch an additional LMCache server.
 KV cache is transferred in the following manner:
diff --git a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py
index 09880a32aa35..6815f420616c 100644
--- a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py
+++ b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py
@@ -65,9 +65,7 @@ async def get_prefiller_info(prefill_clients: list, ready: asyncio.Event):
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """
-    Lifespan context manager to handle startup and shutdown events.
-    """
+    """Lifespan context manager to handle startup and shutdown events."""
     # Startup: Initialize client pools for prefiller and decoder services
     app.state.prefill_clients = []
     app.state.decode_clients = []
@@ -180,10 +178,12 @@ def _parse_prefill_urls(prefill_list):
     """Parse prefill URLs from --prefill arguments.
 
     Format: --prefill URL [BOOTSTRAP_PORT]
+
     Example:
         --prefill http://prefill1:8080 9000  # With bootstrap port
         --prefill http://prefill2:8080 none  # Explicitly no bootstrap port
         --prefill http://prefill3:8080       # Defaults to no bootstrap port
+
     """
     if not prefill_list:
         return []
@@ -228,8 +228,7 @@ def _parse_decode_urls(decode_list):
 
 
 def get_next_client(app, service_type: str):
-    """
-    Get the next client in round-robin fashion.
+    """Get the next client in round-robin fashion.
 
     Args:
         app: The FastAPI app instance
@@ -237,6 +236,7 @@ def get_next_client(app, service_type: str):
 
     Returns:
         The next client to use
+
     """
     if service_type == "prefill":
         return next(app.state.prefill_iterator)
@@ -250,9 +250,7 @@ def get_next_client(app, service_type: str):
 async def send_request_to_service(
     client_info: dict, dp_rank: int, endpoint: str, req_data: dict, request_id: str
 ):
-    """
-    Send a request to a service using a client from the pool.
-    """
+    """Send a request to a service using a client from the pool."""
     req_data = req_data.copy()
     req_data["kv_transfer_params"] = {
         "do_remote_decode": True,
@@ -288,9 +286,7 @@ async def stream_service_response(
     req_data: dict,
     request_id: str,
 ):
-    """
-    Asynchronously stream response from a service using a client from the pool.
-    """
+    """Asynchronously stream response from a service using a client from the pool."""
     headers = {
         "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
         "X-Request-Id": request_id,
diff --git a/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py b/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py
index 801b4b769792..db033e1861e8 100644
--- a/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py
+++ b/examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Demonstration script for Automatic Prefix Caching (APC) in vLLM.
+"""Demonstration script for Automatic Prefix Caching (APC) in vLLM.
 
 Automatic Prefix Caching (APC) allows the vLLM engine to reuse cached
 KV (key-value) pairs from previous prompts if a new query shares the same
diff --git a/examples/features/batch_invariance/reproducibility_offline.py b/examples/features/batch_invariance/reproducibility_offline.py
index 72c1e841dca4..54afcdee9f41 100644
--- a/examples/features/batch_invariance/reproducibility_offline.py
+++ b/examples/features/batch_invariance/reproducibility_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Demonstrates how to achieve reproducibility in vLLM.
+"""Demonstrates how to achieve reproducibility in vLLM.
 
 Main article: https://docs.vllm.ai/en/latest/usage/reproducibility.html
 """
diff --git a/examples/features/context_extension/context_extension_offline.py b/examples/features/context_extension/context_extension_offline.py
index 3874288b5e11..cf270a22305f 100644
--- a/examples/features/context_extension/context_extension_offline.py
+++ b/examples/features/context_extension/context_extension_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This script demonstrates how to extend the context length
+"""This script demonstrates how to extend the context length
 of a Qwen model using the YARN method (rope_parameters)
 and run a simple chat example.
 
diff --git a/examples/features/data_parallel/data_parallel_offline.py b/examples/features/data_parallel/data_parallel_offline.py
index c38ff7297afc..86e61e04f0d7 100644
--- a/examples/features/data_parallel/data_parallel_offline.py
+++ b/examples/features/data_parallel/data_parallel_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Usage:
+"""Usage:
 Single node:
     python examples/features/data_parallel/data_parallel_offline.py \
             --model="ibm-research/PowerMoE-3b" \
diff --git a/examples/features/kv_events/kv_events_subscriber.py b/examples/features/kv_events/kv_events_subscriber.py
index 9c1b5138bc0e..91d3829c54ad 100644
--- a/examples/features/kv_events/kv_events_subscriber.py
+++ b/examples/features/kv_events/kv_events_subscriber.py
@@ -18,7 +18,7 @@ class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
 
 
 class KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True):
-    """Base class for all KV cache-related events"""
+    """Base class for all KV cache-related events."""
 
 
 class BlockStored(KVCacheEvent):
diff --git a/examples/features/logits_processor/custom.py b/examples/features/logits_processor/custom.py
index ce000872dc96..c5028aa797cc 100644
--- a/examples/features/logits_processor/custom.py
+++ b/examples/features/logits_processor/custom.py
@@ -48,7 +48,7 @@ class object.
 
 # Hypothetical custom logits processor
 class DummyLogitsProcessor(LogitsProcessor):
-    """Fake logit processor to support unit testing and examples"""
+    """Fake logit processor to support unit testing and examples."""
 
     @classmethod
     def validate_params(cls, params: SamplingParams):
diff --git a/examples/features/logits_processor/custom_req.py b/examples/features/logits_processor/custom_req.py
index 5763fff5410d..eb57029efa0b 100644
--- a/examples/features/logits_processor/custom_req.py
+++ b/examples/features/logits_processor/custom_req.py
@@ -58,7 +58,7 @@ class DummyPerReqLogitsProcessor:
     token id identified by `target_token`"""
 
     def __init__(self, target_token: int) -> None:
-        """Specify `target_token`"""
+        """Specify `target_token`."""
         self.target_token = target_token
 
     def __call__(
@@ -103,6 +103,7 @@ def new_req_logits_processor(
 
         Returns:
           `Callable` request logits processor, or None
+
         """
         target_token: Any | None = params.extra_args and params.extra_args.get(
             "target_token"
diff --git a/examples/features/logits_processor/custom_req_init.py b/examples/features/logits_processor/custom_req_init.py
index acd2c47f230f..38554f06d8ad 100644
--- a/examples/features/logits_processor/custom_req_init.py
+++ b/examples/features/logits_processor/custom_req_init.py
@@ -59,7 +59,7 @@ class DummyPerReqLogitsProcessor:
     token id identified by `target_token`"""
 
     def __init__(self, target_token: int) -> None:
-        """Specify `target_token`"""
+        """Specify `target_token`."""
         self.target_token = target_token
 
     def __call__(
@@ -111,6 +111,7 @@ def new_req_logits_processor(
 
         Returns:
           `Callable` request logits processor, or None
+
         """
         if (
             not self.is_cuda
diff --git a/examples/features/lora/lora_with_quantization_offline.py b/examples/features/lora/lora_with_quantization_offline.py
index ef6dbd4534fc..23f165cc5f75 100644
--- a/examples/features/lora/lora_with_quantization_offline.py
+++ b/examples/features/lora/lora_with_quantization_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use LoRA with different quantization techniques
+"""This example shows how to use LoRA with different quantization techniques
 for offline inference.
 
 Requires HuggingFace credentials for access.
@@ -73,7 +72,6 @@ def initialize_engine(
     model: str, quantization: str, lora_repo: str | None
 ) -> LLMEngine:
     """Initialize the LLMEngine."""
-
     engine_args = EngineArgs(
         model=model,
         quantization=quantization,
@@ -86,7 +84,6 @@ def initialize_engine(
 
 def main():
     """Main function that sets up and runs the prompt processing."""
-
     test_configs = [
         # QLoRA (https://arxiv.org/abs/2305.14314)
         {
diff --git a/examples/features/lora/multilora_offline.py b/examples/features/lora/multilora_offline.py
index 92021f9fb226..d870ca7562a1 100644
--- a/examples/features/lora/multilora_offline.py
+++ b/examples/features/lora/multilora_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use the multi-LoRA functionality
+"""This example shows how to use the multi-LoRA functionality
 for offline inference.
 
 Requires HuggingFace credentials for access to Llama2.
diff --git a/examples/features/pause_resume/data_parallel_pause_resume.py b/examples/features/pause_resume/data_parallel_pause_resume.py
index 1f11536e5366..fa26d0aa3c66 100644
--- a/examples/features/pause_resume/data_parallel_pause_resume.py
+++ b/examples/features/pause_resume/data_parallel_pause_resume.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test pause/resume with Data Parallel (DP) via HTTP API.
+"""Test pause/resume with Data Parallel (DP) via HTTP API.
 
 This example demonstrates coordinated pause/resume across multiple DP ranks.
 The pause synchronizes across all DP engines via all-reduce.
diff --git a/examples/features/prompt_embed/prompt_embed_offline.py b/examples/features/prompt_embed/prompt_embed_offline.py
index 9e90aa46b5b7..2b5d3121c89b 100644
--- a/examples/features/prompt_embed/prompt_embed_offline.py
+++ b/examples/features/prompt_embed/prompt_embed_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Demonstrates how to generate prompt embeddings using
+"""Demonstrates how to generate prompt embeddings using
 Hugging Face Transformers  and use them as input to vLLM
 for both single and batch inference.
 
diff --git a/examples/features/reset_kv/reset_kv_offline.py b/examples/features/reset_kv/reset_kv_offline.py
index 3fbe7fa7545e..c19d88a35fbc 100644
--- a/examples/features/reset_kv/reset_kv_offline.py
+++ b/examples/features/reset_kv/reset_kv_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file demonstrates preempt requests when using the `LLMEngine`
+"""This file demonstrates preempt requests when using the `LLMEngine`
 for processing prompts with various sampling parameters.
 """
 
diff --git a/examples/features/sharded_state/load_sharded_state_offline.py b/examples/features/sharded_state/load_sharded_state_offline.py
index e867db5d12fe..66d8a75107aa 100644
--- a/examples/features/sharded_state/load_sharded_state_offline.py
+++ b/examples/features/sharded_state/load_sharded_state_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Validates the loading of a model saved with the sharded_state format.
+"""Validates the loading of a model saved with the sharded_state format.
 This script demonstrates how to load a model that was previously saved
 using save_sharded_state_offline.py and validates it by running inference.
 Example usage:
diff --git a/examples/features/sharded_state/save_sharded_state_offline.py b/examples/features/sharded_state/save_sharded_state_offline.py
index 675f2e35a53f..7ed0de1f1054 100644
--- a/examples/features/sharded_state/save_sharded_state_offline.py
+++ b/examples/features/sharded_state/save_sharded_state_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Saves each worker's model state dict directly to a checkpoint, which enables a
+"""Saves each worker's model state dict directly to a checkpoint, which enables a
 fast load path for large tensor-parallel models where each worker only needs to
 read its own shard rather than the entire checkpoint.
 
diff --git a/examples/features/speculative_decoding/mlpspeculator_offline.py b/examples/features/speculative_decoding/mlpspeculator_offline.py
index 6a533eb5c937..b863b28222e2 100644
--- a/examples/features/speculative_decoding/mlpspeculator_offline.py
+++ b/examples/features/speculative_decoding/mlpspeculator_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file demonstrates the usage of text generation with an LLM model,
+"""This file demonstrates the usage of text generation with an LLM model,
 comparing the performance with and without speculative decoding.
 
 Note that this example is out of date and not supported in vLLM v1.
diff --git a/examples/features/structured_diffusion/README.md b/examples/features/structured_diffusion/README.md
new file mode 100644
index 000000000000..9840f9ea3f01
--- /dev/null
+++ b/examples/features/structured_diffusion/README.md
@@ -0,0 +1,74 @@
+# Structured reads on DiffusionGemma
+
+A discrete diffusion model denoises a whole canvas per forward pass. If the
+canvas is seeded with the answer's fixed text and only the answer slots are
+left as noise, one denoise step gives a distribution over each slot. These
+`extra_args` fields (`vllm_xargs` on the OpenAI server) expose that:
+
+| field | type | meaning |
+| --- | --- | --- |
+| `diffusion_seed_canvas` | `list[int]`, exactly `canvas_length` ids | replaces the random initial canvas after prefill |
+| `diffusion_pinned` | `list[int]` of canvas positions | held at their seed value on every denoise step, so a read past one step keeps its template |
+| `diffusion_max_steps` | `int` | denoise steps before the canvas is emitted |
+| `diffusion_read_only` | `bool` | emit the argmax canvas as soon as the cap is reached, end the request there, and return temperature-1 logprobs at every position |
+
+`structured_server.py` turns a question schema into those fields. It serves
+`/v1/chat/completions`: the system message is the schema, the user message
+is the state JSON, and the reply content is one distribution per question
+with a standard error over a few noise draws.
+
+```bash
+vllm serve google/diffusiongemma-26B-A4B-it \
+    --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 --enable-prefix-caching
+python examples/features/structured_diffusion/structured_server.py \
+    --upstream http://127.0.0.1:8000 --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64
+curl -s localhost:8011/v1/chat/completions -H 'content-type: application/json' -d '{
+  "messages": [
+    {"role": "system", "content": "{\"questions\": [{\"id\": \"urgent\", \"type\": \"noul\", \"instructions\": \"Does the customer need a reply within the hour?\"}]}"},
+    {"role": "user", "content": "{\"ticket\": \"Everything is down and we have a demo at noon.\"}"}
+  ]}'
+```
+
+The attention backend is picked as for Gemma 4: FlashAttention 4 on every
+layer when available, otherwise Triton. FlashInfer cannot serve this model (a
+batch mixes causal prefill with bidirectional denoising), and
+`--attention-backend FLASHINFER` is rejected.
+
+Per-request canvas widths smaller than the served canvas require async
+scheduling. The diffusion async scheduler is selected automatically; no
+`--scheduler-cls` argument is needed. Synchronous execution supports full-width
+canvases only.
+
+Question types: `noul` (yes/no), `choice` with `options`, `score` with
+ordered `levels`. Each label must be a single token in the answer template,
+which the server checks with the tokenizer when a request uses the schema.
+
+`POST /v1/systemone` implements the Jev decision API. The body holds
+`state`, `questions` (a map of id to `type`, `instructions` and `criteria`)
+and `model`. Answers come back in that API's shapes: a `noul` probability,
+a `choice` with `probabilities` and `confidence`, or a `score` with a
+0-indexed `legend`. The schema options above go in the same body as
+extensions.
+
+A question may declare `depends_on` (read in a later stage with those
+answers in its prompt), `ask_if` (asked only when a named question's answer
+is among the listed ones, otherwise null) and `alone` (a read of its own).
+
+Images attach as `multipart/form-data`, with the JSON in a part named
+`request` and each image as a file part, or as an `images` array of data
+URLs.
+
+```bash
+curl -s localhost:8011/v1/systemone -H 'content-type: application/json' -d '{
+  "model": "jev-latest",
+  "state": {"ticket": "Everything is down and we have a demo at noon."},
+  "questions": {"urgent": {"type": "noul", "instructions": "Does the customer need a reply within the hour?"}}}'
+```
+
+`"think": N` in the schema lets the model write up to N tokens in its
+thought channel before the read. The thought is an ordinary generation with
+the chat template's thinking marker on, and the read then runs with the
+thought in its prompt, so the answer slots condition on it. The noise draws
+of a decision share one thought. `diagnostics.thought` returns the text, its
+length in tokens, whether the model closed the channel itself and the
+generation time.
diff --git a/examples/features/structured_diffusion/structured_server.py b/examples/features/structured_diffusion/structured_server.py
new file mode 100644
index 000000000000..a8f61da94e41
--- /dev/null
+++ b/examples/features/structured_diffusion/structured_server.py
@@ -0,0 +1,1528 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Structured decisions in front of a vLLM DiffusionGemma server.
+
+POST /v1/systemone takes Jev's request body: {"model", "state", "questions"}.
+"questions" maps an id to {"type", "instructions", "criteria"}, where "type"
+is "noul", "choice" or "score" and the criteria shape follows the type:
+  noul:   optional {"true": ..., "false": ...} descriptions
+  choice: option name -> description or null
+  score:  ordered list of levels
+Answers take Jev's shapes, with this server's diagnostics alongside:
+  noul:   {"noul": p}
+  choice: {"choice", "probabilities", "confidence"}
+  score:  {"score", "legend", "probabilities", "confidence"}
+The request body may also carry the schema keys "instructions", "samples",
+"auto_max", "auto_threshold", "steps", "think", "ask", "chunk_rows",
+"chunk_prompt" and "sequential" as extensions. Images go ahead of the
+state, either as multipart/form-data with the JSON body in a part named
+"request" and each image as a file part, or as an "images" array of data
+URLs in the JSON body.
+
+POST /v1/chat/completions makes the same decision from an OpenAI-shaped
+call. The system message is the schema JSON below and the user message is
+the state. The reply's `content` is the JSON answer set.
+
+POST /v1/raw/chat/completions passes the body to vLLM's chat completions
+unchanged, for plain generation through this port.
+
+With API_KEY set in the environment, every POST needs "Authorization:
+Bearer <key>". --tls-port adds an HTTPS listener with a self-signed
+certificate kept in --cert-dir, for clients that need a secure origin.
+
+Each answer is one distribution per question, from one denoise step over a
+seeded canvas, averaged over a few noise draws. This server handles the
+canvas, tokenizer, slot resolution, noise draws and averaging.
+
+Schema (system message):
+  {"questions": [
+     {"id": "urgent", "type": "noul", "instructions": "..."},
+     {"id": "bucket", "type": "choice", "instructions": "...",
+      "options": [{"name": "billing", "description": "..."}, ...]},
+     {"id": "tone", "type": "score", "instructions": "...",
+      "levels": ["calm", "annoyed", "furious"]}],
+   "instructions": "optional context",
+   "samples": "auto" | N, "auto_threshold": 0.1, "auto_max": 4,
+   "steps": 1, "think": 0}
+
+A question may also declare:
+  "depends_on": [ids]        answered after those, with their answers in
+                             its prompt
+  "ask_if": {id: [answers]}  asked only when that question's answer is
+                             among them (a skipped answer is null)
+  "alone": true              a read of its own
+Questions run in stages by these dependencies. Each stage is one joint
+read. Later stages continue the earlier answers, prefilled for a text
+state and restated for an image.
+
+Up to ten questions answer as "id: label" lines. Past that the id runs
+straight into the label, space separated, one row fewer per question. The
+server splits a schema whose answer template does not fit the canvas into
+chunks that run together, each with its own question list. "chunk_rows"
+sets the rows per chunk, "ask" picks a subset of question ids for one
+read, "sequential": true runs the chunks in order with the earlier answers
+prefilled, and "chunk_prompt": "shared" lists every question in each
+chunk's prompt. "think": N lets the model write up to N tokens in its
+thought channel, as an ordinary generation, and the read then runs with
+that thought in its prompt. With images the model writes the thought with
+the image in view and the server seeds it into the canvas ahead of the
+answer, so the canvas bounds it. The noise draws of a decision share one
+thought.
+
+Serve the model with a canvas that holds the answer template, for example:
+  vllm serve google/diffusiongemma-26B-A4B-it \
+      --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 \
+      --enable-prefix-caching
+then run this in front of it:
+  python structured_server.py --upstream http://127.0.0.1:8000 \
+      --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64 --port 8011
+"""
+
+import argparse
+import json
+import math
+import os
+import random
+import ssl
+import subprocess
+import threading
+import time
+import urllib.error
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import pybase64 as base64
+from transformers import AutoTokenizer
+
+ARGS = None
+TOK = None
+API_KEY = os.environ.get(
+    "API_KEY", ""
+)  # when set, POST routes need "Authorization: Bearer <key>"
+CANVAS_LEN = 64  # the served canvas length. A request may be narrower.
+CANVAS_STEP = 16  # request widths are multiples of this
+VOCAB = 262144
+TURN_CLOSE = 106
+PAD = 0
+TOPK = 20
+MAX_QUESTIONS = 64  # per request
+MAX_SAMPLES = 32  # reads per question, fixed or auto
+MAX_PARALLEL = 16  # question groups read at once
+# the empty thought block the chat template leaves to the model
+SCAFFOLD_TEXT = "<|channel>thought\n<channel|>"
+SCAFFOLD = None
+THOUGHT_OPEN = None
+THOUGHT_CLOSE = None
+
+
+# ----------------------------------------------------------------------------
+# Schema
+# ----------------------------------------------------------------------------
+
+
+class SchemaError(ValueError):
+    pass
+
+
+def parse_schema(value):
+    if (
+        not isinstance(value, dict)
+        or not isinstance(value.get("questions"), list)
+        or not value["questions"]
+    ):
+        raise SchemaError("schema: needs a non-empty questions array")
+    if len(value["questions"]) > MAX_QUESTIONS:
+        raise SchemaError(f"schema: at most {MAX_QUESTIONS} questions")
+    qs = []
+    seen = set()
+    for q in value["questions"]:
+        qid = str(q.get("id", "")).strip()
+        if not qid or ":" in qid or "\n" in qid:
+            raise SchemaError(
+                f"question id {qid!r} must be non-empty, no ':' or newline"
+            )
+        if qid in seen:
+            raise SchemaError(f"duplicate question id {qid!r}")
+        seen.add(qid)
+        kind = q.get("type")
+        if kind in ("noul", "bool", "boolean"):
+            kind = "noul"
+            crit = q.get("criteria") or {}
+            choices = [("yes", crit.get("true")), ("no", crit.get("false"))]
+            labels = ["yes", "no"]
+        elif kind == "choice":
+            opts = q.get("options") or []
+            choices = [
+                (o["name"], o.get("description"))
+                if isinstance(o, dict)
+                else (str(o), None)
+                for o in opts
+            ]
+            labels = [chr(ord("A") + i) for i in range(len(choices))]
+        elif kind == "score":
+            choices = [(str(level), None) for level in (q.get("levels") or [])]
+            labels = (
+                [str(i + 1) for i in range(len(choices))]
+                if len(choices) <= 9
+                else [chr(ord("A") + i) for i in range(len(choices))]
+            )
+        else:
+            raise SchemaError(f"question {qid!r}: unknown type {kind!r}")
+        if len(choices) < 2:
+            raise SchemaError(f"question {qid!r}: needs at least two alternatives")
+        if len(choices) > 26:
+            raise SchemaError(f"question {qid!r}: at most 26 alternatives")
+        deps = q.get("depends_on") or []
+        ask_if = q.get("ask_if") or {}
+        if not isinstance(deps, list) or not all(isinstance(d, str) for d in deps):
+            raise SchemaError(
+                f"question {qid!r}: depends_on must be a list of question ids"
+            )
+        if not isinstance(ask_if, dict) or not all(
+            isinstance(v, list) and v for v in ask_if.values()
+        ):
+            raise SchemaError(
+                f"question {qid!r}: ask_if must map a question id to a "
+                "non-empty list of its answers"
+            )
+        qs.append(
+            {
+                "id": qid,
+                "type": kind,
+                "instructions": str(q.get("instructions", "")),
+                "choices": choices,
+                "labels": labels,
+                "depends_on": list(dict.fromkeys(list(deps) + list(ask_if))),
+                "ask_if": ask_if,
+                "alone": bool(q.get("alone", False)),
+            }
+        )
+    by_id = {q["id"]: q for q in qs}
+    for q in qs:
+        for dep in q["depends_on"]:
+            if dep not in by_id or dep == q["id"]:
+                raise SchemaError(
+                    f"question {q['id']!r}: depends on unknown question {dep!r}"
+                )
+        for dep, vals in q["ask_if"].items():
+            names = [c[0] for c in by_id[dep]["choices"]]
+            if any(v not in names for v in vals):
+                raise SchemaError(
+                    f"question {q['id']!r}: ask_if values for {dep!r} "
+                    f"must be among {names}"
+                )
+    schedule(qs)  # refuses a cycle
+    samples = value.get("samples", "auto")
+    if samples == "auto":
+        policy = {
+            "mode": "auto",
+            "max": max(1, min(int(value.get("auto_max", 4)), MAX_SAMPLES)),
+            "threshold": float(value.get("auto_threshold", 0.1)),
+        }
+    elif isinstance(samples, int) and samples >= 1:
+        policy = {"mode": "fixed", "n": min(samples, MAX_SAMPLES)}
+    else:
+        raise SchemaError('schema: samples must be a positive count or "auto"')
+    ask = value.get("ask")
+    if ask is not None:
+        if not isinstance(ask, list) or not ask or any(a not in seen for a in ask):
+            raise SchemaError("schema: ask must list question ids from this schema")
+        for q in qs:
+            if q["id"] in ask and any(d not in ask for d in q["depends_on"]):
+                raise SchemaError(
+                    f"schema: ask names {q['id']!r} but not everything it depends on"
+                )
+    chunk_rows = value.get("chunk_rows")
+    if chunk_rows is not None and (not isinstance(chunk_rows, int) or chunk_rows < 8):
+        raise SchemaError("schema: chunk_rows must be an integer of at least 8")
+    chunk_prompt = value.get("chunk_prompt", "own")
+    if chunk_prompt not in ("shared", "own"):
+        raise SchemaError('schema: chunk_prompt must be "shared" or "own"')
+    sequential = bool(value.get("sequential", False))
+    think = value.get("think", 0)
+    if isinstance(think, bool) or not isinstance(think, int) or not 0 <= think <= 4096:
+        raise SchemaError("schema: think must be a thought budget in tokens, 0 to 4096")
+    return {
+        "questions": qs,
+        "instructions": value.get("instructions"),
+        "policy": policy,
+        "steps": max(1, min(int(value.get("steps", 1)), 8)),
+        "think": think,
+        "ask": ask,
+        "chunk_rows": chunk_rows,
+        "chunk_prompt": chunk_prompt,
+        "sequential": sequential,
+        "format": "lines" if len(qs) <= 10 else "indexed",
+    }
+
+
+# Answer template shape: (join between questions, what precedes the label,
+# reply instruction). A small schema gets "lines", which is readable.
+# "indexed" ("0yes 1no") costs three tokens a question against four or five
+# and agreed with "lines" on every set tried: 42 booleans, ten 26-way
+# choices, twenty 5-level scores. Past ten questions the saved rows keep a
+# schema in one read. Two tokens a question, or no id at all, loses
+# alignment beyond about twenty questions, because the id ties a label to
+# its question.
+FORMATS = {
+    "lines": (
+        "\n",
+        "{id}: ",
+        'Reply with one line per question, in this order, formatted as "id: label".',
+    ),
+    "indexed": (
+        " ",
+        "{id}",
+        "Reply on one line with each question's id immediately followed by its "
+        "label, separated by single spaces.",
+    ),
+}
+
+
+def system_text(schema, chunked=False):
+    s = (
+        "Answer a fixed set of questions about the state the user provides. "
+        "Each question lists its allowed answers; reply with exactly one label "
+        "per question.\n"
+    )
+    if schema.get("instructions"):
+        s += "\n" + str(schema["instructions"]).strip() + "\n"
+    for q in schema["questions"]:
+        s += f"\nQuestion {q['id']}: {q['instructions'].strip()}\n"
+        for (name, desc), label in zip(q["choices"], q["labels"]):
+            if q["type"] == "noul":
+                s += f"  {label}: {str(desc).strip()}\n" if desc else f"  {label}\n"
+            elif desc:
+                s += f"  {label}: {name} ({str(desc).strip()})\n"
+            else:
+                s += f"  {label}: {name}\n"
+    s += "\n" + FORMATS[schema.get("format", "lines")][2]
+    if chunked:
+        s += (
+            " A reply may cover only some of the questions; answer every line "
+            "that is present."
+        )
+    return s
+
+
+def answer_text(qs, labels, fmt="lines"):
+    join, lead, _ = FORMATS[fmt]
+    return join.join(
+        lead.format(id=q["id"]) + q["labels"][i] for q, i in zip(qs, labels)
+    )
+
+
+def enc(text):
+    return TOK.encode(text, add_special_tokens=False)
+
+
+def init_tokenizer(tok):
+    global TOK, SCAFFOLD, THOUGHT_OPEN, THOUGHT_CLOSE
+    TOK = tok
+    THOUGHT_OPEN = enc("<|channel>thought\n")
+    THOUGHT_CLOSE = enc("<channel|>")
+    SCAFFOLD = enc(SCAFFOLD_TEXT)
+    assert THOUGHT_OPEN + THOUGHT_CLOSE == SCAFFOLD, (
+        "the thought tags must tokenize apart"
+    )
+
+
+def resolve_template(qs, head, lead, fmt):
+    """Tokenize the answer template and find each question's slot. Every label
+    must change exactly one token, at the same position for all of a question's
+    labels, or this raises SchemaError. ``head`` is the token run the canvas
+    starts with: the empty thought block for a plain read, and empty when the
+    prompt already ends the thought channel. ``lead`` is the text before the
+    first answer: the join when earlier answers are in the prompt, so the
+    tokens match one joint template."""
+    base_labels = [0] * len(qs)
+    base = head + enc(lead + answer_text(qs, base_labels, fmt))
+    if len(base) + 1 > CANVAS_LEN:
+        raise SchemaError(
+            f"answer template is {len(base)} tokens; the canvas holds {CANVAS_LEN - 1}"
+        )
+    if len(qs) == 1 and len(base) + 1 > CANVAS_LEN:
+        raise SchemaError(
+            f"question {qs[0]['id']!r} alone needs {len(base) + 1} canvas rows"
+        )
+    slots = []
+    for qi, q in enumerate(qs):
+        pos = None
+        ids = [0] * len(q["labels"])
+        for li in range(1, len(q["labels"])):
+            labels = list(base_labels)
+            labels[qi] = li
+            e = head + enc(lead + answer_text(qs, labels, fmt))
+            if len(e) != len(base):
+                raise SchemaError(
+                    f"question {q['id']!r}: label {q['labels'][li]!r} is not a "
+                    "single token"
+                )
+            diffs = [i for i in range(len(e)) if e[i] != base[i]]
+            if len(diffs) != 1 or (pos is not None and diffs[0] != pos):
+                raise SchemaError(
+                    f"question {q['id']!r}: labels do not share one template slot"
+                )
+            pos = diffs[0]
+            ids[li] = e[pos]
+        ids[0] = base[pos]
+        if len(set(ids)) != len(ids):
+            raise SchemaError(
+                f"question {q['id']!r}: two labels tokenize to the same id"
+            )
+        slots.append({"pos": pos, "label_ids": ids})
+    return base, slots
+
+
+_template_cache = {}
+
+
+def template_for(schema, head, lead):
+    fmt = schema.get("format", "lines")
+    key = json.dumps(
+        [head, lead, fmt] + [(q["id"], q["labels"]) for q in schema["questions"]]
+    )
+    if key not in _template_cache:
+        _template_cache[key] = resolve_template(schema["questions"], head, lead, fmt)
+    return _template_cache[key]
+
+
+# ----------------------------------------------------------------------------
+# Reads
+# ----------------------------------------------------------------------------
+
+
+def canvas_width(template):
+    """Smallest multiple of CANVAS_STEP that holds the template and the turn close."""
+    need = len(template) + 1
+    return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP)
+
+
+def pin_xargs(template, slots, steps):
+    """Past one denoise step the template must be held, or accept/renoise
+    rewrites it: pin every canvas position that is not an answer slot."""
+    if steps <= 1:
+        return {}
+    free = {s["pos"] for s in slots}
+    return {
+        "diffusion_pinned": [p for p in range(canvas_width(template)) if p not in free]
+    }
+
+
+def build_canvas(template, slots, seed):
+    rng = random.Random(seed)
+    canvas = list(template) + [TURN_CLOSE]
+    canvas += [PAD] * (canvas_width(template) - len(canvas))
+    for s in slots:
+        canvas[s["pos"]] = rng.randrange(VOCAB)
+    return canvas
+
+
+def label_id_union(slots):
+    ids = sorted({i for s in slots for i in s["label_ids"]})
+    return ids[:128]  # vLLM's cap per request. A schema needs far fewer.
+
+
+def upstream_chat(body, timeout=600):
+    req = urllib.request.Request(
+        ARGS.upstream.rstrip("/") + "/v1/chat/completions",
+        data=json.dumps(body).encode(),
+        headers={"content-type": "application/json"},
+    )
+    return json.load(urllib.request.urlopen(req, timeout=timeout))
+
+
+def upstream_completions(body, timeout=600):
+    req = urllib.request.Request(
+        ARGS.upstream.rstrip("/") + "/v1/completions",
+        data=json.dumps(body).encode(),
+        headers={"content-type": "application/json"},
+    )
+    return json.load(urllib.request.urlopen(req, timeout=timeout))
+
+
+def chat_prompt_ids(sys_text, state_text, thinking=False):
+    """The prompt the chat endpoint would build, as token ids, ending after
+    the model turn marker. Text states only. ``thinking`` turns the chat
+    template's thinking marker on."""
+    messages = [
+        {"role": "system", "content": sys_text},
+        {"role": "user", "content": state_text},
+    ]
+    out = TOK.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True, enable_thinking=thinking
+    )
+    ids = (
+        out["input_ids"] if hasattr(out, "keys") else out
+    )  # newer transformers return a dict
+    return [int(t) for t in ids]
+
+
+def think(sys_text, state_text, budget):
+    """A read prefix that ends with a thought the model wrote: the chat prompt
+    with thinking on, the open tag, up to ``budget`` generated tokens, the
+    close tag. Returns the prefix and a diagnostics dict for the thought."""
+    prompt = chat_prompt_ids(sys_text, state_text, thinking=True) + THOUGHT_OPEN
+    started = time.time()
+    d = upstream_completions(
+        {
+            "model": ARGS.model,
+            "prompt": prompt,
+            "max_tokens": budget,
+            "logprobs": 0,
+            "return_tokens_as_token_ids": True,
+            "stop_token_ids": THOUGHT_CLOSE,
+        }
+    )
+    ids = [int(t.split(":")[1]) for t in d["choices"][0]["logprobs"]["tokens"]]
+    closed = THOUGHT_CLOSE[0] in ids
+    if closed:
+        ids = ids[: ids.index(THOUGHT_CLOSE[0])]
+    info = {
+        "tokens": len(ids),
+        "closed": closed,
+        "ms": (time.time() - started) * 1e3,
+        "text": TOK.decode(ids),
+    }
+    return prompt + ids + THOUGHT_CLOSE, info
+
+
+def think_chat(sys_text, state_content, budget):
+    """A thought written with the image in view: the chat endpoint with
+    thinking on, capped at ``budget`` tokens and cut at the close tag.
+    Returns the thought's token ids and a diagnostics dict."""
+    body = {
+        "model": ARGS.model,
+        "messages": [
+            {"role": "system", "content": sys_text},
+            {"role": "user", "content": state_content},
+        ],
+        "max_tokens": budget,
+        "logprobs": True,
+        "top_logprobs": 0,
+        "return_tokens_as_token_ids": True,
+        "stop_token_ids": THOUGHT_CLOSE,
+        "chat_template_kwargs": {"enable_thinking": True},
+    }
+    started = time.time()
+    d = upstream_chat(body)
+    choice = d["choices"][0]
+    ids = [
+        int(t["token"].split(":")[1])
+        for t in (choice.get("logprobs") or {}).get("content") or []
+    ]
+    if ids[: len(THOUGHT_OPEN)] == THOUGHT_OPEN:  # the model opens the channel itself
+        ids = ids[len(THOUGHT_OPEN) :]
+    closed = THOUGHT_CLOSE[0] in ids
+    if closed:
+        ids = ids[: ids.index(THOUGHT_CLOSE[0])]
+    return ids, {
+        "tokens": len(ids),
+        "closed": closed,
+        "ms": (time.time() - started) * 1e3,
+        "text": TOK.decode(ids),
+    }
+
+
+def one_read(
+    schema, template, slots, sys_text, state_content, seed, prefix=None, thinking=False
+):
+    if prefix is not None:
+        return one_read_continuation(schema, template, slots, prefix, seed)
+    messages = [
+        {"role": "system", "content": sys_text},
+        {"role": "user", "content": state_content},
+    ]
+    body = {
+        "model": ARGS.model,
+        "messages": messages,
+        "max_tokens": len(template) + 1,
+        "logprobs": True,
+        "top_logprobs": TOPK,
+        # Exact logprobs for every label at every position. With a long
+        # option list most labels never rank in the top-k, and the model's
+        # mass sits on tokens that spell the option name instead.
+        "logprob_token_ids": label_id_union(slots),
+        "return_tokens_as_token_ids": True,
+        "chat_template_kwargs": {"enable_thinking": thinking},
+        "vllm_xargs": {
+            "diffusion_seed_canvas": build_canvas(template, slots, seed),
+            "diffusion_canvas_length": canvas_width(template),
+            "diffusion_max_steps": schema["steps"],
+            "diffusion_read_only": True,
+            **pin_xargs(template, slots, schema["steps"]),
+        },
+    }
+    d = upstream_chat(body)
+    content = d["choices"][0]["logprobs"]["content"]
+    out = []
+    for q, s in zip(schema["questions"], slots):
+        top = {
+            int(t["token"].split(":")[1]): t["logprob"]
+            for t in content[s["pos"]]["top_logprobs"]
+        }
+        out.append(slot_distribution(top, s["label_ids"]))
+    return out, d.get("usage", {})
+
+
+def slot_distribution(top, label_ids):
+    """Label probabilities at one slot from the returned logprobs: every
+    label's own value plus the argmax token. Read-only logprobs are at
+    temperature 1, so the label softmax uses them directly. The entropy is
+    over that returned set."""
+    floor = min(top.values()) - 5.0
+    lp_t = [top.get(i, floor) for i in label_ids]
+    mx = max(lp_t)
+    ex = [math.exp(x - mx) for x in lp_t]
+    probs = [e / sum(ex) for e in ex]
+    top_p = [math.exp(v) for v in top.values()]
+    return {
+        "probs": probs,
+        "label_mass": sum(math.exp(x) for x in lp_t),
+        "entropy": -sum(p * math.log(p) for p in top_p if p > 0),
+        "argmax_is_label": max(top, key=top.get) in label_ids,
+    }
+
+
+def one_read_continuation(schema, template, slots, prompt_ids, seed):
+    """A read whose prompt already holds the thought scaffold and earlier
+    answer lines, sent as token ids so the chat template cannot alter it."""
+    body = {
+        "model": ARGS.model,
+        "prompt": prompt_ids,
+        "max_tokens": len(template) + 1,
+        "logprobs": TOPK,
+        "logprob_token_ids": label_id_union(slots),
+        "return_tokens_as_token_ids": True,
+        "vllm_xargs": {
+            "diffusion_seed_canvas": build_canvas(template, slots, seed),
+            "diffusion_canvas_length": canvas_width(template),
+            "diffusion_max_steps": schema["steps"],
+            "diffusion_read_only": True,
+            **pin_xargs(template, slots, schema["steps"]),
+        },
+    }
+    d = upstream_completions(body)
+    rows = d["choices"][0]["logprobs"]["top_logprobs"]
+    out = []
+    for q, sl in zip(schema["questions"], slots):
+        top = {int(k.split(":")[1]): v for k, v in rows[sl["pos"]].items()}
+        out.append(slot_distribution(top, sl["label_ids"]))
+    return out, d.get("usage", {})
+
+
+def read_many(
+    schema,
+    template,
+    slots,
+    sys_text,
+    state_content,
+    seed,
+    n,
+    prefix=None,
+    thinking=False,
+):
+    results = [None] * n
+    errors = [None] * n
+    usages = [None] * n
+
+    def run(k):
+        try:
+            results[k], usages[k] = one_read(
+                schema,
+                template,
+                slots,
+                sys_text,
+                state_content,
+                seed + k * 7919,
+                prefix,
+                thinking,
+            )
+        except Exception as e:  # raised again below
+            errors[k] = e
+
+    threads = [threading.Thread(target=run, args=(k,)) for k in range(n)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+    for e in errors:
+        if e is not None:
+            raise e
+    return results, usages
+
+
+def schedule(qs):
+    """Questions in stages: a question's stage comes after the stages of
+    everything it depends on. Declaration order is kept within a stage."""
+    ids = {q["id"] for q in qs}
+    pending = list(qs)
+    done: set = set()
+    levels = []
+    while pending:
+        level = [
+            q
+            for q in pending
+            if all(d in done or d not in ids for d in q["depends_on"])
+        ]
+        if not level:
+            raise SchemaError(
+                "schema: dependency cycle among " + ", ".join(q["id"] for q in pending)
+            )
+        levels.append(level)
+        done |= {q["id"] for q in level}
+        pending = [q for q in pending if q["id"] not in done]
+    return levels
+
+
+def chunk_groups(schema, qs):
+    """``qs`` split, in order, into the fewest groups whose answer templates
+    fit ``chunk_rows`` (the canvas by default). A question marked alone gets
+    its own group."""
+    limit = schema.get("chunk_rows") or CANVAS_LEN
+    groups, group = [], []
+    for q in qs:
+        if q["alone"]:
+            if group:
+                groups.append(group)
+                group = []
+            groups.append([q])
+            continue
+        trial = group + [q]
+        rows = (
+            len(SCAFFOLD)
+            + len(
+                enc(answer_text(trial, [0] * len(trial), schema.get("format", "lines")))
+            )
+            + 1
+        )
+        if rows > limit and group:
+            groups.append(group)
+            group = [q]
+        else:
+            group = trial
+    if group:
+        groups.append(group)
+    return groups
+
+
+def answer_name(q, a):
+    """The answer as the name ask_if compares against: yes or no, an option
+    name, or a level name."""
+    if a is None:
+        return None
+    return a["label"] if q["type"] == "noul" else a.get("choice", a.get("level"))
+
+
+def decide(schema, state_content, seed):
+    """One decision. Questions run in stages by their dependencies. A stage
+    is one joint read, chunked by the canvas, with a question marked alone
+    in its own read. Later stages condition on every earlier answer: a
+    prefilled continuation for a text state, or the answers restated in the
+    state for an image. A question whose ask_if condition failed is skipped
+    and its answer is null."""
+    started = time.time()
+    qs = [
+        q
+        for q in schema["questions"]
+        if not schema.get("ask") or q["id"] in schema["ask"]
+    ]
+    levels = schedule(qs)
+    text_state = isinstance(state_content, str)
+    fmt = schema["format"]
+    join = FORMATS[fmt][0]
+    # More than one read in sequence needs the full question list in every
+    # prompt, so that later reads continue one answer.
+    chained = len(levels) > 1 or schema["sequential"]
+    sys_full = system_text(schema, chunked=not text_state and chained)
+    base_ids, thought = None, None
+    if chained and text_state:
+        if schema["think"]:
+            base_ids, thought = think(sys_full, state_content, schema["think"])
+        else:
+            base_ids = chat_prompt_ids(sys_full, state_content) + SCAFFOLD
+    shared = schema["chunk_prompt"] == "shared"
+    answered, lines, earlier = {}, [], []
+    parts, stages, chunks, skipped = [], [], [], {}
+    by_id = {q["id"]: q for q in schema["questions"]}
+
+    def run(group, k, conditioned):
+        sub = dict(schema, questions=group)
+        # The thought is written once: in the prefix of a chained text decision,
+        # or in the first read of anything else.
+        sub["think"] = 0 if (chained and text_state) or conditioned else schema["think"]
+        if text_state:
+            if conditioned:
+                prefix, lead, sys_text = (
+                    base_ids + enc(join.join(lines)),
+                    join,
+                    sys_full,
+                )
+            elif chained:
+                prefix, lead, sys_text = (base_ids if thought else None), "", sys_full
+            else:
+                prefix, lead, sys_text = (
+                    None,
+                    "",
+                    (system_text(schema, chunked=True) if shared else system_text(sub)),
+                )
+            return decide_group(
+                sub, sys_text, state_content, seed + 104729 * k, prefix, lead
+            )
+        state = state_content
+        if conditioned:
+            text = next(
+                (p["text"] for p in state_content if p.get("type") == "text"), ""
+            )
+            text += "\n\nAnswers so far:\n" + "\n".join(earlier)
+            state = [p for p in state_content if p.get("type") != "text"] + [
+                {"type": "text", "text": text}
+            ]
+            sys_text = sys_full
+        else:
+            sys_text = sys_full if (chained or shared) else system_text(sub)
+        return decide_group(sub, sys_text, state, seed + 104729 * k)
+
+    def absorb(group, body, rows):
+        parts.append((body, rows))
+        chunks.append([q["id"] for q in group])
+        answered.update(body["answers"])
+        lines.append(
+            answer_text(
+                group,
+                [q["labels"].index(body["answers"][q["id"]]["label"]) for q in group],
+                fmt,
+            )
+        )
+        earlier.extend(
+            f"{q['id']}: {answer_name(q, body['answers'][q['id']])}" for q in group
+        )
+
+    k = 0
+    for level in levels:
+        asked = []
+        for q in level:
+            failed = next(
+                (
+                    (dep, vals)
+                    for dep, vals in q["ask_if"].items()
+                    if answer_name(by_id[dep], answered.get(dep)) not in vals
+                ),
+                None,
+            )
+            if failed:
+                answered[q["id"]] = None
+                skipped[q["id"]] = {
+                    "because": failed[0],
+                    "was": answer_name(by_id[failed[0]], answered.get(failed[0])),
+                    "wanted": failed[1],
+                }
+                continue
+            asked.append(q)
+        if not asked:
+            continue
+        stages.append([q["id"] for q in asked])
+        groups = chunk_groups(schema, asked)
+        conditioned = bool(lines)
+        if schema["sequential"] or len(groups) == 1:
+            for group in groups:
+                body, rows = run(
+                    group, k, conditioned or (schema["sequential"] and bool(lines))
+                )
+                absorb(group, body, rows)
+                k += 1
+        else:
+            with ThreadPoolExecutor(max_workers=min(len(groups), MAX_PARALLEL)) as ex:
+                results = list(
+                    ex.map(
+                        lambda gk, conditioned=conditioned: run(
+                            gk[1], gk[0], conditioned
+                        ),
+                        [(k + i, g) for i, g in enumerate(groups)],
+                    )
+                )
+            for group, (body, rows) in zip(groups, results):
+                absorb(group, body, rows)
+            k += len(groups)
+
+    answers = {q["id"]: answered.get(q["id"]) for q in qs}
+    diag_q = {}
+    for body, _ in parts:
+        diag_q.update(body["diagnostics"]["questions"])
+    # A thought written here (a chained text decision) is outside every
+    # group's row count. One written inside a group is already counted there.
+    extra_rows = thought["tokens"] if thought else 0
+    if thought is None:
+        thoughts = [b["diagnostics"].get("thought") for b, _ in parts]
+        thought = (
+            thoughts[0] if len(parts) == 1 else ([t for t in thoughts if t] or None)
+        )
+    one = len(parts) == 1 and not skipped
+    diagnostics = {
+        "steps": schema["steps"],
+        "stages": stages,
+        "skipped": skipped,
+        "chunks": chunks,
+        "chunk_prompt": "full" if chained else schema["chunk_prompt"],
+        "sequential": schema["sequential"],
+        "conditioning": (
+            None
+            if len(stages) <= 1 and not schema["sequential"]
+            else ("prefill" if text_state else "restated")
+        ),
+        "thought": thought,
+        "samples": (
+            parts[0][0]["diagnostics"]["samples"]
+            if one
+            else {
+                "n": [b["diagnostics"]["samples"]["n"] for b, _ in parts],
+                "tops": [b["diagnostics"]["samples"]["tops"] for b, _ in parts],
+                "policy": [b["diagnostics"]["samples"]["policy"] for b, _ in parts],
+            }
+        ),
+        "timing": {
+            "total_ms": (time.time() - started) * 1e3,
+            "reads": sum(b["diagnostics"]["timing"]["reads"] for b, _ in parts),
+        },
+        "prompt_tokens": max(
+            (b["diagnostics"].get("prompt_tokens") or 0) for b, _ in parts
+        )
+        or None,
+        "questions": diag_q,
+        "engine": "vllm",
+    }
+    return {"answers": answers, "diagnostics": diagnostics}, sum(
+        rows for _, rows in parts
+    ) + extra_rows
+
+
+def decide_group(schema, sys_text, state_content, seed, prefix=None, lead=""):
+    started = time.time()
+    thought = None
+    head = SCAFFOLD if prefix is None else []
+    thinking = False
+    if prefix is None and schema["think"]:
+        if isinstance(state_content, str):
+            prefix, thought = think(sys_text, state_content, schema["think"])
+            head = []
+        else:
+            # With images the model writes the thought with the image in view, and
+            # the server seeds it into the canvas ahead of the answer, so the read
+            # keeps the image. The canvas bounds the thought.
+            answer = enc(
+                answer_text(
+                    schema["questions"],
+                    [0] * len(schema["questions"]),
+                    schema.get("format", "lines"),
+                )
+            )
+            fits = CANVAS_LEN - 1 - len(THOUGHT_OPEN) - len(THOUGHT_CLOSE) - len(answer)
+            if fits < 8:
+                raise SchemaError(
+                    f"think: the canvas leaves {fits} rows for a thought "
+                    "beside this template"
+                )
+            ids, thought = think_chat(
+                sys_text, state_content, min(schema["think"], fits)
+            )
+            thought["budget"] = min(schema["think"], fits)
+            head = THOUGHT_OPEN + ids + THOUGHT_CLOSE
+            thinking = True
+    template, slots = template_for(schema, head, lead)
+    policy = schema["policy"]
+    if policy["mode"] == "fixed":
+        reads, usages = read_many(
+            schema,
+            template,
+            slots,
+            sys_text,
+            state_content,
+            seed,
+            policy["n"],
+            prefix,
+            thinking,
+        )
+        extended = None
+        first_entropy = None
+    else:
+        reads, usages = read_many(
+            schema, template, slots, sys_text, state_content, seed, 1, prefix, thinking
+        )
+        first_entropy = {
+            q["id"]: r["entropy"] for q, r in zip(schema["questions"], reads[0])
+        }
+        extended = (
+            max(first_entropy.values()) > policy["threshold"] and policy["max"] > 1
+        )
+        if extended:
+            more, more_usages = read_many(
+                schema,
+                template,
+                slots,
+                sys_text,
+                state_content,
+                seed + 1,
+                policy["max"] - 1,
+                prefix,
+                thinking,
+            )
+            reads += more
+            usages += more_usages
+    prompt_tokens = next(
+        (u["prompt_tokens"] for u in usages if u and u.get("prompt_tokens")), None
+    )
+    elapsed_ms = (time.time() - started) * 1e3
+
+    answers = {}
+    diag_q = {}
+    n = len(reads)
+    for qi, q in enumerate(schema["questions"]):
+        per = [r[qi]["probs"] for r in reads]
+        mean = [sum(p[i] for p in per) / n for i in range(len(q["labels"]))]
+        top = max(range(len(mean)), key=lambda i: mean[i])
+        a = {
+            "type": q["type"],
+            "label": q["labels"][top],
+            "confidence": mean[top],
+            "probabilities": {c[0]: m for c, m in zip(q["choices"], mean)},
+        }
+        if q["type"] == "noul":
+            a["noul"] = mean[0]
+        elif q["type"] == "choice":
+            a["choice"] = q["choices"][top][0]
+        else:
+            a["score"] = sum((i + 1) * m for i, m in enumerate(mean))
+            a["level"] = q["choices"][top][0]
+        if n > 1:
+            var = sum((p[top] - mean[top]) ** 2 for p in per) / (n - 1)
+            a["stderr"] = (var / n) ** 0.5
+            a["agreement"] = (
+                sum(1 for p in per if max(range(len(p)), key=lambda i: p[i]) == top) / n
+            )
+        answers[q["id"]] = a
+        diag_q[q["id"]] = {
+            "pos": slots[qi]["pos"],
+            "entropy": [r[qi]["entropy"] for r in reads],
+            "label_mass": reads[0][qi]["label_mass"],
+            "argmax_is_label": reads[0][qi]["argmax_is_label"],
+        }
+    tops = [
+        {
+            q["id"]: [
+                q["labels"][
+                    max(range(len(r[qi]["probs"])), key=lambda i: r[qi]["probs"][i])
+                ],
+                max(r[qi]["probs"]),
+                r[qi]["entropy"],
+            ]
+            for qi, q in enumerate(schema["questions"])
+        }
+        for r in reads
+    ]
+    return {
+        "answers": answers,
+        "diagnostics": {
+            "steps": schema["steps"],
+            "samples": {
+                "n": n,
+                "tops": tops,
+                "policy": dict(
+                    policy, extended=extended, first_read_entropy=first_entropy
+                ),
+            },
+            "timing": {"total_ms": elapsed_ms, "reads": n},
+            "thought": thought,
+            "prompt_tokens": prompt_tokens,
+            "questions": diag_q,
+            "engine": "vllm",
+        },
+    }, len(template) + 1 + (thought["tokens"] if thought else 0)
+
+
+# ----------------------------------------------------------------------------
+# Jev's contract
+# ----------------------------------------------------------------------------
+
+JEV_EXTENSIONS = (
+    "instructions",
+    "samples",
+    "auto_max",
+    "auto_threshold",
+    "steps",
+    "think",
+    "ask",
+    "chunk_rows",
+    "chunk_prompt",
+    "sequential",
+)
+
+
+def jev_schema(body):
+    """This server's schema from a Jev request body."""
+    qs = body.get("questions")
+    if not isinstance(qs, dict) or not qs:
+        raise SchemaError("questions: needs a non-empty map of id -> question")
+    out = []
+    for qid, q in qs.items():
+        if not isinstance(q, dict):
+            raise SchemaError(f"question {qid!r}: must be an object")
+        kind, crit, ins = q.get("type"), q.get("criteria"), q.get("instructions", "")
+        item = {
+            "id": qid,
+            "type": kind,
+            "instructions": ins if isinstance(ins, str) else json.dumps(ins),
+        }
+        if kind == "noul":
+            if crit is not None and not isinstance(crit, dict):
+                raise SchemaError(
+                    f"question {qid!r}: noul criteria must be an object with "
+                    "true and false"
+                )
+            item["criteria"] = crit
+        elif kind == "choice":
+            if not isinstance(crit, dict) or not crit:
+                raise SchemaError(
+                    f"question {qid!r}: choice criteria must map option names "
+                    "to descriptions"
+                )
+            item["options"] = [
+                {"name": str(n), "description": d} for n, d in crit.items()
+            ]
+        elif kind == "score":
+            if not isinstance(crit, list):
+                raise SchemaError(
+                    f"question {qid!r}: score criteria must be an ordered list "
+                    "of levels"
+                )
+            item["levels"] = crit
+        else:
+            raise SchemaError(f"question {qid!r}: unknown type {kind!r}")
+        for key in ("depends_on", "ask_if", "alone"):
+            if key in q:
+                item[key] = q[key]
+        out.append(item)
+    schema = {k: body[k] for k in JEV_EXTENSIONS if k in body}
+    schema["questions"] = out
+    return parse_schema(schema)
+
+
+def image_part(content_type, data):
+    return {
+        "type": "image_url",
+        "image_url": {
+            "url": f"data:{content_type};base64," + base64.b64encode(data).decode()
+        },
+    }
+
+
+def jev_images(value):
+    """Image parts from the body's "images": data URLs, or objects with
+    content_type and base64."""
+    parts = []
+    for i, im in enumerate(value or []):
+        if isinstance(im, str) and im.startswith("data:image/"):
+            parts.append({"type": "image_url", "image_url": {"url": im}})
+        elif (
+            isinstance(im, dict)
+            and str(im.get("content_type", "")).startswith("image/")
+            and isinstance(im.get("base64"), str)
+        ):
+            parts.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:{im['content_type']};base64,{im['base64']}"
+                    },
+                }
+            )
+        else:
+            raise SchemaError(
+                f"images[{i}]: a data:image/... URL or an object with "
+                "content_type and base64"
+            )
+    return parts
+
+
+def jev_state(body, image_parts=()):
+    """The user message: the state as text (as given, or as JSON), with any
+    images ahead of it."""
+    state = body.get("state")
+    if state is None:
+        raise SchemaError("state: required")
+    text = state if isinstance(state, str) else json.dumps(state)
+    if not image_parts:
+        return text
+    return list(image_parts) + [{"type": "text", "text": text}]
+
+
+def jev_answer(q, a):
+    if a is None:
+        return None
+    if q["type"] == "noul":
+        return {"type": "noul", "noul": a["noul"]}
+    if q["type"] == "choice":
+        return {
+            "type": "choice",
+            "choice": a["choice"],
+            "probabilities": a["probabilities"],
+            "confidence": a["confidence"],
+        }
+    names = [c[0] for c in q["choices"]]
+    probs = {str(i): a["probabilities"][n] for i, n in enumerate(names)}
+    return {
+        "type": "score",
+        "score": sum(i * p for i, p in enumerate(probs.values())),
+        "legend": {str(i): n for i, n in enumerate(names)},
+        "probabilities": probs,
+        "confidence": a["confidence"],
+    }
+
+
+# ----------------------------------------------------------------------------
+# HTTP
+# ----------------------------------------------------------------------------
+
+
+def message_text(m):
+    c = m.get("content", "")
+    if isinstance(c, list):
+        return "".join(p.get("text", "") for p in c if isinstance(p, dict))
+    return c if isinstance(c, str) else ""
+
+
+class Handler(BaseHTTPRequestHandler):
+    def log_message(self, fmt, *args):
+        pass
+
+    def _json(self, code, obj):
+        body = json.dumps(obj).encode()
+        self.send_response(code)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_GET(self):
+        if self.path == "/health":
+            return self._json(200, {"status": "ok"})
+        return self._json(404, {"error": {"message": "unknown route"}})
+
+    def _read_request(self):
+        """-> (body, image parts): a JSON body, or multipart/form-data with the
+        JSON in a part named request and each image as a file part, in order."""
+        raw = self.rfile.read(int(self.headers.get("content-length", "0")))
+        ctype = self.headers.get("content-type", "")
+        if not ctype.lower().startswith("multipart/form-data"):
+            return json.loads(raw), []
+        from email.parser import BytesParser
+        from email.policy import HTTP
+
+        msg = BytesParser(policy=HTTP).parsebytes(
+            b"Content-Type: " + ctype.encode() + b"\r\n\r\n" + raw
+        )
+        body, images = None, []
+        for part in msg.iter_parts():
+            name = part.get_param("name", header="content-disposition")
+            data = part.get_payload(decode=True)
+            if name == "request":
+                body = json.loads(data)
+            elif part.get_content_type().startswith("image/"):
+                images.append(image_part(part.get_content_type(), data))
+            else:
+                raise ValueError(
+                    f"part {name!r}: neither the request JSON nor an image"
+                )
+        if body is None:
+            raise ValueError(
+                "multipart needs a part named request holding the JSON body"
+            )
+        return body, images
+
+    def do_POST(self):
+        if API_KEY and self.headers.get("authorization", "") != f"Bearer {API_KEY}":
+            return self._json(
+                401,
+                {
+                    "error": {
+                        "message": "missing or wrong bearer token",
+                        "type": "authentication_error",
+                    }
+                },
+            )
+        if self.path == "/v1/raw/chat/completions":
+            return self._raw_chat()
+        try:
+            req, images = self._read_request()
+        except Exception as e:
+            return self._json(
+                400,
+                {
+                    "error": {
+                        "message": f"invalid body: {e}",
+                        "type": "invalid_request_error",
+                    }
+                },
+            )
+        if self.path == "/v1/systemone":
+            return self._systemone(req, images)
+        if self.path == "/v1/chat/completions":
+            return self._chat(req)
+        return self._json(404, {"error": {"message": "unknown route"}})
+
+    def _raw_chat(self):
+        """Pass the body and status through to vLLM's chat completions."""
+        raw = self.rfile.read(int(self.headers.get("content-length", "0")))
+        req = urllib.request.Request(
+            ARGS.upstream.rstrip("/") + "/v1/chat/completions",
+            data=raw,
+            headers={
+                "content-type": self.headers.get("content-type", "application/json")
+            },
+        )
+        try:
+            with urllib.request.urlopen(req, timeout=600) as r:
+                code, body = r.status, r.read()
+        except urllib.error.HTTPError as e:
+            code, body = e.code, e.read()
+        self.send_response(code)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def _decide(self, schema, state, seed):
+        """-> (status, body) with the error body already shaped."""
+        try:
+            return 200, decide(schema, state, seed)
+        except SchemaError as e:
+            return 422, {"error": {"message": str(e), "type": "validation_error"}}
+        except urllib.error.HTTPError as e:
+            return 502, {
+                "error": {
+                    "message": (
+                        f"upstream {e.code}: {e.read()[:300].decode(errors='replace')}"
+                    ),
+                    "type": "server_error",
+                }
+            }
+        except Exception as e:
+            return 500, {"error": {"message": repr(e), "type": "server_error"}}
+
+    def _systemone(self, req, images):
+        try:
+            schema = jev_schema(req)
+            state = jev_state(req, images + jev_images(req.get("images")))
+        except SchemaError as e:
+            return self._json(
+                422, {"error": {"message": str(e), "type": "validation_error"}}
+            )
+        code, result = self._decide(schema, state, int(req.get("seed", 42)))
+        if code != 200:
+            return self._json(code, result)
+        body, completion_tokens = result
+        answers = {
+            q["id"]: jev_answer(q, body["answers"][q["id"]])
+            for q in schema["questions"]
+        }
+        labels = " ".join(
+            f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items()
+        )
+        print(
+            f"systemone: {labels} "
+            f"reads={body['diagnostics']['timing']['reads']} "
+            f"{body['diagnostics']['timing']['total_ms']:.0f}ms",
+            flush=True,
+        )
+        self._json(
+            200,
+            {
+                "model": ARGS.model,
+                "answers": answers,
+                # vLLM's prompt count when the reads reported one, since it covers
+                # images. Otherwise the tokenizer's count of the text prompt.
+                "usage": {
+                    "input_tokens": body["diagnostics"].get("prompt_tokens")
+                    or (
+                        len(chat_prompt_ids(system_text(schema), state))
+                        if isinstance(state, str)
+                        else 0
+                    ),
+                    "output_tokens": completion_tokens,
+                },
+                "diagnostics": body["diagnostics"],
+            },
+        )
+
+    def _chat(self, req):
+        msgs = req.get("messages") or []
+        if (
+            len(msgs) != 2
+            or msgs[0].get("role") not in ("system", "developer")
+            or msgs[1].get("role") != "user"
+        ):
+            return self._json(
+                400,
+                {
+                    "error": {
+                        "message": (
+                            "a structured request is exactly two messages: the "
+                            "schema (system) and the state JSON (user)"
+                        ),
+                        "type": "invalid_request_error",
+                    }
+                },
+            )
+        try:
+            schema_value = json.loads(message_text(msgs[0]))
+            schema = parse_schema(schema_value)
+            content = msgs[1].get("content", "")
+            has_image = isinstance(content, list) and any(
+                isinstance(p, dict) and p.get("type") in ("image_url", "image")
+                for p in content
+            )
+            if has_image:
+                # image parts pass through to vLLM unchanged, with text parts as context
+                state = content
+            else:
+                state = message_text(msgs[1]).strip()
+                json.loads(state)
+        except SchemaError as e:
+            return self._json(
+                400, {"error": {"message": str(e), "type": "invalid_request_error"}}
+            )
+        except Exception as e:
+            return self._json(
+                400,
+                {
+                    "error": {
+                        "message": (
+                            "system must be a JSON question schema and user "
+                            f"must be JSON state or image parts: {e}"
+                        ),
+                        "type": "invalid_request_error",
+                    }
+                },
+            )
+        code, result = self._decide(schema, state, int(req.get("seed", 42)))
+        if code != 200:
+            if code == 422:
+                result["error"]["type"] = "invalid_request_error"
+                code = 400
+            return self._json(code, result)
+        body, completion_tokens = result
+        content = json.dumps(body, indent=2)
+        labels = " ".join(
+            f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items()
+        )
+        print(
+            f"structured: {labels} "
+            f"reads={body['diagnostics']['timing']['reads']} "
+            f"{body['diagnostics']['timing']['total_ms']:.0f}ms",
+            flush=True,
+        )
+        self._json(
+            200,
+            {
+                "id": f"chatcmpl-{int(time.time() * 1000)}",
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": req.get("model", "dgemma-structured"),
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": content},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 0,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": completion_tokens,
+                },
+            },
+        )
+
+
+def self_signed(cert_dir):
+    """Paths of a self-signed certificate and key in cert_dir, made with
+    openssl on first use."""
+    os.makedirs(cert_dir, exist_ok=True)
+    cert, key = os.path.join(cert_dir, "djev.crt"), os.path.join(cert_dir, "djev.key")
+    if not (os.path.exists(cert) and os.path.exists(key)):
+        subprocess.run(
+            [
+                "openssl",
+                "req",
+                "-x509",
+                "-newkey",
+                "rsa:2048",
+                "-nodes",
+                "-days",
+                "3650",
+                "-subj",
+                "/CN=djev",
+                "-keyout",
+                key,
+                "-out",
+                cert,
+            ],
+            check=True,
+            capture_output=True,
+        )
+    return cert, key
+
+
+def serve_tls(host, port, cert_dir):
+    cert, key = self_signed(cert_dir)
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+    ctx.load_cert_chain(cert, key)
+    srv = ThreadingHTTPServer((host, port), Handler)
+    srv.socket = ctx.wrap_socket(srv.socket, server_side=True)
+    threading.Thread(target=srv.serve_forever, daemon=True).start()
+    return srv
+
+
+def main():
+    global ARGS, CANVAS_LEN, CANVAS_STEP
+    p = argparse.ArgumentParser()
+    p.add_argument("--upstream", default="http://127.0.0.1:8010")
+    p.add_argument("--model", default="dgemma")
+    p.add_argument("--tokenizer", default="/models/dgemma", help="HF id or local path")
+    p.add_argument("--canvas", type=int, default=64, help="the served canvas length")
+    p.add_argument(
+        "--canvas-step",
+        type=int,
+        default=16,
+        help="request widths round up to a multiple of this",
+    )
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=8011)
+    p.add_argument(
+        "--tls-port", type=int, default=0, help="also listen with HTTPS here (0 = off)"
+    )
+    p.add_argument(
+        "--cert-dir",
+        default=os.path.expanduser("~/.cache/djev"),
+        help="directory for the self-signed certificate",
+    )
+    ARGS = p.parse_args()
+    CANVAS_LEN = ARGS.canvas
+    CANVAS_STEP = ARGS.canvas_step
+    init_tokenizer(AutoTokenizer.from_pretrained(ARGS.tokenizer))
+    if ARGS.tls_port:
+        serve_tls(ARGS.host, ARGS.tls_port, ARGS.cert_dir)
+        print(
+            f"structured server https on {ARGS.host}:{ARGS.tls_port} (self-signed)",
+            flush=True,
+        )
+    print(
+        f"structured server on {ARGS.host}:{ARGS.port} -> {ARGS.upstream} "
+        f"(canvas {CANVAS_LEN})",
+        flush=True,
+    )
+    ThreadingHTTPServer((ARGS.host, ARGS.port), Handler).serve_forever()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/features/structured_outputs/structured_outputs_offline.py b/examples/features/structured_outputs/structured_outputs_offline.py
index 6b6099f71b12..e716935755c4 100644
--- a/examples/features/structured_outputs/structured_outputs_offline.py
+++ b/examples/features/structured_outputs/structured_outputs_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file demonstrates the example usage of structured outputs
+"""This file demonstrates the example usage of structured outputs
 in vLLM. It shows how to apply different constraints such as choice,
 regex, json schema, and grammar to produce structured and formatted
 results based on specific prompts.
diff --git a/examples/features/torchrun/torchrun_dp_example_offline.py b/examples/features/torchrun/torchrun_dp_example_offline.py
index f18f6042e9c6..1f8ac603777d 100644
--- a/examples/features/torchrun/torchrun_dp_example_offline.py
+++ b/examples/features/torchrun/torchrun_dp_example_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-experimental support for data-parallel inference with torchrun
+"""experimental support for data-parallel inference with torchrun
 Note the data load balancing and distribution is done out of the vllm engine,
 no internal lb supported in external_launcher mode.
 
diff --git a/examples/features/torchrun/torchrun_example_offline.py b/examples/features/torchrun/torchrun_example_offline.py
index e41bcd420c20..b4b7cbd6a9d2 100644
--- a/examples/features/torchrun/torchrun_example_offline.py
+++ b/examples/features/torchrun/torchrun_example_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-experimental support for tensor-parallel inference with torchrun,
+"""experimental support for tensor-parallel inference with torchrun,
 see https://github.com/vllm-project/vllm/issues/11400 for
 the motivation and use case for this example.
 run the script with `torchrun --nproc-per-node=4 torchrun_example_offline.py`,
diff --git a/examples/generate/multimodal/audio_language_offline.py b/examples/generate/multimodal/audio_language_offline.py
index a2e9b1975602..8b2215d8c1eb 100644
--- a/examples/generate/multimodal/audio_language_offline.py
+++ b/examples/generate/multimodal/audio_language_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM for running offline inference
+"""This example shows how to use vLLM for running offline inference
 with the correct prompt format on audio language models.
 
 For most models, the prompt format should follow corresponding examples
@@ -286,8 +285,7 @@ def run_minicpmo(question: str, audio_count: int) -> ModelRequestData:
 
 # Phi-4-multimodal-instruct
 def run_phi4mm(question: str, audio_count: int) -> ModelRequestData:
-    """
-    Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
+    """Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
     show how to process audio inputs.
     """
     model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct")
diff --git a/examples/generate/multimodal/encoder_decoder_multimodal_offline.py b/examples/generate/multimodal/encoder_decoder_multimodal_offline.py
index 1ab087de3ff2..b4ea893fc03e 100644
--- a/examples/generate/multimodal/encoder_decoder_multimodal_offline.py
+++ b/examples/generate/multimodal/encoder_decoder_multimodal_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM for running offline inference with
+"""This example shows how to use vLLM for running offline inference with
 the explicit/implicit prompt format on enc-dec LMMs for text generation.
 """
 
@@ -56,8 +55,7 @@ def run_whisper():
 
 
 def run_fireredasr2():
-    """
-    FireRedASR2 – Automatic Speech Recognition model.
+    """FireRedASR2 – Automatic Speech Recognition model.
 
     This model uses a Conformer encoder + Qwen2 LLM decoder architecture
     for speech-to-text transcription.  Audio is passed via the implicit
diff --git a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py
index e1c4fd76d7d1..60aba7abd066 100644
--- a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py
+++ b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py
@@ -43,7 +43,6 @@
 
 def encode_base64_content_from_url(content_url: str) -> str:
     """Encode a content retrieved from a remote url to base64 format."""
-
     with requests.get(content_url, headers=headers) as response:
         response.raise_for_status()
         result = base64.b64encode(response.content).decode("utf-8")
@@ -53,7 +52,6 @@ def encode_base64_content_from_url(content_url: str) -> str:
 
 def encode_base64_content_from_file(file_path: str) -> str:
     """Encode a local file content to base64 format."""
-
     with open(file_path, "rb") as file:
         file_content = file.read()
         result = base64.b64encode(file_content).decode("utf-8")
diff --git a/examples/generate/multimodal/qwen2_5_omni/only_thinker.py b/examples/generate/multimodal/qwen2_5_omni/only_thinker.py
index 476ca91d9f5e..66bfd490037f 100644
--- a/examples/generate/multimodal/qwen2_5_omni/only_thinker.py
+++ b/examples/generate/multimodal/qwen2_5_omni/only_thinker.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM for running offline inference
+"""This example shows how to use vLLM for running offline inference
 with the correct prompt format on Qwen2.5-Omni (thinker only).
 """
 
diff --git a/examples/generate/multimodal/qwen3_omni/only_thinker.py b/examples/generate/multimodal/qwen3_omni/only_thinker.py
index 84dd5228e525..57e845ac4f8b 100644
--- a/examples/generate/multimodal/qwen3_omni/only_thinker.py
+++ b/examples/generate/multimodal/qwen3_omni/only_thinker.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM for running offline inference
+"""This example shows how to use vLLM for running offline inference
 with the correct prompt format on Qwen3-Omni (thinker only).
 """
 
diff --git a/examples/generate/multimodal/vision_language_multi_image_offline.py b/examples/generate/multimodal/vision_language_multi_image_offline.py
index 8b12df3991f0..7565f1e42b51 100644
--- a/examples/generate/multimodal/vision_language_multi_image_offline.py
+++ b/examples/generate/multimodal/vision_language_multi_image_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM for running offline inference with
+"""This example shows how to use vLLM for running offline inference with
 multi-image input on vision language models for text generation,
 using the chat template defined by the model.
 """
@@ -56,7 +55,6 @@ def load_aria(question: str, image_urls: list[str]) -> ModelRequestData:
     model_name = "rhymes-ai/Aria"
     engine_args = EngineArgs(
         model=model_name,
-        tokenizer_mode="slow",
         trust_remote_code=True,
         dtype="bfloat16",
         limit_mm_per_prompt={"image": len(image_urls)},
@@ -913,11 +911,9 @@ def load_phi3v(question: str, image_urls: list[str]) -> ModelRequestData:
 
 
 def load_phi4mm(question: str, image_urls: list[str]) -> ModelRequestData:
-    """
-    Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
+    """Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
     show how to process multi images inputs.
     """
-
     model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct")
     # Since the vision-lora and speech-lora co-exist with the base model,
     # we have to manually specify the path of the lora weights.
diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py
index 33e1a8a36b84..95e21090a9a0 100644
--- a/examples/generate/multimodal/vision_language_offline.py
+++ b/examples/generate/multimodal/vision_language_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM for running offline inference with
+"""This example shows how to use vLLM for running offline inference with
 the correct prompt format on vision language models for text generation.
 
 For most models, the prompt format should follow corresponding examples
@@ -1660,8 +1659,7 @@ def run_phi3v(questions: list[str], modality: str) -> ModelRequestData:
 
 # Phi-4-multimodal-instruct
 def run_phi4mm(questions: list[str], modality: str) -> ModelRequestData:
-    """
-    Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
+    """Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
     show how to process image inputs.
     """
     assert modality == "image"
@@ -2254,6 +2252,7 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
     "glm4_5v",
     "glm4_5v_fp8",
     "molmo2",
+    "qwen2_5_vl",
     "qwen3_vl",
     "qwen3_vl_moe",
     "qwen3_5",
@@ -2283,8 +2282,7 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
 
 
 def get_multi_modal_input(args):
-    """
-    return {
+    """Return {
         "data": image or video,
         "question": question,
     }
diff --git a/examples/generate/trace_replay_offline.py b/examples/generate/trace_replay_offline.py
index 2baa03307d19..498a46cbd4ae 100644
--- a/examples/generate/trace_replay_offline.py
+++ b/examples/generate/trace_replay_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Trace-replay with vLLM offline inference.
+"""Trace-replay with vLLM offline inference.
 
 Trace-replay lets you supply a known sequence of decode token IDs alongside
 the prompt.  Instead of sampling from the model distribution, the engine
diff --git a/examples/observability/prometheus_grafana/grafana.json b/examples/observability/prometheus_grafana/grafana.json
index 1c89d4593830..6144f2bfd3ad 100644
--- a/examples/observability/prometheus_grafana/grafana.json
+++ b/examples/observability/prometheus_grafana/grafana.json
@@ -318,7 +318,7 @@
         "type": "prometheus",
         "uid": "${DS_PROMETHEUS}"
       },
-      "description": "Inter token latency in seconds.",
+      "description": "Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains one token, but differs from request-level TPOT under speculative decoding or varying aggregation weights.",
       "fieldConfig": {
         "defaults": {
           "color": {
@@ -479,6 +479,172 @@
       "title": "Inter Token Latency",
       "type": "timeseries"
     },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "${DS_PROMETHEUS}"
+      },
+      "description": "Per-request Time Per Output Token (TPOT) in seconds, computed as (end-to-end latency - TTFT) divided by the number of output tokens minus one. Requests with no more than one output token are recorded as zero; vllm bench serve excludes these requests from its TPOT statistics.",
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "custom": {
+            "axisBorderShow": false,
+            "axisCenteredZero": false,
+            "axisColorMode": "text",
+            "axisLabel": "",
+            "axisPlacement": "auto",
+            "barAlignment": 0,
+            "barWidthFactor": 0.6,
+            "drawStyle": "line",
+            "fillOpacity": 0,
+            "gradientMode": "none",
+            "hideFrom": {
+              "legend": false,
+              "tooltip": false,
+              "viz": false
+            },
+            "insertNulls": false,
+            "lineInterpolation": "linear",
+            "lineWidth": 1,
+            "pointSize": 5,
+            "scaleDistribution": {
+              "type": "linear"
+            },
+            "showPoints": "auto",
+            "spanNulls": false,
+            "stacking": {
+              "group": "A",
+              "mode": "none"
+            },
+            "thresholdsStyle": {
+              "mode": "off"
+            }
+          },
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "red",
+                "value": 80
+              }
+            ]
+          },
+          "unit": "s"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 8,
+        "w": 12,
+        "x": 12,
+        "y": 48
+      },
+      "id": 17,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "single",
+          "sort": "none"
+        }
+      },
+      "targets": [
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))",
+          "fullMetaSearch": false,
+          "includeNullMetadata": false,
+          "instant": false,
+          "legendFormat": "P99",
+          "range": true,
+          "refId": "A",
+          "useBackend": false
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "histogram_quantile(0.95, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))",
+          "fullMetaSearch": false,
+          "hide": false,
+          "includeNullMetadata": false,
+          "instant": false,
+          "legendFormat": "P95",
+          "range": true,
+          "refId": "B",
+          "useBackend": false
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "histogram_quantile(0.9, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))",
+          "fullMetaSearch": false,
+          "hide": false,
+          "includeNullMetadata": false,
+          "instant": false,
+          "legendFormat": "P90",
+          "range": true,
+          "refId": "C",
+          "useBackend": false
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "histogram_quantile(0.5, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))",
+          "fullMetaSearch": false,
+          "hide": false,
+          "includeNullMetadata": false,
+          "instant": false,
+          "legendFormat": "P50",
+          "range": true,
+          "refId": "D",
+          "useBackend": false
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "editorMode": "code",
+          "expr": "rate(vllm:request_time_per_output_token_seconds_sum{model_name=\"$model_name\"}[$__rate_interval])\n/\nrate(vllm:request_time_per_output_token_seconds_count{model_name=\"$model_name\"}[$__rate_interval])",
+          "hide": false,
+          "instant": false,
+          "legendFormat": "Mean",
+          "range": true,
+          "refId": "E"
+        }
+      ],
+      "title": "Time Per Output Token (TPOT)",
+      "type": "timeseries"
+    },
     {
       "datasource": {
         "type": "prometheus",
diff --git a/examples/pooling/classify/classification_with_lora_offline.py b/examples/pooling/classify/classification_with_lora_offline.py
new file mode 100644
index 000000000000..a44066a653ef
--- /dev/null
+++ b/examples/pooling/classify/classification_with_lora_offline.py
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Example of offline classification with a LoRA adapter."""
+
+from vllm import LLM
+from vllm.lora.request import LoRARequest
+from vllm.transformers_utils.repo_utils import hf_api
+
+
+def main():
+    lora_path = hf_api().snapshot_download(
+        repo_id="AmirMohseni/skywork-qwen3-0.6b-reward-lora"
+    )
+    llm = LLM(
+        model="Skywork/Skywork-Reward-V2-Qwen3-0.6B",
+        runner="pooling",
+        enable_lora=True,
+        max_lora_rank=8,
+        max_model_len=512,
+    )
+
+    prompt = "Which response is more helpful and correct?"
+    (output,) = llm.classify(
+        prompt,
+        lora_request=LoRARequest("reward-lora", 1, lora_path),
+    )
+    print(f"Prompt: {prompt!r}")
+    print(f"Class Probabilities: {output.outputs.probs}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/pooling/embed/openai_embedding_long_text/client.py b/examples/pooling/embed/openai_embedding_long_text/client.py
index 7cc33b1f24e2..44e415f670bf 100644
--- a/examples/pooling/embed/openai_embedding_long_text/client.py
+++ b/examples/pooling/embed/openai_embedding_long_text/client.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Example script demonstrating long text embedding with chunked processing in vLLM.
+"""Example script demonstrating long text embedding with chunked processing in vLLM.
 
 This example shows how to use vLLM's chunked processing feature to handle text
 inputs that exceed the model's maximum token length. The feature automatically
diff --git a/examples/pooling/embed/vision_embedding_offline.py b/examples/pooling/embed/vision_embedding_offline.py
index 2c82d0374dfe..5175a6cd46ce 100644
--- a/examples/pooling/embed/vision_embedding_offline.py
+++ b/examples/pooling/embed/vision_embedding_offline.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
-"""
-This example shows how to use vLLM for running offline inference with
+"""This example shows how to use vLLM for running offline inference with
 the correct prompt format on vision language models for multimodal embedding.
 
 For most models, the prompt format should follow corresponding examples
diff --git a/examples/pooling/embed/vision_embedding_online.py b/examples/pooling/embed/vision_embedding_online.py
index 79d6e397c8b6..765528978bda 100644
--- a/examples/pooling/embed/vision_embedding_online.py
+++ b/examples/pooling/embed/vision_embedding_online.py
@@ -37,8 +37,7 @@ def create_chat_embeddings(
     add_special_tokens: bool = False,
     padding: str | NotGiven = NOT_GIVEN,
 ) -> CreateEmbeddingResponse:
-    """
-    Convenience function for accessing vLLM's Chat Embeddings API,
+    """Convenience function for accessing vLLM's Chat Embeddings API,
     which is an extension of OpenAI's existing Embeddings API.
     """
     return client.post(
@@ -56,13 +55,11 @@ def create_chat_embeddings(
 
 
 def run_clip(client: OpenAI, model: str):
-    """
-    Start the server using:
+    """Start the server using:
 
     vllm serve openai/clip-vit-base-patch32 \
         --runner pooling
     """
-
     response = create_chat_embeddings(
         client,
         messages=[
@@ -97,8 +94,7 @@ def run_clip(client: OpenAI, model: str):
 
 
 def run_dse_qwen2_vl(client: OpenAI, model: str):
-    """
-    Start the server using:
+    """Start the server using:
 
     vllm serve MrLight/dse-qwen2-2b-mrl-v1 \
         --runner pooling \
@@ -159,14 +155,12 @@ def run_dse_qwen2_vl(client: OpenAI, model: str):
 
 
 def run_qwen3_vl(client: OpenAI, model: str):
-    """
-    Start the server using:
+    """Start the server using:
 
     vllm serve Qwen/Qwen3-VL-Embedding-2B \
         --runner pooling \
         --max-model-len 8192
     """
-
     default_instruction = "Represent the user's input."
 
     print("Text embedding output:")
@@ -266,14 +260,12 @@ def run_qwen3_vl(client: OpenAI, model: str):
 
 
 def run_siglip(client: OpenAI, model: str):
-    """
-    Start the server using:
+    """Start the server using:
 
     vllm serve google/siglip-base-patch16-224 \
         --runner pooling \
         --chat-template template_basic.jinja
     """
-
     response = create_chat_embeddings(
         client,
         messages=[
@@ -311,8 +303,7 @@ def run_siglip(client: OpenAI, model: str):
 
 
 def run_vlm2vec(client: OpenAI, model: str):
-    """
-    Start the server using:
+    """Start the server using:
 
     vllm serve TIGER-Lab/VLM2Vec-Full \
         --runner pooling \
@@ -320,7 +311,6 @@ def run_vlm2vec(client: OpenAI, model: str):
         --max-model-len 4096 \
         --chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja
     """
-
     response = create_chat_embeddings(
         client,
         messages=[
diff --git a/examples/pooling/plugin/prithvi_geospatial_mae_offline.py b/examples/pooling/plugin/prithvi_geospatial_mae_offline.py
index f7b30d9313ba..6bbde74264d8 100644
--- a/examples/pooling/plugin/prithvi_geospatial_mae_offline.py
+++ b/examples/pooling/plugin/prithvi_geospatial_mae_offline.py
@@ -85,8 +85,7 @@ def generate_datamodule():
 
 
 def process_channel_group(orig_img, channels):
-    """
-    Args:
+    """Args:
         orig_img: torch.Tensor representing original image (reference)
         with shape = (bands, H, W).
         channels: list of indices representing RGB channels.
@@ -94,8 +93,8 @@ def process_channel_group(orig_img, channels):
     Returns:
         torch.Tensor with shape (num_channels, height, width)
         for original image
-    """
 
+    """
     orig_img = orig_img[channels, ...]
     valid_mask = torch.ones_like(orig_img, dtype=torch.bool)
     valid_mask[orig_img == NO_DATA_FLOAT] = False
@@ -121,8 +120,8 @@ def read_geotiff(file_path: str):
     Returns:
         np.ndarray with shape (bands, height, width)
         meta info dict
-    """
 
+    """
     with rasterio.open(file_path) as src:
         img = src.read()
         meta = src.meta
@@ -142,8 +141,8 @@ def save_geotiff(image, output_path: str, meta: dict):
         image: np.ndarray with shape (bands, height, width)
         output_path: path where to save the image
         meta: dict with meta info.
-    """
 
+    """
     with rasterio.open(output_path, "w", **meta) as dest:
         for i in range(image.shape[0]):
             dest.write(image[i, :, :], i + 1)
@@ -172,12 +171,14 @@ def load_example(
               images in *file_paths*.
         std: list containing std values for each band in the
              images in *file_paths*.
+        indices: optional band indices to select from each image. All bands
+             are kept when None.
 
     Returns:
         np.array containing created example
         list of meta info for each image in *file_paths*
-    """
 
+    """
     imgs = []
     metas = []
     temporal_coords = []
diff --git a/examples/pooling/reward/sequence_reward_offline.py b/examples/pooling/reward/sequence_reward_offline.py
index 0727bceee111..4453f4a390b3 100644
--- a/examples/pooling/reward/sequence_reward_offline.py
+++ b/examples/pooling/reward/sequence_reward_offline.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Example offline usage of sequence reward models.
+"""Example offline usage of sequence reward models.
 
 The key distinction between sequence classification and token classification
 lies in their output granularity: sequence classification produces a single
diff --git a/examples/pooling/reward/sequence_reward_online.py b/examples/pooling/reward/sequence_reward_online.py
index 40d8d28e3908..c8362265ddeb 100644
--- a/examples/pooling/reward/sequence_reward_online.py
+++ b/examples/pooling/reward/sequence_reward_online.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example online usage of sequence reward models.
+"""Example online usage of sequence reward models.
 
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.
diff --git a/examples/pooling/reward/token_reward_offline.py b/examples/pooling/reward/token_reward_offline.py
index 4705c0491241..1820c390f8a9 100644
--- a/examples/pooling/reward/token_reward_offline.py
+++ b/examples/pooling/reward/token_reward_offline.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Example offline usage of token reward models.
+"""Example offline usage of token reward models.
 
 The key distinction between sequence classification and token classification
 lies in their output granularity: sequence classification produces a single
diff --git a/examples/pooling/reward/token_reward_online.py b/examples/pooling/reward/token_reward_online.py
index 64ee0c9dfdcc..38f23944515e 100644
--- a/examples/pooling/reward/token_reward_online.py
+++ b/examples/pooling/reward/token_reward_online.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example online usage of token reward models.
+"""Example online usage of token reward models.
 
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.
diff --git a/examples/pooling/score/cohere_rerank_client.py b/examples/pooling/score/cohere_rerank_client.py
index b32209967be9..e7bf720cac6d 100644
--- a/examples/pooling/score/cohere_rerank_client.py
+++ b/examples/pooling/score/cohere_rerank_client.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example of using the OpenAI entrypoint's rerank API which is compatible with
+"""Example of using the OpenAI entrypoint's rerank API which is compatible with
 the Cohere SDK: https://github.com/cohere-ai/cohere-python
 Note that `pip install cohere` is needed to run this example.
 
diff --git a/examples/pooling/score/colbert_rerank_online.py b/examples/pooling/score/colbert_rerank_online.py
index 4cc509b95ac4..b8bc05b9e03c 100644
--- a/examples/pooling/score/colbert_rerank_online.py
+++ b/examples/pooling/score/colbert_rerank_online.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example of using ColBERT late interaction models for reranking and scoring.
+"""Example of using ColBERT late interaction models for reranking and scoring.
 
 ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings
 and MaxSim scoring for document reranking, providing better accuracy than
diff --git a/examples/pooling/score/colmodernvbert_rerank_online.py b/examples/pooling/score/colmodernvbert_rerank_online.py
index de827ae06260..95ca36ed717c 100644
--- a/examples/pooling/score/colmodernvbert_rerank_online.py
+++ b/examples/pooling/score/colmodernvbert_rerank_online.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example of using ColModernVBERT late interaction model for reranking.
+"""Example of using ColModernVBERT late interaction model for reranking.
 
 ColModernVBERT is a multi-modal ColBERT-style model combining a SigLIP
 vision encoder with a ModernBERT text encoder. It produces per-token
diff --git a/examples/pooling/score/colqwen3_5_rerank_online.py b/examples/pooling/score/colqwen3_5_rerank_online.py
index 00746634d5d1..a2cd635361fd 100644
--- a/examples/pooling/score/colqwen3_5_rerank_online.py
+++ b/examples/pooling/score/colqwen3_5_rerank_online.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example of using ColQwen3.5 late interaction model for reranking.
+"""Example of using ColQwen3.5 late interaction model for reranking.
 
 ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
 It produces per-token embeddings and uses MaxSim scoring for retrieval
diff --git a/examples/pooling/score/colqwen3_rerank_online.py b/examples/pooling/score/colqwen3_rerank_online.py
index 0e61531bfd34..22d5fee0795d 100644
--- a/examples/pooling/score/colqwen3_rerank_online.py
+++ b/examples/pooling/score/colqwen3_rerank_online.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
-"""
-Example of using ColQwen3 late interaction model for reranking and scoring.
+"""Example of using ColQwen3 late interaction model for reranking and scoring.
 
 ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL.
 It produces per-token embeddings and uses MaxSim scoring for retrieval
diff --git a/examples/pooling/score/convert_model_to_seq_cls.py b/examples/pooling/score/convert_model_to_seq_cls.py
index 21e294166456..494f683ae709 100644
--- a/examples/pooling/score/convert_model_to_seq_cls.py
+++ b/examples/pooling/score/convert_model_to_seq_cls.py
@@ -2,8 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
 
-"""
-Script to convert Large Language Models (LLMs) to Sequence Classification models.
+"""Script to convert Large Language Models (LLMs) to Sequence Classification models.
 This is particularly useful for converting reranker models that use next-token
 prediction to a sequence classification format for compatibility with standard
 classification and rerank pipelines.
@@ -35,8 +34,7 @@
 
 
 def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device):
-    """
-    This method extracts the difference between weights for 'true' and 'false' tokens
+    """This method extracts the difference between weights for 'true' and 'false' tokens
     from the language model head to create a single classification weight vector.
 
     Args:
@@ -47,6 +45,7 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device):
         device: Target device (cpu/cuda)
 
     Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
+
     """
     assert len(tokens) == 2, (
         "Method requires exactly two tokens for binary classification"
@@ -73,8 +72,7 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device):
 
 
 def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device):
-    """
-    Directly use token weights from the language model head for classification.
+    """Directly use token weights from the language model head for classification.
 
     This method maps each classification label directly to a corresponding token
     in the vocabulary without additional transformation.
@@ -85,6 +83,7 @@ def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device):
         tokenizer: Model tokenizer
         tokens: List of tokens representing class labels
         device: Target device (cpu/cuda)
+
     """
     # Get the language model head weights (vocabulary_size x hidden_size)
     lm_head_weights = causal_lm.lm_head.weight
@@ -110,8 +109,7 @@ def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device):
 def converting(
     model_name, classifier_from_tokens, path, method, use_sep_token=False, device="cpu"
 ):
-    """
-    Main conversion function to transform a CausalLM model to SequenceClassification.
+    """Main conversion function to transform a CausalLM model to SequenceClassification.
 
     Args:
         model_name: Name or path of the pretrained model
@@ -120,6 +118,7 @@ def converting(
         method: Conversion method ('from_2_way_softmax' or 'no_post_processing')
         use_sep_token: Whether to use separating token in the sequence classification model
         device: Device to load the model on ('cpu' or 'cuda')
+
     """
     assert method in method_map, f"Unknown method: {method}"
 
diff --git a/examples/pooling/score/qwen3_reranker_offline.py b/examples/pooling/score/qwen3_reranker_offline.py
index c79ebf97fe64..dbe8d47522c1 100644
--- a/examples/pooling/score/qwen3_reranker_offline.py
+++ b/examples/pooling/score/qwen3_reranker_offline.py
@@ -2,8 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
 
-"""
-What is the difference between the official original version and one
+"""What is the difference between the official original version and one
 that has been converted into a sequence classification model?
 
 Qwen3-Reranker is a language model that doing reranker by using the
@@ -34,8 +33,7 @@
 
 
 def get_llm() -> LLM:
-    """
-    Initializes and returns the LLM model for Qwen3-Reranker.
+    """Initializes and returns the LLM model for Qwen3-Reranker.
 
     Returns:
         LLM: Configured vLLM instance for reranking tasks.
@@ -43,6 +41,7 @@ def get_llm() -> LLM:
     Note:
         This function loads the ORIGINAL Qwen3-Reranker model with specific
         overrides to make it compatible with vLLM's score API.
+
     """
     return LLM(
         # Specify the original model from HuggingFace
diff --git a/examples/pooling/score/qwen3_reranker_online.py b/examples/pooling/score/qwen3_reranker_online.py
index f117b3b5b84d..324e3f5bc349 100644
--- a/examples/pooling/score/qwen3_reranker_online.py
+++ b/examples/pooling/score/qwen3_reranker_online.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
-"""
-What is the difference between the official original version and one
+"""What is the difference between the official original version and one
 that has been converted into a sequence classification model?
 
 Qwen3-Reranker is a language model that doing reranker by using the
diff --git a/examples/pooling/score/rerank_api_online.py b/examples/pooling/score/rerank_api_online.py
index 908d6a9240aa..caf05862825c 100644
--- a/examples/pooling/score/rerank_api_online.py
+++ b/examples/pooling/score/rerank_api_online.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example of using the OpenAI entrypoint's rerank API which is compatible with
+"""Example of using the OpenAI entrypoint's rerank API which is compatible with
 Jina and Cohere https://jina.ai/reranker
 
 run: vllm serve BAAI/bge-reranker-base
diff --git a/examples/pooling/score/score_api_online.py b/examples/pooling/score/score_api_online.py
index af2886f136ea..a60cf04bf8fb 100644
--- a/examples/pooling/score/score_api_online.py
+++ b/examples/pooling/score/score_api_online.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example online usage of Score API.
+"""Example online usage of Score API.
 
 Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
 """
diff --git a/examples/pooling/score/using_template_offline.py b/examples/pooling/score/using_template_offline.py
index f434e699ff0e..abe8157d4c2f 100644
--- a/examples/pooling/score/using_template_offline.py
+++ b/examples/pooling/score/using_template_offline.py
@@ -68,13 +68,13 @@ def get_chat_template(model: str) -> str:
 def get_hf_overrides(model: str) -> dict[str, Any]:
     """Convert Large Language Models (LLMs) to Sequence Classification models.
 
-    note:
+    Note:
         Some reranking models require special configuration overrides to work
         correctly with vLLM's score API.
         Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_offline.py
         Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
-    """
 
+    """
     model_name_to_hf_overrides_map = {
         "BAAI/bge-reranker-v2-gemma": {
             "architectures": ["GemmaForSequenceClassification"],
@@ -122,7 +122,6 @@ def get_hf_overrides(model: str) -> dict[str, Any]:
 
 def main(args: Namespace):
     """Main execution function for the reranking example."""
-
     # Get the overrides for the specified model
     args.hf_overrides = get_hf_overrides(args.model)
 
diff --git a/examples/pooling/score/using_template_online.py b/examples/pooling/score/using_template_online.py
index f0bfa7d15769..f7dc109dd81f 100644
--- a/examples/pooling/score/using_template_online.py
+++ b/examples/pooling/score/using_template_online.py
@@ -1,15 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
-"""
-Example of using the rerank API with template.
+"""Example of using the rerank API with template.
 
 This script demonstrates how to interact with a vLLM server running
 a reranking model via the REST API.
 Before running this script, start the vLLM server with one of the
 supported reranking models using the commands below.
 
-note:
+Note:
     Some reranking models require special configuration overrides to work correctly
     with vLLM's score API.
     Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_online.py
@@ -21,6 +20,7 @@
     vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --chat-template examples/pooling/score/template/mxbai_rerank_v2.jinja
     vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja
     vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja
+
 """
 
 import json
diff --git a/examples/pooling/score/vision_rerank_api_online.py b/examples/pooling/score/vision_rerank_api_online.py
index dce2efd1dc07..cc1f2e97bf2f 100644
--- a/examples/pooling/score/vision_rerank_api_online.py
+++ b/examples/pooling/score/vision_rerank_api_online.py
@@ -2,8 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
 
-"""
-Example Python client for multimodal rerank API which is compatible with
+"""Example Python client for multimodal rerank API which is compatible with
 Jina and Cohere https://jina.ai/reranker
 
 Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
diff --git a/examples/pooling/score/vision_reranker_offline.py b/examples/pooling/score/vision_reranker_offline.py
index ef2b35487771..24c8b9068f45 100644
--- a/examples/pooling/score/vision_reranker_offline.py
+++ b/examples/pooling/score/vision_reranker_offline.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM for running offline inference with
+"""This example shows how to use vLLM for running offline inference with
 vision language reranker models for multimodal scoring tasks.
 
 Vision language rerankers score the relevance between a text query and
diff --git a/examples/pooling/score/vision_score_api_online.py b/examples/pooling/score/vision_score_api_online.py
index 543d4bfa2147..3a3ac62cf7fd 100644
--- a/examples/pooling/score/vision_score_api_online.py
+++ b/examples/pooling/score/vision_score_api_online.py
@@ -2,8 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
 
-"""
-Example online usage of Score API.
+"""Example online usage of Score API.
 
 Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
 e.g.
diff --git a/examples/pooling/token_classify/forced_alignment_offline.py b/examples/pooling/token_classify/forced_alignment_offline.py
index e97a3735767e..13facab4806d 100644
--- a/examples/pooling/token_classify/forced_alignment_offline.py
+++ b/examples/pooling/token_classify/forced_alignment_offline.py
@@ -3,8 +3,7 @@
 # Adapted from Qwen3-ForcedAligner inference:
 # https://github.com/QwenLM/Qwen3-ASR
 
-"""
-Offline forced alignment example using Qwen3-ForcedAligner-0.6B.
+"""Offline forced alignment example using Qwen3-ForcedAligner-0.6B.
 
 Forced alignment takes audio and reference text as input and produces
 word-level timestamps. The model predicts a time bin at each <timestamp>
diff --git a/examples/pooling/token_classify/forced_alignment_online.py b/examples/pooling/token_classify/forced_alignment_online.py
index e967a8ba70d8..ab3026160d5f 100644
--- a/examples/pooling/token_classify/forced_alignment_online.py
+++ b/examples/pooling/token_classify/forced_alignment_online.py
@@ -3,8 +3,7 @@
 # Adapted from Qwen3-ForcedAligner inference:
 # https://github.com/QwenLM/Qwen3-ASR
 
-"""
-Online forced alignment example using Qwen3-ForcedAligner-0.6B-hf.
+"""Online forced alignment example using Qwen3-ForcedAligner-0.6B-hf.
 
 Forced alignment takes audio and reference text as input and produces
 word-level timestamps. The model predicts a time bin at each <timestamp>
diff --git a/examples/pooling/token_classify/ner_online.py b/examples/pooling/token_classify/ner_online.py
index 9ec2bd45a0fe..bd60d6862600 100644
--- a/examples/pooling/token_classify/ner_online.py
+++ b/examples/pooling/token_classify/ner_online.py
@@ -2,8 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Adapted from https://huggingface.co/boltuix/NeuroBERT-NER
 
-"""
-Example online usage of Pooling API for Named Entity Recognition (NER).
+"""Example online usage of Pooling API for Named Entity Recognition (NER).
 
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.
diff --git a/examples/pooling/token_embed/colqwen3_token_embed_online.py b/examples/pooling/token_embed/colqwen3_token_embed_online.py
index cac11188e87e..c1a8f3a56ef8 100644
--- a/examples/pooling/token_embed/colqwen3_token_embed_online.py
+++ b/examples/pooling/token_embed/colqwen3_token_embed_online.py
@@ -2,8 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
 
-"""
-Example online usage of Pooling API for ColQwen3 multi-vector retrieval.
+"""Example online usage of Pooling API for ColQwen3 multi-vector retrieval.
 
 ColQwen3 is a multi-modal late interaction model based on Qwen3-VL that
 produces per-token embeddings (320-dim, L2-normalized) for both text and
diff --git a/examples/pooling/token_embed/jina_reranker_v3_online.py b/examples/pooling/token_embed/jina_reranker_v3_online.py
index 8350aee2f143..b9018f0f23f8 100644
--- a/examples/pooling/token_embed/jina_reranker_v3_online.py
+++ b/examples/pooling/token_embed/jina_reranker_v3_online.py
@@ -2,8 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
 
-"""
-Example online usage of the Jina Reranker v3 score and rerank APIs with a task
+"""Example online usage of the Jina Reranker v3 score and rerank APIs with a task
 instruction.
 
 Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the
diff --git a/examples/pooling/token_embed/multi_vector_retrieval_online.py b/examples/pooling/token_embed/multi_vector_retrieval_online.py
index c2aae88294e8..f36e3dc31d10 100644
--- a/examples/pooling/token_embed/multi_vector_retrieval_online.py
+++ b/examples/pooling/token_embed/multi_vector_retrieval_online.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Example online usage of Pooling API for multi vector retrieval.
+"""Example online usage of Pooling API for multi vector retrieval.
 
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.
diff --git a/examples/ray_serving/batch_llm_inference.py b/examples/ray_serving/batch_llm_inference.py
index 22408dc95033..1b265efb1766 100644
--- a/examples/ray_serving/batch_llm_inference.py
+++ b/examples/ray_serving/batch_llm_inference.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use Ray Data for data parallel batch inference.
+"""This example shows how to use Ray Data for data parallel batch inference.
 
 Ray Data is a data processing framework that can process very large datasets
 with first-class support for vLLM.
diff --git a/examples/ray_serving/ray_serve_deepseek.py b/examples/ray_serving/ray_serve_deepseek.py
index af53443b9101..a1e9348115a4 100644
--- a/examples/ray_serving/ray_serve_deepseek.py
+++ b/examples/ray_serving/ray_serve_deepseek.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Deploy DeepSeek R1 or V3 with Ray Serve LLM.
+"""Deploy DeepSeek R1 or V3 with Ray Serve LLM.
 
 Ray Serve LLM is a scalable and production-grade model serving library built
 on the Ray distributed computing framework and first-class support for the vLLM engine.
diff --git a/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py b/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py
index 1dfc3084646d..53b30f534d47 100644
--- a/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py
+++ b/examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-An example demonstrates how to use tool calling with reasoning models 
+"""An example demonstrates how to use tool calling with reasoning models
 like QwQ-32B. The reasoning will not be parsed by the tool 
 calling process; only the final output will be parsed.
 
diff --git a/examples/reasoning/openai_chat_completion_with_reasoning.py b/examples/reasoning/openai_chat_completion_with_reasoning.py
index 87043897b058..a72606108f8e 100644
--- a/examples/reasoning/openai_chat_completion_with_reasoning.py
+++ b/examples/reasoning/openai_chat_completion_with_reasoning.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-An example shows how to generate chat completions from reasoning models
+"""An example shows how to generate chat completions from reasoning models
 like DeepSeekR1.
 
 To run this example, you need to start the vLLM server
diff --git a/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py b/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py
index 8e262701b720..36f4c1f8e6e2 100644
--- a/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py
+++ b/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-An example shows how to generate chat completions from reasoning models
+"""An example shows how to generate chat completions from reasoning models
 like DeepSeekR1.
 
 To run this example, you need to start the vLLM server with the reasoning
diff --git a/examples/reasoning/openai_responses_client.py b/examples/reasoning/openai_responses_client.py
index b4eb24671507..94ecd84468d5 100644
--- a/examples/reasoning/openai_responses_client.py
+++ b/examples/reasoning/openai_responses_client.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Set up this example by starting a vLLM OpenAI-compatible server.
+"""Set up this example by starting a vLLM OpenAI-compatible server.
 Reasoning models can be used through the Responses API as seen here
 https://platform.openai.com/docs/api-reference/responses
 For example:
diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py
index ea7099962f21..081023156edf 100644
--- a/examples/rl/rlhf_async_new_apis.py
+++ b/examples/rl/rlhf_async_new_apis.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Demonstrates async reinforcement learning using vLLM and Ray,
+"""Demonstrates async reinforcement learning using vLLM and Ray,
 with native weight syncing APIs and batch-invariant generation.
 
 The script separates training and inference workloads onto distinct GPUs
diff --git a/examples/rl/rlhf_http_ipc.py b/examples/rl/rlhf_http_ipc.py
index c688707df984..0ba9350216f6 100644
--- a/examples/rl/rlhf_http_ipc.py
+++ b/examples/rl/rlhf_http_ipc.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-RLHF weight syncing against a `vllm serve` HTTP server, using CUDA IPC for the
+"""RLHF weight syncing against a `vllm serve` HTTP server, using CUDA IPC for the
 data plane.
 
   * OpenAI-compatible API for inference requests
diff --git a/examples/rl/rlhf_http_nccl.py b/examples/rl/rlhf_http_nccl.py
index 97b0db0b0b83..d095b647aafb 100644
--- a/examples/rl/rlhf_http_nccl.py
+++ b/examples/rl/rlhf_http_nccl.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the
+"""RLHF weight syncing against a `vllm serve` HTTP server, using NCCL for the
 data plane.
 
   * OpenAI-compatible API for inference requests
diff --git a/examples/rl/rlhf_ipc_fsdp_ep.py b/examples/rl/rlhf_ipc_fsdp_ep.py
index e3fa1c44f15e..4eb3fbcbbb31 100644
--- a/examples/rl/rlhf_ipc_fsdp_ep.py
+++ b/examples/rl/rlhf_ipc_fsdp_ep.py
@@ -1,7 +1,6 @@
 # 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**
+"""RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC**
 weight transfer and **packed** tensors.
 
 Multi-rank version of `rlhf_http_ipc.py`: the trainer is 4 FSDP2 Ray actors
diff --git a/examples/rl/rlhf_nccl_fsdp_ep.py b/examples/rl/rlhf_nccl_fsdp_ep.py
index e4380bea948d..b9b72c012ec8 100644
--- a/examples/rl/rlhf_nccl_fsdp_ep.py
+++ b/examples/rl/rlhf_nccl_fsdp_ep.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs).
+"""RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs).
 
 8-GPU layout:
   Training  — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors
@@ -74,8 +73,7 @@
 
 @ray.remote(num_gpus=1)
 class FSDPTrainWorker:
-    """
-    One FSDP2 training worker per GPU.  Four of these form the FSDP group.
+    """One FSDP2 training worker per GPU.  Four of these form the FSDP group.
     Rank 0 additionally drives weight transfer to the vLLM server.
     """
 
diff --git a/examples/rl/routed_experts_e2e.py b/examples/rl/routed_experts_e2e.py
index 1666bc3ffe16..04b5871fd1e9 100644
--- a/examples/rl/routed_experts_e2e.py
+++ b/examples/rl/routed_experts_e2e.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-End-to-end example for routed experts capture with hybrid models.
+"""End-to-end example for routed experts capture with hybrid models.
 
 Validates that:
 1. routed_experts is returned in CompletionOutput for MoE models.
diff --git a/examples/scale_out/example_mm_serve.py b/examples/scale_out/example_mm_serve.py
index 157a32d3be8b..558a4f930247 100644
--- a/examples/scale_out/example_mm_serve.py
+++ b/examples/scale_out/example_mm_serve.py
@@ -13,9 +13,9 @@
 
 Launch the server first:
 
-    VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \
-        vllm serve Qwen/Qwen3-VL-2B-Instruct \
-        --dtype bfloat16 --max-model-len 4096 --enforce-eager
+    vllm serve Qwen/Qwen3-VL-2B-Instruct \
+        --dtype bfloat16 --max-model-len 4096 --enforce-eager \
+        --enable-scale-out
 
 Then run this script:
 
diff --git a/examples/scale_out/token_generation_client.py b/examples/scale_out/token_generation_client.py
index 922018f43d16..a4792a68f9a3 100644
--- a/examples/scale_out/token_generation_client.py
+++ b/examples/scale_out/token_generation_client.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Call a token-in/token-out server started with:
 
-VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve Qwen/Qwen3-0.6B
+vllm serve Qwen/Qwen3-0.6B --enable-scale-out
 """
 
 import httpx
diff --git a/examples/speech_to_text/openai/openai_transcription_client.py b/examples/speech_to_text/openai/openai_transcription_client.py
index f928c06d45ea..cc6757bd9635 100644
--- a/examples/speech_to_text/openai/openai_transcription_client.py
+++ b/examples/speech_to_text/openai/openai_transcription_client.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This script demonstrates how to use the vLLM API server to perform audio
+"""This script demonstrates how to use the vLLM API server to perform audio
 transcription with the `openai/whisper-large-v3` model.
 
 Before running this script, you must start the vLLM server with the following command:
@@ -35,8 +34,7 @@ def sync_openai(
     hotwords: str = None,
     prompt: str | None = None,
 ):
-    """
-    Perform synchronous transcription using OpenAI-compatible API.
+    """Perform synchronous transcription using OpenAI-compatible API.
 
     The optional ``prompt`` is the OpenAI-API ``prompt`` field (style /
     vocabulary hint). It is wired through model-by-model: Whisper uses it
@@ -69,9 +67,7 @@ async def stream_openai_response(
     hotwords: str = None,
     prompt: str | None = None,
 ):
-    """
-    Perform asynchronous transcription using OpenAI-compatible API.
-    """
+    """Perform asynchronous transcription using OpenAI-compatible API."""
     print("\ntranscription result [stream]:", end=" ")
     with open(audio_path, "rb") as f:
         transcription = await client.audio.transcriptions.create(
@@ -98,9 +94,7 @@ async def stream_openai_response(
 
 
 def stream_api_response(audio_path: str, model: str, openai_api_base: str):
-    """
-    Perform streaming transcription using raw HTTP requests to the vLLM API server.
-    """
+    """Stream a transcription using raw HTTP requests to the vLLM API server."""
     import json
     import os
 
diff --git a/examples/speech_to_text/realtime/openai_realtime_client.py b/examples/speech_to_text/realtime/openai_realtime_client.py
index fda3d7cb4564..a290b63e83dc 100644
--- a/examples/speech_to_text/realtime/openai_realtime_client.py
+++ b/examples/speech_to_text/realtime/openai_realtime_client.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This script demonstrates how to use the vLLM Realtime WebSocket API to perform
+"""This script demonstrates how to use the vLLM Realtime WebSocket API to perform
 audio transcription by uploading an audio file.
 
 Before running this script, you must start the vLLM server with a realtime-capable
@@ -34,9 +33,7 @@
 
 
 def audio_to_pcm16_base64(audio_path: str) -> str:
-    """
-    Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
-    """
+    """Load an audio file and convert it to base64-encoded PCM16 @ 16kHz."""
     # Load audio and resample to 16kHz mono
     audio, _ = load_audio(audio_path, sr=16000, mono=True)
     # Convert to PCM16
@@ -46,9 +43,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str:
 
 
 async def realtime_transcribe(audio_path: str, host: str, port: int, model: str):
-    """
-    Connect to the Realtime API and transcribe an audio file.
-    """
+    """Connect to the Realtime API and transcribe an audio file."""
     uri = f"ws://{host}:{port}/v1/realtime"
 
     async with websockets.connect(uri) as ws:
diff --git a/examples/speech_to_text/realtime/openai_realtime_microphone_client.py b/examples/speech_to_text/realtime/openai_realtime_microphone_client.py
index a3c07673ffbe..c29486120c03 100644
--- a/examples/speech_to_text/realtime/openai_realtime_microphone_client.py
+++ b/examples/speech_to_text/realtime/openai_realtime_microphone_client.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Minimal Gradio demo for real-time speech transcription using the vLLM Realtime API.
+"""Minimal Gradio demo for real-time speech transcription using the vLLM Realtime API.
 
 Start the vLLM server first:
 
diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools.py b/examples/tool_calling/openai_chat_completion_client_with_tools.py
index 0bd1d05322f8..508c6900b614 100644
--- a/examples/tool_calling/openai_chat_completion_client_with_tools.py
+++ b/examples/tool_calling/openai_chat_completion_client_with_tools.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Set up this example by starting a vLLM OpenAI-compatible server with tool call
+"""Set up this example by starting a vLLM OpenAI-compatible server with tool call
 options enabled. For example:
 
 IMPORTANT: for mistral, you must use one of the provided mistral tool call
diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_required.py b/examples/tool_calling/openai_chat_completion_client_with_tools_required.py
index c00d712b351d..f38011e10d9c 100644
--- a/examples/tool_calling/openai_chat_completion_client_with_tools_required.py
+++ b/examples/tool_calling/openai_chat_completion_client_with_tools_required.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-To run this example, you can start the vLLM server
+"""To run this example, you can start the vLLM server
 without any specific flags:
 
 ```bash
diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py
index f0b0a2db44ed..004b5fc7fcce 100644
--- a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py
+++ b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
-"""
-Set up this example by starting a vLLM OpenAI-compatible server with tool call
+"""Set up this example by starting a vLLM OpenAI-compatible server with tool call
 options enabled for xLAM-2 models:
 
 vllm serve --model Salesforce/Llama-xLAM-2-8b-fc-r --enable-auto-tool-choice --tool-call-parser xlam
@@ -105,8 +104,7 @@ def translate_text(text: str, target_language: str):
 
 
 def process_response(response, tool_functions, original_query):
-    """Process a non-streaming response with possible tool calls"""
-
+    """Process a non-streaming response with possible tool calls."""
     print("\n--- Response Output ---")
 
     # Check if the response has content
@@ -195,7 +193,7 @@ def process_response(response, tool_functions, original_query):
 
 
 def run_test_case(query, test_name):
-    """Run a single test case with the given query"""
+    """Run a single test case with the given query."""
     print(f"\n{'=' * 50}\nTEST CASE: {test_name}\n{'=' * 50}")
     print(f"Query: '{query}'")
 
diff --git a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py
index 94e664c9ec3d..43e4ed21a0d6 100644
--- a/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py
+++ b/examples/tool_calling/openai_chat_completion_client_with_tools_xlam_streaming.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: E501
-"""
-Set up this example by starting a vLLM OpenAI-compatible server with tool call
+"""Set up this example by starting a vLLM OpenAI-compatible server with tool call
 options enabled for xLAM-2 models:
 
 vllm serve --model Salesforce/Llama-xLAM-2-8b-fc-r --enable-auto-tool-choice --tool-call-parser xlam
@@ -107,7 +106,7 @@ def translate_text(text: str, target_language: str):
 
 
 def process_stream(response, tool_functions, original_query):
-    """Process a streaming response with possible tool calls"""
+    """Process a streaming response with possible tool calls."""
     # Track multiple tool calls
     tool_calls = {}  # Dictionary to store tool calls by ID
 
@@ -223,7 +222,7 @@ def process_stream(response, tool_functions, original_query):
 
 
 def run_test_case(query, test_name):
-    """Run a single test case with the given query"""
+    """Run a single test case with the given query."""
     print(f"\n{'=' * 50}\nTEST CASE: {test_name}\n{'=' * 50}")
     print(f"Query: '{query}'")
 
diff --git a/examples/tool_calling/openai_responses_client_with_mcp_tools.py b/examples/tool_calling/openai_responses_client_with_mcp_tools.py
index 1e685a4f3440..0b6ea412bf56 100644
--- a/examples/tool_calling/openai_responses_client_with_mcp_tools.py
+++ b/examples/tool_calling/openai_responses_client_with_mcp_tools.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example demonstrating MCP (Model Context Protocol) tools with the Responses API.
+"""Example demonstrating MCP (Model Context Protocol) tools with the Responses API.
 
 This example shows how to use MCP tools with different allowed_tools configurations:
 1. No filter (allows all tools from the MCP server)
diff --git a/examples/tool_calling/openai_responses_client_with_tools.py b/examples/tool_calling/openai_responses_client_with_tools.py
index a47210f60020..6530f331e45d 100644
--- a/examples/tool_calling/openai_responses_client_with_tools.py
+++ b/examples/tool_calling/openai_responses_client_with_tools.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Set up this example by starting a vLLM OpenAI-compatible server with tool call
+"""Set up this example by starting a vLLM OpenAI-compatible server with tool call
 options enabled.
 Reasoning models can be used through the Responses API as seen here
 https://platform.openai.com/docs/api-reference/responses
@@ -17,8 +16,7 @@
 
 
 def get_weather(latitude: float, longitude: float) -> str:
-    """
-    Mock function to simulate getting weather data.
+    """Mock function to simulate getting weather data.
     In a real application, this would call an external weather API.
     """
     return f"Current temperature at ({latitude}, {longitude}) is 20°C."
diff --git a/pyproject.toml b/pyproject.toml
index 1a77433f9cb8..b5c03b22b7d1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -60,9 +60,13 @@ include = ["vllm*"]
 "vllm/third_party/**" = ["ALL"]
 "vllm/version.py" = ["F401"]
 "vllm/_version.py" = ["ALL"]
+# test args are mostly pytest fixtures, which don't benefit from descriptions
+"tests/**" = ["D417"]
 
 [tool.ruff.lint]
 select = [
+    # pydocstyle
+    "D",
     # pycodestyle
     "E",
     # Pyflakes
@@ -83,6 +87,12 @@ select = [
     "INP",
 ]
 ignore = [
+    # don't enforce that everything has a docstring
+    "D100", "D101", "D102", "D103", "D104", "D105", "D106", "D107",
+    # unwanted pydocstyle rules
+    "D205", "D209", "D301", "D400", "D401", "D404", "D415",
+    # incompatible pydocstyle pairs; pin the half `ruff` already picks
+    "D203", "D213",
     # star imports
     "F405", "F403",
     # lambda expression assignment
@@ -95,6 +105,9 @@ ignore = [
     "UP032",
 ]
 
+[tool.ruff.lint.pydocstyle]
+ignore-var-parameters = true
+
 [tool.ruff.format]
 docstring-code-format = true
 
diff --git a/requirements/common.txt b/requirements/common.txt
index 132979d8e678..b55f680b769a 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -8,10 +8,11 @@ tqdm
 blake3
 py-cpuinfo
 transformers >= 5.10.4
-huggingface_hub >= 1.28.0
+huggingface_hub >= 1.31.0  # Required for the `huggingface_hub.utils.httpx` re-export
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
 safetensors >= 0.6.2  # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
-protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
+protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*; platform_machine != "s390x"
+protobuf == 7.36.1; platform_machine == "s390x"
 fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides.
 starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1
 aiohttp >= 3.13.3
@@ -25,7 +26,7 @@ lm-format-enforcer == 0.11.3
 llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" or platform_machine == "s390x"
 outlines_core == 0.2.14
 lark == 1.2.2
-xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
+xgrammar == 0.2.7; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
 typing_extensions >= 4.10
 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
 partial-json-parser # used for parsing partial JSON outputs
diff --git a/requirements/cpu.txt b/requirements/cpu.txt
index 30e47a8cab1a..06bddd4a678e 100644
--- a/requirements/cpu.txt
+++ b/requirements/cpu.txt
@@ -21,5 +21,10 @@ torchcodec >= 0.14; platform_machine != "s390x" and platform_machine != "riscv64
 # Intel Extension for PyTorch, only for x86_64 CPUs
 intel-openmp==2024.2.1; platform_machine == "x86_64"
 
+# triton-cpu, only for x86_64 CPUs
+triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version == "3.10"
+triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version == "3.11"
+triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version >= "3.12"
+
 # Use this to gather CPU info and optimize based on ARM Neoverse cores
 py-cpuinfo; platform_machine == "aarch64"
diff --git a/requirements/cuda.txt b/requirements/cuda.txt
index 427f86ddbeda..2bbf1e244d94 100644
--- a/requirements/cuda.txt
+++ b/requirements/cuda.txt
@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
 # flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
 # install_requires so the published wheel does not carry an unresolvable pin
 --extra-index-url https://flashinfer.ai/whl/
-flashinfer-python==0.6.18.post1
-flashinfer-cubin==0.6.18.post1
+flashinfer-python==0.7.0
+flashinfer-cubin==0.7.0
 apache-tvm-ffi==0.1.11
 tilelang==0.1.12
 nvidia-cudnn-frontend>=1.19.1
@@ -26,11 +26,11 @@ fastsafetensors >= 0.3.3
 instanttensor >= 0.1.9
 
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
-nvidia-cutlass-dsl[cu13]==4.6.2
-quack-kernels==0.6.4  # Required for CUTLASS DSL 4.6 by MSA
+nvidia-cutlass-dsl[cu13]==4.7.1
+quack-kernels==0.6.5  # Required by MSA
 
 # Tokenspeed_MLA for faster mla with spec decode
 tokenspeed-mla==0.1.8; platform_system == "Linux"
 
 # Humming kernels for quantization gemm
-humming-kernels[cu13]==0.1.12
+humming-kernels[cu13]==0.1.16
diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt
index 1760ed651734..9d9d1278818e 100644
--- a/requirements/kv_connectors.txt
+++ b/requirements/kv_connectors.txt
@@ -2,7 +2,7 @@ lmcache >= 0.3.9
 # CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime
 # images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded.
 cupy-cuda13x != 14.1.0
-nixl == 1.3.2
+nixl == 1.4.1
 # CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for
 # the mooncake-transfer-engine-cuda13 variant of the same version.
 mooncake-transfer-engine >= 0.3.12
diff --git a/requirements/kv_connectors_cu12.txt b/requirements/kv_connectors_cu12.txt
new file mode 100644
index 000000000000..5ae5fba5ff99
--- /dev/null
+++ b/requirements/kv_connectors_cu12.txt
@@ -0,0 +1,5 @@
+# LMCache 0.5.x uses CUDA 13 dependencies; retain the CUDA 12 release here.
+lmcache == 0.4.4
+cupy-cuda12x != 14.1.0
+nixl == 1.4.1
+mooncake-transfer-engine >= 0.3.12
diff --git a/requirements/rubin-prerelease.txt b/requirements/rubin-prerelease.txt
index 18a644f6b1b4..d02fb38884b0 100644
--- a/requirements/rubin-prerelease.txt
+++ b/requirements/rubin-prerelease.txt
@@ -2,15 +2,25 @@
 # Installed when INSTALL_RUBIN_PRERELEASE=true (see docker/Dockerfile).
 # TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels
 # (vllm-project/vllm#49735).
---pre
 --extra-index-url https://download.pytorch.org/whl/nightly/cu134
+--extra-index-url https://flashinfer.ai/whl/cu134
+--extra-index-url https://flashinfer.ai/whl
 
-torch==2.15.0.dev20260903+cu134
-torchvision==0.30.0.dev20260903+cu134
-torchaudio==2.11.0.dev20260903+cu134
+# Match each architecture's torch pin to the latest torchvision wheel.
+torch==2.15.0.dev20260909+cu134; platform_machine == "x86_64"
+torch==2.15.0.dev20260908+cu134; platform_machine == "aarch64"
+torchvision==0.30.0.dev20260909+cu134
+torchaudio==2.11.0.dev20260911+cu134
 
-# docker/Dockerfile strips these two from requirements/cuda.txt whenever
+# docker/Dockerfile strips these packages from requirements/cuda.txt whenever
 # INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to
 # still get installed.
 nvidia-cutlass-dsl[cu13]==4.8.0.dev0
-flashinfer-python==0.6.18
+flashinfer-python==0.7.0
+quack-kernels==0.6.5
+
+cuda-python==13.4.1
+cuda-bindings==13.4.1
+flashinfer-jit-cache==0.7.0+cu134
+flashinfer-cubin==0.7.0
+nixl==1.4.1
diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt
index 81eb1b15cfc4..c1d49139538e 100644
--- a/requirements/test/cpu.txt
+++ b/requirements/test/cpu.txt
@@ -43,7 +43,7 @@ anyio==4.14.1
     #   sse-starlette
     #   starlette
     #   watchfiles
-apache-tvm-ffi==0.1.9
+apache-tvm-ffi==0.1.14.post0
     # via xgrammar
 arctic-inference==0.1.1
     # via -r requirements/test/cuda.in
@@ -127,7 +127,7 @@ cloudpickle==3.1.2
     # via -r requirements/test/../common.txt
 cohere==7.0.8
     # via -r requirements/test/cuda.in
-cohere-melody==0.11.1
+cohere-melody==0.14.0
     # via -r requirements/test/cuda.in
 colorama==0.4.6
     # via
@@ -342,7 +342,7 @@ httpx==0.27.2
     #   schemathesis
 httpx2==2.12.0
     # via mcp
-huggingface-hub==1.28.0
+huggingface-hub==1.31.0
     # via
     #   -r requirements/test/../common.txt
     #   accelerate
@@ -1174,7 +1174,7 @@ tqdm==4.67.3
     #   segmentation-models-pytorch
     #   sentence-transformers
     #   transformers
-transformers==5.16.1
+transformers==5.17.0
     # via
     #   -r requirements/test/../common.txt
     #   -r requirements/test/cuda.in
@@ -1298,7 +1298,7 @@ word2number==1.1
     # via lm-eval
 wrapt==1.17.2
     # via smart-open
-xgrammar==0.2.3
+xgrammar==0.2.7
     # via -r requirements/test/../common.txt
 xxhash==3.5.0
     # via
diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in
index b16c64ed346f..5e3faaebcc92 100644
--- a/requirements/test/cuda.in
+++ b/requirements/test/cuda.in
@@ -39,7 +39,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.16.1
+transformers==5.17.0
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 buildkite-test-collector==0.1.9
diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt
index aee0e7003d1b..3207d395aa24 100644
--- a/requirements/test/cuda.txt
+++ b/requirements/test/cuda.txt
@@ -132,7 +132,7 @@ cloudpickle==3.1.2
     # via -r requirements/test/../common.txt
 cohere==7.0.8
     # via -r requirements/test/cuda.in
-cohere-melody==0.11.1
+cohere-melody==0.14.0
     # via -r requirements/test/cuda.in
 colorama==0.4.6
     # via
@@ -361,7 +361,7 @@ httpx==0.27.2
     #   schemathesis
 httpx2==2.12.0
     # via mcp
-huggingface-hub==1.28.0
+huggingface-hub==1.31.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1281,7 +1281,7 @@ tqdm==4.67.3
     #   segmentation-models-pytorch
     #   sentence-transformers
     #   transformers
-transformers==5.16.1
+transformers==5.17.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1409,7 +1409,7 @@ word2number==1.1
     # via lm-eval
 wrapt==1.17.2
     # via smart-open
-xgrammar==0.2.3
+xgrammar==0.2.7
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt
index 35c4156c5d00..cebf3b782b26 100644
--- a/requirements/test/nightly-torch.txt
+++ b/requirements/test/nightly-torch.txt
@@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.16.1
+transformers==5.17.0
 schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 buildkite-test-collector==0.1.9
diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in
index 8397c9817110..c8c08fabe5f3 100644
--- a/requirements/test/rocm.in
+++ b/requirements/test/rocm.in
@@ -35,7 +35,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==5.16.1
+transformers==5.17.0
 schemathesis>=4.0.0 # Required for openai schema test
 # quantization
 bitsandbytes==0.50.0
diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt
index aa0835cd1b87..c961b6294aa7 100644
--- a/requirements/test/rocm.txt
+++ b/requirements/test/rocm.txt
@@ -134,7 +134,7 @@ cloudpickle==3.1.2
     #   tilelang
 cohere==7.0.8
     # via -r requirements/test/rocm.in
-cohere-melody==0.11.1
+cohere-melody==0.14.0
     # via -r requirements/test/rocm.in
 colorama==0.4.6
     # via
@@ -354,7 +354,7 @@ httpx==0.27.2
     #   schemathesis
 httpx2==2.12.0
     # via mcp
-huggingface-hub==1.28.0
+huggingface-hub==1.31.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1233,7 +1233,7 @@ tqdm==4.67.3
     #   sentence-transformers
     #   tilelang
     #   transformers
-transformers==5.16.1
+transformers==5.17.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1360,7 +1360,7 @@ word2number==1.1
     # via lm-eval
 wrapt==2.1.2
     # via smart-open
-xgrammar==0.2.1
+xgrammar==0.2.7
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in
index 7120b3104302..d2ad5a0cfd53 100644
--- a/requirements/test/xpu.in
+++ b/requirements/test/xpu.in
@@ -1,7 +1,7 @@
 -r ../common.txt
 
 # --- Test Infrastructure ---
-tblib
+tblib==3.1.0
 pytest
 pytest_asyncio
 pytest-timeout
@@ -17,9 +17,12 @@ accelerate
 arctic-inference
 lm_eval[api]>=0.4.12
 modelscope<1.38
-transformers==5.16.1
+runai-model-streamer[s3,gcs,azure]>=0.15.7
+tensorizer==2.10.1
+transformers==5.17.0
 
 # --- Audio Processing ---
+av
 librosa
 audioread
 soxr
@@ -29,16 +32,31 @@ soundfile
 # --- Tool Parsing & Evaluation ---
 blobfile
 rapidfuzz
-gpt-oss
-schemathesis
+gpt-oss >= 0.0.7; python_version > '3.11'
+schemathesis>=4.0.0 # Required for openai schema test.
 jiwer
 bm25s
 pystemmer
 mteb[bm25s]>=2, <3 # required for mteb test
+sentence-transformers==5.2.0 # required for embedding tests
 num2words
 pqdm
 
 # --- Vision & Multimodal ---
-timm
+timm>=1.0.17 # required for internvl and gemma3n-mm test
 albumentations
-mistral-common[image,audio]
\ No newline at end of file
+ftfy
+matplotlib
+decord==0.6.0; platform_machine == "x86_64"
+mistral-common[image,audio]>=1.11.6 # required for voxtral test
+
+open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
+
+# --- Observability ---
+opentelemetry-api>=1.26.0
+opentelemetry-exporter-otlp>=1.26.0
+opentelemetry-sdk>=1.26.0
+opentelemetry-semantic-conventions-ai>=0.4.1
+
+cohere_melody>=0.11.1 # required for cohere command reasoning/tool parser tests
+cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types)
diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt
index f3e8bb1bc678..ed8827f5db0c 100644
--- a/requirements/test/xpu.txt
+++ b/requirements/test/xpu.txt
@@ -54,6 +54,16 @@ attrs==26.1.0
     #   referencing
 audioread==3.1.0
     # via -r requirements/test/xpu.in
+av==18.1.0
+    # via -r requirements/test/xpu.in
+azure-core==1.41.0
+    # via
+    #   azure-identity
+    #   azure-storage-blob
+azure-identity==1.25.3
+    # via runai-model-streamer-azure
+azure-storage-blob==12.30.1
+    # via runai-model-streamer-azure
 blake3==1.0.9
     # via -r requirements/test/../common.txt
 blobfile==3.2.0
@@ -62,6 +72,14 @@ bm25s==0.3.10
     # via
     #   -r requirements/test/xpu.in
     #   mteb
+boto3==1.43.87
+    # via
+    #   runai-model-streamer-s3
+    #   tensorizer
+botocore==1.43.87
+    # via
+    #   boto3
+    #   s3transfer
 bounded-pool-executor==0.0.3
     # via pqdm
 cachetools==7.1.7
@@ -94,16 +112,29 @@ click==8.4.2
     #   uvicorn
 cloudpickle==3.1.2
     # via -r requirements/test/../common.txt
+cohere==7.1.1
+    # via -r requirements/test/xpu.in
+cohere-melody==0.13.4
+    # via -r requirements/test/xpu.in
 colorama==0.4.6
     # via sacrebleu
 compressed-tensors==0.17.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
+contourpy==1.3.3
+    # via matplotlib
 coverage==7.15.4
     # via pytest-cov
 cryptography==50.0.0
-    # via pyjwt
+    # via
+    #   azure-identity
+    #   azure-storage-blob
+    #   google-auth
+    #   msal
+    #   pyjwt
+cycler==0.12.1
+    # via matplotlib
 dataproperty==1.1.1
     # via
     #   pytablewriter
@@ -115,6 +146,8 @@ datasets==5.0.1
     #   mteb
 decorator==5.3.1
     # via librosa
+decord==0.6.0
+    # via -r requirements/test/xpu.in
 defusedxml==0.7.1
     # via nltk
 depyf==0.20.0
@@ -142,7 +175,7 @@ docopt==0.6.2
     # via num2words
 docstring-parser==0.18.0
     # via anthropic
-dpcpp-cpp-rt==2026.0.0
+dpcpp-cpp-rt==2026.1.0
     # via
     #   onemkl-sycl-blas
     #   onemkl-sycl-dft
@@ -172,6 +205,8 @@ fastar==0.11.0
     # via
     #   fastapi
     #   fastapi-cloud-cli
+fastavro==1.12.2
+    # via cohere
 filelock==3.32.3
     # via
     #   -c requirements/common.txt
@@ -181,6 +216,8 @@ filelock==3.32.3
     #   huggingface-hub
     #   modelscope
     #   torch
+fonttools==4.64.0
+    # via matplotlib
 frozenlist==1.8.0
     # via
     #   aiohttp
@@ -191,11 +228,36 @@ fsspec==2026.6.0
     #   evaluate
     #   huggingface-hub
     #   torch
+ftfy==6.3.1
+    # via
+    #   -r requirements/test/xpu.in
+    #   open-clip-torch
+google-api-core==2.34.0
+    # via
+    #   google-cloud-core
+    #   google-cloud-storage
+google-auth==2.57.0
+    # via
+    #   google-api-core
+    #   google-cloud-core
+    #   google-cloud-storage
+    #   runai-model-streamer-gcs
+google-cloud-core==2.7.0
+    # via google-cloud-storage
+google-cloud-storage==3.13.1
+    # via runai-model-streamer-gcs
+google-crc32c==1.8.0
+    # via
+    #   google-cloud-storage
+    #   google-resumable-media
+google-resumable-media==2.10.2
+    # via google-cloud-storage
 googleapis-common-protos==1.75.1
     # via
+    #   google-api-core
     #   opentelemetry-exporter-otlp-proto-grpc
     #   opentelemetry-exporter-otlp-proto-http
-gpt-oss==0.0.9
+gpt-oss==0.0.8
     # via -r requirements/test/xpu.in
 graphql-core==3.2.11
     # via hypothesis-graphql
@@ -210,6 +272,8 @@ harfile==0.5.0
     # via schemathesis
 hf-xet==1.6.0
     # via huggingface-hub
+hiredis==3.4.1
+    # via tensorizer
 html2text==2025.4.15
     # via gpt-oss
 httpcore==1.0.9
@@ -221,6 +285,7 @@ httptools==0.8.0
 httpx==0.28.1
     # via
     #   anthropic
+    #   cohere
     #   datasets
     #   fastapi
     #   fastapi-cloud-cli
@@ -230,17 +295,20 @@ httpx2==2.10.0
     # via
     #   mcp
     #   openai
-huggingface-hub==1.28.0
+huggingface-hub==1.31.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   accelerate
     #   datasets
     #   evaluate
+    #   open-clip-torch
     #   sentence-transformers
     #   timm
     #   tokenizers
     #   transformers
+humanize==4.16.0
+    # via runai-model-streamer
 hypothesis==6.165.10
     # via
     #   hypothesis-graphql
@@ -260,27 +328,27 @@ idna==3.18
     #   yarl
 ijson==3.5.1
     # via -r requirements/test/../common.txt
-impi-rt==2021.18.0
+impi-rt==2021.18.1
     # via
     #   oneccl
     #   torch
 iniconfig==2.3.0
     # via pytest
-intel-cmplr-lib-rt==2026.0.0
+intel-cmplr-lib-rt==2026.1.0
     # via
     #   intel-sycl-rt
     #   torch
-intel-cmplr-lib-ur==2026.0.0
+intel-cmplr-lib-ur==2026.1.0
     # via
     #   intel-openmp
     #   intel-sycl-rt
     #   torch
-intel-cmplr-lic-rt==2026.0.0
+intel-cmplr-lic-rt==2026.1.0
     # via
     #   intel-opencl-rt
     #   intel-sycl-rt
     #   torch
-intel-opencl-rt==2026.0.0
+intel-opencl-rt==2026.1.0
     # via
     #   dpcpp-cpp-rt
     #   onemkl-sycl-blas
@@ -289,20 +357,22 @@ intel-opencl-rt==2026.0.0
     #   onemkl-sycl-rng
     #   onemkl-sycl-sparse
     #   torch
-intel-openmp==2026.0.0
+intel-openmp==2026.1.0
     # via
     #   dpcpp-cpp-rt
     #   mkl
     #   torch
-intel-pti==0.17.0
+intel-pti==1.0.1
     # via torch
-intel-sycl-rt==2026.0.0
+intel-sycl-rt==2026.1.0
     # via
     #   dpcpp-cpp-rt
     #   oneccl
     #   torch
 interegular==0.3.3
     # via lm-format-enforcer
+isodate==0.7.2
+    # via azure-storage-blob
 jinja2==3.1.6
     # via
     #   -c requirements/xpu.txt
@@ -316,7 +386,10 @@ jiter==0.16.0
 jiwer==4.0.0
     # via -r requirements/test/xpu.in
 jmespath==1.1.0
-    # via model-hosting-container-standards
+    # via
+    #   boto3
+    #   botocore
+    #   model-hosting-container-standards
 joblib==1.5.3
     # via
     #   librosa
@@ -333,16 +406,16 @@ jsonschema-rs==0.49.9
     # via schemathesis
 jsonschema-specifications==2025.9.1
     # via jsonschema
-jupyter-client==8.9.1
-    # via gpt-oss
-jupyter-core==5.9.1
-    # via jupyter-client
+kiwisolver==1.5.1
+    # via matplotlib
 lark==1.2.2
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
 lazy-loader==0.5
     # via librosa
+libnacl==2.1.0
+    # via tensorizer
 librosa==1.0.0
     # via -r requirements/test/xpu.in
 llguidance==1.7.6
@@ -370,6 +443,8 @@ markupsafe==3.0.3
     # via
     #   jinja2
     #   werkzeug
+matplotlib==3.11.1
+    # via -r requirements/test/xpu.in
 mbstrdecoder==1.1.5
     # via
     #   dataproperty
@@ -388,7 +463,7 @@ mistral-common==1.11.7
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   -r requirements/test/xpu.in
-mkl==2026.0.0
+mkl==2026.1.0
     # via
     #   onemkl-sycl-blas
     #   onemkl-sycl-dft
@@ -406,6 +481,12 @@ more-itertools==11.1.0
     # via lm-eval
 mpmath==1.3.0
     # via sympy
+msal==1.38.0
+    # via
+    #   azure-identity
+    #   msal-extensions
+msal-extensions==1.3.1
+    # via azure-identity
 msgpack==1.2.1
     # via librosa
 msgspec==0.21.1
@@ -434,17 +515,20 @@ numba==0.65.0
     # via
     #   -c requirements/xpu.txt
     #   librosa
-numpy==2.3.5
+numpy==2.2.6
     # via
     #   -r requirements/test/../common.txt
     #   accelerate
     #   albucore
     #   albumentations
     #   bm25s
+    #   contourpy
     #   datasets
+    #   decord
     #   evaluate
     #   librosa
     #   lm-eval
+    #   matplotlib
     #   mistral-common
     #   mteb
     #   numba
@@ -452,38 +536,41 @@ numpy==2.3.5
     #   pandas
     #   pytrec-eval-terrier
     #   rouge-score
+    #   runai-model-streamer
     #   sacrebleu
     #   scikit-learn
     #   scipy
-    #   sentence-transformers
     #   soundfile
     #   soxr
+    #   tensorizer
     #   torchvision
     #   transformers
     #   xgrammar
-oneccl==2022.0.0
+oneccl==2022.1.1
     # via
     #   oneccl-devel
     #   torch
-oneccl-devel==2022.0.0
+oneccl-devel==2022.1.1
     # via torch
-onemkl-license==2026.0.0
+onemkl-license==2026.1.0
     # via
     #   mkl
     #   torch
-onemkl-sycl-blas==2026.0.0
+onemkl-sycl-blas==2026.1.0
     # via
     #   onemkl-sycl-lapack
     #   onemkl-sycl-sparse
     #   torch
-onemkl-sycl-dft==2026.0.0
+onemkl-sycl-dft==2026.1.0
     # via torch
-onemkl-sycl-lapack==2026.0.0
+onemkl-sycl-lapack==2026.1.0
     # via torch
-onemkl-sycl-rng==2026.0.0
+onemkl-sycl-rng==2026.1.0
     # via torch
-onemkl-sycl-sparse==2026.0.0
+onemkl-sycl-sparse==2026.1.0
     # via torch
+open-clip-torch==2.32.0
+    # via -r requirements/test/xpu.in
 openai==3.2.0
     # via
     #   -c requirements/common.txt
@@ -504,6 +591,7 @@ opentelemetry-api==1.44.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
+    #   -r requirements/test/xpu.in
     #   mcp
     #   opentelemetry-exporter-otlp-proto-grpc
     #   opentelemetry-exporter-otlp-proto-http
@@ -513,6 +601,7 @@ opentelemetry-exporter-otlp==1.44.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
+    #   -r requirements/test/xpu.in
 opentelemetry-exporter-otlp-proto-common==1.44.0
     # via
     #   opentelemetry-exporter-otlp-proto-grpc
@@ -530,6 +619,7 @@ opentelemetry-sdk==1.44.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
+    #   -r requirements/test/xpu.in
     #   opentelemetry-exporter-otlp-proto-grpc
     #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-semantic-conventions-ai
@@ -541,6 +631,7 @@ opentelemetry-semantic-conventions-ai==0.5.1
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
+    #   -r requirements/test/xpu.in
 outlines-core==0.2.14
     # via
     #   -c requirements/common.txt
@@ -554,6 +645,7 @@ packaging==26.3
     #   huggingface-hub
     #   lazy-loader
     #   lm-format-enforcer
+    #   matplotlib
     #   modelscope
     #   pooch
     #   pytest
@@ -571,12 +663,11 @@ pathvalidate==3.3.1
 pillow==12.3.0
     # via
     #   -r requirements/test/../common.txt
+    #   matplotlib
     #   mistral-common
     #   torchvision
 platformdirs==4.11.3
-    # via
-    #   jupyter-core
-    #   pooch
+    # via pooch
 pluggy==1.6.0
     # via
     #   pytest
@@ -606,20 +697,30 @@ propcache==0.5.2
     # via
     #   aiohttp
     #   yarl
+proto-plus==1.28.4
+    # via google-api-core
 protobuf==7.35.1
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
+    #   google-api-core
     #   googleapis-common-protos
     #   opentelemetry-proto
+    #   proto-plus
+    #   tensorizer
 psutil==7.2.2
     # via
     #   -r requirements/test/../common.txt
     #   accelerate
+    #   tensorizer
 py-cpuinfo==9.0.0
     # via -r requirements/test/../common.txt
 pyarrow==25.0.1
     # via datasets
+pyasn1==0.6.4
+    # via pyasn1-modules
+pyasn1-modules==0.4.2
+    # via google-auth
 pybase64==1.5.0
     # via -r requirements/test/../common.txt
 pycountry==26.2.16
@@ -634,6 +735,7 @@ pydantic==2.13.4
     #   -r requirements/test/../common.txt
     #   albumentations
     #   anthropic
+    #   cohere
     #   compressed-tensors
     #   fastapi
     #   fastapi-cloud-cli
@@ -650,7 +752,9 @@ pydantic==2.13.4
     #   pydantic-settings
     #   xgrammar
 pydantic-core==2.46.4
-    # via pydantic
+    # via
+    #   cohere
+    #   pydantic
 pydantic-extra-types==2.11.1
     # via
     #   fastapi
@@ -664,7 +768,11 @@ pygments==2.21.0
     #   pytest
     #   rich
 pyjwt==2.13.0
-    # via mcp
+    # via
+    #   mcp
+    #   msal
+pyparsing==3.3.2
+    # via matplotlib
 pyrate-limiter==4.4.0
     # via schemathesis
 pystemmer==3.1.0
@@ -697,7 +805,8 @@ pytest-timeout==2.4.0
     # via -r requirements/test/xpu.in
 python-dateutil==2.9.0.post0
     # via
-    #   jupyter-client
+    #   botocore
+    #   matplotlib
     #   pandas
     #   typepy
 python-dotenv==1.2.3
@@ -726,17 +835,18 @@ pyyaml==6.0.3
     #   timm
     #   transformers
     #   uvicorn
-pyzes==0.1.1
+pyzes==0.1.2
     # via torch
 pyzmq==27.1.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
-    #   jupyter-client
 rapidfuzz==3.14.5
     # via
     #   -r requirements/test/xpu.in
     #   jiwer
+redis==8.1.0
+    # via tensorizer
 referencing==0.37.0
     # via
     #   jsonschema
@@ -745,6 +855,7 @@ regex==2026.7.19
     # via
     #   -r requirements/test/../common.txt
     #   nltk
+    #   open-clip-torch
     #   sacrebleu
     #   tiktoken
     #   transformers
@@ -752,13 +863,18 @@ requests==2.34.2
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
+    #   azure-core
+    #   cohere
     #   datasets
     #   docker
     #   evaluate
+    #   google-api-core
+    #   google-cloud-storage
     #   gpt-oss
     #   lm-eval
     #   mistral-common
     #   modelscope
+    #   msal
     #   mteb
     #   opentelemetry-exporter-otlp-proto-http
     #   pooch
@@ -783,6 +899,16 @@ rpds-py==2026.6.3
     # via
     #   jsonschema
     #   referencing
+runai-model-streamer==0.16.1
+    # via -r requirements/test/xpu.in
+runai-model-streamer-azure==0.16.1
+    # via runai-model-streamer
+runai-model-streamer-gcs==0.16.1
+    # via runai-model-streamer
+runai-model-streamer-s3==0.16.1
+    # via runai-model-streamer
+s3transfer==0.19.2
+    # via boto3
 sacrebleu==2.6.0
     # via lm-eval
 safetensors==0.8.0
@@ -790,6 +916,7 @@ safetensors==0.8.0
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   accelerate
+    #   open-clip-torch
     #   timm
     #   transformers
 schemathesis==4.24.3
@@ -808,8 +935,10 @@ scipy==1.18.0
     #   pytrec-eval-terrier
     #   scikit-learn
     #   sentence-transformers
-sentence-transformers==5.7.0
-    # via mteb
+sentence-transformers==5.2.0
+    # via
+    #   -r requirements/test/xpu.in
+    #   mteb
 sentencepiece==0.2.1
     # via -r requirements/test/../common.txt
 sentry-sdk==2.68.0
@@ -879,12 +1008,12 @@ tabledata==1.3.5
     # via pytablewriter
 tabulate==0.10.0
     # via sacrebleu
-tbb==2023.0.0
+tbb==2023.1.0
     # via
     #   intel-opencl-rt
     #   mkl
     #   torch
-tblib==3.2.2
+tblib==3.1.0
     # via -r requirements/test/xpu.in
 tcmlib==1.5.0
     # via
@@ -897,6 +1026,8 @@ tenacity==9.1.4
     # via
     #   gpt-oss
     #   lm-eval
+tensorizer==2.10.1
+    # via -r requirements/test/xpu.in
 termcolor==3.3.0
     # via gpt-oss
 threadpoolctl==3.6.0
@@ -909,27 +1040,32 @@ tiktoken==0.14.0
     #   lm-eval
     #   mistral-common
 timm==1.0.28
-    # via -r requirements/test/xpu.in
+    # via
+    #   -r requirements/test/xpu.in
+    #   open-clip-torch
 tokenizers==0.23.1
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
-    #   sentence-transformers
+    #   cohere
     #   transformers
-torch==2.13.0+xpu
+torch==2.14.0+xpu
     # via
     #   -c requirements/xpu.txt
     #   accelerate
     #   compressed-tensors
     #   mteb
+    #   open-clip-torch
+    #   runai-model-streamer
     #   sentence-transformers
+    #   tensorizer
     #   timm
     #   torchvision
     #   xgrammar
-torchvision==0.28.0+xpu
-    # via timm
-tornado==6.5.8
-    # via jupyter-client
+torchvision==0.29.0+xpu
+    # via
+    #   open-clip-torch
+    #   timm
 tqdm==4.70.0
     # via
     #   -r requirements/test/../common.txt
@@ -940,15 +1076,12 @@ tqdm==4.70.0
     #   modelscope
     #   mteb
     #   nltk
+    #   open-clip-torch
     #   openai
     #   pqdm
     #   sentence-transformers
     #   transformers
-traitlets==5.16.1
-    # via
-    #   jupyter-client
-    #   jupyter-core
-transformers==5.16.1
+transformers==5.17.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -957,11 +1090,11 @@ transformers==5.16.1
     #   mteb
     #   sentence-transformers
     #   xgrammar
-triton==3.7.2+xpu
+triton==3.8.0+xpu
     # via
     #   -c requirements/xpu.txt
     #   xgrammar
-triton-xpu==3.7.2
+triton-xpu==3.8.0
     # via
     #   torch
     #   triton
@@ -979,6 +1112,8 @@ typer==0.27.1
     #   fastapi-cli
     #   fastapi-cloud-cli
     #   transformers
+types-requests==2.33.0.20260712
+    # via cohere
 typing-extensions==4.16.0
     # via
     #   -c requirements/common.txt
@@ -988,12 +1123,15 @@ typing-extensions==4.16.0
     #   anthropic
     #   anyio
     #   apache-tvm-ffi
+    #   azure-core
+    #   azure-identity
+    #   azure-storage-blob
     #   chz
+    #   cohere
     #   fastapi
     #   grpcio
     #   httpx2
     #   huggingface-hub
-    #   jupyter-client
     #   lm-eval
     #   mcp
     #   mcp-types
@@ -1032,10 +1170,12 @@ umf==1.1.0
 urllib3==2.7.0
     # via
     #   blobfile
+    #   botocore
     #   docker
     #   modelscope
     #   requests
     #   sentry-sdk
+    #   types-requests
 uvicorn==0.52.3
     # via
     #   fastapi
@@ -1049,13 +1189,15 @@ watchfiles==1.2.0
     # via
     #   -r requirements/test/../common.txt
     #   uvicorn
+wcwidth==0.8.3
+    # via ftfy
 websockets==17.0.1
     # via uvicorn
 werkzeug==3.1.8
     # via schemathesis
 word2number==1.1
     # via lm-eval
-xgrammar==0.2.3
+xgrammar==0.2.7
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
diff --git a/requirements/tpu.txt b/requirements/tpu.txt
index 8eaf8850d877..15e0e9ceabf8 100644
--- a/requirements/tpu.txt
+++ b/requirements/tpu.txt
@@ -12,4 +12,4 @@ ray[data]
 setuptools==78.1.0
 setuptools-rust>=1.9.0
 nixl==0.3.0
-tpu-inference==0.28.0
+tpu-inference==0.29.0
diff --git a/requirements/xpu.txt b/requirements/xpu.txt
index d0e94ca1897e..fd5b3214f0ca 100644
--- a/requirements/xpu.txt
+++ b/requirements/xpu.txt
@@ -1,7 +1,7 @@
 # Common dependencies
 -r common.txt
 --extra-index-url=https://wheels.vllm.ai/xpu/
-triton==3.7.2+xpu
+triton==3.8.0+xpu
 
 ray>=2.9
 cmake>=3.26.1
@@ -14,10 +14,10 @@ jinja2>=3.1.6
 datasets # for benchmark scripts
 numba == 0.65.0 # Required for N-gram speculative decoding
 --extra-index-url=https://download.pytorch.org/whl/xpu
-torch==2.13.0
+torch==2.14.0
 torchaudio
 torchvision
 torchcodec >= 0.14 # Required for the torchcodec video decoding backend
 
 auto_round_lib==0.15.0
-vllm_xpu_kernels==0.1.14.1
+vllm_xpu_kernels==0.1.15.4
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
index 7cd52d6b1c0e..b3f9084bcfaf 100644
--- a/rust/Cargo.lock
+++ b/rust/Cargo.lock
@@ -2290,7 +2290,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
 [[package]]
 name = "llm-multimodal"
 version = "1.7.1"
-source = "git+https://github.com/smg-project/llm-multimodal?rev=24c676dd49d73f05cb4068aade2e60397a5c0afe#24c676dd49d73f05cb4068aade2e60397a5c0afe"
+source = "git+https://github.com/smg-project/llm-multimodal?rev=f0985ef65967615db2c79279aa07818499301bfd#f0985ef65967615db2c79279aa07818499301bfd"
 dependencies = [
  "anyhow",
  "base64 0.22.1",
@@ -5735,6 +5735,11 @@ dependencies = [
  "url",
  "uuid",
  "vllm-build-info",
+ "vllm-chat",
+ "vllm-engine-core-client",
+ "vllm-llm",
+ "vllm-managed-engine",
+ "vllm-text",
  "vllm-tracing",
 ]
 
@@ -5784,6 +5789,7 @@ dependencies = [
  "vllm-parser",
  "vllm-text",
  "vllm-tokenizer",
+ "vllm-tracing",
  "xgrammar-structural-tag",
  "zeromq",
 ]
@@ -5813,6 +5819,7 @@ dependencies = [
  "vllm-engine-core-client",
  "vllm-managed-engine",
  "vllm-server",
+ "vllm-text",
  "vllm-tracing",
 ]
 
@@ -5930,6 +5937,7 @@ dependencies = [
  "openai-protocol",
  "serde",
  "serde_json",
+ "serde_with",
  "thiserror 2.0.18",
  "thiserror-ext",
  "tool-parser",
@@ -5938,6 +5946,18 @@ dependencies = [
  "xgrammar-structural-tag",
 ]
 
+[[package]]
+name = "vllm-proto"
+version = "0.4.0"
+dependencies = [
+ "prost",
+ "prost-types",
+ "protox",
+ "tonic",
+ "tonic-prost",
+ "tonic-prost-build",
+]
+
 [[package]]
 name = "vllm-server"
 version = "0.1.0"
@@ -5962,7 +5982,6 @@ dependencies = [
  "openssl",
  "prost",
  "prost-types",
- "protox",
  "rmp-serde",
  "rmpv",
  "serde",
@@ -5982,8 +6001,6 @@ dependencies = [
  "tokio-util",
  "tonic",
  "tonic-health",
- "tonic-prost",
- "tonic-prost-build",
  "tower",
  "tower-http",
  "tracing",
@@ -5997,6 +6014,7 @@ dependencies = [
  "vllm-engine-core-client",
  "vllm-llm",
  "vllm-metrics",
+ "vllm-proto",
  "vllm-text",
  "vllm-tokenizer",
  "zeromq",
@@ -6585,9 +6603,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
 
 [[package]]
 name = "xgrammar-structural-tag"
-version = "0.2.0+xgrammar.0.2.4.dd729e7"
+version = "0.3.0+xgrammar.0.2.7.82505d0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
+checksum = "3864005575c2ab76057f789898f59fb9fa1378d07c5ceabc24764e012d0e90bb"
 dependencies = [
  "auto_impl",
  "serde",
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index a9421e7c4aca..40b48f85bce0 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -1,5 +1,6 @@
 [workspace]
 members = [
+    "proto",
     "src/bench",
     "src/build-info",
     "src/chat",
@@ -60,7 +61,7 @@ indexmap = "2.13.0"
 indicatif = "0.18.4"
 itertools = "0.14.0"
 libc = "0.2.177"
-llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "24c676dd49d73f05cb4068aade2e60397a5c0afe", default-features = false, features = ["native-tls"] }
+llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "f0985ef65967615db2c79279aa07818499301bfd", default-features = false, features = ["native-tls"] }
 mimalloc = "0.1.52"
 minijinja = { version = "2.24", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
 minijinja-contrib = { version = "2.24", features = ["pycompat"] }
@@ -146,12 +147,13 @@ vllm-llm = { path = "src/llm" }
 vllm-managed-engine = { path = "src/managed-engine" }
 vllm-metrics = { path = "src/metrics" }
 vllm-parser = { path = "src/parser" }
+vllm-proto = { path = "proto" }
 vllm-server = { path = "src/server" }
 vllm-text = { path = "src/text" }
 vllm-tokenizer = { path = "src/tokenizer" }
 vllm-tracing = { path = "src/tracing" }
 winnow = { version = "1.0.2", features = ["simd"] }
-xgrammar-structural-tag = "0.2.0"
+xgrammar-structural-tag = "0.3.0"
 zeromq = { version = "0.6.0", default-features = false, features = [
     "tokio-runtime",
     "all-transport",
diff --git a/rust/README.md b/rust/README.md
index 06b50dbbfe9c..0cc33d64f8d2 100644
--- a/rust/README.md
+++ b/rust/README.md
@@ -43,6 +43,17 @@ For example:
 VLLM_USE_RUST_FRONTEND=1 vllm serve Qwen/Qwen3-0.6B
 ```
 
+### RL weight synchronization
+
+With `VLLM_SERVER_DEV_MODE=1`, the Rust frontend supports the HTTP weight-transfer
+lifecycle used by `HTTPVLLMWeightSyncClient`: initialization, starting an update,
+transferring weights, and finishing the update. It also supports draft-model
+updates and `/update_weight_version` and `/weight_info` for version tracking.
+
+For trainer-side usage, see the [IPC](../examples/rl/rlhf_http_ipc.py) and
+[NCCL](../examples/rl/rlhf_http_nccl.py) examples. Configure the weight-transfer
+backend on the server and pause generation during updates, as shown there.
+
 ### External Engine
 
 `vllm-rs serve` can be run standalone with `--data-parallel-size-local 0` when the Python engines
diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml
new file mode 100644
index 000000000000..492cb5a06d67
--- /dev/null
+++ b/rust/proto/Cargo.toml
@@ -0,0 +1,22 @@
+[package]
+name = "vllm-proto"
+version = "0.4.0"
+edition.workspace = true
+license.workspace = true
+description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend"
+repository = "https://github.com/vllm-project/vllm"
+readme = "README.md"
+include = ["*.proto", "src/**", "build.rs", "Cargo.toml", "README.md"]
+
+[dependencies]
+prost.workspace = true
+prost-types.workspace = true
+tonic.workspace = true
+tonic-prost.workspace = true
+
+[build-dependencies]
+protox.workspace = true
+tonic-prost-build.workspace = true
+
+[lints]
+workspace = true
diff --git a/rust/proto/README.md b/rust/proto/README.md
index a9960f3bb8ab..e792cf79e210 100644
--- a/rust/proto/README.md
+++ b/rust/proto/README.md
@@ -2,11 +2,58 @@
 
 This directory is the canonical source for vLLM's gRPC schema.
 
-The schema is published to `buf.build/vllm-project/vllm`:
+See [context parallel deployment](../../docs/serving/context_parallel_deployment.md)
+for effective attention block-size metadata in Python and `Control.GetServerInfo`.
 
-- A daily workflow publishes the latest Git `main` schema to the `nightly` label.
-- The workflow can be run manually to retry nightly publication.
-- Tags matching `v*` update the Buf `main` label and publish the corresponding release label.
-- Buf commits and generated SDK versions are immutable and can be pinned by consumers.
+Schema updates are no longer published to the Buf Schema Registry. Rust consumers
+should use `vllm-proto` from crates.io; consumers in other languages can generate
+bindings from the `.proto` files in this directory. Buf still builds and lints
+the schemas on pull requests, including on forks.
 
-Repository setup requires a `BUF_TOKEN` GitHub Actions secret with permission to create and push the public Buf module. Register the Prost and Tonic generated SDKs for the `main` and `nightly` labels once so subsequent pushes generate them automatically.
+## Rust crate
+
+`vllm-proto` exposes the Prost message types and Tonic client and server modules for both APIs:
+
+```toml
+[dependencies]
+vllm-proto = "0.3"
+```
+
+For example, import `vllm_proto::inference_client::InferenceClient` or
+`vllm_proto::control_client::ControlClient`. The vLLM frontend uses this same crate.
+The package includes the canonical `.proto` files and generates bindings during
+compilation using `protox` and `tonic-prost-build`. Consumers do not need Buf
+credentials, the vLLM checkout, or a separately installed `protoc`.
+
+### Versioning and releases
+
+The crate has its own version in `Cargo.toml`, independent of the vLLM release
+number and the other Rust workspace crates. Release it when changes to the
+schema, generated Rust API, or dependencies require a new version. Review both
+Rust API and protobuf wire compatibility when choosing the version bump;
+a wire-compatible schema addition can still break Rust callers.
+
+On pull requests and releases, `cargo-semver-checks` compares the crate with its
+latest published version. Include any required version bump in the protocol
+change PR. This check becomes available after the first manual publication.
+
+1. Update the crate version in `rust/proto/Cargo.toml` and update `rust/Cargo.lock`.
+2. Run `cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run`
+   and the frontend gRPC tests. Record the tested vLLM releases or revisions in
+   the release notes; matching crate versions alone do not establish runtime compatibility.
+3. After the change merges, create a `proto-v<version>` tag on that commit.
+   The `proto-crate.yml` workflow checks that the tag matches the crate version,
+   verifies the package, and publishes it to crates.io.
+
+There is no scheduled crate publication. Pull requests and manual workflow runs only
+verify packaging, including on forks. Retry a failed release by rerunning its
+tag workflow; crates.io versions cannot be overwritten.
+
+### Publisher setup
+
+A vLLM maintainer must publish the first version and establish the crate's owners.
+For subsequent releases, configure [crates.io Trusted Publishing](https://crates.io/docs/trusted-publishing)
+for repository `vllm-project/vllm`, workflow `proto-crate.yml`, and environment
+`crates-io`. Configure that GitHub environment to allow `proto-v*` release tags.
+Only tag pushes in the upstream repository can publish; forks do not obtain
+publishing credentials.
diff --git a/rust/proto/buf.md b/rust/proto/buf.md
deleted file mode 100644
index 065124feee65..000000000000
--- a/rust/proto/buf.md
+++ /dev/null
@@ -1,8 +0,0 @@
-# vLLM gRPC API
-
-Protocol definitions for the optional gRPC server exposed by the vLLM Rust frontend.
-
-- `Inference` provides unary and streaming generation.
-- `Control` provides discovery and request control.
-
-The `nightly` label is updated daily from vLLM's `main` branch. The `main` label tracks the latest vLLM release, and version labels match tags such as `v0.27.0`. Pin a Buf commit or generated SDK version for reproducible builds.
diff --git a/rust/proto/buf.yaml b/rust/proto/buf.yaml
index 64764f561ef3..d30dc3297d76 100644
--- a/rust/proto/buf.yaml
+++ b/rust/proto/buf.yaml
@@ -1,7 +1,6 @@
 version: v2
 modules:
   - path: .
-    name: buf.build/vllm-project/vllm
 lint:
   use:
     - MINIMAL
diff --git a/rust/src/server/build.rs b/rust/proto/build.rs
similarity index 88%
rename from rust/src/server/build.rs
rename to rust/proto/build.rs
index 1a2ba8a6b095..27cae507727b 100644
--- a/rust/src/server/build.rs
+++ b/rust/proto/build.rs
@@ -3,7 +3,7 @@
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
-    let proto_dir = std::path::PathBuf::from(format!("{manifest_dir}/../../proto"));
+    let proto_dir = std::path::PathBuf::from(manifest_dir);
     let protos = [
         proto_dir.join("control.proto"),
         proto_dir.join("inference.proto"),
diff --git a/rust/proto/control.proto b/rust/proto/control.proto
index 3050035fdcf6..ef3218d35c91 100644
--- a/rust/proto/control.proto
+++ b/rust/proto/control.proto
@@ -50,6 +50,8 @@ message ServerInfo {
   uint64 max_batched_tokens = 9;
   uint32 max_loras = 10;
   RlCapabilities rl_capabilities = 11;
+  // Full-attention block size in tokens after initialization; absent if unavailable.
+  optional uint64 effective_attention_block_size = 14;
 }
 
 message RlCapabilities {
@@ -66,6 +68,7 @@ message ParallelismInfo {
   uint32 data_parallel_rank = 4;
   uint32 decode_context_parallel_size = 5;
   uint64 world_size = 6;
+  uint32 data_parallel_size_local = 7;
 }
 
 message GetModelInfoRequest {}
diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto
index c854a9abdf1f..2bc35bcb64a2 100644
--- a/rust/proto/inference.proto
+++ b/rust/proto/inference.proto
@@ -53,6 +53,24 @@ message GenerateRequest {
 
   // Loaded LoRA name; empty selects the base model.
   string lora_name = 15;
+
+  // Apply the engine's configured watermark. Defaults to true when omitted.
+  optional bool watermarking = 16;
+
+  KVHintsEnvelope kv_hints = 17;
+}
+
+message KVHintsEnvelope {
+  string protocol_version = 1;
+  string message_id = 2;
+  repeated KVHintAction actions = 3;
+}
+
+message KVHintAction {
+  string action_id = 1;
+  string action_type = 2;
+  string action_version = 3;
+  google.protobuf.Struct payload = 4;
 }
 
 message RandomSampling {
@@ -157,6 +175,8 @@ message SequenceOutput {
 
   // Only present in final output for this sequence
   optional FinishInfo finish_info = 8;
+  // Only present in final output when sampling-mask capture is enabled.
+  repeated TokenIds sampling_mask = 10;
 }
 
 // Prompt info, returned in the first response
@@ -228,7 +248,20 @@ message MediaItem {
     string url = 2;       // http:// or https://
     string data_uri = 3;  // data:
     bytes raw_bytes = 4;
+    PreprocessedMediaFeatures features = 7;
   }
   string mime_type = 5;
   string uuid = 6;
 }
+
+// Engine-ready multimodal features produced by a version-compatible frontend.
+// `kwargs` is the MessagePack wire representation of one MmKwargsItem.
+message PreprocessedMediaFeatures {
+  // Required inline payload. Cache-only references are not accepted.
+  optional bytes kwargs = 1;
+  string identifier = 2;
+  uint64 offset = 3;
+  uint64 length = 4;
+  optional string mm_hash = 5;
+  repeated bool is_embed = 6;
+}
diff --git a/rust/proto/src/lib.rs b/rust/proto/src/lib.rs
new file mode 100644
index 000000000000..d4d44ea099f5
--- /dev/null
+++ b/rust/proto/src/lib.rs
@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Protobuf messages and gRPC clients and servers for vLLM's Inference and Control APIs.
+
+tonic::include_proto!("vllm");
diff --git a/rust/src/bench/AGENTS.md b/rust/src/bench/AGENTS.md
index 1f15165fa429..d9a17f63f29b 100644
--- a/rust/src/bench/AGENTS.md
+++ b/rust/src/bench/AGENTS.md
@@ -23,7 +23,7 @@ cargo test -p vllm-bench -- --ignored
 
 ## Architecture
 
-- `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single)
+- `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single, `mm-processor` subcommand)
 - `src/cli.rs` — clap derive CLI args (~50+ flags)
 - `src/config.rs` — Validated config from CLI; `GoodputConfig`, `RampUpConfig`, sampling param merging
 - `src/error.rs` — `BenchError` enum (Http, Json, Tokenizer, Config, EndpointTimeout, Backend, Io)
@@ -32,6 +32,7 @@ cargo test -p vllm-bench -- --ignored
 - `src/sweep.rs` — Concurrency/rate parameter sweep (`--sweep-max-concurrency`, `--sweep-request-rate`)
 - `src/multi_run.rs` — N-run aggregation with mean/std/min/max/CV (`--num-runs`)
 - `src/compare.rs` — Side-by-side diff of two result JSON files (`--compare`)
+- `src/mm_processor.rs` — `vllm-bench mm-processor` subcommand: offline multimodal preprocessing latency benchmark against a managed headless engine (mirrors `vllm bench mm-processor`); per-stage timings come from `vllm_tracing::timing::RequestTimingLayer` (via `vllm_chat::mm_timing_layer`)
 - `src/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback
 - `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source)
 - `src/hub.rs` — `HubRepo`: thin wrapper over hf-hub's async (reqwest/native-tls) API — file download and repo listing, optional revision; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans
diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml
index 5a3b14894ab8..60436d4251b4 100644
--- a/rust/src/bench/Cargo.toml
+++ b/rust/src/bench/Cargo.toml
@@ -36,6 +36,11 @@ tracing.workspace = true
 url.workspace = true
 uuid.workspace = true
 vllm-build-info.workspace = true
+vllm-chat.workspace = true
+vllm-engine-core-client.workspace = true
+vllm-llm.workspace = true
+vllm-managed-engine.workspace = true
+vllm-text.workspace = true
 vllm-tracing.workspace = true
 
 [lints]
diff --git a/rust/src/bench/README.md b/rust/src/bench/README.md
index d4e0c69b80b0..fb1e5b529dc8 100644
--- a/rust/src/bench/README.md
+++ b/rust/src/bench/README.md
@@ -18,6 +18,7 @@ vllm-bench --backend vllm --base-url http://127.0.0.1:8000 \
 - **Beyond a single run** — concurrency/rate **sweeps**, **multi-run** stats, **multi-turn** conversations, **LoRA** multi-adapter, and result **comparison**.
 - **Steady-state metrics** — throughput/latency measured over the saturated plateau, excluding ramp-up and drain.
 - **Parity** — JSON output schema and timing semantics match Python `vllm bench serve` exactly.
+- **Multimodal preprocessing** — `mm-processor` subcommand measures per-stage VLM preprocessing latency against a managed headless engine (mirrors `vllm bench mm-processor`).
 
 ### Performance vs. Python
 
@@ -404,6 +405,32 @@ vllm-bench \
 
 </details>
 
+<details>
+<summary><b>Multimodal preprocessing (offline)</b></summary>
+
+The `mm-processor` subcommand benchmarks the multimodal preprocessing pipeline
+itself (media fetch → per-modality preprocess → prompt expansion), mirroring
+`vllm bench mm-processor`. It spawns a managed headless Python engine (real
+weights, real encoder) — no serving endpoint needed — and reports per-stage
+timing (mean/median/std/percentiles) plus end-to-end latency.
+
+```bash
+vllm-bench mm-processor \
+  --model Qwen/Qwen2.5-VL-7B-Instruct \
+  --num-prompts 100 --num-warmups 5 \
+  --random-input-len 512 --random-output-len 128 \
+  --random-mm-base-items-per-request 1 \
+  --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \
+  --random-mm-bucket-config '{(1024, 800, 1): 1.0}' \
+  --metric-percentiles "50,90,99" \
+  --output-json mm_processor_stats.json
+```
+
+`--max-concurrency` defaults to `1`, matching Python's serial driver; raise it
+to measure a concurrent serving-style workload.
+
+</details>
+
 ## Supported Backends
 
 ### Generation
@@ -667,6 +694,7 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n
 
 - `--dataset-name random` — synthetic conversations with controllable per-turn token lengths. Auto-sets `min_tokens` to enforce output length without `ignore_eos`.
 - `--dataset-name sharegpt` — loads all turns (not just the first two); filters for entries with ≥ 2 real turns.
+- `--dataset-name hf` — downloads a ShareGPT-format HuggingFace config, then loads it like `sharegpt`.
 
 **Prefix sharing** (random dataset): when `--multi-turn-prefix-global-ratio` or `--multi-turn-prefix-conversation-ratio` is > 0, each turn sends a fixed-length message (no history accumulation) composed of a global prefix + per-conversation prefix + unique suffix. The two ratios must sum to < 1.0.
 
@@ -704,6 +732,36 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n
 
 </details>
 
+<details>
+<summary><b><code>mm-processor</code> subcommand</b></summary>
+
+Offline multimodal preprocessing latency benchmark (mirrors `vllm bench mm-processor`). Spawns a managed headless Python engine and drives the full chat preprocessing pipeline.
+
+| Flag | Default | Description |
+| ------ | --------- | ------------- |
+| `--model` | (required) | Model to serve and benchmark (HF id or local path) |
+| `--num-prompts` | `10` | Prompts to process, excluding warmups |
+| `--num-warmups` | `1` | Warmup prompts processed and discarded before timing |
+| `--metric-percentiles` | `99` | Comma-separated percentiles to report |
+| `--output-json` | — | Path to write aggregate stats as JSON (same schema as Python) |
+| `--random-input-len` | `1024` | Text input length per prompt |
+| `--random-output-len` | `128` | Expected output length per prompt |
+| `--random-prefix-len` | `0` | Prefix token length per prompt |
+| `--random-range-ratio` | `0.0` | Input/output length range ratio (float or `{"input": i, "output": o}`) |
+| `--random-mm-base-items-per-request` | `1` | Base number of multimodal items per request |
+| `--random-mm-num-mm-items-range-ratio` | `0.0` | Range ratio (in [0, 1]) for the number of mm items per request |
+| `--random-mm-limit-mm-per-prompt` | `{"image": 255, "video": 1}` | Per-modality item caps |
+| `--random-mm-bucket-config` | `{(256,256,1): 0.5, (720,1280,1): 0.5}` | Bucket config |
+| `--seed` | `0` | Seed for dataset generation |
+| `--request-id-prefix` | `mm-proc-` | Request-id prefix for per-request timing keys |
+| `--max-concurrency` | `1` | Maximum concurrent requests (1 matches Python's serial driver) |
+| `--trust-remote-code` | `false` | Trust remote code for the tokenizer and the managed engine |
+| `--chat-template` | — | Chat-template override (inline text or path) |
+
+Managed headless-engine options (e.g. `--python`, `--data-parallel-address`, `--data-parallel-rpc-port`, `--data-parallel-size`, `--max-model-len`, `--python-args`) are shared with other managed-engine commands; see `vllm-bench mm-processor --help`.
+
+</details>
+
 ## Tokenizer Support
 
 Tokenizers are loaded with a three-tier fallback chain:
@@ -733,10 +791,11 @@ Use `--append-result` to append multiple runs to the same file in JSONL format.
 
 ```text
 src/
-├── main.rs                  # Entry point, mimalloc, tokio runtime, mode dispatch
+├── main.rs                  # Entry point, mimalloc, tokio runtime, mode/subcommand dispatch
 ├── cli.rs                   # clap CLI argument definitions
 ├── config.rs                # Validated config, goodput/ramp-up parsing
 ├── benchmark.rs             # Core orchestrator (schedule, spawn, collect, verify, profile)
+├── mm_processor.rs          # Offline mm preprocessing benchmark (mm-processor subcommand)
 ├── multi_turn.rs            # Multi-turn conversation orchestrator (channel workers)
 ├── compare.rs               # Result diff (--compare file_a.json file_b.json)
 ├── sweep.rs                 # Parameter sweep (--sweep-max-concurrency, --sweep-request-rate)
diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs
index fc0f5a29e483..302dbb479231 100644
--- a/rust/src/bench/src/cli.rs
+++ b/rust/src/bench/src/cli.rs
@@ -176,7 +176,7 @@ pub struct BenchServeArgs {
     #[arg(long)]
     pub tokenizer: Option<String>,
 
-    /// Tokenizer mode (auto, hf, slow, mistral). Accepted for Python CLI
+    /// Tokenizer mode (auto, hf, mistral). Accepted for Python CLI
     /// compatibility; non-auto values are ignored with a warning.
     #[arg(long, default_value = "auto")]
     pub tokenizer_mode: String,
diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs
index 3a0936b5f4c1..b5a24d942f50 100644
--- a/rust/src/bench/src/config.rs
+++ b/rust/src/bench/src/config.rs
@@ -549,13 +549,15 @@ impl BenchConfig {
                 ));
             }
 
-            // Normalize and validate min/max turns. ShareGPT only consumes max_turns
-            // (the loader walks all available turns up to the cap), so the
-            // min/num/max coupling used for synthetic generation does not apply.
-            if args.dataset_name == DatasetName::ShareGpt {
+            // Normalize and validate min/max turns. Conversation datasets only
+            // consume max_turns (the loader walks all available turns up to the
+            // cap), so the min/num/max coupling used for synthetic generation
+            // does not apply.
+            if matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) {
                 if args.multi_turn_max_turns == 1 {
                     return Err(BenchError::Config(
-                        "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
+                        "--multi-turn-max-turns must be at least 2 for ShareGPT-format multi-turn"
+                            .into(),
                     ));
                 }
             } else {
@@ -722,7 +724,7 @@ impl BenchConfig {
             multi_turn_min_turns,
             multi_turn_max_turns,
             sharegpt_multi_turn_max_turns: if args.multi_turn
-                && args.dataset_name == DatasetName::ShareGpt
+                && matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf)
                 && args.multi_turn_max_turns != 0
             {
                 Some(args.multi_turn_max_turns)
@@ -1067,6 +1069,30 @@ mod tests {
         assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
     }
 
+    #[test]
+    fn test_hf_multi_turn_uses_conversation_turn_cap() {
+        let args = vec![
+            "vllm-bench",
+            "--backend",
+            "openai-chat",
+            "--multi-turn",
+            "--model",
+            "test-model",
+            "--dataset-name",
+            "hf",
+            "--dataset-path",
+            "org/sharegpt-dataset",
+            "--hf-subset",
+            "sharegpt",
+            "--multi-turn-max-turns",
+            "2",
+        ];
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
+
+        assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
+    }
+
     #[test]
     fn test_sweep_summary_percentiles_default_empty() {
         let args = base_multi_turn_args();
diff --git a/rust/src/bench/src/datasets/multi_turn.rs b/rust/src/bench/src/datasets/multi_turn.rs
index 97ec522e29ad..48ff11ff7bf6 100644
--- a/rust/src/bench/src/datasets/multi_turn.rs
+++ b/rust/src/bench/src/datasets/multi_turn.rs
@@ -350,6 +350,27 @@ pub fn load_sharegpt_multi_turn(
         .as_array()
         .ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?;
 
+    load_sharegpt_multi_turn_from_rows(
+        tokenizer,
+        entries,
+        num_conversations,
+        output_len_override,
+        max_turns,
+        seed,
+        request_id_prefix,
+    )
+}
+
+/// Load multi-turn conversations from deserialized ShareGPT rows.
+pub fn load_sharegpt_multi_turn_from_rows(
+    tokenizer: &TokenizerKind,
+    entries: &[serde_json::Value],
+    num_conversations: usize,
+    output_len_override: Option<usize>,
+    max_turns: Option<usize>,
+    seed: u64,
+    request_id_prefix: &str,
+) -> Result<Vec<MultiTurnConversation>> {
     // Filter entries with at least 4 messages (2 turns: user+assistant+user+assistant)
     let mut filtered: Vec<&serde_json::Value> = entries
         .iter()
@@ -507,6 +528,42 @@ fn gen_prompt_to_target_len(
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_load_sharegpt_multi_turn_from_hf_rows() {
+        let tokenizer =
+            TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
+        let rows = vec![serde_json::json!({
+            "conversations": [
+                {"from": "human", "value": "hello"},
+                {"from": "gpt", "value": "hi"},
+                {"from": "human", "value": "how are you?"},
+                {"from": "gpt", "value": "great"}
+            ]
+        })];
+
+        let conversations =
+            load_sharegpt_multi_turn_from_rows(&tokenizer, &rows, 1, Some(7), None, 42, "hf-")
+                .unwrap();
+        let actual = conversations
+            .iter()
+            .map(|conversation| {
+                (
+                    conversation.conversation_id.as_str(),
+                    conversation
+                        .turns
+                        .iter()
+                        .map(|turn| (turn.user_message.as_ref(), turn.expected_output_len))
+                        .collect::<Vec<_>>(),
+                )
+            })
+            .collect::<Vec<_>>();
+
+        assert_eq!(
+            actual,
+            vec![("hf-conv-0", vec![("hello", 7), ("how are you?", 7)])]
+        );
+    }
+
     fn common_prefix_bytes(strings: &[&str]) -> usize {
         if strings.is_empty() {
             return 0;
diff --git a/rust/src/bench/src/lib.rs b/rust/src/bench/src/lib.rs
index 0a4f19c0d6ee..8541c3b31098 100644
--- a/rust/src/bench/src/lib.rs
+++ b/rust/src/bench/src/lib.rs
@@ -10,6 +10,7 @@ mod datasets;
 mod error;
 mod hub;
 mod metrics;
+mod mm_processor;
 mod multi_run;
 mod multi_turn;
 mod output;
@@ -25,6 +26,7 @@ pub use cli::{
     BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
 };
 use config::BenchConfig;
+pub use mm_processor::{MmProcessorArgs, run_mm_processor};
 
 /// Prepare process-wide resources for a benchmark run.
 pub fn prepare_process() {
diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs
index 3485e22bb1e5..4a72b8189279 100644
--- a/rust/src/bench/src/main.rs
+++ b/rust/src/bench/src/main.rs
@@ -6,22 +6,30 @@
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
 use anyhow::Context;
-use clap::Parser;
+use clap::{Parser, Subcommand};
 
 #[derive(Parser)]
 #[command(
     name = "vllm-bench",
-    about = "Benchmark online serving throughput",
+    about = "Benchmark online serving throughput and offline multimodal preprocessing",
     version = vllm_build_info::VERSION
 )]
 struct Cli {
     #[command(flatten)]
     args: vllm_bench::BenchServeArgs,
+
+    /// Optional subcommand; absent by default for the online serving benchmark.
+    #[command(subcommand)]
+    command: Option<Command>,
 }
 
-fn main() -> anyhow::Result<()> {
-    vllm_tracing::init_tracing("Bench");
+#[derive(Subcommand)]
+enum Command {
+    /// Offline multimodal preprocessing latency benchmark.
+    MmProcessor(vllm_bench::MmProcessorArgs),
+}
 
+fn main() -> anyhow::Result<()> {
     let cli = Cli::parse();
     vllm_bench::prepare_process();
 
@@ -30,5 +38,15 @@ fn main() -> anyhow::Result<()> {
         .build()
         .context("Failed to build tokio runtime")?;
 
-    runtime.block_on(vllm_bench::run(cli.args))
+    match cli.command {
+        Some(Command::MmProcessor(args)) => {
+            let (timing_layer, timing_stats) = vllm_chat::mm_timing_layer();
+            vllm_tracing::init_tracing_with("Bench", timing_layer);
+            runtime.block_on(vllm_bench::run_mm_processor(args, timing_stats))
+        }
+        None => {
+            vllm_tracing::init_tracing("Bench");
+            runtime.block_on(vllm_bench::run(cli.args))
+        }
+    }
 }
diff --git a/rust/src/bench/src/metrics/calculator.rs b/rust/src/bench/src/metrics/calculator.rs
index 11b68a3bfeab..b292be42ebaf 100644
--- a/rust/src/bench/src/metrics/calculator.rs
+++ b/rust/src/bench/src/metrics/calculator.rs
@@ -216,20 +216,20 @@ pub fn calculate_metrics(
     (metrics, actual_output_lens)
 }
 
-fn mean(data: &[f64]) -> f64 {
+pub(crate) fn mean(data: &[f64]) -> f64 {
     if data.is_empty() {
         return 0.0;
     }
     data.iter().sum::<f64>() / data.len() as f64
 }
 
-fn sort_clone(data: &[f64]) -> Vec<f64> {
+pub(crate) fn sort_clone(data: &[f64]) -> Vec<f64> {
     let mut sorted = data.to_vec();
     sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
     sorted
 }
 
-fn median_sorted(sorted: &[f64]) -> f64 {
+pub(crate) fn median_sorted(sorted: &[f64]) -> f64 {
     if sorted.is_empty() {
         return 0.0;
     }
@@ -241,7 +241,7 @@ fn median_sorted(sorted: &[f64]) -> f64 {
     }
 }
 
-fn std_dev(data: &[f64]) -> f64 {
+pub(crate) fn std_dev(data: &[f64]) -> f64 {
     if data.is_empty() {
         return 0.0;
     }
@@ -250,7 +250,7 @@ fn std_dev(data: &[f64]) -> f64 {
     variance.sqrt()
 }
 
-fn percentile_sorted(sorted_data: &[f64], p: f64) -> f64 {
+pub(crate) fn percentile_sorted(sorted_data: &[f64], p: f64) -> f64 {
     if sorted_data.is_empty() {
         return 0.0;
     }
diff --git a/rust/src/bench/src/mm_processor.rs b/rust/src/bench/src/mm_processor.rs
new file mode 100644
index 000000000000..5b26181f1615
--- /dev/null
+++ b/rust/src/bench/src/mm_processor.rs
@@ -0,0 +1,573 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Multimodal preprocessing benchmark (`vllm-bench mm-processor`), mirroring
+//! `vllm bench mm-processor`.
+
+use std::collections::{BTreeMap, HashMap};
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use anyhow::{Context, Result};
+use clap::Args;
+use futures::future::join_all;
+use serde::Serialize;
+use tokio::sync::Semaphore;
+
+use vllm_chat::{
+    ChatContent, ChatContentPart, ChatLlm, ChatMessage, ChatRequest, LoadModelBackendsOptions,
+    load_model_backends,
+};
+use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
+use vllm_llm::Llm;
+use vllm_managed_engine::cli::ManagedEngineArgs;
+use vllm_managed_engine::{ManagedEngineConfig, ManagedEngineHandle};
+use vllm_text::TextLlm;
+use vllm_tracing::timing::RequestTimingStats;
+
+use crate::config::RangeRatio;
+use crate::datasets::SampleRequest;
+use crate::datasets::random_mm::{
+    MmBucketKey, MmLimitPerPrompt, generate_random_mm_dataset, parse_bucket_config,
+    parse_limit_mm_per_prompt,
+};
+use crate::error::BenchError;
+use crate::metrics::calculator::{mean, median_sorted, percentile_sorted, sort_clone, std_dev};
+use crate::tokenizer::{TokenizerKind, load_tokenizer};
+
+/// Seconds -> milliseconds.
+const SEC_TO_MS: f64 = 1000.0;
+
+/// Preferred column order; other stages appended alphabetically.
+const STAGE_ORDER: &[&str] = &[
+    "preprocessor_total_ms",
+    "media_fetch_ms",
+    "preprocess_image_ms",
+    "preprocess_video_ms",
+    "preprocess_audio_ms",
+    "prompt_expansion_ms",
+];
+
+/// Time to wait for managed engines to register.
+const ENGINE_READY_TIMEOUT: Duration = Duration::from_secs(600);
+
+/// Maximum time to wait for the managed engine to drain on shutdown.
+const ENGINE_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(120);
+
+#[derive(Args, Debug, Clone)]
+pub struct MmProcessorArgs {
+    /// Model to serve and benchmark (Hugging Face id or local path).
+    #[arg(long)]
+    pub model: String,
+
+    /// Number of prompts to process, excluding warmups.
+    #[arg(long, default_value_t = 10)]
+    pub num_prompts: usize,
+
+    /// Number of warmup prompts processed and discarded before timing.
+    #[arg(long, default_value_t = 1)]
+    pub num_warmups: usize,
+
+    /// Comma-separated percentiles to report, e.g. "50,90,99".
+    #[arg(long, default_value = "99")]
+    pub metric_percentiles: String,
+
+    /// Path to write the aggregate stats as JSON.
+    #[arg(long)]
+    pub output_json: Option<String>,
+
+    /// Text input length per prompt (before range scaling).
+    #[arg(long, default_value_t = 1024)]
+    pub random_input_len: usize,
+
+    /// Expected output length per prompt.
+    #[arg(long, default_value_t = 128)]
+    pub random_output_len: usize,
+
+    /// Prefix token length per prompt.
+    #[arg(long, default_value_t = 0)]
+    pub random_prefix_len: usize,
+
+    /// Input/output length range ratio: a float or `{"input": i, "output": o}`.
+    #[arg(long, default_value = "0.0")]
+    pub random_range_ratio: String,
+
+    /// Base number of multimodal items per request.
+    #[arg(long, default_value_t = 1)]
+    pub random_mm_base_items_per_request: usize,
+
+    /// Range ratio (in [0, 1]) for the number of mm items per request.
+    #[arg(long, default_value_t = 0.0)]
+    pub random_mm_num_mm_items_range_ratio: f64,
+
+    /// Per-modality item caps, e.g. `{"image": 255, "video": 1}`.
+    #[arg(long, default_value = "{\"image\": 255, \"video\": 1}")]
+    pub random_mm_limit_mm_per_prompt: String,
+
+    /// Bucket config, e.g. `{(256,256,1): 0.5, (720,1280,1): 0.5}`.
+    #[arg(long, default_value = "{(256,256,1): 0.5, (720,1280,1): 0.5}")]
+    pub random_mm_bucket_config: String,
+
+    /// Seed for dataset generation.
+    #[arg(long)]
+    pub seed: Option<u64>,
+
+    /// Request-id prefix for per-request timing keys.
+    #[arg(long, default_value = "mm-proc-")]
+    pub request_id_prefix: String,
+
+    /// Maximum concurrent requests (default 1, matching Python's serial driver).
+    #[arg(long, default_value_t = 1)]
+    pub max_concurrency: usize,
+
+    /// Trust remote code for the tokenizer and the managed engine.
+    #[arg(long, default_value_t = false)]
+    pub trust_remote_code: bool,
+
+    /// Optional chat-template override (inline text or path).
+    #[arg(long)]
+    pub chat_template: Option<String>,
+
+    /// Managed headless-engine options.
+    #[command(flatten)]
+    pub engine: ManagedEngineArgs,
+}
+
+/// Entry point for `vllm-bench mm-processor`.
+pub async fn run_mm_processor(args: MmProcessorArgs, timing: RequestTimingStats) -> Result<()> {
+    if args.engine.data_parallel_size_local == Some(0) {
+        anyhow::bail!(
+            "`--data-parallel-size-local 0` is not supported; the benchmark requires \
+             a local managed engine"
+        );
+    }
+
+    // Spawn the managed headless Python engine (real weights, real encoder).
+    let handshake_port = args
+        .engine
+        .resolve_handshake_port()
+        .context("failed to allocate handshake port")?;
+    let mut python_args = args.engine.python_args.clone();
+    if args.trust_remote_code {
+        python_args.push("--trust-remote-code".to_string());
+    }
+    let engine_config = ManagedEngineConfig {
+        python: args.engine.python.clone(),
+        model: args.model.clone(),
+        handshake_host: args.engine.handshake_host.clone(),
+        handshake_port,
+        data_parallel_size: args.engine.data_parallel_size,
+        python_args,
+    };
+    eprintln!(
+        "Spawning managed headless engine: python={} model={} data_parallel_size={}",
+        engine_config.python, engine_config.model, engine_config.data_parallel_size
+    );
+    let engine = ManagedEngineHandle::spawn(engine_config)
+        .await
+        .context("failed to start managed Python headless engine")?;
+
+    let result = run_mm_processor_with_engine(args, handshake_port, &timing).await;
+    let shutdown_result = engine
+        .shutdown(ENGINE_SHUTDOWN_TIMEOUT)
+        .await
+        .context("failed to shut down managed engine");
+    result?;
+    shutdown_result?;
+    Ok(())
+}
+
+async fn run_mm_processor_with_engine(
+    args: MmProcessorArgs,
+    handshake_port: u16,
+    timing: &RequestTimingStats,
+) -> Result<()> {
+    let range_ratio =
+        RangeRatio::parse(&args.random_range_ratio).context("invalid --random-range-ratio")?;
+    let limit: MmLimitPerPrompt = parse_limit_mm_per_prompt(&args.random_mm_limit_mm_per_prompt)
+        .context("invalid --random-mm-limit-mm-per-prompt")?;
+    let buckets: Vec<(MmBucketKey, f64)> = parse_bucket_config(&args.random_mm_bucket_config)
+        .context("invalid --random-mm-bucket-config")?;
+    let percentiles = parse_percentiles(&args.metric_percentiles)?;
+    let seed = args.seed.unwrap_or(0);
+
+    // Offline tokenizer: local Hugging Face tokenizer or built-in tiktoken.
+    let tokenizer = load_tokenizer(&args.model, args.trust_remote_code, None)
+        .await
+        .context("failed to load tokenizer")?;
+
+    // Frontend backends: model configs + multimodal processor metadata.
+    let backends = load_model_backends(
+        &args.model,
+        LoadModelBackendsOptions {
+            chat_template: args.chat_template.clone(),
+            ..LoadModelBackendsOptions::default()
+        },
+    )
+    .await
+    .context("failed to load model backends")?;
+
+    if backends.chat_backend.multimodal_model_info().is_none() {
+        anyhow::bail!(
+            "model `{}` has no multimodal processor registered",
+            args.model
+        );
+    }
+
+    // Connect to the engine over the same transport the serving frontend uses.
+    let client = EngineCoreClient::connect(EngineCoreClientConfig {
+        transport_mode: TransportMode::HandshakeOwner {
+            handshake_address: format!("tcp://{}:{}", args.engine.handshake_host, handshake_port),
+            advertised_host: args.engine.handshake_host.clone(),
+            engine_count: args
+                .engine
+                .data_parallel_size_local
+                .unwrap_or(args.engine.data_parallel_size),
+            ready_timeout: ENGINE_READY_TIMEOUT,
+            local_input_address: None,
+            local_output_address: None,
+        },
+        coordinator_mode: None,
+        model_name: args.model.clone(),
+        client_index: 0,
+    })
+    .await
+    .context("failed to connect to engine core")?;
+
+    let chat = ChatLlm::new(
+        TextLlm::new(Llm::new(client), backends.text_backend),
+        backends.chat_backend,
+    );
+
+    let semaphore = Arc::new(Semaphore::new(args.max_concurrency));
+
+    // Warmup: generate a separate dataset, process it, then drain and discard
+    // the collected timings so cold caches do not skew the results.
+    if args.num_warmups > 0 {
+        let warmup = generate_dataset(
+            &tokenizer,
+            args.num_warmups,
+            seed.wrapping_add(1),
+            &args,
+            range_ratio,
+            &limit,
+            &buckets,
+        )?;
+        eprintln!("Processing {} warmup requests...", warmup.len());
+        run_batch(&chat, &warmup, &semaphore).await;
+        let _ = timing.stat();
+    }
+
+    let samples = generate_dataset(
+        &tokenizer,
+        args.num_prompts,
+        seed,
+        &args,
+        range_ratio,
+        &limit,
+        &buckets,
+    )?;
+    eprintln!("Processing {} requests...", samples.len());
+    let start_time = Instant::now();
+    let outcomes = run_batch(&chat, &samples, &semaphore).await;
+    let total_time = start_time.elapsed();
+
+    let completed = outcomes.iter().filter(|outcome| outcome.is_ok()).count();
+    let failed = outcomes.len() - completed;
+    let e2el_times: Vec<f64> = outcomes
+        .iter()
+        .filter_map(|outcome| outcome.as_ref().ok().map(|elapsed| elapsed.as_secs_f64()))
+        .collect();
+
+    let per_request = timing.stat();
+    let stats = aggregate_stats(per_request, &percentiles);
+
+    eprintln!(
+        "Processed {} requests (+{} warmup) for model `{}`",
+        samples.len(),
+        args.num_warmups,
+        args.model
+    );
+    print_report(&stats, &percentiles);
+    print_e2el_summary(completed, failed, total_time, &e2el_times, &percentiles);
+
+    if let Some(path) = &args.output_json {
+        write_result_json(path, completed, failed, &e2el_times, &percentiles, stats)?;
+        eprintln!("Wrote stats JSON to {path}");
+    }
+
+    chat.shutdown().await.context("failed to shut down chat facade")?;
+    Ok(())
+}
+
+/// Submit every sample through the chat pipeline and drive each stream to
+/// completion, returning per-request wall-clock latencies.
+async fn run_batch(
+    chat: &ChatLlm,
+    samples: &[SampleRequest],
+    semaphore: &Arc<Semaphore>,
+) -> Vec<Result<Duration>> {
+    let submissions = samples.iter().map(|sample| {
+        let request = build_chat_request(sample);
+        let semaphore = semaphore.clone();
+        async move {
+            let _permit = semaphore.acquire().await.context("semaphore closed")?;
+            let request = request?;
+            let start = Instant::now();
+            chat.chat(request)
+                .await
+                .context("failed to submit request")?
+                .collect_message()
+                .await
+                .context("request failed")?;
+            Ok(start.elapsed())
+        }
+    });
+    join_all(submissions).await
+}
+
+fn generate_dataset(
+    tokenizer: &TokenizerKind,
+    num_requests: usize,
+    seed: u64,
+    args: &MmProcessorArgs,
+    range_ratio: RangeRatio,
+    limit: &MmLimitPerPrompt,
+    buckets: &[(MmBucketKey, f64)],
+) -> Result<Vec<SampleRequest>> {
+    generate_random_mm_dataset(
+        tokenizer,
+        num_requests,
+        args.random_input_len,
+        args.random_output_len,
+        args.random_prefix_len,
+        range_ratio,
+        seed,
+        &args.request_id_prefix,
+        args.random_mm_base_items_per_request,
+        args.random_mm_num_mm_items_range_ratio,
+        limit,
+        buckets,
+        false,
+    )
+    .map_err(|e| anyhow::anyhow!("{e}"))
+}
+
+/// Build a chat request from a sample (text prompt + `image_url` parts).
+fn build_chat_request(sample: &SampleRequest) -> Result<ChatRequest> {
+    let mut request = ChatRequest::for_test();
+    request.request_id = sample.request_id.clone().unwrap_or_else(|| "mm-proc".to_string());
+    request.sampling_params.temperature = Some(0.0);
+    request.sampling_params.max_tokens = Some(u32::try_from(sample.expected_output_len)?);
+
+    let mut parts = vec![ChatContentPart::text(sample.prompt.as_ref())];
+    if let Some(mm_items) = &sample.multi_modal_content {
+        for fragment in mm_items.iter() {
+            let url = extract_image_url(fragment)?;
+            parts.push(ChatContentPart::image_url(url));
+        }
+    }
+    request.messages = vec![ChatMessage::user(ChatContent::from(parts))];
+    Ok(request)
+}
+
+/// Extract the `image_url.url` string from a pre-serialized content fragment
+/// (`{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,..."}}`).
+fn extract_image_url(fragment: &str) -> Result<String> {
+    let value: serde_json::Value = serde_json::from_str(fragment)
+        .with_context(|| format!("invalid multimodal fragment: {fragment}"))?;
+    let url = value
+        .get("image_url")
+        .and_then(|v| v.get("url"))
+        .and_then(|v| v.as_str())
+        .ok_or_else(|| BenchError::Config("fragment missing image_url.url".to_string()))?;
+    Ok(url.to_string())
+}
+
+fn parse_percentiles(raw: &str) -> Result<Vec<f64>> {
+    let mut values = Vec::new();
+    for part in raw.split(',') {
+        let part = part.trim();
+        if part.is_empty() {
+            continue;
+        }
+        let value: f64 = part.parse().with_context(|| format!("invalid percentile `{part}`"))?;
+        if !(0.0..=100.0).contains(&value) {
+            anyhow::bail!("percentile `{part}` out of range [0, 100]");
+        }
+        values.push(value);
+    }
+    Ok(values)
+}
+
+/// Aggregate per-request stage timings into `{stage_ms: {statistic: ms}}`,
+/// matching the Python `mm_processor_stats` schema orientation.
+fn aggregate_stats(
+    per_request: HashMap<String, HashMap<String, f64>>,
+    percentiles: &[f64],
+) -> BTreeMap<String, BTreeMap<String, f64>> {
+    let mut stage_values: BTreeMap<String, Vec<f64>> = BTreeMap::new();
+    for stage_map in per_request.into_values() {
+        for (stage, seconds) in stage_map {
+            stage_values.entry(stage).or_default().push(seconds);
+        }
+    }
+
+    stage_values
+        .iter()
+        .map(|(stage, values)| {
+            let sorted = sort_clone(values);
+            let mut stat_map = BTreeMap::new();
+            stat_map.insert("mean".to_string(), mean(values) * SEC_TO_MS);
+            stat_map.insert("median".to_string(), median_sorted(&sorted) * SEC_TO_MS);
+            stat_map.insert("std".to_string(), std_dev(values) * SEC_TO_MS);
+            for p in percentiles {
+                stat_map.insert(format!("p{p}"), percentile_sorted(&sorted, *p) * SEC_TO_MS);
+            }
+            (stage.replace("_secs", "_ms"), stat_map)
+        })
+        .collect()
+}
+
+/// Print a right-aligned plain-text table (pandas `to_string(index=False)` style).
+fn print_table(rows: &[Vec<String>]) {
+    let ncols = rows.first().map_or(0, Vec::len);
+    let mut widths = vec![0usize; ncols];
+    for row in rows {
+        for (i, cell) in row.iter().enumerate() {
+            widths[i] = widths[i].max(cell.len());
+        }
+    }
+    for row in rows {
+        let line = row
+            .iter()
+            .enumerate()
+            .map(|(i, cell)| format!("{cell:>width$}", width = widths[i]))
+            .collect::<Vec<_>>()
+            .join(" ");
+        println!("{line}");
+    }
+}
+
+/// Print a plain-text table with one row per stage and one column per
+/// statistic (mirrors the Python pandas report layout).
+fn print_report(stats: &BTreeMap<String, BTreeMap<String, f64>>, percentiles: &[f64]) {
+    // Column order: mean/median/std then the requested percentiles, like Python.
+    let mut columns: Vec<(String, String)> =
+        [("mean", "Mean"), ("median", "Median"), ("std", "Std")]
+            .into_iter()
+            .map(|(key, label)| (key.to_string(), label.to_string()))
+            .collect();
+    columns.extend(percentiles.iter().map(|p| (format!("p{p}"), format!("P{p:?}"))));
+    for key in stats.values().flat_map(|map| map.keys()) {
+        if !columns.iter().any(|(k, _)| k == key) {
+            columns.push((key.clone(), key.clone()));
+        }
+    }
+
+    let mut stages: Vec<&String> = stats.keys().collect();
+    stages.sort_by_key(|stage| {
+        STAGE_ORDER
+            .iter()
+            .position(|s| s == stage)
+            .map(|i| (i as isize, (*stage).clone()))
+            .unwrap_or((isize::MAX, (*stage).clone()))
+    });
+
+    let mut rows: Vec<Vec<String>> = vec![
+        std::iter::once("Stage".to_string())
+            .chain(columns.iter().map(|(_, label)| label.clone()))
+            .collect(),
+    ];
+    for stage in &stages {
+        let mut row = vec![(*stage).clone()];
+        for (key, _) in &columns {
+            row.push(
+                stats[*stage]
+                    .get(key)
+                    .map(|v| format!("{v:.2}"))
+                    .unwrap_or_else(|| "-".to_string()),
+            );
+        }
+        rows.push(row);
+    }
+    println!("MM Processor Metrics:");
+    print_table(&rows);
+}
+
+/// Print the end-to-end latency summary (mirrors the Python report layout).
+fn print_e2el_summary(
+    completed: usize,
+    failed: usize,
+    total_time: Duration,
+    e2el_times_secs: &[f64],
+    percentiles: &[f64],
+) {
+    println!("Completed: {completed}, Failed: {failed}");
+    if e2el_times_secs.is_empty() {
+        println!("End-to-End Latency (ms): no completed requests");
+        return;
+    }
+
+    let e2el_ms: Vec<f64> = e2el_times_secs.iter().map(|secs| secs * SEC_TO_MS).collect();
+    let sorted = sort_clone(&e2el_ms);
+    let mut rows = vec![
+        vec!["Metric".to_string(), "Value (ms)".to_string()],
+        vec!["Mean".to_string(), format!("{:.2}", mean(&e2el_ms))],
+        vec![
+            "Median".to_string(),
+            format!("{:.2}", median_sorted(&sorted)),
+        ],
+        vec!["Std".to_string(), format!("{:.2}", std_dev(&e2el_ms))],
+    ];
+    for p in percentiles {
+        rows.push(vec![
+            format!("P{p:?}"),
+            format!("{:.2}", percentile_sorted(&sorted, *p)),
+        ]);
+    }
+    println!("End-to-End Latency (ms):");
+    print_table(&rows);
+    let throughput = completed as f64 / total_time.as_secs_f64();
+    println!("Request throughput: {throughput:.3} req/s");
+}
+
+/// JSON result schema mirroring `vllm bench mm-processor`.
+#[derive(Serialize)]
+struct BenchmarkResultJson {
+    completed: usize,
+    failed: usize,
+    mean_e2el_ms: f64,
+    median_e2el_ms: f64,
+    std_e2el_ms: f64,
+    percentiles_e2el_ms: Vec<(f64, f64)>,
+    /// `{stage_ms: {statistic: ms}}`, keyed by stage like the Python schema.
+    mm_processor_stats: BTreeMap<String, BTreeMap<String, f64>>,
+}
+
+fn write_result_json(
+    path: &str,
+    completed: usize,
+    failed: usize,
+    e2el_times_secs: &[f64],
+    percentiles: &[f64],
+    mm_processor_stats: BTreeMap<String, BTreeMap<String, f64>>,
+) -> Result<()> {
+    let e2el_ms: Vec<f64> = e2el_times_secs.iter().map(|secs| secs * SEC_TO_MS).collect();
+    let sorted = sort_clone(&e2el_ms);
+    let result = BenchmarkResultJson {
+        completed,
+        failed,
+        mean_e2el_ms: mean(&e2el_ms),
+        median_e2el_ms: median_sorted(&sorted),
+        std_e2el_ms: std_dev(&e2el_ms),
+        percentiles_e2el_ms: percentiles
+            .iter()
+            .map(|p| (*p, percentile_sorted(&sorted, *p)))
+            .collect(),
+        mm_processor_stats,
+    };
+    let json = serde_json::to_string_pretty(&result).context("failed to serialize result")?;
+    std::fs::write(path, format!("{json}\n")).with_context(|| format!("failed to write {path}"))?;
+    Ok(())
+}
diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs
index e0a31381705f..366fd4f6fbc1 100644
--- a/rust/src/bench/src/multi_turn.rs
+++ b/rust/src/bench/src/multi_turn.rs
@@ -186,9 +186,30 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
             ));
         }
         DatasetName::Hf => {
-            return Err(BenchError::Config(
-                "HF dataset multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(),
-            ));
+            let tok = tokenizer
+                .as_ref()
+                .ok_or_else(|| BenchError::Config("HF dataset requires a tokenizer".into()))?;
+            let dataset_id = config.dataset_path.as_deref().ok_or_else(|| {
+                BenchError::Config("--dataset-path is required for --dataset-name hf".into())
+            })?;
+            let (rows, _config, _split) = crate::datasets::hf_dataset::download_hf_dataset(
+                dataset_id,
+                config.hf_subset.as_deref(),
+                config.hf_split.as_deref(),
+                config.num_prompts,
+                config.seed,
+                config.disable_shuffle,
+            )
+            .await?;
+            crate::datasets::multi_turn::load_sharegpt_multi_turn_from_rows(
+                tok,
+                &rows,
+                config.num_prompts,
+                config.hf_output_len,
+                config.sharegpt_multi_turn_max_turns,
+                config.seed,
+                &config.request_id_prefix,
+            )?
         }
     };
 
diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml
index e92dedae71f5..a110020ec110 100644
--- a/rust/src/chat/Cargo.toml
+++ b/rust/src/chat/Cargo.toml
@@ -28,6 +28,7 @@ thiserror-ext.workspace = true
 time.workspace = true
 tokio.workspace = true
 tracing.workspace = true
+tracing-subscriber.workspace = true
 trait-set.workspace = true
 uuid.workspace = true
 vllm-engine-core-client.workspace = true
@@ -35,6 +36,7 @@ vllm-llm.workspace = true
 vllm-parser.workspace = true
 vllm-text.workspace = true
 vllm-tokenizer.workspace = true
+vllm-tracing.workspace = true
 xgrammar-structural-tag.workspace = true
 
 [dev-dependencies]
@@ -47,7 +49,6 @@ rmp-serde.workspace = true
 serial_test.workspace = true
 tempfile.workspace = true
 tokio.workspace = true
-tracing-subscriber.workspace = true
 uuid.workspace = true
 vllm-engine-core-client = { workspace = true, features = ["test-util"] }
 vllm-tokenizer = { workspace = true, features = ["test-utils"] }
diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs
index 717bde289836..07f0b2dce815 100644
--- a/rust/src/chat/src/backend/hf.rs
+++ b/rust/src/chat/src/backend/hf.rs
@@ -70,12 +70,26 @@ impl HfChatBackend {
                 options,
                 multimodal_render_info,
             )?),
-            RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new()),
-            RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()),
-            RendererSelection::DeepSeekV41 => Arc::new(DeepSeekV41ChatRenderer::new()),
-            RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?),
-            RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?),
-            RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new(tokenizer.clone())),
+            RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new(
+                options.default_chat_template_kwargs,
+            )),
+            RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new(
+                options.default_chat_template_kwargs,
+            )),
+            RendererSelection::DeepSeekV41 => Arc::new(DeepSeekV41ChatRenderer::new(
+                options.default_chat_template_kwargs,
+            )),
+            RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new(
+                options.default_chat_template_kwargs,
+            )?),
+            RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(
+                tokenizer.clone(),
+                options.default_chat_template_kwargs,
+            )?),
+            RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new(
+                tokenizer.clone(),
+                options.default_chat_template_kwargs,
+            )),
         };
 
         info!(
@@ -120,6 +134,7 @@ impl ChatBackend for HfChatBackend {
             self.tokenizer.clone(),
             options.tool_call_parser,
             options.reasoning_parser,
+            options.tool_strict_level,
         )?))
     }
 }
@@ -129,7 +144,8 @@ pub(super) async fn load_model_backends(
     model_id: &str,
     options: LoadModelBackendsOptions,
 ) -> Result<LoadedModelBackends> {
-    let files = ResolvedModelFiles::new(model_id).await?;
+    let mut files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?;
+    files.apply_overrides(&options.hf_overrides)?;
     let text_backend = HfTextBackend::from_resolved_model_files(
         files.clone(),
         model_id.to_string(),
@@ -166,13 +182,11 @@ fn resolve_multimodal_render_info(
 #[cfg(test)]
 mod tests {
     use std::collections::HashMap;
-    use std::path::PathBuf;
     use std::sync::Arc;
 
     use tempfile::tempdir;
     use thiserror_ext::AsReport as _;
     use vllm_text::Prompt;
-    use vllm_text::backend::hf::TokenizerSource;
     use vllm_text::tokenizer::DynTokenizer;
     use vllm_tokenizer::test_utils::TestTokenizer;
 
@@ -206,16 +220,12 @@ mod tests {
         write_json(&config_path, config_json);
         write_json(&tokenizer_config_path, tokenizer_config_json);
 
-        vllm_text::backend::hf::ResolvedModelFiles {
-            tokenizer: TokenizerSource::HuggingFace(PathBuf::from("/tmp/unused-tokenizer.json")),
-            tokenizer_config_path: Some(tokenizer_config_path),
-            generation_config_path: None,
-            preprocessor_config_path: None,
-            video_preprocessor_config_path: None,
-            processor_config_path: None,
-            chat_template_path: None,
-            config_path: Some(config_path),
-        }
+        write_json(&root.join("tokenizer.json"), "{}");
+        futures::executor::block_on(vllm_text::backend::hf::ResolvedModelFiles::new(
+            root.to_str().unwrap(),
+            None,
+        ))
+        .unwrap()
     }
 
     fn test_tokenizer() -> DynTokenizer {
@@ -231,6 +241,8 @@ mod tests {
             resolved_files(config_json, tokenizer_config_json),
             "test-model".to_string(),
             LoadModelBackendsOptions {
+                revision: None,
+                hf_overrides: Default::default(),
                 generation_config: Default::default(),
                 renderer,
                 language_model_only: false,
@@ -283,6 +295,39 @@ mod tests {
         assert_eq!(prompt, "hello");
     }
 
+    #[test]
+    fn native_renderer_inherits_deployment_reasoning_below_request_controls() {
+        let backend = HfChatBackend::from_resolved_model_files(
+            resolved_files(r#"{"model_type":"deepseek_v4"}"#, "{}"),
+            "test-model".to_string(),
+            LoadModelBackendsOptions {
+                default_chat_template_kwargs: [
+                    ("thinking".to_string(), serde_json::json!(false)),
+                    ("reasoning_effort".to_string(), serde_json::json!("low")),
+                ]
+                .into(),
+                ..Default::default()
+            },
+            test_tokenizer(),
+        )
+        .unwrap();
+        let mut request = request_with_user_text("hello");
+        let rendered = backend.chat_renderer().render(&request).unwrap();
+        assert!(rendered.prompt.into_text().unwrap().ends_with("</think>"));
+        assert_eq!(rendered.effective_template_kwargs["enable_thinking"], false);
+
+        request.chat_options.reasoning_effort = Some(crate::EffortValue::from("max"));
+        let rendered = backend.chat_renderer().render(&request).unwrap();
+        let prompt = rendered.prompt.into_text().unwrap();
+        assert!(prompt.contains("Reasoning Effort: Beyond maximum"));
+        assert!(prompt.ends_with("<think>"));
+        assert_eq!(
+            rendered.effective_template_kwargs["reasoning_effort"],
+            "max"
+        );
+        assert_eq!(rendered.effective_template_kwargs["thinking"], true);
+    }
+
     #[test]
     fn auto_uses_harmony_renderer_and_output_processor_for_gpt_oss_model_type() {
         let backend = backend_for_selection(
@@ -299,6 +344,7 @@ mod tests {
         let error = match backend.new_chat_output_processor(
             &mut request,
             NewChatOutputProcessorOptions {
+                tool_strict_level: crate::ToolStrictLevel::Auto,
                 tool_call_parser: &ParserSelection::Explicit("json".to_string()),
                 reasoning_parser: &ParserSelection::Auto,
             },
diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs
index 77b9114a5852..a0343e53f327 100644
--- a/rust/src/chat/src/backend/mod.rs
+++ b/rust/src/chat/src/backend/mod.rs
@@ -5,6 +5,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use serde_json::Value;
+use vllm_text::backend::hf::HfOverrides;
 use vllm_text::{DynTextBackend, GenerationConfigMode, TextBackend};
 
 use crate::error::Result;
@@ -12,7 +13,7 @@ use crate::multimodal::{MmLimitPerPrompt, MultimodalModelInfo};
 use crate::output::DynChatOutputProcessor;
 use crate::renderer::DynChatRenderer;
 use crate::request::ChatRequest;
-use crate::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
+use crate::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection, ToolStrictLevel};
 
 pub mod hf;
 
@@ -20,6 +21,7 @@ pub mod hf;
 pub struct NewChatOutputProcessorOptions<'a> {
     pub tool_call_parser: &'a ParserSelection,
     pub reasoning_parser: &'a ParserSelection,
+    pub tool_strict_level: ToolStrictLevel,
 }
 
 /// Minimal prompt-processing backend needed by `vllm-chat`.
@@ -61,6 +63,10 @@ pub type DynChatTextBackend = Arc<dyn ChatTextBackend>;
 /// Frontend-side chat backend loading options.
 #[derive(Debug, Clone, PartialEq, Eq, Default)]
 pub struct LoadModelBackendsOptions {
+    /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
+    pub revision: Option<String>,
+    /// JSON Merge Patch applied to the model config before loading any backend.
+    pub hf_overrides: HfOverrides,
     /// Which generation-config sampling defaults to inherit.
     pub generation_config: GenerationConfigMode,
     /// Which chat renderer implementation to use.
@@ -73,8 +79,8 @@ pub struct LoadModelBackendsOptions {
     /// Optional server-default chat template override, provided either as an
     /// inline template or as a path to a template file.
     pub chat_template: Option<String>,
-    /// Optional server-default keyword arguments merged into every
-    /// chat-template render before request-level `chat_template_kwargs`.
+    /// Server-default keyword arguments. HF merges these before request kwargs;
+    /// native renderers inherit supported reasoning controls below request controls.
     pub default_chat_template_kwargs: HashMap<String, Value>,
     /// Maximum number of input items allowed per prompt for each modality.
     /// Unspecified modalities are unlimited.
diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs
index e3299dfdd4e1..7d124aaaf05b 100644
--- a/rust/src/chat/src/error.rs
+++ b/rust/src/chat/src/error.rs
@@ -19,6 +19,10 @@ pub enum Error {
     ChatTemplate(String),
     #[error("{0}")]
     InvalidReasoningEffort(String),
+    #[error("{message}")]
+    InvalidReasoningControl { message: String },
+    #[error("chat role `{role}` is not supported by this chat renderer")]
+    UnsupportedChatRole { role: String },
     #[error("multimodal input is not supported by this chat renderer")]
     UnsupportedMultimodalRenderer,
     #[error("unsupported multimodal content: {0}")]
@@ -27,6 +31,8 @@ pub enum Error {
     UnsupportedModality { modality: String },
     #[error("At most {limit} {modality}(s) may be provided in one prompt.")]
     MmLimitExceeded { modality: String, limit: usize },
+    #[error("invalid inline multimodal features: {message}")]
+    InvalidPreprocessedMultimodal { message: String },
     #[error("multimodal preprocessing error: {0}")]
     Multimodal(#[message] String),
     #[error("{kind} parsing is not available for model `{model_id}`")]
@@ -56,6 +62,11 @@ pub enum Error {
         #[source]
         error: BoxedError,
     },
+    #[error("failed to initialize request output parser")]
+    OutputParserInitialization {
+        #[source]
+        error: BoxedError,
+    },
     #[error(
         "gpt_oss uses native Harmony output parsing; generic {kind} parser override `{selection}` is not supported"
     )]
@@ -83,8 +94,11 @@ pub enum Error {
     ToolChoiceRequiresTools,
     #[error("tool_choice function `{name}` was not found in the available tools")]
     ToolChoiceFunctionNotFound { name: String },
-    #[error("failed to build structural tag: {message}")]
-    StructuralTag { message: String },
+    #[error("failed to build output grammar")]
+    OutputGrammar {
+        #[source]
+        error: BoxedError,
+    },
     #[error(transparent)]
     Text(#[from] vllm_text::Error),
     #[error(transparent)]
@@ -99,13 +113,16 @@ impl Error {
         match self {
             Self::PromptTooLong { .. }
             | Self::InvalidReasoningEffort(_)
+            | Self::InvalidReasoningControl { .. }
             | Self::DuplicateToolName { .. }
             | Self::ToolChoiceRequiresTools
-            | Self::ToolChoiceFunctionNotFound { .. } => true,
+            | Self::ToolChoiceFunctionNotFound { .. }
+            | Self::UnsupportedChatRole { .. } => true,
             Self::Text(error) => error.is_request_validation_error(),
             Self::UnsupportedMultimodalRenderer
             | Self::UnsupportedMultimodalContent(_)
             | Self::UnsupportedModality { .. }
+            | Self::InvalidPreprocessedMultimodal { .. }
             | Self::MmLimitExceeded { .. } => true,
 
             _ => false,
diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs
index 6250330e4868..95c4332f7476 100644
--- a/rust/src/chat/src/lib.rs
+++ b/rust/src/chat/src/lib.rs
@@ -22,6 +22,8 @@ pub use event::{
 };
 use futures::{StreamExt, TryStreamExt as _};
 pub use llm_multimodal::MediaContentPart;
+pub use multimodal::{mm_request_span, mm_timing_layer};
+use output::apply_output_grammar;
 pub use output::{
     ChatOutputProcessor, DefaultChatOutputProcessor, DynChatOutputProcessor,
     HarmonyChatOutputProcessor,
@@ -30,16 +32,17 @@ pub use parser::reasoning::{
     ReasoningDelta, ReasoningError, ReasoningParser, ReasoningParserFactory,
 };
 pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory};
-pub use parser::{ParserSelection, validate_parser_overrides};
+pub use parser::{ParserSelection, ToolStrictLevel, validate_parser_overrides};
+pub use reasoning::EffortValue;
 pub use renderer::hf::ChatTemplateContentFormatOption;
 pub use renderer::{
     ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DeepSeekV41ChatRenderer,
-    DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt,
-    RendererSelection,
+    DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, MediaPartSource,
+    RenderedPrompt, RendererSelection,
 };
 pub use request::{
     ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
-    ChatToolChoice, GenerationPromptMode, ReasoningEffort, ResolvedToolContext, SamplingParams,
+    ChatToolChoice, GenerationPromptMode, ResolvedToolContext, SamplingParams,
 };
 pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage};
 pub use vllm_engine_core_client::protocol::multimodal::MmFeatures;
@@ -52,6 +55,7 @@ mod event;
 pub mod multimodal;
 mod output;
 mod parser;
+mod reasoning;
 mod renderer;
 mod request;
 mod stream;
@@ -60,7 +64,7 @@ use vllm_engine_core_client::EngineCoreClient;
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
 use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
 use vllm_llm::Llm;
-use vllm_text::{Prompt, TextLlm, TextRequest};
+use vllm_text::{Prompt, TextLlm, TextRequest, TextRequestProcessor};
 
 /// Chat request preparation shared by inference and render-only frontends.
 pub struct ChatRequestProcessor {
@@ -68,21 +72,28 @@ pub struct ChatRequestProcessor {
     /// Effective model dtype reported by the engine.
     /// Absent for text-only frontends without an engine handshake.
     model_dtype: Option<ModelDtype>,
+    /// Effective runtime context length reported by the engine.
+    /// Absent for render-only frontends without an engine handshake.
+    max_model_len: Option<usize>,
     /// Tool-call parser selection used when preparing generation requests.
     tool_call_parser: ParserSelection,
     /// Reasoning parser selection used when preparing generation requests.
     reasoning_parser: ParserSelection,
+    /// Server-side floor for tool-call structural tags.
+    tool_strict_level: ToolStrictLevel,
 }
 
 impl ChatRequestProcessor {
     /// Create a processor with multimodal support using the effective model
     /// dtype reported by the engine.
-    fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self {
+    fn new(backend: DynChatBackend, model_dtype: ModelDtype, max_model_len: u32) -> Self {
         Self {
             backend,
             model_dtype: Some(model_dtype),
+            max_model_len: Some(max_model_len as usize),
             tool_call_parser: ParserSelection::Auto,
             reasoning_parser: ParserSelection::Auto,
+            tool_strict_level: ToolStrictLevel::Auto,
         }
     }
 
@@ -91,8 +102,10 @@ impl ChatRequestProcessor {
         Self {
             backend,
             model_dtype: None,
+            max_model_len: None,
             tool_call_parser: ParserSelection::Auto,
             reasoning_parser: ParserSelection::Auto,
+            tool_strict_level: ToolStrictLevel::Auto,
         }
     }
 
@@ -107,11 +120,21 @@ impl ChatRequestProcessor {
         self
     }
 
+    /// Configure the server-side floor for tool-call structural tags.
+    pub fn with_tool_strict_level(mut self, tool_strict_level: ToolStrictLevel) -> Self {
+        self.tool_strict_level = tool_strict_level;
+        self
+    }
+
     async fn finalize_rendered_prompt(
         &self,
         request: &ChatRequest,
         rendered: RenderedPrompt,
     ) -> Result<(Prompt, Option<MmFeatures>)> {
+        let media_is_empty = match &rendered.media_order {
+            Some(media_order) => media_order.is_empty(),
+            None => !request.has_multimodal(),
+        };
         match self.model_dtype {
             Some(model_dtype) => {
                 multimodal::finalize_rendered_prompt(
@@ -119,10 +142,11 @@ impl ChatRequestProcessor {
                     rendered,
                     self.backend.multimodal_model_info(),
                     model_dtype,
+                    self.max_model_len,
                 )
                 .await
             }
-            None if !request.has_multimodal() => Ok((rendered.prompt, None)),
+            None if media_is_empty => Ok((rendered.prompt, None)),
             None => Err(Error::UnsupportedMultimodalRenderer),
         }
     }
@@ -130,7 +154,7 @@ impl ChatRequestProcessor {
     /// Prepare media for an already-tokenized request.
     async fn prepare_media(
         &self,
-        media: Vec<MediaContentPart>,
+        media: multimodal::MultimodalInput,
         token_ids: &mut Vec<u32>,
     ) -> Result<Option<MmFeatures>> {
         if media.is_empty() {
@@ -140,23 +164,24 @@ impl ChatRequestProcessor {
             .backend
             .multimodal_model_info()
             .ok_or(Error::UnsupportedMultimodalRenderer)?;
-        let model_dtype = self.model_dtype.ok_or(Error::UnsupportedMultimodalRenderer)?;
-        let features = info.prepare_multimodal(media, token_ids, model_dtype).await?;
+        let features = match media {
+            multimodal::MultimodalInput::Raw(parts) => {
+                let model_dtype = self.model_dtype.ok_or(Error::UnsupportedMultimodalRenderer)?;
+                info.prepare_multimodal(parts, token_ids, model_dtype, self.max_model_len)
+                    .await?
+            }
+            multimodal::MultimodalInput::Preprocessed(features) => {
+                info.prepare_preprocessed(features, token_ids.len())?
+            }
+        };
         Ok(Some(features))
     }
 
     async fn prepare_text_request(&self, request: ChatRequest) -> Result<TextRequest> {
         // Stamp before rendering so render and tokenize count toward TTFT/e2e.
         let arrival_time = vllm_llm::current_unix_timestamp_secs();
-        let rendered = self.backend.chat_renderer().render(&request)?;
-        let reasoning_parser_kwargs =
-            request
-                .sampling_params
-                .structured_outputs
-                .is_some()
-                .then(|| ReasoningParserKwargs {
-                    chat_template_kwargs: rendered.effective_template_kwargs.clone(),
-                });
+        let mut rendered = self.backend.chat_renderer().render(&request)?;
+        let chat_template_kwargs = std::mem::take(&mut rendered.effective_template_kwargs);
         let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
         Ok(TextRequest {
             request_id: request.request_id,
@@ -171,7 +196,11 @@ impl ChatRequestProcessor {
             add_special_tokens: request.add_special_tokens,
             data_parallel_rank: request.data_parallel_rank,
             session_id: request.session_id,
-            reasoning_parser_kwargs,
+            kv_hints: None,
+            reasoning_parser_kwargs: ReasoningParserKwargs {
+                chat_template_kwargs,
+            },
+            reasoning_ended: None,
             lora_request: request.lora_request,
             arrival_time: Some(arrival_time),
         })
@@ -184,19 +213,37 @@ impl ChatRequestProcessor {
     }
 
     /// Prepare one chat request without submitting it to an engine.
+    ///
+    /// The returned text request is already tokenized: the output processor
+    /// is initialized from the final prompt token IDs and its grammar, if any,
+    /// is applied to the request before it is handed to `text_processor` again
+    /// for lowering.
     pub async fn prepare(
         &self,
         mut request: ChatRequest,
+        text_processor: &TextRequestProcessor,
     ) -> Result<(TextRequest, DynChatOutputProcessor)> {
         request.validate()?;
-        let output_processor = self.backend.new_chat_output_processor(
+
+        let mut output_processor = self.backend.new_chat_output_processor(
             &mut request,
             NewChatOutputProcessorOptions {
                 tool_call_parser: &self.tool_call_parser,
                 reasoning_parser: &self.reasoning_parser,
+                tool_strict_level: self.tool_strict_level,
             },
         )?;
-        let text_request = self.prepare_text_request(request).await?;
+        let mut text_request = self.prepare_text_request(request).await?;
+
+        // Initialize the output processor from the final prompt token IDs and apply
+        // the output grammar to the request, if any.
+        let prompt_token_ids = text_processor.tokenize_in_place(&mut text_request)?;
+        let grammar = {
+            output_processor.initialize(prompt_token_ids)?;
+            output_processor.build_output_grammar()?
+        };
+        apply_output_grammar(&mut text_request, grammar)?;
+
         Ok((text_request, output_processor))
     }
 }
@@ -216,10 +263,11 @@ impl ChatLlm {
     /// backend.
     pub fn new(text: TextLlm, backend: DynChatBackend) -> Self {
         let model_dtype = text.engine_core_client().model_dtype();
+        let max_model_len = text.engine_core_client().max_model_len();
 
         Self {
             text,
-            processor: ChatRequestProcessor::new(backend, model_dtype),
+            processor: ChatRequestProcessor::new(backend, model_dtype, max_model_len),
         }
     }
 
@@ -242,6 +290,12 @@ impl ChatLlm {
         self
     }
 
+    /// Set the server-side floor for tool-call structural tags.
+    pub fn with_tool_strict_level(mut self, tool_strict_level: ToolStrictLevel) -> Self {
+        self.processor.tool_strict_level = tool_strict_level;
+        self
+    }
+
     /// Override the effective model dtype used for multimodal tensor encoding.
     pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
         self.processor.model_dtype = Some(model_dtype);
@@ -286,12 +340,14 @@ impl ChatLlm {
     }
 
     /// Prepare media for an already-tokenized request.
+    /// Raw content is preprocessed; inline features are checked against model
+    /// capabilities, modality limits, and the final prompt positions.
     pub async fn prepare_media(
         &self,
-        media: Vec<MediaContentPart>,
+        media: impl Into<multimodal::MultimodalInput>,
         token_ids: &mut Vec<u32>,
     ) -> Result<Option<MmFeatures>> {
-        self.processor.prepare_media(media, token_ids).await
+        self.processor.prepare_media(media.into(), token_ids).await
     }
 
     /// Effective tool-call parser name for this model, if parsing is enabled.
@@ -306,7 +362,8 @@ impl ChatLlm {
 
     /// Render, tokenize, and submit one chat request.
     pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
-        let (text_request, output_processor) = self.processor.prepare(request).await?;
+        let (text_request, output_processor) =
+            self.processor.prepare(request, self.text.request_processor()).await?;
         let request_id = text_request.request_id.clone();
         let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
 
@@ -363,7 +420,7 @@ mod tests {
         )
         .unwrap_err();
 
-        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
     }
 
     #[test]
@@ -374,6 +431,6 @@ mod tests {
         )
         .unwrap_err();
 
-        expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
     }
 }
diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs
index a0743b8769b6..4138d39bdab3 100644
--- a/rust/src/chat/src/multimodal.rs
+++ b/rust/src/chat/src/multimodal.rs
@@ -16,41 +16,50 @@ use std::fs;
 use std::path::Path;
 use std::sync::{Arc, LazyLock};
 
-use itertools::izip;
+use itertools::{Either, izip};
+use llm_multimodal::registry::ModelRegistryError;
 use llm_multimodal::{
     AsyncMultiModalTracker, AudioClip, AudioPreProcessor, EncoderFieldLayouts, FieldLayout,
     ImageFrame, MediaConnector, MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata,
     ModelProcessorSpec, ModelRegistry, PreProcessorConfig, PreprocessedEncoderInputs,
     PromptReplacement, Tokenizer as TokenResolver, TrackedMedia, VideoClip, VisionPreProcessor,
-    VisionProcessorRegistry,
+    VisionPreprocessingContext,
 };
 use serde::{Deserialize, Serialize};
 use thiserror_ext::AsReport as _;
-use tracing::warn;
+use tracing::{Instrument as _, warn};
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
 use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, MmFeatures, MmKwargsItem};
 use vllm_text::Prompt;
 use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
 
 use crate::error::{Error, Result, bail_multimodal, multimodal};
-use crate::renderer::RenderedPrompt;
+use crate::renderer::{MediaPartSource, RenderedPrompt};
 use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
 
 mod audio;
 mod expand;
 mod image;
+mod input;
 mod item;
+mod preprocessed;
 mod tensor;
+mod timing;
 mod video;
 
 use self::expand::expand_prompt_token_ids;
+pub use self::input::MultimodalInput;
+use self::timing::MM_STAGE_TARGET;
+pub use self::timing::{mm_request_span, mm_timing_layer};
 
 /// Resolved multimodal support for one loaded model.
 #[derive(Clone)]
 pub struct MultimodalModelInfo {
     context: MultimodalModelContext,
-    image: Option<ModalitySupport>,
-    video: Option<ModalitySupport>,
+    /// Rendered placeholder marker IDs for all resolved modalities.
+    placeholder_token_ids: Vec<u32>,
+    image: Option<VisionModalitySupport>,
+    video: Option<VisionModalitySupport>,
     audio: Option<AudioModalitySupport>,
     media_connector: Arc<MediaConnector>,
     /// Maximum number of input items allowed per prompt for each modality.
@@ -127,13 +136,18 @@ impl MultimodalModelContext {
         REGISTRY.lookup(&self.metadata())
     }
 
-    /// Resolve a static vision preprocessor for one loaded model.
-    ///
-    /// The vision preprocessor serves both the image and video modalities.
-    fn resolve_vision_processor(&self) -> Option<&'static dyn VisionPreProcessor> {
-        static REGISTRY: LazyLock<VisionProcessorRegistry> =
-            LazyLock::new(VisionProcessorRegistry::with_defaults);
-        REGISTRY.find(&self.model_id, self.model_type.as_deref())
+    /// Build a vision preprocessor for one loaded model and modality.
+    fn resolve_vision_processor(
+        &self,
+        model_spec: &'static dyn ModelProcessorSpec,
+        preprocessor_config: &PreProcessorConfig,
+        modality: Modality,
+    ) -> Result<Option<Arc<dyn VisionPreProcessor>>> {
+        match model_spec.vision_processor(&self.metadata(), preprocessor_config, modality) {
+            Ok(processor) => Ok(Some(Arc::from(processor))),
+            Err(ModelRegistryError::UnsupportedModality { .. }) => Ok(None),
+            Err(error) => Err(error.into()),
+        }
     }
 
     /// Resolve an audio preprocessor for one loaded model.
@@ -142,7 +156,7 @@ impl MultimodalModelContext {
         model_spec: &'static dyn ModelProcessorSpec,
         preprocessor_config: &PreProcessorConfig,
     ) -> Option<Arc<dyn AudioPreProcessor>> {
-        model_spec.audio_processor(&self.config, preprocessor_config).map(Arc::from)
+        model_spec.audio_processor(&self.metadata(), preprocessor_config).map(Arc::from)
     }
 }
 
@@ -236,17 +250,16 @@ impl ResolvedPlaceholder {
     }
 }
 
-/// Static per-modality vision preprocessor plus its loaded config, resolved
-/// placeholder tokens, and the model's shared tensor-layout spec.
+/// Model-owned vision preprocessor plus resolved placeholder tokens and
+/// the model's shared tensor-layout spec.
 #[derive(Clone)]
-struct ModalitySupport {
+struct VisionModalitySupport {
     spec: ResolvedMultimodalSpec,
     placeholder: ResolvedPlaceholder,
-    processor: &'static dyn VisionPreProcessor,
-    config: PreProcessorConfig,
+    processor: Arc<dyn VisionPreProcessor>,
 }
 
-/// Static audio preprocessor plus the resolved model contract for its output.
+/// Model-owned audio preprocessor plus the resolved model contract for its output.
 #[derive(Clone)]
 struct AudioModalitySupport {
     spec: ResolvedMultimodalSpec,
@@ -393,7 +406,7 @@ impl MultimodalModelInfo {
             &context,
             preprocessor_config.clone(),
             video_preprocessor_config,
-        );
+        )?;
         let audio = Self::resolve_audio_lane(&context, &preprocessor_config)?;
 
         if image.is_none() && video.is_none() && audio.is_none() {
@@ -410,8 +423,18 @@ impl MultimodalModelInfo {
             MediaConnectorConfig::default(),
         )?);
 
+        let placeholder_token_ids = [
+            image.as_ref().map(|support| support.placeholder.marker_token_id),
+            video.as_ref().map(|support| support.placeholder.marker_token_id),
+            audio.as_ref().map(|support| support.placeholder.marker_token_id),
+        ]
+        .into_iter()
+        .flatten()
+        .collect();
+
         Ok(Some(Self {
             context,
+            placeholder_token_ids,
             image,
             video,
             audio,
@@ -424,27 +447,19 @@ impl MultimodalModelInfo {
     ///
     /// A missing model spec or vision processor disables both lanes but is
     /// not an error: unsupported model families are logged and skipped.
+    /// Processor construction errors propagate to model loading.
     fn resolve_vision_lanes(
         context: &MultimodalModelContext,
         preprocessor_config: PreProcessorConfig,
         video_preprocessor_config: PreProcessorConfig,
-    ) -> (Option<ModalitySupport>, Option<ModalitySupport>) {
+    ) -> Result<(Option<VisionModalitySupport>, Option<VisionModalitySupport>)> {
         let Some(raw_spec) = context.resolve_model_spec() else {
             warn!(
                 model_id = context.model_id,
                 model_type = context.model_type,
                 "multimodal model spec is not registered; disabling image/video support for this model"
             );
-            return (None, None);
-        };
-
-        let Some(processor) = context.resolve_vision_processor() else {
-            warn!(
-                model_id = context.model_id,
-                model_type = context.model_type,
-                "vision processor is not registered; disabling image/video support for this model"
-            );
-            return (None, None);
+            return Ok((None, None));
         };
 
         // Warn and disable the modality if the placeholder resolution fails.
@@ -462,14 +477,19 @@ impl MultimodalModelInfo {
                 }
             };
 
-        let image = resolve_placeholder(Modality::Image).map(|placeholder| ModalitySupport {
-            spec: ResolvedMultimodalSpec::new(raw_spec, Modality::Image),
-            placeholder,
-            processor,
-            config: preprocessor_config.clone(),
-        });
+        let image = if let Some(placeholder) = resolve_placeholder(Modality::Image) {
+            context
+                .resolve_vision_processor(raw_spec, &preprocessor_config, Modality::Image)?
+                .map(|processor| VisionModalitySupport {
+                    spec: ResolvedMultimodalSpec::new(raw_spec, Modality::Image),
+                    placeholder,
+                    processor,
+                })
+        } else {
+            None
+        };
 
-        let video = resolve_placeholder(Modality::Video).and_then(|placeholder| {
+        let video = if let Some(placeholder) = resolve_placeholder(Modality::Video) {
             // Placeholder expansion attributes markers to modalities by token
             // ID, so a marker shared with the image modality is ambiguous.
             let image_marker = image.as_ref().map(|image| image.placeholder.marker_token_id);
@@ -481,16 +501,23 @@ impl MultimodalModelInfo {
                 );
                 None
             } else {
-                Some(ModalitySupport {
-                    spec: ResolvedMultimodalSpec::new(raw_spec, Modality::Video),
-                    placeholder,
-                    processor,
-                    config: video_preprocessor_config,
-                })
+                context
+                    .resolve_vision_processor(
+                        raw_spec,
+                        &video_preprocessor_config,
+                        Modality::Video,
+                    )?
+                    .map(|processor| VisionModalitySupport {
+                        spec: ResolvedMultimodalSpec::new(raw_spec, Modality::Video),
+                        placeholder,
+                        processor,
+                    })
             }
-        });
+        } else {
+            None
+        };
 
-        (image, video)
+        Ok((image, video))
     }
 
     /// Resolve a spec-backed audio lane when the model advertises audio support.
@@ -533,6 +560,20 @@ impl MultimodalModelInfo {
             Modality::ImageEmbeds => None,
         }
     }
+
+    fn vision_preprocessing_context(
+        &self,
+        prompt_token_ids: &[u32],
+        max_model_len: Option<usize>,
+    ) -> VisionPreprocessingContext {
+        let text_prompt_length = prompt_token_ids
+            .iter()
+            .filter(|token_id| !self.placeholder_token_ids.contains(token_id))
+            .count();
+        VisionPreprocessingContext {
+            token_budget: max_model_len.map(|limit| limit.saturating_sub(text_prompt_length)),
+        }
+    }
 }
 
 /// Finalize a rendered chat prompt into text-generation input.
@@ -545,8 +586,10 @@ pub(crate) async fn finalize_rendered_prompt(
     rendered: RenderedPrompt,
     info: Option<&MultimodalModelInfo>,
     model_dtype: ModelDtype,
+    max_model_len: Option<usize>,
 ) -> Result<(Prompt, Option<MmFeatures>)> {
-    if !request.has_multimodal() {
+    let media_parts = extract_media_parts(request, rendered.media_order.as_deref())?;
+    if media_parts.is_empty() {
         return Ok((rendered.prompt, None));
     }
     let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?;
@@ -558,63 +601,102 @@ pub(crate) async fn finalize_rendered_prompt(
             .map_err(|error| multimodal!("{error}"))?,
         Prompt::TokenIds(token_ids) => token_ids,
     };
-    let media_parts = extract_media_parts(request)?;
-    let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?;
+    let prepared = info
+        .prepare_multimodal(
+            media_parts,
+            &mut prompt_token_ids,
+            model_dtype,
+            max_model_len,
+        )
+        .instrument(mm_request_span(&request.request_id))
+        .await?;
 
     Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared)))
 }
 
-/// Extract media parts from chat messages in message/content order.
-///
-/// Assistant history is skipped because generated assistant blocks are already
-/// represented as text for prompt rendering in this crate.
-fn extract_media_parts(request: &ChatRequest) -> Result<Vec<MediaContentPart>> {
-    let mut all_parts = Vec::new();
-    for message in &request.messages {
-        let content = match message {
-            ChatMessage::System { content }
-            | ChatMessage::Developer { content, .. }
-            | ChatMessage::User { content }
-            | ChatMessage::ToolResponse { content, .. } => content,
-            ChatMessage::Assistant { .. } => continue,
-        };
-        let ChatContent::Parts(parts) = content else {
-            continue;
-        };
-        for part in parts {
-            match part {
-                ChatContentPart::Text { .. } => {}
-                ChatContentPart::ImageUrl {
-                    image_url,
-                    detail,
-                    uuid,
-                } => all_parts.push(MediaContentPart::ImageUrl {
-                    url: image_url.clone(),
-                    detail: *detail,
-                    uuid: uuid.clone(),
-                }),
-                ChatContentPart::VideoUrl { video_url, uuid } => {
-                    all_parts.push(MediaContentPart::VideoUrl {
-                        url: video_url.clone(),
-                        uuid: uuid.clone(),
-                    })
-                }
-                ChatContentPart::InputAudio { data, format, uuid } => {
-                    all_parts.push(MediaContentPart::AudioUrl {
-                        url: input_audio_data_url(data, format.as_deref())?,
-                        uuid: uuid.clone(),
-                    })
-                }
-                ChatContentPart::AudioUrl { audio_url, uuid } => {
-                    all_parts.push(MediaContentPart::AudioUrl {
-                        url: audio_url.clone(),
-                        uuid: uuid.clone(),
-                    })
+/// Resolve media parts in the placeholder order reported by the renderer.
+fn extract_media_parts(
+    request: &ChatRequest,
+    media_order: Option<&[MediaPartSource]>,
+) -> Result<Vec<MediaContentPart>> {
+    let parts = match media_order {
+        // Resolve exactly the sources selected by the renderer, in placeholder order.
+        Some(media_order) => Either::Left(media_order.iter().map(|source| -> Result<_> {
+            let message = request.messages.get(source.message_index).ok_or_else(|| {
+                multimodal!(
+                    "renderer reported missing chat message {} for multimodal content",
+                    source.message_index
+                )
+            })?;
+            let content = match message {
+                ChatMessage::System { content }
+                | ChatMessage::Developer { content, .. }
+                | ChatMessage::User { content }
+                | ChatMessage::ToolResponse { content, .. }
+                | ChatMessage::Custom { content, .. } => content,
+                ChatMessage::Assistant { .. } => {
+                    bail_multimodal!("renderer reported multimodal assistant content")
                 }
+            };
+            let ChatContent::Parts(parts) = content else {
+                bail_multimodal!("renderer reported multimodal content in a text-only chat message")
+            };
+            let part = parts.get(source.content_part_index).ok_or_else(|| {
+                multimodal!(
+                    "renderer reported missing content part {} in chat message {}",
+                    source.content_part_index,
+                    source.message_index
+                )
+            })?;
+            Ok(part)
+        })),
+        // Fall back to all media in request message and content-part order (e.g. Jinja).
+        None => Either::Right(request.messages.iter().flat_map(|message| {
+            let content = match message {
+                ChatMessage::System { content }
+                | ChatMessage::Developer { content, .. }
+                | ChatMessage::User { content }
+                | ChatMessage::ToolResponse { content, .. }
+                | ChatMessage::Custom { content, .. } => Some(content),
+                ChatMessage::Assistant { .. } => None,
+            };
+            match content {
+                Some(ChatContent::Parts(parts)) => parts.as_slice(),
+                _ => &[],
             }
-        }
-    }
-    Ok(all_parts)
+            .iter()
+            .filter(|part| part.is_multimodal())
+            .map(Ok)
+        })),
+    };
+    parts
+        .map(|part| match part? {
+            ChatContentPart::Text { .. } => {
+                bail_multimodal!("renderer reported a text part as multimodal content")
+            }
+            ChatContentPart::ImageUrl {
+                image_url,
+                detail,
+                uuid,
+            } => Ok(MediaContentPart::ImageUrl {
+                url: image_url.clone(),
+                detail: *detail,
+                uuid: uuid.clone(),
+            }),
+            ChatContentPart::VideoUrl { video_url, uuid } => Ok(MediaContentPart::VideoUrl {
+                url: video_url.clone(),
+                uuid: uuid.clone(),
+            }),
+            ChatContentPart::InputAudio { data, format, uuid } => Ok(MediaContentPart::AudioUrl {
+                url: input_audio_data_url(data, format.as_deref())?,
+                uuid: uuid.clone(),
+            }),
+            ChatContentPart::AudioUrl { audio_url, uuid } => Ok(MediaContentPart::AudioUrl {
+                url: audio_url.clone(),
+                uuid: uuid.clone(),
+            }),
+        })
+        .collect()
 }
 
 /// Wrap OpenAI base64 audio in a data URL consumed by `MediaConnector`.
@@ -654,11 +736,16 @@ impl MultimodalModelInfo {
     ///
     /// Modalities without a configured count are unlimited.
     fn validate_mm_limits(&self, media_parts: &[MediaContentPart]) -> Result<()> {
+        self.validate_modality_limits(media_parts.iter().filter_map(media_part_limit_modality))
+    }
+
+    fn validate_modality_limits(
+        &self,
+        modalities: impl IntoIterator<Item = MmModality>,
+    ) -> Result<()> {
         let mut counts: HashMap<MmModality, usize> = HashMap::new();
-        for part in media_parts {
-            if let Some(modality) = media_part_limit_modality(part) {
-                *counts.entry(modality).or_default() += 1;
-            }
+        for modality in modalities {
+            *counts.entry(modality).or_default() += 1;
         }
 
         for (modality, count) in counts {
@@ -677,6 +764,32 @@ impl MultimodalModelInfo {
         Ok(())
     }
 
+    /// Validate inline storage, batching, and placeholder ranges, then check
+    /// this model's supported modalities and item-count limits.
+    ///
+    /// `prompt_len` must include all expanded multimodal placeholders.
+    pub(crate) fn prepare_preprocessed(
+        &self,
+        mut features: MmFeatures,
+        prompt_len: usize,
+    ) -> Result<MmFeatures> {
+        preprocessed::validate_features(&mut features, prompt_len)?;
+        for feature in &features {
+            let supported = match feature.modality {
+                MmModality::Image => self.image.is_some(),
+                MmModality::Video => self.video.is_some(),
+                MmModality::Audio => self.audio.is_some(),
+            };
+            if !supported {
+                return Err(Error::UnsupportedModality {
+                    modality: feature.modality.as_str().to_owned(),
+                });
+            }
+        }
+        self.validate_modality_limits(features.iter().map(|feature| feature.modality))?;
+        Ok(features)
+    }
+
     /// Run media fetch, per-modality preprocessing, prompt expansion, and
     /// feature build.
     ///
@@ -688,18 +801,46 @@ impl MultimodalModelInfo {
         media_parts: Vec<MediaContentPart>,
         prompt_token_ids: &mut Vec<u32>,
         model_dtype: ModelDtype,
+        max_model_len: Option<usize>,
     ) -> Result<MmFeatures> {
-        let media_parts_len = media_parts.len();
-        if media_parts_len == 0 {
+        if media_parts.is_empty() {
             return Ok(Vec::new());
         }
+        self.prepare_multimodal_timed(media_parts, prompt_token_ids, model_dtype, max_model_len)
+            .await
+    }
+
+    /// Timed body of [`Self::prepare_multimodal`]; each stage is a `tracing`
+    /// span aggregated per request by the `vllm-tracing` timing layer.
+    #[tracing::instrument(
+        name = "mm_stage",
+        target = MM_STAGE_TARGET,
+        skip_all,
+        fields(stage = "preprocessor_total")
+    )]
+    async fn prepare_multimodal_timed(
+        &self,
+        media_parts: Vec<MediaContentPart>,
+        prompt_token_ids: &mut Vec<u32>,
+        model_dtype: ModelDtype,
+        max_model_len: Option<usize>,
+    ) -> Result<MmFeatures> {
+        let media_parts_len = media_parts.len();
         self.validate_mm_limits(&media_parts)?;
+        let vision_context = self.vision_preprocessing_context(prompt_token_ids, max_model_len);
         let fetched = self.fetch_media(media_parts).await?;
 
         let mut prepared = Vec::new();
         if !fetched.images.is_empty() {
-            prepared
-                .push(self.prepare_images(fetched.images, fetched.image_uuids, model_dtype).await?);
+            prepared.push(
+                self.prepare_images(
+                    fetched.images,
+                    fetched.image_uuids,
+                    model_dtype,
+                    vision_context,
+                )
+                .await?,
+            );
         }
         if !fetched.videos.is_empty() {
             prepared
@@ -752,6 +893,12 @@ impl MultimodalModelInfo {
 
     /// Fetch all connector-backed media parts and split them per modality,
     /// preserving their request-order UUID metadata.
+    #[tracing::instrument(
+        name = "mm_stage",
+        target = MM_STAGE_TARGET,
+        skip_all,
+        fields(stage = "media_fetch")
+    )]
     async fn fetch_media(&self, media_parts: Vec<MediaContentPart>) -> Result<FetchedMedia> {
         let mut tracker = AsyncMultiModalTracker::new(Arc::clone(&self.media_connector));
         for part in media_parts {
@@ -941,6 +1088,100 @@ mod tests {
         );
     }
 
+    #[test]
+    fn vision_preprocessing_context_uses_runtime_budget_and_text_tokens() {
+        let info = qwen3_vl_info();
+        let context = info.vision_preprocessing_context(
+            &[11, QWEN3_IMAGE_PAD_ID, 12, QWEN3_VIDEO_PAD_ID, 13],
+            Some(4096),
+        );
+
+        assert_eq!(context.token_budget, Some(4096 - 3));
+        assert_eq!(
+            info.vision_preprocessing_context(&[11, 12, 13], Some(2)).token_budget,
+            Some(0)
+        );
+        assert_eq!(
+            info.vision_preprocessing_context(&[11], None).token_budget,
+            None
+        );
+    }
+
+    #[tokio::test]
+    async fn nemotron_budget_controls_preprocessing_and_engine_features() {
+        use vllm_engine_core_client::protocol::multimodal::MmKwargValue;
+
+        let tokenizer = TestTokenizer::new()
+            .with_regular_token("<image>", 1018)
+            .with_regular_token("<img>", 1019)
+            .with_regular_token("</img>", 1020);
+        let config = serde_json::json!({
+            "model_type": "nemotron_h_omni",
+            "img_context_token_id": 1018,
+            "max_position_embeddings": 512,
+            "patch_size": 16,
+            "downsample_ratio": 0.5,
+            "norm_mean": [0.0, 0.0, 0.0],
+            "norm_std": [1.0, 1.0, 1.0],
+            "vision_config": {
+                "args": {"min_num_patches": 1024, "max_num_patches": 13312}
+            }
+        });
+        let info = test_info("nemotron_h_omni", config, tokenizer);
+        let media = MediaContentPart::ImageUrl {
+            url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=".to_string(),
+            detail: None,
+            uuid: None,
+        };
+        let mut prompt = vec![11; 100];
+        prompt.push(1018);
+
+        let mut tokens = prompt.clone();
+        let features = info
+            .prepare_multimodal(
+                vec![media.clone()],
+                &mut tokens,
+                ModelDtype::Float32,
+                Some(512),
+            )
+            .await
+            .unwrap();
+        assert_eq!(features.len(), 1);
+        let feature = &features[0];
+        assert_eq!(
+            (feature.mm_position.offset, feature.mm_position.length),
+            (100, 258)
+        );
+        assert_eq!(tokens.len(), 358);
+        assert_eq!(tokens[100], 1019);
+        assert!(tokens[101..357].iter().all(|&token| token == 1018));
+        assert_eq!(tokens[357], 1020);
+        let data = feature.data.as_ref().unwrap();
+        assert!(matches!(
+            data["pixel_values_flat"].data.as_ref(),
+            Some(MmKwargValue::Tensor(tensor)) if tensor.shape == [3, 512, 512]
+        ));
+        assert_eq!(
+            data["num_tokens_per_image"].data,
+            Some(MmKwargValue::Int(256))
+        );
+        assert_eq!(
+            data["imgs_sizes"].data,
+            Some(MmKwargValue::List(vec![
+                MmKwargValue::Int(512),
+                MmKwargValue::Int(512),
+            ]))
+        );
+
+        let error = info
+            .prepare_multimodal(vec![media], &mut prompt, ModelDtype::Float32, Some(100))
+            .await
+            .unwrap_err();
+        assert!(
+            matches!(error, Error::Multimodal(message) if message.contains("exceeding token_budget"))
+        );
+    }
+
     #[test]
     fn qwen3_vl_without_video_token_id_disables_video_support_only() {
         let config = serde_json::json!({
@@ -1009,6 +1250,75 @@ mod tests {
         assert!(input_audio_data_url("AAAA", Some("flac")).is_err());
     }
 
+    #[test]
+    fn extract_media_parts_uses_request_order_or_explicit_selection() {
+        let request = ChatRequest {
+            messages: vec![
+                ChatMessage::user(vec![
+                    ChatContentPart::text("before image"),
+                    ChatContentPart::ImageUrl {
+                        image_url: "data:image/png;base64,image-b".to_string(),
+                        detail: None,
+                        uuid: Some("image-b".to_string()),
+                    },
+                ]),
+                ChatMessage::user(vec![ChatContentPart::ImageUrl {
+                    image_url: "data:image/png;base64,image-a".to_string(),
+                    detail: None,
+                    uuid: Some("image-a".to_string()),
+                }]),
+                ChatMessage::assistant_text("answer"),
+                ChatMessage::user("follow-up"),
+            ],
+            ..ChatRequest::for_test()
+        };
+        let media_order = [
+            MediaPartSource {
+                message_index: 1,
+                content_part_index: 0,
+            },
+            MediaPartSource {
+                message_index: 0,
+                content_part_index: 1,
+            },
+        ];
+
+        let selections = [None, Some(&media_order[..]), Some(&[][..])];
+        let uuids = selections.map(|order| {
+            extract_media_parts(&request, order)
+                .unwrap()
+                .into_iter()
+                .map(|part| match part {
+                    MediaContentPart::ImageUrl { uuid, .. } => uuid,
+                    _ => panic!("expected image URL"),
+                })
+                .collect::<Vec<_>>()
+        });
+
+        expect_test::expect![[r#"
+            [
+                [
+                    Some(
+                        "image-b",
+                    ),
+                    Some(
+                        "image-a",
+                    ),
+                ],
+                [
+                    Some(
+                        "image-a",
+                    ),
+                    Some(
+                        "image-b",
+                    ),
+                ],
+                [],
+            ]
+        "#]]
+        .assert_debug_eq(&uuids);
+    }
+
     fn image_url_part() -> MediaContentPart {
         MediaContentPart::ImageUrl {
             url: "https://example.com/image.png".to_string(),
diff --git a/rust/src/chat/src/multimodal/audio.rs b/rust/src/chat/src/multimodal/audio.rs
index 897f7aa4927c..6b2fecf78dc1 100644
--- a/rust/src/chat/src/multimodal/audio.rs
+++ b/rust/src/chat/src/multimodal/audio.rs
@@ -8,6 +8,7 @@ use std::sync::Arc;
 use llm_multimodal::{AudioClip, Modality, PreprocessedEncoderInputs};
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
 
+use super::timing::MM_STAGE_TARGET;
 use super::{AudioModalitySupport, MultimodalModelInfo, PreparedMedia, item};
 use crate::error::{Error, Result, bail_multimodal, multimodal};
 
@@ -16,6 +17,12 @@ pub(super) const AUDIO_PRIMARY_KEY: &str = "input_audio_features";
 
 impl MultimodalModelInfo {
     /// Preprocess fetched audio clips as one batch and build per-item features.
+    #[tracing::instrument(
+        name = "mm_stage",
+        target = MM_STAGE_TARGET,
+        skip_all,
+        fields(stage = "preprocess_audio")
+    )]
     pub(super) async fn prepare_audios(
         &self,
         clips: Vec<Arc<AudioClip>>,
@@ -236,6 +243,45 @@ mod tests {
         ));
     }
 
+    #[tokio::test]
+    async fn audio_pipeline_attributes_stage_timings_to_request() {
+        use tracing::Instrument as _;
+        use tracing::instrument::WithSubscriber as _;
+        use tracing_subscriber::layer::SubscriberExt as _;
+
+        use crate::{mm_request_span, mm_timing_layer};
+
+        let info = qwen3_asr_info();
+        let (layer, timings) = mm_timing_layer();
+        let parts = vec![MediaContentPart::AudioData {
+            data: wav_i16_mono(16_000, &[0; 1_600]),
+            mime_type: Some("audio/wav".to_string()),
+            uuid: Some("audio-1".to_string()),
+        }];
+        async {
+            info.prepare_multimodal(parts, &mut vec![AUDIO_PAD_ID], ModelDtype::Float32, None)
+                .instrument(mm_request_span("request-1"))
+                .await
+                .unwrap();
+        }
+        .with_subscriber(tracing_subscriber::registry().with(layer))
+        .await;
+
+        let records = timings.stat();
+        assert_eq!(records.len(), 1);
+        let mut stages = records["request-1"].keys().collect::<Vec<_>>();
+        stages.sort();
+        expect_test::expect![[r#"
+            [
+                "media_fetch_secs",
+                "preprocess_audio_secs",
+                "preprocessor_total_secs",
+                "prompt_expansion_secs",
+            ]
+        "#]]
+        .assert_debug_eq(&stages);
+    }
+
     #[tokio::test]
     async fn tracker_processor_and_lowering_preserve_audio_contract() {
         let info = qwen3_asr_info();
diff --git a/rust/src/chat/src/multimodal/expand.rs b/rust/src/chat/src/multimodal/expand.rs
index 1172d0ccc085..e23c16b45543 100644
--- a/rust/src/chat/src/multimodal/expand.rs
+++ b/rust/src/chat/src/multimodal/expand.rs
@@ -10,6 +10,7 @@ use vllm_engine_core_client::protocol::multimodal::PlaceholderRange;
 use vllm_engine_core_client::protocol::tensor::WireTensor;
 
 use super::PreparedMedia;
+use super::timing::MM_STAGE_TARGET;
 use crate::error::{Error, Result, bail_multimodal};
 
 /// One modality's queue of pending placeholder replacements for prompt
@@ -55,6 +56,12 @@ impl<'a> ExpansionLane<'a> {
 ///
 /// The returned ranges point into the already-expanded prompt, grouped per
 /// modality in item order.
+#[tracing::instrument(
+    name = "mm_stage",
+    target = MM_STAGE_TARGET,
+    skip_all,
+    fields(stage = "prompt_expansion")
+)]
 pub(super) fn expand_prompt_token_ids(
     prompt_token_ids: &mut Vec<u32>,
     prepared: &[PreparedMedia],
diff --git a/rust/src/chat/src/multimodal/image.rs b/rust/src/chat/src/multimodal/image.rs
index 5063ae8d2fa7..818bb72f1af1 100644
--- a/rust/src/chat/src/multimodal/image.rs
+++ b/rust/src/chat/src/multimodal/image.rs
@@ -6,10 +6,11 @@
 
 use std::sync::Arc;
 
-use llm_multimodal::{ImageFrame, Modality, PreprocessedEncoderInputs};
+use llm_multimodal::{ImageFrame, Modality, PreprocessedEncoderInputs, VisionPreprocessingContext};
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
 
-use super::{ModalitySupport, MultimodalModelInfo, PreparedMedia, item};
+use super::timing::MM_STAGE_TARGET;
+use super::{MultimodalModelInfo, PreparedMedia, VisionModalitySupport, item};
 use crate::error::{Error, Result, bail_multimodal, multimodal};
 
 /// Forward-kwargs name of the primary image encoder input.
@@ -18,16 +19,23 @@ pub(super) const IMAGE_PRIMARY_KEY: &str = "pixel_values";
 impl MultimodalModelInfo {
     /// Preprocess all fetched image frames as one batch and build per-item
     /// features.
+    #[tracing::instrument(
+        name = "mm_stage",
+        target = MM_STAGE_TARGET,
+        skip_all,
+        fields(stage = "preprocess_image")
+    )]
     pub(super) async fn prepare_images(
         &self,
         frames: Vec<Arc<ImageFrame>>,
         uuids: Vec<Option<String>>,
         model_dtype: ModelDtype,
+        preprocessing_context: VisionPreprocessingContext,
     ) -> Result<PreparedMedia> {
         let support = self.image.as_ref().ok_or_else(|| Error::UnsupportedModality {
             modality: Modality::Image.to_string(),
         })?;
-        let preprocessed = self.preprocess_images(support, &frames).await?;
+        let preprocessed = self.preprocess_images(support, &frames, preprocessing_context).await?;
         let replacements = support.spec.prompt_replacements_for(&self.context, &preprocessed)?;
         if replacements.len() != frames.len() {
             bail_multimodal!(
@@ -56,16 +64,18 @@ impl MultimodalModelInfo {
     /// conversion.
     async fn preprocess_images(
         &self,
-        support: &ModalitySupport,
+        support: &VisionModalitySupport,
         image_frames: &[Arc<ImageFrame>],
+        preprocessing_context: VisionPreprocessingContext,
     ) -> Result<PreprocessedEncoderInputs> {
-        let config = support.config.clone();
-        let processor = support.processor;
+        let processor = Arc::clone(&support.processor);
         let images = image_frames.iter().map(|frame| frame.data().clone()).collect::<Vec<_>>();
 
         // TODO: is it still necessary given that we've already in a dedicated runtime?
-        tokio::task::spawn_blocking(move || Ok(processor.preprocess(&images, &config)?))
-            .await
-            .map_err(|error| multimodal!("image preprocessing task failed: {error}"))?
+        tokio::task::spawn_blocking(move || {
+            Ok(processor.preprocess_with_context(&images, &preprocessing_context)?)
+        })
+        .await
+        .map_err(|error| multimodal!("image preprocessing task failed: {error}"))?
     }
 }
diff --git a/rust/src/chat/src/multimodal/input.rs b/rust/src/chat/src/multimodal/input.rs
new file mode 100644
index 000000000000..53905f0a2654
--- /dev/null
+++ b/rust/src/chat/src/multimodal/input.rs
@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use llm_multimodal::MediaContentPart;
+use vllm_engine_core_client::protocol::multimodal::MmFeatures;
+
+/// Media preparation stage supplied by a frontend handler.
+#[derive(Debug)]
+pub enum MultimodalInput {
+    /// Raw media to fetch and preprocess, expanding prompt placeholders as needed.
+    Raw(Vec<MediaContentPart>),
+    /// Preprocessed features to validate against storage, batching, model, and prompt constraints.
+    Preprocessed(MmFeatures),
+}
+
+impl MultimodalInput {
+    pub fn is_empty(&self) -> bool {
+        match self {
+            Self::Raw(parts) => parts.is_empty(),
+            Self::Preprocessed(features) => features.is_empty(),
+        }
+    }
+}
+
+impl From<Vec<MediaContentPart>> for MultimodalInput {
+    fn from(parts: Vec<MediaContentPart>) -> Self {
+        Self::Raw(parts)
+    }
+}
diff --git a/rust/src/chat/src/multimodal/preprocessed.rs b/rust/src/chat/src/multimodal/preprocessed.rs
new file mode 100644
index 000000000000..a2ff4687b134
--- /dev/null
+++ b/rust/src/chat/src/multimodal/preprocessed.rs
@@ -0,0 +1,274 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::collections::HashMap;
+
+use vllm_engine_core_client::protocol::dtype::TensorDtype;
+use vllm_engine_core_client::protocol::multimodal::{
+    MmFeatureSpec, MmField, MmKwargValue, MmKwargsItem,
+};
+
+use crate::error::{
+    Error, Result, bail_invalid_preprocessed_multimodal, invalid_preprocessed_multimodal,
+};
+
+const MAX_DEPTH: usize = 32;
+
+/// Validate inline storage, placeholder ranges, and batching compatibility,
+/// then stably sort features by their prompt offsets.
+/// Encoder cache identifiers are supplied by the producer, which must bind
+/// them to the feature content.
+pub(super) fn validate_features(features: &mut [MmFeatureSpec], prompt_len: usize) -> Result<()> {
+    let mut fields = HashMap::new();
+    for feature in features.iter() {
+        if feature.identifier.is_empty() {
+            bail_invalid_preprocessed_multimodal!("feature identifier must be nonempty");
+        }
+        let data = feature.data.as_ref().ok_or_else(|| {
+            invalid_preprocessed_multimodal!(
+                "inline data is required; cache-only features are unsupported"
+            )
+        })?;
+        validate_kwargs(data)?;
+        for (key, elem) in data {
+            if let Some(previous) = fields.insert((feature.modality, key), &elem.field)
+                && !same_batching(previous, &elem.field)
+            {
+                bail_invalid_preprocessed_multimodal!(
+                    "incompatible batching for {} field {key:?}",
+                    feature.modality.as_str()
+                );
+            }
+        }
+        let position = &feature.mm_position;
+        if position.length == 0 || position.offset.checked_add(position.length).is_none() {
+            bail_invalid_preprocessed_multimodal!(
+                "placeholder range must have a positive length and fit in usize"
+            );
+        }
+        if position.offset + position.length > prompt_len {
+            bail_invalid_preprocessed_multimodal!(
+                "{} placeholder range exceeds prompt token IDs",
+                feature.modality.as_str()
+            );
+        }
+        if let Some(mask) = &position.is_embed {
+            if mask.dtype != TensorDtype::Bool || mask.shape != [position.length] {
+                bail_invalid_preprocessed_multimodal!(
+                    "is_embed must be a boolean tensor of shape [placeholder length]"
+                );
+            }
+            mask.validate_inline()
+                .map_err(|message| Error::InvalidPreprocessedMultimodal { message })?;
+        }
+    }
+    // Mirror the Python frontend (`argsort_mm_positions`): features are
+    // ordered by their placeholder position in the prompt.
+    features.sort_by_key(|feature| feature.mm_position.offset);
+    Ok(())
+}
+
+fn validate_kwargs(kwargs: &MmKwargsItem) -> Result<()> {
+    if kwargs.is_empty() {
+        bail_invalid_preprocessed_multimodal!("kwargs must contain at least one field");
+    }
+    for (key, elem) in kwargs {
+        if key.is_empty() {
+            bail_invalid_preprocessed_multimodal!("kwargs field name must be nonempty");
+        }
+        let data = elem.data.as_ref().ok_or_else(|| {
+            invalid_preprocessed_multimodal!("field {key:?} requires inline data")
+        })?;
+        validate_value(data, 0)?;
+        if let MmField::Flat(field) = &elem.field {
+            match data {
+                MmKwargValue::Tensor(tensor) => {
+                    let rank = tensor.shape.len() as i32;
+                    if rank == 0 || field.dim < -rank || field.dim >= rank {
+                        bail_invalid_preprocessed_multimodal!(
+                            "flat field {key:?} dim is outside the tensor rank"
+                        );
+                    }
+                }
+                MmKwargValue::List(_) if field.dim == 0 => {}
+                _ => bail_invalid_preprocessed_multimodal!(
+                    "flat field {key:?} requires a tensor or a list with dim=0"
+                ),
+            }
+        }
+    }
+    Ok(())
+}
+
+fn validate_value(value: &MmKwargValue, depth: usize) -> Result<()> {
+    if depth > MAX_DEPTH {
+        bail_invalid_preprocessed_multimodal!("kwargs nesting exceeds {MAX_DEPTH} levels");
+    }
+    match value {
+        MmKwargValue::Tensor(tensor) => tensor
+            .validate_inline()
+            .map_err(|message| Error::InvalidPreprocessedMultimodal { message }),
+        MmKwargValue::List(values) => {
+            for value in values {
+                validate_value(value, depth + 1)?;
+            }
+            Ok(())
+        }
+        MmKwargValue::Int(_) | MmKwargValue::Float(_) => Ok(()),
+    }
+}
+
+fn same_batching(left: &MmField, right: &MmField) -> bool {
+    // Slices and batch_size describe the producer's original batch. Python
+    // reduce_data combines the already-split item data without using them.
+    match (left, right) {
+        (MmField::Batched(a), MmField::Batched(b)) => a == b,
+        (MmField::Flat(a), MmField::Flat(b)) => a.dim == b.dim && a.keep_on_cpu == b.keep_on_cpu,
+        (MmField::Shared(a), MmField::Shared(b)) => a.keep_on_cpu == b.keep_on_cpu,
+        _ => false,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeMap;
+
+    use vllm_engine_core_client::protocol::multimodal::{
+        MmBatchedField, MmFieldElem, MmFlatField, MmModality, MmSharedField, MmSlice,
+        PlaceholderRange, SliceSpec,
+    };
+    use vllm_engine_core_client::protocol::tensor::{WireArrayData, WireTensor};
+
+    use super::*;
+    use crate::multimodal::tests::qwen3_vl_info;
+
+    fn feature(offset: usize, field: MmField) -> MmFeatureSpec {
+        MmFeatureSpec {
+            data: Some(BTreeMap::from([(
+                "pixels".to_owned(),
+                MmFieldElem {
+                    data: Some(MmKwargValue::Tensor(
+                        WireTensor::from_f32(vec![2], vec![1., 2.]).unwrap(),
+                    )),
+                    field,
+                },
+            )])),
+            modality: MmModality::Image,
+            identifier: format!("image-{offset}"),
+            mm_position: PlaceholderRange {
+                offset,
+                length: 2,
+                is_embed: None,
+            },
+            mm_hash: Some("processor-hash".to_owned()),
+        }
+    }
+
+    fn batched() -> MmField {
+        MmField::Batched(MmBatchedField { keep_on_cpu: false })
+    }
+
+    #[test]
+    fn inline_features_preserve_overlaps_and_stably_sort_prompt_positions() {
+        let info = qwen3_vl_info();
+        let mut a = feature(2, batched());
+        a.mm_position.is_embed = Some(WireTensor::from_bool(vec![2], vec![true, false]).unwrap());
+        let b = feature(1, batched());
+        let mut c = feature(2, batched());
+        c.modality = MmModality::Video;
+        let expected = vec![b.clone(), a.clone(), c.clone()];
+        let features = info.prepare_preprocessed(vec![a, b, c], 4).unwrap();
+        assert_eq!(features, expected);
+    }
+
+    #[test]
+    fn independently_preprocessed_items_keep_their_original_batch_metadata() {
+        let info = qwen3_vl_info();
+        let flat = |length| {
+            MmField::Flat(MmFlatField {
+                slices: vec![MmSlice::Slice(SliceSpec {
+                    start: Some(0),
+                    stop: Some(length),
+                    step: None,
+                })],
+                dim: 0,
+                keep_on_cpu: false,
+            })
+        };
+        let shared = |batch_size| {
+            MmField::Shared(MmSharedField {
+                batch_size,
+                keep_on_cpu: true,
+            })
+        };
+        for (a, b) in [(flat(2), flat(7)), (shared(1), shared(3))] {
+            let features = vec![feature(0, a), feature(2, b)];
+            let result = info.prepare_preprocessed(features.clone(), 4).unwrap();
+            assert_eq!(result, features);
+        }
+    }
+
+    #[test]
+    fn inline_features_reject_cache_references_and_invalid_tensor_storage() {
+        let info = qwen3_vl_info();
+        let original = feature(0, batched());
+        let mut cached = original.clone();
+        cached.data = None;
+        assert!(info.prepare_preprocessed(vec![cached], 4).is_err());
+        let mut field_cached = original.clone();
+        field_cached.data.as_mut().unwrap().get_mut("pixels").unwrap().data = None;
+        assert!(info.prepare_preprocessed(vec![field_cached], 4).is_err());
+        for tensor in [
+            WireTensor {
+                dtype: TensorDtype::F32,
+                shape: vec![2],
+                data: WireArrayData::AuxIndex(1),
+            },
+            WireTensor::from_raw(TensorDtype::F32, vec![2], vec![0; 4]),
+            WireTensor::from_raw(TensorDtype::F32, vec![usize::MAX, 2], vec![]),
+        ] {
+            let mut item = original.clone();
+            item.data.as_mut().unwrap().get_mut("pixels").unwrap().data =
+                Some(MmKwargValue::Tensor(tensor));
+            assert!(info.prepare_preprocessed(vec![item], 4).is_err());
+        }
+    }
+
+    #[test]
+    fn inline_features_validate_placeholder_and_mask_bounds() {
+        let info = qwen3_vl_info();
+        for (offset, length) in [(usize::MAX, 2), (0, 0)] {
+            let mut item = feature(offset, batched());
+            item.mm_position.length = length;
+            assert!(info.prepare_preprocessed(vec![item], 4).is_err());
+        }
+        assert!(info.prepare_preprocessed(vec![feature(2, batched())], 3).is_err());
+        for mask in [
+            WireTensor::from_bool(vec![1], vec![true]).unwrap(),
+            WireTensor::from_u32(vec![2], vec![1, 0]).unwrap(),
+        ] {
+            let mut item = feature(0, batched());
+            item.mm_position.is_embed = Some(mask);
+            assert!(info.prepare_preprocessed(vec![item], 4).is_err());
+        }
+    }
+
+    #[test]
+    fn inline_features_reject_incompatible_batching_and_flat_axes() {
+        let info = qwen3_vl_info();
+        let flat = MmField::Flat(MmFlatField {
+            slices: vec![],
+            dim: 2,
+            keep_on_cpu: false,
+        });
+        assert!(info.prepare_preprocessed(vec![feature(0, flat)], 4).is_err());
+        let shared = MmField::Shared(MmSharedField {
+            batch_size: 1,
+            keep_on_cpu: false,
+        });
+        assert!(
+            info.prepare_preprocessed(vec![feature(0, batched()), feature(2, shared)], 4)
+                .is_err()
+        );
+    }
+}
diff --git a/rust/src/chat/src/multimodal/tensor.rs b/rust/src/chat/src/multimodal/tensor.rs
index 2c564fbda721..39911309896c 100644
--- a/rust/src/chat/src/multimodal/tensor.rs
+++ b/rust/src/chat/src/multimodal/tensor.rs
@@ -18,7 +18,7 @@ use crate::error::{Error, Result, bail_multimodal, multimodal};
 pub(super) enum KwargValue {
     /// Tensor with row-major flat data and shape.
     Tensor { wire: WireTensor },
-    /// Non-tensor kwarg value that is shared or copied as-is.
+    /// Non-tensor kwarg value preserved in protocol form.
     Passthrough(ProtocolKwargValue),
 }
 
@@ -150,7 +150,17 @@ impl KwargValue {
     /// Batched fields use their first axis as media-item index and drop that
     /// axis in the per-feature value, matching vLLM's batched-field semantics.
     pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
-        self.wire_value_range(index, index + 1, true)
+        match self {
+            Self::Passthrough(ProtocolKwargValue::List(values)) => values
+                .get(index)
+                .cloned()
+                .ok_or_else(|| multimodal!("batched passthrough has no item at index {index}")),
+            Self::Passthrough(value) if index == 0 => Ok(value.clone()),
+            Self::Passthrough(_) => Err(multimodal!(
+                "scalar batched passthrough has no item at index {index}"
+            )),
+            Self::Tensor { .. } => self.wire_value_range(index, index + 1, true),
+        }
     }
 
     /// Convert one media item's flat tensor range directly to wire bytes.
@@ -336,6 +346,41 @@ mod tests {
         );
     }
 
+    #[test]
+    fn batched_passthrough_value_at_selects_media_item() {
+        let image_sizes = KwargValue::from_model_specific(
+            ModelSpecificValue::TupleVec(vec![(384, 736), (512, 512)]),
+            ModelDtype::BFloat16,
+        )
+        .unwrap();
+
+        assert_eq!(
+            image_sizes.batched_wire_value_at(1).unwrap(),
+            ProtocolKwargValue::List(vec![
+                ProtocolKwargValue::Int(512),
+                ProtocolKwargValue::Int(512),
+            ])
+        );
+        assert!(image_sizes.batched_wire_value_at(2).is_err());
+
+        let token_counts = KwargValue::from_model_specific(
+            ModelSpecificValue::IntVec(vec![276, 400]),
+            ModelDtype::BFloat16,
+        )
+        .unwrap();
+        assert_eq!(
+            token_counts.batched_wire_value_at(1).unwrap(),
+            ProtocolKwargValue::Int(400)
+        );
+
+        let scalar = KwargValue::Passthrough(ProtocolKwargValue::Int(7));
+        assert_eq!(
+            scalar.batched_wire_value_at(0).unwrap(),
+            ProtocolKwargValue::Int(7)
+        );
+        assert!(scalar.batched_wire_value_at(1).is_err());
+    }
+
     #[test]
     fn flat_wire_value_range_keeps_first_axis() {
         let value = KwargValue::Tensor {
diff --git a/rust/src/chat/src/multimodal/timing.rs b/rust/src/chat/src/multimodal/timing.rs
new file mode 100644
index 000000000000..c57383c7a3e2
--- /dev/null
+++ b/rust/src/chat/src/multimodal/timing.rs
@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Multimodal preprocessing spans and their timing collector configuration.
+
+use tracing::{Span, info_span};
+use tracing_subscriber::filter::{Filtered, LevelFilter};
+use vllm_tracing::timing::{RequestTimingLayer, RequestTimingStats};
+
+/// Target of the multimodal preprocessing stage spans.
+pub(super) const MM_STAGE_TARGET: &str = "mm_processor_timing";
+
+/// Create the timing layer and its stats handle for multimodal preprocessing
+/// stage spans (`vllm-bench mm-processor`), mirroring the Python
+/// `TimingContext` / `MultiModalTimingRegistry`.
+pub fn mm_timing_layer<S>() -> (
+    Filtered<RequestTimingLayer, LevelFilter, S>,
+    RequestTimingStats,
+) {
+    RequestTimingLayer::new(MM_STAGE_TARGET)
+}
+
+/// Span carrying the `request_id` used to attribute multimodal stage timings.
+pub fn mm_request_span(request_id: &str) -> Span {
+    info_span!("mm_request", request_id)
+}
diff --git a/rust/src/chat/src/multimodal/video.rs b/rust/src/chat/src/multimodal/video.rs
index c2a25b27bd59..a6b7856a62a2 100644
--- a/rust/src/chat/src/multimodal/video.rs
+++ b/rust/src/chat/src/multimodal/video.rs
@@ -16,7 +16,8 @@ use vllm_engine_core_client::protocol::multimodal::{
     SliceSpec,
 };
 
-use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor};
+use super::timing::MM_STAGE_TARGET;
+use super::{MultimodalModelInfo, PreparedItem, PreparedMedia, VisionModalitySupport, tensor};
 use crate::error::{Error, Result, bail_multimodal, multimodal};
 
 /// Forward-kwargs name of the primary video encoder input.
@@ -32,6 +33,12 @@ impl MultimodalModelInfo {
     /// Unlike images, each clip runs through the preprocessor independently
     /// (a batch of one), so its tensors are complete per item and need no
     /// cross-item slicing.
+    #[tracing::instrument(
+        name = "mm_stage",
+        target = MM_STAGE_TARGET,
+        skip_all,
+        fields(stage = "preprocess_video")
+    )]
     pub(super) async fn prepare_videos(
         &self,
         clips: Vec<Arc<VideoClip>>,
@@ -76,18 +83,17 @@ impl MultimodalModelInfo {
     /// processor.
     async fn preprocess_video_clip(
         &self,
-        support: &ModalitySupport,
+        support: &VisionModalitySupport,
         clip: Arc<VideoClip>,
     ) -> Result<PreprocessedEncoderInputs> {
-        let config = support.config.clone();
-        let processor = support.processor;
+        let processor = Arc::clone(&support.processor);
 
         tokio::task::spawn_blocking(move || {
             // Prefer the borrowed-RGB fast path, which avoids materializing a
             // `DynamicImage` per sampled frame after media decode.
             if let Some(rgb_video) = clip.rgb_video() {
                 match rgb_video.frame_refs() {
-                    Ok(frame_refs) => match processor.preprocess_video_rgb(&frame_refs, &config) {
+                    Ok(frame_refs) => match processor.preprocess_video_rgb(&frame_refs) {
                         Ok(preprocessed) => return Ok(preprocessed),
                         Err(error) => warn!(
                             error = %error.as_report(),
@@ -102,7 +108,7 @@ impl MultimodalModelInfo {
             }
 
             let frames = clip.materialized_frames().map_err(|error| multimodal!("{error}"))?;
-            Ok(processor.preprocess_video(&frames, &config)?)
+            Ok(processor.preprocess_video(&frames)?)
         })
         .await
         .map_err(|error| multimodal!("video preprocessing task failed: {error}"))?
@@ -117,7 +123,7 @@ impl MultimodalModelInfo {
 /// `flat_from_sizes` treatment of video patches), and batched metadata
 /// tensors drop their singleton batch axis.
 fn build_video_item(
-    support: &ModalitySupport,
+    support: &VisionModalitySupport,
     preprocessed: PreprocessedEncoderInputs,
     hash: String,
     uuid: Option<String>,
diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs
index a73115306342..ee067764cf77 100644
--- a/rust/src/chat/src/output/default/mod.rs
+++ b/rust/src/chat/src/output/default/mod.rs
@@ -3,25 +3,25 @@
 
 //! Default output processing pipeline.
 
-mod structural_tag;
 mod unified;
 
 use std::sync::Once;
 
 use futures::StreamExt as _;
 use tracing::info;
+use vllm_parser::output_grammar::{BuiltOutputGrammar, OutputGrammarContext};
 use vllm_parser::unified::{CombinedParser, UnifiedParser};
 use vllm_text::tokenizer::DynTokenizer;
+use xgrammar_structural_tag::ToolChoice;
 
-use self::structural_tag::apply_structural_tag_constraint;
 use self::unified::unified_event_stream;
 use super::structured::structured_chat_event_stream;
 use crate::error::Result;
 use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream};
-use crate::parser::ParserSelection;
 use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory};
 use crate::parser::tool::{ToolParser, ToolParserFactory};
 use crate::parser::unified::UnifiedParserFactory;
+use crate::parser::{ParserSelection, ToolStrictLevel};
 use crate::request::{ChatRequest, ChatTool};
 use crate::{Error, Result as ChatResult};
 
@@ -34,6 +34,18 @@ use crate::{Error, Result as ChatResult};
 pub struct DefaultChatOutputProcessor {
     parser: Box<dyn UnifiedParser>,
     parallel_tool_calls: bool,
+    /// Request facts the initialized parser needs to build its output grammar.
+    /// Absent for the plain-text-only processor, which never builds one.
+    grammar_inputs: Option<GrammarInputs>,
+}
+
+/// Request-scoped inputs for [`UnifiedParser::build_output_grammar`], captured
+/// at construction because the chat request is consumed before the prompt is
+/// tokenized.
+struct GrammarInputs {
+    tools: Vec<ChatTool>,
+    tool_choice: ToolChoice,
+    tool_strict_level: ToolStrictLevel,
 }
 
 impl DefaultChatOutputProcessor {
@@ -42,13 +54,15 @@ impl DefaultChatOutputProcessor {
     ///
     /// Parser resolution happens here so that request validation, prompt
     /// rendering, and streaming all observe the same parser-adjusted
-    /// request state.
+    /// request state. The parser is initialized and its output grammar built
+    /// later, once the final prompt token IDs are known.
     pub fn new(
         request: &mut ChatRequest,
         model_id: &str,
         tokenizer: DynTokenizer,
         tool_call_parser: &ParserSelection,
         reasoning_parser: &ParserSelection,
+        tool_strict_level: ToolStrictLevel,
     ) -> ChatResult<Self> {
         let parser = if let Some(parser) = Self::resolve_optional_unified_parser(
             request.tools(),
@@ -73,8 +87,6 @@ impl DefaultChatOutputProcessor {
             Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
         };
 
-        apply_structural_tag_constraint(request, parser.structural_tag_builder())?;
-
         if parser.preserve_special_tokens() {
             request.decode_options.skip_special_tokens = false;
         }
@@ -82,6 +94,11 @@ impl DefaultChatOutputProcessor {
         Ok(Self {
             parser,
             parallel_tool_calls: request.parallel_tool_calls(),
+            grammar_inputs: Some(GrammarInputs {
+                tools: request.tools().to_vec(),
+                tool_choice: request.tool_choice().into(),
+                tool_strict_level,
+            }),
         })
     }
 
@@ -94,6 +111,7 @@ impl DefaultChatOutputProcessor {
         Self {
             parser: Box::new(CombinedParser::plain_text_only()),
             parallel_tool_calls: true,
+            grammar_inputs: None,
         }
     }
 
@@ -170,6 +188,30 @@ static REASONING_PARSER_LOG_ONCE: Once = Once::new();
 static UNIFIED_PARSER_LOG_ONCE: Once = Once::new();
 
 impl ChatOutputProcessor for DefaultChatOutputProcessor {
+    fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
+        self.parser.initialize(prompt_token_ids).map_err(|error| {
+            Error::OutputParserInitialization {
+                error: Box::new(error),
+            }
+        })
+    }
+
+    fn build_output_grammar(&self) -> Result<Option<BuiltOutputGrammar>> {
+        let Some(inputs) = &self.grammar_inputs else {
+            return Ok(None);
+        };
+        self.parser
+            .build_output_grammar(&OutputGrammarContext {
+                tools: &inputs.tools,
+                tool_choice: &inputs.tool_choice,
+                tool_strict_level: inputs.tool_strict_level,
+                parallel_tool_calls: self.parallel_tool_calls,
+            })
+            .map_err(|error| Error::OutputGrammar {
+                error: Box::new(error),
+            })
+    }
+
     /// Transforms a raw generate-output token stream into structured chat
     /// events through two sequential stages once text decoding has
     /// already happened:
@@ -191,8 +233,9 @@ mod tests {
     use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::DefaultChatOutputProcessor;
-    use crate::parser::ParserSelection;
-    use crate::request::ChatRequest;
+    use crate::output::ChatOutputProcessor;
+    use crate::parser::{ParserSelection, ToolStrictLevel};
+    use crate::request::{ChatRequest, ChatTool, ChatToolChoice, ResolvedToolContext};
 
     fn tokenizer() -> Arc<TestTokenizer> {
         Arc::new(
@@ -202,6 +245,83 @@ mod tests {
         )
     }
 
+    #[test]
+    fn output_grammar_preserves_tool_strict_level() {
+        let build = |level, strict, choice| {
+            let tools = vec![ChatTool {
+                name: "search".to_string(),
+                description: None,
+                parameters: serde_json::json!({
+                    "type": "object",
+                    "properties": { "query": { "type": "string" } },
+                    "required": ["query"]
+                }),
+                strict,
+            }];
+            let mut request = ChatRequest {
+                tool_context: ResolvedToolContext::new(&[], tools, Some(choice), true).unwrap(),
+                ..ChatRequest::for_test()
+            };
+            let mut processor = DefaultChatOutputProcessor::new(
+                &mut request,
+                "other-model",
+                tokenizer(),
+                &ParserSelection::Explicit("qwen3_coder".to_string()),
+                &ParserSelection::None,
+                level,
+            )
+            .unwrap();
+            processor.initialize(&[]).unwrap();
+            processor.build_output_grammar().unwrap()
+        };
+
+        assert!(build(ToolStrictLevel::Auto, None, ChatToolChoice::Auto).is_none());
+        let envelope = build(ToolStrictLevel::Function, None, ChatToolChoice::Auto).unwrap();
+        let parameters = build(ToolStrictLevel::Parameter, None, ChatToolChoice::Auto).unwrap();
+        let strict = build(ToolStrictLevel::Auto, Some(true), ChatToolChoice::Auto).unwrap();
+        assert_eq!(parameters, strict);
+        assert_ne!(envelope, parameters);
+        assert!(build(ToolStrictLevel::Parameter, None, ChatToolChoice::None).is_none());
+    }
+
+    #[test]
+    fn tool_grammar_honors_parallel_call_policy() {
+        for parallel in [false, true] {
+            let tools = vec![ChatTool {
+                name: "lookup".to_string(),
+                description: None,
+                parameters: serde_json::json!({"type": "object"}),
+                strict: Some(true),
+            }];
+            let mut request = ChatRequest {
+                tool_context: ResolvedToolContext::new(
+                    &[],
+                    tools,
+                    Some(ChatToolChoice::Auto),
+                    parallel,
+                )
+                .unwrap(),
+                ..ChatRequest::for_test()
+            };
+            let mut processor = DefaultChatOutputProcessor::new(
+                &mut request,
+                "other-model",
+                tokenizer(),
+                &ParserSelection::Explicit("qwen3_coder".to_string()),
+                &ParserSelection::None,
+                ToolStrictLevel::Auto,
+            )
+            .unwrap();
+            processor.initialize(&[]).unwrap();
+            let built = processor.build_output_grammar().unwrap().unwrap();
+            let tag =
+                serde_json::to_value(xgrammar_structural_tag::StructuralTag::new(built.format))
+                    .unwrap();
+            assert_eq!(tag["format"]["type"], "triggered_tags");
+            assert_eq!(tag["format"]["stop_after_first"], !parallel);
+        }
+    }
+
     #[test]
     fn equal_explicit_gemma4_uses_unified_parser() {
         let mut request = ChatRequest::for_test();
@@ -213,6 +333,7 @@ mod tests {
             tokenizer(),
             &selection,
             &selection,
+            ToolStrictLevel::Auto,
         )
         .unwrap();
     }
@@ -227,6 +348,7 @@ mod tests {
             tokenizer(),
             &ParserSelection::Auto,
             &ParserSelection::Auto,
+            ToolStrictLevel::Auto,
         )
         .unwrap();
     }
@@ -244,6 +366,7 @@ mod tests {
                 tokenizer(),
                 tool,
                 reasoning,
+                ToolStrictLevel::Auto,
             )
             .unwrap();
         }
@@ -258,6 +381,7 @@ mod tests {
             tokenizer(),
             &ParserSelection::Auto,
             &ParserSelection::Explicit("gemma4".to_string()),
+            ToolStrictLevel::Auto,
         ) {
             Ok(_) => panic!("expected mixed Gemma4 parser selection to fail"),
             Err(error) => error,
diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs
deleted file mode 100644
index 2ef6f0868593..000000000000
--- a/rust/src/chat/src/output/default/structural_tag.rs
+++ /dev/null
@@ -1,316 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-//! Applies xgrammar structural-tag constraints for strict tool calling.
-
-use thiserror_ext::AsReport;
-use vllm_engine_core_client::protocol::structured_outputs::{
-    StructuredOutputBackend, StructuredOutputsParams,
-};
-use vllm_parser::tool::StructuralTagBuilder;
-use xgrammar_structural_tag::builders::StructuralTagOptions;
-use xgrammar_structural_tag::{
-    FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
-    build_structural_tag,
-};
-
-use crate::request::{ChatRequest, ChatToolChoice};
-use crate::{Error, Result as ChatResult};
-
-/// Apply structural tag constraints to the request based on the tool parser's structural tag
-/// support and the request's tool choice.
-pub(super) fn apply_structural_tag_constraint(
-    request: &mut ChatRequest,
-    builder: Option<&dyn StructuralTagBuilder>,
-) -> ChatResult<()> {
-    let Some(builder) = builder else {
-        return Ok(());
-    };
-    let Some(tool_choice) = structural_tag_tool_choice(request) else {
-        return Ok(());
-    };
-
-    let tools = request
-        .tools()
-        .iter()
-        .map(|tool| {
-            ToolParam::Function(FunctionToolParam::new(FunctionDefinition {
-                name: tool.name.clone(),
-                description: tool.description.clone(),
-                parameters: Some(tool.parameters.clone()),
-                strict: tool.strict,
-            }))
-        })
-        .collect::<Vec<_>>();
-
-    let structural_tag = build_structural_tag(
-        builder,
-        &tools,
-        tool_choice,
-        StructuralTagOptions::default().with_reasoning(false),
-    )
-    .and_then(|tag| tag.to_json_string())
-    .map_err(|error| Error::StructuralTag {
-        message: error.to_report_string(),
-    })?;
-
-    // Overwrite any existing structured output settings with the structural tag constraint.
-    request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
-        backend: StructuredOutputBackend::Xgrammar,
-        ..StructuredOutputsParams::structural_tag(structural_tag)
-    });
-
-    Ok(())
-}
-
-/// Resolve the tool choice used for [`xgrammar_structural_tag`] based on the request.
-///
-/// Returns `None` if no structural tag constraints should be applied.
-fn structural_tag_tool_choice(request: &ChatRequest) -> Option<StructuralTagToolChoice> {
-    if request.tools().is_empty() {
-        return None;
-    }
-
-    match request.tool_choice() {
-        // For `Auto`, only apply the structural tag if there's at least one strict tool.
-        ChatToolChoice::Auto if request.tools().iter().any(|tool| tool.strict == Some(true)) => {
-            Some(StructuralTagToolChoice::auto())
-        }
-        ChatToolChoice::Auto | ChatToolChoice::None => None,
-
-        ChatToolChoice::Required => Some(StructuralTagToolChoice::required()),
-        ChatToolChoice::Function { name } => Some(StructuralTagToolChoice::function(name.clone())),
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use serde_json::{Value, json};
-    use vllm_engine_core_client::protocol::structured_outputs::{
-        StructuredOutputBackend, StructuredOutputsParams,
-    };
-    use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser};
-
-    use super::*;
-    use crate::request::{ChatMessage, ResolvedToolContext};
-
-    fn chat_tool(name: &str, strict: Option<bool>) -> Tool {
-        Tool {
-            name: name.to_string(),
-            description: None,
-            parameters: json!({
-                "type": "object",
-                "properties": {
-                    "query": { "type": "string" }
-                },
-                "required": ["query"]
-            }),
-            strict,
-        }
-    }
-
-    fn qwen3_coder_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
-        Qwen3CoderToolParser::create(tools).expect("Qwen3 Coder parser should build")
-    }
-
-    fn request(tool_choice: ChatToolChoice, tools: Vec<Tool>) -> ChatRequest {
-        ChatRequest {
-            tool_context: ResolvedToolContext::new(&[], tools, Some(tool_choice), true)
-                .expect("tool context should resolve"),
-            ..ChatRequest::for_test()
-        }
-    }
-
-    fn structural_tag_value(request: &ChatRequest) -> Value {
-        let params = request
-            .sampling_params
-            .structured_outputs
-            .as_ref()
-            .expect("structured outputs should be set");
-        assert_eq!(params.backend, StructuredOutputBackend::Xgrammar);
-        let structural_tag = params
-            .constraint
-            .as_structural_tag()
-            .expect("structured output constraint should be structural_tag");
-        serde_json::from_str(structural_tag).expect("structural_tag should be valid JSON")
-    }
-
-    fn structured_outputs(request: &ChatRequest) -> &StructuredOutputsParams {
-        request
-            .sampling_params
-            .structured_outputs
-            .as_ref()
-            .expect("structured outputs should be set")
-    }
-
-    #[test]
-    fn auto_strict_tool_choice_builds_structural_tag() {
-        let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag should build");
-
-        let tag = structural_tag_value(&request);
-        assert_eq!(tag["type"], "structural_tag");
-        assert!(tag.to_string().contains("search"));
-    }
-
-    #[test]
-    fn required_dynamic_tool_builds_structural_tag_from_effective_tools() {
-        let messages = vec![ChatMessage::developer(
-            "",
-            Some(vec![chat_tool("lookup", None)]),
-        )];
-        let tool_context =
-            ResolvedToolContext::new(&messages, Vec::new(), Some(ChatToolChoice::Required), true)
-                .expect("dynamic tool context should resolve");
-        let mut request = ChatRequest {
-            messages,
-            tool_context,
-            ..ChatRequest::for_test()
-        };
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag should build");
-
-        let tag = structural_tag_value(&request);
-        assert!(tag.to_string().contains("lookup"));
-    }
-
-    #[test]
-    fn required_initial_and_dynamic_tools_build_one_structural_tag() {
-        let messages = vec![ChatMessage::developer(
-            "",
-            Some(vec![chat_tool("lookup", None)]),
-        )];
-        let tool_context = ResolvedToolContext::new(
-            &messages,
-            vec![chat_tool("search", None)],
-            Some(ChatToolChoice::Required),
-            true,
-        )
-        .expect("tool context should resolve");
-        let mut request = ChatRequest {
-            messages,
-            tool_context,
-            ..ChatRequest::for_test()
-        };
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag should build");
-
-        let tag = structural_tag_value(&request).to_string();
-        assert!(tag.contains("search"));
-        assert!(tag.contains("lookup"));
-    }
-
-    #[test]
-    fn auto_non_strict_tool_choice_skips_structural_tag() {
-        let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag decision should succeed");
-
-        assert!(request.sampling_params.structured_outputs.is_none());
-    }
-
-    #[test]
-    fn auto_strict_tool_choice_overwrites_existing_json_guidance() {
-        let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
-        request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
-            backend: StructuredOutputBackend::Xgrammar,
-            ..StructuredOutputsParams::json(json!({"type": "object"}))
-        });
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag should build");
-
-        let params = structured_outputs(&request);
-        assert!(params.constraint.is_structural_tag());
-        let tag = structural_tag_value(&request);
-        assert_eq!(tag["type"], "structural_tag");
-        assert!(tag.to_string().contains("search"));
-    }
-
-    #[test]
-    fn required_tool_choice_builds_structural_tag_without_strict_tools() {
-        let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag should build");
-
-        let tag = structural_tag_value(&request);
-        assert_eq!(tag["type"], "structural_tag");
-        assert!(tag.to_string().contains("search"));
-    }
-
-    #[test]
-    fn required_tool_choice_overwrites_existing_json_object_guidance() {
-        let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
-        request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
-            backend: StructuredOutputBackend::Xgrammar,
-            ..StructuredOutputsParams::json_object()
-        });
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag should build");
-
-        let params = structured_outputs(&request);
-        assert!(params.constraint.is_structural_tag());
-        let tag = structural_tag_value(&request);
-        assert_eq!(tag["type"], "structural_tag");
-        assert!(tag.to_string().contains("search"));
-    }
-
-    #[test]
-    fn named_tool_choice_builds_structural_tag_for_named_tool_only() {
-        let mut request = request(
-            ChatToolChoice::Function {
-                name: "lookup".to_string(),
-            },
-            vec![chat_tool("search", None), chat_tool("lookup", None)],
-        );
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag should build");
-
-        let tag = structural_tag_value(&request).to_string();
-        assert!(tag.contains("lookup"));
-        assert!(!tag.contains("search"));
-    }
-
-    #[test]
-    fn none_tool_choice_skips_structural_tag() {
-        let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag decision should succeed");
-
-        assert!(request.sampling_params.structured_outputs.is_none());
-    }
-
-    #[test]
-    fn none_tool_choice_preserves_existing_json_object_guidance() {
-        let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
-        request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
-            backend: StructuredOutputBackend::Xgrammar,
-            ..StructuredOutputsParams::json_object()
-        });
-        let parser = qwen3_coder_parser(request.tools());
-
-        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
-            .expect("structural tag decision should succeed");
-
-        let params = structured_outputs(&request);
-        assert!(params.constraint.is_json_object());
-    }
-}
diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs
index 964f0bf35a3a..b4a2dcbe5d30 100644
--- a/rust/src/chat/src/output/default/unified.rs
+++ b/rust/src/chat/src/output/default/unified.rs
@@ -48,25 +48,6 @@ impl UnifiedParserState {
         }
     }
 
-    /// Initialize parser state once prompt token IDs are available.
-    fn initialize(&mut self, prompt_token_ids: &[u32]) {
-        if self.parser_failed {
-            return;
-        }
-
-        match self.parser.initialize(prompt_token_ids) {
-            Ok(()) => {}
-            Err(error) => {
-                warn!(
-                    error = %error.as_report(),
-                    "failed to initialize unified parser; falling back to plain text deltas"
-                );
-                self.parser_failed = true;
-                self.open_call_index = None;
-            }
-        }
-    }
-
     /// Convert one decoded text delta into zero or more parsed assistant events.
     fn process_delta(&mut self, delta: DecodedText) -> Result<Vec<AssistantEvent>> {
         if self.parser_failed {
@@ -262,7 +243,6 @@ pub(crate) async fn unified_event_stream(
                 prompt_token_ids,
                 prompt_logprobs,
             } => {
-                state.initialize(&prompt_token_ids);
                 y.yield_ok(AssistantEvent::Start {
                     prompt_token_ids,
                     prompt_logprobs,
@@ -438,6 +418,7 @@ mod tests {
                 finish_reason: crate::FinishReason::Stop(None),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                sampling_mask: None,
             })),
         }
     }
diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs
index 8dcf84dd9866..035aef3cd4aa 100644
--- a/rust/src/chat/src/output/harmony/tests.rs
+++ b/rust/src/chat/src/output/harmony/tests.rs
@@ -58,6 +58,7 @@ fn finished() -> Finished {
         finish_reason: FinishReason::stop_eos(),
         kv_transfer_params: None,
         ec_transfer_params: None,
+        sampling_mask: None,
     }
 }
 
diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs
index dfccf3ba6f8c..8e62f3e8918b 100644
--- a/rust/src/chat/src/output/mod.rs
+++ b/rust/src/chat/src/output/mod.rs
@@ -7,6 +7,7 @@ use std::sync::Arc;
 use futures::Stream;
 use trait_set::trait_set;
 use uuid::Uuid;
+use vllm_parser::output_grammar::BuiltOutputGrammar;
 use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent};
 
 use crate::FinishReason;
@@ -15,11 +16,13 @@ use crate::event::{AssistantBlockKind, ChatEvent, ChatTokenUsage};
 
 mod default;
 mod harmony;
+mod structural_tag;
 mod structured;
 
 pub use default::DefaultChatOutputProcessor;
 pub use harmony::HarmonyChatOutputProcessor;
 pub(crate) use harmony::validate_harmony_parser_overrides;
+pub(crate) use structural_tag::apply_output_grammar;
 
 /// Internal assistant event before final assembly.
 ///
@@ -65,7 +68,24 @@ pub type DynChatEventStream = Pin<Box<dyn Stream<Item = Result<ChatEvent>> + Sen
 
 /// Request-scoped output processor from decoded text events into structured
 /// chat events.
+///
+/// A processor is created before the prompt is rendered, so that parser
+/// resolution can adjust the request, and then goes through two more phases
+/// once the final prompt token IDs are known: [`Self::initialize`] and
+/// [`Self::build_output_grammar`], in that order, both before
+/// [`Self::process`].
 pub trait ChatOutputProcessor: Send {
+    /// Initialize request-scoped parser state from the final prompt token IDs.
+    fn initialize(&mut self, _prompt_token_ids: &[u32]) -> Result<()> {
+        Ok(())
+    }
+
+    /// Build the structured output grammar implied by the initialized parser
+    /// state, or `None` when this request needs no parser-owned grammar.
+    fn build_output_grammar(&self) -> Result<Option<BuiltOutputGrammar>> {
+        Ok(None)
+    }
+
     /// Consume decoded text stream and return the structured chat-event stream.
     fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream>;
 }
diff --git a/rust/src/chat/src/output/structural_tag.rs b/rust/src/chat/src/output/structural_tag.rs
new file mode 100644
index 000000000000..dba6b6d1b739
--- /dev/null
+++ b/rust/src/chat/src/output/structural_tag.rs
@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Applies parser-built xgrammar structural-tag constraints.
+
+use vllm_engine_core_client::protocol::structured_outputs::{
+    StructuredOutputBackend, StructuredOutputsParams,
+};
+use vllm_parser::output_grammar::{BuiltOutputGrammar, GrammarCoverage};
+use vllm_text::TextRequest;
+use xgrammar_structural_tag::StructuralTag;
+
+use crate::{Error, Result};
+
+/// Apply one parser-built output grammar to the prepared text request.
+///
+/// A parser-owned grammar replaces any structured output constraint the user
+/// supplied, matching the Python frontend. `None` leaves the request untouched.
+pub(crate) fn apply_output_grammar(
+    request: &mut TextRequest,
+    built: Option<BuiltOutputGrammar>,
+) -> Result<()> {
+    let Some(built) = built else {
+        return Ok(());
+    };
+    let structural_tag = StructuralTag::new(built.format).to_json_string().map_err(|error| {
+        Error::OutputGrammar {
+            error: Box::new(error),
+        }
+    })?;
+
+    // Overwrite any existing structured output settings with the structural tag constraint.
+    request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
+        backend: StructuredOutputBackend::Xgrammar,
+        ..StructuredOutputsParams::structural_tag(structural_tag)
+    });
+    request.reasoning_ended = match built.coverage {
+        GrammarCoverage::FromTokenZero => Some(true),
+        GrammarCoverage::FinalOutputOnly => None,
+    };
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+    use vllm_engine_core_client::protocol::structured_outputs::{
+        StructuredOutputBackend, StructuredOutputsParams,
+    };
+    use xgrammar_structural_tag::format::Format;
+
+    use super::*;
+
+    #[test]
+    fn output_grammar_overwrites_answer_constraint() {
+        let mut request = TextRequest::for_test();
+        request.sampling_params.structured_outputs =
+            Some(StructuredOutputsParams::json(json!({ "type": "object" })));
+
+        let built = BuiltOutputGrammar::from_token_zero(Format::const_string("answer"));
+        apply_output_grammar(&mut request, Some(built)).unwrap();
+
+        assert_eq!(request.reasoning_ended, Some(true));
+        let params = request.sampling_params.structured_outputs.unwrap();
+        assert_eq!(params.backend, StructuredOutputBackend::Xgrammar);
+        let serialized = params.constraint.as_structural_tag().unwrap();
+        let value: serde_json::Value = serde_json::from_str(serialized).unwrap();
+        assert_eq!(value["type"], "structural_tag");
+        assert_eq!(value["format"]["type"], "const_string");
+        assert_eq!(value["format"]["value"], "answer");
+    }
+}
diff --git a/rust/src/chat/src/parser/mod.rs b/rust/src/chat/src/parser/mod.rs
index c10a4f11bf3e..632d82aaca1c 100644
--- a/rust/src/chat/src/parser/mod.rs
+++ b/rust/src/chat/src/parser/mod.rs
@@ -12,6 +12,8 @@ use std::str::FromStr;
 
 use serde_with::{DeserializeFromStr, SerializeDisplay};
 
+pub use vllm_parser::output_grammar::ToolStrictLevel;
+
 use self::reasoning::ReasoningParserFactory;
 use self::tool::ToolParserFactory;
 use self::unified::UnifiedParserFactory;
diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs
index c283b3056f88..9dfb4fdb90eb 100644
--- a/rust/src/chat/src/parser/reasoning/mod.rs
+++ b/rust/src/chat/src/parser/reasoning/mod.rs
@@ -8,10 +8,10 @@ use std::sync::{Arc, LazyLock};
 pub use vllm_parser::reasoning::{
     CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
     DeepSeekV4ReasoningParser, DeepSeekV41ReasoningParser, Glm45ReasoningParser,
-    Glm47ReasoningParser, KimiK2ReasoningParser, KimiReasoningParser, MiniMaxM2ReasoningParser,
-    MiniMaxM3ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta,
-    ReasoningError, ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser,
-    Step3p5ReasoningParser,
+    Glm47ReasoningParser, KimiK2ReasoningParser, KimiReasoningParser, MiMoReasoningParser,
+    MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser, NemotronV3ReasoningParser,
+    Qwen3ReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser,
+    Step3ReasoningParser, Step3p5ReasoningParser,
 };
 use vllm_tokenizer::DynTokenizer;
 
@@ -28,6 +28,7 @@ pub mod names {
     pub const GLM47: &str = "glm47";
     pub const KIMI: &str = "kimi";
     pub const KIMI_K2: &str = "kimi_k2";
+    pub const MIMO: &str = "mimo";
     pub const MINIMAX_M2: &str = "minimax_m2";
     pub const MINIMAX_M3: &str = "minimax_m3";
     pub const NEMOTRON_V3: &str = "nemotron_v3";
@@ -69,6 +70,7 @@ impl ReasoningParserFactory {
             .register_parser::<Glm47ReasoningParser>(names::GLM47)
             .register_parser::<KimiReasoningParser>(names::KIMI)
             .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
+            .register_parser::<MiMoReasoningParser>(names::MIMO)
             .register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
             .register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
             .register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
@@ -84,6 +86,7 @@ impl ReasoningParserFactory {
             .register_pattern("deepseek_v4", names::DEEPSEEK_V4)
             .register_pattern("deepseek-v3", names::DEEPSEEK_V3)
             .register_pattern("qwq", names::DEEPSEEK_R1)
+            .register_pattern("mimo-v2", names::MIMO)
             .register_pattern("qwen3", names::QWEN3)
             .register_pattern("glm-5", names::GLM47)
             .register_pattern("glm-4.7", names::GLM47)
diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs
index 26a2ecc4c790..9f3ab4f516a3 100644
--- a/rust/src/chat/src/parser/tool/mod.rs
+++ b/rust/src/chat/src/parser/tool/mod.rs
@@ -8,7 +8,7 @@ use std::sync::{Arc, LazyLock};
 pub use vllm_parser::tool::{
     DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
     DeepSeekV41ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser,
-    HermesToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser,
+    HermesToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MiMoToolParser,
     MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser,
     Qwen3CoderToolParser, Qwen3XmlToolParser, SeedOssToolParser, ToolParser, ToolParserError,
 };
@@ -33,6 +33,7 @@ pub mod names {
     pub const KIMI_K2: &str = "kimi_k2";
     pub const LLAMA3_JSON: &str = "llama3_json";
     pub const LLAMA4_JSON: &str = "llama4_json";
+    pub const MIMO: &str = "mimo";
     pub const MINIMAX_M2: &str = "minimax_m2";
     pub const MINIMAX_M3: &str = "minimax_m3";
     pub const MISTRAL: &str = "mistral";
@@ -76,6 +77,7 @@ impl ToolParserFactory {
             .register_parser::<KimiK2ToolParser>(names::KIMI_K2)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
+            .register_parser::<MiMoToolParser>(names::MIMO)
             .register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
             .register_parser::<MinimaxM3ToolParser>(names::MINIMAX_M3)
             .register_parser::<MistralToolParser>(names::MISTRAL)
@@ -87,6 +89,7 @@ impl ToolParserFactory {
         factory
             .register_pattern("mistral-", names::MISTRAL)
             .register_pattern("mixtral-", names::MISTRAL)
+            .register_pattern("mimo-v2", names::MIMO)
             .register_pattern("qwen3-coder", names::QWEN3_CODER)
             .register_pattern("qwen3.5", names::QWEN3_CODER)
             .register_pattern("qwq", names::HERMES)
diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs
index 24984c073146..69022e550d59 100644
--- a/rust/src/chat/src/parser/tool/tests.rs
+++ b/rust/src/chat/src/parser/tool/tests.rs
@@ -240,3 +240,22 @@ fn factory_new_registers_phi4_mini_json_by_name() {
     assert!(factory.contains(names::PHI4_MINI_JSON));
     factory.create(names::PHI4_MINI_JSON, &[]).unwrap();
 }
+
+#[test]
+fn factory_parses_mimo_parameter_tags() {
+    let factory = ToolParserFactory::new();
+    for model in ["XiaomiMiMo/MiMo-V2.5", "XiaomiMiMo/MiMo-V2.5-Pro"] {
+        assert_eq!(factory.resolve_name_for_model(model), Some(names::MIMO));
+        let mut parser = factory.create_for_model(model, &[]).unwrap();
+        let mut output = ToolParserOutput::default();
+        parser.parse_into("<tool_call>\n<function=lookup>\n<parameter=query>\n杭州\n</parameter>\n</function>\n</tool_call>", &mut output).unwrap();
+        output.append(parser.finish().unwrap());
+        let output = output.coalesce();
+        assert_eq!(output.calls().len(), 1);
+        assert_eq!(output.calls()[0].name.as_deref(), Some("lookup"));
+        assert_eq!(
+            serde_json::from_str::<serde_json::Value>(&output.calls()[0].arguments).unwrap(),
+            serde_json::json!({"query":"\n杭州\n"})
+        );
+    }
+}
diff --git a/rust/src/chat/src/reasoning.rs b/rust/src/chat/src/reasoning.rs
new file mode 100644
index 000000000000..b5a5b17684e6
--- /dev/null
+++ b/rust/src/chat/src/reasoning.rs
@@ -0,0 +1,397 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Shared reasoning-control resolution for native renderers.
+//!
+//! Each input source is parsed independently into [`ReasoningControl`]. Sources
+//! then fall back in priority order: request, deployment defaults, model default.
+//! Within a source, an explicit toggle decides the mode and effort supplies its
+//! intensity. Keeping sources separate lets request effort override a deployment
+//! toggle, and lets an enabled request inherit effort from enabled defaults.
+//!
+//! This module owns input types, precedence, and the standard `"none"` sentinel.
+//! Model adapters own supported effort names, numeric ranges, and prompt encoding.
+//! After model lowering, [`ReasoningControl::template_kwargs`] projects the prompt's
+//! decision for downstream consumers. HF prepares raw template inputs separately.
+
+use std::collections::HashMap;
+use std::fmt;
+
+use serde::{Deserialize, Serialize};
+use serde_json::{Number, Value, json};
+
+use crate::error::{Result, invalid_reasoning_control};
+use crate::request::ChatRequest;
+
+/// A model-specific reasoning effort, before or after renderer lowering.
+///
+/// Strings preserve model-specific names; JSON numbers preserve integer and
+/// floating-point representations for model-local range validation. Missing/null
+/// effort is represented by `Option::None` in [`crate::ChatOptions`].
+/// Supported names and ranges are validated by each renderer or HF template.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum EffortValue {
+    String(String),
+    Number(Number),
+}
+
+impl TryFrom<Value> for EffortValue {
+    type Error = crate::error::Error;
+
+    fn try_from(value: Value) -> Result<Self> {
+        match value {
+            Value::String(value) => Ok(Self::String(value)),
+            Value::Number(value) => Ok(Self::Number(value)),
+            _ => Err(invalid_reasoning_control!(
+                "reasoning effort must be a string or number, got {value}"
+            )),
+        }
+    }
+}
+
+impl TryFrom<f64> for EffortValue {
+    type Error = crate::error::Error;
+
+    /// Construct a numeric effort, rejecting NaN and infinities.
+    fn try_from(value: f64) -> Result<Self> {
+        Number::from_f64(value).map(Self::Number).ok_or_else(|| {
+            invalid_reasoning_control!("reasoning effort must be finite, got {value}")
+        })
+    }
+}
+
+impl From<&str> for EffortValue {
+    fn from(value: &str) -> Self {
+        Self::String(value.to_owned())
+    }
+}
+
+impl fmt::Display for EffortValue {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::String(value) => json!(value).fmt(f),
+            Self::Number(value) => value.fmt(f),
+        }
+    }
+}
+
+impl EffortValue {
+    pub fn as_str(&self) -> Option<&str> {
+        match self {
+            Self::String(value) => Some(value),
+            Self::Number(_) => None,
+        }
+    }
+
+    pub fn as_f64(&self) -> Option<f64> {
+        match self {
+            Self::Number(value) => value.as_f64(),
+            Self::String(_) => None,
+        }
+    }
+}
+
+/// Reasoning intent for one source, or the result of falling back across sources.
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) enum ReasoningControl {
+    /// This source leaves the mode and effort to a lower-priority source.
+    Default,
+    /// Reasoning is explicitly disabled; lower-priority defaults preserve this mode.
+    Disabled,
+    /// Reasoning is enabled. Missing effort inherits from an enabled default.
+    Enabled { effort: Option<EffortValue> },
+}
+
+impl ReasoningControl {
+    /// Parse request controls; the typed effort field takes precedence over kwargs effort.
+    pub(crate) fn from_request(request: &ChatRequest) -> Result<Self> {
+        Self::from_kwargs(
+            &request.chat_options.template_kwargs,
+            request.chat_options.reasoning_effort.as_ref(),
+        )
+    }
+
+    /// Parse a deployment source, whose controls all come from template kwargs.
+    pub(crate) fn from_template_kwargs(kwargs: &HashMap<String, Value>) -> Result<Self> {
+        Self::from_kwargs(kwargs, None)
+    }
+
+    /// Select controls within one source, then derive its mode and optional effort.
+    fn from_kwargs(
+        kwargs: &HashMap<String, Value>,
+        typed_effort: Option<&EffortValue>,
+    ) -> Result<Self> {
+        // Presence sets alias priority: `thinking` wins, including when its value
+        // is invalid. Validate the selected spelling as a boolean.
+        let toggle = ["thinking", "enable_thinking"]
+            .into_iter()
+            .find_map(|key| kwargs.get(key).map(|value| (key, value)))
+            .map(|(key, value)| {
+                value.as_bool().ok_or_else(|| {
+                    invalid_reasoning_control!(
+                        "template kwarg `{key}` must be a boolean, got {value}"
+                    )
+                })
+            })
+            .transpose()?;
+
+        // The typed request field wins over kwargs. Parse kwargs only when used;
+        // disabled mode discards effort, and JSON null means omitted effort.
+        let effort = match (toggle, typed_effort) {
+            (Some(false), _) => None,
+            (_, Some(effort)) => Some(effort.clone()),
+            (_, None) => kwargs
+                .get("reasoning_effort")
+                .filter(|effort| !effort.is_null())
+                .cloned()
+                .map(EffortValue::try_from)
+                .transpose()?,
+        };
+
+        match (toggle, effort) {
+            // 1. Explicit false disables reasoning and discards any effort,
+            //    including values that would fail validation in an enabled mode.
+            (Some(false), _) => Ok(Self::Disabled),
+            // 2. `none` disables reasoning when the toggle is absent. Explicit true
+            //    enables reasoning and leaves effort to enabled deployment/model defaults.
+            (toggle, Some(EffortValue::String(effort))) if effort == "none" => {
+                Ok(if toggle == Some(true) {
+                    Self::Enabled { effort: None }
+                } else {
+                    Self::Disabled
+                })
+            }
+            // 3. Absent toggle and absent/null effort leave the entire decision
+            //    to the next source in the fallback chain.
+            (None, None) => Ok(Self::Default),
+            // 4. Explicit true or a concrete effort enables reasoning;
+            //    missing effort inherits later.
+            (_, effort) => Ok(Self::Enabled { effort }),
+        }
+    }
+
+    /// Parse request and deployment independently, then apply request precedence.
+    ///
+    /// Each source validates its selected input types. Model adapters subsequently
+    /// apply their model default and validate/map the resulting effective effort.
+    pub(crate) fn resolve(
+        request: &ChatRequest,
+        defaults: &HashMap<String, Value>,
+    ) -> Result<Self> {
+        Ok(Self::from_request(request)?.fallback(Self::from_template_kwargs(defaults)?))
+    }
+
+    /// Fill this higher-priority source's missing decision or effort from `defaults`.
+    ///
+    /// `Default` inherits the whole lower-priority state. `Enabled { effort: None }`
+    /// inherits effort only from enabled defaults and retains its enabled mode.
+    /// Explicit disabling and concrete effort are final at this priority level.
+    pub(crate) fn fallback(self, defaults: Self) -> Self {
+        match (self, defaults) {
+            (Self::Default, defaults) => defaults,
+            (Self::Enabled { effort: None }, Self::Enabled { effort }) => Self::Enabled { effort },
+            (request, _) => request,
+        }
+    }
+
+    /// Construct an enabled model default or a model-lowered concrete effort.
+    pub(crate) fn enabled(effort: impl Into<EffortValue>) -> Self {
+        Self::Enabled {
+            effort: Some(effort.into()),
+        }
+    }
+
+    /// Report the resolved mode; callers apply model defaults before rendering.
+    pub(crate) fn is_enabled(&self) -> bool {
+        matches!(self, Self::Enabled { .. })
+    }
+
+    /// Return the active effort, which may still be missing before model fallback.
+    pub(crate) fn effort(&self) -> Option<&EffortValue> {
+        match self {
+            Self::Enabled { effort } => effort.as_ref(),
+            Self::Default | Self::Disabled => None,
+        }
+    }
+
+    /// Project the same model-lowered decision used to render the native prompt.
+    ///
+    /// Preserve request extras and write both toggle aliases plus the effective
+    /// effort. Disabled mode emits `"none"`; enabled mode with no effective effort
+    /// clears the original effort key. `Default` preserves request kwargs.
+    pub(crate) fn template_kwargs(&self, request: &ChatRequest) -> HashMap<String, Value> {
+        let mut kwargs = request.chat_options.template_kwargs.clone();
+        if matches!(self, Self::Default) {
+            return kwargs;
+        }
+        kwargs.insert("thinking".to_string(), json!(self.is_enabled()));
+        kwargs.insert("enable_thinking".to_string(), json!(self.is_enabled()));
+        match self {
+            Self::Disabled => {
+                kwargs.insert("reasoning_effort".to_string(), json!("none"));
+            }
+            Self::Enabled {
+                effort: Some(effort),
+            } => {
+                kwargs.insert("reasoning_effort".to_string(), json!(effort));
+            }
+            Self::Enabled { effort: None } => {
+                kwargs.remove("reasoning_effort");
+            }
+            Self::Default => unreachable!(),
+        }
+        kwargs
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn kwargs(value: Value) -> HashMap<String, Value> {
+        serde_json::from_value(value).unwrap()
+    }
+
+    #[test]
+    fn numeric_effort_rejects_non_finite_values() {
+        for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
+            assert!(EffortValue::try_from(value).unwrap_err().is_request_validation_error());
+        }
+    }
+
+    #[test]
+    fn explicit_toggle_decides_mode_and_none_inherits_enabled_effort() {
+        use ReasoningControl::{Default, Disabled, Enabled};
+
+        for (source, expected) in [
+            (json!({}), Default),
+            (json!({"reasoning_effort": "none"}), Disabled),
+            (
+                json!({"reasoning_effort": "high"}),
+                ReasoningControl::enabled("high"),
+            ),
+            (json!({"enable_thinking": true}), Enabled { effort: None }),
+            (
+                json!({"enable_thinking": true, "reasoning_effort": "none"}),
+                Enabled { effort: None },
+            ),
+            (
+                json!({"enable_thinking": true, "reasoning_effort": "low"}),
+                ReasoningControl::enabled("low"),
+            ),
+            (
+                json!({"enable_thinking": false, "reasoning_effort": "high"}),
+                Disabled,
+            ),
+            (
+                json!({"thinking": true, "enable_thinking": false}),
+                Enabled { effort: None },
+            ),
+            (
+                json!({"thinking": false, "enable_thinking": true}),
+                Disabled,
+            ),
+            (
+                json!({"reasoning_effort": 37}),
+                ReasoningControl::enabled(EffortValue::Number(37.into())),
+            ),
+            (json!({"reasoning_effort": null}), Default),
+        ] {
+            let source = kwargs(source);
+            assert_eq!(
+                ReasoningControl::from_template_kwargs(&source).unwrap(),
+                expected,
+                "{source:?}"
+            );
+        }
+    }
+
+    #[test]
+    fn request_and_deployment_states_fall_back_before_model_defaults() {
+        for (request, defaults, expected) in [
+            (
+                json!({}),
+                json!({"enable_thinking": false}),
+                ReasoningControl::Disabled,
+            ),
+            (
+                json!({"reasoning_effort": "high"}),
+                json!({"thinking": false}),
+                ReasoningControl::enabled("high"),
+            ),
+            (
+                json!({"enable_thinking": false}),
+                json!({"reasoning_effort": "high"}),
+                ReasoningControl::Disabled,
+            ),
+            (
+                json!({"thinking": true, "reasoning_effort": "none"}),
+                json!({"reasoning_effort": "low"}),
+                ReasoningControl::enabled("low"),
+            ),
+            (
+                json!({"thinking": true}),
+                json!({"thinking": false, "reasoning_effort": "low"}),
+                ReasoningControl::enabled("max"),
+            ),
+            (json!({}), json!({}), ReasoningControl::enabled("max")),
+        ] {
+            let request = kwargs(request);
+            let defaults = kwargs(defaults);
+            let actual = ReasoningControl::from_template_kwargs(&request)
+                .unwrap()
+                .fallback(ReasoningControl::from_template_kwargs(&defaults).unwrap())
+                .fallback(ReasoningControl::enabled("max"));
+            assert_eq!(
+                actual, expected,
+                "request={request:?}, defaults={defaults:?}"
+            );
+        }
+    }
+
+    #[test]
+    fn selected_toggle_requires_a_boolean() {
+        for value in [
+            json!("true"),
+            json!(1),
+            json!(null),
+            json!({"type": "enabled"}),
+        ] {
+            assert!(
+                ReasoningControl::from_template_kwargs(&kwargs(json!({
+                    "thinking": value, "enable_thinking": true
+                })))
+                .unwrap_err()
+                .is_request_validation_error()
+            );
+        }
+    }
+
+    #[test]
+    fn invalid_effort_types_are_rejected_only_when_the_source_uses_them() {
+        for value in [json!(true), json!(false), json!([]), json!({})] {
+            for toggle in [None, Some(true), Some(false)] {
+                let mut source = kwargs(json!({"reasoning_effort": value}));
+                if let Some(toggle) = toggle {
+                    source.insert("thinking".to_owned(), json!(toggle));
+                }
+                let control = ReasoningControl::from_template_kwargs(&source);
+                if toggle == Some(false) {
+                    assert_eq!(control.unwrap(), ReasoningControl::Disabled);
+                } else {
+                    assert!(control.unwrap_err().is_request_validation_error());
+                }
+            }
+
+            // A typed request effort shadows malformed kwargs effort.
+            let mut request = ChatRequest::for_test();
+            request.chat_options.reasoning_effort = Some(EffortValue::from("high"));
+            request.chat_options.template_kwargs = kwargs(json!({"reasoning_effort": value}));
+            assert_eq!(
+                ReasoningControl::from_request(&request).unwrap(),
+                ReasoningControl::enabled("high")
+            );
+        }
+    }
+}
diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs
index ab435028ebf3..2683c37ec3b7 100644
--- a/rust/src/chat/src/renderer/deepseek.rs
+++ b/rust/src/chat/src/renderer/deepseek.rs
@@ -11,15 +11,16 @@ use std::collections::HashMap;
 use std::fmt::Write as _;
 
 use serde::Serialize;
-use serde_json::Value;
+use serde_json::{Map, Value};
 use serde_json_fmt::JsonFormat;
 
 use llm_multimodal::DEEPSEEK_V41_IMAGE_PLACEHOLDER;
 
+use super::MediaPartSource;
+use crate::EffortValue;
 use crate::error::{Error, Result};
-use crate::request::{
-    ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ReasoningEffort,
-};
+use crate::reasoning::ReasoningControl;
+use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool};
 use crate::{AssistantContentBlock, AssistantMessageExt, AssistantToolCall};
 
 const BOS_TOKEN: &str = "<｜begin▁of▁sentence｜>";
@@ -86,13 +87,26 @@ struct RenderedToolSchema<'a> {
 }
 
 /// Render one chat request into the final prompt string.
-pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Result<String> {
-    let (thinking_mode, reasoning_effort_prompt) = match dialect {
-        DsDialect::V4 => {
-            resolve_thinking_options(request).map(|(mode, prompt)| (mode, Cow::Borrowed(prompt)))?
-        }
-        DsDialect::V41 => resolve_v41_thinking_options(request)?,
+pub(super) fn render_request(
+    request: &ChatRequest,
+    dialect: DsDialect,
+    reasoning: &ReasoningControl,
+) -> Result<String> {
+    render_request_with_media_order(request, dialect, reasoning).map(|(prompt, _)| prompt)
+}
+
+/// Render one chat request and record media in placeholder order.
+pub(super) fn render_request_with_media_order(
+    request: &ChatRequest,
+    dialect: DsDialect,
+    reasoning: &ReasoningControl,
+) -> Result<(String, Vec<MediaPartSource>)> {
+    let thinking_mode = if reasoning.is_enabled() {
+        ThinkingMode::Thinking
+    } else {
+        ThinkingMode::Chat
     };
+    let reasoning_effort_prompt = reasoning_prompt(reasoning, dialect);
     let request_tools = request_tools(request);
     let synthetic_tool_system = needs_synthetic_tool_system(request, request_tools);
     let drop_thinking = request.parse_template_bool("drop_thinking")?.unwrap_or(true)
@@ -112,6 +126,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
         dialect,
     );
     let mut out = String::from(BOS_TOKEN);
+    let mut media_order = Vec::new();
     if dialect == DsDialect::V41
         && (thinking_mode == ThinkingMode::Thinking
             || synthetic_tool_system
@@ -126,7 +141,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
     let mut request_tools_attached = false;
     let mut render_index = 0isize;
     if synthetic_tool_system {
-        render_system_message(&mut out, None, request_tools, dialect)?;
+        render_system_message(&mut out, None, request_tools, dialect, &mut media_order)?;
         request_tools_attached = true;
         render_index += 1;
     }
@@ -157,14 +172,22 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
                 } else {
                     &[]
                 };
-                render_system_message(&mut out, Some(content), tools, dialect)?;
+                render_system_message(
+                    &mut out,
+                    Some((message_index, content)),
+                    tools,
+                    dialect,
+                    &mut media_order,
+                )?;
             }
             ChatMessage::Developer { content, tools } => {
                 render_developer_message(
                     &mut out,
+                    message_index,
                     content,
                     tools.as_deref().unwrap_or(&[]),
                     dialect,
+                    &mut media_order,
                 )?;
             }
             ChatMessage::User { .. } | ChatMessage::ToolResponse { .. } => {
@@ -174,6 +197,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
                     message_index,
                     &last_tool_call_order,
                     dialect,
+                    &mut media_order,
                 )?;
             }
             ChatMessage::Assistant { content } => {
@@ -203,99 +227,98 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
                     );
                 }
             }
+            ChatMessage::Custom { role, .. } => {
+                return Err(Error::UnsupportedChatRole { role: role.clone() });
+            }
         }
 
-        if (is_user_like_entry(message, current_render_index as usize, dialect)
-            || (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. })))
-            && next_rendered_entry_is_assistant_or_end(
+        if is_user_like_entry(message, current_render_index as usize, dialect)
+            || (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. }))
+        {
+            // Skip the opener only at the last rendered entry when
+            // add_generation_prompt is off. Inspect the next rendered
+            // entry, not the raw index: V4 merges trailing tool results.
+            let should_write_transition = match next_rendered_entry(
                 request.messages.as_slice(),
                 message_index,
                 drop_historical_developers,
                 last_user_like_message_index,
-            )
-        {
-            write_assistant_transition(
-                &mut out,
-                thinking_mode,
-                drop_thinking,
-                current_render_index >= last_user_render_index,
-            );
+            ) {
+                Some(ChatMessage::Assistant { .. }) => true,
+                None => request.chat_options.add_generation_prompt(),
+                Some(_) => false,
+            };
+            if should_write_transition {
+                write_assistant_transition(
+                    &mut out,
+                    thinking_mode,
+                    drop_thinking,
+                    current_render_index >= last_user_render_index,
+                );
+            }
         }
     }
 
-    Ok(out)
+    Ok((out, media_order))
 }
 
-/// Resolve DeepSeek V4's thinking controls. Unlike the Python tokenizer
-/// wrapper, the Rust renderer only consumes the typed top-level
-/// `reasoning_effort`; the generic template-kwargs map is left for HF
-/// templates.
-fn resolve_thinking_options(request: &ChatRequest) -> Result<(ThinkingMode, &'static str)> {
-    let mut thinking_mode = match request.enable_thinking()?.unwrap_or(true) {
-        true => ThinkingMode::Thinking,
-        false => ThinkingMode::Chat,
-    };
-    let mut reasoning_effort_prompt = REASONING_EFFORT_HIGH;
-
-    match request.chat_options.reasoning_effort {
-        Some(ReasoningEffort::None) => thinking_mode = ThinkingMode::Chat,
-        Some(ReasoningEffort::Max) => {
-            reasoning_effort_prompt = REASONING_EFFORT_MAX;
-        }
-        Some(ReasoningEffort::XHigh | ReasoningEffort::High) => {
-            reasoning_effort_prompt = REASONING_EFFORT_HIGH;
+/// Map the shared reasoning decision to the selected DeepSeek effort dialect.
+pub(super) fn resolve_reasoning(
+    control: ReasoningControl,
+    dialect: DsDialect,
+) -> Result<ReasoningControl> {
+    let control = control.fallback(ReasoningControl::enabled("high"));
+    if !control.is_enabled() {
+        return Ok(control);
+    }
+    let effort = control.effort().expect("enabled DeepSeek reasoning has a default effort");
+    let effort = match dialect {
+        DsDialect::V4 => {
+            let effort = match effort.as_str() {
+                Some("minimal" | "medium" | "low") => "low",
+                Some("high" | "xhigh") => "high",
+                Some("max") => "max",
+                _ => {
+                    return Err(Error::InvalidReasoningEffort(format!(
+                        "DeepSeek V4 reasoning_effort must be minimal, low, medium, high, xhigh, or max, got {effort}"
+                    )));
+                }
+            };
+            EffortValue::from(effort)
         }
-        Some(ReasoningEffort::Minimal | ReasoningEffort::Medium | ReasoningEffort::Low) => {
-            reasoning_effort_prompt = "";
+        DsDialect::V41 => {
+            let budget = match effort {
+                EffortValue::String(effort) => match effort.as_str() {
+                    "low" => 25,
+                    "high" => 50,
+                    "xhigh" => 75,
+                    "max" => 100,
+                    _ => return Err(invalid_v41_reasoning_effort()),
+                },
+                EffortValue::Number(budget) => budget
+                    .as_u64()
+                    .filter(|budget| (1..=100).contains(budget))
+                    .ok_or_else(invalid_v41_reasoning_effort)?,
+            };
+            EffortValue::Number(budget.into())
         }
-        None => {}
-    }
-
-    Ok((thinking_mode, reasoning_effort_prompt))
+    };
+    Ok(ReasoningControl::enabled(effort))
 }
 
-/// Resolve V4.1's numeric reasoning effort using the top-level value before template kwargs.
-fn resolve_v41_thinking_options(
-    request: &ChatRequest,
-) -> Result<(ThinkingMode, Cow<'static, str>)> {
-    let mut thinking = request.enable_thinking()?.unwrap_or(true);
-    let budget = match request.chat_options.reasoning_effort {
-        Some(ReasoningEffort::None) => {
-            thinking = false;
-            50
-        }
-        Some(ReasoningEffort::Low) => 25,
-        Some(ReasoningEffort::High) => 50,
-        Some(ReasoningEffort::XHigh) => 75,
-        Some(ReasoningEffort::Max) => 100,
-        Some(ReasoningEffort::Minimal | ReasoningEffort::Medium) => {
-            return Err(invalid_v41_reasoning_effort());
-        }
-        None => match request.chat_options.template_kwargs.get("reasoning_effort") {
-            None | Some(Value::Null) => 50,
-            Some(Value::String(effort)) => match effort.as_str() {
-                "low" => 25,
-                "high" => 50,
-                "xhigh" => 75,
-                "max" => 100,
-                _ => return Err(invalid_v41_reasoning_effort()),
-            },
-            Some(Value::Number(budget)) => budget
-                .as_u64()
-                .filter(|budget| (1..=100).contains(budget))
-                .ok_or_else(invalid_v41_reasoning_effort)?,
-            Some(_) => return Err(invalid_v41_reasoning_effort()),
-        },
+fn reasoning_prompt(control: &ReasoningControl, dialect: DsDialect) -> Cow<'static, str> {
+    let Some(effort) = control.effort() else {
+        return Cow::Borrowed("");
     };
-    if thinking {
-        Ok((
-            ThinkingMode::Thinking,
-            Cow::Owned(format!(
-                "Reasoning Effort: {budget} (range 1-100, the higher the value, the more thorough the reasoning)\n\n"
-            )),
-        ))
-    } else {
-        Ok((ThinkingMode::Chat, Cow::Borrowed("")))
+    match dialect {
+        DsDialect::V4 => Cow::Borrowed(match effort.as_str() {
+            Some("high") => REASONING_EFFORT_HIGH,
+            Some("max") => REASONING_EFFORT_MAX,
+            _ => "",
+        }),
+        DsDialect::V41 => Cow::Owned(format!(
+            "Reasoning Effort: {effort} (range 1-100, the higher the value, the more thorough the reasoning)\n\n"
+        )),
     }
 }
 
@@ -401,14 +424,17 @@ fn is_historical_developer(
         && last_user_like_message_index.is_some_and(|last_index| message_index < last_index)
 }
 
-/// Return whether the next rendered entry is assistant, or there is no next
-/// entry.
-fn next_rendered_entry_is_assistant_or_end(
+/// Return the next rendered message after `message_index`.
+///
+/// Skips merged user-content tails and dropped historical developers so the
+/// caller can tell a following assistant turn apart from the end of the
+/// conversation.
+fn next_rendered_entry(
     messages: &[ChatMessage],
     message_index: usize,
     drop_historical_developers: bool,
     last_user_like_message_index: Option<usize>,
-) -> bool {
+) -> Option<&ChatMessage> {
     let mut next_index = message_index + 1;
     while next_index < messages.len()
         && (is_following_user_content(messages, next_index)
@@ -422,10 +448,7 @@ fn next_rendered_entry_is_assistant_or_end(
         next_index += 1;
     }
 
-    messages
-        .get(next_index)
-        .map(|message| matches!(message, ChatMessage::Assistant { .. }))
-        .unwrap_or(true)
+    messages.get(next_index)
 }
 
 /// Render the tool preamble shown to the model for one DeepSeek dialect.
@@ -488,12 +511,13 @@ fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> {
 /// Render a system turn, optionally followed by the V4 tool preamble.
 fn render_system_message(
     out: &mut String,
-    content: Option<&ChatContent>,
+    content: Option<(usize, &ChatContent)>,
     tools: &[ChatTool],
     dialect: DsDialect,
+    media_order: &mut Vec<MediaPartSource>,
 ) -> Result<()> {
-    if let Some(content) = content {
-        write_chat_content(out, content, dialect)?;
+    if let Some((message_index, content)) = content {
+        write_chat_content(out, content, dialect, message_index, media_order)?;
     }
     if !tools.is_empty() {
         out.push_str("\n\n");
@@ -505,9 +529,11 @@ fn render_system_message(
 /// Developer messages are rendered as user-like turns with optional tools.
 fn render_developer_message(
     out: &mut String,
+    message_index: usize,
     content: &ChatContent,
     tools: &[ChatTool],
     dialect: DsDialect,
+    media_order: &mut Vec<MediaPartSource>,
 ) -> Result<()> {
     if content.is_empty() {
         return Err(Error::ChatTemplate(
@@ -516,7 +542,7 @@ fn render_developer_message(
     }
 
     out.push_str(USER_SP_TOKEN);
-    write_chat_content(out, content, dialect)?;
+    write_chat_content(out, content, dialect, message_index, media_order)?;
     if !tools.is_empty() {
         out.push_str("\n\n");
         render_tools(out, tools, dialect)?;
@@ -531,6 +557,7 @@ fn render_user_content_block(
     message_index: usize,
     tool_call_order: &HashMap<String, usize>,
     dialect: DsDialect,
+    media_order: &mut Vec<MediaPartSource>,
 ) -> Result<()> {
     let (block_start, block_end) = user_content_block_bounds(messages, message_index);
     let mut sorted_tool_indices =
@@ -542,7 +569,9 @@ fn render_user_content_block(
             out.push_str("\n\n");
         }
         match &messages[message_index] {
-            ChatMessage::User { content } => write_chat_content(out, content, dialect)?,
+            ChatMessage::User { content } => {
+                write_chat_content(out, content, dialect, message_index, media_order)?;
+            }
             ChatMessage::ToolResponse { .. } => {
                 let sorted_index = sorted_tool_indices
                     .next()
@@ -550,7 +579,7 @@ fn render_user_content_block(
                 let ChatMessage::ToolResponse { content, .. } = &messages[sorted_index] else {
                     unreachable!("sorted tool response index should reference a tool message");
                 };
-                write_tool_result(out, content, dialect)?;
+                write_tool_result(out, content, dialect, sorted_index, media_order)?;
             }
             _ => unreachable!("user content block should only contain user content messages"),
         }
@@ -597,9 +626,15 @@ fn sorted_tool_response_indices(
 }
 
 /// Render one tool response payload inside a V4 `<tool_result>` block.
-fn write_tool_result(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> {
+fn write_tool_result(
+    out: &mut String,
+    content: &ChatContent,
+    dialect: DsDialect,
+    message_index: usize,
+    media_order: &mut Vec<MediaPartSource>,
+) -> Result<()> {
     out.push_str("<tool_result>");
-    write_chat_content(out, content, dialect)?;
+    write_chat_content(out, content, dialect, message_index, media_order)?;
     out.push_str("</tool_result>");
     Ok(())
 }
@@ -685,19 +720,18 @@ fn encode_arguments_to_dsml(
     dialect: DsDialect,
 ) -> Result<()> {
     let parameter_tag = dialect.parameter_tag();
-    let arguments: Value = serde_json::from_str(&tool_call.arguments).map_err(|error| {
-        Error::ChatTemplate(format!(
-            "assistant tool call has invalid JSON arguments for DeepSeek: {error}"
-        ))
-    })?;
-    let Some(arguments) = arguments.as_object() else {
-        return Err(Error::ChatTemplate(
-            "assistant tool call arguments for DeepSeek must be a JSON object".to_string(),
-        ));
-    };
+    // Match deepseek-recipe's render_tool_arguments for both V4 and V4.1:
+    // https://github.com/deepseek-ai/deepseek-recipe/blob/8cadfede7063c896b944e7bae05daa3549ae97ea/deepseek-recipe-encoding/src/v4/mod.rs#L48-L58
+    let arguments = serde_json::from_str::<Map<String, Value>>(&tool_call.arguments)
+        .unwrap_or_else(|_| {
+            Map::from_iter([(
+                "arguments".to_owned(),
+                Value::String(tool_call.arguments.clone()),
+            )])
+        });
 
     let mut wrote_parameter = false;
-    for (key, value) in arguments {
+    for (key, value) in &arguments {
         if wrote_parameter {
             out.push('\n');
         }
@@ -728,7 +762,13 @@ fn encode_arguments_to_dsml(
 /// The V4.1 dialect inlines the image placeholder at each image part's
 /// position unconditionally (matching the Python encoding's
 /// `IMAGE_PLACEHOLDER`); other dialects reject multimodal parts.
-fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> {
+fn write_chat_content(
+    out: &mut String,
+    content: &ChatContent,
+    dialect: DsDialect,
+    message_index: usize,
+    media_order: &mut Vec<MediaPartSource>,
+) -> Result<()> {
     match content {
         ChatContent::Text(text) => out.push_str(text),
         ChatContent::Parts(parts) => {
@@ -739,6 +779,10 @@ fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialec
                 match part {
                     ChatContentPart::ImageUrl { .. } if dialect == DsDialect::V41 => {
                         out.push_str(DEEPSEEK_V41_IMAGE_PLACEHOLDER);
+                        media_order.push(MediaPartSource {
+                            message_index,
+                            content_part_index: index,
+                        });
                     }
                     _ => out.push_str(part.as_text()?),
                 }
diff --git a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs
index 3ccf4d11db7f..33e6eb99d8b6 100644
--- a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs
+++ b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs
@@ -14,6 +14,7 @@ use serde_json::Value;
 use serde_json_fmt::JsonFormat;
 
 use crate::error::{Error, Result};
+use crate::reasoning::ReasoningControl;
 use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatRole, ChatTool};
 use crate::{AssistantContentBlock, AssistantMessageExt, AssistantToolCall};
 
@@ -42,8 +43,11 @@ struct RenderedToolSchema<'a> {
 }
 
 /// Render one chat request into the final prompt string.
-pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
-    let thinking_mode = match request.enable_thinking()?.unwrap_or(false) {
+pub(super) fn render_request(
+    request: &ChatRequest,
+    reasoning: &ReasoningControl,
+) -> Result<String> {
+    let thinking_mode = match reasoning.is_enabled() {
         true => ThinkingMode::Thinking,
         false => ThinkingMode::Chat,
     };
@@ -61,6 +65,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
         find_last_user_render_index(request.messages.as_slice(), render_offset);
     let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice());
     let continue_final_message = request.chat_options.continue_final_message();
+    let add_generation_prompt = request.chat_options.add_generation_prompt();
     let mut prompt = String::from(BOS_TOKEN);
 
     if !initial_tools.is_empty() {
@@ -79,6 +84,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
             thinking_mode,
             drop_thinking,
             continue_final_message,
+            add_generation_prompt,
         )?;
     }
 
@@ -110,11 +116,13 @@ fn render_message(
     thinking_mode: ThinkingMode,
     drop_thinking: bool,
     continue_final_message: bool,
+    add_generation_prompt: bool,
 ) -> Result<()> {
     let render_index = message_index as isize + render_offset;
     let opens_thinking = render_index == last_user_render_index;
     let after_last_user_turn = render_index > last_user_render_index;
     let after_or_at_last_user_turn = render_index >= last_user_render_index;
+    let add_generation_prompt = message_index + 1 != messages.len() || add_generation_prompt;
 
     match message {
         ChatMessage::System { content } => render_system_message(out, Some(content), &[]),
@@ -123,11 +131,13 @@ fn render_message(
             content,
             tools.as_deref().unwrap_or(&[]),
             thinking_mode == ThinkingMode::Thinking && opens_thinking,
+            add_generation_prompt,
         ),
         ChatMessage::User { content } => render_user_message(
             out,
             content,
             thinking_mode == ThinkingMode::Thinking && opens_thinking,
+            add_generation_prompt,
         ),
         ChatMessage::Assistant { content } => render_assistant_message(
             out,
@@ -148,6 +158,7 @@ fn render_message(
             thinking_mode == ThinkingMode::Thinking && after_or_at_last_user_turn,
             content,
         ),
+        ChatMessage::Custom { role, .. } => Err(Error::UnsupportedChatRole { role: role.clone() }),
     }
 }
 
@@ -195,6 +206,7 @@ fn render_developer_message(
     content: &ChatContent,
     tools: &[ChatTool],
     opens_thinking: bool,
+    add_generation_prompt: bool,
 ) -> Result<()> {
     if content.is_empty() {
         return Err(Error::ChatTemplate(
@@ -209,7 +221,7 @@ fn render_developer_message(
     }
     out.push_str("\n\n# The user's message is: ");
     write_chat_content(out, content)?;
-    write_user_like_suffix(out, opens_thinking);
+    write_user_like_suffix(out, opens_thinking, add_generation_prompt);
     Ok(())
 }
 
@@ -219,17 +231,20 @@ fn render_user_message(
     out: &mut String,
     content: &ChatContent,
     opens_thinking: bool,
+    add_generation_prompt: bool,
 ) -> Result<()> {
     out.push_str("<｜User｜>");
     write_chat_content(out, content)?;
-    write_user_like_suffix(out, opens_thinking);
+    write_user_like_suffix(out, opens_thinking, add_generation_prompt);
     Ok(())
 }
 
 /// Shared trailing wrapper used by both real user turns and native developer
 /// turns after their content has already been written.
-// TODO: respect `add_generation_prompt` option
-fn write_user_like_suffix(out: &mut String, opens_thinking: bool) {
+fn write_user_like_suffix(out: &mut String, opens_thinking: bool, add_generation_prompt: bool) {
+    if !add_generation_prompt {
+        return;
+    }
     out.push_str("<｜Assistant｜>");
     if opens_thinking {
         out.push_str(THINKING_START_TOKEN);
diff --git a/rust/src/chat/src/renderer/deepseek_v32/mod.rs b/rust/src/chat/src/renderer/deepseek_v32/mod.rs
index 1a41e14f8fc1..b4ac1bae1155 100644
--- a/rust/src/chat/src/renderer/deepseek_v32/mod.rs
+++ b/rust/src/chat/src/renderer/deepseek_v32/mod.rs
@@ -3,30 +3,41 @@
 
 mod encoding;
 
+use std::collections::HashMap;
+
+use serde_json::Value;
 use vllm_text::Prompt;
 
-use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
+use super::{ChatRenderer, RenderedPrompt};
 use crate::Result;
+use crate::reasoning::ReasoningControl;
 use crate::request::ChatRequest;
 
 /// Dedicated DeepSeek V3.2 renderer.
-#[derive(Debug, Clone, Copy, Default)]
-pub struct DeepSeekV32ChatRenderer;
+#[derive(Debug, Clone, Default)]
+pub struct DeepSeekV32ChatRenderer {
+    default_template_kwargs: HashMap<String, Value>,
+}
 
 impl DeepSeekV32ChatRenderer {
     /// Create the dedicated DeepSeek V3.2 renderer.
-    pub fn new() -> Self {
-        Self
+    pub fn new(default_template_kwargs: HashMap<String, Value>) -> Self {
+        Self {
+            default_template_kwargs,
+        }
     }
 }
 
 impl ChatRenderer for DeepSeekV32ChatRenderer {
     fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
         request.validate()?;
+        let reasoning = ReasoningControl::resolve(request, &self.default_template_kwargs)?
+            .fallback(ReasoningControl::Disabled);
 
         Ok(RenderedPrompt {
-            prompt: Prompt::Text(encoding::render_request(request)?),
-            effective_template_kwargs: request_template_kwargs(request),
+            prompt: Prompt::Text(encoding::render_request(request, &reasoning)?),
+            media_order: None,
+            effective_template_kwargs: reasoning.template_kwargs(request),
         })
     }
 }
diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs
index 3f93e3653efd..dd9ec13ddc97 100644
--- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs
+++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs
@@ -8,6 +8,7 @@ use serde_json::{Value, json};
 use thiserror_ext::AsReport;
 
 use super::DeepSeekV32ChatRenderer;
+use crate::EffortValue;
 use crate::error::Error;
 use crate::event::{AssistantContentBlock, AssistantToolCall};
 use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
@@ -18,7 +19,7 @@ use crate::request::{
 use crate::{ChatRenderer, ChatRole};
 
 fn render_request(request: &ChatRequest) -> String {
-    DeepSeekV32ChatRenderer::new()
+    DeepSeekV32ChatRenderer::default()
         .render(request)
         .unwrap()
         .prompt
@@ -26,8 +27,30 @@ fn render_request(request: &ChatRequest) -> String {
         .expect("deepseek renderer should return text prompt")
 }
 
+#[test]
+fn effort_enables_binary_thinking_and_none_disables_it() {
+    let mut request = ChatRequest::for_test();
+    for (effort, enabled) in [
+        (EffortValue::from("high"), true),
+        (EffortValue::from("custom"), true),
+        (EffortValue::Number(42.into()), true),
+        (EffortValue::from("none"), false),
+    ] {
+        request.chat_options.reasoning_effort = Some(effort);
+        let rendered = DeepSeekV32ChatRenderer::default().render(&request).unwrap();
+        assert_eq!(
+            rendered.prompt.into_text().unwrap().ends_with("<think>"),
+            enabled
+        );
+        assert_eq!(
+            rendered.effective_template_kwargs["enable_thinking"],
+            enabled
+        );
+    }
+}
+
 fn render_result(request: &ChatRequest) -> Result<String, Error> {
-    DeepSeekV32ChatRenderer::new().render(request).map(|rendered| {
+    DeepSeekV32ChatRenderer::default().render(request).map(|rendered| {
         rendered
             .prompt
             .into_text()
@@ -281,6 +304,44 @@ fn assistant_after_last_user_requires_reasoning_or_tool_calls() {
         .assert_eq(&error.to_report_string());
 }
 
+#[test]
+fn last_user_turn_appends_generation_prompt_by_default() {
+    let request = thinking_request(vec![ChatMessage::user("Hello")]);
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜><think>"));
+}
+
+#[test]
+fn last_user_turn_omits_generation_prompt_when_disabled() {
+    let mut request = thinking_request(vec![ChatMessage::user("Hello")]);
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.ends_with("<｜User｜>Hello"));
+    assert!(!rendered.contains("<｜Assistant｜>"));
+    assert!(!rendered.contains("<think>"));
+}
+
+#[test]
+fn user_then_assistant_still_writes_suffix_without_generation_prompt() {
+    let mut request = ChatRequest {
+        messages: vec![
+            ChatMessage::user("Hello"),
+            ChatMessage::assistant_text("Hi there."),
+        ],
+        ..ChatRequest::for_test()
+    };
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    expect!["<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜></think>Hi there.<｜end▁of▁sentence｜>"]
+        .assert_eq(&rendered);
+}
+
 #[test]
 fn continue_final_assistant_omits_final_eos() {
     let mut request = ChatRequest {
@@ -307,7 +368,7 @@ fn render_rejects_multimodal_input() {
         ..ChatRequest::for_test()
     };
 
-    let error = DeepSeekV32ChatRenderer::new().render(&request).unwrap_err();
+    let error = DeepSeekV32ChatRenderer::default().render(&request).unwrap_err();
 
     assert!(matches!(
         error,
diff --git a/rust/src/chat/src/renderer/deepseek_v4/mod.rs b/rust/src/chat/src/renderer/deepseek_v4/mod.rs
index 155a6ee1403b..136938fbc0c8 100644
--- a/rust/src/chat/src/renderer/deepseek_v4/mod.rs
+++ b/rust/src/chat/src/renderer/deepseek_v4/mod.rs
@@ -1,30 +1,47 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+use std::collections::HashMap;
+
+use serde_json::Value;
 use vllm_text::Prompt;
 
 use super::deepseek::{self, DsDialect};
-use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
+use super::{ChatRenderer, RenderedPrompt};
 use crate::Result;
+use crate::reasoning::ReasoningControl;
 use crate::request::ChatRequest;
 
 /// Dedicated DeepSeek V4 renderer.
-#[derive(Debug, Clone, Copy, Default)]
-pub struct DeepSeekV4ChatRenderer;
+#[derive(Debug, Clone, Default)]
+pub struct DeepSeekV4ChatRenderer {
+    default_template_kwargs: HashMap<String, Value>,
+}
 
 impl DeepSeekV4ChatRenderer {
-    pub fn new() -> Self {
-        Self
+    pub fn new(default_template_kwargs: HashMap<String, Value>) -> Self {
+        Self {
+            default_template_kwargs,
+        }
     }
 }
 
 impl ChatRenderer for DeepSeekV4ChatRenderer {
     fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
         request.validate()?;
+        let reasoning = deepseek::resolve_reasoning(
+            ReasoningControl::resolve(request, &self.default_template_kwargs)?,
+            DsDialect::V4,
+        )?;
 
         Ok(RenderedPrompt {
-            prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V4)?),
-            effective_template_kwargs: request_template_kwargs(request),
+            prompt: Prompt::Text(deepseek::render_request(
+                request,
+                DsDialect::V4,
+                &reasoning,
+            )?),
+            media_order: None,
+            effective_template_kwargs: reasoning.template_kwargs(request),
         })
     }
 }
diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs
index a0ded7bd992c..8702e8d73c08 100644
--- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs
+++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs
@@ -5,15 +5,18 @@ use std::path::PathBuf;
 
 use expect_test::{ExpectFile, expect, expect_file};
 use serde_json::Value;
+use thiserror_ext::AsReport as _;
 
 use super::DeepSeekV4ChatRenderer;
 use crate::ChatRenderer;
+use crate::EffortValue;
 use crate::event::{AssistantContentBlock, AssistantToolCall};
+use crate::renderer::deepseek_v41::DeepSeekV41ChatRenderer;
 use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
-use crate::request::{ChatMessage, ChatRequest, ChatTool, GenerationPromptMode, ReasoningEffort};
+use crate::request::{ChatMessage, ChatRequest, ChatTool, GenerationPromptMode};
 
 fn render_request(request: &ChatRequest) -> String {
-    DeepSeekV4ChatRenderer::new()
+    DeepSeekV4ChatRenderer::default()
         .render(request)
         .unwrap()
         .prompt
@@ -30,13 +33,13 @@ fn thinking_request(messages: Vec<ChatMessage>) -> ChatRequest {
         .chat_options
         .template_kwargs
         .insert("thinking".to_string(), Value::Bool(true));
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::Low);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("low"));
     request
 }
 
 fn fixture_request(input_name: &str) -> ChatRequest {
     let mut request = fixture_chat_request(&fixture_path(input_name), deepseek_fixture_options());
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::Low);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("low"));
     request
 }
 
@@ -202,7 +205,7 @@ fn reasoning_effort_max_adds_prefix_when_thinking_is_enabled() {
         .chat_options
         .template_kwargs
         .insert("thinking".to_string(), Value::Bool(true));
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::Max);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("max"));
 
     let rendered = render_request(&request);
 
@@ -225,7 +228,7 @@ fn reasoning_effort_high_adds_0731_high_prefix() {
         .chat_options
         .template_kwargs
         .insert("thinking".to_string(), Value::Bool(true));
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::High);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("high"));
 
     let rendered = render_request(&request);
 
@@ -257,7 +260,7 @@ fn reasoning_effort_xhigh_maps_to_high() {
         messages: vec![ChatMessage::user("solve it")],
         ..ChatRequest::for_test()
     };
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::XHigh);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("xhigh"));
 
     let rendered = render_request(&request);
 
@@ -267,7 +270,7 @@ fn reasoning_effort_xhigh_maps_to_high() {
 
 #[test]
 fn minimal_and_medium_reasoning_effort_map_to_low() {
-    for effort in [ReasoningEffort::Minimal, ReasoningEffort::Medium] {
+    for effort in [EffortValue::from("minimal"), EffortValue::from("medium")] {
         let mut request = ChatRequest {
             messages: vec![ChatMessage::user("solve it")],
             ..ChatRequest::for_test()
@@ -291,7 +294,7 @@ fn reasoning_effort_low_keeps_the_default_prompt() {
         .chat_options
         .template_kwargs
         .insert("thinking".to_string(), Value::Bool(true));
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::Low);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("low"));
 
     let rendered = render_request(&request);
 
@@ -304,11 +307,7 @@ fn reasoning_effort_none_disables_thinking() {
         messages: vec![ChatMessage::user("answer directly")],
         ..ChatRequest::for_test()
     };
-    request
-        .chat_options
-        .template_kwargs
-        .insert("thinking".to_string(), Value::Bool(true));
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::None);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("none"));
 
     let rendered = render_request(&request);
 
@@ -316,6 +315,49 @@ fn reasoning_effort_none_disables_thinking() {
         .assert_eq(&rendered);
 }
 
+#[test]
+fn normalized_reasoning_drives_both_prompt_and_compatibility_kwargs() {
+    for (kwargs, enabled, effort, prefix) in [
+        (
+            serde_json::json!({"enable_thinking": true, "reasoning_effort": "none"}),
+            true,
+            "high",
+            "Absolute maximum",
+        ),
+        (
+            serde_json::json!({"thinking": false, "enable_thinking": true, "reasoning_effort": "max"}),
+            false,
+            "none",
+            "",
+        ),
+        (
+            serde_json::json!({"reasoning_effort": "xhigh"}),
+            true,
+            "high",
+            "Absolute maximum",
+        ),
+    ] {
+        let mut request = ChatRequest::for_test();
+        request.chat_options.template_kwargs = serde_json::from_value(kwargs).unwrap();
+        let rendered = DeepSeekV4ChatRenderer::default().render(&request).unwrap();
+        let prompt = rendered.prompt.into_text().unwrap();
+        assert_eq!(prompt.ends_with("<think>"), enabled);
+        assert_eq!(prompt.contains("Reasoning Effort:"), enabled);
+        if enabled {
+            assert!(prompt.contains(prefix));
+        }
+        assert_eq!(rendered.effective_template_kwargs["thinking"], enabled);
+        assert_eq!(
+            rendered.effective_template_kwargs["enable_thinking"],
+            enabled
+        );
+        assert_eq!(
+            rendered.effective_template_kwargs["reasoning_effort"],
+            effort
+        );
+    }
+}
+
 #[test]
 fn system_to_assistant_transitions_in_chat_and_thinking_modes() {
     let rendered = [false, true].map(|enable_thinking| {
@@ -337,7 +379,7 @@ fn system_to_assistant_transitions_in_chat_and_thinking_modes() {
             .chat_options
             .template_kwargs
             .insert("enable_thinking".to_string(), Value::Bool(enable_thinking));
-        request.chat_options.reasoning_effort = Some(ReasoningEffort::Low);
+        request.chat_options.reasoning_effort = Some(EffortValue::from("low"));
         render_request(&request)
     });
 
@@ -365,7 +407,7 @@ fn trailing_system_transitions_in_chat_and_thinking_modes() {
             .chat_options
             .template_kwargs
             .insert("enable_thinking".to_string(), Value::Bool(enable_thinking));
-        request.chat_options.reasoning_effort = Some(ReasoningEffort::Low);
+        request.chat_options.reasoning_effort = Some(EffortValue::from("low"));
         render_request(&request)
     });
 
@@ -379,7 +421,7 @@ fn trailing_system_transitions_in_chat_and_thinking_modes() {
 }
 
 #[test]
-fn reasoning_effort_template_kwarg_is_ignored() {
+fn reasoning_effort_template_kwarg_selects_effort() {
     let mut request = ChatRequest {
         messages: vec![ChatMessage::user("solve it")],
         ..ChatRequest::for_test()
@@ -395,10 +437,55 @@ fn reasoning_effort_template_kwarg_is_ignored() {
 
     let rendered = render_request(&request);
 
-    assert!(rendered.starts_with("<｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum"));
+    assert!(rendered.starts_with("<｜begin▁of▁sentence｜>Reasoning Effort: Beyond maximum"));
     assert!(rendered.ends_with("<｜Assistant｜><think>"));
 }
 
+#[test]
+fn tool_call_arguments_preserve_non_objects_like_recipe() {
+    let renderers: [(&dyn ChatRenderer, &str); 2] = [
+        (&DeepSeekV4ChatRenderer::default(), "parameter"),
+        (&DeepSeekV41ChatRenderer::default(), " parameter"),
+    ];
+    for arguments in [
+        "not json",
+        "[1, 2]",
+        "null",
+        "42",
+        "true",
+        "",
+        r#""text""#,
+        r#""{\"query\":\"double\"}""#,
+    ] {
+        let request = ChatRequest {
+            messages: vec![
+                ChatMessage::assistant_blocks(vec![AssistantContentBlock::ToolCall(
+                    AssistantToolCall {
+                        id: "call".to_string(),
+                        name: "search".to_string(),
+                        arguments: arguments.to_string(),
+                    },
+                )]),
+                ChatMessage::tool_response("result", "call"),
+            ],
+            ..ChatRequest::for_test()
+        };
+
+        for (renderer, tag) in renderers {
+            let prompt = renderer.render(&request).unwrap().prompt.into_text().unwrap();
+            let parameters: Vec<_> =
+                prompt.lines().filter(|line| line.contains("parameter name=")).collect();
+            assert_eq!(
+                parameters,
+                [format!(
+                    "<｜DSML｜{tag} name=\"arguments\" string=\"true\">{arguments}</｜DSML｜{tag}>"
+                )],
+                "arguments: {arguments:?}, tag: {tag:?}",
+            );
+        }
+    }
+}
+
 #[test]
 fn tool_results_are_sorted_by_previous_assistant_tool_call_order() {
     let mut request = ChatRequest {
@@ -554,7 +641,7 @@ fn drop_thinking_false_keeps_prior_assistant_reasoning() {
         .chat_options
         .template_kwargs
         .insert("drop_thinking".to_string(), Value::Bool(false));
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::Low);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("low"));
 
     let rendered = render_request(&request);
 
@@ -587,3 +674,111 @@ fn continue_final_assistant_omits_final_eos() {
     expect!["<｜begin▁of▁sentence｜><｜User｜>write<｜Assistant｜></think>partial answer"]
         .assert_eq(&rendered);
 }
+
+#[test]
+fn last_user_turn_appends_generation_prompt_by_default() {
+    let request = thinking_request(vec![ChatMessage::user("Hello")]);
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜><think>"));
+}
+
+#[test]
+fn last_user_turn_omits_generation_prompt_when_disabled() {
+    let mut request = thinking_request(vec![ChatMessage::user("Hello")]);
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    assert!(rendered.ends_with("<｜User｜>Hello"));
+    assert!(!rendered.contains("<｜Assistant｜>"));
+    assert!(!rendered.contains("<think>"));
+}
+
+#[test]
+fn user_then_assistant_still_writes_suffix_without_generation_prompt() {
+    let mut request = ChatRequest {
+        messages: vec![
+            ChatMessage::user("Hello"),
+            ChatMessage::assistant_text("Hi there."),
+        ],
+        ..ChatRequest::for_test()
+    };
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+    request
+        .chat_options
+        .template_kwargs
+        .insert("thinking".to_string(), Value::Bool(false));
+
+    let rendered = render_request(&request);
+
+    expect![
+        "<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜></think>Hi there.<｜end▁of▁sentence｜>"
+    ]
+    .assert_eq(&rendered);
+}
+
+#[test]
+fn trailing_system_omits_generation_prompt_when_disabled() {
+    let mut request = ChatRequest {
+        messages: vec![
+            ChatMessage::system("rules"),
+            ChatMessage::user("question"),
+            ChatMessage::system("final rules"),
+        ],
+        ..ChatRequest::for_test()
+    };
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+    request
+        .chat_options
+        .template_kwargs
+        .insert("enable_thinking".to_string(), Value::Bool(false));
+    request.chat_options.reasoning_effort = Some(EffortValue::from("low"));
+
+    let rendered = render_request(&request);
+
+    expect!["<｜begin▁of▁sentence｜>rules<｜User｜>questionfinal rules"].assert_eq(&rendered);
+}
+
+#[test]
+fn last_merged_tool_result_omits_generation_prompt_when_disabled() {
+    let mut request = thinking_request(vec![
+        ChatMessage::user("question"),
+        ChatMessage::tool_response("result", "call-1"),
+    ]);
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    expect!["<｜begin▁of▁sentence｜><｜User｜>question\n\n<tool_result>result</tool_result>"]
+        .assert_eq(&rendered);
+    assert!(!rendered.contains("<｜Assistant｜>"));
+}
+
+#[test]
+fn last_developer_omits_generation_prompt_when_disabled() {
+    let mut request = thinking_request(vec![ChatMessage::developer("latest instruction", None)]);
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    let rendered = render_request(&request);
+
+    expect!["<｜begin▁of▁sentence｜><｜User｜>latest instruction"].assert_eq(&rendered);
+}
+
+#[test]
+fn custom_roles_are_rejected_as_request_errors() {
+    let request = ChatRequest {
+        messages: vec![
+            ChatMessage::custom("root", "identity"),
+            ChatMessage::user("hello"),
+        ],
+        ..ChatRequest::for_test()
+    };
+
+    let error = DeepSeekV4ChatRenderer::default().render(&request).unwrap_err();
+
+    assert!(error.is_request_validation_error());
+    expect!["chat role `root` is not supported by this chat renderer"]
+        .assert_eq(&error.to_report_string());
+}
diff --git a/rust/src/chat/src/renderer/deepseek_v41/mod.rs b/rust/src/chat/src/renderer/deepseek_v41/mod.rs
index 4e5ddb37257e..54544ec3839f 100644
--- a/rust/src/chat/src/renderer/deepseek_v41/mod.rs
+++ b/rust/src/chat/src/renderer/deepseek_v41/mod.rs
@@ -1,30 +1,45 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+use std::collections::HashMap;
+
+use serde_json::Value;
 use vllm_text::Prompt;
 
 use super::deepseek::{self, DsDialect};
-use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
+use super::{ChatRenderer, RenderedPrompt};
 use crate::Result;
+use crate::reasoning::ReasoningControl;
 use crate::request::ChatRequest;
 
 /// Dedicated DeepSeek V4.1 renderer.
-#[derive(Debug, Clone, Copy, Default)]
-pub struct DeepSeekV41ChatRenderer;
+#[derive(Debug, Clone, Default)]
+pub struct DeepSeekV41ChatRenderer {
+    default_template_kwargs: HashMap<String, Value>,
+}
 
 impl DeepSeekV41ChatRenderer {
-    pub fn new() -> Self {
-        Self
+    pub fn new(default_template_kwargs: HashMap<String, Value>) -> Self {
+        Self {
+            default_template_kwargs,
+        }
     }
 }
 
 impl ChatRenderer for DeepSeekV41ChatRenderer {
     fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
         request.validate()?;
+        let reasoning = deepseek::resolve_reasoning(
+            ReasoningControl::resolve(request, &self.default_template_kwargs)?,
+            DsDialect::V41,
+        )?;
+        let (prompt, media_order) =
+            deepseek::render_request_with_media_order(request, DsDialect::V41, &reasoning)?;
 
         Ok(RenderedPrompt {
-            prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V41)?),
-            effective_template_kwargs: request_template_kwargs(request),
+            prompt: Prompt::Text(prompt),
+            media_order: Some(media_order),
+            effective_template_kwargs: reasoning.template_kwargs(request),
         })
     }
 }
diff --git a/rust/src/chat/src/renderer/deepseek_v41/tests.rs b/rust/src/chat/src/renderer/deepseek_v41/tests.rs
index d7b0c538bb54..62cf6d20247a 100644
--- a/rust/src/chat/src/renderer/deepseek_v41/tests.rs
+++ b/rust/src/chat/src/renderer/deepseek_v41/tests.rs
@@ -8,12 +8,15 @@ use serde_json::{Value, json};
 
 use super::DeepSeekV41ChatRenderer;
 use crate::ChatRenderer;
-use crate::event::AssistantContentBlock;
+use crate::EffortValue;
+use crate::event::{AssistantContentBlock, AssistantToolCall};
 use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
-use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort};
+use crate::request::{
+    ChatContent, ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode,
+};
 
 fn render(request: &ChatRequest) -> String {
-    DeepSeekV41ChatRenderer::new()
+    DeepSeekV41ChatRenderer::default()
         .render(request)
         .unwrap()
         .prompt
@@ -62,7 +65,7 @@ fn inlines_image_placeholder_at_image_part_positions() {
     // The placeholder is unconditional (matching the Python encoding's
     // `IMAGE_PLACEHOLDER`); requests to a multimodal-less backend fail later,
     // at media preparation.
-    let renderer = DeepSeekV41ChatRenderer::new();
+    let renderer = DeepSeekV41ChatRenderer::default();
     let request = ChatRequest {
         messages: vec![ChatMessage::User {
             content: ChatContent::Parts(vec![
@@ -83,14 +86,91 @@ fn inlines_image_placeholder_at_image_part_positions() {
     );
 }
 
+#[test]
+fn media_order_follows_reordered_tool_results() {
+    let renderer = DeepSeekV41ChatRenderer::default();
+    let request = ChatRequest {
+        messages: vec![
+            ChatMessage::assistant_blocks(vec![
+                AssistantContentBlock::ToolCall(AssistantToolCall {
+                    id: "call-a".to_string(),
+                    name: "tool-a".to_string(),
+                    arguments: "{}".to_string(),
+                }),
+                AssistantContentBlock::ToolCall(AssistantToolCall {
+                    id: "call-b".to_string(),
+                    name: "tool-b".to_string(),
+                    arguments: "{}".to_string(),
+                }),
+            ]),
+            ChatMessage::tool_response(
+                vec![ChatContentPart::text("result-b"), image_part("image-b")],
+                "call-b",
+            ),
+            ChatMessage::tool_response(
+                vec![ChatContentPart::text("result-a"), image_part("image-a")],
+                "call-a",
+            ),
+        ],
+        ..ChatRequest::for_test()
+    };
+
+    let rendered = renderer.render(&request).unwrap();
+    let media_order = rendered.media_order.unwrap();
+    let prompt = rendered.prompt.into_text().unwrap();
+
+    assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap());
+    expect![[r#"
+        [
+            MediaPartSource {
+                message_index: 2,
+                content_part_index: 1,
+            },
+            MediaPartSource {
+                message_index: 1,
+                content_part_index: 1,
+            },
+        ]
+    "#]]
+    .assert_debug_eq(&media_order);
+}
+
+#[test]
+fn media_order_excludes_dropped_historical_developer_content() {
+    let renderer = DeepSeekV41ChatRenderer::default();
+    let request = ChatRequest {
+        messages: vec![
+            ChatMessage::developer(vec![image_part("historical-image")], None),
+            ChatMessage::user("question"),
+        ],
+        ..ChatRequest::for_test()
+    };
+
+    let rendered = renderer.render(&request).unwrap();
+    let media_order = rendered.media_order.unwrap();
+    let prompt = rendered.prompt.into_text().unwrap();
+
+    assert!(!prompt.contains("<｜deepseek_image｜>"));
+    assert!(media_order.is_empty());
+}
+
+fn image_part(uuid: &str) -> ChatContentPart {
+    ChatContentPart::ImageUrl {
+        image_url: format!("data:image/png;base64,{uuid}"),
+        detail: None,
+        uuid: Some(uuid.to_string()),
+    }
+}
+
 #[test]
 fn maps_reasoning_effort_to_reference_numeric_budget() {
     for (effort, budget) in [
         (None, 50),
-        (Some(ReasoningEffort::Low), 25),
-        (Some(ReasoningEffort::High), 50),
-        (Some(ReasoningEffort::XHigh), 75),
-        (Some(ReasoningEffort::Max), 100),
+        (Some(EffortValue::from("low")), 25),
+        (Some(EffortValue::from("high")), 50),
+        (Some(EffortValue::from("xhigh")), 75),
+        (Some(EffortValue::from("max")), 100),
+        (Some(EffortValue::Number(37.into())), 37),
     ] {
         let mut request = request();
         request.chat_options.reasoning_effort = effort;
@@ -115,11 +195,43 @@ fn accepts_numeric_template_effort_with_top_level_precedence() {
         request.chat_options.template_kwargs.insert("reasoning_effort".into(), effort);
         assert!(render(&request).contains(&format!("Reasoning Effort: {budget} (range 1-100")));
 
-        request.chat_options.reasoning_effort = Some(ReasoningEffort::XHigh);
+        request.chat_options.reasoning_effort = Some(EffortValue::from("xhigh"));
         assert!(render(&request).contains("Reasoning Effort: 75 (range 1-100"));
     }
 }
 
+#[test]
+fn enabling_with_none_inherits_deployment_effort_and_reports_numeric_control() {
+    let renderer =
+        DeepSeekV41ChatRenderer::new([("reasoning_effort".to_string(), json!(37))].into());
+    let mut request = request();
+    request.chat_options.reasoning_effort = Some(EffortValue::from("none"));
+    request
+        .chat_options
+        .template_kwargs
+        .insert("enable_thinking".into(), json!(true));
+    let rendered = renderer.render(&request).unwrap();
+    let prompt = rendered.prompt.into_text().unwrap();
+    assert!(prompt.contains("Reasoning Effort: 37 (range 1-100"));
+    assert!(prompt.ends_with("<think>"));
+    assert_eq!(rendered.effective_template_kwargs["reasoning_effort"], 37);
+    assert_eq!(rendered.effective_template_kwargs["enable_thinking"], true);
+
+    request.chat_options.template_kwargs.insert("thinking".into(), json!(false));
+    request
+        .chat_options
+        .template_kwargs
+        .insert("reasoning_effort".into(), json!(101));
+    request.chat_options.reasoning_effort = None;
+    let rendered = renderer.render(&request).unwrap();
+    assert!(rendered.prompt.into_text().unwrap().ends_with("</think>"));
+    assert_eq!(
+        rendered.effective_template_kwargs["reasoning_effort"],
+        "none"
+    );
+    assert_eq!(rendered.effective_template_kwargs["enable_thinking"], false);
+}
+
 #[test]
 fn rejects_undefined_effort_names_and_invalid_numeric_budgets() {
     for effort in [
@@ -133,13 +245,13 @@ fn rejects_undefined_effort_names_and_invalid_numeric_budgets() {
     ] {
         let mut request = request();
         request.chat_options.template_kwargs.insert("reasoning_effort".into(), effort);
-        let error = DeepSeekV41ChatRenderer::new().render(&request).unwrap_err();
+        let error = DeepSeekV41ChatRenderer::default().render(&request).unwrap_err();
         assert!(error.is_request_validation_error());
     }
-    for effort in [ReasoningEffort::Minimal, ReasoningEffort::Medium] {
+    for effort in [EffortValue::from("minimal"), EffortValue::from("medium")] {
         let mut request = request();
         request.chat_options.reasoning_effort = Some(effort);
-        let error = DeepSeekV41ChatRenderer::new().render(&request).unwrap_err();
+        let error = DeepSeekV41ChatRenderer::default().render(&request).unwrap_err();
         assert!(error.is_request_validation_error());
     }
 }
@@ -151,8 +263,8 @@ fn chat_mode_and_none_effort_use_closed_thinking_prefix() {
     let expected = expect!["<｜begin▁of▁sentence｜><｜User｜>question<｜Assistant｜></think>"];
     expected.assert_eq(&render(&request));
 
-    request.chat_options.template_kwargs.insert("thinking".into(), json!(true));
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::None);
+    request.chat_options.template_kwargs.remove("thinking");
+    request.chat_options.reasoning_effort = Some(EffortValue::from("none"));
     expected.assert_eq(&render(&request));
 }
 
@@ -205,6 +317,49 @@ fn mid_system_advances_last_user_boundary_and_drops_old_reasoning() {
         .assert_eq(&render(&request));
 }
 
+#[test]
+fn last_user_turn_omits_generation_prompt_when_disabled() {
+    let mut request = request();
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+    expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 50 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question"]
+        .assert_eq(&render(&request));
+}
+
+#[test]
+fn last_mid_system_omits_generation_prompt_when_disabled() {
+    let mut request = ChatRequest {
+        messages: vec![
+            ChatMessage::system("sys"),
+            ChatMessage::user("q1"),
+            ChatMessage::assistant_blocks(vec![
+                AssistantContentBlock::Reasoning { text: "r1".into() },
+                AssistantContentBlock::Text { text: "a1".into() },
+            ]),
+            ChatMessage::developer("old policy", None),
+            ChatMessage::system("mid sys"),
+        ],
+        ..ChatRequest::for_test()
+    };
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+    request
+        .chat_options
+        .template_kwargs
+        .insert("reasoning_effort".into(), json!(88));
+
+    expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 88 (range 1-100, the higher the value, the more thorough the reasoning)\n\nsys<｜User｜>q1<｜Assistant｜></think>a1<｜end▁of▁sentence｜><｜System｜>mid sys"]
+        .assert_eq(&render(&request));
+}
+
+#[test]
+fn chat_mode_omits_generation_prompt_when_disabled() {
+    let mut request = request();
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+    request.chat_options.template_kwargs.insert("thinking".into(), json!(false));
+
+    expect!["<｜begin▁of▁sentence｜><｜User｜>question"].assert_eq(&render(&request));
+}
+
 #[test]
 fn joins_text_parts_with_reference_separator_and_preserves_literal_tags() {
     let mut request = ChatRequest {
diff --git a/rust/src/chat/src/renderer/harmony/mod.rs b/rust/src/chat/src/renderer/harmony/mod.rs
index b07d8458d739..d8dafadca74a 100644
--- a/rust/src/chat/src/renderer/harmony/mod.rs
+++ b/rust/src/chat/src/renderer/harmony/mod.rs
@@ -5,21 +5,26 @@
 
 pub(crate) mod encoding;
 
+use std::collections::HashMap;
+
 use openai_harmony::HarmonyEncoding;
 use openai_harmony::chat::{
     Author, Conversation, DeveloperContent, Message, ReasoningEffort as HarmonyReasoningEffort,
     Role, SystemContent, ToolDescription,
 };
+use serde_json::Value;
 use thiserror_ext::AsReport as _;
 use time::macros::format_description;
 use vllm_text::Prompt;
 
 use self::encoding::harmony_encoding;
-use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
+use super::{ChatRenderer, RenderedPrompt};
+use crate::AssistantMessageExt as _;
+use crate::EffortValue;
 use crate::error::{Error, Result};
 use crate::event::AssistantContentBlock;
+use crate::reasoning::ReasoningControl;
 use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatTool, GenerationPromptMode};
-use crate::{AssistantMessageExt as _, ReasoningEffort};
 
 const SYSTEM_START_DATE_ENV: &str = "VLLM_SYSTEM_START_DATE";
 const HARMONY_SYSTEM_INSTRUCTIONS_ENV: &str = "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS";
@@ -28,6 +33,7 @@ const HARMONY_SYSTEM_INSTRUCTIONS_ENV: &str = "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTR
 pub struct HarmonyChatRenderer {
     encoding: &'static HarmonyEncoding,
     options: Options,
+    default_template_kwargs: HashMap<String, Value>,
 }
 
 struct Options {
@@ -44,8 +50,9 @@ impl HarmonyChatRenderer {
     ///   unset, the renderer uses the current local date with a UTC fallback.
     /// - `VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS` moves leading instructions
     ///   into the system model identity when set to a non-zero integer.
-    pub fn new() -> Result<Self> {
+    pub fn new(default_template_kwargs: HashMap<String, Value>) -> Result<Self> {
         Self::with_options(
+            default_template_kwargs,
             env_system_start_date(),
             env_use_harmony_system_instructions(),
         )
@@ -57,11 +64,13 @@ impl HarmonyChatRenderer {
     /// Production code should call [`Self::new`] so the renderer observes the
     /// same environment contract as the Python Harmony path.
     pub fn with_options(
+        default_template_kwargs: HashMap<String, Value>,
         system_start_date: impl Into<String>,
         use_system_instructions: bool,
     ) -> Result<Self> {
         Ok(Self {
             encoding: harmony_encoding()?,
+            default_template_kwargs,
             options: Options {
                 system_start_date: system_start_date.into(),
                 use_system_instructions,
@@ -73,7 +82,11 @@ impl HarmonyChatRenderer {
     ///
     /// Harmony owns both prompt formatting and tokenization, so the Rust
     /// frontend bypasses the generic HF tokenizer path for GPT-OSS input.
-    fn render_token_ids(&self, request: &ChatRequest) -> Result<Vec<u32>> {
+    fn render_token_ids(
+        &self,
+        request: &ChatRequest,
+        effort: HarmonyReasoningEffort,
+    ) -> Result<Vec<u32>> {
         if request.has_multimodal() {
             return Err(Error::UnsupportedMultimodalContent("image_url"));
         }
@@ -86,7 +99,8 @@ impl HarmonyChatRenderer {
             ));
         }
 
-        let messages = auto_drop_analysis_messages(to_harmony_messages(request, &self.options)?);
+        let messages =
+            auto_drop_analysis_messages(to_harmony_messages(request, &self.options, effort)?);
         let conversation = Conversation::from_messages(messages);
         // Pass `None` so oss-harmony does not apply its narrower built-in
         // analysis-drop policy after the Rust-side Python-parity cleanup above.
@@ -114,9 +128,13 @@ impl ChatRenderer for HarmonyChatRenderer {
     /// Render a chat request as [`Prompt::TokenIds`] with template kwargs echoed
     /// for downstream accounting/debugging.
     fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
+        let reasoning = ReasoningControl::resolve(request, &self.default_template_kwargs)?
+            .fallback(ReasoningControl::enabled("medium"));
+        let effort = to_harmony_reasoning_effort(reasoning.effort())?;
         Ok(RenderedPrompt {
-            prompt: Prompt::TokenIds(self.render_token_ids(request)?),
-            effective_template_kwargs: request_template_kwargs(request),
+            prompt: Prompt::TokenIds(self.render_token_ids(request, effort)?),
+            media_order: None,
+            effective_template_kwargs: reasoning.template_kwargs(request),
         })
     }
 }
@@ -126,12 +144,21 @@ impl ChatRenderer for HarmonyChatRenderer {
 /// This adds the Harmony system/developer preamble, peels at most one leading
 /// system/developer instruction message, and then lowers the remaining chat
 /// history message-by-message.
-fn to_harmony_messages(request: &ChatRequest, options: &Options) -> Result<Vec<Message>> {
+fn to_harmony_messages(
+    request: &ChatRequest,
+    options: &Options,
+    effort: HarmonyReasoningEffort,
+) -> Result<Vec<Message>> {
     let (instructions, leading_developer_tools, remaining_messages) =
         peel_leading_instructions(&request.messages)?;
     let tool_call_names = tool_call_names(&request.messages);
-    let mut messages =
-        build_harmony_preamble(request, instructions, leading_developer_tools, options)?;
+    let mut messages = build_harmony_preamble(
+        request,
+        instructions,
+        leading_developer_tools,
+        options,
+        effort,
+    )?;
 
     for message in remaining_messages {
         messages.extend(to_harmony_message(message, &tool_call_names, options)?);
@@ -161,7 +188,8 @@ fn peel_leading_instructions(
         )),
         ChatMessage::User { .. }
         | ChatMessage::Assistant { .. }
-        | ChatMessage::ToolResponse { .. } => Ok((None, None, messages)),
+        | ChatMessage::ToolResponse { .. }
+        | ChatMessage::Custom { .. } => Ok((None, None, messages)),
     }
 }
 
@@ -176,12 +204,13 @@ fn build_harmony_preamble(
     instructions: Option<String>,
     leading_developer_tools: Option<&[ChatTool]>,
     options: &Options,
+    effort: HarmonyReasoningEffort,
 ) -> Result<Vec<Message>> {
     let mut messages = vec![Message::from_role_and_content(
         Role::System,
         system_content(
             instructions.as_deref().filter(|_| options.use_system_instructions),
-            request.chat_options.reasoning_effort,
+            Some(effort),
             &options.system_start_date,
         )?,
     )];
@@ -226,18 +255,18 @@ fn preamble_tool_descriptions(
 
 /// Construct the Harmony system content for the request preamble.
 ///
-/// Harmony defaults the reasoning effort to `medium` when none is provided, so
-/// this only sets an explicit effort after validating vLLM's request value.
+/// The preamble receives normalized effort; later system messages use Harmony's
+/// own default when their effort is omitted.
 fn system_content(
     instructions: Option<&str>,
-    reasoning_effort: Option<ReasoningEffort>,
+    reasoning_effort: Option<HarmonyReasoningEffort>,
     system_start_date: &str,
 ) -> Result<SystemContent> {
     let mut content =
         SystemContent::new().with_conversation_start_date(system_start_date.to_string());
 
     if let Some(reasoning_effort) = reasoning_effort {
-        content = content.with_reasoning_effort(to_harmony_reasoning_effort(reasoning_effort)?);
+        content = content.with_reasoning_effort(reasoning_effort);
     }
 
     if let Some(instructions) = instructions.filter(|text| !text.is_empty()) {
@@ -300,6 +329,9 @@ fn to_harmony_message(
                 .with_recipient("assistant"),
             ]
         }
+        ChatMessage::Custom { role, .. } => {
+            return Err(Error::UnsupportedChatRole { role: role.clone() });
+        }
     })
 }
 
@@ -444,20 +476,17 @@ fn to_tool_descriptions(tools: &[ChatTool]) -> Vec<ToolDescription> {
         .collect()
 }
 
-/// Map supported OpenAI reasoning-effort values onto Harmony's enum.
+/// Map supported reasoning-effort values onto Harmony's enum.
 fn to_harmony_reasoning_effort(
-    reasoning_effort: ReasoningEffort,
+    reasoning_effort: Option<&EffortValue>,
 ) -> Result<HarmonyReasoningEffort> {
-    match reasoning_effort {
-        ReasoningEffort::Low => Ok(HarmonyReasoningEffort::Low),
-        ReasoningEffort::Medium => Ok(HarmonyReasoningEffort::Medium),
-        ReasoningEffort::High => Ok(HarmonyReasoningEffort::High),
-        ReasoningEffort::None
-        | ReasoningEffort::Minimal
-        | ReasoningEffort::XHigh
-        | ReasoningEffort::Max => Err(Error::ChatTemplate(format!(
-            "reasoning_effort={:?} is not supported by Harmony. Supported values are: low, medium, high.",
-            reasoning_effort.as_str()
+    match reasoning_effort.and_then(EffortValue::as_str) {
+        Some("low") => Ok(HarmonyReasoningEffort::Low),
+        Some("medium") => Ok(HarmonyReasoningEffort::Medium),
+        Some("high") => Ok(HarmonyReasoningEffort::High),
+        _ => Err(Error::InvalidReasoningEffort(format!(
+            "reasoning_effort={} is not supported by Harmony. Supported values are: low, medium, high.",
+            reasoning_effort.unwrap_or(&EffortValue::from("none"))
         ))),
     }
 }
diff --git a/rust/src/chat/src/renderer/harmony/tests.rs b/rust/src/chat/src/renderer/harmony/tests.rs
index 4849c932233b..609c1f115920 100644
--- a/rust/src/chat/src/renderer/harmony/tests.rs
+++ b/rust/src/chat/src/renderer/harmony/tests.rs
@@ -9,12 +9,11 @@ use thiserror_ext::AsReport as _;
 use super::HarmonyChatRenderer;
 use super::encoding::harmony_encoding;
 use crate::ChatRenderer;
+use crate::EffortValue;
 use crate::error::Error;
 use crate::event::{AssistantContentBlock, AssistantToolCall};
 use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
-use crate::request::{
-    ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort,
-};
+use crate::request::{ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode};
 
 const PINNED_DATE: &str = "2025-06-28";
 
@@ -22,7 +21,7 @@ fn fixture_request(input_name: &str) -> ChatRequest {
     fixture_chat_request(
         &fixture_path(input_name),
         FixtureRequestOptions {
-            enable_thinking: Some(false),
+            enable_thinking: None,
             no_generation_prompt_when_last_assistant: false,
         },
     )
@@ -36,7 +35,8 @@ fn fixture_path(name: &str) -> PathBuf {
 }
 
 fn test_renderer(use_system_instructions: bool) -> HarmonyChatRenderer {
-    HarmonyChatRenderer::with_options(PINNED_DATE, use_system_instructions).unwrap()
+    HarmonyChatRenderer::with_options(Default::default(), PINNED_DATE, use_system_instructions)
+        .unwrap()
 }
 
 fn render_token_ids(request: &ChatRequest) -> Vec<u32> {
@@ -138,14 +138,48 @@ fn drops_stale_analysis_fixture() {
 #[test]
 fn rejects_invalid_reasoning_effort() {
     let mut request = ChatRequest::for_test();
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::None);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("none"));
 
     let error = test_renderer(false).render(&request).unwrap_err();
 
-    expect![[r#"chat template error: reasoning_effort="none" is not supported by Harmony. Supported values are: low, medium, high."#]]
+    expect![[r#"reasoning_effort="none" is not supported by Harmony. Supported values are: low, medium, high."#]]
         .assert_eq(&error.to_report_string());
 }
 
+#[test]
+fn normalized_reasoning_respects_harmony_capability_and_deployment_effort() {
+    let renderer = HarmonyChatRenderer::with_options(
+        [("reasoning_effort".to_string(), serde_json::json!("low"))].into(),
+        PINNED_DATE,
+        false,
+    )
+    .unwrap();
+    let mut request = ChatRequest::for_test();
+    request.chat_options.reasoning_effort = Some(EffortValue::from("none"));
+    request
+        .chat_options
+        .template_kwargs
+        .insert("enable_thinking".into(), serde_json::json!(true));
+    let rendered = renderer.render(&request).unwrap();
+    let prompt = harmony_encoding()
+        .unwrap()
+        .tokenizer()
+        .decode_utf8(rendered.prompt.into_token_ids().unwrap())
+        .unwrap();
+    assert!(prompt.contains("Reasoning: low"));
+    assert_eq!(
+        rendered.effective_template_kwargs["reasoning_effort"],
+        "low"
+    );
+    assert_eq!(rendered.effective_template_kwargs["enable_thinking"], true);
+
+    request
+        .chat_options
+        .template_kwargs
+        .insert("thinking".into(), serde_json::json!(false));
+    assert!(renderer.render(&request).is_err());
+}
+
 #[test]
 fn rejects_unknown_tool_response_id() {
     let request = ChatRequest {
diff --git a/rust/src/chat/src/renderer/hf/generation.rs b/rust/src/chat/src/renderer/hf/generation.rs
new file mode 100644
index 000000000000..39e61b5fb6e6
--- /dev/null
+++ b/rust/src/chat/src/renderer/hf/generation.rs
@@ -0,0 +1,187 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Support Hugging Face's `{% generation %}` blocks in MiniJinja chat templates.
+//!
+//! Transformers uses these blocks to track assistant-token masks via scoped call
+//! blocks. Inference only needs their rendered text. MiniJinja lacks this tag, so
+//! we rewrite `{% generation %}...{% endgeneration %}` to
+//! `{% call __hf_generation() %}...{% endcall %}` before compiling the template.
+//! The registered `__hf_generation` callback renders and returns `caller()`,
+//! preserving the block's local scope just as Transformers' AssistantTracker does.
+//!
+//! We use MiniJinja's lexer to identify statement tags and replace only their
+//! keyword spans. This preserves delimiters and whitespace controls (`-` / `+`),
+//! while keeping literals, comments, and raw blocks intact. Original `generation`
+//! and `call` pairings are checked before rewriting both to call blocks.
+
+use minijinja::machinery::{Token, tokenize};
+use minijinja::value::{Kwargs, Value};
+use minijinja::{Error, ErrorKind, State};
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum CallBlockKind {
+    Call,
+    Generation,
+}
+
+/// Rewrite Hugging Face's assistant-mask blocks as scoped MiniJinja call blocks.
+pub(super) fn rewrite_generation_blocks(source: String) -> Result<String, Error> {
+    let tokens = tokenize(&source, false, Default::default(), Default::default())
+        .collect::<Result<Vec<_>, _>>()?;
+    // Track native calls too: rewriting makes both block kinds use `endcall`,
+    // so validate the original pairing before that distinction is lost.
+    let mut blocks = Vec::new();
+    let mut replacements = Vec::new();
+
+    for (index, pair) in tokens.windows(2).enumerate() {
+        let [(Token::BlockStart, _), (Token::Ident(name), span)] = pair else {
+            continue;
+        };
+        match *name {
+            "call" => blocks.push(CallBlockKind::Call),
+            "endcall" if blocks.pop() == Some(CallBlockKind::Generation) => {
+                return Err(Error::new(
+                    ErrorKind::SyntaxError,
+                    "expected endgeneration, got endcall",
+                ));
+            }
+            "generation" | "endgeneration"
+                if matches!(tokens.get(index + 2), Some((Token::BlockEnd, _))) =>
+            {
+                let replacement = if *name == "generation" {
+                    blocks.push(CallBlockKind::Generation);
+                    "call __hf_generation()"
+                } else {
+                    if blocks.pop() != Some(CallBlockKind::Generation) {
+                        return Err(Error::new(
+                            ErrorKind::SyntaxError,
+                            "unexpected endgeneration",
+                        ));
+                    }
+                    "endcall"
+                };
+                replacements.push((
+                    span.start_offset as usize..span.end_offset as usize,
+                    replacement,
+                ));
+            }
+            _ => {}
+        }
+    }
+    if blocks.contains(&CallBlockKind::Generation) {
+        return Err(Error::new(
+            ErrorKind::SyntaxError,
+            "unclosed generation block",
+        ));
+    }
+
+    // Replace only keyword spans, preserving delimiters and their whitespace controls.
+    let mut source = source;
+    for (span, replacement) in replacements.into_iter().rev() {
+        source.replace_range(span, replacement);
+    }
+    Ok(source)
+}
+
+/// Render the caller with the same local scope as Transformers' AssistantTracker.
+pub(super) fn render_generation(state: &State, kwargs: Kwargs) -> Result<Value, Error> {
+    let caller: Value = kwargs.get("caller")?;
+    kwargs.assert_all_used()?;
+    caller.call(state, &[])
+}
+
+#[cfg(test)]
+mod tests {
+    use expect_test::expect;
+    use minijinja::Environment;
+
+    use super::*;
+
+    #[test]
+    fn generation_blocks_reject_invalid_syntax_and_mismatched_end_tags() {
+        let env = Environment::new();
+        for template in [
+            "{% generation %}unclosed",
+            "{% endgeneration %}",
+            "{% generation extra %}{% endgeneration %}",
+            "{% generation %}{% endgeneration extra %}",
+            "{% generation %}{% endcall %}",
+            "{% call f() %}{% endgeneration %}",
+            "{% generation %}{% endcall %}{% call f() %}{% endgeneration %}",
+            "{% generation %}{% if true %}{% endgeneration %}{% endif %}",
+        ] {
+            let result = rewrite_generation_blocks(template.to_string())
+                .and_then(|source| env.template_from_str(&source).map(|_| ()));
+            assert!(result.is_err(), "accepted invalid template: {template}");
+        }
+    }
+
+    #[test]
+    fn rewrite_generation_blocks_preserves_hf_rendering_semantics() {
+        let cases = [
+            (
+                "whitespace",
+                "前  {%- generation -%}  中  {%- endgeneration -%}  后",
+            ),
+            (
+                "keep_newlines",
+                "A{%+ generation +%}\nB{%+ endgeneration +%}\nC",
+            ),
+            ("literal", "{{ '{% generation %}' }}"),
+            ("raw", "{% raw %}{%- generation -%}{% endraw %}"),
+            ("comment", "A{# {% generation %} #}B"),
+            (
+                "scope",
+                "{% set x = 'outer' %}{% generation %}{% set x = 'inner' %}{{ x }}{% endgeneration %}{{ x }}",
+            ),
+            (
+                "nested_calls",
+                "{% macro wrap() %}[{{ caller() }}]{% endmacro %}{% generation %}A{% call wrap() %}{% generation %}B{% endgeneration %}{% endcall %}C{% endgeneration %}",
+            ),
+        ];
+        let results = cases.map(|(name, source)| {
+            let template = rewrite_generation_blocks(source.to_string()).unwrap();
+            let mut env = Environment::new();
+            env.set_trim_blocks(true);
+            env.set_lstrip_blocks(true);
+            env.add_function("__hf_generation", render_generation);
+            let rendered =
+                env.template_from_str(&template).unwrap().render(Value::UNDEFINED).unwrap();
+            (name, rendered)
+        });
+        expect![[r#"
+            [
+                (
+                    "whitespace",
+                    "前中后",
+                ),
+                (
+                    "keep_newlines",
+                    "A\nB\nC",
+                ),
+                (
+                    "literal",
+                    "{% generation %}",
+                ),
+                (
+                    "raw",
+                    "{%- generation -%}",
+                ),
+                (
+                    "comment",
+                    "AB",
+                ),
+                (
+                    "scope",
+                    "innerouter",
+                ),
+                (
+                    "nested_calls",
+                    "A[B]C",
+                ),
+            ]
+        "#]]
+        .assert_debug_eq(&results);
+    }
+}
diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs
index 7461acf2955a..f406e32da64b 100644
--- a/rust/src/chat/src/renderer/hf/mod.rs
+++ b/rust/src/chat/src/renderer/hf/mod.rs
@@ -16,7 +16,7 @@ use self::format::{
     ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption,
 };
 use self::template::{CompiledChatTemplate, TemplateContext};
-use super::{ChatRenderer, RenderedPrompt, effective_template_kwargs};
+use super::{ChatRenderer, RenderedPrompt};
 use crate::error::Result;
 use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
 use crate::{
@@ -25,6 +25,7 @@ use crate::{
 
 mod error;
 mod format;
+mod generation;
 mod template;
 mod tojson;
 
@@ -32,6 +33,31 @@ pub use template::{load_chat_template, resolve_chat_template};
 
 pub use self::format::ChatTemplateContentFormatOption;
 
+/// Extract the effective chat-template kwargs visible to the renderer from the request,
+/// using the provided defaults as the base.
+fn effective_template_kwargs(
+    default_template_kwargs: &HashMap<String, JsonValue>,
+    request: &ChatRequest,
+) -> HashMap<String, JsonValue> {
+    let mut kwargs = default_template_kwargs.clone();
+    kwargs.extend(request.chat_options.template_kwargs.clone());
+
+    if let Some(reasoning_effort) = &request.chat_options.reasoning_effort {
+        kwargs.insert(
+            "reasoning_effort".to_string(),
+            serde_json::json!(reasoning_effort),
+        );
+        if !request.chat_options.template_kwargs.contains_key("enable_thinking") {
+            kwargs.insert(
+                "enable_thinking".to_string(),
+                serde_json::json!(reasoning_effort.as_str() != Some("none")),
+            );
+        }
+    }
+
+    kwargs
+}
+
 /// Template-visible placeholder tokens per supported modality.
 ///
 /// A `None` token means the loaded model does not support that modality, and
@@ -216,6 +242,7 @@ impl HfChatRenderer {
 
         Ok(RenderedPrompt {
             prompt: Prompt::Text(prompt),
+            media_order: None,
             effective_template_kwargs,
         })
     }
@@ -231,8 +258,8 @@ impl ChatRenderer for HfChatRenderer {
 // TODO: borrow more fields directly from the original `ChatMessage`.
 #[serde_with::skip_serializing_none]
 #[derive(Debug, Serialize)]
-struct TemplateMessage {
-    role: &'static str,
+struct TemplateMessage<'a> {
+    role: &'a str,
     content: TemplateContent,
     // Developer-role messages may provide message-local tools in the same shape
     // as top-level request tools.
@@ -295,22 +322,22 @@ struct TemplateToolDefinition {
 }
 
 /// Convert chat messages into the JSON shape expected by Jinja chat templates.
-fn to_template_messages(
-    messages: &[ChatMessage],
+fn to_template_messages<'a>(
+    messages: &'a [ChatMessage],
     content_format: ChatTemplateContentFormat,
     multimodal: Option<&MultimodalRenderInfo>,
-) -> Result<Vec<TemplateMessage>> {
+) -> Result<Vec<TemplateMessage<'a>>> {
     messages
         .iter()
         .map(|message| to_template_message(message, content_format, multimodal))
         .collect()
 }
 
-fn to_template_message(
-    message: &ChatMessage,
+fn to_template_message<'a>(
+    message: &'a ChatMessage,
     content_format: ChatTemplateContentFormat,
     multimodal: Option<&MultimodalRenderInfo>,
-) -> Result<TemplateMessage> {
+) -> Result<TemplateMessage<'a>> {
     Ok(match message {
         ChatMessage::System { content } => TemplateMessage {
             role: "system",
@@ -367,6 +394,15 @@ fn to_template_message(
             tool_calls: None,
             tool_call_id: Some(tool_call_id.clone()),
         },
+        ChatMessage::Custom { role, content } => TemplateMessage {
+            role,
+            content: to_template_content(content, content_format, multimodal)?,
+            tools: None,
+            reasoning: None,
+            reasoning_content: None,
+            tool_calls: None,
+            tool_call_id: None,
+        },
     })
 }
 
@@ -497,7 +533,7 @@ const CONTINUE_FINAL_MESSAGE_TAG: &str = "CONTINUE_FINAL_MESSAGE_TAG ";
 // TODO: transformers v5 also allows continuing a non-`content` field (e.g.
 // `reasoning_content`) by passing a field name; only the boolean form is
 // supported here.
-fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result<String> {
+fn append_continue_final_message_tag(message: &mut TemplateMessage<'_>) -> Result<String> {
     let text = match &mut message.content {
         TemplateContent::String(text) => Some(text),
         // Pick the last text part in the message.
@@ -578,9 +614,10 @@ mod tests {
     use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
 
     use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo};
+    use crate::EffortValue;
     use crate::request::{
         ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice,
-        GenerationPromptMode, ReasoningEffort, ResolvedToolContext,
+        GenerationPromptMode, ResolvedToolContext,
     };
     use crate::{AssistantContentBlock, ChatRenderer, Error, Result};
 
@@ -642,6 +679,22 @@ mod tests {
         .map_err(|_| unreachable!("HF renderer should return text prompt"))
     }
 
+    #[test]
+    fn generation_blocks_allow_content_format_detection_and_request_overrides() {
+        let template = "{% for message in messages %}{% generation %}{% for part in message.content %}{{ part.text }}{% endfor %}{% endgeneration %}{% endfor %}";
+        let mut request = sample_request(vec![ChatMessage::user("hello")]);
+        let default = render(Some(template), &request).unwrap();
+        request.chat_options.chat_template = Some(template.to_string());
+        let overridden = render(Some("unused"), &request).unwrap();
+        expect![[r#"
+            (
+                "hello",
+                "hello",
+            )
+        "#]]
+        .assert_debug_eq(&(default, overridden));
+    }
+
     fn render_mm(
         template: &str,
         request: &ChatRequest,
@@ -957,6 +1010,23 @@ mod tests {
         assert_eq!(rendered, "developer|policy|get_weather|city");
     }
 
+    #[test]
+    fn chat_template_passes_custom_roles_unchanged() {
+        let request = sample_request(vec![
+            ChatMessage::custom("root", "identity"),
+            ChatMessage::system("policy"),
+            ChatMessage::user("hello"),
+        ]);
+
+        let rendered = render(
+            Some("{% for m in messages %}{{ m.role }}={{ m.content }};{% endfor %}"),
+            &request,
+        )
+        .unwrap();
+
+        assert_eq!(rendered, "root=identity;system=policy;user=hello;");
+    }
+
     #[test]
     fn chat_template_keeps_string_text_for_openai_detected_templates() {
         let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
@@ -1131,7 +1201,7 @@ mod tests {
     #[test]
     fn chat_template_reasoning_effort_overrides_template_kwargs() {
         let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
-        request.chat_options.reasoning_effort = Some(ReasoningEffort::Max);
+        request.chat_options.reasoning_effort = Some(EffortValue::from("max"));
         request.chat_options.template_kwargs.insert(
             "reasoning_effort".to_string(),
             Value::String("low".to_string()),
@@ -1163,7 +1233,7 @@ mod tests {
     #[test]
     fn chat_template_reasoning_effort_preserves_request_enable_thinking() {
         let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
-        request.chat_options.reasoning_effort = Some(ReasoningEffort::None);
+        request.chat_options.reasoning_effort = Some(EffortValue::from("none"));
         request
             .chat_options
             .template_kwargs
@@ -1189,6 +1259,70 @@ mod tests {
         );
     }
 
+    #[test]
+    fn chat_template_preserves_typed_effort_values() {
+        for effort in [
+            serde_json::json!(37),
+            serde_json::json!(0.37),
+            serde_json::json!("custom"),
+        ] {
+            let mut request = sample_request(vec![ChatMessage::user("hello")]);
+            request.chat_options.reasoning_effort =
+                Some(serde_json::from_value(effort.clone()).unwrap());
+            request.chat_options.template_kwargs = [
+                ("reasoning_effort".to_string(), serde_json::json!("high")),
+                ("enable_thinking".to_string(), serde_json::json!(true)),
+            ]
+            .into();
+            let rendered = HfChatRenderer::new(
+                Some("{{ reasoning_effort|tojson }}|{{ enable_thinking }}".to_string()),
+                HashMap::new(),
+                ChatTemplateContentFormatOption::Auto,
+            )
+            .unwrap()
+            .render(&request)
+            .unwrap();
+            assert_eq!(rendered.prompt, Prompt::Text(format!("{effort}|True")));
+            assert_eq!(
+                rendered.effective_template_kwargs["reasoning_effort"],
+                effort
+            );
+        }
+    }
+
+    #[test]
+    fn chat_template_keeps_native_control_conflicts_and_raw_effort_values() {
+        for effort in [
+            serde_json::json!(37),
+            serde_json::json!(null),
+            serde_json::json!("none"),
+        ] {
+            let mut request = sample_request(vec![ChatMessage::user("hello")]);
+            request.chat_options.template_kwargs = [
+                ("thinking".to_string(), serde_json::json!("custom-mode")),
+                ("enable_thinking".to_string(), serde_json::json!(true)),
+                ("reasoning_effort".to_string(), effort),
+            ]
+            .into();
+            let rendered = HfChatRenderer::new(
+                Some(
+                    "{{ thinking }}|{{ enable_thinking }}|{{ reasoning_effort is none }}"
+                        .to_string(),
+                ),
+                HashMap::new(),
+                ChatTemplateContentFormatOption::Auto,
+            )
+            .unwrap()
+            .render(&request)
+            .unwrap();
+            assert!(rendered.prompt.into_text().unwrap().starts_with("custom-mode|True|"));
+            assert_eq!(
+                rendered.effective_template_kwargs,
+                request.chat_options.template_kwargs
+            );
+        }
+    }
+
     #[test]
     fn qwen3_template_omits_reasoning_for_historical_assistant_messages() {
         let request = sample_request(vec![
diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs
index b7de0172f36f..de447e2c34c4 100644
--- a/rust/src/chat/src/renderer/hf/template.rs
+++ b/rust/src/chat/src/renderer/hf/template.rs
@@ -20,6 +20,7 @@ use super::error::TemplateError;
 use super::format::{
     ChatTemplateContentFormat, ChatTemplateContentFormatOption, detect_chat_template_content_format,
 };
+use super::generation::{render_generation, rewrite_generation_blocks};
 use super::tojson::hf_tojson_filter;
 use crate::renderer::hf::{TemplateMessage, TemplateTool};
 
@@ -36,6 +37,7 @@ fn build_environment(template: String) -> Result<Environment<'static>> {
 
     env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
     env.add_filter("tojson", hf_tojson_filter);
+    env.add_function("__hf_generation", render_generation);
 
     Ok(env)
 }
@@ -43,7 +45,7 @@ fn build_environment(template: String) -> Result<Environment<'static>> {
 #[serde_with::skip_serializing_none]
 #[derive(Default, Serialize)]
 pub(super) struct TemplateContext<'a> {
-    pub(super) messages: &'a [TemplateMessage],
+    pub(super) messages: &'a [TemplateMessage<'a>],
     pub(super) add_generation_prompt: bool,
     pub(super) continue_final_message: bool,
     pub(super) tools: Option<&'a [TemplateTool]>,
@@ -107,6 +109,7 @@ pub(super) struct CompiledChatTemplate {
 impl CompiledChatTemplate {
     /// Compile the given chat template string into a [`CompiledChatTemplate`].
     pub fn new(template: String, content_format: ChatTemplateContentFormatOption) -> Result<Self> {
+        let template = rewrite_generation_blocks(template)?;
         let content_format = match content_format {
             ChatTemplateContentFormatOption::Auto => detect_chat_template_content_format(&template),
             ChatTemplateContentFormatOption::String => ChatTemplateContentFormat::String,
diff --git a/rust/src/chat/src/renderer/inkling/mod.rs b/rust/src/chat/src/renderer/inkling/mod.rs
index a33e3414f2db..f792f70d3288 100644
--- a/rust/src/chat/src/renderer/inkling/mod.rs
+++ b/rust/src/chat/src/renderer/inkling/mod.rs
@@ -8,8 +8,10 @@ use thiserror_ext::AsReport as _;
 use vllm_text::Prompt;
 use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
 
-use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
+use super::{ChatRenderer, RenderedPrompt};
+use crate::EffortValue;
 use crate::error::{Error, Result};
+use crate::reasoning::ReasoningControl;
 use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool};
 use crate::{AssistantContentBlock, AssistantToolCall};
 
@@ -35,6 +37,7 @@ const MAX_REASONING_EFFORT: f64 = 0.99;
 pub struct InklingChatRenderer {
     tokenizer: DynTokenizer,
     special: InklingSpecialTokenIds,
+    default_template_kwargs: HashMap<String, Value>,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -78,9 +81,16 @@ impl InklingSpecialTokenIds {
 }
 
 impl InklingChatRenderer {
-    pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
+    pub fn new(
+        tokenizer: DynTokenizer,
+        default_template_kwargs: HashMap<String, Value>,
+    ) -> Result<Self> {
         let special = InklingSpecialTokenIds::resolve(tokenizer.as_ref())?;
-        Ok(Self { tokenizer, special })
+        Ok(Self {
+            tokenizer,
+            special,
+            default_template_kwargs,
+        })
     }
 
     fn write_text_tokens(&self, out: &mut Vec<u32>, text: &str) -> Result<()> {
@@ -150,11 +160,6 @@ impl InklingChatRenderer {
     }
 
     fn write_reasoning_effort(&self, out: &mut Vec<u32>, effort: f64) -> Result<()> {
-        if !(0.0..=MAX_REASONING_EFFORT).contains(&effort) {
-            return Err(Error::ChatTemplate(format!(
-                "Inkling reasoning_effort must be in [0.0, 0.99], got {effort}"
-            )));
-        }
         let formatted = format!("{effort:.2}");
         let effort = formatted.trim_end_matches('0').trim_end_matches('.');
         let effort = if matches!(effort, "0" | "-0") {
@@ -294,11 +299,18 @@ impl ChatRenderer for InklingChatRenderer {
 
         let mut out = Vec::new();
         let mut tool_call_id_to_name = HashMap::new();
-        let effective_template_kwargs = request_template_kwargs(request);
+        let reasoning = resolve_reasoning(ReasoningControl::resolve(
+            request,
+            &self.default_template_kwargs,
+        )?)?;
+        let effective_template_kwargs = reasoning.template_kwargs(request);
         let tools = rendered_tools(request);
         self.write_tool_declarations(&mut out, &tools)?;
-        let mut reasoning_effort =
-            resolve_reasoning_effort(effective_template_kwargs.get("reasoning_effort"));
+        let mut reasoning_effort = if reasoning.is_enabled() {
+            reasoning.effort().and_then(EffortValue::as_f64)
+        } else {
+            Some(0.0)
+        };
 
         for message in &request.messages {
             if !matches!(
@@ -333,6 +345,9 @@ impl ChatRenderer for InklingChatRenderer {
                         &tool_call_id_to_name,
                     )?;
                 }
+                ChatMessage::Custom { role, .. } => {
+                    return Err(Error::UnsupportedChatRole { role: role.clone() });
+                }
             }
         }
 
@@ -346,17 +361,39 @@ impl ChatRenderer for InklingChatRenderer {
 
         Ok(RenderedPrompt {
             prompt: Prompt::TokenIds(out),
+            media_order: None,
             effective_template_kwargs,
         })
     }
 }
 
-fn resolve_reasoning_effort(value: Option<&Value>) -> Option<f64> {
+fn resolve_reasoning(control: ReasoningControl) -> Result<ReasoningControl> {
+    let control = control.fallback(ReasoningControl::enabled("high"));
+    if !control.is_enabled() {
+        return Ok(control);
+    }
+    let effort = resolve_reasoning_effort(control.effort());
+    if let Some(effort) = effort {
+        if !(0.0..=MAX_REASONING_EFFORT).contains(&effort) {
+            return Err(Error::InvalidReasoningEffort(format!(
+                "Inkling reasoning_effort must be in [0.0, 0.99], got {effort}"
+            )));
+        }
+        if effort == 0.0 {
+            return Ok(ReasoningControl::Disabled);
+        }
+    }
+    Ok(ReasoningControl::Enabled {
+        effort: effort.map(EffortValue::try_from).transpose()?,
+    })
+}
+
+fn resolve_reasoning_effort(value: Option<&EffortValue>) -> Option<f64> {
     let Some(value) = value else {
         return Some(0.9);
     };
     match value {
-        Value::String(name) => match name.as_str() {
+        EffortValue::String(name) => match name.as_str() {
             "none" => Some(0.0),
             "minimal" => Some(0.1),
             "low" => Some(0.2),
@@ -365,8 +402,7 @@ fn resolve_reasoning_effort(value: Option<&Value>) -> Option<f64> {
             "xhigh" | "max" => Some(0.99),
             _ => None,
         },
-        Value::Number(number) => number.as_f64(),
-        _ => None,
+        EffortValue::Number(number) => number.as_f64(),
     }
 }
 
diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs
index 1c4cd7236dbd..d0dd3055e413 100644
--- a/rust/src/chat/src/renderer/inkling/tests.rs
+++ b/rust/src/chat/src/renderer/inkling/tests.rs
@@ -9,6 +9,7 @@ use serde_json::json;
 use thiserror_ext::AsReport;
 use vllm_text::tokenizer::Tokenizer;
 
+use crate::EffortValue;
 use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
 
 use super::{
@@ -17,7 +18,7 @@ use super::{
     MESSAGE_TOOL, MESSAGE_USER,
 };
 use crate::event::{AssistantContentBlock, AssistantToolCall};
-use crate::request::{ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort};
+use crate::request::{ChatMessage, ChatRequest, GenerationPromptMode};
 use crate::{ChatRenderer, Error};
 
 struct FixtureTokenizer;
@@ -87,7 +88,7 @@ impl Tokenizer for FixtureTokenizer {
 }
 
 fn renderer() -> InklingChatRenderer {
-    InklingChatRenderer::new(Arc::new(FixtureTokenizer)).unwrap()
+    InklingChatRenderer::new(Arc::new(FixtureTokenizer), Default::default()).unwrap()
 }
 
 fn render_token_ids(request: &ChatRequest) -> Vec<u32> {
@@ -105,7 +106,7 @@ fn fixture_request(name: &str) -> ChatRequest {
 
 fn inkling_fixture_options() -> FixtureRequestOptions {
     FixtureRequestOptions {
-        enable_thinking: Some(false),
+        enable_thinking: None,
         no_generation_prompt_when_last_assistant: false,
     }
 }
@@ -206,13 +207,13 @@ fn renders_request_and_developer_tools_as_tool_declare() {
 #[test]
 fn renders_named_reasoning_effort_after_tool_declarations() {
     for (effort, expected) in [
-        (ReasoningEffort::None, "0.0"),
-        (ReasoningEffort::Minimal, "0.1"),
-        (ReasoningEffort::Low, "0.2"),
-        (ReasoningEffort::Medium, "0.7"),
-        (ReasoningEffort::High, "0.9"),
-        (ReasoningEffort::XHigh, "0.99"),
-        (ReasoningEffort::Max, "0.99"),
+        (EffortValue::from("none"), "0.0"),
+        (EffortValue::from("minimal"), "0.1"),
+        (EffortValue::from("low"), "0.2"),
+        (EffortValue::from("medium"), "0.7"),
+        (EffortValue::from("high"), "0.9"),
+        (EffortValue::from("xhigh"), "0.99"),
+        (EffortValue::from("max"), "0.99"),
     ] {
         let mut request = fixture_request("tool_declare_input.json");
         request.chat_options.reasoning_effort = Some(effort);
@@ -284,6 +285,43 @@ fn defaults_reasoning_effort_to_high() {
     );
 }
 
+#[test]
+fn normalized_mode_and_scalar_are_reflected_in_prompt_and_kwargs() {
+    for (kwargs, expected, enabled) in [
+        (
+            json!({"enable_thinking": true, "reasoning_effort": "none"}),
+            "0.9",
+            true,
+        ),
+        (
+            json!({"thinking": false, "reasoning_effort": 10}),
+            "0.0",
+            false,
+        ),
+        (json!({"reasoning_effort": 0}), "0.0", false),
+        (json!({"reasoning_effort": null}), "0.9", true),
+        (json!({"reasoning_effort": 0.37}), "0.37", true),
+    ] {
+        let mut request = ChatRequest::for_test();
+        request.chat_options.template_kwargs = serde_json::from_value(kwargs).unwrap();
+        let rendered = renderer().render(&request).unwrap();
+        let prompt = render_symbolic_tokens(&rendered.prompt.into_token_ids().unwrap());
+        assert!(prompt.contains(&format!("Thinking effort level: {expected}<|end_message|>")));
+        assert_eq!(
+            rendered.effective_template_kwargs["enable_thinking"],
+            enabled
+        );
+        assert_eq!(
+            rendered.effective_template_kwargs["reasoning_effort"],
+            if enabled {
+                json!(expected.parse::<f64>().unwrap())
+            } else {
+                json!("none")
+            }
+        );
+    }
+}
+
 #[test]
 fn renders_numeric_reasoning_effort_template_kwarg() {
     let mut request = ChatRequest::for_test();
@@ -301,19 +339,17 @@ fn renders_numeric_reasoning_effort_template_kwarg() {
 }
 
 #[test]
-fn ignores_unsupported_reasoning_effort_values() {
-    for value in [json!(true), json!("invalid"), json!(null)] {
-        let mut request = ChatRequest::for_test();
-        request
-            .chat_options
-            .template_kwargs
-            .insert("reasoning_effort".to_string(), value);
+fn ignores_unsupported_reasoning_effort_names() {
+    let mut request = ChatRequest::for_test();
+    request
+        .chat_options
+        .template_kwargs
+        .insert("reasoning_effort".to_string(), json!("invalid"));
 
-        assert_eq!(
-            render_symbolic_tokens(&render_token_ids(&request)),
-            "<|message_user|><|content_text|>test<|end_message|><|message_model|>"
-        );
-    }
+    assert_eq!(
+        render_symbolic_tokens(&render_token_ids(&request)),
+        "<|message_user|><|content_text|>test<|end_message|><|message_model|>"
+    );
 }
 
 #[test]
diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs
index 392b6b642d29..242478294794 100644
--- a/rust/src/chat/src/renderer/kimi_k3/encoding.rs
+++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs
@@ -10,7 +10,10 @@ use std::collections::HashMap;
 use serde_json::{Map, Value, json};
 use vllm_tokenizer::Tokenizer;
 
+use super::super::MediaPartSource;
+use crate::EffortValue;
 use crate::error::{Error, Result};
+use crate::reasoning::ReasoningControl;
 use crate::request::{
     ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice,
 };
@@ -29,6 +32,7 @@ const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"];
 pub(super) struct K3TokenWriter<'a> {
     tokenizer: &'a dyn Tokenizer,
     token_ids: Vec<u32>,
+    media_order: Vec<MediaPartSource>,
 }
 
 impl<'a> K3TokenWriter<'a> {
@@ -36,6 +40,7 @@ impl<'a> K3TokenWriter<'a> {
         Self {
             tokenizer,
             token_ids: Vec::new(),
+            media_order: Vec::new(),
         }
     }
 
@@ -55,15 +60,28 @@ impl<'a> K3TokenWriter<'a> {
         Ok(())
     }
 
-    pub(super) fn finish(self) -> Vec<u32> {
-        self.token_ids
+    fn image(&mut self, message_index: usize, content_part_index: usize) -> Result<()> {
+        self.control(IMAGE_PLACEHOLDER)?;
+        self.media_order.push(MediaPartSource {
+            message_index,
+            content_part_index,
+        });
+        Ok(())
+    }
+
+    pub(super) fn finish(self) -> (Vec<u32>, Vec<MediaPartSource>) {
+        (self.token_ids, self.media_order)
     }
 }
 
-/// Render and tokenize one chat request using K3's segment-aware contract.
-pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result<Vec<u32>> {
-    let thinking = thinking_enabled(request)?;
-    let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?;
+/// Render and tokenize one chat request while recording media in placeholder order.
+pub(super) fn render_request_with_media_order(
+    request: &ChatRequest,
+    tokenizer: &dyn Tokenizer,
+    reasoning: &ReasoningControl,
+) -> Result<(Vec<u32>, Vec<MediaPartSource>)> {
+    let thinking = reasoning.is_enabled();
+    let thinking_effort = reasoning.effort().and_then(EffortValue::as_str);
     let tools = request_tools(request);
     let mut out = K3TokenWriter::new(tokenizer);
 
@@ -126,7 +144,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
             resolved.sort_by_key(|item| (item.0, item.1));
         }
 
-        for (position, _, content, name) in resolved {
+        for (position, message_index, content, name) in resolved {
             let tool_name = name.as_deref().ok_or_else(|| {
                 Error::ChatTemplate(
                     "Kimi K3 tool messages need a resolvable tool name: \
@@ -135,7 +153,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
                         .to_string(),
                 )
             })?;
-            write_tool_message(out, tool_name, position, &content)?;
+            write_tool_message(out, tool_name, position, message_index, &content)?;
         }
         Ok(())
     };
@@ -161,14 +179,14 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
             } if !local_tools.is_empty() => {
                 write_tool_declare(&mut out, local_tools, true)?;
                 if !content_is_empty(content) {
-                    write_role_message(&mut out, "system", None, content)?;
+                    write_role_message(&mut out, "system", None, message_index, content)?;
                 }
             }
             ChatMessage::System { content } | ChatMessage::Developer { content, .. } => {
-                write_role_message(&mut out, "system", None, content)?;
+                write_role_message(&mut out, "system", None, message_index, content)?;
             }
             ChatMessage::User { content } => {
-                write_role_message(&mut out, "user", None, content)?;
+                write_role_message(&mut out, "user", None, message_index, content)?;
             }
             ChatMessage::Assistant { content } => {
                 tool_call_id_index.clear();
@@ -187,6 +205,9 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
                 write_assistant_message(&mut out, content, thinking)?;
             }
             ChatMessage::ToolResponse { .. } => unreachable!("handled above"),
+            ChatMessage::Custom { role, .. } => {
+                return Err(Error::UnsupportedChatRole { role: role.clone() });
+            }
         }
     }
     flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?;
@@ -229,58 +250,26 @@ fn request_tools(request: &ChatRequest) -> &[ChatTool] {
     request.initial_tools()
 }
 
-fn thinking_enabled(request: &ChatRequest) -> Result<bool> {
-    if let Some(thinking) = request.parse_template_bool("thinking")? {
-        return Ok(thinking);
-    }
-    if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? {
-        return Ok(enable_thinking);
-    }
-    if let Some(reasoning_effort) = request.chat_options.reasoning_effort {
-        return Ok(reasoning_effort != crate::request::ReasoningEffort::None);
-    }
-    Ok(request
-        .chat_options
-        .template_kwargs
-        .get("reasoning_effort")
-        .and_then(Value::as_str)
-        != Some("none"))
-}
-
-fn thinking_effort(request: &ChatRequest) -> Result<String> {
-    let effort = if let Some(value) = request.chat_options.template_kwargs.get("thinking_effort") {
-        value.as_str().ok_or_else(|| {
-            Error::ChatTemplate(format!(
-                "template kwarg `thinking_effort` must be a string, got {value}"
-            ))
-        })?
-    } else if let Some(effort) = request.chat_options.reasoning_effort {
-        if effort == crate::request::ReasoningEffort::None {
-            DEFAULT_THINKING_EFFORT
-        } else {
-            effort.as_str()
-        }
-    } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") {
+/// Resolve standard reasoning controls and validate K3's supported effort grades.
+pub(super) fn resolve_reasoning(
+    request: &ChatRequest,
+    defaults: &HashMap<String, Value>,
+) -> Result<ReasoningControl> {
+    let reasoning = ReasoningControl::resolve(request, defaults)?
+        .fallback(ReasoningControl::enabled(DEFAULT_THINKING_EFFORT));
+    if let Some(value) = reasoning.effort() {
         let effort = value.as_str().ok_or_else(|| {
-            Error::ChatTemplate(format!(
-                "template kwarg `reasoning_effort` must be a string, got {value}"
+            Error::InvalidReasoningEffort(format!(
+                "Kimi K3 reasoning_effort must be a string, got {value}"
             ))
         })?;
-        if effort == "none" {
-            DEFAULT_THINKING_EFFORT
-        } else {
-            effort
+        if !VALID_THINKING_EFFORTS.contains(&effort) {
+            return Err(Error::InvalidReasoningEffort(format!(
+                "unsupported reasoning_effort={effort:?}; supported values are `low`, `high`, and `max`"
+            )));
         }
-    } else {
-        DEFAULT_THINKING_EFFORT
-    };
-
-    if !VALID_THINKING_EFFORTS.contains(&effort) {
-        return Err(Error::ChatTemplate(format!(
-            "unsupported thinking_effort={effort:?}; supported values are `low`, `high`, and `max`"
-        )));
     }
-    Ok(effort.to_string())
+    Ok(reasoning)
 }
 
 fn content_is_empty(content: &ChatContent) -> bool {
@@ -364,6 +353,7 @@ fn write_role_message(
     out: &mut K3TokenWriter<'_>,
     role: &str,
     name: Option<&str>,
+    message_index: usize,
     content: &ChatContent,
 ) -> Result<()> {
     let mut attrs = vec![("role", role.to_string())];
@@ -372,7 +362,7 @@ fn write_role_message(
     }
     let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect();
     write_open_tag(out, "message", &attr_refs)?;
-    write_content(out, content)?;
+    write_content(out, message_index, content)?;
     write_close_tag(out, "message")?;
     out.control(END_OF_MSG)
 }
@@ -381,6 +371,7 @@ fn write_tool_message(
     out: &mut K3TokenWriter<'_>,
     tool_name: &str,
     index: usize,
+    message_index: usize,
     content: &ChatContent,
 ) -> Result<()> {
     let index_str = index.to_string();
@@ -389,7 +380,7 @@ fn write_tool_message(
         "message",
         &[("role", "tool"), ("tool", tool_name), ("index", &index_str)],
     )?;
-    write_content(out, content)?;
+    write_content(out, message_index, content)?;
     write_close_tag(out, "message")?;
     out.control(END_OF_MSG)
 }
@@ -471,14 +462,20 @@ fn write_assistant_tool_call(
     write_close_tag(out, "call")
 }
 
-fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> {
+fn write_content(
+    out: &mut K3TokenWriter<'_>,
+    message_index: usize,
+    content: &ChatContent,
+) -> Result<()> {
     match content {
         ChatContent::Text(text) => write_text_with_images(out, text),
         ChatContent::Parts(parts) => {
-            for part in parts {
+            for (content_part_index, part) in parts.iter().enumerate() {
                 match part {
                     ChatContentPart::Text { text } => write_text_with_images(out, text)?,
-                    ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?,
+                    ChatContentPart::ImageUrl { .. } => {
+                        out.image(message_index, content_part_index)?;
+                    }
                     ChatContentPart::VideoUrl { .. } => {
                         return Err(Error::UnsupportedMultimodalContent("video_url"));
                     }
diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs
index 09bf356f5d64..c39ea2985a07 100644
--- a/rust/src/chat/src/renderer/kimi_k3/mod.rs
+++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs
@@ -7,10 +7,13 @@ mod encoding;
 #[cfg(test)]
 mod tests;
 
+use std::collections::HashMap;
+
+use serde_json::Value;
 use vllm_text::Prompt;
 use vllm_text::tokenizer::DynTokenizer;
 
-use super::{ChatRenderer, RenderedPrompt, request_template_kwargs};
+use super::{ChatRenderer, RenderedPrompt};
 use crate::Result;
 use crate::request::ChatRequest;
 
@@ -18,22 +21,42 @@ use crate::request::ChatRequest;
 #[derive(Clone)]
 pub struct KimiK3ChatRenderer {
     tokenizer: DynTokenizer,
+    default_template_kwargs: HashMap<String, Value>,
 }
 
 impl KimiK3ChatRenderer {
     /// Create a Kimi K3 renderer.
-    pub fn new(tokenizer: DynTokenizer) -> Self {
-        Self { tokenizer }
+    pub fn new(tokenizer: DynTokenizer, default_template_kwargs: HashMap<String, Value>) -> Self {
+        Self {
+            tokenizer,
+            default_template_kwargs,
+        }
     }
 }
 
 impl ChatRenderer for KimiK3ChatRenderer {
     fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
         request.validate()?;
+        let reasoning = encoding::resolve_reasoning(request, &self.default_template_kwargs)?;
+        let mut effective_template_kwargs = reasoning.template_kwargs(request);
+        // Export K3's native effort key from the resolved standard controls.
+        // Input `thinking_effort` is ignored; clear it when reasoning is disabled.
+        if let Some(effort) = reasoning.effort() {
+            effective_template_kwargs
+                .insert("thinking_effort".to_string(), serde_json::json!(effort));
+        } else {
+            effective_template_kwargs.remove("thinking_effort");
+        }
+        let (token_ids, media_order) = encoding::render_request_with_media_order(
+            request,
+            self.tokenizer.as_ref(),
+            &reasoning,
+        )?;
 
         Ok(RenderedPrompt {
-            prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?),
-            effective_template_kwargs: request_template_kwargs(request),
+            prompt: Prompt::TokenIds(token_ids),
+            media_order: Some(media_order),
+            effective_template_kwargs,
         })
     }
 }
diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs
index 1d9f984dbbd3..5aadb3d403c6 100644
--- a/rust/src/chat/src/renderer/kimi_k3/tests.rs
+++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs
@@ -15,11 +15,10 @@ use vllm_tokenizer::test_utils::TestTokenizer;
 
 use super::KimiK3ChatRenderer;
 use crate::ChatRenderer;
+use crate::EffortValue;
 use crate::renderer::kimi_k3::encoding::{CLOSE, END_OF_MSG, IMAGE_PLACEHOLDER, OPEN, SEP};
 use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
-use crate::request::{
-    ChatContentPart, ChatMessage, ChatTool, GenerationPromptMode, ReasoningEffort,
-};
+use crate::request::{ChatContentPart, ChatMessage, ChatTool, GenerationPromptMode};
 use crate::{AssistantContentBlock, AssistantToolCall};
 
 const OPEN_ID: u32 = 256;
@@ -37,8 +36,70 @@ fn test_tokenizer() -> TestTokenizer {
         .with_special_token(IMAGE_PLACEHOLDER, MEDIA_ID)
 }
 
+#[test]
+fn standard_effort_controls_prompt_and_exported_native_effort() {
+    let tokenizer = Arc::new(test_tokenizer());
+    for (kwargs, typed, defaults, effort) in [
+        (json!({}), None, json!({}), "max"),
+        (
+            json!({}),
+            Some(EffortValue::from("low")),
+            json!({"reasoning_effort": "high"}),
+            "low",
+        ),
+        (
+            json!({"enable_thinking": true, "reasoning_effort": "none"}),
+            None,
+            json!({"reasoning_effort": "low"}),
+            "low",
+        ),
+        (
+            json!({}),
+            Some(EffortValue::from("none")),
+            json!({"reasoning_effort": "low"}),
+            "none",
+        ),
+        (
+            json!({"thinking": true}),
+            Some(EffortValue::from("none")),
+            json!({}),
+            "max",
+        ),
+    ] {
+        let renderer =
+            KimiK3ChatRenderer::new(tokenizer.clone(), serde_json::from_value(defaults).unwrap());
+        let mut request = crate::ChatRequest::for_test();
+        request.chat_options.reasoning_effort = typed;
+        request.chat_options.template_kwargs = serde_json::from_value(kwargs).unwrap();
+        let rendered = renderer.render(&request).unwrap();
+        let prompt = tokenizer.decode(&rendered.prompt.into_token_ids().unwrap(), false).unwrap();
+        let enabled = effort != "none";
+        assert_eq!(prompt.ends_with("<|open|>think<|sep|>"), enabled);
+        assert_eq!(
+            rendered.effective_template_kwargs["reasoning_effort"],
+            effort
+        );
+        assert_eq!(
+            rendered.effective_template_kwargs["enable_thinking"],
+            enabled
+        );
+        if enabled {
+            assert!(prompt.contains(&format!("thinking_effort={effort}")));
+            assert_eq!(
+                rendered.effective_template_kwargs["thinking_effort"],
+                effort
+            );
+        } else {
+            assert!(!rendered.effective_template_kwargs.contains_key("thinking_effort"));
+        }
+    }
+}
+
 fn render_token_ids(request: &crate::request::ChatRequest, tokenizer: DynTokenizer) -> Vec<u32> {
-    let prompt = KimiK3ChatRenderer::new(tokenizer).render(request).unwrap().prompt;
+    let prompt = KimiK3ChatRenderer::new(tokenizer, Default::default())
+        .render(request)
+        .unwrap()
+        .prompt;
     let Prompt::TokenIds(token_ids) = prompt else {
         panic!("kimi k3 renderer should return token IDs")
     };
@@ -67,7 +128,11 @@ fn kimi_k3_fixture_options() -> FixtureRequestOptions {
 
 fn assert_golden(name: &str) {
     let input_name = format!("{name}_input.json");
-    let request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options());
+    let mut request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options());
+    // Reference fixtures use encoder-native keys; Rust accepts standard effort input.
+    if let Some(effort) = request.chat_options.template_kwargs.remove("thinking_effort") {
+        request.chat_options.reasoning_effort = Some(EffortValue::try_from(effort).unwrap());
+    }
     let rendered = render_request(&request);
     expect_file![format!("fixtures/{name}_output.txt")].assert_eq(&rendered);
 }
@@ -284,6 +349,65 @@ fn partial_tool_results_keep_assistant_call_position() {
     );
 }
 
+#[test]
+fn media_order_follows_reordered_tool_results() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.messages = vec![
+        ChatMessage::assistant_blocks(vec![
+            AssistantContentBlock::ToolCall(AssistantToolCall {
+                id: "call-a".to_string(),
+                name: "tool-a".to_string(),
+                arguments: "{}".to_string(),
+            }),
+            AssistantContentBlock::ToolCall(AssistantToolCall {
+                id: "call-b".to_string(),
+                name: "tool-b".to_string(),
+                arguments: "{}".to_string(),
+            }),
+        ]),
+        ChatMessage::tool_response(
+            vec![ChatContentPart::text("result-b"), image_part("image-b")],
+            "call-b",
+        ),
+        ChatMessage::tool_response(
+            vec![ChatContentPart::text("result-a"), image_part("image-a")],
+            "call-a",
+        ),
+    ];
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+    let renderer = KimiK3ChatRenderer::new(Arc::new(test_tokenizer()), Default::default());
+
+    let rendered = renderer.render(&request).unwrap();
+    let media_order = rendered.media_order.unwrap();
+    let Prompt::TokenIds(token_ids) = rendered.prompt else {
+        panic!("kimi k3 renderer should return token IDs")
+    };
+    let prompt = test_tokenizer().decode(&token_ids, false).unwrap();
+
+    assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap());
+    expect![[r#"
+        [
+            MediaPartSource {
+                message_index: 2,
+                content_part_index: 1,
+            },
+            MediaPartSource {
+                message_index: 1,
+                content_part_index: 1,
+            },
+        ]
+    "#]]
+    .assert_debug_eq(&media_order);
+}
+
+fn image_part(uuid: &str) -> ChatContentPart {
+    ChatContentPart::ImageUrl {
+        image_url: format!("data:image/png;base64,{uuid}"),
+        detail: None,
+        uuid: Some(uuid.to_string()),
+    }
+}
+
 #[test]
 fn defaults_thinking_effort_to_max() {
     let rendered = render_request(&crate::request::ChatRequest::for_test());
@@ -312,19 +436,18 @@ fn translates_standard_thinking_kwargs() {
 }
 
 #[test]
-fn native_k3_kwargs_take_precedence() {
+fn thinking_alias_takes_precedence() {
     let mut request = crate::request::ChatRequest::for_test();
     request.chat_options.template_kwargs.extend([
         ("thinking".to_string(), json!(true)),
         ("enable_thinking".to_string(), json!(false)),
-        ("thinking_effort".to_string(), json!("low")),
         ("reasoning_effort".to_string(), json!("high")),
     ]);
 
     let rendered = render_request(&request);
 
-    assert!(rendered.contains("thinking_effort=low"));
-    assert!(!rendered.contains("thinking_effort=high"));
+    assert!(rendered.contains("thinking_effort=high"));
+    assert!(rendered.ends_with("<|open|>think<|sep|>"));
 }
 
 #[test]
@@ -372,7 +495,7 @@ fn enable_thinking_true_overrides_standard_none() {
 #[test]
 fn typed_none_disables_thinking() {
     let mut request = crate::request::ChatRequest::for_test();
-    request.chat_options.reasoning_effort = Some(ReasoningEffort::None);
+    request.chat_options.reasoning_effort = Some(EffortValue::from("none"));
 
     let rendered = render_request(&request);
 
@@ -381,20 +504,17 @@ fn typed_none_disables_thinking() {
 }
 
 #[test]
-fn rejects_removed_medium_thinking_effort() {
+fn rejects_unsupported_standard_reasoning_effort() {
     let mut request = crate::request::ChatRequest::for_test();
-    request
-        .chat_options
-        .template_kwargs
-        .insert("thinking_effort".to_string(), json!("medium"));
+    request.chat_options.reasoning_effort = Some(EffortValue::from("medium"));
 
-    let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer()))
+    let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer()), Default::default())
         .render(&request)
         .unwrap_err();
 
     expect![[r#"
-        ChatTemplate(
-            "unsupported thinking_effort=\"medium\"; supported values are `low`, `high`, and `max`",
+        InvalidReasoningEffort(
+            "unsupported reasoning_effort=\"medium\"; supported values are `low`, `high`, and `max`",
         )
     "#]]
     .assert_debug_eq(&error);
diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs
index a4649a34f70a..347ec43d03e1 100644
--- a/rust/src/chat/src/renderer/mod.rs
+++ b/rust/src/chat/src/renderer/mod.rs
@@ -4,11 +4,11 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use serde_json::{Value, json};
+use serde_json::Value;
 use vllm_text::Prompt;
 
 use crate::error::Result;
-use crate::request::{ChatRequest, ReasoningEffort};
+use crate::request::ChatRequest;
 
 mod deepseek;
 pub mod deepseek_v32;
@@ -30,11 +30,24 @@ pub use inkling::InklingChatRenderer;
 pub use kimi_k3::KimiK3ChatRenderer;
 pub use selection::RendererSelection;
 
+/// Location of one multimodal content part in the source chat request.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct MediaPartSource {
+    /// Zero-based message index in the source request.
+    pub message_index: usize,
+    /// Zero-based content-part index within the source message.
+    pub content_part_index: usize,
+}
+
 /// Rendered chat prompt submitted to the text backend.
 #[derive(Debug, Clone, PartialEq)]
 pub struct RenderedPrompt {
     /// The rendered prompt, either as text or already tokenized.
     pub prompt: Prompt,
+    /// Media sources in rendered placeholder order.
+    /// `None` uses the message and content-part order in the request (e.g. Jinja).
+    /// `Some` uses exactly the listed sources; an empty list omits all media.
+    pub media_order: Option<Vec<MediaPartSource>>,
     /// Effective chat-template kwargs visible to the renderer after applying
     /// server defaults, request overrides, and typed reasoning controls.
     pub effective_template_kwargs: HashMap<String, Value>,
@@ -49,33 +62,3 @@ pub trait ChatRenderer: Send + Sync {
 
 /// Shared trait-object form of [`ChatRenderer`].
 pub type DynChatRenderer = Arc<dyn ChatRenderer>;
-
-/// Extract the effective chat-template kwargs visible to the renderer from the request,
-/// using the provided defaults as the base.
-pub(crate) fn effective_template_kwargs(
-    default_template_kwargs: &HashMap<String, Value>,
-    request: &ChatRequest,
-) -> HashMap<String, Value> {
-    let mut kwargs = default_template_kwargs.clone();
-    kwargs.extend(request.chat_options.template_kwargs.clone());
-
-    if let Some(reasoning_effort) = request.chat_options.reasoning_effort {
-        kwargs.insert(
-            "reasoning_effort".to_string(),
-            Value::String(reasoning_effort.as_str().to_string()),
-        );
-        if !request.chat_options.template_kwargs.contains_key("enable_thinking") {
-            kwargs.insert(
-                "enable_thinking".to_string(),
-                json!(reasoning_effort != ReasoningEffort::None),
-            );
-        }
-    }
-
-    kwargs
-}
-
-/// Extract the effective chat-template kwargs visible to the renderer from the request.
-pub(crate) fn request_template_kwargs(request: &ChatRequest) -> HashMap<String, Value> {
-    effective_template_kwargs(&HashMap::new(), request)
-}
diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs
index 1090a4039f5e..2c59ed7a22c4 100644
--- a/rust/src/chat/src/renderer/test_utils.rs
+++ b/rust/src/chat/src/renderer/test_utils.rs
@@ -8,10 +8,11 @@ use std::path::Path;
 use serde::Deserialize;
 use serde_json::Value;
 
+use crate::EffortValue;
 use crate::event::{AssistantContentBlock, AssistantToolCall};
 use crate::request::{
     ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice,
-    GenerationPromptMode, ReasoningEffort, ResolvedToolContext,
+    GenerationPromptMode, ResolvedToolContext,
 };
 
 /// Options for constructing a [`ChatRequest`] from a fixture file.
@@ -47,7 +48,7 @@ pub(crate) struct FixtureRequest {
     tools: Vec<FixtureTool>,
     messages: Vec<FixtureMessage>,
     add_generation_prompt: Option<bool>,
-    reasoning_effort: Option<ReasoningEffort>,
+    reasoning_effort: Option<EffortValue>,
     /// Standard response format passed to model-specific renderers.
     #[serde(default)]
     response_format: Option<Value>,
diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs
index 50e00bbc43a5..f7ddd9b12843 100644
--- a/rust/src/chat/src/request.rs
+++ b/rust/src/chat/src/request.rs
@@ -11,9 +11,9 @@ pub use vllm_parser::tool::Tool as ChatTool;
 use vllm_text::TextDecodeOptions;
 pub use vllm_text::{PromptTruncation, SamplingParams};
 
-use crate::AssistantMessageExt;
 use crate::error::{Error, Result};
 use crate::event::{AssistantContentBlock, AssistantMessage};
+use crate::{AssistantMessageExt, EffortValue};
 
 /// Role label for one text-only chat message.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -24,6 +24,9 @@ pub enum ChatRole {
     User,
     Assistant,
     ToolResponse,
+    /// Role outside the standard set, passed to chat templates unchanged.
+    #[serde(untagged)]
+    Custom(String),
 }
 
 /// One text-only chat content part in OpenAI-style block format.
@@ -212,6 +215,10 @@ pub enum ChatMessage {
         content: ChatContent,
         tool_call_id: String,
     },
+    /// Message content under a role outside the standard set, passed to chat
+    /// templates unchanged.
+    #[serde(untagged)]
+    Custom { role: String, content: ChatContent },
 }
 
 impl ChatMessage {
@@ -230,6 +237,7 @@ impl ChatMessage {
                      use ChatMessage::tool_response() instead"
                 )
             }
+            ChatRole::Custom(role) => Self::custom(role, content),
         }
     }
 
@@ -286,6 +294,14 @@ impl ChatMessage {
         }
     }
 
+    /// Construct one chat message with a role outside the standard set.
+    pub fn custom(role: impl Into<String>, content: impl Into<ChatContent>) -> Self {
+        Self::Custom {
+            role: role.into(),
+            content: content.into(),
+        }
+    }
+
     /// Return the chat role of this message.
     pub fn role(&self) -> ChatRole {
         match self {
@@ -294,6 +310,7 @@ impl ChatMessage {
             Self::User { .. } => ChatRole::User,
             Self::Assistant { .. } => ChatRole::Assistant,
             Self::ToolResponse { .. } => ChatRole::ToolResponse,
+            Self::Custom { role, .. } => ChatRole::Custom(role.clone()),
         }
     }
 
@@ -303,7 +320,8 @@ impl ChatMessage {
             Self::System { content }
             | Self::Developer { content, .. }
             | Self::User { content }
-            | Self::ToolResponse { content, .. } => content.try_flatten_to_text(),
+            | Self::ToolResponse { content, .. }
+            | Self::Custom { content, .. } => content.try_flatten_to_text(),
             Self::Assistant { content } => Ok(content.text()),
         }
     }
@@ -315,7 +333,8 @@ impl ChatMessage {
             Self::System { .. }
             | Self::Developer { .. }
             | Self::User { .. }
-            | Self::ToolResponse { .. } => None,
+            | Self::ToolResponse { .. }
+            | Self::Custom { .. } => None,
         }
     }
 
@@ -325,7 +344,8 @@ impl ChatMessage {
             Self::System { content }
             | Self::Developer { content, .. }
             | Self::User { content }
-            | Self::ToolResponse { content, .. } => content.has_multimodal(),
+            | Self::ToolResponse { content, .. }
+            | Self::Custom { content, .. } => content.has_multimodal(),
             Self::Assistant { .. } => false,
         }
     }
@@ -362,33 +382,6 @@ pub enum GenerationPromptMode {
     NoGenerationPrompt,
 }
 
-/// Effort level for reasoning models.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
-#[serde(rename_all = "lowercase")]
-pub enum ReasoningEffort {
-    None,
-    Minimal,
-    Low,
-    Medium,
-    High,
-    XHigh,
-    Max,
-}
-
-impl ReasoningEffort {
-    pub fn as_str(self) -> &'static str {
-        match self {
-            Self::None => "none",
-            Self::Minimal => "minimal",
-            Self::Low => "low",
-            Self::Medium => "medium",
-            Self::High => "high",
-            Self::XHigh => "xhigh",
-            Self::Max => "max",
-        }
-    }
-}
-
 /// Chat-template-related request options.
 ///
 /// These are the small subset of chat controls that currently affect prompt
@@ -404,8 +397,10 @@ pub struct ChatOptions {
     /// used instead of the model's default chat template.
     pub chat_template: Option<String>,
 
-    /// Effort level exposed to chat templates for reasoning models.
-    pub reasoning_effort: Option<ReasoningEffort>,
+    /// Model-specific reasoning effort, typically `none`, `minimal`, `low`,
+    /// `medium`, `high`, `xhigh`, or `max`. Supported names and numeric ranges
+    /// are validated by the selected renderer or HF template.
+    pub reasoning_effort: Option<EffortValue>,
 
     /// Standard response format available to model-specific renderers.
     #[serde(default)]
@@ -460,6 +455,17 @@ pub enum ChatToolChoice {
     },
 }
 
+impl From<&ChatToolChoice> for xgrammar_structural_tag::ToolChoice {
+    fn from(tool_choice: &ChatToolChoice) -> Self {
+        match tool_choice {
+            ChatToolChoice::None => Self::none(),
+            ChatToolChoice::Auto => Self::auto(),
+            ChatToolChoice::Required => Self::required(),
+            ChatToolChoice::Function { name } => Self::function(name.clone()),
+        }
+    }
+}
+
 /// Resolved tool state shared by rendering, output parsing, and constraints.
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct ResolvedToolContext {
@@ -673,29 +679,6 @@ impl ChatRequest {
         self.tool_context.parsing_enabled()
     }
 
-    /// Return the request-level thinking toggle when explicitly requested.
-    ///
-    /// We currently accept the two request kwargs `thinking` and
-    /// `enable_thinking`. Both must be booleans when present. If both are
-    /// present, they must have the same value. If neither key is provided,
-    /// return `None`.
-    pub(crate) fn enable_thinking(&self) -> Result<Option<bool>> {
-        let thinking = self.parse_template_bool("thinking")?;
-        let enable_thinking = self.parse_template_bool("enable_thinking")?;
-
-        match (thinking, enable_thinking) {
-            (None, None) => Ok(None),
-            (Some(thinking), Some(enable_thinking)) if thinking != enable_thinking => {
-                Err(Error::ChatTemplate(
-                    "template kwargs `thinking` and `enable_thinking` must match when both are set"
-                        .to_string(),
-                ))
-            }
-            (Some(thinking), _) => Ok(Some(thinking)),
-            (None, Some(enable_thinking)) => Ok(Some(enable_thinking)),
-        }
-    }
-
     pub(crate) fn parse_template_bool(&self, key: &str) -> Result<Option<bool>> {
         match self.chat_options.template_kwargs.get(key) {
             None => Ok(None),
@@ -710,13 +693,14 @@ impl ChatRequest {
 impl ChatRole {
     /// Return the chat-template role string used by the current text-only chat
     /// backend.
-    pub fn as_str(&self) -> &'static str {
+    pub fn as_str(&self) -> &str {
         match self {
             Self::System => "system",
             Self::Developer => "developer",
             Self::User => "user",
             Self::Assistant => "assistant",
             Self::ToolResponse => "tool_response",
+            Self::Custom(role) => role,
         }
     }
 }
@@ -726,7 +710,7 @@ mod tests {
     use serde_json::{json, to_value};
 
     use super::{
-        ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice,
+        ChatContent, ChatContentPart, ChatMessage, ChatRole, ChatTool, ChatToolChoice,
         ResolvedToolContext,
     };
     use crate::Error;
@@ -919,56 +903,4 @@ mod tests {
             Error::ToolChoiceFunctionNotFound { name } if name == "missing"
         ));
     }
-
-    #[test]
-    fn enable_thinking_is_none_when_no_kwargs_are_present() {
-        let request = ChatRequest::for_test();
-        assert_eq!(request.enable_thinking().unwrap(), None);
-    }
-
-    #[test]
-    fn enable_thinking_accepts_matching_duplicate_kwargs() {
-        let mut request = ChatRequest::for_test();
-        request.chat_options.template_kwargs.insert("thinking".to_string(), json!(true));
-        request
-            .chat_options
-            .template_kwargs
-            .insert("enable_thinking".to_string(), json!(true));
-
-        assert_eq!(request.enable_thinking().unwrap(), Some(true));
-    }
-
-    #[test]
-    fn enable_thinking_rejects_non_boolean_kwargs() {
-        let mut request = ChatRequest::for_test();
-        request
-            .chat_options
-            .template_kwargs
-            .insert("thinking".to_string(), json!("yes"));
-
-        assert!(matches!(
-            request.enable_thinking(),
-            Err(Error::ChatTemplate(message))
-                if message.contains("`thinking` must be a boolean")
-        ));
-    }
-
-    #[test]
-    fn enable_thinking_rejects_conflicting_duplicate_kwargs() {
-        let mut request = ChatRequest::for_test();
-        request
-            .chat_options
-            .template_kwargs
-            .insert("thinking".to_string(), json!(false));
-        request
-            .chat_options
-            .template_kwargs
-            .insert("enable_thinking".to_string(), json!(true));
-
-        assert!(matches!(
-            request.enable_thinking(),
-            Err(Error::ChatTemplate(message))
-                if message.contains("`thinking` and `enable_thinking` must match")
-        ));
-    }
 }
diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs
index ab99b4d8ba73..ab366b8827a4 100644
--- a/rust/src/chat/tests/chat.rs
+++ b/rust/src/chat/tests/chat.rs
@@ -227,6 +227,7 @@ impl ChatBackend for FakeChatBackend {
             self.tokenizer(),
             options.tool_call_parser,
             options.reasoning_parser,
+            options.tool_strict_level,
         )?))
     }
 }
@@ -250,6 +251,7 @@ impl ChatRenderer for FakeChatBackend {
 
         Ok(RenderedPrompt {
             prompt: Prompt::Text(prompt),
+            media_order: None,
             effective_template_kwargs: request.chat_options.template_kwargs.clone(),
         })
     }
diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs
index 5669e255e0d5..a7bc24bd2a1a 100644
--- a/rust/src/chat/tests/roundtrip.rs
+++ b/rust/src/chat/tests/roundtrip.rs
@@ -141,6 +141,19 @@ impl RoundtripCase {
         }
     }
 
+    /// MiMo V2.6 uses bare reasoning and Qwen Coder tool tags.
+    fn mimo_v26() -> Self {
+        Self {
+            model_id: "XiaomiMiMo/MiMo-V2.6-Flash-RL",
+            assistant_stop_suffix: "<|im_end|>",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            thinking_behavior: ThinkingBehavior::Toggleable { default: true },
+            json_fmt: compact_json_fmt(),
+            sort_json_keys: false,
+        }
+    }
+
     /// MiniMax M2.5 XML invoke format with `<think>` reasoning tags.
     fn minimax_m25() -> Self {
         Self {
@@ -395,6 +408,7 @@ macro_rules! roundtrip_tests {
 roundtrip_tests! {
     qwen3 => [reasoning_and_content, tool_call_mix],
     qwen35 => [reasoning_and_content, tool_call_mix],
+    mimo_v26 => [reasoning_and_content, tool_call_mix],
     minimax_m25 => [reasoning_and_content, tool_call_mix],
     minimax_m3 => [reasoning_and_content, tool_call_mix],
     deepseek_v4 => [reasoning_and_content, tool_call_mix],
@@ -714,13 +728,15 @@ async fn parse_completion(
     };
 
     let mut request = base_request.clone();
-    let processor = backends.chat_backend.new_chat_output_processor(
+    let mut processor = backends.chat_backend.new_chat_output_processor(
         &mut request,
         NewChatOutputProcessorOptions {
+            tool_strict_level: vllm_chat::ToolStrictLevel::Auto,
             tool_call_parser: &case.tool_call_parser,
             reasoning_parser: &case.reasoning_parser,
         },
     )?;
+    processor.initialize(&prompt_token_ids)?;
 
     let decoded = decoded_completion_stream(
         tokenizer.as_ref(),
@@ -805,6 +821,7 @@ fn decoded_completion_stream(
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 })),
             }
         });
@@ -816,6 +833,7 @@ fn decoded_completion_stream(
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                sampling_mask: None,
             });
             events.push(DecodedTextEvent::TextDelta {
                 decoded: vllm_text::DecodedText::unattributed(chunk.delta),
diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml
index bf92e05ee186..e61c48ee3c83 100644
--- a/rust/src/cmd/Cargo.toml
+++ b/rust/src/cmd/Cargo.toml
@@ -33,6 +33,7 @@ vllm-chat.workspace = true
 vllm-engine-core-client.workspace = true
 vllm-managed-engine.workspace = true
 vllm-server.workspace = true
+vllm-text.workspace = true
 vllm-tracing.workspace = true
 
 [dev-dependencies]
diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs
index c0b4c6d29550..c5a7c75baff8 100644
--- a/rust/src/cmd/src/cli.rs
+++ b/rust/src/cmd/src/cli.rs
@@ -24,6 +24,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as};
 use thiserror_ext::AsReport as _;
 use uuid::Uuid;
 use vllm_chat::GenerationConfigMode;
+use vllm_chat::ToolStrictLevel;
 use vllm_chat::multimodal::MmLimitPerPrompt;
 use vllm_engine_core_client::TransportMode;
 use vllm_managed_engine::ManagedEngineConfig;
@@ -33,6 +34,7 @@ use vllm_server::{
     DEFAULT_KEEP_ALIVE_TIMEOUT, HttpListenerMode, LoraModulePath, ParserSelection, RenderConfig,
     RendererSelection,
 };
+use vllm_text::backend::hf::HfOverrides;
 
 use crate::cli::ssl::SslArgs;
 use crate::cli::unsupported::UnsupportedArgs;
@@ -110,6 +112,12 @@ pub enum BenchCommand {
 pub struct RenderArgs {
     /// Model identifier or local model directory containing tokenizer files.
     model: String,
+    /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
+    #[arg(long)]
+    revision: Option<String>,
+    /// JSON Merge Patch (RFC 7396) for config.json; null removes a field.
+    #[arg(long, value_parser = parse_json::<HfOverrides>, default_value = "{}", value_name = "JSON")]
+    hf_overrides: HfOverrides,
     /// HTTP bind host.
     #[arg(long, default_value = "127.0.0.1")]
     host: String,
@@ -127,6 +135,10 @@ pub struct RenderArgs {
     /// `none` to disable parsing.
     #[arg(long, default_value_t)]
     reasoning_parser: ParserSelection,
+    /// Server-side floor for structural-tag based tool calling: `auto`,
+    /// `function`, or `parameter`.
+    #[arg(long, default_value_t)]
+    tool_strict_level: ToolStrictLevel,
     /// Select the native chat renderer implementation.
     #[arg(long = "tokenizer-mode", default_value_t)]
     renderer: RendererSelection,
@@ -157,11 +169,14 @@ impl RenderArgs {
 
         RenderConfig {
             model: self.model,
+            revision: self.revision,
+            hf_overrides: self.hf_overrides,
             served_model_name: self.served_model_name,
             host: self.host,
             port: self.port,
             tool_call_parser: self.tool_call_parser,
             reasoning_parser: self.reasoning_parser,
+            tool_strict_level: self.tool_strict_level,
             renderer: self.renderer,
             chat_template: self.chat_template,
             default_chat_template_kwargs: self.default_chat_template_kwargs.unwrap_or_default(),
@@ -195,6 +210,17 @@ pub struct SharedRuntimeArgs {
     /// public model ID.
     pub model: String,
 
+    /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
+    #[arg(long)]
+    #[serde(default)]
+    pub revision: Option<String>,
+
+    /// JSON Merge Patch (RFC 7396) for config.json; null removes a field.
+    /// Objects merge recursively and arrays/scalars replace existing values.
+    #[arg(long, value_parser = parse_json::<HfOverrides>, default_value = "{}", value_name = "JSON")]
+    #[serde(default)]
+    pub hf_overrides: HfOverrides,
+
     /// The source of generation-config sampling defaults. `"auto"` loads the
     /// model's defaults, while `"vllm"` uses vLLM's neutral defaults.
     #[arg(long, default_value_t)]
@@ -221,6 +247,14 @@ pub struct SharedRuntimeArgs {
     #[arg(long, default_value_t)]
     #[serde(default = "default_py_bootstrap_parser_selection")]
     pub reasoning_parser: ParserSelection,
+    /// Server-side floor for structural-tag based tool calling, applied on
+    /// top of the per-tool `strict` field: `auto` follows the request's
+    /// tool choice and per-tool strictness,
+    /// `function` constrains the tool-call envelope for every request with
+    /// tools, `parameter` additionally pins argument schemas.
+    #[arg(long, default_value_t)]
+    #[serde(default)]
+    pub tool_strict_level: ToolStrictLevel,
     /// Select the chat renderer implementation.
     #[arg(long = "tokenizer-mode", default_value_t)]
     #[serde(default, rename = "tokenizer_mode")]
@@ -320,6 +354,23 @@ pub struct SharedRuntimeArgs {
     #[serde(default)]
     pub enable_request_id_headers: bool,
 
+    /// Register the scale-out `/inference/v1/generate` endpoint.
+    #[arg(
+        long,
+        default_missing_value = "true",
+        num_args = 0..=1
+    )]
+    #[serde(default)]
+    pub enable_scale_out: bool,
+
+    /// Send an SSE keep-alive comment line every this many seconds when a
+    /// streaming response is idle (queued, prefill, or between tokens), to
+    /// prevent reverse proxies/tunnels with read timeouts from closing the
+    /// connection. Defaults to 0, which disables keep-alive comments entirely.
+    #[arg(long, default_value_t = 0)]
+    #[serde(default)]
+    pub sse_keep_alive_interval: u64,
+
     /// If provided, the server will require one of these keys to be presented
     /// in the Authorization header.
     #[educe(Debug(ignore))]
@@ -481,11 +532,14 @@ impl SharedRuntimeArgs {
                 None => CoordinatorMode::None,
             },
             model: self.model,
+            revision: self.revision,
+            hf_overrides: self.hf_overrides,
             generation_config: self.generation_config,
             served_model_name: self.served_model_name,
             listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd },
             tool_call_parser: self.tool_call_parser,
             reasoning_parser: self.reasoning_parser,
+            tool_strict_level: self.tool_strict_level,
             renderer: self.renderer,
             language_model_only: self.language_model_only,
             chat_template: self.chat_template,
@@ -536,11 +590,14 @@ impl SharedRuntimeArgs {
             },
             coordinator_mode: CoordinatorMode::MaybeInProc,
             model: self.model,
+            revision: self.revision,
+            hf_overrides: self.hf_overrides,
             generation_config: self.generation_config,
             served_model_name: self.served_model_name,
             listener_mode,
             tool_call_parser: self.tool_call_parser,
             reasoning_parser: self.reasoning_parser,
+            tool_strict_level: self.tool_strict_level,
             renderer: self.renderer,
             language_model_only: self.language_model_only,
             chat_template: self.chat_template,
@@ -566,6 +623,9 @@ impl SharedRuntimeArgs {
             enable_log_requests: self.enable_log_requests,
             enable_prompt_tokens_details: self.enable_prompt_tokens_details,
             enable_request_id_headers: self.enable_request_id_headers,
+            enable_scale_out: self.enable_scale_out,
+            sse_keep_alive_interval: (self.sse_keep_alive_interval > 0)
+                .then(|| Duration::from_secs(self.sse_keep_alive_interval)),
         }
     }
 
@@ -740,9 +800,13 @@ impl ServeArgs {
         let reasoning_parser =
             effective_engine_reasoning_parser(&self.runtime.reasoning_parser, &self.runtime.model);
         let profiler_config = self.runtime.profiler_config_json();
+        let hf_overrides = (!self.runtime.hf_overrides.is_empty()).then(|| {
+            serde_json::to_string(&self.runtime.hf_overrides).expect("JSON object serializes")
+        });
 
         self.managed_engine.clone().into_config(
             self.runtime.model.clone(),
+            self.runtime.revision.clone(),
             self.runtime.max_logprobs,
             profiler_config,
             reasoning_parser.as_deref(),
@@ -751,6 +815,7 @@ impl ServeArgs {
             self.runtime.shutdown_timeout,
             handshake_port,
             self.runtime.limit_mm_per_prompt_json(),
+            hf_overrides,
         )
     }
 }
diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs
index e9411891e010..dac00ac9833c 100644
--- a/rust/src/cmd/src/cli/tests.rs
+++ b/rust/src/cmd/src/cli/tests.rs
@@ -36,6 +36,8 @@ fn render_args_build_config_without_tls() {
         "vllm-rs",
         "render",
         "Qwen/Qwen2.5-0.5B-Instruct",
+        "--revision",
+        "release",
         "--port",
         "8080",
         "--max-model-len",
@@ -55,6 +57,7 @@ fn render_args_build_config_without_tls() {
     let config = args.into_config();
 
     assert_eq!(config.model, "Qwen/Qwen2.5-0.5B-Instruct");
+    assert_eq!(config.revision.as_deref(), Some("release"));
     assert_eq!(config.host, "127.0.0.1");
     assert_eq!(config.port, 8080);
     assert_eq!(config.max_model_len, Some(32768));
@@ -173,10 +176,15 @@ fn serve_args_forward_python_flags_with_separator() {
                     uds: None,
                     runtime: SharedRuntimeArgs {
                         model: "Qwen/Qwen3-0.6B",
+                        revision: None,
+                        hf_overrides: HfOverrides(
+                            {},
+                        ),
                         generation_config: Auto,
                         engine_ready_timeout_secs: 600,
                         tool_call_parser: Auto,
                         reasoning_parser: Auto,
+                        tool_strict_level: Auto,
                         renderer: Auto,
                         language_model_only: false,
                         max_logprobs: None,
@@ -191,6 +199,8 @@ fn serve_args_forward_python_flags_with_separator() {
                         enable_log_requests: false,
                         enable_prompt_tokens_details: false,
                         enable_request_id_headers: false,
+                        enable_scale_out: false,
+                        sse_keep_alive_interval: 0,
                         disable_log_stats: false,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
@@ -1007,10 +1017,15 @@ fn frontend_args_accept_json() {
                     data_parallel_size: None,
                     runtime: SharedRuntimeArgs {
                         model: "Qwen/Qwen3-0.6B",
+                        revision: None,
+                        hf_overrides: HfOverrides(
+                            {},
+                        ),
                         generation_config: Auto,
                         engine_ready_timeout_secs: 600,
                         tool_call_parser: None,
                         reasoning_parser: None,
+                        tool_strict_level: Auto,
                         renderer: Auto,
                         language_model_only: false,
                         max_logprobs: None,
@@ -1025,6 +1040,8 @@ fn frontend_args_accept_json() {
                         enable_log_requests: false,
                         enable_prompt_tokens_details: false,
                         enable_request_id_headers: false,
+                        enable_scale_out: false,
+                        sse_keep_alive_interval: 0,
                         disable_log_stats: false,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
@@ -1113,6 +1130,73 @@ fn frontend_args_json_ignores_engine_owned_max_model_len() {
     assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B");
 }
 
+#[test]
+fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() {
+    let patch = r#"{"text_config":{"rope_parameters":{"factor":4}},"sliding_window":null}"#;
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--hf-overrides",
+        patch,
+    ])
+    .unwrap();
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args")
+    };
+    let expected: serde_json::Value = serde_json::from_str(patch).unwrap();
+    assert_eq!(
+        serde_json::to_value(&args.runtime.hf_overrides).unwrap(),
+        expected
+    );
+    let engine = args.to_managed_engine_config(1234);
+    let index = engine.python_args.iter().position(|arg| arg == "--hf-overrides").unwrap();
+    assert_eq!(
+        serde_json::from_str::<serde_json::Value>(&engine.python_args[index + 1]).unwrap(),
+        expected
+    );
+    assert!(!engine.python_args.iter().any(|arg| arg == "--hf-config-path"));
+    let config = args.to_frontend_config("tcp://127.0.0.1:1234".to_string());
+    assert_eq!(serde_json::to_value(config.hf_overrides).unwrap(), expected);
+
+    let payload = serde_json::json!({"model_tag":"Qwen/Qwen3-0.6B", "hf_overrides":expected});
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "frontend",
+        "--listen-fd",
+        "3",
+        "--input-address",
+        "ipc:///tmp/input.sock",
+        "--output-address",
+        "ipc:///tmp/output.sock",
+        "--args-json",
+        &payload.to_string(),
+    ])
+    .unwrap();
+    let Command::Frontend(args) = cli.command else {
+        panic!("expected frontend args")
+    };
+    assert_eq!(
+        serde_json::to_value(args.into_config().hf_overrides).unwrap(),
+        expected
+    );
+}
+
+#[test]
+fn hf_overrides_reject_non_object_inputs() {
+    for patch in ["null", "[]", "1", r#""callable""#] {
+        let error = Cli::try_parse_from([
+            "vllm-rs",
+            "serve",
+            "Qwen/Qwen3-0.6B",
+            "--hf-overrides",
+            patch,
+        ])
+        .unwrap_err();
+        assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation);
+    }
+}
+
 #[test]
 fn frontend_args_json_accepts_supported_non_default_fields() {
     let cli = Cli::try_parse_from([
@@ -1125,7 +1209,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
         "--output-address",
         "ipc:///tmp/output.sock",
         "--args-json",
-        r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#,
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","revision":"release","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#,
     ])
     .unwrap();
 
@@ -1133,6 +1217,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
         panic!("expected frontend args");
     };
     assert_eq!(args.runtime.engine_ready_timeout_secs, 42);
+    assert_eq!(args.runtime.revision.as_deref(), Some("release"));
     assert_eq!(args.runtime.generation_config, GenerationConfigMode::Vllm);
     assert_eq!(
         args.runtime.tool_call_parser,
@@ -1604,10 +1689,15 @@ fn serve_args_accept_handshake_aliases() {
                     uds: None,
                     runtime: SharedRuntimeArgs {
                         model: "Qwen/Qwen3-0.6B",
+                        revision: None,
+                        hf_overrides: HfOverrides(
+                            {},
+                        ),
                         generation_config: Auto,
                         engine_ready_timeout_secs: 600,
                         tool_call_parser: Auto,
                         reasoning_parser: Auto,
+                        tool_strict_level: Auto,
                         renderer: Auto,
                         language_model_only: false,
                         max_logprobs: None,
@@ -1622,6 +1712,8 @@ fn serve_args_accept_handshake_aliases() {
                         enable_log_requests: false,
                         enable_prompt_tokens_details: false,
                         enable_request_id_headers: false,
+                        enable_scale_out: false,
+                        sse_keep_alive_interval: 0,
                         disable_log_stats: false,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
@@ -1752,6 +1844,10 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
             },
             coordinator_mode: MaybeInProc,
             model: "Qwen/Qwen3-0.6B",
+            revision: None,
+            hf_overrides: HfOverrides(
+                {},
+            ),
             generation_config: Auto,
             served_model_name: [],
             listener_mode: BindTcp {
@@ -1760,6 +1856,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
             },
             tool_call_parser: Auto,
             reasoning_parser: Auto,
+            tool_strict_level: Auto,
             renderer: Auto,
             language_model_only: false,
             chat_template: None,
@@ -1772,6 +1869,8 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
                 enable_log_requests: false,
                 enable_prompt_tokens_details: false,
                 enable_request_id_headers: false,
+                enable_scale_out: false,
+                sse_keep_alive_interval: None,
             },
             cors: CorsConfig {
                 allow_origins: [
@@ -1839,6 +1938,10 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
             },
             coordinator_mode: MaybeInProc,
             model: "Qwen/Qwen3-0.6B",
+            revision: None,
+            hf_overrides: HfOverrides(
+                {},
+            ),
             generation_config: Auto,
             served_model_name: [],
             listener_mode: BindTcp {
@@ -1847,6 +1950,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
             },
             tool_call_parser: Auto,
             reasoning_parser: Auto,
+            tool_strict_level: Auto,
             renderer: Auto,
             language_model_only: false,
             chat_template: None,
@@ -1859,6 +1963,8 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
                 enable_log_requests: false,
                 enable_prompt_tokens_details: false,
                 enable_request_id_headers: false,
+                enable_scale_out: false,
+                sse_keep_alive_interval: None,
             },
             cors: CorsConfig {
                 allow_origins: [
@@ -1948,6 +2054,10 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
                 address: "tcp://127.0.0.1:7000",
             },
             model: "Qwen/Qwen3-0.6B",
+            revision: None,
+            hf_overrides: HfOverrides(
+                {},
+            ),
             generation_config: Auto,
             served_model_name: [],
             listener_mode: InheritedFd {
@@ -1955,6 +2065,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
             },
             tool_call_parser: None,
             reasoning_parser: None,
+            tool_strict_level: Auto,
             renderer: Auto,
             language_model_only: false,
             chat_template: None,
@@ -1967,6 +2078,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
                 enable_log_requests: false,
                 enable_prompt_tokens_details: false,
                 enable_request_id_headers: false,
+                enable_scale_out: false,
+                sse_keep_alive_interval: None,
             },
             cors: CorsConfig {
                 allow_origins: [
@@ -2087,3 +2200,18 @@ fn frontend_args_json_disables_profiling_when_profiler_type_is_null() {
     let config = args.into_config();
     assert_eq!(config.profiler, None);
 }
+
+#[test]
+fn serve_revision_reaches_frontend_and_managed_engine() {
+    let cli =
+        Cli::try_parse_from(["vllm-rs", "serve", "test/model", "--revision", "release"]).unwrap();
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args");
+    };
+    assert_eq!(
+        args.to_frontend_config("tcp://localhost:1234".into()).revision.as_deref(),
+        Some("release")
+    );
+    let engine = args.to_managed_engine_config(1234);
+    assert!(engine.python_args.windows(2).any(|args| args == ["--revision", "release"]));
+}
diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs
index 29ea5d06b0e7..5088063914ad 100644
--- a/rust/src/cmd/src/cli/unsupported.rs
+++ b/rust/src/cmd/src/cli/unsupported.rs
@@ -235,11 +235,6 @@ pub struct EngineUnsupportedArgs {
     #[arg(long, default_missing_value = "true", num_args = 0..=1)]
     pub hf_token: Option<Unsupported>,
 
-    /// If a dictionary, contains arguments to be forwarded to the Hugging Face
-    /// config. If a callable, it is called to update the HuggingFace config.
-    #[arg(long)]
-    pub hf_overrides: Option<Unsupported>,
-
     /// Overrides or sets generation config. e.g. `{"temperature": 0.5}`. If
     /// used with `--generation-config auto`, the override parameters will be
     /// merged with the default config from the model. If used with
@@ -251,6 +246,13 @@ pub struct EngineUnsupportedArgs {
     #[arg(long)]
     pub io_processor_plugin: Option<Unsupported>,
 
+    /// Number of worker threads in the renderer thread pool. The pool is
+    /// consumed by the async renderer path to parallelize tokenization, chat
+    /// template rendering, and multimodal preprocessing across concurrent
+    /// requests.
+    #[arg(long)]
+    pub renderer_num_workers: Option<Noop>,
+
     /// Path to a dynamically reasoning parser plugin that can be dynamically
     /// loaded and registered.
     #[arg(long)]
@@ -286,6 +288,16 @@ pub struct EngineUnsupportedArgs {
     )]
     pub data_parallel_external_lb: Option<Unsupported>,
 
+    /// Enable fault tolerance for detailed error recovery, such as scaling
+    /// down fault DPEngineCore.
+    #[arg(
+        long,
+        visible_alias = "no-enable-fault-tolerance",
+        default_missing_value = "true",
+        num_args = 0..=1
+    )]
+    pub enable_fault_tolerance: Option<Unsupported>,
+
     /// This feature is work in progress and no prefill optimization takes place
     /// with this flag enabled currently.
     #[arg(
@@ -313,6 +325,28 @@ pub struct EngineUnsupportedArgs {
     #[arg(long)]
     pub mm_processor_cache_type: Option<Unsupported>,
 
+    /// Hash algorithm to use for multi-modal input caching. Use `"sha256"` or
+    /// `"sha512"` for FIPS-compliant deployments.
+    #[arg(long)]
+    pub mm_hasher_algorithm: Option<Unsupported>,
+
+    /// IPC (inter-process communication) method for multimodal tensors.
+    /// - "direct_rpc": Use msgspec serialization via RPC
+    /// - "torch_shm": Use torch.multiprocessing shared memory for zero-copy IPC
+    #[arg(long)]
+    pub mm_tensor_ipc: Option<Unsupported>,
+
+    /// Device the HF multi-modal processor runs the image/video transform on.
+    /// Convenience for `--mm-processor-kwargs '{"device": ...}'`.
+    #[arg(long)]
+    pub mm_processor_device: Option<Unsupported>,
+
+    /// Amount of GPU memory (in GiB) sequestered on the engine's device for
+    /// GPU-side multimodal work in the API-server (frontend) process, such as
+    /// hardware video decoding.
+    #[arg(long)]
+    pub mm_ipc_gpu_memory_gb: Option<Unsupported>,
+
     /// Dictionary mapping specific modalities to LoRA model paths.
     #[arg(long)]
     pub default_mm_loras: Option<Unsupported>,
@@ -329,6 +363,17 @@ pub struct EngineUnsupportedArgs {
     #[arg(long)]
     pub stream_interval: Option<Unsupported>,
 
+    /// Maximum number of requests that can be in-flight (waiting or running)
+    /// at the same time. When the limit is reached, new requests are rejected
+    /// with HTTP 503 so the client can retry on another instance.
+    #[arg(long)]
+    pub max_num_queued_reqs: Option<Unsupported>,
+
+    /// Maximum total prompt tokens of requests currently in the prefill phase.
+    /// When the limit is reached, new requests are rejected with HTTP 503.
+    #[arg(long)]
+    pub max_num_queued_tokens: Option<Unsupported>,
+
     /// Structured outputs configuration.
     #[arg(long)]
     pub structured_outputs_config: Option<Noop>,
@@ -424,6 +469,15 @@ pub struct ServerUnsupportedArgs {
     #[arg(long)]
     pub max_log_len: Option<Unsupported>,
 
+    /// If set to True, include per-request timing metrics in API responses.
+    #[arg(
+        long,
+        visible_alias = "no-enable-per-request-metrics",
+        default_missing_value = "true",
+        num_args = 0..=1
+    )]
+    pub enable_per_request_metrics: Option<Unsupported>,
+
     /// If set to True, enable tracking server_load_metrics in the app state.
     #[arg(
         long,
@@ -497,6 +551,15 @@ pub struct ServerUnsupportedArgs {
     )]
     pub tokens_only: Option<Unsupported>,
 
+    /// Controls the `system_fingerprint` field on responses: `full`, `hash`,
+    /// `custom`, or `none`.
+    #[arg(long)]
+    pub fingerprint_mode: Option<Unsupported>,
+
+    /// Literal fingerprint string used when `--fingerprint-mode=custom`.
+    #[arg(long)]
+    pub fingerprint_value: Option<Unsupported>,
+
     /// Log level for uvicorn.
     #[arg(long)]
     pub uvicorn_log_level: Option<Unsupported>,
@@ -567,4 +630,14 @@ pub struct ServerUnsupportedArgs {
         num_args = 0..=1
     )]
     pub enable_offline_docs: Option<Unsupported>,
+
+    /// If set, run pooling score MaxSim on GPU in the API server process.
+    /// Can significantly improve late-interaction scoring performance.
+    #[arg(
+        long,
+        visible_alias = "no-enable-flash-late-interaction",
+        default_missing_value = "true",
+        num_args = 0..=1
+    )]
+    pub enable_flash_late_interaction: Option<Noop>,
 }
diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs
index 63689ea3db83..aa6aacf0e017 100644
--- a/rust/src/engine-core-client/src/client.rs
+++ b/rust/src/engine-core-client/src/client.rs
@@ -510,6 +510,15 @@ impl EngineCoreClient {
         self.engines.iter().map(|engine| engine.ready_response.num_gpu_blocks).sum()
     }
 
+    /// Return the effective attention block size if all engines report the same value.
+    pub fn effective_attention_block_size(&self) -> Option<u64> {
+        let size = self.ready_response().effective_attention_block_size?;
+        self.engines
+            .iter()
+            .all(|engine| engine.ready_response.effective_attention_block_size == Some(size))
+            .then_some(size)
+    }
+
     /// Return the minimum engine-reported `max_model_len` across all engines.
     ///
     /// This is the auto-fitted value after KV cache profiling and may differ
@@ -902,13 +911,22 @@ impl EngineCoreClient {
         Ok(())
     }
 
-    /// Wake the engine from sleep, optionally limiting the wake-up to specific
-    /// tags.
-    pub async fn wake_up(&self, tags: Option<Vec<String>>) -> Result<()> {
-        self.call_utility::<(), _>("wake_up", (tags,)).await?;
+    /// Release KV cache memory while keeping model weights resident.
+    pub async fn release_kv_cache_memory(&self) -> Result<()> {
+        self.call_utility::<(), _>("release_kv_cache_memory", ()).await?;
         Ok(())
     }
 
+    /// Wake the engine from sleep, optionally limiting the wake-up to specific
+    /// tags, and return whether every engine is fully awake.
+    pub async fn wake_up(&self, tags: Option<Vec<String>>) -> Result<bool> {
+        Ok(self
+            .call_utility::<bool, _>("wake_up", (tags,))
+            .await?
+            .into_iter()
+            .all(|fully_awake| fully_awake))
+    }
+
     /// Pause the scheduler so generation can be halted
     pub async fn pause_scheduler(&self, mode: PauseMode, clear_cache: bool) -> Result<()> {
         self.call_utility::<(), _>("pause_scheduler", (mode, clear_cache)).await?;
diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs
index 117473ce47cc..7318fa01b8ce 100644
--- a/rust/src/engine-core-client/src/client/imp.rs
+++ b/rust/src/engine-core-client/src/client/imp.rs
@@ -19,7 +19,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti
 use crate::client::stream::EngineCoreStreamOutput;
 use crate::client::{AbortCause, AbortRequest};
 use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
-use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
+use crate::metrics::{IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder};
 use crate::protocol::encode_msgpack;
 use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
 use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
@@ -47,6 +47,7 @@ pub(crate) struct ClientInner {
     /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created.
     msgpack_zero_copy_threshold: usize,
     scheduler_stats_recorder: SchedulerStatsRecorder,
+    iteration_metrics: BTreeMap<u32, IterationMetricHandles>,
     request_reg: Mutex<RequestRegistry>,
     utility_reg: Mutex<UtilityRegistry>,
     health_error: ArcSwapOption<Error>,
@@ -64,12 +65,21 @@ impl ClientInner {
     ) -> Self {
         let scheduler_stats_recorder =
             SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines);
+        let iteration_metrics = engines
+            .iter()
+            .filter_map(|engine| {
+                let engine = engine.engine_id.engine_index()?;
+                let handles = IterationMetricHandles::new(&METRICS.request, &model_name, engine);
+                Some((engine, handles))
+            })
+            .collect();
         Self {
             input_send,
             handle,
             model_name,
             msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(),
             scheduler_stats_recorder,
+            iteration_metrics,
             request_reg: Mutex::new(RequestRegistry::new(engines)),
             utility_reg: Mutex::new(UtilityRegistry::default()),
             health_error: ArcSwapOption::empty(),
@@ -415,6 +425,8 @@ pub(crate) async fn run_output_dispatcher_loop(
 
             match outputs {
                 EngineCoreOutputs::RequestBatch(batch) => {
+                    let has_outputs = !batch.outputs.is_empty();
+                    let mut iteration_tokens = 0_u64;
                     let senders = inner.take_senders_for_outputs(&batch.outputs);
                     for (output, sender) in batch.outputs.into_iter().zip(senders) {
                         let request_id = output.request_id.clone();
@@ -423,6 +435,12 @@ pub(crate) async fn run_output_dispatcher_loop(
                             continue;
                         };
 
+                        iteration_tokens += output.new_token_ids.len() as u64;
+                        if let Some(prefill) = &output.prefill_stats {
+                            // The engine emits prefill_stats once, on the first output.
+                            iteration_tokens += u64::from(prefill.num_computed_tokens);
+                        }
+
                         let wrapped_output = EngineCoreStreamOutput {
                             engine_index: batch.engine_index,
                             timestamp: batch.timestamp,
@@ -433,6 +451,12 @@ pub(crate) async fn run_output_dispatcher_loop(
                         }
                     }
 
+                    if has_outputs
+                        && let Some(handles) = inner.iteration_metrics.get(&batch.engine_index)
+                    {
+                        handles.iteration_tokens_total.observe(iteration_tokens as f64);
+                    }
+
                     // The sender for normally-finished requests should have already been removed
                     // from the registry when their final output was dispatched
                     // above. This serves as a safety net to capture any
diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs
index e15258d6081d..b392d5260185 100644
--- a/rust/src/engine-core-client/src/metrics.rs
+++ b/rust/src/engine-core-client/src/metrics.rs
@@ -8,8 +8,8 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use vllm_metrics::{
     EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames,
     LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily,
-    MooncakeOperationLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
-    WaitingReasonLabels,
+    MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics,
+    U64Counter, U64Gauge, WaitingReasonLabels,
 };
 
 use crate::protocol::stats::{
@@ -20,6 +20,23 @@ use crate::transport::ConnectedEngine;
 const WAITING_REASON_CAPACITY: &str = "capacity";
 const WAITING_REASON_DEFERRED: &str = "deferred";
 
+/// Cached output-batch metric handles for one model and engine index.
+pub(crate) struct IterationMetricHandles {
+    pub iteration_tokens_total: HistogramMetric,
+}
+
+impl IterationMetricHandles {
+    pub(crate) fn new(metrics: &RequestMetrics, model_name: &str, engine: u32) -> Self {
+        let labels = EngineLabels {
+            model_name: model_name.to_string(),
+            engine,
+        };
+        Self {
+            iteration_tokens_total: metrics.iteration_tokens_total.get_or_create_owned(&labels),
+        }
+    }
+}
+
 /// Cached scheduler-stats metric handles for all engines connected to one
 /// frontend client.
 pub(crate) struct SchedulerStatsRecorder {
diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs
index 77d72b525fbd..958fcf8aa1b3 100644
--- a/rust/src/engine-core-client/src/mock_engine.rs
+++ b/rust/src/engine-core-client/src/mock_engine.rs
@@ -72,6 +72,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
         weight_transfer_backend: None,
         enable_sleep_mode: false,
         supports_draft_weight_updates: false,
+        effective_attention_block_size: None,
     }
 }
 
diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs
index c14c19a50373..52d897be52dc 100644
--- a/rust/src/engine-core-client/src/protocol/handshake.rs
+++ b/rust/src/engine-core-client/src/protocol/handshake.rs
@@ -104,6 +104,9 @@ pub struct EngineCoreReadyResponse {
     /// Whether the engine has a speculative draft model that can be updated.
     #[serde(default)]
     pub supports_draft_weight_updates: bool,
+    /// Full-attention block size in tokens after initialization, or unavailable.
+    #[serde(default)]
+    pub effective_attention_block_size: Option<u64>,
 }
 
 /// Frontend-owned ZMQ addresses that are sent to the engine during startup
diff --git a/rust/src/engine-core-client/src/protocol/kv_hints.rs b/rust/src/engine-core-client/src/protocol/kv_hints.rs
new file mode 100644
index 000000000000..418ca2fcedb6
--- /dev/null
+++ b/rust/src/engine-core-client/src/protocol/kv_hints.rs
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::collections::BTreeMap;
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct KvHintAction {
+    pub action_id: String,
+    pub action_type: String,
+    pub action_version: String,
+    pub payload: BTreeMap<String, serde_json::Value>,
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct KvHintsEnvelope {
+    pub protocol_version: String,
+    pub message_id: String,
+    pub actions: Vec<KvHintAction>,
+}
diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs
index d5bc40e06159..72d852ba2ced 100644
--- a/rust/src/engine-core-client/src/protocol/logprobs.rs
+++ b/rust/src/engine-core-client/src/protocol/logprobs.rs
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-mod array;
+pub(crate) mod array;
 #[cfg(test)]
 mod tests;
 mod wire;
@@ -25,10 +25,13 @@ use crate::protocol::tensor::{WireArrayData, WireNdArray};
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct TokenLogprob {
     pub token_id: u32,
+    /// Preserves the engine's value, including NaN and infinities.
     pub logprob: f32,
     /// The sampled/selected token uses its actual vocab rank. Remaining entries
     /// use 1-based top-k ranks matching the engine's returned candidate
     /// order.
+    /// A sampled/selected rank of 0 occurs when its logprob is NaN: the engine's
+    /// `(logprobs >= selected_logprob).sum(-1)` counts no matching values.
     pub rank: u32,
 }
 
@@ -53,10 +56,6 @@ impl PositionLogprobs {
                 logprobs.len()
             );
         }
-        if sampled_rank == 0 {
-            bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs");
-        }
-
         let mut entries = Vec::with_capacity(token_ids.len());
         for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() {
             let rank = if index == 0 {
diff --git a/rust/src/engine-core-client/src/protocol/logprobs/array.rs b/rust/src/engine-core-client/src/protocol/logprobs/array.rs
index 4adf1ae9755f..919ef80f46a3 100644
--- a/rust/src/engine-core-client/src/protocol/logprobs/array.rs
+++ b/rust/src/engine-core-client/src/protocol/logprobs/array.rs
@@ -50,7 +50,7 @@ pub(super) fn decode_array2_u32(
     })
 }
 
-pub(super) fn decode_array1_u32(
+pub(crate) fn decode_array1_u32(
     value: WireNdArray,
     field: &str,
     frames: &[Bytes],
diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs
index 2470460b52a6..f69ff22f1ae0 100644
--- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs
+++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs
@@ -396,3 +396,121 @@ fn rejects_zero_column_logprobs_with_rows() {
         "new_logprobs: zero-column logprobs payload with 2 rows"
     );
 }
+
+fn logprobs_value(ids: &[i64], probs: &[f32], ranks: &[i64]) -> Value {
+    let rows = ranks.len();
+    let cols = ids.len() / rows;
+    Value::Array(vec![
+        ndarray_value(
+            "<i8",
+            &[rows, cols],
+            Value::Ext(3, ids.iter().flat_map(|v| v.to_le_bytes()).collect()),
+        ),
+        ndarray_value(
+            "<f4",
+            &[rows, cols],
+            Value::Ext(3, probs.iter().flat_map(|v| v.to_le_bytes()).collect()),
+        ),
+        ndarray_value(
+            "<i8",
+            &[rows],
+            Value::Ext(3, ranks.iter().flat_map(|v| v.to_le_bytes()).collect()),
+        ),
+        Value::Nil,
+    ])
+}
+
+#[test]
+fn zero_sampled_rank_does_not_fail_frame_mates() {
+    // A batched EngineCoreOutputs frame where one request's row has rank 0
+    // must still resolve; the other requests' logprobs stay intact.
+    let frames = vec![Bytes::from(encode_value(&Value::Array(vec![
+        Value::from(0),
+        Value::Array(vec![
+            Value::Array(vec![
+                Value::from("req-nan"),
+                Value::Array(vec![Value::from(7), Value::from(8)]),
+                logprobs_value(&[1, 2], &[f32::NAN, f32::NEG_INFINITY], &[0]),
+                logprobs_value(&[7, 8], &[f32::NAN, f32::NEG_INFINITY], &[0]),
+                Value::Nil,
+                Value::from(EngineCoreFinishReason::Length as u8),
+            ]),
+            Value::Array(vec![
+                Value::from("req-ok"),
+                Value::Array(vec![Value::from(9)]),
+                logprobs_value(&[3, 4], &[-3.0, -0.5], &[7]),
+                Value::Nil,
+                Value::Nil,
+                Value::from(EngineCoreFinishReason::Length as u8),
+            ]),
+        ]),
+        Value::Nil,
+        Value::from(0.0),
+        Value::Nil,
+        Value::Nil,
+    ])))];
+    let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
+
+    let bad = decoded.outputs[0].new_logprobs.as_ref().unwrap();
+    let prompt = decoded.outputs[0].new_prompt_logprobs_tensors.as_ref().unwrap();
+    let good = decoded.outputs[1].new_logprobs.as_ref().unwrap();
+    expect_test::expect![[r#"
+        (
+            Logprobs {
+                positions: [
+                    PositionLogprobs {
+                        entries: [
+                            TokenLogprob {
+                                token_id: 1,
+                                logprob: NaN,
+                                rank: 0,
+                            },
+                            TokenLogprob {
+                                token_id: 2,
+                                logprob: -inf,
+                                rank: 1,
+                            },
+                        ],
+                    },
+                ],
+            },
+            Logprobs {
+                positions: [
+                    PositionLogprobs {
+                        entries: [
+                            TokenLogprob {
+                                token_id: 7,
+                                logprob: NaN,
+                                rank: 0,
+                            },
+                            TokenLogprob {
+                                token_id: 8,
+                                logprob: -inf,
+                                rank: 1,
+                            },
+                        ],
+                    },
+                ],
+            },
+            Logprobs {
+                positions: [
+                    PositionLogprobs {
+                        entries: [
+                            TokenLogprob {
+                                token_id: 3,
+                                logprob: -3.0,
+                                rank: 7,
+                            },
+                            TokenLogprob {
+                                token_id: 4,
+                                logprob: -0.5,
+                                rank: 1,
+                            },
+                        ],
+                    },
+                ],
+            },
+        )
+    "#]]
+    .assert_debug_eq(&(&**bad, &**prompt, &**good));
+}
diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs
index 23bd09d8d8a2..0586dde135ce 100644
--- a/rust/src/engine-core-client/src/protocol/mod.rs
+++ b/rust/src/engine-core-client/src/protocol/mod.rs
@@ -24,12 +24,15 @@ fn is_false(v: &bool) -> bool {
 
 pub mod dtype;
 pub mod handshake;
+pub mod kv_hints;
 pub mod logprobs;
 pub mod lora;
 pub mod multimodal;
 pub mod output;
 pub mod request;
 pub mod sampling;
+pub mod sampling_mask;
+mod serde_utils;
 pub mod stats;
 pub mod structured_outputs;
 pub mod tensor;
diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs
index 85de8c4e3945..0edd0ce788bd 100644
--- a/rust/src/engine-core-client/src/protocol/output.rs
+++ b/rust/src/engine-core-client/src/protocol/output.rs
@@ -9,10 +9,13 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use serde_default::DefaultFromSerde;
 use serde_repr::{Deserialize_repr, Serialize_repr};
 use serde_tuple::{Deserialize_tuple, Serialize_tuple};
+use serde_with::serde_as;
 
+use super::serde_utils::AllowTrailingFields;
 use super::utility::UtilityOutput;
 use crate::error::{Error, Result, ext_value_decode};
 use crate::protocol::logprobs::MaybeWireLogprobs;
+use crate::protocol::sampling_mask::MaybeWireSamplingMask;
 use crate::protocol::stats::{PrefillStats, SchedulerStats};
 use crate::protocol::{OpaqueValue, decode_msgpack};
 
@@ -125,12 +128,38 @@ pub struct EngineCoreOutput {
     #[serde(default)]
     pub mm_cache_miss_hashes: Option<Vec<String>>,
     #[serde(default)]
-    pub new_sampling_mask: Option<OpaqueValue>,
+    pub new_sampling_mask: Option<MaybeWireSamplingMask>,
     /// Per-request speculative-decoding acceptance metrics, set on the final
-    /// output when `--per-request-spec-decode-metrics` is enabled. Opaque here;
-    /// the Rust frontend does not yet surface it in responses.
+    /// output when `--per-request-spec-decode-metrics` is enabled.
     #[serde(default)]
-    pub spec_decode_metrics: Option<OpaqueValue>,
+    pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
+}
+
+/// Raw per-sequence speculative-decoding accumulator.
+///
+/// Python models this as a plain `@dataclass`, so it is serialized by msgspec
+/// as a map (named fields). Derived values such as acceptance rates are
+/// computed by the frontend, not carried on the wire.
+///
+/// Original Python definition: `RequestSpecDecodeMetrics` in
+/// `vllm/v1/metrics/stats.py`.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RequestSpecDecodeMetrics {
+    /// Configured `num_speculative_tokens`.
+    #[serde(default)]
+    pub num_spec_tokens: u64,
+    /// Verify-step counts indexed by accepted draft-token count.
+    #[serde(default)]
+    pub histogram: Vec<u64>,
+    /// Total proposed draft tokens.
+    #[serde(default)]
+    pub num_draft_tokens: u64,
+    /// Accepted draft count per verify step; empty unless `detailed`.
+    #[serde(default)]
+    pub per_step_accepted: Vec<u64>,
+    /// Proposed draft count per verify step; empty unless `detailed`.
+    #[serde(default)]
+    pub per_step_drafted: Vec<u64>,
 }
 
 impl EngineCoreOutput {
@@ -148,6 +177,9 @@ impl EngineCoreOutput {
         self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
             .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
             .transpose()?;
+        self.new_sampling_mask = (self.new_sampling_mask.take())
+            .map(|value| value.resolve(frames, "new_sampling_mask"))
+            .transpose()?;
         Ok(())
     }
 }
@@ -156,11 +188,14 @@ impl EngineCoreOutput {
 ///
 /// Original Python definition:
 /// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
+#[serde_as]
 #[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
 struct WireEngineCoreOutputs {
     #[serde(default)]
     engine_index: u32,
     /// Outputs grouped for this client in the current engine tick.
+    // Match msgspec's append-only schema evolution, including Omni extensions.
+    #[serde_as(deserialize_as = "Vec<AllowTrailingFields>")]
     #[serde(default)]
     outputs: Vec<EngineCoreOutput>,
     #[serde(default)]
@@ -404,6 +439,61 @@ mod tests {
         );
     }
 
+    #[test]
+    fn engine_core_outputs_ignore_appended_fields() {
+        let output = EngineCoreOutput {
+            request_id: "req-1".into(),
+            new_token_ids: vec![42],
+            finish_reason: Some(EngineCoreFinishReason::Length),
+            ..Default::default()
+        };
+        let mut fields = crate::protocol::decode_value(&encode_msgpack(&output).unwrap())
+            .unwrap()
+            .as_array()
+            .unwrap()
+            .clone();
+        // Omni appends multimodal_output, is_segment_finished, and
+        // new_prompt_len_snapshot. Unknown payloads may contain nested values.
+        fields.extend([
+            OpaqueValue::Map(vec![(
+                "audio".into(),
+                OpaqueValue::Array(vec![OpaqueValue::Ext(1, vec![0, 1])]),
+            )]),
+            false.into(),
+            OpaqueValue::Nil,
+        ]);
+        let wire = (0, vec![OpaqueValue::Array(fields); 2]);
+        let frames = [Bytes::from(encode_msgpack(&wire).unwrap())];
+        let decoded = decode_engine_core_outputs(&frames).unwrap();
+        // A second output also checks that the first output's tail was consumed.
+        assert_eq!(decoded.as_request_batch().unwrap().outputs, vec![output; 2]);
+    }
+
+    #[test]
+    fn engine_core_output_requires_valid_known_fields() {
+        for fields in [
+            vec![],
+            vec!["req-1".into()],
+            vec!["req-1".into(), OpaqueValue::Nil],
+            vec!["req-1".into(), OpaqueValue::Array(vec![]), false.into()],
+        ] {
+            let wire = (0, vec![OpaqueValue::Array(fields)]);
+            let frames = [Bytes::from(encode_msgpack(&wire).unwrap())];
+            assert!(decode_engine_core_outputs(&frames).is_err());
+        }
+        let wire = (0, vec![("req-1", vec![42_u32])]);
+        let frames = [Bytes::from(encode_msgpack(&wire).unwrap())];
+        let decoded = decode_engine_core_outputs(&frames).unwrap();
+        assert_eq!(
+            decoded.as_request_batch().unwrap().outputs,
+            vec![EngineCoreOutput {
+                request_id: "req-1".into(),
+                new_token_ids: vec![42],
+                ..Default::default()
+            }]
+        );
+    }
+
     #[test]
     fn engine_core_outputs_classify_request_batch() {
         let outputs = WireEngineCoreOutputs {
diff --git a/rust/src/engine-core-client/src/protocol/request.rs b/rust/src/engine-core-client/src/protocol/request.rs
index 973ada71c6c1..8f38d7da0849 100644
--- a/rust/src/engine-core-client/src/protocol/request.rs
+++ b/rust/src/engine-core-client/src/protocol/request.rs
@@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize};
 use serde_default::DefaultFromSerde;
 use serde_tuple::{Deserialize_tuple, Serialize_tuple};
 
+use crate::protocol::kv_hints::KvHintsEnvelope;
 use crate::protocol::multimodal::MmFeatures;
 use crate::protocol::sampling::EngineCoreSamplingParams;
 use crate::protocol::{OpaqueValue, lora};
@@ -60,7 +61,7 @@ impl EngineCoreRequestType {
 ///
 /// Original Python construction point:
 /// <https://github.com/vllm-project/vllm/blob/cec2ec11760f9f3beabd4c90451936078bf91533/vllm/entrypoints/openai/chat_completion/serving.py#L367-L369>
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
 pub struct ReasoningParserKwargs {
     /// Effective kwargs visible to the chat template for this request.
     pub chat_template_kwargs: HashMap<String, serde_json::Value>,
@@ -127,6 +128,8 @@ pub struct EngineCoreRequest {
     /// Stable session identity shared by related requests.
     #[serde(default)]
     pub session_id: Option<String>,
+    #[serde(default)]
+    pub kv_hints: Option<KvHintsEnvelope>,
 }
 
 impl EngineCoreRequest {
@@ -159,6 +162,7 @@ mod tests {
 
     use super::*;
     use crate::protocol::dtype::TensorDtype;
+    use crate::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};
     use crate::protocol::multimodal::{
         MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, MmModality,
         PlaceholderRange,
@@ -181,6 +185,16 @@ mod tests {
             arrival_time: 1234.5,
             client_index: 7,
             session_id: Some("session-1".to_string()),
+            kv_hints: Some(KvHintsEnvelope {
+                protocol_version: "0.1".to_string(),
+                message_id: "msg-1".to_string(),
+                actions: vec![KvHintAction {
+                    action_id: "action-1".to_string(),
+                    action_type: "example.action".to_string(),
+                    action_version: "1.0".to_string(),
+                    payload: BTreeMap::from([("key".to_string(), serde_json::json!("value"))]),
+                }],
+            }),
             ..EngineCoreRequest::default()
         };
 
@@ -191,13 +205,14 @@ mod tests {
             other => panic!("expected array, got {other:?}"),
         };
 
-        assert_eq!(array.len(), 21);
+        assert_eq!(array.len(), 22);
         assert_eq!(array[0], Value::from("req-1"));
         assert_eq!(array[2], Value::Nil);
         assert_eq!(array[4], Value::Nil);
         assert_eq!(array[10], Value::Nil);
         assert_eq!(array[11], Value::from(7));
         assert_eq!(array[20], Value::from("session-1"));
+        assert!(matches!(&array[21], Value::Map(_)));
     }
 
     #[test]
diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs
index 19c579ac33a0..92f354be6fe0 100644
--- a/rust/src/engine-core-client/src/protocol/sampling.rs
+++ b/rust/src/engine-core-client/src/protocol/sampling.rs
@@ -20,6 +20,10 @@ fn default_temperature() -> f32 {
     1.0
 }
 
+fn default_watermarking() -> bool {
+    true
+}
+
 fn default_max_tokens() -> u32 {
     16
 }
@@ -68,6 +72,9 @@ pub struct EngineCoreSamplingParams {
     /// greedy sampling.
     #[serde(default = "default_temperature")]
     pub temperature: f32,
+    /// Whether to apply the engine's configured watermark to this request.
+    #[serde(default = "default_watermarking")]
+    pub watermarking: bool,
     /// Cumulative probability threshold for nucleus sampling.
     #[serde(default = "default_top_p")]
     pub top_p: f32,
@@ -155,6 +162,7 @@ impl EngineCoreSamplingParams {
     pub fn for_test() -> Self {
         Self {
             temperature: 1.0,
+            watermarking: true,
             top_p: 1.0,
             top_k: 0,
             seed: None,
@@ -229,6 +237,7 @@ mod tests {
 
         // Omitted fields -> Python defaults.
         assert_eq!(sampling.temperature, 1.0);
+        assert!(sampling.watermarking);
         assert_eq!(sampling.top_p, 1.0);
         assert_eq!(sampling.top_k, 0);
         assert_eq!(sampling.seed, None);
diff --git a/rust/src/engine-core-client/src/protocol/sampling_mask.rs b/rust/src/engine-core-client/src/protocol/sampling_mask.rs
new file mode 100644
index 000000000000..7f5c4592a7d3
--- /dev/null
+++ b/rust/src/engine-core-client/src/protocol/sampling_mask.rs
@@ -0,0 +1,237 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::ops::{Deref, DerefMut};
+
+use bytes::Bytes;
+use enum_as_inner::EnumAsInner;
+use serde::{Deserialize, Deserializer, Serialize};
+use serde_tuple::{Deserialize_tuple, Serialize_tuple};
+
+use crate::error::{Error, Result, bail_ext_value_decode};
+use crate::protocol::logprobs::array::decode_array1_u32;
+use crate::protocol::tensor::WireNdArray;
+
+/// Decoded sampling-mask support sets, one row per generated token position.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
+pub struct SamplingMask {
+    pub rows: Vec<Vec<u32>>,
+}
+
+/// Python `SamplingMaskLists` tuple before ndarray raw views are resolved.
+#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
+pub struct WireSamplingMask {
+    token_ids: WireNdArray,
+    #[serde(default)]
+    offsets: Option<WireNdArray>,
+    #[serde(default)]
+    cu_num_generated_tokens: Option<Vec<usize>>,
+}
+
+/// Sampling-mask field while it transitions from Python wire data to rows.
+#[derive(Debug, Clone, PartialEq, EnumAsInner)]
+pub enum MaybeWireSamplingMask {
+    Wire(Box<WireSamplingMask>),
+    Direct(SamplingMask),
+}
+
+impl Deref for MaybeWireSamplingMask {
+    type Target = SamplingMask;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Self::Wire(_) => panic!("sampling mask is still in wire format"),
+            Self::Direct(value) => value,
+        }
+    }
+}
+
+impl DerefMut for MaybeWireSamplingMask {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        match self {
+            Self::Wire(_) => panic!("sampling mask is still in wire format"),
+            Self::Direct(value) => value,
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for MaybeWireSamplingMask {
+    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        WireSamplingMask::deserialize(deserializer).map(|value| Self::Wire(Box::new(value)))
+    }
+}
+
+impl Serialize for MaybeWireSamplingMask {
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        match self {
+            Self::Wire(value) => value.serialize(serializer),
+            Self::Direct(value) => WireSamplingMask::from_direct(value)
+                .map_err(serde::ser::Error::custom)?
+                .serialize(serializer),
+        }
+    }
+}
+
+impl MaybeWireSamplingMask {
+    pub(super) fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result<Self> {
+        match self {
+            Self::Direct(value) => Ok(Self::Direct(value)),
+            Self::Wire(value) => value.resolve(frames, field_prefix).map(Self::Direct),
+        }
+    }
+}
+
+impl WireSamplingMask {
+    fn from_direct(value: &SamplingMask) -> std::result::Result<Self, String> {
+        let [row] = value.rows.as_slice() else {
+            return Err(format!(
+                "per-request engine-core sampling mask must contain exactly one row, got {}",
+                value.rows.len()
+            ));
+        };
+        let token_ids = row.iter().map(|&token_id| i64::from(token_id)).collect::<Vec<_>>();
+        Ok(Self {
+            token_ids: WireNdArray::from_i64(vec![token_ids.len()], token_ids)?,
+            offsets: None,
+            cu_num_generated_tokens: None,
+        })
+    }
+
+    fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result<SamplingMask> {
+        if let Some(indices) = self.cu_num_generated_tokens {
+            bail_ext_value_decode!(
+                "{field_prefix}.cu_num_generated_tokens: expected None for per-request engine-core sampling-mask payload, got {indices:?}"
+            );
+        }
+
+        let token_ids =
+            decode_array1_u32(self.token_ids, &format!("{field_prefix}.token_ids"), frames)?;
+        let Some(offsets) = self.offsets else {
+            return Ok(SamplingMask {
+                rows: vec![token_ids],
+            });
+        };
+        let offsets = decode_array1_u32(offsets, &format!("{field_prefix}.offsets"), frames)?;
+        if offsets.first().copied() != Some(0) {
+            bail_ext_value_decode!("{field_prefix}.offsets: first offset must be zero");
+        }
+        if offsets.windows(2).any(|pair| pair[0] > pair[1]) {
+            bail_ext_value_decode!("{field_prefix}.offsets: offsets must be nondecreasing");
+        }
+        if offsets.last().copied().map(|value| value as usize) != Some(token_ids.len()) {
+            bail_ext_value_decode!(
+                "{field_prefix}.offsets: final offset {:?} does not match token_ids length {}",
+                offsets.last(),
+                token_ids.len()
+            );
+        }
+
+        let rows = offsets
+            .windows(2)
+            .map(|pair| token_ids[pair[0] as usize..pair[1] as usize].to_vec())
+            .collect();
+        Ok(SamplingMask { rows })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::protocol::dtype::{NumpyDtype, TensorDtype};
+    use crate::protocol::output::{
+        EngineCoreOutput, RequestBatchOutputs, decode_engine_core_outputs,
+    };
+    use crate::protocol::tensor::WireArrayData;
+    use crate::protocol::{decode_msgpack, encode_msgpack};
+
+    #[test]
+    fn decodes_single_position_without_offsets() {
+        let wire = WireSamplingMask {
+            token_ids: WireNdArray::from_i64(vec![3], vec![2, 12, 16]).unwrap(),
+            offsets: None,
+            cu_num_generated_tokens: None,
+        };
+        let encoded = encode_msgpack(&wire).unwrap();
+        let decoded: MaybeWireSamplingMask = decode_msgpack(&encoded).unwrap();
+        let decoded = decoded.resolve(&[], "new_sampling_mask").unwrap();
+
+        assert_eq!(decoded.into_direct().unwrap().rows, vec![vec![2, 12, 16]]);
+    }
+
+    #[test]
+    fn decodes_csr_rows_and_rejects_bad_terminal_offset() {
+        let wire = WireSamplingMask {
+            token_ids: WireNdArray::from_i64(vec![4], vec![2, 12, 16, 18]).unwrap(),
+            offsets: Some(WireNdArray::from_i64(vec![3], vec![0, 2, 4]).unwrap()),
+            cu_num_generated_tokens: None,
+        };
+        let decoded = MaybeWireSamplingMask::Wire(Box::new(wire))
+            .resolve(&[], "new_sampling_mask")
+            .unwrap();
+        assert_eq!(
+            decoded.into_direct().unwrap().rows,
+            vec![vec![2, 12], vec![16, 18]]
+        );
+
+        let malformed = WireSamplingMask {
+            token_ids: WireNdArray::from_i64(vec![2], vec![2, 12]).unwrap(),
+            offsets: Some(WireNdArray::from_i64(vec![2], vec![0, 1]).unwrap()),
+            cu_num_generated_tokens: None,
+        };
+        let error = MaybeWireSamplingMask::Wire(Box::new(malformed))
+            .resolve(&[], "new_sampling_mask")
+            .unwrap_err();
+        assert!(error.to_string().contains("final offset"));
+    }
+
+    #[test]
+    fn decodes_sampling_mask_from_multipart_aux_frames() {
+        let output = EngineCoreOutput {
+            request_id: "req-multipart-mask".to_string(),
+            new_token_ids: vec![16, 18],
+            new_sampling_mask: Some(MaybeWireSamplingMask::Wire(Box::new(WireSamplingMask {
+                token_ids: WireNdArray {
+                    dtype: NumpyDtype::little(TensorDtype::I32),
+                    shape: vec![4],
+                    data: WireArrayData::AuxIndex(1),
+                },
+                offsets: Some(WireNdArray {
+                    dtype: NumpyDtype::little(TensorDtype::I64),
+                    shape: vec![3],
+                    data: WireArrayData::AuxIndex(2),
+                }),
+                cu_num_generated_tokens: None,
+            }))),
+            ..Default::default()
+        };
+        let primary = encode_msgpack(&crate::protocol::output::EngineCoreOutputs::from(
+            RequestBatchOutputs {
+                outputs: vec![output],
+                ..Default::default()
+            },
+        ))
+        .unwrap();
+        let token_ids =
+            [2_i32, 12, 16, 18].into_iter().flat_map(i32::to_le_bytes).collect::<Vec<_>>();
+        let offsets = [0_i64, 2, 4].into_iter().flat_map(i64::to_le_bytes).collect::<Vec<_>>();
+
+        let decoded = decode_engine_core_outputs(&[
+            Bytes::from(primary),
+            Bytes::from(token_ids),
+            Bytes::from(offsets),
+        ])
+        .unwrap();
+        let mask = decoded.as_request_batch().unwrap().outputs[0]
+            .new_sampling_mask
+            .as_ref()
+            .unwrap();
+
+        assert_eq!(mask.rows, vec![vec![2, 12], vec![16, 18]]);
+    }
+}
diff --git a/rust/src/engine-core-client/src/protocol/serde_utils.rs b/rust/src/engine-core-client/src/protocol/serde_utils.rs
new file mode 100644
index 000000000000..708353e42ea0
--- /dev/null
+++ b/rust/src/engine-core-client/src/protocol/serde_utils.rs
@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::fmt;
+use std::marker::PhantomData;
+
+use serde::de::{IgnoredAny, SeqAccess, Visitor, value::SeqAccessDeserializer};
+use serde::{Deserialize, Deserializer};
+use serde_with::DeserializeAs;
+
+/// Decode an array-like struct using its derived field types and defaults,
+/// skipping appended fields as Python's msgspec does. Opt in only for schemas
+/// whose extensions preserve the known field order and semantics.
+pub(super) struct AllowTrailingFields;
+
+impl<'de, T: Deserialize<'de>> DeserializeAs<'de, T> for AllowTrailingFields {
+    fn deserialize_as<D>(deserializer: D) -> Result<T, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        struct ArrayVisitor<T>(PhantomData<T>);
+
+        impl<'de, T: Deserialize<'de>> Visitor<'de> for ArrayVisitor<T> {
+            type Value = T;
+
+            fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+                formatter.write_str("an array-like struct")
+            }
+
+            fn visit_seq<A>(self, mut seq: A) -> Result<T, A::Error>
+            where
+                A: SeqAccess<'de>,
+            {
+                let value = T::deserialize(SeqAccessDeserializer::new(&mut seq))?;
+                while seq.next_element::<IgnoredAny>()?.is_some() {}
+                Ok(value)
+            }
+        }
+
+        deserializer.deserialize_seq(ArrayVisitor(PhantomData))
+    }
+}
diff --git a/rust/src/engine-core-client/src/protocol/structured_outputs.rs b/rust/src/engine-core-client/src/protocol/structured_outputs.rs
index 745ca3150b9b..99fa66ae2b02 100644
--- a/rust/src/engine-core-client/src/protocol/structured_outputs.rs
+++ b/rust/src/engine-core-client/src/protocol/structured_outputs.rs
@@ -104,6 +104,41 @@ impl StructuredOutputsParams {
             backend: StructuredOutputBackend::default(),
         }
     }
+
+    pub fn validate(&self) -> Result<()> {
+        match &self.constraint {
+            StructuredOutputConstraint::Grammar(grammar) if grammar.trim().is_empty() => {
+                Err(Error::InvalidStructuredOutputsParams {
+                    message: "structured_outputs.grammar cannot be an empty string".to_string(),
+                })
+            }
+            StructuredOutputConstraint::Json(Value::String(value)) => {
+                if value.trim().is_empty() {
+                    return Err(Error::InvalidStructuredOutputsParams {
+                        message: "structured_outputs.json cannot be an empty string".to_string(),
+                    });
+                }
+                serde_json::from_str::<Value>(value).map_err(|_| {
+                    Error::InvalidStructuredOutputsParams {
+                        message: "Invalid JSON grammar specification.".to_string(),
+                    }
+                })?;
+                Ok(())
+            }
+            StructuredOutputConstraint::Choice(choice) if choice.is_empty() => {
+                Err(Error::InvalidStructuredOutputsParams {
+                    message: "structured_outputs.choice cannot be empty".to_string(),
+                })
+            }
+            StructuredOutputConstraint::StructuralTag(tag) if tag.trim().is_empty() => {
+                Err(Error::InvalidStructuredOutputsParams {
+                    message: "structured_outputs.structural_tag cannot be an empty string"
+                        .to_string(),
+                })
+            }
+            _ => Ok(()),
+        }
+    }
 }
 
 /// Wire-compatible structured-output payload used by Python engine-core.
@@ -176,7 +211,7 @@ impl TryFrom<WireStructuredOutputsParams> for StructuredOutputsParams {
         }
         insert_constraint!("structural_tag", raw.structural_tag.map(StructuralTag));
 
-        Ok(Self {
+        let params = Self {
             constraint: constraint.map(|(_, c)| c).ok_or_else(|| {
                 Error::InvalidStructuredOutputsParams {
                     message: "missing structured output constraint".to_string(),
@@ -188,7 +223,9 @@ impl TryFrom<WireStructuredOutputsParams> for StructuredOutputsParams {
                 whitespace_pattern: raw.whitespace_pattern,
             },
             backend: raw.backend,
-        })
+        };
+        params.validate()?;
+        Ok(params)
     }
 }
 
@@ -289,6 +326,55 @@ mod tests {
         assert!(error.to_string().contains("json_object must be true"));
     }
 
+    #[test]
+    fn structured_outputs_rejects_empty_json_string() {
+        let error = serde_json::from_value::<StructuredOutputsParams>(serde_json::json!({
+            "json": "  ",
+        }))
+        .unwrap_err();
+
+        assert!(error.to_string().contains("json cannot be an empty string"));
+    }
+
+    #[test]
+    fn validate_rejects_empty_grammar_from_constructor() {
+        let error = StructuredOutputsParams::grammar("  ").validate().unwrap_err();
+        assert!(error.to_string().contains("grammar cannot be an empty string"));
+    }
+
+    #[test]
+    fn validate_rejects_empty_json_string_from_constructor() {
+        let error = StructuredOutputsParams::json(Value::String("  ".to_string()))
+            .validate()
+            .unwrap_err();
+        assert!(error.to_string().contains("json cannot be an empty string"));
+    }
+
+    #[test]
+    fn validate_rejects_malformed_json_string() {
+        let error = StructuredOutputsParams::json(Value::String("not json".to_string()))
+            .validate()
+            .unwrap_err();
+        assert!(error.to_string().contains("Invalid JSON grammar specification."));
+    }
+
+    #[test]
+    fn validate_rejects_empty_choice() {
+        let error = StructuredOutputsParams::choice(vec![]).validate().unwrap_err();
+        assert!(error.to_string().contains("choice cannot be empty"));
+    }
+
+    #[test]
+    fn validate_rejects_empty_structural_tag() {
+        let error = StructuredOutputsParams::structural_tag("\n\t").validate().unwrap_err();
+        assert!(error.to_string().contains("structural_tag cannot be an empty string"));
+    }
+
+    #[test]
+    fn validate_accepts_non_empty_grammar() {
+        assert!(StructuredOutputsParams::grammar("root ::= \"a\"").validate().is_ok());
+    }
+
     #[test]
     fn structured_outputs_serializes_through_raw_shape() {
         let params = StructuredOutputsParams {
diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs
index 13525375f81d..bd3847c483d8 100644
--- a/rust/src/engine-core-client/src/tests/client.rs
+++ b/rust/src/engine-core-client/src/tests/client.rs
@@ -20,6 +20,7 @@ use zeromq::util::PeerIdentity;
 use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage};
 
 use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
+use crate::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};
 use crate::protocol::logprobs::MaybeWireLogprobs;
 use crate::protocol::multimodal::{
     MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmModality, MmSlice,
@@ -151,6 +152,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
         prompt_token_ids: Some(vec![11, 22]),
         sampling_params: Some(EngineCoreSamplingParams {
             temperature: 0.8,
+            watermarking: false,
             top_p: 0.9,
             top_k: 8,
             max_tokens: 32,
@@ -164,6 +166,16 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
         }),
         arrival_time: 42.5,
         session_id: Some("session-1".to_string()),
+        kv_hints: Some(KvHintsEnvelope {
+            protocol_version: "0.1".to_string(),
+            message_id: "msg-1".to_string(),
+            actions: vec![KvHintAction {
+                action_id: "action-1".to_string(),
+                action_type: "example.action".to_string(),
+                action_version: "1.0".to_string(),
+                payload: BTreeMap::from([("key".to_string(), serde_json::json!("value"))]),
+            }],
+        }),
         ..EngineCoreRequest::default()
     }
 }
@@ -2640,6 +2652,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
     let multi_connector_stats_hex =
         lines.next().expect("missing MultiConnector stats fixture line");
     let ready_response_hex = lines.next().expect("missing ready response fixture line");
+    let extended_outputs_hex = lines.next().expect("missing extended outputs fixture line");
 
     let request_bytes = hex::decode(request_hex).unwrap();
     let multimodal_request_bytes = hex::decode(multimodal_request_hex).unwrap();
@@ -2662,6 +2675,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
         sampling,
         EngineCoreSamplingParams {
             temperature: 1.0,
+            watermarking: true,
             top_p: 1.0,
             top_k: 0,
             seed: None,
@@ -2709,17 +2723,25 @@ fn python_msgpack_fixtures_match_rust_encoding() {
         decode_value(&rmp_serde::to_vec_named(&expected_multimodal_request.mm_features).unwrap());
     assert_eq!(python_mm_features, rust_mm_features);
 
-    let decoded_sampling_mask_outputs: EngineCoreOutputs =
-        rmp_serde::from_slice(&sampling_mask_outputs_bytes).unwrap();
+    let decoded_sampling_mask_outputs =
+        decode_engine_core_outputs(&[bytes::Bytes::from(sampling_mask_outputs_bytes)]).unwrap();
     let sampling_mask_output =
         &decoded_sampling_mask_outputs.as_request_batch().unwrap().outputs[0];
     assert!(sampling_mask_output.mm_cache_miss_hashes.is_none());
-    assert!(matches!(
-        sampling_mask_output.new_sampling_mask.as_ref(),
-        Some(Value::Array(fields)) if fields.len() == 3
-    ));
+    assert_eq!(
+        sampling_mask_output.new_sampling_mask.as_ref().unwrap().rows,
+        vec![vec![2, 12, 16, 17, 18]]
+    );
 
     let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap();
+    // Match msgspec's base-schema result for the same extended Python message.
+    let extended_frames = [bytes::Bytes::from(
+        hex::decode(extended_outputs_hex).unwrap(),
+    )];
+    assert_eq!(
+        decode_engine_core_outputs(&extended_frames).unwrap(),
+        decoded_outputs
+    );
     expect_test::expect![[r#"
         RequestBatch(
             RequestBatchOutputs {
@@ -2844,6 +2866,11 @@ fn python_msgpack_fixtures_match_rust_encoding() {
 
     let ready_response: EngineCoreReadyResponse =
         rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap();
+    let mut legacy_ready = serde_json::to_value(&ready_response).unwrap();
+    legacy_ready.as_object_mut().unwrap().remove("effective_attention_block_size");
+    let legacy_ready: EngineCoreReadyResponse =
+        rmp_serde::from_slice(&rmp_serde::to_vec_named(&legacy_ready).unwrap()).unwrap();
+    assert!(legacy_ready.effective_attention_block_size.is_none());
     assert!(ready_response.supports_lora);
     assert_eq!(ready_response.max_loras, 8);
     assert_eq!(
@@ -2852,6 +2879,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
     );
     assert!(ready_response.enable_sleep_mode);
     assert!(ready_response.supports_draft_weight_updates);
+    assert_eq!(ready_response.effective_attention_block_size, Some(64));
     let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode");
     assert!(kv_events_config.enable_kv_cache_events);
     assert_eq!(kv_events_config.publisher, "zmq");
diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py
index fd6cb0c86be3..76978328bb1b 100755
--- a/rust/src/engine-core-client/src/tests/python_compat.py
+++ b/rust/src/engine-core-client/src/tests/python_compat.py
@@ -34,6 +34,7 @@ class FinishReason(IntEnum):
 # Mirror of real SamplingParams; omit_defaults makes fixtures match real maps.
 class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):
     temperature: float = 1.0
+    watermarking: bool = True
     top_p: float = 1.0
     top_k: int = 0
     seed: int | None = None
@@ -51,6 +52,19 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):
     output_kind: RequestOutputKind = RequestOutputKind.DELTA
 
 
+class KvHintAction(msgspec.Struct, frozen=True):
+    action_id: str
+    action_type: str
+    action_version: str
+    payload: dict[str, object]
+
+
+class KvHintsEnvelope(msgspec.Struct, frozen=True):
+    protocol_version: str
+    message_id: str
+    actions: list[KvHintAction]
+
+
 class EngineCoreRequest(
     msgspec.Struct,
     array_like=True,
@@ -77,6 +91,7 @@ class EngineCoreRequest(
     reasoning_parser_kwargs: dict[str, object] | None = None
     abort_immediately: bool = False
     session_id: str | None = None
+    kv_hints: KvHintsEnvelope | None = None
 
 
 class EngineCoreOutput(
@@ -100,6 +115,14 @@ class EngineCoreOutput(
     num_nans_in_logits: int = 0
     mm_cache_miss_hashes: list[str] | None = None
     new_sampling_mask: object | None = None
+    spec_decode_metrics: object | None = None
+
+
+class ExtendedEngineCoreOutput(EngineCoreOutput):
+    # Match Omni's append-only schema without requiring Omni or torch.
+    multimodal_output: dict[str, object] | None = None
+    is_segment_finished: bool = False
+    new_prompt_len_snapshot: int | None = None
 
 
 class EngineCoreOutputs(
@@ -123,6 +146,7 @@ class EngineCoreOutputs(
     mm_features=None,
     sampling_params=EngineCoreSamplingParams(
         temperature=0.8,
+        watermarking=False,
         top_p=0.9,
         top_k=8,
         seed=None,
@@ -143,6 +167,18 @@ class EngineCoreOutputs(
     arrival_time=42.5,
     client_index=0,
     session_id="session-1",
+    kv_hints=KvHintsEnvelope(
+        protocol_version="0.1",
+        message_id="msg-1",
+        actions=[
+            KvHintAction(
+                action_id="action-1",
+                action_type="example.action",
+                action_version="1.0",
+                payload={"key": "value"},
+            )
+        ],
+    ),
 )
 
 # All defaults -> empty map. Regression guard for the sparse-map decode.
@@ -205,18 +241,29 @@ class EngineCoreOutputs(
     finished_requests={"req-1"},
 )
 
+extended_outputs = EngineCoreOutputs(
+    outputs=[
+        ExtendedEngineCoreOutput(
+            **msgspec.structs.asdict(outputs.outputs[0]),
+            multimodal_output={"audio": b"\x00\x01"},
+            is_segment_finished=True,
+            new_prompt_len_snapshot=12,
+        )
+    ],
+    finished_requests=outputs.finished_requests,
+)
+extended_outputs_bytes = msgspec.msgpack.encode(extended_outputs)
+# The ordinary frontend's schema ignores even non-default extension values.
+assert msgspec.msgpack.decode(extended_outputs_bytes, type=EngineCoreOutputs) == outputs
+
 sampling_mask_wire = [
     [
         "<i4",
         [5],
-        msgpack.ExtType(3, np.array([2, 12, 16, 17, 18], dtype=np.int32).tobytes()),
-    ],
-    [
-        "<i8",
-        [2],
-        msgpack.ExtType(3, np.array([0, 5], dtype=np.int64).tobytes()),
+        msgspec.msgpack.Ext(3, np.array([2, 12, 16, 17, 18], dtype=np.int32).tobytes()),
     ],
     None,
+    None,
 ]
 outputs_with_sampling_mask = EngineCoreOutputs(
     outputs=[
@@ -421,6 +468,7 @@ class EngineCoreReadyResponse:
     weight_transfer_backend: str | None = None
     enable_sleep_mode: bool = False
     supports_draft_weight_updates: bool = False
+    effective_attention_block_size: int | None = None
 
 
 ready_response = EngineCoreReadyResponse(
@@ -444,6 +492,7 @@ class EngineCoreReadyResponse:
     weight_transfer_backend="nccl",
     enable_sleep_mode=True,
     supports_draft_weight_updates=True,
+    effective_attention_block_size=64,
     kv_events_config=KVEventsConfig(
         enable_kv_cache_events=True,
         publisher="zmq",
@@ -505,3 +554,4 @@ class EngineCoreReadyResponse:
 print(msgspec.msgpack.encode(mooncake_stats).hex())
 print(msgspec.msgpack.encode(multi_connector_stats).hex())
 print(msgspec.msgpack.encode(ready_response).hex())
+print(extended_outputs_bytes.hex())
diff --git a/rust/src/llm/examples/external_engine_smoke.rs b/rust/src/llm/examples/external_engine_smoke.rs
index 9d318285ae81..a28639e912d2 100644
--- a/rust/src/llm/examples/external_engine_smoke.rs
+++ b/rust/src/llm/examples/external_engine_smoke.rs
@@ -60,7 +60,9 @@ fn build_request(request_id: String, max_tokens: u32) -> GenerateRequest {
         priority: 0,
         data_parallel_rank: None,
         session_id: None,
+        kv_hints: None,
         reasoning_parser_kwargs: None,
+        reasoning_ended: None,
         lora_request: None,
     }
 }
diff --git a/rust/src/llm/src/error.rs b/rust/src/llm/src/error.rs
index 48bd8d18fd8b..ee1f71538962 100644
--- a/rust/src/llm/src/error.rs
+++ b/rust/src/llm/src/error.rs
@@ -12,4 +12,12 @@ pub enum Error {
     EmptyPromptTokenIds { request_id: String },
     #[error("engine-core error")]
     EngineCoreClient(#[from] vllm_engine_core_client::Error),
+    #[error(
+        "sampling mask for request `{request_id}` has {row_count} rows for {token_count} generated tokens"
+    )]
+    SamplingMaskTokenCountMismatch {
+        request_id: String,
+        token_count: usize,
+        row_count: usize,
+    },
 }
diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs
index c8f01b9fb9a3..34eae2568238 100644
--- a/rust/src/llm/src/output.rs
+++ b/rust/src/llm/src/output.rs
@@ -11,7 +11,10 @@ use futures::stream::FusedStream;
 use futures::{Stream, StreamExt as _, pin_mut};
 use serde::{Deserialize, Serialize};
 use vllm_engine_core_client::protocol::logprobs::Logprobs;
-use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason};
+use vllm_engine_core_client::protocol::output::{
+    EngineCoreFinishReason, RequestSpecDecodeMetrics, StopReason,
+};
+use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
 use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
 
 use crate::error::Result;
@@ -44,6 +47,10 @@ pub struct CollectedGenerateOutput {
     /// Connector-specific encoder cache transfer parameters for disaggregated
     /// serving.
     pub ec_transfer_params: Option<serde_json::Value>,
+    /// Sampling support sets aligned one-to-one with generated token positions.
+    pub sampling_mask: Option<SamplingMask>,
+    /// Per-request speculative-decoding metrics from the terminal output.
+    pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
 }
 
 /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for
@@ -155,6 +162,10 @@ pub struct GenerateOutput {
     /// Connector-specific encoder cache transfer parameters for disaggregated
     /// serving.
     pub ec_transfer_params: Option<serde_json::Value>,
+    /// Sampling support sets aligned one-to-one with `token_ids`.
+    pub sampling_mask: Option<SamplingMask>,
+    /// Per-request speculative-decoding metrics, present on terminal outputs.
+    pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>,
 }
 
 impl GenerateOutput {
@@ -202,6 +213,8 @@ impl GenerateOutput {
             cached_token_count: 0,
             kv_transfer_params: None,
             ec_transfer_params: None,
+            sampling_mask: None,
+            spec_decode_metrics: None,
         }
     }
 }
@@ -271,6 +284,16 @@ impl Stream for GenerateOutputStream {
         }
 
         let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap());
+        let sampling_mask = raw.new_sampling_mask.map(|value| value.into_direct().unwrap());
+        if let Some(mask) = sampling_mask.as_ref()
+            && mask.rows.len() != raw.new_token_ids.len()
+        {
+            return Poll::Ready(Some(Err(crate::Error::SamplingMaskTokenCountMismatch {
+                request_id: raw.request_id,
+                token_count: raw.new_token_ids.len(),
+                row_count: mask.rows.len(),
+            })));
+        }
         let cached_token_count = raw
             .prefill_stats
             .as_ref()
@@ -291,6 +314,8 @@ impl Stream for GenerateOutputStream {
             cached_token_count,
             kv_transfer_params: raw.kv_transfer_params,
             ec_transfer_params: raw.ec_transfer_params,
+            sampling_mask,
+            spec_decode_metrics: raw.spec_decode_metrics,
         };
 
         Poll::Ready(Some(Ok(output)))
@@ -341,6 +366,7 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
 
             while let Some(output) = stream.next().await.transpose()? {
                 cached_token_count = cached_token_count.max(output.cached_token_count);
+                let sampling_mask = output.sampling_mask;
                 if let Some(info) = output.prompt_info {
                     if prompt_token_ids.is_none() {
                         prompt_token_ids = Some(info.prompt_token_ids.to_vec());
@@ -359,6 +385,9 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                             existing.logprobs = Some(step_logprobs);
                         }
                     }
+                    if let Some(mut mask) = sampling_mask {
+                        existing.sampling_mask.get_or_insert_default().rows.append(&mut mask.rows);
+                    }
                 } else {
                     collected = Some(CollectedGenerateOutput {
                         request_id: output.request_id,
@@ -374,6 +403,8 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                         },
                         kv_transfer_params: None,
                         ec_transfer_params: None,
+                        sampling_mask,
+                        spec_decode_metrics: None,
                     });
                 }
 
@@ -387,6 +418,16 @@ impl<T: Stream<Item = Result<GenerateOutput>> + Send> T {
                     };
                     collected.kv_transfer_params = output.kv_transfer_params;
                     collected.ec_transfer_params = output.ec_transfer_params;
+                    collected.spec_decode_metrics = output.spec_decode_metrics;
+                    if let Some(mask) = collected.sampling_mask.as_ref()
+                        && mask.rows.len() != collected.token_ids.len()
+                    {
+                        return Err(crate::Error::SamplingMaskTokenCountMismatch {
+                            request_id: collected.request_id,
+                            token_count: collected.token_ids.len(),
+                            row_count: mask.rows.len(),
+                        });
+                    }
                     return Ok(collected);
                 }
             }
diff --git a/rust/src/llm/src/request.rs b/rust/src/llm/src/request.rs
index c647f54d8df4..da5c7ca56f97 100644
--- a/rust/src/llm/src/request.rs
+++ b/rust/src/llm/src/request.rs
@@ -4,6 +4,7 @@
 use std::collections::BTreeMap;
 
 use uuid::Uuid;
+use vllm_engine_core_client::protocol::kv_hints::KvHintsEnvelope;
 use vllm_engine_core_client::protocol::lora::LoraRequest;
 use vllm_engine_core_client::protocol::multimodal::MmFeatures;
 use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningParserKwargs};
@@ -48,9 +49,16 @@ pub struct GenerateRequest {
     pub data_parallel_rank: Option<u32>,
     /// Stable session identity shared by related requests.
     pub session_id: Option<String>,
+    /// Optional orchestrator-originated KV hints.
+    pub kv_hints: Option<KvHintsEnvelope>,
     /// Optional reasoning-parser kwargs forwarded to engine-side structured
     /// output logic.
     pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,
+    /// Optional engine reasoning-gate override.
+    ///
+    /// `Some(true)` means the submitted grammar covers reasoning from the first
+    /// generated token.
+    pub reasoning_ended: Option<bool>,
     /// Optional LoRA adapter request applied to this generation.
     pub lora_request: Option<LoraRequest>,
 }
@@ -79,7 +87,9 @@ impl GenerateRequest {
             priority,
             data_parallel_rank,
             session_id,
+            kv_hints,
             reasoning_parser_kwargs,
+            reasoning_ended,
             lora_request,
         } = self;
 
@@ -109,10 +119,9 @@ impl GenerateRequest {
                 trace_headers,
                 resumable: false,
                 session_id,
+                kv_hints,
                 external_req_id: Some(external_request_id),
-                // Rust parser doesn't expose this information, leave it unset and let the
-                // reasoning logic in engine-sided structured output manager handle it.
-                reasoning_ended: None,
+                reasoning_ended,
                 reasoning_parser_kwargs,
                 abort_immediately: false,
             },
@@ -134,6 +143,7 @@ impl PreparedGenerateRequest {
 mod tests {
     use std::collections::BTreeMap;
 
+    use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};
     use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
     use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
 
@@ -155,6 +165,16 @@ mod tests {
             priority: 3,
             data_parallel_rank: Some(2),
             session_id: Some("session-1".to_string()),
+            kv_hints: Some(KvHintsEnvelope {
+                protocol_version: "0.1".to_string(),
+                message_id: "msg-1".to_string(),
+                actions: vec![KvHintAction {
+                    action_id: "action-1".to_string(),
+                    action_type: "example.action".to_string(),
+                    action_version: "1.0".to_string(),
+                    payload: BTreeMap::new(),
+                }],
+            }),
             reasoning_parser_kwargs: Some(ReasoningParserKwargs {
                 chat_template_kwargs: [(
                     "chat_template_kwargs".to_string(),
@@ -164,6 +184,7 @@ mod tests {
                 )]
                 .into(),
             }),
+            reasoning_ended: None,
             lora_request: None,
         }
     }
@@ -183,6 +204,10 @@ mod tests {
         assert_eq!(request.cache_salt.as_deref(), Some("salt"));
         assert_eq!(request.data_parallel_rank, Some(2));
         assert_eq!(request.session_id.as_deref(), Some("session-1"));
+        assert_eq!(
+            request.kv_hints.as_ref().map(|hints| hints.message_id.as_str()),
+            Some("msg-1")
+        );
         assert_eq!(
             request.trace_headers,
             Some(BTreeMap::from([(
diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs
index 280fa0f10fbb..43802eaf4abe 100644
--- a/rust/src/llm/src/request_metrics.rs
+++ b/rust/src/llm/src/request_metrics.rs
@@ -45,6 +45,7 @@ pub(crate) struct RequestMetricsTracker {
     last_token_ts: f64,
     first_token_latency: f64,
     num_generation_tokens: u32,
+    num_preemptions: u64,
     latest_num_cached_tokens: u32,
 }
 
@@ -69,6 +70,7 @@ struct RequestMetricHandles {
     request_max_num_generation_tokens: HistogramMetric,
     request_params_max_tokens: HistogramMetric,
     request_params_n: HistogramMetric,
+    request_num_preemptions: HistogramMetric,
     request_prefill_kv_computed_tokens: HistogramMetric,
     time_to_first_token_seconds: HistogramMetric,
     inter_token_latency_seconds: HistogramMetric,
@@ -104,6 +106,7 @@ impl RequestMetricsTracker {
             last_token_ts: 0.0,
             first_token_latency: 0.0,
             num_generation_tokens: 0,
+            num_preemptions: 0,
             latest_num_cached_tokens: 0,
         }
     }
@@ -184,6 +187,7 @@ impl RequestMetricsTracker {
             self.handles.request_params_max_tokens.observe(max_tokens_param as f64);
         }
         self.handles.request_params_n.observe(self.n_param as f64);
+        self.handles.request_num_preemptions.observe(self.num_preemptions as f64);
         self.handles
             .request_prefill_kv_computed_tokens
             .observe(prefill_kv_computed_tokens as f64);
@@ -223,6 +227,7 @@ impl RequestMetricsTracker {
                     }
                 }
                 EngineCoreEventType::Preempted => {
+                    self.num_preemptions += 1;
                     self.handles.num_preemptions.inc();
                 }
             }
@@ -276,6 +281,7 @@ fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetri
             .get_or_create_owned(&labels),
         request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels),
         request_params_n: metrics.request_params_n.get_or_create_owned(&labels),
+        request_num_preemptions: metrics.request_num_preemptions.get_or_create_owned(&labels),
         request_prefill_kv_computed_tokens: metrics
             .request_prefill_kv_computed_tokens
             .get_or_create_owned(&labels),
diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs
index 2ff3faa8d459..370d5c046f33 100644
--- a/rust/src/llm/tests/generate.rs
+++ b/rust/src/llm/tests/generate.rs
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use std::collections::BTreeSet;
-use std::sync::Once;
+use std::sync::{Arc, Once};
 use std::time::Duration;
 
 use futures::StreamExt as _;
@@ -18,6 +18,7 @@ use vllm_engine_core_client::protocol::output::{
 };
 use vllm_engine_core_client::protocol::request::EngineCoreRequest;
 use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::sampling_mask::{MaybeWireSamplingMask, SamplingMask};
 use vllm_engine_core_client::protocol::stats::PrefillStats;
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
@@ -166,7 +167,9 @@ fn sample_generate_request(request_id: &str, max_tokens: u32) -> GenerateRequest
         priority: 0,
         data_parallel_rank: None,
         session_id: None,
+        kv_hints: None,
         reasoning_parser_kwargs: None,
+        reasoning_ended: None,
         lora_request: None,
     }
 }
@@ -326,6 +329,11 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
                                     num_local_cached_tokens: 1,
                                     ..Default::default()
                                 }),
+                                new_sampling_mask: Some(MaybeWireSamplingMask::Direct(
+                                    SamplingMask {
+                                        rows: vec![vec![1, 33, 99]],
+                                    },
+                                )),
                                 ..request_output_with_logprobs(
                                     &request.request_id,
                                     vec![33],
@@ -334,15 +342,22 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
                                     Some(prompt_logprobs()),
                                 )
                             },
-                            request_output_with_logprobs_and_kv(
-                                &request.request_id,
-                                vec![44],
-                                Some(EngineCoreFinishReason::Stop),
-                                Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)),
-                                None,
-                                Some(serde_json::json!({"connector": "x"})),
-                                None,
-                            ),
+                            EngineCoreOutput {
+                                new_sampling_mask: Some(MaybeWireSamplingMask::Direct(
+                                    SamplingMask {
+                                        rows: vec![vec![2, 44, 88]],
+                                    },
+                                )),
+                                ..request_output_with_logprobs_and_kv(
+                                    &request.request_id,
+                                    vec![44],
+                                    Some(EngineCoreFinishReason::Stop),
+                                    Some(logprobs_for_position(44, -0.3, 1, 88, -0.4)),
+                                    None,
+                                    Some(serde_json::json!({"connector": "x"})),
+                                    None,
+                                )
+                            },
                         ],
                         ..Default::default()
                     }
@@ -375,6 +390,43 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
         collected.kv_transfer_params,
         Some(serde_json::json!({"connector": "x"}))
     );
+    assert_eq!(
+        collected.sampling_mask,
+        Some(SamplingMask {
+            rows: vec![vec![1, 33, 99], vec![2, 44, 88]],
+        })
+    );
+}
+
+#[tokio::test]
+async fn collect_output_rejects_partial_sampling_mask() {
+    let output = vllm_llm::GenerateOutput {
+        request_id: "req-partial-mask".to_string(),
+        prompt_info: Some(GeneratePromptInfo {
+            prompt_token_ids: Arc::from([11_u32, 22]),
+            prompt_logprobs: None,
+        }),
+        token_ids: vec![33, 44],
+        logprobs: None,
+        finish_reason: Some(FinishReason::Length),
+        cached_token_count: 0,
+        kv_transfer_params: None,
+        ec_transfer_params: None,
+        sampling_mask: Some(SamplingMask {
+            rows: vec![vec![1, 33, 99]],
+        }),
+        spec_decode_metrics: None,
+    };
+
+    let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err();
+    assert!(matches!(
+        error,
+        Error::SamplingMaskTokenCountMismatch {
+            token_count: 2,
+            row_count: 1,
+            ..
+        }
+    ));
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -790,6 +842,12 @@ async fn generate_records_request_metrics_in_prometheus_output() {
     assert!(rendered.contains(&format!(
         "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1"
     )));
+    assert!(rendered.contains(&format!(
+        "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"4\"}} 1.0"
+    )));
+    assert!(rendered.contains(&format!(
+        "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
+    )));
     assert!(rendered.contains(&format!(
         "vllm:time_to_first_token_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
     )));
@@ -883,6 +941,12 @@ async fn dropping_stream_records_abort_terminal_request_metrics() {
     assert!(rendered.contains(&format!(
         "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"5\"}} 1"
     )));
+    assert!(rendered.contains(&format!(
+        "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"5\"}} 0.0"
+    )));
+    assert!(rendered.contains(&format!(
+        "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"5\"}} 1"
+    )));
 
     llm.shutdown().await.unwrap();
 }
diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs
index eadefced236e..0a81ae598e9c 100644
--- a/rust/src/managed-engine/src/cli.rs
+++ b/rust/src/managed-engine/src/cli.rs
@@ -84,6 +84,7 @@ impl ManagedEngineArgs {
     pub fn into_config(
         self,
         model: String,
+        revision: Option<String>,
         max_logprobs: Option<i32>,
         profiler_config: Option<String>,
         reasoning_parser: Option<&str>,
@@ -92,9 +93,14 @@ impl ManagedEngineArgs {
         shutdown_timeout: u64,
         handshake_port: u16,
         limit_mm_per_prompt: Option<String>,
+        hf_overrides: Option<String>,
     ) -> ManagedEngineConfig {
         let mut python_args = self.python_args;
         // Manually forward some args to the Python engine.
+        if let Some(revision) = revision {
+            python_args.push("--revision".to_string());
+            python_args.push(revision);
+        }
         if let Some(max_model_len) = self.max_model_len {
             python_args.push("--max-model-len".to_string());
             python_args.push(max_model_len);
@@ -131,6 +137,10 @@ impl ManagedEngineArgs {
             python_args.push("--limit-mm-per-prompt".to_string());
             python_args.push(limit_mm_per_prompt);
         }
+        if let Some(hf_overrides) = hf_overrides {
+            python_args.push("--hf-overrides".to_string());
+            python_args.push(hf_overrides);
+        }
 
         ManagedEngineConfig {
             python: self.python,
diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs
index c1135781b23c..5bfc32e89534 100644
--- a/rust/src/metrics/src/request.rs
+++ b/rust/src/metrics/src/request.rs
@@ -21,6 +21,10 @@ const REQUEST_LATENCY_BUCKETS: [f64; 21] = [
     480.0, 960.0, 1920.0, 7680.0,
 ];
 const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0];
+const REQUEST_NUM_PREEMPTIONS_BUCKETS: [f64; 7] = [1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 20.0];
+const ITERATION_TOKENS_BUCKETS: [f64; 13] = [
+    1.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, 16384.0,
+];
 
 fn build_1_2_5_buckets(max_value: u32) -> Vec<f64> {
     let mut buckets = Vec::new();
@@ -66,6 +70,14 @@ fn request_params_n_histogram() -> Histogram {
     Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied())
 }
 
+fn request_num_preemptions_histogram() -> Histogram {
+    Histogram::new(REQUEST_NUM_PREEMPTIONS_BUCKETS.iter().copied())
+}
+
+fn iteration_tokens_histogram() -> Histogram {
+    Histogram::new(ITERATION_TOKENS_BUCKETS.iter().copied())
+}
+
 #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
 pub struct FinishedReasonLabels {
     pub model_name: String,
@@ -83,8 +95,12 @@ pub struct PromptTokenSourceLabels {
 pub(crate) type FinishedReasonCounterFamily = Family<FinishedReasonLabels, U64Counter>;
 pub(crate) type PromptTokenSourceCounterFamily = Family<PromptTokenSourceLabels, U64Counter>;
 
-/// Request-lifecycle Prometheus families exported from the `llm` layer.
+/// Request and output-batch Prometheus families.
 pub struct RequestMetrics {
+    /// Computed prompt and generated tokens per nonempty engine output batch,
+    /// recorded by the engine-core client for requests tracked by this frontend.
+    pub iteration_tokens_total: HistogramFamily,
+
     // Request-derived counters.
     pub num_preemptions: Family<EngineLabels, U64Counter>,
     pub prompt_tokens: Family<EngineLabels, U64Counter>,
@@ -92,11 +108,6 @@ pub struct RequestMetrics {
     pub prompt_tokens_cached: Family<EngineLabels, U64Counter>,
     pub generation_tokens: Family<EngineLabels, U64Counter>,
 
-    // We intentionally don't support iteration-level histograms for now, since it seems to make
-    // more sense if the engine maintains these metrics and frontend simply forwards.
-    //
-    // pub iteration_tokens_total: HistogramFamily,
-
     // Request lifecycle counters and histograms.
     pub request_success: FinishedReasonCounterFamily,
     pub request_prompt_tokens: HistogramFamily,
@@ -104,6 +115,7 @@ pub struct RequestMetrics {
     pub request_max_num_generation_tokens: HistogramFamily,
     pub request_params_max_tokens: HistogramFamily,
     pub request_params_n: HistogramFamily,
+    pub request_num_preemptions: HistogramFamily,
     pub request_prefill_kv_computed_tokens: HistogramFamily,
     pub time_to_first_token_seconds: HistogramFamily,
     pub inter_token_latency_seconds: HistogramFamily,
@@ -118,6 +130,14 @@ pub struct RequestMetrics {
 impl RequestMetrics {
     /// Register the request-oriented metric families into the shared registry.
     pub(crate) fn register(registry: &mut Registry) -> Self {
+        let iteration_tokens_total =
+            Family::new_with_constructor(iteration_tokens_histogram as fn() -> Histogram);
+        registry.register(
+            "vllm:iteration_tokens_total",
+            "Histogram of computed prompt and generated tokens per nonempty engine output batch.",
+            iteration_tokens_total.clone(),
+        );
+
         // Request-derived counters.
         let num_preemptions = Family::default();
         registry.register(
@@ -202,6 +222,14 @@ impl RequestMetrics {
             request_params_n.clone(),
         );
 
+        let request_num_preemptions =
+            Family::new_with_constructor(request_num_preemptions_histogram as fn() -> Histogram);
+        registry.register(
+            "vllm:request_num_preemptions",
+            "Histogram of the number of times a request was preempted.",
+            request_num_preemptions.clone(),
+        );
+
         let request_prefill_kv_computed_tokens =
             Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
         registry.register(
@@ -276,6 +304,7 @@ impl RequestMetrics {
         );
 
         Self {
+            iteration_tokens_total,
             num_preemptions,
             prompt_tokens,
             prompt_tokens_by_source,
@@ -287,6 +316,7 @@ impl RequestMetrics {
             request_max_num_generation_tokens,
             request_params_max_tokens,
             request_params_n,
+            request_num_preemptions,
             request_prefill_kv_computed_tokens,
             time_to_first_token_seconds,
             inter_token_latency_seconds,
diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs
index c21682915910..4844aba12018 100644
--- a/rust/src/mock-engine/src/engine.rs
+++ b/rust/src/mock-engine/src/engine.rs
@@ -19,6 +19,7 @@ use vllm_engine_core_client::protocol::output::{
     UtilityCallOutput,
 };
 use vllm_engine_core_client::protocol::request::EngineCoreRequest;
+use vllm_engine_core_client::protocol::stats::PrefillStats;
 use vllm_engine_core_client::protocol::utility::{
     EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope,
 };
@@ -94,12 +95,10 @@ fn utility_response(
         "get_supported_tasks" => utility_envelope(vec!["generate"]),
         "is_sleeping" => utility_envelope(false),
         "reset_prefix_cache" => utility_envelope(true),
-        "reset_mm_cache"
-        | "reset_encoder_cache"
-        | "profile"
-        | "sleep"
-        | "wake_up"
-        | "execute_dummy_batch" => utility_envelope(()),
+        "wake_up" => utility_envelope(true),
+        "reset_mm_cache" | "reset_encoder_cache" | "profile" | "sleep" | "execute_dummy_batch" => {
+            utility_envelope(())
+        }
         _ => utility_envelope(Value::Nil),
     }?;
 
@@ -187,6 +186,11 @@ impl ActiveRequest {
 
     /// Advance this request by one mock engine step.
     fn step(&mut self, opt: &Opt) -> EngineCoreOutput {
+        let prefill_stats = (self.generated == 0).then(|| PrefillStats {
+            num_prompt_tokens: self.prompt_len as u32,
+            num_computed_tokens: self.prompt_len as u32,
+            ..Default::default()
+        });
         let remaining = self.max_tokens - self.generated;
         let chunk_len = remaining.min(opt.output_token_chunk_size);
         let mut new_token_ids = Vec::with_capacity(chunk_len);
@@ -196,11 +200,14 @@ impl ActiveRequest {
         self.generated += chunk_len;
 
         let finished = self.generated >= self.max_tokens;
-        request_output(
-            self.request_id.clone(),
-            new_token_ids,
-            finished.then_some(EngineCoreFinishReason::Length),
-        )
+        EngineCoreOutput {
+            prefill_stats,
+            ..request_output(
+                self.request_id.clone(),
+                new_token_ids,
+                finished.then_some(EngineCoreFinishReason::Length),
+            )
+        }
     }
 }
 
diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs
index 9bfee620ec74..c1cfbb49cb2d 100644
--- a/rust/src/mock-engine/src/tests.rs
+++ b/rust/src/mock-engine/src/tests.rs
@@ -10,6 +10,7 @@ use tokio_util::sync::CancellationToken;
 use vllm_engine_core_client::protocol::output::EngineCoreFinishReason;
 use vllm_engine_core_client::protocol::request::EngineCoreRequest;
 use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::stats::PrefillStats;
 use vllm_engine_core_client::test_utils::IpcNamespace;
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
 
@@ -119,12 +120,22 @@ async fn chunk_size_one_outputs_one_token_per_update() {
     let first = stream.next().await.expect("first").expect("first ok");
     assert_eq!(first.new_token_ids.len(), 1);
     assert_eq!(first.finish_reason, None);
+    assert_eq!(
+        first.prefill_stats,
+        Some(PrefillStats {
+            num_prompt_tokens: 3,
+            num_computed_tokens: 3,
+            ..Default::default()
+        })
+    );
     let second = stream.next().await.expect("second").expect("second ok");
     assert_eq!(second.new_token_ids.len(), 1);
     assert_eq!(second.finish_reason, None);
+    assert!(second.prefill_stats.is_none());
     let third = stream.next().await.expect("third").expect("third ok");
     assert_eq!(third.new_token_ids.len(), 1);
     assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length));
+    assert!(third.prefill_stats.is_none());
     assert!(stream.next().await.is_none());
 
     shutdown_mock(client, shutdown, task).await;
@@ -184,6 +195,7 @@ async fn utility_requests_return_minimal_success_responses() {
     assert!(client.reset_prefix_cache(false, false).await.expect("reset prefix cache"));
     client.reset_mm_cache().await.expect("reset mm cache");
     client.reset_encoder_cache().await.expect("reset encoder cache");
+    assert!(client.wake_up(Some(vec!["kv_cache".to_string()])).await.expect("wake up"));
 
     shutdown_mock(client, shutdown, task).await;
 }
diff --git a/rust/src/parser/Cargo.toml b/rust/src/parser/Cargo.toml
index 3c49362c5d27..92d71ae2272b 100644
--- a/rust/src/parser/Cargo.toml
+++ b/rust/src/parser/Cargo.toml
@@ -11,6 +11,7 @@ test-util = []
 easy-ext.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+serde_with.workspace = true
 thiserror.workspace = true
 thiserror-ext.workspace = true
 vllm-tokenizer.workspace = true
diff --git a/rust/src/parser/benches/utils/adapter.rs b/rust/src/parser/benches/utils/adapter.rs
index a49410f0bf43..d482a011c09f 100644
--- a/rust/src/parser/benches/utils/adapter.rs
+++ b/rust/src/parser/benches/utils/adapter.rs
@@ -3,9 +3,7 @@
 
 use std::sync::Arc;
 
-use vllm_parser::tool::{
-    Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput,
-};
+use vllm_parser::tool::{Result, Tool, ToolParser, ToolParserError, ToolParserOutput};
 use vllm_parser::unified::{
     UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
 };
@@ -89,10 +87,6 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
         self.inner.preserve_special_tokens()
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        self.inner.structural_tag_builder()
-    }
-
     fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
         self.inner.tool_call_id(tool_index)
     }
diff --git a/rust/src/parser/src/lib.rs b/rust/src/parser/src/lib.rs
index c9ddd53896b2..fd4aa0aa7f71 100644
--- a/rust/src/parser/src/lib.rs
+++ b/rust/src/parser/src/lib.rs
@@ -3,6 +3,7 @@
 
 //! Streaming parsers for chat completions.
 
+pub mod output_grammar;
 pub mod reasoning;
 pub mod tool;
 pub mod unified;
diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs
new file mode 100644
index 000000000000..f19254d36482
--- /dev/null
+++ b/rust/src/parser/src/output_grammar.rs
@@ -0,0 +1,393 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Output grammar construction shared by reasoning and tool parsers.
+
+use std::{fmt, str::FromStr};
+
+use serde_with::{DeserializeFromStr, SerializeDisplay};
+use thiserror::Error;
+use xgrammar_structural_tag::builders::{StructuralTagBuilder, StructuralTagOptions};
+use xgrammar_structural_tag::format::Format;
+use xgrammar_structural_tag::tool::ToolChoiceValue;
+use xgrammar_structural_tag::{
+    FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,
+};
+
+use crate::tool::Tool;
+
+/// Result alias for output grammar construction.
+pub type Result<T> = std::result::Result<T, OutputGrammarError>;
+
+/// How much of the generated stream an output grammar covers.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum GrammarCoverage {
+    /// The grammar covers reasoning and advances from the first generated token.
+    FromTokenZero,
+    /// The grammar covers only output emitted after reasoning ends.
+    FinalOutputOnly,
+}
+
+/// One parser-built output grammar and its generated-stream coverage.
+#[derive(Debug, Clone, PartialEq)]
+pub struct BuiltOutputGrammar {
+    /// Typed XGrammar structural-tag format.
+    pub format: Format,
+    /// Portion of the generated stream covered by `format`.
+    pub coverage: GrammarCoverage,
+}
+
+impl BuiltOutputGrammar {
+    /// A grammar that covers reasoning and everything after it.
+    pub fn from_token_zero(format: Format) -> Self {
+        Self {
+            format,
+            coverage: GrammarCoverage::FromTokenZero,
+        }
+    }
+
+    /// A grammar that covers only the output emitted after reasoning ends.
+    pub fn final_output_only(format: Format) -> Self {
+        Self {
+            format,
+            coverage: GrammarCoverage::FinalOutputOnly,
+        }
+    }
+}
+
+/// Server-side floor for tool-call structural tags (`--tool-strict-level`).
+#[derive(
+    Debug,
+    Clone,
+    Copy,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+    Default,
+    DeserializeFromStr,
+    SerializeDisplay,
+)]
+pub enum ToolStrictLevel {
+    /// Derive tool constraints from the request's tool choice and per-tool strictness.
+    #[default]
+    Auto,
+    /// Constrain the tool-call envelope for every request with tools.
+    Function,
+    /// Additionally pin argument schemas, as if every tool were `strict: true`.
+    Parameter,
+}
+
+impl ToolStrictLevel {
+    pub const AUTO_LITERAL: &str = "auto";
+    pub const FUNCTION_LITERAL: &str = "function";
+    pub const PARAMETER_LITERAL: &str = "parameter";
+}
+
+impl FromStr for ToolStrictLevel {
+    type Err = String;
+
+    fn from_str(value: &str) -> std::result::Result<Self, Self::Err> {
+        if value.eq_ignore_ascii_case(Self::AUTO_LITERAL) {
+            Ok(Self::Auto)
+        } else if value.eq_ignore_ascii_case(Self::FUNCTION_LITERAL) {
+            Ok(Self::Function)
+        } else if value.eq_ignore_ascii_case(Self::PARAMETER_LITERAL) {
+            Ok(Self::Parameter)
+        } else {
+            Err(format!(
+                "unknown tool strict level {value:?}; expected one of {}, {}, {}",
+                Self::AUTO_LITERAL,
+                Self::FUNCTION_LITERAL,
+                Self::PARAMETER_LITERAL
+            ))
+        }
+    }
+}
+
+impl fmt::Display for ToolStrictLevel {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            Self::Auto => Self::AUTO_LITERAL,
+            Self::Function => Self::FUNCTION_LITERAL,
+            Self::Parameter => Self::PARAMETER_LITERAL,
+        })
+    }
+}
+
+/// Request facts available while an initialized parser builds an output grammar.
+pub struct OutputGrammarContext<'a> {
+    /// Effective tools available for this generation.
+    pub tools: &'a [Tool],
+    /// Effective tool-choice policy.
+    pub tool_choice: &'a ToolChoice,
+    /// Server-side floor for tool-call constraints.
+    pub tool_strict_level: ToolStrictLevel,
+    /// Whether the request permits more than one tool call.
+    pub parallel_tool_calls: bool,
+}
+
+/// Errors produced while building an output grammar.
+#[derive(Debug, Error)]
+pub enum OutputGrammarError {
+    /// A model-specific structural-tag builder rejected its inputs.
+    #[error("failed to build output grammar")]
+    Build(#[from] xgrammar_structural_tag::Error),
+    /// A model-specific parser could not project a builder format to its
+    /// visible-output language.
+    #[error("unexpected output grammar shape from `{builder}` builder")]
+    UnexpectedBuilderFormat {
+        /// Builder whose output did not match the parser's known shape.
+        builder: &'static str,
+    },
+}
+
+/// Build the visible (post-reasoning) language from a crate structural-tag
+/// builder, or `None` when the request does not ask for strict tool calling.
+///
+/// This is the pre-existing strict-tool-calling grammar: the builder runs with
+/// `reasoning = false`, so the result describes only what follows the reasoning
+/// phase. It backs the default `ToolParser::build_visible_format` and the
+/// native unified parsers that have not yet grown their own builder.
+pub(crate) fn visible_format_from_builder(
+    builder: Option<&dyn StructuralTagBuilder>,
+    ctx: &OutputGrammarContext<'_>,
+) -> Result<Option<Format>> {
+    let Some(builder) = builder else {
+        return Ok(None);
+    };
+
+    format_from_builder(
+        builder,
+        ctx,
+        StructuralTagOptions::default().with_reasoning(false),
+    )
+}
+
+fn format_from_builder(
+    builder: &dyn StructuralTagBuilder,
+    ctx: &OutputGrammarContext<'_>,
+    options: StructuralTagOptions,
+) -> Result<Option<Format>> {
+    if !tool_grammar_applies(ctx) {
+        return Ok(None);
+    }
+
+    let structural_tag = build_structural_tag(
+        builder,
+        &tool_params(ctx.tools, ctx.tool_strict_level),
+        ctx.tool_choice.clone(),
+        options.with_parallel_tool_calls(ctx.parallel_tool_calls),
+    )?;
+
+    Ok(Some(structural_tag.format))
+}
+
+#[cfg(test)]
+pub(crate) fn full_format_from_builder_for_test(
+    builder: &dyn StructuralTagBuilder,
+    ctx: &OutputGrammarContext<'_>,
+    reasoning: xgrammar_structural_tag::builders::ReasoningMode,
+) -> Result<Option<Format>> {
+    format_from_builder(
+        builder,
+        ctx,
+        StructuralTagOptions::default().with_reasoning(reasoning),
+    )
+}
+
+fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec<ToolParam> {
+    tools
+        .iter()
+        .map(|tool| {
+            ToolParam::Function(FunctionToolParam::new(FunctionDefinition {
+                name: tool.name.clone(),
+                description: tool.description.clone(),
+                parameters: Some(tool.parameters.clone()),
+                strict: Some(
+                    strict_level >= ToolStrictLevel::Parameter || tool.strict == Some(true),
+                ),
+            }))
+        })
+        .collect()
+}
+
+/// Whether the request asks for a tool grammar at all.
+fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool {
+    if ctx.tools.is_empty() {
+        return false;
+    }
+
+    match ctx.tool_choice {
+        ToolChoice::Value(ToolChoiceValue::None) => false,
+        ToolChoice::Value(ToolChoiceValue::Auto) => {
+            ctx.tool_strict_level >= ToolStrictLevel::Function
+                || ctx.tools.iter().any(|tool| tool.strict == Some(true))
+        }
+        ToolChoice::Value(ToolChoiceValue::Required)
+        | ToolChoice::NamedFunction(_)
+        | ToolChoice::AllowedTools(_)
+        | ToolChoice::FlatAllowedTools(_)
+        | ToolChoice::Builtin(_) => true,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+    use xgrammar_structural_tag::builders::Qwen3Builder;
+
+    use super::*;
+
+    fn tool(name: &str, strict: Option<bool>) -> Tool {
+        Tool {
+            name: name.to_string(),
+            description: None,
+            parameters: json!({
+                "type": "object",
+                "properties": { "query": { "type": "string" } },
+                "required": ["query"]
+            }),
+            strict,
+        }
+    }
+
+    fn build(tools: &[Tool], tool_choice: &ToolChoice) -> Option<Format> {
+        visible_format_from_builder(
+            Some(&Qwen3Builder),
+            &OutputGrammarContext {
+                tools,
+                tool_choice,
+                tool_strict_level: ToolStrictLevel::Auto,
+                parallel_tool_calls: true,
+            },
+        )
+        .unwrap()
+    }
+
+    #[test]
+    fn tool_grammar_requires_tools_and_a_constraining_choice() {
+        let non_strict = [tool("search", None)];
+        assert!(build(&non_strict, &ToolChoice::none()).is_none());
+        assert!(build(&non_strict, &ToolChoice::auto()).is_none());
+        assert!(build(&non_strict, &ToolChoice::required()).is_some());
+        assert!(build(&non_strict, &ToolChoice::function("search")).is_some());
+
+        // `auto` only constrains once at least one tool is strict.
+        let strict = [tool("search", Some(true)), tool("lookup", None)];
+        assert!(build(&strict, &ToolChoice::auto()).is_some());
+    }
+
+    fn build_at_level(tool_choice: ToolChoice, tools: Vec<Tool>, level: ToolStrictLevel) -> Format {
+        visible_format_from_builder(
+            Some(&Qwen3Builder),
+            &OutputGrammarContext {
+                tools: &tools,
+                tool_choice: &tool_choice,
+                tool_strict_level: level,
+                parallel_tool_calls: true,
+            },
+        )
+        .unwrap()
+        .unwrap()
+    }
+
+    fn tag_string(format: &Format) -> String {
+        serde_json::to_string(format).unwrap()
+    }
+
+    #[test]
+    fn function_level_constrains_auto_without_strict_tools() {
+        let request = build_at_level(
+            ToolChoice::auto(),
+            vec![tool("search", None)],
+            ToolStrictLevel::Function,
+        );
+
+        assert!(tag_string(&request).contains("search"));
+    }
+
+    #[test]
+    fn absent_strict_is_non_strict() {
+        let unset = build_at_level(
+            ToolChoice::required(),
+            vec![tool("search", None)],
+            ToolStrictLevel::Auto,
+        );
+        let non_strict = build_at_level(
+            ToolChoice::required(),
+            vec![tool("search", Some(false))],
+            ToolStrictLevel::Auto,
+        );
+        let strict = build_at_level(
+            ToolChoice::required(),
+            vec![tool("search", Some(true))],
+            ToolStrictLevel::Auto,
+        );
+
+        assert_eq!(tag_string(&unset), tag_string(&non_strict));
+        assert_ne!(tag_string(&unset), tag_string(&strict));
+    }
+
+    #[test]
+    fn strict_tool_does_not_pin_its_neighbour() {
+        let mixed = build_at_level(
+            ToolChoice::auto(),
+            vec![tool("weather", Some(true)), tool("search", None)],
+            ToolStrictLevel::Auto,
+        );
+        let explicit = build_at_level(
+            ToolChoice::auto(),
+            vec![tool("weather", Some(true)), tool("search", Some(false))],
+            ToolStrictLevel::Auto,
+        );
+
+        assert_eq!(tag_string(&mixed), tag_string(&explicit));
+    }
+
+    #[test]
+    fn parameter_level_pins_argument_schemas() {
+        let pinned = build_at_level(
+            ToolChoice::auto(),
+            vec![tool("search", None)],
+            ToolStrictLevel::Parameter,
+        );
+        let strict = build_at_level(
+            ToolChoice::auto(),
+            vec![tool("search", Some(true))],
+            ToolStrictLevel::Auto,
+        );
+        let envelope_only = build_at_level(
+            ToolChoice::auto(),
+            vec![tool("search", None)],
+            ToolStrictLevel::Function,
+        );
+
+        assert_eq!(tag_string(&pinned), tag_string(&strict));
+        assert_ne!(tag_string(&pinned), tag_string(&envelope_only));
+    }
+
+    #[test]
+    fn tool_strict_level_parses_case_insensitively() {
+        assert_eq!("AUTO".parse::<ToolStrictLevel>(), Ok(ToolStrictLevel::Auto));
+        assert_eq!(
+            "PARAMETER".parse::<ToolStrictLevel>(),
+            Ok(ToolStrictLevel::Parameter)
+        );
+        assert_eq!(
+            "function".parse::<ToolStrictLevel>(),
+            Ok(ToolStrictLevel::Function)
+        );
+        assert_eq!(ToolStrictLevel::default(), ToolStrictLevel::Auto);
+        assert_eq!(
+            serde_json::to_value(ToolStrictLevel::Auto).unwrap(),
+            json!("auto")
+        );
+        assert_eq!(
+            serde_json::from_value::<ToolStrictLevel>(json!("auto")).unwrap(),
+            ToolStrictLevel::Auto
+        );
+        assert!("off".parse::<ToolStrictLevel>().is_err());
+        assert!("strict".parse::<ToolStrictLevel>().is_err());
+    }
+}
diff --git a/rust/src/parser/src/reasoning/deepseek_r1.rs b/rust/src/parser/src/reasoning/deepseek_r1.rs
index 16c2e8c8dc6e..5f8c637c1b03 100644
--- a/rust/src/parser/src/reasoning/deepseek_r1.rs
+++ b/rust/src/parser/src/reasoning/deepseek_r1.rs
@@ -2,6 +2,9 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use vllm_tokenizer::{DecodedText, DynTokenizer};
+use xgrammar_structural_tag::format::Format;
+
+use crate::output_grammar::{self, OutputGrammarContext};
 
 use super::{
     DelimitedReasoningParser, DelimitedReasoningParserBuilder, ReasoningDelta, ReasoningParser,
@@ -36,6 +39,14 @@ impl ReasoningParser for DeepSeekR1ReasoningParser {
         self.inner.initialize(prompt_token_ids)
     }
 
+    fn wrap_visible_format(
+        &self,
+        _ctx: &OutputGrammarContext<'_>,
+        visible: &Format,
+    ) -> output_grammar::Result<Option<Format>> {
+        Ok(Some(self.inner.wrap_visible_format(visible)))
+    }
+
     fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {
         Ok(self.inner.push(delta))
     }
diff --git a/rust/src/parser/src/reasoning/deepseek_v3.rs b/rust/src/parser/src/reasoning/deepseek_v3.rs
index 8b053c89f48d..944ff02681ae 100644
--- a/rust/src/parser/src/reasoning/deepseek_v3.rs
+++ b/rust/src/parser/src/reasoning/deepseek_v3.rs
@@ -2,6 +2,9 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use vllm_tokenizer::{DecodedText, DynTokenizer};
+use xgrammar_structural_tag::format::Format;
+
+use crate::output_grammar::{self, OutputGrammarContext};
 
 use super::{
     DelimitedReasoningParser, DelimitedReasoningParserBuilder, ReasoningDelta, ReasoningParser,
@@ -35,6 +38,14 @@ impl ReasoningParser for DeepSeekV3ReasoningParser {
         self.inner.initialize(prompt_token_ids)
     }
 
+    fn wrap_visible_format(
+        &self,
+        _ctx: &OutputGrammarContext<'_>,
+        visible: &Format,
+    ) -> output_grammar::Result<Option<Format>> {
+        Ok(Some(self.inner.wrap_visible_format(visible)))
+    }
+
     fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {
         Ok(self.inner.push(delta))
     }
diff --git a/rust/src/parser/src/reasoning/delimited.rs b/rust/src/parser/src/reasoning/delimited.rs
index 986ce2cc253e..ef7b8f95e08d 100644
--- a/rust/src/parser/src/reasoning/delimited.rs
+++ b/rust/src/parser/src/reasoning/delimited.rs
@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use vllm_tokenizer::{DecodedText, DynTokenizer, Tokenizer};
+use xgrammar_structural_tag::format::Format;
 
 use super::{ReasoningDelta, ReasoningError, Result};
 use crate::utils::partial_prefix_len;
@@ -148,6 +149,31 @@ impl DelimitedReasoningParser {
         self.current_in_reasoning
     }
 
+    /// Wrap visible output from the initialized reasoning and framing state.
+    pub(crate) fn wrap_visible_format(&self, visible: &Format) -> Format {
+        let begin = if self.current_in_reasoning {
+            "".to_string()
+        } else {
+            self.framed_start_token.clone()
+        };
+        let mut reasoning = Format::tag(begin, Format::any_text(), self.end_token.clone());
+        if !self.after_end.is_empty() {
+            reasoning = Format::sequence(vec![reasoning, Format::const_string(self.after_end)]);
+        }
+        if !self.current_in_reasoning {
+            reasoning = Format::optional(reasoning);
+            if !self.pending_after.is_empty() {
+                return Format::sequence(vec![
+                    Format::const_string(self.pending_after),
+                    reasoning,
+                    visible.clone(),
+                ]);
+            }
+        }
+
+        Format::sequence(vec![reasoning, visible.clone()])
+    }
+
     /// Parse one decoded text delta and return its reasoning/content split.
     pub(crate) fn push(&mut self, delta: DecodedText) -> ReasoningDelta {
         self.buffer.append(delta);
diff --git a/rust/src/parser/src/reasoning/glm45.rs b/rust/src/parser/src/reasoning/glm45.rs
index 557e628b268c..41bea717715f 100644
--- a/rust/src/parser/src/reasoning/glm45.rs
+++ b/rust/src/parser/src/reasoning/glm45.rs
@@ -2,6 +2,9 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use vllm_tokenizer::{DecodedText, DynTokenizer};
+use xgrammar_structural_tag::format::Format;
+
+use crate::output_grammar::{self, OutputGrammarContext};
 
 use super::{
     DelimitedReasoningParser, DelimitedReasoningParserBuilder, ReasoningDelta, ReasoningParser,
@@ -38,6 +41,14 @@ impl ReasoningParser for Glm45ReasoningParser {
         self.inner.initialize(prompt_token_ids)
     }
 
+    fn wrap_visible_format(
+        &self,
+        _ctx: &OutputGrammarContext<'_>,
+        visible: &Format,
+    ) -> output_grammar::Result<Option<Format>> {
+        Ok(Some(self.inner.wrap_visible_format(visible)))
+    }
+
     fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {
         Ok(self.inner.push(delta))
     }
diff --git a/rust/src/parser/src/reasoning/hy.rs b/rust/src/parser/src/reasoning/hy.rs
index 4a0ae1223267..ac3132ff7be0 100644
--- a/rust/src/parser/src/reasoning/hy.rs
+++ b/rust/src/parser/src/reasoning/hy.rs
@@ -2,6 +2,9 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use vllm_tokenizer::{DecodedText, DynTokenizer};
+use xgrammar_structural_tag::format::Format;
+
+use crate::output_grammar::{self, OutputGrammarContext};
 
 use super::{
     DelimitedReasoningParser, DelimitedReasoningParserBuilder, ReasoningDelta, ReasoningError,
@@ -43,6 +46,14 @@ impl ReasoningParser for HyReasoningParser {
         self.inner.initialize(prompt_token_ids)
     }
 
+    fn wrap_visible_format(
+        &self,
+        _ctx: &OutputGrammarContext<'_>,
+        visible: &Format,
+    ) -> output_grammar::Result<Option<Format>> {
+        Ok(Some(self.inner.wrap_visible_format(visible)))
+    }
+
     fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {
         Ok(self.inner.push(delta))
     }
diff --git a/rust/src/parser/src/reasoning/mod.rs b/rust/src/parser/src/reasoning/mod.rs
index 670695f7b0b1..ead7c0eeff66 100644
--- a/rust/src/parser/src/reasoning/mod.rs
+++ b/rust/src/parser/src/reasoning/mod.rs
@@ -31,6 +31,9 @@ mod step3p5;
 
 use thiserror::Error;
 use vllm_tokenizer::{DecodedText, DynTokenizer};
+use xgrammar_structural_tag::format::Format;
+
+use crate::output_grammar::{self, OutputGrammarContext};
 
 pub use self::cohere_cmd::CohereCmdReasoningParser;
 pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
@@ -61,6 +64,8 @@ pub type Glm47ReasoningParser = DeepSeekV3ReasoningParser;
 // TODO: kimi k2 may implicitly end reasoning by starting a tool call section
 // using <|tool_calls_section_begin|>, we should support that.
 pub type KimiK2ReasoningParser = DeepSeekV3ReasoningParser;
+/// MiMo V2 uses bare `<think>...</think>` delimiters.
+pub type MiMoReasoningParser = DeepSeekV3ReasoningParser;
 /// MiniMax M2 currently shares the Qwen3 reasoning parser.
 pub type MiniMaxM2ReasoningParser = Qwen3ReasoningParser;
 
@@ -135,6 +140,15 @@ pub trait ReasoningParser: Send {
         false
     }
 
+    /// Prepend the reasoning phase implied by the initialized parser state.
+    fn wrap_visible_format(
+        &self,
+        _ctx: &OutputGrammarContext<'_>,
+        _visible: &Format,
+    ) -> output_grammar::Result<Option<Format>> {
+        Ok(None)
+    }
+
     /// Feed one decoded text delta into the parser.
     fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta>;
 
diff --git a/rust/src/parser/src/reasoning/qwen3.rs b/rust/src/parser/src/reasoning/qwen3.rs
index 9a3f298c4b3d..19786422098a 100644
--- a/rust/src/parser/src/reasoning/qwen3.rs
+++ b/rust/src/parser/src/reasoning/qwen3.rs
@@ -2,6 +2,9 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use vllm_tokenizer::{DecodedText, DynTokenizer};
+use xgrammar_structural_tag::format::Format;
+
+use crate::output_grammar::{self, OutputGrammarContext};
 
 use super::{
     DelimitedReasoningParser, DelimitedReasoningParserBuilder, ReasoningDelta, ReasoningParser,
@@ -40,6 +43,14 @@ impl ReasoningParser for Qwen3ReasoningParser {
         self.inner.initialize(prompt_token_ids)
     }
 
+    fn wrap_visible_format(
+        &self,
+        _ctx: &OutputGrammarContext<'_>,
+        visible: &Format,
+    ) -> output_grammar::Result<Option<Format>> {
+        Ok(Some(self.inner.wrap_visible_format(visible)))
+    }
+
     fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {
         Ok(self.inner.push(delta))
     }
diff --git a/rust/src/parser/src/reasoning/step3p5.rs b/rust/src/parser/src/reasoning/step3p5.rs
index 7f94faf9985d..8e504d7a65ac 100644
--- a/rust/src/parser/src/reasoning/step3p5.rs
+++ b/rust/src/parser/src/reasoning/step3p5.rs
@@ -2,6 +2,9 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use vllm_tokenizer::{DecodedText, DynTokenizer};
+use xgrammar_structural_tag::format::Format;
+
+use crate::output_grammar::{self, OutputGrammarContext};
 
 use super::{
     DelimitedReasoningParser, DelimitedReasoningParserBuilder, ReasoningDelta, ReasoningParser,
@@ -38,6 +41,14 @@ impl ReasoningParser for Step3p5ReasoningParser {
         self.inner.initialize(prompt_token_ids)
     }
 
+    fn wrap_visible_format(
+        &self,
+        _ctx: &OutputGrammarContext<'_>,
+        visible: &Format,
+    ) -> output_grammar::Result<Option<Format>> {
+        Ok(Some(self.inner.wrap_visible_format(visible)))
+    }
+
     fn push(&mut self, delta: DecodedText) -> Result<ReasoningDelta> {
         Ok(self.inner.push(delta))
     }
diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs
index 632eda600f77..9bbe39367b5e 100644
--- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs
+++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs
@@ -193,6 +193,34 @@ mod tests {
         );
     }
 
+    #[test]
+    fn deepseek_v32_parse_complete_keeps_parameter_without_string_attr() {
+        // The model sometimes omits `string="..."`; the parameter must still
+        // be kept and converted through the schema, like `string="false"`.
+        let mut parser = DeepSeekV32ToolParser::new(&test_tools());
+        let output = parser
+            .parse_complete(
+                "<｜DSML｜function_calls>\n\
+                 <｜DSML｜invoke name=\"convert\">\n\
+                 <｜DSML｜parameter name=\"whole\">5.0</｜DSML｜parameter>\n\
+                 <｜DSML｜parameter name=\"flag\" >true</｜DSML｜parameter>\n\
+                 <｜DSML｜parameter name=\"payload\" string=\"true\">{\"nested\":true}</｜DSML｜parameter>\n\
+                 </｜DSML｜invoke>\n\
+                 </｜DSML｜function_calls>",
+            )
+            .unwrap();
+
+        assert_eq!(output.calls().len(), 1);
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
+            json!({
+                "whole": 5.0,
+                "flag": true,
+                "payload": "{\"nested\":true}",
+            })
+        );
+    }
+
     #[test]
     fn deepseek_v32_parse_complete_preserves_raw_closing_tag_text_in_parameter_value() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs
index b65df235ca20..032cfbe5e284 100644
--- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs
+++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs
@@ -147,6 +147,29 @@ mod tests {
         );
     }
 
+    #[test]
+    fn deepseek_v4_streaming_keeps_parameter_without_string_attr() {
+        let mut parser = DeepSeekV4ToolParser::new(&test_tools());
+        let output = collect_stream(
+            &mut parser,
+            &[
+                "<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n",
+                "<｜DSML｜parameter name=\"loc",
+                "ation\">Par",
+                "is</｜DSML｜parameter>\n",
+                "<｜DSML｜parameter name=\"date\" string=\"true\">tomorrow</｜DSML｜parameter>\n",
+                "</｜DSML｜invoke>\n</｜DSML｜tool_calls>",
+            ],
+        );
+
+        assert!(output.normal_text().is_empty());
+        assert_eq!(output.calls().len(), 1);
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls()[0].arguments).unwrap(),
+            json!({ "location": "Paris", "date": "tomorrow" })
+        );
+    }
+
     #[test]
     fn tool_framing_preserves_body_whitespace_across_chunk_boundaries() {
         assert_tool_framing_preserves_body_whitespace::<DeepSeekV4ToolParser>(
diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41.rs
index 44658b1d44f8..be68b6d2e5f3 100644
--- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41.rs
+++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41.rs
@@ -4,8 +4,6 @@
 use super::{DeepSeekDsmlToolParser, DsmlTokens};
 use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
 
-mod structural_tag;
-
 /// Tool parser for DeepSeek V4.1's spaced DSML tags.
 ///
 /// Arguments are emitted only after a full `invoke` block is parsed.
@@ -27,7 +25,7 @@ impl ToolParser for DeepSeekV41ToolParser {
     }
 
     fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(&structural_tag::DeepSeekV41StructuralTagBuilder)
+        Some(xgrammar_structural_tag::Model::DeepSeekV41.builder())
     }
 
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/structural_tag.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/structural_tag.rs
deleted file mode 100644
index d5fb6118fb62..000000000000
--- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/structural_tag.rs
+++ /dev/null
@@ -1,108 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-//! Visible-text grammar for DeepSeek V4.1 tool calls.
-//!
-//! Adapted from xgrammar-structural-tag 0.2.0's DeepSeek V4 builder, with
-//! V4.1's spaced DSML tags and schema-unconstrained parameter values.
-
-use serde_json::Value;
-use xgrammar_structural_tag::Result;
-use xgrammar_structural_tag::builders::{
-    StructuralTagBuilder, StructuralTagContext, StructuralTagOptions,
-};
-use xgrammar_structural_tag::format::{
-    Format, JsonSchemaFormat, StructuralTag, TagFormat, TriggeredTagsFormat,
-};
-use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam};
-
-use super::DsmlTokens;
-
-const TOKENS: DsmlTokens = DsmlTokens::V41;
-
-pub(super) struct DeepSeekV41StructuralTagBuilder;
-
-impl StructuralTagBuilder for DeepSeekV41StructuralTagBuilder {
-    fn build(&self, ctx: StructuralTagContext<'_>) -> Result<StructuralTag> {
-        let tags = ctx
-            .function_tools
-            .iter()
-            .map(|tool| call_tag(tool, ctx.options))
-            .collect::<Vec<_>>();
-        let excludes = if ctx.options.exclude_special_tokens {
-            &["<think>", "</think>"][..]
-        } else {
-            &[]
-        };
-        let format = match ctx.tool_choice {
-            BuilderToolChoice::Auto if tags.is_empty() => Format::any_text_excluding(excludes),
-            BuilderToolChoice::Auto => Format::TriggeredTags(
-                TriggeredTagsFormat::new(
-                    &[TOKENS.tool_calls_start],
-                    vec![TagFormat::new(
-                        format!("{}\n", TOKENS.tool_calls_start),
-                        Format::tags_with_separator(tags, "", true, false),
-                        TOKENS.tool_calls_end,
-                    )],
-                )
-                .with_excludes(excludes),
-            ),
-            BuilderToolChoice::Forced | BuilderToolChoice::Required => Format::sequence(vec![
-                Format::const_string(format!("\n\n{}\n", TOKENS.tool_calls_start)),
-                Format::tags_with_separator(
-                    tags,
-                    "",
-                    true,
-                    ctx.tool_choice == BuilderToolChoice::Forced,
-                ),
-                Format::const_string(TOKENS.tool_calls_end),
-            ]),
-        };
-        // Reasoning boundaries remain owned by the engine reasoner; serving
-        // calls this visible-text builder with reasoning=false.
-        Ok(StructuralTag::new(format))
-    }
-}
-
-fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat {
-    // TODO: enforce tool.parameters and strict schema semantics once V4.1
-    // argument conversion is available. This builder intentionally ignores
-    // parameters and strict: it constrains tool names and DSML/value syntax
-    // only. Parameter names, presence, uniqueness and value schemas remain
-    // unconstrained, including for strict=true. Serving still uses strict to
-    // decide whether auto tool choice activates a grammar.
-    TagFormat::new(
-        format!("{} name=\"{}\">\n", TOKENS.invoke_start, tool.function.name),
-        Format::star(parameter(options)),
-        format!("{}\n", TOKENS.invoke_end),
-    )
-}
-
-fn parameter(options: StructuralTagOptions) -> Format {
-    Format::tag(
-        format!("{} name=\"", TOKENS.parameter_start),
-        Format::sequence(vec![
-            Format::regex(r#"[^"]+"#),
-            Format::const_string("\" string=\""),
-            Format::or(vec![
-                Format::sequence(vec![
-                    Format::const_string("true\">"),
-                    Format::any_text_excluding(&[
-                        TOKENS.parameter_end,
-                        TOKENS.invoke_end,
-                        TOKENS.tool_calls_end,
-                    ]),
-                ]),
-                Format::sequence(vec![
-                    Format::const_string("false\">"),
-                    Format::JsonSchema(
-                        JsonSchemaFormat::new(Value::Bool(true))
-                            .with_any_order(options.any_order)
-                            .with_max_whitespace_cnt(options.max_whitespace_cnt),
-                    ),
-                ]),
-            ]),
-        ]),
-        format!("{}\n", TOKENS.parameter_end),
-    )
-}
diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs
index d40f82d533cd..208d88ea33a8 100644
--- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs
+++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs
@@ -71,6 +71,28 @@ fn spaced_dsml_streaming_preserves_shared_schema_coercion() {
     }
 }
 
+#[test]
+fn spaced_dsml_keeps_parameter_without_string_attr() {
+    let wire = concat!(
+        "<｜DSML｜ calls>\n",
+        "<｜DSML｜ invoke name=\"lookup\">\n",
+        "<｜DSML｜ parameter name=\"query\">value</｜DSML｜ parameter>\n",
+        "<｜DSML｜ parameter name=\"limit\">2</｜DSML｜ parameter>\n",
+        "</｜DSML｜ invoke>\n",
+        "</｜DSML｜ calls>",
+    );
+    let mut parser = DeepSeekV41ToolParser::create(&tools()).unwrap();
+
+    let output = collect_stream(parser.as_mut(), &[wire]);
+
+    assert!(output.normal_text().is_empty());
+    assert_eq!(output.calls().len(), 1);
+    assert_eq!(
+        serde_json::from_str::<serde_json::Value>(&output.calls()[0].arguments).unwrap(),
+        json!({ "query": "value", "limit": 2 })
+    );
+}
+
 #[test]
 fn v4_tags_remain_text_in_v41_dialect() {
     let mut parser = DeepSeekV41ToolParser::create(&tools()).unwrap();
@@ -80,7 +102,7 @@ fn v4_tags_remain_text_in_v41_dialect() {
 }
 
 #[test]
-fn structural_tag_ignores_parameter_schemas_and_strict() {
+fn structural_tag_respects_parameter_schemas_and_strict() {
     let parser = DeepSeekV41ToolParser::create(&tools()).unwrap();
     let builder = parser.structural_tag_builder().unwrap();
     let options = StructuralTagOptions::default().with_reasoning(false);
@@ -106,7 +128,13 @@ fn structural_tag_ignores_parameter_schemas_and_strict() {
             ] {
                 let mut function = FunctionDefinition::new("lookup").with_parameters(parameters);
                 function.strict = strict;
-                assert_eq!(build(function), expected);
+                let tag = build(function);
+                if strict == Some(false) {
+                    assert_eq!(tag, expected);
+                } else {
+                    assert_ne!(tag, expected);
+                }
+                assert!(tag.to_json_string().unwrap().contains("deepseek_v4_1_xml"));
             }
         }
     }
diff --git a/rust/src/parser/src/tool/deepseek_dsml/mod.rs b/rust/src/parser/src/tool/deepseek_dsml/mod.rs
index 01a3fbe282ea..9c7dd203312d 100644
--- a/rust/src/parser/src/tool/deepseek_dsml/mod.rs
+++ b/rust/src/parser/src/tool/deepseek_dsml/mod.rs
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};
-use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated};
+use winnow::combinator::{alt, delimited, eof, opt, preceded, repeat, seq, terminated};
 use winnow::prelude::*;
 use winnow::stream::Partial;
 use winnow::token::{literal, rest, take_until};
@@ -293,13 +293,16 @@ fn parse_invoke_params(invoke_body: &str, tokens: DsmlTokens) -> ModalResult<Vec
 }
 
 /// Parse a DSML parameter block.
+///
+/// The `string` attribute is optional: the model sometimes omits it, and
+/// failing the whole invoke would leak the block into the text output. A
+/// missing or non-`"true"` value is treated like `string="false"`.
 fn parse_parameter(input: &mut &str, tokens: DsmlTokens) -> ModalResult<DsmlParameter> {
     seq! {DsmlParameter {
         _: literal(tokens.parameter_start),
         _: ws1,
         name: name_attr.map(|name: &str| name.to_string()),
-        _: ws1,
-        is_string: string_attr.map(|value| value == "true"),
+        is_string: opt(preceded(ws1, string_attr)).map(|value| value == Some("true")),
         _: ws0,
         _: ">",
         value: take_until(0.., tokens.parameter_end).map(str::to_string),
@@ -315,7 +318,7 @@ fn name_attr<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
 
 /// Parse a string attribute.
 fn string_attr<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
-    delimited("string=\"", alt(("true", "false")), "\"").parse_next(input)
+    delimited("string=\"", take_until(0.., "\""), "\"").parse_next(input)
 }
 
 /// Parse a DSML name attribute.
diff --git a/rust/src/parser/src/tool/hy/structural_tag.rs b/rust/src/parser/src/tool/hy/structural_tag.rs
index cce98d5b0940..e883acbafd50 100644
--- a/rust/src/parser/src/tool/hy/structural_tag.rs
+++ b/rust/src/parser/src/tool/hy/structural_tag.rs
@@ -219,7 +219,7 @@ mod tests {
             ToolChoice::required(),
         );
 
-        expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls:opensource>\n"},{"type":"tags_with_separator","tags":[{"begin":"<tool_call:opensource>get_weather<tool_sep:opensource>\n","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"city"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<tool_sep:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"]},{"type":"const_string","value":"</arg_value:opensource>"},{"type":"const_string","value":"\n"}]},{"type":"star","content":{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"days"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<tool_sep:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"]},{"type":"const_string","value":"</arg_value:opensource>"},{"type":"const_string","value":"\n"}]}}]},"end":"</tool_call:opensource>"}],"separator":"\n","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"\n</tool_calls:opensource>"}]}}"#]].assert_eq(&tag.to_json_string().unwrap());
+        expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls:opensource>\n"},{"type":"tags_with_separator","tags":[{"type":"tag","begin":"<tool_call:opensource>get_weather<tool_sep:opensource>\n","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"city"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<tool_sep:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"],"max_tokens":null,"max_chars":null},{"type":"const_string","value":"</arg_value:opensource>"},{"type":"const_string","value":"\n"}]},{"type":"star","content":{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"days"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<tool_sep:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"],"max_tokens":null,"max_chars":null},{"type":"const_string","value":"</arg_value:opensource>"},{"type":"const_string","value":"\n"}]}}]},"end":"</tool_call:opensource>"}],"separator":"\n","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"\n</tool_calls:opensource>"}]}}"#]].assert_eq(&tag.to_json_string().unwrap());
     }
 
     #[test]
@@ -238,7 +238,7 @@ mod tests {
             ToolChoice::required(),
         );
 
-        expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls:opensource>"},{"type":"tags_with_separator","tags":[{"begin":"<tool_call:opensource>get_weather","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"city"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"]},{"type":"const_string","value":"</arg_value:opensource>"}]}]},"end":"</tool_call:opensource>"}],"separator":"","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"</tool_calls:opensource>"}]}}"#]].assert_eq(&tag.to_json_string().unwrap());
+        expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls:opensource>"},{"type":"tags_with_separator","tags":[{"type":"tag","begin":"<tool_call:opensource>get_weather","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"city"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"],"max_tokens":null,"max_chars":null},{"type":"const_string","value":"</arg_value:opensource>"}]}]},"end":"</tool_call:opensource>"}],"separator":"","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"</tool_calls:opensource>"}]}}"#]].assert_eq(&tag.to_json_string().unwrap());
     }
 
     #[test]
@@ -259,7 +259,7 @@ mod tests {
         );
         let value = serde_json::to_value(tag).unwrap();
 
-        expect![[r#"{"type":"sequence","elements":[{"type":"star","content":{"type":"or","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key>"},{"type":"const_string","value":"query"},{"type":"const_string","value":"</arg_key>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value>"},{"type":"any_text","excludes":["<tool_calls>","</tool_calls>","<tool_call>","</tool_call>","<tool_sep>","<arg_key>","</arg_key>","<arg_value>","</arg_value>"]},{"type":"const_string","value":"</arg_value>"},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key>"},{"type":"const_string","value":"limit"},{"type":"const_string","value":"</arg_key>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value>"},{"type":"any_text","excludes":["<tool_calls>","</tool_calls>","<tool_call>","</tool_call>","<tool_sep>","<arg_key>","</arg_key>","<arg_value>","</arg_value>"]},{"type":"const_string","value":"</arg_value>"},{"type":"const_string","value":"\n"}]}]}}]}"#]].assert_eq(
+        expect![[r#"{"type":"sequence","elements":[{"type":"star","content":{"type":"or","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key>"},{"type":"const_string","value":"query"},{"type":"const_string","value":"</arg_key>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value>"},{"type":"any_text","excludes":["<tool_calls>","</tool_calls>","<tool_call>","</tool_call>","<tool_sep>","<arg_key>","</arg_key>","<arg_value>","</arg_value>"],"max_tokens":null,"max_chars":null},{"type":"const_string","value":"</arg_value>"},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key>"},{"type":"const_string","value":"limit"},{"type":"const_string","value":"</arg_key>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value>"},{"type":"any_text","excludes":["<tool_calls>","</tool_calls>","<tool_call>","</tool_call>","<tool_sep>","<arg_key>","</arg_key>","<arg_value>","</arg_value>"],"max_tokens":null,"max_chars":null},{"type":"const_string","value":"</arg_value>"},{"type":"const_string","value":"\n"}]}]}}]}"#]].assert_eq(
             &serde_json::to_string(
                 &value["format"]["elements"][1]["tags"][0]["content"],
             )
@@ -276,8 +276,8 @@ mod tests {
         let auto = build("", &tools, ToolChoice::auto()).to_json_string().unwrap();
         let forced = build("", &tools, ToolChoice::function("lookup")).to_json_string().unwrap();
 
-        expect![[r#"{"type":"structural_tag","format":{"type":"triggered_tags","triggers":["<tool_calls>"],"tags":[{"begin":"<tool_calls>\n","content":{"type":"tags_with_separator","tags":[{"begin":"<tool_call>search<tool_sep>\n","content":{"type":"any_text","excludes":[]},"end":"</tool_call>"},{"begin":"<tool_call>lookup<tool_sep>\n","content":{"type":"any_text","excludes":[]},"end":"</tool_call>"}],"separator":"\n","at_least_one":true,"stop_after_first":false},"end":"\n</tool_calls>"}],"at_least_one":false,"stop_after_first":false,"excludes":[]}}"#]].assert_eq(&auto);
-        expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls>\n"},{"type":"tag","begin":"<tool_call>lookup<tool_sep>\n","content":{"type":"any_text","excludes":[]},"end":"</tool_call>"},{"type":"const_string","value":"\n</tool_calls>"}]}}"#]].assert_eq(&forced);
+        expect![[r#"{"type":"structural_tag","format":{"type":"triggered_tags","triggers":["<tool_calls>"],"tags":[{"type":"tag","begin":"<tool_calls>\n","content":{"type":"tags_with_separator","tags":[{"type":"tag","begin":"<tool_call>search<tool_sep>\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":"</tool_call>"},{"type":"tag","begin":"<tool_call>lookup<tool_sep>\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":"</tool_call>"}],"separator":"\n","at_least_one":true,"stop_after_first":false},"end":"\n</tool_calls>"}],"at_least_one":false,"stop_after_first":false,"excludes":[]}}"#]].assert_eq(&auto);
+        expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls>\n"},{"type":"tag","begin":"<tool_call>lookup<tool_sep>\n","content":{"type":"any_text","excludes":[],"max_tokens":null,"max_chars":null},"end":"</tool_call>"},{"type":"const_string","value":"\n</tool_calls>"}]}}"#]].assert_eq(&forced);
     }
 
     #[test]
@@ -314,7 +314,7 @@ mod tests {
         );
         let value = serde_json::to_value(tag).unwrap();
 
-        expect![[r#"{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key>"},{"type":"const_string","value":"query"},{"type":"const_string","value":"</arg_key>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value>"},{"type":"any_text","excludes":["<tool_calls>","</tool_calls>","<tool_call>","</tool_call>","<tool_sep>","<arg_key>","</arg_key>","<arg_value>","</arg_value>"]},{"type":"const_string","value":"</arg_value>"},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key>"},{"type":"const_string","value":"tenant"},{"type":"const_string","value":"</arg_key>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value>"},{"type":"any_text","excludes":["<tool_calls>","</tool_calls>","<tool_call>","</tool_call>","<tool_sep>","<arg_key>","</arg_key>","<arg_value>","</arg_value>"]},{"type":"const_string","value":"</arg_value>"},{"type":"const_string","value":"\n"}]}]}"#]].assert_eq(
+        expect![[r#"{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key>"},{"type":"const_string","value":"query"},{"type":"const_string","value":"</arg_key>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value>"},{"type":"any_text","excludes":["<tool_calls>","</tool_calls>","<tool_call>","</tool_call>","<tool_sep>","<arg_key>","</arg_key>","<arg_value>","</arg_value>"],"max_tokens":null,"max_chars":null},{"type":"const_string","value":"</arg_value>"},{"type":"const_string","value":"\n"}]},{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key>"},{"type":"const_string","value":"tenant"},{"type":"const_string","value":"</arg_key>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value>"},{"type":"any_text","excludes":["<tool_calls>","</tool_calls>","<tool_call>","</tool_call>","<tool_sep>","<arg_key>","</arg_key>","<arg_value>","</arg_value>"],"max_tokens":null,"max_chars":null},{"type":"const_string","value":"</arg_value>"},{"type":"const_string","value":"\n"}]}]}"#]].assert_eq(
             &serde_json::to_string(
                 &value["format"]["elements"][1]["tags"][0]["content"],
             )
diff --git a/rust/src/parser/src/tool/mimo.rs b/rust/src/parser/src/tool/mimo.rs
new file mode 100644
index 000000000000..2844efc1f5dd
--- /dev/null
+++ b/rust/src/parser/src/tool/mimo.rs
@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use super::qwen_coder::{Qwen3CoderToolParser, QwenCoderConfig};
+use super::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+
+const MIMO_CONFIG: QwenCoderConfig = QwenCoderConfig {
+    parser_name: "MiMo",
+    tool_call_start: "<tool_call>",
+    first_tool_call_start: "<tool_call>",
+    next_tool_call_start: "<tool_call>",
+    tool_call_end: "</tool_call>",
+    trim_parameter_newlines: false,
+};
+
+/// MiMo V2 uses Qwen Coder's XML grammar with unframed content and parameter values.
+pub struct MiMoToolParser {
+    inner: Qwen3CoderToolParser,
+}
+
+impl ToolParser for MiMoToolParser {
+    fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>> {
+        Ok(Box::new(Self {
+            inner: Qwen3CoderToolParser::with_config(tools, MIMO_CONFIG),
+        }))
+    }
+
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        self.inner.structural_tag_builder()
+    }
+
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.inner.parse_into(chunk, output)
+    }
+
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        self.inner.finish()
+    }
+
+    fn reset(&mut self) -> String {
+        self.inner.reset()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tool::ToolParserTestExt as _;
+    use crate::tool::test_utils::{collect_stream, split_by_chars};
+    use crate::tool::tests::assert_tool_framing_preserves_body_whitespace;
+    use serde_json::{Value, json};
+
+    #[test]
+    fn bare_tool_boundaries_preserve_body_whitespace() {
+        assert_tool_framing_preserves_body_whitespace::<MiMoToolParser>(
+            "",
+            "<tool_call>\n<function=get_weather>\n</function>\n</tool_call>",
+        );
+    }
+
+    #[test]
+    fn adjacent_calls_preserve_parameter_whitespace_and_scalar_types_across_chunks() {
+        let tools = [Tool {
+            name: "convert".into(),
+            description: None,
+            strict: None,
+            parameters: json!({"type": "object", "properties": {
+                "flag": {"type": "boolean"}, "empty": {"type": ["string", "null"]},
+                "text": {"type": "string"}, "literal": {"type": "string"}, "payload": {"type": "object"}
+            }}),
+        }];
+        let call = "<tool_call>\n<function=convert>\n<parameter=flag>True</parameter>\n<parameter=empty>None</parameter>\n<parameter=text>\n value\n</parameter>\n<parameter=literal>None</parameter>\n<parameter=payload>{\"city\":\"杭州\"}</parameter>\n</function>\n</tool_call>";
+        let input = format!("answer\n\n{call}{call}");
+        let expected = MiMoToolParser::create(&tools).unwrap().parse_complete(&input).unwrap();
+        assert_eq!(expected.normal_text(), "answer\n\n");
+        assert_eq!(expected.calls().len(), 2);
+        for call in expected.calls() {
+            assert_eq!(
+                serde_json::from_str::<Value>(&call.arguments).unwrap(),
+                json!({"flag": true, "empty": null, "text": "\n value\n", "literal": "None", "payload": {"city": "杭州"}})
+            );
+        }
+        for size in 1..=input.chars().count() {
+            let mut parser = MiMoToolParser::create(&tools).unwrap();
+            assert_eq!(
+                collect_stream(parser.as_mut(), &split_by_chars(&input, size)),
+                expected
+            );
+        }
+    }
+}
diff --git a/rust/src/parser/src/tool/minimax_m2.rs b/rust/src/parser/src/tool/minimax_m2.rs
index ee48bfe8bc2c..3b868013bb32 100644
--- a/rust/src/parser/src/tool/minimax_m2.rs
+++ b/rust/src/parser/src/tool/minimax_m2.rs
@@ -6,10 +6,12 @@ use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated};
 use winnow::prelude::*;
 use winnow::stream::Partial;
 use winnow::token::{literal, rest, take_until};
+use xgrammar_structural_tag::format::Format;
 
 use super::parameters::ToolSchemas;
 use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker};
 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
+use crate::output_grammar::{self, OutputGrammarContext, OutputGrammarError};
 use crate::tool::{StructuralTagBuilder, Tool};
 
 const TOOL_CALL_START: &str = "<minimax:tool_call>";
@@ -120,6 +122,38 @@ impl ToolParser for MinimaxM2ToolParser {
         Some(xgrammar_structural_tag::Model::Minimax.builder())
     }
 
+    fn build_visible_format(
+        &self,
+        ctx: &OutputGrammarContext<'_>,
+    ) -> output_grammar::Result<Option<Format>> {
+        let Some(format) =
+            output_grammar::visible_format_from_builder(self.structural_tag_builder(), ctx)?
+        else {
+            return Ok(None);
+        };
+
+        // The MiniMax builder's reasoning=false form retains the empty
+        // thinking close and separator used by its chat template. Strip that
+        // fixed prefix so the initialized reasoning parser can supply the
+        // actual reasoning phase.
+        let Format::Sequence(sequence) = format else {
+            return Err(OutputGrammarError::UnexpectedBuilderFormat { builder: "minimax" });
+        };
+        let [
+            Format::ConstString(empty_think),
+            Format::ConstString(separator),
+            visible,
+        ] = sequence.elements.as_slice()
+        else {
+            return Err(OutputGrammarError::UnexpectedBuilderFormat { builder: "minimax" });
+        };
+        if empty_think.value != "\n</think>\n\n" || separator.value != "\n\n" {
+            return Err(OutputGrammarError::UnexpectedBuilderFormat { builder: "minimax" });
+        }
+
+        Ok(Some(visible.clone()))
+    }
+
     fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
 
diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs
index 57bbb5952889..94f0b4671175 100644
--- a/rust/src/parser/src/tool/mod.rs
+++ b/rust/src/parser/src/tool/mod.rs
@@ -11,6 +11,7 @@ mod glm_xml;
 mod hy;
 pub(crate) mod json;
 mod kimi_k2;
+mod mimo;
 mod minimax_m2;
 mod minimax_m3;
 mod parameters;
@@ -30,6 +31,7 @@ pub use json::{
     MistralToolParser, Phi4MiniJsonToolParser, Qwen3XmlToolParser,
 };
 pub use kimi_k2::KimiK2ToolParser;
+pub use mimo::MiMoToolParser;
 pub use minimax_m2::MinimaxM2ToolParser;
 pub use minimax_m3::MinimaxM3ToolParser;
 pub use qwen_coder::Qwen3CoderToolParser;
@@ -37,7 +39,9 @@ pub use seed_oss::SeedOssToolParser;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 pub use xgrammar_structural_tag::builders::StructuralTagBuilder;
+use xgrammar_structural_tag::format::Format;
 
+use crate::output_grammar::{self, OutputGrammarContext};
 use crate::utils;
 
 /// One function-style tool made available to the model.
@@ -187,7 +191,23 @@ pub trait ToolParser: Send {
         false
     }
 
+    /// Build the language of everything the model may emit after reasoning
+    /// ends: visible text plus tool calls. `None` means no tool grammar applies.
+    ///
+    /// The default forwards to [`Self::structural_tag_builder`] with reasoning
+    /// disabled, which is exactly today's strict-tool-calling grammar.
+    fn build_visible_format(
+        &self,
+        ctx: &OutputGrammarContext<'_>,
+    ) -> output_grammar::Result<Option<Format>> {
+        output_grammar::visible_format_from_builder(self.structural_tag_builder(), ctx)
+    }
+
     /// Return the xgrammar structural-tag builder used for strict tool calling.
+    ///
+    /// To be deprecated: this only exists to back the default
+    /// [`Self::build_visible_format`]. New parsers should override that method
+    /// directly instead of exposing a builder.
     fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
         None
     }
diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs
index 913ad815e5f8..2ae26cd894b4 100644
--- a/rust/src/parser/src/tool/parameters.rs
+++ b/rust/src/parser/src/tool/parameters.rs
@@ -284,13 +284,19 @@ impl JsonParamType {
     }
 }
 
+/// Recognize JSON and Python null spellings emitted by chat templates.
+fn is_null_literal(value: &str) -> bool {
+    value.eq_ignore_ascii_case("null") || value.eq_ignore_ascii_case("none")
+}
+
 /// Convert one parameter input to a normalized JSON value.
 fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &ParamInput) -> Value {
     // Coerce the literal text `null` to JSON null, except for `string`-typed
     // params, where it must stay the string "null": a model emitting the literal
     // text "null" for a string field means the string, not a missing value.
+    // Python-style `None` follows the same schema coercion rules.
     if let ParamInput::Text(value) = input
-        && value.eq_ignore_ascii_case("null")
+        && is_null_literal(value)
         && param_type != Some(&JsonParamType::String)
     {
         return Value::Null;
@@ -333,7 +339,7 @@ fn try_convert_text_value(param_type: &JsonParamType, value: &str) -> Option<Val
             // For composite types with string input, simply interpret the string as JSON.
             serde_json::from_str(value).ok()
         }
-        JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null),
+        JsonParamType::Null => is_null_literal(value).then_some(Value::Null),
         JsonParamType::OneOf(types) => {
             types.iter().find_map(|param_type| try_convert_text_value(param_type, value))
         }
@@ -710,11 +716,12 @@ mod tests {
             }
         }));
 
-        assert_eq!(params.convert("name", text("null")), json!("null"));
-        assert_eq!(params.convert("name", text("NULL")), json!("NULL"));
-        // Non-string and schema-less params are unchanged: "null" -> null.
-        assert_eq!(params.convert("count", text("null")), json!(null));
-        assert_eq!(params.convert("anything", text("null")), json!(null));
+        for literal in ["null", "NULL", "None", "none", "NONE"] {
+            assert_eq!(params.convert("name", text(literal)), json!(literal));
+            // Non-string and schema-less params share the same null coercion.
+            assert_eq!(params.convert("count", text(literal)), json!(null));
+            assert_eq!(params.convert("anything", text(literal)), json!(null));
+        }
     }
 
     #[test]
@@ -731,9 +738,11 @@ mod tests {
             }
         }));
 
-        assert_eq!(params.convert("mode", text("null")), json!(null));
+        for literal in ["null", "NULL", "None", "none", "NONE"] {
+            assert_eq!(params.convert("mode", text(literal)), json!(null));
+            assert_eq!(params.convert("color", text(literal)), json!(literal));
+        }
         assert_eq!(params.convert("mode", text("auto")), json!("auto"));
-        assert_eq!(params.convert("color", text("null")), json!("null"));
     }
 
     #[test]
diff --git a/rust/src/parser/src/tool/qwen_coder.rs b/rust/src/parser/src/tool/qwen_coder.rs
index 5bfaa764f695..e233f4720c05 100644
--- a/rust/src/parser/src/tool/qwen_coder.rs
+++ b/rust/src/parser/src/tool/qwen_coder.rs
@@ -21,9 +21,9 @@ const PARAMETER_END: &str = "</parameter>";
 
 /// Model-specific configuration for the shared Qwen Coder grammar.
 ///
-/// Only the tool-call wrapper tokens vary across models that reuse this
-/// grammar; the inner `<function=...>` / `<parameter=...>` tags are always
-/// byte-identical. Seed-OSS, for example, wraps the same body in
+/// Models may vary the tool-call wrapper and parameter newline framing;
+/// the inner `<function=...>` / `<parameter=...>` tags are byte-identical.
+/// Seed-OSS, for example, wraps the same body in
 /// `<seed:tool_call>` / `</seed:tool_call>`.
 #[derive(Debug, Clone, Copy)]
 pub(crate) struct QwenCoderConfig {
@@ -35,6 +35,8 @@ pub(crate) struct QwenCoderConfig {
     pub(crate) next_tool_call_start: &'static str,
     /// Marker that closes a tool-call block.
     pub(crate) tool_call_end: &'static str,
+    /// Consume one template-inserted newline at each parameter value boundary.
+    pub(crate) trim_parameter_newlines: bool,
 }
 
 const QWEN_CODER_CONFIG: QwenCoderConfig = QwenCoderConfig {
@@ -43,6 +45,7 @@ const QWEN_CODER_CONFIG: QwenCoderConfig = QwenCoderConfig {
     first_tool_call_start: "\n\n<tool_call>",
     next_tool_call_start: "\n<tool_call>",
     tool_call_end: TOOL_CALL_END,
+    trim_parameter_newlines: true,
 };
 
 type QwenCoderInput<'i> = Partial<&'i str>;
@@ -200,7 +203,7 @@ fn parse_next_qwen_coder_event(
     match mode {
         QwenCoderMode::Text => parse_text_event(input, config, after_tool_call),
         QwenCoderMode::ToolCall { end_marker_scan } => {
-            tool_call_event(input, end_marker_scan, config.tool_call_end)
+            tool_call_event(input, end_marker_scan, config)
         }
     }
 }
@@ -231,26 +234,26 @@ fn parse_text_event(
 fn tool_call_event(
     input: &mut QwenCoderInput<'_>,
     end_marker_scan: &mut MarkerScanState,
-    tool_call_end: &'static str,
+    config: QwenCoderConfig,
 ) -> ModalResult<QwenCoderEvent> {
     let (body,) = seq!(
         _: ws0,
-        take_until_marker(tool_call_end, end_marker_scan),
-        _: literal(tool_call_end),
+        take_until_marker(config.tool_call_end, end_marker_scan),
+        _: literal(config.tool_call_end),
     )
     .parse_next(input)?;
 
-    parse_tool_call_body(body)
+    parse_tool_call_body(body, config.trim_parameter_newlines)
 }
 
 /// Parse a Qwen Coder function block.
-fn function_event(input: &mut &str) -> ModalResult<QwenCoderEvent> {
+fn function_event(input: &mut &str, trim_parameter_newlines: bool) -> ModalResult<QwenCoderEvent> {
     let (name, raw_params) = seq!(
         _: literal(FUNCTION_START),
         take_until(1.., ">"),
         _: ">",
         _: ws0,
-        repeat(0.., terminated(parameter, ws0)),
+        repeat(0.., terminated(|input: &mut &str| parameter(input, trim_parameter_newlines), ws0)),
         _: literal(FUNCTION_END),
     )
     .parse_next(input)?;
@@ -262,23 +265,33 @@ fn function_event(input: &mut &str) -> ModalResult<QwenCoderEvent> {
 }
 
 /// Parse a Qwen Coder parameter block.
-fn parameter(input: &mut &str) -> ModalResult<(String, String)> {
+fn parameter(input: &mut &str, trim_newlines: bool) -> ModalResult<(String, String)> {
     let (name, value) = seq!(
         _: literal(PARAMETER_START),
         take_until(1.., ">"),
         _: ">",
-        take_until(0.., PARAMETER_END).map(trim_one_wrapping_newline),
+        take_until(0.., PARAMETER_END),
         _: literal(PARAMETER_END),
     )
     .parse_next(input)?;
 
+    let value = if trim_newlines {
+        trim_one_wrapping_newline(value)
+    } else {
+        value
+    };
     Ok((name.to_string(), value.to_string()))
 }
 
 /// Parse a Qwen Coder tool-call body.
-fn parse_tool_call_body(body: &str) -> ModalResult<QwenCoderEvent> {
+fn parse_tool_call_body(body: &str, trim_parameter_newlines: bool) -> ModalResult<QwenCoderEvent> {
     let mut input = body;
-    delimited(ws0, function_event, (ws0, eof)).parse_next(&mut input)
+    delimited(
+        ws0,
+        |input: &mut &str| function_event(input, trim_parameter_newlines),
+        (ws0, eof),
+    )
+    .parse_next(&mut input)
 }
 
 /// Trim a single leading and trailing newline from a parameter value.
diff --git a/rust/src/parser/src/tool/seed_oss.rs b/rust/src/parser/src/tool/seed_oss.rs
index 160297523e7a..3635998a0356 100644
--- a/rust/src/parser/src/tool/seed_oss.rs
+++ b/rust/src/parser/src/tool/seed_oss.rs
@@ -10,6 +10,7 @@ const SEED_OSS_CONFIG: QwenCoderConfig = QwenCoderConfig {
     first_tool_call_start: "\n\n<seed:tool_call>",
     next_tool_call_start: "\n<seed:tool_call>",
     tool_call_end: "</seed:tool_call>",
+    trim_parameter_newlines: true,
 };
 
 /// Tool parser for Seed-OSS XML-style tool calls.
diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs
index a8963517b347..48c77ef90035 100644
--- a/rust/src/parser/src/unified/combined.rs
+++ b/rust/src/parser/src/unified/combined.rs
@@ -6,8 +6,9 @@
 use vllm_tokenizer::{DecodedText, DynTokenizer};
 
 use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
+use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext};
 use crate::reasoning::ReasoningParser;
-use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Tool, ToolParser, ToolParserOutput};
 
 /// Unified parser that composes existing reasoning and tool parsers.
 pub struct CombinedParser {
@@ -79,8 +80,29 @@ impl UnifiedParser for CombinedParser {
             || self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens())
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        self.tool.as_ref().and_then(|parser| parser.structural_tag_builder())
+    fn build_output_grammar(
+        &self,
+        ctx: &OutputGrammarContext<'_>,
+    ) -> output_grammar::Result<Option<BuiltOutputGrammar>> {
+        let Some(tool) = self.tool.as_ref() else {
+            return Ok(None);
+        };
+        let Some(visible) = tool.build_visible_format(ctx)? else {
+            return Ok(None);
+        };
+        // The reasoning parser may wrap that language with the reasoning phase
+        // implied by its initialized (prompt-derived) state. If it does, the
+        // grammar describes the stream from the first generated token and the
+        // engine can skip its own reasoning gate; if it declines, the grammar
+        // covers only the final output and the engine gate stays in charge.
+        let wrapped = match self.reasoning.as_ref() {
+            Some(reasoning) => reasoning.wrap_visible_format(ctx, &visible)?,
+            None => None,
+        };
+        Ok(Some(match wrapped {
+            Some(full) => BuiltOutputGrammar::from_token_zero(full),
+            None => BuiltOutputGrammar::final_output_only(visible),
+        }))
     }
 
     fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -133,11 +155,26 @@ mod tests {
     use std::sync::Arc;
 
     use vllm_tokenizer::test_utils::TestTokenizer;
-    use vllm_tokenizer::{DecodedText, TokenAnchor, TokenAttribution};
+    use vllm_tokenizer::{DecodedText, DynTokenizer, TokenAnchor, TokenAttribution, Tokenizer};
+    use xgrammar_structural_tag::ToolChoice;
+    use xgrammar_structural_tag::builders::ReasoningMode;
+    use xgrammar_structural_tag::format::Format;
 
     use super::CombinedParser;
-    use crate::reasoning::{Qwen3ReasoningParser, ReasoningDelta, ReasoningParser};
-    use crate::tool::{Qwen3XmlToolParser, Tool, ToolParser};
+    use crate::output_grammar::{
+        GrammarCoverage, OutputGrammarContext, full_format_from_builder_for_test,
+    };
+    use crate::reasoning::{
+        DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser, Glm45ReasoningParser,
+        Glm47ReasoningParser, KimiK2ReasoningParser, MiniMaxM2ReasoningParser,
+        NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningParser,
+        Step3ReasoningParser,
+    };
+    use crate::tool::{
+        DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
+        Glm47MoeToolParser, KimiK2ToolParser, MinimaxM2ToolParser, Qwen3CoderToolParser,
+        Qwen3XmlToolParser, Tool, ToolParser,
+    };
     use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
 
     fn tokenizer() -> TestTokenizer {
@@ -169,6 +206,87 @@ mod tests {
         output
     }
 
+    fn assert_builder_parity<R, T>(name: &str)
+    where
+        R: ReasoningParser + 'static,
+        T: ToolParser + 'static,
+    {
+        // All builders' reasoning=true forms begin inside a reasoning block.
+        // Auto starts outside reasoning and allows a generated reasoning block.
+        for (mode, prompt) in [
+            (ReasoningMode::Enabled, &[256][..]),
+            (ReasoningMode::Auto, &[][..]),
+        ] {
+            for tool_choice in [
+                ToolChoice::required(),
+                ToolChoice::function("get_weather"),
+                ToolChoice::auto(),
+            ] {
+                let mut tools = test_tools();
+                if tool_choice == ToolChoice::auto() {
+                    tools[0].strict = Some(true);
+                }
+                let ctx = OutputGrammarContext {
+                    tools: &tools,
+                    tool_choice: &tool_choice,
+                    tool_strict_level: Default::default(),
+                    parallel_tool_calls: true,
+                };
+                let tool = T::create(&tools).unwrap();
+                let expected = full_format_from_builder_for_test(
+                    tool.structural_tag_builder().unwrap(),
+                    &ctx,
+                    mode,
+                )
+                .unwrap()
+                .unwrap();
+                let reasoning = R::create(Arc::new(tokenizer()) as DynTokenizer).unwrap();
+                let mut parser = CombinedParser::new(Some(reasoning), Some(tool));
+                parser.initialize(prompt).unwrap();
+                let actual = parser.build_output_grammar(&ctx).unwrap().unwrap();
+
+                assert_eq!(
+                    actual.coverage,
+                    GrammarCoverage::FromTokenZero,
+                    "{name} {mode:?} {tool_choice:?}"
+                );
+                assert_eq!(
+                    normalize_builder_parity(actual.format),
+                    normalize_builder_parity(expected),
+                    "{name} {mode:?} {tool_choice:?}"
+                );
+            }
+        }
+    }
+
+    fn normalize_builder_parity(format: Format) -> Format {
+        match format {
+            Format::AnyText(mut text) => {
+                text.excludes.clear();
+                Format::AnyText(text)
+            }
+            Format::Tag(mut tag) => {
+                tag.content = Box::new(normalize_builder_parity(*tag.content));
+                Format::Tag(tag)
+            }
+            Format::Optional(mut optional) => {
+                optional.content = Box::new(normalize_builder_parity(*optional.content));
+                Format::Optional(optional)
+            }
+            Format::Sequence(sequence) => {
+                let mut elements = Vec::new();
+                for element in sequence.elements {
+                    match normalize_builder_parity(element) {
+                        Format::Sequence(nested) => elements.extend(nested.elements),
+                        element => elements.push(element),
+                    }
+                }
+                Format::sequence(elements)
+            }
+            format => format,
+        }
+    }
+
     struct PreserveReasoningParser;
 
     impl ReasoningParser for PreserveReasoningParser {
@@ -332,7 +450,6 @@ mod tests {
     fn combined_parser_emits_tool_calls_from_visible_content() {
         let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap();
         let mut parser = CombinedParser::new(None, Some(tool));
-        assert!(parser.structural_tag_builder().is_some());
 
         let output = collect(
             &mut parser,
@@ -380,4 +497,153 @@ mod tests {
         parser = CombinedParser::new(None, Some(Box::new(PreserveToolParser)));
         assert!(parser.preserve_special_tokens());
     }
+
+    #[test]
+    fn split_parser_wrappers_match_builder_full_formats() {
+        assert_builder_parity::<Qwen3ReasoningParser, Qwen3XmlToolParser>("qwen3");
+        assert_builder_parity::<Qwen3ReasoningParser, Qwen3CoderToolParser>("qwen3.5");
+        assert_builder_parity::<DeepSeekR1ReasoningParser, DeepSeekV3ToolParser>("deepseek-r1");
+        assert_builder_parity::<DeepSeekV3ReasoningParser, DeepSeekV31ToolParser>("deepseek-v3.1");
+        assert_builder_parity::<DeepSeekV3ReasoningParser, DeepSeekV32ToolParser>("deepseek-v3.2");
+        assert_builder_parity::<DeepSeekV3ReasoningParser, DeepSeekV4ToolParser>("deepseek-v4");
+        assert_builder_parity::<Glm47ReasoningParser, Glm47MoeToolParser>("glm-4.7");
+        assert_builder_parity::<KimiK2ReasoningParser, KimiK2ToolParser>("kimi-k2");
+        assert_builder_parity::<MiniMaxM2ReasoningParser, MinimaxM2ToolParser>("minimax-m2");
+    }
+
+    #[test]
+    fn outside_reasoning_allows_optional_generated_reasoning() {
+        let tools = test_tools();
+        let tool_choice = ToolChoice::required();
+        let ctx = OutputGrammarContext {
+            tools: &tools,
+            tool_choice: &tool_choice,
+            tool_strict_level: Default::default(),
+            parallel_tool_calls: true,
+        };
+        let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap();
+        let tool = Qwen3XmlToolParser::create(&tools).unwrap();
+        let visible = tool.build_visible_format(&ctx).unwrap().unwrap();
+        let mut parser = CombinedParser::new(Some(reasoning), Some(tool));
+        parser.initialize(&[]).unwrap();
+
+        let actual = parser.build_output_grammar(&ctx).unwrap().unwrap();
+        let expected = Format::sequence(vec![
+            Format::optional(Format::sequence(vec![
+                Format::tag("<think>", Format::any_text(), "</think>"),
+                Format::const_string("\n\n"),
+            ])),
+            visible,
+        ]);
+
+        assert_eq!(actual.coverage, GrammarCoverage::FromTokenZero);
+        assert_eq!(actual.format, expected);
+    }
+
+    #[test]
+    fn glm45_wrapper_keeps_only_generated_start_framing() {
+        let tools = test_tools();
+        let tool_choice = ToolChoice::required();
+        let ctx = OutputGrammarContext {
+            tools: &tools,
+            tool_choice: &tool_choice,
+            tool_strict_level: Default::default(),
+            parallel_tool_calls: true,
+        };
+        let visible = Format::const_string("answer");
+
+        // GLM-4.5/4.6 thinking prompts leave the whole reasoning opener to generation.
+        for (prompt, begin) in [(vec![], "\n<think>"), (vec![256], "")] {
+            let mut parser = Glm45ReasoningParser::new(Arc::new(tokenizer())).unwrap();
+            parser.initialize(&prompt).unwrap();
+            let reasoning = Format::sequence(vec![
+                Format::tag(begin, Format::any_text(), "</think>"),
+                Format::const_string("\n"),
+            ]);
+            let reasoning = if prompt.is_empty() {
+                Format::optional(reasoning)
+            } else {
+                reasoning
+            };
+            assert_eq!(
+                parser.wrap_visible_format(&ctx, &visible).unwrap(),
+                Some(Format::sequence(vec![reasoning, visible.clone()])),
+                "prompt {prompt:?}"
+            );
+        }
+    }
+
+    #[test]
+    fn delimited_wrappers_use_model_specific_suffixes() {
+        let tools = test_tools();
+        let tool_choice = ToolChoice::required();
+        let ctx = OutputGrammarContext {
+            tools: &tools,
+            tool_choice: &tool_choice,
+            tool_strict_level: Default::default(),
+            parallel_tool_calls: true,
+        };
+        let visible = Format::const_string("visible");
+
+        let mut nemotron = NemotronV3ReasoningParser::new(Arc::new(tokenizer())).unwrap();
+        nemotron.initialize(&[256]).unwrap();
+        assert_eq!(
+            nemotron.wrap_visible_format(&ctx, &visible).unwrap(),
+            Some(Format::sequence(vec![
+                Format::sequence(vec![
+                    Format::tag("", Format::any_text(), "</think>"),
+                    Format::const_string("\n"),
+                ]),
+                visible.clone(),
+            ]))
+        );
+
+        let mut step3 = Step3ReasoningParser::new(Arc::new(tokenizer())).unwrap();
+        step3.initialize(&[256]).unwrap();
+        assert_eq!(
+            step3.wrap_visible_format(&ctx, &visible).unwrap(),
+            Some(Format::sequence(vec![
+                Format::tag("", Format::any_text(), "</think>"),
+                visible,
+            ]))
+        );
+    }
+
+    #[test]
+    fn prompt_closed_reasoning_keeps_only_unconsumed_separator() {
+        let tools = test_tools();
+        let tool_choice = ToolChoice::required();
+        let ctx = OutputGrammarContext {
+            tools: &tools,
+            tool_choice: &tool_choice,
+            tool_strict_level: Default::default(),
+            parallel_tool_calls: true,
+        };
+        let tokenizer = Arc::new(tokenizer());
+        let reasoning = Qwen3ReasoningParser::create(tokenizer.clone()).unwrap();
+        let tool = Qwen3XmlToolParser::create(&tools).unwrap();
+        let mut parser = CombinedParser::new(Some(reasoning), Some(tool));
+        parser.initialize(&[]).unwrap();
+        let outside = parser.build_output_grammar(&ctx).unwrap().unwrap().format;
+
+        for (prompt, remaining) in [
+            ("</think>", "\n\n"),
+            ("</think>\n", "\n"),
+            ("</think>\n\n", ""),
+        ] {
+            parser.initialize(&tokenizer.encode(prompt, false).unwrap()).unwrap();
+            let actual = parser.build_output_grammar(&ctx).unwrap().unwrap();
+            let expected = if remaining.is_empty() {
+                outside.clone()
+            } else {
+                Format::sequence(vec![Format::const_string(remaining), outside.clone()])
+            };
+            assert_eq!(actual.coverage, GrammarCoverage::FromTokenZero);
+            assert_eq!(
+                normalize_builder_parity(actual.format),
+                normalize_builder_parity(expected),
+                "{prompt:?}"
+            );
+        }
+    }
 }
diff --git a/rust/src/parser/src/unified/hy/hy_v3.rs b/rust/src/parser/src/unified/hy/hy_v3.rs
index 6aebf7f90463..30e021b5bce5 100644
--- a/rust/src/parser/src/unified/hy/hy_v3.rs
+++ b/rust/src/parser/src/unified/hy/hy_v3.rs
@@ -4,8 +4,9 @@
 use vllm_tokenizer::{DecodedText, DynTokenizer};
 
 use super::detect_hy_token_suffix;
+use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext};
 use crate::reasoning::HyReasoningParser;
-use crate::tool::{HyDialect, HyToolMarkers, HyToolParser, StructuralTagBuilder, Tool};
+use crate::tool::{HyDialect, HyToolMarkers, HyToolParser, Tool};
 use crate::unified::{CombinedParser, Result, UnifiedParser, UnifiedParserOutput, token_id};
 
 /// Unified reasoning and tool parser for HY3 output.
@@ -46,8 +47,11 @@ impl UnifiedParser for HyV3UnifiedParser {
         self.inner.preserve_special_tokens()
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        self.inner.structural_tag_builder()
+    fn build_output_grammar(
+        &self,
+        ctx: &OutputGrammarContext<'_>,
+    ) -> output_grammar::Result<Option<BuiltOutputGrammar>> {
+        self.inner.build_output_grammar(ctx)
     }
 
     fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -73,12 +77,10 @@ mod tests {
 
     use serde_json::json;
     use vllm_tokenizer::{DecodedText, Tokenizer, test_utils::TestTokenizer};
-    use xgrammar_structural_tag::builders::StructuralTagOptions;
-    use xgrammar_structural_tag::{
-        FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,
-    };
+    use xgrammar_structural_tag::{StructuralTag, ToolChoice};
 
     use super::{HyV3UnifiedParser, UnifiedParser};
+    use crate::output_grammar::{GrammarCoverage, OutputGrammarContext};
     use crate::tool::Tool;
     use crate::unified::{UnifiedParserEvent, UnifiedParserOutput};
 
@@ -151,27 +153,58 @@ mod tests {
 
     #[test]
     fn structural_tag_uses_tokenizer_detected_suffix() {
-        let parser = HyV3UnifiedParser::new(&tools(), Arc::new(tokenizer())).unwrap();
-        let structural_tools = [ToolParam::Function(FunctionToolParam::new(
-            FunctionDefinition::new("get_weather").with_parameters(json!({
-                "type": "object",
-                "properties": { "city": { "type": "string" } },
-                "required": ["city"]
-            })),
-        ))];
-
-        let tag = build_structural_tag(
-            parser.structural_tag_builder().unwrap(),
-            &structural_tools,
-            ToolChoice::required(),
-            StructuralTagOptions::default().with_reasoning(false),
-        )
-        .unwrap()
-        .to_json_string()
-        .unwrap();
-
+        let tools = tools();
+        let tokenizer = Arc::new(tokenizer());
+        let mut parser = HyV3UnifiedParser::new(&tools, tokenizer.clone()).unwrap();
+        let prompt = tokenizer.encode("<think:opensource>", false).unwrap();
+        parser.initialize(&prompt).unwrap();
+        let tag = parser
+            .build_output_grammar(&OutputGrammarContext {
+                tools: &tools,
+                tool_strict_level: Default::default(),
+                parallel_tool_calls: true,
+                tool_choice: &ToolChoice::required(),
+            })
+            .unwrap()
+            .unwrap();
+        assert_eq!(tag.coverage, GrammarCoverage::FromTokenZero);
+        let tag = StructuralTag::new(tag.format).to_json_string().unwrap();
+
+        assert!(tag.contains("</think:opensource>"));
         assert!(tag.contains("<tool_calls:opensource>"));
         assert!(tag.contains("<arg_value:opensource>"));
         assert!(!tag.contains("glm_xml"));
     }
+
+    #[test]
+    fn grammar_root_follows_prompt_reasoning_state() {
+        let tools = tools();
+        let tokenizer = Arc::new(tokenizer());
+        let mut parser = HyV3UnifiedParser::new(&tools, tokenizer.clone()).unwrap();
+        for (prompt, inside) in [
+            ("<think:opensource>", true),
+            ("</think:opensource>", false),
+            ("", false),
+        ] {
+            parser.initialize(&tokenizer.encode(prompt, false).unwrap()).unwrap();
+            let grammar = parser
+                .build_output_grammar(&OutputGrammarContext {
+                    tools: &tools,
+                    tool_choice: &ToolChoice::required(),
+                    tool_strict_level: Default::default(),
+                    parallel_tool_calls: true,
+                })
+                .unwrap()
+                .unwrap();
+            let value = serde_json::to_value(grammar.format).unwrap();
+            let root = &value["elements"][0];
+            if inside {
+                assert_eq!(root["begin"], "");
+                assert_eq!(root["end"], "</think:opensource>");
+            } else {
+                assert_eq!(root["type"], "optional");
+                assert_eq!(root["content"]["begin"], "<think:opensource>");
+            }
+        }
+    }
 }
diff --git a/rust/src/parser/src/unified/hy/hy_v4.rs b/rust/src/parser/src/unified/hy/hy_v4.rs
index 58ee042f3cce..777688bcaa05 100644
--- a/rust/src/parser/src/unified/hy/hy_v4.rs
+++ b/rust/src/parser/src/unified/hy/hy_v4.rs
@@ -4,8 +4,9 @@
 use vllm_tokenizer::{DecodedText, DynTokenizer};
 
 use super::detect_hy_token_suffix;
+use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext};
 use crate::reasoning::HyReasoningParser;
-use crate::tool::{HyDialect, HyToolMarkers, HyToolParser, StructuralTagBuilder, Tool};
+use crate::tool::{HyDialect, HyToolMarkers, HyToolParser, Tool};
 use crate::unified::{CombinedParser, Result, UnifiedParser, UnifiedParserOutput, token_id};
 
 /// Unified reasoning and tool parser for HY4 output.
@@ -46,8 +47,11 @@ impl UnifiedParser for HyV4UnifiedParser {
         self.inner.preserve_special_tokens()
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        self.inner.structural_tag_builder()
+    fn build_output_grammar(
+        &self,
+        ctx: &OutputGrammarContext<'_>,
+    ) -> output_grammar::Result<Option<BuiltOutputGrammar>> {
+        self.inner.build_output_grammar(ctx)
     }
 
     fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -73,12 +77,10 @@ mod tests {
 
     use serde_json::json;
     use vllm_tokenizer::{DecodedText, Tokenizer, test_utils::TestTokenizer};
-    use xgrammar_structural_tag::builders::StructuralTagOptions;
-    use xgrammar_structural_tag::{
-        FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,
-    };
+    use xgrammar_structural_tag::{StructuralTag, ToolChoice};
 
     use super::{HyV4UnifiedParser, UnifiedParser};
+    use crate::output_grammar::{GrammarCoverage, OutputGrammarContext};
     use crate::tool::Tool;
     use crate::unified::{UnifiedParserEvent, UnifiedParserOutput};
 
@@ -164,26 +166,25 @@ mod tests {
 
     #[test]
     fn structural_tag_uses_compact_hy4_skeleton() {
-        let parser = HyV4UnifiedParser::new(&tools(), Arc::new(tokenizer())).unwrap();
-        let structural_tools = [ToolParam::Function(FunctionToolParam::new(
-            FunctionDefinition::new("get_weather").with_parameters(json!({
-                "type": "object",
-                "properties": { "city": { "type": "string" } },
-                "required": ["city"]
-            })),
-        ))];
-
-        let tag = build_structural_tag(
-            parser.structural_tag_builder().unwrap(),
-            &structural_tools,
-            ToolChoice::required(),
-            StructuralTagOptions::default().with_reasoning(false),
-        )
-        .unwrap()
-        .to_json_string()
-        .unwrap();
+        let tools = tools();
+        let tokenizer = Arc::new(tokenizer());
+        let mut parser = HyV4UnifiedParser::new(&tools, tokenizer.clone()).unwrap();
+        let prompt = tokenizer.encode("<think:opensource>", false).unwrap();
+        parser.initialize(&prompt).unwrap();
+        let tag = parser
+            .build_output_grammar(&OutputGrammarContext {
+                tools: &tools,
+                tool_strict_level: Default::default(),
+                parallel_tool_calls: true,
+                tool_choice: &ToolChoice::required(),
+            })
+            .unwrap()
+            .unwrap();
+        assert_eq!(tag.coverage, GrammarCoverage::FromTokenZero);
+        let tag = StructuralTag::new(tag.format).to_json_string().unwrap();
 
         assert!(tag.contains("<tool_calls:opensource>"));
+        assert!(tag.contains("</think:opensource>"));
         assert!(tag.contains("<tool_call:opensource>get_weather"));
         assert!(tag.contains("<arg_key:opensource>"));
         assert!(tag.contains("</tool_call:opensource>"));
diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs
index 2f8278dc6251..08ad91193ca8 100644
--- a/rust/src/parser/src/unified/kimi_k3.rs
+++ b/rust/src/parser/src/unified/kimi_k3.rs
@@ -51,7 +51,10 @@ use winnow::token::{literal, rest, take_till, take_until, take_while};
 
 use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER;
 use super::{Result, UnifiedParser, UnifiedParserOutput, token_id};
-use crate::tool::{StructuralTagBuilder, Tool, ToolCallDelta};
+use crate::output_grammar::{
+    self, BuiltOutputGrammar, OutputGrammarContext, visible_format_from_builder,
+};
+use crate::tool::{Tool, ToolCallDelta};
 use crate::unified::parsing_failed;
 use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker};
 
@@ -275,8 +278,14 @@ impl UnifiedParser for KimiK3UnifiedParser {
         true
     }
 
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER)
+    fn build_output_grammar(
+        &self,
+        ctx: &OutputGrammarContext<'_>,
+    ) -> output_grammar::Result<Option<BuiltOutputGrammar>> {
+        Ok(
+            visible_format_from_builder(Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER), ctx)?
+                .map(BuiltOutputGrammar::final_output_only),
+        )
     }
 
     fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> {
diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs
index 0a86a274e1ff..fc30f455e9cc 100644
--- a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs
+++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs
@@ -6,14 +6,14 @@
 use serde_json::{Map, Value};
 use xgrammar_structural_tag::Result;
 use xgrammar_structural_tag::builders::{
-    StructuralTagBuilder, StructuralTagContext, StructuralTagOptions,
+    ReasoningMode, StructuralTagBuilder, StructuralTagContext, StructuralTagOptions,
 };
 use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat};
 use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters};
 
 use super::{
     ARG_CLOSE, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN,
-    RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN,
+    RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN,
 };
 
 pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder =
@@ -50,14 +50,23 @@ impl StructuralTagBuilder for KimiK3StructuralTagBuilder {
     }
 }
 
-fn response_prefix(reasoning: bool) -> Vec<Format> {
+fn response_prefix(reasoning: ReasoningMode) -> Vec<Format> {
     let mut elements = Vec::new();
-    if reasoning {
-        elements.push(Format::tag(
-            "",
+    if reasoning != ReasoningMode::Disabled {
+        let prefix = Format::tag(
+            if reasoning == ReasoningMode::Auto {
+                THINK_OPEN
+            } else {
+                ""
+            },
             Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]),
             THINK_CLOSE,
-        ));
+        );
+        elements.push(if reasoning == ReasoningMode::Auto {
+            Format::optional(prefix)
+        } else {
+            prefix
+        });
         elements.push(Format::const_string(RESPONSE_OPEN));
     } else {
         elements.push(Format::optional(Format::const_string(RESPONSE_OPEN)));
@@ -386,7 +395,7 @@ mod tests {
         )
         .unwrap();
 
-        expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|open|>","<|close|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap());
+        expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|open|>","<|close|>","<|end_of_msg|>"],"max_tokens":null,"max_chars":null},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"type":"tag","begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap());
     }
 
     #[test]
diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs
index 17c3dd1b2430..e65f2de08728 100644
--- a/rust/src/parser/src/unified/mod.rs
+++ b/rust/src/parser/src/unified/mod.rs
@@ -18,10 +18,9 @@ use thiserror::Error;
 use thiserror_ext::Macro;
 use vllm_tokenizer::{DecodedText, DynTokenizer};
 
+use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext};
 use crate::reasoning::ReasoningError;
-use crate::tool::{
-    StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
-};
+use crate::tool::{Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput};
 
 /// Result alias for unified parser operations.
 pub type Result<T> = std::result::Result<T, UnifiedParserError>;
@@ -240,9 +239,12 @@ pub trait UnifiedParser: Send {
         false
     }
 
-    /// Return the xgrammar structural-tag builder used for strict tool calling.
-    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
-        None
+    /// Build the request output grammar after prompt-based initialization.
+    fn build_output_grammar(
+        &self,
+        _ctx: &OutputGrammarContext<'_>,
+    ) -> output_grammar::Result<Option<BuiltOutputGrammar>> {
+        Ok(None)
     }
 
     /// Return the parser-provided ID for a tool call by index, if the model emitted one.
diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml
index 6f5e95e93f44..8990ba2f3bc3 100644
--- a/rust/src/server/Cargo.toml
+++ b/rust/src/server/Cargo.toml
@@ -21,6 +21,7 @@ llm-multimodal.workspace = true
 openssl.workspace = true
 prost.workspace = true
 prost-types.workspace = true
+rmp-serde.workspace = true
 rmpv.workspace = true
 serde.workspace = true
 serde_json.workspace = true
@@ -37,7 +38,6 @@ tokio-stream.workspace = true
 tokio-util.workspace = true
 tonic.workspace = true
 tonic-health.workspace = true
-tonic-prost.workspace = true
 tower.workspace = true
 tower-http.workspace = true
 tracing.workspace = true
@@ -51,19 +51,16 @@ vllm-chat.workspace = true
 vllm-engine-core-client.workspace = true
 vllm-llm.workspace = true
 vllm-metrics.workspace = true
+vllm-proto.workspace = true
 vllm-text.workspace = true
 
-[build-dependencies]
-protox.workspace = true
-tonic-prost-build.workspace = true
-
 [dev-dependencies]
 anyhow.workspace = true
 async-openai = { workspace = true, features = ["full"] }
 bytes.workspace = true
 clap.workspace = true
 expect-test.workspace = true
-rmp-serde.workspace = true
+prost.workspace = true
 serial_test.workspace = true
 tempfile.workspace = true
 tokio = { workspace = true, features = ["test-util"] }
diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs
index d07ba242aed5..3de227babdbe 100644
--- a/rust/src/server/examples/external_engine_openai_qwen.rs
+++ b/rust/src/server/examples/external_engine_openai_qwen.rs
@@ -16,6 +16,7 @@ use clap::Parser;
 use futures::StreamExt as _;
 use tokio_util::sync::CancellationToken;
 use tracing_subscriber::EnvFilter;
+use vllm_chat::ToolStrictLevel;
 use vllm_engine_core_client::TransportMode;
 use vllm_server::{
     ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, CorsConfig,
@@ -60,6 +61,8 @@ async fn main() -> Result<()> {
         },
         coordinator_mode: CoordinatorMode::MaybeInProc,
         model: args.model,
+        revision: None,
+        hf_overrides: Default::default(),
         generation_config: Default::default(),
         served_model_name: vec![],
         listener_mode: HttpListenerMode::BindTcp {
@@ -68,6 +71,7 @@ async fn main() -> Result<()> {
         },
         tool_call_parser: ParserSelection::Auto,
         reasoning_parser: ParserSelection::Auto,
+        tool_strict_level: ToolStrictLevel::Auto,
         renderer: RendererSelection::Auto,
         language_model_only: false,
         chat_template: None,
diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs
index fde53183f5b3..c1c580e75a9e 100644
--- a/rust/src/server/src/config.rs
+++ b/rust/src/server/src/config.rs
@@ -14,8 +14,10 @@ use serde_json::Value;
 use vllm_chat::multimodal::MmLimitPerPrompt;
 use vllm_chat::{
     ChatTemplateContentFormatOption, GenerationConfigMode, ParserSelection, RendererSelection,
+    ToolStrictLevel,
 };
 use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode};
+use vllm_text::backend::hf::HfOverrides;
 
 /// Default keep-alive idle timeout (seconds); also the head-read bound
 /// when keep-alive is disabled (`0`).
@@ -55,6 +57,11 @@ pub struct ApiServerOptions {
     pub enable_prompt_tokens_details: bool,
     /// When `true`, set `X-Request-Id` on every HTTP response.
     pub enable_request_id_headers: bool,
+    /// When `true`, register the scale-out `/inference/v1/generate` route.
+    pub enable_scale_out: bool,
+    /// Idle interval after which streaming SSE responses send a keep-alive
+    /// comment. `None` disables keep-alive comments.
+    pub sse_keep_alive_interval: Option<Duration>,
 }
 
 /// CORS settings mirroring Python's `CORSMiddleware`; the default is permissive.
@@ -210,6 +217,10 @@ pub struct Config {
     pub coordinator_mode: CoordinatorMode,
     /// Backend model identifier used for engine-core loading.
     pub model: String,
+    /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
+    pub revision: Option<String>,
+    /// JSON Merge Patch applied to the model config before backend initialization.
+    pub hf_overrides: HfOverrides,
     /// Which generation-config sampling defaults to inherit.
     pub generation_config: GenerationConfigMode,
     /// Model name(s) exposed to clients via the OpenAI API. When non-empty,
@@ -222,6 +233,8 @@ pub struct Config {
     pub tool_call_parser: ParserSelection,
     /// Reasoning parser selection.
     pub reasoning_parser: ParserSelection,
+    /// Server-side floor for tool-call structural tags.
+    pub tool_strict_level: ToolStrictLevel,
     /// Chat renderer selection.
     pub renderer: RendererSelection,
     /// Disable frontend-side multimodal preprocessing and render the model as
diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs
index b658388cbdd4..3280cd1dd503 100644
--- a/rust/src/server/src/error.rs
+++ b/rust/src/server/src/error.rs
@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use axum::Json;
+use axum::extract::rejection::JsonRejection;
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
 use thiserror_ext::{AsReport as _, Construct, Macro};
@@ -70,6 +71,12 @@ impl ApiError {
     }
 }
 
+impl From<JsonRejection> for ApiError {
+    fn from(error: JsonRejection) -> Self {
+        Self::json_parse_error(error.body_text())
+    }
+}
+
 impl IntoResponse for ApiError {
     fn into_response(self) -> Response {
         (self.status_code(), Json(self.to_error_response())).into_response()
@@ -168,16 +175,22 @@ mod tests {
     }
 
     #[test]
-    fn invalid_reasoning_effort_maps_to_invalid_request() {
-        let error = vllm_chat::Error::InvalidReasoningEffort(
-            "DeepSeek V4.1 reasoning_effort must be within [1, 100]".to_string(),
-        );
-        let api_error = chat_submit_error("failed to submit chat request", error);
-        assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
-        assert_eq!(
-            api_error.to_error_response().error.error_type,
-            "invalid_request_error"
-        );
+    fn invalid_reasoning_parameters_map_to_invalid_request() {
+        for error in [
+            vllm_chat::Error::InvalidReasoningEffort(
+                "DeepSeek V4.1 reasoning_effort must be within [1, 100]".to_string(),
+            ),
+            vllm_chat::Error::InvalidReasoningControl {
+                message: "template kwarg `thinking` must be a boolean".to_string(),
+            },
+        ] {
+            let api_error = chat_submit_error("failed to submit chat request", error);
+            assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
+            assert_eq!(
+                api_error.to_error_response().error.error_type,
+                "invalid_request_error"
+            );
+        }
     }
 
     #[test]
diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs
index ddb74bf40f1d..680ebdfbdaf2 100644
--- a/rust/src/server/src/grpc/control.rs
+++ b/rust/src/server/src/grpc/control.rs
@@ -50,6 +50,7 @@ impl ControlServiceImpl {
             data_parallel_rank: ready.data_parallel_rank,
             decode_context_parallel_size: ready.decode_context_parallel_size,
             world_size: ready.world_size,
+            data_parallel_size_local: self.client().engine_count() as u32,
         }
     }
 
@@ -195,6 +196,7 @@ impl pb::control_server::Control for ControlServiceImpl {
             max_batched_tokens: ready.max_num_batched_tokens,
             max_loras: ready.max_loras,
             rl_capabilities: Some(self.rl_capabilities()),
+            effective_attention_block_size: self.client().effective_attention_block_size(),
         }))
     }
 
diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs
index 2a3f41ffa438..d5dd1ae64df4 100644
--- a/rust/src/server/src/grpc/convert.rs
+++ b/rust/src/server/src/grpc/convert.rs
@@ -4,10 +4,10 @@
 //! Conversion between gRPC protobuf types and internal `vllm-text`
 //! request/response types.
 
+use thiserror_ext::AsReport as _;
 use tonic::Status;
-use url::Url;
 use uuid::Uuid;
-use vllm_chat::MediaContentPart;
+use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};
 use vllm_engine_core_client::protocol::output::StopReason;
 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;
 use vllm_text::{
@@ -17,93 +17,30 @@ use vllm_text::{
 
 use super::pb;
 
-pub fn media_parts_from_request(
-    media: Vec<pb::MediaItem>,
-) -> Result<Vec<MediaContentPart>, Status> {
-    let mut parts = Vec::with_capacity(media.len());
-    for (index, item) in media.into_iter().enumerate() {
-        let modality = item.modality();
-        if modality == pb::Modality::Unspecified {
-            return Err(Status::invalid_argument(format!(
-                "media[{index}].modality is required"
-            )));
-        }
-        let uuid = (!item.uuid.is_empty()).then_some(item.uuid);
-        let mime_type = (!item.mime_type.is_empty()).then_some(item.mime_type);
-        let source = item.source.ok_or_else(|| {
-            Status::invalid_argument(format!("media[{index}].source is required"))
-        })?;
-        match &source {
-            pb::media_item::Source::Url(url) => {
-                validate_media_uri(index, "url", url, &["http", "https"])?;
-            }
-            pb::media_item::Source::DataUri(uri) => {
-                validate_media_uri(index, "data_uri", uri, &["data"])?;
-            }
-            pb::media_item::Source::RawBytes(_) => {}
-        }
-        let part = match (modality, source) {
-            (
-                pb::Modality::Image,
-                pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url),
-            ) => MediaContentPart::ImageUrl {
-                url,
-                detail: None,
-                uuid,
-            },
-            (pb::Modality::Image, pb::media_item::Source::RawBytes(data)) => {
-                MediaContentPart::ImageData {
-                    data,
-                    mime_type,
-                    uuid,
-                    detail: None,
-                }
-            }
-            (
-                pb::Modality::Video,
-                pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url),
-            ) => MediaContentPart::VideoUrl { url, uuid },
-            (pb::Modality::Video, pb::media_item::Source::RawBytes(data)) => {
-                MediaContentPart::VideoData {
-                    data,
-                    mime_type,
-                    uuid,
-                }
-            }
-            (
-                pb::Modality::Audio,
-                pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url),
-            ) => MediaContentPart::AudioUrl { url, uuid },
-            (pb::Modality::Audio, pb::media_item::Source::RawBytes(data)) => {
-                MediaContentPart::AudioData {
-                    data,
-                    mime_type,
-                    uuid,
-                }
-            }
-            (pb::Modality::Unspecified, _) => unreachable!("modality validated above"),
-        };
-        parts.push(part);
+fn kv_hints_from_proto(hints: pb::KvHintsEnvelope) -> KvHintsEnvelope {
+    KvHintsEnvelope {
+        protocol_version: hints.protocol_version,
+        message_id: hints.message_id,
+        actions: hints
+            .actions
+            .into_iter()
+            .map(|action| KvHintAction {
+                action_id: action.action_id,
+                action_type: action.action_type,
+                action_version: action.action_version,
+                payload: action
+                    .payload
+                    .map(|payload| {
+                        payload
+                            .fields
+                            .into_iter()
+                            .map(|(key, value)| (key, proto_value_to_json(&value)))
+                            .collect()
+                    })
+                    .unwrap_or_default(),
+            })
+            .collect(),
     }
-    Ok(parts)
-}
-
-fn validate_media_uri(
-    index: usize,
-    field: &str,
-    value: &str,
-    allowed_schemes: &[&str],
-) -> Result<(), Status> {
-    let uri = Url::parse(value).map_err(|_| {
-        Status::invalid_argument(format!("media[{index}].{field} is not a valid URI"))
-    })?;
-    if !allowed_schemes.contains(&uri.scheme()) {
-        return Err(Status::invalid_argument(format!(
-            "media[{index}].{field} must use the {} scheme",
-            allowed_schemes.join(" or ")
-        )));
-    }
-    Ok(())
 }
 
 // ========================================================================================
@@ -146,6 +83,7 @@ pub fn to_text_request(
         req.request_id
     };
     let session_id = req.session_id.filter(|s| !s.is_empty());
+    let kv_hints = req.kv_hints.map(kv_hints_from_proto);
 
     let sampling = req.sampling.as_ref();
     let decoding = req.decoding.as_ref();
@@ -155,6 +93,7 @@ pub fn to_text_request(
 
     let mut sampling_params =
         build_sampling_params(req.temperature, sampling, decoding, stopping, response)?;
+    sampling_params.watermarking = req.watermarking.unwrap_or(true);
 
     // Thread KVCacheParameters → SamplingParams fields.
     if let Some(kv) = kv {
@@ -197,7 +136,9 @@ pub fn to_text_request(
         add_special_tokens: true,
         data_parallel_rank: None,
         session_id,
-        reasoning_parser_kwargs: None,
+        kv_hints,
+        reasoning_parser_kwargs: Default::default(),
+        reasoning_ended: None,
         lora_request: None,
         arrival_time: None,
     })
@@ -283,7 +224,7 @@ fn build_sampling_params(
     if let Some(r) = response {
         if r.output_logprobs {
             let (count, token_ids) = candidate_logprob_spec(r.output_candidates.as_ref());
-            params.logprobs = Some(count);
+            params.logprobs = count;
             params.logprob_token_ids = token_ids;
         }
         if r.prompt_logprobs {
@@ -300,7 +241,7 @@ fn build_sampling_params(
                 ));
             }
             let (count, _) = candidate_logprob_spec(r.prompt_candidates.as_ref());
-            params.prompt_logprobs = Some(count);
+            params.prompt_logprobs = count;
         }
     }
 
@@ -310,18 +251,22 @@ fn build_sampling_params(
 /// Map the proto `CandidateTokens` selector to a `(logprobs_count,
 /// logprob_token_ids)` pair.
 ///
-/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability
-/// - `all` → `(-1, None)` — return the full vocabulary
-/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the
-///   count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs
-///   are carried via `logprob_token_ids` on `SamplingParams`)
-/// - absent → `(1, None)` — just the sampled/scored token
-fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option<Vec<u32>>) {
+/// - `top_n(k)` → `(Some(k), None)` — return top-k candidates by probability
+/// - `all` → `(Some(-1), None)` — return the full vocabulary
+/// - nonempty `token_ids` → `(None, Some(ids))` — let the engine derive the count
+/// - absent or empty `token_ids` → `(Some(0), None)` — return only the sampled/scored token
+fn candidate_logprob_spec(
+    candidates: Option<&pb::CandidateTokens>,
+) -> (Option<i32>, Option<Vec<u32>>) {
     match candidates.and_then(|c| c.select.as_ref()) {
-        Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None),
-        Some(pb::candidate_tokens::Select::All(true)) => (-1, None),
-        Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())),
-        _ => (1, None),
+        Some(pb::candidate_tokens::Select::TopN(n)) => (Some(*n as i32), None),
+        Some(pb::candidate_tokens::Select::All(true)) => (Some(-1), None),
+        Some(pb::candidate_tokens::Select::TokenIds(ids)) if !ids.ids.is_empty() => {
+            // Match Python HTTP: selected IDs have their own limit, independent
+            // of the numeric top-logprobs count and its max_logprobs cap.
+            (None, Some(ids.ids.clone()))
+        }
+        _ => (Some(0), None),
     }
 }
 
@@ -350,6 +295,9 @@ fn convert_structured_output(
             StructuredOutputsParams::structural_tag(tag.clone())
         }
     };
+    params
+        .validate()
+        .map_err(|error| Status::invalid_argument(error.to_report_string()))?;
     Ok(Some(params))
 }
 
@@ -391,13 +339,19 @@ pub fn to_sequence_output(
     logprobs: Option<&DecodedLogprobs>,
     finished: Option<&Finished>,
     opts: &ResponseOpts,
-) -> pb::SequenceOutput {
+) -> Result<pb::SequenceOutput, Status> {
+    let finish_info = finished.map(|f| to_finish_info(f, token_ids)).transpose()?;
     let (lp_values, rank_values, candidates) = match logprobs {
         Some(lp) if opts.output_logprobs => output_logprobs_to_proto(lp),
         _ => (vec![], vec![], vec![]),
     };
 
-    pb::SequenceOutput {
+    let sampling_mask = finished
+        .and_then(|finished| finished.sampling_mask.as_ref())
+        .map(|mask| mask.rows.iter().map(|row| pb::TokenIds { ids: row.clone() }).collect())
+        .unwrap_or_default();
+
+    Ok(pb::SequenceOutput {
         index: 0, // TODO: multi-sequence (n > 1) not supported
         text: if opts.output_text {
             delta.to_string()
@@ -413,11 +367,12 @@ pub fn to_sequence_output(
         logprobs: lp_values,
         ranks: rank_values,
         candidate_tokens: candidates,
-        finish_info: finished.map(|f| to_finish_info(f, token_ids)),
-    }
+        finish_info,
+        sampling_mask,
+    })
 }
 
-fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
+fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result<pb::FinishInfo, Status> {
     use pb::finish_info::FinishReason as PbFinishReason;
 
     let (finish_reason, stop_reason) = match &finished.finish_reason {
@@ -438,18 +393,17 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
             (PbFinishReason::Stop as i32, sr)
         }
         FinishReason::Length => (PbFinishReason::Length as i32, None),
-        FinishReason::Abort | FinishReason::Error | FinishReason::Repetition(_) => {
-            (PbFinishReason::Aborted as i32, None)
-        }
+        FinishReason::Error => return Err(Status::internal("engine failed during generation")),
+        FinishReason::Abort | FinishReason::Repetition(_) => (PbFinishReason::Aborted as i32, None),
     };
 
-    pb::FinishInfo {
+    Ok(pb::FinishInfo {
         num_output_tokens: finished.usage.output_token_count as u32,
         finish_reason,
         stop_reason,
         kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
         ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct),
-    }
+    })
 }
 
 // ========================================================================================
@@ -604,11 +558,18 @@ impl ResponseOpts {
 
 #[cfg(test)]
 mod tests {
+    use prost::Message as _;
     use vllm_engine_core_client::protocol::output::StopReason;
-    use vllm_text::{FinishReason, Finished, Prompt};
+    use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
+    use vllm_text::{
+        FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params,
+    };
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason};
-    use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request};
+    use super::{
+        ResponseOpts, json_to_proto_struct, pb, to_finish_info, to_sequence_output, to_text_request,
+    };
 
     fn base_request() -> pb::GenerateRequest {
         pb::GenerateRequest {
@@ -619,6 +580,26 @@ mod tests {
         }
     }
 
+    #[test]
+    fn watermarking_defaults_and_opt_out_survive_protobuf_conversion() {
+        for watermarking in [None, Some(true), Some(false)] {
+            let request = pb::GenerateRequest {
+                watermarking,
+                ..base_request()
+            };
+            let encoded = request.encode_to_vec();
+            for stream in [false, true] {
+                let decoded = pb::GenerateRequest::decode(encoded.as_slice()).unwrap();
+                let text = to_text_request(decoded, stream, &["test-model".to_string()])
+                    .expect("convert request");
+                assert_eq!(
+                    text.sampling_params.watermarking,
+                    watermarking.unwrap_or(true)
+                );
+            }
+        }
+    }
+
     #[test]
     fn temperature_propagates_from_top_level_request_field() {
         let req = pb::GenerateRequest {
@@ -629,6 +610,35 @@ mod tests {
         assert_eq!(text.sampling_params.temperature, Some(0.7));
     }
 
+    #[test]
+    fn kv_hints_propagate_from_top_level_request_field() {
+        let req = pb::GenerateRequest {
+            kv_hints: Some(pb::KvHintsEnvelope {
+                protocol_version: "0.1".to_string(),
+                message_id: "msg-1".to_string(),
+                actions: vec![pb::KvHintAction {
+                    action_id: "action-1".to_string(),
+                    action_type: "example.action".to_string(),
+                    action_version: "1.0".to_string(),
+                    payload: json_to_proto_struct(&serde_json::json!({"source": "g2"})),
+                }],
+            }),
+            ..base_request()
+        };
+
+        let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+        let hints = text.kv_hints.expect("kv hints");
+        assert_eq!(hints.protocol_version, "0.1");
+        assert_eq!(hints.message_id, "msg-1");
+        assert_eq!(hints.actions[0].action_id, "action-1");
+        assert_eq!(hints.actions[0].action_type, "example.action");
+        assert_eq!(hints.actions[0].action_version, "1.0");
+        assert_eq!(
+            hints.actions[0].payload.get("source"),
+            Some(&serde_json::json!("g2"))
+        );
+    }
+
     #[test]
     fn unset_temperature_defaults_to_greedy() {
         let text = to_text_request(base_request(), false, &["test-model".to_string()])
@@ -637,6 +647,20 @@ mod tests {
         assert_eq!(text.sampling_params.temperature, Some(0.0));
     }
 
+    #[test]
+    fn grpc_rejects_empty_grammar_before_engine() {
+        use super::pb::decoding_parameters::StructuredOutput;
+        let req = pb::GenerateRequest {
+            decoding: Some(pb::DecodingParameters {
+                structured_output: Some(StructuredOutput::Grammar("  ".to_string())),
+                ..Default::default()
+            }),
+            ..base_request()
+        };
+        let err = to_text_request(req, false, &["test-model".to_string()]).unwrap_err();
+        assert!(err.message().contains("grammar cannot be an empty string"));
+    }
+
     #[test]
     fn absent_seed_is_none() {
         let req = pb::GenerateRequest {
@@ -663,6 +687,69 @@ mod tests {
         assert_eq!(text.sampling_params.seed, Some(0));
     }
 
+    #[test]
+    fn output_logprob_selectors_survive_request_lowering() {
+        use pb::candidate_tokens::Select;
+
+        let selected_ids: Vec<u32> = (100..121).collect();
+        let cases = [
+            (
+                Some(Select::TokenIds(pb::TokenIds {
+                    ids: selected_ids.clone(),
+                })),
+                None,
+                Some(selected_ids),
+            ),
+            (
+                Some(Select::TokenIds(pb::TokenIds {
+                    ids: vec![198, 198],
+                })),
+                None,
+                Some(vec![198, 198]),
+            ),
+            (
+                Some(Select::TokenIds(pb::TokenIds { ids: vec![] })),
+                Some(0),
+                None,
+            ),
+            (None, Some(0), None),
+            (Some(Select::TopN(0)), Some(0), None),
+            (Some(Select::TopN(2)), Some(2), None),
+        ];
+        for (select, expected_count, expected_ids) in cases {
+            let req = pb::GenerateRequest {
+                response: Some(pb::ResponseOptions {
+                    output_logprobs: true,
+                    output_candidates: select.map(|select| pb::CandidateTokens {
+                        select: Some(select),
+                    }),
+                    ..Default::default()
+                }),
+                ..base_request()
+            };
+            let text =
+                to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");
+            let params = lower_sampling_params(
+                text.sampling_params,
+                SamplingHints::default(),
+                SamplingLimits {
+                    max_model_len: 32,
+                    max_logprobs: 20,
+                    model_vocab_size: 512,
+                    tokenizer_vocab_size: 512,
+                },
+                1,
+                &TestTokenizer::new(),
+            )
+            .expect("logprob selector should lower to an engine request");
+
+            assert_eq!(
+                (params.logprobs, params.logprob_token_ids),
+                (expected_count, expected_ids)
+            );
+        }
+    }
+
     #[test]
     fn bypass_prefix_cache_maps_to_skip_reading_prefix_cache() {
         let req = pb::GenerateRequest {
@@ -701,6 +788,7 @@ mod tests {
             finish_reason: reason,
             kv_transfer_params: None,
             ec_transfer_params: None,
+            sampling_mask: None,
         }
     }
 
@@ -709,7 +797,7 @@ mod tests {
         let fin = finished(FinishReason::Stop(None));
         let token_ids = [1_u32, 2, 3, 151643];
 
-        let info = to_finish_info(&fin, &token_ids);
+        let info = to_finish_info(&fin, &token_ids).expect("finish info");
 
         assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
         assert_eq!(info.stop_reason, Some(PbStopReason::EosTokenId(151643)));
@@ -719,7 +807,7 @@ mod tests {
     fn eos_stop_with_empty_token_ids_leaves_stop_reason_unset() {
         let fin = finished(FinishReason::Stop(None));
 
-        let info = to_finish_info(&fin, &[]);
+        let info = to_finish_info(&fin, &[]).expect("finish info");
 
         assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
         assert_eq!(info.stop_reason, None);
@@ -730,7 +818,7 @@ mod tests {
         let fin = finished(FinishReason::Stop(Some(StopReason::TokenId(42))));
         // Terminal token list should be ignored when an explicit stop reason is
         // present.
-        let info = to_finish_info(&fin, &[7, 42]);
+        let info = to_finish_info(&fin, &[7, 42]).expect("finish info");
 
         assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
         assert_eq!(info.stop_reason, Some(PbStopReason::StopTokenId(42)));
@@ -740,7 +828,7 @@ mod tests {
     fn explicit_stop_string_is_preserved() {
         let fin = finished(FinishReason::Stop(Some(StopReason::Text("</stop>".into()))));
 
-        let info = to_finish_info(&fin, &[1, 2, 3]);
+        let info = to_finish_info(&fin, &[1, 2, 3]).expect("finish info");
 
         assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
         assert_eq!(
@@ -753,7 +841,7 @@ mod tests {
     fn length_finish_has_no_stop_reason() {
         let fin = finished(FinishReason::Length);
 
-        let info = to_finish_info(&fin, &[1, 2, 3]);
+        let info = to_finish_info(&fin, &[1, 2, 3]).expect("finish info");
 
         assert_eq!(info.finish_reason, PbFinishReason::Length as i32);
         assert_eq!(info.stop_reason, None);
@@ -763,7 +851,7 @@ mod tests {
     fn abort_finish_is_mapped_to_aborted() {
         let fin = finished(FinishReason::Abort);
 
-        let info = to_finish_info(&fin, &[]);
+        let info = to_finish_info(&fin, &[]).expect("finish info");
 
         assert_eq!(info.finish_reason, PbFinishReason::Aborted as i32);
         assert_eq!(info.stop_reason, None);
@@ -778,10 +866,47 @@ mod tests {
             ..Default::default()
         };
 
-        let out = to_sequence_output("hello", &[10, 20, 30], None, Some(&fin), &opts);
+        let out = to_sequence_output("hello", &[10, 20, 30], None, Some(&fin), &opts)
+            .expect("sequence output");
 
         let finish = out.finish_info.expect("finish_info should be present");
         assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32);
         assert_eq!(finish.stop_reason, Some(PbStopReason::EosTokenId(30)));
     }
+
+    #[test]
+    fn sequence_output_only_carries_sampling_mask_on_terminal_output() {
+        let mut fin = finished(FinishReason::Length);
+        fin.sampling_mask = Some(SamplingMask {
+            rows: vec![vec![1, 10], vec![2, 20]],
+        });
+
+        let terminal =
+            to_sequence_output("", &[10, 20], None, Some(&fin), &ResponseOpts::default())
+                .expect("terminal output");
+        let intermediate = to_sequence_output("", &[9], None, None, &ResponseOpts::default())
+            .expect("intermediate output");
+
+        assert_eq!(
+            terminal.sampling_mask,
+            vec![
+                pb::TokenIds { ids: vec![1, 10] },
+                pb::TokenIds { ids: vec![2, 20] }
+            ]
+        );
+        assert!(intermediate.sampling_mask.is_empty());
+    }
+
+    #[test]
+    fn sampling_mask_uses_additive_proto_tag_ten() {
+        let response = pb::SequenceOutput {
+            sampling_mask: vec![pb::TokenIds { ids: vec![7, 8] }],
+            ..Default::default()
+        };
+
+        assert_eq!(
+            response.encode_to_vec(),
+            vec![0x52, 0x04, 0x0a, 0x02, 0x07, 0x08]
+        );
+    }
 }
diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs
index f864cbde5959..d56415c7475e 100644
--- a/rust/src/server/src/grpc/inference.rs
+++ b/rust/src/server/src/grpc/inference.rs
@@ -149,19 +149,23 @@ impl InferenceServiceImpl {
                 })?);
             }
 
-            let media = convert::media_parts_from_request(media)?;
+            let media = super::media::from_proto(media)?;
             if !media.is_empty() {
                 let Prompt::TokenIds(mut token_ids) = text_request.prompt else {
                     return Err(Status::invalid_argument(
                         "multimodal gRPC requests must provide token_ids input",
                     ));
                 };
-                let mm_features = self
-                    .state
-                    .chat
-                    .prepare_media(media, &mut token_ids)
-                    .await
-                    .map_err(|error| Status::internal(error.to_report_string()))?;
+                let mm_features =
+                    self.state.chat.prepare_media(media, &mut token_ids).await.map_err(
+                        |error| {
+                            if error.is_request_validation_error() {
+                                Status::invalid_argument(error.to_report_string())
+                            } else {
+                                Status::internal(error.to_report_string())
+                            }
+                        },
+                    )?;
                 text_request.prompt = Prompt::TokenIds(token_ids);
                 text_request.mm_features = mm_features;
             }
@@ -241,12 +245,6 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
             .instrument(request_span.clone())
             .await
             .map_err(|error| log_text_error(&request_span, started_at, "collection", error))?;
-        info!(
-            parent: &request_span,
-            elapsed_ms = started_at.elapsed().as_millis() as u64,
-            "gRPC inference request completed"
-        );
-
         // Build the single aggregated response.
         let prompt_info = convert::to_prompt_info(
             &collected.prompt_token_ids,
@@ -259,6 +257,7 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
             finish_reason: collected.finish_reason,
             kv_transfer_params: collected.kv_transfer_params,
             ec_transfer_params: collected.ec_transfer_params,
+            sampling_mask: collected.sampling_mask,
         };
 
         let outputs = convert::to_sequence_output(
@@ -267,6 +266,11 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
             collected.logprobs.as_ref(),
             Some(&finish_info),
             &response_opts,
+        )?;
+        info!(
+            parent: &request_span,
+            elapsed_ms = started_at.elapsed().as_millis() as u64,
+            "gRPC inference request completed"
         );
 
         Ok(Response::new(pb::GenerateResponse {
@@ -331,19 +335,21 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
                                     logprobs,
                                 },
                             finished,
-                        }) => Ok(pb::GenerateResponse {
+                        }) => convert::to_sequence_output(
+                            &decoded.text,
+                            &token_ids,
+                            logprobs.as_ref(),
+                            finished.as_deref(),
+                            &response_opts,
+                        )
+                        .map(|outputs| pb::GenerateResponse {
                             prompt_info: None,
-                            outputs: Some(convert::to_sequence_output(
-                                &decoded.text,
-                                &token_ids,
-                                logprobs.as_ref(),
-                                finished.as_deref(),
-                                &response_opts,
-                            )),
+                            outputs: Some(outputs),
                         }),
                     };
 
-                    if tx.send(response).await.is_err() {
+                    let failed = response.is_err();
+                    if tx.send(response).await.is_err() || failed {
                         break;
                     }
                 }
diff --git a/rust/src/server/src/grpc/media.rs b/rust/src/server/src/grpc/media.rs
new file mode 100644
index 000000000000..ff746b191e04
--- /dev/null
+++ b/rust/src/server/src/grpc/media.rs
@@ -0,0 +1,341 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Protobuf media adaptation into the shared multimodal preparation boundary.
+
+use std::io::Cursor;
+
+use serde::Deserialize as _;
+use thiserror_ext::AsReport as _;
+use tonic::Status;
+use url::Url;
+use vllm_chat::MediaContentPart;
+use vllm_chat::multimodal::MultimodalInput;
+use vllm_engine_core_client::protocol::multimodal::{
+    MmFeatureSpec, MmKwargsItem, MmModality, PlaceholderRange,
+};
+use vllm_engine_core_client::protocol::tensor::WireTensor;
+
+use super::pb;
+
+/// Maximum total serialized kwargs size accepted in one gRPC request.
+const MAX_INLINE_MM_BYTES: usize = 16 * 1024 * 1024;
+const MAX_MSGPACK_DEPTH: usize = 32;
+
+/// Adapt a request containing either raw media or preprocessed features.
+///
+/// Require a single input kind and bound the total kwargs size before decoding.
+/// Shared media preparation validates features, model support, item limits, and prompt bounds.
+pub(super) fn from_proto(media: Vec<pb::MediaItem>) -> Result<MultimodalInput, Status> {
+    if !media
+        .iter()
+        .any(|item| matches!(item.source, Some(pb::media_item::Source::Features(_))))
+    {
+        return raw_parts_from_proto(media).map(MultimodalInput::Raw);
+    }
+    let mut encoded_bytes = 0usize;
+    for item in &media {
+        let Some(pb::media_item::Source::Features(feature)) = &item.source else {
+            return Err(mixed_media_error());
+        };
+        encoded_bytes = encoded_bytes.saturating_add(feature.kwargs.as_ref().map_or(0, Vec::len));
+    }
+    if encoded_bytes > MAX_INLINE_MM_BYTES {
+        return Err(Status::resource_exhausted(format!(
+            "inline multimodal payload exceeds {MAX_INLINE_MM_BYTES} bytes"
+        )));
+    }
+    let features = media
+        .into_iter()
+        .enumerate()
+        .map(|(index, item)| {
+            feature_from_proto(item).map_err(|error| {
+                Status::new(error.code(), format!("media[{index}]: {}", error.message()))
+            })
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+    Ok(MultimodalInput::Preprocessed(features))
+}
+
+fn feature_from_proto(item: pb::MediaItem) -> Result<MmFeatureSpec, Status> {
+    let modality = match item.modality() {
+        pb::Modality::Image => MmModality::Image,
+        pb::Modality::Video => MmModality::Video,
+        pb::Modality::Audio => MmModality::Audio,
+        pb::Modality::Unspecified => return Err(Status::invalid_argument("modality is required")),
+    };
+    let Some(pb::media_item::Source::Features(feature)) = item.source else {
+        return Err(mixed_media_error());
+    };
+    let bytes = feature.kwargs.ok_or_else(|| {
+        Status::invalid_argument("features.kwargs is required; cache-only features are unsupported")
+    })?;
+    let offset = usize::try_from(feature.offset)
+        .map_err(|_| Status::invalid_argument("features.offset exceeds platform limits"))?;
+    let length = usize::try_from(feature.length)
+        .map_err(|_| Status::invalid_argument("features.length exceeds platform limits"))?;
+    let is_embed = if feature.is_embed.is_empty() {
+        None
+    } else {
+        Some(
+            WireTensor::from_bool(vec![length], feature.is_embed)
+                .map_err(Status::invalid_argument)?,
+        )
+    };
+    Ok(MmFeatureSpec {
+        modality,
+        data: Some(decode_inline_mm_kwargs(&bytes)?),
+        identifier: feature.identifier,
+        mm_hash: feature.mm_hash,
+        mm_position: PlaceholderRange {
+            offset,
+            length,
+            is_embed,
+        },
+    })
+}
+
+/// Decode one inline item using the engine protocol's typed serde schema.
+fn decode_inline_mm_kwargs(bytes: &[u8]) -> Result<MmKwargsItem, Status> {
+    let mut decoder = rmp_serde::Deserializer::new(Cursor::new(bytes));
+    decoder.set_max_depth(MAX_MSGPACK_DEPTH);
+    let kwargs = MmKwargsItem::deserialize(&mut decoder).map_err(|error| {
+        Status::invalid_argument(format!(
+            "invalid multimodal kwargs MessagePack: {}",
+            error.as_report()
+        ))
+    })?;
+    if decoder.position() != bytes.len() as u64 {
+        return Err(Status::invalid_argument(
+            "kwargs contains trailing MessagePack data",
+        ));
+    }
+    Ok(kwargs)
+}
+
+fn mixed_media_error() -> Status {
+    Status::invalid_argument("raw media and preprocessed media features cannot be mixed")
+}
+
+fn raw_parts_from_proto(media: Vec<pb::MediaItem>) -> Result<Vec<MediaContentPart>, Status> {
+    let mut parts = Vec::with_capacity(media.len());
+    for (index, item) in media.into_iter().enumerate() {
+        let modality = item.modality();
+        if modality == pb::Modality::Unspecified {
+            return Err(Status::invalid_argument(format!(
+                "media[{index}].modality is required"
+            )));
+        }
+        let uuid = (!item.uuid.is_empty()).then_some(item.uuid);
+        let mime_type = (!item.mime_type.is_empty()).then_some(item.mime_type);
+        let source = item.source.ok_or_else(|| {
+            Status::invalid_argument(format!("media[{index}].source is required"))
+        })?;
+        match &source {
+            pb::media_item::Source::Url(url) => {
+                validate_media_uri(index, "url", url, &["http", "https"])?;
+            }
+            pb::media_item::Source::DataUri(uri) => {
+                validate_media_uri(index, "data_uri", uri, &["data"])?;
+            }
+            pb::media_item::Source::RawBytes(_) => {}
+            pb::media_item::Source::Features(_) => return Err(mixed_media_error()),
+        }
+        let part = match (modality, source) {
+            (
+                pb::Modality::Image,
+                pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url),
+            ) => MediaContentPart::ImageUrl {
+                url,
+                detail: None,
+                uuid,
+            },
+            (pb::Modality::Image, pb::media_item::Source::RawBytes(data)) => {
+                MediaContentPart::ImageData {
+                    data,
+                    mime_type,
+                    uuid,
+                    detail: None,
+                }
+            }
+            (
+                pb::Modality::Video,
+                pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url),
+            ) => MediaContentPart::VideoUrl { url, uuid },
+            (pb::Modality::Video, pb::media_item::Source::RawBytes(data)) => {
+                MediaContentPart::VideoData {
+                    data,
+                    mime_type,
+                    uuid,
+                }
+            }
+            (
+                pb::Modality::Audio,
+                pb::media_item::Source::Url(url) | pb::media_item::Source::DataUri(url),
+            ) => MediaContentPart::AudioUrl { url, uuid },
+            (pb::Modality::Audio, pb::media_item::Source::RawBytes(data)) => {
+                MediaContentPart::AudioData {
+                    data,
+                    mime_type,
+                    uuid,
+                }
+            }
+            (_, pb::media_item::Source::Features(_)) => return Err(mixed_media_error()),
+            (pb::Modality::Unspecified, _) => unreachable!("modality validated above"),
+        };
+        parts.push(part);
+    }
+    Ok(parts)
+}
+
+fn validate_media_uri(
+    index: usize,
+    field: &str,
+    value: &str,
+    allowed_schemes: &[&str],
+) -> Result<(), Status> {
+    let uri = Url::parse(value).map_err(|_| {
+        Status::invalid_argument(format!("media[{index}].{field} is not a valid URI"))
+    })?;
+    if !allowed_schemes.contains(&uri.scheme()) {
+        return Err(Status::invalid_argument(format!(
+            "media[{index}].{field} must use the {} scheme",
+            allowed_schemes.join(" or ")
+        )));
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeMap;
+
+    use super::*;
+    use vllm_engine_core_client::protocol::multimodal::{
+        MmBatchedField, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, SliceSpec,
+    };
+
+    fn preprocessed() -> pb::MediaItem {
+        let kwargs = BTreeMap::from([(
+            "x",
+            MmFieldElem {
+                data: Some(MmKwargValue::Int(1)),
+                field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+            },
+        )]);
+        pb::MediaItem {
+            modality: pb::Modality::Image as i32,
+            source: Some(pb::media_item::Source::Features(
+                pb::PreprocessedMediaFeatures {
+                    kwargs: Some(rmp_serde::to_vec_named(&kwargs).unwrap()),
+                    identifier: "image-id".to_owned(),
+                    length: 1,
+                    ..Default::default()
+                },
+            )),
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn mixed_media_is_rejected_in_both_orders() {
+        let raw = pb::MediaItem {
+            modality: pb::Modality::Image as i32,
+            source: Some(pb::media_item::Source::RawBytes(vec![1])),
+            ..Default::default()
+        };
+        for media in [vec![raw.clone(), preprocessed()], vec![preprocessed(), raw]] {
+            assert_eq!(
+                from_proto(media).unwrap_err().code(),
+                tonic::Code::InvalidArgument
+            );
+        }
+    }
+
+    #[test]
+    fn protobuf_features_require_modality_inline_kwargs_and_a_matching_mask() {
+        let mut missing_modality = preprocessed();
+        missing_modality.modality = 99;
+        let mut missing_kwargs = preprocessed();
+        let Some(pb::media_item::Source::Features(feature)) = &mut missing_kwargs.source else {
+            unreachable!()
+        };
+        feature.kwargs = None;
+        let mut invalid_mask = preprocessed();
+        let Some(pb::media_item::Source::Features(feature)) = &mut invalid_mask.source else {
+            unreachable!()
+        };
+        feature.is_embed = vec![true, false];
+        for media in [missing_modality, missing_kwargs, invalid_mask] {
+            let error = from_proto(vec![media]).unwrap_err();
+            assert_eq!(error.code(), tonic::Code::InvalidArgument);
+            assert!(error.message().starts_with("media[0]:"));
+        }
+    }
+
+    #[test]
+    fn aggregate_payload_limit_is_checked_before_decoding_items() {
+        let mut item = preprocessed();
+        let Some(pb::media_item::Source::Features(feature)) = &mut item.source else {
+            unreachable!()
+        };
+        feature.kwargs = Some(vec![0; MAX_INLINE_MM_BYTES / 2 + 1]);
+        let error = from_proto(vec![item.clone(), item]).unwrap_err();
+        assert_eq!(error.code(), tonic::Code::ResourceExhausted);
+    }
+
+    #[test]
+    fn kwargs_decoder_rejects_trailing_truncated_and_deeply_nested_input() {
+        let kwargs = BTreeMap::from([(
+            "pixels".to_owned(),
+            MmFieldElem {
+                data: Some(MmKwargValue::Int(1)),
+                field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+            },
+        )]);
+        let valid = rmp_serde::to_vec_named(&kwargs).unwrap();
+        assert_eq!(decode_inline_mm_kwargs(&valid).unwrap(), kwargs);
+        let mut trailing = valid.clone();
+        trailing.push(0);
+        let mut nested_kwargs = kwargs.clone();
+        let mut value = MmKwargValue::Int(1);
+        for _ in 0..64 {
+            value = MmKwargValue::List(vec![value]);
+        }
+        nested_kwargs.get_mut("pixels").unwrap().data = Some(value);
+        let nested = rmp_serde::to_vec_named(&nested_kwargs).unwrap();
+        // A declared array length must not drive a proportional allocation.
+        let huge_array = vec![
+            0x81, 0xa1, b'x', 0x81, 0xa4, b'd', b'a', b't', b'a', 0xdd, 0xff, 0xff, 0xff, 0xff,
+        ];
+        for bytes in [&valid[..valid.len() - 1], &trailing, &nested, &huge_array] {
+            assert!(decode_inline_mm_kwargs(bytes).is_err());
+        }
+    }
+
+    #[test]
+    fn kwargs_decode_python_encoder_tensor_and_factory_tuple() {
+        // MsgpackEncoder(size_threshold=1 << 30), vLLM 6cbb3c154e:
+        // pixel_values = tensor([1., 2.]), MultiModalFlatField(slices=[slice(0, 2)]).
+        let bytes = b"\x81\xacpixel_values\x82\xa4data\x93\xa7float32\x91\x02\xd7\x03\
+            \x00\x00\x80\x3f\x00\x00\x00\x40\
+            \xa5field\x92\xa4flat\x83\xabkeep_on_cpu\xc2\
+            \xa6slices\x91\x93\x00\x02\xc0\xa3dim\x00";
+        let kwargs = decode_inline_mm_kwargs(bytes).unwrap();
+        let expected = MmFieldElem {
+            data: Some(MmKwargValue::Tensor(
+                WireTensor::from_f32(vec![2], vec![1., 2.]).unwrap(),
+            )),
+            field: MmField::Flat(MmFlatField {
+                slices: vec![MmSlice::Slice(SliceSpec {
+                    start: Some(0),
+                    stop: Some(2),
+                    step: None,
+                })],
+                dim: 0,
+                keep_on_cpu: false,
+            }),
+        };
+        assert_eq!(kwargs["pixel_values"], expected);
+    }
+}
diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs
index 06c7c9259eb8..c61c1e5e1419 100644
--- a/rust/src/server/src/grpc/mod.rs
+++ b/rust/src/server/src/grpc/mod.rs
@@ -7,11 +7,10 @@ mod control;
 mod convert;
 mod health;
 mod inference;
+mod media;
 
 /// Generated protobuf/gRPC types for the `vllm` package.
-pub mod pb {
-    tonic::include_proto!("vllm");
-}
+pub use vllm_proto as pb;
 
 pub(crate) use control::ControlGrpcService;
 pub use control::ControlServiceImpl;
diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs
index ea4364480cde..18b6e6eb2990 100644
--- a/rust/src/server/src/grpc/tests.rs
+++ b/rust/src/server/src/grpc/tests.rs
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+use std::collections::BTreeMap;
 use std::fs;
 use std::future::Future;
 use std::io;
@@ -23,21 +24,27 @@ use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatu
 use tonic_health::pb::health_client::HealthClient;
 use tonic_health::server::health_reporter;
 use tower::service_fn;
+use vllm_chat::multimodal::{MmLimitPerPrompt, MmLimitSpec, MmModality};
 use vllm_chat::{
     ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
     DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
 };
 use vllm_engine_core_client::mock_engine::{
     DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS,
-    default_ready_response,
+    MockEngineConfig, connect_to_bootstrapped_frontend, default_ready_response,
 };
 use vllm_engine_core_client::protocol::decode_value;
+use vllm_engine_core_client::protocol::dtype::TensorDtype;
 use vllm_engine_core_client::protocol::handshake::{EngineCoreReadyResponse, KvEventsConfig};
+use vllm_engine_core_client::protocol::multimodal::{
+    MmBatchedField, MmField, MmFieldElem, MmKwargValue,
+};
 use vllm_engine_core_client::protocol::output::{
     EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs,
     UtilityCallOutput,
 };
 use vllm_engine_core_client::protocol::request::EngineCoreRequest;
+use vllm_engine_core_client::protocol::tensor::WireTensor;
 use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task_with_ready};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId, TransportMode};
@@ -237,6 +244,7 @@ impl ChatBackend for FakeTextBackend {
             self.tokenizer(),
             options.tool_call_parser,
             options.reasoning_parser,
+            options.tool_strict_level,
         )?))
     }
 }
@@ -245,6 +253,7 @@ impl ChatRenderer for FakeTextBackend {
     fn render(&self, _request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
         Ok(RenderedPrompt {
             prompt: Prompt::Text(String::new()),
+            media_order: None,
             effective_template_kwargs: Default::default(),
         })
     }
@@ -284,6 +293,7 @@ impl ChatBackend for FakeMultimodalBackend {
             self.tokenizer(),
             options.tool_call_parser,
             options.reasoning_parser,
+            options.tool_strict_level,
         )?))
     }
 }
@@ -292,6 +302,12 @@ const QWEN_IMAGE_TOKEN_ID: u32 = 151655;
 const TINY_PNG_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=";
 
 fn multimodal_backend() -> Arc<dyn ChatTextBackend> {
+    multimodal_backend_with_limits(Default::default())
+}
+
+fn multimodal_backend_with_limits(
+    limit_mm_per_prompt: MmLimitPerPrompt,
+) -> Arc<dyn ChatTextBackend> {
     let config_path = std::env::temp_dir().join(format!(
         "vllm-grpc-qwen-config-{}.json",
         uuid::Uuid::new_v4()
@@ -309,7 +325,7 @@ fn multimodal_backend() -> Arc<dyn ChatTextBackend> {
             ..Default::default()
         },
         Arc::new(TestTokenizer::new().with_regular_token("<|image_pad|>", QWEN_IMAGE_TOKEN_ID)),
-        Default::default(),
+        limit_mm_per_prompt,
     )
     .expect("load multimodal info")
     .expect("qwen multimodal info is registered");
@@ -680,6 +696,65 @@ async fn unary_generate_returns_collected_text() {
     server_task.abort();
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn engine_error_finish_returns_internal_for_unary_and_streaming() {
+    for streaming in [false, true] {
+        for prefix in [vec![], vec![b'h' as u32]] {
+            let mut output_specs = Vec::new();
+            if !prefix.is_empty() {
+                output_specs.push((prefix.clone(), None));
+            }
+            output_specs.push((vec![], Some(EngineCoreFinishReason::Error)));
+            let (mut client, server_task, engine_task) =
+                grpc_test_server(b"engine-grpc-error", output_specs).await;
+            let request = pb::GenerateRequest {
+                request_id: "test-engine-error".to_string(),
+                model: "test-model".to_string(),
+                prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
+                stopping: Some(pb::StoppingCriteria {
+                    max_new_tokens: 10,
+                    ..Default::default()
+                }),
+                response: Some(pb::ResponseOptions {
+                    output_token_ids: true,
+                    ..Default::default()
+                }),
+                ..Default::default()
+            };
+
+            let status = if streaming {
+                let mut stream =
+                    client.generate_stream(request).await.expect("start stream").into_inner();
+                let mut received_tokens = Vec::new();
+                let status = loop {
+                    match stream.message().await {
+                        Ok(Some(response)) => {
+                            if let Some(output) = response.outputs {
+                                assert!(
+                                    output.finish_info.is_none(),
+                                    "error became a finish event"
+                                );
+                                received_tokens.extend(output.token_ids);
+                            }
+                        }
+                        Ok(None) => panic!("engine error ended as a successful stream"),
+                        Err(status) => break status,
+                    }
+                };
+                assert_eq!(received_tokens, prefix);
+                status
+            } else {
+                client.generate(request).await.expect_err("engine error must fail the RPC")
+            };
+            assert_eq!(status.code(), tonic::Code::Internal);
+
+            engine_task.await.expect("mock engine task");
+            server_task.abort();
+        }
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn unary_generate_with_token_ids_prompt() {
@@ -818,6 +893,249 @@ async fn unary_generate_prepares_multimodal_input_for_engine_core() {
     server_task.abort();
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn generate_forwards_preprocessed_multimodal_features() {
+    for stream in [false, true] {
+        let kwargs = BTreeMap::from([(
+            "pixel_values".to_string(),
+            MmFieldElem {
+                data: Some(MmKwargValue::Tensor(WireTensor::from_raw(
+                    TensorDtype::U8,
+                    vec![3],
+                    vec![1, 2, 3],
+                ))),
+                field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+            },
+        )]);
+        let encoded_kwargs = rmp_serde::to_vec_named(&kwargs).expect("encode multimodal kwargs");
+        let expected_kwargs = kwargs.clone();
+        let (inference_service, control_service, engine_health, engine_task) =
+            setup_grpc_service_with_backend(
+                b"engine-grpc-preprocessed-multimodal",
+                default_stream_output_specs(),
+                multimodal_backend(),
+                move |request| {
+                    assert_eq!(
+                        request.prompt_token_ids.as_deref(),
+                        Some(&[11, 99, 99, 12][..])
+                    );
+                    let features = request.mm_features.as_ref().expect("multimodal features");
+                    assert_eq!(features.len(), 1);
+                    let feature = &features[0];
+                    assert_eq!(feature.modality.as_str(), "image");
+                    assert_eq!(feature.identifier, "image-hash-a");
+                    assert_eq!(feature.mm_hash.as_deref(), Some("image-hash-a"));
+                    assert_eq!(feature.mm_position.offset, 1);
+                    assert_eq!(feature.mm_position.length, 2);
+                    assert_eq!(feature.data.as_ref(), Some(&expected_kwargs));
+                    assert_eq!(
+                        feature.mm_position.is_embed,
+                        Some(WireTensor::from_bool(vec![2], vec![true, false]).unwrap())
+                    );
+                },
+            )
+            .await;
+        let (channel, server_task) = start_grpc_test_server(
+            inference_service,
+            control_service,
+            engine_health,
+            tokio_util::sync::CancellationToken::new(),
+        )
+        .await;
+        let mut client = InferenceClient::new(channel);
+
+        let request = pb::GenerateRequest {
+            request_id: "test-preprocessed-multimodal".to_string(),
+            model: "test-model".to_string(),
+            prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds {
+                ids: vec![11, 99, 99, 12],
+            })),
+            media: vec![pb::MediaItem {
+                modality: pb::Modality::Image as i32,
+                source: Some(pb::media_item::Source::Features(
+                    pb::PreprocessedMediaFeatures {
+                        kwargs: Some(encoded_kwargs),
+                        identifier: "image-hash-a".to_string(),
+                        offset: 1,
+                        length: 2,
+                        mm_hash: Some("image-hash-a".to_string()),
+                        is_embed: vec![true, false],
+                    },
+                )),
+                mime_type: String::new(),
+                uuid: String::new(),
+            }],
+            stopping: Some(pb::StoppingCriteria {
+                max_new_tokens: 10,
+                ..Default::default()
+            }),
+            ..Default::default()
+        };
+        if stream {
+            let mut response =
+                client.generate_stream(request).await.expect("preprocessed stream").into_inner();
+            while response.message().await.expect("stream message").is_some() {}
+        } else {
+            client.generate(request).await.expect("preprocessed generate");
+        }
+
+        engine_task.await.expect("mock engine task");
+        server_task.abort();
+    }
+}
+
+fn preprocessed_grpc_media(
+    modality: pb::Modality,
+    identifier: &str,
+    offset: u64,
+    length: u64,
+) -> pb::MediaItem {
+    let kwargs = BTreeMap::from([(
+        "pixel_values".to_string(),
+        MmFieldElem {
+            data: Some(MmKwargValue::Tensor(WireTensor::from_raw(
+                TensorDtype::U8,
+                vec![1],
+                vec![1],
+            ))),
+            field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+        },
+    )]);
+    pb::MediaItem {
+        modality: modality as i32,
+        source: Some(pb::media_item::Source::Features(
+            pb::PreprocessedMediaFeatures {
+                kwargs: Some(rmp_serde::to_vec_named(&kwargs).expect("encode kwargs")),
+                identifier: identifier.to_string(),
+                offset,
+                length,
+                mm_hash: Some(identifier.to_string()),
+                is_embed: Vec::new(),
+            },
+        )),
+        mime_type: String::new(),
+        uuid: String::new(),
+    }
+}
+
+fn preprocessed_grpc_request(media: Vec<pb::MediaItem>) -> pb::GenerateRequest {
+    pb::GenerateRequest {
+        request_id: "test-preprocessed-validation".to_string(),
+        model: "test-model".to_string(),
+        prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds {
+            ids: vec![11, 99, 99, 12],
+        })),
+        media,
+        stopping: Some(pb::StoppingCriteria {
+            max_new_tokens: 10,
+            ..Default::default()
+        }),
+        ..Default::default()
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn unary_generate_rejects_preprocessed_features_for_text_only_model() {
+    let (inference_service, control_service, engine_health, _engine_task) =
+        setup_grpc_service_with_backend(
+            b"engine-grpc-preprocessed-text-only",
+            default_stream_output_specs(),
+            Arc::new(FakeTextBackend),
+            |_| {},
+        )
+        .await;
+    let (channel, server_task) = start_grpc_test_server(
+        inference_service,
+        control_service,
+        engine_health,
+        tokio_util::sync::CancellationToken::new(),
+    )
+    .await;
+    let mut client = InferenceClient::new(channel);
+
+    let status = client
+        .generate(preprocessed_grpc_request(vec![preprocessed_grpc_media(
+            pb::Modality::Image,
+            "image-a",
+            1,
+            2,
+        )]))
+        .await
+        .expect_err("text-only models must reject preprocessed features");
+
+    assert_eq!(status.code(), tonic::Code::InvalidArgument);
+    server_task.abort();
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn generate_validates_preprocessed_model_limits_and_prompt() {
+    let limits = MmLimitPerPrompt::from([(MmModality::Image, MmLimitSpec::Count(1))]);
+    let (inference_service, control_service, engine_health, _engine_task) =
+        setup_grpc_service_with_backend(
+            b"engine-grpc-preprocessed-limit",
+            default_stream_output_specs(),
+            multimodal_backend_with_limits(limits),
+            |_| {},
+        )
+        .await;
+    let (channel, server_task) = start_grpc_test_server(
+        inference_service,
+        control_service,
+        engine_health,
+        tokio_util::sync::CancellationToken::new(),
+    )
+    .await;
+    let mut client = InferenceClient::new(channel);
+
+    let over_limit = preprocessed_grpc_request(vec![
+        preprocessed_grpc_media(pb::Modality::Image, "image-a", 0, 1),
+        preprocessed_grpc_media(pb::Modality::Image, "image-b", 2, 1),
+    ]);
+    let bad_range = preprocessed_grpc_request(vec![preprocessed_grpc_media(
+        pb::Modality::Image,
+        "image",
+        3,
+        2,
+    )]);
+    let unsupported = preprocessed_grpc_request(vec![preprocessed_grpc_media(
+        pb::Modality::Audio,
+        "audio",
+        0,
+        1,
+    )]);
+    let mut text = preprocessed_grpc_request(vec![preprocessed_grpc_media(
+        pb::Modality::Image,
+        "image",
+        0,
+        1,
+    )]);
+    let mut truncated = text.clone();
+    text.prompt = Some(pb::generate_request::Prompt::Text("hello".to_owned()));
+    truncated.truncate_prompt_tokens = 1;
+    let empty_identifier =
+        preprocessed_grpc_request(vec![preprocessed_grpc_media(pb::Modality::Image, "", 0, 1)]);
+    for request in [
+        over_limit,
+        bad_range,
+        unsupported,
+        text,
+        truncated,
+        empty_identifier,
+    ] {
+        let status = client.generate(request.clone()).await.expect_err("invalid multimodal input");
+        assert_eq!(status.code(), tonic::Code::InvalidArgument);
+        let status = client
+            .generate_stream(request)
+            .await
+            .expect_err("invalid streaming multimodal input");
+        assert_eq!(status.code(), tonic::Code::InvalidArgument);
+    }
+    server_task.abort();
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn streaming_generate_rejects_text_prompt_with_media() {
@@ -1595,8 +1913,16 @@ async fn control_abort_resolves_external_id_and_empty_is_noop() {
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn control_reports_server_and_model_info() {
+    let mut ready = default_ready_response();
+    ready.effective_attention_block_size = Some(64);
     let (generate_service, control_service, engine_health, _engine_task) =
-        setup_grpc_service(b"engine-grpc-info", default_stream_output_specs()).await;
+        setup_grpc_service_with_engine_script(
+            b"engine-grpc-info".to_vec(),
+            ready,
+            Arc::new(FakeTextBackend),
+            |_, _| boxed_test_future(async {}),
+        )
+        .await;
     let (channel, server_task) = start_grpc_test_server(
         generate_service,
         control_service,
@@ -1612,6 +1938,7 @@ async fn control_reports_server_and_model_info() {
         .expect("get server info")
         .into_inner();
     assert_eq!(server.engine_version, "test-vllm-version");
+    assert_eq!(server.effective_attention_block_size, Some(64));
     assert_eq!(server.api_version, "vllm");
     assert_eq!(server.instance_id, "test-instance");
     assert_eq!(server.max_model_len, DEFAULT_MOCK_MAX_MODEL_LEN as u32);
@@ -1623,6 +1950,7 @@ async fn control_reports_server_and_model_info() {
     assert_eq!(parallelism.tensor_parallel_size, 1);
     assert_eq!(parallelism.pipeline_parallel_size, 1);
     assert_eq!(parallelism.data_parallel_size, 1);
+    assert_eq!(parallelism.data_parallel_size_local, 1);
     assert_eq!(parallelism.data_parallel_rank, 0);
     assert_eq!(parallelism.decode_context_parallel_size, 1);
     let rl = server.rl_capabilities.expect("RL capabilities");
@@ -1866,82 +2194,98 @@ async fn control_forwards_weight_update_without_pause_guard() {
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn control_aggregates_multi_engine_capacity() {
-    let ipc = IpcNamespace::new().expect("create ipc namespace");
-    let handshake_address = ipc.handshake_endpoint();
-
-    let mut ready_0 = default_ready_response();
-    ready_0.max_model_len = 8_192;
-    ready_0.num_gpu_blocks = 10;
-    ready_0.effective_data_parallel_size = 2;
-    ready_0.tensor_parallel_size = 2;
-    ready_0.pipeline_parallel_size = 3;
-    ready_0.world_size = 12;
-    ready_0.weight_transfer_backend = Some("nccl".to_string());
-    ready_0.enable_sleep_mode = true;
-    ready_0.supports_draft_weight_updates = true;
-
-    let mut ready_1 = default_ready_response();
-    ready_1.max_model_len = 4_096;
-    ready_1.num_gpu_blocks = 20;
-    ready_1.effective_data_parallel_size = 2;
-    ready_1.tensor_parallel_size = 2;
-    ready_1.pipeline_parallel_size = 3;
-    ready_1.world_size = 12;
-    ready_1.data_parallel_rank = 1;
-
-    let engine_tasks = [ready_0, ready_1].map(|ready| {
-        let engine_id = EngineId::from_engine_index(
-            ready.data_parallel_rank.try_into().expect("test rank fits engine identity"),
+    // Local capacity and ownership must not expand to the deployment-wide DP size.
+    for (start_rank, global_size) in [(0, 2), (0, 8), (4, 8)] {
+        let ipc = IpcNamespace::new().expect("create ipc namespace");
+        let input_address = ipc.input_endpoint();
+        let output_address = ipc.output_endpoint();
+
+        let mut ready_0 = default_ready_response();
+        ready_0.max_model_len = 8_192;
+        ready_0.num_gpu_blocks = 10;
+        ready_0.effective_data_parallel_size = global_size as u64;
+        ready_0.data_parallel_rank = start_rank;
+        ready_0.tensor_parallel_size = 2;
+        ready_0.pipeline_parallel_size = 3;
+        ready_0.world_size = 12;
+        ready_0.weight_transfer_backend = Some("nccl".to_string());
+        ready_0.enable_sleep_mode = true;
+        ready_0.supports_draft_weight_updates = true;
+        ready_0.effective_attention_block_size = Some(64);
+
+        let mut ready_1 = default_ready_response();
+        ready_1.max_model_len = 4_096;
+        ready_1.num_gpu_blocks = 20;
+        ready_1.effective_data_parallel_size = global_size as u64;
+        ready_1.tensor_parallel_size = 2;
+        ready_1.pipeline_parallel_size = 3;
+        ready_1.world_size = 12;
+        ready_1.data_parallel_rank = start_rank + 1;
+
+        let client_config = EngineCoreClientConfig {
+            transport_mode: TransportMode::Bootstrapped {
+                input_address: input_address.clone(),
+                output_address: output_address.clone(),
+                engine_start_index: start_rank,
+                engine_count: 2,
+                data_parallel_size: global_size,
+                ready_timeout: Duration::from_secs(2),
+            },
+            coordinator_mode: None,
+            model_name: "test-model".to_string(),
+            client_index: 0,
+        };
+        let client_task = tokio::spawn(EngineCoreClient::connect(client_config));
+        let mut engine_sockets = Vec::new();
+        for ready in [ready_0, ready_1] {
+            let engine_id = EngineId::from_engine_index(
+                ready.data_parallel_rank.try_into().expect("test rank fits engine identity"),
+            );
+            engine_sockets.push(
+                connect_to_bootstrapped_frontend(
+                    &input_address,
+                    &output_address,
+                    engine_id,
+                    MockEngineConfig {
+                        ready_response: ready,
+                        ..Default::default()
+                    },
+                )
+                .await
+                .expect("connect mock engine"),
+            );
+        }
+        let client = client_task.await.unwrap().expect("connect multi-engine client");
+        let chat = ChatLlm::from_shared_backend(
+            Llm::new(client),
+            Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
         );
-        MockEngineTask::new(spawn_mock_engine_task_with_ready(
-            handshake_address.clone(),
-            engine_id,
-            ready,
-            |_, _| boxed_test_future(async {}),
-        ))
-    });
+        let state = AppState::new(vec!["test-model".to_string()], chat);
+        let service = ControlServiceImpl::new(Arc::new(state));
 
-    let client = EngineCoreClient::connect(EngineCoreClientConfig {
-        transport_mode: TransportMode::HandshakeOwner {
-            handshake_address,
-            advertised_host: "127.0.0.1".to_string(),
-            engine_count: 2,
-            ready_timeout: Duration::from_secs(2),
-            local_input_address: Some(ipc.input_endpoint()),
-            local_output_address: Some(ipc.output_endpoint()),
-        },
-        coordinator_mode: None,
-        model_name: "test-model".to_string(),
-        client_index: 0,
-    })
-    .await
-    .expect("connect multi-engine client");
-    let chat = ChatLlm::from_shared_backend(
-        Llm::new(client),
-        Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
-    );
-    let state = AppState::new(vec!["test-model".to_string()], chat);
-    let service = ControlServiceImpl::new(Arc::new(state));
-
-    let server = pb::control_server::Control::get_server_info(
-        &service,
-        tonic::Request::new(pb::GetServerInfoRequest {}),
-    )
-    .await
-    .expect("get server info")
-    .into_inner();
-    assert_eq!(server.max_model_len, 4_096);
-    assert_eq!(server.total_kv_blocks, 30);
-    let rl = server.rl_capabilities.expect("RL capabilities");
-    assert!(!rl.weight_transfer_enabled);
-    assert!(rl.weight_transfer_backend.is_empty());
-    assert!(!rl.sleep_mode_enabled);
-    assert!(!rl.draft_weight_updates_enabled);
-    let parallelism = server.parallelism.unwrap();
-    assert_eq!(parallelism.data_parallel_size, 2);
-    assert_eq!(parallelism.world_size, 12);
-
-    drop(engine_tasks);
+        let server = pb::control_server::Control::get_server_info(
+            &service,
+            tonic::Request::new(pb::GetServerInfoRequest {}),
+        )
+        .await
+        .expect("get server info")
+        .into_inner();
+        assert_eq!(server.max_model_len, 4_096);
+        assert_eq!(server.total_kv_blocks, 30);
+        assert!(server.effective_attention_block_size.is_none());
+        let rl = server.rl_capabilities.expect("RL capabilities");
+        assert!(!rl.weight_transfer_enabled);
+        assert!(rl.weight_transfer_backend.is_empty());
+        assert!(!rl.sleep_mode_enabled);
+        assert!(!rl.draft_weight_updates_enabled);
+        let parallelism = server.parallelism.unwrap();
+        assert_eq!(parallelism.data_parallel_size, global_size as u32);
+        assert_eq!(parallelism.data_parallel_rank, start_rank);
+        assert_eq!(parallelism.data_parallel_size_local, 2);
+        assert_eq!(parallelism.world_size, 12);
+
+        drop(engine_sockets);
+    }
 }
 
 #[test]
diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs
index 098c80380557..91564db73c5b 100644
--- a/rust/src/server/src/lib.rs
+++ b/rust/src/server/src/lib.rs
@@ -98,6 +98,8 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
     let loaded = load_model_backends(
         &config.model,
         LoadModelBackendsOptions {
+            revision: config.revision.clone(),
+            hf_overrides: config.hf_overrides.clone(),
             generation_config: config.generation_config,
             renderer: config.renderer,
             language_model_only: config.language_model_only,
@@ -137,7 +139,8 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
 
     let chat = ChatLlm::new(text, chat_backend)
         .with_tool_call_parser(config.tool_call_parser.clone())
-        .with_reasoning_parser(config.reasoning_parser.clone());
+        .with_reasoning_parser(config.reasoning_parser.clone())
+        .with_tool_strict_level(config.tool_strict_level);
 
     let state = Arc::new(
         AppState::new(served_model_names, chat)
diff --git a/rust/src/server/src/middleware/metrics.rs b/rust/src/server/src/middleware/metrics.rs
index 1e6a79a556e1..de07620cf75b 100644
--- a/rust/src/server/src/middleware/metrics.rs
+++ b/rust/src/server/src/middleware/metrics.rs
@@ -26,6 +26,7 @@ const EXCLUDED_HANDLERS: &[&str] = &[
     "/reset_encoder_cache",
     "/collective_rpc",
     "/sleep",
+    "/release_kv_cache_memory",
     "/wake_up",
     "/is_sleeping",
 ];
diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs
index 6942694afbf3..50485e26a408 100644
--- a/rust/src/server/src/render.rs
+++ b/rust/src/server/src/render.rs
@@ -10,9 +10,10 @@ use tokio_util::sync::CancellationToken;
 use tracing::info;
 use vllm_chat::{
     ChatRequestProcessor, ChatTemplateContentFormatOption, LoadModelBackendsOptions,
-    ParserSelection, RendererSelection, load_model_backends,
+    ParserSelection, RendererSelection, ToolStrictLevel, load_model_backends,
 };
 use vllm_text::TextRequestProcessor;
+use vllm_text::backend::hf::HfOverrides;
 
 use crate::{
     HttpListenerMode, TlsConfig,
@@ -24,11 +25,14 @@ use crate::{
 #[derive(Debug)]
 pub struct RenderConfig {
     pub model: String,
+    pub revision: Option<String>,
+    pub hf_overrides: HfOverrides,
     pub served_model_name: Vec<String>,
     pub host: String,
     pub port: u16,
     pub tool_call_parser: ParserSelection,
     pub reasoning_parser: ParserSelection,
+    pub tool_strict_level: ToolStrictLevel,
     pub renderer: RendererSelection,
     pub chat_template: Option<String>,
     pub default_chat_template_kwargs: HashMap<String, Value>,
@@ -67,6 +71,8 @@ async fn build_state(config: &RenderConfig) -> Result<Arc<RenderState>> {
     let loaded = load_model_backends(
         &config.model,
         LoadModelBackendsOptions {
+            revision: config.revision.clone(),
+            hf_overrides: config.hf_overrides.clone(),
             generation_config: Default::default(),
             renderer: config.renderer,
             language_model_only: true,
@@ -83,10 +89,12 @@ async fn build_state(config: &RenderConfig) -> Result<Arc<RenderState>> {
     let max_model_len = config.max_model_len.unwrap_or(u32::MAX);
     let text = TextRequestProcessor::new(loaded.text_backend, max_model_len)
         .with_max_logprobs(config.max_logprobs);
-    let chat = ChatRequestProcessor::render_only(loaded.chat_backend).with_parser_selections(
-        config.tool_call_parser.clone(),
-        config.reasoning_parser.clone(),
-    );
+    let chat = ChatRequestProcessor::render_only(loaded.chat_backend)
+        .with_parser_selections(
+            config.tool_call_parser.clone(),
+            config.reasoning_parser.clone(),
+        )
+        .with_tool_strict_level(config.tool_strict_level);
     Ok(Arc::new(RenderState {
         model: config.model.clone(),
         served_model_names,
@@ -149,7 +157,10 @@ mod tests {
         shutdown.cancel();
         let error = serve_render(
             RenderConfig {
+                tool_strict_level: ToolStrictLevel::Auto,
                 model: "test-model".to_string(),
+                revision: None,
+                hf_overrides: Default::default(),
                 served_model_name: Vec::new(),
                 host: "127.0.0.1".to_string(),
                 port: 8000,
diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs
index 393a08450fbc..4611d0c6d6ce 100644
--- a/rust/src/server/src/routes.rs
+++ b/rust/src/server/src/routes.rs
@@ -6,6 +6,7 @@ mod cache;
 mod collective_rpc;
 mod health;
 mod inference;
+mod kv_event_sources;
 mod load;
 mod lora;
 mod metrics;
@@ -17,6 +18,7 @@ mod server_info;
 mod sleep;
 mod tokenize;
 mod version;
+mod weight_transfer;
 mod world_size;
 
 use std::sync::Arc;
@@ -45,33 +47,14 @@ fn runtime_lora_updating_enabled() -> bool {
         .is_some_and(|value| matches!(value.trim().to_lowercase().as_str(), "1" | "true"))
 }
 
-fn parse_scale_out_endpoints_flag(value: Option<&str>) -> Result<bool, String> {
-    let Some(value) = value else {
-        return Ok(false);
-    };
-    if value.is_empty() {
-        return Ok(false);
-    }
-
-    match value.trim() {
-        "0" => Ok(false),
-        "1" => Ok(true),
-        _ => Err("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1".to_string()),
-    }
-}
-
-fn scale_out_endpoints_enabled() -> bool {
-    let value = std::env::var("VLLM_ENABLE_SCALE_OUT_ENDPOINTS").ok();
-    parse_scale_out_endpoints_flag(value.as_deref()).unwrap_or_else(|message| panic!("{message}"))
-}
-
 /// Build the minimal OpenAI-compatible router for one configured model.
 pub fn build_router(state: Arc<AppState>) -> Router {
+    let scale_out_endpoints_enabled = state.api_server_options.enable_scale_out;
     build_router_with_options(
         state,
         server_dev_mode_enabled(),
         runtime_lora_updating_enabled(),
-        scale_out_endpoints_enabled(),
+        scale_out_endpoints_enabled,
     )
 }
 
@@ -109,6 +92,7 @@ fn build_router_with_options(
         .route("/metrics", get(metrics::scrape))
         .route("/load", get(load::load))
         .route("/version", get(version::version))
+        .route("/kv_event_sources", get(kv_event_sources::kv_event_sources))
         // OpenAI-compatible endpoints
         .route("/v1/models", get(openai::list_models))
         .route("/v1/completions", post(openai::completions))
@@ -120,10 +104,7 @@ fn build_router_with_options(
     if scale_out_endpoints_enabled {
         router = router.route("/inference/v1/generate", post(inference::generate));
     } else {
-        info!(
-            "scale-out endpoints are disabled; set \
-             VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 to enable them"
-        );
+        info!("scale-out endpoints are disabled; pass --enable-scale-out to enable them");
     }
 
     if runtime_lora_updating_enabled {
@@ -139,8 +120,34 @@ fn build_router_with_options(
             .route("/reset_mm_cache", post(cache::reset_mm_cache))
             .route("/reset_encoder_cache", post(cache::reset_encoder_cache))
             .route("/collective_rpc", post(collective_rpc::collective_rpc))
+            .route(
+                "/init_weight_transfer_engine",
+                post(weight_transfer::init_weight_transfer_engine),
+            )
+            .route(
+                "/start_weight_update",
+                post(weight_transfer::start_weight_update),
+            )
+            .route(
+                "/start_draft_weight_update",
+                post(weight_transfer::start_draft_weight_update),
+            )
+            .route("/update_weights", post(weight_transfer::update_weights))
+            .route(
+                "/finish_weight_update",
+                post(weight_transfer::finish_weight_update),
+            )
+            .route(
+                "/update_weight_version",
+                post(weight_transfer::update_weight_version),
+            )
+            .route("/weight_info", get(weight_transfer::weight_info))
             .route("/abort_requests", post(abort_requests::abort_requests))
             .route("/sleep", post(sleep::sleep))
+            .route(
+                "/release_kv_cache_memory",
+                post(sleep::release_kv_cache_memory),
+            )
             .route("/wake_up", post(sleep::wake_up))
             .route("/is_sleeping", get(sleep::is_sleeping))
             .route("/pause", post(pause::pause))
diff --git a/rust/src/server/src/routes/abort_requests.rs b/rust/src/server/src/routes/abort_requests.rs
index 298770316099..d08e5af3fa31 100644
--- a/rust/src/server/src/routes/abort_requests.rs
+++ b/rust/src/server/src/routes/abort_requests.rs
@@ -22,7 +22,7 @@ pub async fn abort_requests(
     State(state): State<Arc<AppState>>,
     body: Result<Json<AbortRequestsRequest>, JsonRejection>,
 ) -> Result<StatusCode, ApiError> {
-    let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
+    let Json(body) = body?;
     // Empty/missing `request_ids` aborts all in-flight requests.
     let request_ids = body.request_ids.unwrap_or_default();
 
diff --git a/rust/src/server/src/routes/collective_rpc.rs b/rust/src/server/src/routes/collective_rpc.rs
index bd769bee43a9..fa3782cbe9bf 100644
--- a/rust/src/server/src/routes/collective_rpc.rs
+++ b/rust/src/server/src/routes/collective_rpc.rs
@@ -36,7 +36,7 @@ pub async fn collective_rpc(
     State(state): State<Arc<AppState>>,
     body: Result<Json<CollectiveRpcRequest>, JsonRejection>,
 ) -> Result<Json<CollectiveRpcResponse>, ApiError> {
-    let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
+    let Json(body) = body?;
     let method = body.method.ok_or_else(|| {
         ApiError::invalid_request(
             "Missing 'method' in request body".to_string(),
diff --git a/rust/src/server/src/routes/http_client_tests.rs b/rust/src/server/src/routes/http_client_tests.rs
index 216670b695b6..762d974b9b6b 100644
--- a/rust/src/server/src/routes/http_client_tests.rs
+++ b/rust/src/server/src/routes/http_client_tests.rs
@@ -168,6 +168,7 @@ impl ChatBackend for FakeChatBackend {
             self.tokenizer(),
             options.tool_call_parser,
             options.reasoning_parser,
+            options.tool_strict_level,
         )?))
     }
 }
@@ -186,6 +187,7 @@ impl ChatRenderer for FakeChatBackend {
         }
         Ok(RenderedPrompt {
             prompt: Prompt::Text(prompt),
+            media_order: None,
             effective_template_kwargs: Default::default(),
         })
     }
diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs
index 308ec3ca3f2e..79b68622cfbd 100644
--- a/rust/src/server/src/routes/inference/generate.rs
+++ b/rust/src/server/src/routes/inference/generate.rs
@@ -14,13 +14,14 @@ use asynk_strim_attr::{TryYielder, try_stream};
 use axum::Json;
 use axum::extract::State;
 use axum::http::HeaderMap;
-use axum::response::sse::{Event, Sse};
+use axum::response::sse::Event;
 use axum::response::{IntoResponse, Response};
 use futures::{Stream, StreamExt as _, pin_mut};
 use thiserror_ext::AsReport as _;
 use tracing::{error, info, trace};
 use tracing_futures::Instrument as _;
 use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
+use vllm_engine_core_client::protocol::multimodal::MmFeatureSpec;
 use vllm_llm::{
     CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _, TokenUsage,
 };
@@ -28,7 +29,8 @@ use vllm_llm::{
 use self::convert::{ResponseOptions, prepare_generate_request};
 use self::types::{
     GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice,
-    GenerateStreamResponse,
+    GenerateStreamResponse, MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo,
+    SpeculativeDecodingMetrics, StreamingSpeculativeDecodingMetrics,
 };
 pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams};
 pub(crate) use self::validate::validate_request_compat;
@@ -38,7 +40,7 @@ use crate::routes::openai::utils::logprobs::clamp_logprob;
 use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
 use crate::routes::openai::utils::validated_json::ValidatedJson;
 use crate::state::AppState;
-use crate::utils::resolve_request_context;
+use crate::utils::{resolve_request_context, sse_response};
 
 /// Validate one token-in/token-out request and proxy it into the shared
 /// `vllm-text` stack.
@@ -51,7 +53,12 @@ pub async fn generate(
     let lora_resolution = state.resolve_model_with_loras(body.model.as_deref()).await;
 
     let mm_features = if let Some(parts) = body.content_parts.take() {
-        match state.chat.prepare_media(parts, &mut body.token_ids).await {
+        match state
+            .chat
+            .prepare_media(parts, &mut body.token_ids)
+            .instrument(vllm_chat::mm_request_span(&request_context.request_id))
+            .await
+        {
             Ok(features) => features,
             Err(e) => {
                 return ApiError::invalid_request(
@@ -64,6 +71,11 @@ pub async fn generate(
     } else {
         None
     };
+    let mm_placeholders = body
+        .return_token_ids
+        .unwrap_or(false)
+        .then(|| extract_mm_placeholders(mm_features.as_deref()))
+        .flatten();
 
     let prepared =
         match prepare_generate_request(body, &lora_resolution, request_context, mm_features) {
@@ -98,10 +110,11 @@ pub async fn generate(
             prepared.request_id,
             api_server_options,
             prepared.options,
+            mm_placeholders,
         );
         let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
 
-        return Sse::new(sse_stream).into_response();
+        return sse_response(sse_stream, api_server_options.sse_keep_alive_interval);
     }
 
     let collected = match raw_stream.collect_output().instrument(request_span.clone()).await {
@@ -120,6 +133,7 @@ pub async fn generate(
         prepared.request_id,
         api_server_options,
         prepared.options,
+        mm_placeholders,
     ) {
         Ok(response) => response,
         Err(error) => return error.into_response(),
@@ -143,19 +157,30 @@ async fn generate_chunk_stream(
         include_logprobs,
         // Ignored: raw generate streaming has no prompt-logprobs wire shape.
         include_prompt_logprobs: _,
+        return_token_ids,
     }: ResponseOptions,
+    mut mm_placeholders: Option<MultiModalPlaceholders>,
     mut y: TryYielder<GenerateStreamResponse, ApiError>,
 ) -> Result<(), ApiError> {
     pin_mut!(stream);
     let mut prompt_tokens = None;
+    let mut prompt_token_ids = None;
     let mut usage = TokenUsage::default();
+    let mut spec_decode_metrics = None;
 
     while let Some(next) = stream.next().await {
         match next {
             Ok(output) => {
-                if prompt_tokens.is_none() {
-                    prompt_tokens =
-                        output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len());
+                if let Some(metrics) = output.spec_decode_metrics {
+                    spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(metrics));
+                }
+                if prompt_tokens.is_none()
+                    && let Some(info) = output.prompt_info.as_ref()
+                {
+                    prompt_tokens = Some(info.prompt_token_ids.len());
+                    if return_token_ids {
+                        prompt_token_ids = Some(info.prompt_token_ids.to_vec());
+                    }
                 }
                 usage.prompt_token_count = prompt_tokens.unwrap_or_default();
                 usage.cached_token_count = usage.cached_token_count.max(output.cached_token_count);
@@ -195,6 +220,7 @@ async fn generate_chunk_stream(
                     None
                 };
 
+                let prompt_token_ids = prompt_token_ids.take();
                 y.yield_ok(GenerateStreamResponse {
                     request_id: request_id.clone(),
                     choices: vec![GenerateResponseStreamChoice {
@@ -205,6 +231,9 @@ async fn generate_chunk_stream(
                     }],
                     usage: include_continuous_usage
                         .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)),
+                    mm_placeholders: prompt_token_ids.as_ref().and_then(|_| mm_placeholders.take()),
+                    prompt_token_ids,
+                    metrics: None,
                 })
                 .await;
             }
@@ -223,6 +252,13 @@ async fn generate_chunk_stream(
             request_id,
             choices: Vec::new(),
             usage: Some(Usage::from_token_usage(usage, enable_prompt_tokens_details)),
+            prompt_token_ids: None,
+            mm_placeholders: None,
+            metrics: spec_decode_metrics.map(|speculative_decoding| PerRequestMetrics {
+                speculative_decoding: StreamingSpeculativeDecodingMetrics::from(
+                    speculative_decoding,
+                ),
+            }),
         })
         .await;
     }
@@ -244,7 +280,9 @@ fn collect_generate(
         include_continuous_usage: _,
         include_logprobs,
         include_prompt_logprobs,
+        return_token_ids,
     }: ResponseOptions,
+    mm_placeholders: Option<MultiModalPlaceholders>,
 ) -> Result<GenerateResponse, ApiError> {
     let logprobs = if include_logprobs {
         let logprobs = collected.logprobs.as_ref().ok_or_else(|| {
@@ -292,11 +330,37 @@ fn collect_generate(
             token_ids: collected.token_ids,
         }],
         prompt_logprobs,
+        prompt_token_ids: return_token_ids.then_some(collected.prompt_token_ids),
+        mm_placeholders: return_token_ids.then_some(mm_placeholders).flatten(),
         kv_transfer_params: collected.kv_transfer_params,
         ec_transfer_params: collected.ec_transfer_params,
+        metrics: collected.spec_decode_metrics.map(|metrics| PerRequestMetrics {
+            speculative_decoding: SpeculativeDecodingMetrics::from(metrics),
+        }),
     })
 }
 
+fn extract_mm_placeholders(features: Option<&[MmFeatureSpec]>) -> Option<MultiModalPlaceholders> {
+    let features = features?;
+    if features.is_empty() {
+        return None;
+    }
+
+    let mut placeholders = MultiModalPlaceholders::new();
+    for feature in features {
+        placeholders.entry(feature.modality.as_str().to_string()).or_default().push(
+            PlaceholderRangeInfo {
+                offset: feature.mm_position.offset,
+                length: feature.mm_position.length,
+            },
+        );
+    }
+    for ranges in placeholders.values_mut() {
+        ranges.sort_by_key(|range| range.offset);
+    }
+    Some(placeholders)
+}
+
 fn raw_logprobs_to_openai_chat(logprobs: &Logprobs) -> Result<ChatLogProbs, ApiError> {
     let content = logprobs
         .positions
@@ -419,6 +483,8 @@ mod tests {
     use std::sync::Arc;
 
     use futures::{TryStreamExt as _, stream};
+    use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange};
+    use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
     use vllm_llm::GeneratePromptInfo;
 
     use super::*;
@@ -435,6 +501,8 @@ mod tests {
                 cached_token_count: 0,
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                sampling_mask: None,
+                spec_decode_metrics: None,
             }),
             Ok(GenerateOutput {
                 request_id: String::new(),
@@ -448,9 +516,18 @@ mod tests {
                 cached_token_count: 2,
                 kv_transfer_params: None,
                 ec_transfer_params: None,
+                sampling_mask: None,
+                spec_decode_metrics: None,
             }),
         ]);
 
+        let mm_placeholders = MultiModalPlaceholders::from([(
+            "image".to_string(),
+            vec![PlaceholderRangeInfo {
+                offset: 1,
+                length: 4,
+            }],
+        )]);
         let chunks: Vec<_> = generate_chunk_stream(
             stream,
             "raw-stream".to_string(),
@@ -461,14 +538,20 @@ mod tests {
             ResponseOptions {
                 include_usage: true,
                 include_continuous_usage: true,
+                return_token_ids: true,
                 ..Default::default()
             },
+            Some(mm_placeholders.clone()),
         )
         .try_collect()
         .await
         .expect("collect chunks");
 
         assert_eq!(chunks.len(), 2);
+        assert_eq!(chunks[0].prompt_token_ids, Some(vec![11, 22]));
+        assert_eq!(chunks[0].mm_placeholders, Some(mm_placeholders));
+        assert!(chunks[1].prompt_token_ids.is_none());
+        assert!(chunks[1].mm_placeholders.is_none());
         assert_eq!(
             chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
             2
@@ -499,6 +582,333 @@ mod tests {
         );
     }
 
+    fn stream_output(
+        prompt_token_ids: Option<&[u32]>,
+        token_ids: Vec<u32>,
+        finish_reason: Option<FinishReason>,
+    ) -> GenerateOutput {
+        GenerateOutput {
+            request_id: String::new(),
+            prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo {
+                prompt_token_ids: Arc::from(ids),
+                prompt_logprobs: None,
+            }),
+            token_ids,
+            logprobs: None,
+            finish_reason,
+            cached_token_count: 0,
+            kv_transfer_params: None,
+            ec_transfer_params: None,
+            sampling_mask: None,
+            spec_decode_metrics: None,
+        }
+    }
+
+    async fn collect_chunks(
+        outputs: Vec<GenerateOutput>,
+        return_token_ids: bool,
+        mm_placeholders: Option<MultiModalPlaceholders>,
+    ) -> Vec<GenerateStreamResponse> {
+        generate_chunk_stream(
+            stream::iter(outputs.into_iter().map(Ok)),
+            "raw-stream".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions {
+                return_token_ids,
+                ..Default::default()
+            },
+            mm_placeholders,
+        )
+        .try_collect()
+        .await
+        .expect("collect chunks")
+    }
+
+    #[tokio::test]
+    async fn generate_chunk_stream_omits_prompt_metadata_by_default() {
+        let chunks = collect_chunks(
+            vec![
+                stream_output(Some(&[11, 22]), vec![33], None),
+                stream_output(Some(&[11, 22]), vec![44], Some(FinishReason::stop_eos())),
+            ],
+            false,
+            Some(MultiModalPlaceholders::from([(
+                "image".to_string(),
+                vec![PlaceholderRangeInfo {
+                    offset: 1,
+                    length: 4,
+                }],
+            )])),
+        )
+        .await;
+
+        assert_eq!(chunks.len(), 2);
+        assert!(chunks.iter().all(|c| c.prompt_token_ids.is_none()));
+        assert!(chunks.iter().all(|c| c.mm_placeholders.is_none()));
+    }
+
+    #[tokio::test]
+    async fn generate_chunk_stream_sends_prompt_metadata_on_first_emitted_chunk() {
+        // An empty first output is skipped and must not consume the metadata.
+        let chunks = collect_chunks(
+            vec![
+                stream_output(Some(&[11, 22]), Vec::new(), None),
+                stream_output(Some(&[11, 22]), vec![33], None),
+                stream_output(Some(&[11, 22]), vec![44], Some(FinishReason::stop_eos())),
+            ],
+            true,
+            None,
+        )
+        .await;
+
+        assert_eq!(chunks.len(), 2);
+        assert_eq!(chunks[0].prompt_token_ids, Some(vec![11, 22]));
+        assert!(chunks[0].mm_placeholders.is_none());
+        assert!(chunks[1].prompt_token_ids.is_none());
+        assert!(chunks[1].mm_placeholders.is_none());
+    }
+
+    #[tokio::test]
+    async fn generate_chunk_stream_returns_terminal_spec_decode_metrics() {
+        let metrics = RequestSpecDecodeMetrics {
+            num_spec_tokens: 3,
+            histogram: vec![0, 1, 1, 0],
+            num_draft_tokens: 6,
+            ..Default::default()
+        };
+        let mut output = stream_output(None, Vec::new(), Some(FinishReason::stop_eos()));
+        output.spec_decode_metrics = Some(metrics.clone());
+
+        for include_usage in [false, true] {
+            let chunks: Vec<_> = generate_chunk_stream(
+                stream::iter(vec![Ok(output.clone())]),
+                "raw-stream".to_string(),
+                ApiServerOptions::default(),
+                ResponseOptions {
+                    include_usage,
+                    ..Default::default()
+                },
+                None,
+            )
+            .try_collect()
+            .await
+            .expect("collect chunks");
+            assert_eq!(chunks.len(), 1 + usize::from(include_usage));
+            assert!(chunks[0].metrics.is_none());
+            if include_usage {
+                assert!(chunks[1].choices.is_empty());
+                assert_eq!(
+                    chunks[1].metrics,
+                    Some(PerRequestMetrics {
+                        speculative_decoding: StreamingSpeculativeDecodingMetrics::from(
+                            SpeculativeDecodingMetrics::from(metrics.clone())
+                        )
+                    })
+                );
+                let json = serde_json::to_value(&chunks[1]).expect("serialize chunk");
+                let stats = &json["metrics"]["speculative_decoding"];
+                assert_eq!(stats.as_object().expect("metrics object").len(), 7);
+                assert_eq!(stats["num_draft_tokens"], 6);
+                assert_eq!(stats["num_accepted_draft_tokens"], 3);
+                assert_eq!(stats["mean_acceptance_length"], 2.5);
+                assert_eq!(
+                    stats["acceptance_histogram"],
+                    serde_json::json!([0, 1, 1, 0])
+                );
+                assert!(stats.get("per_step_accepted").is_none());
+            }
+        }
+    }
+
+    #[test]
+    fn spec_decode_metrics_zero_steps_and_detailed() {
+        let empty = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics {
+            num_spec_tokens: 2,
+            histogram: vec![0, 0, 0],
+            ..Default::default()
+        });
+        assert_eq!(empty.mean_acceptance_length, 1.0);
+        assert_eq!(empty.draft_acceptance_rate, 0.0);
+        assert_eq!(empty.per_step_accepted, None);
+
+        let detailed = SpeculativeDecodingMetrics::from(RequestSpecDecodeMetrics {
+            num_spec_tokens: 2,
+            histogram: vec![0, 1, 1],
+            num_draft_tokens: 4,
+            per_step_accepted: vec![1, 2],
+            per_step_drafted: vec![2, 2],
+        });
+        assert_eq!(detailed.per_step_accepted, Some(vec![1, 2]));
+        assert_eq!(detailed.per_step_drafted, Some(vec![2, 2]));
+    }
+
+    #[test]
+    fn collect_generate_omits_prompt_metadata_by_default() {
+        let output = CollectedGenerateOutput {
+            request_id: "raw-1".to_string(),
+            prompt_logprobs: None,
+            token_ids: vec![30],
+            logprobs: None,
+            finish_reason: FinishReason::stop_eos(),
+            usage: TokenUsage {
+                prompt_token_count: 2,
+                output_token_count: 1,
+                cached_token_count: 0,
+            },
+            kv_transfer_params: None,
+            ec_transfer_params: None,
+            prompt_token_ids: vec![10, 20],
+            sampling_mask: None,
+            spec_decode_metrics: None,
+        };
+
+        let response = collect_generate(
+            output,
+            "raw-1".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions::default(),
+            Some(MultiModalPlaceholders::from([(
+                "image".to_string(),
+                vec![PlaceholderRangeInfo {
+                    offset: 0,
+                    length: 1,
+                }],
+            )])),
+        )
+        .expect("response");
+
+        assert!(response.prompt_token_ids.is_none());
+        assert!(response.mm_placeholders.is_none());
+    }
+
+    #[test]
+    fn collect_generate_returns_spec_decode_metrics() {
+        let metrics = RequestSpecDecodeMetrics {
+            num_spec_tokens: 3,
+            histogram: vec![0, 1, 1, 0],
+            num_draft_tokens: 6,
+            ..Default::default()
+        };
+        let output = CollectedGenerateOutput {
+            request_id: "raw-metrics".to_string(),
+            prompt_logprobs: None,
+            token_ids: vec![30],
+            logprobs: None,
+            finish_reason: FinishReason::stop_eos(),
+            usage: TokenUsage::default(),
+            kv_transfer_params: None,
+            ec_transfer_params: None,
+            prompt_token_ids: vec![10],
+            sampling_mask: None,
+            spec_decode_metrics: Some(metrics.clone()),
+        };
+
+        let response = collect_generate(
+            output,
+            "raw-metrics".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions::default(),
+            None,
+        )
+        .expect("response");
+
+        assert_eq!(
+            response.metrics,
+            Some(PerRequestMetrics {
+                speculative_decoding: SpeculativeDecodingMetrics::from(metrics)
+            })
+        );
+        let json = serde_json::to_value(&response).expect("serialize response");
+        let stats = &json["metrics"]["speculative_decoding"];
+        assert_eq!(stats.as_object().expect("metrics object").len(), 9);
+        assert!(stats["per_step_accepted"].is_null());
+        assert!(stats["per_step_drafted"].is_null());
+    }
+
+    #[test]
+    fn extract_mm_placeholders_returns_none_without_features() {
+        assert!(extract_mm_placeholders(None).is_none());
+        assert!(extract_mm_placeholders(Some(&[])).is_none());
+    }
+
+    #[test]
+    fn extract_mm_placeholders_groups_ranges_in_prompt_order() {
+        let feature = |modality: MmModality, offset: usize, length: usize| MmFeatureSpec {
+            data: None,
+            modality,
+            identifier: format!("{}-{offset}", modality.as_str()),
+            mm_position: PlaceholderRange {
+                offset,
+                length,
+                is_embed: None,
+            },
+            mm_hash: None,
+        };
+        let features = vec![
+            feature(MmModality::Image, 8, 2),
+            feature(MmModality::Audio, 4, 3),
+            feature(MmModality::Image, 1, 5),
+        ];
+
+        let placeholders = extract_mm_placeholders(Some(&features)).expect("metadata");
+
+        assert_eq!(
+            placeholders["image"],
+            vec![
+                PlaceholderRangeInfo {
+                    offset: 1,
+                    length: 5,
+                },
+                PlaceholderRangeInfo {
+                    offset: 8,
+                    length: 2,
+                },
+            ]
+        );
+        assert_eq!(
+            placeholders["audio"],
+            vec![PlaceholderRangeInfo {
+                offset: 4,
+                length: 3,
+            }]
+        );
+    }
+
+    #[test]
+    fn collect_generate_returns_prompt_token_ids_when_requested() {
+        let output = CollectedGenerateOutput {
+            request_id: "raw-1".to_string(),
+            prompt_logprobs: None,
+            token_ids: vec![30],
+            logprobs: None,
+            finish_reason: FinishReason::stop_eos(),
+            usage: TokenUsage {
+                prompt_token_count: 2,
+                output_token_count: 1,
+                cached_token_count: 0,
+            },
+            kv_transfer_params: None,
+            ec_transfer_params: None,
+            prompt_token_ids: vec![10, 20],
+            sampling_mask: None,
+            spec_decode_metrics: None,
+        };
+
+        let response = collect_generate(
+            output,
+            "raw-1".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions {
+                return_token_ids: true,
+                ..Default::default()
+            },
+            None,
+        )
+        .expect("response");
+
+        assert_eq!(response.prompt_token_ids, Some(vec![10, 20]));
+    }
+
     #[test]
     fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
         let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
@@ -514,6 +924,8 @@ mod tests {
             },
             kv_transfer_params: None,
             ec_transfer_params: None,
+            sampling_mask: None,
+            spec_decode_metrics: None,
             prompt_token_ids,
         };
 
@@ -525,6 +937,7 @@ mod tests {
                 include_prompt_logprobs: true,
                 ..Default::default()
             },
+            None,
         )
         .expect("single-token prompt without payload maps to [None]");
         let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
@@ -539,6 +952,7 @@ mod tests {
                 include_prompt_logprobs: true,
                 ..Default::default()
             },
+            None,
         )
         .expect_err("multi-token prompt without payload is an engine failure");
     }
diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs
index cb06cfe33c43..4c54619ed4a7 100644
--- a/rust/src/server/src/routes/inference/generate/convert.rs
+++ b/rust/src/server/src/routes/inference/generate/convert.rs
@@ -30,6 +30,8 @@ pub(super) struct ResponseOptions {
     pub include_logprobs: bool,
     /// Whether the caller requested top-level prompt logprobs.
     pub include_prompt_logprobs: bool,
+    /// Whether the caller requested final prompt token metadata.
+    pub return_token_ids: bool,
 }
 
 /// Validate and lower one raw generate request into the internal
@@ -56,6 +58,7 @@ pub(super) fn prepare_generate_request(
             .unwrap_or(false);
     let include_logprobs = request.sampling_params.inner.logprobs.is_some();
     let include_prompt_logprobs = request.sampling_params.inner.prompt_logprobs.is_some();
+    let return_token_ids = request.return_token_ids.unwrap_or(false);
     let mut sampling_params = request.sampling_params.inner;
     sampling_params.vllm_xargs = merge_kv_transfer_params(
         sampling_params.vllm_xargs,
@@ -79,7 +82,9 @@ pub(super) fn prepare_generate_request(
         add_special_tokens: false,
         data_parallel_rank: ctx.data_parallel_rank,
         session_id: ctx.session_id,
-        reasoning_parser_kwargs: None,
+        kv_hints: None,
+        reasoning_parser_kwargs: Default::default(),
+        reasoning_ended: None,
         lora_request: lora_resolution.lora_request.clone(),
         arrival_time: None,
     };
@@ -93,6 +98,7 @@ pub(super) fn prepare_generate_request(
             include_continuous_usage,
             include_logprobs,
             include_prompt_logprobs,
+            return_token_ids,
         },
     })
 }
@@ -164,6 +170,30 @@ mod tests {
         );
     }
 
+    #[test]
+    fn prepare_generate_request_preserves_watermarking_defaults_and_opt_out() {
+        for watermarking in [None, Some(true), Some(false)] {
+            let mut body = json!({
+                "token_ids": [11, 22],
+                "sampling_params": {}
+            });
+            if let Some(watermarking) = watermarking {
+                body["sampling_params"]["watermarking"] = json!(watermarking);
+            }
+            let prepared = prepare_generate_request(
+                serde_json::from_value(body).expect("parse request"),
+                &served(&["test-model"]),
+                ResolvedRequestContext::default(),
+                None,
+            )
+            .expect("prepare request");
+            assert_eq!(
+                prepared.text_request.sampling_params.watermarking,
+                watermarking.unwrap_or(true)
+            );
+        }
+    }
+
     #[test]
     fn prepare_generate_request_forwards_thinking_token_budget() {
         let request: GenerateRequest = serde_json::from_value(json!({
diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs
index b9220aa7b767..86d8ec00f4fc 100644
--- a/rust/src/server/src/routes/inference/generate/types.rs
+++ b/rust/src/server/src/routes/inference/generate/types.rs
@@ -7,6 +7,7 @@ use llm_multimodal::MediaContentPart;
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 use validator::Validate;
+use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
 use vllm_text::SamplingParams;
 
 use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
@@ -46,6 +47,7 @@ pub struct GenerateRequest {
     pub ec_transfer_params: Option<HashMap<String, Value>>,
     /// Raw multimodal input; server resolves media. Mutually exclusive with `features`.
     pub content_parts: Option<Vec<MediaContentPart>>,
+    pub return_token_ids: Option<bool>,
     #[serde(flatten)]
     pub other: Map<String, Value>,
 }
@@ -81,6 +83,9 @@ pub(super) struct GenerateStreamResponse {
     pub request_id: String,
     pub choices: Vec<GenerateResponseStreamChoice>,
     pub usage: Option<Usage>,
+    pub prompt_token_ids: Option<Vec<u32>>,
+    pub mm_placeholders: Option<MultiModalPlaceholders>,
+    pub metrics: Option<PerRequestMetrics<StreamingSpeculativeDecodingMetrics>>,
 }
 
 /// Mirrors the Python vLLM `GenerateResponse` class.
@@ -89,8 +94,103 @@ pub(super) struct GenerateResponse {
     pub request_id: String,
     pub choices: Vec<GenerateResponseChoice>,
     pub prompt_logprobs: Option<Vec<Option<HashMap<u32, GenerateLogprob>>>>,
+    pub prompt_token_ids: Option<Vec<u32>>,
+    pub mm_placeholders: Option<MultiModalPlaceholders>,
     pub kv_transfer_params: Option<Value>,
     pub ec_transfer_params: Option<Value>,
+    pub metrics: Option<PerRequestMetrics<SpeculativeDecodingMetrics>>,
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize)]
+pub(super) struct PerRequestMetrics<T> {
+    pub speculative_decoding: T,
+}
+
+/// Mirrors the Python vLLM `SpeculativeDecodingMetrics` class.
+///
+/// Derived from the raw engine accumulator the same way as Python
+/// `RequestSpecDecodeMetrics.to_dict`.
+#[derive(Debug, Clone, PartialEq, Serialize)]
+pub(super) struct SpeculativeDecodingMetrics {
+    pub mean_acceptance_length: f64,
+    pub draft_acceptance_rate: f64,
+    pub acceptance_histogram: Vec<u64>,
+    pub num_spec_steps: u64,
+    pub num_accepted_draft_tokens: u64,
+    pub num_draft_tokens: u64,
+    pub num_spec_tokens: u64,
+    pub per_step_accepted: Option<Vec<u64>>,
+    pub per_step_drafted: Option<Vec<u64>>,
+}
+
+impl From<RequestSpecDecodeMetrics> for SpeculativeDecodingMetrics {
+    fn from(raw: RequestSpecDecodeMetrics) -> Self {
+        let num_spec_steps: u64 = raw.histogram.iter().sum();
+        let num_accepted_draft_tokens: u64 =
+            (0u64..).zip(&raw.histogram).map(|(accepted, count)| accepted * count).sum();
+        let ratio = |num: u64, den: u64| {
+            if den == 0 {
+                0.0
+            } else {
+                num as f64 / den as f64
+            }
+        };
+        let detailed = !raw.per_step_accepted.is_empty();
+        Self {
+            mean_acceptance_length: if num_spec_steps == 0 {
+                1.0
+            } else {
+                1.0 + ratio(num_accepted_draft_tokens, num_spec_steps)
+            },
+            draft_acceptance_rate: ratio(num_accepted_draft_tokens, raw.num_draft_tokens),
+            acceptance_histogram: raw.histogram,
+            num_spec_steps,
+            num_accepted_draft_tokens,
+            num_draft_tokens: raw.num_draft_tokens,
+            num_spec_tokens: raw.num_spec_tokens,
+            per_step_accepted: detailed.then_some(raw.per_step_accepted),
+            per_step_drafted: detailed.then_some(raw.per_step_drafted),
+        }
+    }
+}
+
+/// Streaming form omits detailed fields when summary metrics are requested.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, PartialEq, Serialize)]
+pub(super) struct StreamingSpeculativeDecodingMetrics {
+    pub mean_acceptance_length: f64,
+    pub draft_acceptance_rate: f64,
+    pub acceptance_histogram: Vec<u64>,
+    pub num_spec_steps: u64,
+    pub num_accepted_draft_tokens: u64,
+    pub num_draft_tokens: u64,
+    pub num_spec_tokens: u64,
+    pub per_step_accepted: Option<Vec<u64>>,
+    pub per_step_drafted: Option<Vec<u64>>,
+}
+
+impl From<SpeculativeDecodingMetrics> for StreamingSpeculativeDecodingMetrics {
+    fn from(metrics: SpeculativeDecodingMetrics) -> Self {
+        Self {
+            mean_acceptance_length: metrics.mean_acceptance_length,
+            draft_acceptance_rate: metrics.draft_acceptance_rate,
+            acceptance_histogram: metrics.acceptance_histogram,
+            num_spec_steps: metrics.num_spec_steps,
+            num_accepted_draft_tokens: metrics.num_accepted_draft_tokens,
+            num_draft_tokens: metrics.num_draft_tokens,
+            num_spec_tokens: metrics.num_spec_tokens,
+            per_step_accepted: metrics.per_step_accepted,
+            per_step_drafted: metrics.per_step_drafted,
+        }
+    }
+}
+
+pub(super) type MultiModalPlaceholders = HashMap<String, Vec<PlaceholderRangeInfo>>;
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+pub(super) struct PlaceholderRangeInfo {
+    pub offset: usize,
+    pub length: usize,
 }
 
 /// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads.
diff --git a/rust/src/server/src/routes/kv_event_sources.rs b/rust/src/server/src/routes/kv_event_sources.rs
new file mode 100644
index 000000000000..c3d36493a9a0
--- /dev/null
+++ b/rust/src/server/src/routes/kv_event_sources.rs
@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use axum::Json;
+use axum::extract::State;
+use vllm_engine_core_client::protocol::handshake::KvEventsConfig;
+
+use crate::state::AppState;
+
+/// KV-event publisher config of each data-parallel rank, keyed by rank.
+pub async fn kv_event_sources(
+    State(state): State<Arc<AppState>>,
+) -> Json<BTreeMap<u32, KvEventsConfig>> {
+    Json(
+        state
+            .engine_core_client()
+            .ready_responses()
+            .into_iter()
+            .filter_map(|ready| Some((ready.data_parallel_rank, ready.kv_events_config.clone()?)))
+            .collect(),
+    )
+}
diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs
index 474719a875d8..a87c52fb5578 100644
--- a/rust/src/server/src/routes/openai/chat_completions.rs
+++ b/rust/src/server/src/routes/openai/chat_completions.rs
@@ -13,7 +13,7 @@ use asynk_strim_attr::{TryYielder, try_stream};
 use axum::Json;
 use axum::extract::State;
 use axum::http::HeaderMap;
-use axum::response::sse::{Event, Sse};
+use axum::response::sse::Event;
 use axum::response::{IntoResponse, Response};
 use futures::{Stream, StreamExt as _, pin_mut};
 use serde_json::Value;
@@ -45,7 +45,7 @@ use crate::routes::openai::utils::types::{
 use crate::routes::openai::utils::usage::ContinuousUsage;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
 use crate::state::AppState;
-use crate::utils::{ResolvedRequestContext, resolve_request_context, unix_timestamp};
+use crate::utils::{ResolvedRequestContext, resolve_request_context, sse_response, unix_timestamp};
 
 pub(crate) fn lower_chat_request(
     request: ChatCompletionRequest,
@@ -99,7 +99,7 @@ pub async fn chat_completions(
         );
         let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
 
-        Sse::new(sse_stream).into_response()
+        sse_response(sse_stream, api_server_options.sse_keep_alive_interval)
     } else {
         let response = match collect_chat_completion(
             chat_stream,
diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs
index 3aaba8df28dc..d4df61427028 100644
--- a/rust/src/server/src/routes/openai/chat_completions/convert.rs
+++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs
@@ -150,6 +150,7 @@ pub(super) fn prepare_chat_request(
         messages,
         sampling_params: SamplingParams {
             temperature: request.temperature,
+            watermarking: request.watermarking,
             top_p: request.top_p,
             top_k: request.top_k,
             seed: request.seed,
@@ -179,7 +180,7 @@ pub(super) fn prepare_chat_request(
         chat_options: ChatOptions {
             generation_prompt_mode,
             chat_template: request.chat_template,
-            reasoning_effort: request.reasoning_effort,
+            reasoning_effort: request.reasoning_effort.map(|effort| effort.as_str().into()),
             response_format,
             template_kwargs,
         },
@@ -316,6 +317,9 @@ pub(crate) fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, A
             convert_content(content)?,
             convert_message_tools(tools)?,
         )),
+        ChatMessage::Custom { role, content } => {
+            Ok(VllmChatMessage::custom(role, convert_content(content)?))
+        }
     }
 }
 
@@ -437,10 +441,11 @@ mod tests {
     use std::collections::HashMap;
     use std::sync::Arc;
 
-    use axum::http::HeaderMap;
+    use axum::http::{HeaderMap, StatusCode};
     use expect_test::expect;
     use llm_multimodal::ImageDetail;
     use serde_json::json;
+    use thiserror_ext::AsReport as _;
     use validator::Validate;
     use vllm_chat::{
         AssistantContentBlock, AssistantToolCall, ChatContentPart, ChatMessage as VllmChatMessage,
@@ -487,6 +492,50 @@ mod tests {
         }
     }
 
+    #[test]
+    fn chat_http_reasoning_effort_preserves_omission_none_and_kwargs() {
+        for effort in [None, Some(json!(null))].into_iter().chain(
+            ["none", "minimal", "low", "medium", "high", "xhigh", "max"]
+                .map(|name| Some(json!(name))),
+        ) {
+            let kwargs = json!({"reasoning_effort": 37, "thinking": true});
+            let mut value = json!({
+                "messages": [{"role": "user", "content": "hello"}],
+                "chat_template_kwargs": kwargs,
+            });
+            if let Some(effort) = &effort {
+                value["reasoning_effort"] = effort.clone();
+            }
+            let request: ChatCompletionRequest = serde_json::from_value(value).unwrap();
+            let prepared = prepare_chat_request(
+                request,
+                &served(&["test-model"]),
+                request_context(&HeaderMap::new(), None),
+            )
+            .unwrap();
+            let options = prepared.chat_request.chat_options;
+            assert_eq!(
+                serde_json::to_value(options.reasoning_effort).unwrap(),
+                effort.unwrap_or(serde_json::Value::Null)
+            );
+            assert_eq!(
+                serde_json::to_value(options.template_kwargs).unwrap(),
+                kwargs
+            );
+        }
+    }
+
+    #[test]
+    fn chat_http_reasoning_effort_rejects_model_extensions_at_top_level() {
+        for effort in [json!(37), json!("custom")] {
+            let request = json!({
+                "messages": [{"role": "user", "content": "hello"}],
+                "reasoning_effort": effort,
+            });
+            assert!(serde_json::from_value::<ChatCompletionRequest>(request).is_err());
+        }
+    }
+
     #[test]
     fn chat_http_request_defaults_missing_or_null_model() {
         for model in [None, Some(serde_json::Value::Null)] {
@@ -583,6 +632,61 @@ mod tests {
         assert_eq!(prepared.chat_request.decode_options.min_tokens, 0);
     }
 
+    #[test]
+    fn prepare_chat_request_rejects_empty_json_schema() {
+        let mut request = base_request();
+        request.response_format = Some(ResponseFormat::JsonSchema {
+            json_schema: JsonSchemaFormat {
+                name: "answer".to_string(),
+                description: None,
+                schema: json!("  "),
+                strict: None,
+            },
+        });
+
+        let error = prepare_chat_request(
+            request,
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+            ResolvedRequestContext::default(),
+        )
+        .unwrap_err();
+
+        assert_eq!(error.status_code(), StatusCode::BAD_REQUEST);
+        assert!(
+            error
+                .to_error_response()
+                .error
+                .message
+                .contains("json cannot be an empty string")
+        );
+    }
+
+    #[test]
+    fn prepare_chat_request_rejects_whitespace_only_structured_outputs_grammar() {
+        let request: ChatCompletionRequest = serde_json::from_value(json!({
+            "model": "Qwen/Qwen1.5-0.5B-Chat",
+            "messages": [{"role": "user", "content": "hello"}],
+            "structured_outputs": {"grammar": " \t\n"},
+        }))
+        .expect("parse structured_outputs");
+
+        let error = prepare_chat_request(
+            request,
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+            ResolvedRequestContext::default(),
+        )
+        .unwrap_err();
+
+        assert_eq!(error.status_code(), StatusCode::BAD_REQUEST);
+        assert!(
+            error
+                .to_error_response()
+                .error
+                .message
+                .contains("grammar cannot be an empty string")
+        );
+    }
+
     #[test]
     fn prepare_chat_request_defaults_function_tool_fields() {
         let request: ChatCompletionRequest = serde_json::from_value(json!({
@@ -691,7 +795,7 @@ mod tests {
         );
 
         let tokenizer = Arc::new(TestTokenizer::new());
-        let prompt = KimiK3ChatRenderer::new(tokenizer.clone())
+        let prompt = KimiK3ChatRenderer::new(tokenizer.clone(), Default::default())
             .render(&prepared.chat_request)
             .expect("Kimi K3 rendering succeeds")
             .prompt;
@@ -947,6 +1051,43 @@ mod tests {
         );
     }
 
+    #[test]
+    fn prepare_chat_request_accepts_custom_role_messages() {
+        let request: ChatCompletionRequest = serde_json::from_value(json!({
+            "messages": [
+                {"role": "root", "content": "Custom identity."},
+                {"role": "user", "content": "hello"},
+            ],
+        }))
+        .expect("parse custom role message");
+
+        let prepared = prepare_chat_request(
+            request,
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+            ResolvedRequestContext::default(),
+        )
+        .expect("request is valid");
+
+        assert_eq!(
+            prepared.chat_request.messages,
+            vec![
+                VllmChatMessage::custom("root", "Custom identity."),
+                VllmChatMessage::user("hello"),
+            ]
+        );
+    }
+
+    #[test]
+    fn chat_http_message_keeps_standard_role_errors() {
+        let error = serde_json::from_value::<ChatMessage>(json!({
+            "role": "tool",
+            "content": "Sunny",
+        }))
+        .unwrap_err();
+
+        expect!["missing field `tool_call_id`"].assert_eq(&error.to_report_string());
+    }
+
     #[test]
     fn prepare_chat_request_maps_image_url_content_parts() {
         let request = ChatCompletionRequest {
diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs
index e4da86c22d9b..13c7c7233f14 100644
--- a/rust/src/server/src/routes/openai/chat_completions/types.rs
+++ b/rust/src/server/src/routes/openai/chat_completions/types.rs
@@ -9,15 +9,15 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use serde_with::SerializeDisplay;
 use validator::Validate;
-use vllm_chat::ReasoningEffort;
 use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams;
 use vllm_text::TruncationSide;
 
 use crate::routes::openai::utils::structured_outputs::ResponseFormat;
 use crate::routes::openai::utils::types::{
-    ChatLogProbs, ChatMessage, Normalizable, PromptLogprobs, StreamOptions, StreamResponseEnvelope,
-    StringOrArray, Tool, ToolCall, ToolCallDelta, ToolChoice, Usage, default_true,
-    deserialize_request_top_k, validate_messages, validate_stop, validate_top_p_value,
+    ChatLogProbs, ChatMessage, Normalizable, PromptLogprobs, ReasoningEffort, StreamOptions,
+    StreamResponseEnvelope, StringOrArray, Tool, ToolCall, ToolCallDelta, ToolChoice, Usage,
+    default_true, deserialize_request_top_k, validate_messages, validate_stop,
+    validate_top_p_value,
 };
 
 /// vLLM-compatible request type for the Chat Completions API.
@@ -116,6 +116,10 @@ pub struct ChatCompletionRequest {
     pub user: Option<String>,
 
     // -------- vLLM Sampling Parameters --------
+    /// Whether to apply the engine's configured watermark to this request.
+    #[serde(default = "default_true")]
+    pub watermarking: bool,
+
     /// Use beam search instead of sampling
     #[serde(default)]
     pub use_beam_search: bool,
@@ -282,6 +286,7 @@ impl Default for ChatCompletionRequest {
             include_reasoning: true,
             parallel_tool_calls: None,
             user: None,
+            watermarking: true,
             use_beam_search: false,
             top_k: None,
             min_p: None,
diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs
index 7260155799b6..1bfc1d9a091d 100644
--- a/rust/src/server/src/routes/openai/chat_completions/validate.rs
+++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs
@@ -125,8 +125,8 @@ fn validate_function_tools(tools: &[Tool], param: &'static str) -> Result<(), Ap
 mod tests {
     use std::collections::HashMap;
 
+    use crate::routes::openai::utils::types::ReasoningEffort;
     use serde_json::json;
-    use vllm_chat::ReasoningEffort;
 
     use super::validate_request_compat;
     use crate::routes::openai::chat_completions::types::ChatCompletionRequest;
diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs
index a122afbcd97a..7074b8c6293c 100644
--- a/rust/src/server/src/routes/openai/completions.rs
+++ b/rust/src/server/src/routes/openai/completions.rs
@@ -13,7 +13,7 @@ use asynk_strim_attr::{TryYielder, try_stream};
 use axum::Json;
 use axum::extract::State;
 use axum::http::HeaderMap;
-use axum::response::sse::{Event, Sse};
+use axum::response::sse::Event;
 use axum::response::{IntoResponse, Response};
 use futures::{Stream, StreamExt as _, pin_mut};
 use serde_json::Value;
@@ -45,7 +45,7 @@ use crate::routes::openai::utils::types::LogProbs;
 use crate::routes::openai::utils::usage::ContinuousUsage;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
 use crate::state::AppState;
-use crate::utils::{ResolvedRequestContext, resolve_request_context, unix_timestamp};
+use crate::utils::{ResolvedRequestContext, resolve_request_context, sse_response, unix_timestamp};
 
 pub(crate) fn lower_completion_request(
     request: CompletionRequest,
@@ -111,7 +111,7 @@ pub async fn completions(
         );
         let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span);
 
-        Sse::new(sse_stream).into_response()
+        sse_response(sse_stream, api_server_options.sse_keep_alive_interval)
     } else {
         let response = match collect_completion(
             text_stream,
@@ -694,6 +694,7 @@ mod tests {
                     ))),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 }),
             )),
         ]);
@@ -806,6 +807,7 @@ mod tests {
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 }),
             )),
         ]);
@@ -858,6 +860,7 @@ mod tests {
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 }),
             )),
         ]);
@@ -913,6 +916,7 @@ mod tests {
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 }),
             )),
         ]);
@@ -985,6 +989,7 @@ mod tests {
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 }),
             )),
         ]);
@@ -1059,6 +1064,7 @@ mod tests {
                     finish_reason: FinishReason::Length,
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 }),
             )),
         ]);
diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs
index d944dca65740..53a6ff9906ce 100644
--- a/rust/src/server/src/routes/openai/completions/convert.rs
+++ b/rust/src/server/src/routes/openai/completions/convert.rs
@@ -115,6 +115,7 @@ pub(super) fn prepare_completion_request(
         mm_features: None,
         sampling_params: SamplingParams {
             temperature: request.temperature,
+            watermarking: request.watermarking,
             top_p: request.top_p,
             top_k: request.top_k,
             seed: request.seed,
@@ -154,7 +155,9 @@ pub(super) fn prepare_completion_request(
         add_special_tokens: request.add_special_tokens,
         data_parallel_rank: ctx.data_parallel_rank,
         session_id,
-        reasoning_parser_kwargs: None,
+        kv_hints: None,
+        reasoning_parser_kwargs: Default::default(),
+        reasoning_ended: None,
         lora_request: lora_resolution.lora_request.clone(),
         arrival_time: None,
     };
diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs
index 8d93a89ed471..2a6901370682 100644
--- a/rust/src/server/src/routes/openai/completions/types.rs
+++ b/rust/src/server/src/routes/openai/completions/types.rs
@@ -90,6 +90,10 @@ pub struct CompletionRequest {
     pub user: Option<String>,
 
     // -------- vLLM Sampling Parameters --------
+    /// Whether to apply the engine's configured watermark to this request.
+    #[serde(default = "default_true")]
+    pub watermarking: bool,
+
     /// Options for streaming response
     pub stream_options: Option<StreamOptions>,
 
diff --git a/rust/src/server/src/routes/openai/utils/logprobs.rs b/rust/src/server/src/routes/openai/utils/logprobs.rs
index 66c7a8fad9b8..1de2cdc10015 100644
--- a/rust/src/server/src/routes/openai/utils/logprobs.rs
+++ b/rust/src/server/src/routes/openai/utils/logprobs.rs
@@ -369,4 +369,51 @@ mod tests {
         assert_eq!(chat_top_logprobs_len(1), 1);
         assert_eq!(chat_top_logprobs_len(-1), 3);
     }
+
+    #[test]
+    fn chat_logprobs_clamps_nan_and_negative_infinity_for_json() {
+        let mut logprobs = sample_logprobs();
+        logprobs.positions[0].entries[0].logprob = f32::NAN;
+        logprobs.positions[0].entries[0].rank = 0;
+        logprobs.positions[0].entries[1].logprob = f32::NEG_INFINITY;
+
+        let response = decoded_logprobs_to_openai_chat(&logprobs, -1, false).unwrap();
+        let json = serde_json::to_string_pretty(&response).unwrap();
+        expect_test::expect![[r#"
+            {
+              "content": [
+                {
+                  "token": "A",
+                  "logprob": -9999.0,
+                  "bytes": [
+                    65
+                  ],
+                  "top_logprobs": [
+                    {
+                      "token": "A",
+                      "logprob": -9999.0,
+                      "bytes": [
+                        65
+                      ]
+                    },
+                    {
+                      "token": "B",
+                      "logprob": -9999.0,
+                      "bytes": [
+                        66
+                      ]
+                    },
+                    {
+                      "token": "C",
+                      "logprob": -2.0,
+                      "bytes": [
+                        67
+                      ]
+                    }
+                  ]
+                }
+              ]
+            }"#]]
+        .assert_eq(&json);
+    }
 }
diff --git a/rust/src/server/src/routes/openai/utils/structured_outputs.rs b/rust/src/server/src/routes/openai/utils/structured_outputs.rs
index 849870f32aaa..864e00818beb 100644
--- a/rust/src/server/src/routes/openai/utils/structured_outputs.rs
+++ b/rust/src/server/src/routes/openai/utils/structured_outputs.rs
@@ -3,6 +3,7 @@
 
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use thiserror_ext::AsReport as _;
 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;
 
 use crate::error::ApiError;
@@ -80,12 +81,12 @@ pub fn convert_from_response_format(
     let Some(fmt) = response_format else {
         return Ok(None);
     };
-    match fmt {
-        ResponseFormat::Text => Ok(None),
-        ResponseFormat::JsonObject => Ok(Some(StructuredOutputsParams::json_object())),
-        ResponseFormat::JsonSchema { json_schema } => Ok(Some(StructuredOutputsParams::json(
-            json_schema.schema.clone(),
-        ))),
+    let params = match fmt {
+        ResponseFormat::Text => return Ok(None),
+        ResponseFormat::JsonObject => StructuredOutputsParams::json_object(),
+        ResponseFormat::JsonSchema { json_schema } => {
+            StructuredOutputsParams::json(json_schema.schema.clone())
+        }
         ResponseFormat::StructuralTag { .. } => {
             // The Python frontend dumps the entire response_format object (including the
             // `type` field) as a JSON string for the engine-core backend.
@@ -95,9 +96,13 @@ pub fn convert_from_response_format(
                     Some("response_format"),
                 )
             })?;
-            Ok(Some(StructuredOutputsParams::structural_tag(tag_json)))
+            StructuredOutputsParams::structural_tag(tag_json)
         }
-    }
+    };
+    params.validate().map_err(|error| {
+        ApiError::invalid_request(error.to_report_string(), Some("response_format"))
+    })?;
+    Ok(Some(params))
 }
 
 /// Convert raw `response_format` and/or `structured_outputs` JSON blobs into
diff --git a/rust/src/server/src/routes/openai/utils/types.rs b/rust/src/server/src/routes/openai/utils/types.rs
index e9aab2d2b297..85a891341c8e 100644
--- a/rust/src/server/src/routes/openai/utils/types.rs
+++ b/rust/src/server/src/routes/openai/utils/types.rs
@@ -39,6 +39,36 @@ where
     }
 }
 
+/// Effort level for reasoning models.
+///
+/// Fixed OpenAI HTTP request grades. Request conversion maps these names to
+/// the model-independent `vllm_chat::EffortValue` before renderer validation.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum ReasoningEffort {
+    None,
+    Minimal,
+    Low,
+    Medium,
+    High,
+    XHigh,
+    Max,
+}
+
+impl ReasoningEffort {
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Self::None => "none",
+            Self::Minimal => "minimal",
+            Self::Low => "low",
+            Self::Medium => "medium",
+            Self::High => "high",
+            Self::XHigh => "xhigh",
+            Self::Max => "max",
+        }
+    }
+}
+
 // ============================================================================
 // String/Array Utilities
 // ============================================================================
@@ -380,9 +410,17 @@ impl ToolReference {
 // Chat Messages
 // ============================================================================
 
+/// One chat message, tagged by `role`.
+///
+/// The derived serde impls use `remote = "Self"` and cover only the standard
+/// roles; the trait impls below dispatch on the role tag and handle
+/// [`ChatMessage::Custom`] separately. A `#[serde(untagged)]` fallback variant
+/// would also accept malformed standard messages, such as a `tool` message
+/// without `tool_call_id`, and replace their field errors with serde's generic
+/// untagged-enum error.
 #[serde_with::skip_serializing_none]
 #[derive(Debug, Clone, Deserialize, Serialize)]
-#[serde(tag = "role")]
+#[serde(tag = "role", remote = "Self")]
 pub enum ChatMessage {
     #[serde(rename = "system")]
     System {
@@ -417,6 +455,67 @@ pub enum ChatMessage {
         tools: Option<Vec<Tool>>,
         name: Option<String>,
     },
+    /// Message with a role outside the OpenAI set, such as `root`. Chat
+    /// templates receive the role string unchanged.
+    #[serde(skip)]
+    Custom {
+        role: String,
+        content: MessageContent,
+    },
+}
+
+/// Role tags of the standard [`ChatMessage`] variants.
+const STANDARD_CHAT_ROLES: &[&str] = &[
+    "system",
+    "user",
+    "assistant",
+    "tool",
+    "function",
+    "developer",
+];
+
+/// Wire shape of [`ChatMessage::Custom`].
+#[derive(Deserialize, Serialize)]
+struct CustomChatMessage<R, C> {
+    role: R,
+    content: C,
+}
+
+impl Serialize for ChatMessage {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        match self {
+            Self::Custom { role, content } => {
+                CustomChatMessage { role, content }.serialize(serializer)
+            }
+            _ => Self::serialize(self, serializer),
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for ChatMessage {
+    /// Dispatch on the role tag: standard roles keep their derived errors,
+    /// and any other string role becomes [`ChatMessage::Custom`].
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let value = Value::deserialize(deserializer)?;
+        let is_custom = value
+            .get("role")
+            .and_then(Value::as_str)
+            .is_some_and(|role| !STANDARD_CHAT_ROLES.contains(&role));
+
+        if is_custom {
+            let CustomChatMessage { role, content } =
+                CustomChatMessage::deserialize(value).map_err(serde::de::Error::custom)?;
+            Ok(Self::Custom { role, content })
+        } else {
+            Self::deserialize(value).map_err(serde::de::Error::custom)
+        }
+    }
 }
 
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
diff --git a/rust/src/server/src/routes/openai/utils/validated_json.rs b/rust/src/server/src/routes/openai/utils/validated_json.rs
index 1ddec3a9f521..a9b68bc3b407 100644
--- a/rust/src/server/src/routes/openai/utils/validated_json.rs
+++ b/rust/src/server/src/routes/openai/utils/validated_json.rs
@@ -5,7 +5,6 @@
 //! Variation of https://github.com/lightseekorg/smg/blob/main/crates/protocols/src/validated.rs
 
 use axum::Json;
-use axum::extract::rejection::JsonRejection;
 use axum::extract::{FromRequest, Request};
 use serde::de::DeserializeOwned;
 use validator::Validate;
@@ -30,9 +29,7 @@ where
     type Rejection = ApiError;
 
     async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
-        let Json(mut data) = Json::<T>::from_request(req, state)
-            .await
-            .map_err(|err: JsonRejection| ApiError::json_parse_error(err.body_text()))?;
+        let Json(mut data) = Json::<T>::from_request(req, state).await?;
 
         data.normalize();
 
diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs
index a4ae0298e91d..2a6f70ec4b6e 100644
--- a/rust/src/server/src/routes/render.rs
+++ b/rust/src/server/src/routes/render.rs
@@ -105,6 +105,7 @@ fn lower_render_request(
         kv_transfer_params: None,
         ec_transfer_params: None,
         content_parts: None,
+        return_token_ids: None,
         other: Default::default(),
     };
     validate_generate_request(&request, &state.served_model_names)?;
@@ -123,7 +124,7 @@ async fn render_chat(
     let chat_request = lower_chat_request(body, &model_resolution(&state), request_context)?;
     let (text_request, _) = state
         .chat
-        .prepare(chat_request)
+        .prepare(chat_request, &state.text)
         .await
         .map_err(|error| ApiError::invalid_request(error.to_report_string(), None))?;
     Ok(Json(lower_render_request(
diff --git a/rust/src/server/src/routes/sleep.rs b/rust/src/server/src/routes/sleep.rs
index 4bf9981402ca..d64a0850b6e4 100644
--- a/rust/src/server/src/routes/sleep.rs
+++ b/rust/src/server/src/routes/sleep.rs
@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use axum::Json;
 use axum::extract::rejection::QueryRejection;
-use axum::extract::{Query, State};
+use axum::extract::{Query, RawQuery, State};
 use axum::http::StatusCode;
 use serde::{Deserialize, Serialize};
 use vllm_engine_core_client::protocol::utility::PauseMode;
@@ -27,12 +27,6 @@ pub(crate) struct SleepParams {
     mode: PauseMode,
 }
 
-#[derive(Debug, Default, Deserialize)]
-pub(crate) struct WakeUpParams {
-    #[serde(default)]
-    tags: Option<Vec<String>>,
-}
-
 const fn default_sleep_level() -> u32 {
     1
 }
@@ -57,14 +51,35 @@ pub async fn sleep(
     Ok(StatusCode::OK)
 }
 
+/// Release KV cache memory while keeping model weights resident.
+pub async fn release_kv_cache_memory(
+    State(state): State<Arc<AppState>>,
+) -> Result<StatusCode, ApiError> {
+    state
+        .engine_core_client()
+        .release_kv_cache_memory()
+        .await
+        .map_err(|error| utility_call_error("release_kv_cache_memory", error))?;
+
+    Ok(StatusCode::OK)
+}
+
 /// Wake the engine from sleep mode.
 pub async fn wake_up(
     State(state): State<Arc<AppState>>,
-    Query(params): Query<WakeUpParams>,
+    RawQuery(query): RawQuery,
 ) -> Result<StatusCode, ApiError> {
+    let tags = query.and_then(|query| {
+        let tags = url::form_urlencoded::parse(query.as_bytes())
+            .filter(|(key, _)| key == "tags")
+            .map(|(_, value)| value.into_owned())
+            .collect::<Vec<_>>();
+        (!tags.is_empty()).then_some(tags)
+    });
+
     state
         .engine_core_client()
-        .wake_up(params.tags)
+        .wake_up(tags)
         .await
         .map_err(|error| utility_call_error("wake_up", error))?;
 
diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs
index 210ae044a026..5c1fff2cae76 100644
--- a/rust/src/server/src/routes/tests.rs
+++ b/rust/src/server/src/routes/tests.rs
@@ -54,8 +54,7 @@ use zeromq::{DealerSocket, PushSocket, ZmqMessage};
 
 use super::{
     build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora,
-    build_router_with_scale_out_endpoints, parse_scale_out_endpoints_flag,
-    render::build_router as build_render_router,
+    build_router_with_scale_out_endpoints, render::build_router as build_render_router,
 };
 use crate::config::{ApiServerOptions, CorsConfig, LoraModulePath};
 use crate::lora::LoadLoraError;
@@ -518,6 +517,7 @@ impl ChatBackend for FakeChatBackend {
             self.tokenizer(),
             options.tool_call_parser,
             options.reasoning_parser,
+            options.tool_strict_level,
         )?))
     }
 }
@@ -541,6 +541,7 @@ impl ChatRenderer for FakeChatBackend {
         }
         Ok(vllm_chat::RenderedPrompt {
             prompt: Prompt::Text(prompt),
+            media_order: None,
             effective_template_kwargs: Default::default(),
         })
     }
@@ -554,7 +555,8 @@ fn render_fake_message_content(
         ChatMessage::System { content }
         | ChatMessage::Developer { content, .. }
         | ChatMessage::User { content }
-        | ChatMessage::ToolResponse { content, .. } => render_fake_content(content, placeholder),
+        | ChatMessage::ToolResponse { content, .. }
+        | ChatMessage::Custom { content, .. } => render_fake_content(content, placeholder),
         ChatMessage::Assistant { .. } => message.text_content(),
     }
 }
@@ -731,23 +733,36 @@ async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router {
     )
 }
 
-#[test]
-fn scale_out_flag_accepts_unset_empty_zero_one_and_whitespace() {
-    assert_eq!(parse_scale_out_endpoints_flag(None), Ok(false));
-    assert_eq!(parse_scale_out_endpoints_flag(Some("")), Ok(false));
-    assert_eq!(parse_scale_out_endpoints_flag(Some("0")), Ok(false));
-    assert_eq!(parse_scale_out_endpoints_flag(Some("1")), Ok(true));
-    assert_eq!(parse_scale_out_endpoints_flag(Some(" 1 ")), Ok(true));
-}
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn scale_out_generate_route_follows_api_server_options() {
+    let (chat, _engine_task) = test_models_with_engine_outputs_and_backend(
+        b"engine-scale-out-api-server-options",
+        default_stream_output_specs(),
+        Arc::new(FakeChatBackend::new()),
+    )
+    .await;
+    let state = Arc::new(
+        AppState::new(vec!["model".to_string()], chat).with_api_server_options(ApiServerOptions {
+            enable_scale_out: true,
+            ..Default::default()
+        }),
+    );
+    let mut app = build_router(state);
 
-#[test]
-fn scale_out_flag_rejects_values_other_than_zero_or_one() {
-    for value in ["-1", "2", "01", "+1", "invalid", " "] {
-        assert_eq!(
-            parse_scale_out_endpoints_flag(Some(value)),
-            Err("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1".to_string())
-        );
-    }
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/inference/v1/generate")
+                .header("content-type", "application/json")
+                .body(Body::from("{}"))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_ne!(response.status(), StatusCode::NOT_FOUND);
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -2040,6 +2055,56 @@ async fn version_returns_engine_vllm_version() {
     );
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn kv_event_sources_reports_each_ready_engine() {
+    let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse {
+        data_parallel_rank: 1,
+        kv_events_config: Some(
+            vllm_engine_core_client::protocol::handshake::KvEventsConfig {
+                enable_kv_cache_events: true,
+                publisher: "zmq".to_string(),
+                endpoint: "tcp://10.0.0.2:41233".to_string(),
+                replay_endpoint: None,
+                buffer_steps: 10_000,
+                hwm: 100_000,
+                max_queue_size: 100_000,
+                topic: "kv".to_string(),
+            },
+        ),
+        ..default_ready_response()
+    };
+    let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;
+    let response = app
+        .call(
+            Request::builder()
+                .uri("/kv_event_sources")
+                .body(Body::empty())
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::OK);
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert_eq!(
+        json,
+        json!({
+            "1": {
+                "enable_kv_cache_events": true,
+                "publisher": "zmq",
+                "endpoint": "tcp://10.0.0.2:41233",
+                "replay_endpoint": null,
+                "buffer_steps": 10000,
+                "hwm": 100000,
+                "max_queue_size": 100000,
+                "topic": "kv",
+            }
+        })
+    );
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn server_info_endpoint_is_dev_mode_only() {
@@ -2723,6 +2788,49 @@ async fn non_stream_chat_returns_json_response() {
     assert!(!message.contains_key("tool_calls"), "{json}");
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn openai_watermarking_defaults_and_opt_out_reach_engine_core() {
+    for endpoint in ["/v1/chat/completions", "/v1/completions"] {
+        for stream in [false, true] {
+            for watermarking in [None, Some(true), Some(false)] {
+                let (mut app, engine_task) = test_app_with_backend_and_engine_request_check(
+                    Arc::new(FakeChatBackend::new()),
+                    move |request| {
+                        let params = request.sampling_params.as_ref().expect("sampling params");
+                        assert_eq!(params.watermarking, watermarking.unwrap_or(true));
+                    },
+                )
+                .await;
+
+                let mut body = json!({"stream": stream, "max_tokens": 2});
+                if endpoint == "/v1/chat/completions" {
+                    body["messages"] = json!([{"role": "user", "content": "hi"}]);
+                } else {
+                    body["prompt"] = json!("hi");
+                }
+                if let Some(watermarking) = watermarking {
+                    body["watermarking"] = json!(watermarking);
+                }
+                let response = app
+                    .call(
+                        Request::builder()
+                            .method("POST")
+                            .uri(endpoint)
+                            .header("content-type", "application/json")
+                            .body(Body::from(body.to_string()))
+                            .expect("build request"),
+                    )
+                    .await
+                    .expect("call app");
+                assert_eq!(response.status(), StatusCode::OK);
+                to_bytes(response.into_body(), usize::MAX).await.expect("drain response");
+                engine_task.await.expect("mock engine task");
+            }
+        }
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn non_stream_chat_image_url_reaches_engine_mm_features() {
@@ -5749,7 +5857,46 @@ async fn sleep_route_uses_python_compatible_default_query_values() {
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
-async fn wake_up_route_without_tags_sends_none() {
+async fn release_kv_cache_memory_route_sends_expected_utility_call() {
+    let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let utility = recv_engine_message(dealer).await;
+            assert_eq!(utility[0].as_ref(), &[0x03]);
+
+            let payload = decode_value(&utility[1]).expect("decode utility payload");
+            let array = payload.as_array().expect("utility payload array");
+            let call_id = array[1].as_u64().expect("call id");
+
+            assert_eq!(array[2], Value::from("release_kv_cache_memory"));
+            assert_eq!(array[3], Value::Array(Vec::new()));
+
+            send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
+        })
+    })
+    .await;
+
+    let response = app
+        .clone()
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/release_kv_cache_memory")
+                .body(Body::empty())
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    let status = response.status();
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
+    assert!(body.is_empty());
+    engine_task.await.expect("mock engine task");
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn wake_up_route_without_tags_accepts_bool_result() {
     let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
         boxed_test_future(async move {
             let utility = recv_engine_message(dealer).await;
@@ -5762,7 +5909,7 @@ async fn wake_up_route_without_tags_sends_none() {
             assert_eq!(array[2], Value::from("wake_up"));
             assert_eq!(array[3], Value::Array(vec![Value::Nil]));
 
-            send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
+            send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
         })
     })
     .await;
@@ -5786,6 +5933,51 @@ async fn wake_up_route_without_tags_sends_none() {
     engine_task.await.expect("mock engine task");
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn wake_up_route_accepts_repeated_tags() {
+    let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let utility = recv_engine_message(dealer).await;
+            assert_eq!(utility[0].as_ref(), &[0x03]);
+
+            let payload = decode_value(&utility[1]).expect("decode utility payload");
+            let array = payload.as_array().expect("utility payload array");
+            let call_id = array[1].as_u64().expect("call id");
+
+            assert_eq!(array[2], Value::from("wake_up"));
+            assert_eq!(
+                array[3],
+                Value::Array(vec![Value::Array(vec![
+                    Value::from("weights"),
+                    Value::from("kv_cache"),
+                ])])
+            );
+
+            send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
+        })
+    })
+    .await;
+
+    let response = app
+        .clone()
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/wake_up?tags=weights&tags%5B%5D=ignored&tags=kv_cache")
+                .body(Body::empty())
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    let status = response.status();
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
+    assert!(body.is_empty());
+    engine_task.await.expect("mock engine task");
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn is_sleeping_route_returns_json_payload() {
@@ -6104,6 +6296,7 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
     for (method, uri) in [
         ("GET", "/is_sleeping"),
         ("POST", "/sleep"),
+        ("POST", "/release_kv_cache_memory"),
         ("POST", "/wake_up"),
         ("GET", "/is_paused"),
         ("POST", "/pause"),
@@ -6113,6 +6306,13 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
         ("POST", "/reset_prefix_cache"),
         ("POST", "/reset_mm_cache"),
         ("POST", "/reset_encoder_cache"),
+        ("POST", "/init_weight_transfer_engine"),
+        ("POST", "/start_weight_update"),
+        ("POST", "/start_draft_weight_update"),
+        ("POST", "/update_weights"),
+        ("POST", "/finish_weight_update"),
+        ("POST", "/update_weight_version"),
+        ("GET", "/weight_info"),
     ] {
         let response = app
             .clone()
@@ -6132,6 +6332,248 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
     engine_task.abort_and_join().await;
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn weight_transfer_routes_support_the_http_training_lifecycle() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let mut calls = Vec::new();
+            for _ in 0..10 {
+                let utility = recv_engine_message(dealer).await;
+                assert_eq!(utility[0].as_ref(), &[0x03]);
+                let payload = decode_value(&utility[1]).expect("decode utility");
+                let array = payload.as_array().expect("utility array");
+                let call_id = array[1].as_u64().expect("call id");
+                let args: serde_json::Value =
+                    rmpv::ext::from_value(array[3].clone()).expect("decode args");
+                calls.push(format!("{} {args}", array[2].as_str().expect("method")));
+                let result = if array[2] == Value::from("get_weight_version") {
+                    utility_result_value("step-2")
+                } else if array[2] == Value::from("collective_rpc") {
+                    utility_result_value(vec![(), ()])
+                } else {
+                    utility_none_result()
+                };
+                send_outputs(push, utility_outputs(call_id, result)).await;
+            }
+            expect_test::expect![[r#"
+                [
+                    "collective_rpc [\"init_weight_transfer_engine\",null,[],{\"init_info\":{\"master_address\":\"127.0.0.1\",\"master_port\":29500,\"rank_offset\":1,\"world_size\":3}}]",
+                    "collective_rpc [\"start_weight_update\",null,[],{}]",
+                    "collective_rpc [\"update_weights\",null,[],{\"update_info\":{\"names\":[\"model.embed_tokens.weight\"],\"dtype_names\":[\"bfloat16\"],\"shapes\":[[32,16]],\"is_checkpoint_format\":true}}]",
+                    "collective_rpc [\"update_weights\",null,[],{\"update_info\":[{\"rank\":0,\"ipc_handles_pickled\":\"opaque-handle-0\"},{\"rank\":1,\"ipc_handles_pickled\":\"opaque-handle-1\"}]}]",
+                    "collective_rpc [\"finish_weight_update\",null,[],{}]",
+                    "set_weight_version [\"step-1\"]",
+                    "collective_rpc [\"start_draft_weight_update\",null,[],{}]",
+                    "collective_rpc [\"finish_weight_update\",null,[],{}]",
+                    "set_weight_version [\"step-2\"]",
+                    "get_weight_version []",
+                ]
+            "#]].assert_debug_eq(&calls);
+        })
+    })
+    .await;
+
+    let mut responses = Vec::new();
+    for (method, path, body) in [
+        (
+            "POST",
+            "/init_weight_transfer_engine",
+            Some(json!({
+                "init_info": {"master_address": "127.0.0.1", "master_port": 29500,
+                    "rank_offset": 1, "world_size": 3}
+            })),
+        ),
+        ("POST", "/start_weight_update", None),
+        (
+            "POST",
+            "/update_weights",
+            Some(json!({
+                "update_info": {"names": ["model.embed_tokens.weight"],
+                    "dtype_names": ["bfloat16"], "shapes": [[32, 16]],
+                    "is_checkpoint_format": true}
+            })),
+        ),
+        (
+            "POST",
+            "/update_weights",
+            Some(json!({
+                "update_info": [{"rank": 0, "ipc_handles_pickled": "opaque-handle-0"},
+                    {"rank": 1, "ipc_handles_pickled": "opaque-handle-1"}]
+            })),
+        ),
+        (
+            "POST",
+            "/finish_weight_update",
+            Some(json!({"weight_version": "step-1"})),
+        ),
+        ("POST", "/start_draft_weight_update", None),
+        ("POST", "/finish_weight_update", None),
+        (
+            "POST",
+            "/update_weight_version",
+            Some(json!({"new_version": "step-2"})),
+        ),
+        ("GET", "/weight_info", None),
+    ] {
+        let mut request = Request::builder().method(method).uri(path);
+        let body = match body {
+            Some(body) => {
+                request = request.header("content-type", "application/json");
+                Body::from(body.to_string())
+            }
+            None => Body::empty(),
+        };
+        let response =
+            app.call(request.body(body).expect("build request")).await.expect("call app");
+        assert_eq!(response.status(), StatusCode::OK, "{path}");
+        let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+        responses.push(
+            serde_json::from_slice::<serde_json::Value>(&body)
+                .expect("decode json")
+                .to_string(),
+        );
+    }
+    engine_task.await.expect("mock engine task");
+    expect_test::expect![[r#"
+        [
+            "{\"message\":\"Weight transfer initialized\"}",
+            "{\"message\":\"Weight update started\"}",
+            "{\"message\":\"Weights updated\"}",
+            "{\"message\":\"Weights updated\"}",
+            "{\"message\":\"Weight update finished\"}",
+            "{\"message\":\"Draft weight update started\"}",
+            "{\"message\":\"Weight update finished\"}",
+            "{\"success\":true,\"new_version\":\"step-2\"}",
+            "{\"weight_version\":\"step-2\"}",
+        ]
+    "#]]
+    .assert_debug_eq(&responses);
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn weight_transfer_routes_reject_invalid_payloads_before_engine_calls() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, _push| {
+        boxed_test_future(async move {
+            let message = recv_engine_message(dealer).await;
+            panic!("invalid payload reached engine: {message:?}");
+        })
+    })
+    .await;
+
+    for (path, body) in [
+        ("/init_weight_transfer_engine", "{"),
+        ("/init_weight_transfer_engine", "{}"),
+        ("/init_weight_transfer_engine", r#"{"init_info":null}"#),
+        ("/init_weight_transfer_engine", r#"{"init_info":[]}"#),
+        ("/update_weights", "{"),
+        ("/update_weights", "{}"),
+        ("/update_weights", r#"{"update_info":null}"#),
+        ("/update_weights", r#"{"update_info":1}"#),
+        ("/update_weights", r#"{"update_info":[{},null]}"#),
+        ("/finish_weight_update", "{"),
+        ("/finish_weight_update", r#"{"weight_version":1}"#),
+        ("/update_weight_version", "{}"),
+        ("/update_weight_version", r#"{"new_version":null}"#),
+    ] {
+        let response = app
+            .call(
+                Request::builder()
+                    .method("POST")
+                    .uri(path)
+                    .header("content-type", "application/json")
+                    .body(Body::from(body))
+                    .expect("build request"),
+            )
+            .await
+            .expect("call app");
+        assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{path}: {body}");
+    }
+    engine_task.abort_and_join().await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn weight_transfer_routes_do_not_commit_version_when_finish_fails() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let utility = recv_engine_message(dealer).await;
+            let payload = decode_value(&utility[1]).expect("decode utility");
+            let array = payload.as_array().expect("utility array");
+            assert_eq!(array[2], Value::from("collective_rpc"));
+            assert_eq!(
+                array[3].as_array().expect("args")[0],
+                Value::from("finish_weight_update")
+            );
+            send_outputs(
+                push,
+                UtilityCallOutput {
+                    output: UtilityOutput {
+                        call_id: array[1].as_u64().expect("call id").into(),
+                        failure_message: Some("weight reload failed".to_string()),
+                        result: None,
+                    },
+                    ..Default::default()
+                }
+                .into(),
+            )
+            .await;
+            let utility = recv_engine_message(dealer).await;
+            let payload = decode_value(&utility[1]).expect("decode utility");
+            let array = payload.as_array().expect("utility array");
+            assert_eq!(array[2], Value::from("get_weight_version"));
+            send_outputs(
+                push,
+                utility_outputs(
+                    array[1].as_u64().expect("call id"),
+                    utility_result_value("step-0"),
+                ),
+            )
+            .await;
+        })
+    })
+    .await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/finish_weight_update")
+                .header("content-type", "application/json")
+                .body(Body::from(r#"{"weight_version":"step-1"}"#))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    let error: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert!(
+        error["error"]["message"]
+            .as_str()
+            .expect("error message")
+            .contains("weight reload failed")
+    );
+
+    let response = app
+        .call(
+            Request::builder()
+                .uri("/weight_info")
+                .body(Body::empty())
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::OK);
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    assert_eq!(
+        serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
+        json!({"weight_version": "step-0"})
+    );
+    engine_task.await.expect("mock engine task");
+}
+
 // ========================= Stop string tests =========================
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs
index 35a027c1c0f7..d907073c87ae 100644
--- a/rust/src/server/src/routes/tokenize/types.rs
+++ b/rust/src/server/src/routes/tokenize/types.rs
@@ -54,7 +54,9 @@ impl TokenizeCompletionRequest {
             add_special_tokens: self.add_special_tokens,
             data_parallel_rank: None,
             session_id: None,
-            reasoning_parser_kwargs: None,
+            kv_hints: None,
+            reasoning_parser_kwargs: Default::default(),
+            reasoning_ended: None,
             lora_request: None,
             arrival_time: None,
         }
diff --git a/rust/src/server/src/routes/weight_transfer.rs b/rust/src/server/src/routes/weight_transfer.rs
new file mode 100644
index 000000000000..df55e25593af
--- /dev/null
+++ b/rust/src/server/src/routes/weight_transfer.rs
@@ -0,0 +1,190 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::sync::Arc;
+
+use axum::Json;
+use axum::extract::State;
+use axum::extract::rejection::JsonRejection;
+use serde::{Deserialize, Serialize};
+use serde_json::Value as JsonValue;
+
+use crate::error::{ApiError, invalid_request};
+use crate::state::AppState;
+use crate::utils::utility_call_error;
+
+#[derive(Deserialize)]
+pub(crate) struct InitWeightTransferRequest {
+    init_info: Option<JsonValue>,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct UpdateWeightsRequest {
+    update_info: Option<JsonValue>,
+}
+
+#[derive(Default, Deserialize)]
+pub(crate) struct FinishWeightUpdateRequest {
+    weight_version: Option<String>,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct UpdateWeightVersionRequest {
+    new_version: String,
+}
+
+#[derive(Serialize)]
+pub(crate) struct MessageResponse {
+    message: &'static str,
+}
+
+#[derive(Serialize)]
+pub(crate) struct UpdateWeightVersionResponse {
+    success: bool,
+    new_version: String,
+}
+
+#[derive(Serialize)]
+pub(crate) struct WeightInfoResponse {
+    weight_version: String,
+}
+
+/// Initialize weight transfer for RL training.
+pub async fn init_weight_transfer_engine(
+    State(state): State<Arc<AppState>>,
+    body: Result<Json<InitWeightTransferRequest>, JsonRejection>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    let Json(body) = body?;
+    let init_info = body.init_info.filter(JsonValue::is_object).ok_or_else(|| {
+        invalid_request!(
+            param = Some("init_info"),
+            "'init_info' must be a JSON object"
+        )
+    })?;
+
+    state
+        .engine_core_client()
+        .init_weight_transfer_engine(init_info)
+        .await
+        .map_err(|error| utility_call_error("init_weight_transfer_engine", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Weight transfer initialized",
+    }))
+}
+
+/// Start a weight update transaction.
+pub async fn start_weight_update(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    state
+        .engine_core_client()
+        .start_weight_update()
+        .await
+        .map_err(|error| utility_call_error("start_weight_update", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Weight update started",
+    }))
+}
+
+/// Start a draft-model weight update transaction.
+pub async fn start_draft_weight_update(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    state
+        .engine_core_client()
+        .start_draft_weight_update()
+        .await
+        .map_err(|error| utility_call_error("start_draft_weight_update", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Draft weight update started",
+    }))
+}
+
+/// Update model weights with backend-specific transfer metadata.
+pub async fn update_weights(
+    State(state): State<Arc<AppState>>,
+    body: Result<Json<UpdateWeightsRequest>, JsonRejection>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    let Json(body) = body?;
+    let update_info = body
+        .update_info
+        .filter(|info| {
+            info.is_object()
+                || info.as_array().is_some_and(|items| items.iter().all(JsonValue::is_object))
+        })
+        .ok_or_else(|| {
+            invalid_request!(
+                param = Some("update_info"),
+                "'update_info' must be a JSON object or a list of per-worker JSON objects",
+            )
+        })?;
+
+    state
+        .engine_core_client()
+        .update_weights(update_info)
+        .await
+        .map_err(|error| utility_call_error("update_weights", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Weights updated",
+    }))
+}
+
+/// Finish a weight update transaction and optionally set the weight version.
+pub async fn finish_weight_update(
+    State(state): State<Arc<AppState>>,
+    body: Result<Option<Json<FinishWeightUpdateRequest>>, JsonRejection>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    // HTTPVLLMWeightSyncClient omits the body when no version is supplied.
+    let request = body?.map(|Json(request)| request).unwrap_or_default();
+
+    let client = state.engine_core_client();
+    client
+        .finish_weight_update()
+        .await
+        .map_err(|error| utility_call_error("finish_weight_update", error))?;
+    if let Some(version) = request.weight_version {
+        client
+            .set_weight_version(&version)
+            .await
+            .map_err(|error| utility_call_error("set_weight_version", error))?;
+    }
+
+    Ok(Json(MessageResponse {
+        message: "Weight update finished",
+    }))
+}
+
+/// Update the weight version identifier.
+pub async fn update_weight_version(
+    State(state): State<Arc<AppState>>,
+    body: Result<Json<UpdateWeightVersionRequest>, JsonRejection>,
+) -> Result<Json<UpdateWeightVersionResponse>, ApiError> {
+    let Json(body) = body?;
+    state
+        .engine_core_client()
+        .set_weight_version(&body.new_version)
+        .await
+        .map_err(|error| utility_call_error("set_weight_version", error))?;
+
+    Ok(Json(UpdateWeightVersionResponse {
+        success: true,
+        new_version: body.new_version,
+    }))
+}
+
+/// Get current weight metadata.
+pub async fn weight_info(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<WeightInfoResponse>, ApiError> {
+    let weight_version = state
+        .engine_core_client()
+        .get_weight_version()
+        .await
+        .map_err(|error| utility_call_error("get_weight_version", error))?;
+
+    Ok(Json(WeightInfoResponse { weight_version }))
+}
diff --git a/rust/src/server/src/utils.rs b/rust/src/server/src/utils.rs
index b58e89fd3f07..c3a8449617ab 100644
--- a/rust/src/server/src/utils.rs
+++ b/rust/src/server/src/utils.rs
@@ -2,9 +2,13 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use std::collections::HashMap;
-use std::time::{SystemTime, UNIX_EPOCH};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
+use axum::BoxError;
 use axum::http::HeaderMap;
+use axum::response::sse::{Event, KeepAlive, Sse};
+use axum::response::{IntoResponse, Response};
+use futures::Stream;
 use serde_json::Value;
 use thiserror_ext::AsReport;
 use uuid::Uuid;
@@ -27,6 +31,27 @@ pub fn unix_timestamp() -> u64 {
         .unwrap_or_default()
 }
 
+/// Build a streaming SSE response, optionally sending keep-alive comments
+/// while the stream is idle.
+///
+/// With `keep_alive_interval` set, a `: keep-alive` comment is sent whenever no
+/// event has been produced for that long (queued, prefill, between tokens, or
+/// while a parser buffers a long tool call). SSE parsers ignore comments, but
+/// they still count as bytes for read timeouts in reverse proxies and clients.
+pub fn sse_response<S, E>(stream: S, keep_alive_interval: Option<Duration>) -> Response
+where
+    S: Stream<Item = Result<Event, E>> + Send + 'static,
+    E: Into<BoxError>,
+{
+    let sse = Sse::new(stream);
+    match keep_alive_interval {
+        Some(interval) => sse
+            .keep_alive(KeepAlive::new().interval(interval).text("keep-alive"))
+            .into_response(),
+        None => sse.into_response(),
+    }
+}
+
 /// Construct an API error for a failed utility call to the engine core.
 pub fn utility_call_error(method: &str, error: impl AsReport) -> ApiError {
     ApiError::server_error(format!("failed to call {method}: {}", error.as_report()))
@@ -157,3 +182,84 @@ pub fn resolve_base_request_id(
         id
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use std::convert::Infallible;
+
+    use futures::{StreamExt as _, stream};
+    use tokio::time::{Instant, sleep};
+
+    use super::*;
+
+    /// Collect `(elapsed seconds, frame)` pairs from a response whose event
+    /// stream idles for 20 seconds between its two events.
+    async fn idle_stream_frames(keep_alive_interval: Option<Duration>) -> Vec<(u64, String)> {
+        let events = stream::iter([Ok::<_, Infallible>(Event::default().data("first"))]).chain(
+            stream::once(async {
+                sleep(Duration::from_secs(20)).await;
+                Ok(Event::default().data("second"))
+            }),
+        );
+        let start = Instant::now();
+        sse_response(events, keep_alive_interval)
+            .into_body()
+            .into_data_stream()
+            .map(|frame| {
+                let frame = frame.expect("read frame");
+                (
+                    start.elapsed().as_secs(),
+                    String::from_utf8(frame.to_vec()).expect("utf-8 frame"),
+                )
+            })
+            .collect()
+            .await
+    }
+
+    /// An idle gap is filled with `: keep-alive` comments at the configured
+    /// interval, and real events still pass through in order.
+    #[tokio::test(start_paused = true)]
+    async fn sse_response_sends_keep_alive_comments_while_idle() {
+        let frames = idle_stream_frames(Some(Duration::from_secs(8))).await;
+        expect_test::expect![[r#"
+            [
+                (
+                    0,
+                    "data: first\n\n",
+                ),
+                (
+                    8,
+                    ": keep-alive\n\n",
+                ),
+                (
+                    16,
+                    ": keep-alive\n\n",
+                ),
+                (
+                    20,
+                    "data: second\n\n",
+                ),
+            ]
+        "#]]
+        .assert_debug_eq(&frames);
+    }
+
+    /// Without an interval, an idle gap produces no extra bytes.
+    #[tokio::test(start_paused = true)]
+    async fn sse_response_without_keep_alive_interval_sends_no_comments() {
+        let frames = idle_stream_frames(None).await;
+        expect_test::expect![[r#"
+            [
+                (
+                    0,
+                    "data: first\n\n",
+                ),
+                (
+                    20,
+                    "data: second\n\n",
+                ),
+            ]
+        "#]]
+        .assert_debug_eq(&frames);
+    }
+}
diff --git a/rust/src/text/Cargo.toml b/rust/src/text/Cargo.toml
index 7bda7f976e1c..36152d7cc4a1 100644
--- a/rust/src/text/Cargo.toml
+++ b/rust/src/text/Cargo.toml
@@ -16,6 +16,7 @@ reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 serde_with.workspace = true
+tempfile.workspace = true
 thiserror.workspace = true
 thiserror-ext.workspace = true
 tracing.workspace = true
@@ -28,7 +29,6 @@ vllm-tokenizer.workspace = true
 expect-test.workspace = true
 futures.workspace = true
 serial_test.workspace = true
-tempfile.workspace = true
 tokio.workspace = true
 vllm-llm = { workspace = true, features = ["test-util"] }
 vllm-tokenizer = { workspace = true, features = ["test-utils"] }
diff --git a/rust/src/text/src/backend/hf/config.rs b/rust/src/text/src/backend/hf/config.rs
index 849874bedec6..0768adf4ae5d 100644
--- a/rust/src/text/src/backend/hf/config.rs
+++ b/rust/src/text/src/backend/hf/config.rs
@@ -84,10 +84,11 @@ impl HfSpecialTokens {
 
 /// Minimal subset of `config.json` (the model's main HF config).
 ///
-/// This intentionally supports only the two layouts we currently care about in
+/// This intentionally supports only the layouts we currently care about in
 /// the Rust frontend:
 /// - pure text models that keep text metadata at the top level
-/// - composite models that expose a single nested `text_config`
+/// - composite models that expose a single nested `text_config` or its
+///   `llm_config` alias (used by Nemotron)
 ///
 /// We do not support additional entry points such as `decoder`, `generator`, or
 /// `text_encoder`.
@@ -102,6 +103,7 @@ pub struct ModelConfig {
     n_routed_experts: Option<OneOrManyExpertCount>,
     num_local_experts: Option<OneOrManyExpertCount>,
     block_configs: Vec<BlockConfig>,
+    #[serde(alias = "llm_config")]
     text_config: Option<Box<ModelConfig>>,
 }
 
@@ -188,7 +190,7 @@ impl ModelConfig {
     /// Return the config that the Rust frontend treats as the text/LLM config.
     ///
     /// This is deliberately narrower than Python/transformers: we only support
-    /// either the top-level config itself or a single nested `text_config`.
+    /// either the top-level config itself or a single nested text config.
     fn effective_text_config(&self) -> &Self {
         self.text_config.as_deref().unwrap_or(self)
     }
@@ -233,7 +235,8 @@ impl ModelConfig {
     /// config.
     ///
     /// The only intentional simplification here is how we pick the text config:
-    /// Rust only looks at the top level or `text_config`, not the broader
+    /// Rust only looks at the top level or `text_config` (including its
+    /// `llm_config` alias), not the broader
     /// transformers composite-config surface.
     fn num_experts_from_block_configs(&self) -> u32 {
         self.effective_text_config()
@@ -405,6 +408,27 @@ mod tests {
         assert_eq!(config.vocab_size().unwrap(), 151936);
     }
 
+    #[test]
+    fn model_config_uses_llm_config_for_nemotron_composite_models() {
+        let config: ModelConfig = serde_json::from_str(
+            r#"{
+                "model_type": "nemotron_h_omni",
+                "llm_config": {
+                    "model_type": "nemotron_h",
+                    "vocab_size": 131072,
+                    "n_routed_experts": 256,
+                    "eos_token_id": 2
+                }
+            }"#,
+        )
+        .unwrap();
+
+        assert_eq!(config.model_type(), Some("nemotron_h_omni"));
+        assert_eq!(config.vocab_size().unwrap(), 131072);
+        assert_eq!(config.eos_token_ids(), &[2]);
+        assert_eq!(config.num_experts(), 256);
+    }
+
     #[test]
     fn model_config_reads_top_level_eos_token_ids() {
         let single: ModelConfig = serde_json::from_str(r#"{"eos_token_id":151645}"#).unwrap();
diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs
index ae768767d99a..2f32f2cb483d 100644
--- a/rust/src/text/src/backend/hf/mod.rs
+++ b/rust/src/text/src/backend/hf/mod.rs
@@ -3,6 +3,7 @@
 
 mod config;
 mod model_files;
+mod overrides;
 
 use std::collections::BTreeSet;
 use std::sync::Arc;
@@ -18,6 +19,7 @@ pub use self::config::{
     load_tokenizer_config,
 };
 pub use self::model_files::{ResolvedModelFiles, TokenizerSource};
+pub use self::overrides::HfOverrides;
 use crate::backend::{GenerationConfigMode, SamplingHints, TextBackend};
 use crate::error::Result;
 
diff --git a/rust/src/text/src/backend/hf/model_files.rs b/rust/src/text/src/backend/hf/model_files.rs
index d0b0caecce4c..a4707672a2dd 100644
--- a/rust/src/text/src/backend/hf/model_files.rs
+++ b/rust/src/text/src/backend/hf/model_files.rs
@@ -2,11 +2,13 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use std::path::{Path, PathBuf};
+use std::sync::Arc;
 
-use hf_hub::Cache;
 use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo};
+use hf_hub::{Cache, Repo, RepoType};
 use thiserror_ext::AsReport as _;
 
+use super::HfOverrides;
 use super::config::{HfTokenizerConfig, load_tokenizer_config};
 use crate::error::{Error, Result};
 
@@ -50,23 +52,76 @@ pub struct ResolvedModelFiles {
     pub processor_config_path: Option<PathBuf>,
     pub chat_template_path: Option<PathBuf>,
     pub config_path: Option<PathBuf>,
+    /// Keep a patched config alive across clones and backend ownership.
+    config_temp_path: Option<Arc<tempfile::TempPath>>,
 }
 
 impl ResolvedModelFiles {
+    /// Materialize a merge-patched config owned by these resolved model files.
+    pub fn apply_overrides(&mut self, overrides: &HfOverrides) -> Result<()> {
+        if overrides.is_empty() {
+            return Ok(());
+        }
+        let path = materialize_config(self.config_path.as_deref(), overrides).map_err(|error| {
+            Error::Tokenizer(format!(
+                "failed to materialize HF overrides: {}",
+                error.as_report()
+            ))
+        })?;
+        self.config_path = Some(path.to_path_buf());
+        self.config_temp_path = Some(Arc::new(path));
+        Ok(())
+    }
+
     /// Resolve tokenizer/config files from a local model directory first when
     /// `model_id` points to one, otherwise consult the local HF cache and
     /// finally the Hub.
-    pub async fn new(model_id: &str) -> Result<Self> {
+    pub async fn new(model_id: &str, revision: Option<&str>) -> Result<Self> {
         if Path::new(model_id).is_dir() {
             return resolve_local_model_files(Path::new(model_id));
         }
-        if let Some(files) = resolve_cached_model_files(model_id)? {
+        let repo = Repo::with_revision(
+            model_id.to_string(),
+            RepoType::Model,
+            revision.unwrap_or("main").to_string(),
+        );
+        // TODO: Let hf-hub read HF_HUB_CACHE after upgrading to 1.0.
+        let cache = std::env::var_os("HF_HUB_CACHE")
+            .map(|path| Cache::new(path.into()))
+            .unwrap_or_else(Cache::from_env);
+        if let Some(files) = resolve_cached_model_files(&cache, &repo)? {
             return Ok(files);
         }
-        resolve_remote_model_files(model_id).await
+        if is_hf_hub_offline() {
+            return Err(Error::Tokenizer(format!(
+                "model '{model_id}' at revision '{}' has no supported tokenizer in cache '{}' \
+                 and HF_HUB_OFFLINE is enabled",
+                repo.revision(),
+                cache.path().display(),
+            )));
+        }
+        resolve_remote_model_files(model_id, repo, &cache).await
     }
 }
 
+fn materialize_config(
+    path: Option<&Path>,
+    overrides: &HfOverrides,
+) -> anyhow::Result<tempfile::TempPath> {
+    let mut config: serde_json::Value = match path {
+        Some(path) => serde_json::from_slice(&std::fs::read(path)?)?,
+        None => serde_json::json!({}),
+    };
+    anyhow::ensure!(config.is_object(), "HF model config must be a JSON object");
+    overrides.apply(&mut config);
+    let file = tempfile::Builder::new()
+        .prefix("vllm-rs-hf-config-")
+        .suffix(".json")
+        .tempfile()?;
+    serde_json::to_writer(file.as_file(), &config)?;
+    Ok(file.into_temp_path())
+}
+
 fn resolve_local_model_files(model_dir: &Path) -> Result<ResolvedModelFiles> {
     let tokenizer_config_path = local_file_if_exists(model_dir, "tokenizer_config.json");
     let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?;
@@ -84,12 +139,17 @@ fn resolve_local_model_files(model_dir: &Path) -> Result<ResolvedModelFiles> {
         processor_config_path: local_file_if_exists(model_dir, "processor_config.json"),
         chat_template_path: discover_chat_template_in_dir(model_dir),
         config_path: local_file_if_exists(model_dir, "config.json"),
+        config_temp_path: None,
     })
 }
 
-async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles> {
-    let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?;
-    let repo = api.model(model_id.to_string());
+async fn resolve_remote_model_files(
+    model_id: &str,
+    model_repo: Repo,
+    cache: &Cache,
+) -> Result<ResolvedModelFiles> {
+    let api = build_api(cache).map_err(|error| Error::Tokenizer(error.to_report_string()))?;
+    let repo = api.repo(model_repo);
     let info = repo.info().await.map_err(|error| {
         Error::Tokenizer(format!(
             "failed to fetch model '{model_id}': {}",
@@ -123,11 +183,7 @@ async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles
         download_if_present(&repo, model_id, &siblings, "video_preprocessor_config.json").await?;
     let processor_config_path =
         download_if_present(&repo, model_id, &siblings, "processor_config.json").await?;
-    let chat_template_name = siblings
-        .contains("chat_template.json")
-        .then_some("chat_template.json")
-        .or_else(|| siblings.contains("chat_template.jinja").then_some("chat_template.jinja"))
-        .or_else(|| siblings.iter().copied().find(|name| name.ends_with(".jinja")));
+    let chat_template_name = remote_chat_template_name(&siblings);
     let chat_template_path = match chat_template_name {
         Some(name) => Some(download_known_file(&repo, model_id, name).await?),
         None => None,
@@ -143,15 +199,24 @@ async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles
         processor_config_path,
         chat_template_path,
         config_path,
+        config_temp_path: None,
     })
 }
 
-fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFiles>> {
-    let cache_repo = Cache::from_env().model(model_id.to_string());
+fn resolve_cached_model_files(cache: &Cache, repo: &Repo) -> Result<Option<ResolvedModelFiles>> {
+    let cache_repo = cache.repo(repo.clone());
+    // hf-hub 0.5 looks up refs even for commit SHAs; Python caches have no such ref.
+    let snapshot_dir = (repo.revision().len() == 40
+        && repo.revision().bytes().all(|byte| byte.is_ascii_hexdigit()))
+    .then(|| cache.path().join(repo.folder_name()).join("snapshots").join(repo.revision()));
+    let get_file = |name: &str| match &snapshot_dir {
+        Some(dir) => local_file_if_exists(dir, name),
+        None => cache_repo.get(name),
+    };
 
-    let tokenizer_config_path = cache_repo.get("tokenizer_config.json");
+    let tokenizer_config_path = get_file("tokenizer_config.json");
     let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?;
-    let tokenizer = match resolve_cached_tokenizer_source(&cache_repo, &tokenizer_config)? {
+    let tokenizer = match resolve_cached_tokenizer_source(&get_file, &tokenizer_config)? {
         Some(tokenizer) => tokenizer,
         None => return Ok(None),
     };
@@ -159,12 +224,12 @@ fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFile
     let model_dir = tokenizer.path().parent().ok_or_else(|| {
         Error::Tokenizer("resolved tokenizer file has no parent directory".to_string())
     })?;
-    let generation_config_path = cache_repo.get("generation_config.json");
-    let preprocessor_config_path = cache_repo.get("preprocessor_config.json");
-    let video_preprocessor_config_path = cache_repo.get("video_preprocessor_config.json");
-    let processor_config_path = cache_repo.get("processor_config.json");
+    let generation_config_path = get_file("generation_config.json");
+    let preprocessor_config_path = get_file("preprocessor_config.json");
+    let video_preprocessor_config_path = get_file("video_preprocessor_config.json");
+    let processor_config_path = get_file("processor_config.json");
     let chat_template_path = discover_chat_template_in_dir(model_dir);
-    let config_path = cache_repo.get("config.json");
+    let config_path = get_file("config.json");
 
     Ok(Some(ResolvedModelFiles {
         tokenizer,
@@ -175,6 +240,7 @@ fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFile
         processor_config_path,
         chat_template_path,
         config_path,
+        config_temp_path: None,
     }))
 }
 
@@ -207,20 +273,20 @@ async fn resolve_remote_tokenizer_source(
 }
 
 fn resolve_cached_tokenizer_source(
-    cache_repo: &hf_hub::CacheRepo,
+    get_file: &impl Fn(&str) -> Option<PathBuf>,
     tokenizer_config: &HfTokenizerConfig,
 ) -> Result<Option<TokenizerSource>> {
-    let tekken_path = cache_repo.get("tekken.json");
+    let tekken_path = get_file("tekken.json");
 
     if let Some(tekken_path) = tekken_path {
         return Ok(Some(TokenizerSource::Tekken(tekken_path)));
     }
 
-    let Some(tokenizer_path) = cache_repo.get("tokenizer.json").or_else(|| {
+    let Some(tokenizer_path) = get_file("tokenizer.json").or_else(|| {
         // tiktoken.model is the most common name, try it first.
-        cache_repo.get("tiktoken.model").or_else(|| {
+        get_file("tiktoken.model").or_else(|| {
             // Scan for any *.tiktoken file in the cache snapshot directory.
-            let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
+            let snapshot_dir = get_file("config.json")?.parent()?.to_path_buf();
             discover_tiktoken_in_dir(&snapshot_dir)
         })
     }) else {
@@ -315,8 +381,18 @@ async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) ->
     })
 }
 
-fn build_api() -> anyhow::Result<Api> {
-    let mut builder = ApiBuilder::from_env().with_progress(true);
+fn is_hf_hub_offline() -> bool {
+    std::env::var("HF_HUB_OFFLINE").is_ok_and(|value| {
+        matches!(
+            value.to_ascii_uppercase().as_str(),
+            "1" | "ON" | "YES" | "TRUE"
+        )
+    })
+}
+
+fn build_api(cache: &Cache) -> anyhow::Result<Api> {
+    let mut builder =
+        ApiBuilder::from_env().with_cache_dir(cache.path().clone()).with_progress(true);
     if let Ok(token) = std::env::var(HF_TOKEN_ENV)
         && !token.is_empty()
     {
@@ -365,6 +441,22 @@ pub(super) fn is_tiktoken_file(path: &std::path::Path) -> bool {
         .is_some_and(|name| name == "tiktoken.model" || name.ends_with(".tiktoken"))
 }
 
+/// Select a root-level template, keeping auxiliary tokenizer templates separate.
+fn remote_chat_template_name<'a>(
+    siblings: &std::collections::BTreeSet<&'a str>,
+) -> Option<&'a str> {
+    siblings
+        .contains("chat_template.json")
+        .then_some("chat_template.json")
+        .or_else(|| siblings.contains("chat_template.jinja").then_some("chat_template.jinja"))
+        .or_else(|| {
+            siblings
+                .iter()
+                .copied()
+                .find(|name| !name.contains('/') && name.ends_with(".jinja"))
+        })
+}
+
 /// Chat templates are sometimes stored as dedicated .jinja files rather than as
 /// a fixed-name config entry, so we scan the cached model dir.
 fn discover_chat_template_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
@@ -388,12 +480,61 @@ fn discover_chat_template_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
 #[cfg(test)]
 mod tests {
     use std::fs;
+    use std::path::PathBuf;
 
     use tempfile::tempdir;
     use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
 
     use super::{ResolvedModelFiles, TokenizerSource};
 
+    #[test]
+    fn remote_chat_template_selection_stays_with_main_tokenizer() {
+        let mut siblings = std::collections::BTreeSet::from([
+            "audio_tokenizer/chat_template.jinja",
+            "tokenizer_config.json",
+        ]);
+        assert_eq!(super::remote_chat_template_name(&siblings), None);
+        for name in [
+            "template.jinja",
+            "chat_template.jinja",
+            "chat_template.json",
+        ] {
+            siblings.insert(name);
+            assert_eq!(super::remote_chat_template_name(&siblings), Some(name));
+        }
+    }
+
+    #[tokio::test]
+    async fn overrides_keep_configs_alive_across_clones_and_isolate_instances() {
+        let source = tempdir().unwrap();
+        let original = r#"{"vocab_size":10}"#;
+        fs::write(source.path().join("config.json"), original).unwrap();
+        fs::write(source.path().join("tokenizer.json"), "{}").unwrap();
+        let mut first =
+            ResolvedModelFiles::new(source.path().to_str().unwrap(), None).await.unwrap();
+        let patch = serde_json::from_value(serde_json::json!({"vocab_size":20})).unwrap();
+        first.apply_overrides(&patch).unwrap();
+        let first_path = first.config_path.clone().unwrap();
+
+        let mut second = first.clone();
+        second.apply_overrides(&patch).unwrap();
+        let second_path = second.config_path.clone().unwrap();
+        assert_ne!(first_path, second_path);
+
+        let retained = first.clone();
+        drop(first);
+        assert!(first_path.exists());
+        drop(second);
+        assert!(!second_path.exists());
+        assert!(first_path.exists());
+        drop(retained);
+        assert!(!first_path.exists());
+        assert_eq!(
+            fs::read_to_string(source.path().join("config.json")).unwrap(),
+            original
+        );
+    }
+
     #[tokio::test]
     async fn resolved_model_files_prefers_absolute_local_model_dir() {
         let dir = tempdir().expect("create temp dir");
@@ -405,7 +546,7 @@ mod tests {
         .expect("write tokenizer config");
         fs::write(dir.path().join("config.json"), "{}").expect("write config");
 
-        let files = ResolvedModelFiles::new(dir.path().to_str().expect("utf8 path"))
+        let files = ResolvedModelFiles::new(dir.path().to_str().expect("utf8 path"), None)
             .await
             .expect("resolve local model files");
 
@@ -422,10 +563,90 @@ mod tests {
         );
     }
 
+    #[test]
+    fn cached_model_files_respect_branch_tag_and_commit_revisions() {
+        let dir = tempdir().unwrap();
+        let cache = hf_hub::Cache::new(dir.path().to_path_buf());
+        let model = "test/model";
+        let main_sha = "a".repeat(40);
+        let pinned_sha = "b".repeat(40);
+        for (revision, sha) in [("main", &main_sha), ("release/v1", &pinned_sha)] {
+            let repo =
+                hf_hub::Repo::with_revision(model.into(), hf_hub::RepoType::Model, revision.into());
+            cache.repo(repo.clone()).create_ref(sha).unwrap();
+            let snapshot = dir.path().join(repo.folder_name()).join("snapshots").join(sha);
+            fs::create_dir_all(&snapshot).unwrap();
+            fs::write(snapshot.join("tokenizer.json"), "{}").unwrap();
+            fs::write(snapshot.join("config.json"), "{}").unwrap();
+        }
+        for (revision, expected_sha) in [
+            ("main", Some(&main_sha)),
+            ("release/v1", Some(&pinned_sha)),
+            (pinned_sha.as_str(), Some(&pinned_sha)),
+            ("missing", None),
+        ] {
+            let repo =
+                hf_hub::Repo::with_revision(model.into(), hf_hub::RepoType::Model, revision.into());
+            let files = super::resolve_cached_model_files(&cache, &repo).unwrap();
+            match expected_sha {
+                Some(sha) => {
+                    let snapshot = dir.path().join(repo.folder_name()).join("snapshots").join(sha);
+                    let files = files.unwrap();
+                    assert_eq!(files.tokenizer.path(), snapshot.join("tokenizer.json"));
+                    assert_eq!(files.config_path, Some(snapshot.join("config.json")));
+                }
+                None => assert!(files.is_none()),
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn offline_resolution_uses_hub_cache_and_reports_cache_misses() {
+        const CHILD: &str = "VLLM_TEST_HF_OFFLINE_CHILD";
+        if std::env::var_os(CHILD).is_some() {
+            let files = ResolvedModelFiles::new("test/model", Some("release")).await.unwrap();
+            let cache = PathBuf::from(std::env::var_os("HF_HUB_CACHE").unwrap());
+            assert_eq!(
+                files.tokenizer.path(),
+                cache.join("models--test--model/snapshots/commit/tokenizer.json")
+            );
+            let error = ResolvedModelFiles::new("test/model", Some("missing")).await.unwrap_err();
+            let message = thiserror_ext::AsReport::as_report(&error).to_string();
+            assert!(message.contains("HF_HUB_OFFLINE is enabled"), "{message}");
+            assert!(message.contains("revision 'missing'"), "{message}");
+            return;
+        }
+
+        // Isolate environment variables from other tests and Tokio worker threads.
+        let dir = tempdir().unwrap();
+        let cache = dir.path().join("hub-cache");
+        let repo = cache.join("models--test--model");
+        fs::create_dir_all(repo.join("refs")).unwrap();
+        fs::create_dir_all(repo.join("snapshots/commit")).unwrap();
+        fs::write(repo.join("refs/release"), "commit").unwrap();
+        fs::write(repo.join("snapshots/commit/tokenizer.json"), "{}").unwrap();
+        for value in ["1", "ON", "yes", "True"] {
+            let output = std::process::Command::new(std::env::current_exe().unwrap())
+                .args(["--exact", "backend::hf::model_files::tests::offline_resolution_uses_hub_cache_and_reports_cache_misses", "--nocapture"])
+                .env(CHILD, "1")
+                .env("HF_HOME", dir.path().join("hf-home"))
+                .env("HF_HUB_CACHE", &cache)
+                .env("HF_HUB_OFFLINE", value)
+                .env("HF_ENDPOINT", "http://127.0.0.1:1")
+                .output().unwrap();
+            assert!(
+                output.status.success(),
+                "{}\n{}",
+                String::from_utf8_lossy(&output.stdout),
+                String::from_utf8_lossy(&output.stderr)
+            );
+        }
+    }
+
     #[tokio::test]
     #[ignore = "too slow for CI and requires network access to Hugging Face"]
     async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
-        let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
+        let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5", None)
             .await
             .expect("resolve real Kimi K2.5 model files");
 
diff --git a/rust/src/text/src/backend/hf/overrides.rs b/rust/src/text/src/backend/hf/overrides.rs
new file mode 100644
index 000000000000..ec97eb45c38d
--- /dev/null
+++ b/rust/src/text/src/backend/hf/overrides.rs
@@ -0,0 +1,93 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use serde::{Deserialize, Serialize};
+use serde_json::{Map, Value};
+
+/// An object-shaped JSON Merge Patch (RFC 7396) for the model's HF config.
+/// Objects merge recursively, arrays/scalars replace, and null removes a key.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)]
+#[serde(transparent)]
+pub struct HfOverrides(pub Map<String, Value>);
+
+impl HfOverrides {
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    /// Apply the patch to an in-memory config, preserving the source file.
+    pub fn apply(&self, config: &mut Value) {
+        merge_object(config, &self.0);
+    }
+}
+
+/// Apply RFC 7396 object-member updates at every nesting level.
+fn merge_object(target: &mut Value, patch: &Map<String, Value>) {
+    if !target.is_object() {
+        *target = Value::Object(Map::new());
+    }
+    let target = target.as_object_mut().unwrap();
+    for (key, value) in patch {
+        match value {
+            Value::Null => {
+                target.remove(key);
+            }
+            Value::Object(object) => {
+                merge_object(target.entry(key.clone()).or_insert(Value::Null), object)
+            }
+            _ => {
+                target.insert(key.clone(), value.clone());
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+
+    use super::HfOverrides;
+
+    #[test]
+    fn merge_patch_preserves_siblings_replaces_arrays_and_deletes_null_members() {
+        let mut config = json!({
+            "rope_parameters": {"rope_type": "yarn", "factor": 2},
+            "text_config": {"rope_parameters": {"rope_type": "yarn", "factor": 2}},
+            "architectures": ["Old", "Other"],
+            "sliding_window": 4096
+        });
+        let patch: HfOverrides = serde_json::from_value(json!({
+            "rope_parameters": {"factor": 4},
+            "text_config": {"rope_parameters": {"factor": 4}},
+            "architectures": ["New"],
+            "sliding_window": null,
+            "missing": null
+        }))
+        .unwrap();
+        patch.apply(&mut config);
+        assert_eq!(
+            config,
+            json!({
+                "rope_parameters": {"rope_type": "yarn", "factor": 4},
+                "text_config": {"rope_parameters": {"rope_type": "yarn", "factor": 4}},
+                "architectures": ["New"]
+            })
+        );
+    }
+
+    #[test]
+    fn merge_patch_creates_objects_and_preserves_nulls_inside_replacement_arrays() {
+        let mut config = json!({"a": 1, "b": [1, 2]});
+        let patch: HfOverrides = serde_json::from_value(json!({
+            "a": {"nested": {"removed": null}},
+            "b": {},
+            "c": [null, {"literal": null}]
+        }))
+        .unwrap();
+        patch.apply(&mut config);
+        assert_eq!(
+            config,
+            json!({"a": {"nested": {}}, "b": {}, "c": [null, {"literal": null}]})
+        );
+    }
+}
diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs
index b76c2263f2b8..a6aed592b795 100644
--- a/rust/src/text/src/lib.rs
+++ b/rust/src/text/src/lib.rs
@@ -104,6 +104,27 @@ impl TextRequestProcessor {
         Ok(prompt_token_ids)
     }
 
+    /// Tokenize and truncate the prompt in place, so that later preparation
+    /// reuses the final token IDs instead of tokenizing again.
+    ///
+    /// Frontends that need the final prompt tokens before submission (for
+    /// example, to initialize a prompt-aware output parser) call this first
+    /// and then hand the same request to [`TextLlm::generate`].
+    pub fn tokenize_in_place<'r>(&self, request: &'r mut TextRequest) -> Result<&'r [u32]> {
+        request.validate()?;
+        let prompt_token_ids = self.prepare_prompt_tokens(
+            take(&mut request.prompt),
+            request.add_special_tokens,
+            request.prompt_truncation.take(),
+            request.sampling_params.max_tokens,
+        )?;
+        request.prompt = Prompt::TokenIds(prompt_token_ids);
+        match &request.prompt {
+            Prompt::TokenIds(prompt_token_ids) => Ok(prompt_token_ids),
+            Prompt::Text(_) => unreachable!("prompt was just replaced with token IDs"),
+        }
+    }
+
     /// Prepare one request's prompt tokens without generation-specific lowering.
     pub fn tokenize(&self, request: TextRequest) -> Result<Vec<u32>> {
         request.validate()?;
diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs
index 0b18fd423821..e1357ed7f940 100644
--- a/rust/src/text/src/lower.rs
+++ b/rust/src/text/src/lower.rs
@@ -60,7 +60,9 @@ pub fn lower_text_request(
         priority: request.priority,
         data_parallel_rank: request.data_parallel_rank,
         session_id: request.session_id.clone(),
-        reasoning_parser_kwargs: request.reasoning_parser_kwargs.clone(),
+        kv_hints: request.kv_hints.clone(),
+        reasoning_parser_kwargs: Some(request.reasoning_parser_kwargs.clone()),
+        reasoning_ended: request.reasoning_ended,
         lora_request: request.lora_request.clone(),
         arrival_time: request.arrival_time,
         trace_headers: None,
@@ -92,6 +94,7 @@ pub fn lower_sampling_params(
 ) -> Result<EngineCoreSamplingParams> {
     let SamplingParams {
         temperature,
+        watermarking,
         top_p,
         top_k,
         seed,
@@ -164,6 +167,7 @@ pub fn lower_sampling_params(
 
     let params = EngineCoreSamplingParams {
         temperature,
+        watermarking,
         top_p,
         top_k,
         seed,
@@ -318,6 +322,7 @@ mod tests {
     use std::collections::{BTreeSet, HashMap};
 
     use serial_test::file_serial;
+    use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};
     use vllm_engine_core_client::protocol::multimodal::{
         MmFeatureSpec, MmModality, PlaceholderRange,
     };
@@ -632,6 +637,7 @@ mod tests {
         expect_test::expect![[r#"
             EngineCoreSamplingParams {
                 temperature: 1.0,
+                watermarking: true,
                 top_p: 1.0,
                 top_k: 0,
                 seed: None,
@@ -686,6 +692,7 @@ mod tests {
         expect_test::expect![[r#"
             EngineCoreSamplingParams {
                 temperature: 1.0,
+                watermarking: true,
                 top_p: 1.0,
                 top_k: 0,
                 seed: None,
@@ -802,7 +809,8 @@ mod tests {
     #[file_serial(hf_qwen3)]
     async fn lower_text_request_uses_real_qwen_generation_defaults() {
         let model_id = "Qwen/Qwen3-0.6B";
-        let files = ResolvedModelFiles::new(model_id).await.expect("resolve qwen model files");
+        let files =
+            ResolvedModelFiles::new(model_id, None).await.expect("resolve qwen model files");
         let backend = HfTextBackend::from_resolved_model_files(
             files,
             model_id.to_string(),
@@ -853,6 +861,7 @@ mod tests {
         expect_test::expect![[r#"
             EngineCoreSamplingParams {
                 temperature: 0.6,
+                watermarking: true,
                 top_p: 0.95,
                 top_k: 20,
                 seed: None,
@@ -917,6 +926,7 @@ mod tests {
         expect_test::expect![[r#"
             EngineCoreSamplingParams {
                 temperature: 1.0,
+                watermarking: true,
                 top_p: 1.0,
                 top_k: 0,
                 seed: None,
@@ -989,6 +999,7 @@ mod tests {
         expect_test::expect![[r#"
             EngineCoreSamplingParams {
                 temperature: 0.2,
+                watermarking: true,
                 top_p: 0.3,
                 top_k: 4,
                 seed: None,
@@ -1239,6 +1250,7 @@ mod tests {
         expect_test::expect![[r#"
             EngineCoreSamplingParams {
                 temperature: 0.8,
+                watermarking: true,
                 top_p: 0.9,
                 top_k: 12,
                 seed: None,
@@ -1292,6 +1304,35 @@ mod tests {
         assert_eq!(prepared.generate_request.request_id, "text-1");
     }
 
+    #[test]
+    fn lower_text_request_passes_kv_hints_through() {
+        let hints = KvHintsEnvelope {
+            protocol_version: "0.1".to_string(),
+            message_id: "msg-1".to_string(),
+            actions: vec![KvHintAction {
+                action_id: "action-1".to_string(),
+                action_type: "example.action".to_string(),
+                action_version: "1.0".to_string(),
+                payload: Default::default(),
+            }],
+        };
+        let request = TextRequest {
+            kv_hints: Some(hints.clone()),
+            ..sample_request()
+        };
+
+        let prepared = lower_text_request(
+            request,
+            vec![1, 2, 3],
+            sample_sampling_hints(),
+            sample_sampling_limits(),
+            &stub_tokenizer(),
+        )
+        .unwrap();
+
+        assert_eq!(prepared.generate_request.kv_hints, Some(hints));
+    }
+
     #[test]
     fn lower_text_request_passes_arrival_time_through() {
         let request = TextRequest {
diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs
index a438c8e267d6..d3dc42265641 100644
--- a/rust/src/text/src/output/decoded.rs
+++ b/rust/src/text/src/output/decoded.rs
@@ -9,6 +9,7 @@ use serde::{Deserialize, Serialize};
 use tracing::{Level, debug, trace};
 use vllm_engine_core_client::AbortCause;
 use vllm_engine_core_client::protocol::output::StopReason;
+use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
 use vllm_llm::{FinishReason, GenerateOutput, TokenUsage};
 use vllm_tokenizer::{DecodedText, DynTokenizer, IncrementalDecoder};
 
@@ -50,6 +51,8 @@ pub struct Finished {
     /// Connector-specific encoder cache transfer parameters for disaggregated
     /// serving.
     pub ec_transfer_params: Option<serde_json::Value>,
+    /// Sampling support sets aligned with all generated token positions.
+    pub sampling_mask: Option<SamplingMask>,
 }
 
 /// Sample metadata emitted by one engine output update.
@@ -120,6 +123,7 @@ pub async fn decoded_text_event_stream(
     let mut sampled_token_ids = Vec::new();
     let mut output_token_count: usize = 0;
     let mut sampled_logprobs: Option<DecodedLogprobs> = None;
+    let mut sampling_mask: Option<SamplingMask> = None;
 
     while let Some(next) = raw_stream.next().await {
         let output = next?;
@@ -213,6 +217,7 @@ pub async fn decoded_text_event_stream(
 
         let mut new_token_ids = output.token_ids;
         let mut new_logprobs = output.logprobs;
+        let mut new_sampling_mask = output.sampling_mask;
 
         // Trim tokens and logprobs if we matched stop string.
         if let Some(num_tokens) = truncate_tokens_to {
@@ -220,9 +225,15 @@ pub async fn decoded_text_event_stream(
             if let Some(logprobs) = &mut new_logprobs {
                 logprobs.positions.truncate(num_tokens);
             }
+            if let Some(mask) = &mut new_sampling_mask {
+                mask.rows.truncate(num_tokens);
+            }
         }
 
         output_token_count += new_token_ids.len();
+        if let Some(mut mask) = new_sampling_mask {
+            sampling_mask.get_or_insert_default().rows.append(&mut mask.rows);
+        }
 
         let decoded_logprobs = new_logprobs
             .as_ref()
@@ -246,6 +257,16 @@ pub async fn decoded_text_event_stream(
         }
 
         if let Some(reason) = finish_reason {
+            if let Some(mask) = sampling_mask.as_ref()
+                && mask.rows.len() != output_token_count
+            {
+                return Err(vllm_llm::Error::SamplingMaskTokenCountMismatch {
+                    request_id,
+                    token_count: output_token_count,
+                    row_count: mask.rows.len(),
+                }
+                .into());
+            }
             // Flush any remaining buffered text.
             let (last_chunk, full_decoded) = decoder.flush(truncate_output_to)?;
             let text_len = full_decoded.text.len();
@@ -293,6 +314,7 @@ pub async fn decoded_text_event_stream(
                     finish_reason: reason,
                     kv_transfer_params,
                     ec_transfer_params,
+                    sampling_mask,
                 })),
             })
             .await;
@@ -354,6 +376,7 @@ mod tests {
 
     use futures::{Stream, StreamExt as _, stream};
     use vllm_engine_core_client::AbortCause;
+    use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob};
     use vllm_llm::GenerateOutput;
     use vllm_tokenizer::{TokenAnchor, TokenAttribution, test_utils::TestTokenizer};
 
@@ -384,6 +407,73 @@ mod tests {
         s.bytes().map(u32::from).collect()
     }
 
+    #[tokio::test]
+    async fn sampling_masks_survive_text_decode_and_collection() {
+        let prompt: Arc<[u32]> = Arc::from([]);
+        let mut first = GenerateOutput::for_test(Some(prompt), vec![b'a' as u32], None);
+        first.sampling_mask = Some(SamplingMask {
+            rows: vec![vec![1, b'a' as u32]],
+        });
+        let mut terminal =
+            GenerateOutput::for_test(None, vec![b'b' as u32], Some(FinishReason::Length));
+        terminal.sampling_mask = Some(SamplingMask {
+            rows: vec![vec![2, b'b' as u32]],
+        });
+        let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
+
+        let collected = decoded_text_event_stream(
+            "test-mask".into(),
+            tokenizer,
+            stream::iter([Ok(first), Ok(terminal)]),
+            TextDecodeOptions::default(),
+            false,
+        )
+        .collect_output()
+        .await
+        .unwrap();
+
+        assert_eq!(
+            collected.sampling_mask,
+            Some(SamplingMask {
+                rows: vec![vec![1, b'a' as u32], vec![2, b'b' as u32]],
+            })
+        );
+    }
+
+    #[tokio::test]
+    async fn sampling_mask_mismatch_reuses_llm_error() {
+        let prompt: Arc<[u32]> = Arc::from([]);
+        let mut terminal = GenerateOutput::for_test(
+            Some(prompt),
+            vec![b'a' as u32, b'b' as u32],
+            Some(FinishReason::Length),
+        );
+        terminal.sampling_mask = Some(SamplingMask {
+            rows: vec![vec![1, b'a' as u32]],
+        });
+        let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
+
+        let error = decoded_text_event_stream(
+            "test-mask-mismatch".into(),
+            tokenizer,
+            stream::iter([Ok(terminal)]),
+            TextDecodeOptions::default(),
+            false,
+        )
+        .collect_output()
+        .await
+        .unwrap_err();
+
+        assert!(matches!(
+            error,
+            crate::Error::Llm(vllm_llm::Error::SamplingMaskTokenCountMismatch {
+                token_count: 2,
+                row_count: 1,
+                ..
+            })
+        ));
+    }
+
     fn opts(stop: &[&str], min_tokens: u32) -> TextDecodeOptions {
         TextDecodeOptions {
             stop_strings: Some(stop.iter().map(|s| s.to_string()).collect()),
@@ -472,6 +562,7 @@ mod tests {
                             finish_reason: FinishReason::Length,
                             kv_transfer_params: None,
                             ec_transfer_params: None,
+                            sampling_mask: None,
                         }),
                     ),
                 ],
@@ -503,6 +594,7 @@ mod tests {
                             ))),
                             kv_transfer_params: None,
                             ec_transfer_params: None,
+                            sampling_mask: None,
                         }),
                     ),
                 ],
@@ -576,6 +668,60 @@ mod tests {
         assert!(output.finish_reason.is_stop());
     }
 
+    #[tokio::test]
+    async fn stream_stop_string_mid_chunk_keeps_token_logprob_mask_alignment() {
+        let token_ids = ascii_tokens("abcxyz");
+        let mut raw = GenerateOutput::for_test(
+            Some(Arc::from([])),
+            token_ids.clone(),
+            Some(FinishReason::Length),
+        );
+        raw.logprobs = Some(Logprobs {
+            positions: token_ids
+                .iter()
+                .map(|&token_id| PositionLogprobs {
+                    entries: vec![TokenLogprob {
+                        token_id,
+                        logprob: -0.25,
+                        rank: 1,
+                    }],
+                })
+                .collect(),
+        });
+        raw.sampling_mask = Some(SamplingMask {
+            rows: token_ids.iter().map(|&token_id| vec![token_id, token_id + 100]).collect(),
+        });
+        let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
+
+        let output = decoded_text_event_stream(
+            "test-mid-chunk-stop".into(),
+            tokenizer,
+            stream::iter([Ok(raw)]),
+            opts(&["cx"], 0),
+            false,
+        )
+        .collect_output()
+        .await
+        .unwrap();
+
+        assert_eq!(output.text, "ab");
+        assert!(output.finish_reason.is_stop());
+        assert_eq!(output.token_ids, ascii_tokens("abcx"));
+        assert_eq!(output.logprobs.as_ref().unwrap().positions.len(), 4);
+        assert_eq!(
+            output.sampling_mask,
+            Some(SamplingMask {
+                rows: vec![
+                    vec![b'a' as u32, b'a' as u32 + 100],
+                    vec![b'b' as u32, b'b' as u32 + 100],
+                    vec![b'c' as u32, b'c' as u32 + 100],
+                    vec![b'x' as u32, b'x' as u32 + 100],
+                ],
+            })
+        );
+        assert_eq!(output.usage.output_token_count, 4);
+    }
+
     #[tokio::test]
     async fn stream_stop_string_at_end() {
         let output = run_to_completion(ascii_tokens("abcxyz"), opts(&["xyz"], 0)).await;
diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs
index a415bb8900c2..1293ac6025ff 100644
--- a/rust/src/text/src/output/mod.rs
+++ b/rust/src/text/src/output/mod.rs
@@ -17,6 +17,7 @@ mod logprobs;
 use std::sync::Arc;
 
 use futures::{StreamExt as _, pin_mut};
+use vllm_engine_core_client::protocol::sampling_mask::SamplingMask;
 
 use crate::{Error, FinishReason, Result, TextOutputStream};
 
@@ -35,6 +36,8 @@ pub struct CollectedTextOutput {
     /// Connector-specific encoder cache transfer parameters for disaggregated
     /// serving.
     pub ec_transfer_params: Option<serde_json::Value>,
+    /// Sampling support sets aligned with generated token positions.
+    pub sampling_mask: Option<SamplingMask>,
 }
 
 #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")]
@@ -90,6 +93,7 @@ impl<T: TextOutputStream> T {
                                 usage: vllm_llm::TokenUsage::default(),
                                 kv_transfer_params: None,
                                 ec_transfer_params: None,
+                                sampling_mask: None,
                             })
                         };
 
@@ -99,6 +103,7 @@ impl<T: TextOutputStream> T {
                             collected.usage = finished.usage;
                             collected.kv_transfer_params = finished.kv_transfer_params;
                             collected.ec_transfer_params = finished.ec_transfer_params;
+                            collected.sampling_mask = finished.sampling_mask;
                             return Ok(collected);
                         }
                     }
@@ -173,6 +178,7 @@ mod tests {
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 })),
             }),
         ]);
@@ -295,6 +301,7 @@ mod tests {
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
                     ec_transfer_params: None,
+                    sampling_mask: None,
                 })),
             }),
         ]);
diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs
index 182e20c66022..2c768f48554a 100644
--- a/rust/src/text/src/request.rs
+++ b/rust/src/text/src/request.rs
@@ -6,6 +6,7 @@ use std::collections::HashMap;
 use enum_as_inner::EnumAsInner;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use vllm_engine_core_client::protocol::kv_hints::KvHintsEnvelope;
 use vllm_engine_core_client::protocol::lora::LoraRequest;
 use vllm_engine_core_client::protocol::multimodal::MmFeatures;
 use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
@@ -69,6 +70,8 @@ pub struct SamplingParams {
     /// Controls randomness. Lower values are more deterministic; zero means
     /// greedy sampling. `None` means no explicit user override.
     pub temperature: Option<f32>,
+    /// Whether to apply the engine's configured watermark to this request.
+    pub watermarking: bool,
     /// Cumulative probability threshold for nucleus sampling.
     pub top_p: Option<f32>,
     /// Maximum number of top tokens to consider. `Some(0)` means all tokens.
@@ -144,6 +147,7 @@ impl Default for SamplingParams {
     fn default() -> Self {
         Self {
             temperature: None,
+            watermarking: true,
             top_p: None,
             top_k: None,
             seed: None,
@@ -208,10 +212,22 @@ pub struct TextRequest {
     /// Stable session identity shared by related requests.
     #[serde(default)]
     pub session_id: Option<String>,
-    /// Optional reasoning-parser kwargs forwarded to engine-side structured
-    /// output logic.
+    /// Optional orchestrator-originated KV hints.
     #[serde(default)]
-    pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,
+    pub kv_hints: Option<KvHintsEnvelope>,
+    /// Reasoning-parser kwargs forwarded to engine-side structured output
+    /// logic. The engine consults them only when it owns grammar activation;
+    /// see [`Self::reasoning_ended`].
+    #[serde(default)]
+    pub reasoning_parser_kwargs: ReasoningParserKwargs,
+    /// Optional engine reasoning-gate override selected by a higher-level frontend.
+    ///
+    /// `Some(true)` means the structured output grammar covers reasoning from
+    /// the first generated token, so the engine masks and advances immediately
+    /// instead of waiting for its reasoning parser. Unset for final-output-only
+    /// grammars and for requests without a grammar.
+    #[serde(default)]
+    pub reasoning_ended: Option<bool>,
     /// LoRA adapter selected for this request.
     #[serde(default)]
     pub lora_request: Option<LoraRequest>,
@@ -239,7 +255,9 @@ impl TextRequest {
             add_special_tokens: false,
             data_parallel_rank: None,
             session_id: None,
-            reasoning_parser_kwargs: None,
+            kv_hints: None,
+            reasoning_parser_kwargs: Default::default(),
+            reasoning_ended: None,
             lora_request: None,
             arrival_time: None,
         }
diff --git a/rust/src/tracing/src/lib.rs b/rust/src/tracing/src/lib.rs
index eb6da9c0b247..f407b51f5e4a 100644
--- a/rust/src/tracing/src/lib.rs
+++ b/rust/src/tracing/src/lib.rs
@@ -3,13 +3,14 @@
 
 //! Shared tracing subscriber and log formatting for vLLM Rust binaries.
 
+pub mod timing;
+
 use std::{env, fmt, process};
 
 use time::UtcOffset;
 use time::macros::format_description;
 use tracing::level_filters::LevelFilter;
 use tracing::{Event, Level, Subscriber};
-use tracing_subscriber::Layer as _;
 use tracing_subscriber::filter::Targets;
 use tracing_subscriber::fmt::format::{FormatEvent, FormatFields, Writer};
 use tracing_subscriber::fmt::time::FormatTime;
@@ -17,6 +18,7 @@ use tracing_subscriber::fmt::{FmtContext, FormattedFields};
 use tracing_subscriber::layer::SubscriberExt as _;
 use tracing_subscriber::registry::LookupSpan;
 use tracing_subscriber::util::SubscriberInitExt as _;
+use tracing_subscriber::{Layer, Registry};
 
 const CYAN: &str = "\x1b[0;36m";
 const GREY: &str = "\x1b[90m";
@@ -30,20 +32,36 @@ const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
 
 /// Install the process-wide vLLM-style tracing subscriber.
 pub fn init_tracing(process_label: &str) {
+    let _ = tracing_subscriber::registry().with(vllm_fmt_layer(process_label)).try_init();
+}
+
+/// Install the process-wide vLLM-style tracing subscriber, with an
+/// additional layer (e.g. the multimodal preprocessing timing collector).
+pub fn init_tracing_with<L>(process_label: &str, layer: L)
+where
+    L: Layer<Registry> + Send + Sync + 'static,
+{
+    let _ = tracing_subscriber::registry()
+        .with(layer)
+        .with(vllm_fmt_layer(process_label))
+        .try_init();
+}
+
+/// The vLLM-style fmt layer filtered by `VLLM_LOGGING_LEVEL` / `RUST_LOG`.
+fn vllm_fmt_layer<S>(process_label: &str) -> impl Layer<S> + Send + Sync + 'static + use<S>
+where
+    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
+{
     let filter = build_targets_filter(
         env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
         env::var("RUST_LOG").ok().as_deref(),
     );
     let formatter = VllmEventFormatter::new(process_label);
 
-    let _ = tracing_subscriber::registry()
-        .with(
-            tracing_subscriber::fmt::layer()
-                .event_format(formatter)
-                .with_writer(std::io::stderr)
-                .with_filter(filter),
-        )
-        .try_init();
+    tracing_subscriber::fmt::layer()
+        .event_format(formatter)
+        .with_writer(std::io::stderr)
+        .with_filter(filter)
 }
 
 /// Build the CLI log filter by merging the vLLM-style default level with
diff --git a/rust/src/tracing/src/timing.rs b/rust/src/tracing/src/timing.rs
new file mode 100644
index 000000000000..899e9eaea247
--- /dev/null
+++ b/rust/src/tracing/src/timing.rs
@@ -0,0 +1,298 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+//! Request-scoped span timing.
+//!
+//! Pipelines emit stage spans (target = the layer's configured target, a
+//! `stage` field naming the stage) inside a request span carrying a
+//! `request_id` field; [`RequestTimingLayer`] attributes each stage span's
+//! wall-clock time to the nearest ancestor `request_id` and accumulates per
+//! request. Callers that want timings install the layer on their subscriber;
+//! everyone else pays only the cost of untracked spans.
+
+use std::collections::HashMap;
+use std::fmt;
+use std::sync::{Arc, Mutex};
+use std::time::Instant;
+
+use tracing::field::{Field, Visit};
+use tracing::span::Attributes;
+use tracing::{Id, Subscriber};
+use tracing_subscriber::filter::{Filtered, LevelFilter};
+use tracing_subscriber::layer::{Context, Layer};
+use tracing_subscriber::registry::LookupSpan;
+
+/// `"{stage}_secs" => seconds` maps returned by [`RequestTimingStats::stat`].
+pub type StageStats = HashMap<String, f64>;
+
+/// `request_id` extracted from a span, propagated to stage timings.
+struct RequestId(String);
+
+/// Start timestamp of one in-flight stage span.
+struct StageStart {
+    stage: String,
+    start: Instant,
+}
+
+/// Extracts the fields the layer cares about from span attributes.
+#[derive(Default)]
+struct SpanFields {
+    request_id: Option<String>,
+    stage: Option<String>,
+}
+
+impl Visit for SpanFields {
+    fn record_str(&mut self, field: &Field, value: &str) {
+        match field.name() {
+            "request_id" => self.request_id = Some(value.to_string()),
+            "stage" => self.stage = Some(value.to_string()),
+            _ => {}
+        }
+    }
+
+    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
+        match field.name() {
+            "request_id" => self.request_id = Some(format!("{value:?}")),
+            "stage" => self.stage = Some(format!("{value:?}")),
+            _ => {}
+        }
+    }
+}
+
+/// `tracing` layer that aggregates stage-span timings per request, keyed by
+/// the nearest ancestor span with a `request_id` field.
+///
+/// # Span contract
+///
+/// Timed spans must have the configured target and a `stage` field. The
+/// `request_id` is taken from the timed span itself or its nearest ancestor
+/// carrying that field, regardless of the ancestor's target. Spans without
+/// a request ID are ignored. Span names are unrestricted. Both the timed spans
+/// and their `request_id` ancestors must be at `INFO` level or above.
+///
+/// Both fields are read at span creation; subsequent [`tracing::Span::record`]
+/// calls are ignored. String values are used verbatim; other values use their
+/// debug representation. Prefer string fields for stable keys.
+///
+/// # Timing semantics
+///
+/// Each sample measures wall-clock time from span creation to final close,
+/// including async waits and any time the span remains alive through cloned
+/// handles. Nested and concurrent spans are measured independently, so their
+/// durations can overlap. Repeated stages accumulate under `"{stage}_secs"`
+/// for each request. Closing spans are recorded regardless of the operation's
+/// success, failure, or cancellation; callers select the samples to report.
+pub struct RequestTimingLayer {
+    target: &'static str,
+    stats: Arc<Mutex<HashMap<String, StageStats>>>,
+}
+
+/// Drain handle for the timings collected by [`RequestTimingLayer`].
+pub struct RequestTimingStats {
+    stats: Arc<Mutex<HashMap<String, StageStats>>>,
+}
+
+impl RequestTimingLayer {
+    /// Create the layer and its stats handle. Only spans whose target equals
+    /// `target` and that carry a `stage` field are timed.
+    ///
+    /// The layer is filtered to `INFO` so that, composed with a filtered log
+    /// layer, it does not raise the subscriber's max level hint and enable
+    /// `DEBUG`/`TRACE` callsites that the log layer would discard anyway.
+    pub fn new<S>(target: &'static str) -> (Filtered<Self, LevelFilter, S>, RequestTimingStats) {
+        let stats = Arc::new(Mutex::new(HashMap::new()));
+        let layer = Self {
+            target,
+            stats: Arc::clone(&stats),
+        };
+        (
+            Filtered::new(layer, LevelFilter::INFO),
+            RequestTimingStats { stats },
+        )
+    }
+}
+
+impl RequestTimingStats {
+    /// Drain and return `{request_id: {stage_secs}}` records.
+    ///
+    /// Returns accumulated timings since the previous drain. Spans still open
+    /// contribute to a later drain; finish the measured work before calling
+    /// this method to collect complete request timings.
+    pub fn stat(&self) -> HashMap<String, StageStats> {
+        std::mem::take(&mut *self.stats.lock().unwrap())
+    }
+}
+
+impl<S> Layer<S> for RequestTimingLayer
+where
+    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
+{
+    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
+        let Some(span) = ctx.span(id) else {
+            return;
+        };
+        let mut fields = SpanFields::default();
+        attrs.record(&mut fields);
+        let mut extensions = span.extensions_mut();
+        if let Some(request_id) = fields.request_id {
+            extensions.insert(RequestId(request_id));
+        }
+        if span.metadata().target() == self.target
+            && let Some(stage) = fields.stage
+        {
+            extensions.insert(StageStart {
+                stage,
+                start: Instant::now(),
+            });
+        }
+    }
+
+    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
+        let Some(span) = ctx.span(&id) else {
+            return;
+        };
+        let (stage, start) = {
+            let extensions = span.extensions();
+            let Some(stage_start) = extensions.get::<StageStart>() else {
+                return;
+            };
+            (stage_start.stage.clone(), stage_start.start)
+        };
+        let Some(request_id) = span
+            .scope()
+            .find_map(|ancestor| ancestor.extensions().get::<RequestId>().map(|id| id.0.clone()))
+        else {
+            return;
+        };
+        let mut stats = self.stats.lock().unwrap();
+        *stats.entry(request_id).or_default().entry(format!("{stage}_secs")).or_default() +=
+            start.elapsed().as_secs_f64();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    use super::*;
+    use tracing::{debug, debug_span, info_span};
+    use tracing_subscriber::layer::SubscriberExt as _;
+
+    const TARGET: &str = "test_request_timing";
+
+    fn stage_span(stage: &'static str) -> tracing::Span {
+        info_span!(target: TARGET, "stage", stage)
+    }
+
+    #[test]
+    fn records_stage_timings_per_request() {
+        let (layer, stats) = RequestTimingLayer::new(TARGET);
+        let subscriber = tracing_subscriber::registry().with(layer);
+        tracing::subscriber::with_default(subscriber, || {
+            info_span!("request", request_id = "req-1").in_scope(|| {
+                drop(stage_span("media_fetch"));
+                stage_span("prompt_expansion").in_scope(|| {});
+            });
+        });
+        let drained = stats.stat();
+        assert!(drained["req-1"].contains_key("media_fetch_secs"));
+        assert!(drained["req-1"].contains_key("prompt_expansion_secs"));
+        // `stat` drains: a second read is empty.
+        assert!(stats.stat().is_empty());
+    }
+
+    #[test]
+    fn attributes_to_nearest_request_id_ancestor() {
+        let (layer, stats) = RequestTimingLayer::new(TARGET);
+        let subscriber = tracing_subscriber::registry().with(layer);
+        tracing::subscriber::with_default(subscriber, || {
+            info_span!("request", request_id = "req-outer").in_scope(|| {
+                info_span!("request", request_id = "req-inner").in_scope(|| {
+                    drop(stage_span("media_fetch"));
+                });
+            });
+        });
+        let drained = stats.stat();
+        assert!(drained["req-inner"].contains_key("media_fetch_secs"));
+        assert!(!drained.contains_key("req-outer"));
+    }
+
+    #[test]
+    fn stage_without_request_id_is_dropped() {
+        let (layer, stats) = RequestTimingLayer::new(TARGET);
+        let subscriber = tracing_subscriber::registry().with(layer);
+        tracing::subscriber::with_default(subscriber, || {
+            drop(stage_span("media_fetch"));
+        });
+        assert!(stats.stat().is_empty());
+    }
+
+    #[test]
+    fn does_not_enable_debug_at_info() {
+        let (timing_layer, _) = RequestTimingLayer::new(TARGET);
+        let info_layer = tracing_subscriber::fmt::layer()
+            .with_writer(std::io::sink)
+            .with_filter(LevelFilter::INFO);
+        let subscriber = tracing_subscriber::registry().with(timing_layer).with(info_layer);
+        let evaluated_debug_fields = AtomicUsize::new(0);
+
+        tracing::subscriber::with_default(subscriber, || {
+            assert_eq!(LevelFilter::current(), LevelFilter::INFO);
+            debug!(
+                value = evaluated_debug_fields.fetch_add(1, Ordering::Relaxed),
+                "debug event"
+            );
+            assert!(debug_span!("debug span").is_disabled());
+        });
+
+        assert_eq!(evaluated_debug_fields.load(Ordering::Relaxed), 0);
+    }
+
+    #[test]
+    fn does_not_disable_explicit_debug_logging() {
+        let (timing_layer, _) = RequestTimingLayer::new(TARGET);
+        let debug_layer = tracing_subscriber::fmt::layer()
+            .with_writer(std::io::sink)
+            .with_filter(LevelFilter::DEBUG);
+        let subscriber = tracing_subscriber::registry().with(timing_layer).with(debug_layer);
+        let evaluated_debug_fields = AtomicUsize::new(0);
+
+        tracing::subscriber::with_default(subscriber, || {
+            assert_eq!(LevelFilter::current(), LevelFilter::DEBUG);
+            debug!(
+                value = evaluated_debug_fields.fetch_add(1, Ordering::Relaxed),
+                "debug event"
+            );
+            assert!(!debug_span!("debug span").is_disabled());
+        });
+
+        assert_eq!(evaluated_debug_fields.load(Ordering::Relaxed), 1);
+    }
+
+    #[test]
+    fn records_stage_timings_when_logging_at_warn() {
+        let (timing_layer, stats) = RequestTimingLayer::new(TARGET);
+        let warn_layer = tracing_subscriber::fmt::layer()
+            .with_writer(std::io::sink)
+            .with_filter(LevelFilter::WARN);
+        let subscriber = tracing_subscriber::registry().with(timing_layer).with(warn_layer);
+        tracing::subscriber::with_default(subscriber, || {
+            info_span!("request", request_id = "req-1").in_scope(|| {
+                drop(stage_span("media_fetch"));
+            });
+        });
+        assert!(stats.stat()["req-1"].contains_key("media_fetch_secs"));
+    }
+
+    #[test]
+    fn other_targets_are_ignored() {
+        let (layer, stats) = RequestTimingLayer::new(TARGET);
+        let subscriber = tracing_subscriber::registry().with(layer);
+        tracing::subscriber::with_default(subscriber, || {
+            info_span!("request", request_id = "req-1").in_scope(|| {
+                drop(info_span!("unrelated", stage = "media_fetch"));
+            });
+        });
+        assert!(stats.stat().is_empty());
+    }
+}
diff --git a/scripts/autotune_helion_kernels.py b/scripts/autotune_helion_kernels.py
index 864aabfc0964..e6352a40af41 100644
--- a/scripts/autotune_helion_kernels.py
+++ b/scripts/autotune_helion_kernels.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Autotune registered Helion kernels for optimal configurations.
+"""Autotune registered Helion kernels for optimal configurations.
 
 Usage:
     # Autotune all registered kernels
diff --git a/scripts/benchmark_helion_kernels.py b/scripts/benchmark_helion_kernels.py
index fea30d6f9cfa..04350896af8b 100644
--- a/scripts/benchmark_helion_kernels.py
+++ b/scripts/benchmark_helion_kernels.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark a registered Helion kernel against a baseline.
+"""Benchmark a registered Helion kernel against a baseline.
 
 For each input case produced by the kernel's registered input generator, this
 checks the Helion kernel's numerics once against its eager reference, measures
@@ -394,6 +393,7 @@ def check_kernel_correctness(
         ``passed=False`` with the exception text in ``error``; iteration
         continues regardless. An empty input mapping returns an empty list, which
         the CLI reports as a skipped check.
+
     """
     if inputs_dict is None:
         inputs_dict = kernel.get_inputs()
diff --git a/setup.py b/setup.py
index 81ff499be4d5..67f192cdd393 100644
--- a/setup.py
+++ b/setup.py
@@ -22,6 +22,18 @@
 from setuptools_scm import get_version
 from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
 
+# Select vLLM release tags, excluding crate tags such as "proto-v0.3.0".
+VLLM_GIT_DESCRIBE_COMMAND = [
+    "git",
+    "describe",
+    "--dirty",
+    "--tags",
+    "--long",
+    "--abbrev=40",
+    "--match",
+    "v[0-9]*",
+]
+
 
 def load_module_from_path(module_name, path):
     spec = importlib.util.spec_from_file_location(module_name, path)
@@ -512,8 +524,7 @@ def fetch_metadata_for_variant(
         *,
         rocm: bool = False,
     ) -> tuple[list[dict], str]:
-        """
-        Fetches metadata for a specific variant of the precompiled wheel.
+        """Fetches metadata for a specific variant of the precompiled wheel.
 
         For non-ROCm, fetches vllm metadata.
 
@@ -591,7 +602,6 @@ def is_rocm_system() -> bool:
     @staticmethod
     def detect_system_cuda_variant() -> str:
         """Auto-detect CUDA variant from torch, nvidia-smi, or env default."""
-
         # Map CUDA major version to hosted wheel variants on wheels.vllm.ai
         supported = {12: "cu129", 13: "cu130"}
 
@@ -886,8 +896,7 @@ def determine_wheel_url_rocm() -> tuple[str, str | None]:
 
     @staticmethod
     def determine_wheel_url() -> tuple[str, str | None]:
-        """
-        Try to determine the precompiled wheel URL or path to use.
+        """Try to determine the precompiled wheel URL or path to use.
         The order of preference is:
         1. user-specified wheel location (can be either local or remote, via
            VLLM_PRECOMPILED_WHEEL_LOCATION)
@@ -1015,6 +1024,7 @@ def extract_precompiled_and_patch_package(
                             "vllm/_flashmla_extension_C.abi3.so",
                             "vllm/_flashkda_C.abi3.so",
                             "vllm/_sparse_flashmla_C.abi3.so",
+                            "vllm/_deepselect_C.abi3.so",
                             "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
                             "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
                             "vllm/cumem_allocator.abi3.so",
@@ -1251,40 +1261,50 @@ def get_vllm_version() -> str:
     if env_version := os.getenv("VLLM_VERSION_OVERRIDE"):
         print(f"Overriding VLLM version with {env_version} from VLLM_VERSION_OVERRIDE")
         os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version
-        return get_version(write_to="vllm/_version.py")
+        return get_version(
+            write_to="vllm/_version.py",
+            git_describe_command=VLLM_GIT_DESCRIBE_COMMAND,
+        )
 
-    version = get_version(write_to="vllm/_version.py")
+    version = get_version(
+        write_to="vllm/_version.py",
+        git_describe_command=VLLM_GIT_DESCRIBE_COMMAND,
+    )
     sep = "+" if "+" not in version else "."  # dev versions might contain +
 
-    if _no_device():
-        if envs.VLLM_TARGET_DEVICE == "empty":
-            version += f"{sep}empty"
-    elif _is_cuda():
-        if USE_PRECOMPILED_EXTENSIONS and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX:
-            version += f"{sep}precompiled"
+    if not envs.VLLM_SKIP_VERSION_SUFFIX:
+        if _no_device():
+            if envs.VLLM_TARGET_DEVICE == "empty":
+                version += f"{sep}empty"
+        elif _is_cuda():
+            if (
+                USE_PRECOMPILED_EXTENSIONS
+                and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX
+            ):
+                version += f"{sep}precompiled"
+            else:
+                cuda_version = str(get_nvcc_cuda_version())
+                if cuda_version != envs.VLLM_MAIN_CUDA_VERSION:
+                    cuda_version_str = cuda_version.replace(".", "")[:3]
+                    # skip this for source tarball, required for pypi
+                    if "sdist" not in sys.argv:
+                        version += f"{sep}cu{cuda_version_str}"
+        elif _is_hip():
+            # Get the Rocm Version
+            rocm_version = get_rocm_version() or torch.version.hip
+            if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION:
+                version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}"
+        elif _is_tpu():
+            version += f"{sep}tpu"
+        elif _is_cpu():
+            # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above),
+            # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`.
+            if VLLM_TARGET_DEVICE == "cpu":
+                version += f"{sep}cpu"
+        elif _is_xpu():
+            version += f"{sep}xpu"
         else:
-            cuda_version = str(get_nvcc_cuda_version())
-            if cuda_version != envs.VLLM_MAIN_CUDA_VERSION:
-                cuda_version_str = cuda_version.replace(".", "")[:3]
-                # skip this for source tarball, required for pypi
-                if "sdist" not in sys.argv:
-                    version += f"{sep}cu{cuda_version_str}"
-    elif _is_hip():
-        # Get the Rocm Version
-        rocm_version = get_rocm_version() or torch.version.hip
-        if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION:
-            version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}"
-    elif _is_tpu():
-        version += f"{sep}tpu"
-    elif _is_cpu():
-        # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above),
-        # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`.
-        if VLLM_TARGET_DEVICE == "cpu":
-            version += f"{sep}cpu"
-    elif _is_xpu():
-        version += f"{sep}xpu"
-    else:
-        raise RuntimeError("Unknown runtime environment")
+            raise RuntimeError("Unknown runtime environment")
 
     return version
 
@@ -1381,6 +1401,9 @@ def _read_requirements(filename: str) -> list[str]:
         ext_modules.append(
             CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
         )
+        # DeepSelect requires CUDA 12.9 or later (SM100a/SM103a only)
+        # Optional since it won't build on unsupported architectures
+        ext_modules.append(CMakeExtension(name="vllm._deepselect_C", optional=True))
     if USE_PRECOMPILED_EXTENSIONS or (
         CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0")
     ):
@@ -1506,7 +1529,7 @@ def add_vllm_package_data(filename: str) -> None:
     install_requires=get_requirements(),
     extras_require={
         # AMD Zen CPU optimizations via zentorch
-        "zen": ["zentorch==2.13.0.0"],
+        "zen": ["zentorch==2.13.0.1"],
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.3.3"],
diff --git a/tests/basic_correctness/test_basic_correctness.py b/tests/basic_correctness/test_basic_correctness.py
index 3412b30814e3..688d355c318d 100644
--- a/tests/basic_correctness/test_basic_correctness.py
+++ b/tests/basic_correctness/test_basic_correctness.py
@@ -11,8 +11,6 @@
 
 import pytest
 import torch
-from packaging.version import Version
-from transformers import __version__ as TRANSFORMERS_VERSION
 
 import vllm.envs as envs
 from vllm import LLM
@@ -89,7 +87,7 @@ def _resolve_target_test_suite() -> str:
     condition=current_platform.is_rocm(),
 )
 def test_vllm_gc_ed():
-    """Verify vllm instance is GC'ed when it is deleted"""
+    """Verify vllm instance is GC'ed when it is deleted."""
     llm = LLM("hmellor/tiny-random-LlamaForCausalLM")
     weak_llm = weakref.ref(llm)
     del llm
@@ -149,15 +147,6 @@ def test_models(
         if enable_prompt_embeds:
             with torch.no_grad():
                 prompt_embeds = hf_model.get_prompt_embeddings(example_prompts)
-            if model == "hmellor/tiny-random-Gemma2ForCausalLM" and (
-                Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0")
-            ):
-                # For Gemma 1/2 models with Transformers 5.4.0+, the prompt embeddings
-                # are normalised in `get_prompt_embeddings`, like Gemma 3.
-                # For older versions, we need to manually normalise.
-                embed_scale = hf_model.config.hidden_size**0.5
-                normalizer = torch.tensor(embed_scale, dtype=prompt_embeds[0].dtype)
-                prompt_embeds = [p_e * normalizer for p_e in prompt_embeds]
 
     with VllmRunner(
         model,
diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py
index c3f66e24abc5..6a483d26aa9f 100644
--- a/tests/basic_correctness/test_mem.py
+++ b/tests/basic_correctness/test_mem.py
@@ -51,8 +51,7 @@ def create_and_map_with_poison(handle) -> None:
 
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
 def test_python_error():
-    """
-    Test if Python error occurs when there's low-level
+    """Test if Python error occurs when there's low-level
     error happening from the C++ side.
     """
     allocator = get_mem_allocator_instance()
@@ -110,6 +109,100 @@ def test_basic_cumem():
     assert torch.allclose(output, torch.ones_like(output) * 3)
 
 
+@pytest.mark.parametrize("full_sleep", [0, 1, 2], ids=["kv-only", "sleep-1", "sleep-2"])
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
+def test_release_kv_cache_memory_preserves_generation(full_sleep, monkeypatch):
+    """Rejected release and KV-only/full sleep cycles preserve greedy tokens."""
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    kv_cache_memory_bytes = 256 * 1024 * 1024
+    llm = LLM(
+        "Qwen/Qwen3-0.6B",
+        enable_sleep_mode=True,
+        enforce_eager=True,
+        kv_cache_memory_bytes=kv_cache_memory_bytes,
+        max_model_len=1024,
+        max_num_seqs=4,
+        gpu_memory_utilization=0.05,
+    )
+    prompt = "How are you?"
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+    expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
+
+    def get_mapped_bytes_by_tag(worker):
+        allocator = get_mem_allocator_instance()
+        mapped_bytes = mapped_usage(allocator)
+        kv_cache_bytes = sum(
+            data.handle[1]
+            for data in allocator.pointer_to_data.values()
+            if data.tag == "kv_cache" and not data.is_asleep
+        )
+        return mapped_bytes, kv_cache_bytes
+
+    mapped_before, kv_cache_before = llm.collective_rpc(get_mapped_bytes_by_tag)[0]
+    # Utility RPC transports the engine error as a plain Exception.
+    with pytest.raises(Exception, match="requires a completed pause"):
+        llm.release_kv_cache_memory()
+    assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == mapped_before
+    assert not llm.llm_engine.is_sleeping()
+    assert llm.generate(prompt, sampling_params)[0].outputs[0].token_ids == expected
+
+    llm.sleep(level=0)
+    llm.release_kv_cache_memory()
+    mapped_after, kv_cache_after = llm.collective_rpc(get_mapped_bytes_by_tag)[0]
+    assert kv_cache_before > 0
+    assert kv_cache_after == 0
+    assert mapped_before - mapped_after >= kv_cache_before
+    assert mapped_after > 0
+    assert llm.llm_engine.is_sleeping()
+
+    if full_sleep:
+        llm.sleep(level=full_sleep)
+        assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == 0
+    llm.wake_up(tags=None if full_sleep else ["kv_cache"])
+    if full_sleep == 2:
+        llm.collective_rpc("reload_weights")
+    assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == mapped_before
+    assert not llm.llm_engine.is_sleeping()
+    actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
+    assert actual == expected
+
+
+@pytest.mark.parametrize("first_level", [1, 2])
+@pytest.mark.parametrize("second_level", [1, 2])
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
+def test_sleep_with_only_weights_asleep(first_level, second_level, monkeypatch):
+    """Repeated sleep after KV-only wake preserves mappings and recoverability."""
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    llm = LLM(
+        "Qwen/Qwen3-0.6B",
+        enable_sleep_mode=True,
+        enforce_eager=True,
+        kv_cache_memory_bytes=256 * 1024 * 1024,
+        max_model_len=1024,
+        max_num_seqs=4,
+        gpu_memory_utilization=0.05,
+    )
+    prompt = "How are you?"
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+    expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
+
+    def get_mapped_bytes(worker):
+        return mapped_usage(get_mem_allocator_instance())
+
+    llm.sleep(level=first_level)
+    llm.wake_up(tags=["kv_cache"])
+    mapped_before = llm.collective_rpc(get_mapped_bytes)[0]
+    assert mapped_before > 0
+    llm.sleep(level=second_level)
+    assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before
+    llm.wake_up(tags=["weights"])
+    if first_level == 2:
+        llm.collective_rpc("reload_weights")
+    assert not llm.llm_engine.is_sleeping()
+    actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
+    assert actual == expected
+
+
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
 def test_discard_tags():
     """Test that discard(tags) selectively frees GPU memory for specific
diff --git a/tests/benchmarks/test_endpoint_request_func_timing.py b/tests/benchmarks/test_endpoint_request_func_timing.py
new file mode 100644
index 000000000000..57ac5824ab49
--- /dev/null
+++ b/tests/benchmarks/test_endpoint_request_func_timing.py
@@ -0,0 +1,464 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Timing-accounting invariants for the streaming benchmark request functions.
+
+serve.py derives TPOT as (latency - ttft) / (output_len - 1), which only means
+"average inter-token latency" if latency - ttft == sum(itl). These tests pin
+that identity, and pin that the trailing choice-less usage chunk does not
+extend the measured end of the request.
+
+Driven against a fake session and a scripted clock: no GPU, model, or socket.
+"""
+
+import asyncio
+import importlib.util
+import json
+from pathlib import Path
+from typing import Any
+
+import aiohttp
+import numpy as np
+import pytest
+
+import vllm.benchmarks.lib.endpoint_request_func as request_func_module
+from vllm.benchmarks.lib.endpoint_request_func import (
+    RequestFuncInput,
+    async_request_openai_audio,
+    async_request_openai_chat_completions,
+    async_request_openai_completions,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+AUDIO_SAMPLE_RATE = 16_000
+
+
+class _ScriptedClock:
+    """A deterministic stand-in for time.perf_counter.
+
+    Each read advances the clock by exactly one tick, so what a code path
+    observes depends only on how many times it reads the clock. A stray or
+    misattributed read then shows up as a whole tick of drift rather than as
+    sub-microsecond jitter, which keeps these assertions free of flakiness.
+    """
+
+    def __init__(self, tick: float = 1.0, start: float = 1000.0) -> None:
+        self.tick = tick
+        self._now = start
+
+    def __call__(self) -> float:
+        value = self._now
+        self._now += self.tick
+        return value
+
+
+class _FakeTime:
+    """Namespace substituted for the module-level time import."""
+
+    def __init__(self, perf_counter) -> None:
+        self.perf_counter = perf_counter
+
+
+def _sse(payload: dict[str, Any]) -> bytes:
+    return f"data: {json.dumps(payload)}\n\n".encode()
+
+
+def _token_chunk(index: int, *, chat: bool) -> bytes:
+    if chat:
+        choice = {"index": 0, "delta": {"content": f"tok{index} "}}
+    else:
+        choice = {"index": 0, "text": f"tok{index} "}
+    return _sse({"id": "cmpl-test", "object": "x", "choices": [choice]})
+
+
+def _usage_chunk(n_tokens: int) -> bytes:
+    return _sse(
+        {
+            "id": "cmpl-test",
+            "object": "x",
+            "choices": [],
+            "usage": {
+                "prompt_tokens": 7,
+                "completion_tokens": n_tokens,
+                "total_tokens": 7 + n_tokens,
+            },
+        }
+    )
+
+
+def _build_stream(
+    n_tokens: int,
+    *,
+    chat: bool,
+    ttft_delay: float = 0.0,
+    itl_delay: float = 0.0,
+    usage_delay: float = 0.0,
+) -> list[tuple[float, bytes]]:
+    """Build ``(delay_before_chunk, payload)`` pairs for one streamed response.
+
+    ``vllm bench serve`` always requests ``stream_options.include_usage``, so a
+    choice-less usage chunk always trails the final token.
+    """
+    chunks = [
+        (ttft_delay if i == 0 else itl_delay, _token_chunk(i, chat=chat))
+        for i in range(n_tokens)
+    ]
+    chunks.append((usage_delay, _usage_chunk(n_tokens)))
+    chunks.append((0.0, b"data: [DONE]\n\n"))
+    return chunks
+
+
+class _FakeContent:
+    def __init__(self, chunks: list[tuple[float, bytes]]) -> None:
+        self._chunks = chunks
+
+    async def iter_any(self):
+        for delay, payload in self._chunks:
+            if delay:
+                await asyncio.sleep(delay)
+            yield payload
+
+    # The legacy script iterates ``response.content`` directly.
+    def __aiter__(self):
+        return self.iter_any()
+
+
+class _FakeResponse:
+    def __init__(self, chunks: list[tuple[float, bytes]]) -> None:
+        self.status = 200
+        self.reason = "OK"
+        self.content = _FakeContent(chunks)
+
+    async def __aenter__(self) -> "_FakeResponse":
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> bool:
+        return False
+
+
+class _FakeSession:
+    """Returns a scripted SSE stream for any request, without touching a socket."""
+
+    def __init__(self, chunks: list[tuple[float, bytes]]) -> None:
+        self._chunks = chunks
+
+    def post(self, **kwargs: Any) -> _FakeResponse:
+        del kwargs
+        return _FakeResponse(self._chunks)
+
+    async def __aenter__(self) -> "_FakeSession":
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> bool:
+        return False
+
+
+# (id, request function, api url, is_chat_payload)
+STREAMING_ENDPOINTS = [
+    (
+        "completions",
+        async_request_openai_completions,
+        "http://test/v1/completions",
+        False,
+    ),
+    (
+        "chat",
+        async_request_openai_chat_completions,
+        "http://test/v1/chat/completions",
+        True,
+    ),
+    (
+        "audio",
+        async_request_openai_audio,
+        "http://test/v1/audio/transcriptions",
+        True,
+    ),
+]
+
+
+def _make_input(api_url: str, n_tokens: int, endpoint_id: str) -> RequestFuncInput:
+    multi_modal_content = None
+    if endpoint_id == "audio":
+        # A short silent clip; only its duration is read by the request func.
+        multi_modal_content = {
+            "audio": (
+                np.zeros(AUDIO_SAMPLE_RATE // 10, dtype=np.float32),
+                AUDIO_SAMPLE_RATE,
+            )
+        }
+    return RequestFuncInput(
+        prompt="hello",
+        api_url=api_url,
+        prompt_len=7,
+        output_len=n_tokens,
+        model="test-model",
+        multi_modal_content=multi_modal_content,
+    )
+
+
+def _run(
+    request_func,
+    api_url: str,
+    endpoint_id: str,
+    chunks,
+    n_tokens: int,
+    *,
+    expect_success: bool = True,
+):
+    output = asyncio.run(
+        request_func(
+            _make_input(api_url, n_tokens, endpoint_id),
+            _FakeSession(chunks),
+        )
+    )
+    if expect_success:
+        assert output.success, output.error
+    return output
+
+
+@pytest.mark.parametrize(
+    "endpoint_id,request_func,api_url,chat",
+    STREAMING_ENDPOINTS,
+    ids=[e[0] for e in STREAMING_ENDPOINTS],
+)
+def test_decode_span_identity_is_exact(
+    monkeypatch: pytest.MonkeyPatch,
+    endpoint_id: str,
+    request_func,
+    api_url: str,
+    chat: bool,
+) -> None:
+    """Latency - ttft must equal sum(itl) exactly.
+
+    On a scripted clock this can only drift if a code path reads the clock an
+    extra time or attributes a read to the wrong event.
+    """
+    tick = 1.0
+    monkeypatch.setattr(
+        request_func_module, "time", _FakeTime(_ScriptedClock(tick=tick))
+    )
+
+    n_tokens = 5
+    chunks = _build_stream(n_tokens, chat=chat)
+    output = _run(request_func, api_url, endpoint_id, chunks, n_tokens)
+
+    assert len(output.itl) == n_tokens - 1
+    assert output.output_tokens == n_tokens
+
+    residual = (output.latency - output.ttft) - sum(output.itl)
+    assert residual == pytest.approx(0.0, abs=tick / 1000), (
+        f"{endpoint_id}: (latency - ttft) - sum(itl) = {residual!r}; "
+        f"TPOT is derived from latency - ttft, so it must match the measured "
+        f"inter-token latencies"
+    )
+
+
+@pytest.mark.parametrize(
+    "endpoint_id,request_func,api_url,chat",
+    STREAMING_ENDPOINTS,
+    ids=[e[0] for e in STREAMING_ENDPOINTS],
+)
+def test_trailing_usage_chunk_does_not_extend_latency(
+    monkeypatch: pytest.MonkeyPatch,
+    endpoint_id: str,
+    request_func,
+    api_url: str,
+    chat: bool,
+) -> None:
+    """E2E latency must end at the last token, not at the usage trailer.
+
+    The usage chunk carries no token, so counting it stretches E2EL (and
+    therefore TPOT) by the delivery time of one extra SSE frame. The [DONE]
+    sentinel is already excluded on the same grounds.
+    """
+    tick = 1.0
+    clock = _ScriptedClock(tick=tick)
+    monkeypatch.setattr(request_func_module, "time", _FakeTime(clock))
+
+    n_tokens = 5
+    chunks = _build_stream(n_tokens, chat=chat)
+    output = _run(request_func, api_url, endpoint_id, chunks, n_tokens)
+
+    # Clock reads: 1 for the start time, then 1 per non-[DONE] chunk. The last
+    # token is therefore read at start + n_tokens ticks; the usage chunk one
+    # tick later.
+    expected_latency = n_tokens * tick
+    assert output.latency == pytest.approx(expected_latency, abs=tick / 1000), (
+        f"{endpoint_id}: latency {output.latency} should stop at the final "
+        f"token ({expected_latency}), not run on to the usage chunk "
+        f"({expected_latency + tick})"
+    )
+
+
+@pytest.mark.parametrize(
+    "endpoint_id,request_func,api_url,chat",
+    STREAMING_ENDPOINTS,
+    ids=[e[0] for e in STREAMING_ENDPOINTS],
+)
+def test_decode_span_identity_holds_with_real_clock(
+    endpoint_id: str,
+    request_func,
+    api_url: str,
+    chat: bool,
+) -> None:
+    """The identity also holds against the real clock, with real gaps.
+
+    The residual differences out the same recorded timestamps, so wall clock
+    jitter cancels and only float rounding remains; a microsecond bound is
+    ample rather than tight.
+    """
+    n_tokens = 4
+    chunks = _build_stream(
+        n_tokens, chat=chat, ttft_delay=0.02, itl_delay=0.01, usage_delay=0.02
+    )
+    output = _run(request_func, api_url, endpoint_id, chunks, n_tokens)
+
+    residual = (output.latency - output.ttft) - sum(output.itl)
+    assert abs(residual) < 1e-6, (
+        f"{endpoint_id}: (latency - ttft) - sum(itl) drifted by {residual:.9f}s"
+    )
+
+
+@pytest.mark.parametrize(
+    "endpoint_id,request_func,api_url,chat",
+    STREAMING_ENDPOINTS,
+    ids=[e[0] for e in STREAMING_ENDPOINTS],
+)
+def test_usage_only_stream_is_not_reported_as_success(
+    endpoint_id: str,
+    request_func,
+    api_url: str,
+    chat: bool,
+) -> None:
+    """A stream with no token chunk must fail rather than report 0 ms.
+
+    Only token chunks move the end of the request, so a degenerate stream
+    carrying nothing but a usage trailer leaves latency at zero. Reporting that
+    as a success would feed a zero-duration request into the benchmark
+    aggregates.
+    """
+    chunks = [(0.0, _usage_chunk(0)), (0.0, b"data: [DONE]\n\n")]
+    output = _run(request_func, api_url, endpoint_id, chunks, 0, expect_success=False)
+
+    assert not output.success
+    assert output.itl == []
+    assert "TTFT" in output.error
+
+
+# ``benchmarks/backend_request_func.py`` is a standalone script rather than a
+# package module, so it is loaded by path. It also builds its own ClientSession,
+# so the scripted stream is injected by swapping the module's aiohttp reference.
+_LEGACY_PATH = Path(__file__).parents[2] / "benchmarks" / "backend_request_func.py"
+
+
+class _FakeAiohttp:
+    """Stands in for the legacy module's aiohttp import."""
+
+    def __init__(self, chunks: list[tuple[float, bytes]]) -> None:
+        self._chunks = chunks
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(aiohttp, name)
+
+    def ClientSession(self, **kwargs: Any) -> _FakeSession:  # noqa: N802
+        del kwargs
+        return _FakeSession(self._chunks)
+
+
+def _load_legacy():
+    spec = importlib.util.spec_from_file_location(
+        "legacy_backend_request_func", _LEGACY_PATH
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+LEGACY_ENDPOINTS = [
+    (
+        "completions",
+        "async_request_openai_completions",
+        "http://test/v1/completions",
+        False,
+    ),
+    (
+        "chat",
+        "async_request_openai_chat_completions",
+        "http://test/v1/chat/completions",
+        True,
+    ),
+    (
+        "audio",
+        "async_request_openai_audio",
+        "http://test/v1/audio/transcriptions",
+        True,
+    ),
+]
+
+
+def _run_legacy(
+    monkeypatch, func_name, api_url, endpoint_id, chunks, n_tokens, clock=None
+):
+    module = _load_legacy()
+    monkeypatch.setattr(module, "aiohttp", _FakeAiohttp(chunks))
+    if clock is not None:
+        monkeypatch.setattr(module, "time", _FakeTime(clock))
+    request_func = getattr(module, func_name)
+    return asyncio.run(request_func(_make_input(api_url, n_tokens, endpoint_id)))
+
+
+@pytest.mark.parametrize(
+    "endpoint_id,func_name,api_url,chat",
+    LEGACY_ENDPOINTS,
+    ids=[e[0] for e in LEGACY_ENDPOINTS],
+)
+def test_legacy_decode_span_identity_is_exact(
+    monkeypatch: pytest.MonkeyPatch,
+    endpoint_id: str,
+    func_name: str,
+    api_url: str,
+    chat: bool,
+) -> None:
+    """benchmarks/backend_request_func.py must honour the same identity."""
+    tick = 1.0
+    n_tokens = 5
+    chunks = _build_stream(n_tokens, chat=chat)
+    output = _run_legacy(
+        monkeypatch,
+        func_name,
+        api_url,
+        endpoint_id,
+        chunks,
+        n_tokens,
+        clock=_ScriptedClock(tick=tick),
+    )
+
+    assert output.success, output.error
+    assert len(output.itl) == n_tokens - 1
+
+    residual = (output.latency - output.ttft) - sum(output.itl)
+    assert residual == pytest.approx(0.0, abs=tick / 1000), (
+        f"legacy {endpoint_id}: (latency - ttft) - sum(itl) = {residual!r}"
+    )
+    assert output.latency == pytest.approx(n_tokens * tick, abs=tick / 1000)
+
+
+@pytest.mark.parametrize(
+    "endpoint_id,func_name,api_url,chat",
+    LEGACY_ENDPOINTS,
+    ids=[e[0] for e in LEGACY_ENDPOINTS],
+)
+def test_legacy_usage_only_stream_is_not_reported_as_success(
+    monkeypatch: pytest.MonkeyPatch,
+    endpoint_id: str,
+    func_name: str,
+    api_url: str,
+    chat: bool,
+) -> None:
+    """A usage-only stream must fail rather than report a zero-duration success."""
+    chunks = [(0.0, _usage_chunk(0)), (0.0, b"data: [DONE]\n\n")]
+    output = _run_legacy(monkeypatch, func_name, api_url, endpoint_id, chunks, 0)
+
+    assert not output.success
+    assert output.itl == []
diff --git a/tests/benchmarks/test_throughput_cli.py b/tests/benchmarks/test_throughput_cli.py
index 6af4aad4d908..c4a39f24d948 100644
--- a/tests/benchmarks/test_throughput_cli.py
+++ b/tests/benchmarks/test_throughput_cli.py
@@ -17,6 +17,7 @@
 from vllm.benchmarks.datasets import SampleRequest
 from vllm.benchmarks.throughput import (
     _run_vllm_chat_requests,
+    _to_serve_args,
     add_cli_args,
     assign_loras,
     get_requests,
@@ -277,3 +278,91 @@ def sample(self, *, num_requests, tokenizer, **kwargs):  # noqa: ARG002
     requests = get_requests(args, hf_tokenizer)
     assert len(requests) == 3
     assert all(isinstance(r, SampleRequest) for r in requests)
+
+
+def test_to_serve_args_supports_mm_processor_flags() -> None:
+    """The adapter must serve callers without throughput's legacy flags.
+
+    ``bench mm-processor`` exposes neither --input-len/--output-len/--prefix-len
+    nor --backend, and leaves --random-prefix-len at its default of 0. The
+    adapter must not read the legacy flags unconditionally, and must keep the
+    explicit 0 instead of coercing it to None (which later crashed
+    RandomMultiModalDataset.get_prefix on ``None <= 0``).
+    """
+    args = SimpleNamespace(
+        random_input_len=32,
+        random_output_len=1,
+        random_prefix_len=0,
+    )
+    serve_args = _to_serve_args(args)
+
+    assert serve_args.random_input_len == 32
+    assert serve_args.random_output_len == 1
+    assert serve_args.random_prefix_len == 0
+    assert serve_args.backend == "vllm-chat"
+    assert serve_args.enable_multimodal_chat
+    assert serve_args.no_oversample is False
+
+
+def test_to_serve_args_keeps_legacy_flag_fallback() -> None:
+    """Callers that only set the legacy flags still get them mapped."""
+    args = SimpleNamespace(
+        input_len=256,
+        output_len=8,
+        prefix_len=16,
+        backend="vllm",
+    )
+    serve_args = _to_serve_args(args)
+
+    assert serve_args.random_input_len == 256
+    assert serve_args.random_output_len == 8
+    assert serve_args.random_prefix_len == 16
+    assert serve_args.backend == "vllm"
+    assert not serve_args.enable_multimodal_chat
+
+
+@pytest.mark.benchmark
+def test_get_requests_random_mm_with_mm_processor_args(
+    hf_tokenizer: PreTrainedTokenizerBase,
+) -> None:
+    """Regression: bench mm-processor's random-mm sampling must not crash.
+
+    get_requests feeds mm-processor's parsed args (no legacy length/backend
+    flags) through the shared get_samples dispatch. With the default
+    --random-prefix-len 0 this previously crashed inside
+    RandomMultiModalDataset.get_prefix with a TypeError.
+    """
+    from vllm.benchmarks.mm_processor import add_cli_args as add_mm_cli_args
+
+    parser = FlexibleArgumentParser()
+    add_mm_cli_args(parser)
+    args = parser.parse_args(
+        [
+            "--dataset-name",
+            "random-mm",
+            "--num-prompts",
+            "3",
+            "--random-input-len",
+            "32",
+            "--random-output-len",
+            "1",
+            "--random-mm-base-items-per-request",
+            "1",
+            "--random-mm-num-mm-items-range-ratio",
+            "0",
+            "--random-mm-limit-mm-per-prompt",
+            '{"image": 1}',
+            "--random-mm-bucket-config",
+            "{(224, 224, 1): 1.0}",
+        ]
+    )
+
+    requests = get_requests(args, hf_tokenizer)
+    assert len(requests) == 3
+    assert all(isinstance(r, SampleRequest) for r in requests)
+    # mm-processor drives llm.chat, so prompts must be chat-formatted with
+    # multimodal content attached.
+    for req in requests:
+        assert isinstance(req.prompt, list)
+        assert req.prompt[0]["role"] == "user"
+        assert any(item.get("type") == "image_url" for item in req.prompt[0]["content"])
diff --git a/tests/ci_envs.py b/tests/ci_envs.py
index f3a54f308cd8..171fac996a68 100644
--- a/tests/ci_envs.py
+++ b/tests/ci_envs.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-These envs only work for a small part of the tests, fix what you need!
-"""
+"""These envs only work for a small part of the tests, fix what you need!"""
 
 import os
 from collections.abc import Callable
diff --git a/tests/compile/backend.py b/tests/compile/backend.py
index 87f98946a8ad..167de6ad68a3 100644
--- a/tests/compile/backend.py
+++ b/tests/compile/backend.py
@@ -29,8 +29,7 @@
 
 
 class LazyInitPass(InductorPass):
-    """
-    If there's a pass that we want to initialize lazily in a test,
+    """If there's a pass that we want to initialize lazily in a test,
     we can wrap it in LazyInitPass, which will initialize the pass when invoked
     and then immediately invoke it.
     """
@@ -45,8 +44,7 @@ def __call__(self, graph: fx.Graph) -> None:
 
 
 class TestBackend:
-    """
-    This class provides a simple Inductor backend that can be used for testing.
+    """This class provides a simple Inductor backend that can be used for testing.
     It takes a list of custom passes and runs them after Inductor's passes.
     It also saves the graph before and after the custom passes for inspection.
 
diff --git a/tests/compile/conftest.py b/tests/compile/conftest.py
index 7d15b5c47e55..972ec3fef52d 100644
--- a/tests/compile/conftest.py
+++ b/tests/compile/conftest.py
@@ -11,8 +11,7 @@
 
 @pytest.fixture
 def mock_cuda_platform():
-    """
-    Fixture that returns a factory for creating mocked CUDA platforms.
+    """Fixture that returns a factory for creating mocked CUDA platforms.
 
     Usage:
         def test_something(mock_cuda_platform):
@@ -51,8 +50,7 @@ def is_device_capability_family(
 
 @pytest.fixture
 def mock_xpu_platform():
-    """
-    Fixture that returns a factory for creating mocked XPU platforms.
+    """Fixture that returns a factory for creating mocked XPU platforms.
 
     Usage:
         def test_something(mock_xpu_platform):
diff --git a/tests/compile/correctness_e2e/test_async_tp.py b/tests/compile/correctness_e2e/test_async_tp.py
index 97248d153f4c..6b8c26cadc7e 100644
--- a/tests/compile/correctness_e2e/test_async_tp.py
+++ b/tests/compile/correctness_e2e/test_async_tp.py
@@ -2,18 +2,22 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import json
+import logging
 
 import pytest
 
+from tests.compile.fusions_e2e.common import FUSION_LOG_PATTERNS
 from tests.models.registry import HF_EXAMPLE_MODELS
 from tests.utils import (
     compare_two_settings,
     create_new_process_for_each_test,
+    multi_gpu_test,
 )
 from vllm.config import (
     CompilationMode,
 )
 from vllm.platforms import current_platform
+from vllm.transformers_utils.repo_utils import hf_api
 from vllm.utils.flashinfer import has_flashinfer
 
 NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4"
@@ -26,6 +30,78 @@
 }
 
 
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract")
+def test_rocm_async_tp_bf16_output_correctness(caplog_mp_spawn):
+    """Trained-model outputs must agree with unfused TP, including padded decode."""
+    # A snapshot path pins both servers and the comparison helper's tokenizer.
+    model = hf_api().snapshot_download(
+        "Qwen/Qwen3-0.6B", revision="c1899de289a04d12100db370d81485cdf75e47ca"
+    )
+    common_args = [
+        "--dtype",
+        "bfloat16",
+        "--load-format",
+        "auto",
+        "--tensor-parallel-size",
+        "2",
+        "--distributed-executor-backend",
+        "mp",
+        "--max-model-len",
+        "512",
+        "--max-num-batched-tokens",
+        "512",
+        "--max-num-seqs",
+        "8",
+        "--kv-cache-memory-bytes",
+        str(64 * 1024**2),
+        "--gpu-memory-utilization",
+        "0.01",
+        "--attention-backend",
+        "TRITON_ATTN",
+        "--disable-custom-all-reduce",
+    ]
+
+    def settings(enabled):
+        return [
+            *common_args,
+            "--compilation-config",
+            json.dumps(
+                {
+                    "mode": CompilationMode.VLLM_COMPILE,
+                    "cudagraph_mode": "NONE",
+                    "splitting_ops": [],
+                    "inductor_compile_config": {"force_disable_caches": True},
+                    "pass_config": {
+                        "enable_sp": enabled,
+                        "fuse_gemm_comms": enabled,
+                        "fuse_allreduce_rms": False,
+                        "sp_min_token_num": 1,
+                    },
+                }
+            ),
+        ]
+
+    env = {"VLLM_ROCM_USE_AITER": "0", "VLLM_DISABLE_COMPILE_CACHE": "1"}
+    with caplog_mp_spawn(logging.DEBUG) as logs:
+        compare_two_settings(
+            model,
+            settings(False),
+            settings(True),
+            env1=env,
+            env2=env,
+            method="generate",
+            force_v1_runner=True,
+        )
+    matches = [
+        int(count) for count in FUSION_LOG_PATTERNS["async_tp"].findall(logs.text)
+    ]
+    print(f"AsyncTP rewrite counts: {matches}")
+    assert len(matches) >= 2 and all(count > 0 for count in matches), (
+        f"Expected positive AsyncTP rewrites on both TP workers, found {matches}"
+    )
+
+
 @create_new_process_for_each_test()
 @pytest.mark.parametrize(
     "model_id",
diff --git a/tests/compile/correctness_e2e/test_attn_quant.py b/tests/compile/correctness_e2e/test_attn_quant.py
new file mode 100644
index 000000000000..c7fc849ae52f
--- /dev/null
+++ b/tests/compile/correctness_e2e/test_attn_quant.py
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+import logging
+
+import pytest
+import regex as re
+
+from tests.utils import compare_two_settings
+from vllm._aiter_ops import is_aiter_found_and_supported
+from vllm.config import CompilationMode
+from vllm.platforms import current_platform
+from vllm.transformers_utils.repo_utils import hf_api
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm() or not is_aiter_found_and_supported(),
+    reason="Requires ROCm AITER attention and quantization",
+)
+def test_rocm_aiter_static_attention_output_correctness(caplog_mp_spawn):
+    """Trained FP8 outputs agree when actual AITER output quantization is fused."""
+    model = hf_api().snapshot_download(
+        "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
+        revision="870a1177ed082c0ba5dbefd62573ab8f91803f6a",
+        allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt", "*.jinja"],
+    )
+    common_args = [
+        "--dtype",
+        "bfloat16",
+        "--load-format",
+        "auto",
+        "--max-model-len",
+        "512",
+        "--max-num-batched-tokens",
+        "512",
+        "--max-num-seqs",
+        "8",
+        "--kv-cache-memory-bytes",
+        str(256 * 1024**2),
+        "--gpu-memory-utilization",
+        "0.01",
+        "--attention-backend",
+        "ROCM_AITER_UNIFIED_ATTN",
+    ]
+
+    def settings(enabled):
+        return [
+            *common_args,
+            "--compilation-config",
+            json.dumps(
+                {
+                    "mode": CompilationMode.VLLM_COMPILE,
+                    "cudagraph_mode": "NONE",
+                    "splitting_ops": [],
+                    "inductor_compile_config": {"force_disable_caches": True},
+                    "custom_ops": ["+quant_fp8", "-rms_norm"],
+                    "pass_config": {
+                        "fuse_attn_quant": enabled,
+                        "fuse_norm_quant": False,
+                        "fuse_act_quant": False,
+                        "fuse_allreduce_rms": False,
+                    },
+                }
+            ),
+        ]
+
+    env = {
+        "VLLM_ROCM_USE_AITER": "1",
+        "VLLM_ROCM_USE_AITER_LINEAR": "1",
+        "VLLM_DISABLE_COMPILE_CACHE": "1",
+    }
+    with caplog_mp_spawn(logging.DEBUG) as logs:
+        compare_two_settings(
+            model,
+            settings(False),
+            settings(True),
+            env1=env,
+            env2=env,
+            method="generate",
+            force_v1_runner=True,
+        )
+    matches = [
+        int(count) for count in re.findall(r"'attn_quant_fusion': (\d+)", logs.text)
+    ]
+    print(f"Attention output quantization rewrite counts: {matches}")
+    assert matches and all(count == 32 for count in matches), (
+        f"Expected output quantization fused in all 32 Llama layers, found {matches}"
+    )
diff --git a/tests/compile/correctness_e2e/test_sequence_parallel.py b/tests/compile/correctness_e2e/test_sequence_parallel.py
index eed5efb7b36b..e5f8d1870eb0 100644
--- a/tests/compile/correctness_e2e/test_sequence_parallel.py
+++ b/tests/compile/correctness_e2e/test_sequence_parallel.py
@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-WARNING: This test runs in both single-node (4 GPUs) and multi-node
- (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
- important to set the distributed backend to "mp" to avoid Ray scheduling
- all workers in a node other than the head node, which can cause the test
- to fail.
+"""WARNING: This test runs in both single-node (4 GPUs) and multi-node
+(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
+important to set the distributed backend to "mp" to avoid Ray scheduling
+all workers in a node other than the head node, which can cause the test
+to fail.
 """
 
 import json
diff --git a/tests/compile/fullgraph/test_full_cudagraph.py b/tests/compile/fullgraph/test_full_cudagraph.py
index c2af7be048af..33705899a28d 100644
--- a/tests/compile/fullgraph/test_full_cudagraph.py
+++ b/tests/compile/fullgraph/test_full_cudagraph.py
@@ -17,8 +17,7 @@
 
 @contextlib.contextmanager
 def temporary_environ(env_vars):
-    """
-    Temporarily set environment variables and restore them afterward.
+    """Temporarily set environment variables and restore them afterward.
     We have to do this vs monkeypatch because monkeypatch doesn't work
     with "module" scoped fixtures.
     """
@@ -37,7 +36,7 @@ def temporary_environ(env_vars):
 model_backends_full_cudagraph = []
 
 # deepseek-ai/DeepSeek-V2-Lite with MLA
-MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA"]
+MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA", "FlashInferMLA"]
 for mla_backend in MLA_backends:
     model_backends_full_cudagraph.append(
         ("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend])
@@ -54,27 +53,28 @@ def temporary_environ(env_vars):
 @pytest.fixture(scope="class")
 def llm_pair(request):
     model, backend_config, use_inductor_graph_partition = request.param
-    backend_config.comp_config["use_inductor_graph_partition"] = (
-        use_inductor_graph_partition
-    )
-
     if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
         pytest.skip("Inductor graph partition only supported in torch>=2.9")
 
+    backend = AttentionBackendEnum[backend_config.attention_config["backend"]]
+    if current_platform.is_rocm() and backend not in (
+        AttentionBackendEnum.TRITON_ATTN,
+        AttentionBackendEnum.ROCM_ATTN,
+    ):
+        pytest.skip(f"{backend_config.name} is a CUDA-only attention backend")
+    if backend == AttentionBackendEnum.ROCM_ATTN and not current_platform.is_rocm():
+        pytest.skip("ROCM_ATTN requires ROCm")
+
     # Dynamically skip test if GPU capability is not met
-    if (
-        backend_config.specific_gpu_arch
-        and backend_config.specific_gpu_arch != current_platform.get_device_capability()
+    if backend_config.specific_gpu_arch and (
+        not current_platform.is_cuda()
+        or backend_config.specific_gpu_arch != current_platform.get_device_capability()
     ):
         if backend_config.specific_gpu_arch == (9, 0):
             pytest.skip("Only Hopper GPUs support FA3 and FlashMLA")
         elif backend_config.specific_gpu_arch == (10, 0):
             pytest.skip("Only Blackwell GPUs support Cutlass MLA")
 
-    # FlashInfer is not supported on ROCm
-    if backend_config == AttentionBackendEnum.FLASHINFER and current_platform.is_rocm():
-        pytest.skip("FlashInfer is not supported on ROCm")
-
     env_vars = {
         # Force native sampler to avoid potential nondeterminism in FlashInfer
         # when per-request generators are not used in V1.
@@ -87,7 +87,11 @@ def llm_pair(request):
             trust_remote_code=True,
             max_model_len=1024,
             max_num_seqs=128,
-            compilation_config=CompilationConfig(**backend_config.comp_config),
+            attention_config=backend_config.attention_config,
+            compilation_config=CompilationConfig(
+                **backend_config.comp_config,
+                use_inductor_graph_partition=use_inductor_graph_partition,
+            ),
             generation_config="vllm",
             seed=42,
         )
@@ -97,8 +101,10 @@ def llm_pair(request):
             trust_remote_code=True,
             max_model_len=1024,
             max_num_seqs=128,
+            attention_config=backend_config.attention_config,
             compilation_config=CompilationConfig(
-                cudagraph_mode=CUDAGraphMode.PIECEWISE
+                cudagraph_mode=CUDAGraphMode.PIECEWISE,
+                use_inductor_graph_partition=use_inductor_graph_partition,
             ),
             generation_config="vllm",
             seed=42,
@@ -118,15 +124,17 @@ def llm_pair(request):
 @pytest.mark.parametrize(
     "llm_pair",
     [
-        pytest.param((model, backend_config, use_inductor_graph_partition))
+        pytest.param(
+            (model, backend_config, use_inductor_graph_partition),
+            id=f"{backend_config.name}-partition={use_inductor_graph_partition}",
+        )
         for model, backend_config in model_backends_full_cudagraph
         for use_inductor_graph_partition in [True, False]
     ],
     indirect=True,
 )
 class TestFullCUDAGraph:
-    """
-    Use a class such that an llm pair is constructed once for all
+    """Use a class such that an llm pair is constructed once for all
     batch_size/max_tokens combinations and released immediately after.
 
     Module-scope fixtures would stick around the whole time,
@@ -149,11 +157,9 @@ class TestFullCUDAGraph:
         ],
     )
     def test_full_cudagraph(self, batch_size, max_tokens, llm_pair: tuple[LLM, LLM]):
-        """
-        Test various batch sizes and max_tokens to ensure that the
+        """Test various batch sizes and max_tokens to ensure that the
         full cudagraph compilation works for padded cases too.
         """
-
         full_cudagraph_llm, piecewise_llm = llm_pair
 
         prompts = ["the quick brown fox"] * batch_size
diff --git a/tests/compile/fullgraph/test_multiple_graphs.py b/tests/compile/fullgraph/test_multiple_graphs.py
index 6d3788af9de0..e2656633344f 100644
--- a/tests/compile/fullgraph/test_multiple_graphs.py
+++ b/tests/compile/fullgraph/test_multiple_graphs.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test (piecewise) compilation with a simple model where multiple submodules
+"""Test (piecewise) compilation with a simple model where multiple submodules
 are compiled and graph captured separately.
 """
 
diff --git a/tests/compile/fullgraph/test_simple.py b/tests/compile/fullgraph/test_simple.py
index f1ea0e414d76..4d10a92a076f 100644
--- a/tests/compile/fullgraph/test_simple.py
+++ b/tests/compile/fullgraph/test_simple.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test the piecewise compilation with a simple model so that we
+"""Test the piecewise compilation with a simple model so that we
 can exactly calculate the expected output and side effects.
 """
 
@@ -52,8 +51,7 @@ def __init__(
         self.intermediate_unbacked = intermediate_unbacked
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Overall effect:
+        """Overall effect:
         x = 3 * x + 19
         global_counter += 2
         """
diff --git a/tests/compile/fullgraph/test_toy_llama.py b/tests/compile/fullgraph/test_toy_llama.py
index 69c758702e8a..7823c362f6c6 100644
--- a/tests/compile/fullgraph/test_toy_llama.py
+++ b/tests/compile/fullgraph/test_toy_llama.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test the piecewise compilation with a simple model, comparing the output
+"""Test the piecewise compilation with a simple model, comparing the output
 with and without the piecewise compilation.
 
 This is a tractable model, the weights and computation are specially designed
diff --git a/tests/compile/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py
index 77c3ba5112d9..1f92ff5a1988 100644
--- a/tests/compile/fusions_e2e/common.py
+++ b/tests/compile/fusions_e2e/common.py
@@ -128,8 +128,11 @@ def has_cuda_graph_wrapper_metadata() -> bool:
     ),
     "rms_quant_fusion": re.compile(r"rms_quant_fusion.py:\d+] Replaced (\d+) patterns"),
     "act_quant_fusion": re.compile(r"act_quant_fusion.py:\d+] Replaced (\d+) patterns"),
+    # The combined pass can consume these sites before the standalone pass.
     "norm_rope_fusion": re.compile(
-        r"qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on (\d+) sites"
+        r"(?:qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on|"
+        r"qk_norm_rope_kvcache_fusion.py:\d+] "
+        r"QK-Norm\+RoPE(?:/MRoPE)?\+KVCache fusion: replaced) (\d+)"
     ),
     "attn_quant_fusion": re.compile(
         r"attn_quant_fusion.py:\d+] Fused quant onto (\d+) attention nodes"
diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py
index f3df3c358ea2..a2cbe0135b32 100644
--- a/tests/compile/fusions_e2e/conftest.py
+++ b/tests/compile/fusions_e2e/conftest.py
@@ -182,7 +182,11 @@ def run(
             # TODO: Remove log counting in unit tests
             # once all matchers implement VllmFusionPatternMatcherPass
             n_expected = tp_size * num_ranges_activated
-            if match_name not in ("attn_quant_fusion", "act_quant_fusion"):
+            if match_name not in (
+                "attn_quant_fusion",
+                "act_quant_fusion",
+                "norm_rope_fusion",
+            ):
                 assert len(log_matches) == n_expected, (
                     f"Could not find {n_expected} {match_name} "
                     f"(found {len(log_matches)}) in:\n {log_holder.text}"
@@ -190,7 +194,16 @@ def run(
 
             expected_matches = getattr(matches, match_name)
 
-            if match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:
+            if match_name == "norm_rope_fusion":
+                # Opaque LayerName lets the combined pass consume short-range
+                # sites before the standalone pass; count either owner.
+                assert len(log_matches) >= n_expected
+                assert all(m in (0, expected_matches) for m in log_matches)
+                assert sum(log_matches) == expected_matches * n_expected, (
+                    f"{match_name} expected {expected_matches * n_expected} "
+                    f"sites across both passes, found: {log_matches}"
+                )
+            elif match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:
                 # AR+rms+quant takes precedence over rms+quant if activated.
                 # That means we get full matching where ar+rms+quant was not
                 # activated, and less where it was (only the smallest range).
diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py
index 7879cbe65096..eadf1739fdbb 100644
--- a/tests/compile/fusions_e2e/models.py
+++ b/tests/compile/fusions_e2e/models.py
@@ -44,6 +44,15 @@
     ),
 )
 
+ROCM_AITER_MLA_ATTN = pytest.param(
+    AttentionBackendCase(backend=AttentionBackendEnum.ROCM_AITER_MLA),
+    id="ROCM_AITER_MLA",
+    marks=pytest.mark.skipif(
+        not is_aiter_found_and_supported(),
+        reason="ROCM_AITER_MLA only for AMD when AITER is installed",
+    ),
+)
+
 FLASHINFER_MLA_ATTN = pytest.param(
     AttentionBackendCase(backend=AttentionBackendEnum.FLASHINFER_MLA),
     id="FLASHINFER_MLA",
diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py
index 4b3db06bc446..f35bc1bd9cea 100644
--- a/tests/compile/fusions_e2e/test_tp2_ar_rms.py
+++ b/tests/compile/fusions_e2e/test_tp2_ar_rms.py
@@ -19,6 +19,7 @@
 from .models import (
     FLASHINFER_ATTN,
     FLASHINFER_MLA_ATTN,
+    ROCM_AITER_MLA_ATTN,
     ROCM_AITER_UNIFIED_ATTN,
     ROCM_ATTN,
     TRITON_ATTN,
@@ -40,25 +41,35 @@
 )
 
 
-@multi_gpu_test(num_gpus=2)
-@pytest.mark.parametrize(
-    "model_name, matches_fn, model_kwargs, hf_overrides",
-    # qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet supported
+# qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet
+# supported. ROCm carries only the payloads its two backends support.
+TP2_FP8_MODELS = (
     [
         llama3_8b_fp8,
         llama4_scout_fp8,
         qwen3_a3b_fp8,
         deepseek_coder_v2_lite_fp8,
         deepseek_v3_fp8,
-    ],
+    ]
+    if current_platform.is_cuda()
+    else [llama3_8b_fp8, qwen3_a3b_fp8, deepseek_coder_v2_lite_fp8]
+)
+
+TP2_FP8_ATTN_BACKENDS = (
+    [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN]
+    if current_platform.is_cuda()
+    else [ROCM_AITER_UNIFIED_ATTN, ROCM_AITER_MLA_ATTN]
 )
+
+
+@multi_gpu_test(num_gpus=2)
 @pytest.mark.parametrize(
-    "attn_backend", [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN]
+    "model_name, matches_fn, model_kwargs, hf_overrides", TP2_FP8_MODELS
 )
+@pytest.mark.parametrize("attn_backend", TP2_FP8_ATTN_BACKENDS)
 @pytest.mark.parametrize("n_layers", [4])
 @pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
 def test_tp2_ar_rms_fp8_fusions(
     model_name: str,
     matches_fn: Callable[[int], Matches],
@@ -69,10 +80,27 @@ def test_tp2_ar_rms_fp8_fusions(
     custom_ops: str,
     inductor_graph_partition: bool,
     run_e2e_fusion_test,
+    monkeypatch,
 ):
     matches = matches_fn(n_layers)
 
-    block_fp8 = "qwen" in model_name.lower() or "deepseek" in model_name.lower()
+    model_name_lower = model_name.lower()
+    # DeepSeek-Coder on AITER MLA is the one static-FP8 payload here, validated
+    # against the native QuantFP8 pattern; every other DeepSeek stays group
+    # quantized as on CUDA.
+    rocm_static_fp8_mla = (
+        current_platform.is_rocm()
+        and "deepseek-coder" in model_name_lower
+        and attn_backend.backend.name == "ROCM_AITER_MLA"
+    )
+    block_fp8 = "qwen" in model_name_lower or "deepseek" in model_name_lower
+    if rocm_static_fp8_mla:
+        block_fp8 = False
+        if custom_ops != "-quant_fp8,-rms_norm":
+            pytest.skip("ROCm static-FP8 MLA parity uses native quant and RMSNorm")
+        # AITER's MoE GEMM rejects this reduced dummy model shape; keep AITER
+        # for MLA and the all-reduce, and let the MoE run on Triton.
+        monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
     if block_fp8 and "-quant_fp8" in custom_ops:
         # This is why config forces +quant_fp8 by default
         pytest.skip("native QuantFP8 matching not supported for group quant")
@@ -83,6 +111,13 @@ def test_tp2_ar_rms_fp8_fusions(
     model_kwargs["max_model_len"] = 1024
     model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
     model_kwargs["disable_custom_all_reduce"] = False
+    if current_platform.is_rocm() and "llama-3" in model_name_lower:
+        model_kwargs.update(
+            gpu_memory_utilization=0.01,
+            kv_cache_memory_bytes=64 * 1024**2,
+            max_num_batched_tokens=1024,
+            max_num_seqs=32,
+        )
 
     compilation_config = dict(
         use_inductor_graph_partition=inductor_graph_partition,
@@ -103,6 +138,19 @@ def test_tp2_ar_rms_fp8_fusions(
         "ar_rms_fusion",
     ]
 
+    use_aiter = current_platform.is_rocm()
+    if use_aiter:
+        matches = matches._replace(aiter_ar_rms_fusion=matches.ar_rms_fusion)
+        # The AITER all-reduce pass registers the AR+RMS(+group-quant) patterns
+        # before its AR+RMS-only ones, so the standalone AITER RMS-quant pass
+        # does not own the post-AR sites here. ROCm's contract is native
+        # QK-norm+RoPE and attention-output quant fusion plus AITER AR+RMS.
+        matches_check = [
+            "norm_rope_fusion",
+            "attn_quant_fusion",
+            "aiter_ar_rms_fusion",
+        ]
+
     run_e2e_fusion_test(
         model_name,
         matches,
@@ -111,6 +159,7 @@ def test_tp2_ar_rms_fp8_fusions(
         compilation_config,
         matches_check,
         tp_size=2,
+        use_aiter=use_aiter,
     )
 
 
diff --git a/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py b/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py
new file mode 100644
index 000000000000..3143143bec86
--- /dev/null
+++ b/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py
@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bounded ROCm model coverage for actual sequence-parallel AsyncTP rewrites."""
+
+import pytest
+
+from tests.utils import multi_gpu_test
+from vllm.config import PassConfig
+from vllm.platforms import current_platform
+
+from .common import INDUCTOR_GRAPH_PARTITION, custom_ops_combos
+from .models import TRITON_ATTN, llama3_8b
+
+
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract")
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
+@pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm")))
+@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
+def test_tp2_rocm_async_tp_bf16_fusions(
+    inductor_graph_partition, custom_ops, attn_backend, run_e2e_fusion_test
+):
+    n_layers = 4
+    model_kwargs = dict(
+        hf_overrides=dict(
+            num_hidden_layers=n_layers,
+            hidden_size=512,
+            intermediate_size=1024,
+            num_attention_heads=8,
+            num_key_value_heads=2,
+            head_dim=64,
+        ),
+        load_format="dummy",
+        max_model_len=1024,
+        kv_cache_memory_bytes=64 * 1024**2,
+        gpu_memory_utilization=0.01,
+        max_num_batched_tokens=1024,
+        max_num_seqs=32,
+    )
+    run_e2e_fusion_test(
+        llama3_8b.model_name,
+        llama3_8b.matches(n_layers),
+        model_kwargs,
+        attn_backend,
+        dict(
+            use_inductor_graph_partition=inductor_graph_partition,
+            custom_ops=custom_ops.split(","),
+            pass_config=PassConfig(
+                enable_qk_norm_rope_fusion=True,
+                enable_sp=True,
+                fuse_gemm_comms=True,
+                fuse_allreduce_rms=False,
+                sp_min_token_num=512,
+            ),
+        ),
+        ["norm_rope_fusion", "sequence_parallel", "async_tp"],
+        tp_size=2,
+    )
diff --git a/tests/compile/h100/test_startup.py b/tests/compile/h100/test_startup.py
index e57cf85b4a77..a56cf15e5c67 100644
--- a/tests/compile/h100/test_startup.py
+++ b/tests/compile/h100/test_startup.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Cold start and warm start tests for vLLM-compile.
 
-Cold start runs in a forked child (must fork before CUDA init) which
+Cold start runs in a separate child process which
 populates on-disk caches and asserts cold-start counters.  Warm start
 then runs in the parent with clean in-memory state but populated caches.
 """
@@ -18,7 +18,7 @@
 from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
 from vllm.utils.torch_utils import is_torch_equal_or_newer
 
-from ...utils import fork_new_process_for_each_test
+from ...utils import create_new_process_for_each_test, requires_spawn_multiprocessing
 
 MODEL = "microsoft/Phi-tiny-MoE-instruct"
 
@@ -54,17 +54,21 @@ def _cold_start(vllm_runner):
     assert counters["aot_autograd"]["autograd_cache_hit"] == 0
 
 
-@fork_new_process_for_each_test
+@create_new_process_for_each_test()
 @pytest.mark.parametrize("mega_aot_artifact", ["0", "1"])
 def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache, mega_aot_artifact):
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
     monkeypatch.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", mega_aot_artifact)
     monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
 
-    # Cold start in a forked child (must fork before CUDA init).
+    # ROCm and XPU cannot fork safely after device discovery.
     # This model has 32 identical transformer layers which produce
     # 33 subgraphs after splitting on attention — only 3 are unique.
-    ctx = mp.get_context("fork")
+    ctx = (
+        mp.get_context("spawn")
+        if requires_spawn_multiprocessing()
+        else mp.get_context("fork")
+    )
     p = ctx.Process(target=_cold_start, args=(vllm_runner,))
     p.start()
     p.join()
@@ -223,13 +227,17 @@ def _cold_start_model(vllm_runner, spec: ModelStartupSpec):
 
 
 @pytest.mark.parametrize("spec", MODEL_SPECS)
-@fork_new_process_for_each_test
+@create_new_process_for_each_test()
 def test_model_startup(monkeypatch, vllm_runner, fresh_vllm_cache, spec):
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
     monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
 
-    # Cold start in a forked child (must fork before CUDA init).
-    ctx = mp.get_context("fork")
+    # ROCm and XPU cannot fork safely after device discovery.
+    ctx = (
+        mp.get_context("spawn")
+        if requires_spawn_multiprocessing()
+        else mp.get_context("fork")
+    )
     p = ctx.Process(target=_cold_start_model, args=(vllm_runner, spec))
     p.start()
     p.join()
diff --git a/tests/compile/passes/conftest.py b/tests/compile/passes/conftest.py
new file mode 100644
index 000000000000..018747fd8cdc
--- /dev/null
+++ b/tests/compile/passes/conftest.py
@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager
+
+
+@pytest.fixture(autouse=True)
+def _workspace_manager_for_compile_passes(cleanup_fixture):
+    if not current_platform.is_rocm():
+        yield
+        return
+
+    # Release workspace tensors before cleanup_fixture flushes the allocator.
+    reset_workspace_manager()
+    if torch.accelerator.is_available():
+        init_workspace_manager(torch.device(0))
+    yield
+    reset_workspace_manager()
diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py
index 7b547c9720d1..c0fd4c00a8f7 100644
--- a/tests/compile/passes/distributed/test_async_tp.py
+++ b/tests/compile/passes/distributed/test_async_tp.py
@@ -56,10 +56,7 @@ def __init__(self, hidden_size=16, dtype=torch.bfloat16):
         torch.nn.init.normal_(self.gate_proj, std=0.02)
 
     def forward(self, hidden_states):
-        """
-        Forward pass implementing the mm + reduce scatter in the FX graph
-
-        """
+        """Forward pass implementing the mm + reduce scatter in the FX graph."""
         # Reshape input
         view = hidden_states.reshape(-1, self.hidden_size)
 
@@ -88,9 +85,7 @@ def __init__(self, hidden_size=16, dtype=torch.bfloat16):
         torch.nn.init.normal_(self.weight, std=0.02)
 
     def forward(self, hidden_states):
-        """
-        Forward pass implementing the mm + all gather in the FX graph
-        """
+        """Forward pass implementing the mm + all gather in the FX graph."""
         # Reshape input
         view = hidden_states.reshape(-1, self.hidden_size)
         all_gather = tensor_model_parallel_all_gather(view, dim=0)
@@ -122,10 +117,7 @@ def __init__(self, hidden_size=16, dtype=torch.bfloat16):
 
 class TestScaledMMRSModel(_BaseScaledMMModel):
     def forward(self, input: torch.Tensor):
-        """
-        Forward pass implementing the scaled_mm + reduce scatter in the FX graph
-
-        """
+        """Forward pass implementing the scaled_mm + reduce scatter in the FX graph."""
         fp8_input = input.to(FP8_DTYPE)
         scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32)
         scaled_mm = torch._scaled_mm(
@@ -147,9 +139,7 @@ def ops_in_model_after(self):
 
 class TestAGScaledMMModel(_BaseScaledMMModel):
     def forward(self, input: torch.Tensor):
-        """
-        Forward pass implementing the all gather + scaled_mm in the FX graph
-        """
+        """Forward pass implementing the all gather + scaled_mm in the FX graph."""
         # Reshape input
         fp8_input = input.to(FP8_DTYPE)
         all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0)
@@ -173,8 +163,7 @@ def ops_in_model_after(self):
 
 class TestCutlassScaledMMRSModel(_BaseScaledMMModel):
     def forward(self, input: torch.Tensor):
-        """
-        Forward pass implementing the cutlass_scaled_mm + reduce scatter
+        """Forward pass implementing the cutlass_scaled_mm + reduce scatter
         in the FX graph
 
         """
@@ -200,8 +189,7 @@ def ops_in_model_after(self):
 
 class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
     def forward(self, input: torch.Tensor):
-        """
-        Forward pass implementing the all gather + cutlass_scaled_mm
+        """Forward pass implementing the all gather + cutlass_scaled_mm
         in the FX graph
         """
         # Reshape input
@@ -315,6 +303,7 @@ def async_tp_pass_on_test_model(
     dtype: torch.dtype,
     dynamic: bool,
     distributed_init_method: str,
+    check_numerics: bool = False,
 ):
     set_random_seed(0)
 
@@ -375,12 +364,24 @@ def async_tp_pass_on_test_model(
         hidden_states = torch.randn(
             (batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False
         )
+        if check_numerics:
+            hidden_states += local_rank * 0.125
 
         if dynamic:
             torch._dynamo.mark_dynamic(hidden_states, 0)
 
         compiled_model = torch.compile(model, backend=backend)
-        compiled_model(hidden_states)
+        actual = compiled_model(hidden_states)
+
+        if check_numerics:
+            # These are new checks, so no existing tolerance is relaxed.
+            # Allow BF16 rounding vs. eager, but keep the lifetime check exact.
+            expected = model(hidden_states)
+            torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)
+            snapshot = actual.clone()
+            changed = compiled_model(-hidden_states)
+            torch.testing.assert_close(changed, -expected, rtol=1e-2, atol=1e-2)
+            torch.testing.assert_close(actual, snapshot, rtol=0, atol=0)
 
         assert async_tp_pass.matched_count == 1
 
diff --git a/tests/compile/passes/ir/test_clone_cleanup.py b/tests/compile/passes/ir/test_clone_cleanup.py
index b6626a5cb70a..8a51e6d53988 100644
--- a/tests/compile/passes/ir/test_clone_cleanup.py
+++ b/tests/compile/passes/ir/test_clone_cleanup.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Comprehensive tests for UnsafeCloneEliminationPass.
+"""Comprehensive tests for UnsafeCloneEliminationPass.
 
 This test suite exercises all possible valid FX graph patterns involving clones:
 1. Clone with no users (dead code)
diff --git a/tests/compile/passes/ir/test_inplace_functionalization.py b/tests/compile/passes/ir/test_inplace_functionalization.py
index 1e8d5662162f..773ea93ce609 100644
--- a/tests/compile/passes/ir/test_inplace_functionalization.py
+++ b/tests/compile/passes/ir/test_inplace_functionalization.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for IR inplace functionalization pass integration.
+"""Tests for IR inplace functionalization pass integration.
 
 This test suite verifies that the inplace functionalization pass, lowering pass,
 and clone cleanup pass work together correctly with donated buffer tracking.
@@ -108,8 +107,7 @@ def forward(
 
 
 class ModelWithTritonAfterMaybeInplace(nn.Module):
-    """
-    Model using maybe_inplace followed by a Triton kernel.
+    """Model using maybe_inplace followed by a Triton kernel.
     Test clone elimination can handle Triton in the graph
     """
 
@@ -410,8 +408,7 @@ def inner(*args):
     reason="Only test on cuda and rocm platform",
 )
 def test_piecewise_compilation_with_donated_buffers(monkeypatch, fresh_vllm_cache):
-    """
-    Test piecewise compilation with donated buffers across graph splits.
+    """Test piecewise compilation with donated buffers across graph splits.
     Utilizes a custom splitting op. Uses fresh cache to avoid compilation caching.
     """
     torch.set_default_device(current_platform.device_type)
diff --git a/tests/compile/passes/ir/test_lowering.py b/tests/compile/passes/ir/test_lowering.py
index b7ca55e7d1e6..0ac605e7961c 100644
--- a/tests/compile/passes/ir/test_lowering.py
+++ b/tests/compile/passes/ir/test_lowering.py
@@ -11,6 +11,7 @@
 )
 from vllm.config import get_current_vllm_config
 from vllm.ir import ops
+from vllm.model_executor.layers.activation import GeluAndMulSparse
 from vllm.platforms import current_platform
 
 from ...backend import TestBackend
@@ -34,6 +35,15 @@ def forward(self, x):
         return x6 + 3.0
 
 
+class SparseActivationModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.activation = GeluAndMulSparse(0.95, "tanh")
+
+    def forward(self, x):
+        return self.activation(x)
+
+
 @pytest.mark.parametrize("rms_provider", ops.rms_norm.supported_providers())
 def test_lowering_rms_norm(rms_provider, default_vllm_config):
     torch.set_default_device(current_platform.device_type)
@@ -67,3 +77,27 @@ def test_lowering_rms_norm(rms_provider, default_vllm_config):
 
     torch.testing.assert_close(output_unlowered, output)
     torch.testing.assert_close(output_unlowered, output2)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+def test_lowering_gelu_and_mul_sparse(default_vllm_config):
+    torch.set_default_device("cuda")
+    lowering_pass = VllmIRLoweringPass(get_current_vllm_config())
+    backend = TestBackend(lowering_pass)
+    model = SparseActivationModel()
+    x = torch.randn(8, 32768, dtype=torch.bfloat16)
+
+    with (
+        ops.gelu_and_mul_sparse.set_priority(["triton", "native"]),
+        ir.enable_torch_wrap(True),
+    ):
+        compiled_model = torch.compile(model, backend=backend, fullgraph=True)
+        output = compiled_model(x)
+
+    selected = lowering_pass.selected_impls["gelu_and_mul_sparse"]
+    assert selected == {"gelu_and_mul_sparse": "triton"}
+
+    expected = ops.gelu_and_mul_sparse.impls["native"].impl_fn(
+        x, 1.6448536269514722, "tanh"
+    )
+    torch.testing.assert_close(output, expected, rtol=1.6e-2, atol=1e-3)
diff --git a/tests/compile/passes/test_double_aiter_rms_quant_fusion.py b/tests/compile/passes/test_double_aiter_rms_quant_fusion.py
index 6a620d11a492..f05f821852c8 100644
--- a/tests/compile/passes/test_double_aiter_rms_quant_fusion.py
+++ b/tests/compile/passes/test_double_aiter_rms_quant_fusion.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for the DoubleQuant fan-out variants registered by
+"""Unit tests for the DoubleQuant fan-out variants registered by
 ``RocmAiterRMSNormQuantFusionPass``.
 
 Both variants target a 1-to-2 fan-out where one ``rms_norm`` output feeds
@@ -90,8 +89,7 @@ def test_double_aiter_rms_fp8_group_quant_fusion(
     model_cls: type[torch.nn.Module],
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """
-    Both fan-out shapes (with and without an intermediate view) must fuse
+    """Both fan-out shapes (with and without an intermediate view) must fuse
     into ``rocm_aiter_rmsnorm_fp8_group_quant``: the no-view shape via
     ``DoubleAiterRMSFp8GroupQuantPattern`` and the viewed shape via the
     new ``DoubleAiterRMSFp8GroupQuantViewPattern`` sibling.
diff --git a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py
index 6f20d4e15874..0f631b9472ad 100644
--- a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py
+++ b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit test for the MLADualRMSNormFusionPass.
+"""Unit test for the MLADualRMSNormFusionPass.
 
 The pass fuses paired q/kv RMS norms in MLA attention into a single
 fused_mla_dual_rms_norm op backed by AITER's fused_qk_rmsnorm kernel.
@@ -38,13 +37,12 @@
 
 
 class MLADualRMSNormTestModel(torch.nn.Module):
-    """
-    Minimal model reproducing the MLA dual RMS norm pattern:
-        linear -> split([q_dim, kv_dim])
-            +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear
-            +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])
-                    +-- kv_c (getitem 0) -> rms_norm(kv_w, eps)
-                    +-- k_pe
+    """Minimal model reproducing the MLA dual RMS norm pattern:
+    linear -> split([q_dim, kv_dim])
+        +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear
+        +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])
+                +-- kv_c (getitem 0) -> rms_norm(kv_w, eps)
+                +-- k_pe
     """
 
     def __init__(
@@ -155,13 +153,12 @@ def test_fuse_mla_dual_rms_norm(
 
 
 class MLADualRMSNormFp8PerTokenTestModel(torch.nn.Module):
-    """
-    Minimal model reproducing the FP8 MLA attention path with *per-token* quant:
-        linear -> split([q_dim, kv_dim])
-            +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fused_dynamic_quant -> dequant
-            +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])
-                    +-- kv_c (getitem 0) -> rms_norm (bf16)
-                    +-- k_pe
+    """Minimal model reproducing the FP8 MLA attention path with *per-token* quant:
+    linear -> split([q_dim, kv_dim])
+        +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fused_dynamic_quant -> dequant
+        +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])
+                +-- kv_c (getitem 0) -> rms_norm (bf16)
+                +-- k_pe
     """
 
     def __init__(
diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py
index 843052f4afdc..bb15022459ed 100644
--- a/tests/compile/passes/test_fusion.py
+++ b/tests/compile/passes/test_fusion.py
@@ -550,7 +550,9 @@ class TestGatedModel(torch.nn.Module):
     """Model that uses RMSNormGated + reshape + group FP8 quant + linear.
 
     Mimics GatedDeltaNetAttention's output projection path where:
-    - RMSNormGated operates on per-head tensors (N*H, D)
+    - RMSNormGated operates on per-head tensors, either (N, H, D) as
+      QwenGatedDeltaNetAttention passes them or flattened to (N*H, D) as
+      OlmoHybridGatedDeltaNetAttention does
     - Output is reshaped to (N, H*D) before group quantization + linear
     """
 
@@ -563,10 +565,12 @@ def __init__(
         group_shape: GroupShape,
         dtype: torch.dtype,
         use_aiter_quant: bool = True,
+        flatten_heads: bool = True,
     ):
         super().__init__()
         self.num_heads = num_heads
         self.head_dim = head_dim
+        self.flatten_heads = flatten_heads
         hidden_dim = num_heads * head_dim
 
         self.norm = RMSNormGated(
@@ -599,8 +603,11 @@ def forward(self, x, z):
         hidden_dim = num_heads * head_dim
         x = torch.relu(x)
         z = torch.relu(z)
-        x_heads = x.reshape(-1, num_heads, head_dim).reshape(-1, head_dim)
-        z_heads = z.reshape(-1, num_heads, head_dim).reshape(-1, head_dim)
+        x_heads = x.reshape(-1, num_heads, head_dim)
+        z_heads = z.reshape(-1, num_heads, head_dim)
+        if self.flatten_heads:
+            x_heads = x_heads.reshape(-1, head_dim)
+            z_heads = z_heads.reshape(-1, head_dim)
         normed = self.norm(x_heads, z_heads)
         merged = normed.reshape(-1, hidden_dim)
         out = self.fp8_linear(merged)
@@ -638,6 +645,8 @@ def __init__(self, num_v_heads: int, head_v_dim: int, tp_size: int = 1):
 @pytest.mark.parametrize("head_dim", [128])
 @pytest.mark.parametrize("num_tokens", [8])
 @pytest.mark.parametrize("eps", [1e-5, 1e-6])
+@pytest.mark.parametrize("flatten_heads", [True, False])
+@pytest.mark.parametrize("quant_fp8_op", ["-quant_fp8", "+quant_fp8"])
 @pytest.mark.skipif(
     (not current_platform.is_rocm() or not IS_AITER_FOUND),
     reason="Only test on ROCm with aiter package installed",
@@ -648,6 +657,8 @@ def test_aiter_fusion_rmsnorm_gated_quant(
     head_dim: int,
     num_tokens: int,
     eps: float,
+    flatten_heads: bool,
+    quant_fp8_op: str,
     monkeypatch: pytest.MonkeyPatch,
 ):
     group_shape = GroupShape(1, 128)
@@ -655,7 +666,7 @@ def test_aiter_fusion_rmsnorm_gated_quant(
         model_config=ModelConfig(dtype=dtype),
         compilation_config=CompilationConfig(
             mode=CompilationMode.VLLM_COMPILE,
-            custom_ops=["-rms_norm", "-silu_and_mul", "-quant_fp8"],
+            custom_ops=["-rms_norm", "-silu_and_mul", quant_fp8_op],
             pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True),
         ),
     )
@@ -688,6 +699,7 @@ def test_aiter_fusion_rmsnorm_gated_quant(
             group_shape=group_shape,
             dtype=dtype,
             use_aiter_quant=True,
+            flatten_heads=flatten_heads,
         )
 
         noop_pass = NoOpEliminationPass(vllm_config)
@@ -708,7 +720,11 @@ def test_aiter_fusion_rmsnorm_gated_quant(
         model_unfused = torch.compile(model, backend=backend2)
         result_unfused = model_unfused(x, z)
 
-        torch.testing.assert_close(result_fused, result_unfused, atol=1e-2, rtol=1e-2)
+        # The opaque AITER quant op preserves a BF16 intermediate in the
+        # unfused path, while the fused kernel normalizes and quantizes
+        # directly. Allow for that extra rounding before the FP8 GEMM.
+        tol = 2e-2 if quant_fp8_op == "+quant_fp8" else 1e-2
+        torch.testing.assert_close(result_fused, result_unfused, atol=tol, rtol=tol)
 
         assert fusion_pass.matched_count == 1
         backend.check_after_ops(model.ops_in_model_after())
diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py
index 4df02cab9543..ccc7e624dbed 100644
--- a/tests/compile/passes/test_fusion_attn.py
+++ b/tests/compile/passes/test_fusion_attn.py
@@ -12,6 +12,7 @@
     create_common_attn_metadata,
     dense_kv_cache_views,
 )
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
 from vllm.compilation.passes.fusion.attn_quant_fusion import (
     ATTN_OP,
@@ -107,7 +108,6 @@ def __init__(
 
     def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
         """Initialize attention metadata."""
-
         # TODO (Rohan138) reuse utils from vllm/v1/worker/gpu/attn_utils.py
 
         # Create common attn metadata
@@ -190,6 +190,17 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
         return self.fp8_linear(attn_output)
 
 
+class TestAttentionAiterFp8StaticQuantPatternModel(
+    TestAttentionFp8StaticQuantPatternModel
+):
+    aiter_quant = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.fp8_linear.kernel.quant_fp8.use_aiter = True
+        self.fp8_linear.input_scale.fill_(0.125)
+
+
 class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
     """Test model for AttentionNvfp4QuantPattern fusion."""
 
@@ -304,7 +315,7 @@ def test_attention_quant_pattern(
     dist_init,
     disable_vllm_compile_cache,
 ):
-    """Test AttentionStaticQuantPattern fusion pass"""
+    """Test AttentionStaticQuantPattern fusion pass."""
     if backend == AttentionBackendEnum.FLASHINFER and (
         not current_platform.is_device_capability((10, 0)) or not has_flashinfer()
     ):
@@ -456,11 +467,12 @@ def test_attention_quant_pattern(
         )
 
     # Check quantization ops in the graph before and after fusion
-    quant_op = (
-        torch.ops.aten.reciprocal
-        if "-quant_fp8" in custom_ops_list
-        else QUANT_OPS[quant_key]
-    )
+    if "-quant_fp8" in custom_ops_list:
+        quant_op = torch.ops.aten.reciprocal
+    elif getattr(model_class, "aiter_quant", False):
+        quant_op = torch.ops.vllm.rocm_aiter_per_tensor_quant.default
+    else:
+        quant_op = QUANT_OPS[quant_key]
 
     if expect_fusion:
         # Note: for fp8, fully_replaced=False because query quant ops remain in
@@ -524,3 +536,46 @@ def test_attention_quant_pattern(
 
     # Check that results are close
     torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)
+    if getattr(model_class, "aiter_quant", False):
+        # Query quantization remains; exactly the output quant op is removed.
+        assert test_backend.op_count(quant_op, before=True) == (
+            test_backend.op_count(quant_op) + 1
+        )
+        for model in (model_unfused, model_fused):
+            torch.testing.assert_close(
+                model.fp8_linear.input_scale,
+                torch.full_like(model.fp8_linear.input_scale, 0.125),
+                atol=0,
+                rtol=0,
+            )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm() or not is_aiter_found_and_supported(),
+    reason="Requires ROCm AITER attention and quantization",
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_rocm_aiter_static_quant_attention(
+    dtype, dist_init, disable_vllm_compile_cache, monkeypatch
+):
+    """Fuse the actual AITER op without changing its static scale or outputs."""
+    try:
+        with monkeypatch.context() as patch:
+            patch.setenv("VLLM_ROCM_USE_AITER", "1")
+            rocm_aiter_ops.refresh_env_variables()
+            test_attention_quant_pattern(
+                num_qo_heads=32,
+                num_kv_heads=8,
+                head_size=128,
+                batch_size=8,
+                dtype=dtype,
+                kv_cache_dtype="fp8",
+                custom_ops="+quant_fp8",
+                model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
+                model_class=TestAttentionAiterFp8StaticQuantPatternModel,
+                backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+                dist_init=dist_init,
+                disable_vllm_compile_cache=disable_vllm_compile_cache,
+            )
+    finally:
+        rocm_aiter_ops.refresh_env_variables()
diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py
index 11c455bfa2b5..f1de3dde0cbb 100644
--- a/tests/compile/passes/test_mla_attn_quant_fusion.py
+++ b/tests/compile/passes/test_mla_attn_quant_fusion.py
@@ -141,7 +141,6 @@ def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
         the full output[:num_actual_toks] buffer after both forward_mha and
         forward_mqa have written their results.
         """
-
         batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
         common_attn_metadata = create_common_attn_metadata(
             batch_spec, self.block_size, self.device, arange_block_indices=True
@@ -408,7 +407,7 @@ def test_mla_attention_quant_pattern(
     dist_init,
     disable_vllm_compile_cache,
 ):
-    """Test MLA AttentionQuantPattern fusion pass"""
+    """Test MLA AttentionQuantPattern fusion pass."""
     if (
         model_class is TestMLAAttentionNvfp4QuantPatternModel
         and not is_nvfp4_supported()
diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py
index 6965346a5843..ee557235fdf1 100644
--- a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py
+++ b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py
@@ -8,6 +8,7 @@
 import torch
 from packaging.version import Version
 
+import vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion as fusion_module
 import vllm.config
 from tests.compile.backend import TestBackend
 from tests.v1.attention.utils import (
@@ -15,8 +16,12 @@
     create_common_attn_metadata,
     dense_kv_cache_views,
 )
-from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
-from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
+from vllm._aiter_ops import (
+    IS_AITER_FOUND,
+    is_aiter_found_and_supported,
+    rocm_aiter_ops,
+)
+from vllm.compilation.passes.fusion.matcher_utils import MROPE_OP, ROTARY_OP
 from vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion import (
     QkNormRopeKvCacheFusionPass,
 )
@@ -37,8 +42,12 @@
 from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import _USE_LAYERNAME, _encode_layer_name
 from vllm.v1.attention.backend import (
     AttentionBackend,
     CommonAttentionMetadata,
@@ -94,6 +103,7 @@ def __init__(
         device: torch.device,
         rotary_dim: int | None = None,
         prefix: str = "model.layers.0.self_attn.attn",
+        expected_query_quant_group_shape: GroupShape | None = None,
     ):
         super().__init__()
         self.num_heads = num_heads
@@ -133,6 +143,9 @@ def __init__(
             attn_backend=attn_backend.get_class(),
         )
         self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
+        if expected_query_quant_group_shape is not None:
+            assert self.attn.query_quant is not None
+            assert self.attn.query_quant.group_shape == expected_query_quant_group_shape
         assert not self.attn_backend.forward_includes_kv_cache_update, (
             f"Attention backend {self.attn_backend} does not support "
             "fuse_qk_norm_rope_kvcache."
@@ -227,18 +240,15 @@ def forward(
             self.kv_cache_dtype != self.dtype
             and self.attn.impl.supports_quant_query_input
         ):
-            q_fp8 = torch.empty_like(q, dtype=FP8_DTYPE)
-            torch.ops.vllm.rocm_aiter_per_tensor_quant(
-                q_fp8, q, self.attn._q_scale, False
-            )
-            q = q_fp8
+            assert self.attn.query_quant is not None
+            q, _ = self.attn.query_quant(q, self.attn._q_scale)
 
         # Final views + KV cache update
         q = q.view(-1, self.num_heads, self.head_size)
         k = k.view(-1, self.num_kv_heads, self.head_size)
         v = v.view(-1, self.num_kv_heads, self.head_size)
         kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
-            k, v, self.layer_name
+            k, v, _encode_layer_name(self.layer_name)
         )
         return q, k, v, kv_cache_dummy_dep
 
@@ -260,6 +270,64 @@ def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
         return [torch.ops.vllm.fused_qk_norm_rope_and_unified_kv_cache_update.default]
 
 
+class QKNormMRoPEKVCacheTestModel(QKNormRoPEKVCacheTestModel):
+    def __init__(
+        self,
+        *,
+        mrope_section: tuple[int, int, int],
+        mrope_interleaved: bool,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.rotary_emb = MRotaryEmbedding(
+            self.head_size,
+            rotary_dim=self.rotary_dim,
+            max_position_embeddings=4096,
+            base=10000,
+            is_neox_style=self.is_neox,
+            dtype=self.dtype,
+            mrope_section=list(mrope_section),
+            mrope_interleaved=mrope_interleaved,
+        )
+        self.enable_rope_custom_op = self.rotary_emb.enabled()
+
+    def forward(
+        self, qkv: torch.Tensor, positions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        q, k, v, kv_cache_dummy_dep = super().forward(qkv, positions)
+        output = torch.empty(
+            q.shape[0],
+            self.num_heads * self.head_size,
+            device=q.device,
+            dtype=self.dtype,
+        )
+        torch.ops.vllm.unified_attention_with_output(
+            q,
+            k,
+            v,
+            output,
+            _encode_layer_name(self.layer_name),
+            kv_cache_dummy_dep=kv_cache_dummy_dep,
+        )
+        # Production decoder attention consumes K from the cache rather than
+        # its direct K argument. Returning `output` keeps that production user
+        # in the graph without exposing K to an unsupported consumer.
+        return q, output, v, kv_cache_dummy_dep
+
+    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
+        return [MROPE_OP, torch.ops.vllm.unified_kv_cache_update.default]
+
+    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
+        return [torch.ops.vllm.fused_qk_norm_mrope_and_unified_kv_cache_update.default]
+
+
+class QKNormMRoPEObservableKTestModel(QKNormMRoPEKVCacheTestModel):
+    def forward(
+        self, qkv: torch.Tensor, positions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        return QKNormRoPEKVCacheTestModel.forward(self, qkv, positions)
+
+
 def _run_qk_norm_rope_kvcache_fusion_test(
     *,
     attn_backend: AttentionBackendEnum,
@@ -278,18 +346,33 @@ def _run_qk_norm_rope_kvcache_fusion_test(
     rms_norm_eps: float,
     custom_op: str,
     monkeypatch: pytest.MonkeyPatch,
+    mrope_section: tuple[int, int, int] | None = None,
+    mrope_interleaved: bool = False,
+    observe_mrope_k: bool = False,
+    expect_fusion: bool = True,
+    use_quark_scalar_query_scale: bool = False,
 ) -> None:
     device = os.environ.get("VLLM_TEST_CUDA_DEVICE", "cuda")
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)
 
+    # Set AITER toggles before VllmConfig initialization enables env caching.
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1")
+    rocm_aiter_ops.refresh_env_variables()
+
     vllm_config = VllmConfig(
         model_config=ModelConfig(dtype=dtype),
         cache_config=CacheConfig(
             block_size=block_size,
             cache_dtype=kv_cache_dtype,
         ),
+        quant_config=(
+            QuarkConfig(quant_config={"exclude": []})
+            if use_quark_scalar_query_scale
+            else None
+        ),
         compilation_config=CompilationConfig(
             mode=CompilationMode.VLLM_COMPILE,
             custom_ops=[custom_op],
@@ -302,6 +385,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
 
     with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
         m.setenv("VLLM_ROCM_USE_AITER", "1")
+        m.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1")
         m.setenv(
             "VLLM_ROCM_USE_AITER_TRITON_ROPE",
             "1" if enable_aiter_triton_rope else "0",
@@ -309,18 +393,49 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         m.setenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", use_shuffle_kv_layout)
         rocm_aiter_ops.refresh_env_variables()
 
-        model = QKNormRoPEKVCacheTestModel(
-            vllm_config=vllm_config,
-            attn_backend=attn_backend,
-            num_heads=num_heads,
-            num_kv_heads=num_kv_heads,
-            head_size=head_size,
-            rotary_dim=rotary_dim,
-            is_neox=is_neox,
-            rms_norm_eps=rms_norm_eps,
-            dtype=dtype,
-            device=torch.get_default_device(),
-        )
+        expected_query_quant_group_shape = None
+        if use_quark_scalar_query_scale:
+            expected_query_quant_group_shape = (
+                GroupShape(-1, num_heads * head_size)
+                if num_kv_heads == 1
+                else GroupShape.PER_TENSOR
+            )
+
+        model_kwargs = {
+            "vllm_config": vllm_config,
+            "attn_backend": attn_backend,
+            "num_heads": num_heads,
+            "num_kv_heads": num_kv_heads,
+            "head_size": head_size,
+            "rotary_dim": rotary_dim,
+            "is_neox": is_neox,
+            "rms_norm_eps": rms_norm_eps,
+            "dtype": dtype,
+            "device": torch.get_default_device(),
+            "expected_query_quant_group_shape": expected_query_quant_group_shape,
+        }
+        if mrope_section is None:
+            model = QKNormRoPEKVCacheTestModel(**model_kwargs)
+        else:
+            alternate_section = (8, 8, 8)
+            m.setattr(
+                fusion_module,
+                "_discover_mrope_configs",
+                lambda _: (
+                    (mrope_section, mrope_interleaved),
+                    (alternate_section, not mrope_interleaved),
+                ),
+            )
+            model_cls = (
+                QKNormMRoPEObservableKTestModel
+                if observe_mrope_k
+                else QKNormMRoPEKVCacheTestModel
+            )
+            model = model_cls(
+                **model_kwargs,
+                mrope_section=mrope_section,
+                mrope_interleaved=mrope_interleaved,
+            )
 
         fusion_pass = QkNormRopeKvCacheFusionPass(vllm_config)
         passes = [
@@ -337,7 +452,12 @@ def _run_qk_norm_rope_kvcache_fusion_test(
             num_heads * head_size + 2 * num_kv_heads * head_size,
             dtype=dtype,
         )
-        pos = torch.arange(num_tokens, dtype=torch.long)
+        if mrope_section is None:
+            pos = torch.arange(num_tokens, dtype=torch.long)
+        else:
+            # Real Qwen-VL position IDs can be a strided [3, N] view.
+            pos = torch.randint(0, 4096, (3, 2 * num_tokens), dtype=torch.long)[:, ::2]
+            assert not pos.is_contiguous()
 
         qkv_unfused = qkv.clone()
         pos_unfused = pos.clone()
@@ -356,7 +476,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
 
         # Run fused (compiled) forward
         torch._dynamo.mark_dynamic(qkv, 0)
-        torch._dynamo.mark_dynamic(pos, 0)
+        torch._dynamo.mark_dynamic(pos, 0 if mrope_section is None else 1)
         with set_forward_context(None, vllm_config):
             model_fused = torch.compile(model, backend=backend)
             forward_context = get_forward_context()
@@ -369,14 +489,33 @@ def _run_qk_norm_rope_kvcache_fusion_test(
             kv_cache_fused = attn_layer.kv_cache
         del dummy
 
-        assert fusion_pass.matched_count == 1
+        assert fusion_pass.matched_count == int(expect_fusion)
 
-        backend.check_before_ops(model.ops_in_model_before())
-        backend.check_after_ops(model.ops_in_model_after())
+        if expect_fusion:
+            backend.check_before_ops(model.ops_in_model_before())
+            backend.check_after_ops(model.ops_in_model_after())
+        else:
+            for op in model.ops_in_model_before():
+                assert backend.op_count(op, before=True) > 0
+                assert backend.op_count(op) > 0
+
+        if use_quark_scalar_query_scale:
+            expected_quant_op = (
+                torch.ops.vllm.rocm_aiter_per_tensor_quant.default
+                if num_kv_heads != 1
+                else torch.ops._C.static_scaled_fp8_quant.default
+            )
+            assert backend.op_count(expected_quant_op, before=True) > 0
+            # Query quantization remains separate after the prologue fusion.
+            assert backend.op_count(expected_quant_op) > 0
 
         # Sweep-backed (18.2k pts, PR #42749): native-rope ref worst 7.7e-3 -> 1e-2;
         # AITER-triton-rope ref is itself approximate (plateau 1.28e-2) -> 2e-2.
-        ATOL, RTOL = (2e-2, 2e-2) if enable_aiter_triton_rope else (1e-2, 1e-2)
+        ATOL, RTOL = (
+            (2e-2, 2e-2)
+            if enable_aiter_triton_rope or mrope_section is not None
+            else (1e-2, 1e-2)
+        )
         is_fp8_cache = model.kv_cache_dtype != dtype
 
         if q_fused.dtype == FP8_DTYPE:
@@ -387,7 +526,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         else:
             torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
 
-        if not is_fp8_cache:
+        if not expect_fusion or (not is_fp8_cache and mrope_section is None):
             # The AITER PTS kernel populates k_out only for non-FP8 caches.
             # With FP8, the kernel writes quantized K directly to the cache
             # and may leave k_out uninitialised.  In production this is fine
@@ -397,15 +536,16 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         # Should be bit exact since no processing had been done on v for both paths
         torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0)
 
-        # fp8 vs triton-rope ref requires loosening tolerance to 1.25e-1.
-        if is_fp8_cache and enable_aiter_triton_rope:
+        # Fused and unfused arithmetic can straddle an FP8 rounding boundary.
+        # Allow one E4M3 quantization step while still comparing every block.
+        if is_fp8_cache:
             cache_atol = cache_rtol = 1.25e-1
         else:
             cache_atol, cache_rtol = ATOL, RTOL
 
         torch.testing.assert_close(
-            kv_cache_unfused[0].float(),
-            kv_cache_fused[0].float(),
+            kv_cache_unfused.float(),
+            kv_cache_fused.float(),
             atol=cache_atol,
             rtol=cache_rtol,
         )
@@ -494,3 +634,225 @@ def test_qk_norm_rope_kvcache_fusion(
         custom_op=custom_op,
         monkeypatch=monkeypatch,
     )
+
+
+@pytest.mark.parametrize(
+    "mrope_section",
+    [
+        pytest.param((16, 24, 24), id="full_rotary"),
+        pytest.param((4, 6, 6), id="partial_rotary"),
+    ],
+)
+@pytest.mark.parametrize("mrope_interleaved", [False, True])
+@pytest.mark.parametrize("is_neox", [False, True])
+@pytest.mark.parametrize(
+    "kv_layout",
+    [
+        pytest.param(KVCacheLayout.LBHNC, id="head_major"),
+        pytest.param(KVCacheLayout.LBNHC, id="token_major"),
+        pytest.param(KVCacheLayout.LHBNC, id="block_major"),
+    ],
+)
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+@pytest.mark.skipif(
+    not IS_AITER_FOUND,
+    reason="Requires AITER",
+)
+def test_qk_norm_mrope_kvcache_fusion(
+    mrope_section: tuple[int, int, int],
+    mrope_interleaved: bool,
+    is_neox: bool,
+    kv_layout: KVCacheLayout,
+    kv_cache_dtype: str,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _run_qk_norm_rope_kvcache_fusion_test(
+        attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        enable_aiter_triton_rope=False,
+        num_tokens=5,
+        num_heads=8,
+        num_kv_heads=2,
+        head_size=128,
+        rotary_dim=2 * sum(mrope_section),
+        block_size=16,
+        is_neox=is_neox,
+        use_shuffle_kv_layout="0",
+        kv_layout=kv_layout,
+        dtype=torch.bfloat16,
+        kv_cache_dtype=kv_cache_dtype,
+        rms_norm_eps=1e-6,
+        custom_op="+rotary_embedding",
+        monkeypatch=monkeypatch,
+        mrope_section=mrope_section,
+        mrope_interleaved=mrope_interleaved,
+    )
+
+
+@pytest.mark.parametrize(
+    ("num_heads", "num_kv_heads"),
+    [(32, 2), (16, 1), (8, 1)],
+    ids=["tp2", "tp4", "tp8"],
+)
+def test_qk_norm_mrope_kvcache_fusion_quark_scalar_query_scale(
+    num_heads: int,
+    num_kv_heads: int,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _run_qk_norm_rope_kvcache_fusion_test(
+        attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        enable_aiter_triton_rope=False,
+        num_tokens=5,
+        num_heads=num_heads,
+        num_kv_heads=num_kv_heads,
+        head_size=128,
+        rotary_dim=128,
+        block_size=16,
+        is_neox=True,
+        use_shuffle_kv_layout="0",
+        kv_layout=KVCacheLayout.LBHNC,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="fp8",
+        rms_norm_eps=1e-6,
+        custom_op="+rotary_embedding",
+        monkeypatch=monkeypatch,
+        mrope_section=(24, 20, 20),
+        mrope_interleaved=True,
+        use_quark_scalar_query_scale=True,
+    )
+
+
+def test_mrope_fusion_rejects_observable_k(monkeypatch: pytest.MonkeyPatch):
+    _run_qk_norm_rope_kvcache_fusion_test(
+        attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        enable_aiter_triton_rope=False,
+        num_tokens=5,
+        num_heads=8,
+        num_kv_heads=2,
+        head_size=128,
+        rotary_dim=128,
+        block_size=16,
+        is_neox=True,
+        use_shuffle_kv_layout="0",
+        kv_layout=KVCacheLayout.LBHNC,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="auto",
+        rms_norm_eps=1e-6,
+        custom_op="+rotary_embedding",
+        monkeypatch=monkeypatch,
+        mrope_section=(16, 24, 24),
+        observe_mrope_k=True,
+        expect_fusion=False,
+    )
+
+
+def test_mrope_fusion_rejects_unaligned_rotary_dim(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _run_qk_norm_rope_kvcache_fusion_test(
+        attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        enable_aiter_triton_rope=False,
+        num_tokens=5,
+        num_heads=8,
+        num_kv_heads=2,
+        head_size=128,
+        rotary_dim=6,
+        block_size=16,
+        is_neox=True,
+        use_shuffle_kv_layout="0",
+        kv_layout=KVCacheLayout.LBHNC,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="auto",
+        rms_norm_eps=1e-6,
+        custom_op="+rotary_embedding",
+        monkeypatch=monkeypatch,
+        mrope_section=(1, 1, 1),
+        expect_fusion=False,
+    )
+
+
+@pytest.mark.skipif(not _USE_LAYERNAME, reason="requires opaque LayerName")
+def test_opaque_layer_name_patterns_register_once(monkeypatch: pytest.MonkeyPatch):
+    device = os.environ.get("VLLM_TEST_CUDA_DEVICE", "cuda")
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(dtype=torch.bfloat16),
+        cache_config=CacheConfig(block_size=16, cache_dtype="fp8_e4m3"),
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["+rotary_embedding"],
+            pass_config=PassConfig(fuse_qk_norm_rope_kvcache=True),
+        ),
+    )
+
+    with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
+        m.setenv("VLLM_ROCM_USE_AITER", "1")
+        rocm_aiter_ops.refresh_env_variables()
+        model_kwargs = {
+            "vllm_config": vllm_config,
+            "attn_backend": AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+            "num_heads": 8,
+            "num_kv_heads": 2,
+            "head_size": 128,
+            "is_neox": True,
+            "rms_norm_eps": 1e-6,
+            "dtype": torch.bfloat16,
+            "device": torch.device(device),
+        }
+        models = [
+            QKNormRoPEKVCacheTestModel(
+                **model_kwargs,
+                prefix=f"model.layers.{index}.self_attn.attn",
+            )
+            for index in range(2)
+        ]
+
+        # Opaque layer names are wildcards in the pattern graph. Registering
+        # once per physical layer raises "Duplicate pattern" on torch >= 2.11.
+        fusion_pass = QkNormRopeKvCacheFusionPass(vllm_config)
+
+    assert models
+    assert fusion_pass.patterns
+
+
+def _call_fused_mrope_impl(q_out: torch.Tensor) -> None:
+    device = q_out.device
+    fusion_module.fused_qk_norm_mrope_and_unified_kv_cache_update_impl(
+        q_out=q_out,
+        qkv=torch.zeros(2, 16, device=device),
+        positions=torch.arange(2, device=device),
+        q_weight=torch.ones(4, device=device),
+        k_weight=torch.ones(4, device=device),
+        rms_norm_eps=1e-6,
+        cos_sin_cache=torch.zeros(16, 4, device=device),
+        is_neox=True,
+        mrope_section_t=1,
+        mrope_section_h=1,
+        mrope_section_w=0,
+        is_interleaved=False,
+        rotary_dim=4,
+        layer_name="model.layers.0.self_attn.attn",
+    )
+
+
+def test_mrope_profiling_pass_initializes_query(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        fusion_module,
+        "get_attention_context",
+        lambda _: (None, None, None, None),
+    )
+    q_out = torch.full((2, 8), torch.nan, device="cpu")
+
+    _call_fused_mrope_impl(q_out)
+
+    torch.testing.assert_close(q_out, torch.zeros_like(q_out))
+
+
+def test_mrope_attention_context_failure_propagates(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    def fail_context(_):
+        raise RuntimeError("missing attention context")
+
+    monkeypatch.setattr(fusion_module, "get_attention_context", fail_context)
+
+    with pytest.raises(RuntimeError, match="missing attention context"):
+        _call_fused_mrope_impl(torch.empty(2, 8, device="cpu"))
diff --git a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py
index 5ee749659b3f..10dc401a26a1 100644
--- a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py
+++ b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for manual fusion via maybe_fused_act_quant.
+"""Tests for manual fusion via maybe_fused_act_quant.
 
 Tests all fusion paths in _FUSED_ACT_QUANT:
 - kFp8StaticTensorSym: all platforms
@@ -34,6 +33,7 @@
 from vllm.model_executor.layers.fusion.quant_activation import (
     QuantizedActivation,
     expose_input_quant_key,
+    get_input_quant_key,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8Dynamic128Sym,
@@ -45,7 +45,7 @@
 
 # Mock linear layer for testing fusion paths that don't have real kernel support
 class MockLinearForFusion(torch.nn.Module):
-    """Mock linear layer that exposes input_quant_key for fusion testing."""
+    """Mock linear layer that exposes an input key for fusion testing."""
 
     def __init__(
         self,
@@ -55,7 +55,7 @@ def __init__(
         input_global_scale_inv=None,
     ):
         super().__init__()
-        self.input_quant_key = quant_key
+        self._input_quant_key = quant_key
         if input_scale is not None:
             self.input_scale = input_scale
         if input_global_scale is not None:
@@ -118,7 +118,7 @@ def test_manual_fusion_fp8_static_with_linear(
         # Enable fusion
         expose_input_quant_key(fp8_linear, fp8_linear.kernel)
 
-        if not hasattr(fp8_linear, "input_quant_key"):
+        if get_input_quant_key(fp8_linear) is None:
             pytest.skip(
                 f"Kernel {force_kernel.__name__} doesn't support input_quant_key"
             )
@@ -315,7 +315,7 @@ def test_manual_fusion_nvfp4_dynamic(
     envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
 )
 def test_manual_fusion_fallback_no_key(dtype: torch.dtype):
-    """Test that maybe_fused_act_quant falls back when no input_quant_key."""
+    """Test that maybe_fused_act_quant falls back without an input key."""
     torch.set_default_device("cuda")
     torch.set_default_dtype(dtype)
 
@@ -327,7 +327,7 @@ def test_manual_fusion_fallback_no_key(dtype: torch.dtype):
 
     with set_current_vllm_config(config):
         silu_and_mul = SiluAndMul()
-        # Linear without input_quant_key attribute
+        # Linear without an input quantization key
         mock_linear = torch.nn.Linear(128, 128)
 
         result = maybe_fused_act_quant(silu_and_mul, x, mock_linear)
diff --git a/vllm/models/deepseek_v4_1/amd/__init__.py b/tests/compile/rocm/__init__.py
similarity index 100%
rename from vllm/models/deepseek_v4_1/amd/__init__.py
rename to tests/compile/rocm/__init__.py
diff --git a/tests/compile/rocm/test_async_tp.py b/tests/compile/rocm/test_async_tp.py
new file mode 100644
index 000000000000..a4462265f8e7
--- /dev/null
+++ b/tests/compile/rocm/test_async_tp.py
@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""ROCm AsyncTP numerical coverage kept outside the shared distributed suite."""
+
+import pytest
+import torch
+
+from tests.compile.passes.distributed.test_async_tp import (
+    TestAGMMModel,
+    TestMMRSModel,
+    async_tp_pass_on_test_model,
+)
+from tests.utils import multi_gpu_test
+from vllm.distributed import (
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_reduce_scatter,
+)
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_file_store_init_method
+
+
+class TestRocmMMRSModel(TestMMRSModel):
+    def forward(self, hidden_states):
+        view = hidden_states.reshape(-1, self.hidden_size)
+        mm = torch.ops.vllm.rocm_unquantized_gemm(view, self.gate_proj, None)
+        return tensor_model_parallel_reduce_scatter(mm, dim=0)
+
+    def ops_in_model_before(self):
+        return super().ops_in_model_before() + [
+            torch.ops.vllm.rocm_unquantized_gemm.default
+        ]
+
+
+class TestRocmAGMMModel(TestAGMMModel):
+    def __init__(self, hidden_size=16, dtype=torch.bfloat16):
+        super().__init__(hidden_size, dtype)
+        self.weight = torch.nn.Parameter(
+            torch.empty((hidden_size * 2, hidden_size)), requires_grad=False
+        )
+        torch.nn.init.normal_(self.weight, std=0.02)
+
+    def forward(self, hidden_states):
+        view = hidden_states.reshape(-1, self.hidden_size)
+        gathered = tensor_model_parallel_all_gather(view, dim=0)
+        return torch.ops.vllm.rocm_unquantized_gemm(gathered, self.weight, None)
+
+    def ops_in_model_before(self):
+        return super().ops_in_model_before() + [
+            torch.ops.vllm.rocm_unquantized_gemm.default
+        ]
+
+
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend contract")
+@pytest.mark.parametrize(
+    "test_model", [TestMMRSModel, TestAGMMModel, TestRocmMMRSModel, TestRocmAGMMModel]
+)
+@pytest.mark.parametrize("dynamic", [False, True])
+def test_rocm_async_tp_bf16_rewrite_and_correctness(test_model, dynamic):
+    """Real symmetric-memory AG/GEMM and GEMM/RS must replace the collectives."""
+    torch.multiprocessing.spawn(
+        async_tp_pass_on_test_model,
+        args=(
+            2,
+            test_model,
+            8,
+            16,
+            128,
+            torch.bfloat16,
+            dynamic,
+            get_file_store_init_method(),
+            True,
+        ),
+        nprocs=2,
+    )
diff --git a/tests/compile/silly_attention.py b/tests/compile/silly_attention.py
index 29c02f6e6a1d..0d3313512912 100644
--- a/tests/compile/silly_attention.py
+++ b/tests/compile/silly_attention.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shared PyTorch custom silly attention for compilation tests.
+"""Shared PyTorch custom silly attention for compilation tests.
 Centralizes custom operation definitions to avoid duplicate registrations.
 """
 
@@ -21,12 +20,12 @@
 
 
 def get_global_counter():
-    """Get the current global counter value"""
+    """Get the current global counter value."""
     return _global_counter
 
 
 def reset_global_counter():
-    """Reset the global counter to 0"""
+    """Reset the global counter to 0."""
     global _global_counter
     _global_counter = 0
 
@@ -34,8 +33,7 @@ def reset_global_counter():
 def silly_attention(
     q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor
 ) -> None:
-    """
-    Unified attention implementation that depends on
+    """Unified attention implementation that depends on
     all inputs and affects the output.
     Always increments a global counter that tests can use or ignore.
     """
@@ -51,7 +49,7 @@ def silly_attention(
 def silly_attention_fake(
     q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor
 ) -> None:
-    """Fake implementation for testing"""
+    """Fake implementation for testing."""
     return
 
 
diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py
index b3a71b4c3017..e5c11b8acc8a 100644
--- a/tests/compile/test_aot_compile.py
+++ b/tests/compile/test_aot_compile.py
@@ -13,6 +13,8 @@
 import torch
 
 import vllm.envs as envs
+import vllm.kernels  # noqa: F401 to register kernels
+from vllm import ir
 from vllm.compilation.backends import VllmBackend
 from vllm.compilation.caching import (
     StandaloneCompiledArtifacts,
@@ -28,6 +30,9 @@
 )
 from vllm.envs import disable_envs_cache
 from vllm.forward_context import set_forward_context
+from vllm.ir import ops
+from vllm.model_executor.layers.activation import GeluAndMulSparse
+from vllm.platforms import current_platform
 from vllm.utils.torch_utils import is_torch_equal_or_newer
 
 from ..utils import create_new_process_for_each_test
@@ -77,6 +82,18 @@ def forward(self, x: torch.Tensor):
         return reference_fn_tuple(x)
 
 
+@support_torch_compile(dynamic_arg_dims={"x": 0})
+class SparseActivationAotModel(torch.nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.up_proj = torch.nn.Linear(32, 256, bias=False)
+        self.activation = GeluAndMulSparse(0.95, "tanh")
+        self.down_proj = torch.nn.Linear(128, 16, bias=False)
+
+    def forward(self, x: torch.Tensor):
+        return self.down_proj(self.activation(self.up_proj(x)))
+
+
 def make_vllm_config() -> VllmConfig:
     return VllmConfig(
         compilation_config=CompilationConfig(
@@ -175,6 +192,54 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
             assert torch.allclose(ret, expected)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or not is_torch_equal_or_newer("2.10.0"),
+    reason="requires CUDA and torch 2.10",
+)
+def test_sparse_activation_standalone_aot_preserves_dependencies(
+    monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path
+):
+    def native_reference(model: SparseActivationAotModel, x: torch.Tensor):
+        projected = torch.nn.functional.linear(x, model.up_proj.weight)
+        activated = ops.gelu_and_mul_sparse.impls["native"].impl_fn(
+            projected, model.activation.std_multiplier, "tanh"
+        )
+        return torch.nn.functional.linear(activated, model.down_proj.weight)
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_AOT_COMPILE", "1")
+        m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
+        m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
+        m.setenv("VLLM_CACHE_ROOT", str(vllm_tmp_cache / "vllm_cache"))
+
+        vllm_config = make_vllm_config()
+        torch.manual_seed(0)
+        inputs = [
+            torch.randn(4, 32, device="cuda", dtype=torch.bfloat16),
+            torch.randn(7, 32, device="cuda", dtype=torch.bfloat16),
+        ]
+
+        with (
+            use_vllm_config(vllm_config),
+            vllm_config.kernel_config.ir_op_priority.set_priority(),
+            ir.enable_torch_wrap(True),
+        ):
+            compiled = SparseActivationAotModel(vllm_config=vllm_config).to(
+                device="cuda", dtype=torch.bfloat16
+            )
+            compiled.requires_grad_(False)
+            expected = [native_reference(compiled, x) for x in inputs]
+            actual = [compiled(x) for x in inputs]
+
+        for result, reference in zip(actual, expected):
+            assert torch.isfinite(result).all()
+            torch.testing.assert_close(
+                result,
+                reference,
+                **ops.gelu_and_mul_sparse.get_tolerance(result.dtype),
+            )
+
+
 @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch):
     from torch._subclasses import FakeTensorMode
@@ -203,8 +268,7 @@ def foo(x: torch.Tensor):
 def test_cache_load_returns_tuple_consistency_tuple_output(
     monkeypatch: pytest.MonkeyPatch,
 ):
-    """
-    Test that cache loading correctly handles models that return tuples.
+    """Test that cache loading correctly handles models that return tuples.
 
     This verifies that when a model returns a tuple of tensors, the output
     type is preserved as a tuple between fresh compilation and cache load.
@@ -270,8 +334,7 @@ def test_cache_load_returns_tuple_consistency_tuple_output(
 
 @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_shape_env(monkeypatch: pytest.MonkeyPatch):
-    """
-    Test that the shape environment is correctly serialized and preserved
+    """Test that the shape environment is correctly serialized and preserved
     when loading from cache.
     """
     with monkeypatch.context() as m:
@@ -309,8 +372,7 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
 def test_partition_wrapper_applied_on_aot_load(
     monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path, mocker
 ):
-    """
-    Test that partition wrappers are applied when loading AOT cached functions.
+    """Test that partition wrappers are applied when loading AOT cached functions.
 
     This test verifies the fix for GitHub issue #31439 where AOT compile
     caused 2x latency regression when use_inductor_graph_partition=True.
@@ -433,14 +495,12 @@ def test_standalone_compile_correctness():
 @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 @create_new_process_for_each_test("spawn")
 def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
-    """
-    Test that compiling gpt2 twice results in a cache hit.
+    """Test that compiling gpt2 twice results in a cache hit.
 
     Counter values are read from the EngineCore subprocess via
     ``LLM.collective_rpc`` so the test works under default V1
     multiprocessing (no shared memory between test and engine).
     """
-
     from vllm import LLM
 
     def _snap(self):
@@ -576,7 +636,7 @@ def test_num_artifacts_and_entries(self):
 
     @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize")
     def test_load_all_success(self, mock_deserialize):
-        """Test successful loading of all artifacts"""
+        """Test successful loading of all artifacts."""
         cache = StandaloneCompiledArtifacts()
 
         mock_artifact1 = Mock()
@@ -593,7 +653,7 @@ def test_load_all_success(self, mock_deserialize):
 
     @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize")
     def test_load_all_already_loaded(self, mock_deserialize):
-        """Test that load_all skips if already loaded"""
+        """Test that load_all skips if already loaded."""
         cache = StandaloneCompiledArtifacts()
 
         mock_artifact = Mock()
@@ -606,7 +666,7 @@ def test_load_all_already_loaded(self, mock_deserialize):
 
     @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize")
     def test_get_loaded_artifact(self, mock_deserialize):
-        """Test retrieving loaded artifacts"""
+        """Test retrieving loaded artifacts."""
         cache = StandaloneCompiledArtifacts()
 
         mock_artifact = Mock()
diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py
index 5c9132ccbfe2..2d4ef13e3386 100644
--- a/tests/compile/test_config.py
+++ b/tests/compile/test_config.py
@@ -70,6 +70,37 @@ def test_get_raw_stream_patch():
         assert get_raw_stream is _cuda_getCurrentRawStream
 
 
+# Inductor can initialize CUDA even for CPU inputs.
+@pytest.mark.forked
+@pytest.mark.parametrize("use_v2", [False, True])
+def test_e8m0_custom_op_fullgraph(use_v2, use_fresh_inductor_cache):
+    """E8M0 inputs must not prevent mutable custom-op decomposition."""
+    with torch.library._scoped_library("test_e8m0", "FRAGMENT") as lib:
+        lib.define("scale_(Tensor(a!) out, Tensor scale) -> ()")
+
+        @torch.library.impl(lib, "scale_", "CPU")
+        def scale_impl(out, scale):
+            out.mul_(scale.float())
+
+        @torch.library.register_fake("test_e8m0::scale_", lib=lib)
+        def scale_fake(out, scale):
+            return None
+
+        def forward(x, scale):
+            out = x.clone()
+            torch.ops.test_e8m0.scale_(out, scale)
+            return out
+
+        x = torch.randn(4, 4)
+        scale = torch.full((4,), 128, dtype=torch.uint8).view(torch.float8_e8m0fnu)
+        compiled = torch.compile(
+            forward,
+            fullgraph=True,
+            options={"enable_auto_functionalized_v2": use_v2},
+        )
+        torch.testing.assert_close(compiled(x, scale), x * 2)
+
+
 def test_copy_pass():
     vllm_config = VllmConfig()
     inductor_pass = FixFunctionalizationPass(vllm_config)
@@ -541,6 +572,25 @@ def _mock_config_for_cudagraph_sizes(
     return config
 
 
+@pytest.mark.parametrize("max_num_seqs", [100, 101])
+def test_default_cudagraph_capture_sizes_cover_off_stride_max_num_seqs(
+    max_num_seqs: int,
+) -> None:
+    compilation_config = CompilationConfig(
+        cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
+    )
+    config = _mock_config_for_cudagraph_sizes(
+        max_num_seqs=max_num_seqs,
+        num_speculative_tokens=0,
+        max_num_batched_tokens=32768,
+        compilation_config=compilation_config,
+    )
+
+    VllmConfig._set_cudagraph_sizes(config)
+
+    assert max_num_seqs in compilation_config.cudagraph_capture_sizes
+
+
 @pytest.mark.parametrize(
     ("max_num_seqs", "num_speculative_tokens", "widest_is_captured"),
     [
diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py
index 71f66fce4045..ea2ccafbe7e0 100644
--- a/tests/compile/test_dynamic_shapes_compilation.py
+++ b/tests/compile/test_dynamic_shapes_compilation.py
@@ -22,7 +22,7 @@
 
 
 def get_test_models():
-    """Get list of models to test based on PyTorch version"""
+    """Get list of models to test based on PyTorch version."""
     models = [
         "Qwen/Qwen3-0.6B",
         "openai-community/gpt2",
@@ -172,7 +172,6 @@ def test_model_specialization_with_evaluate_guards(
     """Test that evaluate_guards correctly detects shape specialization
     violations.
     """
-
     if (
         use_aot_compile
         and dynamic_shapes_type == DynamicShapesType.BACKED
diff --git a/tests/compile/test_graph_partition.py b/tests/compile/test_graph_partition.py
index bb2a6f2aee55..f5aaffe00e4d 100644
--- a/tests/compile/test_graph_partition.py
+++ b/tests/compile/test_graph_partition.py
@@ -24,8 +24,7 @@
 
 
 def test_getitem_moved_to_producer_subgraph():
-    """
-    Test that getitem operations are moved to the same subgraph as their input,
+    """Test that getitem operations are moved to the same subgraph as their input,
     preventing tuple inputs to submodules.
     """
 
@@ -79,8 +78,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 
 def test_no_tuple_inputs_with_multiple_consumers():
-    """
-    Test that when a tuple is consumed by multiple split operations,
+    """Test that when a tuple is consumed by multiple split operations,
     getitem operations are properly moved to avoid tuple inputs.
     """
 
@@ -138,13 +136,10 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 
 def test_consecutive_ops_in_split():
-    """
-    Test that consecutive splitting operations are grouped into the same subgraph
-    """
+    """Test that consecutive splitting operations are grouped into the same subgraph."""
 
     def model_fn(x: torch.Tensor) -> torch.Tensor:
-        """
-        Define a simple model where consecutive operations create opportunities
+        """Define a simple model where consecutive operations create opportunities
         for splitting subgraphs.
         """
         # Apply silly attention followed by consecutive operations
@@ -211,8 +206,7 @@ def _subgraphs_with_empty_nodes(split_items, *, is_splitting_graph):
 
 
 def test_empty_only_partition_stays_separate_after_splitting_predecessor():
-    """
-    Empty-only subgraphs should not be merged when the only predecessor is
+    """Empty-only subgraphs should not be merged when the only predecessor is
     a splitting-op subgraph.
     """
 
@@ -248,8 +242,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 
 def test_empty_only_partition_is_merged():
-    """
-    Empty-only subgraphs should still be merged when a non-splitting predecessor
+    """Empty-only subgraphs should still be merged when a non-splitting predecessor
     exists. The merged empty node must remain outside splitting-op subgraphs.
     """
 
@@ -294,8 +287,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 
 def test_builtin_empty_only_partition_is_merged():
-    """
-    In Dynamo graphs, torch.empty/empty_like may appear as builtin call targets
+    """In Dynamo graphs, torch.empty/empty_like may appear as builtin call targets
     (not aten OpOverload). Ensure empty-only partitions are still merged.
     """
 
@@ -340,8 +332,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_sym_size_whole_shape_boundary():
-    """
-    Test that using x.size() (whole shape) across a split boundary can be
+    """Test that using x.size() (whole shape) across a split boundary can be
     compiled by standalone_compile.
 
     The dynamo graph looks like:
@@ -400,8 +391,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_symint_crosses_split_boundary():
-    """
-    Test that SymInt placeholders from torch.compile + mark_dynamic
+    """Test that SymInt placeholders from torch.compile + mark_dynamic
     cross split boundaries safely via split_module's natural threading.
 
     SymInt values are threaded through subgraphs by split_module and
@@ -459,8 +449,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_shape_boundary_standalone_compile():
-    """
-    Repro for the original production bug:
+    """Repro for the original production bug:
 
         AssertionError: out_spec mismatch
         TreeSpec(tuple, None, [*, *, TreeSpec(Size, None, [*, *]), *])
@@ -529,8 +518,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_size_used_in_multiple_consumer_subgraphs():
-    """
-    Validates that x.size() (whole shape) used by multiple downstream subgraphs
+    """Validates that x.size() (whole shape) used by multiple downstream subgraphs
     does not cause torch.Size to cross split boundaries.
 
     Model:
@@ -584,8 +572,7 @@ def model_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_sym_size_metadata_propagated():
-    """
-    Validates that new sym_size.int nodes created by the pre-pass have
+    """Validates that new sym_size.int nodes created by the pre-pass have
     example_value metadata set. Without it, placeholder metadata in consumer
     subgraphs would be None, breaking any code that dynamically builds
     example inputs from metadata (e.g. standalone_compile per-submodule).
@@ -634,8 +621,7 @@ def model_fn(x: torch.Tensor) -> torch.Tensor:
 
 
 def test_decompose_size_with_getitem_user():
-    """
-    Regression test: _decompose_size_nodes must handle getitem users of size()
+    """Regression test: _decompose_size_nodes must handle getitem users of size()
     correctly.
 
     When a graph contains x.shape[i], it can appear as:
@@ -704,8 +690,7 @@ def test_decompose_size_with_getitem_user():
 
 
 def test_decompose_size_leaves_scalar_size_with_dim():
-    """
-    Regression test: _decompose_size_nodes must leave x.size(dim) alone.
+    """Regression test: _decompose_size_nodes must leave x.size(dim) alone.
 
     x.size() returns a torch.Size tuple that can't cross split boundaries and
     must be decomposed. x.size(dim), however, already returns a scalar
diff --git a/tests/compile/test_structured_logging.py b/tests/compile/test_structured_logging.py
index 10b0ed139cfe..98f092bfce44 100644
--- a/tests/compile/test_structured_logging.py
+++ b/tests/compile/test_structured_logging.py
@@ -60,6 +60,7 @@ def get(self, event_type: str, name_pattern: str) -> list[dict]:
         Args:
             event_type: The event type to filter by (e.g., "artifact", "graph_dump")
             name_pattern: Regex pattern to match against the artifact name
+
         """
         regex = re.compile(name_pattern)
         return [
diff --git a/tests/config/test_aux_output_config.py b/tests/config/test_aux_output_config.py
new file mode 100644
index 000000000000..3d1a06718aee
--- /dev/null
+++ b/tests/config/test_aux_output_config.py
@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from vllm.config import AuxOutputConfig, VllmConfig
+from vllm.config.kv_transfer import KVRole, KVTransferConfig
+from vllm.engine.arg_utils import EngineArgs
+
+pytestmark = pytest.mark.cpu_test
+
+
+def _config(
+    *,
+    use_v2: bool = True,
+    pp: int = 1,
+    dcp: int = 1,
+    pcp: int = 1,
+    connector: str | None = None,
+    kv_role: KVRole = "kv_both",
+    runner_type: str = "generate",
+    is_moe: bool = True,
+    sliding_window: int | None = None,
+    attention_chunk_size: int | None = None,
+    enable_prefix_caching: bool = True,
+    adaptive_verification: bool = False,
+):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            runner_type=runner_type,
+            is_moe=is_moe,
+            get_sliding_window=lambda: sliding_window,
+            attention_chunk_size=attention_chunk_size,
+        ),
+        use_v2_model_runner=use_v2,
+        parallel_config=SimpleNamespace(
+            pipeline_parallel_size=pp,
+            decode_context_parallel_size=dcp,
+            prefill_context_parallel_size=pcp,
+        ),
+        aux_output_config=AuxOutputConfig(enable_return_routed_experts=True),
+        cache_config=SimpleNamespace(
+            enable_prefix_caching=enable_prefix_caching,
+        ),
+        speculative_config=(
+            SimpleNamespace(enable_adaptive_verification=True)
+            if adaptive_verification
+            else None
+        ),
+        kv_transfer_config=(
+            None
+            if connector is None
+            else KVTransferConfig(
+                kv_connector=connector,
+                kv_role=kv_role,
+            )
+        ),
+    )
+
+
+def test_aux_output_config_defaults():
+    config = AuxOutputConfig()
+
+    assert not config.enabled
+    assert not config.enable_return_routed_experts
+    assert config.max_bytes is None
+
+
+def test_aux_output_capture_changes_compilation_hash():
+    disabled = AuxOutputConfig()
+    enabled = AuxOutputConfig(enable_return_routed_experts=True)
+
+    assert disabled.compute_hash() != enabled.compute_hash()
+
+
+def test_legacy_routed_experts_flag_updates_aux_output_config():
+    args = EngineArgs(enable_return_routed_experts=True)
+
+    assert args.aux_output_config.enabled
+    assert args.aux_output_config.enable_return_routed_experts
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "error"),
+    [
+        ({"use_v2": False}, "requires Model Runner V2"),
+        ({"runner_type": "pooling"}, "only supports generate runners"),
+        ({"is_moe": False}, "only supports MoE models"),
+        ({"enable_prefix_caching": False}, "requires prefix caching"),
+        (
+            {"adaptive_verification": True},
+            "adaptive speculative verification",
+        ),
+        ({"pp": 2}, "pipeline parallelism"),
+        ({"dcp": 2}, "context parallelism"),
+        ({"pcp": 2}, "context parallelism"),
+    ],
+)
+def test_aux_output_connector_rejects_unsupported_configuration(kwargs, error):
+    with pytest.raises(ValueError, match=error):
+        VllmConfig._verify_aux_output_compatibility(_config(**kwargs))
+
+
+@pytest.mark.parametrize(
+    ("connector", "blocked"),
+    [
+        ("NixlConnector", True),
+        ("NixlPullConnector", True),
+        ("NixlPushConnector", True),
+        ("MoRIIOConnector", True),
+        ("MooncakeConnector", True),
+        ("MooncakeStoreConnector", False),
+        ("OffloadingConnector", False),
+        ("LMCacheConnectorV1", False),
+        ("LMCacheMPConnector", False),
+        ("SimpleCPUOffloadConnector", False),
+    ],
+)
+@pytest.mark.parametrize("kv_role", ["kv_both", "kv_producer", "kv_consumer"])
+@pytest.mark.parametrize("multi", [False, True])
+def test_aux_output_connector_policy_is_independent_of_role(
+    connector, blocked, kv_role, multi
+):
+    config = _config(connector=connector, kv_role=kv_role)
+    if multi:
+        config.kv_transfer_config = KVTransferConfig(
+            kv_connector="MultiConnector",
+            kv_role=kv_role,
+            kv_connector_extra_config={
+                "connectors": [
+                    {"kv_connector": "OffloadingConnector", "kv_role": "kv_both"},
+                    {"kv_connector": connector, "kv_role": kv_role},
+                ]
+            },
+        )
+    if blocked:
+        with pytest.raises(ValueError, match=f"incompatible with {connector}"):
+            VllmConfig._verify_aux_output_compatibility(config)
+    else:
+        VllmConfig._verify_aux_output_compatibility(config)
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [{"sliding_window": 4096}, {"attention_chunk_size": 4096}],
+)
+def test_aux_output_config_defers_attention_layout_to_kv_config(kwargs):
+    VllmConfig._verify_aux_output_compatibility(_config(**kwargs))
+
+
+def test_aux_output_guards_are_inactive_when_capture_is_disabled():
+    config: Any = VllmConfig.__new__(VllmConfig)
+    config.model_config = SimpleNamespace()
+    config.aux_output_config = AuxOutputConfig()
+
+    config._verify_aux_output_compatibility()
diff --git a/tests/config/test_bailing_mtp_config.py b/tests/config/test_bailing_mtp_config.py
index bb7cf3ad5631..cd842856a9ba 100644
--- a/tests/config/test_bailing_mtp_config.py
+++ b/tests/config/test_bailing_mtp_config.py
@@ -3,7 +3,7 @@
 
 from typing import get_args
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.speculative import MTPModelTypes, SpeculativeConfig
 from vllm.transformers_utils.model_arch_config_convertor import (
@@ -11,8 +11,8 @@
 )
 
 
-def _bailing_config() -> PretrainedConfig:
-    config = PretrainedConfig(
+def _bailing_config() -> PreTrainedConfig:
+    config = PreTrainedConfig(
         architectures=["BailingMoeV2_5ForCausalLM"],
         hidden_size=4096,
         kv_lora_rank=512,
diff --git a/tests/config/test_config_generation.py b/tests/config/test_config_generation.py
index c7edf2b97174..a6908f35e518 100644
--- a/tests/config/test_config_generation.py
+++ b/tests/config/test_config_generation.py
@@ -109,3 +109,60 @@ def test_unrecognized_env(monkeypatch):
         fail_on_environ_validation=True,
     )
     engine_args.create_engine_config()
+
+
+def test_get_diff_sampling_param_includes_penalties():
+    """``--override-generation-config`` must apply presence_penalty and
+    frequency_penalty as server-side defaults, like the other sampling
+    params. Regression test for the whitelist silently dropping them.
+
+    Called unbound with a lightweight stand-in so no model is loaded:
+    generation_config="vllm" makes the method skip try_get_generation_config
+    and read only override_generation_config.
+    """
+    from types import SimpleNamespace
+
+    from vllm.config.model import ModelConfig
+
+    fake = SimpleNamespace(
+        generation_config="vllm",
+        override_generation_config={
+            "presence_penalty": 1.5,
+            "frequency_penalty": 0.5,
+            "temperature": 0.6,
+        },
+    )
+
+    diff = ModelConfig.get_diff_sampling_param(fake)  # type: ignore[arg-type]
+
+    assert diff["presence_penalty"] == 1.5
+    assert diff["frequency_penalty"] == 0.5
+    assert diff["temperature"] == 0.6
+
+
+def test_generation_config_penalties_survive_diff_sampling_filter():
+    """``generation_config.json`` penalties must survive the diff-sampling
+    whitelist (regression for #50767, auto path).
+
+    Covers the ``generation_config="auto"`` branch, where the override
+    dict is empty and the values come from ``try_get_generation_config``.
+    """
+    from types import SimpleNamespace
+
+    from vllm.config.model import ModelConfig
+
+    fake = SimpleNamespace(
+        generation_config="auto",
+        override_generation_config={},
+        try_get_generation_config=lambda: {
+            "presence_penalty": 1.1,
+            "frequency_penalty": 0.7,
+            "temperature": 0.9,
+        },
+    )
+
+    diff = ModelConfig.get_diff_sampling_param(fake)  # type: ignore[arg-type]
+
+    assert diff["presence_penalty"] == 1.1
+    assert diff["frequency_penalty"] == 0.7
+    assert diff["temperature"] == 0.9
diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py
index 94060d1724a8..1d0f4d915bcf 100644
--- a/tests/config/test_config_utils.py
+++ b/tests/config/test_config_utils.py
@@ -40,7 +40,7 @@ class DummyLogprobsMode(Enum):
 
 
 def test_hash_factors_deterministic():
-    """Test that hash_factors produces consistent SHA-256 hashes"""
+    """Test that hash_factors produces consistent SHA-256 hashes."""
     factors = {"a": 1, "b": "test"}
     hash1 = hash_factors(factors)
     hash2 = hash_factors(factors)
@@ -215,6 +215,9 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs():
     base_hash = CacheConfig().compute_hash()
     assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash
     assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash
+    config = CacheConfig()
+    config.effective_attention_block_size = 64
+    assert config.compute_hash() == base_hash
 
 
 def test_scheduler_config_hash_includes_max_num_seqs():
diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py
index b68820c0d51c..ae4c7220f24c 100644
--- a/tests/config/test_model_arch_config.py
+++ b/tests/config/test_model_arch_config.py
@@ -10,13 +10,15 @@
 from unittest.mock import Mock
 
 import pytest
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
 
 from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
 from vllm.config.model_arch import ModelArchitectureConfig
+from vllm.transformers_utils.configs.diffusion_gemma import DiffusionGemmaConfig
 from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
 from vllm.transformers_utils.model_arch_config_convertor import (
+    MODEL_ARCH_CONFIG_CONVERTORS,
     Gemma4ModelArchConfigConvertor,
     ModelArchConfigConvertorBase,
 )
@@ -123,7 +125,7 @@ def _assert_model_config_methods(
 
 def test_head_size_falls_back_when_head_dim_is_zero():
     """Regression test for configs that materialize missing head_dim as 0."""
-    hf_config = PretrainedConfig(
+    hf_config = PreTrainedConfig(
         model_type="deepseek_vl_v2",
         hidden_size=1280,
         num_attention_heads=10,
@@ -138,7 +140,7 @@ def test_head_size_falls_back_when_head_dim_is_zero():
 
 
 def test_qk_rope_head_dim_recovery_uses_model_revision(monkeypatch: pytest.MonkeyPatch):
-    hf_config = PretrainedConfig(
+    hf_config = PreTrainedConfig(
         model_type="deepseek_v2",
         qk_rope_head_dim=128,
         qk_nope_head_dim=128,
@@ -174,7 +176,7 @@ def test_qk_rope_head_dim_recovery_uses_model_revision(monkeypatch: pytest.Monke
     ],
 )
 def test_num_experts_per_tok_aliases(attribute: str):
-    hf_config = PretrainedConfig(**{attribute: 4})
+    hf_config = PreTrainedConfig(**{attribute: 4})
     convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
     model_config = cast(
         ModelConfig,
@@ -189,7 +191,7 @@ def test_num_experts_per_tok_aliases(attribute: str):
 
 
 def test_num_experts_per_tok_none_is_normalized():
-    hf_config = PretrainedConfig(top_k_experts=None)
+    hf_config = PreTrainedConfig(top_k_experts=None)
     convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
 
     assert convertor.get_num_experts_per_token() == 0
@@ -205,7 +207,7 @@ def test_legacy_modelopt_config_without_producer_is_normalized():
             "modelopt_quant_config": {"quant_cfg": {}},
         }
     }
-    hf_config = PretrainedConfig(quantization_config=quantization_config)
+    hf_config = PreTrainedConfig(quantization_config=quantization_config)
 
     convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
 
@@ -368,7 +370,7 @@ def test_gemma4_uniform_head_dims_are_homogeneous():
     assert arch[3] is arch
 
 
-class _HeterogeneousConfig(PretrainedConfig):
+class _HeterogeneousConfig(PreTrainedConfig):
     """A heterogeneous config with no convertor of its own.
 
     Mirrors the parts vLLM uses: per-layer configs are shallow copies with the
@@ -382,7 +384,7 @@ def __init__(self, per_layer: dict[str, list], **kwargs):
         super().__init__(**kwargs)
 
     @property
-    def per_layer_config(self) -> list[PretrainedConfig]:
+    def per_layer_config(self) -> list[PreTrainedConfig]:
         layers = []
         for i in range(self.num_hidden_layers):
             layer = copy(self)
@@ -496,3 +498,26 @@ def test_deepseek_v4_convertor_splits_vision_architecture():
         conv = DeepseekV4ModelArchConfigConvertor(draft_cfg, draft_cfg)
         assert conv.get_architectures() == [draft_arch]
         assert draft_cfg.architectures == [draft_arch]
+
+
+def test_diffusion_gemma_head_dims_vary_by_layer_type():
+    """Checkpoints carry model_type `diffusion_gemma` at the top level. Under
+    the base convertor every layer reports the sliding head dim, Gemma4Config
+    skips its head-dim handling, and the full-attention KV-cache group falls to
+    FlashInfer, which cannot take a per-request causal tensor."""
+    text_config = dict(
+        num_hidden_layers=6,
+        hidden_size=64,
+        num_attention_heads=8,
+        num_key_value_heads=4,
+        head_dim=16,
+        global_head_dim=32,
+        layer_types=["sliding_attention"] * 5 + ["full_attention"],
+    )
+    hf_config = DiffusionGemmaConfig(text_config=text_config)
+
+    convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS[hf_config.model_type]
+    assert convertor_cls is Gemma4ModelArchConfigConvertor
+    arch = convertor_cls(hf_config, hf_config.text_config).convert()
+
+    assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32]
diff --git a/tests/config/test_mp_reducer.py b/tests/config/test_mp_reducer.py
index 56dc542f1c76..56e4e474740e 100644
--- a/tests/config/test_mp_reducer.py
+++ b/tests/config/test_mp_reducer.py
@@ -9,12 +9,10 @@
 
 
 def test_mp_reducer():
-    """
-    Test that _reduce_config reducer is registered when AsyncLLM is instantiated
+    """Test that _reduce_config reducer is registered when AsyncLLM is instantiated
     without transformers_modules. This is a regression test for
     https://github.com/vllm-project/vllm/pull/18640.
     """
-
     # Ensure transformers_modules is not in sys.modules
     if "transformers_modules" in sys.modules:
         del sys.modules["transformers_modules"]
diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py
index 589eabf91d67..b36f5776da25 100644
--- a/tests/config/test_multimodal_config.py
+++ b/tests/config/test_multimodal_config.py
@@ -1,12 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import copy
 from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
+from tests.models.utils import build_model_context
 from vllm.config.ec_transfer import ECRole, ECTransferConfig
 from vllm.config.model import ModelConfig
 from vllm.config.multimodal import MultiModalConfig
@@ -17,6 +19,26 @@
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 
+@pytest.mark.parametrize(
+    "model_id,limit_mm_per_prompt,expected",
+    [
+        ("Qwen/Qwen2-0.5B-Instruct", {}, False),
+        ("Qwen/Qwen2.5-VL-3B-Instruct", {}, True),
+        ("Qwen/Qwen2.5-VL-3B-Instruct", {"image": 0, "video": 0}, False),
+        ("Qwen/Qwen2.5-VL-3B-Instruct", {"image": 0}, True),
+    ],
+)
+@pytest.mark.core_model
+def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected):
+    """Test supports_multimodal_inputs returns correct boolean for various
+    configs."""
+    ctx = build_model_context(
+        model_id,
+        limit_mm_per_prompt=limit_mm_per_prompt,
+    )
+    assert ctx.model_config.supports_multimodal_inputs is expected
+
+
 def test_mm_encoder_attn_backend_str_conversion():
     config = MultiModalConfig(mm_encoder_attn_backend="FLASH_ATTN")  # type: ignore[arg-type]
     assert config.mm_encoder_attn_backend == AttentionBackendEnum.FLASH_ATTN
@@ -84,63 +106,37 @@ def test_mm_encoder_attn_dtype_hash_updates(tmp_path):
     assert fp8_hash != fp8_static_hash
 
 
+_MULTIMODAL_MODEL = "llava-hf/llava-1.5-7b-hf"
+_TEXT_ONLY_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+
+
 def _make_mm_prefix_model_config(
+    model: str = _MULTIMODAL_MODEL,
     *,
     language_model_only: bool = False,
 ) -> ModelConfig:
-    model_config = MagicMock(spec=ModelConfig)
-    model_config.multimodal_config = MultiModalConfig(
-        language_model_only=language_model_only
-    )
-    # Bind real helper methods onto the mock.
-    model_config._supports_multimodal_for_mm_prefix = (
-        ModelConfig._supports_multimodal_for_mm_prefix.__get__(
-            model_config, ModelConfig
-        )
-    )
-    return model_config
-
-
-@pytest.mark.parametrize("supports_mm", [True, False])
-def test_supports_multimodal_for_mm_prefix_uses_registry(supports_mm: bool):
-    model_config = _make_mm_prefix_model_config()
-
-    with patch(
-        "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
-        return_value=supports_mm,
-    ) as mocked:
-        assert model_config._supports_multimodal_for_mm_prefix() is supports_mm
-        mocked.assert_called_once_with(model_config)
-
-    # Sticky cache — registry must not be consulted again.
-    with patch(
-        "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
-        side_effect=AssertionError("should use cache"),
-    ):
-        assert model_config._supports_multimodal_for_mm_prefix() is supports_mm
+    return ModelConfig(model, language_model_only=language_model_only)
 
 
 def test_supports_multimodal_for_mm_prefix_before_multimodal_config():
-    model_config = _make_mm_prefix_model_config()
-    model_config.multimodal_config = None
+    """A text-only model never builds a multimodal config, so the early
+    return must neither clear mm_prefix nor write the sticky cache."""
+    model_config = _make_mm_prefix_model_config(_TEXT_ONLY_MODEL)
+    assert model_config.multimodal_config is None
 
     assert model_config._supports_multimodal_for_mm_prefix() is True
-    assert not hasattr(model_config, "_supports_multimodal_inputs_cached")
+    assert not getattr(model_config, "_supports_multimodal_inputs_cache", None)
 
 
 def test_language_model_only_disables_via_supports_multimodal_inputs():
     """language_model_only zeros all limits, so registry reports text-only."""
     model_config = _make_mm_prefix_model_config(language_model_only=True)
 
-    with patch(
-        "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
-        return_value=False,
-    ):
-        assert model_config._supports_multimodal_for_mm_prefix() is False
+    assert model_config._supports_multimodal_for_mm_prefix() is False
 
 
 def test_convertor_clears_mm_prefix_when_multimodal_disabled():
-    hf_config = PretrainedConfig(
+    hf_config = PreTrainedConfig(
         model_type="gemma3",
         architectures=["Gemma3ForConditionalGeneration"],
     )
@@ -159,23 +155,19 @@ def test_convertor_clears_mm_prefix_when_multimodal_disabled():
 def test_sticky_cache_survives_text_subconfig_regeneration():
     """with_hf_config deepcopies the cached decision onto text submodules."""
     model_config = _make_mm_prefix_model_config()
-    with patch(
-        "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
-        return_value=False,
-    ):
-        assert model_config._supports_multimodal_for_mm_prefix() is False
+    assert model_config._supports_multimodal_for_mm_prefix() is True
 
-    # Simulate deepcopy onto a Gemma4ForCausalLM-like config that would
-    # otherwise fail registry lookup / return False incorrectly.
-    text_config = _make_mm_prefix_model_config()
-    text_config._supports_multimodal_inputs_cached = (
-        model_config._supports_multimodal_inputs_cached
-    )
-    with patch(
-        "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
-        side_effect=AssertionError("must not re-query registry"),
-    ):
-        assert text_config._supports_multimodal_for_mm_prefix() is False
+    # `with_hf_config` deep-copies this config and swaps `hf_config` for a
+    # text-only submodule (e.g. Gemma4ForCausalLM).
+    text_config = copy.deepcopy(model_config)
+    text_config.hf_config = model_config.hf_text_config
+    assert text_config._supports_multimodal_for_mm_prefix() is True
+
+    # Without the copied cache the submodule architecture has no registered
+    # multimodal processor, so re-querying the registry wrongly settles on
+    # text-only and would clear mm_prefix.
+    del text_config._supports_multimodal_inputs_cache
+    assert text_config._supports_multimodal_for_mm_prefix() is False
 
 
 @pytest.mark.parametrize(
@@ -375,3 +367,85 @@ def test_vllm_config_runs_the_mm_processor_device_check():
         pytest.raises(ValueError, match="also runs the language model"),
     ):
         VllmConfig._validate_mm_processor_device(vllm_config)
+
+
+def _resolve_mm_video_decode_device(
+    *,
+    ec_role: ECRole | None,
+    mm_tensor_ipc: str = "torch_shm",
+    device: str | None = None,
+    video_kwargs: dict | None = None,
+    torchcodec_available: bool = True,
+) -> dict:
+    """Run processor-device then video-decode-device resolution and report
+    the resulting video media IO kwargs."""
+    mm_config = MultiModalConfig(
+        mm_processor_kwargs={} if device is None else {"device": device},
+        mm_tensor_ipc=mm_tensor_ipc,  # type: ignore[arg-type]
+        media_io_kwargs={} if video_kwargs is None else {"video": dict(video_kwargs)},
+    )
+    model_config = MagicMock(spec=ModelConfig)
+    model_config.multimodal_config = mm_config
+    vllm_config = MagicMock(spec=VllmConfig)
+    vllm_config.model_config = model_config
+    vllm_config.ec_transfer_config = (
+        None
+        if ec_role is None
+        else ECTransferConfig(ec_connector="ECExampleConnector", ec_role=ec_role)
+    )
+
+    with (
+        patch("vllm.platforms.current_platform.device_type", "cuda"),
+        patch(
+            "vllm.utils.import_utils.check_torchcodec_available",
+            side_effect=None if torchcodec_available else ImportError("torchcodec"),
+        ),
+    ):
+        VllmConfig._resolve_mm_processor_device(vllm_config)
+        VllmConfig._resolve_mm_video_decode_device(vllm_config)
+    return mm_config.media_io_kwargs.get("video", {})
+
+
+def test_auto_video_decode_uses_nvdec_on_encoder_instance():
+    """The processor runs on the accelerator there, so decoded frames should
+    stay on-device too."""
+    assert _resolve_mm_video_decode_device(ec_role="ec_producer") == {
+        "backend": "torchcodec",
+        "device": "cuda",
+    }
+
+
+@pytest.mark.parametrize("ec_role", [None, "ec_consumer", "ec_both"])
+def test_auto_video_decode_stays_on_cpu_off_encoder_instance(
+    ec_role: ECRole | None,
+):
+    """The processor stays on CPU off encode-only instances, so video
+    decoding should too."""
+    assert _resolve_mm_video_decode_device(ec_role=ec_role) == {}
+
+
+def test_auto_video_decode_follows_explicit_processor_device():
+    assert _resolve_mm_video_decode_device(ec_role="ec_producer", device="cpu") == {}
+    assert _resolve_mm_video_decode_device(ec_role="ec_producer", device="cuda") == {
+        "backend": "torchcodec",
+        "device": "cuda",
+    }
+
+
+def test_auto_video_decode_respects_explicit_media_io_kwargs():
+    assert _resolve_mm_video_decode_device(
+        ec_role="ec_producer", video_kwargs={"backend": "opencv"}
+    ) == {"backend": "opencv"}
+    assert _resolve_mm_video_decode_device(
+        ec_role="ec_producer", video_kwargs={"device": "cpu"}
+    ) == {"device": "cpu"}
+
+
+def test_auto_video_decode_skipped_without_torchcodec():
+    assert (
+        _resolve_mm_video_decode_device(
+            ec_role="ec_producer",
+            torchcodec_available=False,
+        )
+        == {}
+    )
diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py
index 8505172da387..e7014faa47d0 100644
--- a/tests/config/test_speculative_draft_hf_overrides.py
+++ b/tests/config/test_speculative_draft_hf_overrides.py
@@ -15,20 +15,20 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.parallel import ParallelConfig
 from vllm.config.speculative import SpeculativeConfig
 
 
-def _make_hf_config(**kwargs) -> PretrainedConfig:
+def _make_hf_config(**kwargs) -> PreTrainedConfig:
     defaults = dict(
         architectures=["LlamaForCausalLM"],
         model_type="llama",
         num_hidden_layers=64,
     )
     defaults.update(kwargs)
-    return PretrainedConfig(**defaults)
+    return PreTrainedConfig(**defaults)
 
 
 @pytest.mark.cpu_test
@@ -47,12 +47,97 @@ def test_none_overrides_fall_back_to_arch_mapping():
     assert composed is SpeculativeConfig.hf_config_override
 
 
+def _make_speculative_config(
+    hf_kwargs: dict, architecture: str, **speculative_kwargs
+) -> SpeculativeConfig:
+    hf_config = SpeculativeConfig.hf_config_override(_make_hf_config(**hf_kwargs))
+    draft = MagicMock(
+        hf_config=hf_config, architectures=hf_config.architectures, max_model_len=128
+    )
+    draft.registry.inspect_model_cls.return_value = (None, architecture)
+    target = MagicMock(max_model_len=128, quantization=None, hf_overrides={})
+    target.hf_config.model_type = hf_kwargs["model_type"]
+
+    with patch("vllm.config.speculative.ModelConfig", return_value=draft):
+        return SpeculativeConfig(
+            model="draft",
+            **speculative_kwargs,
+            target_model_config=target,
+            target_parallel_config=ParallelConfig(),
+        )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    ("hf_kwargs", "n_predict", "architecture"),
+    [
+        (
+            dict(
+                model_type="deepseek_v4",
+                num_nextn_predict_layers=3,
+                dspark_block_size=5,
+            ),
+            3,
+            "DSparkDraftModel",
+        ),
+        (
+            dict(
+                model_type="deepseek_v41",
+                num_nextn_predict_layers=3,
+                dspark_block_size=5,
+            ),
+            3,
+            "DSparkV41DraftModel",
+        ),
+        (
+            dict(
+                model_type="gemma4_text",
+                architectures=["Gemma4DSparkModel"],
+                block_size=5,
+            ),
+            None,
+            "Gemma4DSparkModel",
+        ),
+    ],
+)
+@pytest.mark.parametrize("speculative_kwargs", [{"num_speculative_tokens": 5}, {}])
+def test_dspark_width_is_independent_of_mtp_stages(
+    hf_kwargs, n_predict, architecture, speculative_kwargs
+):
+    config = _make_speculative_config(
+        hf_kwargs, architecture, method="dspark", **speculative_kwargs
+    )
+
+    assert config.num_speculative_tokens == 5
+    assert getattr(config.draft_model_config.hf_config, "n_predict", None) == n_predict
+    assert config.draft_model_config.hf_config.architectures == [architecture]
+
+
+@pytest.mark.cpu_test
+def test_mtp_stages_are_independent_of_dspark_width():
+    hf_kwargs = dict(
+        model_type="deepseek_v4", num_nextn_predict_layers=3, dspark_block_size=5
+    )
+    config = _make_speculative_config(
+        hf_kwargs, "DeepSeekV4MTPModel", method="mtp", num_speculative_tokens=6
+    )
+
+    assert config.num_speculative_tokens == 6
+    assert config.draft_model_config.hf_config.n_predict == 3
+    assert config.draft_model_config.hf_config.architectures == ["DeepSeekV4MTPModel"]
+
+    with pytest.raises(ValueError, match="must be divisible by n_predict=3"):
+        _make_speculative_config(
+            hf_kwargs, "DeepSeekV4MTPModel", method="mtp", num_speculative_tokens=5
+        )
+
+
 @pytest.mark.cpu_test
 def test_callable_overrides_reach_the_draft_config():
     """A callable override (config-to-config transform) composes with the
     architecture-mapping override and is applied to the draft config."""
 
-    def shrink(hf_config: PretrainedConfig) -> PretrainedConfig:
+    def shrink(hf_config: PreTrainedConfig) -> PreTrainedConfig:
         hf_config.num_hidden_layers = 1
         return hf_config
 
@@ -70,7 +155,7 @@ def test_arch_mapping_applies_before_callable_override():
     observes (and may adjust) the post-mapping config."""
     seen_architectures: list[str] = []
 
-    def record(hf_config: PretrainedConfig) -> PretrainedConfig:
+    def record(hf_config: PreTrainedConfig) -> PreTrainedConfig:
         seen_architectures.append(hf_config.architectures[0])
         return hf_config
 
@@ -117,7 +202,7 @@ def test_inkling_override_exposes_all_mtp_depths():
     assert out.local_layer_ids == [0, 2, 4]
 
 
-def _module_level_shrink(hf_config: PretrainedConfig) -> PretrainedConfig:
+def _module_level_shrink(hf_config: PreTrainedConfig) -> PreTrainedConfig:
     hf_config.num_hidden_layers = 1
     return hf_config
 
diff --git a/tests/config/test_speculative_draft_max_position_embeddings.py b/tests/config/test_speculative_draft_max_position_embeddings.py
index 117ae5a66d53..f5c5524acded 100644
--- a/tests/config/test_speculative_draft_max_position_embeddings.py
+++ b/tests/config/test_speculative_draft_max_position_embeddings.py
@@ -17,7 +17,7 @@
 from typing import Literal
 
 import pytest
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.model import ModelConfig
 from vllm.config.parallel import ParallelConfig
@@ -46,7 +46,7 @@ def _override_logged(caplog: pytest.LogCaptureFixture) -> bool:
 
 @pytest.mark.cpu_test
 def test_override_raises_smaller_value(vllm_caplog: pytest.LogCaptureFixture):
-    hf_config = PretrainedConfig(max_position_embeddings=2048)
+    hf_config = PreTrainedConfig(max_position_embeddings=2048)
     SpeculativeConfig._maybe_override_draft_max_position_embeddings(
         hf_config, target_max_model_len=8192
     )
@@ -56,7 +56,7 @@ def test_override_raises_smaller_value(vllm_caplog: pytest.LogCaptureFixture):
 
 @pytest.mark.cpu_test
 def test_override_keeps_sufficient_value(vllm_caplog: pytest.LogCaptureFixture):
-    hf_config = PretrainedConfig(max_position_embeddings=8192)
+    hf_config = PreTrainedConfig(max_position_embeddings=8192)
     SpeculativeConfig._maybe_override_draft_max_position_embeddings(
         hf_config, target_max_model_len=8192
     )
@@ -66,7 +66,7 @@ def test_override_keeps_sufficient_value(vllm_caplog: pytest.LogCaptureFixture):
 
 @pytest.mark.cpu_test
 def test_override_ignores_missing_attribute(vllm_caplog: pytest.LogCaptureFixture):
-    hf_config = PretrainedConfig()
+    hf_config = PreTrainedConfig()
     hf_config.__dict__.pop("max_position_embeddings", None)
     SpeculativeConfig._maybe_override_draft_max_position_embeddings(
         hf_config, target_max_model_len=8192
diff --git a/tests/conftest.py b/tests/conftest.py
index 5fd594a04a2e..45c629b8755c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -241,7 +241,7 @@ def dist_init():
                 rank=0,
                 distributed_init_method=f"file://{temp_file}",
                 local_rank=0,
-                backend="nccl",
+                backend="gloo" if current_platform.is_cpu() else "nccl",
             )
             initialize_model_parallel(1, 1)
             yield
@@ -504,6 +504,36 @@ def _init(
                     if model_cls is not None:
                         _fix_v4_tied_weights_keys(model_cls)
 
+            from transformers.integrations import is_deepspeed_zero3_enabled
+
+            # On ROCm, avoid a converted CPU copy when downcasting FP32 references.
+            # Same-dtype CPU loads can stay lazily mapped and use less host RAM.
+            # Leave customized, remote, and quantized loading unchanged.
+            if (
+                current_platform.is_rocm()
+                and self.device == "cuda"
+                and auto_cls is AutoModelForCausalLM
+                and revision is None
+                and model_kwargs.keys() == {"dtype"}
+                and self.config.dtype == torch.float32
+                and model_kwargs["dtype"] in (torch.float16, torch.bfloat16)
+                and not (trust_remote_code and hasattr(self.config, "auto_map"))
+                and getattr(self.config, "quantization_config", None) is None
+                and getattr(
+                    self.config.get_text_config(decoder=True),
+                    "quantization_config",
+                    None,
+                )
+                is None
+                and not is_deepspeed_zero3_enabled()
+            ):
+                model_kwargs = {
+                    **model_kwargs,
+                    "device_map": torch.device(
+                        "cuda", torch.accelerator.current_device_index()
+                    ),
+                }
+
             model = cast(
                 nn.Module,
                 auto_cls.from_pretrained(
@@ -901,13 +931,10 @@ def __enter__(self):
             # tests/basic_correctness/test_basic_correctness.py::test_models_distributed
             # where vllm worker processes are still alive and holding GPU
             # memory when hf_runner.__exit__ is called.
-            from tests.utils import (
-                get_physical_device_indices,
-                record_gpu_memory_usage_stats,
-            )
+            from tests.utils import record_gpu_memory_usage_stats
 
             if (device_count := current_platform.device_count()) > 0:
-                devices = get_physical_device_indices(devices=list(range(device_count)))
+                devices = list(range(device_count))
                 mem_usage_stats = record_gpu_memory_usage_stats(devices=devices)
                 self.threshold_ratios = {
                     device: 0.05 + mem_used / mem_tot
@@ -935,7 +962,7 @@ def hf_runner():
 
 
 def _default_block_size() -> int:
-    if torch.xpu.is_available():
+    if current_platform.is_xpu():
         return 64
     if current_platform.is_cpu():
         return 128
@@ -955,6 +982,8 @@ class VllmRunner:
     - `enable_chunked_prefill`: Set to `False` instead of `None` for
       test reproducibility.
     - `enforce_eager`: Set to `False` to test CUDA graph.
+    - `kernel_config.enable_jit_warmup`: Set to `False` to reduce test startup
+      time.
     """
 
     def __init__(
@@ -983,6 +1012,8 @@ def __init__(
             else set_default_torch_num_threads(default_torch_num_threads)
         )
 
+        kwargs.setdefault("kernel_config", {"enable_jit_warmup": False})
+
         if not kwargs.get("compilation_config", None):
             # Note(@tdoublep): This is set to 4 because some tests (e.g., hybrid
             # model tests) may set max_num_seqs=4. If min cudagraph_capture_size is
diff --git a/tests/cuda/test_cuda_context.py b/tests/cuda/test_cuda_context.py
index aca490d15644..cb2ba1b19922 100644
--- a/tests/cuda/test_cuda_context.py
+++ b/tests/cuda/test_cuda_context.py
@@ -11,7 +11,7 @@
 
 
 def check_cuda_context():
-    """Check CUDA driver context status"""
+    """Check CUDA driver context status."""
     try:
         cuda = ctypes.CDLL("libcuda.so")
         device = ctypes.c_int()
@@ -22,7 +22,7 @@ def check_cuda_context():
 
 
 def run_cuda_test_in_thread(device_input, expected_device_id):
-    """Run CUDA context test in separate thread for isolation"""
+    """Run CUDA context test in separate thread for isolation."""
     try:
         # New thread should have no CUDA context initially
         valid_before, device_before = check_cuda_context()
diff --git a/tests/cuda/test_platform_no_cuda_init.py b/tests/cuda/test_platform_no_cuda_init.py
index 697d2cfec8e1..6f50ce606b76 100644
--- a/tests/cuda/test_platform_no_cuda_init.py
+++ b/tests/cuda/test_platform_no_cuda_init.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test that platform imports do not prematurely initialize CUDA.
+"""Test that platform imports do not prematurely initialize CUDA.
 
 This is critical for Ray-based multi-GPU setups where workers need to
 set CUDA_VISIBLE_DEVICES after importing vLLM but before CUDA is initialized.
diff --git a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py
index 44215cb72ae1..3b72d612dd1c 100644
--- a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py
+++ b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py
@@ -44,8 +44,7 @@ def _make_request(stop, include_stop_str_in_output: bool, min_tokens: int = 0):
 
 
 def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: bool):
-    """
-    This test verifies that the detokenizer correctly handles the case where
+    """This test verifies that the detokenizer correctly handles the case where
     the generated token sequence contains both:
     - a stop token
     - an <eos> token
@@ -67,7 +66,6 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo
     where stop-string evaluation occurs before the early-return on
     stop_terminated.
     """
-
     # Generate text "abcdeZ" and tokenize it.
     generated_text = "abcde"
     eos_token = "Z"
diff --git a/tests/distributed/aux_output_connector/test_store.py b/tests/distributed/aux_output_connector/test_store.py
new file mode 100644
index 000000000000..39277f3e7e0e
--- /dev/null
+++ b/tests/distributed/aux_output_connector/test_store.py
@@ -0,0 +1,1854 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+from contextlib import nullcontext
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.distributed.aux_output_connector.connector import (
+    AuxOutputConnectorMetadata,
+    AuxOutputSchedulerConnector,
+    AuxRequestOutput,
+    PackedBlockHashes,
+)
+from vllm.distributed.aux_output_connector.routed_experts import (
+    RoutedExpertsBuffer,
+    materialize_routed_experts,
+    publish_routed_experts,
+    routed_experts_keys,
+)
+from vllm.distributed.aux_output_connector.store import (
+    BackgroundBlockObjectStore,
+    BlockObject,
+    BlockObjectStore,
+    BlockObjectStoreError,
+)
+from vllm.distributed.aux_output_connector.worker import (
+    AuxOutputWorkerConnector,
+)
+from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.worker.gpu import async_utils
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
+
+pytestmark = pytest.mark.cpu_test
+
+_SHAPE = (3, 2)
+_DTYPE = np.dtype("uint8")
+_BLOCK_SIZE = 4
+
+
+def test_background_store_publishes_without_blocking_caller():
+    started = threading.Event()
+    release = threading.Event()
+    underlying = Mock()
+
+    def put(*_args, **_kwargs):
+        started.set()
+        release.wait()
+
+    underlying.put.side_effect = put
+    store = BackgroundBlockObjectStore(underlying, max_pending_batches=1)
+
+    store.put([BlockObject("key", b"value")])
+    assert started.wait(timeout=1)
+    release.set()
+    store.close()
+
+    underlying.close.assert_called_once_with()
+
+
+def _make_connector():
+    return AuxOutputSchedulerConnector()
+
+
+@dataclass
+class _Execution:
+    request_id: str
+    token_start: int
+    num_tokens: int
+    emit_start: int
+    block_hashes: list[bytes] | PackedBlockHashes
+
+
+@dataclass
+class _WorkerStep:
+    metadata: AuxOutputConnectorMetadata
+    executions: list[_Execution]
+
+
+@dataclass
+class _SchedulerRequest:
+    request_id: str
+    block_hashes: list[bytes]
+    num_tokens: int
+    num_output_tokens: int = 1
+    num_computed_tokens: int = 0
+    num_in_flight_tokens: int = 0
+    finished: bool = False
+    sampling_params: SimpleNamespace = field(
+        default_factory=lambda: SimpleNamespace(routed_experts_prompt_start=0)
+    )
+
+    def is_finished(self) -> bool:
+        return self.finished
+
+
+def _request_metadata(
+    request_id: str,
+    token_start: int,
+    num_tokens: int,
+    emit_start: int,
+    block_hashes,
+) -> _Execution:
+    return _Execution(request_id, token_start, num_tokens, emit_start, block_hashes)
+
+
+def _metadata(
+    generation: int,
+    requests: list[_Execution],
+    finished_requests: dict[str, list[bytes] | PackedBlockHashes],
+    hash_updates: dict[str, list[bytes] | PackedBlockHashes] | None = None,
+) -> _WorkerStep:
+    block_hashes = {
+        request.request_id: request.block_hashes
+        for request in requests
+        if request.block_hashes
+    }
+    if hash_updates:
+        block_hashes.update(
+            (request_id, hashes)
+            for request_id, hashes in hash_updates.items()
+            if hashes
+        )
+    block_hashes.update(
+        (request_id, hashes)
+        for request_id, hashes in finished_requests.items()
+        if hashes
+    )
+    return _WorkerStep(
+        AuxOutputConnectorMetadata(
+            generation,
+            {request.request_id: request.emit_start for request in requests},
+            block_hashes,
+            tuple(finished_requests),
+        ),
+        requests,
+    )
+
+
+def _execution_ranges(metadata, request_ids):
+    by_request = {request.request_id: request for request in metadata.executions}
+    token_starts = tuple(
+        by_request[request_id].token_start for request_id in request_ids
+    )
+    num_tokens = tuple(by_request[request_id].num_tokens for request_id in request_ids)
+    return token_starts, num_tokens
+
+
+def _begin_step(worker: AuxOutputWorkerConnector, step: _WorkerStep) -> None:
+    worker.begin_step(step.metadata)
+
+
+def _make_worker(
+    max_num_seqs: int,
+    max_num_batched_tokens: int | None = None,
+    max_concurrent_batches: int = 2,
+    max_store_blocks: int | None = None,
+) -> AuxOutputWorkerConnector:
+    object_nbytes = _BLOCK_SIZE * int(np.prod(_SHAPE))
+    store = _make_store(
+        max_bytes=(
+            1 << 20 if max_store_blocks is None else max_store_blocks * object_nbytes
+        ),
+        object_nbytes=object_nbytes,
+    )
+    worker = object.__new__(AuxOutputWorkerConnector)
+    worker._store = store
+    worker._buffer = RoutedExpertsBuffer(
+        _DTYPE,
+        _SHAPE,
+        _BLOCK_SIZE,
+        max_num_seqs,
+        max_num_batched_tokens or max_num_seqs * _BLOCK_SIZE,
+        max_concurrent_batches,
+    )
+    worker._requests = {}
+    worker._generation = 0
+    worker._step_metadata = None
+    worker._pending_outputs = []
+    worker._lock = threading.Lock()
+    worker._max_concurrent_batches = max_concurrent_batches
+    return worker
+
+
+def test_worker_rejects_metadata_generation_rollback():
+    worker = _make_worker(1)
+    _begin_step(worker, _metadata(1, [], {}))
+
+    with pytest.raises(AssertionError, match="generation moved backwards"):
+        _begin_step(worker, _metadata(0, [], {}))
+
+    worker.close()
+
+
+def test_worker_rejects_run_and_finish_in_one_step():
+    worker = _make_worker(1)
+    request = _request_metadata("request", 0, 1, 0, [])
+    metadata = _metadata(
+        0,
+        [request],
+        {"request": []},
+    )
+
+    with pytest.raises(AssertionError, match="cannot run and finish"):
+        _begin_step(worker, metadata)
+
+    worker.close()
+
+
+def _input_batch(request_ids, token_starts, query_start_loc):
+    return SimpleNamespace(
+        req_ids=request_ids,
+        num_computed_tokens_np=token_starts,
+        query_start_loc_np=query_start_loc,
+    )
+
+
+def test_non_output_rank_skips_capture_snapshot():
+    worker = object.__new__(AuxOutputWorkerConnector)
+    worker._store = None
+    worker._buffer = None
+    worker._capturer = Mock()
+    worker._step_metadata = Mock()
+    assert worker.prepare_output(_input_batch([], np.array([]), np.array([]))) is None
+    worker._capturer.snapshot_routing_data.assert_not_called()
+
+
+def test_worker_skips_aux_outputs_for_internal_warmup_step():
+    worker = _make_worker(1)
+    worker._capturer = Mock()
+
+    worker.begin_step(None)
+
+    assert worker.prepare_output(_input_batch([], np.array([]), np.array([]))) is None
+    worker._capturer.snapshot_routing_data.assert_not_called()
+    worker.close()
+
+
+def test_next_step_does_not_consume_pending_output(monkeypatch):
+    """begin_step must not wait for or consume an unconsumed step output."""
+    event = Mock()
+    monkeypatch.setattr(torch.cuda, "Event", lambda **kwargs: event)
+    monkeypatch.setattr(async_utils, "stream", lambda *args: nullcontext())
+    monkeypatch.setattr(
+        async_utils, "async_copy_to_np", lambda tensor: tensor.numpy().copy()
+    )
+    worker = _make_worker(1)
+    worker.begin_step(
+        _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata
+    )
+    rows = torch.ones((1, *_SHAPE), dtype=torch.uint8)
+    worker._capturer = Mock()
+    worker._capturer.snapshot_routing_data.return_value = rows
+    process_output = Mock(wraps=worker.process_output)
+    monkeypatch.setattr(worker, "process_output", process_output)
+    pending = worker.prepare_output(
+        _input_batch(["request"], np.array([0]), np.array([0, 1]))
+    )
+    assert pending is not None
+    output = async_utils.AsyncOutput(
+        ModelRunnerOutput(["request"], {"request": 0}),
+        SamplerOutput(
+            torch.tensor([[7]]), None, None, None, num_rejected=torch.tensor([0])
+        ),
+        torch.tensor([1]),
+        Mock(),
+        Mock(),
+        check_ep_fault=False,
+        pending_aux_output=pending,
+    )
+    event.record.assert_called_once()
+    assert worker._pending_outputs == [pending]
+
+    worker.begin_step(_metadata(0, [], {}).metadata)
+    process_output.assert_not_called()
+    assert worker._pending_outputs == [pending]
+
+    result = output.get_output()
+
+    assert result.sampled_token_ids == [[7]]
+    np.testing.assert_array_equal(
+        result.aux_output_connector_output["request"].rows, rows.numpy()
+    )
+    process_output.assert_called_once()
+    assert worker._pending_outputs == []
+    worker.close()
+
+
+@pytest.mark.parametrize("num_pending", [1, 2])
+def test_finished_request_teardown_waits_for_pending_output(num_pending):
+    """Terminal cleanup waits for every outstanding output, not just the first."""
+    worker = _make_worker(1)
+    rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, *_SHAPE)
+    worker._capturer = Mock()
+    worker._capturer.snapshot_routing_data.return_value = torch.from_numpy(rows)
+    pending_outputs = []
+    for i in range(num_pending):
+        request = _request_metadata("request", 4 * i, 4, 4 * i, [bytes([i]) * 4])
+        worker.begin_step(_metadata(0, [request], {}).metadata)
+        pending = worker.prepare_output(
+            _input_batch(["request"], np.array([4 * i]), np.array([0, 4]))
+        )
+        assert pending is not None
+        pending.enqueue_cpu_copy(np.array([1]), np.array([0]))
+        pending_outputs.append(pending)
+
+    # The next step finishes the request while its output is unconsumed.
+    worker.begin_step(_metadata(0, [], {"request": []}).metadata)
+    assert "request" in worker._requests
+
+    for i, pending in enumerate(pending_outputs):
+        output = worker.process_output(pending)
+        np.testing.assert_array_equal(output["request"].rows, rows)
+        if i + 1 < num_pending:
+            assert "request" in worker._requests
+            assert worker._store._references
+    assert "request" not in worker._requests
+    # The finished request's store keys must be released after commit.
+    assert not worker._store._references
+    worker.close()
+
+
+def test_pending_outputs_bounded_by_max_concurrent_batches():
+    worker = _make_worker(1, max_concurrent_batches=2)
+    worker._capturer = Mock()
+    worker._capturer.snapshot_routing_data.side_effect = lambda num_rows: torch.zeros(
+        (num_rows, *_SHAPE), dtype=torch.uint8
+    )
+    metadata = _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata
+    batch = _input_batch(["request"], np.array([0]), np.array([0, 1]))
+    for _ in range(2):
+        worker.begin_step(metadata)
+        assert worker.prepare_output(batch) is not None
+
+    worker.begin_step(metadata)
+    with pytest.raises(AssertionError, match="not consumed in order"):
+        worker.prepare_output(batch)
+
+    worker.close()
+
+
+def test_worker_rejects_invalid_rejected_token_count():
+    worker = _make_worker(1)
+    metadata = _metadata(
+        0,
+        [_request_metadata("request", 0, 1, 0, [])],
+        {},
+    )
+    with pytest.raises(AssertionError, match="rejected-token count is invalid"):
+        _process_output(
+            worker,
+            metadata,
+            np.zeros((1, *_SHAPE), dtype=_DTYPE),
+            ["request"],
+            np.array([-1]),
+            token_starts=(0,),
+            num_tokens=(1,),
+        )
+    assert worker._pending_outputs == []
+
+    worker.close()
+
+
+def _process_output(
+    worker,
+    step,
+    rows,
+    request_ids,
+    num_rejected,
+    num_sampled=None,
+    *,
+    token_starts=None,
+    num_tokens=None,
+    query_start_loc=None,
+):
+    if num_sampled is None:
+        num_sampled = np.ones(len(request_ids), dtype=np.int32)
+    if token_starts is None or num_tokens is None:
+        assert isinstance(step, _WorkerStep)
+        token_starts, num_tokens = _execution_ranges(step, request_ids)
+    worker._capturer = Mock()
+    worker._capturer.snapshot_routing_data.return_value = torch.from_numpy(rows)
+    metadata = step.metadata if isinstance(step, _WorkerStep) else step
+    worker.begin_step(metadata)
+    if query_start_loc is None:
+        query_start_loc = np.concatenate(
+            (np.zeros(1, dtype=np.int32), np.cumsum(num_tokens, dtype=np.int32))
+        )
+    pending = worker.prepare_output(
+        _input_batch(request_ids, np.asarray(token_starts), query_start_loc)
+    )
+    assert pending is not None
+    pending.enqueue_cpu_copy(num_sampled, num_rejected)
+    return worker.process_output(pending)
+
+
+def test_worker_ignores_cudagraph_query_padding():
+    worker = _make_worker(2)
+    logical = np.arange(2 * 3 * 2, dtype=np.uint8).reshape(2, 3, 2)
+    metadata = _metadata(
+        0,
+        [
+            _request_metadata("first", 0, 1, 0, []),
+            _request_metadata("second", 0, 1, 0, []),
+        ],
+        {},
+    )
+
+    output = _process_output(
+        worker,
+        metadata,
+        logical,
+        ["first", "second"],
+        np.zeros(2, dtype=np.int32),
+        query_start_loc=np.array([0, 1, 2, 2], dtype=np.int32),
+    )
+
+    np.testing.assert_array_equal(output["first"].rows, logical[:1])
+    np.testing.assert_array_equal(output["second"].rows, logical[1:])
+    worker.close()
+
+
+def test_worker_rejects_scheduler_emit_cursor_ahead():
+    worker = _make_worker(1)
+    _begin_step(
+        worker,
+        _metadata(
+            0,
+            [_request_metadata("request", 0, 1, 0, [])],
+            {},
+        ),
+    )
+
+    with pytest.raises(AssertionError, match="Scheduler emit cursor moved ahead"):
+        _begin_step(
+            worker,
+            _metadata(
+                0,
+                [_request_metadata("request", 1, 1, 1, [])],
+                {},
+            ),
+        )
+
+    worker.close()
+
+
+@pytest.mark.parametrize(
+    "routing",
+    [
+        np.zeros((1, _SHAPE[0], _SHAPE[1] + 1), dtype=_DTYPE),
+        np.zeros((1, *_SHAPE), dtype=np.int32),
+    ],
+    ids=["shape", "dtype"],
+)
+def test_worker_rejects_mismatched_capture_profile(routing):
+    worker = _make_worker(1)
+    metadata = _metadata(
+        0,
+        [_request_metadata("request", 0, 1, 0, [])],
+        {},
+    )
+
+    with pytest.raises(AssertionError, match="capture profile changed"):
+        _process_output(
+            worker,
+            metadata,
+            routing,
+            ["request"],
+            np.array([0]),
+        )
+    worker.close()
+
+
+@pytest.mark.parametrize(
+    ("initial_start", "initial_rows", "invalid_start"),
+    [
+        (4, [4, 5, 6], 6),
+        (0, [0], 2),
+        (None, [], 3),
+    ],
+    ids=["overlap", "gap", "unaligned-initial-capture"],
+)
+def test_logical_buffer_rejects_noncontiguous_capture(
+    initial_start, initial_rows, invalid_start
+):
+    buffer = RoutedExpertsBuffer(np.dtype("uint8"), (1,), 4, 1, 4, 1)
+    if initial_start is not None:
+        rows = np.asarray(initial_rows, dtype=np.uint8).reshape(-1, 1)
+        assert buffer.capture("request", initial_start, rows) == []
+
+    with pytest.raises(AssertionError, match="not contiguous"):
+        buffer.capture(
+            "request", invalid_start, np.array([[invalid_start]], dtype=np.uint8)
+        )
+
+
+def test_logical_buffer_captures_one_row_per_decode_step():
+    buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 8, 2)
+    logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=_DTYPE).reshape(_BLOCK_SIZE, 3, 2)
+
+    completed = []
+    for step in range(_BLOCK_SIZE):
+        completed += buffer.capture("request", step, logical[step : step + 1])
+
+    assert len(completed) == 1
+    np.testing.assert_array_equal(completed[0][1], logical)
+
+
+def test_logical_buffer_copies_borrowed_block_when_retained():
+    buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 4, 2)
+    logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=_DTYPE).reshape(_BLOCK_SIZE, *_SHAPE)
+
+    expected = logical.copy()
+    completed = buffer.capture("request", 0, logical)
+    retained = buffer.retain_block(completed[0][1])
+    logical[:] = 0
+
+    assert not np.shares_memory(retained, logical)
+    np.testing.assert_array_equal(retained, expected)
+    buffer.release_block(retained)
+
+
+def _make_store(
+    *,
+    max_bytes: int = 1 << 20,
+    object_nbytes: int = 4,
+):
+    return BlockObjectStore(
+        max_bytes=max_bytes,
+        object_nbytes=object_nbytes,
+    )
+
+
+def test_publish_routed_experts_publishes_full_blocks():
+    store = _make_store(object_nbytes=_BLOCK_SIZE * int(np.prod(_SHAPE)))
+    buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 8, 2)
+    logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2)
+    hashes = [b"a" * 32, b"b" * 32]
+    keys = routed_experts_keys(hashes, "0")
+    blocks = buffer.capture("request", 0, logical)
+    publish_routed_experts(
+        store,
+        batches=[(keys, blocks)],
+        block_size=_BLOCK_SIZE,
+    )
+    expected = logical.copy()
+    logical[:] = 0
+
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            store,
+            keys,
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        expected,
+    )
+    store.close()
+
+
+def test_worker_data_plane_publishes_blocks_and_reuses_prefix():
+    worker = _make_worker(2)
+
+    hashes = [b"a" * 32, b"b" * 32, b"c" * 32]
+    logical = np.arange(10 * 3 * 2, dtype=np.uint8).reshape(10, 3, 2)
+    first = _metadata(
+        generation=0,
+        requests=[_request_metadata("first", 0, 8, 0, hashes)],
+        finished_requests={},
+    )
+    output = _process_output(worker, first, logical[:8], ["first"], np.array([0]))
+    assert output is not None
+    np.testing.assert_array_equal(output["first"].rows, logical[:8])
+
+    second = _metadata(
+        generation=0,
+        requests=[_request_metadata("second", 8, 2, 0, hashes)],
+        finished_requests={"first": []},
+    )
+    output = _process_output(worker, second, logical[8:], ["second"], np.array([0]))
+    assert output is not None
+    np.testing.assert_array_equal(output["second"].rows, logical)
+    worker.close()
+
+
+def test_worker_matches_kv_eviction_order_after_requests_finish():
+    worker = _make_worker(2, max_store_blocks=3)
+    hashes = [bytes([value]) * 32 for value in range(5)]
+    blocks = [
+        np.full((_BLOCK_SIZE, *_SHAPE), value, dtype=_DTYPE) for value in range(5)
+    ]
+
+    def publish_and_finish(request_id, block_indexes):
+        rows = np.concatenate([blocks[index] for index in block_indexes])
+        step = _metadata(
+            0,
+            [
+                _request_metadata(
+                    request_id,
+                    0,
+                    len(rows),
+                    0,
+                    [hashes[index] for index in block_indexes],
+                )
+            ],
+            {},
+        )
+        _process_output(worker, step, rows, [request_id], np.array([0]))
+        _begin_step(worker, _metadata(0, [], {request_id: []}))
+
+    publish_and_finish("s", [0])
+    publish_and_finish("ab", [1, 2])
+    publish_and_finish("c", [3])
+    publish_and_finish("d", [4])
+
+    assert worker._store is not None
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys([hashes[1], hashes[3], hashes[4]], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        np.concatenate([blocks[1], blocks[3], blocks[4]]),
+    )
+    with pytest.raises(BlockObjectStoreError, match="does not exist"):
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys([hashes[2]], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        )
+    worker.close()
+
+
+def test_worker_rejects_unbacked_capture_gap():
+    worker = _make_worker(1)
+    logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
+    first = _metadata(
+        0,
+        [_request_metadata("request", 0, 1, 0, [])],
+        {},
+    )
+    _process_output(
+        worker,
+        first,
+        logical[:1],
+        ["request"],
+        np.array([0]),
+        np.zeros(1, dtype=np.int32),
+    )
+    second = _metadata(
+        0,
+        [_request_metadata("request", 2, 1, 0, [])],
+        {},
+    )
+
+    with pytest.raises(AssertionError, match="unbacked token gap"):
+        _process_output(
+            worker,
+            second,
+            logical[2:3],
+            ["request"],
+            np.array([0]),
+            np.zeros(1, dtype=np.int32),
+        )
+    worker.close()
+
+
+def test_worker_replays_rejected_speculative_gap():
+    worker = _make_worker(1)
+    logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2)
+    first = _metadata(
+        0,
+        [_request_metadata("request", 0, 5, 0, [])],
+        {},
+    )
+    first_output = _process_output(
+        worker, first, logical[:5], ["request"], np.array([4])
+    )
+
+    queued = _metadata(
+        0,
+        [_request_metadata("request", 5, 4, 0, [])],
+        {},
+    )
+    queued_output = _process_output(
+        worker,
+        queued,
+        logical[1:],
+        ["request"],
+        np.array([3]),
+    )
+
+    rolled_back = _metadata(
+        0,
+        [_request_metadata("request", 6, 3, 0, [b"a" * 32])],
+        {},
+    )
+
+    rolled_back_output = _process_output(
+        worker,
+        rolled_back,
+        logical[2:],
+        ["request"],
+        np.array([0]),
+    )
+
+    np.testing.assert_array_equal(first_output["request"].rows, logical[:1])
+    assert queued_output["request"].token_start == 1
+    np.testing.assert_array_equal(
+        queued_output["request"].rows,
+        logical[1:2],
+    )
+    np.testing.assert_array_equal(
+        rolled_back_output["request"].rows,
+        logical[2:],
+    )
+    assert worker._requests["request"].capture_cursor == len(logical)
+    worker.close()
+
+
+def test_worker_rejects_overlapping_capture_rows():
+    worker = _make_worker(1)
+    logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
+    first = _metadata(
+        0,
+        [_request_metadata("request", 0, 3, 0, [])],
+        {},
+    )
+    _process_output(
+        worker,
+        first,
+        logical[:3],
+        ["request"],
+        np.array([0]),
+        np.zeros(1, dtype=np.int32),
+    )
+    stale_overlap = logical[2:].copy()
+    stale_overlap[0] += 100
+    second = _metadata(
+        0,
+        [_request_metadata("request", 2, 2, 0, [b"a" * 32])],
+        {},
+    )
+    with pytest.raises(AssertionError, match="capture moved backwards"):
+        _process_output(
+            worker,
+            second,
+            stale_overlap,
+            ["request"],
+            np.array([0]),
+        )
+    worker.close()
+
+
+def test_worker_publishes_entire_batch_before_materializing_prefix():
+    worker = _make_worker(2)
+    hashes = [b"a" * 32, b"b" * 32]
+    logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2)
+    metadata = _metadata(
+        generation=0,
+        requests=[
+            _request_metadata("consumer", 4, 4, 0, hashes),
+            _request_metadata("producer", 0, 4, 0, hashes),
+        ],
+        finished_requests={},
+    )
+
+    output = _process_output(
+        worker,
+        metadata,
+        np.concatenate((logical[4:], logical[:4])),
+        ["consumer", "producer"],
+        np.array([0, 0]),
+    )
+
+    assert output is not None
+    np.testing.assert_array_equal(output["consumer"].rows, logical)
+    np.testing.assert_array_equal(output["producer"].rows, logical[:4])
+    worker.close()
+
+
+def test_worker_defers_full_block_until_kv_hash_arrives():
+    worker = _make_worker(1)
+    logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2)
+    first = _metadata(
+        0,
+        [_request_metadata("request", 0, 4, 0, [])],
+        {},
+    )
+
+    output = _process_output(
+        worker,
+        first,
+        logical[:4],
+        ["request"],
+        np.array([0]),
+    )
+
+    assert output is not None
+    np.testing.assert_array_equal(output["request"].rows, logical[:4])
+    assert worker._requests["request"].pending_blocks[0][0] == 0
+
+    block_hash = b"a" * 32
+    second = _metadata(
+        0,
+        [_request_metadata("request", 4, 1, 4, [block_hash])],
+        {},
+    )
+    output = _process_output(worker, second, logical[4:], ["request"], np.array([0]))
+
+    assert output is not None
+    np.testing.assert_array_equal(output["request"].rows, logical[4:])
+    assert not worker._requests["request"].pending_blocks
+    assert worker._store is not None
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys([block_hash], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        logical[:4],
+    )
+    worker.close()
+
+
+def test_worker_releases_newly_keyed_blocks_before_capture():
+    worker = _make_worker(1, max_num_batched_tokens=16)
+    connector = _make_connector()
+    block_hashes = [b"a" * 32]
+    request = _scheduler_request("request", block_hashes, num_tokens=100)
+    logical = np.arange(52 * 3 * 2, dtype=np.uint8).reshape(52, 3, 2)
+
+    for token_start, num_tokens in [(0, 4), (4, 16), (20, 16)]:
+        request.num_computed_tokens = token_start
+        metadata = connector.build_connector_meta(
+            _step_output([request.request_id], [token_start], [num_tokens]),
+            {request.request_id: request},
+        )
+        _process_output(
+            worker,
+            metadata,
+            logical[token_start : token_start + num_tokens],
+            [request.request_id],
+            np.array([0]),
+            token_starts=(token_start,),
+            num_tokens=(num_tokens,),
+        )
+
+    block_hashes.extend(bytes([value]) * 32 for value in range(1, 5))
+    request.num_computed_tokens = 36
+    metadata = connector.build_connector_meta(
+        _step_output([request.request_id], [36], [16]),
+        {request.request_id: request},
+    )
+    _process_output(
+        worker,
+        metadata,
+        logical[36:52],
+        [request.request_id],
+        np.array([0]),
+        token_starts=(36,),
+        num_tokens=(16,),
+    )
+
+    state = worker._requests[request.request_id]
+    assert [start for start, _ in state.pending_blocks] == [
+        20,
+        24,
+        28,
+        32,
+        36,
+        40,
+        44,
+        48,
+    ]
+    worker.close()
+
+
+def test_worker_publishes_pending_block_without_forward():
+    worker = _make_worker(1)
+    _begin_step(
+        worker,
+        _metadata(
+            0,
+            [_request_metadata("request", 0, 1, 0, [])],
+            {},
+        ),
+    )
+    logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(_BLOCK_SIZE, 3, 2)
+    assert worker._buffer is not None
+    worker._requests["request"].pending_blocks = [
+        (0, worker._buffer.retain_block(logical))
+    ]
+    block_hash = b"a" * 32
+
+    _begin_step(worker, _metadata(0, [], {}, {"request": [block_hash]}))
+
+    assert not worker._requests["request"].pending_blocks
+    assert worker._store is not None
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys([block_hash], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        logical,
+    )
+    worker.close()
+
+
+def test_worker_releases_unscheduled_request_blocks_before_capture():
+    worker = _make_worker(2, 20, max_concurrent_batches=1)
+    logical = np.arange(48 * 3 * 2, dtype=np.uint8).reshape(48, 3, 2)
+    prompt_hashes = [b"a" * 32, b"b" * 32]
+    prompt = _metadata(
+        0,
+        [
+            _request_metadata("first", 0, 4, 0, prompt_hashes[:1]),
+            _request_metadata("second", 0, 4, 0, prompt_hashes[1:]),
+        ],
+        {},
+    )
+    _process_output(
+        worker,
+        prompt,
+        logical[:8],
+        ["first", "second"],
+        np.zeros(2),
+        np.zeros(2, dtype=np.int32),
+    )
+
+    first_decode = _metadata(
+        0,
+        [_request_metadata("first", 4, 20, 0, [])],
+        {},
+    )
+    _process_output(
+        worker,
+        first_decode,
+        logical[8:28],
+        ["first"],
+        np.zeros(1),
+        np.zeros(1, dtype=np.int32),
+    )
+    assert len(worker._requests["first"].pending_blocks) == 5
+
+    first_hashes = [bytes([value]) * 32 for value in range(2, 7)]
+    second_decode = _metadata(
+        0,
+        [_request_metadata("second", 4, 20, 0, [])],
+        {},
+        {"first": first_hashes},
+    )
+    _process_output(
+        worker,
+        second_decode,
+        logical[28:48],
+        ["second"],
+        np.zeros(1),
+        np.zeros(1, dtype=np.int32),
+    )
+
+    assert not worker._requests["first"].pending_blocks
+    assert len(worker._requests["second"].pending_blocks) == 5
+    worker.close()
+
+
+def test_worker_does_not_rematerialize_emitted_rows():
+    worker = _make_worker(1)
+    hashes = [b"a" * 32, b"b" * 32]
+    logical = np.arange(8 * 3 * 2, dtype=np.uint8).reshape(8, 3, 2)
+
+    first = _metadata(
+        0,
+        [_request_metadata("request", 0, 4, 0, hashes)],
+        {},
+    )
+    output = _process_output(worker, first, logical[:4], ["request"], np.array([0]))
+    assert output is not None
+
+    assert worker._store is not None
+    worker._store.get_concatenated = Mock(wraps=worker._store.get_concatenated)
+    second = _metadata(
+        0,
+        # Async scheduling may build this before the scheduler consumes first.
+        [_request_metadata("request", 4, 4, 0, [])],
+        {},
+    )
+    output = _process_output(worker, second, logical[4:], ["request"], np.array([0]))
+
+    assert output is not None
+    np.testing.assert_array_equal(output["request"].rows, logical[4:])
+    worker._store.get_concatenated.assert_not_called()
+    worker.close()
+
+
+def test_worker_resumes_from_scheduler_emit_boundary():
+    worker = _make_worker(1)
+    logical = np.arange(5 * 3 * 2, dtype=np.uint8).reshape(5, 3, 2)
+    resumed = _metadata(
+        0,
+        [_request_metadata("request", 4, 1, 4, [b"a" * 32])],
+        {},
+    )
+
+    output = _process_output(worker, resumed, logical[4:], ["request"], np.array([0]))
+
+    assert output["request"].token_start == 4
+    np.testing.assert_array_equal(output["request"].rows, logical[4:])
+    assert worker._requests["request"].emit_cursor == 5
+    worker.close()
+
+
+def test_worker_emits_mid_block_chunked_prefill():
+    worker = _make_worker(1)
+    logical = np.arange(3 * 3 * 2, dtype=np.uint8).reshape(3, 3, 2)
+
+    first = _metadata(
+        0,
+        [_request_metadata("request", 0, 2, 0, [])],
+        {},
+    )
+    output = _process_output(
+        worker,
+        first,
+        logical[:2],
+        ["request"],
+        np.array([0]),
+        np.zeros(1, dtype=np.int32),
+    )
+    assert output is not None and not output
+
+    second = _metadata(
+        0,
+        [_request_metadata("request", 2, 1, 0, [])],
+        {},
+    )
+    output = _process_output(worker, second, logical[2:], ["request"], np.array([0]))
+
+    np.testing.assert_array_equal(output["request"].rows, logical)
+    worker.close()
+
+
+def test_worker_emits_published_block_and_mid_block_tail():
+    worker = _make_worker(1)
+    logical = np.arange(7 * 3 * 2, dtype=np.uint8).reshape(7, 3, 2)
+
+    first = _metadata(
+        0,
+        [_request_metadata("request", 0, 6, 0, [b"a" * 32])],
+        {},
+    )
+    output = _process_output(
+        worker,
+        first,
+        logical[:6],
+        ["request"],
+        np.array([0]),
+        np.zeros(1, dtype=np.int32),
+    )
+    assert output is not None and not output
+
+    second = _metadata(
+        0,
+        [_request_metadata("request", 6, 1, 0, [])],
+        {},
+    )
+    output = _process_output(worker, second, logical[6:], ["request"], np.array([0]))
+
+    np.testing.assert_array_equal(output["request"].rows, logical)
+    worker.close()
+
+
+def test_worker_output_does_not_alias_released_tail_buffer():
+    worker = _make_worker(1)
+    logical = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
+    first = _metadata(
+        0,
+        [_request_metadata("request", 0, 2, 0, [b"a" * 32])],
+        {},
+    )
+    _process_output(
+        worker,
+        first,
+        logical[:2],
+        ["request"],
+        np.array([0]),
+        np.zeros(1, dtype=np.int32),
+    )
+    second = _metadata(
+        0,
+        [_request_metadata("request", 2, 2, 0, [])],
+        {},
+    )
+
+    output = _process_output(worker, second, logical[2:], ["request"], np.array([0]))
+    emitted = output["request"].rows
+    expected = logical.copy()
+
+    replacement = np.full_like(logical, 99)
+    reuse = _metadata(
+        0,
+        [_request_metadata("reuse", 0, 4, 0, [b"b" * 32])],
+        {},
+    )
+    _process_output(
+        worker,
+        reuse,
+        replacement,
+        ["reuse"],
+        np.array([0]),
+        np.zeros(1, dtype=np.int32),
+    )
+
+    np.testing.assert_array_equal(emitted, expected)
+    worker.close()
+
+
+def test_worker_merges_block_hash_deltas():
+    worker = _make_worker(1)
+    _begin_step(
+        worker,
+        _metadata(
+            0,
+            [_request_metadata("request", 0, 1, 0, [b"a" * 32, b"b" * 32])],
+            {},
+        ),
+    )
+    _begin_step(
+        worker,
+        _metadata(
+            0,
+            [
+                _request_metadata(
+                    "request",
+                    1,
+                    1,
+                    0,
+                    [b"c" * 32],
+                )
+            ],
+            {},
+        ),
+    )
+
+    assert worker._requests["request"].aux_output_keys == [
+        "vllm-artifact/0/" + (b"a" * 32).hex(),
+        "vllm-artifact/0/" + (b"b" * 32).hex(),
+        "vllm-artifact/0/" + (b"c" * 32).hex(),
+    ]
+    worker.close()
+
+
+def test_worker_finish_publishes_keyed_block_and_discards_request_state():
+    worker = _make_worker(1)
+    logical = np.arange(2 * _BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(
+        2 * _BLOCK_SIZE, 3, 2
+    )
+    running = _metadata(
+        0,
+        [_request_metadata("request", 0, len(logical), 0, [])],
+        {},
+    )
+    _process_output(
+        worker,
+        running,
+        logical,
+        ["request"],
+        np.array([0]),
+        np.array([0]),
+    )
+    block_hash = b"a" * 32
+
+    _begin_step(
+        worker,
+        _metadata(
+            0,
+            [],
+            {"request": [block_hash]},
+        ),
+    )
+
+    assert "request" not in worker._requests
+    assert worker._buffer is not None
+    assert "request" not in worker._buffer._requests
+    assert worker._store is not None
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys([block_hash], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        logical[:_BLOCK_SIZE],
+    )
+    worker.close()
+
+
+def test_worker_fails_when_cached_aux_output_is_missing():
+    worker = _make_worker(1)
+    worker._generation = 0
+    metadata = _metadata(
+        0,
+        [
+            _request_metadata(
+                "request",
+                8,
+                1,
+                0,
+                [b"a" * 32, b"b" * 32],
+            )
+        ],
+        {},
+    )
+
+    with pytest.raises(BlockObjectStoreError, match="does not exist"):
+        _process_output(
+            worker,
+            metadata,
+            np.zeros((1, *_SHAPE), dtype=_DTYPE),
+            ["request"],
+            np.array([0]),
+        )
+    worker.close()
+
+
+def test_store_rejects_oversized_batch_without_partial_write():
+    store = _make_store(max_bytes=6, object_nbytes=3)
+    store.put([BlockObject("retained", b"rrr")])
+
+    with pytest.raises(BlockObjectStoreError, match="cannot retain"):
+        store.put(
+            [
+                BlockObject("first", b"111"),
+                BlockObject("second", b"222"),
+                BlockObject("third", b"333"),
+            ]
+        )
+
+    assert store.get_concatenated(["retained"]) == b"rrr"
+    with pytest.raises(BlockObjectStoreError, match="does not exist"):
+        store.get_concatenated(["first"])
+    store.close()
+
+
+def test_store_lru_and_immutable_put():
+    store = _make_store(max_bytes=8)
+    store.put([BlockObject("first", b"1111"), BlockObject("second", b"2222")])
+    assert store.get_concatenated(["first"]) == b"1111"
+
+    store.put([BlockObject("first", b"1111"), BlockObject("third", b"3333")])
+
+    assert store.get_concatenated(["first", "third"]) == b"11113333"
+    with pytest.raises(BlockObjectStoreError, match="Increase aux_output_config"):
+        store.get_concatenated(["second"])
+    store.close()
+
+
+def test_store_does_not_evict_referenced_aux_outputs():
+    store = _make_store(max_bytes=8)
+    store.put([BlockObject("first", b"1111"), BlockObject("second", b"2222")])
+    store.put([], retain_keys=["first"])
+    store.put([], retain_keys=["first"], release_keys=["first"])
+
+    store.put([BlockObject("third", b"3333")])
+
+    assert store.get_concatenated(["first", "third"]) == b"11113333"
+    with pytest.raises(BlockObjectStoreError, match="does not exist"):
+        store.get_concatenated(["second"])
+    store.put([], release_keys=["first"])
+    store.close()
+
+
+def test_store_terminal_put_preserves_tail_first_eviction_order():
+    store = _make_store(max_bytes=8)
+    keys = ["first", "second", "third"]
+    store.put(
+        [BlockObject("first", b"1111"), BlockObject("second", b"2222")],
+        retain_keys=keys,
+    )
+
+    store.put(
+        [BlockObject("third", b"3333")],
+        release_keys=reversed(keys),
+    )
+
+    assert store.get_concatenated(keys[:2]) == b"11112222"
+    with pytest.raises(BlockObjectStoreError, match="does not exist"):
+        store.get_concatenated(["third"])
+    store.close()
+
+
+def test_store_terminal_put_orders_survivors_for_later_eviction():
+    store = _make_store(max_bytes=12)
+    keys = ["first", "second", "third"]
+    store.put(
+        [BlockObject("first", b"1111"), BlockObject("second", b"2222")],
+        retain_keys=keys,
+    )
+
+    store.put(
+        [BlockObject("third", b"3333")],
+        release_keys=reversed(keys),
+    )
+    store.put([BlockObject("fourth", b"4444")])
+
+    assert store.get_concatenated(["first", "second", "fourth"]) == b"111122224444"
+    with pytest.raises(BlockObjectStoreError, match="does not exist"):
+        store.get_concatenated(["third"])
+    store.close()
+
+
+def test_store_terminal_put_keeps_a_shared_incoming_key():
+    store = _make_store(max_bytes=8)
+    keys = ["first", "second", "third"]
+    store.put(
+        [BlockObject("first", b"1111"), BlockObject("second", b"2222")],
+        retain_keys=(*keys, "third"),
+    )
+
+    store.put(
+        [BlockObject("third", b"3333")],
+        release_keys=reversed(keys),
+    )
+
+    assert store.get_concatenated(["first", "third"]) == b"11113333"
+    with pytest.raises(BlockObjectStoreError, match="does not exist"):
+        store.get_concatenated(["second"])
+    store.put([], release_keys=["third"])
+    store.close()
+
+
+def test_store_rejects_access_after_close():
+    store = _make_store(max_bytes=8)
+    store.close()
+
+    with pytest.raises(RuntimeError, match="closed"):
+        store.put([BlockObject("first", b"1111")])
+    with pytest.raises(RuntimeError, match="closed"):
+        store.get_concatenated(["first"])
+
+
+def test_store_reuses_evicted_slot_without_moving_live_objects():
+    store = _make_store(max_bytes=12)
+    store.put(
+        [
+            BlockObject("first", b"1111"),
+            BlockObject("second", b"2222"),
+            BlockObject("third", b"3333"),
+        ]
+    )
+    assert store.get_concatenated(["second"]) == b"2222"
+    store.put([BlockObject("fourth", b"4444")])
+
+    assert store.get_concatenated(["second", "third", "fourth"]) == b"222233334444"
+    store.close()
+
+
+def _scheduler_request(
+    request_id: str,
+    block_hashes: list[bytes],
+    *,
+    num_tokens: int = 10,
+    num_output_tokens: int = 1,
+    prompt_start: int = 0,
+):
+    return _SchedulerRequest(
+        request_id=request_id,
+        block_hashes=block_hashes,
+        num_tokens=num_tokens,
+        num_output_tokens=num_output_tokens,
+        num_computed_tokens=num_tokens,
+        sampling_params=SimpleNamespace(routed_experts_prompt_start=prompt_start),
+    )
+
+
+def _step_output(
+    request_ids: list[str], token_starts: list[int], token_counts: list[int]
+) -> SchedulerOutput:
+    output = SchedulerOutput.make_empty()
+    output.scheduled_cached_reqs = CachedRequestData(
+        req_ids=request_ids,
+        resumed_req_ids=set(),
+        new_token_ids=[[] for _ in request_ids],
+        all_token_ids={},
+        new_block_ids=[None for _ in request_ids],
+        num_computed_tokens=token_starts,
+        num_output_tokens=[0 for _ in request_ids],
+    )
+    output.num_scheduled_tokens = dict(zip(request_ids, token_counts, strict=True))
+    output.total_num_scheduled_tokens = sum(token_counts)
+    return output
+
+
+def test_scheduler_connector_builds_worker_metadata_and_forwards_output():
+    connector = _make_connector()
+    request = _scheduler_request("request", [b"a" * 32], num_tokens=5)
+    request.num_computed_tokens = 0
+    scheduler_output = _step_output([request.request_id], [0], [4])
+
+    metadata = connector.build_connector_meta(
+        scheduler_output, {request.request_id: request}
+    )
+
+    assert metadata.generation == 0
+    assert metadata.requests == {request.request_id: 4}
+    assert list(metadata.block_hashes[request.request_id]) == [b"a" * 32]
+
+    routing = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
+    output = {"request": AuxRequestOutput(0, routing)}
+    request.num_computed_tokens = 4
+    np.testing.assert_array_equal(connector.take_output(request, output), routing)
+
+
+def test_scheduler_starts_worker_output_at_requested_prompt_token():
+    connector = _make_connector()
+    request = _scheduler_request(
+        "request",
+        [b"a" * 32],
+        num_tokens=4,
+        num_output_tokens=0,
+        prompt_start=2,
+    )
+
+    metadata = connector.build_connector_meta(
+        _step_output([request.request_id], [0], [4]),
+        {request.request_id: request},
+    )
+
+    assert metadata.requests == {request.request_id: 2}
+
+
+def test_scheduler_connector_preserves_request_finish_order():
+    connector = _make_connector()
+    first = _scheduler_request("first", [b"a" * 32])
+    second = _scheduler_request("second", [b"b" * 32])
+    connector.build_connector_meta(
+        _step_output([first.request_id, second.request_id], [0, 0], [0, 0]),
+        {first.request_id: first, second.request_id: second},
+    )
+    connector.request_finished(second)
+    connector.request_finished(first)
+
+    metadata = connector.build_connector_meta(_step_output([], [], []), {})
+
+    assert metadata.finished_requests == (second.request_id, first.request_id)
+    assert not metadata.block_hashes
+
+
+def test_worker_preemption_terminal_keeps_inflight_hashes_for_publish():
+    worker = _make_worker(1)
+    block_hashes = [b"a" * 32, b"b" * 32]
+    logical = np.arange(2 * _BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(
+        2 * _BLOCK_SIZE, 3, 2
+    )
+    first = _metadata(
+        0,
+        [
+            _request_metadata(
+                "request",
+                0,
+                _BLOCK_SIZE,
+                0,
+                block_hashes[:1],
+            )
+        ],
+        {},
+    )
+    _process_output(
+        worker,
+        first,
+        logical[:_BLOCK_SIZE],
+        ["request"],
+        np.array([0]),
+    )
+    inflight = _metadata(
+        0,
+        [
+            _request_metadata(
+                "request",
+                _BLOCK_SIZE,
+                _BLOCK_SIZE,
+                0,
+                block_hashes[1:],
+            )
+        ],
+        {},
+    )
+    _process_output(
+        worker,
+        inflight,
+        logical[_BLOCK_SIZE:],
+        ["request"],
+        np.array([0]),
+        np.zeros(1, dtype=np.int32),
+        token_starts=(_BLOCK_SIZE,),
+        num_tokens=(_BLOCK_SIZE,),
+    )
+    _begin_step(
+        worker,
+        _metadata(
+            0,
+            [],
+            {"request": []},
+        ),
+    )
+
+    assert worker._store is not None
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys(block_hashes[:1], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        logical[:_BLOCK_SIZE],
+    )
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys(block_hashes[1:], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        logical[_BLOCK_SIZE:],
+    )
+    worker.close()
+
+
+def test_worker_releases_terminal_pins_before_same_step_publish():
+    worker = _make_worker(2, max_store_blocks=1)
+    victim_hash = b"a" * 32
+    producer_hash = b"b" * 32
+    victim_rows = np.zeros((_BLOCK_SIZE, *_SHAPE), dtype=_DTYPE)
+    producer_rows = np.ones((_BLOCK_SIZE, *_SHAPE), dtype=_DTYPE)
+    running = _metadata(
+        0,
+        [
+            _request_metadata("victim", 0, _BLOCK_SIZE, 0, [victim_hash]),
+            _request_metadata("producer", 0, _BLOCK_SIZE, 0, []),
+        ],
+        {},
+    )
+    _process_output(
+        worker,
+        running,
+        np.concatenate((victim_rows, producer_rows)),
+        ["victim", "producer"],
+        np.zeros(2, dtype=np.int32),
+        np.zeros(2, dtype=np.int32),
+    )
+
+    _begin_step(
+        worker,
+        _metadata(
+            0,
+            [],
+            {"victim": []},
+            {"producer": [producer_hash]},
+        ),
+    )
+
+    assert worker._store is not None
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys([producer_hash], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        producer_rows,
+    )
+    with pytest.raises(BlockObjectStoreError, match="does not exist"):
+        worker._store.get_concatenated(routed_experts_keys([victim_hash], "0"))
+    worker.close()
+
+
+def test_terminal_request_publishes_hash_discovered_after_last_schedule():
+    connector = _make_connector()
+    worker = _make_worker(1)
+    block_hashes: list[bytes] = []
+    request = _scheduler_request("request", block_hashes, num_tokens=5)
+    request.num_computed_tokens = 0
+    metadata = connector.build_connector_meta(
+        _step_output([request.request_id], [0], [4]),
+        {request.request_id: request},
+    )
+    logical = np.arange(_BLOCK_SIZE * 3 * 2, dtype=np.uint8).reshape(_BLOCK_SIZE, 3, 2)
+
+    _process_output(
+        worker,
+        metadata,
+        logical,
+        [request.request_id],
+        np.array([0]),
+        token_starts=(0,),
+        num_tokens=(_BLOCK_SIZE,),
+    )
+    assert worker._requests[request.request_id].pending_blocks
+
+    block_hash = b"a" * 32
+    block_hashes.append(block_hash)
+    request.num_computed_tokens = _BLOCK_SIZE
+    connector.request_finished(request)
+    cleanup = connector.build_connector_meta(_step_output([], [], []), {})
+    worker.begin_step(cleanup)
+
+    assert worker._store is not None
+    np.testing.assert_array_equal(
+        materialize_routed_experts(
+            worker._store,
+            routed_experts_keys([block_hash], "0"),
+            shape_per_token=_SHAPE,
+            dtype=_DTYPE,
+        ),
+        logical,
+    )
+    worker.close()
+
+
+def test_scheduler_connector_recreates_preempted_request_state():
+    connector = _make_connector()
+    request = _scheduler_request("request", [b"a" * 32], num_tokens=5)
+    connector.build_connector_meta(
+        _step_output([request.request_id], [0], [4]),
+        {request.request_id: request},
+    )
+    connector.request_finished(request)
+    cleanup = connector.build_connector_meta(
+        _step_output([], [], []),
+        {request.request_id: request},
+    )
+    assert request.request_id in cleanup.finished_requests
+    assert request.request_id not in cleanup.block_hashes
+
+    resumed = connector.build_connector_meta(
+        _step_output([request.request_id], [0], [4]),
+        {request.request_id: request},
+    )
+    assert list(resumed.block_hashes[request.request_id]) == [b"a" * 32]
+
+
+def test_scheduler_consumes_ordered_stale_aux_outputs():
+    connector = _make_connector()
+    request = _scheduler_request("request", [], num_tokens=6)
+    connector.build_connector_meta(
+        _step_output([request.request_id], [0], [4]),
+        {request.request_id: request},
+    )
+    first_rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2)
+    second_rows = np.arange(4 * 3 * 2, dtype=np.uint8).reshape(4, 3, 2) + 40
+    first = {"request": AuxRequestOutput(0, first_rows)}
+    second = {"request": AuxRequestOutput(4, second_rows)}
+    connector.request_finished(request)
+
+    request.num_tokens = 5
+    np.testing.assert_array_equal(connector.take_output(request, first), first_rows)
+    request.num_tokens = 6
+    np.testing.assert_array_equal(
+        connector.take_output(request, second), second_rows[:1]
+    )
+
+
+def test_scheduler_rejects_stale_output_without_aux_outputs():
+    connector = _make_connector()
+    request = _scheduler_request("request", [], num_tokens=1)
+    connector.build_connector_meta(
+        _step_output([request.request_id], [0], [1]),
+        {request.request_id: request},
+    )
+
+    with pytest.raises(
+        AssertionError, match="auxiliary output worker output is missing"
+    ):
+        connector.take_output(
+            request,
+            {},
+        )
+
+
+def test_scheduler_rejects_missing_accepted_aux_output_rows():
+    connector = _make_connector()
+    request = _scheduler_request("request", [], num_tokens=2)
+    connector.build_connector_meta(
+        _step_output([request.request_id], [0], [1]),
+        {request.request_id: request},
+    )
+    output = {
+        "request": AuxRequestOutput(
+            0,
+            np.empty((0, *_SHAPE), dtype=_DTYPE),
+        )
+    }
+
+    with pytest.raises(AssertionError, match="invalid token range"):
+        connector.take_output(request, output)
+
+
+def test_scheduler_rejects_aux_output_past_finished_request():
+    connector = _make_connector()
+    request = _scheduler_request("request", [], num_tokens=1)
+    request.finished = True
+    connector.build_connector_meta(
+        _step_output([request.request_id], [0], [1]),
+        {request.request_id: request},
+    )
+    output = {
+        "request": AuxRequestOutput(
+            1,
+            np.empty((0, *_SHAPE), dtype=_DTYPE),
+        )
+    }
+
+    with pytest.raises(
+        AssertionError, match="finished auxiliary output output has no accepted"
+    ):
+        connector.take_output(request, output)
+
+
+@pytest.mark.parametrize("chunk_size", [2, 4])
+@pytest.mark.parametrize("max_tokens", [1, 2])
+def test_prompt_end_offset_through_worker_and_scheduler(chunk_size, max_tokens):
+    """Omit prompt R3 without losing the empty first output or later decode rows."""
+    connector = _make_connector()
+    worker = _make_worker(1)
+    request = _scheduler_request(
+        "request", [b"a" * 32], num_tokens=4, num_output_tokens=0, prompt_start=4
+    )
+    rows = np.arange(5 * 3 * 2, dtype=_DTYPE).reshape(5, *_SHAPE)
+    try:
+        for start in range(0, 4, chunk_size):
+            metadata = connector.build_connector_meta(
+                _step_output([request.request_id], [start], [chunk_size]),
+                {request.request_id: request},
+            )
+            output = _process_output(
+                worker,
+                metadata,
+                rows[start : start + chunk_size],
+                [request.request_id],
+                np.array([0]),
+                num_sampled=np.array([int(start + chunk_size == 4)]),
+                token_starts=(start,),
+                num_tokens=(chunk_size,),
+            )
+            if start + chunk_size < 4:
+                assert output == {}
+        request.num_tokens = 5
+        request.num_output_tokens = 1
+        request.finished = max_tokens == 1
+        result = connector.take_output(request, output)
+        np.testing.assert_array_equal(result, rows[:0])
+
+        if max_tokens == 2:
+            metadata = connector.build_connector_meta(
+                _step_output([request.request_id], [4], [1]),
+                {request.request_id: request},
+            )
+            output = _process_output(
+                worker,
+                metadata,
+                rows[4:],
+                [request.request_id],
+                np.array([0]),
+                token_starts=(4,),
+                num_tokens=(1,),
+            )
+            request.num_tokens = 6
+            request.finished = True
+            np.testing.assert_array_equal(
+                connector.take_output(request, output), rows[4:]
+            )
+    finally:
+        worker.close()
+
+
+def test_scheduler_connector_sends_each_block_hash_once():
+    connector = _make_connector()
+    request = _scheduler_request("request", [b"a" * 32, b"b" * 32], num_tokens=12)
+    scheduler_output = _step_output([request.request_id], [0], [4])
+
+    first = connector.build_connector_meta(
+        scheduler_output, {request.request_id: request}
+    )
+    request.block_hashes.append(b"c" * 32)
+    second = connector.build_connector_meta(
+        _step_output([request.request_id], [4], [4]),
+        {request.request_id: request},
+    )
+    third = connector.build_connector_meta(
+        _step_output([request.request_id], [8], [4]),
+        {request.request_id: request},
+    )
+
+    assert list(first.block_hashes[request.request_id]) == [b"a" * 32, b"b" * 32]
+    assert list(second.block_hashes[request.request_id]) == [b"c" * 32]
+    assert request.request_id not in third.block_hashes
+
+
+def test_scheduler_connector_sends_unscheduled_hash_update():
+    connector = _make_connector()
+    request = _scheduler_request("request", [b"a" * 32], num_tokens=4)
+    connector.build_connector_meta(
+        _step_output([request.request_id], [0], [4]),
+        {request.request_id: request},
+    )
+
+    request.block_hashes.append(b"b" * 32)
+    request.num_tokens = 9
+    request.num_computed_tokens = 8
+    metadata = connector.build_connector_meta(
+        _step_output([], [], []),
+        {request.request_id: request},
+    )
+
+    assert not metadata.requests
+    assert list(metadata.block_hashes[request.request_id]) == [b"b" * 32]
+
+
+def test_scheduler_connector_reset_derives_emit_start_from_request():
+    connector = _make_connector()
+    request = _scheduler_request("request", [b"a" * 32], num_tokens=5)
+    request.num_computed_tokens = 0
+    before_reset = connector.build_connector_meta(
+        _step_output([request.request_id], [0], [4]),
+        {request.request_id: request},
+    )
+    assert list(before_reset.block_hashes[request.request_id]) == [b"a" * 32]
+
+    connector.reset()
+    empty = connector.build_connector_meta(
+        _step_output([], [], []),
+        {request.request_id: request},
+    )
+    assert not empty.block_hashes
+    request.num_computed_tokens = 4
+    scheduler_output = _step_output([request.request_id], [4], [1])
+    metadata = connector.build_connector_meta(
+        scheduler_output, {request.request_id: request}
+    )
+
+    assert metadata.generation == 1
+    assert metadata.requests[request.request_id] == 4
+    assert list(metadata.block_hashes[request.request_id]) == [b"a" * 32]
diff --git a/tests/distributed/conftest.py b/tests/distributed/conftest.py
index 48df856f5e75..1837d3706e9a 100644
--- a/tests/distributed/conftest.py
+++ b/tests/distributed/conftest.py
@@ -18,13 +18,13 @@
 
 @pytest.fixture
 def random_port():
-    """Generate a random port number for testing"""
+    """Generate a random port number for testing."""
     return random.randint(10000, 59900)
 
 
 @pytest.fixture
 def publisher_config(random_port, request):
-    """Create a publisher config with inproc transport"""
+    """Create a publisher config with inproc transport."""
     how = request.param if hasattr(request, "param") else "inproc"
 
     if how == "inproc":
@@ -47,7 +47,7 @@ def publisher_config(random_port, request):
 
 @pytest.fixture
 def publisher(publisher_config):
-    """Create and return a publisher instance"""
+    """Create and return a publisher instance."""
     pub = EventPublisherFactory.create(publisher_config, DP_RANK)
     yield pub
     pub.shutdown()
@@ -55,7 +55,7 @@ def publisher(publisher_config):
 
 @pytest.fixture
 def subscriber(publisher_config):
-    """Create and return a subscriber for testing"""
+    """Create and return a subscriber for testing."""
     endpoint = publisher_config.endpoint
     replay_endpoint = publisher_config.replay_endpoint
 
@@ -74,7 +74,7 @@ def subscriber(publisher_config):
 
 
 class MockSubscriber:
-    """Helper class to receive and verify published events"""
+    """Helper class to receive and verify published events."""
 
     def __init__(
         self,
@@ -91,8 +91,11 @@ def __init__(
         if isinstance(replay_endpoints, str):
             replay_endpoints = [replay_endpoints]
 
+        ipv6 = any("[" in ep for ep in pub_endpoints + (replay_endpoints or []))
+
         # Set up subscriber socket - connect to all endpoints
         self.sub = self.ctx.socket(zmq.SUB)
+        self.sub.setsockopt(zmq.IPV6, ipv6)
         self.sub.setsockopt(zmq.SUBSCRIBE, topic.encode("utf-8"))
         for endpoint in pub_endpoints:
             self.sub.connect(endpoint)
@@ -103,6 +106,7 @@ def __init__(
         if replay_endpoints:
             for replay_endpoint in replay_endpoints:
                 replay = self.ctx.socket(zmq.DEALER)
+                replay.setsockopt(zmq.IPV6, ipv6)
                 replay.connect(replay_endpoint)
                 self.replay_sockets.append(replay)
 
@@ -113,7 +117,7 @@ def __init__(
         self.decoder = msgspec.msgpack.Decoder(type=decode_type)
 
     def receive_one(self, timeout=1000) -> tuple[int, SampleBatch] | None:
-        """Receive a single message with timeout"""
+        """Receive a single message with timeout."""
         if not self.sub.poll(timeout):
             return None
 
@@ -127,7 +131,7 @@ def receive_one(self, timeout=1000) -> tuple[int, SampleBatch] | None:
         return seq, data
 
     def request_replay(self, start_seq: int, socket_idx: int = 0) -> None:
-        """Request replay of messages starting from start_seq"""
+        """Request replay of messages starting from start_seq."""
         if not self.replay_sockets:
             raise ValueError("Replay sockets not initialized")
         if socket_idx >= len(self.replay_sockets):
@@ -138,7 +142,7 @@ def request_replay(self, start_seq: int, socket_idx: int = 0) -> None:
         )
 
     def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]:
-        """Receive replayed messages from a specific replay socket"""
+        """Receive replayed messages from a specific replay socket."""
         if not self.replay_sockets:
             raise ValueError("Replay sockets not initialized")
         if socket_idx >= len(self.replay_sockets):
@@ -170,7 +174,7 @@ def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]:
         return replayed
 
     def close(self):
-        """Clean up resources"""
+        """Clean up resources."""
         self.sub.close()
         for replay in self.replay_sockets:
             replay.close()
diff --git a/tests/distributed/eplb_utils.py b/tests/distributed/eplb_utils.py
index e6a41842e19d..d868069e0a75 100644
--- a/tests/distributed/eplb_utils.py
+++ b/tests/distributed/eplb_utils.py
@@ -13,6 +13,7 @@
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
 )
+from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 
 mp.set_start_method("spawn", force=True)
@@ -73,13 +74,16 @@ def distributed_run(fn, world_size, *args):
 def set_env_vars_and_device(env: dict[str, str]) -> None:
     update_environment_variables(env)
     local_rank = os.environ["LOCAL_RANK"]
-    device = torch.device(f"cuda:{local_rank}")
+    accelerator_type = torch.accelerator.current_accelerator().type
+    device = torch.device(f"{accelerator_type}:{local_rank}")
     torch.accelerator.set_device_index(device)
 
+    backend = "xccl" if current_platform.is_xpu() else "nccl"
+
     # Create a minimal vllm config for init_distributed_environment
     vllm_config = VllmConfig()
     with set_current_vllm_config(vllm_config):
-        init_distributed_environment()
+        init_distributed_environment(backend=backend)
     atexit.register(_destroy_process_group_if_initialized)
     # Ensure each worker process has the same random seed
     random.seed(42)
diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py
index b5e552af4974..f35ba14c1f6b 100644
--- a/tests/distributed/test_comm_ops.py
+++ b/tests/distributed/test_comm_ops.py
@@ -410,6 +410,21 @@ def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) ->
     nccl_selector.assert_not_called()
 
 
+def test_aiter_all_gather_precedes_pynccl(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Ensure can use aiter all_gather if available even if pynccl is not available."""
+    output = torch.empty(2)
+    aiter_comm = Mock()
+    aiter_comm.should_custom_ag.return_value = True
+    aiter_comm.custom_all_gather.return_value = output
+    communicator = CudaCommunicator.__new__(CudaCommunicator)
+    communicator.world_size = 2
+    communicator.aiter_ar_comm = aiter_comm
+    communicator.pynccl_comm = Mock(disabled=True)
+    monkeypatch.setattr(communicator, "_can_use_aiter_ag_rs", Mock(return_value=True))
+
+    assert communicator.all_gatherv(torch.empty(1)) is output
+
+
 def test_isend_object_posts_size_then_object_and_releases_on_wait(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
diff --git a/tests/distributed/test_context_parallel.py b/tests/distributed/test_context_parallel.py
index 007948aeba82..6d5d454b4a8d 100644
--- a/tests/distributed/test_context_parallel.py
+++ b/tests/distributed/test_context_parallel.py
@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-WARNING: This test runs in both single-node (4 GPUs) and multi-node
- (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
- important to set the distributed backend to "mp" to avoid Ray scheduling
- all workers in a node other than the head node, which can cause the test
- to fail.
+"""WARNING: This test runs in both single-node (4 GPUs) and multi-node
+(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
+important to set the distributed backend to "mp" to avoid Ray scheduling
+all workers in a node other than the head node, which can cause the test
+to fail.
 """
 
 import json
diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py
index 7f91f11d7c33..7119e596a097 100644
--- a/tests/distributed/test_custom_all_reduce.py
+++ b/tests/distributed/test_custom_all_reduce.py
@@ -11,6 +11,7 @@
 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce  # noqa
 from vllm.distributed.device_communicators import custom_all_reduce as car
 from vllm.distributed.parallel_state import get_tp_group, graph_capture
+from vllm.platforms import current_platform
 
 from ..utils import (
     ensure_model_parallel_initialized,
@@ -24,6 +25,116 @@
     test_sizes[i] -= v % 8
 
 
+@ray.remote(num_gpus=1, max_calls=1)
+def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port):
+    from vllm.model_executor.kernels.mhc.tilelang import (
+        mhc_post_tilelang,
+        mhc_pre_delayed_tilelang,
+    )
+
+    with monkeypatch.context() as m:
+        m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+        device = torch.device(f"cuda:{rank}")
+        torch.accelerator.set_device_index(device)
+        init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+        ensure_model_parallel_initialized(tp_size, pp_size)
+        comm = get_tp_group().device_communicator.ca_comm
+        assert comm is not None and comm.mnnvl_lamport_ag_multicast_ptr
+        fn = torch.zeros(24, 20480, device=device)
+        scale = torch.ones(3, device=device)
+        base = torch.zeros(24, device=device)
+
+        # Changing shapes and interleaving all-gather exercise the shared
+        # Lamport stages, including cleanup of a larger previous payload.
+        def run(n):
+            torch.manual_seed(42 + rank)
+            x = torch.randn(n, 5120, device=device, dtype=torch.bfloat16)
+            # Packed +0/-0 pairs collide with the Lamport sentinel.
+            x[:, :16] = 0
+            x[:, 9:16:2] = -0.0
+            torch.manual_seed(123)
+            residual = torch.randn(n, 4, 5120, device=device, dtype=torch.bfloat16)
+            post = torch.rand(n, 4, device=device)
+            comb = torch.randn(n, 4, 4, device=device) * 0.1
+            pre = torch.rand(n, 4, device=device)
+            weight = torch.randn(5120, device=device, dtype=torch.bfloat16)
+            output = torch.empty_like(residual)
+            normalized = torch.empty_like(x)
+
+            def fused():
+                torch.ops._C_custom_ar.all_reduce_mhc(
+                    x,
+                    residual,
+                    post,
+                    comb,
+                    pre,
+                    weight,
+                    output,
+                    normalized,
+                    comm.mnnvl_lamport_ag_local_ptr,
+                    comm.mnnvl_lamport_ag_multicast_ptr,
+                    comm.mnnvl_lamport_epochs[0],
+                    rank,
+                    comm.mnnvl_buffer_size,
+                    1e-6,
+                )
+
+            def check():
+                gathered = comm.custom_all_gather(x)
+                assert gathered is not None
+                peers = gathered.view(tp_size, n, 5120).float()
+                reduced = peers[0].clone()
+                for peer in peers[1:]:
+                    reduced.add_(peer)
+                expected = mhc_post_tilelang(
+                    reduced.bfloat16(), residual, post.unsqueeze(-1), comb
+                )
+                expected_norm = mhc_pre_delayed_tilelang(
+                    expected,
+                    fn,
+                    scale,
+                    base,
+                    1e-6,
+                    1e-6,
+                    1e-6,
+                    2.0,
+                    20,
+                    pre_mix=pre,
+                    norm_weight=weight,
+                )[2]
+                torch.testing.assert_close(output, expected, rtol=0.008, atol=1e-6)
+                torch.testing.assert_close(
+                    normalized, expected_norm, rtol=0.008, atol=0.008
+                )
+
+            fused()
+            check()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                for _ in range(6):
+                    fused()
+            for _ in range(20):
+                graph.replay()
+            check()
+            x.mul_(0.5)
+            residual.mul_(2)
+            graph.replay()
+            check()
+
+        # Cover fixed Q6 and shrinking/growing adaptive verification batches.
+        for n in (1, 6, 12, 8, 16, 3, 5, 2, 4, 1):
+            run(n)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100), reason="Requires SM100"
+)
+def test_all_reduce_mhc_preserves_bf16_boundaries_and_graph_replay(monkeypatch):
+    if torch.accelerator.device_count() < 4:
+        pytest.skip("Requires four GPUs with NVLink multicast")
+    multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc)
+
+
 @pytest.mark.parametrize(
     ("dtype", "expected"),
     [
diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py
index 01c254c2d039..1dc6b5d3bbdd 100644
--- a/tests/distributed/test_elastic_ep.py
+++ b/tests/distributed/test_elastic_ep.py
@@ -10,6 +10,8 @@
 import pytest
 import requests
 
+from vllm.platforms import current_platform
+
 from ..evals.gsm8k.gsm8k_eval import evaluate_gsm8k
 from ..utils import RemoteOpenAIServer, multi_gpu_test
 
@@ -27,7 +29,7 @@ def cleanup_ray_between_tests():
 NUM_GSM8K_QUESTIONS = 256
 EXPECTED_ACCURACY = 0.58
 ACCURACY_TOL = 0.08
-MAX_NUM_SEQS = 32
+DEFAULT_ALL2ALL_BACKEND = "allgather_reducescatter"
 
 
 def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool:
@@ -161,7 +163,12 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:
     return accuracy
 
 
-def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]:
+def _base_serve_args(
+    dp_size: int = 2,
+    max_dp_size: int = 4,
+    enforce_eager: bool = False,
+    all2all_backend: str = DEFAULT_ALL2ALL_BACKEND,
+) -> list[str]:
     args = [
         "--trust-remote-code",
         "--tensor-parallel-size",
@@ -170,12 +177,12 @@ def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]
         "0.8",
         "--max-model-len",
         "4096",
-        "--max-num-seqs",
-        str(MAX_NUM_SEQS),
         "--enable-expert-parallel",
         "--all2all-backend",
-        "allgather_reducescatter",
+        all2all_backend,
         "--enable-elastic-ep",
+        "--elastic-ep-max-dp-size",
+        str(max_dp_size),
         "--enable-eplb",
         "--eplb-config.num_redundant_experts",
         "0",
@@ -205,28 +212,45 @@ def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]
 
 
 @pytest.mark.parametrize(
-    ("enforce_eager", "traffic_mode"),
+    ("enforce_eager", "traffic_mode", "all2all_backend"),
     [
-        pytest.param(True, "none", id="enforce_eager_none"),
-        pytest.param(True, "light", id="enforce_eager_light"),
-        pytest.param(True, "heavy", id="enforce_eager_heavy"),
-        pytest.param(False, "heavy", id="cuda_graphs_heavy"),
+        pytest.param(True, "none", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_none"),
+        pytest.param(True, "light", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_light"),
+        pytest.param(True, "heavy", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_heavy"),
+        pytest.param(
+            False,
+            "heavy",
+            DEFAULT_ALL2ALL_BACKEND,
+            id="cuda_graphs_heavy",
+        ),
+        pytest.param(False, "heavy", "nixl_ep", id="cuda_graphs_heavy_nixl_ep"),
     ],
 )
 @multi_gpu_test(num_gpus=4)
-def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str):
+def test_elastic_ep_scaling(
+    enforce_eager: bool, traffic_mode: str, all2all_backend: str
+):
     from vllm.distributed.eplb.eplb_communicator import has_nixl
+    from vllm.utils.import_utils import has_nixl_ep
 
     if not has_nixl():
         pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
+    if current_platform.is_xpu():
+        pytest.skip("Elastic EP with NIXL is not supported on XPU")
+    if all2all_backend == "nixl_ep" and not has_nixl_ep():
+        pytest.skip("NIXL EP is not installed")
 
     initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2"))
     target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4"))
     assert target_dp_size > initial_dp_size
-    vllm_serve_args = _base_serve_args(initial_dp_size, enforce_eager)
+    vllm_serve_args = _base_serve_args(
+        initial_dp_size, target_dp_size, enforce_eager, all2all_backend
+    )
 
     with RemoteOpenAIServer(
-        MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
+        MODEL_NAME,
+        vllm_serve_args,
+        max_wait_seconds=1200,
     ) as server:
         initial_accuracy = _run_gsm8k_eval(server, "Initial")
 
@@ -255,55 +279,3 @@ def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str):
             f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})"
         )
         print(f"  Tolerance:  {ACCURACY_TOL:.3f}")
-
-
-@multi_gpu_test(num_gpus=4)
-def test_elastic_ep_scaling_uneven():
-    """Test scale up with uneven worker distribution.
-
-    This tests the case where num_new_workers % old_dp_size != 0,
-    specifically 2 -> 3 where remainder = 1 % 2 = 1.
-    This exercises the remainder handling in sender-receiver pairing.
-    """
-    from vllm.distributed.eplb.eplb_communicator import has_nixl
-
-    if not has_nixl():
-        pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
-
-    vllm_serve_args = _base_serve_args()
-
-    with RemoteOpenAIServer(
-        MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
-    ) as server:
-        initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)")
-
-        # Scale 2 -> 3: This has remainder = 1 % 2 = 1
-        # Tests uneven sender-receiver pairing
-        assert _send_scale_command(server, 3)
-        scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)")
-
-        assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, (
-            f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than "
-            f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
-        )
-
-        # Scale back down to 2
-        assert _send_scale_command(server, 2)
-        scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)")
-
-        assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, (
-            f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than "
-            f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
-        )
-
-        print("\nAccuracy Summary (Uneven Scaling):")
-        print(f"  Initial:    {initial_accuracy:.3f}")
-        print(
-            f"  Scale up:   {scale_up_accuracy:.3f} "
-            f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})"
-        )
-        print(
-            f"  Scale down: {scale_down_accuracy:.3f} "
-            f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})"
-        )
-        print(f"  Tolerance:  {ACCURACY_TOL:.3f}")
diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py
new file mode 100644
index 000000000000..ba03387aac63
--- /dev/null
+++ b/tests/distributed/test_engram_dp_shard.py
@@ -0,0 +1,787 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Engram tables shared by co-located DP replicas.
+
+Four GPUs on one node split the hash heads four ways, whether as DEP
+(TP1 x DP4, the single-node MoE layout) or TP2 x DP2. A lookup only returns
+the right rows if the ids of every replica are gathered, each rank looks up
+its own heads, and each replica trades those tokens back for the heads its
+peers own.
+
+With shared host storage, DP replicas instead map the same TP slice and
+prefetch only their own tokens without DP gathers.
+"""
+
+from itertools import product
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from threading import Event
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.multiprocessing as mp
+
+from vllm.config import (
+    EngramConfig,
+    VllmConfig,
+    get_current_vllm_config_or_none,
+    set_current_vllm_config,
+)
+from vllm.config.load import LoadConfig
+from vllm.config.parallel import ParallelConfig
+from vllm.config.scheduler import SchedulerConfig
+from vllm.distributed import cleanup_dist_env_and_memory, parallel_state
+from vllm.distributed.parallel_state import (
+    get_engram_dp_group,
+    get_engram_dp_size,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.forward_context import set_forward_context
+from vllm.models.deepseek_v41.common.engram import EngramLayout
+from vllm.models.deepseek_v41.nvidia import engram as engram_ops
+from vllm.models.deepseek_v41.nvidia.engram import (
+    Engram,
+    ParallelEngramEmbedding,
+    engram_head_shard_rank,
+    gather_engram_hashes,
+)
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.system_utils import update_environment_variables
+
+WORLD_SIZE = 4
+DIM, BLOCK = 64, 32
+# Prime-sized buckets, as the model's layout builds them.
+HEAD_SIZES = (97, 101, 103, 107, 109, 113, 127, 131)
+
+
+def _token_counts(dp_size: int) -> tuple[tuple[int, ...], ...]:
+    """Even loads, then uneven ones with an idle replica to pad around."""
+    return ((8,) * dp_size, (7, 0, 5, 1)[:dp_size], (7,) * dp_size)
+
+
+def _full_table(rows: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """The unsharded table, identical on every rank."""
+    torch.manual_seed(0)
+    weight = (torch.randn(rows, DIM) * 4).to(torch.float8_e4m3fn)
+    scale_inv = torch.randint(120, 134, (rows, DIM // BLOCK), dtype=torch.uint8)
+    return weight, scale_inv
+
+
+def _make_ids(head_sizes: tuple[int, ...], num_tokens: int, seed: int) -> torch.Tensor:
+    """Ids of each head, drawn from that head's own bucket range."""
+    torch.manual_seed(seed)
+    offset = 0
+    columns = []
+    for size in head_sizes:
+        columns.append(torch.randint(offset, offset + size, (num_tokens, 1)))
+        offset += size
+    return torch.cat(columns, dim=1).to(torch.int32).cuda()
+
+
+def _reference(weight, scale_inv, ids) -> torch.Tensor:
+    """Dequantized lookup against the whole table."""
+    valid = ids >= 0
+    safe_ids = ids.clamp_min(0).long()
+    values = torch.nn.functional.embedding(safe_ids, weight).float()
+    scales = torch.nn.functional.embedding(safe_ids, scale_inv)
+    scales = (scales.to(torch.int32) << 23).view(torch.float32)
+    values = values.unflatten(-1, (-1, BLOCK)) * scales.unsqueeze(-1)
+    return values.flatten(-2).to(torch.bfloat16).masked_fill(~valid.unsqueeze(-1), 0)
+
+
+@pytest.mark.parametrize(
+    "architecture,options,edp_ranks,etp_ranks",
+    [
+        ("DeepseekV41ForCausalLM", {}, [0, 2], [0, 1]),
+        ("Qwen4ExpForCausalLM", {}, None, [0, 1]),
+        (
+            "Qwen4ExpForCausalLM",
+            {"embedding_across_dp": True},
+            None,
+            [0, 1, 2, 3],
+        ),
+        ("Qwen4ExpForConditionalGeneration", {}, None, [0, 1]),
+    ],
+)
+def test_engram_group_creation_honors_model_and_options(
+    monkeypatch, architecture, options, edp_ranks, etp_ranks
+):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(architecture=architecture, is_moe=False),
+        parallel_config=ParallelConfig(
+            tensor_parallel_size=2,
+            data_parallel_size=2,
+            distributed_executor_backend="mp",
+        ),
+        engram_config=EngramConfig(**options),
+    )
+    for name in (
+        "_TP",
+        "_ETP",
+        "_ENGRAM_DP",
+        "_DCP",
+        "_PCP",
+        "_PP",
+        "_DP",
+        "_EP",
+        "_EPLB",
+    ):
+        monkeypatch.setattr(parallel_state, name, None)
+    monkeypatch.setattr(parallel_state, "_WORLD", SimpleNamespace(local_rank=0))
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 4)
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0)
+    monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1)
+
+    def make_group(group_ranks, *args, **kwargs):
+        ranks = next(ranks for ranks in group_ranks if 0 in ranks)
+        return SimpleNamespace(ranks=ranks, rank_in_group=0, world_size=len(ranks))
+
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", make_group)
+    with set_current_vllm_config(config):
+        initialize_model_parallel(tensor_model_parallel_size=2, backend="nccl")
+    group = get_engram_dp_group()
+    assert (group.ranks if group is not None else None) == edp_ranks
+    assert parallel_state.get_etp_group().ranks == etp_ranks
+
+
+def _worker(rank: int, tp_size: int, port: int) -> None:
+    torch.accelerator.set_device_index(rank)
+    update_environment_variables(
+        {
+            "RANK": str(rank),
+            "LOCAL_RANK": str(rank),
+            "WORLD_SIZE": str(WORLD_SIZE),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(port),
+        }
+    )
+    dp_size = WORLD_SIZE // tp_size
+    dp_rank, tp_rank = divmod(rank, tp_size)
+    vllm_config = VllmConfig(
+        parallel_config=ParallelConfig(
+            tensor_parallel_size=tp_size,
+            data_parallel_size=dp_size,
+            data_parallel_rank=dp_rank,
+        ),
+        scheduler_config=SchedulerConfig(
+            max_model_len=8,
+            is_encoder_decoder=False,
+            max_num_batched_tokens=8,
+            max_num_seqs=8,
+        ),
+    )
+    # Supply the resolved Engram settings without constructing a full model.
+    vllm_config.engram_config = EngramConfig()
+    vllm_config.model_config = SimpleNamespace(
+        architecture="DeepseekV41ForCausalLM", is_moe=True
+    )
+    try:
+        init_distributed_environment()
+        with set_current_vllm_config(vllm_config):
+            initialize_model_parallel(tensor_model_parallel_size=tp_size)
+        assert get_engram_dp_size() == dp_size
+        assert get_engram_dp_group().ranks == [
+            tp_rank + replica * tp_size for replica in range(dp_size)
+        ]
+        assert engram_head_shard_rank() == tp_rank * dp_size + dp_rank
+
+        for failure in ("create", "open", "mmap", "register", "pinned", None):
+            _check_shared_storage_lifetime_and_failures(failure)
+
+        # Reuse the processes and communication groups across storage modes.
+        for (cpu_offload, dp_shared_memory), n_heads, sequence_parallel in product(
+            [(False, False), (True, False), (True, True)],
+            [8, 7],
+            [False, True] if tp_size > 1 else [False],
+        ):
+            _check_table(
+                vllm_config,
+                cpu_offload,
+                dp_shared_memory,
+                n_heads,
+                sequence_parallel,
+            )
+    finally:
+        cleanup_dist_env_and_memory()
+
+
+def _check_shared_storage_lifetime_and_failures(failure):
+    """One peer's failure must roll back every mapping/registration, without hangs."""
+    import gc
+    import weakref
+
+    group = get_engram_dp_group()
+    runtime = torch.cuda.cudart()
+    make_file = engram_ops.tempfile.NamedTemporaryFile
+    map_file = engram_ops.mmap.mmap
+    paths, mappings, registrations, unregistrations = [], [], [], []
+    failed_rank = 0 if failure == "create" else 1
+
+    def fail():
+        raise OSError(f"injected {failure} failure")
+
+    def temporary_file(*args, **kwargs):
+        if failure == "create":
+            fail()
+        file = make_file(*args, **kwargs)
+        paths.append(Path(file.name))
+        return file
+
+    def open_file(*args, **kwargs):
+        if failure == "open" and group.rank_in_group == failed_rank:
+            fail()
+        return open(*args, **kwargs)
+
+    def mapping(*args, **kwargs):
+        if failure == "mmap" and group.rank_in_group == failed_rank:
+            fail()
+        result = map_file(*args, **kwargs)
+        mappings.append(weakref.ref(result))
+        return result
+
+    def register(pointer, size, flags):
+        if failure == "register" and group.rank_in_group == failed_rank:
+            return SimpleNamespace(value=1)
+        result = runtime.cudaHostRegister(pointer, size, flags)
+        assert result.value == 0
+        registrations.append(pointer)
+        return result
+
+    def unregister(pointer):
+        unregistrations.append(pointer)
+        result = runtime.cudaHostUnregister(pointer)
+        assert result.value == 0
+        return result
+
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(
+            engram_ops,
+            "tempfile",
+            SimpleNamespace(NamedTemporaryFile=temporary_file),
+        )
+        patch.setattr(engram_ops, "open", open_file, raising=False)
+        patch.setattr(
+            engram_ops,
+            "mmap",
+            SimpleNamespace(mmap=mapping, MAP_SHARED=engram_ops.mmap.MAP_SHARED),
+        )
+        patch.setattr(
+            torch.cuda,
+            "cudart",
+            lambda: SimpleNamespace(
+                cudaHostRegister=register,
+                cudaHostUnregister=unregister,
+            ),
+        )
+        if failure == "pinned" and group.rank_in_group == failed_rank:
+            patch.setattr(torch.Tensor, "is_pinned", lambda self: False)
+        if failure is not None:
+            stage = "cudaHostRegister" if failure in ("register", "pinned") else failure
+            with pytest.raises(RuntimeError, match=f"EDP rank {failed_rank}.*{stage}"):
+                engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group)
+            assert unregistrations == registrations
+        else:
+            storage = engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group)
+            assert all(not path.exists() for path in paths)
+            storage_ref = weakref.ref(storage)
+            parameter = torch.nn.Parameter(storage.weight, requires_grad=False)
+            views = storage.get_views(parameter, storage.weight_scale_inv)
+            del storage
+            gc.collect()
+            assert storage_ref() is None
+            assert not unregistrations
+            del parameter
+            gc.collect()
+            assert not unregistrations  # UVA views still retain the CPU allocation.
+            del views
+        gc.collect()
+        assert unregistrations == registrations
+        for ref in mappings:
+            mapped = ref()
+            assert mapped is None or mapped.closed
+        assert all(not path.exists() for path in paths)
+
+
+def _check_table(
+    vllm_config, cpu_offload, dp_shared_memory, n_heads, sequence_parallel
+):
+    parallel = vllm_config.parallel_config
+    dp_size, tp_size = parallel.data_parallel_size, parallel.tensor_parallel_size
+    dp_rank = parallel.data_parallel_rank
+    tp_rank = engram_ops.get_tensor_model_parallel_rank()
+    multithread = (
+        dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1)
+    )
+    load_config = (
+        LoadConfig(
+            load_format="safetensors",
+            use_tqdm_on_load=False,
+            model_loader_extra_config={
+                "enable_multithread_load": True,
+                "num_threads": 2,
+            },
+        )
+        if multithread
+        else vllm_config.load_config
+    )
+    with pytest.MonkeyPatch.context() as m:
+        head_sizes = HEAD_SIZES[:n_heads]
+        rows = sum(head_sizes)
+        config = SimpleNamespace(
+            hidden_size=DIM,
+            hc_mult=2,
+            rms_norm_eps=1e-6,
+            engram_layer_ids=[0],
+            engram_num_embeddings=[rows],
+            engram_max_ngram_size=2,
+            engram_n_heads=n_heads,
+            engram_head_dim=DIM,
+            engram_compressed_vocab_size=32,
+            engram_pad_token_id=0,
+            engram_vocab_size=HEAD_SIZES[0],
+        )
+        with set_current_vllm_config(vllm_config):
+            weight, scale_inv = _full_table(rows)
+            layout = EngramLayout(config)
+            # Stub configuration inputs, not Engram's initialization or execution.
+            component_config = SimpleNamespace(
+                engram_config=EngramConfig(
+                    cpu_offload=cpu_offload,
+                    dp_shared_memory=dp_shared_memory,
+                ),
+                scheduler_config=vllm_config.scheduler_config,
+                load_config=load_config,
+            )
+            with m.context() as init_patch, torch.device("cuda"):
+                init_patch.setattr(
+                    engram_ops, "get_current_vllm_config", lambda: component_config
+                )
+                engram = Engram(
+                    config,
+                    quant_config=None,
+                    layout=layout,
+                    layer_hash_index=0,
+                    use_sequence_parallel=sequence_parallel,
+                    prefix="model.layers.0.engram",
+                )
+            layer = engram.embed_tokens
+            if multithread:
+                _load_shared_multithread(layer, weight, scale_inv, load_config)
+            else:
+                # Only the leader has valid checkpoint payload in shared mode.
+                loader_weight = weight if not dp_shared_memory or dp_rank == 0 else None
+                loader_scale = (
+                    scale_inv if not dp_shared_memory or dp_rank == 0 else None
+                )
+                layer.weight.weight_loader(layer.weight, loader_weight)
+                layer.weight_scale_inv.weight_loader(
+                    layer.weight_scale_inv, loader_scale
+                )
+            if dp_shared_memory:
+                assert layer.dp_size == 1
+                assert layer.head_start == tp_rank * layer.part_n_hash_cols
+                assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned()
+
+                def unexpected_collective(*args, **kwargs):
+                    raise AssertionError("shared lookup must not use DP collectives")
+
+                m.setattr(get_engram_dp_group(), "all_gather", unexpected_collective)
+
+            token_counts = _token_counts(dp_size)
+
+        # Production forward runs after the construction-time config has exited.
+        assert get_current_vllm_config_or_none() is None
+        for batch_idx, counts in enumerate(token_counts):
+            num_tokens = counts[dp_rank]
+            # Ids differ per replica but must agree across its TP ranks.
+            ids = _make_ids(head_sizes, num_tokens, seed=100 + dp_rank)
+            if batch_idx == 2 and dp_rank == 1:
+                ids.fill_(engram_ops.DEAD_ID)
+            with set_forward_context(
+                None,
+                vllm_config,
+                num_tokens=num_tokens,
+                num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32),
+            ):
+                gathered = gather_engram_hashes(
+                    ids, dp_shared_memory=layer.dp_shared_memory
+                )
+                expected_tokens = (
+                    num_tokens if dp_shared_memory else max(counts) * dp_size
+                )
+                assert gathered.shape[0] == expected_tokens
+                engram.prepare_embeddings(gathered)
+                looked_up = engram.embed(ids)
+                standalone = layer(ids)
+            expected = _reference(weight.cuda(), scale_inv.cuda(), ids)
+            torch.testing.assert_close(standalone, expected, atol=0, rtol=0)
+            if sequence_parallel:
+                # SP hands back only this rank's token shard, zero-padded.
+                chunk = -(-num_tokens // tp_size)
+                expected = torch.nn.functional.pad(
+                    expected, (0, 0, 0, 0, 0, chunk * tp_size - num_tokens)
+                )
+                expected = expected[tp_rank * chunk : (tp_rank + 1) * chunk]
+            torch.testing.assert_close(looked_up, expected, atol=0, rtol=0)
+
+        if (
+            cpu_offload
+            and n_heads == 7
+            and not sequence_parallel
+            and (not dp_shared_memory or tp_size == 1)
+        ):
+            _check_dummy_hash_model_forward(
+                vllm_config,
+                engram,
+                head_sizes,
+                weight,
+                scale_inv,
+                dp_shared_memory,
+            )
+
+        replay = n_heads == 7 and (
+            (dp_shared_memory and sequence_parallel == (tp_size > 1))
+            or (
+                cpu_offload
+                and not dp_shared_memory
+                and sequence_parallel == (tp_size > 1)
+            )
+        )
+        if replay:
+            captures = ("full", "breakable")
+            for capture in captures:
+                _check_prefetch_replay(
+                    engram,
+                    head_sizes,
+                    weight,
+                    scale_inv,
+                    dp_shared_memory,
+                    capture,
+                )
+        if dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1):
+            with m.context() as storage_patch:
+                storage_patch.setattr(layer.weight, "data", layer.weight.data.clone())
+                with pytest.raises(RuntimeError, match="storage must not be replaced"):
+                    layer(ids)
+
+
+def _check_dummy_hash_model_forward(
+    vllm_config, engram, head_sizes, weight, scales, dp_shared_memory
+):
+    """A metadata-less replica joins DP sharding but skips a shared lookup."""
+    from vllm.models.deepseek_v41.common.engram import NgramHashState
+    from vllm.models.deepseek_v41.nvidia import model as model_ops
+
+    dp_rank = vllm_config.parallel_config.data_parallel_rank
+    if dp_shared_memory and dp_rank != 1:
+        return
+    tokens = 4
+    ids = _make_ids(head_sizes, tokens, seed=300 + dp_rank)
+    state = NgramHashState.__new__(NgramHashState)
+    torch.nn.Module.__init__(state)
+    state.multipliers = torch.empty(1, 2, dtype=torch.int64, device="cuda")
+    state.primes = torch.empty(1, 1, len(head_sizes), dtype=torch.int64, device="cuda")
+    state.ensure_cache = lambda: True
+    # Hash arithmetic is covered separately; keep the real model's branch and
+    # NgramHashState.dummy_hashes, plus real preparation/lookup collectives.
+    state.forward = lambda *args: ids.unsqueeze(1)
+
+    class Decoder(SimpleNamespace):
+        # Engram hashes and the mask stay the last positional arguments;
+        # **kwargs absorbs the decoder's keyword-only flags.
+        def __call__(self, hidden, positions, input_ids, *args, **kwargs):
+            hashes, keep = args[-2:]
+            if dp_shared_memory and dp_rank == 1:
+                assert hashes is None and keep is None
+                output = hidden.new_zeros(
+                    tokens, len(head_sizes), DIM, dtype=torch.bfloat16
+                )
+            else:
+                assert hashes is not None and keep is not None
+                assert bool(keep.all()) == (dp_rank != 1)
+                if dp_rank == 1:
+                    assert torch.all(hashes == engram_ops.DEAD_ID)
+                output = engram.embed(hashes[:, 0])
+            # Trailing None is previous_aux; this stub captures no aux states.
+            return output, None, None, None, None, None
+
+    model = SimpleNamespace(
+        use_mega_moe=False,
+        use_sequence_parallel=False,
+        fuse_mhc_all_reduce=False,
+        engram_hash=state,
+        engram_swa_prefix="swa",
+        engram_dp_shared_memory=dp_shared_memory,
+        layers=[Decoder(engram=engram)],
+        start_layer=0,
+        end_layer=1,
+        aux_hidden_state_layers=(),
+    )
+    metadata = (
+        None
+        if dp_rank == 1
+        else {
+            "swa": SimpleNamespace(
+                query_start_loc=None, slot_mapping=None, block_table=None
+            )
+        }
+    )
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(
+            model_ops,
+            "get_pp_group",
+            lambda: SimpleNamespace(
+                is_first_rank=True,
+                is_last_rank=False,
+            ),
+        )
+        patch.setattr(model_ops, "mhc_post_tilelang", lambda hidden, *args: hidden)
+        with set_forward_context(
+            metadata,
+            vllm_config,
+            num_tokens=tokens,
+            num_tokens_across_dp=torch.full(
+                (vllm_config.parallel_config.data_parallel_size,),
+                tokens,
+                dtype=torch.int32,
+            ),
+        ):
+            output = model_ops.DeepseekV4Model.forward(
+                model,
+                torch.arange(tokens, device="cuda"),
+                torch.arange(tokens, device="cuda"),
+                intermediate_tensors=None,
+                inputs_embeds=torch.zeros(tokens, DIM, device="cuda"),
+                lookback_token_ids=torch.full((1, 1), -1, device="cuda"),
+            )["hidden_states"]
+    expected_ids = torch.full_like(ids, engram_ops.DEAD_ID) if dp_rank == 1 else ids
+    expected = _reference(weight.cuda(), scales.cuda(), expected_ids)
+    torch.testing.assert_close(output, expected, atol=0, rtol=0)
+
+
+def _load_shared_multithread(layer, weight, scale_inv, load_config):
+    """Opposite shard completion orders must leave every shared parameter ready."""
+    from safetensors.torch import save_file
+
+    from vllm.model_executor.model_loader import weight_utils
+    from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
+
+    tensors = {
+        "weight": weight,
+        "weight_scale_inv": scale_inv.view(torch.float8_e8m0fnu),
+    }
+    order = list(tensors)
+    if get_engram_dp_group().rank_in_group % 2:
+        order.reverse()
+    first_loaded = Event()
+    original_load = weight_utils.load_file
+
+    def load_file(path, **kwargs):
+        if Path(path).stem != order[0]:
+            assert first_loaded.wait(timeout=30), "First shared weight load stalled"
+        return original_load(path, **kwargs)
+
+    with TemporaryDirectory() as directory, pytest.MonkeyPatch.context() as m:
+        for name, tensor in tensors.items():
+            save_file({name: tensor}, str(Path(directory) / f"{name}.safetensors"))
+        m.setattr(weight_utils, "load_file", load_file)
+        loader = DefaultModelLoader(load_config)
+        source = DefaultModelLoader.Source(directory, revision=None)
+        loaded = []
+        for name, tensor in loader._get_weights_iterator(source):
+            param = getattr(layer, name)
+            param.weight_loader(param, tensor)
+            loaded.append(name)
+            first_loaded.set()
+        assert loaded == order
+
+    # Check before any further collective could mask incomplete shared writes.
+    for name, tensor in tensors.items():
+        param = getattr(layer, name)
+        expected = tensor.narrow(0, param.engram_vocab_start, param.shape[0])
+        torch.testing.assert_close(
+            param.view(torch.uint8), expected.view(torch.uint8), atol=0, rtol=0
+        )
+
+
+def _check_prefetch_replay(
+    engram, head_sizes, weight, scale_inv, dp_shared_memory, capture
+):
+    """DP-sharded and shared prefetches must replay with new IDs."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    ids = _make_ids(head_sizes, 8, seed=501)
+    tp_rank = engram_ops.get_tensor_model_parallel_rank()
+    tokens = 8 // engram.embed_tokens.tp_size if engram.use_sequence_parallel else 8
+    output = torch.empty(
+        tokens, len(head_sizes), DIM, device="cuda", dtype=torch.bfloat16
+    )
+
+    def step():
+        gathered = gather_engram_hashes(ids.clone(), dp_shared_memory=dp_shared_memory)
+        engram.prepare_embeddings(gathered)
+        output.copy_(engram.embed(ids))
+
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 8)
+        warmup = torch.cuda.Stream()
+        warmup.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(warmup):
+            step()
+        if capture == "breakable":
+            graph = BreakableCUDAGraphCapture()
+            with torch.cuda.stream(warmup), graph:
+                step()
+            assert graph.num_eager_breaks == 2
+        else:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=warmup):
+                step()
+        torch.cuda.current_stream().wait_stream(warmup)
+        for iteration in range(2):
+            ids.copy_(_make_ids(head_sizes, 8, seed=600 + iteration))
+            graph.replay()
+            expected = _reference(weight.cuda(), scale_inv.cuda(), ids)
+            if engram.use_sequence_parallel:
+                expected = expected[tp_rank * tokens : (tp_rank + 1) * tokens]
+            torch.testing.assert_close(output, expected, atol=0, rtol=0)
+        torch.accelerator.synchronize()
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("tp_size", [1, 2])
+def test_engram_table_shared_across_dp_replicas(tp_size, monkeypatch):
+    """Check all storage modes with one process launch per TP/DP topology."""
+    if torch.accelerator.device_count() < WORLD_SIZE:
+        pytest.skip(f"Need {WORLD_SIZE} GPUs to run the test.")
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
+    mp.spawn(_worker, args=(tp_size, get_open_port()), nprocs=WORLD_SIZE)
+    cleanup_dist_env_and_memory()
+
+
+@pytest.mark.parametrize(
+    "cpu_offload,dp_size,error",
+    [
+        (False, 4, "requires cpu_offload=True"),
+        (True, 1, "effective Engram DP size is 1"),
+    ],
+)
+def test_engram_dp_shared_memory_runtime_requirements(
+    monkeypatch, cpu_offload, dp_size, error
+):
+    """Report unmet storage/topology requirements before allocating CUDA memory."""
+    monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: dp_size)
+    with pytest.raises(ValueError, match=error):
+        ParallelEngramEmbedding(
+            sum(HEAD_SIZES),
+            DIM,
+            HEAD_SIZES,
+            cpu_offload=cpu_offload,
+            dp_shared_memory=True,
+        )
+
+
+def test_engram_tables_too_large_for_shm_are_not_shared(monkeypatch):
+    """A /dev/shm smaller than the tables must fall back instead of failing startup."""
+    monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: 2)
+    layout = SimpleNamespace(num_embeddings=(1 << 50,), head_dim=DIM)
+    assert not engram_ops.can_share_engram_tables(layout)
+
+
+@pytest.mark.parametrize(
+    "node_ids,dp_size,replica_size,expected",
+    [
+        ([0] * 4, 1, 4, 1),
+        ([0] * 4, 4, 1, 4),
+        ([0] * 4, 2, 2, 2),
+        ([0] * 8 + [1] * 8, 8, 2, 4),
+        ([0] * 6 + [1] * 6, 4, 1, 2),
+        ([0] * 3 + [1] * 4, 7, 1, 1),
+        ([0] * 8 + [1] * 8 + [2] * 8, 8, 3, 1),
+        ([0, 1] * 4, 4, 2, 1),
+        ([0] * 3 + [1] * 5, 4, 2, 1),
+    ],
+)
+def test_engram_dp_shard_size_respects_node_boundaries(
+    node_ids, dp_size, replica_size, expected, monkeypatch
+):
+    """Sharing must fall back when replicas or node rank ranges do not align."""
+    monkeypatch.setattr(parallel_state, "get_node_count", lambda: len(set(node_ids)))
+    monkeypatch.setattr(
+        parallel_state, "get_world_group", lambda: SimpleNamespace(cpu_group=None)
+    )
+    monkeypatch.setattr(
+        parallel_state,
+        "in_the_same_node_as",
+        lambda group, src: [node == node_ids[src] for node in node_ids],
+    )
+    assert (
+        parallel_state._engram_dp_shard_size(len(node_ids), dp_size, replica_size)
+        == expected
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [0, 2, 4, 5])
+def test_engram_hash_padding_has_no_valid_rows(num_tokens, monkeypatch):
+    """Idle/padded slots must not query row zero; oversized batches are rejected."""
+    monkeypatch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 4)
+    monkeypatch.setattr(
+        engram_ops,
+        "get_engram_dp_group",
+        lambda: SimpleNamespace(all_gather=lambda ids, dim: torch.cat([ids, ids], dim)),
+    )
+    ids = torch.full((num_tokens, 2, 3), 17, dtype=torch.int32)
+    if num_tokens > 4:
+        with pytest.raises(ValueError, match="exceeds the DP token slot"):
+            gather_engram_hashes(ids)
+        return
+    gathered = gather_engram_hashes(ids).reshape(2, 4, 2, 3)
+    for replica in gathered:
+        torch.testing.assert_close(replica[:num_tokens], ids)
+        assert torch.all(replica[num_tokens:] == engram_ops.DEAD_ID)
+
+
+@pytest.mark.parametrize("tp_size", [1, 2])
+@pytest.mark.parametrize("uniform", [False, True])
+def test_engram_hash_padding_ignores_other_nodes(tp_size, uniform, monkeypatch):
+    """Two four-GPU nodes pad only to their own maximum, including TP x DP."""
+    group_size = 4 // tp_size
+    counts = (
+        [4] * (8 // tp_size)
+        if uniform
+        else [4096] + [2] * (group_size - 1) + [1, 3] + [0] * (group_size - 2)
+    )
+    monkeypatch.setattr(
+        engram_ops,
+        "get_forward_context",
+        lambda: SimpleNamespace(
+            dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=torch.tensor(counts))
+        ),
+    )
+    for dp_rank in range(8 // tp_size):
+        slot = 4 if uniform else (4096 if dp_rank < group_size else 3)
+        ids = torch.full((counts[dp_rank], 2, 3), 17, dtype=torch.int32)
+        monkeypatch.setattr(
+            engram_ops,
+            "get_dp_group",
+            lambda rank=dp_rank: SimpleNamespace(rank_in_group=rank),
+        )
+        monkeypatch.setattr(
+            engram_ops,
+            "get_engram_dp_group",
+            lambda rank=dp_rank: SimpleNamespace(
+                rank_in_group=rank % group_size,
+                world_size=group_size,
+                all_gather=lambda tensor, dim: torch.cat([tensor] * group_size, dim),
+            ),
+        )
+        gathered = gather_engram_hashes(ids)
+        expected = ids.new_full((slot, 2, 3), engram_ops.DEAD_ID)
+        expected[: ids.shape[0]] = ids
+        torch.testing.assert_close(
+            gathered, expected.repeat(group_size, 1, 1), msg=f"DP rank {dp_rank}"
+        )
diff --git a/tests/distributed/test_eplb_algo.py b/tests/distributed/test_eplb_algo.py
index 20360c8c61b9..8cd5cf9e3c05 100644
--- a/tests/distributed/test_eplb_algo.py
+++ b/tests/distributed/test_eplb_algo.py
@@ -10,6 +10,7 @@
     compute_logical_maps,
 )
 from vllm.distributed.eplb.policy.default import DefaultEplbPolicy
+from vllm.platforms import current_platform
 
 
 def test_eplb_load_stats_reduce_across_ranks():
@@ -29,7 +30,7 @@ def test_eplb_load_stats_reduce_across_ranks():
 
 
 def test_basic_rebalance():
-    """Test basic rebalancing functionality"""
+    """Test basic rebalancing functionality."""
     # Example from https://github.com/deepseek-ai/eplb
     weight = torch.tensor(
         [
@@ -92,7 +93,7 @@ def test_basic_rebalance():
 
 
 def test_single_gpu_case():
-    """Test single GPU case"""
+    """Test single GPU case."""
     weight = torch.tensor([[10, 20, 30, 40]])
     num_replicas = 4
     num_groups = 1
@@ -115,7 +116,7 @@ def test_single_gpu_case():
 
 
 def test_equal_weights():
-    """Test case with equal weights"""
+    """Test case with equal weights."""
     weight = torch.tensor([[50, 50, 50, 50, 50, 50, 50, 50]])
     num_replicas = 8
     num_groups = 2
@@ -139,7 +140,7 @@ def test_equal_weights():
 
 
 def test_extreme_weight_imbalance():
-    """Test extreme weight imbalance case"""
+    """Test extreme weight imbalance case."""
     weight = torch.tensor([[1000, 1, 1, 1, 1, 1, 1, 1]])
     num_replicas = 12
     num_groups = 2
@@ -162,7 +163,7 @@ def test_extreme_weight_imbalance():
 
 
 def test_multiple_layers():
-    """Test multiple layers case"""
+    """Test multiple layers case."""
     weight = torch.tensor(
         [
             [10, 20, 30, 40, 50, 60],  # First layer
@@ -195,7 +196,7 @@ def test_multiple_layers():
 
 
 def test_parameter_validation():
-    """Test parameter validation"""
+    """Test parameter validation."""
     weight = torch.tensor([[10, 20, 30, 40]])
 
     # Test non-divisible case - this should handle normally without throwing
@@ -213,7 +214,7 @@ def test_parameter_validation():
 
 
 def test_small_scale_hierarchical():
-    """Test small-scale hierarchical load balancing"""
+    """Test small-scale hierarchical load balancing."""
     weight = torch.tensor(
         [
             [100, 50, 200, 75, 150, 25, 300, 80],  # 8 experts
@@ -243,7 +244,7 @@ def test_small_scale_hierarchical():
 
 
 def test_global_load_balance_fallback():
-    """Test global load balancing fallback case"""
+    """Test global load balancing fallback case."""
     # When num_groups % num_nodes != 0, should fall back to global load
     # balancing
     weight = torch.tensor([[10, 20, 30, 40, 50, 60]])
@@ -263,11 +264,13 @@ def test_global_load_balance_fallback():
     assert torch.sum(logcnt) == num_replicas
 
 
-@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("device", ["cpu", "cuda", "xpu"])
 def test_device_compatibility(device):
-    """Test device compatibility"""
+    """Test device compatibility."""
     if device == "cuda" and not torch.cuda.is_available():
         pytest.skip("CUDA not available")
+    if device == "xpu" and not current_platform.is_xpu():
+        pytest.skip("XPU not available")
 
     weight = torch.tensor([[10, 20, 30, 40]], device=device)
     num_replicas = 6
@@ -287,8 +290,7 @@ def test_device_compatibility(device):
 
 
 def test_additional_cases():
-    """Test more edge cases and different parameter combinations"""
-
+    """Test more edge cases and different parameter combinations."""
     # Test case 1: Large-scale distributed setup
     weight1 = torch.tensor(
         [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]]
@@ -320,8 +322,7 @@ def test_additional_cases():
 
 
 def test_compute_logical_maps_with_negative_indices():
-    """
-    Test that compute_logical_maps correctly handles physical slots containing
+    """Test that compute_logical_maps correctly handles physical slots containing
     -1 (unused slots).
     """
     # 2 layers, 6 physical slots, 4 logical experts.
diff --git a/tests/distributed/test_eplb_events.py b/tests/distributed/test_eplb_events.py
index 40323dc89809..f924b107d354 100644
--- a/tests/distributed/test_eplb_events.py
+++ b/tests/distributed/test_eplb_events.py
@@ -4,15 +4,17 @@
 import threading
 import time
 
+import pytest
 import torch
 
-from vllm.distributed.eplb.eplb_utils import CpuGpuEvent
+from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream
+from vllm.platforms import current_platform
 
 
 def test_wait_blocks_until_record():
     event = CpuGpuEvent()
-    record_stream = torch.cuda.Stream()
-    wait_stream = torch.cuda.Stream()
+    record_stream = torch.Stream()
+    wait_stream = torch.Stream()
     wait_returned = threading.Event()
 
     def waiter():
@@ -33,8 +35,8 @@ def waiter():
 
 def test_reuse_across_multiple_cycles():
     wrapper = CpuGpuEvent()
-    record_stream = torch.cuda.Stream()
-    wait_stream = torch.cuda.Stream()
+    record_stream = torch.Stream()
+    wait_stream = torch.Stream()
     NUM_CYCLES = 8
     completed_cycles = []
     barriers = [threading.Barrier(2) for _ in range(NUM_CYCLES)]
@@ -56,14 +58,18 @@ def waiter():
     assert len(completed_cycles) == NUM_CYCLES
 
 
-def test_producer_consumer():
+@pytest.mark.parametrize("device", ["cuda", "xpu"])
+def test_producer_consumer(device):
+    """This test uses the CpuGpuEvent to synchronize reads and writes to/from
+    a shared GPU tensor on multiple CPU threads.
     """
-    This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared GPU
-    tensor on multiple CPU threads.
-    """
-    worker_stream = torch.cuda.Stream()
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+    if device == "xpu" and not current_platform.is_xpu():
+        pytest.skip("XPU not available")
+    worker_stream = torch.Stream()
     # Create a single element counter that will be shared between two threads
-    buf = torch.zeros(1, device="cuda")
+    buf = torch.zeros(1, device=device)
     NUM_ROUNDS = 5
 
     ready_cpu = [threading.Event() for _ in range(NUM_ROUNDS)]
@@ -77,7 +83,7 @@ def worker():
             if i > 0:
                 events[i - 1].wait(stream=worker_stream)
 
-            with torch.cuda.stream(worker_stream):
+            with device_stream(worker_stream):
                 buf.fill_(float(i + 1))
 
             worker_stream.synchronize()
diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py
index 21fa057fd201..2037535e3ddf 100644
--- a/tests/distributed/test_eplb_execute.py
+++ b/tests/distributed/test_eplb_execute.py
@@ -7,6 +7,8 @@
 import torch
 import torch.distributed
 
+import vllm.distributed.eplb.eplb_communicator as eplb_comm
+import vllm.utils.gpu_sync_debug as gsd
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.distributed.eplb.eplb_communicator import (
     create_eplb_communicator,
@@ -21,18 +23,41 @@
     ensure_model_parallel_initialized,
     get_tp_group,
 )
+from vllm.platforms import current_platform
 
 from .eplb_utils import distributed_run, set_env_vars_and_device
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_gloo_receive_staging_does_not_force_gpu_sync(monkeypatch):
+    """Gloo receives must reach the GPU without an implicit pageable copy."""
+    monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
+    monkeypatch.setattr(gsd, "_sync_check_enabled", True)
+    gsd._install_copy_checkers()
+    monkeypatch.setattr(eplb_comm, "is_local_first_rank", lambda: False)
+
+    monkeypatch.setattr(eplb_comm, "P2POp", lambda op, tensor, peer, group: tensor)
+
+    def receive(tensors):
+        for tensor in tensors:
+            tensor.fill_(7)
+        return []
+
+    monkeypatch.setattr(eplb_comm, "batch_isend_irecv", receive)
+    communicator = eplb_comm.TorchDistGlooStagedEplbCommunicator(cpu_group=None)
+    dst = torch.empty(32, device="cuda")
+    communicator.add_recv([dst], src_rank=1, expert_id=0)
+    gsd.with_gpu_sync_check(communicator.execute)()
+    torch.testing.assert_close(dst.cpu(), torch.full((32,), 7.0))
+
+
 def create_expert_indices_with_redundancy(
     num_layers: int,
     num_logical_experts: int,
     total_physical_experts: int,
     redundancy_config: list[int],  # redundancy for each logical expert
 ) -> torch.Tensor:
-    """
-    Create expert indices with redundancy.
+    """Create expert indices with redundancy.
 
     Args:
         num_layers: number of layers
@@ -42,6 +67,7 @@ def create_expert_indices_with_redundancy(
 
     Returns:
         indices: Shape (num_layers, total_physical_experts)
+
     """
     assert sum(redundancy_config) == total_physical_experts
     assert len(redundancy_config) == num_logical_experts
@@ -70,8 +96,7 @@ def create_expert_weights(
     device: torch.device,
     physical_to_logical_mapping: torch.Tensor,
 ) -> list[list[torch.Tensor]]:
-    """
-    Create fake expert weights tensor for testing.
+    """Create fake expert weights tensor for testing.
 
     Use `arange` to generate predictable weights values, based on logical
     expert ID.
@@ -80,6 +105,7 @@ def create_expert_weights(
     Args:
         physical_to_logical_mapping: Shape (num_layers, num_local_experts)
             mapping[layer, physical_pos] = logical_expert_id
+
     """
     expert_weights = []
 
@@ -183,12 +209,14 @@ def verify_redundant_experts_have_same_weights(
     ep_rank: int,
     world_size: int,
     num_local_experts: int,
+    cpu_group,
 ) -> bool:
-    """
-    Verify that all replicas of the same logical expert have the same weights.
-    """
+    """Verify that all replicas of the same logical expert have the same weights."""
     num_layers = len(expert_weights)
     total_physical_experts = world_size * num_local_experts
+    # XPU: gloo does not support XPU tensors, so gather on CPU then move back.
+    # CUDA: gloo supports CUDA tensors directly, keep original device.
+    use_cpu_gather = current_platform.is_xpu()
 
     ok = True
     for layer in range(num_layers):
@@ -196,31 +224,35 @@ def verify_redundant_experts_have_same_weights(
         all_weights: list[torch.Tensor] = []
 
         for weight_idx, hidden_size in enumerate(hidden_sizes):
+            orig_device = expert_weights[layer][weight_idx].device
+
+            if use_cpu_gather:
+                local_weights = expert_weights[layer][weight_idx].cpu()
+                gather_device = "cpu"
+            else:
+                local_weights = expert_weights[layer][weight_idx]
+                gather_device = orig_device
+
             # Create tensor to store all expert weights
             # Shape: [total_physical_experts, hidden_size]
             gathered_weights = torch.zeros(
                 total_physical_experts,
                 hidden_size,
-                device=expert_weights[layer][weight_idx].device,
-                dtype=expert_weights[layer][weight_idx].dtype,
+                device=gather_device,
+                dtype=local_weights.dtype,
             )
 
-            # Use all_gather to collect expert weights from current node
-            # expert_weights[layer][weight_idx] shape:
-            # [num_local_experts, hidden_size]
-            local_weights = expert_weights[layer][
-                weight_idx
-            ]  # [num_local_experts, hidden_size]
-
             # Split tensor along dim 0 into a list for all_gather
             gathered_weights_list = torch.chunk(gathered_weights, world_size, dim=0)
 
             torch.distributed.all_gather(
-                # Output list: each element corresponds to one rank's weights
                 list(gathered_weights_list),
-                local_weights,  # Input: current rank's local weights
+                local_weights,
+                group=cpu_group,
             )
 
+            if use_cpu_gather:
+                gathered_weights = gathered_weights.to(orig_device)
             all_weights.append(gathered_weights)
 
         # Verify that all replicas of the same logical expert have the same
@@ -271,9 +303,23 @@ def verify_redundant_experts_have_same_weights(
     return ok
 
 
-def assert_verification_synced(local_ok: bool, msg: str) -> None:
-    ok_tensor = torch.tensor([1 if local_ok else 0], device="cuda", dtype=torch.int32)
-    torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN)
+def assert_verification_synced(local_ok: bool, msg: str, cpu_group) -> None:
+    if current_platform.is_xpu():
+        # XPU: default group backend (gloo) does not support XPU tensors;
+        # use CPU tensor with cpu_group instead.
+        ok_tensor = torch.tensor(
+            [1 if local_ok else 0], device="cpu", dtype=torch.int32
+        )
+        torch.distributed.all_reduce(
+            ok_tensor, op=torch.distributed.ReduceOp.MIN, group=cpu_group
+        )
+    else:
+        # CUDA: nccl (default group) supports device tensors directly.
+        device = torch.accelerator.current_accelerator()
+        ok_tensor = torch.tensor(
+            [1 if local_ok else 0], device=device, dtype=torch.int32
+        )
+        torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN)
     assert bool(ok_tensor.item()), msg
 
 
@@ -313,8 +359,10 @@ def _test_async_transfer_layer_without_mtp_worker(
 
         ep_group_coordinator = get_tp_group()
         ep_group = ep_group_coordinator.device_group
+        cpu_group = ep_group_coordinator.cpu_group
         ep_rank = torch.distributed.get_rank()
-        device = torch.device(f"cuda:{ep_rank}")
+        accelerator_type = torch.accelerator.current_accelerator().type
+        device = torch.device(f"{accelerator_type}:{ep_rank}")
 
         total_physical_experts = world_size * num_local_experts
         hidden_sizes = [16, 32]
@@ -353,7 +401,7 @@ def _test_async_transfer_layer_without_mtp_worker(
         new_indices_cpu = new_indices.cpu()
 
         expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
-        cuda_stream = torch.cuda.Stream(device=device)
+        stream = torch.Stream(device=device)
 
         communicator = create_eplb_communicator_or_raise(
             group_coordinator=ep_group_coordinator,
@@ -361,7 +409,7 @@ def _test_async_transfer_layer_without_mtp_worker(
             expert_weights=expert_weights,
             expert_buffer=expert_buffer,
         )
-        communicator.set_stream(cuda_stream)
+        communicator.set_stream(stream)
 
         for layer_idx in range(num_layers):
             transfer_metadata = transfer_layer(
@@ -371,10 +419,10 @@ def _test_async_transfer_layer_without_mtp_worker(
                 expert_weights_buffer=expert_buffer,
                 ep_group=ep_group,
                 communicator=communicator,
-                cuda_stream=cuda_stream,
+                stream=stream,
                 layer_idx=layer_idx,
             )
-            cuda_stream.synchronize()
+            stream.synchronize()
             move_from_buffer(
                 expert_weights=expert_weights[layer_idx],
                 expert_weights_buffers=expert_buffer,
@@ -398,6 +446,7 @@ def _test_async_transfer_layer_without_mtp_worker(
             ep_rank,
             world_size,
             num_local_experts,
+            cpu_group=cpu_group,
         )
         and local_ok
     )
@@ -405,6 +454,7 @@ def _test_async_transfer_layer_without_mtp_worker(
         local_ok,
         "Async transfer verification failed on at least one rank. "
         "See logs for details.",
+        cpu_group=cpu_group,
     )
 
 
@@ -429,9 +479,11 @@ def _test_rearrange_expert_weights_with_redundancy(
         )
 
         ep_group_coordinator = get_tp_group()
-        ep_group = ep_group_coordinator.cpu_group
+        ep_group = ep_group_coordinator.device_group
+        cpu_group = ep_group_coordinator.cpu_group
         ep_rank = torch.distributed.get_rank()
-        device = torch.device(f"cuda:{ep_rank}")
+        accelerator_type = torch.accelerator.current_accelerator().type
+        device = torch.device(f"{accelerator_type}:{ep_rank}")
 
         # Test parameters
         total_physical_experts = world_size * num_local_experts
@@ -500,12 +552,14 @@ def _test_rearrange_expert_weights_with_redundancy(
             ep_rank,
             world_size,
             num_local_experts,
+            cpu_group=cpu_group,
         )
         and local_ok
     )
     assert_verification_synced(
         local_ok,
         "Rearrange verification failed on at least one rank. See logs for details.",
+        cpu_group=cpu_group,
     )
 
 
@@ -533,7 +587,7 @@ def _test_rearrange_expert_weights_with_redundancy(
     ],
 )
 @pytest.mark.parametrize(
-    "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
+    "eplb_communicator", ["torch_nccl", "torch_gloo", "torch_xccl", "pynccl", "nixl"]
 )
 def test_rearrange_expert_weights_with_redundancy(
     world_size,
@@ -543,9 +597,14 @@ def test_rearrange_expert_weights_with_redundancy(
     eplb_communicator,
 ):
     """Test the functionality of rearranging expert weights with redundancy."""
-
     if eplb_communicator == "nixl" and not has_nixl():
         pytest.skip("NIXL is not available")
+    if eplb_communicator == "nixl" and current_platform.is_xpu():
+        pytest.skip("NIXL does not support XPU")
+    if eplb_communicator in ("torch_nccl", "pynccl") and not torch.cuda.is_available():
+        pytest.skip(f"{eplb_communicator} requires CUDA")
+    if eplb_communicator == "torch_xccl" and not current_platform.is_xpu():
+        pytest.skip("torch_xccl requires XPU")
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Need at least {world_size} GPUs to run the test")
     distributed_run(
@@ -570,9 +629,11 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
         )
 
         ep_group_coordinator = get_tp_group()
-        ep_group = ep_group_coordinator.cpu_group
+        ep_group = ep_group_coordinator.device_group
+        cpu_group = ep_group_coordinator.cpu_group
         ep_rank = torch.distributed.get_rank()
-        device = torch.device(f"cuda:{ep_rank}")
+        accelerator_type = torch.accelerator.current_accelerator().type
+        device = torch.device(f"{accelerator_type}:{ep_rank}")
 
         num_layers = 2
         num_local_experts = 2
@@ -601,9 +662,10 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
             original_weights.append(layer_copy)
 
         expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
+        default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl"
         communicator = create_eplb_communicator_or_raise(
             group_coordinator=ep_group_coordinator,
-            backend="torch_nccl",
+            backend=default_backend,
             expert_weights=expert_weights,
             expert_buffer=expert_buffer,
         )
@@ -635,6 +697,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
     assert_verification_synced(
         local_ok,
         "No-change EPLB verification failed on at least one rank.",
+        cpu_group=cpu_group,
     )
 
 
@@ -644,7 +707,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
         (2, 2, 2, 3),
     ],
 )
-@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "nixl"])
+@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "torch_xccl", "nixl"])
 def test_async_transfer_layer_without_mtp(
     world_size: int,
     num_layers: int,
@@ -653,9 +716,12 @@ def test_async_transfer_layer_without_mtp(
     eplb_communicator: str,
 ):
     """Exercise async EPLB transfer path without MTP/spec decode."""
-
     if eplb_communicator == "nixl" and not has_nixl():
         pytest.skip("NIXL is not available")
+    if eplb_communicator == "nixl" and current_platform.is_xpu():
+        pytest.skip("NIXL does not support XPU")
+    if eplb_communicator == "torch_xccl" and not current_platform.is_xpu():
+        pytest.skip("torch_xccl requires XPU")
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Need at least {world_size} GPUs to run the test")
 
@@ -671,11 +737,9 @@ def test_async_transfer_layer_without_mtp(
 
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_rearrange_expert_weights_no_change(world_size):
-    """
-    Test that when the indices do not change, the weights should remain
+    """Test that when the indices do not change, the weights should remain
     unchanged.
     """
-
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Need at least {world_size} GPUs to run the test")
     distributed_run(
@@ -696,9 +760,11 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
         )
 
         ep_group_coordinator = get_tp_group()
-        ep_group = ep_group_coordinator.cpu_group
+        ep_group = ep_group_coordinator.device_group
+        cpu_group = ep_group_coordinator.cpu_group
         ep_rank = torch.distributed.get_rank()
-        device = torch.device(f"cuda:{ep_rank}")
+        accelerator_type = torch.accelerator.current_accelerator().type
+        device = torch.device(f"{accelerator_type}:{ep_rank}")
 
         num_layers = 1
         num_local_experts = 2
@@ -734,9 +800,10 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
             original_weights.append(layer_copy)
 
         expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
+        default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl"
         communicator = create_eplb_communicator_or_raise(
             group_coordinator=ep_group_coordinator,
-            backend="torch_nccl",
+            backend=default_backend,
             expert_weights=expert_weights,
             expert_buffer=expert_buffer,
         )
@@ -769,13 +836,13 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
     assert_verification_synced(
         local_ok,
         "Profile-mode EPLB verification failed on at least one rank.",
+        cpu_group=cpu_group,
     )
 
 
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_rearrange_expert_weights_profile_mode(world_size):
     """Test profile mode (should not copy actual weights)"""
-
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Need at least {world_size} GPUs to run the test")
     distributed_run(
diff --git a/tests/distributed/test_eplb_fused_moe_layer.py b/tests/distributed/test_eplb_fused_moe_layer.py
index d8d870bc98c8..c6ae9e60d6e5 100644
--- a/tests/distributed/test_eplb_fused_moe_layer.py
+++ b/tests/distributed/test_eplb_fused_moe_layer.py
@@ -17,6 +17,7 @@
     get_tp_group,
 )
 from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner
+from vllm.platforms import current_platform
 
 from .eplb_utils import distributed_run, set_env_vars_and_device
 
@@ -81,7 +82,7 @@ def make_fused_moe_layer(
     )
     re = fml.routed_experts
 
-    device = torch.device(f"cuda:{rank}")
+    device = torch.device(f"{torch.accelerator.current_accelerator().type}:{rank}")
 
     from functools import partial
 
@@ -190,13 +191,14 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig):
     vllm_config = VllmConfig()
     vllm_config.parallel_config.tensor_parallel_size = world_size
     vllm_config.parallel_config.enable_expert_parallel = True
+    vllm_config.parallel_config.enable_eplb = True
 
     with set_current_vllm_config(vllm_config):
         ensure_model_parallel_initialized(
             tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
         )
 
-        ep_group = get_tp_group().cpu_group
+        ep_group = get_tp_group().device_group
         ep_rank = torch.distributed.get_rank()
 
         fml_layers = [
@@ -216,9 +218,10 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig):
             shuffled_indices[lidx] = torch.randperm(test_config.num_experts)
 
         expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]]
+        default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl"
         communicator = create_eplb_communicator(
             group_coordinator=get_eplb_group(),
-            backend="torch_nccl",
+            backend=default_backend,
             expert_weights=rank_expert_weights,
             expert_buffer=expert_buffer,
         )
diff --git a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py
index fe11e13cd60b..979742d5014e 100644
--- a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py
+++ b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py
@@ -24,6 +24,7 @@
     ModelOptNvFp4Config,
     ModelOptNvFp4FusedMoE,
 )
+from vllm.platforms import current_platform
 
 from .eplb_utils import distributed_run, set_env_vars_and_device
 
@@ -270,6 +271,8 @@ def test_eplb_fml(
 ):
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Need at least {world_size} GPUs to run the test")
+    if current_platform.is_xpu():
+        pytest.skip("NVFP4 quantization and flashinfer backends are CUDA-only")
 
     num_local_experts = num_experts // world_size
     num_topk = 4
diff --git a/tests/distributed/test_eplb_spec_decode.py b/tests/distributed/test_eplb_spec_decode.py
index b4211c74c171..ced206a3056a 100644
--- a/tests/distributed/test_eplb_spec_decode.py
+++ b/tests/distributed/test_eplb_spec_decode.py
@@ -75,8 +75,7 @@ def test_eplb_spec_decode(
     monkeypatch: pytest.MonkeyPatch,
     model_setup: tuple[str, str, str, int, float],
 ):
-    """
-    Test the correctness of EPLB speculative decoding with GSM8K dataset.
+    """Test the correctness of EPLB speculative decoding with GSM8K dataset.
     Applicable to MoE models with mtp or eagle spec decode.
     """
     method, model_name, spec_model_name, tp_size, expected_gsm8k_value = model_setup
@@ -109,10 +108,7 @@ def test_eplb_spec_decode(
 
 @large_gpu_mark(min_gb=80)
 def test_eplb_spec_decode_qwen3_next_mtp_async() -> None:
-    """
-    Ensure async EPLB works with MTP speculative decoding for Qwen3-Next.
-    """
-
+    """Ensure async EPLB works with MTP speculative decoding for Qwen3-Next."""
     TASK = "gsm8k"
     FILTER = "exact_match,strict-match"
     RTOL = 0.03
diff --git a/tests/distributed/test_eplb_utils.py b/tests/distributed/test_eplb_utils.py
index 53a4ce21af2e..dd2bb2760b1e 100644
--- a/tests/distributed/test_eplb_utils.py
+++ b/tests/distributed/test_eplb_utils.py
@@ -16,26 +16,31 @@ def _make_model_state(
     phy2log: torch.Tensor,
     log2phy: torch.Tensor,
     logcnt: torch.Tensor,
+    phy2log_storage: torch.Tensor | None = None,
 ) -> MagicMock:
-    """Build a minimal EplbModelState mock with only the three map tensors."""
+    """Build a minimal EplbModelState mock with its map tensors."""
     state = MagicMock()
     state.physical_to_logical_map = phy2log
+    state.physical_to_logical_map_buffer = (
+        phy2log if phy2log_storage is None else phy2log_storage
+    )
     state.logical_to_physical_map = log2phy
     state.logical_replica_count = logcnt
     return state
 
 
 def test_commit_eplb_maps_shape_change():
-    """
-    The normal path copies the physical_to_logical map in-place. When the number of
-    physical experts changes, the old map should be replaced entirely.
+    """When the number of physical experts changes, resize the active map within its
+    preallocated storage.
     """
     num_layers, num_logical, num_physical = 2, 4, 6
     max_replicas = 3
 
     # Build current state tensors
+    storage = torch.full((num_layers, num_physical + 2), -1, dtype=torch.long)
     model_state = _make_model_state(
-        phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
+        phy2log=storage[:, :num_physical],
+        phy2log_storage=storage,
         log2phy=torch.full(
             (num_layers, num_logical, max_replicas), -1, dtype=torch.long
         ),
@@ -54,12 +59,12 @@ def test_commit_eplb_maps_shape_change():
     # Check that the number of physical experts has been updated and that the values
     # match
     assert model_state.physical_to_logical_map.shape[1] == num_physical + 2
+    assert model_state.physical_to_logical_map.data_ptr() == storage.data_ptr()
     assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger)
 
 
 def test_commit_eplb_maps_for_layer_logical_padding():
-    """
-    Test that logical_to_physical_map is padded with -1 to fill the
+    """Test that logical_to_physical_map is padded with -1 to fill the
     pre-allocated slots when the new map has fewer replicas than the max.
     """
     num_layers, num_logical, num_physical = 2, 4, 6
diff --git a/tests/distributed/test_events.py b/tests/distributed/test_events.py
index 894f3f7f8b83..dee0b99bb0ec 100644
--- a/tests/distributed/test_events.py
+++ b/tests/distributed/test_events.py
@@ -20,27 +20,26 @@ class EventSample(
     tag=True,  # type: ignore
     array_like=True,  # type: ignore
 ):
-    """Test event for publisher testing"""
+    """Test event for publisher testing."""
 
     id: int
     value: str
 
 
 class SampleBatch(EventBatch):
-    """Test event batch for publisher testing"""
+    """Test event batch for publisher testing."""
 
     events: list[EventSample]
 
 
 def create_test_events(count: int) -> SampleBatch:
-    """Create a batch of test events"""
+    """Create a batch of test events."""
     events = [EventSample(id=i, value=f"test-{i}") for i in range(count)]
     return SampleBatch(ts=time.time(), events=events)
 
 
 def test_basic_publishing(publisher, subscriber):
-    """Test basic event publishing works"""
-
+    """Test basic event publishing works."""
     test_batch = create_test_events(5)
     publisher.publish(test_batch)
 
@@ -58,7 +57,7 @@ def test_basic_publishing(publisher, subscriber):
 
 
 def test_multiple_events(publisher, subscriber):
-    """Test publishing and receiving multiple event batches"""
+    """Test publishing and receiving multiple event batches."""
     for _ in range(10):
         batch = create_test_events(2)
         publisher.publish(batch)
@@ -75,7 +74,7 @@ def test_multiple_events(publisher, subscriber):
 
 
 def test_replay_mechanism(publisher, subscriber):
-    """Test the replay mechanism works correctly"""
+    """Test the replay mechanism works correctly."""
     for _ in range(19):
         batch = create_test_events(1)
         publisher.publish(batch)
@@ -96,7 +95,7 @@ def test_replay_mechanism(publisher, subscriber):
 
 
 def test_replay_includes_topic(publisher, subscriber, publisher_config):
-    """Test that replay responses include the topic, matching PUB format"""
+    """Test that replay responses include the topic, matching PUB format."""
     for _ in range(5):
         publisher.publish(create_test_events(1))
 
@@ -115,7 +114,7 @@ def test_replay_includes_topic(publisher, subscriber, publisher_config):
 
 
 def test_buffer_limit(publisher, subscriber, publisher_config):
-    """Test buffer limit behavior"""
+    """Test buffer limit behavior."""
     buffer_size = publisher_config.buffer_steps
 
     # Publish more events than the buffer can hold
@@ -139,9 +138,7 @@ def test_buffer_limit(publisher, subscriber, publisher_config):
 
 
 def test_topic_filtering(publisher_config):
-    """
-    Test that a subscriber only receives messages matching its topic filter
-    """
+    """Test that a subscriber only receives messages matching its topic filter."""
     publisher_config.replay_endpoint = None
 
     publisher_config.topic = "foo"
@@ -174,7 +171,7 @@ def test_topic_filtering(publisher_config):
 
 
 def test_high_volume(publisher, subscriber):
-    """Test publishing and receiving a high volume of events"""
+    """Test publishing and receiving a high volume of events."""
     num_batches = 10_000
     events_per_batch = 100
 
@@ -210,7 +207,7 @@ def publish_events():
 
 
 def test_null_publisher():
-    """Test that NullEventPublisher can be used without errors"""
+    """Test that NullEventPublisher can be used without errors."""
     publisher = NullEventPublisher(DP_RANK)
 
     # This should not raise any errors
@@ -220,8 +217,7 @@ def test_null_publisher():
 
 
 def test_data_parallel_rank_tagging(publisher_config):
-    """Test that events are properly tagged with their data parallel rank"""
-
+    """Test that events are properly tagged with their data parallel rank."""
     publisher_config.topic = "foo"
     pub_0 = EventPublisherFactory.create(publisher_config, DP_RANK)
     pub_1 = EventPublisherFactory.create(publisher_config, DP_RANK + 1)
@@ -285,7 +281,7 @@ def test_data_parallel_rank_tagging(publisher_config):
 
 
 def test_event_publisher_factory(random_port):
-    """Test event publisher factory creation behavior under different configurations"""
+    """Test event publisher factory creation behavior under different configurations."""
     from vllm.config.kv_events import KVEventsConfig
     from vllm.distributed.kv_events import ZmqEventPublisher
 
@@ -336,3 +332,77 @@ def test_event_publisher_factory(random_port):
     publisher = EventPublisherFactory.create(config, DP_RANK)
     assert isinstance(publisher, ZmqEventPublisher)
     publisher.shutdown()
+
+
+@pytest.mark.parametrize(
+    ("endpoint", "expected"),
+    [("tcp://*:5557", "tcp://*:5558"), ("tcp://*:0", "tcp://*:0")],
+)
+def test_offset_endpoint_port_keeps_port_zero(endpoint, expected):
+    """Port 0 is not offset: every rank asks the OS for its own port."""
+    from vllm.distributed.kv_events import ZmqEventPublisher
+
+    assert ZmqEventPublisher.offset_endpoint_port(endpoint, 1) == expected
+
+
+@pytest.mark.parametrize("node_ip", ["127.0.0.1", "::1"])
+def test_ephemeral_publishers_report_bound_endpoints(monkeypatch, node_ip):
+    """Two ranks binding tcp://*:0 get distinct OS-assigned ports on this
+    node's IP and report endpoints a subscriber can use for events and replay."""
+    from vllm.config.kv_events import KVEventsConfig
+    from vllm.distributed import kv_events
+    from vllm.utils.network_utils import split_zmq_path
+
+    from .conftest import MockSubscriber
+
+    monkeypatch.setattr(kv_events, "get_ip", lambda: node_ip)
+    config = KVEventsConfig(
+        enable_kv_cache_events=True,
+        publisher="zmq",
+        endpoint="tcp://*:0",
+        replay_endpoint="tcp://*:0",
+        topic="kv",
+    )
+    pubs = [EventPublisherFactory.create(config, rank) for rank in range(2)]
+    try:
+        resolved = [pub.get_publisher_config() for pub in pubs]
+        endpoints = [e for c in resolved for e in (c.endpoint, c.replay_endpoint)]
+        hosts_ports = {split_zmq_path(e)[1:] for e in endpoints}
+        assert {h for h, _ in hosts_ports} == {node_ip}
+        assert len(hosts_ports) == 4 and "0" not in {p for _, p in hosts_ports}
+        assert config.endpoint == "tcp://*:0"
+
+        sub = MockSubscriber(endpoints[::2], endpoints[1::2], topic="kv")
+        try:
+            # PUB/SUB drops messages sent before the subscription lands, so
+            # keep publishing until both ranks are heard.
+            ranks: set[int] = set()
+            deadline = time.time() + 10
+            while ranks != {0, 1} and time.time() < deadline:
+                for pub in pubs:
+                    pub.publish(create_test_events(1))
+                while (result := sub.receive_one(timeout=100)) is not None:
+                    ranks.add(result[1].data_parallel_rank)
+            assert ranks == {0, 1}
+            sub.request_replay(0, socket_idx=1)
+            assert sub.receive_replay(socket_idx=1)
+        finally:
+            sub.close()
+    finally:
+        for pub in pubs:
+            pub.shutdown()
+
+
+@pytest.mark.parametrize("node_ip", ["0.0.0.0", "::"])
+def test_ephemeral_bind_needs_a_node_ip(monkeypatch, node_ip):
+    """tcp://*:0 is advertised as this node's IP; with no usable IP the
+    publisher refuses instead of reporting a wildcard."""
+    from vllm.config.kv_events import KVEventsConfig
+    from vllm.distributed import kv_events
+
+    monkeypatch.setattr(kv_events, "get_ip", lambda: node_ip)
+    config = KVEventsConfig(
+        enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://*:0"
+    )
+    with pytest.raises(ValueError, match="VLLM_HOST_IP"):
+        EventPublisherFactory.create(config, DP_RANK)
diff --git a/tests/distributed/test_expert_placement.py b/tests/distributed/test_expert_placement.py
index 46f63408f467..4174b5c6ccff 100644
--- a/tests/distributed/test_expert_placement.py
+++ b/tests/distributed/test_expert_placement.py
@@ -49,7 +49,6 @@ def verify_round_robin_pattern(expert_map, ep_rank, ep_size, global_num_experts)
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_expert_placement_various_sizes(expert_placement_strategy, world_size):
     """Test round_robin expert placement with various expert counts."""
-
     # Test with different global_num_experts values
     # Include both divisible and non-divisible cases
     if world_size == 2:
@@ -116,7 +115,6 @@ def test_expert_placement_various_sizes(expert_placement_strategy, world_size):
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_expert_placement_edge_cases(expert_placement_strategy, world_size):
     """Test edge cases for round_robin expert placement."""
-
     # Test case 1: ep_size = 1 (should return None for expert_map)
     local_num_experts, expert_map, _ = determine_expert_map(
         ep_size=1,
@@ -139,7 +137,6 @@ def test_expert_placement_edge_cases(expert_placement_strategy, world_size):
 
 def test_determine_expert_map_comprehensive():
     """Test of determine_expert_map function with various configurations."""
-
     # Test cases: (ep_size, ep_rank, global_num_experts,
     # expert_placement_strategy, expected_local, expected_map_pattern)
     test_cases = [
diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py
index 636f7cb525f5..362ffe4b1535 100644
--- a/tests/distributed/test_mnnvl_alltoall.py
+++ b/tests/distributed/test_mnnvl_alltoall.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for MNNVL AllToAll operations.
+"""Tests for MNNVL AllToAll operations.
 
 Requires: docker run ... --cap-add=SYS_PTRACE ...
 Run: pytest tests/distributed/test_mnnvl_alltoall.py -v
@@ -15,6 +14,7 @@
 import torch.multiprocessing as mp
 
 from vllm.distributed import get_ep_group
+from vllm.platforms import current_platform
 from vllm.utils.flashinfer import (
     has_flashinfer_nvlink_one_sided,
     has_flashinfer_nvlink_two_sided,
@@ -24,6 +24,8 @@
 
 from ..utils import init_test_distributed_environment
 
+DEVICE = current_platform.device_type
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -87,6 +89,7 @@ def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue):
                  Otherwise use tp=world_size (default for EP-based tests).
         dp_port: Separate port for the DP master (only used when dp_size is set).
         err_queue: Queue for propagating tracebacks to the parent process.
+
     """
     try:
         os.environ.pop("CUDA_VISIBLE_DEVICES", None)
@@ -115,6 +118,7 @@ def _init_dp_environment(world_size, rank, port, dp_size, dp_port):
     Args:
         port: Port for torch.distributed init.
         dp_port: Separate port for the DP master group init.
+
     """
     from vllm.config import VllmConfig, set_current_vllm_config
     from vllm.config.parallel import ParallelConfig
@@ -482,7 +486,7 @@ def _args_dispatch_combine_worker(rank, world_size):
     from vllm.forward_context import get_forward_context
 
     cpu_group = get_ep_group().cpu_group
-    device = torch.device(f"cuda:{rank}")
+    device = torch.device(f"{DEVICE}:{rank}")
 
     hidden_size = 64
     tokens_per_rank = 16
@@ -623,7 +627,7 @@ def _two_sided_data_worker(rank, world_size):
     # Use DP group because MnnvlMoe workspace allocation calls get_dp_group()
     # internally and requires dp_size == ep_size.
     cpu_group = get_dp_group().cpu_group
-    device = torch.device(f"cuda:{rank}")
+    device = torch.device(f"{DEVICE}:{rank}")
     num_gpus = torch.accelerator.device_count()
 
     hidden_size = 128
@@ -768,7 +772,7 @@ def _one_sided_data_worker(rank, world_size):
     from vllm.forward_context import get_forward_context
 
     cpu_group = get_dp_group().cpu_group
-    device = torch.device(f"cuda:{rank}")
+    device = torch.device(f"{DEVICE}:{rank}")
 
     hidden_size = 256
     tokens_per_rank = 32
diff --git a/tests/distributed/test_mq_connect_ip.py b/tests/distributed/test_mq_connect_ip.py
index 4b0cdda3ad9e..05617bbb6037 100644
--- a/tests/distributed/test_mq_connect_ip.py
+++ b/tests/distributed/test_mq_connect_ip.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test that MessageQueue uses the local node's IP for binding,
+"""Test that MessageQueue uses the local node's IP for binding,
 not a remote master_addr. This validates the fix for cross-node
 data-parallel where each DP group leader must bind to its own IP.
 
diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py
index 599025d1c8ec..cd9413865465 100644
--- a/tests/distributed/test_multiproc_executor.py
+++ b/tests/distributed/test_multiproc_executor.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Integration tests for MultiprocExecutor at the executor level.
+"""Integration tests for MultiprocExecutor at the executor level.
 This test directly tests the executor without going through the LLM interface,
 focusing on executor initialization, RPC calls, and distributed execution.
 """
@@ -16,6 +15,7 @@
 from tests.utils import multi_gpu_test
 from vllm.config import VllmConfig
 from vllm.engine.arg_utils import EngineArgs
+from vllm.platforms import current_platform
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.executor import multiproc_executor
 from vllm.v1.executor.multiproc_executor import MultiprocExecutor
@@ -377,11 +377,10 @@ def test_multiproc_executor_properties():
 
 @multi_gpu_test(num_gpus=4)
 def test_multiproc_executor_multi_node():
-    """
-    Test MultiprocExecutor with multi-node configuration.
+    """Test MultiprocExecutor with multi-node configuration.
     This simulates 2 nodes with TP=4:
-    - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2
-    - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
+    - Node 0 (rank 0): Uses GPUs 0,1 (device_control_env_var=0,1) with TP=2
+    - Node 1 (rank 1): Uses GPUs 2,3 (device_control_env_var=2,3) with TP=2
     Total world_size = 4, nnodes = 2
     """
     # Python 3.14+ changed default multiprocessing start method to 'forkserver'
@@ -394,15 +393,20 @@ def test_multiproc_executor_multi_node():
     # symm_mem does not work for simulating multi instance in single node
     os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
 
+    # Use the platform-specific device control env var (e.g.
+    # CUDA_VISIBLE_DEVICES on CUDA, ZE_AFFINITY_MASK on XPU) so that each
+    # node only sees its own devices on different platform.
+    device_control_env_var = current_platform.device_control_env_var
+
     def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int):
         """Run a single node's executor."""
         executor = None
         try:
-            # Set CUDA_VISIBLE_DEVICES for this node
+            # Restrict visible devices for this node
             if node_rank == 0:
-                os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+                os.environ[device_control_env_var] = "0,1"
             else:
-                os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
+                os.environ[device_control_env_var] = "2,3"
 
             # Create config for this node
             vllm_config = create_vllm_config(
diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py
index a464a61683b5..8ba93df4f504 100644
--- a/tests/distributed/test_pipeline_parallel.py
+++ b/tests/distributed/test_pipeline_parallel.py
@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-WARNING: This test runs in both single-node (4 GPUs) and multi-node
- (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
- important to set the distributed backend to "mp" to avoid Ray scheduling
- all workers in a node other than the head node, which can cause the test
- to fail.
+"""WARNING: This test runs in both single-node (4 GPUs) and multi-node
+(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
+important to set the distributed backend to "mp" to avoid Ray scheduling
+all workers in a node other than the head node, which can cause the test
+to fail.
 """
 
 import json
diff --git a/tests/distributed/test_quick_all_reduce.py b/tests/distributed/test_quick_all_reduce.py
index fb2b04abf0d5..107d95fa613a 100644
--- a/tests/distributed/test_quick_all_reduce.py
+++ b/tests/distributed/test_quick_all_reduce.py
@@ -3,12 +3,17 @@
 
 import multiprocessing
 import random
+import sys
+import types
+from types import SimpleNamespace
 
 import pytest
 import ray
 import torch
 import torch.distributed as dist
 
+import vllm._aiter_ops as aiter_ops
+import vllm.distributed as vllm_distributed
 import vllm.envs as envs
 from vllm import _custom_ops as ops
 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce  # noqa
@@ -233,6 +238,173 @@ def fake_qr_all_reduce(
     assert called_quant_level == QuickReduceRegime.FP.value
 
 
+def test_rocm_aiter_fused_rmsnorm_uses_aiter_qr_rmsnorm_for_prefill(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    class FakeAiterAllReduce:
+        aiter_ca = SimpleNamespace(world_size=2, fully_connected=True)
+
+        def use_1stage_fused_ar_rms(self, inp):
+            return False
+
+        def custom_fused_ar_rms(self, *args, **kwargs):
+            raise AssertionError("expected fused QR+RMSNorm dispatch")
+
+    class FakeQuickReduce:
+        disabled = False
+
+        def __init__(self):
+            self.checked_input = None
+
+        def should_quick_allreduce(self, inp):
+            self.checked_input = inp
+            return True
+
+        def _get_qr_quant_level(self, inp):
+            return QuickReduceRegime.INT4.value
+
+    qr_comm = FakeQuickReduce()
+    device_comm = SimpleNamespace(cpu_group=object(), qr_comm=qr_comm)
+    monkeypatch.setattr(
+        vllm_distributed,
+        "get_tp_group",
+        lambda: SimpleNamespace(device_communicator=device_comm),
+    )
+    monkeypatch.setattr(
+        aiter_ops.rocm_aiter_ops,
+        "get_aiter_allreduce",
+        lambda: FakeAiterAllReduce(),
+    )
+    monkeypatch.setattr(aiter_ops.dist, "get_world_size", lambda group: 2)
+    monkeypatch.setattr(aiter_ops.dist, "get_rank", lambda group: 0)
+    monkeypatch.setattr(
+        aiter_ops.dist,
+        "all_gather_object",
+        lambda handles, handle, group: handles.__setitem__(0, handle),
+    )
+
+    calls = []
+    fake_aiter = types.SimpleNamespace(
+        init_custom_qr=lambda rank, world_size, qr_max_size: 123,
+        qr_get_handle=lambda ptr: b"handle",
+        qr_open_handles=lambda ptr, handles: None,
+    )
+
+    def fake_qr_all_reduce_rmsnorm(
+        ptr,
+        inp,
+        residual,
+        residual_out,
+        out,
+        weight,
+        eps,
+        hidden_dim,
+        quant_level,
+        cast_bf2half,
+    ):
+        calls.append(
+            {
+                "ptr": ptr,
+                "inp": inp,
+                "residual": residual,
+                "residual_out": residual_out,
+                "out": out,
+                "weight": weight,
+                "eps": eps,
+                "hidden_dim": hidden_dim,
+                "quant_level": quant_level,
+                "cast_bf2half": cast_bf2half,
+            }
+        )
+        out.copy_(inp)
+        residual_out.copy_(residual)
+
+    fake_aiter.qr_all_reduce_rmsnorm = fake_qr_all_reduce_rmsnorm
+    monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
+    monkeypatch.setattr(envs, "VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None)
+    monkeypatch.setattr(envs, "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", True)
+
+    # token_num > 80 forces vLLM's AITER fused AR+RMS wrapper onto its
+    # non-1stage prefill path. hidden_dim=4096 bf16 gives an 8 KiB row, so four
+    # independent RMSNorm rows fit in one 32 KiB QR tile.
+    inp = torch.randn(81, 4096, dtype=torch.bfloat16)
+    residual = torch.randn_like(inp)
+    weight = torch.randn(4096, dtype=torch.bfloat16)
+
+    out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl(
+        inp, residual, weight, 1e-6
+    )
+
+    assert qr_comm.checked_input is inp
+    assert len(calls) == 1
+    assert calls[0]["ptr"] == 123
+    assert calls[0]["inp"] is inp
+    assert calls[0]["residual"] is residual
+    assert calls[0]["out"] is out
+    assert calls[0]["residual_out"] is residual_out
+    assert calls[0]["weight"] is weight
+    assert calls[0]["hidden_dim"] == 4096
+    assert calls[0]["quant_level"] == QuickReduceRegime.INT4.value
+    assert calls[0]["cast_bf2half"] is True
+    torch.testing.assert_close(out, inp)
+    torch.testing.assert_close(residual_out, residual)
+
+
+def test_rocm_aiter_fused_rmsnorm_keeps_1stage_decode_path(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    calls = []
+
+    class FakeAiterCA:
+        world_size = 2
+        fully_connected = True
+
+        def custom_fused_ar_rms(
+            self, inp, residual, weight, epsilon, *, use_1stage, gemma_norm
+        ):
+            calls.append(
+                {
+                    "inp": inp,
+                    "residual": residual,
+                    "weight": weight,
+                    "eps": epsilon,
+                    "use_1stage": use_1stage,
+                    "gemma_norm": gemma_norm,
+                }
+            )
+            return inp + 1, residual + 1
+
+    class FakeAiterAllReduce:
+        aiter_ca = FakeAiterCA()
+
+        def use_1stage_fused_ar_rms(self, inp):
+            return True
+
+    monkeypatch.setattr(
+        aiter_ops.rocm_aiter_ops,
+        "get_aiter_allreduce",
+        lambda: FakeAiterAllReduce(),
+    )
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+
+    inp = torch.randn(4, 4096, dtype=torch.bfloat16)
+    residual = torch.randn_like(inp)
+    weight = torch.randn(4096, dtype=torch.bfloat16)
+
+    out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl(
+        inp, residual, weight, 1e-6
+    )
+
+    assert len(calls) == 1
+    assert calls[0]["inp"] is inp
+    assert calls[0]["residual"] is residual
+    assert calls[0]["weight"] is weight
+    assert calls[0]["use_1stage"] is True
+    assert calls[0]["gemma_norm"] is False
+    torch.testing.assert_close(out, inp + 1)
+    torch.testing.assert_close(residual_out, residual + 1)
+
+
 @ray.remote(num_gpus=1, max_calls=1)
 def graph_quickreduce(
     monkeypatch: pytest.MonkeyPatch,
@@ -396,8 +568,7 @@ def test_custom_quick_allreduce_bf16_cast(monkeypatch: pytest.MonkeyPatch):
 
 
 def qr_variable_input(rank, world_size):
-    """
-    When the tensor parallelism is set to 4 or 8, frequent changes
+    """When the tensor parallelism is set to 4 or 8, frequent changes
     in the input shape can cause QuickReduce to hang (this issue
     has been observed with the gpt_oss model).
     """
diff --git a/tests/distributed/test_ray_v2_executor.py b/tests/distributed/test_ray_v2_executor.py
index bd4a253b96a4..f9b9e57fc54c 100644
--- a/tests/distributed/test_ray_v2_executor.py
+++ b/tests/distributed/test_ray_v2_executor.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Integration tests for RayExecutorV2 at the executor level.
+"""Integration tests for RayExecutorV2 at the executor level.
 Validates executor initialization, placement group support, RPC calls,
 and distributed execution with various TP/PP configurations.
 """
@@ -299,7 +298,6 @@ def test_ray_v2_run_refs_stored_for_monitoring(executor):
 @pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)])
 def test_ray_v2_single_node_generation(tp_size, pp_size):
     """End-to-end LLM generation with RayExecutorV2."""
-
     llm = LLM(
         model=MODEL,
         tensor_parallel_size=tp_size,
diff --git a/tests/distributed/test_ray_v2_executor_e2e.py b/tests/distributed/test_ray_v2_executor_e2e.py
index fb5830132698..335324c37128 100644
--- a/tests/distributed/test_ray_v2_executor_e2e.py
+++ b/tests/distributed/test_ray_v2_executor_e2e.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Orchestration-level integration tests for RayExecutorV2.
-"""
+"""Orchestration-level integration tests for RayExecutorV2."""
 
 import gc
 import os
@@ -158,8 +156,7 @@ def test_multi_replicas_with_bundle_indices(ray_init):
 
 
 def test_env_var_and_runtime_env_propagation():
-    """
-    Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env
+    """Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env
     propagate to RayWorkerProc actors.
     """
     sentinel_vars = {
diff --git a/tests/distributed/test_rocm_aiter_custom_ar.py b/tests/distributed/test_rocm_aiter_custom_ar.py
index 0b85f36410d7..65018bfc4e5d 100644
--- a/tests/distributed/test_rocm_aiter_custom_ar.py
+++ b/tests/distributed/test_rocm_aiter_custom_ar.py
@@ -8,7 +8,7 @@
 
 from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops
 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce  # noqa
-from vllm.distributed.parallel_state import get_tp_group, graph_capture
+from vllm.distributed.parallel_state import get_dp_group, get_tp_group, graph_capture
 from vllm.envs import disable_envs_cache
 from vllm.platforms import current_platform
 
@@ -132,3 +132,173 @@ def test_rocm_aiter_custom_allreduce(
     test_target,
 ):
     multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
+
+
+def _get_aiter_ag_rs_comm():
+    device_communicator = get_dp_group().device_communicator
+    assert device_communicator.use_aiter_ag_rs, (
+        "AITER custom AG/RS was not enabled on the DP group."
+    )
+    aiter_comm = device_communicator.aiter_ar_comm
+    assert aiter_comm is not None, "AITER custom AG/RS was not initialized."
+    assert not aiter_comm.disabled, "AITER custom AG/RS is disabled."
+    return aiter_comm
+
+
+def _assert_aiter_handles_ag(aiter_comm, inp: torch.Tensor) -> None:
+    assert aiter_comm.should_custom_ag(inp), (
+        f"AITER custom all-gather does not support input shape {inp.shape}."
+    )
+
+
+def _assert_aiter_handles_rs(aiter_comm, inp: torch.Tensor) -> None:
+    assert aiter_comm.should_custom_rs(inp, dim=0), (
+        f"AITER custom reduce-scatter does not support input shape {inp.shape}."
+    )
+
+
+@ray.remote(num_gpus=1, max_calls=1)
+def eager_ag_rs(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size,
+    pp_size,
+    rank,
+    distributed_init_port,
+    data_parallel_size,
+    data_parallel_master_port,
+) -> None:
+    with monkeypatch.context() as m:
+        _configure_aiter_custom_ar_env(m)
+
+        device = torch.device(f"cuda:{rank}")
+        torch.accelerator.set_device_index(device)
+        init_test_distributed_environment(
+            tp_size,
+            pp_size,
+            rank,
+            distributed_init_port,
+            data_parallel_size=data_parallel_size,
+            data_parallel_master_port=data_parallel_master_port,
+        )
+
+        dp_group = get_dp_group()
+        group = dp_group.device_group
+        dp_world = dp_group.world_size
+        aiter_comm = _get_aiter_ag_rs_comm()
+
+        for shape, dtype in test_cases:
+            num_tokens, hidden = shape
+
+            # all-gather: each rank contributes (num_tokens, hidden).
+            inp = torch.ones(shape, dtype=dtype, device=device) * (rank + 1)
+            _assert_aiter_handles_ag(aiter_comm, inp)
+            expected = torch.empty(
+                (num_tokens * dp_world, hidden), dtype=dtype, device=device
+            )
+            dist.all_gather_into_tensor(expected, inp, group=group)
+            out = aiter_comm.custom_all_gather(inp, dim=0)
+            assert out is not None
+            torch.testing.assert_close(out, expected)
+
+            # reduce-scatter: each rank contributes (num_tokens * dp, hidden).
+            rs_in = torch.ones(
+                (num_tokens * dp_world, hidden), dtype=dtype, device=device
+            ) * (rank + 1)
+            _assert_aiter_handles_rs(aiter_comm, rs_in)
+            rs_expected = torch.empty((num_tokens, hidden), dtype=dtype, device=device)
+            dist.reduce_scatter_tensor(rs_expected, rs_in, group=group)
+            rs_out = torch.empty((num_tokens, hidden), dtype=dtype, device=device)
+            aiter_comm.custom_reduce_scatter(rs_in, rs_out, dim=0)
+            torch.testing.assert_close(rs_out, rs_expected)
+
+
+@ray.remote(num_gpus=1, max_calls=1)
+def graph_ag_rs(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size,
+    pp_size,
+    rank,
+    distributed_init_port,
+    data_parallel_size,
+    data_parallel_master_port,
+) -> None:
+    with monkeypatch.context() as m:
+        _configure_aiter_custom_ar_env(m)
+
+        device = torch.device(f"cuda:{rank}")
+        torch.accelerator.set_device_index(device)
+        init_test_distributed_environment(
+            tp_size,
+            pp_size,
+            rank,
+            distributed_init_port,
+            data_parallel_size=data_parallel_size,
+            data_parallel_master_port=data_parallel_master_port,
+        )
+
+        dp_group = get_dp_group()
+        group = dp_group.device_group
+        dp_world = dp_group.world_size
+        aiter_comm = _get_aiter_ag_rs_comm()
+
+        # Warmup so DP comms is initialized before graph capture
+        data = torch.zeros(1, device=device)
+        dist.all_reduce(data, group=group)
+        torch.accelerator.synchronize()
+        del data
+
+        for shape, dtype in test_cases:
+            num_tokens, hidden = shape
+
+            # all-gather under graph capture.
+            inp = torch.ones(shape, dtype=dtype, device=device) * (rank + 1)
+            _assert_aiter_handles_ag(aiter_comm, inp)
+            ag_expected = torch.empty(
+                (num_tokens * dp_world, hidden), dtype=dtype, device=device
+            )
+            dist.all_gather_into_tensor(ag_expected, inp, group=group)
+            with graph_capture(device=device) as graph_capture_context:
+                torch.accelerator.synchronize()
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph, stream=graph_capture_context.stream):
+                    ag_out = aiter_comm.custom_all_gather(inp, dim=0)
+            graph.replay()
+            torch.testing.assert_close(ag_out, ag_expected)
+
+            # reduce-scatter under graph capture.
+            rs_in = torch.ones(
+                (num_tokens * dp_world, hidden), dtype=dtype, device=device
+            ) * (rank + 1)
+            _assert_aiter_handles_rs(aiter_comm, rs_in)
+            rs_expected = torch.empty((num_tokens, hidden), dtype=dtype, device=device)
+            dist.reduce_scatter_tensor(rs_expected, rs_in, group=group)
+            rs_out = torch.empty((num_tokens, hidden), dtype=dtype, device=device)
+            with graph_capture(device=device) as graph_capture_context:
+                torch.accelerator.synchronize()
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph, stream=graph_capture_context.stream):
+                    aiter_comm.custom_reduce_scatter(rs_in, rs_out, dim=0)
+            graph.replay()
+            torch.testing.assert_close(rs_out, rs_expected)
+
+
+@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed")
+@multi_gpu_test(num_gpus=2)
+@pytest.mark.parametrize("tp_size", [1])
+@pytest.mark.parametrize("pipeline_parallel_size", [1])
+@pytest.mark.parametrize("data_parallel_size", [2])
+@pytest.mark.parametrize("test_target", [eager_ag_rs, graph_ag_rs])
+def test_rocm_aiter_custom_ag_rs(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size,
+    pipeline_parallel_size,
+    data_parallel_size,
+    test_target,
+):
+    multi_process_parallel(
+        monkeypatch,
+        tp_size,
+        pipeline_parallel_size,
+        test_target,
+        data_parallel_size=data_parallel_size,
+    )
diff --git a/tests/distributed/test_sharded_rdt_plan.py b/tests/distributed/test_sharded_rdt_plan.py
index 7a75f0b7fcd6..7923c413fc2a 100644
--- a/tests/distributed/test_sharded_rdt_plan.py
+++ b/tests/distributed/test_sharded_rdt_plan.py
@@ -845,7 +845,7 @@ def test_stamps_interleave_across_modules_without_reordering_within(self):
         ]
 
     def test_scatters_carry_their_own_dtype(self):
-        """dtype rides the record from the bake, where it is the fake's dtype
+        """Dtype rides the record from the bake, where it is the fake's dtype
         AFTER its op chain — not a plan-time lookup of the source name, which
         would be wrong for any chain that reinterprets dtype."""
         layer = _FakeLayer("l")
@@ -1110,7 +1110,7 @@ class TestSignalCompleteness:
     signaled parks a producer credit and hangs end_sync."""
 
     def _moe_planner(self, worker_experts, *, ep_size=4, n_experts=8, n_layers=2):
-        """pre / n_layers MoE layers / post. Each layer: one norm (its own
+        """Pre / n_layers MoE layers / post. Each layer: one norm (its own
         module, stamp -1) + n_experts expert names stamped ``e // n_local``.
         The worker's bake covers only ``worker_experts`` (its placement);
         foreign expert names never copied => dropped entirely, as in the real
@@ -1173,7 +1173,7 @@ def test_round_robin_placement_signals_every_group_exactly_once(self):
         assert self._signals(plan) == list(range(len(group_lens)))
 
     def test_placement_changes_only_the_chunk_count_never_the_signals(self):
-        """linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6
+        """Linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6
         hits all 4. More chunks per group, identical signal set."""
         plans = {}
         for label, experts in (("linear", [0, 1, 2, 3]), ("round_robin", [0, 2, 4, 6])):
diff --git a/tests/distributed/test_sharded_rdt_producer.py b/tests/distributed/test_sharded_rdt_producer.py
index 62fc63876262..8147841d9aa3 100644
--- a/tests/distributed/test_sharded_rdt_producer.py
+++ b/tests/distributed/test_sharded_rdt_producer.py
@@ -20,17 +20,23 @@
 
 import contextlib
 import gc
+import os
+import sys
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import TimeoutError as FuturesTimeout
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
 
 import vllm.distributed.weight_transfer.sharded_rdt_trainer as trainer_mod
+from vllm.distributed.nixl_utils import alias_nixl_for_ray
 from vllm.distributed.weight_transfer.sharded_rdt_common import (
     buffer_alloc_bytes,
+    register_nixl_memory,
 )
 from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
     DEFAULT_GATHER_LOOKAHEAD,
@@ -45,6 +51,114 @@
 GROUP_B = ("model.layers.1.w",)
 
 
+@pytest.fixture
+def ray_nixl_transport(monkeypatch):
+    import ray
+    from ray.experimental.rdt import util
+    from ray.experimental.rdt.nixl_tensor_transport import NixlTensorTransport
+
+    from vllm.distributed import nixl_utils
+    from vllm.platforms import current_platform
+
+    transport = NixlTensorTransport()
+    package = ModuleType("nixl_rocm")
+    api = ModuleType("nixl_rocm._api")
+    api.__dict__.update(nixl_agent=Mock(), nixl_agent_config=Mock())
+    package.__dict__["_api"] = api
+    modules = {package.__name__: package, api.__name__: api}
+    imported_rcache_limits = []
+
+    def import_nixl(name):
+        imported_rcache_limits.append(os.environ.get("UCX_RCACHE_MAX_UNRELEASED"))
+        monkeypatch.setitem(sys.modules, name, modules[name])
+        return modules[name]
+
+    nixl_names = ("nixl", "nixl._api", *modules)
+    for name in nixl_names:
+        monkeypatch.delitem(sys.modules, name, raising=False)
+    monkeypatch.delenv("UCX_RCACHE_MAX_UNRELEASED", raising=False)
+    importer = Mock(side_effect=import_nixl)
+    monkeypatch.setattr(
+        nixl_utils, "importlib", SimpleNamespace(import_module=importer)
+    )
+    monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: True)
+    monkeypatch.setattr(util, "transport_managers", {"NIXL": transport})
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr(
+        ray,
+        "get_runtime_context",
+        lambda: SimpleNamespace(get_actor_id=lambda: "rdt-actor"),
+    )
+    try:
+        yield transport, api, importer, imported_rcache_limits
+    finally:
+        # The helper inserts aliases directly; remove them before monkeypatch
+        # restores any modules and environment values present before this test.
+        for name in nixl_names:
+            sys.modules.pop(name, None)
+        os.environ.pop("UCX_RCACHE_MAX_UNRELEASED", None)
+
+
+def test_ray_registers_memory_through_rocm_alias(ray_nixl_transport):
+    """Ray's registration must reach the ROCm agent through the alias."""
+    _, api, importer, _ = ray_nixl_transport
+    register_nixl_memory(torch.empty(4, dtype=torch.uint8))
+    api.nixl_agent.return_value.register_memory.assert_called()
+    assert sys.modules["nixl"] is sys.modules["nixl_rocm"]
+    assert sys.modules["nixl._api"] is api
+
+    importer.reset_mock()
+    alias_nixl_for_ray()
+    importer.assert_not_called()
+
+
+@pytest.mark.parametrize("rcache_limit", [None, "4096"])
+def test_ray_nixl_sets_ucx_limit_before_import(
+    monkeypatch, ray_nixl_transport, rcache_limit
+):
+    """UCX must see the default or the user's override during the first import."""
+    _, _, _, imported_rcache_limits = ray_nixl_transport
+    if rcache_limit is not None:
+        monkeypatch.setenv("UCX_RCACHE_MAX_UNRELEASED", rcache_limit)
+    alias_nixl_for_ray()
+    assert imported_rcache_limits == [rcache_limit or "1024"] * 2
+
+
+@pytest.mark.parametrize("existing_api", [False, True])
+def test_ray_nixl_preserves_non_rocm_or_existing_imports(
+    monkeypatch, ray_nixl_transport, existing_api
+):
+    from vllm.platforms import current_platform
+
+    _, _, importer, _ = ray_nixl_transport
+    if existing_api:
+        api = ModuleType("nixl._api")
+        monkeypatch.setitem(sys.modules, "nixl._api", api)
+    else:
+        monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
+    alias_nixl_for_ray()
+    importer.assert_not_called()
+    assert "UCX_RCACHE_MAX_UNRELEASED" not in os.environ
+    if existing_api:
+        assert sys.modules["nixl._api"] is api
+    else:
+        assert "nixl._api" not in sys.modules
+
+
+def test_ray_nixl_alias_reports_missing_rocm_package(monkeypatch, ray_nixl_transport):
+    from vllm.distributed import nixl_utils
+
+    transport, api, importer, _ = ray_nixl_transport
+    monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: False)
+    with pytest.raises(ImportError, match="nixl_rocm"):
+        register_nixl_memory(torch.empty(1))
+    assert transport._nixl_agent is None
+    api.nixl_agent.assert_not_called()
+    importer.assert_not_called()
+    assert "nixl" not in sys.modules
+    assert "nixl._api" not in sys.modules
+
+
 @pytest.fixture
 def server_factory(monkeypatch):
     """Build real `_RDTProducerServer` instances without a GPU.
@@ -389,7 +503,7 @@ def _run():
     def test_the_residency_bound_scales_with_the_lookahead(
         self, server_factory, gather_engine
     ):
-        """lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing
+        """Lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing
         consumer sees exactly 3 groups gathered before the loop parks."""
         server = server_factory(gather_lookahead=2)
         engine = gather_engine(server, 8, lookahead=2)
diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py
index c3833b6b9f82..7debed4cbbff 100644
--- a/tests/distributed/test_shm_broadcast.py
+++ b/tests/distributed/test_shm_broadcast.py
@@ -43,6 +43,7 @@ def distributed_run(fn, world_size, timeout=60):
         fn: Function to run in each process
         world_size: Number of processes to spawn
         timeout: Maximum time in seconds to wait for processes (default: 60)
+
     """
     number_of_processes = world_size
     processes = []
@@ -634,11 +635,9 @@ def test_acquire_read_releases_slot_when_reader_raises():
 
 
 def test_warning_logs(caplog_vllm):
-    """
-    Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals
+    """Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals
     when indefinite=False, and are not emitted when indefinite=True.
     """
-
     # Patch the warning log interval to every 1 ms during reads
     with mock.patch(
         "vllm.distributed.device_communicators.shm_broadcast.VLLM_RINGBUFFER_WARNING_INTERVAL",
@@ -694,14 +693,43 @@ def test_check_shm_free_space_raises_when_insufficient(tmp_path):
 
 
 def test_check_shm_free_space_passes_when_sufficient(tmp_path):
-    with mock.patch.object(
-        shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(512 << 20)
+    with (
+        mock.patch.object(
+            shm_broadcast.shutil,
+            "disk_usage",
+            return_value=_fake_disk_usage(512 << 20),
+        ),
+        mock.patch.object(shm_broadcast, "check_cgroup_memory_available"),
     ):
         check_shm_free_space(240 << 20, shm_path=str(tmp_path))
 
 
 def test_check_shm_free_space_skipped_when_path_missing(tmp_path):
-    check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist"))
+    with mock.patch.object(shm_broadcast, "check_cgroup_memory_available"):
+        check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist"))
+
+
+def test_check_shm_free_space_checks_cgroup(tmp_path):
+    with (
+        mock.patch.object(
+            shm_broadcast.shutil,
+            "disk_usage",
+            return_value=_fake_disk_usage(512 << 20),
+        ),
+        mock.patch.object(
+            shm_broadcast, "check_cgroup_memory_available"
+        ) as check_cgroup,
+    ):
+        check_shm_free_space(
+            240 << 20,
+            shm_path=str(tmp_path),
+            allocation_name="SHM mmap",
+        )
+
+    check_cgroup.assert_called_once_with(
+        240 << 20,
+        "SHM mmap",
+    )
 
 
 def test_shm_ring_buffer_creation_checks_free_space():
diff --git a/tests/distributed/test_shm_buffer.py b/tests/distributed/test_shm_buffer.py
index e56c71b1c241..6f61c28322f5 100644
--- a/tests/distributed/test_shm_buffer.py
+++ b/tests/distributed/test_shm_buffer.py
@@ -12,20 +12,20 @@
 
 
 class TestSingleWriterShmRingBuffer(unittest.TestCase):
-    """Test suite for the ring buffer implementation"""
+    """Test suite for the ring buffer implementation."""
 
     def setUp(self):
-        """Set up test fixtures"""
+        """Set up test fixtures."""
         self.buffer_size = 4096
         self.ring_buffer = None
 
     def tearDown(self):
-        """Clean up after tests"""
+        """Clean up after tests."""
         if self.ring_buffer:
             self.ring_buffer.close()
 
     def test_buffer_opening(self):
-        """Test opening an existing buffer"""
+        """Test opening an existing buffer."""
         # First create a buffer
         self.ring_buffer = SingleWriterShmRingBuffer(
             data_buffer_size=self.buffer_size, create=True
@@ -39,7 +39,7 @@ def test_buffer_opening(self):
         )
 
     def test_buffer_access(self):
-        """Test accessing allocated buffers"""
+        """Test accessing allocated buffers."""
         self.ring_buffer = SingleWriterShmRingBuffer(
             data_buffer_size=self.buffer_size, create=True
         )
@@ -61,7 +61,7 @@ def test_buffer_access(self):
         self.assertEqual(read_id, monotonic_id)
 
     def test_memory_error_on_full_buffer(self):
-        """Test that MemoryError is raised when buffer is full"""
+        """Test that MemoryError is raised when buffer is full."""
         small_buffer_size = 200
         self.ring_buffer = SingleWriterShmRingBuffer(
             data_buffer_size=small_buffer_size, create=True
@@ -76,7 +76,7 @@ def test_memory_error_on_full_buffer(self):
             self.ring_buffer.allocate_buf(1)  # Would exceed buffer capacity
 
     def test_allocation_and_free(self):
-        """Test allocation and freeing of buffers"""
+        """Test allocation and freeing of buffers."""
         small_buffer_size = 200
         self.ring_buffer = SingleWriterShmRingBuffer(
             data_buffer_size=small_buffer_size, create=True
@@ -96,7 +96,7 @@ def test_allocation_and_free(self):
             self.assertEqual(freed_ids[0], i)
 
     def test_clear_buffer(self):
-        """Test clearing the buffer"""
+        """Test clearing the buffer."""
         self.ring_buffer = SingleWriterShmRingBuffer(
             data_buffer_size=self.buffer_size, create=True
         )
@@ -180,7 +180,7 @@ def ring_allocate(allocate_size):
 
 
 def main():
-    """Main function demonstrating usage and running tests"""
+    """Main function demonstrating usage and running tests."""
     print("=== SingleWriterShmRingBuffer Test Suite ===\n")
 
     # Run unit tests
diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py
index f8948610226c..480c4f1b8a0d 100644
--- a/tests/engine/test_arg_utils.py
+++ b/tests/engine/test_arg_utils.py
@@ -9,7 +9,13 @@
 import pytest
 from pydantic import Field
 
-from vllm.config import AttentionConfig, CompilationConfig, ModelConfig, config
+from vllm.config import (
+    AttentionConfig,
+    CacheConfig,
+    CompilationConfig,
+    ModelConfig,
+    config,
+)
 from vllm.engine.arg_utils import (
     EngineArgs,
     _expand_json_human_readable_numbers,
@@ -53,24 +59,66 @@ def test_watermark_config_cli():
             "--model",
             "dummy",
             "--watermark-config",
-            '{"algorithm":"gumbel","key":42,"prf":"philox"}',
+            '{"algorithm":"dual_key_gumbel","key":42,"prf":"philox","alpha":0.25,'
+            '"allow_target_only_watermarking":true}',
         ]
     )
 
     config = EngineArgs.from_cli_args(args).create_watermark_config()
 
     assert config is not None
-    assert config.algorithm == "gumbel"
+    assert config.algorithm == "dual_key_gumbel"
     assert config.key == 42
+    assert config.alpha == 0.25
     assert config.context_width == 4
+    assert config.deduplicate_contexts == "single_turn"
+    assert config.deduplicate_contexts_max_history == 8192
     assert config.prf == "philox"
-    assert not config.supports_speculative_decoding
+    assert config.allow_target_only_watermarking
+
+    args = parser.parse_args(
+        [
+            "--model",
+            "dummy",
+            "--watermark-config",
+            '{"key":42,"deduplicate_contexts":"none",'
+            '"deduplicate_contexts_max_history":32}',
+        ]
+    )
+    config = EngineArgs.from_cli_args(args).create_watermark_config()
+
+    assert config is not None
+    assert config.deduplicate_contexts == "none"
+    assert config.deduplicate_contexts_max_history == 32
+
+
+@pytest.mark.parametrize(
+    "option",
+    ["--gpu-memory-utilization", "--device-memory-utilization"],
+)
+def test_memory_utilization_cli_aliases(option):
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    args = EngineArgs.from_cli_args(parser.parse_args([option, "0.8"]))
+
+    assert args.gpu_memory_utilization == 0.8
+
+
+def test_device_memory_utilization_property():
+    config = CacheConfig(gpu_memory_utilization=0.8)
+
+    assert config.device_memory_utilization == 0.8
+
+    config.device_memory_utilization = 0.7
+    assert config.gpu_memory_utilization == 0.7
 
 
 @pytest.mark.parametrize(
     "options",
     [
-        ["--engram-config", '{"cpu_offload": false, "embedding_across_dp": true}'],
+        [
+            "--engram-config",
+            '{"cpu_offload": false, "embedding_across_dp": true}',
+        ],
         [
             "--engram-config.cpu_offload",
             "false",
@@ -79,9 +127,8 @@ def test_watermark_config_cli():
         ],
     ],
 )
-def test_engram_config_cli(options, monkeypatch):
-    """CLI settings take precedence over the legacy offload environment."""
-    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
+def test_engram_config_cli(options):
+    """JSON and dotted CLI options independently control Engram settings."""
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
     args = EngineArgs.from_cli_args(parser.parse_args(options))
     assert args.engram_config is not None
@@ -90,14 +137,26 @@ def test_engram_config_cli(options, monkeypatch):
 
 
 @pytest.mark.parametrize(
-    "options, provided",
-    [([], False), (["--engram-config", "{}"], True)],
+    "options,provided,dp_shared_memory",
+    [
+        ([], False, False),
+        (["--engram-config", "{}"], True, None),
+        (
+            ["--engram-config", '{"dp_shared_memory": true}'],
+            True,
+            True,
+        ),
+        (["--engram-config.dp_shared_memory", "true"], True, True),
+    ],
 )
-def test_engram_config_cli_optional(options, provided):
-    """An explicit empty config must remain distinct from an omitted config."""
+def test_engram_config_cli_optional(options, provided, dp_shared_memory):
+    """Explicit configs honor defaults and the JSON/dotted DP shared-memory flag."""
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
     args = EngineArgs.from_cli_args(parser.parse_args(options))
     assert (args.engram_config is not None) == provided
+    if provided:
+        assert args.engram_config.cpu_offload is True
+        assert args.engram_config.dp_shared_memory is dp_shared_memory
 
 
 @pytest.mark.parametrize(
@@ -221,10 +280,13 @@ def test_get_type_hints(type_hint, expected):
     assert get_type_hints(type_hint) == expected
 
 
-def test_get_kwargs():
-    kwargs = get_kwargs(DummyConfig)
-    print(kwargs)
+@pytest.fixture
+def dummy_config_kwargs():
+    return get_kwargs(DummyConfig)
+
 
+def test_get_kwargs(dummy_config_kwargs):
+    kwargs = dummy_config_kwargs
     # bools should not have their type set
     assert kwargs["regular_bool"].get("type") is None
     assert kwargs["optional_bool"].get("type") is None
@@ -234,7 +296,7 @@ def test_get_kwargs():
     assert kwargs["optional_bool_or_str"]["const"] is True
     assert "action" not in kwargs["optional_bool_or_str"]
     # optional literals should have None as a choice
-    assert kwargs["optional_literal"]["choices"] == ["x", "y", "None"]
+    assert kwargs["optional_literal"]["choices"] == ["x", "y", None]
     # tuples should have the correct nargs
     assert kwargs["tuple_n"]["nargs"] == "+"
     assert kwargs["tuple_2"]["nargs"] == 2
@@ -260,6 +322,22 @@ def test_get_kwargs():
     assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2)  # type: ignore[call-arg]
 
 
+@pytest.mark.parametrize(
+    ("args", "expected"),
+    [
+        (["--optional-literal", "None"], None),
+        (["--optional-literal", ""], None),
+        (["--optional-literal", "x"], "x"),
+    ],
+)
+def test_optional_handling(args, expected, dummy_config_kwargs):
+    parser = FlexibleArgumentParser()
+    parser.add_argument("--optional-literal", **dummy_config_kwargs["optional_literal"])
+
+    assert parser.parse_args(args).optional_literal is expected
+    assert "None" in parser.format_help()
+
+
 def test_jit_monitor_verbose_arg():
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
     args = parser.parse_args(["--jit-monitor-verbose"])
@@ -342,8 +420,7 @@ def test_media_io_kwargs_parser(arg, expected):
     ],
 )
 def test_optimization_level(args, expected):
-    """
-    Test space-separated optimization levels (-O 1, -O 2, -O 3) map to
+    """Test space-separated optimization levels (-O 1, -O 2, -O 3) map to
     optimization_level.
     """
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
@@ -362,9 +439,7 @@ def test_optimization_level(args, expected):
     ],
 )
 def test_mode_parser(args, expected):
-    """
-    Test compilation config modes (-cc.mode=int) map to compilation_config.
-    """
+    """Test compilation config modes (-cc.mode=int) map to compilation_config."""
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
     parsed_args = parser.parse_args(args)
     assert parsed_args.compilation_config.mode == expected
@@ -561,24 +636,6 @@ def test_prefix_cache_default():
     assert engine_args.prefix_cache_retention_interval == 64
 
 
-def test_prefix_cache_retention_interval_from_deprecated_env(
-    monkeypatch, caplog, disable_log_dedup
-):
-    monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64")
-
-    engine_args = EngineArgs()
-
-    assert engine_args.prefix_cache_retention_interval == 64
-    assert "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in caplog.text
-    assert "deprecated" in caplog.text
-    assert "prefix_cache_retention_interval" in caplog.text
-
-    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
-    args = parser.parse_args(["--prefix-cache-retention-interval", "32"])
-    engine_args = EngineArgs.from_cli_args(args)
-    assert engine_args.prefix_cache_retention_interval == 32
-
-
 @pytest.mark.parametrize(
     ("arg", "expected", "option"),
     [
diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
index 436db76a9acd..226908d67368 100644
--- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
+++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
@@ -480,7 +480,7 @@ class TestThinkingBlockConversion:
     """
 
     def test_thinking_plus_text_in_assistant_message(self):
-        """thinking + text → reasoning field + plain-string content."""
+        """Thinking + text → reasoning field + plain-string content."""
         request = _make_request(
             [
                 {"role": "user", "content": "Write me some code."},
@@ -542,7 +542,7 @@ def test_thinking_only_in_assistant_message(self):
         assert asst.get("content") is None
 
     def test_thinking_plus_tool_use_in_assistant_message(self):
-        """thinking + tool_use: reasoning field set, tool_calls populated."""
+        """Thinking + tool_use: reasoning field set, tool_calls populated."""
         request = _make_request(
             [
                 {"role": "user", "content": "What is 2+2?"},
@@ -1610,6 +1610,32 @@ async def sse_input():
         assert msg_deltas[0]["delta"]["stop_reason"] == "stop_sequence"
         assert msg_deltas[0]["delta"]["stop_sequence"] == "</tool>"
 
+    @pytest.mark.asyncio
+    async def test_streaming_no_stop_string_emits_explicit_null_stop_sequence(self):
+        """exclude_unset=True drops stop_sequence unless it is set explicitly."""
+
+        async def sse_input():
+            yield _make_stream_chunk(delta=DeltaMessage(role="assistant"))
+            yield _make_stream_chunk(delta=DeltaMessage(content="hi"))
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+
+        events = _parse_sse_events(output)
+        msg_deltas = [data for ev_type, data in events if ev_type == "message_delta"]
+        assert len(msg_deltas) == 1
+        assert msg_deltas[0]["delta"]["stop_reason"] == "end_turn"
+        assert "stop_sequence" in msg_deltas[0]["delta"]
+        assert msg_deltas[0]["delta"]["stop_sequence"] is None
+
 
 # ======================================================================
 # Client-caused errors are 4xx, not 500 (Issue #52088)
diff --git a/tests/entrypoints/cli/test_serve.py b/tests/entrypoints/cli/test_serve.py
new file mode 100644
index 000000000000..9ba9ac7455c3
--- /dev/null
+++ b/tests/entrypoints/cli/test_serve.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the ``vllm serve`` CLI subcommand."""
+
+import argparse
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm.entrypoints.cli.serve import run_headless
+
+
+class _StopAfterEngineConfig(Exception):
+    pass
+
+
+def test_headless_imports_reasoning_parser_plugin_before_engine_config():
+    plugin_path = "/tmp/custom_reasoning_parser.py"
+    args = argparse.Namespace(
+        api_server_count=0,
+        reasoning_parser_plugin=plugin_path,
+    )
+    engine_args = MagicMock()
+
+    with (
+        patch(
+            "vllm.entrypoints.cli.serve.vllm.AsyncEngineArgs.from_cli_args",
+            return_value=engine_args,
+        ),
+        patch(
+            "vllm.reasoning.ReasoningParserManager.import_reasoning_parser"
+        ) as import_plugin,
+    ):
+
+        def create_engine_config(**kwargs):
+            import_plugin.assert_called_once_with(plugin_path)
+            raise _StopAfterEngineConfig
+
+        engine_args.create_engine_config.side_effect = create_engine_config
+
+        with pytest.raises(_StopAfterEngineConfig):
+            run_headless(args)
diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py
index a90e93b6b165..94324dbed76e 100644
--- a/tests/entrypoints/cohere/test_api_router.py
+++ b/tests/entrypoints/cohere/test_api_router.py
@@ -58,7 +58,6 @@ def _enable_cohere_api(monkeypatch):
     flag inside the test body.
     """
     monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "1")
-    monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1")
 
 
 # ----------------------------------------------------------------------
@@ -128,15 +127,19 @@ def _generate_request() -> GenerateRequest:
     )
 
 
-def _build_app(handler: _Handler | None) -> FastAPI:
+def _build_app(handler: _Handler | None, *, enable_scale_out: bool = True) -> FastAPI:
     app = FastAPI()
+    app.state.args = Namespace(enable_scale_out=enable_scale_out, tokens_only=False)
     attach_router(app)
     app.state.cohere_serving_chat_v2 = handler
     return app
 
 
-def _build_render_app(chat_handler, render_handler) -> FastAPI:
+def _build_render_app(
+    chat_handler, render_handler, *, enable_scale_out: bool = True
+) -> FastAPI:
     app = FastAPI()
+    app.state.args = Namespace(enable_scale_out=enable_scale_out, tokens_only=False)
     attach_router(app)
     app.state.cohere_serving_chat_v2 = chat_handler
     app.state.serving_render = render_handler
@@ -151,11 +154,13 @@ def _build_app_with_vllm_handlers(handler: _Handler | None) -> FastAPI:
     into the ``CohereError`` wire shape.
     """
     app = FastAPI()
-    attach_router(app)
-    app.state.cohere_serving_chat_v2 = handler
     # ``validation_exception_handler`` reads ``req.app.state.args``; the
     # real cli builds this via argparse.
-    app.state.args = Namespace(log_error_stack=False)
+    app.state.args = Namespace(
+        log_error_stack=False, enable_scale_out=True, tokens_only=False
+    )
+    attach_router(app)
+    app.state.cohere_serving_chat_v2 = handler
     app.exception_handler(RequestValidationError)(validation_exception_handler)
     app.exception_handler(HTTPException)(http_exception_handler)
     app.exception_handler(VLLMError)(vllm_error_handler)
@@ -402,10 +407,8 @@ def test_route_registered(self):
         paths = [getattr(r, "path", None) for r in app.routes]
         assert "/cohere/v2/chat/render" in paths
 
-    def test_route_not_registered_when_flag_unset(self, monkeypatch):
-        monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False)
-        app = FastAPI()
-        attach_router(app)
+    def test_route_not_registered_when_flag_unset(self):
+        app = _build_app(handler=None, enable_scale_out=False)
         paths = [getattr(r, "path", None) for r in app.routes]
         assert "/cohere/v2/chat/render" not in paths
         assert "/cohere/v2/chat" in paths
diff --git a/tests/entrypoints/cohere/test_chat_v2.py b/tests/entrypoints/cohere/test_chat_v2.py
index 36acecb6224d..4f07d3d87355 100644
--- a/tests/entrypoints/cohere/test_chat_v2.py
+++ b/tests/entrypoints/cohere/test_chat_v2.py
@@ -57,6 +57,8 @@ def server():
         # the conversation to surface a thinking block (SmolLM2 doesn't
         # emit Cohere-style reasoning tokens out of the box).
         "--no-cohere-is-reasoning-model",
+        # Registers /cohere/v2/chat/render alongside /cohere/v2/chat.
+        "--enable-scale-out",
     ]
 
     # Cap the CPU KV-cache pool the vLLM CPU backend reserves at
@@ -71,7 +73,6 @@ def server():
     env_dict = {
         "VLLM_CPU_KVCACHE_SPACE": "1",
         "VLLM_ENABLE_COHERE_API": "1",
-        "VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1",
     }
 
     with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
diff --git a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py
index b67efcc22931..f413869c9009 100644
--- a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py
+++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py
@@ -60,7 +60,10 @@ class MockModelConfig:
     encoder_config = None
     generation_config: str = "auto"
     media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
+    skip_tokenizer_init: bool = False
+    is_encoder_decoder: bool = False
+    is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     vocab_size = 151936
 
     def get_diff_sampling_param(self):
diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py
index 40fefc237404..668f11544858 100644
--- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py
+++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py
@@ -5,6 +5,7 @@
 import socket
 import threading
 import time
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
@@ -14,6 +15,7 @@
     exit_before_report_worker,
 )
 from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
+from vllm.v1.engine.admission_control import SharedAdmissionStats
 from vllm.v1.utils import (
     APIServerProcessManager,
     get_engine_client_zmq_addr,
@@ -32,6 +34,18 @@ def mock_run_api_server_worker(listen_address, sock, args, client_config=None):
     print("Mock worker completed successfully")
 
 
+def update_admission_stats_worker(listen_address, sock, args, client_config):
+    stats = SharedAdmissionStats(
+        client_config,
+        client_config["client_count"],
+        client_config["client_index"],
+    )
+    value = client_config["client_index"] + 1
+    stats.set_num_requests(value)
+    args.barrier.wait()
+    args.results.put(stats.get_num_requests())
+
+
 # Module-level stub for the gather_actual_addresses test. Must be
 # importable by `multiprocessing.spawn` (no closures, no nesting).
 def defer_addresses_stub_worker(listen_address, sock, args, client_config):
@@ -131,6 +145,31 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update):
             assert not proc.is_alive()
 
 
+@pytest.mark.skip_global_cleanup
+def test_api_server_processes_share_admission_stats(api_server_args):
+    args = api_server_args.copy()
+    spawn_context = multiprocessing.get_context("spawn")
+    worker_args = SimpleNamespace(
+        barrier=spawn_context.Barrier(args["num_servers"]),
+        results=spawn_context.Queue(),
+        max_num_queued_reqs=10,
+    )
+    args["args"] = worker_args
+    args["target_server_fn"] = update_admission_stats_worker
+    manager = APIServerProcessManager(**args)
+
+    try:
+        for process in manager.processes:
+            process.join(timeout=30)
+            assert process.exitcode == 0
+
+        assert [worker_args.results.get(timeout=5) for _ in range(3)] == [6] * 3
+    finally:
+        manager.shutdown()
+        worker_args.results.close()
+        worker_args.results.join_thread()
+
+
 @patch("vllm.v1.utils.run_api_server_worker_proc", mock_run_api_server_worker)
 def test_wait_for_completion_or_failure(api_server_args):
     """Test that wait_for_completion_or_failure works with failures."""
diff --git a/tests/entrypoints/launchers/test_cli_args.py b/tests/entrypoints/launchers/test_cli_args.py
index e84abe5d2617..0bcca8cb41ac 100644
--- a/tests/entrypoints/launchers/test_cli_args.py
+++ b/tests/entrypoints/launchers/test_cli_args.py
@@ -152,7 +152,7 @@ def test_multiple_valid_inputs(serve_parser):
 
 ### Tests for serve argument validation that run prior to loading
 def test_enable_auto_choice_passes_without_tool_call_parser(serve_parser):
-    """Ensure validation fails if tool choice is enabled with no call parser"""
+    """Ensure validation fails if tool choice is enabled with no call parser."""
     # If we enable-auto-tool-choice, explode with no tool-call-parser
     args = serve_parser.parse_args(args=["--enable-auto-tool-choice"])
     with pytest.raises(TypeError):
@@ -160,7 +160,7 @@ def test_enable_auto_choice_passes_without_tool_call_parser(serve_parser):
 
 
 def test_enable_auto_choice_passes_with_tool_call_parser(serve_parser):
-    """Ensure validation passes with tool choice enabled with a call parser"""
+    """Ensure validation passes with tool choice enabled with a call parser."""
     args = serve_parser.parse_args(
         args=[
             "--enable-auto-tool-choice",
@@ -172,7 +172,7 @@ def test_enable_auto_choice_passes_with_tool_call_parser(serve_parser):
 
 
 def test_enable_auto_choice_fails_with_enable_reasoning(serve_parser):
-    """Ensure validation fails if reasoning is enabled with auto tool choice"""
+    """Ensure validation fails if reasoning is enabled with auto tool choice."""
     args = serve_parser.parse_args(
         args=[
             "--enable-auto-tool-choice",
@@ -197,7 +197,7 @@ def test_passes_with_reasoning_parser(serve_parser):
 
 
 def test_chat_template_validation_for_happy_paths(serve_parser):
-    """Ensure validation passes if the chat template exists"""
+    """Ensure validation passes if the chat template exists."""
     args = serve_parser.parse_args(
         args=["--chat-template", CHATML_JINJA_PATH.absolute().as_posix()]
     )
@@ -205,7 +205,7 @@ def test_chat_template_validation_for_happy_paths(serve_parser):
 
 
 def test_chat_template_validation_for_sad_paths(serve_parser):
-    """Ensure validation fails if the chat template doesn't exist"""
+    """Ensure validation fails if the chat template doesn't exist."""
     args = serve_parser.parse_args(args=["--chat-template", "does/not/exist"])
     with pytest.raises(VLLMValidationError):
         validate_parsed_serve_args(args)
@@ -304,13 +304,13 @@ def test_sse_keep_alive_interval_non_integer(serve_parser, value):
     ],
 )
 def test_middleware(serve_parser, cli_args, expected_middleware):
-    """Ensure multiple middleware args are parsed properly"""
+    """Ensure multiple middleware args are parsed properly."""
     args = serve_parser.parse_args(args=cli_args)
     assert args.middleware == expected_middleware
 
 
 def test_default_chat_template_kwargs_parsing(serve_parser):
-    """Ensure default_chat_template_kwargs JSON is parsed correctly"""
+    """Ensure default_chat_template_kwargs JSON is parsed correctly."""
     args = serve_parser.parse_args(
         args=["--default-chat-template-kwargs", '{"enable_thinking": false}']
     )
@@ -318,7 +318,7 @@ def test_default_chat_template_kwargs_parsing(serve_parser):
 
 
 def test_default_chat_template_kwargs_complex(serve_parser):
-    """Ensure complex default_chat_template_kwargs JSON is parsed correctly"""
+    """Ensure complex default_chat_template_kwargs JSON is parsed correctly."""
     kwargs_json = '{"enable_thinking": false, "custom_param": "value", "num": 42}'
     args = serve_parser.parse_args(args=["--default-chat-template-kwargs", kwargs_json])
     assert args.default_chat_template_kwargs == {
@@ -329,7 +329,7 @@ def test_default_chat_template_kwargs_complex(serve_parser):
 
 
 def test_default_chat_template_kwargs_default_none(serve_parser):
-    """Ensure default_chat_template_kwargs defaults to None"""
+    """Ensure default_chat_template_kwargs defaults to None."""
     args = serve_parser.parse_args(args=[])
     assert args.default_chat_template_kwargs is None
 
@@ -355,7 +355,7 @@ def test_resolve_default_chat_template_kwargs(default_kwargs, cohere_format, exp
 
 
 def test_default_chat_template_kwargs_invalid_json(serve_parser):
-    """Ensure invalid JSON raises an error"""
+    """Ensure invalid JSON raises an error."""
     with pytest.raises(SystemExit):
         serve_parser.parse_args(
             args=["--default-chat-template-kwargs", "not valid json"]
@@ -402,7 +402,7 @@ def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises):
 
 ### Tests for LoRA target modules parsing
 def test_lora_target_modules_single(serve_parser):
-    """Test parsing single lora-target-modules argument"""
+    """Test parsing single lora-target-modules argument."""
     args = serve_parser.parse_args(
         args=["--enable-lora", "--lora-target-modules", "o_proj"]
     )
@@ -410,7 +410,7 @@ def test_lora_target_modules_single(serve_parser):
 
 
 def test_lora_target_modules_multiple(serve_parser):
-    """Test parsing multiple lora-target-modules arguments"""
+    """Test parsing multiple lora-target-modules arguments."""
     args = serve_parser.parse_args(
         args=[
             "--enable-lora",
@@ -424,6 +424,6 @@ def test_lora_target_modules_multiple(serve_parser):
 
 
 def test_lora_target_modules_default_none(serve_parser):
-    """Test that lora-target-modules defaults to None"""
+    """Test that lora-target-modules defaults to None."""
     args = serve_parser.parse_args(args=[])
     assert args.lora_target_modules is None
diff --git a/tests/entrypoints/launchers/test_dp_supervisor.py b/tests/entrypoints/launchers/test_dp_supervisor.py
index eb5bf0a08875..f96b020d10d8 100644
--- a/tests/entrypoints/launchers/test_dp_supervisor.py
+++ b/tests/entrypoints/launchers/test_dp_supervisor.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Tests for DPSupervisor: unit tests and lifecycle integration tests.
+"""Tests for DPSupervisor: unit tests and lifecycle integration tests.
 
 Lifecycle integration tests replace child vLLM servers with lightweight
 aiohttp "fake" servers controlled by the test, so the suite runs without GPUs.
@@ -463,8 +462,7 @@ async def serve(self):
 
 
 class MockVLLMServer:
-    """
-    Minimal FastAPI server that mimics one vLLM replica.
+    """Minimal FastAPI server that mimics one vLLM replica.
     GET /health returns 200 when healthy, 503 otherwise.
     Health state is toggled by the test via set_healthy().
     """
@@ -582,8 +580,7 @@ def launch_mock_vllm_with_drain(
 
 
 async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -> bool:
-    """
-    GET /health on the supervisor once and check for expected_status.
+    """GET /health on the supervisor once and check for expected_status.
 
     Pass expected_status=-1 to assert the supervisor is not listening yet
     (a connection error is expected). The supervisor only starts its HTTP
@@ -705,8 +702,7 @@ async def _run_supervisor(
 
 @pytest.mark.asyncio
 async def test_basic_lifecycle(monkeypatch):
-    """
-    A) Supervisor is not listening while children are unhealthy.
+    """A) Supervisor is not listening while children are unhealthy.
     B) /health returns 200 once every child reports healthy.
     C) SIGTERM and shutdown
     """
@@ -777,8 +773,7 @@ async def test_basic_lifecycle_with_ssl(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_failed_startup(monkeypatch):
-    """
-    A) One of the vLLM servers crashes during startup.
+    """A) One of the vLLM servers crashes during startup.
     B) DPSupervisor detects this, and cleans up resources.
     """
     args = _make_args()
@@ -801,8 +796,7 @@ async def test_failed_startup(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_becomes_unhealthy(monkeypatch):
-    """
-    A) Supervisor is not listening while children are unhealthy.
+    """A) Supervisor is not listening while children are unhealthy.
     B) /health returns 200 once every child reports healthy.
     C) Child process becomes unhealthy.
     D) Detected and shutdown.
@@ -842,8 +836,7 @@ async def test_becomes_unhealthy(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_dp_server_fails(monkeypatch):
-    """
-    A) Supervisor is not listening while children are unhealthy.
+    """A) Supervisor is not listening while children are unhealthy.
     B) /health returns 200 once every child reports healthy.
     C) Child process fails.
     D) Detected and shutdown.
@@ -886,8 +879,7 @@ async def test_dp_server_fails(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_shutdown_timeout(monkeypatch: pytest.MonkeyPatch):
-    """
-    Child mock servers delay shutdown by 10s on SIGTERM (simulating in-flight
+    """Child mock servers delay shutdown by 10s on SIGTERM (simulating in-flight
     request drain).  The supervisor is configured with shutdown_timeout=10,
     so its total wait budget is 10 + CHILD_EXIT_GRACE_S seconds.  The
     children exit naturally within that window, so no force-kill should occur
diff --git a/tests/entrypoints/unit_tests/test_grpc_health.py b/tests/entrypoints/launchers/test_grpc_health.py
similarity index 100%
rename from tests/entrypoints/unit_tests/test_grpc_health.py
rename to tests/entrypoints/launchers/test_grpc_health.py
diff --git a/tests/entrypoints/launchers/test_launch_cli.py b/tests/entrypoints/launchers/test_launch_cli.py
index 443dd82fdd05..fc8301069858 100644
--- a/tests/entrypoints/launchers/test_launch_cli.py
+++ b/tests/entrypoints/launchers/test_launch_cli.py
@@ -3,15 +3,48 @@
 """Unit tests for the `vllm launch` CLI subcommand."""
 
 import argparse
-from unittest.mock import patch
+import hashlib
+import json
+import os
+import signal
+import subprocess
+import sys
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any
+from unittest.mock import create_autospec, patch
 
 import pytest
 
+import vllm.snapshot.runtime as snapshot_runtime
+import vllm.snapshot.server as snapshot_server
+from vllm.entrypoints.cli import snapshot as snapshot_cli
 from vllm.entrypoints.cli.launch import (
     LaunchSubcommand,
     RenderSubcommand,
     cmd_init,
 )
+from vllm.snapshot.controller import (
+    LocalSnapshotTools,
+    SnapshotCreateError,
+    SnapshotRestoreError,
+    _child_engine_argv,
+    create_snapshot,
+    restore_snapshot,
+)
+from vllm.snapshot.manifest import (
+    SnapshotCompatibilityError,
+    SnapshotManifest,
+    SnapshotRuntimeIdentity,
+    SnapshotSecurityError,
+    read_manifest,
+    validate_artifact_root,
+    validate_identity,
+    write_manifest_atomic,
+)
+from vllm.snapshot.runtime import ProcessInventory, _decode_endpoint, _TcpSocketRecord
+from vllm.snapshot.server import SnapshotCanaryError
+from vllm.snapshot.types import Oracle, oracles_match
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 
@@ -109,3 +142,1011 @@ def test_launch_registered_in_main():
     assert hasattr(launch_module, "cmd_init")
     subcmds = launch_module.cmd_init()
     assert any(s.name == "launch" for s in subcmds)
+
+
+_MODEL_REVISION = "c1899de289a04d12100db370d81485cdf75e47ca"
+
+
+def _run_cli(argv: list[str]) -> None:
+    from vllm.entrypoints.cli import main as cli_main
+
+    with patch.object(sys, "argv", ["vllm", *argv]):
+        cli_main.main()
+
+
+def test_snapshot_environment_contract(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    capsys: pytest.CaptureFixture[str],
+):
+    from vllm.entrypoints.serve.utils import api_utils
+
+    monkeypatch.setattr(api_utils, "cli_env_setup", pytest.fail)
+    secret = "snapshot-secret"
+    monkeypatch.setenv("VLLM_API_KEY", secret)
+    monkeypatch.setenv("VLLM_USER_SETTING", secret)
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", secret)
+    environment = dict(LocalSnapshotTools()._environment_identity())
+    assert (
+        environment["VLLM_USER_SETTING"]
+        == hashlib.sha256(b"VLLM_USER_SETTING\0snapshot-secret").hexdigest()
+    )
+    assert (
+        environment["VLLM_USER_SETTING"] != environment["VLLM_WORKER_MULTIPROC_METHOD"]
+    )
+    assert secret not in str(environment)
+    assert "VLLM_API_KEY" not in environment
+
+    monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD")
+    implicit_worker_default = dict(LocalSnapshotTools()._environment_identity())
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    assert dict(LocalSnapshotTools()._environment_identity()) == implicit_worker_default
+
+    monkeypatch.setenv("VLLM_USER_SETTING", f"{secret}-changed")
+    changed_environment = dict(LocalSnapshotTools()._environment_identity())
+    assert changed_environment["VLLM_USER_SETTING"] != environment["VLLM_USER_SETTING"]
+
+    artifact = tmp_path / "snapshot"
+    artifact.mkdir(mode=0o700)
+    write_manifest_atomic(artifact, _manifest(environment=tuple(environment.items())))
+    before = os.environ.copy()
+
+    _run_cli(["snapshot", "inspect", str(artifact)])
+
+    assert os.environ == before
+    output = capsys.readouterr().out
+    assert secret not in output
+    inspected = json.loads(output)
+    inspected_environment = dict(inspected["environment"])
+    assert inspected_environment == environment
+
+    tools = LocalSnapshotTools()
+    monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
+    manifest = tools.make_manifest(
+        argparse.Namespace(
+            hf_token=secret,
+            model_tag="Qwen/Qwen3-0.6B",
+            revision=_MODEL_REVISION,
+            served_model_name=None,
+            tokenizer_revision=None,
+        ),
+        (
+            "Qwen/Qwen3-0.6B",
+            "--api-key",
+            secret,
+            f"{secret}-alternate",
+            f"--api-key={secret}",
+            "--hf-token",
+            secret,
+            f"--hf-token={secret}",
+        ),
+        ProcessInventory(100, (100, 101), (101,), "GPU-abc"),
+        _oracle(),
+        tmp_path,
+    )
+
+    assert manifest.engine_argv == (
+        "Qwen/Qwen3-0.6B",
+        "--api-key",
+        "***",
+        "***",
+        "--api-key=***",
+        "--hf-token",
+        "***",
+        "--hf-token=***",
+    )
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        (("--hf_token", "SECRET"), ("--hf_token", "***")),
+        (("--hf_token=SECRET",), ("--hf_token=***",)),
+        (("--api_key", "SECRET"), ("--api_key", "***")),
+        (("--api_key=SECRET",), ("--api_key=***",)),
+        (("--hf-tok", "SECRET"), ("--hf-tok", "***")),
+        (("--hf-tok=SECRET",), ("--hf-tok=***",)),
+        (("--hf_tok", "SECRET"), ("--hf_tok", "***")),
+        (("--api-k", "SECRET", "SECOND"), ("--api-k", "***", "***")),
+        (("--api-k=SECRET",), ("--api-k=***",)),
+        (("--hf-overrides", "{}"), ("--hf-overrides", "{}")),
+    ],
+)
+def test_snapshot_manifest_redacts_accepted_secret_options(case):
+    engine_argv, expected = case
+    tools = create_autospec(LocalSnapshotTools, instance=True)
+    tools.current_identity.return_value = _runtime_identity()
+    tools._artifact_bytes.return_value = 0
+    manifest = LocalSnapshotTools.make_manifest(
+        tools,
+        snapshot_server.parse_vllm_args(["model", *engine_argv]),
+        engine_argv,
+        ProcessInventory(100, (100, 101), (101,), "GPU-abc"),
+        _oracle(),
+        Path(),
+    )
+
+    assert manifest.engine_argv == expected
+    if "SECRET" in " ".join(engine_argv):
+        assert "***" in " ".join(manifest.engine_argv)
+        assert "SECRET" not in " ".join(manifest.engine_argv)
+
+
+def parse_snapshot(*argv: str):
+    parser = FlexibleArgumentParser()
+    subparsers = parser.add_subparsers(dest="subparser", required=True)
+    snapshot_cli.SnapshotSubcommand(
+        create_requested=argv[:1] == ("create",)
+    ).subparser_init(subparsers)
+    return parser.parse_args(["snapshot", *argv])
+
+
+@pytest.mark.parametrize("option", ["--snapshot-dir", "--snapshot_dir", "--snapshot-d"])
+@pytest.mark.parametrize("joined", [False, True])
+def test_snapshot_create_forwards_only_engine_options(option, joined, monkeypatch):
+    directory = [f"{option}=/tmp/snapshot"] if joined else [option, "/tmp/snapshot"]
+    argv = ["create", "Qwen/Qwen3-0.6B", "--revision", _MODEL_REVISION, *directory]
+    args = parse_snapshot(*argv)
+    monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv])
+
+    child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None)))
+
+    assert child.model == "Qwen/Qwen3-0.6B"
+    assert child.revision == _MODEL_REVISION
+    assert child.enable_sleep_mode
+
+
+@pytest.mark.parametrize("model_in_cli", [False, True])
+def test_snapshot_create_config_preserves_cli_precedence(
+    tmp_path, monkeypatch, model_in_cli
+):
+    config = tmp_path / "engine.yaml"
+    config.write_text("model: config-model\nmax-model-len: 512\ndtype: float16\n")
+    model = ["Qwen/Qwen3-0.6B"] if model_in_cli else []
+    argv = [
+        "create",
+        *model,
+        "--snapshot-dir",
+        str(tmp_path / "snapshot"),
+        "--config",
+        str(config),
+        "--revision",
+        _MODEL_REVISION,
+        "--max-model-len",
+        "1024",
+    ]
+    args = parse_snapshot(*argv)
+    monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv])
+    child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None)))
+
+    assert (
+        (args.model_tag or args.model)
+        == child.model
+        == ("Qwen/Qwen3-0.6B" if model_in_cli else "config-model")
+    )
+    assert args.max_model_len == child.max_model_len == 1024
+    assert args.dtype == child.dtype == "float16"
+
+
+@pytest.mark.parametrize(
+    ("probe", "status", "output"),
+    [
+        (None, 0, ""),
+        ("nvidia-smi", 1, ""),
+        ("find", 1, ""),
+        ("ss", 1, ""),
+        ("nvidia-smi", 0, "12345"),
+        ("find", 0, "link_remap.fixture"),
+        ("ss", 0, "LISTEN"),
+    ],
+)
+def test_snapshot_e2e_cleanup_requires_successful_empty_probes(probe, status, output):
+    script = (
+        Path(__file__).resolve().parents[3]
+        / ".buildkite/scripts/initialized-snapshot-e2e.sh"
+    )
+    if not script.exists():
+        pytest.skip("snapshot E2E script is not packaged in this test image")
+    source = script.read_text()
+    # Exercise the real probes without entering Docker/CRIU setup.
+    functions = source[source.index("link_remaps() {") : source.index("wait_clean() {")]
+    shell = (
+        """set -Eeuo pipefail
+GPU_UUID=fixture PORT_ONE=18001 PORT_TWO=18002 LINK_REMAP_BASELINE=''
+SNAPSHOT_PIDS=(999999999)
+nvidia-smi() { :; }
+find() { :; }
+ss() { :; }
+"""
+        + functions
+    )
+    if probe:
+        shell += f'\n{probe}() {{ printf %s "{output}"; return {status}; }}\n'
+    shell += "\nif state_clean; then exit 0; else exit 1; fi\n"
+
+    result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10)
+
+    assert result.returncode == (0 if probe is None else 1), result.stderr
+
+
+def test_snapshot_create_cli_accepts_only_pinned_compact_mode(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setattr(snapshot_cli.platform, "system", lambda: "Linux")
+    monkeypatch.setattr(snapshot_cli.platform, "machine", lambda: "x86_64")
+    base = ("create", "Qwen/Qwen3-0.6B", "--snapshot-dir=/tmp/snapshot")
+    compact = parse_snapshot(*base, "--revision", _MODEL_REVISION)
+
+    snapshot_cli.validate_create_args(compact)
+    assert compact.model_tag == "Qwen/Qwen3-0.6B"
+
+    dispatched: list[argparse.Namespace] = []
+    monkeypatch.setattr(snapshot_cli, "run_create", dispatched.append)
+    _run_cli(["snapshot", *base, "--revision", _MODEL_REVISION])
+    assert dispatched[0].revision == _MODEL_REVISION
+
+
+@pytest.mark.parametrize(
+    ("attribute", "value", "message"),
+    [
+        ("prefill_context_parallel_size", 2, "prefill context parallel size 1"),
+        ("model_tag", None, "immutable remote model"),
+        ("logprobs_mode", "raw_logits", "log-probability mode"),
+        ("speculative_config", {}, "does not support speculative decoding"),
+        ("spec_method", "ngram", "does not support speculative decoding"),
+        ("spec_model", "draft-model", "does not support speculative decoding"),
+        ("spec_tokens", 3, "does not support speculative decoding"),
+    ],
+)
+def test_snapshot_create_rejects_invalid_config(
+    attribute: str,
+    value: object,
+    message: str,
+    tmp_path: Path,
+):
+    args = parse_snapshot(
+        "create",
+        "Qwen/Qwen3-0.6B",
+        "--snapshot-dir=/tmp/snapshot",
+        "--revision",
+        _MODEL_REVISION,
+    )
+    if attribute == "model_tag":
+        local_model = tmp_path / "model"
+        local_model.mkdir()
+        value = str(local_model)
+    setattr(args, attribute, value)
+    with pytest.raises(ValueError, match=message):
+        snapshot_cli.validate_create_args(args)
+
+
+def _oracle(
+    token_ids: tuple[int, ...] = (12095,),
+    text: str = " Paris",
+    logprob: float = -0.125,
+) -> Oracle:
+    return Oracle(token_ids, text, logprob)
+
+
+def _manifest(**changes: object) -> SnapshotManifest:
+    manifest = SnapshotManifest(
+        schema_version=1,
+        boundary="post-engine-init-reloadable-state-released",
+        created_at="2026-08-06T00:00:00Z",
+        artifact_bytes=8,
+        vllm_version="0.1.dev1+snapshot",
+        python_version="3.12.3",
+        torch_version="2.9.0",
+        cuda_runtime="12.9",
+        driver_version="575.57.08",
+        criu_version="4.1",
+        cuda_checkpoint_sha256="a" * 64,
+        kernel_release="6.8.0",
+        host_id="host-a",
+        gpu_name="NVIDIA A10",
+        gpu_uuid="GPU-abc",
+        model="Qwen/Qwen3-0.6B",
+        served_model_name="Qwen/Qwen3-0.6B",
+        model_revision="model-sha",
+        tokenizer_revision="tokenizer-sha",
+        engine_argv=("Qwen/Qwen3-0.6B",),
+        environment=(("VLLM_USE_V1", "1"),),
+        process_tree=(100, 101),
+        cuda_holders=(101,),
+        oracle_token_ids=(12095,),
+        oracle_text=" Paris",
+        oracle_sampled_token_logprob=-0.125,
+    )
+    return SnapshotManifest.model_validate(
+        {**manifest.model_dump(mode="python"), **changes}, strict=True
+    )
+
+
+def _runtime_identity(**changes: object) -> Any:
+    fields = SnapshotRuntimeIdentity.model_fields.keys()
+    payload = _manifest().model_dump(mode="python", include=fields)
+    return SnapshotRuntimeIdentity.model_validate({**payload, **changes}, strict=True)
+
+
+def _fake_snapshot_tools() -> Any:
+    tools = create_autospec(LocalSnapshotTools, instance=True)
+    tools.launch_child.return_value = 100
+    tools.wait_ready.return_value = _oracle()
+    tools.inventory.return_value = ProcessInventory(100, (100, 101), (101,), "GPU-abc")
+
+    tools.current_identity.side_effect = lambda _gpu_uuid: _runtime_identity()
+    tools.restore.return_value = 100
+    tools.request_oracle.return_value = _oracle()
+    return tools
+
+
+@pytest.mark.parametrize("value", ["0", "-1", "nan", "inf"])
+def test_snapshot_timeout_must_be_positive_and_finite(
+    value: str, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv("VLLM_SNAPSHOT_TIMEOUT_S", value)
+    with pytest.raises(ValueError, match="positive finite number"):
+        LocalSnapshotTools()
+
+
+def test_snapshot_create_abort_reaps_after_process_group_cleanup(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    tools = LocalSnapshotTools()
+    process = subprocess.Popen(
+        [sys.executable, "-c", "raise SystemExit(1)"],
+        stdin=subprocess.DEVNULL,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        start_new_session=True,
+    )
+    tools._children[process.pid] = process
+    (tmp_path / "child.log").write_text("root failed")
+    killed: list[tuple[int, signal.Signals]] = []
+    monkeypatch.setattr(
+        os,
+        "killpg",
+        lambda pid, sig: killed.append((pid, sig)),
+    )
+    monkeypatch.setattr(os, "waitid", lambda *_args: object(), raising=False)
+
+    try:
+        with pytest.raises(SnapshotCreateError, match="root failed"):
+            tools.wait_ready(tmp_path, process.pid)
+        assert process.returncode is None
+        tools.abort_create(process.pid)
+    finally:
+        if process.pid in tools._children:
+            tools.abort_create(process.pid)
+
+    assert killed == [(process.pid, signal.SIGKILL)]
+    assert process.pid not in tools._children
+
+
+def _supply_waitid(monkeypatch: pytest.MonkeyPatch) -> None:
+    """`wait_ready` binds Linux-only `os.waitid` before it reads the marker."""
+    monkeypatch.setattr(os, "waitid", lambda *_args: None, raising=False)
+
+
+def test_snapshot_wait_ready_reads_back_the_oracle_the_child_wrote(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    _supply_waitid(monkeypatch)
+    snapshot_server.write_ready_atomic(tmp_path / "ready.json", _oracle())
+
+    assert LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) == _oracle()
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [
+        {"token_ids": [12095], "sampled_token_logprob": -0.125},
+        {"token_ids": ["12095"], "text": " Paris", "sampled_token_logprob": -0.125},
+    ],
+)
+def test_snapshot_wait_ready_rejects_a_malformed_marker(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, payload: dict[str, object]
+):
+    _supply_waitid(monkeypatch)
+    (tmp_path / "ready.json").write_text(json.dumps(payload))
+
+    with pytest.raises(SnapshotCreateError, match="ready marker invalid"):
+        LocalSnapshotTools().wait_ready(tmp_path, os.getpid())
+
+
+def test_snapshot_create_abort_rejects_surviving_process_group(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    tools = LocalSnapshotTools()
+    process = create_autospec(subprocess.Popen, instance=True)
+    process.wait.side_effect = subprocess.TimeoutExpired("snapshot child", 10)
+    tools._children[100] = process
+    monkeypatch.setattr(os, "killpg", lambda _pid, _sig: None)
+
+    with pytest.raises(SnapshotCreateError, match="survived SIGKILL"):
+        tools.abort_create(100)
+
+
+@pytest.mark.asyncio
+async def test_snapshot_child_rejects_mismatch_before_publication(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    from vllm.entrypoints.launchers.api_server import entry as api_server
+
+    canaries = iter((_oracle(), _oracle((42,), " Lyon")))
+
+    @asynccontextmanager
+    async def engine_context():
+        yield object()
+
+    async def next_canary(_engine: object) -> Oracle:
+        return next(canaries)
+
+    async def return_none(_engine: object) -> None:
+        return None
+
+    monkeypatch.setattr(
+        api_server, "build_async_engine_client", lambda _args: engine_context()
+    )
+    monkeypatch.setattr(snapshot_server, "run_engine_canary", next_canary)
+    monkeypatch.setattr(snapshot_server, "_release_reloadable_state", return_none)
+    monkeypatch.setattr(snapshot_server, "_restore_reloadable_state", return_none)
+    monkeypatch.setattr(snapshot_server, "detach_snapshot_streams", lambda: None)
+    monkeypatch.setattr(
+        snapshot_server,
+        "write_ready_atomic",
+        lambda *_args: pytest.fail("mismatched snapshot was published"),
+    )
+    control = snapshot_server.ControlArgs(
+        ready_file=tmp_path / "ready.json",
+        release_file=tmp_path / "release.json",
+        release_timeout_s=1,
+    )
+    with pytest.raises(SnapshotCanaryError, match="rehearsal changed"):
+        await snapshot_server.run_vllm_snapshot_child(control, argparse.Namespace())
+
+
+def test_snapshot_release_marker_ignores_unknown_keys(tmp_path: Path):
+    marker = tmp_path / "release.json"
+    marker.write_text(json.dumps({"release": True, "port": 8000, "unknown": 1}))
+
+    listener = snapshot_server.read_release_marker(marker)
+
+    assert listener == snapshot_server.ListenerConfig(host=None, port=8000)
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [
+        ["release"],
+        {"port": 8000},
+        {"release": "true", "port": 8000},
+        {"release": 1, "port": 8000},
+        {"release": True, "port": 0},
+        {"release": True, "port": 70000},
+        {"release": True, "port": True},
+        {"release": True, "port": "8000"},
+        {"release": True, "port": 8000, "host": 5},
+    ],
+)
+def test_snapshot_release_marker_rejects_an_invalid_payload(
+    tmp_path: Path, payload: object
+):
+    marker = tmp_path / "release.json"
+    marker.write_text(json.dumps(payload))
+
+    with pytest.raises(snapshot_server.SnapshotBarrierError, match="release marker"):
+        snapshot_server.read_release_marker(marker)
+
+
+def test_snapshot_create_rolls_back_failed_dump(tmp_path: Path):
+    target = tmp_path / "snapshot"
+    tools = _fake_snapshot_tools()
+    tools.dump.side_effect = RuntimeError("dump failed")
+    with pytest.raises(RuntimeError, match="dump failed"):
+        create_snapshot(
+            argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"),
+            tools=tools,
+        )
+    tools.abort_create.assert_called_once_with(100)
+    assert not target.exists()
+
+
+def _restore_fixture(tmp_path: Path) -> argparse.Namespace:
+    artifact = tmp_path / "snapshot"
+    artifact.mkdir(mode=0o700)
+    write_manifest_atomic(artifact, _manifest())
+    return argparse.Namespace(snapshot_dir=str(artifact), host="127.0.0.1", port=9000)
+
+
+def test_snapshot_restore_rejects_identity_mismatch(tmp_path: Path):
+    args = _restore_fixture(tmp_path)
+
+    identity_mismatch = _fake_snapshot_tools()
+    identity_mismatch.current_identity.side_effect = lambda _gpu_uuid: (
+        _runtime_identity(vllm_version="different")
+    )
+    with pytest.raises(SnapshotCompatibilityError, match="vllm_version"):
+        restore_snapshot(args, tools=identity_mismatch)
+    identity_mismatch.restore.assert_not_called()
+
+
+def test_snapshot_identity_mismatch_names_every_differing_field():
+    with pytest.raises(SnapshotCompatibilityError) as mismatch:
+        validate_identity(
+            _manifest(),
+            _runtime_identity(vllm_version="different", gpu_uuid="GPU-other"),
+        )
+
+    assert str(mismatch.value) == "snapshot mismatch: vllm_version, gpu_uuid"
+
+
+def test_snapshot_restore_cleans_oracle_mismatch(tmp_path: Path):
+    args = _restore_fixture(tmp_path)
+    oracle_mismatch = _fake_snapshot_tools()
+    oracle_mismatch.request_oracle.return_value = _oracle(logprob=-0.5)
+    with pytest.raises(SnapshotRestoreError, match="oracle mismatch"):
+        restore_snapshot(args, tools=oracle_mismatch)
+    oracle_mismatch.cleanup.assert_called_once()
+    oracle_mismatch.complete_restore.assert_not_called()
+    assert oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.0005))
+    assert not oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.002))
+
+
+@pytest.mark.parametrize(
+    ("invalid_update", "diagnostic"),
+    [
+        ({"unexpected": True}, "unexpected"),
+        ({"schema_version": True}, "schema_version"),
+        ({"boundary": "unknown"}, "boundary"),
+        ({"process_tree": [100, 100]}, "root"),
+        ({"cuda_holders": [999]}, "root"),
+        ({"oracle_token_ids": [1, 2]}, "oracle_token_ids"),
+        ({"oracle_sampled_token_logprob": 0}, "oracle_sampled_token_logprob"),
+        (
+            {"oracle_sampled_token_logprob": float("nan")},
+            "oracle_sampled_token_logprob",
+        ),
+    ],
+)
+def test_snapshot_manifest_validation(
+    tmp_path: Path,
+    invalid_update: dict[str, object],
+    diagnostic: str,
+):
+    artifact = tmp_path / "snapshot"
+    artifact.mkdir(mode=0o700)
+    write_manifest_atomic(artifact, _manifest())
+    invalid = _manifest().model_dump(mode="json")
+    invalid.update(invalid_update)
+    (artifact / "manifest.json").write_text(json.dumps(invalid))
+    with pytest.raises(
+        SnapshotCompatibilityError,
+        match=f"invalid snapshot manifest: {diagnostic}",
+    ):
+        read_manifest(artifact)
+
+
+def _local_restore_artifact(tmp_path: Path, name: str = "snapshot") -> Path:
+    artifact = tmp_path / name
+    artifact.mkdir(mode=0o700)
+    (artifact / "child.log").write_bytes(b"startup log")
+    (artifact / "child.log.snapshot-size").write_text("11\n")
+    return artifact
+
+
+def _mark_snapshot_pids_free(monkeypatch: pytest.MonkeyPatch) -> list[int]:
+    probed_pids: list[int] = []
+
+    def missing(pid: int, _signal: int) -> None:
+        probed_pids.append(pid)
+        raise ProcessLookupError
+
+    monkeypatch.setattr(snapshot_runtime.os, "kill", missing)
+    return probed_pids
+
+
+@pytest.mark.parametrize(
+    ("probe_error", "diagnostic"),
+    [
+        (None, "already occupied: 100"),
+        (PermissionError(), "already occupied: 100"),
+        (OSError("probe failed"), "availability probe failed: 100"),
+    ],
+)
+def test_snapshot_restore_rejects_unavailable_pid_before_mutation(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    probe_error: OSError | None,
+    diagnostic: str,
+):
+    args = _restore_fixture(tmp_path)
+    artifact = Path(args.snapshot_dir)
+    (artifact / "child.log").write_bytes(b"startup log")
+    (artifact / "release.json").write_bytes(b"release sentinel")
+    before = {path.name: path.read_bytes() for path in artifact.iterdir()}
+    tools = LocalSnapshotTools()
+
+    def fail_probe(_pid: int, _signal: int) -> None:
+        if probe_error is not None:
+            raise probe_error
+
+    monkeypatch.setattr(snapshot_runtime.os, "kill", fail_probe)
+    monkeypatch.setattr(tools, "preflight", lambda *_args: None)
+    monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
+    monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
+
+    with pytest.raises(SnapshotRestoreError, match=diagnostic):
+        restore_snapshot(args, tools=tools)
+    assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before
+
+
+@pytest.mark.parametrize("mode", ["pinned", "unpinned", "cleanup-fails"])
+def test_snapshot_restore_terminates_only_a_pinned_tree(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    mode: str,
+):
+    _mark_snapshot_pids_free(monkeypatch)
+    artifact = _local_restore_artifact(tmp_path)
+    manifest = _manifest()
+    root_pid = manifest.process_tree[0]
+    tools = LocalSnapshotTools()
+    terminated: list[int] = []
+
+    def terminate_and_wait(pid: int) -> None:
+        terminated.append(pid)
+        if mode == "cleanup-fails":
+            raise SnapshotRestoreError("restored process cleanup is incomplete")
+        tools._restored_processes.pop(pid)
+
+    def pin(
+        _artifact: Path, process_tree: tuple[int, ...], _holders: tuple[int, ...]
+    ) -> None:
+        if mode == "unpinned":
+            raise SnapshotRestoreError(
+                "restored session does not match the captured process tree"
+            )
+        tools._restored_processes[process_tree[0]] = ()
+
+    monkeypatch.setattr(tools, "cleanup", terminate_and_wait)
+    monkeypatch.setattr(tools, "_pin_restored_tree", pin)
+    monkeypatch.setattr(tools, "_criu", lambda *_args: None)
+    monkeypatch.setattr(
+        tools,
+        "_run",
+        lambda command, **_kwargs: subprocess.CompletedProcess(command, 0, "333\n", ""),
+    )
+
+    with pytest.raises(SnapshotRestoreError, match="does not match") as failure:
+        tools.restore(artifact, manifest)
+
+    # A pinned tree is ours: it must be terminated and waited for, not left to
+    # honour the abort marker. An unpinned PID is never signalled.
+    assert terminated == ([] if mode == "unpinned" else [root_pid])
+    if mode == "cleanup-fails":
+        # The failure to terminate is reported beside the primary error, and a
+        # tree that may still be running stays pinned.
+        assert "restore cleanup failed" in str(failure.value)
+        assert tools._restored_processes == {root_pid: ()}
+    else:
+        assert tools._restored_processes == {}
+    assert json.loads((artifact / "release.json").read_text())["release"] is False
+
+
+def test_snapshot_private_path_and_link_remap_security(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    _mark_snapshot_pids_free(monkeypatch)
+    real = tmp_path / "real"
+    real.mkdir(mode=0o700)
+    linked = tmp_path / "linked"
+    linked.symlink_to(real, target_is_directory=True)
+    with pytest.raises(SnapshotSecurityError, match="symlink"):
+        validate_artifact_root(linked, creating=False)
+
+    unsafe = tmp_path / "unsafe"
+    unsafe.mkdir()
+    unsafe.chmod(0o777)
+    with pytest.raises(SnapshotSecurityError, match="world-writable"):
+        validate_artifact_root(unsafe / "snapshot", creating=True)
+
+    artifact = _local_restore_artifact(tmp_path)
+    saved_remaps = artifact / "link-remaps"
+    saved_remaps.mkdir()
+    (saved_remaps / "link_remap.270").write_bytes(b"semaphore state")
+    shm_dir = tmp_path / "shm"
+    shm_dir.mkdir()
+    target = shm_dir / "link_remap.270"
+    replacement = tmp_path / "replacement"
+    replacement.write_bytes(b"semaphore state")
+    target.symlink_to(replacement)
+    tools = LocalSnapshotTools()
+    tools.shm_dir = shm_dir
+    monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
+    with pytest.raises(SnapshotRestoreError, match="conflicting CRIU link remap"):
+        tools.restore(artifact, _manifest())
+
+
+@pytest.mark.parametrize(
+    ("blocked_state", "message"),
+    [("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")],
+)
+def test_snapshot_rejects_unsafe_process_state_before_criu(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, blocked_state: str, message: str
+):
+    tools = LocalSnapshotTools()
+    monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101))
+    monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
+    monkeypatch.setattr(
+        tools,
+        "_descriptor_inventory",
+        lambda _tree: (
+            (101,) if blocked_state == "io_uring" else (),
+            {41: 101},
+        ),
+    )
+    tcp_record = _TcpSocketRecord(
+        family="AF_INET",
+        local_raw="3C00000A:D6F6",
+        remote_raw="01010101:01BB" if blocked_state == "tcp" else "00000000:0000",
+        inode=41,
+    )
+    monkeypatch.setattr(tools, "_tcp_records", lambda: (tcp_record,))
+    with pytest.raises(SnapshotCreateError, match=message) as excinfo:
+        tools.inventory(100, tmp_path)
+    if blocked_state == "tcp":
+        assert "pid 101" in str(excinfo.value)
+        assert "1.1.1.1:443" in str(excinfo.value)
+
+
+def test_decode_endpoint_families():
+    assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030"
+    assert (
+        _decode_endpoint("AF_INET6", "00000000000000000000000001000000:01BB")
+        == "[::1]:443"
+    )
+
+
+def test_tcp_records_skips_missing_tables(tmp_path, monkeypatch):
+    tcp, tcp6 = tmp_path / "tcp", tmp_path / "tcp6"
+    contents = "header\n0: 0100007F:1F90 00000000:0000 01 0 0 0 0 0 41\n"
+    tcp.write_text(contents)
+    tables = (("AF_INET", tcp), ("AF_INET6", tcp6))
+    monkeypatch.setattr(snapshot_runtime, "_TCP_TABLES", tables)
+
+    tools = LocalSnapshotTools()
+    assert tuple(record.family for record in tools._tcp_records()) == ("AF_INET",)
+
+    tcp6.write_text(contents)
+    families = tuple(record.family for record in tools._tcp_records())
+    assert families == ("AF_INET", "AF_INET6")
+
+
+def test_snapshot_manifest_records_external_cache_files(tmp_path: Path):
+    artifact = tmp_path / "snapshot"
+    artifact.mkdir(mode=0o700)
+    recorded = (
+        (str(tmp_path / "kernel.so"), "b" * 64),
+        (str(tmp_path / "launcher.lock"), "locked"),
+    )
+    write_manifest_atomic(artifact, _manifest(external_cache_files=recorded))
+
+    assert read_manifest(artifact).external_cache_files == recorded
+
+    # A manifest written before this field existed still loads.
+    legacy = _manifest().model_dump(mode="json")
+    del legacy["external_cache_files"]
+    (artifact / "manifest.json").write_text(json.dumps(legacy))
+
+    assert read_manifest(artifact).external_cache_files == ()
+
+
+def _pin_generated_cache_roots(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    """Keep the resolver off whatever cache roots the host running this has."""
+    monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache"))
+    monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "triton_cache_default"))
+    monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor_cache"))
+
+
+def _unreadable(_path: Path) -> str:
+    raise OSError("cache file is unreadable")
+
+
+def test_snapshot_inventory_records_open_generated_cache_files(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    cache_root = tmp_path / "triton_cache"
+    cache_root.mkdir()
+    kernel = cache_root / "kernel.so"
+    kernel.write_bytes(b"compiled kernel")
+    lock = cache_root / "launcher.lock"
+    lock.write_bytes(b"held by pid 100")
+    weights = tmp_path / "model.safetensors"
+    weights.write_bytes(b"weights")
+    artifact = tmp_path / "snapshot"
+    artifact.mkdir(mode=0o700)
+    (artifact / "child.log.lock").write_bytes(b"")
+    _pin_generated_cache_roots(monkeypatch, tmp_path)
+    monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root))
+
+    targets = {
+        100: (str(kernel), str(weights), "socket:[41]"),
+        101: (str(lock), str(cache_root), str(artifact / "child.log.lock")),
+        102: (f"{cache_root / 'evicted.so'} (deleted)",),
+    }
+    tools = LocalSnapshotTools()
+    monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101, 102))
+    monkeypatch.setattr(tools, "_descriptor_targets", lambda pid: targets[pid])
+    monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",))
+    monkeypatch.setattr(tools, "_tcp_records", lambda: ())
+    monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity())
+
+    manifest = tools.make_manifest(
+        argparse.Namespace(
+            model_tag="Qwen/Qwen3-0.6B",
+            revision=_MODEL_REVISION,
+            served_model_name=None,
+            tokenizer_revision=None,
+        ),
+        ("Qwen/Qwen3-0.6B",),
+        tools.inventory(100, artifact),
+        _oracle(),
+        artifact,
+    )
+
+    # Weights, sockets, the cache directory itself, files inside the artifact,
+    # and deleted targets that CRIU remaps are all left out.
+    assert manifest.external_cache_files == (
+        (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
+        (str(lock), "locked"),
+    )
+
+
+def test_snapshot_inventory_marks_a_cache_file_it_cannot_read(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    cache_root = tmp_path / "triton_cache"
+    cache_root.mkdir()
+    kernel = cache_root / "kernel.so"
+    kernel.write_bytes(b"compiled kernel")
+    lock = cache_root / "launcher.lock"
+    lock.write_bytes(b"held by pid 100")
+    artifact = tmp_path / "snapshot"
+    artifact.mkdir(mode=0o700)
+    _pin_generated_cache_roots(monkeypatch, tmp_path)
+    monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root))
+
+    tools = LocalSnapshotTools()
+    monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,))
+    monkeypatch.setattr(
+        tools, "_descriptor_targets", lambda _pid: (str(kernel), str(lock))
+    )
+    monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",))
+    monkeypatch.setattr(tools, "_tcp_records", lambda: ())
+    monkeypatch.setattr(tools, "_sha256", _unreadable)
+
+    # A file that cannot be digested keeps its own marker, so the manifest says
+    # which entries restore can check for existence only.
+    assert tools.inventory(100, artifact).external_cache_files == (
+        (str(kernel), "unreadable"),
+        (str(lock), "locked"),
+    )
+
+
+def test_snapshot_inventory_matches_resolved_cache_roots_not_siblings(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    # Descriptor targets come back kernel-canonical, so the fixture must be too.
+    base = Path(os.path.realpath(tmp_path))
+    real_cache = base / "mnt_cache"
+    real_cache.mkdir()
+    kernel = real_cache / "kernel.so"
+    kernel.write_bytes(b"compiled kernel")
+    linked_cache = base / "linked_cache"
+    linked_cache.symlink_to(real_cache, target_is_directory=True)
+    sibling = base / "vllm_cache-backup"
+    sibling.mkdir()
+    stale = sibling / "stale.so"
+    stale.write_bytes(b"stale kernel")
+    artifact = base / "snapshot"
+    artifact.mkdir(mode=0o700)
+
+    _pin_generated_cache_roots(monkeypatch, base)
+    monkeypatch.setenv("TRITON_CACHE_DIR", str(linked_cache))
+
+    tools = LocalSnapshotTools()
+    monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,))
+    monkeypatch.setattr(
+        tools, "_descriptor_targets", lambda _pid: (str(kernel), str(stale))
+    )
+    monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",))
+    monkeypatch.setattr(tools, "_tcp_records", lambda: ())
+
+    # A cache root reached through a symlink still covers the file the kernel
+    # reports by its real path, and a directory that merely shares the root's
+    # name prefix is not part of the cache at all.
+    assert tools.inventory(100, artifact).external_cache_files == (
+        (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
+    )
+
+
+def test_snapshot_restore_rejects_rotated_cache_file_before_mutation(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    _mark_snapshot_pids_free(monkeypatch)
+    kernel = tmp_path / "kernel.so"
+    kernel.write_bytes(b"compiled kernel")
+    lock = tmp_path / "launcher.lock"
+    lock.write_bytes(b"held by pid 100")
+    manifest = _manifest(
+        external_cache_files=(
+            (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
+            (str(lock), "locked"),
+        )
+    )
+    artifact = _local_restore_artifact(tmp_path)
+    before = {path.name: path.read_bytes() for path in artifact.iterdir()}
+    tools = LocalSnapshotTools()
+    monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
+
+    kernel.write_bytes(b"recompiled kernel")
+    with pytest.raises(SnapshotRestoreError, match="changed since capture") as changed:
+        tools.restore(artifact, manifest)
+    assert str(kernel) in str(changed.value)
+
+    # Restoring the digest gets past the first entry, so the lock file is only
+    # checked for existence: the second failure proves the first one passed.
+    kernel.write_bytes(b"compiled kernel")
+    lock.unlink()
+    with pytest.raises(SnapshotRestoreError, match="vanished since capture") as gone:
+        tools.restore(artifact, manifest)
+    assert str(lock) in str(gone.value)
+
+    assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before
+
+
+def test_snapshot_restore_rejects_a_cache_file_it_cannot_read(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    _mark_snapshot_pids_free(monkeypatch)
+    kernel = tmp_path / "kernel.so"
+    kernel.write_bytes(b"compiled kernel")
+    manifest = _manifest(
+        external_cache_files=(
+            (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()),
+        )
+    )
+    artifact = _local_restore_artifact(tmp_path)
+    tools = LocalSnapshotTools()
+    monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called"))
+    monkeypatch.setattr(tools, "_sha256", _unreadable)
+
+    # A recorded digest that cannot be read again is a failed check, not a file
+    # to wave through on existence.
+    with pytest.raises(
+        SnapshotRestoreError, match="unreadable since capture"
+    ) as blocked:
+        tools.restore(artifact, manifest)
+    assert str(kernel) in str(blocked.value)
+
+
+def test_snapshot_runtime_installer_logs_the_failing_exit_code():
+    installer = Path(__file__).parents[3] / "tools" / "install_snapshot_runtime.sh"
+    if not installer.is_file():
+        pytest.skip("installer script is not shipped in this test image")
+    timing = [
+        line
+        for line in installer.read_text().splitlines()
+        if line.startswith(("_snapshot_install_started=", "trap "))
+    ]
+    assert len(timing) == 2
+
+    result = subprocess.run(
+        ["bash", "-c", "\n".join([*timing, "exit 7"])],
+        capture_output=True,
+        text=True,
+    )
+
+    assert result.returncode == 7
+    assert "(exit 7)" in result.stdout
diff --git a/tests/entrypoints/launchers/test_run_batch.py b/tests/entrypoints/launchers/test_run_batch.py
index 7e6ba3d78d72..f7a250580727 100644
--- a/tests/entrypoints/launchers/test_run_batch.py
+++ b/tests/entrypoints/launchers/test_run_batch.py
@@ -509,9 +509,7 @@ def test_completions():
 
 
 def test_completions_invalid_input():
-    """
-    Ensure that we fail when the input doesn't conform to the openai api.
-    """
+    """Ensure that we fail when the input doesn't conform to the openai api."""
     with (
         tempfile.NamedTemporaryFile("w") as input_file,
         tempfile.NamedTemporaryFile("r") as output_file,
@@ -602,9 +600,7 @@ def test_score(input_batch):
 
 
 def test_reasoning_parser():
-    """
-    Test that reasoning_parser parameter works correctly in run_batch.
-    """
+    """Test that reasoning_parser parameter works correctly in run_batch."""
     with (
         tempfile.NamedTemporaryFile("w") as input_file,
         tempfile.NamedTemporaryFile("r") as output_file,
@@ -767,8 +763,7 @@ def test_translation():
 
 
 def test_tool_calling():
-    """
-    Test that tool calling works correctly in run_batch.
+    """Test that tool calling works correctly in run_batch.
     Verifies that requests with tools return tool_calls in the response.
     """
     with (
diff --git a/tests/entrypoints/launchers/test_shutdown.py b/tests/entrypoints/launchers/test_shutdown.py
index 643eff7594de..6d18d825eba9 100644
--- a/tests/entrypoints/launchers/test_shutdown.py
+++ b/tests/entrypoints/launchers/test_shutdown.py
@@ -134,7 +134,6 @@ async def test_shutdown_on_engine_failure(tmp_path: Path):
     Starts a vLLM server, kills it to simulate a crash, then verifies that
     subsequent API calls fail appropriately.
     """
-
     port = get_open_port()
 
     # Redirect server output to a file rather than a pipe: nothing drains the
diff --git a/tests/entrypoints/llm/offline_mode/test_offline_mode.py b/tests/entrypoints/llm/offline_mode/test_offline_mode.py
index 555e8b9b7d07..ab7edd019ba9 100644
--- a/tests/entrypoints/llm/offline_mode/test_offline_mode.py
+++ b/tests/entrypoints/llm/offline_mode/test_offline_mode.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for HF_HUB_OFFLINE mode"""
+"""Tests for HF_HUB_OFFLINE mode."""
 
 import importlib
 import sys
diff --git a/tests/entrypoints/llm/test_accuracy.py b/tests/entrypoints/llm/test_accuracy.py
index af607720c8b0..49d83b9cf6e8 100644
--- a/tests/entrypoints/llm/test_accuracy.py
+++ b/tests/entrypoints/llm/test_accuracy.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file test accuracy of the vLLM server via LMEval.
+"""This file test accuracy of the vLLM server via LMEval.
 It uses local-completions, which interacts with vLLM
 through the OAI API with N concurrent connections.
 This simulates real work usage of the API and makes
@@ -33,7 +32,6 @@
 
 def run_test(model_name, more_args=None):
     """Run the end to end accuracy test."""
-
     model_args = f"pretrained={model_name},max_model_len=4096"
 
     if more_args is not None:
@@ -64,7 +62,6 @@ def run_test(model_name, more_args=None):
 @pytest.mark.parametrize("model", MODEL_NAMES)
 def test_lm_eval_accuracy_v1_engine(model):
     """Run with the V1 Engine."""
-
     more_args = None
     if current_platform.is_tpu():
         # Limit compilation time for TPU V1
@@ -81,7 +78,6 @@ def test_lm_eval_accuracy_v1_engine(model):
 @pytest.mark.parametrize("model", FP8_KV_MODEL_NAMES)
 def test_lm_eval_accuracy_v1_engine_fp8_kv_cache(model):
     """Run with the V1 Engine."""
-
     more_args = None
     if current_platform.is_tpu():
         # Limit compilation time for TPU V1
diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py
index 6ee8d34ce138..276426e692c7 100644
--- a/tests/entrypoints/llm/test_chat.py
+++ b/tests/entrypoints/llm/test_chat.py
@@ -20,8 +20,7 @@ def text_llm(vllm_runner):
 
 @pytest.fixture(scope="function")
 def llm_for_failure_test(vllm_runner):
-    """
-    Fixture for testing issue #26081.
+    """Fixture for testing issue #26081.
     Uses a small max_model_len to easily trigger length errors.
     """
     with vllm_runner(
@@ -67,8 +66,7 @@ def test_multi_chat(text_llm):
 
 
 def test_llm_chat_tokenization_no_double_bos(text_llm):
-    """
-    LLM.chat() should not add special tokens when using chat templates.
+    """LLM.chat() should not add special tokens when using chat templates.
     Check we get a single BOS token for llama chat.
     """
     messages = [
@@ -127,8 +125,7 @@ def test_chat_extra_kwargs(thinking_llm, enable_thinking):
 
 
 def test_chat_batch_failure_cleanup(llm_for_failure_test):
-    """
-    Tests that if a batch call to llm.chat() fails mid-way
+    """Tests that if a batch call to llm.chat() fails mid-way
     (e.g., due to one invalid prompt), the requests that
     were already enqueued are properly aborted and do not
     pollute the queue for subsequent calls.
diff --git a/tests/entrypoints/llm/test_struct_output_generate.py b/tests/entrypoints/llm/test_struct_output_generate.py
index 0f55ad7adc3e..030b0207e4c7 100644
--- a/tests/entrypoints/llm/test_struct_output_generate.py
+++ b/tests/entrypoints/llm/test_struct_output_generate.py
@@ -948,7 +948,7 @@ def test_structured_output_batched_with_non_structured_outputs_requests(
         prompts = [structured_outputs_prompt, non_structured_outputs_prompt]
         sampling_params = [
             SamplingParams(
-                temperature=1.0,
+                temperature=0,
                 max_tokens=400,
                 structured_outputs=StructuredOutputsParams(json=sample_json_schema),
             ),
diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py
index 5c378de35dbe..a3511b9528b1 100644
--- a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py
+++ b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test that ``InputProcessor.inject_into_mm_cache()`` correctly injects
+"""Test that ``InputProcessor.inject_into_mm_cache()`` correctly injects
 pre-processed mm_kwargs into the processor cache and reports MM cache
 hit rate metrics accurately.
 
diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py b/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py
index 14ac7ce0439b..5031cbea1a25 100644
--- a/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py
+++ b/tests/entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py
@@ -47,8 +47,7 @@ async def client(server):
 async def test_online_audio_in_video(
     client: openai.AsyncOpenAI, video_assets: VideoTestAssets
 ):
-    """Test video input with `audio_in_video=True`"""
-
+    """Test video input with `audio_in_video=True`."""
     # we don't use video_urls above because they missed audio stream.
     video_path = video_assets[0].video_path
     with open(video_path, "rb") as f:
@@ -97,8 +96,7 @@ async def test_online_audio_in_video(
 async def test_online_audio_in_video_multi_videos(
     client: openai.AsyncOpenAI, video_assets: VideoTestAssets
 ):
-    """Test multi-video input with `audio_in_video=True`"""
-
+    """Test multi-video input with `audio_in_video=True`."""
     # we don't use video_urls above because they missed audio stream.
     video_path = video_assets[0].video_path
     with open(video_path, "rb") as f:
@@ -151,8 +149,7 @@ async def test_online_audio_in_video_multi_videos(
 async def test_online_audio_in_video_interleaved(
     client: openai.AsyncOpenAI, video_assets: VideoTestAssets
 ):
-    """Test interleaved video/audio input with `audio_in_video=True`"""
-
+    """Test interleaved video/audio input with `audio_in_video=True`."""
     # we don't use video_urls above because they missed audio stream.
     video_path = video_assets[0].video_path
     with open(video_path, "rb") as f:
diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py
index e6b8dc23e9f2..56a7cf6a1312 100644
--- a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py
+++ b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py
@@ -34,8 +34,7 @@
 
 
 def check_output_matches_terms(content: str, term_groups: list[list[str]]) -> bool:
-    """
-    Check if content matches all required term groups.
+    """Check if content matches all required term groups.
     Each term group requires at least one of its terms to be present.
     All term groups must be satisfied.
     """
diff --git a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py
index 20218ac6c875..467e23e7d3d2 100644
--- a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py
+++ b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py
@@ -259,7 +259,7 @@ async def _generator(prompt_idx: int, text: str) -> AsyncGenerator[RequestOutput
 @pytest.mark.asyncio
 @pytest.mark.skip_global_cleanup
 async def test_batched_echo_does_not_prepend_user_prompt() -> None:
-    """echo must only echo the assistant turn, never the user's prompt.
+    """Echo must only echo the assistant turn, never the user's prompt.
 
     With `add_generation_prompt` (the default) the response role is `assistant`
     while the last conversation message is the user's, so there is nothing to
diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py
index aafb899f005c..a1665a71de6d 100644
--- a/tests/entrypoints/openai/chat_completion/test_chat_error.py
+++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py
@@ -56,9 +56,10 @@ class MockModelConfig:
     encoder_config = None
     generation_config: str = "auto"
     media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
+    skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
@@ -121,7 +122,7 @@ async def _fake_preprocess_chat(*args, **kwargs):
 
 @pytest.mark.asyncio
 async def test_chat_error_non_stream():
-    """test finish_reason='error' returns 500 InternalServerError (non-streaming)"""
+    """Test finish_reason='error' returns 500 InternalServerError (non-streaming)."""
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
     mock_engine.model_config = MockModelConfig()
@@ -246,7 +247,7 @@ async def test_renderer_only_chat_request_skips_mm_cache():
 
 @pytest.mark.asyncio
 async def test_chat_error_stream():
-    """test finish_reason='error' returns 500 InternalServerError (streaming)"""
+    """Test finish_reason='error' returns 500 InternalServerError (streaming)."""
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
     mock_engine.model_config = MockModelConfig()
diff --git a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py
index 3e59e7402ea4..61fd2f1576bc 100644
--- a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py
+++ b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py
@@ -451,7 +451,7 @@ async def test_inconsistent_tool_choice_and_tools(
 async def test_max_tokens_with_tool_choice_required(
     client: openai.AsyncOpenAI, tool_choice
 ):
-    """ """
+    """Requesting a single token with a required tool choice must not crash."""
     models = await client.models.list()
     model_name: str = models.data[0].id
 
diff --git a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py
index 8c03fc1daa01..f4036c4374f4 100644
--- a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py
+++ b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py
@@ -52,7 +52,7 @@ def stop_sign_image_url():
 
 @pytest.mark.asyncio
 async def test_translategemma_extra_lang_code_fields(client: openai.AsyncOpenAI):
-    """en -> es translation through TranslateGemma's bundled chat template's
+    """En -> es translation through TranslateGemma's bundled chat template's
     ``text`` content branch, which depends on ``source_lang_code`` /
     ``target_lang_code`` extra fields being preserved on the content part."""
     completion = await client.chat.completions.create(
@@ -92,7 +92,7 @@ async def test_translategemma_image_extra_lang_code_fields(
     client: openai.AsyncOpenAI,
     stop_sign_image_url: str,
 ):
-    """en -> es OCR-translation through TranslateGemma's bundled chat
+    """En -> es OCR-translation through TranslateGemma's bundled chat
     template's ``image`` content branch. Exercises the multimodal branch of
     ``_collect_extra_fields`` in ``_parse_chat_message_content_part``: the
     ``source_lang_code`` / ``target_lang_code`` extras must survive parsing
diff --git a/tests/entrypoints/openai/chat_completion/test_include_reasoning.py b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py
index 1d4e1ebc4055..5eeca438fe64 100644
--- a/tests/entrypoints/openai/chat_completion/test_include_reasoning.py
+++ b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py
@@ -45,7 +45,8 @@ async def test_include_reasoning_true_non_streaming(client: openai.AsyncOpenAI):
     response = await client.chat.completions.create(
         model=MODEL_NAME,
         messages=MESSAGES,
-        max_tokens=200,
+        max_tokens=512,
+        temperature=0,
         extra_body={"include_reasoning": True},
     )
 
@@ -63,7 +64,8 @@ async def test_include_reasoning_false_non_streaming(client: openai.AsyncOpenAI)
     response = await client.chat.completions.create(
         model=MODEL_NAME,
         messages=MESSAGES,
-        max_tokens=200,
+        max_tokens=512,
+        temperature=0,
         extra_body={"include_reasoning": False},
     )
 
@@ -99,7 +101,8 @@ async def test_include_reasoning_true_streaming(client: openai.AsyncOpenAI):
     stream = await client.chat.completions.create(
         model=MODEL_NAME,
         messages=MESSAGES,
-        max_tokens=200,
+        max_tokens=512,
+        temperature=0,
         stream=True,
         extra_body={"include_reasoning": True},
     )
@@ -130,7 +133,8 @@ async def test_include_reasoning_false_streaming(client: openai.AsyncOpenAI):
     stream = await client.chat.completions.create(
         model=MODEL_NAME,
         messages=MESSAGES,
-        max_tokens=200,
+        max_tokens=512,
+        temperature=0,
         stream=True,
         extra_body={"include_reasoning": False},
     )
diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py
index e7ea15877c22..6f580a9c28d6 100644
--- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py
+++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py
@@ -556,6 +556,7 @@ class MockModelConfig:
     skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
     enable_prompt_embeds: bool = False
 
@@ -1232,7 +1233,6 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated():
     """Regression test: when the Mistral tokenizer path returns token IDs
     directly, we must still apply input length + max_tokens validation.
     """
-
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
     mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True)
@@ -1271,7 +1271,6 @@ async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected():
     """Regression test: MistralTokenizer token-id prompts must still enforce
     the max context length for the input itself (token_num >= max_model_len).
     """
-
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
     mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True)
@@ -1512,8 +1511,7 @@ async def _render_chat_prompt_token_ids(
 
 
 class TestServingChatWithHarmony:
-    """
-    These tests ensure Chat Completion requests are being properly converted into
+    """These tests ensure Chat Completion requests are being properly converted into
     Harmony messages and Harmony response messages back into Chat Completion responses.
     These tests are not exhaustive, but each one was created to cover a specific case
     that we got wrong but is now fixed.
@@ -2241,7 +2239,13 @@ async def test_tool_choice_validation_without_parser():
 
 
 @pytest.mark.asyncio
-async def test_streaming_n_gt1_independent_tool_parsers():
+@pytest.mark.parametrize(
+    ("engine_finish_reason", "expected_finish_reason"),
+    [("stop", "tool_calls"), ("length", "length")],
+)
+async def test_streaming_n_gt1_independent_tool_parsers(
+    engine_finish_reason: str, expected_finish_reason: str
+):
     """n>1 streaming must use independent parser instances
     and token-id histories per choice.
     """
@@ -2346,7 +2350,7 @@ async def result_generator():
                     token_ids=[],
                     cumulative_logprob=0.0,
                     logprobs=None,
-                    finish_reason="stop",
+                    finish_reason=engine_finish_reason,
                 )
                 for choice_idx in range(num_choices)
             ],
@@ -2410,8 +2414,8 @@ async def result_generator():
         assert len(reasons) == 1, (
             f"Choice {choice_idx}: expected exactly 1 finish_reason, got {reasons}"
         )
-        assert reasons[0] == "tool_calls", (
-            f"Choice {choice_idx}: expected finish_reason='tool_calls', "
+        assert reasons[0] == expected_finish_reason, (
+            f"Choice {choice_idx}: expected finish_reason={expected_finish_reason!r}, "
             f"got '{reasons[0]}'"
         )
 
diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
index ec7a977f1a77..ecff3f8ee6be 100644
--- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
+++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
@@ -33,7 +33,6 @@ def _count_reasoning_decode_token_ids_between_markers(
     reasoning_end_ids: list[int],
 ) -> int | None:
     """Count decode tokens in the thinking span (after last start, before first end)."""
-
     if not reasoning_start_ids or not reasoning_end_ids:
         raise ValueError("reasoning marker token id lists must be non-empty")
 
@@ -153,7 +152,6 @@ async def client(request, server, server_with_auto_reasoning_config):
 async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):
     """Test that mixed requests (some with thinking_token_budget, some without)
     complete successfully without errors."""
-
     response_with_budget = await client.chat.completions.create(
         model=MODEL_NAME,
         messages=MESSAGES,
@@ -182,7 +180,6 @@ async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI
     grouped into streamed chunks (a single chunk can carry several tokens under
     async scheduling / stream_interval > 1). Counting chunks under-counts.
     """
-
     tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
     start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False))
     end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False))
@@ -227,7 +224,6 @@ async def test_thinking_token_budget_qwen35_fp8_mtp_concurrent_mixed_budget_and_
     Qwen3.5 FP8 + MTP (TP=2) server. Budgeted calls are checked with
     ``_count_reasoning_decode_token_ids_between_markers`` on full token ids.
     """
-
     _batch_spec: list[tuple[Literal["budget"], int] | tuple[Literal["plain"], None]] = [
         ("budget", 1),
         ("budget", 12),
diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py
index af9f1def3614..287664224123 100644
--- a/tests/entrypoints/openai/completion/test_completion.py
+++ b/tests/entrypoints/openai/completion/test_completion.py
@@ -350,10 +350,9 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str
     """Parallel sampling without streaming.
     A single request output contains a list of completions.
     """
-
     prompt = "What is an LLM?"
     n = 3
-    max_tokens = 50  # we want some to finish earlier than others
+    max_tokens = 50
 
     # High temperature to maximize chance of unique completions.
     completion = await client.completions.create(
@@ -371,16 +370,12 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str
     num_completions = len(completion.choices)
     assert num_completions == n, f"Num completions {num_completions} but expected {n}."
     completion_repeats: dict[str, int] = {}
-    output_token_lengths = set()
     for idx, choice in enumerate(completion.choices):
         # Assert correct completion index & some finish reason.
         assert choice.index == idx, f"Index {choice.index} but expected {idx}."
         assert choice.finish_reason is not None, "None finish_reason is invalid."
         text = choice.text
         completion_repeats[text] = completion_repeats.get(text, 0) + 1
-        output_token_lengths.add(len(choice.logprobs.tokens))
-    # Assert subrequests finished at different times
-    assert len(output_token_lengths) > 1
     # Assert `n` unique completions
     num_unique = len(completion_repeats)
     if num_unique != n:
@@ -400,10 +395,9 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):
     The tokens from multiple samples, are flattened into a single stream,
     with an index to indicate which sample the token belongs to.
     """
-
     prompt = "What is an LLM?"
     n = 3
-    max_tokens = 50  # we want some to finish earlier than others
+    max_tokens = 50
 
     stream = await client.completions.create(
         model=model_name,
@@ -427,19 +421,15 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):
         f"Expected {n} completions with valid indices and finish_reason."
     )
     completion_repeats: dict[str, int] = {}
-    chunk_lengths = set()
     for chunk in chunks:
         chunk_len = len(chunk)
         # Assert correct number of completion tokens
-        chunk_lengths.add(chunk_len)
         assert chunk_len <= max_tokens, (
             f"max_tokens={max_tokens} but chunk len is {chunk_len}."
         )
         text = "".join(chunk)
         completion_repeats[text] = completion_repeats.get(text, 0) + 1
         print(text)
-    # Assert subrequests finished at different times
-    assert len(chunk_lengths) > 1
     # Assert `n` unique completions
     num_unique = len(completion_repeats)
     if num_unique != n:
diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py
index 513bde465d71..060ae01a1d6f 100644
--- a/tests/entrypoints/openai/completion/test_completion_error.py
+++ b/tests/entrypoints/openai/completion/test_completion_error.py
@@ -14,8 +14,10 @@
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.scale_out.render.serving import ServingRender
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.exceptions import GenerationError, VLLMValidationError
 from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.renderers.deepseek_v4 import DeepseekV4Renderer
 from vllm.renderers.hf import HfRenderer
 from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.tokenizers.registry import cached_tokenizer_from_config
@@ -24,6 +26,9 @@
 
 MODEL_NAME = "openai-community/gpt2"
 MODEL_NAME_SHORT = "gpt2"
+DEEPSEEK_V4_FIM_BEGIN = "<｜fim▁begin｜>"
+DEEPSEEK_V4_FIM_HOLE = "<｜fim▁hole｜>"
+DEEPSEEK_V4_FIM_END = "<｜fim▁end｜>"
 _PER_REQUEST_STATS = RequestStateStats(
     queued_ts=1.0,
     scheduled_ts=1.5,
@@ -61,9 +66,10 @@ class MockModelConfig:
     encoder_config = None
     generation_config: str = "auto"
     media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
+    skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
@@ -301,9 +307,182 @@ def test_completion_metrics_carries_both_timing_and_spec_decode():
     assert response.metrics.speculative_decoding.num_spec_steps == 2
 
 
+def _build_online_renderer_for_completion(
+    model_config: MockModelConfig,
+    renderer: Any | None = None,
+) -> OnlineRenderer:
+    if renderer is None:
+        renderer = MagicMock()
+        renderer.render_completion_suffix.return_value = None
+
+    return OnlineRenderer(
+        model_config=model_config,
+        renderer=renderer,
+        request_logger=None,
+        chat_template=None,
+        chat_template_content_format="auto",
+    )
+
+
+def _deepseek_v4_suffix_renderer(
+    model_config: MockModelConfig,
+) -> DeepseekV4Renderer:
+    return DeepseekV4Renderer(
+        MockVllmConfig(model_config, parallel_config=MockParallelConfig()),
+        None,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip_global_cleanup
+async def test_completion_suffix_uses_renderer_capability():
+    model_config = MockModelConfig()
+    online_renderer = _build_online_renderer_for_completion(
+        model_config,
+        renderer=_deepseek_v4_suffix_renderer(model_config),
+    )
+    online_renderer.preprocess_completion = AsyncMock(return_value=[{"ok": True}])
+
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt="def fib(n):\n    return ",
+        suffix="\n\nprint(fib(10))",
+        max_tokens=64,
+    )
+
+    result = await online_renderer.render_completion(request)
+
+    assert result == [{"ok": True}]
+    online_renderer.preprocess_completion.assert_awaited_once()
+    assert online_renderer.preprocess_completion.call_args.kwargs["prompt_input"] == (
+        f"{DEEPSEEK_V4_FIM_BEGIN}def fib(n):\n    return "
+        f"{DEEPSEEK_V4_FIM_HOLE}\n\nprint(fib(10))"
+        f"{DEEPSEEK_V4_FIM_END}"
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip_global_cleanup
+async def test_completion_suffix_supports_text_prompt_list():
+    model_config = MockModelConfig()
+    online_renderer = _build_online_renderer_for_completion(
+        model_config,
+        renderer=_deepseek_v4_suffix_renderer(model_config),
+    )
+    online_renderer.preprocess_completion = AsyncMock(return_value=[{"ok": True}])
+
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt=["prefix A", "prefix B"],
+        suffix=" suffix",
+    )
+
+    result = await online_renderer.render_completion(request)
+
+    assert result == [{"ok": True}]
+    assert online_renderer.preprocess_completion.call_args.kwargs["prompt_input"] == [
+        f"{DEEPSEEK_V4_FIM_BEGIN}prefix A"
+        f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}",
+        f"{DEEPSEEK_V4_FIM_BEGIN}prefix B"
+        f"{DEEPSEEK_V4_FIM_HOLE} suffix{DEEPSEEK_V4_FIM_END}",
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip_global_cleanup
+async def test_completion_suffix_rejects_renderer_without_fim_support():
+    model_config = MockModelConfig()
+    online_renderer = _build_online_renderer_for_completion(model_config)
+    online_renderer.preprocess_completion = AsyncMock()
+
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt="prefix",
+        suffix="suffix",
+    )
+
+    result = await online_renderer.render_completion(request)
+
+    assert isinstance(result, ErrorResponse)
+    assert "FIM completion rendering" in result.error.message
+    online_renderer.preprocess_completion.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip_global_cleanup
+async def test_completion_suffix_rejects_echo():
+    model_config = MockModelConfig()
+    online_renderer = _build_online_renderer_for_completion(
+        model_config,
+        renderer=_deepseek_v4_suffix_renderer(model_config),
+    )
+    online_renderer.preprocess_completion = AsyncMock()
+
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt="prefix",
+        suffix="suffix",
+        echo=True,
+    )
+
+    result = await online_renderer.render_completion(request)
+
+    assert isinstance(result, ErrorResponse)
+    assert "Echo is unsupported with suffix" in result.error.message
+    online_renderer.preprocess_completion.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip_global_cleanup
+async def test_completion_suffix_rejects_prompt_embeds():
+    model_config = MockModelConfig()
+    online_renderer = _build_online_renderer_for_completion(
+        model_config,
+        renderer=_deepseek_v4_suffix_renderer(model_config),
+    )
+    online_renderer.preprocess_completion = AsyncMock()
+
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt="prefix",
+        suffix="suffix",
+        prompt_embeds=b"embeds",
+    )
+
+    result = await online_renderer.render_completion(request)
+
+    assert isinstance(result, ErrorResponse)
+    assert "suffix is not supported with prompt_embeds" in result.error.message
+    online_renderer.preprocess_completion.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize("prompt", ([1, 2, 3], [[1, 2, 3]]))
+async def test_completion_suffix_rejects_token_prompt(prompt):
+    model_config = MockModelConfig()
+    online_renderer = _build_online_renderer_for_completion(
+        model_config,
+        renderer=_deepseek_v4_suffix_renderer(model_config),
+    )
+    online_renderer.preprocess_completion = AsyncMock()
+
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt=prompt,
+        suffix="suffix",
+    )
+
+    result = await online_renderer.render_completion(request)
+
+    assert isinstance(result, ErrorResponse)
+    assert "requires text prompt input" in result.error.message
+    online_renderer.preprocess_completion.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_completion_error_non_stream():
-    """test finish_reason='error' returns 500 InternalServerError (non-streaming)"""
+    """Test finish_reason='error' returns 500 InternalServerError (non-streaming)."""
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
     mock_engine.model_config = MockModelConfig()
@@ -439,7 +618,7 @@ async def test_renderer_only_completion_request_skips_mm_cache():
 
 @pytest.mark.asyncio
 async def test_completion_error_stream():
-    """test finish_reason='error' returns 500 InternalServerError (streaming)"""
+    """Test finish_reason='error' returns 500 InternalServerError (streaming)."""
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
     mock_engine.model_config = MockModelConfig()
diff --git a/tests/entrypoints/openai/completion/test_lora_resolvers.py b/tests/entrypoints/openai/completion/test_lora_resolvers.py
index 4d8c3a1fe2a0..8cd874996cd8 100644
--- a/tests/entrypoints/openai/completion/test_lora_resolvers.py
+++ b/tests/entrypoints/openai/completion/test_lora_resolvers.py
@@ -54,6 +54,7 @@ class MockModelConfig:
     skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
diff --git a/tests/entrypoints/openai/completion/test_token_in_token_out.py b/tests/entrypoints/openai/completion/test_token_in_token_out.py
index 8882ae624428..45024d827740 100644
--- a/tests/entrypoints/openai/completion/test_token_in_token_out.py
+++ b/tests/entrypoints/openai/completion/test_token_in_token_out.py
@@ -39,8 +39,7 @@ def server():
 
 @pytest.mark.asyncio
 async def test_token_in_token_out_and_logprobs(server):
-    """
-    Test token-in-token-out and token_ids align with prompt_logprobs
+    """Test token-in-token-out and token_ids align with prompt_logprobs
     & logprobs when return_tokens_as_token_ids is enabled.
     """
     tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
diff --git a/tests/entrypoints/openai/correctness/test_lmeval.py b/tests/entrypoints/openai/correctness/test_lmeval.py
index aad1b5e0624e..c6ce3ebc2b21 100644
--- a/tests/entrypoints/openai/correctness/test_lmeval.py
+++ b/tests/entrypoints/openai/correctness/test_lmeval.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file test accuracy of the vLLM server via LMEval.
+"""This file test accuracy of the vLLM server via LMEval.
 It uses local-completions, which interacts with vLLM
 through the OAI API with N concurrent connections.
 This simulates real work usage of the API and makes
@@ -37,7 +36,6 @@
 
 def run_test(more_args):
     """Run the end to end accuracy test."""
-
     args = list(DEFAULT_ARGS)
     args.extend(more_args)
     print(f"Running with: {args}")
@@ -68,7 +66,6 @@ def run_test(more_args):
 
 def test_lm_eval_accuracy_v1_engine():
     """Run with the V1 Engine."""
-
     more_args = []
 
     # Limit compilation time for V1 on TPU
diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py
index 1b19714ac222..4c1dfe52e033 100644
--- a/tests/entrypoints/openai/parser/test_harmony_utils.py
+++ b/tests/entrypoints/openai/parser/test_harmony_utils.py
@@ -170,7 +170,7 @@ class TestIsFunctionRecipientWithAllowedNames:
     """Tests for is_function_recipient with allowed_function_tool_names."""
 
     def test_prefixed_always_accepted(self):
-        """functions. prefix is always accepted regardless of allowed names."""
+        """Functions. prefix is always accepted regardless of allowed names."""
         fn_names = frozenset({"other_tool"})
         assert is_function_recipient("functions.get_weather", fn_names) is True
 
@@ -252,8 +252,7 @@ def test_dotted_function_name_extraction(self, recipient, expected):
 
 
 class TestCommonParseInputToHarmonyMessage:
-    """
-    Tests for scenarios that are common to both Chat Completion
+    """Tests for scenarios that are common to both Chat Completion
     parse_chat_input_to_harmony_message and Responses API
     response_previous_input_to_harmony functions.
     """
@@ -448,8 +447,7 @@ def test_array_content_with_missing_text(self, parse_function):
 
 
 class TestParseChatInputToHarmonyMessage:
-    """
-    Tests for scenarios that are specific to the Chat Completion API
+    """Tests for scenarios that are specific to the Chat Completion API
     parse_chat_input_to_harmony_message function.
     """
 
diff --git a/tests/entrypoints/openai/responses/test_errors.py b/tests/entrypoints/openai/responses/test_errors.py
index 0ea7833da9a7..1e0182ba6b8c 100644
--- a/tests/entrypoints/openai/responses/test_errors.py
+++ b/tests/entrypoints/openai/responses/test_errors.py
@@ -14,7 +14,7 @@
 
 @pytest.mark.asyncio
 async def test_raise_if_error_raises_generation_error():
-    """test _raise_if_error raises GenerationError"""
+    """Test _raise_if_error raises GenerationError."""
     # create a minimal GenerateBaseServing instance
     mock_engine = MagicMock()
     mock_engine.model_config = MagicMock()
@@ -42,7 +42,7 @@ async def test_raise_if_error_raises_generation_error():
 
 @pytest.mark.asyncio
 async def test_convert_generation_error_to_streaming_response():
-    """test _convert_generation_error_to_streaming_response output"""
+    """Test _convert_generation_error_to_streaming_response output."""
     mock_engine = MagicMock()
     mock_engine.model_config = MagicMock()
     mock_engine.model_config.max_model_len = 100
diff --git a/tests/entrypoints/openai/responses/test_function_call.py b/tests/entrypoints/openai/responses/test_function_call.py
index 9dcbd74c890b..006ee0f6f24e 100644
--- a/tests/entrypoints/openai/responses/test_function_call.py
+++ b/tests/entrypoints/openai/responses/test_function_call.py
@@ -168,8 +168,7 @@ async def test_max_tokens_with_tool_choice_required(
 @pytest.mark.asyncio
 async def test_named_tool_use(client: openai.AsyncOpenAI):
     def get_weather(latitude: float, longitude: float) -> str:
-        """
-        Mock function to simulate getting weather data.
+        """Mock function to simulate getting weather data.
         In a real application, this would call an external weather API.
         """
         return f"Current temperature at ({latitude}, {longitude}) is 20°C."
diff --git a/tests/entrypoints/openai/responses/test_function_call_parsing.py b/tests/entrypoints/openai/responses/test_function_call_parsing.py
index f90a641db7fd..5b70c7819a0e 100644
--- a/tests/entrypoints/openai/responses/test_function_call_parsing.py
+++ b/tests/entrypoints/openai/responses/test_function_call_parsing.py
@@ -57,7 +57,6 @@ def test_direct_function_call_object_preservation():
 
 def test_mixed_input_types_with_function_calls():
     """Test parsing with mixed input types including function calls."""
-
     request_data = {
         "model": "gpt-oss",
         "input": [
@@ -177,7 +176,6 @@ def test_empty_list_input():
 
 def test_function_call_output_not_affected():
     """Test that FunctionCallOutput is not affected by the function_call parsing."""
-
     # Test with FunctionCallOutput as dict (should not be parsed)
     request_data = {
         "model": "gpt-oss",
@@ -265,7 +263,6 @@ def test_function_call_validation_failure_logs_debug(caplog):
 
 def test_validator_handles_iterator_input():
     """Test that validator can handle ValidatorIterator input (Pydantic internal)."""
-
     # This test simulates when Pydantic passes a ValidatorIterator instead of a list
     # This happened with complex nested structures containing reasoning + function_call
 
diff --git a/tests/entrypoints/openai/responses/test_harmony_utils.py b/tests/entrypoints/openai/responses/test_harmony_utils.py
index 14e174dec0bb..b524842467ab 100644
--- a/tests/entrypoints/openai/responses/test_harmony_utils.py
+++ b/tests/entrypoints/openai/responses/test_harmony_utils.py
@@ -20,8 +20,7 @@
 
 
 class TestResponsePreviousInputToHarmony:
-    """
-    Tests for scenarios that are specific to the Responses API
+    """Tests for scenarios that are specific to the Responses API
     response_previous_input_to_harmony function.
     """
 
@@ -273,6 +272,23 @@ def test_browser_search_recipient_respects_incomplete(self, incomplete):
         assert output_items[0].action.type == "search"
         assert output_items[0].action.query == "cursor:weather in San Francisco"
 
+    def test_browser_find_recipient_uses_responses_action_type(self):
+        """browser.find must use the Responses API find_in_page action type."""
+        message = (
+            Message.from_role_and_content(
+                Role.ASSISTANT, '{"pattern": "vLLM", "cursor": 42}'
+            )
+            .with_channel("analysis")
+            .with_recipient("browser.find")
+        )
+
+        output_items = harmony_to_response_output(message, frozenset())
+
+        assert len(output_items) == 1
+        assert isinstance(output_items[0], ResponseFunctionWebSearch)
+        assert output_items[0].action.type == "find_in_page"
+        assert output_items[0].action.pattern == "vLLM"
+
     def test_commentary_with_empty_content_and_no_recipient(self):
         """Test edge case: empty commentary with recipient=None."""
         message = Message.from_role_and_content(Role.ASSISTANT, "")
diff --git a/tests/entrypoints/openai/responses/test_parsable_context_unit.py b/tests/entrypoints/openai/responses/test_parsable_context_unit.py
index 01293f62b222..85bf6f2020a1 100644
--- a/tests/entrypoints/openai/responses/test_parsable_context_unit.py
+++ b/tests/entrypoints/openai/responses/test_parsable_context_unit.py
@@ -36,9 +36,6 @@ class _NoOpParser(DelegatingParser):
     def is_reasoning_end(self, input_ids: list[int]) -> bool:
         return False
 
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        return input_ids
-
     def extract_reasoning(self, model_output, request):
         return None, model_output
 
@@ -63,9 +60,6 @@ class _ReasoningOnlyParser(DelegatingParser):
     def is_reasoning_end(self, input_ids: list[int]) -> bool:
         return False
 
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        return input_ids
-
     def extract_reasoning(self, model_output, request):
         if "<think>" in model_output and "</think>" in model_output:
             start = model_output.index("<think>") + len("<think>")
@@ -131,9 +125,6 @@ def __init__(self, tokenizer, *args, **kwargs):
     def is_reasoning_end(self, input_ids: list[int]) -> bool:
         return False
 
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        return input_ids
-
     def extract_reasoning(self, model_output, request):
         return None, model_output
 
diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py
index 041e517e29dd..1a315902e249 100644
--- a/tests/entrypoints/openai/responses/test_serving_responses.py
+++ b/tests/entrypoints/openai/responses/test_serving_responses.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from contextlib import AsyncExitStack
+from contextlib import AsyncExitStack, asynccontextmanager
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
 
@@ -18,6 +18,9 @@
     ResponseTextConfig,
     ResponseTextDeltaEvent,
 )
+from openai.types.responses import (
+    ResponseUsage as OpenAIResponseUsage,
+)
 from openai.types.responses.response_format_text_json_schema_config import (
     ResponseFormatTextJSONSchemaConfig,
 )
@@ -41,9 +44,11 @@
 from vllm.entrypoints.openai.responses.context import (
     ConversationContext,
     HarmonyContext,
+    ParsableContext,
     SimpleContext,
 )
 from vllm.entrypoints.openai.responses.protocol import (
+    ResponseCompletedEvent,
     ResponseCreatedEvent,
     ResponseRawMessageAndToken,
     ResponsesRequest,
@@ -67,6 +72,7 @@
     _extract_allowed_tools_from_mcp_requests,
 )
 from vllm.sampling_params import SamplingParams
+from vllm.v1.metrics.stats import RequestStateStats
 
 pytestmark = pytest.mark.skip_global_cleanup
 
@@ -80,7 +86,7 @@ def _new_online_renderer() -> OnlineRenderer:
 
 
 class MockConversationContext(ConversationContext):
-    """Mock conversation context for testing"""
+    """Mock conversation context for testing."""
 
     def __init__(self):
         self.init_tool_sessions_called = False
@@ -122,7 +128,7 @@ def test_serialize_message_pydantic_model_returns_dict() -> None:
 
 @pytest.fixture
 def mock_serving_responses():
-    """Create a mock OpenAIServingResponses instance"""
+    """Create a mock OpenAIServingResponses instance."""
     serving_responses = MagicMock(spec=OpenAIServingResponses)
     serving_responses.tool_server = MagicMock(spec=ToolServer)
     return serving_responses
@@ -130,13 +136,13 @@ def mock_serving_responses():
 
 @pytest.fixture
 def mock_context():
-    """Create a mock conversation context"""
+    """Create a mock conversation context."""
     return MockConversationContext()
 
 
 @pytest.fixture
 def mock_exit_stack():
-    """Create a mock async exit stack"""
+    """Create a mock async exit stack."""
     return MagicMock(spec=AsyncExitStack)
 
 
@@ -549,11 +555,92 @@ def test_responses_render_result_rejects_non_single_prompt(engine_inputs):
 
 
 class TestInitializeToolSessions:
-    """Test class for _initialize_tool_sessions method"""
+    """Test class for _initialize_tool_sessions method."""
+
+    @pytest.fixture(params=[ParsableContext, HarmonyContext])
+    def tool_context(self, request):
+        context = request.param.__new__(request.param)
+        context.available_tools = ["browser", "python"]
+        context._tool_sessions = {}
+        context.called_tools = set()
+        return context
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("available_tools", "called_tools"),
+        [
+            ([], []),
+            (["browser"], ["browser"]),
+            (["browser", "python"], []),
+            (["browser", "python"], ["browser"]),
+            (["browser", "python"], ["python"]),
+            (["browser", "python"], ["browser", "python"]),
+        ],
+        ids=["empty", "single", "unused", "first", "last", "both"],
+    )
+    async def test_mcp_sessions_cleaned_once_before_close(
+        self, tool_context, available_tools, called_tools
+    ):
+        events = []
+        sessions = {}
+
+        @asynccontextmanager
+        async def new_session(name, *_args):
+            session = MagicMock()
+            session.call_tool = AsyncMock(
+                side_effect=lambda *_: events.append(("cleanup", name))
+            )
+            sessions[name] = session
+            try:
+                yield session
+            finally:
+                events.append(("close", name))
+
+        tool_server = MagicMock(spec=ToolServer)
+        tool_server.new_session.side_effect = new_session
+        tool_context.available_tools = available_tools
+        async with AsyncExitStack() as stack:
+            await tool_context.init_tool_sessions(tool_server, stack, "req", {})
+            # Reinitializing existing sessions must not duplicate cleanup either.
+            await tool_context.init_tool_sessions(tool_server, stack, "req", {})
+            tool_context.called_tools.update(called_tools)
+
+        for name, session in sessions.items():
+            if name in called_tools:
+                session.call_tool.assert_awaited_once_with("cleanup_session", {})
+                assert events.index(("cleanup", name)) < events.index(("close", name))
+            else:
+                session.call_tool.assert_not_awaited()
+        assert [name for event, name in events if event == "close"] == list(
+            reversed(available_tools)
+        )
+
+    @pytest.mark.asyncio
+    async def test_mcp_partial_initialization_closes_opened_session(self, tool_context):
+        closed = []
+        session = MagicMock(call_tool=AsyncMock())
+
+        @asynccontextmanager
+        async def new_session(name, *_args):
+            if name == "python":
+                raise RuntimeError("session initialization failed")
+            try:
+                yield session
+            finally:
+                closed.append(name)
+
+        tool_server = MagicMock(spec=ToolServer)
+        tool_server.new_session.side_effect = new_session
+        with pytest.raises(RuntimeError, match="session initialization failed"):
+            async with AsyncExitStack() as stack:
+                await tool_context.init_tool_sessions(tool_server, stack, "req", {})
+
+        assert closed == ["browser"]
+        session.call_tool.assert_not_awaited()
 
     @pytest_asyncio.fixture
     async def serving_responses_instance(self):
-        """Create a real OpenAIServingResponses instance for testing"""
+        """Create a real OpenAIServingResponses instance for testing."""
         # Create minimal mocks for required dependencies
         engine_client = MagicMock()
 
@@ -710,6 +797,7 @@ async def generate_output():
             pass
 
         assert serving_responses_instance.engine_client.generate.call_count == 2
+        assert context.request_metrics_cover_all_generation_turns is False
         followup_engine_input = (
             serving_responses_instance.engine_client.generate.call_args_list[1].args[0]
         )
@@ -720,8 +808,7 @@ async def generate_output():
     async def test_initialize_tool_sessions(
         self, serving_responses_instance, mock_context, mock_exit_stack
     ):
-        """Test that method works correctly with only MCP tools"""
-
+        """Test that method works correctly with only MCP tools."""
         request = ResponsesRequest(input="test input", tools=[])
 
         # Call the method
@@ -770,11 +857,11 @@ def test_validate_create_responses_input(
 
 
 class TestValidateGeneratorInput:
-    """Test class for _validate_generator_input method"""
+    """Test class for _validate_generator_input method."""
 
     @pytest_asyncio.fixture
     async def serving_responses_instance(self):
-        """Create a real OpenAIServingResponses instance for testing"""
+        """Create a real OpenAIServingResponses instance for testing."""
         # Create minimal mocks for required dependencies
         engine_client = MagicMock()
 
@@ -802,7 +889,7 @@ async def serving_responses_instance(self):
         return instance
 
     def test_validate_generator_input(self, serving_responses_instance):
-        """Test _validate_generator_input with valid prompt length"""
+        """Test _validate_generator_input with valid prompt length."""
         # Create an engine prompt with valid length (less than max_model_len)
         valid_prompt_token_ids = list(range(5))  # 5 tokens < 100 max_model_len
         engine_input = tokens_input(valid_prompt_token_ids)
@@ -894,6 +981,7 @@ def decode(self, token_ids):
         outputs=[completion],
         finished=True,
         num_cached_tokens=0,
+        num_cache_creation_tokens=3,
     )
     context.append_output(req_output)
 
@@ -914,10 +1002,13 @@ async def dummy_result_generator():
     )
 
     assert response.usage.output_tokens_details.reasoning_tokens == 1
+    assert response.usage.input_tokens_details.cache_write_tokens == 3
+
+    OpenAIResponseUsage.model_validate(response.usage.model_dump(mode="json"))
 
 
 class TestExtractAllowedToolsFromMcpRequests:
-    """Test class for _extract_allowed_tools_from_mcp_requests function"""
+    """Test class for _extract_allowed_tools_from_mcp_requests function."""
 
     def test_extract_allowed_tools_basic_formats(self):
         """Test extraction with list format, object format, and None."""
@@ -1031,7 +1122,7 @@ def _make_previous_item(*, channel, recipient, text="preamble text"):
         return item
 
     def test_preamble_delta_emits_text_events(self) -> None:
-        """commentary + recipient=None should emit output_text.delta events."""
+        """Commentary + recipient=None should emit output_text.delta events."""
         from vllm.entrypoints.openai.responses.streaming_events import (
             emit_content_delta_events,
         )
@@ -1064,7 +1155,7 @@ def test_preamble_delta_second_token_no_added(self) -> None:
         assert "response.output_item.added" not in type_names
 
     def test_commentary_with_function_recipient_not_preamble(self) -> None:
-        """commentary + recipient='functions.X' must NOT use preamble path."""
+        """Commentary + recipient='functions.X' must NOT use preamble path."""
         from vllm.entrypoints.openai.responses.streaming_events import (
             emit_content_delta_events,
         )
@@ -1102,7 +1193,7 @@ def test_preamble_done_emits_text_done_events(self) -> None:
         assert "response.output_item.done" in type_names
 
     def test_commentary_with_recipient_no_preamble_done(self) -> None:
-        """commentary + recipient='functions.X' should route to function call
+        """Commentary + recipient='functions.X' should route to function call
         done, not preamble done."""
         from vllm.entrypoints.openai.responses.streaming_events import (
             emit_previous_item_done_events,
@@ -1160,9 +1251,19 @@ def test_zero_delta_items_should_preserve_streaming_lifecycle(
             assert "response.output_item.done" in type_names
 
 
-def _make_simple_context_with_output(text, token_ids, response_parser=None):
-    """Create a SimpleContext with a RequestOutput containing the given text."""
-    ctx = SimpleContext(response_parser=response_parser)
+_PER_REQUEST_STATS = RequestStateStats(
+    queued_ts=1.0,
+    scheduled_ts=1.5,
+    first_token_ts=2.0,
+    last_token_ts=3.0,
+)
+
+
+def _make_request_output(
+    text,
+    token_ids,
+    metrics: RequestStateStats | None = None,
+):
     completion = CompletionOutput(
         index=0,
         text=text,
@@ -1172,7 +1273,7 @@ def _make_simple_context_with_output(text, token_ids, response_parser=None):
         finish_reason=None,
         stop_reason=None,
     )
-    req_output = RequestOutput(
+    return RequestOutput(
         request_id="req",
         prompt="hi",
         prompt_token_ids=[7, 8],
@@ -1180,13 +1281,29 @@ def _make_simple_context_with_output(text, token_ids, response_parser=None):
         outputs=[completion],
         finished=False,
         num_cached_tokens=0,
+        metrics=metrics,
     )
+
+
+def _make_simple_context_with_output(
+    text,
+    token_ids,
+    response_parser=None,
+    metrics: RequestStateStats | None = None,
+):
+    """Create a SimpleContext with a RequestOutput containing the given text."""
+    ctx = SimpleContext(response_parser=response_parser)
+    req_output = _make_request_output(text, token_ids, metrics)
     ctx.append_output(req_output)
+    ctx.request_metrics = metrics
     return ctx
 
 
-def _make_serving_instance_with_reasoning():
-    """Create an OpenAIServingResponses with a mocked reasoning parser."""
+def _make_serving_instance(
+    *,
+    reasoning_parser: str = "",
+    enable_per_request_metrics: bool = False,
+) -> OpenAIServingResponses:
     engine_client = MagicMock()
     model_config = MagicMock()
     model_config.max_model_len = 100
@@ -1197,18 +1314,112 @@ def _make_serving_instance_with_reasoning():
     engine_client.input_processor = MagicMock()
     engine_client.renderer = MagicMock()
 
-    models = MagicMock()
-
-    serving = OpenAIServingResponses(
+    return OpenAIServingResponses(
         engine_client=engine_client,
-        models=models,
+        models=MagicMock(),
         online_renderer=MagicMock(),
         request_logger=None,
         chat_template=None,
         chat_template_content_format="auto",
-        reasoning_parser="qwen3",
+        reasoning_parser=reasoning_parser,
+        enable_per_request_metrics=enable_per_request_metrics,
+    )
+
+
+async def _empty_context_generator():
+    if False:
+        yield
+
+
+async def _make_full_metrics_response(
+    enable_per_request_metrics: bool,
+    request_metrics_cover_all_generation_turns: bool = True,
+):
+    serving = _make_serving_instance(
+        enable_per_request_metrics=enable_per_request_metrics
+    )
+    request = ResponsesRequest(input="hi", tools=[], stream=False, store=False)
+    sampling_params = SamplingParams(max_tokens=16)
+    context = SimpleContext()
+    context.request_metrics_cover_all_generation_turns = (
+        request_metrics_cover_all_generation_turns
+    )
+
+    async def generate(*args, **kwargs):
+        yield _make_request_output("hello", [10, 20], _PER_REQUEST_STATS)
+
+    serving.engine_client.generate.side_effect = generate
+    result_generator = serving._generate_with_builtin_tools(
+        request_id=request.request_id,
+        engine_input=tokens_input([7, 8]),
+        sampling_params=sampling_params,
+        context=context,
+    )
+    response = await serving.responses_full_generator(
+        request=request,
+        sampling_params=sampling_params,
+        result_generator=result_generator,
+        context=context,
+        model_name="test-model",
+        tokenizer=MagicMock(),
+        request_metadata=RequestResponseMetadata(request_id="req"),
+    )
+    assert isinstance(response, ResponsesResponse)
+    return response
+
+
+@pytest.mark.asyncio
+async def test_responses_per_request_metrics_follow_server_flag():
+    disabled_response = await _make_full_metrics_response(False)
+    assert disabled_response.metrics is None
+    assert "metrics" not in disabled_response.model_dump(mode="json")
+
+    enabled_response = await _make_full_metrics_response(True)
+    assert enabled_response.metrics is not None
+    assert enabled_response.metrics.time_to_first_token_ms == pytest.approx(500.0)
+
+
+@pytest.mark.asyncio
+async def test_responses_streaming_metrics_only_on_completed_event():
+    serving = _make_serving_instance(enable_per_request_metrics=True)
+    request = ResponsesRequest(input="hi", tools=[], stream=True, store=False)
+    context = _make_simple_context_with_output(
+        "hello", [10, 20], metrics=_PER_REQUEST_STATS
+    )
+
+    events = [
+        event
+        async for event in serving.responses_stream_generator(
+            request=request,
+            sampling_params=SamplingParams(max_tokens=16),
+            result_generator=_empty_context_generator(),
+            context=context,
+            model_name="test-model",
+            tokenizer=MagicMock(),
+            request_metadata=RequestResponseMetadata(request_id="req"),
+        )
+    ]
+
+    for event in events[:-1]:
+        assert "metrics" not in event.response.model_dump(mode="json")
+    assert isinstance(events[-1], ResponseCompletedEvent)
+    assert events[-1].response.metrics is not None
+    assert events[-1].response.metrics.time_to_first_token_ms == pytest.approx(500.0)
+
+
+@pytest.mark.asyncio
+async def test_responses_metrics_suppressed_for_multiple_generation_turns():
+    response = await _make_full_metrics_response(
+        True, request_metrics_cover_all_generation_turns=False
     )
-    return serving
+
+    assert response.metrics is None
+    assert "metrics" not in response.model_dump(mode="json")
+
+
+def _make_serving_instance_with_reasoning():
+    """Create an OpenAIServingResponses with a mocked reasoning parser."""
+    return _make_serving_instance(reasoning_parser="qwen3")
 
 
 def _identity_increment(event):
@@ -1260,7 +1471,6 @@ async def test_mixed_delta_reasoning_and_content_emits_reasoning_delta(
         chunk), the trailing reasoning text must be emitted as a
         ResponseReasoningTextDeltaEvent and included in the
         ResponseReasoningTextDoneEvent text."""
-
         monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
         serving = _make_serving_instance_with_reasoning()
 
@@ -1330,7 +1540,6 @@ async def test_transition_without_mixed_delta_no_extra_reasoning_event(
     ):
         """When the transition from reasoning to content is clean (no mixed
         delta), no extra reasoning delta event should be emitted."""
-
         monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
         serving = _make_serving_instance_with_reasoning()
 
@@ -1392,7 +1601,6 @@ async def test_reasoning_only_stream_no_content(self, monkeypatch):
         """When the stream has only reasoning deltas and no content, the
         reasoning done event should be emitted at finalization with the
         full accumulated text, and no text delta events should appear."""
-
         monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
         serving = _make_serving_instance_with_reasoning()
 
diff --git a/tests/entrypoints/openai/responses/test_streaming_events.py b/tests/entrypoints/openai/responses/test_streaming_events.py
index 6b649aadad8b..e1a8bf46f640 100644
--- a/tests/entrypoints/openai/responses/test_streaming_events.py
+++ b/tests/entrypoints/openai/responses/test_streaming_events.py
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from openai.types.responses import ResponseFunctionWebSearch
+from openai_harmony import Message, Role
+
 from vllm.entrypoints.generate.base.protocol import (
     DeltaFunctionCall,
     DeltaMessage,
@@ -8,11 +11,35 @@
 )
 from vllm.entrypoints.openai.responses.streaming_events import (
     SimpleStreamingEventProcessor,
+    StreamingState,
     _StateType,
+    emit_browser_tool_events,
     split_delta,
 )
 
 
+def test_browser_find_uses_responses_action_type():
+    """Both streaming output items must use the Responses find action type."""
+    message = (
+        Message.from_role_and_content(
+            Role.ASSISTANT, '{"pattern": "vLLM", "cursor": 42}'
+        )
+        .with_channel("analysis")
+        .with_recipient("browser.find")
+    )
+
+    events = emit_browser_tool_events(message, StreamingState())
+
+    added, done = events[0], events[-1]
+    assert added.type == "response.output_item.added"
+    assert done.type == "response.output_item.done"
+    for event in (added, done):
+        assert isinstance(event.item, ResponseFunctionWebSearch)
+        assert event.item.action.type == "find_in_page"
+        assert event.item.action.pattern == "vLLM"
+        assert event.item.action.url == "cursor:42"
+
+
 def _make_tool_call(
     index: int, name: str | None = None, arguments: str | None = None
 ) -> DeltaToolCall:
diff --git a/tests/entrypoints/openai/test_async_tokenization.py b/tests/entrypoints/openai/test_async_tokenization.py
deleted file mode 100644
index 1d3d110d3027..000000000000
--- a/tests/entrypoints/openai/test_async_tokenization.py
+++ /dev/null
@@ -1,82 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import asyncio
-import random
-from collections.abc import Callable
-
-import openai
-import pytest
-import pytest_asyncio
-
-from tests.utils import RemoteOpenAIServer
-
-MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
-
-
-@pytest.fixture(scope="module")
-def server():
-    args = [
-        # use half precision for speed and memory savings in CI environment
-        "--dtype",
-        "bfloat16",
-        "--max-model-len",
-        "8192",
-        "--enforce-eager",
-        "--max-num-seqs",
-        "128",
-        "--load-format",
-        "dummy",
-    ]
-
-    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
-        yield remote_server
-
-
-@pytest_asyncio.fixture
-async def client(server):
-    async with server.get_async_client() as async_client:
-        yield async_client
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
-    ids=["completion", "chat"],
-    argnames=["create_func_gen", "content_body"],
-    argvalues=[
-        (lambda x: x.completions.create, {"prompt": " ".join(["A"] * 10_000)}),
-        (
-            lambda x: x.chat.completions.create,
-            {"messages": [{"role": "user", "content": " ".join(["A"] * 10_000)}]},
-        ),
-    ],
-)
-async def test_with_and_without_truncate(
-    server: RemoteOpenAIServer,
-    client: openai.AsyncOpenAI,
-    create_func_gen: Callable,
-    content_body: dict,
-):
-    create_func = create_func_gen(client)
-    body = {"model": MODEL_NAME, **content_body, "max_tokens": 10}
-
-    num_requests = 10
-    truncate_prompt_tokens = [1000] * (num_requests // 2) + [None] * (
-        num_requests - num_requests // 2
-    )
-    random.shuffle(truncate_prompt_tokens)
-
-    bodies = [
-        {**body, "extra_body": {"truncate_prompt_tokens": t}}
-        for t in truncate_prompt_tokens
-    ]
-
-    async def get_status_code(**kwargs):
-        try:
-            await create_func(**kwargs)
-            return 200
-        except openai.APIStatusError as e:
-            return e.status_code
-
-    responses = await asyncio.gather(*[get_status_code(**b) for b in bodies])
-    assert 500 not in responses
diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py
index 6d3fc2f44746..18b01b85a8a3 100644
--- a/tests/entrypoints/openai/test_openai_schema.py
+++ b/tests/entrypoints/openai/test_openai_schema.py
@@ -83,8 +83,7 @@ def before_generate_case(context: schemathesis.HookContext, strategy):
     assert op is not None
 
     def no_invalid_types(case: schemathesis.Case):
-        """
-        Skips tool_calls with `"type": "custom"` which schemathesis incorrectly
+        """Skips tool_calls with `"type": "custom"` which schemathesis incorrectly
         generates instead of the valid `"type": "function"`.
 
         Example test case that is skipped:
diff --git a/tests/entrypoints/openai/test_penalty_default_resolution.py b/tests/entrypoints/openai/test_penalty_default_resolution.py
new file mode 100644
index 000000000000..2c39f5bef68d
--- /dev/null
+++ b/tests/entrypoints/openai/test_penalty_default_resolution.py
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for presence_penalty / frequency_penalty resolution from
+default_sampling_params in ChatCompletionRequest and CompletionRequest.
+
+Regression test for https://github.com/vllm-project/vllm/issues/50767:
+these two penalties defaulted to 0.0 (not None) and to_sampling_params()
+forwarded them straight through, so server-side defaults coming from
+--override-generation-config / generation_config.json were silently
+discarded on the /v1/chat/completions and /v1/completions endpoints.
+"""
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.completion.protocol import (
+    CompletionRequest,
+)
+
+_DEFAULTS = {"presence_penalty": 1.5, "frequency_penalty": 0.5}
+
+
+def _chat(**kwargs):
+    return ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "hi"}],
+        **kwargs,
+    )
+
+
+class TestChatCompletionPenaltyDefaults:
+    def test_defaults_applied_when_client_omits(self):
+        """Server-default penalties are applied when the client sends none."""
+        sp = _chat().to_sampling_params(100, _DEFAULTS)
+        assert sp.presence_penalty == 1.5
+        assert sp.frequency_penalty == 0.5
+
+    def test_client_value_overrides_default(self):
+        """An explicit client penalty wins over the server default."""
+        sp = _chat(presence_penalty=0.2).to_sampling_params(100, _DEFAULTS)
+        assert sp.presence_penalty == 0.2
+        assert sp.frequency_penalty == 0.5
+
+    def test_falls_back_to_zero_without_default(self):
+        """Without a server default the neutral 0.0 is used (no regression)."""
+        sp = _chat().to_sampling_params(100, {})
+        assert sp.presence_penalty == 0.0
+        assert sp.frequency_penalty == 0.0
+
+
+class TestCompletionPenaltyDefaults:
+    def test_defaults_applied_when_client_omits(self):
+        sp = CompletionRequest(model="test-model", prompt="hi").to_sampling_params(
+            16, _DEFAULTS
+        )
+        assert sp.presence_penalty == 1.5
+        assert sp.frequency_penalty == 0.5
+
+    def test_client_value_overrides_default(self):
+        sp = CompletionRequest(
+            model="test-model", prompt="hi", frequency_penalty=0.9
+        ).to_sampling_params(16, _DEFAULTS)
+        assert sp.presence_penalty == 1.5
+        assert sp.frequency_penalty == 0.9
diff --git a/tests/entrypoints/openai/test_render_parity.py b/tests/entrypoints/openai/test_render_parity.py
index 92b7fe71bd80..4fc44c5a7172 100644
--- a/tests/entrypoints/openai/test_render_parity.py
+++ b/tests/entrypoints/openai/test_render_parity.py
@@ -86,9 +86,10 @@ class MockModelConfig:
     generation_config: str = "auto"
     override_generation_config: dict[str, Any] = field(default_factory=dict)
     media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
+    skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
     enable_prompt_embeds: bool = False
 
diff --git a/tests/entrypoints/openai/test_return_routed_experts.py b/tests/entrypoints/openai/test_return_routed_experts.py
index a64e80c39868..46eeed179531 100644
--- a/tests/entrypoints/openai/test_return_routed_experts.py
+++ b/tests/entrypoints/openai/test_return_routed_experts.py
@@ -12,15 +12,13 @@
 MODEL_NAME = "TitanML/tiny-mixtral"
 
 # tiny-mixtral config: 8 local experts, top-2 routing, 2 hidden layers.
-# The published config has sliding_window=4096, which produces
-# SlidingWindowSpec kv-cache groups; RoutedExpertsManager requires a
-# FullAttentionSpec group, so we override sliding_window=null below.
+# Its published sliding window is incompatible with AuxOutput retention.
 NUM_LOCAL_EXPERTS = 8
 NUM_EXPERTS_PER_TOK = 2
 NUM_HIDDEN_LAYERS = 2
 
 
-def assert_valid_routed_experts(encoded: str | None) -> None:
+def decode_routed_experts(encoded: str | None) -> np.ndarray:
     assert encoded is not None
     routed_experts = np.load(io.BytesIO(base64.b64decode(encoded)))
     assert routed_experts.ndim == 3
@@ -30,6 +28,7 @@ def assert_valid_routed_experts(encoded: str | None) -> None:
     assert topk == NUM_EXPERTS_PER_TOK
     assert (routed_experts >= 0).all()
     assert (routed_experts < NUM_LOCAL_EXPERTS).all()
+    return routed_experts
 
 
 @pytest.fixture(scope="module")
@@ -44,7 +43,11 @@ def server():
         "--hf-overrides",
         '{"sliding_window": null}',
     ]
-    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
+    with RemoteOpenAIServer(
+        MODEL_NAME,
+        args,
+        env_dict={"VLLM_USE_V2_MODEL_RUNNER": "1"},
+    ) as remote_server:
         yield remote_server
 
 
@@ -63,4 +66,8 @@ async def test_routed_experts(server):
         choice = result.model_dump()["choices"][0]
 
         assert choice["token_ids"] is not None
-        assert_valid_routed_experts(choice["routed_experts"])
+        assert result.usage is not None
+        routed_experts = decode_routed_experts(choice["routed_experts"])
+        assert len(routed_experts) == (
+            result.usage.prompt_tokens + len(choice["token_ids"]) - 1
+        )
diff --git a/tests/entrypoints/openai/test_return_token_ids.py b/tests/entrypoints/openai/test_return_token_ids.py
index 8537082e3f8d..9fc8c102d4e4 100644
--- a/tests/entrypoints/openai/test_return_token_ids.py
+++ b/tests/entrypoints/openai/test_return_token_ids.py
@@ -189,8 +189,7 @@ async def test_chat_completion_with_tool_use(server):
 
 @pytest.mark.asyncio
 async def test_comparison_with_prompt_logprobs_and_logprobs(server):
-    """
-    Test that token_ids align with prompt_logprobs and
+    """Test that token_ids align with prompt_logprobs and
     logprobs when return_tokens_as_token_ids is enabled.
     """
     async with server.get_async_client() as client:
diff --git a/tests/entrypoints/openai/test_stop_token_ids.py b/tests/entrypoints/openai/test_stop_token_ids.py
index 5e9d79f68649..d6ac2cf1906f 100644
--- a/tests/entrypoints/openai/test_stop_token_ids.py
+++ b/tests/entrypoints/openai/test_stop_token_ids.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Unit tests for stop_token_ids propagation from default_sampling_params
+"""Unit tests for stop_token_ids propagation from default_sampling_params
 to SamplingParams in ChatCompletionRequest and CompletionRequest.
 
 Regression test for https://github.com/vllm-project/vllm/issues/22519
diff --git a/tests/entrypoints/openai/test_tool_choice_content_none.py b/tests/entrypoints/openai/test_tool_choice_content_none.py
index 7db815c45f38..a0a8f38e9f7a 100644
--- a/tests/entrypoints/openai/test_tool_choice_content_none.py
+++ b/tests/entrypoints/openai/test_tool_choice_content_none.py
@@ -31,9 +31,6 @@ class _DummyDelegatingParser(DelegatingParser):
     def is_reasoning_end(self, input_ids: list[int]) -> bool:
         return False
 
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        return input_ids
-
     def extract_reasoning(self, model_output: str, request):
         return None, model_output
 
diff --git a/tests/entrypoints/openai/utils.py b/tests/entrypoints/openai/utils.py
index 4cc4800bf8dd..02a4c51f94e9 100644
--- a/tests/entrypoints/openai/utils.py
+++ b/tests/entrypoints/openai/utils.py
@@ -16,8 +16,7 @@
 async def accumulate_streaming_response(
     stream_generator: AsyncGenerator[str, None],
 ) -> ChatCompletionResponse:
-    """
-    Accumulate streaming SSE chunks into a complete ChatCompletionResponse.
+    """Accumulate streaming SSE chunks into a complete ChatCompletionResponse.
 
     This helper parses the SSE format and builds up the complete response
     by combining all the delta chunks.
diff --git a/tests/entrypoints/pooling/basic/test_tiling_engine.py b/tests/entrypoints/pooling/basic/test_tiling_engine.py
index 88c9b434b38f..c620816dc56a 100644
--- a/tests/entrypoints/pooling/basic/test_tiling_engine.py
+++ b/tests/entrypoints/pooling/basic/test_tiling_engine.py
@@ -30,8 +30,7 @@ def llm(vllm_runner):
 
 @pytest.mark.skip_global_cleanup
 def test_tiling_engine_basic(llm):
-    """
-    Basic test with a small number of prompts (less than max_num_seqs).
+    """Basic test with a small number of prompts (less than max_num_seqs).
     No tiling should be triggered, but the engine still processes correctly.
     """
     prompts = ["Hello", "World"]
@@ -41,8 +40,7 @@ def test_tiling_engine_basic(llm):
 
 @pytest.mark.skip_global_cleanup
 def test_tiling_engine_many_requests(llm):
-    """
-    Test with a large number of prompts that exceeds max_num_seqs.
+    """Test with a large number of prompts that exceeds max_num_seqs.
     This verifies that _run_tiling_engine correctly chunks requests,
     processes all of them, and returns outputs in the correct order.
     """
@@ -54,8 +52,7 @@ def test_tiling_engine_many_requests(llm):
 
 @pytest.mark.skip_global_cleanup
 def test_tiling_engine_with_pooling_params(llm):
-    """
-    Test the tiling engine when different PoolingParams are provided.
+    """Test the tiling engine when different PoolingParams are provided.
     The engine must handle a list of params that matches the number of prompts.
     """
     num_prompts = 10
@@ -77,8 +74,7 @@ def test_tiling_engine_with_pooling_params(llm):
 
 @pytest.mark.skip_global_cleanup
 def test_tiling_engine_abort_on_exception(llm):
-    """
-    Test that abort_request IS called with the correct arguments when an
+    """Test that abort_request IS called with the correct arguments when an
     exception occurs inside the engine's step() loop.
     """
     prompts = ["Prompt 0", "Prompt 1", "Prompt 2"]
diff --git a/tests/entrypoints/pooling/embed/test_online_dimensions.py b/tests/entrypoints/pooling/embed/test_online_dimensions.py
index 7abec99ec355..67c6d7b1345c 100644
--- a/tests/entrypoints/pooling/embed/test_online_dimensions.py
+++ b/tests/entrypoints/pooling/embed/test_online_dimensions.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`.
-"""
+"""Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`."""
 
 import openai
 import pytest
diff --git a/tests/entrypoints/pooling/embed/test_online_long_text.py b/tests/entrypoints/pooling/embed/test_online_long_text.py
index 60cacb17f8fe..ab3633cd3d0f 100644
--- a/tests/entrypoints/pooling/embed/test_online_long_text.py
+++ b/tests/entrypoints/pooling/embed/test_online_long_text.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test cases for long text embedding with automatic chunking mechanism.
+"""Test cases for long text embedding with automatic chunking mechanism.
 
 This test suite validates vLLM's automatic chunking functionality for handling
 text inputs that exceed the model's maximum token length, specifically targeting
@@ -242,7 +241,6 @@ async def test_long_text_embedding_1500_chars(
 ):
     """Test embedding processing for ~1500 character long text
     (~1028 tokens, exceeding 512 token limit)."""
-
     # Verify text length
     # Verify text has sufficient word count (approximately 1500 words)
     word_count = len(LONG_TEXT_1500_WORDS.split())
@@ -291,7 +289,6 @@ async def test_long_text_embedding_2500_chars(
 ):
     """Test embedding processing for ~2500 character long text
     (~2048 tokens, requiring multiple chunks)."""
-
     # Verify text length
     # Verify text has sufficient word count (approximately 2500 words)
     word_count = len(LONG_TEXT_2500_WORDS.split())
@@ -339,7 +336,6 @@ async def test_batch_long_text_embedding(
     client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
 ):
     """Test batch long text embedding processing."""
-
     input_texts = [
         LONG_TEXT_1500_WORDS,
         LONG_TEXT_2500_WORDS,
@@ -385,7 +381,6 @@ async def test_chunked_vs_normal_consistency(
 ):
     """Test consistency between chunked and
     normal processing (using short text)."""
-
     # Use a short text within the 512 token limit
     short_text = (
         "Artificial intelligence technology is changing our world, "
@@ -424,7 +419,6 @@ async def test_chunked_processing_response_format(
     client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
 ):
     """Test response format and structure during chunked processing."""
-
     # Test with long text to trigger chunking
     embedding_response = await client_with_chunked_processing.embeddings.create(
         model=model_name,
diff --git a/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py b/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py
index dac7b6db6cec..16392fd33837 100644
--- a/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py
+++ b/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py
@@ -126,35 +126,6 @@ async def test_score_api_queries_str_n_documents_str_n(
         assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
 
 
-@pytest.mark.asyncio
-async def test_score_api_queries_vs_documents(hf_model, server: RemoteOpenAIServer):
-    text_pairs = [
-        [TEXTS_1[0], TEXTS_2[0]],
-        [TEXTS_1[1], TEXTS_2[1]],
-    ]
-
-    score_response = requests.post(
-        server.url_for("score"),
-        json={
-            "model": MODEL_NAME,
-            "queries": TEXTS_1,
-            "documents": TEXTS_2,
-        },
-    )
-    score_response.raise_for_status()
-    score = ScoreResponse.model_validate(score_response.json())
-
-    assert score.id is not None
-    assert score.data is not None
-    assert len(score.data) == 2
-
-    vllm_outputs = [d.score for d in score.data]
-    hf_outputs = hf_model.predict(text_pairs).tolist()
-
-    for i in range(len(vllm_outputs)):
-        assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
-
-
 @pytest.mark.asyncio
 async def test_score_api_queries_vs_items(hf_model, server: RemoteOpenAIServer):
     text_pairs = [
diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py
index 401bf4ff2268..9228f96fe746 100644
--- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py
+++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py
@@ -143,35 +143,6 @@ async def test_score_api_queries_str_n_documents_str_n(
         assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
 
 
-@pytest.mark.asyncio
-async def test_score_api_queries_vs_documents(hf_model, server: RemoteOpenAIServer):
-    text_pairs = [
-        [TEXTS_1[0], TEXTS_2[0]],
-        [TEXTS_1[1], TEXTS_2[1]],
-    ]
-
-    score_response = requests.post(
-        server.url_for("score"),
-        json={
-            "model": MODEL_NAME,
-            "queries": TEXTS_1,
-            "documents": TEXTS_2,
-        },
-    )
-    score_response.raise_for_status()
-    score = ScoreResponse.model_validate(score_response.json())
-
-    assert score.id is not None
-    assert score.data is not None
-    assert len(score.data) == 2
-
-    vllm_outputs = [d.score for d in score.data]
-    hf_outputs = hf_model.predict(text_pairs).tolist()
-
-    for i in range(len(vllm_outputs)):
-        assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
-
-
 @pytest.mark.asyncio
 async def test_score_api_queries_vs_items(hf_model, server: RemoteOpenAIServer):
     text_pairs = [
diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py
index 9f2d7da80116..bf3c771fdd1b 100644
--- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py
+++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py
@@ -27,31 +27,10 @@
 
 ATTN_BACKENDS = ROCM_ATTN_BACKENDS if current_platform.is_rocm() else ["auto"]
 
-# Per-backend tolerance with explicit entries; "default" is the fallback
-BACKEND_TOL: dict[str, float] = {
-    "default": 0.05,  # 5% tolerance for other backends (e.g. FLASH_ATTN)
-    # Relaxed tolerances for ROCm attn
-    # See: https://github.com/vllm-project/vllm/issues/35569
-    "ROCM_ATTN": 0.09,  # gfx950:~8.45%, gfx942:~3.70%
-    "ROCM_AITER_FA": 0.045,  # gfx950:~2.00%, gfx942:~0.80%
-    "TRITON_ATTN": 0.045,  # gfx950:~3.00%, gfx942:~2.20%
-    "FLEX_ATTENTION": 0.045,  # gfx950:~3.25%, gfx942:~1.10%
-}
-
-# Some ROCm attention backends show small absolute drift on the low
-# text-vs-text probability even though larger scores remain well inside the
-# relative tolerance. The absolute drift is uniform across score magnitudes
-# (~0.005-0.010), so it only exceeds the relative tolerance for the small
-# ~0.10 text-vs-text value. Keep the relative tolerances tight and add only a
-# small absolute floor for the affected backends.
-# TRITON_ATTN: gfx942/ROCm 7.2 drifts ~0.008 abs on text-vs-text (~7.9% rel).
-BACKEND_ABS_TOL: dict[str, float] = {
-    "default": 0.0,
-    "auto": 0.007,
-    "ROCM_AITER_FA": 0.005,
-    "TRITON_ATTN": 0.009,
-    "FLEX_ATTENTION": 0.00875,
-}
+# The absolute floor matters only for the low ~0.10 text-vs-text score, where a
+# few thousandths of drift would otherwise exceed the relative bound.
+REL_TOL = 0.05
+ABS_TOL = 0.007
 
 # ROCm: disable prefix caching and eliminate batch variance to reduce
 # test flakiness.
@@ -62,17 +41,9 @@
 )
 
 
-def get_tol(backend: str) -> float:
-    return BACKEND_TOL.get(backend, BACKEND_TOL["default"])
-
-
-def get_abs_tol(backend: str) -> float:
-    return BACKEND_ABS_TOL.get(backend, BACKEND_ABS_TOL["default"])
-
-
 def assert_score(actual: float, expected: float, backend: str, label: str):
-    tol = get_tol(backend)
-    abs_tol = get_abs_tol(backend)
+    tol = REL_TOL
+    abs_tol = ABS_TOL
     diff = abs(actual - expected)
     rel_diff = diff / abs(expected) if expected != 0 else diff
     print(
@@ -104,9 +75,16 @@ def assert_score(actual: float, expected: float, backend: str, label: str):
     },
 ]
 
-TEXT_VS_TEXT = 0.10040374100208282
-TEXT_VS_IMAGE = 0.7423753142356873
-TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838
+# One bf16 ULP in attention compounds to ~0.1% here, so ROCm backends each
+# settle differently; these are their midpoint across backends and archs.
+if current_platform.is_rocm():
+    TEXT_VS_TEXT = 0.105243
+    TEXT_VS_IMAGE = 0.742472
+    TEXT_VS_TEXT_PLUS_IMAGE = 0.544568
+else:
+    TEXT_VS_TEXT = 0.10040374100208282
+    TEXT_VS_IMAGE = 0.7423753142356873
+    TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838
 
 
 @pytest.fixture(scope="module", params=ATTN_BACKENDS)
@@ -507,6 +485,4 @@ async def test_rerank_api_instruction_field_matches_chat_template_kwargs(
     kwargs_scores = [
         r.relevance_score for r in sorted(kwargs_rerank.results, key=lambda x: x.index)
     ]
-    assert field_scores == pytest.approx(
-        kwargs_scores, rel=get_tol(backend), abs=get_abs_tol(backend)
-    )
+    assert field_scores == pytest.approx(kwargs_scores, rel=REL_TOL, abs=ABS_TOL)
diff --git a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py
index c2fe62f609df..ab5a21c0859b 100644
--- a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py
+++ b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py
@@ -1,15 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for JinaRankingIOProcessor online request building."""
+"""Unit tests for JinaRankingIOProcessor request building."""
 
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
+from tokenizers import Tokenizer, models, pre_tokenizers
+from transformers import TokenizersBackend
 
+from vllm import PoolingParams
 from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
 from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessor
 from vllm.entrypoints.pooling.scoring.protocol import RerankRequest
 from vllm.entrypoints.pooling.scoring.typing import ScoringData
+from vllm.entrypoints.pooling.typing import OfflineScoringInputsContext
+from vllm.renderers import TokenizeParams
 
 pytestmark = pytest.mark.skip_global_cleanup
 
@@ -57,3 +63,82 @@ def _spy_base(self, ctx):
     assert captured["truncation_side"] == "left"
     # The real request is restored after delegating.
     assert ctx.request is request
+
+
+@pytest.fixture
+def offline_processor_and_context():
+    backend = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]"))
+    backend.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
+    proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor)
+    proc.tokenizer = TokenizersBackend(tokenizer_object=backend)
+    proc.model_config = SimpleNamespace(max_model_len=1024, is_encoder_decoder=False)
+    proc.renderer = SimpleNamespace(
+        default_cmpl_tok_params=TokenizeParams(max_total_tokens=1024)
+    )
+    ctx = OfflineScoringInputsContext(
+        pooling_task="token_embed",
+        scoring_data=ScoringData(
+            data_1=["query alpha beta", "query gamma delta"],
+            data_2=["document one two", "document three four"],
+        ),
+        pooling_params=PoolingParams(
+            extra_kwargs={
+                "chat_template_kwargs": {"instruction": "Keep this instruction"}
+            }
+        ),
+        tokenization_kwargs=None,
+        chat_template=None,
+        lora_request=None,
+        priorities=None,
+    )
+    return proc, ctx
+
+
+@pytest.mark.parametrize(
+    ("query_limit", "doc_limit", "n_queries"),
+    [(None, None, 1), (0, 0, 2), (1, None, 1), (None, 2, 1), (1, 2, 1), (1, 2, 2)],
+    ids=["unset", "zero", "query", "doc", "both-1-to-n", "both-n-to-n"],
+)
+def test_offline_truncates_text_before_formatting(
+    offline_processor_and_context, query_limit, doc_limit, n_queries
+):
+    proc, ctx = offline_processor_and_context
+    queries = ctx.scoring_data.data_1[:n_queries]
+    docs = ctx.scoring_data.data_2
+    ctx.scoring_data = ScoringData(data_1=queries, data_2=docs)
+    ctx.pooling_params.extra_kwargs.update(
+        {
+            name: limit
+            for name, limit in (
+                ("max_tokens_per_query", query_limit),
+                ("max_tokens_per_doc", doc_limit),
+            )
+            if limit is not None
+        }
+    )
+
+    factory, num_requests = proc.get_request_factory_offline(ctx)
+    requests = list(factory())
+    assert len(requests) == num_requests == n_queries
+    for i, request in enumerate(requests):
+        prompt = request["prompts"]["prompt"]
+        query = " ".join(queries[i].split()[: query_limit or None])
+        assert f"<query>\n{query}<|rerank_token|>\n</query>" in prompt
+        prompt_docs = docs if n_queries == 1 else [docs[i]]
+        for j, doc in enumerate(prompt_docs):
+            doc = " ".join(doc.split()[: doc_limit or None])
+            assert f'<passage id="{j}">\n{doc}<|embed_token|>\n</passage>' in prompt
+        assert prompt.count("<|embed_token|>") == len(prompt_docs)
+        assert prompt.count("<|rerank_token|>") == 1
+        assert "<instruct>\nKeep this instruction\n</instruct>" in prompt
+
+
+@pytest.mark.parametrize("name", ["max_tokens_per_query", "max_tokens_per_doc"])
+@pytest.mark.parametrize("limit", [-1, 1024])
+def test_offline_rejects_invalid_token_limits(
+    offline_processor_and_context, name, limit
+):
+    proc, ctx = offline_processor_and_context
+    ctx.pooling_params.extra_kwargs[name] = limit
+    with pytest.raises(ValueError, match=name):
+        proc.get_request_factory_offline(ctx)
diff --git a/tests/entrypoints/pooling/test_io_processor.py b/tests/entrypoints/pooling/test_io_processor.py
index 6c64966e588e..9d2207fdf15d 100644
--- a/tests/entrypoints/pooling/test_io_processor.py
+++ b/tests/entrypoints/pooling/test_io_processor.py
@@ -74,3 +74,150 @@ def test_rejects_conflicting_pooling_task(processor: PoolingIOProcessor):
     )
     assert exc_info.value.parameter is None
     assert exc_info.value.value is None
+
+
+def test_qwen3_reranker_warns_without_chat_template(monkeypatch):
+    from unittest.mock import MagicMock
+
+    from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
+
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock()
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.models.interfaces.supports_score_template",
+        lambda *_: False,
+    )
+    monkeypatch.setattr(
+        "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer",
+        lambda *_: False,
+    )
+
+    mock_logger = MagicMock()
+    monkeypatch.setattr(
+        "vllm.entrypoints.pooling.scoring.io_processor.logger",
+        mock_logger,
+    )
+
+    model_config = MagicMock()
+    model_config.architecture = "Qwen3ForSequenceClassification"
+    model_config.model = "Qwen/Qwen3-Reranker-0.6B"
+    model_config.hf_config.is_original_qwen3_reranker = True
+    model_config.use_sep_token = False
+    model_config.is_multimodal_model = False
+    model_config.max_model_len = 8192
+
+    vllm_config = MagicMock(model_config=model_config)
+    tokenizer = MagicMock()
+    tokenizer.pad_token_id = 0
+    renderer = MagicMock()
+    renderer.get_tokenizer.return_value = tokenizer
+    chat_template_config = MagicMock(chat_template=None)
+
+    CrossEncoderIOProcessor(
+        vllm_config=vllm_config,
+        renderer=renderer,
+        chat_template_config=chat_template_config,
+    )
+
+    mock_logger.warning.assert_called_once()
+    warning_call = mock_logger.warning.call_args
+    assert "Qwen/Qwen3-Reranker-0.6B" in warning_call[0][1]
+    assert "qwen3_reranker.jinja" in warning_call[0][2]
+
+
+def test_qwen3_vl_reranker_warns_with_vl_template(monkeypatch):
+    from unittest.mock import MagicMock
+
+    from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
+
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock()
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.models.interfaces.supports_score_template",
+        lambda *_: False,
+    )
+    monkeypatch.setattr(
+        "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer",
+        lambda *_: False,
+    )
+
+    mock_logger = MagicMock()
+    monkeypatch.setattr(
+        "vllm.entrypoints.pooling.scoring.io_processor.logger",
+        mock_logger,
+    )
+
+    model_config = MagicMock()
+    model_config.architecture = "Qwen3VLForSequenceClassification"
+    model_config.model = "Qwen/Qwen3-VL-Reranker-2B"
+    model_config.hf_config.is_original_qwen3_reranker = True
+    model_config.use_sep_token = False
+    model_config.is_multimodal_model = True
+    model_config.max_model_len = 8192
+
+    vllm_config = MagicMock(model_config=model_config)
+    tokenizer = MagicMock()
+    tokenizer.pad_token_id = 0
+    renderer = MagicMock()
+    renderer.get_tokenizer.return_value = tokenizer
+    chat_template_config = MagicMock(chat_template=None)
+
+    CrossEncoderIOProcessor(
+        vllm_config=vllm_config,
+        renderer=renderer,
+        chat_template_config=chat_template_config,
+    )
+
+    mock_logger.warning.assert_called_once()
+    warning_call = mock_logger.warning.call_args
+    assert "Qwen/Qwen3-VL-Reranker-2B" in warning_call[0][1]
+    assert "qwen3_vl_reranker.jinja" in warning_call[0][2]
+
+
+def test_qwen3_reranker_no_warning_when_template_provided(monkeypatch):
+    from unittest.mock import MagicMock
+
+    from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
+
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.get_model_cls", lambda *_: MagicMock()
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.models.interfaces.supports_score_template",
+        lambda *_: False,
+    )
+    monkeypatch.setattr(
+        "vllm.entrypoints.pooling.scoring.io_processor.is_mistral_tokenizer",
+        lambda *_: False,
+    )
+
+    mock_logger = MagicMock()
+    monkeypatch.setattr(
+        "vllm.entrypoints.pooling.scoring.io_processor.logger",
+        mock_logger,
+    )
+
+    model_config = MagicMock()
+    model_config.architecture = "Qwen3ForSequenceClassification"
+    model_config.model = "Qwen/Qwen3-Reranker-0.6B"
+    model_config.hf_config.is_original_qwen3_reranker = True
+    model_config.use_sep_token = False
+    model_config.is_multimodal_model = False
+    model_config.max_model_len = 8192
+
+    vllm_config = MagicMock(model_config=model_config)
+    tokenizer = MagicMock()
+    tokenizer.pad_token_id = 0
+    renderer = MagicMock()
+    renderer.get_tokenizer.return_value = tokenizer
+    chat_template_config = MagicMock(chat_template="custom template")
+
+    CrossEncoderIOProcessor(
+        vllm_config=vllm_config,
+        renderer=renderer,
+        chat_template_config=chat_template_config,
+    )
+
+    mock_logger.warning.assert_not_called()
diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py
index a30e48745018..80ecc0f49cce 100644
--- a/tests/entrypoints/scale_out/derender/test_derender.py
+++ b/tests/entrypoints/scale_out/derender/test_derender.py
@@ -52,6 +52,7 @@ def _make_generate_response(
     logprobs: dict | None = None,
     prompt_logprobs: list | None = None,
     kv_transfer_params: dict | None = None,
+    metrics: dict | None = None,
 ) -> dict:
     choice: dict = {
         "index": 0,
@@ -64,6 +65,7 @@ def _make_generate_response(
         "choices": [choice],
         "prompt_logprobs": prompt_logprobs,
         "kv_transfer_params": kv_transfer_params,
+        "metrics": metrics,
     }
 
 
@@ -266,6 +268,49 @@ async def test_derender_chat_kv_transfer_params_passthrough(client):
     assert response.json()["kv_transfer_params"] == kv
 
 
+@pytest.mark.asyncio
+async def test_derender_chat_metrics_passthrough(client):
+    gen_req = await _render_chat(client)
+    metrics = {
+        "speculative_decoding": {
+            "mean_acceptance_length": 2.0,
+            "draft_acceptance_rate": 0.5,
+            "acceptance_histogram": [0, 1],
+            "num_spec_steps": 1,
+            "num_accepted_draft_tokens": 1,
+            "num_draft_tokens": 2,
+            "num_spec_tokens": 1,
+        }
+    }
+    response = await client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": MODEL_NAME,
+            "generate_response": _make_generate_response(
+                gen_req["token_ids"][:3], metrics=metrics
+            ),
+        },
+    )
+    assert response.status_code == 200
+    actual = response.json()["metrics"]
+    assert set(actual) == {
+        "time_to_first_token_ms",
+        "generation_time_ms",
+        "queue_time_ms",
+        "mean_itl_ms",
+        "tokens_per_second",
+        "speculative_decoding",
+    }
+    assert all(
+        actual[name] is None for name in actual if name != "speculative_decoding"
+    )
+    assert actual["speculative_decoding"] == {
+        **metrics["speculative_decoding"],
+        "per_step_accepted": None,
+        "per_step_drafted": None,
+    }
+
+
 @pytest.mark.asyncio
 async def test_derender_chat_empty_token_ids(client):
     """Empty token_ids list returns 400."""
@@ -555,7 +600,7 @@ async def test_derender_chat_oversized_token_ids_rejected(client):
 
 @pytest.mark.asyncio
 async def test_derender_chat_too_many_choices_rejected(client):
-    """choices count exceeding VLLM_MAX_N_SEQUENCES returns 400."""
+    """Choices count exceeding VLLM_MAX_N_SEQUENCES returns 400."""
     # Default VLLM_MAX_N_SEQUENCES is 16384; use a larger count.
     oversized_choices = [
         {"index": i, "token_ids": [42], "finish_reason": "stop"} for i in range(20_000)
@@ -800,7 +845,12 @@ def _e2e_generate_response(
 
 @pytest.mark.asyncio
 async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer):
-    """Plain text without reasoning markers roundtrips correctly."""
+    """Plain text without reasoning markers roundtrips correctly.
+
+    Markerless output has no ``</think>``, which deepseek_r1 classifies
+    wholly as reasoning, so the text lands there rather than in content.
+    What this pins is detokenization fidelity through the parser path.
+    """
     messages = [{"role": "user", "content": "What is 2+2?"}]
     gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
 
@@ -814,16 +864,22 @@ async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer):
             "model": PARSER_MODEL,
             "generate_response": _e2e_generate_response(output_ids),
             "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {"model": PARSER_MODEL, "messages": messages},
         },
     )
     assert resp.status_code == 200, resp.text
-    content = resp.json()["choices"][0]["message"]["content"]
-    assert content == expected
+    message = resp.json()["choices"][0]["message"]
+    assert message["reasoning"] == expected
+    assert message["content"] is None
 
 
 @pytest.mark.asyncio
 async def test_e2e_token_identity(parser_client, parser_tokenizer):
-    """encode(derender(token_ids)) == token_ids (RL invariant)."""
+    """encode(derender(token_ids)) == token_ids (RL invariant).
+
+    Markerless output comes back as reasoning (see
+    ``test_e2e_plain_roundtrip``), so re-encode that.
+    """
     messages = [{"role": "user", "content": "Hi"}]
     gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
 
@@ -836,17 +892,22 @@ async def test_e2e_token_identity(parser_client, parser_tokenizer):
             "model": PARSER_MODEL,
             "generate_response": _e2e_generate_response(output_ids),
             "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {"model": PARSER_MODEL, "messages": messages},
         },
     )
     assert resp.status_code == 200
-    content = resp.json()["choices"][0]["message"]["content"]
-    re_encoded = _encode(parser_tokenizer, content)
+    reasoning = resp.json()["choices"][0]["message"]["reasoning"]
+    re_encoded = _encode(parser_tokenizer, reasoning)
     assert output_ids == re_encoded
 
 
 @pytest.mark.asyncio
 async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer):
-    """CJK + emoji roundtrip without U+FFFD."""
+    """CJK + emoji roundtrip without U+FFFD.
+
+    Markerless output comes back as reasoning (see
+    ``test_e2e_plain_roundtrip``).
+    """
     messages = [{"role": "user", "content": "Reply in Chinese"}]
     gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
 
@@ -859,11 +920,12 @@ async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer):
             "model": PARSER_MODEL,
             "generate_response": _e2e_generate_response(output_ids),
             "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {"model": PARSER_MODEL, "messages": messages},
         },
     )
     assert resp.status_code == 200
-    content = resp.json()["choices"][0]["message"]["content"]
-    assert "�" not in content
+    reasoning = resp.json()["choices"][0]["message"]["reasoning"]
+    assert "�" not in reasoning
 
 
 @pytest.mark.asyncio
@@ -976,8 +1038,10 @@ async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenize
 
 
 @pytest.mark.asyncio
-async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer):
-    """Without chat_request, derender falls back to plain detokenization."""
+async def test_e2e_no_chat_request_rejected(parser_client, parser_tokenizer):
+    """Without chat_request a parser configured model rejects with 400
+    rather than silently falling back to plain detokenization. This is to
+    prevent the leak of raw reasoning/tool markup into content."""
     messages = [{"role": "user", "content": "Hello"}]
     gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
 
@@ -992,9 +1056,8 @@ async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer):
             "prompt_tokens": len(gen_req["token_ids"]),
         },
     )
-    assert resp.status_code == 200
-    content = resp.json()["choices"][0]["message"]["content"]
-    assert "Hi" in content
+    assert resp.status_code == 400
+    assert "chat_request" in resp.json()["error"]["message"]
 
 
 # ---------------------------------------------------------------------------
diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py
index 5b313464dad1..a89f8d68db53 100644
--- a/tests/entrypoints/scale_out/derender/test_derender_parity.py
+++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py
@@ -3,20 +3,25 @@
 
 """Round trip parity CI: render -> generate -> derender -> render == render.
 
-Pins the coupled parsing path (``/v1/chat/completions`` on a normal GPU
+Pins the coupled parsing path (`/v1/chat/completions` on a normal GPU
 server which parses generated tokens incrementally as they stream out)
-against the disaggregated path (``OnlineDerenderer.derender_chat`` via
-``/v1/chat/completions/derender`` which parses the same tokens all at once,
-out of process). A standard ``vllm serve`` GPU server mounts both, so one
+against the disaggregated path (`OnlineDerenderer.derender_chat` via
+`/v1/chat/completions/derender` which parses the same tokens all at once,
+out of process). A standard `vllm serve` GPU server mounts both, so one
 real generation lets both parsers run on identical input.
 
 Both paths consume the same generated token IDs (extracted from the coupled
-response's ``token_ids`` via ``return_token_ids=True``), so generation
+response's `token_ids` via `return_token_ids=True`), so generation
 nondeterminism is irrelevant. The only variable under test is whether the
 two parsing code paths agree. Parity is asserted unconditionally as only the
 stronger per case assertions (e.g. "a tool call was produced") are gated
 behind the marker actually having been emitted since a 1.5B model is not
-guaranteed to emit ``<think>`` / ``<tool_call>``.
+guaranteed to emit `<think>` / `<tool_call>`.
+
+The streaming cases do the same for `stream=true`. Each coupled chunk's
+`token_ids` becomes one `/inference/v1/generate` stream chunk, fed
+through the chunked derender path with `stream_state` threaded across calls
+the way a client would, so both parsers see the same chunk boundaries.
 """
 
 import json
@@ -24,7 +29,12 @@
 import httpx
 import pytest
 import pytest_asyncio
+import regex as re
 
+from tests.entrypoints.scale_out.derender.utils import (
+    assemble_stream,
+    stream_chat_derender,
+)
 from tests.utils import RemoteOpenAIServer
 
 MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
@@ -34,6 +44,7 @@
     "hermes",
     "--reasoning-parser",
     "deepseek_r1",
+    "--enable-scale-out",
 ]
 
 TOOLS = [
@@ -54,11 +65,7 @@
 
 @pytest.fixture(scope="module")
 def server():
-    with RemoteOpenAIServer(
-        MODEL,
-        ARGS,
-        env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"},
-    ) as remote_server:
+    with RemoteOpenAIServer(MODEL, ARGS) as remote_server:
         yield remote_server
 
 
@@ -231,8 +238,8 @@ async def test_parity_logprobs(client):
     """token_id:N resolution parity vs. the coupled server's real strings.
 
     A real disaggregated worker only has token IDs so it emits logprobs
-    with ``token_id:N`` placeholders (``return_tokens_as_token_ids=True``
-    reproduces that shape here). ``/derender`` must resolve those
+    with `token_id:N` placeholders (`return_tokens_as_token_ids=True`
+    reproduces that shape here). `/derender` must resolve those
     placeholders to the same token strings/bytes the coupled server
     resolves them to directly.
     """
@@ -270,3 +277,166 @@ async def test_parity_logprobs(client):
     for d_entry, r_entry in zip(d_content, r_content):
         assert d_entry["token"] == r_entry["token"]
         assert d_entry["bytes"] == r_entry["bytes"]
+
+
+# ---------------------------------------------------------------------------
+# Streaming helpers
+# ---------------------------------------------------------------------------
+
+
+async def _coupled_stream(
+    client: httpx.AsyncClient, messages: list[dict], **extra
+) -> tuple[list[dict], list[int]]:
+    """Stream `/v1/chat/completions` and return the choices that carry
+    `token_ids` along with the prompt token IDs.
+
+    The leading role only chunk has no `token_ids` and no counterpart on
+    `/inference/v1/generate`, so it is dropped.
+    """
+    choices: list[dict] = []
+    prompt_token_ids: list[int] | None = None
+    async with client.stream(
+        "POST",
+        "/v1/chat/completions",
+        json={
+            "model": MODEL,
+            "messages": messages,
+            "temperature": 0,
+            "max_tokens": 128,
+            "return_token_ids": True,
+            "stream": True,
+            **extra,
+        },
+    ) as resp:
+        assert resp.status_code == 200, await resp.aread()
+        async for line in resp.aiter_lines():
+            if not line.startswith("data: "):
+                continue
+            data = line[len("data: ") :]
+            if data == "[DONE]":
+                break
+            chunk = json.loads(data)
+            assert "error" not in chunk, chunk
+            if chunk.get("prompt_token_ids") is not None:
+                prompt_token_ids = chunk["prompt_token_ids"]
+            choices.extend(
+                ch for ch in chunk["choices"] if ch.get("token_ids") is not None
+            )
+    assert prompt_token_ids is not None
+    return choices, prompt_token_ids
+
+
+def _tool_call_id_shape(tool_call_id: str | None) -> str | None:
+    """Drop the random suffix so IDs compare by format, e.g. `chatcmpl-tool-`.
+    Deterministic IDs such as kimi_k2's `functions.<name>:<idx>` are kept
+    whole."""
+    if tool_call_id is None:
+        return None
+    return re.sub(r"[0-9a-f]{16}$", "", tool_call_id)
+
+
+async def _run_stream_parity_case(
+    client: httpx.AsyncClient,
+    messages: list[dict],
+    finish_only_tail: bool = False,
+    **extra,
+) -> tuple[dict, list[dict]]:
+    """Stream the coupled request, replay its chunks through the streaming
+    derender endpoint and assert the assembled messages agree. Returns the
+    assembled coupled message and the coupled choices.
+
+    Derender gets the engine's finish reason, since `"tool_calls"` is a chat
+    level rewrite of `"stop"` that derender has to reapply itself. With
+    `finish_only_tail` it goes on a trailing chunk with no tokens.
+    """
+    coupled_choices, prompt_token_ids = await _coupled_stream(client, messages, **extra)
+    finish_reason = coupled_choices[-1]["finish_reason"]
+    disagg = await stream_chat_derender(
+        client,
+        [tid for ch in coupled_choices for tid in ch["token_ids"]],
+        [len(ch["token_ids"]) for ch in coupled_choices]
+        + ([0] if finish_only_tail else []),
+        {"model": MODEL, "messages": messages, **extra},
+        len(prompt_token_ids),
+        prompt_token_ids,
+        finish_reason="stop" if finish_reason == "tool_calls" else finish_reason,
+    )
+    coupled = assemble_stream(coupled_choices)
+    for message in (coupled, disagg):
+        for tc in message["tool_calls"]:
+            tc["id"] = _tool_call_id_shape(tc["id"])
+    assert disagg == coupled
+    return coupled, coupled_choices
+
+
+# ---------------------------------------------------------------------------
+# Streaming parity cases
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_stream_parity_plain(client):
+    """Plain streamed detokenization parity."""
+    messages = [
+        {"role": "user", "content": "What is 2+2? Answer in one short sentence."}
+    ]
+    await _run_stream_parity_case(client, messages)
+
+
+@pytest.mark.asyncio
+async def test_stream_parity_reasoning(client):
+    """Streamed reasoning/content split parity."""
+    messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}]
+    coupled, _ = await _run_stream_parity_case(
+        client, messages, include_reasoning=True, max_tokens=256
+    )
+
+    if not coupled["reasoning"]:
+        pytest.skip("Model did not emit a <think> block")
+
+
+@pytest.mark.asyncio
+async def test_stream_parity_tool_call(client):
+    """Streamed tool call parity under `tool_choice="auto"`, where a tool
+    call also rewrites `finish_reason` to `"tool_calls"`."""
+    messages = [{"role": "user", "content": "What's the weather in Paris?"}]
+    coupled, _ = await _run_stream_parity_case(
+        client, messages, tools=TOOLS, tool_choice="auto", max_tokens=512
+    )
+
+    if not coupled["tool_calls"]:
+        pytest.skip("Model did not emit a tool call")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "chunking", ["engine_steps", "multi_token", "finish_only_tail"]
+)
+async def test_stream_parity_reasoning_and_tool_call(client, chunking):
+    """Streamed reasoning + forced tool call parity across chunk shapes.
+
+    `multi_token` batches tokens per chunk with `stream_interval` the way
+    speculative decoding does, so replay has to reuse the original chunk
+    boundaries. `finish_only_tail` sends the finish reason on a chunk with no
+    tokens, which must still flush buffered tool call arguments.
+    """
+    messages = [{"role": "user", "content": "What's the weather in Paris?"}]
+    extra: dict = {
+        "tools": TOOLS,
+        "tool_choice": FORCE_WEATHER_TOOL,
+        "include_reasoning": True,
+        "max_tokens": 512,
+    }
+    if chunking == "multi_token":
+        extra["stream_interval"] = 4
+    coupled, coupled_choices = await _run_stream_parity_case(
+        client,
+        messages,
+        finish_only_tail=chunking == "finish_only_tail",
+        **extra,
+    )
+
+    if chunking == "multi_token":
+        assert any(len(ch["token_ids"]) > 1 for ch in coupled_choices)
+    if not (coupled["reasoning"] and coupled["tool_calls"]):
+        pytest.skip("Model did not emit both a <think> block and a tool call")
diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py
index b783f40ae3c6..8510d63c4b65 100644
--- a/tests/entrypoints/scale_out/derender/test_derender_stream.py
+++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py
@@ -5,29 +5,158 @@
 
 Tests are split into two layers:
 
-1. Unit tests (no server): covers ``_detokenize_delta`` correctness
-   (chunked == one-shot) and ``derender_completion_stream`` /
-   ``derender_chat_stream`` logic via a real tokenizer on a tiny model.
+1. Unit tests (no server): covers `_detokenize_delta` correctness
+   (chunked == one-shot) and `derender_completion_stream` /
+   `derender_chat_stream` logic via a real tokenizer on a tiny model.
+   The parser path is covered both with a deterministic stub parser and
+   with the real `HarmonyParser` (skipped without `openai_harmony`).
 
 2. Integration tests (require a running render server): covers the full
    HTTP round-trip through the streaming endpoint.  Marked with
-   ``@pytest.mark.asyncio`` and gated by the ``server`` / ``client``
-   fixtures from the sibling ``test_derender.py``.
+   `@pytest.mark.asyncio` and gated by the `server` / `client`
+   fixtures from the sibling `test_derender.py`.
 """
 
+import json
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
 import pytest
 import pytest_asyncio
 
+from tests.entrypoints.scale_out.derender.utils import stream_chat_derender
+from vllm.entrypoints.generate.base.protocol import (
+    DeltaFunctionCall,
+    DeltaMessage,
+    DeltaToolCall,
+    PerRequestMetrics,
+    SpeculativeDecodingMetrics,
+)
 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
     DerenderStreamState,
+    GenerateResponse,
+    GenerateResponseChoice,
     GenerateResponseStreamChoice,
     GenerateStreamResponse,
 )
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo
+from vllm.parser import Parser
+from vllm.utils import random_uuid
 
 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
 
 
+def _merge_fake_deltas(
+    accumulated: DeltaMessage | None, new: DeltaMessage | None
+) -> DeltaMessage | None:
+    """Fold one token's worth of `_FakeParser` output into a per call
+    accumulator, the same way a real parser merges multiple tokens'
+    effects into the single `DeltaMessage` it returns per `parse_delta`
+    call.
+    """
+    if new is None:
+        return accumulated
+    if accumulated is None:
+        accumulated = DeltaMessage()
+
+    if new.content:
+        accumulated.content = (accumulated.content or "") + new.content
+    if new.reasoning:
+        accumulated.reasoning = (accumulated.reasoning or "") + new.reasoning
+
+    for tc in new.tool_calls:
+        starts_new_call = tc.id is not None or (
+            tc.function is not None and tc.function.name is not None
+        )
+        existing = (
+            None
+            if starts_new_call
+            else next((t for t in accumulated.tool_calls if t.index == tc.index), None)
+        )
+        if existing is None:
+            accumulated.tool_calls.append(tc)
+            continue
+        if existing.function is None:
+            existing.function = DeltaFunctionCall()
+        if tc.function is not None and tc.function.arguments:
+            existing.function.arguments = (
+                existing.function.arguments or ""
+            ) + tc.function.arguments
+
+    return accumulated
+
+
+class _FakeParser(Parser):
+    """Deterministic `Parser` stub for unit testing the replay/merge/pin
+    mechanics of `OnlineDerenderer._derender_chat_stream_parsed` in
+    isolation from any real reasoning/tool parser's markup grammar.
+
+    Processes every id in `delta_token_ids` in order and folds their
+    effects into one `DeltaMessage`, like a real parser handling a
+    multi-token delta in a single call. Ignores `delta_text`:
+
+    - `TOOL_START` opens a new tool call at index 0 (id + name).
+    - `TOOL_ARG` appends one `"a"` to that tool call's arguments.
+    - `REASON` emits one `"r"` of reasoning (suppressed when the request
+      has `include_reasoning=False`, mirroring the real parsers).
+    - `CONTENT` emits one `"c"` of content.
+    - An empty `delta_token_ids` with `finished=True` (the finish only
+      flush call) emits a `"FLUSH"` content sentinel so tests can confirm
+      it happened.
+    - Anything else emits nothing.
+    """
+
+    TOOL_START = 9
+    TOOL_ARG = 10
+    REASON = 11
+    CONTENT = 12
+
+    def parse_delta(
+        self,
+        delta_text,
+        delta_token_ids,
+        request,
+        prompt_token_ids=None,
+        *,
+        finished,
+    ):
+        if not delta_token_ids:
+            return DeltaMessage(content="FLUSH") if finished else None
+
+        accumulated: DeltaMessage | None = None
+        for tok in delta_token_ids:
+            accumulated = _merge_fake_deltas(
+                accumulated, self._parse_one_token(tok, request)
+            )
+        return accumulated
+
+    def _parse_one_token(self, tok: int, request) -> DeltaMessage | None:
+        if tok == self.TOOL_START:
+            return DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(
+                        id=f"call-{random_uuid()}",
+                        type="function",
+                        function=DeltaFunctionCall(name="get_weather", arguments=""),
+                        index=0,
+                    )
+                ]
+            )
+        if tok == self.TOOL_ARG:
+            return DeltaMessage(
+                tool_calls=[
+                    DeltaToolCall(index=0, function=DeltaFunctionCall(arguments="a"))
+                ]
+            )
+        if tok == self.REASON:
+            if not request.include_reasoning:
+                return None
+            return DeltaMessage(reasoning="r")
+        if tok == self.CONTENT:
+            return DeltaMessage(content="c")
+        return None
+
+
 # ---------------------------------------------------------------------------
 # Helpers shared across tests
 # ---------------------------------------------------------------------------
@@ -39,9 +168,9 @@ def _make_stream_chunk(
     finish_reason: str | None = None,
     request_id: str = "test-req",
     usage: dict | None = None,
+    metrics: PerRequestMetrics | None = None,
 ) -> GenerateStreamResponse:
     """Build a GenerateStreamResponse SSE chunk."""
-
     return GenerateStreamResponse(
         request_id=request_id,
         choices=[
@@ -52,6 +181,7 @@ def _make_stream_chunk(
             )
         ],
         usage=UsageInfo(**usage) if usage else None,
+        metrics=metrics,
     )
 
 
@@ -61,7 +191,6 @@ def _make_usage_chunk(
     request_id: str = "test-req",
 ) -> GenerateStreamResponse:
     """Build a usage only final SSE chunk (empty choices)."""
-
     return GenerateStreamResponse(
         request_id=request_id,
         choices=[],
@@ -87,7 +216,7 @@ def tokenizer():
 
 
 @pytest.fixture(scope="module")
-def derenderer(tokenizer):
+def derenderer(tokenizer, request):
     """Construct a minimal OnlineDerenderer backed by a stub renderer."""
     from unittest.mock import MagicMock
 
@@ -96,6 +225,10 @@ def derenderer(tokenizer):
     renderer = MagicMock()
     renderer.get_tokenizer.return_value = tokenizer
 
+    executor = ThreadPoolExecutor(max_workers=1)
+    request.addfinalizer(executor.shutdown)
+    renderer._executor = executor
+
     model_config = MagicMock()
     model_config.hf_config.model_type = "llama"
     model_config.model = MODEL_NAME
@@ -113,6 +246,38 @@ def derenderer(tokenizer):
     )
 
 
+@pytest.fixture(scope="module")
+def parsed_derenderer(tokenizer):
+    """OnlineDerenderer with `_FakeParser` wired in as `self.parser`.
+
+    The parser-configured path runs `_derender_chat_stream_parsed` via
+    `make_async(..., executor=renderer._executor)`, so
+    `renderer._executor` must be a real `ThreadPoolExecutor`.
+    `loop.run_in_executor` cannot submit work to a `MagicMock`.
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.renderers.online_derenderer import OnlineDerenderer
+
+    renderer = MagicMock()
+    renderer.get_tokenizer.return_value = tokenizer
+    renderer._executor = ThreadPoolExecutor(max_workers=2)
+
+    model_config = MagicMock()
+    model_config.hf_config.model_type = "llama"
+    model_config.model = MODEL_NAME
+
+    dr = OnlineDerenderer(
+        model_config=model_config,
+        renderer=renderer,
+        request_logger=None,
+        chat_template=None,
+        chat_template_content_format="string",
+    )
+    dr.parser = _FakeParser
+    return dr
+
+
 class TestDetokenizeDelta:
     """_detokenize_delta: chunked decode must equal one shot decode."""
 
@@ -195,7 +360,7 @@ def test_state_window_stays_bounded(self, derenderer, tokenizer):
         """prev_tokens must not grow with the number of chunks (bounded transport).
 
         Guards that the carried decode window is a small constant
-        tail, so cumulative ``stream_state`` transport is O(n) and not O(n^2).
+        tail, so cumulative `stream_state` transport is O(n) and not O(n^2).
         """
         token_ids = tokenizer.encode(
             "a reasonably long ascii stream of tokens used to exercise the "
@@ -229,6 +394,46 @@ def test_n_independent_streams_same_result(self, derenderer, tokenizer):
         assert results[0] == self._one_shot(tokenizer, token_ids)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "stream_method",
+    ["derender_chat_stream", "derender_completion_stream"],
+)
+async def test_streaming_detokenization_runs_off_event_loop(
+    derenderer, monkeypatch, stream_method
+):
+    """Streaming detokenization runs on the renderer executor."""
+    import vllm.renderers.online_derenderer as online_derenderer_module
+
+    event_loop_thread_id = threading.get_ident()
+    executor_thread_id = derenderer.renderer._executor.submit(
+        threading.get_ident
+    ).result()
+
+    detokenize_thread_id = None
+    original_detokenize = online_derenderer_module.detokenize_incrementally
+
+    def record_detokenize_thread(*args, **kwargs):
+        nonlocal detokenize_thread_id
+        detokenize_thread_id = threading.get_ident()
+        return original_detokenize(*args, **kwargs)
+
+    monkeypatch.setattr(
+        online_derenderer_module,
+        "detokenize_incrementally",
+        record_detokenize_thread,
+    )
+
+    await getattr(derenderer, stream_method)(
+        model=MODEL_NAME,
+        generate_chunk=_make_stream_chunk([10]),
+    )
+
+    assert detokenize_thread_id is not None
+    assert detokenize_thread_id == executor_thread_id
+    assert detokenize_thread_id != event_loop_thread_id
+
+
 class TestDerenderCompletionStream:
     """derender_completion_stream: streaming output parity with one shot."""
 
@@ -355,6 +560,26 @@ async def test_role_on_first_chunk_only(self, derenderer, tokenizer):
         assert chunk1.choices[0].delta.role == "assistant"
         assert chunk2.choices[0].delta.role is None
 
+    @pytest.mark.asyncio
+    async def test_metrics_passthrough(self, derenderer):
+        metrics = PerRequestMetrics(
+            speculative_decoding=SpeculativeDecodingMetrics(
+                mean_acceptance_length=2.0,
+                draft_acceptance_rate=0.5,
+                acceptance_histogram=[0, 1],
+                num_spec_steps=1,
+                num_accepted_draft_tokens=1,
+                num_draft_tokens=2,
+                num_spec_tokens=1,
+            )
+        )
+        chunk, _ = await derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([], metrics=metrics),
+            state=DerenderStreamState(),
+        )
+        assert chunk.metrics == metrics
+
     @pytest.mark.asyncio
     async def test_chunked_equals_oneshot(self, derenderer, tokenizer):
         """Sum of streaming content deltas == one shot decode."""
@@ -380,47 +605,453 @@ async def test_chunked_equals_oneshot(self, derenderer, tokenizer):
         assert streamed == one_shot
 
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("with_chat_request", [True, False])
-    async def test_parser_raises_not_implemented(self, tokenizer, with_chat_request):
-        """Stream chat with a parser active must fail closed (NotImplementedError).
+    async def test_parser_configured_missing_chat_request_raises(
+        self, parsed_derenderer
+    ):
+        """A parser configured model must never fall through to plain detok
+        even when `chat_request` is omitted. If allowed this would leak raw
+        reasoning/tool markup into `delta.content`. `ServingDerender`
+        pre-checks this too (400 before touching the tokenizer). This pins
+        the `OnlineDerenderer` level backstop."""
+        with pytest.raises(ValueError, match="chat_request"):
+            await parsed_derenderer.derender_chat_stream(
+                model=MODEL_NAME,
+                generate_chunk=_make_stream_chunk([_FakeParser.CONTENT]),
+                state=None,
+                chat_request=None,
+            )
 
-        Checks that a parser configured model must 501 even when ``chat_request`` is
-        omitted, otherwise reasoning/tool markup would leak into ``delta.content`
-          via the plain detok fallback.
-        """
-        from unittest.mock import MagicMock
 
-        from vllm.renderers.online_derenderer import OnlineDerenderer
+def _chat_request(**kwargs):
+    from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 
-        renderer = MagicMock()
-        renderer.get_tokenizer.return_value = tokenizer
+    kwargs.setdefault("messages", [{"role": "user", "content": "hi"}])
+    kwargs.setdefault("model", MODEL_NAME)
+    return ChatCompletionRequest(**kwargs)
 
-        model_config = MagicMock()
-        model_config.hf_config.model_type = "llama"
-        model_config.model = MODEL_NAME
 
-        # Construct a derenderer WITH a tool/reasoning parser active.
-        # ParserManager.get_parser returns None when no parser is named, so
-        # we inject a mock parser class directly.
-        dr = OnlineDerenderer(
-            model_config=model_config,
-            renderer=renderer,
-            request_logger=None,
-            chat_template=None,
-            chat_template_content_format="string",
+class TestDerenderChatStreamParsed:
+    """derender_chat_stream: parser branch (replay + parse_delta) exercised
+    against the deterministic `_FakeParser` so these pin OnlineDerenderer's
+    own replay/merge/pin/finish_reason logic independent of any real
+    parser's markup grammar (covered separately by the parser_server backed
+    integration tests below)."""
+
+    @pytest.mark.asyncio
+    async def test_dispatches_and_emits_content(self, parsed_derenderer):
+        chunk, state = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                [_FakeParser.CONTENT], finish_reason="stop"
+            ),
+            chat_request=_chat_request(),
         )
-        dr.parser = MagicMock()  # simulate a parser being active
+        assert chunk.choices[0].delta.content == "c"
+        assert chunk.choices[0].delta.role == "assistant"
+        assert chunk.choices[0].finish_reason == "stop"
+        assert state.output_token_ids == [_FakeParser.CONTENT]
 
-        chat_request = MagicMock() if with_chat_request else None
-        token_ids = tokenizer.encode("hello")[:3]
+    @pytest.mark.asyncio
+    async def test_role_sent_once(self, parsed_derenderer):
+        chat_request = _chat_request()
+        chunk1, state = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([_FakeParser.CONTENT]),
+            chat_request=chat_request,
+        )
+        chunk2, _ = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                [_FakeParser.CONTENT], finish_reason="stop"
+            ),
+            state=state,
+            chat_request=chat_request,
+        )
+        assert chunk1.choices[0].delta.role == "assistant"
+        assert chunk2.choices[0].delta.role is None
 
-        with pytest.raises(NotImplementedError):
-            await dr.derender_chat_stream(
-                model=MODEL_NAME,
-                generate_chunk=_make_stream_chunk(token_ids),
-                state=None,
-                chat_request=chat_request,
+    @pytest.mark.asyncio
+    async def test_finish_only_chunk_flushes(self, parsed_derenderer):
+        """A finish only chunk (no new tokens) still calls `parse_delta`
+        once with `finished=True` to flush buffered state."""
+        chat_request = _chat_request()
+        _, state = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([_FakeParser.CONTENT]),
+            chat_request=chat_request,
+        )
+        chunk2, _ = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([], finish_reason="stop"),
+            state=state,
+            chat_request=chat_request,
+        )
+        assert chunk2.choices[0].delta.content == "FLUSH"
+        assert chunk2.choices[0].finish_reason == "stop"
+
+    @pytest.mark.asyncio
+    async def test_chunking_invariance(self, parsed_derenderer):
+        """1 token per chunk vs a single whole chunk assemble identically.
+
+        Every chunk (live or replayed) is fed to `parse_delta` at its own
+        producer granularity (see `_derender_chat_stream_parsed`'s
+        docstring). This proves that choice doesn't change the assembled
+        output for a well behaved incremental parser. Each call's effects
+        fold together the same way regardless of how many tokens land in
+        one call vs. many."""
+        token_ids = [
+            _FakeParser.REASON,
+            _FakeParser.REASON,
+            _FakeParser.TOOL_START,
+            _FakeParser.TOOL_ARG,
+            _FakeParser.TOOL_ARG,
+            _FakeParser.CONTENT,
+        ]
+        chat_request = _chat_request(
+            tools=[{"type": "function", "function": {"name": "get_weather"}}],
+            tool_choice="auto",
+            include_reasoning=True,
+        )
+
+        async def _assemble(chunks: list[list[int]]) -> dict:
+            state = None
+            content = ""
+            reasoning = ""
+            tool_args = ""
+            for i, tids in enumerate(chunks):
+                finish = "stop" if i == len(chunks) - 1 else None
+                chunk, state = await parsed_derenderer.derender_chat_stream(
+                    model=MODEL_NAME,
+                    generate_chunk=_make_stream_chunk(tids, finish_reason=finish),
+                    state=state,
+                    chat_request=chat_request,
+                )
+                delta = chunk.choices[0].delta
+                content += delta.content or ""
+                reasoning += delta.reasoning or ""
+                for tc in delta.tool_calls:
+                    if tc.function and tc.function.arguments:
+                        tool_args += tc.function.arguments
+            return {"content": content, "reasoning": reasoning, "tool_args": tool_args}
+
+        whole = await _assemble([token_ids])
+        one_at_a_time = await _assemble([[t] for t in token_ids])
+
+        assert whole == one_at_a_time
+        assert whole == {"content": "c", "reasoning": "rr", "tool_args": "aa"}
+
+    @pytest.mark.asyncio
+    async def test_replay_preserves_chunk_boundaries(
+        self, parsed_derenderer, monkeypatch
+    ):
+        """A multi-token chunk (e.g. a speculative decoding step) must be
+        replayed as one `parse_delta` call, not one per token. Parsers
+        with per-call state such as `history_tool_call_cnt` would
+        otherwise diverge from standard serving."""
+
+        class _CallCountingParser(_FakeParser):
+            """Emits how many `parse_delta` calls it has seen."""
+
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.num_calls = 0
+
+            def parse_delta(
+                self,
+                delta_text,
+                delta_token_ids,
+                request,
+                prompt_token_ids=None,
+                *,
+                finished,
+            ):
+                self.num_calls += 1
+                return DeltaMessage(content=str(self.num_calls))
+
+        monkeypatch.setattr(parsed_derenderer, "parser", _CallCountingParser)
+        chat_request = _chat_request()
+
+        _, state = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([_FakeParser.CONTENT] * 3),
+            chat_request=chat_request,
+        )
+        assert state.output_chunk_lens == [3]
+
+        chunk, _ = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                [_FakeParser.CONTENT], finish_reason="stop"
+            ),
+            state=state,
+            chat_request=chat_request,
+        )
+        assert chunk.choices[0].delta.content == "2"
+
+    @pytest.mark.asyncio
+    async def test_tool_call_id_pinned_across_chunks(self, parsed_derenderer):
+        """Once an index's ID is recorded in `last_tool_call_ids`, a later
+        id bearing delta for that same index is pinned to the recorded
+        value rather than a freshly (re-)generated one."""
+        chat_request = _chat_request(
+            tools=[{"type": "function", "function": {"name": "get_weather"}}],
+            tool_choice="auto",
+        )
+        chunk1, state = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([_FakeParser.TOOL_START]),
+            chat_request=chat_request,
+        )
+        first_id = chunk1.choices[0].delta.tool_calls[0].id
+        assert first_id is not None
+        assert state.last_tool_call_ids == [first_id]
+
+        chunk2, state = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                [_FakeParser.TOOL_START], finish_reason="stop"
+            ),
+            state=state,
+            chat_request=chat_request,
+        )
+        assert chunk2.choices[0].delta.tool_calls[0].id == first_id
+        assert state.last_tool_call_ids == [first_id]
+
+    @pytest.mark.asyncio
+    async def test_finish_reason_rewritten_to_tool_calls(self, parsed_derenderer):
+        chat_request = _chat_request(
+            tools=[{"type": "function", "function": {"name": "get_weather"}}],
+            tool_choice="auto",
+        )
+        chunk, _ = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                [_FakeParser.TOOL_START, _FakeParser.TOOL_ARG], finish_reason="stop"
+            ),
+            chat_request=chat_request,
+        )
+        assert chunk.choices[0].finish_reason == "tool_calls"
+
+    @pytest.mark.asyncio
+    async def test_finish_reason_stop_for_named_tool_choice(self, parsed_derenderer):
+        # ChatCompletionRequest's `check_tool_usage` is a mode="before"
+        # validator, so it sees the raw value and only accepts "auto",
+        # "required" or a dict, never an already built
+        # ChatCompletionNamedToolChoiceParam.
+        chat_request = _chat_request(
+            tools=[{"type": "function", "function": {"name": "get_weather"}}],
+            tool_choice={"type": "function", "function": {"name": "get_weather"}},
+        )
+        chunk, _ = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                [_FakeParser.TOOL_START, _FakeParser.TOOL_ARG], finish_reason="stop"
+            ),
+            chat_request=chat_request,
+        )
+        assert chunk.choices[0].finish_reason == "stop"
+
+    @pytest.mark.asyncio
+    async def test_include_reasoning_false_suppresses_reasoning(
+        self, parsed_derenderer
+    ):
+        chat_request = _chat_request(include_reasoning=False)
+        chunk, _ = await parsed_derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                [_FakeParser.REASON, _FakeParser.CONTENT], finish_reason="stop"
+            ),
+            chat_request=chat_request,
+        )
+        delta = chunk.choices[0].delta
+        assert delta.reasoning is None
+        assert delta.content == "c"
+
+
+# ---------------------------------------------------------------------------
+# Harmony / GPT-OSS replay — unit, no server
+# ---------------------------------------------------------------------------
+
+HARMONY_MODEL = "openai/gpt-oss-20b"
+HARMONY_REASONING = "The user wants 2 plus 3."
+HARMONY_ANSWER = "The answer is 5."
+HARMONY_PROMPT = "<|start|>user<|message|>Add 2 and 3.<|end|><|start|>assistant"
+HARMONY_OUTPUT = (
+    f"<|channel|>analysis<|message|>{HARMONY_REASONING}<|end|>"
+    f"<|start|>assistant<|channel|>final<|message|>{HARMONY_ANSWER}<|return|>"
+)
+
+
+@pytest.fixture(scope="module")
+def harmony_encode():
+    """Encoder for canned GPT-OSS harmony token sequences."""
+    pytest.importorskip("openai_harmony")
+
+    # Pre-caches the o200k_base BPE file that openai-harmony's Rust backend
+    # downloads on first use, same as the sibling suite's E2E harmony tests.
+    from tests.entrypoints.scale_out.derender.test_derender import (
+        _ensure_harmony_vocab,
+    )
+    from vllm.entrypoints.openai.parser.harmony_utils import get_encoding
+
+    _ensure_harmony_vocab()
+    encoding = get_encoding()
+
+    def _encode(harmony_str: str) -> list[int]:
+        return encoding.encode(harmony_str, allowed_special="all")
+
+    return _encode
+
+
+@pytest.fixture(scope="module")
+def harmony_tokenizer():
+    pytest.importorskip("openai_harmony")
+    from vllm.tokenizers import get_tokenizer
+
+    return get_tokenizer(HARMONY_MODEL, trust_remote_code=True)
+
+
+@pytest.fixture(scope="module")
+def harmony_derenderer(harmony_tokenizer):
+    """OnlineDerenderer whose parser resolves to the real `HarmonyParser`.
+
+    Same shape as `parsed_derenderer` (mocked renderer, real executor) but
+    with a real tokenizer and parser, so the replay path is exercised
+    against Harmony's actual channel grammar rather than a stub.
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.parser.harmony import HarmonyParser
+    from vllm.renderers.online_derenderer import OnlineDerenderer
+
+    renderer = MagicMock()
+    renderer.get_tokenizer.return_value = harmony_tokenizer
+    renderer._executor = ThreadPoolExecutor(max_workers=2)
+
+    model_config = MagicMock()
+    model_config.model = HARMONY_MODEL
+    model_config.hf_config.model_type = "gpt_oss"
+    model_config.hf_text_config.model_type = "gpt_oss"
+    model_config.hf_overrides = None
+
+    dr = OnlineDerenderer(
+        model_config=model_config,
+        renderer=renderer,
+        request_logger=None,
+        chat_template=None,
+        chat_template_content_format="string",
+        enable_auto_tools=True,
+        tool_parser="openai",
+        reasoning_parser="openai_gptoss",
+    )
+    assert dr.use_harmony
+    assert dr.parser is HarmonyParser
+    return dr
+
+
+async def _stream_harmony_deltas(
+    derenderer,
+    chat_request,
+    output_ids: list[int],
+    prompt_ids: list[int],
+    chunk_size: int,
+) -> list[DeltaMessage]:
+    """Drive `output_ids` through the parser path in fixed size chunks.
+
+    Returns the per chunk `DeltaMessage`s in order so callers can assert on
+    intermediate emissions, not just the assembled result.
+    """
+    deltas: list[DeltaMessage] = []
+    state = None
+    for start in range(0, len(output_ids), chunk_size):
+        tids = output_ids[start : start + chunk_size]
+        is_last = start + chunk_size >= len(output_ids)
+        chunk, state = await derenderer.derender_chat_stream(
+            model=HARMONY_MODEL,
+            generate_chunk=_make_stream_chunk(
+                tids, finish_reason="stop" if is_last else None
+            ),
+            state=state,
+            chat_request=chat_request,
+            prompt_token_ids=prompt_ids,
+        )
+        deltas.append(chunk.choices[0].delta)
+    return deltas
+
+
+class TestDerenderChatStreamHarmony:
+    """derender_chat_stream: replay + `parse_delta` against real HarmonyParser.
+
+    Skipped where `openai_harmony` is not installed. `prompt_token_ids` is
+    threaded through for call shape fidelity only, since HarmonyParser
+    derives its state from the output tokens alone.
+    """
+
+    @pytest.mark.asyncio
+    async def test_reasoning_never_leaks_as_content_midstream(
+        self, harmony_derenderer, harmony_encode
+    ):
+        """Regression guard for the RFC's original replay + diff design.
+
+        `HarmonyParser.parse()` always flushes to EOS. Mid-stream that
+        raises `HarmonyError` and the recovery branch re-emits the in
+        flight message on the `final` channel, so partial analysis surfaces
+        as content. Replay + `parse_delta` must never do that. Driven one
+        token per chunk, the finest granularity a client can produce.
+        """
+        output_ids = harmony_encode(HARMONY_OUTPUT)
+        prompt_ids = harmony_encode(HARMONY_PROMPT)
+        chat_request = _chat_request(model=HARMONY_MODEL, include_reasoning=True)
+
+        deltas = await _stream_harmony_deltas(
+            harmony_derenderer, chat_request, output_ids, prompt_ids, chunk_size=1
+        )
+
+        reasoning = ""
+        content = ""
+        for i, delta in enumerate(deltas):
+            reasoning += delta.reasoning or ""
+            content += delta.content or ""
+            # Any analysis text emitted as content breaks the prefix
+            # property, since the two channels share no prefix here.
+            assert HARMONY_ANSWER.startswith(content), (
+                f"content after chunk {i} is not a prefix of the final "
+                f"channel text: {content!r}"
             )
+            assert HARMONY_REASONING.startswith(reasoning)
+            assert not delta.tool_calls
+
+        assert reasoning == HARMONY_REASONING
+        assert content == HARMONY_ANSWER
+
+    @pytest.mark.asyncio
+    async def test_stream_matches_batch(self, harmony_derenderer, harmony_encode):
+        """Streamed assembly equals one shot `/derender` over the same IDs."""
+        output_ids = harmony_encode(HARMONY_OUTPUT)
+        prompt_ids = harmony_encode(HARMONY_PROMPT)
+        chat_request = _chat_request(model=HARMONY_MODEL, include_reasoning=True)
+
+        deltas = await _stream_harmony_deltas(
+            harmony_derenderer, chat_request, output_ids, prompt_ids, chunk_size=3
+        )
+        reasoning = "".join(d.reasoning or "" for d in deltas)
+        content = "".join(d.content or "" for d in deltas)
+
+        batch_choices = await harmony_derenderer.derender_chat(
+            GenerateResponse(
+                request_id="test-harmony-batch",
+                choices=[
+                    GenerateResponseChoice(
+                        index=0, token_ids=output_ids, finish_reason="stop"
+                    )
+                ],
+            ),
+            chat_request,
+        )
+        message = batch_choices[0].message
+
+        assert reasoning == message.reasoning
+        assert content == message.content
 
 
 class TestDerenderStreamStateValidation:
@@ -448,6 +1079,24 @@ def test_prev_tokens_at_cap_accepted(self):
         state = DerenderStreamState(prev_tokens=["a"] * 1024)
         assert len(state.prev_tokens) == 1024
 
+    def test_output_chunk_lens_mismatch_rejected(self):
+        from pydantic import ValidationError
+
+        with pytest.raises(ValidationError):
+            DerenderStreamState(output_token_ids=[1, 2, 3], output_chunk_lens=[1, 1])
+
+    def test_output_chunk_lens_zero_entry_rejected(self):
+        from pydantic import ValidationError
+
+        with pytest.raises(ValidationError):
+            DerenderStreamState(output_token_ids=[1, 2], output_chunk_lens=[0, 2])
+
+    def test_output_chunk_lens_matching_accepted(self):
+        state = DerenderStreamState(
+            output_token_ids=[1, 2, 3], output_chunk_lens=[1, 2]
+        )
+        assert state.output_chunk_lens == [1, 2]
+
 
 class TestServingDerenderStreamErrorHandling:
     """Malformed stream_state must surface as 400 and not an unhandled 500."""
@@ -460,8 +1109,10 @@ def _make_serving(self, side_effect: Exception):
         models = MagicMock()
         models.is_base_model.return_value = True
         models.model_config = MagicMock()
+        models.model_config.max_model_len = 100_000
 
         online_derenderer = MagicMock()
+        online_derenderer.parser = None
         online_derenderer.derender_completion_stream = AsyncMock(
             side_effect=side_effect
         )
@@ -506,6 +1157,179 @@ async def test_chat_stream_bad_state_returns_400(self, exc):
         assert result.error.code == 400
 
 
+class TestServingDerenderStreamValidation:
+    """If checks cannot run in `derender_chat_stream_response`: all must
+    reject with 400 before ever calling into `online_derenderer` (i.e.
+    before touching the tokenizer)."""
+
+    def _make_serving(self, *, parser_configured: bool, max_model_len: int = 100_000):
+        from unittest.mock import AsyncMock, MagicMock
+
+        from vllm.entrypoints.openai.chat_completion.protocol import (
+            ChatCompletionStreamResponse,
+        )
+        from vllm.entrypoints.scale_out.derender.serving import ServingDerender
+
+        models = MagicMock()
+        models.is_base_model.return_value = True
+        models.model_config = MagicMock()
+        models.model_config.max_model_len = max_model_len
+
+        online_derenderer = MagicMock()
+        online_derenderer.parser = MagicMock() if parser_configured else None
+        online_derenderer.derender_chat_stream = AsyncMock(
+            return_value=(
+                ChatCompletionStreamResponse(id="t", model=MODEL_NAME, choices=[]),
+                DerenderStreamState(),
+            )
+        )
+        return ServingDerender(models=models, online_derenderer=online_derenderer)
+
+    @pytest.mark.asyncio
+    async def test_missing_chat_request_with_parser_rejected(self):
+        from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
+            DerenderChatStreamRequest,
+        )
+
+        serving = self._make_serving(parser_configured=True)
+        request = DerenderChatStreamRequest(
+            stream=True,
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([1, 2]),
+        )
+        result = await serving.derender_chat_stream_response(request)
+        assert isinstance(result, ErrorResponse)
+        assert result.error.code == 400
+        assert "chat_request" in result.error.message
+        serving.online_derenderer.derender_chat_stream.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_missing_chat_request_without_parser_ok(self):
+        from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
+            DerenderChatStreamRequest,
+        )
+
+        serving = self._make_serving(parser_configured=False)
+        request = DerenderChatStreamRequest(
+            stream=True,
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([1, 2]),
+        )
+        result = await serving.derender_chat_stream_response(request)
+        assert not isinstance(result, ErrorResponse)
+
+    @pytest.mark.asyncio
+    async def test_missing_prompt_token_ids_with_parser_rejected(self):
+        """A parser configured model must reject a missing prompt_token_ids
+        the same way it rejects a missing chat_request. Without it,
+        parse_delta cannot tell whether the prompt left reasoning open and
+        would misclassify reasoning content as plain content."""
+        from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
+            DerenderChatStreamRequest,
+        )
+
+        serving = self._make_serving(parser_configured=True)
+        request = DerenderChatStreamRequest(
+            stream=True,
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([1, 2]),
+            chat_request=_chat_request(),
+        )
+        result = await serving.derender_chat_stream_response(request)
+        assert isinstance(result, ErrorResponse)
+        assert result.error.code == 400
+        assert "prompt_token_ids" in result.error.message
+        serving.online_derenderer.derender_chat_stream.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_prompt_token_ids_present_with_parser_ok(self):
+        from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
+            DerenderChatStreamRequest,
+        )
+
+        serving = self._make_serving(parser_configured=True)
+        request = DerenderChatStreamRequest(
+            stream=True,
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([1, 2]),
+            chat_request=_chat_request(),
+            prompt_token_ids=[1, 2, 3],
+        )
+        result = await serving.derender_chat_stream_response(request)
+        assert not isinstance(result, ErrorResponse)
+
+    @pytest.mark.asyncio
+    async def test_oversized_output_token_ids_rejected(self):
+        from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
+            DerenderChatStreamRequest,
+        )
+
+        serving = self._make_serving(parser_configured=False, max_model_len=4)
+        request = DerenderChatStreamRequest(
+            stream=True,
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([1, 2, 3]),
+            stream_state=DerenderStreamState(
+                output_token_ids=[1, 2], output_chunk_lens=[2]
+            ),
+        )
+        result = await serving.derender_chat_stream_response(request)
+        assert isinstance(result, ErrorResponse)
+        assert result.error.code == 400
+        assert "max_model_len" in result.error.message
+        serving.online_derenderer.derender_chat_stream.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_oversized_prompt_token_ids_rejected(self):
+        """prompt_token_ids is caller supplied and otherwise unbounded. A
+        parser configured deployment rescans it in full on every chunk
+        (is_reasoning_end / adjust_initial_state_from_prompt), so it must be
+        bounded the same way output_token_ids is."""
+        from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
+            DerenderChatStreamRequest,
+        )
+
+        serving = self._make_serving(parser_configured=True, max_model_len=4)
+        request = DerenderChatStreamRequest(
+            stream=True,
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk([1]),
+            chat_request=_chat_request(),
+            prompt_token_ids=[1, 2, 3, 4, 5],
+        )
+        result = await serving.derender_chat_stream_response(request)
+        assert isinstance(result, ErrorResponse)
+        assert result.error.code == 400
+        assert "max_model_len" in result.error.message
+        serving.online_derenderer.derender_chat_stream.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_too_many_choices_rejected(self):
+        """Each streamed chunk contains at most one choice (a single
+        DerenderStreamState is threaded through every choice). The check
+        could never fire since derender_chat_stream itself rejects anything
+        above 1 first."""
+        from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
+            DerenderChatStreamRequest,
+        )
+
+        serving = self._make_serving(parser_configured=False)
+        two_choices = GenerateStreamResponse(
+            request_id="t",
+            choices=[
+                GenerateResponseStreamChoice(index=i, token_ids=[1]) for i in range(2)
+            ],
+        )
+        request = DerenderChatStreamRequest(
+            stream=True, model=MODEL_NAME, generate_chunk=two_choices
+        )
+        result = await serving.derender_chat_stream_response(request)
+        assert isinstance(result, ErrorResponse)
+        assert result.error.code == 400
+        assert "at most one choice" in result.error.message
+        serving.online_derenderer.derender_chat_stream.assert_not_called()
+
+
 # ---------------------------------------------------------------------------
 # Integration tests — require a live render server
 # ---------------------------------------------------------------------------
@@ -531,7 +1355,6 @@ async def client(server):
 
 async def _render_chat(client) -> dict:
     """Render a minimal chat request and return the GenerateRequest dict."""
-
     resp = await client.post(
         "/v1/chat/completions/render",
         json={
@@ -760,3 +1583,336 @@ async def test_streaming_usage_chunk(client):
     assert d2["chunk"]["choices"] == []
     assert d2["chunk"]["usage"]["prompt_tokens"] == 10
     assert d2["chunk"]["usage"]["completion_tokens"] == len(token_ids)
+
+
+# ---------------------------------------------------------------------------
+# Integration tests for parser configured (reasoning + tool calls) require a
+# live render server. Mirrors the parser_server / parser_tokenizer pattern
+# from test_derender.py.
+# ---------------------------------------------------------------------------
+
+PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+
+_PARSER_TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get weather for a city",
+            "parameters": {
+                "type": "object",
+                "properties": {"city": {"type": "string"}},
+            },
+        },
+    }
+]
+
+
+@pytest.fixture(scope="module")
+def parser_server():
+    from tests.utils import RemoteLaunchRenderServer
+
+    args = [
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "hermes",
+        "--reasoning-parser",
+        "deepseek_r1",
+    ]
+    with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def parser_client(parser_server):
+    import httpx
+
+    async with httpx.AsyncClient(
+        base_url=parser_server.url_for(""), timeout=60.0
+    ) as http_client:
+        yield http_client
+
+
+@pytest.fixture(scope="module")
+def parser_tokenizer():
+    from vllm.tokenizers import get_tokenizer
+
+    return get_tokenizer(PARSER_MODEL)
+
+
+def _require_parser_markers(tokenizer, text: str, *markers: str) -> list[int]:
+    """Encode text and skip the test if any marker is lost in roundtrip."""
+    ids = tokenizer.encode(text, add_special_tokens=False)
+    decoded = tokenizer.decode(ids, skip_special_tokens=False)
+    for m in markers:
+        if m not in decoded:
+            pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip")
+    return ids
+
+
+async def _render_parser_chat(client, messages: list[dict]) -> dict:
+    resp = await client.post(
+        "/v1/chat/completions/render",
+        json={"model": PARSER_MODEL, "messages": messages},
+    )
+    assert resp.status_code == 200, resp.text
+    return resp.json()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "chunking",
+    [
+        "whole",
+        "per_token",
+        pytest.param(
+            "triples",
+            marks=pytest.mark.xfail(
+                reason="<think> leaks into reasoning when it shares a "
+                "multi-token delta with reasoning text (vllm#55195)"
+            ),
+        ),
+    ],
+)
+async def test_stream_parsed_matches_batch_reasoning(
+    parser_client, parser_tokenizer, chunking
+):
+    """Streamed == batch: assembled reasoning/content equal the non
+    streaming `/derender` result over the same token IDs, for every
+    chunking of the same output (chunking invariance)."""
+    messages = [{"role": "user", "content": "What is 2+3?"}]
+    gen_req = await _render_parser_chat(parser_client, messages)
+
+    reasoning_text = "The user wants 2 plus 3. That is 5."
+    answer_text = "The answer is 5."
+    output_text = f"<think>{reasoning_text}</think>{answer_text}"
+    output_ids = _require_parser_markers(parser_tokenizer, output_text, "</think>")
+
+    chat_request = {
+        "model": PARSER_MODEL,
+        "messages": messages,
+        "include_reasoning": True,
+    }
+
+    batch_resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": {
+                "request_id": "batch",
+                "choices": [
+                    {"index": 0, "token_ids": output_ids, "finish_reason": "stop"}
+                ],
+            },
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": chat_request,
+        },
+    )
+    assert batch_resp.status_code == 200, batch_resp.text
+    batch_msg = batch_resp.json()["choices"][0]["message"]
+
+    n = len(output_ids)
+    chunk_sizes = {
+        "whole": [n],
+        "per_token": [1] * n,
+        "triples": [3] * (n // 3) + [n % 3],
+    }[chunking]
+    chunk_sizes = [c for c in chunk_sizes if c > 0]
+    streamed = await stream_chat_derender(
+        parser_client,
+        output_ids,
+        chunk_sizes,
+        chat_request,
+        len(gen_req["token_ids"]),
+        gen_req["token_ids"],
+    )
+    assert streamed["content"] == batch_msg["content"]
+    assert streamed["reasoning"] == batch_msg.get("reasoning")
+
+
+@pytest.mark.asyncio
+async def test_stream_parsed_matches_batch_tool_call(parser_client, parser_tokenizer):
+    """Tool call name+id+arguments parity between streamed and batch, plus
+    the finish_reason -> "tool_calls" rewrite for auto tool choice."""
+    messages = [{"role": "user", "content": "Weather in Paris?"}]
+    gen_req = await _render_parser_chat(parser_client, messages)
+
+    output_text = (
+        "<think>Let me check.</think>"
+        '<tool_call>\n{"name": "get_weather", '
+        '"arguments": {"city": "Paris"}}\n</tool_call>'
+    )
+    output_ids = _require_parser_markers(
+        parser_tokenizer, output_text, "</think>", "<tool_call>", "</tool_call>"
+    )
+    chat_request = {
+        "model": PARSER_MODEL,
+        "messages": messages,
+        "tools": _PARSER_TOOLS,
+        "tool_choice": "auto",
+    }
+
+    batch_resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": {
+                "request_id": "batch",
+                "choices": [
+                    {"index": 0, "token_ids": output_ids, "finish_reason": "stop"}
+                ],
+            },
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": chat_request,
+        },
+    )
+    assert batch_resp.status_code == 200, batch_resp.text
+    batch_choice = batch_resp.json()["choices"][0]
+    if not batch_choice["message"]["tool_calls"]:
+        pytest.skip("Model did not emit a tool call")
+
+    for chunk_sizes in ([len(output_ids)], [1] * len(output_ids)):
+        # Snapshot of tool_calls[0]["arguments"] after every chunk, to check
+        # it only ever grows by appending. It never retracts or duplicates
+        # already streamed text. That's the exact failure mode the
+        # replay + parse_delta design (over diff-based streaming) exists to
+        # avoid. Parity with the batch result alone wouldn't catch a
+        # transient mid-stream regression.
+        arg_snapshots: list[str] = []
+
+        def _record(
+            tool_calls: list[dict], _snapshots: list[str] = arg_snapshots
+        ) -> None:
+            if tool_calls and tool_calls[0]["arguments"]:
+                _snapshots.append(tool_calls[0]["arguments"])
+
+        streamed = await stream_chat_derender(
+            parser_client,
+            output_ids,
+            chunk_sizes,
+            chat_request,
+            len(gen_req["token_ids"]),
+            gen_req["token_ids"],
+            on_chunk=_record,
+        )
+        assert streamed["tool_calls"]
+        assert streamed["tool_calls"][0]["name"] == "get_weather"
+        assert streamed["tool_calls"][0]["id"] is not None
+        assert json.loads(streamed["tool_calls"][0]["arguments"]) == json.loads(
+            batch_choice["message"]["tool_calls"][0]["function"]["arguments"]
+        )
+        assert streamed["finish_reason"] == "tool_calls"
+
+        assert arg_snapshots, "expected at least one tool-call argument delta"
+        for prev, curr in zip(arg_snapshots, arg_snapshots[1:]):
+            assert curr.startswith(prev), (
+                f"tool-call arguments regressed: {prev!r} -> {curr!r}"
+            )
+        assert arg_snapshots[-1] == streamed["tool_calls"][0]["arguments"]
+
+
+@pytest.mark.asyncio
+async def test_stream_parsed_cjk_across_chunk_boundaries(
+    parser_client, parser_tokenizer
+):
+    """CJK/emoji tokens split across chunk boundaries with a parser active
+    across the reasoning -> content transition (parser path regression
+    guard analogous to #46159 in the plain detok path)."""
+    messages = [{"role": "user", "content": "Reply in Chinese"}]
+    gen_req = await _render_parser_chat(parser_client, messages)
+
+    reasoning_text = "思考中"
+    answer_text = "你好世界 😀"
+    output_text = f"<think>{reasoning_text}</think>{answer_text}"
+    output_ids = _require_parser_markers(parser_tokenizer, output_text, "</think>")
+
+    chat_request = {
+        "model": PARSER_MODEL,
+        "messages": messages,
+        "include_reasoning": True,
+    }
+    streamed = await stream_chat_derender(
+        parser_client,
+        output_ids,
+        [1] * len(output_ids),
+        chat_request,
+        len(gen_req["token_ids"]),
+        gen_req["token_ids"],
+    )
+    assert "�" not in (streamed["content"] or "")
+    assert "�" not in (streamed["reasoning"] or "")
+    assert answer_text in (streamed["content"] or "")
+
+
+@pytest.mark.asyncio
+async def test_stream_parsed_include_reasoning_false(parser_client, parser_tokenizer):
+    """include_reasoning=False emits no reasoning deltas on the streaming
+    parser path."""
+    messages = [{"role": "user", "content": "What is 2+3?"}]
+    gen_req = await _render_parser_chat(parser_client, messages)
+
+    output_text = "<think>reasoning here</think>The answer is 5."
+    output_ids = _require_parser_markers(parser_tokenizer, output_text, "</think>")
+
+    chat_request = {
+        "model": PARSER_MODEL,
+        "messages": messages,
+        "include_reasoning": False,
+    }
+    streamed = await stream_chat_derender(
+        parser_client,
+        output_ids,
+        [1] * len(output_ids),
+        chat_request,
+        len(gen_req["token_ids"]),
+        gen_req["token_ids"],
+    )
+    assert streamed["reasoning"] is None
+
+
+@pytest.mark.asyncio
+async def test_stream_parsed_missing_chat_request_rejected(parser_client):
+    """Parser configured + no chat_request on the streaming endpoint -> 400
+    (the live-server counterpart to the mocked
+    `TestServingDerenderStreamValidation` checks above)."""
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "stream": True,
+            "model": PARSER_MODEL,
+            "generate_chunk": {
+                "request_id": "reject-test",
+                "choices": [{"index": 0, "token_ids": [1, 2], "finish_reason": None}],
+            },
+        },
+    )
+    assert resp.status_code == 400
+    assert "chat_request" in resp.json()["error"]["message"]
+
+
+@pytest.mark.asyncio
+async def test_stream_parsed_missing_prompt_token_ids_rejected(parser_client):
+    """Parser configured + chat_request given but no prompt_token_ids on the
+    streaming endpoint -> 400 (the live-server counterpart to the mocked
+    `test_missing_prompt_token_ids_with_parser_rejected` above). Without
+    prompt_token_ids, parse_delta cannot tell whether the prompt left
+    reasoning open and would silently misclassify reasoning as content."""
+    messages = [{"role": "user", "content": "Hello"}]
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "stream": True,
+            "model": PARSER_MODEL,
+            "generate_chunk": {
+                "request_id": "reject-test",
+                "choices": [{"index": 0, "token_ids": [1, 2], "finish_reason": None}],
+            },
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 400
+    assert "prompt_token_ids" in resp.json()["error"]["message"]
diff --git a/tests/entrypoints/scale_out/derender/utils.py b/tests/entrypoints/scale_out/derender/utils.py
new file mode 100644
index 000000000000..d8128353bf3c
--- /dev/null
+++ b/tests/entrypoints/scale_out/derender/utils.py
@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Helpers shared by the streaming chat derender tests."""
+
+from collections.abc import Callable
+
+import httpx
+
+
+def assemble_stream(choices: list[dict]) -> dict:
+    """Fold streamed chat choices into one message.
+
+    A tool call index must keep the first ID it was given. On the derender
+    side that is what pinning the IDs in `stream_state` guarantees.
+    """
+    content = ""
+    reasoning = ""
+    tool_calls: list[dict] = []
+    finish_reason = None
+    for choice in choices:
+        delta = choice["delta"]
+        content += delta.get("content") or ""
+        reasoning += delta.get("reasoning") or ""
+        for tc in delta.get("tool_calls") or []:
+            idx = tc["index"]
+            while len(tool_calls) <= idx:
+                tool_calls.append({"id": None, "name": None, "arguments": ""})
+            if tc.get("id"):
+                assert tool_calls[idx]["id"] in (None, tc["id"]), (
+                    f"tool call {idx} changed ID mid stream"
+                )
+                tool_calls[idx]["id"] = tc["id"]
+            fn = tc.get("function") or {}
+            if fn.get("name"):
+                tool_calls[idx]["name"] = fn["name"]
+            tool_calls[idx]["arguments"] += fn.get("arguments") or ""
+        finish_reason = choice.get("finish_reason") or finish_reason
+    return {
+        "content": content or None,
+        "reasoning": reasoning or None,
+        "tool_calls": tool_calls,
+        "finish_reason": finish_reason,
+    }
+
+
+async def stream_chat_derender(
+    client: httpx.AsyncClient,
+    output_ids: list[int],
+    chunk_sizes: list[int],
+    chat_request: dict,
+    prompt_tokens: int,
+    prompt_token_ids: list[int],
+    finish_reason: str = "stop",
+    on_chunk: Callable[[list[dict]], None] | None = None,
+) -> dict:
+    """Feed `output_ids` through the streaming chat derender endpoint in
+    the given `chunk_sizes`, threading `stream_state` across calls and
+    return the assembled message.
+
+    `finish_reason` goes on the last chunk, so a trailing size of 0 sends it
+    on a chunk with no tokens.
+
+    If `on_chunk` is given, it is called after every chunk with the
+    `tool_calls` assembled so far, letting callers assert properties of the
+    intermediate deltas (e.g. monotonic argument growth) rather than only
+    the final assembled result.
+    """
+    state = None
+    choices: list[dict] = []
+    pos = 0
+    for i, size in enumerate(chunk_sizes):
+        tids = output_ids[pos : pos + size]
+        pos += size
+        is_last = i == len(chunk_sizes) - 1
+        resp = await client.post(
+            "/v1/chat/completions/derender",
+            json={
+                "stream": True,
+                "model": chat_request["model"],
+                "generate_chunk": {
+                    "request_id": "stream-test",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "token_ids": tids,
+                            "finish_reason": finish_reason if is_last else None,
+                        }
+                    ],
+                },
+                "stream_state": state,
+                "prompt_tokens": prompt_tokens,
+                "prompt_token_ids": prompt_token_ids,
+                "chat_request": chat_request,
+            },
+        )
+        assert resp.status_code == 200, resp.text
+        data = resp.json()
+        state = data["stream_state"]
+        choices.extend(data["chunk"]["choices"])
+        if on_chunk is not None:
+            on_chunk(assemble_stream(choices)["tool_calls"])
+
+    return assemble_stream(choices)
diff --git a/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh b/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
index b99b1c3ca5b0..462ada22d712 100755
--- a/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
+++ b/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
@@ -181,8 +181,9 @@ run_scale_out_ec() {
     # Encode-only EC producer instance. It runs the vision encoder and
     # publishes embeddings through the ECExampleConnector shared storage.
     echo "Starting encode instance on GPU $GPU_E, port $ENCODE_PORT"
-    env CUDA_VISIBLE_DEVICES="$GPU_E" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \
+    env CUDA_VISIBLE_DEVICES="$GPU_E" \
         vllm serve "$MODEL" \
+        --enable-scale-out \
         --port "$ENCODE_PORT" \
         --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
@@ -204,8 +205,9 @@ run_scale_out_ec() {
     # features plus ec_transfer_params and loads the embeddings that the
     # encode instance published.
     echo "Starting prefill instance on GPU $GPU_PD, port $PREFILL_PORT"
-    env CUDA_VISIBLE_DEVICES="$GPU_PD" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \
+    env CUDA_VISIBLE_DEVICES="$GPU_PD" \
         vllm serve "$MODEL" \
+        --enable-scale-out \
         --port "$PREFILL_PORT" \
         --max-model-len "$MAX_MODEL_LEN" \
         --enforce-eager \
diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py
index 6aaf119f16a0..e65a26a7d140 100644
--- a/tests/entrypoints/scale_out/render/test_render.py
+++ b/tests/entrypoints/scale_out/render/test_render.py
@@ -337,7 +337,6 @@ async def test_chat_completion_render_multi_turn(client):
 @pytest.mark.asyncio
 async def test_chat_completion_render_with_stream_true(client):
     """Render accepts stream params but still returns JSON (non-streamed)."""
-
     response = await client.post(
         "/v1/chat/completions/render",
         json={
@@ -595,176 +594,6 @@ async def test_completion_render_multiple_prompts_token_offsets(client):
             assert 0 <= start <= end <= len(prompt)
 
 
-@pytest.mark.asyncio
-async def test_chat_completion_render_assistant_tokens_mask_default(client):
-    """Without return_assistant_tokens_mask, assistant_tokens_mask should be null."""
-    response = await client.post(
-        "/v1/chat/completions/render",
-        json={
-            "model": MODEL_NAME,
-            "messages": [
-                {"role": "user", "content": "Hello"},
-                {"role": "assistant", "content": "Hi!"},
-                {"role": "user", "content": "How are you?"},
-            ],
-        },
-    )
-
-    assert response.status_code == 200
-    data = response.json()
-    assert data.get("assistant_tokens_mask") is None
-
-
-@pytest.mark.asyncio
-async def test_chat_completion_render_assistant_tokens_mask_false(client):
-    """Explicitly setting return_assistant_tokens_mask=false gives null."""
-    response = await client.post(
-        "/v1/chat/completions/render",
-        json={
-            "model": MODEL_NAME,
-            "messages": [
-                {"role": "user", "content": "Hello"},
-            ],
-            "return_assistant_tokens_mask": False,
-        },
-    )
-
-    assert response.status_code == 200
-    data = response.json()
-    assert data.get("assistant_tokens_mask") is None
-
-
-@pytest.mark.asyncio
-async def test_chat_render_assistant_tokens_mask_null_without_gen_tags(
-    client,
-):
-    """The tiny test model lacks ``{% generation %}`` tags, so the mask is null."""
-    response = await client.post(
-        "/v1/chat/completions/render",
-        json={
-            "model": MODEL_NAME,
-            "messages": [
-                {"role": "user", "content": "Hello"},
-                {"role": "assistant", "content": "Hi!"},
-            ],
-            "return_assistant_tokens_mask": True,
-        },
-    )
-
-    assert response.status_code == 200
-    assert response.json().get("assistant_tokens_mask") is None
-
-
-# A minimal chat template with {% generation %} tags so we can test that
-# the mask correctly marks assistant tokens.
-_TEMPLATE_WITH_GENERATION = (
-    "{% for m in messages %}"
-    "{% if m['role'] == 'user' %}User: {{ m['content'] }}\n"
-    "{% elif m['role'] == 'assistant' %}"
-    "{% generation %}Assistant: {{ m['content'] }}\n{% endgeneration %}"
-    "{% endif %}"
-    "{% endfor %}"
-)
-
-
-@pytest.mark.asyncio
-async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags(
-    client,
-):
-    """With a ``{% generation %}``-enabled template, the mask marks assistant
-    tokens and the masked tokens decode to the assistant content."""
-    response = await client.post(
-        "/v1/chat/completions/render",
-        json={
-            "model": MODEL_NAME,
-            "messages": [
-                {"role": "user", "content": "Hello"},
-                {"role": "assistant", "content": "Hi!"},
-                {"role": "user", "content": "Bye"},
-            ],
-            "chat_template": _TEMPLATE_WITH_GENERATION,
-            "return_assistant_tokens_mask": True,
-        },
-    )
-
-    assert response.status_code == 200
-    data = response.json()
-
-    mask = data["assistant_tokens_mask"]
-    token_ids = data["token_ids"]
-    assert mask is not None
-    assert isinstance(mask, list)
-    assert len(mask) == len(token_ids)
-    assert all(v in (0, 1) for v in mask)
-    assert sum(mask) > 0, "mask should mark at least one assistant token"
-
-    # Detokenize masked (assistant) and unmasked (non-assistant) tokens
-    # separately to verify the mask is correct, not just non-empty.
-    masked_ids = [t for t, m in zip(token_ids, mask, strict=True) if m]
-    unmasked_ids = [t for t, m in zip(token_ids, mask, strict=True) if not m]
-
-    detok = await client.post(
-        "/detokenize",
-        json={"model": MODEL_NAME, "tokens": masked_ids},
-    )
-    assert detok.status_code == 200
-    assert "Hi!" in detok.json()["prompt"]
-
-    detok_rest = await client.post(
-        "/detokenize",
-        json={"model": MODEL_NAME, "tokens": unmasked_ids},
-    )
-    assert detok_rest.status_code == 200
-    assert "Hi!" not in detok_rest.json()["prompt"]
-    assert "Bye" in detok_rest.json()["prompt"]
-
-
-@pytest.mark.asyncio
-async def test_chat_render_assistant_tokens_mask_follows_truncation(client):
-    """The assistant mask must be truncated with the prompt it describes.
-
-    `assistant_tokens_mask` is positional: entry i labels token i. Truncating
-    `token_ids` from the left without truncating the mask leaves the two
-    describing different positions, and the mask ends up marking whichever
-    tokens happen to sit at the old offsets.
-    """
-    messages = [
-        # Deliberately lopsided: a long leading user turn and a short trailing
-        # one, so keeping the head of the mask is distinguishable from keeping
-        # its tail.
-        {"role": "user", "content": "Hello hello hello hello hello hello"},
-        {"role": "assistant", "content": "Hi!"},
-        {"role": "user", "content": "Bye"},
-    ]
-    body = {
-        "model": MODEL_NAME,
-        "messages": messages,
-        "chat_template": _TEMPLATE_WITH_GENERATION,
-        "return_assistant_tokens_mask": True,
-    }
-
-    full = await client.post("/v1/chat/completions/render", json=body)
-    assert full.status_code == 200
-    full_token_ids = full.json()["token_ids"]
-    full_mask = full.json()["assistant_tokens_mask"]
-    assert sum(full_mask) > 0
-
-    keep = len(full_token_ids) - 4
-    # Precondition: with this prompt the head and tail slices of the mask
-    # really do differ, so the assertion below can tell them apart.
-    assert full_mask[-keep:] != full_mask[:keep]
-
-    truncated = await client.post(
-        "/v1/chat/completions/render",
-        json={**body, "truncate_prompt_tokens": keep, "truncation_side": "left"},
-    )
-    assert truncated.status_code == 200
-    data = truncated.json()
-
-    assert data["token_ids"] == full_token_ids[-keep:]
-    assert data["assistant_tokens_mask"] == full_mask[-keep:]
-
-
 @pytest.mark.asyncio
 async def test_messages_render_basic(client):
     """Test basic Anthropic Messages render endpoint."""
diff --git a/tests/entrypoints/scale_out/render/test_render_multimodal.py b/tests/entrypoints/scale_out/render/test_render_multimodal.py
index 70bbc95016ce..eb45ce68ea9b 100644
--- a/tests/entrypoints/scale_out/render/test_render_multimodal.py
+++ b/tests/entrypoints/scale_out/render/test_render_multimodal.py
@@ -16,7 +16,6 @@
 @pytest.fixture(scope="module")
 def vision_server():
     """Vision-capable server used for multimodal /render tests."""
-
     args = [
         "--enforce-eager",
         "--max-model-len",
@@ -27,15 +26,10 @@ def vision_server():
         "1",
         "--limit-mm-per-prompt.video",
         "0",
+        "--enable-scale-out",
     ]
 
-    env_overrides = {"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}
-
-    with RemoteOpenAIServer(
-        VISION_MODEL_NAME,
-        args,
-        env_dict=env_overrides,
-    ) as remote_server:
+    with RemoteOpenAIServer(VISION_MODEL_NAME, args) as remote_server:
         yield remote_server
 
 
@@ -53,7 +47,6 @@ async def test_chat_completion_render_with_base64_image_url(
     local_asset_server,
 ):
     """Render a multimodal chat request and verify tokens are returned."""
-
     image = local_asset_server.get_image_asset("RGBA_comp.png")
     data_url = encode_image_url(image, format="PNG")
 
@@ -117,7 +110,6 @@ async def test_tokenize_matches_render_for_multimodal_input(
     local_asset_server,
 ):
     """`/tokenize` should match `/v1/chat/completions/render` token output."""
-
     image = local_asset_server.get_image_asset("RGBA_comp.png")
     data_url = encode_image_url(image, format="PNG")
 
diff --git a/tests/entrypoints/scale_out/test_factories.py b/tests/entrypoints/scale_out/test_factories.py
index 9811bca8a22c..6d8335395ed8 100644
--- a/tests/entrypoints/scale_out/test_factories.py
+++ b/tests/entrypoints/scale_out/test_factories.py
@@ -3,7 +3,6 @@
 
 from argparse import Namespace
 
-import pytest
 from fastapi import FastAPI
 
 from vllm.entrypoints.launchers.api_server.routers import register_api_routers
@@ -22,69 +21,49 @@
 
 
 def registered_paths(
-    supported_tasks: tuple[SupportedTask, ...], *, tokens_only: bool = False
+    supported_tasks: tuple[SupportedTask, ...],
+    *,
+    tokens_only: bool = False,
+    enable_scale_out: bool = False,
 ) -> set[str]:
     app = FastAPI()
-    args = Namespace(tokens_only=tokens_only, enable_fault_tolerance=False)
+    args = Namespace(
+        tokens_only=tokens_only,
+        enable_scale_out=enable_scale_out,
+        enable_fault_tolerance=False,
+    )
     app.state.args = args
     register_api_routers(args, app, supported_tasks)
     return {route.path for route in app.routes}
 
 
-def test_scale_out_routes_are_disabled_by_default(monkeypatch):
-    monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False)
-
+def test_scale_out_routes_are_disabled_by_default():
     paths = registered_paths(("generate",))
 
     assert SCALE_OUT_PATHS.isdisjoint(paths)
 
 
-def test_disabled_scale_out_routes_are_logged(monkeypatch, caplog):
-    monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False)
-
+def test_disabled_scale_out_routes_are_logged(caplog):
     with caplog.at_level("INFO", logger="vllm.entrypoints.scale_out.factories"):
         registered_paths(("generate",))
 
-    assert any(
-        "VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1" in record.message
-        for record in caplog.records
-    )
+    assert any("--enable-scale-out" in record.message for record in caplog.records)
 
 
-def test_scale_out_routes_can_be_enabled_for_generate_server(monkeypatch):
-    monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1")
-
-    paths = registered_paths(("generate",))
+def test_scale_out_routes_can_be_enabled_for_generate_server():
+    paths = registered_paths(("generate",), enable_scale_out=True)
 
     assert paths >= SCALE_OUT_PATHS
 
 
-def test_render_routes_remain_enabled_for_render_server(monkeypatch):
-    monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False)
-
+def test_render_routes_remain_enabled_for_render_server():
     paths = registered_paths(("render",))
 
     assert paths >= RENDER_PATHS
     assert GENERATE_PATH not in paths
 
 
-def test_render_server_rejects_explicitly_disabled_scale_out_routes(monkeypatch):
-    monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0")
-
-    with pytest.raises(ValueError, match="VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0"):
-        registered_paths(("render",))
-
-
-def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset(monkeypatch):
-    monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False)
-
+def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset():
     paths = registered_paths(("generate",), tokens_only=True)
 
     assert paths >= {GENERATE_PATH, "/abort_requests"}
-
-
-def test_tokens_only_mode_rejects_explicitly_disabled_scale_out_routes(monkeypatch):
-    monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0")
-
-    with pytest.raises(ValueError, match="--tokens-only"):
-        registered_paths(("generate",), tokens_only=True)
diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py
index ad39958702a8..f7c6f05e7fcd 100644
--- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py
+++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py
@@ -12,6 +12,9 @@
 from vllm.entrypoints.generate.base.protocol import StreamOptions
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+from vllm.entrypoints.scale_out.token_in_token_out.mm_features import (
+    placeholder_ranges_from_engine_input,
+)
 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
     GenerateRequest,
     GenerateResponse,
@@ -19,11 +22,13 @@
 from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
 from vllm.exceptions import GenerationError
 from vllm.logprobs import Logprob
+from vllm.multimodal.inputs import PlaceholderRange
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.renderers import renderer_from_config
 from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.metrics.stats import RequestSpecDecodeMetrics
 
 MODEL_NAME = "openai-community/gpt2"
 BASE_MODEL_PATHS = [
@@ -56,9 +61,10 @@ class MockModelConfig:
     encoder_config = None
     generation_config: str = "auto"
     media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
+    skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
@@ -126,6 +132,7 @@ def _make_request_output(
     logprobs: list[dict[int, Any] | None] | None = None,
     num_cached_tokens: int | None = None,
     index: int = 0,
+    spec_decode_metrics: RequestSpecDecodeMetrics | None = None,
 ) -> RequestOutput:
     return RequestOutput(
         request_id=request_id,
@@ -140,6 +147,7 @@ def _make_request_output(
                 cumulative_logprob=None,
                 logprobs=logprobs,
                 finish_reason=finish_reason,
+                spec_decode_metrics=spec_decode_metrics,
             )
         ],
         finished=finished,
@@ -176,6 +184,29 @@ def _parse_sse_chunks(chunks: list[str]) -> list[Any]:
     return parsed
 
 
+def test_placeholder_ranges_from_engine_input():
+    engine_input: Any = {
+        "type": "multimodal",
+        "mm_placeholders": {
+            "image": [
+                PlaceholderRange(offset=1, length=5),
+                PlaceholderRange(offset=8, length=2),
+            ]
+        },
+    }
+
+    placeholders = placeholder_ranges_from_engine_input(engine_input)
+
+    assert placeholders is not None
+    assert [item.model_dump() for item in placeholders["image"]] == [
+        {"offset": 1, "length": 5},
+        {"offset": 8, "length": 2},
+    ]
+
+    text_only: Any = {"type": "token", "prompt_token_ids": [1, 2, 3]}
+    assert placeholder_ranges_from_engine_input(text_only) is None
+
+
 @pytest.mark.asyncio
 async def test_serve_tokens_skips_mm_cache_for_remote_engine_execution():
     engine = _mock_engine()
@@ -294,6 +325,475 @@ async def mock_generate(*args, **kwargs):
     assert data_chunks[2]["choices"][0]["finish_reason"] == "stop"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+async def test_serve_tokens_returns_spec_decode_metrics(stream: bool):
+    engine = _mock_engine()
+    metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
+    metrics.observe(num_draft_tokens=3, num_accepted=2)
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1",
+            token_ids=[10],
+            finish_reason="stop",
+            finished=True,
+            spec_decode_metrics=metrics,
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        stream=stream,
+        stream_options=StreamOptions(include_usage=True) if stream else None,
+    )
+
+    response = await serving.serve_tokens(request)
+    if stream:
+        chunks = [chunk async for chunk in response]
+        payload = next(
+            chunk
+            for chunk in _parse_sse_chunks(chunks)
+            if isinstance(chunk, dict) and chunk.get("metrics")
+        )["metrics"]["speculative_decoding"]
+    else:
+        assert isinstance(response, GenerateResponse)
+        assert response.metrics is not None
+        assert response.metrics.speculative_decoding is not None
+        payload = response.metrics.speculative_decoding.model_dump()
+
+    assert payload["num_draft_tokens"] == 3
+    assert payload["num_accepted_draft_tokens"] == 2
+    assert payload["num_spec_steps"] == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("include_usage", [False, True])
+async def test_stream_returns_spec_decode_metrics_on_empty_terminal_output(
+    include_usage,
+):
+    engine = _mock_engine()
+    metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
+    metrics.observe(num_draft_tokens=3, num_accepted=2)
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10])
+        yield _make_request_output(
+            "req-1",
+            token_ids=[],
+            finish_reason="stop",
+            finished=True,
+            spec_decode_metrics=metrics,
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        stream=True,
+        stream_options=StreamOptions(include_usage=include_usage),
+    )
+
+    chunks = [chunk async for chunk in await serving.serve_tokens(request)]
+    data_chunks = [chunk for chunk in _parse_sse_chunks(chunks) if chunk != "[DONE]"]
+
+    assert "metrics" not in data_chunks[0]
+    assert len(data_chunks) == 1 + int(include_usage)
+    if include_usage:
+        assert data_chunks[1]["choices"] == []
+        assert (
+            data_chunks[1]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+async def test_serve_tokens_omits_spec_decode_metrics_for_parallel_sampling(
+    stream: bool,
+):
+    engine = _mock_engine()
+    metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
+    metrics.observe(num_draft_tokens=3, num_accepted=2)
+    result = _make_request_output(
+        "req-1",
+        token_ids=[10],
+        finish_reason="stop",
+        finished=True,
+        spec_decode_metrics=metrics,
+    )
+    result.outputs.append(
+        CompletionOutput(
+            index=1,
+            text="",
+            token_ids=[20],
+            cumulative_logprob=None,
+            logprobs=None,
+            finish_reason="stop",
+            spec_decode_metrics=metrics,
+        )
+    )
+
+    async def mock_generate(*args, **kwargs):
+        yield result
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1, n=2),
+        model=MODEL_NAME,
+        stream=stream,
+        stream_options=StreamOptions(include_usage=True) if stream else None,
+    )
+
+    response = await serving.serve_tokens(request)
+    if stream:
+        chunks = [chunk async for chunk in response]
+        assert all(
+            "metrics" not in chunk
+            for chunk in _parse_sse_chunks(chunks)
+            if isinstance(chunk, dict)
+        )
+    else:
+        assert isinstance(response, GenerateResponse)
+        assert response.metrics is None
+
+
+@pytest.mark.asyncio
+async def test_non_stream_returns_final_prompt_token_ids_when_requested():
+    engine = _mock_engine()
+    final_prompt_token_ids = [1, 2, 2, 2, 3]
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1",
+            token_ids=[10],
+            finish_reason="stop",
+            finished=True,
+            prompt_token_ids=final_prompt_token_ids,
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        return_token_ids=True,
+    )
+
+    response = await serving.serve_tokens(request)
+
+    assert isinstance(response, GenerateResponse)
+    assert response.prompt_token_ids == final_prompt_token_ids
+
+
+@pytest.mark.asyncio
+async def test_stream_returns_prompt_token_ids_only_once_when_requested():
+    engine = _mock_engine()
+    final_prompt_token_ids = [1, 2, 2, 2, 3]
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[10], prompt_token_ids=final_prompt_token_ids
+        )
+        yield _make_request_output(
+            "req-1",
+            token_ids=[20],
+            finish_reason="stop",
+            finished=True,
+            prompt_token_ids=final_prompt_token_ids,
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=2),
+        model=MODEL_NAME,
+        stream=True,
+        return_token_ids=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    parsed = _parse_sse_chunks([chunk async for chunk in response])
+    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+
+    assert data_chunks[0]["prompt_token_ids"] == final_prompt_token_ids
+    assert "prompt_token_ids" not in data_chunks[1]
+
+
+MM_ENGINE_INPUT: Any = {
+    "type": "multimodal",
+    "prompt_token_ids": [1, 2, 2, 2, 3],
+    "mm_placeholders": {"image": [PlaceholderRange(offset=1, length=3)]},
+}
+MM_PLACEHOLDERS = {"image": [{"offset": 1, "length": 3}]}
+
+
+def _build_mm_serving_tokens(engine: AsyncLLM) -> ServingTokens:
+    serving = _build_serving_tokens(engine)
+
+    async def _fake_preprocess(*args, **kwargs):
+        return [MM_ENGINE_INPUT]
+
+    serving.online_renderer.preprocess_completion = AsyncMock(
+        side_effect=_fake_preprocess
+    )
+    return serving
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("return_token_ids", [False, True])
+async def test_non_stream_prompt_metadata_requires_return_token_ids(
+    return_token_ids: bool,
+):
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[10], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_mm_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        return_token_ids=return_token_ids,
+    )
+
+    response = await serving.serve_tokens(request)
+
+    assert isinstance(response, GenerateResponse)
+    if return_token_ids:
+        assert response.prompt_token_ids == [1, 2, 3]
+        assert response.model_dump()["mm_placeholders"] == MM_PLACEHOLDERS
+    else:
+        assert response.prompt_token_ids is None
+        assert response.mm_placeholders is None
+        assert request._response_mm_placeholders is None
+
+
+@pytest.mark.asyncio
+async def test_non_stream_text_only_prompt_has_no_mm_placeholders():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[10], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        return_token_ids=True,
+    )
+
+    response = await serving.serve_tokens(request)
+
+    assert isinstance(response, GenerateResponse)
+    assert response.prompt_token_ids == [1, 2, 3]
+    assert response.mm_placeholders is None
+
+
+@pytest.mark.asyncio
+async def test_non_stream_encoder_decoder_omits_mm_placeholders():
+    engine = _mock_engine()
+    engine.model_config.is_encoder_decoder = True
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[10], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_mm_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        return_token_ids=True,
+    )
+
+    response = await serving.serve_tokens(request)
+
+    assert isinstance(response, GenerateResponse)
+    assert response.prompt_token_ids == [1, 2, 3]
+    assert response.mm_placeholders is None
+
+
+@pytest.mark.asyncio
+async def test_stream_prompt_metadata_lands_on_first_emitted_chunk():
+    """An empty first output is skipped; metadata goes on the first real chunk,
+    once, and never on the sibling chunk of the same iteration (n=2)."""
+    engine = _mock_engine()
+
+    def _two_outputs(tokens0: list[int], tokens1: list[int], **kwargs):
+        res = _make_request_output("req-1", token_ids=tokens0, **kwargs)
+        res.outputs.append(
+            CompletionOutput(
+                index=1,
+                text="",
+                token_ids=tokens1,
+                cumulative_logprob=None,
+                logprobs=None,
+                finish_reason=kwargs.get("finish_reason"),
+            )
+        )
+        return res
+
+    async def mock_generate(*args, **kwargs):
+        yield _two_outputs([], [])
+        yield _two_outputs([10], [11])
+        yield _two_outputs([20], [21], finish_reason="stop", finished=True)
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_mm_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=2, n=2),
+        model=MODEL_NAME,
+        stream=True,
+        return_token_ids=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    parsed = _parse_sse_chunks([chunk async for chunk in response])
+    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+
+    assert [c["choices"][0]["token_ids"] for c in data_chunks] == [
+        [10],
+        [11],
+        [20],
+        [21],
+    ]
+    assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3]
+    assert data_chunks[0]["mm_placeholders"] == MM_PLACEHOLDERS
+    for chunk in data_chunks[1:]:
+        assert "prompt_token_ids" not in chunk
+        assert "mm_placeholders" not in chunk
+
+
+@pytest.mark.asyncio
+async def test_stream_omits_prompt_metadata_by_default():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[10], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_mm_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        stream=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    parsed = _parse_sse_chunks([chunk async for chunk in response])
+    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+
+    assert data_chunks
+    for chunk in data_chunks:
+        assert "prompt_token_ids" not in chunk
+        assert "mm_placeholders" not in chunk
+
+
+@pytest.mark.asyncio
+async def test_stream_text_only_prompt_omits_mm_placeholders_key():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[10], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        stream=True,
+        return_token_ids=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    parsed = _parse_sse_chunks([chunk async for chunk in response])
+    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+
+    assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3]
+    assert "mm_placeholders" not in data_chunks[0]
+
+
+@pytest.mark.asyncio
+async def test_stream_zero_token_completion_still_delivers_prompt_metadata():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_mm_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        stream=True,
+        return_token_ids=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    parsed = _parse_sse_chunks([chunk async for chunk in response])
+    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+
+    assert len(data_chunks) == 1
+    assert data_chunks[0]["choices"][0]["token_ids"] == []
+    assert data_chunks[0]["choices"][0]["finish_reason"] == "stop"
+    assert data_chunks[0]["prompt_token_ids"] == [1, 2, 3]
+    assert data_chunks[0]["mm_placeholders"] == MM_PLACEHOLDERS
+
+
+@pytest.mark.asyncio
+async def test_stream_zero_token_completion_emits_no_chunk_by_default():
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1", token_ids=[], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        model=MODEL_NAME,
+        stream=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    parsed = _parse_sse_chunks([chunk async for chunk in response])
+
+    assert not [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+    assert parsed[-1] == "[DONE]"
+
+
 @pytest.mark.asyncio
 async def test_stream_error_mid_generation():
     """finish_reason='error' mid-stream yields error chunk then [DONE]."""
diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py
index 8bd1bb9091bf..5cd60a959fb3 100644
--- a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py
+++ b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Roundtrip tests for multimodal serde used by the
+"""Roundtrip tests for multimodal serde used by the
 token_in_token_out generate endpoint.
 """
 
+import pytest
 import torch
 from pydantic import ValidationError
 
@@ -241,6 +241,50 @@ def test_metadata_only_generate_requires_ec_transfer_params():
     assert set(item) == {"image_grid_thw"}
 
 
+@pytest.mark.parametrize("modality", ["image", "video"])
+@pytest.mark.parametrize(
+    ("has_metadata", "has_ec"),
+    [(True, False), (True, True), (False, False)],
+    ids=["metadata-without-ec", "metadata-with-ec", "cache-hit-without-ec"],
+)
+def test_mixed_items_require_ec_only_for_metadata_only(
+    modality: str, has_metadata: bool, has_ec: bool
+):
+    """A full item must not bypass the EC requirement for another item."""
+    engine_input, _, image_grid_thw = _image_engine_input()
+    rendered = extract_mm_features(engine_input)
+    assert rendered is not None
+    features = rendered.model_dump()
+    metadata = (
+        encode_mm_kwargs_item(
+            MultiModalKwargsItem({f"{modality}_grid_thw": image_grid_thw})
+        )
+        if has_metadata
+        else None
+    )
+    second_item = {
+        "mm_hashes": "second-item",
+        "mm_placeholders": {"offset": 2, "length": 1},
+        "kwargs_data": None,
+        "mm_metadata": metadata,
+    }
+    for key, value in second_item.items():
+        features[key].setdefault(modality, []).append(value)
+    payload = {
+        "token_ids": [1, 2, 3],
+        "sampling_params": {"max_tokens": 1},
+        "features": features,
+    }
+    if has_ec:
+        payload["ec_transfer_params"] = {"second-item": {"peer_host": "10.0.0.1"}}
+
+    if has_metadata and not has_ec:
+        with pytest.raises(ValidationError, match="ec_transfer_params"):
+            GenerateRequest.model_validate(payload)
+    else:
+        GenerateRequest.model_validate(payload)
+
+
 def test_kwargs_and_metadata_generate_does_not_require_ec():
     engine_input, pixel_values, image_grid_thw = _image_engine_input()
     rendered = extract_mm_features(engine_input)
diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py
index 963e416e7ea8..9359006562bb 100644
--- a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py
+++ b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py
@@ -15,9 +15,7 @@
 GEN_ENDPOINT = "/inference/v1/generate"
 
 # tiny-mixtral config: 8 local experts, top-2 routing, 2 hidden layers.
-# The published config has sliding_window=4096, which produces
-# SlidingWindowSpec kv-cache groups; RoutedExpertsManager requires a
-# FullAttentionSpec group, so we override sliding_window=null below.
+# Its published sliding window is incompatible with AuxOutput retention.
 NUM_LOCAL_EXPERTS = 8
 NUM_EXPERTS_PER_TOK = 2
 NUM_HIDDEN_LAYERS = 2
@@ -34,11 +32,14 @@ def server():
         "--enable-return-routed-experts",
         "--hf-overrides",
         '{"sliding_window": null}',
+        "--enable-scale-out",
     ]
     with RemoteOpenAIServer(
         MODEL_NAME,
         args,
-        env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"},
+        env_dict={
+            "VLLM_USE_V2_MODEL_RUNNER": "1",
+        },
     ) as remote_server:
         yield remote_server
 
@@ -59,9 +60,10 @@ async def client(server: RemoteOpenAIServer):
 @pytest.mark.asyncio
 async def test_generate_routed_experts(client):
     """Test that /inference/v1/generate returns routed_experts when enabled."""
+    prompt_token_ids = [1, 2, 3]
     payload = {
         "model": MODEL_NAME,
-        "token_ids": [1, 2, 3],
+        "token_ids": prompt_token_ids,
         "sampling_params": {"max_tokens": 10, "temperature": 0.0},
         "stream": False,
     }
@@ -83,3 +85,4 @@ async def test_generate_routed_experts(client):
     assert topk == NUM_EXPERTS_PER_TOK
     assert (routed_experts >= 0).all()
     assert (routed_experts < NUM_LOCAL_EXPERTS).all()
+    assert num_tokens == len(prompt_token_ids) + len(choice["token_ids"]) - 1
diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py
index 9491623fd115..1d017bd3cf1b 100644
--- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py
+++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py
@@ -37,13 +37,10 @@ def server():
         "4096",
         "--enforce-eager",
         "--no-enable-prefix-caching",
+        "--enable-scale-out",
     ]
 
-    with RemoteOpenAIServer(
-        MODEL_NAME,
-        args,
-        env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"},
-    ) as remote_server:
+    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
         yield remote_server
 
 
@@ -157,11 +154,28 @@ async def test_render_to_generate_roundtrip(client, test_image):
     )
 
 
-@pytest.mark.asyncio
-async def test_content_parts_generates_tokens(client, test_image):
-    """content_parts with raw media should produce output tokens."""
-    data_url = encode_image_url(test_image, format="PNG")
+def _collapse_placeholders(
+    token_ids: list[int], mm_placeholders: dict[str, list[dict[str, int]]]
+) -> list[int]:
+    """Undo placeholder expansion: each expanded range becomes one token.
+
+    ``/generate`` with ``content_parts`` expects unexpanded placeholders and
+    expands them itself, so the roundtrip must start from a collapsed prompt.
+    """
+    ranges = sorted(
+        (r["offset"], r["length"]) for rs in mm_placeholders.values() for r in rs
+    )
+    collapsed: list[int] = []
+    pos = 0
+    for offset, length in ranges:
+        collapsed.extend(token_ids[pos:offset])
+        collapsed.append(token_ids[offset])
+        pos = offset + length
+    collapsed.extend(token_ids[pos:])
+    return collapsed
 
+
+async def _render_image_prompt(client, data_url: str) -> dict:
     render_resp = await client.post(
         RENDER_ENDPOINT,
         json={
@@ -178,14 +192,26 @@ async def test_content_parts_generates_tokens(client, test_image):
         },
     )
     render_resp.raise_for_status()
-    token_ids = render_resp.json()["token_ids"]
+    return render_resp.json()
+
+
+@pytest.mark.asyncio
+async def test_content_parts_generates_tokens(client, test_image):
+    """Unexpanded prompt + content_parts must expand exactly like /render."""
+    data_url = encode_image_url(test_image, format="PNG")
+    render_data = await _render_image_prompt(client, data_url)
+    expanded = render_data["token_ids"]
+    mm_placeholders = render_data["features"]["mm_placeholders"]
+    unexpanded = _collapse_placeholders(expanded, mm_placeholders)
+    assert len(unexpanded) < len(expanded)
 
     gen_resp = await client.post(
         GEN_ENDPOINT,
         json={
-            "token_ids": token_ids,
+            "token_ids": unexpanded,
             "content_parts": [{"type": "image_url", "url": data_url}],
             "sampling_params": {"max_tokens": 10, "temperature": 0.0},
+            "return_token_ids": True,
         },
     )
     gen_resp.raise_for_status()
@@ -195,6 +221,8 @@ async def test_content_parts_generates_tokens(client, test_image):
     choice = gen_data["choices"][0]
     assert "token_ids" in choice
     assert len(choice["token_ids"]) > 0
+    assert gen_data["prompt_token_ids"] == expanded
+    assert gen_data["mm_placeholders"] == mm_placeholders
 
 
 @pytest.mark.asyncio
@@ -203,33 +231,20 @@ async def test_content_parts_streaming(client, test_image):
     import json as json_mod
 
     data_url = encode_image_url(test_image, format="PNG")
-
-    render_resp = await client.post(
-        RENDER_ENDPOINT,
-        json={
-            "model": MODEL_NAME,
-            "messages": [
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "image_url", "image_url": {"url": data_url}},
-                        {"type": "text", "text": "Describe this."},
-                    ],
-                }
-            ],
-        },
-    )
-    render_resp.raise_for_status()
-    token_ids = render_resp.json()["token_ids"]
+    render_data = await _render_image_prompt(client, data_url)
+    expanded = render_data["token_ids"]
+    mm_placeholders = render_data["features"]["mm_placeholders"]
+    unexpanded = _collapse_placeholders(expanded, mm_placeholders)
 
     async with client.stream(
         "POST",
         GEN_ENDPOINT,
         json={
-            "token_ids": token_ids,
+            "token_ids": unexpanded,
             "content_parts": [{"type": "image_url", "url": data_url}],
             "sampling_params": {"max_tokens": 8, "temperature": 0.0},
             "stream": True,
+            "return_token_ids": True,
         },
     ) as resp:
         resp.raise_for_status()
@@ -239,3 +254,9 @@ async def test_content_parts_streaming(client, test_image):
                 chunks.append(json_mod.loads(line[6:]))
 
     assert len(chunks) > 0
+    prompt_metadata_chunks = [
+        chunk for chunk in chunks if chunk.get("prompt_token_ids") is not None
+    ]
+    assert len(prompt_metadata_chunks) == 1
+    assert prompt_metadata_chunks[0]["prompt_token_ids"] == expanded
+    assert prompt_metadata_chunks[0]["mm_placeholders"] == mm_placeholders
diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
index fb741df061e1..fe41e651bef0 100644
--- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
+++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
@@ -69,6 +69,7 @@ def server(request):
         # Option B: deterministic prefix caching
         # "--enable-prefix-caching",
         # "--deterministic-prefix-caching",
+        "--enable-scale-out",
     ]
 
     extra_args = getattr(request, "param", None)
@@ -79,11 +80,7 @@ def server(request):
             else [str(extra_args)]
         )
 
-    with RemoteOpenAIServer(
-        MODEL_NAME,
-        args,
-        env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"},
-    ) as remote_server:
+    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
         yield remote_server
 
 
@@ -115,6 +112,29 @@ async def test_generate_endpoint(client):
     assert data["choices"][0].get("sampling_mask") is None
 
 
+@pytest.mark.asyncio
+async def test_generate_rejects_min_tokens_above_filled_max_tokens(client):
+    """Explicit null max_tokens must not skip the min_tokens bound."""
+    payload = {
+        "model": MODEL_NAME,
+        "token_ids": [1, 2, 3],
+        "sampling_params": {"max_tokens": None, "min_tokens": 2147483648},
+        "stream": False,
+    }
+    resp = await client.post(GEN_ENDPOINT, json=payload)
+    assert resp.status_code == 400
+    assert "min_tokens" in resp.json()["error"]["message"]
+
+    followup = {
+        "model": MODEL_NAME,
+        "token_ids": [1, 2, 3],
+        "sampling_params": {"max_tokens": 4},
+        "stream": False,
+    }
+    resp = await client.post(GEN_ENDPOINT, json=followup)
+    resp.raise_for_status()
+
+
 @pytest.mark.asyncio
 @pytest.mark.skipif(
     envs.VLLM_USE_RUST_FRONTEND,
@@ -425,6 +445,31 @@ async def test_stop_string_workflow(client, tokenizer, messages):
     assert generate_res == completions_res
 
 
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    envs.VLLM_USE_RUST_FRONTEND,
+    reason="--tokens-only is not supported by the Rust frontend",
+)
+@pytest.mark.parametrize("server", [["--tokens-only"]], indirect=True)
+async def test_stop_strings_rejected_when_tokens_only(client):
+    """--tokens-only forces detokenize=False after request validation, so
+    stop strings must be rejected explicitly rather than silently ignored."""
+    sampling_params = {"max_tokens": 5, "stop": ["never"]}
+    payload = {
+        "model": MODEL_NAME,
+        "token_ids": [1, 2, 3],
+        "sampling_params": sampling_params,
+        "stream": False,
+    }
+    resp = await client.post(GEN_ENDPOINT, json=payload)
+    assert resp.status_code == 400
+    assert "stop strings" in resp.json()["error"]["message"]
+
+    sampling_params["stop"] = None
+    resp = await client.post(GEN_ENDPOINT, json=payload)
+    resp.raise_for_status()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "server",
diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py
index 80f08078da27..f42a406f41e6 100644
--- a/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py
+++ b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py
@@ -26,3 +26,13 @@ def test_logprobs_zero_emits_sampled_token():
         num_output_top_logprobs=0,
     )
     assert len(result.content[0].top_logprobs) == 1
+
+
+def test_logprobs_minus_one_emits_all_tokens():
+    result = ServingTokens._create_tokens_logprobs(
+        None,
+        token_ids=[7],
+        top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}],
+        num_output_top_logprobs=-1,
+    )
+    assert len(result.content[0].top_logprobs) == 2
diff --git a/tests/entrypoints/serve/dev/rlhf/conftest.py b/tests/entrypoints/serve/dev/rlhf/conftest.py
index 1d2c598b1960..c636b5157291 100644
--- a/tests/entrypoints/serve/dev/rlhf/conftest.py
+++ b/tests/entrypoints/serve/dev/rlhf/conftest.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shared fixtures and helpers for the RL lifecycle test suite.
+"""Shared fixtures and helpers for the RL lifecycle test suite.
 
 All test modules under this directory import from here to avoid duplication.
 
@@ -92,6 +91,7 @@ def server(
         port:            HTTP port to bind (caller is responsible for uniqueness).
         timeout:         Seconds to wait for /health before giving up.
         dummy_weights:   If True, use --load-format dummy (fast, no real weights).
+
     """
     env = {**os.environ, "VLLM_SERVER_DEV_MODE": "1"}
     base = _DUMMY_ARGS if dummy_weights else _BASE_ARGS
diff --git a/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py b/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
index eb9aa7663c93..7cf6f4490c62 100644
--- a/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
+++ b/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
@@ -25,7 +25,7 @@ def echo_args_kwargs(self, *args, **kwargs) -> dict[str, Any]:
         )
 
     def return_none(self, *args, **kwargs) -> None:
-        """Test method that does not return anything"""
+        """Test method that does not return anything."""
         return
 
 
@@ -48,7 +48,7 @@ def server():
 
 
 def test_get_model_name(server):
-    """Test basic response"""
+    """Test basic response."""
     response = requests.post(
         server.url_for("collective_rpc"), json={"method": "get_model_name"}
     )
@@ -59,7 +59,7 @@ def test_get_model_name(server):
 
 
 def test_return_none(server):
-    """Test return none"""
+    """Test return none."""
     response = requests.post(
         server.url_for("collective_rpc"), json={"method": "return_none"}
     )
@@ -69,7 +69,7 @@ def test_return_none(server):
 
 
 def test_echo_args_kwargs(server):
-    """Test args, kwargs, and dict response"""
+    """Test args, kwargs, and dict response."""
     args = ["arg1", "arg2"]
     kwargs = {"key1": "value1", "key2": "value2"}
     response = requests.post(
diff --git a/tests/entrypoints/serve/dev/test_sleep.py b/tests/entrypoints/serve/dev/test_sleep.py
index 260dcd00bae9..34e69c9ecef4 100644
--- a/tests/entrypoints/serve/dev/test_sleep.py
+++ b/tests/entrypoints/serve/dev/test_sleep.py
@@ -1,14 +1,37 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from unittest.mock import AsyncMock
+
+import pytest
 import requests
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
 from prometheus_client.parser import text_string_to_metric_families
 
 from tests.utils import RemoteOpenAIServer
+from vllm.entrypoints.serve.dev.sleep.api_router import attach_router
 
 MODEL_NAME = "meta-llama/Llama-3.2-1B"
 
 
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("fails", [False, True])
+def test_release_kv_cache_memory_route(fails):
+    app = FastAPI()
+    app.state.engine_client = AsyncMock()
+    release = app.state.engine_client.release_kv_cache_memory
+    if fails:
+        release.side_effect = RuntimeError("requires a completed pause first")
+    attach_router(app)
+
+    with TestClient(app, raise_server_exceptions=False) as client:
+        response = client.post("/release_kv_cache_memory")
+
+    assert response.status_code == (500 if fails else 200)
+    release.assert_awaited_once_with()
+
+
 def test_sleep_mode():
     # dtype, max-len etc set so that this can run in CI
     args = [
@@ -63,7 +86,7 @@ def test_sleep_mode():
         )
         assert response.status_code == 200
 
-        # is sleeping should be false after waking up any part of the engine
+        # Partial wake keeps the engine sleeping.
         response = requests.get(remote_server.url_for("is_sleeping"))
         assert response.status_code == 200
         assert response.json().get("is_sleeping") is True
@@ -87,8 +110,7 @@ def test_sleep_mode():
 
 
 def _get_sleep_metrics_from_api(response: requests.Response):
-    """Return (awake, weights_offloaded, discard_all)"""
-
+    """Return (awake, weights_offloaded, discard_all)."""
     awake, weights_offloaded, discard_all = None, None, None
 
     for family in text_string_to_metric_families(response.text):
diff --git a/tests/entrypoints/serve/instrumentator/test_basic.py b/tests/entrypoints/serve/instrumentator/test_basic.py
index 73a97c4fa264..0711cbee957d 100644
--- a/tests/entrypoints/serve/instrumentator/test_basic.py
+++ b/tests/entrypoints/serve/instrumentator/test_basic.py
@@ -21,7 +21,7 @@
 
 @pytest.fixture(scope="module")
 def server_args(request: pytest.FixtureRequest) -> list[str]:
-    """Provide extra arguments to the server via indirect parametrization
+    """Provide extra arguments to the server via indirect parametrization.
 
     Usage:
 
diff --git a/tests/entrypoints/serve/instrumentator/test_metrics.py b/tests/entrypoints/serve/instrumentator/test_metrics.py
index 297fe233b5fb..ac4576fbb95d 100644
--- a/tests/entrypoints/serve/instrumentator/test_metrics.py
+++ b/tests/entrypoints/serve/instrumentator/test_metrics.py
@@ -403,8 +403,7 @@ async def _poll_until(
 
 
 def _get_running_metrics_from_api(server: RemoteOpenAIServer):
-    """Return (running_count, waiting_count, kv_cache_usage)"""
-
+    """Return (running_count, waiting_count, kv_cache_usage)."""
     response = requests.get(server.url_for("metrics"))
     assert response.status_code == HTTPStatus.OK
 
diff --git a/tests/entrypoints/serve/lora/test_lora_adapters.py b/tests/entrypoints/serve/lora/test_lora_adapters.py
index a22f0b38991b..8a860dc6251a 100644
--- a/tests/entrypoints/serve/lora/test_lora_adapters.py
+++ b/tests/entrypoints/serve/lora/test_lora_adapters.py
@@ -26,7 +26,7 @@
     (
         "test_modules_to_save",
         {"modules_to_save": ["lm_head"]},
-        "only supports modules_to_save being None",
+        "Unsupported modules_to_save",
     ),
 ]
 
diff --git a/tests/entrypoints/serve/middleware/test_optional_middleware.py b/tests/entrypoints/serve/middleware/test_optional_middleware.py
index 865bd445c812..160f1e248dcc 100644
--- a/tests/entrypoints/serve/middleware/test_optional_middleware.py
+++ b/tests/entrypoints/serve/middleware/test_optional_middleware.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for middleware that's off by default and can be toggled through
+"""Tests for middleware that's off by default and can be toggled through
 server arguments, mainly --api-key and --enable-request-id-headers.
 """
 
diff --git a/tests/entrypoints/serve/tokenize/test_serving_tokenization.py b/tests/entrypoints/serve/tokenize/test_serving_tokenization.py
index a74962ab0ddd..057d7d1e8a74 100644
--- a/tests/entrypoints/serve/tokenize/test_serving_tokenization.py
+++ b/tests/entrypoints/serve/tokenize/test_serving_tokenization.py
@@ -56,9 +56,10 @@ class MockModelConfig:
     encoder_config = None
     generation_config: str = "auto"
     media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
-    skip_tokenizer_init = False
+    skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
diff --git a/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py b/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py
index 6b226c6999ef..c4317973cb23 100644
--- a/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py
+++ b/tests/entrypoints/serve/tokenize/test_tokenization_vlm.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Regression test: ``/tokenize`` must expand image placeholders for VLM models.
+"""Regression test: ``/tokenize`` must expand image placeholders for VLM models.
 
 Fixed by PR #34560 ("Move InputPreprocessor into Renderer (2/2)").
 Before that change, ``/tokenize`` returned ~26 tokens for a message with an
diff --git a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py
index 50083e345538..1d114bc8b81f 100644
--- a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py
+++ b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Regression test: calling ``/tokenize`` with multimodal data followed by
+"""Regression test: calling ``/tokenize`` with multimodal data followed by
 ``/v1/chat/completions`` with the same data must not cause an error.
 
 Ensures that the ``/tokenize`` endpoint does not pollute internal caches
@@ -53,7 +52,6 @@ async def test_tokenize_then_chat_completion_with_image(
 ):
     """Tokenize a multimodal message, then send the same message to chat
     completions.  The chat completion must succeed (not 500)."""
-
     image_url = local_asset_server.url_for("stop_sign.jpg")
     messages = [
         {
diff --git a/tests/entrypoints/openai/test_sse_keep_alive.py b/tests/entrypoints/serve/utils/test_sse_keep_alive.py
similarity index 99%
rename from tests/entrypoints/openai/test_sse_keep_alive.py
rename to tests/entrypoints/serve/utils/test_sse_keep_alive.py
index 74f5ee3ba234..519b9bb4a129 100644
--- a/tests/entrypoints/openai/test_sse_keep_alive.py
+++ b/tests/entrypoints/serve/utils/test_sse_keep_alive.py
@@ -6,7 +6,7 @@
 import pytest
 from starlette.responses import StreamingResponse
 
-from vllm.entrypoints.openai.sse_keep_alive import (
+from vllm.entrypoints.serve.utils.sse_keep_alive import (
     SSE_KEEP_ALIVE_COMMENT,
     with_sse_keep_alive,
 )
diff --git a/tests/entrypoints/speech_to_text/conftest.py b/tests/entrypoints/speech_to_text/conftest.py
index 94505bd054f4..bf65a72ab57d 100644
--- a/tests/entrypoints/speech_to_text/conftest.py
+++ b/tests/entrypoints/speech_to_text/conftest.py
@@ -11,6 +11,7 @@ def add_attention_backend(server_args, attention_config):
     Args:
         server_args: List of server arguments to extend in-place.
         attention_config: Dict with 'backend' key, or None.
+
     """
     if attention_config and "backend" in attention_config:
         server_args.extend(["--attention-backend", attention_config["backend"]])
diff --git a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py
index 713fa48eee48..855dafaaba09 100644
--- a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py
+++ b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Evaluate Transcription API correctness by computing Word Error Rate (WER)
+"""Evaluate Transcription API correctness by computing Word Error Rate (WER)
 on a given ASR dataset. When provided, it will also compare the WER against
 a baseline.
 This simulates real work usage of the API and makes sure that the frontend and
diff --git a/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py
index cd61507728a6..5534845326ab 100644
--- a/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py
+++ b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py
@@ -18,9 +18,8 @@
 
 @pytest.mark.asyncio
 async def test_chunk_offsets_are_cumulative_not_nominal():
-    """
-    chunk_start_offsets must be cumulative actual chunk lengths, not the old approach of
-    'idx * max_audio_clip_s'.  When split_audio places a boundary before the
+    """chunk_start_offsets must be cumulative actual chunk lengths, not the old approach
+    of 'idx * max_audio_clip_s'.  When split_audio places a boundary before the
     nominal 30 s mark, the old formula drifts; the fixed formula stays exact.
     """
     # Chunks shorter than exactly 30 s, as split_audio produces when a quiet
diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py
index 51599d359e8f..d5658e7cf0b4 100644
--- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py
+++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py
@@ -238,8 +238,7 @@ async def test_stream_options(winning_call, whisper_client):
 
 @pytest.mark.asyncio
 async def test_sampling_params(mary_had_lamb, whisper_client):
-    """
-    Compare sampling with params and greedy sampling to assert results
+    """Compare sampling with params and greedy sampling to assert results
     are different when extreme sampling parameters values are picked.
     """
     transcription = await whisper_client.audio.transcriptions.create(
diff --git a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py
index f03b6e8a3818..79b46bea51ab 100644
--- a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py
+++ b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py
@@ -150,7 +150,6 @@ async def test_non_streaming_tool_call(
     client: openai.AsyncOpenAI, server_config: ServerConfig
 ):
     """Test tool call in non-streaming mode."""
-
     response = await client.chat.completions.create(
         model=server_config["model_arg"],
         messages=MESSAGES,
@@ -183,7 +182,6 @@ async def test_streaming_tool_call(
     client: openai.AsyncOpenAI, server_config: ServerConfig
 ):
     """Test tool call in streaming mode."""
-
     stream = await client.chat.completions.create(
         model=server_config["model_arg"],
         messages=MESSAGES,
@@ -230,7 +228,6 @@ async def test_non_streaming_product_tool_call(
     client: openai.AsyncOpenAI, server_config: ServerConfig
 ):
     """Test tool call integer and boolean parameters in non-streaming mode."""
-
     response = await client.chat.completions.create(
         model=server_config["model_arg"],
         messages=PRODUCT_MESSAGES,
@@ -272,7 +269,6 @@ async def test_streaming_product_tool_call(
     client: openai.AsyncOpenAI, server_config: ServerConfig
 ):
     """Test tool call integer and boolean parameters in streaming mode."""
-
     stream = await client.chat.completions.create(
         model=server_config["model_arg"],
         messages=PRODUCT_MESSAGES,
diff --git a/tests/entrypoints/tool_parsers/test_openai_tool_parser.py b/tests/entrypoints/tool_parsers/test_openai_tool_parser.py
index d99b66d9ac6b..de402b435c23 100644
--- a/tests/entrypoints/tool_parsers/test_openai_tool_parser.py
+++ b/tests/entrypoints/tool_parsers/test_openai_tool_parser.py
@@ -127,8 +127,7 @@ async def client(server):
 # Utility to extract reasoning and tool calls
 # ==========================================================
 def extract_reasoning_and_calls(chunks: list) -> tuple[str, list[str], list[str]]:
-    """
-    Extract accumulated reasoning text and tool call arguments
+    """Extract accumulated reasoning text and tool call arguments
     from streaming chunks.
     """
     reasoning: str = ""
@@ -165,7 +164,6 @@ def extract_reasoning_and_calls(chunks: list) -> tuple[str, list[str], list[str]
 @pytest.mark.asyncio
 async def test_calculator_tool_call_and_argument_accuracy(client: openai.AsyncOpenAI):
     """Verify calculator tool call is made and arguments are accurate."""
-
     response = await client.chat.completions.create(
         model=MODEL_NAME,
         messages=MESSAGES_CALC,
@@ -205,7 +203,6 @@ async def test_calculator_tool_call_and_argument_accuracy(client: openai.AsyncOp
 @pytest.mark.asyncio
 async def test_streaming_tool_call_get_time_with_reasoning(client: openai.AsyncOpenAI):
     """Verify streamed reasoning and tool call behavior for get_time."""
-
     stream = await client.chat.completions.create(
         model=MODEL_NAME,
         messages=MESSAGES_GET_TIME,
@@ -258,8 +255,7 @@ async def test_streaming_multiple_tools(client: openai.AsyncOpenAI):
 
 @pytest.mark.asyncio
 async def test_invalid_tool_call(client: openai.AsyncOpenAI):
-    """
-    Verify that ambiguous instructions that should not trigger a tool
+    """Verify that ambiguous instructions that should not trigger a tool
     do not produce any tool calls.
     """
     response = await client.chat.completions.create(
@@ -283,8 +279,7 @@ async def test_invalid_tool_call(client: openai.AsyncOpenAI):
 
 @pytest.mark.asyncio
 async def test_tool_call_with_temperature(client: openai.AsyncOpenAI):
-    """
-    Verify model produces valid tool or text output
+    """Verify model produces valid tool or text output
     under non-deterministic sampling.
     """
     response = await client.chat.completions.create(
diff --git a/tests/entrypoints/unit_tests/test_chat_utils.py b/tests/entrypoints/unit_tests/test_chat_utils.py
index 575bb904d664..ab57140c2b07 100644
--- a/tests/entrypoints/unit_tests/test_chat_utils.py
+++ b/tests/entrypoints/unit_tests/test_chat_utils.py
@@ -940,7 +940,6 @@ def test_parse_chat_messages_audio_embeds_with_string(
     audio_embeds_model_config,
 ):
     """Test audio_embeds with base64 string embedding data."""
-
     import torch
 
     # Create a sample audio embedding tensor
@@ -983,7 +982,6 @@ async def test_parse_chat_messages_audio_embeds_async(
     audio_embeds_model_config,
 ):
     """Test audio_embeds with async futures."""
-
     import torch
 
     # Create a sample audio embedding tensor
@@ -2384,8 +2382,10 @@ def test_parse_chat_messages_include_thinking_chunk(mistral_model_config):
     assert conversation_with_thinking == expected_conversation
 
 
+@pytest.mark.parametrize("input_audio", [None, {}, {"data": "", "format": "wav"}])
 def test_parse_chat_messages_single_empty_audio_with_uuid(
     qwen2_audio_model_config,
+    input_audio,
 ):
     audio_uuid = "abcd"
     conversation, mm_data, mm_uuids = parse_chat_messages(
@@ -2395,7 +2395,7 @@ def test_parse_chat_messages_single_empty_audio_with_uuid(
                 "content": [
                     {
                         "type": "input_audio",
-                        "input_audio": {},
+                        "input_audio": input_audio,
                         "uuid": audio_uuid,
                     },
                     {"type": "text", "text": "What does the audio say?"},
@@ -2418,8 +2418,10 @@ def test_parse_chat_messages_single_empty_audio_with_uuid(
 
 
 @pytest.mark.asyncio
+@pytest.mark.parametrize("input_audio", [None, {}, {"data": "", "format": "wav"}])
 async def test_parse_chat_messages_single_empty_audio_with_uuid_async(
     qwen2_audio_model_config,
+    input_audio,
 ):
     audio_uuid = "abcd"
     conversation, mm_data, mm_uuids = await parse_chat_messages_async(
@@ -2429,7 +2431,7 @@ async def test_parse_chat_messages_single_empty_audio_with_uuid_async(
                 "content": [
                     {
                         "type": "input_audio",
-                        "input_audio": {},
+                        "input_audio": input_audio,
                         "uuid": audio_uuid,
                     },
                     {"type": "text", "text": "What does the audio say?"},
@@ -2982,7 +2984,6 @@ def test_postprocess_messages_null_arguments_string():
 @pytest.mark.asyncio
 async def test_resolve_items_runs_modalities_concurrently_and_preserves_order():
     """Media fetches overlap while modality and item order are preserved."""
-
     active_fetches = 0
     max_active_fetches = 0
 
diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py
index 1c1f6ed23593..a5b1afe8c50f 100644
--- a/tests/entrypoints/unit_tests/test_context.py
+++ b/tests/entrypoints/unit_tests/test_context.py
@@ -71,8 +71,8 @@ async def generate_mock_outputs(
 class FakeHarmonyParser(HarmonyParser):
     def __init__(self):
         # Skip HarmonyParser initialization and script outputs directly.
-        self.reasoning_parser = None
-        self.tool_parser = None
+        self._reasoning_parser = None
+        self._tool_parser = None
         self._chunk_results: list[ChunkResult] = []
         self._flush_results: list[list[Segment]] = []
         self.processed_chunks: list[list[int]] = []
@@ -544,7 +544,6 @@ async def test_streaming_multi_turn_token_counting():
 @pytest.mark.asyncio
 async def test_streaming_message_synchronization():
     """Completed messages from append-local and flush segments sync into context."""
-
     # Create a streaming context with some initial messages
     initial_messages = [
         Message(
@@ -683,6 +682,7 @@ def create_simple_context_output(
     prompt="Test prompt",
     prompt_token_ids=None,
     num_cached_tokens=0,
+    num_cache_creation_tokens=None,
     logprobs=None,
     finished=True,
 ):
@@ -708,6 +708,7 @@ def create_simple_context_output(
         ],
         finished=finished,
         num_cached_tokens=num_cached_tokens,
+        num_cache_creation_tokens=num_cache_creation_tokens,
     )
 
 
@@ -828,6 +829,7 @@ def test_simple_context_token_counting():
             token_ids=[10, 11],
             prompt_token_ids=[1, 2, 3, 4, 5],
             num_cached_tokens=2,
+            num_cache_creation_tokens=3,
         )
     )
     context.append_output(
@@ -836,12 +838,14 @@ def test_simple_context_token_counting():
             token_ids=[12],
             prompt_token_ids=[1, 2, 3, 4, 5],
             num_cached_tokens=2,
+            num_cache_creation_tokens=3,
         )
     )
 
     assert context.num_prompt_tokens == 5
     assert context.num_output_tokens == 3  # 2 + 1
     assert context.num_cached_tokens == 2
+    assert context.num_cache_creation_tokens == 3
 
 
 def test_simple_context_final_output():
diff --git a/tests/evals/gpt_oss/conftest.py b/tests/evals/gpt_oss/conftest.py
index d35dec4831a3..28fbb88d527e 100644
--- a/tests/evals/gpt_oss/conftest.py
+++ b/tests/evals/gpt_oss/conftest.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Pytest configuration for GPT-OSS evaluation tests.
-"""
+"""Pytest configuration for GPT-OSS evaluation tests."""
 
 from pathlib import Path
 
diff --git a/tests/evals/gpt_oss/test_gpqa_correctness.py b/tests/evals/gpt_oss/test_gpqa_correctness.py
index 63188ec40767..99d799698a37 100644
--- a/tests/evals/gpt_oss/test_gpqa_correctness.py
+++ b/tests/evals/gpt_oss/test_gpqa_correctness.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-GPQA evaluation using vLLM server and GPT-OSS evaluation package.
+"""GPQA evaluation using vLLM server and GPT-OSS evaluation package.
 
 Usage:
 pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \
@@ -48,7 +47,6 @@ def ensure_tiktoken_files():
 
 def run_gpqa_eval(model_name: str, base_url: str, reasoning_effort: str) -> float:
     """Run GPQA evaluation using the gpt-oss evaluation package."""
-
     # Build the command to run the evaluation
     cmd = [
         sys.executable,
diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml
new file mode 100644
index 000000000000..009871925925
--- /dev/null
+++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml
@@ -0,0 +1,20 @@
+model_name: "nvidia/GLM-5.2-NVFP4"
+accuracy_threshold: 0.90
+tolerance: 0.0
+num_questions: 1319
+num_fewshot: 5
+max_tokens: 256
+max_concurrency: 32
+request_timeout_seconds: 1800
+startup_max_wait_seconds: 1200
+server_args: >-
+  --tensor-parallel-size 4
+  --max-model-len 8192
+  --max-num-seqs 32
+  --max-num-batched-tokens 4096
+  --enforce-eager
+  --kv-cache-dtype bfloat16
+  --attention-config '{"hisparse_config":{}}'
+  --kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":32}}'
+env:
+  VLLM_USE_V2_MODEL_RUNNER: "1"
diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml
new file mode 100644
index 000000000000..8f3e9ce918ef
--- /dev/null
+++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml
@@ -0,0 +1,23 @@
+model_name: "nvidia/GLM-5.2-NVFP4"
+accuracy_threshold: 0.90
+num_questions: 1319
+num_fewshot: 5
+max_concurrency: 100
+# Preserve KV-cache headroom after CUDA graph memory profiling and bound the
+# DCP-aware FlashInfer workspace to 2.39 GiB at the 4,096-token scheduler cap.
+server_args: >-
+  --compilation-config '{"cudagraph_mode":"PIECEWISE"}'
+  --gpu-memory-utilization 0.96
+  --max-model-len 4096
+  --max-num-batched-tokens 4096
+  --safetensors-load-strategy prefetch
+  --moe-backend flashinfer_cutlass
+  --prefill-context-parallel-size 4
+  --decode-context-parallel-size 4
+  --dcp-comm-backend ag_rs
+  --enable-expert-parallel
+  --kv-cache-dtype fp8
+env:
+  PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True"
+  VLLM_LOGGING_LEVEL: "DEBUG"
+  VLLM_USE_V2_MODEL_RUNNER: "1"
diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
index 72f7fdf1ae91..141bf68162a3 100644
--- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
+++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
@@ -3,8 +3,12 @@ accuracy_threshold: 0.90
 num_questions: 1319
 num_fewshot: 5
 max_concurrency: 100
+# 30B+ MoE with EP/PCP: download + kernel autotuning can exceed the 1200s default.
+startup_max_wait_seconds: 2400
+# Leave headroom for FlashInfer CUTLASS autotuning after KV-cache allocation.
 server_args: >-
   --compilation-config '{"cudagraph_mode":"NONE"}'
+  --gpu-memory-utilization 0.85
   --max-model-len 4096
   --max-num-batched-tokens 32768
   --safetensors-load-strategy prefetch
diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml
new file mode 100644
index 000000000000..c1298bebfa62
--- /dev/null
+++ b/tests/evals/gsm8k/configs/Qwen3-4B-FP8-DSpark-watermark-TP2.yaml
@@ -0,0 +1,29 @@
+model_name: "Qwen/Qwen3-4B-FP8"
+num_questions: 1319
+num_fewshot: 5
+max_tokens: 256
+use_chat_completions: false
+temperature: 1.0
+seed: 42
+
+# Reference accuracy from the non-watermarked DSpark regression test.
+accuracy_threshold: 0.801
+tolerance: 0.02
+# Measured TP1 acceptance length: 4.115; allow a 0.2 regression.
+min_acceptance_length: 3.915
+
+env:
+  VLLM_USE_V2_MODEL_RUNNER: "1"
+server_args: >-
+  --tensor-parallel-size 2
+  --max-model-len 4096
+  --gpu-memory-utilization 0.92
+  --no-enable-prefix-caching
+  --speculative-config.method dspark
+  --speculative-config.model deepseek-ai/dspark_qwen3_4b_block7
+  --speculative-config.num_speculative_tokens 7
+  --speculative-config.draft_sample_method probabilistic
+  --watermark-config.algorithm dual_key_gumbel
+  --watermark-config.key 42
+  --watermark-config.context_width 4
+  --watermark-config.deduplicate_contexts none
diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc-ROCm.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc-ROCm.yaml
new file mode 100644
index 000000000000..879ad57c2029
--- /dev/null
+++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc-ROCm.yaml
@@ -0,0 +1,6 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.78
+num_questions: 1319
+num_fewshot: 5
+# Triton serves the unquantized boundary layers with TurboQuant-compatible layout.
+server_args: "--kv-cache-dtype turboquant_k3v4_nc --max-model-len 4096 --attention-backend TRITON_ATTN"
diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4-ROCm.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4-ROCm.yaml
new file mode 100644
index 000000000000..aa66561223bb
--- /dev/null
+++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4-ROCm.yaml
@@ -0,0 +1,6 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.80
+num_questions: 1319
+num_fewshot: 5
+# Triton serves the unquantized boundary layers with TurboQuant-compatible layout.
+server_args: "--kv-cache-dtype turboquant_k8v4 --max-model-len 4096 --attention-backend TRITON_ATTN"
diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc-ROCm.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc-ROCm.yaml
new file mode 100644
index 000000000000..103b37880a48
--- /dev/null
+++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc-ROCm.yaml
@@ -0,0 +1,6 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.75
+num_questions: 1319
+num_fewshot: 5
+# Triton serves the unquantized boundary layers with TurboQuant-compatible layout.
+server_args: "--kv-cache-dtype turboquant_3bit_nc --max-model-len 4096 --attention-backend TRITON_ATTN"
diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc-ROCm.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc-ROCm.yaml
new file mode 100644
index 000000000000..1348661b478a
--- /dev/null
+++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc-ROCm.yaml
@@ -0,0 +1,6 @@
+model_name: "Qwen/Qwen3-4B"
+accuracy_threshold: 0.80
+num_questions: 1319
+num_fewshot: 5
+# Triton serves the unquantized boundary layers with TurboQuant-compatible layout.
+server_args: "--kv-cache-dtype turboquant_4bit_nc --max-model-len 4096 --attention-backend TRITON_ATTN"
diff --git a/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml b/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml
new file mode 100644
index 000000000000..bdd42a0f1a78
--- /dev/null
+++ b/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml
@@ -0,0 +1,22 @@
+model_name: "Qwen/Qwen3.5-4B"
+# full GSM8K set is 1319 questions
+num_questions: 1319
+num_fewshot: 5
+max_concurrency: 128
+max_tokens: 1024
+use_chat_completions: true
+temperature: 0.7
+seed: 42
+
+# Calibrated by running non-watermaked 13 times
+# the threshold is the average accuracy, and the tolerance is set to 3 sigma (standard deviations)
+accuracy_threshold: 0.885
+tolerance: 0.015
+
+server_args: >-
+  --language-model-only
+  --max-model-len 8192
+  --max-num-seqs 128
+  --gpu-memory-utilization 0.4
+  --default-chat-template-kwargs '{"enable_thinking": false}'
+  --watermark-config '{"key": 42, "context_width": 4}'
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml
index b83d9e6a9e94..c0bb8d324002 100644
--- a/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml
@@ -3,6 +3,5 @@ accuracy_threshold: 0.39
 num_questions: 1319
 num_fewshot: 5
 server_args: >-
-  --enforce-eager
   --max-model-len 8192
   --linear-backend humming
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml
index 310255f6bf41..2ee157359a34 100644
--- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml
@@ -9,3 +9,5 @@ server_args: >-
   --moe-backend humming
   --linear-backend humming
   --dtype bfloat16
+env:
+  VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml
new file mode 100644
index 000000000000..a6bf4f652528
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml
@@ -0,0 +1,14 @@
+model_name: "nm-testing/Qwen3-30B-A3B-FP8-block"
+accuracy_threshold: 0.86
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --moe-backend humming
+  --linear-backend humming
+  --data-parallel-size 2
+  --enable-expert-parallel
+  --all2all-backend deepep_high_throughput
+env:
+  VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml
new file mode 100644
index 000000000000..cf8e01d70f97
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml
@@ -0,0 +1,14 @@
+model_name: "nm-testing/Qwen3-30B-A3B-FP8-block"
+accuracy_threshold: 0.86
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --moe-backend humming
+  --linear-backend humming
+  --data-parallel-size 2
+  --enable-expert-parallel
+  --all2all-backend deepep_high_throughput
+env:
+  VLLM_HUMMING_MOE_GEMM_TYPE: "indexed"
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml
index 8ca4777ce5e2..357e7a0257e6 100644
--- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml
@@ -3,7 +3,6 @@ accuracy_threshold: 0.86
 num_questions: 1319
 num_fewshot: 5
 server_args: >-
-  --enforce-eager
   --max-model-len 8192
   --moe-backend humming
   --linear-backend humming
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml
index 67725fce64ae..a8b289aaacd4 100644
--- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml
@@ -7,3 +7,5 @@ server_args: >-
   --max-model-len 8192
   --quantization humming
   --linear-backend humming
+env:
+  VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml
index a20c14331036..57923c63b0e4 100644
--- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml
@@ -7,3 +7,5 @@ server_args: >-
   --max-model-len 8192
   --moe-backend humming
   --linear-backend humming
+env:
+  VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml
deleted file mode 100644
index b58c0d710e4b..000000000000
--- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-model_name: "Qwen/Qwen3-30B-A3B"
-accuracy_threshold: 0.80
-num_questions: 1319
-num_fewshot: 5
-server_args: >-
-  --enforce-eager
-  --max-model-len 8192
-  --tensor-parallel-size 1
-  --quantization humming
-  --kernel-config.enable_flashinfer_autotune=False
-env:
-  VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}'
-  VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}'
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-4B-mixed-quant-RTN-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-4B-mixed-quant-RTN-humming.yaml
index 502ec7952268..e41aa130ae2a 100644
--- a/tests/evals/gsm8k/configs/humming/Qwen3-4B-mixed-quant-RTN-humming.yaml
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-4B-mixed-quant-RTN-humming.yaml
@@ -1,5 +1,5 @@
-model_name: "nm-testing/Qwen3-4B-mixed-quant-RTN-wNaM"
-accuracy_threshold: 0.78
+model_name: "nm-testing/Qwen3-4B-mixed-quant-RTN-wNaM-a4"
+accuracy_threshold: 0.70
 num_questions: 1319
 num_fewshot: 5
 server_args: >-
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-4B-mixed-quant-RTN-no-w4a4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-4B-mixed-quant-RTN-no-w4a4-humming.yaml
new file mode 100644
index 000000000000..965a0fe54d77
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/Qwen3-4B-mixed-quant-RTN-no-w4a4-humming.yaml
@@ -0,0 +1,8 @@
+model_name: "nm-testing/Qwen3-4B-mixed-quant-RTN-wNaM-asym"
+accuracy_threshold: 0.85
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 4096
+  --linear-backend humming
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
similarity index 62%
rename from tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml
rename to tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
index c932091dbe35..d2f04348b9b3 100644
--- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml
+++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
@@ -1,7 +1,8 @@
-model_name: "Qwen/Qwen3-30B-A3B"
-accuracy_threshold: 0.80
+model_name: "Qwen/Qwen3.5-35B-A3B"
+accuracy_threshold: 0.90
 num_questions: 1319
 num_fewshot: 5
+gen_prefix: " <think>\n\n</think>\n"
 server_args: >-
   --enforce-eager
   --max-model-len 8192
@@ -9,5 +10,5 @@ server_args: >-
   --quantization humming
   --kernel-config.enable_flashinfer_autotune=False
 env:
-  VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}'
+  VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}'
   VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}'
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
similarity index 57%
rename from tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml
rename to tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
index fa5095cc8824..08687aa6a307 100644
--- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml
+++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
@@ -1,7 +1,8 @@
-model_name: "Qwen/Qwen3-30B-A3B"
-accuracy_threshold: 0.80
+model_name: "Qwen/Qwen3.5-35B-A3B"
+accuracy_threshold: 0.90
 num_questions: 1319
 num_fewshot: 5
+gen_prefix: " <think>\n\n</think>\n"
 server_args: >-
   --enforce-eager
   --max-model-len 8192
@@ -9,4 +10,4 @@ server_args: >-
   --quantization humming
   --kernel-config.enable_flashinfer_autotune=False
 env:
-  VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}'
+  VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}'
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml
new file mode 100644
index 000000000000..067b96924d25
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml
@@ -0,0 +1,13 @@
+model_name: "Qwen/Qwen3.5-9B"
+accuracy_threshold: 0.88
+num_questions: 1319
+num_fewshot: 5
+gen_prefix: " <think>\n\n</think>\n"
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --quantization humming
+  --kernel-config.enable_flashinfer_autotune=False
+env:
+  VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}'
+  VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}'
diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml
new file mode 100644
index 000000000000..24915abbc2b9
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml
@@ -0,0 +1,12 @@
+model_name: "Qwen/Qwen3.5-9B"
+accuracy_threshold: 0.88
+num_questions: 1319
+num_fewshot: 5
+gen_prefix: " <think>\n\n</think>\n"
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --quantization humming
+  --kernel-config.enable_flashinfer_autotune=False
+env:
+  VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}'
diff --git a/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt
new file mode 100644
index 000000000000..0852d14633c2
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt
@@ -0,0 +1,4 @@
+gpt-oss-20b-humming.yaml
+Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml
+Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
+Qwen3-4B-mixed-quant-RTN-no-w4a4-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt
new file mode 100644
index 000000000000..512cbe99ff34
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt
@@ -0,0 +1,4 @@
+Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml
+Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml
+Qwen3-30B-A3B-AWQ-humming.yaml
+Qwen3-4B-mixed-quant-RTN-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt
index 63b428ac0e30..bdfa9033df0c 100644
--- a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt
+++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt
@@ -1,4 +1,5 @@
 Qwen3.5-35B-A3B-experts-int8-humming.yaml
+Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
 gpt-oss-20b-humming.yaml
 Qwen3-30B-A3B-GPTQ-Int4-humming.yaml
 Qwen3-30B-A3B-FP8-block-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-act-int8.txt b/tests/evals/gsm8k/configs/humming/config-act-int8.txt
index a7efbac3d033..7a608c7c5a66 100644
--- a/tests/evals/gsm8k/configs/humming/config-act-int8.txt
+++ b/tests/evals/gsm8k/configs/humming/config-act-int8.txt
@@ -3,4 +3,5 @@ Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml
 Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml
 Qwen3-30B-A3B-AWQ-humming.yaml
 Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml
-Qwen3-4B-mixed-quant-RTN-humming.yaml
+Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
+Qwen3-4B-mixed-quant-RTN-no-w4a4-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-deepep.txt b/tests/evals/gsm8k/configs/humming/config-deepep.txt
new file mode 100644
index 000000000000..2c97531d75f1
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/config-deepep.txt
@@ -0,0 +1,2 @@
+Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml
+Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt
index 60179f1159fe..737a92cf48ae 100644
--- a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt
+++ b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt
@@ -1,4 +1,5 @@
 Qwen3.5-35B-A3B-experts-int8-humming.yaml
+Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
 Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml
 Qwen3-30B-A3B-Fp8-v1-humming.yaml
 Qwen3-30B-A3B-NVFP4-humming.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt b/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt
deleted file mode 100644
index 2c10777a095a..000000000000
--- a/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-Qwen3-30B-A3B-int5wc-hadamard-humming.yaml
-Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml
-Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml
diff --git a/tests/evals/gsm8k/configs/humming/config-l4.txt b/tests/evals/gsm8k/configs/humming/config-l4.txt
new file mode 100644
index 000000000000..d1434b1cab2e
--- /dev/null
+++ b/tests/evals/gsm8k/configs/humming/config-l4.txt
@@ -0,0 +1,9 @@
+Qwen2-1.5B-Instruct-FP8W8-humming.yaml
+Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml
+Qwen3-0.6B-MXFP8-humming.yaml
+Qwen3-0.6B-MXFP8-humming-act-fp8.yaml
+Qwen3-30B-A3B-AWQ-humming.yaml
+Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml
+gpt-oss-20b-humming.yaml
+Qwen3.5-9B-int6wc-hadamard-humming.yaml
+Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml
diff --git a/tests/evals/gsm8k/configs/models-h200-deepseek-r1-dp.txt b/tests/evals/gsm8k/configs/models-h200-deepseek-r1-dp.txt
new file mode 100644
index 000000000000..9e293daf16f4
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-h200-deepseek-r1-dp.txt
@@ -0,0 +1 @@
+DeepSeek-R1-DP.yaml
diff --git a/tests/evals/gsm8k/configs/models-h200-deepseek-r1-tp.txt b/tests/evals/gsm8k/configs/models-h200-deepseek-r1-tp.txt
new file mode 100644
index 000000000000..30826a6f7698
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-h200-deepseek-r1-tp.txt
@@ -0,0 +1 @@
+DeepSeek-R1-TP.yaml
diff --git a/tests/evals/gsm8k/configs/models-h200-deepseek-v32-dp.txt b/tests/evals/gsm8k/configs/models-h200-deepseek-v32-dp.txt
new file mode 100644
index 000000000000..57faba312ffc
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-h200-deepseek-v32-dp.txt
@@ -0,0 +1 @@
+DeepSeek-V3.2-DP.yaml
diff --git a/tests/evals/gsm8k/configs/models-h200-deepseek-v32-tp.txt b/tests/evals/gsm8k/configs/models-h200-deepseek-v32-tp.txt
new file mode 100644
index 000000000000..027f168fd58e
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-h200-deepseek-v32-tp.txt
@@ -0,0 +1 @@
+DeepSeek-V3.2-TP.yaml
diff --git a/tests/evals/gsm8k/configs/models-h200-nemotron-3-super.txt b/tests/evals/gsm8k/configs/models-h200-nemotron-3-super.txt
new file mode 100644
index 000000000000..ec6fc214cdb8
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-h200-nemotron-3-super.txt
@@ -0,0 +1,2 @@
+Nemotron-3-Super-120B-A12B-BF16.yaml
+Nemotron-3-Super-120B-A12B-FP8.yaml
diff --git a/tests/evals/gsm8k/configs/models-h200.txt b/tests/evals/gsm8k/configs/models-h200.txt
deleted file mode 100644
index 255e93b67443..000000000000
--- a/tests/evals/gsm8k/configs/models-h200.txt
+++ /dev/null
@@ -1,6 +0,0 @@
-DeepSeek-R1-TP.yaml
-DeepSeek-R1-DP.yaml
-DeepSeek-V3.2-TP.yaml
-DeepSeek-V3.2-DP.yaml
-Nemotron-3-Super-120B-A12B-BF16.yaml
-Nemotron-3-Super-120B-A12B-FP8.yaml
diff --git a/tests/evals/gsm8k/configs/models-hisparse.txt b/tests/evals/gsm8k/configs/models-hisparse.txt
new file mode 100644
index 000000000000..53ff181f4f66
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-hisparse.txt
@@ -0,0 +1 @@
+GLM-5.2-NVFP4-HiSparse.yaml
diff --git a/tests/evals/gsm8k/configs/models-pcp.txt b/tests/evals/gsm8k/configs/models-pcp.txt
index df645edad57a..e6d5bc9d39e9 100644
--- a/tests/evals/gsm8k/configs/models-pcp.txt
+++ b/tests/evals/gsm8k/configs/models-pcp.txt
@@ -1,2 +1,3 @@
 GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
 GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant-k3v4nc-rocm.txt b/tests/evals/gsm8k/configs/models-turboquant-k3v4nc-rocm.txt
new file mode 100644
index 000000000000..9a1e252a4bec
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-turboquant-k3v4nc-rocm.txt
@@ -0,0 +1 @@
+Qwen3-4B-TQ-k3v4nc-ROCm.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant-k3v4nc.txt b/tests/evals/gsm8k/configs/models-turboquant-k3v4nc.txt
new file mode 100644
index 000000000000..836993e0fbdc
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-turboquant-k3v4nc.txt
@@ -0,0 +1 @@
+Qwen3-4B-TQ-k3v4nc.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant-k8v4-rocm.txt b/tests/evals/gsm8k/configs/models-turboquant-k8v4-rocm.txt
new file mode 100644
index 000000000000..458fe69f0087
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-turboquant-k8v4-rocm.txt
@@ -0,0 +1 @@
+Qwen3-4B-TQ-k8v4-ROCm.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant-k8v4.txt b/tests/evals/gsm8k/configs/models-turboquant-k8v4.txt
new file mode 100644
index 000000000000..3c0f67d9bfaa
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-turboquant-k8v4.txt
@@ -0,0 +1 @@
+Qwen3-4B-TQ-k8v4.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt b/tests/evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt
new file mode 100644
index 000000000000..cdc2462b8ce1
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt
@@ -0,0 +1 @@
+Qwen3-4B-TQ-t3nc-ROCm.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant-t3nc.txt b/tests/evals/gsm8k/configs/models-turboquant-t3nc.txt
new file mode 100644
index 000000000000..49b31f737d5c
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-turboquant-t3nc.txt
@@ -0,0 +1 @@
+Qwen3-4B-TQ-t3nc.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant-t4nc-rocm.txt b/tests/evals/gsm8k/configs/models-turboquant-t4nc-rocm.txt
new file mode 100644
index 000000000000..cd599778b6fb
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-turboquant-t4nc-rocm.txt
@@ -0,0 +1 @@
+Qwen3-4B-TQ-t4nc-ROCm.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant-t4nc.txt b/tests/evals/gsm8k/configs/models-turboquant-t4nc.txt
new file mode 100644
index 000000000000..8cf140fa69d4
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-turboquant-t4nc.txt
@@ -0,0 +1 @@
+Qwen3-4B-TQ-t4nc.yaml
diff --git a/tests/evals/gsm8k/configs/models-turboquant.txt b/tests/evals/gsm8k/configs/models-turboquant.txt
deleted file mode 100644
index 518aac780b90..000000000000
--- a/tests/evals/gsm8k/configs/models-turboquant.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-Qwen3-4B-TQ-k8v4.yaml
-Qwen3-4B-TQ-t4nc.yaml
-Qwen3-4B-TQ-k3v4nc.yaml
-Qwen3-4B-TQ-t3nc.yaml
diff --git a/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt b/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt
new file mode 100644
index 000000000000..127b0b47509f
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-watermark-dspark-tp2.txt
@@ -0,0 +1 @@
+Qwen3-4B-FP8-DSpark-watermark-TP2.yaml
diff --git a/tests/evals/gsm8k/configs/models-watermark.txt b/tests/evals/gsm8k/configs/models-watermark.txt
new file mode 100644
index 000000000000..a761b41c1b3f
--- /dev/null
+++ b/tests/evals/gsm8k/configs/models-watermark.txt
@@ -0,0 +1 @@
+Qwen3.5-4B-watermark.yaml
diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-BF16-moonep.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-BF16-moonep.yaml
new file mode 100644
index 000000000000..8704f29bcf1e
--- /dev/null
+++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-BF16-moonep.yaml
@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-30B-A3B"
+accuracy_threshold: 0.88
+num_questions: 1319
+num_fewshot: 5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend=moonep"
diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml
index 5f10684d2e7c..f81eb1c18824 100644
--- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml
+++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml
@@ -2,6 +2,8 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+# 30B MoE + DeepGEMM/DeepEP JIT compile can exceed the 1200s default startup wait.
+startup_max_wait_seconds: 2400
 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --disable-uvicorn-access-log"
 env:
   VLLM_USE_DEEP_GEMM: "1"
diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml
index 37e6039e9694..3987065f6ffd 100644
--- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml
+++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml
@@ -2,6 +2,8 @@ model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block"
 accuracy_threshold: 0.85
 num_questions: 1319
 num_fewshot: 5
+# 30B MoE + DeepGEMM/DeepEP JIT compile can exceed the 1200s default startup wait.
+startup_max_wait_seconds: 2400
 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --disable-uvicorn-access-log"
 env:
   VLLM_USE_DEEP_GEMM: "1"
diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml
index 6624cea1ef23..004e3e15e302 100644
--- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml
+++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml
@@ -2,4 +2,6 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+# 30B MoE + CuTeDSL/DeepEP JIT compile can exceed the 1200s default startup wait.
+startup_max_wait_seconds: 2400
 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --moe-backend=flashinfer_cutedsl"
diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml
index f2d4588e3aee..55892f21969b 100644
--- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml
+++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml
@@ -2,4 +2,6 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+# 30B MoE + CuTeDSL/DeepEP JIT compile can exceed the 1200s default startup wait.
+startup_max_wait_seconds: 2400
 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --moe-backend=flashinfer_cutedsl"
diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
index 7fb1f4968558..6dcb7414a567 100644
--- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
+++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
@@ -10,3 +10,4 @@ Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ht.yaml
 Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml
 Qwen3-30B-A3B-Fp8-CT-Block-deepgemm.yaml
 Qwen3-30B-A3B-BF16-triton.yaml
+Qwen3-30B-A3B-BF16-moonep.yaml
diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-triton-rocm.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-triton-rocm.yaml
new file mode 100644
index 000000000000..c6da3748b8b8
--- /dev/null
+++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-triton-rocm.yaml
@@ -0,0 +1,5 @@
+model_name: "Qwen/Qwen3-30B-A3B"
+accuracy_threshold: 0.88
+num_questions: 1319
+num_fewshot: 5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=triton"
diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-0.txt b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-0.txt
new file mode 100644
index 000000000000..4ca4943fc78f
--- /dev/null
+++ b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-0.txt
@@ -0,0 +1 @@
+Qwen3-30B-A3B-BF16-triton-rocm.yaml
diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-1.txt b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-1.txt
new file mode 100644
index 000000000000..c46cb2c243fa
--- /dev/null
+++ b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-1.txt
@@ -0,0 +1 @@
+Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml
diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-2.txt b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-2.txt
new file mode 100644
index 000000000000..928653a58039
--- /dev/null
+++ b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-2.txt
@@ -0,0 +1 @@
+Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml
diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-3.txt b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-3.txt
new file mode 100644
index 000000000000..8b175f5b0ef1
--- /dev/null
+++ b/tests/evals/gsm8k/configs/moe-refactor/config-rocm-shard-3.txt
@@ -0,0 +1 @@
+Mixtral-8x7B-Fp8-AutoFp8-triton.yaml
diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py
index b5a1e6a9cc29..7e2534b55a95 100644
--- a/tests/evals/gsm8k/gsm8k_eval.py
+++ b/tests/evals/gsm8k/gsm8k_eval.py
@@ -1,9 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Isolated GSM8K evaluation script for vLLM serve endpoint.
-"""
+"""Isolated GSM8K evaluation script for vLLM serve endpoint."""
 
 import argparse
 import ast
@@ -45,7 +43,7 @@ def download_and_cache_file(url: str, filename: str | None = None) -> str:
 
 
 def load_gsm8k_data() -> tuple[list[dict], list[dict]]:
-    """Load GSM8K train and test data"""
+    """Load GSM8K train and test data."""
     train_url = f"{VLLM_S3_BUCKET_URL}/ci-datasets/gsm8k/train.jsonl"
     test_url = f"{VLLM_S3_BUCKET_URL}/ci-datasets/gsm8k/test.jsonl"
 
@@ -91,6 +89,7 @@ async def call_vllm_api(
 
     Returns:
         Tuple of (response_text, completion_tokens)
+
     """
     data = {
         "prompt": prompt,
@@ -221,8 +220,7 @@ def evaluate_gsm8k(
     gen_prefix: str = "",
     max_concurrency: int | None = None,
 ) -> dict[str, float | int]:
-    """
-    Evaluate GSM8K accuracy using vLLM serve endpoint.
+    """Evaluate GSM8K accuracy using vLLM serve endpoint.
 
     Returns dict with accuracy, invalid_rate, latency, etc.
     """
@@ -364,6 +362,12 @@ def main() -> None:
         type=int,
         help="Maximum number of concurrent requests",
     )
+    parser.add_argument(
+        "--request-timeout-seconds",
+        type=float,
+        default=600,
+        help="Timeout for each request, including time waiting for a connection",
+    )
     parser.add_argument("--save-results", type=str, help="Save results to JSON file")
 
     args = parser.parse_args()
@@ -377,6 +381,7 @@ def main() -> None:
         temperature=args.temperature,
         seed=args.seed,
         max_concurrency=args.max_concurrency,
+        request_timeout_seconds=args.request_timeout_seconds,
     )
 
     # Print results to terminal
diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py
index 5a237cf94764..7411ca0f406a 100644
--- a/tests/evals/gsm8k/test_gsm8k_correctness.py
+++ b/tests/evals/gsm8k/test_gsm8k_correctness.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-GSM8K evaluation using vLLM server and isolated GSM8K script.
+"""GSM8K evaluation using vLLM server and isolated GSM8K script.
 Replacement for lm-eval-harness with better performance and control.
 
 Usage:
@@ -129,6 +128,12 @@ def test_gsm8k_correctness(config_filename):
                 "Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. "
                 "The quantization scheme is not supported on non-GFX950 platforms."
             )
+    if "--all2all-backend=moonep" in eval_config.get("server_args", ""):
+        from vllm.utils.import_utils import has_moonep
+
+        if not has_moonep():
+            pytest.skip("Skipping MoonEP config: the moonep package is not installed.")
+
     # Parse server arguments from config (use shlex to handle quoted strings)
     server_args_str = eval_config.get("server_args", "")
     server_args = shlex.split(server_args_str) if server_args_str else []
diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py
index c140a41e2644..bcc001e2bd43 100644
--- a/tests/evals/gsm8k/test_gsm8k_offloading.py
+++ b/tests/evals/gsm8k/test_gsm8k_offloading.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-GSM8K correctness test for CPU KV offloading connectors.
+"""GSM8K correctness test for CPU KV offloading connectors.
 
 Regression guard for stride computation bugs in the offloading worker
 (e.g. https://github.com/vllm-project/vllm/pull/46888) and silent KV
diff --git a/tests/evals/mrcr/test_mrcr_correctness.py b/tests/evals/mrcr/test_mrcr_correctness.py
index 3adfd4dc8799..f381a592dde6 100644
--- a/tests/evals/mrcr/test_mrcr_correctness.py
+++ b/tests/evals/mrcr/test_mrcr_correctness.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-MRCR long-context accuracy test.
+"""MRCR long-context accuracy test.
 
 Usage:
     pytest -s -v tests/evals/mrcr/test_mrcr_correctness.py \
diff --git a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml
new file mode 100644
index 000000000000..db43869eb4ae
--- /dev/null
+++ b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml
@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+model_name: "Qwen/Qwen3.8-Flash-Next-FP8"
+datasets:
+  gsm8k:
+    metric_threshold: 0.98
+    tolerance: 0.02
+  aime25:
+    metric_threshold: 0.90
+    tolerance: 0.05
+eval_batch_size: 32
+startup_max_wait_seconds: 1800
+generation_config:
+  seed: 1236
+  do_sample: true
+  temperature: 0.7
+  top_p: 0.8
+  top_k: 20
+  repetition_penalty: 1.0
+  presence_penalty: 1.5
+  max_tokens: 32768
+  stream: true
+  extra_body:
+    chat_template_kwargs:
+      enable_thinking: true
+server_args: "--tensor-parallel-size 4 --mamba-cache-mode align"
diff --git a/tests/evals/qwen4_exp/configs/models-rocm.txt b/tests/evals/qwen4_exp/configs/models-rocm.txt
new file mode 100644
index 000000000000..9999574a4d43
--- /dev/null
+++ b/tests/evals/qwen4_exp/configs/models-rocm.txt
@@ -0,0 +1 @@
+Qwen3.8-Flash-Next-FP8-rocm.yaml
diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py
index a794ff1e3a67..15b9835d6cf5 100644
--- a/tests/fusion/test_quant_activation_contract.py
+++ b/tests/fusion/test_quant_activation_contract.py
@@ -9,8 +9,13 @@
     _POSSIBLE_FP8_BLOCK_KERNELS,
     _POSSIBLE_FP8_KERNELS,
     _POSSIBLE_INT8_KERNELS,
+    _POSSIBLE_MXFP8_KERNELS,
     _POSSIBLE_NVFP4_KERNELS,
 )
+from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+    FlashInferCutedslMxfp8LinearKernel,
+    FlashInferCutlassMxfp8LinearKernel,
+)
 from vllm.model_executor.kernels.linear.nvfp4.base import (
     NvFp4LinearKernel,
     NvFp4LinearLayerConfig,
@@ -19,6 +24,11 @@
     FlashInferCutlassNvFp4LinearKernel,
     FlashInferTrtllmNvFp4LinearKernel,
 )
+from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
+    AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
+    AiterPerTokenFp8ScaledMMLinearKernel,
+    AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
+)
 from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
     CutlassFP8ScaledMMLinearKernel,
 )
@@ -37,6 +47,7 @@
     QuantizedActivation,
     as_quantized_activation,
     expose_input_quant_key,
+    get_input_quant_key,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8StaticTensorSym,
@@ -50,6 +61,11 @@
     FlashInferFP8ScaledMMLinearKernel,
     FlashInferCutlassNvFp4LinearKernel,
     PerTensorTorchFP8ScaledMMLinearKernel,
+    AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
+    AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
+    AiterPerTokenFp8ScaledMMLinearKernel,
+    FlashInferCutedslMxfp8LinearKernel,
+    FlashInferCutlassMxfp8LinearKernel,
 }
 
 
@@ -60,6 +76,7 @@ def _all_kernel_classes() -> list[type]:
         _POSSIBLE_FP8_BLOCK_KERNELS,
         _POSSIBLE_INT8_KERNELS,
         _POSSIBLE_NVFP4_KERNELS,
+        _POSSIBLE_MXFP8_KERNELS,
     ):
         for kernels in registry.values():
             for cls in kernels:
@@ -110,13 +127,13 @@ def test_bridge_marks_supporting_and_skips_others():
     supported = _probe(FlashInferCutlassNvFp4LinearKernel)
     layer = torch.nn.Module()
     expose_input_quant_key(layer, supported)
-    assert layer.input_quant_key == kNvfp4Dynamic
+    assert get_input_quant_key(layer) == kNvfp4Dynamic
 
     unsupported = _probe(FlashInferTrtllmNvFp4LinearKernel)
     assert unsupported.input_quant_key() is None
     layer = torch.nn.Module()
     expose_input_quant_key(layer, unsupported)
-    assert not hasattr(layer, "input_quant_key")
+    assert get_input_quant_key(layer) is None
 
 
 def test_as_quantized_activation_validates_key():
diff --git a/tests/ir/ir_test_utils.py b/tests/ir/ir_test_utils.py
index a82206b6f72e..5cd05b9bcf09 100644
--- a/tests/ir/ir_test_utils.py
+++ b/tests/ir/ir_test_utils.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shared test utilities for vLLM IR op correctness tests.
-"""
+"""Shared test utilities for vLLM IR op correctness tests."""
 
 import torch
 
diff --git a/tests/ir/test_op.py b/tests/ir/test_op.py
index ca8185b2d3dc..5033b0ec10f3 100644
--- a/tests/ir/test_op.py
+++ b/tests/ir/test_op.py
@@ -34,7 +34,7 @@ def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
     @_custom_add.register_impl("impl_b", inplace=True)
     def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        """Computes x+y+20"""
+        """Computes x+y+20."""
         x.add_(y)
         x.add_(20)
         return x
@@ -101,7 +101,7 @@ def _custom_kwarg_op(
 class TestIrOpCustomAdd:
     # Registration invariants
     def test_decorated_object(self, custom_add_op):
-        """Make sure that referring directly to an op is correct"""
+        """Make sure that referring directly to an op is correct."""
         _custom_add = custom_add_op
         assert isinstance(_custom_add, IrOp)
         assert "_custom_add" in IrOp.registry
diff --git a/tests/jit_monitor/test_no_runtime_jit_rocm.py b/tests/jit_monitor/test_no_runtime_jit_rocm.py
new file mode 100644
index 000000000000..a99b8e1e1410
--- /dev/null
+++ b/tests/jit_monitor/test_no_runtime_jit_rocm.py
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Keep the MRV2 sampler warmup registry populated on ROCm."""
+
+import pytest
+
+from vllm import LLM
+from vllm.platforms import current_platform
+
+from ..models.utils import dummy_hf_overrides
+from ..utils import create_new_process_for_each_test
+from .test_no_runtime_jit import _run_shape_battery
+
+pytestmark = pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm")
+
+
+def _worker_monitor_state(worker):
+    from vllm.utils import jit_monitor
+
+    return (
+        jit_monitor.is_active(),
+        worker.observability_config.jit_monitor_mode,
+    )
+
+
+@create_new_process_for_each_test("spawn")
+def _run_rocm_shape_battery() -> None:
+    llm = LLM(
+        "Qwen/Qwen3-0.6B",
+        max_model_len=2048,
+        max_num_seqs=8,
+        gpu_memory_utilization=0.03,
+        kv_cache_memory_bytes=256 * 1024 * 1024,
+        load_format="dummy",
+        hf_overrides=dummy_hf_overrides,
+        enforce_eager=False,
+        jit_monitor_mode="error",
+    )
+    try:
+        states = llm.collective_rpc(_worker_monitor_state, timeout=30)
+        assert states and all(state == (True, "error") for state in states)
+        _run_shape_battery(llm)
+    finally:
+        llm.llm_engine.engine_core.shutdown(timeout=30)
+
+
+def test_v2_sampler_warmup_rocm(monkeypatch, tmp_path):
+    """Warmup must cover the existing prefill, decode, and sampler battery.
+
+    A fresh cache and child process prevent previous model executions from
+    hiding a missing warmup specialization.
+    """
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "0")
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "triton"))
+    monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm"))
+    monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor"))
+    _run_rocm_shape_battery()
diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py
index d74a119e654d..cb72c27904b8 100644
--- a/tests/kernels/attention/test_attention_selector.py
+++ b/tests/kernels/attention/test_attention_selector.py
@@ -341,61 +341,6 @@ def test_fp32_fallback(device: str):
                 get_attn_backend(16, torch.float32, None)
 
 
-def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
-    """Test FlashAttn validation."""
-    pytest.skip(
-        "Skipping as current backend selector does not "
-        "handle fallbacks when a backend is explicitly set."
-    )
-
-    attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN)
-    cache_config = CacheConfig(block_size=16)
-    vllm_config = VllmConfig(
-        attention_config=attention_config, cache_config=cache_config
-    )
-
-    with set_current_vllm_config(vllm_config):
-        # Unsupported CUDA arch
-        monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _=None: (7, 5))
-        backend = get_attn_backend(16, torch.float16, None)
-        assert backend.get_name() != "FLASH_ATTN"
-
-        # Reset the monkeypatch for subsequent tests
-        monkeypatch.undo()
-
-        # Unsupported data type
-        backend = get_attn_backend(16, torch.float8_e4m3fn, None)
-        assert backend.get_name() != "FLASH_ATTN"
-
-        # Unsupported kv cache data type
-        backend = get_attn_backend(16, torch.float16, "fp8")
-        assert backend.get_name() != "FLASH_ATTN"
-
-        # Unsupported block size
-        vllm_config.cache_config.block_size = 8
-        backend = get_attn_backend(16, torch.float16, None)
-        assert backend.get_name() != "FLASH_ATTN"
-
-        # flash-attn is not installed
-        import sys
-
-        vllm_config.cache_config.block_size = 16
-        original_module = sys.modules.get("vllm_flash_attn")
-        monkeypatch.setitem(sys.modules, "vllm_flash_attn", None)
-        backend = get_attn_backend(16, torch.float16, None)
-        assert backend.get_name() != "FLASH_ATTN"
-
-        # Restore the original module if it existed
-        if original_module is not None:
-            monkeypatch.setitem(sys.modules, "vllm_flash_attn", original_module)
-        else:
-            monkeypatch.delitem(sys.modules, "vllm_flash_attn", raising=False)
-
-        # Unsupported head size
-        backend = get_attn_backend(17, torch.float16, None)
-        assert backend.get_name() != "FLASH_ATTN"
-
-
 def test_invalid_backend():
     """Test that invalid attention backend names raise ValueError."""
     with (
@@ -808,3 +753,138 @@ def test_fa4_hd256_impl_selection(attn_type, sliding_window, expected):
         )
     assert impl.vllm_flash_attn_version == expected
     assert impl.fa4_hd256 == (expected == 4)
+
+
+@pytest.fixture
+def blackwell_selection():
+    with (
+        patch.object(
+            type(current_platform),
+            "get_device_capability",
+            return_value=DeviceCapability(10, 0),
+        ),
+        patch(
+            "vllm.v1.attention.backends.fa_utils.is_fa_version_supported",
+            return_value=True,
+        ),
+    ):
+        yield
+
+
+@pytest.fixture
+def hopper_selection():
+    with (
+        patch.object(
+            type(current_platform),
+            "get_device_capability",
+            return_value=DeviceCapability(9, 0),
+        ),
+        patch(
+            "vllm.v1.attention.backends.fa_utils.is_fa_version_supported",
+            return_value=True,
+        ),
+    ):
+        yield
+
+
+@blackwell_only
+@pytest.mark.parametrize("use_mm_prefix", [False, True])
+@pytest.mark.parametrize("flash_attn_version", [None, 3, 4])
+def test_hopper_mm_prefix_selects_triton_flash_attn(
+    use_mm_prefix, flash_attn_version, hopper_selection
+):
+    """Hopper selects the composite when its causal route resolves FA3 or FA4."""
+    from vllm.engine.arg_utils import EngineArgs
+
+    config = EngineArgs(
+        model="google/gemma-4-31B-it",
+        dtype="bfloat16",
+        attention_config=AttentionConfig(flash_attn_version=flash_attn_version),
+    ).create_engine_config()
+    with set_current_vllm_config(config):
+        backend = get_attn_backend(
+            256, torch.bfloat16, None, use_mm_prefix=use_mm_prefix
+        )
+    expected = "TRITON_FLASH_ATTN" if use_mm_prefix else "FLASH_ATTN"
+    assert backend.get_name() == expected
+
+
+@blackwell_only
+@pytest.mark.parametrize("use_mm_prefix", [False, True])
+@pytest.mark.parametrize("kv_cache_dtype", [None, "fp8_e4m3"])
+def test_mm_prefix_selects_composite_without_changing_causal_default(
+    use_mm_prefix, kv_cache_dtype, blackwell_selection
+):
+    """The image-mask requirement must reach CUDA's automatic backend priority."""
+    from vllm.engine.arg_utils import EngineArgs
+
+    config = EngineArgs(
+        model="google/gemma-4-31B-it", dtype="bfloat16"
+    ).create_engine_config()
+    with set_current_vllm_config(config):
+        backend = get_attn_backend(
+            256, torch.bfloat16, kv_cache_dtype, use_mm_prefix=use_mm_prefix
+        )
+    expected = "TRITON_FLASHINFER" if use_mm_prefix else "FLASHINFER"
+    assert backend.get_name() == expected
+
+
+@blackwell_only
+@pytest.mark.parametrize(
+    "feature_kwargs,reason",
+    [
+        ({"use_rswa": True}, "R-SWA not supported"),
+        ({"kv_cache_dtype": "int8_per_token_head"}, "kv_cache_dtype not supported"),
+        ({"block_size": 16}, "block_size not supported"),
+        ({"use_dcp": True}, "DCP not supported"),
+        ({"use_pcp": True}, "PCP not supported"),
+        ({"use_adaptive_verification": True}, "device-cpu query lens mismatch"),
+    ],
+)
+def test_composite_rejects_features_not_shared_by_both_routes(feature_kwargs, reason):
+    """Automatic selection must not admit features only one child can execute."""
+    from vllm.engine.arg_utils import EngineArgs
+    from vllm.v1.attention.backends.triton_flashinfer import TritonFlashInferBackend
+
+    config = EngineArgs(
+        model="google/gemma-4-31B-it", dtype="bfloat16"
+    ).create_engine_config()
+    args = dict(
+        head_size=256,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="auto",
+        block_size=128,
+        use_mla=False,
+        has_sink=False,
+        use_sparse=False,
+        use_mm_prefix=True,
+        use_per_head_quant_scales=False,
+        device_capability=DeviceCapability(10, 0),
+        attn_type=AttentionType.DECODER,
+    )
+    args.update(feature_kwargs)
+    with set_current_vllm_config(config):
+        reasons = TritonFlashInferBackend.validate_configuration(**args)
+    assert any(reason in message for message in reasons), reasons
+
+
+@blackwell_only
+def test_rswa_selection_does_not_reuse_causal_result(blackwell_selection):
+    """An R-SWA requirement must invalidate a previously cached causal selection."""
+    from vllm.engine.arg_utils import EngineArgs
+
+    config = EngineArgs(
+        model="google/gemma-4-31B-it",
+        dtype="bfloat16",
+        attention_config={"backend": "TRITON_FLASHINFER"},
+    ).create_engine_config()
+    with set_current_vllm_config(config):
+        assert (
+            get_attn_backend(256, torch.bfloat16, None).get_name()
+            == "TRITON_FLASHINFER"
+        )
+        config.model_config.model_arch_config.rswa_window = 128
+        with pytest.raises(ValueError, match="R-SWA"):
+            get_attn_backend(256, torch.bfloat16, None)
+        config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
+        assert get_attn_backend(256, torch.bfloat16, None).get_name() == "TRITON_ATTN"
diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py
index 6fabcaa789b3..4bc0ce1a67b7 100644
--- a/tests/kernels/attention/test_cache.py
+++ b/tests/kernels/attention/test_cache.py
@@ -805,6 +805,94 @@ def test_concat_and_cache_mla(
         torch.testing.assert_close(kv_cache, ref_kv_cache)
 
 
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+@pytest.mark.parametrize("shared_slot_mapping", [False, True])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e5m2"])
+@torch.inference_mode()
+def test_concat_and_cache_mla_grouped(
+    device: str,
+    shared_slot_mapping: bool,
+    kv_cache_dtype: str,
+) -> None:
+    if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm():
+        pytest.skip("fp8_e5m2 KV cache is not supported on ROCm/HIP")
+
+    set_random_seed(0)
+    torch.set_default_device(device)
+    torch.accelerator.set_device_index(device)
+
+    num_layers = 5
+    num_tokens = 42
+    num_blocks = 8
+    block_size = 16
+    kv_lora_rank = 512
+    qk_rope_head_dim = 64
+    entry_size = kv_lora_rank + qk_rope_head_dim
+    total_slots = num_blocks * block_size
+
+    kv_c = torch.randn(num_layers, num_tokens, kv_lora_rank, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_layers, num_tokens, qk_rope_head_dim, dtype=torch.bfloat16)
+    kv_caches = torch.zeros(
+        num_layers,
+        num_blocks,
+        block_size,
+        entry_size,
+        dtype=torch.bfloat16 if kv_cache_dtype == "auto" else torch.uint8,
+    )
+    reference = torch.zeros_like(kv_caches)
+    scales = torch.linspace(0.05, 0.25, num_layers, dtype=torch.float32)
+
+    slot_mapping = torch.stack(
+        [torch.randperm(total_slots)[:num_tokens] for _ in range(num_layers)]
+    )
+    slot_mapping[:, -1] = -1
+    if shared_slot_mapping:
+        slot_mapping = slot_mapping[:1].expand(num_layers, -1)
+
+    for layer_idx in range(num_layers):
+        ops.concat_and_cache_mla(
+            kv_c[layer_idx],
+            k_pe[layer_idx],
+            reference[layer_idx],
+            slot_mapping[layer_idx],
+            kv_cache_dtype,
+            scales[layer_idx],
+        )
+
+    cache_ptrs = torch.tensor(
+        [kv_caches[layer_idx].data_ptr() for layer_idx in range(num_layers)],
+        dtype=torch.int64,
+    )
+    ref_cache = kv_caches[0]
+
+    def run_grouped(kv_scales: torch.Tensor | None) -> None:
+        ops.concat_and_cache_mla_grouped(
+            kv_c,
+            k_pe,
+            cache_ptrs,
+            slot_mapping,
+            ref_cache.size(1),
+            ref_cache.stride(0),
+            ref_cache.stride(1),
+            kv_scales,
+            kv_cache_dtype,
+        )
+
+    run_grouped(None if kv_cache_dtype == "auto" else scales)
+
+    torch.testing.assert_close(kv_caches, reference, rtol=0, atol=0)
+
+    if kv_cache_dtype == "fp8" and not shared_slot_mapping:
+        noncontiguous_scales = torch.ones(num_layers, 2)[:, 0]
+        assert not noncontiguous_scales.is_contiguous()
+        for invalid_scales, error in (
+            (scales.cpu(), "same CUDA device"),
+            (noncontiguous_scales, "must be contiguous"),
+        ):
+            with pytest.raises(RuntimeError, match=error):
+                run_grouped(invalid_scales)
+
+
 @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
 @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
@@ -872,14 +960,14 @@ def test_concat_and_cache_ds_mla(
             tile_end = (tile_idx + 1) * 128
             tile_data[:] = kv_c_data[tile_start:tile_end]
 
-            # tile_scale = tile_data.amax().to(torch.float32) / 448.
-            # NOTE: Using torch's amax() gives different results,
-            # so this must be manually computed.
+            # Using torch's amax() gives different results, so this must be
+            # manually computed.
             tile_data_float = tile_data.to(torch.float32)
             manual_max = abs(tile_data_float[0])
             for j in range(1, 128):
                 manual_max = max(manual_max, abs(tile_data_float[j]))
-            tile_scale = manual_max / 448.0
+            raw_scale = torch.clamp(manual_max / 448.0, min=1e-4)
+            tile_scale = torch.exp2(torch.ceil(torch.log2(raw_scale)))
 
             ref_cache_32bit[kv_lora_rank // 4 + tile_idx] = tile_scale
 
@@ -920,10 +1008,54 @@ def test_concat_and_cache_ds_mla(
         ref_rope = ref_cache_slice.view(dtype)[kv_lora_rank // 2 + 8 :]
 
         torch.testing.assert_close(kv_nope, ref_nope, atol=0.001, rtol=0.1)
-        torch.testing.assert_close(kv_scales, ref_scales, atol=0.001, rtol=0.1)
+        torch.testing.assert_close(kv_scales, ref_scales, atol=0, rtol=0)
         torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1)
 
 
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+@pytest.mark.parametrize("block_size", [64, 256])
+@torch.inference_mode()
+def test_concat_and_cache_ds_mla_nope(device: str, block_size: int) -> None:
+    """NoPE matches zero RoPE, clears valid tails, and preserves unused slots."""
+    dtype = torch.bfloat16
+    if current_platform.is_rocm():
+        pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm")
+    num_tokens, num_blocks = 3, 2
+    set_random_seed(0)
+    torch.set_default_device(device)
+    torch.accelerator.set_device_index(device)
+
+    slot_mapping = torch.tensor(
+        [block_size - 1, block_size, -1], dtype=torch.long, device=device
+    )
+    kv_c = torch.randn(num_tokens, 512, dtype=dtype, device=device)
+    k_pe = torch.empty(num_tokens, 0, dtype=dtype, device=device)
+    zero_k_pe = torch.zeros(num_tokens, 64, dtype=dtype, device=device)
+    scale = torch.tensor(1.0, dtype=torch.float32, device=device)
+    kv_cache = torch.full(
+        (num_blocks, block_size, 656), 0xFF, dtype=torch.uint8, device=device
+    )
+    ref_cache = kv_cache.clone()
+
+    opcheck(
+        torch.ops._C_cache_ops.concat_and_cache_mla,
+        (kv_c, k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale),
+        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
+    )
+    kv_cache.fill_(0xFF)
+    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale)
+    ops.concat_and_cache_mla(
+        kv_c, zero_k_pe, ref_cache, slot_mapping, "fp8_ds_mla", scale
+    )
+
+    torch.testing.assert_close(kv_cache, ref_cache, atol=0, rtol=0)
+    rows = kv_cache.view(num_blocks * block_size, 656)
+    assert (rows[slot_mapping[:2], 528:] == 0).all()
+    untouched = torch.ones(num_blocks * block_size, dtype=torch.bool, device=device)
+    untouched[slot_mapping[:2]] = False
+    assert (rows[untouched] == 0xFF).all()
+
+
 # Bytes per token for the nvfp4_ds_mla cache layout (see flashmla_sparse.py).
 NVFP4_DS_MLA_ENTRY_SIZE = 352
 
diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py
index 7c4b7811c78b..f79b1b15d2d9 100644
--- a/tests/kernels/attention/test_cpu_attn.py
+++ b/tests/kernels/attention/test_cpu_attn.py
@@ -3,13 +3,20 @@
 
 import functools
 import math
+from types import SimpleNamespace
 
 import pytest
 import torch
 
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
 from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.attention.backend import AttentionType
 from vllm.v1.attention.backends.cpu_attn import _get_attn_isa
+from vllm.v1.attention.backends.zentorch_sdpa import (
+    should_use_zentorch_sdpa,
+    zentorch_sdpa_attn,
+)
 
 if not current_platform.is_cpu():
     pytest.skip("skipping CPU-only tests", allow_module_level=True)
@@ -47,6 +54,24 @@
 ENCODER_SEQ_LENS = [
     [1, 678, 2367, 145, 4162, 36, 7812],
 ]
+# Uniform (dense zentorch_sdpa call) and ragged (per-sequence loop).
+ZEN_ENCODER_SEQ_LENS = [
+    [128, 128, 128],
+    [1, 67, 233, 5],
+]
+_skip_no_zentorch_sdpa = pytest.mark.skipif(
+    not current_platform.is_zen_cpu() or not has_zentorch_op(["zentorch_sdpa_attn"]),
+    reason="zentorch_sdpa requires a Zen CPU with the op registered",
+)
+# should_use_zentorch_sdpa additionally gates each dtype on its ISA.
+_skip_no_avx512_bf16 = pytest.mark.skipif(
+    not torch.cpu._is_avx512_bf16_supported(),
+    reason="zentorch_sdpa bfloat16 requires AVX512-BF16",
+)
+_skip_no_avx512 = pytest.mark.skipif(
+    not torch.cpu._is_avx512_supported(),
+    reason="zentorch_sdpa float32 requires AVX512",
+)
 
 
 def get_attn_isa(
@@ -205,6 +230,7 @@ def ref_varlen_encoder_attn(
     seq_lens: list[int],
     scale: float,
     sliding_window: int | None = None,
+    alibi_slopes: torch.Tensor | None = None,
 ) -> torch.Tensor:
     num_seqs = len(seq_lens)
     dtype = query.dtype
@@ -232,6 +258,11 @@ def ref_varlen_encoder_attn(
         else:
             mask = empty_mask.logical_not()
 
+        if alibi_slopes is not None:
+            q_pos = torch.arange(0, seq_len)[None, :, None]
+            kv_pos = torch.arange(0, seq_len)[None, None, :]
+            attn += -alibi_slopes[:, None, None].float() * (q_pos - kv_pos)
+
         attn.masked_fill_(mask, float("-inf"))
         attn = torch.softmax(attn, dim=-1)
         out = torch.einsum("hqk,khd->qhd", attn, v).to(dtype=dtype)
@@ -428,6 +459,7 @@ def varlen_with_paged_kv(
     k_scale: float = 1.0,
     v_scale: float = 1.0,
     dynamic_causal: list[bool] | None = None,
+    s_aux_dtype: torch.dtype = torch.bfloat16,
 ) -> None:
     set_random_seed(0)
     num_seqs = len(seq_lens)
@@ -449,9 +481,7 @@ def varlen_with_paged_kv(
     # 2^(-8/n)
     alibi_slopes = _get_alibi_slopes(num_query_heads) if use_alibi else None
 
-    s_aux = (
-        15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None
-    )
+    s_aux = 15 * torch.rand((num_query_heads,), dtype=s_aux_dtype) if use_sink else None
 
     is_fp8 = kv_cache_dtype != "auto"
     if is_fp8 and current_platform.get_cpu_architecture() != CpuArchEnum.X86:
@@ -538,6 +568,8 @@ def varlen_with_paged_kv(
     )
 
     out_without_split = torch.empty_like(query)
+    if s_aux is not None and s_aux.dtype != torch.bfloat16:
+        s_aux = s_aux.to(torch.float32)
     cpu_attention_with_kv_cache(
         query=query,
         key_cache=packed_key_cache,
@@ -771,6 +803,108 @@ def test_varlen_encoder_attention_amx(
     )
 
 
+@torch.inference_mode()
+def varlen_encoder_zentorch_sdpa(
+    seq_lens: list[int],
+    num_heads: tuple[int, int],
+    head_size: int,
+    sliding_window: int | None,
+    dtype: torch.dtype,
+    use_alibi: bool,
+) -> None:
+    set_random_seed(0)
+    num_query_heads, num_kv_heads = num_heads
+    assert num_query_heads % num_kv_heads == 0
+    scale = head_size**-0.5
+    token_num = sum(seq_lens)
+
+    query = tensor_cache(
+        token_num * num_query_heads * head_size, dtype, tag="query"
+    ).view(token_num, num_query_heads, head_size)
+    key_value = tensor_cache(
+        2 * token_num * num_kv_heads * head_size, dtype, tag="kv"
+    ).view(2, token_num, num_kv_heads, head_size)
+    key, value = key_value.unbind(0)
+
+    query_start_loc = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)
+    torch.cumsum(torch.tensor(seq_lens, dtype=torch.int32), 0, out=query_start_loc[1:])
+    alibi_slopes = _get_alibi_slopes(num_query_heads) if use_alibi else None
+    output = torch.empty_like(query)
+
+    zentorch_sdpa_attn(
+        query,
+        key,
+        value,
+        output,
+        SimpleNamespace(query_start_loc=query_start_loc, causal=False),
+        scale,
+        sliding_window if sliding_window is not None else -1,
+        alibi_slopes,
+    )
+    ref_output = ref_varlen_encoder_attn(
+        query=query,
+        key=key,
+        value=value,
+        seq_lens=seq_lens,
+        scale=scale,
+        sliding_window=sliding_window,
+        alibi_slopes=alibi_slopes,
+    )
+    torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
+
+
+_ZENTORCH_SDPA_DTYPES = [
+    pytest.param(torch.bfloat16, marks=_skip_no_avx512_bf16),
+    pytest.param(torch.float32, marks=_skip_no_avx512),
+]
+
+
+@_skip_no_zentorch_sdpa
+@pytest.mark.parametrize("seq_lens", ZEN_ENCODER_SEQ_LENS)
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
+@pytest.mark.parametrize("head_size", [128])
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
+@pytest.mark.parametrize("dtype", _ZENTORCH_SDPA_DTYPES)
+@pytest.mark.parametrize("use_alibi", [False, True])
+def test_zentorch_sdpa_attn(
+    seq_lens: list[int],
+    num_heads: tuple[int, int],
+    head_size: int,
+    sliding_window: int | None,
+    dtype: torch.dtype,
+    use_alibi: bool,
+) -> None:
+    varlen_encoder_zentorch_sdpa(
+        seq_lens=seq_lens,
+        num_heads=num_heads,
+        head_size=head_size,
+        sliding_window=sliding_window,
+        dtype=dtype,
+        use_alibi=use_alibi,
+    )
+
+
+@_skip_no_zentorch_sdpa
+@pytest.mark.parametrize(
+    "attn_type", [AttentionType.ENCODER_ONLY, AttentionType.ENCODER]
+)
+@pytest.mark.parametrize("dtype", _ZENTORCH_SDPA_DTYPES)
+def test_should_use_zentorch_sdpa(attn_type: str, dtype: torch.dtype) -> None:
+    assert should_use_zentorch_sdpa(attn_type, dtype)
+
+
+@_skip_no_zentorch_sdpa
+@pytest.mark.parametrize(
+    "attn_type,dtype",
+    [
+        (AttentionType.DECODER, torch.bfloat16),
+        (AttentionType.ENCODER_ONLY, torch.float16),
+    ],
+)
+def test_should_not_use_zentorch_sdpa(attn_type: str, dtype: torch.dtype) -> None:
+    assert not should_use_zentorch_sdpa(attn_type, dtype)
+
+
 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3", "fp8_e5m2"])
 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
 @pytest.mark.parametrize("num_heads", NUM_HEADS)
@@ -1128,6 +1262,51 @@ def test_varlen_with_paged_kv_sink(
     )
 
 
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
+@pytest.mark.parametrize("seq_lens", SEQ_LENS)
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
+@pytest.mark.parametrize("head_size", [96])
+@pytest.mark.parametrize("block_size", [128])
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("soft_cap", [None])
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
+@pytest.mark.parametrize("use_alibi", [False])
+@pytest.mark.parametrize("use_sink", [True])
+@pytest.mark.parametrize("isa", [get_attn_isa()])
+@pytest.mark.parametrize("s_aux_dtype", [torch.float16])
+def test_varlen_with_paged_kv_sink_fp16(
+    seq_lens: list[tuple[int, int]],
+    num_heads: tuple[int, int],
+    head_size: int,
+    sliding_window: int | None,
+    dtype: torch.dtype,
+    block_size: int,
+    soft_cap: float | None,
+    num_blocks: int,
+    use_alibi: bool,
+    use_sink: bool,
+    isa: str,
+    kv_cache_dtype: str,
+    s_aux_dtype: torch.dtype,
+) -> None:
+    varlen_with_paged_kv(
+        seq_lens=seq_lens,
+        num_heads=num_heads,
+        head_size=head_size,
+        sliding_window=sliding_window,
+        dtype=dtype,
+        block_size=block_size,
+        soft_cap=soft_cap,
+        num_blocks=num_blocks,
+        use_alibi=use_alibi,
+        use_sink=use_sink,
+        isa=isa,
+        kv_cache_dtype=kv_cache_dtype,
+        s_aux_dtype=s_aux_dtype,
+    )
+
+
 @pytest.mark.parametrize(
     "kv_cache_dtype",
     [
diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py
index 109f59bc8508..c6cba8d530e4 100644
--- a/tests/kernels/attention/test_flashinfer.py
+++ b/tests/kernels/attention/test_flashinfer.py
@@ -96,6 +96,7 @@ def _make_paged_kv_metadata(
         kv_indices         – CUDA int32, shape [total_blocks]
         kv_last_page_lens  – CPU int32, shape [num_seqs]
         block_tables       – CUDA int32, shape [num_seqs, max_blocks_per_seq]
+
     """
     num_seqs = len(kv_lens)
     max_blocks = (max(kv_lens) + block_size - 1) // block_size
diff --git a/tests/kernels/attention/test_flashinfer_mla_decode.py b/tests/kernels/attention/test_flashinfer_mla_decode.py
index a49ae3ae1ad3..c3f7b4c96f78 100644
--- a/tests/kernels/attention/test_flashinfer_mla_decode.py
+++ b/tests/kernels/attention/test_flashinfer_mla_decode.py
@@ -130,69 +130,123 @@ def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
     torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2)
 
 
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize("mode", ["tp", "dcp", "hisparse"])
 @requires_sm10x
-def test_flashinfer_trtllm_sparse_mla_decode_without_rope():
-    """The native sparse MLA path supports a zero-width rotary tail."""
+def test_flashinfer_trtllm_sparse_mla_decode_without_rope(dtype, mode, monkeypatch):
+    """TP skips preparation; empty DCP contributions and HiSparse outputs vanish."""
+    import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as backend
+    from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
+
     torch.set_default_device("cuda")
     torch.manual_seed(42)
-
-    batch_size = 2
-    block_size = 64
-    num_blocks = 4
-    sparse_topk = 128
-    valid_lens = torch.tensor([17, 73], dtype=torch.int32)
-
-    query = torch.randn(
-        batch_size,
-        1,
-        NUM_HEADS,
-        KV_LORA_RANK,
-        dtype=torch.bfloat16,
-    )
-    kv_cache = torch.randn(
-        num_blocks,
-        block_size,
-        KV_LORA_RANK,
-        dtype=torch.bfloat16,
-    )
-
+    batch_size, num_heads, block_size, num_blocks, sparse_topk = 3, 8, 64, 4, 2176
+    valid_lens = [1 if mode == "tp" else 0, 17, 73]
+    counts = torch.tensor(valid_lens, dtype=torch.int32)
+    query = torch.randn(batch_size, 1, num_heads, KV_LORA_RANK).to(dtype)
+    kv_cache = torch.randn(num_blocks, block_size, KV_LORA_RANK).to(dtype)
     num_slots = num_blocks * block_size
-    slot_tables = torch.stack(
-        [torch.randperm(num_slots)[:sparse_topk] for _ in range(batch_size)]
-    ).to(torch.int32)
-    for row, valid_len in zip(slot_tables, valid_lens.tolist()):
-        row[valid_len:] = -1
-
-    workspace_buffer = torch.empty(
-        FLASHINFER_WORKSPACE_BUFFER_SIZE,
-        dtype=torch.int8,
+    indices = torch.full((batch_size, sparse_topk), -1, dtype=torch.int32)
+    for token, length in enumerate(valid_lens):
+        indices[token, :length] = torch.randperm(num_slots)[:length].int()
+    empty = counts == 0
+    reference_indices = indices.clone()
+    reference_indices[:, 0].masked_fill_(empty, 0)
+    workspace = torch.empty(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.int8)
+
+    impl = object.__new__(backend.FlashInferMLASparseImpl)
+    impl.qk_nope_head_dim = QK_NOPE_HEAD_DIM
+    impl.kv_lora_rank = KV_LORA_RANK
+    impl.qk_rope_head_dim = 0
+    impl.index_group = (
+        object.__new__(HiSparseMLAIndexGroup) if mode == "hisparse" else None
     )
-    out = trtllm_batch_decode_with_kv_cache_mla(
+    impl._workspace_buffer = workspace
+    impl.bmm1_scale = QK_NOPE_HEAD_DIM**-0.5
+    impl.bmm2_scale = 1.0
+    impl.is_nope_mla = True
+    impl.need_to_return_lse_for_decode = mode == "dcp"
+    if mode == "tp":
+
+        def unexpected_helper(*args, **kwargs):
+            pytest.fail("TP must not launch empty-query preparation")
+
+        monkeypatch.setattr(
+            backend, "prepare_sparse_mla_safe_lengths", unexpected_helper
+        )
+
+    out, lse = impl._run_mqa_kernel(query.squeeze(1), kv_cache, indices, counts)
+    reference = trtllm_batch_decode_with_kv_cache_mla(
         query=query,
         kv_cache=kv_cache.unsqueeze(1),
-        workspace_buffer=workspace_buffer,
+        workspace_buffer=workspace,
         qk_nope_head_dim=QK_NOPE_HEAD_DIM,
         kv_lora_rank=KV_LORA_RANK,
         qk_rope_head_dim=0,
-        block_tables=slot_tables.unsqueeze(1),
-        seq_lens=valid_lens,
+        block_tables=reference_indices.unsqueeze(1),
+        seq_lens=counts,
         max_seq_len=sparse_topk,
         sparse_mla_top_k=sparse_topk,
-        sparse_mla_top_k_lens=valid_lens,
+        sparse_mla_top_k_lens=counts.clamp(min=1),
         bmm1_scale=QK_NOPE_HEAD_DIM**-0.5,
         bmm2_scale=1.0,
-    ).squeeze(1)
+        return_lse=mode == "dcp",
+    )
+    if mode == "dcp":
+        reference, reference_lse = reference
+        reference_lse = reference_lse.reshape(batch_size, num_heads)
+        reference_lse.masked_fill_(empty[:, None], float("-inf"))
+        torch.testing.assert_close(lse, reference_lse, atol=0, rtol=0)
+
+        from vllm.v1.attention.ops.dcp import CPTritonContext, correct_attn_out
+
+        # The peer owns only our empty selections. Check our contribution after
+        # the production merge, without zeroing the local attention output.
+        peer_lse = torch.full_like(lse, float("-inf")).masked_fill(empty[:, None], 0)
+        out, _ = correct_attn_out(
+            out, torch.stack((lse, peer_lse)), cp_rank=0, ctx=CPTritonContext()
+        )
+    else:
+        assert lse is None
+    reference = reference.squeeze(1).masked_fill(empty[:, None, None], 0)
+    torch.testing.assert_close(out, reference, atol=0, rtol=0)
+
+    if dtype == torch.bfloat16:
+        flat_cache = kv_cache.view(num_slots, KV_LORA_RANK).float()
+        refs = []
+        for token, length in enumerate(valid_lens):
+            selected_kv = flat_cache[reference_indices[token, :length].long()]
+            scores = torch.einsum("hd,kd->hk", query[token, 0].float(), selected_kv)
+            probs = torch.softmax(scores * QK_NOPE_HEAD_DIM**-0.5, dim=-1)
+            refs.append(torch.einsum("hk,kd->hd", probs, selected_kv))
+        ref = torch.stack(refs).to(torch.bfloat16)
+        torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2)
 
-    flat_cache = kv_cache.view(num_slots, KV_LORA_RANK).float()
-    refs = []
-    for batch_idx, valid_len in enumerate(valid_lens.tolist()):
-        selected_kv = flat_cache[slot_tables[batch_idx, :valid_len].long()]
-        scores = torch.einsum("hd,kd->hk", query[batch_idx, 0].float(), selected_kv)
-        probs = torch.softmax(scores * QK_NOPE_HEAD_DIM**-0.5, dim=-1)
-        refs.append(torch.einsum("hk,kd->hd", probs, selected_kv))
-    ref = torch.stack(refs).to(torch.bfloat16)
 
-    torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2)
+@requires_sm10x
+def test_sparse_mla_safe_lengths():
+    from vllm.v1.attention.backends.mla.sparse_utils import (
+        prepare_sparse_mla_safe_lengths,
+    )
+
+    device = "cuda"
+    num_tokens = 259
+    # Include a partial tile and expose accidental writes into adjacent storage.
+    counts_storage = torch.full((num_tokens, 2), 17, dtype=torch.int32, device=device)
+    counts = counts_storage[:, 0]
+    counts[::3] = 0
+    original_counts = counts_storage.clone()
+    empty = counts == 0
+    indices_storage = torch.full(
+        (num_tokens, 256), 42, dtype=torch.int32, device=device
+    )
+    indices = indices_storage[:, ::2]
+    expected_indices = indices_storage.clone()
+    expected_indices[:, 0].masked_fill_(empty, 0)
+    safe = prepare_sparse_mla_safe_lengths(indices, counts)
+    torch.testing.assert_close(safe, counts.clamp(min=1), rtol=0, atol=0)
+    torch.testing.assert_close(indices_storage, expected_indices, rtol=0, atol=0)
+    torch.testing.assert_close(counts_storage, original_counts, rtol=0, atol=0)
 
 
 @requires_sm90
@@ -286,7 +340,7 @@ def run(
 
 @requires_sm10x
 def test_flashinfer_mla_decode_workspace_supports_autotune():
-    """vLLM's FlashInfer MLA decode workspace must be int8 for autotuning.
+    """VLLM's FlashInfer MLA decode workspace must be int8 for autotuning.
 
     Model Runner V2's warmup autotunes ``trtllm_batch_decode_mla``, which makes
     the FlashInfer autotuner enumerate the CuteDSL tactic. That tactic asserts
diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py
index fb93557a32c0..7c1a455094d9 100644
--- a/tests/kernels/attention/test_flashmla_sparse.py
+++ b/tests/kernels/attention/test_flashmla_sparse.py
@@ -109,8 +109,9 @@ def test_sparse_flashmla_metadata_smoke():
     assert num_splits is None
 
 
-def test_sparse_flashmla_decode_smoke():
+def test_sparse_flashmla_decode_matches_cache_writer_scales():
     import vllm.v1.attention.ops.flashmla as fm
+    from vllm import _custom_ops as ops
 
     ok, reason = fm.is_flashmla_sparse_supported()
     if not ok:
@@ -130,7 +131,7 @@ def test_sparse_flashmla_decode_smoke():
     # Metadata
     q_seq_per_hk = seqlen_q * num_heads_q // num_heads_k
     # q_heads_per_hk = num_heads_q // num_heads_k
-    cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)
+    cache_seqlens = torch.ones(batch_size, dtype=torch.int32, device=device)
     tile_md, num_splits = fm.get_mla_metadata(
         cache_seqlens,
         q_seq_per_hk,
@@ -146,14 +147,33 @@ def test_sparse_flashmla_decode_smoke():
         dtype=torch.bfloat16,
         device=device,
     )
-    k_cache = torch.zeros(
-        (1, page_block_size, num_heads_k, bytes_per_token),
-        dtype=torch.uint8,
-        device=device,
+    cache_rows = torch.zeros(
+        (1, page_block_size, bytes_per_token), dtype=torch.uint8, device=device
+    )
+    # 336 / 448 = 0.75 distinguishes arbitrary FP32 from power-of-two scales.
+    kv_c = torch.full((1, 512), 336.0, dtype=torch.bfloat16, device=device)
+    kv_c[:, -128:] = 0
+    ops.concat_and_cache_mla(
+        kv_c,
+        torch.zeros((1, 64), dtype=torch.bfloat16, device=device),
+        cache_rows,
+        torch.zeros(1, dtype=torch.int64, device=device),
+        "fp8_ds_mla",
+        torch.ones(1, dtype=torch.float32, device=device),
     )
-    indices = torch.zeros(
-        (batch_size, seqlen_q, topk), dtype=torch.int32, device=device
+    scales = cache_rows[0, 0].view(torch.float32)[128:132]
+    expected_scales = torch.tensor([1.0, 1.0, 1.0, 2**-13], device=device)
+    torch.testing.assert_close(scales, expected_scales, rtol=0, atol=0)
+    cached_nope = cache_rows[0, 0, :512].view(torch.float8_e4m3fn).float()
+    expected = (cached_nope * scales.repeat_interleave(128)).to(torch.bfloat16)
+
+    k_cache = cache_rows.unsqueeze(2)
+    indices = torch.full(
+        (batch_size, seqlen_q, topk), -1, dtype=torch.int32, device=device
     )
+    indices[..., 0] = 0
+    # SM90 sparse decode only supports a fixed top-k width. Padding with -1
+    # exercises the same cache data on SM90 and SM100 without topk_length.
 
     block_table = torch.zeros((batch_size, 128), dtype=torch.int32, device=device)
     out, lse = fm.flash_mla_with_kvcache(
@@ -170,6 +190,9 @@ def test_sparse_flashmla_decode_smoke():
     assert out.shape[0] == batch_size
     assert out.shape[-1] == head_dim_v
     assert lse.shape[0] == batch_size
+    torch.testing.assert_close(
+        out, expected.view(1, 1, 1, 512).expand_as(out), rtol=0, atol=0
+    )
 
 
 @pytest.mark.parametrize("h_q", [64, 128])
@@ -486,3 +509,59 @@ def test_flashinfer_mixed_sparse_indices_separates_window_and_padded_width():
     assert sparse_indices.shape == (1, padded_width)
     assert sparse_indices[0].cpu().tolist() == [0, 1, 2, 3, -1, -1, -1, -1]
     assert sparse_lens.cpu().tolist() == [padded_width]
+
+
+@pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"])
+@pytest.mark.parametrize("page_padding", [0, 64])
+def test_flashinfer_sparse_forward_reads_packed_kv_rows(model, page_padding):
+    """Decode and prefill must read storage rows after vLLM remaps packed pages."""
+    import importlib
+
+    from vllm.platforms import current_platform
+    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
+
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("Requires FlashInfer TRTLLM sparse MLA on SM100/SM103")
+    mod = importlib.import_module(f"vllm.models.{model}.nvidia.flashinfer_sparse")
+    device = "cuda"
+    block_size = 64
+    storage = torch.full(
+        (3, block_size + page_padding, 512), 99, dtype=torch.bfloat16, device=device
+    )
+    cache = storage[:, :block_size]
+    cache[1].fill_(2)
+    cache[2].fill_(3)
+    attn = object.__new__(mod.DeepseekV4FlashInferMLAAttention)
+    attn.compress_ratio = 1
+    attn.window_size = 128
+    attn.scale = 512**-0.5
+    attn.kv_cache_torch_dtype = torch.bfloat16
+    attn.attn_sink = torch.full((64,), -float("inf"), device=device)
+    attn.topk_indices_buffer = torch.empty((3, 0), dtype=torch.int32, device=device)
+    decode_indices = torch.full((1, 128), -1, dtype=torch.int32, device=device)
+    decode_indices[0, :4] = torch.arange(64, 68, dtype=torch.int32, device=device)
+    query_start = torch.tensor([0, 1, 3], dtype=torch.int32)
+    metadata = DeepseekSparseSWAMetadata(
+        block_table=torch.tensor([[1], [2]], dtype=torch.int32, device=device),
+        slot_mapping=torch.tensor([67, 130, 131], device=device),
+        block_size=block_size,
+        seq_lens=torch.tensor([4, 4], dtype=torch.int32, device=device),
+        query_start_loc=query_start.to(device),
+        query_start_loc_cpu=query_start,
+        token_to_req_indices=torch.tensor([0, 1, 1], dtype=torch.int32, device=device),
+        decode_swa_indices=decode_indices,
+        decode_swa_lens=torch.tensor([4], dtype=torch.int32, device=device),
+        decode_swa_width=128,
+        is_valid_token=torch.ones(3, dtype=torch.bool, device=device),
+        replay_start=torch.zeros(2, dtype=torch.int32, device=device),
+        num_decodes=1,
+        num_prefills=1,
+        num_decode_tokens=1,
+        num_prefill_tokens=2,
+        max_decode_query_len=1,
+    )
+    query = torch.zeros((3, 64, 512), dtype=torch.bfloat16, device=device)
+    output = torch.empty_like(query)
+    attn._forward(query, None, cache, metadata, None, True, output)
+    expected = torch.tensor([2, 3, 3], dtype=output.dtype, device=device)
+    torch.testing.assert_close(output, expected[:, None, None].expand_as(output))
diff --git a/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py b/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py
index e1582b7621a4..af41722605c6 100644
--- a/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py
+++ b/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py
@@ -222,7 +222,8 @@ def test_decode_concat_ignores_negative_slots_for_cache(cache_format: str) -> No
     _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache)
 
 
-def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None:
+@pytest.mark.parametrize("decode", [False, True])
+def test_ds_mla_cache_insert_bit_compatible_with_reference(decode: bool) -> None:
     """Fused ds_mla insertion must match the reference bit-for-bit."""
     torch.manual_seed(0)
     num_tokens, num_blocks = 33, 16
@@ -242,5 +243,18 @@ def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None:
     got = ref.clone()
     scale = torch.ones(1, device="cuda", dtype=torch.float32)
     ops.concat_and_cache_mla(kv_c, k_pe.squeeze(1), ref, slots, "fp8_ds_mla", scale)
-    fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots)
+    if decode:
+        ql_nope = torch.randn(
+            num_tokens, NUM_HEADS, KV_LORA_RANK, device="cuda", dtype=dt
+        )
+        fused_mla_decode_q_concat_kv_cache_insert(
+            ql_nope, q[..., -ROPE_HEAD_DIM:], kv_c, k_pe, got, slots, ds_mla=True
+        )
+    else:
+        fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots)
     assert torch.equal(ref, got)
+    rows = got.view(-1, DS_MLA_CACHE_ENTRY)[slots[slots >= 0]]
+    scales = rows[:, KV_LORA_RANK : KV_LORA_RANK + 16].view(torch.float32)
+    torch.testing.assert_close(
+        torch.log2(scales), torch.log2(scales).round(), rtol=0, atol=0
+    )
diff --git a/tests/kernels/attention/test_lightning_attn.py b/tests/kernels/attention/test_lightning_attn.py
index 61e13166808c..d4b1a359d28e 100644
--- a/tests/kernels/attention/test_lightning_attn.py
+++ b/tests/kernels/attention/test_lightning_attn.py
@@ -23,7 +23,7 @@
 
 
 def reference_lightning_attention(q, k, v, ed, block_size, kv_history):
-    """Reference implementation of lightning attention core algorithm
+    """Reference implementation of lightning attention core algorithm.
 
     The difference from the main implementation is that this processes
     each step sequentially, instead of using parallelized triton kernels
@@ -72,7 +72,7 @@ def reference_lightning_attention(q, k, v, ed, block_size, kv_history):
 
 
 def reference_linear_decode(q, k, v, kv_caches, slope_rate, slot_idx):
-    """Reference implementation: linear attention decode function"""
+    """Reference implementation: linear attention decode function."""
     B, H, _, D = q.shape
     output = torch.zeros(B, H * D, dtype=q.dtype, device=q.device)
 
diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py
index b6578fa2f915..8fa4963a6a19 100644
--- a/tests/kernels/attention/test_mha_attn.py
+++ b/tests/kernels/attention/test_mha_attn.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test:
+"""Test:
 
 * Tests for MMEncoderAttention layer
 """
@@ -49,9 +48,7 @@ def reset_test_state():
 
 @pytest.mark.parametrize("device", devices)
 def test_mha_attn_platform(default_vllm_config, device: str):
-    """
-    Test the attention selector between different platform and device.
-    """
+    """Test the attention selector between different platform and device."""
     torch.set_default_dtype(torch.float16)
 
     if device == "cpu":
@@ -113,8 +110,7 @@ def ref_attention(
     value: torch.Tensor,
     scale: float,
 ) -> torch.Tensor:
-    """
-    Native implementation of scaled dot product attention without mask:
+    """Native implementation of scaled dot product attention without mask:
     - query, key, value: [batch_size, seq_len, num_heads, head_size]
     - attn_mask: [batch_size, seq_len, seq_len]
     """
diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py
index c45b87d04b1f..9d3b1500e1d2 100644
--- a/tests/kernels/attention/test_minimax_m3.py
+++ b/tests/kernels/attention/test_minimax_m3.py
@@ -19,7 +19,10 @@
     minimax_m3_sparse_attn,
     minimax_m3_sparse_attn_decode,
 )
-from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseTritonImpl
+from vllm.models.minimax_m3.common.sparse_attention import (
+    MiniMaxM3SparseTritonImpl,
+    minimax_m3_rebase_slots_to_page16,
+)
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.attention.backends.utils import record_kv_cache_layout
@@ -28,6 +31,7 @@
     KVCacheLayout,
     MLAAttentionSpec,
     compute_layer_kv_cache_shape_bytes,
+    get_kv_quant_mode,
 )
 
 if not (current_platform.is_cuda() or current_platform.is_rocm()):
@@ -1025,17 +1029,20 @@ def test_decode_index_topk_correctness(
         "expected",
     ),
     [
-        (0, 24, 1, 16, True, (2, False, False)),
-        (1, 1, 1, 16, True, (16, True, True)),
-        (512, 4, 1, 16, True, (16, True, True)),
-        (513, 24, 4, 16, True, (16, True, True)),
-        (2760, 5, 2, 16, True, (16, True, True)),
-        (8192, 64, 4, 16, True, (16, True, True)),
-        (8193, 4, 1, 16, True, (16, False, False)),
-        (8193, 16, 1, 16, True, (4, False, False)),
-        (65534, 64, 1, 16, True, (1, False, False)),
-        (513, 16, 1, 13, True, (4, False, False)),
-        (513, 16, 4, 16, False, (1, False, False)),
+        (0, 24, 1, 16, True, (2, 512, 8, 2, False, False)),
+        (1, 1, 1, 16, True, (1, 128, 2, 1, True, True)),
+        (74, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
+        (128, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
+        (129, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
+        (512, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
+        (513, 24, 4, 16, True, (16, 512, 4, 2, True, True)),
+        (2760, 5, 2, 16, True, (16, 512, 4, 2, True, True)),
+        (8192, 64, 4, 16, True, (16, 512, 4, 2, True, True)),
+        (8193, 4, 1, 16, True, (16, 512, 8, 2, False, False)),
+        (8193, 16, 1, 16, True, (4, 512, 8, 2, False, False)),
+        (65534, 64, 1, 16, True, (1, 512, 8, 2, False, False)),
+        (513, 16, 1, 13, True, (4, 512, 8, 2, False, False)),
+        (513, 16, 4, 16, False, (1, 512, 8, 2, False, False)),
     ],
 )
 def test_amd_decode_topk_launch_policy(
@@ -1044,9 +1051,9 @@ def test_amd_decode_topk_launch_policy(
     num_idx_heads: int,
     topk: int,
     is_gfx950: bool,
-    expected: tuple[int, bool, bool],
+    expected: tuple[int, int, int, int, bool, bool],
 ):
-    """All supported head and query layouts use adaptive selection."""
+    """The short-context specialization has strict dispatch boundaries."""
     from vllm.models.minimax_m3.amd.ops.index_topk import (
         _decode_topk_launch_policy,
     )
@@ -1073,19 +1080,23 @@ def test_amd_decode_topk_launch_policy(
         "single_tile_guaranteed",
         "adaptive_final_merge",
         "max_blocks",
+        "block_size_k",
+        "num_warps",
+        "num_stages",
         "topk",
         "num_idx_heads",
     ),
     [
-        (1, False, False, 65534, 16, 1),
-        (2, False, False, 8193, 16, 1),
-        (4, False, False, 8193, 16, 1),
-        (8, False, False, 8193, 16, 1),
-        (16, False, False, 8193, 16, 1),
-        (16, False, False, 65534, 16, 1),
-        (2, True, False, 513, 13, 1),
-        (4, False, True, 8192, 16, 2),
-        (16, True, True, 8192, 16, 4),
+        (1, False, False, 65534, 512, 8, 2, 16, 1),
+        (2, False, False, 8193, 512, 8, 2, 16, 1),
+        (4, False, False, 8193, 512, 8, 2, 16, 1),
+        (8, False, False, 8193, 512, 8, 2, 16, 1),
+        (16, False, False, 8193, 512, 8, 2, 16, 1),
+        (16, False, False, 65534, 512, 8, 2, 16, 1),
+        (2, True, False, 513, 512, 8, 2, 13, 1),
+        (4, False, True, 8192, 512, 4, 2, 16, 2),
+        (16, True, True, 8192, 512, 4, 2, 16, 4),
+        (1, True, True, 74, 128, 2, 1, 16, 1),
     ],
 )
 def test_amd_decode_fused_topk_total_order_and_replay(
@@ -1093,10 +1104,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
     single_tile_guaranteed: bool,
     adaptive_final_merge: bool,
     max_blocks: int,
+    block_size_k: int,
+    num_warps: int,
+    num_stages: int,
     topk: int,
     num_idx_heads: int,
 ):
-    """Atomic multi-chunk selection is ordered and resets graph state."""
+    """Selector specializations preserve ordering and graph state."""
     from vllm.models.minimax_m3.amd.ops.index_topk import (
         SPARSE_BLOCK_SIZE,
         _decode_topk_fused_kernel,
@@ -1104,18 +1118,25 @@ def test_amd_decode_fused_topk_total_order_and_replay(
 
     set_random_seed(0)
     block_size_t = 1 << (topk - 1).bit_length()
+    block_counts = tuple(
+        count
+        for count in (
+            15,
+            16,
+            17,
+            block_size_k // 2 - 1,
+            block_size_k // 2,
+            block_size_k // 2 + 1,
+            block_size_k - 1,
+            block_size_k,
+            block_size_k + 1,
+        )
+        if count <= max_blocks
+    )
     long_seq_lens_list = [
         0,
         1,
-        15 * SPARSE_BLOCK_SIZE,
-        16 * SPARSE_BLOCK_SIZE,
-        17 * SPARSE_BLOCK_SIZE,
-        255 * SPARSE_BLOCK_SIZE,
-        256 * SPARSE_BLOCK_SIZE,
-        257 * SPARSE_BLOCK_SIZE,
-        511 * SPARSE_BLOCK_SIZE,
-        512 * SPARSE_BLOCK_SIZE,
-        513 * SPARSE_BLOCK_SIZE,
+        *(blocks * SPARSE_BLOCK_SIZE for blocks in block_counts),
         max_blocks * SPARSE_BLOCK_SIZE,
     ]
     short_seq_lens_list = [0, 1, 127, 128, 129, 0, 1, 128, 129, 0, 1, 129]
@@ -1126,12 +1147,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
         boundary_limit = 4
     if boundary_limit:
         boundary_blocks = [
-            (active_chunks - 1) * 512 + 1
+            (active_chunks - 1) * block_size_k + 1
             for active_chunks in range(3, boundary_limit + 1)
         ]
-        assert [(num_blocks + 511) // 512 for num_blocks in boundary_blocks] == list(
-            range(3, boundary_limit + 1)
-        )
+        assert [
+            (num_blocks + block_size_k - 1) // block_size_k
+            for num_blocks in boundary_blocks
+        ] == list(range(3, boundary_limit + 1))
         long_seq_lens_list.extend(
             num_blocks * SPARSE_BLOCK_SIZE for num_blocks in boundary_blocks
         )
@@ -1139,6 +1161,8 @@ def test_amd_decode_fused_topk_total_order_and_replay(
             num_blocks * SPARSE_BLOCK_SIZE for num_blocks in reversed(boundary_blocks)
         )
     batch = len(long_seq_lens_list)
+    repeats = (batch + len(short_seq_lens_list) - 1) // len(short_seq_lens_list)
+    short_seq_lens_list = (short_seq_lens_list * repeats)[:batch]
     score = torch.randn((num_idx_heads, batch, max_blocks), device="cuda")
     score[:, :, 0] = float("nan")
     score[:, :, 1] = float("inf")
@@ -1147,8 +1171,9 @@ def test_amd_decode_fused_topk_total_order_and_replay(
     score[:, :, 4] = -0.0
     score[:, :, 5] = 7.0
     score[:, :, 6] = -1e30
-    score[:, :, 256] = 7.0
-    score[:, :, 512] = 7.0
+    for index in (block_size_k // 2, block_size_k):
+        if index < max_blocks:
+            score[:, :, index] = 7.0
 
     seq_lens = torch.empty(batch, device="cuda", dtype=torch.int32)
     partial = torch.full(
@@ -1198,13 +1223,13 @@ def launch() -> None:
             pages_per_sparse_block=8,
             block_page_stride=8,
             NUM_TOPK_CHUNKS=num_topk_chunks,
-            BLOCK_SIZE_K=512,
+            BLOCK_SIZE_K=block_size_k,
             BLOCK_SIZE_T=block_size_t,
             EMIT_SPARSE_TABLE=False,
             SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
             ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
-            num_warps=4 if adaptive_final_merge else 8,
-            num_stages=2,
+            num_warps=num_warps,
+            num_stages=num_stages,
         )
 
     score_cpu = score.cpu()
@@ -1251,6 +1276,7 @@ def reference(seq_lens_list: list[int]) -> torch.Tensor:
         (16, False, False, 8193),
         (4, False, True, 8192),
         (16, True, True, 8192),
+        (1, True, True, 74),
     ):
         graph = torch.cuda.CUDAGraph()
         torch.accelerator.synchronize()
@@ -1650,8 +1676,9 @@ def test_main_cache_layout_contract():
         assert set(order) == set(range(len(order)))
 
 
-def test_aiter_sparse_pa_layout_contract(monkeypatch):
-    """The shuffle-only AITER path retains separately contiguous K/V storage."""
+@pytest.mark.parametrize("layout", [KVCacheLayout.LHBNC, KVCacheLayout.LBHNC])
+def test_aiter_sparse_pa_layout_contract(monkeypatch, layout):
+    """Build the exact separate and packed page-16 views used by the model."""
     import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
 
     monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
@@ -1663,21 +1690,43 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch):
 
     backend = sparse_attn_mod.MiniMaxM3SparseBackend
     assert KVCacheLayout.LHBNC in backend.supported_kv_cache_layouts()
+    assert KVCacheLayout.LBHNC in backend.supported_kv_cache_layouts()
 
-    nb, h = 7, 1
+    nb, h = 3, 1
+    cache_dtype = current_platform.fp8_dtype()
     spec = FullAttentionSpec(
         block_size=BLOCK_SIZE,
         num_kv_heads=h,
         head_size=HEAD_DIM,
-        dtype=DTYPE,
+        dtype=cache_dtype,
+        kv_quant_mode=get_kv_quant_mode("fp8"),
         num_head_slots=2,
-        state_content_bytes=h * HEAD_DIM * DTYPE.itemsize,
+        state_content_bytes=h * HEAD_DIM * cache_dtype.itemsize,
     )
     raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
-    view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LHBNC)[0]
-    key_cache, value_cache = view.unbind(1)
-    assert key_cache.is_contiguous()
-    assert value_cache.is_contiguous()
+    kv_cache = dense_kv_cache_views(raw, spec, nb, 1, layout)[0].view(cache_dtype)
+    key_cache, value_cache = kv_cache.unbind(1)
+    x = 16 // cache_dtype.itemsize
+    pages_per_side = BLOCK_SIZE // 16
+    if key_cache.is_contiguous() and value_cache.is_contiguous():
+        k_src, v_src = key_cache, value_cache
+        block_pages, v_page_offset = pages_per_side, 0
+    else:
+        assert kv_cache.is_contiguous()
+        k_src = v_src = kv_cache
+        block_pages, v_page_offset = 2 * pages_per_side, pages_per_side
+
+    num_phys16 = nb * block_pages
+    k_view = k_src.view(num_phys16, h, HEAD_DIM // x, 16, x)
+    v_view = v_src.view(num_phys16, h, 16 // x, HEAD_DIM, x)[v_page_offset:]
+    assert k_view.is_contiguous()
+    assert v_view.is_contiguous()
+    if layout is KVCacheLayout.LHBNC:
+        assert k_view.shape == (24, 1, 8, 16, 16)
+        assert v_view.shape == (24, 1, 1, 128, 16)
+    else:
+        assert k_view.shape == (48, 1, 8, 16, 16)
+        assert v_view.shape == (40, 1, 1, 128, 16)
 
 
 def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
@@ -1747,6 +1796,330 @@ def test_aiter_indexer_requires_the_aiter_attend(monkeypatch):
     assert reason is None or "AITER sparse PA attend" not in reason
 
 
+def test_aiter_consolidated_qknorm_enabled(monkeypatch):
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
+    monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
+
+    scale = torch.tensor([0.25])
+    assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
+    assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8", scale, scale)
+    assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e4m3", scale, scale)
+    assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8")
+    assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8", scale, None)
+    assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e5m2", scale, scale)
+
+    monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", False)
+    assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
+
+    monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
+    monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: False)
+    assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
+
+
+def _fake_aiter_qknorm_args(kv_cache_dtype: str = "auto") -> dict:
+    tensor = torch.empty(1)
+    return {
+        "qkv": tensor,
+        "q_norm_weight": tensor,
+        "k_norm_weight": tensor,
+        "cos_sin_cache": tensor,
+        "positions": tensor,
+        "num_heads": 16,
+        "num_kv_heads": 1,
+        "rotary_dim": 64,
+        "eps": 1e-6,
+        "slot_mapping": tensor,
+        "kv_cache_k": tensor,
+        "kv_cache_v": tensor,
+        "q_out": tensor,
+        "kv_cache_dtype": kv_cache_dtype,
+    }
+
+
+def _install_fake_aiter_qknorm(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
+    """Stub ``aiter.fused_qknorm_idxrqknorm`` so mapping tests run without AITER."""
+    import sys
+    import types
+    from typing import Any
+
+    calls: list[dict] = []
+
+    def fake_op(*args, **kwargs):
+        calls.append(kwargs)
+
+    existing = sys.modules.get("aiter")
+    if existing is not None:
+        monkeypatch.setattr(existing, "fused_qknorm_idxrqknorm", fake_op, raising=False)
+    else:
+        fake_mod: Any = types.ModuleType("aiter")
+        fake_mod.fused_qknorm_idxrqknorm = fake_op
+        monkeypatch.setitem(sys.modules, "aiter", fake_mod)
+    return calls
+
+
+def test_aiter_consolidated_qknorm_packed_shuffle_unequal_spans(monkeypatch):
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    calls = _install_fake_aiter_qknorm(monkeypatch)
+    call_args = _fake_aiter_qknorm_args()
+    call_args["kv_cache_k"] = torch.empty(2, 1)
+    call_args["kv_cache_v"] = torch.empty(1, 1)
+
+    rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
+    assert len(calls) == 1
+
+
+@pytest.mark.parametrize(
+    ("kv_cache_dtype", "aiter_dtype", "passes_scales"),
+    [
+        ("auto", "auto", False),
+        ("fp8", "fp8_e4m3_static", True),
+        ("fp8_e4m3", "fp8_e4m3_static", True),
+    ],
+)
+def test_aiter_consolidated_qknorm_dtype_and_layout_mapping(
+    monkeypatch, kv_cache_dtype, aiter_dtype, passes_scales
+):
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    calls = _install_fake_aiter_qknorm(monkeypatch)
+    call_args = _fake_aiter_qknorm_args(kv_cache_dtype)
+    k_scale = torch.tensor([0.25])
+    v_scale = torch.tensor([0.5])
+    call_args.update(k_scale=k_scale, v_scale=v_scale)
+
+    rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
+    assert len(calls) == 1
+    kwargs = calls[0]
+    assert kwargs["kv_cache_dtype"] == aiter_dtype
+    assert kwargs["index_cache_dtype"] == "auto"
+    assert kwargs["block_size"] == 16
+    assert kwargs["asm_layout"] is True
+    if passes_scales:
+        assert kwargs["k_scale"] is k_scale
+        assert kwargs["v_scale"] is v_scale
+    else:
+        assert kwargs["k_scale"] is None
+        assert kwargs["v_scale"] is None
+
+
+def test_aiter_consolidated_qknorm_fp8_index_dtype_mapping(monkeypatch):
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    calls = _install_fake_aiter_qknorm(monkeypatch)
+    fp8 = current_platform.fp8_dtype()
+    call_args = _fake_aiter_qknorm_args("fp8")
+    call_args.update(
+        k_scale=torch.tensor([0.25]),
+        v_scale=torch.tensor([0.5]),
+        index_cache=torch.empty(1, dtype=fp8),
+        index_q_out=torch.empty(1, dtype=fp8),
+        num_index_heads=1,
+    )
+
+    rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
+    assert calls[0]["index_cache_dtype"] == "fp8"
+    assert calls[0]["kv_cache_dtype"] == "fp8_e4m3_static"
+
+
+@pytest.mark.parametrize("skip_index_branch", [False, True])
+def test_aiter_consolidated_qknorm_full_and_skip_index_args(
+    monkeypatch, skip_index_branch
+):
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    calls = _install_fake_aiter_qknorm(monkeypatch)
+    call_args = _fake_aiter_qknorm_args()
+    if skip_index_branch:
+        call_args.update(skip_index_branch=True)
+    else:
+        tensor = torch.empty(1)
+        call_args.update(
+            index_q_norm_weight=tensor,
+            index_k_norm_weight=tensor,
+            num_index_heads=1,
+            index_cache=tensor,
+            index_q_out=tensor,
+            index_slot_mapping=tensor,
+        )
+
+    rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
+    kwargs = calls[0]
+    assert kwargs["skip_index_branch"] is skip_index_branch
+    if skip_index_branch:
+        assert kwargs["index_q_norm_weight"] is None
+        assert kwargs["index_k_norm_weight"] is None
+        assert kwargs["index_cache"] is None
+        assert kwargs["index_q_out"] is None
+        assert kwargs["index_slot_mapping"] is None
+    else:
+        assert kwargs["index_q_norm_weight"] is not None
+        assert kwargs["index_k_norm_weight"] is not None
+        assert kwargs["index_cache"] is not None
+        assert kwargs["index_q_out"] is not None
+        assert kwargs["index_slot_mapping"] is not None
+
+
+def test_aiter_consolidated_qknorm_rejects_fp8_e5m2(monkeypatch):
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
+    monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
+    scale = torch.tensor([0.25])
+    assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e5m2", scale, scale)
+
+
+def test_aiter_sparse_pa_slot_rebase_preserves_page_offsets():
+    slots = torch.tensor([0, 15, 16, 127, 128, 129, -1])
+    actual = minimax_m3_rebase_slots_to_page16(slots, block_size=128)
+    expected = torch.tensor([0, 15, 16, 127, 256, 257, -1])
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="The consolidated packed-SHUFFLE AITER op is ROCm-only",
+)
+@pytest.mark.parametrize("skip_index_branch", [False, True])
+@pytest.mark.parametrize("fp8_index", [False, True])
+def test_aiter_consolidated_qknorm_real_packed_shuffle(skip_index_branch, fp8_index):
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    torch.manual_seed(20260831)
+    num_tokens = 4
+    num_heads = 16
+    num_kv_heads = 1
+    num_index_heads = 1
+    rotary_dim = 64
+    head_dim = 128
+    logical_block_size = 128
+    num_blocks = 3
+    fp8_dtype = current_platform.fp8_dtype()
+
+    qkv = torch.randn(
+        num_tokens,
+        (num_heads + 2 * num_kv_heads + num_index_heads + 1) * head_dim,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    weights = [
+        torch.randn(head_dim, dtype=torch.bfloat16, device="cuda") for _ in range(4)
+    ]
+    cos_sin_cache = torch.randn(16, rotary_dim, dtype=torch.bfloat16, device="cuda")
+    positions = torch.arange(num_tokens, dtype=torch.int64, device="cuda")
+    logical_slots = torch.tensor([0, 127, 128, -1], dtype=torch.int64, device="cuda")
+    slot_mapping = minimax_m3_rebase_slots_to_page16(logical_slots, logical_block_size)
+    index_slot_mapping = logical_slots.clone()
+
+    packed = torch.zeros(
+        num_blocks,
+        2,
+        logical_block_size,
+        head_dim,
+        dtype=fp8_dtype,
+        device="cuda",
+    )
+    x = 16 // packed.element_size()
+    pages_per_side = logical_block_size // 16
+    num_phys16 = num_blocks * 2 * pages_per_side
+    kv_cache_k = packed.view(num_phys16, num_kv_heads, head_dim // x, 16, x)
+    kv_cache_v = packed.view(num_phys16, num_kv_heads, 16 // x, head_dim, x)[
+        pages_per_side:
+    ]
+    assert kv_cache_k.shape == (48, 1, 8, 16, 16)
+    assert kv_cache_v.shape == (40, 1, 1, 128, 16)
+
+    q_out = torch.zeros(
+        num_tokens, num_heads * head_dim, dtype=torch.bfloat16, device="cuda"
+    )
+    index_dtype = fp8_dtype if fp8_index else torch.bfloat16
+    if fp8_index:
+        index_q_out = torch.empty(
+            (num_tokens, num_index_heads * head_dim),
+            dtype=index_dtype,
+            device="cuda",
+        )
+        index_cache = torch.empty(
+            (num_blocks, logical_block_size, head_dim),
+            dtype=index_dtype,
+            device="cuda",
+        )
+        index_q_out.view(torch.uint8).fill_(7)
+        index_cache.view(torch.uint8).fill_(7)
+    else:
+        index_q_out = torch.full(
+            (num_tokens, num_index_heads * head_dim),
+            7,
+            dtype=torch.bfloat16,
+            device="cuda",
+        )
+        index_cache = torch.full(
+            (num_blocks, logical_block_size, head_dim),
+            7,
+            dtype=torch.bfloat16,
+            device="cuda",
+        )
+    scale = torch.tensor([0.02], dtype=torch.float32, device="cuda")
+
+    kwargs = {}
+    if not skip_index_branch:
+        kwargs = {
+            "index_q_norm_weight": weights[2],
+            "index_k_norm_weight": weights[3],
+            "index_cache": index_cache,
+            "index_q_out": index_q_out,
+            "index_slot_mapping": index_slot_mapping,
+        }
+    rocm_aiter_ops.fused_qknorm_idxrqknorm(
+        qkv,
+        weights[0],
+        weights[1],
+        cos_sin_cache,
+        positions,
+        num_heads,
+        num_kv_heads,
+        rotary_dim,
+        1e-6,
+        slot_mapping,
+        kv_cache_k,
+        kv_cache_v,
+        q_out,
+        "fp8",
+        scale,
+        scale,
+        num_index_heads=num_index_heads,
+        skip_index_branch=skip_index_branch,
+        **kwargs,
+    )
+    torch.accelerator.synchronize()
+
+    assert torch.isfinite(q_out).all()
+    assert torch.count_nonzero(q_out) > 0
+    assert torch.count_nonzero(packed) > 0
+    if fp8_index:
+        index_q_untouched = torch.all(index_q_out.view(torch.uint8) == 7)
+        index_cache_untouched = torch.all(index_cache.view(torch.uint8) == 7)
+        index_cache_t0 = torch.all(index_cache[0, 0].view(torch.uint8) == 7)
+        index_cache_t127 = torch.all(index_cache[0, 127].view(torch.uint8) == 7)
+        index_cache_t128 = torch.all(index_cache[1, 0].view(torch.uint8) == 7)
+    else:
+        index_q_untouched = torch.all(index_q_out == 7)
+        index_cache_untouched = torch.all(index_cache == 7)
+        index_cache_t0 = torch.all(index_cache[0, 0] == 7)
+        index_cache_t127 = torch.all(index_cache[0, 127] == 7)
+        index_cache_t128 = torch.all(index_cache[1, 0] == 7)
+    if skip_index_branch:
+        assert index_q_untouched
+        assert index_cache_untouched
+    else:
+        assert not index_q_untouched
+        assert not index_cache_t0
+        assert not index_cache_t127
+        assert not index_cache_t128
+
+
 def test_indexer_cache_squeezes_to_contiguous_3d():
     """The indexer side cache is standardized 4D with H=1: under both layouts
     the allocator's logical view stays contiguous and squeezes (as
diff --git a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
index ac4607355999..ffdb76e176f8 100644
--- a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
+++ b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
@@ -217,7 +217,8 @@ def run(kv):
     assert (out_ref - out_view).abs().max().item() == 0.0
 
 
-def test_flashmla_fp8_sparse_decode_unified_slot_view():
+@pytest.mark.parametrize("next_n", [1, 2, 4])
+def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int):
     """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
     unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
     cache, with finite nonzero output."""
@@ -240,7 +241,10 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view():
     n_layers = 3
     layer = 1
 
-    q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
+    q = (
+        torch.randn(batch, next_n, h_q, head_dim, device=dev, dtype=torch.bfloat16)
+        * 0.1
+    )
 
     # Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B
     # bf16 rope (random bytes corrupt the scale region and yield NaNs).
@@ -273,13 +277,13 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view():
 
     # Sparse indices: each batch uses its own disjoint blocks.
     blocks_per_batch = num_blocks // batch
-    idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32)
+    idx = torch.full((batch, next_n, topk), -1, device=dev, dtype=torch.int32)
     for b in range(batch):
         slots: list[int] = []
         for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch):
             slots.extend(blk * page + off for off in range(page))
         slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32)
-        idx[b, 0, : slots_t.numel()] = slots_t
+        idx[b, :, : slots_t.numel()] = slots_t
 
     def run(kc):
         meta, num_splits = fm.get_mla_metadata()
diff --git a/tests/kernels/attention/test_rocm_aiter_candidate_mask.py b/tests/kernels/attention/test_rocm_aiter_candidate_mask.py
new file mode 100644
index 000000000000..85e742a09636
--- /dev/null
+++ b/tests/kernels/attention/test_rocm_aiter_candidate_mask.py
@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""ROCm correctness tests for the strided DSA candidate mask.
+
+``_apply_candidate_mask_strided`` is a decode-path variant of the shared
+``apply_candidate_mask``. It differs in two ways that need pinning down:
+
+* it launches a fixed number of programs that stride the column axis, rather
+  than one program per 1024-column tile, so at ``max_model_len`` widths the
+  grid no longer scales with a workspace that is mostly padding;
+* it stops at each row's ``end`` instead of sanitizing the whole width.
+
+The second is only sound because every consumer bounds its scan by the same
+ends -- ``top_k_per_row_decode`` takes ``rowEnd`` from ``seq_lens``. So the
+contract is agreement with the shared kernel over ``[0, end)`` and nothing
+beyond it, which is what these tests assert rather than comparing whole rows.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-specific tests"
+)
+
+BLOCK_SIZE = 8
+TOPK_BLOCKS = 64
+# Above 128 tiles the grid saturates and programs start striding; below it the
+# min() clamp applies instead. Both regimes need covering, so the widths here
+# deliberately straddle 128 * 1024.
+NARROW = 8192
+WIDE = 262144
+
+
+def _inputs(rows, width, ends, starts, *, block_size, col_stride, dtype):
+    """Logits, bounds and candidate blocks shared by both kernels."""
+    torch.manual_seed(0)
+    # A column stride > 1 keeps the kernels honest about stride_col; the decode
+    # caller slices a workspace, so contiguous rows are not guaranteed.
+    logits = torch.randn(rows, width * col_stride, device="cuda")[:, ::col_stride]
+    logits[0, : min(16, width)] = 0.0
+    if rows > 2:
+        logits[2, 5] = float("nan")
+    row_ke = torch.tensor(ends, device="cuda", dtype=dtype)
+    row_ks = torch.tensor(starts, device="cuda", dtype=dtype) if starts else None
+    nblocks = (width + block_size - 1) // block_size
+    candidates = torch.randint(
+        0, nblocks, (rows, TOPK_BLOCKS), device="cuda", dtype=torch.int32
+    )
+    # -1 is the "no block" sentinel the selector emits for short rows.
+    candidates[:, ::7] = -1
+    return logits, row_ks, row_ke, candidates
+
+
+def _run_both(logits, row_ks, row_ke, candidates, block_size, row_repeat):
+    from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+        apply_candidate_mask,
+    )
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        _apply_candidate_mask_strided,
+    )
+
+    ref = logits.clone()
+    got = logits.clone()
+    apply_candidate_mask(ref, row_ks, row_ke, candidates, block_size, row_repeat)
+    _apply_candidate_mask_strided(
+        got, row_ks, row_ke, candidates, block_size, row_repeat
+    )
+    return ref, got
+
+
+def _assert_agrees_within_ends(ref, got, row_ke, row_repeat):
+    """Compare only the region a consumer can read."""
+    for row in range(ref.shape[0]):
+        end = int(row_ke[row // row_repeat])
+        if end <= 0:
+            continue
+        torch.testing.assert_close(
+            got[row, :end],
+            ref[row, :end],
+            rtol=0,
+            atol=0,
+            equal_nan=True,
+            msg=lambda m, r=row, e=end: f"row {r} differs within [0, {e}): {m}",
+        )
+
+
+@pytest.mark.parametrize("width", [NARROW, WIDE])
+@pytest.mark.parametrize("col_stride", [1, 2])
+def test_matches_shared_kernel_ragged_ends(width, col_stride):
+    """Ends on and around tile boundaries, plus the empty and full rows."""
+    ends = [0, 1, 1023, 1024, 1025, width - 1, width, width // 2]
+    logits, row_ks, row_ke, candidates = _inputs(
+        len(ends),
+        width,
+        ends,
+        None,
+        block_size=BLOCK_SIZE,
+        col_stride=col_stride,
+        dtype=torch.int32,
+    )
+    ref, got = _run_both(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1)
+    _assert_agrees_within_ends(ref, got, row_ke, 1)
+
+
+@pytest.mark.parametrize("block_size", [8, 16])
+def test_matches_shared_kernel_with_starts(block_size):
+    """Non-zero starts: the decode caller passes zeros, prefill does not."""
+    ends = [0, 64, 4096, 8192]
+    starts = [0, 7, 1024, 3]
+    logits, row_ks, row_ke, candidates = _inputs(
+        len(ends),
+        NARROW,
+        ends,
+        starts,
+        block_size=block_size,
+        col_stride=1,
+        dtype=torch.int64,
+    )
+    ref, got = _run_both(logits, row_ks, row_ke, candidates, block_size, 1)
+    _assert_agrees_within_ends(ref, got, row_ke, 1)
+
+
+def test_matches_shared_kernel_row_repeat():
+    """Speculative decode: next_n logit rows share one bound."""
+    row_repeat = 3
+    ends = [0, 1025, WIDE]
+    logits, row_ks, row_ke, candidates = _inputs(
+        len(ends) * row_repeat,
+        WIDE,
+        ends,
+        None,
+        block_size=BLOCK_SIZE,
+        col_stride=1,
+        dtype=torch.int32,
+    )
+    ref, got = _run_both(logits, row_ks, row_ke, candidates, BLOCK_SIZE, row_repeat)
+    _assert_agrees_within_ends(ref, got, row_ke, row_repeat)
+
+
+def test_leaves_columns_past_end_untouched():
+    """The deliberate difference from the shared kernel, pinned explicitly.
+
+    Skipping the tail is where the speedup comes from, so a change that
+    quietly restores full-width sanitizing should fail here rather than just
+    get slower. Work stops at the tile boundary containing ``end``, not at
+    ``end`` itself, so that is the bound asserted.
+    """
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        _MASK_TILE,
+        _apply_candidate_mask_strided,
+    )
+
+    ends = [0, 1, 1025, WIDE // 2]
+    logits, row_ks, row_ke, candidates = _inputs(
+        len(ends),
+        WIDE,
+        ends,
+        None,
+        block_size=BLOCK_SIZE,
+        col_stride=1,
+        dtype=torch.int32,
+    )
+    sentinel = 1.5
+    logits.fill_(sentinel)
+    _apply_candidate_mask_strided(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1)
+    for row, end in enumerate(ends):
+        touched = min(-(-end // _MASK_TILE) * _MASK_TILE, WIDE)
+        tail = logits[row, touched:]
+        assert torch.equal(tail, torch.full_like(tail, sentinel)), (
+            f"row {row} (end {end}) was written past column {touched}"
+        )
+        # And the bound is tight: a row that ends early must not have cost a
+        # full-width pass, which is the whole point of the kernel.
+        assert touched <= end + _MASK_TILE
+
+
+def test_cudagraph_replay_tracks_changing_ends():
+    """The grid is static so a FULL capture stays valid when ends change.
+
+    Capturing bakes the grid, so a data-dependent program count would freeze
+    whatever lengths were resident at capture. Replaying against longer and
+    shorter contexts than the captured ones is the property that matters.
+    """
+    from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+        apply_candidate_mask,
+    )
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        _apply_candidate_mask_strided,
+    )
+
+    rows = 4
+    captured_ends = [1024, 4096, 8192, 2048]
+    logits, row_ks, row_ke, candidates = _inputs(
+        rows,
+        WIDE,
+        captured_ends,
+        None,
+        block_size=BLOCK_SIZE,
+        col_stride=1,
+        dtype=torch.int32,
+    )
+    base = logits.clone()
+
+    # Warm up outside the capture: Triton compiles on first call, and autotune
+    # launches would otherwise happen inside the graph.
+    _apply_candidate_mask_strided(
+        logits.clone(), row_ks, row_ke, candidates, BLOCK_SIZE, 1
+    )
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        _apply_candidate_mask_strided(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1)
+
+    for replay_ends in ([1, 131072, WIDE, 0], [WIDE, 1025, 3, 65536]):
+        row_ke.copy_(torch.tensor(replay_ends, device="cuda", dtype=row_ke.dtype))
+        logits.copy_(base)
+        graph.replay()
+        torch.accelerator.synchronize()
+
+        ref = base.clone()
+        apply_candidate_mask(ref, row_ks, row_ke, candidates, BLOCK_SIZE, 1)
+        _assert_agrees_within_ends(ref, logits, row_ke, 1)
diff --git a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
index d67df8ee2b4d..d5c37724aacd 100644
--- a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
+++ b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
@@ -30,7 +30,10 @@ def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 def _install_fake_aiter_modules(
-    monkeypatch: pytest.MonkeyPatch, *, supports_fp8: bool
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    supports_fp8: bool,
+    supports_causal: bool = False,
 ) -> None:
     aiter_mod: Any = types.ModuleType("aiter")
     mla_mod: Any = types.ModuleType("aiter.mla")
@@ -65,6 +68,14 @@ def mla_decode_fwd_without_fp8(
 
         mla_decode_fwd = mla_decode_fwd_without_fp8
 
+    if supports_causal:
+        inner = mla_decode_fwd
+
+        def mla_decode_fwd_with_causal(*args, causal=True, **kwargs):
+            return inner(*args, **kwargs)
+
+        mla_decode_fwd = mla_decode_fwd_with_causal
+
     mla_mod.mla_decode_fwd = mla_decode_fwd
     aiter_mod.mla = mla_mod
 
@@ -135,3 +146,107 @@ def test_aiter_mla_fp8_support_result_is_cached(monkeypatch):
         assert _check_aiter_mla_fp8_support() is True
         assert _check_aiter_mla_fp8_support() is True
         assert signature_mock.call_count == 1
+
+
+def test_missing_causal_arg_fails_closed(monkeypatch):
+    """Old aiter wheels have no causal=; do not report that as a valid draft path."""
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
+
+    _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=False)
+    with pytest.raises(RuntimeError, match="causal-only"):
+        rocm_aiter_ops.mla_decode_supports_non_causal()
+    with pytest.raises(RuntimeError, match="causal-only"):
+        AiterMLABackend.supports_non_causal()
+
+
+@pytest.mark.parametrize("has_kernels", [False, True])
+def test_supports_non_causal_requires_arch_kernels(monkeypatch, has_kernels):
+    """The Python argument alone does not guarantee an arch has the kernel."""
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.backends.mla import rocm_aiter_mla
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
+
+    _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True)
+    monkeypatch.setattr(
+        rocm_aiter_mla,
+        "_aiter_mla_non_causal_asm_kernels",
+        lambda: has_kernels,
+    )
+    assert rocm_aiter_ops.mla_decode_supports_non_causal() is True
+    assert AiterMLABackend.supports_non_causal() is has_kernels
+
+
+@pytest.mark.parametrize("causal", [True, False])
+def test_decode_fwd_forwards_causal_to_aiter(monkeypatch, causal):
+    """The decode op must pass the block mask through; omitting it is causal-only."""
+    seen: list = []
+
+    def mla_decode_fwd(*args, causal=None, **kwargs):
+        seen.append(causal)
+
+    mla_mod: Any = types.ModuleType("aiter.mla")
+    mla_mod.mla_decode_fwd = mla_decode_fwd
+    aiter_mod: Any = types.ModuleType("aiter")
+    aiter_mod.mla = mla_mod
+    monkeypatch.setitem(sys.modules, "aiter", aiter_mod)
+    monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod)
+
+    import torch
+
+    import vllm._aiter_ops as aiter_ops
+
+    aiter_ops._rocm_aiter_mla_decode_fwd_impl(
+        torch.zeros(2, 1, 8),
+        torch.zeros(2, 1, 1, 8),
+        torch.zeros(2, 1, 8),
+        torch.zeros(3, dtype=torch.int32),
+        1,
+        causal=causal,
+    )
+    assert seen == [causal]
+
+
+def _aiter_mla_validate(
+    monkeypatch,
+    *,
+    dtype,
+    kv_cache_dtype,
+    use_non_causal,
+):
+    from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.attention.backend import AttentionType
+    from vllm.v1.attention.backends.mla import rocm_aiter_mla
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend
+
+    _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True)
+    monkeypatch.setattr(
+        rocm_aiter_mla, "_aiter_mla_non_causal_asm_kernels", lambda: True
+    )
+    return AiterMLABackend.validate_configuration(
+        head_size=576,
+        dtype=dtype,
+        kv_cache_dtype=kv_cache_dtype,
+        block_size=1,
+        use_mla=True,
+        has_sink=False,
+        use_sparse=False,
+        use_mm_prefix=False,
+        use_per_head_quant_scales=False,
+        device_capability=DeviceCapability(9, 5),
+        attn_type=AttentionType.DECODER,
+        use_non_causal=use_non_causal,
+    )
+
+
+def test_non_causal_fp16_auto_cache_is_rejected(monkeypatch):
+    """Pinned AITER aborts fp16 Q; auto-select must fall through to Triton."""
+    import torch
+
+    reasons = _aiter_mla_validate(
+        monkeypatch,
+        dtype=torch.float16,
+        kv_cache_dtype="auto",
+        use_non_causal=True,
+    )
+    assert any("non-causal fp16" in reason for reason in reasons)
diff --git a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
index 659211521a07..0944be8937c8 100644
--- a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
+++ b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
@@ -35,7 +35,7 @@ def _require_aiter():
 
 @torch.inference_mode()
 def test_mla_decode_fwd_op_schema() -> None:
-    """opcheck validates registration, schema, fake-tensor, and ``o`` aliasing.
+    """Opcheck validates registration, schema, fake-tensor, and ``o`` aliasing.
 
     A single opcheck call covers that the op is registered/callable, that its
     fake implementation matches the real op (torch.compile tracing), and that
diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
index 99da0af7d57b..aabce97fac9e 100644
--- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
+++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
@@ -81,7 +81,7 @@ def _make_common_metadata():
         query_start_loc=query_start_loc,
         query_start_loc_cpu=query_start_loc,
         seq_lens=seq_lens,
-        _seq_lens_cpu=seq_lens,
+        seq_lens_cpu_upper_bound=seq_lens,
         num_reqs=2,
         num_actual_tokens=2,
         max_query_len=1,
@@ -99,7 +99,7 @@ def _make_mixed_common_metadata():
         query_start_loc=query_start_loc,
         query_start_loc_cpu=query_start_loc,
         seq_lens=seq_lens,
-        _seq_lens_cpu=seq_lens,
+        seq_lens_cpu_upper_bound=seq_lens,
         num_reqs=2,
         num_actual_tokens=5,
         max_query_len=4,
diff --git a/tests/kernels/attention/test_rocm_attention_selector.py b/tests/kernels/attention/test_rocm_attention_selector.py
deleted file mode 100644
index 2a684ed70d60..000000000000
--- a/tests/kernels/attention/test_rocm_attention_selector.py
+++ /dev/null
@@ -1,73 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config
-from vllm.platforms.rocm import RocmPlatform
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
-
-
-@pytest.fixture(autouse=True)
-def clear_cache():
-    """Clear lru cache to ensure each test case runs without caching."""
-    _cached_get_attn_backend.cache_clear()
-
-
-@pytest.mark.skip(reason="Skipped for now. Should be revisited.")
-def test_selector(monkeypatch: pytest.MonkeyPatch):
-    # Set the current platform to ROCm using monkeypatch
-    monkeypatch.setattr("vllm.v1.attention.selector.current_platform", RocmPlatform())
-
-    # Test standard ROCm attention
-    attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_ATTN)
-    vllm_config = VllmConfig(attention_config=attention_config)
-
-    with set_current_vllm_config(vllm_config):
-        backend = get_attn_backend(16, torch.float16, torch.float16, 16, False)
-        assert backend.get_name() == "ROCM_FLASH" or backend.get_name() == "TRITON_ATTN"
-
-    # MLA test for deepseek related
-    # Change the attention backend to triton MLA
-    attention_config = AttentionConfig(backend=AttentionBackendEnum.TRITON_MLA)
-    vllm_config = VllmConfig(attention_config=attention_config)
-
-    with set_current_vllm_config(vllm_config):
-        backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
-        assert backend.get_name() == "TRITON_MLA"
-
-    # If attention backend is None
-    # If use_mla is true
-    # The selected backend is triton MLA
-    attention_config = AttentionConfig(backend=None)
-    vllm_config = VllmConfig(attention_config=attention_config)
-
-    with set_current_vllm_config(vllm_config):
-        backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
-        assert backend.get_name() == "TRITON_MLA"
-
-    # Change the attention backend to AITER MLA
-    attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_AITER_MLA)
-    vllm_config = VllmConfig(attention_config=attention_config)
-
-    with set_current_vllm_config(vllm_config):
-        backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False, use_mla=True)
-        assert backend.get_name() == "ROCM_AITER_MLA"
-
-    # If attention backend is None
-    # If use_mla is true
-    # If VLLM_ROCM_USE_AITER is enabled
-    # The selected backend is ROCM_AITER_MLA
-    with monkeypatch.context() as m:
-        m.setenv("VLLM_ROCM_USE_AITER", "1")
-
-        attention_config = AttentionConfig(backend=None)
-        vllm_config = VllmConfig(attention_config=attention_config)
-
-        with set_current_vllm_config(vllm_config):
-            backend = get_attn_backend(
-                576, torch.bfloat16, "auto", 1, False, use_mla=True
-            )
-            assert backend.get_name() == "ROCM_AITER_MLA"
diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
index fb696edd840f..6c10f169bee4 100644
--- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import functools
 from types import SimpleNamespace
 
 import pytest
@@ -188,6 +189,7 @@ def _ref_sparse_decode_ragged(
     extra_rows: list[list[int]] | None = None,
     main_use_fnuz: bool = False,
     extra_use_fnuz: bool = False,
+    extra_block_size: int | None = None,
 ) -> torch.Tensor:
     q_f32 = q.float()
     out = torch.empty_like(q_f32)
@@ -209,7 +211,10 @@ def _ref_sparse_decode_ragged(
             )
             row_kv.append(
                 _read_fp8_ds_mla_cache_rows(
-                    extra_cache, extra_slots, block_size, extra_use_fnuz
+                    extra_cache,
+                    extra_slots,
+                    extra_block_size or block_size,
+                    extra_use_fnuz,
                 )
             )
 
@@ -244,18 +249,32 @@ def _ragged_from_rows(
     )
 
 
+def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]:
+    ends = indptr.cpu().tolist()
+    values = indices[: ends[-1]].cpu().tolist()
+    return [values[start:end] for start, end in zip(ends, ends[1:])]
+
+
 def _launch_sparse_decode_reduce(
     part_m: torch.Tensor,
     part_l: torch.Tensor,
     part_acc: torch.Tensor,
     adaptive_splits: bool,
+    positions: torch.Tensor | None = None,
+    cos_sin_cache: torch.Tensor | None = None,
+    out_dtype: torch.dtype = torch.bfloat16,
+    out_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
+    """Run the reduce kernel, with the inverse-RoPE epilogue if given positions.
+
+    Passing ``out_scale`` selects the MXFP8 epilogue (``out_dtype`` e4m3).
+    """
     from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
 
     num_queries, num_splits, num_heads = part_m.shape
     out = torch.empty(
         (num_queries, num_heads, HEAD_DIM),
-        dtype=torch.bfloat16,
+        dtype=out_dtype,
         device=part_m.device,
     )
     attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device)
@@ -265,13 +284,19 @@ def _launch_sparse_decode_reduce(
         part_acc,
         attn_sink,
         out,
+        positions,
+        cos_sin_cache,
+        out_scale,
         out.stride(0),
         out.stride(1),
+        out_scale.stride(0) if out_scale is not None else 0,
+        HEAD_DIM // 32,
         part_m.stride(0),
         part_m.stride(1),
         part_acc.stride(0),
         part_acc.stride(1),
         part_acc.stride(2),
+        cos_sin_cache.stride(0) if cos_sin_cache is not None else 0,
         num_heads,
         HAS_ATTN_SINK=False,
         ADAPTIVE_SPLITS=adaptive_splits,
@@ -279,6 +304,10 @@ def _launch_sparse_decode_reduce(
         BLOCK_H=1,
         NUM_SPLITS=num_splits,
         SPLITS_PAD=1 << (num_splits - 1).bit_length(),
+        FUSE_INV_ROPE=positions is not None,
+        NOPE=NOPE_HEAD_DIM,
+        HALF=ROPE_HEAD_DIM // 2,
+        QUANT_OUT=out_scale is not None,
         num_warps=4,
     )
     return out
@@ -336,6 +365,101 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None:
     torch.testing.assert_close(actual_lens, expected_lens)
 
 
+@torch.inference_mode()
+def test_combine_topk_swa_indices_adds_image_visibility() -> None:
+    from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices
+
+    device = torch.device("cuda")
+    num_tokens = 8
+    topk_indices = torch.full((num_tokens, 1), -1, dtype=torch.int32, device=device)
+    query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
+    seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
+    gather_lens = seq_lens.clone()
+    left_visible = torch.tensor(
+        [0, 0, 0, 1, 2, 3, 4, 0], dtype=torch.int32, device=device
+    )
+    right_visible = torch.tensor(
+        [0, 0, 4, 3, 2, 1, 0, 0], dtype=torch.int32, device=device
+    )
+
+    indices, lens = combine_topk_swa_indices(
+        topk_indices,
+        query_start_loc,
+        seq_lens,
+        gather_lens,
+        window_size=4,
+        compress_ratio=1,
+        topk=0,
+        M=16,
+        N=0,
+        max_image_tokens=5,
+        left_visible=left_visible,
+        right_visible=right_visible,
+    )
+
+    expected_rows = [
+        [0],
+        [0, 1],
+        [0, 1, 2, 3, 4, 5, 6],
+        [0, 1, 2, 3, 4, 5, 6],
+        [1, 2, 3, 4, 5, 6],
+        [2, 3, 4, 5, 6],
+        [2, 3, 4, 5, 6],
+        [4, 5, 6, 7],
+    ]
+    for token_idx, expected in enumerate(expected_rows):
+        actual = indices[token_idx, : lens[token_idx]].cpu().tolist()
+        assert actual == expected
+
+
+@torch.inference_mode()
+def test_combine_topk_swa_indices_apc_hit_inside_image() -> None:
+    from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices
+
+    device = torch.device("cuda")
+    indices, lens = combine_topk_swa_indices(
+        torch.full((2, 1), -1, dtype=torch.int32, device=device),
+        torch.tensor([0, 2], dtype=torch.int32, device=device),
+        torch.tensor([10], dtype=torch.int32, device=device),
+        # Only positions [5, 10) exist in the gathered SWA workspace.
+        torch.tensor([5], dtype=torch.int32, device=device),
+        window_size=4,
+        compress_ratio=1,
+        topk=0,
+        M=10,
+        N=0,
+        max_image_tokens=10,
+        left_visible=torch.tensor([8, 9], dtype=torch.int32, device=device),
+        # Deliberately extends beyond seq_len to exercise the upper clamp too.
+        right_visible=torch.tensor([5, 5], dtype=torch.int32, device=device),
+    )
+
+    assert lens.cpu().tolist() == [5, 5]
+    assert indices[:, :5].cpu().tolist() == [list(range(5)), list(range(5))]
+
+
+@torch.inference_mode()
+def test_combine_topk_swa_indices_keeps_vision_row_width_without_images() -> None:
+    from vllm.models.deepseek_v4.amd.rocm import combine_topk_swa_indices
+
+    device = torch.device("cuda")
+    indices, lens = combine_topk_swa_indices(
+        torch.full((1, 1), -1, dtype=torch.int32, device=device),
+        torch.tensor([0, 1], dtype=torch.int32, device=device),
+        torch.tensor([1], dtype=torch.int32, device=device),
+        torch.tensor([1], dtype=torch.int32, device=device),
+        window_size=120,
+        compress_ratio=1,
+        topk=0,
+        M=256,
+        N=0,
+        max_image_tokens=16,
+    )
+
+    assert indices.shape == (1, 256)
+    assert lens.item() == 1
+
+
 def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None:
     from vllm.models.deepseek_v4.amd import rocm as mod
 
@@ -903,6 +1027,211 @@ def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> No
     assert torch.equal(actual, torch.full_like(actual, 2))
 
 
+@requires_gfx950
+@pytest.mark.parametrize("adaptive_splits", [False, True])
+@torch.inference_mode()
+def test_sparse_attn_decode_reduce_inverse_rope_epilogue(adaptive_splits: bool) -> None:
+    """FUSE_INV_ROPE must land where the standalone rotation would.
+
+    The epilogue rotates a whole [H, nope + rope] row with one expression by
+    feeding cos=1/sin=0 to the NoPE lanes, which only holds if the pair split
+    lands on ``nope_head_dim // 2``; the NoPE lanes therefore have to come
+    back untouched, not merely close.
+    """
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
+
+    device = torch.device("cuda")
+    torch.manual_seed(7)
+    num_queries, num_splits, num_heads, max_pos = 5, 4, 8, 64
+    shape = (num_queries, num_splits, num_heads)
+    part_m = torch.randn(shape, dtype=torch.float32, device=device)
+    part_l = torch.rand(shape, dtype=torch.float32, device=device) + 0.5
+    part_acc = torch.randn((*shape, HEAD_DIM), dtype=torch.float32, device=device)
+    positions = torch.randint(
+        0, max_pos, (num_queries,), dtype=torch.int64, device=device
+    )
+    angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device)
+    cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous()
+
+    unfused = _launch_sparse_decode_reduce(part_m, part_l, part_acc, adaptive_splits)
+    expected = _fused_inverse_rope_gptj(
+        unfused, positions, cos_sin_cache, ROPE_HEAD_DIM
+    )
+    actual = _launch_sparse_decode_reduce(
+        part_m, part_l, part_acc, adaptive_splits, positions, cos_sin_cache
+    )
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+    assert torch.equal(actual[..., :NOPE_HEAD_DIM], unfused[..., :NOPE_HEAD_DIM])
+    # An epilogue that quietly did nothing would satisfy everything above.
+    assert not torch.equal(actual[..., NOPE_HEAD_DIM:], unfused[..., NOPE_HEAD_DIM:])
+
+
+def _mxfp8_dequant(data: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    blocks = data.float().view(*data.shape[:-1], -1, 32)
+    return (blocks * torch.exp2(scale.float() - 127.0)[..., None]).view(data.shape)
+
+
+def _random_cos_sin_cache(max_pos: int, device: torch.device) -> torch.Tensor:
+    angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device)
+    return torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous()
+
+
+@requires_gfx950
+@torch.inference_mode()
+def test_sparse_attn_decode_reduce_mxfp8_epilogue() -> None:
+    """QUANT_OUT must equal MXFP8-quantizing the reduce's fp32 output."""
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        _mxfp8_e4m3_quantize_torch,
+    )
+
+    device = torch.device("cuda")
+    torch.manual_seed(7)
+    num_queries, num_splits, num_heads = 5, 4, 8
+    shape = (num_queries, num_splits, num_heads)
+    part_m = torch.randn(shape, device=device)
+    part_l = torch.rand(shape, device=device) + 0.5
+    part_acc = torch.randn((*shape, HEAD_DIM), device=device)
+    positions = torch.randint(0, 64, (num_queries,), device=device)
+    cos_sin_cache = _random_cos_sin_cache(64, device)
+    args = (part_m, part_l, part_acc, False, positions, cos_sin_cache)
+
+    rotated = _launch_sparse_decode_reduce(*args, out_dtype=torch.float32)
+    expected_data, expected_scale = _mxfp8_e4m3_quantize_torch(
+        rotated.view(num_queries, -1)
+    )
+    scale = torch.empty_like(expected_scale)
+    data = _launch_sparse_decode_reduce(
+        *args, out_dtype=torch.float8_e4m3fn, out_scale=scale
+    )
+
+    assert torch.equal(scale, expected_scale)
+    assert torch.equal(
+        data.view(num_queries, -1).view(torch.uint8), expected_data.view(torch.uint8)
+    )
+
+
+@requires_gfx950
+@torch.inference_mode()
+def test_inverse_rope_mxfp8_rows() -> None:
+    """Prefill rows get the same inverse RoPE + MXFP8 as the decode epilogue."""
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        _mxfp8_e4m3_quantize_torch,
+    )
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        rocm_inverse_rope_mxfp8_rows,
+    )
+
+    device = torch.device("cuda")
+    torch.manual_seed(5)
+    num_tokens, num_heads = 7, 4
+    o = torch.randn(num_tokens, num_heads, HEAD_DIM, device=device).bfloat16()
+    positions = torch.randint(0, 64, (num_tokens,), device=device)
+    cos_sin_cache = _random_cos_sin_cache(64, device)
+
+    ref = o.float()
+    cos, sin = cos_sin_cache[positions, None].chunk(2, dim=-1)
+    even, odd = ref[..., NOPE_HEAD_DIM::2].clone(), ref[..., NOPE_HEAD_DIM + 1 :: 2]
+    ref[..., NOPE_HEAD_DIM::2] = even * cos + odd * sin
+    ref[..., NOPE_HEAD_DIM + 1 :: 2] = odd * cos - even * sin
+    ref = ref.view(num_tokens, -1)
+    _, expected_scale = _mxfp8_e4m3_quantize_torch(ref)
+
+    data = torch.empty_like(ref, dtype=torch.float8_e4m3fn)
+    scale = torch.empty_like(expected_scale)
+    rocm_inverse_rope_mxfp8_rows(
+        o, positions, cos_sin_cache, ROPE_HEAD_DIM, data, scale
+    )
+
+    assert torch.equal(scale, expected_scale)
+    # Within half an e4m3 step at the top of each block's range.
+    half_step = 16.0 * torch.exp2(scale.float() - 127.0).repeat_interleave(32, -1)
+    assert bool(((_mxfp8_dequant(data, scale) - ref).abs() <= half_step).all())
+
+
+@requires_gfx950
+@torch.inference_mode()
+def test_sparse_attn_decode_mxfp8_output() -> None:
+    """``out_mxfp8`` routes the real decode through the MXFP8 epilogue."""
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        _rocm_sparse_attn_decode_ragged_triton,
+    )
+
+    device = torch.device("cuda")
+    torch.manual_seed(1)
+    num_heads = 16
+    q = torch.randn(2, num_heads, HEAD_DIM, device=device).bfloat16() * 0.125
+    kv = torch.randn(6, HEAD_DIM, device=device).bfloat16() * 0.125
+    decode = functools.partial(
+        _rocm_sparse_attn_decode_ragged_triton,
+        q=q,
+        main_cache=_pack_fp8_ds_mla_cache(kv, 4, current_platform.is_fp8_fnuz()),
+        main_indices=torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device),
+        main_indptr=torch.tensor([0, 2, 4], dtype=torch.int32, device=device),
+        scale=HEAD_DIM**-0.5,
+        attn_sink=None,
+        nope_head_dim=NOPE_HEAD_DIM,
+        rope_head_dim=ROPE_HEAD_DIM,
+        inv_rope_positions=torch.tensor([3, 9], device=device),
+        inv_rope_cos_sin_cache=_random_cos_sin_cache(16, device),
+    )
+
+    expected = decode().float().view(2, -1)
+    data = torch.empty(
+        2, num_heads * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device
+    )
+    scale = torch.empty(2, num_heads * HEAD_DIM // 32, dtype=torch.uint8, device=device)
+    decode(out_mxfp8=(data, scale))
+
+    torch.testing.assert_close(
+        _mxfp8_dequant(data, scale), expected, atol=1e-3, rtol=2**-4
+    )
+
+
+@requires_gfx950
+@pytest.mark.parametrize(
+    "num_tokens, n_groups",
+    # One case per tile tier of _mxfp8_wo_a_bmm_config, with partial M tiles.
+    [
+        (1, 4),
+        (20, 4),
+        (48, 4),
+        (77, 4),
+        (130, 4),
+        (300, 4),
+        (700, 4),
+        (1000, 4),
+        (1100, 8),
+    ],
+)
+@torch.inference_mode()
+def test_rocm_mxfp8_wo_a_bmm(num_tokens: int, n_groups: int) -> None:
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        _mxfp8_e4m3_quantize_torch,
+    )
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_mxfp8_wo_a_bmm
+
+    device = torch.device("cuda")
+    torch.manual_seed(3)
+    o_lora_rank, group_dim = 128, 2048
+    a, a_scale = _mxfp8_e4m3_quantize_torch(
+        torch.randn(num_tokens, n_groups * group_dim, device=device)
+    )
+    w, w_scale = _mxfp8_e4m3_quantize_torch(
+        torch.randn(n_groups * o_lora_rank, group_dim, device=device)
+    )
+    wo_a = SimpleNamespace(weight=w, weight_scale=w_scale)
+
+    out = rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank)
+
+    expected = torch.einsum(
+        "tgd,grd->tgr",
+        _mxfp8_dequant(a, a_scale).view(num_tokens, n_groups, group_dim),
+        _mxfp8_dequant(w, w_scale).view(n_groups, o_lora_rank, group_dim),
+    )
+    torch.testing.assert_close(out.float(), expected.flatten(1), atol=5e-2, rtol=1e-2)
+
+
 @requires_gfx950
 @pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65])
 @torch.inference_mode()
@@ -1097,6 +1426,254 @@ def run_decode() -> torch.Tensor:
     torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2)
 
 
+@requires_gfx950
+@torch.inference_mode()
+def test_dsv4_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None:
+    """Replay the production ROCm decode path after device-only reallocation."""
+    from tests.v1.attention.utils import create_vllm_config
+    from vllm.models.deepseek_v4.amd.rocm import (
+        DeepseekV4ROCMAiterMLASparseMetadataBuilder,
+        DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
+    )
+    from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
+    from vllm.v1.kv_cache_interface import MLAAttentionSpec, SlidingWindowMLASpec
+
+    device = torch.device("cuda")
+    num_reqs = 3
+    upper_query_len = 8
+    graph_tokens = num_reqs * upper_query_len
+    block_size = 256
+    compressed_block_size = block_size // 128
+    window_size = 32
+
+    vllm_config = create_vllm_config(
+        model_name="facebook/opt-125m",
+        max_model_len=1024,
+        block_size=block_size,
+        max_num_seqs=num_reqs,
+        max_num_batched_tokens=graph_tokens,
+        hf_config_override={
+            "compress_ratios": [128],
+            "index_topk": 2048,
+            "sliding_window": window_size,
+        },
+    )
+    vllm_config.speculative_config = SimpleNamespace(
+        num_speculative_tokens=upper_query_len - 1,
+        parallel_drafting=False,
+        enable_adaptive_verification=True,
+        use_dspark=lambda: True,
+    )
+    mla_spec = MLAAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=HEAD_DIM,
+        dtype=torch.bfloat16,
+        tokens_per_state=128,
+        cache_dtype_str="fp8_ds_mla",
+        model_version="deepseek_v4",
+    )
+    swa_spec = SlidingWindowMLASpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=HEAD_DIM,
+        dtype=torch.bfloat16,
+        sliding_window=window_size,
+        cache_dtype_str="fp8_ds_mla",
+        model_version="deepseek_v4",
+    )
+    mla_builder = DeepseekV4ROCMAiterMLASparseMetadataBuilder(
+        mla_spec, ["c128a"], vllm_config, device
+    )
+    swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder(
+        swa_spec, ["c128a"], vllm_config, device
+    )
+    assert (
+        mla_builder.get_cudagraph_support(vllm_config, mla_spec)
+        == AttentionCGSupport.ALWAYS
+    )
+    assert (
+        swa_builder.get_cudagraph_support(vllm_config, swa_spec)
+        == AttentionCGSupport.ALWAYS
+    )
+
+    seq_lens = torch.tensor([520, 528, 536], dtype=torch.int32, device=device)
+    seq_lens_cpu = seq_lens.cpu()
+    cpu_query_start_loc = torch.arange(
+        0,
+        graph_tokens + 1,
+        upper_query_len,
+        dtype=torch.int32,
+    )
+    max_blocks = (int(seq_lens_cpu.max()) + block_size - 1) // block_size
+    block_table = torch.arange(
+        num_reqs * max_blocks, dtype=torch.int32, device=device
+    ).view(num_reqs, max_blocks)
+
+    def build_metadata(query_lens: list[int]):
+        query_lens_tensor = torch.tensor(query_lens, dtype=torch.int32, device=device)
+        query_start_loc = torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int32, device=device),
+                query_lens_tensor.cumsum(0),
+            ]
+        )
+        active_tokens = sum(query_lens)
+        positions = torch.zeros(graph_tokens, dtype=torch.int64, device=device)
+        position_rows = [
+            torch.arange(
+                seq_len - query_len,
+                seq_len,
+                dtype=torch.int64,
+                device=device,
+            )
+            for seq_len, query_len in zip(seq_lens_cpu.tolist(), query_lens)
+        ]
+        positions[:active_tokens] = torch.cat(position_rows)
+        slot_mapping = torch.full((graph_tokens,), -1, dtype=torch.int64, device=device)
+        slot_mapping[:active_tokens] = torch.arange(
+            active_tokens, dtype=torch.int64, device=device
+        )
+        common = CommonAttentionMetadata(
+            query_start_loc=query_start_loc,
+            query_start_loc_cpu=cpu_query_start_loc,
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=seq_lens_cpu,
+            num_reqs=num_reqs,
+            num_actual_tokens=graph_tokens,
+            max_query_len=upper_query_len,
+            max_seq_len=int(seq_lens_cpu.max()),
+            block_table_tensor=block_table,
+            slot_mapping=slot_mapping,
+            positions=positions,
+            causal=True,
+        )
+        return (
+            mla_builder.build_for_cudagraph_capture(common),
+            swa_builder.build_for_cudagraph_capture(common),
+        )
+
+    mla_metadata, swa_metadata = build_metadata([8, 8, 8])
+    assert mla_metadata.for_cudagraph_capture
+    metadata_ptrs = (
+        mla_metadata.c128a_decode_topk_ragged_indices.data_ptr(),
+        mla_metadata.c128a_decode_topk_ragged_indptr.data_ptr(),
+        swa_metadata.decode_swa_ragged_indices.data_ptr(),
+        swa_metadata.decode_swa_ragged_indptr.data_ptr(),
+    )
+
+    torch.manual_seed(19)
+    num_heads = 16
+    q = (
+        torch.randn(
+            graph_tokens,
+            num_heads,
+            HEAD_DIM,
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        * 0.125
+    )
+    num_cache_blocks = num_reqs * max_blocks
+    swa_cache = _pack_fp8_ds_mla_cache(
+        torch.randn(
+            num_cache_blocks * block_size,
+            HEAD_DIM,
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        * 0.125,
+        block_size,
+        use_fnuz=False,
+    )
+    compressed_cache = _pack_fp8_ds_mla_cache(
+        torch.randn(
+            num_cache_blocks * compressed_block_size,
+            HEAD_DIM,
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        * 0.125,
+        compressed_block_size,
+        use_fnuz=False,
+    )
+    attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
+    out = torch.empty_like(q)
+    monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1)
+
+    def run_decode(mla_md, swa_md) -> None:
+        mod.rocm_sparse_attn_decode(
+            q=q,
+            kv_cache=compressed_cache,
+            swa_k_cache=swa_cache,
+            swa_only=False,
+            topk_indices=mla_md.c128a_global_decode_topk_indices,
+            topk_lens=mla_md.c128a_decode_topk_lens,
+            swa_indices=swa_md.decode_swa_indices,
+            swa_lens=swa_md.decode_swa_lens,
+            swa_ragged_indices=swa_md.decode_swa_ragged_indices,
+            swa_ragged_indptr=swa_md.decode_swa_ragged_indptr,
+            topk_ragged_indices=mla_md.c128a_decode_topk_ragged_indices,
+            topk_ragged_indptr=mla_md.c128a_decode_topk_ragged_indptr,
+            attn_sink=attn_sink,
+            scale=HEAD_DIM**-0.5,
+            head_dim=HEAD_DIM,
+            nope_head_dim=NOPE_HEAD_DIM,
+            rope_head_dim=ROPE_HEAD_DIM,
+            output=out,
+            extra_cache_nan_free=True,
+            adaptive_splits=True,
+        )
+
+    def reference(mla_md, swa_md) -> torch.Tensor:
+        return _ref_sparse_decode_ragged(
+            q=q,
+            main_cache=swa_cache,
+            main_rows=_rows_from_ragged(
+                swa_md.decode_swa_ragged_indices,
+                swa_md.decode_swa_ragged_indptr,
+            ),
+            scale=HEAD_DIM**-0.5,
+            attn_sink=attn_sink,
+            block_size=block_size,
+            extra_cache=compressed_cache,
+            extra_rows=_rows_from_ragged(
+                mla_md.c128a_decode_topk_ragged_indices,
+                mla_md.c128a_decode_topk_ragged_indptr,
+            ),
+            extra_block_size=compressed_block_size,
+        )
+
+    run_decode(mla_metadata, swa_metadata)
+    torch.accelerator.synchronize()
+    expected_full = reference(mla_metadata, swa_metadata)
+    torch.testing.assert_close(out, expected_full, atol=2e-2, rtol=2e-2)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run_decode(mla_metadata, swa_metadata)
+    torch.accelerator.synchronize()
+    captured_full = out.clone()
+
+    reallocated_mla, reallocated_swa = build_metadata([3, 8, 5])
+    assert metadata_ptrs == (
+        reallocated_mla.c128a_decode_topk_ragged_indices.data_ptr(),
+        reallocated_mla.c128a_decode_topk_ragged_indptr.data_ptr(),
+        reallocated_swa.decode_swa_ragged_indices.data_ptr(),
+        reallocated_swa.decode_swa_ragged_indptr.data_ptr(),
+    )
+    expected_reallocated = reference(reallocated_mla, reallocated_swa)
+    run_decode(reallocated_mla, reallocated_swa)
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
+    assert not torch.equal(captured_full, expected_reallocated)
+
+    graph.replay()
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
+
+
 # ---------------------------------------------------------------------------
 # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum)
 # ---------------------------------------------------------------------------
@@ -1268,6 +1845,31 @@ def test_get_cached_wo_a_bf16_plain_caches() -> None:
     torch.testing.assert_close(out2, expected, atol=0, rtol=0)
 
 
+@pytest.mark.parametrize("scale_attr", ["weight_scale", "weight_scale_inv"])
+@torch.inference_mode()
+def test_get_cached_wo_a_bf16_dequantized_ignores_retained_scale(
+    scale_attr: str,
+) -> None:
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
+
+    device = torch.device("cuda")
+    torch.manual_seed(6)
+    n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
+    weight = torch.randn(
+        n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device
+    )
+    retained_scale = torch.full(
+        (n_local_groups, 1), 0.25, dtype=torch.float32, device=device
+    )
+    wo_a = _FakeWoA(weight)
+    setattr(wo_a, scale_attr, retained_scale)
+
+    out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
+
+    expected = weight.view(n_local_groups, o_lora_rank, hidden_dim)
+    torch.testing.assert_close(out, expected, atol=0, rtol=0)
+
+
 @torch.inference_mode()
 def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:
     from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
@@ -1310,3 +1912,189 @@ def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:
 
     # Second call returns the same cached object.
     assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out
+
+
+@requires_split_decode_arch
+@torch.inference_mode()
+def test_fp8_paged_mqa_logits_triton_matches_torch_ref() -> None:
+    """Block-flat Triton decode logits vs the eager torch reference.
+
+    Compare only valid key positions: the kernel skips the -inf pad that
+    the torch ref writes past seq_len.
+    """
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
+    from vllm.v1.worker.workspace import init_workspace_manager
+
+    device = torch.device("cuda")
+    init_workspace_manager(device)
+    torch.manual_seed(0)
+
+    fp8_dtype = current_platform.fp8_dtype()
+    batch_size = 2
+    next_n = 1
+    num_heads = 4
+    head_size = 128
+    block_size = 64
+    max_model_len = 256
+    seq_lens = [180, 64]
+    num_pages = [(seq_len + block_size - 1) // block_size for seq_len in seq_lens]
+    num_blocks = sum(num_pages)
+
+    q = torch.randn(
+        batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16
+    ).to(fp8_dtype)
+    weights = torch.rand(
+        batch_size * next_n, num_heads, device=device, dtype=torch.float32
+    )
+    context_lens = torch.tensor(seq_lens, device=device, dtype=torch.int32)
+    block_tables = torch.full(
+        (batch_size, max(num_pages)), 0, device=device, dtype=torch.int32
+    )
+    page = 0
+    for i, n_pages in enumerate(num_pages):
+        block_tables[i, :n_pages] = torch.arange(
+            page, page + n_pages, device=device, dtype=torch.int32
+        )
+        page += n_pages
+
+    kv_cache = torch.empty(
+        num_blocks, block_size, 1, head_size + 4, device=device, dtype=torch.uint8
+    )
+    values = torch.randn(
+        num_blocks, block_size, head_size, device=device, dtype=torch.bfloat16
+    ).to(fp8_dtype)
+    scales = torch.rand(
+        num_blocks, block_size, device=device, dtype=torch.float32
+    ).clamp(min=1e-3)
+    kv_flat = kv_cache.view(num_blocks, -1)
+    scale_off = block_size * head_size
+    kv_flat[:, :scale_off] = values.reshape(num_blocks, -1).view(torch.uint8)
+    kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_blocks, -1)
+
+    got = mod.rocm_fp8_paged_mqa_logits_triton(
+        q, kv_cache, weights, context_lens, block_tables, max_model_len
+    )
+    ref = mod.fp8_paged_mqa_logits_torch(
+        q, kv_cache, weights, context_lens, block_tables, max_model_len
+    )
+    for i, seq_len in enumerate(seq_lens):
+        torch.testing.assert_close(
+            got[i, :seq_len], ref[i, :seq_len], atol=2e-3, rtol=2e-3
+        )
+
+
+def test_indexer_k_is_c4a_block_flat_gate() -> None:
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        _indexer_k_is_c4a_block_flat,
+    )
+
+    assert not _indexer_k_is_c4a_block_flat(1)
+    assert not _indexer_k_is_c4a_block_flat(2)
+    assert _indexer_k_is_c4a_block_flat(4)
+    assert not _indexer_k_is_c4a_block_flat(128)
+
+
+@requires_split_decode_arch
+@torch.inference_mode()
+def test_paged_mqa_logits_gate_v41_shuffle_vs_c4a_flat() -> None:
+    """Ratio 1/2 stay on AITER; ratio 4 takes Triton."""
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        fp8_paged_mqa_logits_torch,
+        indexer_k_quant_and_cache_triton,
+        rocm_fp8_paged_mqa_logits,
+        rocm_fp8_paged_mqa_logits_triton,
+    )
+    from vllm.v1.worker.workspace import init_workspace_manager
+
+    device = torch.device("cuda")
+    init_workspace_manager(device)
+    torch.manual_seed(1)
+
+    fp8_dtype = current_platform.fp8_dtype()
+    batch_size = 1
+    next_n = 1
+    num_heads = 32
+    head_size = 128
+    block_size = 64
+    seq_len = 180
+    max_model_len = 256
+    num_pages = (seq_len + block_size - 1) // block_size
+    num_tokens = num_pages * block_size
+
+    q = torch.randn(
+        batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16
+    ).to(fp8_dtype)
+    weights = torch.rand(
+        batch_size * next_n, num_heads, device=device, dtype=torch.float32
+    )
+    context_lens = torch.tensor([seq_len], device=device, dtype=torch.int32)
+    block_tables = torch.arange(num_pages, device=device, dtype=torch.int32).view(1, -1)
+    dummy_sched = torch.zeros(8, 2, device=device, dtype=torch.int32)
+
+    k_bf16 = torch.randn(num_tokens, head_size, device=device, dtype=torch.bfloat16)
+    slot = torch.arange(num_tokens, device=device, dtype=torch.int32)
+    cache_shuffle = torch.zeros(
+        num_pages, block_size, head_size + 4, device=device, dtype=torch.uint8
+    )
+    indexer_k_quant_and_cache_triton(k_bf16, cache_shuffle, slot, 128, "ue8m0")
+    cache_shuffle = cache_shuffle.unsqueeze(2)
+
+    aiter = rocm_fp8_paged_mqa_logits(
+        q,
+        cache_shuffle,
+        weights,
+        context_lens,
+        block_tables,
+        dummy_sched,
+        max_model_len,
+        compress_ratio=1,
+    )
+    for ratio in (1, 2):
+        got = rocm_fp8_paged_mqa_logits(
+            q,
+            cache_shuffle,
+            weights,
+            context_lens,
+            block_tables,
+            dummy_sched,
+            max_model_len,
+            compress_ratio=ratio,
+        )
+        torch.testing.assert_close(got[:, :seq_len], aiter[:, :seq_len], atol=0, rtol=0)
+
+    cache_flat = torch.empty(
+        num_pages, block_size, 1, head_size + 4, device=device, dtype=torch.uint8
+    )
+    values = torch.randn(
+        num_pages, block_size, head_size, device=device, dtype=torch.bfloat16
+    ).to(fp8_dtype)
+    scales = torch.rand(
+        num_pages, block_size, device=device, dtype=torch.float32
+    ).clamp(min=1e-3)
+    kv_flat = cache_flat.view(num_pages, -1)
+    scale_off = block_size * head_size
+    kv_flat[:, :scale_off] = values.reshape(num_pages, -1).view(torch.uint8)
+    kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_pages, -1)
+
+    gated_c4a = rocm_fp8_paged_mqa_logits(
+        q,
+        cache_flat,
+        weights,
+        context_lens,
+        block_tables,
+        dummy_sched,
+        max_model_len,
+        compress_ratio=4,
+    )
+    triton = rocm_fp8_paged_mqa_logits_triton(
+        q, cache_flat, weights, context_lens, block_tables, max_model_len
+    )
+    ref = fp8_paged_mqa_logits_torch(
+        q, cache_flat, weights, context_lens, block_tables, max_model_len
+    )
+    torch.testing.assert_close(
+        gated_c4a[:, :seq_len], triton[:, :seq_len], atol=0, rtol=0
+    )
+    torch.testing.assert_close(
+        gated_c4a[:, :seq_len], ref[:, :seq_len], atol=2e-3, rtol=2e-3
+    )
diff --git a/tests/kernels/attention/test_triton_mla_causal_verify_flatten.py b/tests/kernels/attention/test_triton_mla_causal_verify_flatten.py
new file mode 100644
index 000000000000..a27552c90e3e
--- /dev/null
+++ b/tests/kernels/attention/test_triton_mla_causal_verify_flatten.py
@@ -0,0 +1,215 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the multi-token decode flatten in ``TritonMLAImpl.forward_mqa``.
+
+``decode_attention_fwd`` has no causal flag: it launches one program per row of
+``q`` and reads that row's KV extent from ``B_Seqlen``. So a ``query_len``-token
+decode block is flattened to one row per query token, and intra-block causality
+has to be expressed through the per-row sequence lengths -- row ``t`` gets
+``seq_len - (query_len - 1) + t``, i.e. the committed prefix plus block tokens
+``0..t``. A non-causal draft block instead gives every row the same full extent.
+
+The kernel is replaced by a spy, so these tests assert on the ``block_table``
+and ``seq_lens`` the backend actually submits and need no GPU or Triton runtime.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.attention.mla_attention import QueryLenSupport
+from vllm.v1.attention.backend import AttentionCGSupport
+from vllm.v1.attention.backends.mla.triton_mla import (
+    TritonMLAImpl,
+    TritonMLAMetadataBuilder,
+)
+
+NUM_HEADS = 16
+KV_LORA_RANK = 512
+QK_ROPE_HEAD_DIM = 64
+HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM
+PAGE_SIZE = 16
+
+# Committed context per request, including a fresh request at zero context
+# where the whole KV extent is the block itself. A decode block's seq_len
+# already counts the block, so seq_len = context + query_len.
+CONTEXT_LENS = [0, 1, 37, 512]
+
+
+def _seq_lens(query_len: int) -> list[int]:
+    return [c + query_len for c in CONTEXT_LENS]
+
+
+def _make_impl(dcp_world_size: int = 1) -> TritonMLAImpl:
+    """Only the attributes ``forward_mqa`` reads; __init__ wants a full config."""
+    impl = object.__new__(TritonMLAImpl)
+    impl.kv_lora_rank = KV_LORA_RANK
+    impl.scale = HEAD_SIZE**-0.5
+    impl._sm_count = 304
+    impl.dcp_world_size = dcp_world_size
+    return impl
+
+
+def _make_metadata(seq_lens: list[int], query_len: int, causal: bool):
+    num_decodes = len(seq_lens)
+    return SimpleNamespace(
+        num_decodes=num_decodes,
+        num_decode_tokens=num_decodes * query_len,
+        max_seq_len=max(seq_lens),
+        causal=causal,
+        decode=SimpleNamespace(
+            # arange page ids so each request's window is an identifiable slice.
+            block_table=torch.arange(num_decodes * 64, dtype=torch.int32).reshape(
+                num_decodes, 64
+            ),
+            seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
+        ),
+    )
+
+
+def _run_forward_mqa(
+    seq_lens: list[int],
+    query_len: int,
+    causal: bool,
+    q_rows=None,
+    dcp_world_size: int = 1,
+):
+    """Drive forward_mqa with the decode kernel spied out.
+
+    Returns the ``(block_table, seq_lens)`` handed to ``decode_attention_fwd``.
+    """
+    captured: dict = {}
+
+    def spy(q, kv_cache, kv_c_cache, o, lse, block_table, b_seq_len, *args, **kwargs):
+        captured["block_table"] = block_table.detach().clone()
+        captured["seq_lens"] = b_seq_len.detach().clone()
+
+    metadata = _make_metadata(seq_lens, query_len, causal)
+    rows = metadata.num_decode_tokens if q_rows is None else q_rows
+    q = torch.zeros(rows, NUM_HEADS, HEAD_SIZE, dtype=torch.bfloat16)
+    kv_cache = torch.zeros(
+        len(seq_lens) * 64, PAGE_SIZE, HEAD_SIZE, dtype=torch.bfloat16
+    )
+    layer = SimpleNamespace(_k_scale=torch.ones(1), layer_name="test")
+
+    with patch("vllm.v1.attention.backends.mla.triton_mla.decode_attention_fwd", spy):
+        _make_impl(dcp_world_size).forward_mqa(q, kv_cache, metadata, layer)
+
+    assert captured, "forward_mqa did not reach the decode kernel"
+    return captured
+
+
+def test_multi_token_decode_flags_are_declared_together():
+    """Capture probes max_query_len > 1, which only clears the assert in
+    build_for_cudagraph_capture once query_len_support raised the threshold."""
+    assert (
+        TritonMLAMetadataBuilder._cudagraph_support == AttentionCGSupport.UNIFORM_BATCH
+    )
+    assert TritonMLAMetadataBuilder.query_len_support == QueryLenSupport.UNIFORM
+
+
+@pytest.mark.parametrize("query_len", [2, 3, 4, 5, 8])
+def test_causal_block_rows_see_prefix_plus_own_position(query_len):
+    """Causal verify: row t sees seq_len - (query_len - 1) + t KV entries.
+
+    Regression guard: the causal branch used to be skipped entirely, so a verify
+    token could attend the draft siblings it is supposed to be checking. Fails
+    on unmodified upstream, where no flatten happens at all for causal blocks.
+    """
+    seq_lens = _seq_lens(query_len)
+    captured = _run_forward_mqa(seq_lens, query_len, causal=True)
+
+    got = captured["seq_lens"].tolist()
+    want = [s - (query_len - 1) + t for s in seq_lens for t in range(query_len)]
+    assert got == want, (
+        "causal flatten is not causal: verify row r*query_len+t must get "
+        f"seq_len_r - {query_len - 1} + t entries.\n"
+        f"  seq_lens {seq_lens}\n  got      {got}\n  expected {want}"
+    )
+
+    # Per-request invariants the arithmetic must hold regardless of query_len.
+    for r, seq_len in enumerate(seq_lens):
+        rows = got[r * query_len : (r + 1) * query_len]
+        assert rows == sorted(rows) and len(set(rows)) == query_len, (
+            f"request {r} rows must strictly increase, got {rows}"
+        )
+        assert rows[-1] == seq_len, "last row must see the full sequence"
+        assert rows[0] == seq_len - (query_len - 1)
+        assert all(n > 0 for n in rows), f"non-positive KV extent in {rows}"
+
+
+def test_cudagraph_padding_rows_present_no_kv_extent():
+    """Causal flattening pins cudagraph padding rows (seq_len 0) to exactly
+    zero KV extent, keeping them in the kernel's skipped branch."""
+    query_len = 4
+    seq_lens = _seq_lens(query_len) + [0]
+    captured = _run_forward_mqa(seq_lens, query_len, causal=True)
+
+    padding = captured["seq_lens"].tolist()[-query_len:]
+    assert padding == [0] * query_len, (
+        f"padding rows must present exactly zero KV extent, got {padding}"
+    )
+
+
+@pytest.mark.parametrize("query_len", [2, 4])
+def test_non_causal_block_rows_all_see_the_full_prefix(query_len):
+    """The DSpark draft block is generated in one pass, so no row is masked."""
+    seq_lens = _seq_lens(query_len)
+    captured = _run_forward_mqa(seq_lens, query_len, causal=False)
+
+    want = [s for s in seq_lens for _ in range(query_len)]
+    assert captured["seq_lens"].tolist() == want
+
+
+@pytest.mark.parametrize("causal", [True, False])
+def test_block_table_is_expanded_to_match_q_rows(causal):
+    """The kernel indexes block_table by program_id, so it needs one row per
+    query token -- a short block_table is read out of bounds."""
+    query_len = 4
+    captured = _run_forward_mqa(_seq_lens(query_len), query_len, causal=causal)
+
+    num_rows = len(CONTEXT_LENS) * query_len
+    assert captured["block_table"].shape[0] == num_rows
+    assert captured["seq_lens"].shape[0] == num_rows
+    # repeat_interleave, not repeat: a request's rows must stay adjacent.
+    for r in range(len(CONTEXT_LENS)):
+        block = captured["block_table"][r * query_len : (r + 1) * query_len]
+        assert torch.equal(block, block[:1].expand_as(block))
+
+
+@pytest.mark.parametrize("causal", [True, False])
+def test_single_token_decode_is_untouched(causal):
+    """query_len == 1 must short-circuit: ordinary decode sees no expansion."""
+    seq_lens = _seq_lens(1)
+    captured = _run_forward_mqa(seq_lens, query_len=1, causal=causal)
+
+    assert captured["seq_lens"].tolist() == seq_lens
+    assert captured["block_table"].shape[0] == len(seq_lens)
+
+
+def test_expansion_factor_follows_the_q_row_count():
+    """The kernel indexes block_table and seq_lens by program_id, so the
+    expansion must follow q rows, not the token/request counts."""
+    num_rows = len(CONTEXT_LENS) * 4
+    captured = _run_forward_mqa(_seq_lens(1), query_len=1, causal=True, q_rows=num_rows)
+
+    assert captured["block_table"].shape[0] == num_rows
+    assert captured["seq_lens"].shape[0] == num_rows
+
+
+def test_non_uniform_block_is_rejected():
+    """Q rows must divide evenly across the decode requests."""
+    with pytest.raises(AssertionError, match="non-uniform decode block"):
+        _run_forward_mqa(_seq_lens(1), query_len=1, causal=True, q_rows=6)
+
+
+def test_causal_multi_token_decode_is_rejected_under_dcp():
+    """Per-row extents offset the global seq_len, but DCP passes the rank-local
+    slice, so the causal arithmetic would silently address the wrong KV."""
+    with pytest.raises(AssertionError, match="not supported with DCP"):
+        _run_forward_mqa(_seq_lens(4), query_len=4, causal=True, dcp_world_size=2)
+
+    # Non-causal rows all take the same extent, which stays correct when local.
+    _run_forward_mqa(_seq_lens(4), query_len=4, causal=False, dcp_world_size=2)
diff --git a/tests/kernels/attention/test_triton_prefill_attention.py b/tests/kernels/attention/test_triton_prefill_attention.py
index 6316a926ae33..e9b90251a492 100644
--- a/tests/kernels/attention/test_triton_prefill_attention.py
+++ b/tests/kernels/attention/test_triton_prefill_attention.py
@@ -5,6 +5,7 @@
 import torch
 import torch.nn.functional as F
 
+import vllm.v1.attention.ops.triton_prefill_attention as prefill_ops
 from vllm.platforms import current_platform
 from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
 
@@ -230,3 +231,87 @@ def test_context_attention_sliding_window(
 
     # Compare outputs
     torch.testing.assert_close(o, o_ref, rtol=2e-2, atol=2e-2)
+
+
+class _LaunchCapture:
+    """Records the launch configuration in place of ``_fwd_kernel``."""
+
+    grid: tuple
+    kwargs: dict
+
+    def __getitem__(self, grid):
+        self.grid = grid
+        return self._record
+
+    def _record(self, *args, **kwargs) -> None:
+        self.kwargs = kwargs
+
+
+def _capture_tile_config(
+    monkeypatch, *, is_rocm: bool, on_gfx1x: bool, dtype=torch.bfloat16
+) -> _LaunchCapture:
+    """Capture the tile configuration with the platform predicates mocked.
+
+    Both tile widths are numerically correct, so the tests above pass whichever
+    one is selected. This needs no RDNA part and allocates no device memory, so
+    it covers the RDNA branch on the CDNA and NVIDIA agents CI actually runs.
+    """
+    platform = prefill_ops.current_platform
+    monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm)
+    # Every device this kernel targets is cuda-alike at capability 80 or better,
+    # so the stock tile is 128 for 16-bit dtypes.
+    monkeypatch.setattr(platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(platform, "has_device_capability", lambda *a, **k: True)
+    if is_rocm:
+        import vllm.platforms.rocm as rocm_platform
+
+        monkeypatch.setattr(rocm_platform, "on_gfx1x", lambda: on_gfx1x)
+
+    capture = _LaunchCapture()
+    monkeypatch.setattr(prefill_ops, "_fwd_kernel", capture)
+
+    def meta(*shape):
+        return torch.empty(shape, dtype=dtype, device="meta")
+
+    seq_lens = torch.empty(2, dtype=torch.int32, device="meta")
+    context_attention_fwd(
+        meta(256, 8, 128),
+        meta(256, 2, 128),
+        meta(256, 2, 128),
+        meta(256, 8, 128),
+        seq_lens,
+        seq_lens,
+        128,
+    )
+    return capture
+
+
+@pytest.mark.parametrize(
+    ("is_rocm", "on_gfx1x", "dtype", "expected_block_n"),
+    [
+        pytest.param(True, True, torch.bfloat16, 32, id="rdna"),
+        # on_gfx1x() is what excludes gfx10xx, CDNA and gfx1250.
+        pytest.param(True, False, torch.bfloat16, 128, id="rocm-not-rdna"),
+        pytest.param(False, False, torch.bfloat16, 128, id="not-rocm"),
+        # get_block_size already returns 32 here, so min() must not widen it.
+        pytest.param(True, True, torch.float32, 32, id="rdna-float32"),
+    ],
+)
+def test_kv_tile_width_is_gated_by_platform(
+    monkeypatch, is_rocm: bool, on_gfx1x: bool, dtype, expected_block_n: int
+) -> None:
+    capture = _capture_tile_config(
+        monkeypatch, is_rocm=is_rocm, on_gfx1x=on_gfx1x, dtype=dtype
+    )
+    assert capture.kwargs["BLOCK_N"] == expected_block_n
+
+
+def test_rdna_narrows_the_kv_tile_and_nothing_else(monkeypatch) -> None:
+    with monkeypatch.context() as m:
+        tuned = _capture_tile_config(m, is_rocm=True, on_gfx1x=True)
+    with monkeypatch.context() as m:
+        stock = _capture_tile_config(m, is_rocm=True, on_gfx1x=False)
+
+    assert tuned.kwargs.pop("BLOCK_N") != stock.kwargs.pop("BLOCK_N")
+    assert tuned.kwargs == stock.kwargs
+    assert tuned.grid == stock.grid
diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py
index 772fee137aa6..e30d9225eac8 100644
--- a/tests/kernels/attention/test_triton_unified_attention.py
+++ b/tests/kernels/attention/test_triton_unified_attention.py
@@ -10,6 +10,7 @@
 from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.attention.ops.triton_attention_helpers import (
+    apply_softcap,
     compute_tile_loop_bounds,
 )
 from vllm.v1.attention.ops.triton_unified_attention import unified_attention
@@ -135,6 +136,7 @@ def ref_paged_attn(
     num_seqs = len(query_lens)
     block_tables = block_tables.cpu().numpy()
     _, block_size, num_kv_heads, head_size = key_cache.shape
+    head_size_v = value_cache.shape[-1]
 
     outputs: list[torch.Tensor] = []
     start_idx = 0
@@ -149,7 +151,7 @@ def ref_paged_attn(
 
         k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
         k = k[:kv_len]
-        v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
+        v = value_cache[block_indices].view(-1, num_kv_heads, head_size_v)
         v = v[:kv_len]
 
         if q.shape[1] != k.shape[1]:
@@ -898,3 +900,40 @@ def test_triton_unified_attn_use_td_tile_clamp(
         soft_cap=None,
         seq_threshold_3D=0,
     )
+
+
+@triton.jit
+def _softcap_probe_kernel(
+    scores_ptr,
+    out_ptr,
+    numel,
+    soft_cap,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < numel
+    scores = tl.load(scores_ptr + offs, mask=mask, other=0.0)
+    tl.store(out_ptr + offs, apply_softcap(scores, soft_cap), mask=mask)
+
+
+@torch.inference_mode()
+def test_softcap_does_not_overflow_on_large_scores() -> None:
+    """Scores above ~88 * soft_cap must not overflow to inf/NaN.
+
+    The exp-based softcap computes ``(exp(y) - exp(-y)) / (exp(y) + exp(-y))``
+    with ``y = S / soft_cap``. For ``|y| > ~88`` the exponentials overflow to
+    ``inf`` and the ratio becomes ``inf / inf = NaN``, poisoning the whole
+    attention row. Gemma-2 style models use ``attn_logit_softcapping = 50``,
+    so scores above 4400 are in range for their large attention logits.
+    """
+    soft_cap = 50.0
+    scores = torch.tensor(
+        [1.0, 100.0, 1000.0, 3000.0, 5000.0, 10000.0, -1.0, -10000.0, 4400.0],
+        device=DEVICE_TYPE,
+        dtype=torch.float32,
+    )
+    out = torch.empty_like(scores)
+    _softcap_probe_kernel[(1,)](scores, out, scores.numel(), soft_cap, BLOCK=16)
+    ref = soft_cap * torch.tanh(scores / soft_cap)
+    assert torch.isfinite(out).all(), out
+    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py
index 513811a09c04..ea90b60a7514 100644
--- a/tests/kernels/attention/test_triton_unified_attention_diffkv.py
+++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py
@@ -1,12 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for the Triton DiffKV unified-attention kernel.
-"""
+"""Unit tests for the Triton DiffKV unified-attention kernel."""
 
 import pytest
 import torch
 
+from tests.kernels.attention.test_triton_unified_attention import ref_paged_attn
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import (
@@ -83,7 +82,7 @@ def _alloc_segm_buffers(seq_threshold_3D: int, num_query_heads: int, head_size_v
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
 @torch.inference_mode()
-def test_triton_unified_attn_diffkv_vs_fa(
+def test_triton_unified_attn_diffkv_vs_reference(
     seq_lens: list[tuple[int, int]],
     num_heads: tuple[int, int],
     head_sizes: tuple[int, int],
@@ -95,12 +94,13 @@ def test_triton_unified_attn_diffkv_vs_fa(
 ) -> None:
     head_size_qk, head_size_v = head_sizes
 
-    # DiffKV requires FA3 (Hopper) / FA4 (Blackwell) as the reference.
-    fa_version = get_flash_attn_version(head_size=head_size_qk, head_size_v=head_size_v)
-    if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4):
-        pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).")
-
-    from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
+    # Keep the FA3/FA4 comparison on NVIDIA; ROCm uses the PyTorch oracle.
+    if not current_platform.is_rocm():
+        fa_version = get_flash_attn_version(
+            head_size=head_size_qk, head_size_v=head_size_v
+        )
+        if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4):
+            pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).")
 
     torch.set_default_device(DEVICE_TYPE)
     set_random_seed(0)
@@ -137,28 +137,44 @@ def test_triton_unified_attn_diffkv_vs_fa(
         0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
     )
 
-    # ---- FlashAttention DiffKV (ground truth) ---------------------------
-    # Mirror the backend: fix degenerate strides on size-1 dims so FA's
-    # TMA path sees ≥16-byte-aligned strides (matters for num_kv_heads==1).
-    fa_k = canonicalize_singleton_dim_strides(key_cache)
-    fa_v = canonicalize_singleton_dim_strides(value_cache)
-    fa_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype)
-    flash_attn_varlen_func(
-        q=query,
-        k=fa_k,
-        v=fa_v,
-        out=fa_out,
-        cu_seqlens_q=cu_query_lens,
-        max_seqlen_q=max_query_len,
-        seqused_k=kv_lens_t,
-        max_seqlen_k=max_kv_len,
-        softmax_scale=scale,
-        causal=True,
-        window_size=list(window_size),
-        block_table=block_tables,
-        softcap=soft_cap if soft_cap is not None else 0,
-        fa_version=fa_version,
-    )
+    if current_platform.is_rocm():
+        # FP32 also keeps the helper's in-place scaling off the kernel input.
+        ref_out = ref_paged_attn(
+            query.float(),
+            key_cache.float(),
+            value_cache.float(),
+            query_lens,
+            kv_lens,
+            block_tables,
+            scale,
+            sliding_window=sliding_window,
+            soft_cap=soft_cap,
+        ).to(dtype)
+    else:
+        from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
+
+        # FA's TMA path needs aligned singleton strides (num_kv_heads == 1).
+        fa_k = canonicalize_singleton_dim_strides(key_cache)
+        fa_v = canonicalize_singleton_dim_strides(value_cache)
+        ref_out = torch.empty(
+            sum(query_lens), num_query_heads, head_size_v, dtype=dtype
+        )
+        flash_attn_varlen_func(
+            q=query,
+            k=fa_k,
+            v=fa_v,
+            out=ref_out,
+            cu_seqlens_q=cu_query_lens,
+            max_seqlen_q=max_query_len,
+            seqused_k=kv_lens_t,
+            max_seqlen_k=max_kv_len,
+            softmax_scale=scale,
+            causal=True,
+            window_size=list(window_size),
+            block_table=block_tables,
+            softcap=soft_cap if soft_cap is not None else 0,
+            fa_version=fa_version,
+        )
 
     # ---- Triton DiffKV --------------------------------------------------
     segm_output, segm_max, segm_expsum = _alloc_segm_buffers(
@@ -185,7 +201,4 @@ def test_triton_unified_attn_diffkv_vs_fa(
         softmax_segm_expsum=segm_expsum,
     )
 
-    (
-        torch.testing.assert_close(triton_out, fa_out, atol=2e-2, rtol=2e-2),
-        f"triton vs FA max abs diff: {torch.max(torch.abs(triton_out - fa_out))}",
-    )
+    torch.testing.assert_close(triton_out, ref_out, atol=2e-2, rtol=2e-2)
diff --git a/tests/kernels/attention/test_trtllm_kvfp8_dequant.py b/tests/kernels/attention/test_trtllm_kvfp8_dequant.py
index bd267c43bebe..8368869b65e3 100644
--- a/tests/kernels/attention/test_trtllm_kvfp8_dequant.py
+++ b/tests/kernels/attention/test_trtllm_kvfp8_dequant.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
+"""Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
 
 Tests both contiguous and non-contiguous (cross-layer unified) KV cache
 layouts against a pure-PyTorch reference implementation.
@@ -43,8 +42,7 @@ def make_random_kv_cache(
     layout=KVCacheLayout.LBHNC,
     num_layers=None,
 ):
-    """
-    Create a random fp8 KV cache mimicking the production allocation.
+    """Create a random fp8 KV cache mimicking the production allocation.
 
     Physical layout: [L, B, H, N, 2 * hs] permuted by ``layout.stride_order``
     Returned view:   (num_blocks, 2, num_kv_heads, block_size, head_size)
@@ -382,8 +380,7 @@ def test_large_block_size():
 
 @torch.inference_mode()
 def test_cross_layer_many_layers():
-    """
-    Non-contiguous with 36 layers -- matches real gpt-oss-120b.
+    """Non-contiguous with 36 layers -- matches real gpt-oss-120b.
     Cross-layer (BLHNC) strides are far from contiguous (factor of 36 in the gaps).
     """
     from vllm.v1.attention.backends.flashinfer import (
diff --git a/tests/kernels/attention/test_xpu_mla_sparse.py b/tests/kernels/attention/test_xpu_mla_sparse.py
index 965c6f735f0f..47139fc465d5 100644
--- a/tests/kernels/attention/test_xpu_mla_sparse.py
+++ b/tests/kernels/attention/test_xpu_mla_sparse.py
@@ -30,12 +30,12 @@ def reference_mla_sparse_prefill(
     topk_length: torch.Tensor | None = None,
     attn_sink: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Returns:
+    """Returns:
     - o: [s_q, h_q, dv]
     - o_fp32: [s_q, h_q, dv]
     - max_logits: [s_q, h_q]
     - lse: [s_q, h_q]
+
     """
     s_q, h_q, d_qk = q.shape
     s_kv, _, _ = kv.shape
diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py
index b1a52a5d557d..85ec4903ee06 100644
--- a/tests/kernels/core/test_activation.py
+++ b/tests/kernels/core/test_activation.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import random
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -236,6 +237,10 @@ def test_swiglu_limit_func_without_routing_uses_output_buffer() -> None:
     torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
 
 
+@pytest.mark.parametrize(
+    ("alpha", "beta"),
+    [(1.0, 0.0), (1.702, 0.0), (1.0, 1.0)],
+)
 @pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
 @pytest.mark.parametrize("d", D)
@@ -245,6 +250,8 @@ def test_swiglu_limit_func_without_routing_uses_output_buffer() -> None:
 @torch.inference_mode()
 def test_silu_and_mul_with_clamp(
     default_vllm_config,
+    alpha: float,
+    beta: float,
     swiglu_limit: float,
     num_tokens: int,
     d: int,
@@ -258,8 +265,38 @@ def test_silu_and_mul_with_clamp(
     # Use large values to ensure clamping is exercised.
     x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
 
-    layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
-    out = layer(x)
+    default_vllm_config.compilation_config.custom_ops = [
+        "none",
+        "+silu_and_mul_with_clamp",
+    ]
+    layer = SiluAndMulWithClamp(
+        swiglu_limit,
+        alpha=alpha,
+        beta=beta,
+        compile_native=False,
+    )
+    if current_platform.is_rocm():
+        # forward_hip is always dispatched; the alpha/beta gate is checked
+        # inside it at call time rather than picked at construction time, so
+        # verify the actual routing by spying on the two candidate methods.
+        assert layer._forward_method == layer.forward_hip
+        with (
+            patch.object(layer, "forward_cuda", wraps=layer.forward_cuda) as cuda_spy,
+            patch.object(
+                layer, "forward_native", wraps=layer.forward_native
+            ) as native_spy,
+        ):
+            out = layer(x)
+        if alpha == 1.0 and beta == 0.0:
+            cuda_spy.assert_called_once()
+            native_spy.assert_not_called()
+        else:
+            native_spy.assert_called_once()
+            cuda_spy.assert_not_called()
+    else:
+        assert layer._forward_method == layer.forward_cuda
+        out = layer(x)
+
     ref_out = layer.forward_native(x)
 
     rtol = {
@@ -273,10 +310,11 @@ def test_silu_and_mul_with_clamp(
 
     # Verify clamping is actually being applied: the clamped output should
     # differ from the unclamped SiluAndMul output when inputs are large.
-    unclamped_out = SiluAndMul.forward_native(x)
-    assert not torch.equal(ref_out.float(), unclamped_out.float()), (
-        "Input was not large enough to exercise the clamp; increase scale"
-    )
+    if alpha == 1.0 and beta == 0.0:
+        unclamped_out = SiluAndMul.forward_native(x)
+        assert not torch.equal(ref_out.float(), unclamped_out.float()), (
+            "Input was not large enough to exercise the clamp; increase scale"
+        )
 
     # Verify gate clamping semantics with a controlled scalar case.
     # gate=large_val is clamped to limit first, then silu(limit) * 1.0.
@@ -309,7 +347,10 @@ def test_silu_and_mul_with_clamp(
 
     # opcheck
     out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
-    opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
+    opcheck(
+        torch.ops._C.silu_and_mul_with_clamp,
+        (out_buf, x, swiglu_limit, layer.alpha, layer.beta),
+    )
 
 
 @pytest.mark.parametrize("linear_beta", [-1.0, 2.0])
@@ -352,49 +393,46 @@ def test_masked_situ_and_mul(
     )
 
 
-@pytest.mark.parametrize(
-    ("activation", "activation_config"),
-    [
-        (MoEActivation.SILU, ApplyMoEActivationConfig()),
-        (
-            MoEActivation.SILU,
-            ApplyMoEActivationConfig(clamp_limit=3.0),
+MOE_ACTIVATION_CASES = [
+    pytest.param(MoEActivation.SILU, ApplyMoEActivationConfig(), id="silu"),
+    pytest.param(
+        MoEActivation.SILU, ApplyMoEActivationConfig(clamp_limit=3.0), id="silu_clamp"
+    ),
+    pytest.param(MoEActivation.GELU, ApplyMoEActivationConfig(), id="gelu"),
+    pytest.param(MoEActivation.GELU_TANH, ApplyMoEActivationConfig(), id="gelu_tanh"),
+    pytest.param(
+        MoEActivation.SITU,
+        ApplyMoEActivationConfig(
+            activation_situ_beta=1.5,
+            activation_situ_linear_beta=2.0,
         ),
-        (MoEActivation.GELU, ApplyMoEActivationConfig()),
-        (MoEActivation.GELU_TANH, ApplyMoEActivationConfig()),
-        (
-            MoEActivation.SITU,
-            ApplyMoEActivationConfig(
-                activation_situ_beta=1.5,
-                activation_situ_linear_beta=2.0,
-            ),
-        ),
-        (MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig()),
-        (
-            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
-            ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5),
-        ),
-        (MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig()),
-        (MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig()),
-        (MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig()),
-        (MoEActivation.GELU_TANH_NO_MUL, ApplyMoEActivationConfig()),
-        (MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig()),
-    ],
-    ids=[
-        "silu",
-        "silu_clamp",
-        "gelu",
-        "gelu_tanh",
-        "situ",
-        "swigluoai",
-        "swigluoai_uninterleave",
-        "swiglustep",
-        "silu_no_mul",
-        "gelu_no_mul",
-        "gelu_tanh_no_mul",
-        "relu2_no_mul",
-    ],
-)
+        id="situ",
+    ),
+    pytest.param(MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig(), id="swigluoai"),
+    pytest.param(
+        MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+        ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5),
+        id="swigluoai_uninterleave",
+    ),
+    pytest.param(MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig(), id="swiglustep"),
+    pytest.param(
+        MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig(), id="silu_no_mul"
+    ),
+    pytest.param(
+        MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig(), id="gelu_no_mul"
+    ),
+    pytest.param(
+        MoEActivation.GELU_TANH_NO_MUL,
+        ApplyMoEActivationConfig(),
+        id="gelu_tanh_no_mul",
+    ),
+    pytest.param(
+        MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig(), id="relu2_no_mul"
+    ),
+]
+
+
+@pytest.mark.parametrize(("activation", "activation_config"), MOE_ACTIVATION_CASES)
 @torch.inference_mode()
 def test_masked_moe_activation_dispatch(
     default_vllm_config,
@@ -497,3 +535,73 @@ def test_activation(
 
     out = torch.empty_like(x)
     opcheck(fn, (out, x))
+
+
+HUMMING_ACTIVATION_CASES = MOE_ACTIVATION_CASES + [
+    pytest.param(MoEActivation.RELU2, ApplyMoEActivationConfig(), id="relu2"),
+]
+HUMMING_ACTIVATION_CASES += [
+    pytest.param(
+        MoEActivation.SITU,
+        ApplyMoEActivationConfig(
+            activation_situ_beta=1.5,
+            activation_situ_linear_beta=linear_beta,
+        ),
+        id=f"situ-linear-beta-{linear_beta}",
+    )
+    for linear_beta in (None, 0.0, -1.0)
+]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Humming requires CUDA")
+@pytest.mark.parametrize(("activation", "activation_config"), HUMMING_ACTIVATION_CASES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize(("num_tokens", "d"), [(1, 512), (7, 768), (83, 512)])
+@torch.inference_mode()
+def test_humming_activation_matches_framework(
+    activation: MoEActivation,
+    activation_config: ApplyMoEActivationConfig,
+    dtype: torch.dtype,
+    num_tokens: int,
+    d: int,
+) -> None:
+    """Compare activation math/layouts without quantization or Hadamard error."""
+    pytest.importorskip("humming")
+    from humming.ops import process_input
+
+    from vllm.model_executor.layers.quantization.utils.humming.activation import (
+        get_humming_activation,
+    )
+
+    set_random_seed(0)
+    width = d * 2 if activation.is_gated else d
+    x = 4 * torch.randn(num_tokens, width, dtype=dtype, device=CUDA_DEVICES[0])
+    # Exercise zero, saturation, and both sides of the clamp limits (3 and 7).
+    edges = x.new_tensor([0, 0.001, 1, 2.99, 3, 3.01, 6.99, 7, 7.01, 8, 16])
+    edges = torch.cat((-edges[1:].flip(0), edges))
+    x[0] = edges[torch.arange(width, device=x.device) % edges.numel()]
+
+    actual, _, _ = process_input(
+        x,
+        quant_mode="none",
+        hadamard_block_size=0,
+        **get_humming_activation(activation, activation_config),
+    )
+    expected = torch.empty(num_tokens, d, dtype=dtype, device=x.device)
+    if activation == MoEActivation.RELU2:
+        # apply_moe_activation has only the non-gated ReLU2 variant.
+        gate, up = x.float().chunk(2, dim=-1)
+        activated_gate = torch.empty_like(gate)
+        apply_moe_activation(MoEActivation.RELU2_NO_MUL, activated_gate, gate.clone())
+        expected.copy_(activated_gate * up)
+    else:
+        # The framework's non-gated ReLU2 path modifies its input in place.
+        apply_moe_activation(
+            activation, expected, x.clone(), activation_config=activation_config
+        )
+    torch.testing.assert_close(
+        actual,
+        expected,
+        atol=get_default_atol(expected),
+        rtol=get_default_rtol(expected),
+    )
diff --git a/tests/kernels/core/test_apply_rotary_emb.py b/tests/kernels/core/test_apply_rotary_emb.py
index 5c765bbb4a21..a576199663ea 100644
--- a/tests/kernels/core/test_apply_rotary_emb.py
+++ b/tests/kernels/core/test_apply_rotary_emb.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for ApplyRotaryEmb CustomOp dispatch behavior.
+"""Tests for ApplyRotaryEmb CustomOp dispatch behavior.
 
 This test ensures that RotaryEmbedding classes correctly call the appropriate
 ApplyRotaryEmb methods based on the calling context:
@@ -179,10 +178,50 @@ def test_rotary_embedding_dispatch(
     test_case: RotaryEmbeddingTestCase,
     device: str,
 ):
-    """
-    Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method.
+    """Test that RotaryEmbedding classes dispatch to the correct ApplyRotaryEmb method.
 
     - forward_native methods should call ApplyRotaryEmb.forward_native()
     - forward_cuda/forward methods should call ApplyRotaryEmb.forward()
     """
     run_dispatch_test(test_case, device)
+
+
+@pytest.mark.skipif(not current_platform.is_xpu(), reason="XPU only test.")
+@pytest.mark.parametrize(
+    "num_tokens,num_heads,head_size,rot_dim,is_neox_style",
+    [
+        (128, 32, 128, 128, True),  # typical text shape
+        (4096, 16, 80, 80, False),  # GPT-J style (interleaved)
+        (4096, 32, 128, 64, True),  # partial rotary (rot_dim < head_size)
+    ],
+    ids=["text", "gptj_interleaved", "partial_rotary"],
+)
+def test_apply_rotary_emb_xpu_matches_native(
+    num_tokens: int,
+    num_heads: int,
+    head_size: int,
+    rot_dim: int,
+    is_neox_style: bool,
+):
+    """ApplyRotaryEmb.forward_xpu (vllm_xpu_kernels SYCL op) must match
+    forward_native (the CustomOp default fallback) exactly."""
+    from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
+
+    vllm_config = VllmConfig(
+        compilation_config=CompilationConfig(custom_ops=["all", "+apply_rotary_emb"])
+    )
+    get_cached_compilation_config.cache_clear()
+
+    with set_current_vllm_config(vllm_config):
+        op = ApplyRotaryEmb(enforce_enable=True, is_neox_style=is_neox_style)
+
+        x = torch.randn(
+            num_tokens, num_heads, head_size, device="xpu", dtype=torch.bfloat16
+        )
+        x = x[..., :rot_dim].contiguous()
+        cos = torch.randn(num_tokens, rot_dim // 2, device="xpu", dtype=torch.bfloat16)
+        sin = torch.randn(num_tokens, rot_dim // 2, device="xpu", dtype=torch.bfloat16)
+
+        expected = op.forward_native(x, cos, sin)
+        actual = op.forward_xpu(x, cos, sin)
+        torch.testing.assert_close(actual, expected, atol=0, rtol=0)
diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py
index 8debe7c3a35b..3d0e2ac11e35 100644
--- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py
+++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py
@@ -112,3 +112,192 @@ def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
             rtol=1e-2,
             atol=1e-2,
         )
+
+
+@pytest.mark.skipif(
+    not current_platform.has_device_capability(100), reason="MXFP8 needs Blackwell"
+)
+@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 129, 1024])
+@pytest.mark.parametrize("q_size", [384, 1280, 1344])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_fused_q_kv_rmsnorm_quant_matches_separate(num_tokens, q_size, dtype):
+    """Preserve native MXFP8 rounding and all swizzled scale padding bytes."""
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+    )
+    from vllm.models.deepseek_v41.common.ops.query_quant import (
+        fused_q_kv_rmsnorm_quant,
+    )
+
+    torch.manual_seed(0)
+    kv_size = 512
+    x = torch.randn(num_tokens, q_size + kv_size, device="cuda", dtype=dtype)
+    x[:1] = 0
+    x[1:2] *= 1e-30
+    qr, kv = x.split([q_size, kv_size], -1)
+    qw = torch.randn(q_size, device="cuda", dtype=x.dtype)
+    kvw = torch.randn(kv_size, device="cuda", dtype=x.dtype)
+    eps = 1e-20
+    result = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, eps)
+    assert result[0].orig_dtype == dtype
+    assert result[0].orig_shape == qr.shape
+    if not num_tokens:
+        assert result[0].data.shape == qr.shape
+        assert result[0].scale.numel() == 0
+        assert result[1].shape == kv.shape
+        return
+    qr_ref, kv_ref = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
+    q_ref, scale_ref = mxfp8_e4m3_quantize(qr_ref, is_sf_swizzled_layout=True)
+    for replay in (False, True):
+        if replay:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                result = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, eps)
+            graph.replay()
+        torch.testing.assert_close(
+            result[0].data.view(torch.uint8), q_ref.view(torch.uint8), rtol=0, atol=0
+        )
+        torch.testing.assert_close(result[0].scale, scale_ref, rtol=0, atol=0)
+        torch.testing.assert_close(result[1], kv_ref, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.has_device_capability(100), reason="MXFP8 needs Blackwell"
+)
+@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass"])
+@pytest.mark.parametrize("num_tokens", [1, 17, 256])
+@pytest.mark.parametrize("projection", ["attention", "indexer"])
+@pytest.mark.parametrize("bias", [False, True])
+def test_shared_query_quant_preserves_projection(
+    num_tokens, backend, projection, bias, monkeypatch
+):
+    """The native-checkpoint projection keeps identical GEMM outputs."""
+    from dataclasses import replace
+    from types import SimpleNamespace
+    from typing import Any, cast
+
+    from vllm.config import (
+        CompilationConfig,
+        CompilationMode,
+        VllmConfig,
+        set_current_vllm_config,
+    )
+    from vllm.model_executor.kernels.linear.mxfp8 import Mxfp8LinearLayerConfig
+    from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+        FlashInferCutedslMxfp8LinearKernel,
+        FlashInferCutlassMxfp8LinearKernel,
+    )
+    from vllm.model_executor.layers.fusion.quant_activation import (
+        expose_input_quant_key,
+    )
+    from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear
+    from vllm.model_executor.layers.quantization.modelopt import (
+        ModelOptLinearMethod,
+    )
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+    )
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic
+    from vllm.models.deepseek_v41.attention import (
+        DeepseekV4Attention,
+        DeepseekV4Indexer,
+    )
+    from vllm.models.deepseek_v41.common.ops.query_quant import (
+        can_fuse_query_quant,
+        fused_q_kv_rmsnorm_quant,
+    )
+    from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config
+
+    cls: (
+        type[FlashInferCutedslMxfp8LinearKernel]
+        | type[FlashInferCutlassMxfp8LinearKernel]
+    )
+    if backend == "cute-dsl":
+        cls = FlashInferCutedslMxfp8LinearKernel
+    else:
+        cls = FlashInferCutlassMxfp8LinearKernel
+    supported, reason = cls.is_supported()
+    if not supported:
+        pytest.skip(reason)
+    torch.manual_seed(0)
+    with set_current_vllm_config(
+        VllmConfig(compilation_config=CompilationConfig(mode=CompilationMode.NONE))
+    ):
+        kernel = cls(Mxfp8LinearLayerConfig())
+        # Match the checkpoint's TP4 attention and replicated indexer widths.
+        if projection == "attention":
+            linear = ColumnParallelLinear.__new__(ColumnParallelLinear)
+            output_size = 8192
+        else:
+            linear = ReplicatedLinear.__new__(ReplicatedLinear)
+            output_size = 4096
+        torch.nn.Module.__init__(linear)
+        linear.bias = (
+            torch.nn.Parameter(
+                torch.randn(output_size, device="cuda", dtype=torch.bfloat16)
+            )
+            if bias
+            else None
+        )
+        linear.skip_bias_add = False
+        linear.return_bias = projection == "indexer"
+        linear.gather_output = False
+        w, ws = mxfp8_e4m3_quantize(
+            torch.randn(output_size, 1280, device="cuda", dtype=torch.bfloat16)
+        )
+        linear.weight = torch.nn.Parameter(w, requires_grad=False)
+        linear.weight_scale = torch.nn.Parameter(ws, requires_grad=False)
+        quant_config = DeepseekV4FP8Config.from_config(
+            {
+                "quant_method": "fp8",
+                "activation_scheme": "dynamic",
+                "weight_block_size": [32, 32],
+                "scale_fmt": "ue8m0",
+                "expert_dtype": "fp4",
+            }
+        )
+        method = quant_config.get_quant_method(linear, "model.layers.2.attn.wq_b")
+        assert type(method) is ModelOptLinearMethod
+        method.out_dtype = torch.bfloat16
+        method.kernel = kernel
+        linear.quant_method = method
+        kernel.process_weights_after_loading(linear)
+        expose_input_quant_key(linear, kernel)
+        assert can_fuse_query_quant([linear, linear])
+        x = torch.randn(num_tokens, 1792, device="cuda", dtype=torch.bfloat16)
+        qr, kv = x.split([1280, 512], -1)
+        qw = torch.randn(1280, device="cuda", dtype=x.dtype)
+        kvw = torch.randn(512, device="cuda", dtype=x.dtype)
+        qr_ref, _ = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-20)
+        q, _ = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, 1e-20)
+        owner = cast(Any, SimpleNamespace(wq_b=linear))
+        project = (
+            DeepseekV4Attention._wq_b_proj
+            if projection == "attention"
+            else DeepseekV4Indexer._wq_b_proj
+        )
+        expected = project(owner, qr_ref)
+
+        def reject_requantization(*args, **kwargs):
+            raise AssertionError("pre-quantized input must skip quantization")
+
+        monkeypatch.setattr(
+            "vllm.model_executor.kernels.linear.mxfp8.flashinfer.mxfp8_e4m3_quantize",
+            reject_requantization,
+        )
+        actual = project(owner, q)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = project(owner, q)
+        graph.replay()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        assert not can_fuse_query_quant(
+            [linear, torch.nn.Linear(1280, 256, bias=False)]
+        )
+        reshaped = project(
+            owner, replace(q, orig_shape=torch.Size((1, num_tokens, 1280)))
+        )
+        torch.testing.assert_close(reshaped, expected.unsqueeze(0), rtol=0, atol=0)
+        with pytest.raises(AssertionError, match="consumer kernel"):
+            linear(replace(q, quant_key=kNvfp4Dynamic))
diff --git a/tests/kernels/core/test_fused_quant_layernorm.py b/tests/kernels/core/test_fused_quant_layernorm.py
index 201830fe2ff2..7a6e53314f06 100644
--- a/tests/kernels/core/test_fused_quant_layernorm.py
+++ b/tests/kernels/core/test_fused_quant_layernorm.py
@@ -318,10 +318,13 @@ def test_rms_norm(
     # Per-block bf16 scales: allow a small relative tolerance for a few groups
     # whose abs-max flips by one ULP between the fused and reference paths. The
     # per-token and fp32 paths stay strict.
-    relax_block_rocm = (
+    # The same one-ULP group-scale flip also occurs on CUDA (H200), so extend
+    # the block relaxation there too — the fused groupwise reduction rounds
+    # differently than the reference for isolated groups.
+    relax_block = (
         group_size is not None
         and dtype == torch.bfloat16
-        and current_platform.is_rocm()
+        and (current_platform.is_rocm() or current_platform.is_cuda())
     )
     use_gfx950_fp8_allclose = (
         current_platform.is_rocm()
@@ -340,7 +343,7 @@ def test_rms_norm(
     def scales_close(rtol: float, atol: float) -> bool:
         if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol):
             return True
-        return relax_block_rocm and torch.allclose(
+        return relax_block and torch.allclose(
             ref_scales, ops_scales, rtol=1e-2, atol=atol
         )
 
@@ -354,7 +357,7 @@ def scales_close(rtol: float, atol: float) -> bool:
         b = ops_out.to(dtype=torch.float32)
         ok = torch.allclose(a, b, atol=1e-6)
         if not ok:
-            if relax_block_rocm:
+            if relax_block:
                 # ULP-flipped group scale can cross an E4M3 tie; tolerate a
                 # bounded count of isolated fp8 outliers.
                 ulp = fp8_ulp_distance(ref_out, ops_out)
diff --git a/tests/kernels/core/test_minimax_reduce_rms.py b/tests/kernels/core/test_minimax_reduce_rms.py
index cd8efc14a3c2..b92a08a317dc 100644
--- a/tests/kernels/core/test_minimax_reduce_rms.py
+++ b/tests/kernels/core/test_minimax_reduce_rms.py
@@ -36,7 +36,6 @@ def _worker_forward_qk(
     eps,
 ):
     """Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel."""
-
     if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
         cleanup_dist_env_and_memory()
         return
diff --git a/tests/kernels/core/test_mrope.py b/tests/kernels/core/test_mrope.py
index 81e0a916a304..0e355d7cd352 100644
--- a/tests/kernels/core/test_mrope.py
+++ b/tests/kernels/core/test_mrope.py
@@ -229,7 +229,7 @@ def test_mrope_torch_compile_tracing(
 
     # Create a wrapper that makes the in-place function appear functional
     def functional_forward_cuda(pos, q, k):
-        """Wrapper that converts in-place operation to functional style
+        """Wrapper that converts in-place operation to functional style.
 
         CUDA Graph does not support in-place operations.
         This wrapper creates working copies of the
diff --git a/tests/kernels/core/test_opcheck.py b/tests/kernels/core/test_opcheck.py
index 40ced08b933a..45dd01d44e50 100644
--- a/tests/kernels/core/test_opcheck.py
+++ b/tests/kernels/core/test_opcheck.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for miscellaneous utilities
-"""
+"""Tests for miscellaneous utilities."""
 
 import torch
 
diff --git a/tests/kernels/core/test_rotary_embedding.py b/tests/kernels/core/test_rotary_embedding.py
index 1cbb5dbd1881..353aefa11aae 100644
--- a/tests/kernels/core/test_rotary_embedding.py
+++ b/tests/kernels/core/test_rotary_embedding.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for miscellaneous utilities
-"""
+"""Tests for miscellaneous utilities."""
 
 import pytest
 import torch
diff --git a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
index 2a2525c307bc..3f192ba2e7f1 100644
--- a/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
+++ b/tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for fused MLA KV-cache write and RoPE fused kernel
-"""
+"""Tests for fused MLA KV-cache write and RoPE fused kernel."""
 
 import random
 
@@ -17,6 +15,88 @@
 from vllm.utils.torch_utils import set_random_seed
 
 
+def assert_rope_close_rocm(
+    actual: torch.Tensor,
+    ref: torch.Tensor,
+    rope_input: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    is_neox: bool,
+    *,
+    max_fail_fraction: float | None = None,
+    atol: float | None = None,
+    rtol: float | None = None,
+) -> None:
+    """assert_close for RoPE fused-vs-reference on ROCm, with a per-element
+    product-ULP allowance for catastrophic-cancellation elements.
+
+    The fused and reference kernels may contract the RoPE multiply-add into an
+    FMA with either product pre-rounded to the working dtype. The two valid
+    orderings differ by at most 2 ULPs of the product at that element, which
+    cancellation (x*cos ~= -y*sin) can amplify past the default tolerance.
+
+    Tier 1: |actual - ref| <= atol + rtol * |ref|  (default tolerance).
+    Tier 2: elements failing tier 1 must each satisfy
+    |actual_e - ref_e| <= 2 * eps * |product_e|, where product_e is the larger
+    of the two products combined at that element, and their fraction must be
+    <= max_fail_fraction (1e-4 for fp16; 1e-2 for bf16, whose coarser eps
+    puts 1-ULP divergences past the default tolerance far more often).
+    """
+    if atol is None:
+        atol = get_default_atol(ref)
+    if rtol is None:
+        rtol = get_default_rtol(ref)
+    if max_fail_fraction is None:
+        max_fail_fraction = 1e-2 if ref.dtype == torch.bfloat16 else 1e-4
+    eps = torch.finfo(ref.dtype).eps
+
+    # per-element product magnitude, in the (seq, heads, rot_dim) layout
+    seq_len, rot_dim = rope_input.shape[0], rope_input.shape[-1]
+    half = rot_dim // 2
+    inp = rope_input.reshape(seq_len, -1, rot_dim).to(torch.float32)
+    cs = cos_sin_cache[positions].to(torch.float32)  # (seq_len, 2*half)
+    cos = cs[:, :half].unsqueeze(1)  # broadcast over heads
+    sin = cs[:, half:].unsqueeze(1)
+    x, y = inp[..., :half], inp[..., half:]
+    if is_neox:
+        # out = [x*cos - y*sin, y*cos + x*sin]
+        prod = torch.cat(
+            [
+                torch.maximum((x * cos).abs(), (y * sin).abs()),
+                torch.maximum((y * cos).abs(), (x * sin).abs()),
+            ],
+            dim=-1,
+        )
+    else:
+        # out[2i] = x[2i]*cos[i] - x[2i+1]*sin[i], interleaved
+        xe, xo = inp[..., 0::2], inp[..., 1::2]
+        prod = torch.stack(
+            [
+                torch.maximum((xe * cos).abs(), (xo * sin).abs()),
+                torch.maximum((xo * cos).abs(), (xe * sin).abs()),
+            ],
+            dim=-1,
+        ).flatten(-2)
+    prod = prod.reshape(rope_input.shape)
+
+    diff = (actual.to(torch.float32) - ref.to(torch.float32)).abs()
+    tol = atol + rtol * ref.to(torch.float32).abs()
+    fail = diff > tol
+    n_fail = int(fail.sum())
+    if n_fail == 0:
+        return
+    assert n_fail <= max_fail_fraction * diff.numel(), (
+        f"{n_fail} elements beyond default tolerance "
+        f"({n_fail / diff.numel():.2%} > {max_fail_fraction:.2%}); "
+        f"max diff {diff.max().item()}"
+    )
+    excess = diff[fail] - 2 * eps * prod[fail]
+    assert (excess <= 0).all(), (
+        f"{int((excess > 0).sum())} cancellation elements exceed 2 product-ULPs; "
+        f"worst excess {excess.max().item()}"
+    )
+
+
 @pytest.fixture
 def default_vllm_config(monkeypatch):
     """Enable the AITER triton rope on ROCm for fp16-consistent numerics.
@@ -99,6 +179,8 @@ def test_concat_and_cache_mla_rope_fused(
 
     k_pe = torch.flatten(key[..., :qk_rope_head_dim], start_dim=1).to(device=device)
     kv_c = torch.flatten(key[..., qk_rope_head_dim:], start_dim=1).to(device=device)
+    k_pe_orig = k_pe.clone()
+    query_orig = query.clone()
 
     if current_platform.is_rocm():
         # We use forward_hip for the same numerics as the fused custom kernel on ROCm
@@ -181,12 +263,9 @@ def test_concat_and_cache_mla_rope_fused(
         kv_cache_scale,
     )
 
-    # On ROCm the AITER Triton rope diverges by ~1 ULP from the fused kernel,
-    # which the tight CUDA tolerances don't cover for the low-precision paths:
-    #  - bf16: up to ~1 bf16 ULP (0.0156-0.03125 at values ~2-3), so relax the
-    #    kv-cache atol (bounded ~0.02) and the query atol (~0.04).
-    #  - neox-style fp8: one e4m3 ULP (~12.5%), so relax rtol.
-    # Other paths use the CUDA defaults.
+    # On ROCm the AITER Triton rope diverges by up to 2 product-ULPs from the
+    # fused kernel, which the default tolerance misses at the rare
+    # catastrophic-cancellation elements. fp8 keeps its calibrated tolerance.
     rocm_neox = current_platform.is_rocm() and is_neox_style
     rocm_bf16 = current_platform.is_rocm() and dtype == torch.bfloat16
     if kv_cache_dtype == "fp8":
@@ -207,16 +286,40 @@ def test_concat_and_cache_mla_rope_fused(
             atol=0.004 if rocm_bf16 else 0.001,
             rtol=0.15 if rocm_neox or rocm_bf16 else 0.1,
         )
-    elif rocm_bf16:
-        torch.testing.assert_close(kv_cache, ref_kv_cache, atol=0.02, rtol=1e-3)
-    elif rocm_neox:
-        torch.testing.assert_close(kv_cache, ref_kv_cache, atol=1e-3, rtol=1e-3)
+    elif current_platform.is_rocm() and dtype in (torch.half, torch.bfloat16):
+        # gather the rope region of the cache back to token order
+        block_idx = slot_mapping // block_size
+        block_offset = slot_mapping % block_size
+        fused_k_rope = kv_cache[block_idx, block_offset, kv_lora_rank:]
+        ref_k_rope_gathered = ref_kv_cache[block_idx, block_offset, kv_lora_rank:]
+        assert_rope_close_rocm(
+            fused_k_rope,
+            ref_k_rope_gathered,
+            k_pe_orig,
+            rope.cos_sin_cache,
+            positions,
+            is_neox_style,
+        )
+        # the nope region is a byte-exact copy
+        torch.testing.assert_close(
+            kv_cache[..., :kv_lora_rank], ref_kv_cache[..., :kv_lora_rank]
+        )
     else:
         torch.testing.assert_close(kv_cache, ref_kv_cache)
 
-    torch.testing.assert_close(
-        query,
-        ref_q_pe,
-        atol=0.04 if rocm_bf16 else get_default_atol(query),
-        rtol=get_default_rtol(query),
-    )
+    if current_platform.is_rocm() and dtype in (torch.half, torch.bfloat16):
+        assert_rope_close_rocm(
+            query,
+            ref_q_pe,
+            query_orig,
+            rope.cos_sin_cache,
+            positions,
+            is_neox_style,
+        )
+    else:
+        torch.testing.assert_close(
+            query,
+            ref_q_pe,
+            atol=get_default_atol(query),
+            rtol=get_default_rtol(query),
+        )
diff --git a/tests/kernels/helion/test_config_manager.py b/tests/kernels/helion/test_config_manager.py
index f8e5eae6f106..63f813084926 100644
--- a/tests/kernels/helion/test_config_manager.py
+++ b/tests/kernels/helion/test_config_manager.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for Helion ConfigManager and ConfigSet.
+"""Unit tests for Helion ConfigManager and ConfigSet.
 
 Tests the simplified configuration management system for Helion custom kernels.
 """
diff --git a/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py b/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py
index 6f5a40fad861..14a4356e14a2 100644
--- a/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py
+++ b/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel
+"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel.
 
 Run `pytest tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py`.
 """
diff --git a/tests/kernels/helion/test_fused_qk_norm_rope.py b/tests/kernels/helion/test_fused_qk_norm_rope.py
index 4ced522d9193..a973533169a9 100644
--- a/tests/kernels/helion/test_fused_qk_norm_rope.py
+++ b/tests/kernels/helion/test_fused_qk_norm_rope.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the fused_qk_norm_rope helion kernel
+"""Tests for the fused_qk_norm_rope helion kernel.
 
 Run `pytest tests/kernels/helion/test_fused_qk_norm_rope.py`.
 """
diff --git a/tests/kernels/helion/test_helion_available.py b/tests/kernels/helion/test_helion_available.py
index 560a671f149b..a4de68bd6716 100644
--- a/tests/kernels/helion/test_helion_available.py
+++ b/tests/kernels/helion/test_helion_available.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for Helion kernel availability and basic functionality.
+"""Tests for Helion kernel availability and basic functionality.
 
 This module demonstrates the pattern for testing optional Helion kernels.
 Tests in this directory will be skipped if Helion is not installed.
diff --git a/tests/kernels/helion/test_per_token_group_fp8_quant.py b/tests/kernels/helion/test_per_token_group_fp8_quant.py
index d9d96f9287c2..d9dd8a881763 100644
--- a/tests/kernels/helion/test_per_token_group_fp8_quant.py
+++ b/tests/kernels/helion/test_per_token_group_fp8_quant.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the per_token_group_fp8_quant helion kernel
+"""Tests for the per_token_group_fp8_quant helion kernel.
 
 Run `pytest tests/kernels/helion/test_per_token_group_fp8_quant.py`.
 """
diff --git a/tests/kernels/helion/test_register.py b/tests/kernels/helion/test_register.py
index 9876135056b2..dd4274bfafd7 100644
--- a/tests/kernels/helion/test_register.py
+++ b/tests/kernels/helion/test_register.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for Helion kernel registration.
+"""Unit tests for Helion kernel registration.
 
 Tests ConfiguredHelionKernel, HelionKernelWrapper, and PresetConfigSearch
 including config picker registration and custom autotuner integration.
@@ -1020,7 +1019,6 @@ def f(x, y):
     )
     def test_inductor_backend_compiles_helion_hop(self):
         """Test torch.compile with inductor backend and Helion fusion enabled."""
-
         configs: dict[CaseKey, helion.Config] = {
             CaseKey.default(): helion.Config(block_sizes=[4, 4])
         }
diff --git a/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py b/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py
index 50591ae34a61..d0650dfc7afe 100644
--- a/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py
+++ b/tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the rms_norm_dynamic_per_token_quant helion kernel
+"""Tests for the rms_norm_dynamic_per_token_quant helion kernel.
 
 Run `pytest tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py`.
 """
diff --git a/tests/kernels/helion/test_rms_norm_per_block_quant.py b/tests/kernels/helion/test_rms_norm_per_block_quant.py
index 4ffc759c63d2..e9c003f4e303 100644
--- a/tests/kernels/helion/test_rms_norm_per_block_quant.py
+++ b/tests/kernels/helion/test_rms_norm_per_block_quant.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the rms_norm_per_block_quant helion kernel
+"""Tests for the rms_norm_per_block_quant helion kernel.
 
 Run `pytest tests/kernels/helion/test_rms_norm_per_block_quant.py`.
 """
diff --git a/tests/kernels/helion/utils.py b/tests/kernels/helion/utils.py
index 38893fc8fecb..4a282f275d10 100644
--- a/tests/kernels/helion/utils.py
+++ b/tests/kernels/helion/utils.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Helion Kernel test utils"""
+"""Helion Kernel test utils."""
 
 import pytest
 import torch
diff --git a/tests/kernels/ir/test_activation.py b/tests/kernels/ir/test_activation.py
new file mode 100644
index 000000000000..a36ff8ee339b
--- /dev/null
+++ b/tests/kernels/ir/test_activation.py
@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+import vllm.kernels  # noqa: F401
+from tests.ir.ir_test_utils import assert_close, clone_args
+from vllm import ir
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
+
+gelu_and_mul_sparse_native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn
+
+
+def test_gelu_and_mul_sparse_registration():
+    expected = {
+        "native": True,
+        "triton": HAS_TRITON and current_platform.is_cuda(),
+    }
+    actual = {
+        provider: impl.supported
+        for provider, impl in ir.ops.gelu_and_mul_sparse.impls.items()
+    }
+    assert actual == expected
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("num_tokens", [1, 8, 64])
+@pytest.mark.parametrize("intermediate_size", [2048, 4096, 8192, 13824, 16384, 32768])
+def test_gelu_and_mul_sparse_triton(
+    dtype: torch.dtype, num_tokens: int, intermediate_size: int
+):
+    torch.manual_seed(0)
+    torch.set_default_device("cuda")
+    args = ir.ops.gelu_and_mul_sparse.generate_inputs(
+        num_tokens=num_tokens,
+        hidden_size=intermediate_size,
+        dtype=dtype,
+    )
+    impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
+    assert impl.supports_args(*args)
+
+    expected = gelu_and_mul_sparse_native(*clone_args(args))
+    actual = impl.impl_fn(*clone_args(args))
+    assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
+
+    mask_mismatch = torch.count_nonzero((actual == 0) != (expected == 0))
+    assert mask_mismatch / actual.numel() < 1e-3
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+def test_gelu_and_mul_sparse_triton_second_seed():
+    torch.manual_seed(1)
+    torch.set_default_device("cuda")
+    args = ir.ops.gelu_and_mul_sparse.generate_inputs(
+        num_tokens=32,
+        hidden_size=8192,
+        dtype=torch.bfloat16,
+    )
+
+    actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args))
+    expected = gelu_and_mul_sparse_native(*clone_args(args))
+    assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("intermediate_size", [1, 31, 128])
+def test_gelu_and_mul_sparse_triton_small_sizes(
+    dtype: torch.dtype, intermediate_size: int
+):
+    torch.manual_seed(0)
+    torch.set_default_device("cuda")
+    args = ir.ops.gelu_and_mul_sparse.generate_inputs(
+        num_tokens=5,
+        hidden_size=intermediate_size,
+        dtype=dtype,
+    )
+
+    actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args))
+    expected = gelu_and_mul_sparse_native(*clone_args(args))
+    assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
+    assert torch.equal(actual == 0, expected == 0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+def test_gelu_and_mul_sparse_triton_nd_and_empty():
+    torch.set_default_device("cuda")
+    impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
+
+    x = torch.randn(2, 3, 16384, dtype=torch.bfloat16)
+    args = (x, 1.6448536269514722, "tanh")
+    actual = impl.impl_fn(*args)
+    expected = gelu_and_mul_sparse_native(*args)
+    assert actual.shape == (2, 3, 8192)
+    assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
+
+    empty = torch.empty(0, 16384, dtype=torch.bfloat16)
+    empty_out = impl.impl_fn(empty, 1.6448536269514722, "tanh")
+    assert empty_out.shape == (0, 8192)
+    assert empty_out.dtype == empty.dtype
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_gelu_and_mul_sparse_triton_constant_rows(dtype: torch.dtype):
+    torch.set_default_device("cuda")
+    gate = torch.full((4, 8192), 3.0, dtype=dtype)
+    up = torch.randn_like(gate)
+    x = torch.cat((gate, up), dim=-1)
+    args = (x, 1.6448536269514722, "tanh")
+
+    actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args)
+    expected = gelu_and_mul_sparse_native(*args)
+    torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
+    assert torch.count_nonzero(actual) == 0
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")])
+def test_gelu_and_mul_sparse_triton_nonfinite_gate(nonfinite: float):
+    torch.set_default_device("cuda")
+    x = torch.randn(2, 16384, dtype=torch.bfloat16)
+    x[0, 0] = nonfinite
+    args = (x, 1.6448536269514722, "tanh")
+
+    actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args)
+    expected = gelu_and_mul_sparse_native(*args)
+    torch.testing.assert_close(actual, expected, equal_nan=True)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+def test_gelu_and_mul_sparse_triton_fallback_conditions():
+    torch.set_default_device("cuda")
+    impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
+    x = torch.randn(4, 16384, dtype=torch.bfloat16)
+
+    assert impl.supports_args(x, 1.0, "tanh")
+    assert not impl.supports_args(x, 1.0, "none")
+    assert not impl.supports_args(x[:, ::2], 1.0, "tanh")
+    assert not impl.supports_args(x[:, :-1], 1.0, "tanh")
+    assert not impl.supports_args(
+        torch.randn(1, 65538, dtype=torch.bfloat16), 1.0, "tanh"
+    )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+def test_gelu_and_mul_sparse_triton_dispatch_and_opcheck():
+    torch.set_default_device("cuda")
+    args = ir.ops.gelu_and_mul_sparse.generate_inputs(
+        num_tokens=8,
+        hidden_size=4096,
+        dtype=torch.bfloat16,
+    )
+    impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
+
+    with ir.ops.gelu_and_mul_sparse.set_priority(["triton", "native"]):
+        dispatched = ir.ops.gelu_and_mul_sparse(*args)
+        torch.library.opcheck(torch.ops.vllm_ir.gelu_and_mul_sparse, args)
+
+    direct = impl.impl_fn(*args)
+    torch.testing.assert_close(dispatched, direct, rtol=0.0, atol=0.0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
+def test_gelu_and_mul_sparse_triton_cudagraph():
+    torch.set_default_device("cuda")
+    x = torch.randn(16, 32768, dtype=torch.bfloat16)
+    impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
+
+    for _ in range(3):
+        expected = impl.impl_fn(x, 1.6448536269514722, "tanh")
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = impl.impl_fn(x, 1.6448536269514722, "tanh")
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
diff --git a/tests/kernels/ir/test_aiter_norm_dispatch.py b/tests/kernels/ir/test_aiter_norm_dispatch.py
new file mode 100644
index 000000000000..f84a7615b167
--- /dev/null
+++ b/tests/kernels/ir/test_aiter_norm_dispatch.py
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the AITER norm support predicates.
+
+The AITER norm wrappers flatten >2-D activations with ``Tensor.reshape``, which
+silently copies when the flattened shape is not expressible with the input's
+strides. The support predicates must reject those inputs so dispatch falls
+through to a provider that handles arbitrary strides.
+"""
+
+import pytest
+import torch
+
+from vllm._aiter_ops import is_aiter_found_and_supported
+from vllm.kernels.aiter_ops import flatten_to_2d_is_free
+
+pytestmark = pytest.mark.skipif(
+    not is_aiter_found_and_supported(),
+    reason="Only test on ROCm with AITER installed and supported",
+)
+
+
+def _qkv_slice_by_head(num_tokens, num_q_heads, num_kv_heads, head_dim):
+    """Q viewed per-head from a fused QKV projection, as Qwen3-style QK-norm does."""
+    q_size, kv_size = num_q_heads * head_dim, num_kv_heads * head_dim
+    qkv = torch.empty(num_tokens, q_size + 2 * kv_size)
+    q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0]
+    return q.view(*q.shape[:-1], num_q_heads, head_dim)
+
+
+@pytest.mark.parametrize(
+    "x",
+    [
+        torch.empty(8, 16),
+        torch.empty(16),
+        torch.empty(2, 4, 16),
+        torch.empty(2, 1, 16),
+        # A contiguous tensor stays flattenable after a leading-dim slice.
+        torch.empty(8, 4, 16)[2:6],
+    ],
+)
+def test_flattenable(x):
+    assert flatten_to_2d_is_free(x)
+    assert x.reshape(-1, x.shape[-1]).data_ptr() == x.data_ptr()
+
+
+@pytest.mark.parametrize(
+    "x",
+    [
+        _qkv_slice_by_head(8, 32, 8, 128),
+        # Non-unit last-dim stride.
+        torch.empty(4, 8, 32).transpose(-1, -2),
+        # Leading dims cannot be merged: a slice along the middle dim.
+        torch.empty(4, 8, 32)[:, :4],
+    ],
+)
+def test_not_flattenable(x):
+    assert not flatten_to_2d_is_free(x)
+    # reshape has to copy, which is exactly what the predicate guards against.
+    assert x.reshape(-1, x.shape[-1]).data_ptr() != x.data_ptr()
diff --git a/tests/kernels/ir/test_ir_ops.py b/tests/kernels/ir/test_ir_ops.py
index 730db9833b1d..560d047c0f01 100644
--- a/tests/kernels/ir/test_ir_ops.py
+++ b/tests/kernels/ir/test_ir_ops.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Meta-tests for vLLM IR op infrastructure.
+"""Meta-tests for vLLM IR op infrastructure.
 
 Ensures all registered ops have input generators defined.
 Per-op correctness tests live alongside their op definitions
diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
index c239f831c614..1ea264239cb9 100644
--- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
@@ -473,6 +473,20 @@ def _conv_inputs(total_tokens: int):
     return x, weight, bias
 
 
+def _conv_states(
+    num_slots: int,
+    state_len: int,
+    dim: int = CONV_DIM,
+    layout: str = "SD",
+) -> torch.Tensor:
+    if layout == "SD":
+        return torch.zeros(num_slots, state_len, dim, dtype=torch.bfloat16).transpose(
+            1, 2
+        )
+    assert layout == "DS"
+    return torch.zeros(num_slots, dim, state_len, dtype=torch.bfloat16)
+
+
 def _sd_conv_states(
     num_slots: int, state_len: int, dim: int = CONV_DIM
 ) -> torch.Tensor:
@@ -484,6 +498,39 @@ def _maybe_pack_conv_weight(weight: torch.Tensor, is_vnni: bool) -> torch.Tensor
     return ops.causal_conv1d_weight_pack(weight) if is_vnni else weight
 
 
+@torch.inference_mode()
+def test_causal_conv1d_update_cpu_preserves_sd_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """SD state and indices reach the native operator unchanged."""
+    num_slots = 2
+    state_len = CONV_KERNEL - 1
+    conv_state = _sd_conv_states(num_slots, state_len)
+    state_indices = torch.tensor([1], dtype=torch.int32)
+    x = torch.zeros(1, CONV_DIM, dtype=torch.bfloat16)
+    recorded = {}
+
+    def record_update(**kwargs):
+        recorded.update(kwargs)
+        return x
+
+    monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False)
+    monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_update_cpu", record_update)
+    out = gdn_attention._causal_conv1d_update_cpu(
+        x=x,
+        conv_states=conv_state,
+        weight=torch.zeros(CONV_DIM, CONV_KERNEL, dtype=torch.bfloat16),
+        bias=None,
+        silu_activation=False,
+        conv_state_indices=state_indices,
+        is_vnni=False,
+    )
+
+    assert out is x
+    assert recorded["conv_states"] is conv_state
+    assert recorded["conv_state_indices"] is state_indices
+
+
 @torch.inference_mode()
 def test_spec_aware_mixed_routing_preserves_token_order(
     monkeypatch: pytest.MonkeyPatch,
@@ -1123,6 +1170,145 @@ def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False):
     return out.transpose(0, 1).contiguous(), conv_state
 
 
+def _python_native_fwd(**kwargs):
+    from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
+        causal_conv1d_fn_cpu,
+    )
+
+    return causal_conv1d_fn_cpu(
+        x=kwargs["x"],
+        weight=kwargs["weight"],
+        bias=kwargs["bias"],
+        conv_states=kwargs["conv_states"],
+        query_start_loc=kwargs["query_start_loc"],
+        cache_indices=kwargs["cache_indices"],
+        has_initial_state=kwargs["has_initial_state"],
+        activation=kwargs["silu_activation"],
+    )
+
+
+@torch.inference_mode()
+def test_causal_conv1d_fwd_cpu_ds_adapter_varlen_mixed_initial_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    dim = 8
+    seq_lens = [2, 3]
+    state_len = _STATE_LEN
+    state_indices = torch.tensor([2, 0], dtype=torch.int32)
+    query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32)
+    has_initial_state = torch.tensor([False, True])
+    x_tokens, weight, bias = _conv_inputs(sum(seq_lens))
+    x_tokens = x_tokens[:, :dim]
+    weight = weight[:dim]
+    bias = bias[:dim]
+    x = x_tokens.transpose(0, 1).contiguous()
+    logical_state = (
+        torch.arange(3 * dim * state_len, dtype=torch.float32)
+        .view(3, dim, state_len)
+        .to(torch.bfloat16)
+    )
+    state_sd = _conv_states(3, state_len, dim, "SD")
+    state_ds = _conv_states(3, state_len, dim, "DS")
+    state_sd.copy_(logical_state)
+    state_ds.copy_(logical_state)
+    recorded = []
+
+    def record_fwd(**kwargs):
+        recorded.append(
+            (
+                kwargs["conv_states"].shape,
+                kwargs["conv_states"].stride(),
+                kwargs["cache_indices"].clone(),
+            )
+        )
+        return _python_native_fwd(**kwargs)
+
+    monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_fwd_cpu", record_fwd)
+    monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False)
+    out_sd = gdn_attention._causal_conv1d_fwd_cpu(
+        x=x,
+        weight=weight,
+        bias=bias,
+        conv_states=state_sd,
+        query_start_loc=query_start_loc,
+        cache_indices=state_indices,
+        has_initial_state=has_initial_state,
+        silu_activation=True,
+        is_vnni=False,
+    )
+    monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True)
+    out_ds = gdn_attention._causal_conv1d_fwd_cpu(
+        x=x,
+        weight=weight,
+        bias=bias,
+        conv_states=state_ds,
+        query_start_loc=query_start_loc,
+        cache_indices=state_indices,
+        has_initial_state=has_initial_state,
+        silu_activation=True,
+        is_vnni=False,
+    )
+
+    torch.testing.assert_close(out_ds, out_sd)
+    torch.testing.assert_close(state_ds, state_sd)
+    assert recorded[0][2].equal(state_indices)
+    assert recorded[1][0] == (2, dim, state_len)
+    assert recorded[1][1] == (dim * state_len, 1, dim)
+    assert recorded[1][2].equal(torch.arange(2, dtype=torch.int32))
+
+
+@torch.inference_mode()
+def test_causal_conv1d_update_cpu_ds_adapter_wide_multi_token_copyback(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    dim = 8
+    seq_len = 4
+    state_len = seq_len + 5
+    state = _conv_states(4, state_len, dim, "DS")
+    state.copy_(
+        torch.arange(state.numel(), dtype=torch.float32).view_as(state).to(state.dtype)
+    )
+    original_state = state.clone()
+    state_indices = torch.tensor([3, 1, 0], dtype=torch.int32)
+    num_accepted_tokens = torch.tensor([1, 4, 2], dtype=torch.int32)
+    mutated_scratch = torch.arange(
+        state_indices.numel() * dim * state_len, dtype=torch.float32
+    ).view(state_indices.numel(), dim, state_len)
+    mutated_scratch = (mutated_scratch + 200).to(state.dtype)
+    recorded = {}
+
+    def mutate_update(**kwargs):
+        recorded.update(kwargs)
+        kwargs["conv_states"].copy_(mutated_scratch)
+        return torch.zeros_like(kwargs["x"])
+
+    monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True)
+    monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_update_cpu", mutate_update)
+    gdn_attention._causal_conv1d_update_cpu(
+        x=torch.zeros(state_indices.numel(), seq_len, dim, dtype=state.dtype),
+        conv_states=state,
+        weight=torch.zeros(dim, CONV_KERNEL, dtype=state.dtype),
+        bias=None,
+        silu_activation=False,
+        conv_state_indices=state_indices,
+        is_vnni=False,
+        num_accepted_tokens=num_accepted_tokens,
+    )
+
+    expected_state = original_state.clone()
+    expected_state.index_copy_(0, state_indices.to(torch.long), mutated_scratch)
+    torch.testing.assert_close(state, expected_state)
+    torch.testing.assert_close(
+        state.index_select(0, state_indices)[:, :, _STATE_LEN:],
+        mutated_scratch[:, :, _STATE_LEN:],
+    )
+    torch.testing.assert_close(
+        recorded["conv_state_indices"],
+        torch.arange(state_indices.numel(), dtype=torch.int32),
+    )
+    torch.testing.assert_close(recorded["num_accepted_tokens"], num_accepted_tokens)
+
+
 @pytest.mark.skipif(
     not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
 )
diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py
index d2a2981edf2c..d4f5fac59ca6 100644
--- a/tests/kernels/mamba/test_causal_conv1d.py
+++ b/tests/kernels/mamba/test_causal_conv1d.py
@@ -36,8 +36,7 @@ def causal_conv1d_ref(
     final_states_out: torch.Tensor | None = None,
     activation: str | None = "silu",
 ):
-    """
-    x: (batch, dim, seqlen)
+    """x: (batch, dim, seqlen)
     weight: (dim, width)
     bias: (dim,)
     initial_states: (batch, dim, width - 1)
@@ -72,8 +71,7 @@ def causal_conv1d_ref(
 def causal_conv1d_update_ref(
     x, conv_state, weight, bias=None, activation=None, cache_seqlens=None
 ):
-    """
-    x: (batch, dim) or (batch, dim, seqlen)
+    """x: (batch, dim) or (batch, dim, seqlen)
     conv_state: (batch, dim, state_len), where state_len >= width - 1
     weight: (dim, width)
     bias: (dim,)
@@ -136,8 +134,7 @@ def causal_conv1d_opcheck_fn(
     activation: str | None = "silu",
     null_block_id: int = NULL_BLOCK_ID,
 ):
-    """
-    x: (batch, dim, seqlen)
+    """x: (batch, dim, seqlen)
     weight: (dim, width)
     bias: (dim,)
     seq_idx: (batch, seqlen)
diff --git a/tests/kernels/mamba/test_gdn_output_projection.py b/tests/kernels/mamba/test_gdn_output_projection.py
new file mode 100644
index 000000000000..f9abfb4ac666
--- /dev/null
+++ b/tests/kernels/mamba/test_gdn_output_projection.py
@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shape contract for ``QwenGatedDeltaNetAttention._output_projection``.
+
+Every ``forward_*`` already allocates ``core_attn_out`` and ``z`` as
+``(N, H, D)``. The helper now norms those tensors as-is and flattens once
+for ``out_proj``; it used to be rank-agnostic via ``reshape(z_shape_og)``.
+"""
+
+from __future__ import annotations
+
+import types
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.layernorm import RMSNormGated
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+    QwenGatedDeltaNetAttention,
+)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32])
+@torch.inference_mode()
+def test_output_projection_norms_per_head_and_flattens(
+    default_vllm_config,
+    dtype: torch.dtype,
+) -> None:
+    num_tokens, num_heads, head_dim = 3, 4, 8
+    hidden = num_heads * head_dim
+
+    layer = types.SimpleNamespace()
+    layer.norm = RMSNormGated(
+        head_dim,
+        eps=1e-5,
+        group_size=None,
+        norm_before_gate=True,
+        device="cpu",
+        dtype=dtype,
+    )
+    layer.out_proj = lambda x: (x, None)
+    layer._output_projection = types.MethodType(
+        QwenGatedDeltaNetAttention._output_projection, layer
+    )
+
+    core_attn_out = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype)
+    z = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype)
+    out = layer._output_projection(core_attn_out, z)
+
+    assert core_attn_out.shape == (num_tokens, num_heads, head_dim)
+    assert z.shape == (num_tokens, num_heads, head_dim)
+    assert out.shape == (num_tokens, hidden)
+    assert out.dtype == dtype
diff --git a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py
index 861b686e32da..d86cfde2c7c7 100644
--- a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py
+++ b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py
@@ -192,3 +192,86 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
     assert no_buffer_state_error.mean().item() < 6e-4
     assert buffer_o_error.max().item() == 0
     assert buffer_state_error.max().item() == 0
+
+
+def test_gdn_chunk_cutedsl_correlated_keys_stability():
+    """Cover a KKT block where three BF16 Newton rounds diverge."""
+    total_tokens = 64
+    num_k_heads = 4
+    num_v_heads = 8
+    head_dim = 128
+    rng = torch.Generator("cuda").manual_seed(3456)
+
+    q = torch.zeros(
+        1,
+        total_tokens,
+        num_k_heads,
+        head_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    k = torch.zeros_like(q)
+    q[..., 0] = 1
+    k[..., 0] = 1
+    v = torch.randn(
+        1,
+        total_tokens,
+        num_v_heads,
+        head_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+        generator=rng,
+    )
+    g = torch.zeros(
+        1,
+        total_tokens,
+        num_v_heads,
+        device="cuda",
+        dtype=torch.float32,
+    )
+    beta = torch.full_like(g, 0.9)
+    initial_state = torch.zeros(
+        1,
+        num_v_heads,
+        head_dim,
+        head_dim,
+        device="cuda",
+        dtype=torch.float32,
+    )
+    cu_seqlens = torch.tensor([0, total_tokens], device="cuda", dtype=torch.int32)
+    chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, total_tokens)
+
+    ref_o, ref_state = chunk_gated_delta_rule(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=initial_state,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+        use_qk_l2norm_in_kernel=False,
+    )
+    actual_o, actual_state = chunk_gated_delta_rule_cutedsl(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=initial_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_offsets=chunk_offsets,
+    )
+    torch.accelerator.synchronize()
+
+    assert torch.isfinite(actual_o).all()
+    assert torch.isfinite(actual_state).all()
+    output_error = (actual_o.float() - ref_o.float()).abs()
+    state_error = (
+        actual_state.float() - ref_state.to(actual_state.dtype).float()
+    ).abs()
+    assert output_error.max().item() < 1e-1
+    assert output_error.mean().item() < 2e-3
+    assert state_error.max().item() < 2e-1
+    assert state_error.mean().item() < 6e-4
diff --git a/tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py b/tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py
new file mode 100644
index 000000000000..851d6fb1d34c
--- /dev/null
+++ b/tests/kernels/mamba/test_gdn_rocm_layout_dispatch.py
@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Layout dispatch tests for ``QwenGatedDeltaNetAttention._forward_core_rocm``.
+
+The AITER fused reshape+conv kernel learned Qwen3.5's flat ``[q|k|v|z]``
+packing in https://github.com/ROCm/aiter/pull/3251, so flat-layout models take
+the decode fast path too instead of falling back to the generic path.
+
+The dispatch tests run host-side with recording stubs. The prefill numerical
+tests exercise the production ChunkGatedDeltaRule wrapper on ROCm against a
+token-wise Torch reference, including ragged state and output-buffer layouts.
+"""
+
+from __future__ import annotations
+
+import types
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+    QwenGatedDeltaNetAttention,
+)
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-specific tests"
+)
+
+PREFIX = "model.layers.0.linear_attn"
+H = 2  # num key heads
+HV = 4  # num value heads
+K = 8  # head_k_dim
+V = 8  # head_v_dim
+
+
+def _make_metadata(
+    *,
+    num_prefills: int = 0,
+    num_decodes: int = 2,
+    spec_sequence_masks: torch.Tensor | None = None,
+):
+    """Only the fields the dispatch condition reads carry meaningful values."""
+    return GDNAttentionMetadata(
+        num_prefills=num_prefills,
+        num_prefill_tokens=num_prefills,
+        num_decodes=num_decodes,
+        num_decode_tokens=num_decodes,
+        num_spec_decodes=0,
+        num_spec_decode_tokens=0,
+        num_actual_tokens=num_prefills + num_decodes,
+        spec_sequence_masks=spec_sequence_masks,
+    )
+
+
+def _make_layer(gqa_interleaved_layout: bool):
+    """Stub layer running the real ``_forward_core_rocm``, recording dispatch."""
+    layer = types.SimpleNamespace()
+    layer.prefix = PREFIX
+    layer.gqa_interleaved_layout = gqa_interleaved_layout
+    layer.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat"
+
+    layer.calls = []
+    layer._forward_core_decode_aiter = lambda **kw: layer.calls.append("aiter")
+    layer._forward_core = lambda **kw: layer.calls.append("generic")
+    layer.prepare_gdn_attention_core_inputs = lambda qkvz, ba, n: (
+        torch.zeros(n, H * K * 2 + HV * V),
+        torch.zeros(n, HV, V),
+        torch.zeros(n, HV),
+        torch.zeros(n, HV),
+    )
+    layer._forward_core_rocm = types.MethodType(
+        QwenGatedDeltaNetAttention._forward_core_rocm, layer
+    )
+    return layer
+
+
+def _run(layer, meta) -> str:
+    num_tokens = max(meta.num_actual_tokens, 1)
+    ctx = types.SimpleNamespace(attn_metadata={PREFIX: meta})
+    with patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=ctx):
+        layer._forward_core_rocm(
+            qkvz=torch.zeros(num_tokens, 2 * H * K + 2 * HV * V),
+            ba=torch.zeros(num_tokens, 2 * HV),
+            z_out=torch.zeros(num_tokens, HV, V),
+            core_attn_out=torch.zeros(num_tokens, HV, V),
+        )
+    assert len(layer.calls) == 1
+    return layer.calls[0]
+
+
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
+def test_pure_decode_takes_the_fast_path(gqa_interleaved_layout: bool) -> None:
+    """Both packings reach the fast path; flat used to be guarded out."""
+    layer = _make_layer(gqa_interleaved_layout)
+    assert _run(layer, _make_metadata()) == "aiter"
+
+
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
+def test_layout_string_matches_the_packing(gqa_interleaved_layout: bool) -> None:
+    """The value handed to the kernel's ``qkvz_layout`` parameter."""
+    layer = _make_layer(gqa_interleaved_layout)
+    expected = "interleaved" if gqa_interleaved_layout else "flat"
+    assert layer.qkvz_layout == expected
+
+
+@pytest.mark.parametrize("gqa_interleaved_layout", [True, False])
+@pytest.mark.parametrize(
+    "meta_kwargs",
+    [
+        {"num_prefills": 1},
+        {"num_decodes": 0},
+        {"spec_sequence_masks": torch.zeros(1, dtype=torch.bool)},
+    ],
+    ids=["has_prefill", "no_decode", "spec_decode"],
+)
+def test_non_pure_decode_batches_use_the_generic_path(
+    gqa_interleaved_layout: bool, meta_kwargs: dict
+) -> None:
+    layer = _make_layer(gqa_interleaved_layout)
+    assert _run(layer, _make_metadata(**meta_kwargs)) == "generic"
+
+
+def _prefill_reference(q, k, v, g, beta, initial_state, lengths):
+    """Independent token-wise recurrence in the layer's [value, key] layout."""
+    q, k, v, g, beta, initial_state = (
+        x.cpu().float() for x in (q, k, v, g, beta, initial_state)
+    )
+    q = q.repeat_interleave(v.shape[2] // q.shape[2], dim=2)
+    k = k.repeat_interleave(v.shape[2] // k.shape[2], dim=2)
+    q = q * q.shape[-1] ** -0.5
+    output = torch.empty_like(v)
+    final_state = torch.empty_like(initial_state)
+    previous_threads = torch.get_num_threads()
+    torch.set_num_threads(1)
+    try:
+        start = 0
+        for sequence, length in enumerate(lengths):
+            state = initial_state[sequence].clone()
+            for token in range(start, start + length):
+                state = state * g[0, token].exp()[:, None, None]
+                key = k[0, token]
+                prediction = (state * key[:, None, :]).sum(dim=-1)
+                delta = (v[0, token] - prediction) * beta[0, token, :, None]
+                state = state + delta[:, :, None] * key[:, None, :]
+                output[0, token] = (state * q[0, token, :, None, :]).sum(dim=-1)
+            final_state[sequence] = state
+            start += length
+    finally:
+        torch.set_num_threads(previous_threads)
+    return output, final_state
+
+
+@pytest.mark.parametrize(
+    "state_dtype,correlated",
+    [(torch.bfloat16, False), (torch.float32, False), (torch.float32, True)],
+    ids=["ragged-bf16-state", "ragged-fp32-state", "correlated-keys"],
+)
+@torch.inference_mode()
+def test_rocm_prefill_numerics_and_output_layout(state_dtype, correlated):
+    from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+        ChunkGatedDeltaRule,
+    )
+    from vllm.third_party.flash_linear_attention.ops.index import (
+        prepare_chunk_indices,
+        prepare_chunk_offsets,
+    )
+
+    if not torch.cuda.is_available():
+        pytest.skip("ROCm prefill numerics require a GPU")
+    # Only the config is needed for production backend selection; no weights.
+    config = VllmConfig(
+        model_config=ModelConfig(
+            model="Qwen/Qwen3.5-0.8B",
+            revision="2fc06364715b967f1860aea9cf38778875588b17",
+            max_model_len=1024,
+        )
+    )
+    with set_current_vllm_config(config):
+        prefill = ChunkGatedDeltaRule()
+    assert prefill.gdn_prefill_backend == "triton"
+    assert prefill._forward_method == prefill.forward_native
+
+    torch.manual_seed(1234)
+    lengths = [64] if correlated else [1, 63, 64, 65, 129]
+    tokens, key_heads, value_heads, dim = sum(lengths), 4, 8, 128
+    q = torch.randn(1, tokens, key_heads, dim, device="cuda")
+    q = torch.nn.functional.normalize(q, dim=-1).to(torch.bfloat16)
+    k = torch.randn_like(q)
+    k = torch.nn.functional.normalize(k.float(), dim=-1).to(q.dtype)
+    v = torch.randn(1, tokens, value_heads, dim, device="cuda", dtype=q.dtype)
+    g = -torch.rand(1, tokens, value_heads, device="cuda") * 0.1
+    beta = torch.rand_like(g)
+    state = (torch.randn(len(lengths), value_heads, dim, dim, device="cuda") * 0.05).to(
+        state_dtype
+    )
+    if correlated:
+        q.zero_()
+        k.zero_()
+        q[..., 0] = 1
+        k[..., 0] = 1
+        g.zero_()
+        beta.fill_(0.9)
+        state.zero_()
+    expected, expected_state = _prefill_reference(q, k, v, g, beta, state, lengths)
+    initial_state_copy = state.clone()
+    cu_seqlens = torch.tensor(
+        [0, *torch.tensor(lengths).cumsum(0).tolist()],
+        device="cuda",
+        dtype=torch.int32,
+    )
+    kwargs = dict(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=state,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=prepare_chunk_indices(cu_seqlens, 64),
+        chunk_offsets=prepare_chunk_offsets(cu_seqlens, 64),
+        use_qk_l2norm_in_kernel=False,
+    )
+    storage = torch.full(
+        (tokens + 3, value_heads, dim), float("nan"), device="cuda", dtype=q.dtype
+    )
+    output, final_state = prefill(**kwargs, core_attn_out=storage[:tokens])
+    assert output.data_ptr() == storage.data_ptr()
+    assert output.shape == v.shape
+    assert final_state.shape == state.shape
+    assert torch.isnan(storage[tokens:]).all()
+    torch.testing.assert_close(state, initial_state_copy, atol=0, rtol=0)
+
+    # Preserve the upstream prefill suite's numerical budgets. The reference
+    # here is a token-wise recurrence, independent of the chunk implementation.
+    output_error = (output.cpu().float() - expected).abs()
+    state_error = (final_state.cpu().float() - expected_state).abs()
+    assert torch.isfinite(output).all() and torch.isfinite(final_state).all()
+    assert output_error.max() < (1e-1 if correlated else 2e-3)
+    assert output_error.mean() < (2e-3 if correlated else 6e-5)
+    assert state_error.max() < (2e-1 if correlated else 2e-2)
+    assert state_error.mean() < 6e-4
+
+    no_buffer_output, no_buffer_state = prefill(**kwargs)
+    torch.testing.assert_close(no_buffer_output, output, atol=0, rtol=0)
+    torch.testing.assert_close(no_buffer_state, final_state, atol=0, rtol=0)
diff --git a/tests/kernels/mamba/test_mamba_ssm.py b/tests/kernels/mamba/test_mamba_ssm.py
index f8af3db56be6..c24f4c0f8719 100644
--- a/tests/kernels/mamba/test_mamba_ssm.py
+++ b/tests/kernels/mamba/test_mamba_ssm.py
@@ -51,8 +51,7 @@ def selective_scan_ref(
     prev_state=None,
     final_state_out=None,
 ):
-    """
-    u: r(B D L)
+    """u: r(B D L)
     delta: r(B D L)
     A: c(D N) or r(D N)
     B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
@@ -135,7 +134,7 @@ def selective_scan_opcheck_fn(
     cu_chunk_seqlen=None,
     last_chunk_indices=None,
 ):
-    """if return_last_state is True, returns (out, last_state)
+    """If return_last_state is True, returns (out, last_state)
     last_state has shape (batch, dim, dstate).
     """
     if u.stride(-1) != 1:
diff --git a/tests/kernels/mamba/test_mamba_ssm_configs.py b/tests/kernels/mamba/test_mamba_ssm_configs.py
index d35ce95746c7..7ea1529f2e8b 100644
--- a/tests/kernels/mamba/test_mamba_ssm_configs.py
+++ b/tests/kernels/mamba/test_mamba_ssm_configs.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for the JSON-based config loader added to selective_state_update.
+"""Unit tests for the JSON-based config loader added to selective_state_update.
 
 Tests cover:
   - Flat MoE-style filename generation
diff --git a/tests/kernels/mamba/test_mamba_ssm_ssd.py b/tests/kernels/mamba/test_mamba_ssm_ssd.py
index 1de25780eac4..22abeab1ed8f 100644
--- a/tests/kernels/mamba/test_mamba_ssm_ssd.py
+++ b/tests/kernels/mamba/test_mamba_ssm_ssd.py
@@ -41,14 +41,15 @@ def segsum(x):
 
 
 def ssd_minimal_discrete(X, A, B, C, block_len, initial_states=None):
-    """
-    Arguments:
+    """Arguments:
         X: (batch, length, n_heads, d_head)
         A: (batch, length, n_heads)
         B: (batch, length, n_heads, d_state)
         C: (batch, length, n_heads, d_state)
+
     Return:
         Y: (batch, length, n_heads, d_head)
+
     """
     assert X.dtype == A.dtype == B.dtype == C.dtype
     assert X.shape[1] % block_len == 0
diff --git a/tests/kernels/mamba/utils.py b/tests/kernels/mamba/utils.py
index b82c4d82c33e..c7cf47342778 100644
--- a/tests/kernels/mamba/utils.py
+++ b/tests/kernels/mamba/utils.py
@@ -9,8 +9,7 @@
 def selective_state_update_ref(
     state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
 ):
-    """
-    Argument:
+    """Argument:
         state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
         x: (batch, dim) or (batch, nheads, dim)
         dt: (batch, dim) or (batch, nheads, dim)
@@ -20,8 +19,10 @@ def selective_state_update_ref(
         D: (dim,) or (nheads, dim)
         z: (batch, dim) or (batch, nheads, dim)
         dt_bias: (dim,) or (nheads, dim)
+
     Return:
         out: (batch, dim) or (batch, nheads, dim)
+
     """
     has_heads = state.dim() > 3
     if state.dim() == 3:
diff --git a/tests/kernels/moe/conftest.py b/tests/kernels/moe/conftest.py
index 1405991c9ca5..af99dc756b96 100644
--- a/tests/kernels/moe/conftest.py
+++ b/tests/kernels/moe/conftest.py
@@ -3,6 +3,7 @@
 import pytest
 
 from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.v1.worker.workspace import reset_workspace_manager
 
 NEEDS_CLEAN_ENTRY = frozenset(
     {
@@ -57,5 +58,6 @@ def pytest_runtest_setup(item):
     # Not tryfirst: the skipping plugin evaluates skip marks in a tryfirst
     # hook, so tests about to be skipped never reach this, while fixture setup
     # still runs afterwards.
-    if item.path.name in NEEDS_CLEAN_ENTRY:
+    if item.path.name in NEEDS_CLEAN_ENTRY or "dist_init" in item.fixturenames:
+        reset_workspace_manager()
         cleanup_dist_env_and_memory()
diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py
index 8041db68d752..1944aaac5733 100644
--- a/tests/kernels/moe/modular_kernel_tools/common.py
+++ b/tests/kernels/moe/modular_kernel_tools/common.py
@@ -15,6 +15,7 @@
     dequantize_nvfp4_to_dtype,
 )
 from tests.kernels.utils import torch_experts
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.distributed import (
     get_dp_group,
@@ -33,7 +34,9 @@
     FusedMoEQuantConfig,
     RoutingMethodType,
 )
+from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
     kFp8Dynamic128Sym,
     kFp8DynamicTensorSym,
     kFp8DynamicTokenSym,
@@ -81,6 +84,8 @@ class Config:
 
     world_size: int
 
+    activation: MoEActivation = MoEActivation.SILU
+
     torch_trace_dir_path: str | None = None
 
     def __post_init__(self):
@@ -148,9 +153,7 @@ def num_local_experts(self) -> int:
         return self.E // self.world_size
 
     def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
-        """
-        make env data for vllm launch.
-        """
+        """Make env data for vllm launch."""
         vllm_config = VllmConfig()
         vllm_config.model_config = SimpleNamespace(
             enforce_eager=True,
@@ -167,31 +170,34 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
 
         return vllm_config, env_dict
 
+    def fp8_quant_key_pair(self) -> tuple[QuantKey, QuantKey]:
+        """Derive the (weight_quant_key, activation_quant_key) pair an FP8
+        quant config of this shape corresponds to (either OCP or FNUZ FP8,
+        see ``current_platform.fp8_dtype()``)."""
+        if self.quant_block_shape is not None:
+            return kFp8Static128BlockSym, kFp8Dynamic128Sym
+        if self.is_per_out_ch_quant:
+            return (
+                kFp8StaticChannelSym,
+                kFp8DynamicTokenSym
+                if self.is_per_act_token_quant
+                else kFp8StaticTensorSym,
+            )
+        return (
+            kFp8StaticTensorSym,
+            kFp8DynamicTensorSym
+            if self.is_per_act_token_quant
+            else kFp8StaticTensorSym,
+        )
+
     def fe_supports_quant_scheme(self) -> bool:
         """Check if the fused experts class supports this quant config.
         See https://github.com/ROCm/aiter/issues/2419 for AITER gaps."""
         if self.quant_config is None or self.quant_dtype is None:
             return True
-        if self.quant_dtype != torch.float8_e4m3fn:
+        if not is_fp8(self.quant_dtype):
             return True
-        # Derive QuantKeys from test config
-        if self.quant_block_shape is not None:
-            w_key = kFp8Static128BlockSym
-            a_key = kFp8Dynamic128Sym
-        elif self.is_per_out_ch_quant:
-            w_key = kFp8StaticChannelSym
-            a_key = (
-                kFp8DynamicTokenSym
-                if self.is_per_act_token_quant
-                else kFp8StaticTensorSym
-            )
-        else:
-            w_key = kFp8StaticTensorSym
-            a_key = (
-                kFp8DynamicTensorSym
-                if self.is_per_act_token_quant
-                else kFp8StaticTensorSym
-            )
+        w_key, a_key = self.fp8_quant_key_pair()
         fe_cls = self.fused_experts_type
         if hasattr(fe_cls, "_supports_quant_scheme"):
             try:
@@ -201,10 +207,7 @@ def fe_supports_quant_scheme(self) -> bool:
         return True
 
     def is_fp8_block_quantized(self):
-        return (
-            self.quant_dtype == torch.float8_e4m3fn
-            and self.quant_block_shape is not None
-        )
+        return is_fp8(self.quant_dtype) and self.quant_block_shape is not None
 
     def is_batched_prepare_finalize(self):
         info = prepare_finalize_info(self.prepare_finalize_type)
@@ -310,6 +313,16 @@ def is_valid(self) -> tuple[bool, str | None]:
                 f"block={self.quant_block_shape})"
             )
 
+        # Check activation support; NotImplementedError means no opinion.
+        try:
+            if not self.fused_experts_type._supports_activation(self.activation):
+                return False, (
+                    f"FE {self.fused_experts_type.__name__} does not support "
+                    f"activation {self.activation}"
+                )
+        except NotImplementedError:
+            pass
+
         # Check block quantization support
         is_block_quantized = self.quant_block_shape is not None
         if is_block_quantized and self.quant_dtype is None:
@@ -336,6 +349,11 @@ def is_valid(self) -> tuple[bool, str | None]:
             return False, "Needs Aiter, but Aiter not available."
         if self.needs_mori() and not has_mori():  # noqa: SIM103
             return False, "Needs MoRI, but MoRI not available."
+        if self.needs_mori() and not rocm_aiter_ops.is_fused_moe_enabled():
+            return False, (
+                "Mori requires AITER's fused-moe backend to be enabled "
+                "(VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_MOE=1)."
+            )
 
         try:
             if not self.fused_experts_type._supports_current_device():
@@ -372,7 +390,7 @@ def describe(self):
     def is_quantized(self) -> bool:
         # or w1_scale is not None?
         return (
-            self.w1.dtype == torch.float8_e4m3fn
+            is_fp8(self.w1.dtype)
             or self.w1.dtype == torch.uint8
             or self.w1.dtype == torch.int8
         )
@@ -444,9 +462,7 @@ def describe(self):
     def make_hidden_states(
         config: Config,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """
-        Return hidden_states
-        """
+        """Return hidden_states."""
         m, k, dtype = (config.M, config.K, config.dtype)
         device = torch.accelerator.current_device_index()
         a = torch.randn((m, k), device=device, dtype=dtype) / 15.0
@@ -601,6 +617,7 @@ def reference_moe_impl(
         topk_ids=rank_tensors.topk_ids,
         global_num_experts=config.E,
         expert_map=None,
+        activation=config.activation,
         w1_scale=w1_scale,
         w2_scale=w2_scale,
         a1_scale=a_scale,
@@ -644,7 +661,7 @@ def make_modular_kernel(
         moe_parallel_config=moe_parallel_config,
         in_dtype=config.dtype,
         max_num_tokens=next_power_of_2(config.M),
-        activation=MoEActivation.SILU,
+        activation=config.activation,
         device=vllm_config.device_config.device,
         routing_method=RoutingMethodType.DeepSeekV3,
     )
@@ -770,7 +787,7 @@ def run_modular_kernel(
         "w2": rank_weights.w2,
         "topk_weights": rank_tensors.topk_weights,
         "topk_ids": topk_ids,
-        "activation": MoEActivation.SILU,
+        "activation": config.activation,
         "expert_map": rank_tensors.expert_map,
         "global_num_experts": config.E,
         "apply_router_weight_on_input": config.topk == 1
diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py
index 7f1924bac5ab..e0f37f54360a 100644
--- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py
+++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py
@@ -86,14 +86,14 @@ class ExpertInfo:
 standard_format = mk.FusedMoEActivationFormat.Standard
 batched_format = mk.FusedMoEActivationFormat.BatchedExperts
 common_float_types: list[torch.dtype | str] = [
-    torch.float8_e4m3fn,
+    current_platform.fp8_dtype(),
     torch.bfloat16,
     torch.float16,
     torch.float32,
 ]
 common_float_and_int_types = common_float_types + [torch.int8]
 nvfp4_types = ["nvfp4"]
-fp8_types = [torch.float8_e4m3fn]
+fp8_types = [current_platform.fp8_dtype()]
 
 
 def register_prepare_and_finalize(
@@ -324,7 +324,9 @@ def expert_info(kind) -> ExpertInfo:
     register_experts(
         AiterExperts,
         standard_format,
-        fp8_types,
+        # AiterExperts also supports the fully-unquantized (None, None)
+        # scheme (see SUPPORTED_W_A in rocm_aiter_moe.py), not just fp8.
+        common_float_types,
         blocked_quantization_support=True,
         needs_aiter=True,
     )
@@ -397,35 +399,35 @@ def expert_info(kind) -> ExpertInfo:
     None,
     # per-channel / per-column weights and per-tensor activations
     TestMoEQuantConfig(
-        quant_dtype=torch.float8_e4m3fn,
+        quant_dtype=current_platform.fp8_dtype(),
         per_out_ch_quant=True,
         per_act_token_quant=False,
         block_shape=None,
     ),
     # per-channel / per-column weights and per-token activations
     TestMoEQuantConfig(
-        quant_dtype=torch.float8_e4m3fn,
+        quant_dtype=current_platform.fp8_dtype(),
         per_out_ch_quant=True,
         per_act_token_quant=True,
         block_shape=None,
     ),
     # per-tensor weights and per-tensor activations
     TestMoEQuantConfig(
-        quant_dtype=torch.float8_e4m3fn,
+        quant_dtype=current_platform.fp8_dtype(),
         per_out_ch_quant=False,
         per_act_token_quant=False,
         block_shape=None,
     ),
     # per-tensor weights and per-token activations
     TestMoEQuantConfig(
-        quant_dtype=torch.float8_e4m3fn,
+        quant_dtype=current_platform.fp8_dtype(),
         per_out_ch_quant=False,
         per_act_token_quant=True,
         block_shape=None,
     ),
     # block-quantized weights and 128 block per-token activations
     TestMoEQuantConfig(
-        quant_dtype=torch.float8_e4m3fn,
+        quant_dtype=current_platform.fp8_dtype(),
         per_out_ch_quant=False,
         per_act_token_quant=False,
         block_shape=[128, 128],
diff --git a/tests/kernels/moe/parallel_utils.py b/tests/kernels/moe/parallel_utils.py
index f926b034b3c2..972ebe6512eb 100644
--- a/tests/kernels/moe/parallel_utils.py
+++ b/tests/kernels/moe/parallel_utils.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-DeepEP test utilities
-"""
+"""DeepEP test utilities."""
 
 import dataclasses
 import os
diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py
index 3fca9190dc31..4c4c82d54e17 100644
--- a/tests/kernels/moe/test_batched_deepgemm.py
+++ b/tests/kernels/moe/test_batched_deepgemm.py
@@ -34,7 +34,6 @@ def test_batched_deepgemm_vs_triton(
     E: int, T: int, K: int, N: int, topk: int, monkeypatch, workspace_init
 ):
     """Compare BatchedDeepGemmExperts to BatchedTritonExperts."""
-
     monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1")
 
     device = "cuda"
diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py
index 1cbbbb45fb09..68f6b3b0b312 100644
--- a/tests/kernels/moe/test_block_fp8.py
+++ b/tests/kernels/moe/test_block_fp8.py
@@ -25,6 +25,7 @@
 from vllm.model_executor.layers.fused_moe.config import (
     fp8_w8a8_moe_quant_config,
 )
+from vllm.model_executor.layers.fused_moe.experts import triton_moe
 from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
     _valid_deep_gemm_shape,
 )
@@ -106,7 +107,16 @@
 
 
 def torch_w8a8_block_fp8_moe(
-    a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, block_shape, silu_fp32=False
+    a,
+    w1,
+    w2,
+    w1_s,
+    w2_s,
+    topk_weight,
+    topk_ids,
+    block_shape,
+    silu_fp32=False,
+    expert_map=None,
 ):
     """Fused MoE with block-wise fp8 quantization using native torch.
 
@@ -122,6 +132,9 @@ def torch_w8a8_block_fp8_moe(
 
     topk_weight = topk_weight.view(-1)
     topk_ids = topk_ids.view(-1)
+    if expert_map is not None:
+        # w1/w2 hold this rank's experts; slots mapped to -1 contribute nothing.
+        topk_ids = expert_map[topk_ids]
 
     _, block_k = block_shape[0], block_shape[1]
     # Quantize with the production per-token-group fp8 kernel (same HIP/CUDA op the
@@ -246,6 +259,108 @@ def test_w8a8_block_fp8_fused_moe(
     torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol)
 
 
+@pytest.mark.parametrize("clamp_limit", [None, 10.0])
+@torch.inference_mode()
+def test_w8a8_block_fp8_clamp_skips_fused_silu_quant(
+    clamp_limit, monkeypatch, workspace_init, disable_deepgemm_ue8m0
+):
+    """A SwiGLU clamp must bypass the fused SiLU+quant kernel, which cannot clamp."""
+    torch.manual_seed(0)
+    M, N, K, E, topk = 16, 256, 512, 8, 2
+    block_size = [128, 128]
+
+    a = torch.randn((M, K), dtype=torch.bfloat16) / 10
+    score = torch.randn((M, E), dtype=torch.bfloat16)
+    (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
+        E, N, K, torch.bfloat16, torch.float8_e4m3fn, block_shape=block_size
+    )
+    quant_config = fp8_w8a8_moe_quant_config(
+        w1_scale=w1_s,
+        w2_scale=w2_s,
+        block_shape=block_size,
+        gemm1_clamp_limit=clamp_limit,
+    )
+    m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)
+    topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
+
+    fused_calls = 0
+    fused_op = triton_moe.ops.silu_and_mul_per_block_quant
+
+    def counting_fused_op(*args, **kwargs):
+        nonlocal fused_calls
+        fused_calls += 1
+        return fused_op(*args, **kwargs)
+
+    monkeypatch.setattr(
+        triton_moe.ops, "silu_and_mul_per_block_quant", counting_fused_op
+    )
+
+    with set_current_vllm_config(vllm_config):
+        m_fused_moe.apply(
+            a,
+            w1,
+            w2,
+            topk_weights,
+            topk_ids,
+            activation=MoEActivation.SILU,
+            apply_router_weight_on_input=False,
+            expert_map=None,
+            global_num_experts=E,
+        )
+
+    assert (fused_calls > 0) == (clamp_limit is None), (
+        f"fused fast path ran {fused_calls} times with clamp_limit={clamp_limit}"
+    )
+
+
+@pytest.mark.parametrize(("M", "N", "K"), [(83, 512, 512), (2048, 4608, 512)])
+@pytest.mark.parametrize("topk", [2, 6])
+@pytest.mark.parametrize("ep_size", [2, 4])
+@torch.inference_mode()
+def test_w8a8_block_fp8_fused_moe_expert_map(M, N, K, topk, ep_size, workspace_init):
+    """One EP rank: local expert shard plus an expert_map with -1 entries."""
+    E, block_size, dtype = 16, [128, 128], torch.bfloat16
+    torch.manual_seed(0)
+
+    a = torch.randn((M, K), dtype=dtype) / 10
+    score = torch.randn((M, E), dtype=dtype)
+    w1, w2, quant_config = make_test_quant_config(
+        E, N, K, dtype, quant_dtype=torch.float8_e4m3fn, block_shape=block_size
+    )
+
+    local_e = E // ep_size
+    e_ids = torch.randperm(E, dtype=torch.int32)[:local_e]
+    e_map = torch.full((E,), -1, dtype=torch.int32)
+    e_map[e_ids] = torch.arange(local_e, dtype=torch.int32)
+    w1, w2 = w1[e_ids].contiguous(), w2[e_ids].contiguous()
+    w1_s = quant_config.w1_scale[e_ids].contiguous()
+    w2_s = quant_config.w2_scale[e_ids].contiguous()
+    quant_config = fp8_w8a8_moe_quant_config(
+        w1_scale=w1_s, w2_scale=w2_s, block_shape=block_size
+    )
+
+    m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)
+    topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
+
+    with set_current_vllm_config(vllm_config):
+        ref_out = torch_w8a8_block_fp8_moe(
+            a, w1, w2, w1_s, w2_s, topk_weights, topk_ids, block_size, expert_map=e_map
+        )
+        m_out = m_fused_moe.apply(
+            a,
+            w1,
+            w2,
+            topk_weights,
+            topk_ids,
+            activation=MoEActivation.SILU,
+            apply_router_weight_on_input=False,
+            expert_map=e_map,
+            global_num_experts=E,
+        )
+
+    torch.testing.assert_close(m_out, ref_out, atol=0.035, rtol=0.035)
+
+
 @pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS_DG)
 @pytest.mark.parametrize("E", E)
 @pytest.mark.parametrize("topk", TOP_KS)
diff --git a/tests/kernels/moe/test_count_expert_num_tokens.py b/tests/kernels/moe/test_count_expert_num_tokens.py
index 8014076192e2..b2b6abc5b489 100644
--- a/tests/kernels/moe/test_count_expert_num_tokens.py
+++ b/tests/kernels/moe/test_count_expert_num_tokens.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests compute_expert_num_tokens kernels
-"""
+"""Tests compute_expert_num_tokens kernels."""
 
 import dataclasses
 
diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py
index dccda70c0b9e..d4c543722131 100644
--- a/tests/kernels/moe/test_cpu_fused_moe.py
+++ b/tests/kernels/moe/test_cpu_fused_moe.py
@@ -306,17 +306,28 @@ def test_cpu_fused_moe(
 
 
 @pytest.mark.skipif(
-    current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
-    reason="Requires Arm CPU",
+    current_platform.get_cpu_architecture()
+    not in (CpuArchEnum.ARM, CpuArchEnum.POWERPC),
+    reason="Requires Arm or POWER CPU",
 )
 @pytest.mark.parametrize("batch_size", BATCH_SIZE)
 @pytest.mark.parametrize("expert_num", EXPERT_NUM)
 @pytest.mark.parametrize("hidden_size", HIDDEN_DIM)
 @pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM)
 @pytest.mark.parametrize("use_bias", USE_BIAS)
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize(
+    "dtype",
+    (
+        [torch.float32, torch.float16, torch.bfloat16]
+        if current_platform.get_cpu_architecture() == CpuArchEnum.ARM
+        else [torch.float32, torch.bfloat16]
+    ),
+)
 @pytest.mark.parametrize("act", ACT)
-@pytest.mark.parametrize("isa", ["neon"])
+@pytest.mark.parametrize(
+    "isa",
+    ["neon"] if current_platform.get_cpu_architecture() == CpuArchEnum.ARM else ["vsx"],
+)
 def test_cpu_fused_moe_int8(
     batch_size: int,
     expert_num: int,
@@ -386,7 +397,8 @@ def test_cpu_fused_moe_int8(
         isa,
     )
 
-    torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2)
+    atol = rtol = 1e-1 if isa == "vsx" else 2e-2
+    torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
 
 
 # moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel
@@ -501,7 +513,6 @@ def test_cpu_fused_moe_unaligned_intermediate_size(
 ):
     """CPU kernels handle unaligned intermediate sizes by zero-padding the
     weights before prepacking."""
-
     set_random_seed(0)
     batch_size = 64
     intermediate_size = UNALIGNED_INTERMEDIATE_DIM
diff --git a/tests/kernels/moe/test_cpu_int4_moe.py b/tests/kernels/moe/test_cpu_int4_moe.py
index 04931e386f14..2d5534ba086e 100644
--- a/tests/kernels/moe/test_cpu_int4_moe.py
+++ b/tests/kernels/moe/test_cpu_int4_moe.py
@@ -62,6 +62,7 @@ def _make_int4_moe_weights(
     Returns:
         (w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias)
         where *_ref are the dequantized float reference weights
+
     """
     # Generate INT4 weights as int8 values in [-8, 7]
     w13_int4 = torch.randint(-8, 8, (E, 2 * N, K), dtype=torch.int8)
diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py
index 84ff1ae266db..e89fbb55a089 100644
--- a/tests/kernels/moe/test_cpu_quant_fused_moe.py
+++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py
@@ -826,6 +826,7 @@ def _make_int4_moe_weights(E, N, K, group_size, quant_algo):
         w1_zeros, w2_zeros,
         w1_zeros_packed, w2_zeros_packed,
         w1_s, w2_s
+
     """
     w1_int4 = torch.randint(0, 16, (E, K, 2 * N), dtype=torch.int32)
     w2_int4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32)
@@ -1113,5 +1114,203 @@ def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace):
     )
 
 
+# ===========================================================================
+# FP8 W8A8 MoE
+# ===========================================================================
+
+
+requires_cpu_fp8_w8a8 = pytest.mark.skipif(
+    not ops.cpu_has_amx_fp8(),
+    reason="requires native AMX-FP8 hardware",
+)
+
+FP8_W8A8_MAX = torch.finfo(torch.float8_e4m3fn).max
+FP8_W8A8_QUANT_GROUP = 128  # quantization group size for both K and N dimensions
+
+
+def _make_fp8_w8a8_weight_w13(
+    E: int, two_n: int, K: int, group_K: int = FP8_W8A8_QUANT_GROUP
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Create w13 [E, 2N, K] FP8 with group scales [E, 2N, G] where G = K//group_K."""
+    G = K // group_K
+    w_list, s_list = [], []
+    for _ in range(E):
+        w_f32 = torch.randn(two_n, K)
+        # Per-row, per-K-group quantization: scale [2N, G]
+        w_re = w_f32.view(two_n, G, group_K)
+        abs_max = w_re.abs().amax(dim=2, keepdim=True).clamp(min=1e-7)  # [2N, G, 1]
+        scale = (abs_max / FP8_W8A8_MAX).squeeze(2)  # [2N, G]
+        w_q = (
+            (w_re / abs_max).clamp(-FP8_W8A8_MAX, FP8_W8A8_MAX).to(torch.float8_e4m3fn)
+        )
+        w_list.append(w_q.view(two_n, K).contiguous())
+        s_list.append(scale.float())
+    return torch.stack(w_list), torch.stack(s_list)  # [E, 2N, K], [E, 2N, G]
+
+
+def _make_fp8_w8a8_weight_w2(
+    E: int,
+    K: int,
+    N: int,
+    group_K: int = FP8_W8A8_QUANT_GROUP,
+    group_N: int = FP8_W8A8_QUANT_GROUP,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Create w2 [E, K, N] FP8 with block scales [E, K//gK, N//gN]."""
+    nK = K // group_K
+    nN = N // group_N
+    w_list, s_list = [], []
+    for _ in range(E):
+        w_f32 = torch.randn(K, N)
+        w_re = w_f32.view(nK, group_K, nN, group_N)
+        abs_max = w_re.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-7)
+        scale = (abs_max / FP8_W8A8_MAX).squeeze(1).squeeze(2)  # [nK, nN]
+        w_q = (
+            (w_re / abs_max).clamp(-FP8_W8A8_MAX, FP8_W8A8_MAX).to(torch.float8_e4m3fn)
+        )
+        w_list.append(w_q.view(K, N).contiguous())
+        s_list.append(scale.float())
+    return torch.stack(w_list), torch.stack(s_list)  # [E, K, N], [E, K//gK, N//gN]
+
+
+def _pack_fp8_w8a8_w13(
+    w13: torch.Tensor,  # [E, 2N, K] FP8
+    w13_scale: torch.Tensor,  # [E, 2N, G] float32
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack w13 for CPU FP8 W8A8 MoE kernel using float8_linear_prepack_cpu."""
+    E = w13.size(0)
+    packed_list, scale_list = [], []
+    for i in range(E):
+        pw, ps = torch.ops._C.float8_linear_prepack_cpu(
+            w13[i].contiguous(), w13_scale[i].contiguous()
+        )
+        packed_list.append(pw)
+        scale_list.append(ps)
+    return torch.stack(packed_list), torch.stack(scale_list)
+
+
+def _ref_fp8_w8a8_moe(
+    hidden_states: torch.Tensor,  # BF16 [M, K_hidden] (after FP8 round-trip dequant)
+    w13: torch.Tensor,  # FP8 [E, 2N, K_hidden]
+    w13_scale: torch.Tensor,  # float32 [E, 2N, G] (G = K_hidden//group_K)
+    w2: torch.Tensor,  # FP8 [E, K_hidden, N]
+    w2_scale: torch.Tensor,  # float32 [E, K_hidden//gK, N//gN]
+    topk_weights: torch.Tensor,  # float32 [M, top_k]
+    topk_ids: torch.Tensor,  # int64 [M, top_k]
+    N: int,
+    group_K: int = FP8_W8A8_QUANT_GROUP,
+    group_N: int = FP8_W8A8_QUANT_GROUP,
+) -> torch.Tensor:
+    """Reference implementation: dequantize weights then compute MoE."""
+    M, K_hidden = hidden_states.shape
+    top_k = topk_ids.shape[1]
+    output = torch.zeros(M, K_hidden, dtype=torch.bfloat16)
+
+    for tok in range(M):
+        x = hidden_states[tok].float()  # [K_hidden]
+        for k_idx in range(top_k):
+            expert_id = topk_ids[tok, k_idx].item()
+            w = topk_weights[tok, k_idx].item()
+
+            # Stage 1: x @ w13.T — gate and up projections
+            w13_e = w13[expert_id].float()  # [2N, K_hidden]
+            ws_e = w13_scale[expert_id].float()  # [2N, G]
+            G = ws_e.shape[1]
+            w13_dq = torch.zeros_like(w13_e)
+            for g in range(G):
+                c0, c1 = g * group_K, (g + 1) * group_K
+                w13_dq[:, c0:c1] = w13_e[:, c0:c1] * ws_e[:, g : g + 1]
+            gate_up = (x.unsqueeze(0) @ w13_dq.T).squeeze(0)  # [2N]
+
+            gate = gate_up[:N]
+            up = gate_up[N:]
+            act = F.silu(gate) * up  # [N]
+
+            # Stage 2: act @ w2.T — down projection
+            w2_e = w2[expert_id].float()  # [K_hidden, N]
+            ws2_e = w2_scale[expert_id].float()  # [K_hidden//gK, N//gN]
+            nK2 = K_hidden // group_K
+            nN2 = N // group_N
+            w2_dq = torch.zeros_like(w2_e)
+            for gi in range(nK2):
+                for gj in range(nN2):
+                    r0, r1 = gi * group_K, (gi + 1) * group_K
+                    c0, c1 = gj * group_N, (gj + 1) * group_N
+                    w2_dq[r0:r1, c0:c1] = w2_e[r0:r1, c0:c1] * ws2_e[gi, gj]
+            out_e = (act.unsqueeze(0) @ w2_dq.T).squeeze(0)  # [K_hidden]
+
+            output[tok] += (w * out_e).to(torch.bfloat16)
+
+    return output
+
+
+@requires_cpu_fp8_w8a8
+@pytest.mark.parametrize(
+    "E,N,K,M,top_k",
+    [
+        # N and K must be multiples of FP8_W8A8_QUANT_GROUP=128
+        (4, 128, 256, 1, 2),  # single token (decode phase)
+        (4, 128, 256, 8, 2),
+        (8, 256, 512, 16, 2),
+        (4, 256, 512, 4, 1),
+    ],
+)
+def test_fp8_w8a8_cpu_fused_moe(E: int, N: int, K: int, M: int, top_k: int):
+    """Test fused_experts_cpu FP8_W8A8 shape/dtype and accuracy vs a
+    dequantized BF16 reference."""
+    set_random_seed(42)
+
+    w13, w13_scale = _make_fp8_w8a8_weight_w13(E, 2 * N, K)
+    w2, w2_scale = _make_fp8_w8a8_weight_w2(E, K, N)
+
+    packed_w13, packed_w13_scale = _pack_fp8_w8a8_w13(w13, w13_scale)
+    packed_w2 = _prepack_experts(w2)
+
+    hidden_states = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+    x_fp8, x_scales = torch.ops._C.quantize_fp8e4m3_vec(hidden_states, True, None)
+    assert x_fp8.dtype == torch.float8_e4m3fn
+    assert x_scales.shape == (M,)
+
+    gen = torch.Generator().manual_seed(0)
+    topk_weights = torch.softmax(torch.randn(M, top_k, generator=gen), dim=-1).float()
+    topk_ids = torch.zeros(M, top_k, dtype=torch.int32)
+    for i in range(M):
+        topk_ids[i] = torch.randperm(E, generator=torch.Generator().manual_seed(i))[
+            :top_k
+        ].int()
+
+    output = torch.empty(M, K, dtype=torch.bfloat16)
+    ops.fused_experts_cpu(
+        output,
+        x_fp8,
+        packed_w13,
+        packed_w2,
+        topk_weights,
+        topk_ids,
+        ops.CPUQuantMethod.FP8_W8A8,
+        packed_w13_scale,  # w1_scale: [E, Nc, G, BLOCK_N]
+        w2_scale,  # w2_scale: [E, K//128, N//128]
+        None,  # w1_zero
+        None,  # w2_zero
+        [FP8_W8A8_QUANT_GROUP, FP8_W8A8_QUANT_GROUP],  # block_size
+        None,  # w1_bias
+        None,  # w2_bias
+        None,  # alpha
+        None,  # limit
+        True,  # is_vnni (weights already packed)
+        x_scales,  # a1_scale: per-token FP8 activation scales
+    )
+
+    assert output.shape == (M, K)
+    assert output.dtype == torch.bfloat16
+    assert not torch.isnan(output).any()
+    assert not torch.isinf(output).any()
+
+    x_dequant = (x_fp8.float() * x_scales.unsqueeze(1)).bfloat16()
+    ref = _ref_fp8_w8a8_moe(
+        x_dequant, w13, w13_scale, w2, w2_scale, topk_weights, topk_ids.long(), N
+    )
+    torch.testing.assert_close(output.float(), ref.float(), atol=1e-2, rtol=1e-2)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))
diff --git a/tests/kernels/moe/test_cutedsl_moe.py b/tests/kernels/moe/test_cutedsl_moe.py
index 3a371968fc51..09a7256a5cdc 100644
--- a/tests/kernels/moe/test_cutedsl_moe.py
+++ b/tests/kernels/moe/test_cutedsl_moe.py
@@ -91,11 +91,9 @@ def break_fp4_bytes(a, dtype):
 def generate_balanced_routing(
     hidden_states: torch.Tensor, num_experts: int, top_k: int
 ):
-    """
-    Generate routing weights and topk indices such that every expert is active.
+    """Generate routing weights and topk indices such that every expert is active.
     Returns routing_weights, topk_idx
     """
-
     num_tokens, hidden_dim = hidden_states.shape
     #   num_tokens = batch_size * seq_len
 
@@ -232,13 +230,13 @@ def grouped_gemm_ref(
     *,
     block_size: int = 16,
 ) -> torch.Tensor:
-    """
-    Computes the reference grouped GEMM (fp4 quantized per-expert loop),
+    """Computes the reference grouped GEMM (fp4 quantized per-expert loop),
     computes flashinfer grouped GEMM (for scale consistency),
     and returns ONLY the repacked reference output: out_ref.
 
     Returns:
         out_ref: Tensor [num_experts, max_m, n_out]
+
     """
     device_hs = hidden_states_expanded.device
     device_w = weights.device
diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py
index efb1e2f2969d..b755741380ed 100644
--- a/tests/kernels/moe/test_deepep_deepgemm_moe.py
+++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test DeepEP + DeepGEMM integration
+"""Test DeepEP + DeepGEMM integration
 DeepGEMM are gemm kernels specialized for the
 fp8 block-quantized case.
 """
@@ -92,9 +91,7 @@ def make_block_quant_fp8_weights(
     k: int,
     block_size: list[int],
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Return weights w1q, w2q, w1_scale, w2_scale
-    """
+    """Return weights w1q, w2q, w1_scale, w2_scale."""
     (_, w1q, w1_scale, _), (_, w2q, w2_scale, _) = make_test_weights(
         e, n, k, torch.bfloat16, torch.float8_e4m3fn, block_shape=block_size
     )
@@ -459,10 +456,7 @@ def test_ht_deepep_deepgemm_moe(
     disable_deepgemm_ue8m0,
     workspace_init,
 ):
-    """
-    Tests for High-Throughput DeepEP + DeepGemm integration.
-    """
-
+    """Tests for High-Throughput DeepEP + DeepGemm integration."""
     m, n, k = mnk
     set_random_seed(7)
 
@@ -533,9 +527,7 @@ def test_ll_deepep_deepgemm_moe(
     disable_deepgemm_ue8m0,
     workspace_init,
 ):
-    """
-    Tests for Low-Latency DeepEP + DeepGemm integration.
-    """
+    """Tests for Low-Latency DeepEP + DeepGemm integration."""
     assert not is_deep_gemm_e8m0_used()
 
     m, n, k = mnk
diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py
index ce9e4f343d62..ea2113d64b43 100644
--- a/tests/kernels/moe/test_deepep_moe.py
+++ b/tests/kernels/moe/test_deepep_moe.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test deepep dispatch-combine logic
-"""
+"""Test deepep dispatch-combine logic."""
 
 import dataclasses
 
@@ -56,9 +54,7 @@
 def make_weights(
     e, n, k, dtype
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Return weights w1, w2, w1_scale, w2_scale
-    """
+    """Return weights w1, w2, w1_scale, w2_scale."""
     if dtype in [torch.float16, torch.bfloat16]:
         w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
         w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
diff --git a/tests/kernels/moe/test_deepep_v2_async_finalize.py b/tests/kernels/moe/test_deepep_v2_async_finalize.py
new file mode 100644
index 000000000000..193d8906f37c
--- /dev/null
+++ b/tests/kernels/moe/test_deepep_v2_async_finalize.py
@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Contract tests for DeepEPV2PrepareAndFinalize.finalize_async.
+
+Hermetic: fake buffer/event, single process, no GPU collectives — only
+requires deep_ep to be importable.
+
+  T1 finalize_async  -> combine is issued with async_with_compute_stream=True
+                        and the join is deferred to the receiver, which waits
+                        on the combine event (device-side) before the copy.
+  T2 DBO active      -> combine falls back to async_with_compute_stream=False;
+                        the receiver still produces the correct output.
+  T3 finalize (sync) -> combine is issued synchronously and the output is
+                        copied immediately.
+"""
+
+import pytest
+import torch
+
+from vllm.utils.import_utils import has_deep_ep_v2
+
+requires_deep_ep_v2 = pytest.mark.skipif(
+    not has_deep_ep_v2(),
+    reason="Requires DeepEP v2 (ElasticBuffer)",
+)
+
+if has_deep_ep_v2():
+    import vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 as _dv2
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
+        DeepEPV2PrepareAndFinalize,
+    )
+    from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+        TopKWeightAndReduceContiguous,
+    )
+
+
+class _FakeEvent:
+    def __init__(self, has_event: bool):
+        self.event = object() if has_event else None
+        self.waited = False
+
+    def current_stream_wait(self):
+        assert self.event is not None
+        self.waited = True
+
+
+class _FakeBuffer:
+    def __init__(self, out: torch.Tensor):
+        self.out = out
+        self.calls: list[dict] = []
+        self.last_event: _FakeEvent | None = None
+
+    def combine(
+        self,
+        x,
+        handle,
+        topk_weights,
+        async_with_compute_stream,
+        allocate_on_comm_stream,
+    ):
+        self.calls.append({"async_with_compute_stream": async_with_compute_stream})
+        # DeepEP returns a completion event only for an async combine.
+        self.last_event = _FakeEvent(has_event=async_with_compute_stream)
+        return self.out, None, self.last_event
+
+
+def _make_pf(out: torch.Tensor):
+    pf = DeepEPV2PrepareAndFinalize(
+        buffer=_FakeBuffer(out),
+        num_dispatchers=1,
+        dp_size=1,
+        rank_expert_offset=0,
+        num_experts=8,
+        num_topk=2,
+    )
+    pf.handles[0] = object()
+    return pf
+
+
+def _run(pf, output: torch.Tensor, do_async: bool):
+    empty = torch.empty(0, dtype=torch.bfloat16)
+    weights = torch.empty(0, 2)
+    ids = torch.empty(0, 2, dtype=torch.int64)
+    fn = pf.finalize_async if do_async else pf.finalize
+    return fn(output, empty, weights, ids, False, TopKWeightAndReduceContiguous())
+
+
+@requires_deep_ep_v2
+def test_finalize_async_defers_join_to_receiver():
+    out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16)
+    pf = _make_pf(out_src)
+    dst = torch.zeros_like(out_src)
+    recv = _run(pf, dst, do_async=True)
+    assert pf.buffer.calls[0]["async_with_compute_stream"] is True
+    assert not torch.equal(dst, out_src), "join must not happen before receiver"
+    recv()
+    assert pf.buffer.last_event.waited, "receiver must join via event wait"
+    assert torch.equal(dst, out_src)
+
+
+@requires_deep_ep_v2
+def test_finalize_async_under_dbo_uses_sync_combine(monkeypatch):
+    monkeypatch.setattr(_dv2, "dbo_enabled", lambda: True)
+    out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16)
+    pf = _make_pf(out_src)
+    dst = torch.zeros_like(out_src)
+    recv = _run(pf, dst, do_async=True)
+    assert pf.buffer.calls[0]["async_with_compute_stream"] is False
+    recv()
+    assert not pf.buffer.last_event.waited, "no event to wait on for sync combine"
+    assert torch.equal(dst, out_src)
+
+
+@requires_deep_ep_v2
+def test_finalize_sync():
+    out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16)
+    pf = _make_pf(out_src)
+    dst = torch.zeros_like(out_src)
+    ret = _run(pf, dst, do_async=False)
+    assert ret is None
+    assert pf.buffer.calls[0]["async_with_compute_stream"] is False
+    assert torch.equal(dst, out_src)
diff --git a/tests/kernels/moe/test_deepep_v2_moe.py b/tests/kernels/moe/test_deepep_v2_moe.py
index f1e34d5bd780..2fbde2f86e06 100644
--- a/tests/kernels/moe/test_deepep_v2_moe.py
+++ b/tests/kernels/moe/test_deepep_v2_moe.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test DeepEP v2 (ElasticBuffer) dispatch-combine logic.
+"""Test DeepEP v2 (ElasticBuffer) dispatch-combine logic.
 Compares against a pure-PyTorch reference MoE implementation.
 """
 
diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py
index ae6c9685a459..cee0faa59229 100644
--- a/tests/kernels/moe/test_deepgemm.py
+++ b/tests/kernels/moe/test_deepgemm.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP).
+"""Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP).
 Compare DeepGEMM path against the Triton fallback inside vLLM's fused_experts.
 """
 
@@ -27,6 +26,10 @@
 )
 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
     deepgemm_moe_permute,
+    ep_gather,
+)
+from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
+    _fp8_workspace_shape,
 )
 from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
     TritonOrDeepGemmExperts,
@@ -44,6 +47,42 @@
 BLOCK_SIZE = [128, 128]
 
 
+def test_ep_gather_uses_64_bit_row_offsets():
+    hidden_size = 4096
+    source_row = 1 << 18
+    input_tensor = torch.empty(
+        (source_row + 1, hidden_size),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    input_tensor[source_row].fill_(1)
+    output = torch.empty((1, hidden_size), device="cuda", dtype=torch.bfloat16)
+
+    ep_gather(
+        input_tensor=input_tensor,
+        recv_topk_ids=torch.zeros((1, 1), device="cuda", dtype=torch.int64),
+        recv_topk_weight=torch.ones((1, 1), device="cuda", dtype=torch.float32),
+        input_index=torch.full((1, 1), source_row, device="cuda", dtype=torch.int32),
+        expert_map=None,
+        output_tensor=output,
+    )
+
+    torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("workspace_dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("num_columns", [2048, 6144, 6145])
+def test_fp8_workspace_shape(workspace_dtype, num_columns):
+    num_rows = 17
+    shape = _fp8_workspace_shape(num_rows, num_columns, workspace_dtype)
+
+    allocated_bytes = math.prod(shape) * workspace_dtype.itemsize
+    required_bytes = num_rows * num_columns * torch.float8_e4m3fn.itemsize
+
+    assert allocated_bytes >= required_bytes
+    assert allocated_bytes - required_bytes < num_rows * workspace_dtype.itemsize
+
+
 @pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
 def test_deepgemm_moe_permute_initializes_padding_scales(workspace_init):
     hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16)
@@ -79,8 +118,7 @@ def make_block_quant_fp8_weights(
     k: int,
     block_size: list[int],
 ):
-    """
-    Generate (w1, w2) expert weights and their per-block scale tensors
+    """Generate (w1, w2) expert weights and their per-block scale tensors
     in FP8 block-quantized format.
 
       w1 shape: (E, 2N, K)
@@ -121,8 +159,7 @@ def make_block_quant_fp8_weights(
 
 
 def run_single_case(m, n, k, topk, num_experts, block_size):
-    """
-    Run one (M,N,K) configuration on a single GPU and assert DeepGEMM ==
+    """Run one (M,N,K) configuration on a single GPU and assert DeepGEMM ==
     Triton baseline within tolerance.
     """
     tokens_bf16 = (
@@ -248,8 +285,7 @@ def make_mxfp4_weights(
     n: int,
     k: int,
 ):
-    """
-    Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales,
+    """Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales,
     plus BF16 reference weights for validation.
 
       w1 shape: (E, 2N, K//2) uint8    — packed FP4
@@ -310,8 +346,7 @@ def _bf16_moe_reference(x, w1, w2, topk_weights, topk_ids):
 
 
 def run_single_fp4_case(m, n, k, topk, num_experts):
-    """
-    Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert
+    """Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert
     DeepGEMM FP4 == BF16 reference within tolerance.
     """
     tokens_bf16 = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py
index 2e022ede3627..6329780f9dc9 100644
--- a/tests/kernels/moe/test_flashinfer.py
+++ b/tests/kernels/moe/test_flashinfer.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import multiprocessing
 from dataclasses import dataclass
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -22,6 +24,9 @@
 from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
     FlashInferExperts,
 )
+from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
+    view_as_block_major_k,
+)
 from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
     TrtLlmFp8ExpertsModular,
     TrtLlmFp8ExpertsMonolithic,
@@ -433,6 +438,276 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(
 
     assert w13_converted.shape[0] == num_experts
     assert w2_converted.shape[0] == num_experts
+    assert w13_converted.data_ptr() == w13.data_ptr()
+    assert w2_converted.data_ptr() == w2.data_ptr()
+
+
+@pytest.mark.parametrize("is_gated_act_gemm", [True, False])
+def test_convert_moe_weights_to_flashinfer_trtllm_block_layout_values(
+    is_gated_act_gemm,
+):
+    from flashinfer.fused_moe.core import (
+        _maybe_get_cached_w3_w1_permute_indices,
+        get_w2_permute_indices_with_cache,
+    )
+
+    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+        convert_moe_weights_to_flashinfer_trtllm_block_layout,
+    )
+
+    num_experts, intermediate, hidden = 2, 256, 256
+    w13_multiplier = 2 if is_gated_act_gemm else 1
+    w13 = torch.randn(
+        (num_experts, w13_multiplier * intermediate, hidden),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    w2 = torch.randn(
+        (num_experts, hidden, intermediate),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+
+    def _reference_block_layout(
+        weight: torch.Tensor,
+        is_w13: bool,
+        cache: dict[torch.Size, torch.Tensor],
+    ) -> torch.Tensor:
+        outputs = []
+        for expert in weight:
+            expert_uint8 = expert.view(torch.uint8)
+            if is_w13:
+                indices = _maybe_get_cached_w3_w1_permute_indices(
+                    cache,
+                    expert_uint8,
+                    128,
+                    is_gated_act_gemm=is_gated_act_gemm,
+                )
+                if is_gated_act_gemm:
+                    indices = (indices + expert_uint8.shape[0] // 2) % (
+                        expert_uint8.shape[0]
+                    )
+            else:
+                indices = get_w2_permute_indices_with_cache(
+                    cache,
+                    expert_uint8,
+                    128,
+                )
+            rows, cols = expert_uint8.shape
+            blocks = expert_uint8.view(rows, cols // 128, 128).permute(1, 0, 2)
+            outputs.append(torch.index_select(blocks, 1, indices.to(weight.device)))
+        return torch.stack(outputs).view(torch.bfloat16)
+
+    reference_cache: dict[torch.Size, torch.Tensor] = {}
+    expected_w13 = _reference_block_layout(w13, is_w13=True, cache=reference_cache)
+    expected_w2 = _reference_block_layout(w2, is_w13=False, cache=reference_cache)
+    w13_ptr = w13.data_ptr()
+    w2_ptr = w2.data_ptr()
+
+    actual_w13, actual_w2 = convert_moe_weights_to_flashinfer_trtllm_block_layout(
+        {},
+        w13,
+        w2,
+        is_gated_act_gemm=is_gated_act_gemm,
+    )
+
+    assert actual_w13.data_ptr() == w13_ptr
+    assert actual_w2.data_ptr() == w2_ptr
+    assert torch.equal(actual_w13, expected_w13)
+    assert torch.equal(actual_w2, expected_w2)
+
+
+def _make_unquantized_flashinfer_test_layer(
+    monkeypatch, intermediate, is_gated, *, device="cuda"
+):
+    from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
+        UnquantizedMoeBackend,
+    )
+    from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+        UnquantizedFusedMoEMethod,
+    )
+
+    moe_config = make_dummy_moe_config(
+        num_experts=2,
+        hidden_dim=256,
+        intermediate_size=intermediate,
+        activation=MoEActivation.SILU if is_gated else MoEActivation.RELU2_NO_MUL,
+    )
+    moe_config.intermediate_size_per_partition = (intermediate + 127) // 128 * 128
+    method = object.__new__(UnquantizedFusedMoEMethod)
+    method.moe = moe_config
+    method.unquantized_backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM
+    method.moe_kernel = None
+    mock_kernel = MagicMock()
+    monkeypatch.setattr(
+        method,
+        "_init_moe_kernel",
+        lambda _: setattr(method, "moe_kernel", mock_kernel),
+    )
+
+    layer = torch.nn.Module()
+    layer.moe_config = moe_config
+    with torch.device(device):
+        method.create_weights(
+            layer,
+            num_experts=2,
+            hidden_size=256,
+            intermediate_size_per_partition=moe_config.intermediate_size_per_partition,
+            params_dtype=torch.bfloat16,
+        )
+    return method, layer
+
+
+@pytest.mark.parametrize("intermediate", [192, 256])
+@pytest.mark.parametrize("is_gated", [True, False])
+def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed(
+    monkeypatch, intermediate, is_gated
+):
+    """Each raw reload clears padding left dirty by prior in-place conversion."""
+    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+        convert_moe_weights_to_flashinfer_trtllm_block_layout,
+    )
+
+    method, layer = _make_unquantized_flashinfer_test_layer(
+        monkeypatch, intermediate, is_gated
+    )
+    w13_shape, w2_shape = layer.w13_weight.shape, layer.w2_weight.shape
+    padded = w2_shape[-1]
+    w13_ptr = layer.w13_weight.data_ptr()
+    w2_ptr = layer.w2_weight.data_ptr()
+    load_numel = (
+        layer.w13_weight.weight_loader_numel,
+        layer.w2_weight.weight_loader_numel,
+    )
+
+    for _ in range(2):
+        reloaded_w13 = torch.randn_like(layer.w13_weight)
+        reloaded_w2 = torch.randn_like(layer.w2_weight)
+        reloaded_w13[:, intermediate:padded].zero_()
+        if is_gated:
+            reloaded_w13[:, padded + intermediate :].zero_()
+        reloaded_w2[:, :, intermediate:].zero_()
+        expected_w13, expected_w2 = (
+            convert_moe_weights_to_flashinfer_trtllm_block_layout(
+                {},
+                reloaded_w13.clone(),
+                reloaded_w2.clone(),
+                is_gated_act_gemm=is_gated,
+            )
+        )
+
+        # The loader writes only logical checkpoint slices, not padding.
+        layer.w13_weight.fill_(float("nan"))
+        layer.w2_weight.fill_(float("nan"))
+        layer.w13_weight[:, :intermediate].copy_(reloaded_w13[:, :intermediate])
+        if is_gated:
+            layer.w13_weight[:, padded : padded + intermediate].copy_(
+                reloaded_w13[:, padded : padded + intermediate]
+            )
+        layer.w2_weight[:, :, :intermediate].copy_(reloaded_w2[:, :, :intermediate])
+        method.process_weights_after_loading(layer)
+
+        assert layer.w13_weight.shape == w13_shape
+        assert layer.w2_weight.shape == w2_shape
+        assert layer.w13_weight.data_ptr() == w13_ptr
+        assert layer.w2_weight.data_ptr() == w2_ptr
+        assert (
+            layer.w13_weight.weight_loader_numel,
+            layer.w2_weight.weight_loader_numel,
+        ) == load_numel
+        kernel_w13 = view_as_block_major_k(layer.w13_weight)
+        kernel_w2 = view_as_block_major_k(layer.w2_weight)
+        assert torch.equal(kernel_w13, expected_w13)
+        assert torch.equal(kernel_w2, expected_w2)
+
+
+def _check_flashinfer_ipc_weights(entries, expected, mode, device_index):
+    from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader
+    from vllm.model_executor.model_loader.weight_cache.protocol import (
+        WeightCacheState,
+    )
+    from vllm.model_executor.utils import weights_already_processed
+
+    torch.accelerator.set_device_index(device_index)
+    with pytest.MonkeyPatch.context() as monkeypatch:
+        method, layer = _make_unquantized_flashinfer_test_layer(
+            monkeypatch, 192, is_gated=True, device="meta"
+        )
+        loader = object.__new__(IpcModelLoader)
+        loader.mode = mode
+        loader._apply_entries(layer, WeightCacheState(entries, {}, {}), device_index)
+        pointers = (layer.w13_weight.data_ptr(), layer.w2_weight.data_ptr())
+
+        # Like ipc_cache: a fresh method, only tensor metadata, no _setup_kernel.
+        with weights_already_processed():
+            method.process_weights_after_loading(layer)
+        assert method.moe_kernel is not None
+        actual = (
+            view_as_block_major_k(layer.w13_weight),
+            view_as_block_major_k(layer.w2_weight),
+        )
+        for weight, reference, pointer in zip(actual, expected, pointers):
+            assert weight.data_ptr() == pointer
+            assert torch.equal(weight.cpu(), reference)
+        # Cross-process addresses need not match; writes prove sharing/isolation.
+        with torch.no_grad():
+            for weight in actual:
+                weight.fill_(1.0)
+        torch.accelerator.synchronize()
+
+
+@pytest.mark.parametrize("cache_ndim", [3, 4])
+@pytest.mark.parametrize("mode", ["copy", "zero_copy"])
+def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state(
+    monkeypatch, cache_ndim, mode
+):
+    """A spawned IPC consumer restores views without transient method state."""
+    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+        convert_moe_weights_to_flashinfer_trtllm_block_layout,
+    )
+    from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry
+
+    method, layer = _make_unquantized_flashinfer_test_layer(
+        monkeypatch, 192, is_gated=True, device="meta"
+    )
+    packed = convert_moe_weights_to_flashinfer_trtllm_block_layout(
+        {},
+        torch.randn_like(layer.w13_weight, device="cuda"),
+        torch.randn_like(layer.w2_weight, device="cuda"),
+    )
+    expected = [weight.cpu() for weight in packed]
+    entries = {}
+    for name, weight in zip(("w13_weight", "w2_weight"), packed):
+        cached = weight.view_as(getattr(layer, name)) if cache_ndim == 3 else weight
+        entries[name] = TensorEntry.from_tensor(cached, kind="param")
+    consumer = multiprocessing.get_context("spawn").Process(
+        target=_check_flashinfer_ipc_weights,
+        args=(
+            entries,
+            expected,
+            mode,
+            torch.accelerator.current_device_index(),
+        ),
+    )
+    consumer.start()
+    try:
+        # Keep producer allocations alive until the consumer releases its views.
+        consumer.join(timeout=180)
+        assert not consumer.is_alive(), "IPC consumer timed out"
+        assert consumer.exitcode == 0, "IPC consumer failed; see child traceback"
+        for source, reference in zip(packed, expected):
+            expected_source = (
+                torch.ones_like(reference) if mode == "zero_copy" else reference
+            )
+            assert torch.equal(source.cpu(), expected_source)
+    finally:
+        if consumer.is_alive():
+            consumer.terminate()
+            consumer.join(timeout=10)
+            if consumer.is_alive():
+                consumer.kill()
+                consumer.join(timeout=10)
+        consumer.close()
 
 
 @pytest.mark.parametrize(
diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py
index d8e449ec0aff..3d637258d3f2 100644
--- a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py
+++ b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py
@@ -83,7 +83,7 @@ def _reference_activation(
     beta: float | None,
     limit: float | None,
 ):
-    """vLLM's own op for this activation, so the reference is not re-derived."""
+    """VLLM's own op for this activation, so the reference is not re-derived."""
     if activation == MoEActivation.RELU2_NO_MUL:
         return lambda x: torch.square(torch.relu(x))
     if activation == MoEActivation.SWIGLUOAI:
diff --git a/tests/kernels/moe/test_fp8_tp_loading.py b/tests/kernels/moe/test_fp8_tp_loading.py
new file mode 100644
index 000000000000..680ccb812cca
--- /dev/null
+++ b/tests/kernels/moe/test_fp8_tp_loading.py
@@ -0,0 +1,311 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm.distributed as distributed
+from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm.model_executor.layers.fused_moe import (
+    unquantized_fused_moe_method as unquantized_module,
+)
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
+from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+from vllm.model_executor.layers.quantization import fp8 as fp8_module
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+
+def _make_fp8_tp_experts(
+    monkeypatch,
+    tp_size,
+    tp_rank,
+    backend="flashinfer_trtllm",
+    *,
+    mock_backend=True,
+    num_experts=2,
+    hidden_dim=256,
+    intermediate_size=640,
+    quant_config=None,
+    config_overrides=None,
+):
+    # Exercise allocation and the public weight loader without selecting a GPU kernel.
+    if mock_backend:
+        monkeypatch.setattr(
+            fp8_module, "select_fp8_moe_backend", lambda **kwargs: (None, None)
+        )
+        monkeypatch.setattr(
+            unquantized_module,
+            "select_unquantized_moe_backend",
+            lambda **kwargs: (None, None),
+        )
+    monkeypatch.setattr(
+        fp8_module, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    # The block-shape validator imports this function from vllm.distributed.
+    monkeypatch.setattr(
+        distributed, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    config = make_dummy_moe_config(
+        num_experts=num_experts,
+        hidden_dim=hidden_dim,
+        intermediate_size=intermediate_size,
+        experts_per_token=min(num_experts, 10),
+    )
+    config.routing_method = RoutingMethodType.RenormalizeNaive
+    config.moe_parallel_config.tp_size = tp_size
+    config.moe_parallel_config.tp_rank = tp_rank
+    config.intermediate_size_per_partition = intermediate_size // tp_size
+    config.intermediate_size_per_partition_unpadded = intermediate_size // tp_size
+    config.moe_backend = backend
+    for field, value in (config_overrides or {}).items():
+        setattr(config, field, value)
+    return RoutedExperts(
+        "model.layers.0.mlp.experts",
+        torch.bfloat16,
+        config,
+        quant_config
+        or Fp8Config(is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128]),
+        expert_map_manager=SimpleNamespace(
+            local_num_experts=num_experts,
+            placement_strategy="linear",
+            expert_map=None,
+            expert_mask=None,
+            routing_tables=None,
+            map_global_to_local=lambda index: index,
+        ),
+    )
+
+
+@pytest.mark.parametrize("tp_size", [2, 4, 8])
+@pytest.mark.parametrize("batched", [False, True])
+@pytest.mark.parametrize("intermediate_size", [640, 896])
+def test_fp8_block_aligned_tp_preserves_checkpoint(
+    monkeypatch, tp_size, batched, intermediate_size
+):
+    """All ranks reconstruct the original dequantized projections, including
+    padding-only ranks. Reloading must clear stale weights and scales.
+    """
+    generator = torch.Generator().manual_seed(42)
+    weights = {
+        "w1": torch.randn(2, intermediate_size, 256, generator=generator).to(
+            torch.float8_e4m3fn
+        ),
+        "w3": torch.randn(2, intermediate_size, 256, generator=generator).to(
+            torch.float8_e4m3fn
+        ),
+        "w2": torch.randn(2, 256, intermediate_size, generator=generator).to(
+            torch.float8_e4m3fn
+        ),
+    }
+    scales = {
+        name: torch.rand(2, w.shape[1] // 128, w.shape[2] // 128, generator=generator)
+        + 0.1
+        for name, w in weights.items()
+    }
+    reconstructed: dict[str, list[torch.Tensor]] = {name: [] for name in weights}
+    num_blocks = intermediate_size // 128
+    for rank in range(tp_size):
+        layer = _make_fp8_tp_experts(
+            monkeypatch, tp_size, rank, intermediate_size=intermediate_size
+        )
+        assert layer.quant_method.weight_scale_refine is None
+        assert layer.quant_method.moe_block_shape == [128, 128]
+        width = layer.moe_config.intermediate_size_per_partition
+        assert width == ((num_blocks + tp_size - 1) // tp_size) * 128
+        for name in weights:
+            prefix = "w2" if name == "w2" else "w13"
+            for suffix, checkpoint in (
+                ("weight", weights[name]),
+                ("weight_scale_inv", scales[name]),
+            ):
+                param_name = f"{prefix}_{suffix}"
+                param = getattr(layer, param_name)
+                # Also reload after dirtying the destination, to catch padding leaks.
+                for _ in range(2):
+                    if name != "w3":
+                        param.data.fill_(7)
+                    if batched:
+                        assert param.weight_loader(
+                            param, checkpoint, param_name, name, 0, return_success=True
+                        )
+                    else:
+                        for expert in range(2):
+                            assert param.weight_loader(
+                                param,
+                                checkpoint[expert],
+                                param_name,
+                                name,
+                                expert,
+                                return_success=True,
+                            )
+            w = getattr(layer, f"{prefix}_weight").float()
+            s = getattr(layer, f"{prefix}_weight_scale_inv")
+            if name in ("w1", "w3"):
+                half = 0 if name == "w1" else 1
+                w = w.chunk(2, dim=1)[half]
+                s = s.chunk(2, dim=1)[half]
+            dim = 2 if name == "w2" else 1
+            valid = max(0, min(width, intermediate_size - rank * width))
+            assert torch.count_nonzero(w.narrow(dim, valid, width - valid)) == 0
+            assert torch.all(s.narrow(dim, valid // 128, (width - valid) // 128) == 1)
+            dequant = w * s.repeat_interleave(128, dim=1).repeat_interleave(128, dim=2)
+            reconstructed[name].append(dequant.narrow(dim, 0, valid))
+    for name, parts in reconstructed.items():
+        expected = weights[name].float() * scales[name].repeat_interleave(
+            128, dim=1
+        ).repeat_interleave(128, dim=2)
+        assert torch.equal(torch.cat(parts, dim=2 if name == "w2" else 1), expected)
+
+
+@pytest.mark.parametrize("backend", ["auto", "triton"])
+@pytest.mark.parametrize("tp_size,block", [(2, 64), (4, 32)])
+def test_fp8_tp_default_keeps_refined_layout(monkeypatch, backend, tp_size, block):
+    layer = _make_fp8_tp_experts(monkeypatch, tp_size, 0, backend)
+    assert layer.moe_config.intermediate_size_per_partition == 640 // tp_size
+    assert layer.quant_method.moe_block_shape == [block, block]
+
+
+@pytest.mark.parametrize(
+    "field,value",
+    [
+        ("hidden_dim", 192),
+        ("intermediate_size", 672),
+        ("is_lora_enabled", True),
+        ("has_bias", True),
+    ],
+)
+def test_fp8_block_aligned_tp_rejects_unsupported_layout(monkeypatch, field, value):
+    with pytest.raises(ValueError, match="Block-aligned FP8 TP sharding requires"):
+        _make_fp8_tp_experts(monkeypatch, 4, 0, config_overrides={field: value})
+
+
+@pytest.mark.parametrize("tp_size", [1, 5])
+def test_fp8_aligned_tp_keeps_original_layout(monkeypatch, tp_size):
+    layer = _make_fp8_tp_experts(monkeypatch, tp_size, 0)
+    assert layer.moe_config.intermediate_size_per_partition == 640 // tp_size
+    assert not layer.moe_config.tp_shard_with_padding
+    assert layer.quant_method.weight_scale_refine is None
+
+
+def test_fp8_skipped_layer_keeps_original_tp_layout(monkeypatch, default_vllm_config):
+    layer = _make_fp8_tp_experts(
+        monkeypatch,
+        4,
+        0,
+        quant_config=Fp8Config(
+            is_checkpoint_fp8_serialized=True,
+            weight_block_size=[128, 128],
+            ignored_layers=["model.layers.0.mlp.experts"],
+        ),
+    )
+    assert isinstance(layer.quant_method, unquantized_module.UnquantizedFusedMoEMethod)
+    assert not layer.moe_config.tp_shard_with_padding
+    assert layer.moe_config.intermediate_size_per_partition == 160
+    assert layer.w13_weight.shape == (2, 320, 256)
+    assert layer.w2_weight.shape == (2, 256, 160)
+
+
+def test_fp8_block_aligned_tp_rejects_presharded_weights(monkeypatch):
+    layer = _make_fp8_tp_experts(monkeypatch, 4, 0)
+    with pytest.raises(ValueError, match="unsharded checkpoint"):
+        layer.w13_weight.weight_loader(
+            layer.w13_weight,
+            torch.zeros(160, 256, dtype=torch.float8_e4m3fn),
+            "w13_weight",
+            "w1",
+            0,
+        )
+
+
+@pytest.mark.parametrize("tp_size", [2, 4, 8])
+@pytest.mark.parametrize("num_tokens", [1, 17])
+@torch.inference_mode()
+def test_fp8_block_aligned_tp_flashinfer_matches_unsharded(
+    monkeypatch,
+    tp_size,
+    num_tokens,
+    dist_init,
+    workspace_init,
+):
+    """Compare sharded FI with unsharded FI, reporting the cross-backend
+    difference separately because activation quantization can differ.
+    """
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
+    from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("Requires datacenter Blackwell and FlashInfer TRTLLM")
+    torch.manual_seed(42)
+    e, h, n, topk = 16, 2560, 640, 10
+    with torch.device("cuda"), set_current_vllm_config(VllmConfig()):
+        x = torch.randn(num_tokens, h, dtype=torch.bfloat16) / 10
+        logits = torch.randn(num_tokens, e, dtype=torch.float32)
+        probabilities = logits.softmax(dim=-1)
+        topk_weights, topk_ids = probabilities.topk(topk, dim=-1)
+        topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
+        weights = {
+            "w1": torch.randn(e, n, h).to(torch.float8_e4m3fn),
+            "w3": torch.randn(e, n, h).to(torch.float8_e4m3fn),
+            "w2": torch.randn(e, h, n).to(torch.float8_e4m3fn),
+        }
+        scales = {
+            name: (torch.rand(e, w.shape[1] // 128, w.shape[2] // 128) + 0.5) / 32
+            for name, w in weights.items()
+        }
+        ref = fused_experts(
+            x,
+            torch.cat([weights["w1"], weights["w3"]], dim=1),
+            weights["w2"],
+            topk_weights,
+            topk_ids,
+            quant_config=fp8_w8a8_moe_quant_config(
+                w1_scale=torch.cat([scales["w1"], scales["w3"]], dim=1),
+                w2_scale=scales["w2"],
+                block_shape=[128, 128],
+            ),
+        )
+        total = torch.zeros_like(ref, dtype=torch.float32)
+        fi_reference = None
+        for rank in range(-1, tp_size):
+            layer = _make_fp8_tp_experts(
+                monkeypatch,
+                1 if rank == -1 else tp_size,
+                max(rank, 0),
+                mock_backend=False,
+                num_experts=e,
+                hidden_dim=h,
+            )
+            for name in weights:
+                prefix = "w2" if name == "w2" else "w13"
+                for suffix, checkpoint in (
+                    ("weight", weights[name]),
+                    ("weight_scale_inv", scales[name]),
+                ):
+                    param_name = f"{prefix}_{suffix}"
+                    param = getattr(layer, param_name)
+                    param.weight_loader(param, checkpoint, param_name, name, 0)
+            method = layer.quant_method
+            method.process_weights_after_loading(layer)
+            assert method.is_monolithic
+            out = method.apply_monolithic(layer, x, logits)
+            assert torch.isfinite(out).all()
+            if rank == -1:
+                fi_reference = out.float().clone()
+            else:
+                total += out.float()
+        # Different FP8 GEMM reduction orders need not be bitwise identical.
+        assert fi_reference is not None
+        relative_l2 = (total - fi_reference).norm() / fi_reference.norm().clamp_min(
+            1e-8
+        )
+        cross_backend_l2 = (
+            fi_reference - ref.float()
+        ).norm() / ref.float().norm().clamp_min(1e-8)
+        print(
+            f"TP{tp_size} M={num_tokens}: shard_l2={relative_l2.item():.6f}, "
+            f"unsharded_fi_vs_triton_l2={cross_backend_l2.item():.6f}"
+        )
+        assert relative_l2 < 0.01, relative_l2.item()
diff --git a/tests/kernels/moe/test_fused_topk.py b/tests/kernels/moe/test_fused_topk.py
index 9c196a884d9f..18d1d9b4dfb0 100644
--- a/tests/kernels/moe/test_fused_topk.py
+++ b/tests/kernels/moe/test_fused_topk.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the MoE fused topk kernel
+"""Tests for the MoE fused topk kernel.
 
 Run `pytest tests/kernels/moe/test_fused_topk.py`.
 """
diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py
index 4afcba706ccd..c1d83d8349bf 100644
--- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py
+++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py
@@ -7,7 +7,7 @@
 import torch.nn.functional as F
 
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
 
 if not has_triton_kernels():
     pytest.skip(
@@ -15,9 +15,15 @@
         allow_module_level=True,
     )
 
-import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
 import triton_kernels.swiglu
-from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
+
+if get_triton_kernels_version() == "3.8":
+    # 3.8: matmul_ogs -> matmul, matmul_ogs_details -> matmul_details
+    import triton_kernels.matmul_details.opt_flags as opt_flags
+    from triton_kernels.matmul import FlexCtx, PrecisionConfig
+else:
+    import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
+    from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
 from triton_kernels.numerics import InFlexData
 from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp
 from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
@@ -28,9 +34,10 @@
 from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
     triton_kernel_moe_forward,
 )
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs
 from vllm.utils.math_utils import round_up
 
-from .utils import shuffle_weight
+from .utils import mxfp4_w_layouts, shuffle_weight
 
 
 def deshuffle(w: torch.Tensor):
@@ -133,12 +140,12 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int):
 
         x_tri = F.pad(x_tri, (0, x_pad, 0, 0), mode="constant", value=0)
 
-        w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
-        w_scale_layout, w_scale_layout_opts = (
-            layout.make_default_matmul_mxfp4_w_scale_layout(
-                mx_axis=1, num_warps=num_warps
-            )
-        )
+        (
+            w_layout,
+            w_layout_opts,
+            w_scale_layout,
+            w_scale_layout_opts,
+        ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps)
 
         w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1)
         w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, torch.bfloat16, axis=1)
@@ -165,10 +172,10 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int):
         )
 
         pc1 = PrecisionConfig(
-            weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
+            **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
         )
         pc2 = PrecisionConfig(
-            weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
+            **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
         )
 
         # tucuate so the rest can run properly
@@ -281,8 +288,6 @@ class Case:
 @pytest.mark.parametrize("num_token", [2])
 @pytest.mark.parametrize("tp", [1, 2, 4, 8])
 def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init):
-    from triton_kernels.tensor_details import layout
-
     if not hasattr(layout, "make_default_matmul_mxfp4_w_layout"):
         pytest.skip("make_default_matmul_mxfp4_w_layout not available")
 
@@ -384,10 +389,14 @@ def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk):
 
     make_routing_data = gptoss_moe.make_routing_data
     routing_data_from_sparse_topk = gptoss_moe.routing_data_from_sparse_topk
-    use_legacy_triton_kernels = gptoss_moe.use_legacy_triton_kernels
 
-    if use_legacy_triton_kernels:
+    if gptoss_moe.triton_kernels_version == "3.5.1":
         pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+")
+    if gptoss_moe.triton_kernels_version == "3.8":
+        pytest.skip(
+            "3.8 rebuilds routing from raw ids, so routing_data_from_sparse_topk "
+            "just forwards to make_routing_data and the comparison is vacuous"
+        )
 
     from triton_kernels.topk import topk as topk_fn
 
diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py
index 1ab8550d64e8..2004feae6595 100644
--- a/tests/kernels/moe/test_grouped_topk.py
+++ b/tests/kernels/moe/test_grouped_topk.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the MoE grouped topk kernel
+"""Tests for the MoE grouped topk kernel.
 
 Run `pytest tests/kernels/moe/test_grouped_topk.py`.
 """
@@ -15,8 +15,10 @@
     get_cached_compilation_config,
     set_current_vllm_config,
 )
+from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
     GroupedTopk,
+    GroupedTopKRouter,
     fused_grouped_topk,
 )
 from vllm.platforms import current_platform
@@ -70,6 +72,36 @@ def _single_group_reference(
     return values, indices.to(torch.int32)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+@pytest.mark.parametrize("skip_padding", [False, True])
+def test_kimi_grouped_topk_masks_padding_when_enabled(skip_padding: bool):
+    hidden_states = torch.randn((2, 16), device="cuda")
+    gating_output = torch.randn((2, 896), device="cuda")
+    is_padding = torch.tensor([False, True], device="cuda")
+    router = GroupedTopKRouter(
+        top_k=16,
+        global_num_experts=896,
+        num_expert_group=1,
+        topk_group=1,
+        skip_padding=skip_padding,
+    )
+
+    with set_forward_context(None, VllmConfig(), is_padding=is_padding):
+        topk_weights, topk_ids = router._compute_routing(
+            hidden_states, gating_output, torch.int32
+        )
+
+    assert torch.all(topk_ids[0] >= 0)
+    if skip_padding:
+        assert torch.all(topk_ids[1] == -1)
+        assert torch.all(topk_weights[1] == 0)
+    else:
+        assert torch.all(topk_ids[1] >= 0)
+        assert torch.all(topk_weights[1] > 0)
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
 )
diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py
index 6756d479febb..97041e24b9bb 100644
--- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py
+++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py
@@ -31,6 +31,7 @@ def mxint4_quantize(
     Returns:
         - uint8 packed (2 INT4/byte): [..., k//2] - stores SIGNED INT4 [-8, 7]
         - scales in BF16: [..., k//sf_vec_size]
+
     """
     x_reshaped = x.reshape(-1, sf_vec_size)
     x_max = x_reshaped.max(dim=-1, keepdim=True)[0].to(torch.float32)
@@ -61,6 +62,7 @@ def mxint4_quantize_moe_weights(
     Returns:
         - weights_mxint4: Quantized weights [e, n, k//2] uint8
         - scales_mxint4: Quantization scales [e, n, k//group_size] bf16
+
     """
     e = weights_bf16.shape[0]
     weight_list = []
@@ -101,6 +103,7 @@ def marlin_quantize_moe_weights(
     Returns:
         - weights_marlin: Marlin quantized weights [e, k//8, n] int32
         - scales_marlin: Marlin quantization scales [e, k//group_size, n] bf16
+
     """
     from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
         marlin_quantize,
diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py
index 0c0e1d61f90e..9bfd3b5a2e19 100644
--- a/tests/kernels/moe/test_modular_kernel_combinations.py
+++ b/tests/kernels/moe/test_modular_kernel_combinations.py
@@ -12,6 +12,7 @@
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm
@@ -37,6 +38,7 @@
     ProcessGroupInfo,
     parallel_launch_with_config,
 )
+from .utils import check_accuracy
 
 has_any_multi_gpu_package = (
     has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe()
@@ -47,12 +49,6 @@
     reason="Requires deep_ep or deep_gemm or flashinfer packages",
 )
 
-if current_platform.is_fp8_fnuz():
-    pytest.skip(
-        "Tests in this file require float8_e4m3fn and platform does not support",
-        allow_module_level=True,
-    )
-
 
 def format_result(verbose, msg, ex=None):
     if ex is not None:
@@ -70,6 +66,44 @@ def format_result(verbose, msg, ex=None):
         print(".", end="")
 
 
+def assert_aiter_quant_scheme_case(config: Config) -> None:
+    """Make the AITER (weight_quant_key, activation_quant_key) pair this
+    config exercises explicit, instead of AiterExperts being reached only
+    indirectly through the general quant-config sweep.
+    See https://github.com/vllm-project/vllm/issues/54966."""
+    fe_cls = config.fused_experts_type
+    if getattr(fe_cls, "__name__", "") != "AiterExperts":
+        return
+
+    if config.quant_config is None:
+        w_key, a_key = None, None
+    else:
+        w_key, a_key = config.fp8_quant_key_pair()
+
+    assert fe_cls._supports_quant_scheme(w_key, a_key), (
+        f"AITER case (weight_key={w_key}, activation_key={a_key}) reached "
+        "the modular-kernel harness, but AiterExperts._supports_quant_scheme "
+        "does not declare it supported."
+    )
+    print(f"[AITER case] weight_key={w_key}, activation_key={a_key}")
+
+
+def assert_aiter_activation_case(config: Config) -> None:
+    """Make the AITER activation this config exercises explicit, instead of
+    AiterExperts being reached only indirectly through the general
+    activation sweep. See https://github.com/vllm-project/vllm/issues/54966."""
+    fe_cls = config.fused_experts_type
+    if getattr(fe_cls, "__name__", "") != "AiterExperts":
+        return
+
+    assert fe_cls._supports_activation(config.activation), (
+        f"AITER case activation={config.activation} reached the "
+        "modular-kernel harness, but AiterExperts._supports_activation "
+        "does not declare it supported."
+    )
+    print(f"[AITER case] activation={config.activation}")
+
+
 def rank_worker(
     pgi: ProcessGroupInfo,
     vllm_config: VllmConfig,
@@ -125,6 +159,26 @@ def rank_worker(
                 count -= 1
                 continue
 
+            # Skip unsupported: AITER x DeepEP-HT/Mori dispatch crashes
+            # with an illegal memory access at world_size>1 on gfx942.
+            # https://github.com/vllm-project/vllm/issues/57029
+            if (
+                config.world_size > 1
+                and getattr(config.fused_experts_type, "__name__", "") == "AiterExperts"
+                and getattr(config.prepare_finalize_type, "__name__", "")
+                in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize")
+            ):
+                print(
+                    f"Skipping[{pgi.rank}]: m={m}, topk={topk}"
+                    " (AITER x DeepEP-HT/Mori illegal memory access,"
+                    " https://github.com/vllm-project/vllm/issues/57029)"
+                )
+                count -= 1
+                continue
+
+            assert_aiter_quant_scheme_case(config)
+            assert_aiter_activation_case(config)
+
             # modular kernel out
             mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors)
 
@@ -152,32 +206,7 @@ def rank_worker(
                 and config.quant_config is not None
             )
             if is_aiter_fp8:
-                diff = (ref_out - mk_out).abs()
-                n_total = diff.numel()
-                max_diff = diff.max().item()
-                n_exceed = int((diff > atol).sum().item())
-                pct_exceed = n_exceed / n_total * 100
-                # FP8 hw matmul vs f32 reference: up to ~4% of
-                # elements may exceed base tolerance, but max
-                # error should stay within 3x base tolerance.
-                max_pct_allowed = 5.0
-                relaxed_atol = atol * 4
-                print(
-                    f"[AITER FP8 precision] "
-                    f"max_diff={max_diff:.6f}, "
-                    f"exceed_atol={n_exceed}/{n_total} "
-                    f"({pct_exceed:.4f}%), "
-                    f"max_pct_allowed={max_pct_allowed}%, "
-                    f"relaxed_limit={relaxed_atol}"
-                )
-                assert pct_exceed <= max_pct_allowed, (
-                    f"AITER FP8: {pct_exceed:.2f}% elements exceed "
-                    f"atol={atol} (max allowed {max_pct_allowed}%)"
-                )
-                assert max_diff <= relaxed_atol, (
-                    f"AITER FP8: max_diff={max_diff:.6f} exceeds "
-                    f"relaxed limit {relaxed_atol}"
-                )
+                check_accuracy(ref_out, mk_out, atol=atol, rtol=rtol, percent=0.9)
             else:
                 torch.testing.assert_close(ref_out, mk_out, atol=atol, rtol=rtol)
             format_result(verbose, config.describe())
@@ -214,6 +243,10 @@ def run(config: Config, verbose: bool):
 TOPKs = [4, 1]
 Es = [32]
 DTYPEs = [torch.bfloat16]
+MK_ACTIVATIONS = [
+    MoEActivation.SILU,
+    MoEActivation.GELU,
+]
 
 
 def is_nyi_config(config: Config) -> bool:
@@ -225,7 +258,14 @@ def is_nyi_config(config: Config) -> bool:
         unsupported_quant_config = (
             config.is_per_act_token_quant + config.is_per_out_ch_quant
         ) == 1
-        return unsupported_quant_config
+        if unsupported_quant_config:
+            return True
+
+    if config.activation != MoEActivation.SILU:
+        if getattr(config.fused_experts_type, "__name__", "") != "AiterExperts":
+            return True  # AITER-only for this axis, for now
+        if config.quant_dtype is not None:
+            return True  # unquantized-only for this axis, for now
 
     return False
 
@@ -236,12 +276,13 @@ def generate_valid_test_cases(
     cases = []
     total = 0
 
-    for k, n, e, dtype, quant_config, combination in product(
+    for k, n, e, dtype, quant_config, activation, combination in product(
         Ks,
         Ns,
         Es,
         DTYPEs,
         MK_QUANT_CONFIGS,
+        MK_ACTIVATIONS,
         product(prepare_finalize_types, MK_FUSED_EXPERT_TYPES),
     ):
         total = total + 1
@@ -254,6 +295,7 @@ def generate_valid_test_cases(
             topks=TOPKs,
             dtype=dtype,
             quant_config=quant_config,
+            activation=activation,
             prepare_finalize_type=combination[0],
             fused_experts_type=combination[1],
             world_size=world_size,
@@ -281,6 +323,7 @@ def generate_valid_test_cases(
                 e,
                 dtype,
                 quant_config,
+                activation,
                 combination[0],
                 combination[1],
                 world_size,
@@ -293,7 +336,8 @@ def generate_valid_test_cases(
 
 
 @pytest.mark.parametrize(
-    "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,world_size",
+    "k,n,e,dtype,quant_config,activation,"
+    "prepare_finalize_type,fused_experts_type,world_size",
     generate_valid_test_cases(
         world_size=2, prepare_finalize_types=MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
     ),
@@ -305,6 +349,7 @@ def test_modular_kernel_combinations_multigpu(
     e: int,
     dtype: torch.dtype,
     quant_config: TestMoEQuantConfig | None,
+    activation: MoEActivation,
     prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
     fused_experts_type: mk.FusedMoEExperts,
     world_size: int,
@@ -325,6 +370,7 @@ def test_modular_kernel_combinations_multigpu(
         topks=TOPKs,
         dtype=dtype,
         quant_config=quant_config,
+        activation=activation,
         prepare_finalize_type=prepare_finalize_type,
         fused_experts_type=fused_experts_type,
         world_size=world_size,
@@ -334,7 +380,8 @@ def test_modular_kernel_combinations_multigpu(
 
 
 @pytest.mark.parametrize(
-    "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,world_size",
+    "k,n,e,dtype,quant_config,activation,"
+    "prepare_finalize_type,fused_experts_type,world_size",
     generate_valid_test_cases(
         world_size=1, prepare_finalize_types=MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES
     ),
@@ -345,6 +392,7 @@ def test_modular_kernel_combinations_singlegpu(
     e: int,
     dtype: torch.dtype,
     quant_config: TestMoEQuantConfig | None,
+    activation: MoEActivation,
     prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
     fused_experts_type: mk.FusedMoEExperts,
     world_size: int,
@@ -361,6 +409,7 @@ def test_modular_kernel_combinations_singlegpu(
         topks=TOPKs,
         dtype=dtype,
         quant_config=quant_config,
+        activation=activation,
         prepare_finalize_type=prepare_finalize_type,
         fused_experts_type=fused_experts_type,
         world_size=world_size,
diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py
index d0fd4df23ccd..9665079e373d 100644
--- a/tests/kernels/moe/test_modular_oai_triton_moe.py
+++ b/tests/kernels/moe/test_modular_oai_triton_moe.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test modular OAI Triton MoE
-"""
+"""Test modular OAI Triton MoE."""
 
 from __future__ import annotations
 
@@ -12,7 +10,7 @@
 
 from tests.utils import wait_for_gpu_memory_to_clear
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
-from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
 
 if not has_triton_kernels():
     pytest.skip(
@@ -20,11 +18,14 @@
         allow_module_level=True,
     )
 
-from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
+if get_triton_kernels_version() == "3.8":
+    # 3.8: matmul_ogs -> matmul
+    from triton_kernels.matmul import FlexCtx, PrecisionConfig
+else:
+    from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
 from triton_kernels.numerics import InFlexData
 from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp
 from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
-from triton_kernels.tensor_details import layout
 from triton_kernels.testing import assert_close
 
 from vllm.config import VllmConfig, set_current_vllm_config
@@ -37,11 +38,12 @@
     UnfusedOAITritonExperts,
 )
 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import round_up
 from vllm.utils.torch_utils import set_random_seed
 
-from .utils import make_dummy_moe_config, shuffle_weight
+from .utils import make_dummy_moe_config, mxfp4_w_layouts, shuffle_weight
 
 MNK = [
     (1, 512, 384),
@@ -132,10 +134,12 @@ def make_weights(dtype, k, n, e):
     w2 = w2[..., :n, :k]
 
     num_warps = 8
-    w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
-    w_scale_layout, w_scale_layout_opts = (
-        layout.make_default_matmul_mxfp4_w_scale_layout(mx_axis=1, num_warps=num_warps)
-    )
+    (
+        w_layout,
+        w_layout_opts,
+        w_scale_layout,
+        w_scale_layout_opts,
+    ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps)
 
     w1_tri = convert_layout(wrap_torch_tensor(w1_tri, FP4), w_layout, **w_layout_opts)
     w1_scale_tri = convert_layout(
@@ -152,10 +156,10 @@ def make_weights(dtype, k, n, e):
     )
 
     w1_precision_config = PrecisionConfig(
-        weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
+        **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
     )
     w2_precision_config = PrecisionConfig(
-        weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
+        **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData())
     )
 
     return (
diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py
index ea732d68876a..8d6deac4f025 100644
--- a/tests/kernels/moe/test_moe.py
+++ b/tests/kernels/moe/test_moe.py
@@ -119,8 +119,7 @@ def iterative_moe(
     expert_map: torch.Tensor = None,
     renormalize: bool = False,
 ) -> torch.Tensor:
-    """
-    Baseline implementation of fused moe.
+    """Baseline implementation of fused moe.
 
     Args:
         hidden_states: [*, hidden_size]
@@ -128,6 +127,7 @@ def iterative_moe(
         w2: [num_experts, hidden_size, intermediate_size]
         gating_output: [*, num_experts]
         expert_map: [num_experts]
+
     """
     orig_shape = hidden_states.shape
     hidden_size = hidden_states.shape[-1]
@@ -1245,7 +1245,7 @@ def _make_humming_indexed_experts(activation: MoEActivation):
     from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
         HummingIndexedExperts,
     )
-    from vllm.model_executor.layers.quantization.utils import humming_utils
+    from vllm.model_executor.layers.quantization.utils import humming as humming_utils
     from vllm.utils import humming
 
     top_k, num_experts = 6, 12
@@ -1306,6 +1306,56 @@ def _make_humming_indexed_experts(activation: MoEActivation):
     return experts, layer
 
 
+@pytest.mark.parametrize("valid_shape_m", [8192, 8193])
+@pytest.mark.parametrize("w2_block_size", [64, 96, 128])
+def test_humming_indexed_routing_matches_each_projection_tile(
+    valid_shape_m: int, w2_block_size: int
+):
+    """Independently tuned projections preserve routed rows and expert identity."""
+    from types import SimpleNamespace
+
+    from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
+        HummingIndexedExperts,
+    )
+
+    topk_ids = torch.tensor([[0, 1], [1, 2], [2, 0]], device="cuda")
+    expert_map = torch.tensor([1, -1, 0], dtype=torch.int32, device="cuda")
+    experts = SimpleNamespace(
+        global_num_experts=3,
+        estimate_local_valid_shape_m=lambda _: valid_shape_m,
+        w13_tuning_config=[
+            [0, 8192, {"block_shape": [64, 128, 64]}],
+            [8192, 16384, {"block_shape": [128, 128, 64]}],
+        ],
+        w2_tuning_config=[
+            [0, 8192, {"block_shape": [64, 128, 64]}],
+            [8192, 16384, {"block_shape": [w2_block_size, 128, 64]}],
+        ],
+        compute_config_str="compute",
+        w13_tuning_config_str="w13",
+        w2_tuning_config_str="w2",
+    )
+    first, second, _ = HummingIndexedExperts.prepare_humming_moe_kwargs(
+        experts, topk_ids, expert_map, None
+    )
+    block_sizes = (64, 64) if valid_shape_m == 8192 else (128, w2_block_size)
+    expected_rows = torch.where(expert_map[topk_ids.flatten()] >= 0)[0]
+    for kwargs, block_size in zip((first, second), block_sizes):
+        padded = kwargs["num_tokens_padded"].item()
+        assert padded % block_size == 0
+        rows = kwargs["sorted_ids"][:padded].long()
+        valid = rows < topk_ids.numel()
+        torch.testing.assert_close(rows[valid].sort().values, expected_rows)
+        routed_experts = kwargs["expert_ids"][: padded // block_size]
+        torch.testing.assert_close(
+            routed_experts.repeat_interleave(block_size)[valid],
+            expert_map[topk_ids.flatten()[rows[valid]]],
+        )
+    assert (first["top_k"], second["top_k"]) == (2, 1)
+    for key in ("sorted_ids", "expert_ids", "num_tokens_padded"):
+        assert (first[key] is second[key]) == (block_sizes[0] == block_sizes[1])
+
+
 @pytest.mark.parametrize("activation", list(MoEActivation))
 def test_humming_activation_metadata_tracks_shared_apply(activation: MoEActivation):
     from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
@@ -1384,24 +1434,33 @@ def test_humming_global_valid_shape_m(
     assert result == expected
 
 
-def test_humming_permute_scratch_is_keyed_by_runtime_topk(
+def test_humming_permute_scratch_is_shared_by_config(
     monkeypatch: pytest.MonkeyPatch,
 ):
     from types import SimpleNamespace
     from unittest.mock import Mock
 
     import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
+    import vllm.model_executor.layers.fused_moe.moe_permute_unpermute as permute
+    import vllm.v1.worker.workspace as workspace
+
+    manager = workspace.WorkspaceManager(torch.device("cpu"))
+    monkeypatch.setattr(workspace, "_manager", manager)
 
     scratch_topk6 = Mock()
     scratch_topk1 = Mock()
     scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1])
     monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True)
-    monkeypatch.setattr(humming, "MoEPermuteScratch", scratch_type)
+    monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type)
     moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6)
     moe_config.moe_parallel_config.dp_size = 2
-    experts = SimpleNamespace(_permute_scratch={}, moe_config=moe_config)
+    experts = SimpleNamespace(moe_config=moe_config)
+    other_layer = SimpleNamespace(moe_config=moe_config)
 
     assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
+    assert (
+        humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6
+    )
     assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1
     assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
 
@@ -1412,6 +1471,11 @@ def test_humming_permute_scratch_is_keyed_by_runtime_topk(
     assert second_call.kwargs["max_num_tokens"] == 6144
     assert second_call.kwargs["topk"] == 1
 
+    manager.lock()
+    assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
+    with pytest.raises(AssertionError, match="was not allocated during warmup"):
+        humming.HummingExpertsBase._get_permute_scratch(experts, 2)
+
 
 def test_humming_delegates_to_instance_activation():
     from types import SimpleNamespace
@@ -1450,54 +1514,46 @@ def test_humming_delegates_to_instance_activation():
 def test_humming_grouped_apply_forwards_valid_prefix(
     monkeypatch: pytest.MonkeyPatch,
 ):
-    from types import MethodType, SimpleNamespace
+    from types import SimpleNamespace
     from unittest.mock import Mock
 
     import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
     from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
-        HummingExpertsBase,
         HummingGroupedExperts,
     )
 
-    activation_func = Mock()
     buffers = {
         "gate_up_output": torch.empty(6, 4),
-        "activation_output": torch.empty(6, 2),
+        "quanted_down_input": torch.empty(6, 2),
         "down_output": torch.empty(6, 4),
     }
     experts = SimpleNamespace(
-        activation=activation_func,
         num_experts=2,
         estimate_local_valid_shape_m=lambda _: 6,
         prepare_buffers=lambda *_: buffers,
-        _get_permute_scratch=lambda _: None,
-        quantize_input=lambda _, *, inputs, input_scale=None, quanted_input: (
-            inputs,
-            input_scale,
+        _get_permute_scratch=lambda _, *, indices_only=False: object(),
+        process_input=Mock(
+            side_effect=lambda _, **kwargs: (
+                kwargs["inputs"],
+                kwargs.get("input_scale"),
+                None,
+            )
         ),
         humming_forward=Mock(),
         compute_config_str="",
         w13_tuning_config_str="",
         w2_tuning_config_str="",
     )
-    experts.apply_activation = MethodType(HummingExpertsBase.apply_activation, experts)
 
     hidden_states = torch.empty(3, 4)
     output = torch.empty_like(hidden_states)
     topk_ids = torch.zeros((3, 2), dtype=torch.int64)
     topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
-    expected_counts = torch.tensor([4], dtype=torch.int32)
     expert_offsets = torch.tensor([0, 2, 4], dtype=torch.int64)
     monkeypatch.setattr(
         humming,
-        "moe_permute",
-        lambda **_: (
-            torch.empty(6, 4),
-            None,
-            expert_offsets,
-            torch.arange(6),
-            None,
-        ),
+        "moe_prepare_scatter",
+        lambda *_: (expert_offsets, torch.arange(6).view(3, 2)),
     )
     monkeypatch.setattr(humming, "moe_unpermute", lambda **_: None)
 
@@ -1520,12 +1576,11 @@ def test_humming_grouped_apply_forwards_valid_prefix(
         apply_router_weight_on_input=False,
     )
 
-    activation_func.assert_called_once()
-    call_kwargs = activation_func.call_args.kwargs
+    call_kwargs = experts.process_input.call_args.kwargs
     assert call_kwargs["activation"] == MoEActivation.SITU
-    assert call_kwargs["input"].shape == (6, 4)
-    assert call_kwargs["output"].shape == (6, 2)
-    torch.testing.assert_close(call_kwargs["valid_token_counts"], expected_counts)
+    assert call_kwargs["inputs"].shape == (6, 4)
+    assert call_kwargs["quanted_input"].shape == (6, 2)
+    torch.testing.assert_close(call_kwargs["num_valid_tokens"], expert_offsets[-1:])
 
 
 def test_batched_marlin_activation_uses_expert_token_counts(
@@ -1639,7 +1694,8 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
         activation=moe_config.activation,
     )
     assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size
-    assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size
+    assert "activation_output" not in buffer_metas
+    assert buffer_metas["quanted_down_input"]["shape"][-1] == intermediate_size
     assert experts.moe_problem_size(
         a1=torch.empty(1, hidden_size),
         w1=torch.empty(num_experts, 1),
@@ -1647,6 +1703,13 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
         topk_ids=torch.empty(1, top_k, dtype=torch.long),
     ) == (num_experts, 1, intermediate_size, hidden_size, top_k)
 
+    # Quantized batched input storage must cover every dispatcher's token slots.
+    experts.is_batched = lambda: True
+    experts.max_num_tokens = 17
+    experts.num_dispatchers = 2
+    batched_metas, _ = experts.get_buffer_metas(17, top_k, activation)
+    assert batched_metas["quanted_gate_up_input"]["shape"][0] == num_experts * 17 * 2
+
 
 def test_humming_indexed_writes_supplied_output_buffer():
     from vllm.forward_context import set_forward_context
@@ -1956,8 +2019,7 @@ def __init__(
             )
 
         def is_valid(self):
-            """
-            Return True only if the input can be represented in a Batched
+            """Return True only if the input can be represented in a Batched
             format.
             """
             return torch.all(self.expert_num_tokens_cpu <= self.max_tokens_per_batch)
@@ -2078,9 +2140,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
     dtype: torch.dtype,
     workspace_init,
 ):
-    """
-    Test BF16 unquantized MoE with FlashInfer TRTLLM backend.
-    """
+    """Test BF16 unquantized MoE with FlashInfer TRTLLM backend."""
     set_random_seed(7)
 
     from vllm.model_executor.layers.fused_moe.config import (
diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py
index 7d07af4080c4..ca339eeed8bd 100644
--- a/tests/kernels/moe/test_moe_align_block_size.py
+++ b/tests/kernels/moe/test_moe_align_block_size.py
@@ -22,6 +22,46 @@
 set_random_seed(0)
 
 
+@pytest.mark.parametrize("num_experts,num_tokens", [(32, 8), (256, 8), (256, 1024)])
+@pytest.mark.parametrize("use_expert_map", [False, True])
+def test_moe_align_scatter_masks_invalid_routes(
+    num_experts, num_tokens, use_expert_map
+):
+    topk_ids = torch.arange(num_tokens * 4, device="cuda").reshape(-1, 4)
+    topk_ids %= num_experts
+    topk_ids[0] = torch.tensor([-1, num_experts, 0, 1], device="cuda")
+    expert_map = None
+    valid = (topk_ids >= 0) & (topk_ids < num_experts)
+    if use_expert_map:
+        expert_map = torch.full((num_experts,), -1, dtype=torch.int32, device="cuda")
+        expert_map[::2] = torch.arange(
+            num_experts // 2, device="cuda", dtype=torch.int32
+        )
+        valid &= topk_ids % 2 == 0
+
+    sorted_ids, expert_ids, num_padded, scatter_idx = moe_align_block_size(
+        topk_ids,
+        32,
+        num_experts,
+        expert_map,
+        ignore_invalid_experts=True,
+        return_scatter_idx=True,
+    )
+    expected = torch.arange(topk_ids.numel(), device="cuda", dtype=torch.int32)
+    expected[~valid.flatten()] = -1
+    torch.testing.assert_close(scatter_idx, expected[:, None])
+    active_sorted = sorted_ids[: num_padded.item()]
+    active_sorted = active_sorted[active_sorted < topk_ids.numel()]
+    torch.testing.assert_close(active_sorted.sort().values, expected[expected >= 0])
+
+    baseline = moe_align_block_size(
+        topk_ids, 32, num_experts, expert_map, ignore_invalid_experts=True
+    )
+    assert len(baseline) == 3
+    torch.testing.assert_close(num_padded, baseline[2])
+    torch.testing.assert_close(expert_ids, baseline[1])
+
+
 def _group_tokens_by_expert(
     sorted_ids: torch.Tensor,
     expert_ids: torch.Tensor,
@@ -54,8 +94,7 @@ def _verify_expert_level_sorting(
     valid_length: int,
     total_tokens: int,
 ):
-    """
-    Verify that actual_sorted_ids follows the correct expert-level sorting.
+    """Verify that actual_sorted_ids follows the correct expert-level sorting.
     The kerne limplementation may or may not preserve original token order
     in topk_ids in the final sorted_ids however this does not impact quality.
     """
@@ -96,8 +135,7 @@ def torch_moe_align_block_size(
     expert_map: torch.Tensor | None = None,
     pad_sorted_ids: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Golden torch implementation of moe_align_block_size.
+    """Golden torch implementation of moe_align_block_size.
 
     This function aligns the token distribution across experts to be compatible
     with block size for matrix multiplication by sorting tokens by expert and
@@ -189,7 +227,7 @@ def torch_moe_align_block_size(
 def test_moe_align_block_size(
     m: int, topk: int, num_experts: int, block_size: int, pad_sorted_ids: bool
 ):
-    """Test moe_align_block_size without expert mapping"""
+    """Test moe_align_block_size without expert mapping."""
     topk_ids = torch.zeros((m, topk), device="cuda", dtype=torch.int32)
     for i in range(m):
         experts = torch.randperm(num_experts, device="cuda")[:topk]
@@ -254,7 +292,7 @@ def test_moe_align_block_size_with_expert_map(
     block_size: int,
     mask_inactive_experts: bool,
 ):
-    """Test moe_align_block_size with expert mapping (EP scenario)"""
+    """Test moe_align_block_size with expert mapping (EP scenario)."""
     expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
     local_experts = list(range(0, num_experts, 2))
     for i, expert_id in enumerate(local_experts):
@@ -329,13 +367,15 @@ def test_moe_align_block_size_deterministic():
 
 @pytest.mark.parametrize("max_tokens_per_batch", [13, 16, 512])
 @pytest.mark.parametrize("num_experts", [8, 16, 32, 64])
-@pytest.mark.parametrize("block_size", [8, 16, 32, 64])
-@pytest.mark.parametrize("simulate_empty_batches", [False, True])
+@pytest.mark.parametrize("block_size", [8, 16, 32, 48, 64])
+@pytest.mark.parametrize(
+    "token_distribution", ["random", "alternating_empty", "skewed"]
+)
 def test_batched_moe_align_block_size(
     max_tokens_per_batch: int,
     num_experts: int,
     block_size: int,
-    simulate_empty_batches: bool,
+    token_distribution: str,
 ):
     def ref_outputs(
         expert_num_tokens: torch.Tensor,
@@ -384,15 +424,17 @@ def ref_outputs(
     # Compute expert_num_tokens
     expert_num_tokens = torch.randint(
         low=0,
-        high=max_tokens_per_batch,
+        high=max_tokens_per_batch + 1,
         size=(num_experts,),
         device="cpu",
         dtype=torch.int32,
     )
-    if simulate_empty_batches:
-        # mark half the batches to have 0 tokens
-        zero_batches = torch.randperm(num_experts)[: num_experts // 2]
-        expert_num_tokens[zero_batches] = 0
+    if token_distribution == "alternating_empty":
+        expert_num_tokens[::2] = 0
+    elif token_distribution == "skewed":
+        expert_num_tokens.zero_()
+        expert_num_tokens[0] = max_tokens_per_batch
+        expert_num_tokens[-1] = 1
 
     # ref outputs
     ref_sorted_ids, ref_expert_ids, ref_num_tokens_post_pad = ref_outputs(
diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py
index e1a152a23325..5fe5ed1ab2a6 100644
--- a/tests/kernels/moe/test_moe_layer.py
+++ b/tests/kernels/moe/test_moe_layer.py
@@ -240,8 +240,7 @@ def maybe_roundup_layer_hidden_size(
     act_dtype: torch.dtype,
     backend: str | None,
 ) -> int:
-    """
-    Given layer hidden size and MoE configurations, round up hidden_size
+    """Given layer hidden size and MoE configurations, round up hidden_size
     if necessary.
 
     Args:
@@ -253,6 +252,7 @@ def maybe_roundup_layer_hidden_size(
         Rounded up hidden_size if rounding up is required based on the configs
         and all2all backend.
         Original hidden size otherwise.
+
     """
     if backend == "deepep_high_throughput":
         from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
@@ -945,6 +945,7 @@ def create_shared_experts_from_config(
 
     Returns:
         TestMLP instance or None if config is None
+
     """
     if shared_experts_config is None:
         return None
@@ -991,6 +992,7 @@ def setup_moe_test_data(
 
     Returns:
         MoETestData containing all test data and transforms
+
     """
     # For latent MoE: latent_size = k // 2
     latent_size = k // 2
@@ -1482,6 +1484,7 @@ def _run_one_config(
     use_shared_experts: bool,
     use_gate: bool,
     use_routed_input_transform: bool,
+    input_scale: float = 1.0,
     **kwargs,
 ) -> None:
     """Generic test loop that sets up environment and delegates to test_body_fn.
@@ -1528,6 +1531,8 @@ def _run_one_config(
 
         # Extract data from test_data
         hidden_states = test_data.hidden_states
+        if input_scale != 1.0:
+            hidden_states = hidden_states * input_scale
         router_logits = test_data.router_logits
         w1 = test_data.w1
         w2 = test_data.w2
@@ -1710,7 +1715,6 @@ def test_moe_layer_no_parallel(
     monkeypatch,
 ):
     """Test MoE layer without parallelism (dp_size=1, tp_size=1, use_ep=False)."""
-
     if os.environ.get("VLLM_LOGGING_LEVEL") is None:
         monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR")
 
@@ -1981,7 +1985,140 @@ def _parallel_worker_rocm_deepep(
             os._exit(exit_code)
 
 
-# TODO: add cudagraphs/torch.compile tests
+def _test_body_mori_graph(
+    moe_layer: MoERunner,
+    hidden_states: torch.Tensor,
+    router_logits: torch.Tensor,
+    vllm_config: VllmConfig,
+    num_tokens: int,
+    num_tokens_across_dp: torch.Tensor,
+    baseline_output: torch.Tensor,
+    **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
+        AiterExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
+        MoriPrepareAndFinalize,
+    )
+
+    kernel = moe_layer._quant_method.moe_kernel
+    assert kernel is not None
+    assert isinstance(kernel.prepare_finalize, MoriPrepareAndFinalize)
+    assert isinstance(kernel.fused_experts, AiterExperts)
+    assert not kernel.prepare_finalize.use_fp8_dispatch
+
+    def check(actual: torch.Tensor, expected: torch.Tensor) -> None:
+        torch.testing.assert_close(actual, expected, atol=3.5e-2, rtol=3.5e-2)
+        # These inputs produce small outputs; an absolute-only budget can
+        # accept all-zero output. Also bound error relative to the reference.
+        ref_norm = expected.float().norm()
+        assert ref_norm > 0
+        assert (actual.float() - expected.float()).norm() < 3.5e-2 * ref_norm
+
+    # Forward-context metadata belongs on CPU, outside graph capture.
+    with set_forward_context(
+        None,
+        vllm_config,
+        num_tokens=num_tokens,
+        num_tokens_across_dp=num_tokens_across_dp.cpu(),
+    ):
+        initial_hidden = hidden_states.clone()
+        initial_router = router_logits.clone()
+        for _ in range(2):
+            eager_output = moe_layer(hidden_states, router_logits)
+        saved_eager_output = eager_output.clone()
+        torch.accelerator.synchronize()
+
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured_output = moe_layer(hidden_states, router_logits)
+        graph.replay()
+        original_output = captured_output.clone()
+        check(original_output, baseline_output)
+
+        for scale, shift in [(-0.5, 1), (1.75, 3)]:
+            hidden_states.copy_(initial_hidden * scale)
+            router_logits.copy_(initial_router.roll(shift, dims=-1))
+            expected = moe_layer(hidden_states, router_logits).clone()
+            graph.replay()
+            check(captured_output, expected)
+            assert (expected.float() - original_output.float()).norm() > (
+                0.1 * original_output.float().norm()
+            )
+
+        # Later calls must not overwrite an earlier eager return tensor.
+        check(eager_output, saved_eager_output)
+        check(original_output, baseline_output)
+        return baseline_output, original_output
+
+
+def _mori_graph_worker(
+    pgi: ProcessGroupInfo,
+    vllm_config: VllmConfig,
+    cpu_group,
+    m: int,
+    n: int,
+) -> None:
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    rocm_aiter_ops.refresh_env_variables()
+    _run_one_config(
+        vllm_config=vllm_config,
+        ep_size=2,
+        dp_size=2,
+        tp_size=1,
+        dp_rank=pgi.rank,
+        tp_rank=0,
+        is_sequence_parallel=False,
+        m=m,
+        n=n,
+        k=512,
+        num_experts=8,
+        top_k=2,
+        quantization=None,
+        backend="mori_high_throughput",
+        test_body_fn=_test_body_mori_graph,
+        use_shared_experts=False,
+        use_gate=False,
+        use_routed_input_transform=False,
+        input_scale=1.0 + pgi.rank,
+    )
+
+
+@pytest.mark.parametrize("m,n", [(1, 256), (32, 512), (45, 256)])
+def test_moe_layer_mori_graph(m: int, n: int, monkeypatch):
+    """BF16 EP dispatch/combine preserves routing and outputs across replays."""
+    from vllm._aiter_ops import is_aiter_found_and_supported
+
+    if not current_platform.is_rocm() or not has_mori():
+        pytest.skip("Requires ROCm and MoRI")
+    if not is_aiter_found_and_supported():
+        pytest.skip("Requires supported AITER kernels")
+    if current_platform.device_count() < 2:
+        pytest.skip("Requires two GPUs")
+
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0")
+    compilation_config = CompilationConfig()
+    compilation_config.pass_config.fuse_allreduce_rms = False
+    vllm_config = VllmConfig(
+        parallel_config=ParallelConfig(
+            data_parallel_size=2,
+            tensor_parallel_size=1,
+            enable_expert_parallel=True,
+            all2all_backend="mori_high_throughput",
+        ),
+        compilation_config=compilation_config,
+        scheduler_config=SchedulerConfig.default_factory(
+            max_num_batched_tokens=64, max_num_seqs=64
+        ),
+    )
+    parallel_launch_with_config(2, _mori_graph_worker, vllm_config, None, m, n)
+
+
+# TODO: add torch.compile tests and graph coverage for other backends.
 @pytest.mark.parametrize("dp_size, tp_size, use_ep", PARALLEL_COMBOS)
 @pytest.mark.parametrize("backend", BACKENDS)
 @pytest.mark.parametrize("enable_eplb", [False, True])
diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py
index 4089bbfe843a..76f4b57cf07f 100644
--- a/tests/kernels/moe/test_moe_permute_unpermute.py
+++ b/tests/kernels/moe/test_moe_permute_unpermute.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the MOE permute/unpermute kernel
+"""Tests for the MOE permute/unpermute kernel.
 
 Run `pytest tests/kernels/test_moe_permute_unpermute.py`.
 """
@@ -15,12 +15,15 @@
 )
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
     MoEPermuteScratch,
+    get_moe_permute_scratch,
     moe_permute,
     moe_permute_unpermute_supported,
+    moe_prepare_scatter,
     moe_unpermute,
 )
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.worker.workspace import current_workspace_manager
 
 NUM_EXPERTS = [16, 64, 256]
 TOP_KS = [2, 6, 8]
@@ -120,6 +123,7 @@ def torch_unpermute(
 @pytest.mark.parametrize("topk", TOP_KS)
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("ep_size", EP_SIZE)
+@pytest.mark.parametrize("use_scratch", [False, True])
 def test_moe_permute_unpermute(
     n_token: int,
     n_hidden: int,
@@ -127,6 +131,8 @@ def test_moe_permute_unpermute(
     n_expert: int,
     ep_size: int,
     dtype: torch.dtype,
+    use_scratch: bool,
+    workspace_init,
 ):
     if not moe_permute_unpermute_supported():
         pytest.skip("moe_permute_unpermute is not supported on this platform.")
@@ -160,6 +166,18 @@ def test_moe_permute_unpermute(
         expert_map=expert_map,
     )
 
+    scratch = None
+    if use_scratch:
+        scratch = get_moe_permute_scratch(
+            max_num_tokens=n_token,
+            topk=topk,
+            num_experts=n_expert,
+            num_local_experts=n_local_expert,
+            device=hidden_states.device,
+            hidden_size=n_hidden,
+            hidden_dtype=dtype,
+        )
+
     (
         permuted_hidden_states,
         _,
@@ -173,6 +191,7 @@ def test_moe_permute_unpermute(
         n_expert=n_expert,
         n_local_expert=n_local_expert,
         expert_map=expert_map,
+        scratch=scratch,
     )
 
     # check expert_first_token_offset
@@ -212,20 +231,22 @@ def test_moe_permute_unpermute(
     torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0)
 
 
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("n_token", [1, 33, 128])
+@pytest.mark.parametrize("topk", [1, 6, 8])
+def test_moe_permute_reuses_scratch_buffers(
+    dtype: torch.dtype, n_token: int, topk: int, workspace_init
+):
     if not moe_permute_unpermute_supported():
         pytest.skip("moe_permute_unpermute is not supported on this platform.")
 
-    n_token = 64
     n_hidden = 2048
     n_expert = 16
-    topk = 4
     hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype)
     gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype)
     _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False)
 
-    scratch = MoEPermuteScratch(
+    scratch_config = dict(
         max_num_tokens=n_token,
         topk=topk,
         num_experts=n_expert,
@@ -234,6 +255,7 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
         hidden_size=n_hidden,
         hidden_dtype=hidden_states.dtype,
     )
+    scratch = get_moe_permute_scratch(**scratch_config)
 
     first = moe_permute(
         hidden_states=hidden_states,
@@ -242,13 +264,26 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
         n_expert=n_expert,
         scratch=scratch,
     )
-    second = moe_permute(
-        hidden_states=hidden_states,
-        a1q_scale=None,
-        topk_ids=topk_ids,
-        n_expert=n_expert,
-        scratch=scratch,
-    )
+    current_workspace_manager().lock()
+    assert get_moe_permute_scratch(**scratch_config) is scratch
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        second = moe_permute(
+            hidden_states=hidden_states,
+            a1q_scale=None,
+            topk_ids=topk_ids,
+            n_expert=n_expert,
+            scratch=get_moe_permute_scratch(**scratch_config),
+        )
+
+    for _ in range(2):
+        hidden_states.add_(1)
+        topk_ids.copy_(topk_ids.roll(1, 0))
+        graph.replay()
+        expected = moe_permute(hidden_states, None, topk_ids, n_expert)
+        for actual, reference in zip(second, expected):
+            if actual is not None:
+                torch.testing.assert_close(actual, reference)
 
     (
         permuted_hidden_states_1,
@@ -288,6 +323,123 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
     )
 
 
+def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None:
+    """Switching graph sizes must not expose metadata left by a larger batch."""
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    device = torch.device("cuda:0")
+    scratch = get_moe_permute_scratch(
+        max_num_tokens=128,
+        topk=6,
+        num_experts=16,
+        num_local_experts=16,
+        device=device,
+        hidden_size=2048,
+        hidden_dtype=torch.bfloat16,
+    )
+    current_workspace_manager().lock()
+    runs = []
+    for n_token in (128, 1, 33):
+        hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device)
+        topk_ids = torch.randint(16, (n_token, 6), device=device)
+        moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+        runs.append((graph, hidden, topk_ids, output))
+
+    for graph, hidden, topk_ids, output in runs + runs[::-1]:
+        hidden.add_(1)
+        topk_ids.copy_((topk_ids + 1) % 16)
+        graph.replay()
+        expected = moe_permute(hidden, None, topk_ids, 16)
+        for actual, reference in zip(output, expected):
+            if actual is not None:
+                torch.testing.assert_close(actual, reference)
+
+
+def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> None:
+    """Concurrent graph replays in different ubatches/lanes cannot share scratch."""
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    from vllm.v1.worker import workspace
+
+    device = torch.device("cuda:0")
+    manager = workspace.WorkspaceManager(device, num_ubatches=2, num_lanes=2)
+    monkeypatch.setattr(workspace, "_manager", manager)
+    ubatch = 0
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch)
+    scratch_config = dict(
+        max_num_tokens=64,
+        topk=6,
+        num_experts=16,
+        num_local_experts=16,
+        device=device,
+        hidden_size=2048,
+        hidden_dtype=torch.bfloat16,
+    )
+    runs = []
+    for i, n_token in enumerate((1, 7, 33, 64)):
+        ubatch = i // 2
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream):
+            hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device)
+            topk_ids = torch.randint(16, (n_token, 6), device=device)
+            scratch = get_moe_permute_scratch(**scratch_config)
+            moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+            runs.append((stream, graph, hidden, topk_ids, output, scratch))
+        torch.cuda.current_stream().wait_stream(stream)
+
+    assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4
+    manager.lock()
+    for _ in range(3):
+        for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs):
+            ubatch = i // 2
+            with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream):
+                assert get_moe_permute_scratch(**scratch_config) is scratch
+                hidden.add_(1)
+                topk_ids.copy_((topk_ids + 1) % 16)
+                graph.replay()
+        for stream, _, hidden, topk_ids, output, _ in runs:
+            torch.cuda.current_stream().wait_stream(stream)
+            expected = moe_permute(hidden, None, topk_ids, 16)
+            for actual, reference in zip(output, expected):
+                if actual is not None:
+                    torch.testing.assert_close(actual, reference)
+
+
+def test_moe_permute_scratch_without_manager(monkeypatch) -> None:
+    """Standalone calls get independent scratch with initialized row indices."""
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    from vllm.v1.worker import workspace
+
+    monkeypatch.setattr(workspace, "_manager", None)
+    config = dict(
+        max_num_tokens=4,
+        topk=2,
+        num_experts=4,
+        num_local_experts=4,
+        device=torch.device("cuda"),
+    )
+    first = get_moe_permute_scratch(**config)
+    second = get_moe_permute_scratch(**config)
+
+    assert (
+        first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr()
+    )
+    torch.testing.assert_close(
+        first.token_expert_indices, torch.arange(8, dtype=torch.int32, device="cuda")
+    )
+
+
 def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
     if not moe_permute_unpermute_supported():
         pytest.skip("moe_permute_unpermute is not supported on this platform.")
@@ -307,7 +459,7 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
         hidden_dtype=hidden_states.dtype,
     )
 
-    permuted, _, expert_offsets, _, _ = moe_permute(
+    permuted, _, expert_offsets, inverse, _ = moe_permute(
         hidden_states=hidden_states,
         a1q_scale=None,
         topk_ids=topk_ids,
@@ -322,3 +474,27 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
         torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"),
     )
     torch.testing.assert_close(permuted[:2], hidden_states[[0, 2]])
+    assert torch.all(inverse[[1, 3, 4]] >= expert_offsets[-1])
+    expected = torch.zeros_like(hidden_states)
+    expected[[0, 2]] = hidden_states[[0, 2]]
+    output = torch.empty_like(hidden_states)
+    moe_unpermute(
+        output, permuted, torch.ones(5, 1, device="cuda"), inverse, expert_offsets
+    )
+    torch.testing.assert_close(output, expected)
+
+    expected_inverse = inverse.clone()
+    expected_offsets = expert_offsets.clone()
+    expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch)
+    torch.testing.assert_close(indices.flatten(), expected_inverse)
+    torch.testing.assert_close(expert_offsets, expected_offsets)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch)
+    topk_ids.zero_()
+    graph.replay()
+
+    expected_inverse = torch.arange(5, dtype=torch.int32, device="cuda")
+    expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda")
+    torch.testing.assert_close(indices.flatten(), expected_inverse)
+    torch.testing.assert_close(expert_offsets, expected_offsets)
diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py
index d3716a196b02..ef31de06d14a 100644
--- a/tests/kernels/moe/test_moe_weight_loading_padded.py
+++ b/tests/kernels/moe/test_moe_weight_loading_padded.py
@@ -12,7 +12,13 @@
 import pytest
 import torch
 
+from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
+    UnquantizedMoeBackend,
+)
 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+    UnquantizedFusedMoEMethod,
+)
 
 from .utils import make_dummy_moe_config
 
@@ -249,6 +255,57 @@ def test_narrowed_data_shares_storage(self):
 class TestWeightLoadingWithPaddedHiddenSize:
     """Integration-style tests that simulate padded weight loading."""
 
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+    @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+    @pytest.mark.parametrize("tp_rank", [0, 1])
+    def test_load_w2_chunks_noncontiguous_cpu_source_to_padded_cuda(
+        self, dtype, tp_rank
+    ):
+        hidden = 4096
+        intermediate = 1024
+        loaded_weight = (
+            torch.arange(hidden * intermediate * 2)
+            .remainder(251)
+            .to(dtype)
+            .reshape(hidden, intermediate * 2)
+        )
+        tp_source = loaded_weight[
+            :, intermediate * tp_rank : intermediate * (tp_rank + 1)
+        ]
+        expert_data_full = torch.zeros(
+            hidden + 8, intermediate + 8, device="cuda", dtype=dtype
+        )
+        destination = expert_data_full[:hidden, :intermediate]
+
+        assert not tp_source.is_contiguous()
+        assert tp_source.nbytes > 1 << 20
+        assert not destination.is_contiguous()
+
+        experts = object.__new__(RoutedExperts)
+        torch.nn.Module.__init__(experts)
+        experts.moe_config = make_dummy_moe_config()
+        experts.moe_config.moe_parallel_config.tp_size = 2
+
+        torch.accelerator.synchronize()
+        allocated_before = torch.accelerator.memory_allocated()
+        torch.accelerator.reset_peak_memory_stats()
+        experts._load_w2(
+            expert_data=expert_data_full,
+            shard_dim=1,
+            loaded_weight=loaded_weight,
+            tp_rank=tp_rank,
+        )
+        torch.accelerator.synchronize()
+        peak_extra = torch.accelerator.max_memory_allocated() - allocated_before
+
+        # A single copy into the strided CUDA view allocates a full-shard
+        # temporary. Correct values alone would not catch that regression.
+        assert peak_extra < tp_source.nbytes
+
+        torch.testing.assert_close(destination.cpu(), tp_source)
+        assert torch.count_nonzero(expert_data_full[hidden:, :]) == 0
+        assert torch.count_nonzero(expert_data_full[:hidden, intermediate:]) == 0
+
     def test_load_w2_with_padding(self):
         """Simulate loading w2 weights when hidden_size is padded."""
         padded_hidden = 3072
@@ -364,6 +421,50 @@ def test_narrow_shard_dim(self):
         )
 
 
+class TestUnquantizedTrtLlmPrePadding:
+    @staticmethod
+    def _make_method(
+        intermediate: int,
+        backend: UnquantizedMoeBackend = UnquantizedMoeBackend.FLASHINFER_TRTLLM,
+    ) -> UnquantizedFusedMoEMethod:
+        moe_config = make_dummy_moe_config(
+            num_experts=2,
+            hidden_dim=64,
+            intermediate_size=intermediate,
+        )
+        method = object.__new__(UnquantizedFusedMoEMethod)
+        method.moe = moe_config
+        method.unquantized_backend = backend
+        method.moe_kernel = None
+        return method
+
+    @pytest.mark.parametrize(
+        "backend,original,expected",
+        [
+            (UnquantizedMoeBackend.FLASHINFER_TRTLLM, 1344, 1408),
+            (UnquantizedMoeBackend.FLASHINFER_TRTLLM, 1408, 1408),
+            (UnquantizedMoeBackend.TRITON, 1344, 1344),
+        ],
+    )
+    def test_rounds_intermediate_before_weight_allocation(
+        self,
+        backend: UnquantizedMoeBackend,
+        original: int,
+        expected: int,
+    ):
+        method = self._make_method(original, backend)
+
+        hidden, intermediate = method.maybe_roundup_sizes(
+            hidden_size=64,
+            intermediate_size_per_partition=original,
+            act_dtype=torch.bfloat16,
+            moe_parallel_config=method.moe.moe_parallel_config,
+        )
+
+        assert hidden == 64
+        assert intermediate == expected
+
+
 class TestLoadWeightsExpertBias:
     """Some quantized exports (e.g. GPTQ, llm-compressor NVFP4) materialize
     all-zero per-expert `.bias` tensors for models whose experts have no bias
diff --git a/tests/kernels/moe/test_moonep_bf16_poc.py b/tests/kernels/moe/test_moonep_bf16_poc.py
new file mode 100644
index 000000000000..ba1e78fb7bbf
--- /dev/null
+++ b/tests/kernels/moe/test_moonep_bf16_poc.py
@@ -0,0 +1,412 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test MoonEP dispatch / prefetch / combine logic (BF16 PoC).
+
+Two paths, both compared against the pure-PyTorch reference MoE:
+- MoonEPPrepareAndFinalize + a reference segment-loop expert runner over
+  MoonEP's expert-grouped ``[NvS, H]`` layout;
+- the full modular kernel, MoonEPPrepareAndFinalize + MoonEPExperts through
+  FusedMoEKernel.
+Requires NVSwitch multicast capable GPUs.
+"""
+
+import dataclasses
+import types
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch.distributed import ProcessGroup
+
+from tests.kernels.moe.utils import make_test_weights
+from tests.kernels.utils import torch_experts
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceNoOP,
+)
+from vllm.utils.import_utils import has_moonep
+from vllm.utils.torch_utils import set_random_seed
+
+from ...utils import multi_gpu_test
+from .parallel_utils import ProcessGroupInfo, parallel_launch
+
+if has_moonep():
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
+        MoonEPExpertWeightLayout,
+        MoonEPPrepareAndFinalize,
+        make_moonep_weight_layout,
+    )
+
+requires_moonep = pytest.mark.skipif(
+    not has_moonep(),
+    reason="Requires MoonEP",
+)
+
+
+class MulticastNotAvailableError(RuntimeError):
+    pass
+
+
+@dataclasses.dataclass
+class TestConfig:
+    topk: int
+    m: int
+    k: int
+    n: int
+    num_experts: int
+    router_skew: float
+    apply_router_weight_on_input: bool = False
+
+
+@dataclasses.dataclass
+class TestTensors:
+    rank_tokens: torch.Tensor
+    topk: torch.Tensor
+    topk_weights: torch.Tensor
+    config: TestConfig
+
+    @staticmethod
+    def make(config: TestConfig) -> "TestTensors":
+        rank_tokens = (
+            torch.randn((config.m, config.k), device="cuda", dtype=torch.bfloat16) / 10
+        )
+        # Skewed routing via a per-expert bias: scaling logits by a scalar
+        # preserves their ordering and would not change the selected topk
+        # ids, but a shared bias makes a few experts hot on every token so
+        # the planner has to fill redundant-expert prefetch slots.
+        expert_bias = config.router_skew * torch.randn(
+            1, config.num_experts, device="cuda", dtype=torch.float32
+        )
+        logits = (
+            torch.randn(
+                config.m, config.num_experts, device="cuda", dtype=torch.float32
+            )
+            + expert_bias
+        )
+        topk_weights, topk = torch.topk(logits, config.topk, dim=-1)
+        topk_weights = torch.softmax(topk_weights, dim=-1)
+        return TestTensors(
+            rank_tokens=rank_tokens,
+            topk=topk.to(dtype=torch.int64),
+            topk_weights=topk_weights,
+            config=config,
+        )
+
+
+def reference_moonep_experts(
+    hidden_nvsh: torch.Tensor,
+    route_weights_nvs: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    weight_layout: "MoonEPExpertWeightLayout",
+) -> torch.Tensor:
+    """Segment loop over MoonEP's expert-grouped layout.
+
+    ``prefetch_weight`` has already materialized redundant experts' weights
+    in rows ``[E, E+B)``, so every segment reads its own row. Route weights
+    are applied here; MoonEP's combine does the K-sum.
+    """
+    output = torch.empty_like(hidden_nvsh)
+    prev = 0
+    for row, cur in enumerate(cu_seqlens.tolist()):
+        if cur == prev:
+            continue
+        x = hidden_nvsh[prev:cur]
+        gate = F.linear(x, weight_layout.full_gate_weight[row])
+        up = F.linear(x, weight_layout.full_up_weight[row])
+        y = F.linear(F.silu(gate) * up, weight_layout.full_down_weight[row])
+        y = y * route_weights_nvs[prev:cur].to(dtype=y.dtype).unsqueeze(-1)
+        output[prev:cur].copy_(y)
+        prev = cur
+    if prev < hidden_nvsh.shape[0]:
+        output[prev:].zero_()
+    return output
+
+
+def make_moonep_prepare_finalize(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    hidden_size: int,
+    num_experts: int,
+    topk: int,
+    max_tokens_per_rank: int,
+    weight_layout: "MoonEPExpertWeightLayout",
+    pass_layout_to_pf: bool = True,
+):
+    from moonep._C import nvl_multicast_supported
+
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
+        MoonEPBufferPool,
+    )
+
+    if not nvl_multicast_supported():
+        raise MulticastNotAvailableError("NVSwitch multicast not available")
+
+    pool = MoonEPBufferPool(
+        dict(
+            H=hidden_size,
+            K=topk,
+            E=num_experts,
+            num_ep_ranks=pgi.world_size,
+            B=weight_layout.num_prefetch_slots,
+            group=pg,
+            explicitly_destroy=True,
+        ),
+        max_tokens_per_rank=max_tokens_per_rank,
+    )
+    return pool, MoonEPPrepareAndFinalize(
+        buffer_pool=pool,
+        max_tokens_per_rank=max_tokens_per_rank,
+        num_dispatchers=pgi.world_size,
+        num_global_experts=num_experts,
+        weight_layout=weight_layout if pass_layout_to_pf else None,
+    )
+
+
+def moonep_moe_impl(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    test_tensors: TestTensors,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    num_prefetch_slots: int,
+) -> torch.Tensor:
+    config = test_tensors.config
+    hidden_size = test_tensors.rank_tokens.shape[1]
+    max_tokens_per_rank = 128 * ((config.m + 127) // 128)
+
+    weight_layout = make_moonep_weight_layout(w1, w2, num_prefetch_slots)
+    pool, pf = make_moonep_prepare_finalize(
+        pg,
+        pgi,
+        hidden_size,
+        config.num_experts,
+        config.topk,
+        max_tokens_per_rank,
+        weight_layout,
+    )
+    try:
+        hidden_nvsh, _, _, _, route_weights_nvs = pf.prepare(
+            test_tensors.rank_tokens,
+            test_tensors.topk_weights,
+            test_tensors.topk,
+            num_experts=config.num_experts,
+            expert_map=None,
+            apply_router_weight_on_input=config.apply_router_weight_on_input,
+            quant_config=_no_quant_config(),
+        )
+        assert route_weights_nvs is not None
+        if config.apply_router_weight_on_input:
+            # prepare() already applied the weights to the inputs; the
+            # runner must not apply them again.
+            route_weights_nvs = torch.ones_like(route_weights_nvs)
+        if config.router_skew >= 8 and config.m >= 100:
+            # Heavy skew must actually engage the redundant-expert planner.
+            assert (pf.plan.experts_to_copy >= 0).any(), (
+                "no prefetch slot used despite heavy router skew"
+            )
+        expert_out = reference_moonep_experts(
+            hidden_nvsh, route_weights_nvs, pf.cu_seqlens, weight_layout
+        )
+        output = torch.empty_like(test_tensors.rank_tokens)
+        pf.finalize(
+            output,
+            expert_out,
+            test_tensors.topk_weights,
+            test_tensors.topk,
+            apply_router_weight_on_input=False,
+            weight_and_reduce_impl=TopKWeightAndReduceNoOP(),
+        )
+        torch.accelerator.synchronize()
+        return output
+    finally:
+        pool.destroy()
+
+
+def _no_quant_config():
+    from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+
+    return FusedMoEQuantConfig.make(quant_dtype=None)
+
+
+def moonep_modular_kernel_impl(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    test_tensors: TestTensors,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    num_prefetch_slots: int,
+) -> torch.Tensor:
+    """Full modular-kernel path: MoonEPPrepareAndFinalize + MoonEPExperts."""
+    from tests.kernels.moe.utils import make_dummy_moe_config
+    from vllm.model_executor.layers.fused_moe.experts.moonep_experts import (
+        MoonEPExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+
+    config = test_tensors.config
+    hidden_size = test_tensors.rank_tokens.shape[1]
+    max_tokens_per_rank = 128 * ((config.m + 127) // 128)
+
+    weight_layout = make_moonep_weight_layout(w1, w2, num_prefetch_slots)
+    # weight_layout deliberately not passed to the P/F: the engine path
+    # resolves it through post_init_setup + the experts' weight hook.
+    pool, pf = make_moonep_prepare_finalize(
+        pg,
+        pgi,
+        hidden_size,
+        config.num_experts,
+        config.topk,
+        max_tokens_per_rank,
+        weight_layout,
+        pass_layout_to_pf=False,
+    )
+    try:
+        moe_config = make_dummy_moe_config(
+            num_experts=config.num_experts,
+            experts_per_token=config.topk,
+            hidden_dim=hidden_size,
+            intermediate_size=config.n,
+            max_num_tokens=max_tokens_per_rank,
+        )
+        experts = MoonEPExperts(moe_config=moe_config, quant_config=_no_quant_config())
+        # Stand-in for the layer carrying the layout, as
+        # convert_to_unquantized_kernel_format leaves it in the engine path.
+        fake_layer = types.SimpleNamespace(_moonep_weight_layout=weight_layout)
+        experts.process_weights_after_loading(fake_layer)
+        kernel = FusedMoEKernel(prepare_finalize=pf, fused_experts=experts)
+        out = kernel.apply(
+            hidden_states=test_tensors.rank_tokens,
+            w1=weight_layout.full_gate_weight,
+            w2=weight_layout.full_down_weight,
+            topk_weights=test_tensors.topk_weights,
+            topk_ids=test_tensors.topk,
+            activation=MoEActivation.SILU,
+            global_num_experts=config.num_experts,
+            expert_map=None,
+            apply_router_weight_on_input=config.apply_router_weight_on_input,
+        )
+        torch.accelerator.synchronize()
+        return out
+    finally:
+        pool.destroy()
+
+
+def _moonep_moe(
+    pgi: ProcessGroupInfo,
+    config: TestConfig,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    num_prefetch_slots: int,
+    use_modular_kernel: bool,
+):
+    from vllm.v1.worker.workspace import init_workspace_manager
+
+    device_idx = torch.accelerator.current_device_index()
+    init_workspace_manager(torch.device("cuda", device_idx))
+    w1 = w1.to(device=device_idx)
+    w2 = w2.to(device=device_idx)
+
+    pg = torch.distributed.new_group(list(range(pgi.world_size)))
+    set_random_seed(7 + pgi.rank)
+    test_tensors = TestTensors.make(config)
+
+    with set_current_vllm_config(VllmConfig()):
+        torch_combined = torch_experts(
+            test_tensors.rank_tokens,
+            w1,
+            w2,
+            test_tensors.topk_weights,
+            test_tensors.topk,
+            apply_router_weights_on_input=config.apply_router_weight_on_input,
+        )
+        impl = moonep_modular_kernel_impl if use_modular_kernel else moonep_moe_impl
+        moonep_combined = impl(pg, pgi, test_tensors, w1, w2, num_prefetch_slots)
+
+    torch.testing.assert_close(
+        torch_combined,
+        moonep_combined,
+        atol=6e-2,
+        rtol=6e-2,
+    )
+
+
+MNKs = [
+    (1, 256, 512),
+    (37, 256, 512),
+    (100, 512, 1024),
+    (512, 768, 2048),
+]
+
+
+@pytest.mark.parametrize("m,n,k", MNKs)
+@pytest.mark.parametrize("num_experts", [32])
+@pytest.mark.parametrize("topk", [4])
+@pytest.mark.parametrize("router_skew", [1.0, 8.0])
+@pytest.mark.parametrize("num_prefetch_slots", [4])
+@pytest.mark.parametrize("world_size", [2])
+@pytest.mark.parametrize("use_modular_kernel", [False, True])
+@multi_gpu_test(num_gpus=2)
+@requires_moonep
+def test_moonep_bf16_moe(
+    m: int,
+    n: int,
+    k: int,
+    num_experts: int,
+    topk: int,
+    router_skew: float,
+    num_prefetch_slots: int,
+    world_size: int,
+    use_modular_kernel: bool,
+):
+    set_random_seed(7)
+    config = TestConfig(
+        topk=topk,
+        m=m,
+        k=k,
+        n=n,
+        num_experts=num_experts,
+        router_skew=router_skew,
+    )
+    (_, w1, _, _), (_, w2, _, _) = make_test_weights(num_experts, n, k)
+
+    try:
+        parallel_launch(
+            world_size,
+            _moonep_moe,
+            config,
+            w1,
+            w2,
+            num_prefetch_slots,
+            use_modular_kernel,
+        )
+    except Exception as exc:
+        if "MulticastNotAvailableError" in str(exc):
+            pytest.skip("NVSwitch multicast not available")
+        raise
+
+
+@pytest.mark.parametrize("use_modular_kernel", [False, True])
+@multi_gpu_test(num_gpus=2)
+@requires_moonep
+def test_moonep_input_weighted_moe(use_modular_kernel: bool):
+    """Llama-4-style routing applies route weights to the expert inputs."""
+    set_random_seed(7)
+    config = TestConfig(
+        topk=1,
+        m=100,
+        k=1024,
+        n=512,
+        num_experts=32,
+        router_skew=1.0,
+        apply_router_weight_on_input=True,
+    )
+    (_, w1, _, _), (_, w2, _, _) = make_test_weights(
+        config.num_experts, config.n, config.k
+    )
+    try:
+        parallel_launch(2, _moonep_moe, config, w1, w2, 4, use_modular_kernel)
+    except Exception as exc:
+        if "MulticastNotAvailableError" in str(exc):
+            pytest.skip("NVSwitch multicast not available")
+        raise
diff --git a/tests/kernels/moe/test_mori_moe.py b/tests/kernels/moe/test_mori_moe.py
new file mode 100644
index 000000000000..9869b6bc5467
--- /dev/null
+++ b/tests/kernels/moe/test_mori_moe.py
@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Two-rank MoRI/AITER numerical counterpart to the DeepEP v2 MoE matrix."""
+
+from dataclasses import replace
+
+import pytest
+import torch
+
+from tests.kernels.moe.modular_kernel_tools.parallel_utils import (
+    parallel_launch_with_config,
+)
+from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights
+from tests.kernels.utils import torch_experts
+from tests.utils import multi_gpu_test
+from vllm.config import ParallelConfig, SchedulerConfig, VllmConfig
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
+    maybe_make_prepare_finalize,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_mori
+from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.worker.workspace import init_workspace_manager
+
+MNKS = [
+    (1, 256, 256),
+    (2, 256, 512),
+    (3, 1024, 2048),
+    (32, 256, 1024),
+    (45, 512, 2048),
+    (64, 1024, 1024),
+    (222, 1024, 2048),
+]
+
+
+def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph):
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
+        MoriPrepareAndFinalize,
+    )
+
+    rocm_aiter_ops.refresh_env_variables()
+    init_workspace_manager(pgi.device)
+    set_random_seed(47)
+    dtype = current_platform.fp8_dtype() if fp8 else None
+    (_, w1, s1, _), (_, w2, s2, _) = make_test_weights(
+        32, n, k, quant_dtype=dtype, per_out_ch_quant=fp8
+    )
+    # Share global expert weights, but give each rank different token inputs.
+    set_random_seed(101 + pgi.rank)
+    x = torch.randn(m, k, device=pgi.device, dtype=torch.bfloat16) / 10
+    routing = torch.randn(m, 32, device=pgi.device)
+    # Guarantee a route to both ranks, including the one-token case.
+    routing[:, 0] = routing[:, 16] = float("inf")
+    ids = routing.topk(6, dim=-1).indices.to(torch.int32)
+    weights = torch.randn(m, 6, device=pgi.device, dtype=torch.float32)
+    parallel = replace(
+        FusedMoEParallelConfig.make_no_parallel(),
+        dp_size=2,
+        dp_rank=pgi.rank,
+        ep_size=2,
+        ep_rank=pgi.rank,
+        use_ep=True,
+        all2all_backend="mori_high_throughput",
+    )
+    config = replace(
+        make_dummy_moe_config(
+            num_experts=32,
+            num_local_experts=16,
+            experts_per_token=6,
+            hidden_dim=k,
+            intermediate_size=n,
+            max_num_tokens=256,
+        ),
+        moe_parallel_config=parallel,
+    )
+    start, end = pgi.rank * 16, (pgi.rank + 1) * 16
+    quant = FusedMoEQuantConfig.make(
+        dtype,
+        per_act_token_quant=fp8,
+        per_out_ch_quant=fp8,
+        w1_scale=None if s1 is None else s1[start:end].contiguous(),
+        w2_scale=None if s2 is None else s2[start:end].contiguous(),
+    )
+    prepare = maybe_make_prepare_finalize(config, quant)
+    assert isinstance(prepare, MoriPrepareAndFinalize)
+    assert prepare.use_fp8_dispatch == fp8
+    assert prepare.num_dispatchers() == 2
+    experts = AiterExperts(config, quant)
+    assert not experts.expects_unquantized_inputs
+    kernel = FusedMoEKernel(prepare, experts)
+    local_w1, local_w2 = rocm_aiter_ops.shuffle_weights(
+        w1[start:end].contiguous(), w2[start:end].contiguous()
+    )
+    local_w1.is_shuffled = local_w2.is_shuffled = True
+    expert_mask = torch.zeros(32, dtype=torch.int32, device=pgi.device)
+    expert_mask[start:end] = 1
+
+    def reference(route_weights):
+        return torch_experts(
+            x,
+            w1,
+            w2,
+            route_weights,
+            ids,
+            w1_scale=s1,
+            w2_scale=s2,
+            quant_dtype=dtype,
+            per_act_token_quant=fp8,
+        )
+
+    def run():
+        with set_forward_context(None, vllm_config):
+            return kernel.apply(
+                x,
+                local_w1,
+                local_w2,
+                weights,
+                ids,
+                activation=MoEActivation.SILU,
+                global_num_experts=32,
+                expert_map=expert_mask,
+                apply_router_weight_on_input=False,
+            )
+
+    def check(actual):
+        expected = reference(weights)
+        assert actual.shape == expected.shape
+        assert actual.dtype == expected.dtype
+        local_weights = weights * ((ids >= start) & (ids < end))
+        missing_remote = reference(local_weights)
+        norm = expected.float().norm()
+        assert norm > 0
+        norm_budget = 0.2 if fp8 else 0.06
+        assert (actual.float() - expected.float()).norm() / norm < norm_budget
+        # The norm check also rejects low-signal zero or rank-local results.
+        for incorrect in (torch.zeros_like(expected), missing_remote):
+            assert (incorrect.float() - expected.float()).norm() / norm > norm_budget
+        if fp8:
+            assert (
+                torch.isclose(actual, expected, atol=0.2, rtol=0.2)
+                .float()
+                .mean()
+                .item()
+                > 0.99
+            )
+        else:
+            torch.testing.assert_close(actual, expected, atol=0.06, rtol=0.06)
+        assert torch.isfinite(actual).all()
+
+    actual = run()
+    check(actual)
+    if graph:
+        saved = actual.clone()
+        run()
+        capture = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(capture):
+            captured = run()
+        for factor in (-0.5, 1.5):
+            x.mul_(factor)
+            weights.mul_(-0.75)
+            capture.replay()
+            torch.accelerator.synchronize()
+            check(captured)
+            torch.testing.assert_close(actual, saved, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm() or not has_mori(), reason="Requires ROCm MoRI"
+)
+@pytest.mark.parametrize("m,n,k", MNKS)
+@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8_token_channel"])
+@multi_gpu_test(num_gpus=2)
+def test_mori_moe(m, n, k, fp8):
+    config = VllmConfig(
+        parallel_config=ParallelConfig(
+            data_parallel_size=2,
+            enable_expert_parallel=True,
+            all2all_backend="mori_high_throughput",
+        ),
+        scheduler_config=SchedulerConfig.default_factory(max_num_batched_tokens=256),
+    )
+    env = {
+        "VLLM_ROCM_USE_AITER": "1",
+        "VLLM_ROCM_USE_AITER_MOE": "1",
+        "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS": "0",
+    }
+    parallel_launch_with_config(2, _mori_worker, config, env, m, n, k, fp8, m == 32)
diff --git a/tests/kernels/moe/test_mxfp4_moe.py b/tests/kernels/moe/test_mxfp4_moe.py
index 16b233b935e0..7c7ca1088574 100644
--- a/tests/kernels/moe/test_mxfp4_moe.py
+++ b/tests/kernels/moe/test_mxfp4_moe.py
@@ -67,8 +67,7 @@ def compute_ref_output(
 @pytest.mark.parametrize("num_experts", [8, 16, 32])
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16])
 def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype):
-    """
-    Test the MXFP4 grouped GEMM kernel by:
+    """Test the MXFP4 grouped GEMM kernel by:
     1. Creating random per-expert inputs and weights
     2. Quantizing both to MXFP4 using the CUDA kernel
     3. Running the CUTLASS grouped GEMM
@@ -205,9 +204,7 @@ def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype):
     reason="mxfp4_experts_quant requires CUDA SM100",
 )
 def test_mxfp4_experts_quant_basic():
-    """
-    Basic smoke test for the MXFP4 experts quantization kernel.
-    """
+    """Basic smoke test for the MXFP4 experts quantization kernel."""
     device = "cuda"
     num_experts = 4
     k = 256
@@ -295,8 +292,7 @@ def compute_reference_e8m0_scale(block_max: float) -> int:
 @pytest.mark.parametrize("k", [256, 7168])
 @pytest.mark.parametrize("m", [16, 64])
 def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k):
-    """
-    Test that mxfp4_experts_quant computes E8M0 block scales correctly
+    """Test that mxfp4_experts_quant computes E8M0 block scales correctly
     per OCP MX spec (not the NVFP4 formula).
 
     The old buggy kernel used: floor(log2(max/6)) + 127
@@ -416,8 +412,7 @@ def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k):
     reason="mxfp4_experts_quant requires CUDA SM100",
 )
 def test_mxfp4_experts_quant_no_saturation():
-    """
-    Test that the E8M0 scale is large enough to avoid unexpected saturation.
+    """Test that the E8M0 scale is large enough to avoid unexpected saturation.
 
     With the buggy NVFP4 formula, the scale was too small causing most values
     to saturate to ±6 in FP4. The fixed OCP MX formula should ensure that
diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py
index 830ff48734e1..61b26af62c09 100644
--- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py
+++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py
@@ -2,10 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """MXFP8 MoE backend selection for the AITER FlyDSL kernel (gfx950).
 
-GPU-free: mocks the platform (gfx950) and the ``flydsl`` package check, then
+GPU-free: mocks the platform (gfx950) and the aiter fused-MoE enable flag, then
 exercises the oracle so the FlyDSL backend is auto-picked when usable (including
 under expert parallelism, since apply() forwards the expert_map as aiter's
-expert_mask) and skipped (native fallback) when the device/package is missing.
+expert_mask) and skipped (native fallback) when the device is unsupported or the
+aiter runtime is disabled.
 """
 
 import dataclasses
@@ -48,6 +49,8 @@
     kMxfp8Static,
 )
 
+from vllm._aiter_ops import rocm_aiter_ops  # noqa: E402  # isort: skip
+
 _AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe"
 
 
@@ -82,8 +85,8 @@ def _gfx950():
     )
 
 
-def _flydsl_installed(present: bool):
-    return patch(f"{_AITER_MOD}.is_aiter_mxfp8_moe_available", return_value=present)
+def _aiter_moe_enabled(present: bool):
+    return patch.object(rocm_aiter_ops, "is_fused_moe_enabled", return_value=present)
 
 
 def test_aiter_mxfp8_registered():
@@ -109,13 +112,18 @@ def test_ep_supported(ep_size):
 @pytest.mark.parametrize(
     "present,ep_size,supported,reason_substr",
     [
-        (True, 1, True, None),  # gfx950 + flydsl + TP -> selectable
-        (True, 2, True, None),  # gfx950 + flydsl + EP -> selectable (expert_mask)
-        (False, 1, False, "flydsl package"),  # package missing -> clear reason
+        (True, 1, True, None),  # gfx950 + aiter MoE + TP -> selectable
+        (True, 2, True, None),  # gfx950 + aiter MoE + EP -> selectable (expert_mask)
+        (
+            False,
+            1,
+            False,
+            "does not support current device",
+        ),  # disabled -> not selected
     ],
 )
 def test_is_supported_config(present, ep_size, supported, reason_substr):
-    with _gfx950(), _flydsl_installed(present):
+    with _gfx950(), _aiter_moe_enabled(present):
         ok, reason = AiterMxfp8Experts.is_supported_config(
             AiterMxfp8Experts,
             _config(ep_size),
@@ -130,8 +138,8 @@ def test_is_supported_config(present, ep_size, supported, reason_substr):
 
 def test_explicit_moe_backend_aiter():
     """--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear
-    ValueError when the flydsl package is missing."""
-    with _gfx950(), _flydsl_installed(True):
+    ValueError when the aiter MoE runtime is disabled."""
+    with _gfx950(), _aiter_moe_enabled(True):
         assert (
             _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
             is AiterMxfp8Experts
@@ -142,18 +150,17 @@ def test_explicit_moe_backend_aiter():
         )
     with (
         _gfx950(),
-        _flydsl_installed(False),
-        pytest.raises(ValueError, match="flydsl package"),
+        _aiter_moe_enabled(False),
+        pytest.raises(ValueError, match="does not support current device"),
     ):
         _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
 
 
 def test_gfx950_picks_aiter():
-    """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins."""
-    # NOTE: Fp8MoeBackend.AITER_MXFP8 does not require VLLM_ROCM_USE_AITER=1
+    """Auto-select on real ROCm hardware with aiter MoE enabled -> FlyDSL wins."""
     with (
         patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=True),
-        _flydsl_installed(True),
+        _aiter_moe_enabled(True),
     ):
         backend, experts_cls = select_mxfp8_moe_backend(_config())
     assert backend is Fp8MoeBackend.AITER_MXFP8
@@ -161,7 +168,7 @@ def test_gfx950_picks_aiter():
 
 
 def test_gfx942_picks_emulation():
-    """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton
+    """Flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton
     dot_scaled backend wins instead."""
     with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False):
         backend, experts_cls = select_mxfp8_moe_backend(_config())
diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py
index 9157d4eaa254..1cd204d4b00b 100644
--- a/tests/kernels/moe/test_ocp_mx_moe.py
+++ b/tests/kernels/moe/test_ocp_mx_moe.py
@@ -214,8 +214,7 @@ def reference_moe(
     activation: str = "swiglu",
     use_interleaved_layout: bool = False,
 ):
-    """
-    Reference MoE implementation for accuracy testing.
+    """Reference MoE implementation for accuracy testing.
 
     Args:
         activation: One of "swiglu", "silu", "relu2". Controls the activation
@@ -224,6 +223,7 @@ def reference_moe(
             (gate=x[..., ::2], up=x[..., 1::2]) as used by SWIGLUOAI.
             If False, uses chunked layout (gate, up = chunk(x, 2)) as used
             by standard swiglu/silu.
+
     """
     # renormalize routing
     experts = torch.topk(roouting_logits, k=topk, dim=-1, sorted=True)
@@ -1441,8 +1441,7 @@ def test_rocm_mxfp4_moe_oracle(
     intermediate_size: int,
     monkeypatch: pytest.MonkeyPatch,
 ):
-    """
-    Test ROCm MXFP4 MoE using oracle functions.
+    """Test ROCm MXFP4 MoE using oracle functions.
 
     This test validates that the oracle functions work end-to-end:
     - select_mxfp4_moe_backend() selects a valid backend
@@ -1825,8 +1824,7 @@ def test_mxfp4_emulation_rounds_up_to_block_size(
 def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons(
     monkeypatch: pytest.MonkeyPatch,
 ):
-    """
-    select_mxfp4_moe_backend() must raise NotImplementedError, with the
+    """select_mxfp4_moe_backend() must raise NotImplementedError, with the
     collected per-backend unsupported reasons in the message, when no
     backend supports the requested deployment configuration.
     """
@@ -1870,6 +1868,61 @@ def is_supported_config(
         mxfp4_oracle.select_mxfp4_moe_backend(moe_config)
 
 
+@pytest.mark.parametrize("backend_name", ["AITER_TRITON_MXFP4_BF16", "AITER_MXFP4_FP8"])
+@pytest.mark.parametrize("use_ep", [False, True])
+def test_aiter_mxfp4_monolithic_rejects_expert_parallel(
+    monkeypatch: pytest.MonkeyPatch, backend_name: str, use_ep: bool
+):
+    """Global routing cannot index local expert weights, even without all-to-all."""
+    from tests.kernels.moe.utils import make_dummy_moe_config
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+    from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
+    from vllm.model_executor.layers.fused_moe.modular_kernel import (
+        FusedMoEActivationFormat,
+    )
+    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+        Mxfp4MoeBackend,
+        backend_to_kernel_cls,
+    )
+    from vllm.model_executor.layers.quantization.utils.quant_utils import (
+        kFp8StaticTensorSym,
+        kMxfp4Static,
+    )
+
+    experts_cls = backend_to_kernel_cls(Mxfp4MoeBackend[backend_name])[0]
+    monkeypatch.setattr(experts_cls, "_supports_current_device", lambda: True)
+    config = make_dummy_moe_config(
+        num_experts=256,
+        num_local_experts=64 if use_ep else 256,
+        experts_per_token=6,
+        hidden_dim=4096,
+        intermediate_size=2048,
+        activation=MoEActivation.SWIGLUOAI,
+    )
+    config = replace(
+        config,
+        routing_method=RoutingMethodType.Renormalize,
+        moe_parallel_config=replace(
+            config.moe_parallel_config,
+            tp_size=1 if use_ep else 4,
+            ep_size=4 if use_ep else 1,
+            use_ep=use_ep,
+        ),
+    )
+    assert not config.moe_parallel_config.use_all2all_kernels
+    supported, reason = experts_cls.is_supported_config(
+        experts_cls,
+        config,
+        kMxfp4Static,
+        kFp8StaticTensorSym if backend_name == "AITER_MXFP4_FP8" else None,
+        FusedMoEActivationFormat.Standard,
+    )
+
+    assert supported is not use_ep
+    if use_ep:
+        assert "parallel config" in reason
+
+
 # Every activation-quantizing OCP MX scheme must map to a `quant_dtype` that
 # `moe_kernel_quantize_input` actually dispatches on. Its final `else` returns
 # the activation untouched, so a name it does not know (e.g. "mxfp6" instead of
diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py
index 89244bed1b0f..0967f9499e0f 100644
--- a/tests/kernels/moe/test_rocm_aiter_moe.py
+++ b/tests/kernels/moe/test_rocm_aiter_moe.py
@@ -565,6 +565,74 @@ def test_aiter_moe_shared_experts_enablement_follows_env(
         assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected
 
 
+def test_aiter_moe_situv2_syncs_aiter_a4w4_env(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """SiTUv2 MoE should route to AITER a4w4 and clear legacy a8w4 overrides."""
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    _assert_aiter_supported()
+
+    with monkeypatch.context() as mp:
+        mp.delenv("AITER_SITUV2_A8W4", raising=False)
+        mp.delenv("AITER_SITUV2_A4W4", raising=False)
+        mp.setenv("VLLM_ROCM_USE_AITER", "1")
+        mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
+        mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "1")
+        _reload_envs()
+        rocm_aiter_ops.refresh_env_variables()
+
+        import os
+
+        assert os.environ.get("AITER_SITUV2_A4W4") == "1"
+        assert "AITER_SITUV2_A8W4" not in os.environ
+
+
+def test_aiter_moe_situv2_legacy_a8w4_alias_enables_a4w4(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Deprecated VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 still enables SiTUv2 a4w4."""
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    _assert_aiter_supported()
+
+    with monkeypatch.context() as mp:
+        mp.delenv("AITER_SITUV2_A8W4", raising=False)
+        mp.delenv("AITER_SITUV2_A4W4", raising=False)
+        mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", raising=False)
+        mp.setenv("VLLM_ROCM_USE_AITER", "1")
+        mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
+        mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "1")
+        _reload_envs()
+        rocm_aiter_ops.refresh_env_variables()
+
+        import os
+
+        assert os.environ.get("AITER_SITUV2_A4W4") == "1"
+        assert "AITER_SITUV2_A8W4" not in os.environ
+
+
+def test_aiter_moe_situv2_clears_aiter_env_when_disabled(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    _assert_aiter_supported()
+
+    with monkeypatch.context() as mp:
+        mp.setenv("AITER_SITUV2_A4W4", "1")
+        mp.setenv("VLLM_ROCM_USE_AITER", "1")
+        mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
+        mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", raising=False)
+        mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0")
+        _reload_envs()
+        rocm_aiter_ops.refresh_env_variables()
+
+        import os
+
+        assert "AITER_SITUV2_A4W4" not in os.environ
+
+
 @pytest.mark.parametrize("moe_padding", [True, False])
 def test_aiter_moe_padding_env_var(
     moe_padding: bool,
diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py
index 0eded09b1ae5..4baad24cc3cd 100644
--- a/tests/kernels/moe/test_routing.py
+++ b/tests/kernels/moe/test_routing.py
@@ -87,6 +87,21 @@ def test_multiple_expert_groups_use_grouped_topk() -> None:
     )
 
     assert isinstance(router, GroupedTopKRouter)
+    assert not router.skip_padding
+
+
+def test_grouped_topk_padding_skip_must_be_enabled() -> None:
+    router = create_fused_moe_router(
+        top_k=4,
+        global_num_experts=128,
+        use_grouped_topk=True,
+        num_expert_group=8,
+        topk_group=4,
+        skip_padding=True,
+    )
+
+    assert isinstance(router, GroupedTopKRouter)
+    assert router.skip_padding
 
 
 def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None:
@@ -214,8 +229,7 @@ def assert_routing_results_close(
     rtol: float = 1e-3,
     atol: float = 1e-3,
 ):
-    """
-    Compare routing results, sorting by expert ID first to handle non-deterministic
+    """Compare routing results, sorting by expert ID first to handle non-deterministic
     ordering from sorted=False in topk.
     """
     # Sort both results by expert IDs for consistent comparison
@@ -294,8 +308,7 @@ def assert_aiter_routing_valid(
 def baseline_fused_topk(
     router_logits: torch.Tensor, top_k: int, renormalize: bool
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Baseline for standard fused top-k routing.
+    """Baseline for standard fused top-k routing.
 
     Algorithm:
     1. Apply softmax to router logits
@@ -320,8 +333,7 @@ def baseline_fused_topk_bias(
     e_score_correction_bias: torch.Tensor,
     routed_scaling_factor: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Baseline for fused top-k with bias correction.
+    """Baseline for fused top-k with bias correction.
 
     Algorithm:
     1. Apply softmax to router logits
@@ -364,8 +376,7 @@ def baseline_grouped_topk(
     e_score_correction_bias: torch.Tensor | None,
     routed_scaling_factor: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Baseline for grouped top-k routing (e.g., DeepSeek).
+    """Baseline for grouped top-k routing (e.g., DeepSeek).
 
     Algorithm:
     1. Apply scoring function (softmax or sigmoid)
@@ -436,8 +447,7 @@ def baseline_grouped_topk(
 def baseline_custom_llama4(
     router_logits: torch.Tensor, top_k: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Baseline for Llama4 custom routing.
+    """Baseline for Llama4 custom routing.
 
     Algorithm:
     1. Select top-k expert indices (without softmax)
diff --git a/tests/kernels/moe/test_routing_simulator.py b/tests/kernels/moe/test_routing_simulator.py
index 35b1af618cf7..2cb122122fcf 100644
--- a/tests/kernels/moe/test_routing_simulator.py
+++ b/tests/kernels/moe/test_routing_simulator.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test script for the token-to-expert routing simulator.
+"""Test script for the token-to-expert routing simulator.
 
 This script demonstrates how to use the routing simulator to test
 different routing strategies and analyze their performance, including
diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py
index af9455b96039..895e168b2da3 100644
--- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py
+++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for FusedMoEFactory with routed_input_transform.
+"""Tests for FusedMoEFactory with routed_input_transform.
 
 Verifies that applying routed_input_transform inside FusedMoEFactory
 produces the same results as applying the transform manually outside.
diff --git a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
index 85b992ad6544..0ca08c50b324 100644
--- a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+++ b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
@@ -142,9 +142,7 @@ def silu_mul_quant(
 
 
 def pack_scales(x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor:
-    """
-    pack float32 scales into a int32 tensor
-    """
+    """Pack float32 scales into a int32 tensor."""
     assert x.dtype == torch.float32
     E, T, G = x.size()
 
@@ -170,8 +168,7 @@ def ref_with_scale_fmt(
     up: torch.Tensor,
     scale_fmt: DeepGemmQuantScaleFMT,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    The precision types of the operations triggered by this function
+    """The precision types of the operations triggered by this function
     match closely with the kernel implementation so we compare more
     accurately.
     """
@@ -204,8 +201,7 @@ def ref_with_scale_fmt(
 
 
 def token_random(E, T, H2, tokens_per_expert):
-    """
-    Initialize each token in a random range so we test a range of
+    """Initialize each token in a random range so we test a range of
     scale values.
     """
     y = torch.empty((E, T, H2), dtype=torch.bfloat16, device=DEVICE)
diff --git a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py
index bfe763b44ab2..204954d1fe46 100644
--- a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py
+++ b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py
@@ -19,11 +19,9 @@
 
 
 def reference_quant(x: torch.Tensor, use_ue8m0: bool):
-    """
-    Reference triton quant kernel from,
+    """Reference triton quant kernel from,
     vllm.model_executor.layers.quantization.utils.fp8_utils
     """
-
     x_q = torch.empty_like(x, device=x.device, dtype=FLOAT8_DTYPE)
 
     # Allocate the scale tensor in column-major format.
diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py
index ab9072576d92..bc438c66931b 100644
--- a/tests/kernels/moe/test_topk_softplus_sqrt.py
+++ b/tests/kernels/moe/test_topk_softplus_sqrt.py
@@ -81,6 +81,65 @@ def _torch_topk_softplus_sqrt(
     return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection")
+@pytest.mark.parametrize("has_hash_routing", [False, True])
+@pytest.mark.parametrize(
+    "backend,expected_experts",
+    [
+        ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"),
+        ("aiter", "AiterExperts"),
+        ("triton_unfused", "UnfusedOAITritonExperts"),
+    ],
+)
+def test_hash_routing_backend_selection(
+    dist_init,
+    default_vllm_config,
+    monkeypatch,
+    has_hash_routing,
+    backend,
+    expected_experts,
+):
+    """Hash tables must reach backend selection; modular AITER remains eligible."""
+    from dataclasses import replace
+
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory
+    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+        select_deepseek_v4_mxfp4_moe_backend,
+    )
+    from vllm.platforms import rocm
+
+    # Exercise both native MX and monolithic AITER selection on any ROCm device.
+    monkeypatch.setattr(rocm, "on_gfx950", lambda: True)
+    monkeypatch.setattr(rocm, "on_gfx1250", lambda: False)
+    monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
+    monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: True)
+    default_vllm_config.kernel_config.moe_backend = "triton"
+    layer = FusedMoEFactory(
+        num_experts=8,
+        top_k=2,
+        hidden_size=256,
+        intermediate_size=256,
+        params_dtype=torch.bfloat16,
+        scoring_func="sqrtsoftplus",
+        e_score_correction_bias=torch.zeros(8),
+        hash_indices_table=(
+            torch.tensor([[0, 1], [2, 3]], dtype=torch.int32)
+            if has_hash_routing
+            else None
+        ),
+        prefix="hash_routing_selection",
+    )
+    config = replace(layer.routed_experts.moe_config, moe_backend=backend)
+
+    if has_hash_routing and backend == "aiter_triton_mxfp4_bf16":
+        with pytest.raises(ValueError, match="hash routing"):
+            select_deepseek_v4_mxfp4_moe_backend(config)
+    else:
+        _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config)
+        assert experts_cls.__name__ == expected_experts
+
+
 def test_torch_topk_softplus_sqrt_breaks_ties_by_expert_id():
     gating_output = torch.tensor([[2.0, 1.0, 1.0, 0.0]])
 
diff --git a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py
index 01246325c803..8b10ff75e6db 100644
--- a/tests/kernels/moe/test_triton_moe_ptpc_fp8.py
+++ b/tests/kernels/moe/test_triton_moe_ptpc_fp8.py
@@ -57,7 +57,6 @@ def fp8_mask(a, mask):
 def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
     """This function performs fused moe with per-column int8
     quantization using native torch."""
-
     B, D = a.shape
     # Perform per-token quantization
     a_q, a_s = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True)
diff --git a/tests/kernels/moe/test_trtllm_bf16_moe.py b/tests/kernels/moe/test_trtllm_bf16_moe.py
index 7b6b7d81ffcf..5e569e153a7c 100644
--- a/tests/kernels/moe/test_trtllm_bf16_moe.py
+++ b/tests/kernels/moe/test_trtllm_bf16_moe.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for the FlashInfer TRTLLM BF16 MoE backend
+"""Tests for the FlashInfer TRTLLM BF16 MoE backend
 (`TrtLlmBf16ExpertsModular`).
 
 This mirrors the TRTLLM NvFP4 modular test shape: construct the modular
@@ -13,6 +12,7 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from tests.kernels.moe.utils import make_dummy_moe_config
 from tests.kernels.utils import torch_moe
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -28,6 +28,9 @@
 from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
     TrtLlmBf16ExpertsModular,
 )
+from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
+    TrtLlmBf16LoRAExperts,
+)
 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
     UnquantizedMoeBackend,
     convert_to_unquantized_kernel_format,
@@ -76,6 +79,10 @@ def test_trtllm_bf16_moe_modular_no_graph(
         a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
         w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
         w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
+        # The FlashInfer conversion may rewrite unpadded input storage in place.
+        # Preserve the original layout for the independent torch reference.
+        reference_w1 = w1.clone()
+        reference_w2 = w2.clone()
         score = torch.randn((m, e), device="cuda", dtype=dtype)
         scores = torch.softmax(score, dim=-1, dtype=torch.float32)
         topk_weights, topk_ids = torch.topk(scores, topk)
@@ -134,8 +141,8 @@ def test_trtllm_bf16_moe_modular_no_graph(
 
         torch_output = torch_moe(
             a,
-            w1,
-            w2,
+            reference_w1,
+            reference_w2,
             score,
             topk,
             activation=MoEActivation.SILU,
@@ -147,3 +154,48 @@ def test_trtllm_bf16_moe_modular_no_graph(
             atol=1e-1,
             rtol=2e-1,
         )
+
+
+@pytest.mark.parametrize("has_lora_delta", [False, True])
+@torch.inference_mode()
+def test_trtllm_bf16_lora_accepts_checkpoint_shaped_weights(has_lora_delta):
+    """LoRA dispatch accepts both 3D parameters and legacy 4D packed views."""
+    config = make_dummy_moe_config(
+        num_experts=128, hidden_dim=256, intermediate_size=128
+    )
+    experts = TrtLlmBf16LoRAExperts(config, FUSED_MOE_UNQUANTIZED_CONFIG)
+    x = torch.randn(2, 256, device="cuda", dtype=torch.bfloat16) / 10
+    w1, w2 = convert_to_unquantized_kernel_format(
+        UnquantizedMoeBackend.FLASHINFER_TRTLLM,
+        config,
+        torch.randn(128, 256, 256, device="cuda", dtype=torch.bfloat16) / 10,
+        torch.randn(128, 256, 128, device="cuda", dtype=torch.bfloat16) / 10,
+    )
+    assert w1.ndim == w2.ndim == 3
+    topk = (
+        torch.tensor([[0], [1]], device="cuda", dtype=torch.int32),
+        torch.ones(2, 1, device="cuda", dtype=torch.float32),
+    )
+    delta = torch.randn(2, 1, 256, device="cuda", dtype=torch.bfloat16)
+
+    def invoke(w1, w2):
+        result = experts.invoke_routed_moe(
+            hidden_states=x,
+            w1=w1,
+            w2=w2,
+            topk_ids_and_weights=topk,
+            gemm1_lora_delta=delta if has_lora_delta else None,
+            global_num_experts=128,
+            a1q_scale=None,
+            output=torch.empty_like(x),
+        )
+        if has_lora_delta:
+            # Only compare rows belonging to real tokens, excluding padding.
+            indices = result[2].flatten().long()
+            return result[0][indices], result[3][indices]
+        return result
+
+    actual = invoke(w1, w2)
+    expected = invoke(w1.view(128, 4, 256, 64), w2.view(128, 2, 256, 64))
+    for actual_tensor, expected_tensor in zip(actual, expected):
+        torch.testing.assert_close(actual_tensor, expected_tensor)
diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py
index 59eaa4c1390c..b9684396b245 100644
--- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py
+++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for the FlashInfer TRTLLM NvFP4 MoE backend
+"""Tests for the FlashInfer TRTLLM NvFP4 MoE backend
 (`TrtLlmNvFp4ExpertsModular`).
 
 Covers the activations the wrapper claims to support — SiLU, RELU^2 (non-gated),
diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py
index d4c037d978a3..9d0f4d6900c4 100644
--- a/tests/kernels/moe/test_unquantized_backend_selection.py
+++ b/tests/kernels/moe/test_unquantized_backend_selection.py
@@ -6,6 +6,7 @@
 import torch
 
 from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
     UnquantizedMoeBackend,
@@ -228,6 +229,65 @@ def test_select_rocm_aiter_backend(mock_aiter_enabled, mock_has_flashinfer):
         assert expert_cls is not None
 
 
+@patch(
+    "vllm.utils.flashinfer.has_flashinfer",
+    return_value=False,
+)
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops."
+    "is_fused_moe_enabled",
+    return_value=True,
+)
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops."
+    "is_rdna_aiter_enabled",
+    return_value=False,
+)
+@pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-specific backend selection test"
+)
+def test_select_rocm_aiter_backend_non_gated_activation_falls_back(
+    mock_rdna_disabled, mock_aiter_enabled, mock_has_flashinfer, monkeypatch
+):
+    """Test ROCm backend selection falls back (not raises) for non-gated MoE."""
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
+
+    with patch(
+        "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
+    ) as mock_platform:
+        mock_platform.is_cuda.return_value = False
+        mock_platform.is_rocm.return_value = True
+        mock_platform.is_cpu.return_value = False
+        mock_platform.is_xpu.return_value = False
+        mock_platform.is_tpu.return_value = False
+        mock_platform.is_out_of_tree.return_value = False
+
+        moe_config = make_dummy_moe_config(activation=MoEActivation.SILU_NO_MUL)
+        assert moe_config.is_act_and_mul is False
+
+        selected_backend, expert_cls = select_unquantized_moe_backend(
+            moe_config=moe_config,
+        )
+
+        assert selected_backend != UnquantizedMoeBackend.AITER
+        assert expert_cls is not None
+
+
+@pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-specific backend selection test"
+)
+def test_explicit_aiter_backend_non_gated_activation_still_raises():
+    """Explicit `--moe-backend aiter` still raises for non-gated MoE;
+    only the env-var opt-in path falls back."""
+    moe_config = make_dummy_moe_config(activation=MoEActivation.SILU_NO_MUL)
+    moe_config.moe_backend = "aiter"  # explicit pin, not "auto"
+    assert moe_config.is_act_and_mul is False
+
+    with pytest.raises(ValueError):
+        select_unquantized_moe_backend(moe_config=moe_config)
+
+
 @patch(
     "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config",
     return_value=(True, None),
diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py
index c2084800f979..0202b380c526 100644
--- a/tests/kernels/moe/test_zero_expert_moe.py
+++ b/tests/kernels/moe/test_zero_expert_moe.py
@@ -83,8 +83,7 @@ def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_token
 
 @pytest.mark.parametrize("num_tokens", [1, 32])
 def test_zero_expert_moe_forward(zero_expert_moe, num_tokens):
-    """
-    Run a forward pass through FusedMoEFactory with zero experts
+    """Run a forward pass through FusedMoEFactory with zero experts
     and verify output shape.
     """
     layer, vllm_config = zero_expert_moe
diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py
index e91fa58f3e81..d92e4d6129ba 100644
--- a/tests/kernels/moe/utils.py
+++ b/tests/kernels/moe/utils.py
@@ -64,8 +64,7 @@ def make_dummy_moe_config(
     max_num_tokens: int = 512,
     activation: MoEActivation = MoEActivation.SILU,
 ) -> FusedMoEConfig:
-    """
-    This is a dummy config for the mk constructor interface
+    """This is a dummy config for the mk constructor interface
     as most kernels like DeepGEMM, CUTLASSFp4, Triton, MARLIN
     do not actually use this config.
 
@@ -706,3 +705,27 @@ def check_accuracy(a, b, atol, rtol, percent):
             f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} "
             f"(threshold: {1 - percent:.4f})"
         )
+
+
+def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8):
+    """Weight/scale layouts for mxfp4 MoE, as (layout, opts) pairs.
+
+    triton_kernels 3.8 returns layout instances; earlier versions return a
+    (layout, opts) tuple.
+    """
+    from triton_kernels.tensor_details import layout
+
+    from vllm.utils.import_utils import get_triton_kernels_version
+
+    if get_triton_kernels_version() == "3.8":
+        w = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis)
+        s = layout.make_default_matmul_mxfp4_w_scale_layout(
+            mx_axis=mx_axis, num_warps=num_warps
+        )
+        return w, {}, s, {}
+
+    w, w_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis)
+    s, s_opts = layout.make_default_matmul_mxfp4_w_scale_layout(
+        mx_axis=mx_axis, num_warps=num_warps
+    )
+    return w, w_opts, s, s_opts
diff --git a/tests/kernels/quantization/nvfp4_utils.py b/tests/kernels/quantization/nvfp4_utils.py
index e0120301f7e3..04f7943ae630 100644
--- a/tests/kernels/quantization/nvfp4_utils.py
+++ b/tests/kernels/quantization/nvfp4_utils.py
@@ -118,6 +118,7 @@ def dequant_nvfp4_kv_cache(
 
     Returns:
         [..., num_heads, block_size, head_size] float32.
+
     """
     data_dim = head_size // 2
     scale_dim = head_size // 16
diff --git a/tests/kernels/quantization/test_allspark_gemm.py b/tests/kernels/quantization/test_allspark_gemm.py
deleted file mode 100644
index 7e4ce1afaf96..000000000000
--- a/tests/kernels/quantization/test_allspark_gemm.py
+++ /dev/null
@@ -1,128 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import pytest
-import torch
-
-from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
-from vllm import _custom_ops as ops
-from vllm.model_executor.layers.quantization.utils.allspark_utils import (
-    ALLSPARK_AMPERE_K_ALIGN,
-    ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
-    ALLSPARK_AMPERE_N_ALIGN,
-)
-from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
-from vllm.platforms import current_platform
-from vllm.scalar_type import scalar_types
-from vllm.utils.platform_utils import num_compute_units
-
-
-def is_gptq_allspark_supported(min_capability: int, max_capability: int) -> bool:
-    if not current_platform.is_cuda():
-        return False
-
-    capability = current_platform.get_device_capability()
-    assert capability is not None
-
-    return (
-        capability.to_int() >= min_capability and capability.to_int() <= max_capability
-    )
-
-
-MNK_FACTORS = [
-    (1, 4, 8),
-    (13, 17, 67),
-    (26, 37, 13),
-    (48, 16, 24),
-    (67, 13, 88),
-    (257, 13, 11),
-    (658, 13, 11),
-    (1033, 9, 17),
-]
-
-DTYPES = [torch.float16, torch.bfloat16]
-HAS_ZP_OPTS = [False, True]
-
-
-def compute_max_diff(output, output_ref):
-    return torch.mean(torch.abs(output - output_ref)) / torch.mean(
-        torch.abs(output_ref)
-    )
-
-
-def rand_data(shape, dtype=torch.float16):
-    return torch.randn(shape, dtype=dtype, device="cuda")
-
-
-@pytest.mark.skipif(
-    not is_gptq_allspark_supported(80, 89),
-    reason="AllSpark Ampere kernel is not supported on this GPU type.",
-)
-@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
-@pytest.mark.parametrize("group_size", [-1])
-@pytest.mark.parametrize("has_zp", HAS_ZP_OPTS)
-@pytest.mark.parametrize("dtype", DTYPES)
-def test_gptq_allspark_gemm_ampere(mnk_factors, group_size, has_zp, dtype):
-    m_factor, n_factor, k_factor = mnk_factors
-    m = m_factor
-    n = n_factor * ALLSPARK_AMPERE_N_ALIGN
-    k = k_factor * ALLSPARK_AMPERE_K_ALIGN
-
-    input = rand_data((m, k), dtype=dtype)
-    weight = rand_data((k, n), dtype=dtype)
-
-    # Quantize weights.
-    w_ref, qw, s, zp = quantize_weights(
-        weight, scalar_types.uint8b128, group_size, has_zp
-    )
-
-    qw = qw.to(torch.uint8)
-    if has_zp:
-        zp = zp.to(dtype)
-    properties = torch.cuda.get_device_properties(qw.device.index)
-    sm_count = num_compute_units(qw.device.index)
-    sm_version = properties.major * 10 + properties.minor
-
-    n_32align = (n + 32 - 1) // 32 * 32
-
-    qw_reorder, s_reorder, zp_reorder = ops.allspark_repack_weight(qw, s, zp, has_zp)
-    opcheck(
-        torch.ops._C.rearrange_kn_weight_as_n32k16_order,
-        (qw, s, zp, has_zp, qw_reorder, s_reorder, zp_reorder, k, n, n_32align),
-    )
-
-    opcheck(
-        torch.ops._C.allspark_w8a16_gemm,
-        (
-            input,
-            qw_reorder,
-            s_reorder,
-            zp_reorder,
-            n,
-            group_size,
-            sm_count,
-            sm_version,
-            ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
-            has_zp,
-            True,
-        ),
-        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
-    )
-    output = ops.allspark_w8a16_gemm(
-        input,
-        qw_reorder,
-        s_reorder,
-        zp_reorder,
-        n,
-        group_size,
-        sm_count,
-        sm_version,
-        ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
-        has_zp,
-        True,
-    )
-
-    output_ref = torch.matmul(input, w_ref)
-    torch.accelerator.synchronize()
-    max_diff = compute_max_diff(output, output_ref)
-
-    assert max_diff < 0.04
diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py
index fd530e7ffb53..038b94bd2d74 100644
--- a/tests/kernels/quantization/test_block_fp8.py
+++ b/tests/kernels/quantization/test_block_fp8.py
@@ -27,6 +27,7 @@
 from vllm.utils.deep_gemm import (
     fp8_gemm_nt,
     get_tma_aligned_size,
+    is_deep_gemm_e8m0_used,
     per_block_cast_to_fp8,
     should_use_deepgemm_for_fp8_linear,
 )
@@ -161,11 +162,16 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
 @pytest.mark.skipif(
     not current_platform.is_cuda(), reason="CUTLASS only supported on CUDA platform."
 )
+@pytest.mark.parametrize(
+    # 65/66/67 cover all M%4 residue classes above the SM100 swapAB
+    # threshold (m <= 64); 1026 crosses multiple 128-row SF atoms.
+    "M",
+    [32, 65, 66, 67, 1026],
+)
 @torch.inference_mode()
-def test_w8a8_block_fp8_cutlass_matmul():
+def test_w8a8_block_fp8_cutlass_matmul(M):
     # Test simple case where weight.shape % 128 != 0,
     # like in DSV3 kv_a_proj_with_mqa
-    M = 32
     N = 576
     K = 7168
     block_size = [128, 128]
@@ -291,7 +297,11 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
     A_fp8, As_fp8 = per_token_group_quant_fp8(
         A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True
     )
-    B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size)
+    B_fp8, Bs_fp8 = per_block_cast_to_fp8(
+        B_fp32,
+        block_size=block_size,
+        use_ue8m0=is_deep_gemm_e8m0_used(),
+    )
 
     As = As_fp8.to(torch.float32)
     Bs = Bs_fp8.to(torch.float32)
diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
index 3154e2cb98bb..0b0d12a66c85 100644
--- a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
+++ b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for CPU FP8 W8A16 block-scaled GEMM kernel (fp8_scaled_mm_cpu).
+"""Tests for CPU FP8 scaled-mm GEMM kernels.
+
+Covers:
+  * fp8_scaled_mm_cpu (W8A16)
+  * fp8_scaled_mm_with_quant (W8A8)
 
 Run `pytest tests/kernels/quantization/test_cpu_fp8_scaled_mm.py -v`.
 """
@@ -9,21 +13,45 @@
 import torch
 
 from vllm import _custom_ops as ops
+from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
+    CPUFp8PerTensorScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
+    FP8ScaledMMLinearLayerConfig,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8StaticTensorSym,
+)
 from vllm.platforms import current_platform
 
 if not current_platform.is_cpu():
     pytest.skip("skipping CPU-only tests", allow_module_level=True)
 
-if not ops._supports_cpu_fp8_w8a16:
-    pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True)
+requires_cpu_fp8_w8a16 = pytest.mark.skipif(
+    not ops._supports_cpu_fp8_w8a16, reason="fp8_scaled_mm_cpu op not available"
+)
+requires_cpu_fp8_w8a8 = pytest.mark.skipif(
+    not ops.cpu_has_amx_fp8(),
+    reason="requires native AMX-FP8 hardware",
+)
+
+if not torch.cpu._is_amx_tile_supported():
+    pytest.skip("requires AMX tile support", allow_module_level=True)
 
+FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
 BLOCK_SIZE = [128, 128]
+GROUP_SIZE = 128
 
 
 def cdiv(a: int, b: int) -> int:
     return -(a // -b)
 
 
+# ---------------------------------------------------------------------------
+# W8A16 helpers
+# ---------------------------------------------------------------------------
+
+
 def quantize_weight_block_fp8(
     weight: torch.Tensor,
     block_size: list[int],
@@ -33,6 +61,7 @@ def quantize_weight_block_fp8(
     Returns:
         fp8_weight: [N, K] float8_e4m3fn
         scales: [n_tiles, k_tiles] float32
+
     """
     N, K = weight.shape
     block_n, block_k = block_size
@@ -113,7 +142,7 @@ def ref_fp8_block_scaled_mm(
 
 
 # ---------------------------------------------------------------------------
-# Test parameters
+# W8A16 test: fp8_scaled_mm_cpu
 # ---------------------------------------------------------------------------
 M_SIZES = [1, 4, 16, 64, 128]
 # (N, K) — weight shape is [N, K], output has N columns.
@@ -128,6 +157,7 @@ def ref_fp8_block_scaled_mm(
 ]
 
 
+@requires_cpu_fp8_w8a16
 @pytest.mark.parametrize("M", M_SIZES)
 @pytest.mark.parametrize("N,K", NK_SIZES)
 @pytest.mark.parametrize("use_bias", [False, True])
@@ -160,3 +190,308 @@ def test_cpu_fp8_scaled_mm(M: int, N: int, K: int, use_bias: bool):
 
     assert kernel_out.dtype == out_dtype
     torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01)
+
+
+def quantize_weight_per_tensor_fp8(
+    weight: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize weight [N, K] to FP8 with a single per-tensor scale."""
+    fp8_max = torch.finfo(torch.float8_e4m3fn).max
+    scale = weight.abs().amax() / fp8_max
+    scale = torch.where(scale == 0, torch.ones_like(scale), scale)
+    q = (weight / scale).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn)
+    return q, scale
+
+
+def ref_fp8_per_tensor_scaled_mm(
+    x: torch.Tensor,
+    fp8_weight: torch.Tensor,
+    scale: torch.Tensor,
+    bias: torch.Tensor | None,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    w_dq = fp8_weight.float() * scale
+    out = torch.mm(x.float(), w_dq.t())
+    if bias is not None:
+        out = out + bias.float()
+    return out.to(out_dtype)
+
+
+# ---------------------------------------------------------------------------
+# W8A8 helpers
+# ---------------------------------------------------------------------------
+
+
+def quantize_weight_per_tensor(
+    weight: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize weight [N, K] to FP8 with a single per-tensor scale, broadcast
+    to [N, 1] (the shape float8_linear_prepack_cpu expects)."""
+    N = weight.shape[0]
+    abs_max = weight.abs().max()
+    scale = (abs_max / FP8_MAX).clamp(min=1e-7)
+    q = (weight.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+    return q, scale.expand(N).contiguous().view(N, 1)
+
+
+def quantize_weight_per_channel(
+    weight: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize weight [N, K] to FP8 with per-channel (per-row) scales [N, 1]."""
+    abs_max = weight.abs().amax(dim=1, keepdim=True)
+    scale = (abs_max / FP8_MAX).clamp(min=1e-7)
+    q = (weight.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+    return q, scale.float()
+
+
+def quantize_weight_per_group(
+    weight: torch.Tensor,
+    group_size: int = GROUP_SIZE,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize weight [N, K] to FP8 with per-group (128-block) scales [N, G]."""
+    N, K = weight.shape
+    G = K // group_size
+    w_groups = weight.view(N, G, group_size)
+    abs_max = w_groups.abs().amax(dim=2, keepdim=True)
+    scale = (abs_max / FP8_MAX).clamp(min=1e-7)
+    q = (w_groups.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+    return q.view(N, K), scale.squeeze(2).float()
+
+
+def quantize_act_per_token(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Reference per-token (row) dynamic FP8 quantization of activations."""
+    x_f = x.float()
+    scale = (x_f.abs().amax(dim=1, keepdim=True) / FP8_MAX).clamp(min=1e-7)
+    scaled = (x_f / scale).clamp(-FP8_MAX, FP8_MAX)
+    # Round via fp16 to match the AVX10.2 kernel path.
+    q = scaled.half().to(torch.float8_e4m3fn)
+    return q, scale
+
+
+def reference_static_fp8_linear(
+    x: torch.Tensor,  # BF16 [M, K]
+    weight_fp8: torch.Tensor,  # FP8 [N, K]
+    weight_scale: torch.Tensor,  # float32 [N, 1]
+    act_scale: torch.Tensor,  # float32 scalar
+    bias: torch.Tensor | None,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Static per-tensor: quantize/dequant x with act_scale, dequant weight, matmul."""
+    x_scaled = (x.float() / act_scale.item()).clamp(-FP8_MAX, FP8_MAX)
+    # Round via fp16 to match the AVX10.2 kernel path.
+    x_q = x_scaled.half().to(torch.float8_e4m3fn)
+    x_dq = x_q.float() * act_scale.item()
+    w_dq = weight_fp8.float() * weight_scale.view(-1, 1)
+    out = torch.mm(x_dq, w_dq.t())
+    if bias is not None:
+        out = out + bias.float()
+    return out.to(out_dtype)
+
+
+NK_SIZES_PER_TENSOR = [
+    (32, 64),
+    (5120, 5120),
+]
+
+
+@pytest.mark.parametrize("M", [1, 64])
+@pytest.mark.parametrize("N,K", NK_SIZES_PER_TENSOR)
+@pytest.mark.parametrize("use_bias", [False, True])
+def test_cpu_fp8_per_tensor_scaled_mm_kernel(
+    M: int, N: int, K: int, use_bias: bool, default_vllm_config
+):
+    """CPUFp8PerTensorScaledMMLinearKernel correctness against float reference.
+
+    Exercises the full kernel class (not just the raw op), including the
+    weight-orientation fixup: Fp8LinearMethod stores `layer.weight` as
+    [K, N] (torch._scaled_mm convention) before calling
+    process_weights_after_loading, so the kernel must transpose back to
+    [N, K] before VNNI-packing.
+    """
+    torch.manual_seed(0)
+    out_dtype = torch.bfloat16
+
+    x = torch.randn(M, K, dtype=out_dtype) / (K**0.5)
+    w_f32 = torch.randn(N, K, dtype=torch.float32) / (K**0.5)
+    fp8_weight, scale = quantize_weight_per_tensor_fp8(w_f32)
+    bias = torch.randn(N, dtype=torch.float32) * 0.1 if use_bias else None
+
+    ref_out = ref_fp8_per_tensor_scaled_mm(x, fp8_weight, scale, bias, out_dtype)
+
+    config = FP8ScaledMMLinearLayerConfig(
+        weight_quant_key=kFp8StaticTensorSym,
+        activation_quant_key=kFp8StaticTensorSym,
+        weight_shape=(N, K),
+        input_dtype=out_dtype,
+        out_dtype=out_dtype,
+    )
+    kernel = CPUFp8PerTensorScaledMMLinearKernel(
+        config,
+        layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"],
+    )
+
+    layer = torch.nn.Module()
+    # Fp8LinearMethod stores the weight transposed to [K, N] before calling
+    # process_weights_after_loading.
+    layer.register_parameter(
+        "weight", torch.nn.Parameter(fp8_weight.t().contiguous(), requires_grad=False)
+    )
+    layer.register_parameter(
+        "weight_scale", torch.nn.Parameter(scale.clone(), requires_grad=False)
+    )
+    kernel.process_weights_after_loading(layer)
+
+    kernel_out = kernel.apply_weights(layer, x, bias)
+
+    assert kernel_out.dtype == out_dtype
+    torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01)
+
+
+@pytest.mark.parametrize("n", [16, 48])
+def test_cpu_fp8_per_tensor_kernel_rejects_non_multiple_of_32_n(n: int):
+    """The AMX tinygemm kernel tiles N in chunks of 32 and cannot handle a
+    remainder tile smaller than that, so can_implement must reject any N
+    that isn't a multiple of 32 rather than let it crash the process.
+    """
+    config = FP8ScaledMMLinearLayerConfig(
+        weight_quant_key=kFp8StaticTensorSym,
+        activation_quant_key=kFp8StaticTensorSym,
+        weight_shape=(n, 64),
+        input_dtype=torch.bfloat16,
+        out_dtype=torch.bfloat16,
+    )
+    supported, _ = CPUFp8PerTensorScaledMMLinearKernel.can_implement(config)
+    assert not supported
+
+
+def reference_dynamic_fp8_linear(
+    x: torch.Tensor,  # BF16 [M, K]
+    weight_fp8: torch.Tensor,  # FP8 [N, K]
+    weight_scale: torch.Tensor,  # float32 [N, G]
+    bias: torch.Tensor | None,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Dynamic per-token act x per-group weight: quantize x per-token,
+    dequant weight per-group, matmul."""
+    N, K = weight_fp8.shape
+    G = weight_scale.shape[1]
+    group_size = K // G
+    x_q, x_scale = quantize_act_per_token(x)
+    x_dq = x_q.float() * x_scale
+    w_dq = (
+        weight_fp8.float().view(N, G, group_size) * weight_scale.view(N, G, 1)
+    ).view(N, K)
+    out = torch.mm(x_dq, w_dq.t())
+    if bias is not None:
+        out = out + bias.float()
+    return out.to(out_dtype)
+
+
+# ---------------------------------------------------------------------------
+# W8A8 tests: fp8_scaled_mm_with_quant
+# ---------------------------------------------------------------------------
+W8A8_M_SIZES = [1, 4, 16, 64]
+W8A8_NK_SIZES = [(128, 256), (256, 512), (64, 128)]
+
+
+@requires_cpu_fp8_w8a8
+@pytest.mark.parametrize("static_activation", [True, False])
+@pytest.mark.parametrize("M", W8A8_M_SIZES)
+@pytest.mark.parametrize("N,K", W8A8_NK_SIZES)
+@pytest.mark.parametrize("use_bias", [False, True])
+def test_cpu_fp8_w8a8_per_tensor(
+    M: int, N: int, K: int, use_bias: bool, static_activation: bool
+):
+    """CPUFP8W8A8ScaledMMLinearKernel with per-tensor weight scale."""
+    torch.manual_seed(42)
+    out_dtype = torch.bfloat16
+
+    x = torch.randn(M, K, dtype=torch.bfloat16)
+    weight = torch.randn(N, K, dtype=torch.bfloat16)
+    bias = torch.randn(N, dtype=torch.float32) if use_bias else None
+
+    w_fp8, w_scale = quantize_weight_per_tensor(weight)
+    packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale)
+
+    if static_activation:
+        act_scale = (x.float().abs().max() / FP8_MAX).clamp(min=1e-7).view(1)
+        kernel_out = ops.fp8_scaled_mm_with_quant(
+            x, act_scale, False, packed_weight, packed_scale, bias, out_dtype
+        )
+        ref_out = reference_static_fp8_linear(
+            x, w_fp8, w_scale, act_scale, bias, out_dtype
+        )
+        atol = 0.05
+    else:
+        kernel_out = ops.fp8_scaled_mm_with_quant(
+            x, None, True, packed_weight, packed_scale, bias, out_dtype
+        )
+        ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype)
+        atol = 0.05
+
+    assert kernel_out.dtype == out_dtype
+    torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=atol)
+
+
+@requires_cpu_fp8_w8a8
+@pytest.mark.parametrize("static_activation", [True, False])
+@pytest.mark.parametrize("M", W8A8_M_SIZES)
+@pytest.mark.parametrize("N,K", W8A8_NK_SIZES)
+@pytest.mark.parametrize("use_bias", [False, True])
+def test_cpu_fp8_w8a8_per_channel(
+    M: int, N: int, K: int, use_bias: bool, static_activation: bool
+):
+    """CPUFP8W8A8ScaledMMLinearKernel with per-channel weight scale."""
+    torch.manual_seed(42)
+    out_dtype = torch.bfloat16
+
+    x = torch.randn(M, K, dtype=torch.bfloat16)
+    weight = torch.randn(N, K, dtype=torch.bfloat16)
+    bias = torch.randn(N, dtype=torch.float32) if use_bias else None
+
+    w_fp8, w_scale = quantize_weight_per_channel(weight)
+    packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale)
+
+    if static_activation:
+        act_scale = (x.float().abs().max() / FP8_MAX).clamp(min=1e-7).view(1)
+        kernel_out = ops.fp8_scaled_mm_with_quant(
+            x, act_scale, False, packed_weight, packed_scale, bias, out_dtype
+        )
+        ref_out = reference_static_fp8_linear(
+            x, w_fp8, w_scale, act_scale, bias, out_dtype
+        )
+        atol = 0.05
+    else:
+        kernel_out = ops.fp8_scaled_mm_with_quant(
+            x, None, True, packed_weight, packed_scale, bias, out_dtype
+        )
+        ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype)
+        atol = 0.05
+
+    assert kernel_out.dtype == out_dtype
+    torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=atol)
+
+
+@requires_cpu_fp8_w8a8
+@pytest.mark.parametrize("M", W8A8_M_SIZES)
+@pytest.mark.parametrize("N,K", W8A8_NK_SIZES)
+@pytest.mark.parametrize("use_bias", [False, True])
+def test_cpu_fp8_w8a8_per_group(M: int, N: int, K: int, use_bias: bool):
+    """CPUFp8W8A8BlockScaledMMKernel: per-group (128-block) weight scale +
+    dynamic per-token activation scale."""
+    torch.manual_seed(42)
+    out_dtype = torch.bfloat16
+
+    x = torch.randn(M, K, dtype=torch.bfloat16)
+    weight = torch.randn(N, K, dtype=torch.bfloat16)
+    bias = torch.randn(N, dtype=torch.float32) if use_bias else None
+
+    w_fp8, w_scale = quantize_weight_per_group(weight)
+    packed_weight, packed_scale = torch.ops._C.float8_linear_prepack_cpu(w_fp8, w_scale)
+    kernel_out = ops.fp8_scaled_mm_with_quant(
+        x, None, True, packed_weight, packed_scale, bias, out_dtype
+    )
+    ref_out = reference_dynamic_fp8_linear(x, w_fp8, w_scale, bias, out_dtype)
+
+    assert kernel_out.dtype == out_dtype
+    torch.testing.assert_close(kernel_out, ref_out, rtol=0.05, atol=0.05)
diff --git a/tests/kernels/quantization/test_cutlass_scaled_mm.py b/tests/kernels/quantization/test_cutlass_scaled_mm.py
index a522d3171450..9e1d6d2cc76c 100644
--- a/tests/kernels/quantization/test_cutlass_scaled_mm.py
+++ b/tests/kernels/quantization/test_cutlass_scaled_mm.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for cutlass kernels
+"""Tests for cutlass kernels.
 
 Run `pytest tests/kernels/quantization/test_cutlass_scaled_mm.py`.
 """
diff --git a/tests/kernels/quantization/test_cutlass_w4a8.py b/tests/kernels/quantization/test_cutlass_w4a8.py
index 8cfc993fe8e8..33bb7bccc74e 100644
--- a/tests/kernels/quantization/test_cutlass_w4a8.py
+++ b/tests/kernels/quantization/test_cutlass_w4a8.py
@@ -300,8 +300,7 @@ def test_w4a8_cuda_graph():
 )
 @pytest.mark.parametrize("shape", MNK_SHAPES)
 def test_convert_packed_uint4b8_to_signed_int4_inplace(shape):
-    """
-    The W4A16 checkpoints encode the weights as int4b8 packed to int32.
+    """The W4A16 checkpoints encode the weights as int4b8 packed to int32.
     The CUTLASS kernels expect signed int4 packed to int32.
     This tests checks that the runtime int4b8 -> signed int4 conversion
     matches the offline conversion step exactly.
diff --git a/tests/kernels/quantization/test_cutlass_w4a8_moe.py b/tests/kernels/quantization/test_cutlass_w4a8_moe.py
index 5e6c170db644..10a76e553e25 100644
--- a/tests/kernels/quantization/test_cutlass_w4a8_moe.py
+++ b/tests/kernels/quantization/test_cutlass_w4a8_moe.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for the CUTLASS-based W4A8 grouped GEMM kernel and the full MoE layer.
-"""
+"""Tests for the CUTLASS-based W4A8 grouped GEMM kernel and the full MoE layer."""
 
 import random
 from dataclasses import dataclass
@@ -37,8 +35,7 @@ def cutlass_quantize(
     group_size: int | None,
     zero_points: bool = False,
 ):
-    """
-    Quantize weights into W4 and compute reference dequantized weights.
+    """Quantize weights into W4 and compute reference dequantized weights.
 
     Encoding/reordering of weights and packing of scales is deferred
     until after all experts are combined.
@@ -66,13 +63,13 @@ def cutlass_quantize(
 def cutlass_preprocess(
     w_q_experts: list[torch.Tensor], w_s_experts: list[torch.Tensor]
 ):
-    """
-    Reorder/encode expert weights and pack scales.
+    """Reorder/encode expert weights and pack scales.
 
     Returns:
         w_q_packed: Packed/encoded int4 weights for all experts.
         w_s_packed: Packed fp8 scales for all experts.
         packed_layout: Layout/stride metadata for grouped GEMM.
+
     """
     w_s_packed = ops.cutlass_pack_scale_fp8(torch.stack(w_s_experts))
     w_q_packed, packed_layout = ops.cutlass_encode_and_reorder_int4b_grouped(
@@ -130,7 +127,6 @@ def make_moe_test_setup(
     random_zero: bool = False,
 ) -> MoETestSetup:
     """Create a full set of tensors for testing cutlass_w4a8_moe_mm."""
-
     assert K % GROUP_SIZE == 0
     # Token counts per expert (multiples of `alignment`).
     Ms = [alignment * random.randint(1, max_blocks) for _ in range(num_experts)]
@@ -276,9 +272,7 @@ def test_cutlass_w4a8_moe_mm_end_to_end(shape, random_zero):
 
 
 class W4A8MoELayer(torch.nn.Module):
-    """
-    Minimal wrapper module to test cuda graphs
-    """
+    """Minimal wrapper module to test cuda graphs."""
 
     def __init__(self, setup: MoETestSetup):
         super().__init__()
diff --git a/tests/kernels/quantization/test_fp8_min_max_helper.py b/tests/kernels/quantization/test_fp8_min_max_helper.py
index 8cd68a3fef7e..4f1beb857674 100644
--- a/tests/kernels/quantization/test_fp8_min_max_helper.py
+++ b/tests/kernels/quantization/test_fp8_min_max_helper.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for the get_fp8_min_max() helper function.
+"""Unit tests for the get_fp8_min_max() helper function.
 
 These tests verify the FP8 min/max value logic for both standard
 and fnuz (ROCm MI300) dtype handling.
diff --git a/tests/kernels/quantization/test_int8_kernel.py b/tests/kernels/quantization/test_int8_kernel.py
index 704bb90f84a8..4858ea47adbf 100644
--- a/tests/kernels/quantization/test_int8_kernel.py
+++ b/tests/kernels/quantization/test_int8_kernel.py
@@ -54,7 +54,6 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
 def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids):
     """This function performs fused moe with per-column int8 quantization
     using native torch."""
-
     B, D = a.shape
     # Perform per-token quantization
     a_q, a_s = per_token_quant_int8(a)
diff --git a/tests/kernels/quantization/test_marlin_gemm.py b/tests/kernels/quantization/test_marlin_gemm.py
index 4c290dada6b5..29910390bb9a 100644
--- a/tests/kernels/quantization/test_marlin_gemm.py
+++ b/tests/kernels/quantization/test_marlin_gemm.py
@@ -17,6 +17,7 @@
     per_token_quant_int8,
 )
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    get_marlin_workspace,
     marlin_make_empty,
     marlin_make_workspace_new,
     marlin_permute_bias,
@@ -43,6 +44,7 @@
 )
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
+from vllm.v1.worker.workspace import current_workspace_manager
 
 if current_platform.is_rocm():
     pytest.skip(
@@ -477,6 +479,62 @@ def test_marlin_gemm(
     assert max_diff < 0.04
 
 
+def test_marlin_persistent_locks_cuda_graph_streams(workspace_init):
+    """Compiled GEMMs select stream-local locks at runtime and survive replay."""
+    size_m, size_k, size_n = 32, 1024, 2048
+    inputs = [rand_data((size_m, size_k)) for _ in range(2)]
+    weight = rand_data((size_k, size_n))
+    quant_type = scalar_types.uint4b8
+    w_ref, qweight, scales = marlin_quantize(weight, quant_type, 128)
+    zeros = marlin_make_empty(scales.device)
+    streams = [torch.cuda.Stream() for _ in inputs]
+    graphs = [torch.cuda.CUDAGraph() for _ in inputs]
+    locks = []
+    outputs = []
+
+    @torch.compile(backend="eager", fullgraph=True)
+    def run(x):
+        return torch.ops.vllm.marlin_gemm(
+            x,
+            qweight,
+            None,
+            scales,
+            None,
+            None,
+            zeros,
+            None,
+            quant_type.id,
+            size_m,
+            size_n,
+            size_k,
+            use_atomic_add=False,
+            use_fp32_reduce=True,
+        )
+
+    for stream, graph, x in zip(streams, graphs, inputs):
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            locks.append(get_marlin_workspace(x.device))
+            run(x)
+            with torch.cuda.graph(graph, stream=stream):
+                outputs.append(run(x))
+        torch.cuda.current_stream().wait_stream(stream)
+
+    assert locks[0].data_ptr() != locks[1].data_ptr()
+    current_workspace_manager().lock()
+    for _ in range(3):
+        for stream, graph, x, lock in zip(streams, graphs, inputs, locks):
+            with torch.cuda.stream(stream):
+                assert get_marlin_workspace(x.device) is lock
+                x.add_(0.1)
+                graph.replay()
+        for stream in streams:
+            torch.cuda.current_stream().wait_stream(stream)
+        for x, output, lock in zip(inputs, outputs, locks):
+            assert compute_max_diff(output, x @ w_ref) < 0.04
+            assert torch.count_nonzero(lock) == 0
+
+
 def test_marlin_gemm_subset_input():
     quant_type = scalar_types.uint4b8
     group_size = 128
diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py
index 027ffbf10c96..d224b0a2b257 100644
--- a/tests/kernels/quantization/test_marlin_tile_padding.py
+++ b/tests/kernels/quantization/test_marlin_tile_padding.py
@@ -14,6 +14,7 @@
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     GPTQ_MARLIN_TILE,
     apply_gptq_marlin_linear,
+    get_marlin_workspace,
     marlin_make_empty,
     marlin_make_workspace_new,
     marlin_moe_padded_intermediate,
@@ -225,7 +226,7 @@ def test_fp8_marlin_padded_round_trip(shape, use_bias):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
         bias=layer.bias if use_bias else None,
@@ -288,7 +289,7 @@ def test_nvfp4_marlin_padded_round_trip(shape):
         weight=layer.weight,
         weight_scale=layer.weight_scale,
         weight_global_scale=layer.weight_global_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
     )
@@ -392,7 +393,7 @@ def test_fp8_block_marlin_padded_round_trip(shape):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale_inv,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
         bias=None,
@@ -440,7 +441,7 @@ def test_mxfp8_marlin_padded_round_trip(shape):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
     )
@@ -756,7 +757,7 @@ def q(w):  # (out, in) -> fp8 weight, scale, dequant reference
         topk_ids,
         quant_type_id=scalar_types.float8_e4m3fn.id,
         global_num_experts=e,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(device),
     )
     with set_current_vllm_config(VllmConfig()):
         ref = torch_experts(
@@ -833,7 +834,7 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape):
         topk_ids,
         quant_type_id=scalar_types.float8_e4m3fn.id,
         global_num_experts=e,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(device),
     )
     with set_current_vllm_config(VllmConfig()):
         ref = torch_experts(
diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py
index 28ce615ce655..46ff139fc76d 100644
--- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py
+++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py
@@ -1,20 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for MXFP4 linear kernel selection logic (CPU-only)
+"""Tests for MXFP4 linear kernel selection logic (CPU-only).
 
 Run `pytest tests/kernels/quantization/test_mxfp4_kernel_selection.py`.
 """
 
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 
+import vllm.envs as envs
 from vllm.model_executor.kernels.linear import (
     AiterMxfp4LinearKernel,
     EmulationMxfp4LinearKernel,
     FlashInferMxFp4LinearKernel,
-    HummingMxFp4LinearKernel,
     MarlinMxFp4LinearKernel,
     MxFp4LinearKernel,
     MxFp4LinearLayerConfig,
@@ -44,7 +44,20 @@
 
 # Weight-only (A16) kernels: they never quantize activations. They still accept
 # MXFP4 activation keys as an intentional compatibility fallback.
-_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel]
+_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel]
+
+
+def _make_emulation_layer():
+    layer = torch.nn.Module()
+    layer.register_parameter(
+        "weight",
+        torch.nn.Parameter(torch.ones((2, 2), dtype=torch.uint8), False),
+    )
+    layer.register_parameter(
+        "weight_scale",
+        torch.nn.Parameter(torch.ones((2, 1), dtype=torch.uint8), False),
+    )
+    return layer
 
 
 def test_can_implement_is_abstract():
@@ -156,6 +169,77 @@ def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch):
         assert kernel.quant_dequant_func is quant_dequant_mxfp4
 
 
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_emulation_kernel_dequantizes_at_load_and_keeps_activation_qdq(
+    dtype, monkeypatch
+):
+    monkeypatch.setattr(
+        "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark",
+        lambda: True,
+    )
+    kernel = EmulationMxfp4LinearKernel(
+        MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic)
+    )
+    kernel.quant_dequant_func = MagicMock(side_effect=lambda x: x * 0.5)
+    layer = _make_emulation_layer()
+    dequantized = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.bfloat16)
+    x = torch.tensor([[2.0, 4.0]], dtype=dtype)
+
+    with (
+        patch.object(envs, "VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD", True),
+        patch(
+            "vllm.model_executor.kernels.linear.mxfp4.emulation.dequant_mxfp4",
+            return_value=dequantized,
+        ) as dequant,
+    ):
+        kernel.process_weights_after_loading(layer)
+        actual = kernel.apply_weights(layer, x)
+
+    expected = torch.nn.functional.linear(x * 0.5, dequantized.to(dtype))
+    torch.testing.assert_close(actual, expected)
+    dequant.assert_called_once()
+    kernel.quant_dequant_func.assert_called_once_with(x)
+    assert layer.weight.dtype == torch.bfloat16
+    assert layer.weight.device == dequantized.device
+    assert not layer.weight.requires_grad
+    assert layer.weight_scale is not None
+    assert not layer.weight_scale.requires_grad
+
+
+def test_emulation_kernel_opt_out_dequantizes_per_invocation(monkeypatch):
+    monkeypatch.setattr(
+        "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark",
+        lambda: True,
+    )
+    kernel = EmulationMxfp4LinearKernel(
+        MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic)
+    )
+    kernel.quant_dequant_func = MagicMock(side_effect=lambda x: x * 0.5)
+    layer = _make_emulation_layer()
+    dequantized = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.bfloat16)
+    x = torch.tensor([[2.0, 4.0]], dtype=torch.bfloat16)
+
+    with (
+        patch.object(envs, "VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD", False),
+        patch(
+            "vllm.model_executor.kernels.linear.mxfp4.emulation.dequant_mxfp4",
+            return_value=dequantized,
+        ) as dequant,
+    ):
+        kernel.process_weights_after_loading(layer)
+        dequant.assert_not_called()
+        first = kernel.apply_weights(layer, x)
+        second = kernel.apply_weights(layer, x)
+
+    expected = torch.nn.functional.linear(x * 0.5, dequantized)
+    torch.testing.assert_close(first, expected)
+    torch.testing.assert_close(second, expected)
+    assert dequant.call_count == 2
+    assert kernel.quant_dequant_func.call_count == 2
+    assert layer.weight.dtype == torch.uint8
+    assert layer.weight_scale is not None
+
+
 def test_aiter_kernel_is_supported_requires_native_mx_support():
     """AiterMxfp4LinearKernel must not be selected on platforms without
     native MX compute, even if AITER itself is importable."""
diff --git a/tests/kernels/quantization/test_mxfp4_triton_ep.py b/tests/kernels/quantization/test_mxfp4_triton_ep.py
index 93d30f459816..1edaa86cac61 100644
--- a/tests/kernels/quantization/test_mxfp4_triton_ep.py
+++ b/tests/kernels/quantization/test_mxfp4_triton_ep.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests that triton_kernel_moe_forward correctly applies expert_map
+"""Tests that triton_kernel_moe_forward correctly applies expert_map
 remapping when expert parallelism (EP) is enabled.
 
 When expert_map is provided, global expert IDs are remapped to local IDs
diff --git a/tests/kernels/quantization/test_mxfp6_kernel_selection.py b/tests/kernels/quantization/test_mxfp6_kernel_selection.py
index 2c7b565bea57..5c5b1b3b9e50 100644
--- a/tests/kernels/quantization/test_mxfp6_kernel_selection.py
+++ b/tests/kernels/quantization/test_mxfp6_kernel_selection.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for MXFP6 linear kernel selection logic (CPU-only)
+"""Tests for MXFP6 linear kernel selection logic (CPU-only).
 
 Run `pytest tests/kernels/quantization/test_mxfp6_kernel_selection.py`.
 """
diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py
index eb057cebd912..3a136484acb3 100644
--- a/tests/kernels/quantization/test_nvfp4_emulation.py
+++ b/tests/kernels/quantization/test_nvfp4_emulation.py
@@ -4,9 +4,21 @@
 
 import pytest
 import torch
+from nvfp4_utils import (
+    FLOAT4_E2M1_MAX,
+    FLOAT8_E4M3_MAX,
+    break_fp4_bytes,
+    kE2M1ToFloat,
+)
 from safetensors import safe_open
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.config.kernel import KernelConfig
+from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel
+from vllm.model_executor.kernels.linear.nvfp4.emulation import (
+    EmulationNvFp4LinearKernel,
+)
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
@@ -34,6 +46,7 @@
 from vllm.platforms import current_platform
 from vllm.transformers_utils.repo_utils import hf_api
 from vllm.triton_utils import triton
+from vllm.utils.torch_utils import set_random_seed
 
 if current_platform.is_rocm():
     from vllm.platforms.rocm import on_gfx950
@@ -64,8 +77,7 @@ def loaded_model_files():
 
 
 class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts):
-    """
-    Extension of TritonExperts to support emulated NVFP4 MoE experts.
+    """Extension of TritonExperts to support emulated NVFP4 MoE experts.
 
     It may be used for NVFP4 models when the device does not have
     native support for this dtype.
@@ -770,3 +782,85 @@ def _make_quant_config():
         atol=0.0 if on_gfx950() else 0.02,
         rtol=0,
     )
+
+
+def _dense_nvfp4_reference(
+    layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
+) -> torch.Tensor:
+    """Decode E2M1 nibbles and both scale levels without production helpers."""
+    levels = kE2M1ToFloat.cpu()
+    weight = break_fp4_bytes(layer.weight.cpu(), torch.float32)
+    weight *= (
+        layer.weight_scale.cpu().float() * layer.weight_global_scale.cpu()
+    ).repeat_interleave(16, dim=-1)
+    blocks = x.cpu().float().reshape(x.shape[0], -1, 16)
+    global_scale = layer.input_global_scale_inv.cpu()
+    scales = (
+        (blocks.abs().amax(-1, keepdim=True) * (1.0 / FLOAT4_E2M1_MAX) * global_scale)
+        .clamp(max=FLOAT8_E4M3_MAX)
+        .to(torch.float8_e4m3fn)
+        .float()
+    )
+    inverse = torch.where(scales == 0, 0, global_scale / scales)
+    scaled = (blocks * inverse).abs().clamp(max=FLOAT4_E2M1_MAX)
+    distances = (scaled.unsqueeze(-1) - levels).abs()
+    # E2M1 ties round to the even significand (even encoding).
+    nearest_mask = distances == distances.amin(-1, keepdim=True)
+    even_mask = nearest_mask & (torch.arange(8, device="cpu") % 2 == 0)
+    candidates = torch.where(even_mask.any(-1, keepdim=True), even_mask, nearest_mask)
+    nearest = candidates.to(torch.int8).argmax(-1)
+    activation = levels[nearest] * blocks.sign() * (scales / global_scale)
+    expected = activation.reshape_as(x).to(x.dtype) @ weight.to(x.dtype).T
+    if bias is not None:
+        expected += bias.cpu()
+    return expected
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm NVFP4 fallback")
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("m,n,k", [(1, 32, 64), (17, 128, 256)])
+@pytest.mark.parametrize("with_bias", [False, True])
+def test_nvfp4_dense_emulation(
+    dtype: torch.dtype, m: int, n: int, k: int, with_bias: bool
+) -> None:
+    """W4A4 fallback preserves packed scales, bias and outputs across replay."""
+    set_random_seed(47)
+    layer = torch.nn.Module()
+    # Every signed E2M1 code appears, including both zero encodings.
+    codes = (torch.randperm(n * k, device="cuda") % 16).reshape(n, k)
+    layer.weight = (codes[:, ::2] | (codes[:, 1::2] << 4)).to(torch.uint8)
+    layer.weight_scale = (
+        torch.arange(n * k // 16, device="cuda").reshape(n, k // 16) % 7 / 8
+    ).to(torch.float8_e4m3fn)
+    layer.weight_global_scale = torch.tensor(0.25, device="cuda")
+    layer.input_global_scale_inv = torch.tensor(3.0, device="cuda")
+    x = torch.randn(m, k, device="cuda", dtype=dtype)
+    x[:, :16] = 0
+    bias = torch.randn(n, device="cuda", dtype=dtype) if with_bias else None
+    with set_current_vllm_config(
+        VllmConfig(kernel_config=KernelConfig(linear_backend="emulation"))
+    ):
+        kernel = init_nvfp4_linear_kernel()
+    assert isinstance(kernel, EmulationNvFp4LinearKernel)
+    kernel.process_weights_after_loading(layer)
+    original_weight = layer.weight.clone()
+    original_scale = layer.weight_scale.view(torch.uint8).clone()
+    actual = kernel.apply_weights(layer, x, bias)
+    torch.testing.assert_close(
+        actual.cpu(), _dense_nvfp4_reference(layer, x, bias), atol=0.03, rtol=0.01
+    )
+    assert actual.abs().max() > 0.1
+    saved = actual.clone()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = kernel.apply_weights(layer, x, bias)
+    for sign in (-1, 1):
+        x.mul_(sign * 0.75)
+        graph.replay()
+        torch.accelerator.synchronize()
+        torch.testing.assert_close(
+            captured.cpu(), _dense_nvfp4_reference(layer, x, bias), atol=0.03, rtol=0.01
+        )
+        torch.testing.assert_close(actual, saved, atol=0, rtol=0)
+    assert torch.equal(layer.weight, original_weight)
+    assert torch.equal(layer.weight_scale.view(torch.uint8), original_scale)
diff --git a/tests/kernels/quantization/test_nvfp4_quant.py b/tests/kernels/quantization/test_nvfp4_quant.py
index 120dc2774a5d..b8a1cdda4612 100644
--- a/tests/kernels/quantization/test_nvfp4_quant.py
+++ b/tests/kernels/quantization/test_nvfp4_quant.py
@@ -176,8 +176,7 @@ def test_python_util_matches_cpp_allocation(
     shape: tuple[int, int],
     is_sf_swizzled_layout: bool,
 ) -> None:
-    """
-    Verify that the Python utility (create_fp4_output_tensors) allocates
+    """Verify that the Python utility (create_fp4_output_tensors) allocates
     tensors with the same shapes and dtypes as the C++ functional variant
     (scaled_fp4_quant_func).
     """
diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py
index a2788d058304..6d3c91022ed8 100644
--- a/tests/kernels/quantization/test_per_token_group_quant.py
+++ b/tests/kernels/quantization/test_per_token_group_quant.py
@@ -11,8 +11,24 @@
 
 
 @pytest.mark.parametrize(
-    "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)]
+    "shape",
+    [
+        (0, 128),
+        (1, 128),
+        (3, 384),
+        (5, 384),
+        (7, 384),
+        (15, 384),
+        (31, 128),
+        (32, 128),
+        (63, 256),
+        (64, 256),
+        (16, 512),
+        (3, 2688),
+        (3, 5376),
+    ],
 )
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
 @pytest.mark.parametrize("column_major", [False, True])
 @pytest.mark.parametrize("tma_aligned", [False, True])
 @pytest.mark.parametrize("scale_ue8m0", [False, True])
@@ -22,14 +38,20 @@
     reason="Only test on CUDA/ROCm/XPU.",
 )
 def test_per_token_group_quant_fp8(
-    shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int
+    shape,
+    dtype,
+    column_major: bool,
+    tma_aligned: bool,
+    scale_ue8m0: bool,
+    group_size: int,
 ):
+    """Tail groups preserve FP8 values and scales in every supported layout."""
     device = current_platform.device_type
 
     torch.manual_seed(42)
     num_tokens, hidden_dim = shape
 
-    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
+    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8
 
     # native kernel path
     out_q, scale = fp8_utils.per_token_group_quant_fp8(
@@ -97,7 +119,6 @@ def test_per_token_group_quant_fp8_packed(
 ):
     """Test the packed DeepGEMM quantization kernel against the Triton
     reference (row-major, UE8M0 scales)."""
-
     device = "cuda"
     torch.manual_seed(42)
 
@@ -184,7 +205,6 @@ def test_per_token_group_quant_fp8_packed_all_zero():
     For all-zero input, eps/fp8_max < 1e-10, so the inner fmax clamps back to
     1e-10, giving exp2(ceil(log2(1e-10))) = exp2(-33) => UE8M0 byte 0x5E (94).
     """
-
     device = "cuda"
     num_tokens, hidden_dim, group_size = 4, 7168, 128
     x = torch.zeros((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16)
@@ -237,7 +257,6 @@ def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():
     """Inputs whose absmax/max_8bit produces a non-power-of-2 force the
     mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior
     before optimization."""
-
     device = "cuda"
     num_tokens, hidden_dim, group_size = 4, 7168, 128
 
@@ -310,7 +329,6 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
     """When output_q is allocated with shape (tma_aligned_mn, k) instead of
     (mn, k), the kernel must overwrite the padded mn rows with zeros so
     callers can use ``torch.empty`` instead of ``torch.zeros``."""
-
     device = "cuda"
     group_size = 128
     torch.manual_seed(42)
@@ -374,7 +392,6 @@ def test_per_token_group_quant_fp8_packed_large_mn():
     This is a differential test that compares fp8 output against Triton output
     reference when token size sits just above the gridDim.y 2^16 - 1 limit.
     """
-
     device = "cuda"
     group_size = 128
     # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
@@ -411,16 +428,34 @@ def test_per_token_group_quant_fp8_packed_large_mn():
     assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
 
 
-@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (0, 128),
+        (1, 128),
+        (3, 384),
+        (5, 384),
+        (7, 384),
+        (15, 384),
+        (32, 128),
+        (64, 256),
+        (16, 512),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
 @pytest.mark.parametrize("group_size", [64, 128])
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_per_token_group_quant_int8(shape, group_size: int):
-    device = "cuda"
+@pytest.mark.skipif(
+    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
+    reason="Only test on CUDA/ROCm/XPU.",
+)
+def test_per_token_group_quant_int8(shape, dtype, group_size: int):
+    """The shared launch policy also handles INT8 tail groups."""
+    device = current_platform.device_type
 
     torch.manual_seed(42)
     num_tokens, hidden_dim = shape
 
-    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
+    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8
 
     # cuda path
     out_q, scale = int8_utils.per_token_group_quant_int8(
diff --git a/tests/kernels/quantization/test_rdna3_compile_guards.py b/tests/kernels/quantization/test_rdna3_compile_guards.py
index 3c1e226d2c07..194441f0f909 100644
--- a/tests/kernels/quantization/test_rdna3_compile_guards.py
+++ b/tests/kernels/quantization/test_rdna3_compile_guards.py
@@ -34,6 +34,13 @@
 if not current_platform.is_rocm():
     pytest.skip("RDNA3 compile-guard tests are ROCm-only", allow_module_level=True)
 
+from vllm.model_executor.layers.quantization.utils.quant_utils import (  # noqa: E402
+    kInt4Static,
+    kInt4Static32,
+    kInt4Static32Asym,
+    kInt4StaticAsym,
+    kInt8Static,
+)
 from vllm.platforms.rocm import on_gfx1100  # noqa: E402
 
 gfx1100_only = pytest.mark.skipif(
@@ -166,15 +173,13 @@ def test_op_absent_on_non_gfx1100(op_name):
 
 @not_gfx1100
 def test_rocm_moe_not_supported_on_non_gfx1100():
-    """rocm_moe_rdna.is_supported() must return False on non-gfx1100 hardware."""
-    from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa: E501
-        rocm_moe_rdna,
+    """The RDNA3 MoE experts must not be selectable on non-gfx1100 hardware."""
+    from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+        Rdna3WNA16Experts,
     )
 
-    wq = type("WQ", (), {"num_bits": 4})()
-    assert rocm_moe_rdna.is_supported(wq) is False, (
-        "rocm_moe_rdna.is_supported() returned True on non-gfx1100 — "
-        "dispatch guard is broken"
+    assert Rdna3WNA16Experts._supports_current_device() is False, (
+        "Rdna3WNA16Experts reported support on non-gfx1100 — dispatch guard is broken"
     )
 
 
@@ -359,54 +364,58 @@ def test_no_toplevel_rocm_c_import(self):
 # ============================================================================
 
 
-class _FakeWeightQuant:
-    """Minimal stand-in for a weight quantization config."""
-
-    def __init__(self, num_bits):
-        self.num_bits = num_bits
-
-
 class TestMoEDispatchMocked:
     """Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable."""
 
-    def test_is_supported_false_when_mocked_cdna(self):
-        """rocm_moe_rdna.is_supported() must return False when not on gfx1100."""
-        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa: E501
-            rocm_moe_rdna,
+    def test_kernel_unavailable_when_mocked_cdna(self):
+        """The device gate must reject when not on gfx1100."""
+        from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+            rdna3_moe_kernel_available,
         )
 
         with patch("vllm.platforms.rocm.on_gfx1100", return_value=False):
-            assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
+            assert rdna3_moe_kernel_available() is False
 
-    @pytest.mark.parametrize("num_bits", [2, 3, 8, 16])
-    def test_is_supported_rejects_non_w4(self, num_bits):
-        """is_supported() rejects non-4-bit even before checking arch."""
-        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa: E501
-            rocm_moe_rdna,
+    @pytest.mark.parametrize(
+        "weight_key",
+        [kInt8Static, kInt4StaticAsym, kInt4Static32Asym, None],
+    )
+    def test_quant_scheme_rejects_non_symmetric_int4(self, weight_key):
+        """Only symmetric int4 weight-only schemes reach the kernel."""
+        from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+            Rdna3WNA16Experts,
         )
 
-        assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False
+        assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is False
 
-    def test_is_supported_false_when_op_missing(self):
-        """is_supported() returns False when the C++ op doesn't exist."""
-        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa: E501
-            rocm_moe_rdna,
+    @pytest.mark.parametrize("weight_key", [kInt4Static, kInt4Static32])
+    def test_quant_scheme_accepts_symmetric_int4(self, weight_key):
+        from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+            Rdna3WNA16Experts,
+        )
+
+        assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is True
+
+    def test_kernel_unavailable_when_op_missing(self):
+        """The device gate returns False when the C++ op doesn't exist."""
+        from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+            rdna3_moe_kernel_available,
         )
 
         fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})()
         with patch.object(torch, "ops", create=True) as mock_ops:
             mock_ops._rocm_C = fake_rocm_c
-            assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
+            assert rdna3_moe_kernel_available() is False
 
-    def test_is_supported_false_when_rocm_c_absent(self):
-        """is_supported() returns False when _rocm_C doesn't exist at all."""
-        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa: E501
-            rocm_moe_rdna,
+    def test_kernel_unavailable_when_rocm_c_absent(self):
+        """The device gate returns False when _rocm_C doesn't exist at all."""
+        from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+            rdna3_moe_kernel_available,
         )
 
         fake_ops = type("FakeOps", (), {})()
         with patch.object(torch, "ops", fake_ops):
-            assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
+            assert rdna3_moe_kernel_available() is False
 
 
 class TestDenseKernelSelectionMocked:
@@ -471,30 +480,30 @@ def test_chooser_skips_rdna3_when_mocked_cdna(self):
             )
 
 
-class TestCompressedTensorsMoEDispatchGuard:
-    """Verify compressed_tensors_moe.py only enters rocm_moe_rdna under is_rocm()."""
+class TestWNA16OracleWiring:
+    """The RDNA3 backend must reach the kernel only through the oracle."""
 
-    def test_rocm_guard_in_dispatch_source(self):
-        """The rocm_moe_rdna import and call must be inside an is_rocm() check."""
-        src = _read_pkg_source_or_skip(
-            "model_executor",
-            "layers",
-            "quantization",
-            "compressed_tensors",
-            "compressed_tensors_moe",
-            "compressed_tensors_moe.py",
+    def test_backend_maps_to_rdna3_experts(self):
+        from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+            Rdna3WNA16Experts,
+        )
+        from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
+            WNA16MoEBackend,
+            backend_to_kernel_cls,
+            map_wna16_backend,
         )
-        lines = src.splitlines()
 
-        for i, line in enumerate(lines, 1):
-            stripped = line.strip()
-            if "rocm_moe" in stripped and not stripped.startswith("#"):
-                found_guard = False
-                for j in range(i - 1, max(0, i - 15), -1):
-                    if "is_rocm()" in lines[j - 1]:
-                        found_guard = True
-                        break
-                assert found_guard, (
-                    f"L{i}: rocm_moe_rdna reference not protected by "
-                    f"is_rocm() guard: {stripped}"
-                )
+        assert backend_to_kernel_cls(WNA16MoEBackend.RDNA3) == [Rdna3WNA16Experts]
+        assert map_wna16_backend("rdna3") == WNA16MoEBackend.RDNA3
+
+    def test_backend_is_offered_before_the_triton_fallback(self):
+        """Priority order: the native kernel outranks Triton when supported."""
+        from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
+            WNA16MoEBackend,
+            _get_priority_backends,
+        )
+
+        backends = _get_priority_backends()
+        assert backends.index(WNA16MoEBackend.RDNA3) < backends.index(
+            WNA16MoEBackend.TRITON
+        )
diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
index 562d1ef2ef71..fe3db8dab360 100644
--- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
+++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py
@@ -37,6 +37,20 @@
 # ---------------------------------------------------------------------------
 
 
+def _pack_zp_rows_for_kernel(zp_nkg: torch.Tensor) -> torch.Tensor:
+    """Pack raw uint4 zero points along N: [N, G] int32 -> [N//8, G] int32.
+
+    Row n's nibble lands in word[n//8] at bits 4*(n%8).
+    """
+    assert zp_nkg.dtype == torch.int32
+    N, G = zp_nkg.shape
+    assert N % 8 == 0
+    shifts = (torch.arange(8, device=zp_nkg.device, dtype=torch.int32) * 4)[:, None]
+    return torch.sum(
+        (zp_nkg.view(N // 8, 8, G) & 0xF) << shifts, dim=1, dtype=torch.int32
+    ).contiguous()
+
+
 def _rdna_hybrid_w4a16_reference(
     x_mk: torch.Tensor,
     w_int4_nk: torch.Tensor,
@@ -50,8 +64,8 @@ def _rdna_hybrid_w4a16_reference(
     x_mk: [M, K] fp16/bf16
     w_int4_nk: [N, K] int32 with raw uint4 values in [0, 15]
     scales_nkg: [N, K//G] fp16/bf16
-    zp_nkg: [N, K//G] fp16/bf16 raw zero points (already in act dtype),
-            or None for symmetric (uint4b8, dequant subtracts 8)
+    zp_nkg: [N, K//G] int32 raw zero points in [0, 15], or None for
+            symmetric (uint4b8, dequant subtracts 8)
     """
     G = group_size
     N, K = w_int4_nk.shape
@@ -110,13 +124,16 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M)
         0.05 * torch.rand((N, K // group_size), device=device, dtype=torch.float32)
     ).to(dtype)
 
-    # Optional raw zero points [N, K//G] in act dtype.
+    # Optional raw zero points, [N, K//G] int32 for the reference and packed
+    # along N for the op.
     if has_zp:
         zp_nkg = torch.randint(
             0, 16, (N, K // group_size), device=device, dtype=torch.int32
-        ).to(dtype)
+        )
+        w_zp = _pack_zp_rows_for_kernel(zp_nkg)
     else:
         zp_nkg = None
+        w_zp = None
 
     from vllm.utils.platform_utils import num_compute_units
 
@@ -124,7 +141,7 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M)
         x_mk,
         w_q,
         scales_nkg,
-        zp_nkg,
+        w_zp,
         None,  # bias
         num_compute_units(),
         group_size,
@@ -332,7 +349,7 @@ def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_ini
 
 @pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)
 def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_init):
-    """uint4 (asymmetric): zero points unpacked to raw values in act dtype."""
+    """uint4 (asymmetric): zero points are packed [N//8, K//G] int32."""
     if not torch.cuda.is_available():
         pytest.skip("CUDA/HIP device not available")
 
@@ -374,10 +391,11 @@ def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_in
     )
     kernel.process_weights_after_loading(layer)
 
-    # Zero-points: unpacked to [N, K//G], cast to act dtype, raw values [0..15].
-    assert layer.weight_zero_point.dtype == torch.float16
-    assert tuple(layer.weight_zero_point.shape) == (N, K // G)
-    expected_zp = zeros_int4_gn.t().to(torch.float16)  # [N, K//G]
+    # Zero-points: [N//8, K//G] int32, row n's nibble at word[n//8]
+    # bits 4*(n%8).
+    assert layer.weight_zero_point.dtype == torch.int32
+    assert tuple(layer.weight_zero_point.shape) == (N // 8, K // G)
+    expected_zp = _pack_zp_rows_for_kernel(zeros_int4_gn.t().contiguous())
     torch.testing.assert_close(layer.weight_zero_point, expected_zp)
 
     # Quantized weights match symmetric path's layout regardless of zp.
@@ -518,3 +536,30 @@ def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G):
     ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8)
 
     torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2)
+
+
+@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
+@pytest.mark.skipif(
+    not hasattr(torch.ops, "_rocm_C")
+    or not hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"),
+    reason="wvSplitK_int4_g not built",
+)
+def test_wvsplitk_int4_g_rejects_unpacked_zero_points():
+    """Act-dtype zero points raise instead of being read as packed words.
+
+    Both layouts are 2D with compatible extents, so only the dtype separates
+    them; misreading one as the other returns wrong numbers silently.
+    """
+    import vllm._custom_ops as ops
+    from vllm.utils.platform_utils import num_compute_units
+
+    K, N, G, M = 256, 64, 128, 1
+    a = torch.randn((M, K), device=device, dtype=torch.float16)
+    w = torch.randint(0, 255, (N, K // 2), device=device, dtype=torch.uint8).view(
+        torch.int8
+    )
+    scales = torch.rand((N, K // G), device=device, dtype=torch.float16)
+    zp_unpacked = torch.zeros((N, K // G), device=device, dtype=torch.float16)
+
+    with pytest.raises(RuntimeError, match="Zero points must be int32 or uint32"):
+        ops.wvSplitK_int4_g(w, a, scales, num_compute_units(), G, zp_unpacked, None)
diff --git a/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py b/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py
index a9b35a4ea642..8e4563ab8004 100644
--- a/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py
+++ b/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py
@@ -17,7 +17,6 @@
 @pytest.mark.parametrize(
     "model_path",
     [
-        # Listed in tests/weight_loading/models.txt
         "nm-testing/tinyllama-oneshot-w4a16-group128-v2",
     ],
 )
diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py
index 43b5daa45a7e..d4e16629afac 100644
--- a/tests/kernels/quantization/test_rocm_mxfp4.py
+++ b/tests/kernels/quantization/test_rocm_mxfp4.py
@@ -50,16 +50,7 @@
 }
 PRESHUFFLED_SHAPES = [
     (64, 4096, 8192),
-    # aiter 0.1.20 Triton preshuffled fp4 GEMM OOBs (GPU memory access fault ->
-    # SIGABRT, crashing the process) for (M=32, N=8192, K=8192) on gfx950, despite
-    # advertising the (N, K) as tuned. A hard GPU fault can't be xfail'd, so skip it
-    # until the aiter kernel is fixed (ROCm/aiter#4867).
-    pytest.param(
-        (32, 8192, 8192),
-        marks=pytest.mark.skip(
-            reason="aiter 0.1.20 preshuffled fp4 GEMM OOB (M32,N8192,K8192, gfx950)"
-        ),
-    ),
+    (32, 8192, 16384),
 ]
 
 
@@ -156,22 +147,26 @@ def test_fp4_env_defaults():
     """ROCm FP4 env defaults should stay stable for the AITER gates."""
     import vllm.envs as envs
 
+    assert envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM is False
     assert envs.VLLM_ROCM_USE_AITER_FP4BMM is True
 
 
 @pytest.mark.parametrize(
     (
         "use_aiter",
+        "use_fp4_asm_gemm",
         "use_fp4bmm",
     ),
     [
-        (True, True),
-        (True, False),
-        (False, True),
+        (True, True, True),
+        (True, True, False),
+        (True, False, True),
+        (False, True, True),
     ],
 )
 def test_rocm_aiter_fp4_enablement_follows_env_and_arch(
     use_aiter,
+    use_fp4_asm_gemm,
     use_fp4bmm,
     monkeypatch,
 ):
@@ -187,11 +182,15 @@ def test_rocm_aiter_fp4_enablement_follows_env_and_arch(
     _assert_aiter_supported()
 
     on_gfx950_value = on_gfx950()
-    expected_asm_gemm = use_aiter and on_gfx950_value
+    expected_asm_gemm = use_aiter and use_fp4_asm_gemm and on_gfx950_value
     expected_fp4bmm = use_aiter and use_fp4bmm and on_gfx950_value
 
     with monkeypatch.context() as mp:
         mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
+        mp.setenv(
+            "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM",
+            "1" if use_fp4_asm_gemm else "0",
+        )
         mp.setenv("VLLM_ROCM_USE_AITER_FP4BMM", "1" if use_fp4bmm else "0")
         _reload_envs()
         rocm_aiter_ops.refresh_env_variables()
@@ -422,7 +421,8 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape):
     M, K, N = shape
 
     assert M <= 64
-    assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K)
+    # The predicate takes K as packed bytes, matching `weight.shape[1]`.
+    assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K // 2)
 
     A = torch.randn(M, K, dtype=torch.bfloat16)
     B = torch.randn(N, K, dtype=torch.bfloat16)
@@ -504,7 +504,7 @@ def run_gemm():
     ],
 )
 def test_aiter_hardware_fp4_dynamic_quant_format(shape):
-    """aiter hardware FP4 dynamic quant produces correct output format.
+    """Aiter hardware FP4 dynamic quant produces correct output format.
 
     Tests gfx950 hardware-accelerated FP4 quantization (OCP MXFP4 E2M1).
     Parity with B200 scaled_fp4_quant: block_size=32, packed uint8 output.
diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py
new file mode 100644
index 000000000000..1d5e62c6c4b1
--- /dev/null
+++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Native CDNA4 MXFP8 linear (``RocmDotScaledMxfp8LinearKernel``).
+
+DeepSeek V4.1 routes every attention linear through this kernel. It used to
+carry two BF16-dequantize workarounds on ROCm (one for all ``.attn.`` linears,
+one for the row-parallel ``wo_b``, whose comment blamed NaNs at small token
+counts). These tests pin the two properties those workarounds stood in for --
+finite output and MXFP8-level accuracy -- across the V4.1 attention shapes at
+TP=4 and the small token counts decode actually runs.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    _mxfp8_e4m3_quantize_torch,
+    _mxfp8_e4m3_quantize_triton,
+    dequant_mxfp8_to_bf16,
+    mxfp8_e4m3_quantize,
+)
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm() or not current_platform.supports_mx(),
+    reason="native MX linear requires CDNA4 (gfx95x)",
+)
+
+# (name, N, K) for the DeepSeek V4.1 attention linears at tensor_parallel=4:
+# fused wq_a|wkv is replicated, wq_b/wo_a are column-parallel, wo_b is
+# row-parallel (so K, not N, is sharded).
+V41_ATTN_SHAPES = [
+    ("fused_wqa_wkv", 1280 + 512, 5120),
+    ("wq_b", 64 * 512 // 4, 1280),
+    ("wo_a", 8 * 1024 // 4, 4096),
+    ("wo_b", 5120, 8192 // 4),
+]
+
+
+def _make_layer(n: int, k: int, device: str) -> torch.nn.Module:
+    weight_bf16 = torch.randn(n, k, device=device, dtype=torch.bfloat16) * 0.05
+    weight, weight_scale = _mxfp8_e4m3_quantize_torch(weight_bf16)
+    layer = torch.nn.Module()
+    layer.weight = torch.nn.Parameter(weight, requires_grad=False)
+    layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+    return layer
+
+
+@pytest.mark.parametrize("name,n,k", V41_ATTN_SHAPES, ids=lambda v: str(v))
+@pytest.mark.parametrize("num_tokens", [1, 2, 7, 16, 31, 32, 128, 1024])
+@torch.inference_mode()
+def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens):
+    """Output is finite and matches an exact MXFP8 GEMM of the same inputs.
+
+    The reference quantizes the activation with the same quantizer and then
+    contracts in FP32, so the tolerance covers only kernel/accumulation error,
+    not the MXFP8 quantization error both sides share.
+    """
+    from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+        Mxfp8LinearLayerConfig,
+    )
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        RocmDotScaledMxfp8LinearKernel,
+    )
+
+    torch.manual_seed(num_tokens)
+    device = "cuda"
+    layer = _make_layer(n, k, device)
+    kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
+    kernel.process_weights_after_loading(layer)
+
+    x = torch.randn(num_tokens, k, device=device, dtype=torch.bfloat16) * 0.5
+    out = kernel.apply_weights(layer, x)
+
+    assert out.shape == (num_tokens, n)
+    assert torch.isfinite(out).all(), f"{name}: non-finite output"
+
+    x_q, x_scale = mxfp8_e4m3_quantize(x)
+    expected = (
+        dequant_mxfp8_to_bf16(x_q, x_scale).float()
+        @ dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale).float().T
+    )
+    rel = (out.float() - expected).norm() / expected.norm()
+    assert rel < 5e-3, f"{name}: relative error {rel:.4f}"
+
+
+@torch.inference_mode()
+def test_rocm_mxfp8_linear_handles_extreme_magnitudes():
+    """Rows that are all-zero or near the E4M3 range must not produce NaNs.
+
+    A degenerate block makes the dynamic activation quantizer pick an extreme
+    E8M0 scale; a zero row makes it pick the tiny-clamped one. Both feed
+    ``tl.dot_scaled`` directly.
+    """
+    from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+        Mxfp8LinearLayerConfig,
+    )
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        RocmDotScaledMxfp8LinearKernel,
+    )
+
+    torch.manual_seed(0)
+    device = "cuda"
+    n, k = 5120, 2048  # wo_b at TP=4
+    layer = _make_layer(n, k, device)
+    kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
+    kernel.process_weights_after_loading(layer)
+
+    x = torch.randn(8, k, device=device, dtype=torch.bfloat16)
+    x[0] = 0.0
+    x[1] = 1e-30
+    x[2] = 400.0
+    x[3, ::2] = 0.0
+
+    out = kernel.apply_weights(layer, x)
+    assert torch.isfinite(out).all()
+    assert (out[0] == 0).all(), "an all-zero activation row must give a zero row"
+
+
+@pytest.mark.parametrize("num_rows", [1, 64, 65])
+@torch.inference_mode()
+def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows):
+    """The fused ROCm quantizer must agree with the torch one, zeros included.
+
+    An all-zero 32-element block clamps the E8M0 scale to 0, i.e. 2**-127 --
+    subnormal in fp32 and flushed to zero on CDNA. Dividing by it turned the
+    block into 0/0 == NaN, which then poisoned every output column of the GEMM
+    that consumed it.
+    """
+    torch.manual_seed(num_rows)
+    device = "cuda"
+    x = torch.randn(num_rows, 128, device=device, dtype=torch.bfloat16)
+    x[0] = 0.0  # all-zero row: every block clamps to scale byte 0
+    x[-1, :32] = 0.0  # single zero block in an otherwise normal row
+    x[-1, 32:64] = 1e-38  # denormal-ish block
+
+    q_triton, s_triton = _mxfp8_e4m3_quantize_triton(x)
+    q_torch, s_torch = _mxfp8_e4m3_quantize_torch(x)
+
+    assert not torch.isnan(q_triton.float()).any()
+    torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
+    torch.testing.assert_close(
+        q_triton.view(torch.uint8), q_torch.view(torch.uint8), rtol=0, atol=0
+    )
diff --git a/tests/kernels/quantization/test_rocm_skinny_gemms.py b/tests/kernels/quantization/test_rocm_skinny_gemms.py
index f0de9c834512..7a7fe326bfa4 100644
--- a/tests/kernels/quantization/test_rocm_skinny_gemms.py
+++ b/tests/kernels/quantization/test_rocm_skinny_gemms.py
@@ -53,6 +53,17 @@
 
 DTYPES = [torch.bfloat16, torch.float16]
 
+# kWvSlots in csrc/rocm/skinny_gemms.cu; beyond it streams take overflow
+# workspaces.
+WVSPLITKRC_SLOTS = 8
+
+# The CU count is the grid size. Production passes num_compute_units(); a small
+# value manufactures better concurrency for these tests.
+WVSPLITKRC_TEST_CU = 16
+
+# Gate spin, ~125 ms on gfx950 -- must outlast the host enqueue loop it gates.
+_GATE_CYCLES = 300_000_000
+
 # Specific (N, K, M) combinations for targeted testing
 NKM_FACTORS_LLMM1 = [
     # Small, medium, large cases
@@ -116,6 +127,15 @@
 # M tiles are 64 rows, +16 for a partial tile
 M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]
 
+# (N, K, M) with more K-shards than the readback can stage in one LDS pass.
+# Spans both CHUNKK values and both N-tile counts.
+NKM_FACTORS_WVSPLITKRC_LARGE_K = [
+    (128, 6144, 128),
+    (96, 8192, 128),
+    (128, 12288, 128),
+    (32, 12288, 128),
+]
+
 NKM_FACTORS_WVSPLITK_FP8 = [
     # FP8-specific cases with K % 16 == 0
     (1, 16, 16),
@@ -222,6 +242,83 @@ def test_rocm_wvsplitkrc_kernel(n, k, m, dtype, padded_a, bias_mode, xnorm, seed
         torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)
 
 
+@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC_LARGE_K)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
+@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
+def test_rocm_wvsplitkrc_large_k(n, k, m, dtype, seed):
+    """K large enough that the split-K readback must stage LDS in batches."""
+    torch.manual_seed(seed)
+
+    xavier = math.sqrt(2 / k)
+    A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
+    B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
+
+    ref_out = torch.nn.functional.linear(A, B, None)
+    out = ops.wvSplitKrc(A, B, num_compute_units(), None)
+
+    torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)
+
+
+@pytest.mark.parametrize("n_streams", [WVSPLITKRC_SLOTS, WVSPLITKRC_SLOTS + 4])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
+@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
+def test_rocm_wvsplitkrc_multistream(n_streams, dtype, seed):
+    """Concurrent streams must not share a split-K workspace.
+
+    The workspace is cleared by protocol rather than per invocation, so sharing
+    one corrupts the reduction and returns wrong numbers rather than failing.
+    n_streams > WVSPLITKRC_SLOTS additionally exercises the overflow path.
+    """
+    torch.manual_seed(seed)
+
+    n, k, m = 16, 2048, 256
+    iters = 4
+    xavier = math.sqrt(2 / k)
+
+    B = torch.randn(m, k, dtype=dtype, device="cuda") * xavier
+    As = [
+        torch.randn(n, k, dtype=dtype, device="cuda") * xavier for _ in range(n_streams)
+    ]
+    refs = [torch.nn.functional.linear(A, B, None) for A in As]
+    torch.accelerator.synchronize()
+
+    streams = [torch.Stream() for _ in range(n_streams)]
+    outs: list[list[torch.Tensor]] = [[] for _ in range(n_streams)]
+    _release_together(streams)
+    # Round-robin the enqueues so kernels from different streams are in flight
+    # at the same time.
+    for _ in range(iters):
+        for i, s in enumerate(streams):
+            with torch.cuda.stream(s):
+                outs[i].append(ops.wvSplitKrc(As[i], B, WVSPLITKRC_TEST_CU, None))
+    for s in streams:
+        torch.accelerator.current_stream().wait_stream(s)
+    torch.accelerator.synchronize()
+
+    for i in range(n_streams):
+        for out in outs[i]:
+            torch.testing.assert_close(out, refs[i], atol=1e-3, rtol=1e-2)
+
+
+def _release_together(streams):
+    """Release every stream at once, once everything below is enqueued.
+
+    An event recorded on an idle stream is already complete, so waiting on it
+    gates nothing -- the spin, not the event, is what holds the streams.
+    """
+    blocker = torch.Stream()
+    with torch.cuda.stream(blocker):
+        torch.cuda._sleep(_GATE_CYCLES)
+        gate = torch.Event()
+        gate.record()
+    for s in streams:
+        s.wait_event(gate)
+
+
 @pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("rows_per_block", [2, 4, 8, 16])
diff --git a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py
index bedebdb59b85..d682ed05b0d8 100644
--- a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py
+++ b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for ScaledMM kernel selection logic (CPU-only)
+"""Tests for ScaledMM kernel selection logic (CPU-only).
 
 Run `pytest tests/kernels/quantization/test_scaled_mm_kernel_selection.py`.
 """
@@ -12,20 +12,90 @@
 import pytest
 import torch
 
+from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config
 from vllm.model_executor.kernels.linear import (
     AiterInt8ScaledMMLinearKernel,
     CPUInt8ScaledMMLinearKernel,
+    HummingFP8ScaledMMLinearKernel,
     Int8ScaledMMLinearKernel,
     Int8ScaledMMLinearLayerConfig,
     ScaledMMLinearKernel,
+    _get_linear_backend,
+    _resolve_backend_kernels,
+    init_fp8_linear_kernel,
     init_int8_linear_kernel,
     register_linear_kernel,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8StaticTensorSym,
+)
 from vllm.platforms import PlatformEnum
 
 pytestmark = pytest.mark.cpu_test
 
 
+def test_auto_linear_backend_is_unchanged_without_overrides():
+    kernels = [object]
+    quantizations = (
+        "fp8_block_w8a8",
+        "fp8_w8a8",
+        "int8_w8a8",
+        "mixed_precision",
+        "mxfp8",
+        "mxfp4",
+        "mxfp6",
+        "w8a16_fp8",
+        "nvfp4_w4a4",
+        "nvfp4_w4a16",
+    )
+    config = VllmConfig(kernel_config=KernelConfig(linear_backend="auto"))
+
+    with set_current_vllm_config(config):
+        for quantization in quantizations:
+            assert _get_linear_backend(quantization=quantization) == "auto"
+            assert (
+                _resolve_backend_kernels(kernels, "test", quantization=quantization)
+                is kernels
+            )
+
+
+def test_linear_backend_override_is_quantization_specific():
+    config = VllmConfig(
+        kernel_config=KernelConfig(
+            linear_backend="cutlass",
+            linear_backend_per_quant={"nvfp4_w4a16": "humming"},
+        )
+    )
+
+    with set_current_vllm_config(config):
+        assert _get_linear_backend(quantization="nvfp4_w4a16") == "humming"
+        assert _get_linear_backend(quantization="nvfp4_w4a4") == "cutlass"
+        assert _get_linear_backend(quantization="fp8_w8a8") == "cutlass"
+
+
+@patch.object(HummingFP8ScaledMMLinearKernel, "is_supported", return_value=(True, None))
+@patch("vllm.model_executor.kernels.linear.current_platform")
+def test_fp8_linear_backend_override(platform_mock, _):
+    platform_mock._enum = PlatformEnum.CUDA
+    config = VllmConfig(
+        kernel_config=KernelConfig(
+            linear_backend="cutlass",
+            linear_backend_per_quant={"fp8_w8a8": "humming"},
+        )
+    )
+
+    with set_current_vllm_config(config):
+        kernel = init_fp8_linear_kernel(
+            activation_quant_key=kFp8StaticTensorSym,
+            weight_quant_key=kFp8StaticTensorSym,
+            input_dtype=torch.bfloat16,
+            out_dtype=torch.bfloat16,
+            weight_shape=(128, 128),
+        )
+
+    assert type(kernel) is HummingFP8ScaledMMLinearKernel
+
+
 def test_is_supported_is_abstract():
     """Test that is_supported() is properly defined as abstract."""
     assert issubclass(ScaledMMLinearKernel, ABC)
diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py
index e829c63a8462..55d1f76b2d1e 100644
--- a/tests/kernels/quantization/test_triton_scaled_mm.py
+++ b/tests/kernels/quantization/test_triton_scaled_mm.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the triton_scaled_mm kernel
+"""Tests for the triton_scaled_mm kernel.
 
 Run `pytest tests/kernels/quantization/test_triton_scaled_mm.py`.
 """
diff --git a/tests/kernels/test_apply_repetition_penalties.py b/tests/kernels/test_apply_repetition_penalties.py
index 8270cf885f60..c1551e2e82e6 100644
--- a/tests/kernels/test_apply_repetition_penalties.py
+++ b/tests/kernels/test_apply_repetition_penalties.py
@@ -35,8 +35,7 @@ def test_apply_repetition_penalties(
     dtype: torch.dtype,
     seed: int,
 ) -> None:
-    """
-    Test the apply_repetition_penalties custom op
+    """Test the apply_repetition_penalties custom op
     against a reference implementation.
     """
     set_random_seed(seed)
@@ -86,8 +85,7 @@ def test_apply_repetition_penalties(
 )
 @torch.inference_mode()
 def test_apply_repetition_penalties_zero_seqs() -> None:
-    """
-    Test the apply_repetition_penalties custom op with num_seqs=0
+    """Test the apply_repetition_penalties custom op with num_seqs=0
     against a reference implementation.
     """
     num_seqs = 0
diff --git a/tests/kernels/test_awq_int4_to_int8.py b/tests/kernels/test_awq_int4_to_int8.py
index f94214785d22..dc7124e9ae70 100644
--- a/tests/kernels/test_awq_int4_to_int8.py
+++ b/tests/kernels/test_awq_int4_to_int8.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for AWQ INT4 W4A8 GEMM pipeline (SGLang kernel migration).
+"""Unit tests for AWQ INT4 W4A8 GEMM pipeline (SGLang kernel migration).
 
 Part 1: Weight packing tests
   - convert_weight_packed_scale_zp correctness
@@ -44,6 +43,7 @@ def make_awq_checkpoint_data(K, N, group_size, seed=42):
         float_ref:      [K, N] float32, reference dequantized weights
         weight_int4_orig: [K, N] int32, original int4 values (0-15)
         zeros_int4_orig:  [num_groups, N] int32, original zero points (0-15)
+
     """
     rng = np.random.RandomState(seed)
     num_groups = K // group_size
diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py
index 98a49c753c27..c16b4d44e8e8 100644
--- a/tests/kernels/test_bf16_skinny_gemm.py
+++ b/tests/kernels/test_bf16_skinny_gemm.py
@@ -269,6 +269,70 @@ def test_kda_overlap_configs_match_measured_table() -> None:
     }
 
 
+@pytest.mark.parametrize(
+    "compute_capability,supported",
+    [((9, 0), False), ((10, 0), True), ((10, 3), True), ((12, 0), True)],
+)
+def test_kda_mixed_precision_bf16_fma_capability(
+    monkeypatch: pytest.MonkeyPatch,
+    compute_capability: tuple[int, int],
+    supported: bool,
+) -> None:
+    from vllm.models.kimi_k3.nvidia.ops.cute_dsl import kda_skinny_gemm
+
+    capability = compute_capability[0] * 10 + compute_capability[1]
+    monkeypatch.setattr(
+        kda_skinny_gemm.current_platform,
+        "has_device_capability",
+        lambda target, device_id=0: capability >= target,
+    )
+
+    assert kda_skinny_gemm._has_mixed_precision_bf16_fma() is supported
+
+
+@pytest.mark.parametrize("num_tokens", [2, 8, 14])
+def test_kda_skinny_gemms_sm90_cuda_graph(num_tokens: int) -> None:
+    """Hopper must compile and capture the pre-SM100 FMA specialization."""
+    _require_capability_and_cute((9, 0))
+    from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
+        KdaSkinnyGemm,
+    )
+
+    torch.manual_seed(43 + num_tokens)
+    hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda")
+    f_ab_weight = torch.randn(144, 7168, dtype=torch.bfloat16, device="cuda")
+    f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda")
+    gemm = KdaSkinnyGemm()
+
+    gemm.run_k(gemm.run_n(hidden_states, f_ab_weight), f_b_weight)
+    torch.accelerator.synchronize()
+
+    capture_stream = torch.cuda.Stream()
+    capture_stream.wait_stream(torch.cuda.current_stream())
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=capture_stream):
+        projected_fab = gemm.run_n(hidden_states, f_ab_weight)
+        projected_fb = gemm.run_k(projected_fab, f_b_weight)
+    torch.cuda.current_stream().wait_stream(capture_stream)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    expected_fab = torch.nn.functional.linear(
+        hidden_states.float(), f_ab_weight.float()
+    )
+    expected_fb = torch.nn.functional.linear(
+        expected_fab[:, :128].to(torch.bfloat16).float(), f_b_weight.float()
+    )
+    for actual, expected in (
+        (projected_fab, expected_fab),
+        (projected_fb, expected_fb),
+    ):
+        cosine = torch.nn.functional.cosine_similarity(
+            actual.float().flatten(), expected.flatten(), dim=0
+        ).item()
+        assert cosine > 0.999
+
+
 @pytest.mark.parametrize("tp_size", [1, 2, 4, 16])
 def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None:
     from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
@@ -282,8 +346,34 @@ def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None:
     assert kda._projection_overlap_max_tokens == 0
 
 
+@pytest.mark.parametrize(
+    ("capability", "expected"),
+    [
+        ((9, 0), None),
+        ((10, 0), "cute-dsl"),
+        ((10, 3), "cute-dsl"),
+        ((10, 7), None),
+        ((12, 0), None),
+    ],
+)
+def test_kda_qkvg_flashinfer_backend(
+    monkeypatch: pytest.MonkeyPatch,
+    capability: tuple[int, int],
+    expected: str | None,
+) -> None:
+    monkeypatch.setattr(
+        current_platform,
+        "is_device_capability",
+        lambda candidate: candidate == capability,
+    )
+
+    assert k3_gemm._kda_qkvg_flashinfer_backend() == expected
+
+
+@pytest.mark.parametrize("backend", ["cute-dsl", None])
 def test_kda_qkvg_autotune_enables_full_overlap(
     monkeypatch: pytest.MonkeyPatch,
+    backend: str | None,
 ) -> None:
     flashinfer_gemm = pytest.importorskip("flashinfer.gemm")
 
@@ -304,10 +394,20 @@ def fake_mm_bf16(a, b, *, pdl, backend):
         return torch.empty(a.shape[0], b.shape[1], dtype=a.dtype)
 
     monkeypatch.setattr(flashinfer_gemm, "mm_bf16", fake_mm_bf16)
+    monkeypatch.setattr(
+        k3_gemm,
+        "_kda_qkvg_flashinfer_backend",
+        lambda: backend,
+    )
 
     k3_gemm.autotune_kda_qkvg(kda)
 
-    assert calls == [(torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")]
+    expected_calls = []
+    if backend is not None:
+        expected_calls.append(
+            (torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")
+        )
+    assert calls == expected_calls
     assert (
         kda._projection_overlap_max_tokens == k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS
     )
diff --git a/tests/kernels/test_cache_kernels.py b/tests/kernels/test_cache_kernels.py
index eff4508b0d83..a3cb91242af4 100644
--- a/tests/kernels/test_cache_kernels.py
+++ b/tests/kernels/test_cache_kernels.py
@@ -15,12 +15,10 @@
 
 @pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="Need CUDA device")
 def test_gather_cache_oob():
-    """
-    Tests for OOB read in gather_and_maybe_dequant_cache (Issue #27909).
+    """Tests for OOB read in gather_and_maybe_dequant_cache (Issue #27909).
     This test constructs a boundary case identified in the issue where
     seq_starts causes the block_table offset to read out of bounds.
     """
-
     block_size = 64
     # The kernel only supports the MLA entry sizes.
     entry_size = 576
diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py
index c1049eb65811..06897f04c0b7 100644
--- a/tests/kernels/test_compressor_kv_cache.py
+++ b/tests/kernels/test_compressor_kv_cache.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
+"""Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
 
 These tests cover:
   A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
@@ -42,6 +41,852 @@
 from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.parametrize(
+    "cache_dtype,kv_mxfp8",
+    [
+        (torch.uint8, False),
+        (torch.uint8, True),
+        (torch.bfloat16, False),
+        (torch.float8_e4m3fn, False),
+    ],
+)
+@pytest.mark.parametrize("num_tokens", [1, 17, 1023, 1024])
+@pytest.mark.parametrize("use_graph", [False, True])
+def test_dspark_context_kv_matches_query_insert(
+    cache_dtype, kv_mxfp8, num_tokens, use_graph
+):
+    """KV-only insertion must preserve every cache byte, including graph replay."""
+    from vllm.models.deepseek_v41.nvidia.dspark import _insert_context_kv
+
+    torch.manual_seed(42)
+    block_size = 256
+    num_blocks = math.ceil((num_tokens + 7) / block_size)
+    row_size = (528 if kv_mxfp8 else 584) if cache_dtype == torch.uint8 else 512
+    page_stride = math.ceil(block_size * row_size / 576) * 576 + 576
+    backing = torch.full((num_blocks, page_stride), 3, device="cuda").to(cache_dtype)
+    cache = backing.as_strided(
+        (num_blocks, block_size, row_size), (page_stride, row_size, 1)
+    )
+    reference_backing = backing.clone()
+    reference = reference_backing.as_strided(cache.shape, cache.stride())
+    kv = torch.randn(num_tokens + 3, 512, device="cuda", dtype=torch.bfloat16)
+    positions = torch.arange(num_tokens + 3, device="cuda") + 7
+    angles = torch.randn(num_tokens + 16, 32, device="cuda")
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+    slots = torch.randperm(num_blocks * block_size, device="cuda")[:num_tokens]
+    slots[::5] = -1
+    scale = torch.tensor([0.7], device="cuda")
+    # Only what _insert_context_kv actually reads: the record width comes off
+    # the cache tensor, so no per-record width attribute is stubbed here.
+    attn = SimpleNamespace(
+        swa_cache_layer=SimpleNamespace(kv_cache=cache, block_size=block_size),
+        kv_mxfp8=kv_mxfp8,
+        rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin),
+        _flashinfer_fp8_kv_scale=scale,
+    )
+
+    def legacy_insert():
+        q = torch.zeros(kv.shape[0], 8, 512, dtype=kv.dtype, device="cuda")
+        if cache_dtype == torch.uint8:
+            torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+                q,
+                kv,
+                reference.view(num_blocks, -1),
+                slots,
+                positions,
+                cos_sin,
+                8,
+                1e-20,
+                block_size,
+                True,
+                kv_mxfp8,
+            )
+        elif cache_dtype == torch.bfloat16:
+            torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
+                q,
+                kv,
+                reference,
+                slots,
+                positions,
+                cos_sin,
+                1e-20,
+                block_size,
+            )
+        else:
+            q_fp8 = torch.empty_like(q, dtype=cache_dtype)
+            torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
+                q,
+                kv,
+                q_fp8,
+                reference,
+                slots,
+                positions,
+                cos_sin,
+                scale,
+                scale,
+                1e-20,
+                block_size,
+            )
+
+    def insert():
+        _insert_context_kv(attn, kv, positions, slots)
+
+    insert()
+    graph = None
+    if use_graph:
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            insert()
+    for iteration in range(3):
+        kv.normal_()
+        positions.add_(1)
+        if iteration == 2:
+            slots.fill_(-1)
+        backing.view(torch.uint8).fill_(165)
+        reference_backing.copy_(backing)
+        before = kv.clone()
+        legacy_insert()
+        insert() if graph is None else graph.replay()
+        torch.testing.assert_close(
+            backing.view(torch.uint8),
+            reference_backing.view(torch.uint8),
+            rtol=0,
+            atol=0,
+        )
+        torch.testing.assert_close(kv, before, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="graph capture coverage"
+)
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+@pytest.mark.parametrize("use_graph", [False, True])
+@pytest.mark.parametrize(
+    "lengths",
+    [(1, 5, 13), (1, 1, 1), (1, 2, 0), (7, 8, 8), (2, 3, 3)],
+    ids=["prefill", "decode", "empty", "wrap_threshold", "packed_boundary"],
+)
+def test_v41_fused_save_compress_and_insert(
+    compress_ratio: int, use_graph: bool, lengths: tuple[int, ...]
+):
+    """Ring states: a group boundary whose previous token is outside the chunk
+    pools the ring row that chunk left behind, even when this chunk overwrites
+    that row; only the last ``capacity`` tokens of a chunk are stored.
+
+    Exercise historical group starts, adjacent requests, a chunk longer than
+    the ring, single-token decode pairs, empty requests, wrap thresholds at
+    odd and even starts, padded graph rows, and padded physical pages against
+    a Torch reference with BF16 rounding before RoPE. Replays change both
+    inputs and history so stale scratch/cache values cannot satisfy the test.
+    """
+    from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
+        fused_save_compress_norm,
+        rope_quant_insert,
+    )
+
+    torch.manual_seed(42)
+    device = "cuda"
+    capacity = 8
+    starts = [7, 8, 13]
+    actual = sum(lengths)
+    num_tokens = actual + 5
+    positions = torch.zeros(num_tokens, dtype=torch.int64, device=device)
+    req_ids = torch.zeros(actual + 2, dtype=torch.int32, device=device)
+    state_slots = torch.full((actual + 2,), -1, dtype=torch.int64, device=device)
+    cache_slots = torch.full_like(state_slots, -1)
+    query_start_loc = torch.tensor(
+        [0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32
+    ).to(device)
+    ring_blocks = torch.randperm(6, device=device)[:3]
+    raw = torch.randn(num_tokens, 512 * compress_ratio, device=device)
+    norm = torch.randn(512, dtype=torch.bfloat16, device=device)
+    latent = torch.empty(num_tokens, 512, dtype=torch.bfloat16, device=device)
+    angles = torch.randn(32, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+    # capacity x 1024 floats rounded up to 576-byte alignment.
+    state_backing = torch.empty(6, 8208, device=device)
+    state = state_backing.as_strided((6, capacity, 1024), (8208, 1024, 1))
+    cache_block = 256 // compress_ratio
+    cache_stride = math.ceil(cache_block * 584 / 576) * 576
+    cache_backing = torch.empty(3, cache_stride, dtype=torch.uint8, device=device)
+    cache = cache_backing.as_strided((3, cache_block, 584), (cache_stride, 584, 1))
+    cursor = 0
+    for req, (start, length) in enumerate(zip(starts, lengths)):
+        pos = torch.arange(start, start + length, device=device)
+        rows = slice(cursor, cursor + length)
+        positions[rows] = pos
+        req_ids[rows] = req
+        state_slots[rows] = ring_blocks[req] * capacity + pos % capacity
+        cache_slots[rows] = (2 - req) * cache_block + pos // compress_ratio
+        cursor += length
+
+    has_ring = compress_ratio == 2
+
+    def run():
+        fused_save_compress_norm(
+            raw,
+            positions,
+            state if has_ring else None,
+            state_slots,
+            query_start_loc if has_ring else None,
+            req_ids if has_ring else None,
+            norm,
+            1e-20,
+            compress_ratio,
+            latent,
+        )
+        rope_quant_insert(
+            latent, positions, cos_sin, cache, cache_slots, compress_ratio
+        )
+
+    state_backing.normal_()
+    run()
+    graph = None
+    if use_graph:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            run()
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            run()
+
+    for _ in range(3):
+        raw.normal_()
+        state_backing.normal_()
+        history = state.clone()
+        expected_state = state_backing.clone()
+        expected_rows = expected_state.as_strided((6, capacity, 1024), (8208, 1024, 1))
+        cache_backing.fill_(165)
+        expected_cache = cache_backing.clone()
+        latent.fill_(float("nan"))
+        # Ratio 1 has no ring; ratio 2 stores each chunk's last `capacity` rows.
+        for t in range(actual if has_ring else 0):
+            req = req_ids[t].item()
+            if query_start_loc[req + 1].item() - t > capacity:
+                continue
+            slot = state_slots[t].item()
+            expected_rows[slot // capacity, slot % capacity, :512] = raw[t, :512]
+            expected_rows[slot // capacity, slot % capacity, 512:] = raw[t, 512:]
+
+        if graph is None:
+            run()
+        else:
+            graph.replay()
+
+        torch.testing.assert_close(state_backing, expected_state, rtol=0, atol=0)
+        for t in range(num_tokens):
+            pos = positions[t].item()
+            if t >= actual or (pos + 1) % compress_ratio:
+                assert torch.isnan(latent[t]).all()
+                continue
+            req = req_ids[t].item()
+            group_rows = []
+            for k in range(compress_ratio - 1, -1, -1):
+                if t - k >= query_start_loc[req].item():
+                    kv_score = raw[t - k]
+                    if compress_ratio == 1:
+                        kv_score = torch.cat((kv_score, torch.zeros_like(kv_score)))
+                    group_rows.append(kv_score)
+                else:
+                    ring = ring_blocks[req].item()
+                    group_rows.append(history[ring, (pos - k) % capacity])
+            group = torch.stack(group_rows)
+            pooled = (group[:, :512] * group[:, 512:].softmax(0)).sum(0)
+            normed = pooled * torch.rsqrt(pooled.square().mean() + 1e-20) * norm
+            torch.testing.assert_close(latent[t].float(), normed, rtol=0.004, atol=1e-6)
+
+            slot = cache_slots[t].item()
+            page, row = divmod(slot, cache_block)
+            values = expected_cache[page, row * 576 : (row + 1) * 576]
+            quantized, scales = _ue8m0_reference(latent[t, :448], 64, 448.0)
+            values[:448] = quantized.view(torch.uint8)
+            c, s = cos_sin[pos // compress_ratio * compress_ratio].chunk(2)
+            rope_input = latent[t, 448:].float()
+            rotated = (
+                torch.stack(
+                    (
+                        rope_input[0::2] * c - rope_input[1::2] * s,
+                        rope_input[1::2] * c + rope_input[0::2] * s,
+                    ),
+                    dim=-1,
+                )
+                .flatten()
+                .to(torch.bfloat16)
+            )
+            actual_rope = cache_backing[page, row * 576 + 448 : (row + 1) * 576].view(
+                torch.bfloat16
+            )
+            torch.testing.assert_close(actual_rope, rotated, rtol=0.008, atol=1e-6)
+            values[448:] = actual_rope.view(torch.uint8)
+            scale_offset = cache_block * 576 + row * 8
+            expected_cache[page, scale_offset : scale_offset + 7] = (
+                scales.log2() + 127
+            ).to(torch.uint8)
+            expected_cache[page, scale_offset + 7] = 0
+        # Includes unwritten rows, page padding, quantized values and scales.
+        torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0)
+
+
+def _rotate_rope_tail(
+    latent_row: torch.Tensor, cos_sin_row: torch.Tensor
+) -> torch.Tensor:
+    """GPT-J RoPE of one latent row's last 64 dims, in fp32."""
+    row = latent_row.float()
+    c, s = cos_sin_row.chunk(2)
+    rotated = row.clone()
+    even, odd = row[448::2], row[449::2]
+    rotated[448::2] = even * c - odd * s
+    rotated[449::2] = odd * c + even * s
+    return rotated
+
+
+def _mxfp8_record_reference(
+    latent_row: torch.Tensor, cos_sin_row: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Reference V4.1 record for one token: 512 fp8 bytes + 16 UE8M0 scales.
+
+    Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the
+    RoPE tail first, then scale every 32-dim tile, RoPE tiles included.
+    """
+    rotated = _rotate_rope_tail(latent_row, cos_sin_row)
+    quantized, scales = _ue8m0_reference(rotated, 32, 448.0)
+    return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+def test_v41_rope_insert_mxfp8_record(compress_ratio: int):
+    """The 528-byte V4.1 record quantizes the RoPE dims instead of keeping them
+    in bf16, so a page is [block_size x 512 fp8][block_size x 16 UE8M0].
+
+    Byte-exact against the reference quantizer, including the rows and page
+    padding the kernel must leave alone.
+    """
+    from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
+        rope_quant_insert,
+    )
+
+    torch.manual_seed(7)
+    device = "cuda"
+    cache_block = 64
+    cache_stride = math.ceil(cache_block * 528 / 512) * 512
+    num_tokens = 12
+
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
+    angles = torch.randn(64, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+
+    cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device)
+    cache = cache_backing.as_strided((2, cache_block, 528), (cache_stride, 528, 1))
+    # Leave one token unmapped so the kernel's negative-slot skip is covered.
+    slots = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    slots[3] = -1
+
+    rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio)
+
+    expected = torch.full_like(cache_backing, 165)
+    for t in range(num_tokens):
+        slot = slots[t].item()
+        pos = positions[t].item()
+        if slot < 0 or (pos + 1) % compress_ratio:
+            continue
+        page, row = divmod(slot, cache_block)
+        values, scales = _mxfp8_record_reference(
+            latent[t], cos_sin[pos // compress_ratio * compress_ratio]
+        )
+        expected[page, row * 512 : (row + 1) * 512] = values
+        scale_offset = cache_block * 512 + row * 16
+        expected[page, scale_offset : scale_offset + 16] = scales
+
+    torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+def test_v41_mxfp8_cache_round_trip():
+    """quantize_and_insert -> dequantize_and_gather recovers the V4.1 record.
+
+    The RoPE dims now go through fp8, so they carry quantization error too --
+    the check is that every dim survives within one MXFP8 tile step, and that
+    the bytes on the way match the reference quantizer.
+    """
+    from vllm.models.deepseek_v41.common.ops import (
+        dequantize_and_gather_k_cache,
+        quantize_and_insert_k_cache,
+    )
+
+    torch.manual_seed(11)
+    device = "cuda"
+    block_size = 64
+    num_tokens = 70
+    num_blocks = 4
+    page_bytes = math.ceil(block_size * 528 / 512) * 512
+
+    k = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
+    cache = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device)
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    quantize_and_insert_k_cache(
+        k, cache, slot_mapping, block_size=block_size, bytes_per_token=528
+    )
+
+    for t in (0, 1, block_size, num_tokens - 1):
+        values, scales = _ue8m0_reference(k[t].float(), 32, 448.0)
+        page, row = divmod(t, block_size)
+        torch.testing.assert_close(
+            cache[page, row * 512 : (row + 1) * 512],
+            values.view(torch.uint8),
+            rtol=0,
+            atol=0,
+        )
+        scale_offset = block_size * 512 + row * 16
+        torch.testing.assert_close(
+            cache[page, scale_offset : scale_offset + 16],
+            (scales.log2() + 127).to(torch.uint8),
+            rtol=0,
+            atol=0,
+        )
+
+    out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device)
+    dequantize_and_gather_k_cache(
+        out,
+        cache.view(num_blocks, block_size, 528),
+        seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device),
+        gather_lens=None,
+        block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view(
+            1, -1
+        ),
+        block_size=block_size,
+        offset=0,
+    )
+    # Half an e4m3 ULP at the top of the range (16 units) bounds the error.
+    tile_amax = k.float().abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4)
+    tolerance = 16.0 * (tile_amax / 448.0).log2().ceil().exp2()
+    error = (out[0].float() - k.float()).abs()
+    assert (error <= tolerance.repeat_interleave(32, dim=-1)).all()
+
+
+# e2m1 magnitudes indexed by the low 3 bits of a code; bit 3 is the sign.
+_E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
+
+
+def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+    """Unpack e2m1 rows and apply their e4m3 tile scales."""
+    mags = torch.tensor(_E2M1_MAGNITUDES, device=packed.device)
+    codes = torch.empty(
+        (*packed.shape[:-1], 512), dtype=torch.uint8, device=packed.device
+    )
+    codes[..., 0::2] = packed & 0xF  # even element in the low nibble
+    codes[..., 1::2] = packed >> 4
+    vals = mags[(codes & 7).long()] * torch.where(codes >= 8, -1.0, 1.0)
+    return (
+        vals.reshape(*packed.shape[:-1], 32, 16)
+        * scales.view(torch.float8_e4m3fn).float().reshape(*packed.shape[:-1], 32, 1)
+    ).flatten(-2)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+def test_v41_rope_insert_nvfp4_record(compress_ratio: int):
+    """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales.
+
+    Scales are byte-exact against the reference (``amax / 6`` clamped to the
+    e4m3 range); values are checked as round-to-nearest on the e2m1 grid, whose
+    coarsest step is 2 between 4 and 6, so half a step is ``1.0 * scale``.
+    """
+    from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
+        rope_quant_insert,
+    )
+
+    torch.manual_seed(13)
+    device = "cuda"
+    cache_block = 64
+    cache_stride = math.ceil(cache_block * 288 / 512) * 512
+    num_tokens = 12
+
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
+    angles = torch.randn(64, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+
+    cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device)
+    cache = cache_backing.as_strided((2, cache_block, 288), (cache_stride, 288, 1))
+    slots = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    slots[3] = -1  # covers the negative-slot skip
+
+    rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio)
+
+    untouched = torch.full_like(cache_backing, 165)
+    for t in range(num_tokens):
+        slot = slots[t].item()
+        pos = positions[t].item()
+        if slot < 0 or (pos + 1) % compress_ratio:
+            continue
+        page, row = divmod(slot, cache_block)
+        cs = cos_sin[pos // compress_ratio * compress_ratio]
+        rotated = _rotate_rope_tail(latent[t], cs)
+        amax = rotated.view(32, 16).abs().amax(-1)
+        scale = (amax / 6.0).clamp(2.0**-9, 448.0).to(torch.float8_e4m3fn)
+
+        got_scales = cache_backing[page, cache_block * 256 + row * 32 :][:32]
+        torch.testing.assert_close(got_scales, scale.view(torch.uint8), rtol=0, atol=0)
+        decoded = _decode_nvfp4_row(
+            cache_backing[page, row * 256 : (row + 1) * 256], got_scales
+        )
+        step = scale.float().repeat_interleave(16)
+        assert ((decoded - rotated).abs() <= step).all()
+        untouched[page, row * 256 : (row + 1) * 256] = cache_backing[
+            page, row * 256 : (row + 1) * 256
+        ]
+        untouched[page, cache_block * 256 + row * 32 :][:32] = got_scales
+    # Rows the kernel must not have touched, page padding included.
+    torch.testing.assert_close(cache_backing, untouched, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.parametrize("num_reqs", [1, 8, 32])
+@pytest.mark.parametrize("num_tokens", [70, 128, 129, 2051])
+@pytest.mark.parametrize("partial", [False, True])
+def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial):
+    """The NVFP4 gather dequantizes exactly what the insert kernel wrote."""
+    from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
+    from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
+        rope_quant_insert,
+    )
+
+    torch.manual_seed(17)
+    device = "cuda"
+    block_size = 64
+    num_blocks = math.ceil(num_tokens / block_size)
+    page_bytes = math.ceil(block_size * 288 / 512) * 512
+
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
+    angles = torch.randn(num_tokens, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+    backing = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device)
+    cache = backing.as_strided((num_blocks, block_size, 288), (page_bytes, 288, 1))
+    slots = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    rope_quant_insert(latent, positions, cos_sin, cache, slots, 1)
+
+    offset = 7 if partial else 0
+    seq_lens = num_tokens - torch.arange(num_reqs, dtype=torch.int32, device=device)
+    gather_lens = seq_lens - 3 if partial else None
+    storage = torch.full(
+        (num_reqs, num_tokens + offset + 16, 512),
+        -123.0,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    out = storage[:, : num_tokens + offset]
+    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).repeat(
+        num_reqs, 1
+    )
+
+    def gather():
+        dequantize_and_gather_k_cache(
+            out,
+            cache,
+            seq_lens=seq_lens,
+            gather_lens=gather_lens,
+            block_table=block_table,
+            block_size=block_size,
+            offset=offset,
+        )
+
+    expected = _decode_nvfp4_row(
+        backing[:, : block_size * 256].reshape(-1, 256),
+        backing[:, block_size * 256 : block_size * 288].reshape(-1, 32),
+    )
+
+    def check(shorter_by):
+        assert (storage[:, num_tokens + offset :] == -123).all()
+        for req in range(num_reqs):
+            start = 3 if partial else 0
+            length = num_tokens - req - start - shorter_by
+            assert (out[req, :offset] == -123).all()
+            assert (out[req, offset + length :] == -123).all()
+            torch.testing.assert_close(
+                out[req, offset : offset + length].float(),
+                expected[start : start + length],
+                rtol=0,
+                atol=0,
+            )
+
+    gather()
+    check(0)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        gather()
+    storage.fill_(-123)
+    seq_lens.sub_(1)
+    if gather_lens is not None:
+        gather_lens.sub_(1)
+    graph.replay()
+    check(1)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+@pytest.mark.parametrize("store_fp8", [False, True])
+def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool):
+    """The FlashInfer plain-row compressed cache gets [448 NoPE | 64 RoPE] rows.
+
+    bf16 caches store the latent verbatim with the RoPE tail rotated; per-tensor
+    fp8 caches scale the bf16-rounded row by 1/fp8_scale and clamp to e4m3.
+    Non-boundary positions, negative slots and untouched rows must keep their
+    prior contents.
+    """
+    from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
+        rope_quant_insert,
+    )
+
+    torch.manual_seed(3)
+    device = "cuda"
+    num_tokens = 9
+    cache_block = 4
+    positions = torch.arange(5, 5 + num_tokens, dtype=torch.int64, device=device)
+    cache_slots = torch.randperm(3 * cache_block, device=device)[:num_tokens]
+    cache_slots[2] = -1
+    latent = torch.randn(num_tokens, 512, device=device).to(torch.bfloat16)
+    angles = torch.randn(32, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+    dtype = torch.float8_e4m3fn if store_fp8 else torch.bfloat16
+    # Pad the page stride so the kernel must honour strides, not shape.
+    cache_backing = torch.full((3, (cache_block + 1) * 512), 3.0, device=device)
+    cache_backing = cache_backing.to(dtype)
+    cache = cache_backing.as_strided(
+        (3, cache_block, 512), ((cache_block + 1) * 512, 512, 1)
+    )
+    fp8_scale = torch.tensor([0.5], dtype=torch.float32, device=device)
+    expected = cache_backing.clone().float()
+
+    rope_quant_insert(
+        latent,
+        positions,
+        cos_sin,
+        cache,
+        cache_slots,
+        compress_ratio,
+        fp8_scale=fp8_scale if store_fp8 else None,
+    )
+
+    for t in range(num_tokens):
+        pos = positions[t].item()
+        slot = cache_slots[t].item()
+        if slot < 0 or (pos + 1) % compress_ratio:
+            continue
+        c, s = cos_sin[pos // compress_ratio * compress_ratio].chunk(2)
+        rope_input = latent[t, 448:].float()
+        rotated = torch.stack(
+            (
+                rope_input[0::2] * c - rope_input[1::2] * s,
+                rope_input[1::2] * c + rope_input[0::2] * s,
+            ),
+            dim=-1,
+        ).flatten()
+        row = torch.cat((latent[t, :448], rotated.to(torch.bfloat16))).float()
+        if store_fp8:
+            row = (row * (1.0 / fp8_scale)).clamp(-448.0, 448.0)
+            row = row.to(torch.float8_e4m3fn).float()
+        page, idx = divmod(slot, cache_block)
+        expected[page, idx * 512 : (idx + 1) * 512] = row
+
+    actual = cache_backing.float()
+    if store_fp8:
+        # One e4m3 ulp of slack for FMA-vs-separate rounding in the RoPE tail.
+        torch.testing.assert_close(actual, expected, rtol=0.13, atol=1e-2)
+    else:
+        nope = torch.arange(512, device=device) < 448
+        nope = nope.repeat(cache_block + 1)
+        torch.testing.assert_close(actual[:, nope], expected[:, nope], rtol=0, atol=0)
+        torch.testing.assert_close(
+            actual[:, ~nope], expected[:, ~nope], rtol=0.008, atol=1e-6
+        )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Triton kernel")
+def test_v41_compressor_metadata_maps_tokens_to_their_ring():
+    """The ring group's generic slot mapping is disabled (all PAD), so the
+    builder must map every real token to ``ring_block * capacity + pos %
+    capacity`` and keep padding tokens at PAD."""
+    from unittest.mock import MagicMock
+
+    from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder
+    from vllm.v1.attention.backend import CommonAttentionMetadata
+    from vllm.v1.kv_cache_interface import CircularBufferSpec
+
+    capacity = 8
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_batched_tokens = 16
+    spec = CircularBufferSpec(
+        block_size=capacity,
+        num_kv_heads=1,
+        head_size=1024,
+        head_size_v=0,
+        dtype=torch.float32,
+    )
+    device = torch.device("cuda")
+    builder = CompressorMetadataBuilder(spec, ["state"], vllm_config, device)
+
+    # Two requests: 3 tokens at positions 13..15 on ring block 5, then 2
+    # tokens at positions 7..8 on ring block 2; three padding tokens.
+    query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
+    positions = torch.tensor([13, 14, 15, 7, 8, 0, 0, 0], device=device)
+    block_table = torch.tensor([[5], [2]], dtype=torch.int32, device=device)
+    common = CommonAttentionMetadata(
+        query_start_loc=query_start_loc,
+        query_start_loc_cpu=query_start_loc.cpu(),
+        seq_lens=torch.tensor([16, 9], dtype=torch.int32, device=device),
+        num_reqs=2,
+        num_actual_tokens=5,
+        max_query_len=3,
+        max_seq_len=16,
+        block_table_tensor=block_table,
+        slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device),
+        positions=positions,
+    )
+    metadata = builder.build(0, common)
+
+    expected = [5 * 8 + 5, 5 * 8 + 6, 5 * 8 + 7, 2 * 8 + 7, 2 * 8 + 0, -1, -1, -1]
+    assert metadata.slot_mapping.tolist() == expected
+    assert metadata.query_start_loc is query_start_loc
+    assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage")
+@pytest.mark.parametrize(
+    "use_aux,use_graph", [(False, False), (True, False), (True, True)]
+)
+@torch.inference_mode()
+def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph):
+    """Both reused-event joins must publish this forward's states and cache rows."""
+    from vllm.forward_context import ForwardContext, override_forward_context
+    from vllm.models.deepseek_v41.attention import (
+        DeepseekV4Attention,
+        DeepseekV4Indexer,
+    )
+    from vllm.models.deepseek_v41.compressor import DeepseekCompressor
+
+    torch.manual_seed(43)
+    raw = torch.randn(19, 1024, device="cuda")
+    q = torch.zeros(19, 512, dtype=torch.bfloat16, device="cuda")
+    positions = torch.arange(7, 26, device="cuda")
+    state_slots = positions.clone()
+    state_slots[-2:] = -1
+    cache_slots = torch.where(state_slots >= 0, positions // 2, -1)
+    state = torch.empty(4, 8, 1024, device="cuda")
+    main = torch.empty(1, 128, 584, dtype=torch.uint8, device="cuda")
+    index = torch.empty(1, 128, 132, dtype=torch.uint8, device="cuda")
+    caches = (state, main, index)
+    observed = tuple(torch.empty_like(cache) for cache in caches)
+    rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda"))
+    metadata = {
+        "state": SimpleNamespace(
+            slot_mapping=state_slots,
+            query_start_loc=torch.tensor([0, 19], dtype=torch.int32, device="cuda"),
+            token_to_req_indices=torch.zeros(19, dtype=torch.int32, device="cuda"),
+        ),
+        "main": SimpleNamespace(slot_mapping=cache_slots),
+        "index": SimpleNamespace(slot_mapping=cache_slots),
+    }
+    context = ForwardContext({}, metadata, {})
+    compressor = DeepseekCompressor.__new__(DeepseekCompressor)
+    torch.nn.Module.__init__(compressor)
+    compressor.head_dim, compressor.rope_head_dim, compressor.compress_ratio = (
+        512,
+        64,
+        2,
+    )
+    compressor.rms_norm_eps = 1e-20
+    compressor.norm = SimpleNamespace(
+        weight=torch.ones(512, dtype=torch.bfloat16, device="cuda")
+    )
+    compressor.state_cache = SimpleNamespace(prefix="state", kv_cache=state)
+    compressor.k_cache_prefix = "main"
+    compressor._static_forward_context = {"main": SimpleNamespace(kv_cache=main)}
+    indexer_weight = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda")
+    indexer = SimpleNamespace(
+        owns_k=True,
+        wk=lambda latent: (torch.nn.functional.linear(latent, indexer_weight), None),
+        k_norm=SimpleNamespace(
+            weight=torch.ones(128, dtype=torch.bfloat16, device="cuda"),
+            variance_epsilon=1e-20,
+        ),
+        k_cache=SimpleNamespace(prefix="index", kv_cache=index),
+        compress_ratio=2,
+        use_fp4_kv=False,
+    )
+
+    def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale):
+        DeepseekV4Indexer._produce_k(indexer, latent, positions, rotary)
+        return None, None, None
+
+    def observe(*args):
+        for output, cache in zip(observed, caches):
+            output.copy_(cache)
+
+    auxiliary = [torch.cuda.Stream()] if use_aux else None
+    attention = SimpleNamespace(
+        compressor=compressor,
+        indexer=prepare_indexer,
+        aux_stream_list=auxiliary,
+        ln_events=[torch.cuda.Event(), torch.cuda.Event()],
+        rotary_emb=rotary,
+        indexer_rotary_emb=rotary,
+        n_local_heads=1,
+        head_dim=512,
+        _wq_b_proj=lambda qr, scale: qr.clone(),
+        _fused_qnorm_rope_kv_insert=lambda q, kv, pos, meta: q,
+        _sparse_indexer_and_attn=observe,
+    )
+
+    def run():
+        with override_forward_context(context):
+            DeepseekV4Attention._prepare_and_attn(
+                attention, q, q, q, None, raw, q, positions, q
+            )
+
+    def reset():
+        state.zero_()
+        main.fill_(165)
+        index.fill_(165)
+
+    reset()
+    run()
+    graph = None
+    if use_graph:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            run()
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            run()
+
+    previous = None
+    for _ in range(3):
+        raw.normal_()
+        reset()
+        attention.aux_stream_list = None
+        run()
+        expected = tuple(output.clone() for output in observed)
+        reset()
+        attention.aux_stream_list = auxiliary
+        if graph is None:
+            run()
+        else:
+            graph.replay()
+        for actual, reference in zip(observed, expected):
+            torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+        if previous is not None:
+            assert all(not torch.equal(a, b) for a, b in zip(expected, previous))
+        previous = expected
+
+
 def _on_gfx950() -> bool:
     if not current_platform.is_rocm():
         return False
@@ -273,9 +1118,13 @@ def test_gfx950_compressed_cache_canonicalizes_nonfinite(writer: str) -> None:
     ],
 )
 def test_get_c128_boundary(starts, query_start_loc, expected):
+    query_start_loc_tensor = torch.tensor(query_start_loc)
+    query_lens = query_start_loc_tensor[1:] - query_start_loc_tensor[:-1]
     metadata = SimpleNamespace(
-        _num_computed_tokens_cpu=None if starts is None else torch.tensor(starts),
-        query_start_loc_cpu=torch.tensor(query_start_loc),
+        seq_lens_cpu_upper_bound=(
+            None if starts is None else torch.tensor(starts) + query_lens
+        ),
+        query_start_loc_cpu=query_start_loc_tensor,
     )
     assert _get_c128_boundary(metadata) is expected
 
@@ -288,7 +1137,6 @@ def test_get_c128_boundary(starts, query_start_loc, expected):
 def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size: int):
     """compressed_kv → quantize_and_insert_k_cache → dequantize_and_gather_k_cache
     → compare against original."""
-
     HEAD_DIM = 512
     NOPE_DIM = 448
     HEAD_BYTES = 584  # 448 fp8 + 128 bf16 + 8 uint8 scale
@@ -499,9 +1347,8 @@ def test_dequantize_and_gather_k_cache(
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
 @pytest.mark.parametrize("block_size", [16, 64])
 def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int):
-    """k → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache
+    """K → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache
     → manual dequant → compare against original."""
-
     HEAD_DIM = 128
     QUANT_BLOCK_SIZE = 128
     # cache_stride = head_dim + (head_dim * 4 / quant_block_size) = 128 + 4 = 132
@@ -565,7 +1412,6 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int):
 
 def test_indexer_gather_accepts_upper_bound_output():
     """Gather only exact cu_seq_lens even when dst is over-allocated."""
-
     head_dim = 128
     quant_block_size = 128
     cache_stride = head_dim + head_dim * 4 // quant_block_size
@@ -665,7 +1511,6 @@ def test_indexer_gather_accepts_upper_bound_output():
 
 def test_deepseek_v4_quant_magnitude_range():
     """Test that quantization handles a range of magnitudes correctly."""
-
     HEAD_DIM = 512
     NOPE_DIM = 448
     HEAD_BYTES = 584
diff --git a/tests/kernels/test_concat_mla_q.py b/tests/kernels/test_concat_mla_q.py
index 6ace34e503e4..6d9a1021064e 100644
--- a/tests/kernels/test_concat_mla_q.py
+++ b/tests/kernels/test_concat_mla_q.py
@@ -148,3 +148,38 @@ def test_concat_mla_q_values_preserved(num_tokens):
     assert torch.equal(out_bits[..., :nope_dim], ql_nope_bits)
 
     assert torch.equal(out_bits[..., nope_dim:], q_pe_bits)
+
+
+@pytest.mark.parametrize("rope_dim,transposed", [(0, False), (0, True), (64, False)])
+@pytest.mark.parametrize("enabled", [False, True])
+def test_concat_mla_q_fp8_nope_and_rope(rope_dim, transposed, enabled):
+    """FP8 query packing preserves values with empty RoPE and strided BMM output."""
+    from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.attention.mla_attention import _DecodeConcatQuantFP8
+    from vllm.model_executor.layers.quantization.utils.quant_utils import (
+        GroupShape,
+        get_fp8_min_max,
+    )
+    from vllm.platforms import current_platform
+
+    torch.manual_seed(42)
+    shape = (8, 17, 512) if transposed else (17, 8, 512)
+    ql = torch.randn(shape, dtype=torch.bfloat16, device="cuda") * 128
+    if transposed:
+        ql = ql.transpose(0, 1)
+    qp = torch.randn(17, 8, rope_dim, dtype=torch.bfloat16, device="cuda") * 128
+    scale = torch.tensor(0.3, device="cuda", dtype=torch.float32)
+    joined = torch.cat((ql, qp), dim=-1)
+    fp8_min, fp8_max = get_fp8_min_max()
+    expected = (joined.float() * scale.reciprocal()).clamp(fp8_min, fp8_max)
+    expected = expected.to(current_platform.fp8_dtype())
+    custom_ops = ["none"] + (["+mla_decode_concat_quant_fp8"] if enabled else [])
+    config = VllmConfig(
+        compilation_config=CompilationConfig(mode=0, custom_ops=custom_ops)
+    )
+    with set_current_vllm_config(config):
+        quant = _DecodeConcatQuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
+        actual = quant(ql, qp, scale)
+    assert actual.dtype == current_platform.fp8_dtype()
+    torch.testing.assert_close(actual.float(), expected.float(), atol=0, rtol=0)
+    assert actual.is_contiguous()
diff --git a/tests/kernels/test_cp_gather_fp8.py b/tests/kernels/test_cp_gather_fp8.py
index 35afe33808cc..40c5ef5fdbd0 100644
--- a/tests/kernels/test_cp_gather_fp8.py
+++ b/tests/kernels/test_cp_gather_fp8.py
@@ -29,6 +29,7 @@ def _build_test_case(seq_lens, block_size, seed=42):
     Returns:
         Tuple of (cache, block_table, workspace_starts_t, num_reqs,
                   total_tokens, expected_output).
+
     """
     torch.manual_seed(seed)
 
@@ -316,6 +317,48 @@ def test_cp_gather_fp8_shuffled_blocks():
     assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
 
 
+def test_cp_gather_fp8_mixes_resident_and_pinned_host_rows():
+    seq_lens = [7, 5]
+    block_size = 4
+    (
+        cache,
+        block_table,
+        workspace_starts,
+        num_reqs,
+        total_tokens,
+        expected,
+    ) = _build_test_case(seq_lens, block_size)
+
+    device_cache = cache.clone()
+    host_cache = cache.cpu().pin_memory().view(-1, ENTRY_BYTES)
+    device_rows = torch.arange(host_cache.shape[0], dtype=torch.int32, device="cuda")
+    use_device = device_rows % 2 == 0
+    host_rows = device_rows.clone()
+    device_rows.masked_fill_(~use_device, -1)
+
+    host_cache[use_device.cpu()] = 0
+    device_cache.view(-1, ENTRY_BYTES)[~use_device] = 0
+    dst = torch.empty(
+        total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
+    )
+
+    ops.cp_gather_and_upconvert_fp8_kv_cache(
+        device_cache,
+        dst,
+        block_table,
+        workspace_starts,
+        num_reqs,
+        host_cache=host_cache,
+        host_row_ids=host_rows,
+        device_row_ids=device_rows,
+    )
+
+    torch.testing.assert_close(
+        dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
+    )
+    assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
+
+
 @pytest.mark.parametrize(
     "gather_seq_lens,seq_starts",
     [
diff --git a/tests/kernels/test_dsv41_mega_attn_layouts.py b/tests/kernels/test_dsv41_mega_attn_layouts.py
new file mode 100644
index 000000000000..6f9ff56bd236
--- /dev/null
+++ b/tests/kernels/test_dsv41_mega_attn_layouts.py
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Layout plumbing for DeepSeek V4.1 mega attention.
+
+The kernel's Q and O layouts are produced by permuting ``wq_b`` rows and
+``wo_a`` columns once at load, so what has to hold is that the permuted GEMMs
+agree with a reference that permutes the activation instead.
+"""
+
+import torch
+
+from vllm.models.deepseek_v41.common.ops.fused_layout import (
+    WV_GROUP_SIZE,
+    o_fused_chunk_permutation,
+    o_fused_permutation,
+    permute_q_to_fused,
+    permute_wo_a_,
+    permute_wq_b_,
+    q_fused_permutation,
+)
+
+HEAD_DIM = 512
+
+
+def test_permuted_wq_b_gemm_matches_permuted_activation():
+    """Permuting wq_b's rows makes the Q GEMM emit the fused layout directly."""
+    torch.manual_seed(0)
+    num_heads, q_lora_rank, num_tokens = 16, 64, 5
+    weight = torch.randn(num_heads * HEAD_DIM, q_lora_rank)
+    # One scale row per weight row, as an MXFP8 wq_b shard carries.
+    scale = torch.randint(
+        0, 256, (num_heads * HEAD_DIM, q_lora_rank // 32), dtype=torch.uint8
+    )
+    qr = torch.randn(num_tokens, q_lora_rank)
+    orig_weight, orig_scale = weight.clone(), scale.clone()
+
+    standard = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM)
+    expected = permute_q_to_fused(standard)
+
+    permute_wq_b_(weight, scale, num_heads)
+    fused = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM)
+    torch.testing.assert_close(fused, expected)
+    # The scale has to follow its row, or dequant reads another row's scale.
+    perm = q_fused_permutation(num_heads, HEAD_DIM)
+    torch.testing.assert_close(weight, orig_weight[perm])
+    torch.testing.assert_close(scale, orig_scale[perm])
+
+
+def test_permuted_wo_a_consumes_fused_output():
+    """Permuting wo_a's input columns lets it read the kernel's O layout."""
+    torch.manual_seed(1)
+    out_features, num_tokens = 32, 4
+    in_features = WV_GROUP_SIZE * HEAD_DIM
+    weight = torch.randn(out_features, in_features)
+    scale = torch.randint(0, 256, (out_features, in_features // 32), dtype=torch.uint8)
+    standard_o = torch.randn(num_tokens, in_features)
+    orig_weight, orig_scale = weight.clone(), scale.clone()
+
+    expected = standard_o @ weight.T
+    # The kernel emits the same values in the fused chunk order.
+    perm = o_fused_permutation(WV_GROUP_SIZE, HEAD_DIM)
+    fused_o = standard_o[:, perm]
+
+    permute_wo_a_(weight, scale, WV_GROUP_SIZE)
+    torch.testing.assert_close(weight, orig_weight[:, perm])
+    torch.testing.assert_close(
+        scale, orig_scale[:, o_fused_chunk_permutation(WV_GROUP_SIZE, HEAD_DIM)]
+    )
+    # Unlike wq_b, this permutation sits inside the 4096-term reduction, so the
+    # summation order changes and the result differs in the last few ulps.
+    torch.testing.assert_close(fused_o @ weight.T, expected, rtol=1e-3, atol=1e-3)
diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py
new file mode 100644
index 000000000000..45cebcfe5711
--- /dev/null
+++ b/tests/kernels/test_engram.py
@@ -0,0 +1,994 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import bisect
+import inspect
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v41.common import engram as engram_ops
+from vllm.models.deepseek_v41.common.engram import (
+    Engram as CommonEngram,
+)
+from vllm.models.deepseek_v41.common.engram import (
+    NgramHashState,
+)
+from vllm.models.deepseek_v41.nvidia import engram as nvidia_engram_ops
+from vllm.models.deepseek_v41.nvidia.engram import Engram, ParallelEngramEmbedding
+from vllm.platforms import current_platform
+
+
+def _reference_engram_post_wkv(
+    hidden_states,
+    kv,
+    q_weight,
+    k_weight,
+    token_mask,
+    eps,
+    clamp_value,
+    kv_start,
+):
+    num_tokens, hc_mult, dim = hidden_states.shape
+    source = torch.arange(num_tokens, device=kv.device) + kv_start
+    valid = source < kv.shape[0]
+    local_kv = kv.new_zeros(num_tokens, kv.shape[1])
+    local_kv[valid] = kv[source[valid]]
+
+    key = local_kv[:, : hc_mult * dim].float().view(num_tokens, hc_mult, dim)
+    value = local_kv[:, hc_mult * dim :].float()
+    h = hidden_states.float()
+    weight = q_weight.float() * k_weight.float()
+    rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(
+        key.square().mean(-1) + eps
+    )
+    dot = (h * weight * key).sum(-1) * rstd * dim**-0.5
+    gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(clamp_value).sqrt(), dot))
+    if token_mask is not None:
+        local_mask = torch.zeros(num_tokens, dtype=torch.bool, device=kv.device)
+        local_mask[valid] = token_mask[source[valid]]
+        gate = gate.masked_fill(~local_mask.unsqueeze(-1), 0)
+    return (h + gate.unsqueeze(-1) * value.unsqueeze(-2)).to(hidden_states.dtype)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "num_tokens,num_kv_tokens,hc_mult,dim,tp_size,tp_rank,use_mask",
+    [
+        (1, 1, 4, 5120, 1, 0, False),
+        (3, 5, 4, 96, 2, 0, True),
+        (3, 5, 2, 257, 2, 1, False),
+        (3, 5, 4, 96, 2, 1, True),
+        (1, 1, 4, 96, 4, 3, True),
+        (0, 0, 4, 96, 4, 0, False),
+    ],
+)
+def test_fused_engram_post_wkv_matches_reference(
+    num_tokens,
+    num_kv_tokens,
+    hc_mult,
+    dim,
+    tp_size,
+    tp_rank,
+    use_mask,
+    monkeypatch,
+):
+    """Preserve gate math, WKV layout, masking, and SP padding semantics."""
+    torch.manual_seed(0)
+    device = "cuda"
+    hidden_states = torch.randn(
+        num_tokens, hc_mult, dim, dtype=torch.bfloat16, device=device
+    )
+    kv = torch.randn(
+        num_kv_tokens,
+        (hc_mult + 1) * dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    q_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device)
+    k_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device)
+    token_mask = None
+    if use_mask:
+        token_mask = torch.arange(num_kv_tokens, device=device) % 2 == 0
+    eps = 1e-20
+    clamp_value = 1e-6
+    use_sequence_parallel = tp_size > 1
+    shard_size = (num_kv_tokens + tp_size - 1) // tp_size
+    kv_start = tp_rank * shard_size if use_sequence_parallel else 0
+
+    expected = _reference_engram_post_wkv(
+        hidden_states,
+        kv,
+        q_weight,
+        k_weight,
+        token_mask,
+        eps,
+        clamp_value,
+        kv_start,
+    )
+    monkeypatch.setattr(
+        engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank)
+
+    module = CommonEngram.__new__(CommonEngram)
+    torch.nn.Module.__init__(module)
+    module.dim = dim
+    module.hc_mult = hc_mult
+    module.eps = eps
+    module.clamp_value = clamp_value
+    module.use_sequence_parallel = use_sequence_parallel
+    module.embed_tokens = torch.nn.Identity()
+    # `forward` reads rows staged by `prepare_embeddings`, so inject kv there.
+    module.embed_tokens.tp_size = module.embed_tokens.n_hash_cols = 1
+    module.staged_rows = kv.unsqueeze(1)
+    if use_sequence_parallel:
+        padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size))
+        module.staged_rows = padded[
+            tp_rank * shard_size : (tp_rank + 1) * shard_size
+        ].unsqueeze(1)
+    module.wkv = torch.nn.Identity()
+
+    def check_local_projection(_, args):
+        assert args[0].shape[0] == num_tokens
+
+    module.wkv.register_forward_pre_hook(check_local_projection)
+    module.q_weight = torch.nn.Parameter(q_weight, requires_grad=False)
+    module.k_weight = torch.nn.Parameter(k_weight, requires_grad=False)
+
+    actual = module(hidden_states, kv.unsqueeze(1), token_mask)
+
+    assert actual.shape == hidden_states.shape
+    assert actual.dtype == hidden_states.dtype
+    torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)
+
+
+def _hash_state(use_slot_cache: bool) -> NgramHashState:
+    state = NgramHashState.__new__(NgramHashState)
+    torch.nn.Module.__init__(state)
+    state.block_size = 64
+    state.use_slot_cache = use_slot_cache
+    state._cache = None
+    state._kv_cache_ref = None
+    state.swa_cache_module = torch.nn.Module()
+    state.swa_cache_module.kv_cache = torch.empty(0)
+    return state
+
+
+@pytest.mark.parametrize("num_tokens", [0, 7])
+def test_engram_dummy_hashes_leave_history_untouched(num_tokens):
+    """Dummy forwards have no valid table IDs and do not alter cached history."""
+    state = _hash_state(use_slot_cache=True)
+    state.multipliers = torch.empty(2, 4, dtype=torch.int64)
+    state.primes = torch.empty(2, 3, 8, dtype=torch.int64)
+    state._cache = torch.arange(16, dtype=torch.int32)
+    history = state._cache.clone()
+
+    hashes, keep = state.dummy_hashes(torch.arange(num_tokens))
+
+    assert hashes.shape == (num_tokens, 2, 24)
+    assert hashes.dtype == torch.int32
+    assert torch.all(hashes == engram_ops.DEAD_ID)
+    assert keep.shape == (num_tokens,) and keep.dtype == torch.bool
+    assert not keep.any()
+    torch.testing.assert_close(state._cache, history)
+
+
+@pytest.mark.parametrize("num_blocks", [4, 100])
+def test_engram_cache_follows_kv_cache_binding(num_blocks):
+    """Discard profiling history when the real KV cache replaces its storage."""
+    state = _hash_state(use_slot_cache=True)
+    assert not state.ensure_cache()
+
+    state.swa_cache_module.kv_cache = torch.empty(4, 64, 1)
+    assert state.ensure_cache()
+    profiling_cache = state._cache
+    profiling_cache.fill_(7)
+    assert state.ensure_cache()
+    assert state._cache is profiling_cache
+    assert torch.all(state._cache == 7)
+
+    state.swa_cache_module.kv_cache = torch.empty(num_blocks, 64, 1)
+    assert state.ensure_cache()
+    assert state._cache is not profiling_cache
+    assert state._cache.shape == (num_blocks * 64,)
+    assert torch.count_nonzero(state._cache) == 0
+
+    state.swa_cache_module.kv_cache = torch.empty(0)
+    assert not state.ensure_cache()
+    assert state._cache is None
+
+
+def test_engram_without_slot_cache_only_gates_on_kv_binding():
+    """The V2 runner supplies every lookback, so no slot cache is allocated."""
+    state = _hash_state(use_slot_cache=False)
+    assert not state.ensure_cache()
+    state.swa_cache_module.kv_cache = torch.empty(4, 64, 1)
+    assert state.ensure_cache()
+    assert state._cache is None
+
+
+def _reference_hashes(
+    ids,
+    positions,
+    slots,
+    starts,
+    tables,
+    dead,
+    cache,
+    token_map,
+    multipliers,
+    primes,
+    offsets,
+):
+    """Scalar n-gram oracle with persistent, physically addressed token history."""
+    for token, slot, is_dead in zip(ids, slots, dead):
+        if slot >= 0:
+            cache[slot] = -1 if is_dead else token_map[token]
+    result = []
+    for i, position in enumerate(positions):
+        req = min(bisect.bisect_right(starts[1:], i), len(tables) - 1)
+        history, blocked = [], False
+        for shift in range(4):
+            lookback = position - shift
+            p = max(0, min(lookback, len(tables[req]) * 16 - 1))
+            source = cache[tables[req][p // 16] * 16 + p % 16]
+            blocked |= lookback < 0 or source == -1
+            history.append(0 if blocked else source)
+        layers = []
+        for layer, row in enumerate(multipliers):
+            rolling, hashes = history[0] * row[0], []
+            for j, head_primes in enumerate(primes[layer], start=1):
+                rolling ^= history[j] * row[j]
+                for head, prime in enumerate(head_primes):
+                    col = (j - 1) * len(head_primes) + head
+                    hashes.append(rolling % prime + offsets[layer][col])
+            layers.append(hashes)
+        result.append(layers)
+    return torch.tensor(result, dtype=torch.int32)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("query_len", [1, 17, 257])
+@pytest.mark.parametrize("capture", [False, True])
+@pytest.mark.parametrize("strided_slots", [False, True])
+@pytest.mark.parametrize("large_hashes", [False, True])
+def test_engram_hash_cache_replay(query_len, capture, strided_slots, large_hashes):
+    """Preserve history, dead tokens and padding as inputs change across replay."""
+    num_tokens = 2 * query_len + 3
+    blocks = (3 * query_len + 15) // 16
+    tables = [list(range(blocks, 2 * blocks)), list(range(blocks))]
+    starts = [0, query_len, 2 * query_len]
+
+    def tensor(data, dtype=torch.int64):
+        return torch.tensor(data, dtype=dtype, device="cuda")
+
+    input_ids = tensor([0] * num_tokens)
+    positions = torch.zeros_like(input_ids)
+    slots = tensor([-1] * (num_tokens * (2 if strided_slots else 1)))
+    if strided_slots:
+        slots = slots[::2]
+    dead_mask = tensor([False] * num_tokens, torch.bool)
+    cache = tensor([-1] * (2 * blocks * 16), torch.int32)
+    token_map = [i * 3091 if large_hashes else i % 7 for i in range(32)]
+    multipliers = (
+        engram_ops.compute_hash_multipliers((1, 14), 4, 99092).tolist()
+        if large_hashes
+        else [[3, 5, 7, 9], [11, 13, 15, 17]]
+    )
+    primes = (
+        [
+            [
+                [16000057 + 2 * (layer * 24 + j * 8 + h) for h in range(8)]
+                for j in range(3)
+            ]
+            for layer in range(2)
+        ]
+        if large_hashes
+        else [[[19], [23], [29]]] * 2
+    )
+    offsets = []
+    for layer in primes:
+        row, offset = [], 0
+        for heads in layer:
+            for prime in heads:
+                row.append(offset)
+                offset += prime
+        offsets.append(row)
+    state = NgramHashState.__new__(NgramHashState)
+    torch.nn.Module.__init__(state)
+    state.token_map = tensor(token_map, torch.int32)
+    state.multipliers = tensor(multipliers)
+    state.primes = tensor(primes)
+    state.offsets = tensor(offsets)
+    state._cache = cache
+    state.use_slot_cache = True
+    state.pad_id = 0
+    state.block_size = 16
+    lookback = tensor([[-1] * 3] * 2, torch.int32)
+    args = (
+        input_ids,
+        positions,
+        tensor(starts, torch.int32),
+        dead_mask,
+        lookback,
+        torch.zeros_like(lookback, dtype=torch.bool),
+        slots,
+        tensor(tables, torch.int32),
+    )
+    op = state
+    graph = None
+    if capture:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                op(*args)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            output = op(*args)
+        torch.cuda.current_stream().wait_stream(stream)
+    cache.fill_(-1)
+    expected_cache = [-1] * cache.numel()
+
+    for step in range(3):
+        ids = [(i + 3 * step) % 32 for i in range(num_tokens)]
+        pos = list(range(step * query_len, (step + 1) * query_len)) * 2
+        physical_slots = [
+            tables[i // query_len][p // 16] * 16 + p % 16 for i, p in enumerate(pos)
+        ]
+        pos += [0] * 3
+        physical_slots += [-1] * 3
+        # A fully padded replay must leave every cache slot untouched.
+        if step == 2:
+            physical_slots = [-1] * num_tokens
+        dead = [(i + step) % 5 == 0 for i in range(num_tokens)]
+        input_ids.copy_(tensor(ids))
+        positions.copy_(tensor(pos))
+        slots.copy_(tensor(physical_slots))
+        dead_mask.copy_(tensor(dead, torch.bool))
+        if graph is None:
+            output = op(*args)
+        else:
+            graph.replay()
+        expected = _reference_hashes(
+            ids,
+            pos,
+            physical_slots,
+            starts,
+            tables,
+            dead,
+            expected_cache,
+            token_map,
+            multipliers,
+            primes,
+            offsets,
+        )
+        if step != 2:
+            real = 2 * query_len
+            torch.testing.assert_close(
+                output[:real].cpu(), expected[:real], rtol=0, atol=0
+            )
+        assert cache.cpu().tolist() == expected_cache
+
+    empty_output = state(
+        input_ids[:0], positions[:0], args[2], dead_mask[:0], *args[4:]
+    )
+    assert empty_output.shape == (0, 2, len(offsets[0]))
+    assert cache.cpu().tolist() == expected_cache
+
+
+def _hash_ids(requests, tables_meta, block_table, cache=None, capture=False):
+    """Run the op on a batch of (token_ids, start, window) requests; `window`
+    is the runner's [depth] lookback for that request (None = all unknown).
+    Token IDs 14 and 22 stand in for image tokens that break n-grams.
+    Returns one [len(token_ids), ...] hash tensor per request."""
+    token_map, multipliers, primes, offsets, block_size = tables_meta
+    depth = multipliers.shape[1] - 1
+
+    def tensor(data, dtype=torch.int64):
+        return torch.tensor(data, dtype=dtype, device="cuda")
+
+    input_ids, positions, slots, windows, starts = [], [], [], [], [0]
+    for req, (token_ids, start, window) in enumerate(requests):
+        pos = list(range(start, start + len(token_ids)))
+        input_ids += token_ids
+        positions += pos
+        slots += [
+            int(block_table[req, p // block_size]) * block_size + p % block_size
+            for p in pos
+        ]
+        windows.append(window or [-1] * depth)
+        starts.append(starts[-1] + len(token_ids))
+    window = tensor(windows, torch.int32)
+    state = _hash_state(use_slot_cache=cache is not None)
+    state.token_map = token_map
+    state.multipliers = multipliers
+    state.primes = primes
+    state.offsets = offsets
+    state.pad_id = 0
+    state.block_size = block_size
+    state._cache = cache
+    input_tensor = tensor(input_ids, torch.int32)
+    args = (
+        input_tensor,
+        tensor(positions),
+        tensor(starts, torch.int32),
+        (input_tensor == 14) | (input_tensor == 22),
+        window,
+        (window == 14) | (window == 22),
+        tensor(slots) if cache is not None else None,
+        block_table if cache is not None else None,
+    )
+    out = state(*args)
+    if capture:
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            out = state(*args)
+        graph.replay()
+    return [out[a:b] for a, b in zip(starts, starts[1:])]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("runner", ["v2_window_only", "v1_prompt_window_plus_cache"])
+@pytest.mark.parametrize("capture", [False, True])
+def test_lookback_window_reproduces_single_instance(runner, capture):
+    """Decoding on a fresh instance whose prompt KV came from elsewhere (P/D,
+    offload) must hash like the instance that processed the whole sequence,
+    with two requests batched and chunks of 1-2 tokens so in-batch, window
+    and (V1) cache lookbacks mix within one step.
+
+    V2 supplies every lookback from its device token history and needs no
+    slot cache. V1 supplies prompt positions only and reads generated
+    positions from the slot cache it fills itself. Without any window the
+    first decode token hashes stale slots (negative control)."""
+    block_size = 4
+    block_table = torch.tensor(
+        [[3, 1, 5, 0], [2, 4, 6, 7]], dtype=torch.int32, device="cuda"
+    )
+    tables_meta = (
+        torch.arange(50, dtype=torch.int32, device="cuda"),
+        torch.tensor([[3, 5, 7, 9]], device="cuda"),
+        torch.tensor([[[97, 89], [83, 79], [73, 71]]], device="cuda"),
+        torch.tensor([[0, 97, 186, 269, 348, 421]], device="cuda"),
+        block_size,
+    )
+    depth = 3
+    prompts = [[11, 12, 13, 14, 15, 16], [31, 32, 33, 34, 35]]
+    decodes = [[21, 22, 23, 24], [41, 42, 43, 44]]
+    chunk_sizes = [1, 2, 1]
+    v2 = runner.startswith("v2")
+
+    def fresh_cache():
+        return torch.zeros(8 * block_size, dtype=torch.int32, device="cuda")
+
+    reference = _hash_ids(
+        [(p + d, 0, None) for p, d in zip(prompts, decodes)],
+        tables_meta,
+        block_table,
+        cache=fresh_cache(),
+    )
+
+    histories = [list(p) for p in prompts]
+    cache = None if v2 else fresh_cache()
+    consumed = 0
+    for chunk in chunk_sizes:
+        batch = []
+        for history, prompt, decode in zip(histories, prompts, decodes):
+            start = len(history)
+            tokens = decode[consumed : consumed + chunk]
+            window = [
+                history[start - 1 - j] if v2 or start - 1 - j < len(prompt) else -1
+                for j in range(depth)
+            ]
+            batch.append((tokens, start, window))
+        got = _hash_ids(batch, tables_meta, block_table, cache=cache, capture=capture)
+        for req, (tokens, start, _) in enumerate(batch):
+            assert torch.equal(got[req], reference[req][start : start + chunk]), (
+                f"{runner} request {req} chunk at {start}"
+            )
+            histories[req] += tokens
+        consumed += chunk
+
+    unseeded = _hash_ids(
+        [(decodes[0][:1], len(prompts[0]), None)],
+        tables_meta,
+        block_table,
+        cache=fresh_cache(),
+    )
+    assert not torch.equal(unseeded[0][0], reference[0][len(prompts[0])])
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_v2_model_state_gathers_lookback_window():
+    """The window is gathered on device from the runner's token history."""
+    from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
+
+    depth, max_num_reqs, max_model_len = 3, 4, 16
+    state = DeepseekV41ModelState.__new__(DeepseekV41ModelState)
+    state.rope_state = None
+    state.supports_mm_inputs = False
+    state.prompt_embeds_state = None
+    state.lookback_token_ids = torch.full(
+        (max_num_reqs, depth), -1, dtype=torch.int32, device="cuda"
+    )
+
+    all_token_ids = torch.zeros(max_num_reqs, max_model_len, dtype=torch.int32)
+    all_token_ids[2, :8] = torch.tensor([11, 12, 13, 14, 15, 16, 21, 22])
+    all_token_ids[0, :5] = torch.arange(1, 6)
+    req_states = SimpleNamespace(
+        num_computed_tokens=SimpleNamespace(
+            gpu=torch.tensor([0, 9, 8, 6], dtype=torch.int32, device="cuda")
+        ),
+        all_token_ids=SimpleNamespace(gpu=all_token_ids.cuda()),
+    )
+    # Batch rows -> request state rows: a request two decode steps in, a fresh
+    # request, and one at the end of a 5-token prompt.
+    input_batch = SimpleNamespace(
+        idx_mapping=torch.tensor([2, 0, 3], dtype=torch.int64, device="cuda")
+    )
+    window = state.prepare_inputs(input_batch, req_states)["lookback_token_ids"]
+    assert window.cpu().tolist() == [
+        [22, 21, 16],
+        [-1, -1, -1],
+        [0, 0, 0],
+        [-1, -1, -1],
+    ]
+    # Graph capture runs on the dummy inputs and replays read the same buffer.
+    dummy = state.prepare_dummy_inputs(num_reqs=3, num_tokens=3)["lookback_token_ids"]
+    assert dummy is window and torch.all(dummy == -1)
+
+
+def _reference_lookup(weight, scale_inv, ids, start, end, block=32):
+    """The torch expression the fused kernel replaces."""
+    mask = (ids < start) | (ids >= end)
+    local = (ids - start).masked_fill(mask, 0).long()
+    values = torch.nn.functional.embedding(local, weight)
+    scales = torch.nn.functional.embedding(local, scale_inv)
+    scales = (scales.to(torch.int32) << 23).view(torch.float32)
+    values = values.float().unflatten(-1, (-1, block))
+    values = (values * scales.unsqueeze(-1)).flatten(-2).to(torch.bfloat16)
+    return values.masked_fill(mask.unsqueeze(-1), 0)
+
+
+def _make_embedding(cpu_offload, rows=4096, dim=256, block=32):
+    layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding)
+    torch.nn.Module.__init__(layer)
+    layer.dim, layer.block_size, layer.tp_size = dim, block, 1
+    layer.dp_size = 1
+    layer.n_hash_cols = layer.part_n_hash_cols = 24
+    layer.head_start = 0
+    layer.cpu_offload = cpu_offload
+    layer.dp_shared_memory = False
+    layer.part_num_embeddings = rows
+    # A window strictly inside the table, so unowned rows are exercised too.
+    layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2
+    layer._views = layer._view_src = None
+    layer._num_sms = torch.cuda.get_device_properties(
+        torch.accelerator.current_device_index()
+    ).multi_processor_count
+    kwargs = (
+        {"device": "cpu", "pin_memory": True} if cpu_offload else {"device": "cuda"}
+    )
+    owned = layer.vocab_end_idx - layer.vocab_start_idx
+    layer.weight = torch.empty(owned, dim, dtype=torch.float8_e4m3fn, **kwargs)
+    layer.weight_scale_inv = torch.empty(
+        owned, dim // block, dtype=torch.uint8, **kwargs
+    )
+    torch.manual_seed(0)
+    layer.weight.copy_((torch.randn(owned, dim) * 4).to(torch.float8_e4m3fn))
+    layer.weight_scale_inv.copy_(
+        torch.randint(120, 134, (owned, dim // block), dtype=torch.uint8)
+    )
+    return layer
+
+
+@pytest.mark.parametrize(
+    "tp_size,dp_size,n_heads", [(1, 4, 5), (2, 2, 5), (4, 1, 6), (8, 1, 6)]
+)
+def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch):
+    """Reject empty owners before allocating weights or accessing CUDA."""
+    monkeypatch.setattr(
+        engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    monkeypatch.setattr(nvidia_engram_ops, "get_engram_dp_size", lambda: dp_size)
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(nvidia_engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
+    with pytest.raises(AssertionError, match="ranks without hash heads"):
+        ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="CUDA or ROCm required"
+)
+@pytest.mark.parametrize("cpu_offload", [False, True])
+@pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
+def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch):
+    """Keep complete buckets and reconstruct head order, including TP padding."""
+    head_sizes = (17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73)
+    num_rows, dim = sum(head_sizes), 64
+    torch.manual_seed(0)
+    weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn)
+    scales = torch.randint(120, 134, (num_rows + 7, dim // 32), dtype=torch.uint8)
+    hashes = torch.empty(7, 2, len(head_sizes), dtype=torch.int32, device="cuda")
+    start = 0
+    for head, size in enumerate(head_sizes):
+        hashes[:, :, head].random_(start, start + size)
+        hashes[0, :, head] = start
+        hashes[-1, :, head] = start + size - 1
+        start += size
+    ids = hashes[:, 1]
+    expected = _reference_lookup(weight.cuda(), scales.cuda(), ids, 0, num_rows)
+    monkeypatch.setattr(
+        engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    shards = []
+    layers = []
+    for rank in range(tp_size):
+        monkeypatch.setattr(
+            engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank
+        )
+        with torch.device("cuda"):
+            layer = ParallelEngramEmbedding(
+                num_rows + 7, dim, head_sizes, cpu_offload=cpu_offload
+            )
+        layer.weight.weight_loader(layer.weight, weight)
+        layer.weight_scale_inv.weight_loader(
+            layer.weight_scale_inv, scales.view(torch.float8_e8m0fnu)
+        )
+        out = torch.empty(
+            len(ids),
+            layer.part_n_hash_cols,
+            dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        layer.lookup(ids, out)
+        shards.append(out)
+        layers.append(layer)
+    gathered = torch.cat(shards, dim=1)
+    torch.testing.assert_close(gathered[:, : len(head_sizes)], expected, rtol=0, atol=0)
+    assert torch.count_nonzero(gathered[:, len(head_sizes) :]) == 0
+    assert sum(layer.part_num_embeddings for layer in layers) == num_rows
+
+    def gather(local, dim):
+        torch.testing.assert_close(local, shards[0], rtol=0, atol=0)
+        return torch.cat(shards, dim=dim)
+
+    monkeypatch.setattr(engram_ops, "tensor_model_parallel_all_gather", gather)
+    torch.testing.assert_close(layers[0](ids), expected, rtol=0, atol=0)
+    module = Engram.__new__(Engram)
+    torch.nn.Module.__init__(module)
+    module.embed_tokens = layers[0]
+    module.use_sequence_parallel = False
+    module.staged_rows = torch.empty_like(shards[0])
+    module.prepare_embeddings(ids)
+    torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0)
+    # Slicing before head reordering must preserve padded heads and tokens.
+    module.use_sequence_parallel = True
+    chunk = (len(ids) + tp_size - 1) // tp_size
+    padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size))
+    for rank in range(tp_size):
+        monkeypatch.setattr(
+            engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank
+        )
+        torch.testing.assert_close(
+            module.embed(ids), padded[rank * chunk : (rank + 1) * chunk], rtol=0, atol=0
+        )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_engram_lookup_reuses_jit_across_token_shapes():
+    """Runtime token counts and launch grids must share one JIT variant."""
+    layer = _make_embedding(cpu_offload=False)
+    kernel = engram_ops._engram_lookup_kernel
+    kernel_cache = kernel.device_caches[torch.accelerator.current_device_index()][0]
+    kernel_cache.clear()
+
+    cache_sizes = []
+    for background in (False, True):
+        for num_tokens in (1, 7, 256):
+            ids = torch.zeros(num_tokens, 24, dtype=torch.int32, device="cuda")
+            out = torch.empty(
+                num_tokens,
+                24,
+                layer.dim,
+                dtype=torch.bfloat16,
+                device="cuda",
+            )
+            layer.lookup(ids, out, background=background)
+            cache_sizes.append(len(kernel_cache))
+
+    assert cache_sizes == [1] * len(cache_sizes)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="CUDA or ROCm required"
+)
+@pytest.mark.parametrize("cpu_offload", [False, True])
+@pytest.mark.parametrize("background", [False, True])
+@pytest.mark.parametrize("num_tokens", [1, 7, 256])
+def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
+    """The fused gather must be bit-exact with the torch dequant path it
+    replaces, from HBM and from pinned host memory alike, and must contribute
+    zeros for rows another TP rank owns."""
+    layer = _make_embedding(cpu_offload)
+    cols, rows = 24, layer.part_num_embeddings
+    ids = torch.randint(0, rows, (num_tokens, cols), dtype=torch.int32, device="cuda")
+    expected = _reference_lookup(
+        layer.weight.cuda(),
+        layer.weight_scale_inv.cuda(),
+        ids,
+        layer.vocab_start_idx,
+        layer.vocab_end_idx,
+    )
+    out = torch.empty(num_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda")
+    layer.lookup(ids, out, background=background)
+    assert torch.equal(out, expected)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "backend,cpu_offload", [("nvidia", None), ("nvidia", False), ("common", False)]
+)
+def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
+    """Default offload uses pinned storage and a side stream; False uses HBM."""
+    from vllm.config import EngramConfig
+
+    config = SimpleNamespace(hidden_size=16, hc_mult=1, rms_norm_eps=1e-6)
+    layout = SimpleNamespace(
+        num_embeddings=(72,),
+        head_dim=32,
+        max_ngram_size=2,
+        n_heads=24,
+        primes=(((3,) * 24,),),
+    )
+    vllm_config = SimpleNamespace(
+        engram_config=EngramConfig()
+        if cpu_offload is None
+        else EngramConfig(cpu_offload=cpu_offload),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+    )
+    offloaded = backend == "nvidia" and vllm_config.engram_config.cpu_offload
+    if backend == "common":
+        vllm_config.engram_config = None
+    monkeypatch.setattr(engram_ops, "get_current_vllm_config", lambda: vllm_config)
+    monkeypatch.setattr(
+        nvidia_engram_ops, "get_current_vllm_config", lambda: vllm_config
+    )
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        engram_ops, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity()
+    )
+    with torch.device("cuda"):
+        cls = Engram if backend == "nvidia" else CommonEngram
+        module = cls(config, None, layout, 0, False, "engram")
+    layer = module.embed_tokens
+    if offloaded:
+        assert layer.weight.device.type == "cpu" and layer.weight.is_pinned()
+    else:
+        assert layer.weight.device.type == "cuda"
+    layer.weight.data.copy_(torch.randn(layer.weight.shape).to(torch.float8_e4m3fn))
+    layer.weight_scale_inv.data.fill_(127)
+    streams = []
+    lookup = layer.lookup
+
+    def track_lookup(ids, out, background=False):
+        streams.append(torch.cuda.current_stream())
+        lookup(ids, out, background=background)
+
+    layer.lookup = track_lookup
+    main = torch.cuda.current_stream()
+    ids = torch.randint(
+        layer.vocab_start_idx,
+        layer.vocab_end_idx,
+        (8, 24),
+        device="cuda",
+        dtype=torch.int32,
+    )
+    module.prepare_embeddings(ids)
+    expected = _reference_lookup(
+        layer.weight.cuda(),
+        layer.weight_scale_inv.cuda(),
+        ids,
+        layer.vocab_start_idx,
+        layer.vocab_end_idx,
+    )
+    torch.testing.assert_close(module.embed(ids), expected, atol=0, rtol=0)
+    if offloaded:
+        assert streams == [module._prefetch_stream]
+        assert streams[0] != main
+    else:
+        assert streams == [main]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "cpu_offload,delay",
+    [(False, None), (True, None), (True, "producer"), (True, "lookup")],
+)
+@pytest.mark.parametrize("capture", ["eager", "full", "breakable"])
+def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture, delay):
+    """Temporary lookup IDs survive allocator reuse and graph replay."""
+    _run_engram_prepared_rows(cpu_offload, capture, delay=delay)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "missing_dependency", ["record_stream", "producer_wait", "consumer_wait"]
+)
+def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_dependency):
+    """Negative controls must expose incorrect rows when a dependency is removed."""
+    if (
+        missing_dependency == "record_stream"
+        and torch.cuda.memory.get_allocator_backend() != "native"
+    ):
+        pytest.skip(
+            "The allocator reuse negative control requires the native allocator"
+        )
+
+    def start(self, hash_ids, rows, stream):
+        if missing_dependency != "producer_wait":
+            stream.wait_stream(torch.cuda.current_stream())
+        if missing_dependency != "record_stream":
+            hash_ids.record_stream(stream)
+        with torch.cuda.stream(stream):
+            self.embed_tokens.lookup(hash_ids, rows, background=True)
+
+    def finish(self, stream):
+        pass
+
+    monkeypatch.setattr(Engram, "_start_prefetch", start)
+    if missing_dependency == "consumer_wait":
+        monkeypatch.setattr(Engram, "_finish_prefetch", finish)
+    delay = "producer" if missing_dependency == "producer_wait" else "lookup"
+    try:
+        with pytest.raises(
+            AssertionError, match=r"Tensor-likes are not (?:equal|close)!"
+        ):
+            _run_engram_prepared_rows(True, "eager", delay=delay)
+    finally:
+        torch.accelerator.synchronize()
+
+
+def _run_engram_prepared_rows(
+    cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False, delay=None
+):
+    from vllm.compilation.breakable_cudagraph import (
+        BreakableCUDAGraphCapture,
+        eager_break_during_capture,
+    )
+
+    layer = _make_embedding(cpu_offload)
+    if delay == "lookup":
+        lookup = layer.lookup
+
+        def delayed_lookup(indices, out, background=False):
+            torch.cuda._sleep(2_000_000)
+            lookup(indices, out, background=background)
+
+        layer.lookup = delayed_lookup
+    cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64)
+    layer.n_hash_cols = cols
+    layer.tp_size = tp_size
+    layer.part_n_hash_cols = (cols + tp_size - 1) // tp_size
+    layer.head_start = rank * layer.part_n_hash_cols
+    engram = Engram.__new__(Engram)
+    torch.nn.Module.__init__(engram)
+    engram.embed_tokens = layer
+    engram._prefetch_stream = torch.cuda.Stream() if cpu_offload else None
+    # Exercise the production eager boundaries even when the test process
+    # imported Engram before breakable graphs were enabled.
+    if capture == "breakable":
+        import vllm.envs as envs
+
+        with pytest.MonkeyPatch.context() as patch:
+            patch.setattr(envs, "VLLM_USE_BREAKABLE_CUDAGRAPH", True)
+            for name in ("_start_prefetch", "_finish_prefetch"):
+                fn = eager_break_during_capture(inspect.unwrap(getattr(Engram, name)))
+                setattr(engram, name, fn.__get__(engram, Engram))
+    engram.use_sequence_parallel = use_sequence_parallel
+    engram.staged_rows = torch.empty(
+        num_tokens,
+        layer.part_n_hash_cols,
+        layer.dim,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    # Match the non-contiguous per-layer slice of the model hash tensor.
+    hashes = torch.randint(
+        0,
+        layer.part_num_embeddings,
+        (num_tokens, 2, cols),
+        dtype=torch.int32,
+        device="cuda",
+    )
+    src = hashes[:, 1]
+    local_tokens = (
+        (num_tokens + tp_size - 1) // tp_size if use_sequence_parallel else num_tokens
+    )
+    out = torch.empty(
+        local_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda"
+    )
+    embed = engram.embed
+    if capture == "compiled":
+        embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True)
+
+    def step(cap=None):
+        if delay == "producer":
+            torch.cuda._sleep(2_000_000)
+        # Drop the last reference to a non-contiguous input after launching lookup.
+        engram.prepare_embeddings(hashes.clone()[:, 1])
+        # Exercise same-size allocator reuse before consuming the prefetched rows.
+        torch.empty_like(hashes).fill_(layer.part_num_embeddings - 1)
+        if cap is not None:
+            cap.add_eager(lambda: None)
+        out.copy_(embed(src))
+
+    warmup = torch.cuda.Stream()
+    warmup.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(warmup):
+        step()
+    torch.cuda.current_stream().wait_stream(warmup)
+
+    graph = None
+    if capture == "full":
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            step()
+    elif capture == "breakable":
+        graph = BreakableCUDAGraphCapture()
+        with torch.cuda.stream(warmup), graph:
+            step(graph)
+        torch.cuda.current_stream().wait_stream(warmup)
+        assert graph.num_eager_breaks == (3 if cpu_offload else 1)
+
+    for _ in range(3):
+        # Eager prefill can overwrite staging rows between decode replays.
+        hashes.random_(0, layer.part_num_embeddings)
+        step()
+        hashes.random_(0, layer.part_num_embeddings)
+        if graph is None:
+            step()
+        else:
+            graph.replay()
+        expected = _reference_lookup(
+            layer.weight.cuda(),
+            layer.weight_scale_inv.cuda(),
+            src,
+            layer.vocab_start_idx,
+            layer.vocab_end_idx,
+        )
+        if use_sequence_parallel:
+            expected = torch.nn.functional.pad(
+                expected, (0, 0, 0, 0, 0, (-num_tokens) % tp_size)
+            )[rank * local_tokens : (rank + 1) * local_tokens]
+        torch.testing.assert_close(out, expected, rtol=0, atol=0)
+
+
+def _engram_tp_worker(rank, tp_size, port):
+    from tests.utils import init_test_distributed_environment
+    from vllm.distributed import cleanup_dist_env_and_memory
+
+    torch.accelerator.set_device_index(rank)
+    init_test_distributed_environment(tp_size, 1, rank, str(port), local_rank=rank)
+    try:
+        for cpu_offload in (False, True):
+            for sp in (False, True):
+                for capture in ("eager", "compiled", "full", "breakable"):
+                    _run_engram_prepared_rows(cpu_offload, capture, tp_size, rank, sp)
+    finally:
+        cleanup_dist_env_and_memory()
+
+
+@pytest.mark.distributed(num_gpus=2)
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_engram_head_collectives_survive_graph_breaks():
+    """All-gather preserves head order and local SP tokens across graph replay."""
+    from vllm.utils.network_utils import get_open_port
+
+    if torch.accelerator.device_count() < 2:
+        pytest.skip("Requires two GPUs")
+    torch.multiprocessing.spawn(_engram_tp_worker, args=(2, get_open_port()), nprocs=2)
diff --git a/tests/kernels/test_fla_layernorm_guard.py b/tests/kernels/test_fla_layernorm_guard.py
index 492c099bf1c5..9223fb7110ff 100644
--- a/tests/kernels/test_fla_layernorm_guard.py
+++ b/tests/kernels/test_fla_layernorm_guard.py
@@ -501,6 +501,50 @@ def test_rmsnorm_gated_forward_native_dtype(
     torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
 
 
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@torch.inference_mode()
+def test_rmsnorm_gated_3d_matches_flattened_2d(
+    default_vllm_config,
+    dtype: torch.dtype,
+) -> None:
+    """Qwen GDN now norms (N, H, D) in place of a flatten to (N*H, D).
+
+    Last-dim RMS + gate must match the flattened layout on both the native
+    and CUDA/HIP kernels, including a 3D z view that is not contiguous.
+    """
+    from vllm.model_executor.layers.layernorm import RMSNormGated
+
+    set_random_seed(42)
+    device = torch.device(DEVICE)
+    num_tokens, num_heads, head_dim = 5, 4, 64
+    layer = RMSNormGated(
+        head_dim,
+        eps=1e-5,
+        group_size=None,
+        norm_before_gate=True,
+        device=device,
+        dtype=dtype,
+    )
+
+    packed = torch.randn(num_tokens, num_heads, 2, head_dim, dtype=dtype, device=device)
+    x3 = packed[:, :, 0, :]
+    z3 = packed[:, :, 1, :]
+    assert not x3.is_contiguous()
+    assert not z3.is_contiguous()
+    x2 = x3.reshape(-1, head_dim).contiguous()
+    z2 = z3.reshape(-1, head_dim).contiguous()
+
+    native_3d = layer.forward_native(x3, z3)
+    native_2d = layer.forward_native(x2, z2).reshape(num_tokens, num_heads, head_dim)
+    torch.testing.assert_close(native_3d, native_2d, atol=0, rtol=0)
+
+    cuda_3d = layer.forward_cuda(x3, z3)
+    cuda_2d = layer.forward_cuda(x2, z2).reshape(num_tokens, num_heads, head_dim)
+    torch.testing.assert_close(cuda_3d, cuda_2d, atol=1e-3, rtol=1e-2)
+    torch.testing.assert_close(cuda_3d, native_3d, atol=1e-2, rtol=1e-2)
+    assert cuda_3d.shape == (num_tokens, num_heads, head_dim)
+
+
 if __name__ == "__main__":
     # Run a quick smoke test
     test_layer_norm_fwd_basic(128, 1024, torch.float16, 42, False)
diff --git a/tests/kernels/test_flex_attention.py b/tests/kernels/test_flex_attention.py
index 12d00ec4054b..6862f8ffb181 100644
--- a/tests/kernels/test_flex_attention.py
+++ b/tests/kernels/test_flex_attention.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Integration tests for FlexAttention backend vs default backend"""
+"""Integration tests for FlexAttention backend vs default backend."""
 
 import pytest
 import torch
diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py
index d0562f6d2d25..df9d45a99db1 100644
--- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py
+++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py
@@ -24,6 +24,8 @@
 rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses).
 """
 
+from typing import NamedTuple
+
 import pytest
 import torch
 
@@ -33,19 +35,30 @@
 FP8 = torch.float8_e4m3fn
 FP8_MAX = 448.0
 
-# GLM-5.2 / DeepSeek-V3.2 shapes (TP8 local heads).
-Q_LORA = 2048
-KV_LORA = 512
-ROPE_DIM = 64
-NUM_HEADS = 8
-INDEX_HEADS = 32
-INDEX_HEAD_DIM = 128
-HIDDEN = 6144
 EPS = 1e-6
 
+
+class ModelConfig(NamedTuple):
+    """Per-model attention dims (num_heads is the TP8 local count)."""
+
+    q_lora: int
+    kv_lora: int
+    rope_dim: int
+    num_heads: int
+    index_heads: int
+    index_head_dim: int
+    hidden: int
+
+
+MODEL_CONFIGS = [
+    ModelConfig(2048, 512, 64, 8, 32, 128, 6144),  # GLM-5.2
+    ModelConfig(1536, 512, 64, 16, 64, 128, 7168),  # DeepSeek-V3.2
+]
+MODEL_IDS = ["glm5.2", "dsv3.2"]
+
 pytestmark = pytest.mark.skipif(
-    not current_platform.is_cuda() or not current_platform.has_device_capability(89),
-    reason="deepseek_v32 fused kernels require CUDA with fp8 (SM89+)",
+    not current_platform.supports_fp8(),
+    reason="deepseek_v32 fused kernels require fp8 support",
 )
 
 
@@ -144,29 +157,32 @@ def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str):
 # ── fused_norm_rope ──────────────────────────────────────────────────────────
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
 @pytest.mark.parametrize("index_interleave", [True, False])
 @pytest.mark.parametrize("mla_dtype", ["auto", "bfloat16", "fp8"])
-def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str):
+def test_fused_norm_rope(
+    num_tokens: int, index_interleave: bool, mla_dtype: str, cfg: ModelConfig
+):
     torch.manual_seed(0)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
-    ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
-    ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
+    ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
 
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)  # MLA k_pe: interleaved
-    idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)  # MLA k_pe: interleaved
+    idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     bs = max_pos  # single block covering all tokens
-    mla_dim = KV_LORA + ROPE_DIM
+    mla_dim = cfg.kv_lora + cfg.rope_dim
     mla_fp8 = mla_dtype == "fp8"
     if mla_fp8:
         mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8)
@@ -174,7 +190,7 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
     else:
         mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16)
         mla_k_scale = None
-    idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4  # 132
+    idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4  # 132
     idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8)
     slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
     topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
@@ -213,20 +229,20 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
     if mla_fp8:
         cache = mla_cache.view(FP8)[0, :num_tokens]
         s = mla_k_scale.item()
-        assert_fp8(cache[:, :KV_LORA], (kv_ref / s).to(FP8), "MLA kv fp8")
-        assert_fp8(cache[:, KV_LORA:], (kpe_ref / s).to(FP8), "MLA k_pe fp8")
+        assert_fp8(cache[:, : cfg.kv_lora], (kv_ref / s).to(FP8), "MLA kv fp8")
+        assert_fp8(cache[:, cfg.kv_lora :], (kpe_ref / s).to(FP8), "MLA k_pe fp8")
     else:
         cache = mla_cache[0, :num_tokens]
-        assert_bf16(cache[:, :KV_LORA], kv_ref, "MLA kv bf16")
-        assert_bf16(cache[:, KV_LORA:], kpe_ref, "MLA k_pe bf16")
+        assert_bf16(cache[:, : cfg.kv_lora], kv_ref, "MLA kv bf16")
+        assert_bf16(cache[:, cfg.kv_lora :], kpe_ref, "MLA k_pe bf16")
 
     # Indexer-K cache (packed [bs*head_dim fp8 | bs*4 fp32 scale]).
     ik_ref = layer_norm(ik, ikw, ikb)
     ik_ref = rope(ik_ref, pos, idx_cos_sin, interleave=index_interleave)
     q_ref, s_ref = ue8m0_quant(ik_ref)
     flat = idx_cache[0].reshape(-1)
-    vals = flat[: bs * INDEX_HEAD_DIM].view(FP8).reshape(bs, INDEX_HEAD_DIM)
-    scales = flat[bs * INDEX_HEAD_DIM :].view(torch.float32)
+    vals = flat[: bs * cfg.index_head_dim].view(FP8).reshape(bs, cfg.index_head_dim)
+    scales = flat[bs * cfg.index_head_dim :].view(torch.float32)
     assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8")
     torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0)
 
@@ -234,7 +250,83 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
     assert (topk == -1).all(), "topk buffer not cleared on indexer layer"
 
 
-def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
+def test_fused_norm_rope_packed_indexer_block_stride(cfg: ModelConfig):
+    """Indexer writes use their HMA block stride and independent slot map."""
+    torch.manual_seed(7)
+    dev = "cuda"
+    num_tokens = block_size = 2
+    pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
+    ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    cos_sin = make_cos_sin(32, cfg.rope_dim, dev)
+
+    idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4
+    packed_block_stride = block_size * idx_row + 64
+    backing = torch.zeros(2 * packed_block_stride, device=dev, dtype=torch.uint8)
+    idx_cache = torch.as_strided(
+        backing,
+        (2, block_size, idx_row),
+        (packed_block_stride, idx_row, 1),
+    )
+    mla_cache = torch.zeros(
+        1,
+        block_size,
+        cfg.kv_lora + cfg.rope_dim,
+        device=dev,
+        dtype=torch.bfloat16,
+    )
+    topk = torch.zeros(num_tokens, 8, device=dev, dtype=torch.int32)
+
+    K.fused_norm_rope(
+        pos,
+        q_c,
+        qw,
+        EPS,
+        kv_c,
+        kvw,
+        EPS,
+        k_pe,
+        cos_sin,
+        ik,
+        ikw,
+        ikb,
+        EPS,
+        cos_sin,
+        topk,
+        slot_mapping=torch.arange(num_tokens, device=dev, dtype=torch.int64),
+        indexer_slot_mapping=torch.arange(
+            block_size, block_size + num_tokens, device=dev, dtype=torch.int64
+        ),
+        indexer_k_cache=idx_cache,
+        mla_kv_cache=mla_cache,
+        has_indexer=True,
+    )
+
+    ik_ref = rope(layer_norm(ik, ikw, ikb), pos, cos_sin, interleave=False)
+    q_ref, s_ref = ue8m0_quant(ik_ref)
+    packed = idx_cache[1].reshape(-1)
+    values = (
+        packed[: block_size * cfg.index_head_dim]
+        .view(FP8)
+        .view(block_size, cfg.index_head_dim)
+    )
+    scales = packed[block_size * cfg.index_head_dim :].view(torch.float32)
+    assert_fp8(values, q_ref, "packed indexer-K fp8")
+    torch.testing.assert_close(scales, s_ref, rtol=0, atol=0)
+    assert (backing[block_size * idx_row : packed_block_stride] == 0).all()
+
+
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
+def test_fused_norm_rope_materializes_qk_without_local_cache_slots(
+    cfg: ModelConfig,
+):
     """DCP non-owner ranks still need valid query shards for query AllGather."""
     torch.manual_seed(7)
     dev = "cuda"
@@ -242,20 +334,20 @@ def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
     max_pos = 16
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
-    ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
-    ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
-    cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
+    ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     mla_cache = torch.zeros(
-        1, max_pos, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16
+        1, max_pos, cfg.kv_lora + cfg.rope_dim, device=dev, dtype=torch.bfloat16
     )
-    idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
+    idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4
     idx_cache = torch.zeros(1, max_pos, idx_row, device=dev, dtype=torch.uint8)
     no_local_slots = torch.full((num_tokens,), -1, device=dev, dtype=torch.int64)
     topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
@@ -301,23 +393,26 @@ def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
     assert (topk == -1).all(), "topk buffer not cleared on non-owner rank"
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 17, 512])
-def test_fused_norm_rope_no_indexer(num_tokens: int):
+def test_fused_norm_rope_no_indexer(num_tokens: int, cfg: ModelConfig):
     """Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched."""
     torch.manual_seed(1)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     bs = max_pos
-    mla_cache = torch.zeros(1, bs, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16)
+    mla_cache = torch.zeros(
+        1, bs, cfg.kv_lora + cfg.rope_dim, device=dev, dtype=torch.bfloat16
+    )
     slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
     topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
 
@@ -348,9 +443,9 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
 
     assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (no-indexer)")
     cache = mla_cache[0, :num_tokens]
-    assert_bf16(cache[:, :KV_LORA], rms_norm(kv_c, kvw), "MLA kv (no-indexer)")
+    assert_bf16(cache[:, : cfg.kv_lora], rms_norm(kv_c, kvw), "MLA kv (no-indexer)")
     assert_bf16(
-        cache[:, KV_LORA:],
+        cache[:, cfg.kv_lora :],
         rope(k_pe.float(), pos, mla_cos_sin, interleave=True),
         "MLA k_pe (no-indexer)",
     )
@@ -358,36 +453,39 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
     assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("has_indexer", [False, True])
-def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool):
+def test_fused_norm_rope_materializes_pcp_cache_inputs(
+    has_indexer: bool, cfg: ModelConfig
+):
     """PCP gets local normalized/rotated K rows without direct cache writes."""
     torch.manual_seed(6)
     dev = "cuda"
     num_tokens = 17
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
     ik = (
-        torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
+        torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
         if has_indexer
         else None
     )
     ikw = (
-        torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
+        torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
         if has_indexer
         else None
     )
     ikb = (
-        torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
+        torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
         if has_indexer
         else None
     )
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
-    idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) if has_indexer else None
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
+    idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) if has_indexer else None
     q_out = torch.empty_like(q_c)
     kv_out = torch.empty_like(kv_c)
     kpe_out = torch.empty_like(k_pe)
@@ -433,24 +531,25 @@ def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool):
         assert_bf16(ik_out, ik_ref, "PCP indexer-K")
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
-def test_fused_norm_rope_ds_mla(num_tokens: int):
+def test_fused_norm_rope_ds_mla(num_tokens: int, cfg: ModelConfig):
     """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).
 
-    Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, dynamic float32
-    scale) | 4 float32 scales | 64 bf16 (unquantized) RoPE.
+    Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, power-of-two scales
+    stored as float32) | 4 float32 scales | 64 bf16 (unquantized) RoPE.
     """
     torch.manual_seed(5)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     bs = max_pos
     mla_cache = torch.zeros(1, bs, 656, device=dev, dtype=torch.uint8)
@@ -487,22 +586,144 @@ def test_fused_norm_rope_ds_mla(num_tokens: int):
     kv_ref = rms_norm(kv_c, kvw)  # [N, 512] fp32
     kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True)  # [N, 64]
     tiles = kv_ref.view(num_tokens, 4, 128)
-    ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1.1754944e-38)
-    ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8)
+    ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1e-4)
+    ref_scale = torch.exp2(torch.ceil(torch.log2(ref_scale)))
+    ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, cfg.kv_lora).to(FP8)
 
     cache = mla_cache[0, :num_tokens]  # [N, 656] uint8
-    nope = cache[:, :KV_LORA].view(FP8)
-    scales = cache.view(torch.float32)[:, KV_LORA // 4 : KV_LORA // 4 + 4]
-    rope_off = KV_LORA // 2 + 8
-    rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM]
+    nope = cache[:, : cfg.kv_lora].view(FP8)
+    scales = cache.view(torch.float32)[:, cfg.kv_lora // 4 : cfg.kv_lora // 4 + 4]
+    rope_off = cfg.kv_lora // 2 + 8
+    rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + cfg.rope_dim]
 
-    torch.testing.assert_close(scales, ref_scale, rtol=1e-2, atol=1e-6)
+    torch.testing.assert_close(scales, ref_scale, rtol=0, atol=0)
     assert_fp8(nope, ref_nope, "ds_mla NoPE fp8")
     assert_bf16(rope_vals, kpe_ref, "ds_mla RoPE bf16")
     # No indexer on this call: top-k buffer must be untouched.
     assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"
 
 
+E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
+
+
+def quantize_to_e2m1(x: torch.Tensor) -> torch.Tensor:
+    """Round to nearest e2m1, saturating to +-6 like cvt.rn.satfinite.e2m1x2.f32."""
+    mags = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x.device)
+    a = x.float().abs().clamp_max(6.0)
+    mids = (mags[:-1] + mags[1:]) / 2
+    code = torch.bucketize(a, mids, right=True)
+    on_tie = (a.unsqueeze(-1) == mids).any(dim=-1)
+    tie_code = torch.bucketize(a, mids, right=False)
+    code = torch.where(on_tie, tie_code + (tie_code & 1), code)
+    return (torch.signbit(x).to(torch.uint8) << 3) | code.to(torch.uint8)
+
+
+def nvfp4_sf_byte(s: torch.Tensor) -> torch.Tensor:
+    """Scale-factor byte permutation shared with FlashMLA: 8*(s&3) + (s>>2)."""
+    return 8 * (s % 4) + (s // 4)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_device_capability(100),
+    reason="nvfp4_ds_mla requires SM100 (Blackwell)",
+)
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
+def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int, cfg: ModelConfig):
+    """nvfp4_ds_mla MLA cache layout (FlashMLA sparse, SM100 only).
+
+    Per-token 352-byte entry: 256 B of 512 e2m1 NoPE packed 2/byte (low nibble
+    = even element) | 64 B unscaled e4m3 RoPE | 32 B byte-permuted e4m3 tile
+    scales, one per 16 NoPE elements.
+    """
+    torch.manual_seed(5)
+    dev = "cuda"
+    max_pos = 8192
+    pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
+
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    # Tile 0 is deliberately low-magnitude: amax/6 lands between two e4m3
+    # subnormals (spaced a flat 2^-9 there), so round-to-nearest picks a scale
+    # BELOW amax/6 and the tile's peak saturates at +-6. Asserted below, so a
+    # future change to the scale rule cannot silently stop covering this.
+    # amax must land in (6, 9) * 2^-9 after rms_norm; 7.5 sits mid-band, with
+    # margin for the per-token RMS to vary.
+    kv_c[:, :16] = (
+        7.5 * (2.0**-9) * torch.linspace(0.2, 1.0, 16, device=dev, dtype=torch.bfloat16)
+    )
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.ones(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
+
+    bs = max_pos
+    mla_cache = torch.zeros(1, bs, 352, device=dev, dtype=torch.uint8)
+    slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
+    topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
+
+    q_out = K.fused_norm_rope(
+        pos,
+        q_c,
+        qw,
+        EPS,
+        kv_c,
+        kvw,
+        EPS,
+        k_pe,
+        mla_cos_sin,
+        None,
+        None,
+        None,
+        EPS,
+        None,
+        topk,
+        slot_mapping=slot,
+        indexer_k_cache=None,
+        mla_kv_cache=mla_cache,
+        mla_kv_cache_dtype="nvfp4_ds_mla",
+        mla_k_scale=None,
+        has_indexer=False,
+        index_rope_interleave=False,
+    )
+
+    assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (nvfp4_ds_mla)")
+
+    kv_ref = rms_norm(kv_c, kvw)  # [N, 512] fp32
+    kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True)  # [N, 64]
+    tiles = kv_ref.view(num_tokens, 32, 16)
+    amax = tiles.abs().amax(dim=-1)
+    scale_target = torch.clamp_min(amax / 6.0, 2.0**-9)
+    ref_scale = scale_target.to(FP8)  # round-to-nearest
+
+    # The crafted tile must actually exercise the saturating path.
+    assert (ref_scale[:, 0].float() < scale_target[:, 0]).all(), (
+        "tile 0 should round its scale DOWN; adjust the crafted magnitude"
+    )
+    assert (amax[:, 0] / ref_scale[:, 0].float() > 6.0).all(), (
+        "tile 0 should saturate e2m1; adjust the crafted magnitude"
+    )
+
+    ref_codes = quantize_to_e2m1(tiles / ref_scale.float().unsqueeze(-1))
+    ref_codes = ref_codes.reshape(num_tokens, cfg.kv_lora)
+
+    cache = mla_cache[0, :num_tokens]  # [N, 352] uint8
+    packed = cache[:, : cfg.kv_lora // 2]
+    got_codes = torch.stack([packed & 0xF, packed >> 4], dim=-1)
+    got_codes = got_codes.reshape(num_tokens, cfg.kv_lora)  # low nibble = even elem
+    torch.testing.assert_close(got_codes, ref_codes, rtol=0, atol=0)
+
+    got_rope = cache[:, cfg.kv_lora // 2 : cfg.kv_lora // 2 + cfg.rope_dim].view(FP8)
+    assert_fp8(got_rope, kpe_ref.to(FP8), "nvfp4_ds_mla RoPE e4m3")
+
+    perm = nvfp4_sf_byte(torch.arange(32, device=dev))
+    got_scale = cache[:, cfg.kv_lora // 2 + cfg.rope_dim :].view(FP8)[:, perm]
+    torch.testing.assert_close(got_scale.float(), ref_scale.float(), rtol=0, atol=0)
+
+    # No indexer on this call: top-k buffer must be untouched.
+    assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"
+
+
 def test_fused_norm_rope_supports_large_token_count():
     """Keep the token count off CUDA grid-y at its 65,536-block boundary."""
     num_tokens = 65536
@@ -546,39 +767,34 @@ def test_fused_norm_rope_supports_large_token_count():
 # ── fused_q ──────────────────────────────────────────────────────────────────
 
 
-@pytest.mark.parametrize(
-    ("num_tokens", "num_q_heads"),
-    [
-        (1, NUM_HEADS),
-        (4, NUM_HEADS),
-        (17, NUM_HEADS),
-        (512, NUM_HEADS),
-        (4096, NUM_HEADS),
-        (17, 64),
-    ],
-)
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
 @pytest.mark.parametrize("index_interleave", [True, False])
-def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
+def test_fused_q(num_tokens: int, index_interleave: bool, cfg: ModelConfig):
     torch.manual_seed(2)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
     q_pe = torch.randn(
-        num_tokens, num_q_heads, ROPE_DIM, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
     )
     ql_nope = torch.randn(
-        num_tokens, num_q_heads, KV_LORA, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
     )
     index_q = torch.randn(
-        num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
+        num_tokens,
+        cfg.index_heads,
+        cfg.index_head_dim,
+        device=dev,
+        dtype=torch.bfloat16,
     )
-    index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
+    index_w = torch.randn(num_tokens, cfg.index_heads, device=dev, dtype=torch.float32)
     q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
-    softmax_scale = INDEX_HEAD_DIM**-0.5
-    head_scale = INDEX_HEADS**-0.5
-    q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)  # q_pe: interleaved
-    idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    softmax_scale = cfg.index_head_dim**-0.5
+    head_scale = cfg.index_heads**-0.5
+    q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)  # q_pe: interleaved
+    idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     iq_fp8, iw_out, mqa = K.fused_q(
         pos,
@@ -600,18 +816,18 @@ def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
     mqa_nope_ref = (ql_nope.float() / s).to(FP8)
     qpe_ref = rope(
         q_pe.float(),
-        pos.unsqueeze(-1).expand(num_tokens, num_q_heads),
+        pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
         q_cos_sin,
         interleave=True,
     )
     mqa_pe_ref = (qpe_ref / s).to(FP8)
-    assert_fp8(mqa[:, :, :KV_LORA], mqa_nope_ref, "mqa ql_nope")
-    assert_fp8(mqa[:, :, KV_LORA:], mqa_pe_ref, "mqa q_pe")
+    assert_fp8(mqa[:, :, : cfg.kv_lora], mqa_nope_ref, "mqa ql_nope")
+    assert_fp8(mqa[:, :, cfg.kv_lora :], mqa_pe_ref, "mqa q_pe")
 
     # Indexer-Q: RoPE + UE8M0 fp8 quant; index weights fold in q-scale.
     iq_ref = rope(
         index_q.float(),
-        pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
+        pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads),
         idx_cos_sin,
         interleave=index_interleave,
     )
@@ -621,20 +837,21 @@ def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
     torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 17, 512])
-def test_fused_q_no_indexer(num_tokens: int):
+def test_fused_q_no_indexer(num_tokens: int, cfg: ModelConfig):
     torch.manual_seed(3)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
     q_pe = torch.randn(
-        num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
     )
     ql_nope = torch.randn(
-        num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
     )
     q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32)
-    q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
     _, _, mqa = K.fused_q(
         pos,
         q_pe,
@@ -650,19 +867,20 @@ def test_fused_q_no_indexer(num_tokens: int):
         index_rope_interleave=False,
     )
     s = q_scale.item()
-    assert_fp8(mqa[:, :, :KV_LORA], (ql_nope.float() / s).to(FP8), "mqa ql_nope")
+    assert_fp8(mqa[:, :, : cfg.kv_lora], (ql_nope.float() / s).to(FP8), "mqa ql_nope")
     qpe_ref = rope(
         q_pe.float(),
-        pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
+        pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
         q_cos_sin,
         interleave=True,
     )
-    assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe")
+    assert_fp8(mqa[:, :, cfg.kv_lora :], (qpe_ref / s).to(FP8), "mqa q_pe")
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 17, 512])
 @pytest.mark.parametrize("has_indexer", [True, False])
-def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
+def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool, cfg: ModelConfig):
     """bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is
     produced (bf16, unquantized); ql_nope is consumed directly by the caller."""
     torch.manual_seed(6)
@@ -671,21 +889,27 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
     q_pe = torch.randn(
-        num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
     )
     ql_nope = torch.randn(
-        num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
     )
     q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
-    q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     index_q = index_w = idx_cos_sin = None
     if has_indexer:
         index_q = torch.randn(
-            num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
+            num_tokens,
+            cfg.index_heads,
+            cfg.index_head_dim,
+            device=dev,
+            dtype=torch.bfloat16,
+        )
+        index_w = torch.randn(
+            num_tokens, cfg.index_heads, device=dev, dtype=torch.float32
         )
-        index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
-        idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+        idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     iq_fp8, iw_out, q_pe_out = K.fused_q(
         pos,
@@ -696,8 +920,8 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
         ql_nope,
         q_scale,
         index_w,
-        INDEX_HEAD_DIM**-0.5,
-        INDEX_HEADS**-0.5,
+        cfg.index_head_dim**-0.5,
+        cfg.index_heads**-0.5,
         has_indexer=has_indexer,
         index_rope_interleave=False,
         quantize_mqa=False,
@@ -705,10 +929,10 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
 
     # MQA query: only the RoPE'd q_pe, bf16, unquantized.
     assert q_pe_out.dtype == torch.bfloat16
-    assert q_pe_out.shape == (num_tokens, NUM_HEADS, ROPE_DIM)
+    assert q_pe_out.shape == (num_tokens, cfg.num_heads, cfg.rope_dim)
     qpe_ref = rope(
         q_pe.float(),
-        pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
+        pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
         q_cos_sin,
         interleave=True,
     )
@@ -719,13 +943,15 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
         assert index_q is not None
         iq_ref = rope(
             index_q.float(),
-            pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
+            pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads),
             idx_cos_sin,
             interleave=False,
         )
         q_ref, scale_ref = ue8m0_quant(iq_ref)
         assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)")
-        iw_ref = index_w * scale_ref * (INDEX_HEAD_DIM**-0.5) * (INDEX_HEADS**-0.5)
+        iw_ref = (
+            index_w * scale_ref * (cfg.index_head_dim**-0.5) * (cfg.index_heads**-0.5)
+        )
         torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
 
 
@@ -765,21 +991,22 @@ def test_fused_q_triton_supports_large_token_count():
 # ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
-def test_fused_eh_norm(num_tokens: int):
+def test_fused_eh_norm(num_tokens: int, cfg: ModelConfig):
     torch.manual_seed(4)
     dev = "cuda"
     # Mix in a position-0 token to exercise the embeds-zeroing branch.
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
     pos[0] = 0
-    embeds = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16)
-    prev = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16)
-    ew = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16)
-    hw = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16)
+    embeds = torch.randn(num_tokens, cfg.hidden, device=dev, dtype=torch.bfloat16)
+    prev = torch.randn(num_tokens, cfg.hidden, device=dev, dtype=torch.bfloat16)
+    ew = torch.randn(cfg.hidden, device=dev, dtype=torch.bfloat16)
+    hw = torch.randn(cfg.hidden, device=dev, dtype=torch.bfloat16)
 
     out = K.fused_eh_norm(pos, embeds, prev, ew, hw, EPS)
 
     masked = torch.where(pos.unsqueeze(-1) == 0, torch.zeros_like(embeds), embeds)
     ref = torch.cat([rms_norm(masked, ew), rms_norm(prev, hw)], dim=-1)
-    assert out.shape == (num_tokens, 2 * HIDDEN)
+    assert out.shape == (num_tokens, 2 * cfg.hidden)
     assert_bf16(out, ref, "eh_norm")
diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
index 0e25884368c1..c41af1e4b9ca 100644
--- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel:
+"""Standalone unit test for the horizontally-fused DeepseekV4-MLA kernel:
 
   fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
     - Q side:  optional per-head RMSNorm + GPT-J RoPE on last 64 dims
@@ -97,7 +96,7 @@ def apply_rope_gptj_last_k(
     # `e*c - o*s` fused form. This keeps the reference close to the kernel, but
     # the fp32 reference and the fp32 GPU kernel can still round to bf16 on
     # opposite sides of a round-to-nearest tie for a tiny number of elements at
-    # high positions, so callers compare the RoPE region within 1 bf16 ULP.
+    # high positions, so callers allow for bf16 rounding in the RoPE region.
     new_even = torch.addcmul(-odd * sin, even, cos)
     new_odd = torch.addcmul(odd * cos, even, sin)
     rope_rotated = torch.stack((new_even, new_odd), dim=-1).reshape(shape)
@@ -315,7 +314,19 @@ def test_q_path_without_qnorm_matches_rope_only_reference(num_tokens: int):
         apply_q_norm=False,
     )
 
-    torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=0, atol=0)
+    torch.testing.assert_close(
+        q_out[:, :n_heads, :NOPE_DIM], q_ref[..., :NOPE_DIM], rtol=0, atol=0
+    )
+    # HIP and PyTorch can contract different FP32 products into an FMA.
+    # Allow bf16 rounding and small cancellation errors only in the RoPE region
+    # on ROCm; the unrotated dimensions and padded heads must remain exact.
+    rope_rtol, rope_atol = (1e-2, 1e-6) if current_platform.is_rocm() else (0, 0)
+    torch.testing.assert_close(
+        q_out[:, :n_heads, NOPE_DIM:],
+        q_ref[..., NOPE_DIM:],
+        rtol=rope_rtol,
+        atol=rope_atol,
+    )
     assert q_out[:, n_heads:].count_nonzero().item() == 0
 
 
@@ -424,6 +435,287 @@ def _dequant(k_cache_2d):
     )
 
 
+# ── Test 2c: DeepSeek V4.1 MXFP8 record (kv_mxfp8=True) ──────────────────────
+
+V41_HEAD_BYTES = HEAD_DIM + HEAD_DIM // 32  # 512 fp8 + 16 UE8M0 scales = 528
+
+
+@pytest.mark.parametrize("num_tokens", [1, 17, 64])
+@pytest.mark.parametrize("block_size", [32, 64])
+def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int):
+    """``kv_mxfp8=True`` writes DeepSeek's V4.1 record: all 512 dims as fp8 e4m3
+    with one UE8M0 scale per 32 dims, no bf16 RoPE tail.
+
+    Both sides quantize the same bf16-rounded rotated row, so the bytes must
+    match the Triton encoder exactly apart from the documented 1-ULP bf16
+    rounding split on the rotation itself.
+    """
+    from vllm.models.deepseek_v41.common.ops import (
+        dequantize_and_gather_k_cache as v41_dequantize_and_gather_k_cache,
+    )
+    from vllm.models.deepseek_v41.common.ops import (
+        quantize_and_insert_k_cache as v41_quantize_and_insert_k_cache,
+    )
+
+    torch.manual_seed(3)
+    device = "cuda"
+    dtype = torch.bfloat16
+    max_pos = 4096
+
+    kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
+
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
+
+    k_cache_ref = torch.zeros(
+        num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device
+    )
+    v41_quantize_and_insert_k_cache(
+        apply_rope_gptj_last_k(kv, positions, cos_sin_cache),
+        k_cache_ref,
+        slot_mapping,
+        block_size=block_size,
+        use_fnuz=USE_FNUZ,
+        bytes_per_token=V41_HEAD_BYTES,
+    )
+
+    k_cache_fused = torch.zeros_like(k_cache_ref)
+    q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device)
+    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        q_dummy,
+        kv,
+        k_cache_fused,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        64,
+        1e-6,
+        block_size,
+        True,
+        True,  # kv_mxfp8
+    )
+
+    # A pre-quantization bf16 tie can land the two encoders on adjacent e4m3
+    # codes, so compare after dequantizing and allow one step at the tile scale.
+    def _dequant(cache_2d):
+        out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=dtype, device=device)
+        v41_dequantize_and_gather_k_cache(
+            out,
+            cache_2d.view(num_blocks, block_size, V41_HEAD_BYTES),
+            torch.tensor([num_tokens], dtype=torch.int32, device=device),
+            None,
+            torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(0),
+            block_size,
+            offset=0,
+            use_fnuz=USE_FNUZ,
+        )
+        return out[0]
+
+    rec_ref = _dequant(k_cache_ref).float()
+    rec_fused = _dequant(k_cache_fused).float()
+    tile_scale = (
+        (rec_ref.abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) / FP8_MAX)
+        .log2()
+        .ceil()
+        .exp2()
+    )
+    # One e4m3 ULP at the top of the range is 32 quantized units.
+    step = 32.0 * tile_scale.repeat_interleave(32, dim=-1)
+    assert ((rec_fused - rec_ref).abs() <= step).all()
+    # The NoPE dims see no rotation at all, so they must be byte-identical.
+    nope_bytes = slice(0, block_size * HEAD_DIM)
+    ref_rows = k_cache_ref[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM)
+    fused_rows = k_cache_fused[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM)
+    torch.testing.assert_close(
+        fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0
+    )
+
+
+# ── Test 2d: mega-attention Q layout (is_q_interleaved) ──────────────────────
+
+Q_CHUNK = 16
+NUM_Q_CHUNKS = HEAD_DIM // Q_CHUNK  # 32
+
+
+@pytest.mark.parametrize("num_tokens", [1, 17, 2048])
+@pytest.mark.parametrize("n_heads,q_head_padded", [(16, 64), (64, 128), (64, 0)])
+def test_q_interleaved_pads_without_rope(
+    num_tokens: int, n_heads: int, q_head_padded: int
+):
+    """``is_q_interleaved`` with no norm or RoPE only zero-pads Q.
+
+    The mega-attention kernel norms and rotates Q itself, so every live
+    element must survive bit-exact while the padding heads -- which interleave
+    with the live ones inside each of the 32 head-dim chunks -- come back
+    zero. ``q_head_padded=0`` drops the Q pass entirely and returns nothing,
+    for a shard already at the kernel's head count.
+    """
+    torch.manual_seed(5)
+    device = "cuda"
+    dtype = torch.bfloat16
+    block_size = 64
+
+    q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
+    kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    k_cache = torch.zeros(
+        num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device
+    )
+
+    q_out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        q,
+        kv,
+        k_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        q_head_padded,
+        1e-6,
+        block_size,
+        False,  # apply_q_norm
+        True,  # kv_mxfp8
+        False,  # apply_q_rope
+        True,  # is_q_interleaved
+    )
+
+    # The KV half is unaffected by the Q mode: it must still be inserted.
+    assert k_cache.any()
+
+    if q_head_padded == 0:
+        assert q_out.numel() == 0
+        return
+    assert q_out.shape == (num_tokens, q_head_padded, HEAD_DIM)
+    got = q_out.view(num_tokens, NUM_Q_CHUNKS, q_head_padded, Q_CHUNK)
+    want = q.view(num_tokens, NUM_Q_CHUNKS, n_heads, Q_CHUNK)
+    torch.testing.assert_close(got[:, :, :n_heads], want, rtol=0, atol=0)
+    assert (got[:, :, n_heads:] == 0).all()
+
+
+@pytest.mark.parametrize("n_heads", [8, 16, 32])
+def test_q_interleaved_survives_graph_replay(n_heads: int):
+    """A captured fused insert refills its padded Q buffer on every replay.
+
+    The op allocates ``q_out`` itself, so the tensor a capture hands back has
+    to keep receiving the new padding when the graph is replayed over changed
+    input -- that is what lets the layer prepare Q inside the captured region
+    rather than in ``forward_mqa``.
+    """
+    torch.manual_seed(6)
+    device = "cuda"
+    dtype = torch.bfloat16
+    block_size, num_tokens, q_head_padded = 64, 7, 64
+
+    q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
+    kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    k_cache = torch.zeros(
+        num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device
+    )
+
+    def prepare():
+        return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+            q,
+            kv,
+            k_cache,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            q_head_padded,
+            1e-6,
+            block_size,
+            False,  # apply_q_norm
+            True,  # kv_mxfp8
+            False,  # apply_q_rope
+            True,  # is_q_interleaved
+        )
+
+    # Capture takes its warmup on a side stream.
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        for _ in range(3):
+            prepare()
+    torch.cuda.current_stream().wait_stream(side)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        q_out = prepare()
+
+    for factor in (0.25, -0.75):
+        q.copy_(torch.randn_like(q) * factor)
+        graph.replay()
+        torch.accelerator.synchronize()
+        got = q_out.view(num_tokens, NUM_Q_CHUNKS, q_head_padded, Q_CHUNK)
+        want = q.view(num_tokens, NUM_Q_CHUNKS, n_heads, Q_CHUNK)
+        torch.testing.assert_close(got[:, :, :n_heads], want, rtol=0, atol=0)
+        assert (got[:, :, n_heads:] == 0).all()
+
+
+@pytest.mark.parametrize("n_heads", [16, 64])
+@pytest.mark.parametrize("apply_q_norm", [False, True])
+@pytest.mark.parametrize("apply_q_rope", [False, True])
+def test_q_interleaved_is_orthogonal_to_norm_and_rope(
+    n_heads: int, apply_q_norm: bool, apply_q_rope: bool
+):
+    """The Q layout composes with whatever norm and RoPE it is given.
+
+    A lane owns the same 512-dim slice of its head in either layout -- only
+    the address differs -- so the head-major result permuted into the
+    interleaved layout must equal running the interleaved path on a permuted
+    input, bit for bit, for every norm/RoPE combination.
+    """
+    from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused
+
+    torch.manual_seed(7)
+    device = "cuda"
+    dtype = torch.bfloat16
+    block_size, num_tokens, q_head_padded = 64, 5, 64
+
+    q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
+    kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
+
+    def run(q_in, is_q_interleaved):
+        cache = torch.zeros(
+            num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device
+        )
+        out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+            q_in,
+            kv,
+            cache,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            q_head_padded,
+            1e-6,
+            block_size,
+            apply_q_norm,
+            True,  # kv_mxfp8
+            apply_q_rope,
+            is_q_interleaved,
+        )
+        return out, cache
+
+    head_major, cache_major = run(q, False)
+    interleaved, cache_inter = run(permute_q_to_fused(q), True)
+
+    torch.testing.assert_close(
+        interleaved, permute_q_to_fused(head_major), rtol=0, atol=0
+    )
+    # The Q layout must not leak into the KV half.
+    torch.testing.assert_close(cache_inter, cache_major, rtol=0, atol=0)
+
+
 # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ─────────────────────
 
 
diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py
index 5b568c938eec..0ce3a918513f 100644
--- a/tests/kernels/test_fused_indexer_q_rope_quant.py
+++ b/tests/kernels/test_fused_indexer_q_rope_quant.py
@@ -14,6 +14,7 @@
 """
 
 import contextlib
+import importlib
 from unittest import mock
 
 import pytest
@@ -29,9 +30,10 @@
 
 HEAD_DIM = 128
 ROPE_DIM = 64
-N_HEAD = 64
 MAX_POS = 4096
 
+CUTEDSL_MODULE = "vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl"
+
 
 def quantize_to_mxfp4(
     x: torch.Tensor,
@@ -43,6 +45,7 @@ def quantize_to_mxfp4(
     Returns:
         packed: [..., head_dim//2]  uint8   2 E2M1 nibbles/byte, low nibble = even index
         scales: [..., head_dim//32] uint8   1 ue8m0 byte
+
     """
     MXFP4_BLOCK_SIZE = 32
     orig_shape = x.shape
@@ -90,6 +93,7 @@ def _reference(
     weights: torch.Tensor,
     softmax_scale: float,
     head_scale: float,
+    n_head: int,
     use_fp4: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     q_rot = q.clone()
@@ -105,7 +109,7 @@ def _reference(
     )
 
     if use_fp4:
-        q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
+        q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, n_head, HEAD_DIM))
         # Pack 4 ue8m0 bytes into 1 int32
         q_scale = ue8m0.view(torch.int32).squeeze(-1)
         # FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
@@ -118,8 +122,8 @@ def _reference(
             HEAD_DIM,
             use_ue8m0=True,
         )
-        q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
-        q_scale = q_scale.view(-1, N_HEAD)
+        q_fp8 = q_fp8.view(-1, n_head, HEAD_DIM)
+        q_scale = q_scale.view(-1, n_head)
 
         weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
         return q_fp8, weights_out
@@ -144,9 +148,10 @@ def _reference(
     ],
 )
 @pytest.mark.parametrize("use_cutedsl", [False, True])
+@pytest.mark.parametrize("n_head", [32, 64])
 @torch.inference_mode()
 def test_fused_indexer_q_rope_quant_matches_unfused(
-    num_tokens, cache_dtype, use_fp4, use_cutedsl
+    num_tokens, cache_dtype, use_fp4, use_cutedsl, n_head
 ):
     if use_cutedsl and not has_cutedsl():
         pytest.skip("cutedsl (cutlass) not installed")
@@ -154,17 +159,29 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
     device = "cuda"
     torch.manual_seed(0)
 
-    q = torch.randn(num_tokens, N_HEAD, HEAD_DIM, dtype=torch.bfloat16, device=device)
+    q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device)
     positions = torch.randint(
         0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device
     )
     cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=cache_dtype, device=device)
-    weights = torch.randn(num_tokens, N_HEAD, dtype=torch.bfloat16, device=device)
+    weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device)
     softmax_scale = HEAD_DIM**-0.5
-    head_scale = N_HEAD**-0.5
+    # head_scale must be an exact power of two for bit-exactness to be a fair
+    # expectation: the triton kernel multiplies by softmax_scale and head_scale
+    # separately (two roundings) while the cutedsl kernels fold them into one
+    # fp32 scalar first (one rounding). Those agree only when one factor is a
+    # pow2. n_head**-0.5 is a pow2 at 64 heads but not at 32.
+    head_scale = 0.125
 
     q_quant_ref, weights_ref = _reference(
-        positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
+        positions,
+        q,
+        cos_sin_cache,
+        weights,
+        softmax_scale,
+        head_scale,
+        n_head,
+        use_fp4,
     )
     # use_cutedsl=False: force the triton path even when cutedsl is installed
     # by patching the dispatcher's has_cutedsl() binding to return False.
@@ -210,3 +227,389 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
         f"weights mismatch: max abs diff "
         f"{(weights_ref - weights_fused).abs().max().item()}"
     )
+
+
+@pytest.mark.skipif(not has_cutedsl(), reason="cutedsl (cutlass) not installed")
+@pytest.mark.parametrize(
+    "use_fp4",
+    [
+        False,
+        pytest.param(
+            True,
+            marks=pytest.mark.skipif(
+                not (
+                    current_platform.is_cuda()
+                    and current_platform.is_device_capability_family(100)
+                ),
+                reason="MXFP4 indexer cache requires an SM100-family GPU",
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize("n_head", [32, 64])
+@torch.inference_mode()
+def test_cutedsl_indexer_q_writes_stay_within_num_tokens(use_fp4, n_head):
+    """The cutedsl kernels must not store past ``num_tokens``.
+
+    The launcher rounds ``num_tokens * threads_per_token`` up to whole 128-thread
+    blocks, so the trailing block is only partially in range unless
+    ``threads_per_token`` is itself a multiple of the block size. With 32 heads
+    and coarsen=4 (picked for >= 512 tokens) a token needs 64 threads, so an odd
+    token count leaves half of the last block addressing token ``num_tokens``.
+    Every output gets one extra sentinel row that the kernel must leave alone.
+    """
+    mod = importlib.import_module(CUTEDSL_MODULE)
+    device = "cuda"
+    torch.manual_seed(0)
+    num_tokens = 935
+
+    sentinel_u8, sentinel_f32 = 0x55, -12345.0
+    # positions/q are padded too, so an out-of-bounds *load* stays addressable
+    # and the assertions below are about stores only.
+    positions = torch.zeros(num_tokens + 1, dtype=torch.int64, device=device)
+    q = torch.randn(
+        num_tokens + 1, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device
+    )
+    cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device)
+    weights = torch.randn(num_tokens + 1, n_head, dtype=torch.bfloat16, device=device)
+    weights_out = torch.full(
+        (num_tokens + 1, n_head), sentinel_f32, dtype=torch.float32, device=device
+    )
+    outputs = {"weights_out": weights_out}
+
+    if use_fp4:
+        outputs["q_packed"] = torch.full(
+            (num_tokens + 1, n_head, HEAD_DIM // 2),
+            sentinel_u8,
+            dtype=torch.uint8,
+            device=device,
+        )
+        outputs["q_scale"] = torch.full(
+            (num_tokens + 1, n_head, HEAD_DIM // 32),
+            sentinel_u8,
+            dtype=torch.uint8,
+            device=device,
+        )
+        mod._INDEXER_Q_MXFP4_KERNEL(
+            positions=positions[:num_tokens],
+            q=q[:num_tokens],
+            cos_sin_cache=cos_sin_cache,
+            weights=weights[:num_tokens],
+            weights_softmax_scale=1.0,
+            weights_head_scale=1.0,
+            q_packed=outputs["q_packed"][:num_tokens],
+            q_scale=outputs["q_scale"][:num_tokens],
+            weights_out=weights_out[:num_tokens],
+        )
+    else:
+        outputs["q_fp8"] = torch.full(
+            (num_tokens + 1, n_head, HEAD_DIM),
+            sentinel_u8,
+            dtype=torch.uint8,
+            device=device,
+        )
+        mod._INDEXER_Q_FP8_KERNEL(
+            positions=positions[:num_tokens],
+            q=q[:num_tokens],
+            cos_sin_cache=cos_sin_cache,
+            weights=weights[:num_tokens],
+            weights_softmax_scale=1.0,
+            weights_head_scale=1.0,
+            q_fp8=outputs["q_fp8"][:num_tokens].view(torch.float8_e4m3fn),
+            weights_out=weights_out[:num_tokens],
+        )
+    torch.accelerator.synchronize()
+
+    clobbered = {}
+    for name, tensor in outputs.items():
+        sentinel = sentinel_f32 if tensor.dtype == torch.float32 else sentinel_u8
+        n_changed = (tensor[num_tokens] != sentinel).sum().item()
+        if n_changed:
+            clobbered[name] = n_changed
+    assert not clobbered, (
+        f"kernel wrote past token {num_tokens - 1} into the sentinel row: {clobbered}"
+    )
+
+
+def _indexer_k_reference(
+    k_pre, positions, cos_sin_cache, rms_norm_weight, compress_ratio, use_fp4
+):
+    k = k_pre.float()
+    k = k * torch.rsqrt(k.square().mean(dim=-1, keepdim=True) + 1e-20)
+    k = (k * rms_norm_weight.float()).to(torch.bfloat16).float()
+    group_positions = positions // compress_ratio * compress_ratio
+    cos, sin = cos_sin_cache[group_positions].float().chunk(2, dim=-1)
+    even, odd = k[:, 64::2], k[:, 65::2]
+    rotated = torch.stack((even * cos - odd * sin, odd * cos + even * sin), dim=-1)
+    k = torch.cat((k[:, :64], rotated.flatten(1)), dim=-1).to(torch.bfloat16)
+    if use_fp4:
+        return quantize_to_mxfp4(k)
+    exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448
+    scale = exponent.log2().ceil().exp2()
+    values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
+    return values.view(torch.uint8), scale.view(torch.uint8)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel")
+@pytest.mark.parametrize("num_tokens", [1, 17, 257])
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+@pytest.mark.parametrize("use_fp4", [False, True])
+@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
+@torch.inference_mode()
+def test_indexer_k_inserts_only_valid_groups_into_padded_pages(
+    num_tokens, compress_ratio, use_fp4, cache_dtype
+):
+    """Insert group ends and preserve skipped slots, graph rows, and page padding."""
+    from vllm.models.deepseek_v41.common.ops.indexer_k_store import (
+        indexer_k_norm_rope_store,
+    )
+
+    if use_fp4 and not current_platform.is_device_capability_family(100):
+        pytest.skip("MXFP4 indexer cache requires an SM100-family GPU")
+    torch.manual_seed(17)
+    device = "cuda"
+    k_pre = torch.randn(num_tokens + 3, 128, device=device, dtype=torch.bfloat16)
+    norm_weight = torch.randn(128, device=device, dtype=torch.bfloat16)
+    positions = torch.arange(7, num_tokens + 10, device=device)
+    slots = torch.randperm(2 * num_tokens, device=device)[:num_tokens].clone()
+    slots[5::7] = -1
+    valid = (slots >= 0) & ((positions[:num_tokens] + 1) % compress_ratio == 0)
+    k_pre[3::7] = 0
+    k_pre[8::11] *= 1e-10
+    k_pre[:num_tokens][~valid] = float("nan")
+    k_pre[num_tokens:] = float("nan")
+    positions[num_tokens:] = MAX_POS + 10
+    angles = torch.randn(MAX_POS, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).to(cache_dtype)
+    block_size = 16
+    value_bytes = 64 if use_fp4 else 128
+    page_bytes = block_size * (value_bytes + 4)
+    num_blocks = (2 * num_tokens + block_size - 1) // block_size + 1
+    storage = torch.full(
+        (num_blocks, page_bytes + 128), 0xA5, device=device, dtype=torch.uint8
+    )
+    cache = storage.as_strided(
+        (num_blocks, block_size, value_bytes + 4),
+        (storage.stride(0), value_bytes + 4, 1),
+    )
+    expected = storage.clone()
+    values, scales = _indexer_k_reference(
+        k_pre[:num_tokens][valid],
+        positions[:num_tokens][valid],
+        cos_sin,
+        norm_weight,
+        compress_ratio,
+        use_fp4,
+    )
+    block_ids = slots[valid] // block_size
+    block_offsets = slots[valid] % block_size
+    value_offsets = block_offsets[:, None] * value_bytes + torch.arange(
+        value_bytes, device=device
+    )
+    scale_offsets = block_size * value_bytes + block_offsets[:, None] * 4
+    scale_offsets = scale_offsets + torch.arange(4, device=device)
+    expected[block_ids[:, None], value_offsets] = values
+    expected[block_ids[:, None], scale_offsets] = scales
+
+    indexer_k_norm_rope_store(
+        k_pre,
+        positions,
+        cos_sin,
+        norm_weight,
+        1e-20,
+        cache,
+        slots,
+        compress_ratio,
+        use_fp4,
+    )
+    # Fused RoPE can change zero signs without changing their numerical values.
+    actual_values = storage[block_ids[:, None], value_offsets]
+    if use_fp4:
+        for shift in (0, 4):
+            zero_mask = (actual_values >> shift) & 7 == 0
+            actual_values &= ~(zero_mask.to(torch.uint8) << (shift + 3))
+            zero_mask = (values >> shift) & 7 == 0
+            values &= ~(zero_mask.to(torch.uint8) << (shift + 3))
+    else:
+        actual_values[actual_values == 128] = 0
+        values[values == 128] = 0
+    storage[block_ids[:, None], value_offsets] = actual_values
+    expected[block_ids[:, None], value_offsets] = values
+    torch.testing.assert_close(storage, expected, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel")
+@pytest.mark.parametrize("use_fp4", [False, True])
+@torch.inference_mode()
+def test_indexer_k_cuda_graph_replay_reads_current_projection(use_fp4):
+    """Replay must consume updated keys and slot mapping through the captured API."""
+    from vllm.models.deepseek_v41.common.ops.indexer_k_store import (
+        indexer_k_norm_rope_store,
+    )
+
+    if use_fp4 and not current_platform.is_device_capability_family(100):
+        pytest.skip("MXFP4 indexer cache requires an SM100-family GPU")
+    torch.manual_seed(18)
+    k_pre = torch.randn(19, 128, device="cuda", dtype=torch.bfloat16)
+    norm_weight = torch.ones(128, device="cuda", dtype=torch.bfloat16)
+    positions = torch.arange(19, device="cuda")
+    cos_sin = torch.zeros(32, 64, device="cuda")
+    cos_sin[:, :32] = 1
+    slots = torch.arange(17, device="cuda")
+    cache = torch.full(
+        (2, 16, (64 if use_fp4 else 128) + 4),
+        0xA5,
+        device="cuda",
+        dtype=torch.uint8,
+    )
+
+    def run():
+        indexer_k_norm_rope_store(
+            k_pre,
+            positions,
+            cos_sin,
+            norm_weight,
+            1e-20,
+            cache,
+            slots,
+            2,
+            use_fp4,
+        )
+
+    run()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run()
+    k_pre.neg_()
+    slots.add_(4)
+    cache.fill_(0xA5)
+    graph.replay()
+    replayed = cache.clone()
+    cache.fill_(0xA5)
+    run()
+    torch.testing.assert_close(cache, replayed, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm cache layout")
+@pytest.mark.parametrize("block_size", [16, 64, 128])
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+@torch.inference_mode()
+def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_ratio):
+    """ROCm reads the indexer K cache back with the 16x16-tiled value layout.
+
+    ``cp_gather_indexer_k_quant_cache_triton`` (prefill) and aiter's
+    ``Preshuffle=True`` paged kernel (decode) both pick that layout whenever
+    ``block_size > 1``, so a row-major store would feed the indexer permuted
+    key bytes and randomize its top-k. Assert the write/read pair is exact.
+    """
+    from vllm.models.deepseek_v41.common.ops.indexer_k_store import (
+        indexer_k_norm_rope_store,
+    )
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        cp_gather_indexer_k_quant_cache_triton,
+    )
+
+    torch.manual_seed(0)
+    device = "cuda"
+    num_tokens = 3 * block_size + 5
+    num_blocks = num_tokens // block_size + 2
+    k_pre = torch.randn(num_tokens, HEAD_DIM, device=device, dtype=torch.bfloat16)
+    norm_weight = torch.randn(HEAD_DIM, device=device, dtype=torch.bfloat16)
+    positions = torch.arange(num_tokens, device=device)
+    angles = torch.randn(MAX_POS, ROPE_DIM // 2, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).float()
+    cache = torch.zeros(
+        num_blocks, block_size, HEAD_DIM + 4, device=device, dtype=torch.uint8
+    )
+    slots = torch.arange(num_tokens, device=device)
+
+    indexer_k_norm_rope_store(
+        k_pre,
+        positions,
+        cos_sin,
+        norm_weight,
+        1e-20,
+        cache,
+        slots,
+        compress_ratio,
+        False,
+    )
+
+    k_fp8 = torch.zeros(num_tokens, HEAD_DIM, device=device, dtype=torch.float8_e4m3fn)
+    k_scale = torch.zeros(num_tokens, 4, device=device, dtype=torch.uint8)
+    cp_gather_indexer_k_quant_cache_triton(
+        cache,
+        k_fp8,
+        k_scale,
+        torch.arange(num_blocks, device=device, dtype=torch.int32).view(1, -1),
+        torch.tensor([0, num_tokens], device=device, dtype=torch.int32),
+        torch.zeros(num_tokens, device=device, dtype=torch.int32),
+    )
+
+    expected_values, expected_scale = _indexer_k_reference(
+        k_pre, positions, cos_sin, norm_weight, compress_ratio, False
+    )
+    emitted = (positions + 1) % compress_ratio == 0
+    gathered = k_fp8.view(torch.uint8)
+    gathered = torch.where(gathered == 128, gathered.new_zeros(()), gathered)
+    expected_values = torch.where(
+        expected_values == 128, expected_values.new_zeros(()), expected_values
+    )
+    torch.testing.assert_close(
+        gathered[emitted], expected_values[emitted], rtol=0, atol=0
+    )
+    torch.testing.assert_close(
+        k_scale[emitted], expected_scale[emitted], rtol=0, atol=0
+    )
+
+
+@pytest.mark.parametrize("use_cutedsl", [False, True])
+@pytest.mark.skipif(
+    not (
+        current_platform.is_cuda() and current_platform.is_device_capability_family(100)
+    ),
+    reason="MXFP4 indexer cache requires an SM100-family GPU",
+)
+@torch.inference_mode()
+def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl):
+    """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's
+    sparse MQA-logits kernels take) instead of fp32; the values are the fp32
+    result rounded once, so the scoring kernel needs no cast."""
+    if use_cutedsl and not has_cutedsl():
+        pytest.skip("cutedsl (cutlass) not installed")
+
+    device = "cuda"
+    torch.manual_seed(0)
+    num_tokens, n_head = 257, 32
+    q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device)
+    positions = torch.randint(
+        0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device
+    )
+    cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device)
+    weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device)
+    softmax_scale, head_scale = HEAD_DIM**-0.5, 0.125  # pow2 head_scale, see above
+
+    _, weights_ref = _reference(
+        positions, q, cos_sin_cache, weights, softmax_scale, head_scale, n_head, True
+    )
+    cutedsl_patch = (
+        mock.patch(
+            "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl",
+            return_value=False,
+        )
+        if not use_cutedsl
+        else contextlib.nullcontext()
+    )
+    with cutedsl_patch:
+        _, weights_fused = fused_indexer_q_rope_quant(
+            positions,
+            q.clone(),
+            cos_sin_cache,
+            weights,
+            softmax_scale,
+            head_scale,
+            use_fp4=True,
+            weights_out_dtype=torch.bfloat16,
+        )
+    assert weights_fused.dtype == torch.bfloat16
+    assert torch.equal(weights_fused, weights_ref.to(torch.bfloat16))
diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py
index 5b71e59baa9e..03d5b7921fa2 100644
--- a/tests/kernels/test_fused_inv_rope_fp8_quant.py
+++ b/tests/kernels/test_fused_inv_rope_fp8_quant.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel.
+"""Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel.
 
 Tests compare the fused kernel against a reference implementation built from
 the existing separate operations (inverse RoPE via rotate_neox + FP8 quant
@@ -132,6 +131,7 @@ def reference_inv_rope(
 
     Returns:
         o with inverse RoPE applied on the rope portion (bf16).
+
     """
     assert cos_sin_cache.dtype == torch.float32
     cos_sin = cos_sin_cache[positions]  # [T, rope_dim] fp32
@@ -159,13 +159,13 @@ def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Ten
       absmax -> 2^ceil(log2(absmax / fp8_max)) -> clamp(x / scale) -> fp8
 
     Args:
-        x_f32: [..., quant_group_size] float32 — one or more 128-element blocks.
+        x_f32: [..., quant_group_size] float32 quantization blocks.
 
     Returns:
         x_fp8: same shape, float8_e4m3fn
         scales: [...] float32, one scale per block
+
     """
-    assert x_f32.shape[-1] == QUANT_GROUP_SIZE
     absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS)
     scale_raw = absmax * (1.0 / FP8_MAX)
     scale = torch.exp2(torch.ceil(torch.log2(scale_raw)))
@@ -193,6 +193,7 @@ def reference_inv_rope_fp8_quant(
     Returns:
         o_fp8: [T, G, D] FP8 with strides (D, T*D, 1)
         o_scale: [T, G, S] FP32 with strides (S, T*S, 1)
+
     """
     assert cos_sin_cache.dtype == torch.float32
     T, _H, head_dim = o.shape
@@ -253,6 +254,56 @@ def reference_inv_rope_fp8_quant(
 # =========================================================================
 
 
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 4096])
+@pytest.mark.parametrize("quant_group_size", [32, 128])
+@torch.inference_mode()
+def test_packed_quantization_preserves_rotary_dimensions(num_tokens, quant_group_size):
+    """RoPE may span multiple scale blocks; pack every scale and padding row."""
+    from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
+        fused_inv_rope_fp8_quant as quantize,
+    )
+
+    torch.manual_seed(42)
+    groups, heads_per_group = 2, 2
+    x = torch.randn(
+        num_tokens,
+        groups * heads_per_group,
+        HEAD_DIM,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    cache = make_cos_sin_cache(4096)
+    positions = torch.randint(4096, (num_tokens,), device="cuda")
+    q, packed_scales = quantize(
+        x,
+        positions,
+        cache,
+        groups,
+        heads_per_group,
+        quant_group_size=quant_group_size,
+        tma_aligned_scales=True,
+    )
+    scale_bytes = torch.stack(
+        [(packed_scales >> (8 * i)) & 0xFF for i in range(4)], dim=-1
+    ).flatten(-2)
+    scales = torch.exp2(scale_bytes.float() - 127)
+    rotated = reference_inv_rope(x.float(), positions, cache)
+    ref_q, ref_scales = _ref_ue8m0_quant_block(
+        rotated.reshape(num_tokens, groups, -1, quant_group_size)
+    )
+    torch.testing.assert_close(scales, ref_scales, rtol=0, atol=0)
+    actual = q.float().unflatten(-1, (-1, quant_group_size)) * scales.unsqueeze(-1)
+    expected = ref_q.float() * ref_scales.unsqueeze(-1)
+    assert (actual - expected).norm() / expected.norm() < 1e-3
+
+    aligned_tokens = (num_tokens + 3) // 4 * 4
+    full_scales = packed_scales.transpose(0, 1).as_strided(
+        (groups, aligned_tokens, packed_scales.shape[-1]),
+        (packed_scales.stride(1), 1, packed_scales.stride(2)),
+    )
+    assert torch.count_nonzero(full_scales[:, num_tokens:]) == 0
+
+
 @pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
 @pytest.mark.parametrize(
     "num_heads,n_groups",
@@ -260,8 +311,9 @@ def reference_inv_rope_fp8_quant(
     ids=["H64_G8", "H32_G4", "H128_G8"],
 )
 @pytest.mark.parametrize("seed", [0, 42])
+@pytest.mark.parametrize("quantize", [True, False], ids=["fp8", "bf16"])
 @torch.inference_mode()
-def test_correctness(num_tokens, num_heads, n_groups, seed):
+def test_correctness(num_tokens, num_heads, n_groups, seed, quantize):
     """Compare fused kernel against reference for FP8 values and scales."""
     torch.manual_seed(seed)
 
@@ -280,6 +332,26 @@ def test_correctness(num_tokens, num_heads, n_groups, seed):
         max_pos, ROPE_DIM, dtype=torch.float32, device=device
     )
 
+    if not quantize:
+        ref_output = reference_inv_rope(o, positions, cos_sin_cache).view(
+            num_tokens, n_groups, -1
+        )
+        fused_output, fused_scale = fused_inv_rope_fp8_quant(
+            o.clone(),
+            positions,
+            cos_sin_cache,
+            n_groups,
+            heads_per_group,
+            quantize=False,
+        )
+        d = heads_per_group * HEAD_DIM
+        assert fused_output.shape == (num_tokens, n_groups, d)
+        assert fused_scale.numel() == 0
+        assert fused_output.dtype == torch.bfloat16
+        assert fused_output.stride() == (d, num_tokens * d, 1)
+        torch.testing.assert_close(fused_output, ref_output, rtol=1e-2, atol=1e-2)
+        return
+
     # Reference
     ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
         o.clone(),
@@ -332,7 +404,6 @@ def test_output_strides(num_tokens, num_heads, n_groups):
     - FP8: logical [T, G, D] backed by contiguous [G, T, D].
     - Scale: MN-major TMA-aligned (column-major: T-stride=1).
     """
-
     heads_per_group = num_heads // n_groups
     max_pos = 4096
     device = "cuda"
@@ -726,6 +797,7 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
     This catches stride/layout bugs that only manifest when the einsum
     kernel actually consumes the quantized activations.
     """
+    from vllm.models.deepseek_v4.nvidia.ops.o_proj import deep_gemm_fp8_o_proj
     from vllm.utils.deep_gemm import (
         fp8_einsum,
         is_deep_gemm_supported,
@@ -793,23 +865,23 @@ def ceil_div(a: int, b: int) -> int:
     )
 
     # -- FUSED path --
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
+    wo_a = torch.nn.Module()
+    wo_a.register_buffer("weight", w_fp8)
+    wo_a.register_buffer("weight_scale", w_scale_t)
+    z_fused = deep_gemm_fp8_o_proj(
         o.clone(),
         positions,
         cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-    z_fused = torch.empty(
-        num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16
-    )
-    fp8_einsum(
-        "bhr,hdr->bhd",
-        (fused_fp8, fused_scale),
-        (w_fp8, w_scale_t),
-        z_fused,
-        recipe=recipe,
-    )
+        wo_a,
+        torch.nn.Identity(),
+        n_groups=n_groups,
+        heads_per_group=heads_per_group,
+        nope_dim=NOPE_DIM,
+        rope_dim=ROPE_DIM,
+        o_lora_rank=o_lora_rank,
+        einsum_recipe=recipe,
+        tma_aligned_scales=True,
+    ).view_as(z_ref)
 
     # -- Checks --
     # Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
@@ -832,7 +904,6 @@ def calc_diff(x, y):
 def test_with_real_deepseek_v4_rope(num_tokens, default_vllm_config):
     """Test with real DeepseekV4ScalingRotaryEmbedding (GPT-J style,
     mscale=0, YaRN scaling) matching the production config."""
-
     num_heads = 64
     n_groups = 8
     heads_per_group = num_heads // n_groups
diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py
index 267484e43361..321cb6978e3d 100644
--- a/tests/kernels/test_fused_quant_activation.py
+++ b/tests/kernels/test_fused_quant_activation.py
@@ -91,7 +91,7 @@ class MockLinearFp8Static(torch.nn.Module):
 
     def __init__(self, input_scale: torch.Tensor):
         super().__init__()
-        self.input_quant_key = kFp8StaticTensorSym
+        self._input_quant_key = kFp8StaticTensorSym
         self.input_scale = input_scale
 
 
@@ -100,15 +100,21 @@ class MockLinearFp8Dynamic128(torch.nn.Module):
 
     def __init__(self):
         super().__init__()
-        self.input_quant_key = kFp8Dynamic128Sym
+        self._input_quant_key = kFp8Dynamic128Sym
 
 
 class MockLinearNoQuant(torch.nn.Module):
-    """Mock linear layer with no input_quant_key (no fusion)."""
+    """Mock linear layer with no input quantization key."""
 
     pass
 
 
+class MockLinearRequiresUnquantized(MockLinearFp8Static):
+    """Mock consumer that needs the original activation for another branch."""
+
+    requires_unquantized_input = True
+
+
 @pytest.mark.parametrize("num_tokens", [1, 16, 128])
 @pytest.mark.parametrize("hidden_size", [128, 512, 1024])
 @pytest.mark.parametrize("dtype", DTYPES)
@@ -199,7 +205,7 @@ def test_maybe_fused_act_quant_fallback(
     hidden_size: int,
     dtype: torch.dtype,
 ) -> None:
-    """Test maybe_fused_act_quant falls back when no input_quant_key."""
+    """Test maybe_fused_act_quant falls back without an input quantization key."""
     device = "cuda:0"
     torch.set_default_device(device)
 
@@ -216,3 +222,19 @@ def test_maybe_fused_act_quant_fallback(
 
     ref_out = act_fn(x)
     torch.testing.assert_close(result, ref_out)
+
+
+@torch.inference_mode()
+def test_maybe_fused_act_quant_preserves_required_unquantized_input(
+    default_vllm_config,
+) -> None:
+    device = "cuda:0"
+    act_fn = SiluAndMul()
+    scale = torch.tensor([0.5], device=device, dtype=torch.float32)
+    linear = MockLinearRequiresUnquantized(scale)
+    x = torch.randn(17, 256, dtype=torch.bfloat16, device=device)
+
+    result = maybe_fused_act_quant(act_fn, x, linear)
+
+    assert isinstance(result, torch.Tensor)
+    torch.testing.assert_close(result, act_fn(x))
diff --git a/tests/kernels/test_kimi_k3_gemm_rs_ar.py b/tests/kernels/test_gemm_rs_ar.py
similarity index 63%
rename from tests/kernels/test_kimi_k3_gemm_rs_ar.py
rename to tests/kernels/test_gemm_rs_ar.py
index 7a887c501f10..b1c79af2b925 100644
--- a/tests/kernels/test_kimi_k3_gemm_rs_ar.py
+++ b/tests/kernels/test_gemm_rs_ar.py
@@ -6,7 +6,7 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from tests.utils import ensure_current_vllm_config
+from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.distributed.parallel_state import (
     init_distributed_environment,
@@ -17,6 +17,7 @@
 from vllm.utils.system_utils import update_environment_variables
 
 _SHAPES = (
+    (129, 512),
     (129, 768),
     (257, 768),
     (1023, 4224),
@@ -29,6 +30,7 @@
 def _reference(
     x: torch.Tensor,
     weight: torch.Tensor,
+    weight_scale: torch.Tensor | None,
     world_size: int,
     group: dist.ProcessGroup,
     all_reduce: bool,
@@ -40,7 +42,21 @@ def _reference(
         dtype=x.dtype,
         device=x.device,
     )
-    torch.mm(x, weight.T, out=partial[:M])
+    if weight_scale is None:
+        torch.mm(x, weight.T, out=partial[:M])
+    else:
+        from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+            mxfp8_e4m3_quantize,
+        )
+
+        x_q, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+        partial[:M] = torch._scaled_mm(
+            x_q,
+            weight.T,
+            x_sf.view(torch.float8_e8m0fnu),
+            weight_scale.view(torch.float8_e8m0fnu),
+            out_dtype=x.dtype,
+        )
     if padded_M > M:
         partial[M:].zero_()
 
@@ -75,6 +91,10 @@ def _assert_valid_rows_close(
         rtol=5e-2,
         atol=4.0,
     )
+    # Sequence-parallel callers may pass an unpadded M and read the padding
+    # rows back (they flow through per-token kernels before being dropped),
+    # so RS must leave them zero like ``sp_reduce_scatter`` does.
+    assert not actual[valid_rows:].any()
 
 
 def _run_mode(
@@ -85,19 +105,52 @@ def _run_mode(
     rank: int,
     world_size: int,
     weights: dict[int, torch.Tensor],
+    backend: str,
 ) -> None:
     # cute_dsl is unavailable off CUDA, so import it only inside the GPU worker.
-    from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr
+    from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
 
     gemm_rs_ar = GemmRsAr(
         max_M=max(M for M, _ in _SHAPES),
         N=_N,
         all_reduce=all_reduce,
     )
+    from vllm.config.quantization import QuantizationConfigArgs
+    from vllm.model_executor.layers.linear import RowParallelLinear
+    from vllm.model_executor.layers.quantization.online.base import (
+        OnlineQuantizationConfig,
+    )
+
+    quant_config = (
+        None
+        if backend == "bf16"
+        else OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8"))
+    )
+    projections = {}
+    operands = {}
+    for K, w in weights.items():
+        with torch.device(device):
+            linear = RowParallelLinear(
+                K * world_size,
+                _N,
+                bias=False,
+                params_dtype=torch.bfloat16,
+                quant_config=quant_config,
+            )
+        # Select the fused path before online quantization replaces meta weights.
+        assert gemm_rs_ar.can_run(linear)
+        linear.weight = torch.nn.Parameter(w.clone(), requires_grad=False)
+        linear.quant_method.process_weights_after_loading(linear)
+        projections[K] = linear
+        weight = linear.weight
+        if backend == "flashinfer_cutedsl":
+            weight = weight.t()
+        operands[K] = (weight, getattr(linear, "weight_scale", None))
+
     input_generator = torch.Generator(device=device)
 
-    # Alternating shapes exercise producer-flag reuse across different grids,
-    # CTA-group choices, and both BN=128 and BN=256 dispatches.
+    # Alternating shapes exercise producer-flag reuse across different grids
+    # and CTA-group choices.
     for M, K in (*_SHAPES, *_SHAPES[::-1]):
         input_generator.manual_seed(2000 + M + K)
         x = torch.randn(
@@ -107,8 +160,8 @@ def _run_mode(
             device=device,
             generator=input_generator,
         )
-        expected = _reference(x, weights[K], world_size, group, all_reduce)
-        actual = gemm_rs_ar(x, weights[K])
+        expected = _reference(x, *operands[K], world_size, group, all_reduce)
+        actual = gemm_rs_ar.apply(x, projections[K])
         torch.accelerator.synchronize(device)
         _assert_valid_rows_close(actual, expected, M, rank, all_reduce)
 
@@ -123,9 +176,9 @@ def _run_mode(
             device=device,
             generator=input_generator,
         )
-        first_output = gemm_rs_ar(lifetime_x, weights[lifetime_K])
+        first_output = gemm_rs_ar.apply(lifetime_x, projections[lifetime_K])
         first_snapshot = first_output.clone()
-        second_output = gemm_rs_ar(-lifetime_x, weights[lifetime_K])
+        second_output = gemm_rs_ar.apply(-lifetime_x, projections[lifetime_K])
         torch.accelerator.synchronize(device)
         assert first_output.data_ptr() != second_output.data_ptr()
         torch.testing.assert_close(first_output, first_snapshot, rtol=0, atol=0)
@@ -141,7 +194,7 @@ def _run_mode(
     )
     graph_expected = _reference(
         graph_x,
-        weights[graph_K],
+        *operands[graph_K],
         world_size,
         group,
         all_reduce,
@@ -152,13 +205,13 @@ def _run_mode(
     capture_stream.wait_stream(torch.cuda.current_stream())
     with torch.cuda.stream(capture_stream):
         for _ in range(3):
-            graph_output.copy_(gemm_rs_ar(graph_x, weights[graph_K]))
+            graph_output.copy_(gemm_rs_ar.apply(graph_x, projections[graph_K]))
     capture_stream.synchronize()
     dist.barrier(group=group)
 
     graph = torch.cuda.CUDAGraph()
     with torch.cuda.graph(graph, stream=capture_stream):
-        graph_output.copy_(gemm_rs_ar(graph_x, weights[graph_K]))
+        graph_output.copy_(gemm_rs_ar.apply(graph_x, projections[graph_K]))
     torch.cuda.current_stream().wait_stream(capture_stream)
     dist.barrier(group=group)
 
@@ -177,7 +230,7 @@ def _run_mode(
     del gemm_rs_ar, capture_stream, graph, graph_output
 
 
-def _worker(local_rank: int, world_size: int, master_port: int) -> None:
+def _worker(local_rank: int, world_size: int, master_port: int, backend: str) -> None:
     # This module pulls in cute_dsl, which is unavailable off CUDA, so importing
     # it at module level would fail collection. Import it where it is used, as
     # `kda.py` does.
@@ -195,7 +248,10 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None:
     )
 
     init_distributed_environment()
-    with ensure_current_vllm_config():
+    config = VllmConfig()
+    config.model_config = ModelConfig(dtype="bfloat16")
+    config.kernel_config.linear_backend = "auto" if backend == "bf16" else backend
+    with set_current_vllm_config(config):
         initialize_model_parallel(tensor_model_parallel_size=world_size)
 
     group = dist.group.WORLD
@@ -215,24 +271,29 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None:
 
     # Production binds one mode per worker; exercise both mode-bound instances
     # sequentially in this test without implying that both are initialized.
-    for all_reduce in (False, True):
-        _run_mode(
-            all_reduce=all_reduce,
-            device=device,
-            group=group,
-            rank=rank,
-            world_size=world_size,
-            weights=weights,
-        )
+    with set_current_vllm_config(config):
+        for all_reduce in (False, True):
+            _run_mode(
+                all_reduce=all_reduce,
+                device=device,
+                group=group,
+                rank=rank,
+                world_size=world_size,
+                weights=weights,
+                backend=backend,
+            )
     cleanup_dist_env_and_memory()
 
 
+@pytest.mark.parametrize(
+    "backend", ["bf16", "flashinfer_cutlass", "flashinfer_cutedsl"]
+)
 @pytest.mark.distributed(num_gpus=2)
 @pytest.mark.skipif(
     not current_platform.is_device_capability_family(100),
-    reason="Kimi-K3 GEMM-RS/AR requires SM100",
+    reason="GEMM-RS/AR requires SM100",
 )
-def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, backend: str) -> None:
     world_size = 2
     if torch.accelerator.device_count() < world_size:
         pytest.skip("GEMM-RS/AR requires two GPUs")
@@ -242,7 +303,7 @@ def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch) -> None:
     try:
         mp.spawn(
             _worker,
-            args=(world_size, get_open_port()),
+            args=(world_size, get_open_port(), backend),
             nprocs=world_size,
         )
     finally:
diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py
index afc0c1f836b7..37ba4b58c858 100644
--- a/tests/kernels/test_kpool_decode_update_batched.py
+++ b/tests/kernels/test_kpool_decode_update_batched.py
@@ -384,9 +384,16 @@ def test_leading_invalid_tail_slot():
     _assert_eq(r_ref, r_kern)
 
 
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
-def test_amd_prefill_seed_honors_padded_tail_block_stride():
-    """The tail shares a padded indexer allocation in production."""
+def test_prefill_seed_honors_padded_tail_block_stride():
+    """The tail shares a padded indexer allocation in production.
+
+    ``get_kv_cache_config_from_groups`` aliases each tail tensor onto its
+    indexer tensor with the indexer's block stride (38016 B for GLM-5.3-Flash
+    vs a dense 2048 B tail block), so a seed kernel that addresses blocks
+    densely writes into an unrelated indexer block and leaves the request's
+    tail block untouched. Runs on every platform; the NVIDIA kernel had this
+    bug while the AMD kernel did not.
+    """
     kpool = 4
     num_blocks = 6
     logical_block_elems = 2 * kpool * HEAD_DIM
diff --git a/tests/kernels/test_mhc_jit_warmup.py b/tests/kernels/test_mhc_jit_warmup.py
new file mode 100644
index 000000000000..512b3e6ba53a
--- /dev/null
+++ b/tests/kernels/test_mhc_jit_warmup.py
@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Validate JIT dispatch against pre-contract behavior."""
+
+from typing import Any
+
+import pytest
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_cuda_alike():
+    pytest.skip("NVIDIA dispatch tests require CUDA", allow_module_level=True)
+
+from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+    HcHeadFusedTileLangKernel,
+    HcPrenormGemmTileLangKernel,
+    MhcFusedTileLangKernel,
+    MhcPostTileLangKernel,
+    MhcPreBigFuseTileLangKernel,
+    mhc_fused_post_pre_split_config,
+    require_fused_post_pre_config,
+)
+from vllm.model_executor.warmup import jit_warmup_tilelang_helper
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "expected"),
+    [
+        (
+            dict(
+                num_tokens=64,
+                hc_hidden_size=4096,
+                hidden_size=2048,
+                hc_mult=2,
+                n_out=128,
+            ),
+            (2048, 2, 128, 1024, 4, 1, False, 1),
+        ),
+        (
+            dict(
+                num_tokens=1024,
+                hc_hidden_size=4096,
+                hidden_size=2048,
+                hc_mult=2,
+                n_out=128,
+            ),
+            (2048, 2, 128, 512, 12, 1, True, 2),
+        ),
+        (
+            dict(
+                num_tokens=64,
+                hc_hidden_size=4096,
+                hidden_size=2048,
+                hc_mult=2,
+                n_out=128,
+                n_thr=256,
+                tile_n=8,
+                n_splits=4,
+            ),
+            (2048, 2, 128, 256, 8, 4, False, 1),
+        ),
+    ],
+)
+def test_hc_prenorm_gemm_dispatch_matches_legacy_runtime_config(
+    kwargs: dict[str, Any],
+    expected: tuple[int, int, int, int, int, int, bool, int],
+) -> None:
+    kernel = HcPrenormGemmTileLangKernel()
+
+    assert kernel.dispatch(**kwargs) == kernel.CompileKey(*expected)
+
+
+@pytest.mark.parametrize(
+    ("is_broadcast", "use_norm_weight", "expected_use_norm", "expected_eps"),
+    [
+        (False, False, False, 0.0),
+        (False, True, True, 1.0e-5),
+        (True, False, True, 2.0e-5),
+    ],
+)
+def test_mhc_pre_big_fuse_dispatch_matches_legacy_runtime_config(
+    is_broadcast: bool,
+    use_norm_weight: bool,
+    expected_use_norm: bool,
+    expected_eps: float,
+) -> None:
+    kernel = MhcPreBigFuseTileLangKernel()
+
+    assert kernel.dispatch(
+        hidden_size=4096,
+        hc_mult=4,
+        n_splits=2,
+        is_broadcast=is_broadcast,
+        use_norm_weight=use_norm_weight,
+        rms_eps=1.0e-6,
+        hc_pre_eps=2.0e-6,
+        hc_sinkhorn_eps=3.0e-6,
+        hc_post_mult_value=0.5,
+        sinkhorn_repeat=3,
+        norm_eps=1.0e-5,
+        broadcast_norm_eps=2.0e-5,
+    ) == kernel.CompileKey(
+        hidden_size=4096,
+        hc_mult=4,
+        n_splits=2,
+        use_norm_weight=expected_use_norm,
+        is_broadcast=is_broadcast,
+        rms_eps=1.0e-6,
+        hc_pre_eps=2.0e-6,
+        hc_sinkhorn_eps=3.0e-6,
+        hc_post_mult_value=0.5,
+        sinkhorn_repeat=3,
+        norm_eps=expected_eps,
+    )
+
+
+@pytest.mark.parametrize(
+    ("num_tokens", "hidden_size", "expected_tile_n", "expected_n_splits"),
+    [
+        (1, 4096, 2, 8),
+        (4, 8192, 2, 8),
+        (8, 5120, 2, 8),
+        (16, 5120, 6, 8),
+        (32, 7168, 6, 8),
+    ],
+)
+def test_mhc_fused_dispatch_matches_the_launch_config(
+    num_tokens: int,
+    hidden_size: int,
+    expected_tile_n: int,
+    expected_n_splits: int,
+) -> None:
+    """The compile key must be the config the launch will actually use.
+
+    Both read mhc_fused_post_pre_split_config, so this pins the tuned bands
+    and guards against the two drifting apart again.
+    """
+    kernel = MhcFusedTileLangKernel()
+    tile_n, n_splits, n_thr = require_fused_post_pre_config(
+        num_tokens, hidden_size, hc_mult=4
+    )
+
+    assert (tile_n, n_splits) == (expected_tile_n, expected_n_splits)
+    assert kernel.dispatch(
+        num_tokens=num_tokens,
+        hidden_size=hidden_size,
+        hc_mult=4,
+    ) == kernel.CompileKey(
+        hidden_size=hidden_size,
+        hc_mult=4,
+        n_splits=n_splits,
+        tile_n=tile_n,
+        n_thr=n_thr,
+    )
+
+
+def test_mhc_fused_declines_shapes_it_cannot_tile() -> None:
+    """Above the token cutoff, and for a hidden size the block cannot split."""
+    assert mhc_fused_post_pre_split_config(33, 5120, 4) is None
+    assert mhc_fused_post_pre_split_config(1, 5137, 4) is None
+    with pytest.raises(ValueError, match="does not cover num_tokens"):
+        require_fused_post_pre_config(33, 5120, 4)
+
+
+@pytest.mark.parametrize(
+    ("kernel", "compile_key"),
+    [
+        (
+            HcPrenormGemmTileLangKernel(),
+            HcPrenormGemmTileLangKernel.CompileKey(
+                hidden_size=2048,
+                hc_mult=2,
+                n_out=128,
+                n_thr=1024,
+                tile_n=4,
+                n_splits=1,
+                use_block_m=False,
+                block_m=1,
+            ),
+        ),
+        (
+            HcPrenormGemmTileLangKernel(),
+            HcPrenormGemmTileLangKernel.CompileKey(
+                hidden_size=2048,
+                hc_mult=2,
+                n_out=128,
+                n_thr=512,
+                tile_n=12,
+                n_splits=1,
+                use_block_m=False,
+                block_m=1,
+            ),
+        ),
+        (
+            HcPrenormGemmTileLangKernel(),
+            HcPrenormGemmTileLangKernel.CompileKey(
+                hidden_size=2048,
+                hc_mult=2,
+                n_out=128,
+                n_thr=512,
+                tile_n=12,
+                n_splits=1,
+                use_block_m=True,
+                block_m=2,
+            ),
+        ),
+        (
+            MhcPreBigFuseTileLangKernel(),
+            MhcPreBigFuseTileLangKernel.CompileKey(
+                hidden_size=4096,
+                hc_mult=4,
+                n_splits=2,
+                use_norm_weight=True,
+                is_broadcast=False,
+                rms_eps=1.0e-6,
+                hc_pre_eps=2.0e-6,
+                hc_sinkhorn_eps=3.0e-6,
+                hc_post_mult_value=0.5,
+                sinkhorn_repeat=3,
+                norm_eps=1.0e-5,
+            ),
+        ),
+        (
+            MhcPreBigFuseTileLangKernel(),
+            MhcPreBigFuseTileLangKernel.CompileKey(
+                hidden_size=4096,
+                hc_mult=4,
+                n_splits=1,
+                use_norm_weight=False,
+                is_broadcast=False,
+                rms_eps=1.0e-6,
+                hc_pre_eps=2.0e-6,
+                hc_sinkhorn_eps=3.0e-6,
+                hc_post_mult_value=0.5,
+                sinkhorn_repeat=3,
+                norm_eps=0.0,
+            ),
+        ),
+        (
+            MhcPreBigFuseTileLangKernel(),
+            MhcPreBigFuseTileLangKernel.CompileKey(
+                hidden_size=4096,
+                hc_mult=4,
+                n_splits=2,
+                use_norm_weight=True,
+                is_broadcast=True,
+                rms_eps=1.0e-6,
+                hc_pre_eps=2.0e-6,
+                hc_sinkhorn_eps=3.0e-6,
+                hc_post_mult_value=0.5,
+                sinkhorn_repeat=3,
+                norm_eps=1.0e-5,
+            ),
+        ),
+        (
+            MhcPostTileLangKernel(),
+            MhcPostTileLangKernel.CompileKey(hidden_size=4096, hc_mult=4),
+        ),
+        (
+            MhcFusedTileLangKernel(),
+            MhcFusedTileLangKernel.CompileKey(
+                hidden_size=4096,
+                hc_mult=4,
+                n_splits=8,
+                tile_n=6,
+                n_thr=128,
+            ),
+        ),
+        (
+            MhcFusedTileLangKernel(),
+            MhcFusedTileLangKernel.CompileKey(
+                hidden_size=4096,
+                hc_mult=4,
+                n_splits=8,
+                tile_n=2,
+                n_thr=128,
+            ),
+        ),
+        (
+            MhcFusedTileLangKernel(),
+            MhcFusedTileLangKernel.CompileKey(
+                hidden_size=8192,
+                hc_mult=4,
+                n_splits=8,
+                tile_n=2,
+                n_thr=128,
+            ),
+        ),
+        (
+            HcHeadFusedTileLangKernel(),
+            HcHeadFusedTileLangKernel.CompileKey(
+                hidden_size=4096,
+                hc_mult=4,
+                rms_eps=1.0e-6,
+                hc_eps=2.0e-6,
+            ),
+        ),
+    ],
+)
+def test_tilelang_warmup_inputs_reproduce_compile_key(
+    monkeypatch: pytest.MonkeyPatch,
+    kernel: Any,
+    compile_key: Any,
+) -> None:
+    compiled: list[Any] = []
+    single_kernel = isinstance(
+        kernel,
+        (MhcPostTileLangKernel, MhcFusedTileLangKernel, HcHeadFusedTileLangKernel),
+    )
+    expected_kernel = kernel.kernel() if single_kernel else kernel.kernel(compile_key)
+    if single_kernel:
+        monkeypatch.setattr(
+            kernel,
+            "dispatch",
+            lambda **kwargs: pytest.fail("single-kernel launch must not dispatch"),
+        )
+    monkeypatch.setattr(
+        jit_warmup_tilelang_helper,
+        "compile_tilelang",
+        lambda jit_impl, *args, **kwargs: compiled.append(jit_impl),
+    )
+
+    kernel.compile(compile_key)
+
+    assert compiled == [expected_kernel]
diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py
index a146c1a5c973..4735fb96d18d 100644
--- a/tests/kernels/test_mhc_kernels.py
+++ b/tests/kernels/test_mhc_kernels.py
@@ -10,9 +10,19 @@
 import vllm.model_executor.kernels.mhc  # noqa: F401
 import vllm.model_executor.layers.mhc as mhc_layers
 from vllm.model_executor.kernels.mhc.tilelang import (
-    _tilelang_hc_prenorm_gemm,
     _torch_hc_prenorm_gemm,
+    mhc_post_tilelang,
+    mhc_pre_delayed_tilelang,
 )
+from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+    _HC_PRENORM_GEMM_TILELANG_KERNEL,
+    mhc_fused_post_pre_split_config,
+)
+from vllm.model_executor.kernels.mhc.torch import (
+    mhc_post_torch,
+    mhc_pre_delayed_torch,
+)
+from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton
 from vllm.model_executor.layers.mhc import (
     HAS_AITER_MHC,
     HAS_AITER_MHC_FUSED,
@@ -20,18 +30,832 @@
     HAS_AITER_MHC_PRE_NORM,
     HAS_TILELANG_MHC,
     MHCFusedPostPreOp,
+    MHCPreDelayedOp,
     MHCPreOp,
 )
 from vllm.models.deepseek_v4.nvidia.model import (
     DeepseekV4DecoderLayer,
     DeepseekV4Model,
 )
+from vllm.models.deepseek_v41.nvidia.model import (
+    DeepseekV4DecoderLayer as DeepseekV41DecoderLayer,
+)
+from vllm.models.deepseek_v41.nvidia.ops.mega_mhc import (
+    is_mega_mhc_supported,
+    mhc_shifted_post_pre_deep_gemm,
+)
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
 DEVICE = current_platform.device_type
 
 
+@pytest.mark.parametrize(
+    "tp,ep,hidden,hc,multicast,expected",
+    [
+        (4, False, 5120, 4, 1, True),
+        (2, False, 5120, 4, 1, False),
+        (4, True, 5120, 4, 1, False),
+        (4, False, 4096, 4, 1, False),
+        (4, False, 5120, 2, 1, False),
+        (4, False, 5120, 4, 0, False),
+    ],
+)
+def test_deepseek_v41_all_reduce_fusion_requires_kernel_support(
+    monkeypatch, tp, ep, hidden, hc, multicast, expected
+):
+    """Eligibility depends on kernel inputs, not serving or attention settings."""
+    from vllm.models.deepseek_v41.nvidia.ops import mhc
+
+    config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=tp, enable_expert_parallel=ep
+        ),
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(hidden_size=hidden, hc_mult=hc)
+        ),
+    )
+    group = SimpleNamespace(
+        device_communicator=SimpleNamespace(
+            ca_comm=SimpleNamespace(mnnvl_lamport_ag_multicast_ptr=multicast)
+        )
+    )
+    monkeypatch.setattr(mhc, "get_tp_group", lambda: group)
+    assert mhc.supports_mhc_all_reduce(config) is expected
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "num_tokens,hc_mult,hidden_size",
+    [
+        (0, 4, 5120),
+        (1, 4, 5120),
+        (33, 4, 5137),
+        (128, 4, 5120),
+        (1024, 4, 5120),
+        (2, 1, 513),
+        (7, 8, 5137),
+        (65537, 4, 65),
+    ],
+)
+@pytest.mark.parametrize("strided", [False, True])
+def test_hc_collapse_preserves_weighted_residual_sum(
+    num_tokens, hc_mult, hidden_size, strided
+):
+    """Handle tile tails and strided streams without changing FP32 mixing."""
+    set_random_seed(0)
+    x = torch.randn(
+        num_tokens,
+        hc_mult,
+        hidden_size * (2 if strided else 1),
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    pre = torch.rand(num_tokens, hc_mult, device=DEVICE)
+    if strided:
+        x = x[..., ::2]
+        pre = pre.t().contiguous().t()
+    if num_tokens:
+        pre[0].zero_()
+        pre[0, -1] = 1
+    if num_tokens > 1 and hc_mult == 4:
+        pre[1].fill_(1)
+        x[1, 1] = -x[1, 0]
+        x[1, 3] = -x[1, 2]
+    expected = (pre.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype)
+    actual = hc_collapse_triton(x, pre)
+    assert actual.is_contiguous()
+    # BF16 rounding only: ptxas <13.1 contracts the FP32 mix into FFMA despite
+    # enable_fp_fusion=False, so compare within a ULP and pin the exact cases below.
+    torch.testing.assert_close(actual, expected, atol=1.6e-2, rtol=1e-2)
+    if num_tokens:
+        torch.testing.assert_close(actual[0], x[0, -1], atol=0, rtol=0)
+    if num_tokens > 1 and hc_mult == 4:
+        torch.testing.assert_close(
+            actual[1], torch.zeros_like(actual[1]), atol=0, rtol=0
+        )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_hc_collapse_custom_op_supports_compile():
+    x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
+    pre = torch.rand(2, 4, device=DEVICE)
+    torch.library.opcheck(torch.ops.vllm.hc_collapse_triton.default, (x, pre))
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("num_tokens", [1, 7])
+def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch):
+    """Match reference DSparkBlock.forward_head's hc_pre before its RMSNorm."""
+    from vllm.models.deepseek_v41.nvidia import dspark
+
+    set_random_seed(0)
+    hidden_size, hc_mult = 5120, 4
+    streams = torch.randn(
+        num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
+    )
+    mixes = [torch.rand(num_tokens, hc_mult, device=DEVICE) for _ in range(2)]
+    carried_mixes = []
+
+    def make_layer(mix):
+        def forward(hidden, positions, ids, pre, post, res, residual, **kwargs):
+            carried_mixes.append(pre)
+            return hidden, None, None, None, mix, None
+
+        return forward
+
+    monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams)
+    draft = SimpleNamespace(
+        use_mega_moe=False,
+        use_sequence_parallel=False,
+        hc_mult=hc_mult,
+        layers=[make_layer(mix) for mix in mixes],
+    )
+    actual = dspark.DSparkDeepseekV4Model.forward(
+        draft,
+        input_ids=torch.zeros(num_tokens, dtype=torch.long, device=DEVICE),
+        positions=torch.arange(num_tokens, device=DEVICE),
+        inputs_embeds=torch.zeros(
+            num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE
+        ),
+    )
+    expected = (streams.float() * mixes[-1].unsqueeze(-1)).sum(dim=1)
+    assert carried_mixes[0] is None
+    assert carried_mixes[1] is mixes[0]
+    torch.testing.assert_close(
+        actual, expected.to(streams.dtype), atol=1.6e-2, rtol=1e-2
+    )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("num_tokens", [1, 7, 32])
+@pytest.mark.parametrize("tp_size", [4, 8])
+def test_hc_collapse_before_gather_is_bitwise_equal(num_tokens, tp_size):
+    """Simulate SP on one GPU, including padding and ranks without valid tokens."""
+    set_random_seed(0)
+    padded_tokens = num_tokens + (-num_tokens) % tp_size
+    streams = torch.randn(padded_tokens, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
+    pre_mix = torch.rand(padded_tokens, 4, device=DEVICE)
+    stream_shards = streams.chunk(tp_size)
+    mix_shards = pre_mix.chunk(tp_size)
+    before = hc_collapse_triton(
+        torch.cat(stream_shards)[:num_tokens], torch.cat(mix_shards)[:num_tokens]
+    )
+    after = torch.cat(
+        [hc_collapse_triton(x, pre) for x, pre in zip(stream_shards, mix_shards)]
+    )[:num_tokens]
+    assert after.shape == (num_tokens, 5120)
+    torch.testing.assert_close(after, before, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize(
+    "num_tokens,sinkhorn_repeat,norm_eps",
+    [
+        (0, 20, 1e-6),
+        (1, 1, 1e-6),
+        (1, 20, 1e-6),
+        (33, 20, 1e-20),
+        (128, 20, 1e-6),
+        (1024, 20, 1e-6),
+    ],
+)
+@pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"])
+@pytest.mark.parametrize("use_deep_gemm", [False, True])
+@pytest.mark.parametrize("fused_norm", [False, True])
+def test_deepseek_v41_mhc_pre_delayed(
+    num_tokens,
+    sinkhorn_repeat,
+    norm_eps,
+    entry,
+    use_deep_gemm,
+    fused_norm,
+    monkeypatch,
+):
+    """Collapse with the carried pre-mix and save the new one for the next sublayer."""
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    if use_deep_gemm and not is_deep_gemm_supported():
+        pytest.skip("DeepGEMM not available")
+    monkeypatch.setattr(
+        "vllm.utils.deep_gemm.is_deep_gemm_supported", lambda: use_deep_gemm
+    )
+    set_random_seed(0)
+    hc_mult, hidden_size = 4, 5120
+    residual = torch.randn(
+        num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
+    )
+    if num_tokens > 1:
+        residual[0].zero_()
+    fn = torch.randn(24, hc_mult, hidden_size, device=DEVICE) * 0.02
+    x = None
+    pre_mix = None
+    if entry == "broadcast":
+        x = residual[:, 0].contiguous()
+        residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()
+        fn = fn.sum(1)
+    else:
+        fn = fn.flatten(1)
+        if entry == "carried":
+            pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE)
+            if num_tokens > 1:
+                # A constant collapse just above 1 rounds to BF16 1 before RMSNorm.
+                residual[1].fill_(1)
+                pre_mix[1].zero_()
+                pre_mix[1, 0] = 1.003
+    scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE)
+    base = torch.randn(24, device=DEVICE)
+    args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, sinkhorn_repeat)
+    expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix, x=x)
+    norm_kwargs = {}
+    if fused_norm:
+        weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE)
+        weight.uniform_(0.5, 1.5)
+        norm_kwargs = dict(norm_weight=weight, norm_eps=norm_eps)
+        # Compare to the original collapse + CUDA RMSNorm, including BF16 rounding.
+        unfused = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x)
+        normalized = torch.empty_like(unfused[2])
+        if num_tokens:
+            torch.ops._C.rms_norm(normalized, unfused[2], weight, norm_eps)
+        expected = (*unfused[:2], normalized, unfused[3])
+    actual = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x, **norm_kwargs)
+    if fused_norm and num_tokens:
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = mhc_pre_delayed_tilelang(
+                *args, pre_mix=pre_mix, x=x, **norm_kwargs
+            )
+        graph.replay()
+        for eager, replayed in zip(actual, captured, strict=True):
+            torch.testing.assert_close(eager, replayed, atol=0, rtol=0)
+    if fused_norm and entry == "carried" and num_tokens > 1:
+        torch.testing.assert_close(actual[2][1], expected[2][1], atol=0, rtol=0)
+    # DeepGEMM uses TF32 for the projection; the TileLang fallback uses FP32.
+    atol, rtol = (1e-3, 1e-3) if use_deep_gemm else (1e-5, 1e-4)
+    for i in (0, 1, 3):
+        torch.testing.assert_close(actual[i], expected[i], atol=atol, rtol=rtol)
+    if pre_mix is None and not fused_norm:
+        torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0)
+    else:
+        torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not HAS_TILELANG_MHC or not current_platform.is_device_capability_family(100),
+    reason="SM100 TileLang mHC required",
+)
+@pytest.mark.parametrize("num_tokens", [0, 1, 8, 9, 32, 128, 256])
+@pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"])
+@pytest.mark.parametrize("hidden_size,hc_mult", [(4096, 2), (5120, 4), (7168, 4)])
+def test_deepseek_v41_mhc_overlap_preserves_outputs_on_replay(
+    num_tokens, entry, hidden_size, hc_mult
+):
+    """The input is usable before joining; coefficients stay exact on graph replay."""
+    from vllm.models.deepseek_v41.nvidia.ops.mhc import mhc_pre_delayed_overlap
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    if not is_deep_gemm_supported():
+        pytest.skip("DeepGEMM required")
+    set_random_seed(42)
+    mix_size = hc_mult * (hc_mult + 2)
+    residual = torch.randn(
+        num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
+    )
+    fn = torch.randn(mix_size, hc_mult, hidden_size, device=DEVICE) * 0.02
+    x = None
+    pre_mix = None
+    if entry == "broadcast":
+        x = residual[:, 0].contiguous()
+        residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()
+        fn = fn.sum(1)
+    else:
+        fn = fn.flatten(1)
+        if entry == "carried":
+            pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE)
+    scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE)
+    base = torch.randn(mix_size, device=DEVICE)
+    weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE).uniform_(
+        0.5, 1.5
+    )
+    args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20)
+    kwargs = dict(pre_mix=pre_mix, x=x, norm_weight=weight, norm_eps=1e-20)
+    side = torch.cuda.Stream()
+
+    def run():
+        outputs = mhc_pre_delayed_overlap(*args, **kwargs, stream=side)
+        consumed_input = outputs[2].clone()
+        torch.cuda.current_stream().wait_stream(side)
+        return (*outputs, consumed_input)
+
+    def check(actual):
+        expected = mhc_pre_delayed_tilelang(*args, **kwargs)
+        for a, b in zip(actual[:4], expected, strict=True):
+            torch.testing.assert_close(a, b, atol=0, rtol=0)
+        torch.testing.assert_close(actual[4], expected[2], atol=0, rtol=0)
+
+    check(run())
+    if num_tokens == 0:
+        return
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = run()
+    for _ in range(5):
+        if x is not None:
+            x.normal_()
+            residual.copy_(x.unsqueeze(1).expand_as(residual))
+        else:
+            residual.normal_()
+        if pre_mix is not None:
+            pre_mix.uniform_()
+        graph.replay()
+        check(captured)
+
+
+def mhc_fused_post_pre_delayed_ref(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    *mix_args,
+    pre_mix: torch.Tensor | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Project the post result before it is rounded to BF16.
+
+    The fused kernel keeps the updated residual streams in registers, so its
+    projection sees the FP32 post result while the streams it stores are the
+    BF16 rounding of it.
+    """
+    post = torch.einsum("tij,tih->tjh", comb_res_mix, residual.float())
+    post = post + post_layer_mix * x.float().unsqueeze(-2)
+    return mhc_pre_delayed_torch(
+        post.bfloat16(), *mix_args, pre_mix=pre_mix, x=post.flatten(1)
+    )
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("num_tokens", [0, 1, 8, 128])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+@pytest.mark.parametrize("carried", [False, True])
+def test_deepseek_v41_mhc_fused_post_pre_delayed(num_tokens, hidden_size, carried):
+    """Fold the post block into the delayed pre without changing its outputs."""
+    set_random_seed(0)
+    hc_mult = 4
+    mix_size = hc_mult * (hc_mult + 2)
+    x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE)
+    residual = torch.randn(
+        num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
+    )
+    post_layer_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE)
+    comb_res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE)
+    fn = torch.randn(mix_size, hc_mult * hidden_size, device=DEVICE) * 0.02
+    scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE)
+    base = torch.randn(mix_size, device=DEVICE)
+    pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) if carried else None
+    weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE)
+    weight.uniform_(0.5, 1.5)
+    mix_args = (fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20)
+
+    residual_cur, post, comb, layer_input, next_pre, aux = (
+        torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang(
+            x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6
+        )
+    )
+    assert aux.shape == (0, hidden_size)
+    assert [
+        residual_cur.shape,
+        post.shape,
+        comb.shape,
+        layer_input.shape,
+        next_pre.shape,
+    ] == [
+        (num_tokens, hc_mult, hidden_size),
+        (num_tokens, hc_mult, 1),
+        (num_tokens, hc_mult, hc_mult),
+        (num_tokens, hidden_size),
+        (num_tokens, hc_mult),
+    ]
+    if num_tokens == 0:
+        return
+
+    # The post mapping and the collapse are what the unfused kernels produce.
+    residual_ref = torch.ops.vllm.mhc_post_tilelang(
+        x, residual, post_layer_mix, comb_res_mix
+    )
+    layer_input_ref = mhc_pre_delayed_tilelang(
+        residual_ref, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6
+    )[2]
+    torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=0)
+    torch.testing.assert_close(layer_input, layer_input_ref, atol=0, rtol=0)
+
+    post_ref, comb_ref, _, next_pre_ref = mhc_fused_post_pre_delayed_ref(
+        x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix=pre_mix
+    )
+    # Above the fused kernel's token count the projection falls back to a
+    # split-k TF32 GEMM over the stored BF16 streams.
+    tol = (
+        1e-5
+        if mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) is not None
+        else 5e-3
+    )
+    for actual, expected in (
+        (post, post_ref),
+        (comb, comb_ref),
+        (next_pre, next_pre_ref),
+    ):
+        torch.testing.assert_close(actual, expected, atol=tol, rtol=tol)
+
+    # Folding the stream mean into the collapse must not change it, or anything
+    # else the op returns.
+    captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+        *mix_args,
+        pre_mix,
+        weight,
+        1e-6,
+        True,
+    )
+    torch.testing.assert_close(captured[5], residual_cur.mean(dim=1), atol=0, rtol=0)
+    unchanged = (residual_cur, post, comb, layer_input, next_pre)
+    for with_aux, without in zip(captured[:5], unchanged, strict=True):
+        torch.testing.assert_close(with_aux, without, atol=0, rtol=0)
+
+    # The unnormalized epilogue has its own fused aux path.
+    bare = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang(
+        x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, None, 1e-6, True
+    )
+    torch.testing.assert_close(bare[5], bare[0].mean(dim=1), atol=0, rtol=0)
+
+    # Decode replays this op from a captured graph.
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang(
+            x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6
+        )
+    graph.replay()
+    eager = (residual_cur, post, comb, layer_input, next_pre, aux)
+    for replayed, expected in zip(captured, eager, strict=True):
+        torch.testing.assert_close(replayed, expected, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("carried", [False, True])
+@pytest.mark.parametrize("capture_aux", [False, True])
+def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_aux):
+    set_random_seed(0)
+    x = torch.randn(2, 5120, dtype=torch.bfloat16, device=DEVICE)
+    residual = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
+    post_layer_mix = torch.rand(2, 4, 1, device=DEVICE)
+    comb_res_mix = torch.rand(2, 4, 4, device=DEVICE)
+    fn = torch.randn(24, 20480, device=DEVICE) * 0.02
+    pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None
+    scale = torch.ones(3, device=DEVICE)
+    base = torch.zeros(24, device=DEVICE)
+    torch.library.opcheck(
+        torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang.default,
+        (
+            x,
+            residual,
+            post_layer_mix,
+            comb_res_mix,
+            fn,
+            scale,
+            base,
+            1e-20,
+            1e-6,
+            1e-6,
+            2.0,
+            20,
+            pre_mix,
+            None,
+            1e-6,
+            capture_aux,
+        ),
+    )
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"])
+@pytest.mark.parametrize(
+    "mhc_mode", ["disabled", "eager", "overlap", "piecewise", "large_batch"]
+)
+def test_deepseek_v41_decoder_mixes_match_torch(
+    entry, mhc_mode, monkeypatch, default_vllm_config
+):
+    """Preserve carried pre-mixes and Engram ordering at every decoder entry."""
+    from vllm.models.deepseek_v41.nvidia.ops.mhc import MHC_OVERLAP_MAX_TOKENS
+
+    num_tokens = MHC_OVERLAP_MAX_TOKENS + 1 if mhc_mode == "large_batch" else 3
+    set_random_seed(0)
+    decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer)
+    nn.Module.__init__(decoder)
+    decoder.hc_mult = 4
+    decoder.hc_sinkhorn_iters = 20
+    decoder.hc_eps = 1e-6
+    decoder.rms_norm_eps = 1e-20
+    decoder.hc_post_alpha = 2.0
+    decoder.use_sequence_parallel = False
+    decoder.fuse_mhc_all_reduce = False
+    decoder.mhc_stream = None
+    if mhc_mode != "disabled":
+        from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+        if (
+            not current_platform.is_device_capability_family(100)
+            or not is_deep_gemm_supported()
+        ):
+            pytest.skip("SM100 DeepGEMM required for overlap")
+        decoder.mhc_stream = torch.cuda.Stream()
+    if mhc_mode in ("eager", "piecewise", "large_batch"):
+
+        def unexpected_overlap(*args, **kwargs):
+            pytest.fail("unsupported execution must retain native mHC")
+
+        monkeypatch.setattr(
+            "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_overlap",
+            unexpected_overlap,
+        )
+        monkeypatch.setattr(
+            "vllm.models.deepseek_v41.nvidia.ops.mhc.mhc_pre_delayed_overlap",
+            unexpected_overlap,
+        )
+    decoder.engram = None
+    from vllm.model_executor.layers.layernorm import RMSNorm
+
+    decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to(
+        device=DEVICE, dtype=torch.bfloat16
+    )
+    decoder.attn = lambda positions, x, _: x * 0.5
+    decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25
+    with torch.device(DEVICE):
+        decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02
+        decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02
+        decoder.hc_attn_fn_broadcast = decoder.hc_attn_fn.view(24, 4, 5120).sum(1)
+        decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3)
+        decoder.hc_attn_base = torch.randn(24)
+        decoder.hc_ffn_base = torch.randn(24)
+        x = torch.randn(num_tokens, 5120, dtype=torch.bfloat16)
+        positions = torch.arange(num_tokens)
+        kwargs = {}
+        if entry != "broadcast":
+            kwargs["pre_mix"] = torch.rand(num_tokens, 4)
+        if entry == "pipeline":
+            x = torch.randn(num_tokens, 4, 5120, dtype=torch.bfloat16)
+        if entry in ("residual", "engram"):
+            kwargs.update(
+                residual=torch.randn(num_tokens, 4, 5120, dtype=torch.bfloat16),
+                post_mix=torch.rand(num_tokens, 4, 1),
+                res_mix=torch.rand(num_tokens, 4, 4),
+            )
+        if entry == "engram":
+
+            class FakeEngram(nn.Module):
+                layer_hash_index = 0
+
+                def forward(self, residual, hashes, mask):
+                    return residual + 0.125
+
+            decoder.engram = FakeEngram()
+            kwargs["engram_hashes"] = torch.zeros(num_tokens, 1, 1, dtype=torch.int32)
+
+    from vllm.config import CUDAGraphMode
+    from vllm.forward_context import set_forward_context
+
+    mode = CUDAGraphMode.PIECEWISE if mhc_mode == "piecewise" else CUDAGraphMode.NONE
+    overlap_calls = []
+    if mhc_mode == "overlap":
+        from vllm.models.deepseek_v41.nvidia.ops.mhc import mhc_pre_delayed_overlap
+
+        def checked_overlap(*args, **kwargs):
+            overlap_calls.append(torch.cuda.is_current_stream_capturing())
+            return mhc_pre_delayed_overlap(*args, **kwargs)
+
+        for module in ("model", "ops.mhc"):
+            monkeypatch.setattr(
+                f"vllm.models.deepseek_v41.nvidia.{module}.mhc_pre_delayed_overlap",
+                checked_overlap,
+            )
+    with set_forward_context(None, default_vllm_config, cudagraph_runtime_mode=mode):
+        actual = decoder(x, positions, None, **kwargs)
+        if mhc_mode == "overlap":
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                actual = decoder(x, positions, None, **kwargs)
+            graph.replay()
+            assert overlap_calls and all(overlap_calls)
+
+    def reference(*args, norm_weight, norm_eps, **kwargs):
+        post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs)
+        return post, res, decoder.attn_norm(collapsed), pre
+
+    def fused_reference(
+        x,
+        residual,
+        post_mix,
+        res_mix,
+        *args,
+        capture_aux=False,
+        stream=None,
+        reduce_results=False,
+        **kw,
+    ):
+        residual = mhc_post_torch(x, residual, post_mix, res_mix)
+        aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0)
+        return residual, *reference(residual, *args, **kw), aux
+
+    monkeypatch.setattr(
+        "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang",
+        reference,
+    )
+    decoder.mhc_stream = None
+    monkeypatch.setattr(
+        "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre",
+        fused_reference,
+    )
+    expected = decoder(x, positions, None, **kwargs)
+    assert actual[-1] is None and expected[-1] is None
+    for result, ref in zip(actual[:-1], expected[:-1], strict=True):
+        torch.testing.assert_close(result, ref, atol=2e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("entry", ["fused", "engram"])
+def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_config):
+    """Read the aux hidden state back out of the seam that already computed it.
+
+    The draft model's input is the mean over hc streams of the previous
+    layer's post. That post now runs inside this layer's fused pre, so the
+    captured value must equal a standalone ``mhc_post_tilelang`` on the same
+    inputs -- and Engram must not be folded into it, since the injection
+    happens after the aux consumers read the stream.
+    """
+    set_random_seed(0)
+    decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer)
+    nn.Module.__init__(decoder)
+    decoder.hc_mult = 4
+    decoder.hc_sinkhorn_iters = 20
+    decoder.hc_eps = 1e-6
+    decoder.rms_norm_eps = 1e-6
+    decoder.hc_post_alpha = 2.0
+    decoder.use_sequence_parallel = False
+    decoder.fuse_mhc_all_reduce = False
+    decoder.mhc_stream = None
+    decoder.engram = None
+    from vllm.model_executor.layers.layernorm import RMSNorm
+
+    decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to(
+        device=DEVICE, dtype=torch.bfloat16
+    )
+    decoder.attn = lambda positions, x, _: x * 0.5
+    decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25
+    with torch.device(DEVICE):
+        decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02
+        decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02
+        decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3)
+        decoder.hc_attn_base = torch.randn(24)
+        decoder.hc_ffn_base = torch.randn(24)
+        x = torch.randn(3, 5120, dtype=torch.bfloat16)
+        kwargs = dict(
+            pre_mix=torch.rand(3, 4),
+            residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16),
+            post_mix=torch.rand(3, 4, 1),
+            res_mix=torch.rand(3, 4, 4),
+        )
+        if entry == "engram":
+
+            class FakeEngram(nn.Module):
+                layer_hash_index = 0
+
+                def forward(self, residual, hashes, mask):
+                    return residual + 0.125
+
+            decoder.engram = FakeEngram()
+            kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32)
+        positions = torch.arange(3)
+
+    expected = torch.ops.vllm.mhc_post_tilelang(
+        x, kwargs["residual"], kwargs["post_mix"], kwargs["res_mix"]
+    ).mean(dim=1)
+    previous_aux = decoder(x, positions, None, **kwargs, capture_previous_aux=True)[-1]
+    torch.testing.assert_close(previous_aux, expected, atol=0, rtol=0)
+    assert decoder(x, positions, None, **kwargs)[-1] is None
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("carried", [False, True])
+def test_mhc_pre_delayed_custom_op_supports_compile(carried):
+    set_random_seed(0)
+    x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
+    fn = torch.randn(24, 20480, device=DEVICE) * 0.02
+    pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None
+    scale = torch.ones(3, device=DEVICE)
+    base = torch.zeros(24, device=DEVICE)
+    torch.library.opcheck(
+        torch.ops.vllm.mhc_pre_delayed_tilelang.default,
+        (x, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20, pre_mix),
+    )
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="DeepGEMM Mega mHC requires SM100-family CUDA",
+)
+@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 1024])
+@pytest.mark.parametrize("hidden_size", [5120, 7168])
+def test_deep_gemm_mega_mhc_correctness(num_tokens, hidden_size):
+    set_random_seed(0)
+    hc_mult = 4
+    if not is_mega_mhc_supported(hidden_size, 4):
+        pytest.skip("DeepGEMM Mega mHC is not available")
+
+    x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE)
+    residual = torch.randn(
+        num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
+    )
+    previous_mix = torch.rand(num_tokens, hc_mult, device=DEVICE)
+    post_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE)
+    res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE)
+    fn = torch.randn(24, hc_mult * hidden_size, device=DEVICE) * 0.01
+    scale = torch.randn(3, device=DEVICE) * 0.1
+    base = torch.randn(24, device=DEVICE) * 0.1
+    norm_weight = torch.empty(
+        hidden_size, dtype=torch.bfloat16, device=DEVICE
+    ).uniform_(0.9, 1.1)
+
+    expected_residual = (
+        mhc_post_tilelang(x, residual, post_mix, res_mix) if num_tokens else residual
+    )
+    (
+        expected_post_mix,
+        expected_res_mix,
+        expected_y_bf16,
+        expected_previous_mix,
+    ) = mhc_pre_delayed_tilelang(
+        expected_residual,
+        fn,
+        scale,
+        base,
+        2e-5,
+        3e-4,
+        2e-6,
+        1.25,
+        10,
+        pre_mix=previous_mix,
+        norm_weight=norm_weight,
+        norm_eps=7e-6,
+    )
+    args = (
+        x,
+        residual,
+        previous_mix,
+        post_mix,
+        res_mix,
+        fn,
+        scale,
+        base,
+        2e-5,
+        3e-4,
+        1.25,
+        2e-6,
+        10,
+        norm_weight,
+        7e-6,
+    )
+    actual = mhc_shifted_post_pre_deep_gemm(*args)
+    expected = (
+        expected_residual,
+        expected_post_mix,
+        expected_res_mix,
+        expected_y_bf16,
+        expected_previous_mix,
+    )
+    for i, (result, ref) in enumerate(zip(actual, expected, strict=True)):
+        # Match the existing mHC BF16 tolerance: post rounding differences
+        # can be amplified by cancellation in the carried collapse.
+        atol, rtol = (1.6e-2, 1e-2) if i in (0, 3) else (1e-6, 1e-3)
+        torch.testing.assert_close(result, ref, atol=atol, rtol=rtol)
+
+    if num_tokens:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            # DeepGEMM initializes barriers per stream before capture.
+            mhc_shifted_post_pre_deep_gemm(*args)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                captured = mhc_shifted_post_pre_deep_gemm(*args)
+            for multiplier in (0.5, -1.0):
+                x.mul_(multiplier)
+                graph.replay()
+                eager = mhc_shifted_post_pre_deep_gemm(*args)
+                for result, ref in zip(captured, eager, strict=True):
+                    torch.testing.assert_close(result, ref, atol=0, rtol=0)
+        torch.cuda.current_stream().wait_stream(stream)
+
+
 def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor:
     x = x.softmax(-1) + eps
     x = x / (x.sum(-2, keepdim=True) + eps)
@@ -52,7 +876,7 @@ def mhc_pre_ref(
     hc_post_mult_value: float,
     sinkhorn_repeat: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """mHC pre reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_pre.py#L303"""
+    """MHC pre reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_pre.py#L303."""
     hc_mult = residual.shape[-2]
 
     residual_flat = residual.flatten(-2, -1).float()
@@ -91,7 +915,7 @@ def mhc_post_ref(
     post_layer_mix: torch.Tensor,
     comb_res_mix: torch.Tensor,
 ) -> torch.Tensor:
-    """mHC post reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_post.py#L68"""
+    """MHC post reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_post.py#L68."""
     term2 = torch.bmm(comb_res_mix.mT, residual.float())
     return (x.float().unsqueeze(-2) * post_layer_mix + term2).bfloat16()
 
@@ -120,7 +944,8 @@ def hc_head_ref(
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
 @pytest.mark.parametrize("hc_mult", [4])
-def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
+@pytest.mark.parametrize("fused_norm", [False, True])
+def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult, fused_norm):
     torch.set_default_device(DEVICE)
     set_random_seed(0)
 
@@ -150,6 +975,12 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
         hc_post_alpha,
         sinkhorn_repeat,
     )
+    norm_kwargs = {}
+    if fused_norm:
+        # Non-shifted mHC must still use its freshly computed pre-mix.
+        weight = torch.ones(hidden_size, dtype=torch.bfloat16)
+        norm_kwargs = dict(norm_weight=weight, norm_eps=rms_eps)
+        ref = (*ref[:2], F.rms_norm(ref[2], (hidden_size,), weight, rms_eps))
     out = torch.ops.vllm.mhc_pre_tilelang(
         residual,
         fn,
@@ -160,6 +991,7 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
         hc_sinkhorn_eps,
         hc_post_alpha,
         sinkhorn_repeat,
+        **norm_kwargs,
     )
 
     for actual, expected in zip(out, ref, strict=True):
@@ -200,10 +1032,17 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
     sqrsum = torch.empty_like(sqrsum_ref)
 
     _torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref)
-    _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult)
+    _HC_PRENORM_GEMM_TILELANG_KERNEL(
+        x,
+        fn,
+        out,
+        sqrsum,
+        hidden_size,
+        hc_mult,
+    )
 
     torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4)
-    torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4)
+    torch.testing.assert_close(sqrsum, sqrsum_ref, atol=1e-2, rtol=1e-6)
 
 
 @pytest.mark.skipif(
@@ -413,6 +1252,109 @@ def test_mhc_fused_rocm_fallback_applies_norm(monkeypatch):
     torch.testing.assert_close(out[3], expected_layer_input)
 
 
+@pytest.mark.skipif(
+    not (current_platform.is_rocm() and HAS_AITER_MHC),
+    reason="AITER mHC required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128, 1024])
+@pytest.mark.parametrize("carried", [False, True])
+def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried):
+    """AITER must reproduce the delayed reference on both seam variants.
+
+    ``num_tokens`` spans the split-k choices AITER makes for the projection,
+    since the pre-mix is recovered from that unreduced output.
+    """
+    set_random_seed(0)
+    hc_mult, hidden_size = 4, 5120
+    residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs(
+        num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult
+    )
+    pre_mix = (
+        torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5
+        if carried
+        else None
+    )
+    rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6
+    args = (
+        residual,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        1.0,
+        20,
+    )
+
+    expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix)
+    residual_out, *actual = object.__new__(MHCPreDelayedOp).forward_hip(
+        *args, pre_mix=pre_mix
+    )
+
+    # No post was requested, so the residual comes straight back.
+    assert residual_out is residual
+    for i in (0, 1, 3):
+        torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3)
+    # The collapse is the same FP32 multiply-and-sum in both paths.
+    torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_rocm() and HAS_AITER_MHC),
+    reason="AITER mHC required",
+)
+def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch):
+    """The broadcast seam and a fused norm must not take the AITER path.
+
+    Neither is expressible with AITER's pre kernels: the broadcast projects a
+    narrower ``x``, and ``mhc_pre_gemm_sqrsum`` folds no RMSNorm. What is
+    asserted here is the routing decision, which is what this gate owns; the
+    numerics of whichever fallback it lands on are covered by
+    ``test_deepseek_v41_mhc_pre_delayed``.
+    """
+    set_random_seed(0)
+    hc_mult, hidden_size = 4, 5120
+    residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs(
+        num_tokens=4, hidden_size=hidden_size, hc_mult=hc_mult
+    )
+    args = (residual, fn, hc_scale, hc_base, 1e-6, 1e-6, 1e-6, 1.0, 20)
+    op = object.__new__(MHCPreDelayedOp)
+
+    aiter_op = torch.ops.vllm.mhc_pre_delayed_aiter
+    took_aiter = False
+
+    def spy(*spy_args, **spy_kwargs):
+        nonlocal took_aiter
+        took_aiter = True
+        return aiter_op(*spy_args, **spy_kwargs)
+
+    monkeypatch.setattr(torch.ops.vllm, "mhc_pre_delayed_aiter", spy)
+
+    x = residual[:, 0].contiguous()
+    broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1)
+    broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()
+    expected = mhc_pre_delayed_torch(broadcast_residual, broadcast_fn, *args[2:], x=x)
+    # forward_hip leads with the residual; no post was requested here, so it is
+    # the caller's own tensor and the reference's four outputs follow it.
+    residual_out, *actual = op.forward_hip(
+        broadcast_residual, broadcast_fn, *args[2:], x=x
+    )
+    assert residual_out is broadcast_residual
+    assert not took_aiter, "the broadcast seam must not reach AITER"
+    for i in range(4):
+        torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3)
+
+    op.forward_hip(*args, norm_weight=norm_weight, norm_eps=1e-6)
+    assert not took_aiter, "a fused norm must not reach AITER"
+
+    # Positive control: the same inputs without a norm do take the AITER path,
+    # so the two declines above are the gate discriminating rather than the
+    # path being unavailable in this environment.
+    op.forward_hip(*args)
+    assert took_aiter
+
+
 @pytest.mark.skipif(
     not (current_platform.is_rocm() and HAS_AITER_MHC and HAS_AITER_MHC_PRE_NORM),
     reason="AITER mHC with fused RMSNorm required",
@@ -640,3 +1582,68 @@ def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch):
     assert layer.hc_attn_fn_broadcast is buffer
     expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1)
     assert torch.equal(layer.hc_attn_fn_broadcast, expected)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_rocm() and HAS_AITER_MHC),
+    reason="AITER mHC required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128])
+@pytest.mark.parametrize("carried", [False, True])
+def test_mhc_fused_post_pre_delayed_rocm_aiter(
+    num_tokens, carried, default_vllm_config
+):
+    """Folding the post into the pre projection must match doing them apart.
+
+    The residual is bfloat16, so the two accumulation orders are allowed to
+    land a rounding step apart; the tolerance here is one bfloat16 ULP.
+    """
+    set_random_seed(0)
+    hc_mult, hidden_size = 4, 5120
+    residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs(
+        num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult
+    )
+    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=DEVICE)
+    post_layer_mix = (
+        torch.rand(num_tokens, hc_mult, 1, dtype=torch.float32, device=DEVICE) + 0.5
+    )
+    comb_res_mix = (
+        torch.rand(num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=DEVICE)
+        + 0.2
+    )
+    pre_mix = (
+        torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5
+        if carried
+        else None
+    )
+    rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6
+    pre_args = (fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, 1.0, 20)
+
+    next_residual = mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
+    expected = (
+        next_residual,
+        *mhc_pre_delayed_torch(next_residual, *pre_args, pre_mix=pre_mix),
+    )
+    actual = MHCPreDelayedOp().forward_hip(
+        residual,
+        *pre_args,
+        pre_mix=pre_mix,
+        sublayer_out=x,
+        post_layer_mix=post_layer_mix,
+        comb_res_mix=comb_res_mix,
+    )
+
+    for got, want in zip(actual, expected, strict=True):
+        torch.testing.assert_close(got, want.reshape(got.shape), atol=2e-2, rtol=8e-3)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_rocm() and HAS_AITER_MHC),
+    reason="AITER mHC required",
+)
+def test_mhc_fused_post_pre_delayed_falls_back_for_large_batches():
+    """AITER's heuristic hands large batches back to the unfused path."""
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    assert not rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1)
+    assert rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1 << 20)
diff --git a/tests/kernels/test_mhc_tilelang_jit.py b/tests/kernels/test_mhc_tilelang_jit.py
index 4f39fae02058..1c430e7ae338 100644
--- a/tests/kernels/test_mhc_tilelang_jit.py
+++ b/tests/kernels/test_mhc_tilelang_jit.py
@@ -24,19 +24,26 @@ class _PassConfigKey:
 def _install_tilelang_stub(
     monkeypatch: pytest.MonkeyPatch,
 ) -> dict[str, int]:
-    calls = {"jit_decorate": 0, "compiled_call": 0}
+    calls = {"jit_decorate": 0, "compiled_call": 0, "compiled_compile": 0}
 
     tilelang: Any = ModuleType("tilelang")
 
+    class _JitImpl:
+        def __init__(self, func: Any) -> None:
+            self.func = func
+
+        def __call__(self, *args: Any, **kw: Any) -> Any:
+            calls["compiled_call"] += 1
+            return self.func.__name__
+
+        def compile(self, *args: Any, **kw: Any) -> Any:
+            calls["compiled_compile"] += 1
+            return self.func.__name__
+
     def jit(**kwargs: Any) -> Any:
         def decorate(func: Any) -> Any:
             calls["jit_decorate"] += 1
-
-            def compiled(*args: Any, **kw: Any) -> Any:
-                calls["compiled_call"] += 1
-                return func.__name__
-
-            return compiled
+            return _JitImpl(func)
 
         return decorate
 
@@ -78,6 +85,21 @@ def test_tilelang_jit_decorator_is_lazy_only_on_rocm(
     assert calls["compiled_call"] == 1
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm")
+def test_tilelang_jit_proxies_compile_only_warmup(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _install_tilelang_stub(monkeypatch)
+    module_name = "vllm.model_executor.kernels.mhc.tilelang_kernels"
+    monkeypatch.delitem(sys.modules, module_name, raising=False)
+    module = importlib.import_module(module_name)
+
+    assert module.mhc_post_tilelang.compile() == "mhc_post_tilelang"
+    assert calls["jit_decorate"] == 1
+    assert calls["compiled_compile"] == 1
+    assert calls["compiled_call"] == 0
+
+
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm")
 def test_deepseek_v4_import_and_jit_monitor_do_not_hijack_hip_symbols() -> None:
     if importlib.util.find_spec("tilelang") is None:
diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py
index 182672597a9b..27bd74ce41dc 100644
--- a/tests/kernels/test_minimax_m3_amd_ops.py
+++ b/tests/kernels/test_minimax_m3_amd_ops.py
@@ -10,6 +10,7 @@
   * Fused MXFP8 activation quant (Triton)        -> _mxfp8_e4m3_quantize_torch
   * Native MXFP8 linear (dot_scaled)             -> dequant-to-bf16 @ matmul
   * Native MXFP8 MoE (dot_scaled grouped GEMM)   -> dequant-to-bf16 MoE math
+  * Strided K/V sparse-PA cache insert           -> contiguous-copy insert
 
 The native MXFP8 GEMMs also guard the ``dot_scaled`` rhs-scale orientation: the
 scale is loaded ``[N, K//32]`` and passed WITHOUT transpose; a stray ``.T``
@@ -39,9 +40,11 @@
     _mxfp8_e4m3_quantize_triton,
     dequant_mxfp8_to_bf16,
 )
+from vllm.models.minimax_m3.amd.model import _kv_insert_operand  # noqa: E402
 from vllm.models.minimax_m3.amd.ops import (  # noqa: E402
     gemma_fused_add_rmsnorm,
     gemma_rmsnorm,
+    swiglu_oai_quantize_mxfp8,
     swiglu_oai_split,
 )
 from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps  # noqa: E402
@@ -167,6 +170,25 @@ def test_swiglu_oai_split(m, inter, limit, dtype):
     assert _relerr(got, ref) < 5e-3
 
 
+@torch.inference_mode()
+def test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale():
+    # Keep a small value in two MX blocks next to their maxima. The scale must
+    # use E4M3's finite range (448), otherwise that value underflows to zero.
+    # Keep the third block empty to cover the reference's finite tiny clamp.
+    m, inter = 3, 96
+    gate = torch.full((m, inter), 2.0, device=DEVICE, dtype=torch.float16)
+    up = torch.full((m, inter), 2**-10, device=DEVICE, dtype=torch.float16)
+    up[:, :64:32] = 1.0
+    up[:, 64:] = 0.0
+    gate_up = torch.cat((gate, up), dim=-1)
+
+    got_q, got_s = swiglu_oai_quantize_mxfp8(gate_up, alpha=0.0, beta=0.0, limit=None)
+    ref_q, ref_s = _mxfp8_e4m3_quantize_torch(up, is_sf_swizzled_layout=False)
+
+    assert torch.equal(got_s, ref_s)
+    assert torch.equal(got_q, ref_q)
+
+
 # --------------------------------------------------------------------------- #
 # Fused MXFP8 activation quant (Triton vs torch reference)
 # --------------------------------------------------------------------------- #
@@ -179,8 +201,8 @@ def test_mxfp8_quant_triton_matches_torch(shape, dtype):
     xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False)
     xq_k, s_k = _mxfp8_e4m3_quantize_triton(x)
     assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32)
-    # E8M0 block exponents share the floor(log2(amax))+127 algorithm; allow at
-    # most a 1-step difference at exact powers of two.
+    # Both paths use the E4M3-aware scale calculation; allow a 1-step
+    # difference at exact powers of two due to floating-point rounding.
     assert (s_k.int() - s_t.int()).abs().max().item() <= 1
     # Dequantized values agree to fp8 granularity.
     deq_t = dequant_mxfp8_to_bf16(xq_t, s_t)
@@ -424,6 +446,69 @@ def test_mxfp8_linear_emulation_bf16_at_load(
     assert _relerr(out.float(), out_ref.float()) < 2e-2
 
 
+# --------------------------------------------------------------------------- #
+# Native MXFP8 linear: the K % 128 != 0 weight is dequantized once at load
+# --------------------------------------------------------------------------- #
+@requires_gfx950
+# (5120, 576) and (5120, 288): DSv4.1 shared-expert down_proj at TP4 and TP8.
+@pytest.mark.parametrize("shape", [(5120, 576), (5120, 288), (256, 2048)])
+@torch.inference_mode()
+def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape):
+    """``dot_scaled`` tiles K by 128, so an unaligned weight must not reach it.
+
+    ``process_weights_after_loading`` converts those to BF16 once; a weight that
+    never went through it has to fall back per step rather than take the
+    invalid path, which is what keying the dispatch on the weight dtype alone
+    would do.
+    """
+    from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
+    from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+        Mxfp8LinearLayerConfig,
+    )
+
+    N, K = shape
+    aligned = K % 128 == 0
+    torch.manual_seed(0)
+    w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
+    x = torch.randn(7, K, device=DEVICE, dtype=torch.bfloat16) * 0.5
+    out_ref = torch.nn.functional.linear(x, dequant_mxfp8_to_bf16(w_fp8, w_scale))
+    # dot_scaled re-quantizes the activation, the dequant paths do not.
+    tol = 5e-2 if aligned else 2e-3
+
+    def make_layer():
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
+        layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
+        return layer
+
+    kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig())
+    layer = make_layer()
+    kernel.process_weights_after_loading(layer)
+    assert (layer.weight.element_size() >= 2) is not aligned
+    assert _relerr(kernel.apply_weights(layer, x).float(), out_ref.float()) < tol
+
+    # An unaligned weight that is still FP8 must not reach dot_scaled.
+    unprocessed = make_layer()
+    unprocessed.weight_scale = torch.nn.Parameter(
+        w_scale[:N, : K // 32].contiguous(), requires_grad=False
+    )
+    calls = []
+    real = rocm_native._mxfp8_dot_scaled_linear
+
+    def _spy(*a, **k):
+        calls.append(1)
+        return real(*a, **k)
+
+    rocm_native._mxfp8_dot_scaled_linear = _spy
+    try:
+        out = kernel.apply_weights(unprocessed, x)
+    finally:
+        rocm_native._mxfp8_dot_scaled_linear = real
+    assert bool(calls) is aligned
+    assert _relerr(out.float(), out_ref.float()) < tol
+
+
 # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ────────────────
 # The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER
 # experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and
@@ -525,3 +610,77 @@ def resolve(consumes_mask, *, has_moe_kernel=True):
     assert torch.equal(resolve(True), mask)  # AITER kernel -> 0/1 mask
     assert torch.equal(resolve(False), canonical)  # non-AITER -> canonical map
     assert torch.equal(resolve(False, has_moe_kernel=False), canonical)  # non-modular
+
+
+# --------------------------------------------------------------------------- #
+# Strided K/V operand for the AITER sparse-PA cache insert
+# --------------------------------------------------------------------------- #
+def _fused_qkv_kv_slices(num_tokens, num_kv_heads, head_dim, num_q_heads=16):
+    """K/V as the sparse-PA insert sees them: column slices of a fused qkv."""
+    kv = num_kv_heads * head_dim
+    row = num_q_heads * head_dim + 2 * kv + 2 * head_dim
+    qkv = torch.randn((num_tokens, row), dtype=torch.bfloat16, device=DEVICE)
+    k_start = num_q_heads * head_dim
+    k = qkv[:, k_start : k_start + kv].view(num_tokens, num_kv_heads, head_dim)
+    v = qkv[:, k_start + kv : k_start + 2 * kv].view(num_tokens, num_kv_heads, head_dim)
+    return k, v
+
+
+def _asm_kv_cache(pages, num_kv_heads, head_dim):
+    """Page-16 K/V cache views matching ``_ensure_aiter_sparse_pa_kv_cache``."""
+    x = 16 // torch.tensor([], dtype=torch.bfloat16).element_size()
+    key = torch.zeros(
+        (pages, num_kv_heads, head_dim // x, 16, x),
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    value = torch.zeros(
+        (pages, num_kv_heads, 16 // x, head_dim, x),
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    return key, value
+
+
+@pytest.mark.parametrize("num_tokens", [1, 8, 129, 1024])
+@pytest.mark.parametrize("num_kv_heads", [1, 2])
+def test_kv_insert_operand_matches_contiguous(num_tokens, num_kv_heads):
+    """A strided K/V slice must insert bit-identically to a contiguous copy."""
+    reshape_and_cache = pytest.importorskip("aiter").reshape_and_cache
+    head_dim = 128
+    pages = num_tokens // 16 + 4
+    slots = torch.randperm(pages * 16, device=DEVICE)[:num_tokens].to(torch.int32)
+
+    torch.manual_seed(num_tokens)
+    k, v = _fused_qkv_kv_slices(num_tokens, num_kv_heads, head_dim)
+    assert num_tokens == 1 or not k.is_contiguous()
+
+    caches = []
+    for operand in (lambda t: t.contiguous(), _kv_insert_operand):
+        key_cache, value_cache = _asm_kv_cache(pages, num_kv_heads, head_dim)
+        reshape_and_cache(
+            operand(k),
+            operand(v),
+            key_cache,
+            value_cache,
+            slots,
+            kv_cache_dtype="auto",
+            asm_layout=True,
+        )
+        caches.append((key_cache, value_cache))
+
+    (ref_k, ref_v), (got_k, got_v) = caches
+    assert torch.equal(ref_k, got_k)
+    assert torch.equal(ref_v, got_v)
+
+
+def test_kv_insert_operand_copies_unsupported_layouts():
+    """Only a row-strided slice passes through; anything else is copied."""
+    k, _ = _fused_qkv_kv_slices(32, 2, 128)
+    assert _kv_insert_operand(k) is k
+
+    inner_strided = k[:, :, ::2]
+    assert _kv_insert_operand(inner_strided).is_contiguous()
+
+    flat = k.reshape(32, -1)
+    assert _kv_insert_operand(flat).is_contiguous()
diff --git a/tests/kernels/test_onednn.py b/tests/kernels/test_onednn.py
index 4e681ca6ac83..4e41791a5408 100644
--- a/tests/kernels/test_onednn.py
+++ b/tests/kernels/test_onednn.py
@@ -15,6 +15,7 @@
     (256, 128),
     (4096, 4096),
     (16384, 4096),
+    (64128, 4096),
     (1023, 491),
     (1001, 15),
 ]
diff --git a/tests/kernels/test_relu2_fp8_quant.py b/tests/kernels/test_relu2_fp8_quant.py
new file mode 100644
index 000000000000..cb605914c7f5
--- /dev/null
+++ b/tests/kernels/test_relu2_fp8_quant.py
@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
+from vllm.config.vllm import OptimizationLevel
+from vllm.model_executor.layers.activation import ReLUSquaredActivation
+from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    kFp8StaticTensorSym,
+)
+from vllm.platforms import current_platform
+
+requires_sm90 = pytest.mark.skipif(
+    not (current_platform.is_cuda() and current_platform.has_device_capability(90)),
+    reason="This test requires SM90 or newer",
+)
+
+
+class _StaticFp8Linear(torch.nn.Module):
+    def __init__(self, input_scale: torch.Tensor) -> None:
+        super().__init__()
+        self._input_quant_key = kFp8StaticTensorSym
+        self.input_scale = input_scale
+
+
+def _assert_fp8_bitwise_equal(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    assert actual.dtype == expected.dtype == current_platform.fp8_dtype()
+    torch.testing.assert_close(
+        actual.view(torch.uint8), expected.view(torch.uint8), rtol=0, atol=0
+    )
+
+
+def _assert_quantized_activation(
+    result: QuantizedActivation, x: torch.Tensor, scale: torch.Tensor
+) -> None:
+    assert result.scale is scale
+    assert result.orig_dtype == x.dtype
+    assert result.orig_shape == x.shape
+    assert result.quant_key == kFp8StaticTensorSym
+
+
+@pytest.fixture(scope="module")
+def o2_relu2_fp8_ops():
+    config = VllmConfig(
+        optimization_level=OptimizationLevel.O2,
+        compilation_config=CompilationConfig(custom_ops=["none"]),
+    )
+    with set_current_vllm_config(config):
+        yield (
+            ReLUSquaredActivation(),
+            QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR),
+        )
+
+
+@requires_sm90
+@pytest.mark.parametrize("shape", [(1, 1), (1, 16), (17, 5120), (128, 10240)])
+@torch.inference_mode()
+def test_relu2_static_fp8_quant(o2_relu2_fp8_ops, shape) -> None:
+    torch.manual_seed(0)
+    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    scale = torch.tensor(0.1495535671710968, device="cuda", dtype=torch.float32)
+
+    relu2, quant_fp8 = o2_relu2_fp8_ops
+    expected, _ = quant_fp8(relu2(x), scale)
+    result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
+
+    assert isinstance(result, QuantizedActivation)
+    _assert_quantized_activation(result, x, scale)
+    _assert_fp8_bitwise_equal(result.data, expected)
+
+
+@requires_sm90
+@pytest.mark.parametrize(
+    "enabled_ops",
+    [(), ("relu2",), ("quant_fp8",), ("relu2", "quant_fp8")],
+)
+@torch.inference_mode()
+def test_relu2_static_fp8_quant_ignores_custom_op_dispatch(enabled_ops) -> None:
+    custom_ops = ["none", *(f"+{op}" for op in enabled_ops)]
+    config = VllmConfig(
+        optimization_level=OptimizationLevel.O2,
+        compilation_config=CompilationConfig(custom_ops=custom_ops),
+    )
+    with set_current_vllm_config(config):
+        torch.manual_seed(0)
+        relu2 = ReLUSquaredActivation()
+        quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
+        x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16)
+        scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32)
+
+        expected, _ = quant_fp8(relu2(x), scale)
+        result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
+
+    assert isinstance(result, QuantizedActivation)
+    _assert_quantized_activation(result, x, scale)
+    _assert_fp8_bitwise_equal(result.data, expected)
+
+
+@requires_sm90
+@torch.inference_mode()
+def test_relu2_static_fp8_quant_without_compilation() -> None:
+    config = VllmConfig(optimization_level=OptimizationLevel.O0)
+    with set_current_vllm_config(config):
+        torch.manual_seed(0)
+        relu2 = ReLUSquaredActivation()
+        quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
+        x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16)
+        scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32)
+
+        expected, _ = quant_fp8(relu2(x), scale)
+        result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
+
+    assert isinstance(result, QuantizedActivation)
+    _assert_quantized_activation(result, x, scale)
+    _assert_fp8_bitwise_equal(result.data, expected)
+
+
+@requires_sm90
+@pytest.mark.parametrize(
+    "scale_value",
+    [
+        0.00435965,
+        0.007149832788854837,
+        0.03125,
+        0.0714285746216774,
+        0.1495535671710968,
+        0.1517857164144516,
+        6.03571415,
+    ],
+)
+@torch.inference_mode()
+def test_relu2_static_fp8_quant_all_bf16_values(o2_relu2_fp8_ops, scale_value) -> None:
+    bits = torch.arange(1 << 16, dtype=torch.int32).to(torch.uint16)
+    x = bits.view(torch.bfloat16).reshape(1, -1).to("cuda")
+    scale = torch.tensor(scale_value, device="cuda", dtype=torch.float32)
+
+    relu2, quant_fp8 = o2_relu2_fp8_ops
+    expected, _ = quant_fp8(relu2(x), scale)
+    result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
+
+    assert isinstance(result, QuantizedActivation)
+    _assert_quantized_activation(result, x, scale)
+    _assert_fp8_bitwise_equal(result.data, expected)
+
+
+@requires_sm90
+@torch.inference_mode()
+def test_relu2_static_fp8_quant_empty(o2_relu2_fp8_ops) -> None:
+    x = torch.empty(0, device="cuda", dtype=torch.bfloat16)
+    scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32)
+
+    relu2, _ = o2_relu2_fp8_ops
+    result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
+
+    assert isinstance(result, QuantizedActivation)
+    _assert_quantized_activation(result, x, scale)
+    assert result.data.shape == x.shape
+    assert result.data.dtype == current_platform.fp8_dtype()
+
+
+@requires_sm90
+@pytest.mark.parametrize("unsupported", ["dtype", "layout", "scale"])
+@torch.inference_mode()
+def test_relu2_static_fp8_quant_falls_back(o2_relu2_fp8_ops, unsupported: str) -> None:
+    x = torch.randn((17, 32), device="cuda", dtype=torch.bfloat16)
+    scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32)
+    if unsupported == "dtype":
+        x = x.to(torch.float16)
+    elif unsupported == "layout":
+        x = x.T
+    else:
+        scale = scale.to(torch.float16)
+
+    relu2, _ = o2_relu2_fp8_ops
+    expected = relu2(x)
+    result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
+
+    assert isinstance(result, torch.Tensor)
+    assert not isinstance(result, QuantizedActivation)
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+@requires_sm90
+@torch.inference_mode()
+def test_relu2_static_fp8_quant_torch_compile_dynamic(
+    o2_relu2_fp8_ops,
+) -> None:
+    torch.manual_seed(1)
+    scale = torch.tensor(0.00435965, device="cuda", dtype=torch.float32)
+    relu2, quant_fp8 = o2_relu2_fp8_ops
+    linear = _StaticFp8Linear(scale)
+
+    def fused(x: torch.Tensor) -> torch.Tensor:
+        result = maybe_fused_act_quant(relu2, x, linear)
+        assert isinstance(result, QuantizedActivation)
+        return result.data
+
+    compile_count = 0
+
+    def counting_inductor_backend(gm, example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return torch._inductor.compile(gm, example_inputs)
+
+    torch._dynamo.reset()
+    try:
+        x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16)
+        torch._dynamo.mark_dynamic(x, 0)
+        compiled_fused = torch.compile(
+            fused, backend=counting_inductor_backend, fullgraph=True
+        )
+
+        for candidate in (
+            x,
+            torch.randn((31, 5120), device="cuda", dtype=torch.bfloat16),
+            torch.randn((63, 5120), device="cuda", dtype=torch.bfloat16),
+        ):
+            expected, _ = quant_fp8(relu2(candidate), scale)
+            _assert_fp8_bitwise_equal(compiled_fused(candidate), expected)
+
+        assert compile_count == 1
+    finally:
+        torch._dynamo.reset()
diff --git a/tests/kernels/test_shuffle_rows.py b/tests/kernels/test_shuffle_rows.py
index 08d5883a1e70..4f32949eb765 100644
--- a/tests/kernels/test_shuffle_rows.py
+++ b/tests/kernels/test_shuffle_rows.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the shuffle_rows function
+"""Tests for the shuffle_rows function.
 
 Run `pytest tests/kernels/test_shuffle_rows.py`.
 """
diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py
index a1e27d2e50a3..b6c469181f09 100644
--- a/tests/kernels/test_top_k_per_row.py
+++ b/tests/kernels/test_top_k_per_row.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+
 import numpy as np
 import pytest
 import torch
@@ -21,6 +22,24 @@ def _has_device_capability(major: int) -> bool:
     return current_platform.is_cuda() and current_platform.has_device_capability(major)
 
 
+def _on_gfx950() -> bool:
+    if not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import on_gfx950
+
+    return on_gfx950()
+
+
+# DeepSelect is compiled for sm_100a/sm_103a only.
+requires_sm100 = pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="DeepSelect requires SM100a/SM103a",
+)
+requires_gfx950 = pytest.mark.skipif(
+    not _on_gfx950(), reason="This test exercises the gfx950 launch configuration"
+)
+
+
 COOPERATIVE_TOPK_BACKEND = pytest.param(
     "cooperative_topk",
     marks=pytest.mark.skipif(
@@ -138,8 +157,7 @@ def compare_top_k_results(
     top_k: int,
     tolerance: float = 1e-5,
 ) -> bool:
-    """
-    Compare results from CUDA top_k_per_row with torch.topk.
+    """Compare results from CUDA top_k_per_row with torch.topk.
     Both results should be sorted and contain the same top-k elements.
     """
     num_rows = cuda_indices.shape[0]
@@ -194,8 +212,7 @@ def validate_topk_against_reference(
     top_k: int,
     kernel_name: str,
 ) -> None:
-    """
-    Validate CUDA top-k results against PyTorch reference implementation.
+    """Validate CUDA top-k results against PyTorch reference implementation.
 
     Args:
         logits: Input logits tensor
@@ -204,6 +221,7 @@ def validate_topk_against_reference(
         row_ends: Row end positions
         top_k: Number of top elements to select
         kernel_name: Name of the kernel being tested (for error messages)
+
     """
     num_rows = cuda_indices.shape[0]
     torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
@@ -229,9 +247,7 @@ def test_top_k_per_row(
     top_k: int,
     clean_logits: bool,
 ) -> None:
-    """
-    Test top_k_per_row.
-    """
+    """Test top_k_per_row."""
     set_random_seed(0)
     torch.set_default_device("cuda:0")
 
@@ -279,9 +295,7 @@ def _run_top_k_per_row_decode_test(
     clean_logits: bool,
     data_generation: str,
 ) -> None:
-    """
-    Helper function to run top_k_per_row_decode test with given parameters.
-    """
+    """Helper function to run top_k_per_row_decode test with given parameters."""
     torch.set_default_device("cuda:0")
 
     # Create test data
@@ -346,9 +360,7 @@ def test_top_k_per_row_decode(
     clean_logits: bool,
     data_generation: str,
 ) -> None:
-    """
-    Test top_k_per_row with seq_lens tensor.
-    """
+    """Test top_k_per_row with seq_lens tensor."""
     set_random_seed(0)
     vocab_size = 20000
     _run_top_k_per_row_decode_test(
@@ -360,9 +372,7 @@ def test_top_k_per_row_decode(
 @pytest.mark.parametrize("clean_logits", [True, False])
 @torch.inference_mode()
 def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
-    """
-    Test top_k_per_row_decode with large vocabulary size.
-    """
+    """Test top_k_per_row_decode with large vocabulary size."""
     set_random_seed(0)
     top_k = 2048
     batch_size = 2
@@ -479,6 +489,150 @@ def test_top_k_per_row_decode_gfx950_long_c4a_1d_seq_lens() -> None:
     )
 
 
+def _assert_exact_topk(
+    logits: torch.Tensor, indices: torch.Tensor, row_ends: torch.Tensor
+) -> None:
+    ends = row_ends.reshape(-1, 1).clamp_min(0)
+    valid = (indices >= 0) & (indices < ends)
+    assert torch.equal(valid.sum(1), ends.clamp_max(indices.shape[1]).flatten())
+    assert torch.all(valid | (indices == -1))
+    ordered = indices.sort(1).values
+    assert not ((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any()
+
+    width = int(ends.max())
+    visible = logits[:, :width].clone()
+    columns = torch.arange(width, device=logits.device)
+    visible.masked_fill_(columns[None, :] >= ends, float("-inf"))
+    expected = visible.topk(min(indices.shape[1], width), dim=1).values
+    selected = logits.gather(1, indices.clamp_min(0).long())
+    selected.masked_fill_(~valid, float("-inf"))
+    selected = selected.sort(1, descending=True).values
+    assert torch.equal(selected[:, : expected.shape[1]], expected)
+
+
+@requires_gfx950
+@torch.inference_mode()
+def test_top_k_per_row_decode_gfx950_k512_1d_seq_lens() -> None:
+    """Honor causal offsets and exact compressed lengths, including short rows."""
+    next_n, width, top_k = 6, 16_384, 512
+    lengths = torch.tensor(
+        [0, 1, 511, 512, 513, 1023, 1024, 16_000],
+        dtype=torch.int32,
+        device="cuda",
+    )
+    offsets = torch.arange(next_n, dtype=torch.int32, device="cuda")
+    row_ends = (lengths[:, None] - next_n + offsets + 1).clamp_min(0).flatten()
+    rows = row_ends.numel()
+    storage = torch.full((rows, width + 8), 1e20, dtype=torch.float32, device="cuda")
+    logits = storage[:, 1 : width + 1]
+    logits[:, :16_000] = torch.arange(16_000, dtype=torch.float32, device="cuda")
+    indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda")
+    _run_topk_backend(
+        "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n
+    )
+    _assert_exact_topk(logits, indices, row_ends)
+
+
+@pytest.mark.parametrize(
+    ("rows", "row_length"),
+    [
+        pytest.param(128, 65_536, id="single-before-64k"),
+        pytest.param(128, 65_537, id="four-splits-after-64k"),
+        pytest.param(129, 262_144, id="single-before-three-split-length"),
+        pytest.param(129, 262_145, id="three-splits-after-256k"),
+        pytest.param(160, 262_145, id="three-splits-upper-row-bound"),
+        pytest.param(161, 262_144, id="single-before-two-split-length"),
+        pytest.param(161, 262_145, id="two-splits-after-256k"),
+        pytest.param(256, 262_145, id="two-splits-upper-row-bound"),
+        pytest.param(257, 262_145, id="large-row-single-block"),
+        pytest.param(384, 262_145, id="optimized-row-limit"),
+        pytest.param(385, 262_145, id="native-fallback"),
+    ],
+)
+@requires_gfx950
+@torch.inference_mode()
+def test_top_k_per_row_decode_gfx950_k512_split_policy_boundaries(
+    rows: int, row_length: int
+) -> None:
+    width, top_k = 262_145, 512
+    row = torch.arange(width, dtype=torch.float32, device="cuda")
+    logits = row.expand(rows, -1)
+    lengths = torch.full((rows,), row_length, dtype=torch.int32, device="cuda")
+    indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda")
+
+    _run_topk_backend("top_k_per_row_decode", logits, lengths, indices, top_k, width)
+
+    expected = torch.arange(
+        row_length - top_k, row_length, dtype=torch.int32, device="cuda"
+    ).expand_as(indices)
+    assert torch.equal(indices.sort(1).values, expected)
+
+
+@requires_gfx950
+@torch.inference_mode()
+def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None:
+    """Preserve visible top-k scores when masks and lengths change on replay."""
+    rows, next_n, width, top_k = 24, 6, 262_145, 512
+    logits = torch.full((rows, width), 1e20, dtype=torch.float32, device="cuda")
+    lengths = torch.full(
+        (rows // next_n, next_n), 70_000, dtype=torch.int32, device="cuda"
+    )
+    indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda")
+    logits[:, :70_000] = float("-inf")
+    bin_sizes = [128, 129, 1024, 1025]
+    for row in range(rows):
+        pattern = row % 9
+        if pattern < 4:
+            count = [0, 1, 511, 513][pattern]
+            logits[row, :count] = torch.arange(
+                count, dtype=torch.float32, device="cuda"
+            )
+        elif pattern == 4:
+            logits[row, :511] = 0
+            logits[row, 511] = torch.finfo(torch.float32).min
+        else:
+            count = bin_sizes[pattern - 5]
+            higher = max(0, top_k - count // 2)
+            logits[row, :higher] = 2
+            logits[row, higher : higher + count] = 1
+            logits[row, higher] = 1 + 2**-23
+
+    def run() -> None:
+        _run_topk_backend(
+            "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n
+        )
+
+    run()
+    _assert_exact_topk(logits, indices, lengths)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run()
+    indices.fill_(-777)
+    graph.replay()
+    _assert_exact_topk(logits, indices, lengths)
+
+    bounds = [
+        0,
+        511,
+        512,
+        513,
+        16_384,
+        16_385,
+        131_072,
+        131_073,
+        262_144,
+        262_145,
+    ]
+    row_bounds = torch.tensor(bounds, dtype=torch.int32, device="cuda")
+    repeats = (rows + len(bounds) - 1) // len(bounds)
+    lengths.flatten().copy_(row_bounds.repeat(repeats)[:rows])
+    logits[:] = torch.arange(width, dtype=torch.float32, device="cuda")
+    indices.fill_(-777)
+    graph.replay()
+    _assert_exact_topk(logits, indices, lengths)
+
+
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm")
 @torch.inference_mode()
 def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None:
@@ -612,6 +766,31 @@ def decode_kernel(*args, **kwargs) -> None:
                 compress_ratio=4,
                 num_rows=257,
                 max_valid_seq_len=250_000,
+                num_columns=524_288,
+                on_gfx950=True,
+            ),
+            decode_kernel,
+        ),
+        (
+            dict(
+                is_prefill=False,
+                compress_ratio=2,
+                num_rows=385,
+                max_valid_seq_len=250_000,
+                num_columns=524_288,
+                topk_tokens=512,
+                on_gfx950=True,
+            ),
+            decode_kernel,
+        ),
+        (
+            dict(
+                is_prefill=False,
+                compress_ratio=2,
+                num_rows=384,
+                max_valid_seq_len=50_000,
+                num_columns=1_048_577,
+                topk_tokens=512,
                 on_gfx950=True,
             ),
             decode_kernel,
@@ -630,6 +809,15 @@ def fail_if_imported() -> None:
     )
     ineligible_cases = [
         dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True),
+        dict(
+            is_prefill=False,
+            compress_ratio=2,
+            num_rows=384,
+            max_valid_seq_len=1_048_576,
+            num_columns=1_048_576,
+            topk_tokens=512,
+            on_gfx950=True,
+        ),
         dict(
             is_prefill=False,
             compress_ratio=4,
@@ -689,8 +877,7 @@ def test_deepseek_workspace_topk(
     next_n: int,
     backend: str,
 ) -> None:
-    """
-    Test workspace top-k backends with varying sequence lengths and speculative
+    """Test workspace top-k backends with varying sequence lengths and speculative
     decoding.
     Supports speculative decoding with next_n > 1.
     """
@@ -735,6 +922,106 @@ def test_deepseek_workspace_topk(
     )
 
 
+@pytest.mark.skipif(not _has_device_capability(80), reason="Requires SM80+")
+@pytest.mark.parametrize("rows", [65, 128])
+@pytest.mark.parametrize("top_k", [512, 1024, 2048])
+@pytest.mark.parametrize(
+    "distribution", ["random", "10LSBits", "ascending", "constant", "sampled_peaks"]
+)
+@torch.inference_mode()
+def test_persistent_topk_sampled_graph(
+    rows: int, top_k: int, distribution: str
+) -> None:
+    """Sampling and both fallbacks preserve exact values across graph replays."""
+    if torch.cuda.get_device_properties(0).shared_memory_per_block_optin < 144 * 1024:
+        pytest.skip("Sampled top-k requires at least 144 KiB of shared memory")
+    set_random_seed(42)
+    min_sampled_length = 98304 if top_k == 512 else 65536
+    width = min_sampled_length + 3
+    lengths = torch.full((rows,), width, dtype=torch.int32, device="cuda")
+    logits = create_random_logits(
+        torch.zeros_like(lengths),
+        lengths,
+        torch.float32,
+        42,
+        False,
+        "10LSBits" if distribution == "10LSBits" else "random",
+    )
+    if distribution == "ascending":
+        logits.copy_(torch.arange(width, device="cuda", dtype=torch.float32))
+    elif distribution == "constant":
+        logits.fill_(1.0)
+    elif distribution == "sampled_peaks":
+        # A biased sample leaves fewer than k survivors and must fall back.
+        logits.zero_()
+        sample = torch.arange(4096, device="cuda")
+        positions = (sample // 32) * (width // 128) + sample % 32
+        logits[:, positions] = 1 + sample.float() / 4096
+    original = logits.clone()
+    indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda")
+    workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda")
+
+    def run() -> None:
+        torch.ops._C.persistent_topk(
+            logits,
+            lengths.view(-1, 4 if rows == 128 else 1),
+            indices,
+            workspace,
+            top_k,
+            width,
+        )
+
+    run()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run()
+    positions = torch.arange(width, device="cuda")
+    slots = torch.arange(top_k, device="cuda")
+    for step in range(3):
+        bounds = torch.tensor(
+            [
+                -1,
+                0,
+                1,
+                top_k - 1,
+                top_k,
+                32769,
+                min_sampled_length - 1,
+                min_sampled_length,
+                width,
+                width + 17,
+            ],
+            dtype=torch.int32,
+            device="cuda",
+        )
+        lengths.copy_(
+            bounds[(torch.arange(rows, device="cuda") + step) % bounds.numel()]
+        )
+        logits.copy_(original if step != 1 else original.flip(1))
+        logits.masked_fill_(positions[None] >= lengths[:, None], float("nan"))
+        indices.fill_(-2)
+        graph.replay()
+        valid = slots[None] < lengths.clamp(max=top_k)[:, None]
+        assert torch.all(indices[~valid] == -1)
+        assert torch.all(((indices >= 0) & (indices < lengths[:, None])) == valid)
+        ordered_indices = indices.sort(dim=1).values
+        assert torch.all(
+            (ordered_indices[:, 1:] != ordered_indices[:, :-1])
+            | (ordered_indices[:, 1:] == -1)
+        )
+        selected = logits.gather(1, indices.clamp_min(0).long())
+        selected.masked_fill_(~valid, -float("inf"))
+        expected = logits.masked_fill(
+            positions[None] >= lengths[:, None], -float("inf")
+        )
+        torch.testing.assert_close(
+            selected.sort(dim=1, descending=True).values,
+            expected.topk(top_k, dim=1).values,
+            atol=0,
+            rtol=0,
+        )
+
+
 def run_large_context_topk_test(
     batch_size: int,
     seq_lens: list[int],
@@ -743,8 +1030,7 @@ def run_large_context_topk_test(
     seed: int = 42,
     backend: str = "cooperative_topk",
 ) -> None:
-    """
-    Helper to run a top-k backend test with given parameters.
+    """Helper to run a top-k backend test with given parameters.
 
     Args:
         batch_size: Number of rows/sequences
@@ -753,6 +1039,7 @@ def run_large_context_topk_test(
         data_type: Type of test data to generate
         seed: Random seed for reproducibility
         backend: Top-k backend to test
+
     """
     torch.set_default_device("cuda:0")
     set_random_seed(seed)
@@ -977,8 +1264,7 @@ def test_cooperative_topk_cs2(top_k: int) -> None:
 @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
 def test_workspace_topk_correctness(test_config: dict, backend: str) -> None:
-    """
-    Comprehensive correctness tests covering:
+    """Comprehensive correctness tests covering:
     - Sequence length edge cases (trivial, boundary, varied)
     - Very small sequences (< 100 elements)
     - Mixed sequence lengths in same batch
@@ -1042,8 +1328,7 @@ def test_workspace_topk_correctness(test_config: dict, backend: str) -> None:
 @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
 def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None:
-    """
-    Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2):
+    """Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2):
     - Batch size scalability (1, 4, 32, 256)
     - Single-CTA vs Multi-CTA execution
     - Extreme configurations (large batch, max context length)
@@ -1060,8 +1345,7 @@ def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None
 @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
 def test_workspace_topk_stress(backend: str) -> None:
-    """
-    Stress test with random configurations to catch edge cases.
+    """Stress test with random configurations to catch edge cases.
     Capped at 163840 (DeepSeek V3.2 max context) for realistic testing.
     """
     torch.set_default_device("cuda:0")
@@ -1208,8 +1492,7 @@ def test_cooperative_topk_512_tie_workspace_is_per_row() -> None:
 @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
 def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None:
-    """
-    Tests specific to workspace top-k backends:
+    """Tests specific to workspace top-k backends:
     - Mixed medium/large rows in the same batch (dynamic per-row dispatch)
     - Boundary around LARGE_THRESHOLD (32K)
     - Trivial + medium + large rows in a single batch
@@ -1270,8 +1553,7 @@ def test_persistent_topk_reused_group_after_short_row() -> None:
 @pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
 @torch.inference_mode()
 def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None:
-    """
-    Test workspace top-k backends with padded logits (large stride, small seq_len)
+    """Test workspace top-k backends with padded logits (large stride, small seq_len)
     to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits
     returns [B, max_model_len] with max_model_len=163840.
     """
@@ -1315,3 +1597,333 @@ def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None:
                 f"Row {i}: {backend} with padded stride doesn't match. "
                 f"seq_len={sl}, stride={padded_stride}"
             )
+
+
+@requires_sm100
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("batch_size", [1, 8, 64])
+@pytest.mark.parametrize("vocab_size", [4096, 131072])
+@pytest.mark.parametrize("top_k", [512, 2048])
+@torch.inference_mode()
+def test_deep_select_topk(
+    dtype: torch.dtype,
+    batch_size: int,
+    vocab_size: int,
+    top_k: int,
+) -> None:
+    """DeepSelect wrapper vs torch.topk with variable per-row ends."""
+    from vllm.model_executor.layers import indexer_topk
+
+    set_random_seed(0)
+    torch.set_default_device("cuda:0")
+
+    row_ends = torch.randint(
+        1, vocab_size + 1, (batch_size,), dtype=torch.int32, device="cuda"
+    )
+    # Force one short row to exercise the -1 out-of-bounds fill.
+    row_ends[0] = min(top_k - 1, vocab_size)
+    logits = torch.randn(batch_size, vocab_size, dtype=dtype, device="cuda")
+    # Values beyond each row's end must never be read by the kernel: fill
+    # them with NaN (production tails are uninitialized; DeepSelect traps on
+    # NaN with abort_when_nan_found=True, so a wrong read fails loudly).
+    col_idx = torch.arange(vocab_size, device="cuda")
+    logits[col_idx[None, :] >= row_ends[:, None]] = float("nan")
+
+    indices = indexer_topk.deep_select_topk(logits, top_k, end=row_ends)
+    torch.accelerator.synchronize()
+
+    assert indices.shape == (batch_size, top_k)
+    row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
+    torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
+    for i in range(batch_size):
+        row_end = int(row_ends[i])
+        k_i = min(top_k, row_end)
+        idx = logits[i, :row_end].topk(k_i, dim=-1)[1]
+        torch_indices[i, :k_i] = idx
+        assert torch.all(indices[i, k_i:] == -1), (
+            f"Row {i}: expected -1 fill after {k_i} valid indices"
+        )
+        # All selected indices must lie inside [0, row_end).
+        assert torch.all(indices[i, :k_i] < row_end)
+
+    assert compare_top_k_results(
+        logits, indices, torch_indices, row_starts, row_ends, top_k
+    ), "DeepSelect topk results don't match torch.topk"
+
+
+@requires_sm100
+@torch.inference_mode()
+def test_deep_select_topk_preallocated_output() -> None:
+    """DeepSelect writes into a preallocated aligned buffer (no `end`)."""
+    from vllm.model_executor.layers import indexer_topk
+
+    set_random_seed(0)
+    torch.set_default_device("cuda:0")
+
+    batch_size = 8
+    vocab_size = 131072
+    top_k = 2048
+    logits = torch.randn(batch_size, vocab_size, dtype=torch.float32, device="cuda")
+
+    # Mimic vLLM's topk_indices_buffer: slice of a wider contiguous buffer,
+    # whose stride(0) is 32B-aligned.
+    buffer = torch.full((batch_size, 4096), -2, dtype=torch.int32, device="cuda")
+    output_idx = buffer[:, :top_k]
+    result = indexer_topk.deep_select_topk(logits, top_k, output_idx=output_idx)
+    torch.accelerator.synchronize()
+
+    assert result.data_ptr() == output_idx.data_ptr()
+    assert torch.all(buffer[:, top_k:] == -2), "kernel wrote out of bounds"
+
+    row_ends = torch.full((batch_size,), vocab_size, dtype=torch.int32, device="cuda")
+    row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
+    torch_indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
+    for i in range(batch_size):
+        torch_indices[i] = logits[i].topk(top_k, dim=-1)[1]
+
+    assert compare_top_k_results(
+        logits, output_idx, torch_indices, row_starts, row_ends, top_k
+    ), "DeepSelect topk (preallocated output) results don't match torch.topk"
+
+
+def _has_flashinfer_topk() -> bool:
+    if not current_platform.is_cuda():
+        return False
+    # find_spec only: importing flashinfer initializes CUDA at import time.
+    import importlib.util
+
+    return importlib.util.find_spec("flashinfer") is not None
+
+
+def _has_cooperative_topk() -> bool:
+    return _has_device_capability(
+        90
+    ) and not current_platform.is_device_capability_family(120)
+
+
+SPARSE_INDEXER_EXPLICIT_BACKENDS = [
+    pytest.param(
+        "per_row",
+        marks=pytest.mark.skipif(
+            not current_platform.is_cuda(), reason="requires CUDA"
+        ),
+    ),
+    "torch",
+    pytest.param(
+        "persistent",
+        marks=pytest.mark.skipif(
+            not current_platform.is_cuda(), reason="requires CUDA"
+        ),
+    ),
+    pytest.param(
+        "cooperative",
+        marks=pytest.mark.skipif(
+            not _has_cooperative_topk(),
+            reason="cooperative_topk requires SM90+ (non-SM12x)",
+        ),
+    ),
+    pytest.param(
+        "flashinfer",
+        marks=pytest.mark.skipif(
+            not _has_flashinfer_topk(), reason="requires flashinfer top-k"
+        ),
+    ),
+    pytest.param("deep_select", marks=requires_sm100),
+]
+
+
+@pytest.mark.parametrize("next_n", [1, 4])
+@pytest.mark.parametrize("backend", SPARSE_INDEXER_EXPLICIT_BACKENDS)
+@torch.inference_mode()
+def test_sparse_indexer_decode_topk_explicit_backends(
+    backend: str, next_n: int, workspace_init
+) -> None:
+    """Every explicit sparse-indexer decode top-k backend must restrict each
+    row to its seq_len (dirty data past the end must never be selected) and
+    match torch.topk on the valid region."""
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
+
+    set_random_seed(0)
+    torch.set_default_device("cuda:0")
+
+    top_k = 2048
+    batch_size = 4
+    num_rows = batch_size * next_n
+
+    # Per-row effective lens in the (B, next_n) "native spec decode" form;
+    # cooperative/persistent_topk require lengths.numel() == num_rows.
+    row_ends = torch.randint(4000, 16000, (num_rows,), dtype=torch.int32, device="cuda")
+    seq_lens = row_ends.reshape(batch_size, next_n)
+    max_seq_len = int(row_ends.max())
+    # Pad the stride: multiple of 4 for cooperative_topk (TMA) and
+    # 1024B-aligned (256 fp32) for DeepSelect.
+    stride = (max_seq_len + 255) // 256 * 256
+
+    row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
+
+    logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda")
+    col_idx = torch.arange(stride, device="cuda")
+    logits[col_idx[None, :] >= row_ends[:, None]] = 1e30  # dirty tail
+
+    indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda")
+    cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
+    with set_current_vllm_config(cfg):
+        SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, max_seq_len)
+    torch.accelerator.synchronize()
+
+    # k_i == top_k for every row here (row_ends >= 3997 > top_k).
+    assert torch.all(indices >= 0)
+    torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
+    for i in range(num_rows):
+        row_end = int(row_ends[i])
+        torch_indices[i] = logits[i, :row_end].topk(top_k, dim=-1)[1]
+        assert torch.all(indices[i] < row_end), (
+            f"{backend}: row {i} selected indices past its seq_len"
+        )
+
+    assert compare_top_k_results(
+        logits, indices, torch_indices, row_starts, row_ends, top_k
+    ), f"{backend} results don't match torch.topk"
+
+
+@pytest.mark.parametrize(
+    "backend",
+    [
+        "torch",
+        pytest.param(
+            "per_row",
+            marks=pytest.mark.skipif(
+                not current_platform.is_cuda(), reason="requires CUDA"
+            ),
+        ),
+        pytest.param(
+            "flashinfer",
+            marks=pytest.mark.skipif(
+                not _has_flashinfer_topk(), reason="requires flashinfer top-k"
+            ),
+        ),
+        pytest.param("deep_select", marks=requires_sm100),
+    ],
+)
+@torch.inference_mode()
+def test_sparse_indexer_decode_topk_short_seq_lens(
+    backend: str, workspace_init
+) -> None:
+    """seq_len < next_n: derived per-row ends must clamp at 0 (the reference
+    kernels do max(0, ...)); negative ends are OOB for the ragged kernels."""
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
+
+    set_random_seed(0)
+    next_n = 4
+    top_k = 512
+    vocab = 8192  # 1024B-aligned rows for DeepSelect
+    # 1D (B,) seq_lens: rows of a request share its len and each kernel
+    # derives per-row ends (with clamping) itself.
+    seq_lens = torch.tensor([2, 6], dtype=torch.int32, device="cuda")
+    num_rows = 2 * next_n
+    logits = torch.randn(num_rows, vocab, dtype=torch.float32, device="cuda")
+    indices = torch.full((num_rows, top_k), -2, dtype=torch.int32, device="cuda")
+
+    cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
+    with set_current_vllm_config(cfg):
+        SparseIndexerTopk()(logits, seq_lens, next_n, indices, top_k, vocab)
+    torch.accelerator.synchronize()
+
+    # Derived ends: [2,6] - next_n + 1 + arange(next_n) -> [-1,0,1,2] and
+    # [3,4,5,6], clamped to [0,0,1,2] and [3,4,5,6].
+    for row, end in enumerate([0, 0, 1, 2, 3, 4, 5, 6]):
+        k_i = min(top_k, end)
+        assert torch.all(indices[row, k_i:] == -1), f"{backend}: row {row}"
+        if k_i:
+            ref = logits[row, :end].topk(k_i).indices.sort().values
+            got = indices[row, :k_i].sort().values
+            assert torch.equal(got, ref), f"{backend}: row {row}"
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA")
+@torch.inference_mode()
+def test_sparse_indexer_topk_backend_resolution() -> None:
+    """Auto heuristic chain and fail-fast validation of explicit backends."""
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
+
+    # 1024B-aligned stride: 131072 * 4 bytes.
+    logits = torch.randn(64, 131072, dtype=torch.float32, device="cuda")
+    # 4-divisible (cooperative TMA ok) but not 1024B-aligned (DeepSelect no):
+    # 131076 * 4 = 524304 bytes.
+    unaligned_logits = torch.randn(128, 131076, dtype=torch.float32, device="cuda")[
+        :, :131072
+    ]
+
+    def resolve(
+        backend: str,
+        t: torch.Tensor = logits,
+        k: int = 2048,
+        num_rows: int = 64,
+    ) -> str:
+        cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
+        with set_current_vllm_config(cfg):
+            return SparseIndexerTopk().resolve_backend(t, k, num_rows)
+
+    is_sm100 = current_platform.is_device_capability_family(100)
+    has_coop = _has_cooperative_topk()
+    has_fi = _has_flashinfer_topk()
+
+    # "auto" is the chain cooperative -> persistent -> per_row, with the
+    # cooperative/persistent switch at cooperative_topk's cluster-wave row
+    # limit. It must never select the opt-in backends by itself.
+    if has_coop:
+        from vllm.model_executor.layers.indexer_topk import (
+            AUTO_COOPERATIVE_MAX_ROWS,
+        )
+
+        # cooperative for every batch within its cluster-wave row limit...
+        assert (
+            resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS) == "cooperative"
+        )
+        assert resolve("auto", k=512, num_rows=8) == "cooperative"
+        # ...persistent past the row limit, even where DeepSelect would be
+        # applicable.
+        assert (
+            resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS + 1)
+            == "persistent"
+        )
+        assert resolve("auto", k=512, num_rows=128) == "persistent"
+        # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: the
+        # same row-limit logic applies.
+        assert (
+            resolve("auto", unaligned_logits, num_rows=AUTO_COOPERATIVE_MAX_ROWS)
+            == "cooperative"
+        )
+    # ...and persistent past the row limit.
+    assert resolve("auto", unaligned_logits, num_rows=128) == "persistent"
+    # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels)
+    # -> per_row fallback.
+    assert resolve("auto", k=5000) == "per_row"
+
+    # Explicit values are returned as-is when constraints are met.
+    assert resolve("per_row") == "per_row"
+    assert resolve("torch") == "torch"
+    assert resolve("persistent") == "persistent"
+    if has_coop:
+        assert resolve("cooperative") == "cooperative"
+    if has_fi:
+        assert resolve("flashinfer") == "flashinfer"
+    if is_sm100:
+        # Explicit deep_select has no row-count requirement.
+        assert resolve("deep_select", num_rows=1) == "deep_select"
+        with pytest.raises(RuntimeError, match="constraints"):
+            resolve("deep_select", unaligned_logits)
+
+    # Fail-fast on unmet constraints.
+    with pytest.raises(RuntimeError, match="num_rows must be <= 64"):
+        resolve("cooperative", num_rows=128)
+    with pytest.raises(RuntimeError, match="topk_tokens must be in"):
+        resolve("persistent", k=3000)
+    with pytest.raises(RuntimeError, match="topk_tokens must be in"):
+        resolve("cooperative", k=3000)
+    if not is_sm100:
+        with pytest.raises(RuntimeError, match="SM100a/SM103a"):
+            resolve("deep_select")
diff --git a/tests/kernels/utils.py b/tests/kernels/utils.py
index 8bd63872f3d6..cf68565e9d5b 100644
--- a/tests/kernels/utils.py
+++ b/tests/kernels/utils.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Kernel test utils"""
+"""Kernel test utils."""
 
 import itertools
 import random
@@ -94,16 +94,15 @@ def _collect(result: Any) -> list[torch.Tensor]:
 
 
 class QKVInputs(NamedTuple):
-    """
-    Data structure for representing unpacked attention inputs,
+    """Data structure for representing unpacked attention inputs,
     query/key/values and their sequence lengths.
 
     Attributes:
-
         * {query,key,value}: unpacked (batch_size x padded_seq_len x
                              num_heads x head_size) attention inputs
         * q_seq_lens: query sequence lengths list
         * kv_seq_lens: shared key/value sequence lengths list
+
     """
 
     query: torch.Tensor
@@ -114,16 +113,15 @@ class QKVInputs(NamedTuple):
 
 
 class QKVO(NamedTuple):
-    """
-    Data structure for representing unpacked attention inputs,
+    """Data structure for representing unpacked attention inputs,
     alongside unpacked known-correct attention output
 
     Attributes:
-
         * qkv: unpacked (batch_size x padded_seq_len x
                              num_heads x head_size) attention inputs
         * ideal_output: unpacked (batch_size x padded_seq_len x
                         num_heads x head_size) known-correct attention output
+
     """
 
     qkv: QKVInputs
@@ -131,11 +129,9 @@ class QKVO(NamedTuple):
 
 
 class PackedQKVInputs(NamedTuple):
-    """
-    Data structure for representing packed attention inputs
+    """Data structure for representing packed attention inputs.
 
     Attributes:
-
         * {query,key,value}: packed (number_of_tokens x num_heads
                              x head_size) attention inputs
         * q_start_loc_list: list of query start locations within packed tensor
@@ -143,6 +139,7 @@ class PackedQKVInputs(NamedTuple):
                              packed tensor
         * q_seq_lens: query sequence lengths list
         * kv_seq_lens: shared key/value sequence lengths list
+
     """
 
     query: torch.Tensor
@@ -155,16 +152,15 @@ class PackedQKVInputs(NamedTuple):
 
 
 class PackedQKVO(NamedTuple):
-    """
-    Data structure for representing packed attention inputs,
+    """Data structure for representing packed attention inputs,
     alongside packed known-correct attention output
 
     Attributes:
-
         * packed_qkv: packed (number_of_tokens x num_heads
                       x head_size) attention inputs
         * ideal_output: packed (number_of_tokens x num_heads
                         x head_size) known-correct attention output
+
     """
 
     packed_qkv: PackedQKVInputs | None
@@ -172,13 +168,12 @@ class PackedQKVO(NamedTuple):
 
 
 class KVMemoryMap(NamedTuple):
-    """
-    Data structure for encapsulating KV cache memory mapping.
+    """Data structure for encapsulating KV cache memory mapping.
 
     Attributes:
-
         * block_tables: KV cache block tables
         * slot_mapping: mapping of sequence offset to physical address
+
     """
 
     block_tables: torch.Tensor
@@ -186,18 +181,17 @@ class KVMemoryMap(NamedTuple):
 
 
 class PhaseTestParameters(NamedTuple):
-    """
-    Data structure for encapsulating the test parameters
+    """Data structure for encapsulating the test parameters
     for a given test "phase" (prefill or decode phase) and attention
     scenario (encoder, decoder-self, encoder/decoder-cross)
 
     Attributes:
-
         * packed_qkvo: packed (number_of_tokens x num_heads
                        x head_size) attention inputs & known-correct
                        output
         * kv_mmap: KV cache memory mapping, specific to this test phase &
                    attention scenario
+
     """
 
     packed_qkvo: PackedQKVO
@@ -208,13 +202,11 @@ def maybe_make_int_tensor(
     _list: list[int] | None,
     device: torch.device | str,
 ) -> torch.Tensor:
-    """
-    Convert Python int list to a 1D int torch.Tensor on `device`
+    """Convert Python int list to a 1D int torch.Tensor on `device`.
 
     Returns:
+        If _list is not None: 1D int torch.Tensor on `device`, `None` otherwise.
 
-    * If _list is not None: 1D int torch.Tensor on `device`
-    * None otherwise
     """
     return (
         None if _list is None else torch.tensor(_list, dtype=torch.int, device=device)
@@ -225,13 +217,11 @@ def maybe_make_long_tensor(
     _list: list[int] | None,
     device: torch.device | str,
 ) -> torch.Tensor:
-    """
-    Convert Python int list to a 1D long torch.Tensor on `device`
+    """Convert Python int list to a 1D long torch.Tensor on `device`.
 
     Returns:
+        If _list is not None: 1D long torch.Tensor on `device`, `None` otherwise.
 
-    * If _list is not None: 1D long torch.Tensor on `device`
-    * None otherwise
     """
     return (
         None if _list is None else torch.tensor(_list, dtype=torch.long, device=device)
@@ -239,11 +229,11 @@ def maybe_make_long_tensor(
 
 
 def maybe_max(_list: list | None) -> Number | None:
-    """
+    """Compute the maximum of a list, if it exists.
+
     Returns:
+        If _list is not None: max(_list), `None` otherwise.
 
-    * If _list is not None: max(_list)
-    * None otherwise
     """
     return None if _list is None else max(_list)
 
@@ -252,19 +242,16 @@ def make_causal_mask(
     q_max_seq_len: int,
     kv_max_seq_len: int,
 ) -> torch.Tensor:
-    """
-    Create a q_max_seq_len x kv_max_seq_len causal mask
+    """Create a q_max_seq_len x kv_max_seq_len causal mask.
 
     Arguments:
-
-    * q_max_seq_len: query max seq len
-    * kv_max_seq_len: key/value max seq len
+        q_max_seq_len: query max seq len
+        kv_max_seq_len: key/value max seq len
 
     Returns:
+        2D tensor, q_max_seq_len x kv_max_seq_len
 
-    * 2D tensor, q_max_seq_len x kv_max_seq_len
     """
-
     # Create a matrix where entry (i, j) is True if i >= j
     mask = torch.triu(torch.ones(q_max_seq_len, kv_max_seq_len), diagonal=1)
     # Replace True with float('-inf') and False with 0
@@ -281,8 +268,7 @@ def ref_masked_attention(
     q_seq_lens: list | None = None,
     kv_seq_lens: list | None = None,
 ) -> torch.Tensor:
-    """
-    "Golden" masked attention reference. Supports two types of masking:
+    """Reference ("golden") masked attention. Supports two types of masking:
 
     * Basic attention mask, utilizing {q,kv}_seq_lens args to mask out
       padding elements
@@ -290,21 +276,18 @@ def ref_masked_attention(
       causal
 
     Arguments:
-
-    * query: batch_size x q_padded_seq_len x num_heads x head_size
-    * key: batch_size x kv_padded_seq_len x num_heads x head_size
-    * value: batch_size x kv_padded_seq_len x num_heads x head_size
-    * scale: Attention scale factor
-    * custom_mask: custom attention mask; good place to inject a causal
-      attention mask
-    * q_seq_lens: list of unpadded query seq_lens for each batch index
-    * kv_seq_lens: list of unpadded key/value seq_lens for each batch index
+        query: batch_size x q_padded_seq_len x num_heads x head_size
+        key: batch_size x kv_padded_seq_len x num_heads x head_size
+        value: batch_size x kv_padded_seq_len x num_heads x head_size
+        scale: Attention scale factor
+        custom_mask: custom attention mask; good place to inject a causal attention mask
+        q_seq_lens: list of unpadded query seq_lens for each batch index
+        kv_seq_lens: list of unpadded key/value seq_lens for each batch index
 
     Returns:
+        Attention result, batch_size x q_padded_seq_len x num_heads x head_size
 
-    * Attention result, batch_size x q_padded_seq_len x num_heads x head_size
     """
-
     assert q_seq_lens is not None
     assert kv_seq_lens is not None
 
@@ -346,8 +329,7 @@ def make_qkv(
     attn_type: AttentionType = AttentionType.ENCODER_DECODER,
     force_max_len: bool = False,
 ) -> tuple[QKVInputs, QKVInputs, QKVInputs]:
-    """
-    Construct QKV test tensors for self- and cross-attention.
+    """Construct QKV test tensors for self- and cross-attention.
 
     Generates three query/key/value triplets:
 
@@ -361,30 +343,28 @@ def make_qkv(
     seqlens
 
     Arguments:
-
-    * batch_size
-    * max_q_seq_len: max query seq len
-    * max_kv_seq_len: max key/value seq len
-    * num_heads
-    * head_size
-    * is_encoder_decoder_attn: if True, query seqlen may differ from
-      key/value seqlen (as is often the case for cross-attention);
-      o/w, query/key/value seqlens match at each batch index
-      (max_kv_seq_len is unused)
-    * force_kv_seq_lens: if not None, overrides kv sequence lengths
-    * attn_type: encoder, decoder self, or enc/dec cross attention
-    * force_max_len: if True, all query seqlens are max_q_seq_len; o/w query
-      seqlens are random in [2,max_q_seq_lens]. Same for key/value seqlens
-      and max_kv_seq_len, unless forced by is_encoder_decoder_attn=False
-    * device: CPU or CUDA device
+        batch_size: number of sequences in the batch
+        max_q_seq_len: max query seq len
+        max_kv_seq_len: max key/value seq len
+        num_heads: number of attention heads
+        head_size: head dimension
+        is_encoder_decoder_attn: if True, query seqlen may differ from
+          key/value seqlen (as is often the case for cross-attention);
+          o/w, query/key/value seqlens match at each batch index
+          (max_kv_seq_len is unused)
+        force_kv_seq_lens: if not None, overrides kv sequence lengths
+        attn_type: encoder, decoder self, or enc/dec cross attention
+        force_max_len: if True, all query seqlens are max_q_seq_len; o/w query
+          seqlens are random in [2,max_q_seq_lens]. Same for key/value seqlens
+          and max_kv_seq_len, unless forced by is_encoder_decoder_attn=False
+        device: CPU or CUDA device
 
     Returns:
+        Overall QKVInputs structure (containing full unpacked Q/K/V tensors)
+        Prefill QKVInputs structure (containing all but the last sequence offset)
+        Decode QKVInputs structure (containing all only the last sequence offset)
 
-    * Overall QKVInputs structure (containing full unpacked Q/K/V tensors)
-    * Prefill QKVInputs structure (containing all but the last sequence offset)
-    * Decode QKVInputs structure (containing all only the last sequence offset)
     """
-
     if force_max_len:
         q_seq_lens = [max_q_seq_len for _ in range(batch_size)]
     else:
@@ -474,24 +454,21 @@ def make_qkv(
 def pack_tensor(
     unpacked_tensor: torch.Tensor, seq_lens: list[int], device: torch.device | str
 ) -> tuple[torch.Tensor, list[int]]:
-    """
-    Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an
+    """Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an
     unpadded number_of_tokens x num_heads x head_size tensor, where
     number_of_tokens = sum(seq_lens)
 
     Arguments:
+        unpacked_tensor: batch_size x padded_seq_len x num_heads x head_size
+        seq_lens: list of token counts for each seq
+        device: CPU or CUDA device
 
-    * unpacked_tensor: batch_size x padded_seq_len x num_heads x head_size
-    * seq_lens: list of token counts for each seq
-    * device: CPU or CUDA device
-
-    Returns
+    Returns:
+        packed_tensor: number_of_tokens x num_heads x head_size
+        start_loc_list: start idx of each batch elt in packed_tensor; [0] +
+          list(itertools.accumulate(seq_lens))
 
-    * packed_tensor: number_of_tokens x num_heads x head_size
-    * start_loc_list: start idx of each batch elt in packed_tensor; [0] +
-      list(itertools.accumulate(seq_lens))
     """
-
     num_tok = sum(seq_lens)
     num_heads = unpacked_tensor.shape[-2]
     head_size = unpacked_tensor.shape[-1]
@@ -507,8 +484,7 @@ def pack_tensor(
 
 
 def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs:
-    """
-    Individually pack each of Q, K and V, each with dimensions batch_size x
+    """Individually pack each of Q, K and V, each with dimensions batch_size x
     padded_seq_len x num_heads x head_size, into respective number_of_tokens x
     num_heads x head_size tensors.
 
@@ -517,17 +493,15 @@ def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs:
     For K and V, number_of_tokens = sum(kv_seq_lens)
 
     Arguments:
+        qkv: Unpacked (batch_size x padded_seq_len x num_heads x head_size)
+               attention inputs
+        device: CPU or CUDA device
 
-    * qkv: Unpacked (batch_size x padded_seq_len x num_heads x head_size)
-           attention inputs
-    * device: CPU or CUDA device
-
-    Returns
+    Returns:
+        Packed (number_of_tokens x num_heads x head_size) QKV inputs
+          derived from unpacked inputs
 
-    * Packed (number_of_tokens x num_heads x head_size) QKV inputs
-      derived from unpacked inputs
     """
-
     if qkv.query is None:
         packed_query = None
         q_start_loc_list = None
@@ -563,26 +537,24 @@ def _make_metadata_tensors(
     torch.Tensor,
     int | None,
 ]:
-    """
-    Build scalar & tensor values required to build attention metadata structure.
+    """Build scalar & tensor values required to build attention metadata structure.
 
     Arguments:
-
-    * seq_lens: list of token-counts for each decoder input seq
-    * context_lens: list of context length values for each seq
-    * encoder_seq_lens: list of token-counts for each encoder input seq
-    * device: CPU or CUDA device
+        seq_lens: list of token-counts for each decoder input seq
+        context_lens: list of context length values for each seq
+        encoder_seq_lens: list of token-counts for each encoder input seq
+        device: CPU or CUDA device
 
     Returns:
+        seq_lens_tensor: decoder seq_lens list, as tensor
+        context_lens_tensor: context_lens list, as tensor
+        max_context_len: max(context_lens)
+        max_seq_len: max(seq_lens)
+        seq_start_loc: start idx of each sequence
+        encoder_seq_lens_tensor: encoder seq_lens list, as tensor
+        encoder_seq_start_loc: start idx of each encoder sequence
+        max_encoder_seq_len: encoder seq_lens list, as tensor
 
-    * seq_lens_tensor: decoder seq_lens list, as tensor
-    * context_lens_tensor: context_lens list, as tensor
-    * max_context_len: max(context_lens)
-    * max_seq_len: max(seq_lens)
-    * seq_start_loc: start idx of each sequence
-    * encoder_seq_lens_tensor: encoder seq_lens list, as tensor
-    * encoder_seq_start_loc: start idx of each encoder sequence
-    * max_encoder_seq_len: encoder seq_lens list, as tensor
     """
     seq_lens_tensor = maybe_make_int_tensor(seq_lens, device)
     context_lens_tensor = maybe_make_int_tensor(context_lens, device)
@@ -637,22 +609,21 @@ def make_kv_cache(
     backend: str,
     default_val: float = 0.0,
 ) -> torch.Tensor:
-    """
-    Create a fake KV cache.
+    """Create a fake KV cache.
 
     Arguments:
-
-    * num_blocks: number of blocks in the KV cache
-    * num_heads: number of attention heads
-    * head_size: head dimension
-    * block_size: number of offsets within a block
-    * device: CPU or CUDA device
-    * default_val: initialization value for KV cache elements
+        num_blocks: number of blocks in the KV cache
+        num_heads: number of attention heads
+        head_size: head dimension
+        block_size: number of offsets within a block
+        device: CPU or CUDA device
+        backend: attention backend name; only 'FLASH_ATTN' is supported
+        default_val: initialization value for KV cache elements
 
     Returns:
+        kv_cache: 2 x num_blocks x block_size x num_heads x head_size
+            for backend 'FLASH_ATTN'
 
-    * kv_cache: 2 x num_blocks x block_size x num_heads x head_size
-    *     for backend 'FLASH_ATTN'
     """
     if backend != "FLASH_ATTN":
         raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
@@ -663,8 +634,7 @@ def make_kv_cache(
 
 
 def _num_tokens_to_min_blocks(num_tokens: int, block_size: int) -> int:
-    """
-    Compute the minimum number of blocks required to hold num_tokens tokens,
+    """Compute the minimum number of blocks required to hold num_tokens tokens,
     given block_size
     """
     return (num_tokens + block_size - 1) // block_size
@@ -683,8 +653,7 @@ def split_slot_mapping(
     seq_lens: list[int],
     device: torch.device | str,
 ):
-    """
-    Split a slot mapping into valid prefill- and decode-phase slot mappings.
+    """Split a slot mapping into valid prefill- and decode-phase slot mappings.
 
     Context:
     * Your goal is to test (1) prefill of N prompts, with prompt-lengths
@@ -708,21 +677,19 @@ def split_slot_mapping(
     The N excised entries are appended to obtain the decode-phase slot mapping
 
     Arguments:
-
-    * slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N
-      post-decode sequences
-    * seq_lens: list of N post-decode sequence lengths (K_i + 1 in the
-      description above)
-    * device: cuda, cpu, etc.
+        slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N
+          post-decode sequences
+        seq_lens: list of N post-decode sequence lengths (K_i + 1 in the
+          description above)
+        device: cuda, cpu, etc.
 
     Returns:
+        prefill_slot_mapping: Length-M 1D slot mapping (as Tensor)
+          reflecting all N prefill prompts
+        decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting
+          all N decoded tokens
 
-    * prefill_slot_mapping: Length-M 1D slot mapping (as Tensor)
-      reflecting all N prefill prompts
-    * decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting
-      all N decoded tokens
     """
-
     prefill_slot_mapping = []
     decode_slot_mapping = []
 
@@ -746,8 +713,7 @@ def make_block_tables_slot_mapping(
     device: torch.device | str,
     block_base_addr: int = 0,
 ) -> tuple[torch.Tensor, list[int], int]:
-    """
-    Construct fake block tables & slot mappings.
+    """Construct fake block tables & slot mappings.
 
     For a sequence with num_tokens tokens the minimum number
     of required KV cache blocks is
@@ -772,19 +738,17 @@ def make_block_tables_slot_mapping(
     i.e. the total of prefill prompt tokens + decoded tokens.
 
     Arguments:
-
-    * block_size: number of offsets per block
-    * seq_lens: list of token-counts for each sequence
-    * block_base_addr: the block table base address
-    * device: CPU or CUDA device
+        block_size: number of offsets per block
+        seq_lens: list of token-counts for each sequence
+        block_base_addr: the block table base address
+        device: CPU or CUDA device
 
     Return:
-
     * block_tables_tensor: block table for sequence
     * slot_mapping_list: slot mapping for sequence
     * max_block_idx: the highest block address within this block table
-    """
 
+    """
     # Provision minimum number of KV cache blocks
     num_blocks_list = [
         _num_tokens_to_min_blocks(num_tokens, block_size) for num_tokens in seq_lens
@@ -824,14 +788,14 @@ def make_block_tables_slot_mapping(
 def assert_actual_matches_ideal(
     test_params: PhaseTestParameters, output_under_test: torch.Tensor, backend: str
 ) -> None:
-    """
-    Assert that observed output matches the ideal output
+    """Assert that observed output matches the ideal output
     contained in the test parameters data structure.
 
     Arguments:
+        test_params: Test parameters including packed ideal output
+        output_under_test: actually observed output value
+        backend: attention backend name; only 'FLASH_ATTN' is supported
 
-    * test_params: Test parameters including packed ideal output
-    * output_under_test: actually observed output value
     """
     ideal_output = test_params.packed_qkvo.ideal_output
     if backend != "FLASH_ATTN":
@@ -852,9 +816,7 @@ def fp8_allclose(
     atol: float = 1e-08,
     equal_nan: bool = False,
 ) -> bool:
-    """
-    Reference implementation of torch.allclose
-    """
+    """Reference implementation of torch.allclose."""
     torch._refs._check_close_args(name="torch.allclose", a=a, b=b, rtol=rtol, atol=atol)
 
     return bool(
diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py
index 2c64fb224534..f5ca44f174e2 100644
--- a/tests/lora/conftest.py
+++ b/tests/lora/conftest.py
@@ -33,7 +33,6 @@ def should_do_global_cleanup_after_test(request) -> bool:
     This can provide a ~10x speedup for non-GPU unit tests since they don't need
     to initialize torch.
     """
-
     return not request.node.get_closest_marker("skip_global_cleanup")
 
 
@@ -270,6 +269,30 @@ def qwen3_woofing_lora_files():
     return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA")
 
 
+@pytest.fixture(scope="session")
+def qwen3_guard_star_trek_lora_files():
+    return hf_api().snapshot_download(
+        repo_id="geoffmunn/Qwen3Guard-StarTrek-Classification-0.6B",
+        allow_patterns=["adapter_config.json", "adapter_model.safetensors"],
+    )
+
+
+@pytest.fixture(scope="session")
+def qwen3_guard_new_zealand_lora_files():
+    return hf_api().snapshot_download(
+        repo_id="geoffmunn/Qwen3Guard-NewZealand-Classification-0.6B",
+        allow_patterns=["adapter_config.json", "adapter_model.safetensors"],
+    )
+
+
+@pytest.fixture(scope="session")
+def skywork_qwen3_reward_lora_files():
+    return hf_api().snapshot_download(
+        repo_id="AmirMohseni/skywork-qwen3-0.6b-reward-lora",
+        allow_patterns=["adapter_config.json", "adapter_model.safetensors"],
+    )
+
+
 @pytest.fixture(scope="session")
 def deepseekv2_lora_files():
     return hf_api().snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA")
@@ -343,8 +366,7 @@ def qwen36_moe_3d_lora_files():
 
 @pytest.fixture
 def reset_default_device():
-    """
-    Some tests, such as `test_punica_ops.py`, explicitly set the
+    """Some tests, such as `test_punica_ops.py`, explicitly set the
     default device, which can affect subsequent tests. Adding this fixture
     helps avoid this problem.
     """
diff --git a/tests/lora/test_add_lora.py b/tests/lora/test_add_lora.py
index 628f5bc28eb5..97c37c189522 100644
--- a/tests/lora/test_add_lora.py
+++ b/tests/lora/test_add_lora.py
@@ -55,8 +55,7 @@ async def requests_processing_time(llm, lora_requests: list[LoRARequest]) -> flo
 
 @pytest.mark.asyncio
 async def test_add_lora(chatglm3_lora_files):
-    """
-    The add_lora function is used to preload some LoRA adapters into the
+    """The add_lora function is used to preload some LoRA adapters into the
     engine in anticipation of future requests using these adapters. To test
     this functionality, we use the async engine to process some requests - We
     do it twice, once with add_lora() preloading and once without.
diff --git a/tests/lora/test_default_mm_loras.py b/tests/lora/test_default_mm_loras.py
index 5e3c762f79b9..e858331c9364 100644
--- a/tests/lora/test_default_mm_loras.py
+++ b/tests/lora/test_default_mm_loras.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for applying default registered multimodal loras.
-"""
+"""Tests for applying default registered multimodal loras."""
 
 import os
 import unittest.mock as mock
diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py
index 444b609492ba..8b5b3d577a32 100644
--- a/tests/lora/test_fused_moe_lora_kernel.py
+++ b/tests/lora/test_fused_moe_lora_kernel.py
@@ -41,8 +41,7 @@ def CEILDIV(x, y):
 
 
 def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int):
-    """
-    Split `num_tokens` into `num_sequences` sequences.
+    """Split `num_tokens` into `num_sequences` sequences.
     Each sequence randomly selects 1 LoRA index from [0, max_loras),
     and all tokens in that sequence are assigned this LoRA index.
 
@@ -54,6 +53,7 @@ def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int):
     Returns:
         torch.Tensor: 1D tensor of shape [num_tokens], where each value
                       is the LoRA index assigned to that token.
+
     """
     assert num_sequences > 0 and max_loras > 0
     assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences"
@@ -81,9 +81,8 @@ def assign_loras_to_tokens(num_tokens: int, num_sequences: int, max_loras: int):
 
 
 def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int):
-    """
-    For each token, randomly select `top_k_num` distinct experts out of `num_experts`,
-    and assign normalized random weights that sum to 1.
+    """For each token, randomly select `top_k_num` distinct experts out of
+    `num_experts`, and assign normalized random weights that sum to 1.
 
     Args:
         num_tokens (int): Total number of tokens.
@@ -95,6 +94,7 @@ def assign_experts_to_tokens(num_tokens: int, num_experts: int, top_k_num: int):
                                        expert index for each token.
         expert_weights (torch.Tensor): shape [num_tokens, top_k_num],
                                        normalized weights (sum = 1 per row).
+
     """
     assert top_k_num <= num_experts, "top_k_num must be <= num_experts"
 
@@ -379,8 +379,7 @@ def use_fused_moe_lora_kernel_naive(
     offset=0,
     add_inputs=True,
 ):
-    """
-    Test helper for naive_block_assignment path.
+    """Test helper for naive_block_assignment path.
     Skips moe_lora_align_block_size and uses flattened topk_ids as expert_ids.
     """
     config = {
@@ -472,8 +471,7 @@ def test_fused_moe_lora_kernel_naive_block_assignment(
     device,
     seed,
 ):
-    """
-    Test the naive_block_assignment path of the fused_moe_lora kernel.
+    """Test the naive_block_assignment path of the fused_moe_lora kernel.
     This path is triggered when batch_size * top_k is much smaller than
     num_experts * max_loras, and skips the moe_lora_align_block_size kernel.
     """
diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py
index bce91510d428..e6f84cf0bc6c 100644
--- a/tests/lora/test_layers.py
+++ b/tests/lora/test_layers.py
@@ -31,6 +31,9 @@
 )
 from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
 from vllm.lora.punica_wrapper import get_punica_wrapper
+from vllm.model_executor.layers.fusion.quant_activation import (
+    get_input_quant_key,
+)
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
@@ -56,6 +59,13 @@
     torch.bfloat16: (3e-2, 2e-2),
 }
 
+
+def test_lora_linear_requires_unquantized_input() -> None:
+    layer = RowParallelLinearWithLoRA.__new__(RowParallelLinearWithLoRA)
+    layer._input_quant_key = object()
+    assert get_input_quant_key(layer) is None
+
+
 pytestmark = [
     pytest.mark.skipif(
         not (
@@ -116,8 +126,7 @@ def clean_cache_reset_device(reset_default_device):
 
 @pytest.fixture(autouse=True)
 def skip_cuda_with_stage_false(request):
-    """
-    On cuda-like platforms, we use the same kernels for prefill and decode
+    """On cuda-like platforms, we use the same kernels for prefill and decode
     stage, and 'stage' is generally ignored, so we only need to test once.
     """
     if current_platform.is_cuda_alike() or current_platform.is_xpu():
@@ -143,8 +152,8 @@ def get_random_id_to_index(
         num_slots: The number of slots in the mapping. Must be larger
             than num_loras.
         log: Whether to log the output.
-    """
 
+    """
     if num_loras > num_slots:
         raise ValueError(
             f"num_loras is higher than num_slots: {num_loras} > {num_slots}. "
@@ -180,8 +189,8 @@ def populate_loras(
         repeats: must only be set for column parallel packed
             layers. Indicates the number of loras to compose
             together to create a single lora layer.
-    """
 
+    """
     # Dictionary that maps the lora ID to the
     # corresponding lora weights.
     lora_dict: dict[int, LoRALayerWeights] = dict()
@@ -236,8 +245,8 @@ def create_random_inputs(
         input_range: the range of values to include in the input.
             input_range[0] <= possible input values < input_range[1]
         input_type: the type of values in the input.
-    """
 
+    """
     low, high = input_range
 
     inputs: list[torch.Tensor] = []
diff --git a/tests/lora/test_lora_checkpoints.py b/tests/lora/test_lora_checkpoints.py
index 8db529223b24..9dd0296a88b6 100644
--- a/tests/lora/test_lora_checkpoints.py
+++ b/tests/lora/test_lora_checkpoints.py
@@ -147,10 +147,10 @@ def test_gemma4_lora_weights_mapping():
 def test_gemma4_moe_lora_weights_mapping():
     mapper = Gemma4ForCausalLM.hf_to_vllm_mapper
     name = (
-        "base_model.model.model.language_model.layers.9.moe.experts."
+        "base_model.model.model.language_model.layers.9.experts."
         "gate_up_proj.lora_B.weight"
     )
     assert parse_fine_tuned_lora_name(name, mapper) == (
-        "model.layers.9.moe.gate_up_proj",
+        "model.layers.9.experts.gate_up_proj",
         False,
     )
diff --git a/tests/lora/test_lora_functions.py b/tests/lora/test_lora_functions.py
index ecf2ad9ee654..1b841b70cea2 100644
--- a/tests/lora/test_lora_functions.py
+++ b/tests/lora/test_lora_functions.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Script to test add_lora, remove_lora, pin_lora, list_loras functions.
-"""
+"""Script to test add_lora, remove_lora, pin_lora, list_loras functions."""
 
 import pytest
 
diff --git a/tests/lora/test_lora_weights.py b/tests/lora/test_lora_weights.py
new file mode 100644
index 000000000000..ff132147933e
--- /dev/null
+++ b/tests/lora/test_lora_weights.py
@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+
+from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
+
+RANK, ALPHA = 16, 32
+# PEFTHelper stores alpha/sqrt(r) for use_rslora adapters.
+RSLORA_SCALING = ALPHA / RANK**0.5
+
+
+@pytest.fixture(autouse=True)
+def should_do_global_cleanup_after_test():
+    # Pure-python weight packing; skip GPU/ray cleanup.
+    return False
+
+
+def moe_loras(scaling: float, stacked: bool = False) -> list[LoRALayerWeights]:
+    a_shape = (3, RANK, 4) if stacked else (RANK, 4)
+    b_shape = (3, 6, RANK) if stacked else (6, RANK)
+    return [
+        LoRALayerWeights(
+            module_name=f"experts.w{i}",
+            rank=RANK,
+            lora_alpha=ALPHA,
+            lora_a=torch.randn(*a_shape),
+            lora_b=torch.randn(*b_shape),
+            scaling=scaling,
+        )
+        for i in (1, 2, 3)
+    ]
+
+
+@pytest.mark.parametrize(
+    "packer,stacked",
+    [
+        (PackedLoRALayerWeights.pack_moe, False),
+        (PackedLoRALayerWeights.pack_moe_stacked, True),
+    ],
+)
+def test_moe_packing_uses_stored_scaling(packer, stacked: bool):
+    packed = packer(moe_loras(RSLORA_SCALING, stacked), "experts")
+    assert packed.scaling == pytest.approx([RSLORA_SCALING] * 3)
+
+
+def test_non_gated_moe_keeps_w3_unscaled():
+    packed = PackedLoRALayerWeights.pack_moe(
+        moe_loras(RSLORA_SCALING), "experts", is_non_gated_moe=True
+    )
+    assert packed.scaling[:2] == pytest.approx([RSLORA_SCALING] * 2)
+    assert packed.scaling[2] == 1.0
diff --git a/tests/lora/test_mixtral.py b/tests/lora/test_mixtral.py
index 3868bff79663..a8f1fdda64c3 100644
--- a/tests/lora/test_mixtral.py
+++ b/tests/lora/test_mixtral.py
@@ -32,7 +32,7 @@ def do_sample(
 
 @pytest.mark.parametrize("tp_size", [4])
 def test_mixtral_lora(mixtral_lora_files, tp_size):
-    """Original test, the LoRA model has the common target modules, not all"""
+    """Original test, the LoRA model has the common target modules, not all."""
     if (
         torch.accelerator.device_count() < tp_size
         and tp_size > 1
diff --git a/tests/lora/test_moe_lora_ep_load.py b/tests/lora/test_moe_lora_ep_load.py
index ead40b25522b..60d0c2c9ba7e 100644
--- a/tests/lora/test_moe_lora_ep_load.py
+++ b/tests/lora/test_moe_lora_ep_load.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-End-to-end correctness test for 2D MoE LoRA expert-parallel
+"""End-to-end correctness test for 2D MoE LoRA expert-parallel
 load-time slicing
 """
 
diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py
index 23a8aad1a58c..9fb52ea7275c 100644
--- a/tests/lora/test_peft_helper.py
+++ b/tests/lora/test_peft_helper.py
@@ -20,7 +20,7 @@
     (
         "test_modules_to_save",
         {"modules_to_save": ["lm_head"]},
-        "only supports modules_to_save being None",
+        "Unsupported modules_to_save",
     ),
     ("test_rank_zero", {"r": 0}, "must be a positive integer"),
     ("test_rank_negative", {"r": -8}, "must be a positive integer"),
diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py
index 9e1f99b3f6a3..34cbe061dc42 100644
--- a/tests/lora/test_punica_ops.py
+++ b/tests/lora/test_punica_ops.py
@@ -160,8 +160,7 @@ def check_lora_shrink_kernel(
     seq_length: int,
     scaling: float,
 ):
-    """
-    Compare outputs of torch_ops.sgmv_shrink and triton_ops.lora_shrink
+    """Compare outputs of torch_ops.sgmv_shrink and triton_ops.lora_shrink
     kernels.
     """
     data: PunicaTensors = generate_data_for_nslices(
@@ -234,8 +233,7 @@ def check_lora_expand_kernel(
     seq_length: int,
     add_inputs: bool,
 ):
-    """
-    Compare outputs of torch_ops.sgmv_expand and triton_ops.lora_expand
+    """Compare outputs of torch_ops.sgmv_expand and triton_ops.lora_expand
     kernels.
     """
     data: PunicaTensors = generate_data_for_nslices(
@@ -438,9 +436,7 @@ def test_kernels(
     seed: int,
     op_type: str,
 ):
-    """
-    Tests LoRA kernels.
-    """
+    """Tests LoRA kernels."""
     torch.set_default_device(device)
     torch.accelerator.set_device_index(device)
     set_random_seed(seed)
@@ -491,9 +487,7 @@ def test_kernels_hidden_size(
     seed: int,
     op_type: str,
 ):
-    """
-    Tests SGMV and LoRA kernels.
-    """
+    """Tests SGMV and LoRA kernels."""
     torch.set_default_device(device)
     torch.accelerator.set_device_index(device)
     set_random_seed(seed)
@@ -526,8 +520,7 @@ def test_kernels_hidden_size(
 
 @pytest.mark.parametrize("device", DEVICES)
 def test_add_lora_fused_moe_early_exit(device):
-    """
-    Ensures add_lora_fused_moe does not invoke the LoRA kernel or
+    """Ensures add_lora_fused_moe does not invoke the LoRA kernel or
     modify the output tensor when no_lora_flag_cpu is True
     """
     from types import SimpleNamespace
diff --git a/tests/lora/test_punica_xpu_ops.py b/tests/lora/test_punica_xpu_ops.py
index 585c97cfa547..36966809ebce 100644
--- a/tests/lora/test_punica_xpu_ops.py
+++ b/tests/lora/test_punica_xpu_ops.py
@@ -86,9 +86,7 @@ def check_bgmv_shrink(
     device: str,
     scaling: float,
 ):
-    """
-    Compare vllm.bgmv_shrink against a reference implementation.
-    """
+    """Compare vllm.bgmv_shrink against a reference implementation."""
     seq_length = 1
     data: PunicaTensors = generate_data(
         batches,
@@ -130,9 +128,7 @@ def check_bgmv_expand(
     device: str,
     add_inputs: bool,
 ):
-    """
-    Compare vllm.bgmv_expand against a reference implementation.
-    """
+    """Compare vllm.bgmv_expand against a reference implementation."""
     seq_length = 1
     data: PunicaTensors = generate_data(
         batches,
@@ -172,9 +168,7 @@ def check_bgmv_expand_slice(
     device: str,
     add_inputs: bool,
 ):
-    """
-    Compare vllm.bgmv_expand_slice against a reference implementation.
-    """
+    """Compare vllm.bgmv_expand_slice against a reference implementation."""
     seq_length = 1
     data: PunicaTensors = generate_data_for_expand_nslices(
         batches,
diff --git a/tests/lora/test_qwen3_unembed.py b/tests/lora/test_qwen3_unembed.py
index 4064b443ac6a..a70df902f719 100644
--- a/tests/lora/test_qwen3_unembed.py
+++ b/tests/lora/test_qwen3_unembed.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for Qwen3 unembed LoRA support.
+"""Tests for Qwen3 unembed LoRA support.
 
 This test creates synthetic LoRA weights that include lm_head (output embedding)
 to verify that Qwen3 properly supports LoRA on the unembed/lm_head layer.
diff --git a/tests/lora/test_qwen3_with_multi_loras.py b/tests/lora/test_qwen3_with_multi_loras.py
index 9870da6959db..185ab91c7a3d 100644
--- a/tests/lora/test_qwen3_with_multi_loras.py
+++ b/tests/lora/test_qwen3_with_multi_loras.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This script contains:
+"""This script contains:
 1. test multi loras service with tp >= 2
 2. test multi loras request
 """
@@ -125,8 +124,7 @@ def call_llm_get_outputs(prompt: str, lora_name: str):
         return output_text
 
     def reload_lora(name: str):
-        """
-        reload a lora to simulate the case:
+        """Reload a lora to simulate the case:
         setting `VLLM_ALLOW_RUNTIME_LORA_UPDATING=true`
         for dynamic lora loading and unloading
         """
@@ -205,8 +203,7 @@ def test_multiple_lora_requests():
 def test_load_inplace_offline_reload(
     qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str
 ) -> None:
-    """
-    Test that load_inplace=True allows reloading LoRA adapters with the same ID
+    """Test that load_inplace=True allows reloading LoRA adapters with the same ID
     in offline mode (using LLM class directly).
     """
     llm = LLM(
@@ -255,8 +252,7 @@ def test_load_inplace_offline_reload(
 def test_load_inplace_false_no_reload(
     qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str
 ) -> None:
-    """
-    Test that load_inplace=False prevents reloading when an adapter
+    """Test that load_inplace=False prevents reloading when an adapter
     with the same ID already exists.
     """
     llm = LLM(
diff --git a/tests/lora/test_qwenvl.py b/tests/lora/test_qwenvl.py
index 3362aa46ed26..360b9aff05b0 100644
--- a/tests/lora/test_qwenvl.py
+++ b/tests/lora/test_qwenvl.py
@@ -3,8 +3,6 @@
 from dataclasses import dataclass
 
 import pytest
-from packaging.version import Version
-from transformers import __version__ as TRANSFORMERS_VERSION
 
 import vllm
 from tests.conftest import VllmRunner
@@ -29,24 +27,16 @@ class TestConfig:
 
     def __post_init__(self):
         if self.mm_processor_kwargs is None:
-            # There is a bug in transformers v4 where size is ignored by
-            # `Qwen2VLProcessor.__call__`
-            if Version(TRANSFORMERS_VERSION) < Version("5.2.0"):
-                self.mm_processor_kwargs = {
-                    "min_pixels": 28 * 28,
-                    "max_pixels": 1280 * 28 * 28,
-                }
-            else:
-                self.mm_processor_kwargs = {
-                    "size": {
-                        "shortest_edge": 28 * 28,
-                        "longest_edge": 1280 * 28 * 28,
-                    }
+            self.mm_processor_kwargs = {
+                "size": {
+                    "shortest_edge": 28 * 28,
+                    "longest_edge": 1280 * 28 * 28,
                 }
+            }
 
 
 class Qwen2VLTester:
-    """Test helper for Qwen2 VL models with LoRA"""
+    """Test helper for Qwen2 VL models with LoRA."""
 
     PROMPT_TEMPLATE = (
         "<|im_start|>system\nYou are a helpful assistant.<|im_end|>"
@@ -199,7 +189,7 @@ def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 def test_qwen2vl_lora(qwen2vl_lora_files):
-    """Test Qwen 2.0 VL model with LoRA"""
+    """Test Qwen 2.0 VL model with LoRA."""
     config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
     with Qwen2VLTester(config) as tester:
         # Test with different LoRA IDs
@@ -230,7 +220,7 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
     current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
 )
 def test_qwen25vl_lora(qwen25vl_lora_files):
-    """Test Qwen 2.5 VL model with LoRA"""
+    """Test Qwen 2.5 VL model with LoRA."""
     config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files)
     with Qwen2VLTester(config) as tester:
         # Test with different LoRA IDs
@@ -292,8 +282,7 @@ def test_qwen2vl_multiple_lora_types(
     qwen2vl_vision_tower_lora_files,
     monkeypatch: pytest.MonkeyPatch,
 ):
-    """
-    Test multiple LoRA adapter types (language, vision tower + connector,
+    """Test multiple LoRA adapter types (language, vision tower + connector,
     vision tower only) using the same LLM instance to verify mm_encoder_cache
     behavior with different LoRA requests.
 
diff --git a/tests/lora/test_sequence_classification.py b/tests/lora/test_sequence_classification.py
new file mode 100644
index 000000000000..86721bfb6ef0
--- /dev/null
+++ b/tests/lora/test_sequence_classification.py
@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from tests.utils import multi_gpu_test
+from vllm import LLM, PoolingParams
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.lora.request import LoRARequest
+
+MODEL_NAME = "Qwen/Qwen3-0.6B"
+NATIVE_MODEL_NAME = "Skywork/Skywork-Reward-V2-Qwen3-0.6B"
+PROMPTS = {
+    "star_trek": "Does warp drive appear in Star Trek?",
+    "new_zealand": "Wellington is the capital of New Zealand.",
+}
+EXPECTED_OUT = {
+    "star_trek": [13.6640625, -13.6640625],
+    "new_zealand": [-4.80078125, 8.6171875],
+    "native": [-1.52734375],
+}
+
+
+def _classify_logits(
+    llm: LLM,
+    prompts: list[str],
+    requests: LoRARequest | list[LoRARequest],
+) -> torch.Tensor:
+    params = [PoolingParams(task="classify", use_activation=False) for _ in prompts]
+    outputs = llm.classify(
+        prompts,
+        pooling_params=params,
+        lora_request=requests,
+        use_tqdm=False,
+    )
+    return torch.tensor([output.outputs.probs for output in outputs])
+
+
+def test_converted_model_with_modules_to_save(
+    qwen3_guard_star_trek_lora_files: str,
+) -> None:
+    prompt = PROMPTS["star_trek"]
+    llm = LLM(
+        model=MODEL_NAME,
+        runner="pooling",
+        convert="classify",
+        hf_overrides={"num_labels": 2},
+        dtype="float16",
+        enable_lora=True,
+        max_lora_rank=16,
+        enforce_eager=True,
+        max_model_len=512,
+        gpu_memory_utilization=0.5,
+    )
+    actual = _classify_logits(
+        llm,
+        [prompt],
+        LoRARequest("star-trek", 1, qwen3_guard_star_trek_lora_files),
+    )[0]
+
+    torch.testing.assert_close(
+        actual,
+        torch.tensor(EXPECTED_OUT["star_trek"]),
+        atol=2e-2,
+        rtol=2e-2,
+    )
+
+    del llm
+    cleanup_dist_env_and_memory()
+
+
+def test_native_classification_model_with_modules_to_save(
+    skywork_qwen3_reward_lora_files: str,
+) -> None:
+    prompt = "Which response is more helpful and correct?"
+    llm = LLM(
+        model=NATIVE_MODEL_NAME,
+        runner="pooling",
+        dtype="float16",
+        enable_lora=True,
+        max_lora_rank=8,
+        enforce_eager=True,
+        max_model_len=512,
+        gpu_memory_utilization=0.5,
+    )
+    actual = _classify_logits(
+        llm,
+        [prompt],
+        LoRARequest("native", 1, skywork_qwen3_reward_lora_files),
+    )[0]
+
+    torch.testing.assert_close(
+        actual,
+        torch.tensor(EXPECTED_OUT["native"]),
+        atol=2e-2,
+        rtol=2e-2,
+    )
+
+    del llm
+    cleanup_dist_env_and_memory()
+
+
+def _test_batched_loras(
+    qwen3_guard_star_trek_lora_files: str,
+    qwen3_guard_new_zealand_lora_files: str,
+    tp_size: int,
+) -> None:
+    names = ["star_trek", "new_zealand"]
+    prompts = [PROMPTS[name] for name in names]
+    adapter_paths = [
+        qwen3_guard_star_trek_lora_files,
+        qwen3_guard_new_zealand_lora_files,
+    ]
+    requests = [
+        LoRARequest(name, index, adapter_path)
+        for index, (name, adapter_path) in enumerate(zip(names, adapter_paths), start=1)
+    ]
+
+    llm = LLM(
+        model=MODEL_NAME,
+        runner="pooling",
+        convert="classify",
+        hf_overrides={"num_labels": 2},
+        dtype="float16",
+        enable_lora=True,
+        max_loras=2,
+        max_cpu_loras=2,
+        max_lora_rank=16,
+        enforce_eager=True,
+        max_model_len=512,
+        gpu_memory_utilization=0.5,
+        tensor_parallel_size=tp_size,
+    )
+    actual = _classify_logits(llm, prompts, requests)
+
+    torch.testing.assert_close(
+        actual,
+        torch.tensor([EXPECTED_OUT[name] for name in names]),
+        atol=2e-2,
+        rtol=2e-2,
+    )
+
+    del llm
+    cleanup_dist_env_and_memory()
+
+
+def test_batched_loras(
+    qwen3_guard_star_trek_lora_files: str,
+    qwen3_guard_new_zealand_lora_files: str,
+) -> None:
+    _test_batched_loras(
+        qwen3_guard_star_trek_lora_files,
+        qwen3_guard_new_zealand_lora_files,
+        tp_size=1,
+    )
+
+
+@multi_gpu_test(num_gpus=2)
+def test_batched_loras_tp(
+    qwen3_guard_star_trek_lora_files: str,
+    qwen3_guard_new_zealand_lora_files: str,
+) -> None:
+    _test_batched_loras(
+        qwen3_guard_star_trek_lora_files,
+        qwen3_guard_new_zealand_lora_files,
+        tp_size=2,
+    )
diff --git a/tests/lora/test_whisper.py b/tests/lora/test_whisper.py
index 6f1a894cf916..73cc550b294b 100644
--- a/tests/lora/test_whisper.py
+++ b/tests/lora/test_whisper.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Integration tests for Whisper models with LoRA adapters.
+"""Integration tests for Whisper models with LoRA adapters.
 
 These tests verify that Whisper models can correctly load and use LoRA adapters
 for speech-to-text transcription tasks.
diff --git a/tests/lora/utils.py b/tests/lora/utils.py
index e5ce7a88464c..8ee2f8af8224 100644
--- a/tests/lora/utils.py
+++ b/tests/lora/utils.py
@@ -116,8 +116,7 @@ class PunicaTensors:
     token_lora_mapping: torch.Tensor
 
     def meta(self) -> tuple[int, int]:
-        """
-        Infer max_seq_length and token_nums from the tensors
+        """Infer max_seq_length and token_nums from the tensors
         and return them.
         """
         max_seq_length = self.seq_len_tensor.max()
diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py
index 27368ead27dd..9d4f9beeb9a3 100644
--- a/tests/model_executor/kernels/test_b12x_linear.py
+++ b/tests/model_executor/kernels/test_b12x_linear.py
@@ -115,7 +115,9 @@ def test_b12x_backend_registration_priority_and_selection(
     assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after)
 
     monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
-    monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
+    monkeypatch.setattr(
+        linear_mod, "_get_linear_backend", lambda *, quantization: "b12x"
+    )
     monkeypatch.setattr(
         kernel_cls,
         "is_supported",
@@ -784,7 +786,9 @@ def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None:
     import vllm.model_executor.kernels.linear as linear_mod
 
     monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
-    monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
+    monkeypatch.setattr(
+        linear_mod, "_get_linear_backend", lambda *, quantization: "b12x"
+    )
     monkeypatch.setitem(
         linear_mod._POSSIBLE_NVFP4_KERNELS,
         PlatformEnum.CUDA,
diff --git a/tests/model_executor/layers/test_activation.py b/tests/model_executor/layers/test_activation.py
new file mode 100644
index 000000000000..d0ef42c99afe
--- /dev/null
+++ b/tests/model_executor/layers/test_activation.py
@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from functools import partial
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul
+from vllm.platforms import CpuArchEnum, current_platform
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cpu()
+    or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC,
+    reason="Native CPU activations require a non-PowerPC CPU platform",
+)
+@pytest.mark.parametrize(
+    "activation_factory,reference",
+    [
+        (SiluAndMul, nn.SiLU()),
+        (GeluAndMul, nn.GELU()),
+        (partial(GeluAndMul, approximate="tanh"), nn.GELU(approximate="tanh")),
+    ],
+    ids=["silu", "gelu", "gelu_tanh"],
+)
+def test_cpu_act_and_mul_without_compiled_ops(
+    default_vllm_config, monkeypatch, activation_factory, reference
+):
+    """CPU activations must construct and run without compiled wheel operators."""
+    default_vllm_config.compilation_config.custom_ops = ["all"]
+    monkeypatch.setattr(torch.ops, "_C", object())
+    x = torch.tensor([[-2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0]], device="cpu")
+
+    layer = activation_factory()
+
+    torch.testing.assert_close(layer(x), reference(x[:, :4]) * x[:, 4:])
diff --git a/tests/model_executor/layers/test_conv.py b/tests/model_executor/layers/test_conv.py
new file mode 100644
index 000000000000..f7b23fbdcd45
--- /dev/null
+++ b/tests/model_executor/layers/test_conv.py
@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.layers.conv import Conv3dLayer
+
+
+@pytest.mark.parametrize(
+    "input_shape,kernel_size,out_channels,bias",
+    [
+        pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1280, False, id="minimax_m3"),
+        pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1152, False, id="qwen2_vl"),
+        pytest.param((4, 3, 2, 14, 14), (2, 14, 14), 1152, True, id="qwen3_vl"),
+        pytest.param((4, 3, 1, 14, 14), (1, 14, 14), 1536, True, id="glm_vl"),
+        pytest.param((1, 3, 4, 28, 28), (2, 14, 14), 32, True, id="multiple_patches"),
+    ],
+)
+def test_conv3d_patch_embedding_correctness(
+    default_vllm_config, input_shape, kernel_size, out_channels, bias
+):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+    reference = nn.Conv3d(
+        input_shape[1], out_channels, kernel_size, stride=kernel_size, bias=bias
+    ).to(device=device, dtype=dtype)
+    layer = Conv3dLayer(
+        input_shape[1],
+        out_channels,
+        kernel_size,
+        stride=kernel_size,
+        bias=bias,
+        params_dtype=dtype,
+    ).to(device)
+    layer.load_state_dict(reference.state_dict(), strict=True)
+    assert layer.enable_linear
+
+    x = torch.randn(input_shape, device=device, dtype=dtype)
+    expected = reference(x)
+    atol = rtol = 0.02 if dtype == torch.bfloat16 else 1e-5
+    torch.testing.assert_close(layer(x), expected, atol=atol, rtol=rtol)
diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py
index 37693f18c436..08fb6818e3ab 100644
--- a/tests/model_executor/layers/test_fused_shared_expert.py
+++ b/tests/model_executor/layers/test_fused_shared_expert.py
@@ -808,7 +808,6 @@ def test_models_fse_init(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """Model construction resolves FSE consistently with Quark quantization."""
-
     quantization_config: dict[str, Any] = (
         get_deepseek_v4_quark_config(["layers.0.ffn.shared_experts"] if exclude else [])
         if model_type == "deepseek_v4"
@@ -832,6 +831,7 @@ def test_models_fse_init(
         runner_type="generate",
         is_moe=True,
         logits_processors=None,
+        rswa_window=None,
     )
     vllm_config.parallel_config.enable_expert_parallel = False
     if model_type == "deepseek_v4":
diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py
index 8bd97a3ab47f..77c52fb291dd 100644
--- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py
+++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py
@@ -7,6 +7,7 @@
 import torch
 
 import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer
+import vllm.models.deepseek_v32.attention as deepseek_attention
 from vllm.config import CUDAGraphMode
 from vllm.models.deepseek_v32 import attention as deepseek_v32_attention
 from vllm.models.deepseek_v32.attention import DeepseekV32Attention
@@ -16,6 +17,51 @@
 MLA_LAYER = "model.layers.0.self_attn.attn"
 
 
+def test_sparse_attention_refreshes_batch_state_inside_eager_segment(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    metadata = SimpleNamespace(num_actual_tokens=1)
+
+    class BatchStateRefreshed(Exception):
+        pass
+
+    def refresh_batch_state(attn_metadata):
+        assert attn_metadata is metadata
+        raise BatchStateRefreshed
+
+    layer = SimpleNamespace(
+        indexer=None,
+        skip_topk=True,
+        layer_name=MLA_LAYER,
+        impl=SimpleNamespace(
+            prepare_for_batch=refresh_batch_state,
+            record_logical_topk_ready=lambda: None,
+        ),
+    )
+    monkeypatch.setattr(
+        deepseek_attention,
+        "get_attention_context",
+        lambda layer_name: (metadata, None, torch.empty(1), None),
+    )
+
+    with pytest.raises(BatchStateRefreshed):
+        DeepseekV32Attention._sparse_indexer_and_attn(
+            layer,
+            torch.empty(1, dtype=torch.long),
+            torch.empty(1, 1),
+            torch.empty(1, 1, 1),
+            torch.empty(1, 1, 1),
+            None,
+            None,
+            None,
+            None,
+            None,
+            torch.empty(1, 1, 1),
+            torch.empty(1, 1, 1),
+            torch.empty(1, 1),
+        )
+
+
 def make_indexer_metadata(
     *,
     num_decodes: int = 0,
@@ -246,6 +292,10 @@ def record_forward_impl(*args):
         _fp8_query=fp8_query,
         _use_sparse_mha=lambda _: True,
         rotary_emb=lambda _positions, q: (q + 1, None),
+        impl=SimpleNamespace(
+            record_logical_topk_ready=lambda: None,
+            prepare_for_batch=lambda _: None,
+        ),
         forward_impl=record_forward_impl,
     )
     q_nope = torch.randn(2, 1, 2)
@@ -290,3 +340,714 @@ def record_forward_impl(*args):
             expected_args[1:],
         )
     )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_select_candidate_blocks_tolerates_empty_rows():
+    """Full-cudagraph decode pads the batch with seq_len-0 rows. The newest
+    block pin must not index -1 for them (device-side assert); they select
+    no candidate blocks and real rows still pin their newest block."""
+    block_size, topk_blocks = 8, 3
+    logits = torch.zeros(3, 64, device="cuda")
+    logits[0, 3] = 5.0  # block 0 scores highest for row 0
+    logits[2, 9] = 5.0  # block 1 scores highest for row 2
+    row_ks = torch.zeros(3, dtype=torch.int64, device="cuda")
+    row_ke = torch.tensor([40, 0, 17], device="cuda")
+    out = torch.empty(3, topk_blocks, dtype=torch.int32, device="cuda")
+
+    sparse_indexer._select_candidate_blocks(
+        logits, row_ks, row_ke, topk_blocks, block_size, out
+    )
+
+    assert out[1].tolist() == [-1, -1, -1]
+    assert out[0, 0].item() == 4 and 0 in out[0].tolist()  # newest block pinned
+    assert out[2, 0].item() == 2 and 1 in out[2].tolist()
+    assert (out[0] >= 0).all() and (out[2, :2] >= 0).all()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("score", [0.0, float("nan")])
+@pytest.mark.parametrize("start", [0, 3])
+def test_candidate_selection_keeps_valid_blocks_with_nan_scores(score, start):
+    """NaN warmup scores must not turn real blocks into repeated sparse padding."""
+    logits = torch.full((3, 64), score, device="cuda")
+    starts = torch.full((3,), start, device="cuda", dtype=torch.int32)
+    lengths = torch.tensor([0, 17, 35], device="cuda", dtype=torch.int32)
+    out = torch.empty(3, 8, device="cuda", dtype=torch.int32)
+
+    sparse_indexer._select_candidate_blocks(logits, starts, starts + lengths, 8, 8, out)
+
+    for row, count in enumerate([0, 3, 5]):
+        valid = out[row][out[row] >= 0].sort().values
+        torch.testing.assert_close(
+            valid, torch.arange(count, device="cuda", dtype=torch.int32)
+        )
+        assert (out[row] == -1).sum().item() == 8 - count
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "width,block_size,k,decode",
+    [
+        (73, 8, 16, False),
+        (97, 3, 7, False),
+        (32768, 8, 2048, False),
+        (32768, 8, 2048, True),
+    ],
+)
+def test_candidate_kernels_preserve_packed_bounds_and_padding(
+    width, block_size, k, decode
+):
+    """Preserve top-k ties, newest blocks, empty rows and candidate clamping."""
+    from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+        apply_candidate_mask,
+        select_candidate_blocks,
+    )
+
+    torch.manual_seed(42)
+    rows = 6
+    logits = torch.randn(rows, width * 2, device="cuda")[:, ::2]
+    logits[0, :16] = 0
+    logits[2, 5] = float("nan")
+    starts = torch.tensor([0, 5, 3, 0, 7, 1], device="cuda", dtype=torch.int32)
+    ends = torch.tensor([0, width, width - 1, 1, 11, width], device="cuda")
+    repeat = 1
+    if decode:
+        starts = None
+        ends = torch.tensor([0, 1, width], device="cuda", dtype=torch.int32)
+        repeat = 2
+    ks = torch.zeros(rows, device="cuda", dtype=torch.int64) if decode else starts
+    ke = ends.repeat_interleave(repeat)
+    cols = torch.arange(width, device="cuda")
+    valid = (cols >= ks[:, None]) & (cols < ke[:, None])
+    scores = logits.masked_fill(~valid, -torch.inf)
+    blocks = ((cols - ks[:, None]) // block_size).clamp(min=0).long()
+    nblocks = (width + block_size - 1) // block_size
+    reduced = logits.new_full((rows, nblocks), -torch.inf)
+    reduced.scatter_reduce_(1, blocks, scores, reduce="amax", include_self=True)
+    lengths = ke - ks
+    last = ((lengths - 1) // block_size).clamp(min=0).long()
+    reduced.scatter_(
+        1, last[:, None], torch.where(lengths > 0, torch.inf, -torch.inf)[:, None]
+    )
+    top = reduced.topk(min(k, nblocks), dim=-1)
+    expected = torch.full((rows, k), -1, device="cuda", dtype=torch.int32)
+    expected[:, : top.indices.shape[1]] = torch.where(
+        top.values != -torch.inf, top.indices, -1
+    ).int()
+    actual = torch.empty(rows, k * 2, device="cuda", dtype=torch.int32)[:, ::2]
+    select_candidate_blocks(logits, starts, ends, k, block_size, actual, repeat)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    candidates = actual.clone()
+    candidates[:, 0] = nblocks + 10
+    candidates[:, 1] = 0
+    candidates[:, 2] = 0
+    positions = ks[:, None, None] + candidates.long()[:, :, None] * block_size
+    positions = positions + torch.arange(block_size, device="cuda")
+    keep = torch.zeros(rows, width, device="cuda", dtype=torch.int8)
+    keep.scatter_reduce_(
+        1,
+        positions.clamp(0, width - 1).reshape(rows, -1),
+        (candidates >= 0)[:, :, None]
+        .expand(-1, -1, block_size)
+        .reshape(rows, -1)
+        .to(torch.int8),
+        reduce="amax",
+        include_self=True,
+    )
+    reference = logits.masked_fill((keep == 0) | ~valid, -torch.inf)
+    apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat)
+    torch.testing.assert_close(logits, reference, rtol=0, atol=0, equal_nan=True)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat)
+    for block in (1, -1):
+        candidates.fill_(-1)
+        candidates[:, 0] = block
+        logits.fill_(3.0)
+        graph.replay()
+        keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block)
+        reference = torch.where(keep, 3.0, -torch.inf)
+        torch.testing.assert_close(logits, reference, rtol=0, atol=0)
+
+
+# --- DeepGEMM sparse-MQA indexer path (V4.1 two-level candidate filtering) ---
+
+
+def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int):
+    """Production-faithful candidates: top min(avail, K) blocks per row —
+    full coverage for lens <= K*cbk, a fixed-size subset beyond that."""
+    avail = ((lens.cpu() + cbk - 1) // cbk).tolist()
+    rows = []
+    for a in avail:
+        n = min(a, num_candidates)
+        picks = torch.randperm(a)[:n].sort().values[:num_candidates].tolist()
+        rows.append(picks + [-1] * (num_candidates - len(picks)))
+    return torch.tensor(rows, dtype=torch.int32)
+
+
+def _candidate_token_mask(
+    candidates: torch.Tensor, cbk: int, total: int, ks: torch.Tensor | None = None
+):
+    """[rows, total] bool mask of the candidate-covered token positions.
+
+    ``ks`` converts request-local candidate blocks to absolute positions
+    (prefill's packed workspace); pass None when already request-local.
+    """
+    rows = candidates.shape[0]
+    device = candidates.device
+    # Sentinel block keeps -1 padding from scattering onto real columns.
+    width = max(total, (candidates.shape[1] + 1) * cbk)
+    if ks is not None:
+        width += int(ks.max()) + cbk
+    in_cand = torch.zeros(rows, width, dtype=torch.bool, device=device)
+    block_id = torch.where(candidates >= 0, candidates.long(), candidates.shape[1])
+    tok = block_id.unsqueeze(2) * cbk + torch.arange(cbk, device=device)
+    if ks is not None:
+        tok = tok + ks.unsqueeze(1).unsqueeze(2)
+    keep = (candidates >= 0).unsqueeze(2).expand(-1, -1, cbk)
+    in_cand.scatter_(1, tok.flatten(1), keep.flatten(1))
+    return in_cand[:, :total]
+
+
+def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor:
+    """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid
+    count are -1 (the indexer's "no token" sentinel)."""
+    values, indices = masked.topk(topk, dim=1)
+    return torch.where(values > float("-inf"), indices.int(), -1)
+
+
+def _assert_same_selection(
+    logits: torch.Tensor,
+    got: torch.Tensor,
+    expected: torch.Tensor,
+    row_starts: torch.Tensor | None = None,
+) -> None:
+    """Both top-k index sets must pick the same multiset of logit values per
+    row (-1 padding must match too). Indices themselves may differ: the
+    top-k kernels break ties among equal values in nondeterministic order."""
+    for r in range(logits.shape[0]):
+        base = int(row_starts[r]) if row_starts is not None else 0
+        g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0]
+        assert g.numel() == e.numel(), f"row {r}: {g.numel()} vs {e.numel()} selected"
+        vg = logits[r][(base + g).long()].sort().values
+        ve = logits[r][(base + e).long()].sort().values
+        assert torch.equal(vg, ve), f"row {r}: selected values differ"
+
+
+def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict:
+    """Caller-owned buffers the sparse path writes into (the metadata builder
+    owns these in production)."""
+    i32 = dict(dtype=torch.int32, device="cuda")
+    return dict(
+        sparse_indices=torch.zeros(rows, num_sparse, **i32),
+        end=torch.zeros(rows, **i32),
+        col_indices=torch.zeros(rows, topk, **i32),
+    )
+
+
+def _expand_candidates_reference(
+    candidates: torch.Tensor, ks: torch.Tensor, ke: torch.Tensor, cbk: int, sbk: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand
+    each candidate block into ``cbk // sbk`` sparse blocks, keep those inside
+    ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by
+    repeating the last valid block, and count the valid sparse columns."""
+    ratio = cbk // sbk
+    rows, num_candidates = candidates.shape
+    num_sparse = num_candidates * ratio
+    big = torch.iinfo(torch.int32).max
+    offsets = torch.arange(ratio, device=candidates.device, dtype=torch.int32)
+    blocks = (candidates.unsqueeze(-1) * ratio + offsets).flatten(1)
+    max_block = ((ke - ks).clamp(min=0) + sbk - 1) // sbk
+    valid = (blocks >= 0) & (blocks < max_block.unsqueeze(1))
+    sorted_blocks, _ = torch.where(valid, blocks, big).sort(dim=1)
+    n_valid = (sorted_blocks != big).sum(dim=1, dtype=torch.int32)
+    base = (ks // sbk).unsqueeze(1)
+    kernel_blocks = torch.where(sorted_blocks != big, sorted_blocks, 0) + base
+    last_slot = (n_valid - 1).clamp(min=0).unsqueeze(1).long()
+    last_block = kernel_blocks.gather(1, last_slot)
+    pad = torch.arange(num_sparse, device=candidates.device) >= n_valid.unsqueeze(1)
+    sparse_indices = torch.where(pad, last_block, kernel_blocks)
+    last_start = last_block.squeeze(1) * sbk + ks % sbk
+    last_fill = (ke - last_start).clamp(min=1, max=sbk)
+    end = torch.where(n_valid > 0, (n_valid - 1) * sbk + last_fill, 0)
+    return sparse_indices.int(), end.int()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel")
+def test_candidate_blocks_to_sparse_indices_math():
+    """Candidate->sparse-block expansion: ks anchoring, range filtering,
+    repeat-last padding, and the per-row valid column count (``end``)."""
+    from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+        candidate_blocks_to_sparse_indices,
+    )
+
+    # Rows cover: exactly one block, an unaligned ks with a partial tail
+    # block, an empty row, duplicates / -1 padding / out-of-range garbage
+    # (warmup runs touch uninitialized buffers), and int32-overflow junk.
+    ks = torch.tensor([64, 13, 100, 8, 0], dtype=torch.int32, device="cuda")
+    ke = torch.tensor([72, 30, 100, 40, 1000], dtype=torch.int32, device="cuda")
+    candidates = torch.tensor(
+        [
+            [0, -1, -1, -1],
+            [2, 0, -1, -1],
+            [3, -1, -1, -1],
+            [1, 1, 7, 0],
+            [2**30, -(2**30), 0, -1],
+        ],
+        dtype=torch.int32,
+        device="cuda",
+    )
+    indices, end = candidate_blocks_to_sparse_indices(
+        candidates, ks, ke, candidate_block_size=8, sparse_block_kv=8
+    )
+    assert indices.tolist() == [
+        [8, 8, 8, 8],
+        [1, 3, 3, 3],
+        [12, 12, 12, 12],
+        [1, 2, 2, 2],
+        [0, 0, 0, 0],
+    ]
+    # end: row 1's last block covers [29, 37) but ke = 30 -> one valid token;
+    # row 3 keeps the duplicated block (production candidates are unique).
+    assert end.tolist() == [8, 9, 0, 24, 8]
+
+    # Wider candidate blocks expand into multiple sparse blocks.
+    idx2, end2 = candidate_blocks_to_sparse_indices(
+        torch.tensor([[1, 0]], dtype=torch.int32, device="cuda"),
+        torch.tensor([0], dtype=torch.int32, device="cuda"),
+        torch.tensor([48], dtype=torch.int32, device="cuda"),
+        candidate_block_size=16,
+        sparse_block_kv=8,
+    )
+    assert idx2.tolist() == [[0, 1, 2, 3]]
+    assert end2.tolist() == [32]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel")
+@pytest.mark.parametrize("cbk,sbk", [(8, 8), (16, 8), (16, 16)])
+@pytest.mark.parametrize("end_offset", [0, 1])
+def test_candidate_blocks_to_sparse_indices_matches_reference(
+    cbk: int, sbk: int, end_offset: int
+):
+    """The Triton expansion equals the PyTorch reference on production-shaped
+    random inputs, including strided (column-sliced) candidate rows and
+    strided output buffers."""
+    from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+        candidate_blocks_to_sparse_indices,
+    )
+
+    torch.manual_seed(0)
+    rows, num_candidates = 41, 2048
+    ratio = cbk // sbk
+    lens = torch.randint(1, 40000, (rows,), dtype=torch.int32)
+    lens[0], lens[1] = 1, 3
+    ks = torch.zeros(rows, dtype=torch.int32)
+    ks[1:] = lens.cumsum(0)[:-1]
+    ks, ke = ks.cuda(), (ks + lens).cuda()
+    candidates = _make_candidates(lens, num_candidates, cbk).cuda()
+    # Row 2 gets out-of-range garbage past the padding, as a CUDA-graph
+    # warm-up on uninitialized buffers would.
+    candidates[2, -3:] = torch.tensor([2**30, -(2**29), 10**6], dtype=torch.int32)
+
+    ref_indices, ref_end = _expand_candidates_reference(candidates, ks, ke, cbk, sbk)
+
+    cand_buf = torch.full((rows, num_candidates + 64), -7, dtype=torch.int32).cuda()
+    cand_buf[:, :num_candidates] = candidates
+    si_buf = torch.zeros(rows, num_candidates * ratio + 32, dtype=torch.int32).cuda()
+    end_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda")
+    ks_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda")
+    ks_buf[::2] = ks
+    got_indices, got_end = candidate_blocks_to_sparse_indices(
+        cand_buf[:, :num_candidates],
+        ks_buf[::2],
+        ke,
+        cbk,
+        sbk,
+        out=(si_buf[:, : num_candidates * ratio], end_buf[end_offset::2]),
+    )
+    assert torch.equal(got_indices, ref_indices)
+    assert torch.equal(got_end, ref_end)
+    assert torch.equal(
+        cand_buf[:, num_candidates:], torch.full_like(cand_buf[:, num_candidates:], -7)
+    )
+    assert not si_buf[:, num_candidates * ratio :].any()
+    assert not end_buf[1 - end_offset :: 2].any()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_sparse_topk_remap_bounds_nan_rows():
+    """DeepSelect writes an out-of-range sentinel to slot 0 of NaN rows (the
+    rest is left untouched); the remap turns it into -1."""
+    from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+        has_deep_select,
+        sparse_topk_remap,
+    )
+
+    if not has_deep_select():
+        pytest.skip("DeepSelect extension (vllm._deepselect_C) required")
+    rows, num_sparse, sbk, topk = 4, 128, 8, 512
+    width = num_sparse * sbk
+    logits = torch.randn(rows, width, dtype=torch.bfloat16, device="cuda")
+    logits[1:, :16] = float("nan")
+    sparse_indices = torch.arange(num_sparse, dtype=torch.int32, device="cuda")
+    end = torch.full((rows,), width, dtype=torch.int32, device="cuda")
+    out = torch.empty(rows, topk, dtype=torch.int32, device="cuda")
+    cols = torch.zeros_like(out)
+
+    sparse_topk_remap(
+        logits,
+        sparse_indices.repeat(rows, 1),
+        end,
+        torch.zeros_like(end),
+        sbk,
+        topk,
+        out,
+        col_indices=cols,
+    )
+    assert (cols[1:, 0] >= width).all() and (out[1:, 0] == -1).all()
+    assert (out[0] >= 0).all()
+
+
+def _skip_unless_sm100_sparse_kernels():
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("SM100 required")
+    deep_gemm = pytest.importorskip("deep_gemm")
+    if not hasattr(deep_gemm, "fp8_fp4_sparse_mqa_logits"):
+        pytest.skip("DeepGEMM >= 2.8 required for sparse MQA logits")
+    from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+        has_deep_select,
+    )
+
+    if not has_deep_select():
+        pytest.skip("DeepSelect extension (vllm._deepselect_C) required")
+    return deep_gemm
+
+
+def _quant_fp4(x: torch.Tensor):
+    from deep_gemm.utils import per_token_cast_to_fp4
+
+    return per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
+
+
+@pytest.mark.parametrize("capture", [False, True])
+def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool):
+    """The sparse-MQA candidate path selects exactly the same logit values as
+    ``torch.topk`` over dense bf16 logits masked to the candidates (the two
+    DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA
+    graph capture, and a later layer reusing the step's expansion/schedule
+    selects the same values."""
+    deep_gemm = _skip_unless_sm100_sparse_kernels()
+    from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+        sparse_mqa_logits_prefill_chunk,
+    )
+
+    torch.manual_seed(0)
+    rows, num_heads, head_dim = 37, 32, 128
+    cbk, topk, num_candidates = 8, 512, 2048  # V4.1-Flash production values
+    lens = torch.randint(1, 30000, (rows,), dtype=torch.int32)
+    lens[0], lens[1], lens[2] = 1, 5, 2048  # single token, sub-block, long
+    ks = torch.zeros(rows, dtype=torch.int32)
+    ks[1:] = lens.cumsum(0)[:-1]
+    ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda()
+    total = int(lens.sum())
+    q_fp, q_sf = _quant_fp4(
+        torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    )
+    k_fp, k_sf = _quant_fp4(
+        torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16)
+    )
+    weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16)
+    candidates = _make_candidates(lens, num_candidates, cbk).cuda()
+
+    # Reference: dense bf16 logits masked to the candidates, then torch.topk.
+    dense = deep_gemm.fp8_fp4_mqa_logits(
+        (
+            q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2),
+            q_sf.view(rows, num_heads),
+        ),
+        (k_fp.view(torch.int8), k_sf.view(total)),
+        weights,
+        ks,
+        ke,
+        clean_logits=False,
+        logits_dtype=torch.bfloat16,
+    )
+    cols = torch.arange(total, device="cuda")
+    in_range = (cols - ks.unsqueeze(1) >= 0) & (
+        cols - ks.unsqueeze(1) < lens.unsqueeze(1)
+    )
+    masked = dense.masked_fill(
+        ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)),
+        float("-inf"),
+    ).float()
+    ref_abs = _reference_topk(masked, topk)  # packed-workspace positions
+    ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1)
+
+    out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda")
+    scratch = _sparse_scratch(rows, num_candidates, topk)
+
+    def run_sparse(kernel_metadata=None):
+        return sparse_mqa_logits_prefill_chunk(
+            q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2),
+            q_sf.view(rows, num_heads),
+            k_fp.view(torch.int8),
+            k_sf.view(total),
+            weights,
+            ks,
+            ke,
+            candidates,
+            cbk,
+            cbk,
+            topk,
+            out,
+            kernel_metadata=kernel_metadata,
+            **scratch,
+        )
+
+    if capture:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            run_sparse()  # warm up JIT + the DeepGEMM workspace pre-capture
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                run_sparse()
+        out.fill_(-1)
+        graph.replay()
+    else:
+        kernel_metadata = run_sparse()
+        _assert_same_selection(dense, out, ref, row_starts=ks)
+        out.fill_(-1)
+        run_sparse(kernel_metadata)
+
+    _assert_same_selection(dense, out, ref, row_starts=ks)
+
+
+def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int):
+    """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens'
+    fp4 data first, then their packed UE8M0 scales."""
+    device = req_lens.device
+    pages_per = (req_lens + page_kv - 1) // page_kv
+    num_pages, max_pages = int(pages_per.sum()), int(pages_per.max())
+    kv_cache = torch.zeros(
+        num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device
+    )
+    block_table = torch.zeros(
+        len(req_lens), max_pages, dtype=torch.int32, device=device
+    )
+    page_pool = torch.randperm(num_pages, device=device)
+    offset = 0
+    for b in range(len(req_lens)):
+        n = int(pages_per[b])
+        block_table[b, :n] = page_pool[offset : offset + n].to(torch.int32)
+        offset += n
+        k_fp, k_sf = _quant_fp4(
+            torch.randn(int(req_lens[b]), head_dim, device=device, dtype=torch.bfloat16)
+        )
+        rows68 = torch.zeros(n * page_kv, 68, dtype=torch.uint8, device=device)
+        rows68[: int(req_lens[b]), : head_dim // 2] = k_fp.view(torch.uint8)
+        rows68[: int(req_lens[b]), head_dim // 2 :] = k_sf.view(torch.uint8).view(-1, 4)
+        by_page = rows68.view(n, page_kv, 68)
+        split = torch.cat(
+            [by_page[..., :64].reshape(n, -1), by_page[..., 64:].reshape(n, -1)], dim=1
+        ).view(n, page_kv, 1, head_dim // 2 + 4)
+        kv_cache[block_table[b, :n].long()] = split
+    return kv_cache, block_table
+
+
+@pytest.mark.parametrize("next_n", [1, 2])
+@pytest.mark.parametrize("capture", [False, True])
+def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool):
+    """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks
+    from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to
+    one query row each), CUDA graph capture/replay and schedule reuse by a
+    later layer."""
+    deep_gemm = _skip_unless_sm100_sparse_kernels()
+    from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+        sparse_mqa_logits_paged_decode,
+    )
+
+    torch.manual_seed(0)
+    batch, num_heads, head_dim, page_kv = 9, 32, 128, 128
+    cbk, topk, num_candidates = 8, 512, 2048
+    rows = batch * next_n
+    req_lens = torch.randint(1, 30000, (batch,), dtype=torch.int32, device="cuda")
+    req_lens[0], req_lens[1] = 1, 30000  # tiny row + a filtered (subset) row
+    # Row (b, j) sees L_b - next_n + j + 1 tokens (native MTP layout).
+    row_lens = (
+        (req_lens.unsqueeze(1) - next_n + torch.arange(1, next_n + 1, device="cuda"))
+        .clamp(min=0)
+        .to(torch.int32)
+    )
+    kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim)
+    assert kv_cache.stride(0) % 512 == 0  # the paged sparse kernel requires it
+
+    q_fp, q_sf = _quant_fp4(
+        torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    )
+    q = q_fp.view(torch.int8).view(batch, next_n, num_heads, head_dim // 2)
+    q_scale = q_sf.view(batch, next_n, num_heads)
+    weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16)
+    candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda()
+
+    def reference_topk():
+        context_lens_2d = row_lens.view(batch, next_n)
+        sched = deep_gemm.get_paged_mqa_logits_metadata(
+            context_lens_2d, page_kv, deep_gemm.get_num_sms(), None
+        )
+        logits = deep_gemm.fp8_fp4_paged_mqa_logits(
+            (q, q_scale),
+            kv_cache,
+            weights,
+            context_lens_2d,
+            block_table,
+            sched,
+            int(req_lens.max()),
+            False,
+            torch.bfloat16,
+            None,
+        )
+        lens_flat = row_lens.flatten()
+        cols = torch.arange(logits.shape[1], device="cuda")
+        masked = logits.masked_fill(
+            ~(
+                (cols.unsqueeze(0) < lens_flat.unsqueeze(1))
+                & _candidate_token_mask(candidates, cbk, logits.shape[1])
+            ),
+            float("-inf"),
+        ).float()
+        return _reference_topk(masked, topk), logits
+
+    expected, dense_logits = reference_topk()
+    out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda")
+
+    # The metadata builder flattens every query to one row: per-row context
+    # lengths, an expanded block table and a row -> request map.
+    context_lens = row_lens.flatten().contiguous()
+    row_block_table = block_table.repeat_interleave(next_n, dim=0)
+    row_indices = torch.arange(batch, dtype=torch.int32, device="cuda")
+    row_indices = row_indices.repeat_interleave(next_n)
+    row_ks = torch.zeros(rows, dtype=torch.int32, device="cuda")
+    scratch = _sparse_scratch(rows, num_candidates, topk)
+
+    def run_sparse(o: torch.Tensor, kernel_metadata=None):
+        return sparse_mqa_logits_paged_decode(
+            q.view(rows, 1, num_heads, head_dim // 2),
+            q_scale.view(rows, 1, num_heads),
+            kv_cache,
+            weights,
+            context_lens,
+            row_block_table,
+            row_indices,
+            candidates,
+            cbk,
+            cbk,
+            topk,
+            o,
+            row_ks=row_ks,
+            kernel_metadata=kernel_metadata,
+            **scratch,
+        )
+
+    if capture:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            run_sparse(out)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                run_sparse(out)
+        out.fill_(-1)
+        graph.replay()
+    else:
+        kernel_metadata = run_sparse(out)
+        _assert_same_selection(dense_logits, out, expected)
+        out.fill_(-1)
+        run_sparse(out, kernel_metadata)
+
+    _assert_same_selection(dense_logits, out, expected)
+
+
+def test_sparse_mqa_logits_accuracy_vs_torch():
+    """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized
+    MXFP4 inputs."""
+    _skip_unless_sm100_sparse_kernels()
+    from deep_gemm.utils import cast_back_from_fp4
+
+    from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+        candidate_blocks_to_sparse_indices,
+    )
+    from vllm.utils.deep_gemm import (
+        fp8_fp4_sparse_mqa_logits,
+        get_sparse_mqa_logits_metadata,
+    )
+
+    torch.manual_seed(0)
+    rows, num_heads, head_dim = 17, 32, 128
+    sbk, num_candidates = 8, 2048
+    lens = torch.randint(1, 1500, (rows,), dtype=torch.int32)
+    lens[0] = 1
+    ks = torch.zeros(rows, dtype=torch.int32)
+    ks[1:] = lens.cumsum(0)[:-1]
+    ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda()
+    total = int(lens.sum())
+    q_fp, q_sf = _quant_fp4(
+        torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    )
+    k_fp, k_sf = _quant_fp4(
+        torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16)
+    )
+    weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16)
+    candidates = _make_candidates(lens, num_candidates, sbk).cuda()
+
+    sparse_idx, end = candidate_blocks_to_sparse_indices(candidates, ks, ke, sbk, sbk)
+    metadata = get_sparse_mqa_logits_metadata(
+        ks, ke, total, sparse_idx, torch.int8, sbk
+    )
+    logits = fp8_fp4_sparse_mqa_logits(
+        (
+            q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2),
+            q_sf.view(rows, num_heads),
+        ),
+        (k_fp.view(torch.int8), k_sf.view(total)),
+        weights,
+        metadata,
+        sparse_idx.shape[1],
+        sbk,
+    )
+
+    # PyTorch fp32 reference on the dequantized inputs.
+    q_deq = cast_back_from_fp4(q_fp, q_sf, gran_k=32, use_packed_ue8m0=True)
+    k_deq = cast_back_from_fp4(k_fp, k_sf, gran_k=32, use_packed_ue8m0=True)
+    dots = torch.einsum(
+        "rhd,td->rht",
+        q_deq.view(rows, num_heads, head_dim).float(),
+        k_deq.float(),
+    )
+    ref = (dots.relu() * weights.float().unsqueeze(-1)).sum(1)  # [rows, total]
+
+    # Compare on valid sparse columns (absolute packed-workspace positions).
+    width = sparse_idx.shape[1] * sbk
+    cols = torch.arange(width, device="cuda")
+    valid = cols.unsqueeze(0) < end.unsqueeze(1).long()
+    pos = (
+        sparse_idx.repeat_interleave(sbk, dim=1).long() * sbk
+        + (ks % sbk).unsqueeze(1)
+        + (cols % sbk).unsqueeze(0)
+    ).clamp(min=0, max=total - 1)
+    abs_err = (logits.float() - ref.gather(1, pos)).abs()[valid]
+    row_scale = ref.abs().amax(dim=1)
+    norm_err = ((logits.float() - ref.gather(1, pos)).abs() / row_scale.unsqueeze(1))[
+        valid
+    ]
+    print(
+        f"sparse vs torch fp32: max_abs={abs_err.max().item():.4f}; "
+        f"rowmax-normalized p50={norm_err.median().item():.5f} "
+        f"p99={norm_err.quantile(0.99).item():.5f} max={norm_err.max().item():.5f}"
+    )
+    # bf16-level numerics: far below the value gaps that change top-k selection.
+    assert norm_err.quantile(0.99).item() < 1e-2
+    assert norm_err.max().item() < 2e-2
diff --git a/tests/model_executor/layers/test_pooler_methods.py b/tests/model_executor/layers/test_pooler_methods.py
index 4196b99820dc..2abf3d72802c 100644
--- a/tests/model_executor/layers/test_pooler_methods.py
+++ b/tests/model_executor/layers/test_pooler_methods.py
@@ -9,7 +9,7 @@
 import numpy as np
 import pytest
 import torch
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import PoolerConfig, set_current_vllm_config
 from vllm.model_executor.layers.pooler import PoolingParamsUpdate
@@ -315,7 +315,7 @@ def test_dispatch_seq_cls_honors_token_pooling_type(tok_pooling_type):
         model_config=SimpleNamespace(
             pooler_config=pooler_config,
             head_dtype=None,
-            hf_config=PretrainedConfig(num_labels=3),
+            hf_config=PreTrainedConfig(num_labels=3),
         ),
     )
     classifier = torch.nn.Linear(4, 3)
diff --git a/tests/model_executor/layers/test_rocm_unquantized_gemm.py b/tests/model_executor/layers/test_rocm_unquantized_gemm.py
index 66b4af772773..491cfa54350e 100644
--- a/tests/model_executor/layers/test_rocm_unquantized_gemm.py
+++ b/tests/model_executor/layers/test_rocm_unquantized_gemm.py
@@ -17,9 +17,10 @@
 from vllm.model_executor.layers import utils
 
 
-def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
+@pytest.mark.parametrize("m", [1, 128])
+def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch, m):
     x = torch.randn(1, 64, dtype=torch.float16)
-    weight = torch.randn(128, 64, dtype=torch.float16)
+    weight = torch.randn(m, 64, dtype=torch.float16)
 
     monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
     monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
@@ -42,6 +43,20 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
     assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("n", range(1, 6))
+def test_rocm_unquantized_gemm_single_output_real_kernel(dtype, n):
+    torch.manual_seed(0)
+    x = torch.randn(n, 2048, device="cuda", dtype=dtype)
+    weight = torch.randn(1, 2048, device="cuda", dtype=dtype)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
 def test_rocm_unquantized_gemm_makes_skinny_activation_contiguous(monkeypatch):
     x = torch.randn(64, 4, dtype=torch.float16).t()
     weight = torch.randn(128, 64, dtype=torch.float16)
@@ -149,11 +164,12 @@ def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch,
     torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
 
 
-def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch):
+@pytest.mark.parametrize("m", [1, 128])
+def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m):
     # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must
     # fall back to torch.nn.functional.linear.
     x = torch.randn(6, 64, dtype=torch.float16)
-    weight = torch.randn(128, 64, dtype=torch.float16)
+    weight = torch.randn(m, 64, dtype=torch.float16)
 
     monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
     monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
diff --git a/tests/model_executor/model_loader/test_mm_encoder_only_weight_filter.py b/tests/model_executor/model_loader/test_mm_encoder_only_weight_filter.py
new file mode 100644
index 000000000000..210cd28d6f56
--- /dev/null
+++ b/tests/model_executor/model_loader/test_mm_encoder_only_weight_filter.py
@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for --mm-encoder-only safetensors shard filtering."""
+
+import json
+import os
+import tempfile
+
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
+
+from vllm.model_executor.model_loader.weight_utils import (
+    filter_mm_encoder_only_safetensors_files,
+    resolve_mm_encoder_only_lm_prefixes,
+)
+from vllm.model_executor.models.utils import WeightsMapper
+
+
+def _write_index(folder: str, weight_map: dict[str, str]) -> None:
+    with open(os.path.join(folder, SAFE_WEIGHTS_INDEX_NAME), "w") as f:
+        json.dump({"weight_map": weight_map}, f)
+
+
+def _qwen25_mapper() -> WeightsMapper:
+    """Qwen2.5-VL dual layout mapper (nested + flat catch-all)."""
+    return WeightsMapper(
+        orig_to_new_prefix={
+            "model.language_model.": "language_model.model.",
+            "model.visual.": "visual.",
+            "lm_head.": "language_model.lm_head.",
+            "model.": "language_model.model.",
+        }
+    )
+
+
+def _qwen3_mapper() -> WeightsMapper:
+    return WeightsMapper(
+        orig_to_new_prefix={
+            "model.visual.": "visual.",
+            "lm_head.": "language_model.lm_head.",
+            "model.language_model.": "language_model.model.",
+        }
+    )
+
+
+def _molmo_mapper() -> WeightsMapper:
+    return WeightsMapper(
+        orig_to_new_prefix={
+            "model.vision_backbone.": "vision_backbone.",
+            "model.transformer.blocks.": "model.layers.",
+            "model.transformer.ln_f.": "model.norm.",
+            "model.transformer.mlp.down_proj.": "lm_head.",
+        }
+    )
+
+
+def test_keeps_vision_shards_drops_language_only():
+    with tempfile.TemporaryDirectory() as folder:
+        files = [
+            os.path.join(folder, "model-00001-of-000003.safetensors"),
+            os.path.join(folder, "model-00002-of-000003.safetensors"),
+            os.path.join(folder, "model-00003-of-000003.safetensors"),
+        ]
+        for path in files:
+            open(path, "wb").close()
+
+        _write_index(
+            folder,
+            {
+                "language_model.layers.0.weight": "model-00001-of-000003.safetensors",
+                "language_model.layers.1.weight": "model-00002-of-000003.safetensors",
+                "vision_tower.blocks.0.weight": "model-00003-of-000003.safetensors",
+                "mm_projector.weight": "model-00003-of-000003.safetensors",
+            },
+        )
+
+        kept = filter_mm_encoder_only_safetensors_files(
+            files,
+            folder,
+            SAFE_WEIGHTS_INDEX_NAME,
+            ("language_model.",),
+        )
+        assert kept == [files[2]]
+
+
+def test_keeps_mixed_shard_with_vision_and_lm():
+    with tempfile.TemporaryDirectory() as folder:
+        mixed = os.path.join(folder, "model-00001-of-000001.safetensors")
+        open(mixed, "wb").close()
+        _write_index(
+            folder,
+            {
+                "language_model.embed.weight": "model-00001-of-000001.safetensors",
+                "vision_tower.patch.weight": "model-00001-of-000001.safetensors",
+            },
+        )
+
+        kept = filter_mm_encoder_only_safetensors_files(
+            [mixed],
+            folder,
+            SAFE_WEIGHTS_INDEX_NAME,
+            ("language_model.",),
+        )
+        assert kept == [mixed]
+
+
+def test_no_index_returns_unchanged():
+    with tempfile.TemporaryDirectory() as folder:
+        files = [os.path.join(folder, "model.safetensors")]
+        open(files[0], "wb").close()
+        kept = filter_mm_encoder_only_safetensors_files(
+            files,
+            folder,
+            SAFE_WEIGHTS_INDEX_NAME,
+            ("language_model.",),
+        )
+        assert kept == files
+
+
+def test_resolve_returns_module_prefixes_without_hf_hardcode():
+    """No global Qwen HF nest list — just vLLM module prefixes."""
+    prefixes = resolve_mm_encoder_only_lm_prefixes(["language_model"])
+    assert prefixes == ("language_model.",)
+    # With Qwen mapper, still module prefixes (keys classified via mapper).
+    assert resolve_mm_encoder_only_lm_prefixes(
+        ["language_model"], weights_mapper=_qwen25_mapper()
+    ) == ("language_model.",)
+
+
+def test_resolve_empty_language_model_names_fail_closed():
+    """Missing/empty ``_language_model_names`` → None (no language_model. default)."""
+    assert resolve_mm_encoder_only_lm_prefixes(None) is None
+    assert resolve_mm_encoder_only_lm_prefixes([]) is None
+    assert resolve_mm_encoder_only_lm_prefixes(()) is None
+
+
+def test_resolve_skips_filter_for_shared_hf_root():
+    """Shared HF root (Molmo-shaped mapper) → fail-closed, not magic name list."""
+    assert (
+        resolve_mm_encoder_only_lm_prefixes(["model"], weights_mapper=_molmo_mapper())
+        is None
+    )
+    # Without a mapper, identity layout: module prefix is returned as-is.
+    assert resolve_mm_encoder_only_lm_prefixes(["model"]) == ("model.",)
+    assert resolve_mm_encoder_only_lm_prefixes(["model."]) == ("model.",)
+
+
+def test_resolve_keeps_llm_prefix_without_model_broadening():
+    prefixes = resolve_mm_encoder_only_lm_prefixes(["llm"])
+    assert prefixes == ("llm.",)
+
+
+def test_qwen3_nested_index_skips_lm_keeps_visual():
+    with tempfile.TemporaryDirectory() as folder:
+        files = [
+            os.path.join(folder, "model-00001-of-000002.safetensors"),
+            os.path.join(folder, "model-00002-of-000002.safetensors"),
+        ]
+        for path in files:
+            open(path, "wb").close()
+        _write_index(
+            folder,
+            {
+                "model.language_model.layers.0.weight": (
+                    "model-00001-of-000002.safetensors"
+                ),
+                "lm_head.weight": "model-00001-of-000002.safetensors",
+                "model.visual.blocks.0.weight": "model-00002-of-000002.safetensors",
+            },
+        )
+        prefixes = resolve_mm_encoder_only_lm_prefixes(
+            ["language_model"], weights_mapper=_qwen3_mapper()
+        )
+        assert prefixes is not None
+        kept = filter_mm_encoder_only_safetensors_files(
+            files,
+            folder,
+            SAFE_WEIGHTS_INDEX_NAME,
+            prefixes,
+            weights_mapper=_qwen3_mapper(),
+        )
+        assert kept == [files[1]]
+
+
+def test_qwen25_flat_index_skips_lm_keeps_visual():
+    with tempfile.TemporaryDirectory() as folder:
+        files = [
+            os.path.join(folder, "model-00001-of-000002.safetensors"),
+            os.path.join(folder, "model-00002-of-000002.safetensors"),
+        ]
+        for path in files:
+            open(path, "wb").close()
+        _write_index(
+            folder,
+            {
+                "model.layers.0.self_attn.q_proj.weight": (
+                    "model-00001-of-000002.safetensors"
+                ),
+                "model.embed_tokens.weight": "model-00001-of-000002.safetensors",
+                "model.norm.weight": "model-00001-of-000002.safetensors",
+                "lm_head.weight": "model-00001-of-000002.safetensors",
+                "visual.blocks.0.weight": "model-00002-of-000002.safetensors",
+            },
+        )
+        prefixes = resolve_mm_encoder_only_lm_prefixes(
+            ["language_model"], weights_mapper=_qwen25_mapper()
+        )
+        assert prefixes is not None
+        kept = filter_mm_encoder_only_safetensors_files(
+            files,
+            folder,
+            SAFE_WEIGHTS_INDEX_NAME,
+            prefixes,
+            weights_mapper=_qwen25_mapper(),
+        )
+        assert kept == [files[1]]
+
+
+def test_molmo_shaped_index_not_filtered_when_deny_disabled():
+    """Shared-root LM attr → resolve None; callers leave the file list alone."""
+    assert (
+        resolve_mm_encoder_only_lm_prefixes(["model"], weights_mapper=_molmo_mapper())
+        is None
+    )
+    with tempfile.TemporaryDirectory() as folder:
+        files = [
+            os.path.join(folder, "vision.safetensors"),
+            os.path.join(folder, "lm.safetensors"),
+        ]
+        for path in files:
+            open(path, "wb").close()
+        _write_index(
+            folder,
+            {
+                "model.vision_backbone.patch.weight": "vision.safetensors",
+                "model.transformer.blocks.0.weight": "lm.safetensors",
+            },
+        )
+        # Document why raw HF startswith("model.") is unsafe without fail-closed:
+        # whole-shard deny would also drop vision.
+        if_applied_naively = filter_mm_encoder_only_safetensors_files(
+            files, folder, SAFE_WEIGHTS_INDEX_NAME, ("model.",)
+        )
+        assert if_applied_naively == []
diff --git a/tests/model_executor/model_loader/test_registry.py b/tests/model_executor/model_loader/test_registry.py
index 95b797bb5146..a795497a11b6 100644
--- a/tests/model_executor/model_loader/test_registry.py
+++ b/tests/model_executor/model_loader/test_registry.py
@@ -79,7 +79,7 @@ def test_default_loader_hf_still_falls_back_to_pt(tmp_path):
     # Control: load_format="hf" still picks up .pt weights via fallback.
     (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00")
     loader = DefaultModelLoader(LoadConfig(load_format="hf"))
-    _, files, use_safetensors = loader._prepare_weights(
+    _, files, use_safetensors, _ = loader._prepare_weights(
         str(tmp_path),
         None,
         None,
diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py
index f19ec35560f1..155957160fb6 100644
--- a/tests/model_executor/model_loader/test_reload.py
+++ b/tests/model_executor/model_loader/test_reload.py
@@ -1,7 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import gc
+import importlib.machinery
 import inspect
+import sys
+import types
+from types import SimpleNamespace
 from unittest.mock import Mock
 from weakref import WeakKeyDictionary, ref
 
@@ -326,8 +330,8 @@ def _load_marlin_checkpoint_format_weights(layer):
     )
 
 
-def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init):
-    """Marlin must reuse the workspace storage captured by CUDA graphs."""
+def test_marlin_post_load_does_not_own_workspace(monkeypatch, dist_init):
+    """Weight reload must not create layer-owned Marlin lock storage."""
     _stub_marlin_ops(monkeypatch)
     kernel = _make_marlin_kernel()
 
@@ -335,19 +339,17 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i
     _load_marlin_checkpoint_format_weights(layer)
     kernel.process_weights_after_loading(layer)
 
-    workspace_ptr = kernel.workspace.data_ptr()
+    assert not hasattr(kernel, "workspace")
 
     _load_marlin_checkpoint_format_weights(layer)
     kernel.process_weights_after_loading(layer)
 
-    assert kernel.workspace.data_ptr() == workspace_ptr
-    assert torch.all(kernel.workspace == 0)
+    assert not hasattr(kernel, "workspace")
 
 
 @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"])
-def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant):
-    """The Marlin fallback prepare_* functions rerun on weight reload and must
-    reuse the workspace storage whose address captured CUDA graphs hold."""
+def test_marlin_prepare_layer_does_not_own_workspace(monkeypatch, variant):
+    """Weight preparation must not attach runtime workspace to model layers."""
     from vllm import _custom_ops as ops
     from vllm.model_executor.layers.quantization.utils import (
         marlin_utils,
@@ -413,34 +415,57 @@ def load_checkpoint_format_weights():
 
     load_checkpoint_format_weights()
     prepare(layer)
-    workspace_ptr = layer.workspace.data_ptr()
+    assert not hasattr(layer, "workspace")
 
     # Reload: fresh checkpoint-format tensors, prepare runs again
     load_checkpoint_format_weights()
     prepare(layer)
 
-    assert layer.workspace.data_ptr() == workspace_ptr
-    assert torch.all(layer.workspace == 0)
+    assert not hasattr(layer, "workspace")
 
 
-def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch):
-    """An incompatible existing workspace means the address captured by CUDA
-    graphs is already unusable; allocating a replacement would hide that."""
+def test_marlin_workspace_uses_persistent_workspace_manager(monkeypatch):
+    """Calls reuse initialized locks; independent streams get separate storage."""
     from vllm.model_executor.layers.quantization.utils import marlin_utils
+    from vllm.utils import torch_utils
+    from vllm.v1.worker import workspace as workspace_module
 
     monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)
     device = torch.device("cpu")
+    manager = workspace_module.WorkspaceManager(device)
+    monkeypatch.setattr(workspace_module, "_manager", manager)
+    stream = "main"
+    monkeypatch.setattr(torch_utils, "current_stream", lambda: stream)
 
-    workspace = marlin_utils.marlin_make_workspace_new(device)
-    reused = marlin_utils.marlin_make_workspace_new(device, existing=workspace)
-    assert reused is workspace
+    workspace = marlin_utils.get_marlin_workspace(device)
+    assert workspace.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,)
+    assert workspace.dtype == torch.int32
+    assert torch.count_nonzero(workspace) == 0
 
-    with pytest.raises(ValueError, match="incompatible"):
-        marlin_utils.marlin_make_workspace_new(device, 4, existing=workspace)
-    with pytest.raises(ValueError, match="incompatible"):
-        marlin_utils.marlin_make_workspace_new(
-            device, existing=workspace.to(torch.int64)
-        )
+    workspace.fill_(1)
+    assert marlin_utils.get_marlin_workspace(device) is workspace
+    assert torch.all(workspace == 1)
+
+    stream = "aux"
+    aux_workspace = marlin_utils.get_marlin_workspace(device)
+    assert aux_workspace.data_ptr() != workspace.data_ptr()
+    assert torch.count_nonzero(aux_workspace) == 0
+
+    manager.lock()
+    assert marlin_utils.get_marlin_workspace(device) is aux_workspace
+
+
+def test_marlin_workspace_without_manager(monkeypatch):
+    from vllm.model_executor.layers.quantization.utils import marlin_utils
+    from vllm.v1.worker import workspace as workspace_module
+
+    monkeypatch.setattr(workspace_module, "_manager", None)
+    monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)
+    first = marlin_utils.get_marlin_workspace(torch.device("cpu"))
+    second = marlin_utils.get_marlin_workspace(torch.device("cpu"))
+    assert first.data_ptr() != second.data_ptr()
+    assert first.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,)
+    assert torch.count_nonzero(first) == 0
 
 
 def test_model_cleanup(dist_init, default_vllm_config):
@@ -464,6 +489,111 @@ def test_model_cleanup(dist_init, default_vllm_config):
     assert len(mock_info_dict) == 0
 
 
+@pytest.mark.parametrize("is_gated", [False, True])
+@pytest.mark.parametrize("has_bias", [False, True])
+@pytest.mark.parametrize("tp_rank", [0, 1])
+@pytest.mark.parametrize("padded", [False, True])
+def test_padded_moe_reload_releases_each_layer(
+    monkeypatch, is_gated, has_bias, tp_rank, padded
+):
+    """Checkpoint-sized copies finish each layer without global finalization."""
+    from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+    from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+        UnquantizedFusedMoEMethod,
+    )
+
+    hidden, intermediate, experts = 4, 3, 2
+    stored_hidden, stored_intermediate = (8, 8) if padded else (hidden, intermediate)
+    config = SimpleNamespace(
+        hidden_dim_unpadded=hidden,
+        intermediate_size_per_partition_unpadded=intermediate,
+        is_act_and_mul=is_gated,
+        has_bias=has_bias,
+        tp_rank=tp_rank,
+        tp_shard_with_padding=False,
+        moe_parallel_config=SimpleNamespace(tp_size=2),
+    )
+    model = torch.nn.ModuleList()
+    processed: list[torch.nn.Module] = []
+    for _ in range(2):
+        method = object.__new__(UnquantizedFusedMoEMethod)
+        torch.nn.Module.__init__(method)
+        method.moe = config
+        # The regression concerns streaming reload, not kernel conversion.
+        monkeypatch.setattr(method, "process_weights_after_loading", processed.append)
+        layer = object.__new__(RoutedExperts)
+        torch.nn.Module.__init__(layer)
+        layer.moe_config = config
+        layer.quant_config = None
+        layer.quant_method = method
+        layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i)
+        layer._loaded_expert_biases = set()
+        method.create_weights(
+            layer,
+            experts,
+            stored_hidden,
+            stored_intermediate,
+            torch.float32,
+            weight_loader=layer.weight_loader,
+        )
+        model.append(layer)
+
+    record_metadata_for_reloading(model)
+    original_params = [dict(layer.named_parameters()) for layer in model]
+    shards = ["w1", "w3", "w2"] if is_gated else ["w1", "w2"]
+    for cycle in range(2):
+        initialize_layerwise_reload(model)
+        for layer_index, layer in enumerate(model):
+            info = reload_layerwise.get_layerwise_info(layer)
+            inputs = []
+            expected = {
+                name: torch.full_like(p, float("nan"))
+                for name, p in original_params[layer_index].items()
+            }
+            params = dict(layer.named_parameters())
+            calls = [
+                (e, s, b)
+                for e in range(experts)
+                for s in shards
+                for b in ([False, True] if has_bias else [False])
+            ]
+            for call_index, (expert, shard, bias) in enumerate(calls):
+                name = ("w2" if shard == "w2" else "w13") + (
+                    "_bias" if bias else "_weight"
+                )
+                shape = (
+                    ((hidden,) if bias else (hidden, 2 * intermediate))
+                    if shard == "w2"
+                    else ((2 * intermediate,) if bias else (2 * intermediate, hidden))
+                )
+                weight = torch.arange(
+                    torch.Size(shape).numel(), dtype=torch.float32
+                ).reshape(shape)
+                weight = weight + 100 * (1 + call_index + cycle)
+                inputs.append(ref(weight))
+                # Direct checkpoint loading is the reference for deferred reload.
+                layer.weight_loader(expected[name], weight, name, shard, expert)
+                param = params[name]
+                param.weight_loader(param, weight, name, shard, expert)
+                del weight
+                if call_index != len(calls) - 1:
+                    assert info.can_load(), (
+                        "Layer processed before its final checkpoint shard"
+                    )
+
+            assert not info.can_load(), (
+                "Padding must not defer the layer until finalization"
+            )
+            assert not info.loaded_weights
+            assert len(processed) == cycle * len(model) + layer_index + 1
+            assert all(source() is None for source in inputs)
+            for name, original in original_params[layer_index].items():
+                assert getattr(layer, name) is original
+                # Kernel-specific tests cover padding, which is not checkpoint data.
+                mask = torch.isfinite(expected[name])
+                assert torch.equal(original[mask], expected[name][mask])
+
+
 def test_get_numel_loaded():
     param = torch.empty(10, device="meta")
     loaded_weight = torch.empty(10)
@@ -817,6 +947,110 @@ def materialize_with_sentinel(meta_tensor):
     assert "0.scale" not in model.state_dict()
 
 
+@pytest.fixture
+def hpc_rope_norm(monkeypatch, default_vllm_config):
+    """Import HpcRopeNorm with the external ``hpc`` package stubbed out."""
+    if "hpc" not in sys.modules:
+        stub = types.ModuleType("hpc")
+        stub.__spec__ = importlib.machinery.ModuleSpec("hpc", loader=None)
+        stub.QuantType = types.SimpleNamespace(  # type: ignore[attr-defined]
+            QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR=types.SimpleNamespace(value=0)
+        )
+        monkeypatch.setitem(sys.modules, "hpc", stub)
+    from vllm.model_executor.layers.hpc import rope_norm
+
+    monkeypatch.setattr(rope_norm, "_hpc_rope_norm_instances", {})
+    return rope_norm
+
+
+def test_hpc_rope_norm_kernel_sees_refit_norm_weights(monkeypatch, hpc_rope_norm):
+    """The fused HPC kernel is handed the live QK-norm weights after a refit.
+
+    Drives the production ``_forward_impl`` with a recording ``hpc`` stub. The
+    Q/K norm weights it receives must be the model's own float32 parameters,
+    so a layerwise reload that rewrites them in place (same storage) is what
+    the kernel sees. Previously the kernel read separate mirrors that no
+    reload path refreshed.
+    """
+    from vllm.model_executor.layers.layernorm import RMSNorm
+
+    head_dim, num_heads, num_kv_heads, block_size = 128, 8, 1, 4
+    layer = torch.nn.Module()
+    layer.q_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32)
+    layer.k_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32)
+    layer.hpc_rope_norm = hpc_rope_norm.HpcRopeNorm(
+        num_heads=num_heads,
+        num_kv_heads=num_kv_heads,
+        head_dim=head_dim,
+        cos_sin_cache=torch.ones(16, head_dim),
+        use_qk_norm=True,
+        fallback_qnorm=layer.q_norm,
+        fallback_knorm=layer.k_norm,
+        kv_cache_dtype="auto",
+        layer_name="hpc_test_layer",
+    )
+    model = torch.nn.Sequential(layer)
+    rnorm = layer.hpc_rope_norm
+
+    calls: list[dict] = []
+
+    def record(*args, **kwargs):
+        calls.append(kwargs)
+
+    monkeypatch.setattr(sys.modules["hpc"], "rope_norm_store_kv", record, raising=False)
+
+    def kernel_norm_weights():
+        q_size, kv_size = num_heads * head_dim, num_kv_heads * head_dim
+        qkv = torch.zeros(1, q_size + 2 * kv_size, dtype=torch.bfloat16)
+        kv_cache = torch.zeros(
+            2, num_kv_heads, block_size, 2 * head_dim, dtype=torch.bfloat16
+        )
+        attn_layer = types.SimpleNamespace(
+            _k_scale=torch.ones(1), _v_scale=torch.ones(1)
+        )
+        attn_metadata = types.SimpleNamespace(
+            num_actual_tokens=1,
+            num_decodes=1,
+            num_decode_tokens=1,
+            num_prefills=0,
+            num_prefill_tokens=0,
+            max_query_len=1,
+            decode_query_len=1,
+            qo_indptr=None,
+            qo_indptr_decode=None,
+            slot_mapping=torch.tensor([4]),
+            seq_lens=torch.tensor([1]),
+            block_table_tensor=torch.tensor([[1]]),
+            hpc_kv_written=False,
+        )
+        output = torch.zeros(1, q_size, dtype=torch.bfloat16)
+        rnorm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output)
+        return calls[-1]["q_norm_weight"], calls[-1]["k_norm_weight"]
+
+    def loaded(value):
+        return torch.full((head_dim,), value, dtype=torch.bfloat16)
+
+    default_weight_loader(layer.q_norm.weight, loaded(0.5))
+    default_weight_loader(layer.k_norm.weight, loaded(0.25))
+    q, k = kernel_norm_weights()
+    assert torch.equal(q, loaded(0.5).float())
+    assert torch.equal(k, loaded(0.25).float())
+    q_ptr, k_ptr = q.data_ptr(), k.data_ptr()
+
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    layer.q_norm.weight.weight_loader(layer.q_norm.weight, loaded(2.0))
+    layer.k_norm.weight.weight_loader(layer.k_norm.weight, loaded(3.0))
+    finalize_layerwise_reload(model, model_config=None)
+
+    q, k = kernel_norm_weights()
+    assert q.dtype == k.dtype == torch.float32
+    assert torch.equal(q, loaded(2.0).float())
+    assert torch.equal(k, loaded(3.0).float())
+    assert (q.data_ptr(), k.data_ptr()) == (q_ptr, k_ptr)
+    assert not hasattr(rnorm, "qnorm_weight")
+
+
 @pytest.mark.parametrize(
     "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
 )
diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py
index fa0b6b418343..7b60aaf04b26 100644
--- a/tests/model_executor/model_loader/test_weight_cache.py
+++ b/tests/model_executor/model_loader/test_weight_cache.py
@@ -7,6 +7,7 @@
 identical outputs.
 """
 
+import glob
 import os
 import shutil
 import subprocess
@@ -21,7 +22,6 @@
 
 from vllm import SamplingParams
 from vllm.assets.image import ImageAsset
-from vllm.model_executor.model_loader.weight_cache.protocol import get_socket_path
 from vllm.platforms import current_platform
 
 DAEMON_TIMEOUT_S = 600
@@ -30,10 +30,19 @@
 class WeightCacheDaemon:
     """Context manager running the real weight cache daemon as a subprocess."""
 
-    def __init__(self, model: str, tp_size: int, extra_args: list[str] | None = None):
+    def __init__(
+        self,
+        model: str,
+        tp_size: int,
+        extra_args: list[str] | None = None,
+        num_groups: int = 1,
+    ):
         # Short base path: Unix socket paths are limited to ~107 characters.
         self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_")
         self.tp_size = tp_size
+        # Each daemon group (target plus cached drafts) binds one socket per
+        # local rank.
+        self.num_sockets = tp_size * num_groups
         self._cmd = [
             sys.executable,
             "-m",
@@ -79,9 +88,7 @@ def _wait_ready(self, timeout_s: float) -> None:
         # Poll for the socket files rather than a log line: model loading can
         # pull in JIT compilers that swap the process's stderr and swallow
         # everything logged afterwards, making log-based readiness flaky.
-        expected = [
-            get_socket_path(gpu_id, self.socket_dir) for gpu_id in range(self.tp_size)
-        ]
+        pattern = os.path.join(self.socket_dir, "vllm_weight_cache_*.sock")
         deadline = time.monotonic() + timeout_s
         while time.monotonic() < deadline:
             if self._proc.poll() is not None:
@@ -89,7 +96,7 @@ def _wait_ready(self, timeout_s: float) -> None:
                     f"Weight cache daemon exited with {self._proc.returncode}:\n"
                     f"{self._logs()}"
                 )
-            if all(os.path.exists(path) for path in expected):
+            if len(glob.glob(pattern)) >= self.num_sockets:
                 return
             time.sleep(1.0)
         raise TimeoutError(
@@ -114,6 +121,9 @@ class ModelCase:
     images: list | None = None
     llm_kwargs: dict[str, Any] = field(default_factory=dict)
     daemon_args: list[str] = field(default_factory=list)
+    # Daemon groups the launcher starts; 2 when a cached draft group joins the
+    # target group (MTP/EAGLE speculative decoding).
+    daemon_groups: int = 1
 
 
 def generate(
@@ -174,14 +184,41 @@ def generate(
     daemon_args=["--trust-remote-code"],
 )
 
+# Qwen3.5-0.8B ships one MTP layer in the target checkpoint, so method="mtp"
+# loads the draft from the same model. The daemon must cache it in its draft
+# group for the warm runs (fallback=False) to succeed.
+QWEN_MTP_CASE = ModelCase(
+    model="Qwen/Qwen3.5-0.8B",
+    prompts=[
+        "Hello, my name is",
+        "The capital of France is",
+    ],
+    llm_kwargs=dict(
+        gpu_memory_utilization=0.3,
+        enforce_eager=True,
+        enable_chunked_prefill=True,
+        speculative_config={"method": "mtp", "num_speculative_tokens": 1},
+    ),
+    daemon_args=[
+        "--speculative-config",
+        '{"method": "mtp", "num_speculative_tokens": 1}',
+    ],
+    daemon_groups=2,
+)
+
 
-@pytest.mark.parametrize("case", [QWEN_CASE, K3_CASE], ids=["qwen3.5", "kimi-k3"])
+@pytest.mark.parametrize(
+    "case",
+    [QWEN_CASE, K3_CASE, QWEN_MTP_CASE],
+    ids=["qwen3.5", "kimi-k3", "qwen3.5-mtp"],
+)
 def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
     """Cold start falls back to disk; warm restarts load weights via CUDA IPC.
 
     All runs must produce outputs identical to a default-loader baseline. The
     warm runs disable the disk fallback, so they only pass if the weights
-    really came from the daemon.
+    really came from the daemon — for the MTP case, both the target's and the
+    draft's daemon groups.
     """
     if not current_platform.is_cuda_alike():
         pytest.skip("Weight cache IPC sharing requires CUDA or ROCm")
@@ -201,7 +238,12 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
     with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir:
         cold_outputs = generate(vllm_runner, case, empty_socket_dir, fallback=True)
 
-    with WeightCacheDaemon(case.model, tp_size=1, extra_args=case.daemon_args) as d:
+    with WeightCacheDaemon(
+        case.model,
+        tp_size=1,
+        extra_args=case.daemon_args,
+        num_groups=case.daemon_groups,
+    ) as d:
         warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
         # Warm restart: a second engine lifetime against the same daemon.
         restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
@@ -209,3 +251,116 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
     assert cold_outputs == baseline_outputs
     assert warm_outputs == baseline_outputs
     assert restart_outputs == baseline_outputs
+
+
+def _parallel(**kw):
+    from types import SimpleNamespace
+
+    base = dict(
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+        data_parallel_size=1,
+        data_parallel_size_local=1,
+        data_parallel_rank=0,
+        nnodes=1,
+        node_rank=0,
+    )
+    base.update(kw)
+    return SimpleNamespace(**base)
+
+
+def test_daemon_places_tp_and_dp_ranks_on_local_gpus():
+    """Each node serves a contiguous global-rank block (DP-major), so without
+    DP local GPU i is TP rank r*local+i, and single-node DP launchers offset
+    by --data-parallel-start-rank."""
+    from vllm.model_executor.model_loader.weight_cache.daemon import plan_local_ranks
+
+    tp_second_node = _parallel(tensor_parallel_size=8, nnodes=2, node_rank=1)
+    assert plan_local_ranks(tp_second_node) == [(i, 0, 4 + i) for i in range(4)]
+
+    dp_third_node = _parallel(
+        data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=8
+    )
+    assert plan_local_ranks(dp_third_node) == [(i, 8 + i, 0) for i in range(4)]
+
+    dp_tp = _parallel(
+        tensor_parallel_size=2,
+        data_parallel_size=4,
+        data_parallel_size_local=2,
+        data_parallel_rank=2,
+    )
+    assert plan_local_ranks(dp_tp) == [(0, 2, 0), (1, 2, 1), (2, 3, 0), (3, 3, 1)]
+
+    # DP + nnodes: node-local TP replicas
+    dp_tp_node0 = _parallel(tensor_parallel_size=8, data_parallel_size=2, nnodes=2)
+    assert plan_local_ranks(dp_tp_node0) == [(i, 0, i) for i in range(8)]
+    dp_tp_node1 = _parallel(
+        tensor_parallel_size=8, data_parallel_size=2, nnodes=2, node_rank=1
+    )
+    assert plan_local_ranks(dp_tp_node1) == [(i, 1, i) for i in range(8)]
+
+    # DP + nnodes: TP group spans nodes (TP8 x DP2 on 4 nodes)
+    tp_span_node1 = _parallel(
+        tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=1
+    )
+    assert plan_local_ranks(tp_span_node1) == [(i, 0, 4 + i) for i in range(4)]
+    tp_span_node2 = _parallel(
+        tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=2
+    )
+    assert plan_local_ranks(tp_span_node2) == [(i, 1, i) for i in range(4)]
+
+    # DP + nnodes: several DP replicas per node (TP4 x DP4 on 2 nodes)
+    dp_multi_node0 = _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2)
+    assert plan_local_ranks(dp_multi_node0) == [(i, i // 4, i % 4) for i in range(8)]
+
+
+def test_daemon_rejects_unmappable_parallelism():
+    from vllm.model_executor.model_loader.weight_cache.daemon import (
+        _reject_unsupported_parallelism,
+    )
+
+    _reject_unsupported_parallelism(
+        _parallel(
+            data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=12
+        )
+    )
+    # DP combines with --nnodes when the world size divides evenly.
+    _reject_unsupported_parallelism(
+        _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2)
+    )
+    with pytest.raises(ValueError, match="pipeline"):
+        _reject_unsupported_parallelism(_parallel(pipeline_parallel_size=2))
+    with pytest.raises(ValueError, match="evenly divide"):
+        _reject_unsupported_parallelism(_parallel(tensor_parallel_size=3, nnodes=2))
+    with pytest.raises(ValueError, match="evenly divide"):
+        _reject_unsupported_parallelism(
+            _parallel(tensor_parallel_size=2, data_parallel_size=3, nnodes=4)
+        )
+    with pytest.raises(ValueError, match="exceeds"):
+        _reject_unsupported_parallelism(
+            _parallel(
+                data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=13
+            )
+        )
+
+
+def test_weight_cache_key_distinguishes_dp_ranks():
+    from dataclasses import replace
+
+    from vllm.model_executor.model_loader.weight_cache.protocol import WeightCacheKey
+
+    key = WeightCacheKey(
+        checkpoint="ckpt",
+        model_arch="Arch",
+        tp_size=1,
+        tp_rank=0,
+        dtype="bf16",
+        quantization=None,
+        quant_config_hash="h",
+        revision=None,
+        vllm_version="v",
+        dp_size=16,
+        dp_rank=3,
+    )
+    assert key.mismatched_fields(replace(key, dp_rank=4)) == ["dp_rank"]
+    assert key.mismatched_fields(replace(key, dp_size=8, dp_rank=3)) == ["dp_size"]
diff --git a/tests/model_executor/test_cpu_unquantized_gemm_dispatch.py b/tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
index f18780cf6b57..014a9ca25dfe 100644
--- a/tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
+++ b/tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
@@ -7,6 +7,7 @@
 
 from vllm.model_executor.layers import utils
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_default_torch_dtype
 
 
 @pytest.fixture(scope="module")
@@ -89,3 +90,25 @@ def test_dispatch_cpu_unquantized_gemm_logs_zentorch_dispatch(monkeypatch):
             expected_prepacked,
         )
     ]
+
+
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
+@pytest.mark.parametrize("weight_dtype", [torch.bfloat16, torch.float16, torch.float32])
+def test_dispatch_cpu_unquantized_gemm_remove_weight_keeps_dtype(
+    monkeypatch, weight_dtype
+):
+    monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
+
+    layer = torch.nn.Linear(16, 8, bias=False, dtype=weight_dtype)
+    loading_dtype = (
+        torch.float32 if weight_dtype is not torch.float32 else torch.bfloat16
+    )
+    with set_default_torch_dtype(loading_dtype):
+        utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
+
+    assert layer.weight.numel() == 0
+    assert layer.weight.dtype is weight_dtype
+
+    x = torch.randn(4, 16, dtype=weight_dtype)
+    output = layer.cpu_linear(x, layer.weight, None)
+    assert not output.isnan().any()
diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py
new file mode 100644
index 000000000000..e71061e73317
--- /dev/null
+++ b/tests/model_executor/test_deep_gemm_warmup.py
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import (
+    DeepGemmFp8BlockScaledMMKernel,
+)
+from vllm.model_executor.warmup import deep_gemm_warmup
+
+
+def _block_fp8_layer(n: int, k: int, scale_name: str = "weight_scale"):
+    layer = SimpleNamespace(
+        weight=torch.empty((n, k), dtype=torch.float8_e4m3fn),
+        weight_block_size=[128, 128],
+        deep_gemm_warmup_provider=object.__new__(DeepGemmFp8BlockScaledMMKernel),
+    )
+    setattr(layer, scale_name, torch.empty((n // 128, k // 128), dtype=torch.float32))
+    return layer
+
+
+def _run_warmup(monkeypatch, layers) -> list[dict]:
+    calls: list[dict] = []
+    monkeypatch.setattr(
+        deep_gemm_warmup,
+        "get_mk_alignment_for_contiguous_layout",
+        lambda: [128, 128],
+    )
+    monkeypatch.setattr(
+        deep_gemm_warmup,
+        "_deepgemm_fp8_gemm_nt_warmup",
+        lambda **kwargs: calls.append(kwargs),
+    )
+    model = SimpleNamespace(modules=lambda: iter(layers))
+    deep_gemm_warmup.deepgemm_fp8_gemm_nt_warmup(model, max_tokens=16)
+    return calls
+
+
+@pytest.mark.parametrize("scale_name", ["weight_scale", "weight_scale_inv"])
+def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> None:
+    """Any layer stamped by the DeepGEMM kernel is warmed, regardless of the
+    quantization method that owns it or the name of its scale parameter."""
+    layer = _block_fp8_layer(256, 128, scale_name)
+
+    calls = _run_warmup(monkeypatch, [layer])
+
+    assert len(calls) == 1
+    assert calls[0]["w"] is layer.weight
+    assert calls[0]["ws"] is getattr(layer, scale_name)
+
+
+def test_n_multiple_of_64_matches_kernel_selection(monkeypatch) -> None:
+    """DeepGEMM accepts N % 64 == 0 (e.g. DeepSeek kv_a_proj_with_mqa, N=576),
+    so warmup must not require N % 128 == 0."""
+    assert len(_run_warmup(monkeypatch, [_block_fp8_layer(576, 256)])) == 1
+
+
+def test_unstamped_and_mismatched_block_layers_are_skipped(monkeypatch) -> None:
+    mismatched = _block_fp8_layer(256, 128)
+    mismatched.weight_block_size = [1, 32]
+    layers = [SimpleNamespace(), mismatched, _block_fp8_layer(256, 128)]
+
+    calls = _run_warmup(monkeypatch, layers)
+
+    assert len(calls) == 1
+    assert calls[0]["w"] is layers[-1].weight
+
+
+@pytest.mark.parametrize("is_bmm", [False, True])
+def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None:
+    """Bmm layers bypass the kernel at runtime, so they must not be warmed."""
+    kernel = object.__new__(DeepGemmFp8BlockScaledMMKernel)
+    kernel.is_deep_gemm_supported = False
+    layer = torch.nn.Module()
+    layer.weight = torch.nn.Parameter(
+        torch.empty((256, 128), dtype=torch.float8_e4m3fn), requires_grad=False
+    )
+    layer.weight_scale = torch.nn.Parameter(
+        torch.empty((2, 1), dtype=torch.float32), requires_grad=False
+    )
+    layer.weight_block_size = [128, 128]
+    layer.is_bmm = is_bmm
+
+    kernel.process_weights_after_loading(layer)
+
+    provider = getattr(layer, "deep_gemm_warmup_provider", None)
+    assert provider is (None if is_bmm else kernel)
diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py
index a8ecd1276e16..dd162629b8f4 100644
--- a/tests/model_executor/test_flashinfer_autotune_warmup.py
+++ b/tests/model_executor/test_flashinfer_autotune_warmup.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from inspect import signature
 from types import SimpleNamespace
 from typing import Any
 from unittest.mock import Mock, call, patch
@@ -76,8 +77,16 @@ def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated():
     assert token_counts == (32,)
 
 
-def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
+@pytest.mark.parametrize("skip_attn", [False, True])
+def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn):
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2Runner
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner as V1Runner
+
     runner = _make_runner([])
+    dummy_run_signature = signature((V2Runner if skip_attn else V1Runner)._dummy_run)
+    runner._dummy_run.side_effect = lambda **kwargs: dummy_run_signature.bind(
+        None, **kwargs
+    )
     max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192)
     deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128)
 
@@ -93,7 +102,7 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
         ) as get_buckets,
         patch("vllm.utils.flashinfer.autotune") as autotune,
     ):
-        _run_flashinfer_autotune_dummy_runs(runner)
+        _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn)
 
     assert get_buckets.call_args_list == [call(8192), call(128)]
     assert autotune.call_args_list == [
@@ -106,11 +115,13 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run():
             skip_eplb=True,
             is_profile=True,
             randomize_inputs=True,
+            **({"skip_attn": True} if skip_attn else {}),
         ),
         call(
             num_tokens=128,
             skip_eplb=True,
             is_profile=True,
             randomize_inputs=True,
+            **({"skip_attn": True} if skip_attn else {}),
         ),
     ]
diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py
index 62ecb05f6b78..7d38f4b40d42 100644
--- a/tests/model_executor/test_jit_warmup.py
+++ b/tests/model_executor/test_jit_warmup.py
@@ -171,6 +171,19 @@ def test_warmup_range_validates_custom_advancement() -> None:
         )
 
 
+def test_warmup_cases_support_lambda_advancement() -> None:
+    def warmup_inputs() -> dict[str, Any]:
+        value: Any = WarmupIntRange(1, 9, advance=lambda value: value * 2)
+        return dict(value=value)
+
+    assert list(ToyKernel()._expand_warmup_cases(warmup_inputs)) == [
+        {"value": 1},
+        {"value": 2},
+        {"value": 4},
+        {"value": 8},
+    ]
+
+
 def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
     cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
 
@@ -577,6 +590,25 @@ def test_registry_records_only_inside_model_setup_context() -> None:
     assert kernel.compiled == []
 
 
+def test_registry_capture() -> None:
+    class Owner:
+        @JitWarmupRegistry.capture
+        def __init__(
+            self,
+            vllm_config: Any,
+            kernel: RecordingToyKernel,
+        ) -> None:
+            kernel.register_warmup(3, vllm_config)
+
+    kernel = RecordingToyKernel()
+    config = _config()
+    owner = Owner(config, kernel)
+
+    assert len(owner.jit_warmup_registry) == 1  # type: ignore[attr-defined]
+    kernel.register_warmup(5, config)
+    assert len(owner.jit_warmup_registry) == 1  # type: ignore[attr-defined]
+
+
 def test_registry_expands_requests_and_deduplicates_owner_keys() -> None:
     registry = JitWarmupRegistry(_config())
     kernel = RecordingToyKernel()
diff --git a/tests/model_executor/test_jit_warmup_cutedsl_launcher.py b/tests/model_executor/test_jit_warmup_cutedsl_launcher.py
index 2524c60d2e37..76ec4d1200b1 100644
--- a/tests/model_executor/test_jit_warmup_cutedsl_launcher.py
+++ b/tests/model_executor/test_jit_warmup_cutedsl_launcher.py
@@ -5,10 +5,10 @@
 from typing import Any
 
 from vllm.model_executor.warmup import jit_warmup_cutedsl_helper
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
     CuTeDSLLaunchSpec,
     VllmCuTeDSLJitKernel,
-    cutedsl_kernel_launcher,
 )
 
 
@@ -30,7 +30,7 @@ def get_warmup_keys(self) -> list[CompileKey]:
     def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
         return (f"fake-{compile_key.variant}",)
 
-    @cutedsl_kernel_launcher
+    @kernel_launcher
     def __call__(
         self,
         payload: str,
@@ -38,7 +38,7 @@ def __call__(
         variant: int,
     ) -> CuTeDSLLaunchSpec[CompileKey]:
         compile_key = self.dispatch(variant=variant)
-        return compile_key, (payload,), {"variant": variant}
+        return compile_key, (payload,)
 
 
 def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None:
@@ -74,3 +74,38 @@ def executor(*runtime_args: Any) -> tuple[int, tuple[Any, ...]]:
         (("entry", 2), ("fake-2",)),
     ]
     assert launch_calls == [(1, ("runtime",)), (2, ("other",))]
+
+
+def test_cutedsl_launcher_runs_post_launch_finalizer(monkeypatch) -> None:
+    events: list[str] = []
+
+    def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any:
+        def executor(*runtime_args: Any) -> None:
+            events.append("launch")
+
+        return executor
+
+    class FinalizedKernel(_TestCuTeDSLKernel):
+        @kernel_launcher
+        def __call__(
+            self,
+            payload: str,
+            *,
+            variant: int,
+        ) -> CuTeDSLLaunchSpec[_TestCuTeDSLKernel.CompileKey]:
+            compile_key = self.dispatch(variant=variant)
+
+            def finalize() -> str:
+                events.append("finalize")
+                return "output"
+
+            return compile_key, (payload,), "unused", finalize
+
+    monkeypatch.setattr(
+        jit_warmup_cutedsl_helper,
+        "compile_cutedsl",
+        compile_cutedsl,
+    )
+
+    assert FinalizedKernel()("runtime", variant=1) == "output"
+    assert events == ["launch", "finalize"]
diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py
index 2e624cc7f614..725a383833a6 100644
--- a/tests/model_executor/test_jit_warmup_triton_launcher.py
+++ b/tests/model_executor/test_jit_warmup_triton_launcher.py
@@ -4,22 +4,69 @@
 from dataclasses import dataclass
 from typing import Any
 
+import pytest
+import torch
+
+from vllm.model_executor.warmup import jit_warmup_triton_helper
+from vllm.model_executor.warmup.jit_warmup import (
+    WarmupChoices,
+    WarmupIntRange,
+    kernel_launcher,
+)
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
+    TritonJitKey,
+    TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
+    triton_kernel_dispatcher_with_warmup,
+    triton_warmup_inputs,
 )
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _binder_test_kernel(x, value, BLOCK: tl.constexpr):
+    pass
+
+
+@triton.jit
+def _second_binder_test_kernel(x, value, BLOCK: tl.constexpr):
+    pass
+
+
+@triton.jit
+def _pointer_group_test_kernel(x_ptr, value, BLOCK: tl.constexpr):
+    pass
 
 
 class _FakeTritonKernel:
-    arg_names = ("first", "second", "CONST")
+    arg_names: tuple[str, ...] = ("first", "second", "CONST")
 
     def __init__(self) -> None:
         self.warmup_calls: list[dict[str, Any]] = []
+        self.runtime_calls: list[tuple[Any, tuple[Any, ...], dict[str, Any]]] = []
 
     def warmup(self, **kwargs: Any) -> None:
         self.warmup_calls.append(kwargs)
 
+    def __getitem__(self, grid: Any) -> Any:
+        def launch(*args: Any, **kwargs: Any) -> None:
+            self.runtime_calls.append((grid, args, kwargs))
+
+        return launch
+
+
+def _patch_key_deriver(
+    monkeypatch: pytest.MonkeyPatch,
+    derive: Any,
+) -> None:
+    monkeypatch.setattr(
+        jit_warmup_triton_helper,
+        "_triton_key_deriver",
+        lambda kernel: lambda kwargs: derive(kernel, kwargs),
+    )
+
 
 class _TestTritonKernel(VllmTritonJitKernel["_TestTritonKernel.CompileKey"]):
     kernel = _FakeTritonKernel()
@@ -91,6 +138,312 @@ class TestCpuKernel(_TestTritonKernel):
     assert calls == [("runtime", 2, 7)]
 
 
+def test_triton_kernel_decorator_returns_launcher(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+    kernel.arg_names = (
+        "FIRST_PTR",
+        "FIRST_STRIDE0",
+        "aliased_ptr",
+        "aliased_stride_0",
+        "SECOND",
+        "CONST",
+    )
+
+    def warmup_inputs() -> dict[str, Any]:
+        return dict(
+            first=TritonWarmupTensor(torch.float32, shape=(2, 3), strides=(5, 1)),
+            second=WarmupChoices(1, 2),
+        )
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def launch(first: str, second: int = 2, config: int = 7) -> LaunchSpec:
+        return (2,), dict(aliased_ptr=first, CONST=config)
+
+    def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]:
+        return {TritonJitKey(id(kernel), "fake", 0, kwargs["SECOND"])}
+
+    _patch_key_deriver(monkeypatch, fake_keys)
+
+    keys = launch.get_warmup_keys()
+    keys_with_config = launch.get_warmup_keys(vllm_config=object())
+    assert keys_with_config == keys
+    assert [dict(key.inputs)["second"] for key in keys] == [1, 2]
+    launch.compile(keys[0])
+    assert kernel.warmup_calls == [
+        {
+            "grid": (1,),
+            "FIRST_PTR": TritonWarmupTensor(
+                torch.float32, shape=(2, 3), strides=(5, 1)
+            ),
+            "FIRST_STRIDE0": 5,
+            "aliased_ptr": TritonWarmupTensor(
+                torch.float32, shape=(2, 3), strides=(5, 1)
+            ),
+            "aliased_stride_0": 5,
+            "SECOND": 1,
+            "CONST": 7,
+        }
+    ]
+    first = torch.empty_strided((2, 3), (5, 1))
+    launch(first)
+    assert kernel.runtime_calls == [
+        (
+            (2,),
+            (),
+            {
+                "FIRST_PTR": first,
+                "FIRST_STRIDE0": 5,
+                "aliased_ptr": first,
+                "aliased_stride_0": 5,
+                "SECOND": 2,
+                "CONST": 7,
+            },
+        )
+    ]
+    assert launch.__name__ == "launch"
+    with pytest.raises(TypeError, match="unexpected keyword"):
+        launch(first, 1, 7, stale_constexpr=True)
+
+
+def test_triton_kernel_decorator_without_triton(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(jit_warmup_triton_helper, "HAS_TRITON", False)
+    monkeypatch.setattr(jit_warmup_triton_helper, "triton", object())
+    kernel = _FakeTritonKernel()
+
+    def warmup_inputs() -> dict[str, Any]:
+        return dict(first="warmup", second=1)
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def launch(first: str, second: int) -> LaunchSpec:
+        return (2,), dict(CONST=7)
+
+    launch("runtime", 2)
+    assert kernel.runtime_calls == [
+        ((2,), (), {"first": "runtime", "second": 2, "CONST": 7})
+    ]
+
+
+def test_triton_kernel_decorator_exhausts_large_ranges_before_deduplication(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+    dispatched: list[int] = []
+
+    def warmup_inputs() -> dict[str, Any]:
+        tokens: Any = WarmupIntRange(1, 8193)
+        return dict(first="warmup", second=tokens)
+
+    def specialization(second: int) -> int:
+        return 7 if second <= 97 else 8
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def dispatch(first: str, second: int) -> LaunchSpec:
+        dispatched.append(second)
+        return (second,), dict(CONST=specialization(second))
+
+    def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]:
+        return {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])}
+
+    _patch_key_deriver(monkeypatch, fake_keys)
+
+    keys = dispatch.get_warmup_keys()
+    assert dispatched == list(range(1, 8193))
+    assert len(keys) == 2
+    assert {specialization(dict(key.inputs)["second"]) for key in keys} == {7, 8}
+
+
+def test_triton_kernel_decorator_propagates_dispatch_assertions(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+
+    def warmup_inputs() -> dict[str, Any]:
+        return dict(first="warmup", second=WarmupChoices(1, 2))
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def dispatch(first: str, second: int) -> LaunchSpec:
+        assert second != 2, "broken dispatch"
+        return (1,), dict(CONST=second)
+
+    _patch_key_deriver(
+        monkeypatch,
+        lambda kernel, kwargs: {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])},
+    )
+
+    with pytest.raises(AssertionError, match="broken dispatch"):
+        dispatch.get_warmup_keys()
+
+
+def test_triton_kernel_dispatch_uses_device_fake_tensors(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+    device_type = current_platform.device_type
+
+    def warmup_inputs() -> dict[str, Any]:
+        return dict(
+            first=TritonWarmupTensor(torch.float32, shape=(2, 3)),
+            second=1,
+        )
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def dispatch(first: torch.Tensor, second: int) -> LaunchSpec:
+        assert isinstance(first, torch.Tensor)
+        assert first.device.type == device_type
+        assert first[0].is_contiguous()
+        return (first.shape[0],), dict(CONST=first[0].numel())
+
+    _patch_key_deriver(
+        monkeypatch,
+        lambda kernel, kwargs: {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])},
+    )
+
+    keys = dispatch.get_warmup_keys()
+    assert len(keys) == 1
+    assert dict(keys[0].inputs)["first"].device.type == device_type
+
+
+def test_triton_kernel_decorates_native_launchers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+
+    def warmup_inputs(kernel: Any = kernel) -> dict[str, Any]:
+        return triton_warmup_inputs(
+            kernel,
+            "warmup",
+            WarmupChoices(1, 2),
+            grid=(2,),
+            CONST=7,
+        )
+
+    launcher = triton_kernel_dispatcher_with_warmup(warmup_inputs=warmup_inputs)(kernel)
+    launcher[(3,)]("runtime", 4, CONST=8)
+
+    assert kernel.runtime_calls == [((3,), ("runtime", 4), {"CONST": 8})]
+
+    _patch_key_deriver(
+        monkeypatch,
+        lambda kernel, kwargs: {
+            TritonJitKey(id(kernel), "fake", 0, (kwargs["second"], kwargs["CONST"]))
+        },
+    )
+
+    keys = launcher.get_warmup_keys()
+    assert len(keys) == 2
+    launcher.compile(keys[0])
+    assert kernel.warmup_calls == [
+        {"grid": (1,), "first": "warmup", "second": 1, "CONST": 7}
+    ]
+
+
+def test_triton_warmup_inputs_expands_explicit_pointer_dtypes() -> None:
+    inputs = triton_warmup_inputs(
+        _pointer_group_test_kernel,
+        grid=(1,),
+        pointer_dtypes={torch.float32: ("x_ptr",)},
+        value=2,
+        BLOCK=16,
+    )
+
+    assert inputs["x_ptr"] == TritonWarmupTensor(torch.float32)
+    assert inputs["value"] == 2
+
+    with pytest.raises(ValueError, match="not a pointer: value"):
+        triton_warmup_inputs(
+            _pointer_group_test_kernel,
+            grid=(1,),
+            pointer_dtypes={torch.float32: ("value",)},
+            x_ptr=TritonWarmupTensor(torch.float32),
+            BLOCK=16,
+        )
+    with pytest.raises(ValueError, match="Missing Triton pointer inputs: x_ptr"):
+        triton_warmup_inputs(
+            _pointer_group_test_kernel,
+            grid=(1,),
+            pointer_dtypes={},
+            value=2,
+            BLOCK=16,
+        )
+
+
+def test_triton_key_derivation_applies_wrappers_and_runtime_options(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    knobs = triton.knobs
+    driver = triton.runtime.driver
+
+    device = "test-device"
+    monkeypatch.setattr(driver.active, "get_current_device", lambda: device)
+    calls: list[dict[str, Any]] = []
+
+    def binder(x: Any, value: int, BLOCK: int, **options: Any) -> Any:
+        calls.append(dict(BLOCK=BLOCK, **options))
+        specialization = [("pointer", x.dtype), ("i32", value), ("constexpr", BLOCK)]
+        return {}, specialization, options
+
+    cache: dict[Any, Any] = {}
+    _binder_test_kernel.device_caches[device] = ({}, cache, None, None, binder)
+    heuristic_kernel = triton.heuristics(
+        {"BLOCK": lambda args: 16 if args["value"] <= 16 else 32}
+    )(_binder_test_kernel)
+
+    keys = jit_warmup_triton_helper._triton_key_deriver(heuristic_kernel)(
+        {"x": TritonWarmupTensor(tl.float32), "value": 2}
+    )
+
+    assert len(keys) == 1
+    expected_options = {
+        "debug": knobs.runtime.debug,
+        "instrumentation_mode": knobs.compilation.instrumentation_mode,
+    }
+    if hasattr(knobs.compilation, "fpsan_homomorphic_casts"):
+        expected_options["fpsan_homomorphic_casts"] = (
+            knobs.compilation.fpsan_homomorphic_casts
+        )
+    assert calls == [{"BLOCK": 16, **expected_options}]
+
+
+def test_triton_key_derivation_covers_autotune_configs_and_jit_identity(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    driver = triton.runtime.driver
+
+    device = "test-device"
+    monkeypatch.setattr(driver.active, "get_current_device", lambda: device)
+
+    def binder(x: Any, value: int, BLOCK: int, **options: Any) -> Any:
+        specialization = [("pointer", x.dtype), ("i32", value), ("constexpr", BLOCK)]
+        return {}, specialization, options
+
+    for kernel in (_binder_test_kernel, _second_binder_test_kernel):
+        kernel.device_caches[device] = ({}, {}, None, None, binder)
+    autotuned = triton.autotune(
+        configs=[
+            triton.Config({"BLOCK": 16}, num_warps=2),
+            triton.Config({"BLOCK": 32}, num_warps=4),
+        ],
+        key=["value"],
+    )(_binder_test_kernel)
+    inputs = {"x": TritonWarmupTensor(tl.float32), "value": 2}
+
+    autotune_keys = jit_warmup_triton_helper._triton_key_deriver(autotuned)(inputs)
+    other_kernel_keys = jit_warmup_triton_helper._triton_key_deriver(
+        _second_binder_test_kernel
+    )(inputs | {"BLOCK": 16})
+
+    assert len(autotune_keys) == 2
+    assert {key.jit_function_key for key in autotune_keys} == {
+        _binder_test_kernel.cache_key
+    }
+    assert {key.jit_function_id for key in autotune_keys} == {id(_binder_test_kernel)}
+    assert autotune_keys.isdisjoint(other_kernel_keys)
+
+
 def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None:
     from vllm.v1.attention.backends.utils import PAD_SLOT_ID
     from vllm.v1.worker.block_table import ComputeSlotMappingKernel
@@ -107,13 +460,13 @@ def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None:
     )
     launches: list[tuple[Any, ...]] = []
 
-    def launch(grid: Any, inputs: Any, **kwargs: Any) -> None:
-        launches.append((grid, inputs, kwargs))
+    def launch(launch_spec: Any, inputs: Any) -> None:
+        launches.append((launch_spec, inputs))
 
     monkeypatch.setattr(owner, "launch", launch)
     owner.compile(compile_key)
 
-    grid, inputs, kwargs = launches[0]
+    (grid, kwargs), inputs = launches[0]
     assert grid == (2,)
     assert inputs["num_tokens"] == 2
     assert inputs["block_table_stride"] == 128
diff --git a/tests/model_executor/test_model_load_with_params.py b/tests/model_executor/test_model_load_with_params.py
index 785335e99111..7d1bddab0db2 100644
--- a/tests/model_executor/test_model_load_with_params.py
+++ b/tests/model_executor/test_model_load_with_params.py
@@ -23,9 +23,7 @@
     current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
 )
 def test_model_loading_with_params(vllm_runner, monkeypatch):
-    """
-    Test parameter weight loading with tp>1.
-    """
+    """Test parameter weight loading with tp>1."""
     # to use apply_model
     monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
     with vllm_runner(
@@ -68,9 +66,7 @@ def check_model(model):
     current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
 )
 def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):
-    """
-    Test parameter weight loading with tp>1.
-    """
+    """Test parameter weight loading with tp>1."""
     # to use apply_model
     monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
     with vllm_runner(
@@ -113,9 +109,7 @@ def check_model(model):
     current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
 )
 def test_facebook_roberta_model_loading_with_params(vllm_runner, monkeypatch):
-    """
-    Test loading roberta-base model with no lm_head.
-    """
+    """Test loading roberta-base model with no lm_head."""
     # to use apply_model
     monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
     model_name = "FacebookAI/roberta-base"
diff --git a/tests/model_executor/test_nemotron_h_quantization.py b/tests/model_executor/test_nemotron_h_quantization.py
index 81a3fb5b13fc..e4e1c87db411 100644
--- a/tests/model_executor/test_nemotron_h_quantization.py
+++ b/tests/model_executor/test_nemotron_h_quantization.py
@@ -3,6 +3,8 @@
 
 from unittest.mock import Mock, patch
 
+import torch
+
 
 def test_nemotron_h_lm_head_receives_quant_config():
     from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
@@ -32,3 +34,28 @@ def test_nemotron_h_lm_head_receives_quant_config():
         MockLMHead.assert_called_once()
         call_kwargs = MockLMHead.call_args.kwargs
         assert call_kwargs["quant_config"] is mock_quant_config
+
+
+def test_relu2_fp8_fusion_uses_registry():
+    from vllm.model_executor.models.nemotron_h import NemotronHMLP
+
+    projected = torch.empty((1, 1), dtype=torch.bfloat16)
+    fused = Mock()
+    act_fn = Mock()
+    down_proj = Mock(side_effect=lambda x: (x, None))
+
+    mlp = NemotronHMLP.__new__(NemotronHMLP)
+    torch.nn.Module.__init__(mlp)
+    mlp.up_proj = Mock(return_value=(projected, None))
+    mlp.down_proj = down_proj
+    mlp.act_fn = act_fn
+
+    with patch(
+        "vllm.model_executor.models.nemotron_h.maybe_fused_act_quant",
+        return_value=fused,
+    ) as maybe_fused:
+        result = mlp(Mock())
+
+    maybe_fused.assert_called_once_with(act_fn, projected, down_proj)
+    act_fn.assert_not_called()
+    assert result is fused
diff --git a/tests/model_executor/test_oink_integration.py b/tests/model_executor/test_oink_integration.py
index 2f37472b73ef..b16d71641369 100644
--- a/tests/model_executor/test_oink_integration.py
+++ b/tests/model_executor/test_oink_integration.py
@@ -62,7 +62,6 @@ def test_oink_availability_checks(
     expected_fused: bool,
 ):
     """Test OINK support detection with clean import state for each parameter set."""
-
     # Use spawn to run function in fresh process with clean imports
     # TODO migrate to spawn utility:
     # https://github.com/vllm-project/vllm/issues/41415
diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py
index e507b82d9c5e..e96dae818644 100644
--- a/tests/model_executor/test_qwen3_omni.py
+++ b/tests/model_executor/test_qwen3_omni.py
@@ -8,15 +8,15 @@
 import pytest
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
+from vllm.config import ParallelConfig
 from vllm.multimodal.processing import InputProcessingContext
 
 
 # Helper function to print input IDs with coalesced audio/video tokens.
 def print_input_ids(input_ids):
-    """
-    Print input IDs, compressing consecutive special tokens.
+    """Print input IDs, compressing consecutive special tokens.
     - 151675: <|audio_pad|>
     - 151656: <|video_pad|>
     """
@@ -52,7 +52,7 @@ def print_input_ids(input_ids):
 @pytest.fixture
 def mock_qwen3_omni_config():
     """Create a mock Qwen3OmniMoeThinker config."""
-    config = Mock(spec=PretrainedConfig)
+    config = Mock(spec=PreTrainedConfig)
     # Token IDs from https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
     config.audio_token_id = 151675  # <|audio_pad|>
     config.video_token_id = 151656  # <|video_pad|>
@@ -136,7 +136,6 @@ def test_qwen3_omni_get_updates_use_audio_in_video(
     mock_image_processor,
 ):
     """Test the get_updates_use_audio_in_video method directly."""
-
     from vllm.model_executor.models.qwen3_omni_moe_thinker import (
         Qwen3OmniMoeThinkerMultiModalProcessor,
         Qwen3OmniMoeThinkerProcessingInfo,
@@ -145,6 +144,11 @@ def test_qwen3_omni_get_updates_use_audio_in_video(
     # Create a mock context
     mock_ctx = Mock(spec=InputProcessingContext)
     mock_ctx.tokenizer = mock_tokenizer
+    # `model_config` is an instance attribute, so it is not covered by the
+    # spec; the data parser reads it for `allow_missing_mm_embeddings`.
+    mock_ctx.model_config = SimpleNamespace(
+        multimodal_config=SimpleNamespace(allow_missing_mm_embeddings=False)
+    )
 
     # Create processing info
     info = Qwen3OmniMoeThinkerProcessingInfo(mock_ctx)
@@ -293,6 +297,25 @@ def forward(self, hidden_states, residual):
     torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[15.0]]))
 
 
+def _dspark_vocab_stub(input_vocab_size: int, draft_vocab_size: int):
+    """A Qwen3 DSpark model with only what `load_weights` reads."""
+    from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM
+
+    model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM)
+    nn.Module.__init__(model)
+    object.__setattr__(
+        model,
+        "config",
+        SimpleNamespace(
+            vocab_size=input_vocab_size,
+            draft_vocab_size=draft_vocab_size,
+        ),
+    )
+    object.__setattr__(model, "target_vocab_size", 100)
+    object.__setattr__(model, "model", Mock(confidence_head=None))
+    return model
+
+
 @pytest.mark.skip_global_cleanup
 @pytest.mark.parametrize(
     ("input_vocab_size", "draft_vocab_size", "weights", "error"),
@@ -311,24 +334,25 @@ def forward(self, hidden_states, residual):
 def test_qwen3_dspark_rejects_incomplete_vocab_weights(
     input_vocab_size, draft_vocab_size, weights, error
 ):
-    from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM
-
-    model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM)
-    nn.Module.__init__(model)
-    object.__setattr__(
-        model,
-        "config",
-        SimpleNamespace(
-            vocab_size=input_vocab_size,
-            draft_vocab_size=draft_vocab_size,
-        ),
-    )
-    object.__setattr__(model, "target_vocab_size", 100)
+    model = _dspark_vocab_stub(input_vocab_size, draft_vocab_size)
 
     with pytest.raises(ValueError, match=error):
         model.load_weights(weights)
 
 
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize("draft_vocab_size", [100, 101])
+def test_qwen3_dspark_accepts_draft_vocab_at_least_target(draft_vocab_size):
+    """A draft vocab >= the target's maps identically, so needs no lm_head/d2t.
+
+    Guards the padded case (draft > target): the draft's unembedding is padded
+    to its physical size, which is not a reduced vocabulary.
+    """
+    model = _dspark_vocab_stub(100, draft_vocab_size)
+
+    model.load_weights([])
+
+
 @pytest.mark.skip_global_cleanup
 def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
     from vllm.v1.worker.gpu.spec_decode.dspark import utils as dspark_utils
@@ -350,10 +374,12 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
             draft_parallel_config=SimpleNamespace(tensor_parallel_size=1),
             attention_backend=None,
             kv_cache_dtype=None,
+            draft_load_config=None,
         ),
-        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        parallel_config=ParallelConfig(),
         attention_config=SimpleNamespace(backend=None),
         cache_config=SimpleNamespace(),
+        load_config=SimpleNamespace(load_format="auto"),
         model_config=SimpleNamespace(get_vocab_size=Mock(return_value=100)),
     )
 
@@ -368,6 +394,10 @@ def fake_replace(config, **changes):
             "vllm.v1.worker.gpu.spec_decode.eagle.utils.get_pp_group",
             return_value=SimpleNamespace(world_size=1),
         ),
+        patch(
+            "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group",
+            return_value=SimpleNamespace(world_size=1),
+        ),
         patch(
             "vllm.compilation.backends.set_model_tag",
             return_value=nullcontext(),
diff --git a/tests/model_executor/test_qwen3_vl_mrope.py b/tests/model_executor/test_qwen3_vl_mrope.py
index 90d9fd6e4ff8..a50bf319c938 100644
--- a/tests/model_executor/test_qwen3_vl_mrope.py
+++ b/tests/model_executor/test_qwen3_vl_mrope.py
@@ -51,8 +51,7 @@ class DummyConfig:
 def make_video_embedding(
     t, h, w, interleave_text_tokens: tuple[int, int], video_pruning_rate: float = 0.0
 ):
-    """
-    Helper function to make a video embedding for a given video size and pruning rate.
+    """Make a video embedding for a given video size and pruning rate.
 
     Args:
         t: Number of frames.
@@ -64,6 +63,7 @@ def make_video_embedding(
 
     Returns:
         Tuple of (unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask)
+
     """
     unpruned_tokens_sequence = []
     population = list(range(1, 100))
diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py
index 49986a991e2b..ff683ae44c67 100644
--- a/tests/model_executor/test_routed_experts_capture.py
+++ b/tests/model_executor/test_routed_experts_capture.py
@@ -7,26 +7,19 @@
 import pytest
 import torch
 
-from vllm.config import ModelConfig, VllmConfig
+from vllm.config import VllmConfig
 from vllm.config.compilation import CompilationMode
 from vllm.distributed.eplb.eplb_state import EplbLayerState
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
     RoutedExpertsCapturer,
-    RoutedExpertsManager,
     bind_routed_experts_capturer,
-    get_routed_experts_attn_gid,
 )
 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
-from vllm.transformers_utils.model_arch_config_convertor import (
-    ModelArchConfigConvertorBase,
-)
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
-    KVCacheConfig,
     KVCacheGroupSpec,
     MLAAttentionSpec,
-    SlidingWindowMLASpec,
     UniformTypeKVCacheSpecs,
 )
 
@@ -42,6 +35,7 @@ def _capturer_with_buffer(
     num_experts_per_tok: int = 2,
     dp_rank: int = 0,
     tp_size: int = 1,
+    dtype: torch.dtype = torch.int32,
 ) -> RoutedExpertsCapturer:
     # Bypass __init__ so the test can use a CPU buffer and skip the
     # VllmConfig dependency. The CUDA device-tensor allocation in the
@@ -49,10 +43,11 @@ def _capturer_with_buffer(
     c = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
     c.dp_rank = dp_rank
     c.tp_size = tp_size
+    c.output_dtype = torch.uint8
     c.device_buffer = torch.full(
         (max_tokens, num_layers, num_experts_per_tok),
         -1,
-        dtype=torch.int32,
+        dtype=dtype,
     )
     return c
 
@@ -88,74 +83,36 @@ def _make_modular_routed_experts():
     )
 
 
-def _make_model_config(hf_config):
-    num_experts_per_token = ModelArchConfigConvertorBase(
-        hf_config, hf_config
-    ).get_num_experts_per_token()
-    model_config = SimpleNamespace(
-        hf_text_config=hf_config,
-        model_arch_config=SimpleNamespace(
-            num_experts_per_token=num_experts_per_token,
-        ),
-    )
-    model_config.get_num_experts = lambda: hf_config.num_experts
-    model_config.get_num_experts_per_tok = lambda: (
-        ModelConfig.get_num_experts_per_tok(model_config)
-    )
-    model_config.get_total_num_hidden_layers = lambda: hf_config.num_hidden_layers
-    return model_config
-
-
-def test_routed_experts_manager_uses_gemma4_top_k_experts():
-    hf_config = SimpleNamespace(
-        num_experts=8,
-        top_k_experts=2,
-        num_hidden_layers=3,
-    )
-    vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config))
-    kv_cache_spec = FullAttentionSpec(
-        block_size=4,
-        num_kv_heads=1,
-        head_size=1,
-        dtype=torch.float32,
-    )
-    kv_cache_config = KVCacheConfig(
-        num_blocks=2,
-        kv_cache_tensors=[],
-        kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
-    )
-
-    manager = RoutedExpertsManager(vllm_config, kv_cache_config)
-
-    assert manager.routed_experts_by_slot.shape == (8, 3, 2)
-
-
-def test_routed_experts_manager_uses_kimi_k3_experts_per_token():
-    hf_config = SimpleNamespace(
-        num_experts=8,
-        num_experts_per_token=2,
-        num_hidden_layers=3,
-    )
-    vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config))
-    kv_cache_spec = FullAttentionSpec(
-        block_size=4,
+def _full_attention_kv_group(
+    spec_type: type[FullAttentionSpec] = FullAttentionSpec,
+) -> KVCacheGroupSpec:
+    full_attention = spec_type(
+        block_size=16,
         num_kv_heads=1,
         head_size=1,
         dtype=torch.float32,
     )
-    kv_cache_config = KVCacheConfig(
-        num_blocks=2,
-        kv_cache_tensors=[],
-        kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
+    return KVCacheGroupSpec(
+        ["layer"],
+        UniformTypeKVCacheSpecs(
+            block_size=16,
+            kv_cache_specs={"layer": full_attention},
+        ),
     )
 
-    manager = RoutedExpertsManager(vllm_config, kv_cache_config)
-
-    assert manager.routed_experts_by_slot.shape == (8, 3, 2)
 
+@pytest.mark.parametrize("eplb_enabled", [False, True])
+def test_base_router_capture_pre_eplb_mapping(eplb_enabled):
+    eplb_state = None
+    if eplb_enabled:
+        eplb_state = EplbLayerState()
+        eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
+        eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
+        eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
+        eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)
+        eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)]
+    router = _make_router(eplb_state)
 
-def test_base_router_capture_pre_eplb_mapping():
-    router = _make_router()
     captured = []
 
     def capture_fn(ids):
@@ -173,34 +130,7 @@ def capture_fn(ids):
     assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
 
 
-def test_base_router_capture_with_eplb_enabled():
-    eplb_state = EplbLayerState()
-    eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
-    eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
-    eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
-    eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)
-    eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)]
-    router = _make_router(eplb_state=eplb_state)
-
-    captured = []
-
-    def capture_fn(ids):
-        captured.append(ids.clone())
-
-    router.set_capture_fn(capture_fn)
-    _, topk_ids = router.select_experts(
-        hidden_states=torch.empty(1),
-        router_logits=torch.empty(1),
-    )
-
-    assert len(captured) == 1
-    # Capture should see logical ids pre-EPLB mapping.
-    assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
-    # Our DummyRouter mapping adds +10.
-    assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
-
-
-def test_public_binding_only_visits_target_model(monkeypatch):
+def test_public_binding_binds_target_model_router(monkeypatch):
     class DummyFusedMoE:
         def __init__(self, layer_id):
             self.layer_id = layer_id
@@ -208,7 +138,6 @@ def __init__(self, layer_id):
             self._quant_method = _make_modular_routed_experts().quant_method
 
     target_module = DummyFusedMoE(layer_id=7)
-    draft_module = DummyFusedMoE(layer_id=0)
 
     import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
 
@@ -221,12 +150,30 @@ def __init__(self, layer_id):
     )
 
     assert target_module.router.capture_fn is not None
-    assert draft_module.router.capture_fn is None
     topk_ids = torch.tensor([[5, 6]])
     target_module.router.capture_fn(topk_ids)
     assert calls == [(7, topk_ids)]
 
 
+def test_public_binding_supports_direct_capture_source():
+    class DirectCaptureSource(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layer_id = 3
+            self.capture_fn = None
+
+    source = DirectCaptureSource()
+    calls = []
+    capturer = types.SimpleNamespace(capture=lambda *args: calls.append(args))
+
+    bind_routed_experts_capturer(source, capturer)
+
+    assert source.capture_fn is not None
+    topk_ids = torch.tensor([[1, 2]])
+    source.capture_fn(topk_ids)
+    assert calls == [(3, topk_ids)]
+
+
 def test_public_binding_rejects_monolithic_without_replay_support(monkeypatch):
     class DummyFusedMoE:
         def __init__(self):
@@ -277,8 +224,53 @@ def test_routed_experts_capturer_single_dp_no_metadata():
     assert capturer.device_buffer[3, 0, 0].item() == -1
 
 
+@pytest.mark.parametrize(
+    ("num_experts", "dtype_name", "torch_dtype"),
+    [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)],
+)
+def test_routed_experts_capturer_exposes_output_profile(
+    monkeypatch, num_experts, dtype_name, torch_dtype
+):
+    import vllm.model_executor.layers.fused_moe.routed_experts_capturer as module
+
+    monkeypatch.setattr(module, "current_platform", SimpleNamespace(device_type="cpu"))
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            get_total_num_hidden_layers=lambda: 3,
+            get_num_experts=lambda: num_experts,
+            get_num_experts_per_tok=lambda: 2,
+            hf_text_config=SimpleNamespace(model_type="test"),
+        ),
+        parallel_config=SimpleNamespace(data_parallel_rank=0, tensor_parallel_size=1),
+    )
+
+    capturer = RoutedExpertsCapturer(8, config)
+
+    assert capturer.shape_per_token == (3, 2)
+    assert capturer.output_dtype_name == dtype_name
+    assert capturer.output_dtype == torch_dtype
+    assert capturer.device_buffer.shape == (8, 3, 2)
+
+
+@pytest.mark.parametrize("output_dtype", [torch.uint8, torch.uint16])
+def test_routed_experts_capturer_narrows_snapshot(output_dtype):
+    capturer = _capturer_with_buffer(dtype=torch.int32)
+    capturer.output_dtype = output_dtype
+    topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64)
+    ctx = SimpleNamespace(dp_metadata=None)
+    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
+        capturer.capture(layer_id=0, topk_ids=topk)
+
+    output = capturer.snapshot_routing_data(2)
+    assert capturer.device_buffer.dtype == torch.int32
+    assert capturer.device_buffer[:2, 0, :].tolist() == topk.tolist()
+    assert output.dtype == output_dtype
+    assert output[:, 0, :].tolist() == topk.tolist()
+    assert output.data_ptr() != capturer.device_buffer.data_ptr()
+
+
 def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
-    """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
+    """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
     capturer = _capturer_with_buffer(dp_rank=1)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
     ctx = SimpleNamespace(
@@ -295,7 +287,7 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
 
 
 def test_routed_experts_capturer_dp_modular_local_tokens():
-    """n == token_num_per_dp: topk is already local to this DP rank."""
+    """N == token_num_per_dp: topk is already local to this DP rank."""
     capturer = _capturer_with_buffer(dp_rank=1)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
     ctx = SimpleNamespace(
@@ -307,6 +299,47 @@ def test_routed_experts_capturer_dp_modular_local_tokens():
     assert torch.equal(capturer.device_buffer[:3, 0, :], topk)
 
 
+def test_routed_experts_capturer_dp_ep_gathered_shards():
+    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+    num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32)
+    ctx = SimpleNamespace(
+        dp_metadata=SimpleNamespace(
+            num_tokens_across_dp_cpu=num_tokens_dp,
+            local_sizes=[2, 2, 1, 1],
+        )
+    )
+    topk = torch.tensor(
+        [[0, 1], [2, 3], [4, 5], [-1, -1], [10, 11], [12, 13]],
+        dtype=torch.int32,
+    )
+
+    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
+        capturer.capture(layer_id=0, topk_ids=topk)
+
+    assert torch.equal(capturer.device_buffer[:2, 0], topk[4:])
+
+
+def test_routed_experts_capturer_sp_modular_gathers_tp_shards():
+    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+    num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
+    ctx = SimpleNamespace(
+        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    )
+    local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32)
+    gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32)
+    tp_group = Mock()
+    tp_group.all_gather.return_value = gathered
+
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
+    ):
+        capturer.capture(layer_id=0, topk_ids=local_shard)
+
+    tp_group.all_gather.assert_called_once_with(local_shard, dim=0)
+    assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3])
+
+
 def test_routed_experts_capturer_dp_unexpected_batch_raises():
     """Mismatch between topk batch dim and DP layout: fail fast."""
     capturer = _capturer_with_buffer(dp_rank=0)
@@ -324,128 +357,119 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises():
     assert capturer.device_buffer[0, 0, 0].item() == -1
 
 
-def test_routed_experts_attention_group_is_shared_and_fail_closed():
-    """Both sides key routing data by this gid, so it must skip non-full-attention
-    groups rather than defaulting to 0, and fail closed when none exists."""
-    common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
-    config = SimpleNamespace(
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["swa_layer"],
-                SlidingWindowMLASpec(block_size=4, sliding_window=8, **common),
-            ),
-            KVCacheGroupSpec(["full_layer"], FullAttentionSpec(block_size=4, **common)),
-        ]
-    )
-    assert get_routed_experts_attn_gid(config) == 1
-
-    with pytest.raises(ValueError, match="requires a full-attention KV cache group"):
-        get_routed_experts_attn_gid(SimpleNamespace(kv_cache_groups=[]))
-
-
-def test_routed_experts_attention_group_unwraps_uniform_type_specs():
-    """DeepSeek-V4-shaped groups wrap their specs in ``UniformTypeKVCacheSpecs``.
+def test_get_aux_output_connector_passes_config(monkeypatch):
+    import vllm.distributed.aux_output_connector.worker as aux_output_worker
 
-    The wrapper is not a ``FullAttentionSpec``, so a bare isinstance check finds
-    no group and fails closed on every worker. Unwrap semantics themselves are
-    covered by ``test_is_full_attention_spec_*`` in tests/v1/core.
-    """
-    common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
-    swa_spec = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common)
-    mla_spec = MLAAttentionSpec(block_size=4, **common)
+    connector = Mock()
+    constructor = Mock(return_value=connector)
+    monkeypatch.setattr(aux_output_worker, "AuxOutputWorkerConnector", constructor)
     config = SimpleNamespace(
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["swa_layer"],
-                UniformTypeKVCacheSpecs(
-                    block_size=4, kv_cache_specs={"swa_layer": swa_spec}
-                ),
-            ),
-            KVCacheGroupSpec(
-                ["mla_layer"],
-                UniformTypeKVCacheSpecs(
-                    block_size=4, kv_cache_specs={"mla_layer": mla_spec}
-                ),
-            ),
-        ]
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32)
     )
+    model = Mock()
+    kv_cache_config = Mock()
 
-    assert get_routed_experts_attn_gid(config) == 1
+    result = aux_output_worker.get_aux_output_connector(model, config, kv_cache_config)
 
-    swa_only = SimpleNamespace(kv_cache_groups=[config.kv_cache_groups[0]])
-    with pytest.raises(ValueError, match="requires a full-attention KV cache group"):
-        get_routed_experts_attn_gid(swa_only)
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_mrv2_async_output_returns_existing_routed_experts_field():
-    from vllm.v1.outputs import ModelRunnerOutput, RoutedExpertsTensors
-    from vllm.v1.worker.gpu.async_utils import AsyncOutput
-    from vllm.v1.worker.gpu.sample.output import SamplerOutput
-
-    routed_experts = RoutedExpertsTensors(
-        routing_data=torch.arange(6, dtype=torch.int32, device="cuda").reshape(3, 1, 2),
-        slot_mapping=torch.tensor([11, 12, 13], device="cuda"),
-    )
-    num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda")
-    sampler_output = SamplerOutput(
-        sampled_token_ids=torch.tensor([[1]], device="cuda"),
-        logprobs_tensors=None,
-        num_nans=None,
-        num_sampled=num_sampled,
-        num_rejected=torch.tensor([0], dtype=torch.int32, device="cuda"),
+    constructor.assert_called_once_with(
+        model=model,
+        kv_cache_config=kv_cache_config,
+        vllm_config=config,
     )
-    output = AsyncOutput(
-        model_runner_output=ModelRunnerOutput(req_ids=["req"], req_id_to_index={}),
-        sampler_output=sampler_output,
-        num_sampled_tokens=num_sampled,
-        main_stream=torch.cuda.current_stream(),
-        copy_stream=torch.cuda.Stream(),
-        routed_experts=routed_experts,
-    ).get_output()
-
-    assert output.routed_experts is not None
-    assert output.routed_experts.routing_data[:, 0, 0].tolist() == [0, 2, 4]
-    assert output.routed_experts.slot_mapping.tolist() == [11, 12, 13]
+    assert result is connector
 
 
-@pytest.mark.parametrize("rank", [0, 1])
-def test_all_tp_ranks_initialize_capture(monkeypatch, rank):
-    pytest.importorskip("vllm.vllm_flash_attn", exc_type=ImportError)
-    import vllm.v1.worker.gpu.model_runner as model_runner
+def test_aux_output_worker_connector_binds_capture_on_non_output_rank(monkeypatch):
+    import vllm.distributed.aux_output_connector.worker as aux_output_worker
 
     capturer = Mock()
     constructor = Mock(return_value=capturer)
     bind = Mock()
-    monkeypatch.setattr(model_runner, "RoutedExpertsCapturer", constructor)
-    monkeypatch.setattr(model_runner, "bind_routed_experts_capturer", bind)
-
-    runner = model_runner.GPUModelRunner.__new__(model_runner.GPUModelRunner)
-    runner.max_num_tokens = 32
-    runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(rank=rank))
-    runner.kv_cache_config = SimpleNamespace()
-    runner.model = Mock()
+    monkeypatch.setattr(aux_output_worker, "RoutedExpertsCapturer", constructor)
+    monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", bind)
+    monkeypatch.setattr(
+        aux_output_worker,
+        "get_tp_group",
+        lambda: SimpleNamespace(is_first_rank=False, world_size=1),
+    )
 
-    runner.init_routed_experts_capturer()
+    config = SimpleNamespace(
+        aux_output_config=SimpleNamespace(enable_return_routed_experts=True),
+        kv_transfer_config=None,
+        max_concurrent_batches=2,
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32),
+    )
+    model = Mock()
+    connector = aux_output_worker.AuxOutputWorkerConnector(
+        vllm_config=config,
+        model=model,
+        kv_cache_config=SimpleNamespace(kv_cache_groups=[_full_attention_kv_group()]),
+    )
 
     constructor.assert_called_once_with(
         max_num_batched_tokens=32,
-        vllm_config=runner.vllm_config,
-        kv_cache_config=runner.kv_cache_config,
+        vllm_config=config,
+    )
+    bind.assert_called_once_with(model, capturer)
+    connector.begin_step(Mock())
+    assert connector.prepare_output(Mock()) is None
+    capturer.snapshot_routing_data.assert_not_called()
+
+
+def test_aux_output_worker_connector_default_capacity(monkeypatch):
+    import vllm.distributed.aux_output_connector.worker as aux_output_worker
+
+    tp_group = SimpleNamespace(is_first_rank=True, world_size=1)
+    store_constructor = Mock()
+    background_store_constructor = Mock(side_effect=lambda store, **_: store)
+    capturer = SimpleNamespace(shape_per_token=(2,), output_dtype_name="int32")
+    monkeypatch.setattr(aux_output_worker, "get_tp_group", lambda: tp_group)
+    monkeypatch.setattr(
+        aux_output_worker, "RoutedExpertsCapturer", Mock(return_value=capturer)
+    )
+    monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", Mock())
+    monkeypatch.setattr(
+        aux_output_worker,
+        "resolve_kv_cache_block_sizes",
+        lambda *_: (32, 16),
+    )
+    monkeypatch.setattr(aux_output_worker, "BlockObjectStore", store_constructor)
+    monkeypatch.setattr(
+        aux_output_worker, "BackgroundBlockObjectStore", background_store_constructor
     )
-    bind.assert_called_once_with(runner.model, capturer)
-    assert runner.routed_experts_capturer is capturer
+    monkeypatch.setattr(aux_output_worker, "RoutedExpertsBuffer", Mock())
+
+    config = SimpleNamespace(
+        aux_output_config=SimpleNamespace(max_bytes=None),
+        kv_transfer_config=None,
+        cache_config=SimpleNamespace(enable_prefix_caching=True),
+        scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=32),
+        max_concurrent_batches=2,
+    )
+    kwargs = dict(
+        vllm_config=config,
+        model=Mock(),
+        kv_cache_config=SimpleNamespace(
+            num_blocks=10,
+            kv_cache_groups=[_full_attention_kv_group(MLAAttentionSpec)],
+        ),
+    )
+
+    aux_output_worker.AuxOutputWorkerConnector(**kwargs)
+    assert store_constructor.call_args.kwargs["max_bytes"] == 2560
+    assert store_constructor.call_args.kwargs["object_nbytes"] == 128
+    assert background_store_constructor.call_args.kwargs["max_pending_batches"] == 16
 
 
 def test_v2_model_runner_accepts_routed_experts(monkeypatch):
     monkeypatch.setattr("importlib.metadata.entry_points", lambda **_: ())
     config = SimpleNamespace(
         model_config=SimpleNamespace(
-            enable_return_routed_experts=True,
             use_mla=False,
             logits_processors=None,
             enable_prompt_embeds=False,
         ),
+        aux_output_config=SimpleNamespace(enable_return_routed_experts=True),
         speculative_config=None,
         parallel_config=SimpleNamespace(
             prefill_context_parallel_size=1,
diff --git a/tests/model_executor/test_watermark_sample_warmup.py b/tests/model_executor/test_watermark_sample_warmup.py
new file mode 100644
index 000000000000..73ecb95ee36a
--- /dev/null
+++ b/tests/model_executor/test_watermark_sample_warmup.py
@@ -0,0 +1,300 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The warmup must cover every specialization the sampler will launch.
+
+``_philox_gumbel_kernel`` is keyed on the Philox key halves, on the logits
+dtype (model dtype when the sampler passes logits through, fp32 when logits
+processing copies them) and on whether a skip mask is passed, so a missing
+combination leaves a request paying the compile this module exists to remove.
+"""
+
+import logging
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.config.watermarking import WatermarkConfig, derive_watermark_key
+from vllm.model_executor.warmup.watermark_sample_warmup import watermark_sample_warmup
+from vllm.platforms import current_platform
+from vllm.v1.watermarking.factory import create_watermarker
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
+from vllm.v1.watermarking.spec_decode import create_speculative_target_watermarker
+
+MASTER_KEY = 20250910
+KEY_A = derive_watermark_key(MASTER_KEY, b"key_a")
+KEY_B = derive_watermark_key(MASTER_KEY, b"key_b")
+
+
+@pytest.fixture(autouse=True)
+def cuda_alike_platform(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The warmup returns early off CUDA, so pin the platform it reads.
+
+    Everything here runs on the fake worker's cpu device; without this the
+    tests below would pass only on a CUDA-alike host and silently warm nothing
+    anywhere else.
+    """
+    monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
+
+
+def _sampler(
+    watermark_config: WatermarkConfig,
+    speculative: bool,
+    use_fp64_gumbel: bool,
+) -> GPUWatermarkSampler:
+    """The sampler ``GPUModelRunner`` builds for this config.
+
+    Mirrors ``vllm/v1/worker/gpu/model_runner.py:473-486``. ``__new__`` skips
+    ``Sampler.__init__``, which allocates UVA-backed device buffers; the warmup
+    only reads the three attributes set here.
+    """
+    watermarker = create_watermarker(watermark_config)
+    if speculative:
+        watermarker = create_speculative_target_watermarker(watermarker)
+    sampler = GPUWatermarkSampler.__new__(GPUWatermarkSampler)
+    sampler.watermarker = watermarker
+    sampler.deduplicate_contexts = watermark_config.deduplicate_contexts
+    sampler.use_fp64_gumbel = use_fp64_gumbel
+    return sampler
+
+
+def _worker(
+    watermark_config: WatermarkConfig | None,
+    speculative: bool = False,
+    use_fp64_gumbel: bool = False,
+    with_sampler: bool = True,
+) -> SimpleNamespace:
+    sampler = (
+        _sampler(watermark_config, speculative, use_fp64_gumbel)
+        if watermark_config is not None and with_sampler
+        else None
+    )
+    return SimpleNamespace(
+        device=torch.device("cpu"),
+        model_runner=SimpleNamespace(sampler=sampler),
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                get_vocab_size=lambda: 128,
+                dtype=torch.bfloat16,
+            ),
+            watermark_config=watermark_config,
+        ),
+    )
+
+
+def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
+    calls: list[dict] = []
+
+    def fake_philox_gumbel_sample(logits, contexts, key, **kwargs):
+        calls.append({"logits": logits, "contexts": contexts, "key": key, **kwargs})
+        return torch.zeros(logits.shape[0], dtype=torch.int64)
+
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample",
+        fake_philox_gumbel_sample,
+    )
+    return calls
+
+
+def _variants(calls: list[dict]) -> set[tuple[int, torch.dtype, bool]]:
+    return {
+        (call["key"], call["logits"].dtype, call.get("skip_mask") is not None)
+        for call in calls
+    }
+
+
+def test_no_warmup_without_a_watermark_config(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(_worker(None))
+
+    assert calls == []
+
+
+def test_no_warmup_without_a_watermarked_sampler(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A rank that does not sample launches nothing to warm.
+
+    Only the final pipeline stage of a generative model holds a
+    ``GPUWatermarkSampler``; on any other rank the kernel is never launched.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(WatermarkConfig(key=MASTER_KEY), with_sampler=False)
+    )
+
+    assert calls == []
+
+
+def test_no_warmup_off_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    """``GumbelWatermarker`` launches the kernel only on a CUDA-alike device.
+
+    Elsewhere -- an XPU worker inherits ``kernel_warmup`` from the CUDA worker
+    -- the sampler takes the torch path, so warming would compile a kernel
+    nothing launches.
+    """
+    calls = _capture_calls(monkeypatch)
+    monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: False)
+    watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
+
+    assert calls == []
+
+
+def test_single_key_warms_both_logits_dtypes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
+
+    assert _variants(calls) == {
+        (MASTER_KEY, torch.bfloat16, True),
+        (MASTER_KEY, torch.float32, True),
+    }
+    assert len(calls) == len(_variants(calls))
+
+
+def test_dual_key_warms_both_derived_keys(monkeypatch: pytest.MonkeyPatch) -> None:
+    """``DualKeyGumbelWatermarker.sample`` launches the kernel once per key.
+
+    The halves of a derived key exceed 2**31, so each key is typed i64/i32 by
+    magnitude and a wrong key compiles a specialization the engine never uses.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel"))
+    )
+
+    assert _variants(calls) == {
+        (KEY_A, torch.bfloat16, True),
+        (KEY_A, torch.float32, True),
+        (KEY_B, torch.bfloat16, True),
+        (KEY_B, torch.float32, True),
+    }
+
+
+@pytest.mark.parametrize("alpha,key", [(0.0, KEY_A), (1.0, KEY_B)])
+def test_dual_key_alpha_bounds_warm_one_key(
+    monkeypatch: pytest.MonkeyPatch, alpha: float, key: int
+) -> None:
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(
+            WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel", alpha=alpha)
+        )
+    )
+
+    assert _variants(calls) == {
+        (key, torch.bfloat16, True),
+        (key, torch.float32, True),
+    }
+
+
+def test_speculative_config_warms_the_target_role_only(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Under spec decode the sampler holds ``create_target_watermarker()``.
+
+    Rows without draft tokens sample through it, so the sampler path launches
+    key B alone; the draft key runs in the speculator's own sampler.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(
+            WatermarkConfig(key=MASTER_KEY, algorithm="dual_key_gumbel"),
+            speculative=True,
+        )
+    )
+
+    assert _variants(calls) == {
+        (KEY_B, torch.bfloat16, True),
+        (KEY_B, torch.float32, True),
+    }
+
+
+def test_dedup_none_also_warms_the_mask_free_specialization(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Without deduplication an all-watermarked batch passes no skip mask.
+
+    A later batch mixing in an opted-out request passes one, so both
+    specializations are reachable and both must be warmed.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none"))
+    )
+
+    assert _variants(calls) == {
+        (MASTER_KEY, torch.bfloat16, True),
+        (MASTER_KEY, torch.float32, True),
+        (MASTER_KEY, torch.bfloat16, False),
+        (MASTER_KEY, torch.float32, False),
+    }
+    for call in calls:
+        if call.get("skip_mask") is None:
+            assert "temperatures" not in call
+
+
+def test_warmup_argument_dtypes_match_the_runtime_buffers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Contexts are int32, the sampling state int64 indices and fp32 temperatures.
+
+    ``GPUWatermarkSampler._get_contexts`` gathers from the int32 request-state
+    token ids, ``expanded_idx_mapping``/``positions``/``seeds`` are int64 and
+    ``SamplingStates.temperature`` is fp32; warming other dtypes compiles
+    specializations no engine launches.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY, context_width=3)))
+
+    assert calls
+    for call in calls:
+        num_tokens = call["logits"].shape[0]
+        assert call["contexts"].dtype == torch.int32
+        assert call["contexts"].shape == (num_tokens, 3)
+        assert call["skip_mask"].dtype == torch.bool
+        assert call["expanded_idx_mapping"].dtype == torch.int64
+        assert call["positions"].dtype == torch.int64
+        assert call["seeds"].dtype == torch.int64
+        assert call["temperatures"].dtype == torch.float32
+
+
+def test_use_fp64_follows_the_path_that_forwards_it(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Only the mixed path forwards the sampler's ``use_fp64_gumbel``.
+
+    ``GumbelWatermarker._sample_watermarked`` launches the mask-free variant
+    with the ``use_fp64=False`` default, so warming it with fp64 would compile
+    a specialization the engine never launches and miss the one it does.
+    """
+    calls = _capture_calls(monkeypatch)
+    watermark_sample_warmup(
+        _worker(
+            WatermarkConfig(key=MASTER_KEY, deduplicate_contexts="none"),
+            use_fp64_gumbel=True,
+        )
+    )
+
+    assert calls
+    for call in calls:
+        assert call["use_fp64"] is (call.get("skip_mask") is not None)
+
+
+def test_warmup_failure_does_not_abort_startup(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    def failing_philox_gumbel_sample(*args, **kwargs):
+        raise RuntimeError("no triton")
+
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample",
+        failing_philox_gumbel_sample,
+    )
+    with caplog.at_level(logging.WARNING):
+        watermark_sample_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
+
+    assert "Skipping watermark sampler warmup." in caplog.text
diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py
index e577690c92fc..3dea587075f2 100644
--- a/tests/model_executor/test_weight_utils.py
+++ b/tests/model_executor/test_weight_utils.py
@@ -72,28 +72,28 @@ def test_qkv_proj_v_scale(self):
         assert result == "model.layers.0.self_attn.attn.v_scale"
 
     def test_modelopt_k_proj_k_scale(self):
-        """ModelOpt format: k_proj.k_scale -> attn.k_scale"""
+        """ModelOpt format: k_proj.k_scale -> attn.k_scale."""
         result = maybe_remap_kv_scale_name(
             "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT
         )
         assert result == "model.layers.0.self_attn.attn.k_scale"
 
     def test_modelopt_v_proj_v_scale(self):
-        """ModelOpt format: v_proj.v_scale -> attn.v_scale"""
+        """ModelOpt format: v_proj.v_scale -> attn.v_scale."""
         result = maybe_remap_kv_scale_name(
             "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT
         )
         assert result == "model.layers.0.self_attn.attn.v_scale"
 
     def test_deprecated_kv_scale(self):
-        """Old format: kv_scale -> attn.k_scale (deprecated)"""
+        """Old format: kv_scale -> attn.k_scale (deprecated)."""
         result = maybe_remap_kv_scale_name(
             "model.layers.0.self_attn.kv_scale", self.PARAMS_DICT
         )
         assert result == "model.layers.0.self_attn.attn.k_scale"
 
     def test_default_bare_k_scale(self):
-        """Default format: .k_scale -> .attn.k_scale"""
+        """Default format: .k_scale -> .attn.k_scale."""
         result = maybe_remap_kv_scale_name(
             "model.layers.0.self_attn.k_scale", self.PARAMS_DICT
         )
diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py
index 68cf493b443a..234a63fb73e8 100644
--- a/tests/models/glm5next/test_kda_recurrent.py
+++ b/tests/models/glm5next/test_kda_recurrent.py
@@ -4,18 +4,23 @@
 
 The decode path hands the kernel column slices of the merged ``q|k|v`` conv
 output and of the fused ``qkvbfg_a`` projection (beta), so q/k/v/beta are
-token-strided rather than contiguous. The kernel must read them in place,
-match a pure-PyTorch recurrence, and reject layouts it cannot address.
+token-strided rather than contiguous. Both backends must match a pure-PyTorch
+recurrence. CUDA reads these slices in place and rejects unsupported layouts;
+the ROCm wrapper makes contiguous copies before launching its kernel.
 """
 
 import pytest
 import torch
 
-from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
 from vllm.platforms import current_platform
 
+if current_platform.is_rocm():
+    from vllm.models.glm5next.amd.ops.third_party.kda import fused_recurrent_kda
+else:
+    from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
+
 pytestmark = pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="CUDA-only Triton kernel"
+    not current_platform.is_cuda_alike(), reason="Requires CUDA or ROCm"
 )
 
 H, D = 16, 128
@@ -160,6 +165,10 @@ def test_fused_recurrent_kda_strided_inputs_bit_identical_to_contiguous(
     torch.testing.assert_close(state, state_ref, rtol=0, atol=0)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="Only the CUDA implementation requires zero-copy token-strided layouts",
+)
 @torch.inference_mode()
 def test_fused_recurrent_kda_rejects_unaddressable_layouts():
     """Layouts the token-stride addressing cannot express must fail loudly
diff --git a/tests/models/glm5next/test_sequence_parallel.py b/tests/models/glm5next/test_sequence_parallel.py
new file mode 100644
index 000000000000..2769f786bb09
--- /dev/null
+++ b/tests/models/glm5next/test_sequence_parallel.py
@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-5.3-Flash sequence-parallel MoE layout.
+
+With DP > 1, TP > 1 and expert parallelism, ``Glm5NextModel`` shards the token
+dimension across the TP group once at the model entry and every layer runs its
+MLP on that shard. A module that still does tensor-parallel collectives there
+sums the partial results of *different* tokens held by different ranks, so the
+dense MLP of the first ``first_k_dense_replace`` layers must hold replicated
+weights and run without any collective, like the shared experts already do.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor import parameter
+from vllm.model_executor.layers import linear
+from vllm.models.glm5next.common import model as glm_model
+from vllm.transformers_utils.configs.glm5_next import Glm5NextTextConfig
+
+HIDDEN, INTERMEDIATE, TP_SIZE = 8, 16, 2
+
+
+class _Attention(nn.Module):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+        self.o_proj = SimpleNamespace(reduce_results=True)
+
+
+def _fake_tensor_parallel_world(monkeypatch) -> None:
+    """Pretend to be rank 1 of a TP=2 group whose all-reduce must never run."""
+
+    def _forbidden_all_reduce(*args, **kwargs):
+        raise AssertionError("tensor-parallel all-reduce on a sequence-parallel shard")
+
+    for module in (linear, parameter):
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_world_size", lambda: TP_SIZE
+        )
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
+    monkeypatch.setattr(
+        linear, "tensor_model_parallel_all_reduce", _forbidden_all_reduce
+    )
+
+
+def _reference_mlp(mlp: glm_model.Glm5NextMLP, x: torch.Tensor) -> torch.Tensor:
+    gate_up = x @ mlp.gate_up_proj.weight.T
+    gate, up = gate_up.split(INTERMEDIATE, dim=-1)
+    return (torch.nn.functional.silu(gate) * up) @ mlp.down_proj.weight.T
+
+
+def _assert_runs_locally_on_shard(mlp: glm_model.Glm5NextMLP) -> None:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    mlp.to(device)
+    assert mlp.gate_up_proj.weight.shape == (2 * INTERMEDIATE, HIDDEN)
+    assert mlp.down_proj.weight.shape == (HIDDEN, INTERMEDIATE)
+    torch.manual_seed(0)
+    for param in mlp.parameters():
+        param.data.normal_()
+    shard = torch.randn(3, HIDDEN, device=device)
+    torch.testing.assert_close(mlp(shard), _reference_mlp(mlp, shard))
+
+
+def test_dense_mlp_runs_locally_on_the_sequence_parallel_shard(
+    monkeypatch, default_vllm_config
+):
+    _fake_tensor_parallel_world(monkeypatch)
+
+    sharded = glm_model.Glm5NextMLP(HIDDEN, INTERMEDIATE, "silu")
+    assert sharded.down_proj.tp_size == TP_SIZE
+    assert sharded.gate_up_proj.weight.shape == (2 * INTERMEDIATE // TP_SIZE, HIDDEN)
+
+    _assert_runs_locally_on_shard(
+        glm_model.Glm5NextMLP(HIDDEN, INTERMEDIATE, "silu", is_sequence_parallel=True)
+    )
+
+
+@pytest.mark.parametrize("use_sequence_parallel_moe", [True, False])
+def test_dense_layer_mlp_follows_the_sequence_parallel_layout(
+    monkeypatch, default_vllm_config, use_sequence_parallel_moe: bool
+):
+    _fake_tensor_parallel_world(monkeypatch)
+    monkeypatch.setattr(glm_model, "Glm5NextLinearAttention", _Attention)
+    monkeypatch.setattr(glm_model, "Glm5NextMLAAttention", _Attention)
+
+    config = Glm5NextTextConfig(
+        hidden_size=HIDDEN,
+        intermediate_size=INTERMEDIATE,
+        num_hidden_layers=2,
+        first_k_dense_replace=1,
+        n_routed_experts=4,
+        layer_types=["linear_attention", "deepseek_sparse_attention"],
+        mhc=False,
+    )
+    vllm_config = SimpleNamespace(
+        cache_config=None,
+        quant_config=None,
+        parallel_config=SimpleNamespace(
+            use_sequence_parallel_moe=use_sequence_parallel_moe
+        ),
+        kernel_config=SimpleNamespace(enable_jit_warmup=False),
+    )
+
+    layer = glm_model.Glm5NextDecoderLayer(vllm_config, config, layer_idx=0)
+
+    assert isinstance(layer.mlp, glm_model.Glm5NextMLP)
+    assert layer.self_attn.o_proj.reduce_results is not use_sequence_parallel_moe
+    if use_sequence_parallel_moe:
+        _assert_runs_locally_on_shard(layer.mlp)
+    else:
+        assert layer.mlp.down_proj.tp_size == TP_SIZE
diff --git a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
new file mode 100644
index 000000000000..9f94da459c15
--- /dev/null
+++ b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the GLM kpool indexer's top-k backend wiring."""
+
+import pytest
+import torch
+
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="CUDA-only dispatch"
+)
+
+
+def _require_deep_gemm() -> None:
+    from vllm.utils.deep_gemm import has_deep_gemm
+
+    if not has_deep_gemm():
+        pytest.skip("kpool indexer requires DeepGEMM")
+
+
+@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"])
+def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None:
+    """The kpool indexer must read kernel_config.sparse_indexer_topk_backend
+    and hand it to the shared SparseIndexerTopk dispatcher, rather than
+    hard-coding its own cooperative/persistent/per_row choice."""
+    _require_deep_gemm()
+    from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool
+
+    cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
+    with set_current_vllm_config(cfg):
+        op = SparseAttnIndexerKpool(
+            k_cache=None,
+            quant_block_size=128,
+            scale_fmt="ue8m0",
+            topk_tokens=2048,
+            head_dim=128,
+            max_pool_len=4096,
+            max_total_seq_len=8192,
+            topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"),
+        )
+    assert op.topk_backend == backend
diff --git a/tests/models/kimi_k3/test_amd_kda_direct_return.py b/tests/models/kimi_k3/test_amd_kda_direct_return.py
new file mode 100644
index 000000000000..e46efe7681a1
--- /dev/null
+++ b/tests/models/kimi_k3/test_amd_kda_direct_return.py
@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention
+from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer
+
+
+class _ReturningAttention(nn.Module):
+    def __init__(self, result: torch.Tensor) -> None:
+        super().__init__()
+        self.result = result
+
+    def forward(
+        self, hidden_states: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor:
+        return self.result
+
+
+class _TupleProjection(nn.Module):
+    def __init__(self, result: torch.Tensor) -> None:
+        super().__init__()
+        self.result = result
+
+    def forward(self, _: torch.Tensor) -> tuple[torch.Tensor, None]:
+        return self.result, None
+
+
+def _make_layer(self_attn: nn.Module) -> KimiDecoderLayer:
+    layer = object.__new__(KimiDecoderLayer)
+    nn.Module.__init__(layer)
+    layer.self_attn = self_attn
+    return layer
+
+
+def test_kda_self_attention_returns_projection_storage_directly() -> None:
+    hidden_states = torch.randn(4, 8)
+    projected = torch.randn_like(hidden_states)
+    layer = _make_layer(_ReturningAttention(projected))
+
+    output = layer._run_self_attn(torch.arange(4), hidden_states)
+
+    assert output is projected
+    assert output.data_ptr() == projected.data_ptr()
+
+
+def test_k3_kda_forward_returns_output_projection_storage() -> None:
+    hidden_states = torch.randn(4, 8)
+    projected = torch.randn_like(hidden_states)
+
+    def fill_core_output(**kwargs: torch.Tensor) -> None:
+        kwargs["core_attn_out"].fill_(1)
+
+    attention = SimpleNamespace(
+        local_projection_size=4,
+        head_dim=2,
+        local_num_heads=2,
+        in_proj_padding=0,
+        in_proj_qkvgfab=_TupleProjection(torch.randn(4, 20)),
+        f_b_proj=_TupleProjection(torch.randn(4, 4)),
+        o_proj=_TupleProjection(projected),
+        _forward=fill_core_output,
+    )
+
+    output = KimiK3DeltaAttention.forward(attention, hidden_states, torch.arange(4))
+
+    assert output is projected
+    assert output.data_ptr() == projected.data_ptr()
diff --git a/tests/models/kimi_k3/test_amd_mla_direct_return.py b/tests/models/kimi_k3/test_amd_mla_direct_return.py
new file mode 100644
index 000000000000..a1c120737df8
--- /dev/null
+++ b/tests/models/kimi_k3/test_amd_mla_direct_return.py
@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch import nn
+
+from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer, KimiMLAAttention
+
+
+class _DirectMLA(KimiMLAAttention):
+    def __init__(self, result: torch.Tensor) -> None:
+        nn.Module.__init__(self)
+        self.result = result
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.result
+
+
+def _make_layer(self_attn: nn.Module) -> KimiDecoderLayer:
+    layer = object.__new__(KimiDecoderLayer)
+    nn.Module.__init__(layer)
+    layer.self_attn = self_attn
+    return layer
+
+
+def test_mla_self_attention_returns_projection_storage_directly() -> None:
+    hidden_states = torch.randn(4, 8)
+    projected = torch.randn_like(hidden_states)
+    layer = _make_layer(_DirectMLA(projected))
+
+    output = layer._run_self_attn(torch.arange(4), hidden_states)
+
+    assert output is projected
+    assert output.data_ptr() == projected.data_ptr()
diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py
index 574f7a592500..366abe022159 100644
--- a/tests/models/kimi_k3/test_attn_res.py
+++ b/tests/models/kimi_k3/test_attn_res.py
@@ -17,6 +17,26 @@
 attn_res_module = importlib.import_module("vllm.models.kimi_k3.nvidia.ops.attn_res")
 
 
+def _on_rocm_below_10() -> bool:
+    if not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import get_rocm_version
+
+    return (get_rocm_version() or (0,)) < (10,)
+
+
+# The Triton bundled with ROCm < 10 (3.7.x) crashes in the AMD
+# CanonicalizePointers pass on the kernel's tl.where over pointer tensors. Only
+# the num_blocks > 0 loop contains it. ROCm serves AttnRes from
+# vllm/models/kimi_k3/amd/ops/attn_res.py instead, so nothing real is lost.
+_OLD_ROCM = _on_rocm_below_10()
+
+
+def _skip_on_old_rocm(num_blocks: int) -> None:
+    if _OLD_ROCM and num_blocks > 0:
+        pytest.skip("Triton on ROCm < 10 cannot compile this kernel's pointer select")
+
+
 def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
     storage = torch.randn(
         *shape[:-1],
@@ -79,6 +99,7 @@ def test_attn_res(
 ):
     if backend == "nvidia" and not current_platform.is_device_capability_family(100):
         pytest.skip("NVIDIA AttnRes requires the SM100 family")
+    _skip_on_old_rocm(num_blocks)
 
     prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
     delta = (
@@ -133,6 +154,7 @@ def test_attn_res(
 
 @pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1))
 def test_attn_res_block_counts(num_blocks: int):
+    _skip_on_old_rocm(num_blocks)
     prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
     blocks = torch.randn(
         1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
@@ -169,6 +191,7 @@ def test_attn_res_block_counts(num_blocks: int):
 
 
 def test_attn_res_without_output_norm():
+    _skip_on_old_rocm(MAX_BLOCKS)
     prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
     delta = torch.randn_like(prefix)
     blocks = torch.randn(
diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py
index 1fd4982e767e..f03d2cdcae3f 100644
--- a/tests/models/kimi_k3/test_kda.py
+++ b/tests/models/kimi_k3/test_kda.py
@@ -14,6 +14,10 @@
 
 from vllm import _custom_ops as ops
 from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+from vllm.model_executor.layers.mamba.checkpoint import MambaPrefillCheckpointMetadata
+from vllm.model_executor.layers.mamba.kda_checkpoint import (
+    FlashKDAPrefillCheckpointExporter,
+)
 from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update
 from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
     gather_initial_states,
@@ -25,7 +29,6 @@
 from vllm.models.kimi_k3.nvidia.kda import (
     _flashinfer_kda_prefill,
     _flashkda_prefill,
-    _store_cache_checkpoints_kernel,
     is_flashinfer_fused_kda_decode_supported,
     is_flashinfer_recurrent_kda_prefill_supported,
     is_flashkda_supported,
@@ -1421,14 +1424,21 @@ def test_flashinfer_kda_prefill_breakable_graph_cross_stream():
     torch.testing.assert_close(graph_value, torch.full_like(graph_value, 2))
 
 
+@pytest.mark.parametrize(
+    ("state_dtype", "tolerance"),
+    [
+        pytest.param(torch.bfloat16, 0.03, id="bf16"),
+        pytest.param(torch.float32, 0.01, id="fp32"),
+    ],
+)
 @torch.inference_mode()
-def test_flashkda_checkpoint_correctness():
+def test_flashkda_checkpoint_correctness(state_dtype: torch.dtype, tolerance: float):
     lower_bound = -3.0
-    _require_kda_prefill_backend("flashkda", torch.float32, lower_bound)
+    _require_kda_prefill_backend("flashkda", state_dtype, lower_bound)
 
     import vllm._flashkda_C  # noqa: F401
 
-    inputs = _make_kda_prefill_inputs(torch.float32, lower_bound=lower_bound)
+    inputs = _make_kda_prefill_inputs(state_dtype, lower_bound=lower_bound)
     q, k, v = inputs.q, inputs.k, inputs.v
     raw_g, raw_beta = inputs.raw_g, inputs.raw_beta
     A_log, dt_bias = inputs.A_log, inputs.dt_bias
@@ -1480,47 +1490,24 @@ def test_flashkda_checkpoint_correctness():
         checkpoint_offsets=checkpoint_offsets,
     )
 
-    assert_close("checkpoint_o", expected_out, checkpoint_out, 0.01)
-    assert_close("checkpoint_ht", expected_state, checkpoint_final_state, 0.01)
-    assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], 0.01)
+    assert_close("checkpoint_o", expected_out, checkpoint_out, tolerance)
+    assert_close("checkpoint_ht", expected_state, checkpoint_final_state, tolerance)
+    assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], tolerance)
 
     conv_state = torch.zeros(2, H * D, 3, dtype=q.dtype, device=DEVICE)
-    recurrent_storage = torch.zeros(2, H * D * D + 8, device=DEVICE)
+    recurrent_storage = torch.zeros(2, H * D * D + 8, dtype=state_dtype, device=DEVICE)
     recurrent_state = recurrent_storage[:, : H * D * D].view(2, H, D, D)
     conv_input = q[0].flatten(1)
     checkpoint_state_indices = torch.tensor(
         [1, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE
     )
-    state_len = conv_state.shape[-1]
-    width = H * D
-    recurrent_row_size = checkpoint_state[0].numel()
-    block_size = 256
-    _store_cache_checkpoints_kernel[
-        (
-            checkpoint_state_indices.numel(),
-            (max(width * state_len, recurrent_row_size) + block_size - 1) // block_size,
-        )
-    ](
-        conv_input,
-        conv_state,
-        checkpoint_state,
-        recurrent_state,
-        cu_seqlens,
-        checkpoint_offsets,
-        checkpoint_state_indices,
-        conv_input.stride(0),
-        conv_input.stride(1),
-        conv_state.stride(0),
-        conv_state.stride(1),
-        conv_state.stride(2),
-        checkpoint_state.stride(0),
-        recurrent_state.stride(0),
-        checkpoint_offsets.stride(0),
-        state_len,
-        width,
-        recurrent_row_size,
-        NULL_BLOCK_ID,
-        block_size,
+    FlashKDAPrefillCheckpointExporter().export(
+        MambaPrefillCheckpointMetadata(checkpoint_offsets, checkpoint_state_indices),
+        raw_qkv=conv_input,
+        conv_state=conv_state,
+        recurrent_checkpoint=checkpoint_state,
+        recurrent_state=recurrent_state,
+        cu_seqlens=cu_seqlens,
     )
     torch.testing.assert_close(conv_state[1], q[0, 13:16].flatten(1).transpose(0, 1))
     torch.testing.assert_close(recurrent_state[1], checkpoint_state[0])
diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py
index c3f423804a90..fa810217b188 100644
--- a/tests/models/kimi_k3/test_kda_metadata.py
+++ b/tests/models/kimi_k3/test_kda_metadata.py
@@ -934,3 +934,91 @@ def test_aligned_block_table_matches_shared_gdn():
     )
 
     torch.testing.assert_close(actual, expected)
+
+
+def _build_non_spec(batch, is_prefilling, full_cuda_graph=False):
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE
+    ).replace(
+        is_prefilling=None
+        if is_prefilling is None
+        else torch.tensor(is_prefilling, dtype=torch.bool)
+    )
+    builder = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=0,
+        full_cuda_graph=full_cuda_graph,
+    )
+    return builder, common_attn_metadata, builder.build(0, common_attn_metadata)
+
+
+def test_one_token_first_chunk_excludes_padding():
+    """Neither padding requests nor padding tokens count as prefill work."""
+    common = create_common_attn_metadata(
+        BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]),
+        BLOCK_SIZE,
+        DEVICE,
+    ).replace(
+        is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool),
+        num_actual_tokens=4,
+    )
+    builder = _make_builder(
+        KimiK3KDAMetadataBuilder, num_speculative_tokens=0, full_cuda_graph=False
+    )
+    actual = builder.build(0, common)
+
+    assert actual.num_decodes == 1
+    assert actual.num_prefills == 1
+    assert actual.num_decode_tokens == 1
+    assert actual.num_prefill_tokens == 1
+
+
+@pytest.mark.parametrize(
+    ("seq_len", "query_len", "is_prefilling", "num_prefills"),
+    [
+        pytest.param(1, 1, True, 1, id="first-chunk"),
+        pytest.param(65, 1, True, 0, id="resumed-chunk"),
+        pytest.param(0, 0, True, 0, id="padding"),
+        pytest.param(1, 1, None, 0, id="missing-prefill-flag"),
+    ],
+)
+def test_one_token_chunk_classification(
+    seq_len, query_len, is_prefilling, num_prefills
+):
+    """Only a real first chunk with a prefill flag needs state initialization."""
+    _, _, actual = _build_non_spec(
+        BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]),
+        is_prefilling=None if is_prefilling is None else [False, is_prefilling],
+    )
+
+    assert actual.num_prefills == num_prefills
+    assert actual.num_decodes == 2 - num_prefills
+    assert actual.num_prefill_tokens == num_prefills
+    assert actual.num_decode_tokens == 1 + query_len - num_prefills
+    if num_prefills:
+        assert actual.has_initial_state is not None
+        assert actual.has_initial_state.tolist() == [True, False]
+    else:
+        assert actual.has_initial_state is None
+
+
+def test_cudagraph_capture_batch_stays_decode_only():
+    """Capture rows have no history, but must still select decode kernels."""
+    batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4)
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE
+    ).replace(is_prefilling=torch.zeros(4, dtype=torch.bool))
+    builder = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=0,
+        full_cuda_graph=True,
+    )
+    actual = builder.build_for_cudagraph_capture(common_attn_metadata)
+
+    assert actual.num_prefills == 0
+    assert actual.num_decodes == 4
+    assert actual.has_initial_state is None
+    staged = actual.non_spec_state_indices_tensor
+    assert staged is not None
+    assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr()
+    torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0])
diff --git a/tests/models/kimi_k3/test_nvidia_kda_direct_return.py b/tests/models/kimi_k3/test_nvidia_kda_direct_return.py
new file mode 100644
index 000000000000..5b40d7ef6fa1
--- /dev/null
+++ b/tests/models/kimi_k3/test_nvidia_kda_direct_return.py
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch import nn
+
+from vllm.models.kimi_k3.nvidia.model import KimiDecoderLayer
+
+
+class _ReturningSharedKDA(nn.Module):
+    def __init__(self, result: torch.Tensor) -> None:
+        super().__init__()
+        self.result = result
+
+    def forward(
+        self, hidden_states: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor:
+        return self.result
+
+
+def test_low_rank_kda_caller_returns_projection_storage_directly() -> None:
+    hidden_states = torch.randn(4, 8)
+    projected = torch.randn_like(hidden_states)
+    layer = object.__new__(KimiDecoderLayer)
+    nn.Module.__init__(layer)
+    layer.self_attn = _ReturningSharedKDA(projected)
+
+    output = layer._run_self_attn(torch.arange(4), hidden_states)
+
+    assert output is projected
+    assert output.data_ptr() == projected.data_ptr()
diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py
index 33ff4e7ccb38..450e8badcebe 100644
--- a/tests/models/language/generation/test_common.py
+++ b/tests/models/language/generation/test_common.py
@@ -3,8 +3,6 @@
 
 import pytest
 import torch
-from packaging.version import Version
-from transformers import __version__ as TRANSFORMERS_VERSION
 
 from vllm.platforms import current_platform
 
@@ -26,7 +24,7 @@
     "meta-llama/Llama-3.2-1B-Instruct",
     "openbmb/MiniCPM3-4B",
     "Qwen/Qwen2.5-0.5B-Instruct",
-    "TitanML/tiny-mixtral",
+    "axolotl-ai-co/tiny-mixtral-30m",
     "Qwen/Qwen3-8B",
 ]
 
@@ -95,7 +93,7 @@
         pytest.param("stabilityai/stablelm-3b-4e1t"),  # stablelm
         pytest.param("bigcode/starcoder2-3b"),  # starcoder2
         pytest.param(
-            "TitanML/tiny-mixtral",  # mixtral
+            "axolotl-ai-co/tiny-mixtral-30m",  # mixtral
             marks=[pytest.mark.core_model],
         ),
         pytest.param("swiss-ai/Apertus-8B-Instruct-2509"),  # apertus
@@ -126,12 +124,6 @@ def test_models(
     model_info.check_available_online(on_fail="skip")
     model_info.check_transformers_version(on_fail="skip")
 
-    if current_platform.is_rocm() and model == "TitanML/tiny-mixtral":
-        # Its single-token router selects LLMM1, whose low-precision
-        # accumulation can change the top-2 experts. Keep the optimized kernel
-        # enabled generally, but use the reference GEMM for this accuracy test.
-        monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
-
     if use_rocm_aiter and (model in AITER_MODEL_LIST):
         monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
     elif use_rocm_aiter and model not in AITER_MODEL_LIST:
@@ -166,17 +158,9 @@ def test_models(
                 hf_model.model.device
             )
             if prompt_embeds is not None:
-                embed = hf_model.model.get_input_embeddings()(token_ids)
-
-                if "gemma" in model.lower() and (
-                    Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0")
-                ):
-                    # For Gemma 1/2 models with Transformers 5.4.0+, the prompt
-                    # embeddings are normalised in `get_prompt_embeddings`,
-                    # like Gemma 3. For older versions, we need to manually normalise.
-                    embed_scale = hf_model.config.hidden_size**0.5
-                    normalizer = torch.tensor(embed_scale, dtype=embed.dtype)
-                    embed *= normalizer
+                # Retained prompt values must not keep reference weights alive.
+                with torch.no_grad():
+                    embed = hf_model.model.get_input_embeddings()(token_ids)
 
                 # MiniCPM models apply scale_emb to embeddings internally.
                 # vLLM expects pre-scaled embeddings when using inputs_embeds.
diff --git a/tests/models/language/generation/test_gemma.py b/tests/models/language/generation/test_gemma.py
index c2c8e88e427c..5cffe6fb192e 100644
--- a/tests/models/language/generation/test_gemma.py
+++ b/tests/models/language/generation/test_gemma.py
@@ -14,7 +14,10 @@
     Gemma3nTextModel,
     _kv_sharing_weights_mapper,
 )
-from vllm.model_executor.models.gemma4 import Gemma4Model
+from vllm.model_executor.models.gemma4 import (
+    Gemma4ForCausalLM,
+    _gemma4_layer_weights_mapper,
+)
 
 MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"]
 
@@ -57,36 +60,70 @@ def __init__(self, *, vllm_config, prefix):
 
 
 @pytest.mark.cpu_test
-def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None:
-    """KV-shared layers have q_proj instead of a packed qkv_proj; their
-    redundant K/V tensors in original checkpoints have no parameter and are
-    skipped rather than failing the load."""
-    model = torch.nn.Module()
-    model.config = SimpleNamespace(num_experts=0)
-    model.start_layer, model.end_layer = 0, 2
-    model.layers = torch.nn.ModuleList([torch.nn.Module(), torch.nn.Module()])
-    for layer in model.layers:
-        layer.self_attn = torch.nn.Module()
-    model.layers[0].self_attn.qkv_proj = torch.nn.Linear(2, 6, bias=False)
-    model.layers[1].self_attn.q_proj = torch.nn.Linear(2, 2, bias=False)
-    shards: list[str] = []
-    model.layers[0].self_attn.qkv_proj.weight.weight_loader = (
-        lambda param, weight, shard_id: shards.append(shard_id)
+def test_gemma4_attention_mapper() -> None:
+    """Layers with a qkv_proj pack q/k/v; `attention_k_eq_v` full-attention
+    layers also load K as the V shard, and leave any v_proj such a checkpoint
+    ships unmapped so it fails the load rather than silently overwriting V;
+    KV-shared layers keep q_proj and drop the K/V tensors original checkpoints
+    still ship for them."""
+    config = SimpleNamespace(
+        num_hidden_layers=3,
+        num_kv_shared_layers=1,
+        attention_k_eq_v=True,
+        layer_types=["sliding_attention", "full_attention", "sliding_attention"],
     )
+    weights = [
+        (f"model.layers.{i}.self_attn.{tensor}.weight", torch.full((2, 2), i + 1.0))
+        for i in range(3)
+        for tensor in ("q_proj", "k_proj", "k_norm")
+    ] + [
+        ("model.layers.0.mlp.up_proj.weight", torch.empty(0)),
+        ("model.layers.1.self_attn.v_proj.weight", torch.empty(0)),
+    ]
 
-    q_weight = torch.full((2, 2), 2.0)
+    mapper = _gemma4_layer_weights_mapper(config)
+    mapped = list(mapper.apply(weights))
+
+    assert [(name, getattr(w, "shard_id", None)) for name, w in mapped] == [
+        ("model.layers.0.self_attn.qkv_proj.weight", "q"),
+        ("model.layers.0.self_attn.qkv_proj.weight", "k"),
+        ("model.layers.0.self_attn.k_norm.weight", None),
+        ("model.layers.1.self_attn.qkv_proj.weight", "q"),
+        ("model.layers.1.self_attn.qkv_proj.weight", "k"),
+        ("model.layers.1.self_attn.qkv_proj.weight", "v"),
+        ("model.layers.1.self_attn.k_norm.weight", None),
+        ("model.layers.2.self_attn.q_proj.weight", None),
+        ("model.layers.0.mlp.gate_up_proj.weight", 1),
+        ("model.layers.1.self_attn.v_proj.weight", None),
+    ]
+    k_weight, v_weight = weights[4][1], mapped[5][1]
+    assert torch.equal(v_weight, k_weight) and v_weight is not k_weight
+
+
+@pytest.mark.cpu_test
+def test_gemma4_expert_names_strip_language_model_prefix() -> None:
+    """The text-only path reuses the conditional wrapper's checkpoint naming,
+    so fused and per-expert tensors reach the experts under `model.*`."""
+    prefix = "model.language_model.layers.0."
     weights = [
-        (f"layers.{i}.self_attn.{tensor}.weight", q_weight)
-        for i in (0, 1)
-        for tensor in ("q_proj", "k_proj", "v_proj")
-    ] + [("layers.1.self_attn.k_norm.weight", torch.ones(2))]
-    loaded = Gemma4Model.load_weights(cast(Gemma4Model, model), weights)
-
-    assert shards == ["q", "k", "v"]
-    assert "layers.1.self_attn.q_proj.weight" in loaded
-    assert not any(name.startswith("layers.1.self_attn.k") for name in loaded)
-    assert not any(name.startswith("layers.1.self_attn.v") for name in loaded)
-    assert torch.equal(model.layers[1].self_attn.q_proj.weight, q_weight)
+        (prefix + name, torch.empty(0))
+        for name in (
+            "experts.gate_up_proj",
+            "experts.3.down_proj.weight_packed",
+            "router.per_expert_scale",
+        )
+    ]
+
+    mapped = [
+        (name, getattr(w, "shard_id", None))
+        for name, w in Gemma4ForCausalLM.hf_to_vllm_mapper.apply(weights)
+    ]
+
+    assert mapped == [
+        ("model.layers.0.experts.gate_up_proj", None),
+        ("model.layers.0.experts.3.down_proj.weight_packed", None),
+        ("model.layers.0.router.per_expert_scale", None),
+    ]
 
 
 @pytest.mark.cpu_test
diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py
index c5406fae48e0..b77670956aaf 100644
--- a/tests/models/language/generation/test_hybrid.py
+++ b/tests/models/language/generation/test_hybrid.py
@@ -164,8 +164,7 @@ def test_chunked_prefill_with_parallel_sampling(
     max_tokens: int,
     conv_state_layout: str,
 ) -> None:
-    """
-    Tests chunked prefill in conjunction with n > 1.
+    """Tests chunked prefill in conjunction with n > 1.
 
     In this case, prefill is populated with decoding tokens and
     we test that it doesn't fail.
@@ -199,8 +198,7 @@ def test_mamba_cache_cg_padding(
     max_tokens: int,
     conv_state_layout: str,
 ) -> None:
-    """
-    This test is for verifying that mamba cache is padded to CG captured
+    """This test is for verifying that mamba cache is padded to CG captured
     batch size. If it's not, a torch RuntimeError will be raised because
     tensor dimensions aren't compatible.
     """
@@ -236,8 +234,7 @@ def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
     example_prompts,
     model: str,
 ) -> None:
-    """
-    This test is for verifying that the hybrid inner state management doesn't
+    """This test is for verifying that the hybrid inner state management doesn't
     collapse in case where the number of incoming requests and
     finished_requests_ids is larger than the maximum mamba block capacity.
 
@@ -263,8 +260,7 @@ def test_state_cleanup(
     example_prompts,
     model: str,
 ) -> None:
-    """
-    This test is for verifying that the Hybrid state is cleaned up between
+    """This test is for verifying that the Hybrid state is cleaned up between
     steps.
 
     If it's not cleaned, an error would be expected.
diff --git a/tests/models/language/pooling/conftest.py b/tests/models/language/pooling/conftest.py
index 0571f1cd6d75..ceb21b5a409e 100644
--- a/tests/models/language/pooling/conftest.py
+++ b/tests/models/language/pooling/conftest.py
@@ -2,12 +2,42 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Pytest configuration for vLLM language generation tests."""
 
+import gc
+
+import pytest
 import torch
 
+from tests.utils import wait_for_memory_to_settle
+from vllm.logger import init_logger
 from vllm.platforms import current_platform
 
+logger = init_logger(__name__)
+
 
 def pytest_sessionstart(session):
     """Configure ROCm-specific settings before test session starts."""
     if current_platform.is_rocm():
         torch.set_float32_matmul_precision("high")
+
+
+@pytest.fixture(autouse=True)
+def release_gpu_memory_between_tests():
+    """Reclaim cached GPU memory between tests in this shard.
+
+    These tests start many vLLM/HF engines sequentially in one pytest
+    process; without explicit reclamation between tests, VRAM fragmentation
+    accumulates and later tests fail at engine startup ("Engine core
+    initialization failed").
+    """
+    yield
+    gc.collect()
+    if torch.accelerator.is_available():
+        torch.accelerator.empty_cache()
+    if current_platform.is_rocm():
+        return
+    try:
+        wait_for_memory_to_settle()
+    except ValueError as e:
+        # Longer-lived (class/module-scoped) engine fixtures may legitimately
+        # still hold VRAM; nothing further to reclaim in that case.
+        logger.info("Failed to clean GPU memory: %s", e)
diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py
index 75323d54bf8d..8829558a065d 100644
--- a/tests/models/language/pooling/test_colbert.py
+++ b/tests/models/language/pooling/test_colbert.py
@@ -128,9 +128,12 @@ def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device
     """Download and return the ColBERT linear projection weight."""
     from safetensors.torch import load_file
 
-    from vllm.transformers_utils.repo_utils import hf_api
+    from vllm.transformers_utils.repo_utils import hf_api, with_retry
 
-    path = hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"])
+    path = with_retry(
+        lambda: hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"]),
+        f"Error downloading {hf_spec['weights_file']} from {model_name}",
+    )
     weights = load_file(path)
     return weights[hf_spec["weights_key"]].to(device)
 
diff --git a/tests/models/language/pooling/test_jina_embeddings_v5.py b/tests/models/language/pooling/test_jina_embeddings_v5.py
index fbfebcc5cedb..21c82e141981 100644
--- a/tests/models/language/pooling/test_jina_embeddings_v5.py
+++ b/tests/models/language/pooling/test_jina_embeddings_v5.py
@@ -14,7 +14,7 @@
 from typing import cast
 
 import pytest
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import ModelConfig
 from vllm.model_executor.models.config import (
@@ -23,7 +23,7 @@
 )
 
 
-def _model_config(hf_config: PretrainedConfig) -> ModelConfig:
+def _model_config(hf_config: PreTrainedConfig) -> ModelConfig:
     """Minimal stand-in for ModelConfig; only hf_config is read."""
     return cast(ModelConfig, SimpleNamespace(hf_config=hf_config))
 
@@ -42,7 +42,7 @@ def test_encoder_backbone_enables_bidirectional_attention():
     EncoderOnlyAttention; JinaEmbeddingsV5Model then dispatches to the encoder
     implementation.
     """
-    hf_config = PretrainedConfig(is_decoder=False)
+    hf_config = PreTrainedConfig(is_decoder=False)
 
     JinaEmbeddingsV5ModelConfig.verify_and_update_model_config(_model_config(hf_config))
 
@@ -55,13 +55,13 @@ def test_supported_decoder_backbone_is_accepted():
 
     `-small` omits `is_decoder` entirely, so an absent attribute has to be
     treated as a decoder. The first assertion pins that assumption: if
-    PretrainedConfig ever gains an `is_decoder=False` default, this fails here
+    PreTrainedConfig ever gains an `is_decoder=False` default, this fails here
     rather than silently rejecting a supported checkpoint.
     """
-    absent = PretrainedConfig()
+    absent = PreTrainedConfig()
     assert not hasattr(absent, "is_decoder")
 
     JinaEmbeddingsV5ModelConfig.verify_and_update_model_config(_model_config(absent))
     JinaEmbeddingsV5ModelConfig.verify_and_update_model_config(
-        _model_config(PretrainedConfig(is_decoder=True))
+        _model_config(PreTrainedConfig(is_decoder=True))
     )
diff --git a/tests/models/language/pooling/test_max_tokens_per_doc.py b/tests/models/language/pooling/test_max_tokens_per_doc.py
index b6d6ba145478..b109b935dff6 100644
--- a/tests/models/language/pooling/test_max_tokens_per_doc.py
+++ b/tests/models/language/pooling/test_max_tokens_per_doc.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for max_tokens_per_doc and max_tokens_per_query.
-"""
+"""Tests for max_tokens_per_doc and max_tokens_per_query."""
 
 import json
 import os
@@ -80,10 +78,10 @@ class TestConfig:
         ],
         # 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,
+        without_truncated_prompt_tokens=(284, 285),
+        with_max_tokens_per_query_prompt_tokens=(154, 155),
+        with_max_tokens_per_doc_prompt_tokens=154,
+        with_max_tokens_per_query_and_doc_prompt_tokens=24,
     ),
     # 4. late-interaction
     TestConfig(
@@ -94,10 +92,10 @@ class TestConfig:
             "512",
             "--trust-remote-code",
         ],
-        without_truncated_prompt_tokens=285,
-        with_max_tokens_per_query_prompt_tokens=155,
-        with_max_tokens_per_doc_prompt_tokens=155,
-        with_max_tokens_per_query_and_doc_prompt_tokens=25,
+        without_truncated_prompt_tokens=284,
+        with_max_tokens_per_query_prompt_tokens=154,
+        with_max_tokens_per_doc_prompt_tokens=154,
+        with_max_tokens_per_query_and_doc_prompt_tokens=24,
     ),
     # 5. jinaai/jina-reranker-v3
     TestConfig(
diff --git a/tests/models/language/pooling/test_nomic_max_model_len.py b/tests/models/language/pooling/test_nomic_max_model_len.py
index 6ea29f6d0677..4ecbe77bc634 100644
--- a/tests/models/language/pooling/test_nomic_max_model_len.py
+++ b/tests/models/language/pooling/test_nomic_max_model_len.py
@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: SIM117
 
+from functools import partial
+
 import pytest
 
 from ...utils import EmbedModelInfo
@@ -26,6 +28,22 @@
 max_model_len = int(original_max_position_embeddings * factor)
 
 
+@pytest.fixture(scope="module", autouse=True)
+def predownload_models():
+    """Warm the HF cache so engine startup never hits transient hub errors."""
+    from vllm.transformers_utils.repo_utils import hf_api, with_retry
+
+    for model_info in MODELS:
+        with_retry(
+            partial(
+                hf_api().snapshot_download,
+                model_info.name,
+                revision=model_info.revision,
+            ),
+            f"Error downloading {model_info.name}",
+        )
+
+
 @pytest.mark.parametrize("model_info", MODELS)
 def test_default(model_info, vllm_runner):
     with vllm_runner(
diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py
index d8d813f404fb..c205a3c42a8a 100644
--- a/tests/models/language/pooling/test_token_classification.py
+++ b/tests/models/language/pooling/test_token_classification.py
@@ -115,6 +115,7 @@ def test_modernbert_models(
 
 @pytest.mark.parametrize("model", ["Davlan/xlm-roberta-base-ner-hrl"])
 @pytest.mark.parametrize("dtype", ["float"])
+@pytest.mark.flaky(reruns=3)
 @torch.inference_mode
 def test_xlm_roberta_models(
     hf_runner,
@@ -197,6 +198,7 @@ def test_openai_privacy_filter(
 
 @pytest.mark.parametrize("model", ["bd2lcco/Qwen3-0.6B-finetuned"])
 @pytest.mark.parametrize("dtype", ["float"])
+@pytest.mark.flaky(reruns=3)
 @torch.inference_mode
 def test_auto_conversion(
     hf_runner,
diff --git a/tests/models/language/pooling_mteb_test/conftest.py b/tests/models/language/pooling_mteb_test/conftest.py
new file mode 100644
index 000000000000..ce34d82563a8
--- /dev/null
+++ b/tests/models/language/pooling_mteb_test/conftest.py
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pytest configuration for vLLM pooling MTEB tests."""
+
+import gc
+
+import pytest
+import torch
+
+from tests.utils import wait_for_memory_to_settle
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+@pytest.fixture(autouse=True)
+def release_gpu_memory_between_tests():
+    """Reclaim cached GPU memory between tests in this shard.
+
+    These tests start many vLLM/HF engines sequentially in one pytest
+    process; without explicit reclamation between tests, VRAM fragmentation
+    accumulates and later tests fail at engine startup ("Engine core
+    initialization failed").
+    """
+    yield
+    gc.collect()
+    if torch.accelerator.is_available():
+        torch.accelerator.empty_cache()
+    if current_platform.is_rocm():
+        return
+    try:
+        wait_for_memory_to_settle()
+    except ValueError as e:
+        # Longer-lived (class/module-scoped) engine fixtures may legitimately
+        # still hold VRAM; nothing further to reclaim in that case.
+        logger.info("Failed to clean GPU memory: %s", e)
diff --git a/tests/models/language/pooling_mteb_test/mteb_embed_utils.py b/tests/models/language/pooling_mteb_test/mteb_embed_utils.py
index b9deecff8d1b..0423e94d2165 100644
--- a/tests/models/language/pooling_mteb_test/mteb_embed_utils.py
+++ b/tests/models/language/pooling_mteb_test/mteb_embed_utils.py
@@ -161,10 +161,14 @@ def mteb_test_embed_models(
     model_info: EmbedModelInfo,
     vllm_extra_kwargs=None,
     hf_model_callback=None,
-    atol=MTEB_EMBED_TOL,
+    atol: float | None = None,
     prompt_prefix: str | None = None,
     vllm_model_callback=None,
 ):
+    if atol is None:
+        atol = (
+            model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_EMBED_TOL
+        )
     vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
 
     # Test embed_dims, isnan and whether to use normalize
@@ -268,4 +272,8 @@ def mteb_test_embed_models(
 
     # We are not concerned that the vllm mteb results are better
     # than SentenceTransformers, so we only perform one-sided testing.
-    assert st_main_score - vllm_main_score < atol
+    diff = st_main_score - vllm_main_score
+    assert diff < atol, (
+        f"diff={diff:.6g} tol={atol} model={model_info.name} "
+        f"(st={st_main_score}, vllm={vllm_main_score})"
+    )
diff --git a/tests/models/language/pooling_mteb_test/mteb_score_utils.py b/tests/models/language/pooling_mteb_test/mteb_score_utils.py
index 16081cbe238b..2d83a2ebaab6 100644
--- a/tests/models/language/pooling_mteb_test/mteb_score_utils.py
+++ b/tests/models/language/pooling_mteb_test/mteb_score_utils.py
@@ -230,8 +230,12 @@ def mteb_test_rerank_models(
     hf_runner=HFMtebCrossEncoder,
     vllm_extra_kwargs=None,
     vllm_mteb_encoder=VllmMtebCrossEncoder,
-    atol=MTEB_RERANK_TOL,
+    atol: float | None = None,
 ):
+    if atol is None:
+        atol = (
+            model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_RERANK_TOL
+        )
     vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
 
     # Maybe load chat_template.
@@ -308,4 +312,8 @@ def mteb_test_rerank_models(
 
     # We are not concerned that the vllm mteb results are better
     # than SentenceTransformers, so we only perform one-sided testing.
-    assert st_main_score - vllm_main_score < atol
+    diff = st_main_score - vllm_main_score
+    assert diff < atol, (
+        f"diff={diff:.6g} tol={atol} model={model_info.name} "
+        f"(st={st_main_score}, vllm={vllm_main_score})"
+    )
diff --git a/tests/models/language/pooling_mteb_test/test_baai.py b/tests/models/language/pooling_mteb_test/test_baai.py
index ec11960fda07..a3a4ffaef196 100644
--- a/tests/models/language/pooling_mteb_test/test_baai.py
+++ b/tests/models/language/pooling_mteb_test/test_baai.py
@@ -17,6 +17,7 @@
         "BAAI/bge-base-en",
         architecture="BertModel",
         mteb_score=0.779336792,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -54,6 +55,7 @@
         "BAAI/bge-m3",
         architecture="XLMRobertaModel",
         mteb_score=0.787343078,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -82,6 +84,7 @@
         "BAAI/bge-reranker-base",
         architecture="XLMRobertaForSequenceClassification",
         mteb_score=0.32398,
+        mteb_tol=1e-2,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -101,6 +104,7 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
@@ -113,6 +117,7 @@ def test_embed_models_correctness(
     correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)
diff --git a/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py b/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
index 23bc95548bcd..b750b622b9fc 100644
--- a/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
+++ b/tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
@@ -26,6 +26,7 @@
             "method": "no_post_processing",
         },
         mteb_score=0.33757,
+        mteb_tol=1e-2,
         seq_pooling_type="LAST",
         attn_type="decoder",
         is_prefix_caching_supported=True,
@@ -127,6 +128,7 @@ def get_inputs(pairs, tokenizer, prompt=None):
         return torch.Tensor(scores)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(
diff --git a/tests/models/language/pooling_mteb_test/test_cross_encoder.py b/tests/models/language/pooling_mteb_test/test_cross_encoder.py
index 0d1067d5e2c5..8f5d86fa8d1f 100644
--- a/tests/models/language/pooling_mteb_test/test_cross_encoder.py
+++ b/tests/models/language/pooling_mteb_test/test_cross_encoder.py
@@ -17,6 +17,7 @@
         is_prefix_caching_supported=False,
         is_chunked_prefill_supported=False,
         mteb_score=0.32898,
+        mteb_tol=1e-2,
     ),
     RerankModelInfo(
         "tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
@@ -27,10 +28,12 @@
         is_chunked_prefill_supported=True,
         chat_template_name="qwen3_reranker.jinja",
         mteb_score=0.33459,
+        mteb_tol=1e-2,
     ),
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)
diff --git a/tests/models/language/pooling_mteb_test/test_gte.py b/tests/models/language/pooling_mteb_test/test_gte.py
index 0a54262e124f..c6b9ebc274fc 100644
--- a/tests/models/language/pooling_mteb_test/test_gte.py
+++ b/tests/models/language/pooling_mteb_test/test_gte.py
@@ -18,6 +18,7 @@
     EmbedModelInfo(
         "thenlper/gte-large",
         mteb_score=0.76807651,
+        mteb_tol=2e-3,
         architecture="BertModel",
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
@@ -44,6 +45,7 @@
         "Alibaba-NLP/gte-multilingual-base",
         architecture="GteNewModel",
         mteb_score=0.775074696,
+        mteb_tol=2e-3,
         hf_overrides={"architectures": ["GteNewModel"]},
         seq_pooling_type="CLS",
         attn_type="encoder_only",
@@ -79,6 +81,7 @@
     EmbedModelInfo(
         "Alibaba-NLP/gte-modernbert-base",
         mteb_score=0.748193353,
+        mteb_tol=2e-3,
         architecture="ModernBertModel",
         seq_pooling_type="CLS",
         attn_type="encoder_only",
@@ -90,6 +93,7 @@
     EmbedModelInfo(
         "Qwen/Qwen3-Embedding-0.6B",
         mteb_score=0.771163695,
+        mteb_tol=2e-3,
         architecture="Qwen3ForCausalLM",
         seq_pooling_type="LAST",
         attn_type="decoder",
@@ -109,6 +113,7 @@
         # classifier_pooling: mean
         "Alibaba-NLP/gte-reranker-modernbert-base",
         mteb_score=0.33386,
+        mteb_tol=1e-2,
         architecture="ModernBertForSequenceClassification",
         seq_pooling_type="CLS",
         attn_type="encoder_only",
@@ -119,6 +124,7 @@
     RerankModelInfo(
         "Alibaba-NLP/gte-multilingual-reranker-base",
         mteb_score=0.33062,
+        mteb_tol=1e-2,
         architecture="GteNewForSequenceClassification",
         hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
         seq_pooling_type="CLS",
@@ -130,6 +136,7 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
@@ -142,6 +149,7 @@ def test_embed_models_correctness(
     correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     vllm_extra_kwargs = {}
diff --git a/tests/models/language/pooling_mteb_test/test_intfloat.py b/tests/models/language/pooling_mteb_test/test_intfloat.py
index adadb60eeefe..805a3f91edb1 100644
--- a/tests/models/language/pooling_mteb_test/test_intfloat.py
+++ b/tests/models/language/pooling_mteb_test/test_intfloat.py
@@ -13,6 +13,7 @@
         "intfloat/e5-small",
         architecture="BertModel",
         mteb_score=0.742285423,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -29,6 +30,7 @@
         "intfloat/multilingual-e5-base",
         architecture="XLMRobertaModel",
         mteb_score=0.779325955,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -48,6 +50,7 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
diff --git a/tests/models/language/pooling_mteb_test/test_jina.py b/tests/models/language/pooling_mteb_test/test_jina.py
index 6c9308745e3f..753e6125101f 100644
--- a/tests/models/language/pooling_mteb_test/test_jina.py
+++ b/tests/models/language/pooling_mteb_test/test_jina.py
@@ -23,6 +23,7 @@
     EmbedModelInfo(
         "jinaai/jina-embeddings-v3",
         mteb_score=0.824413164,
+        mteb_tol=2e-3,
         architecture="XLMRobertaModel",
         is_matryoshka=True,
         seq_pooling_type="MEAN",
@@ -33,6 +34,7 @@
     EmbedModelInfo(
         "jinaai/jina-embeddings-v5-text-small",
         mteb_score=0.794535707854956,
+        mteb_tol=2e-3,
         architecture="JinaEmbeddingsV5Model",
         seq_pooling_type="LAST",
         attn_type="decoder",
@@ -42,6 +44,7 @@
     EmbedModelInfo(
         "jinaai/jina-embeddings-v5-text-nano",
         architecture="JinaEmbeddingsV5Model",
+        mteb_tol=2e-3,
         dtype="bfloat16" if current_platform.is_rocm() else "auto",
         seq_pooling_type="LAST",
         attn_type="encoder_only",
@@ -54,6 +57,7 @@
     RerankModelInfo(
         "jinaai/jina-reranker-v2-base-multilingual",
         mteb_score=0.33643,
+        mteb_tol=1e-2,
         architecture="XLMRobertaForSequenceClassification",
         seq_pooling_type="CLS",
         attn_type="encoder_only",
@@ -63,6 +67,7 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     task = "retrieval" if "v5" in model_info.name else "text-matching"
@@ -98,6 +103,7 @@ def hf_model_callback(model):
     )
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)
diff --git a/tests/models/language/pooling_mteb_test/test_modernbert_fp8.py b/tests/models/language/pooling_mteb_test/test_modernbert_fp8.py
index 8e062f0d76db..f97057f11fc1 100644
--- a/tests/models/language/pooling_mteb_test/test_modernbert_fp8.py
+++ b/tests/models/language/pooling_mteb_test/test_modernbert_fp8.py
@@ -56,6 +56,7 @@ def _assert_modernbert_online_fp8(model) -> None:
         )
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.skipif(
     _fp8_scaled_mm_unsupported(),
     reason="No FP8 ScaledMM kernel is available on this GPU type.",
diff --git a/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py b/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
index 74fe760e7839..1ad1f68b9adf 100644
--- a/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
+++ b/tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
@@ -30,6 +30,7 @@
         is_chunked_prefill_supported=True,
         chat_template_name="mxbai_rerank_v2.jinja",
         mteb_score=0.33651,
+        mteb_tol=1e-2,
         enable_test=True,
     ),
     RerankModelInfo(
@@ -104,6 +105,7 @@ def compute_logits(inputs):
         return torch.Tensor(scores)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info, hf_runner=MxbaiRerankerHfRunner)
diff --git a/tests/models/language/pooling_mteb_test/test_nemotron.py b/tests/models/language/pooling_mteb_test/test_nemotron.py
index 79fae2833990..711224a1a29f 100644
--- a/tests/models/language/pooling_mteb_test/test_nemotron.py
+++ b/tests/models/language/pooling_mteb_test/test_nemotron.py
@@ -19,6 +19,7 @@
         "nvidia/llama-nemotron-embed-1b-v2",
         architecture="LlamaBidirectionalModel",
         mteb_score=0.689164662128673,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -32,6 +33,7 @@
         architecture="LlamaBidirectionalForSequenceClassification",
         chat_template_name="nemotron-rerank.jinja",
         mteb_score=0.33994,
+        mteb_tol=1e-2,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -40,11 +42,13 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)
diff --git a/tests/models/language/pooling_mteb_test/test_nomic.py b/tests/models/language/pooling_mteb_test/test_nomic.py
index d7947f73f95e..38e34b78b024 100644
--- a/tests/models/language/pooling_mteb_test/test_nomic.py
+++ b/tests/models/language/pooling_mteb_test/test_nomic.py
@@ -17,6 +17,7 @@
         #  HF version and remove revision set.
         revision="720244025c1a7e15661a174c63cce63c8218e52b",
         mteb_score=0.737568559,
+        mteb_tol=2e-3,
         enable_test=True,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
@@ -35,6 +36,7 @@
         "nomic-ai/nomic-embed-text-v2-moe",
         architecture="NomicBertModel",
         mteb_score=0.715488912,
+        mteb_tol=2e-3,
         enable_test=True,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
@@ -44,6 +46,7 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
diff --git a/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py b/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
index 3c182cb046b5..10c90e733e6a 100644
--- a/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
+++ b/tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
@@ -32,6 +32,7 @@
         is_prefix_caching_supported=True,
         is_chunked_prefill_supported=True,
         mteb_score=0.33459,
+        mteb_tol=1e-2,
         enable_test=True,
     ),
     RerankModelInfo(
@@ -108,11 +109,13 @@ def compute_logits(inputs):
         return torch.Tensor(scores)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info, hf_runner=Qwen3RerankerHfRunner)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 @multi_gpu_test(num_gpus=2)
 def test_rerank_models_mteb_tp(vllm_runner, model_info: RerankModelInfo) -> None:
diff --git a/tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py b/tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py
index f3afbe84fa93..b68b2778b687 100644
--- a/tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py
+++ b/tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py
@@ -14,6 +14,7 @@
         is_matryoshka=False,
         architecture="BertModel",
         mteb_score=0.714927797,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -37,6 +38,7 @@
         is_matryoshka=False,
         architecture="NomicBertModel",
         mteb_score=0.681146831,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -54,6 +56,7 @@
         is_matryoshka=True,
         architecture="BertModel",
         mteb_score=0.649088363,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -65,6 +68,7 @@
         is_matryoshka=True,
         architecture="XLMRobertaModel",
         mteb_score=0.712258299,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -76,6 +80,7 @@
         is_matryoshka=True,
         architecture="GteModel",
         mteb_score=0.706622444,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -85,6 +90,7 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
diff --git a/tests/models/language/pooling_mteb_test/test_st_projector.py b/tests/models/language/pooling_mteb_test/test_st_projector.py
index 395846347fb3..7397db8535ff 100644
--- a/tests/models/language/pooling_mteb_test/test_st_projector.py
+++ b/tests/models/language/pooling_mteb_test/test_st_projector.py
@@ -14,6 +14,7 @@
         "TencentBAC/Conan-embedding-v1",
         architecture="BertModel",
         mteb_score=0.688611955,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -24,6 +25,7 @@
         "google/embeddinggemma-300m",
         architecture="Gemma3TextModel",
         mteb_score=0.7473819294684156,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -33,6 +35,7 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
diff --git a/tests/models/language/pooling_mteb_test/test_voyage.py b/tests/models/language/pooling_mteb_test/test_voyage.py
index 99ef1de9adfd..628ea31ca0e9 100644
--- a/tests/models/language/pooling_mteb_test/test_voyage.py
+++ b/tests/models/language/pooling_mteb_test/test_voyage.py
@@ -22,6 +22,7 @@
             "num_labels": 2048,
         },
         mteb_score=0.7054,
+        mteb_tol=2e-3,
         # === MTEB Results ===
         # STS12: 0.6613
         # STS13: 0.6906
@@ -34,6 +35,7 @@
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     # Encoder-only attention models need enforce_eager=True to avoid
diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py
index c364bab3828f..269dad05f403 100644
--- a/tests/models/multimodal/generation/test_common.py
+++ b/tests/models/multimodal/generation/test_common.py
@@ -151,10 +151,12 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model):
         img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>",
         video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>",
         enforce_eager=False,
+        needs_video_metadata=True,
         max_model_len=4096,
         max_num_seqs=2,
         auto_cls=AutoModelForImageTextToText,
         vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output,
+        patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner,
         image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
         marks=[pytest.mark.core_model, pytest.mark.cpu_model],
     ),
@@ -317,8 +319,6 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model):
         img_idx_to_prompt=lambda idx: "",
         auto_cls=AutoModelForImageTextToText,
         vllm_output_post_proc=model_utils.blip2_vllm_to_hf_output,
-        # FIXME: https://github.com/huggingface/transformers/pull/38510
-        marks=[pytest.mark.skip("Model is broken")],
     ),
     "cosmos3": VLMTestInfo(
         models=["nvidia/Cosmos3-Nano"],
@@ -819,10 +819,6 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model):
         patch_hf_runner=model_utils.paddleocr_vl_patch_hf_runner,
         image_size_factors=[(0.25,)],
         marks=[
-            pytest.mark.skipif(
-                Version(TRANSFORMERS_VERSION) == Version("4.57.3"),
-                reason="This model is broken in Transformers v4.57.3",
-            ),
             pytest.mark.skipif(
                 Version(TRANSFORMERS_VERSION) >= Version("5.0.0"),
                 reason="Model's custom code uses ROPE_INIT_FUNCTIONS"
diff --git a/tests/models/multimodal/generation/test_interleaved.py b/tests/models/multimodal/generation/test_interleaved.py
index a773db19825e..07ce235853b6 100644
--- a/tests/models/multimodal/generation/test_interleaved.py
+++ b/tests/models/multimodal/generation/test_interleaved.py
@@ -23,11 +23,9 @@ def base_prompt(modalities_str: str) -> str:
 @pytest.mark.parametrize("dtype", ["float16"])
 @pytest.mark.parametrize("max_tokens", [128])
 def test_models(vllm_runner, model, dtype: str, max_tokens: int) -> None:
-    """
-    This is a simple test to check if interleaved and non-interleaved prompts
+    """This is a simple test to check if interleaved and non-interleaved prompts
     give the same result.
     """
-
     image_cherry = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
     image_stop = convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB")
     images = [image_cherry, image_stop]
diff --git a/tests/models/multimodal/generation/test_maverick.py b/tests/models/multimodal/generation/test_maverick.py
index ff6e523e5b25..0f6bec107fc2 100644
--- a/tests/models/multimodal/generation/test_maverick.py
+++ b/tests/models/multimodal/generation/test_maverick.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Create a reduced-layer version of the Maverick model for testing purposes.
+"""Create a reduced-layer version of the Maverick model for testing purposes.
 
 This script creates a new model with fewer layers by:
 1. Loading the original Maverick model configuration
@@ -39,7 +38,6 @@ def run_maverick_serving(model: str):
     """Test Llama-4-Maverick model with vLLM LLM class using CLI equivalent
     options with reduced layers.
     """
-
     try:
         sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
@@ -71,8 +69,7 @@ def run_maverick_serving(model: str):
 
 
 def get_rope_layers_config(model_path: str) -> list[int]:
-    """
-    Get the interleaved RoPE configuration from HuggingFace config
+    """Get the interleaved RoPE configuration from HuggingFace config.
 
     Args:
         model_path: Path to the local directory containing the reduced
@@ -81,6 +78,7 @@ def get_rope_layers_config(model_path: str) -> list[int]:
     Returns:
         List of 0 or 1 indicating whether each layer uses RoPE and local attn
         0 indicates that RoPE is not used while 1 indicates that RoPE is used.
+
     """
     config_path = Path(model_path) / "config.json"
     model_config = json.loads(config_path.read_text())
@@ -98,8 +96,7 @@ def create_reduced_maverick_model(
     vision_layers: int = 2,
     force_recreate: bool = False,
 ) -> str:
-    """
-    Create a reduced-layer version of the Maverick model.
+    """Create a reduced-layer version of the Maverick model.
 
     Args:
         original_model_name: Name of the original Maverick model
@@ -111,8 +108,8 @@ def create_reduced_maverick_model(
 
     Returns:
         Path to the created reduced model directory
-    """
 
+    """
     print(
         f"Creating reduced Maverick model with {text_layers} text layers and "
         f"{vision_layers} vision layers..."
@@ -178,7 +175,6 @@ def create_reduced_config(
     original_config: Any, text_layers: int, num_experts: int, vision_layers: int
 ) -> dict[str, Any]:
     """Create a reduced configuration based on the original."""
-
     # Convert config to dictionary
     config_dict = original_config.to_dict()
 
@@ -220,7 +216,6 @@ def create_reduced_config(
 
 def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None:
     """Copy tokenizer files from the original model."""
-
     try:
         tokenizer = AutoTokenizer.from_pretrained(
             original_model_name, trust_remote_code=True
@@ -233,7 +228,6 @@ def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None:
 
 def create_preprocessor_config(original_config: Any, output_path: Path) -> None:
     """Create preprocessor_config.json for multimodal model."""
-
     # Try to load the original preprocessor config
     try:
         processor = AutoProcessor.from_pretrained(
@@ -253,7 +247,6 @@ def create_reduced_safetensors(
     original_config: Any, reduced_config: dict[str, Any], output_path: Path
 ) -> None:
     """Create safetensors files with weights for the reduced model."""
-
     print("Generating synthetic weights for reduced model...")
 
     text_config = reduced_config["text_config"]
@@ -276,7 +269,6 @@ def create_reduced_safetensors(
 
 def create_text_model_weights(text_config: dict[str, Any]) -> dict[str, torch.Tensor]:
     """Create synthetic weights for the text model with MoE structure."""
-
     weights = {}
 
     vocab_size = text_config["vocab_size"]
@@ -405,7 +397,6 @@ def create_vision_model_weights(
     vision_config: dict[str, Any],
 ) -> dict[str, torch.Tensor]:
     """Create synthetic weights for the vision model."""
-
     weights = {}
 
     hidden_size = vision_config["hidden_size"]
@@ -473,8 +464,7 @@ def create_vision_model_weights(
 def create_shared_weights(
     text_config: dict[str, Any], vision_config: dict[str, Any]
 ) -> dict[str, torch.Tensor]:
-    """Create weights for shared components (vision-language connector)"""
-
+    """Create weights for shared components (vision-language connector)."""
     weights = {}
 
     text_hidden_size = text_config["hidden_size"]
@@ -492,7 +482,6 @@ def save_weights_to_safetensors(
     weights: dict[str, torch.Tensor], output_path: Path
 ) -> None:
     """Save weights to safetensors files and create index."""
-
     # Determine how to shard the weights
     max_shard_size = 5 * 1024 * 1024 * 1024  # 5GB per shard
 
@@ -616,6 +605,8 @@ def test_dummy_maverick(
 ) -> None:
     # Disable multiprocessing allows us to access model executor from LLM engine
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    # The forked test process may inherit CUDA state from test collection.
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
 
     model_path = create_reduced_maverick_model(
         original_model_name=original_model_name,
@@ -653,7 +644,6 @@ def test_dummy_maverick(
 
 def main():
     """Main function to create and test the reduced model."""
-
     import argparse
 
     parser = argparse.ArgumentParser(
diff --git a/tests/models/multimodal/generation/test_nemotron_parse.py b/tests/models/multimodal/generation/test_nemotron_parse.py
index db0d3d06840a..69e41b21e342 100644
--- a/tests/models/multimodal/generation/test_nemotron_parse.py
+++ b/tests/models/multimodal/generation/test_nemotron_parse.py
@@ -31,8 +31,7 @@ def mask_bbox_tokens(
     output: tuple[list[int], str, SampleLogprobs],
     tokenizer: TokenizerLike,
 ) -> tuple[list[int], str, SampleLogprobs]:
-    """
-    Always pass check_logprobs_close check for bounding box tokens
+    """Always pass check_logprobs_close check for bounding box tokens
     because it is reasonable for them to differ slightly.
     """
     ignore_pattern = r"<[xy]_[\d.]+>"
diff --git a/tests/models/multimodal/generation/test_qwen2_5_vl.py b/tests/models/multimodal/generation/test_qwen2_5_vl.py
index 05a9b0299839..975a8b339d85 100644
--- a/tests/models/multimodal/generation/test_qwen2_5_vl.py
+++ b/tests/models/multimodal/generation/test_qwen2_5_vl.py
@@ -4,10 +4,10 @@
 import pytest
 
 from vllm.assets.image import ImageAsset
-from vllm.multimodal.video import sample_frames_from_video
 from vllm.platforms import current_platform
 
 from ....conftest import VIDEO_ASSETS
+from .vlm_utils.builders import sample_frames_with_video_metadata
 
 models = ["Qwen/Qwen2.5-VL-3B-Instruct"]
 target_dtype = "bfloat16"
@@ -81,7 +81,9 @@ def test_qwen2_5_vl_evs_functionality(
 
     # Sample frames from video assets
     sampled_vids = [
-        sample_frames_from_video(asset.np_ndarrays, num_frames)
+        sample_frames_with_video_metadata(
+            (asset.np_ndarrays, asset.metadata), num_frames
+        )
         for asset in video_assets
     ]
 
@@ -145,7 +147,9 @@ def test_qwen2_5_vl_evs_batched_videos(
     monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
     # Sample frames from video assets
     sampled_vids = [
-        sample_frames_from_video(asset.np_ndarrays, num_frames)
+        sample_frames_with_video_metadata(
+            (asset.np_ndarrays, asset.metadata), num_frames
+        )
         for asset in video_assets
     ]
 
diff --git a/tests/models/multimodal/generation/test_qwen2_vl.py b/tests/models/multimodal/generation/test_qwen2_vl.py
index e163a25eb82f..4508135b2bb9 100644
--- a/tests/models/multimodal/generation/test_qwen2_vl.py
+++ b/tests/models/multimodal/generation/test_qwen2_vl.py
@@ -97,7 +97,7 @@ def batch_make_image_embeddings(
     processor,
     llm: VllmRunner,
 ) -> list[Qwen2VLPromptImageEmbeddingInput]:
-    """batched image embeddings for Qwen2-VL
+    """Batched image embeddings for Qwen2-VL.
 
     This will infer all images' embeddings in a single batch,
       and split the result according to input batches.
@@ -108,7 +108,6 @@ def batch_make_image_embeddings(
 
     returns: `list[Qwen2VLPromptImageEmbeddingInput]`
     """
-
     image_batches_: list[Any] = image_batches[:]
 
     # convert single-image batches to multiple-image batches
@@ -183,7 +182,7 @@ def get_image_embeds(model):
 def batch_make_video_embeddings(
     video_batches: PromptVideoInput, processor, llm: VllmRunner
 ) -> list[Qwen2VLPromptVideoEmbeddingInput]:
-    """batched video embeddings for Qwen2-VL
+    """Batched video embeddings for Qwen2-VL.
 
     A NDArray represents a single video's all frames.
 
@@ -194,7 +193,6 @@ def batch_make_video_embeddings(
       - Single-video batches: `list[NDArray]`
       - Multiple-video batches: `list[list[NDArray]]`
     """
-
     video_batches_: list[Any] = video_batches[:]
 
     for idx in range(len(video_batches_)):
diff --git a/tests/models/multimodal/generation/test_ultravox.py b/tests/models/multimodal/generation/test_ultravox.py
index 964454331232..72f473835800 100644
--- a/tests/models/multimodal/generation/test_ultravox.py
+++ b/tests/models/multimodal/generation/test_ultravox.py
@@ -204,7 +204,6 @@ def test_variable_length_audio_batching(
 @pytest.mark.asyncio
 async def test_online_serving(client, audio_assets: AudioTestAssets):
     """Exercises online serving with/without chunked prefill enabled."""
-
     messages = [
         {
             "role": "user",
diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py
index f8136ae8f3a7..da472a164dab 100644
--- a/tests/models/multimodal/generation/test_vit_cudagraph.py
+++ b/tests/models/multimodal/generation/test_vit_cudagraph.py
@@ -55,6 +55,10 @@ def kimi_vl_chat_template(content: str) -> str:
     )
 
 
+def deepseek_v41_chat_template(content: str) -> str:
+    return f"<｜begin▁of▁sentence｜><｜User｜>{content}<｜Assistant｜></think>"
+
+
 def step3_vl_chat_template(content: str) -> str:
     return (
         "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n "
@@ -140,7 +144,7 @@ def minicpmv_chat_template(content: str) -> str:
             "<|vision_start|><|video_pad|><|vision_end|>"
             "Describe this video in one sentence."
         ),
-        needs_video_metadata=False,
+        needs_video_metadata=True,
         marks=[pytest.mark.core_model],
     ),
     "kimi_vl": VitCudagraphTestConfig(
@@ -221,6 +225,29 @@ def minicpmv_chat_template(content: str) -> str:
         vllm_runner_kwargs={"gpu_memory_utilization": 0.80},
         marks=[pytest.mark.core_model],
     ),
+    "deepseek_v41": VitCudagraphTestConfig(
+        model="deepseek-ai/DeepSeek-V4.1-Flash",
+        modalities=["image"],
+        image_prompt=deepseek_v41_chat_template(
+            "<｜deepseek_image｜>\n\nWhat is in this image?"
+        ),
+        compilation_config_overrides={
+            "encoder_cudagraph_token_budgets": [1024],
+        },
+        vllm_runner_kwargs={
+            "load_format": "dummy",
+            "attention_backend": (
+                "ROCM_FLASHMLA_SPARSE_DSV4"
+                if current_platform.is_rocm()
+                else "FLASHMLA_SPARSE_DSV41"
+            ),
+            "hf_overrides": partial(
+                dummy_hf_overrides,
+                model_arch="DeepseekV41ForCausalLM",
+                exist_overrides={"vision_n_layers": 1},
+            ),
+        },
+    ),
     "step3_vl": VitCudagraphTestConfig(
         model="stepfun-ai/Step3-VL-10B",
         modalities=["image"],
@@ -499,3 +526,118 @@ def test_vit_cudagraph_video(model_id, vllm_runner, video_assets):
 
         # Ensure the output is a string
         assert isinstance(output_text, str)
+
+
+def _check_eonly_encoder_outputs(worker, batches):
+    """Compare real E-only replay with eager, retaining outputs across replay."""
+    import torch
+
+    from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+
+    runner = worker.model_runner
+    assert isinstance(runner, MMEncoderModelRunner)
+    manager = runner.model_state.encoder_runner.cudagraph_manager
+    assert manager is not None and manager.is_captured()
+    retained: list[tuple[torch.Tensor, torch.Tensor]] = []
+    max_error = 0.0
+    stats = []
+    with torch.inference_mode():
+        for batch in batches:
+            kwargs = {key: value.to(runner.device) for key, value in batch.items()}
+            actual = manager.execute(kwargs)
+            expected = runner.model.embed_multimodal(**kwargs)
+            assert len(actual) == len(expected)
+            for output, reference in zip(actual, expected):
+                assert torch.isfinite(output).all()
+                # Two BF16 ulps at unit scale, fixed before model validation.
+                torch.testing.assert_close(output, reference, rtol=0.016, atol=0.016)
+                max_error = max(max_error, (output - reference).abs().max().item())
+            for output, snapshot in retained:
+                torch.testing.assert_close(output, snapshot, rtol=0, atol=0)
+            retained.extend((output, output.clone()) for output in actual)
+            stats.append(manager.get_cumulative_stats())
+    assert stats[0]["graph_hits"] > 0
+    assert stats[1]["graph_hits"] > stats[0]["graph_hits"]
+    assert stats[2]["graph_misses"] > stats[1]["graph_misses"]
+    assert stats[3]["graph_hits"] > stats[2]["graph_hits"]
+    assert stats[4]["graph_hits"] > stats[3]["graph_hits"]
+    assert stats[5]["graph_hits"] > stats[4]["graph_hits"]
+    assert stats[6]["graph_hits"] > stats[5]["graph_hits"]
+    return {"max_abs_error": max_error, "stats": stats}
+
+
+@pytest.mark.parametrize(
+    "model_id",
+    params_with_marks({key: MODEL_CONFIGS[key] for key in ("qwen3_vl", "qwen3_5")}),
+)
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA graphs")
+def test_eonly_vit_cudagraph_outputs(
+    model_id, vllm_runner, image_assets, video_assets, monkeypatch
+):
+    """Exercise startup capture, image/video replay, fallback and output ownership."""
+    import numpy as np
+    from PIL import Image
+    from transformers import AutoProcessor
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    config = MODEL_CONFIGS[model_id]
+    token_budgets = [64, 256, 2048]
+    processor = AutoProcessor.from_pretrained(config.model)
+    first, second = [asset.pil_image for asset in image_assets]
+    batches = [
+        [first.resize((224, 224))],
+        [second.resize((448, 224)), first.resize((224, 224))],
+        [second.resize((1792, 1792))],
+        [second.resize((224, 224))],
+        [first.resize((1280, 720))],
+        [first.resize((224, 224))],
+    ]
+    frames = sample_frames_from_video(video_assets[0].np_ndarrays, 2)
+    video = np.stack(
+        [np.asarray(Image.fromarray(frame).resize((224, 224))) for frame in frames]
+    )
+    with vllm_runner(
+        config.model,
+        dtype=config.dtype,
+        mm_encoder_only=True,
+        # FA padding NaNs are tracked separately in #57136.
+        mm_encoder_attn_backend="FLASHINFER",
+        enable_prefix_caching=False,
+        max_model_len=4096,
+        max_num_seqs=2,
+        limit_mm_per_prompt={"image": 2, "video": 1},
+        compilation_config={
+            "cudagraph_mode": "NONE",
+            "cudagraph_mm_encoder": True,
+            "encoder_cudagraph_token_budgets": token_budgets,
+            "encoder_cudagraph_max_vision_items_per_batch": 2,
+            "encoder_cudagraph_max_frames_per_batch": 2,
+        },
+    ) as model:
+        mm_config = model.llm.llm_engine.vllm_config.model_config.multimodal_config
+        processor_kwargs = mm_config.merge_mm_processor_kwargs({})
+        inputs = [
+            dict(
+                processor.image_processor(
+                    images=images, return_tensors="pt", **processor_kwargs
+                )
+            )
+            for images in batches
+        ]
+        inputs.append(
+            dict(
+                processor.video_processor(
+                    videos=[video],
+                    do_sample_frames=False,
+                    return_tensors="pt",
+                    **processor_kwargs,
+                )
+            )
+        )
+        results = model.llm.collective_rpc(
+            partial(_check_eonly_encoder_outputs, batches=inputs)
+        )
+        for result in results:
+            assert result["stats"][0]["token_budgets"] == token_budgets
+        print(f"E-only {model_id}: {results}")
diff --git a/tests/models/multimodal/generation/test_voxtral.py b/tests/models/multimodal/generation/test_voxtral.py
index a6e8f3ff18a4..544c06020f66 100644
--- a/tests/models/multimodal/generation/test_voxtral.py
+++ b/tests/models/multimodal/generation/test_voxtral.py
@@ -78,7 +78,6 @@ def test_online_serving(vllm_runner, audio_assets: AudioTestAssets):
 
     Steps run sequentially so each releases the GPU before the next starts.
     """
-
     question = f"What's happening in these {len(audio_assets)} audio clips?"
     max_tokens = 10
     audio_data = [asset.audio_and_sample_rate for asset in audio_assets]
diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py
index 0ef3c33aebc0..2c13e07ff367 100644
--- a/tests/models/multimodal/generation/test_voxtral_realtime.py
+++ b/tests/models/multimodal/generation/test_voxtral_realtime.py
@@ -14,12 +14,14 @@
 
 from vllm import LLM, SamplingParams
 from vllm.assets.audio import AudioAsset
+from vllm.config import CUDAGraphMode
 from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.kv_cache_interface import SlidingWindowSpec
 
-from ....utils import ROCM_ENGINE_KWARGS
+from ....utils import ROCM_ENGINE_KWARGS, create_new_process_for_each_test
 
 MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
 AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn"
@@ -130,6 +132,45 @@ async def async_engine():
         wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
 
 
+def _build_inputs(
+    audio_assets: list[AudioAsset], tokenizer: MistralTokenizer
+) -> tuple[list[dict], list[SamplingParams]]:
+    audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
+    inputs = []
+    sampling_params = []
+    for audio_asset in audio_assets:
+        audio = Audio.from_file(audio_asset.get_local_path(), strict=False)
+        req = TranscriptionRequest(
+            audio=audio.to_base64(audio.format),
+            streaming=StreamingMode.OFFLINE,
+            language=None,
+        )
+        tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
+        tokens = tokenized.tokens
+        audio_array = tokenized.audios[0].audio_array
+
+        num_samples = audio_array.shape[0]
+        max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1
+        sampling_params.append(SamplingParams(temperature=0.0, max_tokens=max_tokens))
+        inputs.append(
+            {
+                "multi_modal_data": {"audio": [(audio_array, None)]},
+                "prompt_token_ids": tokens,
+            }
+        )
+    return inputs, sampling_params
+
+
+def _assert_expected_text(outputs) -> None:
+    texts = _normalize([out.outputs[0].text for out in outputs])
+    for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
+        assert got == expected, (
+            f"Output mismatch at index {i}:\n"
+            f"  got:      {got!r}\n"
+            f"  expected: {expected!r}"
+        )
+
+
 def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch):
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
 
@@ -138,51 +179,39 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypa
 
     with vllm_runner(**vllm_kwargs) as vllm_model:
         assert_encoder_kv_cache_spec(vllm_model.llm)
-        audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
-
-        def from_file(file_path: str):
-            audio = Audio.from_file(file_path, strict=False)
-            req = TranscriptionRequest(
-                audio=audio.to_base64(audio.format),
-                streaming=StreamingMode.OFFLINE,
-                language=None,
-            )
-            tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
+        inputs, sampling_params = _build_inputs(audio_assets, tokenizer)
 
-            return (tokenized.tokens, tokenized.audios[0].audio_array)
-
-        tokenized_list = [
-            from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
-        ]
+        outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params)
+        _assert_expected_text(outputs)
 
-        inputs = []
-        sampling_params = []
 
-        for tokens, audio_array in tokenized_list:
-            num_samples = audio_array.shape[0]
-            max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1
-            sampling_params.append(
-                SamplingParams(temperature=0.0, max_tokens=max_tokens)
-            )
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="Full cudagraphs require CUDA"
+)
+@pytest.mark.parametrize(
+    "cudagraph_mode",
+    [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE],
+)
+@create_new_process_for_each_test(method="spawn")
+def test_voxtral_realtime_cudagraph(
+    audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode
+):
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
 
-            input_dict = {
-                "multi_modal_data": {"audio": [(audio_array, None)]},
-                "prompt_token_ids": tokens,
-            }
-            inputs.append(input_dict)
+    vllm_kwargs = {**ENGINE_CONFIG}
+    vllm_kwargs["model_name"] = vllm_kwargs.pop("model")
+    vllm_kwargs["enforce_eager"] = False
+    vllm_kwargs["compilation_config"] = {
+        "cudagraph_mode": cudagraph_mode,
+        # max_num_seqs is 4, so this is the only decode size worth capturing.
+        "cudagraph_capture_sizes": [4],
+    }
 
-        outputs = vllm_model.llm.generate(
-            inputs,
-            sampling_params=sampling_params,
-        )
+    with vllm_runner(**vllm_kwargs) as vllm_model:
+        inputs, sampling_params = _build_inputs(audio_assets, tokenizer)
 
-        texts = _normalize([out.outputs[0].text for out in outputs])
-        for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
-            assert got == expected, (
-                f"Output mismatch at index {i}:\n"
-                f"  got:      {got!r}\n"
-                f"  expected: {expected!r}"
-            )
+        outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params)
+        _assert_expected_text(outputs)
 
 
 @pytest.mark.asyncio
diff --git a/tests/models/multimodal/generation/vlm_utils/case_filtering.py b/tests/models/multimodal/generation/vlm_utils/case_filtering.py
index cbe660425751..39d9d119bfa4 100644
--- a/tests/models/multimodal/generation/vlm_utils/case_filtering.py
+++ b/tests/models/multimodal/generation/vlm_utils/case_filtering.py
@@ -156,6 +156,7 @@ def get_wrapped_test_sizes(
     Args:
         test_info: Test configuration to be expanded.
         test_type: The type of test being filtered for.
+
     """
     # If it is an embedding test, we always use the EMBEDDING_SIZE_FACTORS
     if test_type == VLMTestType.EMBEDDING:
diff --git a/tests/models/multimodal/generation/vlm_utils/core.py b/tests/models/multimodal/generation/vlm_utils/core.py
index 207d3a3202a1..d665b6ad353b 100644
--- a/tests/models/multimodal/generation/vlm_utils/core.py
+++ b/tests/models/multimodal/generation/vlm_utils/core.py
@@ -200,7 +200,7 @@ def process_runner_outputs(
 
 
 def process_outputs(output_processor, model, outputs_per_image):
-    """Applies a model specific post-processor function to a runner's output"""
+    """Applies a model specific post-processor function to a runner's output."""
     return [
         [output_processor(res, model) for res in outputs]
         for outputs in outputs_per_image
diff --git a/tests/models/multimodal/generation/vlm_utils/custom_inputs.py b/tests/models/multimodal/generation/vlm_utils/custom_inputs.py
index 84109233685b..11ef99ea6443 100644
--- a/tests/models/multimodal/generation/vlm_utils/custom_inputs.py
+++ b/tests/models/multimodal/generation/vlm_utils/custom_inputs.py
@@ -22,6 +22,7 @@ def multi_image_multi_aspect_ratio_inputs(formatter: Callable[[str], str]):
 
     Args:
         formatter: model-specific prompt formatter.
+
     """
     stop_sign = IMAGE_ASSETS[0].pil_image
     cherry_blossom = IMAGE_ASSETS[1].pil_image
@@ -65,6 +66,7 @@ def multi_video_multi_aspect_ratio_inputs(
 
     Args:
         formatter: model-specific prompt formatter.
+
     """
     video = sample_frames_from_video(VIDEO_ASSETS[0].np_ndarrays, num_frames)
     # Apply the selected formatter to the base prompts
diff --git a/tests/models/multimodal/generation/vlm_utils/model_utils.py b/tests/models/multimodal/generation/vlm_utils/model_utils.py
index e85a49311a77..ae3f376e1f7f 100644
--- a/tests/models/multimodal/generation/vlm_utils/model_utils.py
+++ b/tests/models/multimodal/generation/vlm_utils/model_utils.py
@@ -252,6 +252,7 @@ def qwen_prompt_path_encoder(
         tmp_path: Tempdir for test under consideration.
         prompt: Prompt with image placeholders.
         assets: list of image assets whose len equals the num placeholders.
+
     """
     # Ensure that the number of placeholders matches the number of assets;
     # If this is not true, the test is probably written incorrectly.
@@ -307,16 +308,6 @@ def processor(*args, **kwargs):
 
     hf_model.processor = processor
 
-    orig_generate = hf_model.model.generate
-
-    def _generate(self, *args, **kwargs):
-        # FIXME: https://github.com/huggingface/transformers/issues/38333
-        kwargs["disable_compile"] = True
-
-        return orig_generate(*args, **kwargs)
-
-    hf_model.model.generate = types.MethodType(_generate, hf_model.model)
-
     return hf_model
 
 
@@ -507,15 +498,12 @@ def isaac_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
     2) Ensure IsaacModel.forward returns hidden_states
     for compatibility with hidden_states_to_seq_logprobs()
     """
-
     from perceptron.tensorstream import TextType
     from perceptron.tensorstream.ops import compute_mrope_pos_tensor, modality_mask
     from transformers.modeling_outputs import BaseModelOutputWithPast
 
     def compute_position_ids_input_ids(input_ids: torch.Tensor) -> torch.Tensor:
-        """
-        Create 3D positional indices for token input.
-        """
+        """Create 3D positional indices for token input."""
         batch_size, seq_length = input_ids.shape
         position_ids = torch.arange(seq_length, device=input_ids.device)
         position_ids = position_ids.view(1, -1).expand(batch_size, -1)
@@ -570,8 +558,7 @@ def patched_forward(
         cache_position=None,
         **kwargs,
     ):
-        """
-        Forward pass with MRoPE position embeddings.
+        """Forward pass with MRoPE position embeddings.
         Computes position embeddings once and passes them through all layers.
         """
         output_hidden_states = (
@@ -1184,7 +1171,6 @@ def voxtral_patch_hf_runner(hf_model: "HfRunner") -> "HfRunner":
     HfRunner.generate calls batch_decode on the full sequence (prompt +
     generated).
     """
-
     import io
 
     import pybase64 as base64
diff --git a/tests/models/multimodal/pooling/conftest.py b/tests/models/multimodal/pooling/conftest.py
new file mode 100644
index 000000000000..fdd58e2fba81
--- /dev/null
+++ b/tests/models/multimodal/pooling/conftest.py
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pytest configuration for vLLM multimodal pooling tests."""
+
+import gc
+
+import pytest
+import torch
+
+from tests.utils import wait_for_memory_to_settle
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+@pytest.fixture(autouse=True)
+def release_gpu_memory_between_tests():
+    """Reclaim cached GPU memory between tests in this shard.
+
+    These tests start many vLLM/HF engines sequentially in one pytest
+    process; without explicit reclamation between tests, VRAM fragmentation
+    accumulates and later tests fail at engine startup ("Engine core
+    initialization failed").
+    """
+    yield
+    gc.collect()
+    if torch.accelerator.is_available():
+        torch.accelerator.empty_cache()
+    if current_platform.is_rocm():
+        return
+    try:
+        wait_for_memory_to_settle()
+    except ValueError as e:
+        # Longer-lived (class/module-scoped) engine fixtures may legitimately
+        # still hold VRAM; nothing further to reclaim in that case.
+        logger.info("Failed to clean GPU memory: %s", e)
diff --git a/tests/models/multimodal/pooling/test_dse_qwen2_vl.py b/tests/models/multimodal/pooling/test_dse_qwen2_vl.py
index ac3eb6e61723..0fff8f40ac74 100644
--- a/tests/models/multimodal/pooling/test_dse_qwen2_vl.py
+++ b/tests/models/multimodal/pooling/test_dse_qwen2_vl.py
@@ -87,7 +87,7 @@ def _run_test(
     *,
     dtype: str,
 ) -> None:
-    """SET PYTHONPATH"""
+    """SET PYTHONPATH."""
     # NOTE: take care of the order. run vLLM first, and then run HF.
     # vLLM needs a fresh new process without cuda initialization.
     # if we run HF first, the cuda initialization will be done and it
diff --git a/tests/models/multimodal/pooling/test_intern_vit.py b/tests/models/multimodal/pooling/test_intern_vit.py
index 72d2a7018bf9..4d5e8728589c 100644
--- a/tests/models/multimodal/pooling/test_intern_vit.py
+++ b/tests/models/multimodal/pooling/test_intern_vit.py
@@ -3,7 +3,7 @@
 import pytest
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModel, CLIPImageProcessor
+from transformers import AutoConfig, AutoModel, CLIPImageProcessor, PreTrainedModel
 
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.platforms import current_platform
@@ -12,11 +12,6 @@
 
 from ....conftest import ImageTestAssets
 
-pytestmark = pytest.mark.skip(
-    reason="InternVisionModel's custom code is incompatible with "
-    "transformers v5 (missing all_tied_weights_keys)"
-)
-
 # we use snapshot_download to prevent conflicts between
 # dynamic_module and trust_remote_code for hf_runner
 DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]
@@ -45,9 +40,23 @@ def run_intern_vit_test(
     if not getattr(config, "norm_type", None):
         config.norm_type = "rms_norm"
 
-    hf_model = AutoModel.from_pretrained(
-        model, dtype=torch_dtype, trust_remote_code=True
-    ).to(DEVICE_TYPE)
+    # Monkey-patch for transformers v5 compatibility: InternVisionModel's
+    # custom code doesn't call post_init(), so all_tied_weights_keys is
+    # never set. Remove once https://github.com/OpenGVLab/InternVL fixes it.
+    _orig_init = PreTrainedModel.__init__
+
+    def _patched_init(self, *args, **kwargs):
+        _orig_init(self, *args, **kwargs)
+        self.all_tied_weights_keys = {}
+
+    PreTrainedModel.__init__ = _patched_init
+
+    try:
+        hf_model = AutoModel.from_pretrained(
+            model, dtype=torch_dtype, trust_remote_code=True
+        ).to(DEVICE_TYPE)
+    finally:
+        PreTrainedModel.__init__ = _orig_init
     hf_outputs_per_image = [
         hf_model(pixel_value.to(DEVICE_TYPE)).last_hidden_state
         for pixel_value in pixel_values
diff --git a/tests/models/multimodal/pooling/test_jinavl_reranker.py b/tests/models/multimodal/pooling/test_jinavl_reranker.py
index aa039735ced7..dca2403e3510 100644
--- a/tests/models/multimodal/pooling/test_jinavl_reranker.py
+++ b/tests/models/multimodal/pooling/test_jinavl_reranker.py
@@ -2,8 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
-import transformers
-from packaging import version
 from transformers import AutoModel
 
 from vllm.assets.base import VLLM_S3_BUCKET_URL
@@ -16,11 +14,6 @@
 
 from ....conftest import HfRunner, VllmRunner
 
-pytestmark = pytest.mark.skip(
-    reason="jinaai/jina-reranker-m0 custom code is incompatible with "
-    "transformers v5 (missing all_tied_weights_keys)"
-)
-
 MODELS = ["jinaai/jina-reranker-m0"]
 
 MM_PROCESSOR_KWARGS = {
@@ -31,8 +24,8 @@
 LIMIT_MM_PER_PROMPT = {"image": 2}
 
 CHECKPOINT_TO_HF_MAPPER = {
-    "visual.": "model.visual.",
-    "model.": "model.language_model.",
+    r"^visual\.": "model.visual.",
+    r"^model\.(?!language_model\.|visual\.)": "model.language_model.",
 }
 
 HANDELSBLATT_IMAGE_URL = (
@@ -108,8 +101,7 @@ def _normalize_image(image_val: str) -> str:
 def create_score_multimodal_param(
     content_parts: list[dict],
 ) -> list[ScoreMultiModalParam]:
-    """
-    Create a ScoreMultiModalParam from a list of content dictionaries.
+    """Create a ScoreMultiModalParam from a list of content dictionaries.
 
     Each dict supports the following formats:
     - Text: {'text': 'content'}
@@ -247,17 +239,13 @@ def _run_test(
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
-@pytest.mark.skipif(
-    version.parse(transformers.__version__) == version.parse("4.57.5"),
-    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
-)
 def test_model_text_image(
     hf_runner,
     vllm_runner,
     model: str,
     dtype: str,
 ) -> None:
-    """Visual Documents Reranking"""
+    """Visual Documents Reranking."""
     _run_test(
         hf_runner,
         vllm_runner,
@@ -270,17 +258,13 @@ def test_model_text_image(
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
-@pytest.mark.skipif(
-    version.parse(transformers.__version__) == version.parse("4.57.5"),
-    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
-)
 def test_model_text_text(
     hf_runner,
     vllm_runner,
     model: str,
     dtype: str,
 ) -> None:
-    """Textual Documents Reranking"""
+    """Textual Documents Reranking."""
     _run_test(
         hf_runner,
         vllm_runner,
@@ -293,17 +277,13 @@ def test_model_text_text(
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
-@pytest.mark.skipif(
-    version.parse(transformers.__version__) == version.parse("4.57.5"),
-    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
-)
 def test_model_image_text(
     hf_runner,
     vllm_runner,
     model: str,
     dtype: str,
 ) -> None:
-    """Image Querying for Textual Documents"""
+    """Image Querying for Textual Documents."""
     _run_test(
         hf_runner,
         vllm_runner,
@@ -316,17 +296,13 @@ def test_model_image_text(
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
-@pytest.mark.skipif(
-    version.parse(transformers.__version__) == version.parse("4.57.5"),
-    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
-)
 def test_model_image_image(
     hf_runner,
     vllm_runner,
     model: str,
     dtype: str,
 ) -> None:
-    """Image Querying for Image Documents"""
+    """Image Querying for Image Documents."""
     _run_test(
         hf_runner,
         vllm_runner,
@@ -339,17 +315,13 @@ def test_model_image_image(
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
-@pytest.mark.skipif(
-    version.parse(transformers.__version__) == version.parse("4.57.5"),
-    reason="Skipped for transformers==4.57.5, https://github.com/huggingface/transformers/issues/43295",
-)
 def test_model_text_mixed_documents(
     hf_runner,
     vllm_runner,
     model: str,
     dtype: str,
 ) -> None:
-    """Text Query for Mixed Text and Image Documents"""
+    """Text Query for Mixed Text and Image Documents."""
     _run_test(
         hf_runner,
         vllm_runner,
diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py
index 9516ab95d52d..52802939cea5 100644
--- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py
+++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for the LlamaNemotronVL model family:
+"""Tests for the LlamaNemotronVL model family:
   - nvidia/llama-nemotron-embed-vl-1b-v2  (LlamaNemotronVLForCausalLM / embed)
   - nvidia/llama-nemotron-rerank-vl-1b-v2
       (LlamaNemotronVLForSequenceClassification / rerank)
diff --git a/tests/models/multimodal/processing/test_audio_in_video.py b/tests/models/multimodal/processing/test_audio_in_video.py
index df6eea0f3865..bde7bd51ab06 100644
--- a/tests/models/multimodal/processing/test_audio_in_video.py
+++ b/tests/models/multimodal/processing/test_audio_in_video.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Regression tests for Qwen2.5-Omni and Qwen3-Omni audio-in-video processor
+"""Regression tests for Qwen2.5-Omni and Qwen3-Omni audio-in-video processor
 caching.
 
 Tests the use_audio_in_video feature where audio is extracted from video and
@@ -23,7 +22,7 @@
 import pytest
 
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.cache import MultiModalProcessorSenderCache
+from vllm.multimodal.cache import LruKeyReplicatedSenderCache
 
 from ....multimodal.utils import random_audio, random_video
 from ...utils import build_model_context
@@ -50,8 +49,7 @@ def create_mm_data(num_videos: int) -> dict[str, list]:
 @pytest.mark.parametrize("model_id", MODELS)
 @pytest.mark.parametrize("num_videos", [1, 2])
 def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> None:
-    """
-    Regression test for https://github.com/vllm-project/vllm/pull/36800
+    """Regression test for https://github.com/vllm-project/vllm/pull/36800.
 
     MultiModalProcessorSenderCache.get_and_update_item returns (None, updates)
     on a cache hit, so mm_kwargs["video"] items become None on the second call.
@@ -65,36 +63,32 @@ def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> Non
         mm_processor_cache_gb=1,
     )
 
-    # Baseline: no cache, always processes from scratch.
-    baseline_processor = MULTIMODAL_REGISTRY.create_processor(
-        ctx.model_config, cache=None
-    )
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+
     # Sender cache: on a cache hit returns (None, prompt_updates) for each
     # item, setting mm_kwargs["video"] = [None] – the exact condition that
     # triggered the original bug.
-    sender_cache = MultiModalProcessorSenderCache(ctx.model_config)
-    cached_processor = MULTIMODAL_REGISTRY.create_processor(
-        ctx.model_config, cache=sender_cache
-    )
+    sender_cache = LruKeyReplicatedSenderCache(ctx.model_config)
 
-    video_token_id = baseline_processor.info.get_hf_config().video_token_id
+    video_token_id = processor.info.get_hf_config().video_token_id
 
     mm_data = create_mm_data(num_videos)
     hf_processor_mm_kwargs = {"use_audio_in_video": True}
 
-    def run(processor):
+    def run(processor, cache):
         return processor(
             [video_token_id] * num_videos,
-            mm_items=baseline_processor.info.parse_mm_data(mm_data),
+            mm_items=processor.info.parse_mm_data(mm_data),
             hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            cache=cache,
         )["prompt_token_ids"]
 
-    baseline_ids = run(baseline_processor)
+    baseline_ids = run(processor, cache=None)
 
     # First call on the sender-cache processor: cache miss.
     # mm_kwargs["video"] items are real tensors; use_audio_in_video is
     # detected normally from the item data.
-    first_ids = run(cached_processor)
+    first_ids = run(processor, cache=sender_cache)
     assert first_ids == baseline_ids, (
         "Cache-miss call produced different prompt_token_ids than baseline.\n"
         f"  baseline  : {baseline_ids}\n"
@@ -105,7 +99,7 @@ def run(processor):
     # MultiModalProcessorSenderCache.get_and_update_item returns (None, …),
     # so mm_kwargs["video"] = [None].  Before the fix, use_audio_in_video was
     # not detected, yielding wrong token ids.
-    second_ids = run(cached_processor)
+    second_ids = run(processor, cache=sender_cache)
     assert second_ids == baseline_ids, (
         "Cache-hit call produced different prompt_token_ids than baseline.\n"
         "This is the regression introduced when use_audio_in_video detection\n"
@@ -117,8 +111,7 @@ def run(processor):
 
 @pytest.mark.parametrize("model_id", MODELS)
 def test_use_audio_in_video_without_audio_track(model_id: str) -> None:
-    """
-    A video with no audio track, combined with `use_audio_in_video=True`,
+    """A video with no audio track, combined with `use_audio_in_video=True`,
     must raise a clear `ValueError` naming the mismatch, rather than an
     opaque `StopIteration`-derived error from the underlying HF processor.
     """
@@ -126,7 +119,7 @@ def test_use_audio_in_video_without_audio_track(model_id: str) -> None:
         model_id,
         limit_mm_per_prompt={"audio": 1, "image": 0, "video": 1},
     )
-    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config, cache=None)
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
     video_token_id = processor.info.get_hf_config().video_token_id
 
     rng = np.random.RandomState(0)
diff --git a/tests/models/multimodal/processing/test_audioflamingo3.py b/tests/models/multimodal/processing/test_audioflamingo3.py
index f5920a3be4b8..c3adccb27d59 100644
--- a/tests/models/multimodal/processing/test_audioflamingo3.py
+++ b/tests/models/multimodal/processing/test_audioflamingo3.py
@@ -22,18 +22,19 @@
 import numpy as np
 import pytest
 import torch
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from tests.models.registry import HF_EXAMPLE_MODELS
+from vllm.config.multimodal import MultiModalDummyOptions
 
 
-class MockAudioFlamingo3Config(PretrainedConfig):
+class MockAudioFlamingo3Config(PreTrainedConfig):
     model_type = "audioflamingo3"
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.audio_config = PretrainedConfig()
-        self.text_config = PretrainedConfig()
+        self.audio_config = PreTrainedConfig()
+        self.text_config = PreTrainedConfig()
 
 
 class MockAudioFlamingo3Processor:
@@ -127,7 +128,7 @@ def test_dummy_data_generation(mock_ctx):
     builder = AudioFlamingo3DummyInputsBuilder(info)
 
     mm_counts = {"audio": 2}
-    dummy_data = builder.get_dummy_mm_data(100, mm_counts, {})
+    dummy_data = builder.get_dummy_mm_data(100, mm_counts, MultiModalDummyOptions())
 
     assert "audio" in dummy_data
     assert len(dummy_data["audio"]) == 2
diff --git a/tests/models/multimodal/processing/test_bailing_moe_v3_vl.py b/tests/models/multimodal/processing/test_bailing_moe_v3_vl.py
index 0ebc61909fb3..5ac47190ddb7 100644
--- a/tests/models/multimodal/processing/test_bailing_moe_v3_vl.py
+++ b/tests/models/multimodal/processing/test_bailing_moe_v3_vl.py
@@ -8,7 +8,7 @@
 import torch
 from PIL import Image
 from tokenizers import Tokenizer, decoders, models, pre_tokenizers
-from transformers import PreTrainedTokenizerFast
+from transformers import TokenizersBackend
 
 from vllm.model_executor.models.bailing_moe_v3_vl import (
     IMAGE_PLACEHOLDER,
@@ -34,7 +34,7 @@ def processor():
         add_prefix_space=False, use_regex=False
     )
     backend.decoder = decoders.ByteLevel()
-    tokenizer = PreTrainedTokenizerFast(
+    tokenizer = TokenizersBackend(
         tokenizer_object=backend,
         additional_special_tokens=[
             "<|vision_start|>",
diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py
index bbc39ca171d3..20a5121e6008 100644
--- a/tests/models/multimodal/processing/test_common.py
+++ b/tests/models/multimodal/processing/test_common.py
@@ -1,24 +1,33 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from collections.abc import Set as AbstractSet
 from functools import partial
 
 import numpy as np
 import pytest
+import torch
 from PIL import Image
 
+from tests.models.utils import build_model_context
 from vllm.config import ModelConfig
 from vllm.config.multimodal import (
     AudioDummyOptions,
     BaseDummyOptions,
     ImageDummyOptions,
+    MultiModalDummyOptions,
     VideoDummyOptions,
 )
+from vllm.distributed.ec_transfer.ec_connector.utils import (
+    PlaceholderMetadataResolver,
+    collect_ec_item_metadata,
+)
+from vllm.entrypoints.chat_utils import _get_embeds_data
 from vllm.inputs import MultiModalDataDict, MultiModalInput
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.cache import MultiModalProcessorOnlyCache
-from vllm.multimodal.inputs import batched_tensors_equal
+from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal
 from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
 from vllm.platforms import current_platform
 from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
@@ -33,9 +42,7 @@
 
 
 def add_video_metadata(mm_data: MultiModalDataDict) -> MultiModalDataDict:
-    """
-    Add metadata to video mm_data
-    """
+    """Add metadata to video mm_data."""
 
     def create_metadata(frames: np.ndarray):
         num_frames = len(frames)
@@ -61,8 +68,7 @@ def create_metadata(frames: np.ndarray):
 
 
 def glmasr_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict:
-    """
-    Patch the multimodal data for GLM-ASR model.
+    """Patch the multimodal data for GLM-ASR model.
     GLM-ASR requires text and audio to match 1:1, so we limit audio to 1.
     """
     if "audio" in mm_data:
@@ -74,6 +80,8 @@ def glmasr_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict:
 
 
 _IGNORE_MM_KEYS = {
+    # Dithering causes minor divergence
+    "cohere_asr": {"input_features"},
     # In Ultravox, the audio_features can be different depending on padding
     # The slight difference should not be a problem though, since
     # attention_mask lets us ignore the difference.
@@ -159,7 +167,6 @@ def get_token_prompt(
     processor: BaseMultiModalProcessor,
     mm_data: MultiModalDataDict,
 ) -> list[int]:
-    dummy_inputs = processor.dummy_inputs
     tokenizer: TokenizerLike = processor.info.get_tokenizer()
     model_config = processor.info.ctx.model_config
 
@@ -174,18 +181,18 @@ def get_token_prompt(
     mm_counts = {k: len(vs) for k, vs in parsed_data.items()}
 
     if is_mistral_tokenizer(tokenizer):
-        inputs = dummy_inputs.get_dummy_processor_inputs(
+        inputs = processor.get_dummy_inputs(
             model_config.max_model_len,
             mm_counts,
-            mm_options={},
+            mm_options=MultiModalDummyOptions(),
             # Assume all Mistral models define this extra argument
             mm_data=mm_data,  # type: ignore[call-arg]
         )
     else:
-        inputs = dummy_inputs.get_dummy_processor_inputs(
+        inputs = processor.get_dummy_inputs(
             model_config.max_model_len,
             mm_counts,
-            mm_options={},
+            mm_options=MultiModalDummyOptions(),
         )
 
     if not isinstance(inputs.prompt, list):
@@ -194,6 +201,79 @@ def get_token_prompt(
     return inputs.prompt
 
 
+@pytest.mark.parametrize(
+    "model_id,durations",
+    [
+        ("Qwen/Qwen2-Audio-7B-Instruct", [1.0, 2.3]),
+        ("nvidia/audio-flamingo-3-hf", [31.0]),
+        ("fixie-ai/ultravox-v0_5-llama-3_2-1b", [1.0, 31.0]),
+        ("Qwen/Qwen2.5-Omni-3B", [1.0, 2.3]),
+        ("Qwen/Qwen3-Omni-30B-A3B-Instruct", [1.0, 2.3]),
+    ],
+)
+def test_audio_metadata_only_roundtrip(model_id, durations, monkeypatch):
+    """EC metadata preserves the raw-audio prompt without running the HF processor."""
+    ctx = build_model_context(model_id, limit_mm_per_prompt={"audio": len(durations)})
+    mm_config = ctx.model_config.multimodal_config
+    mm_config.enable_mm_embeds = True
+    mm_config.allow_missing_mm_embeddings = True
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    sample_rate = processor.info.get_feature_extractor().sampling_rate
+    audios = [
+        np.zeros(int(duration * sample_rate), dtype=np.float32)
+        for duration in durations
+    ]
+    prompt = processor.dummy_inputs.get_dummy_text({"audio": len(audios)})
+    raw = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data({"audio": audios}),
+        hf_processor_mm_kwargs={},
+    )
+    features = [
+        MultiModalFeatureSpec(
+            data=item,
+            modality="audio",
+            identifier=str(index),
+            mm_position=position,
+        )
+        for index, (item, position) in enumerate(
+            zip(raw["mm_kwargs"]["audio"], raw["mm_placeholders"]["audio"])
+        )
+    ]
+    published = collect_ec_item_metadata(
+        features, PlaceholderMetadataResolver(ctx.model_config)
+    )
+    metadata = json.loads(
+        json.dumps([entry["metadata"] for entry in published.values()])
+    )
+    assert len(metadata) == len(audios)
+    assert all(metadata)
+    audio_data = _get_embeds_data("audio", metadata, processor)
+
+    def unexpected_preprocessing(*args, **kwargs):
+        pytest.fail("The EC consumer must not preprocess audio")
+
+    monkeypatch.setattr(processor, "_call_hf_processor", unexpected_preprocessing)
+    remote = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data({"audio": audio_data}),
+        hf_processor_mm_kwargs={},
+    )
+    assert remote["prompt_token_ids"] == raw["prompt_token_ids"]
+    assert remote["mm_placeholders"] == raw["mm_placeholders"]
+    # Omni also consumes these lengths when constructing M-RoPE positions.
+    if "audio_feature_lengths" in audio_data:
+        assert torch.equal(
+            remote["mm_kwargs"].get_data()["audio_feature_lengths"],
+            raw["mm_kwargs"].get_data()["audio_feature_lengths"],
+        )
+
+    mm_config.allow_missing_mm_embeddings = False
+    ordinary_processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    with pytest.raises(ValueError, match="should contain the fields"):
+        ordinary_processor.info.parse_mm_data({"audio": audio_data})
+
+
 def random_vision_chunk(
     rng: np.random.RandomState,
     min_wh: int,
@@ -281,13 +361,14 @@ def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions:
         return BaseDummyOptions(count=count)
 
     # Assign normalized DummyOptions to the model config
-    model_config.get_multimodal_config().limit_per_prompt = {
-        modality: _to_dummy_options(modality, count)
-        for modality, count in limit_mm_per_prompt_ints.items()
-    }
+    model_config.get_multimodal_config().limit_per_prompt = MultiModalDummyOptions(
+        {
+            modality: _to_dummy_options(modality, count)
+            for modality, count in limit_mm_per_prompt_ints.items()
+        }
+    )
 
-    baseline_processor = factories.build_processor(ctx, cache=None)
-    cached_processor = factories.build_processor(ctx, cache=cache)
+    processor = factories.build_processor(ctx)
 
     rng = np.random.RandomState(0)
 
@@ -336,41 +417,42 @@ def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions:
         _test_processing_correctness_one(
             model_config,
             mm_data,
-            baseline_processor,
-            cached_processor,
+            processor,
             batch_idx,
             hit_rate,
             num_batches,
             simplify_rate,
+            cache=cache,
         )
 
 
 def _test_processing_correctness_one(
     model_config: ModelConfig,
     mm_data: MultiModalDataDict,
-    baseline_processor: BaseMultiModalProcessor,
-    cached_processor: BaseMultiModalProcessor,
+    processor: BaseMultiModalProcessor,
     batch_idx: int,
     hit_rate: float,
     num_batches: int,
     simplify_rate: float,
+    cache: MultiModalProcessorOnlyCache,
 ):
     model_type = model_config.hf_config.model_type
 
-    token_prompt = get_token_prompt(baseline_processor, mm_data)
-    mm_items = baseline_processor.info.parse_mm_data(mm_data)
+    token_prompt = get_token_prompt(processor, mm_data)
+    mm_items = processor.info.parse_mm_data(mm_data)
     ignore_mm_keys = _IGNORE_MM_KEYS.get(model_type, set[str]())
 
-    baseline_tokenized_result = baseline_processor(
+    baseline_tokenized_result = processor(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
     )
 
-    cached_tokenized_result = cached_processor(
+    cached_tokenized_result = processor(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
+        cache=cache,
     )
 
     _assert_inputs_equal(
diff --git a/tests/models/multimodal/processing/test_deepseek_ocr.py b/tests/models/multimodal/processing/test_deepseek_ocr.py
index 860f15ef3800..2bcf5e05d37e 100644
--- a/tests/models/multimodal/processing/test_deepseek_ocr.py
+++ b/tests/models/multimodal/processing/test_deepseek_ocr.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Regression test for DeepSeek-OCR TensorSchema validation with empty images_crop.
+"""Regression test for DeepSeek-OCR TensorSchema validation with empty images_crop.
 
 When using the Gundam preset (BASE_SIZE=1024, IMAGE_SIZE=640, CROP_MODE=True),
 images that are small enough to not require cropping produce an empty
diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py
index f30afe47dde2..6430f40d8ba7 100644
--- a/tests/models/multimodal/processing/test_gemma4.py
+++ b/tests/models/multimodal/processing/test_gemma4.py
@@ -3,6 +3,7 @@
 
 from collections.abc import Mapping
 
+import numpy as np
 import pytest
 import torch
 from PIL import Image as PILImage
@@ -13,6 +14,7 @@
     Gemma4ImagePixelInputs,
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.cache import MultiModalProcessorOnlyCache
 from vllm.multimodal.inputs import MultiModalFieldConfig
 from vllm.utils.mem_constants import GiB_bytes
 
@@ -199,6 +201,73 @@ def test_get_prompt_updates_respects_nested_max_soft_tokens(model_id: str):
     assert replacement == expected
 
 
+@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
+@pytest.mark.parametrize("kwargs_on_init", [False, True])
+@pytest.mark.parametrize(
+    "image_kwargs", [{"rescale_factor": 1 / 127.5}, {"max_soft_tokens": 560}]
+)
+@pytest.mark.parametrize("video_uuid", [None, "same-video"])
+def test_video_cache_is_independent_of_image_kwargs(
+    model_id: str,
+    kwargs_on_init: bool,
+    image_kwargs: dict[str, object],
+    video_uuid: str | None,
+):
+    """Image overrides must not change video frames or depend on cache warmth."""
+    kwargs = {"images_kwargs": image_kwargs}
+    ctx = build_model_context(
+        model_id,
+        limit_mm_per_prompt={"image": 1, "video": 1},
+        mm_processor_cache_gb=1,
+    )
+    cache = MultiModalProcessorOnlyCache(ctx.model_config)
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    hf_processor = processor.info.get_hf_processor()
+    image = PILImage.new("RGB", (48, 48), color=(128, 128, 128))
+    frames = np.stack([np.asarray(image)] * 2)
+    metadata = {"fps": 2.0, "frames_indices": [0, 1]}
+    mm_items = processor.info.parse_mm_data(
+        {"image": image, "video": [(frames, metadata)]}
+    )
+
+    def process(mm_kwargs, cache):
+        return processor(
+            hf_processor.image_token + hf_processor.video_token,
+            mm_items,
+            mm_uuid_items={"video": [video_uuid]},
+            hf_processor_mm_kwargs=mm_kwargs,
+            cache=cache,
+        )
+
+    baseline = process({}, cache=None)
+    if kwargs_on_init:
+        ctx = build_model_context(
+            model_id,
+            mm_processor_kwargs=kwargs,
+            limit_mm_per_prompt={"image": 1, "video": 1},
+            mm_processor_cache_gb=1,
+        )
+        cache = MultiModalProcessorOnlyCache(ctx.model_config)
+        processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    request_kwargs = {} if kwargs_on_init else kwargs
+    process(request_kwargs, cache=cache)
+    cached = process(request_kwargs, cache=cache)
+    cache.clear_cache()
+    fresh = process(request_kwargs, cache=cache)
+
+    def pixels(result, modality, field):
+        return result["mm_kwargs"][modality][0][field].data
+
+    video_pixels = pixels(fresh, "video", "pixel_values_videos")
+    assert torch.equal(pixels(cached, "video", "pixel_values_videos"), video_pixels)
+    assert torch.equal(pixels(baseline, "video", "pixel_values_videos"), video_pixels)
+    assert baseline["mm_hashes"]["video"] == fresh["mm_hashes"]["video"]
+    assert not torch.equal(
+        pixels(baseline, "image", "pixel_values"),
+        pixels(fresh, "image", "pixel_values"),
+    )
+
+
 @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
 def test_limit_mm_per_prompt(
     image_assets: ImageTestAssets,
diff --git a/tests/models/multimodal/processing/test_gemma4_unified.py b/tests/models/multimodal/processing/test_gemma4_unified.py
index 67a81ddb7b41..fe544e233884 100644
--- a/tests/models/multimodal/processing/test_gemma4_unified.py
+++ b/tests/models/multimodal/processing/test_gemma4_unified.py
@@ -6,9 +6,13 @@
 import pytest
 import torch
 from PIL import Image as PILImage
+from transformers.models.gemma4 import Gemma4AudioFeatureExtractor
 
 from vllm.exceptions import VLLMValidationError
-from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs
+from vllm.model_executor.models.gemma4_mm import (
+    Gemma4ImagePixelInputs,
+    _dummy_audio_num_samples,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import MultiModalFieldConfig
 
@@ -204,3 +208,41 @@ def test_limit_mm_per_prompt(
             mm_items=processor.info.parse_mm_data(mm_data),
             hf_processor_mm_kwargs={},
         )
+
+
+def test_dummy_audio_length_falls_back_for_the_unified_extractor():
+    """The unified audio feature extractor does not define ``fft_length``.
+
+    ``Gemma4DummyInputsBuilder.get_dummy_mm_data`` read that attribute
+    unconditionally, so memory profiling raised ``AttributeError`` for every
+    unified Gemma 4 deployment in which audio is the modality with the largest
+    per-item token count -- which is the case as soon as ``audio_seq_length``
+    is raised above the video budget.
+    """
+
+    class _UnifiedFeatureExtractor:
+        sampling_rate = 16000
+
+    class _UnifiedProcessor:
+        feature_extractor = _UnifiedFeatureExtractor()
+        audio_seq_length = 3000
+        audio_ms_per_token = 40
+
+    # 3000 tokens * 40 ms = 120 s at 16 kHz
+    assert _dummy_audio_num_samples(_UnifiedProcessor()) == 120 * 16000
+
+
+def test_dummy_audio_length_still_uses_fft_length_for_the_tower_extractor():
+    """The mel/tower extractor keeps its existing behaviour."""
+
+    class _TowerFeatureExtractor(Gemma4AudioFeatureExtractor):
+        def __init__(self):  # noqa: D107 - bypass the real feature-extractor init
+            self.fft_length = 512
+            self.sampling_rate = 16000
+
+    class _TowerProcessor:
+        feature_extractor = _TowerFeatureExtractor()
+        audio_seq_length = 750
+        audio_ms_per_token = 40
+
+    assert _dummy_audio_num_samples(_TowerProcessor()) == 512
diff --git a/tests/models/multimodal/processing/test_glm4_1v.py b/tests/models/multimodal/processing/test_glm4_1v.py
index ce2ab2adfb41..adf7219db1b7 100644
--- a/tests/models/multimodal/processing/test_glm4_1v.py
+++ b/tests/models/multimodal/processing/test_glm4_1v.py
@@ -152,8 +152,7 @@ def test_video_loader_consistency(
     fps: int,
     backend: str,
 ):
-    """
-    Ensure dynamic video loader (pre-sampled by loader) and normal video
+    """Ensure dynamic video loader (pre-sampled by loader) and normal video
     loader (post-sampled by processor) produce same video processing outputs.
     """
     ctx = build_model_context(
diff --git a/tests/models/multimodal/processing/test_h2ovl.py b/tests/models/multimodal/processing/test_h2ovl.py
index 3ba256f3c798..0c2e2ff3a932 100644
--- a/tests/models/multimodal/processing/test_h2ovl.py
+++ b/tests/models/multimodal/processing/test_h2ovl.py
@@ -6,7 +6,7 @@
 
 import pytest
 from PIL import Image
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.image import rescale_image_size
@@ -17,7 +17,7 @@
 
 
 def _get_expected_num_patches(
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
     image: Image.Image,
     num_imgs: int,
     min_num: int,
diff --git a/tests/models/multimodal/processing/test_idefics3.py b/tests/models/multimodal/processing/test_idefics3.py
index 7365db59f2bc..feffbf6568a1 100644
--- a/tests/models/multimodal/processing/test_idefics3.py
+++ b/tests/models/multimodal/processing/test_idefics3.py
@@ -3,9 +3,7 @@
 """Tests for Idefics3's multimodal preprocessing kwargs."""
 
 import pytest
-from packaging.version import Version
 from transformers import Idefics3Config
-from transformers import __version__ as TRANSFORMERS_VERSION
 
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
@@ -13,10 +11,6 @@
 from ...utils import build_model_context
 
 
-@pytest.mark.skipif(
-    Version(TRANSFORMERS_VERSION) < Version("5.2.0"),
-    reason="See https://github.com/huggingface/transformers/pull/43948",
-)
 @pytest.mark.parametrize("model_id", ["HuggingFaceM4/Idefics3-8B-Llama3"])
 @pytest.mark.parametrize(
     ("mm_processor_kwargs", "expected_toks_per_img"),
diff --git a/tests/models/multimodal/processing/test_internvl.py b/tests/models/multimodal/processing/test_internvl.py
index 7954dd6b5004..ab5fbe9a0130 100644
--- a/tests/models/multimodal/processing/test_internvl.py
+++ b/tests/models/multimodal/processing/test_internvl.py
@@ -6,7 +6,7 @@
 
 import pytest
 from PIL import Image
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.image import rescale_image_size
@@ -17,7 +17,7 @@
 
 
 def _get_expected_num_patches(
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
     image: Image.Image,
     num_imgs: int,
     min_num: int,
diff --git a/tests/models/multimodal/processing/test_llava_next.py b/tests/models/multimodal/processing/test_llava_next.py
index d7eaa35dd8d9..dc1ad425fb7d 100644
--- a/tests/models/multimodal/processing/test_llava_next.py
+++ b/tests/models/multimodal/processing/test_llava_next.py
@@ -116,8 +116,7 @@ def _test_image_prompt_replacements(
     num_imgs: int,
     image_sizes: list[ImageSize],
 ) -> None:
-    """
-    Ensure LlavaNextMultiModalProcessor
+    """Ensure LlavaNextMultiModalProcessor
     handles prompt replacement properly for input images.
     """
     failed_size_excs = list[tuple[ImageSize, Exception]]()
diff --git a/tests/models/multimodal/processing/test_llava_next_video.py b/tests/models/multimodal/processing/test_llava_next_video.py
index 2dcbb7d4b55f..c4fdecd21ac7 100644
--- a/tests/models/multimodal/processing/test_llava_next_video.py
+++ b/tests/models/multimodal/processing/test_llava_next_video.py
@@ -14,7 +14,7 @@
 
 
 class _StubModel:
-    """Carries only the two attributes the token helpers read from
+    """Carries only the attributes the token-count method reads from
     `self`, so the real methods can be exercised without constructing
     the full `nn.Module` (vision tower, language model, etc.)."""
 
@@ -22,18 +22,15 @@ class _StubModel:
     pooled_grid_length: int
 
 
-get_num_mm_encoder_tokens = (
-    LlavaNextVideoForConditionalGeneration.get_num_mm_encoder_tokens
-)
-get_num_mm_connector_tokens = (
-    LlavaNextVideoForConditionalGeneration.get_num_mm_connector_tokens
+get_mm_lora_token_counts = (
+    LlavaNextVideoForConditionalGeneration.get_mm_lora_token_counts
 )
 
 
 @pytest.mark.parametrize("model_id", ["llava-hf/LLaVA-NeXT-Video-7B-hf"])
 def test_num_mm_tokens_match_real_config(model_id):
     """The stored grid lengths must match what `__init__` derives from
-    the real HF config, and the two helpers must invert each other's
+    the real HF config, and the LoRA token counts must preserve the
     frame-level scaling exactly."""
     ctx = build_model_context(model_id, limit_mm_per_prompt={"video": 1})
     config = ctx.model_config.hf_config
@@ -49,10 +46,13 @@ def test_num_mm_tokens_match_real_config(model_id):
     for num_frames in (1, 2, 8, 16, 32):
         num_video_tokens = num_frames * pooled_grid_length**2
 
-        encoder_tokens = get_num_mm_encoder_tokens(stub, num_video_tokens)
+        encoder_tokens, connector_tokens = get_mm_lora_token_counts(
+            stub,
+            modality="video",
+            mm_kwargs=None,
+            num_mm_embeds=num_video_tokens,
+        )
         assert encoder_tokens == num_frames * patch_grid_length**2
-
-        connector_tokens = get_num_mm_connector_tokens(stub, encoder_tokens)
         assert connector_tokens == num_video_tokens
 
 
@@ -72,10 +72,13 @@ def test_num_mm_tokens_roundtrip(patch_grid_length, pooled_grid_length, num_fram
 
     num_video_tokens = num_frames * pooled_grid_length**2
 
-    encoder_tokens = get_num_mm_encoder_tokens(stub, num_video_tokens)
+    encoder_tokens, connector_tokens = get_mm_lora_token_counts(
+        stub,
+        modality="video",
+        mm_kwargs=None,
+        num_mm_embeds=num_video_tokens,
+    )
     assert encoder_tokens == num_frames * patch_grid_length**2
-
-    connector_tokens = get_num_mm_connector_tokens(stub, encoder_tokens)
     assert connector_tokens == num_video_tokens
 
 
@@ -84,5 +87,9 @@ def test_num_mm_tokens_zero():
     stub.patch_grid_length = 24
     stub.pooled_grid_length = 12
 
-    assert get_num_mm_encoder_tokens(stub, 0) == 0
-    assert get_num_mm_connector_tokens(stub, 0) == 0
+    assert get_mm_lora_token_counts(
+        stub,
+        modality="video",
+        mm_kwargs=None,
+        num_mm_embeds=0,
+    ) == (0, 0)
diff --git a/tests/models/multimodal/processing/test_llava_onevision.py b/tests/models/multimodal/processing/test_llava_onevision.py
index 4e8a8103a592..388316aca289 100644
--- a/tests/models/multimodal/processing/test_llava_onevision.py
+++ b/tests/models/multimodal/processing/test_llava_onevision.py
@@ -122,8 +122,7 @@ def _test_image_prompt_replacements(
     num_imgs: int,
     image_sizes: list[ImageSize],
 ) -> None:
-    """
-    Ensure LlavaOnevisionMultiModalProcessor
+    """Ensure LlavaOnevisionMultiModalProcessor
     handles prompt replacement properly for input images.
     """
     failed_size_excs = list[tuple[ImageSize, Exception]]()
diff --git a/tests/models/multimodal/processing/test_mistral3.py b/tests/models/multimodal/processing/test_mistral3.py
index fcb783f22f6e..4d900e627fdf 100644
--- a/tests/models/multimodal/processing/test_mistral3.py
+++ b/tests/models/multimodal/processing/test_mistral3.py
@@ -5,7 +5,7 @@
 import pytest
 import torch
 from PIL import Image
-from transformers import BatchFeature
+from transformers import AutoProcessor, BatchFeature
 
 from vllm.model_executor.models.lightonocr import LightOnOCRProcessingInfo
 from vllm.model_executor.models.mistral3 import Mistral3HFEncoderInfo
@@ -22,25 +22,41 @@
 _LIGHTON_MODEL_ID = "lightonai/LightOnOCR-1B-1025"
 
 
-def _processed_pixel_values(
+def _process_images_with_hf(
     hf_processor,
     images: list[Image.Image],
     mm_processor_kwargs: dict[str, object],
-) -> list[torch.Tensor]:
-    """Resize via the HF image processor and un-pad to per-image H×W."""
-    image_processor = hf_processor.image_processor
-    hf_out = image_processor(images=images, return_tensors="pt", **mm_processor_kwargs)
+) -> tuple[list[torch.Tensor], list[int]]:
+    """Process images and placeholders through the public HF processor."""
+    hf_out = hf_processor(
+        text=hf_processor.image_token * len(images),
+        images=images,
+        return_tensors="pt",
+        **mm_processor_kwargs,
+    )
     pixel_values = hf_out["pixel_values"]
     image_sizes = hf_out["image_sizes"]
-    return [p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)]
-
-
-def _placeholder_count_from_prompt_updates(
+    unpadded = [p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)]
+
+    special_token_ids = {
+        hf_processor.image_token_id,
+        hf_processor.image_break_token_id,
+        hf_processor.image_end_token_id,
+    }
+    placeholder_tokens = [
+        token_id
+        for token_id in hf_out["input_ids"][0].tolist()
+        if token_id in special_token_ids
+    ]
+    return unpadded, placeholder_tokens
+
+
+def _placeholder_tokens_from_prompt_updates(
     processor,
     images: list[Image.Image],
     pixel_values: list[torch.Tensor],
     mm_processor_kwargs: dict[str, object],
-) -> int:
+) -> list[int]:
     hf_inputs = BatchFeature({"pixel_values": pixel_values})
     fields_config = processor._get_mm_fields_config(hf_inputs, mm_processor_kwargs)
     out_mm_kwargs = MultiModalKwargsItems.from_hf_inputs(hf_inputs, fields_config)
@@ -49,13 +65,11 @@ def _placeholder_count_from_prompt_updates(
     updates = processor._get_prompt_updates(
         mm_items, mm_processor_kwargs, out_mm_kwargs
     )
-    image_token_id = processor.info.get_hf_config().image_token_index
-
-    total = 0
+    placeholder_tokens: list[int] = []
     for item_idx in range(len(images)):
         details = updates[0].resolve(item_idx).content
-        total += details.full.count(image_token_id)
-    return total
+        placeholder_tokens.extend(details.full)
+    return placeholder_tokens
 
 
 def _expected_placeholder_tokens_per_image(
@@ -72,12 +86,12 @@ def _expected_placeholder_tokens_per_image(
         patch_h = patch_w = int(patch_size)
 
     spatial_merge_size = getattr(hf_processor, "spatial_merge_size", 1)
-    assert image_h % patch_h == 0
-    assert image_w % patch_w == 0
+    merged_patch_h = patch_h * spatial_merge_size
+    merged_patch_w = patch_w * spatial_merge_size
+    assert image_h % merged_patch_h == 0
+    assert image_w % merged_patch_w == 0
 
-    return (image_h // (patch_h * spatial_merge_size)) * (
-        image_w // (patch_w * spatial_merge_size)
-    )
+    return (image_h // merged_patch_h) * (image_w // merged_patch_w)
 
 
 @pytest.mark.parametrize("model_id", [_MODEL_ID])
@@ -87,8 +101,8 @@ def _expected_placeholder_tokens_per_image(
         ({}, (448, 448), 256),
         ({"size": {"longest_edge": 1008}}, (1540, 1540), 1296),
         ({"size": {"longest_edge": 1288}}, (1536, 1187), 1656),
-        ({"size": {"longest_edge": 1008}}, (29, 29), 1),
-        ({"size": {"longest_edge": 1000}}, (1540, 1700), 1152),
+        ({"size": {"longest_edge": 1008}}, (29, 29), 4),
+        ({"size": {"longest_edge": 1000}}, (1540, 1700), 1188),
     ],
 )
 @pytest.mark.parametrize("num_imgs", [1, 2])
@@ -109,21 +123,29 @@ def test_processor_size_override(
     )
     processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
     hf_processor_mm_kwargs = {} if kwargs_on_init else mm_processor_kwargs
-    hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
+    hf_processor = AutoProcessor.from_pretrained(
+        model_id,
+        fix_mistral_regex=True,
+    )
 
     dummy_image = Image.new("RGB", image_size, color=(127, 127, 127))
     images = [dummy_image] * num_imgs
     merged_mm_kwargs = processor.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs)
-    pixel_values = _processed_pixel_values(hf_processor, images, merged_mm_kwargs)
+    pixel_values, hf_placeholder_tokens = _process_images_with_hf(
+        hf_processor, images, merged_mm_kwargs
+    )
 
-    image_token_count = _placeholder_count_from_prompt_updates(
+    prompt_update_tokens = _placeholder_tokens_from_prompt_updates(
         processor, images, pixel_values, hf_processor_mm_kwargs
     )
     expected_from_pixel_values = _expected_placeholder_tokens_per_image(
         hf_processor, pixel_values[0]
     )
     assert expected_from_pixel_values == expected_toks_per_img
-    assert image_token_count == expected_from_pixel_values * num_imgs
+    assert hf_placeholder_tokens.count(hf_processor.image_token_id) == (
+        expected_from_pixel_values * num_imgs
+    )
+    assert prompt_update_tokens == hf_placeholder_tokens
 
 
 def test_lightonocr_keeps_vision_config_image_size():
diff --git a/tests/models/multimodal/processing/test_mllama4.py b/tests/models/multimodal/processing/test_mllama4.py
index e5ba6ae24107..6cf946eff279 100644
--- a/tests/models/multimodal/processing/test_mllama4.py
+++ b/tests/models/multimodal/processing/test_mllama4.py
@@ -23,11 +23,9 @@ def test_profiling(model_id: str, max_model_len: int):
         model_config_kwargs=model_config_kwargs,
         limit_mm_per_prompt=mm_counts,
     )
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
 
-    mm_inputs = MULTIMODAL_REGISTRY.get_dummy_mm_inputs(
-        ctx.model_config,
-        mm_counts=mm_counts,
-    )
+    mm_inputs = processor.get_dummy_mm_inputs(mm_counts)
 
     hf_config = ctx.get_hf_config(Llama4Config)
     image_size = hf_config.vision_config.image_size
diff --git a/tests/models/multimodal/processing/test_molmo2.py b/tests/models/multimodal/processing/test_molmo2.py
index 5633d47a97dc..68d3c8163cb7 100644
--- a/tests/models/multimodal/processing/test_molmo2.py
+++ b/tests/models/multimodal/processing/test_molmo2.py
@@ -6,7 +6,7 @@
 import pytest
 import torch
 
-from vllm.config.multimodal import VideoDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions
 from vllm.model_executor.models.molmo2 import (
     Molmo2DummyInputsBuilder,
     build_flat_image_bool_length,
@@ -83,7 +83,9 @@ def test_dummy_video_num_frames_override_honors_min_of_two(
     data = builder.get_dummy_mm_data(
         seq_len=128,
         mm_counts={"image": 0, "video": 1},
-        mm_options={"video": VideoDummyOptions(num_frames=num_frames_override)},
+        mm_options=MultiModalDummyOptions(
+            video=VideoDummyOptions(num_frames=num_frames_override)
+        ),
     )
 
     video, _metadata = data["video"][0]
diff --git a/tests/models/multimodal/processing/test_moss_audio.py b/tests/models/multimodal/processing/test_moss_audio.py
index 87e7079a810b..17863b791ca0 100644
--- a/tests/models/multimodal/processing/test_moss_audio.py
+++ b/tests/models/multimodal/processing/test_moss_audio.py
@@ -168,7 +168,6 @@ def _build_moss_audio_processor(cache=None):
         MossAudioMultiModalProcessor(
             info,
             MossAudioDummyInputsBuilder(info),
-            cache=cache,
         ),
         ctx,
     )
@@ -259,21 +258,24 @@ def test_moss_audio_multimodal_processor_handles_token_and_cache_paths():
     )
 
     cache = MultiModalProcessorOnlyCache(ctx.model_config)
-    cached_processor, _ = _build_moss_audio_processor(cache=cache)
+    cached_processor, _ = _build_moss_audio_processor()
     cached_text_miss = cached_processor(
         prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
+        cache=cache,
     )
     cached_text_hit = cached_processor(
         prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
+        cache=cache,
     )
     cached_token_hit = cached_processor(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
+        cache=cache,
     )
 
     expected_audio_tokens = MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
diff --git a/tests/models/multimodal/processing/test_nemotron_vl.py b/tests/models/multimodal/processing/test_nemotron_vl.py
index be5c222fd213..b1c1e11892c7 100644
--- a/tests/models/multimodal/processing/test_nemotron_vl.py
+++ b/tests/models/multimodal/processing/test_nemotron_vl.py
@@ -6,7 +6,7 @@
 
 import pytest
 from PIL import Image
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.image import rescale_image_size
@@ -17,7 +17,7 @@
 
 
 def _get_expected_num_patches(
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
     image: Image.Image,
     num_imgs: int,
     min_num: int,
diff --git a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py
index af12a14a45a0..ebbc7fab421d 100644
--- a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py
+++ b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for Qwen2.5-Omni embed_input_ids to verify embeddings are
+"""Unit tests for Qwen2.5-Omni embed_input_ids to verify embeddings are
 correctly assigned to audio/image/video token positions.
 
 Regression test for: https://github.com/vllm-project/vllm/issues/34506
@@ -40,8 +39,7 @@ def _mm_embed(shape: tuple[int, ...], value: float, modality: str) -> torch.Tens
 def make_token_seq(
     audio_n: int, image_n: int, video_n: int, text_prefix: int = 3, text_sep: int = 2
 ):
-    """
-    Build a flat token sequence:
+    """Build a flat token sequence:
       [text_prefix] [AUDIO * audio_n] [text_sep] [IMAGE * image_n]
       [text_sep] [VIDEO * video_n] [text_sep]
     Returns (input_ids tensor, is_multimodal mask, positions dict).
@@ -67,9 +65,8 @@ def make_token_seq(
 def make_interleaved_seq(
     video_chunks: list[int], audio_chunks: list[int], text_prefix: int = 2
 ):
-    """
-    Build an interleaved sequence like use_audio_in_video:
-      [text] [V*v0] [A*a0] [V*v1] [A*a1] ...
+    """Build an interleaved sequence like use_audio_in_video:
+    [text] [V*v0] [A*a0] [V*v1] [A*a1] ...
     """
     tokens = [TEXT_TOKEN_ID] * text_prefix
     for v, a in zip(video_chunks, audio_chunks):
@@ -140,8 +137,7 @@ def test_multi_video_with_boundary_tokens(self):
         )
 
     def test_batched_non_interleaved_no_false_positive(self):
-        """
-        Regression test for https://github.com/vllm-project/vllm/issues/35394.
+        """Regression test for https://github.com/vllm-project/vllm/issues/35394.
 
         5 identical non-interleaved mixed-modality requests batched together:
         each has [audio][image][video] in separate blocks with text between them.
@@ -172,8 +168,7 @@ def test_batched_non_interleaved_no_false_positive(self):
 
 
 def make_mock_model(hidden: int = 8):
-    """
-    Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration
+    """Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration
     that has enough structure to run embed_input_ids.
     """
     from vllm.model_executor.models.qwen2_5_omni_thinker import (
@@ -236,8 +231,7 @@ def fake_super_embed(
 def build_mm_embeds(
     audio_n, image_n, video_n, hidden, audio_val=10.0, image_val=20.0, video_val=30.0
 ):
-    """
-    Build multimodal_embeddings list in position order (audio, image, video).
+    """Build multimodal_embeddings list in position order (audio, image, video).
     Each embedding is filled with a distinct constant so we can verify placement.
     """
     embs = []
@@ -252,8 +246,7 @@ def build_mm_embeds(
 
 class TestEmbedInputIds:
     def _run(self, audio_n, image_n, video_n, hidden=8):
-        """
-        Run embed_input_ids for a non-interleaved mixed-modality sequence.
+        """Run embed_input_ids for a non-interleaved mixed-modality sequence.
         Returns (result_embeds, input_ids, is_multimodal).
         """
         input_ids, is_multimodal = make_token_seq(audio_n, image_n, video_n)
@@ -288,8 +281,7 @@ def test_video_only(self):
         )
 
     def test_mixed_modalities_audio_goes_to_audio_pos(self):
-        """
-        Regression test for GitHub issue #34506:
+        """Regression test for GitHub issue #34506:
         With audio + image + video (non-interleaved), audio positions must
         receive audio embeddings (not image or video embeddings).
         """
@@ -329,8 +321,7 @@ def test_text_positions_unchanged(self):
         )
 
     def test_interleaved_use_audio_in_video(self):
-        """
-        Interleaved (use_audio_in_video): video chunks interleaved with audio.
+        """Interleaved (use_audio_in_video): video chunks interleaved with audio.
         Video embeddings must go to video positions, audio to audio positions.
         """
         hidden = 8
diff --git a/tests/models/multimodal/processing/test_qwen2_vl.py b/tests/models/multimodal/processing/test_qwen2_vl.py
index 5f3f5b05e106..51f65188d7b7 100644
--- a/tests/models/multimodal/processing/test_qwen2_vl.py
+++ b/tests/models/multimodal/processing/test_qwen2_vl.py
@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from typing import Any
+
+import numpy as np
 import pytest
 import torch
-from packaging.version import Version
 from PIL import Image
-from transformers import __version__ as TRANSFORMERS_VERSION
 
 from vllm.model_executor.models.vision import FusedInputNorm
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -28,7 +29,6 @@ def test_jina_vl_processing_order() -> None:
     processor = MULTIMODAL_REGISTRY.create_processor(
         ctx.model_config,
         tokenizer=ctx.tokenizer,
-        cache=cache,
     )
 
     placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
@@ -39,6 +39,7 @@ def process(images: list[Image.Image]):
         return processor(
             placeholder * len(images),
             mm_items=processor.info.parse_mm_data({"image": images}),
+            cache=cache,
         )
 
     query = process([query_image])
@@ -133,12 +134,6 @@ def test_processor_override(
     kwargs_on_init: bool,
 ):
     """Ensure Qwen2VLMultiModalProcessor handles min/max pixels properly."""
-    if (
-        Version(TRANSFORMERS_VERSION) < Version("5.2.0")
-        and "size" in mm_processor_kwargs
-    ):
-        pytest.skip("`size` ignored by `Qwen2VLProcessor.__call__`")
-
     ctx = build_model_context(
         model_id,
         mm_processor_kwargs=mm_processor_kwargs if kwargs_on_init else None,
@@ -184,12 +179,6 @@ def test_get_image_size_with_most_features(
     model_id: str,
     mm_processor_kwargs: dict[str, object],
 ):
-    if (
-        Version(TRANSFORMERS_VERSION) < Version("5.2.0")
-        and "size" in mm_processor_kwargs
-    ):
-        pytest.skip("`size` ignored by `Qwen2VLProcessor.__call__`")
-
     ctx = build_model_context(
         model_id,
         mm_processor_kwargs=mm_processor_kwargs,
@@ -222,6 +211,35 @@ def test_get_image_size_with_most_features(
         assert tokens < max_tokens
 
 
+def _build_qwen2_5_vl_video_mm_data(num_frames: int, fps: float) -> dict[str, Any]:
+    video = np.zeros((num_frames, 56, 56, 3), dtype=np.uint8)
+    metadata = {
+        "fps": fps,
+        "duration": num_frames / fps,
+        "total_num_frames": num_frames,
+        "frames_indices": list(range(num_frames)),
+        "video_backend": "opencv",
+        "do_sample_frames": False,
+    }
+    return {"video": [(video, metadata)]}
+
+
+@pytest.mark.parametrize("model_id", ["Qwen/Qwen2.5-VL-3B-Instruct"])
+def test_qwen2_5_vl_video_fps_to_second_per_grid_ts(model_id: str) -> None:
+    ctx = build_model_context(model_id, limit_mm_per_prompt={"image": 0, "video": 1})
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    temporal_patch_size = (
+        processor.info.get_hf_processor().video_processor.temporal_patch_size
+    )
+
+    prompt = "<|vision_start|><|video_pad|><|vision_end|>"
+    for fps in (1.0, 30.0):
+        mm_data = _build_qwen2_5_vl_video_mm_data(num_frames=32, fps=fps)
+        processed = processor(prompt, mm_items=processor.info.parse_mm_data(mm_data))
+        spg = processed["mm_kwargs"]["video"][0].get_data()["second_per_grid_ts"]
+        assert float(spg) == pytest.approx(temporal_patch_size / fps, rel=2e-2)
+
+
 @pytest.mark.parametrize(
     "model_id", ["Qwen/Qwen2-VL-2B-Instruct", "Qwen/Qwen2.5-VL-3B-Instruct"]
 )
@@ -232,7 +250,6 @@ def test_mm_device_do_normalize(
     num_imgs: int,
 ):
     """Ensure that enable mm_device_do_normalize yields the correct result."""
-
     ctx = build_model_context(
         model_id,
         limit_mm_per_prompt={"image": num_imgs},
diff --git a/tests/models/multimodal/processing/test_qwen3_omni.py b/tests/models/multimodal/processing/test_qwen3_omni.py
index e7a7e2de87a0..301471f45a7b 100644
--- a/tests/models/multimodal/processing/test_qwen3_omni.py
+++ b/tests/models/multimodal/processing/test_qwen3_omni.py
@@ -25,8 +25,7 @@ def test_processor_with_audio_sample_rate(
     audio_sample_rate: int,
     audio_duration_sec: float,
 ) -> None:
-    """
-    Test that vLLM's processor generates expected outputs with audio_sample_rate.
+    """Test that vLLM's processor generates expected outputs with audio_sample_rate.
 
     This validates that the processor correctly handles audio_sample_rate
     passed via hf_processor_mm_kwargs and generates audio tokens.
@@ -70,8 +69,7 @@ def test_processor_with_audio_sample_rate(
 
 @pytest.mark.parametrize("model_id", ["Qwen/Qwen3-Omni-30B-A3B-Instruct"])
 def test_longer_audio_generates_more_tokens(model_id: str) -> None:
-    """
-    Test that longer audio generates more tokens than shorter audio.
+    """Test that longer audio generates more tokens than shorter audio.
 
     This validates that audio_sample_rate is being used correctly by checking
     that audio duration affects token count as expected.
diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py
index c58904a9e2c9..04f87049c961 100644
--- a/tests/models/multimodal/processing/test_qwen3_vl.py
+++ b/tests/models/multimodal/processing/test_qwen3_vl.py
@@ -12,6 +12,7 @@
 import pytest
 
 from vllm.config import ModelConfig
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
 from ...registry import HF_EXAMPLE_MODELS
@@ -339,7 +340,9 @@ def test_dummy_video_spreads_budget_when_frame_cap_enabled(model_id: str) -> Non
         limit_mm_per_prompt={"image": 0, "video": 1},
     )
     capped = MULTIMODAL_REGISTRY.create_processor(capped_ctx.model_config)
-    capped_dummy = capped.dummy_inputs.get_dummy_mm_data(1024, {"video": 1}, {})
+    capped_dummy = capped.dummy_inputs.get_dummy_mm_data(
+        1024, {"video": 1}, MultiModalDummyOptions()
+    )
     capped_frames = capped_dummy["video"][0][0].shape[0]
     assert capped_frames == 16, (
         f"Expected the dummy to spread the budget over 16 frames, got {capped_frames}"
@@ -351,7 +354,9 @@ def test_dummy_video_spreads_budget_when_frame_cap_enabled(model_id: str) -> Non
         limit_mm_per_prompt={"image": 0, "video": 1},
     )
     uncapped = MULTIMODAL_REGISTRY.create_processor(uncapped_ctx.model_config)
-    uncapped_dummy = uncapped.dummy_inputs.get_dummy_mm_data(1024, {"video": 1}, {})
+    uncapped_dummy = uncapped.dummy_inputs.get_dummy_mm_data(
+        1024, {"video": 1}, MultiModalDummyOptions()
+    )
     uncapped_frames = uncapped_dummy["video"][0][0].shape[0]
     assert uncapped_frames == 2, (
         f"Expected the uncapped dummy to keep 2 frames, got {uncapped_frames}"
diff --git a/tests/models/multimodal/processing/test_smolvlm.py b/tests/models/multimodal/processing/test_smolvlm.py
index 678b3fd39db1..f27c4eb7b8f5 100644
--- a/tests/models/multimodal/processing/test_smolvlm.py
+++ b/tests/models/multimodal/processing/test_smolvlm.py
@@ -3,9 +3,7 @@
 """Tests for smolvlm's multimodal preprocessing kwargs."""
 
 import pytest
-from packaging.version import Version
 from transformers import SmolVLMConfig
-from transformers import __version__ as TRANSFORMERS_VERSION
 
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
@@ -13,10 +11,6 @@
 from ...utils import build_model_context
 
 
-@pytest.mark.skipif(
-    Version(TRANSFORMERS_VERSION) < Version("5.2.0"),
-    reason="See https://github.com/huggingface/transformers/pull/43948",
-)
 @pytest.mark.parametrize("model_id", ["HuggingFaceTB/SmolVLM2-2.2B-Instruct"])
 @pytest.mark.parametrize(
     ("mm_processor_kwargs", "expected_toks_per_img"),
diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py
index 0ba18e9f0315..73be3ab76c0b 100644
--- a/tests/models/multimodal/processing/test_tensor_schema.py
+++ b/tests/models/multimodal/processing/test_tensor_schema.py
@@ -18,6 +18,7 @@
     AudioDummyOptions,
     BaseDummyOptions,
     ImageDummyOptions,
+    MultiModalDummyOptions,
     VideoDummyOptions,
 )
 from vllm.distributed import (
@@ -89,16 +90,15 @@ def create_batched_mm_kwargs(
     size_factors: tuple[float, ...] = (1.0, 0.5, 0.25),
 ) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
     processing_info = processor.info
-    dummy_inputs = processor.dummy_inputs
     supported_mm_limits = processing_info.get_supported_mm_limits()
     mm_counts = {
         modality: 3 if limit is None else limit
         for modality, limit in supported_mm_limits.items()
     }
-    processor_inputs = dummy_inputs.get_dummy_processor_inputs(
+    processor_inputs = processor.get_dummy_inputs(
         seq_len=model_config.max_model_len,
         mm_counts=mm_counts,
-        mm_options={},
+        mm_options=MultiModalDummyOptions(),
     )
     mm_items = processor_inputs.mm_data_items
     resized_mm_data = {
@@ -164,10 +164,11 @@ def test_model_tensor_schema(model_id: str):
             "Kimi-K2.5's offline inference has issues about vision chunks. Fix later."
         )
 
-    if model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" and not (
-        current_platform.is_cuda()
+    if (
+        model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp"
+        and not current_platform.is_cuda_alike()
     ):
-        pytest.skip("Deepseek V4 is only supported on CUDA")
+        pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm")
 
     if model_id == "zai-org/GLM-5.3-Flash" and (current_platform.is_xpu()):
         pytest.skip("GLM-5.3-Flash is not supported on XPU")
@@ -216,8 +217,14 @@ def test_model_tensor_schema(model_id: str):
 
     factories = model_cls._processor_factory
 
+    # Capture helpers return capture-buffer containers (e.g.
+    # EncoderCudaGraphCaptureInputs), not TensorSchema mm inputs.
+    capture_helpers = {"prepare_encoder_cudagraph_capture_inputs"}
+
     inputs_parse_methods = []
     for attr_name in dir(model_cls):
+        if attr_name in capture_helpers:
+            continue
         attr = getattr(model_cls, attr_name)
         if hasattr(attr, "__annotations__"):
             return_type = attr.__annotations__.get("return", None)
@@ -247,11 +254,13 @@ def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions:
             return AudioDummyOptions(count=count)
         return BaseDummyOptions(count=count)
 
-    model_config.get_multimodal_config().limit_per_prompt = {
-        modality: _to_dummy_options(modality, count)
-        for modality, count in limit_mm_per_prompt.items()
-    }
-    processor = factories.build_processor(ctx, cache=None)
+    model_config.get_multimodal_config().limit_per_prompt = MultiModalDummyOptions(
+        {
+            modality: _to_dummy_options(modality, count)
+            for modality, count in limit_mm_per_prompt.items()
+        }
+    )
+    processor = factories.build_processor(ctx)
 
     with initialize_dummy_model(model_cls, model_config) as model:
         for modality, _, mm_kwargs in create_batched_mm_kwargs(model_config, processor):
diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py
index f5de7d66f4ff..8f5a47be7b94 100644
--- a/tests/models/multimodal/processing/test_transformers_image.py
+++ b/tests/models/multimodal/processing/test_transformers_image.py
@@ -245,6 +245,7 @@ def process():
             prompt="<image>\nWhat is this?",
             mm_items=mm_processor.info.parse_mm_data({"image": image}),
             hf_processor_mm_kwargs={},
+            cache=cache,
         )
 
     first, second = process(), process()
diff --git a/tests/models/multimodal/processing/test_whisper.py b/tests/models/multimodal/processing/test_whisper.py
new file mode 100644
index 000000000000..1627b48914c3
--- /dev/null
+++ b/tests/models/multimodal/processing/test_whisper.py
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for Whisper's multimodal preprocessing."""
+
+import numpy as np
+import pytest
+
+from vllm.multimodal import MULTIMODAL_REGISTRY
+
+from ...utils import build_model_context
+
+
+@pytest.mark.parametrize("model_id", ["openai/whisper-large-v3-turbo"])
+@pytest.mark.parametrize("audio_duration_s", [5, 30, 35])
+def test_audio_features_fit_encoder_window(
+    model_id: str,
+    audio_duration_s: int,
+) -> None:
+    """The encoder only has positions for one 30 s window, so a longer clip
+    has to be cut to that window like the HF feature extractor does by default.
+    """
+    ctx = build_model_context(model_id, limit_mm_per_prompt={"audio": 1})
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    feature_extractor = processor.info.get_feature_extractor()
+
+    sampling_rate = feature_extractor.sampling_rate
+    rng = np.random.RandomState(0)
+    audio = rng.rand(sampling_rate * audio_duration_s).astype(np.float32)
+
+    processed_inputs = processor(
+        "<|startoftranscript|>",
+        mm_items=processor.info.parse_mm_data({"audio": [(audio, sampling_rate)]}),
+        hf_processor_mm_kwargs={},
+    )
+
+    input_features = processed_inputs["mm_kwargs"].get_data()["input_features"]
+    assert input_features.shape[-1] == feature_extractor.nb_max_frames
diff --git a/tests/models/multimodal/processing/transformers_backend.py b/tests/models/multimodal/processing/transformers_backend.py
index dc49036e8352..869afe945653 100644
--- a/tests/models/multimodal/processing/transformers_backend.py
+++ b/tests/models/multimodal/processing/transformers_backend.py
@@ -48,4 +48,4 @@ def create_cached_processor(model_id: str, processor_cls):
     )
     info = MultiModalProcessingInfo(ctx)
     cache = MultiModalProcessorOnlyCache(model_config)
-    return processor_cls(info, MultiModalDummyInputsBuilder(info), cache=cache), cache
+    return processor_cls(info, MultiModalDummyInputsBuilder(info)), cache
diff --git a/tests/models/multimodal/test_aria.py b/tests/models/multimodal/test_aria.py
new file mode 100644
index 000000000000..fed7b2064660
--- /dev/null
+++ b/tests/models/multimodal/test_aria.py
@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+from transformers import AriaTextConfig
+from transformers.models.aria.modeling_aria import AriaTextMoELayer as HFMoE
+
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (  # noqa: E501
+    CompressedTensorsConfig,
+)
+from vllm.model_executor.models.aria import AriaTextModel, AriaTextMoELayer
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+
+@pytest.mark.cpu_test
+@pytest.mark.usefixtures("dist_init")
+@pytest.mark.parametrize(
+    "checkpoint_name,loaded_name,param_name",
+    [
+        (
+            "experts.fc1.weight",
+            "experts.w13_weight",
+            "experts.routed_experts.w13_weight",
+        ),
+        (
+            "experts.fc2.weight",
+            "experts.w2_weight",
+            "experts.routed_experts.w2_weight",
+        ),
+        (
+            "shared_experts.down_proj.weight",
+            "shared_experts.down_proj.weight",
+            "shared_experts.down_proj.weight",
+        ),
+    ],
+)
+def test_aria_expert_weights_load_with_checkpoint_layout(
+    checkpoint_name: str, loaded_name: str, param_name: str
+):
+    """Real Aria checkpoint weights must reach the right parameter and layout."""
+    config = AriaTextConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        moe_num_experts=2,
+        moe_topk=1,
+        moe_num_shared_experts=1,
+    )
+    checkpoint_weight = HFMoE(config).state_dict()[checkpoint_name]
+    checkpoint_weight.copy_(
+        torch.arange(
+            checkpoint_weight.numel(), dtype=checkpoint_weight.dtype
+        ).reshape_as(checkpoint_weight)
+    )
+    expected = checkpoint_weight.clone()
+    if checkpoint_name.startswith("experts."):
+        expected = expected.transpose(-1, -2)
+
+    layer = AriaTextMoELayer(config, quant_config=None, prefix="mlp")
+    loaded = AutoWeightsLoader(layer).load_weights(
+        [(checkpoint_name, checkpoint_weight)], mapper=AriaTextModel.hf_to_vllm_mapper
+    )
+
+    assert loaded == {loaded_name}
+    torch.testing.assert_close(
+        dict(layer.named_parameters())[param_name], expected, rtol=0, atol=0
+    )
+
+
+@pytest.mark.cpu_test
+def test_aria_quant_config_renames_expert_modules():
+    quant_config = CompressedTensorsConfig(
+        target_scheme_map={},
+        ignore=[
+            "model.layers.0.mlp.experts.fc1",
+            "model.layers.0.mlp.experts.fc2",
+        ],
+        quant_format="pack-quantized",
+    )
+    AriaTextModel.__new__(
+        AriaTextModel, vllm_config=VllmConfig(quant_config=quant_config)
+    )
+
+    assert quant_config.ignore == [
+        "model.layers.0.mlp.experts.gate_up_proj",
+        "model.layers.0.mlp.experts.down_proj",
+    ]
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "name,expected", [("fc1", "gate_up_proj"), ("fc2", "down_proj")]
+)
+def test_aria_expert_scale_names(name: str, expected: str):
+    suffixes = ["weight_scale", "input_scale"]
+    names = [f"model.layers.0.mlp.experts.{name}.{suffix}" for suffix in suffixes]
+    assert AriaTextModel.hf_to_vllm_mapper.apply_list(names) == [
+        f"model.layers.0.mlp.experts.{expected}.{suffix}" for suffix in suffixes
+    ]
diff --git a/tests/models/multimodal/test_cohere_asr.py b/tests/models/multimodal/test_cohere_asr.py
index 85a9b59dd513..950542f90019 100644
--- a/tests/models/multimodal/test_cohere_asr.py
+++ b/tests/models/multimodal/test_cohere_asr.py
@@ -4,9 +4,65 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 
 from vllm.config import SpeechToTextConfig
-from vllm.model_executor.models.cohere_asr import CohereAsrForConditionalGeneration
+from vllm.model_executor.models.cohere_asr import (
+    CohereAsrForConditionalGeneration,
+    RelPositionMultiHeadAttention,
+)
+
+
+class _ReferenceRelPositionMultiHeadAttention(RelPositionMultiHeadAttention):
+    """Relative attention implementation before score accumulation is fused."""
+
+    def forward(self, query, key, value, mask, pos_emb=None):
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)
+
+        assert pos_emb is not None
+        n_batch_pos = pos_emb.size(0)
+        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
+        p = p.transpose(1, 2)
+
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        matrix_bd = self.rel_shift(matrix_bd)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+        matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)]
+        scores = (matrix_ac + matrix_bd) / self.s_d_k
+        return self.forward_attention(v, scores, mask)
+
+
+@pytest.mark.parametrize(("query_len", "key_len"), [(2, 2), (3, 5)])
+def test_relative_attention_fuses_score_accumulation(monkeypatch, query_len, key_len):
+    torch.manual_seed(7)
+    batch, heads, model_dim = 2, 4, 32
+    optimized = RelPositionMultiHeadAttention(heads, model_dim, None, None)
+    reference = _ReferenceRelPositionMultiHeadAttention(heads, model_dim, None, None)
+    reference.load_state_dict(optimized.state_dict())
+
+    query = torch.randn(batch, query_len, model_dim)
+    key = torch.randn(batch, key_len, model_dim)
+    value = torch.randn(batch, key_len, model_dim)
+    pos_emb = torch.randn(1, query_len + key_len - 1, model_dim)
+    expected = reference(query, key, value, None, pos_emb)
+
+    original_baddbmm = torch.baddbmm
+    calls = 0
+
+    def tracked_baddbmm(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        return original_baddbmm(*args, **kwargs)
+
+    monkeypatch.setattr(torch, "baddbmm", tracked_baddbmm)
+    actual = optimized(query, key, value, None, pos_emb)
+
+    torch.testing.assert_close(actual, expected)
+    assert calls == 1
 
 
 @pytest.mark.parametrize(
diff --git a/tests/models/multimodal/test_mapping.py b/tests/models/multimodal/test_mapping.py
index b762ab2a13b6..95d60914218a 100644
--- a/tests/models/multimodal/test_mapping.py
+++ b/tests/models/multimodal/test_mapping.py
@@ -215,7 +215,7 @@ def test_bailing_vl_mapper_handles_module_and_parameter_names():
 
 
 def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]:
-    """Create weights from safetensors checkpoint metadata"""
+    """Create weights from safetensors checkpoint metadata."""
     metadata = try_get_safetensors_metadata(repo)
     weight_names = list(metadata.weight_map.keys())
     with torch.device("meta"):
@@ -223,9 +223,7 @@ def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]:
 
 
 def create_dummy_base_model(repo: str, model_arch: str) -> PreTrainedModel:
-    """
-    Create weights from a dummy meta deserialized hf base model with name conversion
-    """
+    """Create weights from a dummy meta-deserialized HF base model, with renaming."""
     config = AutoConfig.from_pretrained(repo)
     with torch.device("meta"):
         model = AutoModel.from_config(config)
@@ -233,9 +231,7 @@ def create_dummy_base_model(repo: str, model_arch: str) -> PreTrainedModel:
 
 
 def create_dummy_model(repo: str, model_arch: str) -> PreTrainedModel:
-    """
-    Create weights from a dummy meta deserialized hf model with name conversion
-    """
+    """Create weights from a dummy meta deserialized hf model with name conversion."""
     model_cls: PreTrainedModel = getattr(transformers, model_arch)
     config = AutoConfig.from_pretrained(repo)
     with torch.device("meta"):
diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py
index 0aef83a8271b..75900cde53a3 100644
--- a/tests/models/multimodal/test_nano_nemotron_vl.py
+++ b/tests/models/multimodal/test_nano_nemotron_vl.py
@@ -5,13 +5,29 @@
 
 import numpy as np
 import pytest
+import torch
 
 from vllm import envs
 from vllm.model_executor.models.nano_nemotron_vl import (
     NanoNemotronVLMultiModalProcessor,
     NemotronH_Nano_VL_V2,
 )
-from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems
+from vllm.multimodal.parse import (
+    MultiModalDataItems,
+    MultiModalDataParser,
+    VideoProcessorItems,
+)
+
+
+@pytest.mark.parametrize("input_key", ["image_embeds", "video_embeds"])
+def test_precomputed_multimodal_embeddings(input_key: str):
+    model = object.__new__(NemotronH_Nano_VL_V2)
+    embeds = torch.randn(2, 4, 8)
+
+    outputs = model.embed_multimodal(**{input_key: embeds})
+
+    assert len(outputs) == len(embeds)
+    assert all(torch.equal(output, embed) for output, embed in zip(outputs, embeds))
 
 
 class _TextOnlyMultiModalConfig:
@@ -154,6 +170,9 @@ def test_extract_audio_from_videos_passes_max_duration():
     mm_items = _make_mm_items_with_video_bytes(b"\x00" * 64)
 
     processor = object.__new__(NanoNemotronVLMultiModalProcessor)
+    processor.data_parser = MultiModalDataParser(
+        target_sr=dummy_audio[1], target_channels=1
+    )
 
     target = "vllm.model_executor.models.nano_nemotron_vl.load_audio_pyav"
     with patch(target, return_value=dummy_audio) as mock_load:
diff --git a/tests/models/quantization/conftest.py b/tests/models/quantization/conftest.py
new file mode 100644
index 000000000000..261f7934f4cf
--- /dev/null
+++ b/tests/models/quantization/conftest.py
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared BF16 baselines for the quantization comparison tests.
+
+Several tests here quantize the same model and compare against an
+unquantized run of it. Computing that reference once per session avoids
+booting an identical engine for every quantization scheme under test.
+
+Only the outputs are cached; each engine is still torn down as usual, so
+tests keep the isolation ``VllmRunner.__exit__`` provides.
+"""
+
+import pytest
+
+
+@pytest.fixture(scope="session")
+def quant_baseline_logprobs(vllm_runner):
+    """Greedy logprobs for an unquantized model, computed once per config."""
+    cache: dict = {}
+
+    def baseline(
+        model: str,
+        prompts: list[str],
+        *,
+        max_model_len: int,
+        max_tokens: int,
+        num_logprobs: int,
+    ):
+        key = (model, max_model_len, max_tokens, num_logprobs, tuple(prompts))
+        if key not in cache:
+            with vllm_runner(
+                model,
+                max_model_len=max_model_len,
+                enforce_eager=True,
+            ) as vllm_model:
+                cache[key] = vllm_model.generate_greedy_logprobs(
+                    prompts, max_tokens, num_logprobs
+                )
+        return cache[key]
+
+    return baseline
diff --git a/tests/models/quantization/test_fp8_per_channel.py b/tests/models/quantization/test_fp8_per_channel.py
index f003ccd5f8dc..ef11b17968c6 100644
--- a/tests/models/quantization/test_fp8_per_channel.py
+++ b/tests/models/quantization/test_fp8_per_channel.py
@@ -37,6 +37,7 @@
 @pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
 def test_fp8_per_channel_logprobs(
     vllm_runner,
+    quant_baseline_logprobs,
     example_prompts,
     model: str,
     monkeypatch: pytest.MonkeyPatch,
@@ -52,14 +53,13 @@ def test_fp8_per_channel_logprobs(
     with monkeypatch.context() as m:
         m.setenv("TOKENIZERS_PARALLELISM", "true")
 
-        with vllm_runner(
+        baseline_outputs = quant_baseline_logprobs(
             model,
+            example_prompts,
             max_model_len=MAX_MODEL_LEN,
-            enforce_eager=True,
-        ) as vllm_model:
-            baseline_outputs = vllm_model.generate_greedy_logprobs(
-                example_prompts, MAX_TOKENS, NUM_LOG_PROBS
-            )
+            max_tokens=MAX_TOKENS,
+            num_logprobs=NUM_LOG_PROBS,
+        )
 
         with vllm_runner(
             model,
diff --git a/tests/models/quantization/test_gpt_oss.py b/tests/models/quantization/test_gpt_oss.py
index 209f15275902..10b0777f5e3a 100644
--- a/tests/models/quantization/test_gpt_oss.py
+++ b/tests/models/quantization/test_gpt_oss.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-End-to-end accuracy test for GPT-OSS model quantization.
+"""End-to-end accuracy test for GPT-OSS model quantization.
 
 Config:
     Task:   gsm8k_platinum
@@ -170,8 +169,7 @@ def test_aiter_mxfp4_moe_ignores_padded_rows(
     monkeypatch: pytest.MonkeyPatch,
     dist_init,
 ) -> None:
-    """
-    Garbage in cudagraph padding rows must not reach the unpadded outputs.
+    """Garbage in cudagraph padding rows must not reach the unpadded outputs.
 
     A cudagraph replay of a size-`TOKENS_PADDED` graph driven by
     `TOKENS_UNPADDED` real tokens leaves the trailing padding row containing
diff --git a/tests/models/quantization/test_mimo_v2_qkv_shard.py b/tests/models/quantization/test_mimo_v2_qkv_shard.py
new file mode 100644
index 000000000000..68afe7fdecfc
--- /dev/null
+++ b/tests/models/quantization/test_mimo_v2_qkv_shard.py
@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit test for MiMo-V2 fused fp8 ``qkv_proj`` sharding.
+
+The checkpoint stores the fused QKV pre-sharded at ``num_key_value_heads``
+chunks (``ckpt_tp``), each holding that slice's ``[Q_c | K_c | V_c]`` rows with
+its own block scale grid. ``ckpt_tp`` is not the layer's KV-head count: a
+MiMo-V2.5 SWA layer has 8 KV heads but 4 chunks of 3712 rows (= 29 blocks), so
+each chunk carries two KV heads. Reading the scales as one slice per KV head
+expands 14 x 128 = 1792 rows for an 1856-row slice and fails.
+
+These tests build such checkpoints for the GA, SWA and Pro geometries and check
+that every rank gets exactly the weights it should own, for TP sizes below,
+equal to and above ``ckpt_tp`` (the last case replicating KV heads).
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.models.mimo_v2 import _shard_fp8_qkv_proj
+from vllm.utils.math_utils import cdiv
+from vllm.utils.torch_utils import set_default_torch_num_threads
+
+pytestmark = pytest.mark.cpu_test
+
+BLOCK = 128
+FP8_DTYPE = torch.float8_e4m3fn
+FP8_MAX = torch.finfo(FP8_DTYPE).max
+COLS = 128  # one column block is enough to exercise the row mapping
+# fp8 e4m3 has 3 mantissa bits, so a single block-quantization round trip costs
+# ~2-3% relative L2; a wrong row -> scale pairing costs >= 10%.
+TOL = 0.06
+
+# (num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp): MiMo-V2.5 GA, its
+# SWA/MTP layers, and MiMo-V2.5-Pro.
+GEOMETRIES = {
+    "ga": (64, 4, 192, 128, 4),
+    "swa": (64, 8, 192, 128, 4),
+    "pro": (128, 8, 192, 128, 8),
+}
+
+
+@pytest.fixture(scope="module", autouse=True)
+def single_threaded_cpu_ops():
+    # Do not prime the parent's thread pool before later engine forks.
+    with set_default_torch_num_threads(1):
+        yield
+
+
+def _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp):
+    return (
+        num_heads // ckpt_tp * head_dim
+        + num_kv_heads // ckpt_tp * head_dim
+        + num_kv_heads // ckpt_tp * v_head_dim
+    )
+
+
+def _quantize_chunks(truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp):
+    """Block-quantize each checkpoint chunk on its own grid, like the export."""
+    rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
+    weights, scales = [], []
+    for chunk in range(ckpt_tp):
+        rows = truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk]
+        quantized = torch.zeros_like(rows)
+        scale = torch.zeros(
+            cdiv(rows_per_chunk, BLOCK), rows.shape[1] // BLOCK, dtype=torch.float32
+        )
+        for row in range(0, rows_per_chunk, BLOCK):
+            for col in range(0, rows.shape[1], BLOCK):
+                block = rows[row : row + BLOCK, col : col + BLOCK]
+                amax = block.abs().max().clamp(min=1e-12) / FP8_MAX
+                scale[row // BLOCK, col // BLOCK] = amax
+                quantized[row : row + BLOCK, col : col + BLOCK] = block / amax
+        weights.append(quantized.to(FP8_DTYPE))
+        scales.append(scale)
+    return torch.cat(weights, dim=0), torch.cat(scales, dim=0)
+
+
+def _owned_rows(
+    truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp, tp_rank, tp_size
+):
+    """Rows the rank must own, in the forward's `[Q | K | V]` order."""
+    rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
+    q_per_chunk = num_heads // ckpt_tp * head_dim
+    k_per_chunk = num_kv_heads // ckpt_tp * head_dim
+    kv_heads: list[int] = []
+    if tp_size <= num_kv_heads:
+        kv_heads = list(
+            range(
+                tp_rank * (num_kv_heads // tp_size),
+                (tp_rank + 1) * (num_kv_heads // tp_size),
+            )
+        )
+    else:
+        kv_heads = [tp_rank // (tp_size // num_kv_heads)]
+    index = []
+    for head in range(
+        tp_rank * (num_heads // tp_size), (tp_rank + 1) * (num_heads // tp_size)
+    ):
+        chunk, offset = divmod(head, num_heads // ckpt_tp)
+        start = chunk * rows_per_chunk + offset * head_dim
+        index += list(range(start, start + head_dim))
+    for head in kv_heads:
+        chunk, offset = divmod(head, num_kv_heads // ckpt_tp)
+        start = chunk * rows_per_chunk + q_per_chunk + offset * head_dim
+        index += list(range(start, start + head_dim))
+    for head in kv_heads:
+        chunk, offset = divmod(head, num_kv_heads // ckpt_tp)
+        start = chunk * rows_per_chunk + q_per_chunk + k_per_chunk + offset * v_head_dim
+        index += list(range(start, start + v_head_dim))
+    return truth[torch.tensor(index)]
+
+
+def _dequantize(weight, scale):
+    """What the kernel computes: scale[row // block] applied per weight row."""
+    per_row = scale.repeat_interleave(BLOCK, dim=1)
+    return weight.to(torch.float32) * per_row[torch.arange(weight.shape[0]) // BLOCK]
+
+
+def _relative_l2(got, expected):
+    return ((got - expected).norm() / expected.norm()).item()
+
+
+@pytest.mark.parametrize("geometry", list(GEOMETRIES))
+@pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
+def test_fused_qkv_proj_sharding(geometry, tp_size):
+    num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES[geometry]
+    if num_heads % tp_size or (tp_size > num_kv_heads and tp_size % num_kv_heads):
+        pytest.skip("TP size does not divide the head counts")
+
+    rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
+    truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
+    for chunk in range(ckpt_tp):  # per-chunk magnitude spread, as in the export
+        truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] *= 1.0 + chunk
+    weight, scale = _quantize_chunks(
+        truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
+    )
+
+    kv_per_rank = (num_kv_heads // tp_size) if tp_size <= num_kv_heads else 1
+    rows_rank = (num_heads // tp_size) * head_dim + kv_per_rank * (
+        head_dim + v_head_dim
+    )
+    for tp_rank in range(tp_size):
+        w_rank, s_rank = _shard_fp8_qkv_proj(
+            weight,
+            scale,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            v_head_dim=v_head_dim,
+            tp_rank=tp_rank,
+            tp_size=tp_size,
+            ckpt_tp=ckpt_tp,
+        )
+        assert w_rank.shape == (rows_rank, COLS)
+        assert s_rank.shape == (cdiv(rows_rank, BLOCK), scale.shape[1])
+        expected = _owned_rows(
+            truth,
+            num_heads,
+            num_kv_heads,
+            head_dim,
+            v_head_dim,
+            ckpt_tp,
+            tp_rank,
+            tp_size,
+        )
+        assert _relative_l2(_dequantize(w_rank, s_rank), expected) <= TOL
+
+
+@pytest.mark.parametrize("geometry", ["ga", "swa"])
+@pytest.mark.parametrize("tp_size", [1, 2, 4])
+def test_exact_chunk_matches_checkpoint(geometry, tp_size):
+    """At tp_size <= ckpt_tp the chunk count is the only thing that matters."""
+    num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES[geometry]
+    rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
+    truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
+    weight, scale = _quantize_chunks(
+        truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
+    )
+    if tp_size == ckpt_tp:
+        w_rank, s_rank = _shard_fp8_qkv_proj(
+            weight,
+            scale,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            v_head_dim=v_head_dim,
+            tp_rank=0,
+            tp_size=tp_size,
+            ckpt_tp=ckpt_tp,
+        )
+        # The checkpoint chunk *is* the rank's slice: no re-quantization.
+        # (torch.equal has no fp8 kernel, so compare the underlying bytes.)
+        expected_w = weight.chunk(tp_size, dim=0)[0]
+        assert torch.equal(w_rank.view(torch.uint8), expected_w.view(torch.uint8))
+        assert torch.equal(s_rank, scale.chunk(tp_size, dim=0)[0])
+
+
+def test_wrong_chunk_count_is_detected():
+    """Reading SWA scales as one slice per KV head must not look correct.
+
+    This is the regression the fix addresses: 116 scale rows are 4 x 29 (29 =
+    ceil(3712 / 128)), not 8 x 15, so ``ckpt_tp = num_kv_heads`` pairs every
+    weight row with the wrong scale.
+    """
+    num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES["swa"]
+    rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
+    truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
+    for chunk in range(ckpt_tp):
+        truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] *= 1.0 + chunk
+    weight, scale = _quantize_chunks(
+        truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
+    )
+    assert scale.shape[0] == ckpt_tp * cdiv(rows_per_chunk, BLOCK)
+    assert scale.shape[0] != num_kv_heads * cdiv(rows_per_chunk, BLOCK)
+
+    kw = dict(
+        num_heads=num_heads,
+        num_kv_heads=num_kv_heads,
+        head_dim=head_dim,
+        v_head_dim=v_head_dim,
+        tp_rank=0,
+        tp_size=1,
+    )
+    correct = _shard_fp8_qkv_proj(weight, scale, ckpt_tp=ckpt_tp, **kw)
+    wrong = _shard_fp8_qkv_proj(weight, scale, ckpt_tp=num_kv_heads, **kw)
+    assert correct[0].shape == wrong[0].shape
+
+    expected = _owned_rows(
+        truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp, 0, 1
+    )
+    assert _relative_l2(_dequantize(*correct), expected) <= TOL
+    assert _relative_l2(_dequantize(*wrong), expected) > 5 * TOL
diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py
index c12a72a09c0a..06b2b3cccd20 100644
--- a/tests/models/quantization/test_mxfp8.py
+++ b/tests/models/quantization/test_mxfp8.py
@@ -42,6 +42,7 @@
 @pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
 def test_mxfp8_logprobs(
     vllm_runner,
+    quant_baseline_logprobs,
     example_prompts,
     model: str,
     monkeypatch: pytest.MonkeyPatch,
@@ -56,14 +57,13 @@ def test_mxfp8_logprobs(
     with monkeypatch.context() as m:
         m.setenv("TOKENIZERS_PARALLELISM", "true")
 
-        with vllm_runner(
+        baseline_outputs = quant_baseline_logprobs(
             model,
+            example_prompts,
             max_model_len=MAX_MODEL_LEN,
-            enforce_eager=True,
-        ) as vllm_model:
-            baseline_outputs = vllm_model.generate_greedy_logprobs(
-                example_prompts, MAX_TOKENS, NUM_LOG_PROBS
-            )
+            max_tokens=MAX_TOKENS,
+            num_logprobs=NUM_LOG_PROBS,
+        )
 
         with vllm_runner(
             model,
@@ -111,11 +111,6 @@ def test_mxfp8_aiter_requires_swigluoai_activation(
         "_supports_current_device",
         staticmethod(lambda: True),
     )
-    monkeypatch.setattr(
-        aiter_mxfp8_moe,
-        "is_aiter_mxfp8_moe_available",
-        lambda: True,
-    )
 
     config = FusedMoEConfig(
         num_experts=8,
@@ -164,11 +159,6 @@ def test_mxfp8_aiter_requires_swigluoai_params(
         "_supports_current_device",
         staticmethod(lambda: True),
     )
-    monkeypatch.setattr(
-        aiter_mxfp8_moe,
-        "is_aiter_mxfp8_moe_available",
-        lambda: True,
-    )
 
     config = FusedMoEConfig(
         num_experts=8,
@@ -218,11 +208,6 @@ def test_mxfp8_aiter_accepts_swigluoai_params(
         "_supports_current_device",
         staticmethod(lambda: True),
     )
-    monkeypatch.setattr(
-        aiter_mxfp8_moe,
-        "is_aiter_mxfp8_moe_available",
-        lambda: True,
-    )
 
     config = FusedMoEConfig(
         num_experts=8,
diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py
index 020e1cc8fc78..e7d0e72d8a3e 100644
--- a/tests/models/qwen4_exp/test_ple.py
+++ b/tests/models/qwen4_exp/test_ple.py
@@ -1610,6 +1610,12 @@ def test_fused_conv_correctness(
         dtype=torch.bfloat16,
         generator=rng,
     )
+    outer_residual = torch.randn(
+        inputs.shape,
+        device=device,
+        dtype=torch.bfloat16,
+        generator=rng,
+    )
     null_state = conv_state[NULL_BLOCK_ID].clone()
     residual_kernel = residual.clone()
     residual_reference = residual.clone()
@@ -1617,6 +1623,7 @@ def test_fused_conv_correctness(
     module._short_conv_dilated_dispatch(
         inputs=inputs,
         residual=residual_kernel,
+        outer_residual=outer_residual,
         metadata=metadata,
         conv_state=conv_state,
         conv_weights=weights,
@@ -1630,15 +1637,17 @@ def test_fused_conv_correctness(
         conv_state_len=module.conv_state_len,
         dilation=module.short_conv_dilation,
     )
+    residual_reference = outer_residual + residual_reference
 
-    torch.testing.assert_close(
-        residual_kernel.float(), residual_reference.float(), atol=3e-2, rtol=3e-2
+    assert torch.equal(
+        residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens]
     )
     assert torch.equal(conv_state, state_reference)
     assert torch.equal(conv_state[NULL_BLOCK_ID], null_state)
     if case.graph_padding:
         assert torch.equal(
-            residual_kernel[num_real_tokens:], residual[num_real_tokens:]
+            residual_kernel[num_real_tokens:],
+            (outer_residual + residual)[num_real_tokens:],
         )
 
 
diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py
index 9259972ac9eb..b3bc0d1cae6d 100644
--- a/tests/models/qwen4_exp/test_qsa_reference.py
+++ b/tests/models/qwen4_exp/test_qsa_reference.py
@@ -203,7 +203,14 @@ def _qsa_sparse_paged_attention_reference(
     block_table: torch.Tensor,
     token_to_req: torch.Tensor,
     softmax_scale: float,
+    k_scale: float = 1.0,
+    v_scale: float = 1.0,
 ) -> torch.Tensor:
+    """Dense reference for QSA sparse paged attention.
+
+    Mirrors the kernel's dequant: fp8-e4m3 K/V caches are dequantized with the
+    per-tensor k_scale/v_scale host floats; bf16 caches use unit scales.
+    """
     output = torch.zeros_like(q)
     repeats = q.shape[1] // k_cache.shape[2]
     page_size = k_cache.shape[1]
@@ -215,13 +222,15 @@ def _qsa_sparse_paged_attention_reference(
         request = token_to_req[row].long()
         pages = block_table[request, logical // page_size].long()
         offsets = logical % page_size
-        keys = k_cache[pages, offsets].repeat_interleave(repeats, dim=1)
-        values = v_cache[pages, offsets].repeat_interleave(repeats, dim=1)
-        scores = torch.einsum("hd,khd->hk", q[row].float(), keys.float())
-        probabilities = torch.softmax(scores * softmax_scale, dim=-1)
-        output[row] = torch.einsum("hk,khd->hd", probabilities, values.float()).to(
-            q.dtype
+        keys = (k_cache[pages, offsets].float() * k_scale).repeat_interleave(
+            repeats, dim=1
         )
+        values = (v_cache[pages, offsets].float() * v_scale).repeat_interleave(
+            repeats, dim=1
+        )
+        scores = torch.einsum("hd,khd->hk", q[row].float(), keys)
+        probabilities = torch.softmax(scores * softmax_scale, dim=-1)
+        output[row] = torch.einsum("hk,khd->hd", probabilities, values).to(q.dtype)
     return output
 
 
@@ -906,22 +915,28 @@ def test_qsa_block_expansion_correctness() -> None:
         "page_size",
         "use_prefill_config",
         "num_requests",
+        "fp8",
     ),
     [
         # Production page sizes from hybrid-cache block alignment: 784/800
         # at TP4 and 1568/1600 at TP1/TP2 (no-MTP / MTP num_spec=3). Head
         # splits are per-rank TP1/TP2/TP4; the largest batch runs both
         # use_prefill_config variants.
-        pytest.param(1, 24, 2, 1600, True, 2, id="tp1_r1"),
-        pytest.param(16, 12, 1, 1600, True, 3, id="tp2_r16"),
-        pytest.param(32, 6, 1, 800, True, 5, id="tp4_r32"),
-        pytest.param(128, 24, 2, 1568, True, 7, id="tp1_r128"),
-        pytest.param(257, 6, 1, 800, True, 13, id="tp4_r257"),
-        pytest.param(513, 6, 1, 784, True, 17, id="tp4_r513"),
-        pytest.param(700, 6, 1, 800, True, 23, id="tp4_r700"),
-        pytest.param(1024, 24, 2, 1600, True, 33, id="tp1_r1024"),
-        pytest.param(2048, 24, 2, 1600, True, 63, id="tp1_r2048_prefill"),
-        pytest.param(2048, 24, 2, 1600, False, 63, id="tp1_r2048_uniform"),
+        pytest.param(1, 24, 2, 1600, True, 2, False, id="tp1_r1"),
+        pytest.param(16, 12, 1, 1600, True, 3, False, id="tp2_r16"),
+        pytest.param(32, 6, 1, 800, True, 5, False, id="tp4_r32"),
+        pytest.param(128, 24, 2, 1568, True, 7, False, id="tp1_r128"),
+        pytest.param(257, 6, 1, 800, True, 13, False, id="tp4_r257"),
+        pytest.param(513, 6, 1, 784, True, 17, False, id="tp4_r513"),
+        pytest.param(700, 6, 1, 800, True, 23, False, id="tp4_r700"),
+        pytest.param(1024, 24, 2, 1600, True, 33, False, id="tp1_r1024"),
+        pytest.param(2048, 24, 2, 1600, True, 63, False, id="tp1_r2048_prefill"),
+        pytest.param(2048, 24, 2, 1600, False, 63, False, id="tp1_r2048_uniform"),
+        # fp8_e4m3 K/V caches on the TP1 head split.
+        pytest.param(1, 24, 2, 1600, True, 2, True, id="tp1_r1_fp8"),
+        pytest.param(128, 24, 2, 1568, True, 7, True, id="tp1_r128_fp8"),
+        pytest.param(2048, 24, 2, 1600, True, 63, True, id="tp1_r2048_prefill_fp8"),
+        pytest.param(2048, 24, 2, 1600, False, 63, True, id="tp1_r2048_uniform_fp8"),
     ],
 )
 def test_qsa_sparse_paged_attention_correctness(
@@ -931,8 +946,18 @@ def test_qsa_sparse_paged_attention_correctness(
     page_size: int,
     use_prefill_config: bool,
     num_requests: int,
+    fp8: bool,
 ) -> None:
+    """QSA sparse paged attention matches the dense reference.
+
+    fp8 only changes the K/V cache dtype (e4m3 with a per-tensor scale pair) and
+    the scales; the reference dequantizes the same cache with those scales, so
+    both paths compare the production kernel against the reference on identical
+    inputs. fp8=True additionally covers the host-side scale folding.
+    """
     torch.manual_seed(2)
+    # One QSA attention problem: bf16 Q and paged K/V, a packed selection with
+    # the trailing count column, block table and row-to-request map.
     head_dim = 256
     num_selected_pages = 64
     # Keep the newest page outside the synthetic top-k as causal headroom.
@@ -944,6 +969,7 @@ def test_qsa_sparse_paged_attention_correctness(
     q = torch.randn(
         num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16
     )
+    output_gate = torch.randn_like(q)
     kv_cache = torch.randn(
         num_cache_blocks,
         page_size,
@@ -1012,8 +1038,18 @@ def test_qsa_sparse_paged_attention_correctness(
         indexer_budget,
         logical_indices,
     )
-    assert logical_indices.shape == (num_rows, selection_width + 1)
-    scale = q.shape[-1] ** -0.5
+
+    scale = head_dim**-0.5
+
+    if fp8:
+        # A fixed non-unit pair (k != v) exercises the host-side scale folding
+        # and catches a k/v swap; scales are host floats, as the layer exposes
+        # them. Stored values are the scaled ones, as reshape_and_cache does.
+        k_scale, v_scale = 0.5, 2.0
+        k_cache = (k_cache / k_scale).to(torch.float8_e4m3fn)
+        v_cache = (v_cache / v_scale).to(torch.float8_e4m3fn)
+    else:
+        k_scale, v_scale = 1.0, 1.0
 
     actual = qsa_ops.qsa_sparse_paged_attention(
         q,
@@ -1023,6 +1059,9 @@ def test_qsa_sparse_paged_attention_correctness(
         block_table,
         token_to_req,
         use_prefill_config=use_prefill_config,
+        k_scale=k_scale,
+        v_scale=v_scale,
+        output_gate=output_gate,
     )
     expected = _qsa_sparse_paged_attention_reference(
         q,
@@ -1032,7 +1071,10 @@ def test_qsa_sparse_paged_attention_correctness(
         block_table,
         token_to_req,
         scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
     )
+    expected = expected * torch.sigmoid(output_gate)
 
     torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
 
diff --git a/tests/models/registry.py b/tests/models/registry.py
index c3973c4d0b94..0511d7ef7dfd 100644
--- a/tests/models/registry.py
+++ b/tests/models/registry.py
@@ -7,7 +7,7 @@
 
 import pytest
 from packaging.version import Version
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers import __version__ as TRANSFORMERS_VERSION
 
 from vllm.config.model import ModelDType, TokenizerMode
@@ -127,8 +127,7 @@ def check_transformers_version(
         check_min_version: bool = True,
         check_max_version: bool = True,
     ) -> str | None:
-        """
-        If the installed transformers version does not meet the requirements,
+        """If the installed transformers version does not meet the requirements,
         perform the given action.
         """
         if (
@@ -185,9 +184,7 @@ def check_available_online(
         *,
         on_fail: Literal["error", "skip"],
     ) -> None:
-        """
-        If the model is not available online, perform the given action.
-        """
+        """If the model is not available online, perform the given action."""
         if not self.is_available_online:
             msg = "Model is not available online"
 
@@ -450,7 +447,10 @@ def check_available_online(
     ),
     "MixtralForCausalLM": _HfExamplesInfo(
         "mistralai/Mixtral-8x7B-Instruct-v0.1",
-        {"tiny": "TitanML/tiny-mixtral"},
+        {
+            "tiny": "axolotl-ai-co/tiny-mixtral-30m",
+            "tiny-random": "TitanML/tiny-mixtral",
+        },
     ),
     "NemotronForCausalLM": _HfExamplesInfo("nvidia/Minitron-8B-Base"),
     "NemotronHForCausalLM": _HfExamplesInfo(
@@ -461,6 +461,10 @@ def check_available_online(
         trust_remote_code=True,
         is_available_online=False,
     ),
+    "NanbeigeForCausalLM": _HfExamplesInfo(
+        "Nanbeige/Nanbeige4.2-3B",
+        trust_remote_code=True,
+    ),
     "OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"),
     "Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),
     "Olmo3ForCausalLM": _HfExamplesInfo("allenai/Olmo-3-7B-Instruct"),
@@ -858,6 +862,9 @@ def check_available_online(
     "DeepseekV4ForConditionalGeneration": _HfExamplesInfo(
         "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp",
     ),
+    "DeepseekV41ForCausalLM": _HfExamplesInfo(
+        "deepseek-ai/DeepSeek-V4.1-Flash",
+    ),
     "Dots3NoteForCausalLM": _HfExamplesInfo(
         "dots-studio/dots3-note-prev",
         is_available_online=False,
@@ -963,6 +970,7 @@ def check_available_online(
     "HCXVisionV2ForCausalLM": _HfExamplesInfo(
         "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
         trust_remote_code=True,
+        revision="a6cdfd3464d1b767259cad23e164eaf39d3e3960",
     ),
     "HunYuanVLForConditionalGeneration": _HfExamplesInfo(
         "tencent/HunyuanOCR",
@@ -1225,7 +1233,7 @@ def check_available_online(
         # NemotronH layers are constructed via `hybrid_override_pattern`
         use_original_num_layers=True,
         hf_overrides={
-            "vision_config": PretrainedConfig(
+            "vision_config": PreTrainedConfig(
                 args={
                     "min_num_patches": 1,
                     "max_num_patches": 12,
@@ -1536,6 +1544,12 @@ def check_available_online(
         is_available_online=False,
         use_original_num_layers=True,  # DSpark has >1 draft block
     ),
+    "DSparkV41DraftModel": _HfExamplesInfo(
+        "deepseek-ai/DeepSeek-V4.1-Flash",
+        speculative_model="deepseek-ai/DeepSeek-V4.1-Flash",  # draft in mtp.*
+        is_available_online=False,
+        use_original_num_layers=True,  # DSpark has >1 draft block
+    ),
     "K3DSparkModel": _HfExamplesInfo(
         "moonshotai/Kimi-K3",
         speculative_model="Inferact/Kimi-K3-DSpark",
diff --git a/tests/models/test_adapters.py b/tests/models/test_adapters.py
index 7547dbd4ef43..62442e332e8c 100644
--- a/tests/models/test_adapters.py
+++ b/tests/models/test_adapters.py
@@ -4,7 +4,7 @@
 
 import pytest
 import torch
-from transformers import Gemma3Config, PretrainedConfig, Qwen2Config
+from transformers import Gemma3Config, PreTrainedConfig, Qwen2Config
 
 from vllm.model_executor.models.adapters import (
     _create_pooling_model_cls,
@@ -266,14 +266,14 @@ def test_resolve_num_labels_defaults_when_undeclared():
     config = _composite_config()
     assert (
         _resolve_num_labels(config, config.get_text_config())
-        == PretrainedConfig().num_labels
+        == PreTrainedConfig().num_labels
     )
 
 
 def test_resolve_num_labels_declared_on_outer_config():
     """Multimodal checkpoints keep id2label/problem_type on the top-level config."""
     config = _composite_config(outer_labels=20)
-    assert config.get_text_config().num_labels == PretrainedConfig().num_labels
+    assert config.get_text_config().num_labels == PreTrainedConfig().num_labels
     assert _resolve_num_labels(config, config.get_text_config()) == 20
 
 
diff --git a/tests/models/test_deepseek_v41_replay_start.py b/tests/models/test_deepseek_v41_replay_start.py
new file mode 100644
index 000000000000..7d4e45fe8170
--- /dev/null
+++ b/tests/models/test_deepseek_v41_replay_start.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SWA bounded replay in DeepseekV41ModelState: the batch's replay starts reach
+the sliding-window builders, and the replayed tokens' slots are padded in the
+prefix-cacheable groups only."""
+
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
+from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+DEVICE = torch.device("cuda")
+WINDOW = 4
+# Group 0 is prefix-cacheable (compressed KV), group 1 is the replayed window.
+KV_CACHE_CONFIG = SimpleNamespace(
+    kv_cache_groups=[
+        SimpleNamespace(
+            kv_cache_spec=SimpleNamespace(prefix_cacheable=True, prefix_replay_tokens=0)
+        ),
+        SimpleNamespace(
+            kv_cache_spec=SimpleNamespace(
+                prefix_cacheable=False, prefix_replay_tokens=WINDOW
+            )
+        ),
+    ]
+)
+
+
+@pytest.fixture
+def state(monkeypatch):
+    cfg = MagicMock()
+    cfg.model_config.enable_prompt_embeds = False
+    cfg.model_config.uses_mrope = False
+    cfg.model_config.is_multimodal_model = False
+    cfg.scheduler_config.max_num_seqs = 8
+    cfg.scheduler_config.max_num_batched_tokens = 64
+    seen: dict = {}
+
+    def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw):
+        seen["slot_mappings"] = slot_mappings.clone()
+        seen["extra"] = kw["model_specific_attn_metadata"]
+        return {}
+
+    monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn)
+    state = DeepseekV41ModelState(
+        cfg, SimpleNamespace(token_lookback_depth=0), None, DEVICE
+    )
+    state.seen = seen  # type: ignore[attr-defined]
+    return state
+
+
+def _batch(query_lens, seq_lens, idx_mapping, is_prefilling):
+    num_reqs, num_tokens = len(query_lens), sum(query_lens)
+    batch = InputBatch.make_dummy(num_reqs, num_tokens, InputBuffers(8, 64, DEVICE))
+    return replace(
+        batch,
+        query_start_loc=torch.tensor(
+            [0, *np.cumsum(query_lens)], dtype=torch.int32, device=DEVICE
+        ),
+        idx_mapping_np=np.array(idx_mapping, dtype=np.intp),
+        is_prefilling_np=np.array(is_prefilling, dtype=np.bool_),
+        positions=torch.cat(
+            [torch.arange(s - q, s) for q, s in zip(query_lens, seq_lens)]
+        ).to(DEVICE),
+    )
+
+
+def _prepare(state, batch):
+    slot_mappings = torch.arange(2 * batch.num_tokens, device=DEVICE).view(2, -1)
+    state.prepare_attn(
+        batch, CUDAGraphMode.NONE, (), slot_mappings, [], KV_CACHE_CONFIG
+    )
+    swa_builder = MagicMock(spec=DeepseekSparseSWAMetadataBuilder)
+    extra = state.seen["extra"]
+    replay_start = extra.get_extra_attn_kwargs(swa_builder, batch.num_reqs)[
+        "replay_start"
+    ]
+    assert extra.get_extra_attn_kwargs(MagicMock(), batch.num_reqs) == {}
+    return replay_start.tolist(), state.seen["slot_mappings"]
+
+
+def test_replayed_tokens_write_only_the_window_group(state):
+    # Request state 3 resumes a hit at 16: positions 16..19 are replayed.
+    state.add_request(3, SimpleNamespace(replay_start=16))
+    state.add_request(1, SimpleNamespace(replay_start=0))
+    batch = _batch([6, 2], [22, 2], idx_mapping=[3, 1], is_prefilling=[True, True])
+    replay_start, slots = _prepare(state, batch)
+    assert replay_start == [16, 0]
+    padded = [PAD_SLOT_ID] * WINDOW + [4, 5, 6, 7]
+    assert slots[0].tolist() == padded  # compressed KV keeps its cached rows
+    assert slots[1].tolist() == list(range(8, 16))  # the window is rebuilt
+
+
+def test_only_prefills_carry_a_replay_start(state):
+    state.add_request(0, SimpleNamespace(replay_start=16))
+    # The same request, decoding: its rows sit above the hit.
+    batch = _batch([1], [40], idx_mapping=[0], is_prefilling=[False])
+    replay_start, slots = _prepare(state, batch)
+    assert replay_start == [0]
+    assert slots.tolist() == [[0], [1]]
diff --git a/tests/models/test_deepseek_v4_dspark_rocm.py b/tests/models/test_deepseek_v4_dspark_rocm.py
new file mode 100644
index 000000000000..efa97e10d7f9
--- /dev/null
+++ b/tests/models/test_deepseek_v4_dspark_rocm.py
@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v4.amd import dspark as dspark_module
+from vllm.models.deepseek_v4.amd.dspark import DSparkDeepseekV4ForCausalLM
+
+
+def _make_uninitialized_model(confidence_head):
+    model = DSparkDeepseekV4ForCausalLM.__new__(DSparkDeepseekV4ForCausalLM)
+    object.__setattr__(
+        model,
+        "model",
+        SimpleNamespace(confidence_head=confidence_head),
+    )
+    return model
+
+
+def _prepare_loader_model(model, named_parameters):
+    object.__setattr__(
+        model,
+        "config",
+        SimpleNamespace(
+            n_routed_experts=1,
+            expert_dtype="fp4",
+            num_attention_heads=1,
+        ),
+    )
+    object.__setattr__(model, "named_parameters", lambda: named_parameters)
+
+
+def _disable_distributed_loader_paths(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        dspark_module,
+        "fused_moe_make_expert_params_mapping",
+        lambda *args, **kwargs: [],
+    )
+    monkeypatch.setattr(
+        dspark_module, "get_tensor_model_parallel_world_size", lambda: 1
+    )
+    monkeypatch.setattr(dspark_module, "get_tensor_model_parallel_rank", lambda: 0)
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_maps_enabled_confidence_head():
+    model = _make_uninitialized_model(object())
+
+    assert (
+        model._remap_dspark_name("mtp.2.confidence_head.proj.weight")
+        == "model.confidence_head.proj.weight"
+    )
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_skips_disabled_confidence_head():
+    model = _make_uninitialized_model(None)
+
+    assert model._remap_dspark_name("mtp.2.confidence_head.proj.weight") is None
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_disables_unloaded_confidence_head(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    model = _make_uninitialized_model(object())
+    _prepare_loader_model(model, [])
+    _disable_distributed_loader_paths(monkeypatch)
+
+    assert model.load_weights([]) == set()
+    assert model.model.confidence_head is None
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_loads_complete_confidence_head(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    class FakeParameter:
+        loaded_weight = None
+
+        def weight_loader(self, param, loaded_weight):
+            assert param is self
+            self.loaded_weight = loaded_weight
+
+    confidence_head = object()
+    parameter = FakeParameter()
+    model = _make_uninitialized_model(confidence_head)
+    _prepare_loader_model(
+        model,
+        [("model.confidence_head.proj.weight", parameter)],
+    )
+    _disable_distributed_loader_paths(monkeypatch)
+    loaded_weight = torch.tensor([[1.0, 2.0]])
+
+    loaded = model.load_weights([("mtp.2.confidence_head.proj.weight", loaded_weight)])
+
+    assert loaded == {"model.confidence_head.proj.weight"}
+    assert parameter.loaded_weight is loaded_weight
+    assert model.model.confidence_head is confidence_head
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_confidence_is_probability():
+    class ConfidenceHead:
+        def __call__(self, head_hidden, markov_embed):
+            return (head_hidden[:, 0] + markov_embed[:, 0]).float()
+
+    model = _make_uninitialized_model(ConfidenceHead())
+    head_hidden = torch.tensor([[0.0], [1.0]], dtype=torch.bfloat16)
+    markov_embed = torch.tensor([[0.0], [-2.0]], dtype=torch.bfloat16)
+
+    confidence = model.compute_confidence(head_hidden, markov_embed)
+
+    torch.testing.assert_close(
+        confidence,
+        torch.sigmoid(torch.tensor([0.0, -1.0])),
+    )
+    assert torch.all((confidence >= 0) & (confidence <= 1))
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_confidence_requires_a_head():
+    model = _make_uninitialized_model(None)
+    empty = torch.zeros((1, 1))
+
+    with pytest.raises(RuntimeError, match="confidence_head"):
+        model.compute_confidence(empty, empty)
diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py
index 77b05101d685..99d661e75598 100644
--- a/tests/models/test_deepseek_v4_fi_moe_ep.py
+++ b/tests/models/test_deepseek_v4_fi_moe_ep.py
@@ -118,7 +118,7 @@ def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend):
     [["KimiK3ForConditionalGeneration"], ["MixtralForCausalLM"]],
 )
 def test_native_deep_gemm_mega_moe_not_arch_gated(architectures):
-    """vLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it);
+    """VLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it);
     models validate their own constraints at construction time."""
     validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures)
 
diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py
index 4e95904f6b2d..f87dd75fb220 100644
--- a/tests/models/test_deepseek_v4_mega_moe.py
+++ b/tests/models/test_deepseek_v4_mega_moe.py
@@ -15,10 +15,15 @@
     DeepseekV4MegaMoEExperts,
     DeepseekV4MoE,
     make_deepseek_v4_expert_params_mapping,
+    prepare_mega_gate_routing_metadata,
 )
 from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP
 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
+from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
+from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE
 from vllm.platforms import current_platform
+from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
+from vllm.utils.torch_utils import set_default_torch_dtype
 
 pytestmark = pytest.mark.skipif(
     not current_platform.is_cuda(),
@@ -26,6 +31,253 @@
 )
 
 
+@pytest.fixture
+def v41_moe_config(dist_init):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            dtype=torch.bfloat16,
+            hf_config=DeepseekV41Config(
+                text_config=dict(
+                    hidden_size=128,
+                    num_hidden_layers=2,
+                    n_routed_experts=8,
+                    num_experts_per_tok=2,
+                    dspark_n_routed_experts=4,
+                    dspark_num_experts_per_tok=3,
+                    n_shared_experts=1,
+                    moe_intermediate_size=128,
+                    hidden_act="silu",
+                    swiglu_limit=10.0,
+                    norm_topk_prob=True,
+                    topk_method="noaux_tc",
+                    routed_scaling_factor=1.5,
+                ),
+            ),
+        ),
+        quant_config=None,
+        kernel_config=SimpleNamespace(moe_backend="deep_gemm_mega_moe"),
+        parallel_config=SimpleNamespace(
+            enable_expert_parallel=True,
+            enable_eplb=False,
+            eplb_config=SimpleNamespace(num_redundant_experts=0),
+        ),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
+        compilation_config=SimpleNamespace(static_forward_context={}),
+    )
+
+
+@pytest.mark.parametrize("use_cudagraph", [False, True])
+@pytest.mark.parametrize("above_threshold", [False, True])
+@pytest.mark.parametrize("vision", [False, True])
+@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 384, 6), (2, 128, 3)])
+def test_deepseek_v41_moe_routes_without_hash_table(
+    v41_moe_config,
+    monkeypatch,
+    vision,
+    layer_id,
+    num_experts,
+    top_k,
+    above_threshold,
+    use_cudagraph,
+):
+    """Main and draft layers select experts by score and preserve image routing."""
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("DeepGEMM Mega Gate requires SM100")
+
+    config = v41_moe_config.model_config.hf_config
+    config.hidden_size = 256
+    config.n_routed_experts = 384
+    config.num_experts_per_tok = 6
+    config.dspark_n_routed_experts = 128
+    num_tokens = (1 if num_experts == 128 else 16) + int(above_threshold)
+    config.vision_n_layers = int(vision)
+    with (
+        set_default_torch_dtype(v41_moe_config.model_config.dtype),
+        torch.device("cuda"),
+    ):
+        moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn")
+        hidden_states = torch.randn(
+            num_tokens, config.hidden_size, dtype=torch.bfloat16
+        )
+        input_ids = (
+            torch.tensor(
+                [42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257]
+            ).repeat((num_tokens + 3) // 4)[:num_tokens]
+            if vision
+            else None
+        )
+    assert moe.gate.tid2eid is None
+    assert moe.gate.weight.shape == (num_experts, config.hidden_size)
+    assert isinstance(moe.experts, DeepseekV4MegaMoEExperts)
+    assert moe.experts.w13_weight.shape == (num_experts, 256, 128)
+    assert moe.experts.top_k == top_k
+    assert (config.n_routed_experts, config.num_experts_per_tok) == (384, 6)
+
+    with torch.no_grad():
+        moe.gate.weight.normal_(std=0.01)
+        moe.gate.e_score_correction_bias.copy_(torch.arange(num_experts, device="cuda"))
+        if vision:
+            moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias)
+
+    scores = torch.nn.functional.softplus(
+        torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32)
+    ).sqrt()
+    bias = moe.gate.e_score_correction_bias
+    if vision:
+        image_mask = input_ids == IMAGE_SENTINEL_BASE_ID
+        bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias)
+    expected_ids = (scores + bias).topk(top_k, dim=-1).indices
+    expected_weights = scores.gather(1, expected_ids)
+    expected_weights *= config.routed_scaling_factor / expected_weights.sum(
+        dim=-1, keepdim=True
+    )
+
+    routed = {}
+
+    def check_routing(x, weights, ids, *, activation_clamp):
+        routed["ids"], routed["weights"] = ids, weights
+        assert activation_clamp == config.swiglu_limit
+        return x.clone()
+
+    monkeypatch.setattr(moe.experts, "forward", check_routing)
+    monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x)
+    routing_input_ids = (
+        input_ids
+        if input_ids is not None
+        else torch.zeros(num_tokens, dtype=torch.int64, device="cuda")
+    )
+    metadata = prepare_mega_gate_routing_metadata(
+        routing_input_ids,
+        has_hash_routing=False,
+        image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID if vision else None,
+    )
+    if use_cudagraph:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                moe(hidden_states, input_ids, metadata)
+        stream.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            output = moe(hidden_states, input_ids, metadata)
+        graph.replay()
+        torch.accelerator.synchronize()
+    else:
+        output = moe(hidden_states, input_ids, metadata)
+
+    ids, weights = routed["ids"], routed["weights"]
+    torch.testing.assert_close(ids, expected_ids)
+    torch.testing.assert_close(
+        weights,
+        expected_weights,
+        rtol=1e-3,
+        atol=1e-4,
+    )
+    torch.testing.assert_close(output, 3 * hidden_states)
+
+
+@pytest.mark.parametrize(
+    "draft_experts,draft_top_k,expected", [(0, 0, (8, 2)), (4, 3, (4, 3))]
+)
+def test_deepseek_v41_fused_moe_uses_draft_counts_or_main_defaults(
+    v41_moe_config, monkeypatch, draft_experts, draft_top_k, expected
+):
+    config = v41_moe_config.model_config.hf_config
+    config.dspark_n_routed_experts = draft_experts
+    config.dspark_num_experts_per_tok = draft_top_k
+    v41_moe_config.kernel_config.moe_backend = "auto"
+    captured = {}
+
+    def make_experts(**kwargs):
+        captured.update(kwargs)
+        return torch.nn.Identity()
+
+    monkeypatch.setattr(
+        "vllm.models.deepseek_v4.nvidia.model.FusedMoEFactory", make_experts
+    )
+    moe = DeepseekV41MoE(v41_moe_config, prefix="model.layers.2.ffn")
+    assert moe.gate.weight.shape == (expected[0], 128)
+    assert (captured["num_experts"], captured["top_k"]) == expected
+    assert captured["hash_indices_table"] is None
+
+
+def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config):
+    config = v41_moe_config.model_config.hf_config
+    config.num_hash_layers = 1
+    config.vocab_size = 32
+    moe = DeepseekV4MoE(
+        v41_moe_config,
+        prefix="model.layers.0.ffn",
+        num_hash_layers=config.num_hash_layers,
+    )
+    assert moe.gate.tid2eid.shape == (32, 2)
+    assert moe.gate.e_score_correction_bias is None
+    with pytest.raises(ValueError, match="hash MoE routing requires input_ids"):
+        moe(torch.zeros(1, 128))
+
+
+@pytest.mark.parametrize("num_tokens", [16, 17])
+def test_deepseek_v4_mega_gate_hash_routing_correctness(
+    v41_moe_config, monkeypatch, num_tokens
+):
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("DeepGEMM Mega Gate requires SM100")
+
+    config = v41_moe_config.model_config.hf_config
+    config.hidden_size = 256
+    config.num_hash_layers = 1
+    config.vocab_size = 32
+    with (
+        set_default_torch_dtype(v41_moe_config.model_config.dtype),
+        torch.device("cuda"),
+    ):
+        moe = DeepseekV4MoE(
+            v41_moe_config,
+            prefix="model.layers.0.ffn",
+            num_hash_layers=config.num_hash_layers,
+        )
+        hidden_states = torch.randn(
+            num_tokens, config.hidden_size, dtype=torch.bfloat16
+        )
+        input_ids = torch.arange(num_tokens)
+
+    token_ids = torch.arange(config.vocab_size, device="cuda")
+    fixed_ids = torch.stack(
+        (
+            token_ids % config.n_routed_experts,
+            (token_ids + 3) % config.n_routed_experts,
+        ),
+        dim=1,
+    )
+    with torch.no_grad():
+        moe.gate.weight.normal_(std=0.01)
+        moe.gate.tid2eid.copy_(fixed_ids)
+
+    expected_ids = fixed_ids[input_ids]
+    scores = torch.nn.functional.softplus(
+        torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32)
+    ).sqrt()
+    expected_weights = scores.gather(1, expected_ids)
+    expected_weights *= config.routed_scaling_factor / expected_weights.sum(
+        dim=-1, keepdim=True
+    )
+
+    def check_routing(x, weights, ids, *, activation_clamp):
+        torch.testing.assert_close(ids, expected_ids)
+        torch.testing.assert_close(weights, expected_weights, rtol=1e-3, atol=1e-4)
+        return x.clone()
+
+    monkeypatch.setattr(moe.experts, "forward", check_routing)
+    monkeypatch.setattr(moe.shared_experts, "forward", torch.zeros_like)
+    metadata = prepare_mega_gate_routing_metadata(
+        input_ids,
+        has_hash_routing=True,
+        image_sentinel_base_id=None,
+    )
+    torch.testing.assert_close(moe(hidden_states, input_ids, metadata), hidden_states)
+
+
 def test_deepseek_v4_mega_moe_expert_mapping():
     mapping = make_deepseek_v4_expert_params_mapping(2)
 
@@ -169,21 +421,76 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
     assert torch.count_nonzero(experts.w13_weight[1]) == 0
 
 
+@pytest.mark.parametrize("intermediate_size", [512, 2304])
+def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions(
+    monkeypatch, intermediate_size
+):
+    """Keep native widths so V4.1 routed and shared experts can fuse."""
+    experts = DeepseekV4MegaMoEExperts(
+        SimpleNamespace(
+            scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
+            compilation_config=SimpleNamespace(static_forward_context={}),
+        ),
+        num_experts=2,
+        num_local_experts=2,
+        experts_start_idx=0,
+        top_k=1,
+        hidden_size=128,
+        intermediate_size=intermediate_size,
+    )
+    originals = []
+    for param in (
+        experts.w13_weight,
+        experts.w13_weight_scale,
+        experts.w2_weight,
+        experts.w2_weight_scale,
+    ):
+        param.data.random_(1, 128)
+        originals.append(param.data.clone())
+    monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None)
+    monkeypatch.setattr(
+        "vllm.utils.deep_gemm._import_deep_gemm",
+        lambda: SimpleNamespace(
+            transform_sf_into_required_layout=lambda sf, *args: sf,
+            transform_weights_for_mega_moe=lambda l1, l2: (l1, l2),
+        ),
+    )
+
+    experts.finalize_weights()
+    assert experts.intermediate_size == intermediate_size
+    transformed = (*experts._transformed_l1_weights, *experts._transformed_l2_weights)
+    for actual, original in zip(transformed, originals):
+        expected = (
+            original
+            if actual.dtype == torch.int8
+            else experts._ue8m0_uint8_to_float(original)
+        )
+        assert torch.equal(actual.view(expected.dtype), expected)
+
+    transformed_l1 = experts._transformed_l1_weights
+    experts.finalize_weights()
+    assert experts._transformed_l1_weights is transformed_l1
+
+
 @pytest.mark.parametrize(
-    "hidden_size,intermediate_size,block_size,fused",
+    "hidden_size,intermediate_size,block_size,mxfp8",
     [
-        pytest.param(128, 512, 128, True, id="aligned-block128"),
+        pytest.param(128, 512, 128, False, id="aligned-block128"),
+        pytest.param(128, 512, 128, True, id="aligned-block128-mxfp8"),
         pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"),
+        pytest.param(128, 2304, 128, False, id="native-block128"),
+        pytest.param(5120, 2304, 32, True, id="deepseek-v41-flash-mxfp8"),
     ],
 )
 def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
-    monkeypatch, hidden_size, intermediate_size, block_size, fused
+    monkeypatch, hidden_size, intermediate_size, block_size, mxfp8
 ):
-    """V4.1-Flash's routed padding must preserve the serial shared-expert fallback."""
+    """Shared fusion preserves checkpoint weights and scales at native widths."""
 
     class FakeDeepGemm:
         transformed_dims: list[tuple[int, int]] = []
         scale_inputs: list[tuple[int, ...]] = []
+        scale_values: list[torch.Tensor] = []
 
         @staticmethod
         def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs):
@@ -208,6 +515,7 @@ def fp8_fp4_mega_moe(
         @classmethod
         def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs):
             cls.scale_inputs.append(tuple(sf.shape))
+            cls.scale_values.append(sf)
             return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32)
 
         @classmethod
@@ -235,7 +543,10 @@ def transform_weights_for_mega_moe(cls, l1_weights, l2_weights):
 
     def fp8_parameter(*shape):
         return torch.nn.Parameter(
-            torch.empty(*shape, dtype=torch.float8_e4m3fn), requires_grad=False
+            torch.empty(*shape, dtype=torch.uint8)
+            .random_(1, 119)
+            .view(torch.float8_e4m3fn),
+            requires_grad=False,
         )
 
     def scale_parameter(*shape, dtype=torch.int32):
@@ -261,31 +572,49 @@ def scale_parameter(*shape, dtype=torch.int32):
             ),
         ),
     )
+    if mxfp8:
+        for linear in (shared_experts.gate_up_proj, shared_experts.down_proj):
+            linear.weight_block_size = (1, 32)
+            linear.weight_scale = torch.nn.Parameter(
+                linear.weight_scale_inv.view(torch.uint8)
+                .repeat_interleave(block_size, dim=0)
+                .repeat_interleave(block_size // 32, dim=1),
+                requires_grad=False,
+            )
+            del linear.weight_scale_inv
+    originals = []
+    for linear in (shared_experts.gate_up_proj, shared_experts.down_proj):
+        scale = linear.weight_scale if mxfp8 else linear.weight_scale_inv
+        scale.data.view(torch.uint8).random_(124, 130)
+        block_m, block_k = linear.weight_block_size
+        scale_1x32 = (
+            experts._ue8m0_uint8_to_float(scale.view(torch.uint8))
+            .repeat_interleave(block_m, dim=0)
+            .repeat_interleave(block_k // 32, dim=1)
+        )
+        originals.append((linear.weight.view(torch.uint8).clone(), scale_1x32))
     monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
 
     original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()
-    original_down_ptr = shared_experts.down_proj.weight.data_ptr()
     experts.finalize_weights(shared_experts)
 
+    assert experts.has_fused_shared_experts
+    assert experts.intermediate_size == intermediate_size
+    assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
     assert FakeDeepGemm.scale_inputs[-2:] == [
         (1, 2 * intermediate_size, hidden_size // 32),
         (1, hidden_size, intermediate_size // 32),
     ]
-    assert experts.has_fused_shared_experts is fused
-    if not fused:
-        assert experts.intermediate_size == 2560
-        assert FakeDeepGemm.transformed_dims == [(3, 3)]
-        assert experts.num_shared_experts == 0
-        assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr
-        assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr
-        assert shared_experts.gate_up_proj.weight.shape == (
-            2 * intermediate_size,
-            hidden_size,
-        )
-        assert shared_experts.down_proj.weight.shape == (hidden_size, intermediate_size)
-        return
-
-    assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
+    for transformed, actual_scale, (weight, scale) in zip(
+        (
+            experts._transformed_shared_l1_weights,
+            experts._transformed_shared_l2_weights,
+        ),
+        FakeDeepGemm.scale_values[-2:],
+        originals,
+    ):
+        assert torch.equal(transformed[0].view(torch.uint8), weight)
+        assert torch.equal(actual_scale[0], scale)
     assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
     assert (
         experts._transformed_shared_l1_weights[0].data_ptr()
@@ -296,12 +625,17 @@ def scale_parameter(*shape, dtype=torch.int32):
         == shared_experts.down_proj.weight.data_ptr()
     )
 
+    transformed_l1 = experts._transformed_shared_l1_weights
+    experts.finalize_weights(shared_experts)
+    assert experts._transformed_shared_l1_weights is transformed_l1
+
 
 @pytest.mark.parametrize("fused", [False, True])
 def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert(
     monkeypatch, fused
 ):
     class FakeGate(torch.nn.Module):
+        weight = torch.empty(2, 128)
         tid2eid = None
         e_score_correction_bias = None
 
@@ -352,12 +686,14 @@ def forward(self, hidden_states):
     not torch.cuda.is_available(),
     reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
 )
-def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
+@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384])
+@pytest.mark.parametrize("hidden_size", [256, 5120])
+def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact(
+    num_tokens, hidden_size
+):
     from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
 
     device = torch.device("cuda")
-    num_tokens = 7
-    hidden_size = 256
     top_k = 8
 
     generator = torch.Generator(device=device)
@@ -512,6 +848,7 @@ def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m):
     assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
     assert torch.equal(fused_x_sf, ref_x_sf)
     assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated])
+    assert torch.all(fused_shared_x_sf[~populated] == -1)
 
 
 def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast():
@@ -549,11 +886,14 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe():
     not torch.cuda.is_available(),
     reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
 )
-def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding():
+@pytest.mark.parametrize("nonfinite_padding", [False, True])
+@pytest.mark.parametrize("num_tokens", [7, 65, 135])
+def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(
+    nonfinite_padding, num_tokens
+):
     from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
 
     device = torch.device("cuda")
-    num_tokens = 7
     hidden_size = 256
     top_k = 8
 
@@ -581,10 +921,11 @@ def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding():
         dtype=torch.float32,
         generator=generator,
     )
-    is_padding = torch.tensor(
-        [False, True, False, False, True, False, True],
-        device=device,
-    )
+    is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool)
+    is_padding[1::3] = True
+    is_padding[-1] = True
+    if nonfinite_padding:
+        topk_weights[is_padding] = float("nan")
 
     ref_x, ref_x_sf = per_token_cast_to_fp8(
         hidden_states,
diff --git a/tests/models/test_deepseek_v4_rocm_wo_a.py b/tests/models/test_deepseek_v4_rocm_wo_a.py
new file mode 100644
index 000000000000..69c6ed95307c
--- /dev/null
+++ b/tests/models/test_deepseek_v4_rocm_wo_a.py
@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v4.amd.rocm import _wo_a_block_scale_to_e8m0
+
+
+def test_wo_a_block_scale_to_e8m0_from_float():
+    scale = torch.tensor([[0.5, 1.0, 2.0, 4.0]], dtype=torch.float32)
+
+    encoded = _wo_a_block_scale_to_e8m0(scale)
+
+    assert encoded is not None
+    torch.testing.assert_close(
+        encoded,
+        torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8),
+    )
+    assert encoded.is_contiguous()
+
+
+def test_wo_a_block_scale_to_e8m0_preserves_encoded_scales():
+    raw = torch.tensor([[125, 127, 131]], dtype=torch.uint8)
+    encoded = raw.view(torch.float8_e8m0fnu)
+
+    converted = _wo_a_block_scale_to_e8m0(encoded)
+
+    assert converted is not None
+    torch.testing.assert_close(converted, raw)
+
+
+@pytest.mark.parametrize(
+    "scale",
+    [
+        torch.tensor([[0.0, 1.0]]),
+        torch.tensor([[-1.0, 1.0]]),
+        torch.tensor([[0.75, 1.0]]),
+        torch.tensor([[float("inf"), 1.0]]),
+        torch.ones(1, dtype=torch.int32),
+    ],
+)
+def test_wo_a_block_scale_to_e8m0_rejects_invalid_scales(scale: torch.Tensor):
+    assert _wo_a_block_scale_to_e8m0(scale) is None
diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py
new file mode 100644
index 000000000000..098e5185fcea
--- /dev/null
+++ b/tests/models/test_deepseek_v4_vl_rocm.py
@@ -0,0 +1,354 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.config import ParallelConfig
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-specific tests"
+)
+
+
+def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None:
+    from vllm.config import CacheConfig
+    from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype
+
+    cache_config = CacheConfig()
+
+    resolved_dtype, torch_dtype = _resolve_dsv4_kv_cache_dtype(
+        use_fp8_ds_mla_layout=True,
+        kv_cache_dtype=cache_config.cache_dtype,
+        cache_config=cache_config,
+    )
+
+    assert resolved_dtype == "fp8_ds_mla"
+    assert torch_dtype is torch.uint8
+    assert cache_config.cache_dtype == "fp8_ds_mla"
+
+
+def test_rocm_packed_kv_cache_rejects_unquantized_dtype() -> None:
+    from vllm.config import CacheConfig
+    from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype
+
+    cache_config = CacheConfig(cache_dtype="bfloat16")
+
+    with pytest.raises(ValueError, match="only supports fp8 kv-cache"):
+        _resolve_dsv4_kv_cache_dtype(
+            use_fp8_ds_mla_layout=True,
+            kv_cache_dtype=cache_config.cache_dtype,
+            cache_config=cache_config,
+        )
+
+
+def test_vl_mapper_preserves_rocm_weight_mapping() -> None:
+    from vllm.models.deepseek_v4.amd.model import _make_deepseek_v4_weights_mapper
+    from vllm.models.deepseek_v4.common.vl_model import (
+        _make_deepseek_v4_vl_weights_mapper,
+    )
+
+    text_mapper = _make_deepseek_v4_weights_mapper("fp4", fuse_shared_experts=True)
+    mapper = _make_deepseek_v4_vl_weights_mapper(text_mapper, image_enabled=True)
+
+    assert mapper._map_name("layers.3.attn.wq_a.input_scale") == (
+        "language_model.model.layers.3.attn.wq_a.input_scale_2"
+    )
+    assert mapper._map_name("layers.3.ffn.shared_experts.w2.weight") == (
+        "language_model.model.layers.3.ffn.shared_experts.w2.weight"
+    )
+    assert mapper._map_name("head.weight") == "language_model.lm_head.weight"
+
+
+def test_rocm_moe_wires_vision_routing_on_hash_and_regular_layers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm.models.deepseek_v4.amd import model as rocm_model
+    from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
+
+    captured: list[dict] = []
+
+    class FakeGate(nn.Module):
+        def __init__(self, **kwargs) -> None:
+            super().__init__()
+
+    def fake_factory(**kwargs):
+        captured.append(kwargs)
+        return nn.Identity()
+
+    monkeypatch.setattr(rocm_model, "GateLinear", FakeGate)
+    monkeypatch.setattr(rocm_model, "FusedMoEFactory", fake_factory)
+    monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_rank", lambda: 0)
+
+    config = SimpleNamespace(
+        hidden_size=16,
+        n_routed_experts=4,
+        num_experts_per_tok=2,
+        moe_intermediate_size=8,
+        swiglu_limit=None,
+        norm_topk_prob=True,
+        scoring_func="sqrtsoftplus",
+        num_hash_layers=1,
+        vocab_size=32,
+        topk_method="noaux_tc",
+        vision_n_layers=1,
+        n_shared_experts=None,
+    )
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=config),
+        quant_config=None,
+        parallel_config=ParallelConfig(),
+    )
+
+    hash_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.0.ffn")
+    regular_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.1.ffn")
+
+    assert hash_moe.gate.tid2eid is not None
+    assert regular_moe.gate.tid2eid is None
+    for moe, factory_kwargs in zip((hash_moe, regular_moe), captured, strict=True):
+        assert moe.gate.e_score_correction_bias is not None
+        assert moe.gate.bias_vl is not None
+        assert factory_kwargs["bias_vl"] is moe.gate.bias_vl
+        assert factory_kwargs["image_sentinel_lo"] == IMAGE_SENTINEL_BASE_ID
+
+
+def test_rocm_fse_rejects_data_parallel_deployment(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm.models.deepseek_v4.amd.model import _fuse_shared_experts_enabled
+
+    config = SimpleNamespace(n_shared_experts=1)
+
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1")
+    with pytest.raises(ValueError, match="data parallelism"):
+        _fuse_shared_experts_enabled(config, ParallelConfig(data_parallel_size=2))
+
+    # Expert parallelism already disables FSE: no raise, FSE stays off.
+    assert not _fuse_shared_experts_enabled(
+        config, ParallelConfig(data_parallel_size=2, enable_expert_parallel=True)
+    )
+    assert _fuse_shared_experts_enabled(config, ParallelConfig())
+
+
+def test_rocm_mtp_forwards_input_ids_for_vision_routing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm.models.deepseek_v4.amd import mtp as rocm_mtp
+
+    hidden_size = 4
+    hc_mult = 2
+
+    class FakeNorm(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(hidden_size))
+            self.variance_epsilon = 1e-6
+
+    class FakeMTPBlock(nn.Module):
+        use_fused_mhc = False
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.seen_input_ids: torch.Tensor | None = None
+
+        def forward(
+            self,
+            *,
+            positions: torch.Tensor,
+            x: torch.Tensor,
+            input_ids: torch.Tensor | None,
+        ):
+            self.seen_input_ids = input_ids
+            return x, None, None, None
+
+    def passthrough_mtp_input(
+        inputs_embeds: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        *args,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return inputs_embeds, previous_hidden_states
+
+    monkeypatch.setattr(
+        rocm_mtp, "_FUSED_MTP_INPUT_RMSNORM_KERNEL", passthrough_mtp_input
+    )
+
+    layer = object.__new__(rocm_mtp.DeepSeekV4MultiTokenPredictorLayer)
+    nn.Module.__init__(layer)
+    layer.config = SimpleNamespace(hidden_size=hidden_size)
+    layer.hc_mult = hc_mult
+    layer.enorm = FakeNorm()
+    layer.hnorm = FakeNorm()
+    layer.e_proj = nn.Identity()
+    layer.h_proj = nn.Identity()
+    layer.mtp_block = FakeMTPBlock()
+
+    input_ids = torch.tensor([11, 12])
+    positions = torch.tensor([3, 4])
+    inputs_embeds = torch.arange(8, dtype=torch.float32).view(2, hidden_size)
+    previous_hidden_states = torch.arange(16, dtype=torch.float32).view(2, -1)
+
+    output = layer(
+        input_ids,
+        positions,
+        previous_hidden_states,
+        inputs_embeds,
+    )
+
+    assert layer.mtp_block.seen_input_ids is input_ids
+    expected = previous_hidden_states.view(2, hc_mult, hidden_size)
+    expected = expected + inputs_embeds.unsqueeze(-2)
+    torch.testing.assert_close(output, expected.flatten(1))
+
+
+def test_rocm_compute_logits_local_skips_gather() -> None:
+    from vllm.models.deepseek_v4.amd.model import DeepseekV4ForCausalLM
+
+    calls: list[tuple[nn.Module, torch.Tensor, bool]] = []
+
+    def logits_processor(
+        lm_head: nn.Module, hidden_states: torch.Tensor, *, skip_gather: bool = False
+    ) -> torch.Tensor:
+        calls.append((lm_head, hidden_states, skip_gather))
+        return hidden_states + 1
+
+    model = object.__new__(DeepseekV4ForCausalLM)
+    nn.Module.__init__(model)
+    model.lm_head = nn.Identity()
+    model.logits_processor = logits_processor
+    hidden_states = torch.tensor([4.0])
+
+    result = model.compute_logits_local(hidden_states)
+
+    assert torch.equal(result, torch.tensor([5.0]))
+    assert calls == [(model.lm_head, hidden_states, True)]
+
+
+class _FakeLanguageModel(nn.Module):
+    finalizes_weights_during_load = False
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.tensor_a = nn.Parameter(torch.zeros(1))
+        self.tensor_c = nn.Parameter(torch.zeros(1))
+        self.finalized_values: list[tuple[float, float]] = []
+
+    def process_weights_after_loading(self) -> None:
+        self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item()))
+
+    def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states + 1
+
+
+def test_vl_wrapper_streams_then_delegates_finalization() -> None:
+    from vllm.models.deepseek_v4.common.vl_model import (
+        DeepseekV4ForConditionalGeneration,
+    )
+
+    model = object.__new__(DeepseekV4ForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.language_model = _FakeLanguageModel()
+    model.vision = nn.Module()
+    model.vision.tensor_b = nn.Parameter(torch.zeros(1))
+    model.hf_to_vllm_mapper = WeightsMapper()
+
+    def interleaved_weights():
+        yield "language_model.tensor_a", torch.tensor([1.0])
+        assert model.language_model.tensor_a.item() == 1.0
+        yield "vision.tensor_b", torch.tensor([2.0])
+        assert model.vision.tensor_b.item() == 2.0
+        yield "language_model.tensor_c", torch.tensor([3.0])
+
+    loaded = model.load_weights(interleaved_weights())
+
+    assert loaded == {
+        "language_model.tensor_a",
+        "vision.tensor_b",
+        "language_model.tensor_c",
+    }
+    assert model.language_model.finalized_values == []
+
+    model.process_weights_after_loading()
+
+    assert model.language_model.finalized_values == [(1.0, 3.0)]
+    assert torch.equal(
+        model.compute_logits_local(torch.tensor([4.0])), torch.tensor([5.0])
+    )
+    model.process_weights_after_loading()
+    assert model.language_model.finalized_values == [(1.0, 3.0)]
+
+
+class _FakeFinalizingLanguageModel(_FakeLanguageModel):
+    finalizes_weights_during_load = True
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.load_calls = 0
+
+    def load_weights(self, weights) -> set[str]:
+        self.load_calls += 1
+        loaded = set()
+        for name, value in weights:
+            getattr(self, name).data.copy_(value)
+            loaded.add(name)
+        self.process_weights_after_loading()
+        return loaded
+
+
+def test_vl_wrapper_groups_child_that_finalizes_during_load() -> None:
+    from vllm.models.deepseek_v4.common.vl_model import (
+        DeepseekV4ForConditionalGeneration,
+    )
+
+    model = object.__new__(DeepseekV4ForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.language_model = _FakeFinalizingLanguageModel()
+    model.vision = nn.Module()
+    model.vision.tensor_b = nn.Parameter(torch.zeros(1))
+    model.hf_to_vllm_mapper = WeightsMapper()
+
+    loaded = model.load_weights(
+        iter(
+            (
+                ("language_model.tensor_a", torch.tensor([1.0])),
+                ("vision.tensor_b", torch.tensor([2.0])),
+                ("language_model.tensor_c", torch.tensor([3.0])),
+            )
+        )
+    )
+
+    assert loaded == {
+        "language_model.tensor_a",
+        "vision.tensor_b",
+        "language_model.tensor_c",
+    }
+    assert model.language_model.load_calls == 1
+    assert model.language_model.finalized_values == [(1.0, 3.0)]
+
+    # The framework's later model-level hook must not double-finalize a child
+    # which already completed this work in load_weights.
+    model.process_weights_after_loading()
+    assert model.language_model.finalized_values == [(1.0, 3.0)]
+
+
+def test_vl_wrapper_dummy_load_delegates_finalization() -> None:
+    from vllm.models.deepseek_v4.common.vl_model import (
+        DeepseekV4ForConditionalGeneration,
+    )
+
+    model = object.__new__(DeepseekV4ForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.language_model = _FakeFinalizingLanguageModel()
+
+    # DummyModelLoader bypasses model.load_weights(), so no finalized marker
+    # exists and the framework-level hook must still delegate to the child.
+    model.process_weights_after_loading()
+    assert model.language_model.finalized_values == [(0.0, 0.0)]
+    model.process_weights_after_loading()
+    assert model.language_model.finalized_values == [(0.0, 0.0)]
diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py
index 01ee0b0533e5..3e70450f2ce0 100644
--- a/tests/models/test_dspark_mla.py
+++ b/tests/models/test_dspark_mla.py
@@ -11,6 +11,7 @@
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
 from vllm.model_executor.models.registry import ModelRegistry
+from vllm.models.deepseek_v4.nvidia import dspark as dsv4_dspark
 from vllm.models.kimi_k3.nvidia import dspark_mla
 from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM, K3DSparkModel
 
@@ -186,3 +187,151 @@ def test_context_kv_weights_are_loaded_as_merged_linear_shards():
     assert [weight.shard_id for _, weight in mapped] == [1, 0, 1, 1]
     assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr()
     assert mapped[2][1].data_ptr() == mapped[3][1].data_ptr()
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "scale_dtype", [torch.uint8, torch.float8_e8m0fnu, torch.float32]
+)
+@pytest.mark.parametrize(
+    ("checkpoint_name", "runtime_module", "shard_id"),
+    [
+        ("mtp.0.attn.wq_a.scale", "model.layers.0.attn.fused_wqa_wkv", 0),
+        ("mtp.0.attn.wkv.scale", "model.layers.0.attn.fused_wqa_wkv", 1),
+        ("mtp.0.main_proj.scale", "model.main_proj", None),
+        (
+            "mtp.0.ffn.shared_experts.w1.scale",
+            "model.layers.0.ffn.shared_experts.gate_up_proj",
+            0,
+        ),
+        (
+            "mtp.0.ffn.shared_experts.w2.scale",
+            "model.layers.0.ffn.shared_experts.down_proj",
+            None,
+        ),
+    ],
+)
+def test_v41_dspark_loads_linear_scales(
+    monkeypatch, scale_dtype, checkpoint_name, runtime_module, shard_id
+):
+    """Checkpoint ``.scale`` maps to the quant method's scale parameter and
+    loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static
+    loader (see tests/quantization/test_modelopt.py), not in load_weights."""
+    from vllm.models.deepseek_v41.nvidia import dspark
+
+    mxfp8 = scale_dtype != torch.float32
+    scale_name = "weight_scale" if mxfp8 else "weight_scale_inv"
+    runtime_name = f"{runtime_module}.{scale_name}"
+    raw = torch.tensor([[120, 127], [128, 130]], dtype=torch.uint8)
+    checkpoint_scale = raw.view(scale_dtype) if mxfp8 else raw.float()
+    param = nn.Parameter(torch.empty_like(checkpoint_scale), requires_grad=False)
+    shards = []
+
+    def load_scale(param, weight, *args):
+        shards.append(args)
+        assert weight.dtype == checkpoint_scale.dtype
+        param.copy_(weight)
+
+    param.weight_loader = load_scale
+    draft = SimpleNamespace(
+        config=SimpleNamespace(num_attention_heads=4, n_routed_experts=1),
+        quant_config=SimpleNamespace(
+            weight_block_size=[32, 32] if mxfp8 else [128, 128]
+        ),
+        linear_scale_name=scale_name,
+        pad_shared_expert=False,
+        model=SimpleNamespace(
+            layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))],
+            confidence_head=None,
+        ),
+        named_parameters=lambda: [(runtime_name, param)],
+        process_weights_after_loading=lambda: None,
+    )
+    draft._remap_dspark_name = lambda name: (
+        dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name)
+    )
+    monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        dspark, "fused_moe_make_expert_params_mapping", lambda *a, **kw: []
+    )
+
+    loaded = dspark.DSparkDeepseekV4ForCausalLM.load_weights(
+        draft, [(checkpoint_name, checkpoint_scale)]
+    )
+
+    assert loaded == {runtime_name}
+    assert shards == [() if shard_id is None else (shard_id,)]
+    torch.testing.assert_close(param, checkpoint_scale)
+
+
+def test_dsv4_context_wkv_weights_are_duplicated_by_draft_layer():
+    weights = [
+        ("mtp.0.attn.wkv.weight", torch.arange(4)),
+        ("mtp.1.attn.wq_a.weight", torch.arange(3)),
+        ("mtp.2.attn.wkv.scale", torch.tensor(0.5)),
+        ("mtp.3.attn.wkv.weight", torch.arange(2)),
+    ]
+
+    duplicated = list(dsv4_dspark._duplicate_context_wkv_weights(weights, 3))
+
+    assert [name for name, _ in duplicated] == [
+        "mtp.0.attn.wkv.weight",
+        "context_wkv_proj.weight",
+        "mtp.1.attn.wq_a.weight",
+        "mtp.2.attn.wkv.scale",
+        "context_wkv_proj.scale",
+        "mtp.3.attn.wkv.weight",
+    ]
+    assert duplicated[1][1].shard_id == 0
+    assert duplicated[4][1].shard_id == 2
+    assert duplicated[0][1].data_ptr() == duplicated[1][1].data_ptr()
+    assert duplicated[3][1].data_ptr() == duplicated[4][1].data_ptr()
+
+
+def test_dsv4_context_kv_uses_one_stacked_wkv_projection(monkeypatch):
+    calls = []
+    stacked_output = torch.arange(24, dtype=torch.float32).view(2, 12)
+
+    class StackedProjection:
+        def __init__(self):
+            self.calls = 0
+
+        def __call__(self, main_x):
+            self.calls += 1
+            assert main_x.shape == (2, 5)
+            return stacked_output
+
+    projection = StackedProjection()
+    layers = [
+        SimpleNamespace(attn=SimpleNamespace(kv_norm=lambda kv, offset=i: kv + offset))
+        for i in range(3)
+    ]
+    model = SimpleNamespace(
+        config=SimpleNamespace(head_dim=4),
+        context_wkv_proj=projection,
+        layers=layers,
+        num_dspark_layers=3,
+    )
+    slot_mappings = [torch.tensor([0, 1]), None, torch.tensor([4, 5])]
+    monkeypatch.setattr(
+        dsv4_dspark,
+        "_insert_context_kv",
+        lambda attn, kv, positions, slots: calls.append(
+            (attn, kv.clone(), positions, slots)
+        ),
+    )
+
+    dsv4_dspark.DSparkDeepseekV4Model.precompute_and_store_context_kv(
+        model,
+        torch.zeros(2, 5),
+        torch.tensor([7, 8]),
+        slot_mappings,
+    )
+
+    assert projection.calls == 1
+    assert len(calls) == 2
+    assert torch.equal(calls[0][1], stacked_output.view(2, 3, 4)[:, 0])
+    assert torch.equal(calls[1][1], stacked_output.view(2, 3, 4)[:, 2] + 2)
+    assert calls[0][3] is slot_mappings[0]
+    assert calls[1][3] is slot_mappings[2]
diff --git a/tests/models/test_hyv4_rocm.py b/tests/models/test_hyv4_rocm.py
new file mode 100644
index 000000000000..ea740c5cb099
--- /dev/null
+++ b/tests/models/test_hyv4_rocm.py
@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import pytest
+
+from tests.utils import multi_gpu_marks
+from vllm import SamplingParams
+from vllm.config import CompilationConfig
+from vllm.config.compilation import CompilationMode, CUDAGraphMode
+from vllm.platforms import current_platform
+from vllm.v1.metrics.reader import Counter
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-specific tests"
+)
+
+
+def _small_hyv4_config(config):
+    # Apply to both target and MTP draft while retaining MLA/indexer dimensions.
+    config.update(
+        {
+            "hidden_size": 512,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 64,
+            "q_lora_rank": 128,
+            "intermediate_size": 512,
+            "moe_intermediate_size": 256,
+            "n_routed_experts": 4,
+            "num_experts_per_tok": 2,
+            "n_shared_experts": 1,
+            "enable_ihc": True,
+            "index_n_heads": 32,
+            "index_topk": 2048,
+            "layer_types": ["deepseek_sparse_attention"] * 2,
+            "indexer_types": ["full", "shared"],
+            "num_nextn_predict_layers": 1,
+            "vocab_size": 256,
+            "bos_token_id": 1,
+            "eos_token_id": 2,
+            "pad_token_id": 0,
+        }
+    )
+    return config
+
+
+def _num_captured_graphs(_worker):
+    from vllm.compilation.counter import compilation_counter
+
+    return compilation_counter.num_cudagraph_captured
+
+
+def _piecewise_compile_state(worker):
+    from vllm.compilation.counter import compilation_counter
+
+    # `cudagraph_mode` is returned by name: the enum does not survive the
+    # collective_rpc round trip, it arrives as its bare value.
+    return (
+        compilation_counter.num_models_seen,
+        compilation_counter.num_cudagraph_captured,
+        worker.vllm_config.compilation_config.cudagraph_mode.name,
+    )
+
+
+@pytest.mark.parametrize(
+    "tp,enforce_eager,mtp,dtype,kv_cache_dtype",
+    [
+        pytest.param(1, True, False, "bfloat16", "auto", id="eager-tp1"),
+        pytest.param(1, True, False, "float16", "auto", id="eager-fp16"),
+        pytest.param(1, True, False, "bfloat16", "fp8", id="eager-fp8-kv"),
+        pytest.param(
+            2,
+            False,
+            False,
+            "bfloat16",
+            "auto",
+            id="graphs-tp2",
+            marks=multi_gpu_marks(num_gpus=2),
+        ),
+        pytest.param(1, True, True, "bfloat16", "auto", id="eager-mtp"),
+    ],
+)
+def test_hyv4_dummy_generation(
+    vllm_runner,
+    enable_pickle,
+    tp: int,
+    enforce_eager: bool,
+    mtp: bool,
+    dtype: str,
+    kv_cache_dtype: str,
+) -> None:
+    """Exercise prefill, decode, and batch consistency with dummy HY weights."""
+    with vllm_runner(
+        "tencent/Hy4-preview",
+        trust_remote_code=True,
+        skip_tokenizer_init=True,
+        load_format="dummy",
+        disable_log_stats=not mtp,
+        dtype=dtype,
+        kv_cache_dtype=kv_cache_dtype,
+        hf_overrides=_small_hyv4_config,
+        max_model_len=256,
+        max_num_batched_tokens=128,
+        max_num_seqs=8,
+        enable_chunked_prefill=True,
+        enable_prefix_caching=False,
+        kv_cache_memory_bytes=128 * 1024 * 1024,
+        enforce_eager=enforce_eager,
+        tensor_parallel_size=tp,
+        speculative_config={"method": "mtp", "num_speculative_tokens": 2}
+        if mtp
+        else None,
+        seed=17,
+        # Exercise supported decode graphs independently of default graph policy.
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.NONE,
+            cudagraph_mode=CUDAGraphMode.NONE
+            if enforce_eager
+            else CUDAGraphMode.FULL_DECODE_ONLY,
+            cudagraph_capture_sizes=[] if enforce_eager else [1, 2, 4],
+        ),
+    ) as model:
+        if not enforce_eager:
+            assert model.collective_rpc(_num_captured_graphs) == [3] * tp
+        prompts = [
+            {"prompt_token_ids": [11, 12, 13, 14]},
+            {"prompt_token_ids": list(range(20, 84))},
+        ]
+        params = SamplingParams(
+            temperature=0, max_tokens=4, ignore_eos=True, logprobs=1
+        )
+        outputs = model.llm.generate(prompts, params)
+        assert len(outputs) == len(prompts)
+        for prompt, output in zip(prompts, outputs):
+            completion = output.outputs[0]
+            assert len(completion.token_ids) == 4
+            assert completion.cumulative_logprob is not None
+            assert math.isfinite(completion.cumulative_logprob)
+            single = model.llm.generate([prompt], params)[0].outputs[0]
+            assert single.token_ids == completion.token_ids
+
+        if mtp:
+            assert any(
+                isinstance(metric, Counter)
+                and metric.name == "vllm:spec_decode_num_drafts"
+                and metric.value > 0
+                for metric in model.llm.get_metrics()
+            ), "MTP drafter did not run"
+
+
+def test_hyv4_backbone_reaches_the_compiler(vllm_runner, enable_pickle) -> None:
+    """Guard against the backbone losing `@support_torch_compile`.
+
+    An undecorated backbone is not a hard failure: vLLM downgrades
+    `cudagraph_mode` to NONE and serves the whole model eager, which on TP8
+    gfx942 cost ~4x decode latency.
+    """
+    with vllm_runner(
+        "tencent/Hy4-preview",
+        trust_remote_code=True,
+        skip_tokenizer_init=True,
+        load_format="dummy",
+        dtype="bfloat16",
+        hf_overrides=_small_hyv4_config,
+        max_model_len=256,
+        max_num_batched_tokens=128,
+        max_num_seqs=8,
+        enable_chunked_prefill=True,
+        enable_prefix_caching=False,
+        kv_cache_memory_bytes=128 * 1024 * 1024,
+        enforce_eager=False,
+        seed=17,
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            cudagraph_mode=CUDAGraphMode.PIECEWISE,
+            cudagraph_capture_sizes=[1, 2, 4],
+        ),
+    ) as model:
+        for models_seen, captured, mode in model.collective_rpc(
+            _piecewise_compile_state
+        ):
+            assert models_seen > 0, "backbone was never handed to the compiler"
+            assert mode == CUDAGraphMode.PIECEWISE.name
+            assert captured > 0
diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py
index 3b68f424788f..073f616af875 100644
--- a/tests/models/test_initialization.py
+++ b/tests/models/test_initialization.py
@@ -1,12 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import gc
+from collections.abc import Generator
 from functools import partial
 from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm import LLM
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
 from vllm.v1.core.kv_cache_utils import (
@@ -15,7 +20,12 @@
 )
 from vllm.v1.engine.core import EngineCore as V1EngineCore
 
-from ..utils import create_new_process_for_each_test, requires_spawn_multiprocessing
+from ..utils import (
+    create_new_process_for_each_test,
+    requires_spawn_multiprocessing,
+    wait_for_gpu_memory_to_clear,
+    wait_for_memory_to_settle,
+)
 from .registry import (
     _TRANSFORMERS_BACKEND_MODELS,
     AUTO_EXAMPLE_MODELS,
@@ -24,6 +34,8 @@
 )
 from .utils import dummy_hf_overrides
 
+logger = init_logger(__name__)
+
 # This minimal list of model architectures is smaller than the total list of
 # supported models. The intention is that in the "typical" regression testing
 # scenario, we only test initializing these models. This subset was chosen
@@ -53,6 +65,37 @@
 )
 
 
+@pytest.fixture(autouse=True)
+def cleanup_gpu_memory_between_tests() -> Generator[None, None, None]:
+    """Reclaim GPU memory left over by each test's forked subprocess.
+
+    These tests are known to leave behind GPU memory (see the comment at
+    ``gpu_memory_utilization=0.80`` below). Each parametrized case builds an
+    LLM in a forked subprocess whose VRAM is only reclaimed asynchronously
+    when the subprocess exits, so the leftovers can OOM a later test in this
+    module (observed with the heavier multimodal archs, e.g.
+    InternVLChatModel). Mirrors ``clean_gpu_memory_between_tests`` in the
+    top-level conftest, but unconditionally for this module.
+    """
+    yield
+    gc.collect()
+    if not torch.cuda.is_available():
+        return
+    torch.accelerator.empty_cache()
+    if current_platform.is_rocm() or current_platform.is_xpu():
+        wait_for_memory_to_settle()
+        return
+    try:
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(torch.accelerator.device_count())),
+            threshold_ratio=0.1,
+        )
+    except Exception as e:
+        # NVML is not usable on MIG slices (raises NVMLError_NoPermission);
+        # cleanup is best-effort, so never fail a test over it.
+        logger.info("Failed to clean GPU memory: %s", e)
+
+
 @create_new_process_for_each_test()
 def can_initialize(
     model_arch: str, monkeypatch: pytest.MonkeyPatch, EXAMPLE_MODELS: HfExampleModels
@@ -64,7 +107,6 @@ def can_initialize(
     The spawn process causes the _initialize_kv_caches_v1 function below to
     become ineffective.
     """
-
     model_info = EXAMPLE_MODELS.get_hf_info(model_arch)
     model_info.check_available_online(on_fail="skip")
     model_info.check_transformers_version(
@@ -135,8 +177,8 @@ def _initialize_kv_caches_v1(self, vllm_config):
     if model_arch == "DeepseekV4ForConditionalGeneration":
         from vllm.platforms import current_platform
 
-        if not current_platform.is_cuda():
-            pytest.skip("Deepseek V4 is only supported on CUDA")
+        if not (current_platform.is_cuda() or current_platform.is_rocm()):
+            pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm")
 
     with (
         patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1),
@@ -199,13 +241,13 @@ def _initialize_kv_caches_v1(self, vllm_config):
 
 @pytest.mark.parametrize("model_arch", MINIMAL_MODEL_ARCH_LIST)
 def test_can_initialize_small_subset(model_arch: str, monkeypatch: pytest.MonkeyPatch):
-    """Test initializing small subset of supported models"""
+    """Test initializing small subset of supported models."""
     can_initialize(model_arch, monkeypatch, HF_EXAMPLE_MODELS)
 
 
 @pytest.mark.parametrize("model_arch", OTHER_MODEL_ARCH_LIST)
 def test_can_initialize_large_subset(model_arch: str, monkeypatch: pytest.MonkeyPatch):
-    """Test initializing large subset of supported models
+    """Test initializing large subset of supported models.
 
     This test covers the complement of the tests covered in the "small subset"
     test.
diff --git a/tests/models/test_qwen3_5_mtp_config.py b/tests/models/test_qwen3_5_mtp_config.py
index 6cf54dbc95a8..885995b906f4 100644
--- a/tests/models/test_qwen3_5_mtp_config.py
+++ b/tests/models/test_qwen3_5_mtp_config.py
@@ -5,7 +5,7 @@
 from typing import Any
 
 import pytest
-from transformers import AutoConfig, PretrainedConfig
+from transformers import AutoConfig, PreTrainedConfig
 
 from vllm.config.speculative import SpeculativeConfig
 
@@ -15,19 +15,19 @@
 }
 
 
-def _mtp_config(model_type: str) -> PretrainedConfig:
+def _mtp_config(model_type: str) -> PreTrainedConfig:
     """Create a top-level MTP configuration with mtp_num_hidden_layers."""
     kwargs: dict[str, Any] = {
         "model_type": model_type,
         "architectures": ["SomeArch"],
         "mtp_num_hidden_layers": 1,
     }
-    return PretrainedConfig(**kwargs)
+    return PreTrainedConfig(**kwargs)
 
 
 def _multimodal_wrapper_mtp_config(
     model_type: str, mtp_layers: int = 1
-) -> PretrainedConfig:
+) -> PreTrainedConfig:
     """Download a multimodal wrapper config via AutoConfig.from_pretrained
     and configure mtp_num_hidden_layers in text_config.
 
@@ -36,7 +36,7 @@ def _multimodal_wrapper_mtp_config(
     - MoE   (qwen3_5_moe): Qwen/Qwen3.6-35B-A3B
     """
     repo = _CHECKPOINTS[model_type]
-    config: PretrainedConfig = AutoConfig.from_pretrained(repo)
+    config: PreTrainedConfig = AutoConfig.from_pretrained(repo)
     text_config = config.get_text_config()
     text_config.mtp_num_hidden_layers = mtp_layers
     return config
diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py
index 30eb80a3f88f..b7b13397866d 100644
--- a/tests/models/test_registry.py
+++ b/tests/models/test_registry.py
@@ -64,10 +64,11 @@ def test_registry_imports(model_arch):
     ):
         pytest.skip("HY V4 is only supported on CUDA and ROCm")
 
-    if model_arch == "DeepseekV4ForConditionalGeneration" and not (
-        current_platform.is_cuda()
+    if (
+        model_arch == "DeepseekV4ForConditionalGeneration"
+        and not current_platform.is_cuda_alike()
     ):
-        pytest.skip("Deepseek V4 is only supported on CUDA")
+        pytest.skip("Deepseek V4 vision is only supported on CUDA and ROCm")
 
     # Ensure all model classes can be imported successfully
     model_cls = ModelRegistry._try_load_model_cls(model_arch)
@@ -93,6 +94,7 @@ def test_registry_imports(model_arch):
     [
         ("LlamaForCausalLM", False, False, "bi-encoder"),
         ("LlavaForConditionalGeneration", True, True, "bi-encoder"),
+        ("DeepseekV41ForCausalLM", True, False, "bi-encoder"),
         ("BertForSequenceClassification", False, False, "cross-encoder"),
         ("RobertaForSequenceClassification", False, False, "cross-encoder"),
         ("XLMRobertaForSequenceClassification", False, False, "cross-encoder"),
@@ -128,8 +130,6 @@ def test_registry_model_property(model_arch, is_mm, init_cuda, score_type):
         # in V1.
         # ("MLPSpeculatorPreTrainedModel", False, False),
         ("DeepseekV2ForCausalLM", True, False),
-        ("KimiLinearForCausalLM", True, False),
-        ("KimiK3ForConditionalGeneration", True, False),
         ("Qwen2VLForConditionalGeneration", True, True),
     ],
 )
diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py
index 02903504c684..0dac2dbf2a6d 100644
--- a/tests/models/test_utils.py
+++ b/tests/models/test_utils.py
@@ -217,3 +217,25 @@ def test_get_rename_mapper_keeps_only_renames():
     for name in ("drop_regex.w", "drop_substr.w", "drop_prefix.w", "w.drop_suffix"):
         assert mapper._map_name(name) is None
         assert renames._map_name(name) == name
+
+
+def test_weights_mapper_stacks_one_weight_into_several_shards():
+    """One checkpoint tensor can feed more than one shard of a stacked
+    parameter, and a dropped name stays dropped rather than reappearing as a
+    shard of the stacked one."""
+    mapper = WeightsMapper(
+        orig_to_new_substr={"layers.1.k_proj.": None},
+        orig_to_new_stacked={".k_proj.": (".qkv_proj.", ["k", "v"])},
+    )
+    weight = torch.ones(2)
+    weights = [(f"layers.{i}.k_proj.weight", weight) for i in (0, 1)]
+
+    mapped = list(mapper.apply(weights))
+
+    assert [(name, w.shard_id) for name, w in mapped] == [
+        ("layers.0.qkv_proj.weight", "k"),
+        ("layers.0.qkv_proj.weight", "v"),
+    ]
+    # Each shard needs its own tensor object, but they alias one allocation.
+    assert mapped[0][1] is not mapped[1][1]
+    assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr() == weight.data_ptr()
diff --git a/tests/models/test_vision.py b/tests/models/test_vision.py
index 24260e2ef804..94edce45bf53 100644
--- a/tests/models/test_vision.py
+++ b/tests/models/test_vision.py
@@ -41,9 +41,7 @@
 def test_resolve_visual_encoder_outputs(
     select_layers, num_layers_loaded, max_possible_layers, expected_features
 ):
-    """
-    Test that offsets are correctly handled for vision feature layers.
-    """
+    """Test that offsets are correctly handled for vision feature layers."""
     encoder_outputs = [torch.tensor([idx]) for idx in range(num_layers_loaded + 1)]
     output_tensor = resolve_visual_encoder_outputs(
         encoder_outputs=encoder_outputs,
@@ -94,11 +92,9 @@ def test_run_dp_sharded_vision_model(batch_size: int):
 def run_dp_sharded_vision_model_vs_direct(
     local_rank: int, world_size: int, batch_size: int, master_port: int
 ):
-    """
-    Test that run_dp_sharded_vision_model produces the same results as
+    """Test that run_dp_sharded_vision_model produces the same results as
     calling the model directly.
     """
-
     # Set random seed for reproducibility
     set_random_seed(0)
 
@@ -282,8 +278,7 @@ def test_run_dp_sharded_mrope_vision_model(batch_size: int):
 def run_dp_sharded_mrope_vision_model_vs_direct(
     local_rank: int, world_size: int, batch_size: int, master_port: int
 ):
-    """
-    Test that run_dp_sharded_mrope_vision_model produces the same results as
+    """Test that run_dp_sharded_mrope_vision_model produces the same results as
     calling the model directly.
     """
     # Set random seed for reproducibility
diff --git a/tests/models/transformers/fusers/test_mla.py b/tests/models/transformers/fusers/test_mla.py
index 9f51ab56621b..625fc4dc0815 100644
--- a/tests/models/transformers/fusers/test_mla.py
+++ b/tests/models/transformers/fusers/test_mla.py
@@ -53,7 +53,6 @@ def _match(q_lora_rank: int | None) -> MLAFuser | None:
 def test_discovers_modules_without_q_lora():
     fuser = _match(q_lora_rank=None)
     assert isinstance(fuser, MLAFuser)
-    assert not fuser.has_q_lora
     assert fuser.q_proj_name == "q_proj"
     assert fuser.kv_a_proj_name == "kv_a_proj_with_mqa"
     assert fuser.kv_a_layernorm_name == "kv_a_layernorm"
@@ -68,7 +67,6 @@ def test_discovers_modules_without_q_lora():
 def test_discovers_modules_with_q_lora():
     fuser = _match(q_lora_rank=64)
     assert isinstance(fuser, MLAFuser)
-    assert fuser.has_q_lora
     assert fuser.q_a_proj_name == "q_a_proj"
     assert fuser.q_a_layernorm_name == "q_a_layernorm"
     assert fuser.q_b_proj_name == "q_b_proj"
@@ -107,7 +105,7 @@ def test_update_forward_rewrites_real_attention(q_lora_rank):
     assert isinstance(fuser, MLAFuser)
     fuser.update_forward(_attention(q_lora_rank))
     names = set(fuser.fused_forward.__code__.co_names)
-    if fuser.has_q_lora:
+    if fuser.q_a_proj_name is not None:
         assert _FUSED_QKV_A_PROJ in names
         assert not {"q_a_proj", "kv_a_proj_with_mqa"} & names
 
@@ -158,7 +156,7 @@ def test_discovers_modules_under_arbitrary_names():
         module = RenamedMLA()
         fuser = MLAFuser.match(trace(module), module)
     assert isinstance(fuser, MLAFuser)
-    assert not fuser.has_q_lora
+    assert fuser.q_a_proj_name is None
     assert fuser.q_proj_name == "alpha"
     assert fuser.kv_a_proj_name == "beta"
     assert fuser.kv_a_layernorm_name == "gamma"
diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py
index fd0d980deffc..fa5e14eddd21 100644
--- a/tests/models/transformers/test_backend.py
+++ b/tests/models/transformers/test_backend.py
@@ -13,7 +13,7 @@
 import pytest
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModel, PretrainedConfig
+from transformers import AutoConfig, AutoModel, PreTrainedConfig
 
 from vllm.config import ModelConfig, VllmConfig
 from vllm.model_executor.models.interfaces import SupportsMultiModal
@@ -524,7 +524,7 @@ def replace_vocab_embeddings(model, **config_kwargs):
 
     stub = nn.Module()
     stub.model = model
-    stub.config = PretrainedConfig(
+    stub.config = PreTrainedConfig(
         vocab_size=VOCAB_SIZE, num_positions=NUM_POSITIONS, **config_kwargs
     )
     embeddings = Base._vocab_embeddings(stub)
diff --git a/tests/models/utils.py b/tests/models/utils.py
index aff9fb0f5a66..003ef1a078ae 100644
--- a/tests/models/utils.py
+++ b/tests/models/utils.py
@@ -8,7 +8,7 @@
 
 import torch
 import torch.nn.functional as F
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.model import AttnTypeStr, ModelConfig, ModelDType, RunnerOption
 from vllm.config.pooler import SequencePoolingType, TokenPoolingType
@@ -29,8 +29,7 @@ def check_outputs_equal(
     name_0: str,
     name_1: str,
 ):
-    """
-    Compare the two sequences generated by different models,
+    """Compare the two sequences generated by different models,
     which should be equal.
     """
     assert len(outputs_0_lst) == len(outputs_1_lst)
@@ -131,6 +130,7 @@ def check_logprobs_close(
       warn_on_mismatch: Issue a warning if there is token-wise or text-wise
                         mismatch between the two sequences
       always_check_logprobs: If true, check logprobs even when tokens match
+
     """
     assert len(outputs_0_lst) == len(outputs_1_lst)
 
@@ -290,6 +290,7 @@ def build_model_context(
 
     Returns:
         InputProcessingContext for the model being considered.
+
     """
     model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
     model_info.check_available_online(on_fail="skip")
@@ -395,6 +396,8 @@ class ModelInfo:
 @dataclass
 class EmbedModelInfo(ModelInfo):
     mteb_score: float | None = None
+    # Per-model MTEB score tolerance; falls back to MTEB_EMBED_TOL.
+    mteb_tol: float | None = None
     is_matryoshka: bool = False
     matryoshka_dimensions: list[int] | None = None
 
@@ -402,6 +405,8 @@ class EmbedModelInfo(ModelInfo):
 @dataclass
 class RerankModelInfo(ModelInfo):
     mteb_score: float | None = None
+    # Per-model MTEB score tolerance; falls back to MTEB_RERANK_TOL.
+    mteb_tol: float | None = None
     chat_template_name: str | None = None
 
 
@@ -441,14 +446,13 @@ def get_vllm_extra_kwargs(model_info: ModelInfo, vllm_extra_kwargs):
 
 
 def dummy_hf_overrides(
-    hf_config: PretrainedConfig,
+    hf_config: PreTrainedConfig,
     *,
     model_arch: str = "",
     exist_overrides: dict[str, Any] | None = None,
     use_original_num_layers: bool = False,
-) -> PretrainedConfig:
-    """
-    Dummy HF overrides function used to create dummy model
+) -> PreTrainedConfig:
+    """Dummy HF overrides function used to create dummy model
     with only minimum nums of layer.
     """
     # Copy because this helper is called more than once
diff --git a/tests/multimodal/media/test_audio.py b/tests/multimodal/media/test_audio.py
index 3473a911eab5..0404535e1b25 100644
--- a/tests/multimodal/media/test_audio.py
+++ b/tests/multimodal/media/test_audio.py
@@ -218,7 +218,7 @@ def test_load_audio_auto_falls_back_without_torchcodec(dummy_audio_bytes):
 
 
 def test_load_audio_auto_falls_back_without_ffmpeg(dummy_audio_bytes):
-    """torchcodec installed but system ffmpeg missing (`AudioDecoder is None`)
+    """Torchcodec installed but system ffmpeg missing (`AudioDecoder is None`)
     must surface as ImportError so `auto` falls back to soundfile → PyAV."""
     ref_audio, ref_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None)
     with patch.object(audio_module, "AudioDecoder", None):
@@ -232,7 +232,7 @@ def test_load_audio_auto_falls_back_without_ffmpeg(dummy_audio_bytes):
 def test_load_audio_auto_falls_back_when_libtorchcodec_unloadable(
     dummy_audio_bytes,
 ):
-    """torchcodec loads its ffmpeg-backed core lazily at decoder construction;
+    """Torchcodec loads its ffmpeg-backed core lazily at decoder construction;
     a "Could not load libtorchcodec" RuntimeError there (no system ffmpeg)
     must also surface as ImportError so `auto` falls back."""
 
diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py
index b222d57d3c24..5949c66b562a 100644
--- a/tests/multimodal/media/test_connector.py
+++ b/tests/multimodal/media/test_connector.py
@@ -299,6 +299,7 @@ async def test_fetch_video_http(video_url: str, num_frames: int):
     assert metadata_sync == metadata_async
 
 
+@pytest.mark.flaky(reruns=3, reruns_delay=5)
 @pytest.mark.asyncio
 @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
 @pytest.mark.parametrize("max_duration", [1, 60, 1800])
@@ -320,8 +321,11 @@ async def test_fetch_video_http_with_dynamic_loader(
             }
         )
 
-        video_sync, metadata_sync = connector.fetch_video(video_url)
-        video_async, metadata_async = await connector.fetch_video_async(video_url)
+        try:
+            video_sync, metadata_sync = connector.fetch_video(video_url)
+            video_async, metadata_async = await connector.fetch_video_async(video_url)
+        except (TimeoutError, asyncio.TimeoutError) as e:
+            pytest.skip(f"Timeout fetching video (CI network flakiness): {e}")
 
         assert np.array_equal(video_sync, video_async)
         assert metadata_sync == metadata_async
@@ -379,6 +383,7 @@ def test_placeholder_range_extract_embeds_range(offset, is_embed, expected):
     assert pr.extract_embeds_range() == expected
 
 
+@pytest.mark.flaky(reruns=3, reruns_delay=5)
 @pytest.mark.asyncio
 @pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
 @pytest.mark.parametrize("num_frames", [-1, 32, 1800])
@@ -394,8 +399,12 @@ async def test_allowed_media_domains(video_url: str, num_frames: int):
         ],
     )
 
-    video_sync, metadata_sync = connector.fetch_video(video_url)
-    video_async, metadata_async = await connector.fetch_video_async(video_url)
+    try:
+        video_sync, metadata_sync = connector.fetch_video(video_url)
+        video_async, metadata_async = await connector.fetch_video_async(video_url)
+    except (TimeoutError, asyncio.TimeoutError) as e:
+        pytest.skip(f"Timeout fetching video (CI network flakiness): {e}")
+
     assert np.array_equal(video_sync, video_async)
     assert metadata_sync == metadata_async
 
diff --git a/tests/multimodal/media/test_image.py b/tests/multimodal/media/test_image.py
index 5bbf54b7872e..6b7e4ad407b7 100644
--- a/tests/multimodal/media/test_image.py
+++ b/tests/multimodal/media/test_image.py
@@ -105,7 +105,6 @@ def test_image_media_io_no_mode_conversion(tmp_path):
 
 def test_image_media_io_rgba_background_color_validation():
     """Test that invalid rgba_background_color values are properly rejected."""
-
     # Test invalid types
     with pytest.raises(
         ValueError, match="rgba_background_color must be a list or tuple"
diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py
index a3ea0890a20b..90310d6e6d48 100644
--- a/tests/multimodal/media/test_video.py
+++ b/tests/multimodal/media/test_video.py
@@ -77,8 +77,7 @@ def test_video_media_io_kwargs(monkeypatch: pytest.MonkeyPatch):
 @pytest.mark.parametrize("is_color", [True, False])
 @pytest.mark.parametrize("fourcc, ext", [("mp4v", "mp4"), ("XVID", "avi")])
 def test_opencv_video_io_colorspace(tmp_path, is_color: bool, fourcc: str, ext: str):
-    """
-    Test all functions that use OpenCV for video I/O return RGB format.
+    """Test all functions that use OpenCV for video I/O return RGB format.
     Both RGB and grayscale videos are tested.
     """
     image_path = get_vllm_public_assets(
@@ -168,8 +167,7 @@ def load_bytes(
 
 
 def test_video_media_io_backend_kwarg_override(monkeypatch: pytest.MonkeyPatch):
-    """
-    Test that video_backend kwarg can override the VLLM_VIDEO_LOADER_BACKEND
+    """Test that video_backend kwarg can override the VLLM_VIDEO_LOADER_BACKEND
     environment variable.
 
     This allows users to dynamically select a different video backend
@@ -201,8 +199,7 @@ def test_video_media_io_backend_kwarg_override(monkeypatch: pytest.MonkeyPatch):
 def test_video_media_io_backend_kwarg_not_passed_to_loader(
     monkeypatch: pytest.MonkeyPatch,
 ):
-    """
-    Test that video_backend kwarg is consumed by VideoMediaIO and NOT passed
+    """Test that video_backend kwarg is consumed by VideoMediaIO and NOT passed
     through to the underlying video loader's load_bytes method.
 
     This ensures the kwarg is properly popped from kwargs before forwarding.
@@ -245,8 +242,7 @@ def load_bytes(
 
 
 def test_video_media_io_backend_env_var_fallback(monkeypatch: pytest.MonkeyPatch):
-    """
-    Test that when video_backend kwarg is None or not provided,
+    """Test that when video_backend kwarg is None or not provided,
     VideoMediaIO falls back to VLLM_VIDEO_LOADER_BACKEND env var.
     """
     with monkeypatch.context() as m:
@@ -285,7 +281,6 @@ def test_load_base64_jpeg_returns_metadata():
     metadata, which broke downstream consumers that rely on fields like
     total_num_frames and fps. See PR #37301.
     """
-
     num_test_frames = 3
 
     b64_frames = _make_jpeg_b64_frames(num_test_frames)
@@ -532,6 +527,50 @@ def test_strips_pool_size_from_runtime(self):
     def test_unknown_backend_not_treated_as_gpu(self):
         assert not VIDEO_LOADER_REGISTRY.backend_requires_gpu("totally_unknown")
 
+    def test_strips_request_level_device(self):
+        """The decode device is a startup-only knob: a request must not move
+        decoding onto the GPU when the startup config did not opt in, nor off
+        it when it did."""
+        result = VideoMediaIO.merge_kwargs(
+            default_kwargs={"backend": "torchcodec"},
+            runtime_kwargs={"device": "cuda"},
+        )
+        assert "device" not in result
+
+        result = VideoMediaIO.merge_kwargs(
+            default_kwargs={"backend": "torchcodec", "device": "cuda"},
+            runtime_kwargs={"device": "cpu", "num_frames": 8},
+        )
+        assert result["device"] == "cuda"
+        assert result["num_frames"] == 8
+
+
+@pytest.mark.parametrize(
+    "default_kwargs",
+    [
+        {"backend": "torchcodec", "device": "cuda", "seek_mode": "approximate"},
+        {"backend": "pynvvideocodec", "hw_decoders": 2},
+    ],
+)
+def test_switching_backend_drops_stale_codec_options(default_kwargs):
+    """Codec-specific options from the static config must not leak into a
+    different codec backend selected per-request, where they would fail the
+    new backend's option validation."""
+    result = VideoMediaIO.merge_kwargs(
+        default_kwargs={**default_kwargs, "num_frames": 8},
+        runtime_kwargs={"backend": "opencv"},
+    )
+    assert result == {"backend": "opencv", "num_frames": 8}
+
+
+def test_same_backend_keeps_codec_options():
+    result = VideoMediaIO.merge_kwargs(
+        default_kwargs={"backend": "torchcodec", "device": "cuda"},
+        runtime_kwargs={"backend": "torchcodec", "num_frames": 8},
+    )
+    assert result["device"] == "cuda"
+    assert result["num_frames"] == 8
+
 
 @pytest.mark.parametrize("layout", ["nhwc", "nchw"])
 def test_pynvvc_frames_normalized_to_nhwc(layout: str):
diff --git a/tests/multimodal/test_audio.py b/tests/multimodal/test_audio.py
index 21b39710d4f4..c816e979abce 100644
--- a/tests/multimodal/test_audio.py
+++ b/tests/multimodal/test_audio.py
@@ -695,7 +695,6 @@ class TestAudioChunking:
 
     def test_split_audio_short_clip(self):
         """Audio shorter than max_clip_duration_s should not be split."""
-
         # 10 seconds of audio at 16kHz
         audio = np.linspace(-1.0, 1.0, 160000, dtype=np.float32)
 
@@ -712,7 +711,6 @@ def test_split_audio_short_clip(self):
 
     def test_split_audio_exact_length(self):
         """Audio exactly at max_clip_duration_s should not be split."""
-
         # Exactly 30 seconds at 16kHz
         audio = np.linspace(-1.0, 1.0, 480000, dtype=np.float32)
 
@@ -729,7 +727,6 @@ def test_split_audio_exact_length(self):
 
     def test_split_audio_long_clip(self):
         """Long audio should be split into multiple chunks."""
-
         # 65 seconds of audio at 16kHz
         audio = np.linspace(-1.0, 1.0, 1040000, dtype=np.float32)
 
@@ -749,7 +746,6 @@ def test_split_audio_long_clip(self):
 
     def test_split_audio_chunks_have_correct_length(self):
         """Each chunk (except last) should be approximately max_clip_duration_s."""
-
         # 65 seconds of audio at 16kHz
         audio = np.linspace(-1.0, 1.0, 1040000, dtype=np.float32)
 
@@ -841,7 +837,6 @@ def test_split_audio_rejects_multi_channel(self):
 
     def test_split_audio_preserves_boundaries(self):
         """Verify first and last samples are preserved when chunking."""
-
         audio = np.arange(1120000, dtype=np.float32)  # 70s at 16kHz
 
         chunks = split_audio(
@@ -893,7 +888,6 @@ def test_split_audio_nan_input_terminates(self):
 
     def test_split_audio_with_different_sample_rates(self):
         """Test chunking works with different sample rates."""
-
         # 40 seconds at 8kHz
         audio_8k = np.linspace(-1.0, 1.0, 320000, dtype=np.float32)
 
diff --git a/tests/multimodal/test_cache.py b/tests/multimodal/test_cache.py
index f77dacd341b8..ffd7d61cf25b 100644
--- a/tests/multimodal/test_cache.py
+++ b/tests/multimodal/test_cache.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import multiprocessing as mp
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import numpy as np
 import pytest
@@ -9,20 +10,23 @@
 
 from vllm.config import ModelConfig, ParallelConfig, VllmConfig
 from vllm.config.multimodal import MultiModalConfig
-from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.cache import (
     BaseMultiModalProcessorCache,
     BaseMultiModalReceiverCache,
-    MultiModalCache,
+    LruKeyReplicatedReceiverCache,
+    LruKeyReplicatedSenderCache,
     MultiModalCacheMissError,
-    MultiModalProcessorCacheInItem,
-    MultiModalProcessorCacheItem,
-    MultiModalProcessorCacheItemMetadata,
     MultiModalProcessorOnlyCache,
-    MultiModalProcessorSenderCache,
-    MultiModalReceiverCache,
     ShmObjectStoreReceiverCache,
     ShmObjectStoreSenderCache,
+    engine_receiver_cache_from_config,
+    processor_cache_from_config,
+)
+from vllm.multimodal.cache.base import (
+    MultiModalCache,
+    MultiModalProcessorCacheInItem,
+    MultiModalProcessorCacheItem,
+    MultiModalProcessorCacheItemMetadata,
 )
 from vllm.multimodal.hasher import MultiModalHasher
 from vllm.multimodal.inputs import (
@@ -34,8 +38,13 @@
     PlaceholderRange,
 )
 from vllm.multimodal.processing import PromptInsertion
+from vllm.renderers import renderer_from_config
+from vllm.utils.async_utils import make_async
 from vllm.utils.mem_constants import GiB_bytes, MiB_bytes
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.llm_engine import LLMEngine
 
+from ..models.utils import build_model_context
 from ..utils import create_new_process_for_each_test
 
 pytestmark = pytest.mark.cpu_test
@@ -131,10 +140,10 @@ def _compare_caches(
     n_iter: int = 100,
     seed: int = 0,
 ):
-    cache_0_p0 = MULTIMODAL_REGISTRY.processor_cache_from_config(config_0)
-    cache_0_p1 = MULTIMODAL_REGISTRY.engine_receiver_cache_from_config(config_0)
-    cache_1_p0 = MULTIMODAL_REGISTRY.processor_cache_from_config(config_1)
-    cache_1_p1 = MULTIMODAL_REGISTRY.engine_receiver_cache_from_config(config_1)
+    cache_0_p0 = processor_cache_from_config(config_0)
+    cache_0_p1 = engine_receiver_cache_from_config(config_0)
+    cache_1_p0 = processor_cache_from_config(config_1)
+    cache_1_p1 = engine_receiver_cache_from_config(config_1)
 
     cache_size_gb = max(
         config_0.model_config.multimodal_config.mm_processor_cache_gb,
@@ -269,8 +278,8 @@ def test_oversized_item_is_served_uncached():
     assert p0_only.get_and_update_item((small, []), "small")[0] is small
     assert p0_only.is_cached_item("small")
 
-    p0 = MultiModalProcessorSenderCache(model_config)  # type: ignore[arg-type]
-    p1 = MultiModalReceiverCache(model_config)  # type: ignore[arg-type]
+    p0 = LruKeyReplicatedSenderCache(model_config)  # type: ignore[arg-type]
+    p1 = LruKeyReplicatedReceiverCache(model_config)  # type: ignore[arg-type]
     assert p0.get_and_update_item((item, []), "big")[0] is item
     assert not p0.is_cached_item("big")
     assert p1.get_and_update_item(item, "big") is item
@@ -290,8 +299,8 @@ def test_mm_cache_miss_raises_and_recovers():
     model or network.
     """
     model_config = _StubModelConfig(mm_processor_cache_gb=1)
-    p0 = MultiModalProcessorSenderCache(model_config)  # type: ignore[arg-type]
-    p1 = MultiModalReceiverCache(model_config)  # type: ignore[arg-type]
+    p0 = LruKeyReplicatedSenderCache(model_config)  # type: ignore[arg-type]
+    p1 = LruKeyReplicatedReceiverCache(model_config)  # type: ignore[arg-type]
 
     mm_hash = "image_A"
     item = MultiModalKwargsItem.dummy(nbytes=64)
@@ -324,6 +333,51 @@ def test_mm_cache_miss_raises_and_recovers():
     assert p1.get_and_update_item(None, mm_hash) == item
 
 
+def test_receiver_cache_replaces_stale_item_when_payload_resent():
+    """A P0 miss can resend a different item under the same identity.
+
+    Independent LRU eviction leaves P1 holding the old tensor. Substituting it
+    for the new payload pairs the request's placeholders with the wrong item
+    and kills EngineCore. Prefer the fresh payload and keep it for later hits.
+    """
+    model_config = _StubModelConfig(mm_processor_cache_gb=1)
+    p1 = LruKeyReplicatedReceiverCache(model_config)  # type: ignore[arg-type]
+    small = MultiModalKwargsItem.dummy(nbytes=64)
+    large = MultiModalKwargsItem.dummy(nbytes=256)
+    mm_hash = "shared-id"
+
+    assert p1.get_and_update_item(small, mm_hash) is small
+    assert p1.get_and_update_item(large, mm_hash) is large
+    assert p1.get_and_update_item(None, mm_hash) is large
+
+
+def test_receiver_cache_features_keep_resent_payload():
+    """EngineCore updates features through get_and_update_features."""
+    model_config = _StubModelConfig(mm_processor_cache_gb=1)
+    p1 = LruKeyReplicatedReceiverCache(model_config)  # type: ignore[arg-type]
+    small = MultiModalKwargsItem.dummy(nbytes=64)
+    large = MultiModalKwargsItem.dummy(nbytes=256)
+    mm_hash = "shared-id"
+
+    def _feature(
+        data: MultiModalKwargsItem | None, length: int
+    ) -> MultiModalFeatureSpec:
+        return MultiModalFeatureSpec(
+            data=data,
+            modality="image",
+            identifier=mm_hash,
+            mm_position=PlaceholderRange(offset=0, length=length),
+            mm_hash=mm_hash,
+        )
+
+    seeded = p1.get_and_update_features([_feature(small, length=4)])
+    assert seeded[0].data is small
+
+    updated = p1.get_and_update_features([_feature(large, length=2048)])
+    assert updated[0].data is large
+    assert updated[0].mm_position.length == 2048
+
+
 def test_mm_cache_miss_batches_all_drifted_hashes():
     """All hashes drifted within one request must surface in a single error.
 
@@ -333,7 +387,7 @@ def test_mm_cache_miss_batches_all_drifted_hashes():
     non-drifted items in the same request are still ingested.
     """
     model_config = _StubModelConfig(mm_processor_cache_gb=1)
-    p1 = MultiModalReceiverCache(model_config)  # type: ignore[arg-type]
+    p1 = LruKeyReplicatedReceiverCache(model_config)  # type: ignore[arg-type]
     item = MultiModalKwargsItem.dummy(nbytes=64)
 
     def _feature(
@@ -531,8 +585,8 @@ def test_cache_eviction_lru_cache():
         model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf",
         mm_processor_cache_gb=6 / GiB_bytes,
     )
-    sender_cache = MultiModalProcessorSenderCache(model_config)
-    receiver_cache = MultiModalReceiverCache(model_config)
+    sender_cache = LruKeyReplicatedSenderCache(model_config)
+    receiver_cache = LruKeyReplicatedReceiverCache(model_config)
 
     _run_test_cache_eviction_lru(sender_cache, receiver_cache, base_item_size=1)
 
@@ -724,7 +778,7 @@ def test_processor_cache_shared_across_loras():
         model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf",
         mm_processor_cache_gb=1,
     )
-    receiver_cache = MultiModalReceiverCache(model_config)
+    receiver_cache = LruKeyReplicatedReceiverCache(model_config)
 
     base_mm_hash = "image_hash_abc123"
     lora_a_identifier = f"12345:{base_mm_hash}"
@@ -755,6 +809,74 @@ def test_processor_cache_shared_across_loras():
     assert feature_lora_b.data == item_data
 
 
+@pytest.mark.parametrize("use_async", [False, True])
+@pytest.mark.parametrize(
+    "release_error",
+    [
+        None,
+        "requires a completed pause first",
+        "requires all executor memory to be resident",
+    ],
+    ids=["released", "not-paused", "nonresident-memory"],
+)
+@pytest.mark.asyncio
+async def test_release_kv_cache_resends_mm_payload(use_async, release_error):
+    """Release must not leave a sender hit pointing at a cleared receiver."""
+    ctx = build_model_context(
+        "llava-hf/llava-v1.6-mistral-7b-hf",
+        mm_processor_kwargs=None,
+        limit_mm_per_prompt={"image": 1},
+        mm_processor_cache_gb=1,
+    )
+    model_config = ctx.model_config
+
+    sender = LruKeyReplicatedSenderCache(model_config)
+    receiver = LruKeyReplicatedReceiverCache(model_config)
+    item = _dummy_item({"pixel_values": 16})
+    mm_hash = "image_A"
+    payload, _ = sender.get_and_update_item((item, []), mm_hash)
+    assert receiver.get_and_update_item(payload, mm_hash) is item
+    assert sender.get_and_update_item(None, mm_hash)[0] is None
+
+    renderer = renderer_from_config(VllmConfig(model_config=model_config))
+    renderer._mm_processor_cache = sender
+
+    def release():
+        if release_error:
+            raise RuntimeError(release_error)
+        receiver.clear_cache()
+
+    engine = SimpleNamespace(
+        renderer=renderer,
+        engine_core=SimpleNamespace(
+            release_kv_cache_memory=release,
+            release_kv_cache_memory_async=make_async(
+                release,
+                executor=renderer._executor,
+            ),
+        ),
+        logger_manager=Mock(),
+    )
+
+    async def call_release():
+        if use_async:
+            await AsyncLLM.release_kv_cache_memory(engine)
+        else:
+            LLMEngine.release_kv_cache_memory(engine)
+
+    if release_error:
+        with pytest.raises(RuntimeError, match=release_error):
+            await call_release()
+        engine.logger_manager.record_sleep_state.assert_not_called()
+    else:
+        await call_release()
+        engine.logger_manager.record_sleep_state.assert_called_once_with(1, 0)
+
+    payload, _ = sender.get_and_update_item((item, []), mm_hash)
+    assert payload is item
+    assert receiver.get_and_update_item(payload, mm_hash) is item
+
+
 _SLEEP_VISION_PROMPT = (
     "<|im_start|>system\nYou are a helpful assistant.<|im_end|>"
     "\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
diff --git a/tests/multimodal/test_embedding_shape_validation.py b/tests/multimodal/test_embedding_shape_validation.py
index fcac5be5b5a7..3d7151271065 100644
--- a/tests/multimodal/test_embedding_shape_validation.py
+++ b/tests/multimodal/test_embedding_shape_validation.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Embedding shape validation in multimodal APIs.
+"""Embedding shape validation in multimodal APIs.
 
 Tests verify that embeddings with correct ndim but incorrect hidden_size
 are rejected before they can cause crashes during model inference.
@@ -178,8 +177,7 @@ class TestShapeValidationIntegration:
     """Integration tests verifying attack scenarios are blocked."""
 
     def test_attack_scenario_multimodal_image(self):
-        """
-        Simulate attack through Chat API with image embeddings.
+        """Simulate attack through Chat API with image embeddings.
 
         Verifies validation occurs in multimodal parser path.
         """
@@ -193,8 +191,7 @@ def test_attack_scenario_multimodal_image(self):
             parser.parse_mm_data({"image": attack_tensor})
 
     def test_attack_scenario_multimodal_audio(self):
-        """
-        Simulate attack through Chat API with audio embeddings.
+        """Simulate attack through Chat API with audio embeddings.
 
         Verifies validation occurs in multimodal parser path.
         """
@@ -208,8 +205,7 @@ def test_attack_scenario_multimodal_audio(self):
             parser.parse_mm_data({"audio": attack_tensor})
 
     def test_attack_scenario_multimodal_video(self):
-        """
-        Simulate attack through Chat API with video embeddings.
+        """Simulate attack through Chat API with video embeddings.
 
         Verifies validation occurs in multimodal parser path.
         """
diff --git a/tests/multimodal/test_embedding_shape_validation_unit.py b/tests/multimodal/test_embedding_shape_validation_unit.py
index 7966aad4e988..c8f785046eb5 100644
--- a/tests/multimodal/test_embedding_shape_validation_unit.py
+++ b/tests/multimodal/test_embedding_shape_validation_unit.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for embedding shape validation.
+"""Unit tests for embedding shape validation.
 
 Simple, fast unit tests that can run without server fixtures.
 Run with: pytest tests/multimodal/test_embedding_shape_validation_unit.py -v
@@ -167,16 +166,14 @@ def test_hidden_size_validation_wrong_size_rejected(self):
 
 
 class TestShapeValidationDoSPrevention:
-    """
-    Tests for DoS prevention through shape validation.
+    """Tests for DoS prevention through shape validation.
 
     Verifies that embeddings with incorrect shapes are rejected early,
     preventing crashes during model inference.
     """
 
     def test_prevent_crash_from_wrong_shape_image_embeds(self):
-        """
-        Prevent crash scenario: wrong hidden size in image embeddings.
+        """Prevent crash scenario: wrong hidden size in image embeddings.
 
         Without validation, this would pass initial checks but crash later
         during model forward pass when dimensions don't match.
@@ -198,9 +195,7 @@ def test_prevent_crash_from_wrong_shape_image_embeds(self):
         assert str(wrong_hidden_size) in error_msg  # Received
 
     def test_prevent_crash_from_wrong_shape_audio_embeds(self):
-        """
-        Prevent crash scenario: wrong hidden size in audio embeddings.
-        """
+        """Prevent crash scenario: wrong hidden size in audio embeddings."""
         expected_hidden_size = 768
         wrong_hidden_size = 4096
 
diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py
index 4810081ad40c..971d97a2ca94 100644
--- a/tests/multimodal/test_hasher.py
+++ b/tests/multimodal/test_hasher.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
 import hashlib
+import struct
 import uuid
 from io import BytesIO
 from pathlib import Path
@@ -8,7 +10,7 @@
 import numpy as np
 import pytest
 import torch
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageOps
 
 from vllm.config.multimodal import MMHasherAlgorithm
 from vllm.multimodal.hasher import MultiModalHasher
@@ -137,6 +139,33 @@ def item_for_hash(num_frames: int):
     )
 
 
+def test_hash_video_tensor_frames():
+    """Videos holding tensor frames (e.g. NVDEC-decoded) hash like
+    array-framed ones, from the original bytes without a D2H copy."""
+    source = b"x" * 100
+
+    def item_for_hash(frames):
+        metadata = {
+            "total_num_frames": 2,
+            "fps": 2.0,
+            "duration": 1.0,
+            "video_backend": "torchcodec",
+            "frames_indices": [0, 1],
+            "do_sample_frames": False,
+        }
+        video = MediaWithBytes((frames, metadata), source)
+        items = MultiModalDataParser()._parse_video_data([video])
+        return items.get_all_items_for_hash()[0]
+
+    np_frames = np.zeros((2, 8, 8, 3), dtype=np.uint8)
+    torch_frames = torch.zeros((2, 8, 8, 3), dtype=torch.uint8)
+
+    hasher = MultiModalHasher
+    assert hasher.hash_kwargs("blake3", video=item_for_hash(np_frames)) == (
+        hasher.hash_kwargs("blake3", video=item_for_hash(torch_frames))
+    )
+
+
 def test_hash_non_contiguous_array():
     arr = np.arange(24).reshape(4, 6).T
     assert not arr.flags.c_contiguous
@@ -172,6 +201,33 @@ def test_hash_image_exif_id():
     )
 
 
+def test_hash_image_malformed_exif():
+    # Test that images with malformed EXIF headers (e.g. invalid TIFF header)
+    # do not raise an unhandled exception during hashing and fall back to image data.
+    buf = BytesIO()
+    Image.new("RGB", (64, 48)).save(buf, "JPEG")
+    jpg = buf.getvalue()
+    rest = jpg[2:]
+    rest = rest[2 + struct.unpack(">H", rest[2:4])[0] :]
+    payload = b"Exif\x00\x00XXXX\x00\x00\x00\x08" + bytes(32)
+    data = b"\xff\xd8\xff\xe1" + struct.pack(">H", len(payload) + 2) + payload + rest
+
+    image = Image.open(BytesIO(data))
+    with contextlib.suppress(Exception):
+        image = ImageOps.exif_transpose(image)
+    image.load()
+
+    hasher = MultiModalHasher
+    # Should hash without raising SyntaxError or any other exception
+    hash_val = hasher.hash_kwargs("blake3", image=image)
+    assert isinstance(hash_val, str) and len(hash_val) > 0
+
+    # Also verify MediaWithBytes wrapping the image with malformed EXIF
+    media_item = MediaWithBytes(image, data)
+    hash_media = hasher.hash_kwargs("blake3", image=media_item)
+    assert isinstance(hash_media, str) and len(hash_media) > 0
+
+
 def _rgba_png_bytes() -> bytes:
     image = Image.new("RGBA", (8, 8), (255, 0, 0, 128))
     buf = BytesIO()
@@ -212,3 +268,43 @@ def test_hash_media_io_noop_config_preserves_hash():
     assert hasher.hash_kwargs("blake3", image=loaded) == hasher.hash_kwargs(
         "blake3", image=plain
     )
+
+
+# The digest input is a concatenation of byte chunks, so it has to be uniquely
+# decodable. Each case below is a pair of distinct processor kwargs that used to
+# serialize to the same bytes, which made two requests share an mm hash -- and
+# therefore share both the processor cache entry and the prefix-cache block key.
+IMAGE = b"\x89PNG\r\n\x1a\n"
+
+
+def _hash(**mm_processor_kwargs: object) -> str:
+    return MultiModalHasher.hash_kwargs(
+        "blake3", model_id="m", image=IMAGE, **mm_processor_kwargs
+    )
+
+
+def test_hash_collision_kwargs_key_value_boundary():
+    # Both used to flatten to b"ab" + b"c".
+    assert _hash(**{"ab": "c"}) != _hash(**{"a": "bc"})
+
+
+def test_hash_collision_nested_vs_flattened_key():
+    nested = _hash(size={"shortest_edge": 224})
+    flattened = _hash(**{"size.shortest_edge": 224})
+    assert nested != flattened
+
+
+def test_hash_collision_sequence_vs_mapping():
+    assert _hash(fps=[2, 4]) != _hash(fps={"0": 2, "1": 4})
+
+
+@pytest.mark.parametrize("empty", ["", b"", []])
+def test_hash_collision_none_vs_empty(empty):
+    assert _hash(video_pruning_rate=None) != _hash(video_pruning_rate=empty)
+
+
+def test_hash_collision_empty_container_vs_omitted():
+    omitted = _hash()
+    assert _hash(size={}) != omitted
+    assert _hash(size=[]) != omitted
+    assert _hash(size={}) != _hash(size=[])
diff --git a/tests/multimodal/test_parse.py b/tests/multimodal/test_parse.py
index 685e615f7832..c20677fc25c6 100644
--- a/tests/multimodal/test_parse.py
+++ b/tests/multimodal/test_parse.py
@@ -6,6 +6,7 @@
 from PIL import Image
 
 from vllm.multimodal.parse import (
+    AudioProcessorItems,
     ImageProcessorItems,
     MultiModalDataParser,
     VideoProcessorItems,
@@ -14,6 +15,45 @@
 H, W = 480, 640
 
 
+class AudioMetadataParser(MultiModalDataParser):
+    embedding_fields = {
+        "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
+    }
+
+
+@pytest.mark.parametrize("allow_missing", [False, True])
+def test_audio_metadata_requires_ec_consumer(allow_missing):
+    parser = AudioMetadataParser(allow_missing_mm_embeddings=allow_missing)
+    data = {"audio_num_tokens": torch.tensor([[3], [5]])}
+    if not allow_missing:
+        with pytest.raises(ValueError, match="audio_embeds"):
+            parser.parse_mm_data({"audio": data})
+    else:
+        items = parser.parse_mm_data({"audio": data})["audio"]
+        assert len(items) == 2
+        assert items.get(1)["audio_num_tokens"].item() == 5
+        assert items.get_processor_data() == {}
+
+
+@pytest.mark.parametrize("counts", [[0], [-1], [1.5], [True], [[1, 2]]])
+def test_audio_metadata_rejects_invalid_token_counts(counts):
+    parser = AudioMetadataParser(allow_missing_mm_embeddings=True)
+    with pytest.raises(ValueError, match="positive integer"):
+        parser.parse_mm_data({"audio": {"audio_num_tokens": torch.tensor(counts)}})
+
+
+def test_audio_metadata_checks_supplied_embedding_lengths():
+    parser = AudioMetadataParser()
+    data = {
+        "audio_num_tokens": torch.tensor([3, 5]),
+        "audio_embeds": [torch.zeros(3, 8), torch.zeros(5, 8)],
+    }
+    assert len(parser.parse_mm_data({"audio": data})["audio"]) == 2
+    data["audio_num_tokens"] = torch.tensor([3, 4])
+    with pytest.raises(ValueError, match="does not match"):
+        parser.parse_mm_data({"audio": data})
+
+
 @pytest.mark.parametrize(
     "image",
     [
@@ -55,9 +95,49 @@ def test_frame_size_hwc_chw(frame):
     assert items.get_frame_size(0) == (W, H)
 
 
+def test_video_with_metadata_tensor_passthrough():
+    """Tensor frames pass through unchanged regardless of device: HF video
+    processors accept tensors, and device-resident frames (e.g. NVDEC-decoded)
+    must not be copied back to host."""
+    frames = torch.zeros((4, H, W, 3), dtype=torch.uint8)
+    video, metadata = MultiModalDataParser()._get_video_with_metadata(frames)
+
+    assert video is frames
+    assert metadata is None
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+def test_video_with_metadata_keeps_device_tensor():
+    """Device-resident frames (e.g. NVDEC-decoded) pass through as tensors,
+    so a device-side HF processor can consume them without a D2H copy."""
+    frames = torch.zeros((4, H, W, 3), dtype=torch.uint8, device="cuda")
+    video, metadata = MultiModalDataParser()._get_video_with_metadata(frames)
+
+    assert video is frames
+    assert metadata is None
+
+
+@pytest.mark.parametrize(
+    "frames",
+    [
+        [np.zeros((H, W, 3), dtype=np.uint8) for _ in range(2)],
+        [torch.zeros((H, W, 3), dtype=torch.uint8) for _ in range(2)],
+    ],
+)
+def test_parse_video_frame_list_as_single_video(frames):
+    """A list of decoded frames must represent one video item."""
+    items = MultiModalDataParser().parse_mm_data({"video": frames})["video"]
+
+    assert items.get_count() == 1
+    video = items.get(0)
+    assert isinstance(video, np.ndarray)
+    np.testing.assert_array_equal(video, np.stack([np.asarray(f) for f in frames]))
+
+
 @pytest.mark.parametrize(
     "modality,processor_cls",
     [
+        ("audio", AudioProcessorItems),
         ("image", ImageProcessorItems),
         ("video", VideoProcessorItems),
     ],
@@ -68,3 +148,19 @@ def test_parse_mm_data_accepts_none_cached_item(modality, processor_cls):
     assert isinstance(items, processor_cls)
     assert len(items) == 1
     assert items.get(0) is None
+
+
+def test_cached_audio_items_preserve_positions_during_resampling():
+    waveform = np.arange(16, dtype=np.float32)
+    parser = MultiModalDataParser(
+        target_sr=16000, target_channels=1, audio_resample_method="scipy"
+    )
+    items = parser.parse_mm_data(
+        {"audio": [None, (waveform, 8000), None, (waveform, 16000)]}
+    )["audio"]
+
+    assert len(items) == 4
+    assert items.get(0) is None
+    assert items.get(2) is None
+    assert len(items.get(1)) == 32
+    np.testing.assert_array_equal(items.get(3), waveform)
diff --git a/tests/multimodal/test_processing.py b/tests/multimodal/test_processing.py
index 4a41013b21a5..0bd19a0ea0dc 100644
--- a/tests/multimodal/test_processing.py
+++ b/tests/multimodal/test_processing.py
@@ -7,7 +7,7 @@
 import numpy as np
 import pytest
 
-from vllm.config import ModelConfig
+from vllm.config import ModelConfig, SchedulerConfig
 from vllm.exceptions import VLLMValidationError
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.hasher import MultiModalHasher
@@ -32,6 +32,7 @@
 )
 from vllm.utils.collection_utils import flatten_2d_lists
 
+from ..models.utils import build_model_context
 from .utils import random_image
 
 pytestmark = pytest.mark.cpu_test
@@ -968,8 +969,7 @@ def test_hf_processor_call_kwargs(
 
 
 def test_apply_matches_no_match_exits_quickly():
-    """
-    Test that _apply_matches exits quickly when no matches are found.
+    """Test that _apply_matches exits quickly when no matches are found.
 
     Previously, _apply_matches had O(n²) behavior when no match was found
     because it would increment start_idx by 1 each iteration while
@@ -1027,8 +1027,7 @@ def test_iter_token_matches_rejects_negative_start_idx():
 
 
 def test_find_mm_placeholders_avoids_quadratic_false_prefixes():
-    """
-    Test that placeholder scanning stays linear under adversarial candidates.
+    """Test that placeholder scanning stays linear under adversarial candidates.
 
     The fast-forward scan must not rescan the prompt tail per position when
     one candidate's first token never occurs (forcing a full search) while
@@ -1061,8 +1060,7 @@ def test_find_mm_placeholders_avoids_quadratic_false_prefixes():
     ],
 )
 def test_find_mm_placeholders_stops_at_missing_item(prompt):
-    """
-    Test that the scan returns no placeholders once it fails to find
+    """Test that the scan returns no placeholders once it fails to find
     an item's placeholder, leaving later items unresolved.
     """
     result = find_mm_placeholders(
@@ -1079,8 +1077,7 @@ def test_find_mm_placeholders_stops_at_missing_item(prompt):
 
 
 class _FakeTokenizer:
-    """
-    Character-level tokenizer where "foo" merges into one token differently
+    """Character-level tokenizer where "foo" merges into one token differently
     depending on whether it is followed by "d", like BPE merging "foo" in
     "food" across the search-text boundary.
     """
@@ -1145,8 +1142,7 @@ def _text_fallback_processor() -> BaseMultiModalProcessor:
 
 
 def test_apply_prompt_updates_falls_back_to_text_matching():
-    """
-    Test that the fallback in `_apply_prompt_updates` finds targets that
+    """Test that the fallback in `_apply_prompt_updates` finds targets that
     tokenize differently inside the prompt ("foo" in "food").
     """
     processor = _text_fallback_processor()
@@ -1166,8 +1162,7 @@ def test_apply_prompt_updates_falls_back_to_text_matching():
 
 
 def test_apply_prompt_updates_falls_back_with_prefix_target():
-    """
-    Test that `PromptIndexTargets.prefix` targets are resolved against the
+    """Test that `PromptIndexTargets.prefix` targets are resolved against the
     decoded text in the fallback path of `_apply_prompt_updates`.
     """
     processor = _text_fallback_processor()
@@ -1192,8 +1187,7 @@ def test_apply_prompt_updates_falls_back_with_prefix_target():
 
 
 def test_apply_prompt_updates_falls_back_with_index_targets():
-    """
-    Test that the text resolvers of `PromptIndexTargets.start`/`end`
+    """Test that the text resolvers of `PromptIndexTargets.start`/`end`
     match against the decoded text when another item forces the
     fallback in `_apply_prompt_updates`.
     """
@@ -1331,3 +1325,74 @@ def test_processor_inputs_hashes_ignore_unrelated_kwargs():
     )
 
     assert inputs.get_mm_hashes("test-model", "blake3") == {"image": ["image-uuid"]}
+
+
+@pytest.mark.parametrize(
+    ("left", "right"),
+    [
+        # Shifting the key/value boundary: both flatten to the dotted key
+        # "mm_processor_kwargs.abc" followed by no value bytes.
+        ({"ab": "c"}, {"a": "bc"}),
+        # A nested mapping and a caller-supplied dotted key flatten alike.
+        ({"size": {"shortest_edge": 224}}, {"size.shortest_edge": 224}),
+        # A sequence and a mapping keyed by stringified indices flatten alike.
+        ({"fps": [2, 4]}, {"fps": {"0": 2, "1": 4}}),
+        # None contributes the key alone, which a zero-byte value also does.
+        ({"video_pruning_rate": None}, {"video_pruning_rate": ""}),
+        # An empty container contributes nothing, as does omitting the key.
+        ({"size": {}}, {}),
+    ],
+)
+def test_processor_inputs_hashes_distinguish_kwargs_shapes(left, right):
+    """Distinct processor kwargs must not share a multi-modal hash.
+
+    ``hf_processor_mm_kwargs`` is per-request input, so both the keys and the
+    values here are caller-controlled. The hash is the identity of the
+    processor cache entry and is mixed into the prefix-cache block key, so two
+    requests sharing one is a cross-request cache hit.
+    """
+    image = random_image(np.random.RandomState(0), min_wh=8, max_wh=9)
+    mm_data_items = MultiModalDataParser().parse_mm_data({"image": [image]})
+
+    def hash_with(hf_processor_mm_kwargs):
+        return ProcessorInputs(
+            prompt=[],
+            mm_data_items=mm_data_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+        ).get_mm_hashes("test-model", "blake3")["image"][0]
+
+    assert hash_with(left) != hash_with(right)
+
+
+@pytest.mark.parametrize(
+    ("chunked_prefill", "max_model_len", "expected_seq_len"),
+    [(None, 491520, 491520), (True, 491520, 8192), (True, 128, 128), (False, 128, 128)],
+)
+def test_dummy_inputs_scheduler_budget(
+    chunked_prefill, max_model_len, expected_seq_len
+):
+    ctx = build_model_context(
+        "llava-hf/llava-v1.6-mistral-7b-hf",
+        mm_processor_kwargs=None,
+        limit_mm_per_prompt={"image": 1},
+    )
+    ctx.model_config.max_model_len = max_model_len
+
+    processor = MULTIMODAL_REGISTRY.create_processor(
+        ctx.model_config,
+        tokenizer=ctx.tokenizer,
+    )
+    processor.apply = lambda *args, **kwargs: {"prompt_token_ids": [7]}
+
+    kwargs = {}
+    if chunked_prefill is not None:
+        kwargs["scheduler_config"] = SchedulerConfig(
+            max_model_len=max_model_len,
+            is_encoder_decoder=False,
+            max_num_batched_tokens=8192,
+            max_num_seqs=1,
+            enable_chunked_prefill=chunked_prefill,
+        )
+
+    result = processor.get_dummy_mm_inputs({"image": 1}, **kwargs)
+    assert len(result["prompt_token_ids"]) == expected_seq_len
diff --git a/tests/multimodal/test_registry.py b/tests/multimodal/test_registry.py
index 99f0689c7da3..f395d3693334 100644
--- a/tests/multimodal/test_registry.py
+++ b/tests/multimodal/test_registry.py
@@ -1,71 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for MultiModalRegistry.supports_multimodal_inputs and
+"""Unit tests for MultiModalRegistry.supports_multimodal_inputs and
 Qwen2.5-VL visual component loading behavior.
 """
 
 from types import SimpleNamespace
-from unittest.mock import MagicMock
 
 import pytest
 
-from vllm.config import SchedulerConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
-from ..models.utils import build_model_context
-
 pytestmark = pytest.mark.cpu_test
 
 
-@pytest.mark.parametrize(
-    ("chunked_prefill", "max_model_len", "expected_seq_len"),
-    [(None, 491520, 491520), (True, 491520, 8192), (True, 128, 128), (False, 128, 128)],
-)
-def test_dummy_inputs_scheduler_budget(
-    chunked_prefill, max_model_len, expected_seq_len
-):
-    model_config = MagicMock()
-    model_config.max_model_len = max_model_len
-    processor = MagicMock()
-    processor.apply.return_value = {"prompt_token_ids": [7]}
-    kwargs = {}
-    if chunked_prefill is not None:
-        kwargs["scheduler_config"] = SchedulerConfig(
-            max_model_len=max_model_len,
-            is_encoder_decoder=False,
-            max_num_batched_tokens=8192,
-            max_num_seqs=1,
-            enable_chunked_prefill=chunked_prefill,
-        )
-    result = MULTIMODAL_REGISTRY.get_dummy_mm_inputs(
-        model_config, {"image": 1}, processor=processor, **kwargs
-    )
-    get_inputs = processor.dummy_inputs.get_dummy_processor_inputs
-    assert get_inputs.call_args.kwargs["seq_len"] == expected_seq_len
-    assert len(result["prompt_token_ids"]) == expected_seq_len
-
-
-@pytest.mark.parametrize(
-    "model_id,limit_mm_per_prompt,expected",
-    [
-        ("Qwen/Qwen2-0.5B-Instruct", {}, False),
-        ("Qwen/Qwen2.5-VL-3B-Instruct", {}, True),
-        ("Qwen/Qwen2.5-VL-3B-Instruct", {"image": 0, "video": 0}, False),
-        ("Qwen/Qwen2.5-VL-3B-Instruct", {"image": 0}, True),
-    ],
-)
-@pytest.mark.core_model
-def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected):
-    """Test supports_multimodal_inputs returns correct boolean for various
-    configs."""
-    ctx = build_model_context(
-        model_id,
-        limit_mm_per_prompt=limit_mm_per_prompt,
-    )
-    assert MULTIMODAL_REGISTRY.supports_multimodal_inputs(ctx.model_config) is expected
-
-
 def test_create_processor_error_uses_served_model_name():
     model_config = SimpleNamespace(
         is_multimodal_model=False,
diff --git a/tests/multimodal/test_sparse_tensor_validation_unit.py b/tests/multimodal/test_sparse_tensor_validation_unit.py
index 2eec8ea8283a..4e35e4abc81c 100644
--- a/tests/multimodal/test_sparse_tensor_validation_unit.py
+++ b/tests/multimodal/test_sparse_tensor_validation_unit.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for sparse tensor validation.
+"""Unit tests for sparse tensor validation.
 
 Simple, fast unit tests that can run without server fixtures.
 Run with: pytest tests/multimodal/test_sparse_tensor_validation_unit.py -v
@@ -56,8 +55,7 @@ def test_negative_indices_rejected(self):
                 tensor.to_dense()
 
     def test_without_context_manager_allows_invalid(self):
-        """
-        WITHOUT validation, invalid tensors may not immediately error.
+        """WITHOUT validation, invalid tensors may not immediately error.
 
         This demonstrates the vulnerability: PyTorch 2.8.0+ doesn't validate
         by default, which can lead to memory corruption.
diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py
index 9c24b5f96932..432521c11f74 100644
--- a/tests/multimodal/test_video.py
+++ b/tests/multimodal/test_video.py
@@ -11,6 +11,7 @@
 import numpy as np
 import numpy.typing as npt
 import pytest
+import torch
 from transformers import AutoVideoProcessor
 from transformers.video_utils import VideoMetadata
 
@@ -173,9 +174,14 @@ def fake_import_module(name: str, package: str):
     [
         (
             "torchcodec",
-            {"min_frames": 4, "num_ffmpeg_threads": 2, "seek_mode": "approximate"},
+            {
+                "min_frames": 4,
+                "num_ffmpeg_threads": 2,
+                "seek_mode": "approximate",
+                "device": "cuda",
+            },
             {"min_frames": 4},
-            {"num_ffmpeg_threads": 2, "seek_mode": "approximate"},
+            {"num_ffmpeg_threads": 2, "seek_mode": "approximate", "device": "cuda"},
         ),
         (
             "deepstream",
@@ -205,6 +211,11 @@ def test_video_backend_rejects_options_for_another_decoder():
     ):
         resolve_video_backend_kwargs("opencv", {"num_ffmpeg_threads": 2})
 
+    with pytest.raises(
+        ValueError, match="device is not supported by the 'opencv' backend"
+    ):
+        resolve_video_backend_kwargs("opencv", {"device": "cuda"})
+
 
 @pytest.mark.parametrize(
     ("backend", "error"),
@@ -831,8 +842,7 @@ def test_video_processor_from_model_repo(
 
 
 def test_video_backend_handles_broken_frames(monkeypatch: pytest.MonkeyPatch):
-    """
-    Regression test for handling videos with broken frames.
+    """Regression test for handling videos with broken frames.
     This test uses a pre-corrupted video file (assets/corrupted.mp4) that
     contains broken frames to verify the video loader handles
     them gracefully without crashing and returns accurate metadata.
@@ -871,8 +881,7 @@ def test_video_backend_handles_broken_frames(monkeypatch: pytest.MonkeyPatch):
 def test_video_backend_handles_edit_list_trimmed_video(
     monkeypatch: pytest.MonkeyPatch,
 ):
-    """
-    An mp4 edit list (e.g. from a lossless ``ffmpeg -ss ... -c copy`` cut)
+    """An mp4 edit list (e.g. from a lossless ``ffmpeg -ss ... -c copy`` cut)
     hides the decode lead-in: the header still counts every physical sample
     while sequential decode only yields the visible frames. Sampling over the
     header count used to collapse such videos to the few indices below the
@@ -912,8 +921,7 @@ def test_video_backend_handles_edit_list_trimmed_video(
 
 
 def test_video_recovery_simulated_failures(monkeypatch: pytest.MonkeyPatch):
-    """
-    Test that frame recovery correctly uses the next valid frame when
+    """Test that frame recovery correctly uses the next valid frame when
     target frames fail to load.
 
     Uses corrupted.mp4 and mocks VideoCapture.grab() to fail on specific
@@ -1008,8 +1016,7 @@ def release(self):
 
 
 def test_video_recovery_with_corrupted_file(monkeypatch: pytest.MonkeyPatch):
-    """
-    Test frame recovery with an actual corrupted video file using sparse sampling.
+    """Test frame recovery with an actual corrupted video file using sparse sampling.
 
     This test uses corrupted.mp4 which has genuine H.264 codec errors on
     frame 17. With num_frames=8, the target frames are [0, 3, 7, 10, 14, 17, 21, 25].
@@ -1077,8 +1084,7 @@ def test_video_recovery_with_corrupted_file(monkeypatch: pytest.MonkeyPatch):
 
 
 def test_video_recovery_dynamic_backend(monkeypatch: pytest.MonkeyPatch):
-    """
-    Test that frame_recovery works with the dynamic video backend.
+    """Test that frame_recovery works with the dynamic video backend.
 
     The dynamic backend samples frames based on fps/duration rather than
     loading all frames. This test verifies recovery works in that context.
@@ -1205,6 +1211,30 @@ def test_torchcodec_backend_rejects_frame_recovery(dummy_video_path):
         )
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+def test_torchcodec_backend_cuda_decodes_on_gpu(dummy_video_path):
+    """With device="cuda", torchcodec decodes via NVDEC and keeps the frames
+    on the GPU instead of returning a host-side numpy array."""
+    pytest.importorskip("torchcodec")
+
+    with open(dummy_video_path, "rb") as f:
+        video_data = f.read()
+
+    loader = VIDEO_LOADER_REGISTRY.load("opencv")
+    frames, metadata = loader.load_bytes(
+        video_data, num_frames=8, backend="torchcodec", device="cuda"
+    )
+
+    assert isinstance(frames, torch.Tensor)
+    assert frames.device.type == "cuda"
+    assert frames.dtype == torch.uint8
+    assert frames.ndim == 4
+    assert frames.shape[3] == 3  # RGB
+    assert frames.shape[0] == 8
+    assert frames.shape[0] == len(metadata["frames_indices"])
+    assert metadata["video_backend"] == "torchcodec"
+
+
 def test_torchcodec_backend_returns_target_frames_not_keyframes():
     """Regression test: torchcodec must return the requested frames, not the
     GOP keyframe they seek back to.
diff --git a/vllm/models/deepseek_v4_1/common/__init__.py b/tests/parser/cohere/__init__.py
similarity index 100%
rename from vllm/models/deepseek_v4_1/common/__init__.py
rename to tests/parser/cohere/__init__.py
diff --git a/tests/parser/cohere/conftest.py b/tests/parser/cohere/conftest.py
new file mode 100644
index 000000000000..4477cacf43a3
--- /dev/null
+++ b/tests/parser/cohere/conftest.py
@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+
+from .utils import MockCohereTokenizer
+
+
+@pytest.fixture(scope="package")
+def tokenizer() -> MockCohereTokenizer:
+    return MockCohereTokenizer()
+
+
+@pytest.fixture
+def request_obj() -> ChatCompletionRequest:
+    return ChatCompletionRequest(messages=[], model="test-model")
diff --git a/tests/parser/cohere/test_citations.py b/tests/parser/cohere/test_citations.py
new file mode 100644
index 000000000000..14ee1612ead4
--- /dev/null
+++ b/tests/parser/cohere/test_citations.py
@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Citation source resolution and the streamed citation wire shape."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+from vllm.parser.cohere_command import (
+    _melody_citations_to_vllm,
+    _melody_sources_to_vllm,
+)
+
+from .utils import make_parser, stream_parser
+
+
+class TestMelodySourceResolution:
+    """Pin the parser-side source resolution (see
+    ``_melody_sources_to_vllm``). The parser receives melody's numeric
+    ``(tool_call_index, tool_result_indices)`` addressing and resolves
+    it against a ``position_to_source`` map handed in by
+    ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same
+    map that would otherwise live in the serving layer's resolver.
+    """
+
+    @staticmethod
+    def _fake_melody_source(bucket: int, indices: list[int]) -> Any:
+        return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices)
+
+    def test_multi_index_source_fans_out(self):
+        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
+
+        position_map: dict[tuple[int, int], CitationSource] = {
+            (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}),
+            (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}),
+        }
+        raw = [self._fake_melody_source(0, [0, 1])]
+        out = _melody_sources_to_vllm(raw, position_map)
+        assert [s.id for s in out] == ["d0", "d1"]
+        # Verify type / payload were plumbed through, not just the id.
+        assert out[0].type == "document"
+        assert out[0].document == {"id": "d0"}
+
+    def test_unresolvable_position_skipped(self):
+        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
+
+        position_map: dict[tuple[int, int], CitationSource] = {
+            (0, 0): CitationSource(type="document", id="d0"),
+        }
+        raw = [self._fake_melody_source(9, [0])]
+        assert _melody_sources_to_vllm(raw, position_map) == []
+
+    def test_missing_position_map_drops_all_sources(self):
+        # A parser instance without a position map (parser wired
+        # outside of ``CohereServingChatV2``) can't attribute anything,
+        # so every source is dropped. Callers downstream will see the
+        # citation with empty ``sources`` and drop it entirely.
+        raw = [self._fake_melody_source(0, [0])]
+        assert _melody_sources_to_vllm(raw, None) == []
+
+    def test_citations_pass_through_is_thinking_tag(self):
+        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
+
+        position_map: dict[tuple[int, int], CitationSource] = {
+            (0, 0): CitationSource(type="document", id="d0"),
+        }
+        raw = [
+            SimpleNamespace(
+                start_index=0,
+                end_index=5,
+                text="hello",
+                is_thinking=True,
+                sources=[self._fake_melody_source(0, [0])],
+            )
+        ]
+        out = _melody_citations_to_vllm(raw, position_map)
+        assert out is not None
+        assert out[0].type == "THINKING_CONTENT"
+        assert out[0].sources[0].id == "d0"
+
+
+class TestParserStreamingEndToEnd:
+    """End-to-end streaming shape: raw model output text is tokenized
+    with the mock byte-level tokenizer and fed one token at a time
+    through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`.
+    We collect the ordered sequence of ``DeltaMessage``s the parser
+    hands back -- i.e. what the client would receive on the wire --
+    and pin the exact shape.
+
+    Combines the three flavors that reach ``delta``: a thinking block
+    (populates ``reasoning``), a text block (populates ``content``),
+    and a resolved citation (populates ``citations``) referencing
+    ``(start, end)`` offsets into the already-streamed content.
+    """
+
+    _MODEL_OUTPUT = (
+        "<|START_THINKING|>Let me check.<|END_THINKING|>"
+        "<|START_RESPONSE|>The capital is <co>Paris</co: 0:[1]>.<|END_RESPONSE|>"
+    )
+
+    def test_full_wire_stream_shape(self, tokenizer, request_obj):
+        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
+        from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
+
+        # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs``
+        # installs on a request whose tool_call_index=0 tool result at
+        # index 1 is the "France" document.
+        position_to_source = {
+            (0, 1): CitationSource(
+                type="document",
+                id="doc-paris",
+                document={"id": "doc-paris", "title": "France"},
+            ),
+        }
+
+        parser = make_parser(
+            tokenizer,
+            "cohere_command3",
+            chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source},
+        )
+        deltas = stream_parser(parser, request_obj, tokenizer, self._MODEL_OUTPUT)
+
+        # -- 1. Reasoning and content stream separately and reach
+        # the wire only as their real payload. Framing tokens
+        # (``<|START_THINKING|>`` etc.) are consumed by the parser's
+        # state machine and never produce a delta.
+        reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None]
+        content_stream = [d.content for d in deltas if d.content is not None]
+        citation_deltas = [d for d in deltas if getattr(d, "citations", None)]
+
+        assert "".join(reasoning_stream) == "Let me check."
+        assert "".join(content_stream) == "The capital is Paris."
+
+        # -- 2. Exactly one citation reaches the wire, resolved to the
+        # source from ``position_to_source`` (not the raw
+        # ``(bucket, idx)`` coordinates).
+        assert len(citation_deltas) == 1
+        (cite_delta,) = citation_deltas
+        assert cite_delta.citations is not None
+        assert len(cite_delta.citations) == 1
+        cite = cite_delta.citations[0]
+        assert cite.text == "Paris"
+        # Char offsets are into the accumulated ``content``: "The
+        # capital is " is 15 chars, "Paris" is 5.
+        assert cite.start == 15
+        assert cite.end == 20
+        assert len(cite.sources) == 1
+        assert cite.sources[0].id == "doc-paris"
+        assert cite.sources[0].type == "document"
+
+        # -- 3. The citation's ``(start, end)`` offsets point into
+        # bytes that were already emitted as content-deltas earlier
+        # in the stream -- i.e. the citation is anchored to real
+        # already-shipped content.
+        cite_delta_idx = deltas.index(cite_delta)
+        earlier_content = "".join(
+            d.content for d in deltas[:cite_delta_idx] if d.content is not None
+        )
+        assert earlier_content[cite.start : cite.end] == "Paris"
+
+        # Everything after the citation delta is the trailing period.
+        later_content = "".join(
+            d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None
+        )
+        assert later_content == "."
+
+        # -- 4. Reasoning and content are never mixed on the same
+        # delta: the parser flips modes on the ``<|END_THINKING|>``
+        # boundary and never emits a delta with both fields set.
+        for d in deltas:
+            assert not (d.reasoning is not None and d.content is not None)
diff --git a/tests/parser/cohere/test_reasoning.py b/tests/parser/cohere/test_reasoning.py
new file mode 100644
index 000000000000..8c3311e8388c
--- /dev/null
+++ b/tests/parser/cohere/test_reasoning.py
@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reasoning extraction, reasoning-end gating and parser selection for the
+unified Cohere Command parser."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.parser import ParserManager
+from vllm.reasoning.cohere_command_reasoning_parser import (
+    CohereCommand3ReasoningParser,
+    CohereCommand4ReasoningParser,
+)
+
+from .utils import make_parser, stream_parser
+
+
+@dataclass
+class ReasoningCase:
+    parser_name: str
+    model_output: str
+    expected_reasoning: str | None
+    expected_content: str | None
+
+
+REASONING_CASES = [
+    pytest.param(
+        ReasoningCase(
+            parser_name="cohere_command3",
+            model_output="""\
+<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
+[
+    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
+]
+<|END_ACTION|>""",
+            expected_reasoning="i will call foo with query1",
+            expected_content=None,
+        ),
+        id="cmd3-single_tool_call",
+    ),
+    pytest.param(
+        ReasoningCase(
+            parser_name="cohere_command4",
+            model_output="""\
+<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
+[
+    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
+]
+<|END_ACTION|>""",
+            expected_reasoning="i will call foo with query1",
+            expected_content=None,
+        ),
+        id="cmd4-single_tool_call",
+    ),
+    pytest.param(
+        ReasoningCase(
+            parser_name="cohere_command3",
+            model_output="""\
+<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
+<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
+            expected_reasoning="This is a rainbow emoji: 🌈",
+            expected_content="foo bar",
+        ),
+        id="cmd3-citations_with_emoji",
+    ),
+    pytest.param(
+        ReasoningCase(
+            parser_name="cohere_command4",
+            model_output="""\
+<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
+<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
+            expected_reasoning="This is a rainbow emoji: 🌈",
+            expected_content="foo bar",
+        ),
+        id="cmd4-citations_with_emoji",
+    ),
+]
+
+
+@pytest.mark.parametrize("case", REASONING_CASES)
+class TestExtractReasoning:
+    def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase):
+        parser = make_parser(tokenizer, case.parser_name)
+        reasoning, content, tool_calls = parser.parse(
+            case.model_output, request_obj, enable_auto_tools=True
+        )
+        assert reasoning == case.expected_reasoning
+        assert content == case.expected_content
+        # No tools on the request, so any tool call is dropped, not surfaced.
+        assert not tool_calls
+
+    def test_streaming(self, tokenizer, request_obj, case: ReasoningCase):
+        parser = make_parser(tokenizer, case.parser_name)
+        deltas = stream_parser(parser, request_obj, tokenizer, case.model_output)
+
+        reasoning = "".join(d.reasoning for d in deltas if d.reasoning) or None
+        content = "".join(d.content for d in deltas if d.content) or None
+        assert reasoning == case.expected_reasoning
+        assert content == case.expected_content
+        assert all(not d.tool_calls for d in deltas)
+
+
+class TestFramingTokensStripped:
+    @pytest.mark.parametrize(
+        ("parser_name", "content_tags"),
+        [
+            pytest.param(
+                "cohere_command4", ("<|START_TEXT|>", "<|END_TEXT|>"), id="cmd4"
+            ),
+            pytest.param(
+                "cohere_command3", ("<|START_RESPONSE|>", "<|END_RESPONSE|>"), id="cmd3"
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("chunk_size", [1, 4], ids=["per_token", "batched"])
+    @pytest.mark.parametrize("with_tools", [False, True], ids=["no_tools", "tools"])
+    def test_content_framing_tokens_stripped(
+        self, tokenizer, parser_name, content_tags, chunk_size, with_tools
+    ):
+        parser = make_parser(tokenizer, parser_name)
+        tools = [{"type": "function", "function": {"name": "foo"}}]
+        request = ChatCompletionRequest(
+            messages=[], model="test-model", tools=tools if with_tools else None
+        )
+        parser.adjust_request(request)
+
+        start_tag, end_tag = content_tags
+        generation = (
+            f"<|START_THINKING|>Think deeply. The user greets us.<|END_THINKING|>"
+            f"{start_tag}I'm doing well, thank you{end_tag}"
+        )
+        deltas = stream_parser(parser, request, tokenizer, generation, chunk_size)
+
+        assert "".join(d.reasoning or "" for d in deltas) == (
+            "Think deeply. The user greets us."
+        )
+        assert not any(d.tool_calls for d in deltas)
+        assert "".join(d.content or "" for d in deltas) == "I'm doing well, thank you"
+
+
+class TestIsReasoningEnd:
+    @pytest.mark.parametrize(
+        "parser_cls",
+        [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser],
+        ids=["cmd3", "cmd4"],
+    )
+    def test_is_reasoning_end(self, tokenizer, parser_cls):
+        parser = parser_cls(tokenizer)
+        start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
+        end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
+        chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
+        content_ids = [99, 100]
+
+        # Generation-only tokens have no chatbot marker, so the whole sequence
+        # is considered.
+        assert parser.is_reasoning_end([end_id])
+        assert parser.is_reasoning_end([start_id, *content_ids, end_id])
+        assert not parser.is_reasoning_end([start_id, *content_ids])
+
+        # Full prompt/history tokens are scoped to the latest chatbot marker,
+        # so stray thinking tokens from the preamble or previous turns are ignored.
+        assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids])
+        assert parser.is_reasoning_end(
+            [start_id, end_id, chatbot_id, start_id, *content_ids, end_id]
+        )
+
+
+@pytest.mark.parametrize("parser_name", ["cohere_command3", "cohere_command4"])
+def test_count_reasoning_tokens(tokenizer, parser_name):
+    parser = make_parser(tokenizer, parser_name)
+    start = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
+    end = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
+
+    assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2
+    assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3
+    assert parser.count_reasoning_tokens([start, 1, 2]) == 2
+    assert parser.count_reasoning_tokens([1, 2, end]) == 0
+
+
+class TestParserSelection:
+    @pytest.mark.parametrize("with_tool_parser", [True, False])
+    def test_adjust_request_keeps_framing_tokens(self, tokenizer, with_tool_parser):
+        """Melody needs the special tokens, with or without a tool parser."""
+        cls = ParserManager.get_parser(
+            "cohere_command4" if with_tool_parser else None,
+            "cohere_command4",
+            enable_auto_tools=True,
+        )
+        request = ChatCompletionRequest(messages=[], model="m")
+        assert request.skip_special_tokens
+        cls(tokenizer).adjust_request(request)
+        assert request.skip_special_tokens is False
+
+    def test_non_cohere_selection_does_not_import_melody_parser(self, monkeypatch):
+        """``cohere_melody`` is optional; other parsers must not depend on it."""
+        import sys
+
+        monkeypatch.setitem(sys.modules, "vllm.parser.cohere_command", None)
+        assert ParserManager.get_parser("hermes", "deepseek_r1", True) is not None
diff --git a/tests/parser/cohere/test_structural_tags.py b/tests/parser/cohere/test_structural_tags.py
new file mode 100644
index 000000000000..27440d52af81
--- /dev/null
+++ b/tests/parser/cohere/test_structural_tags.py
@@ -0,0 +1,362 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``adjust_request``: folding tools and JSON schemas into Cohere structural tags."""
+
+from __future__ import annotations
+
+import json
+from collections import UserDict
+from types import SimpleNamespace
+
+import pytest
+from pydantic import ValidationError
+
+from vllm.entrypoints.generate.base.protocol import (
+    JsonSchemaResponseFormat,
+    ResponseFormat,
+    StructuralTagResponseFormat,
+)
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.parser.cohere_command import (
+    CohereCommandParser,
+    _has_effective_tools,
+    _response_format_type,
+    _schema_dict_from_structured_outputs,
+    convert_schema_to_structural_tags,
+)
+from vllm.sampling_params import StructuredOutputsParams
+
+from .utils import MockCohereTokenizer, make_parser
+
+SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}}
+SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}}
+GET_WEATHER_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "get_weather",
+        "parameters": {
+            "type": "object",
+            "properties": {"city": {"type": "string"}},
+        },
+    },
+}
+VALID_STRUCTURAL_TAG = {
+    "type": "structural_tag",
+    "format": {
+        "type": "triggered_tags",
+        "tags": [
+            {
+                "begin": "<tool>",
+                "content": {"type": "any_text"},
+                "end": "</tool>",
+            }
+        ],
+        "triggers": ["<tool>"],
+    },
+}
+
+
+def _model_config(arch: str) -> SimpleNamespace:
+    return SimpleNamespace(
+        architecture=arch,
+        architectures=[arch],
+        hf_text_config=SimpleNamespace(architectures=[arch]),
+    )
+
+
+def _make_chat_request(**kwargs) -> ChatCompletionRequest:
+    data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"}
+    data.update(kwargs)
+    return ChatCompletionRequest.model_validate(data)
+
+
+def _first_json_schema(tag_json: str) -> dict | None:
+    outer = json.loads(tag_json)
+    for t in (outer.get("format") or {}).get("tags") or []:
+        c = t.get("content") or {}
+        if c.get("type") == "json_schema":
+            js = c.get("json_schema")
+            return js if isinstance(js, dict) else None
+    return None
+
+
+def _content_types(tag_json: str) -> set[str]:
+    outer = json.loads(tag_json)
+    out: set[str] = set()
+    for t in (outer.get("format") or {}).get("tags") or []:
+        ty = (t.get("content") or {}).get("type")
+        if isinstance(ty, str):
+            out.add(ty)
+    return out
+
+
+@pytest.fixture(scope="module")
+def parser(tokenizer: MockCohereTokenizer) -> CohereCommandParser:
+    """Parser configured with a supported Cohere architecture."""
+    return make_parser(
+        tokenizer, "cohere_command4", model_config=_model_config("Cohere2ForCausalLM")
+    )
+
+
+@pytest.fixture(scope="module")
+def parser_no_model_config(tokenizer: MockCohereTokenizer) -> CohereCommandParser:
+    """Parser with no ``model_config`` (cannot resolve architecture)."""
+    return make_parser(tokenizer, "cohere_command4")
+
+
+@pytest.fixture(scope="module")
+def parser_unsupported_arch(tokenizer: MockCohereTokenizer) -> CohereCommandParser:
+    """Parser configured with an architecture that has no structural tag style."""
+    return make_parser(
+        tokenizer, "cohere_command4", model_config=_model_config("LlamaForCausalLM")
+    )
+
+
+class TestAdjustRequestPassthrough:
+    def test_structured_outputs_structural_tag_not_modified(self, parser) -> None:
+        tag = json.dumps(VALID_STRUCTURAL_TAG)
+        r = _make_chat_request(structured_outputs={"structural_tag": tag})
+        o = parser.adjust_request(r)
+        assert o.structured_outputs.structural_tag == tag
+
+    def test_response_format_structural_tag_short_circuit(self, parser) -> None:
+        # ``ChatCompletionRequest`` validates ``response_format`` as a union;
+        # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model).
+        rf = StructuralTagResponseFormat(
+            type="structural_tag",
+            format=VALID_STRUCTURAL_TAG["format"],
+        )
+        r = _make_chat_request(response_format=rf)
+        o = parser.adjust_request(r)
+        assert _response_format_type(o.response_format) == "structural_tag"
+        assert o.structured_outputs is None
+
+
+class TestAdjustRequestNoOp:
+    def test_no_schema_no_tools(self, parser) -> None:
+        o = parser.adjust_request(_make_chat_request())
+        assert o.structured_outputs is None
+        assert o.response_format is None
+
+    def test_no_model_config(self, parser_no_model_config) -> None:
+        inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
+        r = _make_chat_request(
+            response_format=ResponseFormat(type="json_schema", json_schema=inner),
+        )
+        o = parser_no_model_config.adjust_request(r)
+        assert o.response_format is not None
+        assert o.structured_outputs is None
+
+
+class TestAdjustRequestUnsupportedArchitecture:
+    def test_json_schema_raises(self, parser_unsupported_arch) -> None:
+        inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
+        r = _make_chat_request(
+            response_format=ResponseFormat(type="json_schema", json_schema=inner),
+        )
+        with pytest.raises(ValueError, match="does not support"):
+            parser_unsupported_arch.adjust_request(r)
+
+
+class TestAdjustRequestFoldFromResponseFormat:
+    @pytest.mark.parametrize(
+        "response_format, expected_schema",
+        [
+            pytest.param(
+                ResponseFormat(
+                    type="json_schema",
+                    json_schema=JsonSchemaResponseFormat(
+                        name="n", json_schema=SCHEMA_A
+                    ),
+                ),
+                SCHEMA_A,
+                id="json_schema_pydantic",
+            ),
+            pytest.param(
+                {
+                    "type": "json_schema",
+                    "json_schema": {"name": "n", "schema": SCHEMA_A},
+                },
+                SCHEMA_A,
+                id="json_schema_dict",
+            ),
+            pytest.param(
+                {"type": "json_object"},
+                {"type": "object"},
+                id="json_object",
+            ),
+        ],
+    )
+    def test_response_format_cleared(
+        self, parser, response_format, expected_schema
+    ) -> None:
+        r = _make_chat_request(response_format=response_format)
+        o = parser.adjust_request(r)
+        assert o.response_format is None
+        assert (
+            _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
+        )
+
+
+class TestHasEffectiveTools:
+    @pytest.mark.parametrize(
+        "tools, expected",
+        [
+            pytest.param(None, False, id="none"),
+            pytest.param([], False, id="empty_list"),
+            pytest.param("   ", False, id="blank_str"),
+            pytest.param(
+                [{"type": "function", "function": {"name": "f"}}],
+                True,
+                id="non_empty_list",
+            ),
+            pytest.param('{"x": 1}', True, id="non_empty_str"),
+        ],
+    )
+    def test_has_effective_tools(self, tools, expected) -> None:
+        assert _has_effective_tools(tools) is expected
+
+    def test_convert_schema_json_only_with_empty_tools_list(self) -> None:
+        tag = convert_schema_to_structural_tags(
+            schema=SCHEMA_B,
+            tools=[],
+            model_architecture="Cohere2ForCausalLM",
+        )
+        assert tag is not None
+        assert _first_json_schema(tag) == SCHEMA_B
+
+
+class TestAdjustRequestFoldFromStructuredOutputs:
+    @pytest.mark.parametrize(
+        "structured_outputs, expected_schema",
+        [
+            pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"),
+            pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"),
+            pytest.param(
+                {"json_object": True}, {"type": "object"}, id="json_object_flag"
+            ),
+            pytest.param(
+                StructuredOutputsParams(json=SCHEMA_B),
+                SCHEMA_B,
+                id="structured_outputs_dataclass",
+            ),
+            pytest.param(
+                {"json": {"name": "n", "schema": SCHEMA_A}},
+                SCHEMA_A,
+                id="openai_wrapper_dict_unwrapped",
+            ),
+        ],
+    )
+    def test_structured_outputs_folded(
+        self, parser, structured_outputs, expected_schema
+    ) -> None:
+        o = parser.adjust_request(
+            _make_chat_request(structured_outputs=structured_outputs),
+        )
+        assert (
+            _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
+        )
+
+    def test_responses_request_default_empty_tools(self, parser) -> None:
+        """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``."""
+        r = ResponsesRequest.model_validate(
+            {
+                "input": "hi",
+                "model": "m",
+                "structured_outputs": {"json": SCHEMA_B},
+            }
+        )
+        assert r.tools == []
+        o = parser.adjust_request(r)
+        assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B
+
+    def test_json_userdict_mapping_unwrapped(self) -> None:
+        inner = {"type": "object", "properties": {"u": {"type": "number"}}}
+        so = StructuredOutputsParams(json=UserDict(inner))
+        assert _schema_dict_from_structured_outputs(so) == inner
+
+    @pytest.mark.parametrize(
+        "json_value, match",
+        [
+            pytest.param("{not json}", "valid JSON", id="invalid_json_string"),
+            pytest.param(
+                json.dumps(["a", "b"]), "JSON object", id="non_object_json_string"
+            ),
+            pytest.param("   ", "empty", id="empty_json_string"),
+        ],
+    )
+    def test_structured_outputs_json_string_raises(
+        self, parser, json_value, match
+    ) -> None:
+        with pytest.raises(ValueError, match=match):
+            parser.adjust_request(
+                _make_chat_request(structured_outputs={"json": json_value}),
+            )
+
+    @pytest.mark.parametrize(
+        "construct",
+        [
+            pytest.param(
+                lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}),
+                id="chat_completion_request",
+            ),
+            pytest.param(
+                lambda: StructuredOutputsParams(json=[1, 2, 3]),  # type: ignore[arg-type]
+                id="structured_outputs_params",
+            ),
+        ],
+    )
+    def test_json_wrong_type_raises(self, construct) -> None:
+        """Non-str / non-dict ``json`` fails at Pydantic validation."""
+        with pytest.raises(ValidationError):
+            construct()
+
+
+class TestAdjustRequestPrecedence:
+    def test_response_format_over_structured_outputs_json(self, parser) -> None:
+        s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}}
+        s_so = {"type": "object", "properties": {"so": {"type": "number"}}}
+        inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf)
+        r = _make_chat_request(
+            response_format=ResponseFormat(type="json_schema", json_schema=inner),
+            structured_outputs={"json": s_so},
+        )
+        o = parser.adjust_request(r)
+        assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf
+
+
+class TestAdjustRequestTextPlusStructuredOutputs:
+    def test_text_response_format_preserved(self, parser) -> None:
+        sch = {"type": "object", "properties": {"k": {"type": "string"}}}
+        r = _make_chat_request(
+            response_format=ResponseFormat(type="text"),
+            structured_outputs={"json": sch},
+        )
+        o = parser.adjust_request(r)
+        assert o.response_format is not None
+        assert o.response_format.type == "text"
+        assert _first_json_schema(o.structured_outputs.structural_tag) == sch
+
+
+class TestAdjustRequestTools:
+    def test_tools_only_command_a_grammar(self, parser) -> None:
+        o = parser.adjust_request(
+            _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"),
+        )
+        assert "grammar" in _content_types(o.structured_outputs.structural_tag)
+
+    def test_tools_plus_json_schema_both_kinds(self, parser) -> None:
+        inner = JsonSchemaResponseFormat(
+            name="n",
+            json_schema={"type": "object", "properties": {"r": {"type": "string"}}},
+        )
+        r = _make_chat_request(
+            response_format=ResponseFormat(type="json_schema", json_schema=inner),
+            tools=[GET_WEATHER_TOOL],
+            tool_choice="auto",
+        )
+        o = parser.adjust_request(r)
+        types = _content_types(o.structured_outputs.structural_tag)
+        assert "grammar" in types
+        assert "json_schema" in types
diff --git a/tests/parser/cohere/test_tool_calls.py b/tests/parser/cohere/test_tool_calls.py
new file mode 100644
index 000000000000..a9fe7edb8db6
--- /dev/null
+++ b/tests/parser/cohere/test_tool_calls.py
@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tool-call parsing for the unified Cohere Command parser.
+
+The ``cohere_command3`` / ``cohere_command4`` tool parsers are registry shims;
+``CohereCommandParser`` (resolved through ``ParserManager``) does the parsing.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.parser.abstract_parser import Parser
+
+from .utils import MockCohereTokenizer, drive_parser, make_parser, token_deltas
+
+SPECIAL_TOKEN_MARKERS = (
+    "<|START_THINKING|>",
+    "<|END_THINKING|>",
+    "<|START_RESPONSE|>",
+    "<|END_RESPONSE|>",
+    "<|START_ACTION|>",
+    "<|END_ACTION|>",
+    "<|START_TEXT|>",
+    "<|END_TEXT|>",
+)
+
+DUMMY_TOOLS = [
+    {"type": "function", "function": {"name": name, "parameters": {"type": "object"}}}
+    for name in ("foo", "bar")
+]
+
+# (registry name, answer framing tag)
+VARIANTS = (("cohere_command3", "RESPONSE"), ("cohere_command4", "TEXT"))
+
+
+@dataclass
+class ExpectedToolCall:
+    id: str
+    name: str
+    arguments: dict
+
+
+@dataclass
+class ToolCallCase:
+    parser_name: str
+    model_output: str
+    expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list)
+    expected_reasoning: str | None = None
+    expected_content: str | None = None
+
+
+SINGLE_TOOL_CALL = """\
+<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
+[
+    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
+]
+<|END_ACTION|>"""
+MULTI_TOOL_CALL = """\
+<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|>
+[
+    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}},
+    {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}}
+]
+<|END_ACTION|>"""
+CITATIONS_NO_TOOL_CALL = (
+    "<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>\n"
+    "<|START_{tag}|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_{tag}|>"
+)
+
+# (case id, model output template, expectations)
+_CASE_TABLE = (
+    (
+        "single_tool_call",
+        SINGLE_TOOL_CALL,
+        dict(
+            expected_tool_calls=[ExpectedToolCall("0", "foo", {"query": "query1"})],
+            expected_reasoning="i will call foo with query1",
+        ),
+    ),
+    (
+        "citations_no_tool_calls",
+        CITATIONS_NO_TOOL_CALL,
+        dict(
+            expected_reasoning="This is a rainbow emoji: 🌈", expected_content="foo bar"
+        ),
+    ),
+    (
+        "multiple_tool_calls",
+        MULTI_TOOL_CALL,
+        dict(
+            expected_tool_calls=[
+                ExpectedToolCall("0", "foo", {"query": "query1"}),
+                ExpectedToolCall("1", "bar", {"x": 42}),
+            ],
+            expected_reasoning="first I think about foo",
+        ),
+    ),
+    (
+        "reasoning_only",
+        "<|START_THINKING|>just think, no response<|END_THINKING|>",
+        dict(expected_reasoning="just think, no response"),
+    ),
+)
+
+TOOL_CALL_CASES = [
+    pytest.param(
+        ToolCallCase(name, output.replace("{tag}", tag), **expected),
+        id=f"{name}-{case_id}",
+    )
+    for name, tag in VARIANTS
+    for case_id, output, expected in _CASE_TABLE
+]
+
+PARSER_NAMES = pytest.mark.parametrize("parser_name", [name for name, _ in VARIANTS])
+
+
+def _tool_parser(tokenizer: MockCohereTokenizer, parser_name: str) -> Parser:
+    return make_parser(tokenizer, parser_name, DUMMY_TOOLS)
+
+
+def _tool_request() -> ChatCompletionRequest:
+    return ChatCompletionRequest(
+        messages=[], model="test-model", tools=DUMMY_TOOLS, tool_choice="auto"
+    )
+
+
+@dataclass
+class StreamingResult:
+    tool_calls: dict[int, dict]
+    reasoning: str | None
+    content: str | None
+
+
+def run_streaming(
+    tokenizer: MockCohereTokenizer,
+    parser_name: str,
+    model_output: str,
+    chunk_size: int = 1,
+) -> StreamingResult:
+    deltas = drive_parser(
+        _tool_parser(tokenizer, parser_name),
+        _tool_request(),
+        token_deltas(tokenizer, model_output, chunk_size),
+    )
+    accumulated: dict[int, dict] = {}
+    for delta in deltas:
+        for tc in delta.tool_calls or []:
+            slot = accumulated.setdefault(
+                tc.index, {"id": "", "name": "", "arguments": ""}
+            )
+            if tc.id:
+                slot["id"] = tc.id
+            if tc.function and tc.function.name:
+                slot["name"] = tc.function.name
+            if tc.function and tc.function.arguments:
+                slot["arguments"] += tc.function.arguments
+    return StreamingResult(
+        tool_calls=accumulated,
+        reasoning="".join(d.reasoning for d in deltas if d.reasoning) or None,
+        content="".join(d.content for d in deltas if d.content) or None,
+    )
+
+
+def run_nonstreaming(
+    tokenizer: MockCohereTokenizer, parser_name: str, model_output: str
+) -> tuple[str | None, str | None, list]:
+    reasoning, content, tool_calls = _tool_parser(tokenizer, parser_name).parse(
+        model_output, _tool_request(), enable_auto_tools=True
+    )
+    return reasoning, content, tool_calls or []
+
+
+@pytest.mark.parametrize("case", TOOL_CALL_CASES)
+class TestExtractToolCalls:
+    def test_streaming(self, tokenizer: MockCohereTokenizer, case: ToolCallCase):
+        streamed = run_streaming(tokenizer, case.parser_name, case.model_output)
+        assert len(streamed.tool_calls) == len(case.expected_tool_calls)
+        for i, expected in enumerate(case.expected_tool_calls):
+            tc = streamed.tool_calls[i]
+            assert tc["id"] == expected.id
+            assert tc["name"] == expected.name
+            assert json.loads(tc["arguments"]) == expected.arguments
+
+    def test_streaming_reasoning(
+        self, tokenizer: MockCohereTokenizer, case: ToolCallCase
+    ):
+        streamed = run_streaming(tokenizer, case.parser_name, case.model_output)
+        assert streamed.reasoning == case.expected_reasoning
+
+    def test_streaming_content(
+        self, tokenizer: MockCohereTokenizer, case: ToolCallCase
+    ):
+        streamed = run_streaming(tokenizer, case.parser_name, case.model_output)
+        assert streamed.content == case.expected_content
+
+    def test_nonstreaming(self, tokenizer: MockCohereTokenizer, case: ToolCallCase):
+        reasoning, content, tool_calls = run_nonstreaming(
+            tokenizer, case.parser_name, case.model_output
+        )
+        assert reasoning == case.expected_reasoning
+        assert content == case.expected_content
+        assert len(tool_calls) == len(case.expected_tool_calls)
+        for actual, expected in zip(tool_calls, case.expected_tool_calls):
+            assert actual.name == expected.name
+            assert json.loads(actual.arguments) == expected.arguments
+
+    def test_streaming_nonstreaming_agree(
+        self, tokenizer: MockCohereTokenizer, case: ToolCallCase
+    ):
+        streamed = run_streaming(tokenizer, case.parser_name, case.model_output)
+        _, _, tool_calls = run_nonstreaming(
+            tokenizer, case.parser_name, case.model_output
+        )
+        assert len(streamed.tool_calls) == len(tool_calls)
+        for i, actual in enumerate(tool_calls):
+            assert streamed.tool_calls[i]["name"] == actual.name
+            assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads(
+                actual.arguments
+            )
+
+
+class TestSpeculativeDecodingMultiTokenDelta:
+    @PARSER_NAMES
+    @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6])
+    def test_no_special_token_leak_in_streaming_deltas(
+        self, tokenizer: MockCohereTokenizer, parser_name: str, chunk_size: int
+    ):
+        deltas = drive_parser(
+            _tool_parser(tokenizer, parser_name),
+            _tool_request(),
+            token_deltas(tokenizer, SINGLE_TOOL_CALL, chunk_size),
+        )
+        for delta in deltas:
+            fields: list[tuple[str, str | None]] = [
+                ("reasoning", delta.reasoning),
+                ("content", delta.content),
+            ]
+            for tc in delta.tool_calls or []:
+                if tc.function:
+                    fields.append(("tool_call.name", tc.function.name))
+                    fields.append(("tool_call.arguments", tc.function.arguments))
+            for marker in SPECIAL_TOKEN_MARKERS:
+                for field_name, value in fields:
+                    assert value is None or marker not in value, (
+                        f"special token {marker!r} leaked into {field_name} "
+                        f"with chunk_size={chunk_size} delta={delta!r}"
+                    )
+
+    @PARSER_NAMES
+    @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6])
+    def test_multi_token_chunks_still_produce_correct_tool_call(
+        self, tokenizer: MockCohereTokenizer, parser_name: str, chunk_size: int
+    ):
+        streamed = run_streaming(tokenizer, parser_name, SINGLE_TOOL_CALL, chunk_size)
+        assert len(streamed.tool_calls) == 1
+        tc = streamed.tool_calls[0]
+        assert tc["id"] == "0"
+        assert tc["name"] == "foo"
+        assert json.loads(tc["arguments"]) == {"query": "query1"}
+
+
+class TestStreamingDeltaShape:
+    @PARSER_NAMES
+    def test_reasoning_and_tool_calls_are_separate_deltas(
+        self, tokenizer: MockCohereTokenizer, parser_name: str
+    ):
+        deltas = drive_parser(
+            _tool_parser(tokenizer, parser_name),
+            _tool_request(),
+            token_deltas(tokenizer, SINGLE_TOOL_CALL),
+        )
+        for delta in deltas:
+            populated = [
+                delta.content is not None,
+                delta.reasoning is not None,
+                bool(delta.tool_calls),
+            ]
+            assert sum(populated) == 1, (
+                "A single streaming delta must carry exactly one of "
+                f"content/reasoning/tool_calls, got {delta!r}"
+            )
+        assert any(d.reasoning is not None for d in deltas)
+        assert any(d.tool_calls for d in deltas)
diff --git a/tests/parser/cohere/utils.py b/tests/parser/cohere/utils.py
new file mode 100644
index 000000000000..f26df5296c6d
--- /dev/null
+++ b/tests/parser/cohere/utils.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared helpers for the unified Cohere Command parser tests."""
+
+from __future__ import annotations
+
+import itertools
+
+import regex as re
+
+from vllm.entrypoints.generate.base.protocol import DeltaMessage
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.parser import ParserManager
+from vllm.parser.abstract_parser import Parser
+
+_SPECIAL_TOKEN_RE = re.compile(r"(<\|[A-Z_]+\|>)")
+REPLACEMENT_CHAR = "�"
+
+
+class MockCohereTokenizer:
+    """Byte-level stand-in for the Cohere tokenizer.
+
+    Text round-trips through UTF-8 bytes so a multi-byte character can split
+    across "tokens" (reproducing trailing U+FFFD buffering), while each
+    ``<|...|>`` marker is a single id so token-id reasoning-end gating behaves
+    as with the real tokenizer.
+    """
+
+    _SPECIAL_TOKEN_IDS = {
+        tok: 256 + i
+        for i, tok in enumerate(
+            (
+                "<|START_THINKING|>",
+                "<|END_THINKING|>",
+                "<|CHATBOT_TOKEN|>",
+                "<|START_RESPONSE|>",
+                "<|END_RESPONSE|>",
+                "<|START_TEXT|>",
+                "<|END_TEXT|>",
+                "<|START_ACTION|>",
+                "<|END_ACTION|>",
+            )
+        )
+    }
+    _ID_TO_SPECIAL_TOKEN = {v: k for k, v in _SPECIAL_TOKEN_IDS.items()}
+
+    def convert_tokens_to_ids(self, token: str) -> int:
+        return self._SPECIAL_TOKEN_IDS.get(token, 0)
+
+    def get_vocab(self) -> dict[str, int]:
+        return {}
+
+    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+        ids: list[int] = []
+        for part in _SPECIAL_TOKEN_RE.split(text):
+            if part in self._SPECIAL_TOKEN_IDS:
+                ids.append(self._SPECIAL_TOKEN_IDS[part])
+            else:
+                ids.extend(part.encode("utf-8"))
+        return ids
+
+    def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
+        out: list[str] = []
+        for special, run in itertools.groupby(
+            ids, self._ID_TO_SPECIAL_TOKEN.__contains__
+        ):
+            if not special:
+                out.append(bytes(run).decode("utf-8", errors="replace"))
+            elif not skip_special_tokens:
+                out.extend(self._ID_TO_SPECIAL_TOKEN[i] for i in run)
+        return "".join(out)
+
+
+def token_deltas(
+    tokenizer: MockCohereTokenizer, text: str, chunk_size: int = 1
+) -> list[tuple[str, list[int]]]:
+    """Split ``text`` into ``(delta_text, delta_token_ids)`` steps of
+    ``chunk_size`` tokens, buffering incomplete multi-byte sequences (trailing
+    U+FFFD) into the next step as real streaming does."""
+    ids = tokenizer.encode(text)
+    deltas: list[tuple[str, list[int]]] = []
+    prev = ""
+    pending: list[int] = []
+    for start in range(0, len(ids), chunk_size):
+        end = start + chunk_size
+        pending.extend(ids[start:end])
+        current = tokenizer.decode(ids[:end])
+        if current.endswith(REPLACEMENT_CHAR) and end < len(ids):
+            continue
+        deltas.append((current[len(prev) :], pending))
+        prev, pending = current, []
+    return deltas
+
+
+def make_parser(
+    tokenizer: MockCohereTokenizer,
+    name: str,
+    tools: list[dict] | None = None,
+    chat_template_kwargs: dict | None = None,
+    model_config=None,
+) -> Parser:
+    """Resolve the unified Cohere parser through ``ParserManager`` by registry
+    name (``cohere_command3`` / ``cohere_command4``)."""
+    cls = ParserManager.get_parser(name, name, enable_auto_tools=True)
+    assert cls is not None
+    return cls(
+        tokenizer,
+        tools,
+        model_config=model_config,
+        chat_template_kwargs=chat_template_kwargs or {},
+    )
+
+
+def drive_parser(
+    parser: Parser,
+    request: ChatCompletionRequest,
+    deltas: list[tuple[str, list[int]]],
+) -> list[DeltaMessage]:
+    """Feed ``deltas`` through ``parse_delta`` (flushing on the last one) and
+    return the non-``None`` messages in order."""
+    out: list[DeltaMessage] = []
+    for i, (text, ids) in enumerate(deltas):
+        delta = parser.parse_delta(text, ids, request, finished=i == len(deltas) - 1)
+        if delta is not None:
+            out.append(delta)
+    return out
+
+
+def stream_parser(parser, request, tokenizer, text: str, chunk_size: int = 1):
+    """Stream ``text`` through ``parser`` in ``chunk_size``-token deltas."""
+    return drive_parser(parser, request, token_deltas(tokenizer, text, chunk_size))
diff --git a/tests/parser/engine/conftest.py b/tests/parser/engine/conftest.py
index 2522d146c895..1c482434eb5f 100644
--- a/tests/parser/engine/conftest.py
+++ b/tests/parser/engine/conftest.py
@@ -30,6 +30,7 @@ def make_mock_tokenizer(
     The returned mock supports get_vocab(), encode(), and decode().
     decode() maps known token IDs back to their text and falls back to
     chr(id) for ASCII IDs or ``<id>`` for others.
+
     """
     id_to_text = {v: k for k, v in vocab.items()}
     tokenizer = MagicMock()
diff --git a/tests/parser/engine/replay_harness.py b/tests/parser/engine/replay_harness.py
index 824a1ed86d88..5eec8d213a02 100644
--- a/tests/parser/engine/replay_harness.py
+++ b/tests/parser/engine/replay_harness.py
@@ -183,6 +183,7 @@ def replay_streaming(
 
     Returns:
         List of ``DeltaMessage`` results from each ``parse_delta()`` call.
+
     """
     if chunk_size is None:
         chunk_size = len(tokens)
diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py
index 1c1029b51a01..2cf3a8bb85f3 100644
--- a/tests/parser/engine/test_deepseek_v4.py
+++ b/tests/parser/engine/test_deepseek_v4.py
@@ -119,6 +119,35 @@ def test_empty_args(self):
         result = json.loads(_dsml_arg_converter("", partial=False))
         assert result == {}
 
+    def _bare(self, name: str, value: str) -> str:
+        return f'<｜DSML｜parameter name="{name}">{value}{_PARAM_CLOSE}'
+
+    def test_missing_string_attribute_keeps_literal(self):
+        """The model sometimes omits ``string``; the parameter must survive."""
+        raw = self._bare("name", "alpha")
+        assert json.loads(_dsml_arg_converter(raw, partial=False)) == {"name": "alpha"}
+
+    def test_missing_string_attribute_parses_json(self):
+        raw = self._bare("count", "7") + self._bare("opts", '{"k": [1]}')
+        assert json.loads(_dsml_arg_converter(raw, partial=False)) == {
+            "count": 7,
+            "opts": {"k": [1]},
+        }
+
+    def test_missing_string_attribute_alongside_annotated(self):
+        raw = self._bare("a", "x") + _param("b", "true", "42")
+        assert json.loads(_dsml_arg_converter(raw, partial=False)) == {
+            "a": "x",
+            "b": "42",
+        }
+
+    def test_missing_string_attribute_partial_value(self):
+        raw = _param("a", "true", "x") + '<｜DSML｜parameter name="n">12'
+        assert json.loads(_dsml_arg_converter(raw, partial=True)) == {
+            "a": "x",
+            "n": 12,
+        }
+
     def test_invalid_json_fallback(self):
         raw = self._raw(("data", "false", "[broken"))
         result = json.loads(_dsml_arg_converter(raw, partial=False))
@@ -232,6 +261,47 @@ def test_streaming_split_next_parameter_tag_is_buffered(
         }
 
 
+class TestMissingStringAttribute:
+    def test_non_streaming_keeps_parameter(self, mock_tokenizer, mock_request):
+        text = (
+            f"{DSML_TOOL_START}"
+            f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n"
+            f'<｜DSML｜parameter name="location">Paris{_PARAM_CLOSE}\n'
+            f"{_param('date', 'true', 'tomorrow')}\n"
+            f"{DSML_INVOKE_END}{DSML_TOOL_END}"
+        )
+
+        result = DeepSeekV4Parser(mock_tokenizer).extract_tool_calls(text, mock_request)
+
+        assert result.tools_called is True
+        assert json.loads(result.tool_calls[0].function.arguments) == {
+            "location": "Paris",
+            "date": "tomorrow",
+        }
+
+    def test_streaming_keeps_parameter(self, mock_tokenizer, mock_request):
+        chunks = [
+            DSML_TOOL_START,
+            f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n",
+            '<｜DSML｜parameter name="loc',
+            'ation">Par',
+            "is",
+            _PARAM_CLOSE,
+            f"\n{_param('date', 'true', 'tomorrow')}\n",
+            DSML_INVOKE_END,
+            DSML_TOOL_END,
+        ]
+
+        results = simulate_tool_streaming(
+            DeepSeekV4Parser(mock_tokenizer), mock_request, chunks
+        )
+
+        assert json.loads(collect_tool_arguments(results)) == {
+            "location": "Paris",
+            "date": "tomorrow",
+        }
+
+
 # ── Bare </think> absorption and duplicate <think> absorption ─────────
 
 
@@ -574,7 +644,7 @@ def test_reasoning_end_token_ids_drop_multi_token_tool_opener(
         # the text lexer.
         assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID}
         assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1
-        assert thinking_parser.find_reasoning_end_offset([7, 8]) is None
+        assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2
 
     def test_streaming_reasoning_implicit_end(self, thinking_parser):
         chunks = [
diff --git a/tests/parser/engine/test_deepseek_v41.py b/tests/parser/engine/test_deepseek_v41.py
index cc7d6e0aedf4..d99ca6b2a17b 100644
--- a/tests/parser/engine/test_deepseek_v41.py
+++ b/tests/parser/engine/test_deepseek_v41.py
@@ -137,10 +137,12 @@ def test_python_argument_conversion_and_partial_values():
         '<｜DSML｜ parameter name="object" string="false">'
         '{"a": [true, null]}</｜DSML｜ parameter>'
         '<｜DSML｜ parameter name="bad" string="false">[broken</｜DSML｜ parameter>'
+        '<｜DSML｜ parameter name="bare">7</｜DSML｜ parameter>'
         '<｜DSML｜ parameter name="text" string="true">  a<b'
     )
     assert json.loads(converter(raw, True)) == {
         "object": {"a": [True, None]},
         "bad": "[broken",
+        "bare": 7,
         "text": "  a<b",
     }
diff --git a/tests/parser/engine/test_delegating_replay.py b/tests/parser/engine/test_delegating_replay.py
index 2bc4338785cb..5cd48dd870f2 100644
--- a/tests/parser/engine/test_delegating_replay.py
+++ b/tests/parser/engine/test_delegating_replay.py
@@ -99,6 +99,10 @@ def _discover_pairings() -> list[_PairingInfo]:
             # terminal matching; combined-parser replay coverage lives in
             # test_inkling.py.
             continue
+        if cfg.name == "granite":
+            # Granite has a JSON-array tool body with no TOOL_END terminal;
+            # its replay coverage lives in test_granite.py.
+            continue
 
         parser_cls = type(
             f"_Delegating{engine_cls.__name__}",
diff --git a/tests/parser/engine/test_gemma4_streaming_reasoning.py b/tests/parser/engine/test_gemma4_streaming_reasoning.py
index 74224c15b749..ae21d69fc303 100644
--- a/tests/parser/engine/test_gemma4_streaming_reasoning.py
+++ b/tests/parser/engine/test_gemma4_streaming_reasoning.py
@@ -1616,3 +1616,36 @@ def test_non_streaming_multiple_commas(self, multi_comma_parser, request_obj):
         assert result.tools_called is True
         args = json.loads(result.tool_calls[0].function.arguments)
         assert args["destination"] == "456 Oakwood Avenue, Rivermist, 83214"
+
+
+class TestGemma4IsReasoningEnd:
+    @staticmethod
+    def _parser(thinking: bool):
+        return Gemma4Parser(
+            _make_tokenizer(_PLAIN_ANSWER_TOKENS),
+            chat_template_kwargs={"enable_thinking": thinking},
+        )
+
+    @pytest.mark.parametrize(
+        ("ids", "thinking", "ended"),
+        [
+            ([CHANNEL_START_ID, 3000, CHANNEL_END_ID], True, True),
+            ([CHANNEL_START_ID, 3000, TOOL_CALL_START_ID], True, True),
+            ([CHANNEL_START_ID, 3000], True, False),
+            ([NEW_TURN_ID, 9100], True, False),
+            ([CHANNEL_END_ID, NEW_TURN_ID, 9100], True, False),
+            ([NEW_TURN_ID, 9100], False, True),
+            # The bundled template's thinking-off prompt: closed channel.
+            ([NEW_TURN_ID, 9100, CHANNEL_START_ID, 3000, CHANNEL_END_ID], False, True),
+            # A marker still wins with thinking off: the grammar honours it.
+            ([CHANNEL_START_ID, 3000], False, False),
+            # A tool-call opener after the channel closed: still ended.
+            (
+                [CHANNEL_START_ID, 3000, CHANNEL_END_ID, TOOL_CALL_START_ID, 9200],
+                True,
+                True,
+            ),
+        ],
+    )
+    def test_is_reasoning_end(self, ids, thinking, ended):
+        assert self._parser(thinking).is_reasoning_end(ids) is ended
diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py
new file mode 100644
index 000000000000..58fe7aa28cbd
--- /dev/null
+++ b/tests/parser/engine/test_glm47_moe.py
@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reasoning-boundary tests for the GLM-4.7 parser.
+
+GLM can transition straight from reasoning into a tool call without ever
+emitting ``</think>``, so ``<tool_call>`` is an implicit reasoning
+terminator. These cover that path plus the multi-turn prompts where an
+earlier turn's markers must not be read as the current turn's state.
+"""
+
+import pytest
+
+from tests.parser.engine.conftest import make_mock_tokenizer
+from vllm.parser.glm47_moe import (
+    THINK_END,
+    THINK_START,
+    TOOL_CALL_END,
+    TOOL_CALL_START,
+    Glm47MoeParser,
+)
+
+THINK_S, THINK_E = 1000, 1001
+TOOL_S, TOOL_E = 1002, 1003
+ASSISTANT, OBSERVATION, USER = 1004, 1005, 1006
+TEXT = 42  # stand-in for an ordinary reasoning/content token
+
+VOCAB = {
+    THINK_START: THINK_S,
+    THINK_END: THINK_E,
+    TOOL_CALL_START: TOOL_S,
+    TOOL_CALL_END: TOOL_E,
+    "<|assistant|>": ASSISTANT,
+    "<|observation|>": OBSERVATION,
+    "<|user|>": USER,
+}
+
+
+@pytest.fixture
+def parser():
+    return Glm47MoeParser(make_mock_tokenizer(VOCAB))
+
+
+@pytest.fixture
+def no_thinking_parser():
+    return Glm47MoeParser(
+        make_mock_tokenizer(VOCAB),
+        chat_template_kwargs={"thinking": False},
+    )
+
+
+class TestIsReasoningEnd:
+    def test_open_reasoning(self, parser):
+        assert not parser.is_reasoning_end([THINK_S, TEXT])
+
+    def test_think_end(self, parser):
+        assert parser.is_reasoning_end([THINK_S, TEXT, THINK_E, TEXT])
+
+    def test_tool_call_without_think_end(self, parser):
+        """Reasoning -> tool call with no ``</think>`` still ends reasoning."""
+        assert parser.is_reasoning_end([THINK_S, TEXT, TOOL_S, TEXT])
+
+    def test_previous_turn_tool_call_ignored(self, parser):
+        """A finished tool call from an earlier turn says nothing about the
+        turn currently being generated."""
+        prompt = [THINK_S, TEXT, TOOL_S, TEXT, TOOL_E, OBSERVATION, TEXT, ASSISTANT]
+        assert not parser.is_reasoning_end(prompt)
+
+    def test_previous_turn_think_end_ignored(self, parser):
+        prompt = [THINK_S, TEXT, THINK_E, TEXT, USER, TEXT, ASSISTANT]
+        assert not parser.is_reasoning_end(prompt)
+
+    def test_empty_input(self, parser):
+        assert not parser.is_reasoning_end([])
+
+    def test_thinking_disabled(self, no_thinking_parser):
+        assert no_thinking_parser.is_reasoning_end([THINK_S, TEXT])
+
+
+class TestExtractContentIds:
+    def test_think_end_wins_over_later_tool_call(self, parser):
+        """Content between ``</think>`` and a tool call must survive."""
+        ids = [THINK_S, TEXT, THINK_E, 20, 21, TOOL_S, 30, TOOL_E]
+        assert parser.extract_content_ids(ids) == [20, 21, TOOL_S, 30, TOOL_E]
+
+    def test_every_tool_call_after_think_end_kept(self, parser):
+        ids = [THINK_E, 20, TOOL_S, 30, TOOL_E, 21, TOOL_S, 31, TOOL_E]
+        assert parser.extract_content_ids(ids) == ids[1:]
+
+    def test_falls_back_to_tool_call(self, parser):
+        """Without ``</think>``, content starts at the opener itself so the
+        tool parser still receives a well-formed call."""
+        ids = [THINK_S, TEXT, TOOL_S, 30, TOOL_E]
+        assert parser.extract_content_ids(ids) == [TOOL_S, 30, TOOL_E]
+
+    def test_falls_back_to_first_tool_call_of_turn(self, parser):
+        ids = [THINK_S, TEXT, TOOL_S, 30, TOOL_E, TOOL_S, 31, TOOL_E]
+        assert parser.extract_content_ids(ids) == ids[2:]
+
+    def test_previous_turn_tool_call_ignored(self, parser):
+        ids = [TOOL_S, 30, TOOL_E, OBSERVATION, TEXT, ASSISTANT, THINK_S, TEXT]
+        assert parser.extract_content_ids(ids) == ids
+
+    def test_no_markers_returns_input_ids(self, parser):
+        assert parser.extract_content_ids([20, 21]) == [20, 21]
+
+    def test_thinking_disabled(self, no_thinking_parser):
+        ids = [THINK_S, TEXT, TOOL_S]
+        assert no_thinking_parser.extract_content_ids(ids) == ids
diff --git a/tests/parser/engine/test_granite.py b/tests/parser/engine/test_granite.py
new file mode 100644
index 000000000000..ddc7e2bfe0cd
--- /dev/null
+++ b/tests/parser/engine/test_granite.py
@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the engine-based Granite parser.
+
+Granite emits tool calls as a JSON array following a marker
+(``<|tool_call|>`` for 3.0, ``<tool_call>`` for 3.1) with no closing
+terminal; the ``]`` ends the region and surrounding prose is content. The
+engine's ``tool_call_body_array`` mode splits the array into one call per
+element and ``_granite_arg_converter`` carves each element's ``arguments``.
+"""
+
+import json
+
+import pytest
+
+from tests.parser.engine.conftest import make_mock_tokenizer
+from tests.parser.engine.trace_builder import _GRANITE_SCENARIOS, _build_granite
+from vllm.parser.granite import GraniteParser, _granite_arg_converter
+
+TOOL_TOKEN = "<|tool_call|>"
+TOOL_STRING = "<tool_call>"
+
+_GRANITE_VOCAB = {TOOL_TOKEN: 49154}
+
+
+@pytest.fixture
+def mock_tokenizer():
+    return make_mock_tokenizer(_GRANITE_VOCAB)
+
+
+@pytest.fixture
+def parser(mock_tokenizer):
+    return GraniteParser(mock_tokenizer)
+
+
+def _tokenize(text: str) -> list[tuple[int, str]]:
+    """Tokenize like the real stream: ``<|tool_call|>`` is one special token,
+    plain text becomes one token per character."""
+    tokens: list[tuple[int, str]] = []
+    i = 0
+    while i < len(text):
+        if text.startswith(TOOL_TOKEN, i):
+            tokens.append((_GRANITE_VOCAB[TOOL_TOKEN], TOOL_TOKEN))
+            i += len(TOOL_TOKEN)
+        else:
+            tokens.append((ord(text[i]), text[i]))
+            i += 1
+    return tokens
+
+
+def _stream(parser, request, text: str, chunk_size: int):
+    tokens = _tokenize(text)
+    results = []
+    previous_text = ""
+    previous_token_ids: list[int] = []
+    for start in range(0, len(tokens), chunk_size):
+        batch = tokens[start : start + chunk_size]
+        delta_text = "".join(t for _, t in batch)
+        delta_token_ids = [tid for tid, _ in batch]
+        current_text = previous_text + delta_text
+        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=delta_text,
+            previous_token_ids=tuple(previous_token_ids),
+            current_token_ids=tuple(current_token_ids),
+            delta_token_ids=tuple(delta_token_ids),
+            request=request,
+        )
+        results.append(delta)
+        previous_text = current_text
+        previous_token_ids = current_token_ids
+    finish = parser.finish_streaming()
+    if finish is not None:
+        results.append(finish)
+    return results
+
+
+def _collect_content(results) -> str:
+    return "".join(d.content for d in results if d and d.content)
+
+
+def _collect_names(results) -> list[str]:
+    names: dict[int, str] = {}
+    for d in results:
+        if not (d and d.tool_calls):
+            continue
+        for tc in d.tool_calls:
+            if tc.function and tc.function.name:
+                names[tc.index] = names.get(tc.index, "") + tc.function.name
+    return [names[i] for i in sorted(names)]
+
+
+def _collect_args(results) -> dict[int, str]:
+    args: dict[int, str] = {}
+    for d in results:
+        if not (d and d.tool_calls):
+            continue
+        for tc in d.tool_calls:
+            if tc.function and tc.function.arguments:
+                args[tc.index] = args.get(tc.index, "") + tc.function.arguments
+    return args
+
+
+class TestArgConverter:
+    def test_complete_wrapper(self):
+        raw = '{"name": "get_weather", "arguments": {"city": "SF"}}'
+        assert _granite_arg_converter(raw, False) == '{"city": "SF"}'
+
+    def test_arguments_before_name(self):
+        raw = '{"arguments": {"city": "SF"}, "name": "get_weather"}'
+        assert _granite_arg_converter(raw, False) == '{"city": "SF"}'
+
+    def test_partial_before_arguments(self):
+        assert _granite_arg_converter('{"name": "get_w', True) == ""
+
+    def test_partial_inside_arguments(self):
+        raw = '{"name": "x", "arguments": {"a": 1'
+        assert _granite_arg_converter(raw, True) == '{"a": 1'
+
+    def test_prefix_stability(self):
+        # Each growing prefix must extend the previous converter output.
+        full = '{"name": "x", "arguments": {"city": "Tokyo"}}'
+        prev = ""
+        for i in range(1, len(full) + 1):
+            out = _granite_arg_converter(full[:i], True)
+            assert out.startswith(prev) or prev.startswith(out) or out == ""
+            if out:
+                prev = out
+
+    def test_missing_arguments_defaults(self):
+        assert _granite_arg_converter('{"name": "x"}', False) == "{}"
+        assert _granite_arg_converter('{"name": "x"}', True) == ""
+
+    def test_non_object_arguments_rejected(self):
+        with pytest.raises(ValueError, match="JSON object"):
+            _granite_arg_converter('{"name": "x", "arguments": [1]}', False)
+
+
+class TestNonStreaming:
+    def test_plain_text(self, parser, mock_request):
+        reasoning, content, tools = parser.parse("Just a reply.", mock_request)
+        assert reasoning is None
+        assert content == "Just a reply."
+        assert tools is None
+
+    def test_single_tool_token_marker(self, parser, mock_request):
+        text = (
+            f'{TOOL_TOKEN} [{{"name": "get_weather", "arguments": {{"city": "SF"}}}}]'
+        )
+        _, content, tools = parser.parse(text, mock_request)
+        assert content is None
+        assert [t.name for t in tools] == ["get_weather"]
+        assert json.loads(tools[0].arguments) == {"city": "SF"}
+
+    def test_single_tool_string_marker(self, parser, mock_request):
+        # Granite 3.1 uses the plain-text ``<tool_call>`` marker.
+        text = f'{TOOL_STRING} [{{"name": "get_time", "arguments": {{}}}}]'
+        _, content, tools = parser.parse(text, mock_request)
+        assert [t.name for t in tools] == ["get_time"]
+        assert json.loads(tools[0].arguments) == {}
+
+    def test_parallel_calls_one_array(self, parser, mock_request):
+        text = (
+            f'{TOOL_TOKEN} [{{"name": "a", "arguments": {{"x": 1}}}}, '
+            f'{{"name": "b", "arguments": {{"y": [1, 2]}}}}]'
+        )
+        _, _, tools = parser.parse(text, mock_request)
+        assert [t.name for t in tools] == ["a", "b"]
+        assert json.loads(tools[0].arguments) == {"x": 1}
+        assert json.loads(tools[1].arguments) == {"y": [1, 2]}
+
+    def test_surrounding_text_is_content(self, parser, mock_request):
+        text = (
+            "Let me check.\n"
+            f'{TOOL_TOKEN} [{{"name": "get_weather", "arguments": {{"city": "SF"}}}}]'
+        )
+        _, content, tools = parser.parse(text, mock_request)
+        assert content == "Let me check."
+        assert [t.name for t in tools] == ["get_weather"]
+
+    def test_arguments_before_name(self, parser, mock_request):
+        text = (
+            f'{TOOL_TOKEN} [{{"arguments": {{"city": "SF"}}, "name": "get_weather"}}]'
+        )
+        _, _, tools = parser.parse(text, mock_request)
+        assert [t.name for t in tools] == ["get_weather"]
+        assert json.loads(tools[0].arguments) == {"city": "SF"}
+
+    def test_marker_without_array_is_not_a_tool_call(self, parser, mock_request):
+        # A non-array body after the marker (malformed) yields no tool call.
+        text = f'{TOOL_TOKEN} {{"name": "func", "arguments": {{}}}}'
+        _, _, tools = parser.parse(text, mock_request)
+        assert not tools
+
+
+class TestStreaming:
+    @pytest.mark.parametrize("chunk_size", [1, 2, 3, 5, 11])
+    def test_parallel_calls_chunk_invariance(self, parser, mock_request, chunk_size):
+        text = (
+            f'{TOOL_TOKEN} [{{"name": "get_weather", '
+            f'"arguments": {{"city": "Tokyo"}}}}, '
+            f'{{"name": "get_time", "arguments": {{"timezone": "Asia/Tokyo"}}}}]'
+        )
+        results = _stream(parser, mock_request, text, chunk_size)
+        assert _collect_names(results) == ["get_weather", "get_time"]
+        args = _collect_args(results)
+        assert json.loads(args[0]) == {"city": "Tokyo"}
+        assert json.loads(args[1]) == {"timezone": "Asia/Tokyo"}
+
+    @pytest.mark.parametrize("chunk_size", [1, 3, 7])
+    def test_surrounding_text_streaming(self, parser, mock_request, chunk_size):
+        text = (
+            "Let me check.\n"
+            f'{TOOL_TOKEN} [{{"name": "get_weather", "arguments": {{"city": "SF"}}}}]'
+        )
+        results = _stream(parser, mock_request, text, chunk_size)
+        assert _collect_names(results) == ["get_weather"]
+        assert "Let me check." in _collect_content(results)
+
+    def test_incomplete_array_at_eos(self, parser, mock_request):
+        # Truncated mid-element must not raise and must not leak the marker.
+        text = f'{TOOL_TOKEN} [{{"name": "func", "arguments": {{"a": 1'
+        results = _stream(parser, mock_request, text, 3)
+        assert TOOL_TOKEN not in _collect_content(results)
+
+
+class TestReplay:
+    @pytest.mark.parametrize("scenario", _GRANITE_SCENARIOS, ids=lambda s: s.id)
+    def test_replay_scenarios(self, scenario):
+        # ``_build_granite`` replays the sample through GraniteParser at
+        # chunk_size=1 and asserts reasoning/content/tool_calls (validate=True).
+        _build_granite(scenario)
diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py
index f27ce52cb37b..8fd7c4d2623f 100644
--- a/tests/parser/engine/test_parser_engine.py
+++ b/tests/parser/engine/test_parser_engine.py
@@ -185,8 +185,8 @@ def test_config_without_reasoning_has_empty_set(self):
     def test_find_reasoning_end_offset_returns_first_match(self):
         engine = _make_engine()
         assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1
-        assert engine.find_reasoning_end_offset([5, 6]) is None
-        assert engine.find_reasoning_end_offset([]) is None
+        assert engine.find_reasoning_end_offset([5, 6]) == 2
+        assert engine.find_reasoning_end_offset([]) == 0
         # Rejected-draft placeholders never match.
         assert engine.find_reasoning_end_offset([-1, 201]) == 1
 
diff --git a/tests/parser/engine/test_qwen3_reasoning.py b/tests/parser/engine/test_qwen3_reasoning.py
index 37e6b5d13090..3f285e4d0315 100644
--- a/tests/parser/engine/test_qwen3_reasoning.py
+++ b/tests/parser/engine/test_qwen3_reasoning.py
@@ -180,9 +180,8 @@ def test_prompt_tool_example_before_generation_think_not_end(self, parser):
         """Tool examples before the generation <think> must not end reasoning."""
         assert not parser.is_reasoning_end([_TOOL_CALL_ID, _TEXT_ID, _THINK_START_ID])
 
-    def test_paired_tool_call_not_end(self, parser):
-        """Paired <tool_call>...</tool_call> (from template) is NOT end."""
-        assert not parser.is_reasoning_end(
+    def test_completed_tool_call_after_open_think_is_end(self, parser):
+        assert parser.is_reasoning_end(
             [_THINK_START_ID, 1, _TOOL_CALL_ID, 2, _TOOL_CALL_END_ID]
         )
 
@@ -196,6 +195,43 @@ def test_empty_ids(self, parser):
         assert not parser.is_reasoning_end([])
 
 
+class TestIsReasoningEndDerivedFromGrammar:
+    @pytest.mark.parametrize(
+        ("ids", "thinking", "ended"),
+        [
+            # Explicit close, and the implicit close on <tool_call>.
+            ([_THINK_START_ID, 1, _THINK_END_ID], True, True),
+            ([_THINK_START_ID, 1, _TOOL_CALL_ID], True, True),
+            # Still inside the block.
+            ([_THINK_START_ID, 1], True, False),
+            ([_THINK_END_ID, _THINK_START_ID, 1], True, False),
+            # Uncommitted: a bare turn, an empty sequence. Thinking on starts
+            # in REASONING, so reasoning is expected.
+            ([_IM_END_ID, _IM_START_ID, 1], True, False),
+            ([], True, False),
+            # Thinking off starts in CONTENT with no way into REASONING, so
+            # a <think> the template left behind is inert, and a silent tail
+            # is ended because no block is expected.
+            ([_THINK_START_ID, 1], False, True),
+            ([_THINK_END_ID, _THINK_START_ID, 1], False, True),
+            ([_IM_END_ID, _IM_START_ID, 1], False, True),
+            ([], False, True),
+        ],
+    )
+    def test_is_reasoning_end(self, mock_tokenizer, ids, thinking, ended):
+        parser = Qwen3Parser(
+            mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking}
+        )
+        assert parser.is_reasoning_end(ids) is ended
+
+    @pytest.mark.parametrize("thinking", [True, False])
+    def test_wait_for_reasoning_follows_thinking(self, mock_tokenizer, thinking):
+        parser = Qwen3Parser(
+            mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking}
+        )
+        assert parser.parser_engine_config.wait_for_reasoning is thinking
+
+
 class TestIsReasoningEndTurnBoundaries:
     """Reasoning markers from earlier turns must not mark the new turn's
     reasoning as finished (think-marker injection / interleaved-thinking
diff --git a/tests/parser/engine/test_replay.py b/tests/parser/engine/test_replay.py
index fee74f8d2cca..fbb911d7aea1 100644
--- a/tests/parser/engine/test_replay.py
+++ b/tests/parser/engine/test_replay.py
@@ -78,6 +78,11 @@ def _discover_parsers() -> list[_ParserInfo]:
             # Inkling opts out of token-id terminal matching and has typed
             # structural blocks; its replay coverage lives in test_inkling.py.
             continue
+        if cfg.name == "granite":
+            # Granite has a JSON-array tool body with no TOOL_END terminal, so
+            # it does not fit this token-terminal harness; its replay coverage
+            # lives in test_granite.py.
+            continue
         tool_end = cfg.token_id_terminals.get("TOOL_END")
         if not tool_end:
             raise RuntimeError(
diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py
index 36e347fc572f..2851edd4d0d9 100644
--- a/tests/parser/engine/trace_builder.py
+++ b/tests/parser/engine/trace_builder.py
@@ -34,6 +34,7 @@
     DeepSeekV41Parser,
     Gemma4Parser,
     Glm47MoeParser,
+    GraniteParser,
     InklingParser,
     KimiK2Parser,
     MinimaxM2Parser,
@@ -1042,6 +1043,89 @@ def _build_inkling(scenario: Scenario, validate: bool = True) -> Sample:
     return sample
 
 
+# ── Granite (JSON-array tool bodies, no reasoning) ───────────────────────
+
+_GRANITE_VOCAB: dict[str, int] = {
+    "<|tool_call|>": 49154,
+}
+
+
+def _granite_segments(scenario: Scenario) -> list[tuple[str, bool]]:
+    segs: list[tuple[str, bool]] = []
+    if scenario.content:
+        # Granite has no reasoning; prose is plain content preceding the marker.
+        segs.append((scenario.content, False))
+    if scenario.tool_calls:
+        segs.append(("<|tool_call|>", True))
+        payload = json.dumps(
+            [
+                {"name": tc.name, "arguments": tc.arguments}
+                for tc in scenario.tool_calls
+            ],
+            ensure_ascii=False,
+            separators=(", ", ": "),
+        )
+        segs.append((" " + payload, False))
+    return segs
+
+
+def _granite_expected_content(scenario: Scenario) -> str | None:
+    if scenario.tool_calls:
+        if not scenario.content:
+            return None
+        return scenario.content.strip() or None
+    return scenario.content
+
+
+def _build_granite(scenario: Scenario, validate: bool = True) -> Sample:
+    sample = _make_sample(
+        sample_id=f"granite-{scenario.id}",
+        description=scenario.description,
+        vocab=_GRANITE_VOCAB,
+        segments=_granite_segments(scenario),
+        expected_reasoning=None,
+        expected_content=_granite_expected_content(scenario),
+        expected_tool_calls=_expected_tc(scenario),
+        tools=_expected_tools(scenario),
+    )
+    if validate:
+        _validate_sample(sample, GraniteParser)
+    return sample
+
+
+# Granite has no reasoning, so the shared reasoning-centric SCENARIOS do not
+# apply; these exercise the JSON-array tool body (single, parallel, surrounding
+# text) instead.
+_GRANITE_SCENARIOS: list[Scenario] = [
+    Scenario(
+        id="single-tool",
+        description="Single tool call",
+        tool_calls=[_READ_TOOL],
+    ),
+    Scenario(
+        id="parallel-tools",
+        description="Parallel tool calls in one JSON array",
+        tool_calls=[_BASH_TOOL, _WEATHER_TOOL],
+    ),
+    Scenario(
+        id="complex-json-args",
+        description="Tool call with nested objects, arrays, numbers, booleans",
+        tool_calls=[_COMPLEX_TOOL],
+    ),
+    Scenario(
+        id="content-only",
+        description="Plain content response without tool calls",
+        content="Hello! How can I help you today?",
+    ),
+    Scenario(
+        id="surrounding-text",
+        description="Prose content preceding the tool call",
+        content="Let me check the weather.",
+        tool_calls=[_WEATHER_TOOL],
+    ),
+]
+
+
 # ── Registry and public API ──────────────────────────────────────────
 
 _BUILDERS: dict[str, Any] = {
@@ -1049,6 +1133,7 @@ def _build_inkling(scenario: Scenario, validate: bool = True) -> Sample:
     "deepseek_v4": _build_deepseek_v4,
     "deepseek_v41": functools.partial(_build_deepseek_v4, v41=True),
     "gemma4": _build_gemma4,
+    "granite": _build_granite,
     "minimax_m2": _build_minimax_m2,
     "nemotron_v3": _build_nemotron_v3,
     "seed_oss": _build_seed_oss,
diff --git a/tests/parser/mistral/test_tool_calls.py b/tests/parser/mistral/test_tool_calls.py
index d5e09453a035..25b41f44c2d2 100644
--- a/tests/parser/mistral/test_tool_calls.py
+++ b/tests/parser/mistral/test_tool_calls.py
@@ -2045,7 +2045,6 @@ def test_content_tool_calls_transition_emits_reasoning_end(
     reasoning_encoding,
 ):
     """(CONTENT, TOOL_CALLS) must emit REASONING_END when reasoning is enabled."""
-
     cfg = mistral_config(reasoning_encoding=reasoning_encoding)
     engine = StreamingParserEngine(
         config=cfg, tokenizer=mistral_tokenizer, vocab=mistral_tokenizer.get_vocab()
@@ -2219,7 +2218,7 @@ def test_adjust_request_pre_v11_required_clears_response_format(
     mistral_pre_v11_tool_parser: MistralToolParser,
     response_format: dict,
 ) -> None:
-    """required + response_format must inject the tool schema and clear
+    """Required + response_format must inject the tool schema and clear
     response_format so the tool schema is the sole structured-output
     constraint, matching the base ToolParser (otherwise the request either
     hits the "multiple constraints" engine error or, with no constraint,
diff --git a/tests/platforms/test_cpu.py b/tests/platforms/test_cpu.py
index fd270315eb11..9c7f0b188bc8 100644
--- a/tests/platforms/test_cpu.py
+++ b/tests/platforms/test_cpu.py
@@ -1,12 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
 from types import SimpleNamespace
 
 import pytest
 import torch
 
-from vllm.config import CacheConfig, VllmConfig
+from vllm.config import CacheConfig, KVTransferConfig, VllmConfig
 from vllm.platforms.cpu import CpuPlatform
 
 
@@ -44,6 +45,7 @@ def _cpu_config(
         architecture=architecture or model_type,
         has_inner_state=True,
         use_mla=False,
+        hf_config=SimpleNamespace(model_type=model_type),
         hf_text_config=SimpleNamespace(
             model_type=model_type,
             layer_types=layer_types,
@@ -150,3 +152,74 @@ def test_cpu_accelerated_gdn_dtype_policy(
 
     CpuPlatform.check_and_update_config(config)
     assert cache_config.mamba_ssm_cache_dtype == expected_dtype
+
+
+@pytest.mark.parametrize(
+    (
+        "connector",
+        "extra_config",
+        "explicit_layout",
+        "avx512_bf16_supported",
+        "expected_layout",
+    ),
+    [
+        pytest.param("NixlConnector", None, None, True, "DS", id="nixl"),
+        pytest.param("NixlPullConnector", None, None, True, "DS", id="nixl-pull"),
+        pytest.param("NixlPushConnector", None, None, True, "DS", id="nixl-push"),
+        pytest.param(
+            "MultiConnector",
+            {
+                "connectors": [
+                    {
+                        "kv_connector": "ExampleConnector",
+                        "kv_connector_extra_config": {},
+                    },
+                    {
+                        "kv_connector": "NixlConnector",
+                        "kv_connector_extra_config": {},
+                    },
+                ]
+            },
+            None,
+            True,
+            "DS",
+            id="multi-nixl",
+        ),
+        pytest.param("NixlConnector", None, "SD", True, "SD", id="explicit-override"),
+        pytest.param("OffloadingConnector", None, None, True, "SD", id="non-nixl"),
+        pytest.param(
+            "OffloadingConnector", None, None, False, None, id="non-nixl-no-avx"
+        ),
+    ],
+)
+def test_cpu_conv_state_layout_selection(
+    monkeypatch: pytest.MonkeyPatch,
+    connector: str,
+    extra_config: dict | None,
+    explicit_layout: str | None,
+    avx512_bf16_supported: bool,
+    expected_layout: str | None,
+) -> None:
+    layout_env = "VLLM_SSM_CONV_STATE_LAYOUT"
+    monkeypatch.setattr(
+        "torch.cpu._is_avx512_bf16_supported",
+        lambda: avx512_bf16_supported,
+    )
+
+    kv_transfer_config = KVTransferConfig(
+        kv_connector=connector,
+        kv_connector_extra_config=extra_config or {},
+        kv_role="kv_both",
+    )
+    config = _cpu_config(
+        CacheConfig(mamba_ssm_cache_dtype="float32"),
+        model_type="qwen3_5",
+        resolved_dtype="float32",
+    )
+    config.kv_transfer_config = kv_transfer_config
+    monkeypatch.delenv(layout_env, raising=False)
+    if explicit_layout is not None:
+        monkeypatch.setenv(layout_env, explicit_layout)
+
+    CpuPlatform.check_and_update_config(config)
+    assert os.environ.get(layout_env) == expected_layout
diff --git a/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py b/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py
index a35cb91bdb00..b2cf6f23455a 100644
--- a/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py
+++ b/tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py
@@ -121,7 +121,7 @@ def _build_document_prompt(
         tokenizer,
         content_ids: list[int],
     ) -> TokensPrompt:
-        """[CLS] [DocumentMarker] <tokens> [SEP]"""
+        """[CLS] [DocumentMarker] <tokens> [SEP]."""
         _, document_marker_id = self._resolve_marker_ids(tokenizer)
 
         content_ids = content_ids[: self.max_model_len - 3]
diff --git a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py
index d83399e07aa2..3c540f480aaa 100644
--- a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py
+++ b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py
@@ -58,6 +58,7 @@ def save_geotiff(image: torch.Tensor, meta: dict, out_format: str) -> str | byte
         image: np.ndarray with shape (bands, height, width)
         output_path: path where to save the image
         meta: dict with meta info.
+
     """
     if out_format == "path":
         # create temp file
@@ -100,8 +101,8 @@ def read_geotiff(
     Returns:
         np.ndarray with shape (bands, height, width)
         meta info dict
-    """
 
+    """
     if all([x is None for x in [file_path, path_type, file_data]]):
         raise Exception("All input fields to read_geotiff are None")
     write_to_file: bytes | None = None
@@ -168,8 +169,8 @@ def load_image(
     Returns:
         np.array containing created example
         list of meta info for each image in *file_paths*
-    """
 
+    """
     imgs: list[np.ndarray] = []
     metas = []
     temporal_coords = []
diff --git a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py
index 3a1a9c3be41e..c75c5626161c 100644
--- a/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py
+++ b/tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py
@@ -39,13 +39,13 @@ class ImagePrompt(BaseModel):
 
 
 class ImageRequestOutput(BaseModel):
-    """
-    The output data of an image request to vLLM.
+    """The output data of an image request to vLLM.
 
     Args:
         type (str): The data content type [path, object]
         format (str): The image format (e.g., jpeg, png, etc.)
         data (Any): The resulting data.
+
     """
 
     type: Literal["path", "b64_json"]
diff --git a/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py b/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py
index 7ba2c09b283c..6acd2a9f6dbf 100644
--- a/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py
+++ b/tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py
@@ -5,8 +5,7 @@
 
 
 class DummyStatLogger(StatLoggerBase):
-    """
-    A dummy stat logger for testing purposes.
+    """A dummy stat logger for testing purposes.
     Implements the minimal interface expected by StatLoggerManager.
     """
 
diff --git a/tests/plugins_tests/gguf/test_gguf_plugin_generate.py b/tests/plugins_tests/gguf/test_gguf_plugin_generate.py
index fbda46527530..2f35a272c0da 100644
--- a/tests/plugins_tests/gguf/test_gguf_plugin_generate.py
+++ b/tests/plugins_tests/gguf/test_gguf_plugin_generate.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E tests for GGUF plugin functionality.
-"""
+"""E2E tests for GGUF plugin functionality."""
 
 import os
 from typing import NamedTuple
diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py
index 2bc463b593cb..ef5f29602fde 100644
--- a/tests/quantization/test_auto_round.py
+++ b/tests/quantization/test_auto_round.py
@@ -146,6 +146,21 @@
     },
 }
 
+# These checkpoints fail during XPU graph capture, so run them eagerly there.
+XPU_EAGER_ONLY_MODELS = {
+    "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc",
+    "Intel/Qwen2-0.5B-Instruct-int4-sym-AutoRound",
+    "Intel/Qwen3-8B-w2g64-for-ut",
+    "INCModel/Qwen3-30B-A3B-12L-W4A16-test",
+}
+
+
+def _runner_kwargs(model: str) -> dict[str, Any]:
+    kwargs = dict(MODEL_RUNNER_KWARGS.get(model, {}))
+    if current_platform.is_xpu() and model in XPU_EAGER_ONLY_MODELS:
+        kwargs["enforce_eager"] = True
+    return kwargs
+
 
 @pytest.mark.skipif(
     not (
@@ -157,7 +172,7 @@
 )
 @pytest.mark.parametrize("model", MODELS + QWEN3_AUTOROUND_MODELS)
 def test_auto_round_model(vllm_runner, model):
-    with vllm_runner(model, **MODEL_RUNNER_KWARGS.get(model, {})) as llm:
+    with vllm_runner(model, **_runner_kwargs(model)) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=8)
 
     assert output
@@ -1638,19 +1653,6 @@ def test_wna16_linear_gptq_unsupported_config_raises() -> None:
         INCWNA16LinearScheme(make_layer_config(sym=False))
 
 
-def test_wna16_xpu_unsupported_config_still_raises(monkeypatch) -> None:
-    monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
-    monkeypatch.setattr(current_platform, "is_cpu", lambda: False)
-
-    with pytest.raises(NotImplementedError, match="unsupported config"):
-        INCWna16Scheme().get_linear_method(
-            make_config(weight_bits=2, sym=False),
-            object(),
-            "layer",
-            make_layer_config(bits=2, sym=False),
-        )
-
-
 def test_inc_get_quant_method_unquantized_linear_returns_unquantized() -> None:
     config = make_config(extra_config={"layer": {"bits": 16}})
     layer = object.__new__(LinearBase)
diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py
index ef80287e5ff2..126ac92eefe2 100644
--- a/tests/quantization/test_compressed_tensors.py
+++ b/tests/quantization/test_compressed_tensors.py
@@ -574,8 +574,7 @@ def test_compressed_tensors_fp8_block_enabled(
     reason="This test is not for non-CUDA platforms",
 )
 def test_compressed_tensors_moe_ignore_with_model(vllm_runner):
-    """
-    Integration test for MoE layer ignore functionality with a real model.
+    """Integration test for MoE layer ignore functionality with a real model.
 
     This test would verify that when loading a compressed-tensors quantized
     MoE model where some MoE layers are in the ignore list, those layers
@@ -587,7 +586,6 @@ def test_compressed_tensors_moe_ignore_with_model(vllm_runner):
     - Config with ignore list containing specific MoE layers
     - Multiple MoE layers where some are quantized and some are not
     """
-
     # model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only" # CT 12.3
     model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable"  # CT 12.2
 
diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py
index 1c73573d8dc0..65e119b083fd 100644
--- a/tests/quantization/test_fp8.py
+++ b/tests/quantization/test_fp8.py
@@ -37,6 +37,7 @@
     Fp8MoEMethod,
 )
 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
+from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod
 from vllm.model_executor.layers.quantization.online.fp8 import (
     Fp8PerTensorOnlineLinearMethod,
 )
@@ -63,6 +64,383 @@
 ]
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA")
+@pytest.mark.parametrize("tp_rank", [0, 1])
+@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
+def test_deepseek_v41_mxfp8_scale_loading(
+    dist_init, default_vllm_config, monkeypatch, tp_rank, scale_dtype
+):
+    """Expand linear scales before TP slicing, including padded shared experts."""
+    from vllm.model_executor.layers import linear as linear_module
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        RowParallelLinear,
+    )
+    from vllm.models.deepseek_v41 import quant_config as quant_module
+    from vllm.models.deepseek_v41.nvidia import model as model_module
+
+    default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
+    quant_config = quant_module.DeepseekV4FP8Config(
+        is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]
+    )
+    quant_config._resolved_expert_dtype = "fp4"
+    for module in (model_module, linear_module):
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: tp_rank)
+    tp_args = dict(quant_config=quant_config, bias=False)
+    block = torch.nn.Module()
+    block.attn = torch.nn.Module()
+    block.attn.wq_b = ColumnParallelLinear(128, 128, **tp_args)
+    block.attn.fused_wqa_wkv = MergedColumnParallelLinear(128, [128, 64], **tp_args)
+    block.ffn = torch.nn.Module()
+    block.ffn.shared_experts = torch.nn.Module()
+    block.ffn.shared_experts.gate_up_proj = MergedColumnParallelLinear(
+        128, [128, 128], **tp_args
+    )
+    block.ffn.shared_experts.down_proj = RowParallelLinear(128, 128, **tp_args)
+    block.ffn.experts = torch.nn.Module()
+    block.ffn.experts.register_parameter(
+        "weight_scale", torch.nn.Parameter(torch.empty(3, 4, dtype=torch.uint8), False)
+    )
+    block.engram = torch.nn.Module()
+    block.engram.register_parameter(
+        "weight_scale_inv",
+        torch.nn.Parameter(torch.empty(3, 4, dtype=scale_dtype), False),
+    )
+
+    def load_expert_scale(param, weight, *args, **kwargs):
+        param.data.copy_(weight)
+        return True
+
+    block.ffn.experts.weight_scale.weight_loader = load_expert_scale
+    model = model_module.DeepseekV4Model.__new__(model_module.DeepseekV4Model)
+    torch.nn.Module.__init__(model)
+    model.layers = torch.nn.ModuleList([block])
+    model.config = SimpleNamespace(num_attention_heads=2)
+    model.quant_config = quant_config
+    model.use_sequence_parallel = False
+    model.get_expert_mapping = lambda: [
+        ("experts.weight_scale", "experts.0.w1.weight_scale", 0, "w1")
+    ]
+
+    checkpoints = [
+        ("attn.wq_b", "attn.wq_b", 128, 128, 0, None),
+        ("attn.wq_a", "attn.fused_wqa_wkv", 128, 128, 0, slice(0, 64)),
+        ("attn.wkv", "attn.fused_wqa_wkv", 64, 128, 0, slice(64, 96)),
+        (
+            "ffn.shared_experts.w1",
+            "ffn.shared_experts.gate_up_proj",
+            96,
+            128,
+            0,
+            slice(0, 64),
+        ),
+        (
+            "ffn.shared_experts.w3",
+            "ffn.shared_experts.gate_up_proj",
+            96,
+            128,
+            0,
+            slice(64, 128),
+        ),
+        (
+            "ffn.shared_experts.down_proj",
+            "ffn.shared_experts.down_proj",
+            128,
+            96,
+            1,
+            None,
+        ),
+    ]
+    weights = []
+    expected = []
+    for source, target, n, k, axis, shard in checkpoints:
+        weight = torch.randint(-4, 5, (n, k)).to(torch.float8_e4m3fn)
+        scale_bytes = torch.randint(124, 131, (n // 32, k // 32), dtype=torch.uint8)
+        weights.extend(
+            [
+                (f"layers.0.{source}.weight", weight),
+                (f"layers.0.{source}.scale", scale_bytes.view(scale_dtype)),
+            ]
+        )
+        dequant = (
+            weight.float().reshape(n // 32, 32, k // 32, 32)
+            * torch.exp2(scale_bytes.float() - 127)[:, None, :, None]
+        ).reshape(n, k)
+        if "shared_experts" in source:
+            padded = torch.zeros(128, 128)
+            padded[:n, :k] = dequant
+            dequant = padded
+        expected.append((target, shard, dequant.chunk(2, dim=axis)[tp_rank]))
+
+    expert_scale = torch.full((3, 4), 125, dtype=torch.uint8)
+    weights.append(
+        ("layers.0.ffn.experts.0.w1.weight_scale", expert_scale.view(scale_dtype))
+    )
+    weights.append(("layers.0.engram.weight_scale_inv", expert_scale.view(scale_dtype)))
+    mapper = model_module._make_deepseek_v4_weights_mapper("fp4", "weight_scale")
+    loaded = model.load_weights(
+        (name.removeprefix("model."), weight) for name, weight in mapper.apply(weights)
+    )
+    assert "layers.0.ffn.experts.weight_scale" in loaded
+    assert torch.equal(block.ffn.experts.weight_scale, expert_scale)
+    assert torch.equal(block.engram.weight_scale_inv.view(torch.uint8), expert_scale)
+    for target, shard, reference in expected:
+        linear = block.get_submodule(target)
+        assert isinstance(linear.quant_method, ModelOptLinearMethod)
+        assert f"layers.0.{target}.weight_scale" in loaded
+        weight = linear.weight if shard is None else linear.weight[shard]
+        scale = linear.weight_scale if shard is None else linear.weight_scale[shard]
+        actual = (
+            weight.float().unflatten(-1, (-1, 32))
+            * torch.exp2(scale.float() - 127).unsqueeze(-1)
+        ).flatten(-2)
+        torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA")
+@pytest.mark.parametrize(
+    "weight_block_size,expert_dtype,scale_name",
+    [
+        ([32, 32], "fp4", "weight_scale"),
+        ([128, 128], "fp4", "weight_scale_inv"),
+        ([32, 32], "fp8", "weight_scale_inv"),
+    ],
+)
+def test_deepseek_v41_vl_mapper_routes_linear_scales(
+    weight_block_size, expert_dtype, scale_name
+):
+    """The VL wrapper must map ``.scale`` keys to the parameter the linear
+    quant method registers, as the text model does. A hardcoded
+    ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints."""
+    from vllm.models.deepseek_v41.nvidia import model as model_module
+    from vllm.models.deepseek_v41.nvidia import vl_model as vl_module
+
+    vllm_config = SimpleNamespace(
+        quant_config=SimpleNamespace(weight_block_size=weight_block_size)
+    )
+    resolved = model_module._linear_scale_param_name(vllm_config, expert_dtype)
+    assert resolved == scale_name
+
+    mapper = vl_module._make_deepseek_v4_vl_weights_mapper(expert_dtype, resolved)
+    weight = torch.empty(0)
+    mapped = [
+        name
+        for name, _ in mapper.apply(
+            [("layers.0.attn.wq_a.scale", weight), ("layers.0.attn.wkv.weight", weight)]
+        )
+    ]
+    assert mapped == [
+        f"language_model.model.layers.0.attn.wq_a.{scale_name}",
+        "language_model.model.layers.0.attn.wkv.weight",
+    ]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
+@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
+@pytest.mark.parametrize(
+    "weight_shape,is_bmm",
+    [
+        ((129, 160), False),
+        ((387, 160), True),
+        ((3, 129, 160), False),
+        ((3, 128, 512), False),
+    ],
+)
+def test_deepgemm_mxfp8_preserves_weight_and_scale_values(
+    scale_dtype, weight_shape, is_bmm
+):
+    """DeepGEMM layout conversion preserves native weights and E8M0 scales."""
+    from vllm.model_executor.layers.quantization.utils import fp8_utils
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    if not is_deep_gemm_supported() or not current_platform.is_device_capability_family(
+        100
+    ):
+        pytest.skip("DeepGEMM MXFP8 requires Blackwell")
+
+    weight = torch.randn(weight_shape, device="cuda").to(torch.float8_e4m3fn)
+    scales = (
+        torch.randint(
+            1,
+            255,
+            (weight.numel() // 32 + 1,),
+            device="cuda",
+            dtype=torch.uint8,
+        )[1:]
+        .view(*weight_shape[:-1], weight_shape[-1] // 32)
+        .view(scale_dtype)
+    )
+    original_weight = weight.view(torch.uint8).clone()
+    original_scales = scales.view(torch.uint8).clone()
+    processed_weight, packed = fp8_utils.deepgemm_post_process_fp8_weight_block(
+        weight,
+        scales,
+        quant_block_shape=(1, 32),
+        use_e8m0=True,
+        is_bmm=is_bmm,
+        bmm_batch_size=3 if is_bmm else 0,
+    )
+    assert packed.dtype == torch.int32
+    unpacked = torch.stack(
+        [(packed >> (8 * i)) & 0xFF for i in range(4)], dim=-1
+    ).flatten(-2)
+    unpacked = unpacked[..., : scales.shape[-1]].to(torch.uint8)
+    torch.testing.assert_close(
+        unpacked.reshape_as(original_scales), original_scales, rtol=0, atol=0
+    )
+    torch.testing.assert_close(
+        processed_weight.view(torch.uint8).reshape_as(original_weight),
+        original_weight,
+        rtol=0,
+        atol=0,
+    )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
+@pytest.mark.parametrize("prequantized", [False, True])
+@pytest.mark.parametrize("config_source", ["deepseek", "mxfp8"])
+@pytest.mark.parametrize("num_tokens", [1, 7, 128])
+def test_mxfp8_bmm_loads_and_projects_grouped_weights(
+    dist_init, default_vllm_config, prequantized, config_source, num_tokens
+):
+    """BMM metadata set after construction selects grouped weight processing."""
+    from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import (
+        DeepGemmMxfp8BmmLinearKernel,
+    )
+    from vllm.model_executor.layers.linear import ColumnParallelLinear
+    from vllm.model_executor.layers.quantization import get_quantization_config
+    from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
+        fused_inv_rope_fp8_quant,
+    )
+    from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
+        compute_fp8_einsum_recipe,
+        deep_gemm_fp8_o_proj,
+    )
+    from vllm.models.deepseek_v41.nvidia.model import DeepseekV4Model
+    from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    if not is_deep_gemm_supported() or not current_platform.is_device_capability_family(
+        100
+    ):
+        pytest.skip("DeepGEMM MXFP8 BMM requires Blackwell")
+
+    default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
+    quant_config = DeepseekV4FP8Config(
+        is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]
+    )
+    quant_config._resolved_expert_dtype = "fp4"
+    if config_source == "mxfp8":
+        quant_config = get_quantization_config("mxfp8").from_config(
+            {"quant_method": "mxfp8"}
+        )
+    with torch.device("cuda"):
+        linear = ColumnParallelLinear(
+            512,
+            256,
+            bias=False,
+            quant_config=quant_config,
+            return_bias=False,
+        )
+    linear.is_bmm = True
+    linear.bmm_batch_size = 2
+    weight = torch.randn(256, 512, device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.randint(124, 131, (8, 16), device="cuda", dtype=torch.uint8)
+    checkpoint_scales = (
+        scales.view(torch.float8_e8m0fnu)
+        if config_source == "deepseek"
+        else scales.repeat_interleave(32, dim=0)
+    )
+    model = DeepseekV4Model.__new__(DeepseekV4Model)
+    torch.nn.Module.__init__(model)
+    block = torch.nn.Module()
+    block.attn = torch.nn.Module()
+    block.attn.wo_a = linear
+    model.layers = torch.nn.ModuleList([block])
+    model.config = SimpleNamespace(num_attention_heads=2)
+    model.quant_config = quant_config
+    model.use_sequence_parallel = False
+    model.get_expert_mapping = lambda: []
+    model.load_weights(
+        [
+            ("layers.0.attn.wo_a.weight", weight),
+            ("layers.0.attn.wo_a.weight_scale", checkpoint_scales),
+        ]
+    )
+    reference_weight = (
+        weight.float().reshape(8, 32, 16, 32)
+        * torch.exp2(scales.float() - 127)[:, None, :, None]
+    ).reshape(2, 128, 512)
+    linear.quant_method.process_weights_after_loading(linear)
+    linear.quant_method.process_weights_after_loading(linear)
+    assert isinstance(linear.quant_method.kernel, DeepGemmMxfp8BmmLinearKernel)
+    assert linear.weight.shape == (2, 128, 512)
+    torch.testing.assert_close(
+        linear.weight.view(torch.uint8).flatten(),
+        weight.view(torch.uint8).flatten(),
+        rtol=0,
+        atol=0,
+    )
+    unpacked_scales = torch.stack(
+        [(linear.weight_scale >> (8 * i)) & 0xFF for i in range(4)], dim=-1
+    ).flatten(-2)
+    torch.testing.assert_close(
+        unpacked_scales.to(torch.uint8).reshape(256, 16),
+        scales.repeat_interleave(32, dim=0),
+        rtol=0,
+        atol=0,
+    )
+
+    x = torch.randn(2, num_tokens, 512, device="cuda", dtype=torch.bfloat16)
+    reference = torch.einsum("gmk,gnk->mgn", x.float(), reference_weight)
+    inputs = x.transpose(0, 1)
+    if prequantized:
+        cache = torch.cat(
+            (
+                torch.ones(num_tokens, 32, device="cuda"),
+                torch.zeros(num_tokens, 32, device="cuda"),
+            ),
+            dim=1,
+        )
+        inputs = fused_inv_rope_fp8_quant(
+            x.permute(1, 0, 2),
+            torch.arange(num_tokens, device="cuda"),
+            cache,
+            n_groups=2,
+            heads_per_group=1,
+            nope_dim=448,
+            rope_dim=64,
+            quant_group_size=32,
+            tma_aligned_scales=current_platform.has_device_capability(100),
+        )
+    with torch.no_grad():
+        output = linear(inputs)
+    assert output.shape == reference.shape
+    assert (output.float() - reference).norm() / reference.norm() < 0.06
+    if prequantized:
+        recipe, tma_aligned_scales = compute_fp8_einsum_recipe(block_size=32)
+        projected = deep_gemm_fp8_o_proj(
+            x.permute(1, 0, 2),
+            torch.arange(num_tokens, device="cuda"),
+            cache,
+            linear,
+            torch.nn.Identity(),
+            n_groups=2,
+            heads_per_group=1,
+            nope_dim=448,
+            rope_dim=64,
+            o_lora_rank=128,
+            einsum_recipe=recipe,
+            tma_aligned_scales=tma_aligned_scales,
+        )
+        torch.testing.assert_close(projected, output.flatten(1), rtol=0, atol=0)
+    compiled = torch.compile(linear, backend="eager", fullgraph=True)
+    with torch.no_grad():
+        torch.testing.assert_close(compiled(inputs), output, rtol=0, atol=0)
+
+
 def test_prepare_gated_trtllm_fp8_moe_weights_pads_each_projection(monkeypatch):
     monkeypatch.setattr(
         flashinfer_utils,
diff --git a/tests/quantization/test_gfx950_moe.py b/tests/quantization/test_gfx950_moe.py
index 67728a2da8aa..b71e952b769a 100644
--- a/tests/quantization/test_gfx950_moe.py
+++ b/tests/quantization/test_gfx950_moe.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for MXFP4 MoE oracle backend selection on mi355x (GFX950).
+"""Tests for MXFP4 MoE oracle backend selection on mi355x (GFX950).
 
 These tests run on real hardware — no mocks. Skipped on non-GFX950 platforms.
 """
diff --git a/tests/quantization/test_gptq_v2.py b/tests/quantization/test_gptq_v2.py
deleted file mode 100644
index a10088f40390..000000000000
--- a/tests/quantization/test_gptq_v2.py
+++ /dev/null
@@ -1,106 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests whether vllm correctly load and run gptq_v2 format checkpoints.
-
-Run `pytest tests/quantization/test_gptq_v2.py --forked`.
-
-Note: 2/3-bit GPTQ models are no longer supported after the consolidation
-to Marlin kernels. Only 4/8-bit symmetric GPTQ models are supported.
-"""
-
-import pytest
-import torch
-from transformers import AutoTokenizer
-
-from vllm import SamplingParams
-from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod
-
-# A dummy small model quantized by GPTQModel, stored in GPTQ v2 format
-# Note: This is a 2-bit model which is no longer supported with Marlin kernels
-MODELS = ["XXXXyu/Qwen3-1.7B-w2g64-gptq_v2"]
-
-# Generate multiple sequences for testing, because an 1.7B 2-bit model
-# cannot always generate normal texts.
-N_SEQ = 5
-
-
-@pytest.mark.skip(reason="2-bit GPTQ is no longer supported after Marlin consolidation")
-@pytest.mark.parametrize("model_id", MODELS)
-def test_model_load(vllm_runner, model_id, monkeypatch):
-    # `LLM.apply_model` requires pickling a function.
-    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
-
-    with vllm_runner(model_id, dtype=torch.float16, max_model_len=512) as llm:
-
-        def check_model(model_id):
-            for name, submodule in model_id.named_modules():
-                # Could check more modules if necessary
-                if name == "model_id.layers.0.self_attn.qkv_proj":
-                    assert isinstance(submodule.quant_method, AutoGPTQLinearMethod)
-                    # Just break since currently we only check 1 module
-                    break
-
-        # Check if gptq_v2 format is correctly loaded
-        llm.apply_model(check_model)
-
-
-@pytest.mark.skip(reason="2-bit GPTQ is no longer supported after Marlin consolidation")
-@pytest.mark.parametrize("model_id", MODELS)
-def test_model_inference(vllm_runner, model_id):
-    # Prepare prompt to test the model's generation result.
-    prompt = "What is the meaning of life?"
-    messages = [
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": prompt},
-    ]
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True,
-        enable_thinking=False,  # If thinking model, set it to false
-    )
-    sampling_params = SamplingParams(
-        n=N_SEQ,
-        max_tokens=128,
-        temperature=0.7,
-        top_p=0.8,
-        top_k=20,
-        min_p=0,
-        presence_penalty=2,
-    )
-
-    with vllm_runner(model_id, dtype=torch.float16, max_model_len=512) as llm:
-        # Generate a response to verify inference correctness
-        output = llm.generate(text, sampling_params)
-
-    # Make sure the output exists
-    assert output
-    assert output[0][1]
-    assert len(output[0][1]) == N_SEQ
-
-    def has_normal_char_distribution(texts, min_len):
-        for text in texts:
-            # Response too short
-            if len(text) < min_len:
-                return False
-
-            # Basic ratio checks
-            letters = sum(c.isalpha() for c in text)
-            spaces = sum(c.isspace() for c in text)
-            total = len(text)
-
-            letter_ratio = letters / total
-            space_ratio = spaces / total
-
-            # At least 1 normal text should exist within output sequences
-            # Normal text should be mostly letters with reasonable spacing
-            # Some magic numbers, could be adjusted
-            if 0.5 <= letter_ratio <= 0.9 and 0.01 <= space_ratio <= 0.3:
-                return True
-        # No sequence contains normal text, output might be broken
-        return False
-
-    # Apply some simple checks for giberish output
-    # Print the output sequences if failed
-    assert has_normal_char_distribution(output[0][1], 5), output[0][1]
diff --git a/tests/quantization/test_humming_mxfp4_block_fp8.py b/tests/quantization/test_humming_mxfp4_block_fp8.py
index dbec989df3f9..74c266e0bfe5 100644
--- a/tests/quantization/test_humming_mxfp4_block_fp8.py
+++ b/tests/quantization/test_humming_mxfp4_block_fp8.py
@@ -15,6 +15,7 @@
 """
 
 import pytest
+import torch
 
 from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
     HummingExpertsBase,
@@ -56,8 +57,8 @@ def test_mxfp4_weight_with_per_token_fp8_activation_still_supported():
     ],
 )
 def test_humming_fp8_input_schema_to_quant_key(group_size, expected_key):
-    from vllm.model_executor.layers.quantization.utils.humming_utils import (
-        _humming_input_schema_to_quant_key,
+    from vllm.model_executor.layers.quantization.utils.humming import (
+        input_schema_to_quant_key,
     )
     from vllm.utils.humming import HummingInputSchema
     from vllm.utils.humming import dtypes as humming_dtypes
@@ -66,18 +67,21 @@ def test_humming_fp8_input_schema_to_quant_key(group_size, expected_key):
         a_dtype=humming_dtypes.float8e4m3,
         input_scale_group_size=group_size,
     )
-    assert _humming_input_schema_to_quant_key(schema) == expected_key
+    assert input_schema_to_quant_key(schema, param_dtype=torch.bfloat16) == expected_key
 
 
 @pytest.mark.skipif(not has_humming(), reason="humming is not installed")
 def test_humming_bf16_input_schema_is_unquantized():
-    from vllm.model_executor.layers.quantization.utils.humming_utils import (
-        _humming_input_schema_to_quant_key,
+    from vllm.model_executor.layers.quantization.utils.humming import (
+        input_schema_to_quant_key,
     )
     from vllm.utils.humming import HummingInputSchema
 
     # No a_dtype -> unquantized (bf16/fp16) inputs -> None.
-    assert _humming_input_schema_to_quant_key(HummingInputSchema()) is None
+    assert (
+        input_schema_to_quant_key(HummingInputSchema(), param_dtype=torch.bfloat16)
+        is None
+    )
 
 
 def test_block_fp8_activation_quant_config_is_block_quantized():
@@ -86,7 +90,7 @@ def test_block_fp8_activation_quant_config_is_block_quantized():
     DeepEP prepare/finalize step quantize activations to block FP8 *before* the
     all-to-all dispatch instead of deferring to Humming -- and what makes
     HummingExpertsBase.expects_unquantized_inputs return False."""
-    from vllm.model_executor.layers.quantization.utils.humming_utils import (
+    from vllm.model_executor.layers.quantization.utils.humming import (
         make_humming_moe_quant_config,
     )
     from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
@@ -108,7 +112,7 @@ def test_default_activation_quant_config_defers_to_humming():
     """Without a block activation shape the config stays per-token (the deferred
     path): Humming quantizes internally, preserving pre-existing behavior for
     every non-block-FP8 scheme."""
-    from vllm.model_executor.layers.quantization.utils.humming_utils import (
+    from vllm.model_executor.layers.quantization.utils.humming import (
         make_humming_moe_quant_config,
     )
     from vllm.platforms import current_platform
diff --git a/tests/quantization/test_int8_moe_oracle.py b/tests/quantization/test_int8_moe_oracle.py
index 2eb6fa92a547..9580eb9ab328 100644
--- a/tests/quantization/test_int8_moe_oracle.py
+++ b/tests/quantization/test_int8_moe_oracle.py
@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for INT8 (W8A8) fused-MoE oracle backend selection.
+"""Tests for INT8 (W8A8) fused-MoE oracle backend selection.
 
-These exercise ``select_int8_moe_backend`` only (no kernels are launched), so
-they run on any platform where the Triton INT8 MoE kernel is available — CUDA
-(SM >= 7.5) or ROCm — not just gfx950.
+These exercise ``select_int8_moe_backend`` only (no MoE kernels are launched),
+so they run on any platform where the Triton INT8 MoE kernel is available —
+CUDA (SM >= 7.5), ROCm, or XPU — not just gfx950.
 """
 
 import pytest
@@ -28,17 +27,22 @@
 )
 from vllm.platforms import current_platform
 
-# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+) and on
-# ROCm CDNA GPUs. Gate on that rather than on a specific arch.
+# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+), on
+# ROCm CDNA GPUs and on XPU. Gate on that rather than on a specific arch.
 INT8_MOE_SUPPORTED = (
-    current_platform.is_cuda() and current_platform.has_device_capability((7, 5))
-) or current_platform.is_rocm()
+    (current_platform.is_cuda() and current_platform.has_device_capability((7, 5)))
+    or current_platform.is_rocm()
+    or current_platform.is_xpu()
+)
 
 requires_int8_moe = pytest.mark.skipif(
     not INT8_MOE_SUPPORTED,
-    reason="Requires a GPU with Triton INT8 MoE support (CUDA SM>=7.5 or ROCm)",
+    reason="Requires Triton INT8 MoE support (CUDA SM>=7.5, ROCm, or XPU)",
 )
 
+# So FusedMoEConfig.device names the device the test actually runs on.
+DEVICE = current_platform.device_type
+
 
 def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
     from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -53,7 +57,7 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
         moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
         activation=MoEActivation.SILU,
         in_dtype=torch.bfloat16,
-        device="cuda",
+        device=DEVICE,
         routing_method=RoutingMethodType.Renormalize,
         moe_backend=moe_backend,
     )
@@ -63,15 +67,20 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
 @pytest.mark.parametrize(
     "weight_key,activation_key",
     [
-        # per-channel weight + dynamic per-token activation
-        (kInt8StaticChannelSym, kInt8DynamicTokenSym),
-        # per-tensor weight + dynamic per-tensor activation
-        (kInt8StaticTensorSym, kInt8DynamicTensorSym),
+        pytest.param(
+            kInt8StaticChannelSym,
+            kInt8DynamicTokenSym,
+            id="per_channel_weight-per_token_act",
+        ),
+        pytest.param(
+            kInt8StaticTensorSym,
+            kInt8DynamicTensorSym,
+            id="per_tensor_weight-per_tensor_act",
+        ),
     ],
 )
 def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key):
-    """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor
-    weights) select the Triton backend."""
+    """Both dynamic-activation INT8 MoE schemes select the Triton backend."""
     config = _make_int8_moe_config()
     backend, experts_cls = select_int8_moe_backend(
         config, weight_key=weight_key, activation_key=activation_key
@@ -80,6 +89,28 @@ def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key):
     assert experts_cls is not None
 
 
+@pytest.mark.skipif(not current_platform.is_xpu(), reason="XPU-only behaviour")
+def test_scaled_int8_quant_is_available_on_xpu():
+    """The per-tensor scheme's activation quantization works on XPU.
+
+    ``_int8_quantize`` sends per-tensor activations through
+    ``ops.scaled_int8_quant``, which has an XPU branch. This is what lets the
+    per-tensor scheme be enabled here alongside CUDA and ROCm; pin it so a
+    future change cannot silently remove that branch.
+    """
+    from vllm import _custom_ops as ops
+
+    x = torch.randn(4, 16, device=DEVICE, dtype=torch.bfloat16)
+    scale = torch.full((1,), 0.05, device=DEVICE, dtype=torch.float32)
+
+    q, returned_scale, azp = ops.scaled_int8_quant(x, scale=scale)
+
+    assert q.dtype == torch.int8
+    assert q.shape == x.shape
+    assert azp is None
+    torch.testing.assert_close(returned_scale, scale)
+
+
 @requires_int8_moe
 def test_int8_explicit_moe_backend_triton():
     """An explicit --moe-backend triton selects the Triton INT8 backend."""
diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py
index 59a333bfe46b..ae8548c37058 100644
--- a/tests/quantization/test_modelopt.py
+++ b/tests/quantization/test_modelopt.py
@@ -127,6 +127,53 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head():
     assert method.spec.activation is kNvfp4Dynamic
 
 
+def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypatch):
+    """Standard MXFP8 checkpoints already have one scale row per weight row."""
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    kernel = Mock()
+    kernel.input_quant_key.return_value = None
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel",
+        lambda **kwargs: kernel,
+    )
+    config = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})
+    linear = ReplicatedLinear(64, 64, bias=False, quant_config=config)
+    scales = torch.arange(128, dtype=torch.uint8).reshape(64, 2)
+    linear.weight_scale.weight_loader(linear.weight_scale, scales)
+    assert torch.equal(linear.weight_scale, scales)
+
+
+def test_modelopt_mxfp8_pre_processed_weights_follow_kernel(monkeypatch):
+    """The weight cache daemon can only serve MXFP8 layers whose kernel needs
+    no post-load state beyond the parameters it exports."""
+    from vllm.config.quantization import QuantSpec
+    from vllm.model_executor.kernels.linear.mxfp8 import Mxfp8LinearKernel
+
+    method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
+    method.spec = QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic)
+    method.kernel = None
+    assert not method.supports_pre_processed_weights
+
+    method.kernel = Mock(spec=Mxfp8LinearKernel)
+    method.kernel.supports_pre_processed_weights = False
+    assert not method.supports_pre_processed_weights
+    method.kernel.supports_pre_processed_weights = True
+    assert method.supports_pre_processed_weights
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.quantization.modelopt.is_weights_pre_processed",
+        lambda: True,
+    )
+    layer = torch.nn.Module()
+    method.process_weights_after_loading(layer)
+    method.kernel.process_weights_after_loading.assert_not_called()
+
+    method.kernel.supports_pre_processed_weights = False
+    with pytest.raises(RuntimeError, match="pre-processed"):
+        method.process_weights_after_loading(layer)
+
+
 def test_modelopt_fp8_updates_weight_dims_after_transpose():
     """Humming reads weight.input_dim/output_dim. Swapping the
     ModelWeightParameter for a plain Parameter drops them, so the per-tensor
@@ -324,7 +371,7 @@ def test_modelopt_mixed_precision_composes_gemma4_mappers():
                 "quant_algo": "NVFP4",
                 "group_size": 16,
             },
-            "model.language_model.layers.1.moe.experts.gate_up_proj": {
+            "model.language_model.layers.1.experts.gate_up_proj": {
                 "quant_algo": "NVFP4",
                 "group_size": 16,
             },
@@ -336,10 +383,10 @@ def test_modelopt_mixed_precision_composes_gemma4_mappers():
     )
     config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_rename_mapper())
 
-    expected_prefix = "language_model.model.layers.0.moe.experts"
+    expected_prefix = "language_model.model.layers.0.experts"
     assert set(config.quantized_layers) == {
         expected_prefix,
-        "language_model.model.layers.1.moe.gate_up_proj",
+        "language_model.model.layers.1.experts.gate_up_proj",
     }
     assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
 
diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py
index 475ca84f6a17..989bf967fafd 100644
--- a/tests/quantization/test_moe_wna16.py
+++ b/tests/quantization/test_moe_wna16.py
@@ -93,6 +93,40 @@ def test_moe_wna16_accepts_channelwise_gptq_activation_order():
             False,
             "MoeWNA16 checkpoint layout",
         ),
+        (
+            WNA16MoEBackend.RDNA3,
+            AutoGPTQConfig(4, 128, False, True, False, {}, {}),
+            False,
+            False,
+            "compressed-tensors",
+        ),
+        (
+            WNA16MoEBackend.RDNA3,
+            QuantizationArgs(
+                num_bits=4,
+                type=QuantizationType.INT,
+                strategy=QuantizationStrategy.GROUP,
+                symmetric=False,
+                dynamic=False,
+                group_size=128,
+            ),
+            True,
+            False,
+            "asymmetric",
+        ),
+        (
+            WNA16MoEBackend.RDNA3,
+            QuantizationArgs(
+                num_bits=4,
+                type=QuantizationType.INT,
+                strategy=QuantizationStrategy.CHANNEL,
+                symmetric=True,
+                dynamic=False,
+            ),
+            False,
+            False,
+            "group-wise scales",
+        ),
     ],
 )
 def test_wna16_oracle_rejects_incompatible_quant_structures(
@@ -195,7 +229,7 @@ def test_moe_wna16_humming_adapter_repacks_uint8_tensors():
 
 
 def test_moe_wna16_uses_humming_quant_config(monkeypatch):
-    from vllm.model_executor.layers.quantization.utils import humming_utils
+    from vllm.model_executor.layers.quantization.utils import humming as humming_utils
 
     method = object.__new__(MoeWNA16Method)
     method.wna16_backend = WNA16MoEBackend.HUMMING
diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py
index fbd5cbc56392..4fc908f4840c 100644
--- a/tests/quantization/test_online.py
+++ b/tests/quantization/test_online.py
@@ -36,6 +36,7 @@
 )
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     LinearBase,
@@ -449,7 +450,6 @@ def test_online_shorthand_selects_checkpoint_or_online_config(
     tmp_path, quantization, checkpoint_quantization_config, is_config_file
 ) -> None:
     """Online shorthands defer to checkpoint metadata when available."""
-
     if is_config_file:
         (tmp_path / "hf_quant_config.json").write_text(
             '{"quant_method": "mxfp8", "ignored_layers": []}'
@@ -504,6 +504,76 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None:
         ModelConfig(model=str(tmp_path), quantization="fp8_per_channel")
 
 
+@pytest.mark.parametrize("per_token_activation", [False, True])
+@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
+def test_nvfp4_one_sided_sizes_dispatched_activations(
+    monkeypatch, per_token_activation: bool, input_dtype: torch.dtype
+) -> None:
+    from tests.kernels.moe.utils import make_dummy_moe_config
+    from vllm.model_executor.layers.fused_moe.all2all_utils import (
+        flashinfer_one_sided_dispatch_layout,
+    )
+    from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+    from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
+        TrtLlmNvFp4ExpertsModular,
+    )
+
+    config = make_dummy_moe_config(hidden_dim=6144, in_dtype=input_dtype)
+    config.moe_parallel_config.dp_size = 2
+    config.moe_parallel_config.ep_size = 2
+    config.moe_parallel_config.use_ep = True
+    config.moe_parallel_config.all2all_backend = "flashinfer_nvlink_one_sided"
+    expert_quant_config = FusedMoEQuantConfig.make("nvfp4")
+
+    def check_dispatch_layout(*, quant_config, input_dtype, **kwargs):
+        assert quant_config is expert_quant_config
+        layout = flashinfer_one_sided_dispatch_layout(
+            config.hidden_dim, quant_config, input_dtype=input_dtype
+        )
+        if per_token_activation:
+            assert input_dtype == config.in_dtype
+            assert (
+                layout.x_bytes_per_token == config.hidden_dim * config.in_dtype.itemsize
+            )
+            assert layout.x_sf_bytes_per_token == 0
+        else:
+            assert input_dtype is None
+            assert layout.x_bytes_per_token == config.hidden_dim // 2
+            assert layout.x_sf_bytes_per_token == config.hidden_dim // 16
+        assert expert_quant_config.quant_dtype == "nvfp4"
+        raise RuntimeError("dispatch layout verified before allocation")
+
+    monkeypatch.setattr(
+        nvfp4_oracle, "maybe_make_prepare_finalize", check_dispatch_layout
+    )
+    with pytest.raises(
+        RuntimeError, match="dispatch layout verified before allocation"
+    ):
+        nvfp4_oracle.make_nvfp4_moe_kernel(
+            moe_quant_config=expert_quant_config,
+            moe_config=config,
+            experts_cls=TrtLlmNvFp4ExpertsModular,
+            backend=nvfp4_oracle.NvFp4MoeBackend.FLASHINFER_TRTLLM,
+            per_token_activation=per_token_activation,
+        )
+
+
+@pytest.mark.parametrize(
+    "input_dtype",
+    [torch.float32, torch.float8_e4m3fn, torch.float8_e5m2, torch.int8],
+)
+def test_nvfp4_one_sided_rejects_unsupported_input_dtype(input_dtype: torch.dtype):
+    from vllm.model_executor.layers.fused_moe.all2all_utils import (
+        flashinfer_one_sided_dispatch_layout,
+    )
+    from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+
+    with pytest.raises(ValueError, match="unpacked inputs must be float16 or bfloat16"):
+        flashinfer_one_sided_dispatch_layout(
+            6144, FusedMoEQuantConfig.make("nvfp4"), input_dtype=input_dtype
+        )
+
+
 @pytest.mark.skipif(
     not is_quant_method_supported("fp8"),
     reason="FP8 is not supported on this GPU type.",
@@ -609,14 +679,12 @@ def test_online_quantization(
     dist_init,
     workspace_init,
 ) -> None:
-    """
-    Tests that online quantization frontend configuration works -
+    """Tests that online quantization frontend configuration works -
     selecting quant schemes, overriding quant schemes by type, ignoring
     layers.
 
     Does not test performance, peak memory usage, etc.
     """
-
     # TODO: Relax this condition once there is a native MXFP4_MXFP4
     # linear/moe backend supported on cuda.
     if quant_scheme == "mxfp4" and not (on_gfx950() or on_gfx942()):
diff --git a/tests/quantization/test_online_mxfp4.py b/tests/quantization/test_online_mxfp4.py
index c435790bce39..9495a15ba681 100644
--- a/tests/quantization/test_online_mxfp4.py
+++ b/tests/quantization/test_online_mxfp4.py
@@ -135,8 +135,7 @@ def _skip_reason_if_unavailable(backend: str, dtype: torch.dtype) -> str | None:
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
 @pytest.mark.parametrize("backend", ["triton", "aiter", "xpu", "quark"])
 def test_mxfp4_quantization_correctness(backend: str, dtype: torch.dtype):
-    """
-    Tests that the different implementations of mxfp4_quantize
+    """Tests that the different implementations of mxfp4_quantize
     in mxfp4_utils.py all match.
     """
     skip_reason = _skip_reason_if_unavailable(backend, dtype)
@@ -290,8 +289,7 @@ def test_online_mxfp4_moe_matches_quark(
     dist_init,
     monkeypatch,
 ):
-    """
-    Ensures `Mxfp4OnlineMoEMethod` (online quantization)
+    """Ensures `Mxfp4OnlineMoEMethod` (online quantization)
     and `QuarkOCP_MX_MoEMethod` (AMD Quark checkpoints) produce the same weights,
     with same MOE backend used.
 
@@ -495,8 +493,7 @@ def scatter_into_padded(
 def test_online_mxfp4_dense_matches_quark(
     linear_backend: str, default_vllm_config, dist_init, monkeypatch
 ):
-    """
-    Ensures `Mxfp4OnlineLinearMethod` (online quantization)
+    """Ensures `Mxfp4OnlineLinearMethod` (online quantization)
     and `QuarkOCP_MX` (AMD Quark checkpoints) produce the same weights,
     with same linear backend used.
     """
diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py
index ae236d71f805..f2f8ed6bd6b6 100644
--- a/tests/quantization/test_quark.py
+++ b/tests/quantization/test_quark.py
@@ -21,9 +21,18 @@
 
 from tests.quantization.utils import load_model_without_vllm_runner
 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
-from vllm.config import set_current_vllm_config
+from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.config.cache import CacheConfig
 from vllm.forward_context import set_forward_context
+from vllm.model_executor import parameter
+from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
+    AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
+    AiterPerTokenFp8ScaledMMLinearKernel,
+    AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
+    FP8ScaledMMLinearLayerConfig,
+)
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoeWeightScaleSupported,
@@ -49,12 +58,17 @@
 from vllm.model_executor.layers.quantization.quark.quark_moe import (  # noqa: E501
     QuarkMoEMethod,
     QuarkW4A8Fp8MoEMethod,
+    QuarkW4A16Int4MoEMethod,
     QuarkW8A8Fp8MoEMethod,
     QuarkW8A8Int8MoEMethod,
 )
-from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
+from vllm.model_executor.layers.quantization.quark.schemes import (
+    QuarkScheme,
+    QuarkW4A16Int4,
+)
 from vllm.model_executor.layers.quantization.quark.utils import (
     QuarkQTensorHint,
+    canonicalize_quark_packed_int4,
     should_ignore_layer,
 )
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
@@ -88,6 +102,7 @@
     kNvfp4Static,
 )
 from vllm.model_executor.models.llama import LlamaForCausalLM
+from vllm.model_executor.models.utils import WeightsMapper
 from vllm.platforms import current_platform
 from vllm.transformers_utils.repo_utils import hf_api
 
@@ -113,6 +128,12 @@ def on_gfx950() -> bool:
 
 AITER_AVAILABLE = is_aiter_found_and_supported()
 
+AITER_PTPC_KERNELS = (
+    AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
+    AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
+    AiterPerTokenFp8ScaledMMLinearKernel,
+)
+
 DEVICE_TYPE = current_platform.device_type
 
 
@@ -834,7 +855,7 @@ def __init__(self):
         UnquantizedFusedMoEMethod,
     )
 
-    dynamic_mxfp4_config = _make_qtensor_config(
+    mxfp4_config = _make_qtensor_config(
         {
             "dtype": "fp4",
             "qscheme": "per_group",
@@ -845,24 +866,23 @@ def __init__(self):
         None,
         exclude=["self_attn.q_proj", "mlp.down_proj"],
     )
-    dynamic_mxfp4_config.dynamic_mxfp4_quant = True
-
-    assert dynamic_mxfp4_config.get_quant_method_target(
-        "self_attn.q_proj", LinearBase
-    ) == (kMxfp4Static, None, QuarkLinearMethod)
-    attention_proj = TestLinear()
+    assert mxfp4_config.get_quant_method_target("self_attn.q_proj", LinearBase) == (
+        None,
+        None,
+        UnquantizedLinearMethod,
+    )
     assert isinstance(
-        dynamic_mxfp4_config.get_quant_method(attention_proj, "self_attn.q_proj"),
-        QuarkLinearMethod,
+        mxfp4_config.get_quant_method(TestLinear(), "self_attn.q_proj"),
+        UnquantizedLinearMethod,
     )
-    assert isinstance(attention_proj.scheme, QuarkOCP_MX)
-    assert attention_proj.scheme.dynamic_mxfp4_quant
 
-    assert dynamic_mxfp4_config.get_quant_method_target(
-        "mlp.down_proj", LinearBase
-    ) == (None, None, UnquantizedLinearMethod)
+    assert mxfp4_config.get_quant_method_target("mlp.down_proj", LinearBase) == (
+        None,
+        None,
+        UnquantizedLinearMethod,
+    )
     assert isinstance(
-        dynamic_mxfp4_config.get_quant_method(TestLinear(), "mlp.down_proj"),
+        mxfp4_config.get_quant_method(TestLinear(), "mlp.down_proj"),
         UnquantizedLinearMethod,
     )
 
@@ -1095,6 +1115,54 @@ def test_quark_w8a8_fp8_moe_per_block_rejects_misaligned_partition():
         )
 
 
+def test_quark_fp8_ptpc_exposes_kernel_input_quant_key(monkeypatch):
+    """QuarkW8A8Fp8 must advertise the key its kernel consumes pre-quantized.
+
+    Only kernel selection runs, no GEMM. The shape matters: AITER PTPC
+    kernels decline untuned shapes, leaving a torch kernel that has no key.
+    """
+    # Llama-3.1-70B qkv_proj at TP1 (aiter-tuned N, K on gfx950).
+    N, K = 10240, 8192
+    dtype = torch.bfloat16
+    kernel_config = FP8ScaledMMLinearLayerConfig(
+        weight_quant_key=kFp8StaticChannelSym,
+        activation_quant_key=kFp8DynamicTokenSym,
+        weight_shape=(N, K),
+        input_dtype=dtype,
+        out_dtype=dtype,
+    )
+    if not any(
+        cls.is_supported()[0] and cls.can_implement(kernel_config)[0]
+        for cls in AITER_PTPC_KERNELS
+    ):
+        pytest.skip("no AITER PTPC kernel is usable for this shape and environment")
+
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
+
+    scheme = QuarkW8A8Fp8.__new__(QuarkW8A8Fp8)
+    scheme.weight_qscheme = "per_channel"
+    scheme.is_static_input_scheme = False
+    scheme.input_qscheme = "per_channel"
+    scheme.activation_quant_key = kFp8DynamicTokenSym
+    scheme.weight_quant_key = kFp8StaticChannelSym
+    scheme.out_dtype = dtype
+    scheme.input_dtype = dtype
+
+    layer = torch.nn.Module()
+    with set_current_vllm_config(VllmConfig()):
+        scheme.create_weights(
+            layer,
+            output_partition_sizes=[N],
+            input_size_per_partition=K,
+            params_dtype=dtype,
+            weight_loader=lambda *args, **kwargs: None,
+        )
+
+    assert isinstance(scheme.fp8_linear, AITER_PTPC_KERNELS)
+    assert layer.input_quant_key == kFp8DynamicTokenSym
+
+
 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
 def test_quark_fp8_w_per_tensor_a_per_tensor(
     kv_cache_dtype: str, monkeypatch, dist_init, workspace_init
@@ -1668,3 +1736,580 @@ def test_suffix_match_at_module_boundary(prefix, ignored_layer, expected):
         )
         is expected
     )
+
+
+_GLM5_MXFP4_WEIGHT = {
+    "dtype": "fp4",
+    "qscheme": "per_group",
+    "group_size": 32,
+    "scale_format": "e8m0",
+    "is_dynamic": False,
+}
+_GLM5_BLOCK_FP8_WEIGHT = {
+    "dtype": "fp8_e4m3",
+    "qscheme": "per_block",
+    "is_dynamic": False,
+    "block_size": [128, 128],
+    "symmetric": True,
+}
+_GLM5_BLOCK_FP8_INPUT = {
+    "dtype": "fp8_e4m3",
+    "qscheme": "per_group",
+    "is_dynamic": True,
+    "group_size": 128,
+    "symmetric": True,
+}
+_GLM5_GATE_UP = "language_model.model.layers.0.mlp.gate_up_proj"
+
+
+def _glm5_mixed_precision_config() -> QuarkConfig:
+    """Global MXFP4 with the dense-MLP gate/up shards marked block-FP8."""
+    fp8 = {"weight": _GLM5_BLOCK_FP8_WEIGHT, "input_tensors": _GLM5_BLOCK_FP8_INPUT}
+    return QuarkConfig(
+        {
+            "global_quant_config": {
+                "weight": _GLM5_MXFP4_WEIGHT,
+                "input_tensors": None,
+            },
+            "layer_type_quant_config": {},
+            "layer_quant_config": {"*mlp.gate_proj": fp8, "*mlp.up_proj": fp8},
+            "exclude": [],
+        }
+    )
+
+
+class _GLM5RecordingParam:
+    """Stand-in for a BF16 target param that records what gets loaded."""
+
+    def __init__(self):
+        self.loaded_weight: torch.Tensor | None = None
+        self.loaded_shard: object = "unset"
+
+    def weight_loader(self, param, loaded_weight, shard_id=None):
+        self.loaded_weight = loaded_weight
+        self.loaded_shard = shard_id
+
+
+def _glm5_block_fp8(out_dim: int, in_dim: int, block: int = 128):
+    """Return a (fp8 weight, f32 per-block scale) pair."""
+    weight = (torch.randn(out_dim, in_dim) * 0.1).to(torch.float8_e4m3fn)
+    scale = torch.rand(out_dim // block, in_dim // block, dtype=torch.float32) + 0.5
+    return weight, scale
+
+
+def test_glm5next_gate_up_proj_mapping_is_genuine_fusion():
+    from vllm.models.glm5next.common.model import Glm5NextForConditionalGeneration
+
+    mapping = Glm5NextForConditionalGeneration.packed_modules_mapping
+    assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"]
+
+
+def test_glm5next_genuine_fusion_resolves_gate_up_to_block_fp8():
+    # The fused module must expand to its real shard names so each resolves to
+    # the per-layer block-FP8 entry (rather than the global MXFP4 scheme).
+    config = _glm5_mixed_precision_config()
+    config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
+    _, _, scheme_cls = config.get_scheme_cls(LinearBase, _GLM5_GATE_UP)
+    assert scheme_cls is QuarkW8A8Fp8PerBlock
+
+
+def test_glm5next_attn_loader_accepts_quark_weight_scale():
+    # Regression for KeyError on '...kv_a_proj_with_mqa.weight_scale': the fused
+    # q_a/kv_a projection is kept BF16 and dequantized on load; the Quark scale
+    # name must be recognized and routed to the correct fused shard.
+    from vllm.models.glm5next.common.model import (
+        _dequant_fp8_block,
+        _try_load_fp8_attn_proj,
+    )
+
+    prefix = "layers.3.self_attn"
+    target = _GLM5RecordingParam()
+    params_dict = {f"{prefix}.fused_qkv_a_proj.weight": target}
+    buf: dict = {}
+    loaded: set = set()
+
+    weight, scale = _glm5_block_fp8(256, 256)
+    # fp8 weight arrives first -> buffered, nothing loaded yet.
+    assert (
+        _try_load_fp8_attn_proj(
+            f"{prefix}.kv_a_proj_with_mqa.weight", weight, buf, params_dict, loaded, 0
+        )
+        is True
+    )
+    assert target.loaded_weight is None
+    # Quark-named scale completes the pair -> dequantize + load.
+    assert (
+        _try_load_fp8_attn_proj(
+            f"{prefix}.kv_a_proj_with_mqa.weight_scale",
+            scale,
+            buf,
+            params_dict,
+            loaded,
+            0,
+        )
+        is True
+    )
+    assert target.loaded_weight is not None
+    assert target.loaded_weight.dtype == torch.bfloat16
+    assert target.loaded_shard == 1  # kv_a is shard 1 of fused_qkv_a_proj
+    assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128))
+    assert f"{prefix}.fused_qkv_a_proj.weight" in loaded
+
+
+def test_glm5next_attn_loader_accepts_deepseek_weight_scale_inv():
+    # DeepSeek scale name keeps working unchanged.
+    from vllm.models.glm5next.common.model import (
+        _dequant_fp8_block,
+        _try_load_fp8_attn_proj,
+    )
+
+    prefix = "layers.7.self_attn"
+    target = _GLM5RecordingParam()
+    params_dict = {f"{prefix}.o_proj.weight": target}
+    buf: dict = {}
+    loaded: set = set()
+
+    weight, scale = _glm5_block_fp8(128, 256)
+    _try_load_fp8_attn_proj(
+        f"{prefix}.o_proj.weight", weight, buf, params_dict, loaded, 0
+    )
+    assert (
+        _try_load_fp8_attn_proj(
+            f"{prefix}.o_proj.weight_scale_inv", scale, buf, params_dict, loaded, 0
+        )
+        is True
+    )
+    assert target.loaded_shard is None  # o_proj is a direct (non-fused) proj
+    assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128))
+
+
+_REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
+
+
+def _quark_int4_config(
+    *,
+    pack_method: str = "reorder",
+    symmetric: bool = True,
+    exclude: list[str] | None = None,
+) -> dict:
+    return {
+        "quant_method": "quark",
+        "export": {"pack_method": pack_method, "kv_cache_group": []},
+        "global_quant_config": {
+            "weight": {
+                "dtype": "int4",
+                "group_size": 128,
+                "symmetric": symmetric,
+            }
+        },
+        "exclude": exclude or [],
+    }
+
+
+def _sign_extend_int4_nibbles(t: torch.Tensor) -> torch.Tensor:
+    mask = (t & 0x8).bool()
+    t = t.clone()
+    t[mask] = t[mask] | 0xF0
+    return t
+
+
+def _dequantize_quark_signed_awq_torch(
+    qweight: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor,
+    group_size: int,
+    *,
+    pack_reorder: bool = True,
+) -> torch.Tensor:
+    bits = 4
+    shifts = torch.arange(0, 32, bits, device=qweight.device)
+    iweights = ((qweight[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8)
+    iweights = iweights.view(qweight.shape[0], -1)
+    zeros = ((qzeros[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8)
+    zeros = zeros.view(qzeros.shape[0], -1)
+
+    if pack_reorder:
+        order = torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=qweight.device)
+    else:
+        order = torch.arange(8, device=qweight.device)
+    iweights = iweights.view(qweight.shape[0], -1, 8)[:, :, order].reshape(
+        qweight.shape[0], -1
+    )
+    zeros = zeros.view(qzeros.shape[0], -1, 8)[:, :, order].reshape(qzeros.shape[0], -1)
+    iweights = _sign_extend_int4_nibbles(iweights & 0xF)
+    zeros = _sign_extend_int4_nibbles(zeros & 0xF)
+
+    scales = scales.repeat_interleave(group_size, dim=0)
+    zeros = zeros.repeat_interleave(group_size, dim=0)
+    return (iweights - zeros) * scales
+
+
+def _dequantize_awq_unsigned_torch(
+    qweight: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor,
+    group_size: int,
+    *,
+    pack_reorder: bool = True,
+) -> torch.Tensor:
+    bits = 4
+    shifts = torch.arange(0, 32, bits, device=qweight.device)
+    iweights = ((qweight[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8)
+    iweights = iweights.view(qweight.shape[0], -1)
+    zeros = ((qzeros[:, :, None] >> shifts[None, None, :]) & 0xF).to(torch.int8)
+    zeros = zeros.view(qzeros.shape[0], -1)
+
+    if pack_reorder:
+        order = torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=qweight.device)
+    else:
+        order = torch.arange(8, device=qweight.device)
+    iweights = iweights.view(qweight.shape[0], -1, 8)[:, :, order].reshape(
+        qweight.shape[0], -1
+    )
+    zeros = zeros.view(qzeros.shape[0], -1, 8)[:, :, order].reshape(qzeros.shape[0], -1)
+
+    scales = scales.repeat_interleave(group_size, dim=0)
+    zeros = zeros.repeat_interleave(group_size, dim=0)
+    return (iweights - zeros) * scales
+
+
+def _pack_int4_nibbles(nibbles: torch.Tensor, *, pack_reorder: bool) -> torch.Tensor:
+    pack_order = (
+        torch.tensor(_REVERSE_AWQ_PACK_ORDER, device=nibbles.device)
+        if pack_reorder
+        else torch.arange(8, device=nibbles.device)
+    )
+    shifts = pack_order * 4
+    return ((nibbles.to(torch.int64) & 0xF) << shifts).sum(dim=-1).to(torch.int32)
+
+
+class TestQuarkInt4Format:
+    """Tests for Quark INT4 export format compatibility."""
+
+    def test_quark_order_pack_method_uses_native_int4_scheme(self):
+        quant_config = QuarkConfig.from_config(_quark_int4_config(pack_method="order"))
+        weight_config = quant_config.quant_config["global_quant_config"]
+        weight_key, act_key, scheme_cls = quant_config._get_scheme_cls_from_config(
+            weight_config
+        )
+        assert scheme_cls is QuarkW4A16Int4
+
+        scheme = quant_config.init_scheme(
+            scheme_cls,
+            weight_quant_key=weight_key,
+            activation_quant_key=act_key,
+            weight_config=weight_config["weight"],
+        )
+        assert isinstance(scheme, QuarkW4A16Int4)
+        assert not scheme.pack_reorder
+
+    def test_quark_int4_scheme_supports_asymmetric_weights(self):
+        quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False))
+        weight_config = quant_config.quant_config["global_quant_config"]
+        weight_key, act_key, scheme_cls = quant_config._get_scheme_cls_from_config(
+            weight_config
+        )
+        assert scheme_cls is QuarkW4A16Int4
+
+        scheme = quant_config.init_scheme(
+            scheme_cls,
+            weight_quant_key=weight_key,
+            activation_quant_key=act_key,
+            weight_config=weight_config["weight"],
+        )
+        assert isinstance(scheme, QuarkW4A16Int4)
+        assert not scheme.is_symmetric
+
+    @pytest.mark.parametrize("missing_field", ["group_size", "symmetric"])
+    def test_quark_int4_scheme_requires_weight_config_fields(self, missing_field):
+        from vllm.model_executor.layers.quantization.utils.quant_utils import (
+            kInt4Static,
+        )
+
+        weight_config = {
+            "dtype": "int4",
+            "group_size": 128,
+            "symmetric": True,
+        }
+        weight_config.pop(missing_field)
+        quant_config = QuarkConfig.from_config(
+            {
+                "quant_method": "quark",
+                "export": {"pack_method": "reorder", "kv_cache_group": []},
+                "global_quant_config": {"weight": weight_config},
+                "exclude": [],
+            }
+        )
+
+        with pytest.raises(ValueError, match=missing_field):
+            quant_config.init_scheme(
+                QuarkW4A16Int4,
+                weight_quant_key=kInt4Static,
+                activation_quant_key=None,
+                weight_config=weight_config,
+            )
+
+    def test_quark_int4_moe_uses_native_moe_method(self):
+        from unittest.mock import MagicMock, patch
+
+        from vllm.model_executor.layers.quantization.utils.quant_utils import (
+            kInt4StaticAsym,
+        )
+
+        quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False))
+        moe_config = type("MoeConfig", (), {})()
+        moe_config.has_bias = False
+
+        mock_backend = MagicMock()
+        mock_experts_cls = MagicMock()
+        with patch(
+            "vllm.model_executor.layers.quantization.quark.quark_moe"
+            ".select_wna16_moe_backend",
+            return_value=(mock_backend, mock_experts_cls),
+        ):
+            method = QuarkW4A16Int4MoEMethod(
+                kInt4StaticAsym,
+                None,
+                quant_config.quant_config["global_quant_config"]["weight"],
+                quant_config.pack_method,
+                moe_config,
+            )
+
+        assert method.group_size == 128
+        assert method.pack_reorder
+        assert method.use_wna16_backend
+
+    def test_triton_wna16_experts_supports_asymmetric_int4(self):
+        from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
+            TritonWNA16Experts,
+        )
+        from vllm.model_executor.layers.quantization.utils.quant_utils import (
+            kInt4Static,
+            kInt4Static32,
+            kInt4Static32Asym,
+            kInt4StaticAsym,
+        )
+
+        for key in [kInt4Static, kInt4Static32, kInt4StaticAsym, kInt4Static32Asym]:
+            assert TritonWNA16Experts._supports_quant_scheme(key, None), (
+                f"{key} should be supported by TritonWNA16Experts"
+            )
+
+    def test_quark_moe_loader_shards_w2_zero_point_across_tp_ranks(self):
+        """w2 zero-points are sharded along the intermediate dim per TP rank.
+
+        Loading the same Quark checkpoint tensor at tensor_parallel_size=2 for
+        both ranks must jointly reconstruct the tensor_parallel_size=1 (full)
+        result, with no overlap or gap. This guards the custom loader's TP
+        split for the packed zero-point layout.
+        """
+
+        class _FakeMoEConfig:
+            has_bias = False
+
+            def __init__(self, tp_size, tp_rank):
+                self.tp_size = tp_size
+                self.tp_rank = tp_rank
+
+        class _FakeLayer:
+            def __init__(self, moe_config):
+                self.moe_config = moe_config
+                self.intermediate_size_per_partition = 64
+                self.group_size_div_factor = 1
+
+        def load_w2_zero_point(raw_zp, tp_size, tp_rank):
+            from unittest.mock import MagicMock, patch
+
+            from vllm.model_executor.layers.quantization.utils.quant_utils import (
+                kInt4StaticAsym,
+            )
+
+            quant_config = QuarkConfig.from_config(_quark_int4_config(symmetric=False))
+            moe_config = _FakeMoEConfig(tp_size, tp_rank)
+            with patch(
+                "vllm.model_executor.layers.quantization.quark.quark_moe"
+                ".select_wna16_moe_backend",
+                return_value=(MagicMock(), MagicMock()),
+            ):
+                method = QuarkW4A16Int4MoEMethod(
+                    kInt4StaticAsym,
+                    None,
+                    quant_config.quant_config["global_quant_config"]["weight"],
+                    quant_config.pack_method,
+                    moe_config,
+                )
+
+            layer = _FakeLayer(moe_config)
+            # The loader converts the packed zero-point, then writes the
+            # per-rank slice into param.data[expert_id]. A single expert dim is
+            # enough; the last dim shrinks by tp_size after sharding.
+            param = torch.nn.Parameter(
+                torch.zeros(1, 16, 8 // tp_size, dtype=torch.uint8),
+                requires_grad=False,
+            )
+            loader = method.get_weight_loader(layer, weight_loader=None)
+            loader(
+                param,
+                raw_zp.clone(),
+                weight_name="w2_weight_zero_point",
+                shard_id="w2",
+                expert_id=0,
+            )
+            return param.data[0]
+
+        # Quark-packed zero-point for one expert: rows = hidden // pack_factor,
+        # cols carry the (packed) intermediate dim, split evenly across ranks.
+        raw_zp = torch.randint(0, 256, (8, 16), dtype=torch.uint8)
+
+        full = load_w2_zero_point(raw_zp, tp_size=1, tp_rank=0)
+        shard0 = load_w2_zero_point(raw_zp, tp_size=2, tp_rank=0)
+        shard1 = load_w2_zero_point(raw_zp, tp_size=2, tp_rank=1)
+
+        expected = full.view(full.size(0), 2, -1)
+        torch.testing.assert_close(shard0, expected[:, 0])
+        torch.testing.assert_close(shard1, expected[:, 1])
+
+    def test_quark_shared_expert_gate_keeps_quantized_tensors(self):
+        quant_config = QuarkConfig.from_config(_quark_int4_config())
+        mapper = quant_config.get_cache_scale_mapper()
+
+        output_names = {
+            name
+            for name, _ in mapper.apply(
+                [
+                    (
+                        "model.language_model.layers.0.mlp.shared_expert_gate.weight",
+                        torch.zeros(1),
+                    ),
+                    (
+                        "model.language_model.layers.0.mlp.shared_expert_gate"
+                        ".weight_scale",
+                        torch.zeros(1),
+                    ),
+                    (
+                        "model.language_model.layers.0.mlp.shared_expert_gate"
+                        ".weight_zero_point",
+                        torch.zeros(1),
+                    ),
+                ]
+            )
+        }
+
+        assert (
+            "model.language_model.layers.0.mlp.shared_expert_gate.weight"
+            in output_names
+        )
+        assert (
+            "model.language_model.layers.0.mlp.shared_expert_gate.weight_scale"
+            in output_names
+        )
+        assert (
+            "model.language_model.layers.0.mlp.shared_expert_gate.weight_zero_point"
+            in output_names
+        )
+
+    def test_quark_apply_mapper_updates_exclude_and_layer_quant_config(self):
+        quant_config = QuarkConfig.from_config(
+            {
+                **_quark_int4_config(),
+                "exclude": ["lm_head"],
+                "layer_quant_config": {
+                    "model.language_model.layers.0.mlp.gate": {
+                        "weight": {"dtype": "float16"},
+                    },
+                },
+            }
+        )
+        quant_config.apply_vllm_mapper(
+            WeightsMapper(
+                orig_to_new_prefix={
+                    "lm_head": "language_model.lm_head",
+                    "model.language_model.": "language_model.model.",
+                },
+            )
+        )
+
+        layer_quant_config = quant_config.quant_config["layer_quant_config"]
+        assert "language_model.model.layers.0.mlp.gate" in layer_quant_config
+        assert "model.language_model.layers.0.mlp.gate" not in layer_quant_config
+        assert quant_config.quant_config["exclude"] == [
+            "language_model.lm_head",
+        ]
+
+    def test_quark_apply_mapper_ignores_non_string_list_entries(self):
+        quant_config = QuarkConfig.from_config(
+            {
+                **_quark_int4_config(),
+                "algo_config": [
+                    {"name": "qronos", "inside_layer_modules": ["self_attn.q_proj"]}
+                ],
+            }
+        )
+        quant_config.apply_vllm_mapper(WeightsMapper())
+
+        assert quant_config.quant_config["algo_config"] == [
+            {"name": "qronos", "inside_layer_modules": ["self_attn.q_proj"]}
+        ]
+
+    def test_quark_exclude_matches_via_mapper_prefix(self):
+        quant_config = QuarkConfig.from_config(_quark_int4_config(exclude=["lm_head"]))
+        quant_config.apply_vllm_mapper(
+            WeightsMapper(orig_to_new_prefix={"lm_head": "language_model.lm_head"})
+        )
+
+        exclude_layers = quant_config.quant_config["exclude"]
+        assert should_ignore_layer("language_model.lm_head", ignore=exclude_layers)
+        assert not should_ignore_layer(
+            "language_model.model.layers.0.mlp.gate", ignore=exclude_layers
+        )
+
+
+@pytest.mark.parametrize("symmetric", [False, True])
+@pytest.mark.parametrize("pack_method", ["order", "reorder"])
+def test_quark_int4_canonicalizes_pack_for_kernel_layout(pack_method, symmetric):
+    """Quark pack order is normalized to the layout expected by awq_* ops.
+
+    This reuses the existing AWQ dequant/gemm kernels for compute only; loading
+    still goes through the native Quark quantization path, not AutoAWQ.
+    """
+    pack_reorder = pack_method == "reorder"
+    group_size = 2
+    packed_values = torch.tensor(
+        [
+            [0, 1, 7, 8, 9, 15, 2, 14],
+            [15, 8, 0, 3, 12, 7, 1, 9],
+        ],
+        dtype=torch.int32,
+    )
+    packed_zeros = torch.zeros((1, 8), dtype=torch.int32)
+    qweight = _pack_int4_nibbles(packed_values, pack_reorder=pack_reorder).view(2, 1)
+    qzeros = _pack_int4_nibbles(packed_zeros, pack_reorder=pack_reorder).view(1, 1)
+    scales = torch.ones((1, 8), dtype=torch.float16)
+
+    dequantize = (
+        _dequantize_quark_signed_awq_torch
+        if symmetric
+        else _dequantize_awq_unsigned_torch
+    )
+    expected = dequantize(
+        qweight,
+        scales,
+        qzeros,
+        group_size,
+        pack_reorder=pack_reorder,
+    )
+    canonical_weight = canonicalize_quark_packed_int4(
+        qweight,
+        pack_reorder=pack_reorder,
+        is_symmetric=symmetric,
+    )
+    canonical_zero = canonicalize_quark_packed_int4(
+        qzeros,
+        pack_reorder=pack_reorder,
+        is_symmetric=symmetric,
+    )
+    actual = _dequantize_awq_unsigned_torch(
+        canonical_weight, scales, canonical_zero, group_size
+    )
+
+    assert torch.equal(actual, expected)
diff --git a/tests/quantization/test_register_quantization_config.py b/tests/quantization/test_register_quantization_config.py
index 5529a65b778f..227c604ccf87 100644
--- a/tests/quantization/test_register_quantization_config.py
+++ b/tests/quantization/test_register_quantization_config.py
@@ -48,7 +48,6 @@ def apply(
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Perform fake quantization before the linear layer."""
-
         # Calculate the scales dynamically
         max_val = torch.amax(x, dim=(0, -1), keepdims=True)
         min_val = torch.amin(x, dim=(0, -1), keepdims=True)
@@ -108,7 +107,6 @@ def get_quant_method(
 
 def test_register_quantization_config(caplog_vllm):
     """Test register custom quantization config."""
-
     # The quantization method `custom_quant` should be registered.
     assert get_quantization_config("custom_quant") == CustomQuantConfig
 
diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py
index 50817ce3303d..3c839a687e3a 100644
--- a/tests/quantization/test_torchao.py
+++ b/tests/quantization/test_torchao.py
@@ -9,7 +9,6 @@
 
 from tests.quantization.utils import load_model_without_vllm_runner
 from vllm.config import set_current_vllm_config
-from vllm.config.load import LoadConfig
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.model_loader import get_model_loader
@@ -67,65 +66,6 @@ def test_pre_quantized_model(monkeypatch, dist_init, workspace_init):
     assert torch.isfinite(logits).all()
 
 
-@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
-@pytest.mark.parametrize(
-    "pt_load_map_location",
-    [
-        f"{DEVICE_TYPE}:0",
-        # {"": "cuda"},
-    ],
-)
-def test_opt_125m_int8wo_model_loading_with_params(
-    pt_load_map_location, monkeypatch, dist_init, workspace_init
-):
-    torch._dynamo.reset()
-    model_name = "jerryzh168/opt-125m-int8wo-partial-quant"
-    model, vllm_config = load_model_without_vllm_runner(
-        model_name,
-        quantization="torchao",
-        vllm_config_kwargs={
-            "load_config": LoadConfig(pt_load_map_location=pt_load_map_location)
-        },
-    )
-
-    monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
-    input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
-    positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
-    with (
-        set_current_vllm_config(vllm_config),
-        set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
-    ):
-        hidden_states = model(input_ids, positions, None)
-        logits = model.compute_logits(hidden_states)
-    assert torch.isfinite(logits).all()
-
-
-@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
-def test_qwenvl_int8wo_model_loading_with_params(
-    monkeypatch, dist_init, workspace_init
-):
-    torch._dynamo.reset()
-    model_name = "mobicham/Qwen2.5-VL-3B-Instruct_int8wo_ao"
-    model, vllm_config = load_model_without_vllm_runner(
-        model_name,
-        quantization="torchao",
-        vllm_config_kwargs={
-            "load_config": LoadConfig(pt_load_map_location=f"{DEVICE_TYPE}:0")
-        },
-    )
-
-    monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
-    input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
-    positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE).expand(3, -1)
-    with (
-        set_current_vllm_config(vllm_config),
-        set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
-    ):
-        hidden_states = model(input_ids, positions, None)
-        logits = model.compute_logits(hidden_states)
-    assert torch.isfinite(logits).all()
-
-
 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
 @pytest.mark.skip(
     reason="since torchao nightly is only compatible with torch nightly"
diff --git a/tests/reasoning/test_base_thinking_reasoning_parser.py b/tests/reasoning/test_base_thinking_reasoning_parser.py
index f4d74ceeec0a..fb506860621c 100644
--- a/tests/reasoning/test_base_thinking_reasoning_parser.py
+++ b/tests/reasoning/test_base_thinking_reasoning_parser.py
@@ -51,8 +51,7 @@ def test_tokenizer():
 
 
 class TestBaseThinkingReasoningParserInit:
-    """
-    Test initialization and basic properties of
+    """Test initialization and basic properties of
     BaseThinkingReasoningParser.
     """
 
@@ -274,8 +273,7 @@ class TestBaseThinkingReasoningParserStreaming:
 
     @pytest.mark.parametrize("streaming", [True, False])
     def test_simple_reasoning_extraction(self, test_tokenizer, streaming):
-        """
-        Test basic reasoning extraction in both
+        """Test basic reasoning extraction in both
         streaming and non-streaming modes.
         """
         parser = TestThinkingReasoningParser(test_tokenizer)
@@ -369,14 +367,12 @@ def test_streaming_only_end_token(self, test_tokenizer):
 
 
 class TestBaseThinkingReasoningParserMultipleImplementations:
-    """
-    Test that multiple implementations of
+    """Test that multiple implementations of
     BaseThinkingReasoningParser work correctly.
     """
 
     def test_different_token_implementations(self, test_tokenizer):
-        """
-        Test that different implementations
+        """Test that different implementations
         with different tokens work independently.
         """
         parser1 = TestThinkingReasoningParser(test_tokenizer)
diff --git a/tests/reasoning/test_cohere_command_reasoning_parser.py b/tests/reasoning/test_cohere_command_reasoning_parser.py
deleted file mode 100644
index ea9ab63aef72..000000000000
--- a/tests/reasoning/test_cohere_command_reasoning_parser.py
+++ /dev/null
@@ -1,796 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import json
-from collections import UserDict
-from dataclasses import dataclass
-from types import SimpleNamespace
-from typing import Any
-
-import pytest
-from pydantic import ValidationError
-
-from vllm.entrypoints.generate.base.protocol import (
-    JsonSchemaResponseFormat,
-    ResponseFormat,
-    StructuralTagResponseFormat,
-)
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionRequest,
-)
-from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-from vllm.reasoning.cohere_command_reasoning_parser import (
-    CohereCommand3ReasoningParser,
-    CohereCommand4ReasoningParser,
-    _has_effective_tools,
-    _melody_citations_to_vllm,
-    _melody_sources_to_vllm,
-    _response_format_type,
-    _schema_dict_from_structured_outputs,
-    convert_schema_to_structural_tags,
-)
-from vllm.sampling_params import StructuredOutputsParams
-
-
-@dataclass
-class ReasoningCase:
-    parser_cls: Any
-    model_output: str
-    expected_reasoning: str | None
-    expected_content: str | None
-
-
-REASONING_CASES = [
-    pytest.param(
-        ReasoningCase(
-            parser_cls=CohereCommand3ReasoningParser,
-            model_output="""\
-<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
-[
-    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
-]
-<|END_ACTION|>""",
-            expected_reasoning="i will call foo with query1",
-            expected_content="""\
-<|START_ACTION|>
-[
-    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
-]
-<|END_ACTION|>""",
-        ),
-        id="cmd3-single_tool_call",
-    ),
-    pytest.param(
-        ReasoningCase(
-            parser_cls=CohereCommand4ReasoningParser,
-            model_output="""\
-<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
-[
-    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
-]
-<|END_ACTION|>""",
-            expected_reasoning="i will call foo with query1",
-            expected_content="""\
-<|START_ACTION|>
-[
-    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
-]
-<|END_ACTION|>""",
-        ),
-        id="cmd4-single_tool_call",
-    ),
-    pytest.param(
-        ReasoningCase(
-            parser_cls=CohereCommand3ReasoningParser,
-            model_output="""\
-<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
-<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
-            expected_reasoning="This is a rainbow emoji: 🌈",
-            expected_content="foo bar",
-        ),
-        id="cmd3-citations_with_emoji",
-    ),
-    pytest.param(
-        ReasoningCase(
-            parser_cls=CohereCommand4ReasoningParser,
-            model_output="""\
-<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
-<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
-            expected_reasoning="This is a rainbow emoji: 🌈",
-            expected_content="foo bar",
-        ),
-        id="cmd4-citations_with_emoji",
-    ),
-]
-
-
-class MockCohereTokenizer:
-    """Minimal byte-level stand-in for the Cohere tokenizer.
-
-    ``encode``/``decode`` round-trip through UTF-8 bytes so splitting a
-    multi-byte character (e.g. an emoji) across "tokens" reproduces the
-    trailing U+FFFD buffering that real streaming exhibits. Cohere special
-    tokens map to distinct synthetic ids; everything else shares a default id.
-    ``adjust_request`` only needs the token ids, not real tokenization.
-    """
-
-    _SPECIAL_TOKEN_IDS = {
-        "<|START_THINKING|>": -1,
-        "<|END_THINKING|>": -2,
-        "<|CHATBOT_TOKEN|>": -3,
-    }
-
-    def convert_tokens_to_ids(self, token: str) -> int:
-        return self._SPECIAL_TOKEN_IDS.get(token, 0)
-
-    def get_vocab(self) -> dict[str, int]:
-        return {}
-
-    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
-        return list(text.encode("utf-8"))
-
-    def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
-        return bytes(ids).decode("utf-8", errors="replace")
-
-
-@pytest.fixture(scope="module")
-def tokenizer() -> MockCohereTokenizer:
-    return MockCohereTokenizer()
-
-
-@pytest.fixture
-def request_obj():
-    return ChatCompletionRequest(messages=[], model="test-model")
-
-
-REPLACEMENT_CHAR = "\ufffd"
-
-
-def _token_deltas(tokenizer, text: str) -> list[str]:
-    """Progressively decode the token sequence and return per-step string
-    deltas.  Incomplete multi-byte sequences (trailing U+FFFD) are buffered
-    until the next token completes them, matching real streaming behaviour."""
-    ids = tokenizer.encode(text, add_special_tokens=False)
-    deltas: list[str] = []
-    prev = ""
-    for i in range(1, len(ids) + 1):
-        current = tokenizer.decode(ids[:i], skip_special_tokens=False)
-        if current.endswith(REPLACEMENT_CHAR):
-            continue
-        delta = current[len(prev) :]
-        if delta:
-            deltas.append(delta)
-        prev = current
-    return deltas
-
-
-@pytest.mark.parametrize("case", REASONING_CASES)
-class TestExtractReasoning:
-    def test_nonstreaming(self, tokenizer, request_obj, case: ReasoningCase):
-        parser = case.parser_cls(tokenizer)
-        reasoning, content = parser.extract_reasoning(case.model_output, request_obj)
-
-        assert reasoning == case.expected_reasoning
-        assert content == case.expected_content
-
-    def test_streaming(self, tokenizer, case: ReasoningCase):
-        parser = case.parser_cls(tokenizer)
-        token_strings = _token_deltas(tokenizer, case.model_output)
-
-        reasoning_parts: list[str] = []
-        content_parts: list[str] = []
-        tool_call_deltas: list[dict] = []
-
-        previous_text = ""
-        previous_token_ids: list[int] = []
-
-        for token_str in token_strings:
-            current_text = previous_text + token_str
-            current_token_ids = previous_token_ids + [0]
-
-            delta = parser.extract_reasoning_streaming(
-                previous_text=previous_text,
-                current_text=current_text,
-                delta_text=token_str,
-                previous_token_ids=previous_token_ids,
-                current_token_ids=current_token_ids,
-                delta_token_ids=[0],
-            )
-            if delta is not None:
-                if delta.reasoning is not None:
-                    reasoning_parts.append(delta.reasoning)
-                if delta.content is not None:
-                    content_parts.append(delta.content)
-                for tc in delta.tool_calls:
-                    tool_call_deltas.append(
-                        {
-                            "id": tc.id,
-                            "index": tc.index,
-                            "name": tc.function.name if tc.function else None,
-                            "arguments": (
-                                tc.function.arguments if tc.function else None
-                            ),
-                        }
-                    )
-
-            previous_text = current_text
-            previous_token_ids = current_token_ids
-
-        reasoning = "".join(reasoning_parts) if reasoning_parts else None
-        assert reasoning == case.expected_reasoning
-
-        content = "".join(content_parts) if content_parts else None
-        assert content == case.expected_content
-        assert tool_call_deltas == []
-
-
-class TestIsReasoningEnd:
-    @pytest.mark.parametrize(
-        "parser_cls",
-        [CohereCommand3ReasoningParser, CohereCommand4ReasoningParser],
-        ids=["cmd3", "cmd4"],
-    )
-    def test_is_reasoning_end(self, tokenizer, parser_cls):
-        parser = parser_cls(tokenizer)
-        start_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
-        end_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
-        chatbot_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
-        content_ids = [99, 100]
-
-        # Generation-only tokens have no chatbot marker, so the whole sequence
-        # is considered.
-        assert parser.is_reasoning_end([end_id])
-        assert parser.is_reasoning_end([start_id, *content_ids, end_id])
-        assert not parser.is_reasoning_end([start_id, *content_ids])
-
-        # Full prompt/history tokens are scoped to the latest chatbot marker,
-        # so stray thinking tokens from the preamble or previous turns are ignored.
-        assert not parser.is_reasoning_end([start_id, end_id, chatbot_id, *content_ids])
-        assert parser.is_reasoning_end(
-            [start_id, end_id, chatbot_id, start_id, *content_ids, end_id]
-        )
-
-
-@pytest.mark.parametrize(
-    "parser_cls",
-    [
-        CohereCommand3ReasoningParser,
-        CohereCommand4ReasoningParser,
-    ],
-    ids=["cmd3", "cmd4"],
-)
-def test_count_reasoning_tokens(tokenizer, parser_cls):
-    parser = parser_cls(tokenizer)
-    start = parser.start_token_id
-    end = parser.end_token_id
-
-    assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2
-    assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3
-    assert parser.count_reasoning_tokens([start, 1, 2]) == 2
-    assert parser.count_reasoning_tokens([1, 2, end]) == 0
-
-
-SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}}
-SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}}
-GET_WEATHER_TOOL = {
-    "type": "function",
-    "function": {
-        "name": "get_weather",
-        "parameters": {
-            "type": "object",
-            "properties": {"city": {"type": "string"}},
-        },
-    },
-}
-VALID_STRUCTURAL_TAG = {
-    "type": "structural_tag",
-    "format": {
-        "type": "triggered_tags",
-        "tags": [
-            {
-                "begin": "<tool>",
-                "content": {"type": "any_text"},
-                "end": "</tool>",
-            }
-        ],
-        "triggers": ["<tool>"],
-    },
-}
-
-
-def _model_config(arch: str) -> SimpleNamespace:
-    return SimpleNamespace(
-        architecture=arch,
-        architectures=[arch],
-        hf_text_config=SimpleNamespace(architectures=[arch]),
-    )
-
-
-def _make_chat_request(**kwargs) -> ChatCompletionRequest:
-    data = {"messages": [{"role": "user", "content": "hi"}], "model": "m"}
-    data.update(kwargs)
-    return ChatCompletionRequest.model_validate(data)
-
-
-def _first_json_schema(tag_json: str) -> dict | None:
-    outer = json.loads(tag_json)
-    for t in (outer.get("format") or {}).get("tags") or []:
-        c = t.get("content") or {}
-        if c.get("type") == "json_schema":
-            js = c.get("json_schema")
-            return js if isinstance(js, dict) else None
-    return None
-
-
-def _content_types(tag_json: str) -> set[str]:
-    outer = json.loads(tag_json)
-    out: set[str] = set()
-    for t in (outer.get("format") or {}).get("tags") or []:
-        ty = (t.get("content") or {}).get("type")
-        if isinstance(ty, str):
-            out.add(ty)
-    return out
-
-
-@pytest.fixture(scope="module")
-def parser(tokenizer: MockCohereTokenizer) -> CohereCommand4ReasoningParser:
-    """Parser configured with a supported Cohere architecture."""
-    return CohereCommand4ReasoningParser(
-        tokenizer,
-        model_config=_model_config("Cohere2ForCausalLM"),
-    )
-
-
-@pytest.fixture(scope="module")
-def parser_no_model_config(
-    tokenizer: MockCohereTokenizer,
-) -> CohereCommand4ReasoningParser:
-    """Parser with no ``model_config`` (cannot resolve architecture)."""
-    return CohereCommand4ReasoningParser(tokenizer, model_config=None)
-
-
-@pytest.fixture(scope="module")
-def parser_unsupported_arch(
-    tokenizer: MockCohereTokenizer,
-) -> CohereCommand4ReasoningParser:
-    """Parser configured with an architecture that has no structural tag style."""
-    return CohereCommand4ReasoningParser(
-        tokenizer,
-        model_config=_model_config("LlamaForCausalLM"),
-    )
-
-
-class TestAdjustRequestPassthrough:
-    def test_structured_outputs_structural_tag_not_modified(self, parser) -> None:
-        tag = json.dumps(VALID_STRUCTURAL_TAG)
-        r = _make_chat_request(structured_outputs={"structural_tag": tag})
-        o = parser.adjust_request(r)
-        assert o.structured_outputs.structural_tag == tag
-
-    def test_response_format_structural_tag_short_circuit(self, parser) -> None:
-        # ``ChatCompletionRequest`` validates ``response_format`` as a union;
-        # bare ``{"type": "structural_tag"}`` is invalid (use pydantic model).
-        rf = StructuralTagResponseFormat(
-            type="structural_tag",
-            format=VALID_STRUCTURAL_TAG["format"],
-        )
-        r = _make_chat_request(response_format=rf)
-        o = parser.adjust_request(r)
-        assert _response_format_type(o.response_format) == "structural_tag"
-        assert o.structured_outputs is None
-
-
-class TestAdjustRequestNoOp:
-    def test_no_schema_no_tools(self, parser) -> None:
-        o = parser.adjust_request(_make_chat_request())
-        assert o.structured_outputs is None
-        assert o.response_format is None
-
-    def test_no_model_config(self, parser_no_model_config) -> None:
-        inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
-        r = _make_chat_request(
-            response_format=ResponseFormat(type="json_schema", json_schema=inner),
-        )
-        o = parser_no_model_config.adjust_request(r)
-        assert o.response_format is not None
-        assert o.structured_outputs is None
-
-
-class TestAdjustRequestUnsupportedArchitecture:
-    def test_json_schema_raises(self, parser_unsupported_arch) -> None:
-        inner = JsonSchemaResponseFormat(name="n", json_schema=SCHEMA_A)
-        r = _make_chat_request(
-            response_format=ResponseFormat(type="json_schema", json_schema=inner),
-        )
-        with pytest.raises(ValueError, match="does not support"):
-            parser_unsupported_arch.adjust_request(r)
-
-
-class TestAdjustRequestFoldFromResponseFormat:
-    @pytest.mark.parametrize(
-        "response_format, expected_schema",
-        [
-            pytest.param(
-                ResponseFormat(
-                    type="json_schema",
-                    json_schema=JsonSchemaResponseFormat(
-                        name="n", json_schema=SCHEMA_A
-                    ),
-                ),
-                SCHEMA_A,
-                id="json_schema_pydantic",
-            ),
-            pytest.param(
-                {
-                    "type": "json_schema",
-                    "json_schema": {"name": "n", "schema": SCHEMA_A},
-                },
-                SCHEMA_A,
-                id="json_schema_dict",
-            ),
-            pytest.param(
-                {"type": "json_object"},
-                {"type": "object"},
-                id="json_object",
-            ),
-        ],
-    )
-    def test_response_format_cleared(
-        self, parser, response_format, expected_schema
-    ) -> None:
-        r = _make_chat_request(response_format=response_format)
-        o = parser.adjust_request(r)
-        assert o.response_format is None
-        assert (
-            _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
-        )
-
-
-class TestHasEffectiveTools:
-    @pytest.mark.parametrize(
-        "tools, expected",
-        [
-            pytest.param(None, False, id="none"),
-            pytest.param([], False, id="empty_list"),
-            pytest.param("   ", False, id="blank_str"),
-            pytest.param(
-                [{"type": "function", "function": {"name": "f"}}],
-                True,
-                id="non_empty_list",
-            ),
-            pytest.param('{"x": 1}', True, id="non_empty_str"),
-        ],
-    )
-    def test_has_effective_tools(self, tools, expected) -> None:
-        assert _has_effective_tools(tools) is expected
-
-    def test_convert_schema_json_only_with_empty_tools_list(self) -> None:
-        tag = convert_schema_to_structural_tags(
-            schema=SCHEMA_B,
-            tools=[],
-            model_architecture="Cohere2ForCausalLM",
-        )
-        assert tag is not None
-        assert _first_json_schema(tag) == SCHEMA_B
-
-
-class TestAdjustRequestFoldFromStructuredOutputs:
-    @pytest.mark.parametrize(
-        "structured_outputs, expected_schema",
-        [
-            pytest.param({"json": SCHEMA_B}, SCHEMA_B, id="json_dict"),
-            pytest.param({"json": json.dumps(SCHEMA_B)}, SCHEMA_B, id="json_string"),
-            pytest.param(
-                {"json_object": True}, {"type": "object"}, id="json_object_flag"
-            ),
-            pytest.param(
-                StructuredOutputsParams(json=SCHEMA_B),
-                SCHEMA_B,
-                id="structured_outputs_dataclass",
-            ),
-            pytest.param(
-                {"json": {"name": "n", "schema": SCHEMA_A}},
-                SCHEMA_A,
-                id="openai_wrapper_dict_unwrapped",
-            ),
-        ],
-    )
-    def test_structured_outputs_folded(
-        self, parser, structured_outputs, expected_schema
-    ) -> None:
-        o = parser.adjust_request(
-            _make_chat_request(structured_outputs=structured_outputs),
-        )
-        assert (
-            _first_json_schema(o.structured_outputs.structural_tag) == expected_schema
-        )
-
-    def test_responses_request_default_empty_tools(self, parser) -> None:
-        """``ResponsesRequest.tools`` defaults to ``[]``, not ``None``."""
-        r = ResponsesRequest.model_validate(
-            {
-                "input": "hi",
-                "model": "m",
-                "structured_outputs": {"json": SCHEMA_B},
-            }
-        )
-        assert r.tools == []
-        o = parser.adjust_request(r)
-        assert _first_json_schema(o.structured_outputs.structural_tag) == SCHEMA_B
-
-    def test_json_userdict_mapping_unwrapped(self) -> None:
-        inner = {"type": "object", "properties": {"u": {"type": "number"}}}
-        so = StructuredOutputsParams(json=UserDict(inner))
-        assert _schema_dict_from_structured_outputs(so) == inner
-
-    @pytest.mark.parametrize(
-        "json_value, match",
-        [
-            pytest.param("{not json}", "valid JSON", id="invalid_json_string"),
-            pytest.param(
-                json.dumps(["a", "b"]), "JSON object", id="non_object_json_string"
-            ),
-            pytest.param("   ", "empty", id="empty_json_string"),
-        ],
-    )
-    def test_structured_outputs_json_string_raises(
-        self, parser, json_value, match
-    ) -> None:
-        with pytest.raises(ValueError, match=match):
-            parser.adjust_request(
-                _make_chat_request(structured_outputs={"json": json_value}),
-            )
-
-    @pytest.mark.parametrize(
-        "construct",
-        [
-            pytest.param(
-                lambda: _make_chat_request(structured_outputs={"json": [1, 2, 3]}),
-                id="chat_completion_request",
-            ),
-            pytest.param(
-                lambda: StructuredOutputsParams(json=[1, 2, 3]),  # type: ignore[arg-type]
-                id="structured_outputs_params",
-            ),
-        ],
-    )
-    def test_json_wrong_type_raises(self, construct) -> None:
-        """Non-str / non-dict ``json`` fails at Pydantic validation."""
-        with pytest.raises(ValidationError):
-            construct()
-
-
-class TestAdjustRequestPrecedence:
-    def test_response_format_over_structured_outputs_json(self, parser) -> None:
-        s_rf = {"type": "object", "properties": {"rf": {"type": "string"}}}
-        s_so = {"type": "object", "properties": {"so": {"type": "number"}}}
-        inner = JsonSchemaResponseFormat(name="n", json_schema=s_rf)
-        r = _make_chat_request(
-            response_format=ResponseFormat(type="json_schema", json_schema=inner),
-            structured_outputs={"json": s_so},
-        )
-        o = parser.adjust_request(r)
-        assert _first_json_schema(o.structured_outputs.structural_tag) == s_rf
-
-
-class TestAdjustRequestTextPlusStructuredOutputs:
-    def test_text_response_format_preserved(self, parser) -> None:
-        sch = {"type": "object", "properties": {"k": {"type": "string"}}}
-        r = _make_chat_request(
-            response_format=ResponseFormat(type="text"),
-            structured_outputs={"json": sch},
-        )
-        o = parser.adjust_request(r)
-        assert o.response_format is not None
-        assert o.response_format.type == "text"
-        assert _first_json_schema(o.structured_outputs.structural_tag) == sch
-
-
-class TestAdjustRequestTools:
-    def test_tools_only_command_a_grammar(self, parser) -> None:
-        o = parser.adjust_request(
-            _make_chat_request(tools=[GET_WEATHER_TOOL], tool_choice="auto"),
-        )
-        assert "grammar" in _content_types(o.structured_outputs.structural_tag)
-
-    def test_tools_plus_json_schema_both_kinds(self, parser) -> None:
-        inner = JsonSchemaResponseFormat(
-            name="n",
-            json_schema={"type": "object", "properties": {"r": {"type": "string"}}},
-        )
-        r = _make_chat_request(
-            response_format=ResponseFormat(type="json_schema", json_schema=inner),
-            tools=[GET_WEATHER_TOOL],
-            tool_choice="auto",
-        )
-        o = parser.adjust_request(r)
-        types = _content_types(o.structured_outputs.structural_tag)
-        assert "grammar" in types
-        assert "json_schema" in types
-
-
-class TestMelodySourceResolution:
-    """Pin the parser-side source resolution (see
-    ``_melody_sources_to_vllm``). The parser receives melody's numeric
-    ``(tool_call_index, tool_result_indices)`` addressing and resolves
-    it against a ``position_to_source`` map handed in by
-    ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same
-    map that would otherwise live in the serving layer's resolver.
-    """
-
-    @staticmethod
-    def _fake_melody_source(bucket: int, indices: list[int]) -> Any:
-        return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices)
-
-    def test_multi_index_source_fans_out(self):
-        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
-
-        position_map: dict[tuple[int, int], CitationSource] = {
-            (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}),
-            (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}),
-        }
-        raw = [self._fake_melody_source(0, [0, 1])]
-        out = _melody_sources_to_vllm(raw, position_map)
-        assert [s.id for s in out] == ["d0", "d1"]
-        # Verify type / payload were plumbed through, not just the id.
-        assert out[0].type == "document"
-        assert out[0].document == {"id": "d0"}
-
-    def test_unresolvable_position_skipped(self):
-        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
-
-        position_map: dict[tuple[int, int], CitationSource] = {
-            (0, 0): CitationSource(type="document", id="d0"),
-        }
-        raw = [self._fake_melody_source(9, [0])]
-        assert _melody_sources_to_vllm(raw, position_map) == []
-
-    def test_missing_position_map_drops_all_sources(self):
-        # A parser instance without a position map (parser wired
-        # outside of ``CohereServingChatV2``) can't attribute anything,
-        # so every source is dropped. Callers downstream will see the
-        # citation with empty ``sources`` and drop it entirely.
-        raw = [self._fake_melody_source(0, [0])]
-        assert _melody_sources_to_vllm(raw, None) == []
-
-    def test_citations_pass_through_is_thinking_tag(self):
-        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
-
-        position_map: dict[tuple[int, int], CitationSource] = {
-            (0, 0): CitationSource(type="document", id="d0"),
-        }
-        raw = [
-            SimpleNamespace(
-                start_index=0,
-                end_index=5,
-                text="hello",
-                is_thinking=True,
-                sources=[self._fake_melody_source(0, [0])],
-            )
-        ]
-        out = _melody_citations_to_vllm(raw, position_map)
-        assert out is not None
-        assert out[0].type == "THINKING_CONTENT"
-        assert out[0].sources[0].id == "d0"
-
-
-class TestParserStreamingEndToEnd:
-    """End-to-end streaming shape: raw model output text is tokenized
-    with the mock byte-level tokenizer and fed one token at a time
-    through :meth:`CohereCommand3ReasoningParser.extract_reasoning_streaming`.
-    We collect the ordered sequence of ``DeltaMessage``s the parser
-    hands back -- i.e. what the client would receive on the wire --
-    and pin the exact shape.
-
-    Combines the three flavors that reach ``delta``: a thinking block
-    (populates ``reasoning``), a text block (populates ``content``),
-    and a resolved citation (populates ``citations``) referencing
-    ``(start, end)`` offsets into the already-streamed content.
-    """
-
-    _MODEL_OUTPUT = (
-        "<|START_THINKING|>Let me check.<|END_THINKING|>"
-        "<|START_RESPONSE|>The capital is <co>Paris</co: 0:[1]>.<|END_RESPONSE|>"
-    )
-
-    @staticmethod
-    def _drive_parser(
-        parser: CohereCommand3ReasoningParser, text: str, tok: MockCohereTokenizer
-    ) -> list[Any]:
-        """Feed ``text`` through the parser one token at a time
-        (byte-level) and return the ordered list of non-``None``
-        ``DeltaMessage``s.
-        """
-        deltas: list[Any] = []
-        previous_text = ""
-        previous_token_ids: list[int] = []
-        for token_str in _token_deltas(tok, text):
-            current_text = previous_text + token_str
-            current_token_ids = previous_token_ids + [0]
-            delta = parser.extract_reasoning_streaming(
-                previous_text=previous_text,
-                current_text=current_text,
-                delta_text=token_str,
-                previous_token_ids=previous_token_ids,
-                current_token_ids=current_token_ids,
-                delta_token_ids=[0],
-            )
-            if delta is not None:
-                deltas.append(delta)
-            previous_text = current_text
-            previous_token_ids = current_token_ids
-        return deltas
-
-    def test_full_wire_stream_shape(self, tokenizer):
-        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
-        from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
-
-        # Mimic what ``CohereServingChatV2._apply_cohere_template_kwargs``
-        # installs on a request whose tool_call_index=0 tool result at
-        # index 1 is the "France" document.
-        position_to_source = {
-            (0, 1): CitationSource(
-                type="document",
-                id="doc-paris",
-                document={"id": "doc-paris", "title": "France"},
-            ),
-        }
-
-        parser = CohereCommand3ReasoningParser(
-            tokenizer,
-            chat_template_kwargs={POSITION_TO_SOURCE_KEY: position_to_source},
-        )
-
-        deltas = self._drive_parser(parser, self._MODEL_OUTPUT, tokenizer)
-
-        # -- 1. Reasoning and content stream separately and reach
-        # the wire only as their real payload. Framing tokens
-        # (``<|START_THINKING|>`` etc.) are consumed by the parser's
-        # state machine and never produce a delta.
-        reasoning_stream = [d.reasoning for d in deltas if d.reasoning is not None]
-        content_stream = [d.content for d in deltas if d.content is not None]
-        citation_deltas = [d for d in deltas if getattr(d, "citations", None)]
-
-        assert "".join(reasoning_stream) == "Let me check."
-        assert "".join(content_stream) == "The capital is Paris."
-
-        # -- 2. Exactly one citation reaches the wire, resolved to the
-        # source from ``position_to_source`` (not the raw
-        # ``(bucket, idx)`` coordinates).
-        assert len(citation_deltas) == 1
-        (cite_delta,) = citation_deltas
-        assert cite_delta.citations is not None
-        assert len(cite_delta.citations) == 1
-        cite = cite_delta.citations[0]
-        assert cite.text == "Paris"
-        # Char offsets are into the accumulated ``content``: "The
-        # capital is " is 15 chars, "Paris" is 5.
-        assert cite.start == 15
-        assert cite.end == 20
-        assert len(cite.sources) == 1
-        assert cite.sources[0].id == "doc-paris"
-        assert cite.sources[0].type == "document"
-
-        # -- 3. The citation's ``(start, end)`` offsets point into
-        # bytes that were already emitted as content-deltas earlier
-        # in the stream -- i.e. the citation is anchored to real
-        # already-shipped content.
-        cite_delta_idx = deltas.index(cite_delta)
-        earlier_content = "".join(
-            d.content for d in deltas[:cite_delta_idx] if d.content is not None
-        )
-        assert earlier_content[cite.start : cite.end] == "Paris"
-
-        # Everything after the citation delta is the trailing period.
-        later_content = "".join(
-            d.content for d in deltas[cite_delta_idx + 1 :] if d.content is not None
-        )
-        assert later_content == "."
-
-        # -- 4. Reasoning and content are never mixed on the same
-        # delta: the parser flips modes on the ``<|END_THINKING|>``
-        # boundary and never emits a delta with both fields set.
-        for d in deltas:
-            assert not (d.reasoning is not None and d.content is not None)
diff --git a/tests/renderers/test_cohere.py b/tests/renderers/test_cohere.py
index 76b38734c212..ae2dfe6c7a88 100644
--- a/tests/renderers/test_cohere.py
+++ b/tests/renderers/test_cohere.py
@@ -997,6 +997,7 @@ class _MockModelConfig:
     skip_tokenizer_init: bool = True
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
 
 
diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py
index 6ae8c4d5510a..99862219b97e 100644
--- a/tests/renderers/test_completions.py
+++ b/tests/renderers/test_completions.py
@@ -39,6 +39,7 @@ class MockModelConfig:
     skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
     hidden_size: int = 768
     dtype: torch.dtype = torch.float32
@@ -506,7 +507,7 @@ def test_explicit_side_text_pretokenization_guard(self):
 
 class TestRenderEmbedPrompt:
     def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes:
-        """Helper to create base64-encoded tensor bytes"""
+        """Helper to create base64-encoded tensor bytes."""
         buffer = io.BytesIO()
         torch.save(tensor, buffer)
         buffer.seek(0)
diff --git a/tests/renderers/test_gemma4_chat_template.py b/tests/renderers/test_gemma4_chat_template.py
index 2c1312a84c6c..7d6b0b05f814 100644
--- a/tests/renderers/test_gemma4_chat_template.py
+++ b/tests/renderers/test_gemma4_chat_template.py
@@ -228,7 +228,7 @@ def test_generation_prompt_not_after_tool_response(self, gemma4_template):
         assert not result.strip().endswith("<|turn>model\n")
 
     def test_reasoning_in_tool_chains(self, gemma4_template):
-        """reasoning field on assistant with tool_calls after last user
+        """Reasoning field on assistant with tool_calls after last user
         message emits <|channel>thought\\n...<channel|>."""
         messages = [
             {"role": "user", "content": "Calculate something"},
@@ -252,7 +252,7 @@ def test_reasoning_in_tool_chains(self, gemma4_template):
         assert "<channel|>" in result
 
     def test_reasoning_not_before_last_user(self, gemma4_template):
-        """reasoning on assistant BEFORE the last user message is dropped."""
+        """Reasoning on assistant BEFORE the last user message is dropped."""
         messages = [
             {"role": "user", "content": "First"},
             {
diff --git a/tests/renderers/test_hf.py b/tests/renderers/test_hf.py
index 736478fbfafe..833656aa3c31 100644
--- a/tests/renderers/test_hf.py
+++ b/tests/renderers/test_hf.py
@@ -123,7 +123,7 @@ def test_no_load_chat_template_literallike():
 )
 @pytest.mark.parametrize("use_tools", [True, False])
 def test_resolve_chat_template(sample_json_schema, model, use_tools):
-    """checks that chat_template is a dict type for HF models."""
+    """Checks that chat_template is a dict type for HF models."""
     model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
     model_info.check_available_online(on_fail="skip")
 
@@ -196,7 +196,7 @@ def test_resolve_chat_template(sample_json_schema, model, use_tools):
     ],
 )
 def test_resolve_chat_template_kwargs(sample_json_schema, model, expected_kwargs):
-    """checks that chat_template is a dict type for HF models."""
+    """Checks that chat_template is a dict type for HF models."""
     model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
     model_info.check_available_online(on_fail="skip")
 
diff --git a/tests/renderers/test_kimi_k3.py b/tests/renderers/test_kimi_k3.py
index af8723f801a2..de43c0f6bbcc 100644
--- a/tests/renderers/test_kimi_k3.py
+++ b/tests/renderers/test_kimi_k3.py
@@ -40,6 +40,7 @@ class MockHFConfig:
 class MockModelConfig:
     runner_type: str = "generate"
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     multimodal_config: Any = None
     hf_config: MockHFConfig = field(default_factory=MockHFConfig)
     allowed_local_media_path: str = ""
diff --git a/tests/renderers/test_mistral.py b/tests/renderers/test_mistral.py
index 3a8b49a4db29..6ad55f2fb27f 100644
--- a/tests/renderers/test_mistral.py
+++ b/tests/renderers/test_mistral.py
@@ -37,6 +37,7 @@ class MockModelConfig:
     skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    supports_multimodal_inputs: bool = False
     renderer_num_workers: int = 1
 
 
diff --git a/tests/renderers/test_process_multi_modal_uuids.py b/tests/renderers/test_process_multi_modal_uuids.py
index 76b394897ee1..a9eb0396b241 100644
--- a/tests/renderers/test_process_multi_modal_uuids.py
+++ b/tests/renderers/test_process_multi_modal_uuids.py
@@ -15,7 +15,11 @@
 
 cherry_pil_image = ImageAsset("cherry_blossom").pil_image
 stop_pil_image = ImageAsset("stop_sign").pil_image
-baby_reading_np_ndarrays = VideoAsset("baby_reading").np_ndarrays
+baby_reading_video_asset = VideoAsset("baby_reading")
+baby_reading_video = (
+    baby_reading_video_asset.np_ndarrays,
+    baby_reading_video_asset.metadata,
+)
 
 
 def _build_renderer(
@@ -127,7 +131,7 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
 
     mm_data = {
         "image": [cherry_pil_image, stop_pil_image],
-        "video": baby_reading_np_ndarrays,
+        "video": baby_reading_video,
     }
 
     # Use a consistent two-image scenario across all configurations
@@ -176,17 +180,39 @@ def test_multi_modal_uuids_accepts_empty(
     assert processed_mm_uuids == mm_uuids
 
 
-def test_multi_modal_uuids_ignored_when_caching_disabled():
-    # When both processor cache is 0 and prefix caching disabled, the
-    # processor builds overrides from request id instead of using user UUIDs.
+@pytest.mark.parametrize(
+    "mm_uuids, expected",
+    [
+        (
+            {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]},
+            {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]},
+        ),
+        (
+            {"image": [None, "hash_stop"], "video": None},
+            {"image": ["req-42-image-0", "hash_stop"], "video": ["req-42-video-0"]},
+        ),
+        (
+            {"image": ["", None]},
+            {"image": ["", "req-42-image-1"], "video": ["req-42-video-0"]},
+        ),
+        (
+            {},
+            {
+                "image": ["req-42-image-0", "req-42-image-1"],
+                "video": ["req-42-video-0"],
+            },
+        ),
+    ],
+)
+def test_multi_modal_uuids_preserved_when_caching_disabled(mm_uuids, expected):
+    """Only missing UUIDs get request-local IDs when both caches are disabled."""
     renderer = _build_renderer(mm_cache_gb=0.0, enable_prefix_caching=False)
 
     request_id = "req-42"
     mm_data = {
         "image": [cherry_pil_image, stop_pil_image],
-        "video": baby_reading_np_ndarrays,
+        "video": baby_reading_video,
     }
-    mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}
 
     mm_processor = renderer.get_mm_processor()
     mm_data_items = mm_processor.info.parse_mm_data(mm_data)
@@ -196,16 +222,4 @@ def test_multi_modal_uuids_ignored_when_caching_disabled():
         mm_data, mm_data_items, mm_uuid_items, request_id
     )
 
-    # Expect request-id-based overrides are passed through
-    assert set(mm_uuids.keys()) == {"image", "video"}
-    assert len(mm_uuids["image"]) == 2
-    assert len(mm_uuids["video"]) == 1
-    assert processed_mm_uuids["image"][0].startswith(
-        f"{request_id}-image-"
-    ) and processed_mm_uuids["image"][0].endswith("-0")
-    assert processed_mm_uuids["image"][1].startswith(
-        f"{request_id}-image-"
-    ) and processed_mm_uuids["image"][1].endswith("-1")
-    assert processed_mm_uuids["video"][0].startswith(
-        f"{request_id}-video-"
-    ) and processed_mm_uuids["video"][0].endswith("-0")
+    assert processed_mm_uuids == expected
diff --git a/tests/renderers/test_prompt_embeds_payload_validation.py b/tests/renderers/test_prompt_embeds_payload_validation.py
new file mode 100644
index 000000000000..bec28c4c49ac
--- /dev/null
+++ b/tests/renderers/test_prompt_embeds_payload_validation.py
@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""`prompt_embeds` payloads that are not tensor files must be client errors.
+
+`prompt_embeds` is a base64 blob the caller writes, and every other way of
+getting it wrong -- wrong rank, wrong hidden_size, wrong dtype, not a tensor at
+all -- already answers 400. A payload that `torch.load` cannot read at all was
+the exception: torch has no single error type for that (the zip reader raises
+`RuntimeError`, the legacy pickle path raises an unpickling error, `KeyError` or
+`IndexError`, an empty payload raises `EOFError`), none of them is a
+`ValueError`, and the entrypoints' fallback therefore mapped them to 500.
+
+The assertions go through `create_error_response`, which is what actually turns
+an exception into a status code, so they pin the HTTP answer rather than the
+exception class alone.
+"""
+
+import io
+import zipfile
+from unittest.mock import patch
+
+import pybase64 as base64
+import pytest
+import torch
+
+from vllm.entrypoints.serve.exception_handling.error_response import (
+    create_error_response,
+)
+from vllm.exceptions import VLLMValidationError
+from vllm.renderers.embed_utils import (
+    _MAX_EMBED_ERROR_REASON_CHARS,
+    safe_load_prompt_embeds,
+)
+
+
+@pytest.fixture
+def model_config():
+    from vllm.config import ModelConfig
+
+    return ModelConfig(
+        model="facebook/opt-125m",
+        tokenizer="facebook/opt-125m",
+        tokenizer_mode="auto",
+        trust_remote_code=False,
+        dtype="float32",
+        seed=0,
+        enable_prompt_embeds=True,
+    )
+
+
+def _encode(raw: bytes) -> bytes:
+    return base64.b64encode(raw)
+
+
+def _zip_without_a_tensor() -> bytes:
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, "w") as archive:
+        archive.writestr("junk.txt", "hello")
+    return buffer.getvalue()
+
+
+def _saved_tensor(tensor: torch.Tensor) -> bytes:
+    buffer = io.BytesIO()
+    torch.save(tensor, buffer)
+    return buffer.getvalue()
+
+
+UNPARSABLE_PAYLOADS = [
+    pytest.param(b"", id="empty"),
+    pytest.param(b"\x80", id="single-byte"),
+    pytest.param(b"\x00\x01\x02\x03not-a-tensor", id="random-bytes"),
+    pytest.param(b"hello world this is not a torch file", id="ascii-text"),
+    pytest.param(b"PK\x03\x04" + b"\x00" * 20, id="truncated-zip"),
+    pytest.param(_zip_without_a_tensor(), id="zip-without-a-tensor"),
+    # pickle.dumps(42): a valid pickle stream that is not a torch checkpoint.
+    # Spelled out rather than built, so this file does not import pickle.
+    pytest.param(b"\x80\x04K*.", id="legacy-pickle-of-an-int"),
+]
+
+
+@pytest.mark.parametrize("payload", UNPARSABLE_PAYLOADS)
+def test_unreadable_prompt_embeds_payload_is_a_client_error(model_config, payload):
+    with pytest.raises(VLLMValidationError) as excinfo:
+        safe_load_prompt_embeds(model_config, _encode(payload))
+
+    assert excinfo.value.parameter == "prompt_embeds"
+    assert int(create_error_response(excinfo.value).error.code) == 400
+
+
+@pytest.mark.parametrize("payload", UNPARSABLE_PAYLOADS)
+def test_the_error_body_stays_bounded(model_config, payload):
+    """Torch's reason is built from the caller's bytes, so it must be truncated.
+
+    An unpickler that reports the offending global, for instance, quotes a name
+    the caller chose; without a bound a large one would be reflected whole.
+    """
+    padding = b"A" * 100_000
+    with pytest.raises(VLLMValidationError) as excinfo:
+        safe_load_prompt_embeds(model_config, _encode(payload + padding))
+
+    assert len(str(excinfo.value)) < 500
+
+
+def test_a_long_reason_is_truncated_and_says_how_much_was_left_out(model_config):
+    """The cap is pinned against a reason of known length, not a real payload.
+
+    Which of torch's messages happen to exceed the cap is torch's business and
+    changes between releases, so the reason is injected here: that keeps the
+    assertion about *this* code's truncation rule rather than about whichever
+    error text the installed torch produces.
+    """
+    reason = "E" * (_MAX_EMBED_ERROR_REASON_CHARS + 37)
+
+    with (
+        patch(
+            "vllm.renderers.embed_utils.torch.load", side_effect=RuntimeError(reason)
+        ),
+        pytest.raises(VLLMValidationError) as excinfo,
+    ):
+        safe_load_prompt_embeds(model_config, _encode(b"ignored"))
+
+    message = str(excinfo.value)
+    assert "E" * _MAX_EMBED_ERROR_REASON_CHARS in message
+    assert "E" * (_MAX_EMBED_ERROR_REASON_CHARS + 1) not in message
+    assert "... (37 more characters truncated)" in message
+
+
+def test_a_short_reason_is_passed_through_whole(model_config):
+    """A reason that fits is not annotated -- the count only appears when it is
+    telling the reader something."""
+    reason = "invalid magic number; corrupt file?"
+    assert len(reason) <= _MAX_EMBED_ERROR_REASON_CHARS
+
+    with (
+        patch(
+            "vllm.renderers.embed_utils.torch.load", side_effect=RuntimeError(reason)
+        ),
+        pytest.raises(VLLMValidationError) as excinfo,
+    ):
+        safe_load_prompt_embeds(model_config, _encode(b"ignored"))
+
+    message = str(excinfo.value)
+    assert reason in message
+    assert "truncated" not in message
+
+
+def test_a_well_formed_payload_still_loads(model_config):
+    """Positive control: the healthy path is untouched."""
+    hidden_size = model_config.get_hidden_size()
+    tensor = torch.zeros(3, hidden_size, dtype=torch.float32)
+
+    loaded = safe_load_prompt_embeds(model_config, _encode(_saved_tensor(tensor)))
+
+    assert loaded.shape == (3, hidden_size)
+    assert loaded.dtype == model_config.dtype
+
+
+def test_a_payload_that_is_not_base64_is_still_a_client_error(model_config):
+    """Positive control: the base64 stage already answered 400 and still does."""
+    with pytest.raises(ValueError) as excinfo:
+        safe_load_prompt_embeds(model_config, b"!!!!not base64!!!!")
+
+    assert int(create_error_response(excinfo.value).error.code) == 400
diff --git a/tests/renderers/test_sparse_tensor_concurrent_race.py b/tests/renderers/test_sparse_tensor_concurrent_race.py
index 22925a9f1961..a1683eee0bff 100644
--- a/tests/renderers/test_sparse_tensor_concurrent_race.py
+++ b/tests/renderers/test_sparse_tensor_concurrent_race.py
@@ -11,6 +11,7 @@
 import torch
 
 from vllm.config import ModelConfig
+from vllm.exceptions import VLLMValidationError
 from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO
 from vllm.renderers.embed_utils import safe_load_prompt_embeds
 from vllm.utils.sparse_utils import check_sparse_tensor_invariants_threadsafe
@@ -54,7 +55,7 @@ class TestNegativeControlWithoutRace:
 
     def test_malicious_sparse_rejected_by_prompt_loader(self, model_config):
         encoded = _encode_tensor(_create_malicious_sparse_tensor())
-        with pytest.raises((RuntimeError, ValueError)):
+        with pytest.raises((RuntimeError, ValueError, VLLMValidationError)):
             safe_load_prompt_embeds(model_config, encoded)
 
     def test_malicious_sparse_rejected_by_image_loader(self):
@@ -147,7 +148,7 @@ def attempt_load(_):
             try:
                 safe_load_prompt_embeds(model_config, malicious_encoded)
                 bypassed.append(True)
-            except (RuntimeError, ValueError):
+            except (RuntimeError, ValueError, VLLMValidationError):
                 rejected.append(True)
 
         with ThreadPoolExecutor(max_workers=num_workers) as pool:
@@ -192,7 +193,7 @@ def test_flag_restored_after_exception(self, model_config):
         initial = torch.sparse.check_sparse_tensor_invariants.is_enabled()
         malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor())
 
-        with pytest.raises((RuntimeError, ValueError)):
+        with pytest.raises((RuntimeError, ValueError, VLLMValidationError)):
             safe_load_prompt_embeds(model_config, malicious_encoded)
 
         assert torch.sparse.check_sparse_tensor_invariants.is_enabled() == initial
@@ -203,7 +204,7 @@ def test_flag_restored_after_concurrent_exceptions(self, model_config):
         malicious_encoded = _encode_tensor(_create_malicious_sparse_tensor())
 
         def attempt_load(_):
-            with contextlib.suppress(RuntimeError, ValueError):
+            with contextlib.suppress(RuntimeError, ValueError, VLLMValidationError):
                 safe_load_prompt_embeds(model_config, malicious_encoded)
 
         with ThreadPoolExecutor(max_workers=4) as pool:
@@ -244,7 +245,7 @@ def load_image():
         # serialization — one finishes before the other starts), AND invalid
         # tensors are still rejected.
         malicious = _encode_tensor(_create_malicious_sparse_tensor())
-        with pytest.raises((RuntimeError, ValueError)):
+        with pytest.raises((RuntimeError, ValueError, VLLMValidationError)):
             safe_load_prompt_embeds(model_config, malicious)
         with pytest.raises((RuntimeError, ValueError)):
             io_handler.load_base64("", malicious.decode("utf-8"))
diff --git a/tests/renderers/test_sparse_tensor_validation.py b/tests/renderers/test_sparse_tensor_validation.py
index c119b090c0d0..6cb26a8a1de8 100644
--- a/tests/renderers/test_sparse_tensor_validation.py
+++ b/tests/renderers/test_sparse_tensor_validation.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests verify that malicious sparse tensors are rejected before they can trigger
+"""Tests verify that malicious sparse tensors are rejected before they can trigger
 out-of-bounds memory writes, or an unbounded allocation, during to_dense()
 operations.
 
@@ -51,8 +50,7 @@ def _encode_tensor(tensor: torch.Tensor) -> bytes:
 
 
 def _create_malicious_sparse_tensor() -> torch.Tensor:
-    """
-    Create a malicious sparse COO tensor with out-of-bounds indices.
+    """Create a malicious sparse COO tensor with out-of-bounds indices.
 
     This tensor has indices that point beyond the declared shape, which would
     cause an out-of-bounds write when converted to dense format without
@@ -117,8 +115,10 @@ def test_malicious_sparse_tensor_rejected(self, model_config):
         malicious_tensor = _create_malicious_sparse_tensor()
         encoded = _encode_tensor(malicious_tensor)
 
-        # Should raise RuntimeError due to invalid sparse tensor
-        with pytest.raises((RuntimeError, ValueError)) as exc_info:
+        # The invariant check raises RuntimeError from inside torch.load;
+        # safe_load_prompt_embeds now reports that as a client error so the
+        # caller gets a 400 rather than the 500 a bare RuntimeError produced.
+        with pytest.raises(VLLMValidationError) as exc_info:
             safe_load_prompt_embeds(model_config, encoded)
 
         # Error should indicate sparse tensor validation failure
@@ -137,7 +137,7 @@ def test_extremely_large_indices_rejected(self, model_config):
         )
         encoded = _encode_tensor(malicious_tensor)
 
-        with pytest.raises((RuntimeError, ValueError)):
+        with pytest.raises(VLLMValidationError):
             safe_load_prompt_embeds(model_config, encoded)
 
     def test_negative_indices_rejected(self, model_config):
@@ -152,7 +152,7 @@ def test_negative_indices_rejected(self, model_config):
         )
         encoded = _encode_tensor(malicious_tensor)
 
-        with pytest.raises((RuntimeError, ValueError)):
+        with pytest.raises(VLLMValidationError):
             safe_load_prompt_embeds(model_config, encoded)
 
     def test_hidden_size_mismatch_rejected(self, model_config):
@@ -264,7 +264,7 @@ def test_load_bytes_validates(self):
             io_handler.load_bytes(buffer.read())
 
     def test_valid_numpy_tensor_accepted(self):
-        """numpy .npy format should load and return correct tensor."""
+        """Numpy .npy format should load and return correct tensor."""
         io_handler = ImageEmbeddingMediaIO()
 
         arr = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)
@@ -279,8 +279,7 @@ def test_valid_numpy_tensor_accepted(self):
         assert torch.allclose(result, torch.from_numpy(arr))
 
     def test_numpy_int32_tensor_accepted(self):
-        """numpy int32 arrays should round-trip correctly."""
-
+        """Numpy int32 arrays should round-trip correctly."""
         io_handler = ImageEmbeddingMediaIO()
 
         arr = np.arange(280, dtype=np.int32)
@@ -294,8 +293,7 @@ def test_numpy_int32_tensor_accepted(self):
         assert (result == torch.from_numpy(arr)).all()
 
     def test_load_file_numpy_tensor_accepted(self, tmp_path):
-        """numpy .npy files should load correctly via load_file."""
-
+        """Numpy .npy files should load correctly via load_file."""
         io_handler = ImageEmbeddingMediaIO()
 
         arr = np.array([[1.5, 2.5], [3.5, 4.5]], dtype=np.float32)
@@ -362,13 +360,10 @@ def test_load_bytes_validates(self):
 
 
 class TestSparseTensorValidationIntegration:
-    """
-    These tests verify the complete attack chain is blocked at all entry points.
-    """
+    """These tests verify the complete attack chain is blocked at all entry points."""
 
     def test_attack_scenario_completions_api(self, model_config):
-        """
-        Simulate a complete attack through the Completions API.
+        """Simulate a complete attack through the Completions API.
 
         Attack scenario:
         1. Attacker crafts malicious sparse tensor
@@ -379,13 +374,12 @@ def test_attack_scenario_completions_api(self, model_config):
         # Step 1-2: Attacker creates malicious payload
         attack_payload = _encode_tensor(_create_malicious_sparse_tensor())
 
-        # Step 3-4: Server processes and should reject
-        with pytest.raises((RuntimeError, ValueError)):
+        # Step 3-4: Server processes and should reject, as a client error
+        with pytest.raises(VLLMValidationError):
             safe_load_prompt_embeds(model_config, attack_payload)
 
     def test_attack_scenario_chat_api_image(self):
-        """
-        Simulate attack through Chat API with image_embeds.
+        """Simulate attack through Chat API with image_embeds.
 
         Verifies the image embeddings path is protected.
         """
@@ -396,8 +390,7 @@ def test_attack_scenario_chat_api_image(self):
             io_handler.load_base64("", attack_payload.decode("utf-8"))
 
     def test_attack_scenario_chat_api_audio(self):
-        """
-        Simulate attack through Chat API with audio_embeds.
+        """Simulate attack through Chat API with audio_embeds.
 
         Verifies the audio embeddings path is protected.
         """
diff --git a/tests/renderers/test_token_offsets.py b/tests/renderers/test_token_offsets.py
index fc9192bb43e7..aad91655ba03 100644
--- a/tests/renderers/test_token_offsets.py
+++ b/tests/renderers/test_token_offsets.py
@@ -20,7 +20,7 @@ def fast_tokenizer():
     """gpt2 ships a Fast tokenizer; use it to test the offsets happy path."""
     from transformers import AutoTokenizer
 
-    return AutoTokenizer.from_pretrained("openai-community/gpt2", use_fast=True)
+    return AutoTokenizer.from_pretrained("openai-community/gpt2")
 
 
 def _make_base_renderer_with(tokenizer):
diff --git a/tests/renderers/test_truncate_after_mm_expansion.py b/tests/renderers/test_truncate_after_mm_expansion.py
new file mode 100644
index 000000000000..96cdfb12c2d7
--- /dev/null
+++ b/tests/renderers/test_truncate_after_mm_expansion.py
@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Truncation must bound the prompt *after* multimodal placeholders expand.
+
+``TokenizeParams.apply_post_tokenization`` truncates the unexpanded prompt, so
+one placeholder token can expand into many and push the result back over
+``truncate_prompt_tokens`` (vllm-project/vllm#57038).
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config import ModelConfig, VllmConfig
+from vllm.exceptions import VLLMValidationError
+from vllm.multimodal.inputs import PlaceholderRange
+from vllm.renderers import renderer_from_config
+from vllm.renderers.base import BaseRenderer
+from vllm.renderers.hf import _ensure_prompt_embeds_placeholder_token
+from vllm.renderers.params import TokenizeParams
+
+
+def _renderer():
+    """A BaseRenderer stub: only the tokenizer attribute is needed here."""
+
+    class _StubRenderer(BaseRenderer):
+        def __init__(self):
+            self.tokenizer = None
+
+        def render_messages(self, *args, **kwargs):
+            raise NotImplementedError
+
+    return _StubRenderer()
+
+
+def _mm_input(num_tokens: int, offset: int, length: int, **extra):
+    return {
+        "type": "multimodal",
+        "prompt_token_ids": list(range(num_tokens)),
+        "mm_kwargs": {},
+        "mm_hashes": {},
+        "mm_placeholders": {"audio": [PlaceholderRange(offset=offset, length=length)]},
+        **extra,
+    }
+
+
+def _params(limit: int | None, side: str = "right"):
+    return TokenizeParams(
+        max_total_tokens=4096,
+        truncate_prompt_tokens=limit,
+        truncation_side=side,
+    )
+
+
+def test_expanded_prompt_is_bounded():
+    """The reported case: audio fits, trailing text is cut to the limit."""
+    engine_input = _mm_input(num_tokens=147, offset=1, length=52)
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+    assert engine_input["prompt_token_ids"] == list(range(96))
+    assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange(
+        offset=1, length=52
+    )
+
+
+def test_prompt_within_limit_is_untouched():
+    engine_input = _mm_input(num_tokens=64, offset=1, length=52)
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+    assert engine_input["prompt_token_ids"] == list(range(64))
+
+
+def test_split_multimodal_span_is_rejected():
+    """Cutting mid-span would desync features from token ids, so reject."""
+    engine_input = _mm_input(num_tokens=147, offset=80, length=50)
+
+    with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
+        _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+
+def test_dropped_multimodal_span_is_rejected():
+    """A span entirely past the cut would be discarded silently."""
+    engine_input = _mm_input(num_tokens=147, offset=100, length=20)
+
+    with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
+        _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+
+def test_left_truncation_shifts_placeholder_offsets():
+    engine_input = _mm_input(num_tokens=147, offset=100, length=20)
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(96, side="left"))
+
+    assert engine_input["prompt_token_ids"] == list(range(51, 147))
+    # The span started at 100 and 51 leading tokens were dropped.
+    assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange(
+        offset=49, length=20
+    )
+
+
+def test_text_only_input_is_untouched():
+    """Only multimodal prompts grow after tokenization, so nothing else runs."""
+    engine_input = {
+        "type": "token",
+        "prompt_token_ids": list(range(147)),
+        "prompt_token_offsets": [(i, i + 1) for i in range(147)],
+    }
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(96))
+
+    assert engine_input["prompt_token_ids"] == list(range(147))
+    assert len(engine_input["prompt_token_offsets"]) == 147
+
+
+def test_no_truncation_requested_is_noop():
+    engine_input = _mm_input(num_tokens=147, offset=1, length=52)
+
+    _renderer()._truncate_expanded_prompt(engine_input, _params(None))
+
+    assert engine_input["prompt_token_ids"] == list(range(147))
+
+
+def test_mixed_prompt_embeds_truncation_covers_embeds_spans():
+    """Bound in _process_singleton, where prompt_embeds spans are already attached."""
+    model_config = ModelConfig(
+        "Qwen/Qwen3-ASR-0.6B-hf",
+        revision="7f1569a48a89f3e3f4dc3a5c9d28bddd903bc76c",
+        dtype="float32",
+        max_model_len=4096,
+        mm_processor_cache_gb=0,
+        enable_prompt_embeds=True,
+    )
+    renderer = renderer_from_config(VllmConfig(model_config=model_config))
+    tok = renderer.tokenizer
+    pe_id = _ensure_prompt_embeds_placeholder_token(tok)
+    hidden = model_config.get_hidden_size()
+    audio = np.zeros(4 * 16000, dtype=np.float32)
+
+    pe_len = 20
+    tensors = [torch.zeros(pe_len, hidden, dtype=torch.float32)]
+
+    def _build_prompt():
+        text = "<|audio_start|><|audio_pad|><|audio_end|>"
+        ids = tok.encode(text) + [pe_id] + tok.encode(" tail" * 40)
+        return {
+            "prompt_token_ids": ids,
+            "multi_modal_data": {"audio": (audio, 16000)},
+            "_prompt_embeds": (tensors, pe_id),
+        }
+
+    # A limit past every span (audio ends at 52, prompt_embeds at 74):
+    # truncation should succeed and bound the prompt.
+    out = renderer._process_singleton(
+        _build_prompt(),
+        tok_params=_params(79),
+    )
+    assert len(out["prompt_token_ids"]) == 79
+
+    # A limit inside the prompt_embeds span (54..74): must be rejected.
+    with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
+        renderer._process_singleton(
+            _build_prompt(),
+            tok_params=_params(64),
+        )
diff --git a/tests/renderers/test_warmup.py b/tests/renderers/test_warmup.py
index f6f9f7a4db6c..b0e3dcb54796 100644
--- a/tests/renderers/test_warmup.py
+++ b/tests/renderers/test_warmup.py
@@ -4,7 +4,9 @@
 
 These tests exercise:
   - Zero-limit modalities are filtered from mm_counts passed to
-    get_dummy_processor_inputs (e.g. --limit-mm-per-prompt image=0 ...)
+    get_dummy_inputs (e.g. --limit-mm-per-prompt image=0 ...)
+  - The warmup runs through the processor-only cache and clears it afterwards,
+    so its dummy inputs never land in the sender cache used by the serving path
   - MM warmup is skipped entirely when mm_processor is None
   - The multimodal warmup is launched as a task on the single-worker
     _mm_executor to overlap engine-core init (future lifecycle, join by
@@ -21,10 +23,13 @@
 """
 
 from concurrent.futures import Future, ThreadPoolExecutor
+from types import MethodType, SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 
+from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
+from vllm.multimodal.processing import BaseMultiModalProcessor
 from vllm.renderers.base import BaseRenderer
 from vllm.renderers.params import ChatParams
 
@@ -34,9 +39,11 @@ def _make_renderer_mock(mm_limits: dict[str, int]) -> MagicMock:
 
     render_chat is mocked to raise ChatTemplateResolutionError so the chat
     warmup block is skipped cleanly, keeping the test focused on MM warmup.
-    """
-    from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
 
+    The processor is a mock at its I/O boundaries (``dummy_inputs`` and
+    ``apply``) but keeps the real ``get_dummy_mm_inputs``, since that is where
+    the warmup's seq_len capping and cache plumbing actually live.
+    """
     renderer = MagicMock()
 
     # chat warmup: make render_chat raise so we skip past it cleanly
@@ -45,13 +52,17 @@ def _make_renderer_mock(mm_limits: dict[str, int]) -> MagicMock:
     # MM processor with configurable limits
     mm_processor = MagicMock()
     mm_processor.info.allowed_mm_limits = mm_limits
+    mm_processor.info.ctx.model_config = renderer.model_config
     mm_processor.apply.return_value = {"prompt_token_ids": [1]}
+    mm_processor.get_dummy_mm_inputs = MethodType(
+        BaseMultiModalProcessor.get_dummy_mm_inputs, mm_processor
+    )
     renderer.mm_processor = mm_processor
-    renderer._readonly_mm_processor = None
+    renderer._mm_processor_cache = MagicMock(name="sender_cache")
+    renderer._mm_processor_only_cache = MagicMock(name="processor_only_cache")
     renderer._warmup_mm_processor = BaseRenderer._warmup_mm_processor.__get__(
         renderer, BaseRenderer
     )
-    renderer._clear_processor_cache = BaseRenderer._clear_processor_cache
     renderer.warmup_mm = BaseRenderer.warmup_mm.__get__(renderer, BaseRenderer)
     renderer.start_mm_warmup_in_background = (
         BaseRenderer.start_mm_warmup_in_background.__get__(renderer, BaseRenderer)
@@ -64,7 +75,6 @@ def _make_renderer_mock(mm_limits: dict[str, int]) -> MagicMock:
     renderer._mm_warmup_future = None
     # MM warmup has not run yet; warmup_mm must actually execute on the mock.
     renderer._mm_warmup_done = False
-    renderer.clear_mm_cache = MagicMock()
     renderer.model_config.max_model_len = 128
     renderer.config.scheduler_config.max_num_batched_tokens = 8192
     renderer.config.scheduler_config.enable_chunked_prefill = True
@@ -80,10 +90,9 @@ def test_zero_limit_modality_excluded_from_mm_counts(self):
         """A modality with limit=0 must not appear in mm_counts."""
         renderer = _make_renderer_mock({"image": 1, "video": 0})
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())
+        BaseRenderer.warmup(renderer, ChatParams())
 
-        get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs
+        get_inputs = renderer.mm_processor.get_dummy_inputs
         get_inputs.assert_called_once()
         _, kwargs = get_inputs.call_args
         assert "video" not in kwargs["mm_counts"]
@@ -93,10 +102,9 @@ def test_all_zero_limits_passes_empty_mm_counts(self):
         """When all limits are 0, mm_counts must be empty."""
         renderer = _make_renderer_mock({"image": 0, "video": 0})
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())
+        BaseRenderer.warmup(renderer, ChatParams())
 
-        get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs
+        get_inputs = renderer.mm_processor.get_dummy_inputs
         get_inputs.assert_called_once()
         _, kwargs = get_inputs.call_args
         assert kwargs["mm_counts"] == {}
@@ -105,10 +113,9 @@ def test_positive_limits_all_included_in_mm_counts(self):
         """All modalities with limit > 0 must be present in mm_counts."""
         renderer = _make_renderer_mock({"image": 2, "video": 1})
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())
+        BaseRenderer.warmup(renderer, ChatParams())
 
-        get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs
+        get_inputs = renderer.mm_processor.get_dummy_inputs
         get_inputs.assert_called_once()
         _, kwargs = get_inputs.call_args
         assert kwargs["mm_counts"] == {"image": 1, "video": 1}
@@ -125,35 +132,52 @@ def test_warmup_sequence_length_is_bounded(self, max_model_len, expected_seq_len
         renderer = _make_renderer_mock({"image": 1, "video": 1})
         renderer.model_config.max_model_len = max_model_len
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())
+        BaseRenderer.warmup(renderer, ChatParams())
 
-        get_inputs = renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs
+        get_inputs = renderer.mm_processor.get_dummy_inputs
         assert get_inputs.call_args.kwargs["seq_len"] == expected_seq_len
 
     def test_processor_apply_called(self):
         renderer = _make_renderer_mock({"image": 1})
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())
+        BaseRenderer.warmup(renderer, ChatParams())
 
         renderer.mm_processor.apply.assert_called_once()
 
-    def test_mm_cache_cleared_after_warmup(self):
+    def test_render_chat_called_with_warmup_message(self):
         renderer = _make_renderer_mock({"image": 1})
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())
+        BaseRenderer.warmup(renderer, ChatParams())
 
-        renderer.clear_mm_cache.assert_called_once()
+        renderer.render_chat.assert_called_once()
 
-    def test_render_chat_called_with_warmup_message(self):
+
+class TestMmWarmupCacheRouting:
+    # The warmup runs through the processor-only cache and drops it afterwards,
+    # so the dummy inputs never reach the sender cache used by the serving path.
+
+    def test_warmup_passes_processor_only_cache(self):
         renderer = _make_renderer_mock({"image": 1})
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
+        with patch.object(
+            renderer.mm_processor,
+            "get_dummy_mm_inputs",
+            wraps=renderer.mm_processor.get_dummy_mm_inputs,
+        ) as spy:
             BaseRenderer.warmup(renderer, ChatParams())
 
-        renderer.render_chat.assert_called_once()
+        assert spy.call_args.kwargs["cache"] is renderer._mm_processor_only_cache
+        assert (
+            spy.call_args.kwargs["scheduler_config"] is renderer.config.scheduler_config
+        )
+
+    def test_only_cache_cleared_and_sender_cache_untouched(self):
+        renderer = _make_renderer_mock({"image": 1})
+
+        BaseRenderer.warmup(renderer, ChatParams())
+
+        renderer._mm_processor_only_cache.clear_cache.assert_called_once()
+        renderer._mm_processor_cache.clear_cache.assert_not_called()
 
 
 class TestMmWarmupSkippedWhenNoProcessor:
@@ -168,51 +192,28 @@ def test_no_warmup_without_processor(self):
         renderer.model_config.get_multimodal_config.assert_not_called()
 
 
-class TestReadonlyMmWarmup:
-    """Readonly MM processor warmup must mirror the render path behavior."""
-
-    def test_readonly_processor_apply_called_and_cache_cleared(self):
-        renderer = _make_renderer_mock({"image": 1})
-        readonly_mm_processor = MagicMock()
-        readonly_mm_processor.info.allowed_mm_limits = {"image": 1}
-        renderer._readonly_mm_processor = readonly_mm_processor
-
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())
-
-        readonly_mm_processor.apply.assert_called_once()
-        readonly_mm_processor.cache.clear_cache.assert_called_once()
-
-
 class TestWarmupFaultIsolation:
-    # A failure during a multimodal processor warmup is caught so it does not
-    # abort the remaining warmup steps; warmup itself must not raise.
+    # A failure during a warmup step is caught so it does not abort the other
+    # step; warmup itself must not raise.
 
     def test_chat_failure_does_not_abort_mm_warmup(self):
         renderer = _make_renderer_mock({"image": 1})
         renderer.render_chat.side_effect = RuntimeError("chat boom")
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())  # must not raise
+        BaseRenderer.warmup(renderer, ChatParams())  # must not raise
 
         renderer.mm_processor.apply.assert_called_once()
 
-    def test_mm_failure_does_not_abort_readonly_warmup(self):
+    def test_mm_failure_is_swallowed_and_cache_cleared(self):
         renderer = _make_renderer_mock({"image": 1})
-        readonly_mm_processor = MagicMock()
-        readonly_mm_processor.info.allowed_mm_limits = {"image": 1}
-        renderer._readonly_mm_processor = readonly_mm_processor
-        # main processor warmup blows up before apply()
-        renderer.mm_processor.dummy_inputs.get_dummy_processor_inputs.side_effect = (
-            RuntimeError("mm boom")
-        )
+        # the warmup blows up before apply()
+        renderer.mm_processor.get_dummy_inputs.side_effect = RuntimeError("mm boom")
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            BaseRenderer.warmup(renderer, ChatParams())  # must not raise
+        BaseRenderer.warmup(renderer, ChatParams())  # must not raise
 
-        readonly_mm_processor.apply.assert_called_once()
-        # cache is still cleared in the failed task's finally
-        renderer.clear_mm_cache.assert_called_once()
+        renderer.mm_processor.apply.assert_not_called()
+        # the cache is still cleared in the failed warmup's finally
+        renderer._mm_processor_only_cache.clear_cache.assert_called_once()
 
 
 class TestBackgroundMmWarmup:
@@ -234,16 +235,15 @@ def _make_bg_renderer(self, mm_limits: dict[str, int]):
     def test_start_submits_future_to_mm_executor(self):
         renderer = self._make_bg_renderer({"image": 1})
         try:
-            with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-                # Spy on the real executor's submit to confirm the warmup is
-                # dispatched through _mm_executor (not a separate Thread).
-                with patch.object(
-                    renderer._mm_executor, "submit", wraps=renderer._mm_executor.submit
-                ) as spy:
-                    renderer.start_mm_warmup_in_background()
-                    assert isinstance(renderer._mm_warmup_future, Future)
-                    assert spy.called
-                renderer._join_mm_warmup()
+            # Spy on the real executor's submit to confirm the warmup is
+            # dispatched through _mm_executor (not a separate Thread).
+            with patch.object(
+                renderer._mm_executor, "submit", wraps=renderer._mm_executor.submit
+            ) as spy:
+                renderer.start_mm_warmup_in_background()
+                assert isinstance(renderer._mm_warmup_future, Future)
+                assert spy.called
+            renderer._join_mm_warmup()
             assert renderer._mm_warmup_future is None
             renderer.mm_processor.apply.assert_called_once()
         finally:
@@ -252,7 +252,6 @@ def test_start_submits_future_to_mm_executor(self):
     def test_start_noop_for_text_only_model(self):
         renderer = _make_renderer_mock({})
         renderer.mm_processor = None
-        # _readonly_mm_processor is already None
 
         renderer.start_mm_warmup_in_background()
 
@@ -261,12 +260,11 @@ def test_start_noop_for_text_only_model(self):
     def test_start_is_run_at_most_once(self):
         renderer = self._make_bg_renderer({"image": 1})
         try:
-            with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-                renderer.start_mm_warmup_in_background()
-                first = renderer._mm_warmup_future
-                renderer.start_mm_warmup_in_background()  # must not spawn a second
-                assert renderer._mm_warmup_future is first
-                renderer._join_mm_warmup()
+            renderer.start_mm_warmup_in_background()
+            first = renderer._mm_warmup_future
+            renderer.start_mm_warmup_in_background()  # must not spawn a second
+            assert renderer._mm_warmup_future is first
+            renderer._join_mm_warmup()
         finally:
             renderer._mm_executor.shutdown(wait=True)
 
@@ -275,9 +273,8 @@ def test_warmup_joins_background_and_does_not_rerun_mm(self):
         # run only the chat warmup — the MM warmup must not run twice.
         renderer = self._make_bg_renderer({"image": 1})
         try:
-            with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-                renderer.start_mm_warmup_in_background()
-                BaseRenderer.warmup(renderer, ChatParams())
+            renderer.start_mm_warmup_in_background()
+            BaseRenderer.warmup(renderer, ChatParams())
 
             # MM apply called exactly once (by the background warmup task).
             renderer.mm_processor.apply.assert_called_once()
@@ -294,13 +291,12 @@ def test_warmup_does_not_rerun_mm_after_reset_joins_background(self):
         # re-run the MM warmup — the _mm_warmup_done flag survives the join.
         renderer = self._make_bg_renderer({"image": 1})
         try:
-            with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-                renderer.start_mm_warmup_in_background()
-                # Simulate reset_mm_cache joining the background warmup first
-                # (clears _mm_warmup_future, sets _mm_warmup_done via the join).
-                renderer._join_mm_warmup()
-                assert renderer._mm_warmup_future is None  # joined & cleared
-                BaseRenderer.warmup(renderer, ChatParams())
+            renderer.start_mm_warmup_in_background()
+            # Simulate reset_mm_cache joining the background warmup first
+            # (clears _mm_warmup_future, sets _mm_warmup_done via the join).
+            renderer._join_mm_warmup()
+            assert renderer._mm_warmup_future is None  # joined & cleared
+            BaseRenderer.warmup(renderer, ChatParams())
 
             # MM apply called exactly once (by the background warmup task);
             # warmup() did not re-run warmup_mm despite the future being cleared.
@@ -313,9 +309,8 @@ def test_warmup_mm_runs_at_most_once(self):
         # Direct repeated calls to warmup_mm run the MM warmup only once.
         renderer = _make_renderer_mock({"image": 1})
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            renderer.warmup_mm()
-            renderer.warmup_mm()  # second call must be a no-op
+        renderer.warmup_mm()
+        renderer.warmup_mm()  # second call must be a no-op
 
         renderer.mm_processor.apply.assert_called_once()
 
@@ -324,10 +319,9 @@ def test_shutdown_joins_background_warmup(self):
         # so the mm_processor_cache is never touched concurrently.
         renderer = self._make_bg_renderer({"image": 1})
 
-        with patch("vllm.multimodal.processing.TimingContext", autospec=True):
-            renderer.start_mm_warmup_in_background()
-            future = renderer._mm_warmup_future
-            renderer.shutdown()
+        renderer.start_mm_warmup_in_background()
+        future = renderer._mm_warmup_future
+        renderer.shutdown()
 
         assert renderer._mm_warmup_future is None
         # The background warmup was allowed to complete (apply ran) and the
@@ -344,10 +338,6 @@ class TestEngineStartWarmupHook:
     # These tests pin the renderer plumbing contract on EngineCoreClient.
 
     def _mock_config(self):
-        from types import SimpleNamespace
-
-        from vllm.v1.engine.core_client import EngineCoreClient  # noqa: F401
-
         cfg = SimpleNamespace(
             parallel_config=SimpleNamespace(
                 data_parallel_size=1,
diff --git a/tests/samplers/conftest.py b/tests/samplers/conftest.py
new file mode 100644
index 000000000000..4585e98898f7
--- /dev/null
+++ b/tests/samplers/conftest.py
@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm import envs
+from vllm.platforms import current_platform
+
+
+@pytest.fixture(autouse=True)
+def skip_unsupported_flashinfer_sampler():
+    # CI runs this suite with FlashInfer both explicitly enabled and disabled.
+    if (
+        current_platform.is_cuda()
+        and envs.is_set("VLLM_USE_FLASHINFER_SAMPLER")
+        and envs.VLLM_USE_FLASHINFER_SAMPLER
+        and current_platform.num_compute_units(torch.accelerator.current_device_index())
+        <= 16
+    ):
+        pytest.skip("FlashInfer top-k masking requires more than 16 SMs")
diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py
index 91dee5634b52..8bcf4c48ab59 100644
--- a/tests/samplers/test_beam_search.py
+++ b/tests/samplers/test_beam_search.py
@@ -6,13 +6,17 @@
 """
 
 import json
+from types import SimpleNamespace
+from unittest.mock import Mock
 
 import jsonschema
 import pytest
 from transformers import AutoModelForSeq2SeqLM
 
+from vllm import CompletionOutput, RequestOutput
 from vllm.assets.audio import AudioAsset
 from vllm.entrypoints.llm import LLM
+from vllm.logprobs import Logprob, SampleLogprobs
 from vllm.platforms import current_platform
 from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams
 
@@ -45,6 +49,87 @@
 MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
 
 
+@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)])
+@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]])
+def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts(
+    monkeypatch,
+    abort_after: int,
+    prompt_token: int,
+    terminal_logprobs: SampleLogprobs | None,
+) -> None:
+    """Return tokens and scores from the last completed step for aborted prompts.
+
+    Other prompts in the batch continue generating.
+    """
+
+    def run_requests(prompts, **kwargs):
+        results = []
+        for prompt in prompts:
+            tokens = prompt["prompt_token_ids"]
+            if tokens[0] == prompt_token:
+                assert len(tokens) <= abort_after + 1, "Continued after abort"
+            aborted = (
+                tokens[0] == prompt_token
+                and len(tokens) == abort_after + 1
+                and tokens[-1] != 12
+            )
+            results.append(
+                RequestOutput(
+                    request_id="inner",
+                    prompt=None,
+                    prompt_token_ids=tokens,
+                    prompt_logprobs=None,
+                    finished=True,
+                    outputs=[
+                        CompletionOutput(
+                            index=0,
+                            text="",
+                            token_ids=[] if aborted else [11],
+                            cumulative_logprob=None,
+                            logprobs=terminal_logprobs
+                            if aborted
+                            else [{11: Logprob(-0.1), 12: Logprob(-0.2)}],
+                            finish_reason="abort" if aborted else "length",
+                        )
+                    ],
+                )
+            )
+        return results
+
+    llm = LLM.__new__(LLM)
+    llm.llm_engine = Mock()
+    tokenizer = SimpleNamespace(
+        eos_token_id=0,
+        decode=lambda tokens, skip_special_tokens=False: " ".join(
+            str(token) for token in tokens if not skip_special_tokens or token != 0
+        ),
+    )
+    llm.renderer = Mock(get_tokenizer=Mock(return_value=tokenizer))
+    monkeypatch.setattr(llm, "_preprocess_cmpl", lambda prompts: prompts)
+    monkeypatch.setattr(llm, "_render_and_run_requests", run_requests)
+    prompts = [
+        {"type": "token", "prompt_token_ids": [token]} for token in [prompt_token, 2]
+    ]
+    params = BeamSearchParams(beam_width=2, max_tokens=3)
+    outputs = llm.beam_search(prompts, params)
+
+    aborted, normal = outputs
+    expected_tokens = (
+        [[prompt_token, 11], [prompt_token, 12]] if abort_after else [[prompt_token]]
+    )
+    assert [beam.tokens for beam in aborted.sequences] == expected_tokens
+    for beam in aborted.sequences:
+        assert beam.finish_reason == "abort"
+        assert beam.text == tokenizer.decode(
+            beam.tokens, skip_special_tokens=params.skip_special_tokens
+        )
+        assert len(beam.logprobs) == abort_after
+    assert aborted.sequences[0].cum_logprob == pytest.approx(-0.1 * abort_after)
+    assert len(normal.sequences) == 2
+    assert all(len(beam.tokens) == 4 for beam in normal.sequences)
+    assert all(beam.finish_reason != "abort" for beam in normal.sequences)
+
+
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
 @pytest.mark.parametrize("max_tokens", MAX_TOKENS)
diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py
index 14481f79c68a..ccb40cc65906 100644
--- a/tests/samplers/test_beam_search_online.py
+++ b/tests/samplers/test_beam_search_online.py
@@ -5,14 +5,19 @@
 
 from vllm import CompletionOutput, RequestOutput
 from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin
-from vllm.logprobs import Logprob
+from vllm.logprobs import Logprob, SampleLogprobs
 from vllm.sampling_params import BeamSearchParams
 
 
 class _Tokenizer:
     eos_token_id = 0
+    special_token_id = eos_token_id
 
-    def decode(self, token_ids: list[int]) -> str:
+    def decode(self, token_ids: list[int], skip_special_tokens: bool = False) -> str:
+        if skip_special_tokens:
+            token_ids = [
+                token_id for token_id in token_ids if token_id != self.special_token_id
+            ]
         return " ".join(str(token_id) for token_id in token_ids)
 
 
@@ -72,3 +77,86 @@ async def test_beam_search_handles_extra_logprob_candidates() -> None:
     assert outputs[0].outputs[0].finish_reason == "stop"
     assert outputs[0].outputs[0].token_ids == []
     assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("skip_special_tokens", "expected_text"),
+    [
+        pytest.param(True, "", id="skip"),
+        pytest.param(False, "0", id="keep"),
+    ],
+)
+async def test_beam_search_respects_skip_special_tokens(
+    skip_special_tokens: bool, expected_text: str
+) -> None:
+    prompt = {
+        "type": "token",
+        "prompt": "prompt",
+        "prompt_token_ids": [1],
+    }
+    params = BeamSearchParams(
+        beam_width=1,
+        max_tokens=1,
+        ignore_eos=True,
+        skip_special_tokens=skip_special_tokens,
+    )
+
+    outputs = [
+        output async for output in _Serving().beam_search(prompt, "request", params)
+    ]
+
+    assert outputs[0].outputs[0].text == expected_text
+    assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)])
+@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]])
+async def test_beam_search_abort_returns_partial_outputs(
+    monkeypatch,
+    abort_after: int,
+    prompt_token: int,
+    terminal_logprobs: SampleLogprobs | None,
+) -> None:
+    """Abort returns each beam's tokens and scores from the last completed step."""
+    calls = 0
+
+    async def generate(prompt, *args, **kwargs):
+        nonlocal calls
+        calls += 1
+        depth = len(prompt["prompt_token_ids"]) - 1
+        assert depth <= abort_after, "Beam search continued after abort"
+        result = await anext(_EngineClient().generate(prompt, *args, **kwargs))
+        output = result.outputs[0]
+        output.token_ids = [11]
+        output.logprobs = [{11: Logprob(-0.1), 12: Logprob(-0.2)}]
+        output.finish_reason = "length"
+        if depth == abort_after and prompt["prompt_token_ids"][-1] != 12:
+            output.token_ids = []
+            output.logprobs = terminal_logprobs
+            output.finish_reason = "abort"
+        yield result
+
+    serving = _Serving()
+    monkeypatch.setattr(serving.engine_client, "generate", generate)
+    prompt = {"type": "token", "prompt": "prompt", "prompt_token_ids": [prompt_token]}
+    outputs = [
+        output
+        async for output in serving.beam_search(
+            prompt, "request", BeamSearchParams(beam_width=2, max_tokens=4)
+        )
+    ]
+
+    assert calls == 1 + 2 * abort_after
+    assert len(outputs) == 1
+    assert outputs[0].finished
+    assert outputs[0].request_id == "request"
+    expected_tokens = [[11], [12]] if abort_after else [[]]
+    assert [output.token_ids for output in outputs[0].outputs] == expected_tokens
+    for output in outputs[0].outputs:
+        assert output.finish_reason == "abort"
+        assert output.text == _Tokenizer().decode(list(output.token_ids))
+        assert output.logprobs is not None
+        assert len(output.logprobs) == abort_after
+    assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1 * abort_after)
diff --git a/tests/spec_decode/test_custom_proposer.py b/tests/spec_decode/test_custom_proposer.py
index 1e42bde8104d..2e9282af744b 100755
--- a/tests/spec_decode/test_custom_proposer.py
+++ b/tests/spec_decode/test_custom_proposer.py
@@ -29,6 +29,7 @@ def __init__(self, vllm_config: VllmConfig):
 
         Args:
             vllm_config: vLLM configuration containing model and speculative settings.
+
         """
         assert vllm_config.speculative_config is not None
         self.num_speculative_tokens = (
@@ -57,6 +58,7 @@ def propose(
 
         Returns:
             List of draft token sequences for each request.
+
         """
         # Cross-process flag to prove this method was executed
         with open("proposer_called.flag", "w") as f:
diff --git a/tests/test_access_log_filter.py b/tests/test_access_log_filter.py
index a28771ce59ab..4f03892e7753 100644
--- a/tests/test_access_log_filter.py
+++ b/tests/test_access_log_filter.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for the UvicornAccessLogFilter class.
-"""
+"""Tests for the UvicornAccessLogFilter class."""
 
 import logging
 
diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py
index dd0c6261d349..5d217433db71 100644
--- a/tests/test_cmake_utils.py
+++ b/tests/test_cmake_utils.py
@@ -38,7 +38,7 @@ def test_exact_family_arch_precedes_generic_family_fallback(tmp_path: Path):
 def test_extract_archs_prefers_sass_target_over_corrupted_virtual_arch(
     tmp_path: Path,
 ):
-    """torch's autodetection can emit a bogus arch=compute_* half (e.g.
+    """Torch's autodetection can emit a bogus arch=compute_* half (e.g.
     capability 12.1 corrupted to arch=compute_20,code=sm_121); the SASS
     target must win, while PTX-only entries keep the virtual arch."""
     repo_root = Path(__file__).parents[1]
diff --git a/tests/test_config.py b/tests/test_config.py
index 96ca8f36e4ed..c391216d9ab0 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -12,6 +12,7 @@
 
 import pydantic
 import pytest
+import torch
 from huggingface_hub import ResolvedRevision
 from pydantic import ValidationError
 
@@ -19,10 +20,12 @@
 import vllm.envs as envs
 from vllm.compilation.backends import VllmBackend
 from vllm.config import (
+    AttentionConfig,
     CacheConfig,
     CompilationConfig,
     DeviceConfig,
     EngramConfig,
+    HiSparseConfig,
     KernelConfig,
     KVTransferConfig,
     ModelConfig,
@@ -35,7 +38,7 @@
     WatermarkConfig,
     update_config,
 )
-from vllm.config.compilation import CompilationMode, CUDAGraphMode
+from vllm.config.compilation import CompilationMode, CUDAGraphMode, PassConfig
 from vllm.config.kernel import IrOpPriorityConfig
 from vllm.config.load import LoadConfig
 from vllm.config.mamba import MambaBackendEnum
@@ -44,6 +47,7 @@
 from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel
 from vllm.platforms import current_platform
 from vllm.transformers_utils.config import (
+    _patch_hf_transformers_nested_rope_validation,
     get_pooling_config,
     try_get_dense_modules,
 )
@@ -52,11 +56,86 @@
 DEVICE_TYPE = current_platform.device_type
 
 
+def test_nested_rope_validation_patch_preserves_flat_rope_parameters(monkeypatch):
+    calls = []
+
+    def original_validate_rope(config, *args, **kwargs):
+        calls.append(config)
+
+    from transformers import PreTrainedConfig
+
+    monkeypatch.setattr(PreTrainedConfig, "validate_rope", original_validate_rope)
+    _patch_hf_transformers_nested_rope_validation()
+
+    nested_rope_parameters = {
+        "full_attention": {"rope_type": "default"},
+        "original_max_position_embeddings": 32768,
+    }
+    PreTrainedConfig.validate_rope(
+        SimpleNamespace(rope_parameters=nested_rope_parameters)
+    )
+    assert nested_rope_parameters == {"full_attention": {"rope_type": "default"}}
+
+    flat_rope_parameters = {
+        "rope_type": "linear",
+        "factor": 8.0,
+        "rope_theta": 500000.0,
+    }
+    PreTrainedConfig.validate_rope(
+        SimpleNamespace(rope_parameters=flat_rope_parameters)
+    )
+    assert flat_rope_parameters == {
+        "rope_type": "linear",
+        "factor": 8.0,
+        "rope_theta": 500000.0,
+    }
+    assert len(calls) == 2
+
+
 def _write_json(path: Path, value: object) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
     path.write_text(json.dumps(value), encoding="utf-8")
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test")
+@pytest.mark.parametrize(
+    ("is_mm_prefix_lm", "is_multimodal_model", "expected"),
+    [
+        pytest.param(True, True, True, id="multimodal-prefix-lm"),
+        pytest.param(False, True, False, id="multimodal-causal"),
+        pytest.param(True, False, False, id="text-prefix-lm"),
+        pytest.param(None, True, False, id="missing-model-config"),
+    ],
+)
+def test_rocm_mm_prefix_lm_disables_chunked_mm_input(
+    is_mm_prefix_lm: bool | None,
+    is_multimodal_model: bool,
+    expected: bool,
+) -> None:
+    from vllm.platforms.rocm import RocmPlatform
+
+    config = SimpleNamespace(
+        compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE),
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=1,
+            worker_cls="test-worker",
+        ),
+        model_config=(
+            None
+            if is_mm_prefix_lm is None
+            else SimpleNamespace(is_mm_prefix_lm=is_mm_prefix_lm)
+        ),
+        scheduler_config=SimpleNamespace(
+            is_multimodal_model=is_multimodal_model,
+            disable_chunked_mm_input=False,
+        ),
+    )
+
+    RocmPlatform.check_and_update_config(config)
+
+    assert config.scheduler_config.disable_chunked_mm_input is expected
+
+
 def test_kda_recoverssm_derivation_is_revalidated():
     config = SimpleNamespace(
         cache_config=SimpleNamespace(
@@ -211,6 +290,46 @@ def test_kv_offloading_does_not_skip_dcp_interleave_validation():
         VllmConfig.validate_block_size(config)
 
 
+def test_nixl_dcp_check_skipped_for_submodel_config():
+    # with_hf_config() builds a submodule view of the config (e.g. a
+    # multimodal model's text stack) whose architecture list is empty.
+    # Re-running VllmConfig validation on it is unsafe: the NIXL DCP check
+    # reads use_mla, which resolves the architecture registry.
+    model_config = ModelConfig("Qwen/Qwen2-VL-2B-Instruct", max_model_len=2048)
+    vllm_config = VllmConfig(
+        model_config=model_config,
+        device_config=DeviceConfig(device="cpu"),
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="NixlConnector",
+            kv_role="kv_both",
+        ),
+    )
+    submodel_config = vllm_config.with_hf_config(model_config.hf_text_config)
+    assert submodel_config.model_config.is_submodel_config
+    assert submodel_config.model_config.architectures == []
+    assert submodel_config.model_config.use_mla is False
+
+
+def test_nixl_dcp_check_rejects_non_mla_model_with_dcp(monkeypatch):
+    # Pretend the model has a single KV head so the DCP feasibility checks
+    # pass and the MLA-only assert is what actually fires.
+    monkeypatch.setattr(ModelConfig, "get_total_num_kv_heads", lambda self: 1)
+    with pytest.raises(ValidationError, match="only supported for MLA models"):
+        VllmConfig(
+            model_config=ModelConfig("Qwen/Qwen3-0.6B", max_model_len=2048),
+            device_config=DeviceConfig(device="cpu"),
+            parallel_config=ParallelConfig(
+                tensor_parallel_size=2,
+                decode_context_parallel_size=2,
+                distributed_executor_backend="mp",
+            ),
+            kv_transfer_config=KVTransferConfig(
+                kv_connector="NixlConnector",
+                kv_role="kv_both",
+            ),
+        )
+
+
 def test_compile_config_repr_succeeds():
     # setup: VllmBackend mutates the config object
     config = VllmConfig()
@@ -240,6 +359,164 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
     assert envs.VLLM_USE_V2_MODEL_RUNNER is expected
 
 
+def test_hisparse_requires_v2_model_runner():
+    config = object.__new__(VllmConfig)
+    config.attention_config = AttentionConfig(hisparse_config=HiSparseConfig())
+
+    with patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", None):
+        assert config.use_v2_model_runner
+    with (
+        patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", False),
+        pytest.raises(ValueError, match="requires Model Runner V2"),
+    ):
+        _ = config.use_v2_model_runner
+
+
+def test_hisparse_rejects_decode_context_parallelism(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(current_platform, "device_count", lambda: 2)
+    with pytest.raises(ValueError, match="decode context parallelism"):
+        VllmConfig(
+            attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
+            parallel_config=ParallelConfig(
+                tensor_parallel_size=2,
+                decode_context_parallel_size=2,
+            ),
+        )
+
+
+def test_hisparse_rejects_pipeline_parallelism(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(current_platform, "device_count", lambda: 2)
+    with pytest.raises(ValueError, match="pipeline parallelism"):
+        VllmConfig(
+            attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
+            parallel_config=ParallelConfig(pipeline_parallel_size=2),
+        )
+
+
+def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
+    with pytest.raises(ValueError, match="requires the hybrid KV cache manager"):
+        VllmConfig(
+            attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
+            scheduler_config=SchedulerConfig(
+                max_model_len=2048,
+                is_encoder_decoder=False,
+                disable_hybrid_kv_cache_manager=True,
+            ),
+        )
+
+
+def test_hisparse_rejects_non_cuda(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: False)
+    with pytest.raises(ValueError, match="requires NVIDIA CUDA"):
+        VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig()))
+
+
+@pytest.mark.parametrize(
+    "kv_transfer_config",
+    [
+        KVTransferConfig(
+            kv_connector="HiSparseConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={"host_pool_gib": 128},
+        ),
+        KVTransferConfig(
+            kv_connector="MultiConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={
+                "connectors": [
+                    {
+                        "kv_connector": "OffloadingConnector",
+                        "kv_role": "kv_both",
+                        "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
+                    },
+                    {
+                        "kv_connector": "HiSparseConnector",
+                        "kv_role": "kv_both",
+                        "kv_connector_extra_config": {"host_pool_gib": 128},
+                    },
+                ]
+            },
+        ),
+        KVTransferConfig(
+            kv_connector="MultiConnector",
+            kv_role="kv_consumer",
+            kv_connector_extra_config={
+                "connectors": [
+                    {"kv_connector": "NixlConnector", "kv_role": "kv_consumer"},
+                    {
+                        "kv_connector": "HiSparseConnector",
+                        "kv_role": "kv_both",
+                        "kv_connector_extra_config": {"host_pool_gib": 128},
+                    },
+                ]
+            },
+        ),
+    ],
+    ids=["standalone", "multi-connector", "pd-decode"],
+)
+def test_hisparse_connector_implies_attention_config(monkeypatch, kv_transfer_config):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
+    config = VllmConfig(
+        kv_transfer_config=kv_transfer_config,
+        # Skip the HMA auto-detect block: it imports the connector class,
+        # which pulls the platform's compiled attention extensions.
+        scheduler_config=SchedulerConfig(
+            max_model_len=2048,
+            is_encoder_decoder=False,
+            disable_hybrid_kv_cache_manager=False,
+        ),
+    )
+    assert isinstance(config.attention_config.hisparse_config, HiSparseConfig)
+
+
+def test_hisparse_connector_preserves_explicit_attention_config(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
+    config = VllmConfig(
+        attention_config=AttentionConfig(
+            hisparse_config=HiSparseConfig(device_buffer_size=512)
+        ),
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="HiSparseConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={"host_pool_gib": 128},
+        ),
+        scheduler_config=SchedulerConfig(
+            max_model_len=2048,
+            is_encoder_decoder=False,
+            disable_hybrid_kv_cache_manager=False,
+        ),
+    )
+    assert config.attention_config.hisparse_config.device_buffer_size == 512
+
+
+def test_hisparse_connector_without_cuda_still_rejected(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: False)
+    with pytest.raises(ValueError, match="requires NVIDIA CUDA"):
+        VllmConfig(
+            kv_transfer_config=KVTransferConfig(
+                kv_connector="HiSparseConnector",
+                kv_role="kv_both",
+                kv_connector_extra_config={"host_pool_gib": 128},
+            )
+        )
+
+
+def test_no_hisparse_connector_keeps_attention_config_unset(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    config = VllmConfig(
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="OffloadingConnector", kv_role="kv_both"
+        )
+    )
+    assert config.attention_config.hisparse_config is None
+
+
 def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):
     """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled
     MRV1 paths and off breakable cudagraphs by default."""
@@ -261,12 +538,17 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):
         assert "DeepseekV32ForCausalLM" not in breakable_architectures
         assert "DeepseekV32MTPModel" not in breakable_architectures
         assert "GlmMoeDsaForCausalLM" not in breakable_architectures
+        # V4.1 cannot torch.compile and the ROCm sparse SWA backend only
+        # supports uniform-batch CUDA graphs, so it must opt in or the
+        # default FULL_AND_PIECEWISE serve path cannot start.
+        assert "DeepseekV41ForCausalLM" in breakable_architectures
 
         # The carve-out takes effect via the runner-selection property
         # (warning_once args must be hashable for its lru_cache).
         monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
         config = SimpleNamespace(
-            model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"])
+            model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]),
+            attention_config=AttentionConfig(),
         )
         assert VllmConfig.use_v2_model_runner.fget(config) is False
     finally:
@@ -308,6 +590,7 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(
     )
     config = SimpleNamespace(
         model_config=model_config,
+        attention_config=AttentionConfig(),
         speculative_config=SimpleNamespace(method="mtp") if with_mtp else None,
         parallel_config=SimpleNamespace(prefill_context_parallel_size=1),
         compilation_config=CompilationConfig(
@@ -374,6 +657,70 @@ def test_breakable_cudagraph_platform_default(
         default_breakable_cudagraph_architectures.cache_clear()
 
 
+@pytest.mark.parametrize(
+    "case,hidden,heads,intermediate,tp,expected",
+    [
+        ("bf16", 2048, (16, 8), 6144, 1, True),
+        ("bi-off", 2048, (16, 8), 6144, 1, False),
+        ("opt-out", 2048, (16, 8), 6144, 1, False),
+        ("eager", 2048, (16, 8), 6144, 1, False),
+        ("sp", 2048, (16, 8), 6144, 1, False),
+        ("fp16", 2048, (16, 8), 6144, 1, False),
+        ("quantized", 2048, (16, 8), 6144, 1, False),
+        ("untuned", 4096, (32, 32), 11008, 1, False),
+        ("tp4-tuned", 4096, (8, 2), 12288, 4, True),
+        ("list-intermediate", 2048, (16, 8), [2048, 4096], 1, False),
+        ("no-table", 2048, None, 6144, 1, False),
+    ],
+)
+def test_batch_invariant_breakable_cudagraph(
+    monkeypatch, case, hidden, heads, intermediate, tp, expected
+):
+    from vllm.config.vllm import default_breakable_cudagraph_architectures
+    from vllm.model_executor.determinism import batch_invariant_configs as bi_configs
+
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0" if case == "bi-off" else "1")
+    monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False)
+    if case == "opt-out":
+        monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "0")
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
+    table = None if case == "no-table" else {(12288, 2048): None, (6144, 4096): None}
+    monkeypatch.setattr(current_platform, "get_device_capability", lambda: None)
+    monkeypatch.setattr(bi_configs, "_get_tuned_matmul_arch_family", lambda cap: "test")
+    monkeypatch.setattr(
+        bi_configs,
+        "_BATCH_INVARIANT_MATMUL_TUNED_CONFIGS",
+        {"test": table} if table else {},
+    )
+    default_breakable_cudagraph_architectures.cache_clear()
+    config = object.__new__(VllmConfig)
+    config.model_config = SimpleNamespace(
+        architectures=["Qwen3ForCausalLM"],
+        enforce_eager=case == "eager",
+        dtype=torch.float16 if case == "fp16" else torch.bfloat16,
+        quantization="fp8" if case == "quantized" else None,
+        hf_text_config=SimpleNamespace(intermediate_size=intermediate),
+        get_hidden_size=lambda: hidden,
+        get_head_size=lambda: 128,
+        get_num_attention_heads=lambda pc: heads[0],
+        get_num_kv_heads=lambda pc: heads[1],
+    )
+    config.parallel_config = SimpleNamespace(tensor_parallel_size=tp)
+    config.compilation_config = (
+        CompilationConfig(pass_config=PassConfig(enable_sp=True))
+        if case == "sp"
+        else CompilationConfig()
+    )
+    try:
+        assert config._maybe_enable_breakable_cudagraph() is expected
+        if expected:
+            assert config.compilation_config.mode == CompilationMode.NONE
+    finally:
+        os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None)
+        default_breakable_cudagraph_architectures.cache_clear()
+
+
 @pytest.mark.parametrize(
     ("model_type", "expected_architecture"),
     [
@@ -383,9 +730,9 @@ def test_breakable_cudagraph_platform_default(
     ],
 )
 def test_dsa_models_select_matching_mtp(model_type, expected_architecture):
-    from transformers import PretrainedConfig
+    from transformers import PreTrainedConfig
 
-    hf_config = PretrainedConfig(
+    hf_config = PreTrainedConfig(
         architectures=["DeepseekV32ForCausalLM"],
         num_nextn_predict_layers=1,
     )
@@ -410,6 +757,15 @@ def test_v2_model_runner_supports_extract_hidden_states():
     assert config._get_v2_model_runner_unsupported_features() == []
 
 
+def test_v2_model_runner_supports_custom_logits_processors():
+    config = VllmConfig()
+    config.model_config = cast(
+        ModelConfig, SimpleNamespace(logits_processors=["a.b:C"])
+    )
+
+    assert config._get_v2_model_runner_unsupported_features() == []
+
+
 def test_dflash2_draft_forces_v2_model_runner():
     """A DFlash2 draft must reach the V2 speculator, the only one that runs its
     candidate selector; on V1 it would draft as DFlash1 without raising."""
@@ -463,6 +819,66 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(
     assert compilation_config.cudagraph_capture_sizes == expected_capture_sizes
 
 
+@pytest.mark.parametrize(
+    ("mode", "piecewise_capture_available", "attention_support", "expected"),
+    [
+        ("PIECEWISE", False, "ALWAYS", "NONE"),
+        ("FULL_AND_PIECEWISE", False, "ALWAYS", "FULL_DECODE_ONLY"),
+        ("FULL_DECODE_ONLY", False, "ALWAYS", "FULL_DECODE_ONLY"),
+        ("FULL_DECODE_ONLY", False, "NEVER", "NONE"),
+    ],
+)
+def test_resolve_cudagraph_mode_uses_loaded_piecewise_provider(
+    mode, piecewise_capture_available, attention_support, expected
+):
+    compilation_config = CompilationConfig(
+        mode=CompilationMode.VLLM_COMPILE,
+        cudagraph_mode=CUDAGraphMode[mode],
+        use_inductor_graph_partition=True,
+    )
+
+    resolved = compilation_config.resolve_cudagraph_mode_and_sizes(
+        AttentionCGSupport[attention_support],
+        "FakeAttentionBackend",
+        piecewise_capture_available=piecewise_capture_available,
+    )
+
+    assert resolved.name == expected
+    assert compilation_config.cudagraph_mode == resolved
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="Requires CUDA graph support"
+)
+@pytest.mark.parametrize(
+    "engine_kwargs",
+    [
+        {"runner": "pooling", "convert": "embed"},
+        pytest.param(
+            {"prefill_context_parallel_size": 2, "tensor_parallel_size": 2},
+            marks=pytest.mark.skipif(
+                not current_platform.is_rocm(), reason="ROCm PCP graph restriction"
+            ),
+        ),
+    ],
+)
+def test_late_piecewise_restrictions_without_compilation(monkeypatch, engine_kwargs):
+    """Late compatibility overrides must not restore unavailable piecewise graphs."""
+    from vllm.engine.arg_utils import EngineArgs
+
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "0")
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    config = EngineArgs(
+        model="facebook/opt-125m",
+        compilation_config=CompilationConfig(mode=CompilationMode.NONE),
+        **engine_kwargs,
+    ).create_engine_config()
+
+    assert config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
+    assert config.compilation_config.cudagraph_capture_sizes == []
+    assert config.compilation_config.max_cudagraph_capture_size == 0
+
+
 def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():
     """Cudagraph memory profiling uses a minimal KV cache, so the Mamba
     block-count guard must only fire for the real cache sizing."""
@@ -498,6 +914,30 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():
     assert cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE
 
 
+@pytest.mark.parametrize(
+    ("graph_mode", "should_raise"),
+    [
+        (CUDAGraphMode.FULL_DECODE_ONLY, False),
+        (CUDAGraphMode.FULL_AND_PIECEWISE, False),
+        (CUDAGraphMode.NONE, True),
+        (CUDAGraphMode.PIECEWISE, True),
+    ],
+)
+def test_adaptive_verification_requires_full_cudagraphs(graph_mode, should_raise):
+    config = SimpleNamespace(
+        speculative_config=SimpleNamespace(enable_adaptive_verification=True),
+        lora_config=None,
+        compilation_config=CompilationConfig(cudagraph_mode=graph_mode),
+        parallel_config=SimpleNamespace(pipeline_parallel_size=1),
+    )
+
+    if should_raise:
+        with pytest.raises(ValueError, match="requires full CUDA graphs"):
+            VllmConfig._validate_adaptive_verification(config)
+    else:
+        VllmConfig._validate_adaptive_verification(config)
+
+
 @pytest.mark.parametrize(
     ("model_config", "expected"),
     [
@@ -736,7 +1176,10 @@ def test_models_default_to_v2_model_runner(model_config, expected, monkeypatch):
     monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
     monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True)
     monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
-    config = SimpleNamespace(model_config=model_config)
+    config = SimpleNamespace(
+        model_config=model_config,
+        attention_config=AttentionConfig(),
+    )
     config._get_v2_model_runner_unsupported_features = lambda: []
 
     assert VllmConfig.use_v2_model_runner.fget(config) is expected
@@ -744,10 +1187,11 @@ def test_models_default_to_v2_model_runner(model_config, expected, monkeypatch):
 
 def test_v1_model_runner_rejects_v2_only_features():
     config = SimpleNamespace(
-        parallel_config=SimpleNamespace(
+        parallel_config=ParallelConfig(
             prefill_context_parallel_size=2,
-            enable_batch_sharded_sampling=False,
+            distributed_executor_backend="mp",
         ),
+        scheduler_config=SchedulerConfig.default_factory(async_scheduling=False),
         speculative_config=None,
         model_config=None,
     )
@@ -907,6 +1351,47 @@ def test_async_scheduling_with_pipeline_parallelism_is_allowed():
     assert cfg.scheduler_config.async_scheduling is True
 
 
+def test_v1_model_runner_drops_async_scheduling_with_pipeline_parallelism(monkeypatch):
+    """PP>1 must stay buildable whenever the V1 model runner is selected, such
+    as the external_launcher fallback; async scheduling is dropped instead."""
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
+
+    cfg = VllmConfig(
+        scheduler_config=SchedulerConfig(
+            max_model_len=8192,
+            is_encoder_decoder=False,
+        ),
+        parallel_config=ParallelConfig(
+            pipeline_parallel_size=2,
+            distributed_executor_backend="mp",
+            nnodes=2,
+        ),
+    )
+    assert cfg.scheduler_config.async_scheduling is False
+
+
+def test_v1_model_runner_rejects_pipeline_parallelism_with_async_scheduling():
+    """Only the async combination desyncs the grammar FSM (#45014), so plain
+    PP>1 must stay usable on the V1 model runner."""
+    config = SimpleNamespace(
+        parallel_config=ParallelConfig(
+            pipeline_parallel_size=2,
+            distributed_executor_backend="mp",
+        ),
+        scheduler_config=SchedulerConfig.default_factory(async_scheduling=False),
+        speculative_config=None,
+        model_config=None,
+    )
+    config._dflash_needs_multi_kv_group = lambda: False
+    config._is_dflash2_draft = lambda: False
+
+    assert VllmConfig._get_v1_model_runner_unsupported_features(config) == []
+
+    config.scheduler_config.async_scheduling = True
+    unsupported = VllmConfig._get_v1_model_runner_unsupported_features(config)
+    assert "pipeline parallelism with async scheduling" in unsupported
+
+
 def test_data_parallel_rpc_port_has_fixed_default():
     assert ParallelConfig().data_parallel_rpc_port == 29550
 
@@ -929,31 +1414,80 @@ def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: in
     assert config.get_parallel_size(parallel) == expected
 
 
-def test_engram_rejects_elastic_cross_dp():
+@pytest.mark.parametrize("option", ["embedding_across_dp", "dp_shared_memory"])
+def test_engram_rejects_elastic_cross_dp(option):
     parallel = ParallelConfig(
         tensor_parallel_size=4,
         data_parallel_size=2,
         distributed_executor_backend="mp",
     )
     parallel.enable_elastic_ep = True
-    with pytest.raises(ValueError, match="embedding_across_dp.*elastic EP"):
-        EngramConfig(embedding_across_dp=True).verify_parallel_config(parallel)
+    with pytest.raises(ValueError, match=f"{option}.*elastic EP"):
+        EngramConfig(**{option: True}).verify_parallel_config(parallel)
 
 
-@pytest.mark.parametrize("legacy", [None, "0", "1"])
-def test_engram_cpu_offload_environment_fallback(monkeypatch, legacy):
-    """Explicit settings must override the legacy environment fallback."""
-    monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False)
-    if legacy is not None:
-        monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
-    assert EngramConfig().cpu_offload == (legacy == "1")
-    assert EngramConfig(cpu_offload=False).cpu_offload is False
-    assert EngramConfig(cpu_offload=True).cpu_offload is True
+def test_engram_dp_shared_memory_requires_cpu_offload():
+    with pytest.raises(ValueError, match="requires cpu_offload"):
+        EngramConfig(cpu_offload=False, dp_shared_memory=True)
+
+
+@pytest.mark.parametrize(
+    "cpu_offload,dp_size,elastic_ep,expected",
+    [(True, 2, False, True), (False, 2, False, False), (True, 1, False, False)],
+)
+def test_engram_dp_shared_memory_defaults_when_supported(
+    cpu_offload, dp_size, elastic_ep, expected
+):
+    """Unset dp_shared_memory enables sharing only for offloaded, non-elastic DP."""
+    parallel = ParallelConfig(data_parallel_size=dp_size)
+    parallel.enable_elastic_ep = elastic_ep
+    config = EngramConfig(cpu_offload=cpu_offload)
+    config.resolve_dp_shared_memory(parallel)
+    assert config.dp_shared_memory is expected
+
+
+@pytest.mark.parametrize(
+    "dp_size,load_format,multithread,error",
+    [
+        (1, "auto", False, "requires data_parallel_size > 1"),
+        (2, "auto", False, None),
+        (2, "safetensors", True, None),
+    ],
+)
+def test_engram_dp_shared_memory_config_validation(
+    monkeypatch, dp_size, load_format, multithread, error
+):
+    """Reject invalid shared configs before distributed init; allow threaded loads."""
+    monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
+    config = cast(
+        VllmConfig,
+        SimpleNamespace(
+            model_config=SimpleNamespace(
+                architecture="DeepseekV41ForCausalLM",
+                hf_text_config=SimpleNamespace(engram_layer_ids=[1]),
+            ),
+            speculative_config=None,
+            engram_config=EngramConfig(cpu_offload=True, dp_shared_memory=True),
+            parallel_config=ParallelConfig(data_parallel_size=dp_size),
+            load_config=LoadConfig(
+                load_format=load_format,
+                model_loader_extra_config={"enable_multithread_load": multithread},
+            ),
+        ),
+    )
+    if error:
+        with pytest.raises(ValueError, match=error):
+            VllmConfig._resolve_and_verify_engram_config(config)
+    else:
+        VllmConfig._resolve_and_verify_engram_config(config)
 
 
 @pytest.mark.parametrize(
-    "architecture, ple_layers, cuda, supported",
+    "architecture, ple_layers, accelerator, supported",
     [
+        ("DeepseekV41ForCausalLM", [1], True, True),
+        ("DeepseekV41ForCausalLM", [], True, False),
+        ("DeepseekV41ForCausalLM", [1], False, False),
         ("Qwen4ExpForCausalLM", [1], True, True),
         ("Qwen4ExpForConditionalGeneration", [1], True, True),
         ("Qwen4ExpForCausalLM", [], True, False),
@@ -964,15 +1498,19 @@ def test_engram_cpu_offload_environment_fallback(monkeypatch, legacy):
         (None, None, True, False),
     ],
 )
-def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, supported):
+def test_engram_model_support(
+    monkeypatch, architecture, ple_layers, accelerator, supported
+):
     """A similarly named HF field must not enable unsupported implementations."""
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
+    monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: accelerator)
     model = (
         cast(
             ModelConfig,
             SimpleNamespace(
                 architecture=architecture,
-                hf_text_config=SimpleNamespace(ple_layer_ids=ple_layers),
+                hf_text_config=SimpleNamespace(
+                    ple_layer_ids=ple_layers, engram_layer_ids=ple_layers
+                ),
             ),
         )
         if architecture is not None
@@ -985,57 +1523,67 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo
         with pytest.raises(ValueError, match="requires a model with supported Engram"):
             config.verify_model_config(model)
 
+    resolved = cast(
+        VllmConfig,
+        SimpleNamespace(
+            model_config=model,
+            speculative_config=None,
+            parallel_config=ParallelConfig(),
+            load_config=LoadConfig(load_format="dummy"),
+            engram_config=None,
+        ),
+    )
+    VllmConfig._resolve_and_verify_engram_config(resolved)
+    assert (resolved.engram_config is not None) == supported
+    if supported:
+        assert resolved.engram_config.cpu_offload is True
+
+
+def test_engram_cpu_offload_default():
+    assert EngramConfig().cpu_offload is True
 
-def test_engram_config_defaults_to_none(monkeypatch):
-    monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False)
+
+def test_engram_config_defaults_to_none():
     config = VllmConfig()
     assert config.engram_config is None
     assert config.compute_hash()
 
 
-@pytest.mark.parametrize("legacy", ["0", "1"])
-def test_engram_none_resolves_legacy_offload(monkeypatch, legacy):
-    """Legacy enablement materializes a config before model validation."""
-    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+@pytest.mark.parametrize("explicit", [False, True])
+@pytest.mark.parametrize("enable_dbo,ubatch_size", [(True, 0), (False, 2)])
+def test_deepseek_engram_rejects_microbatching(explicit, enable_dbo, ubatch_size):
+    """Reject overlapping Engram staging even without an explicit EngramConfig."""
     config = cast(
         VllmConfig,
         SimpleNamespace(
             model_config=SimpleNamespace(
-                architecture="Qwen4ExpForCausalLM",
-                hf_text_config=SimpleNamespace(ple_layer_ids=[1]),
+                architecture="DeepseekV41ForCausalLM",
+                hf_text_config=SimpleNamespace(engram_layer_ids=[1]),
             ),
             speculative_config=None,
-            engram_config=None,
-            parallel_config=ParallelConfig(),
+            engram_config=EngramConfig(cpu_offload=False) if explicit else None,
+            parallel_config=ParallelConfig(
+                enable_dbo=enable_dbo, ubatch_size=ubatch_size
+            ),
         ),
     )
-    VllmConfig._resolve_and_verify_engram_config(config)
-    if legacy == "1":
-        assert config.engram_config is not None
-        assert config.engram_config.cpu_offload is True
-    else:
-        assert config.engram_config is None
+    with pytest.raises(
+        ValueError, match="Engram does not support DBO or microbatching"
+    ):
+        VllmConfig._resolve_and_verify_engram_config(config)
 
 
-@pytest.mark.parametrize("legacy", ["0", "1"])
-def test_engram_explicit_config_requires_supported_model(monkeypatch, legacy):
+def test_engram_explicit_config_requires_supported_model():
     """Explicit all-false settings still opt into model validation."""
-    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
     with pytest.raises(ValueError, match="requires a model with supported Engram"):
         VllmConfig(engram_config=EngramConfig(cpu_offload=False))
 
 
-def test_engram_legacy_offload_requires_supported_model(monkeypatch):
-    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
-    with pytest.raises(ValueError, match="requires a model with supported Engram"):
-        VllmConfig()
-
-
 @pytest.mark.parametrize("target_has_ple", [False, True])
-def test_engram_draft_config_validates_target(monkeypatch, target_has_ple):
+@pytest.mark.parametrize("explicit", [False, True])
+def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, explicit):
     """MTP may inherit cross-DP sharding without having its own PLE layers."""
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
     target = SimpleNamespace(
         architecture="Qwen4ExpForCausalLM",
         hf_text_config=SimpleNamespace(ple_layer_ids=[1] if target_has_ple else []),
@@ -1048,24 +1596,29 @@ def test_engram_draft_config_validates_target(monkeypatch, target_has_ple):
             speculative_config=SimpleNamespace(
                 draft_model_config=draft, target_model_config=target
             ),
-            engram_config=EngramConfig(embedding_across_dp=True),
+            engram_config=EngramConfig(embedding_across_dp=True) if explicit else None,
             parallel_config=ParallelConfig(),
+            load_config=LoadConfig(load_format="dummy"),
         ),
     )
-    if target_has_ple:
-        VllmConfig._resolve_and_verify_engram_config(config)
-    else:
+    if explicit and not target_has_ple:
         with pytest.raises(ValueError, match="requires a model with supported Engram"):
             VllmConfig._resolve_and_verify_engram_config(config)
+    else:
+        VllmConfig._resolve_and_verify_engram_config(config)
+        assert (config.engram_config is not None) == target_has_ple
+        if target_has_ple:
+            assert config.engram_config.cpu_offload is True
 
 
-def test_engram_hash_tracks_storage_and_sharding():
+def test_engram_hash_tracks_execution_options():
     configs = [
-        EngramConfig(cpu_offload=offload, embedding_across_dp=across_dp)
-        for offload in (False, True)
-        for across_dp in (False, True)
+        EngramConfig(cpu_offload=True),
+        EngramConfig(cpu_offload=False),
+        EngramConfig(cpu_offload=True, embedding_across_dp=True),
+        EngramConfig(cpu_offload=True, dp_shared_memory=True),
     ]
-    assert len({config.compute_hash() for config in configs}) == 4
+    assert len({config.compute_hash() for config in configs}) == len(configs)
 
 
 @pytest.mark.parametrize("port", [1, 29550, 65535])
@@ -1104,6 +1657,41 @@ def test_reconfigure_for_independent_dp_rank_on_multinode_dense_model():
     assert parallel_config.world_size == 8
 
 
+@pytest.mark.parametrize("data_parallel_size", [4, 8])
+def test_nnodes_within_dp_when_replicas_outnumber_nodes(data_parallel_size):
+    """A replica that fits on one node spans one node, never zero.
+
+    External LB pins ``data_parallel_size_local`` to 1, so the ratio rounds
+    down to 0 as soon as there are more DP replicas than nodes. Anything
+    dividing by ``nnodes_within_dp`` then raises ZeroDivisionError.
+    """
+    parallel_config = ParallelConfig(
+        data_parallel_size=data_parallel_size,
+        data_parallel_size_local=1,
+        data_parallel_external_lb=True,
+        distributed_executor_backend="mp",
+        nnodes=2,
+        node_rank=1,
+    )
+
+    assert parallel_config.nnodes_within_dp == 1
+    assert parallel_config.node_rank_within_dp == 0
+    assert parallel_config.local_world_size == parallel_config.world_size
+
+
+@pytest.mark.parametrize("nnodes", [4, 9])
+def test_nnodes_within_dp_rejects_uneven_internal_lb(nnodes):
+    parallel_config = ParallelConfig(
+        data_parallel_size=8,
+        data_parallel_size_local=1,
+        distributed_executor_backend="mp",
+        nnodes=nnodes,
+    )
+
+    with pytest.raises(ValueError, match="Invalid data parallel configuration"):
+        _ = parallel_config.nnodes_within_dp
+
+
 def test_draft_model_enables_async_scheduling_by_default():
     parallel_config = ParallelConfig(distributed_executor_backend="uni")
     model_config = ModelConfig("Qwen/Qwen3-0.6B", max_model_len=2048)
@@ -1127,6 +1715,180 @@ def test_draft_model_enables_async_scheduling_by_default():
     assert cfg.scheduler_config.async_scheduling is True
 
 
+def test_dflash_allows_async_scheduling(tmp_path: Path):
+    """DFlash must stay async-schedulable: it was the only parallel-drafting
+    method excluded from the allowlist."""
+    from transformers import LlamaConfig
+
+    common = dict(
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        vocab_size=256,
+        max_position_embeddings=2048,
+    )
+    target_path = tmp_path / "target"
+    draft_path = tmp_path / "draft"
+    _write_json(
+        target_path / "config.json",
+        LlamaConfig(architectures=["LlamaForCausalLM"], **common).to_dict(),
+    )
+    _write_json(
+        draft_path / "config.json",
+        dict(
+            common,
+            architectures=["DFlash2DraftModel"],
+            model_type="qwen3",
+            num_hidden_layers=1,
+            layer_types=["sliding_attention"],
+            sliding_window=128,
+            dflash_config=dict(
+                block_size=4,
+                mask_token_id=255,
+                target_layer_ids=[0],
+                conv_group_size=2,
+                conv_kernel_size=2,
+                selector_rank=8,
+                selector_top_k=4,
+            ),
+        ),
+    )
+    parallel_config = ParallelConfig(distributed_executor_backend="uni")
+    model_config = ModelConfig(
+        model=str(target_path), tokenizer_mode="skip", max_model_len=2048
+    )
+    speculative_config = SpeculativeConfig(
+        method="dflash",
+        model=str(draft_path),
+        num_speculative_tokens=3,
+        target_model_config=model_config,
+        target_parallel_config=parallel_config,
+    )
+    cfg = VllmConfig(
+        model_config=model_config,
+        scheduler_config=SchedulerConfig(
+            max_model_len=2048,
+            is_encoder_decoder=False,
+        ),
+        parallel_config=parallel_config,
+        speculative_config=speculative_config,
+    )
+
+    assert speculative_config.method == "dflash"
+    assert cfg.scheduler_config.async_scheduling is True
+
+
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize("tp_size", [1, 2])
+@pytest.mark.parametrize("target_ep", [False, True], ids=["ep-off", "ep-on"])
+@pytest.mark.parametrize(
+    ("method", "draft_is_moe"),
+    [
+        pytest.param("draft_model", False, id="dense-draft"),
+        pytest.param("eagle", False, id="eagle"),
+        pytest.param("eagle3", False, id="eagle3"),
+        pytest.param("draft_model", True, id="moe-draft"),
+        pytest.param("mtp", True, id="mtp"),
+        pytest.param("dspark", True, id="dspark"),
+    ],
+)
+def test_draft_inherits_ep_only_for_moe(
+    tmp_path: Path,
+    tp_size: int,
+    target_ep: bool,
+    method: str,
+    draft_is_moe: bool,
+):
+    """Validate final draft configs without loading weights or mocking validation."""
+    from transformers import LlamaConfig, MixtralConfig
+
+    from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config
+
+    common = dict(
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        vocab_size=128,
+        max_position_embeddings=2048,
+    )
+    if method in ("mtp", "dspark"):
+        target_hf_config = DeepseekV4Config(
+            architectures=["DeepseekV4ForCausalLM"],
+            n_routed_experts=4,
+            num_experts_per_tok=2,
+            num_nextn_predict_layers=1,
+            compress_ratios=[1, 1],
+            head_dim=32,
+            **common,
+        )
+    else:
+        target_hf_config = MixtralConfig(
+            architectures=["MixtralForCausalLM"], num_local_experts=4, **common
+        )
+    target_path = tmp_path / "target"
+    draft_path = tmp_path / "draft"
+    _write_json(target_path / "config.json", target_hf_config.to_dict())
+    _write_json(
+        draft_path / "config.json",
+        LlamaConfig(architectures=["LlamaForCausalLM"], **common).to_dict(),
+    )
+    target_model_config = ModelConfig(
+        model=str(target_path), tokenizer_mode="skip", max_model_len=2048
+    )
+    target_parallel_config = ParallelConfig(
+        tensor_parallel_size=tp_size,
+        enable_expert_parallel=target_ep,
+        distributed_executor_backend="mp",
+    )
+    target_model_config.verify_with_parallel_config(target_parallel_config)
+    speculative_config = SpeculativeConfig(
+        method=method,
+        model=str(target_path if draft_is_moe else draft_path),
+        num_speculative_tokens=1,
+        target_model_config=target_model_config,
+        target_parallel_config=target_parallel_config,
+    )
+
+    assert speculative_config.method == method
+    assert speculative_config.draft_model_config.is_moe is draft_is_moe
+    assert speculative_config.draft_parallel_config.enable_expert_parallel is (
+        target_ep and draft_is_moe
+    )
+    assert speculative_config.draft_parallel_config.tensor_parallel_size == tp_size
+    assert target_parallel_config.enable_expert_parallel is target_ep
+    assert target_parallel_config.tensor_parallel_size == tp_size
+
+
+@pytest.mark.skip_global_cleanup
+@pytest.mark.parametrize("target_ep", [False, True], ids=["ep-off", "ep-on"])
+@pytest.mark.parametrize("pass_none", [False, True], ids=["omitted", "explicit-none"])
+def test_draft_parallel_config_preserves_ep_without_model(
+    target_ep: bool, pass_none: bool
+):
+    """Legacy callers without draft model information keep EP inheritance."""
+    target_parallel_config = ParallelConfig(
+        tensor_parallel_size=2,
+        enable_expert_parallel=target_ep,
+        distributed_executor_backend="mp",
+    )
+    if pass_none:
+        draft_parallel_config = SpeculativeConfig.create_draft_parallel_config(
+            target_parallel_config, 2, draft_model_config=None
+        )
+    else:
+        draft_parallel_config = SpeculativeConfig.create_draft_parallel_config(
+            target_parallel_config, 2
+        )
+
+    assert draft_parallel_config.enable_expert_parallel is target_ep
+    assert draft_parallel_config.tensor_parallel_size == 2
+    assert target_parallel_config.enable_expert_parallel is target_ep
+
+
 @pytest.mark.parametrize(
     ("method", "parallel_drafting", "expected_slots"),
     [
@@ -1701,14 +2463,14 @@ def test_get_and_verify_max_len_with_nope_layers(
     max_model_len, rope_parameters, expected_max_len
 ):
     """NoPE layers do not prevent deriving or scaling the context length."""
-    from transformers import PretrainedConfig
+    from transformers import PreTrainedConfig
 
     from vllm.config.model import _get_and_verify_max_len
     from vllm.transformers_utils.model_arch_config_convertor import (
         ModelArchConfigConvertorBase,
     )
 
-    hf_config = PretrainedConfig(
+    hf_config = PreTrainedConfig(
         max_position_embeddings=4096,
         original_max_position_embeddings=2048,
     )
@@ -1731,8 +2493,56 @@ def test_get_and_verify_max_len_with_nope_layers(
     assert hf_config.rope_parameters["full_attention"] is None
 
 
+@pytest.mark.parametrize(
+    ("rope_type", "factor", "expected_max_len"),
+    [
+        # TeleChat3-36B-Thinking: 32768 already scaled from 8192 by 4
+        ("yarn", 4.0, 32768),
+        # sarvam-105b: declares factor 40 but only serves 131072 of it
+        ("deepseek_yarn", 40.0, 32768),
+        ("deepseek_llama_scaling", 40.0, 32768),
+        # Non-YaRN scaling still multiplies
+        ("linear", 4.0, 131072),
+    ],
+)
+def test_get_and_verify_max_len_yarn_is_already_scaled(
+    rope_type, factor, expected_max_len
+):
+    """YaRN variants must not re-apply `factor` to max_position_embeddings.
+
+    Transformers treats max_position_embeddings as the final context length
+    for every YaRN variant, so scaling it again overstates the limit and lets
+    requests past the end of the cos/sin cache.
+    """
+    from transformers import PreTrainedConfig
+
+    from vllm.config.model import _get_and_verify_max_len
+    from vllm.transformers_utils.model_arch_config_convertor import (
+        ModelArchConfigConvertorBase,
+    )
+
+    hf_config = PreTrainedConfig(max_position_embeddings=32768)
+    hf_config.rope_parameters = {
+        "rope_type": rope_type,
+        "factor": factor,
+        "original_max_position_embeddings": 8192,
+    }
+    model_arch_config = ModelArchConfigConvertorBase(hf_config, hf_config).convert()
+
+    actual_max_len = _get_and_verify_max_len(
+        hf_config=hf_config,
+        model_arch_config=model_arch_config,
+        tokenizer_config=None,
+        max_model_len=None,
+        disable_sliding_window=False,
+        sliding_window=None,
+    )
+
+    assert actual_max_len == expected_max_len
+
+
 class MockConfig:
-    """Simple mock object for testing maybe_pull_model_tokenizer_for_runai"""
+    """Simple mock object for testing maybe_pull_model_tokenizer_for_runai."""
 
     def __init__(self, model: str, tokenizer: str):
         self.model = model
@@ -2210,7 +3020,6 @@ def test_validate_mamba_align_subblock_prefill():
 )
 def test_vllm_config_defaults(model_id, compilation_config, optimization_level):
     """Test that optimization-level defaults are correctly applied."""
-
     model_config = None
     if model_id is not None:
         model_config = ModelConfig(model_id)
@@ -2395,7 +3204,6 @@ def test_vllm_config_explicit_overrides():
 def test_fusion_pass_op_priority():
     """This test checks that custom op enablement & IR op priority
     correctly control default fusions"""
-
     # Default config, O2, rms_norm+quant fusion disabled
     cfg1 = VllmConfig()
     assert not cfg1.compilation_config.pass_config.fuse_norm_quant
@@ -2578,18 +3386,188 @@ def test_draft_sample_method_gumbel_is_rejected():
 def _watermarked_vllm_config() -> VllmConfig:
     config = object.__new__(VllmConfig)
     config.watermark_config = WatermarkConfig(key=42)
+    config.attention_config = AttentionConfig()
     config.speculative_config = None
     return config
 
 
-def test_gumbel_watermark_rejects_speculative_decoding():
+def test_target_only_gumbel_allows_speculative_decoding(caplog_vllm, disable_log_dedup):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        key=42, allow_target_only_watermarking=True
+    )
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with caplog_vllm.at_level(logging.WARNING):
+        config._check_watermarking_unsupported()
+
+    assert "Target-only watermarking leaves accepted draft tokens" in caplog_vllm.text
+    assert "Context deduplication is not supported" in caplog_vllm.text
+
+
+def test_speculative_watermarking_without_context_dedup_does_not_warn(
+    caplog_vllm, disable_log_dedup
+):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        algorithm="dual_key_gumbel", key=42, deduplicate_contexts="none"
+    )
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with caplog_vllm.at_level(logging.WARNING):
+        config._check_watermarking_unsupported()
+
+    assert "Context deduplication is not supported" not in caplog_vllm.text
+
+
+def test_gumbel_rejects_speculative_decoding_without_target_only():
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        algorithm="gumbel", key=42, allow_target_only_watermarking=False
+    )
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with pytest.raises(ValueError, match="'gumbel'.*allow_target_only_watermarking"):
+        config._check_watermarking_unsupported()
+
+
+def test_dual_key_gumbel_warns_that_configured_alpha_is_unused(
+    caplog_vllm, disable_log_dedup
+):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        algorithm="dual_key_gumbel", key=42, alpha=0.25
+    )
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with caplog_vllm.at_level(logging.WARNING):
+        config._check_watermarking_unsupported()
+
+    assert "The configured alpha=0.25 is not used" in caplog_vllm.text
+
+
+@pytest.mark.parametrize(
+    ("alpha", "speculative"),
+    [(0.1, True), (0.25, False)],
+    ids=["default-alpha", "no-specdec"],
+)
+def test_dual_key_gumbel_alpha_warning_is_not_emitted(
+    caplog_vllm, disable_log_dedup, alpha, speculative
+):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(
+        algorithm="dual_key_gumbel", key=42, alpha=alpha
+    )
+    if speculative:
+        config.speculative_config = SimpleNamespace(
+            method="mtp",
+            draft_sample_method="probabilistic",
+            rejection_sample_method="standard",
+            parallel_drafting=False,
+        )
+
+    with caplog_vllm.at_level(logging.WARNING):
+        config._check_watermarking_unsupported()
+
+    assert "is not used" not in caplog_vllm.text
+
+
+def test_dual_key_gumbel_requires_probabilistic_drafting():
     config = _watermarked_vllm_config()
-    config.speculative_config = SpeculativeConfig(
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    config.speculative_config = SimpleNamespace(
+        method="mtp",
+        draft_sample_method="greedy",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    with pytest.raises(ValueError, match="draft_sample_method='probabilistic'"):
+        config._check_watermarking_unsupported()
+
+
+@pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"])
+def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    config.speculative_config = SimpleNamespace(
+        method=method,
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
+    )
+
+    config._check_watermarking_unsupported()
+
+
+def test_dual_key_gumbel_supports_dspark():
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    config.speculative_config = SimpleNamespace(
+        method="dspark",
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=True,
+    )
+
+    config._check_watermarking_unsupported()
+
+
+def test_dual_key_gumbel_rejects_non_autoregressive_speculation():
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    config.speculative_config = SimpleNamespace(
         method="ngram",
-        num_speculative_tokens=1,
+        draft_sample_method="probabilistic",
+        rejection_sample_method="standard",
+        parallel_drafting=False,
     )
 
-    with pytest.raises(ValueError, match="does not support speculative decoding"):
+    with pytest.raises(ValueError, match="autoregressive model-based"):
+        config._check_watermarking_unsupported()
+
+
+@pytest.mark.parametrize(
+    ("overrides", "match"),
+    [
+        ({"rejection_sample_method": "synthetic"}, "rejection_sample_method"),
+        ({"rejection_sample_method": "block"}, "rejection_sample_method"),
+        ({"parallel_drafting": True}, "Parallel speculative drafting"),
+    ],
+)
+def test_dual_key_gumbel_rejects_incompatible_speculative_modes(overrides, match):
+    config = _watermarked_vllm_config()
+    config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+    values = {
+        "method": "mtp",
+        "draft_sample_method": "probabilistic",
+        "rejection_sample_method": "standard",
+        "parallel_drafting": False,
+        **overrides,
+    }
+    config.speculative_config = SimpleNamespace(**values)
+
+    with pytest.raises(ValueError, match=match):
         config._check_watermarking_unsupported()
 
 
@@ -2608,6 +3586,12 @@ def test_watermark_key_must_fit_in_64_bits():
         WatermarkConfig(key=2**64)
 
 
+@pytest.mark.parametrize("alpha", [-0.1, 1.1])
+def test_watermark_alpha_must_be_a_probability(alpha):
+    with pytest.raises(ValidationError):
+        WatermarkConfig(key=42, algorithm="dual_key_gumbel", alpha=alpha)
+
+
 def test_unknown_watermark_prf_is_rejected():
     with pytest.raises(ValidationError):
         pydantic.TypeAdapter(WatermarkConfig).validate_python(
@@ -2807,8 +3791,24 @@ def test_ir_op_priority_default():
     assert priority_config.fused_add_rms_norm == ["native"]
 
 
+@pytest.mark.parametrize("mode", [CompilationMode.NONE, CompilationMode.VLLM_COMPILE])
+@pytest.mark.parametrize("backend", ["inductor", "eager"])
+def test_ir_op_platform_defaults_support_sparse_gelu(mode, backend):
+    """Worker initialization must not select an unregistered sparse GELU provider."""
+    from vllm import ir
+
+    config = SimpleNamespace(
+        compilation_config=CompilationConfig(mode=mode, backend=backend)
+    )
+    priority = current_platform.get_default_ir_op_priority(config)
+    expected = ["triton", "native"] if current_platform.is_cuda() else ["native"]
+
+    with priority.set_priority():
+        assert ir.ops.gelu_and_mul_sparse.get_priority() == expected
+
+
 def test_ir_op_priority_str():
-    """Test that passing a comma-delimited string works"""
+    """Test that passing a comma-delimited string works."""
     from vllm.config.kernel import IrOpPriorityConfig
 
     priority_config = IrOpPriorityConfig(rms_norm="vllm_c")
@@ -2867,19 +3867,12 @@ def test_load_config_rejects_non_string_load_format(bad_load_format):
         LoadConfig(load_format=bad_load_format)
 
 
-# A real Qwen3-0.6B model revision that is used in the tests below.
+# A real Qwen3-0.6B model revision that is used in the test below.
 REVISION = "c1899de289a04d12100db370d81485cdf75e47ca"
 
 
 @patch("vllm.config.model.resolve_revision", return_value=ResolvedRevision(REVISION))
-def test_revision_not_resolved_when_weights_differ_from_model(mock_resolve):
-    model_weights = "unsloth/Qwen3-0.6B-GGUF:Q8_0"
-    config = ModelConfig("Qwen/Qwen3-0.6B", model_weights=model_weights)
-    assert config.revision is None
-
-
-@patch("vllm.config.model.resolve_revision", return_value=ResolvedRevision(REVISION))
-def test_revision_resolved_when_weights_match_model(mock_resolve):
+def test_revision_resolved_for_model(mock_resolve):
     model = "Qwen/Qwen3-0.6B"
     config = ModelConfig(model)
     assert isinstance(config.revision, ResolvedRevision)
diff --git a/tests/test_envs.py b/tests/test_envs.py
index aaa53d2a5448..dfd08a9ec287 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -18,6 +18,11 @@
 from vllm.exceptions import VLLMValidationError
 
 
+def test_object_storage_shm_default_name():
+    """The generated name must fit macOS's shared-memory name limit."""
+    assert len(envs._generate_shm_name()) <= 30
+
+
 def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch):
     assert envs.VLLM_HOST_IP == ""
     assert envs.VLLM_PORT is None
@@ -43,41 +48,6 @@ def test_api_key_is_not_compile_factor(monkeypatch: pytest.MonkeyPatch):
     assert "VLLM_API_KEY" not in envs.compile_factors()
 
 
-def test_scale_out_endpoints_flag_is_runtime_only(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1")
-
-    assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is True
-    assert "VLLM_ENABLE_SCALE_OUT_ENDPOINTS" not in envs.compile_factors()
-
-
-def test_scale_out_endpoints_flag_distinguishes_unset_from_disabled(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False)
-    assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is None
-
-    monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0")
-    assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is False
-
-
-def test_scale_out_endpoints_flag_treats_empty_as_unset(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "")
-
-    assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is None
-
-
-@pytest.mark.parametrize("value", ["-1", "2", "01", "+1", "invalid", " "])
-def test_scale_out_endpoints_flag_rejects_values_other_than_zero_or_one(
-    monkeypatch: pytest.MonkeyPatch, value: str
-):
-    monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", value)
-
-    with pytest.raises(ValueError, match="must be 0 or 1"):
-        _ = envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS
-
-
 def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False)
     monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False)
diff --git a/tests/test_logger.py b/tests/test_logger.py
index 2ff100151b29..9dcc4a1ea825 100644
--- a/tests/test_logger.py
+++ b/tests/test_logger.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import enum
+import io
 import json
 import logging
 import os
@@ -19,6 +20,7 @@
     _DATE_FORMAT,
     _FORMAT,
     _configure_vllm_root_logger,
+    _use_color,
     enable_trace_function_call,
     init_logger,
 )
@@ -72,6 +74,27 @@ def test_default_vllm_root_logger_configuration(monkeypatch):
     assert formatter.datefmt == _DATE_FORMAT
 
 
+def test_use_color_force_color(monkeypatch):
+    """FORCE_COLOR forces colored logs without a TTY, while NO_COLOR and an
+    explicit VLLM_LOGGING_COLOR=0 take precedence over it."""
+    monkeypatch.setattr(sys, "stdout", io.StringIO())
+    monkeypatch.setattr(sys, "stderr", io.StringIO())
+    for var in ("NO_COLOR", "FORCE_COLOR", "VLLM_LOGGING_COLOR"):
+        monkeypatch.delenv(var, raising=False)
+
+    assert not _use_color()
+
+    monkeypatch.setenv("FORCE_COLOR", "1")
+    assert _use_color()
+
+    monkeypatch.setenv("VLLM_LOGGING_COLOR", "0")
+    assert not _use_color()
+    monkeypatch.delenv("VLLM_LOGGING_COLOR")
+
+    monkeypatch.setenv("NO_COLOR", "1")
+    assert not _use_color()
+
+
 def test_descendent_loggers_depend_on_and_propagate_logs_to_root_logger(monkeypatch):
     """This test presumes that VLLM_CONFIGURE_LOGGING (default: True) and
     VLLM_LOGGING_CONFIG_PATH (default: None) are not configured and default
diff --git a/tests/test_pcp_dp.py b/tests/test_pcp_dp.py
new file mode 100644
index 000000000000..f6eed57a6f70
--- /dev/null
+++ b/tests/test_pcp_dp.py
@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.config import ParallelConfig
+from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager
+from vllm.forward_context import DPMetadata
+from vllm.v1.attention.ops.pcp import maybe_gather_mla_latent_cache_inputs
+
+
+@pytest.mark.parametrize(
+    "pcp_size,sp_size,enable_ep,expected",
+    [
+        (1, 1, True, [5, 7]),
+        (1, 2, True, [3, 3, 4, 4]),
+        (2, 1, False, [10, 14]),
+        (2, 1, True, [5, 5, 7, 7]),
+        (2, 2, True, [3, 3, 3, 3, 4, 4, 4, 4]),
+    ],
+)
+def test_dispatch_sizes_expand_pcp_before_tp(pcp_size, sp_size, enable_ep, expected):
+    config = ParallelConfig(
+        distributed_executor_backend="mp",
+        data_parallel_size=2,
+        prefill_context_parallel_size=pcp_size,
+        tensor_parallel_size=sp_size,
+        enable_expert_parallel=enable_ep,
+    )
+    metadata = DPMetadata.make(config, 5, torch.tensor([5, 7]))
+    with metadata.sp_local_sizes(sp_size, pcp_size=pcp_size, use_ep=enable_ep) as sizes:
+        assert sizes == expected
+    assert metadata.local_sizes is None
+    assert metadata.num_tokens_across_dp_cpu.tolist() == [5, 7]
+
+
+@pytest.mark.parametrize(
+    "dp_size,pcp_size,tp_size,use_ep,is_sp,expected",
+    [
+        (2, 2, 1, True, False, "ep"),
+        (2, 2, 2, True, True, "ep"),
+        (2, 2, 2, False, False, "dp"),
+        (2, 1, 2, True, False, "dp"),
+        (1, 2, 2, True, False, "pcp"),
+        (2, 2, 2, True, False, None),
+    ],
+)
+def test_dispatch_reuses_existing_groups(
+    dp_size, pcp_size, tp_size, use_ep, is_sp, expected, monkeypatch
+):
+    groups = {
+        "dp": SimpleNamespace(world_size=dp_size),
+        "pcp": SimpleNamespace(world_size=pcp_size),
+        "ep": object(),
+    }
+    for name, group in groups.items():
+        monkeypatch.setattr(
+            f"vllm.distributed.device_communicators.all2all.get_{name}_group",
+            lambda group=group: group,
+        )
+    manager = AgRsAll2AllManager.__new__(AgRsAll2AllManager)
+    manager.dp_world_size = dp_size
+    manager.tp_group = SimpleNamespace(world_size=tp_size)
+    manager.use_ep = use_ep
+    if expected is None:
+        with pytest.raises(AssertionError, match="requires sequence-parallel MoE"):
+            manager._get_comm_group(is_sp)
+    else:
+        assert manager._get_comm_group(is_sp) is groups[expected]
+
+
+@pytest.mark.parametrize("enable_ep", [False, True])
+def test_ag_rs_dispatch_and_combine_use_dp_pcp_sizes(monkeypatch, enable_ep):
+    calls = []
+    local_tokens = [20, 21] if enable_ep else [20, 21, 22, 23]
+
+    class FakeGroup:
+        world_size = 4 if enable_ep else 2
+        rank_in_group = 2 if enable_ep else 1
+
+        def all_gatherv(self, tensors, dim, sizes):
+            calls.append(("gather", sizes))
+            return [torch.tensor([10, 11, 20, 21, 22, 23]) for _ in tensors]
+
+        def reduce_scatterv(self, tensor, dim, sizes):
+            calls.append(("scatter", sizes))
+            return tensor[2 : 2 + len(local_tokens)]
+
+    config = ParallelConfig(
+        distributed_executor_backend="mp",
+        data_parallel_size=2,
+        data_parallel_rank=1,
+        prefill_context_parallel_size=2,
+        enable_expert_parallel=enable_ep,
+    )
+    metadata = DPMetadata.make(config, 2, torch.tensor([1, 2]))
+    manager = AgRsAll2AllManager.__new__(AgRsAll2AllManager)
+    manager.dp_world_size = 2
+    manager.use_ep = enable_ep
+
+    manager.tp_group = SimpleNamespace(world_size=1)
+    for name in ("dp", "ep"):
+        monkeypatch.setattr(
+            f"vllm.distributed.device_communicators.all2all.get_{name}_group",
+            FakeGroup,
+        )
+    monkeypatch.setattr(
+        "vllm.distributed.device_communicators.all2all.get_pcp_group",
+        lambda: SimpleNamespace(world_size=2),
+    )
+    monkeypatch.setattr(
+        "vllm.distributed.device_communicators.all2all.get_forward_context",
+        lambda: SimpleNamespace(dp_metadata=metadata),
+    )
+    with metadata.sp_local_sizes(1, pcp_size=2, use_ep=enable_ep):
+        hidden_states, _, _ = manager.dispatch(
+            torch.tensor(local_tokens),
+            torch.ones(len(local_tokens)),
+            torch.zeros(len(local_tokens)),
+        )
+        combined = manager.combine(hidden_states)
+    assert combined.tolist() == local_tokens
+    sizes = [1, 1, 2, 2] if enable_ep else [2, 4]
+    assert calls == [("gather", sizes), ("scatter", sizes)]
+
+
+@pytest.mark.parametrize("slots", [[3, 4], [3, 4, -1, -1, -1, -1]])
+def test_decode_cache_write_ignores_dp_padding(slots):
+    kv = torch.arange(12).reshape(6, 2)
+    pe = torch.arange(6).reshape(6, 1, 1)
+    slots = torch.tensor(slots)
+    cache_kv, cache_pe, cache_slots = maybe_gather_mla_latent_cache_inputs(
+        kv, pe, slots, num_decode_tokens=2, use_pcp=True
+    )
+    torch.testing.assert_close(cache_kv, kv[:2])
+    torch.testing.assert_close(cache_pe, pe[:2])
+    torch.testing.assert_close(cache_slots, slots[:2])
+
+
+def test_expanded_slot_mapping_keeps_pcp_prefill_padding(monkeypatch):
+    calls = []
+
+    def all_gather(tensor, dim):
+        calls.append(tensor.shape[0])
+        return torch.cat((tensor, tensor), dim=dim)
+
+    monkeypatch.setattr(
+        "vllm.v1.attention.ops.pcp.get_pcp_group",
+        lambda: SimpleNamespace(world_size=2, all_gather=all_gather),
+    )
+    kv = torch.zeros(3, 2)  # Two decodes followed by one PCP padding row.
+    pe = torch.zeros(3, 1, 1)
+    slots = torch.tensor([3, 4, 8, 3, 4, -1])
+    cache_kv, _, cache_slots = maybe_gather_mla_latent_cache_inputs(
+        kv, pe, slots, num_decode_tokens=2, use_pcp=True
+    )
+    assert calls == [1, 1]
+    assert cache_kv.shape == (4, 2)
+    assert cache_slots.tolist() == [3, 4, 8, -1]
diff --git a/tests/test_request_input_bounds.py b/tests/test_request_input_bounds.py
index 44f536e7fbe7..7f79d4ba47dc 100644
--- a/tests/test_request_input_bounds.py
+++ b/tests/test_request_input_bounds.py
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Bounds on request-controlled inputs that would otherwise amplify work.
+"""Bounds and alignment checks for request-controlled inputs.
 
-Consolidates the regression tests for the request-input amplification fixes:
-stop-string caps, bad-words dedup/tokenization limit, stop-token-id dedup,
-beam-width/sequence caps, and the DeepSeek history-scan bound.
+Covers prompt-mask alignment, stop-string caps, bad-words dedup/tokenization
+limits, stop-token-id dedup, beam-width/sequence caps, and the DeepSeek
+history-scan bound.
 """
 
 import os
@@ -12,8 +12,10 @@
 import sys
 from collections.abc import Callable
 from typing import Protocol
+from unittest.mock import Mock
 
 import pytest
+import torch
 from pydantic import ValidationError
 
 import vllm.envs as envs
@@ -25,12 +27,136 @@
 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.exceptions import VLLMValidationError
+from vllm.inputs.engine import embeds_input
 from vllm.sampling_params import BeamSearchParams
 from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding
+from vllm.v1.engine.input_processor import InputProcessor
 
 pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
 
 
+@pytest.fixture
+def prompt_embeds_processor():
+    model_config = Mock(
+        max_model_len=16,
+        max_logprobs=20,
+        logits_processors=None,
+        is_diffusion=False,
+        return_sampling_mask=False,
+    )
+    model_config.get_vocab_size.return_value = 32
+    model_config.try_get_generation_config.return_value = {}
+    config = Mock(
+        model_config=model_config,
+        speculative_config=None,
+        structured_outputs_config=None,
+    )
+    renderer = Mock(tokenizer=None)
+    renderer.get_eos_token_id.return_value = None
+    registry = Mock()
+    registry.supports_multimodal_inputs.return_value = False
+    return InputProcessor(config, renderer, mm_registry=registry)
+
+
+@pytest.mark.parametrize("mask_len", [0, 1, 2, 4])
+def test_prompt_embeds_mask_length_rejected_before_engine_submission(
+    prompt_embeds_processor, mask_len
+):
+    """Reject short, long, and broadcastable masks before they reach a worker."""
+    prompt = embeds_input(
+        torch.zeros(3, 4),
+        prompt_token_ids=[1, 2, 3],
+        is_token_ids=[False] * mask_len,
+    )
+
+    with pytest.raises(VLLMValidationError, match="prompt_is_token_ids") as exc_info:
+        prompt_embeds_processor.process_inputs(
+            "invalid", prompt, SamplingParams(max_tokens=1), ("generate",)
+        )
+
+    assert exc_info.value.parameter == "prompt_is_token_ids"
+
+
+@pytest.mark.parametrize(
+    ("token_ids", "mask"),
+    [(None, None), ([1, 2, 3], None), ([1, 2, 3], [True, False, True])],
+    ids=["pure-embeds", "token-ids-without-mask", "mixed"],
+)
+def test_prompt_embeds_valid_inputs_preserved(prompt_embeds_processor, token_ids, mask):
+    embeds = torch.zeros(3, 4)
+    prompt = embeds_input(embeds, prompt_token_ids=token_ids, is_token_ids=mask)
+
+    request = prompt_embeds_processor.process_inputs(
+        "valid", prompt, SamplingParams(max_tokens=1), ("generate",)
+    )
+
+    assert request.prompt_embeds is embeds
+    assert request.prompt_token_ids == token_ids
+    assert request.prompt_is_token_ids == mask
+
+
+@pytest.mark.parametrize("offset", [-1, 0, 2, 3, 4, 2**127])
+def test_routed_experts_offset_validated_before_engine_submission(offset):
+    """Reject invalid offsets even when routed-expert output is disabled."""
+    processor = Mock(spec=InputProcessor)
+    processor.tokenizer = None
+    processor.generation_config_fields = {}
+    processor.renderer = Mock()
+    processor.renderer.get_eos_token_id.return_value = None
+    processor.vllm_config = Mock()
+    params = SamplingParams(routed_experts_prompt_start=offset)
+    prompt = {"type": "token", "prompt_token_ids": [1, 2, 3]}
+
+    if not 0 <= offset <= 3:
+        with pytest.raises(VLLMValidationError, match="routed_experts_prompt_start"):
+            InputProcessor.process_inputs(
+                processor, "invalid", prompt, params, ("generate",)
+            )
+        params.routed_experts_prompt_start = 0
+
+    request = InputProcessor.process_inputs(
+        processor, "valid", prompt, params, ("generate",)
+    )
+    assert (
+        request.sampling_params.routed_experts_prompt_start
+        == params.routed_experts_prompt_start
+    )
+
+
+def _process_inputs_with_max_model_len(
+    params: SamplingParams, max_model_len: int = 2048
+):
+    processor = Mock(spec=InputProcessor)
+    processor.tokenizer = None
+    processor.generation_config_fields = {}
+    processor.renderer = Mock()
+    processor.renderer.get_eos_token_id.return_value = None
+    processor.vllm_config = Mock()
+    processor.model_config = Mock()
+    processor.model_config.max_model_len = max_model_len
+    prompt = {"type": "token", "prompt_token_ids": [1, 2, 3]}
+    return InputProcessor.process_inputs(
+        processor, "req", prompt, params, ("generate",)
+    )
+
+
+def test_unset_max_tokens_rejects_oversized_min_tokens():
+    """Reject min_tokens that only becomes invalid after max_tokens is filled."""
+    params = SamplingParams(max_tokens=None, min_tokens=2147483648)
+    with pytest.raises(
+        VLLMValidationError,
+        match="min_tokens must be less than or equal to max_tokens",
+    ):
+        _process_inputs_with_max_model_len(params)
+
+
+def test_unset_max_tokens_accepts_min_tokens_within_filled_max():
+    params = SamplingParams(max_tokens=None, min_tokens=10)
+    request = _process_inputs_with_max_model_len(params)
+    assert request.sampling_params.max_tokens == 2045
+    assert request.sampling_params.min_tokens == 10
+
+
 # --- Stop strings: public requests cap the number of stop strings ---------
 
 
diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py
index cd23ea993816..84ce45caa958 100644
--- a/tests/test_sampling_params.py
+++ b/tests/test_sampling_params.py
@@ -1,11 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
+from types import SimpleNamespace
 
 import pytest
 
 from vllm import SamplingParams
+from vllm.config.diffusion import DiffusionConfig
 from vllm.exceptions import VLLMValidationError
+from vllm.utils.diffusion import validate_diffusion_sampling_params
+from vllm.v1.engine.input_processor import InputProcessor
 
 
 @dataclass
@@ -13,6 +17,7 @@ class MockModelConfig:
     is_diffusion: bool = False
     max_logprobs: int = 20
     logits_processors: list | None = None
+    return_sampling_mask: bool = False
 
     def get_vocab_size(self) -> int:
         return 1024
@@ -51,6 +56,191 @@ def test_non_diffusion_models_unaffected():
     params.verify(MockModelConfig(), None, None, None)
 
 
+def test_verify_leaves_logits_processors_to_admission():
+    """verify() is runner-agnostic; LP validation lives in the admission
+    layer, so an unimportable FQCN must not fail verify()."""
+    SamplingParams().verify(
+        MockModelConfig(logits_processors=["no.such:Cls"]), None, None, None
+    )
+
+
+def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None):
+    model_config = MockModelConfig(is_diffusion=True)
+    params.verify(model_config, None, None, None)
+    validate_diffusion_sampling_params(
+        params,
+        canvas_length=canvas_length,
+        vocab_size=model_config.get_vocab_size(),
+        async_scheduling=True,
+    )
+
+
+def test_diffusion_extra_args_are_validated_without_a_served_canvas():
+    # No --diffusion-config: the canvas is unknown, the ids are still checked
+    # and a read-only request is still normalised.
+    params = SamplingParams(
+        max_tokens=64,
+        extra_args={"diffusion_seed_canvas": [0, 1], "diffusion_read_only": True},
+    )
+    _verify_diffusion(params, canvas_length=None)
+    assert params.ignore_eos is True
+
+    bad = SamplingParams(extra_args={"diffusion_seed_canvas": [0, 10**9]})
+    with pytest.raises(VLLMValidationError, match="ids must be in"):
+        _verify_diffusion(bad, canvas_length=None)
+
+
+@pytest.mark.parametrize("async_scheduling", [False, True])
+@pytest.mark.parametrize(
+    "extra_args",
+    [
+        {},
+        {"diffusion_canvas_length": None},
+        {"diffusion_canvas_length": 4},
+        {"diffusion_canvas_length": 8},
+    ],
+)
+def test_narrow_diffusion_canvas_requires_async_scheduling(
+    async_scheduling, extra_args
+):
+    processor = SimpleNamespace(
+        model_config=MockModelConfig(is_diffusion=True),
+        vllm_config=SimpleNamespace(
+            scheduler_config=SimpleNamespace(async_scheduling=async_scheduling)
+        ),
+        speculative_config=None,
+        structured_outputs_config=None,
+        diffusion_config=DiffusionConfig(canvas_length=8),
+        tokenizer=None,
+        validate_logits_processors_params=lambda params: None,
+    )
+    params = SamplingParams(extra_args=extra_args)
+    if not async_scheduling and extra_args.get("diffusion_canvas_length") == 4:
+        with pytest.raises(VLLMValidationError, match="requires --async-scheduling"):
+            InputProcessor._validate_params(processor, params, ("generate",))
+    else:
+        InputProcessor._validate_params(processor, params, ("generate",))
+
+
+@pytest.mark.parametrize(
+    "extra_args, match",
+    [
+        ({"diffusion_seed_canvas": "abc"}, "list of token ids"),
+        ({"diffusion_seed_canvas": [1, 2.5]}, "list of token ids"),
+        ({"diffusion_seed_canvas": [1, True]}, "list of token ids"),
+        ({"diffusion_seed_canvas": [0, 1024]}, r"in \[0, 1024\)"),
+        ({"diffusion_seed_canvas": [0, -1]}, r"in \[0, 1024\)"),
+        ({"diffusion_max_steps": 0}, "positive integer"),
+        ({"diffusion_max_steps": "1"}, "positive integer"),
+        ({"diffusion_max_steps": True}, "positive integer"),
+        ({"diffusion_read_only": "yes"}, "boolean"),
+        ({"diffusion_read_only": 2}, "boolean"),
+        ({"diffusion_pinned": "0,1"}, "list of canvas positions"),
+        ({"diffusion_pinned": [0, True]}, "list of canvas positions"),
+        ({"diffusion_pinned": [0]}, "needs a diffusion_seed_canvas"),
+    ],
+)
+def test_diffusion_rejects_bad_extra_args(extra_args: dict, match: str):
+    with pytest.raises(VLLMValidationError, match=match):
+        _verify_diffusion(SamplingParams(extra_args=extra_args))
+
+
+def test_diffusion_seed_canvas_must_fill_the_canvas():
+    params = SamplingParams(extra_args={"diffusion_seed_canvas": [0] * 7})
+    with pytest.raises(VLLMValidationError, match="exactly 8 ids, got 7"):
+        _verify_diffusion(params, canvas_length=8)
+    # Without the served diffusion config the canvas length is unknown.
+    _verify_diffusion(params)
+
+
+@pytest.mark.parametrize("max_tokens, expected", [(100, 8), (5, 5), (None, 8)])
+@pytest.mark.parametrize("flag", [True, 1])
+def test_diffusion_read_only_ends_after_one_canvas(max_tokens, expected, flag):
+    params = SamplingParams(
+        max_tokens=max_tokens, extra_args={"diffusion_read_only": flag}
+    )
+    _verify_diffusion(params, canvas_length=8)
+    assert params.max_tokens == expected
+    assert params.ignore_eos
+
+
+@pytest.mark.parametrize(
+    "width, match",
+    [
+        (0, "positive integer"),
+        ("4", "positive integer"),
+        (True, "positive integer"),
+        (9, "no larger"),
+    ],
+)
+def test_diffusion_rejects_bad_canvas_length(width, match):
+    with pytest.raises(VLLMValidationError, match=match):
+        _verify_diffusion(
+            SamplingParams(extra_args={"diffusion_canvas_length": width}),
+            canvas_length=8,
+        )
+
+
+def test_diffusion_canvas_length_sizes_the_seed_and_the_read():
+    params = SamplingParams(
+        max_tokens=100,
+        extra_args={
+            "diffusion_canvas_length": 4,
+            "diffusion_seed_canvas": [1, 2, 3, 4],
+            "diffusion_read_only": True,
+        },
+    )
+    _verify_diffusion(params, canvas_length=8)
+    assert params.max_tokens == 4
+
+    params = SamplingParams(
+        extra_args={"diffusion_canvas_length": 4, "diffusion_seed_canvas": [0] * 8}
+    )
+    with pytest.raises(VLLMValidationError, match="exactly 4 ids, got 8"):
+        _verify_diffusion(params, canvas_length=8)
+
+
+def test_diffusion_pinned_positions_stay_inside_the_canvas():
+    seed = [0] * 8
+    with pytest.raises(VLLMValidationError, match="inside the canvas"):
+        _verify_diffusion(
+            SamplingParams(
+                extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [7, 8]}
+            ),
+            canvas_length=8,
+        )
+    with pytest.raises(VLLMValidationError, match="inside the canvas"):
+        _verify_diffusion(
+            SamplingParams(
+                extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [-1]}
+            )
+        )
+    # A narrower request canvas bounds the positions.
+    with pytest.raises(VLLMValidationError, match="inside the canvas"):
+        _verify_diffusion(
+            SamplingParams(
+                extra_args={
+                    "diffusion_canvas_length": 4,
+                    "diffusion_seed_canvas": seed[:4],
+                    "diffusion_pinned": [4],
+                }
+            ),
+            canvas_length=8,
+        )
+
+
+def test_diffusion_accepts_extra_args():
+    params = SamplingParams(
+        extra_args={
+            "diffusion_seed_canvas": list(range(8)),
+            "diffusion_pinned": [0, 1, 7],
+            "diffusion_max_steps": 4,
+            "diffusion_read_only": True,
+        }
+    )
+    _verify_diffusion(params, canvas_length=8)
+
+
 @pytest.mark.parametrize("value", [-(2**63) - 1, 2**64])
 def test_extra_args_rejects_nested_integer_overflow(value):
     """Reject extension values before they reach the engine transport."""
diff --git a/tests/test_triton_utils.py b/tests/test_triton_utils.py
index 0ae4c0d6d69d..7f6422466f5e 100644
--- a/tests/test_triton_utils.py
+++ b/tests/test_triton_utils.py
@@ -97,7 +97,8 @@ def test_triton_placeholder_language():
     lang = TritonLanguagePlaceholder()
     assert isinstance(lang, types.ModuleType)
     assert lang.__name__ == "triton.language"
-    assert lang.constexpr is None
+    assert lang.constexpr(2**31 - 1) == 2**31 - 1
+    assert lang.constexpr(1.5) == 1.5
     assert lang.dtype is None
     assert lang.int64 is None
     assert lang.int32 is None
@@ -143,3 +144,4 @@ def test_no_triton_fallback():
         assert triton.__class__.__name__ == "TritonPlaceholder"
         assert triton.language.__class__.__name__ == "TritonLanguagePlaceholder"
         assert tl.__class__.__name__ == "TritonLanguagePlaceholder"
+        assert tl.constexpr(2**31 - 1) == 2**31 - 1
diff --git a/tests/test_zen_cpu_platform_detection.py b/tests/test_zen_cpu_platform_detection.py
index 9f386f70cd98..0c99c6568a2b 100644
--- a/tests/test_zen_cpu_platform_detection.py
+++ b/tests/test_zen_cpu_platform_detection.py
@@ -65,7 +65,7 @@ def test_cpu_target_selects_cpu_platform_from_non_cpu_wheel(
     rocm_plugin.assert_not_called()
 
 
-def test_platform_detection_logs_zentorch_import_failure(caplog):
+def test_broken_zentorch_falls_back_to_cpu_platform(caplog):
     original_import = builtins.__import__
 
     def import_with_broken_zentorch(name, *args, **kwargs):
@@ -88,6 +88,54 @@ def import_with_broken_zentorch(name, *args, **kwargs):
     ):
         platform = resolve_current_platform_cls_qualname()
 
+    assert platform == "vllm.platforms.cpu.CpuPlatform"
+    assert "zentorch failed to import" in caplog.text
+    assert "OSError: incompatible shared library" in caplog.text
+
+
+def test_zentorch_failure_other_than_oserror_is_not_recovered(caplog):
+    original_import = builtins.__import__
+
+    def import_with_broken_zentorch(name, *args, **kwargs):
+        if name == "zentorch":
+            raise RuntimeError("unknown zentorch failure")
+        return original_import(name, *args, **kwargs)
+
+    with (
+        patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"),
+        patch.dict(
+            "vllm.platforms.builtin_platform_plugins",
+            {"cpu": cpu_platform_plugin},
+            clear=True,
+        ),
+        patch("vllm.platforms.load_plugins_by_group", return_value={}),
+        patch("vllm.platforms.vllm_version_matches_substr", return_value=True),
+        patch("vllm.platforms._is_amd_zen_cpu", return_value=True),
+        patch.object(builtins, "__import__", side_effect=import_with_broken_zentorch),
+        caplog.at_level(logging.DEBUG, logger="vllm.platforms"),
+    ):
+        platform = resolve_current_platform_cls_qualname()
+
+    assert platform == "vllm.platforms.interface.UnspecifiedPlatform"
+    assert "RuntimeError: unknown zentorch failure" in caplog.text
+
+
+def test_platform_plugin_failure_is_logged(caplog):
+    def failing_plugin():
+        raise RuntimeError("plugin exploded")
+
+    with (
+        patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"),
+        patch.dict(
+            "vllm.platforms.builtin_platform_plugins",
+            {"cpu": failing_plugin},
+            clear=True,
+        ),
+        patch("vllm.platforms.load_plugins_by_group", return_value={}),
+        caplog.at_level(logging.DEBUG, logger="vllm.platforms"),
+    ):
+        platform = resolve_current_platform_cls_qualname()
+
     assert platform == "vllm.platforms.interface.UnspecifiedPlatform"
     assert "Platform plugin cpu failed during detection" in caplog.text
-    assert "OSError: incompatible shared library" in caplog.text
+    assert "RuntimeError: plugin exploded" in caplog.text
diff --git a/tests/tokenizers_/fixtures/deepseek_v4/test_input_5.json b/tests/tokenizers_/fixtures/deepseek_v4/test_input_5.json
new file mode 100644
index 000000000000..2a2fa6a89e81
--- /dev/null
+++ b/tests/tokenizers_/fixtures/deepseek_v4/test_input_5.json
@@ -0,0 +1,39 @@
+{
+    "thinking_mode": "thinking",
+    "reasoning_effort": "max",
+    "messages": [
+        {
+            "role": "system",
+            "content": "You are a helpful vision assistant."
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": "请按“第一张、第二张”的顺序回答：第一张图"
+                },
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": "examples/images/carrots.jpeg"
+                    }
+                },
+                {
+                    "type": "text",
+                    "text": "和第二张图"
+                },
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": "examples/images/corn.jpeg"
+                    }
+                },
+                {
+                    "type": "text",
+                    "text": "中分别是什么食材？它们通常食用的部位分别是什么？"
+                }
+            ]
+        }
+    ]
+}
diff --git a/tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt b/tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt
index dbd823476c1c..7e3c9bd5a394 100644
--- a/tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt
+++ b/tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt
@@ -1,4 +1,4 @@
-<｜begin▁of▁sentence｜>
+<｜begin▁of▁sentence｜>You are a helpful assistant.
 
 ## Tools
 
@@ -26,7 +26,7 @@ Otherwise, output directly after </think> with tool calls or final response.
 {"name": "search", "description": "Search the web for information", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}, "num_results": {"type": "integer", "description": "Number of results to return"}}, "required": ["query"]}}
 
 You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
-You are a helpful assistant.<｜User｜>What's the weather in Beijing?<｜Assistant｜><think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>
+<｜User｜>What's the weather in Beijing?<｜Assistant｜><think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>
 
 <｜DSML｜tool_calls>
 <｜DSML｜invoke name="get_weather">
diff --git a/tests/tokenizers_/fixtures/deepseek_v4/test_output_5.txt b/tests/tokenizers_/fixtures/deepseek_v4/test_output_5.txt
new file mode 100644
index 000000000000..1203dc8db114
--- /dev/null
+++ b/tests/tokenizers_/fixtures/deepseek_v4/test_output_5.txt
@@ -0,0 +1,9 @@
+<｜begin▁of▁sentence｜>You are a helpful vision assistant.<｜User｜>请按“第一张、第二张”的顺序回答：第一张图
+
+<｜deepseek_image｜>
+
+和第二张图
+
+<｜deepseek_image｜>
+
+中分别是什么食材？它们通常食用的部位分别是什么？<｜Assistant｜></think>
\ No newline at end of file
diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py
index 4da6381a0bc9..4a1a2b156050 100644
--- a/tests/tokenizers_/test_basic.py
+++ b/tests/tokenizers_/test_basic.py
@@ -23,7 +23,7 @@ def _assert_tokenizer_like(tokenizer: object):
 
 
 def test_tokenizer_like_protocol():
-    tokenizer = get_tokenizer("openai-community/gpt2", use_fast=True)
+    tokenizer = get_tokenizer("openai-community/gpt2")
     assert isinstance(tokenizer, TokenizersBackend)
     _assert_tokenizer_like(tokenizer)
 
diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py
index a55aa57666dd..4403c2a877d0 100644
--- a/tests/tokenizers_/test_deepseek_v4.py
+++ b/tests/tokenizers_/test_deepseek_v4.py
@@ -54,13 +54,9 @@ def _load_reference_case(case_id: int):
 
 def _render_reference_case(case_id: int, **kwargs):
     messages, tools = _load_reference_case(case_id)
-    conversation, _, _ = parse_chat_messages(
-        messages,
-        _model_config(),
-        content_format="string",
-    )
+    # Preserve content blocks for encoding without fetching fixture media.
     return _tokenizer().apply_chat_template(
-        conversation=conversation,
+        conversation=messages,
         messages=messages,
         tools=tools,
         tokenize=False,
@@ -397,6 +393,7 @@ def test_deepseek_v4_maps_xhigh_to_high_reasoning_effort():
         (2, {"thinking": True, "reasoning_effort": "low"}),
         (3, {"thinking": True, "reasoning_effort": "low"}),
         (4, {"thinking": False}),
+        (5, {"thinking": False}),
     ],
 )
 def test_deepseek_v4_matches_reference_golden_fixtures(case_id, kwargs):
@@ -447,7 +444,10 @@ def test_deepseek_v4_image_blocks_become_placeholders():
         thinking=False,
     )
 
-    assert "<｜User｜>first:<｜deepseek_image｜>second:<｜deepseek_image｜>" in prompt
+    assert (
+        "<｜User｜>first:\n\n<｜deepseek_image｜>\n\n"
+        "second:\n\n<｜deepseek_image｜>" in prompt
+    )
 
 
 def test_deepseek_v4_image_sentinel_ids_match_tokenizer():
@@ -485,3 +485,146 @@ def convert_tokens_to_ids(self, token: str) -> int:
         ]
     )
     assert image_sentinel_mask(ids).tolist() == [False, True, True, False, False]
+
+
+def _request_tools():
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "description": "Get weather for a city",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                    "required": ["city"],
+                },
+            },
+        }
+    ]
+
+
+def _encode_reference(messages):
+    from vllm.tokenizers.deepseek_v4_encoding import encode_messages
+
+    return encode_messages(messages, thinking_mode="thinking", reasoning_effort="low")
+
+
+def test_deepseek_v4_attaches_request_tools_to_existing_system_message():
+    messages = [
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "Weather in Paris?"},
+    ]
+
+    prompt = _tokenizer().apply_chat_template(
+        messages,
+        tools=_request_tools(),
+        tokenize=False,
+        thinking=True,
+        reasoning_effort="low",
+    )
+
+    expected = _encode_reference(
+        [{**messages[0], "tools": _request_tools()}, messages[1]]
+    )
+    assert prompt == expected
+    assert prompt.startswith("<｜begin▁of▁sentence｜>You are helpful.\n\n## Tools")
+    assert prompt.count("## Tools") == 1
+
+
+def test_deepseek_v4_synthetic_system_only_when_no_system_message():
+    """Without a system message, request tools still get a synthetic leading
+    system entry."""
+    messages = [{"role": "user", "content": "Weather in Paris?"}]
+
+    prompt = _tokenizer().apply_chat_template(
+        messages,
+        tools=_request_tools(),
+        tokenize=False,
+        thinking=True,
+        reasoning_effort="low",
+    )
+
+    expected = _encode_reference(
+        [{"role": "system", "tools": _request_tools()}, messages[0]]
+    )
+    assert prompt == expected
+    assert prompt.startswith("<｜begin▁of▁sentence｜>\n\n## Tools")
+
+
+def test_deepseek_v4_request_tools_override_system_message_tools():
+    message_tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "message_level_tool",
+                "description": "Should not be rendered",
+                "parameters": {"type": "object", "properties": {}},
+            },
+        }
+    ]
+    system_message = {
+        "role": "system",
+        "content": "You are helpful.",
+        "tools": message_tools,
+    }
+    user_message = {"role": "user", "content": "Weather in Paris?"}
+
+    prompt = _tokenizer().apply_chat_template(
+        [system_message, user_message],
+        tools=_request_tools(),
+        tokenize=False,
+        thinking=True,
+        reasoning_effort="low",
+    )
+
+    expected = _encode_reference(
+        [{**system_message, "tools": _request_tools()}, user_message]
+    )
+    assert prompt == expected
+    assert '"name": "get_weather"' in prompt
+    assert "message_level_tool" not in prompt
+
+
+def test_deepseek_v4_attaches_request_tools_to_first_system_message_only():
+    """Only the first system message, at any index, gets the request tools."""
+    messages = [
+        {"role": "user", "content": "Hi"},
+        {"role": "system", "content": "First system."},
+        {"role": "system", "content": "Second system."},
+        {"role": "user", "content": "Weather in Paris?"},
+    ]
+
+    prompt = _tokenizer().apply_chat_template(
+        messages,
+        tools=_request_tools(),
+        tokenize=False,
+        thinking=True,
+        reasoning_effort="low",
+    )
+
+    expected = _encode_reference(
+        [
+            messages[0],
+            {**messages[1], "tools": _request_tools()},
+            messages[2],
+            messages[3],
+        ]
+    )
+    assert prompt == expected
+    assert "First system.\n\n## Tools" in prompt
+    assert prompt.count("## Tools") == 1
+
+
+def test_deepseek_v4_request_tools_do_not_mutate_caller_messages():
+    import copy
+
+    messages = [
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "Weather in Paris?"},
+    ]
+    snapshot = copy.deepcopy(messages)
+
+    _tokenizer().apply_chat_template(messages, tools=_request_tools(), tokenize=False)
+
+    assert messages == snapshot
diff --git a/tests/tokenizers_/test_deepseek_v41.py b/tests/tokenizers_/test_deepseek_v41.py
index 493fdd88ba5b..13c0375fbb0a 100644
--- a/tests/tokenizers_/test_deepseek_v41.py
+++ b/tests/tokenizers_/test_deepseek_v41.py
@@ -99,6 +99,24 @@ def test_raw_text_parts_preserve_reference_separator():
     assert messages == original
 
 
+@pytest.mark.parametrize(
+    ("role", "responses_type"),
+    [("user", "input_text"), ("assistant", "output_text")],
+)
+def test_responses_text_parts_match_chat_text_parts(role, responses_type):
+    responses_message = {
+        "role": role,
+        "content": [{"type": responses_type, "text": "hello"}],
+    }
+    chat_message = {
+        "role": role,
+        "content": [{"type": "text", "text": "hello"}],
+    }
+    assert render([responses_message], thinking=False) == render(
+        [chat_message], thinking=False
+    )
+
+
 def test_mid_system_gets_its_own_marker_and_generation_header():
     assert render(
         [
diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py
index 0a47426dc974..a4ee4a755e34 100644
--- a/tests/tokenizers_/test_registry.py
+++ b/tests/tokenizers_/test_registry.py
@@ -10,6 +10,7 @@
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING
 
 from vllm.tokenizers import TokenizerLike
+from vllm.tokenizers.hf import CachedHfTokenizer
 from vllm.tokenizers.registry import (
     TokenizerRegistry,
     cached_get_tokenizer,
@@ -68,6 +69,78 @@ def test_resolve_tokenizer_args_idempotent(runner_type):
     )
 
 
+@pytest.mark.parametrize("input_kwargs", [{}, {"mistral_format": False}])
+def test_resolve_tokenizer_args_forces_hf_mistral_format_false(input_kwargs):
+    resolved_mode, _, _, kwargs = resolve_tokenizer_args(
+        "mistralai/Mistral-Nemo-Instruct-2407",
+        tokenizer_mode="hf",
+        **input_kwargs,
+    )
+
+    assert resolved_mode == "hf"
+    assert kwargs["mistral_format"] is False
+
+
+def test_resolve_tokenizer_args_rejects_hf_mistral_format_true():
+    with pytest.raises(
+        ValueError,
+        match="mistral_format=True is not supported with tokenizer_mode='hf'",
+    ):
+        resolve_tokenizer_args(
+            "mistralai/Mistral-Nemo-Instruct-2407",
+            tokenizer_mode="hf",
+            mistral_format=True,
+        )
+
+
+@pytest.mark.parametrize(
+    ("transformers_version", "model_type", "native_filename", "expect_error"),
+    [
+        pytest.param("5.14.0", "mistral", "tekken.json", True),
+        pytest.param("5.15.0", "mistral", "tekken.json", False),
+        pytest.param("5.14.0", "mistral3", "tekken.json", True),
+        pytest.param("5.15.0", "mistral3", "tekken.json", False),
+        pytest.param("5.14.0", "mixtral", "tokenizer.model.v1", True),
+        pytest.param("5.15.0", "mixtral", "tokenizer.model.v1", False),
+        pytest.param("5.14.0", "mistral", None, False),
+        pytest.param("5.14.0", "qwen3_5_moe", "tekken.json", False),
+    ],
+)
+def test_get_tokenizer_handles_hf_mistral_transformers_compatibility(
+    tmp_path: Path,
+    transformers_version: str,
+    model_type: str,
+    native_filename: str | None,
+    expect_error: bool,
+):
+    if native_filename is not None:
+        (tmp_path / native_filename).touch()
+
+    with (
+        patch(
+            "vllm.tokenizers.registry.get_config",
+            return_value=SimpleNamespace(model_type=model_type),
+        ),
+        patch(
+            "vllm.tokenizers.registry.transformers.__version__",
+            transformers_version,
+        ),
+        patch.object(
+            CachedHfTokenizer,
+            "from_pretrained",
+            return_value=SimpleNamespace(is_fast=True),
+        ) as from_pretrained,
+    ):
+        if expect_error:
+            with pytest.raises(ValueError, match="requires transformers>=5.15.0"):
+                get_tokenizer(str(tmp_path), tokenizer_mode="hf")
+        else:
+            tokenizer = get_tokenizer(str(tmp_path), tokenizer_mode="hf")
+            assert tokenizer.is_fast is True
+
+    assert from_pretrained.called == (not expect_error)
+
+
 def test_customized_tokenizer():
     TokenizerRegistry.register("test_tokenizer", __name__, TestTokenizer.__name__)
 
diff --git a/tests/tool_parsers/common_tests.py b/tests/tool_parsers/common_tests.py
index 925506aa73d4..b3f9ca7907ea 100644
--- a/tests/tool_parsers/common_tests.py
+++ b/tests/tool_parsers/common_tests.py
@@ -49,6 +49,7 @@ class ToolParserTestConfig:
         Special flags:
         allow_empty_or_json_empty_args: True if "" or "{}" both valid for empty args
         supports_typed_arguments: True if the parser supports typed function arguments
+
     """
 
     # Parser identification
@@ -110,6 +111,7 @@ def test_config(self) -> ToolParserTestConfig:
             def test_mistral_specific_feature(self, tool_parser):
                 # Custom test logic
                 pass
+
     """
 
     @pytest.fixture
diff --git a/tests/tool_parsers/test_cohere_command_tool_parser.py b/tests/tool_parsers/test_cohere_command_tool_parser.py
deleted file mode 100644
index f102474b6289..000000000000
--- a/tests/tool_parsers/test_cohere_command_tool_parser.py
+++ /dev/null
@@ -1,510 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import json
-from dataclasses import dataclass, field
-from typing import Any
-
-import pytest
-
-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-from vllm.tool_parsers.cohere_command_tool_parser import (
-    CohereCommand3ToolParser,
-    CohereCommand4ToolParser,
-)
-
-
-@dataclass
-class ExpectedToolCall:
-    id: str
-    name: str
-    arguments: dict
-
-
-@dataclass
-class ToolCallCase:
-    parser_cls: Any
-    model_output: str
-    expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list)
-    expected_reasoning: str | None = None
-    expected_content: str | None = None
-
-
-TOOL_CALL_CASES = [
-    pytest.param(
-        ToolCallCase(
-            parser_cls=CohereCommand3ToolParser,
-            model_output="""\
-<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
-[
-    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
-]
-<|END_ACTION|>""",
-            expected_tool_calls=[
-                ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
-            ],
-            expected_reasoning="i will call foo with query1",
-        ),
-        id="cmd3-single_tool_call",
-    ),
-    pytest.param(
-        ToolCallCase(
-            parser_cls=CohereCommand4ToolParser,
-            model_output="""\
-<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|>
-[
-    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}
-]
-<|END_ACTION|>""",
-            expected_tool_calls=[
-                ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
-            ],
-            expected_reasoning="i will call foo with query1",
-        ),
-        id="cmd4-single_tool_call",
-    ),
-    pytest.param(
-        ToolCallCase(
-            parser_cls=CohereCommand3ToolParser,
-            model_output="""\
-<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
-<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
-            expected_reasoning="This is a rainbow emoji: 🌈",
-            expected_content="foo bar",
-        ),
-        id="cmd3-citations_no_tool_calls",
-    ),
-    pytest.param(
-        ToolCallCase(
-            parser_cls=CohereCommand4ToolParser,
-            model_output="""\
-<|START_THINKING|>This is a rainbow <co>emoji: 🌈</co: 0:[1]><|END_THINKING|>
-<|START_RESPONSE|>foo <co>bar</co: 0:[1,2],1:[3,4]><|END_RESPONSE|>""",
-            expected_reasoning="This is a rainbow emoji: 🌈",
-            expected_content="foo bar",
-        ),
-        id="cmd4-citations_no_tool_calls",
-    ),
-    pytest.param(
-        ToolCallCase(
-            parser_cls=CohereCommand3ToolParser,
-            model_output="""\
-<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|>
-[
-    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}},
-    {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}}
-]
-<|END_ACTION|>""",
-            expected_tool_calls=[
-                ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
-                ExpectedToolCall(id="1", name="bar", arguments={"x": 42}),
-            ],
-            expected_reasoning="first I think about foo",
-        ),
-        id="cmd3-multiple_tool_calls",
-    ),
-    pytest.param(
-        ToolCallCase(
-            parser_cls=CohereCommand4ToolParser,
-            model_output="""\
-<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|>
-[
-    {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}},
-    {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}}
-]
-<|END_ACTION|>""",
-            expected_tool_calls=[
-                ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}),
-                ExpectedToolCall(id="1", name="bar", arguments={"x": 42}),
-            ],
-            expected_reasoning="first I think about foo",
-        ),
-        id="cmd4-multiple_tool_calls",
-    ),
-    pytest.param(
-        ToolCallCase(
-            parser_cls=CohereCommand3ToolParser,
-            model_output="""\
-<|START_THINKING|>just think, no response<|END_THINKING|>""",
-            expected_reasoning="just think, no response",
-        ),
-        id="cmd3-reasoning_only",
-    ),
-    pytest.param(
-        ToolCallCase(
-            parser_cls=CohereCommand4ToolParser,
-            model_output="""\
-<|START_THINKING|>just think, no response<|END_THINKING|>""",
-            expected_reasoning="just think, no response",
-        ),
-        id="cmd4-reasoning_only",
-    ),
-]
-
-
-class MockCohereTokenizer:
-    """Minimal byte-level stand-in for the Cohere tokenizer."""
-
-    def get_vocab(self) -> dict[str, int]:
-        return {}
-
-    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
-        return list(text.encode("utf-8"))
-
-    def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
-        return bytes(ids).decode("utf-8", errors="replace")
-
-
-@pytest.fixture(scope="module")
-def tokenizer() -> MockCohereTokenizer:
-    return MockCohereTokenizer()
-
-
-@pytest.fixture
-def request_obj() -> ChatCompletionRequest:
-    return ChatCompletionRequest(messages=[], model="test-model")
-
-
-REPLACEMENT_CHAR = "\ufffd"
-
-
-def _token_deltas(tokenizer: MockCohereTokenizer, text: str) -> list[str]:
-    """Decode per-token string deltas, buffering incomplete multi-byte chars."""
-    ids = tokenizer.encode(text, add_special_tokens=False)
-    deltas: list[str] = []
-    prev = ""
-    for i in range(1, len(ids) + 1):
-        current = tokenizer.decode(ids[:i], skip_special_tokens=False)
-        if current.endswith(REPLACEMENT_CHAR):
-            continue
-        delta = current[len(prev) :]
-        if delta:
-            deltas.append(delta)
-        prev = current
-    return deltas
-
-
-@dataclass
-class StreamingResult:
-    tool_calls: dict[int, dict]
-    reasoning: str | None
-    content: str | None
-
-
-def _run_streaming_over_deltas(
-    parser,
-    deltas: list[str],
-    request_obj: ChatCompletionRequest,
-) -> StreamingResult:
-    accumulated: dict[int, dict] = {}
-    reasoning_parts: list[str] = []
-    content_parts: list[str] = []
-    previous_text = ""
-
-    for token_str in deltas:
-        current_text = previous_text + token_str
-        delta = parser.extract_tool_calls_streaming(
-            previous_text=previous_text,
-            current_text=current_text,
-            delta_text=token_str,
-            previous_token_ids=[],
-            current_token_ids=[],
-            delta_token_ids=[],
-            request=request_obj,
-        )
-        if delta is not None:
-            if delta.reasoning is not None:
-                reasoning_parts.append(delta.reasoning)
-            if delta.content is not None:
-                content_parts.append(delta.content)
-            for tc in delta.tool_calls:
-                idx = tc.index
-                if idx not in accumulated:
-                    accumulated[idx] = {"id": "", "name": "", "arguments": ""}
-                if tc.id:
-                    accumulated[idx]["id"] = tc.id
-                if tc.function and tc.function.name:
-                    accumulated[idx]["name"] = tc.function.name
-                if tc.function and tc.function.arguments:
-                    accumulated[idx]["arguments"] += tc.function.arguments
-        previous_text = current_text
-
-    return StreamingResult(
-        tool_calls=accumulated,
-        reasoning="".join(reasoning_parts) if reasoning_parts else None,
-        content="".join(content_parts) if content_parts else None,
-    )
-
-
-def _run_streaming(
-    parser,
-    tokenizer: MockCohereTokenizer,
-    model_output: str,
-    request_obj: ChatCompletionRequest,
-) -> StreamingResult:
-    return _run_streaming_over_deltas(
-        parser,
-        _token_deltas(tokenizer, model_output),
-        request_obj,
-    )
-
-
-@pytest.mark.parametrize("case", TOOL_CALL_CASES)
-class TestExtractToolCalls:
-    def test_streaming(
-        self,
-        tokenizer: MockCohereTokenizer,
-        request_obj: ChatCompletionRequest,
-        case: ToolCallCase,
-    ):
-        parser = case.parser_cls(tokenizer)
-        streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj)
-
-        assert len(streamed.tool_calls) == len(case.expected_tool_calls)
-
-        for i, expected_tc in enumerate(case.expected_tool_calls):
-            tc = streamed.tool_calls[i]
-            assert tc["id"] == expected_tc.id
-            assert tc["name"] == expected_tc.name
-            assert json.loads(tc["arguments"]) == expected_tc.arguments
-
-    def test_streaming_reasoning(
-        self,
-        tokenizer: MockCohereTokenizer,
-        request_obj: ChatCompletionRequest,
-        case: ToolCallCase,
-    ):
-        parser = case.parser_cls(tokenizer)
-        streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj)
-
-        assert streamed.reasoning == case.expected_reasoning
-
-    def test_streaming_content(
-        self,
-        tokenizer: MockCohereTokenizer,
-        request_obj: ChatCompletionRequest,
-        case: ToolCallCase,
-    ):
-        parser = case.parser_cls(tokenizer)
-        streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj)
-
-        assert streamed.content == case.expected_content
-
-    def test_nonstreaming(
-        self,
-        request_obj: ChatCompletionRequest,
-        tokenizer: MockCohereTokenizer,
-        case: ToolCallCase,
-    ):
-        parser = case.parser_cls(tokenizer)
-        result = parser.extract_tool_calls(case.model_output, request_obj)
-
-        assert result.tools_called == (len(case.expected_tool_calls) > 0)
-        assert len(result.tool_calls) == len(case.expected_tool_calls)
-
-        for actual_tc, expected_tc in zip(result.tool_calls, case.expected_tool_calls):
-            assert actual_tc.type == "function"
-            assert actual_tc.function.name == expected_tc.name
-            assert json.loads(actual_tc.function.arguments) == expected_tc.arguments
-
-    def test_streaming_nonstreaming_agree(
-        self,
-        tokenizer: MockCohereTokenizer,
-        request_obj: ChatCompletionRequest,
-        case: ToolCallCase,
-    ):
-        parser_streaming = case.parser_cls(tokenizer)
-        parser_nonstreaming = case.parser_cls(tokenizer)
-
-        streamed = _run_streaming(
-            parser_streaming,
-            tokenizer,
-            case.model_output,
-            request_obj,
-        )
-        result = parser_nonstreaming.extract_tool_calls(
-            case.model_output,
-            request_obj,
-        )
-
-        assert len(streamed.tool_calls) == len(result.tool_calls)
-
-        for i, actual_tc in enumerate(result.tool_calls):
-            assert streamed.tool_calls[i]["name"] == actual_tc.function.name
-            assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads(
-                actual_tc.function.arguments
-            )
-
-
-SPECIAL_TOKEN_MARKERS = (
-    "<|START_THINKING|>",
-    "<|END_THINKING|>",
-    "<|START_RESPONSE|>",
-    "<|END_RESPONSE|>",
-    "<|START_ACTION|>",
-    "<|END_ACTION|>",
-    "<|START_TEXT|>",
-    "<|END_TEXT|>",
-)
-
-
-def _multi_token_deltas(
-    tokenizer: MockCohereTokenizer,
-    text: str,
-    chunk_size: int,
-) -> list[str]:
-    ids = tokenizer.encode(text, add_special_tokens=False)
-    deltas: list[str] = []
-    prev = ""
-    i = 0
-    while i < len(ids):
-        end = min(len(ids), i + chunk_size)
-        current = tokenizer.decode(ids[:end], skip_special_tokens=False)
-        i = end
-        if current.endswith(REPLACEMENT_CHAR):
-            continue
-        delta = current[len(prev) :]
-        if delta:
-            deltas.append(delta)
-        prev = current
-    return deltas
-
-
-class TestSpeculativeDecodingMultiTokenDelta:
-    MODEL_OUTPUT = (
-        "<|START_THINKING|> i will call foo with query1<|END_THINKING|>"
-        "<|START_ACTION|>\n"
-        '[\n    {"tool_call_id": "0", "tool_name": "foo", '
-        '"parameters": {"query": "query1"}}\n]\n'
-        "<|END_ACTION|>"
-    )
-
-    @pytest.mark.parametrize(
-        "parser_cls",
-        [CohereCommand3ToolParser, CohereCommand4ToolParser],
-        ids=["cmd3", "cmd4"],
-    )
-    @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6])
-    def test_no_special_token_leak_in_streaming_deltas(
-        self,
-        tokenizer: MockCohereTokenizer,
-        request_obj: ChatCompletionRequest,
-        parser_cls,
-        chunk_size: int,
-    ):
-        parser = parser_cls(tokenizer)
-        chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size)
-
-        previous_text = ""
-        for token_str in chunked:
-            current_text = previous_text + token_str
-            delta = parser.extract_tool_calls_streaming(
-                previous_text=previous_text,
-                current_text=current_text,
-                delta_text=token_str,
-                previous_token_ids=[],
-                current_token_ids=[],
-                delta_token_ids=[],
-                request=request_obj,
-            )
-            previous_text = current_text
-            if delta is None:
-                continue
-
-            fields: list[tuple[str, str | None]] = [
-                ("reasoning", delta.reasoning),
-                ("content", delta.content),
-            ]
-            for tc in delta.tool_calls or []:
-                if tc.function:
-                    fields.append(("tool_call.name", tc.function.name))
-                    fields.append(("tool_call.arguments", tc.function.arguments))
-
-            for marker in SPECIAL_TOKEN_MARKERS:
-                for field_name, value in fields:
-                    assert value is None or marker not in value, (
-                        f"special token {marker!r} leaked into {field_name} "
-                        f"with chunk_size={chunk_size} delta={delta!r}"
-                    )
-
-    @pytest.mark.parametrize(
-        "parser_cls",
-        [CohereCommand3ToolParser, CohereCommand4ToolParser],
-        ids=["cmd3", "cmd4"],
-    )
-    @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6])
-    def test_multi_token_chunks_still_produce_correct_tool_call(
-        self,
-        tokenizer: MockCohereTokenizer,
-        request_obj: ChatCompletionRequest,
-        parser_cls,
-        chunk_size: int,
-    ):
-        parser = parser_cls(tokenizer)
-        chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size)
-        streamed = _run_streaming_over_deltas(parser, chunked, request_obj)
-
-        assert len(streamed.tool_calls) == 1
-        tc = streamed.tool_calls[0]
-        assert tc["id"] == "0"
-        assert tc["name"] == "foo"
-        assert json.loads(tc["arguments"]) == {"query": "query1"}
-
-
-class TestStreamingDeltaShape:
-    @pytest.mark.parametrize(
-        "parser_cls",
-        [CohereCommand3ToolParser, CohereCommand4ToolParser],
-        ids=["cmd3", "cmd4"],
-    )
-    def test_reasoning_and_tool_calls_are_separate_deltas(
-        self,
-        tokenizer: MockCohereTokenizer,
-        request_obj: ChatCompletionRequest,
-        parser_cls,
-    ):
-        parser = parser_cls(tokenizer)
-        model_output = (
-            "<|START_THINKING|> i will call foo with query1<|END_THINKING|>"
-            "<|START_ACTION|>\n"
-            '[\n    {"tool_call_id": "0", "tool_name": "foo", '
-            '"parameters": {"query": "query1"}}\n]\n'
-            "<|END_ACTION|>"
-        )
-
-        token_strings = _token_deltas(tokenizer, model_output)
-        previous_text = ""
-        saw_reasoning = False
-        saw_tool_call = False
-
-        for token_str in token_strings:
-            current_text = previous_text + token_str
-            delta = parser.extract_tool_calls_streaming(
-                previous_text=previous_text,
-                current_text=current_text,
-                delta_text=token_str,
-                previous_token_ids=[],
-                current_token_ids=[],
-                delta_token_ids=[],
-                request=request_obj,
-            )
-            if delta is not None:
-                populated = [
-                    delta.content is not None,
-                    delta.reasoning is not None,
-                    bool(delta.tool_calls),
-                ]
-                assert sum(populated) == 1, (
-                    "A single streaming delta must carry exactly one of "
-                    f"content/reasoning/tool_calls, got {delta!r}"
-                )
-                if delta.reasoning is not None:
-                    saw_reasoning = True
-                if delta.tool_calls:
-                    saw_tool_call = True
-            previous_text = current_text
-
-        assert saw_reasoning, "expected at least one reasoning delta"
-        assert saw_tool_call, "expected at least one tool-call delta"
diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py
index 8d74f0431934..97ed0cd69d4c 100644
--- a/tests/tool_parsers/test_gemma4_tool_parser.py
+++ b/tests/tool_parsers/test_gemma4_tool_parser.py
@@ -897,3 +897,111 @@ def test_streaming_trailing_bare_bool_not_duplicated(self, parser, mock_request)
         }
 
         assert args_text.count("replace_all") == 1
+
+
+class TestBareAndMalformedOpeners:
+    """Tests for bare and malformed tool call openers.
+
+    See https://github.com/vllm-project/vllm/issues/53431.
+    """
+
+    def test_bare_opener_non_streaming(self, parser, mock_request):
+        model_output = '<|tool_call>:get_weather{location:<|"|>London<|"|>}<tool_call|>'
+        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_documented_form_still_works(self, parser, mock_request):
+        model_output = (
+            '<|tool_call>call:get_weather{location:<|"|>London<|"|>}<tool_call|>'
+        )
+        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"
+
+    def test_bare_opener_multiple_args(self, parser, mock_request):
+        model_output = (
+            "<|tool_call>:get_weather{"
+            'location:<|"|>Tokyo<|"|>,unit:<|"|>celsius<|"|>}'
+            "<tool_call|>"
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert result.tool_calls[0].function.name == "get_weather"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"location": "Tokyo", "unit": "celsius"}
+
+    def test_malformed_no_brace_bounded(self, parser, mock_request):
+        model_output = "<|tool_call>call:bad_func no brace<tool_call|>"
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        name = result.tool_calls[0].function.name
+        assert "<tool_call|>" not in name
+
+    def test_malformed_then_well_formed(self, parser, mock_request):
+        model_output = (
+            "<|tool_call>call:bad_func no brace<tool_call|>"
+            '<|tool_call>call:get_weather{location:<|"|>London<|"|>}'
+            "<tool_call|>"
+        )
+        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 == "bad_func no brace"
+        assert result.tool_calls[1].function.name == "get_weather"
+        assert "London" in result.tool_calls[1].function.arguments
+
+    def test_bare_malformed_then_well_formed(self, parser, mock_request):
+        model_output = (
+            "<|tool_call>:bad_func no brace<tool_call|>"
+            '<|tool_call>call:get_weather{location:<|"|>London<|"|>}'
+            "<tool_call|>"
+        )
+        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 == "bad_func no brace"
+        assert result.tool_calls[1].function.name == "get_weather"
+
+    def test_three_mixed_calls(self, parser, mock_request):
+        model_output = (
+            "<|tool_call>:bad no brace<tool_call|>"
+            '<|tool_call>call:search{input:<|"|>hello<|"|>}<tool_call|>'
+            "<|tool_call>:set{flag:true,count:5}<tool_call|>"
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 3
+        assert result.tool_calls[0].function.name == "bad no brace"
+        assert result.tool_calls[1].function.name == "search"
+        assert result.tool_calls[2].function.name == "set"
+
+    def test_colon_in_content_not_broken(self, parser, mock_request):
+        model_output = (
+            "Hello: world"
+            '<|tool_call>:get_weather{location:<|"|>London<|"|>}<tool_call|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert result.content is not None
+        assert "Hello" in result.content
+
+    def test_bare_opener_empty_args(self, parser, mock_request):
+        model_output = "<|tool_call>:set_status{}<tool_call|>"
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert result.tool_calls[0].function.name == "set_status"
diff --git a/tests/tool_parsers/test_gigachat3_tool_parser.py b/tests/tool_parsers/test_gigachat3_tool_parser.py
index a011fbf080d8..e5a1692ffa39 100644
--- a/tests/tool_parsers/test_gigachat3_tool_parser.py
+++ b/tests/tool_parsers/test_gigachat3_tool_parser.py
@@ -333,9 +333,7 @@ def test_streaming_tool_call_with_large_steps(
     model_output_deltas: list[str],
     gigachat_tokenizer: TokenizerLike,
 ):
-    """
-    Test that the closing braces are streamed correctly.
-    """
+    """Test that the closing braces are streamed correctly."""
     tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")(
         gigachat_tokenizer
     )
diff --git a/tests/tool_parsers/test_granite_tool_parser.py b/tests/tool_parsers/test_granite_tool_parser.py
index af3386112f50..5c51f389bee7 100644
--- a/tests/tool_parsers/test_granite_tool_parser.py
+++ b/tests/tool_parsers/test_granite_tool_parser.py
@@ -16,7 +16,7 @@
     split_string_into_token_deltas,
 )
 from vllm.tokenizers import get_tokenizer
-from vllm.tool_parsers.granite_tool_parser import GraniteToolParser
+from vllm.tool_parsers.granite_engine_tool_parser import GraniteEngineToolParser
 
 
 class TestGraniteToolParser(ToolParserTests):
@@ -76,23 +76,6 @@ def test_config(self) -> ToolParserTestConfig:
             single_tool_call_expected_content=None,
             parallel_tool_calls_count=2,
             parallel_tool_calls_names=["get_weather", "get_time"],
-            # xfail markers
-            xfail_streaming={
-                "test_malformed_input": (
-                    "Streaming mode incorrectly creates tool call from malformed JSON"
-                ),
-                "test_surrounding_text": (
-                    "Parser doesn't handle surrounding text correctly in streaming"
-                ),
-                "test_streaming_reconstruction": (
-                    "Streaming mode doesn't strip <|tool_call|> marker from content"
-                ),
-            },
-            xfail_nonstreaming={
-                "test_surrounding_text": (
-                    "Parser doesn't handle surrounding text correctly in non-streaming"
-                ),
-            },
         )
 
     # Granite-Specific Tests
@@ -140,7 +123,7 @@ def test_streaming_parallel_calls_batched_deltas(granite_tokenizer, chunk_size):
     parallel calls must not drop the first call's name. granite streams
     arguments before name, so the name only completes as the next call appears.
     """
-    parser = GraniteToolParser(granite_tokenizer)
+    parser = GraniteEngineToolParser(granite_tokenizer)
     model_output = (
         '<|tool_call|> [{"arguments": {"city": "Tokyo"}, "name": "get_weather"}, '
         '{"arguments": {"timezone": "Asia/Tokyo"}, "name": "get_time"}]'
diff --git a/tests/tool_parsers/test_internlm2_tool_parser.py b/tests/tool_parsers/test_internlm2_tool_parser.py
index 7fd3860ef719..881933890dfb 100644
--- a/tests/tool_parsers/test_internlm2_tool_parser.py
+++ b/tests/tool_parsers/test_internlm2_tool_parser.py
@@ -18,7 +18,6 @@ class TestInternLM2ToolParser(ToolParserTests):
     @pytest.fixture
     def tokenizer(self, default_tokenizer: TokenizerLike) -> TokenizerLike:
         """Add some internlm2 specific tokens to the default vocab."""
-
         tokenizer_vocab = default_tokenizer.get_vocab()
         default_tokenizer.get_vocab = MagicMock()
         tokenizer_vocab.update(
diff --git a/tests/tool_parsers/test_kimi_k2_tool_parser.py b/tests/tool_parsers/test_kimi_k2_tool_parser.py
index b56032b91c17..69c5abfd78a6 100644
--- a/tests/tool_parsers/test_kimi_k2_tool_parser.py
+++ b/tests/tool_parsers/test_kimi_k2_tool_parser.py
@@ -218,6 +218,7 @@ def _split_tool_output_to_deltas(
     Args:
         content: text before tool section
         tool_strs: list of (tool_id, args_json)
+
     """
     deltas = [content, SECTION_BEGIN]
     for tool_id, args_json in tool_strs:
diff --git a/tests/tool_parsers/test_llama3_json_tool_parser.py b/tests/tool_parsers/test_llama3_json_tool_parser.py
index b8e194448184..e11434c0b0fd 100644
--- a/tests/tool_parsers/test_llama3_json_tool_parser.py
+++ b/tests/tool_parsers/test_llama3_json_tool_parser.py
@@ -252,7 +252,7 @@ def test_extract_tool_calls_missing_parameters_and_arguments_key(parser):
 
 
 def test_regex_timeout_handling(parser):
-    """Test regex timeout is handled gracefully"""
+    """Test regex timeout is handled gracefully."""
     fake_problematic_input = "{hello world[A(A=" + "\t)A(A=,\t" * 2
 
     # create a mock regex that raises TimeoutError
diff --git a/tests/tool_parsers/test_llama4_pythonic_tool_parser.py b/tests/tool_parsers/test_llama4_pythonic_tool_parser.py
index a3995c9f4bea..351e4115a6fc 100644
--- a/tests/tool_parsers/test_llama4_pythonic_tool_parser.py
+++ b/tests/tool_parsers/test_llama4_pythonic_tool_parser.py
@@ -247,7 +247,7 @@ def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike):
 
 @pytest.mark.parametrize("streaming", [False])
 def test_regex_timeout_handling(streaming: bool, default_tokenizer: TokenizerLike):
-    """test regex timeout is handled gracefully"""
+    """Test regex timeout is handled gracefully."""
     tool_parser: ToolParser = ToolParserManager.get_tool_parser("llama4_pythonic")(
         default_tokenizer
     )
diff --git a/tests/tool_parsers/test_olmo3_tool_parser.py b/tests/tool_parsers/test_olmo3_tool_parser.py
index 1928607222f9..02b5c621f61a 100644
--- a/tests/tool_parsers/test_olmo3_tool_parser.py
+++ b/tests/tool_parsers/test_olmo3_tool_parser.py
@@ -229,7 +229,7 @@ def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike):
 
 @pytest.mark.parametrize("streaming", [False])
 def test_regex_timeout_handling(streaming: bool, default_tokenizer: TokenizerLike):
-    """test regex timeout is handled gracefully"""
+    """Test regex timeout is handled gracefully."""
     tool_parser: ToolParser = ToolParserManager.get_tool_parser("olmo3")(
         default_tokenizer
     )
diff --git a/tests/tool_parsers/test_phi4mini_tool_parser.py b/tests/tool_parsers/test_phi4mini_tool_parser.py
index eff9fa9bb8ff..98237ff5b271 100644
--- a/tests/tool_parsers/test_phi4mini_tool_parser.py
+++ b/tests/tool_parsers/test_phi4mini_tool_parser.py
@@ -16,7 +16,6 @@ class TestPhi4MiniToolParser(ToolParserTests):
     @pytest.fixture
     def tokenizer(self, default_tokenizer: TokenizerLike) -> TokenizerLike:
         """Add some phi4mini specific tokens to the default vocab."""
-
         tokenizer = default_tokenizer
         tokenizer_vocab = tokenizer.get_vocab()
         tokenizer.get_vocab = MagicMock()
diff --git a/tests/tool_parsers/test_pythonic_tool_parser.py b/tests/tool_parsers/test_pythonic_tool_parser.py
index 05ddc411f142..fec713d5a78d 100644
--- a/tests/tool_parsers/test_pythonic_tool_parser.py
+++ b/tests/tool_parsers/test_pythonic_tool_parser.py
@@ -209,7 +209,7 @@ def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike):
 
 @pytest.mark.parametrize("streaming", [False])
 def test_regex_timeout_handling(streaming: bool, default_tokenizer: TokenizerLike):
-    """test regex timeout is handled gracefully"""
+    """Test regex timeout is handled gracefully."""
     tool_parser: ToolParser = ToolParserManager.get_tool_parser("pythonic")(
         default_tokenizer
     )
diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py
index 946d490f94aa..c397d0ffadcf 100644
--- a/tests/tool_parsers/test_qwen3coder_tool_parser.py
+++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py
@@ -423,7 +423,7 @@ def test_extract_tool_calls(
 def test_extract_tool_calls_fallback_no_tags(
     qwen3_tool_parser,
 ):
-    """Test fallback parsing when XML tags are missing"""
+    """Test fallback parsing when XML tags are missing."""
     model_output = """<function=get_current_weather>
 <parameter=city>
 Dallas
@@ -444,7 +444,7 @@ def test_extract_tool_calls_fallback_no_tags(
 
 
 def test_extract_tool_calls_type_conversion(qwen3_tokenizer):
-    """Test parameter type conversion based on tool schema"""
+    """Test parameter type conversion based on tool schema."""
     tools = [
         ChatCompletionToolsParam(
             type="function",
@@ -862,7 +862,7 @@ def test_extract_tool_calls_streaming(
     expected_tool_calls,
     expected_content,
 ):
-    """Test incremental streaming behavior including typed parameters"""
+    """Test incremental streaming behavior including typed parameters."""
     request = ChatCompletionRequest(model=MODEL, messages=[])
 
     other_content = ""
@@ -932,7 +932,7 @@ def test_extract_tool_calls_streaming(
 def test_extract_tool_calls_missing_closing_parameter_tag(
     qwen3_tool_parser,
 ):
-    """Test handling of missing closing </parameter> tag"""
+    """Test handling of missing closing </parameter> tag."""
     # Using get_current_weather from sample_tools but with malformed XML
     model_output = """Let me check the weather for you:
 <tool_call>
@@ -974,7 +974,7 @@ def test_extract_tool_calls_missing_closing_parameter_tag(
 def test_extract_tool_calls_streaming_missing_closing_tag(
     qwen3_tool_parser, qwen3_tokenizer
 ):
-    """Test streaming with missing closing </parameter> tag"""
+    """Test streaming with missing closing </parameter> tag."""
     # Using get_current_weather from sample_tools but with malformed XML
     model_output = """Let me check the weather for you:
 <tool_call>
@@ -1046,7 +1046,7 @@ def test_extract_tool_calls_streaming_missing_closing_tag(
 
 
 def test_extract_tool_calls_streaming_incremental(qwen3_tool_parser, qwen3_tokenizer):
-    """Test that streaming is truly incremental"""
+    """Test that streaming is truly incremental."""
     model_output = """I'll check the weather.<tool_call>
 <function=get_current_weather>
 <parameter=city>
@@ -1106,7 +1106,7 @@ def test_extract_tool_calls_streaming_incremental(qwen3_tool_parser, qwen3_token
 def test_extract_tool_calls_streaming_missing_opening_tag(
     qwen3_tool_parser, qwen3_tokenizer
 ):
-    """Test streaming with missing opening <tool_call> tag
+    """Test streaming with missing opening <tool_call> tag.
 
     This tests that the streaming parser correctly handles
     tool calls that start directly with <function=...>
diff --git a/tests/tool_parsers/test_step3p5_tool_parser.py b/tests/tool_parsers/test_step3p5_tool_parser.py
index d8f1f9332db6..de21aab79762 100644
--- a/tests/tool_parsers/test_step3p5_tool_parser.py
+++ b/tests/tool_parsers/test_step3p5_tool_parser.py
@@ -366,7 +366,7 @@ def test_extract_tool_calls(
 
 
 def test_extract_tool_calls_fallback_no_tags(step3p5_tool_parser, sample_tools):
-    """Test fallback parsing when XML tags are missing"""
+    """Test fallback parsing when XML tags are missing."""
     model_output = """<function=get_current_weather>
 <parameter=city>
 Dallas
@@ -387,7 +387,7 @@ def test_extract_tool_calls_fallback_no_tags(step3p5_tool_parser, sample_tools):
 
 
 def test_extract_tool_calls_type_conversion(step3p5_tokenizer):
-    """Test parameter type conversion based on tool schema"""
+    """Test parameter type conversion based on tool schema."""
     tools = [
         ChatCompletionToolsParam(
             type="function",
@@ -622,7 +622,7 @@ def test_extract_tool_calls_streaming(
     expected_tool_calls,
     expected_content,
 ):
-    """Test incremental streaming behavior including typed parameters"""
+    """Test incremental streaming behavior including typed parameters."""
     request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools)
 
     other_content = ""
@@ -692,7 +692,7 @@ def test_extract_tool_calls_streaming(
 def test_extract_tool_calls_missing_closing_parameter_tag(
     step3p5_tool_parser, sample_tools
 ):
-    """Test handling of missing closing </parameter> tag"""
+    """Test handling of missing closing </parameter> tag."""
     # Using get_current_weather from sample_tools but with malformed XML
     model_output = """Let me check the weather for you:
 <tool_call>
@@ -734,7 +734,7 @@ def test_extract_tool_calls_missing_closing_parameter_tag(
 def test_extract_tool_calls_streaming_missing_closing_tag(
     step3p5_tool_parser, step3p5_tokenizer, sample_tools
 ):
-    """Test streaming with missing closing </parameter> tag"""
+    """Test streaming with missing closing </parameter> tag."""
     # Using get_current_weather from sample_tools but with malformed XML
     model_output = """Let me check the weather for you:
 <tool_call>
@@ -808,7 +808,7 @@ def test_extract_tool_calls_streaming_missing_closing_tag(
 def test_extract_tool_calls_streaming_incremental(
     step3p5_tool_parser, step3p5_tokenizer, sample_tools
 ):
-    """Test that streaming is truly incremental"""
+    """Test that streaming is truly incremental."""
     model_output = """I'll check the weather.<tool_call>
 <function=get_current_weather>
 <parameter=city>
@@ -864,7 +864,7 @@ def test_extract_tool_calls_streaming_incremental(
 
 
 def test_extract_tool_calls_complex_type_with_single_quote(step3p5_tokenizer):
-    """Test parameter type conversion based on tool schema"""
+    """Test parameter type conversion based on tool schema."""
     tools = [
         ChatCompletionToolsParam(
             type="function",
diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py
index b77b4efaffe6..3221d1fb2c41 100644
--- a/tests/tool_parsers/test_structural_tag_registry.py
+++ b/tests/tool_parsers/test_structural_tag_registry.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -22,6 +23,7 @@
 from vllm.tool_parsers.deepseekv32_engine_tool_parser import (
     DeepSeekV32EngineToolParser,
 )
+from vllm.tool_parsers.deepseekv41_engine_tool_parser import DeepSeekV41EngineToolParser
 from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser
 from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
 from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser
@@ -33,9 +35,11 @@
     SUPPORTED_STRUCTURAL_TAG_MODELS,
     VLLM_BUILTIN_STRUCTURAL_TAG_MODELS,
     XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS,
+    ToolChoice,
     get_function_parameters,
     get_model_structural_tag,
 )
+from vllm.tool_parsers.tool_strict_level import ToolStrictLevel
 
 
 @pytest.fixture
@@ -80,11 +84,38 @@ def test_supported_structural_tag_models_include_vllm_builtins():
     assert "hermes" in VLLM_BUILTIN_STRUCTURAL_TAG_MODELS
 
 
+def test_deepseek_v41_named_choice_emits_one_call(sample_tools):
+    tag = get_model_structural_tag(
+        "deepseek_v4_1",
+        sample_tools,
+        ChatCompletionNamedToolChoiceParam(
+            function=ChatCompletionNamedFunction(name="get_weather")
+        ),
+        reasoning=False,
+    )
+    grammar = Grammar.from_structural_tag(tag)
+    call = (
+        '<｜DSML｜ invoke name="get_weather">\n'
+        '<｜DSML｜ parameter name="city" string="true">Paris</｜DSML｜ parameter>\n'
+        "</｜DSML｜ invoke>\n"
+    )
+    assert _is_grammar_accept_string(
+        grammar, "\n\n<｜DSML｜ calls>\n" + call + "</｜DSML｜ calls>"
+    )
+    assert not _is_grammar_accept_string(
+        grammar, "\n\n<｜DSML｜ calls>\n" + call * 2 + "</｜DSML｜ calls>"
+    )
+    assert (
+        get_model_structural_tag("deepseek_v4_1", sample_tools, "auto", reasoning=False)
+        is None
+    )
+
+
 @pytest.mark.parametrize("choice", ["auto", "required", "get_weather"])
-def test_deepseek_v41_format_constraints_ignore_parameter_schema(
-    choice,
-    sample_tools_strict,
+def test_deepseek_v41_constrains_parameters_after_reasoning(
+    sample_tools_strict, choice
 ):
+    """Strict calls must enforce the schema without requiring another think close."""
     tool_choice = (
         ChatCompletionNamedToolChoiceParam(
             function=ChatCompletionNamedFunction(name=choice)
@@ -92,52 +123,48 @@ def test_deepseek_v41_format_constraints_ignore_parameter_schema(
         if choice == "get_weather"
         else choice
     )
+    sample_tools_strict[0].function.parameters["additionalProperties"] = False
     tag = get_model_structural_tag(
-        "deepseek_v41", sample_tools_strict, tool_choice, reasoning=False
+        "deepseek_v4_1", sample_tools_strict, tool_choice, reasoning=False
     )
     grammar = Grammar.from_structural_tag(tag)
     begin = '\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name="get_weather">\n'
     end = "</｜DSML｜ invoke>\n</｜DSML｜ calls>"
     parameter = (
-        '<｜DSML｜ parameter name="undeclared" string="false">'
-        '{"nested": [true, null, 1.5]}</｜DSML｜ parameter>\n'
-    )
-    # Required city is omitted; extra and repeated fields remain unconstrained.
-    assert _is_grammar_accept_string(grammar, begin + end)
-    assert _is_grammar_accept_string(grammar, begin + parameter * 2 + end)
-    assert not _is_grammar_accept_string(
-        grammar,
-        begin + parameter.replace('{"nested": [true, null, 1.5]}', "invalid") + end,
-    )
-    assert not _is_grammar_accept_string(
-        grammar, (begin + end).replace('name="get_weather"', 'name="unknown"')
+        '<｜DSML｜ parameter name="city" string="true">Paris</｜DSML｜ parameter>\n'
     )
+    assert _is_grammar_accept_string(grammar, begin + parameter + end)
+    for invalid in (
+        "",  # missing required city
+        parameter * 2,
+        parameter.replace('name="city"', 'name="unknown"'),
+        parameter.replace('string="true">Paris', 'string="false">42'),
+        parameter.replace("｜ parameter", "｜parameter"),
+        '{"city":"Paris"}',  # the encoder emits DSML parameters, not a JSON body
+    ):
+        assert not _is_grammar_accept_string(grammar, begin + invalid + end)
     assert not _is_grammar_accept_string(
-        grammar, begin + parameter.replace("｜ parameter", "｜parameter") + end
+        grammar, "reason</think>" + begin + parameter + end
     )
     assert _is_grammar_accept_string(grammar, "Hello") == (choice == "auto")
 
 
-def test_deepseek_v41_named_choice_emits_one_call(sample_tools):
+def test_deepseek_v41_non_strict_parallel_calls_keep_typed_dsml(sample_tools):
+    sample_tools[0].function.strict = False
     tag = get_model_structural_tag(
-        "deepseek_v41",
-        sample_tools,
-        ChatCompletionNamedToolChoiceParam(
-            function=ChatCompletionNamedFunction(name="get_weather")
-        ),
-        reasoning=False,
+        "deepseek_v4_1", sample_tools, "required", reasoning=False
     )
     grammar = Grammar.from_structural_tag(tag)
-    call = '<｜DSML｜ invoke name="get_weather">\n</｜DSML｜ invoke>\n'
-    assert _is_grammar_accept_string(
-        grammar, "\n\n<｜DSML｜ calls>\n" + call + "</｜DSML｜ calls>"
+    call = (
+        '<｜DSML｜ invoke name="get_weather">\n'
+        '<｜DSML｜ parameter name="extra" string="false">'
+        '{"nested":[true,null,1.5]}</｜DSML｜ parameter>\n'
+        "</｜DSML｜ invoke>\n"
     )
+    output = "\n\n<｜DSML｜ calls>\n" + call * 2 + "</｜DSML｜ calls>"
+    assert _is_grammar_accept_string(grammar, output)
     assert not _is_grammar_accept_string(
-        grammar, "\n\n<｜DSML｜ calls>\n" + call * 2 + "</｜DSML｜ calls>"
-    )
-    assert (
-        get_model_structural_tag("deepseek_v41", sample_tools, "auto", reasoning=False)
-        is None
+        grammar, output.replace('{"nested":[true,null,1.5]}', "invalid")
     )
 
 
@@ -157,11 +184,11 @@ def test_get_model_structural_tag_supports_all_xgrammar_builtins(
 
 
 def test_get_model_structural_tag_supports_vllm_hermes(
-    sample_tools: list[ChatCompletionToolsParam],
+    sample_tools_strict: list[ChatCompletionToolsParam],
 ):
     tag = get_model_structural_tag(
         model="hermes",
-        tools=sample_tools,
+        tools=sample_tools_strict,
         tool_choice="required",
         reasoning=False,
     )
@@ -199,6 +226,22 @@ def test_get_model_structural_tag_supports_vllm_hermes(
         assert content["json_schema"] == expected_schema
 
 
+def test_hermes_required_without_strict_leaves_arguments_free(
+    sample_tools: list[ChatCompletionToolsParam],
+):
+    tag = get_model_structural_tag(
+        model="hermes",
+        tools=sample_tools,
+        tool_choice="required",
+        reasoning=False,
+    )
+
+    assert isinstance(tag, StructuralTag)
+    tags = tag.model_dump()["format"]["tags"]
+    assert tags
+    assert all(tag_dump["content"]["json_schema"] is True for tag_dump in tags)
+
+
 def test_hermes_required_tool_calls_use_empty_separator():
     tools = [
         ChatCompletionToolsParam(
@@ -287,10 +330,19 @@ def _k3_tools(*calls: str) -> str:
     return f"{_K3_TOOLS_OPEN}{''.join(calls)}{_K3_TOOLS_CLOSE}"
 
 
-def _k3_grammar(tool_choice, tools=None):
+def _as_strict(tools: list[ChatCompletionToolsParam]) -> list[ChatCompletionToolsParam]:
+    for tool in tools:
+        tool.function.strict = True
+    return tools
+
+
+def _k3_grammar(tool_choice, tools=None, *, strict: bool = True):
+    """Build the K3 grammar; tools are strict by default so the typed
+    argument channel is exercised (an unset ``strict`` leaves arguments free)."""
+    tools = tools if tools is not None else _k3_tools_by_name()
     tag = get_model_structural_tag(
         model="kimi_k3",
-        tools=tools if tools is not None else _k3_tools_by_name(),
+        tools=_as_strict(tools) if strict else tools,
         tool_choice=tool_choice,
         reasoning=False,
     )
@@ -393,6 +445,18 @@ def test_kimi_k3_required_rejects_invalid(body: str):
     assert not _is_grammar_accept_string(_k3_grammar("required"), body)
 
 
+def test_kimi_k3_required_without_strict_leaves_arguments_free():
+    # An unset ``strict`` pins the call envelope only: a value the schema
+    # would reject is accepted, and the declared-required argument may be
+    # omitted.
+    grammar = _k3_grammar("required", strict=False)
+    body = _k3_response() + _k3_tools(
+        _k3_call("get_weather", _k3_arg("days", "number", "abc"))
+    )
+    assert _is_grammar_accept_string(grammar, body)
+    assert not _is_grammar_accept_string(_k3_grammar("required"), body)
+
+
 def test_kimi_k3_schema_without_required_accepts_empty_call():
     tools = [
         ChatCompletionToolsParam(
@@ -457,6 +521,7 @@ def test_get_model_structural_tag_supports_named_tool_choice(
         (DeepSeekV31ToolParser, "deepseek_v3_1"),
         (DeepSeekV32EngineToolParser, "deepseek_v3_2"),
         (DeepSeekV4EngineToolParser, "deepseek_v4"),
+        (DeepSeekV41EngineToolParser, "deepseek_v4_1"),
         (Glm47MoeModelToolParser, "glm_4_7"),
         (Hermes2ProToolParser, "hermes"),
         (KimiK2ToolParser, "kimi"),
@@ -524,7 +589,11 @@ def fake_get_model_structural_tag(*, reasoning: bool, **kwargs):
     assert captured == [False]
 
 
+@pytest.mark.parametrize(
+    "parser_cls", [Qwen3EngineToolParser, DeepSeekV41EngineToolParser]
+)
 def test_unified_parser_get_structural_tag_disables_reasoning(
+    parser_cls,
     monkeypatch: pytest.MonkeyPatch,
     sample_tools_strict: list[ChatCompletionToolsParam],
 ):
@@ -540,7 +609,7 @@ def fake_get_model_structural_tag(*, reasoning: bool, **kwargs):
     )
 
     class TestParser(DelegatingParser):
-        tool_parser_cls = Qwen3EngineToolParser
+        tool_parser_cls = parser_cls
 
     request = ChatCompletionRequest(
         messages=[],
@@ -549,7 +618,7 @@ class TestParser(DelegatingParser):
         tool_choice="auto",
     )
     parser = TestParser(MagicMock(), tools=sample_tools_strict)
-    parser.reasoning_parser = MagicMock(adjust_request=lambda request: request)
+    parser._reasoning_parser = MagicMock(adjust_request=lambda request: request)
 
     parser.adjust_request(request)
 
@@ -755,3 +824,209 @@ def test_kimi_k3_forced_tool_choice_builds_single_mandatory_call():
     response_only = _k3_response("no call here")
     assert _is_grammar_accept_string(grammar, ok)
     assert not _is_grammar_accept_string(grammar, response_only)
+
+
+def _pins_argument_schema(tag: StructuralTag) -> bool:
+    return '"json_schema": {' in json.dumps(tag.model_dump(), ensure_ascii=False)
+
+
+def _dumped(tag: StructuralTag) -> str:
+    return json.dumps(tag.model_dump(), ensure_ascii=False)
+
+
+def test_tool_strict_level_auto_is_the_default(
+    sample_tools: list[ChatCompletionToolsParam],
+):
+    """Auto + no strict tool gets no tag unless the operator raises the floor."""
+    assert (
+        get_model_structural_tag(
+            model="deepseek_v4",
+            tools=sample_tools,
+            tool_choice="auto",
+            reasoning=False,
+        )
+        is None
+    )
+
+
+@pytest.mark.parametrize("model", sorted(SUPPORTED_STRUCTURAL_TAG_MODELS))
+def test_tool_strict_level_function_lifts_auto_gate(
+    model: str,
+    sample_tools: list[ChatCompletionToolsParam],
+):
+    tag = get_model_structural_tag(
+        model=model,
+        tools=sample_tools,
+        tool_choice="auto",
+        reasoning=False,
+        strict_level=ToolStrictLevel.FUNCTION,
+    )
+
+    assert tag is not None
+
+
+@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
+def test_tool_strict_level_function_pins_envelope_only(
+    model: str,
+    sample_tools: list[ChatCompletionToolsParam],
+):
+    """The request's own tools stay untouched; unset ``strict`` stays free."""
+    tag = get_model_structural_tag(
+        model=model,
+        tools=sample_tools,
+        tool_choice="auto",
+        reasoning=False,
+        strict_level=ToolStrictLevel.FUNCTION,
+    )
+
+    assert tag is not None
+    assert not _pins_argument_schema(tag)
+    assert all(tool.function.strict is None for tool in sample_tools)
+
+
+@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
+def test_tool_strict_level_parameter_pins_argument_schemas(
+    model: str,
+    sample_tools: list[ChatCompletionToolsParam],
+):
+    tag = get_model_structural_tag(
+        model=model,
+        tools=sample_tools,
+        tool_choice="auto",
+        reasoning=False,
+        strict_level=ToolStrictLevel.PARAMETER,
+    )
+
+    assert tag is not None
+    assert _pins_argument_schema(tag)
+
+
+def test_tool_strict_level_function_keeps_client_strict_tools_strict(
+    sample_tools_strict: list[ChatCompletionToolsParam],
+):
+    tag = get_model_structural_tag(
+        model="deepseek_v4",
+        tools=sample_tools_strict,
+        tool_choice="auto",
+        reasoning=False,
+        strict_level=ToolStrictLevel.FUNCTION,
+    )
+
+    assert tag is not None
+    assert _pins_argument_schema(tag)
+
+
+@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
+def test_absent_strict_is_not_pinned_next_to_a_strict_tool(model: str):
+    """A strict tool must not drag a neighbour with unset ``strict`` into
+    schema enforcement."""
+    weather = ChatCompletionToolsParam(
+        type="function",
+        function={
+            "name": "get_weather",
+            "strict": True,
+            "parameters": {
+                "type": "object",
+                "properties": {"city": {"type": "string"}},
+                "required": ["city"],
+            },
+        },
+    )
+    search = ChatCompletionToolsParam(
+        type="function",
+        function={
+            "name": "search",
+            "parameters": {
+                "type": "object",
+                "properties": {"query": {"type": "string"}},
+                "required": ["query"],
+            },
+        },
+    )
+    tag = get_model_structural_tag(
+        model=model,
+        tools=[weather, search],
+        tool_choice="auto",
+        reasoning=False,
+    )
+
+    assert tag is not None
+    dumped = _dumped(tag)
+    assert '"city"' in dumped
+    assert '"query"' not in dumped
+
+
+@pytest.mark.parametrize(
+    "tool_choice",
+    [
+        "required",
+        ChatCompletionNamedToolChoiceParam(
+            function=ChatCompletionNamedFunction(name="get_weather")
+        ),
+    ],
+)
+@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
+def test_forced_tool_choice_pins_schema_only_for_strict_tools(
+    model: str,
+    tool_choice: ToolChoice,
+    sample_tools: list[ChatCompletionToolsParam],
+    sample_tools_strict: list[ChatCompletionToolsParam],
+):
+    """Required / named always constrain the call; the argument schema
+    follows the per-tool ``strict`` (or the parameter level)."""
+    free = get_model_structural_tag(
+        model=model,
+        tools=sample_tools,
+        tool_choice=tool_choice,
+        reasoning=False,
+        strict_level=ToolStrictLevel.FUNCTION,
+    )
+    pinned = get_model_structural_tag(
+        model=model,
+        tools=sample_tools_strict,
+        tool_choice=tool_choice,
+        reasoning=False,
+    )
+    floored = get_model_structural_tag(
+        model=model,
+        tools=sample_tools,
+        tool_choice=tool_choice,
+        reasoning=False,
+        strict_level=ToolStrictLevel.PARAMETER,
+    )
+
+    assert free is not None and not _pins_argument_schema(free)
+    assert pinned is not None and _pins_argument_schema(pinned)
+    assert floored is not None and _pins_argument_schema(floored)
+
+
+def test_tool_strict_level_parameter_overrides_client_strict_false():
+    tools = [
+        ChatCompletionToolsParam(
+            type="function",
+            function={
+                "name": "get_weather",
+                "parameters": {"type": "object", "properties": {}},
+                "strict": False,
+            },
+        )
+    ]
+    tag = get_model_structural_tag(
+        model="deepseek_v4",
+        tools=tools,
+        tool_choice="required",
+        reasoning=False,
+        strict_level=ToolStrictLevel.PARAMETER,
+    )
+
+    assert tag is not None
+    assert _pins_argument_schema(tag)
+
+
+def test_tool_strict_level_from_name():
+    assert ToolStrictLevel.from_name("AUTO") is ToolStrictLevel.AUTO
+    assert ToolStrictLevel.from_name("Parameter") is ToolStrictLevel.PARAMETER
+    with pytest.raises(ValueError, match="expected one of auto, function, parameter"):
+        ToolStrictLevel.from_name("strict")
+    with pytest.raises(ValueError, match="expected one of auto, function, parameter"):
+        ToolStrictLevel.from_name("off")
diff --git a/tests/tool_use/test_chat_completions.py b/tests/tool_use/test_chat_completions.py
index e5bb475875ac..585ca9314b0d 100644
--- a/tests/tool_use/test_chat_completions.py
+++ b/tests/tool_use/test_chat_completions.py
@@ -166,8 +166,7 @@ async def test_chat_completion_with_tools(
 async def test_response_format_with_tool_choice_required(
     client: openai.AsyncOpenAI, server_config: ServerConfig
 ):
-    """
-    Test that combining response_format: json_object with tool_choice: required
+    """Test that combining response_format: json_object with tool_choice: required
     doesn't crash the engine.
 
     Before the fix, this would cause a validation error:
diff --git a/tests/tool_use/test_gemma4_responses_adjust_request.py b/tests/tool_use/test_gemma4_responses_adjust_request.py
index 937be9d016a4..fabfa7476660 100644
--- a/tests/tool_use/test_gemma4_responses_adjust_request.py
+++ b/tests/tool_use/test_gemma4_responses_adjust_request.py
@@ -171,7 +171,7 @@ def test_tool_parser_adjust_request_builds_valid_response_text_config() -> None:
 
 
 def test_gemma4_required_skips_structured_outputs_chatcompletion() -> None:
-    """required + ChatCompletion: ``Gemma4EngineToolParser`` must skip the
+    """Required + ChatCompletion: ``Gemma4EngineToolParser`` must skip the
     forced JSON ``structured_outputs`` so the model emits its native
     ``<|tool_call>`` syntax. The base parser constrained output to JSON the
     native parser cannot read, leaking it as content with empty
@@ -187,7 +187,7 @@ def test_gemma4_required_skips_structured_outputs_chatcompletion() -> None:
 
 
 def test_gemma4_named_skips_structured_outputs_chatcompletion() -> None:
-    """named + ChatCompletion: the forced single-function JSON schema must be
+    """Named + ChatCompletion: the forced single-function JSON schema must be
     skipped, same as ``required``.
     """
     parser = Gemma4ToolParser(_StubTokenizer())
@@ -202,7 +202,7 @@ def test_gemma4_named_skips_structured_outputs_chatcompletion() -> None:
 
 
 def test_gemma4_required_skips_structured_outputs_responses() -> None:
-    """required + Responses: the forced JSON schema (``request.text``) must be
+    """Required + Responses: the forced JSON schema (``request.text``) must be
     skipped so the native delimiters reach the extractor.
     """
     parser = Gemma4ToolParser(_StubTokenizer())
@@ -215,7 +215,7 @@ def test_gemma4_required_skips_structured_outputs_responses() -> None:
 
 
 def test_gemma4_named_skips_structured_outputs_responses() -> None:
-    """named (``ToolChoiceFunction``) + Responses: the forced single-function
+    """Named (``ToolChoiceFunction``) + Responses: the forced single-function
     JSON schema must be skipped.
     """
     parser = Gemma4ToolParser(_StubTokenizer())
@@ -230,7 +230,7 @@ def test_gemma4_named_skips_structured_outputs_responses() -> None:
 
 
 def test_gemma4_keeps_special_tokens_with_tools_thinking_disabled() -> None:
-    """tools active + thinking disabled: ``skip_special_tokens`` must stay
+    """Tools active + thinking disabled: ``skip_special_tokens`` must stay
     False so ``<|tool_call>`` delimiters reach the extractor. The merged
     enable_thinking early-return stripped them, breaking tool calling when
     thinking is off.
diff --git a/tests/tool_use/test_parallel_tool_calls.py b/tests/tool_use/test_parallel_tool_calls.py
index 4cfd165f1a8e..d1dc6e47921c 100644
--- a/tests/tool_use/test_parallel_tool_calls.py
+++ b/tests/tool_use/test_parallel_tool_calls.py
@@ -237,10 +237,7 @@ async def test_parallel_tool_calls_with_results(
 async def test_parallel_tool_calls_false(
     client: openai.AsyncOpenAI, server_config: ServerConfig
 ):
-    """
-    Ensure only one tool call is returned when parallel_tool_calls is False.
-    """
-
+    """Ensure only one tool call is returned when parallel_tool_calls is False."""
     models = await client.models.list()
     model_name: str = models.data[0].id
     messages = apply_parallel_tool_system_prompt(
diff --git a/tests/tools/test_check_test_tethering.py b/tests/tools/test_check_test_tethering.py
new file mode 100644
index 000000000000..34f4fd179c8e
--- /dev/null
+++ b/tests/tools/test_check_test_tethering.py
@@ -0,0 +1,879 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the test-tethering pre-commit check.
+
+The check answers one question - "does any Buildkite job (in
+``.buildkite/test_areas/`` or the legacy ``.buildkite/test-amd.yaml``) collect
+this test file?" - by parsing job ``commands`` as shell. The risk that
+matters is a *false tether*: a misparse that invents coverage makes the gate
+pass while the test still never runs, which is the exact failure the check
+exists to catch. So the parser cases below pin both directions, and
+:func:`test_no_selection_matches_everything` guards the catastrophic version of
+it.
+"""
+
+import textwrap
+from pathlib import Path
+
+import pytest
+
+import tools.pre_commit.check_test_tethering as checker
+from tools.pre_commit.check_test_tethering import (
+    FindSelection,
+    PytestSelection,
+    _parse_command,
+    _to_repo_relative,
+    all_test_modules,
+    allowlist_entries_missing_reason,
+    is_test_module,
+    is_tethered,
+    load_allowlist,
+    load_selections,
+    main,
+    normalize_test_path,
+    run_changed_files_check,
+    run_full_scan,
+)
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def runs(command: str, test_file: str) -> bool:
+    """True if `command` collects `test_file` (path relative to ``tests/``)."""
+    return any(sel.runs(test_file) for sel in _parse_command(command))
+
+
+# --------------------------------------------------------------------------- #
+# pytest path arguments
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    ("command", "test_file", "expected"),
+    [
+        # Directory sweeps cover everything nested underneath.
+        ("pytest -v -s kernels", "kernels/test_attention.py", True),
+        ("pytest -v -s kernels/", "kernels/moe/test_moe.py", True),
+        ("pytest -v -s kernels", "lora/test_lora.py", False),
+        # An exact file argument covers that file and nothing beside it.
+        ("pytest -v -s kernels/test_a.py", "kernels/test_a.py", True),
+        ("pytest -v -s kernels/test_a.py", "kernels/test_b.py", False),
+        # Commands run from the repo root and from tests/ alike.
+        ("pytest -v -s tests/kernels/test_a.py", "kernels/test_a.py", True),
+        ("pytest -v -s ./kernels/test_a.py", "kernels/test_a.py", True),
+        (
+            "pytest -v -s /vllm-workspace/tests/kernels/test_a.py",
+            "kernels/test_a.py",
+            True,
+        ),
+        # A ::nodeid selector still collects the file.
+        ("pytest -v -s kernels/test_a.py::test_one", "kernels/test_a.py", True),
+        ("pytest -v -s kernels/test_a.py::Cls::test_one", "kernels/test_a.py", True),
+        # Several positional paths, including a bare root-level file.
+        ("pytest -v -s test_envs.py test_outputs.py", "test_outputs.py", True),
+        ("pytest -v -s test_envs.py test_outputs.py", "test_regression.py", False),
+        # A bare word is a path only when it is really a directory under tests/.
+        ("pytest -v -s samplers", "samplers/test_sampler.py", True),
+        # Glob path args do not cross a directory separator, matching the shell.
+        ("pytest -v -s kernels/test_*.py", "kernels/test_a.py", True),
+        ("pytest -v -s kernels/test_*.py", "kernels/moe/test_a.py", False),
+        ("pytest -v -s kernels/test_*.py", "kernels/helper.py", False),
+    ],
+)
+def test_pytest_path_arguments(command, test_file, expected):
+    assert runs(command, test_file) is expected
+
+
+# --------------------------------------------------------------------------- #
+# --ignore / --deselect
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    ("command", "test_file", "expected"),
+    [
+        # --ignore removes a file from an otherwise covering sweep, in both the
+        # `=` and space-separated spellings.
+        ("pytest -v -s kernels --ignore=kernels/test_a.py", "kernels/test_a.py", False),
+        ("pytest -v -s kernels --ignore kernels/test_a.py", "kernels/test_a.py", False),
+        # ...and only that file.
+        ("pytest -v -s kernels --ignore=kernels/test_a.py", "kernels/test_b.py", True),
+        # --ignore of a directory removes the whole subtree.
+        ("pytest -v -s kernels --ignore=kernels/moe", "kernels/moe/test_moe.py", False),
+        (
+            "pytest -v -s lora --ignore-glob=lora/*_tp.py",
+            "lora/test_llama_tp.py",
+            False,
+        ),
+        # --ignore-glob is a whole-path fnmatch (pytest semantics): its `*`
+        # crosses `/`, so a bare `*_tp.py` drops a nested match too. A positional
+        # path glob would not - that is the distinction the parser has to keep.
+        ("pytest -v -s kernels --ignore-glob=*_tp.py", "kernels/test_a_tp.py", False),
+        (
+            "pytest -v -s kernels --ignore-glob=*_tp.py",
+            "kernels/nested/test_a_tp.py",
+            False,
+        ),
+        ("pytest -v -s kernels --ignore-glob '*_tp.py'", "kernels/test_a.py", True),
+        # A node-id deselect leaves the file collected; a bare-path one does not.
+        (
+            "pytest -v -s kernels --deselect kernels/test_a.py::test_one",
+            "kernels/test_a.py",
+            True,
+        ),
+        (
+            "pytest -v -s kernels --deselect kernels/test_a.py",
+            "kernels/test_a.py",
+            False,
+        ),
+    ],
+)
+def test_ignore_and_deselect(command, test_file, expected):
+    assert runs(command, test_file) is expected
+
+
+# --------------------------------------------------------------------------- #
+# Options, env prefixes and other token noise
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    ("command", "test_file", "expected"),
+    [
+        # A -k / -m filter narrows which tests run but does not un-run the file.
+        ("pytest -v -s kernels -k 'not slow'", "kernels/test_a.py", True),
+        ("pytest -v -s kernels -m core_model", "kernels/test_a.py", True),
+        # A value-taking option's value is never mistaken for a path.
+        ("pytest -m core_model", "core_model", False),
+        ("pytest -k kernels", "kernels/test_a.py", False),
+        ("pytest -n 4 kernels/test_a.py", "kernels/test_a.py", True),
+        # Buildkite shard flags carry a $$VAR and must not swallow the path.
+        (
+            (
+                "pytest -v -s kernels --shard-id=$$BUILDKITE_PARALLEL_JOB "
+                "--num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT"
+            ),
+            "kernels/test_a.py",
+            True,
+        ),
+        # A leading env-var assignment is not a path.
+        (
+            "VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s kernels",
+            "kernels/test_a.py",
+            True,
+        ),
+        # `python -m pytest` is the same invocation.
+        ("python3 -m pytest -v -s kernels/test_a.py", "kernels/test_a.py", True),
+        # Quoted paths, wrapper commands and redirects all survive tokenizing.
+        ("pytest -v -s 'kernels/test_a.py'", "kernels/test_a.py", True),
+        ("timeout 600 pytest kernels/test_a.py", "kernels/test_a.py", True),
+        ("pytest kernels/test_a.py > out.log 2>&1", "kernels/test_a.py", True),
+        # Each sub-command of a compound line is classified on its own.
+        ("pytest kernels/test_a.py && pytest lora/test_b.py", "lora/test_b.py", True),
+        ("pytest kernels/test_a.py ; pytest lora/test_b.py", "lora/test_b.py", True),
+        # An unresolvable $VAR path yields no coverage rather than a guess.
+        ("pytest -v -s $TEST_TARGET", "kernels/test_a.py", False),
+    ],
+)
+def test_option_and_token_handling(command, test_file, expected):
+    assert runs(command, test_file) is expected
+
+
+# --------------------------------------------------------------------------- #
+# find | xargs pytest pipelines
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    ("command", "test_file", "expected"),
+    [
+        ("find kernels -name 'test_*.py' | xargs pytest -v", "kernels/test_a.py", True),
+        (
+            "find kernels -name 'test_*.py' | xargs pytest -v",
+            "kernels/moe/test_b.py",
+            True,
+        ),
+        # A non-matching filename is not collected by the find expression.
+        (
+            "find kernels -name 'test_*.py' | xargs pytest -v",
+            "kernels/helper.py",
+            False,
+        ),
+        # -maxdepth 1 keeps direct children and drops nested ones.
+        (
+            "find kernels -maxdepth 1 -name 'test_*.py' | xargs pytest",
+            "kernels/test_a.py",
+            True,
+        ),
+        (
+            "find kernels -maxdepth 1 -name 'test_*.py' | xargs pytest",
+            "kernels/moe/test_b.py",
+            False,
+        ),
+        # -not -name excludes its target and leaves siblings alone.
+        (
+            "find kernels -name 'test_*.py' -not -name 'test_skip.py' | xargs pytest",
+            "kernels/test_skip.py",
+            False,
+        ),
+        (
+            "find kernels -name 'test_*.py' -not -name 'test_skip.py' | xargs pytest",
+            "kernels/test_a.py",
+            True,
+        ),
+        # The root confines the sweep.
+        ("find kernels -name 'test_*.py' | xargs pytest", "lora/test_a.py", False),
+        # -maxdepth is honored even when the find root is the tests/ root itself
+        # ("." or "tests/"), not just a subdir.
+        (
+            "find . -maxdepth 1 -name 'test_*.py' | xargs pytest",
+            "test_regression.py",
+            True,
+        ),
+        (
+            "find . -maxdepth 1 -name 'test_*.py' | xargs pytest",
+            "v1/test_scheduler.py",
+            False,
+        ),
+        (
+            "find tests/ -maxdepth 1 -name 'test_*.py' | xargs pytest",
+            "v1/test_scheduler.py",
+            False,
+        ),
+        # -exec pytest wires the sweep just as `| xargs pytest` does.
+        (
+            "find compile -name 'test_*.py' -exec pytest -s -v {} \\;",
+            "compile/test_a.py",
+            True,
+        ),
+        # A pytest in a *separate* command on the line (not fed by the find) does
+        # not arm the sweep - the find here is piped to `wc`.
+        (
+            "find kernels -name 'test_*.py' | wc -l ; pytest test_regression.py",
+            "kernels/test_a.py",
+            False,
+        ),
+        (
+            "find kernels -name 'test_*.py' -delete && pytest test_regression.py",
+            "kernels/test_a.py",
+            False,
+        ),
+        # ...but only a *non-filtering* intermediate stage. `sort` reorders,
+        # every match still reaches pytest.
+        (
+            "find kernels -name 'test_*.py' | sort | xargs pytest",
+            "kernels/test_a.py",
+            True,
+        ),
+        # `head` drops all but one match - the sweep is not fully tethered.
+        (
+            "find kernels -name 'test_*.py' | head -n 1 | xargs pytest",
+            "kernels/test_a.py",
+            False,
+        ),
+        (
+            "find kernels -name 'test_*.py' | grep attention | xargs pytest",
+            "kernels/test_a.py",
+            False,
+        ),
+        # `xargs` runs `echo` here, not pytest - the trailing `pytest` is echo's
+        # argument.
+        (
+            "find kernels -name 'test_*.py' | xargs echo pytest",
+            "kernels/test_a.py",
+            False,
+        ),
+        # An `-exec` with no `\;` / `+` terminator is a `find` syntax error - it
+        # never runs pytest.
+        (
+            "find kernels -name 'test_*.py' -exec pytest {}",
+            "kernels/test_a.py",
+            False,
+        ),
+        # An `-exec pytest` with no `{}` runs only the named file for every
+        # match - the matched files never reach pytest.
+        (
+            "find kernels -name 'test_*.py' -exec pytest test_regression.py \\;",
+            "kernels/test_a.py",
+            False,
+        ),
+        # `{}` as an option *value* (`--ignore {}`) doesn't reach pytest as a
+        # test path either - only `test_regression.py` runs.
+        (
+            (
+                "find kernels -name 'test_*.py' -exec pytest test_regression.py "
+                "--ignore {} \\;"
+            ),
+            "kernels/test_a.py",
+            False,
+        ),
+        # xargs value-flags with a space don't swallow the command.
+        (
+            "find kernels -name 'test_*.py' | xargs -0 -I {} pytest {}",
+            "kernels/test_a.py",
+            True,
+        ),
+        # `xargs env FOO=1 pytest` - xargs runs `env`, which runs pytest.
+        (
+            "find kernels -name 'test_*.py' | xargs env FOO=1 pytest",
+            "kernels/test_a.py",
+            True,
+        ),
+    ],
+)
+def test_find_pipelines(command, test_file, expected):
+    assert runs(command, test_file) is expected
+
+
+def test_find_pipeline_parses_as_find_selection():
+    (selection,) = [
+        s
+        for s in _parse_command("find kernels -name 'test_*.py' | xargs pytest -v")
+        if isinstance(s, FindSelection)
+    ]
+    assert selection.root == "kernels"
+    assert selection.name_globs == ["test_*.py"]
+
+
+# --------------------------------------------------------------------------- #
+# Direct file runners
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    ("command", "test_file", "expected"),
+    [
+        (
+            "python3 standalone_tests/lazy_imports.py",
+            "standalone_tests/lazy_imports.py",
+            True,
+        ),
+        (
+            "python standalone_tests/lazy_imports.py",
+            "standalone_tests/lazy_imports.py",
+            True,
+        ),
+        # `timeout -k <grace> <dur>` wrapping a runner: the grace value must not
+        # be consumed as the duration (which would hide `python`).
+        (
+            "timeout -k 5s 30s python standalone_tests/lazy_imports.py",
+            "standalone_tests/lazy_imports.py",
+            True,
+        ),
+        (
+            "torchrun --nproc-per-node 2 distributed/test_a.py",
+            "distributed/test_a.py",
+            True,
+        ),
+        # The option's value is a count, not a path.
+        ("torchrun --nproc-per-node 2 distributed/test_a.py", "2", False),
+        (
+            "VLLM_X=1 python3 standalone_tests/lazy_imports.py | grep -q ok",
+            "standalone_tests/lazy_imports.py",
+            True,
+        ),
+        # Only the executed script counts; the tokens after it are that script's
+        # own argv. `--suffix v1` must not tether the whole tests/v1 subtree.
+        (
+            (
+                "python3 features/tensorize_vllm_model.py --model facebook/opt-125m "
+                "serialize --serialized-directory /tmp/ --suffix v1"
+            ),
+            "v1/test_scheduler.py",
+            False,
+        ),
+        (
+            (
+                "python3 features/tensorize_vllm_model.py --model facebook/opt-125m "
+                "serialize --serialized-directory /tmp/ --suffix v1"
+            ),
+            "features/tensorize_vllm_model.py",
+            True,
+        ),
+        # A script named by an absolute `$VAR` path still resolves.
+        (
+            'python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/e2e/test_x.py"',
+            "v1/e2e/test_x.py",
+            True,
+        ),
+        # `python -m <module>` runs the module, not a script - only `-m pytest`
+        # collects tests. `compileall kernels` byte-compiles, it does not run.
+        ("python -m compileall kernels", "kernels/test_a.py", False),
+        ("python -m mypy kernels/test_a.py", "kernels/test_a.py", False),
+        # `coverage run -m pytest` still resolves through to pytest.
+        ("coverage run -m pytest kernels/test_a.py", "kernels/test_a.py", True),
+        # `python -c '<code>' file.py` runs the code string; `file.py` is argv.
+        ("python -c 'pass' kernels/test_a.py", "kernels/test_a.py", False),
+        ("python -c 'import sys' kernels/test_a.py", "kernels/test_a.py", False),
+        # a `-m pytest` *after* the `-c` code operand is that program's argv too.
+        ("python -c 'pass' -m pytest kernels/test_a.py", "kernels/test_a.py", False),
+        # ...but `-c` as the *script's* own arg still tethers the script.
+        ("python kernels/test_a.py -c config.py", "kernels/test_a.py", True),
+        # ...and pytest's own `-c <cfg>` (after `-m pytest`) is fine.
+        ("python -m pytest kernels/test_a.py -c setup.cfg", "kernels/test_a.py", True),
+    ],
+)
+def test_direct_runners(command, test_file, expected):
+    assert runs(command, test_file) is expected
+
+
+# --------------------------------------------------------------------------- #
+# Commands that must contribute no coverage
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    "command",
+    [
+        "export VLLM_WORKER_MULTIPROC_METHOD=spawn",
+        "pip install pytest-timeout pytest-forked",
+        "uv pip install -r requirements/test/cuda.in",
+        "bash tests/does_not_exist.sh",
+        "echo 'pytest kernels/test_a.py'",
+        # `pytest` as an argument to another command - `echo` runs, not pytest.
+        "echo pytest kernels/test_a.py",
+        'pytest -v -s "unbalanced',
+        # A `.sh` that is only an *argument* (Buildkite printing / copying a
+        # filename) is not a script invocation.
+        "echo tests/foo.sh",
+        "cp scripts/run_tests.sh /tmp/",
+        # A `find` whose output never reaches pytest is not a test selection.
+        "find kernels -name 'test_*.py'",
+        "find kernels -name 'test_*.py' | wc -l",
+        "find kernels -name 'test_*.py' -delete",
+        "find kernels -name 'test_*.py' | xargs echo pytest",
+        "find kernels -name 'test_*.py' -exec pytest {}",
+    ],
+)
+def test_non_test_commands_contribute_nothing(command):
+    assert _parse_command(command) == [] or not any(
+        sel.runs("kernels/test_a.py") for sel in _parse_command(command)
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Job-level yaml parsing
+# --------------------------------------------------------------------------- #
+
+
+def _write_area(tmp_path, monkeypatch, body: str):
+    mod = checker
+    area_dir = tmp_path / "test_areas"
+    area_dir.mkdir(exist_ok=True)
+    (area_dir / "synthetic.yaml").write_text(textwrap.dedent(body))
+    monkeypatch.setattr(mod, "TEST_AREAS_DIR", area_dir)
+    # Keep the real test-amd.yaml out of these isolated parsing tests.
+    monkeypatch.setattr(mod, "TEST_AMD_YAML", tmp_path / "no-such-test-amd.yaml")
+    # The yaml-error path reports a repo-relative name, so REPO_ROOT has to
+    # contain the synthetic area too.
+    monkeypatch.setattr(mod, "REPO_ROOT", tmp_path)
+    return load_selections()
+
+
+def test_plain_step_commands_are_collected(tmp_path, monkeypatch):
+    selections = _write_area(
+        tmp_path,
+        monkeypatch,
+        """
+        steps:
+        - label: ":nvidia: (H200) Kernels"
+          commands:
+          - pytest -v -s kernels
+        """,
+    )
+    assert is_tethered("kernels/test_a.py", selections)
+
+
+def test_singular_command_key_is_collected(tmp_path, monkeypatch):
+    selections = _write_area(
+        tmp_path,
+        monkeypatch,
+        """
+        steps:
+        - label: "Singular"
+          command: pytest -v -s kernels/test_a.py
+        """,
+    )
+    assert is_tethered("kernels/test_a.py", selections)
+
+
+def test_test_amd_yaml_coverage_counts_as_tethered(tmp_path, monkeypatch):
+    """A test wired only into the legacy test-amd.yaml is still run by CI."""
+    mod = checker
+    (tmp_path / "test_areas").mkdir()
+    amd_yaml = tmp_path / "test-amd.yaml"
+    amd_yaml.write_text(
+        textwrap.dedent(
+            """
+            steps:
+            - label: ":amd: (MI300) Quantization"
+              commands:
+              - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/
+              - pytest -v -s rocm/test_moe_weight_replay.py
+            """
+        )
+    )
+    monkeypatch.setattr(mod, "TEST_AREAS_DIR", tmp_path / "test_areas")
+    monkeypatch.setattr(mod, "TEST_AMD_YAML", amd_yaml)
+    monkeypatch.setattr(mod, "REPO_ROOT", tmp_path)
+    selections = load_selections()
+    assert is_tethered("rocm/test_moe_weight_replay.py", selections)
+    assert is_tethered("quantization/test_foo.py", selections)
+
+
+def test_missing_test_amd_yaml_is_not_fatal(tmp_path, monkeypatch):
+    """When the migration finally deletes test-amd.yaml, the checker still runs."""
+    mod = checker
+    (tmp_path / "test_areas").mkdir()
+    monkeypatch.setattr(mod, "TEST_AREAS_DIR", tmp_path / "test_areas")
+    monkeypatch.setattr(mod, "TEST_AMD_YAML", tmp_path / "gone.yaml")
+    monkeypatch.setattr(mod, "REPO_ROOT", tmp_path)
+    assert load_selections() == []
+
+
+def test_job_dirs_come_from_ci_config(tmp_path, monkeypatch):
+    """The pipeline directories are whatever `.buildkite/ci_config*.yaml` lists
+    under `job_dirs:` - so a test wired only into e.g. `hardware_tests/` counts,
+    and an unreadable config falls back to `test_areas/` alone."""
+    mod = checker
+    bk = tmp_path / ".buildkite"
+    (bk / "test_areas").mkdir(parents=True)
+    (bk / "hardware_tests").mkdir()
+    (bk / "ci_config.yaml").write_text("job_dirs:\n  - .buildkite/hardware_tests\n")
+    (bk / "ci_config_broken.yaml").write_text("job_dirs: [oops\n")  # ignored
+    (bk / "hardware_tests" / "cpu.yaml").write_text(
+        textwrap.dedent(
+            """
+            steps:
+            - label: "CPU"
+              commands:
+              - pytest -v -s v1/e2e/test_cpu_thing.py
+            """
+        )
+    )
+    monkeypatch.setattr(mod, "REPO_ROOT", tmp_path)
+    monkeypatch.setattr(mod, "TEST_AREAS_DIR", bk / "test_areas")
+    monkeypatch.setattr(mod, "TEST_AMD_YAML", tmp_path / "no-amd.yaml")
+
+    assert (bk / "hardware_tests") in checker._pipeline_job_dirs()
+    assert is_tethered("v1/e2e/test_cpu_thing.py", load_selections())
+
+
+def test_step_level_ignore_is_honored(tmp_path, monkeypatch):
+    selections = _write_area(
+        tmp_path,
+        monkeypatch,
+        """
+        steps:
+        - label: "Kernels"
+          commands:
+          - pytest -v -s kernels --ignore=kernels/test_a.py
+        """,
+    )
+    assert not is_tethered("kernels/test_a.py", selections)
+    assert is_tethered("kernels/test_b.py", selections)
+
+
+def test_amd_mirror_only_command_is_collected(tmp_path, monkeypatch):
+    """A mirror that overrides `commands` runs tests the base step never does."""
+    selections = _write_area(
+        tmp_path,
+        monkeypatch,
+        """
+        steps:
+        - label: ":nvidia: (H200) Base"
+          commands:
+          - pytest -v -s kernels/test_base.py
+          mirror:
+            amd:
+              label: ":amd: (MI300) Base"
+              commands:
+              - pytest -v -s kernels/test_amd_only.py
+        """,
+    )
+    assert is_tethered("kernels/test_base.py", selections)
+    assert is_tethered("kernels/test_amd_only.py", selections)
+
+
+def test_shell_script_continuation_vars_and_nesting(tmp_path, monkeypatch):
+    """The v1 P/D integration scripts name their test through a `\\`-continued
+    line, a literal shell var, and a nested `bash "$SCRIPT"`. Each indirection
+    has to be followed or the check falsely reports the test untethered."""
+    scripts = tmp_path / "tests" / "v1" / "kv_connector" / "nixl_integration"
+    scripts.mkdir(parents=True)
+    (scripts / "inner.sh").write_text(
+        "python3 -m pytest -s -v \\\n"
+        '  "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_inner.py"\n'
+    )
+    (scripts / "sweep.sh").write_text(
+        'SCRIPT="v1/kv_connector/nixl_integration/inner.sh"\n'
+        'IMPORT_CANARY="v1/kv_connector/nixl_integration/test_canary.py"\n'
+        'LITERAL="v1/kv_connector/nixl_integration/test_not_run.py"\n'
+        'python3 -m pytest -s -x "${IMPORT_CANARY}"\n'
+        # single-quoted: the shell passes a literal $LITERAL, collecting nothing
+        "python3 -m pytest -s -x '${LITERAL}'\n"
+        'if ! env FOO=1 bash "${SCRIPT}"; then exit 1; fi\n'
+    )
+    monkeypatch.setattr(checker, "REPO_ROOT", tmp_path)
+    monkeypatch.setattr(checker, "TESTS_DIR", tmp_path / "tests")
+
+    selections = _parse_command("bash v1/kv_connector/nixl_integration/sweep.sh")
+
+    def tethered(path):
+        return any(s.runs(path) for s in selections)
+
+    # `\`-continuation + $VAR-prefixed absolute path, resolved through a nested
+    # `env FOO=1 bash "$SCRIPT"`:
+    assert tethered("v1/kv_connector/nixl_integration/test_inner.py")
+    # literal `VAR=path` + double-quoted `"$VAR"` (the import-canary pattern):
+    assert tethered("v1/kv_connector/nixl_integration/test_canary.py")
+    # single-quoted `'$VAR'` is NOT expanded by the shell, so not tethered:
+    assert not tethered("v1/kv_connector/nixl_integration/test_not_run.py")
+
+
+def test_bash_dash_c_does_not_follow_a_script_operand(tmp_path, monkeypatch):
+    """`bash -c '<code>' foo.sh` runs the `-c` string; `foo.sh` is `$0` and is
+    never executed, so its contents must not count as coverage."""
+    (tmp_path / "tests").mkdir()
+    (tmp_path / "tests" / "decoy.sh").write_text("pytest -v -s lora/test_decoy.py\n")
+    monkeypatch.setattr(checker, "REPO_ROOT", tmp_path)
+    monkeypatch.setattr(checker, "TESTS_DIR", tmp_path / "tests")
+
+    with_c = _parse_command("bash -c 'echo hi' tests/decoy.sh")
+    assert not any(s.runs("lora/test_decoy.py") for s in with_c)
+    # ...but a plain `bash tests/decoy.sh` still follows it.
+    plain = _parse_command("bash tests/decoy.sh")
+    assert any(s.runs("lora/test_decoy.py") for s in plain)
+    # ...and a `-c` that is the *script's* own flag (not bash's) is not `bash -c`.
+    script_flag = _parse_command("bash tests/decoy.sh -c models.txt")
+    assert any(s.runs("lora/test_decoy.py") for s in script_flag)
+    # A `bash -c` trailing arg that looks like a command string is `$0`, not run.
+    cmd_string_arg = _parse_command("bash -c 'echo ok' 'pytest lora/test_x.py'")
+    assert not any(s.runs("lora/test_x.py") for s in cmd_string_arg)
+    # ...but the `-c` code string *itself* is parsed - it is what runs.
+    code_runs = _parse_command("bash -c 'pytest lora/test_x.py'")
+    assert any(s.runs("lora/test_x.py") for s in code_runs)
+    # a combined short-option group resolves the `-c` operand with `c` in any
+    # position (`-ec`, `-ce`), and its trailing args are still `$0`, not run.
+    for grp in ("-ec", "-ce", "-euxc"):
+        assert any(
+            s.runs("lora/test_x.py")
+            for s in _parse_command(f"bash {grp} 'pytest lora/test_x.py'")
+        )
+    assert not any(
+        s.runs("lora/test_x.py")
+        for s in _parse_command("bash -ce 'echo ok' 'pytest lora/test_x.py'")
+    )
+
+
+def test_unparsable_yaml_is_fatal(tmp_path, monkeypatch):
+    """Silently skipping a bad yaml would drop its coverage and report false
+    untethered files, so it must fail loudly instead."""
+    with pytest.raises(SystemExit):
+        _write_area(tmp_path, monkeypatch, "steps: [oops\n")
+
+
+def test_non_mapping_yaml_is_fatal(tmp_path, monkeypatch):
+    """A yaml that parses but isn't a pipeline mapping (a bare list or scalar)
+    must raise the actionable error, not crash on ``.get``."""
+    with pytest.raises(SystemExit):
+        _write_area(tmp_path, monkeypatch, "- just\n- a\n- list\n")
+
+
+# --------------------------------------------------------------------------- #
+# is_test_module / path normalization
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    ("path", "expected"),
+    [
+        ("tests/kernels/test_a.py", True),
+        ("tests/kernels/a_test.py", True),
+        ("tests/kernels/conftest.py", False),
+        ("tests/kernels/__init__.py", False),
+        ("tests/kernels/utils.py", False),
+        ("tests/kernels/test_a.txt", False),
+        ("vllm/kernels/test_a.py", False),
+    ],
+)
+def test_is_test_module(path, expected):
+    assert is_test_module(path) is expected
+
+
+@pytest.mark.parametrize(
+    ("raw", "expected"),
+    [
+        ("tests/kernels/test_a.py", "kernels/test_a.py"),
+        ("./kernels/test_a.py", "kernels/test_a.py"),
+        ("/vllm-workspace/tests/kernels/test_a.py", "kernels/test_a.py"),
+        ("'kernels/test_a.py'", "kernels/test_a.py"),
+        ("kernels/test_a.py::test_one", "kernels/test_a.py"),
+        ("kernels/", "kernels"),
+        ("${GIT_ROOT}/tests/kernels/test_a.py", "kernels/test_a.py"),
+        ("$GIT_ROOT/tests/kernels/test_a.py", "kernels/test_a.py"),
+    ],
+)
+def test_normalize_test_path(raw, expected):
+    assert normalize_test_path(raw) == expected
+
+
+# --------------------------------------------------------------------------- #
+# Whole-tree invariants
+# --------------------------------------------------------------------------- #
+
+
+def test_real_tree_parses_into_selections():
+    assert len(load_selections()) > 200
+    assert len(all_test_modules()) > 1000
+
+
+def test_no_selection_matches_everything():
+    """A path arg that normalizes to '' or '.' would tether every test in the
+    repo and silently disarm the gate - the worst possible misparse."""
+    for selection in load_selections():
+        if isinstance(selection, PytestSelection):
+            offenders = [
+                p
+                for p in selection.included_paths
+                if normalize_test_path(p) in ("", ".")
+            ]
+            assert not offenders, f"match-everything pytest arg: {offenders}"
+        else:
+            assert normalize_test_path(selection.root) not in ("", ".")
+
+
+def test_tethering_is_deterministic():
+    selections = load_selections()
+    modules = all_test_modules()[:200]
+    first = [is_tethered(m, selections) for m in modules]
+    assert first == [is_tethered(m, load_selections()) for m in modules]
+
+
+def test_allowlist_entries_are_well_formed():
+    allowlist = load_allowlist()
+    assert allowlist, "allowlist should not be empty while gaps remain"
+    assert all(is_test_module(path) for path in allowlist)
+    assert all(not path.startswith("/") for path in allowlist)
+
+
+def test_real_allowlist_entries_all_have_a_reason():
+    assert allowlist_entries_missing_reason() == []
+
+
+def test_bare_allowlist_entry_is_flagged(tmp_path, monkeypatch):
+    allowlist = tmp_path / "allowlist.txt"
+    allowlist.write_text(
+        "tests/a/test_ok.py  # torchrun multi-GPU, no job\n"
+        "tests/b/test_bare.py\n"
+        "  # a standalone comment line is fine\n"
+    )
+    monkeypatch.setattr(checker, "ALLOWLIST_PATH", allowlist)
+    assert allowlist_entries_missing_reason() == ["tests/b/test_bare.py"]
+
+
+def test_real_tree_has_no_unallowlisted_gaps():
+    """Every test module is either collected by a job or allowlisted.
+
+    Deliberately non-strict: a stale allowlist entry is advisory by design, so
+    asserting `strict=True` here would make unrelated allowlist cleanup block
+    CI - the exact behaviour the two-mode design avoids.
+    """
+    assert run_full_scan(load_selections(), load_allowlist(), strict=False) == 0
+
+
+# --------------------------------------------------------------------------- #
+# CLI behaviour
+# --------------------------------------------------------------------------- #
+
+
+@pytest.fixture
+def fake_repo(tmp_path, monkeypatch):
+    """Point the module at a throwaway tree so the changed-files path can be
+    exercised without writing test files into the real repo."""
+    mod = checker
+    (tmp_path / "tests" / "kernels").mkdir(parents=True)
+    (tmp_path / "tests" / "kernels" / "test_new.py").write_text("def test_x(): pass\n")
+    monkeypatch.setattr(mod, "REPO_ROOT", tmp_path)
+    monkeypatch.setattr(mod, "TESTS_DIR", tmp_path / "tests")
+    monkeypatch.setattr(mod, "ALLOWLIST_PATH", tmp_path / "allowlist.txt")
+    # The throwaway tree is not a git checkout, so the tracked-file inventory
+    # has to be stubbed rather than shelled out.
+    monkeypatch.setattr(mod, "all_test_modules", lambda: ["tests/kernels/test_new.py"])
+    return tmp_path
+
+
+def test_untethered_changed_file_fails(fake_repo, capsys):
+    rc = run_changed_files_check(["tests/kernels/test_new.py"], [], set())
+    assert rc == 1
+    assert "is not run by any Buildkite job" in capsys.readouterr().out
+
+
+def test_tethered_changed_file_passes(fake_repo):
+    selections = _parse_command("pytest -v -s kernels")
+    assert (
+        run_changed_files_check(["tests/kernels/test_new.py"], selections, set()) == 0
+    )
+
+
+def test_allowlisted_changed_file_passes(fake_repo):
+    allowlist = {"tests/kernels/test_new.py"}
+    assert run_changed_files_check(["tests/kernels/test_new.py"], [], allowlist) == 0
+
+
+def test_non_test_and_missing_paths_are_ignored(fake_repo):
+    paths = ["tests/kernels/conftest.py", "vllm/config.py", "tests/kernels/gone.py"]
+    assert run_changed_files_check(paths, [], set()) == 0
+
+
+def test_stale_allowlist_entries_are_advisory_but_reported(fake_repo, capsys):
+    """Dangling entries must not block an unrelated PR, but they have to be
+    printed - the hook sets `verbose: true` so a passing run still shows them."""
+    allowlist = {"tests/kernels/gone.py"}
+    rc = run_full_scan(_parse_command("pytest -v -s kernels"), allowlist, strict=False)
+    out = capsys.readouterr().out
+    assert rc == 0
+    assert "no longer exists" in out
+
+
+def test_stale_allowlist_entries_are_fatal_under_all(fake_repo, capsys):
+    allowlist = {"tests/kernels/gone.py"}
+    rc = run_full_scan(_parse_command("pytest -v -s kernels"), allowlist, strict=True)
+    assert rc == 1
+    assert "no longer exists" in capsys.readouterr().out
+
+
+def test_now_tethered_allowlist_entry_is_reported(fake_repo, capsys):
+    allowlist = {"tests/kernels/test_new.py"}
+    rc = run_full_scan(_parse_command("pytest -v -s kernels"), allowlist, strict=True)
+    assert rc == 1
+    assert "is now tethered" in capsys.readouterr().out
+
+
+def test_main_passes_a_tethered_file(monkeypatch):
+    """The CLI's normal pre-commit path: a tethered file exits 0.
+
+    `--all` is not asserted here on purpose - it implies `strict=True`, so a
+    stale allowlist entry would turn unrelated cleanup into a CI failure.
+    """
+    monkeypatch.setattr(
+        "sys.argv",
+        ["check_test_tethering.py", "tests/tools/test_check_test_tethering.py"],
+    )
+    assert main() == 0
+
+
+def test_main_flags_an_untethered_file(monkeypatch):
+    # Must be named like a test module, or the check correctly skips it.
+    orphan = REPO_ROOT / "tests" / "test_tethering_probe_orphan.py"
+    orphan.write_text("def test_x(): pass\n")
+    monkeypatch.setattr("sys.argv", ["check_test_tethering.py", str(orphan)])
+    try:
+        assert main() == 1
+    finally:
+        orphan.unlink()
+
+
+def test_to_repo_relative_accepts_absolute_paths():
+    absolute = str(REPO_ROOT / "tests" / "kernels" / "test_a.py")
+    assert _to_repo_relative(absolute) == "tests/kernels/test_a.py"
diff --git a/tests/tools/test_docker_build_metadata_args.py b/tests/tools/test_docker_build_metadata_args.py
index 8095ad34aa0e..127768afc733 100644
--- a/tests/tools/test_docker_build_metadata_args.py
+++ b/tests/tools/test_docker_build_metadata_args.py
@@ -155,10 +155,15 @@ def test_vllm_openai_image_embeds_metadata_contract() -> None:
 
 
 def test_rust_build_cache_excludes_git_metadata() -> None:
+    import torch
+
     from vllm.platforms import current_platform
 
     dockerfile_names = ["Dockerfile", "Dockerfile.cpu"]
-    if not current_platform.is_rocm():
+    # CPU jobs can reuse ROCm artifacts, which omit the XPU Dockerfile.
+    if not current_platform.is_rocm() and (
+        not current_platform.is_cpu() or torch.version.hip is None
+    ):
         dockerfile_names.append("Dockerfile.xpu")
     for name in dockerfile_names:
         dockerfile = (REPO_ROOT / "docker" / name).read_text()
diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py
index ba1f0ef8ddd8..4dee8fab6081 100644
--- a/tests/transformers_utils/test_config.py
+++ b/tests/transformers_utils/test_config.py
@@ -1,17 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This test file includes some cases where it is inappropriate to
+"""This test file includes some cases where it is inappropriate to
 only get the `eos_token_id` from the tokenizer as defined by
 `BaseRenderer.get_eos_token_id`.
 """
 
+import math
 from types import SimpleNamespace
 from typing import cast
 from unittest.mock import MagicMock, patch
 
 import pytest
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.model import ModelConfig
 from vllm.tokenizers import get_tokenizer
@@ -28,6 +28,7 @@
     Glm5NextTextConfig,
     Glm5NextVisionConfig,
 )
+from vllm.transformers_utils.configs.mistral import adapt_config_dict
 
 
 def test_patch_legacy_rope_type_preserves_nope_layers():
@@ -52,6 +53,82 @@ def test_patch_legacy_rope_type_preserves_nope_layers():
     }
 
 
+def test_patch_legacy_rope_type_normalizes_telechat3_yarn():
+    """TeleChat3's RoPE is YaRN with 0.07 in place of the usual 0.1.
+
+    Encoding that as a precomputed attention_factor keeps the config
+    plain YaRN, which Transformers and every "yarn" guard understand.
+    `mscale` cannot express it: Transformers only applies mscale when
+    mscale_all_dim is also truthy.
+    """
+    rope_parameters = {
+        "type": "telechat3-yarn",
+        "rope_type": "telechat3-yarn",
+        "factor": 4.0,
+        "original_max_position_embeddings": 8192,
+    }
+
+    patch_legacy_rope_type(rope_parameters)
+
+    assert rope_parameters == {
+        "rope_type": "yarn",
+        "factor": 4.0,
+        "original_max_position_embeddings": 8192,
+        "attention_factor": pytest.approx(0.07 * math.log(4.0) + 1.0),
+    }
+
+
+def test_mistral_yarn_apply_scale_false_disables_yarn_magnitude_scaling():
+    """`yarn.apply_scale: false` must reach the DeepSeek-style attentions.
+
+    Transformers spells it `attention_factor = 1.0`, which DeepseekV2Attention
+    and its siblings read to select `deepseek_llama_scaling` over
+    `deepseek_yarn`; without it Mistral-Large-3 runs with a spurious
+    yarn_get_mscale(factor)^2 attention scaling.
+    """
+    params = {
+        "dim": 7168,
+        "n_layers": 61,
+        "head_dim": 192,
+        "hidden_dim": 16384,
+        "n_heads": 128,
+        "n_kv_heads": 128,
+        "norm_eps": 1e-5,
+        "vocab_size": 131072,
+        "rope_theta": 10000.0,
+        "max_position_embeddings": 294912,
+        "q_lora_rank": 1536,
+        "kv_lora_rank": 512,
+        "qk_nope_head_dim": 128,
+        "qk_rope_head_dim": 64,
+        "v_head_dim": 128,
+        "moe": {
+            "num_experts": 128,
+            "num_experts_per_tok": 4,
+            "num_shared_experts": 1,
+            "expert_hidden_dim": 4096,
+            "first_k_dense_replace": 3,
+            "route_every_n": 1,
+            "routed_scale": 1.0,
+            "num_expert_groups": 1,
+            "num_expert_groups_per_tok": 1,
+        },
+        "llama_4_scaling": {"beta": 0.1, "original_max_position_embeddings": 8192},
+        "yarn": {
+            "alpha": 1,
+            "apply_scale": False,
+            "beta": 32,
+            "factor": 36,
+            "original_max_position_embeddings": 8192,
+        },
+    }
+
+    config = adapt_config_dict(params, defaults={})
+
+    assert config.architectures == ["MistralLarge3ForCausalLM"]
+    assert config.rope_parameters["attention_factor"] == 1.0
+
+
 def test_glm5_next_accepts_deepseek_sparse_attention_layers():
     layer_types = ["linear_attention", "deepseek_sparse_attention"]
 
@@ -131,6 +208,7 @@ def test_model_config_generation_fallback_forwards_code_revision():
             model="org/model",
             trust_remote_code=True,
             revision="model-pin",
+            _hf_config_revision=None,
             code_revision="code-pin",
             config_format="auto",
             hf_token=None,
@@ -146,7 +224,7 @@ def test_model_config_generation_fallback_forwards_code_revision():
         patch.object(
             config_module,
             "get_config",
-            return_value=PretrainedConfig(),
+            return_value=PreTrainedConfig(),
         ) as get_config,
     ):
         ModelConfig.try_get_generation_config(model_config)
@@ -192,7 +270,7 @@ def test_safetensors_metadata_of_repo_without_safetensors():
     ],
 )
 def test_mrope_num_dims(section_key, mrope_section, expected_num_dims):
-    config = PretrainedConfig()
+    config = PreTrainedConfig()
     config.rope_parameters = {"rope_type": "default", section_key: mrope_section}
 
     assert uses_mrope(config)
@@ -203,7 +281,7 @@ def test_mrope_num_dims(section_key, mrope_section, expected_num_dims):
 def test_mrope_num_dims_from_config_attribute(section_name):
     """Some configs expose the section as an attribute rather than under
     `rope_parameters`."""
-    config = PretrainedConfig()
+    config = PreTrainedConfig()
     setattr(config, section_name, [16, 16, 16, 16])
 
     assert uses_mrope(config)
@@ -212,7 +290,7 @@ def test_mrope_num_dims_from_config_attribute(section_name):
 
 def test_mrope_num_dims_from_nested_rope_parameters():
     """Sections nested by layer type must be found, not silently defaulted."""
-    config = PretrainedConfig()
+    config = PreTrainedConfig()
     config.rope_parameters = {
         "full_attention": {"mrope_section": [16, 16, 16, 16]},
         "linear_attention": {"rope_type": "default"},
@@ -223,4 +301,4 @@ def test_mrope_num_dims_from_nested_rope_parameters():
 
 
 def test_mrope_num_dims_without_mrope():
-    assert mrope_num_dims(PretrainedConfig()) == 0
+    assert mrope_num_dims(PreTrainedConfig()) == 0
diff --git a/tests/transformers_utils/test_config_parser_registry.py b/tests/transformers_utils/test_config_parser_registry.py
index 0931bd734f8f..41dcf4078243 100644
--- a/tests/transformers_utils/test_config_parser_registry.py
+++ b/tests/transformers_utils/test_config_parser_registry.py
@@ -4,7 +4,7 @@
 from pathlib import Path
 
 import pytest
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.transformers_utils.config import get_config_parser, register_config_parser
 from vllm.transformers_utils.config_parser_base import ConfigParserBase
@@ -19,7 +19,7 @@ def parse(
         revision: str | None = None,
         code_revision: str | None = None,
         **kwargs,
-    ) -> tuple[dict, PretrainedConfig]:
+    ) -> tuple[dict, PreTrainedConfig]:
         raise NotImplementedError
 
 
diff --git a/tests/transformers_utils/test_hf_overrides_model_type.py b/tests/transformers_utils/test_hf_overrides_model_type.py
index cce10d3b084e..71356edee092 100644
--- a/tests/transformers_utils/test_hf_overrides_model_type.py
+++ b/tests/transformers_utils/test_hf_overrides_model_type.py
@@ -5,12 +5,12 @@
 import json
 import tempfile
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.transformers_utils.config import _CONFIG_REGISTRY, get_config
 
 
-class _TestCustomConfig(PretrainedConfig):
+class _TestCustomConfig(PreTrainedConfig):
     model_type = "test_custom_model"
 
     def __init__(self, custom_attr=42, **kw):
@@ -23,7 +23,6 @@ def test_hf_overrides_model_type_returns_correct_config_class():
     checkpoint has a *different* model_type on disk, get_config() must return
     an instance of the registered config class — not the class that matches
     the on-disk model_type."""
-
     # Register the custom config
     _CONFIG_REGISTRY["test_custom_model"] = _TestCustomConfig
 
diff --git a/tests/transformers_utils/test_speculators_override.py b/tests/transformers_utils/test_speculators_override.py
index 35330f327f37..07f607cb88a7 100644
--- a/tests/transformers_utils/test_speculators_override.py
+++ b/tests/transformers_utils/test_speculators_override.py
@@ -40,7 +40,7 @@
 def _patched_get_config_dict(config_dict):
     """Return a patch context that bypasses HF Hub for get_config_dict."""
     return patch(
-        "vllm.transformers_utils.config.PretrainedConfig.get_config_dict",
+        "vllm.transformers_utils.config.PreTrainedConfig.get_config_dict",
         return_value=(config_dict, {}),
     )
 
diff --git a/tests/utils.py b/tests/utils.py
index 0c2cfeb6a0b9..399a741df011 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -87,25 +87,7 @@ def prewarm_hf_cache(assets: list[tuple[str, str]]) -> None:
             )
 
 
-if current_platform.is_rocm():
-    from amdsmi import (
-        amdsmi_get_gpu_vram_usage,
-        amdsmi_get_processor_handles,
-        amdsmi_init,
-        amdsmi_shut_down,
-    )
-
-    _amdsmi_lock = threading.Lock()
-
-    @contextmanager
-    def _nvml():
-        with _amdsmi_lock:
-            try:
-                amdsmi_init()
-                yield
-            finally:
-                amdsmi_shut_down()
-elif current_platform.is_cuda():
+if current_platform.is_cuda():
     from vllm.third_party.pynvml import (
         nvmlDeviceGetHandleByIndex,
         nvmlDeviceGetHandleByUUID,
@@ -245,7 +227,7 @@ def _create_cli_subcommand(self):
     def _start_server(
         self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
     ) -> None:
-        """Subclasses override this method to customize server process launch"""
+        """Subclasses override this method to customize server process launch."""
         raise NotImplementedError
 
     def _pre_download_model(self, model: str, args) -> None:
@@ -595,27 +577,20 @@ def _find_pgid_members(pgid: int) -> list[int]:
         return members
 
     def _get_gpu_memory_used(self) -> float | None:
-        """Get total GPU memory used across all visible devices in bytes."""
+        """Get device-wide usage across visible devices in bytes.
+
+        On ROCm this initializes HIP in the caller; subsequent multiprocessing
+        GPU workers must use spawn, as the ROCm server helpers already do.
+        """
         try:
             if current_platform.is_rocm():
-                with _nvml():
-                    handles = amdsmi_get_processor_handles()
-                    devices = get_physical_device_indices(
-                        list(range(current_platform.device_count()))
-                    )
-                    total_used_mib = 0
-                    for device in devices:
-                        handle = handles[device]
-                        vram_info = amdsmi_get_gpu_vram_usage(handle)
-                        total_used_mib += vram_info["vram_used"]
-                    # amdsmi reports VRAM in MiB; convert to bytes so this
-                    # matches the CUDA/nvml branch (already bytes) and the
-                    # byte-based target in _wait_for_gpu_memory_release. Without
-                    # this, that wait compares MiB against a ~2e9-byte target,
-                    # is always satisfied instantly, and returns "released to
-                    # 0.00 GB" while the previous server's VRAM is still
-                    # resident -- OOMing the next server's startup on ROCm.
-                    return total_used_mib * 1024 * 1024
+                # Use HIP logical devices. On MI355 DPX/NPS2, the primary
+                # AMD SMI device can expose whole-card memory counters.
+                total_used = 0
+                for i in range(current_platform.device_count()):
+                    free, total = torch.accelerator.get_memory_info(i)
+                    total_used += total - free
+                return float(total_used)
             elif current_platform.is_cuda():
                 with _nvml():
                     total_used = 0
@@ -653,9 +628,12 @@ def _wait_for_gpu_memory_release(
             # Can't query GPU memory - nothing to do
             return
 
-        # Allow up to 2 GiB overhead above baseline for driver/context state
-        # that may persist between server instances.
-        headroom_bytes = 2 * 1024 * 1024 * 1024
+        # Allow aggregate driver/context growth above the baseline.
+        headroom_bytes = (
+            4 * 1024 * 1024 * 1024
+            if current_platform.is_rocm()
+            else 2 * 1024 * 1024 * 1024
+        )
         target = baseline + headroom_bytes
 
         start = time.time()
@@ -703,7 +681,7 @@ def _wait_for_gpu_memory_release(
         )
 
     def _poll(self) -> int | None:
-        """Subclasses override this method to customize process polling"""
+        """Subclasses override this method to customize process polling."""
         return self.proc.poll()
 
     def _wait_for_server(self, *, url: str, timeout: float):
@@ -866,7 +844,7 @@ def _wait_for_gpu_memory_release(
 
 
 class RemoteOpenAIServerCustom(RemoteOpenAIServer):
-    """Launch test server with custom child process"""
+    """Launch test server with custom child process."""
 
     def _start_server(
         self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
@@ -1213,8 +1191,7 @@ def compare_two_settings(
     include_seeded_sampling: bool = True,
     force_v1_runner: bool = False,
 ) -> None:
-    """
-    Launch API server with two different sets of arguments/environments
+    """Launch API server with two different sets of arguments/environments
     and compare the results of the API calls.
 
     Args:
@@ -1228,8 +1205,8 @@ def compare_two_settings(
         force_v1_runner: Whether to pin all compared settings to the v1 model
             runner to avoid mixing model runner differences into correctness
             tests.
-    """
 
+    """
     compare_all_settings(
         model,
         [arg1, arg2],
@@ -1251,9 +1228,9 @@ def compare_all_settings(
     include_seeded_sampling: bool = True,
     force_v1_runner: bool = False,
 ) -> None:
-    """
-    Launch API server with several different sets of arguments/environments
+    """Launch API server with several different sets of arguments/environments
     and compare the results of the API calls with the first set of arguments.
+
     Args:
         model: The model to test.
         all_args: A list of argument lists to pass to the API server.
@@ -1263,8 +1240,8 @@ def compare_all_settings(
         force_v1_runner: Whether to pin all compared settings to the v1 model
             runner to avoid mixing model runner differences into correctness
             tests.
-    """
 
+    """
     if force_v1_runner:
         all_envs = [
             {"VLLM_USE_V2_MODEL_RUNNER": "0", **(env or {})} for env in all_envs
@@ -1394,6 +1371,7 @@ def ensure_current_vllm_config():
         with ensure_current_vllm_config():
             init_distributed_environment(...)
             ensure_model_parallel_initialized(...)
+
     """
     from vllm.config import (
         VllmConfig,
@@ -1416,6 +1394,8 @@ def init_test_distributed_environment(
     rank: int,
     distributed_init_port: str,
     local_rank: int = -1,
+    data_parallel_size: int = 1,
+    data_parallel_master_port: int | None = None,
 ) -> None:
     # Note: This function is often called from Ray worker processes, so we
     # can't rely on pytest fixtures to set the config. We check if the config
@@ -1425,8 +1405,34 @@ def init_test_distributed_environment(
         get_current_vllm_config_or_none,
         set_current_vllm_config,
     )
+    from vllm.platforms import current_platform
 
     distributed_init_method = f"tcp://localhost:{distributed_init_port}"
+    backend = current_platform.dist_backend
+
+    if data_parallel_size > 1:
+        # For DP we need to set a common DP master port
+        from vllm.config.parallel import ParallelConfig
+
+        assert data_parallel_master_port is not None, (
+            "data_parallel_master_port is required when data_parallel_size > 1"
+        )
+        tp_pp_world = tp_size * pp_size
+        parallel_config = ParallelConfig(
+            data_parallel_size=data_parallel_size,
+            data_parallel_rank=rank // tp_pp_world,
+            _data_parallel_master_port_list=[int(data_parallel_master_port)],
+        )
+        with set_current_vllm_config(VllmConfig(parallel_config=parallel_config)):
+            init_distributed_environment(
+                world_size=tp_pp_world,
+                rank=rank % tp_pp_world,
+                distributed_init_method=distributed_init_method,
+                local_rank=local_rank if local_rank >= 0 else rank,
+                backend=backend,
+            )
+            ensure_model_parallel_initialized(tp_size, pp_size)
+        return
 
     if get_current_vllm_config_or_none() is not None:
         # Config already set, use it directly
@@ -1435,6 +1441,7 @@ def init_test_distributed_environment(
             rank=rank,
             distributed_init_method=distributed_init_method,
             local_rank=local_rank,
+            backend=backend,
         )
         ensure_model_parallel_initialized(tp_size, pp_size)
     else:
@@ -1445,6 +1452,7 @@ def init_test_distributed_environment(
                 rank=rank,
                 distributed_init_method=distributed_init_method,
                 local_rank=local_rank,
+                backend=backend,
             )
             ensure_model_parallel_initialized(tp_size, pp_size)
 
@@ -1454,6 +1462,7 @@ def multi_process_parallel(
     tp_size: int,
     pp_size: int,
     test_target: Any,
+    data_parallel_size: int = 1,
 ) -> None:
     import ray
 
@@ -1474,18 +1483,34 @@ def multi_process_parallel(
     )
 
     distributed_init_port = get_open_port()
+    # Separate port for the DP master group; only used when data_parallel_size > 1.
+    data_parallel_master_port = get_open_port() if data_parallel_size > 1 else None
+    world_size = data_parallel_size * tp_size * pp_size
     try:
         refs = []
-        for rank in range(tp_size * pp_size):
-            refs.append(
-                test_target.remote(
-                    monkeypatch,
-                    tp_size,
-                    pp_size,
-                    rank,
-                    distributed_init_port,
-                ),
-            )
+        for rank in range(world_size):
+            if data_parallel_size > 1:
+                refs.append(
+                    test_target.remote(
+                        monkeypatch,
+                        tp_size,
+                        pp_size,
+                        rank,
+                        distributed_init_port,
+                        data_parallel_size,
+                        data_parallel_master_port,
+                    ),
+                )
+            else:
+                refs.append(
+                    test_target.remote(
+                        monkeypatch,
+                        tp_size,
+                        pp_size,
+                        rank,
+                        distributed_init_port,
+                    ),
+                )
         ray.get(refs)
     finally:
         ray.shutdown()
@@ -1533,8 +1558,7 @@ def assert_rocm_custom_allreduce_backend_state_on_worker(
 
 @contextmanager
 def error_on_warning(category: type[Warning] = Warning):
-    """
-    Within the scope of this context manager, tests will fail if any warning
+    """Within the scope of this context manager, tests will fail if any warning
     of the given category is emitted.
     """
     with warnings.catch_warnings():
@@ -1570,13 +1594,17 @@ def record_gpu_memory_usage_stats(
     *,
     devices: list[int],
 ) -> dict[int, tuple[float, float]]:
+    """Return device-wide used/total GiB; ROCm IDs are HIP logical devices.
+
+    ROCm queries initialize HIP in the caller and include other processes'
+    allocations on the same memory partition.
+    """
     output: dict[int, tuple[float, float]] = {}
     for device in devices:
         if current_platform.is_rocm():
-            dev_handle = amdsmi_get_processor_handles()[device]
-            mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
-            gb_used = mem_info["vram_used"] / 2**10
-            gb_total = mem_info["vram_total"] / 2**10
+            free, total = torch.accelerator.get_memory_info(device)
+            gb_used = (total - free) / 2**30
+            gb_total = total / 2**30
         elif current_platform.is_xpu():
             # nvml/amdsmi are unavailable on XPU. Query device memory through
             # torch.accelerator.get_memory_info, which the XPU platform patches
@@ -1604,7 +1632,9 @@ def wait_for_gpu_memory_to_clear(
     poll_interval_s: float = 5,
 ) -> None:
     assert threshold_bytes is not None or threshold_ratio is not None
-    devices = get_physical_device_indices(devices)
+    # HIP already applies device visibility; keep logical IDs and threshold keys.
+    if not current_platform.is_rocm():
+        devices = get_physical_device_indices(devices)
     if isinstance(threshold_bytes, int):
         threshold_bytes = {device: threshold_bytes for device in devices}
     elif isinstance(threshold_bytes, dict):
@@ -1626,8 +1656,6 @@ def wait_for_gpu_memory_to_clear(
                 threshold_bytes.get(device, 0), min_threshold_b if ratio < 0.05 else 0
             )
 
-    # Use nvml instead of pytorch to reduce measurement error from torch cuda
-    # context.
     start_time = time.time()
     stable_since: float | None = None
     stable_used_bytes: dict[int, int] | None = None
@@ -1987,6 +2015,7 @@ def create_new_process_for_each_test(
 
     Returns:
         A decorator to run test functions in separate processes.
+
     """
     if method is None:
         method = "spawn" if requires_spawn_multiprocessing() else "fork"
@@ -2000,8 +2029,7 @@ def create_new_process_for_each_test(
 
 
 def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
-    """
-    Get a pytest mark, which skips the test if the GPU doesn't meet
+    """Get a pytest mark, which skips the test if the GPU doesn't meet
     a minimum memory requirement in GB.
 
     This can be leveraged via `@large_gpu_test` to skip tests in environments
@@ -2036,8 +2064,7 @@ def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
 
 
 def large_gpu_test(*, min_gb: int):
-    """
-    Decorate a test to be skipped if no GPU is available or it does not have
+    """Decorate a test to be skipped if no GPU is available or it does not have
     sufficient memory.
 
     Currently, the CI machine uses L4 GPU which has 24 GB VRAM.
@@ -2062,9 +2089,7 @@ def multi_gpu_marks(*, num_gpus: int):
 
 
 def multi_gpu_test(*, num_gpus: int):
-    """
-    Decorate a test to be run only when multiple GPUs are available.
-    """
+    """Decorate a test to be run only when multiple GPUs are available."""
     marks = multi_gpu_marks(num_gpus=num_gpus)
 
     def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
@@ -2078,13 +2103,13 @@ def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
 
 
 def gpu_tier_mark(*, min_gpus: int = 1, max_gpus: int | None = None):
-    """
-    Mark a test to only run when the GPU count falls within [min_gpus, max_gpus].
+    """Mark a test to only run when the GPU count falls within [min_gpus, max_gpus].
 
     Examples:
         @gpu_tier_mark(min_gpus=2)          # only on multi-GPU
         @gpu_tier_mark(max_gpus=1)          # only on single-GPU
         @gpu_tier_mark(min_gpus=2, max_gpus=4)  # 2-4 GPUs only
+
     """
     gpu_count = current_platform.device_count()
     marks = []
@@ -2152,8 +2177,8 @@ async def completions_with_server_args(
 
     Returns:
       OpenAI Completion instance
-    """
 
+    """
     if isinstance(max_tokens, int):
         max_tokens = [max_tokens] * len(prompts)
 
@@ -2211,9 +2236,7 @@ def get_client_text_logprob_generations(
 
 
 def has_module_attribute(module_name, attribute_name):
-    """
-    Helper function to check if a module has a specific attribute.
-    """
+    """Helper function to check if a module has a specific attribute."""
     try:
         module = importlib.import_module(module_name)
         return hasattr(module, attribute_name)
@@ -2285,14 +2308,15 @@ def _native_plain(
 
 
 def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)):
-    """
-    Generate prompts which a bunch of assignments,
+    """Generate prompts which a bunch of assignments,
     then asking for the value of one of them.
     The prompt is just under 10k tokens; sliding window is 4k
     so the answer is outside sliding window, but should still be correct.
+
     Args:
         batch_size: number of prompts to generate
         ln_range: an argument to control the length of the prompt
+
     """
     prompts: list[str] = []
     answer: list[int] = []
@@ -2332,8 +2356,7 @@ def check_answers(
 
 
 def flat_product(*iterables: Iterable[Any]):
-    """
-    Flatten lists of tuples of the cartesian product.
+    """Flatten lists of tuples of the cartesian product.
     Useful when we want to avoid nested tuples to allow
     test params to be unpacked directly from the decorator.
 
@@ -2345,6 +2368,7 @@ def flat_product(*iterables: Iterable[Any]):
       (3, 4, "a"),
       (3, 4, "b"),
     ]
+
     """
     for element in itertools.product(*iterables):
         normalized = (e if isinstance(e, tuple) else (e,) for e in element)
@@ -2352,8 +2376,7 @@ def flat_product(*iterables: Iterable[Any]):
 
 
 class TestFP8Layer(torch.nn.Module):
-    """
-    Test helper for FP8 linear operations. Creates random weights and scales
+    """Test helper for FP8 linear operations. Creates random weights and scales
     based on quantization configuration.
 
     Args:
@@ -2362,6 +2385,7 @@ class TestFP8Layer(torch.nn.Module):
         weight_quant_key: Weight quantization configuration.
         out_dtype: Output dtype. Defaults to current default dtype.
         force_kernel: Optional kernel to force use of specific implementation.
+
     """
 
     def __init__(
diff --git a/tests/utils_/__init__.py b/tests/utils_/__init__.py
index e6b4c3f6364c..57451f3bd857 100644
--- a/tests/utils_/__init__.py
+++ b/tests/utils_/__init__.py
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This module is named `utils_` instead of `utils` to avoid obscuring
+"""This module is named `utils_` instead of `utils` to avoid obscuring
 `tests/utils.py`.
 """
diff --git a/tests/utils_/test_argparse_utils.py b/tests/utils_/test_argparse_utils.py
index baf57d994ff9..89e0f9977d0c 100644
--- a/tests/utils_/test_argparse_utils.py
+++ b/tests/utils_/test_argparse_utils.py
@@ -524,3 +524,20 @@ def test_flat_product():
         (3, 4, "a", 5, 6),
         (3, 4, "b", 5, 6),
     ]
+
+
+def test_group_description_is_summary_only():
+    """Group descriptions are config docstrings, too long for the terminal."""
+    parser = FlexibleArgumentParser()
+    group = parser.add_argument_group(
+        title="MyConfig",
+        description="Summary line.\n\nDetails which only belong in the docs.",
+    )
+    group.add_argument("--my-arg")
+
+    assert "Summary line." in parser.format_help()
+    assert "only belong in the docs" not in parser.format_help()
+
+    parser._search_keyword = "myconfig"
+    assert "Summary line." in parser.format_help()
+    assert "only belong in the docs" not in parser.format_help()
diff --git a/tests/utils_/test_cpu_resource_utils.py b/tests/utils_/test_cpu_resource_utils.py
new file mode 100644
index 000000000000..ff314d2231d2
--- /dev/null
+++ b/tests/utils_/test_cpu_resource_utils.py
@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cgroup memory readers and SHM allocation preflight checks."""
+
+import sys
+from io import StringIO
+from unittest import mock
+
+import pytest
+
+from vllm.utils import cpu_resource_utils as cru
+from vllm.utils.mem_constants import GiB_bytes
+
+_V2_LIMIT_PATH = "/sys/fs/cgroup/memory.max"
+_V2_USAGE_PATH = "/sys/fs/cgroup/memory.current"
+_V1_LIMIT_PATH = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
+_V1_USAGE_PATH = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
+
+
+@pytest.fixture(autouse=True)
+def _clear_cgroup_cache():
+    cru.get_cgroup_memory_limit.cache_clear()
+    yield
+    cru.get_cgroup_memory_limit.cache_clear()
+
+
+def _stub_files(monkeypatch, files: dict):
+    """Stub ``open()`` for a fixed set of paths; ``None`` -> OSError."""
+    real_open = open
+
+    def fake_open(path, *args, **kwargs):
+        if path in files:
+            content = files[path]
+            if content is None:
+                raise OSError(f"no such file: {path}")
+            return StringIO(content)
+        return real_open(path, *args, **kwargs)
+
+    monkeypatch.setattr("builtins.open", fake_open)
+
+
+def test_cgroup_v2_limit_and_usage(monkeypatch):
+    monkeypatch.setattr(sys, "platform", "linux")
+    _stub_files(
+        monkeypatch,
+        {
+            _V2_LIMIT_PATH: f"{20 * GiB_bytes}\n",
+            _V2_USAGE_PATH: f"{5 * GiB_bytes}\n",
+        },
+    )
+
+    assert cru.get_cgroup_memory_limit() == 20 * GiB_bytes
+    assert cru.get_cgroup_memory_usage() == 5 * GiB_bytes
+
+
+def test_cgroup_v2_unlimited_falls_back_to_v1(monkeypatch):
+    monkeypatch.setattr(sys, "platform", "linux")
+    _stub_files(
+        monkeypatch,
+        {
+            _V2_LIMIT_PATH: "max\n",
+            _V1_LIMIT_PATH: f"{8 * GiB_bytes}\n",
+            _V1_USAGE_PATH: f"{1 * GiB_bytes}\n",
+        },
+    )
+
+    assert cru.get_cgroup_memory_limit() == 8 * GiB_bytes
+    assert cru.get_cgroup_memory_usage() == 1 * GiB_bytes
+
+
+def test_cgroup_v1_limit_and_usage(monkeypatch):
+    monkeypatch.setattr(sys, "platform", "linux")
+    _stub_files(
+        monkeypatch,
+        {
+            _V2_LIMIT_PATH: None,
+            _V1_LIMIT_PATH: f"{8 * GiB_bytes}\n",
+            _V1_USAGE_PATH: f"{2 * GiB_bytes}\n",
+        },
+    )
+
+    assert cru.get_cgroup_memory_limit() == 8 * GiB_bytes
+    assert cru.get_cgroup_memory_usage() == 2 * GiB_bytes
+
+
+def test_cgroup_v1_unlimited_sentinel_is_ignored(monkeypatch):
+    """An unlimited cgroup v1 sentinel must not be treated as a real limit."""
+    monkeypatch.setattr(sys, "platform", "linux")
+    _stub_files(
+        monkeypatch,
+        {_V2_LIMIT_PATH: None, _V1_LIMIT_PATH: f"{(1 << 63) - 1}\n"},
+    )
+
+    assert cru.get_cgroup_memory_limit() is None
+    assert cru.get_cgroup_memory_usage() is None
+
+
+def test_cgroup_no_limit_files_present(monkeypatch):
+    monkeypatch.setattr(sys, "platform", "linux")
+    _stub_files(monkeypatch, {_V2_LIMIT_PATH: None, _V1_LIMIT_PATH: None})
+
+    assert cru.get_cgroup_memory_limit() is None
+    assert cru.get_cgroup_memory_usage() is None
+
+
+def test_cgroup_skipped_on_non_linux(monkeypatch):
+    monkeypatch.setattr(sys, "platform", "darwin")
+    _stub_files(monkeypatch, {_V2_LIMIT_PATH: f"{1 * GiB_bytes}\n"})
+
+    assert cru.get_cgroup_memory_limit() is None
+    assert cru.get_cgroup_memory_usage() is None
+
+
+def test_cgroup_limit_is_cached_until_explicitly_cleared(monkeypatch):
+    monkeypatch.setattr(sys, "platform", "linux")
+    files = {
+        _V2_LIMIT_PATH: f"{1 * GiB_bytes}\n",
+        _V2_USAGE_PATH: f"{512 << 20}\n",
+    }
+    _stub_files(monkeypatch, files)
+
+    assert cru.get_cgroup_memory_limit() == 1 * GiB_bytes
+    files[_V2_LIMIT_PATH] = f"{2 * GiB_bytes}\n"
+    assert cru.get_cgroup_memory_limit() == 1 * GiB_bytes
+
+    cru.get_cgroup_memory_limit.cache_clear()
+    assert cru.get_cgroup_memory_limit() == 2 * GiB_bytes
+
+
+def test_cgroup_usage_is_read_without_cache(monkeypatch):
+    monkeypatch.setattr(sys, "platform", "linux")
+    files = {
+        _V2_LIMIT_PATH: f"{1 * GiB_bytes}\n",
+        _V2_USAGE_PATH: f"{512 << 20}\n",
+    }
+    _stub_files(monkeypatch, files)
+
+    assert cru.get_cgroup_memory_usage() == 512 << 20
+    files[_V2_USAGE_PATH] = f"{768 << 20}\n"
+    assert cru.get_cgroup_memory_usage() == 768 << 20
+
+
+def test_check_cgroup_memory_available_warns_on_low_headroom(monkeypatch):
+    monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: 1 << 30)
+    monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: 512 << 20)
+
+    with (
+        mock.patch.object(cru.logger, "debug") as log_debug,
+        mock.patch.object(cru.logger, "warning") as log_warning,
+    ):
+        cru.check_cgroup_memory_available(600 << 20, "mmap")
+
+    log_debug.assert_not_called()
+    log_warning.assert_called_once()
+    assert "current headroom is below" in log_warning.call_args.args[-1]
+
+
+def test_check_cgroup_memory_available_logs_success_at_debug(monkeypatch):
+    monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: 1 << 30)
+    monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: 512 << 20)
+
+    with (
+        mock.patch.object(cru.logger, "debug") as log_debug,
+        mock.patch.object(cru.logger, "warning") as log_warning,
+    ):
+        cru.check_cgroup_memory_available(256 << 20, "mmap")
+
+    log_debug.assert_called_once()
+    log_warning.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("limit", "usage"),
+    [(None, 512 << 20), (1 << 30, None)],
+)
+def test_check_cgroup_memory_available_skips_without_snapshot(
+    monkeypatch, limit, usage
+):
+    monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: limit)
+    monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: usage)
+
+    with (
+        mock.patch.object(cru.logger, "debug") as log_debug,
+        mock.patch.object(cru.logger, "warning") as log_warning,
+    ):
+        cru.check_cgroup_memory_available(1 << 60, "mmap")
+
+    log_debug.assert_not_called()
+    log_warning.assert_not_called()
diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py
index dc8b12280010..10b550c0a174 100644
--- a/tests/utils_/test_gpu_sync_debug.py
+++ b/tests/utils_/test_gpu_sync_debug.py
@@ -123,6 +123,24 @@ def test_suppressing_works_while_compiling(monkeypatch):
     with_gpu_sync_check(suppressed)()
 
 
+@create_new_process_for_each_test()
+def test_copy_checker_is_traceable_while_compiling(monkeypatch):
+    """The copy wrapper must bypass ContextVars while Dynamo is tracing."""
+    monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
+    monkeypatch.setattr(gsd, "_sync_check_enabled", True)
+    gsd._install_copy_checkers()
+
+    def convert_like(x, like):
+        return x.to(device=like.device, dtype=like.dtype)
+
+    compiled = torch.compile(convert_like, fullgraph=True)
+    result = compiled(
+        torch.ones(4, device="cuda", dtype=torch.float64),
+        torch.ones(4, device="cuda", dtype=torch.float32),
+    )
+    assert result.dtype == torch.float32
+
+
 @create_new_process_for_each_test()
 def test_sync_debug_mode_restored_after_checked_call(monkeypatch):
     """The mode is armed only for the duration of a checked call. Leaving it
diff --git a/tests/utils_/test_jsontree.py b/tests/utils_/test_jsontree.py
index 0af2751b2638..49e75c3477ef 100644
--- a/tests/utils_/test_jsontree.py
+++ b/tests/utils_/test_jsontree.py
@@ -5,7 +5,6 @@
 
 def test_json_count_leaves():
     """Test json_count_leaves function from jsontree utility."""
-
     # Single leaf values
     assert json_count_leaves(42) == 1
     assert json_count_leaves("hello") == 1
diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py
index 6901a162a407..79dd3e320ae6 100644
--- a/tests/utils_/test_torch_utils.py
+++ b/tests/utils_/test_torch_utils.py
@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import numpy as np
 import pytest
 import torch
 
 from vllm.utils.torch_utils import (
     OMP_NUM_THREADS_SET_BY_VLLM,
+    async_tensor_h2d,
     available_cpu_count,
     common_broadcastable_dtype,
     current_stream,
@@ -192,3 +194,26 @@ def test_runtime_threads_override_vllm_set_omp_num_threads(
     torch.set_num_threads(3)
     set_torch_threads_for_runtime()
     assert torch.get_num_threads() == 1
+
+
+@pytest.mark.parametrize(
+    "device", ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]
+)
+def test_async_tensor_h2d_staging(device):
+    """Staging must preserve values while normalizing dtype and layout."""
+    # Non-contiguous source with a dtype conversion: would previously sync.
+    src = torch.arange(12, dtype=torch.int64).reshape(3, 4).T
+    result = async_tensor_h2d(src, device=device, dtype=torch.int32)
+    assert result.dtype == torch.int32
+    assert result.is_contiguous()
+    assert torch.equal(result.cpu(), src.to(torch.int32))
+
+    # numpy and list sources still work.
+    assert torch.equal(
+        async_tensor_h2d(np.arange(4, dtype=np.int64), device=device),
+        torch.arange(4, device=device),
+    )
+    assert torch.equal(
+        async_tensor_h2d([1, 2, 3], device=device, dtype=torch.int32),
+        torch.tensor([1, 2, 3], dtype=torch.int32, device=device),
+    )
diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py
index 6609705c2202..b02c85f68202 100644
--- a/tests/v1/attention/test_attention_backends.py
+++ b/tests/v1/attention/test_attention_backends.py
@@ -151,6 +151,7 @@ def create_and_prepopulate_kv_cache(
     Returns:
         A 4D tensor in logical ``(num_blocks, num_kv_heads, block_size,
         2 * head_size)`` order with strides determined by ``layout``.
+
     """
     batch_size = len(k_contexts)
     seq_lens = common_attn_metadata.seq_lens.cpu()
@@ -242,14 +243,19 @@ def create_and_prepopulate_kv_cache(
 class MockAttentionLayer:
     """A mock attention layer for testing."""
 
-    def __init__(self, device: torch.device):
+    def __init__(
+        self,
+        device: torch.device,
+        k_scale: float = 1.0,
+        v_scale: float = 1.0,
+    ):
         self._q_scale = torch.tensor(1.0, device=device)
-        self._k_scale = torch.tensor(1.0, device=device)
-        self._v_scale = torch.tensor(1.0, device=device)
+        self._k_scale = torch.tensor(k_scale, device=device)
+        self._v_scale = torch.tensor(v_scale, device=device)
         # Add float versions for flashinfer
         self._q_scale_float = 1.0
-        self._k_scale_float = 1.0
-        self._v_scale_float = 1.0
+        self._k_scale_float = k_scale
+        self._v_scale_float = v_scale
 
 
 def _clone_kv_cache_in_layout(
@@ -284,9 +290,10 @@ def run_attention_backend(
     kv_cache_dtype: str = "auto",
     sinks: torch.Tensor | None = None,
     use_cuda_graph: bool = False,
+    layer_k_scale: float = 1.0,
+    layer_v_scale: float = 1.0,
 ) -> torch.Tensor:
     """Run attention computation using the specified backend's AttentionImpl."""
-
     use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
     if backend == "FLEX_ATTENTION_SLOW":
         use_direct_block_mask = False
@@ -356,7 +363,7 @@ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
         )
 
     # Create mock layer and output buffer
-    mock_layer = MockAttentionLayer(device)
+    mock_layer = MockAttentionLayer(device, layer_k_scale, layer_v_scale)
     output = torch.empty_like(query)
 
     if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:
@@ -386,9 +393,11 @@ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
             )
         graph.replay()
     else:
-        output = impl.forward(
+        backend_output = impl.forward(
             mock_layer, query, key, value, kv_cache, attn_metadata, output=output
         )
+        if backend_output is not None:
+            output = backend_output
 
     return output
 
@@ -413,9 +422,10 @@ def _test_backend_correctness(
     model_dtype: torch.dtype | None = None,
     max_num_seqs: int | None = None,
     max_num_batched_tokens: int | None = None,
+    layer_k_scale: float = 1.0,
+    layer_v_scale: float = 1.0,
 ):
-    """
-    Test that all backends produce similar outputs to a reference implementation
+    """Test that all backends produce similar outputs to a reference implementation
     using FlexAttention or an explicit attention-sink reference.
 
     This test works by:
@@ -688,6 +698,8 @@ def _test_backend_correctness(
             kv_cache_dtype=kv_cache_dtype,
             sinks=sinks,
             use_cuda_graph=use_cuda_graph,
+            layer_k_scale=layer_k_scale,
+            layer_v_scale=layer_v_scale,
         )
 
         # Check shape and dtype consistency
@@ -889,7 +901,7 @@ def test_flashinfer_xqa_query_lens_preserve_cudagraph_padding():
 
     device = torch.device("cpu")
     builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder)
-    builder.use_dedicated_xqa = True
+    builder.use_xqa = True
     qo_indptr = torch.tensor([0, 3, 9, 15, 15], dtype=torch.int32, device=device)
 
     q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens(
@@ -905,6 +917,95 @@ def test_flashinfer_xqa_query_lens_preserve_cudagraph_padding():
     assert q_cu_seq_lens.tolist() == [0, 3, 9, 15, 15]
 
 
+@pytest.mark.skipif(
+    AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
+    reason="FlashInfer is not available.",
+)
+def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypatch):
+    """Non-ragged XQA keeps one query/output row per padded request."""
+    import unittest.mock
+
+    from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+
+    impl = object.__new__(flashinfer_backend.FlashInferImpl)
+    impl.scale = 1.0
+    impl.bmm1_scale = None
+    impl.bmm2_scale = None
+    impl.kv_cache_dtype = "auto"
+    impl.is_kvcache_nvfp4 = False
+    impl.head_size = 16
+    impl.dcp_world_size = 1
+    impl.o_sf_scale = None
+    impl.window_left = -1
+    impl.sinks = None
+    impl.cache_config = unittest.mock.Mock()
+    impl.cache_config.get_resolved_kv_cache_layout.return_value = KVCacheLayout.LBHNC
+
+    layer = unittest.mock.Mock(
+        _q_scale=torch.tensor(1.0),
+        _q_scale_float=1.0,
+        _k_scale_float=1.0,
+        _v_scale_float=1.0,
+    )
+    decode = flashinfer_backend.FlashInferTrtllmAPIDecode(
+        kernel=flashinfer_backend.FlashInferDecodeKernel.XQA,
+        block_tables=torch.zeros((4, 1), dtype=torch.int32),
+        seq_lens=torch.tensor([8, 8, 0, 0], dtype=torch.int32),
+        max_seq_len=8,
+        q_len_per_req=1,
+    )
+    attn_metadata = flashinfer_backend.FlashInferMetadata(
+        num_actual_tokens=2,
+        slot_mapping=torch.empty(0, dtype=torch.int64),
+        q_data_type_prefill=torch.bfloat16,
+        q_data_type_decode=torch.bfloat16,
+        num_decodes=4,
+        num_decode_tokens=2,
+        num_prefills=0,
+        num_prefill_tokens=0,
+        causal=True,
+        prefill=None,
+        decode=decode,
+        use_cascade=False,
+        cascade_wrapper=None,
+    )
+
+    seen_shapes = {}
+
+    def mock_xqa(**kwargs):
+        seen_shapes["query"] = kwargs["query"].shape
+        seen_shapes["out"] = kwargs["out"].shape
+        seen_shapes["block_tables"] = kwargs["block_tables"].shape
+
+    monkeypatch.setattr(
+        flashinfer_backend,
+        "_get_trtllm_workspace_buffer",
+        lambda: torch.empty(1, dtype=torch.uint8),
+    )
+    monkeypatch.setattr(
+        flashinfer_backend, "flashinfer_xqa_batch_decode_with_kv_cache", mock_xqa
+    )
+
+    query = torch.zeros((4, 1, 16), dtype=torch.bfloat16)
+    output = torch.empty_like(query)
+    result = impl.forward(
+        layer,
+        query,
+        torch.empty_like(query),
+        torch.empty_like(query),
+        torch.zeros((1, 1, 1, 32), dtype=torch.bfloat16),
+        attn_metadata,
+        output,
+    )
+
+    assert result is output
+    assert seen_shapes == {
+        "query": torch.Size([4, 1, 16]),
+        "out": torch.Size([4, 1, 16]),
+        "block_tables": torch.Size([4, 1]),
+    }
+
+
 @pytest.mark.skipif(
     AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
     reason="FlashInfer is not available.",
@@ -913,7 +1014,7 @@ def test_flashinfer_xqa_query_lens_require_exact_uniform_product():
     from vllm.v1.attention.backends import flashinfer as flashinfer_backend
 
     builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder)
-    builder.use_dedicated_xqa = True
+    builder.use_xqa = True
     qo_indptr = torch.tensor([0, 3, 3, 3], dtype=torch.int32)
 
     q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens(
@@ -929,6 +1030,42 @@ def test_flashinfer_xqa_query_lens_require_exact_uniform_product():
     assert q_cu_seq_lens.tolist() == [0, 3, 3, 3]
 
 
+@pytest.mark.skipif(
+    AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
+    reason="FlashInfer is not available.",
+)
+@pytest.mark.parametrize(
+    ("qo_indptr_values", "num_decode_tokens", "expected_q_len"),
+    [
+        ([0, 3, 3], 3, 3),
+        ([0, 2, 4, 4, 4], 4, 2),
+    ],
+)
+def test_flashinfer_trtllm_gen_padded_decode_uses_varlen_offsets(
+    qo_indptr_values: list[int],
+    num_decode_tokens: int,
+    expected_q_len: int,
+):
+    """Padded speculative decode keeps its actual packed query width."""
+    from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+
+    builder = object.__new__(flashinfer_backend.FlashInferMetadataBuilder)
+    builder.use_xqa = False
+    qo_indptr = torch.tensor(qo_indptr_values, dtype=torch.int32)
+
+    q_len, q_cu_seq_lens, q_lens = builder._compute_decode_query_lens(
+        qo_indptr,
+        qo_indptr,
+        num_decodes=len(qo_indptr_values) - 1,
+        num_decode_tokens=num_decode_tokens,
+    )
+
+    assert q_len == expected_q_len
+    assert q_lens is None
+    assert q_cu_seq_lens is not None
+    assert q_cu_seq_lens.tolist() == qo_indptr_values
+
+
 @pytest.mark.skipif(
     AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
     reason="FlashInfer is not available.",
@@ -961,10 +1098,12 @@ def test_flashinfer_attention_sinks_refreshed_after_reload(dtype):
     reason="FlashInfer is not available.",
 )
 def test_flashinfer_native_prefill_with_sinks(default_vllm_config):
-    if not (
-        current_platform.is_cuda() and current_platform.is_device_capability_family(120)
-    ):
-        pytest.skip("Native FlashInfer prefill with sinks requires SM12x.")
+    supported = current_platform.is_cuda() and (
+        current_platform.is_device_capability(90)
+        or current_platform.is_device_capability_family(120)
+    )
+    if not supported:
+        pytest.skip("Native FlashInfer prefill with sinks requires SM90 or SM12x.")
 
     from vllm.v1.attention.backends.flashinfer import FlashInferBackend
 
@@ -995,7 +1134,7 @@ def causal_mask_mod(
     reason="FlashInfer is not available.",
 )
 def test_flashinfer_xqa_decode_correctness(default_vllm_config):
-    """FlashInfer should route supported decode through XQA and match SDPA."""
+    """FlashInfer should route SM90/SM12x decode through XQA and match SDPA."""
     supported = current_platform.is_cuda() and (
         current_platform.is_device_capability(90)
         or current_platform.is_device_capability_family(120)
@@ -1066,16 +1205,23 @@ def causal_mask_mod(
             )
             attn_metadata = builder.build(0, common_attn_metadata)
 
-    expected_cg_support = (
-        AttentionCGSupport.UNIFORM_BATCH
-        if current_platform.is_device_capability_family(120)
-        else AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
-    )
     assert (
         flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support(
             vllm_config, kv_cache_spec
         )
-        == expected_cg_support
+        == AttentionCGSupport.UNIFORM_BATCH
+    )
+    wide_head_kv_cache_spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=kv_cache_spec.num_kv_heads,
+        head_size=512,
+        dtype=vllm_config.model_config.dtype,
+    )
+    assert (
+        flashinfer_backend.FlashInferMetadataBuilder.get_cudagraph_support(
+            vllm_config, wide_head_kv_cache_spec
+        )
+        == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
     )
     assert isinstance(
         attn_metadata.decode,
@@ -1184,6 +1330,45 @@ def sliding_window_mask_mod(
         )
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test")
+def test_rocm_aiter_fa_unquantized_cache_ignores_kv_scales(
+    default_vllm_config,
+):
+    """Stale scale tensors must not affect an unquantized KV cache."""
+    from vllm._aiter_ops import is_aiter_found_and_supported
+
+    if not is_aiter_found_and_supported():
+        pytest.skip("AITER is required")
+
+    def sliding_window_mask_mod(
+        b: torch.Tensor,
+        h: torch.Tensor,
+        q_idx: torch.Tensor,
+        kv_idx: torch.Tensor,
+        *,
+        context_len: int,
+        sliding_window: int,
+    ):
+        causal_mask = q_idx + context_len >= kv_idx
+        window_mask = q_idx + context_len - kv_idx < sliding_window
+        return causal_mask & window_mask
+
+    model = "microsoft/Phi-tiny-MoE-instruct"
+    batch_spec = BATCH_SPECS["small_decode"]
+    model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens))
+    sliding_window = model_config.get_sliding_window()
+    assert sliding_window is not None
+
+    _test_backend_correctness(
+        batch_spec,
+        model,
+        [AttentionBackendEnum.ROCM_AITER_FA],
+        partial(sliding_window_mask_mod, sliding_window=sliding_window),
+        layer_k_scale=0.25,
+        layer_v_scale=0.5,
+    )
+
+
 @pytest.mark.parametrize(
     "batch_spec_name",
     [
diff --git a/tests/v1/attention/test_attention_splitting.py b/tests/v1/attention/test_attention_splitting.py
index acbc22a02793..bad0f61b7c45 100644
--- a/tests/v1/attention/test_attention_splitting.py
+++ b/tests/v1/attention/test_attention_splitting.py
@@ -20,12 +20,12 @@
 
 @pytest.fixture
 def sample_query_start_loc():
-    """Sample query_start_loc tensor for testing"""
+    """Sample query_start_loc tensor for testing."""
     return torch.tensor([0, 5, 12, 20, 35, 50])
 
 
 def test_basic_slice_middle(sample_query_start_loc):
-    """Test slicing from middle of tensor"""
+    """Test slicing from middle of tensor."""
     req_slice = slice(1, 3)  # slice from index 1 to 3
     result = slice_query_start_locs(sample_query_start_loc, req_slice)
 
@@ -34,7 +34,7 @@ def test_basic_slice_middle(sample_query_start_loc):
 
 
 def test_slice_from_beginning(sample_query_start_loc):
-    """Test slicing from the beginning of tensor"""
+    """Test slicing from the beginning of tensor."""
     req_slice = slice(0, 2)  # slice from index 0 to 2
     result = slice_query_start_locs(sample_query_start_loc, req_slice)
 
@@ -43,7 +43,7 @@ def test_slice_from_beginning(sample_query_start_loc):
 
 
 def test_slice_to_end(sample_query_start_loc):
-    """Test slicing to the end of tensor"""
+    """Test slicing to the end of tensor."""
     req_slice = slice(3, 5)  # slice from index 3 to 5 (last index)
     result = slice_query_start_locs(sample_query_start_loc, req_slice)
 
@@ -52,7 +52,7 @@ def test_slice_to_end(sample_query_start_loc):
 
 
 def test_single_element_slice(sample_query_start_loc):
-    """Test slice that results in single element"""
+    """Test slice that results in single element."""
     req_slice = slice(2, 3)  # slice from index 2 to 3
     result = slice_query_start_locs(sample_query_start_loc, req_slice)
 
@@ -61,7 +61,7 @@ def test_single_element_slice(sample_query_start_loc):
 
 
 def test_full_tensor_slice(sample_query_start_loc):
-    """Test slicing the entire tensor"""
+    """Test slicing the entire tensor."""
     req_slice = slice(0, 5)  # slice entire tensor
     result = slice_query_start_locs(sample_query_start_loc, req_slice)
 
@@ -80,7 +80,7 @@ def test_slice_bounds_edge_cases(sample_query_start_loc):
 
 @pytest.fixture
 def small_decode_metadata():
-    """Create metadata for small decode batch"""
+    """Create metadata for small decode batch."""
     batch_spec = BATCH_SPECS["small_decode"]
     device = torch.device("cpu")
     return create_common_attn_metadata(batch_spec, block_size=16, device=device)
@@ -88,7 +88,7 @@ def small_decode_metadata():
 
 @pytest.fixture
 def large_decode_metadata():
-    """Create metadata for small decode batch"""
+    """Create metadata for small decode batch."""
     batch_spec = BATCH_SPECS["large_decode"]
     device = torch.device("cpu")
     return create_common_attn_metadata(batch_spec, block_size=16, device=device)
@@ -96,7 +96,7 @@ def large_decode_metadata():
 
 @pytest.fixture
 def mixed_small_metadata():
-    """Create metadata for mixed small batch"""
+    """Create metadata for mixed small batch."""
     batch_spec = BATCH_SPECS["mixed_small"]
     device = torch.device("cpu")
     return create_common_attn_metadata(batch_spec, block_size=16, device=device)
@@ -104,7 +104,7 @@ def mixed_small_metadata():
 
 # Tests for _make_metadata_with_slice
 def test_make_metadata_with_slice_decode_batch(small_decode_metadata):
-    """Test slicing decode batch metadata"""
+    """Test slicing decode batch metadata."""
     # Split first request only
     ubatch_slice = UBatchSlice(slice(0, 1), slice(0, 1))
 
@@ -119,7 +119,7 @@ def test_make_metadata_with_slice_decode_batch(small_decode_metadata):
 
 
 def test_make_metadata_with_slice_mixed_batch(mixed_small_metadata):
-    """Test slicing mixed batch metadata"""
+    """Test slicing mixed batch metadata."""
     ubatch_slice = UBatchSlice(slice(1, 3), slice(1, 7))  # Requests 1-3, tokens 1-7
 
     result = _make_metadata_with_slice(ubatch_slice, mixed_small_metadata)
@@ -132,7 +132,7 @@ def test_make_metadata_with_slice_mixed_batch(mixed_small_metadata):
 
 
 def test_split_attn_metadata_decode_batch(large_decode_metadata):
-    """Test splitting decode batch into two equal parts"""
+    """Test splitting decode batch into two equal parts."""
     num_tokens = large_decode_metadata.num_reqs
     mid_point = num_tokens // 2
     ubatch_slices = [
@@ -297,7 +297,7 @@ def test_split_decodes_and_prefills_uniform_mixed_batch_non_uniform_decodes():
 
 
 def test_split_decodes_and_prefills_uniform_padded_batch_all_same():
-    """uniform batch where all query lengths are identical with 0 length padded reqs."""
+    """Uniform batch where all query lengths are identical with 0 length padded reqs."""
     # All query lengths are 2, with decode_threshold=3 (so 2 <= 3)
     # This triggers the padded uniform path at line 891
     query_lens = [2, 2, 2, 0]
@@ -324,7 +324,7 @@ def test_split_decodes_and_prefills_uniform_padded_batch_all_same():
 def test_prefill_split_across_ubatches(
     seq_lens, query_lens, split_point, expected_first_reqs, expected_second_reqs
 ):
-    """Test splitting a prefill across ubatches"""
+    """Test splitting a prefill across ubatches."""
     import numpy as np
 
     device = torch.device("cpu")
diff --git a/tests/v1/attention/test_backend_per_kind.py b/tests/v1/attention/test_backend_per_kind.py
index 9d180d3aee53..53dedf3d6bce 100644
--- a/tests/v1/attention/test_backend_per_kind.py
+++ b/tests/v1/attention/test_backend_per_kind.py
@@ -1,16 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for per-KV-group attention backend selection (backend_per_kind)."""
+"""Tests for attention configuration and per-KV-group backend selection."""
+
+from types import SimpleNamespace
 
 import pytest
 
-from vllm.config.attention import AttentionConfig
+from vllm.config.attention import AttentionConfig, HiSparseConfig
 from vllm.model_executor.layers.attention.attention import (
     _largest_kernel_block_within,
 )
 from vllm.v1.attention.backend import AttentionType, MultipleOf
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.selector import get_attn_spec_kind
+from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig
 from vllm.v1.kv_cache_interface import KVCacheSpecKind
 
 
@@ -26,7 +29,7 @@
 def test_largest_kernel_block_within(supported_sizes, expected):
     class Backend:
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(kv_cache_spec=None):
             return supported_sizes
 
     assert _largest_kernel_block_within(Backend, 1024, 1024 * 1536, 2048) == expected
@@ -86,3 +89,52 @@ def test_backend_per_kind_rejects_unknown_kind():
 
 def test_backend_per_kind_defaults_empty():
     assert AttentionConfig().backend_per_kind == {}
+
+
+def test_hisparse_device_buffer_size_boundaries():
+    vllm_config = SimpleNamespace(
+        attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
+        speculative_config=None,
+    )
+    resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
+    assert resolved is not None
+    assert resolved.device_buffer_size == 256
+
+    vllm_config.attention_config.hisparse_config = HiSparseConfig(
+        device_buffer_size=127
+    )
+    with pytest.raises(ValueError, match="expected at least 128"):
+        ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
+
+    vllm_config.attention_config.hisparse_config = HiSparseConfig(
+        device_buffer_size=32768
+    )
+    resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
+    assert resolved is not None
+    assert resolved.device_buffer_size == 32768
+
+    vllm_config.attention_config.hisparse_config = HiSparseConfig(
+        device_buffer_size=32769
+    )
+    with pytest.raises(ValueError, match="int16 slot-index limit"):
+        ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
+
+
+def test_hisparse_device_buffer_covers_speculative_window():
+    vllm_config = SimpleNamespace(
+        attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
+        speculative_config=SimpleNamespace(
+            num_speculative_tokens=3,
+            parallel_drafting=False,
+        ),
+    )
+
+    resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
+    assert resolved is not None
+    assert resolved.device_buffer_size == 5 * 128
+
+    vllm_config.attention_config.hisparse_config = HiSparseConfig(
+        device_buffer_size=4 * 128 - 1
+    )
+    with pytest.raises(ValueError, match="expected at least 512"):
+        ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
diff --git a/tests/v1/attention/test_chunked_local_attention.py b/tests/v1/attention/test_chunked_local_attention.py
index c2798c8f2fa4..0826b7073080 100644
--- a/tests/v1/attention/test_chunked_local_attention.py
+++ b/tests/v1/attention/test_chunked_local_attention.py
@@ -181,7 +181,8 @@ def test_local_attention_virtual_batches(test_data: LocalAttentionTestData):
 
     # Convert to numpy for easier comparison
     actual_q_seqlens = np.diff(result.query_start_loc_cpu.numpy())
-    actual_k_seqlens = result.seq_lens_cpu.numpy()
+    assert result.seq_lens_cpu_upper_bound is not None
+    actual_k_seqlens = result.seq_lens_cpu_upper_bound.numpy()
 
     # Check that all query lengths are less than or equal to attn_chunk_size
     assert all(q_len <= attn_chunk_size for q_len in actual_q_seqlens)
diff --git a/tests/v1/attention/test_dcp_a2a_pack_mask.py b/tests/v1/attention/test_dcp_a2a_pack_mask.py
new file mode 100644
index 000000000000..cff1228baf03
--- /dev/null
+++ b/tests/v1/attention/test_dcp_a2a_pack_mask.py
@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The empty-shard mask fused into the DCP A2A pack kernel must be bit-exact
+with the eager ``mask_dcp_empty_shards_`` pass it replaces."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.attention.ops.dcp import (
+    _dcp_a2a_lse_pack_dim,
+    _dcp_a2a_pack_send,
+    mask_dcp_empty_shards_,
+)
+
+requires_accelerator = pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="needs a CUDA or ROCm device for the Triton pack kernel",
+)
+
+
+def _pack(out, lse, world_size, h_per_rank, head_dim, seq_lens, query_start_loc):
+    lse_pack_dim = _dcp_a2a_lse_pack_dim(out.dtype)
+    send = torch.zeros(
+        (world_size, out.shape[0], h_per_rank, head_dim + lse_pack_dim),
+        device=out.device,
+        dtype=out.dtype,
+    )
+    _dcp_a2a_pack_send(
+        out,
+        lse,
+        send,
+        world_size,
+        h_per_rank,
+        head_dim,
+        lse_pack_dim,
+        seq_lens=seq_lens,
+        query_start_loc=query_start_loc,
+    )
+    return send
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("world_size,h_per_rank", [(8, 2), (8, 16), (4, 4), (2, 1)])
+@pytest.mark.parametrize("tokens_per_req", [1, 2, 3])
+@pytest.mark.parametrize("num_pad_rows", [0, 2])
+@requires_accelerator
+@pytest.mark.parametrize(
+    "seq_lens_list",
+    [
+        [16, 16, 16, 16],  # no empty shards (the long-context decode case)
+        [16, 0, 16, 0],  # interior empty shards
+        [0, 0, 0, 0],  # every shard empty
+        [16, 16, 0, 0],  # trailing empty shards, i.e. cudagraph padding
+        [0, 16, 16, 16],  # leading empty shard
+        [16] * 52,  # the profiled conc-52 decode batch
+        [7],  # single request
+    ],
+)
+def test_fused_mask_matches_eager(
+    dtype, world_size, h_per_rank, tokens_per_req, num_pad_rows, seq_lens_list
+):
+    torch.manual_seed(0)
+    device = "cuda"
+    head_dim = 512
+    num_seqs = len(seq_lens_list)
+    # tokens_per_req > 1 is the MTP / multi-token-verify shape.
+    num_rows = num_seqs * tokens_per_req + num_pad_rows
+    num_heads = world_size * h_per_rank
+
+    out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype)
+    lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32)
+    seq_lens = torch.tensor(seq_lens_list, device=device, dtype=torch.int32)
+    query_start_loc = torch.arange(
+        0,
+        (num_seqs + 1) * tokens_per_req,
+        tokens_per_req,
+        device=device,
+        dtype=torch.int32,
+    )  # rows past query_start_loc[-1] are cudagraph padding
+
+    # Reference: mask eagerly, then pack with masking disabled.
+    ref_lse = lse.clone()
+    mask_dcp_empty_shards_(ref_lse, seq_lens, query_start_loc)
+    expected = _pack(out, ref_lse, world_size, h_per_rank, head_dim, None, None)
+
+    # Under test: pack with the mask fused in.
+    actual = _pack(
+        out, lse.clone(), world_size, h_per_rank, head_dim, seq_lens, query_start_loc
+    )
+
+    # The packed LSE slots hold halves of an fp32, not meaningful floats, so
+    # compare the send buffers bitwise: -inf's high half is a NaN in fp16.
+    assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), (
+        "fused mask is not bit-exact with the eager mask"
+    )
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("seed", range(8))
+@requires_accelerator
+def test_fused_mask_matches_eager_ragged(dtype, seed):
+    """Non-uniform per-request query lengths.
+
+    The uniform-stride cases above never exercise an unequal gap between
+    consecutive ``query_start_loc`` entries, which is what the in-kernel
+    boundary search has to get right. MTP with partial acceptance produces
+    exactly this: some requests contribute one row, others several, in the
+    same batch.
+    """
+    torch.manual_seed(seed)
+    device = "cuda"
+    world_size, h_per_rank, head_dim = 8, 2, 512
+    num_seqs = 17
+
+    query_lens = torch.randint(1, 5, (num_seqs,), dtype=torch.int32)
+    query_start_loc = torch.zeros(num_seqs + 1, dtype=torch.int32)
+    query_start_loc[1:] = torch.cumsum(query_lens, 0)
+    # Mix real, empty, and (via padding rows) out-of-range shards.
+    seq_lens = torch.randint(0, 2, (num_seqs,), dtype=torch.int32) * 16
+    num_rows = int(query_start_loc[-1]) + 3
+
+    query_start_loc = query_start_loc.to(device)
+    seq_lens = seq_lens.to(device)
+    num_heads = world_size * h_per_rank
+    out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype)
+    lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32)
+
+    ref_lse = lse.clone()
+    mask_dcp_empty_shards_(ref_lse, seq_lens, query_start_loc)
+    expected = _pack(out, ref_lse, world_size, h_per_rank, head_dim, None, None)
+    actual = _pack(
+        out, lse.clone(), world_size, h_per_rank, head_dim, seq_lens, query_start_loc
+    )
+
+    assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), (
+        f"ragged query_start_loc mismatch: query_lens={query_lens.tolist()}, "
+        f"seq_lens={seq_lens.tolist()}"
+    )
+
+
+@requires_accelerator
+def test_mask_disabled_is_unmasked():
+    """Passing no seq_lens/query_start_loc must leave the LSE untouched."""
+    torch.manual_seed(0)
+    device = "cuda"
+    world_size, h_per_rank, head_dim = 8, 2, 512
+    out = torch.randn(
+        4, world_size * h_per_rank, head_dim, device=device, dtype=torch.bfloat16
+    )
+    lse = torch.randn(4, world_size * h_per_rank, device=device, dtype=torch.float32)
+
+    packed = _pack(out, lse, world_size, h_per_rank, head_dim, None, None)
+    assert torch.isfinite(packed.float()).all()
+
+
+def test_eager_mask_handles_rank_with_no_local_sequences():
+    """Padded graph rows are all empty when a DCP rank has no sequences."""
+    lse = torch.randn(8, 4, dtype=torch.float32)
+    seq_lens = torch.empty(0, dtype=torch.int32)
+    query_start_loc = torch.tensor([0], dtype=torch.int32)
+
+    mask_dcp_empty_shards_(lse, seq_lens, query_start_loc)
+
+    assert torch.isneginf(lse).all()
diff --git a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py
new file mode 100644
index 000000000000..914cd249fea2
--- /dev/null
+++ b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py
@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v4.amd.rocm import (
+    DeepseekV4ROCMAiterMLASparseMetadataBuilder,
+    DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
+)
+from vllm.v1.attention.backend import AttentionCGSupport
+from vllm.v1.attention.backends.mla import indexer
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV4IndexerBackend,
+    DeepseekV32IndexerMetadataBuilder,
+    DeepseekV41IndexerBackend,
+)
+
+
+def _make_indexer_config(*, architecture: str, adaptive: bool):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(architectures=[architecture]),
+        num_speculative_tokens=1,
+        speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive),
+    )
+
+
+def _mock_rocm_platform(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Platform predicates are mutually exclusive in production. Override both
+    # so these ROCm policy tests do not inherit CUDA capabilities from the CI
+    # host running them.
+    monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(indexer.current_platform, "is_rocm", lambda: True)
+
+
+def _make_indexer_builder(*, adaptive: bool, capacity: int = 12):
+    builder = DeepseekV32IndexerMetadataBuilder.__new__(
+        DeepseekV32IndexerMetadataBuilder
+    )
+    builder.vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive)
+    )
+    builder.supports_varlen = False
+    builder.decode_seq_lens_buffer = torch.zeros(capacity, dtype=torch.int32)
+    builder.expanded_block_table_buffer = torch.zeros((capacity, 2), dtype=torch.int32)
+    builder.decode_lens_buffer = torch.zeros(capacity, dtype=torch.int32)
+    builder.arange_buffer = torch.arange(capacity, dtype=torch.int32)
+    return builder
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs():
+    adaptive_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(enable_adaptive_verification=True)
+    )
+    fixed_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(enable_adaptive_verification=False)
+    )
+
+    for builder_cls in (
+        DeepseekV4ROCMAiterMLASparseMetadataBuilder,
+        DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
+    ):
+        assert (
+            builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace())
+            == AttentionCGSupport.ALWAYS
+        )
+        assert (
+            builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace())
+            == AttentionCGSupport.UNIFORM_BATCH
+        )
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_adaptive_indexer_support(monkeypatch: pytest.MonkeyPatch):
+    _mock_rocm_platform(monkeypatch)
+    adaptive_config = _make_indexer_config(
+        architecture="DeepseekV4ForCausalLM", adaptive=True
+    )
+    fixed_config = _make_indexer_config(
+        architecture="DeepseekV4ForCausalLM", adaptive=False
+    )
+
+    assert DeepseekV4IndexerBackend.supports_device_cpu_query_lens_mismatch()
+    assert indexer._rocm_supports_flattened_device_query_lens(adaptive_config)
+    assert (
+        DeepseekV4IndexerBackend.get_builder_cls() is DeepseekV32IndexerMetadataBuilder
+    )
+    assert indexer._use_flattening(adaptive_config)
+    assert (
+        DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
+            adaptive_config, SimpleNamespace()
+        )
+        == AttentionCGSupport.ALWAYS
+    )
+    assert not indexer._use_flattening(fixed_config)
+    assert (
+        DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
+            fixed_config, SimpleNamespace()
+        )
+        == AttentionCGSupport.UNIFORM_BATCH
+    )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "architecture",
+    ["DeepseekV32ForCausalLM", "DeepseekV41ForCausalLM", "GlmMoeDsaForCausalLM"],
+)
+def test_rocm_adaptive_indexer_flattening_is_scoped_to_deepseek_v4(
+    monkeypatch: pytest.MonkeyPatch,
+    architecture: str,
+):
+    _mock_rocm_platform(monkeypatch)
+    adaptive_config = _make_indexer_config(architecture=architecture, adaptive=True)
+
+    assert not indexer._rocm_supports_flattened_device_query_lens(adaptive_config)
+    assert not indexer._use_flattening(adaptive_config)
+    assert (
+        DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
+            adaptive_config, SimpleNamespace()
+        )
+        == AttentionCGSupport.UNIFORM_BATCH
+    )
+
+    if architecture == "DeepseekV41ForCausalLM":
+        assert not DeepseekV41IndexerBackend.supports_device_cpu_query_lens_mismatch()
+
+
+@pytest.mark.cpu_test
+def test_rocm_adaptive_indexer_preserves_single_request_uniform_path(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    builder = _make_indexer_builder(adaptive=True, capacity=8)
+
+    class FakeUniformKernel:
+        called = False
+
+        def __call__(
+            self,
+            seq_lens,
+            decode_seq_lens,
+            block_table,
+            expanded_block_table,
+            decode_lens,
+            num_decode_tokens,
+            max_decode_len,
+        ):
+            self.called = True
+            assert num_decode_tokens == 4
+            assert block_table.stride(0) == expanded_block_table.stride(0) == 2
+            decode_seq_lens[:max_decode_len] = torch.arange(
+                seq_lens[0] - max_decode_len + 1,
+                seq_lens[0] + 1,
+                dtype=torch.int32,
+            )
+            expanded_block_table[:max_decode_len] = block_table[0]
+            decode_lens[:max_decode_len] = 1
+
+    fake_kernel = FakeUniformKernel()
+    monkeypatch.setattr(indexer, "_PREPARE_UNIFORM_DECODE_KERNEL", fake_kernel)
+
+    seq_lens, block_table, decode_lens, batch_size, requires_padding = (
+        builder._prepare_decode_tensors(
+            seq_lens=torch.tensor([10], dtype=torch.int32),
+            block_table=torch.tensor([[1, 2]], dtype=torch.int32),
+            decode_lens=torch.tensor([4], dtype=torch.int32),
+            decode_lens_cpu=torch.tensor([4], dtype=torch.int32),
+            query_start_loc=torch.tensor([0], dtype=torch.int32),
+            num_decodes=1,
+            num_decode_tokens=4,
+            use_native=False,
+            next_n=8,
+            max_decode_len=4,
+        )
+    )
+
+    assert fake_kernel.called
+    torch.testing.assert_close(seq_lens, torch.tensor([7, 8, 9, 10], dtype=torch.int32))
+    torch.testing.assert_close(
+        block_table,
+        torch.tensor([[1, 2], [1, 2], [1, 2], [1, 2]], dtype=torch.int32),
+    )
+    torch.testing.assert_close(decode_lens, torch.ones(4, dtype=torch.int32))
+    assert batch_size == 4
+    assert not requires_padding
+
+
+@pytest.mark.cpu_test
+def test_rocm_adaptive_indexer_replays_changed_allocations_with_stable_buffers():
+    builder = _make_indexer_builder(adaptive=True)
+    buffer_ptrs = (
+        builder.decode_seq_lens_buffer.data_ptr(),
+        builder.expanded_block_table_buffer.data_ptr(),
+        builder.decode_lens_buffer.data_ptr(),
+    )
+    block_table = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
+    decode_lens_cpu = torch.tensor([4, 4, 0], dtype=torch.int32)
+
+    first = builder._prepare_decode_tensors(
+        seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32),
+        block_table=block_table,
+        decode_lens=torch.tensor([7, 1, 0], dtype=torch.int32),
+        decode_lens_cpu=decode_lens_cpu,
+        query_start_loc=torch.tensor([0, 7, 8], dtype=torch.int32),
+        num_decodes=3,
+        num_decode_tokens=10,
+        use_native=False,
+        next_n=8,
+        max_decode_len=4,
+    )
+    torch.testing.assert_close(
+        first[0],
+        torch.tensor([14, 15, 16, 17, 18, 19, 20, 20, 0, 0], dtype=torch.int32),
+    )
+    torch.testing.assert_close(
+        first[1][:, 0],
+        torch.tensor([1, 1, 1, 1, 1, 1, 1, 3, 0, 0], dtype=torch.int32),
+    )
+
+    second = builder._prepare_decode_tensors(
+        seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32),
+        block_table=block_table,
+        decode_lens=torch.tensor([1, 7, 0], dtype=torch.int32),
+        decode_lens_cpu=decode_lens_cpu,
+        query_start_loc=torch.tensor([0, 1, 8], dtype=torch.int32),
+        num_decodes=3,
+        num_decode_tokens=10,
+        use_native=False,
+        next_n=8,
+        max_decode_len=4,
+    )
+    torch.testing.assert_close(
+        second[0],
+        torch.tensor([20, 14, 15, 16, 17, 18, 19, 20, 0, 0], dtype=torch.int32),
+    )
+    torch.testing.assert_close(
+        second[1][:, 0],
+        torch.tensor([1, 3, 3, 3, 3, 3, 3, 3, 0, 0], dtype=torch.int32),
+    )
+    torch.testing.assert_close(second[2], torch.ones(10, dtype=torch.int32))
+    assert second[3:] == (10, False)
+    assert buffer_ptrs == (
+        builder.decode_seq_lens_buffer.data_ptr(),
+        builder.expanded_block_table_buffer.data_ptr(),
+        builder.decode_lens_buffer.data_ptr(),
+    )
diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py
index d18b7c671a54..7a74ce8c81be 100644
--- a/tests/v1/attention/test_deepseek_v4_swa_visible.py
+++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py
@@ -11,6 +11,8 @@
 ``pos + right`` (inclusive).
 """
 
+import inspect
+
 import pytest
 import torch
 from typing_extensions import TypedDict
@@ -57,10 +59,13 @@ def ref_left_right(
     return lefts, rights
 
 
-def ref_swa_bounds(pos: int, window: int, left: int, right: int) -> tuple[int, int]:
-    """Reference [start, end) window bounds for one query token."""
+def ref_swa_bounds(
+    pos: int, window: int, left: int, right: int, replay_start: int = 0
+) -> tuple[int, int]:
+    """Reference [start, end) window bounds for one query token. Under SWA
+    bounded replay no window KV exists below ``replay_start``."""
     left_add = max(left - (window - 1), 0)
-    start = max(pos - (window - 1) - left_add, 0)
+    start = max(pos - (window - 1) - left_add, 0, replay_start)
     return start, pos + right + 1
 
 
@@ -97,12 +102,14 @@ def ref_swa_slot_rows(
     window: int,
     max_image_tokens: int,
     width: int,
+    replay_starts: list[int] | None = None,
 ) -> tuple[list[list[int]], list[int]]:
     """Reference paged slot-id rows and lens for every token in the batch."""
     block_table_cpu = block_table.cpu()
     lefts, rights = ref_left_right(
         seq_lens, query_lens, spans_per_req, max_image_tokens
     )
+    replay_starts = replay_starts or [0] * len(seq_lens)
     rows: list[list[int]] = []
     lens: list[int] = []
     token = 0
@@ -110,7 +117,9 @@ def ref_swa_slot_rows(
         prefix_len = seq_len - query_len
         for i in range(query_len):
             pos = prefix_len + i
-            start, end = ref_swa_bounds(pos, window, lefts[token], rights[token])
+            start, end = ref_swa_bounds(
+                pos, window, lefts[token], rights[token], replay_starts[req]
+            )
             row = []
             for p in range(start, end):
                 blk = int(block_table_cpu[req, p // BLOCK_SIZE])
@@ -139,6 +148,7 @@ def run_swa_kernel(
     swa_indices = torch.zeros(num_tokens, 1, width, dtype=torch.int32, device=device)
     swa_lens = torch.zeros(num_tokens, dtype=torch.int32, device=device)
     is_valid = slot_mapping >= 0
+    replay_start_t = torch.zeros(len(seq_lens), dtype=torch.int32, device=device)
 
     if with_image:
         lefts, rights = ref_left_right(
@@ -164,6 +174,7 @@ def run_swa_kernel(
         block_table,
         block_table.stride(0),
         BLOCK_SIZE,
+        replay_start_t,
         token_offset=0,
         HAS_IMAGE=with_image,
         TRITON_BLOCK_SIZE=1024,
@@ -289,18 +300,25 @@ def combine_case(
     query_lens: list[int],
     spans: list[list[tuple[int, int]]],
     with_image: bool,
+    replay_starts: list[int] | None = None,
+    combine_fn=combine_topk_swa_indices,
 ):
     """Run combine_topk_swa_indices and return (indices, lens, expected)."""
     device = torch.device("cuda")
     num_reqs = len(seq_lens)
+    replay_starts = replay_starts or [0] * num_reqs
     query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
     query_start_loc[1:] = torch.tensor(
         query_lens, dtype=torch.int32, device=device
     ).cumsum(0)
     num_tokens = int(query_start_loc[-1])
     seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
+    # The builder's gather covers only the context above the replay start.
     gather_lens = torch.tensor(
-        [q + min(s - q, WINDOW - 1) for s, q in zip(seq_lens, query_lens)],
+        [
+            q + min(max(s - q - r, 0), WINDOW - 1)
+            for s, q, r in zip(seq_lens, query_lens, replay_starts)
+        ],
         dtype=torch.int32,
         device=device,
     )
@@ -319,7 +337,14 @@ def combine_case(
     else:
         left_t = right_t = None
 
-    combined_indices, combined_lens = combine_topk_swa_indices(
+    # V4.1's combine has no image windows (causal SWA only).
+    has_image = "left_visible" in inspect.signature(combine_fn).parameters
+    image_kwargs = (
+        dict(left_visible=left_t, right_visible=right_t, max_image_tokens=MAX_IMG)
+        if has_image
+        else {}
+    )
+    combined_indices, combined_lens = combine_fn(
         topk_indices,
         query_start_loc,
         seq_lens_t,
@@ -329,9 +354,7 @@ def combine_case(
         topk,
         M,
         N,
-        left_visible=left_t,
-        right_visible=right_t,
-        max_image_tokens=MAX_IMG,
+        **image_kwargs,
     )
 
     # Reference rows.
@@ -342,7 +365,7 @@ def combine_case(
         MAX_IMG,
     )
     topk_cpu = topk_indices.cpu()
-    width = WINDOW + MAX_IMG
+    width = WINDOW + (MAX_IMG if has_image else 0)
     combined_topk = (topk + width + 127) // 128 * 128
     rows = []
     lens = []
@@ -354,6 +377,8 @@ def combine_case(
             pos = prefix_len + i
             topk_len = min((pos + 1) // compress_ratio, topk)
             start, end = ref_swa_bounds(pos, WINDOW, lefts[token], rights[token])
+            # The window never reaches below the gathered buffer.
+            start = max(start, gather_start)
             swa_len = end - start
             row = [-1] * combined_topk
             for j in range(topk_len):
@@ -388,6 +413,31 @@ def test_combine_topk_swa_with_image_spans(cfg):
     assert indices.cpu().tolist() == rows
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("cfg", COMBINE_CASES)
+def test_v41_combine_topk_swa_stops_at_replay_start(cfg):
+    """SWA bounded replay: the gathered buffer starts at replay_start, so the
+    window never indexes below it."""
+    from vllm.models.deepseek_v41.common.ops.cache_utils import (
+        combine_topk_swa_indices as combine_v41,
+    )
+
+    # Request 0 replays [16, 40): the windows of its first rows would
+    # otherwise reach below 16.
+    indices, lens, rows, exp_lens = combine_case(
+        cfg["compress_ratio"],
+        cfg["topk"],
+        seq_lens=[40, 12],
+        query_lens=[24, 12],
+        spans=[[], []],
+        with_image=False,
+        replay_starts=[16, 0],
+        combine_fn=combine_v41,
+    )
+    assert lens.cpu().tolist() == exp_lens
+    assert indices.cpu().tolist() == rows
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 @pytest.mark.parametrize("cfg", COMBINE_CASES)
 def test_combine_topk_swa_without_image_unchanged(cfg):
@@ -462,13 +512,32 @@ def test_flashinfer_mixed_sparse_indices_with_image_spans():
     assert actual == exp_rows[1:]
 
 
-def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder:
-    overrides: dict = {"sliding_window": WINDOW}
+def make_builder(
+    vision: bool,
+    window: int = WINDOW,
+    max_image_tokens: int = MAX_IMG,
+    max_num_batched_tokens: int = 64,
+    model_name: str = "meta-llama/Meta-Llama-3-8B",
+    builder_cls: type[
+        DeepseekSparseSWAMetadataBuilder
+    ] = DeepseekSparseSWAMetadataBuilder,
+    mm_prefix_clamp_sliding_window: bool = True,
+    bounded_replay: bool = False,
+) -> DeepseekSparseSWAMetadataBuilder:
+    overrides: dict = {"sliding_window": window}
     if vision:
-        overrides.update(vision_n_layers=2, vision_max_n_token=MAX_IMG)
+        # Emulate the V4 vision config, whose mm_prefix_clamp_sliding_window
+        # gates the in-kernel SWA widening. V4.1 (causal image tokens) leaves
+        # it off.
+        overrides.update(
+            vision_n_layers=2,
+            vision_max_n_token=max_image_tokens,
+            mm_prefix_clamp_sliding_window=mm_prefix_clamp_sliding_window,
+        )
     vllm_config = create_vllm_config(
+        model_name=model_name,
         max_model_len=4096,
-        max_num_batched_tokens=64,
+        max_num_batched_tokens=max_num_batched_tokens,
         max_num_seqs=8,
         hf_config_override=overrides,
     )
@@ -477,11 +546,12 @@ def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder:
         num_kv_heads=1,
         head_size=512,
         dtype=torch.bfloat16,
-        sliding_window=WINDOW,
+        sliding_window=window,
         cache_dtype_str="auto",
         model_version="deepseek_v4",
+        bounded_replay=bounded_replay,
     )
-    return DeepseekSparseSWAMetadataBuilder(
+    return builder_cls(
         kv_cache_spec=spec,
         layer_names=["layer0"],
         vllm_config=vllm_config,
@@ -494,11 +564,17 @@ def build_metadata(
     seq_lens: list[int],
     query_lens: list[int],
     mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None,
+    replay_starts: list[int] | None = None,
 ):
     device = torch.device("cuda")
     query_start_loc, seq_lens_t, _, slot_mapping, block_table = make_batch(
         seq_lens, query_lens, device
     )
+    replay_start = (
+        None
+        if replay_starts is None
+        else torch.tensor(replay_starts, dtype=torch.int32, device=device)
+    )
     return builder.build(
         0,
         CommonAttentionMetadata(
@@ -515,6 +591,7 @@ def build_metadata(
             causal=True,
             mm_req_doc_ranges=mm_req_doc_ranges,
         ),
+        replay_start=replay_start,
     )
 
 
@@ -578,3 +655,165 @@ def test_builder_text_model_unchanged():
     )
     assert md.prefill_swa_lens.cpu().tolist() == lens
     assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("compress_ratio", [0, 1, 2])
+@pytest.mark.parametrize("query_len", [100, 4000])
+def test_v41_image_prefill_uses_causal_swa(compress_ratio, query_len):
+    """V4.1 image tokens use the same causal SWA as text, including in chunks."""
+    from vllm.models.deepseek_v41.common.ops.cache_utils import (
+        combine_topk_swa_indices as combine_v41,
+    )
+    from vllm.models.deepseek_v41.sparse_mla import DeepseekV41SparseSWAMetadataBuilder
+
+    seq_len, window, max_image_tokens = 4000, 128, 2048
+    spans = [(1900, 3947)]
+    builder = make_builder(
+        True,
+        window,
+        max_image_tokens,
+        query_len,
+        builder_cls=DeepseekV41SparseSWAMetadataBuilder,
+        mm_prefix_clamp_sliding_window=False,
+    )
+    assert builder.max_image_tokens == 0
+    md = build_metadata(builder, [seq_len], [query_len], {0: spans})
+    assert md.prefill_left_visible is None
+    assert md.prefill_right_visible is None
+    assert md.prefill_swa_indices.shape[-1] == window
+    plan = md.get_prefill_chunk_plan(
+        compress_ratio, prefill_chunk_size=4, has_compressed=compress_ratio > 0
+    )
+    [(start, end, n, m)] = plan
+    gather_len = int(md.prefill_gather_lens[0])
+    assert (start, end, m - n) == (0, 1, gather_len)
+    assert gather_len == query_len + min(seq_len - query_len, window - 1)
+    indices, lens = combine_v41(
+        torch.empty(query_len, 0, dtype=torch.int32, device="cuda"),
+        md.query_start_loc,
+        md.prefill_seq_lens,
+        md.prefill_gather_lens,
+        window,
+        compress_ratio,
+        0,
+        m,
+        n,
+    )
+    indices, lens = indices.cpu(), lens.cpu().tolist()
+    paged_indices = md.prefill_swa_indices.cpu()
+    paged_lens = md.prefill_swa_lens.cpu().tolist()
+    for token, pos in enumerate(range(seq_len - query_len, seq_len)):
+        lo, hi = max(0, pos - window + 1), pos + 1
+        row = indices[token, : lens[token]]
+        assert lens[token] == hi - lo
+        assert torch.all((row >= n) & (row < m))
+        absolute_keys = row - n + seq_len - gather_len
+        assert absolute_keys.tolist() == list(range(lo, hi))
+        assert paged_lens[token] == hi - lo
+        assert paged_indices[token, 0, : hi - lo].tolist() == list(range(lo, hi))
+
+
+# SWA bounded replay: every prefill index path clamps the window at replay_start.
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_builder_replay_start_bounds_prefill_window_and_gather():
+    """Paged-direct prefill indices stop at replay_start and the FlashMLA
+    gather only covers the context the request may see."""
+    seq_lens = [40, 12, 30]
+    query_lens = [24, 12, 10]
+    replay_starts = [12, 0, 0]
+    builder = make_builder(vision=False)
+    md = build_metadata(builder, seq_lens, query_lens, None, replay_starts)
+    assert md.replay_start is not None
+    assert md.replay_start.cpu().tolist() == replay_starts
+
+    _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda"))
+    rows, lens = ref_swa_slot_rows(
+        seq_lens,
+        query_lens,
+        [[], [], []],
+        block_table,
+        WINDOW,
+        MAX_IMG,
+        WINDOW,
+        replay_starts=replay_starts,
+    )
+    assert md.prefill_swa_lens.cpu().tolist() == lens
+    assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows
+    # Request 0: the first replayed row (pos 16) sees [12, 16]; its last row
+    # (pos 39) is a full window above the replay start and sees all of it.
+    assert md.prefill_swa_lens[0].item() == 16 - 12 + 1
+    assert md.prefill_swa_lens[23].item() == WINDOW
+    # gather_len = query_len + min(prefix_len - replay_start, WINDOW - 1).
+    assert md.prefill_gather_lens.cpu().tolist() == [
+        24 + min(16 - 12, WINDOW - 1),
+        12,
+        10 + min(20, WINDOW - 1),
+    ]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_builder_replaying_group_builds_without_replay_start():
+    """Graph captures build without model-specific kwargs: a replaying group's
+    metadata then bounds nothing."""
+    builder = make_builder(vision=False, bounded_replay=True)
+    md = build_metadata(builder, [40, 12], [24, 12], None)
+    assert not md.replay_start.any()
+    assert md.prefill_swa_lens[0].item() == WINDOW
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_v41_flashinfer_mixed_sparse_indices_respect_replay_start():
+    from vllm.models.deepseek_v41.common.ops.cache_utils import (
+        build_flashinfer_mixed_sparse_indices as build_v41,
+    )
+
+    device = torch.device("cuda")
+    # 1 decode token (req 0) + two prefill requests; req 1 resumes at 16.
+    seq_lens = [20, 40, 9]
+    query_lens = [1, 24, 9]
+    replay_starts = [0, 16, 0]
+    query_start_loc = torch.tensor([0, 1, 25, 34], dtype=torch.int32, device=device)
+    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
+    token_to_req = torch.tensor(
+        [0] + [1] * 24 + [2] * 9, dtype=torch.int32, device=device
+    )
+    block_table = torch.arange(3, dtype=torch.int32, device=device).view(3, 1)
+    kwargs = dict(
+        decode_swa_indices=torch.zeros((1, WINDOW), dtype=torch.int32, device=device),
+        decode_compressed_indices=None,
+        decode_compressed_topk_lens=None,
+        prefill_topk_indices=torch.empty((33, 0), dtype=torch.int32, device=device),
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens_t,
+        token_to_req_indices=token_to_req,
+        swa_block_table=block_table,
+        swa_block_size=BLOCK_SIZE,
+        compressed_block_table=None,
+        compressed_block_size=BLOCK_SIZE,
+        window_size=WINDOW,
+        compress_ratio=1,
+        topk=0,
+    )
+    plain, _ = build_v41(
+        replay_start=torch.zeros(3, dtype=torch.int32, device=device), **kwargs
+    )
+    bounded, _ = build_v41(
+        replay_start=torch.tensor(replay_starts, dtype=torch.int32, device=device),
+        **kwargs,
+    )
+    rows, _ = ref_swa_slot_rows(
+        seq_lens,
+        query_lens,
+        [[], [], []],
+        block_table,
+        WINDOW,
+        MAX_IMG,
+        WINDOW,
+        replay_starts=replay_starts,
+    )
+    # Decode row is copied through; prefill rows follow the bounded window.
+    assert bounded[0].cpu().tolist() == plain[0].cpu().tolist()
+    assert bounded[1:].cpu().tolist() == rows[1:]
diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py
index 8fd74314d1da..3d4af1741362 100644
--- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py
+++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py
@@ -168,6 +168,7 @@ def _run_sparse_backend_vs_sdpa(
     )
     model_config = vllm_config.model_config
     model_config.hf_text_config = SimpleNamespace(
+        index_topk=topk_tokens,
         q_lora_rank=None,
         kv_lora_rank=kv_lora_rank,
         qk_nope_head_dim=qk_nope_head_dim,
@@ -175,6 +176,7 @@ def _run_sparse_backend_vs_sdpa(
         v_head_dim=v_head_dim,
         model_type="deepseek_v2",
     )
+    del model_config.hf_config.index_topk  # Composite configs only nest this field.
     model_config.dtype = dtype
     model_config.get_num_attention_heads = MethodType(
         lambda self, parallel_config: num_heads, model_config
@@ -287,7 +289,7 @@ def _sparse_sdpa(idx_tensor, q_tok, k_mqa=k_mqa, v_mqa=v_mqa):
     causal_reference = torch.cat(causal_reference_outputs, dim=0)
 
     vllm_config.cache_config.cache_dtype = kv_cache_dtype
-    vllm_config.model_config.hf_config.index_topk = topk_tokens
+    vllm_config.model_config.hf_text_config.index_topk = topk_tokens
 
     common_attn_metadata = create_common_attn_metadata(
         batch_spec, block_size, device, arange_block_indices=True
@@ -602,3 +604,154 @@ def test_dspark_noncausal_differs_from_causal(
         f"non-causal backend output matches the causal reference "
         f"(max abs diff={causal_err}); future-pointing indices are not attended to"
     )
+
+
+@pytest.mark.parametrize("context_len", [20, 128, 900])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize("num_heads", [16, 64])
+def test_dsv41_flashinfer_dspark_window_matches_reference(
+    context_len, dtype, num_heads, monkeypatch
+):
+    """Draft queries see the full block without treating padded slots as keys."""
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("DSV4 TRTLLM sparse attention requires SM100")
+    from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import (
+        DeepseekSparseSWAFlashInferMetadataBuilder,
+        DeepseekV4FlashInferMLAAttention,
+    )
+    from vllm.models.deepseek_v41.sparse_mla import (
+        DeepseekV41SparseSWAMetadataBuilder,
+    )
+
+    torch.manual_seed(123)
+    device = "cuda"
+    query_lens = [5, 3]
+    context_lens = [context_len, context_len + 137]
+    num_real_tokens, num_tokens, head_dim = 8, 9, 512
+    cache = torch.randn(32, 128, head_dim, device=device, dtype=torch.bfloat16).to(
+        dtype
+    )
+    query = torch.randn(
+        num_tokens, num_heads, head_dim, device=device, dtype=torch.bfloat16
+    ).to(dtype)
+    indices = torch.full((num_tokens, 256), -1, device=device, dtype=torch.int32)
+    visible_indices = []
+    visible_lens = []
+    for req, (context, query_len) in enumerate(zip(context_lens, query_lens)):
+        visible = (
+            torch.arange(max(context - 128, 0), context + query_len, device=device)
+            + req * 16 * 128
+        )
+        visible_indices.extend([visible] * query_len)
+        visible_lens.extend([visible.numel()] * query_len)
+    for token, visible in enumerate(visible_indices):
+        indices[token, : visible.numel()] = visible.to(torch.int32)
+    query_start_loc = torch.tensor([0, 5, 8, 9], dtype=torch.int32)
+    metadata = SimpleNamespace(
+        num_decodes=3,
+        num_prefills=0,
+        num_decode_tokens=num_tokens,
+        num_prefill_tokens=0,
+        seq_lens=torch.tensor(
+            [context_lens[0] + 5, context_lens[1] + 3, 1],
+            device=device,
+            dtype=torch.int32,
+        ),
+        query_start_loc=query_start_loc.to(device),
+        query_start_loc_cpu=query_start_loc,
+        token_to_req_indices=torch.tensor(
+            [0] * 5 + [1] * 3 + [2], device=device, dtype=torch.int32
+        ),
+        decode_swa_indices=indices,
+        decode_swa_width=256,
+        decode_swa_lens=torch.tensor(
+            visible_lens + [0], device=device, dtype=torch.int32
+        ),
+        block_table=torch.arange(48, device=device, dtype=torch.int32).view(3, -1),
+        block_size=128,
+        replay_start=torch.zeros(3, device=device, dtype=torch.int32),
+        flashinfer_sparse_index_cache={},
+        max_decode_query_len=5,
+    )
+
+    # Exercise FlashInfer preparation without constructing a model/config.
+    def init_parent(builder):
+        builder._max_tokens = num_tokens
+        builder.device = device
+        builder.window_size = 128
+
+    monkeypatch.setattr(DeepseekV41SparseSWAMetadataBuilder, "__init__", init_parent)
+    monkeypatch.setattr(
+        DeepseekV41SparseSWAMetadataBuilder,
+        "build",
+        lambda *args: metadata,
+    )
+    builder = DeepseekSparseSWAFlashInferMetadataBuilder()
+    common_metadata = SimpleNamespace(causal=False)
+    builder.build(0, common_metadata)
+    prepared = (
+        metadata.flashinfer_decode_topk_lens,
+        metadata.flashinfer_decode_seq_lens,
+    )
+    attention = SimpleNamespace(
+        kv_cache_torch_dtype=dtype,
+        window_size=128,
+        compress_ratio=0,
+        topk_indices_buffer=torch.empty(
+            num_tokens, 0, device=device, dtype=torch.int32
+        ),
+        scale=1 / math.sqrt(head_dim),
+        _flashinfer_fp8_bmm1_scale=1 / math.sqrt(head_dim),
+        _flashinfer_fp8_bmm2_scale=1.0,
+        attn_sink=None,
+    )
+    attention._build_sparse_index_metadata = MethodType(
+        DeepseekV4FlashInferMLAAttention._build_sparse_index_metadata, attention
+    )
+    output = torch.empty_like(query, dtype=torch.bfloat16)
+
+    def forward():
+        metadata.flashinfer_sparse_index_cache.clear()
+        DeepseekV4FlashInferMLAAttention._forward(
+            attention, query, None, cache, metadata, None, True, output
+        )
+
+    def check_output():
+        references = []
+        for token, visible in enumerate(visible_indices):
+            keys = cache.flatten(0, 1)[visible].float()
+            weights = torch.softmax(
+                query[token].float() @ keys.T / math.sqrt(head_dim), -1
+            )
+            references.append(weights @ keys)
+        atol = 0.01 if dtype == torch.bfloat16 else 0.05
+        torch.testing.assert_close(
+            output[:num_real_tokens].float(),
+            torch.stack(references),
+            atol=atol,
+            rtol=0.05,
+        )
+
+    forward()
+    check_output()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        forward()
+    # Replay must consume the current inputs, including KV written since capture.
+    query.copy_(torch.randn_like(query, dtype=torch.bfloat16).to(dtype))
+    cache.copy_(torch.randn_like(cache, dtype=torch.bfloat16).to(dtype))
+    # A new step changes visibility as well as Q/KV, without recapturing.
+    metadata.seq_lens.sub_(1)
+    metadata.decode_swa_lens[:num_real_tokens].sub_(1)
+    for token, visible in enumerate(visible_indices):
+        indices[token, visible.numel() - 1] = -1
+        visible_indices[token] = visible[:-1]
+    builder.build(0, common_metadata)
+    assert metadata.flashinfer_decode_topk_lens.data_ptr() == prepared[0].data_ptr()
+    assert metadata.flashinfer_decode_seq_lens.data_ptr() == prepared[1].data_ptr()
+    torch.testing.assert_close(prepared[0], metadata.decode_swa_lens.clamp_min(128))
+    torch.testing.assert_close(
+        prepared[1], metadata.seq_lens[metadata.token_to_req_indices.long()]
+    )
+    graph.replay()
+    check_output()
diff --git a/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py b/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py
index 868d59247af5..94d9ecd9a0cc 100644
--- a/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py
+++ b/tests/v1/attention/test_flashinfer_dcp_spec_reorder.py
@@ -13,6 +13,7 @@
 
 from tests.v1.attention.utils import create_vllm_config
 from vllm.config import SpeculativeConfig, set_current_vllm_config
+from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.attention.backends import flashinfer as flashinfer_backend
 from vllm.v1.attention.backends.flashinfer import (
     FlashInferDecodeKernel,
@@ -73,3 +74,7 @@ def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch):
         builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
     )
     assert builder.reorder_batch_threshold == 1
+    assert (
+        FlashInferMetadataBuilder.get_cudagraph_support(vllm_config, kv_cache_spec)
+        == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
+    )
diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py
index 039d30c9aa9d..a72dbc21d145 100644
--- a/tests/v1/attention/test_flashinfer_mla_dcp.py
+++ b/tests/v1/attention/test_flashinfer_mla_dcp.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import MethodType
 from unittest.mock import MagicMock
 
 import pytest
@@ -61,6 +62,7 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
     attn_metadata.max_seq_len = 1
     attn_metadata.decode.block_table = torch.zeros(2, 1, dtype=torch.int32)
     attn_metadata.decode.seq_lens = torch.ones(2, dtype=torch.int32)
+    attn_metadata.decode.max_query_len = 1
 
     query = torch.ones(2, 24, 576, dtype=torch.bfloat16)
     kv_cache = torch.ones(1, 128, 576, dtype=torch.bfloat16)
@@ -100,6 +102,9 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
     impl.dcp_world_size = 8
     impl.dcp_rank = 3
     impl.cp_kv_cache_interleave_size = 1
+    impl._flattened_decode_metadata = MethodType(
+        flashinfer_mla.FlashInferMLAImpl._flattened_decode_metadata, impl
+    )
 
     num_reqs, query_len = 2, 3
     num_tokens = num_reqs * query_len
@@ -114,6 +119,11 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):
     attn_metadata.decode.block_table = block_table
     attn_metadata.decode.seq_lens = seq_lens
     attn_metadata.decode.dcp_tot_seq_lens = global_causal_seq_lens
+    attn_metadata.decode.max_query_len = query_len
+    attn_metadata.decode.query_start_loc = torch.arange(
+        0, num_tokens + 1, query_len, dtype=torch.int32
+    )
+    attn_metadata.decode.query_len = 0
 
     query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16)
     kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16)
diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py
index b20c16330659..72b2de05c274 100644
--- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py
+++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""CPU tests for the FlashInfer SM90 sparse MLA backend wiring (no GPU).
+"""Tests for the FlashInfer SM90 sparse MLA backend wiring and index packing.
 
 The FlashInfer wrapper and top-k conversion are replaced by CPU recorders;
 the tests pin the contract between the impl and the kernel API: page_size=1
@@ -65,12 +65,18 @@ class FakeState:
     def __init__(self, width, max_tokens=64):
         self.kv_indices = torch.zeros(max_tokens * width, dtype=torch.int32)
         self.kv_len_arr = torch.zeros(max_tokens, dtype=torch.int32)
+        self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32)
         self.wrapper = FakeWrapper()
         self.plan_calls = []
 
     def plan(self, num_tokens, kv_lens):
         self.plan_calls.append((num_tokens, kv_lens))
 
+    def pack_indices(self, slots):
+        for row in range(slots.shape[0]):
+            start, end = self.kv_indptr[row : row + 2].tolist()
+            self.kv_indices[start:end] = slots[row, : end - start].clamp(min=0)
+
 
 def make_impl(qk_rope, kv_dtype="fp8_e4m3", num_heads=2, topk_width=TOPK):
     impl = object.__new__(FlashInferMLASparseSM90Impl)
@@ -113,6 +119,9 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
         [6, 5, 4] + [-1] * (TOPK - 3),
         [2] + [-1] * (TOPK - 1),
     ]
+    impl.topk_indices_buffer.copy_(torch.tensor(topk_rows, dtype=torch.int32))
+    state.kv_indptr[1:5] = torch.tensor([8, 13, 16, 17], dtype=torch.int32)
+    state.kv_indptr[5:] = 17
     meta = make_batch(rows, topk_rows, [3])
     meta.state = state
     q_nope = torch.randn(rows, impl.num_heads, HEAD)
@@ -133,11 +142,13 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
     ref_slots, ref_counts = ref_convert(
         meta.req_id_per_token, meta.block_table, impl.topk_indices_buffer
     )
-    width = TOPK
-    got_slots = state.kv_indices[: rows * width].view(rows, width)
+    offset = 0
     for t in range(rows):
         k = int(ref_counts[t])
-        assert got_slots[t, :k].tolist() == ref_slots[t, :k].tolist()
+        assert (
+            state.kv_indices[offset : offset + k].tolist() == ref_slots[t, :k].tolist()
+        )
+        offset += k
     assert state.plan_calls == []
 
     assert state.wrapper.run_args is not None
@@ -151,32 +162,44 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
         assert kwargs == {}
 
 
-def test_builder_attaches_its_state(monkeypatch):
+@pytest.mark.parametrize("use_mha", [False, True])
+@pytest.mark.parametrize("num_decodes", [0, 1])
+def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_decodes):
+    """MHA prefill rows must not make the MQA kernel read beyond its query."""
     builder = object.__new__(FlashInferMLASparseSM90Builder)
     builder._index_topk = 2048
     builder._index_kpool = 4
     builder._async_scheduling = False
     builder.state = FakeState(TOPK)
+    builder._attention_layer = SimpleNamespace(_use_sparse_mha=lambda _: use_mha)
     metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata)
     metadata.state = None
+    metadata.num_prefills = 1
+    metadata.num_decode_tokens = num_decodes
     monkeypatch.setattr(
         sm90_mod.FlashInferMLASparseMetadataBuilder,
         "build",
         lambda *_args, **_kwargs: metadata,
     )
     cam = SimpleNamespace(
-        num_reqs=1,
-        query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32),
-        seq_lens=torch.tensor([1], dtype=torch.int32),
-        seq_lens_cpu_upper_bound=torch.tensor([1], dtype=torch.int32),
+        num_reqs=num_decodes + 1,
+        query_start_loc_cpu=torch.tensor(
+            [0, 1, 6] if num_decodes else [0, 5], dtype=torch.int32
+        ),
+        seq_lens_cpu_upper_bound=torch.tensor(
+            [1402, 5] if num_decodes else [5], dtype=torch.int32
+        ),
         positions=None,
     )
 
     result = builder.build(0, cam)
 
     assert result.state is builder.state
-    assert builder.state.plan_calls[0][0] == 1
-    assert builder.state.plan_calls[0][1].tolist() == [1]
+    expected_lens = [1402] if num_decodes else []
+    if not use_mha:
+        expected_lens += [1, 2, 3, 4, 5]
+    assert builder.state.plan_calls[0][0] == len(expected_lens)
+    assert builder.state.plan_calls[0][1].tolist() == expected_lens
 
 
 def test_plan_uses_state_params(monkeypatch):
@@ -209,13 +232,51 @@ def test_plan_uses_state_params(monkeypatch):
     args, kwargs = wrapper.plan_args
     (qo, kv, indices, kv_len, heads, ckv, kpe, page, causal, scale) = args
     assert qo.tolist() == [0, 1, 2, 3, 3]  # clamp: rows past 3 have no queries
-    assert kv.tolist() == [i * TOPK for i in (0, 1, 2, 3, 3)]
-    assert kv_len.tolist() == [2, 5, 7, TOPK]  # padded row keeps full width
+    assert kv.tolist() == [0, 2, 7, 14, 14]
+    assert kv_len.tolist() == [2, 5, 7, 0]
     assert (heads, ckv, kpe, page, causal) == (4, HEAD, 64, 1, False)
     assert scale == 576**-0.5
     assert kwargs["q_data_type"] == torch.bfloat16
     assert kwargs["kv_data_type"] == torch.bfloat16
 
+    # Replanning a smaller batch must clear the previous rows' lengths.
+    state.plan(1, torch.tensor([TOPK], dtype=torch.int32))
+    assert state._kv_cpu.tolist() == [0, TOPK, TOPK, TOPK, TOPK]
+    assert state._lens_cpu.tolist() == [TOPK, 0, 0, 0]
+    state.plan(0, torch.empty(0, dtype=torch.int32))
+    assert state._kv_cpu.tolist() == [0, 0, 0, 0, 0]
+    assert state._lens_cpu.tolist() == [0, 0, 0, 0]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+@pytest.mark.parametrize("width", [128, 2048, 2176])
+def test_pack_indices_replays_with_updated_offsets(width):
+    """Graph replay packs exact prefixes after row lengths and slots change."""
+    state = sm90_mod._SM90State.__new__(sm90_mod._SM90State)
+    state.kv_indptr = torch.zeros(5, dtype=torch.int32, device="cuda")
+    state.kv_indices = torch.full((4 * width,), -99, dtype=torch.int32, device="cuda")
+    # Exercise a contiguous view with an unaligned starting address.
+    slots = torch.arange(4 * width + 1, dtype=torch.int32, device="cuda")[1:]
+    slots = slots.view(4, width)
+    state.kv_indptr.copy_(torch.tensor([0, 1, 4, 4, 4]))
+    state.pack_indices(slots)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        state.pack_indices(slots)
+
+    for lengths in ([1, 3, 0, 0], [width, 0, 17, width - 1], [0, 0, 0, 0]):
+        offsets = torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()])
+        state.kv_indptr.copy_(offsets)
+        state.kv_indices.fill_(-99)
+        slots.add_(1)
+        graph.replay()
+        expected = torch.cat(
+            [slots[row, :length] for row, length in enumerate(lengths)]
+        )
+        torch.testing.assert_close(state.kv_indices[: expected.numel()], expected)
+        assert (state.kv_indices[expected.numel() :] == -99).all()
+
 
 def test_kv_lens_host_formula():
     """Per-row host lengths: context == position + 1; capped at
diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py
new file mode 100644
index 000000000000..5b63c2c609f9
--- /dev/null
+++ b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py
@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU-only tests for FlashInfer sparse MLA backend constraints."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
+    _DEFAULT_WORKSPACE_BUFFER_SIZE,
+    FlashInferMLASparseTRTLLMBackend,
+    _required_workspace_bytes,
+    compute_trtllm_sparse_mla_workspace_bytes,
+)
+
+# Reporter config from vllm#50781: GLM-5.2 (64 q heads), TP=8, DCP=8,
+# max_num_batched_tokens=16384, FlashInfer 0.6.14.
+_REPORTER_HEADS_PER_RANK = 64 // 8
+_REPORTER_DCP = 8
+_REPORTER_MAX_TOKENS = 16384
+# From the reported crash: trtllm-gen requested exactly this many bytes for
+# trtllm_gen_softmax_workspace after an 8 MiB counter carve had left
+# 404,750,336 of the 413,138,944-byte default buffer.
+_REPORTER_OBSERVED_SOFTMAX_BYTES = 1_611_661_312
+_REPORTER_OBSERVED_PRECARVE_BYTES = 8_388_608
+
+
+def test_reporter_softmax_carve_is_reproduced_byte_exactly():
+    # The crashing step scheduled 12288 tokens (a single 12K-token request);
+    # the softmax slab formula must reproduce the kernel's request exactly:
+    # 8 * (8 heads/rank * 8 dcp) * 12288 tokens * 256 + 1 MiB guard.
+    softmax_bytes = compute_trtllm_sparse_mla_workspace_bytes(
+        base_workspace_bytes=0,
+        dcp_world_size=_REPORTER_DCP,
+        num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
+        max_num_batched_tokens=12288,
+    )
+    assert softmax_bytes == _REPORTER_OBSERVED_SOFTMAX_BYTES
+
+
+def test_reporter_config_covers_observed_overflow():
+    computed = compute_trtllm_sparse_mla_workspace_bytes(
+        base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE,
+        dcp_world_size=_REPORTER_DCP,
+        num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
+        max_num_batched_tokens=_REPORTER_MAX_TOKENS,
+    )
+    assert computed >= (
+        _REPORTER_OBSERVED_SOFTMAX_BYTES + _REPORTER_OBSERVED_PRECARVE_BYTES
+    )
+
+
+def test_non_dcp_size_is_unchanged():
+    computed = compute_trtllm_sparse_mla_workspace_bytes(
+        base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE,
+        dcp_world_size=1,
+        num_heads_per_rank=128,
+        max_num_batched_tokens=65536,
+    )
+    assert computed == _DEFAULT_WORKSPACE_BUFFER_SIZE
+
+
+def test_combined_pcp_dcp_rejects_flashinfer_sparse():
+    config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=4,
+            decode_context_parallel_size=4,
+        ),
+        model_config=None,
+    )
+    with patch("vllm.config.get_current_vllm_config", return_value=config):
+        reason = FlashInferMLASparseTRTLLMBackend.supports_combination(
+            head_size=576,
+            dtype=torch.bfloat16,
+            kv_cache_dtype="fp8",
+            block_size=64,
+            use_mla=True,
+            has_sink=False,
+            use_sparse=True,
+            use_mm_prefix=False,
+            device_capability=SimpleNamespace(major=10),
+        )
+
+    assert reason is not None
+    assert "use FLASHMLA_SPARSE" in reason
+
+
+@pytest.mark.parametrize("pcp_size,dcp_size", [(4, 1), (1, 4)])
+def test_single_context_parallel_mode_allows_flashinfer_sparse(
+    pcp_size: int, dcp_size: int
+):
+    config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=pcp_size,
+            decode_context_parallel_size=dcp_size,
+        ),
+        model_config=None,
+    )
+    with patch("vllm.config.get_current_vllm_config", return_value=config):
+        reason = FlashInferMLASparseTRTLLMBackend.supports_combination(
+            head_size=576,
+            dtype=torch.bfloat16,
+            kv_cache_dtype="fp8",
+            block_size=64,
+            use_mla=True,
+            has_sink=False,
+            use_sparse=True,
+            use_mm_prefix=False,
+            device_capability=SimpleNamespace(major=10),
+        )
+
+    assert reason is None
+
+
+def test_default_constant_matches_envs_default(monkeypatch):
+    from vllm import envs
+
+    monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False)
+    assert _DEFAULT_WORKSPACE_BUFFER_SIZE == envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
+
+
+def test_env_unset_returns_computed(monkeypatch):
+    monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False)
+    required = _required_workspace_bytes(
+        dcp_world_size=_REPORTER_DCP,
+        num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
+        max_num_batched_tokens=_REPORTER_MAX_TOKENS,
+    )
+    assert required == compute_trtllm_sparse_mla_workspace_bytes(
+        _DEFAULT_WORKSPACE_BUFFER_SIZE,
+        _REPORTER_DCP,
+        _REPORTER_HEADS_PER_RANK,
+        _REPORTER_MAX_TOKENS,
+    )
+
+
+def test_env_override_below_computed_is_respected_with_warning(monkeypatch):
+    override = 100 * 1024 * 1024
+    monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override))
+    with patch(
+        "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger"
+    ) as mock_logger:
+        required = _required_workspace_bytes(
+            dcp_world_size=_REPORTER_DCP,
+            num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
+            max_num_batched_tokens=_REPORTER_MAX_TOKENS,
+        )
+    assert required == override
+    mock_logger.warning_once.assert_called_once()
+
+
+def test_env_override_above_computed_is_respected_without_warning(monkeypatch):
+    override = 8 * 1024 * 1024 * 1024
+    monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override))
+    with patch(
+        "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger"
+    ) as mock_logger:
+        required = _required_workspace_bytes(
+            dcp_world_size=_REPORTER_DCP,
+            num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
+            max_num_batched_tokens=_REPORTER_MAX_TOKENS,
+        )
+    assert required == override
+    mock_logger.warning_once.assert_not_called()
diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py
new file mode 100644
index 000000000000..db2c1becc411
--- /dev/null
+++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py
@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU tests for the FlashMLA sparse backend NoPE-512 (GLM5Next) gates.
+
+GLM5Next is rope-free MLA: qk_nope_head_dim=256, qk_rope_head_dim=0,
+v_head_dim=256, kv_lora_rank=512 -> head_size 512. These tests pin the
+selection contract only (no GPU required):
+
+* ``get_supported_head_sizes`` advertises 512 so the backend is a candidate;
+* ``supports_combination`` admits exactly two shapes for 512 -- a quantized
+  DS-MLA cache format (zero-padded 576/656B envelope) and a rope-free bf16
+  cache on SM90 -- and rejects everything else;
+* adding 512 must NOT reorder the SM100 priority list for bf16, which is why
+  the bf16 arm is restricted to SM90.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.platforms.interface import DeviceCapability
+from vllm.v1.attention.backends.mla.flashmla_sparse import (
+    QUANTIZED_DS_MLA_CACHE_FORMATS,
+    FlashMLASparseBackend,
+)
+
+SM90 = DeviceCapability(major=9, minor=0)
+SM100 = DeviceCapability(major=10, minor=0)
+
+BF16_CACHE_DTYPES = (None, "auto", "bfloat16", "float16")
+QUANTIZED_CACHE_DTYPES = ("fp8_ds_mla", "nvfp4_ds_mla")
+
+
+def _supports_combination(head_size, kv_cache_dtype, capability):
+    return FlashMLASparseBackend.supports_combination(
+        head_size,
+        torch.bfloat16,
+        kv_cache_dtype,
+        64,
+        use_mla=True,
+        has_sink=False,
+        use_sparse=True,
+        use_mm_prefix=False,
+        device_capability=capability,
+    )
+
+
+@pytest.fixture
+def rope_free_model(monkeypatch):
+    """Patch get_current_vllm_config so the model looks like GLM5Next NoPE."""
+    import vllm.config as cfg
+
+    monkeypatch.setattr(
+        cfg,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_text_config=SimpleNamespace(qk_rope_head_dim=0)
+            )
+        ),
+    )
+
+
+@pytest.fixture
+def rope_carrying_model(monkeypatch):
+    """Patch get_current_vllm_config with a rope-carrying (576-style) model."""
+    import vllm.config as cfg
+
+    monkeypatch.setattr(
+        cfg,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_text_config=SimpleNamespace(qk_rope_head_dim=64)
+            )
+        ),
+    )
+
+
+def test_supported_head_sizes_include_512():
+    assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512]
+
+
+def test_quantized_ds_mla_formats_are_the_envelope_set():
+    assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS
+
+
+@pytest.mark.parametrize("kv_cache_dtype", QUANTIZED_CACHE_DTYPES)
+def test_nope_512_quantized_rejected_here(kv_cache_dtype):
+    # Quantized DS-MLA NoPE-512 is served by the zero-padded 576/656B
+    # envelope, wired up in a separate change; this backend must reject it.
+    reason = _supports_combination(512, kv_cache_dtype, SM90)
+    assert reason is not None
+
+
+@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
+def test_nope_512_bf16_sm90_rope_free_accepted(kv_cache_dtype, rope_free_model):
+    assert _supports_combination(512, kv_cache_dtype, SM90) is None
+
+
+@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
+def test_nope_512_bf16_sm90_rope_carrying_rejected(kv_cache_dtype, rope_carrying_model):
+    reason = _supports_combination(512, kv_cache_dtype, SM90)
+    assert reason is not None and "rope-free" in reason
+
+
+@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
+def test_nope_512_bf16_sm100_rejected(kv_cache_dtype, rope_free_model):
+    # The SM90 restriction is what keeps the SM100 priority order unchanged.
+    reason = _supports_combination(512, kv_cache_dtype, SM100)
+    assert reason is not None
+
+
+@pytest.mark.parametrize("kv_cache_dtype", ("fp8", "fp8_e4m3"))
+def test_nope_512_plain_fp8_rejected(kv_cache_dtype, rope_free_model):
+    reason = _supports_combination(512, kv_cache_dtype, SM90)
+    assert reason is not None
+
+
+@pytest.mark.parametrize("capability", [SM90, SM100], ids=["sm90", "sm100"])
+@pytest.mark.parametrize(
+    "kv_cache_dtype",
+    ("auto", "bfloat16", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"),
+)
+def test_576_unchanged_for_all_dtypes(kv_cache_dtype, capability):
+    # DeepSeek 576 must be untouched by the NoPE-512 gates, except for the
+    # pre-existing SM100-only nvfp4 rule.
+    reason = _supports_combination(576, kv_cache_dtype, capability)
+    if kv_cache_dtype == "nvfp4_ds_mla" and capability.major != 10:
+        assert reason is not None and "SM100" in reason
+    else:
+        assert reason is None, reason
+
+
+def _sparse_order(kv_cache_dtype, head_size, num_heads=32, capability=SM100):
+    try:
+        from vllm.platforms.cuda import _get_backend_priorities
+    except ImportError as e:  # pragma: no cover - depends on host
+        pytest.skip(f"CUDA platform unavailable on this host: {e}")
+
+    priorities = _get_backend_priorities(
+        use_mla=True,
+        device_capability=capability,
+        num_heads=num_heads,
+        kv_cache_dtype=kv_cache_dtype,
+        head_size=head_size,
+    )
+    sparse = {
+        "FLASH_ATTN_MLA_SPARSE",
+        "FLASHMLA_SPARSE",
+        "FLASHINFER_MLA_SPARSE",
+        "FLASHINFER_MLA_SPARSE_SM90",
+    }
+    return [b.name for b in priorities if b.name in sparse]
+
+
+@pytest.mark.parametrize("num_heads", [32, 64], ids=["tp2", "tp1"])
+def test_sm100_bf16_512_priority_unchanged(num_heads):
+    """Adding head_size 512 must not reorder the SM100 bf16 sparse list.
+
+    On SM100 with >16 heads the hardcoded bf16 order is
+    [FLASHMLA_SPARSE, FLASHINFER_MLA_SPARSE] (FlashMLA first). head_size 512
+    must leave this untouched (identical to the 576 order); the guard that
+    keeps bf16-512 from *serving* through FlashMLA on SM100 lives in
+    supports_combination (gated to SM90), tested separately.
+    """
+    order512 = _sparse_order("auto", 512, num_heads=num_heads)
+    order576 = _sparse_order("auto", 576, num_heads=num_heads)
+    assert order512 == order576, order512
+    assert order512[0] == "FLASHMLA_SPARSE", order512
+
+
+def test_sm100_576_priority_unchanged():
+    order = _sparse_order("auto", 576, num_heads=32)
+    assert order[0] == "FLASHMLA_SPARSE", order
diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py
index 221f933d8942..b4d3acf07cbb 100644
--- a/tests/v1/attention/test_gdn_metadata_builder.py
+++ b/tests/v1/attention/test_gdn_metadata_builder.py
@@ -221,3 +221,86 @@ def test_full_cudagraph_spec_metadata_uses_request_count():
     assert meta.spec_query_start_loc.shape == (batch.batch_size + 1,)
     assert meta.num_accepted_tokens is not None
     assert meta.num_accepted_tokens.shape == (batch.batch_size,)
+
+
+def _build_non_spec(
+    batch: BatchSpec,
+    is_prefilling: list[bool] | None,
+    full_cuda_graph: bool = False,
+):
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE
+    ).replace(
+        is_prefilling=None
+        if is_prefilling is None
+        else torch.tensor(is_prefilling, dtype=torch.bool)
+    )
+    builder = _create_gdn_builder(full_cuda_graph=full_cuda_graph)
+    return builder, common_attn_metadata, builder.build(0, common_attn_metadata)
+
+
+@pytest.mark.parametrize(
+    ("seq_len", "query_len", "is_prefilling", "num_prefills"),
+    [
+        pytest.param(1, 1, True, 1, id="first-chunk"),
+        pytest.param(65, 1, True, 0, id="resumed-chunk"),
+        pytest.param(0, 0, True, 0, id="padding"),
+        pytest.param(1, 1, None, 0, id="missing-prefill-flag"),
+    ],
+)
+def test_one_token_chunk_classification(
+    seq_len: int,
+    query_len: int,
+    is_prefilling: bool | None,
+    num_prefills: int,
+):
+    """Only a real first chunk with a prefill flag needs state initialization."""
+    _, _, meta = _build_non_spec(
+        BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]),
+        is_prefilling=None if is_prefilling is None else [False, is_prefilling],
+    )
+
+    assert meta.num_prefills == num_prefills
+    assert meta.num_decodes == 2 - num_prefills
+    assert meta.num_prefill_tokens == num_prefills
+    assert meta.num_decode_tokens == 1 + query_len - num_prefills
+    if num_prefills:
+        assert meta.has_initial_state is not None
+        assert meta.has_initial_state.tolist() == [True, False]
+    else:
+        assert meta.has_initial_state is None
+
+
+def test_one_token_first_chunk_excludes_padding():
+    """Neither padding requests nor padding tokens count as prefill work."""
+    common = create_common_attn_metadata(
+        BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]),
+        BLOCK_SIZE,
+        DEVICE,
+    ).replace(
+        is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool),
+        num_actual_tokens=4,
+    )
+    meta = _create_gdn_builder().build(0, common)
+
+    assert meta.num_decodes == 1
+    assert meta.num_prefills == 1
+    assert meta.num_decode_tokens == 1
+    assert meta.num_prefill_tokens == 1
+
+
+def test_cudagraph_capture_batch_stays_decode_only():
+    """Capture rows have no history, but must still select decode kernels."""
+    batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4)
+    builder, common_attn_metadata, _ = _build_non_spec(
+        batch, [False] * 4, full_cuda_graph=True
+    )
+    meta = builder.build_for_cudagraph_capture(common_attn_metadata)
+
+    assert meta.num_prefills == 0
+    assert meta.num_decodes == 4
+    assert meta.has_initial_state is None
+    staged = meta.non_spec_state_indices_tensor
+    assert staged is not None
+    assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr()
+    torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0])
diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py
index 7f7c20d53e96..b099159bc766 100644
--- a/tests/v1/attention/test_indexer_dcp_localize.py
+++ b/tests/v1/attention/test_indexer_dcp_localize.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import numpy as np
 import pytest
 import torch
 
@@ -306,6 +307,23 @@ def test_get_dcp_local_seq_lens_preserves_mtp_bounds_shape():
     torch.testing.assert_close(actual, expected)
 
 
+def test_get_dcp_local_seq_lens_reuses_device_rank_tensor(monkeypatch):
+    seq_lens = torch.tensor([0, 1, 7, 8, 17], dtype=torch.int32)
+    world = 4
+    rank = 2
+    rank_tensor = torch.tensor(rank, dtype=torch.int32)
+    expected = get_dcp_local_seq_lens(seq_lens, world, rank, 2)
+
+    def fail_tensor_construction(*args, **kwargs):
+        raise AssertionError("rank tensor must be reused")
+
+    with monkeypatch.context() as patch:
+        patch.setattr(torch, "tensor", fail_tensor_construction)
+        actual = get_dcp_local_seq_lens(seq_lens, world, rank_tensor, 2)
+
+    torch.testing.assert_close(actual, expected)
+
+
 def test_get_dcp_local_seq_lens_must_run_after_decode_expansion():
     world = 2
     rank = 1
@@ -682,9 +700,8 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel(
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
 @pytest.mark.parametrize("interleave", [1, 2])
 @pytest.mark.parametrize("dcp_rank", [0, 1])
-# 384 is not a power of two, so it exercises the multi-tile atomic allocator
-# rather than the single-tile path 1024 takes.
-@pytest.mark.parametrize("num_topk", [1024, 384])
+# Include padded single-tile compaction and the multi-tile atomic fallback.
+@pytest.mark.parametrize("num_topk", [1024, 2176, 4224])
 def test_dcp_filter_compaction_matches_reference(
     interleave: int, dcp_rank: int, num_topk: int
 ):
@@ -711,7 +728,7 @@ def test_dcp_filter_compaction_matches_reference(
     for r in range(num_rows):
         # Ids are distinct, so a row holds at most `seq` valid entries.
         hi = min(num_topk * 3 // 4, seq)
-        n_valid = int(torch.randint(num_topk // 4, hi, (1,)).item())
+        n_valid = int(torch.randint(hi // 3, hi, (1,)).item())
         perm = torch.randperm(seq, device=device)[:n_valid].to(torch.int32)
         token_indices[r, :n_valid] = perm
 
@@ -962,3 +979,58 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp():
     dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses)
     torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5)
     torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("dcp_world_size", [2, 4, 8])
+@pytest.mark.parametrize("interleave", [1, 64])
+@pytest.mark.parametrize("rows_per_req", [1, 2])
+@pytest.mark.parametrize(
+    "req_lens", [[7], [8, 8], [1, 9], [63, 64, 65], [255, 256, 257], [256, 1, 2730]]
+)
+def test_pcp_plan_deinterleave_restores_global_order(
+    monkeypatch, dcp_world_size, interleave, rows_per_req, req_lens
+):
+    """Gathering block-interleaved shards must restore each request's token order."""
+    from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan
+
+    monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False)
+    monkeypatch.setattr(
+        "vllm.v1.attention.backends.mla.indexer.async_tensor_h2d",
+        lambda x, device: x.to(device),
+    )
+    scheduled = np.array(req_lens, dtype=np.int64)
+    shard_rows = np.array(
+        [
+            max(
+                _local_count(g, r, dcp_world_size, interleave)
+                for r in range(dcp_world_size)
+            )
+            for g in req_lens
+        ]
+    )
+    rows = np.repeat(np.arange(len(scheduled)), rows_per_req)
+    plan = build_pcp_global_chunk_plan(
+        rows, shard_rows[rows], dcp_world_size, torch.device("cpu"), interleave
+    )
+
+    starts = np.concatenate([[0], np.cumsum(scheduled)])
+
+    # Build each rank's padded shard exactly as the cache gather would.
+    padded_cu = plan.padded_local_cu[::rows_per_req].tolist()
+    shards = torch.zeros(dcp_world_size, plan.padded_local_total, 1)
+    for r in range(dcp_world_size):
+        for i, g in enumerate(req_lens):
+            owned = [t for t in range(g) if (t // interleave) % dcp_world_size == r]
+            for local, t in enumerate(owned):
+                shards[r, padded_cu[i] + local, 0] = starts[i] + t
+
+    gathered = shards.reshape(dcp_world_size * plan.padded_local_total, 1)
+    restored = gathered[plan.deinterleave_idx]
+    # The layout is padded per request; only the real positions are read.
+    row_start = plan.row_start_cu.tolist()
+    for row, i in enumerate(rows):
+        g = req_lens[i]
+        torch.testing.assert_close(
+            restored[row_start[row] : row_start[row] + g, 0],
+            torch.arange(starts[i], starts[i] + g, dtype=torch.float32),
+        )
diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
index c0ce30d2144a..83fceef0181e 100644
--- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
@@ -7,29 +7,180 @@
 import torch
 
 from tests.v1.attention.utils import create_vllm_config
+from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend
+from vllm.models.deepseek_v41.sparse_mla import (
+    DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend,
+)
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.mla.compressor_utils import (
     CompressedSlotMappingKernel,
+    get_compressed_slot_mapping,
 )
 from vllm.v1.attention.backends.mla.indexer import (
     BuildPrefillChunkMetadataKernel,
+    DeepseekV4IndexerBackend,
     DeepseekV32IndexerMetadataBuilder,
+    DeepseekV41IndexerBackend,
 )
 from vllm.v1.attention.backends.mla.sparse_utils import (
     ConvertReqIndexToGlobalIndexKernel,
 )
-from vllm.v1.kv_cache_interface import MLAAttentionSpec
+from vllm.v1.kv_cache_interface import (
+    MLAAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
+)
 from vllm.v1.worker.block_table import get_block_table_width
+from vllm.v1.worker.utils import select_common_block_size
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("query_lens", [[1], [6] * 16, [0, 2, 6, 0, 4], [0, 0, 0]])
+@pytest.mark.parametrize("padding", [0, 17])
+def test_fused_indexer_decode_metadata(query_lens, padding):
+    """Flatten device query boundaries and clear graph padding on every replay."""
+    from vllm.v1.attention.ops.metadata import _indexer_decode_metadata_kernel
+
+    device = "cuda"
+    lengths = torch.tensor(query_lens, device=device, dtype=torch.int32)
+    qsl = torch.cat(
+        [torch.zeros(1, device=device, dtype=torch.int32), lengths.cumsum(0).int()]
+    )
+    reqs = len(query_lens)
+    n = sum(query_lens)
+    tokens = n + padding
+    capacity = max(tokens, reqs) + 31
+    seq = lengths + 4096
+    # Exercise block-table row padding and masked columns.
+    bt = torch.arange(reqs * 74, device=device, dtype=torch.int32).view(reqs, 74)[
+        :, :65
+    ]
+    out_bt = torch.full((capacity, 69), -99, device=device, dtype=torch.int32)
+    outputs = [
+        torch.full((capacity,), -99, device=device, dtype=torch.int32) for _ in range(4)
+    ]
+    out_seq, out_lens, indices, per_req = outputs
+    grid = max(reqs, tokens + (capacity - tokens + 255) // 256)
+    _indexer_decode_metadata_kernel[(grid,)](
+        qsl,
+        seq,
+        bt,
+        out_seq,
+        out_bt,
+        out_lens,
+        indices,
+        per_req,
+        reqs,
+        tokens,
+        capacity,
+        bt.stride(0),
+        out_bt.stride(0),
+        BLOCK_COLS=bt.shape[1],
+        num_warps=4,
+    )
+    expected_req = torch.repeat_interleave(torch.arange(reqs, device=device), lengths)
+    expected_seq = torch.cat(
+        [torch.arange(4097, 4097 + q, device=device) for q in query_lens]
+    )
+    torch.testing.assert_close(out_seq[:n], expected_seq.int(), rtol=0, atol=0)
+    assert torch.count_nonzero(out_seq[n:]) == 0
+    torch.testing.assert_close(out_bt[:n, :65], bt[expected_req], rtol=0, atol=0)
+    assert torch.count_nonzero(out_bt[n:tokens, :65]) == 0
+    torch.testing.assert_close(indices[:n], expected_req.int(), rtol=0, atol=0)
+    torch.testing.assert_close(
+        indices[n:tokens],
+        torch.arange(reqs, reqs + padding, device=device, dtype=torch.int32),
+    )
+    assert torch.all(out_lens[:tokens] == 1)
+    torch.testing.assert_close(per_req[:reqs], lengths)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0]])
+def test_device_token_request_mapping(query_lens):
+    """Graph replay follows device boundaries even when CPU lengths are stale."""
+    lengths = torch.tensor(query_lens, device="cuda", dtype=torch.int32)
+    qsl = torch.cat(
+        [torch.zeros(1, device="cuda", dtype=torch.int32), lengths.cumsum(0).int()]
+    )
+    n = sum(query_lens)
+    output = torch.full((n + 7,), -99, device="cuda", dtype=torch.int32)
+    common = CommonAttentionMetadata(
+        query_start_loc=qsl,
+        query_start_loc_cpu=qsl.cpu(),
+        seq_lens=lengths,
+        num_reqs=len(query_lens),
+        num_actual_tokens=output.numel(),
+        max_query_len=max(query_lens),
+        max_seq_len=max(query_lens),
+        block_table_tensor=torch.empty(
+            (len(query_lens), 1), device="cuda", dtype=torch.int32
+        ),
+        slot_mapping=torch.full((n + 7,), -1, device="cuda", dtype=torch.int64),
+    )
+    result = common.token_to_req_indices(output)
+    assert result.data_ptr() == output.data_ptr()
+    expected = torch.repeat_interleave(
+        torch.arange(len(query_lens), device="cuda", dtype=torch.int32), lengths
+    )
+    torch.testing.assert_close(output[:n], expected)
+    assert torch.count_nonzero(output[n:]) == 0
+    graph = torch.cuda.CUDAGraph()
+    common._token_to_req_indices_cache = None
+    with torch.cuda.graph(graph):
+        common.token_to_req_indices(output)
+    reversed_lens = lengths.flip(0)
+    qsl[1:].copy_(reversed_lens.cumsum(0))
+    graph.replay()
+    expected = torch.repeat_interleave(
+        torch.arange(len(query_lens), device="cuda", dtype=torch.int32), reversed_lens
+    )
+    torch.testing.assert_close(output[:n], expected)
+    assert torch.count_nonzero(output[n:]) == 0
+
+
+def test_indexer_shares_uncompressed_block_size_with_deepseek_v4_mla():
+    """Packed MLA/indexer groups must retain 64 compressed rows per page."""
+    kernel_block_size = select_common_block_size(
+        256, [DeepseekV4SparseMLABackend, DeepseekV4IndexerBackend]
+    )
+    spec = MLAAttentionSpec(
+        block_size=256,
+        num_kv_heads=1,
+        head_size=132,
+        dtype=torch.uint8,
+        tokens_per_state=4,
+    )
+    assert compute_layer_kv_cache_shape_bytes(spec, 2, kernel_block_size) == (
+        2,
+        1,
+        64,
+        132,
+    )
+
+
+def test_indexer_preserves_deepseek_v41_mla_block_size():
+    """V4.1 retains its smaller pages independently of V4's C4 indexer."""
+    block_size = DeepseekV41SparseMLABackend.get_supported_kernel_block_sizes()[0]
+    assert isinstance(block_size, int)
+    assert (
+        select_common_block_size(
+            block_size, [DeepseekV41SparseMLABackend, DeepseekV41IndexerBackend]
+        )
+        == block_size
+    )
 
 
 def test_indexer_warmup_normalizes_zero_compress_ratios():
     config = SimpleNamespace(
         scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
         model_config=SimpleNamespace(
-            hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0], index_kpool=32)
+            hf_text_config=SimpleNamespace(
+                compress_ratios=[0, 0, 4, 128, 0], index_kpool=32
+            )
         ),
         parallel_config=SimpleNamespace(
             decode_context_parallel_size=1,
+            prefill_context_parallel_size=1,
             cp_kv_cache_interleave_size=1,
         ),
     )
@@ -44,25 +195,79 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
     }
 
 
+def test_indexer_warmup_includes_pcp_normalized_dcp_key(monkeypatch):
+    monkeypatch.setattr(
+        "vllm.v1.attention.backends.mla.indexer.get_dcp_group",
+        lambda: SimpleNamespace(rank_in_group=2),
+    )
+    config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+        model_config=SimpleNamespace(
+            hf_text_config=SimpleNamespace(compress_ratios=[32], index_kpool=32)
+        ),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=4,
+            prefill_context_parallel_size=4,
+            cp_kv_cache_interleave_size=1,
+        ),
+    )
+
+    keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config)
+
+    # Triton's compile key normalizes generic i32 values to 2 and divisible
+    # i32 values (including zero) to 16.
+    assert {(key.dcp_rank, key.dcp_world) for key in keys} == {(2, 2), (16, 1)}
+    assert {
+        (
+            key.input_variant.is_aligned("uncompressed_seq_lens"),
+            key.input_variant.is_aligned("cu_compressed_seq_lens"),
+        )
+        for key in keys
+    } == {(False, False), (False, True), (True, False), (True, True)}
+
+
 def test_compressed_slot_mapping_warmup_includes_index_kpool():
     config = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=256),
-        model_config=SimpleNamespace(hf_config=SimpleNamespace(index_kpool=32)),
+        model_config=SimpleNamespace(hf_text_config=SimpleNamespace(index_kpool=32)),
     )
 
     keys = CompressedSlotMappingKernel().get_warmup_keys(config)
     assert {(key.compress_ratio, key.block_size) for key in keys} == {(32, 2)}
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_compressed_slot_mapping_inherits_padded_token_slots():
+    """A token whose own slot is padded (SWA bounded replay) closes no
+    compressed state either."""
+    device = torch.device("cuda")
+    query_start_loc = torch.tensor([0, 8], dtype=torch.int32, device=device)
+    seq_lens = torch.tensor([8], dtype=torch.int32, device=device)
+    block_table = torch.tensor([[3]], dtype=torch.int32, device=device)
+    slot_mapping = torch.arange(8, dtype=torch.int64, device=device)
+    slot_mapping[:4] = -1
+    compressed = get_compressed_slot_mapping(
+        8,
+        slot_mapping,
+        query_start_loc,
+        seq_lens,
+        block_table,
+        block_size=4,
+        compress_ratio=2,
+    )
+    assert compressed.tolist() == [-1, -1, -1, -1, -1, 3 * 4 + 2, -1, 3 * 4 + 3]
+
+
 def test_index_conversion_warmup_uses_physical_block_stride():
     config = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=64),
         model_config=SimpleNamespace(
             max_model_len=1024,
-            hf_config=SimpleNamespace(index_topk=2048),
+            hf_text_config=SimpleNamespace(index_topk=2048),
         ),
         parallel_config=SimpleNamespace(
             decode_context_parallel_size=1,
+            prefill_context_parallel_size=1,
             cp_kv_cache_interleave_size=1,
         ),
     )
diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py
index 0966f9b20a6b..ed0e250b1811 100644
--- a/tests/v1/attention/test_kpool_tail_slot_mapping.py
+++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py
@@ -32,6 +32,7 @@
 )
 from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides
 from vllm.v1.kv_cache_layout import KVCacheLayout
+from vllm.v1.worker.block_table import get_block_table_width
 
 KPOOL = 4
 
@@ -55,6 +56,29 @@ def test_tail_backend_layout_matches_kernel_pointer_arithmetic():
     assert content_stride == 1
 
 
+def test_tail_spec_opts_out_of_generic_slot_mapping():
+    """The tail row is one block wide (padded to the block-table alignment), so
+    the generic kernel's ``pos // kpool`` column index runs off the end of the
+    allocation for long prompts. The spec must opt out of it entirely."""
+    spec = KpoolTailSpec(
+        block_size=KPOOL,
+        num_kv_heads=2,
+        head_size=128,
+        head_size_v=0,
+        dtype=torch.bfloat16,
+        sliding_window=KPOOL,
+    )
+    max_len = 1 << 20
+    width = get_block_table_width(
+        spec.max_num_blocks_per_req(None, max_len),
+        spec.block_size,
+        token_alignment=spec.block_table_token_alignment,
+    )
+
+    assert width * KPOOL < max_len
+    assert spec.uses_slot_mapping is False
+
+
 def make_tail_block_table(own_blocks, width=64):
     """Tail-group block table as BlockTables produces it: column 0 holds the
     request's single KpoolTailManager block, the remaining columns are never
@@ -367,3 +391,66 @@ def test_interleaved_decode_pollution_legacy_vs_circular():
 
     # The circular mapping keeps the rings isolated under interleaving.
     torch.testing.assert_close(circular, ground_truth)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+@pytest.mark.parametrize(
+    "per_req,num_actual,padded_len",
+    [
+        ([list(range(10)), list(range(12))], 22, 22),
+        ([list(range(10)), list(range(12))], 22, 30),
+        ([[3, 4], [0], [7, 8, 9]], 6, 8),
+        ([[5]], 1, 1),
+    ],
+)
+def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len):
+    """The CUDA (Triton) path must match the CPU torch reference, including
+    tokens between the last request boundary and num_actual_tokens (mapped to
+    the last request) and untouched padding beyond num_actual."""
+    positions, qsl, slot_mapping, _, num_reqs = make_batch(
+        per_req, padded_len=padded_len
+    )
+    # Replace the all--1 placeholder slots with sentinel values to check the
+    # padding range is copied through untouched.
+    slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000
+    bt = make_tail_block_table(list(range(5, 5 + num_reqs)))
+
+    ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
+    got = circular_tail_slots(
+        slot_mapping.cuda(),
+        bt.cuda(),
+        qsl.cuda().to(torch.int32),
+        positions.cuda(),
+        num_actual,
+        num_reqs,
+    )
+    torch.testing.assert_close(got.cpu(), ref)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
+def test_triton_mapping_reads_strided_block_table():
+    """The kernel must address the block table through its real stride(0),
+    not a dense assumption (the #57477 class of bug): a tail block table that
+    is a column view of a wider allocation must still read column 0."""
+    per_req = [list(range(10)), list(range(12))]
+    positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req)
+    own_blocks = [11, 22]
+    wide = torch.zeros(num_reqs, 8, dtype=torch.int32)
+    bt = wide[:, 1:4]  # non-contiguous view, stride(0) == 8
+    bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32)
+
+    ref = circular_tail_slots(
+        slot_mapping, bt.contiguous(), qsl, positions, num_actual, num_reqs
+    )
+    got = circular_tail_slots(
+        slot_mapping.cuda(),
+        bt.cuda(),
+        qsl.cuda().to(torch.int32),
+        positions.cuda(),
+        num_actual,
+        num_reqs,
+    )
+    torch.testing.assert_close(got.cpu(), ref)
+    for req, blk in enumerate(own_blocks):
+        start, end = int(qsl[req]), int(qsl[req + 1])
+        assert (got[start:end] // KPOOL == blk).all()
diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py
index 4ec138270203..f4451fe7af82 100644
--- a/tests/v1/attention/test_mamba_update_block_table.py
+++ b/tests/v1/attention/test_mamba_update_block_table.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Regression test for https://github.com/vllm-project/vllm/issues/34865
+"""Regression test for https://github.com/vllm-project/vllm/issues/34865.
 
 When multiple KV cache groups share the same MambaSpec (as in Nemotron
 hybrid models), the metadata caching optimization reuses metadata from
@@ -76,7 +75,6 @@ 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
     builder's buffers."""
-
     block_size = 16
     max_model_len = 256
     num_reqs = 4
@@ -181,7 +179,6 @@ def test_state_indices_tensor_d_includes_num_speculative_blocks():
     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
@@ -215,7 +212,6 @@ def test_block_idx_cudagraph_capture_padded_by_num_reqs():
     _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
diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py
index edc6cbaded32..2a9cd7736b74 100644
--- a/tests/v1/attention/test_mla_backends.py
+++ b/tests/v1/attention/test_mla_backends.py
@@ -25,6 +25,7 @@
 from vllm.model_executor.layers.attention import mla_attention as mla_attention_module
 from vllm.model_executor.layers.attention.mla_attention import (
     MLAAttention,
+    MLACommonBaseImpl,
     QueryLenSupport,
     _DecodeConcatQuantFP8,
     _use_masked_mha,
@@ -166,6 +167,25 @@ def test_glm5_flashinfer_masked_mha_routing(
     )
 
 
+@pytest.mark.parametrize("qk_rope_head_dim", [64, 0], ids=["rope", "nope"])
+def test_concat_k_nope_k_pe_matches_torch_cat(qk_rope_head_dim):
+    """The K concat used by the MLA prefill context loop must equal torch.cat of
+    k_nope with the broadcast k_pe; with no RoPE part it returns k_nope itself
+    instead of allocating and copying."""
+    torch.manual_seed(0)
+    num_tokens, num_heads, qk_nope_head_dim = 5, 4, 256
+    k_nope = torch.randn(num_tokens, num_heads, qk_nope_head_dim, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, 1, qk_rope_head_dim, dtype=torch.bfloat16)
+    impl = SimpleNamespace(_use_flashinfer_concat_mla_k=False)
+
+    k = MLACommonBaseImpl._concat_k_nope_k_pe(impl, k_nope, k_pe)
+
+    expected = torch.cat([k_nope, k_pe.expand(-1, num_heads, -1)], dim=-1)
+    assert k.shape == (num_tokens, num_heads, qk_nope_head_dim + qk_rope_head_dim)
+    torch.testing.assert_close(k, expected, rtol=0, atol=0)
+    assert (k.data_ptr() == k_nope.data_ptr()) == (qk_rope_head_dim == 0)
+
+
 def test_masked_mha_routing_is_dimension_specific():
     assert _use_masked_mha(
         backend_name="FLASHMLA_SPARSE",
@@ -200,6 +220,7 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
     layer = SimpleNamespace(
         kv_cache_dtype=cache_dtype,
         head_size=576,
+        indexer=None,
         non_causal_multi_token_decode=False,
         sliding_window=None,
     )
@@ -494,6 +515,7 @@ def create_and_prepopulate_kv_cache(
 
     Returns:
         MLA KV cache tensor
+
     """
     batch_size = len(kv_c_contexts)
     seq_lens = common_attn_metadata.seq_lens.cpu()
@@ -1245,6 +1267,7 @@ def fake_get_mla_metadata_dense_fp8(
         query_start_loc_cpu=query_start_loc,
         query_start_loc_device=query_start_loc,
         num_decode_tokens=2,
+        max_query_len=1,
         dcp_tot_seq_lens_device=None,
     )
 
@@ -1286,7 +1309,6 @@ def run_attention_backend(
     chunked_prefill_workspace_size: int | None = None,
 ) -> torch.Tensor:
     """Run attention computation using the specified backend's AttentionImpl."""
-
     builder_cls, impl_cls = try_get_attention_backend(backend)
 
     # Force the prefill backend selection (None means auto-select).
@@ -1413,8 +1435,7 @@ def _run_backend_correctness(
     v_head_dim: int,
     chunked_prefill_workspace_size: int | None = None,
 ):
-    """
-    Test that all backends produce similar outputs to a reference implementation
+    """Test that all backends produce similar outputs to a reference implementation
     using torch.nn.functional.scaled_dot_product_attention.
 
     This test works by:
@@ -1433,7 +1454,6 @@ def _run_backend_correctness(
     multiple GPUs. This tests that backends work correctly with different
     head counts.
     """
-
     # Filter backends to those that support the requested kv_cache_dtype
     backends_to_test = [
         b
diff --git a/tests/v1/attention/test_mla_context_chunks.py b/tests/v1/attention/test_mla_context_chunks.py
index 091079acc639..2d63b165a055 100644
--- a/tests/v1/attention/test_mla_context_chunks.py
+++ b/tests/v1/attention/test_mla_context_chunks.py
@@ -8,14 +8,20 @@
 longer needs an empty-span masking pass).
 """
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
+import vllm.utils.gpu_sync_debug as gsd
 from vllm.model_executor.layers.attention.mla_attention import (
     build_mla_chunked_context_metadata,
     init_mla_context_partial,
     reorg_kvcache,
 )
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SparseMLACommonMetadataBuilder,
+)
 
 BLOCK_SIZE = 16
 
@@ -45,6 +51,41 @@ def build_chunked_context(
     )
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_sparse_context_lengths_do_not_force_gpu_sync(monkeypatch):
+    """Sparse MLA must pin computed context lengths before the H2D copy."""
+    builder = SimpleNamespace(
+        _context_lens_cpu=torch.empty(2, dtype=torch.int32, pin_memory=True),
+        chunked_prefill_workspace=torch.empty((2048, 1)),
+        chunked_prefill_workspace_size=1024,
+        kv_cache_spec=SimpleNamespace(block_size=BLOCK_SIZE),
+        device=torch.device("cuda"),
+        dcp_world_size=1,
+        dcp_local_block_size=1,
+        dcp_virtual_block_size=1,
+        dcp_manager=None,
+    )
+    common_metadata = SimpleNamespace(
+        seq_lens_cpu_upper_bound=torch.tensor([17, 2052, 324], dtype=torch.int32),
+        query_start_loc_cpu=torch.tensor([0, 1, 5, 9], dtype=torch.int32),
+    )
+    monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
+    monkeypatch.setattr(gsd, "_sync_check_enabled", True)
+    gsd._install_copy_checkers()
+    metadata = gsd.with_gpu_sync_check(
+        SparseMLACommonMetadataBuilder._build_chunked_context_fields
+    )(
+        builder,
+        common_metadata,
+        num_decodes=1,
+        num_prefills=2,
+        prefill_query_lens_cpu=torch.tensor([4, 4], dtype=torch.int32),
+    )
+    assert metadata is not None
+    assert metadata.context_lens.device.type == "cuda"
+    assert metadata.context_lens.tolist() == [2048, 320]
+
+
 @pytest.mark.parametrize(
     "context_lens,workspace_size",
     [
@@ -73,6 +114,7 @@ def test_chunks_gather_every_context_row_exactly_once(context_lens, workspace_si
     query_lens = [4] * len(context_lens)
     metadata = build_chunked_context(context_lens, query_lens, workspace_size)
     assert metadata is not None
+    assert metadata.context_lens.tolist() == context_lens
 
     gathered: dict[int, list[tuple[int, int]]] = {}
     previous_request_start = -1
@@ -213,6 +255,7 @@ def test_dcp_chunks_fit_the_per_rank_row_budget():
         dcp_local_block_size=interleave,
     )
     assert metadata is not None
+    assert metadata.context_lens.tolist() == context_lens
     virtual_block_size = interleave * dcp_world_size
 
     local_cursor: dict[int, int] = {}
@@ -290,3 +333,23 @@ def test_dcp_reorg_uses_each_chunks_local_starts():
         )
 
         torch.testing.assert_close(reorganized, torch.cat(expected))
+
+
+def test_chunked_context_metadata_owns_context_lengths():
+    """Reusing CPU scratch must not overwrite an earlier batch's context lengths."""
+    context_lens = torch.tensor([16, 32], dtype=torch.int32)
+    metadata = build_mla_chunked_context_metadata(
+        context_lens_cpu=context_lens,
+        prefill_query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32),
+        chunked_prefill_workspace=torch.empty((64, 1)),
+        chunked_prefill_workspace_size=64,
+        block_size=BLOCK_SIZE,
+        align_chunk_to_block=True,
+        device=torch.device("cpu"),
+        dcp_world_size=1,
+        dcp_local_block_size=1,
+        dcp_virtual_block_size=1,
+    )
+    context_lens.zero_()
+    assert metadata is not None
+    assert metadata.context_lens.tolist() == [16, 32]
diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py
index d5856ab6bdb8..098eddfbf6eb 100644
--- a/tests/v1/attention/test_mla_prefill_selector.py
+++ b/tests/v1/attention/test_mla_prefill_selector.py
@@ -313,6 +313,40 @@ def test_backend_supported_dimension_validation(self):
             )
             assert invalid_reasons == []
 
+    def test_flash_attn_accepts_glm53_flash_nope_dimensions(self):
+        """(256, 0, 256) runs the same FA kernels as GLM-5's (192, 64, 256); a
+        RoPE-carrying 320-wide query does not."""
+        try:
+            from vllm.v1.attention.backends.mla.prefill.flash_attn import (
+                FlashAttnPrefillBackend,
+            )
+        except ImportError:
+            pytest.skip("MLA prefill backend not available")
+            return
+
+        capability = DeviceCapability(major=10, minor=0)
+
+        def validate(qk_rope_head_dim: int) -> list[str]:
+            selector_config = MLAPrefillSelectorConfig(
+                dtype=torch.bfloat16,
+                mla_dimensions=MLADimensions(
+                    qk_nope_head_dim=256,
+                    qk_rope_head_dim=qk_rope_head_dim,
+                    v_head_dim=256,
+                ),
+            )
+            with patch.object(
+                FlashAttnPrefillBackend, "is_available", return_value=True
+            ):
+                return FlashAttnPrefillBackend.validate_configuration(
+                    capability, selector_config
+                )
+
+        assert validate(qk_rope_head_dim=0) == []
+        invalid_reasons = validate(qk_rope_head_dim=64)
+        assert len(invalid_reasons) == 1
+        assert "supported MLA dimensions" in invalid_reasons[0]
+
 
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike(),
diff --git a/tests/v1/attention/test_mm_prefix.py b/tests/v1/attention/test_mm_prefix.py
index 6b8717389f45..f878ede4865f 100644
--- a/tests/v1/attention/test_mm_prefix.py
+++ b/tests/v1/attention/test_mm_prefix.py
@@ -638,3 +638,249 @@ def test_mm_prefix_kv_cache_path(head_size: int):
     assert not torch.allclose(expected, causal_only, atol=2e-2, rtol=2e-2), (
         "test batch does not actually exercise the mm_prefix branch"
     )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA backend imports")
+@pytest.mark.parametrize(
+    "query_len,seq_len,spans,expected,unclamped_expected",
+    [
+        (16, 32, [(20, 28)], True, True),
+        (16, 32, [(0, 16)], False, True),
+        (16, 32, [(0, 17)], True, True),
+        (16, 32, [(32, 40)], False, False),
+        (16, 32, [(31, 40)], False, True),
+        (1, 24, [(20, 28)], False, True),
+        (0, 24, [(20, 28)], False, False),
+        (16, 32, [(20, 20)], False, False),
+        (16, 64, [(20, 28)], False, False),
+    ],
+)
+def test_composite_routes_queries_that_need_image_masking(
+    query_len, seq_len, spans, expected, unclamped_expected
+):
+    """Historical images and single-query steps must use the causal graph path."""
+    from types import SimpleNamespace
+
+    from vllm.v1.attention.backends.composite import requires_mm_prefix
+
+    common = SimpleNamespace(
+        max_query_len=query_len,
+        mm_req_doc_ranges={0: spans},
+        query_start_loc_cpu=torch.tensor([0, query_len]),
+        seq_lens_cpu_upper_bound=torch.tensor([seq_len]),
+    )
+    assert requires_mm_prefix(common) is expected
+    assert requires_mm_prefix(common, unclamped_window=True) is unclamped_expected
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA attention kernels")
+@pytest.mark.parametrize("head_size", [256, 512])
+@pytest.mark.parametrize(
+    "reverse_children,cache_dtype",
+    [(False, "auto"), (True, "auto"), (False, "fp8_e4m3")],
+)
+@pytest.mark.parametrize("causal_backend_name", ["FLASHINFER", "FLASH_ATTN"])
+def test_composite_shared_cache_across_image_and_causal_steps(
+    head_size, reverse_children, cache_dtype, causal_backend_name
+):
+    """Changing routes must preserve KV writes, image masking and graph replay."""
+    if causal_backend_name == "FLASHINFER":
+        if not current_platform.is_device_capability_family(100):
+            pytest.skip("The Triton/FlashInfer composite supports Blackwell")
+    else:
+        if not current_platform.is_device_capability(90):
+            pytest.skip("The Triton/FlashAttention composite supports Hopper")
+        if reverse_children or cache_dtype != "auto":
+            pytest.skip("Reverse-child and FP8 coverage is backend-independent")
+
+    from vllm.config import set_current_vllm_config
+    from vllm.engine.arg_utils import EngineArgs
+    from vllm.model_executor.layers.attention import Attention
+    from vllm.utils.torch_utils import set_default_torch_dtype
+    from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
+    from vllm.v1.attention.backends.composite import (
+        MMPrefixAttentionRouting,
+        create_composite_attention_backend,
+    )
+    from vllm.v1.attention.backends.flash_attn import (
+        FlashAttentionMetadata,
+    )
+    from vllm.v1.attention.backends.flashinfer import (
+        FlashInferBackend,
+        FlashInferMetadata,
+        FlashInferTrtllmAPIDecode,
+        TRTLLMPrefill,
+    )
+    from vllm.v1.attention.backends.triton_attn import (
+        TritonAttentionBackend,
+        TritonAttentionMetadata,
+    )
+    from vllm.v1.attention.backends.triton_flash_attn import (
+        TritonFlashAttentionBackend,
+    )
+    from vllm.v1.attention.backends.triton_flashinfer import (
+        TritonFlashInferBackend,
+    )
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
+
+    if causal_backend_name == "FLASHINFER":
+        backend = TritonFlashInferBackend
+        causal_metadata_type = FlashInferMetadata
+        block_size = backend.get_supported_kernel_block_sizes()[0]
+    else:
+        backend = TritonFlashAttentionBackend
+        causal_metadata_type = FlashAttentionMetadata
+        block_size = backend.get_preferred_block_size(16)
+    num_blocks = (209 + block_size - 1) // block_size
+    torch.manual_seed(42)
+    attention_config = {"backend": backend.get_name()}
+    if causal_backend_name == "FLASH_ATTN":
+        attention_config["flash_attn_version"] = 4
+    cfg = EngineArgs(
+        model="google/gemma-4-31B-it",
+        dtype="bfloat16",
+        max_model_len=512,
+        max_num_seqs=4,
+        max_num_batched_tokens=512,
+        block_size=block_size,
+        kv_cache_dtype=cache_dtype,
+        enforce_eager=True,
+        attention_config=attention_config,
+    ).create_engine_config()
+    cfg.cache_config.kv_cache_layout = "LBHNC"
+    window = 128 if head_size == 256 else None
+    storage_dtype = torch.uint8 if cache_dtype == "fp8_e4m3" else DTYPE
+    spec_kwargs = dict(
+        block_size=block_size,
+        num_kv_heads=16,
+        head_size=head_size,
+        dtype=storage_dtype,
+        kv_quant_mode=get_kv_quant_mode(cache_dtype),
+    )
+    spec = (
+        SlidingWindowSpec(**spec_kwargs, sliding_window=window)
+        if window
+        else FullAttentionSpec(**spec_kwargs)
+    )
+    if reverse_children:
+        # Swapping the children and policy must preserve outputs and KV writes;
+        # dispatch cannot depend on a particular child's class or position.
+        class ReversedRouting(MMPrefixAttentionRouting):
+            capture_variant = 0
+
+            def select(self, metadata):
+                return 1 - super().select(metadata)
+
+            @staticmethod
+            def variant_uses_mm_prefix(variant):
+                return variant == 1
+
+        backend = create_composite_attention_backend(
+            FlashInferBackend,
+            TritonAttentionBackend,
+            name="ReversedCompositeBackend",
+            backend_name="CUSTOM",
+            module=__name__,
+            routing_policy=ReversedRouting,
+        )
+    with set_current_vllm_config(cfg), set_default_torch_dtype(DTYPE):
+        layer = Attention(
+            32,
+            head_size,
+            head_size**-0.5,
+            num_kv_heads=16,
+            cache_config=cfg.cache_config,
+            prefix="composite_test",
+            per_layer_sliding_window=window,
+            attn_backend=backend,
+        ).to(DEVICE)
+        layer.mm_prefix_clamp_sliding_window = window is not None
+        assert backend.get_builder_cls().get_cudagraph_support(cfg, spec) == (
+            AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
+        )
+        if window:
+            layer.mm_prefix_clamp_sliding_window = False
+            assert backend.get_builder_cls().get_cudagraph_support(cfg, spec) == (
+                AttentionCGSupport.NEVER
+            )
+            layer.mm_prefix_clamp_sliding_window = True
+        builder = backend.get_builder_cls()(spec, ["composite_test"], cfg, DEVICE)
+        builder.set_kernel_block_size(block_size)
+        # Match production, where unwritten page tails are zero-initialized.
+        cache = torch.zeros(
+            num_blocks,
+            16,
+            block_size,
+            2 * head_size,
+            dtype=storage_dtype,
+            device=DEVICE,
+        )
+        keys = torch.randn(209, 16, head_size, dtype=DTYPE, device=DEVICE)
+        values = torch.randn_like(keys)
+        if cache_dtype == "fp8_e4m3":
+            keys = keys.to(torch.float8_e4m3fn).to(DTYPE)
+            values = values.to(torch.float8_e4m3fn).to(DTYPE)
+        table = torch.arange(num_blocks, dtype=torch.int32, device=DEVICE).unsqueeze(0)
+        previous = 0
+        for end, spans, metadata_type in [
+            (128, [(32, 111)], TritonAttentionMetadata),
+            (129, [(32, 111)], causal_metadata_type),
+            (145, [(32, 111)], causal_metadata_type),
+            (209, [(32, 111), (160, 207)], TritonAttentionMetadata),
+        ]:
+            qlen = end - previous
+            starts = torch.tensor([0, qlen], dtype=torch.int32)
+            lengths = torch.tensor([end], dtype=torch.int32)
+            common = CommonAttentionMetadata(
+                query_start_loc=starts.to(DEVICE),
+                query_start_loc_cpu=starts,
+                seq_lens=lengths.to(DEVICE),
+                seq_lens_cpu_upper_bound=lengths,
+                num_reqs=1,
+                num_actual_tokens=qlen,
+                max_query_len=qlen,
+                max_seq_len=end,
+                block_table_tensor=table,
+                slot_mapping=torch.arange(previous, end, device=DEVICE),
+                mm_req_doc_ranges={0: spans},
+            )
+            metadata = builder.build(0, common)
+            assert isinstance(metadata, metadata_type)
+            if isinstance(metadata, FlashInferMetadata):
+                if qlen == 1:
+                    assert isinstance(metadata.decode, FlashInferTrtllmAPIDecode)
+                else:
+                    assert isinstance(metadata.prefill, TRTLLMPrefill)
+            query = torch.randn(qlen, 32, head_size, dtype=DTYPE, device=DEVICE)
+            if cache_dtype == "fp8_e4m3":
+                query = query.to(torch.float8_e4m3fn).to(DTYPE)
+            key, value = keys[previous:end], values[previous:end]
+            output = torch.empty_like(query)
+            layer.impl.do_kv_cache_update(layer, key, value, cache, common.slot_mapping)
+            layer.impl.forward(layer, query, key, value, cache, metadata, output)
+            expected = _dense_reference(
+                query,
+                [keys[:end]],
+                [values[:end]],
+                [qlen],
+                [end],
+                {0: spans},
+                window,
+                head_size**-0.5,
+                mm_clamp_sw=window or 0,
+            )
+            torch.testing.assert_close(output.float(), expected, atol=2e-2, rtol=2e-2)
+            if qlen == 1:
+                metadata = builder.build_for_cudagraph_capture(common)
+                assert isinstance(metadata, metadata_type)
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph):
+                    layer.impl.forward(
+                        layer, query, key, value, cache, metadata, output
+                    )
+                output.zero_()
+                graph.replay()
+                torch.testing.assert_close(
+                    output.float(), expected, atol=2e-2, rtol=2e-2
+                )
+            previous = end
diff --git a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
index 06684302cf39..e6390ce36ed7 100644
--- a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
+++ b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
@@ -151,6 +151,25 @@ def test_decode_and_verify_are_disjoint(kv_cache_dtype, num_heads):
         )
 
 
+@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
+def test_oversized_kv_cache_refuses_gluon(gluon_available, kv_cache_dtype):
+    """Past 2 GiB per layer Gluon drops its KV bounds mask and reads out of range.
+
+    ``None`` is the profiling case, before the cache is sized: nothing to
+    refuse yet, so the usual route stands.
+    """
+    bound = rocm_aiter_mla._GLUON_MAX_KV_CACHE_BYTES
+    for kv_cache_bytes in (None, bound):
+        assert AiterMLAHelper.use_gluon_decode(8, 1, kv_cache_dtype, kv_cache_bytes)
+        assert AiterMLAHelper.use_gluon_verify(
+            8, 8, kv_cache_dtype, kv_cache_bytes=kv_cache_bytes
+        )
+    assert not AiterMLAHelper.use_gluon_decode(8, 1, kv_cache_dtype, bound + 1)
+    assert not AiterMLAHelper.use_gluon_verify(
+        8, 8, kv_cache_dtype, kv_cache_bytes=bound + 1
+    )
+
+
 @pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15])
 def test_padded_query_is_contiguous(num_heads):
     """asm_mla.cu:805 requires Q.is_contiguous().
@@ -176,3 +195,8 @@ def test_pad_unpad_round_trip_preserves_head_order(num_heads):
 
     assert unpadded.shape == q.shape
     torch.testing.assert_close(unpadded, q)
+
+
+def test_a_non_causal_block_never_routes_to_gluon(gluon_available):
+    """A small-head block that normally uses Gluon must use ASM when non-causal."""
+    assert not AiterMLAHelper.use_gluon_verify(12, 8, "auto", causal=False)
diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
index 112cdff324df..1970dbb54fb9 100644
--- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
@@ -107,6 +107,8 @@ def _builder(
         _dcp_verify_buffers=None,
         _graph_seq_lens=None,
         _kv_cache_dtype_str=kv_cache_dtype,
+        # Unsized as during profiling, so the Gluon KV bound stays out of the way.
+        _kv_cache_bytes=None,
         paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32),
         paged_kv_indices=torch.empty(1024, dtype=torch.int32),
         paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32),
@@ -130,6 +132,7 @@ def _builder(
         _mla_q_dtype=torch.bfloat16,
         _mla_kv_dtype=torch.bfloat16,
         decode_attn_out_dtype=torch.bfloat16,
+        _decode_causal=True,
     )
     # Bound to the stub rather than faked: the verify flatten's per-row view is
     # part of what _build_decode is being tested for.
@@ -155,6 +158,25 @@ def test_backend_declares_uniform_batch_support():
     )
 
 
+@pytest.mark.parametrize("causal", [True, False])
+def test_build_hands_common_causal_flag_to_aiter_builder(monkeypatch, causal):
+    seen = []
+
+    def fake_common_build(self, *args):
+        seen.append(self._decode_causal)
+        return SimpleNamespace(decode=None, prefill=None)
+
+    monkeypatch.setattr(
+        rocm_aiter_mla.MLACommonMetadataBuilder, "build", fake_common_build
+    )
+    builder = object.__new__(AiterMLAMetadataBuilder)
+    builder._fp8_prefill_enabled = False
+
+    builder.build(0, SimpleNamespace(causal=causal))
+
+    assert seen == [causal]
+
+
 def test_dcp_verify_row_view_is_causal_per_row():
     """Row t must cover this rank's share of positions [0, seq_len - qlen + t].
 
@@ -258,6 +280,7 @@ def test_dcp_fp8_verify_build_uses_segmented(monkeypatch):
         query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
         query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
         num_decode_tokens=2 * qlen,
+        max_query_len=qlen,
         dcp_tot_seq_lens_device=torch.tensor([10, 12], dtype=torch.int32),
     )
 
@@ -494,7 +517,6 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(
     Sweeping num_heads asserts metadata is sized for the padded decode shape,
     covering Kimi-K3 TP4's 24 -> 32 head path and native fp8 nhead=32 folding.
     """
-
     dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16")
     info_calls = []
 
@@ -565,7 +587,7 @@ def init_common_builder(self, *args, **kwargs):
             get_num_attention_heads=lambda parallel_config: num_heads,
         ),
         scheduler_config=SimpleNamespace(max_num_seqs=2),
-        cache_config=SimpleNamespace(cache_dtype="fp8_e4m3"),
+        cache_config=SimpleNamespace(cache_dtype="fp8_e4m3", num_gpu_blocks=None),
         compilation_config=SimpleNamespace(
             cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False),
             # Empty: the per-layer head-count probe finds no attention layer and
@@ -614,6 +636,7 @@ def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch):
         query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32),
         query_start_loc_device=torch.tensor([0, 4, 8], dtype=torch.int32),
         num_decode_tokens=8,
+        max_query_len=4,
         dcp_tot_seq_lens_device=None,
     )
 
@@ -664,6 +687,7 @@ def test_min_kv_seq_len_ignores_cudagraph_padding_rows(monkeypatch):
         query_start_loc_cpu=query_start_loc,
         query_start_loc_device=query_start_loc,
         num_decode_tokens=num_reqs * mtp_qlen,
+        max_query_len=mtp_qlen,
         dcp_tot_seq_lens_device=None,
     )
 
@@ -674,7 +698,6 @@ def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr(
     monkeypatch,
 ):
     """Full-CG zero-qo rows follow rocm_aiter_mla.py:608-657,717-759."""
-
     get_mla_metadata_v1 = mock.MagicMock()
     monkeypatch.setitem(
         sys.modules,
@@ -715,6 +738,7 @@ def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr(
         query_start_loc_cpu=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32),
         query_start_loc_device=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32),
         num_decode_tokens=seq_lens.numel() * mtp_qlen,
+        max_query_len=mtp_qlen,
         dcp_tot_seq_lens_device=None,
     )
 
@@ -736,7 +760,6 @@ def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr(
 
 def test_decode_expands_kernel_block_page_indices(monkeypatch):
     """kernel_block_size>1 expands b -> b*K+offset at rocm_aiter_mla.py:696-704."""
-
     expand_kernel = _ExpandPageIndicesKernel()
     monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand_kernel)
     # qlen==1 now takes the persistent-metadata path, which imports
@@ -784,6 +807,7 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch):
         query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.int32),
         query_start_loc_device=torch.tensor([0, 1, 2], dtype=torch.int32),
         num_decode_tokens=seq_lens.numel(),
+        max_query_len=1,
         dcp_tot_seq_lens_device=None,
     )
 
@@ -875,6 +899,7 @@ def test_persistent_metadata_gate(
         query_start_loc_cpu=query_start_loc,
         query_start_loc_device=query_start_loc,
         num_decode_tokens=num_reqs * qo_len,
+        max_query_len=qo_len,
         dcp_tot_seq_lens_device=None,
     )
 
@@ -944,8 +969,88 @@ def test_persistent_metadata_gate_without_gluon_build(
         query_start_loc_cpu=query_start_loc,
         query_start_loc_device=query_start_loc,
         num_decode_tokens=num_reqs * qo_len,
+        max_query_len=qo_len,
         dcp_tot_seq_lens_device=None,
     )
 
     assert metadata.has_persistent_metadata is expect_persistent
     assert get_mla_metadata_v1.called is expect_persistent
+
+
+def _build_non_causal(monkeypatch, *, num_heads, kv_cache_dtype, qlen, mtp_qlen):
+    """Drive _build_decode with causal=False and hand back the aiter mock."""
+    get_mla_metadata_v1 = mock.MagicMock()
+    monkeypatch.setitem(
+        sys.modules,
+        "aiter",
+        SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
+    )
+    monkeypatch.setattr(
+        rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
+    )
+    builder = _builder(
+        mtp_decode_qlen=mtp_qlen,
+        kv_cache_dtype=kv_cache_dtype,
+        num_heads=num_heads,
+    )
+    builder._decode_causal = False
+    metadata = AiterMLAMetadataBuilder._build_decode(
+        builder,
+        block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
+        seq_lens_device=torch.tensor([7, 5], dtype=torch.int32),
+        max_seq_len=7,
+        query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
+        query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
+        num_decode_tokens=2 * qlen,
+        max_query_len=qlen,
+        dcp_tot_seq_lens_device=None,
+    )
+    return metadata, get_mla_metadata_v1
+
+
+def test_non_causal_build_hands_the_mask_to_aiter(monkeypatch):
+    """The schedule carries the mask, so is_causal must follow the block."""
+    metadata, get_mla_metadata_v1 = _build_non_causal(
+        monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=8, mtp_qlen=8
+    )
+    assert metadata.has_persistent_metadata
+    # 6th positional arg of get_mla_metadata_v1 is is_causal.
+    assert get_mla_metadata_v1.call_args.args[5] is False
+
+
+@pytest.mark.parametrize("num_heads", [12, 16, 48])
+def test_a_two_token_fp8_block_is_refused_without_a_fold(monkeypatch, num_heads):
+    """16-head (and padded-to-16) qlen-2 fp8 has no non-causal kernel.
+
+    48/80/112 (and 33 padded to 48) fold to H16 while keeping Q2, so they
+    must be refused the same way as native 16-head qlen 2.
+    """
+    with pytest.raises(ValueError, match="2-token"):
+        _build_non_causal(
+            monkeypatch,
+            num_heads=num_heads,
+            kv_cache_dtype="fp8",
+            qlen=2,
+            mtp_qlen=8,
+        )
+
+
+def test_a_two_token_fp8_block_is_allowed_when_folded(monkeypatch):
+    """32 heads at qlen 2 folds onto the 16-head / 4-token kernel."""
+    metadata, get_mla_metadata_v1 = _build_non_causal(
+        monkeypatch,
+        num_heads=32,
+        kv_cache_dtype="fp8",
+        qlen=2,
+        mtp_qlen=8,
+    )
+    assert metadata.has_persistent_metadata
+    assert get_mla_metadata_v1.called
+
+
+def test_a_two_token_bf16_block_is_allowed(monkeypatch):
+    """The bf16 fold has no such hole, so the guard must not fire there."""
+    metadata, _ = _build_non_causal(
+        monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8
+    )
+    assert metadata.has_persistent_metadata
diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py
index 8668fa26bcc8..b99a33f9f7a9 100644
--- a/tests/v1/attention/test_rocm_attention_backends_selection.py
+++ b/tests/v1/attention/test_rocm_attention_backends_selection.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for attention backend selectors."""
 
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -34,6 +35,15 @@ def mock_get_cdna_version():
         yield
 
 
+@pytest.fixture
+def cleared_attention_selector_cache():
+    from vllm.v1.attention.selector import _cached_get_attn_backend
+
+    _cached_get_attn_backend.cache_clear()
+    yield
+    _cached_get_attn_backend.cache_clear()
+
+
 def test_aiter_unified_attention_uses_dedicated_metadata_builder():
     from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
         RocmAiterUnifiedAttentionBackend,
@@ -434,6 +444,139 @@ def test_aiter_fa_requires_mi3xx(mock_vllm_config):
         )
 
 
+@pytest.fixture
+def turboquant_run_config():
+    """Current config of a run whose KV cache dtype is a turboquant_* preset."""
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(cache_dtype="turboquant_k8v4")
+    )
+    with patch(
+        "vllm.config.get_current_vllm_config_or_none",
+        return_value=config,
+    ):
+        yield config
+
+
+def test_turboquant_boundary_selection_is_not_cached_from_ordinary_run(
+    cleared_attention_selector_cache,
+):
+    from vllm.config import CacheConfig, VllmConfig, set_current_vllm_config
+    from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
+    from vllm.v1.attention.selector import get_attn_backend
+    from vllm.v1.kv_cache_layout import KVCacheLayout
+
+    ordinary_config = VllmConfig(cache_config=CacheConfig(cache_dtype="auto"))
+    turboquant_config = VllmConfig(
+        cache_config=CacheConfig(cache_dtype="turboquant_k8v4")
+    )
+
+    backend_priorities = [
+        AttentionBackendEnum.ROCM_ATTN,
+        AttentionBackendEnum.TRITON_ATTN,
+        AttentionBackendEnum.TURBOQUANT,
+    ]
+    with patch(
+        "vllm.platforms.rocm._get_backend_priorities",
+        return_value=backend_priorities,
+    ):
+        with set_current_vllm_config(ordinary_config):
+            ordinary_backend = get_attn_backend(128, torch.float16, "auto")
+
+        with set_current_vllm_config(turboquant_config):
+            boundary_backend = get_attn_backend(128, torch.float16, "auto")
+            quantized_backend = get_attn_backend(128, torch.float16, "turboquant_k8v4")
+
+    assert ordinary_backend is AttentionBackendEnum.ROCM_ATTN.get_class()
+    assert quantized_backend is AttentionBackendEnum.TURBOQUANT.get_class()
+    assert get_supported_kv_cache_layouts([boundary_backend, quantized_backend]) == [
+        KVCacheLayout.LBNHC
+    ]
+
+
+def test_turboquant_run_does_not_mask_unsupported_selected_backend(
+    turboquant_run_config,
+):
+    from vllm.platforms.rocm import RocmPlatform
+
+    attn_selector_config = AttentionSelectorConfig(
+        head_size=128,
+        dtype=torch.float16,
+        kv_cache_dtype="auto",
+        block_size=16,
+    )
+
+    with (
+        patch("vllm.platforms.rocm.get_cdna_version", return_value=1),
+        pytest.raises(ValueError, match="compute capability not supported"),
+    ):
+        RocmPlatform.get_attn_backend_cls(
+            selected_backend=AttentionBackendEnum.ROCM_AITER_FA,
+            attn_selector_config=attn_selector_config,
+        )
+
+
+def test_turboquant_layout_check_respects_backend_override():
+    from vllm.platforms.rocm import _shares_layout_with_turboquant
+    from vllm.v1.kv_cache_layout import KVCacheLayout
+
+    boundary_backend = MagicMock()
+    boundary_backend.supported_kv_cache_layouts.return_value = (KVCacheLayout.LBHNC,)
+    turboquant_override = MagicMock()
+    turboquant_override.supported_kv_cache_layouts.return_value = (KVCacheLayout.LBHNC,)
+
+    with patch.object(
+        AttentionBackendEnum.TURBOQUANT,
+        "get_class",
+        return_value=turboquant_override,
+    ) as get_turboquant_class:
+        assert _shares_layout_with_turboquant(boundary_backend)
+    get_turboquant_class.assert_called_once_with()
+
+
+@pytest.mark.parametrize(
+    "selected_backend",
+    [None, "ROCM_ATTN", "ROCM_AITER_FA", "TURBOQUANT"],
+)
+def test_turboquant_boundary_layers_share_a_layout(
+    selected_backend,
+    turboquant_run_config,
+    mock_get_cdna_version,
+):
+    """A turboquant_* run must resolve to backends with a common KV layout.
+
+    The boundary layers keep the native dtype and pick their own backend while
+    every other layer picks TURBOQUANT. Engine startup hard-errors when the two
+    share no layout, so the boundary layers must not land on a native ROCm or
+    AITER backend, whichever backend the run asked for.
+    """
+    from vllm.platforms.rocm import RocmPlatform
+    from vllm.utils.import_utils import resolve_obj_by_qualname
+    from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
+
+    backend_enum = (
+        getattr(AttentionBackendEnum, selected_backend) if selected_backend else None
+    )
+
+    def resolve(kv_cache_dtype):
+        path = RocmPlatform.get_attn_backend_cls(
+            selected_backend=backend_enum,
+            attn_selector_config=AttentionSelectorConfig(
+                head_size=128,
+                dtype=torch.float16,
+                kv_cache_dtype=kv_cache_dtype,
+                block_size=16,
+            ),
+        )
+        return resolve_obj_by_qualname(path)
+
+    boundary_backend = resolve("auto")
+    quantized_backend = resolve("turboquant_k8v4")
+
+    assert quantized_backend is AttentionBackendEnum.TURBOQUANT.get_class()
+    # Raises when the intersection is empty, which is the startup failure.
+    assert get_supported_kv_cache_layouts([boundary_backend, quantized_backend])
+
+
 def test_sparse_not_supported(mock_vllm_config):
     """Test that sparse MLA without use_mla flag raises an error."""
     from vllm.platforms.rocm import RocmPlatform
@@ -455,3 +598,115 @@ def test_sparse_not_supported(mock_vllm_config):
         RocmPlatform.get_attn_backend_cls(
             selected_backend=None, attn_selector_config=attn_selector_config
         )
+
+
+def _kv_connector_selector_config() -> AttentionSelectorConfig:
+    return AttentionSelectorConfig(
+        head_size=128,
+        dtype=torch.float16,
+        kv_cache_dtype="auto",
+        block_size=16,
+        use_mla=False,
+        has_sink=False,
+        use_sparse=False,
+        use_kv_connector=True,
+    )
+
+
+def test_unified_attn_declares_kv_connector_support():
+    """ROCM_AITER_UNIFIED_ATTN opts into KV connectors and ROCM_ATTN does not."""
+    from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
+        RocmAiterUnifiedAttentionBackend,
+    )
+    from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
+
+    assert RocmAiterUnifiedAttentionBackend.supports_kv_connector() is True
+    assert RocmAttentionBackend.supports_kv_connector() is False
+
+
+def test_unified_attn_supports_kv_connector(mock_vllm_config, mock_get_cdna_version):
+    """ROCM_AITER_UNIFIED_ATTN can be selected with KV connectors."""
+    from vllm.platforms.rocm import RocmPlatform
+
+    backend_path = RocmPlatform.get_attn_backend_cls(
+        selected_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        attn_selector_config=_kv_connector_selector_config(),
+    )
+
+    assert backend_path == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
+
+
+def test_rocm_attn_rejects_kv_connector(mock_vllm_config, mock_get_cdna_version):
+    """Selecting ROCM_ATTN with a KV connector is illegal."""
+    from vllm.platforms.rocm import RocmPlatform
+
+    attn_selector_config = _kv_connector_selector_config()
+
+    with pytest.raises(ValueError, match="KV connector not supported"):
+        RocmPlatform.get_attn_backend_cls(
+            selected_backend=AttentionBackendEnum.ROCM_ATTN,
+            attn_selector_config=attn_selector_config,
+        )
+
+
+@pytest.mark.parametrize(
+    "aiter_found, expected_backend",
+    [
+        (True, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN),
+        (False, AttentionBackendEnum.TRITON_ATTN),
+    ],
+)
+def test_auto_selection_for_kv_connector(
+    aiter_found, expected_backend, mock_vllm_config, mock_get_cdna_version
+):
+    """Auto-selection with a KV connector and AITER enabled resolves to unified attn,
+    and to triton attn if AITER not enabled."""
+    from vllm.platforms.rocm import RocmPlatform
+
+    with patch(
+        "vllm._aiter_ops.is_aiter_found_and_supported", return_value=aiter_found
+    ):
+        backend_path = RocmPlatform.get_attn_backend_cls(
+            selected_backend=None,
+            attn_selector_config=_kv_connector_selector_config(),
+        )
+
+    assert backend_path == expected_backend.get_path()
+
+
+def test_unified_attn_prefers_block_contiguous_layout():
+    """Unified attn prefers a block-first KV layout, hence ok with kv connectors."""
+    from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
+        RocmAiterUnifiedAttentionBackend,
+    )
+    from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
+
+    unified_preferred = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()[0]
+    rocm_attn_preferred = RocmAttentionBackend.supported_kv_cache_layouts()[0]
+
+    assert unified_preferred.is_block_contiguous is True
+    assert rocm_attn_preferred.is_block_contiguous is False
+
+
+def test_unified_attn_drops_lhbnc_with_kv_connector():
+    """Connectors move a block as one contiguous byte range, which LHBNC breaks."""
+    from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
+    from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
+        RocmAiterUnifiedAttentionBackend,
+    )
+    from vllm.v1.kv_cache_interface import KVCacheLayout
+
+    assert KVCacheLayout.LHBNC in (
+        RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()
+    )
+
+    config = VllmConfig(
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="ExampleConnector", kv_role="kv_both"
+        )
+    )
+    with set_current_vllm_config(config):
+        layouts = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()
+
+    assert layouts
+    assert all(layout.is_block_compact for layout in layouts)
diff --git a/tests/v1/attention/test_rocm_glm5next_sparse.py b/tests/v1/attention/test_rocm_glm5next_sparse.py
index 509d8f1faea4..ccc9e5c6178e 100644
--- a/tests/v1/attention/test_rocm_glm5next_sparse.py
+++ b/tests/v1/attention/test_rocm_glm5next_sparse.py
@@ -76,7 +76,9 @@ def test_fit_kpool_indices_rejects_narrow_input():
         ("auto", 512, 0, 2, 2, 1, True),
         ("fp8", 512, 1, 0, 0, 32, False),
         ("auto", 576, 1, 0, 0, 32, False),
-        ("auto", 512, 0, 2, 4, 2, False),
+        ("auto", 512, 0, 2, 4, 2, True),
+        ("auto", 512, 0, 2, 12, 6, True),
+        ("auto", 512, 0, 0, 0, 0, False),
     ],
 )
 def test_rocm_sparse_triton_route(
@@ -88,6 +90,7 @@ def test_rocm_sparse_triton_route(
     max_query_len,
     expected,
 ):
+    """Validate Triton routing for prefill, decode, and MTP verification."""
     assert (
         _use_rocm_sparse_triton(
             kv_cache_dtype=kv_cache_dtype,
diff --git a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
index 94fdadd6e95d..d87a82c37b2b 100644
--- a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
+++ b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
@@ -20,10 +20,10 @@
 
 import torch
 
-# Bootstrap the glm5next package before entering the indexer module: its
-# kpool_compress import runs glm5next/__init__, which pulls model ->
-# attention -> back into sparse_attn_indexer_kpool (attention.py imports the
-# class at module scope). Production always enters via attention.py first.
+# Bootstrap the glm5next package before importing kpool_compress directly: its
+# import runs glm5next/__init__, which pulls model -> attention -> the indexer
+# dispatcher (sparse_indexer) -> back into the nvidia/amd sparse_indexer
+# modules. Production always enters via attention.py first.
 # isort: off
 import vllm.models.glm5next  # noqa: F401
 from vllm.models.glm5next.nvidia.ops.kpool_compress import (  # noqa: E402
@@ -32,13 +32,17 @@
 )
 # isort: on
 
-import vllm.model_executor.layers.sparse_attn_indexer_kpool as indexer_mod
-from vllm.model_executor.layers.sparse_attn_indexer_kpool import (
+from vllm.models.glm5next.common.sparse_indexer import (
     _decode_topk_seq_lens,
     _fill_short_decode_causal_indices,
 )
 from vllm.platforms import current_platform
 
+if current_platform.is_rocm():
+    from vllm.models.glm5next.amd import sparse_indexer as indexer_mod
+else:
+    from vllm.models.glm5next.nvidia import sparse_indexer as indexer_mod
+
 KPOOL = 4
 TOPK_TOKENS = 16
 SELECT_K = TOPK_TOKENS // KPOOL
diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py
index 002aa5a4b72c..1d7fc389ef69 100644
--- a/tests/v1/attention/test_sparse_mla_backends.py
+++ b/tests/v1/attention/test_sparse_mla_backends.py
@@ -3,7 +3,9 @@
 """Unit tests for the sparse MLA backends and utilities."""
 
 import math
+from collections import deque
 from types import MethodType, SimpleNamespace
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -20,12 +22,27 @@
     create_vllm_config,
 )
 from vllm import _custom_ops as ops
-from vllm.config import set_current_vllm_config
+from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+from vllm.config import (
+    CUDAGraphMode,
+    HiSparseConfig,
+    SpeculativeConfig,
+    set_current_vllm_config,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
+    HiSparseConnectorWorker,
+)
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.attention import mla_attention
 from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha
 from vllm.model_executor.layers.attention.sparse_mla_attention import (
     GLOBAL_TOPK_MASK_MAX_BYTES,
+    SparseMLACommonImpl,
+    SparseMLAPrefillMetadata,
+    _is_masked_mha_available,
     _masked_mha_workspace_fits,
     _topk_mask_shape,
+    _use_dense_mha_prefill,
 )
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.platforms import current_platform
@@ -40,15 +57,26 @@
         allow_module_level=True,
     )
 
-import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as flashinfer_sparse_mod
 from vllm.model_executor.layers.attention.mla_attention import (
     _canonicalize_sparse_mla_kv_cache_dtype,
 )
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SharedTopkIndicesBuffer,
+)
 from vllm.utils.math_utils import cdiv
+from vllm.utils.torch_utils import current_stream
+from vllm.v1.attention.backends.mla import index_group as index_group_module
+from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
+    FlashAttnMLASparseImpl,
+)
 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
     FlashInferMLASparseImpl,
+    FlashInferMLASparseMetadataBuilder,
     FlashInferMLASparseTRTLLMBackend,
 )
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
+    FlashInferMLASparseSM120Impl,
+)
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
     FlashMLASparseBackend,
     FlashMLASparseImpl,
@@ -56,12 +84,29 @@
     FlashMLASparseMetadataBuilder,
     triton_convert_req_index_to_global_index,
 )
-from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder
+from vllm.v1.attention.backends.mla.index_group import (
+    HiSparseMLAIndexGroup,
+    SparseMLAIndexGroupBuilder,
+)
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV32IndexerMetadataBuilder,
+)
+from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
 from vllm.v1.attention.backends.utils import (
     split_decodes_and_prefills,
     split_prefill_chunks,
 )
 from vllm.v1.attention.ops import flashmla
+from vllm.v1.hisparse import runtime as hisparse_runtime
+from vllm.v1.hisparse.runtime import (
+    HiSparseCacheHandle,
+    HiSparseRuntime,
+    ResolvedHiSparseConfig,
+    _has_hisparse_ops,
+    build_hisparse_prefill_staging_plan,
+    hisparse_prefill_staging_remap,
+)
+from vllm.v1.hisparse.types import SparseKVRowMirror
 
 SPARSE_BACKEND_BATCH_SPECS = {
     name: BATCH_SPECS[name]
@@ -109,10 +154,11 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch):
     impl.bmm2_scale = None
     impl.is_nope_mla = True
     impl.need_to_return_lse_for_decode = False
-    monkeypatch.setattr(
-        flashinfer_sparse_mod,
-        "triton_convert_req_index_to_global_index",
-        lambda *args, **kwargs: (topk, torch.ones(1, dtype=torch.int32)),
+    impl.index_group = MagicMock()
+    impl.index_group_index = 0
+    impl.index_group.convert_logical_to_physical_topk.return_value = (
+        topk,
+        torch.ones(1, dtype=torch.int32),
     )
 
     import flashinfer.decode
@@ -138,6 +184,41 @@ def fake_flashinfer(**kwargs):
     assert recorded_scale != kv_lora_rank**-0.5
 
 
+def test_hisparse_routes_prefill_to_sparse_mqa():
+    layer = SimpleNamespace(hisparse_cache=object())
+
+    assert not mla_attention.MLAAttention._use_sparse_mha(layer, SimpleNamespace())
+
+
+def test_hisparse_metadata_keeps_short_prefill_indexer_enabled():
+    config = SimpleNamespace(
+        attention_config=SimpleNamespace(
+            hisparse_config=object(), sparse_mla_force_mqa=False
+        )
+    )
+
+    assert not _use_dense_mha_prefill(config, prefill_max_seq_len=512, topk_tokens=2048)
+
+
+def test_unified_mla_prepares_batch_without_slot_mapping(monkeypatch):
+    layer = SimpleNamespace(
+        prepare_kv_cache_update=MagicMock(),
+        update_kv_cache=MagicMock(),
+    )
+    metadata = object()
+    monkeypatch.setattr(
+        mla_attention,
+        "get_attention_context",
+        lambda _: (metadata, layer, torch.empty(0), None),
+    )
+    tensor = torch.empty(0, device=DEVICE_TYPE)
+
+    mla_attention.unified_mla_kv_cache_update(tensor, tensor, "layer", "auto", tensor)
+
+    layer.prepare_kv_cache_update.assert_called_once_with(metadata)
+    layer.update_kv_cache.assert_not_called()
+
+
 def _float_to_e8m0_truncate(f: float) -> float:
     """Simulate SM100's float -> e8m0 -> bf16 scale conversion.
     e8m0 format only stores the exponent (power of 2).
@@ -164,8 +245,8 @@ def _dequantize_fp8_ds_mla_entry(
     Args:
         simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's
             float -> e8m0 -> bf16 scale conversion path.
-    """
 
+    """
     # The first kv_lora_rank bytes store FP8 latent values with one scale per
     # 128 element tile written as float32 right after the latent payload.
     scales = cache_slice.view(torch.float32)[kv_lora_rank // 4 : kv_lora_rank // 4 + 4]
@@ -202,8 +283,8 @@ def _quantize_dequantize_fp8_ds_mla(
     Args:
         simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's
             float -> e8m0 -> bf16 scale conversion in dequantization.
-    """
 
+    """
     if kv_c.numel() == 0:
         return kv_c.clone(), k_pe.clone()
 
@@ -340,7 +421,18 @@ def test_sparse_backend_decode_correctness(
     workspace_init,
     q_scale: float,
     k_scale: float,
+    monkeypatch,
 ):
+    if (
+        batch_name == "large_q_pure_prefill"
+        and backend_cls == FlashMLASparseBackend
+        and kv_cache_dtype == "fp8_ds_mla"
+        and tensor_parallel_size == 4
+    ):
+        monkeypatch.setattr(
+            "vllm.v1.attention.backends.mla.flashmla_sparse.split_prefill_chunks",
+            lambda rows, capacity: [(i, i + 1) for i in range(len(rows))],
+        )
     if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
         pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
 
@@ -419,6 +511,7 @@ def test_sparse_backend_decode_correctness(
     )
     model_config = vllm_config.model_config
     model_config.hf_text_config = SimpleNamespace(
+        index_topk=topk_tokens,
         q_lora_rank=None,
         kv_lora_rank=kv_lora_rank,
         qk_nope_head_dim=qk_nope_head_dim,
@@ -426,6 +519,7 @@ def test_sparse_backend_decode_correctness(
         v_head_dim=v_head_dim,
         model_type="deepseek_v2",
     )
+    del model_config.hf_config.index_topk  # Composite configs only nest this field.
     model_config.dtype = dtype
     model_config.get_num_attention_heads = MethodType(
         lambda self, parallel_config: num_heads,
@@ -587,7 +681,7 @@ def test_sparse_backend_decode_correctness(
     sdpa_reference = torch.cat(reference_outputs, dim=0)
 
     vllm_config.cache_config.cache_dtype = kv_cache_dtype
-    vllm_config.model_config.hf_config.index_topk = topk_tokens
+    vllm_config.model_config.hf_text_config.index_topk = topk_tokens
 
     common_attn_metadata = create_common_attn_metadata(
         batch_spec,
@@ -612,8 +706,6 @@ def test_sparse_backend_decode_correctness(
 
     # The sparse builder clones the layer's dense-MHA prefill backend from
     # static_forward_context; register a mock layer carrying one.
-    from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
-
     prefill_backend = get_mla_prefill_backend(vllm_config)(
         num_heads=num_heads,
         scale=scale,
@@ -693,6 +785,10 @@ def test_sparse_backend_decode_correctness(
     )
 
     with torch.inference_mode():
+        if backend_cls == FlashMLASparseBackend and kv_cache_dtype == "fp8_ds_mla":
+            from vllm.v1.worker.workspace import current_workspace_manager
+
+            current_workspace_manager().lock()
         backend_output = mock_layer.forward_impl(
             query_vllm,
             kv_c_vllm,
@@ -823,6 +919,75 @@ def test_triton_convert_req_index_to_global_index_decode_only(
     torch.testing.assert_close(result, reference_result, rtol=0, atol=0)
 
 
+def test_index_group_convert_during_piecewise_capture():
+    """In piecewise cudagraph mode the indexer is captured while the convert
+    runs in the following eager break. The convert's side stream must not
+    wait on ``logical_topk_ready`` there: an event recorded inside a captured
+    segment is graph-local, and an eager wait on it raises
+    cudaErrorInvalidValue."""
+    device = torch.device(DEVICE_TYPE)
+    num_tokens, num_topk, num_requests, blocks_per_req, block_size = 8, 128, 4, 4, 16
+
+    logical_topk_indices = torch.randint(
+        0,
+        block_size * blocks_per_req,
+        (num_tokens, num_topk),
+        dtype=torch.int32,
+        device=device,
+    )
+    builder = SparseMLAIndexGroupBuilder(logical_topk_indices)
+    group, layer_index = builder.register_layer(is_index_producing_layer=True)
+    assert layer_index == 0 and group.has_indexer
+
+    req_id = torch.arange(num_tokens, dtype=torch.int32, device=device) % num_requests
+    block_table = torch.arange(
+        num_requests * blocks_per_req, dtype=torch.int32, device=device
+    ).view(num_requests, blocks_per_req)
+    attn_metadata = SimpleNamespace(
+        req_id_per_token=req_id, block_table=block_table, block_size=block_size
+    )
+    expected = triton_convert_req_index_to_global_index(
+        req_id,
+        block_table,
+        logical_topk_indices,
+        BLOCK_SIZE=block_size,
+        NUM_TOPK_TOKENS=num_topk,
+    )
+
+    vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B")
+    marker = torch.zeros(1, device=device)
+    capture_stream = torch.cuda.Stream(device=device)
+    # torch.cuda.stream routes through vLLM's patched torch.cuda.set_stream,
+    # so current_stream() tracks the capture stream as in production.
+    with (
+        torch.cuda.stream(capture_stream),
+        set_forward_context(
+            None, vllm_config, cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE
+        ),
+    ):
+        capture = BreakableCUDAGraphCapture()
+        with capture:
+            marker.add_(1)
+            # Record inside a captured segment, convert in the eager
+            # break — mirrors mla.py -> unified_mla_attention_with_output.
+            group.set_logical_topk_ready(layer_index)
+            capture.add_eager(
+                lambda: group.convert_logical_to_physical_topk(
+                    layer_index,
+                    logical_topk_indices,
+                    attn_metadata,
+                    block_stride_rows=None,
+                    return_valid_counts=False,
+                )
+            )
+            marker.add_(1)
+        capture.replay()
+        capture_stream.synchronize()
+
+    result = group.physical_topk_indices[:num_tokens]
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
 @pytest.mark.parametrize("block_size", [16])
 @pytest.mark.skipif(
     torch.cuda.get_device_capability() < (9, 0),
@@ -966,6 +1131,8 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
         ),
         block_table=torch.arange(40, dtype=torch.int32, device=device).view(4, 10),
         block_size=block_size,
+        num_decode_tokens=num_mqa_tokens,
+        num_actual_tokens=num_mqa_tokens,
     )
     assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens
 
@@ -984,11 +1151,31 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
     def _stub_kernel(q, kv, indices, lengths, actual_num_heads):
         captured["indices"] = indices
         captured["actual_num_heads"] = actual_num_heads
-        return torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device)
+        return (
+            torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device),
+            None,
+        )
+
+    def _convert_topk(indices, metadata, **kwargs):
+        return triton_convert_req_index_to_global_index(
+            metadata.req_id_per_token[: indices.shape[0]],
+            metadata.block_table,
+            indices,
+            BLOCK_SIZE=metadata.block_size,
+            BLOCK_STRIDE_ROWS=kwargs["block_stride_rows"],
+            NUM_TOPK_TOKENS=indices.shape[1],
+            return_valid_counts=kwargs["return_valid_counts"],
+        )
 
-    stub_impl = SimpleNamespace(_bf16_flash_mla_kernel=_stub_kernel)
+    stub_impl = SimpleNamespace(
+        _bf16_flash_mla_kernel=_stub_kernel,
+        _convert_logical_to_physical_topk=_convert_topk,
+        index_group=None,
+        index_group_index=0,
+        pcp_dcp_kv_gather=False,
+    )
 
-    out = FlashMLASparseImpl._forward_bf16_kv(
+    out, _ = FlashMLASparseImpl._forward_bf16_kv(
         stub_impl, q, kv_cache, topk_indices, attn_metadata, q.shape[1]
     )
 
@@ -1070,6 +1257,33 @@ def test_masked_mha_workspace_guards_long_routing_policy(
     )
 
 
+@pytest.mark.parametrize(
+    ("model_dims", "kv_cache_dtype", "fa_version", "expected"),
+    [
+        pytest.param((128, 512, 128, 64, 128), "auto", 4, True, id="deepseek_v32"),
+        pytest.param((64, 512, 192, 64, 256), "auto", 4, True, id="glm5"),
+        pytest.param((64, 512, 256, 0, 256), "auto", 4, True, id="glm53_flash_nope"),
+        pytest.param((64, 512, 256, 0, 256), "fp8", 4, False, id="quantized_kv"),
+        pytest.param((64, 512, 256, 0, 256), "auto", 3, False, id="no_fa4"),
+        pytest.param((64, 512, 256, 64, 256), "auto", 4, False, id="rope_320"),
+        pytest.param((128, 512, 256, 0, 256), "auto", 4, False, id="wrong_heads"),
+    ],
+)
+def test_is_masked_mha_available_model_dims(
+    monkeypatch, model_dims, kv_cache_dtype, fa_version, expected
+):
+    """The allow-list gates masked MHA per exact model geometry: the DeepSeek-V3.2,
+    GLM-5 and NoPE GLM-5.3-Flash layouts on an SM100-family GPU with FA4 and an
+    unquantized KV cache, nothing else."""
+    import vllm.model_executor.layers.attention.sparse_mla_attention as mod
+
+    monkeypatch.setattr(
+        mod.current_platform, "is_device_capability_family", lambda family: True
+    )
+    monkeypatch.setattr(mod, "get_flash_attn_version", lambda **kwargs: fa_version)
+    assert _is_masked_mha_available(*model_dims, kv_cache_dtype) is expected
+
+
 def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
     """Request count and context chunk length are independent multipliers."""
     base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048)
@@ -1103,6 +1317,7 @@ def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
     [
         pytest.param(128, 128, 64, 128, id="deepseek_hd192_v128"),
         pytest.param(64, 192, 64, 256, id="glm5_hd256_v256"),
+        pytest.param(64, 256, 0, 256, id="glm53_flash_nope_hd256_v256"),
     ],
 )
 def test_sparse_backend_prefill_correctness(
@@ -1145,6 +1360,7 @@ def test_sparse_backend_prefill_correctness(
     )
     model_config = vllm_config.model_config
     model_config.hf_text_config = SimpleNamespace(
+        index_topk=topk_tokens,
         q_lora_rank=None,
         kv_lora_rank=kv_lora_rank,
         qk_nope_head_dim=qk_nope_head_dim,
@@ -1152,6 +1368,7 @@ def test_sparse_backend_prefill_correctness(
         v_head_dim=v_head_dim,
         model_type="deepseek_v2",
     )
+    del model_config.hf_config.index_topk  # Composite configs only nest this field.
     model_config.dtype = dtype
     model_config.model_arch_config.total_num_attention_heads = num_heads
     model_config.get_num_attention_heads = MethodType(
@@ -1251,7 +1468,7 @@ def test_sparse_backend_prefill_correctness(
     ref_output = torch.cat(reference_outputs, dim=0)
 
     vllm_config.cache_config.cache_dtype = kv_cache_dtype
-    vllm_config.model_config.hf_config.index_topk = topk_tokens
+    vllm_config.model_config.hf_text_config.index_topk = topk_tokens
 
     common_attn_metadata = create_common_attn_metadata(
         batch_spec,
@@ -1275,8 +1492,6 @@ def test_sparse_backend_prefill_correctness(
 
     # The sparse builder clones the layer's dense-MHA prefill backend from
     # static_forward_context; register a mock layer carrying one.
-    from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
-
     prefill_backend = get_mla_prefill_backend(vllm_config)(
         num_heads=num_heads,
         scale=scale,
@@ -1448,10 +1663,12 @@ def test_split_indexer_prefill_chunks_single_request_overflow():
     assert out == expected
 
 
-# 384 is not a power of two, so it counts via the tiled atomic accumulation
-# rather than the single-tile path 128 takes.
-@pytest.mark.parametrize("num_topk_tokens", [128, 384])
-def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
+# Power-of-two, GLM's padded tile, and atomic fallback, with reused buffers.
+@pytest.mark.parametrize(
+    "num_topk_tokens,reuse_buffers",
+    [(128, False), (2176, False), (2176, True), (4224, False), (4224, True)],
+)
+def test_triton_convert_returns_valid_counts(num_topk_tokens: int, reuse_buffers: bool):
     """Test that return_valid_counts correctly counts non-negative indices."""
     device = torch.device(DEVICE_TYPE)
     num_tokens = 8
@@ -1464,8 +1681,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
         num_requests * max_blocks_per_req, dtype=torch.int32, device=device
     ).view(num_requests, max_blocks_per_req)
 
-    # Create token indices with varying numbers of valid entries: half the row,
-    # a quarter of it, the whole row, then a single valid entry -- twice over.
+    # Include partly valid, fully valid, and empty selections for each request.
     token_indices = torch.full(
         (num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device
     )
@@ -1473,7 +1689,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
         num_topk_tokens // 2,
         num_topk_tokens // 4,
         num_topk_tokens,
-        1,
+        0,
     ] * 2
     expected_valid = []
     for i in range(num_tokens):
@@ -1487,6 +1703,15 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
         expected_valid, dtype=torch.int32, device=device
     )
 
+    buffers = {}
+    if reuse_buffers:
+        buffers = dict(
+            out=torch.full_like(token_indices, -99),
+            valid_counts_out=torch.full(
+                (num_tokens,), 99, dtype=torch.int32, device=device
+            ),
+        )
+
     # Test with return_valid_counts=True
     result, valid_counts = triton_convert_req_index_to_global_index(
         req_id,
@@ -1495,6 +1720,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
         BLOCK_SIZE=block_size,
         NUM_TOPK_TOKENS=num_topk_tokens,
         return_valid_counts=True,
+        **buffers,
     )
 
     torch.testing.assert_close(valid_counts, expected_valid_tensor, rtol=0, atol=0)
@@ -1516,143 +1742,1719 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
         assert torch.all(result[row, num_valid:] == -1)
 
 
-def test_flashmla_cache_dtype_aliases_use_ds_layout():
-    # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer
-    # stores kv_cache_dtype, so they cannot bypass the gate.
-    for alias in ("fp8", "fp8_e4m3"):
-        assert (
-            _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias)
-            == "fp8_ds_mla"
-        )
+def test_sparse_mla_index_group_converts_decode_indices_once(monkeypatch):
+    device = torch.device(DEVICE_TYPE)
+    logical = torch.full((2, 128), -1, dtype=torch.int32, device=device)
+    logical[:, :4] = torch.tensor(
+        [[0, 2, 5, -1], [1, 3, 7, -1]], dtype=torch.int32, device=device
+    )
+    builder = SparseMLAIndexGroupBuilder(logical)
+    group, leader_index = builder.register_layer(True)
+    follower_group, follower_index = builder.register_layer(False)
+    assert follower_group is group
+
+    block_table = torch.tensor([[2, 3], [4, 5]], dtype=torch.int32, device=device)
+    metadata = SimpleNamespace(
+        block_table=block_table,
+        block_size=4,
+        req_id_per_token=torch.tensor([0, 1], dtype=torch.int32, device=device),
+        num_decode_tokens=2,
+        num_actual_tokens=2,
+    )
+    calls = []
+    convert = index_group_module.triton_convert_req_index_to_global_index
 
+    def convert_spy(*args, **kwargs):
+        calls.append((args, kwargs))
+        return convert(*args, **kwargs)
 
-def test_flashmla_fp8_metadata_reuses_common_batch_split():
-    builder = SimpleNamespace(
-        device=torch.device(DEVICE_TYPE),
-        vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=8)),
+    monkeypatch.setattr(
+        index_group_module,
+        "triton_convert_req_index_to_global_index",
+        convert_spy,
     )
-    common_metadata = SimpleNamespace(
-        num_actual_tokens=1,
-        seq_lens_cpu_upper_bound=torch.tensor([1]),
-        query_start_loc_cpu=torch.tensor([0, 1]),
-        block_table_tensor=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
+    group.set_logical_topk_ready(leader_index)
+    leader_result = group.convert_logical_to_physical_topk(
+        leader_index,
+        logical,
+        metadata,
+        block_stride_rows=None,
+        return_valid_counts=True,
     )
-    metadata = FlashMLASparseMetadata(
-        num_reqs=1,
-        max_query_len=1,
-        max_seq_len=1,
-        num_actual_tokens=1,
-        query_start_loc=torch.tensor([0, 1], device=DEVICE_TYPE),
-        slot_mapping=torch.tensor([0], device=DEVICE_TYPE),
-        block_table=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
-        req_id_per_token=torch.zeros(1, dtype=torch.int32, device=DEVICE_TYPE),
-        num_decodes=0,
-        num_prefills=1,
-        num_decode_tokens=0,
+    follower_result = group.convert_logical_to_physical_topk(
+        follower_index,
+        logical,
+        metadata,
+        block_stride_rows=None,
+        return_valid_counts=True,
     )
+    torch.accelerator.synchronize()
 
-    fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(
-        builder, common_metadata, metadata
+    assert len(calls) == 1
+    leader_indices, leader_counts = leader_result
+    follower_indices, follower_counts = follower_result
+    torch.testing.assert_close(follower_indices, leader_indices)
+    torch.testing.assert_close(follower_counts, leader_counts)
+
+
+def test_sparse_mla_index_group_falls_back_for_prefill_sized_batch(monkeypatch):
+    device = torch.device(DEVICE_TYPE)
+    logical = torch.full((4, 128), -1, dtype=torch.int32, device=device)
+    logical[:, 0] = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device)
+    builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2)
+    group, leader_index = builder.register_layer(True)
+    follower_group, follower_index = builder.register_layer(False)
+    assert follower_group is group
+
+    metadata = SimpleNamespace(
+        block_table=torch.tensor(
+            [[2], [3], [4], [5]], dtype=torch.int32, device=device
+        ),
+        block_size=4,
+        req_id_per_token=torch.arange(4, dtype=torch.int32, device=device),
+    )
+    calls = 0
+    convert = index_group_module.triton_convert_req_index_to_global_index
+
+    def convert_spy(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        return convert(*args, **kwargs)
+
+    monkeypatch.setattr(
+        index_group_module,
+        "triton_convert_req_index_to_global_index",
+        convert_spy,
+    )
+    leader_result = group.convert_logical_to_physical_topk(
+        leader_index,
+        logical,
+        metadata,
+        block_stride_rows=None,
+        return_valid_counts=False,
+    )
+    follower_result = group.convert_logical_to_physical_topk(
+        follower_index,
+        logical,
+        metadata,
+        block_stride_rows=None,
+        return_valid_counts=False,
     )
+    torch.accelerator.synchronize()
 
-    assert fp8_metadata.num_decodes == 0
-    assert fp8_metadata.num_prefills == 1
-    assert fp8_metadata.num_decode_tokens == 0
-    assert fp8_metadata.num_prefill_tokens == 1
+    assert calls == 2
+    torch.testing.assert_close(follower_result, leader_result)
 
 
-def test_flashmla_common_metadata_requires_uniform_decodes():
-    common_metadata = SimpleNamespace(
-        max_query_len=3,
-        num_reqs=3,
-        num_actual_tokens=6,
-        query_start_loc_cpu=torch.tensor([0, 1, 3, 6]),
-        is_prefilling=None,
+def test_sparse_mla_index_groups_own_distinct_physical_buffers():
+    logical = torch.empty((8, 128), dtype=torch.int32, device=DEVICE_TYPE)
+    builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2)
+
+    first, _ = builder.register_layer(True)
+    follower, _ = builder.register_layer(False)
+    second, _ = builder.register_layer(True)
+
+    assert follower is first
+    assert second is not first
+    assert first.physical_topk_indices[:2].shape == (2, 128)
+    assert second.physical_topk_indices[:2].shape == (2, 128)
+    assert first.physical_topk_indices.data_ptr() != (
+        second.physical_topk_indices.data_ptr()
     )
 
-    split = split_decodes_and_prefills(
-        common_metadata,
-        decode_threshold=128,
-        require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes,
+
+# HiSparse is host-resident-only and kernel-only: runtime construction
+# raises without the compiled CUDA ops.
+requires_hisparse_ops = pytest.mark.skipif(
+    not _has_hisparse_ops(),
+    reason="HiSparse CUDA ops not compiled",
+)
+
+
+def _make_hisparse_runtime(
+    *,
+    top_k: int = 4,
+    device_buffer_size: int = 5,
+    max_num_reqs: int = 2,
+    row_width: int = 8,
+    block_size: int = 64,
+    max_swap_rows: int | None = None,
+    index_group: hisparse_runtime.HiSparseIndexGroup | None = None,
+) -> HiSparseRuntime:
+    runtime = HiSparseRuntime(
+        config=ResolvedHiSparseConfig(
+            top_k=top_k,
+            device_buffer_size=device_buffer_size,
+        ),
+        max_num_reqs=max_num_reqs,
+        row_width=row_width,
+        kv_dtype=torch.float32,
+        device=DEVICE_TYPE,
+        max_swap_rows=max_swap_rows,
+        index_group=index_group,
+    )
+    blocks_per_request = cdiv(runtime.region_stride, block_size)
+    # Leave one extra physical block for resident-tier tests. HMA group layouts
+    # alias this backing, but distinct group allocations use distinct block IDs.
+    num_blocks = max_num_reqs * blocks_per_request + 1
+    raw = torch.zeros(
+        num_blocks * block_size * row_width,
+        dtype=torch.float32,
+        device=DEVICE_TYPE,
+    ).view(torch.int8)
+    block_table = torch.arange(
+        max_num_reqs * blocks_per_request,
+        dtype=torch.int32,
+        device=DEVICE_TYPE,
+    ).view(max_num_reqs, blocks_per_request)
+    runtime.bind_hot_cache(
+        raw,
+        byte_offset=0,
+        block_stride=block_size * row_width * torch.float32.itemsize,
+        num_blocks=num_blocks,
+        block_size=block_size,
+        block_table=block_table,
+    )
+    runtime.request_state_indices = torch.arange(
+        max_num_reqs, dtype=torch.int32, device=runtime.device
+    )
+    return runtime
+
+
+def _make_hisparse_cache_handle(
+    *,
+    top_k: int = 4,
+    device_buffer_size: int = 5,
+    max_num_reqs: int = 2,
+    row_width: int = 8,
+    block_size: int = 64,
+    max_swap_rows: int | None = None,
+) -> HiSparseCacheHandle:
+    runtime = _make_hisparse_runtime(
+        top_k=top_k,
+        device_buffer_size=device_buffer_size,
+        max_num_reqs=max_num_reqs,
+        row_width=row_width,
+        block_size=block_size,
+        max_swap_rows=max_swap_rows,
     )
+    return HiSparseCacheHandle(runtime)
 
-    assert split == (1, 2, 1, 5)
 
+@requires_hisparse_ops
+def test_hisparse_uses_graph_stable_request_state_mapping():
+    device = torch.device(DEVICE_TYPE)
+    block_size, row_width = 64, 8
+    runtime = _make_hisparse_runtime(
+        top_k=1,
+        device_buffer_size=2,
+        max_num_reqs=2,
+        row_width=row_width,
+        block_size=block_size,
+    )
+    runtime.bind_source_cache(
+        torch.arange(2 * block_size * row_width, dtype=torch.float32)
+        .view(2, block_size, row_width)
+        .pin_memory()
+    )
+    runtime.request_state_indices = torch.tensor([1], dtype=torch.int32, device=device)
 
-def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch):
-    monkeypatch.setattr(
-        "vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata",
-        lambda: (object(), None),
+    cache_handle = HiSparseCacheHandle(runtime)
+    cache_handle.all_context_pages_resident = False
+    runtime.begin_forward()
+    cache_handle.swap_in(
+        req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device),
+        block_table=torch.tensor([[1]], dtype=torch.int32, device=device),
+        logical_topk_indices=torch.tensor([[0]], dtype=torch.int32, device=device),
+        block_size=block_size,
     )
-    builder = SimpleNamespace(
-        device=torch.device(DEVICE_TYPE),
-        dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE),
-        max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE),
+    torch.accelerator.synchronize()
+
+    assert (runtime.index_group.device_global_indices[0] == -1).all()
+    assert (runtime.index_group.device_global_indices[1] == block_size).any()
+
+
+@requires_hisparse_ops
+def test_hisparse_maps_speculative_rows_through_request_state():
+    """Multiple verification rows for one request share its persistent state."""
+    device = torch.device(DEVICE_TYPE)
+    block_size, row_width = 64, 8
+    runtime = _make_hisparse_runtime(
+        top_k=1,
+        device_buffer_size=2,
+        max_num_reqs=1,
+        row_width=row_width,
+        block_size=block_size,
+        max_swap_rows=4,
     )
-    query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660])
-    common_metadata = SimpleNamespace(
-        num_actual_tokens=660,
-        query_start_loc_cpu=query_start_loc_cpu,
-        seq_lens=torch.arange(7, device=DEVICE_TYPE),
+    runtime.bind_source_cache(
+        torch.zeros(1, block_size, row_width, dtype=torch.float32).pin_memory()
     )
-    metadata = FlashMLASparseMetadata(
-        num_reqs=7,
-        max_query_len=110,
-        max_seq_len=110,
-        num_actual_tokens=660,
-        query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE),
-        slot_mapping=torch.arange(660, device=DEVICE_TYPE),
-        block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE),
-        req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE),
-        num_decodes=7,
-        num_prefills=0,
-        num_decode_tokens=660,
+
+    cache_handle = HiSparseCacheHandle(runtime)
+    resident_table = torch.tensor([[1]], dtype=torch.int32, device=device)
+    cache_handle.bind_cache(
+        runtime.hot.cache.view(torch.int8),
+        byte_offset=0,
+        block_stride=block_size * row_width * torch.float32.itemsize,
+        num_blocks=2,
+        block_size=block_size,
+        block_table=resident_table,
+        slot_mapping=torch.arange(4, dtype=torch.int64, device=device),
+    )
+    cache_handle.all_context_pages_resident = True
+    runtime.begin_forward()
+    physical = cache_handle.swap_in(
+        req_id_per_token=torch.zeros(4, dtype=torch.int32, device=device),
+        block_table=torch.zeros((1, 1), dtype=torch.int32, device=device),
+        logical_topk_indices=torch.zeros((4, 1), dtype=torch.int32, device=device),
+        block_size=block_size,
     )
+    torch.accelerator.synchronize()
 
-    fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(
-        builder, common_metadata, metadata
+    assert physical.tolist() == [[block_size]] * 4
+
+
+@requires_hisparse_ops
+def test_hisparse_resident_rows_bypass_hot_lru():
+    device = torch.device(DEVICE_TYPE)
+    block_size, row_width = 64, 8
+    cache_handle = _make_hisparse_cache_handle(
+        top_k=4,
+        device_buffer_size=5,
+        max_num_reqs=1,
+        row_width=row_width,
+        block_size=block_size,
     )
+    raw = cache_handle.runtime.hot.cache.view(torch.int8)
+    resident_table = torch.tensor([[1]], dtype=torch.int32, device=device)
+    resident_slots = torch.tensor([block_size], dtype=torch.int64, device=device)
+    cache_handle.bind_cache(
+        raw,
+        byte_offset=0,
+        block_stride=block_size * row_width * torch.float32.itemsize,
+        num_blocks=2,
+        block_size=block_size,
+        block_table=resident_table,
+        slot_mapping=resident_slots,
+    )
+
+    host = torch.randn(1, block_size, row_width).pin_memory()
+    cache_handle.runtime.bind_source_cache(host)
+    resident_rows = host[0].to(device).add(7)
+    cache_handle.runtime.hot.cache[1].copy_(resident_rows)
+    source_table = torch.tensor([[5]], dtype=torch.int32, device=device)
+    request_ids = torch.zeros(1, dtype=torch.int32, device=device)
+    topk = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=device)
+    lru_before = cache_handle.runtime.index_group.lru_slots.clone()
+
+    cache_handle.runtime.begin_forward()
+    indices, valid_counts = cache_handle.swap_in(
+        req_id_per_token=request_ids,
+        block_table=source_table,
+        logical_topk_indices=topk,
+        block_size=block_size,
+        return_valid_counts=True,
+    )
+    torch.accelerator.synchronize()
 
-    assert fp8_metadata.num_decodes == 6
-    assert fp8_metadata.num_decode_tokens == 660
-    assert fp8_metadata.decode is not None
-    assert fp8_metadata.decode.decode_query_len == 110
+    expected_indices = topk + block_size
+    torch.testing.assert_close(indices, expected_indices)
     torch.testing.assert_close(
-        fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE)
+        valid_counts, torch.tensor([topk.shape[1]], dtype=torch.int32, device=device)
+    )
+    torch.testing.assert_close(cache_handle.runtime.index_group.lru_slots, lru_before)
+    torch.testing.assert_close(
+        cache_handle.runtime.hot.attention_cache.view(-1, row_width)[
+            indices.to(torch.long)
+        ],
+        resident_rows[topk.to(torch.long)],
     )
 
 
-@pytest.mark.parametrize("use_mixed_batch", [False, True])
-def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool):
-    num_decode_tokens = 2
-    num_batch_tokens = 5
-    q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE)
-    topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE)
-    kernel_q_shapes = []
+@requires_hisparse_ops
+def test_hisparse_swap_in_preserves_rows_across_eviction():
+    device = torch.device(DEVICE_TYPE)
+    block_size = 64
+    row_width = 64
+    num_blocks = 64
+    top_k = 128
+    buf = 256
+    num_reqs = 4
+
+    kv_pool = torch.randn(
+        (num_blocks + 1, block_size, row_width), dtype=torch.float32
+    ).pin_memory()
+    flat_pool = kv_pool.reshape(-1, row_width)
+
+    runtime = _make_hisparse_runtime(
+        top_k=top_k,
+        device_buffer_size=buf,
+        max_num_reqs=num_reqs,
+        row_width=row_width,
+    )
+    runtime.bind_source_cache(kv_pool)
+    cache = HiSparseCacheHandle(runtime)
+    cache.all_context_pages_resident = False
+    cache.runtime.begin_forward()
+
+    blocks_per_req = num_blocks // num_reqs
+    block_table = torch.arange(
+        1, num_blocks + 1, dtype=torch.int32, device=device
+    ).view(num_reqs, blocks_per_req)
+    req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
+    seq_len = blocks_per_req * block_size
+    base = torch.arange(top_k, dtype=torch.int32, device=device)
+    for step in range(3):
+        cache.runtime.begin_forward()
+        topk = torch.stack(
+            [(base + step * top_k + row * 17) % seq_len for row in range(num_reqs)]
+        )
+        topk[:, -1] = -1
+
+        hot_indices, valid_counts = cache.swap_in(
+            req_id_per_token=req_ids,
+            block_table=block_table,
+            logical_topk_indices=topk.clone(),
+            block_size=block_size,
+            return_valid_counts=True,
+        )
+        torch.accelerator.synchronize()
 
-    def convert_indices(*args, **kwargs):  # noqa: ARG001
-        assert not kwargs.get("HAS_PREFILL_WORKSPACE", False)
-        if not kwargs.get("return_valid_counts", False):
-            return topk_indices
-        valid_counts = torch.full(
-            (num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE
+        global_ref = _triton_convert_reference_impl(
+            req_ids, block_table, topk, block_size, top_k
         )
-        return topk_indices, valid_counts
+        valid = hot_indices >= 0
+        torch.testing.assert_close(
+            valid_counts,
+            (global_ref >= 0).sum(dim=1, dtype=torch.int32),
+            rtol=0,
+            atol=0,
+        )
+        gathered = cache.runtime.hot.attention_cache.reshape(-1, row_width)[
+            hot_indices[valid].to(torch.long)
+        ].cpu()
+        expected = flat_pool[global_ref[valid].cpu().to(torch.long)]
+        torch.testing.assert_close(gathered, expected)
 
-    monkeypatch.setattr(
-        "vllm.v1.attention.backends.mla.flashmla_sparse."
-        "triton_convert_req_index_to_global_index",
-        convert_indices,
+
+@requires_hisparse_ops
+def test_hisparse_multi_step_swaps_match_independent():
+    """The shared LRU must retain every speculative step until followers run."""
+    device = torch.device(DEVICE_TYPE)
+    torch.manual_seed(0)
+    block_size = 64
+    row_width = 64
+    num_blocks = 64
+    top_k = 128
+    buf = 3 * top_k
+    num_reqs = 4
+
+    kv_pool = torch.randn(
+        (num_blocks, block_size, row_width), dtype=torch.float32
+    ).pin_memory()
+
+    def make(
+        index_group: hisparse_runtime.HiSparseIndexGroup | None = None,
+    ) -> HiSparseCacheHandle:
+        runtime = _make_hisparse_runtime(
+            top_k=top_k,
+            max_num_reqs=num_reqs,
+            device_buffer_size=buf,
+            row_width=row_width,
+            block_size=block_size,
+            max_swap_rows=2 * num_reqs,
+            index_group=index_group,
+        )
+        runtime.bind_source_cache(kv_pool)
+        cache = HiSparseCacheHandle(runtime)
+        cache.all_context_pages_resident = False
+        return cache
+
+    blocks_per_req = num_blocks // num_reqs
+    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view(
+        num_reqs, blocks_per_req
+    )
+    req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
+    seq_len = blocks_per_req * block_size
+    producer = make()
+    shared = make(producer.runtime.index_group)
+    indep = make()
+    base = torch.arange(top_k, dtype=torch.int32, device=device)
+    for iteration in range(2):
+        producer.runtime.begin_forward()
+        shared.runtime.begin_forward()
+        indep.runtime.begin_forward()
+        topks = [
+            torch.stack(
+                [
+                    (base + (iteration * 2 + step) * top_k + row * 17) % seq_len
+                    for row in range(num_reqs)
+                ]
+            )
+            for step in range(2)
+        ]
+        kw = dict(
+            req_id_per_token=req_ids,
+            block_table=block_table,
+            block_size=block_size,
+        )
+        producer_indices = [
+            producer.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
+        ]
+        independent_indices = [
+            indep.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
+        ]
+        follower_indices = [
+            shared.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
+        ]
+        torch.accelerator.synchronize()
+
+        for idx_full, idx_indep, idx_shared in zip(
+            producer_indices, independent_indices, follower_indices
+        ):
+            torch.testing.assert_close(idx_shared, idx_full, rtol=0, atol=0)
+            torch.testing.assert_close(idx_indep, idx_full, rtol=0, atol=0)
+        torch.testing.assert_close(shared.runtime.hot.cache, indep.runtime.hot.cache)
+
+        shared_rows = shared.runtime.hot.attention_cache.reshape(-1, row_width)
+        host_rows = kv_pool.reshape(-1, row_width)
+        for physical_rows, logical_rows in zip(follower_indices, topks):
+            global_rows = _triton_convert_reference_impl(
+                req_ids, block_table, logical_rows, block_size, top_k
+            )
+            valid = physical_rows >= 0
+            torch.testing.assert_close(
+                shared_rows[physical_rows[valid].to(torch.long)].cpu(),
+                host_rows[global_rows[valid].cpu().to(torch.long)],
+            )
+
+
+@requires_hisparse_ops
+def test_hisparse_multi_step_writes_request_major_output():
+    device = torch.device(DEVICE_TYPE)
+    block_size = 64
+    row_width = 8
+    num_reqs = 2
+    query_len = 3
+    top_k = 4
+    cache = _make_hisparse_cache_handle(
+        top_k=top_k,
+        device_buffer_size=query_len * top_k,
+        max_num_reqs=num_reqs,
+        row_width=row_width,
+        block_size=block_size,
+        max_swap_rows=num_reqs * query_len,
+    )
+    host = torch.arange(
+        (num_reqs + 1) * block_size * row_width, dtype=torch.float32
+    ).view(num_reqs + 1, block_size, row_width)
+    cache.runtime.bind_source_cache(host.pin_memory())
+    block_table = torch.arange(1, num_reqs + 1, dtype=torch.int32, device=device).view(
+        -1, 1
+    )
+    request_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
+    logical = torch.tensor(
+        [
+            [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]],
+            [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]],
+        ],
+        dtype=torch.int32,
+        device=device,
     )
+    physical = torch.full_like(logical, -1)
+    valid_counts = torch.zeros((num_reqs, query_len), dtype=torch.int32, device=device)
 
-    def run_kernel(**kwargs):
-        kernel_q_shapes.append(kwargs["q"].shape)
-        return kwargs["q"][..., :1], None
+    cache.runtime.begin_forward()
+    for step in range(query_len):
+        cache.swap_in(
+            request_ids,
+            block_table,
+            logical[:, step],
+            block_size=block_size,
+            return_valid_counts=True,
+            attention_indices_out=physical[:, step],
+            valid_counts_out=valid_counts[:, step],
+        )
+    torch.accelerator.synchronize()
 
-    if use_mixed_batch:
-        fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
-            scheduler_metadata=object(),  # type: ignore[arg-type]
-            dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
+    torch.testing.assert_close(
+        valid_counts,
+        torch.full_like(valid_counts, top_k),
+    )
+    global_rows = _triton_convert_reference_impl(
+        torch.arange(num_reqs, device=device, dtype=torch.int32).repeat_interleave(
+            query_len
+        ),
+        block_table,
+        logical.view(num_reqs * query_len, top_k),
+        block_size,
+        top_k,
+    )
+    hot_rows = cache.runtime.hot.attention_cache.reshape(-1, row_width)
+    host_rows = host.reshape(-1, row_width)
+    torch.testing.assert_close(
+        hot_rows[physical.view(num_reqs * query_len, top_k).to(torch.long)].cpu(),
+        host_rows[global_rows.cpu().to(torch.long)],
+    )
+
+
+@requires_hisparse_ops
+def test_hisparse_kv_update_writes_resident_and_staging_caches():
+    device = torch.device(DEVICE_TYPE)
+    block_size = 4
+    row_width = 8
+    cache_handle = _make_hisparse_cache_handle(
+        max_num_reqs=1,
+        row_width=row_width,
+        block_size=block_size,
+    )
+    resident_slots = torch.tensor([4, 5, -1], dtype=torch.int64, device=device)
+    cache_handle.bind_cache(
+        cache_handle.runtime.hot.cache.view(torch.int8),
+        byte_offset=0,
+        block_stride=block_size * row_width * torch.float32.itemsize,
+        num_blocks=cache_handle.runtime.hot.cache.shape[0],
+        block_size=block_size,
+        block_table=torch.tensor([[1]], dtype=torch.int32, device=device),
+        slot_mapping=resident_slots,
+    )
+    cache_handle.mirror_staging_cache = torch.empty(
+        (1, block_size, row_width), dtype=torch.float32, device=device
+    )
+    cache_handle.mirror_staging_slots = torch.arange(
+        block_size, dtype=torch.int64, device=device
+    )
+    slots = torch.tensor([3, 7, -1], dtype=torch.int64, device=device)
+    kv_c = torch.randn(8, row_width - 2, device=device)
+    k_pe = torch.randn(8, 1, 2, device=device)
+    cache_handle.num_actual_tokens = slots.numel()
+    cache_handle.decode_batch = False
+    cache_handle.host_mirror_required = True
+    source_cache = torch.zeros_like(cache_handle.view.cache)
+    impl = object.__new__(FlashMLASparseImpl)
+    layer = SimpleNamespace(
+        hisparse_cache=cache_handle,
+        use_pcp=False,
+        impl=impl,
+    )
+    mla_attention.MLAAttention.update_kv_cache(
+        layer,
+        kv_c,
+        k_pe,
+        source_cache,
+        slots,
+        SimpleNamespace(num_decode_tokens=0),
+        "auto",
+        torch.tensor(1.0, device=device),
+    )
+    torch.accelerator.synchronize()
+
+    expected = torch.cat([kv_c[:2], k_pe[:2, 0]], dim=-1).cpu()
+    torch.testing.assert_close(
+        cache_handle.view.cache.view(-1, row_width)[resident_slots[:2]],
+        expected.to(device),
+    )
+    torch.testing.assert_close(source_cache, torch.zeros_like(source_cache))
+    staged = cache_handle.mirror_staging_cache.view(-1, row_width)
+    staged_expected = torch.cat([kv_c[:3], k_pe[:3, 0]], dim=-1)
+    torch.testing.assert_close(staged[:3], staged_expected)
+
+
+@requires_hisparse_ops
+def test_hisparse_row_dma_copies_discontiguous_spans_across_layers():
+    device = torch.device(DEVICE_TYPE)
+    block_size = 4
+    row_width = 8
+    num_rows = 16
+    resident_caches = tuple(
+        (
+            torch.arange(num_rows * row_width, dtype=torch.uint8, device=device)
+            .add_(layer_index * 17)
+            .view(num_rows // block_size, block_size, row_width)
+        )
+        for layer_index in range(2)
+    )
+    host_caches = tuple(
+        torch.zeros((num_rows, row_width), dtype=torch.uint8, pin_memory=True)
+        for _ in resident_caches
+    )
+    worker = object.__new__(HiSparseConnectorWorker)
+    worker.is_host_writer = True
+    worker.kernel_block_size = block_size
+    worker.resident_caches = resident_caches
+    worker.host_caches = host_caches
+    worker.cache_handles = [
+        SimpleNamespace(
+            decode_batch=True,
+            runtime=SimpleNamespace(resident_source_index=0),
+        )
+        for _ in resident_caches
+    ]
+    worker.dma_stream = torch.cuda.Stream(device=device)
+    worker.host_write_event = torch.Event()
+    worker._dma_free_descriptors = []
+    worker._pending_dma_descriptors = deque()
+    worker._pending_transfer_events = deque()
+    worker._enqueued_transfer_ids = []
+    worker._dma_submitted = False
+    worker._set_row_mirrors(
+        (
+            SparseKVRowMirror((1,), 5, 2),
+            SparseKVRowMirror((8,), 10, 4),
+        )
+    )
+
+    worker._enqueue_row_dma(range(2))
+    current_stream().wait_event(worker.host_write_event)
+    torch.accelerator.synchronize()
+
+    for resident, host in zip(resident_caches, host_caches):
+        flat_resident = resident.view(num_rows, row_width).cpu()
+        torch.testing.assert_close(host[5:7], flat_resident[1:3])
+        torch.testing.assert_close(host[10:14], flat_resident[8:12])
+
+
+@requires_hisparse_ops
+def test_hisparse_remaps_strided_hma_rows_for_attention():
+    device = torch.device(DEVICE_TYPE)
+    block_size, row_width = 4, 8
+    runtime = _make_hisparse_runtime(
+        max_num_reqs=1,
+        block_size=block_size,
+        row_width=row_width,
+    )
+    page_elements = block_size * row_width
+    stride_elements = 3 * page_elements
+    byte_offset = page_elements * torch.float32.itemsize
+    raw = torch.zeros(2 * stride_elements, dtype=torch.float32, device=device).view(
+        torch.int8
+    )
+    runtime.bind_hot_cache(
+        raw,
+        byte_offset=byte_offset,
+        block_stride=stride_elements * torch.float32.itemsize,
+        num_blocks=2,
+        block_size=block_size,
+        block_table=torch.tensor([[0, 1]], dtype=torch.int32, device=device),
+    )
+
+    source = (
+        torch.arange(36 * row_width, dtype=torch.float32)
+        .view(9, block_size, row_width)
+        .pin_memory()
+    )
+    runtime.bind_source_cache(source)
+    cache = HiSparseCacheHandle(runtime)
+    cache.all_context_pages_resident = False
+    runtime.begin_forward()
+    attention_indices = cache.swap_in(
+        req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device),
+        block_table=torch.arange(1, 9, dtype=torch.int32, device=device).view(1, 8),
+        logical_topk_indices=torch.tensor(
+            [[0, 1, 2, 3]], dtype=torch.int32, device=device
+        ),
+        block_size=block_size,
+    )
+
+    assert runtime.hot.attention_cache is not None
+    assert runtime.hot.attention_cache.is_contiguous()
+    physical_indices = runtime.index_group.shared_topk.device_topk_rows[0]
+    expected = (
+        physical_indices // block_size * (stride_elements // row_width)
+        + physical_indices % block_size
+    )
+    torch.testing.assert_close(attention_indices[0], expected)
+
+
+@requires_hisparse_ops
+def test_hisparse_gather_zeroes_unaligned_destination():
+    """Invalid host rows are safely zeroed even at a byte-unaligned address."""
+    device = torch.device(DEVICE_TYPE)
+    row_bytes = 16
+    host_cache = torch.ones(1, row_bytes, dtype=torch.uint8).pin_memory()
+    storage = torch.full((2 * row_bytes + 1,), 7, dtype=torch.uint8, device=device)
+    hot_cache = torch.as_strided(
+        storage,
+        size=(2, 1, row_bytes),
+        stride=(row_bytes + 1, row_bytes, 1),
+    )
+    global_indices = torch.tensor([[1]], dtype=torch.int32, device=device)
+    hot_indices = torch.tensor([[1]], dtype=torch.int32, device=device)
+    miss_mask = torch.ones_like(global_indices)
+
+    torch.ops._C_cache_ops.hisparse_gather_plan(
+        host_cache,
+        hot_cache,
+        global_indices,
+        hot_indices,
+        miss_mask,
+        None,
+    )
+    torch.accelerator.synchronize()
+
+    torch.testing.assert_close(hot_cache[1, 0], torch.zeros_like(hot_cache[1, 0]))
+    assert storage[row_bytes].item() == 7
+
+
+@requires_hisparse_ops
+def test_hisparse_gather_rejects_short_attention_block_stride():
+    """Reject a stride that would let the gather kernel cross block bounds."""
+    device = torch.device(DEVICE_TYPE)
+    host_cache = torch.ones(1, 16, dtype=torch.uint8).pin_memory()
+    hot_cache = torch.empty((1, 2, 16), dtype=torch.uint8, device=device)
+    global_indices = torch.zeros((1, 1), dtype=torch.int32, device=device)
+    hot_indices = torch.zeros_like(global_indices)
+    miss_mask = torch.ones_like(global_indices)
+    attention_indices = torch.empty_like(global_indices)
+
+    with pytest.raises(RuntimeError, match="stride must cover one hot block"):
+        torch.ops._C_cache_ops.hisparse_gather_plan(
+            host_cache,
+            hot_cache,
+            global_indices,
+            hot_indices,
+            miss_mask,
+            None,
+            attention_indices,
+            1,
+        )
+
+
+@requires_hisparse_ops
+def test_hisparse_newest_write_and_recycled_slot_invalidation():
+    """Newest writes clamp padding and invalidated rows are loaded again."""
+    device = torch.device(DEVICE_TYPE)
+    block_size = 4
+    row_width = 8
+    num_blocks = 8
+
+    kv_pool = (
+        torch.arange(num_blocks * block_size * row_width, dtype=torch.float32)
+        .view(num_blocks, block_size, row_width)
+        .pin_memory()
+    )
+    flat_pool = kv_pool.reshape(-1, row_width)
+
+    cache_handle = _make_hisparse_cache_handle(block_size=block_size)
+    cache_handle.runtime.eager_host_mirror = True
+
+    block_table = torch.tensor([[2, 0, 4]], dtype=torch.int32, device=device)
+    resident_slots = torch.tensor([4, -1], dtype=torch.int64, device=device)
+    cache_handle.bind_cache(
+        cache_handle.runtime.hot.cache.view(torch.int8),
+        byte_offset=0,
+        block_stride=block_size * row_width * torch.float32.itemsize,
+        num_blocks=cache_handle.runtime.hot.cache.shape[0],
+        block_size=block_size,
+        block_table=torch.tensor([[0, 0, 1]], dtype=torch.int32, device=device),
+        slot_mapping=resident_slots,
+    )
+    req_ids = torch.tensor([0], dtype=torch.int32, device=device)
+    newest_global = 4 * block_size
+    padded_global = 1
+    padded_row = flat_pool[padded_global].clone()
+    slot_mapping = torch.tensor(
+        [newest_global, padded_global], dtype=torch.int64, device=device
+    )
+    request_state_indices = torch.tensor([0, -1], dtype=torch.int32, device=device)
+    cache_handle.runtime.request_state_indices = request_state_indices
+
+    kv_c = torch.randn(3, row_width - 2, device=device)
+    k_pe = torch.randn(3, 1, 2, device=device)
+    cache_handle.runtime.bind_source_cache(kv_pool)
+    cache_handle.num_actual_tokens = slot_mapping.numel()
+    ops.concat_and_cache_mla(
+        kv_c[:1],
+        k_pe[:1, 0],
+        cache_handle.view.cache,
+        resident_slots[:1],
+        kv_cache_dtype="auto",
+        scale=torch.tensor(1.0, device=device),
+    )
+    torch.accelerator.synchronize()
+    expected_row = torch.cat([kv_c[0], k_pe[0, 0]]).cpu()
+    flat_pool[newest_global].copy_(expected_row)
+    torch.testing.assert_close(flat_pool[newest_global], expected_row)
+    torch.testing.assert_close(flat_pool[padded_global], padded_row)
+    flat_hot = cache_handle.runtime.hot.cache.view(-1, row_width)
+    torch.testing.assert_close(flat_hot[resident_slots[0]].cpu(), expected_row)
+
+    topk = torch.tensor([[0, -1, -1, -1]], dtype=torch.int32, device=device)
+    cache_handle.all_context_pages_resident = False
+    cache_handle.runtime.begin_forward()
+    hot_indices = cache_handle.swap_in(
+        req_id_per_token=req_ids,
+        block_table=block_table,
+        logical_topk_indices=topk,
+        block_size=block_size,
+    )
+    torch.accelerator.synchronize()
+    stale_hot_slot = hot_indices[0, 0].item()
+    stale_row = flat_hot[stale_hot_slot].clone()
+
+    flat_pool[8] += 1000
+    cache_handle.runtime.invalidate_written_slots(
+        torch.tensor([8], dtype=torch.int64, device=device), req_ids
+    )
+    cache_handle.runtime.begin_forward()
+    hot_indices = cache_handle.swap_in(
+        req_id_per_token=req_ids,
+        block_table=block_table,
+        logical_topk_indices=topk,
+        block_size=block_size,
+    )
+    torch.accelerator.synchronize()
+    idx = hot_indices.cpu().tolist()[0][0]
+    assert not torch.equal(flat_hot[idx], stale_row)
+    torch.testing.assert_close(flat_hot[idx].cpu(), flat_pool[8])
+
+
+@requires_hisparse_ops
+def test_hisparse_mixed_batch_bf16_row_split(
+    default_vllm_config, dist_init, workspace_init
+):
+    """Host-resident mixed batch on the bf16 path is row-split.
+
+    Two long-context speculative-decode requests + one short local-prefill
+    chunk: every decode step must be served from the bounded hot buffer before
+    it is reused, while only the prefill rows' blocks are staged host->GPU.
+    """
+    ok, reason = flashmla.is_flashmla_sparse_supported()
+    if not ok:
+        pytest.skip(reason)
+
+    device = torch.device(DEVICE_TYPE)
+    dtype = torch.bfloat16
+    torch.manual_seed(0)
+
+    num_heads = 64
+    kv_lora_rank = 512
+    qk_nope_head_dim = 128
+    qk_rope_head_dim = 64
+    v_head_dim = 128
+    head_size = kv_lora_rank + qk_rope_head_dim
+    topk_tokens = 128
+    block_size = 64
+
+    # Long decode contexts + a short prefill chunk (router shortcut shape).
+    batch_spec = BatchSpec(seq_lens=[2048, 2048, 192], query_lens=[2, 2, 64])
+    max_seqlen = max(batch_spec.seq_lens)
+    total_cache_tokens = sum(batch_spec.seq_lens)
+    total_tokens = batch_spec.compute_num_tokens()
+
+    vllm_config = create_vllm_config(
+        model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
+        tensor_parallel_size=1,
+        max_model_len=max_seqlen,
+        num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
+        block_size=block_size,
+        hf_config_override={
+            "index_topk": topk_tokens,
+            "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
+        },
+    )
+    vllm_config.attention_config.hisparse_config = HiSparseConfig(
+        device_buffer_size=2 * topk_tokens,
+    )
+    vllm_config.speculative_config = SpeculativeConfig(
+        method="ngram", num_speculative_tokens=1
+    )
+    model_config = vllm_config.model_config
+    model_config.hf_text_config = SimpleNamespace(
+        index_topk=topk_tokens,
+        q_lora_rank=None,
+        kv_lora_rank=kv_lora_rank,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+        v_head_dim=v_head_dim,
+        model_type="deepseek_v2",
+    )
+    model_config.dtype = dtype
+    model_config.get_num_attention_heads = MethodType(
+        lambda self, parallel_config: num_heads, model_config
+    )
+    model_config.get_num_kv_heads = MethodType(
+        lambda self, parallel_config: 1, model_config
+    )
+    model_config.get_head_size = MethodType(lambda self: head_size, model_config)
+    model_config.get_sliding_window = MethodType(lambda self: None, model_config)
+    # create_hisparse_cache_handle sizes hot buffers per layer.
+    model_config.get_num_layers = MethodType(
+        lambda self, parallel_config: 1, model_config
+    )
+
+    kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        vllm_config.cache_config.block_size,
+        device,
+        arange_block_indices=True,
+    )
+    # Prepopulate every position of every sequence so the forward needs no
+    # KV-cache update (the row split itself is what is under test).
+    kv_c_contexts = [
+        torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
+        for s_len in batch_spec.seq_lens
+    ]
+    k_pe_contexts = [
+        torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
+        for s_len in batch_spec.seq_lens
+    ]
+    kv_cache = create_and_prepopulate_kv_cache(
+        kv_c_contexts=kv_c_contexts,
+        k_pe_contexts=k_pe_contexts,
+        block_size=vllm_config.cache_config.block_size,
+        head_size=head_size,
+        dtype=dtype,
+        device=device,
+        num_blocks=vllm_config.cache_config.num_gpu_blocks,
+        common_attn_metadata=common_attn_metadata,
+        randomize_blocks=False,
+        kv_cache_dtype="auto",
+    )
+
+    prefill_backend = get_mla_prefill_backend(vllm_config)(
+        num_heads=num_heads,
+        scale=1.0 / math.sqrt(head_size),
+        kv_lora_rank=kv_lora_rank,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+        v_head_dim=v_head_dim,
+        vllm_config=vllm_config,
+    )
+    vllm_config.compilation_config.static_forward_context["placeholder"] = (
+        SimpleNamespace(prefill_backend=prefill_backend)
+    )
+
+    builder_cls = FlashMLASparseBackend.get_builder_cls()
+    builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
+    metadata = builder.build(
+        common_prefix_len=0, common_attn_metadata=common_attn_metadata
+    )
+    assert isinstance(metadata.prefill, SparseMLAPrefillMetadata)
+
+    # Per-token sparse indices bounded by each token's position, with unique
+    # offsets and -1 padding (same construction as the decode parity test).
+    positions: list[int] = []
+    for i in range(batch_spec.batch_size):
+        ctx_len = batch_spec.seq_lens[i] - batch_spec.query_lens[i]
+        positions.extend(ctx_len + q_idx for q_idx in range(batch_spec.query_lens[i]))
+    sparse_indices = torch.empty(
+        total_tokens, topk_tokens, dtype=torch.int32, device=device
+    )
+    for tok_idx in range(total_tokens):
+        max_valid_idx = positions[tok_idx]
+        offset = tok_idx * 7  # Prime number for varied offsets
+        num_valid = min(topk_tokens // 2, max_valid_idx + 1)
+        valid_range = torch.arange(num_valid, device=device, dtype=torch.int32)
+        sparse_indices[tok_idx] = torch.cat(
+            [
+                (valid_range + offset) % (max_valid_idx + 1),
+                torch.full(
+                    (topk_tokens - num_valid,), -1, device=device, dtype=torch.int32
+                ),
+            ]
+        )
+
+    q = torch.rand(total_tokens, num_heads, head_size, dtype=dtype, device=device)
+    mock_indexer = SimpleNamespace(
+        topk_indices_buffer=sparse_indices, topk_tokens=topk_tokens
+    )
+
+    impl_cls = FlashMLASparseBackend.get_impl_cls()
+    with set_current_vllm_config(vllm_config):
+        impl = impl_cls(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=1.0 / math.sqrt(head_size),
+            num_kv_heads=1,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="auto",
+            logits_soft_cap=None,
+            attn_type="decoder",
+            kv_sharing_target_layer_name=None,
+            q_lora_rank=None,
+            kv_lora_rank=kv_lora_rank,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            kv_b_proj=None,
+            indexer=mock_indexer,
+        )
+    assert isinstance(impl.index_group, HiSparseMLAIndexGroup)
+    cache_handle = impl.index_group.cache(impl.index_group_index)
+    blocks_per_request = cdiv(cache_handle.runtime.region_stride, block_size)
+    num_hot_blocks = cache_handle.runtime.max_num_reqs * blocks_per_request
+    hot_cache = torch.zeros(
+        num_hot_blocks * block_size * head_size,
+        dtype=dtype,
+        device=device,
+    ).view(torch.int8)
+    cache_handle.runtime.bind_hot_cache(
+        hot_cache,
+        byte_offset=0,
+        block_stride=block_size * head_size * dtype.itemsize,
+        num_blocks=num_hot_blocks,
+        block_size=block_size,
+        block_table=torch.arange(num_hot_blocks, dtype=torch.int32, device=device).view(
+            cache_handle.runtime.max_num_reqs, blocks_per_request
+        ),
+    )
+    cache_handle.runtime.request_state_indices = torch.arange(
+        cache_handle.runtime.max_num_reqs, dtype=torch.int32, device=device
+    )
+    impl.prepare_for_batch(metadata)
+
+    # Device-resident reference: the whole batch converted against the full
+    # block table and run as one kernel call over the GPU cache.
+    ref_topk, ref_topk_length = triton_convert_req_index_to_global_index(
+        metadata.req_id_per_token,
+        metadata.block_table,
+        sparse_indices,
+        BLOCK_SIZE=metadata.block_size,
+        NUM_TOPK_TOKENS=topk_tokens,
+        return_valid_counts=True,
+    )
+    reference, _ = impl._bf16_flash_mla_kernel(q, kv_cache, ref_topk, ref_topk_length)
+
+    # Host-resident pool with identical contents.
+    kv_pool = kv_cache.squeeze(1).cpu().pin_memory()
+    cache_handle.runtime.bind_source_cache(kv_pool)
+    cache_handle.source_block_table = metadata.block_table
+
+    staging_calls = []
+    original_gather = cache_handle.runtime.gather_prefill_cache
+
+    def spy_gather(self, kv, plan, resident_cache=None):
+        staged = original_gather(kv, plan, resident_cache)
+        staging_calls.append((plan, staged.shape))
+        return staged
+
+    cache_handle.runtime.gather_prefill_cache = MethodType(
+        spy_gather, cache_handle.runtime
+    )
+
+    backend_output, _ = impl._forward_bf16_kv(
+        q, kv_pool, sparse_indices, metadata, q.shape[1]
+    )
+    torch.accelerator.synchronize()
+
+    # Only the prefill rows' blocks were staged: the decode rows' 2048-token
+    # contexts (32 blocks each) must stay off the staging gather.
+    assert len(staging_calls) == 1
+    plan, staged_shape = staging_calls[0]
+    assert plan is metadata.prefill.host_staging_plan
+    prefill_blocks = cdiv(batch_spec.seq_lens[-1], block_size)
+    assert staged_shape[0] <= prefill_blocks + 1  # +1: block-0 tail padding
+
+    torch.testing.assert_close(backend_output, reference, rtol=0.01, atol=0.01)
+
+
+def test_hisparse_prefill_staging_remap():
+    """Compacted staging references the same host rows as direct indexing."""
+    block_size = 4
+    block_table = torch.tensor(
+        [[5, 2, -1, -1], [2, 7, 3, -1], [9, -1, -1, -1]], dtype=torch.int32
+    )
+
+    new_bt, row_ids = hisparse_prefill_staging_remap(block_table, block_size)
+
+    assert new_bt.shape == block_table.shape
+    assert row_ids.dtype == torch.int32 and new_bt.dtype == torch.int32
+    n_unique = len({0, 2, 3, 5, 7, 9})  # -1 clamps to block 0
+    flat_rows = row_ids.flatten()
+    assert (flat_rows >= 0).sum() == n_unique * block_size
+    for i in range(block_table.shape[0]):
+        for j in range(block_table.shape[1]):
+            orig = max(int(block_table[i, j]), 0)
+            staged = int(new_bt[i, j])
+            for k in range(block_size):
+                assert int(flat_rows[staged * block_size + k]) == orig * block_size + k
+
+
+def test_hisparse_prefill_staging_plan_masks_unused_blocks():
+    block_table = torch.tensor([[5, 2, 7], [9, 3, 4]], dtype=torch.int32)
+    plan = build_hisparse_prefill_staging_plan(
+        block_table,
+        seq_lens=torch.tensor([5, 8], dtype=torch.int32),
+        block_size=4,
+        staging_block_capacity=4,
+    )
+
+    assert (plan.block_table[:, 2] == 0).all()
+    valid_rows = plan.row_ids[plan.row_ids >= 0]
+    assert set((valid_rows[::4] // 4).tolist()) == {0, 2, 3, 5, 9}
+    for row, used in enumerate((2, 2)):
+        for column in range(used):
+            staged_block = int(plan.block_table[row, column])
+            host_block = int(block_table[row, column])
+            assert int(valid_rows[staged_block * 4]) == host_block * 4
+
+
+def test_hisparse_prefill_staging_plan_resolves_resident_sources():
+    """Per-page resident hits become device sources; misses stay host DMAs."""
+    block_size = 4
+    resident_block_size = 2
+    block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32)
+    seq_lens = torch.tensor([5, 8], dtype=torch.int32)
+    plan = build_hisparse_prefill_staging_plan(
+        block_table, seq_lens, block_size, staging_block_capacity=4
+    )
+    # Two resident pages per host block; 0 entries are null (not resident).
+    resident_table = torch.tensor(
+        [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], dtype=torch.int32
+    )
+
+    plan.ensure_gpu_sources(resident_table, resident_block_size)
+
+    assert plan.gpu_row_ids is not None
+    unique_hosts = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist()
+    gpu_rows = plan.gpu_row_ids[0].view(-1, block_size)
+    reps = {5: (0, 0), 2: (0, 1), 9: (1, 0), 3: (1, 1)}
+    for u, host_id in enumerate(unique_hosts):
+        for t in range(block_size):
+            if host_id < 0:
+                assert int(gpu_rows[u, t]) == -1
+                continue
+            if host_id == 0:
+                assert int(gpu_rows[u, t]) == -1
+                continue
+            row, col = reps[host_id]
+            page = t // resident_block_size
+            res_block = int(resident_table[row, col * 2 + page])
+            expected = (
+                res_block * resident_block_size + t % resident_block_size
+                if res_block > 0
+                else -1
+            )
+            assert int(gpu_rows[u, t]) == expected
+    valid_rows = plan.row_ids >= 0
+    torch.testing.assert_close(
+        plan.miss_mask,
+        ((plan.gpu_row_ids < 0) & valid_rows).int(),
+        check_dtype=False,
+    )
+
+    plan.ensure_gpu_sources(torch.zeros_like(resident_table), resident_block_size)
+
+    assert plan.gpu_row_ids is not None
+    assert (plan.gpu_row_ids == -1).all()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_hisparse_gather_prefill_cache_prefers_resident_rows():
+    """Staged rows come from the resident cache when a shadow page exists."""
+    if not _has_hisparse_ops():
+        pytest.skip("hisparse CUDA ops unavailable")
+    device = torch.device("cuda")
+    block_size, resident_block_size, row_width = 4, 2, 16
+    block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32, device=device)
+    seq_lens = torch.tensor([5, 8], dtype=torch.int32, device=device)
+    plan = build_hisparse_prefill_staging_plan(
+        block_table, seq_lens, block_size, staging_block_capacity=4
+    )
+    resident_table = torch.tensor(
+        [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]],
+        dtype=torch.int32,
+        device=device,
+    )
+    plan.ensure_gpu_sources(resident_table, resident_block_size)
+
+    num_host_blocks, num_res_blocks = 10, 24
+    host_cache = (
+        torch.arange(num_host_blocks * block_size * row_width, dtype=torch.float32)
+        .view(num_host_blocks, block_size, row_width)
+        .pin_memory()
+    )
+    resident_cache = (
+        (
+            -torch.arange(
+                num_res_blocks * resident_block_size * row_width, dtype=torch.float32
+            )
+            - 1.0
+        )
+        .view(num_res_blocks, resident_block_size, row_width)
+        .to(device)
+    )
+
+    staged = HiSparseRuntime.gather_prefill_cache(
+        None, host_cache, plan, resident_cache=resident_cache
+    )
+
+    staged_flat = staged.view(-1, row_width).cpu()
+    host_flat = host_cache.view(-1, row_width)
+    resident_flat = resident_cache.reshape(-1, row_width).cpu()
+    gpu_rows = plan.gpu_row_ids[0].cpu()
+    host_rows = plan.row_ids[0].cpu()
+    for i in range(staged_flat.shape[0]):
+        if int(host_rows[i]) < 0:
+            assert int(plan.miss_mask[0, i]) == 0
+            continue
+        expected = (
+            resident_flat[int(gpu_rows[i])]
+            if int(gpu_rows[i]) >= 0
+            else host_flat[int(host_rows[i])]
+        )
+        torch.testing.assert_close(staged_flat[i], expected)
+
+
+def test_hisparse_fp8_decode_resolves_steps_then_runs_batched_attention(monkeypatch):
+    """FP8 decode must use active dimensions when common metadata is padded."""
+    device = torch.device("cpu")
+    num_decodes = 2
+    query_len = 3
+    num_tokens = num_decodes * query_len
+    q = torch.randn(num_tokens, 2, 4, device=device)
+    topk = (
+        torch.arange(num_tokens, dtype=torch.int32, device=device)
+        .view(-1, 1)
+        .expand(-1, 4)
+    )
+    steps: list[int] = []
+    kernel_shapes: list[torch.Size] = []
+
+    def swap_in(req_ids, *, logical_topk_indices, **kwargs):  # noqa: ARG001
+        steps.append(int(logical_topk_indices[0, 0]))
+        output = kwargs["attention_indices_out"]
+        if output is not None:
+            output.copy_(logical_topk_indices + 10)
+        return topk[:num_decodes]
+
+    def run_kernel(self, *, q, **kwargs):  # noqa: ARG001
+        kernel_shapes.append(q.shape)
+        return q[..., :1], None
+
+    runtime = SimpleNamespace(
+        hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
+    )
+    leader_cache = SimpleNamespace(
+        runtime=runtime,
+        source_block_table=torch.empty(
+            num_decodes, 1, dtype=torch.int32, device=device
+        ),
+        swap_in=MagicMock(side_effect=swap_in),
+    )
+    follower_cache = SimpleNamespace(
+        runtime=runtime,
+        source_block_table=torch.empty(
+            num_decodes, 1, dtype=torch.int32, device=device
+        ),
+        swap_in=MagicMock(side_effect=swap_in),
+    )
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.caches = [leader_cache, follower_cache]
+    index_group.physical_topk_indices = torch.empty(
+        (num_tokens + 1, 4), dtype=torch.int32, device=device
+    )
+    index_group.valid_topk_counts = torch.empty(
+        num_tokens + 1, dtype=torch.int32, device=device
+    )
+    index_group.request_ids = torch.arange(
+        num_decodes, dtype=torch.int32, device=device
+    )
+    impl = SimpleNamespace(
+        kv_lora_rank=1,
+        index_group=index_group,
+        index_group_index=0,
+    )
+    impl._fp8_flash_mla_kernel = MethodType(run_kernel, impl)
+    metadata = SimpleNamespace(
+        num_decodes=num_tokens,
+        num_decode_tokens=num_tokens,
+        decode_max_query_len=1,
+        query_start_loc=torch.arange(
+            0,
+            num_tokens + 1,
+            query_len,
+            dtype=torch.int32,
+            device=device,
+        ),
+        block_table=torch.empty(num_decodes, 1, dtype=torch.int32, device=device),
+        block_size=64,
+    )
+
+    output = FlashMLASparseImpl._host_backed_fp8_decode(
+        impl,
+        q,
+        topk,
+        metadata,
+        SimpleNamespace(),
+        num_decodes,
+        query_len,
+    )
+
+    assert steps == [0, 1, 2]
+    assert kernel_shapes == [(num_decodes, query_len, 2, 4)]
+    assert output.shape == (num_tokens, 2, 1)
+    torch.testing.assert_close(
+        index_group.physical_topk_indices[:num_tokens], topk + 10
+    )
+
+    follower_indices = index_group.convert_decode_logical_to_physical_topk(
+        1,
+        topk,
+        metadata,
+        return_valid_counts=False,
+        num_decodes=num_decodes,
+        decode_query_len=query_len,
+    )
+    assert steps == [0, 1, 2, 0, 1, 2]
+    torch.testing.assert_close(follower_indices, topk + 10)
+
+
+def test_hisparse_varlen_decode_uses_canonical_request_rows():
+    """Graph metadata's query upper bound must preserve packed token rows."""
+    num_decodes = 2
+    topk = torch.arange(18, dtype=torch.int32).view(9, 2)
+
+    def swap_in(_request_ids, *, logical_topk_indices, **_kwargs):
+        return logical_topk_indices + 10
+
+    cache = SimpleNamespace(
+        source_block_table=torch.zeros((num_decodes, 1), dtype=torch.int32),
+        swap_in=MagicMock(side_effect=swap_in),
+    )
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.caches = [cache]
+    index_group.physical_topk_indices = torch.empty((10, 2), dtype=torch.int32)
+    index_group.valid_topk_counts = torch.empty(10, dtype=torch.int32)
+    index_group.request_ids = torch.arange(num_decodes, dtype=torch.int32)
+    metadata = SimpleNamespace(
+        num_decodes=num_decodes,
+        num_decode_tokens=topk.shape[0],
+        decode_max_query_len=15,
+        query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32),
+        block_size=64,
+    )
+
+    physical_topk = index_group.convert_decode_logical_to_physical_topk(
+        0,
+        topk,
+        metadata,
+        return_valid_counts=False,
+    )
+
+    torch.testing.assert_close(physical_topk, topk + 10)
+    request_rows = cache.swap_in.call_args_list[0].args[0]
+    torch.testing.assert_close(request_rows, index_group.request_ids)
+
+
+def test_flashinfer_hisparse_decode_runs_batched_attention():
+    device = torch.device("cpu")
+    num_tokens = 6
+    q = torch.randn(num_tokens, 2, 4, device=device)
+    topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device)
+    physical_topk = topk + 10
+    valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device)
+    kernel_shapes: list[torch.Size] = []
+
+    def convert_decode(self, *args, **kwargs):  # noqa: ARG001
+        return physical_topk, valid_counts
+
+    def prepare_kernel(self, *args, **kwargs):  # noqa: ARG001
+        pass
+
+    def run_kernel(self, q, cache, indices, counts):  # noqa: ARG001
+        kernel_shapes.append(q.shape)
+        return q[..., :1], None
+
+    cache_handle = SimpleNamespace(
+        runtime=SimpleNamespace(
+            hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
+        ),
+    )
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.caches = [cache_handle]
+    index_group.convert_decode_logical_to_physical_topk = MethodType(
+        convert_decode, index_group
+    )
+    impl = object.__new__(FlashInferMLASparseImpl)
+    impl.topk_indices_buffer = topk
+    impl.index_group = index_group
+    impl.index_group_index = 0
+    impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl)
+    impl._run_mqa_kernel = MethodType(run_kernel, impl)
+    metadata = SimpleNamespace(num_decode_tokens=num_tokens)
+
+    output, lse = FlashInferMLASparseImpl.forward_mqa(
+        impl,
+        q,
+        torch.empty(1, device=device),
+        metadata,
+        SimpleNamespace(),
+    )
+
+    assert kernel_shapes == [q.shape]
+    assert output.shape == (num_tokens, 2, 1)
+    assert lse is None
+
+
+def test_flashattn_hisparse_decode_uses_index_group():
+    num_tokens = 4
+    q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE)
+    q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE)
+    topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
+    physical = topk.clone()
+    counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE)
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.convert_decode_logical_to_physical_topk = MagicMock(
+        return_value=(physical, counts)
+    )
+    index_group.physical_kv_cache = MagicMock(
+        return_value=torch.empty(1, device=DEVICE_TYPE)
+    )
+    impl = object.__new__(FlashAttnMLASparseImpl)
+    impl.topk_indices_buffer = topk
+    impl.index_group = index_group
+    impl.index_group_index = 0
+    impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1])
+    metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64)
+
+    output, lse = FlashAttnMLASparseImpl.forward_mqa(
+        impl,
+        (q_nope, q_rope),
+        torch.empty(1, device=DEVICE_TYPE),
+        metadata,
+        SimpleNamespace(),
+    )
+
+    assert output.shape == (num_tokens, 2, 1)
+    assert lse is None
+    index_group.convert_decode_logical_to_physical_topk.assert_called_once()
+    impl._run_mqa_kernel.assert_called_once()
+
+
+def test_flashinfer_sm120_hisparse_decode_uses_index_group():
+    num_tokens = 4
+    q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE)
+    topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
+    physical = topk.clone()
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.convert_decode_logical_to_physical_topk = MagicMock(
+        return_value=physical
+    )
+    index_group.physical_kv_cache = MagicMock(
+        return_value=torch.empty(1, device=DEVICE_TYPE)
+    )
+    impl = object.__new__(FlashInferMLASparseSM120Impl)
+    impl.topk_indices_buffer = topk
+    impl.index_group = index_group
+    impl.index_group_index = 0
+    impl._run_mqa_kernel = MagicMock(return_value=q[..., :1])
+    metadata = SimpleNamespace(
+        num_decode_tokens=num_tokens,
+        topk_tokens=topk.shape[1],
+    )
+
+    output, lse = FlashInferMLASparseSM120Impl.forward_mqa(
+        impl,
+        q,
+        torch.empty(1, device=DEVICE_TYPE),
+        metadata,
+        SimpleNamespace(),
+    )
+
+    assert output.shape == (num_tokens, 2, 1)
+    assert lse is None
+    index_group.convert_decode_logical_to_physical_topk.assert_called_once()
+    impl._run_mqa_kernel.assert_called_once()
+
+
+def test_hisparse_resident_prefill_uses_attention_block_stride():
+    expected = torch.tensor([[19]], dtype=torch.int32)
+    cache_handle = SimpleNamespace(
+        all_context_pages_resident=True,
+        view=SimpleNamespace(block_size=64, attention_block_stride=832),
+        block_table=torch.tensor([[3]], dtype=torch.int32),
+    )
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.caches = [cache_handle]
+    index_group.physical_topk_indices = torch.empty((1, 1), dtype=torch.int32)
+    index_group._convert_once = MagicMock(return_value=expected)
+    topk = torch.zeros((2, 1), dtype=torch.int32)
+    metadata = SimpleNamespace(
+        req_id_per_token=torch.zeros(2, dtype=torch.int32),
+    )
+
+    result = index_group.convert_logical_to_physical_topk(
+        0,
+        topk,
+        metadata,
+        block_stride_rows=None,
+        return_valid_counts=False,
+    )
+
+    assert result is expected
+    assert index_group._convert_once.call_args.kwargs["block_stride_rows"] == 832
+
+
+def test_hisparse_shared_sparse_builder_routes_multi_token_chunks_to_prefill():
+    """Default sparse backends must route multi-token chunks to prefill."""
+    builder = object.__new__(FlashInferMLASparseMetadataBuilder)
+    builder.vllm_config = SimpleNamespace(
+        attention_config=SimpleNamespace(hisparse_config=object()),
+        speculative_config=SimpleNamespace(
+            num_speculative_tokens=4,
+            parallel_drafting=False,
+        ),
+        parallel_config=SimpleNamespace(decode_context_parallel_size=1),
+    )
+
+    builder._init_reorder_batch_threshold(
+        1024,
+        supports_spec_as_decode=True,
+        supports_dcp_with_varlen=True,
+    )
+
+    assert builder.reorder_batch_threshold == 1
+
+
+def test_hisparse_flashmla_reorders_full_speculative_window_as_decode():
+    builder = object.__new__(FlashMLASparseMetadataBuilder)
+    builder.vllm_config = SimpleNamespace(
+        attention_config=SimpleNamespace(hisparse_config=object()),
+        speculative_config=SimpleNamespace(
+            num_speculative_tokens=3,
+            parallel_drafting=False,
+        ),
+        parallel_config=SimpleNamespace(decode_context_parallel_size=1),
+    )
+
+    builder._init_reorder_batch_threshold(
+        256,
+        supports_spec_as_decode=True,
+        supports_dcp_with_varlen=True,
+    )
+
+    assert builder.reorder_batch_threshold == 4
+
+
+def test_hisparse_mixed_single_token_batch_is_not_decode_only():
+    cache = SimpleNamespace(
+        runtime=SimpleNamespace(
+            begin_forward=lambda: None,
+            eager_host_mirror=True,
+        ),
+    )
+    metadata = SimpleNamespace(
+        num_decode_tokens=1,
+        num_actual_tokens=2,
+        max_query_len=1,
+        num_reqs=2,
+        req_id_per_token=torch.arange(2),
+    )
+
+    HiSparseCacheHandle._prepare_for_batch(cache, metadata)
+
+    assert not cache.decode_batch
+
+
+def test_flashmla_cache_dtype_aliases_use_ds_layout():
+    # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer
+    # stores kv_cache_dtype, so they cannot bypass the gate.
+    for alias in ("fp8", "fp8_e4m3"):
+        assert (
+            _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias)
+            == "fp8_ds_mla"
+        )
+
+
+def test_flashmla_common_metadata_requires_uniform_decodes():
+    common_metadata = SimpleNamespace(
+        max_query_len=3,
+        num_reqs=3,
+        num_actual_tokens=6,
+        query_start_loc_cpu=torch.tensor([0, 1, 3, 6]),
+        is_prefilling=None,
+    )
+
+    split = split_decodes_and_prefills(
+        common_metadata,
+        decode_threshold=128,
+        require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes,
+    )
+
+    assert split == (1, 2, 1, 5)
+
+
+def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch):
+    monkeypatch.setattr(
+        "vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata",
+        lambda: (object(), None),
+    )
+    builder = SimpleNamespace(
+        device=torch.device(DEVICE_TYPE),
+        dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE),
+        max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE),
+        pcp_dcp_kv_gather=False,
+    )
+    query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660])
+    common_metadata = SimpleNamespace(
+        num_actual_tokens=660,
+        query_start_loc_cpu=query_start_loc_cpu,
+        seq_lens=torch.arange(7, device=DEVICE_TYPE),
+    )
+    metadata = FlashMLASparseMetadata(
+        num_reqs=7,
+        max_query_len=110,
+        max_seq_len=110,
+        num_actual_tokens=660,
+        query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE),
+        slot_mapping=torch.arange(660, device=DEVICE_TYPE),
+        block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE),
+        req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE),
+        num_decodes=7,
+        num_prefills=0,
+        num_decode_tokens=660,
+    )
+
+    fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(
+        builder, common_metadata, metadata
+    )
+
+    assert fp8_metadata.num_decodes == 6
+    assert fp8_metadata.num_decode_tokens == 660
+    assert fp8_metadata.decode is not None
+    assert fp8_metadata.decode.decode_query_len == 110
+    torch.testing.assert_close(
+        fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE)
+    )
+
+
+@pytest.mark.parametrize("use_mixed_batch", [False, True])
+def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool):
+    num_decode_tokens = 2
+    num_batch_tokens = 5
+    q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE)
+    topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE)
+    kernel_q_shapes = []
+
+    def convert_indices(*args, **kwargs):  # noqa: ARG001
+        assert not kwargs.get("HAS_PREFILL_WORKSPACE", False)
+        if not kwargs.get("return_valid_counts", False):
+            return topk_indices
+        valid_counts = torch.full(
+            (num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE
+        )
+        return topk_indices, valid_counts
+
+    monkeypatch.setattr(
+        "vllm.v1.attention.backends.mla.flashmla_sparse."
+        "triton_convert_req_index_to_global_index",
+        convert_indices,
+    )
+
+    def run_kernel(**kwargs):
+        kernel_q_shapes.append(kwargs["q"].shape)
+        return kwargs["q"][..., :1], None
+
+    if use_mixed_batch:
+        fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
+            scheduler_metadata=object(),  # type: ignore[arg-type]
+            dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
             cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
         )
     else:
@@ -1679,6 +3481,7 @@ def run_kernel(**kwargs):
         fp8_extra_metadata=fp8_metadata,
         fp8_use_mixed_batch=use_mixed_batch,
         num_actual_tokens=num_batch_tokens,
+        num_decode_tokens=num_decode_tokens,
         req_id_per_token=torch.empty(
             num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE
         ),
@@ -1690,9 +3493,20 @@ def run_kernel(**kwargs):
         topk_indices_buffer=topk_indices,
         num_heads=2,
         kv_lora_rank=1,
+        index_group=None,
+        index_group_index=0,
         dcp_world_size=1,
+        pcp_dcp_kv_gather=False,
         need_to_return_lse_for_decode=False,
         _fp8_flash_mla_kernel=run_kernel,
+        _convert_logical_to_physical_topk=(
+            lambda indices, metadata, **kwargs: convert_indices(
+                metadata.req_id_per_token[: indices.shape[0]],
+                metadata.block_table,
+                indices,
+                return_valid_counts=kwargs["return_valid_counts"],
+            )
+        ),
     )
     impl._forward_fp8_kv_mixed_batch = MethodType(
         FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl
@@ -1743,6 +3557,7 @@ def _build_sparse_dcp_vllm_config(
     model_config = vllm_config.model_config
     model_config.dtype = torch.bfloat16
     model_config.hf_text_config = SimpleNamespace(
+        index_topk=topk_tokens,
         q_lora_rank=None,
         kv_lora_rank=kv_lora_rank,
         qk_nope_head_dim=qk_nope_head_dim,
@@ -1750,6 +3565,7 @@ def _build_sparse_dcp_vllm_config(
         v_head_dim=v_head_dim,
         model_type="deepseek_v2",
     )
+    del model_config.hf_config.index_topk  # Composite configs only nest this field.
     model_config.get_num_attention_heads = MethodType(
         lambda self, parallel_config: local_heads, model_config
     )
@@ -1846,6 +3662,8 @@ def run_kernel(**kwargs):
         cp_kv_cache_interleave_size=1,
     )
     impl = SimpleNamespace(
+        index_group=None,
+        index_group_index=0,
         dcp_world_size=2,
         dcp_rank=0,
         need_to_return_lse_for_decode=True,
@@ -1864,3 +3682,155 @@ def run_kernel(**kwargs):
     assert out.is_contiguous()
     assert not out.isnan().any()
     assert not lse.isnan().any()
+
+
+def test_hisparse_prefill_reuses_builder_staging_plan():
+    """Every layer must reuse the batch plan instead of synchronizing to dedupe."""
+    plan = SimpleNamespace(
+        block_table=torch.tensor([[0]], dtype=torch.int32),
+        ensure_gpu_sources=MagicMock(),
+    )
+    staged = torch.empty((1, 1, 8))
+    calls = []
+
+    def gather(kv_cache, staging_plan, resident_cache=None):
+        calls.append((kv_cache, staging_plan, resident_cache))
+        return staged
+
+    resident_cache = torch.empty((1, 1, 8))
+    resident_block_table = torch.tensor([[1]], dtype=torch.int32)
+    cache = SimpleNamespace(
+        runtime=SimpleNamespace(
+            gather_prefill_cache=gather,
+        ),
+        view=SimpleNamespace(cache=resident_cache, block_size=1),
+        block_table=resident_block_table,
+    )
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.caches = [cache]
+    source = torch.empty((1, 1, 8))
+    metadata = SimpleNamespace(
+        num_decodes=0,
+        num_decode_tokens=0,
+        seq_lens=torch.tensor([1], dtype=torch.int32),
+        prefill=SimpleNamespace(host_staging_plan=plan),
+        req_id_per_token=torch.tensor([0], dtype=torch.int32),
+    )
+
+    result, block_table, request_ids = index_group.stage_prefill_rows(
+        0, source, metadata
+    )
+
+    assert result is staged
+    assert block_table is plan.block_table
+    torch.testing.assert_close(request_ids, metadata.req_id_per_token)
+    plan.ensure_gpu_sources.assert_called_once()
+    args = plan.ensure_gpu_sources.call_args.args
+    torch.testing.assert_close(args[0], resident_block_table)
+    assert args[1] == 1
+    assert calls == [(source, plan, resident_cache)]
+
+
+def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch):
+    compute_stream = object()
+    prefill_stream = MagicMock()
+    gather = MagicMock()
+    monkeypatch.setattr(index_group_module, "current_stream", lambda: compute_stream)
+    monkeypatch.setattr(
+        index_group_module.ops, "cp_gather_and_upconvert_fp8_kv_cache", gather
+    )
+
+    resident_cache = torch.empty((2, 4, 656), dtype=torch.uint8)
+    ready = MagicMock()
+    cache = SimpleNamespace(
+        view=SimpleNamespace(cache=resident_cache, block_size=4),
+        block_table=torch.tensor([[2, 3]], dtype=torch.int32),
+    )
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.caches = [cache]
+    index_group.prefill_stream = prefill_stream
+    index_group.prefill_ready_events = [ready]
+    source = torch.empty((3, 4, 656), dtype=torch.uint8)
+    dst = torch.empty((8, 576), dtype=torch.bfloat16)
+    block_table = torch.tensor([[0, 1]], dtype=torch.int32)
+    workspace_starts = torch.tensor([0], dtype=torch.int32)
+    plan = SimpleNamespace(
+        block_table=torch.tensor([[0, 1]], dtype=torch.int32),
+        row_ids=torch.arange(8, dtype=torch.int32).view(1, -1),
+        gpu_row_ids=torch.full((1, 8), -1, dtype=torch.int32),
+        ensure_gpu_sources=MagicMock(),
+    )
+    metadata = SimpleNamespace(
+        prefill=SimpleNamespace(host_staging_plan=plan),
+        num_decodes=0,
+    )
+
+    result = index_group.gather_fp8_prefill(
+        0,
+        source,
+        dst,
+        block_table,
+        workspace_starts,
+        1,
+        metadata,
+        0,
+    )
+
+    prefill_stream.wait_stream.assert_called_once_with(compute_stream)
+    ready.record.assert_called_once_with(prefill_stream)
+    assert result is ready
+    gather.assert_called_once()
+    args = gather.call_args
+    assert args.args[0] is resident_cache
+    assert args.args[1] is dst
+    torch.testing.assert_close(args.args[2], plan.block_table)
+    assert args.args[3] is workspace_starts
+    assert args.args[4] == 1
+    assert args.kwargs["host_cache"].data_ptr() == source.data_ptr()
+    assert args.kwargs["host_row_ids"] is plan.row_ids
+    assert args.kwargs["device_row_ids"] is plan.gpu_row_ids
+    plan.ensure_gpu_sources.assert_called_once()
+    ensure_args = plan.ensure_gpu_sources.call_args.args
+    torch.testing.assert_close(ensure_args[0], cache.block_table)
+    assert ensure_args[1] == 4
+
+
+def test_sparse_impl_observes_repointed_indexer_buffer():
+    """The MTP proposer repoints the draft's indexer at the target model's buffer
+    after the backend impl is built, so the impl must resolve the buffer per read.
+    Snapshotting it in __init__ leaves the layer reading indices nothing writes."""
+    impl = object.__new__(FlashInferMLASparseImpl)
+    own = torch.zeros(4, 8, dtype=torch.int32)
+    target = torch.ones(4, 8, dtype=torch.int32)
+    indexer = SimpleNamespace(topk_indices_buffer=own)
+    impl.init_topk_indices_buffer(indexer, None)
+
+    assert impl.topk_indices_buffer is own
+
+    indexer.topk_indices_buffer = target
+
+    assert impl.topk_indices_buffer is target
+
+
+def test_explicit_topk_buffer_supersedes_indexer():
+    """Backbone skip-topk layers have no indexer, and the proposer also assigns the
+    shared buffer directly onto draft submodules."""
+    impl = object.__new__(FlashInferMLASparseImpl)
+    indexer = SimpleNamespace(
+        topk_indices_buffer=torch.zeros(2, 2, dtype=torch.int32),
+    )
+    impl.init_topk_indices_buffer(indexer, None)
+    explicit = torch.ones(2, 2, dtype=torch.int32)
+
+    impl.topk_indices_buffer = explicit
+
+    assert impl.topk_indices_buffer is explicit
+
+
+def test_sparse_mla_common_impl_resolves_buffer_lazily():
+    """Guards the whole SparseMLACommonImpl family at once: a backend that
+    snapshots the buffer instead silently loses MTP buffer sharing."""
+    assert issubclass(SparseMLACommonImpl, SharedTopkIndicesBuffer)
+    assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), (
+        "topk_indices_buffer must stay a lazily-resolved property"
+    )
diff --git a/tests/v1/attention/test_trtllm_attention_integration.py b/tests/v1/attention/test_trtllm_attention_integration.py
index d81bc2de4ea1..42dd8f9f480f 100644
--- a/tests/v1/attention/test_trtllm_attention_integration.py
+++ b/tests/v1/attention/test_trtllm_attention_integration.py
@@ -9,6 +9,7 @@
 import torch
 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 
+from tests.kernels.attention.test_trtllm_kvfp8_dequant import make_random_kv_cache
 from tests.v1.attention.test_attention_backends import create_and_prepopulate_kv_cache
 from tests.v1.attention.utils import (
     BatchSpec,
@@ -17,7 +18,11 @@
 )
 from vllm.config import set_current_vllm_config
 from vllm.platforms import current_platform
-from vllm.utils.torch_utils import nvfp4_kv_cache_full_dim, set_random_seed
+from vllm.utils.torch_utils import (
+    is_strictly_contiguous,
+    nvfp4_kv_cache_full_dim,
+    set_random_seed,
+)
 from vllm.v1.attention.backends.utils import PerLayerParameters
 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode
 
@@ -33,6 +38,7 @@
     FlashInferMetadataBuilder,
     FlashInferTrtllmAPIDecode,
     TRTLLMPrefill,
+    trtllm_prefill_attn_kvfp8_dequant,
 )
 
 
@@ -83,6 +89,70 @@ def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
     }
 
 
+@pytest.mark.parametrize(
+    "block_size,num_prefills,has_host_lengths,quantize_query",
+    [
+        (16, 1, True, False),
+        (64, 2, True, False),
+        (128, 2, True, False),
+        (16, 2, False, False),
+        (16, 2, True, True),
+    ],
+)
+@torch.inference_mode()
+def test_prefill_dequant_scratch_ignores_unused_capacity(
+    block_size, num_prefills, has_host_lengths, quantize_query, monkeypatch
+):
+    """Bound scratch by attended prefill KV, including cached prefixes."""
+    device = torch.device("cuda")
+    config = create_vllm_config(model_name=MODEL, max_model_len=32768)
+    config.cache_config.cache_dtype = "fp8"
+    config.cache_config.kv_cache_layout = "LBHNC"
+    config.attention_config.use_trtllm_attention = True
+    config.attention_config.disable_flashinfer_q_quantization = not quantize_query
+    monkeypatch.setattr(
+        "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
+        _mock_get_per_layer_parameters,
+    )
+    spec = FullAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=2,
+        head_size=64,
+        dtype=torch.uint8,
+        kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,
+    )
+    kv_cache, scale = make_random_kv_cache(32, 2, block_size, 64)
+    batch = BatchSpec(
+        seq_lens=[1024] + [17, 65][-num_prefills:],
+        query_lens=[1] + [2, 8][-num_prefills:],
+    )
+    common = create_common_attn_metadata(batch, block_size, device)
+    if not has_host_lengths:
+        common.seq_lens_cpu_upper_bound = None
+    live_pages = ((65 if has_host_lengths else 1024) + block_size - 1) // block_size
+    with set_current_vllm_config(config):
+        builder = FlashInferMetadataBuilder(spec, ["layer.0"], config, device)
+        for capacity in (64, 256):
+            # Unused columns can contain nonzero page IDs.
+            common.block_table_tensor = torch.randint(
+                1, 32, (batch.batch_size, capacity), dtype=torch.int32, device=device
+            )
+            metadata = builder.build(0, common)
+            assert isinstance(metadata.prefill, TRTLLMPrefill)
+            table = metadata.prefill.block_tables
+            assert is_strictly_contiguous(table)
+            assert (metadata.q_data_type_prefill == current_platform.fp8_dtype()) == (
+                quantize_query
+            )
+            if not quantize_query:
+                converted, _ = trtllm_prefill_attn_kvfp8_dequant(
+                    kv_cache, table, scale, scale, metadata.q_data_type_prefill
+                )
+                assert converted.shape[0] == num_prefills * live_pages + 1
+            width = capacity if quantize_query else live_pages
+            torch.testing.assert_close(table, common.block_table_tensor[1:, :width])
+
+
 def _create_nvfp4_hnd_kv_cache(
     k_contexts,
     v_contexts,
@@ -124,6 +194,7 @@ def _create_nvfp4_hnd_kv_cache(
 
     Returns:
         ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, LBHNC-strided).
+
     """
     # First create a bf16 cache so block tables are populated.
     bf16_cache = create_and_prepopulate_kv_cache(
diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py
index 590c9efb8ad6..cd062cb486b8 100644
--- a/tests/v1/attention/utils.py
+++ b/tests/v1/attention/utils.py
@@ -85,13 +85,6 @@ def create_common_attn_metadata(
     seq_lens_cpu = seq_lens.cpu()
     max_seq_len = int(seq_lens_cpu.max())
 
-    # Create computed tokens (context length for each sequence)
-    context_lens = [
-        batch_spec.seq_lens[i] - batch_spec.query_lens[i]
-        for i in range(batch_spec.batch_size)
-    ]
-    num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
-
     # Create block table and slot mapping
     max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size
     if arange_block_indices:
@@ -122,8 +115,6 @@ def create_common_attn_metadata(
         query_start_loc_cpu=query_start_loc_cpu,
         seq_lens=seq_lens,
         seq_lens_cpu_upper_bound=seq_lens_cpu,
-        _seq_lens_cpu=seq_lens_cpu,
-        _num_computed_tokens_cpu=num_computed_tokens_cpu,
         num_reqs=batch_spec.batch_size,
         num_actual_tokens=num_tokens,
         max_query_len=max_query_len,
@@ -202,7 +193,6 @@ def create_vllm_config(
     hf_config_override: dict | None = None,
 ) -> VllmConfig:
     """Create a VllmConfig for testing with reasonable defaults."""
-
     model_config = ModelConfig(
         model=model_name,
         tokenizer=model_name,
diff --git a/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py b/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py
new file mode 100644
index 000000000000..ff62abab8c86
--- /dev/null
+++ b/tests/v1/core/prefix_cache/test_hybrid_prefix_cache_hit_rate.py
@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Replay fixed chat histories and check per-turn cache hits and metrics."""
+
+import asyncio
+import json
+from math import lcm
+from uuid import uuid4
+
+import httpx
+import pytest
+from openai.types.chat import ChatCompletionMessageParam
+from prometheus_client.parser import text_string_to_metric_families
+
+from tests.utils import RemoteOpenAIServer
+from vllm.utils.math_utils import round_down
+
+MODEL = "Qwen/Qwen3.5-0.8B"
+NUM_CONVERSATIONS = 8
+MAX_TOKENS = 32
+# Long first turns leave room for hybrid cache alignment and MTP recomputation.
+MIN_PROMPT_TOKENS = 2048
+METRICS_TIMEOUT = 30
+
+
+def _user_turns(conversation: int) -> list[str]:
+    records = "\n".join(
+        f"Sensor {conversation}-{i}: temperature {20 + i % 10} C, "
+        f"humidity {40 + i % 20} percent."
+        for i in range(160 + conversation)
+    )
+    return [
+        f"Session {conversation}. Summarize these sensor readings briefly:\n{records}",
+        "Which readings should an operator investigate first? Explain briefly.",
+        "Suggest two checks the operator should perform next.",
+    ]
+
+
+async def _scheduler_block_size(client: httpx.AsyncClient) -> int:
+    """Return the coarsest boundary an expected cache hit can land on.
+
+    Mirrors ``scheduler_block_size`` in ``resolve_kv_cache_block_sizes``, which
+    the engine keeps internal: the LCM of the resolved group block sizes. A
+    finer ``prefix_match_unit`` only adds hits, so this stays conservative.
+    """
+    response = await client.get("/metrics")
+    response.raise_for_status()
+    for family in text_string_to_metric_families(response.text):
+        for sample in family.samples:
+            if sample.name != "vllm:cache_config_info":
+                continue
+            sizes = [
+                int(value)
+                for name in ("block_size", "mamba_block_size")
+                if (value := sample.labels.get(name, "None")) != "None"
+            ]
+            assert sizes, f"no block size in {sample.labels}"
+            return lcm(*sizes)
+    raise AssertionError("missing vllm:cache_config_info metric")
+
+
+async def _prefix_cache_counters(
+    client: httpx.AsyncClient,
+) -> tuple[float, float]:
+    response = await client.get("/metrics")
+    response.raise_for_status()
+    names = ("vllm:prefix_cache_queries_total", "vllm:prefix_cache_hits_total")
+    counters: dict[str, float] = {}
+    for family in text_string_to_metric_families(response.text):
+        for sample in family.samples:
+            if sample.name in names:
+                counters[sample.name] = counters.get(sample.name, 0) + sample.value
+    assert all(name in counters for name in names), (
+        f"missing prefix-cache metrics: {set(names) - counters.keys()}"
+    )
+    return counters[names[0]], counters[names[1]]
+
+
+async def _wait_for_counters(
+    client: httpx.AsyncClient, expected: tuple[float, float]
+) -> None:
+    """Allow engine stats to reach the frontend after completions are returned."""
+    deadline = asyncio.get_running_loop().time() + METRICS_TIMEOUT
+    while True:
+        actual = await _prefix_cache_counters(client)
+        if actual == expected:
+            return
+        assert all(a <= e for a, e in zip(actual, expected)), (
+            f"prefix-cache counters exceeded usage: {actual=}, {expected=}"
+        )
+        assert asyncio.get_running_loop().time() < deadline, (
+            f"prefix-cache counters did not match usage: {actual=}, {expected=}"
+        )
+        await asyncio.sleep(0.1)
+
+
+@pytest.fixture(scope="module")
+def server(request):
+    args = ["--enable-prompt-tokens-details"]
+    if request.param:
+        args += [
+            "--speculative-config",
+            json.dumps({"method": "mtp", "num_speculative_tokens": 1}),
+        ]
+    with RemoteOpenAIServer(MODEL, args) as remote_server:
+        yield remote_server
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server, dropped_blocks",
+    # MTP's prefill lookahead pollutes the tail block, so EAGLE drops one.
+    [pytest.param(False, 0, id="base"), pytest.param(True, 1, id="mtp")],
+    indirect=["server"],
+)
+async def test_prefix_cache_hit_rate(
+    server: RemoteOpenAIServer, dropped_blocks: int
+) -> None:
+    # Isolate sessions so another conversation cannot hide a cache miss.
+    salts = [uuid4().hex for _ in range(NUM_CONVERSATIONS)]
+    histories: list[list[ChatCompletionMessageParam]] = [
+        [] for _ in range(NUM_CONVERSATIONS)
+    ]
+    prompt_lengths: list[list[int]] = [[] for _ in range(NUM_CONVERSATIONS)]
+
+    async with (
+        server.get_async_client() as client,
+        httpx.AsyncClient(base_url=server.url_root, timeout=METRICS_TIMEOUT) as metrics,
+    ):
+        block_size = await _scheduler_block_size(metrics)
+        print(f"{block_size=}")
+
+        async def run_conversation(session: int, *, replay: bool) -> tuple[int, int]:
+            history = histories[session]
+            queries = hits = 0
+            for turn, user_text in enumerate(_user_turns(session)):
+                if not replay:
+                    history.append({"role": "user", "content": user_text})
+                # Reuse the original replies, even if replay generates different ones.
+                response = await client.chat.completions.create(
+                    model=MODEL,
+                    messages=history[: 2 * turn + 1],
+                    temperature=0,
+                    max_tokens=MAX_TOKENS,
+                    extra_body={"cache_salt": salts[session]},
+                )
+                usage = response.usage
+                assert usage is not None
+                assert usage.prompt_tokens_details is not None
+                prompt = usage.prompt_tokens
+                cached = usage.prompt_tokens_details.cached_tokens
+                context = f"{session=}, {turn=}, {replay=}: {cached=}, {prompt=}"
+                assert cached is not None, context
+                assert 0 <= cached < prompt, context
+                assert prompt >= MIN_PROMPT_TOKENS, context
+                if replay:
+                    assert prompt == prompt_lengths[session][turn], context
+                    # The identical request computed this whole prompt, but at
+                    # least one token is always recomputed.
+                    computed = prompt - 1
+                elif turn:
+                    computed = prompt_lengths[session][turn - 1]
+                else:
+                    assert cached == 0, context
+                    computed = 0
+                # Every token below the last reachable boundary must be reused.
+                floor = max(
+                    0, round_down(computed, block_size) - dropped_blocks * block_size
+                )
+                print(
+                    f"{session=} {turn=} {replay=} {prompt=} {cached=} {floor=} "
+                    f"hit_rate={cached / prompt:.4f}"
+                )
+                assert cached >= floor, (
+                    f"{context}, {computed=}, {block_size=}, "
+                    f"hit_rate={cached / prompt:.4f}, "
+                    f"expected_hit_rate={floor / prompt:.4f}"
+                )
+                if not replay:
+                    prompt_lengths[session].append(prompt)
+                    history.append(
+                        {
+                            "role": "assistant",
+                            "content": response.choices[0].message.content or "",
+                        }
+                    )
+                queries += prompt
+                hits += cached
+            return queries, hits
+
+        queries, hits = await _prefix_cache_counters(metrics)
+        for replay in (False, True):
+            counts = await asyncio.gather(
+                *(run_conversation(i, replay=replay) for i in range(NUM_CONVERSATIONS))
+            )
+            queries += sum(prompt for prompt, _ in counts)
+            hits += sum(cached for _, cached in counts)
+            await _wait_for_counters(metrics, (queries, hits))
diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py
index fd003fdce6ff..4cb9afe87138 100644
--- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py
+++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py
@@ -6,7 +6,7 @@
 candidate and drops it. Mamba materializes state only on its own block grid, so
 the pruned position holds nothing and the hit floors back a whole Mamba block.
 
-Gated by ``CacheConfig.enable_mamba_fine_grained_prefix_cache``. Every test here
+Gated by ``CacheConfig.enable_mamba_shared_prefix_checkpoint``. Every test here
 drives the real ``Scheduler._mamba_block_aligned_split`` and the real
 ``KVCacheManager``; no chunk boundary is hard-coded.
 """
@@ -32,7 +32,7 @@ def _manager(
     block_size,
     hash_block_size,
     *,
-    fine_grained=True,
+    shared_prefix_checkpoint=True,
     num_blocks=8192,
     eagle_group=None,
     num_prefill_lookahead=0,
@@ -52,7 +52,7 @@ def _manager(
         hash_block_size=hash_block_size,
         use_eagle=True,
         num_prefill_lookahead=num_prefill_lookahead,
-        enable_mamba_fine_grained_prefix_cache=fine_grained,
+        enable_mamba_shared_prefix_checkpoint=shared_prefix_checkpoint,
     )
 
 
@@ -76,8 +76,8 @@ def _stub(manager, block_size, hash_block_size, *, block_drop=True):
         hash_block_size=hash_block_size,
         mamba_has_prefill_checkpoint_blocks=False,  # forced False under eagle
         mamba_partial_cache_hit=partial_hit,
-        mamba_fine_grained_prefix_cache=(
-            partial_hit and manager.mamba_fine_grained_prefix_cache
+        mamba_shared_prefix_checkpoint=(
+            partial_hit and manager.mamba_shared_prefix_checkpoint
         ),
     )
 
@@ -151,7 +151,7 @@ def test_sibling_resumes_from_the_observed_junction():
     block_size, hash_block_size = 512, 32
     manager = _manager(block_size, hash_block_size)
     stub = _stub(manager, block_size, hash_block_size)
-    assert stub.mamba_fine_grained_prefix_cache, "the feature must be armed"
+    assert stub.mamba_shared_prefix_checkpoint, "the feature must be armed"
     _orphaned_full_attention_tail(manager, stub, 2020)
 
     consumer = make_request(
@@ -299,7 +299,14 @@ def test_a_junction_past_the_prompt_falls_back_and_registers_nothing():
     running = request.num_computed_tokens
     assert not mamba.req_to_blocks["r"][running // block_size].is_null
     request.shared_prefix_boundary = running
-    assert mamba._cache_partial_tail_block(request, running) is None
+    assert (
+        mamba._cache_partial_tail_block(
+            request,
+            running,
+            retention_interval=manager.coordinator.retention_interval,
+        )
+        is None
+    )
 
 
 @pytest.mark.parametrize("eagle_group", [None, 0])
@@ -318,11 +325,11 @@ def test_enabling_never_reduces_reuse(eagle_group, num_prefill_lookahead):
     """
     block_size, hash_block_size = 512, 32
 
-    def sibling_hit(fine_grained):
+    def sibling_hit(shared_prefix_checkpoint):
         manager = _manager(
             block_size,
             hash_block_size,
-            fine_grained=fine_grained,
+            shared_prefix_checkpoint=shared_prefix_checkpoint,
             eagle_group=eagle_group,
             num_prefill_lookahead=num_prefill_lookahead,
         )
@@ -341,7 +348,7 @@ def sibling_hit(fine_grained):
 def test_disabled_by_default():
     """Off, the junction is block-floored and no resume point is registered."""
     block_size, hash_block_size = 512, 32
-    manager = _manager(block_size, hash_block_size, fine_grained=False)
+    manager = _manager(block_size, hash_block_size, shared_prefix_checkpoint=False)
     stub = _stub(manager, block_size, hash_block_size)
     _orphaned_full_attention_tail(manager, stub, 2020)
 
diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
index c020a5db5483..f1301d928837 100644
--- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
@@ -41,7 +41,9 @@ def test_connector_without_divergent_hit_support_uses_common_lookup():
     manager = MagicMock()
     manager.get_computed_blocks.return_value = (common_blocks, 0, 0)
     scheduler = SimpleNamespace(
-        connector=SimpleNamespace(supports_divergent_local_hybrid_hits=False),
+        connector=SimpleNamespace(
+            supports_divergent_local_hybrid_hits=False,
+        ),
         kv_cache_manager=manager,
     )
 
@@ -61,7 +63,9 @@ def test_capable_connector_uses_divergent_partial_hit_lookup():
         True,
     )
     scheduler = SimpleNamespace(
-        connector=SimpleNamespace(supports_divergent_local_hybrid_hits=True),
+        connector=SimpleNamespace(
+            supports_divergent_local_hybrid_hits=True,
+        ),
         kv_cache_manager=manager,
     )
 
@@ -184,7 +188,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
         dcp_world_size=dcp_world_size,
         scheduler_block_size=scheduler_block_size,
         mamba_partial_cache_hit=True,
-        mamba_fine_grained_prefix_cache=False,
+        mamba_shared_prefix_checkpoint=False,
         mamba_has_prefill_checkpoint_blocks=False,
     )
     split = Scheduler._mamba_block_aligned_split
@@ -232,7 +236,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():
         use_eagle_block_drop=False,
         hash_block_size=32,
         mamba_partial_cache_hit=False,
-        mamba_fine_grained_prefix_cache=False,
+        mamba_shared_prefix_checkpoint=False,
         mamba_has_prefill_checkpoint_blocks=False,
     )
     req = make_request("0", [0] * prompt_length, 32, sha256)
@@ -273,7 +277,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
         use_eagle_block_drop=False,
         hash_block_size=32,
         mamba_partial_cache_hit=False,
-        mamba_fine_grained_prefix_cache=False,
+        mamba_shared_prefix_checkpoint=False,
         mamba_has_prefill_checkpoint_blocks=False,
     )
     req = make_request("0", [0] * prompt_length, 32, sha256)
@@ -611,6 +615,77 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle():
     assert num_computed == 112
 
 
+@pytest.mark.parametrize(
+    "retention_interval,transient_published",
+    [(None, True), (0, False)],
+)
+def test_internal_checkpoint_publication_respects_retention(
+    retention_interval, transient_published
+):
+    """With retention_interval=0, a mid-prompt internal checkpoint is
+    request-local and must not enter the prefix cache; the prompt-end
+    checkpoint keeps the existing publication behavior. Dense retention
+    (None) is unchanged.
+    """
+    hash_block_size = 16
+    manager = make_full_mamba_manager(
+        dcp_world_size=1,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=32,
+        num_prefill_checkpoint_blocks=1,
+    )
+    manager.coordinator.retention_interval = retention_interval
+
+    request = make_request("producer", list(range(240)), hash_block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
+
+    # Mid-prompt chunk: exports a transient checkpoint at 112.
+    assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
+    transient_hash = request.block_hashes[112 // hash_block_size - 1]
+    transient_hit = manager.block_pool.get_cached_block(transient_hash, [1])
+    assert (transient_hit is not None) == transient_published
+
+    request.num_computed_tokens = 128
+    manager.new_step_starts()
+
+    # Final chunk: the prompt-end checkpoint at 224 stays published.
+    assert manager.allocate_slots(request, 112) is not None
+    end_hash = request.block_hashes[224 // hash_block_size - 1]
+    assert manager.block_pool.get_cached_block(end_hash, [1]) is not None
+
+
+def test_transient_checkpoint_evicts_retained_boundary_hash():
+    """With retention_interval=0, the checkpoint slot can coincide with a
+    reachable boundary block that this step's full-block pass just hashed.
+    The checkpoint state overwrites the slot, so the stale boundary hash
+    must be evicted rather than left pointing at the wrong state.
+    """
+    hash_block_size = 16
+    manager = make_full_mamba_manager(
+        dcp_world_size=1,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=32,
+        num_prefill_checkpoint_blocks=1,
+    )
+    manager.coordinator.retention_interval = 0
+
+    request = make_request("producer", list(range(240)), hash_block_size, sha256)
+    # Block 2 ends at the shared-prefix junction 96, so the full-block pass
+    # hashes it as a retained boundary.
+    request.shared_prefix_boundary = 96
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
+    assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
+
+    # The same slot is the transient checkpoint block (state@112): both the
+    # boundary hash and the checkpoint itself must stay unpublished.
+    boundary_hash = request.block_hashes[96 // hash_block_size - 1]
+    assert manager.block_pool.get_cached_block(boundary_hash, [1]) is None
+    checkpoint_hash = request.block_hashes[112 // hash_block_size - 1]
+    assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None
+
+
 def test_eagle_block_aligned_checkpoint_replaces_newer_hash():
     hash_block_size = mamba_block_size = 32
     manager = make_full_mamba_manager(
@@ -922,71 +997,6 @@ def request_finished_all_groups(self, request, block_ids):
     scheduler.kv_cache_manager.remove_skipped_blocks.assert_called_once()
 
 
-def test_block_pool_touch_pins_released_cow_target():
-    """The connector can rescue an offered CoW target after its step-scoped
-    retention is released by using the bound BlockPool's touch method."""
-    hash_block_size = 2
-    block_size = 2 * hash_block_size
-    kv_cache_config = KVCacheConfig(
-        num_blocks=24,
-        kv_cache_tensors=[],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["full"],
-                FullAttentionSpec(
-                    block_size=hash_block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                ),
-            ),
-            KVCacheGroupSpec(
-                ["mamba"],
-                MambaSpec(
-                    block_size=block_size,
-                    shapes=(1, 1),
-                    dtypes=(torch.float32,),
-                    mamba_cache_mode="align",
-                ),
-            ),
-        ],
-    )
-    manager = make_kv_cache_manager(
-        kv_cache_config=kv_cache_config,
-        max_model_len=8192,
-        enable_caching=True,
-        hash_block_size=hash_block_size,
-    )
-    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
-    computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
-    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
-    req0.num_computed_tokens = 6
-    req0.append_output_token_ids([3])
-    assert manager.allocate_slots(req0, 1) is not None
-
-    # Retention released before the drain (defer_block_free=False ordering).
-    _copies, retained = manager.take_kv_cache_block_copies()
-    manager.block_pool.free_blocks(retained)
-
-    offloads = drain_boundary_state_offloads(manager)
-    ((_group_id, block_id, boundary_tokens),) = offloads["0"]
-    assert boundary_tokens == 6
-    cow_block = manager.block_pool.blocks[block_id]
-    assert cow_block.ref_cnt == 0
-    assert block_id in _free_block_ids(manager)
-
-    manager.block_pool.touch([cow_block])
-    assert cow_block.ref_cnt == 1
-    assert block_id not in _free_block_ids(manager)
-
-    # The connector-pinned block is out of the free queue: draining every free block
-    # neither trips the allocator's ref_cnt assert nor hands it out.
-    new_blocks = manager.block_pool.get_new_blocks(
-        manager.block_pool.get_num_free_blocks()
-    )
-    assert block_id not in {b.block_id for b in new_blocks}
-
-
 def test_boundary_state_offload_dropped_when_request_freed_before_drain():
     """A hand-off recorded in the same scheduling pass as the request's death
     must not be drained: its release hook has already run, so draining would
@@ -1816,7 +1826,7 @@ def test_mamba_align_split_stops_below_eagle_proof_boundary():
         use_eagle_block_drop=True,
         hash_block_size=hash_block_size,
         mamba_partial_cache_hit=True,
-        mamba_fine_grained_prefix_cache=False,
+        mamba_shared_prefix_checkpoint=False,
         mamba_has_prefill_checkpoint_blocks=False,
     )
     split = Scheduler._mamba_block_aligned_split
@@ -2386,8 +2396,8 @@ def test_retention_snapshots_handed_off_with_exact_block_ids():
 def test_snapshot_handoff_dense_default_retention():
     """With dense (default) retention, every materialized mamba boundary
     state block is handed off — regular mamba-align + prefix-match-unit
-    deployments get store-able boundary snapshots without setting
-    VLLM_PREFIX_CACHE_RETENTION_INTERVAL."""
+    deployments get store-able boundary snapshots without setting a
+    prefix-cache retention interval."""
     hash_block_size = 2
     block_size = 4
     manager = make_kv_cache_manager(
diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py
index 7f256f025e6d..52167a507317 100644
--- a/tests/v1/core/test_async_scheduler.py
+++ b/tests/v1/core/test_async_scheduler.py
@@ -13,7 +13,7 @@
 from vllm.v1.structured_output import StructuredOutputGrammar
 from vllm.v1.utils import ConstantList
 
-from .utils import create_requests, create_scheduler
+from .utils import create_requests, create_scheduler, mock_kv
 
 pytestmark = pytest.mark.cpu_test
 
@@ -145,40 +145,30 @@ def abort_request():
         assert req.num_output_tokens == abort_order_copy.index(i)
 
 
-def test_preempt():
-    scheduler = create_scheduler(async_scheduling=True)
-    requests = create_requests(num_requests=10, max_tokens=20)
-
-    for req in requests:
-        scheduler.add_request(req)
-
-    sched_outputs: deque[SchedulerOutput] = deque()
-    sched_outputs.append(scheduler.schedule())
-    sched_outputs.append(scheduler.schedule())
+def test_connector_metadata_precedes_async_placeholder_advance(monkeypatch):
+    """Mirroring must see earlier unresolved outputs, excluding the current step."""
+    scheduler = create_scheduler(
+        async_scheduling=True,
+        use_kv_connector=mock_kv(matched_tokens=0, is_async=False),
+    )
+    (request,) = create_requests(num_requests=1, num_tokens=4, max_tokens=4)
+    scheduler.add_request(request)
 
-    abort_order = [0, 8, 3, 1, 6, 4, 2, 5, 7, 9]
-    abort_order_copy = abort_order.copy()
+    assert scheduler.connector is not None
+    build_connector_meta = scheduler.connector.build_connector_meta
+    observed_placeholders = []
 
-    def abort_request():
-        if not abort_order:
-            return
-        req = requests[abort_order.pop(0)]
-        scheduler.finish_requests(req.request_id, RequestStatus.FINISHED_ABORTED)
+    def record_placeholders(scheduler_output):
+        observed_placeholders.append(request.num_output_placeholders)
+        return build_connector_meta(scheduler_output)
 
-    while sched_outputs:
-        # Abort a scheduled request.
-        abort_request()
-        sched_output = sched_outputs.popleft()
-        model_runner_output = _make_model_runner_output(sched_output)
-        scheduler.update_from_output(sched_output, model_runner_output)
-
-        sched_output = scheduler.schedule()
-        if sched_output.num_scheduled_tokens:
-            sched_outputs.append(sched_output)
+    monkeypatch.setattr(
+        scheduler.connector, "build_connector_meta", record_placeholders
+    )
+    scheduler.schedule()
+    scheduler.schedule()
 
-    for i, req in enumerate(requests):
-        assert req.status == RequestStatus.FINISHED_ABORTED
-        assert req.num_output_tokens == abort_order_copy.index(i)
+    assert observed_placeholders == [0, 1]
 
 
 def test_prefix_caching_for_prefill_dedup():
@@ -308,7 +298,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     request = create_requests(num_requests=1, num_tokens=1)[0]
     request.structured_output_request = Mock()
     request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
-    request.structured_output_request.grammar.accept_tokens.return_value = False
     request.status = RequestStatus.RUNNING
     request.num_computed_tokens = request.num_tokens
     request.num_output_placeholders = 1
@@ -317,10 +306,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.connector = None
     scheduler.ec_connector = None
     scheduler.structured_output_manager = Mock()
-    scheduler.structured_output_manager.should_advance.return_value = True
-    scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = (
-        lambda request, new_token_ids: new_token_ids
-    )
+    scheduler.structured_output_manager.accept_tokens.return_value = False
     scheduler.requests = {request.request_id: request}
     scheduler.running = [request]
     scheduler.waiting = Mock()
@@ -331,9 +317,9 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.finished_req_ids = set()
     scheduler.finished_req_ids_dict = None
     scheduler.grammar_compile_error_reqs = set()
+    scheduler.encoder_cache_mismatch_reqs = set()
     scheduler.vllm_config = Mock()
-    scheduler.vllm_config.model_config.enable_return_routed_experts = False
-    scheduler.enable_return_routed_experts = False
+    scheduler.aux_output_connector = None
     scheduler.return_sampling_mask = False
     scheduler.recompute_kv_load_failures = False
     scheduler.defer_block_free = False
@@ -701,16 +687,29 @@ def test_requires_kv_delivery_defaults_to_producer_role():
         assert scheduler.requires_kv_delivery is expected, role
 
 
-@pytest.mark.parametrize("kv_role", ["kv_producer", "kv_consumer"])
-def test_kv_pressure_preempt_mid_handoff(kv_role: str):
-    """P/D race: a request is KV-pressure preempted while the output of its
-    final prefill chunk -- the hand-off token that would finish it -- is in
-    flight.
-
-    On a producer, that output must be dropped so the request recomputes;
-    delivering it would finish the request and hand off blocks the preemption
-    already freed, so the consumer pulls garbage. A consumer hands nothing off,
-    so it keeps the lossless deliver-stale path.
+@pytest.mark.parametrize(
+    ("kv_role", "defer_free"),
+    [
+        ("kv_producer", False),
+        ("kv_consumer", False),
+        ("kv_consumer", True),
+    ],
+)
+def test_kv_pressure_preempt_mid_handoff(kv_role: str, defer_free: bool):
+    """P/D race: KV pressure hits while the output of a request's final
+    prefill chunk -- the hand-off token that would finish it -- is in flight.
+
+    When the victim's blocks free immediately, the request is preempted. On a
+    producer, that output must be dropped so the request recomputes; delivering
+    it would finish the request and hand off blocks the preemption already
+    freed, so the consumer pulls garbage. A consumer hands nothing off, so it
+    keeps the lossless deliver-stale path.
+
+    A consumer with overlapping batches (async scheduling or PP) instead fences
+    the victim's free behind its in-flight output, so the allocation retry
+    stops instead of preempting; the request then finishes from that output
+    once it lands. The gate depends on the platform (async scheduling is
+    force-disabled on CPU), so force the flag to cover both paths everywhere.
     """
     is_producer = kv_role == "kv_producer"
     scheduler = create_scheduler(
@@ -722,10 +721,13 @@ def test_kv_pressure_preempt_mid_handoff(kv_role: str):
         max_num_batched_tokens=512,
     )
     assert scheduler.requires_kv_delivery is is_producer
+    # The production gate requires overlapping batches, which async
+    # scheduling only provides off-CPU.
+    scheduler.defer_block_free = defer_free
 
     # 32-token prompts fill 2 blocks each, exhausting the usable pool, so the
-    # next decode allocation preempts the tail of the running queue (the handoff
-    # request) while its prefill output is still in flight.
+    # next decode allocation targets the tail of the running queue (the
+    # handoff request) while its prefill output is still in flight.
     decoder = create_requests(
         num_requests=1, num_tokens=32, max_tokens=8, req_ids=["decoder"]
     )[0]
@@ -739,9 +741,17 @@ def test_kv_pressure_preempt_mid_handoff(kv_role: str):
     assert handoff.num_output_placeholders == 1
 
     scheduler.schedule()
-    assert handoff.status == RequestStatus.PREEMPTED
-    assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens
-    assert handoff.drop_stale_output is is_producer
+    if defer_free:
+        # The victim's blocks are fenced behind its in-flight output, so
+        # preempting it could not satisfy the allocation; the retry stops
+        # instead of preempting.
+        assert handoff.status == RequestStatus.RUNNING
+        assert handoff.num_stale_output_tokens == 0
+    else:
+        # Blocks free immediately, so the handoff request is preempted.
+        assert handoff.status == RequestStatus.PREEMPTED
+        assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens
+        assert handoff.drop_stale_output is is_producer
 
     scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output))
 
diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py
index 5249de7f8c11..58d4e20bf70c 100644
--- a/tests/v1/core/test_contiguous_kv_packing.py
+++ b/tests/v1/core/test_contiguous_kv_packing.py
@@ -75,6 +75,7 @@ def _mock_vllm_config(layout: str | None):
     config.cache_config = CacheConfig()
     config.cache_config.num_gpu_blocks_override = None
     config.cache_config.kv_cache_layout = layout
+    config.attention_config.hisparse_config = None
     return config
 
 
@@ -398,6 +399,56 @@ def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wid
         assert len(gdn) == 10
 
 
+class TestSlidingWindowBucketCap:
+    def test_sliding_window_bucket_is_capped_at_the_main_page(self):
+        """DeepSeek-V4.1 shape: 43 SlidingWindowMLASpec SWA caches beside an
+        unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would
+        set the block width and the MLA group would fill under a third of
+        every block; capped, it splits into groups no wider than the MLA
+        page."""
+        config = _shared_layout_config()
+        specs: dict[str, KVCacheSpec] = {}
+        for layer, ratio in ((2, 2), (8, 2), (14, 2), (20, 1)):
+            specs[f"layers.{layer}.attn"] = MLAAttentionSpec(
+                block_size=128,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                tokens_per_state=ratio,
+                alignment=576,
+            )
+            specs[f"layers.{layer}.idx"] = MLAAttentionSpec(
+                block_size=128,
+                num_kv_heads=1,
+                head_size=132,
+                dtype=torch.uint8,
+                tokens_per_state=ratio,
+                alignment=576,
+            )
+        for layer in range(43):
+            specs[f"layers.{layer}.swa"] = SlidingWindowMLASpec(
+                block_size=64,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                sliding_window=128,
+                alignment=576,
+            )
+        groups = get_kv_cache_groups(config, specs)
+        pages = _pages(groups)
+        main = next(g for g in groups if "layers.2.attn" in g.layer_names)
+        main_bytes = sum(pages[n] for n in main.layer_names)
+        assert _get_kv_cache_bytes_per_block(groups) == main_bytes
+        swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")]
+        per_group = main_bytes // pages["layers.0.swa"]
+        assert len(swa_groups) == -(-43 // per_group)
+        for g in swa_groups:
+            assert sum(pages[n] for n in g.layer_names) <= main_bytes
+        assert sorted(n for g in swa_groups for n in g.layer_names) == sorted(
+            f"layers.{i}.swa" for i in range(43)
+        )
+
+
 class TestDensePacking:
     def test_bytes_per_block_is_largest_group(self):
         groups, g1, g2 = _mixed_page_groups()
@@ -548,3 +599,119 @@ def test_bound_views_round_trip(self, layout):
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])
+
+
+class TestCompressorRingGroup:
+    def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self):
+        """DeepSeek-V4.1 compressor state: a CircularBufferSpec ring (capacity
+        8) beside 128-token paged MLA groups and 64-token SWA groups. The ring
+        claims one block per request, never hashes, and does not disturb the
+        scheduler/hash block sizes the paged groups imply."""
+        config = _shared_layout_config()
+        config.cache_config.enable_prefix_caching = True
+        specs: dict[str, KVCacheSpec] = {
+            "layers.2.attn": MLAAttentionSpec(
+                block_size=128,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                tokens_per_state=2,
+                alignment=576,
+            ),
+            "layers.2.attn.compressor.state_cache": CircularBufferSpec(
+                block_size=8,
+                num_kv_heads=1,
+                head_size=1024,
+                head_size_v=0,
+                dtype=torch.float32,
+            ),
+        }
+        for layer in range(4):
+            specs[f"layers.{layer}.attn.swa_cache"] = SlidingWindowMLASpec(
+                block_size=64,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                sliding_window=128,
+                alignment=576,
+            )
+        groups = get_kv_cache_groups(config, specs)
+        ring_groups = [g for g in groups if "state_cache" in g.layer_names[0]]
+        assert len(ring_groups) == 1
+        assert not ring_groups[0].kv_cache_spec.prefix_cacheable
+        kv_cache_config = get_kv_cache_config_from_groups(
+            config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups)
+        )
+        assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64)
+        manager = KVCacheManager(
+            generate_scheduler_kv_cache_config([kv_cache_config]),
+            max_model_len=8192,
+            enable_caching=True,
+            hash_block_size=64,
+            scheduler_block_size=128,
+        )
+        assert len(manager.coordinator.single_type_managers) == len(groups)
+
+
+class TestSWABoundedReplayGrouping:
+    def test_non_cacheable_swa_leaves_hits_to_the_paged_group(self):
+        """Under SWA bounded replay only the paged MLA group hashes (hash
+        block = its block) and the coordinator probes one spec group."""
+        config = _shared_layout_config()
+        config.cache_config.enable_prefix_caching = True
+        specs: dict[str, KVCacheSpec] = {
+            "layers.2.attn": MLAAttentionSpec(
+                block_size=128,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                tokens_per_state=2,
+                alignment=576,
+            ),
+        }
+        swa = SlidingWindowMLASpec(
+            block_size=32,
+            num_kv_heads=1,
+            head_size=584,
+            dtype=torch.uint8,
+            sliding_window=128,
+            alignment=576,
+            bounded_replay=True,
+        )
+        for layer in range(4):
+            specs[f"layers.{layer}.attn.swa_cache"] = swa
+
+        groups = get_kv_cache_groups(config, specs)
+        # The worker reads the replay window off the (packed) group specs to
+        # arm the window clamp; the wrapper must forward it like cacheability.
+        assert sorted(
+            (g.kv_cache_spec.prefix_cacheable, g.kv_cache_spec.prefix_replay_tokens)
+            for g in groups
+        ) == [(False, 128)] * 4 + [(True, 0)]
+        kv_cache_config = get_kv_cache_config_from_groups(
+            config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups)
+        )
+        manager = KVCacheManager(
+            generate_scheduler_kv_cache_config([kv_cache_config]),
+            max_model_len=8192,
+            enable_caching=True,
+            hash_block_size=128,
+            scheduler_block_size=128,
+        )
+        assert len(manager.coordinator.attention_groups) == 1
+        assert isinstance(
+            manager.coordinator.attention_groups[0].spec, MLAAttentionSpec
+        )
+
+        # Allocating and caching a request goes through every group's manager;
+        # the 32-token SWA group must skip hashing rather than assert on the
+        # 128-token hash block.
+        from tests.v1.core.utils import create_requests
+
+        request = create_requests(num_requests=1, num_tokens=256, block_size=128)[0]
+        assert manager.allocate_slots(request, 256) is not None
+        request.num_computed_tokens = 256
+        manager.cache_blocks(request, 256)
+        for single in manager.coordinator.single_type_managers:
+            cached = single.num_cached_block.get(request.request_id, 0)
+            assert cached == (2 if single.kv_cache_spec.prefix_cacheable else 0)
diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py
index 13789a396015..44cdf567f7c9 100644
--- a/tests/v1/core/test_deferred_block_free.py
+++ b/tests/v1/core/test_deferred_block_free.py
@@ -48,13 +48,17 @@ def _make_model_runner_output(
     )
 
 
-def _create_deferring_scheduler():
+def _create_deferring_scheduler(scheduling_policy="fcfs"):
     """Async scheduler with deferred block freeing forced on.
 
     The production gate additionally requires a PD KV-consumer connector;
     the mechanism itself is independent of it.
     """
-    scheduler = create_scheduler(model=MODEL, async_scheduling=True)
+    scheduler = create_scheduler(
+        model=MODEL,
+        async_scheduling=True,
+        scheduling_policy=scheduling_policy,
+    )
     scheduler.defer_block_free = True
     return scheduler
 
@@ -79,6 +83,24 @@ def _setup_request_with_inflight_step(scheduler, max_tokens: int = 5):
     return request, out0, out1
 
 
+def _fail_one_allocation(scheduler, call_number: int):
+    real_allocate = scheduler.kv_cache_manager.allocate_slots
+    calls = 0
+
+    def allocate(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        if calls == call_number:
+            return None
+        return real_allocate(*args, **kwargs)
+
+    return patch.object(
+        scheduler.kv_cache_manager,
+        "allocate_slots",
+        side_effect=allocate,
+    )
+
+
 def test_gate_enabled_for_async_consumer():
     # Overlapping batches + consumer-side connector enables the gate. Async
     # scheduling (which would give >1 concurrent batches) is force-disabled on
@@ -224,6 +246,54 @@ def test_preempt_defers_free_and_clears_bookkeeping():
     assert pool.get_num_free_blocks() == num_free_initially
 
 
+@pytest.mark.parametrize(
+    ("policy", "failed_call"),
+    [
+        ("fcfs", 1),
+        ("priority", 2),
+    ],
+)
+def test_allocation_retry_waits_for_fence_then_succeeds(policy, failed_call):
+    scheduler = _create_deferring_scheduler(policy)
+    requests = create_requests(
+        num_requests=3,
+        num_tokens=NUM_PROMPT_TOKENS,
+        max_tokens=10,
+        stop_token_ids=[STOP_TOKEN_ID],
+    )
+    for request in requests:
+        scheduler.add_request(request)
+    out0 = scheduler.schedule()
+    out1 = scheduler.schedule()
+
+    worst, trigger, tail = requests
+    worst.priority, trigger.priority, tail.priority = 9, 0, 1
+
+    with _fail_one_allocation(scheduler, failed_call) as allocate:
+        blocked = scheduler.schedule()
+    assert allocate.call_count == failed_call
+    assert not blocked.preempted_req_ids
+    retry_trigger = requests[failed_call - 1]
+    assert retry_trigger.request_id not in blocked.num_scheduled_tokens
+
+    for output in (out0, out1, blocked):
+        if output.total_num_scheduled_tokens:
+            scheduler.update_from_output(
+                output,
+                _make_model_runner_output(output),
+            )
+
+    expected_victim = worst if policy == "priority" else tail
+    with _fail_one_allocation(scheduler, failed_call) as allocate:
+        resumed = scheduler.schedule()
+
+    assert allocate.call_count > failed_call
+    assert expected_victim.request_id in resumed.preempted_req_ids
+    assert retry_trigger.request_id in resumed.num_scheduled_tokens
+    if policy == "priority":
+        assert tail.request_id in resumed.num_scheduled_tokens
+
+
 def test_multiple_deferred_frees_drain_in_order():
     scheduler = _create_deferring_scheduler()
     pool = scheduler.kv_cache_manager.block_pool
@@ -347,7 +417,7 @@ def test_abort_mid_prefill_defers_free():
     chunk is still in flight must withhold its blocks.
     """
     scheduler = create_scheduler(
-        model=MODEL, async_scheduling=True, long_prefill_token_threshold=16
+        model=MODEL, async_scheduling=True, max_num_batched_tokens=16
     )
     scheduler.defer_block_free = True
     pool = scheduler.kv_cache_manager.block_pool
diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py
index 3948f7c6e300..3c61f3367f98 100644
--- a/tests/v1/core/test_encoder_cache_manager.py
+++ b/tests/v1/core/test_encoder_cache_manager.py
@@ -114,6 +114,7 @@ def test_eviction_when_cache_is_full():
 
     # 'x' should have been evicted.
     assert "x" not in manager.cached
+    assert "x" not in manager.num_embeds_by_hash
     assert "req1" not in manager.request_cached_ids
     assert "x" in manager.get_freed_mm_hashes()
 
@@ -273,6 +274,31 @@ def get_num_encoder_embeds(self, input_id: int) -> int:
     assert savings_ratio == 12.5
 
 
+def test_encoder_cache_tracks_embedding_count_by_identifier():
+    class MockRequestWithMask(MockRequest):
+        def get_num_encoder_embeds(self, input_id: int) -> int:
+            return self.mm_features[input_id].mm_position.get_num_embeds()
+
+    is_embed = torch.tensor([False, True, True, False, True, False])
+    request = MockRequestWithMask("owner", ["img1"], [len(is_embed)])
+    request.mm_features[0] = MultiModalFeatureSpec(
+        data=None,
+        modality="image",
+        identifier="img1",
+        mm_position=PlaceholderRange(offset=0, length=6, is_embed=is_embed),
+    )
+
+    manager = EncoderCacheManager(cache_size=10)
+    manager.allocate(request, 0)
+    assert manager.get_cached_num_encoder_embeds(request, 0) == 3
+
+    manager.free(request)
+    assert manager.get_cached_num_encoder_embeds(request, 0) == 3
+
+    manager.reset()
+    assert manager.get_cached_num_encoder_embeds(request, 0) is None
+
+
 def test_encoder_cache_mask_based_retrieval():
     class MockRequestWithMask(MockRequest):
         def get_num_encoder_embeds(self, input_id: int) -> int:
@@ -336,6 +362,7 @@ def test_reset_clears_all_state():
     manager.reset()
 
     assert len(manager.cached) == 0
+    assert len(manager.num_embeds_by_hash) == 0
     assert len(manager.request_cached_ids) == 0
     assert len(manager.freeable) == 0
     assert len(manager.freed) == 0
diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py
index f0bd296b7183..bc2a761d86ec 100644
--- a/tests/v1/core/test_kv_cache_utils.py
+++ b/tests/v1/core/test_kv_cache_utils.py
@@ -3,6 +3,7 @@
 import copy
 import hashlib
 import importlib
+import mmap
 import subprocess
 import sys
 from collections.abc import Callable
@@ -15,6 +16,7 @@
 import torch
 
 import vllm.v1.core.kv_cache_utils as kv_cache_utils
+import vllm.v1.hisparse.runtime as hisparse_runtime_module
 from vllm.config import (
     CacheConfig,
     KVTransferConfig,
@@ -22,6 +24,7 @@
     SchedulerConfig,
     VllmConfig,
 )
+from vllm.config.attention import HiSparseConfig
 from vllm.config.kv_events import KVEventsConfig
 from vllm.lora.request import LoRARequest
 from vllm.multimodal.inputs import (
@@ -32,7 +35,7 @@
 from vllm.sampling_params import SamplingParams
 from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor
 from vllm.utils.mem_constants import GiB_bytes
-from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
     FreeKVCacheBlockQueue,
@@ -52,10 +55,17 @@
     make_block_hash_with_group_id,
     tensor_data,
 )
+from vllm.v1.hisparse.layout import (
+    create_hisparse_layout,
+    get_hisparse_gpu_memory_usage,
+)
 from vllm.v1.kv_cache_interface import (
     ChunkedLocalAttentionSpec,
+    CircularBufferSpec,
     FullAttentionSpec,
     HiddenStateCacheSpec,
+    HiSparseHotSpec,
+    HiSparseResidentSpec,
     KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -68,18 +78,159 @@
     SinkFullAttentionSpec,
     SlidingWindowMLASpec,
     SlidingWindowSpec,
+    SparseCacheRole,
     UniformTypeKVCacheSpecs,
     get_kv_cache_spec_kind,
     get_kv_cache_spec_sliding_window,
     is_full_attention_spec,
     iter_layer_specs,
 )
+from vllm.v1.kv_cache_layout import KVCacheLayout
 from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
 from vllm.v1.request import Request
 
 pytestmark = pytest.mark.cpu_test
 
 
+@pytest.mark.parametrize("gpu_block_size", [32, 64])
+@pytest.mark.parametrize("shared_host_pool", [False, True])
+def test_hisparse_hma_uses_resolved_gpu_block_size(
+    monkeypatch, gpu_block_size, shared_host_pool
+):
+    monkeypatch.setattr(
+        hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
+    )
+    specs = {
+        "model.layers.0.self_attn": MLAAttentionSpec(
+            block_size=gpu_block_size,
+            num_kv_heads=1,
+            head_size=576,
+            dtype=torch.bfloat16,
+            is_index_group_leader=True,
+        ),
+        "model.layers.0.self_attn.indexer": MLAAttentionSpec(
+            block_size=gpu_block_size,
+            num_kv_heads=1,
+            head_size=128,
+            dtype=torch.bfloat16,
+            cache_role=SparseCacheRole.INDEXER,
+        ),
+    }
+    group_spec = UniformTypeKVCacheSpecs.from_specs(specs)
+    assert group_spec is not None
+    group = KVCacheGroupSpec(list(specs), group_spec)
+    config = SimpleNamespace(
+        attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(index_topk=128),
+            max_model_len=gpu_block_size,
+        ),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=2 if shared_host_pool else 1,
+            pipeline_parallel_size=1,
+            prefill_context_parallel_size=1,
+            decode_context_parallel_size=1,
+            world_size=2 if shared_host_pool else 1,
+            distributed_executor_backend="mp",
+            nnodes_within_dp=1,
+        ),
+        cache_config=SimpleNamespace(
+            num_gpu_blocks_override=7,
+            prefix_cache_retention_interval=None,
+            get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC,
+        ),
+    )
+    indexer_spec = specs["model.layers.0.self_attn.indexer"]
+    assert get_hisparse_gpu_memory_usage(config, [group]) == (
+        indexer_spec.max_memory_usage_bytes(config)
+    )
+
+    monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30)
+    cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
+        config, [group], available_memory=2**30
+    )
+    assert cache_config.num_blocks == 7
+    assert cache_config.hisparse_host_num_blocks is not None
+    assert cache_config.hisparse_host_num_blocks > 7
+
+    host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups
+    assert cache_config.hisparse_shared_host_pool is shared_host_pool
+    host_page = host_group.kv_cache_spec.page_size_bytes
+    alignment = mmap.PAGESIZE if shared_host_pool else 1
+    expected_host_stride = (host_page + alignment - 1) // alignment * alignment
+    assert cache_config.hisparse_host_block_stride == expected_host_stride
+    assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride
+    assert host_group.host_resident
+    assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups])
+    host_layers = set(host_group.layer_names)
+    for tensor in cache_config.kv_cache_tensors:
+        assert all(
+            (name in host_layers) == tensor.host_resident for name in tensor.layers
+        )
+    host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident)
+    host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]]
+    assert host_tensor.block_stride == host_spec.page_size_bytes
+    assert host_tensor.layer_stride == (
+        host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks
+    )
+    assert host_group.kv_cache_spec.block_size == gpu_block_size
+    assert indexer_group.kv_cache_spec.block_size == gpu_block_size
+    host_specs = host_group.kv_cache_spec.kv_cache_specs
+    gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs
+    assert set(host_specs) == {"model.layers.0.self_attn"}
+    assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"}
+    assert indexer_group.kv_cache_spec.prefix_cacheable
+    assert host_group.enable_kv_transfer
+    auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups]
+    assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs)
+    assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs)
+    assert all(
+        spec.block_size == gpu_block_size
+        for spec in auxiliary_specs
+        if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec))
+    )
+    scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes(
+        cache_config,
+        SimpleNamespace(
+            cache_config=SimpleNamespace(
+                block_size=16,
+                enable_prefix_caching=True,
+                prefix_match_unit=None,
+            ),
+            parallel_config=SimpleNamespace(decode_context_parallel_size=1),
+            kv_transfer_config=object(),
+        ),
+    )
+    assert scheduler_block_size == hash_block_size == gpu_block_size
+
+
+def test_hisparse_rejects_deepseek_v4():
+    full_specs = {
+        "model.layers.0.attn": MLAAttentionSpec(
+            block_size=256,
+            num_kv_heads=1,
+            head_size=512,
+            dtype=torch.uint8,
+            model_version="deepseek_v4",
+        )
+    }
+    full_uniform = UniformTypeKVCacheSpecs.from_specs(full_specs)
+    assert full_uniform is not None
+    group = KVCacheGroupSpec(list(full_specs), full_uniform)
+    config = SimpleNamespace(
+        attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()),
+        model_config=SimpleNamespace(hf_config=SimpleNamespace(index_topk=512)),
+        cache_config=SimpleNamespace(num_gpu_blocks_override=7),
+    )
+
+    with pytest.raises(ValueError, match="does not support DeepSeek V4"):
+        create_hisparse_layout(
+            config,
+            [group],
+            host_budget=2**30,
+        )
+
+
 @pytest.fixture(autouse=True)
 def _auto_init_hash_fn(request):
     hash_fn: Callable
@@ -128,6 +279,44 @@ def make_request(
     )
 
 
+@pytest.mark.parametrize("dcp", [1, 4])
+def test_effective_attention_block_size_matches_events(dcp):
+    from vllm.distributed.kv_events import BlockStored
+    from vllm.v1.engine.core import EngineCore
+
+    config = KVCacheConfig(
+        num_blocks=32,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["attention"], new_kv_cache_spec()),
+        ],
+    )
+    manager = KVCacheManager(
+        generate_scheduler_kv_cache_config([config]),
+        max_model_len=256,
+        scheduler_block_size=16 * dcp,
+        hash_block_size=16 * dcp,
+        dcp_world_size=dcp,
+        enable_kv_cache_events=True,
+    )
+    core = EngineCore.__new__(EngineCore)
+    core.vllm_config = SimpleNamespace(cache_config=CacheConfig(block_size=16))
+    core.scheduler = SimpleNamespace(kv_cache_manager=manager)
+    core._initialize_effective_attention_block_size()
+    block_size = core.vllm_config.cache_config.effective_attention_block_size
+    assert block_size == 16 * dcp
+
+    request = make_request(
+        "block-size", list(range(64)), block_size=16 * dcp, hash_fn=sha256
+    )
+    assert manager.allocate_slots(request, 64) is not None
+    assert [
+        event.block_size
+        for event in manager.take_events()
+        if isinstance(event, BlockStored)
+    ] == [block_size]
+
+
 def new_kv_cache_spec(
     block_size=16,
     num_kv_heads=2,
@@ -174,6 +363,45 @@ def test_kv_cache_config_selects_only_transferable_groups():
     )
 
 
+def test_kv_cache_config_selects_prefix_cacheable_groups():
+    """Prefix stores exclude scratch state without changing transfer groups."""
+    full_group = KVCacheGroupSpec(["full"], new_kv_cache_spec())
+    qsa_group = KVCacheGroupSpec(
+        ["qsa"],
+        CircularBufferSpec(
+            block_size=4,
+            num_kv_heads=1,
+            head_size=64,
+            head_size_v=0,
+            dtype=torch.float16,
+        ),
+    )
+    disabled_group = KVCacheGroupSpec(
+        ["disabled"], new_kv_cache_spec(), enable_kv_transfer=False
+    )
+    config = KVCacheConfig(
+        num_blocks=1,
+        kv_cache_tensors=[],
+        kv_cache_groups=[full_group, qsa_group, disabled_group],
+    )
+    assert config.transfer_group_ids == (0, 1)
+    assert config.select_transfer_block_ids(([1], [2], [3])) == ([1], [2])
+    assert config.prefix_cacheable_group_ids == (0,)
+    assert config.prefix_cacheable_groups == (full_group,)
+
+
+def test_kv_cache_blocks_selects_requested_groups():
+    blocks = KVCacheBlocks(
+        (
+            [KVCacheBlock(1)],
+            [KVCacheBlock(2)],
+            [KVCacheBlock(3)],
+        )
+    )
+
+    assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3])
+
+
 def new_sliding_window_spec(
     block_size=16,
     num_kv_heads=2,
@@ -840,11 +1068,25 @@ def _stats(requests: int, queries: int, hits: int) -> PrefixCacheStats:
     return PrefixCacheStats(requests=requests, queries=queries, hits=hits)
 
 
-def test_metrics():
-    """
-    Test the prefix caching metrics.
+def test_metrics_empty_distinguishes_no_queries_from_no_hits():
+    """`hit_rate` alone cannot tell the two apart; `empty` can.
+
+    Both an unobserved window and a genuine all-miss window report a hit
+    rate of 0.0, so anything surfacing that number to a human has to check
+    `empty` first - which is what the prefix-cache log line does.
     """
     metrics = CachingMetrics(max_recent_requests=5)
+    assert metrics.empty
+    assert metrics.hit_rate == 0.0
+
+    metrics.observe(_stats(1, 20, 0))
+    assert not metrics.empty
+    assert metrics.hit_rate == 0.0
+
+
+def test_metrics():
+    """Test the prefix caching metrics."""
+    metrics = CachingMetrics(max_recent_requests=5)
     assert metrics.hit_rate == 0.0
 
     metrics.observe(_stats(1, 20, 9))
@@ -873,9 +1115,7 @@ def test_metrics():
 
 
 def test_metrics_empty_stats():
-    """
-    Test the prefix caching metrics with empty stats.
-    """
+    """Test the prefix caching metrics with empty stats."""
     metrics = CachingMetrics(max_recent_requests=5)
     metrics.observe(_stats(0, 0, 0))
     metrics.observe(_stats(1, 20, 9))
@@ -1675,7 +1915,7 @@ def test_get_max_concurrency_for_kv_cache_config():
 
 
 def test_allocate_with_lookahead():
-    """Verify that lookahead tokens correctly affect block allocation"""
+    """Verify that lookahead tokens correctly affect block allocation."""
     block_size = 4
     config = KVCacheConfig(
         num_blocks=10,
@@ -3018,6 +3258,45 @@ def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm):
     assert "attention compute is unchanged" in caplog_vllm.text
 
 
+def test_hidden_states_with_tp_scales_page_size():
+    """When TP shrinks KV pages below the hidden-state per-token cost,
+    get_kv_cache_groups must scale up target block sizes so that the
+    common page accommodates the unsharded hidden states."""
+    # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank.
+    # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2)
+    #         = 16 * 2 * 64 * 2 = 4096 bytes.
+    kv_spec = new_kv_cache_spec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=64,
+        dtype=torch.bfloat16,
+    )
+    # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512)
+    #   * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page.
+    hs_spec = HiddenStateCacheSpec(
+        block_size=16,
+        num_kv_heads=6,
+        head_size=512,
+        dtype=torch.bfloat16,
+    )
+    specs = {
+        "target.0.attn": kv_spec,
+        "target.1.attn": kv_spec,
+        "cache_only_layers.48": hs_spec,
+    }
+
+    groups = get_kv_cache_groups(_grouping_config(), specs)
+
+    # The hidden-state layer should be present and no assertion should fire.
+    all_layers = {name for g in groups for name in g.layer_names}
+    assert "cache_only_layers.48" in all_layers
+
+    # The target group block sizes must have been scaled up.
+    for g in groups:
+        if "cache_only_layers.48" not in g.layer_names:
+            assert g.kv_cache_spec.block_size > kv_spec.block_size
+
+
 def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():
     assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION
 
@@ -3655,8 +3934,7 @@ def test_unify_kv_cache_spec_page_size_mamba():
 
 
 def test_hma_not_disabled_when_kv_events_enabled():
-    """
-    Test enabling KV events must not force disable_hybrid_kv_cache_manager to True.
+    """Test enabling KV events must not force disable_hybrid_kv_cache_manager to True.
 
     This test guards against that regression by verifying that a VllmConfig
     with kv_events_config set still resolves disable_hybrid_kv_cache_manager
@@ -3848,6 +4126,22 @@ def test_iter_layer_specs_returns_group_members():
     assert list(iter_layer_specs(wrapped)) == [full, mla]
 
 
+def test_wrapped_mamba_group_requires_block_zeroing():
+    mamba = MambaSpec(
+        block_size=4,
+        shapes=((4, 1),),
+        dtypes=(torch.float32,),
+    )
+    wrapped = UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs={"mamba": mamba})
+    config = KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[],
+        kv_cache_groups=[KVCacheGroupSpec(["mamba"], wrapped)],
+    )
+
+    assert config.needs_kv_cache_zeroing
+
+
 def _spec_decode_grouping_config(method="dspark", model_type=None):
     """Grouping config with an EAGLE-family speculative method enabled."""
     return SimpleNamespace(
@@ -3929,9 +4223,9 @@ def test_draft_group_not_annotated_without_spec_decode():
 
 
 def test_unidentifiable_draft_with_mamba_warns(caplog_vllm):
-    # No group carries the draft marker, so every consumer falls back to
-    # flagging all groups -- including Mamba ones, which then can never report
-    # a hit. That is silent today; it must at least be visible.
+    # No group carries the draft marker, so consumers fall back to
+    # conservative behavior that silently breaks reuse for Mamba groups.
+    # That must at least be visible.
     groups = get_kv_cache_groups(
         _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False)
     )
@@ -3940,7 +4234,20 @@ def test_unidentifiable_draft_with_mamba_warns(caplog_vllm):
     assert "no KV cache group could be identified as the draft model's" in (
         caplog_vllm.text
     )
-    assert "Mamba groups" in caplog_vllm.text
+
+
+def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm):
+    # Pure-attention models degrade gracefully under the consumers'
+    # conservative fallback (a one-block hit drop at most), so the warning
+    # stays silent to avoid noise on every unannotated EAGLE deployment.
+    specs = {
+        "target.attn.0": new_mla_spec(block_size=64),
+        "target.attn.1": new_mla_spec(block_size=64),
+    }
+    groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs)
+
+    assert not any(g.is_eagle_group for g in groups)
+    assert "could be identified as the draft model's" not in caplog_vllm.text
 
 
 def test_no_warning_when_draft_group_is_identified(caplog_vllm):
@@ -3968,12 +4275,16 @@ def _deepseek_v4_specs(model_version="deepseek_v4"):
     }
 
 
-def test_deepseek_v4_draft_group_annotated_on_packed_path():
+@pytest.mark.parametrize(
+    ("method", "model_type"),
+    [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")],
+)
+def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type):
     # DeepseekV4's MTP block reuses the target's decoder layer, so its spec
     # carries no draft marker and only the positional rule can find it. This
     # pins the pre-existing behaviour that the unified annotator must preserve.
     groups = get_kv_cache_groups(
-        _spec_decode_grouping_config(method="mtp", model_type="deepseek_v4"),
+        _spec_decode_grouping_config(method=method, model_type=model_type),
         _deepseek_v4_specs(model_version=None),
     )
 
@@ -3982,13 +4293,93 @@ def test_deepseek_v4_draft_group_annotated_on_packed_path():
     assert "model.layers.3.self_attn.attn" in flagged[0].layer_names
 
 
-def test_deepseek_v4_annotation_requires_model_type():
-    # The positional rule is only sound for DeepseekV4, where the draft layer
-    # is known to be registered last. Without that model gate nothing may be
-    # flagged, however the grouping happens to fall out.
+def test_trailing_layer_fallback_applies_to_any_mtp_model():
+    # The positional rule is sound for every MTP drafter, not just DeepseekV4:
+    # MTP blocks reuse the target's decoder layer (no spec marker) and always
+    # register after every target layer. The model_type must not gate it.
     groups = get_kv_cache_groups(
         _spec_decode_grouping_config(method="mtp", model_type="other"),
         _deepseek_v4_specs(),
     )
 
+    flagged = [g for g in groups if g.is_eagle_group]
+    assert len(flagged) == 1
+    assert "model.layers.3.self_attn.attn" in flagged[0].layer_names
+
+
+def _qwen3_5_hybrid_specs(with_mtp_layer: bool):
+    """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with
+    the MTP drafter's full-attn layer (spec-identical to the target's)
+    registered last."""
+    specs = {}
+    idx = 0
+    for _ in range(2):
+        for _ in range(3):
+            specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec(
+                mamba_cache_mode="align"
+            )
+            idx += 1
+        specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec()
+        idx += 1
+    if with_mtp_layer:
+        specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec()
+    return specs
+
+
+def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm):
+    # A hybrid mamba + full-attention model with an MTP drafter that is
+    # spec-indistinguishable from the target reaches the general multi-group
+    # path. The trailing-layer rule must locate the draft group there so the
+    # Mamba groups are not swept up by the flag-all consumer fallback.
+    groups = get_kv_cache_groups(
+        _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"),
+        _qwen3_5_hybrid_specs(with_mtp_layer=True),
+    )
+
+    flagged = [g for g in groups if g.is_eagle_group]
+    assert len(flagged) == 1
+    assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names
+    for group in groups:
+        if any(
+            isinstance(spec, MambaSpec)
+            for spec in iter_layer_specs(group.kv_cache_spec)
+        ):
+            assert not group.is_eagle_group
+    assert "could be identified as the draft model's" not in caplog_vllm.text
+
+
+@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"])
+def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method):
+    # Other drafters are not covered by the trailing-layer rule, so an
+    # unidentifiable hybrid draft must still warn.
+    groups = get_kv_cache_groups(
+        _spec_decode_grouping_config(method=method, model_type="qwen3_5"),
+        _qwen3_5_hybrid_specs(with_mtp_layer=True),
+    )
+
     assert not any(g.is_eagle_group for g in groups)
+    assert "could be identified as the draft model's" in caplog_vllm.text
+
+
+def test_trailing_layer_fallback_requires_exact_partition():
+    # If the groups do not partition the layers exactly (e.g. a caller that
+    # dropped or duplicated layers), the positional rule is meaningless and
+    # must not fire.
+    from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups
+
+    specs = _qwen3_5_hybrid_specs(with_mtp_layer=True)
+    config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5")
+    groups = get_kv_cache_groups(config, specs)
+    for g in groups:
+        g.is_eagle_group = False
+    # Remove one layer from its group: no longer an exact partition.
+    trimmed = [
+        KVCacheGroupSpec(
+            [n for n in g.layer_names if n != "model.layers.0.linear_attn"],
+            g.kv_cache_spec,
+        )
+        for g in groups
+    ]
+    _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True)
+
+    assert not any(g.is_eagle_group for g in trimmed)
diff --git a/tests/v1/core/test_kv_sharing.py b/tests/v1/core/test_kv_sharing.py
index db0e8dae3e78..4a030c674bbf 100644
--- a/tests/v1/core/test_kv_sharing.py
+++ b/tests/v1/core/test_kv_sharing.py
@@ -16,42 +16,8 @@ def new_kv_cache_spec():
     )
 
 
-def test_initialize_kv_cache_for_kv_sharing_different_attn_groups():
-    """
-    Test initializing KV cache sharing with different attention groups.
-    Layers in the same KV cache group might be placed in different attn groups
-    if they have different attention backends.
-    """
-    shared_kv_cache_layers = {
-        "model.layers.2": "model.layers.0",
-        "model.layers.3": "model.layers.1",
-    }
-
-    # Layers 0 and 1 both belong in KV cache group 0
-    # However, if they have different attention backends, they will be
-    # placed in different attention groups for KV cache group 0
-    kv_cache_groups = [
-        KVCacheGroupSpec(["model.layers.0", "model.layers.1"], new_kv_cache_spec()),
-    ]
-
-    add_kv_sharing_layers_to_kv_cache_groups(
-        shared_kv_cache_layers=shared_kv_cache_layers,
-        kv_cache_groups=kv_cache_groups,
-    )
-
-    # Check that the layers were added to the correct KV cache group
-    assert len(kv_cache_groups) == 1
-    assert kv_cache_groups[0].layer_names == [
-        "model.layers.0",
-        "model.layers.1",
-        "model.layers.2",
-        "model.layers.3",
-    ]
-
-
 def test_initialize_kv_cache_for_kv_sharing_same_attn_groups():
-    """
-    Test case assuming that all layers in the same KV cache group have the same
+    """Test case assuming that all layers in the same KV cache group have the same
     attention backends. This is true for most models.
     """
     shared_kv_cache_layers = {
@@ -79,8 +45,7 @@ def test_initialize_kv_cache_for_kv_sharing_same_attn_groups():
 
 
 def test_initialize_kv_cache_for_kv_sharing_no_attn_groups():
-    """
-    Test KV sharing set up when no attention groups are provided.
+    """Test KV sharing set up when no attention groups are provided.
     This is the case for the TPU model runner, which doesn't have
     support for attention groups yet.
     """
diff --git a/tests/v1/core/test_mamba_align_chunk_split.py b/tests/v1/core/test_mamba_align_chunk_split.py
index 1bd7f13b5935..b1575239158d 100644
--- a/tests/v1/core/test_mamba_align_chunk_split.py
+++ b/tests/v1/core/test_mamba_align_chunk_split.py
@@ -144,7 +144,7 @@ def _split(
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         # `prefix_match_unit` finer than the block size (#46384).
         mamba_partial_cache_hit=partial_hit,
-        mamba_fine_grained_prefix_cache=False,
+        mamba_shared_prefix_checkpoint=False,
         hash_block_size=ATTN_BLOCK_SIZE,
         mamba_has_prefill_checkpoint_blocks=(num_prefill_checkpoint_blocks > 0),
         mamba_prefill_checkpoint_alignment=(
diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py
index 1760ca14d0d6..0a34a6131dc3 100644
--- a/tests/v1/core/test_prefix_caching.py
+++ b/tests/v1/core/test_prefix_caching.py
@@ -6,7 +6,8 @@
 from collections.abc import Callable
 from dataclasses import replace
 from math import lcm
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -14,11 +15,21 @@
 import vllm.v1.core.kv_cache_manager as kv_cache_manager
 import vllm.v1.core.kv_cache_utils as kv_cache_utils
 from vllm.distributed.kv_events import (
+    MEDIUM_CPU,
     MEDIUM_GPU,
     AllBlocksCleared,
     BlockRemoved,
     BlockStored,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
+    HiSparseConnector,
+    HiSparseConnectorScheduler,
+    HiSparseConnectorWorkerMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
+    OffloadingConnector,
+)
 from vllm.lora.request import LoRARequest
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
@@ -40,16 +51,23 @@
     init_none_hash,
     make_block_hash_with_group_id,
 )
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
+    HiSparseHotSpec,
+    HiSparseResidentSpec,
     KVCacheConfig,
+    KVCacheGroupRole,
     KVCacheGroupSpec,
     KVCacheSpecKind,
     MambaSpec,
     MLAAttentionSpec,
     SlidingWindowSpec,
 )
+from vllm.v1.outputs import KVConnectorOutput
+from vllm.v1.request import RequestStatus
 
 pytestmark = pytest.mark.cpu_test
 
@@ -118,7 +136,10 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa
             kv_cache_config,
             prefix_cache_retention_interval=kwargs.pop("retention_interval"),
         )
-    return KVCacheManager(kv_cache_config, **kwargs)
+    manager = KVCacheManager(kv_cache_config, **kwargs)
+    if kv_cache_config.hisparse_host_num_blocks is not None:
+        get_hisparse_coordinator(manager)
+    return manager
 
 
 def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
@@ -139,6 +160,1087 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
     )
 
 
+HISPARSE_BLOCK_SIZE = 16
+
+
+def make_hisparse_kv_cache_config(
+    num_blocks: int,
+    host_num_blocks: int,
+    *,
+    transfer_device_cache: bool = False,
+) -> KVCacheConfig:
+    source_spec = FullAttentionSpec(
+        block_size=HISPARSE_BLOCK_SIZE,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+    )
+    groups = [
+        KVCacheGroupSpec(
+            ["source"],
+            source_spec,
+            role=KVCacheGroupRole.HISPARSE_SOURCE,
+            host_resident=True,
+        ),
+        KVCacheGroupSpec(
+            ["indexer"],
+            source_spec,
+            enable_kv_transfer=transfer_device_cache,
+            role=KVCacheGroupRole.HISPARSE_INDEXER,
+        ),
+    ]
+    groups.append(
+        KVCacheGroupSpec(
+            ["resident"],
+            HiSparseResidentSpec(
+                block_size=HISPARSE_BLOCK_SIZE,
+                page_size=HISPARSE_BLOCK_SIZE * 4,
+            ),
+            enable_kv_transfer=transfer_device_cache,
+        )
+    )
+    groups.append(
+        KVCacheGroupSpec(
+            ["hot"],
+            HiSparseHotSpec(
+                block_size=HISPARSE_BLOCK_SIZE,
+                page_size=HISPARSE_BLOCK_SIZE * 4,
+                blocks_per_request=2,
+            ),
+            enable_kv_transfer=False,
+        )
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        hisparse_host_num_blocks=host_num_blocks,
+        kv_cache_tensors=[],
+        kv_cache_groups=groups,
+    )
+
+
+def make_hisparse_kv_cache_manager(
+    num_blocks: int,
+    host_num_blocks: int,
+    *,
+    max_model_len: int = 128,
+    max_in_flight_tokens: int | None = None,
+    enable_caching: bool = False,
+    **config_kwargs,
+) -> KVCacheManager:
+    config = make_hisparse_kv_cache_config(num_blocks, host_num_blocks, **config_kwargs)
+    return make_kv_cache_manager(
+        config,
+        max_model_len=max_model_len,
+        max_in_flight_tokens=max_in_flight_tokens,
+        enable_caching=enable_caching,
+        hash_block_size=HISPARSE_BLOCK_SIZE,
+    )
+
+
+def test_hisparse_builds_dma_row_mirrors_across_pages():
+    manager = make_hisparse_kv_cache_manager(32, 16)
+    request = make_request(
+        "request",
+        list(range(2 * HISPARSE_BLOCK_SIZE)),
+        HISPARSE_BLOCK_SIZE,
+        sha256,
+    )
+    assert manager.allocate_slots(request, num_new_tokens=32) is not None
+    coordinator = get_hisparse_coordinator(manager)
+    resident_blocks = coordinator.resident_managers[0].req_to_blocks[request.request_id]
+    host_blocks = coordinator.host_manager.req_to_blocks[request.request_id]
+
+    mirrors = coordinator.build_row_mirrors([(request.request_id, 14, 4)])
+
+    assert [mirror.source_starts for mirror in mirrors] == [
+        (resident_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14,),
+        (resident_blocks[1].block_id * HISPARSE_BLOCK_SIZE,),
+    ]
+    assert [mirror.destination_start for mirror in mirrors] == [
+        host_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14,
+        host_blocks[1].block_id * HISPARSE_BLOCK_SIZE,
+    ]
+    assert [mirror.num_rows for mirror in mirrors] == [2, 2]
+
+
+def test_hisparse_async_speculation_mirrors_uncertain_position_range():
+    """Unresolved drafts must not leave gaps in the eager host mirror."""
+    coordinator = MagicMock(host_group_id=0)
+    coordinator.take_block_table_updates.return_value = {}
+    coordinator.build_offload_command.return_value = None
+    coordinator.build_row_mirrors.return_value = ()
+    scheduler = HiSparseConnectorScheduler(
+        async_speculative=True,
+        draft_kv_lookahead=4,
+    )
+    scheduler.bind_coordinator(coordinator)
+    request = SimpleNamespace(request_id="request", num_output_placeholders=3)
+    connector = object.__new__(HiSparseConnector)
+    connector.connector_scheduler = scheduler
+    connector.update_state_after_alloc(request, None, 0)
+    scheduler_output = SimpleNamespace(
+        block_table_updates=None,
+        kv_cache_block_copies=None,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=SimpleNamespace(
+            new_block_ids=[],
+            req_ids=["request"],
+            num_computed_tokens=[103],
+        ),
+        num_scheduled_tokens={"request": 4},
+    )
+
+    scheduler.build_connector_meta(scheduler_output)
+
+    coordinator.build_row_mirrors.assert_called_once_with((("request", 100, 11),))
+
+    # Acceptance arrives between steps; the connector must read the live count.
+    request.num_output_placeholders = 1
+    scheduler.build_connector_meta(scheduler_output)
+    coordinator.build_row_mirrors.assert_called_with((("request", 102, 9),))
+
+    connector.request_finished_all_groups(request, ())
+    assert not scheduler.requests
+
+
+def test_hisparse_reports_when_context_is_fully_resident():
+    """A page the pool reuses under a host-reading request stops being resident."""
+    manager = make_hisparse_kv_cache_manager(32, 16)
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
+    coordinator = get_hisparse_coordinator(manager)
+    scheduled = ((request.request_id, len(tokens) - 1, 1),)
+    assert coordinator.all_context_pages_resident(scheduled)
+
+    for hot_manager in coordinator.hot_managers:
+        hot_manager.require_hot(request.request_id)
+    assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
+    _publish_hisparse_pages(manager)
+    pool = manager.block_pool
+    pool.get_new_blocks(pool.get_num_free_blocks())
+
+    assert not coordinator.all_context_pages_resident(scheduled)
+    assert coordinator.take_block_table_updates().keys() == {request.request_id}
+
+
+def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():
+    """Restore the final host page when completing an indexer-only import."""
+    manager = make_hisparse_kv_cache_manager(
+        32,
+        16,
+        enable_caching=True,
+    )
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
+    spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
+    spill_counts = {spill.transfer_id: 1 for spill in spills}
+    get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts)
+    _, indexer_blocks, _, _ = manager.get_blocks(original.request_id).blocks
+    evicted_indexer_id = indexer_blocks[2].block_id
+    manager.free(original)
+    manager.block_pool.evict_blocks({evicted_indexer_id})
+
+    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    blocks, num_local, _ = manager.get_computed_blocks(resumed)
+    max_completion = HISPARSE_BLOCK_SIZE
+    assert num_local == 2 * HISPARSE_BLOCK_SIZE
+    assert [len(group_blocks) for group_blocks in blocks.blocks] == [3, 2, 0, 0]
+
+    allocated = manager.allocate_slots(
+        resumed,
+        num_new_tokens=1,
+        num_new_computed_tokens=num_local,
+        new_computed_blocks=blocks,
+        num_external_computed_tokens=max_completion,
+    )
+
+    assert allocated is not None
+    source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks
+    assert len(source) == len(indexer) == len(resident) == 4
+    assert len(hot) == 2
+    assert not any(block.is_null for block in resident[:2])
+    assert not resident[2].is_null
+    assert not resident[3].is_null
+    coordinator = get_hisparse_coordinator(manager)
+    assert not coordinator.build_offload_command().page_transfers
+    coordinator.finish_host_import(resumed.request_id, failed=False)
+    transfers = coordinator.build_offload_command().page_transfers
+    assert len(transfers) == 1 and transfers[0].restore
+    assert transfers[0].resident_block_ids == (resident[2].block_id,)
+
+
+def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix():
+    manager = make_hisparse_kv_cache_manager(
+        32,
+        16,
+        enable_caching=True,
+    )
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
+    spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
+    spill_counts = {spill.transfer_id: 1 for spill in spills}
+    get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts)
+    host_blocks, _, _, _ = manager.get_blocks(original.request_id).blocks
+    evicted_host_id = host_blocks[0].block_id
+    manager.free(original)
+    get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
+        {evicted_host_id}
+    )
+
+    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    _, num_local, _ = manager.get_computed_blocks(resumed)
+    connector = object.__new__(OffloadingConnector)
+    connector._kv_cache_config = manager.kv_cache_config
+    connector.bind_kv_cache_manager(manager)
+    assert num_local == 0
+    assert connector._max_loadable_tokens(resumed, num_local) == 0
+
+
+@pytest.mark.parametrize("nested", [False, True])
+@pytest.mark.parametrize(
+    "available_tokens,imported_tokens,host_blocks",
+    [(0, 0, 2), (16, 16, 2), (64, 32, 3)],
+)
+def test_connector_completes_partial_prefix_without_importing_missing_host_kv(
+    nested, available_tokens, imported_tokens, host_blocks
+):
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    tokens = list(range(64))
+    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(original, num_new_tokens=64) is not None
+    _publish_hisparse_pages(manager)
+    host, indexer, _, _ = manager.get_blocks(original.request_id).blocks
+    missing_indexer = indexer[1].block_id
+    missing_host = host[0].block_id
+    manager.free(original)
+    manager.block_pool.evict_blocks({missing_indexer})
+
+    defer = True
+
+    def lookup(request, num_computed_tokens, max_num_new_tokens=None):
+        nonlocal defer
+        if defer:
+            defer = False
+            return None, False
+        available = available_tokens
+        if max_num_new_tokens is not None:
+            available = min(available, max_num_new_tokens)
+        return available, available > 0
+
+    connector = object.__new__(OffloadingConnector)
+    connector._kv_cache_config = manager.kv_cache_config
+    connector.connector_scheduler = SimpleNamespace(get_num_new_matched_tokens=lookup)
+    if nested:
+        for _ in range(2):
+            outer = object.__new__(MultiConnector)
+            outer._connectors = (connector,)
+            outer._requests_to_connector = {}
+            connector = outer
+    connector.bind_kv_cache_manager(manager)
+    scheduler = SimpleNamespace(connector=connector, kv_cache_manager=manager)
+    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
+
+    blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit(
+        scheduler, resumed
+    )
+    assert local == 16 and not diverged
+    assert connector.get_num_new_matched_tokens(resumed, local) == (None, False)
+
+    blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit(
+        scheduler, resumed
+    )
+    external, asynchronous = connector.get_num_new_matched_tokens(resumed, local)
+    assert external == imported_tokens
+    assert asynchronous == (external > 0)
+    assert [len(group) for group in blocks.blocks] == [3, 1, 0, 0]
+    assert (
+        manager.allocate_slots(
+            resumed,
+            int(external == 0),
+            delay_cache_blocks=True,
+            num_new_computed_tokens=local,
+            new_computed_blocks=blocks,
+            num_external_computed_tokens=external,
+        )
+        is not None
+    )
+    allocated = manager.get_blocks(resumed.request_id).blocks
+    assert len(allocated[0]) == host_blocks
+    assert allocated[2][1].is_null == (host_blocks > 2)
+    assert not allocated[2][-1].is_null
+    manager.free(resumed)
+
+    # A surviving indexer offload must not make absent host KV look computed.
+    get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
+        {missing_host}
+    )
+    retry = make_request("retry", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    _, local, _, _ = Scheduler._get_local_prefix_cache_hit(scheduler, retry)
+    assert local == 0
+    assert connector.get_num_new_matched_tokens(retry, local) == (0, False)
+
+
+def allocate_external_prefix(
+    manager: KVCacheManager, request: Request, num_tokens: int
+) -> KVCacheBlocks | None:
+    return manager.allocate_slots(
+        request,
+        num_new_tokens=0,
+        num_external_computed_tokens=num_tokens,
+        delay_cache_blocks=True,
+        full_sequence_must_fit=True,
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [1, 15, 16, 17, 31, 32, 33, 127])
+def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens):
+    """Import every real token and restore the page used by last-token replay."""
+    from tests.v1.kv_connector.unit.utils import create_vllm_config
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
+        NixlPullConnectorScheduler,
+    )
+
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    config = create_vllm_config(block_size=HISPARSE_BLOCK_SIZE)
+    connector = NixlPullConnectorScheduler(config, "test", manager.kv_cache_config)
+    request = make_request(
+        "import", list(range(num_tokens)), HISPARSE_BLOCK_SIZE, sha256
+    )
+    request.kv_transfer_params = {"do_remote_prefill": True}
+    count, is_async = connector.get_num_new_matched_tokens(request, 0)
+    assert count == num_tokens and is_async
+    assert allocate_external_prefix(manager, request, count) is not None
+    source, _, resident, _ = manager.get_blocks(request.request_id).blocks
+    assert len(source) == len(resident) == (num_tokens + 15) // 16
+    assert all(block.is_null for block in resident[:-1])
+    assert not resident[-1].is_null
+    coordinator = get_hisparse_coordinator(manager)
+    assert not coordinator.build_offload_command().page_transfers
+
+    coordinator.finish_host_import(request.request_id, failed=False)
+    transfers = coordinator.build_offload_command().page_transfers
+    assert len(transfers) == 1
+    restore = transfers[0]
+    assert restore.restore and not restore.after_forward
+    assert restore.host_block_id == source[-1].block_id
+    assert restore.resident_block_ids == (resident[-1].block_id,)
+    request.num_computed_tokens = count
+    scheduler = SimpleNamespace(
+        connector=connector,
+        kv_cache_manager=manager,
+        failed_recving_kv_req_ids=set(),
+        finished_recving_kv_req_ids={request.request_id},
+        prefix_replay_tokens=0,
+    )
+    scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler)
+    Scheduler._update_waiting_for_remote_kv(scheduler, request)
+    assert request.num_tokens - request.num_computed_tokens == 1
+    assert manager.allocate_slots(request, num_new_tokens=1) is not None
+    assert coordinator.build_row_mirrors([(request.request_id, num_tokens - 1, 1)])
+    # A pending restore must not expose uninitialized GPU copies to prefix hits.
+    assert all(resident[-1] not in copies for copies in coordinator.copies.values())
+
+    coordinator.update_spills({restore.transfer_id: 2}, {restore.transfer_id: 1})
+    assert resident[-1].ref_cnt == 2
+    assert all(resident[-1] not in copies for copies in coordinator.copies.values())
+    coordinator.update_spills({}, {restore.transfer_id: 1})
+    assert resident[-1].ref_cnt == 1
+    assert coordinator.request_states[request.request_id].valid_pages == set(
+        range(num_tokens // HISPARSE_BLOCK_SIZE)
+    )
+    if num_tokens % HISPARSE_BLOCK_SIZE:
+        assert source[-1].block_hash is None
+    else:
+        assert coordinator.copies[source[-1].block_hash] == (resident[-1],)
+
+
+def test_hisparse_aborted_tail_restore_retains_both_endpoints():
+    """An abort cannot recycle storage still used by a queued tail restore."""
+    manager = make_hisparse_kv_cache_manager(16, 16)
+    request = make_request("import", list(range(17)), HISPARSE_BLOCK_SIZE, sha256)
+    assert allocate_external_prefix(manager, request, 17) is not None
+    source, _, resident, _ = manager.get_blocks(request.request_id).blocks
+    host_tail, gpu_tail = source[-1], resident[-1]
+    coordinator = get_hisparse_coordinator(manager)
+    coordinator.finish_host_import(request.request_id, failed=False)
+    restore = coordinator.build_offload_command().page_transfers[0]
+    manager.free(request)
+    assert host_tail.ref_cnt == gpu_tail.ref_cnt == 1
+    coordinator.update_spills({restore.transfer_id: 2}, {restore.transfer_id: 1})
+    assert host_tail.ref_cnt == gpu_tail.ref_cnt == 1
+    coordinator.update_spills({}, {restore.transfer_id: 1})
+    assert host_tail.ref_cnt == gpu_tail.ref_cnt == 0
+    assert not coordinator.has_pending_work()
+
+
+@pytest.mark.parametrize("enable_caching", [False, True])
+def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_caching):
+    """Later admissions must not remove a scheduled request's only readable tier."""
+    manager = make_hisparse_kv_cache_manager(
+        32, 32, max_model_len=256, enable_caching=enable_caching
+    )
+    first = make_request("first", list(range(160)), HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(first, num_new_tokens=128) is not None
+    _publish_hisparse_pages(manager)
+    first.num_computed_tokens = 128
+    pool = manager.block_pool
+    held = pool.get_new_blocks(pool.get_num_free_blocks() - 3)
+
+    assert manager.allocate_slots(first, num_new_tokens=1) is not None
+    second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256)
+    manager.allocate_slots(second, num_new_tokens=16)
+    coordinator = get_hisparse_coordinator(manager)
+    scheduled = (("first", 128, 1),)
+    assert coordinator.all_context_pages_resident(scheduled)
+    assert all(not m.has_hot("first") for m in coordinator.hot_managers)
+
+    # Once another allocation finishes, A can acquire its hot buffer and
+    # safely release its clean pages for other requests to reuse.
+    pool.free_blocks(held)
+    assert manager.allocate_slots(first, num_new_tokens=1) is not None
+    assert all(m.has_hot("first") for m in coordinator.hot_managers)
+    pool.get_new_blocks(pool.get_num_free_blocks())
+    assert not coordinator.all_context_pages_resident(scheduled)
+    assert "first" in coordinator.take_block_table_updates()
+
+
+def test_hisparse_full_pool_keeps_pages_pinned_until_preemption():
+    """Without room for a hot buffer, admission must defer instead of losing KV."""
+    manager = make_hisparse_kv_cache_manager(18, 18, max_model_len=160)
+    first = make_request("first", list(range(128)), HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(first, num_new_tokens=128) is not None
+    _publish_hisparse_pages(manager)
+    assert manager.block_pool.get_num_free_blocks() == 1
+    first.num_computed_tokens = 128
+    assert manager.allocate_slots(first, num_new_tokens=1) is None
+    assert get_hisparse_coordinator(manager).all_context_pages_resident(
+        (("first", 127, 1),)
+    )
+
+    manager.free(first)
+    second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(second, num_new_tokens=16) is not None
+
+
+def test_hisparse_materializes_prefix_without_allocating_hot_blocks():
+    """A host prefix becomes visible only when every page is durable.
+
+    Completing a later page first must not expose a prefix with a hole.
+    """
+    manager = make_hisparse_kv_cache_manager(
+        32,
+        16,
+        enable_caching=True,
+    )
+    tokens = list(range(2 * HISPARSE_BLOCK_SIZE))
+    request = make_request("resident", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
+    blocks = manager.get_block_ids(request.request_id)
+    assert len(blocks[2]) == 2
+    assert blocks[3] == []
+
+    spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
+    assert len(spills) == 2
+    assert all(spill.after_forward for spill in spills)
+    resident_blocks = manager.get_blocks(request.request_id).blocks[2]
+    assert [block.ref_cnt for block in resident_blocks] == [2, 2]
+    spill_counts = {spill.transfer_id: 1 for spill in spills}
+    duplicate = make_request(
+        "duplicate",
+        list(range(3 * HISPARSE_BLOCK_SIZE)),
+        HISPARSE_BLOCK_SIZE,
+        sha256,
+    )
+    _, num_computed, _ = manager.get_computed_blocks(duplicate)
+    assert num_computed == 0
+
+    get_hisparse_coordinator(manager).update_spills(spill_counts, {})
+    assert [block.ref_cnt for block in resident_blocks] == [2, 2]
+    _, num_computed, _ = manager.get_computed_blocks(duplicate)
+    assert num_computed == 0
+
+    get_hisparse_coordinator(manager).update_spills({}, {spills[1].transfer_id: 1})
+    assert [block.ref_cnt for block in resident_blocks] == [2, 1]
+    _, num_computed, _ = manager.get_computed_blocks(duplicate)
+    assert num_computed == 0
+
+    get_hisparse_coordinator(manager).update_spills({}, {spills[0].transfer_id: 1})
+    _, num_computed, _ = manager.get_computed_blocks(duplicate)
+    assert num_computed == 2 * HISPARSE_BLOCK_SIZE
+    blocks = manager.get_block_ids(request.request_id)
+    assert len(blocks[2]) == 2
+    assert blocks[3] == []
+
+
+def test_hisparse_materialization_respects_per_step_spill_budget():
+    """Prefix publication must not bypass the configured spill batch limit."""
+    manager = make_hisparse_kv_cache_manager(
+        32,
+        16,
+        enable_caching=True,
+    )
+    coordinator = get_hisparse_coordinator(manager)
+    coordinator.max_spill_pages = 1
+    tokens = list(range(2 * HISPARSE_BLOCK_SIZE))
+    request = make_request("bounded", tokens, HISPARSE_BLOCK_SIZE, sha256)
+
+    assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
+    first = coordinator.build_offload_command().page_transfers
+    assert len(first) == 1
+
+    coordinator.plan_prefix_materialization(request.request_id, len(tokens))
+    second = coordinator.build_offload_command().page_transfers
+    assert len(second) == 1
+    assert first[0].transfer_id != second[0].transfer_id
+
+
+def test_hisparse_deferred_free_retains_host_blocks_until_fence():
+    """Detaching a request must not make in-flight host pages reusable."""
+    manager = make_hisparse_kv_cache_manager(16, 16)
+    host_pool = get_hisparse_coordinator(manager).get_host_block_pool()
+    assert host_pool is not None
+    host_free = host_pool.get_num_free_blocks()
+    device_free = manager.block_pool.get_num_free_blocks()
+    request = make_request("free", list(range(16)), HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, 16, delay_cache_blocks=True) is not None
+    blocks = manager.pop_blocks_for_free(request)
+    assert host_pool.get_num_free_blocks() < host_free
+    assert blocks and any(host_pool.blocks[b.block_id] is b for b in blocks)
+    manager.block_pool.free_blocks(reversed(blocks))
+    assert host_pool.get_num_free_blocks() == host_free
+    assert manager.block_pool.get_num_free_blocks() == device_free
+
+
+def test_hisparse_host_cow_copy_is_drained_without_a_gpu_pool():
+    """Host-only copy-on-write work must reach the worker copy queue."""
+    manager = make_hisparse_kv_cache_manager(16, 16)
+    source_manager = manager.coordinator.single_type_managers[0]
+    source_block = source_manager.block_pool.get_new_blocks(1)[0]
+    source_manager.req_to_blocks["cow"] = [source_block]
+    source_manager._partial_hit_reqs["cow"] = (0, source_block)
+
+    new_blocks = source_manager.allocate_new_blocks(
+        "cow", HISPARSE_BLOCK_SIZE, HISPARSE_BLOCK_SIZE
+    )
+    new_block_ids = manager.take_new_block_ids()
+    coordinator = get_hisparse_coordinator(manager)
+    assert manager.take_kv_cache_block_copies() == ([], [])
+    connector = HiSparseConnectorScheduler(async_speculative=False)
+    connector.bind_coordinator(coordinator)
+    output = SchedulerOutput.make_empty()
+    copies = connector.build_connector_meta(output).host_block_copies
+    assert output.has_sync_kv_loads
+
+    assert new_blocks and new_block_ids == []
+    assert len(copies) == 1
+    assert copies[0].src_block_id == source_block.block_id
+    assert copies[0].dst_block_id == new_blocks[0].block_id
+    source_manager.free("cow")
+    assert source_block.ref_cnt > 0 and new_blocks[0].ref_cnt > 0
+    assert coordinator.has_pending_work()
+    connector.update_connector_output(
+        KVConnectorOutput(
+            kv_connector_worker_meta=HiSparseConnectorWorkerMetadata(
+                {}, {}, completed_host_copy_dst_ids=(copies[0].dst_block_id,)
+            )
+        )
+    )
+    assert source_block.ref_cnt == new_blocks[0].ref_cnt == 0
+    assert not coordinator.has_pending_work()
+
+
+def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages():
+    """An in-flight host import must reserve its unwritten resident pages."""
+    manager = make_hisparse_kv_cache_manager(
+        12,
+        16,
+        transfer_device_cache=True,
+    )
+    request = make_request(
+        "partial-import",
+        list(range(6 * HISPARSE_BLOCK_SIZE)),
+        HISPARSE_BLOCK_SIZE,
+        sha256,
+    )
+    imported_tokens = 4 * HISPARSE_BLOCK_SIZE
+
+    assert allocate_external_prefix(manager, request, imported_tokens) is not None
+    required = manager.coordinator.get_num_blocks_to_allocate(
+        request_id=request.request_id,
+        num_tokens=request.num_tokens,
+        new_computed_blocks=manager.empty_kv_cache_blocks.blocks,
+        num_encoder_tokens=0,
+        total_computed_tokens=imported_tokens,
+        num_local_computed_tokens=0,
+        num_tokens_main_model=request.num_tokens,
+        apply_admission_cap=True,
+    )
+
+    assert required == 4
+
+
+@pytest.mark.parametrize(
+    "full_sequence_must_fit,host_num_blocks,admitted",
+    [(True, 2, False), (True, 3, True), (False, 2, False), (False, 3, True)],
+)
+def test_hisparse_async_admission_requires_only_import_destinations(
+    full_sequence_must_fit, host_num_blocks, admitted, tmp_path
+):
+    """Imports need host destinations, but do not reserve future prefill pages."""
+    from .utils import create_scheduler, mock_kv
+
+    (tmp_path / "config.json").write_text(
+        '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
+    )
+    scheduler = create_scheduler(
+        model=str(tmp_path),
+        skip_tokenizer_init=True,
+        max_model_len=128,
+        use_kv_connector=mock_kv(matched_tokens=HISPARSE_BLOCK_SIZE, is_async=True),
+    )
+    manager = make_hisparse_kv_cache_manager(32, host_num_blocks)
+    scheduler.kv_cache_manager = manager
+    scheduler.kv_cache_config = manager.kv_cache_config
+    scheduler.scheduler_reserve_full_isl = full_sequence_must_fit
+    inflight = make_request(
+        "inflight", list(range(3 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
+    )
+    scheduler.add_request(inflight)
+    scheduler.schedule()
+    assert inflight in scheduler._inflight_prefills
+
+    request = make_request(
+        "waiting", list(range(4 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
+    )
+    scheduler.add_request(request)
+    scheduler.schedule()
+
+    assert (request in scheduler._inflight_prefills) == admitted
+    assert bool(manager.get_blocks(request.request_id).blocks[0]) == admitted
+
+
+@pytest.mark.parametrize(
+    "evictable,local_tokens,free_blocks,admitted",
+    [
+        (False, HISPARSE_BLOCK_SIZE, 1, True),
+        (True, HISPARSE_BLOCK_SIZE, 1, False),
+        (False, HISPARSE_BLOCK_SIZE - 1, 1, False),
+        (False, HISPARSE_BLOCK_SIZE - 1, 2, True),
+    ],
+)
+def test_hisparse_import_capacity_includes_hits_and_cow(
+    evictable, local_tokens, free_blocks, admitted
+):
+    """Touching an evictable hit and copying a partial hit consume capacity."""
+    manager = make_hisparse_kv_cache_manager(32, free_blocks + 2 - int(evictable))
+    coordinator = get_hisparse_coordinator(manager)
+    host_pool = coordinator.get_host_block_pool()
+    assert host_pool is not None
+    hit = host_pool.get_new_blocks(1)[0]
+    if evictable:
+        host_pool.free_blocks([hit])
+    assert host_pool.get_num_free_blocks() == free_blocks
+    assert (
+        coordinator.host_manager.get_num_blocks_to_allocate(
+            "waiting", 64, [hit], 32, local_tokens, 64
+        )
+        <= host_pool.get_num_free_blocks()
+    ) == admitted
+
+
+@pytest.mark.parametrize("enable_caching", [False, True])
+def test_hisparse_host_exhaustion_keeps_gpu_pages_readable(enable_caching):
+    """Pages without a host destination must survive GPU cache reclamation."""
+    manager = make_hisparse_kv_cache_manager(32, 2, enable_caching=enable_caching)
+    request = make_request("resident", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, 64) is not None
+    coordinator = get_hisparse_coordinator(manager)
+    host = coordinator.host_manager
+    assert host is not None
+    assert [b.is_null for b in host.req_to_blocks[request.request_id]] == [
+        False,
+        True,
+        True,
+        True,
+    ]
+    _publish_hisparse_pages(manager)
+    for hot in coordinator.hot_managers:
+        hot.require_hot(request.request_id)
+        hot.allocate_new_blocks(request.request_id, 64, 64)
+    coordinator.update_residency(request.request_id)
+    for resident in coordinator.resident_managers:
+        for block in resident.req_to_blocks[request.request_id][1:]:
+            assert not block.is_null and block.ref_cnt > 0
+    assert coordinator.request_states[request.request_id].valid_pages == {0}
+    pool = coordinator.get_host_block_pool()
+    assert pool is not None
+    manager.free(request)
+    assert pool.get_num_free_blocks() == 1
+
+
+def test_hisparse_host_cow_takes_priority_over_new_pages():
+    """The last host block must hold the private tail, not a fresh page."""
+    manager = make_hisparse_kv_cache_manager(32, 3)
+    coordinator = get_hisparse_coordinator(manager)
+    host = coordinator.host_manager
+    assert host is not None
+    source = host.block_pool.get_new_blocks(1)[0]
+    host.req_to_blocks["cow"] = [source]
+    host._partial_hit_reqs["cow"] = (0, source)
+
+    host.allocate_new_blocks("cow", 32, 32)
+
+    private, missing = host.req_to_blocks["cow"]
+    assert not private.is_null and private is not source
+    assert missing.is_null
+    assert host.take_host_cow_copies() == [(source, private)]
+
+
+@pytest.mark.parametrize("failed", [False, True])
+def test_host_receive_completion_without_spill_metadata(failed):
+    """Only successful external receives make host pages readable."""
+    manager = make_hisparse_kv_cache_manager(32, 16)
+    request = make_request("import", list(range(32)), HISPARSE_BLOCK_SIZE, sha256)
+    assert allocate_external_prefix(manager, request, 32) is not None
+    request.num_computed_tokens = 32
+    request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
+    scheduler = HiSparseConnectorScheduler(async_speculative=False)
+    scheduler.bind_coordinator(get_hisparse_coordinator(manager))
+    scheduler.requests[request.request_id] = request
+    connector = object.__new__(HiSparseConnector)
+    connector.connector_scheduler = scheduler
+    composite = object.__new__(MultiConnector)
+    composite._connectors = (connector,)
+    composite.update_connector_output(
+        KVConnectorOutput(
+            finished_recving={request.request_id},
+            failed_recving={request.request_id} if failed else set(),
+        )
+    )
+    state = get_hisparse_coordinator(manager).request_states.get(request.request_id)
+    if failed:
+        assert state is None
+    else:
+        assert state is not None and state.valid_pages == {0, 1}
+
+
+def test_hisparse_recomputation_does_not_complete_failed_host_import():
+    """Allocating recomputation is not evidence that imported host KV is valid."""
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    request = make_request(
+        "partial-failure", list(range(48)), HISPARSE_BLOCK_SIZE, sha256
+    )
+    assert allocate_external_prefix(manager, request, 32) is not None
+    # A partial load failure leaves only the first page computed. The next
+    # allocation reaches the original import boundary before the worker runs.
+    request.num_computed_tokens = 16
+    manager.cache_blocks(request, 16)
+    assert manager.allocate_slots(request, 16) is not None
+    state = get_hisparse_coordinator(manager).request_states.get(request.request_id)
+    assert state is None or 1 not in state.valid_pages
+
+
+def test_hisparse_host_import_ignores_unsealed_tail():
+    """A partial imported page must not become readable from stale host data."""
+    manager = make_hisparse_kv_cache_manager(16, 16)
+    coordinator = get_hisparse_coordinator(manager)
+
+    coordinator.complete_host_import("partial", HISPARSE_BLOCK_SIZE + 1)
+
+    state = coordinator.request_states["partial"]
+    assert state.valid_pages == {0}
+    assert state.ready_prefix_pages == 1
+
+
+def test_hisparse_resident_request_can_grow_without_hot_capacity():
+    """Running resident history must not be mistaken for a new host-prefix hit."""
+    manager = make_hisparse_kv_cache_manager(8, 16)
+    request = make_request("resident", list(range(48)), HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, num_new_tokens=32) is not None
+    request.num_computed_tokens = 32
+
+    # Three free blocks fit the next indexer/resident pages, but not a hot region.
+    assert manager.block_pool.get_num_free_blocks() == 3
+    assert manager.allocate_slots(request, num_new_tokens=1) is not None
+    assert manager.get_block_ids(request.request_id)[3] == []
+
+
+def test_hisparse_capacity_query_does_not_require_hot_blocks():
+    """A read-only capacity query must not mutate hot-block requirements."""
+    manager = make_hisparse_kv_cache_manager(16, 16)
+    coordinator = get_hisparse_coordinator(manager)
+    host_pool = coordinator.get_host_block_pool()
+    assert host_pool is not None
+    computed: tuple[list[KVCacheBlock], ...] = (
+        [host_pool.blocks[1]],
+        [],
+        [],
+        [],
+    )
+
+    manager.coordinator.get_num_blocks_to_allocate(
+        request_id="query-only",
+        num_tokens=HISPARSE_BLOCK_SIZE,
+        new_computed_blocks=computed,
+        num_encoder_tokens=0,
+        total_computed_tokens=HISPARSE_BLOCK_SIZE,
+        num_local_computed_tokens=HISPARSE_BLOCK_SIZE,
+        num_tokens_main_model=HISPARSE_BLOCK_SIZE,
+    )
+
+    assert all(
+        "query-only" not in hot_manager.hot_required
+        for hot_manager in coordinator.hot_managers
+    )
+
+
+def _publish_hisparse_pages(manager: KVCacheManager) -> None:
+    """Ack all planned spills so host pages publish to the prefix cache."""
+    command = get_hisparse_coordinator(manager).build_offload_command()
+    counts = {spill.transfer_id: 1 for spill in command.page_transfers}
+    get_hisparse_coordinator(manager).update_spills(counts, counts)
+
+
+def test_hisparse_events_report_host_and_device_placement():
+    """Host publication and eviction must not advertise GPU-resident KV."""
+    manager = make_kv_cache_manager(
+        make_hisparse_kv_cache_config(32, 16),
+        max_model_len=128,
+        enable_caching=True,
+        hash_block_size=HISPARSE_BLOCK_SIZE,
+        enable_kv_cache_events=True,
+    )
+    request = make_request("request", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, 64) is not None
+    _publish_hisparse_pages(manager)
+
+    stored = [
+        event for event in manager.take_events() if isinstance(event, BlockStored)
+    ]
+    assert {(event.group_idx, event.medium) for event in stored} == {
+        (0, MEDIUM_CPU),
+        (1, MEDIUM_GPU),
+    }
+
+    host_block = manager.get_blocks(request.request_id).blocks[0][0]
+    get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
+        {host_block.block_id}
+    )
+    [removed] = manager.take_events()
+    assert isinstance(removed, BlockRemoved)
+    assert removed.medium == MEDIUM_CPU
+    assert removed.group_idx == 0
+
+
+def test_mamba_boundary_handoffs_do_not_pin_obsolete_blocks():
+    """Draining unused connector offers must not retain old recurrent states."""
+    manager = make_kv_cache_manager(
+        KVCacheConfig(
+            num_blocks=16,
+            kv_cache_tensors=[],
+            kv_cache_groups=[
+                KVCacheGroupSpec(
+                    ["mamba"],
+                    MambaSpec(
+                        block_size=16,
+                        shapes=((1,),),
+                        dtypes=(torch.float32,),
+                        mamba_cache_mode="align",
+                    ),
+                )
+            ],
+        ),
+        max_model_len=128,
+        enable_caching=True,
+        hash_block_size=16,
+    )
+    request = make_request("request", list(range(64)), 16, sha256)
+    old_blocks = []
+    for end in (16, 32, 48):
+        assert manager.allocate_slots(request, 16) is not None
+        request.num_computed_tokens = end
+        offers = manager.take_boundary_state_offloads()
+        [(_, block_id, boundary)] = offers[request.request_id]
+        assert boundary == end
+        old_blocks.append(manager.block_pool.blocks[block_id])
+        manager.remove_skipped_blocks(request.request_id, end, 64)
+        manager.new_step_starts()
+
+    assert all(block.ref_cnt == 0 for block in old_blocks[:-1])
+
+
+def test_hisparse_prefix_hit_adopts_gpu_shadow_pages():
+    """A host prefix hit must come back GPU-resident while shadows survive."""
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
+    _publish_hisparse_pages(manager)
+    original_resident_ids = [
+        block.block_id for block in manager.get_blocks("original").blocks[2]
+    ]
+    manager.free(original)
+
+    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    computed, num_computed, _ = manager.get_computed_blocks(resumed)
+    assert num_computed == 3 * HISPARSE_BLOCK_SIZE
+    assert manager.allocate_slots(
+        resumed,
+        num_new_tokens=len(tokens) - num_computed,
+        num_new_computed_tokens=num_computed,
+        new_computed_blocks=computed,
+    )
+    resident_blocks = manager.get_blocks("resumed").blocks[2]
+    assert [block.block_id for block in resident_blocks[:3]] == (
+        original_resident_ids[:3]
+    )
+    assert not any(block.is_null for block in resident_blocks[:3])
+    assert get_hisparse_coordinator(manager).all_context_pages_resident(
+        ((resumed.request_id, num_computed, len(tokens) - num_computed),)
+    )
+
+
+def test_hisparse_host_backed_request_accepts_local_prefix_hit():
+    """Local group-completion hits need no external host import allocation."""
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
+    _publish_hisparse_pages(manager)
+    manager.free(original)
+
+    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    computed, num_computed, _ = manager.get_computed_blocks(resumed)
+    assert num_computed == 3 * HISPARSE_BLOCK_SIZE
+    assert manager.allocate_slots(
+        resumed,
+        num_new_tokens=len(tokens) - num_computed,
+        num_new_computed_tokens=num_computed,
+        new_computed_blocks=computed,
+        num_external_computed_tokens=0,
+    )
+
+
+def test_hisparse_finished_request_leaves_free_copies():
+    """Clean pages of a finished request are free-queue blocks, not pins."""
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
+    _publish_hisparse_pages(manager)
+    manager.free(original)
+
+    coordinator = get_hisparse_coordinator(manager)
+    assert coordinator.copies
+    pool = manager.block_pool
+    assert pool.get_num_free_blocks() == pool.num_gpu_blocks - 1
+    pool.get_new_blocks(pool.get_num_free_blocks())
+    assert not coordinator.copies
+    assert not coordinator.spills_to_send
+
+
+def test_hisparse_reset_prefix_cache_drops_copies():
+    """Reset must forget GPU copies along with the host hashes they mirror."""
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None
+    _publish_hisparse_pages(manager)
+    manager.free(request)
+
+    coordinator = get_hisparse_coordinator(manager)
+    assert coordinator.copies
+    assert manager.reset_prefix_cache()
+    assert not coordinator.copies
+
+
+def test_hisparse_reused_copy_is_not_adopted():
+    """A host prefix hit must not adopt a GPU copy the pool has handed out."""
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None
+    _publish_hisparse_pages(manager)
+    indexer_blocks = list(manager.get_blocks("original").blocks[1])
+    manager.free(original)
+    pool = manager.block_pool
+    pool.touch(indexer_blocks)
+    reused = pool.get_new_blocks(pool.get_num_free_blocks())
+    pool.free_blocks(reused)
+    pool.free_blocks(indexer_blocks)
+
+    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    computed, num_computed, _ = manager.get_computed_blocks(resumed)
+    assert num_computed == 3 * HISPARSE_BLOCK_SIZE
+    assert manager.allocate_slots(
+        resumed,
+        num_new_tokens=len(tokens) - num_computed,
+        num_new_computed_tokens=num_computed,
+        new_computed_blocks=computed,
+    )
+    resident_blocks = manager.get_blocks("resumed").blocks[2]
+    assert all(block.is_null for block in resident_blocks[:3])
+
+
+def test_hisparse_external_import_uses_hard_gpu_footprint():
+    """An external prefix larger than resident capacity must remain admissible.
+
+    This guards the admission boundary where indexer + one writable resident
+    page + the fixed hot allocation fit, but indexer + full resident history do
+    not. Requiring one more resident page would silently strand long P/D inputs.
+    """
+    num_prompt_blocks = 4
+    num_prompt_tokens = num_prompt_blocks * HISPARSE_BLOCK_SIZE
+    manager = make_hisparse_kv_cache_manager(
+        8,
+        9,
+        transfer_device_cache=True,
+    )
+    request = make_request(
+        "host-import",
+        list(range(num_prompt_tokens)),
+        HISPARSE_BLOCK_SIZE,
+        sha256,
+    )
+
+    allocated = allocate_external_prefix(manager, request, num_prompt_tokens)
+
+    assert allocated is not None
+    source, indexer, resident, hot = manager.get_blocks(request.request_id).blocks
+    assert len(source) == num_prompt_blocks
+    assert len(indexer) == num_prompt_blocks
+    assert len(resident) == num_prompt_blocks
+    assert all(block.is_null for block in resident[:-1])
+    assert not resident[-1].is_null
+    assert len(hot) == 2
+    assert manager.block_pool.get_num_free_blocks() == 0
+
+
+def test_hisparse_external_import_survives_capacity_retry():
+    manager = make_hisparse_kv_cache_manager(
+        9,
+        7,
+        transfer_device_cache=True,
+    )
+    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
+    first = make_request("first", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    second = make_request("second", tokens, HISPARSE_BLOCK_SIZE, sha256)
+
+    assert allocate_external_prefix(manager, first, len(tokens)) is not None
+
+    assert allocate_external_prefix(manager, second, len(tokens)) is None
+
+    manager.free(first)
+    assert allocate_external_prefix(manager, second, len(tokens)) is not None
+    _, _, resident, hot = manager.get_blocks(second.request_id).blocks
+    assert all(block.is_null for block in resident[:-1])
+    assert not resident[-1].is_null
+    assert len(hot) == 2
+
+
 def make_kv_cache_config_hybrid_model(
     block_size: int,
     num_blocks: int,
@@ -812,8 +1914,7 @@ def _test_partial_request_hit(
 def _make_hybrid_kv_cache_config(
     block_size: int, num_blocks: int, spec_types: list[str]
 ) -> KVCacheConfig:
-    """
-    Create a KVCacheConfig with the specified spec types.
+    """Create a KVCacheConfig with the specified spec types.
 
     Args:
         block_size: The block size for KV cache.
@@ -823,6 +1924,7 @@ def _make_hybrid_kv_cache_config(
             - "sliding_window": SlidingWindowSpec with window=2*block_size
             - "sliding_window_large": SlidingWindowSpec with window=4*block_size
             - "mamba": MambaSpec
+
     """
     spec_map = {
         "full": lambda: FullAttentionSpec(
@@ -924,8 +2026,7 @@ def _make_hybrid_kv_cache_config(
 
 @pytest.mark.parametrize("spec_types", _HYBRID_MODEL_TEST_CASES)
 def test_prefill_hybrid_model_combinations(spec_types: list[str]):
-    """
-    Test prefix caching with hybrid models containing various combinations of
+    """Test prefix caching with hybrid models containing various combinations of
     KV cache spec types.
 
     This unified test covers:
@@ -1007,8 +2108,7 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):
 def test_prefill_hybrid_model_combinations_eagle(
     spec_types: list[str], expect_hit_length: int
 ):
-    """
-    Test prefix caching with hybrid models (1 full attn + 1 other) with EAGLE.
+    """Test prefix caching with hybrid models (1 full attn + 1 other) with EAGLE.
     More complex hybrid models with EAGLE are not yet supported (see issue #32802).
     """
     block_size = 16
@@ -1169,7 +2269,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
         use_eagle_block_drop=False,
         hash_block_size=block_size,
         mamba_partial_cache_hit=False,
-        mamba_fine_grained_prefix_cache=False,
+        mamba_shared_prefix_checkpoint=False,
         mamba_has_prefill_checkpoint_blocks=False,
     )
     req_2.shared_prefix_boundary = shared_prefix_boundary
@@ -1475,8 +2575,7 @@ def test_evict():
 
 
 def test_hash_block_correct_reuse():
-    """
-    This tests when a previously cached block is reused as a new block,
+    """This tests when a previously cached block is reused as a new block,
     its hash metadata should be correctly reset.
     """
     block_size = 16
@@ -1516,8 +2615,7 @@ def test_hash_block_correct_reuse():
 
 
 def test_computed_blocks_not_evicted():
-    """
-    Test that the computed blocks are not evicted when getting new blocks
+    """Test that the computed blocks are not evicted when getting new blocks
     for a request if there are any other free blocks.
     """
     block_size = 16
@@ -1576,9 +2674,7 @@ def test_computed_blocks_not_evicted():
 
 
 def test_basic_prefix_caching_disabled():
-    """
-    This tests that the prefix caching is disabled.
-    """
+    """This tests that the prefix caching is disabled."""
     block_size = 4
     manager = make_kv_cache_manager(
         make_kv_cache_config(block_size, 5),
@@ -1625,11 +2721,9 @@ def test_basic_prefix_caching_disabled():
 
 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
 def test_cache_blocks(hash_fn):
-    """
-    This is a unit test that tests the correctness of the _cache_full_blocks
+    """This is a unit test that tests the correctness of the _cache_full_blocks
     function of KVCacheManager.
     """
-
     block_size = 4
     block_pool = BlockPool(
         num_gpu_blocks=5,
@@ -1674,9 +2768,7 @@ def test_cache_blocks(hash_fn):
 
 
 def test_cache_blocks_multi_group():
-    """
-    This tests that blocks are cached correctly for different kv cache groups.
-    """
+    """This tests that blocks are cached correctly for different kv cache groups."""
     block_size = 4
     block_pool = BlockPool(
         num_gpu_blocks=10, enable_caching=True, hash_block_size=block_size
@@ -1759,10 +2851,7 @@ def test_cache_blocks_multi_group():
 
 
 def test_mm_prefix_caching():
-    """
-    This tests that the multi-modal prefix caching is correct.
-    """
-
+    """This tests that the multi-modal prefix caching is correct."""
     block_size = 16
     manager = make_kv_cache_manager(
         make_kv_cache_config(block_size, 11),
@@ -1867,8 +2956,7 @@ def test_mm_prefix_caching():
 
 
 def test_cache_key_salting():
-    """
-    This tests that cache salts are applied during hashing and the cache
+    """This tests that cache salts are applied during hashing and the cache
     is separated cache as expected.
     """
     block_size = 16
@@ -1947,8 +3035,7 @@ def test_cache_key_salting():
 
 
 def test_prefill_not_enough_free_blocks_with_computed_blocks():
-    """
-    This is a unit test that tests the correctness of the allocate_slots
+    """This is a unit test that tests the correctness of the allocate_slots
     when there is not enough free blocks. Specifically, when a request
     has computed blocks but cannot be allocated due to not enough free blocks,
     the computed blocks should not be touched.
@@ -2448,8 +3535,7 @@ def test_session_id_does_not_affect_prefix_cache_identity():
 
 
 def test_block_stored_event_group_idx_multiple_groups():
-    """
-    Test BlockStored events for separate HMA groups that each carry the
+    """Test BlockStored events for separate HMA groups that each carry the
     correct group_idx.
 
     Simulates the HMA scenario where full-attention blocks (group 0) and
@@ -2571,8 +3657,7 @@ def collect_warning(message, *args, **kwargs):
 
 @pytest.mark.parametrize("group_id", [0, 1, 2])
 def test_block_removed_event_group_idx(group_id: int):
-    """
-    Test BlockRemoved events emitted on eviction carry the group_idx extracted
+    """Test BlockRemoved events emitted on eviction carry the group_idx extracted
     from the evicted block's BlockHashWithGroupId via get_group_id().
     """
     block_size = 4
@@ -3874,6 +4959,42 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
     )
 
 
+@pytest.mark.parametrize(("reserved_blocks", "fits"), [(4, True), (10, False)])
+def test_full_sequence_admission_reserves_device_blocks_for_immediate_chunk(
+    reserved_blocks, fits
+):
+    """Reservations constrain the current allocation, not the entire prompt."""
+    block_size = 16
+    config = KVCacheConfig(
+        num_blocks=11,  # Ten usable blocks after the null block.
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["layer"],
+                FullAttentionSpec(
+                    block_size=block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            )
+        ],
+    )
+    manager = make_kv_cache_manager(
+        config, max_model_len=128, enable_caching=True, hash_block_size=block_size
+    )
+    request = make_request("chunk", list(range(128)), block_size, sha256)
+
+    blocks = manager.allocate_slots(
+        request,
+        block_size,
+        full_sequence_must_fit=True,
+        reserved_blocks=reserved_blocks,
+    )
+
+    assert (blocks is not None) is fits
+
+
 def test_can_fit_full_sequence_full_attention_still_gates_oversized():
     """The cap only loosens the SWA group; a prompt that exceeds the
     full-attention pool capacity must still be rejected."""
@@ -3926,6 +5047,30 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized():
     assert manager.allocate_slots(req, block_size, full_sequence_must_fit=True) is None
 
 
+def test_can_fit_full_sequence_hisparse_caps_resident_pages():
+    manager = make_hisparse_kv_cache_manager(
+        num_blocks=11,
+        host_num_blocks=9,
+        max_model_len=8 * HISPARSE_BLOCK_SIZE,
+        max_in_flight_tokens=2 * HISPARSE_BLOCK_SIZE,
+    )
+    request = make_request(
+        "long-hisparse",
+        list(range(8 * HISPARSE_BLOCK_SIZE)),
+        HISPARSE_BLOCK_SIZE,
+        sha256,
+    )
+
+    assert (
+        manager.allocate_slots(
+            request,
+            num_new_tokens=HISPARSE_BLOCK_SIZE,
+            full_sequence_must_fit=True,
+        )
+        is not None
+    )
+
+
 def test_cache_hit_local_and_external():
     # Regression test for #33775: when a request hits the local prefix cache
     # in one KV cache group and needs external (connector) blocks in another,
@@ -4771,3 +5916,50 @@ def distinct(v):
     assert last_req_hit(retention=0, pin=False) == 0
     # retention=0 with the pin keeps the junction window -> reuse restored.
     assert last_req_hit(retention=0, pin=True) == 4 * block_size
+
+
+def test_device_eviction_keeps_host_prefix():
+    """Device eviction must not target a same-numbered host block."""
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    request = make_request("producer", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, 64) is not None
+    _publish_hisparse_pages(manager)
+    host, indexer, _, _ = manager.get_blocks(request.request_id).blocks
+    device_block = indexer[0]
+    host_block = host[0]
+    assert device_block.block_id == host_block.block_id
+    manager.free(request)
+    assert host_block.block_hash is not None and device_block.block_hash is not None
+    manager.evict_blocks({device_block.block_id})
+    assert device_block.block_hash is None, (
+        "Device eviction left the indexer prefix cached"
+    )
+    assert host_block.block_hash is not None, (
+        "Device eviction invalidated unrelated host KV"
+    )
+
+
+def test_get_unhashed_block_ids_all_groups():
+    """Unhashed, non-null block ids are reported per KV cache group.
+    A group with no unhashed blocks reports an empty list, not a missing entry.
+    """
+
+    def hashed(block_id: int, group_id: int) -> KVCacheBlock:
+        block = KVCacheBlock(block_id=block_id)
+        block.set_block_hash(make_block_hash_with_group_id(BlockHash(b"h"), group_id))
+        return block
+
+    blocks = KVCacheBlocks(
+        (
+            [
+                KVCacheBlock(block_id=1),  # unhashed
+                hashed(2, group_id=0),  # cached
+                KVCacheBlock(block_id=3, is_null=True),  # null padding
+                KVCacheBlock(block_id=4),  # unhashed
+            ],
+            # No unhashed blocks in this group.
+            [hashed(5, group_id=1), KVCacheBlock(block_id=6, is_null=True)],
+        )
+    )
+
+    assert blocks.get_unhashed_block_ids_all_groups() == [[1, 4], []]
diff --git a/tests/v1/core/test_prefix_replay.py b/tests/v1/core/test_prefix_replay.py
new file mode 100644
index 000000000000..b2e78e1f925e
--- /dev/null
+++ b/tests/v1/core/test_prefix_replay.py
@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SWA bounded replay: after a prefix hit the scheduler recomputes the tail of
+the hit to rebuild the non-cacheable sliding-window group, keeps the hit's
+blocks, and hands the worker the replay start."""
+
+import pytest
+import torch
+
+from tests.v1.kv_connector.unit.utils import MockKVConfig, create_model_runner_output
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    SlidingWindowMLASpec,
+)
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import RequestStatus
+from vllm.v1.structured_output import StructuredOutputManager
+
+from .utils import create_requests, create_scheduler
+
+BLOCK_SIZE = 16
+WINDOW = 32
+NUM_PROMPT_TOKENS = 100
+# 100 tokens -> 6 full blocks cached -> 96-token hit.
+HIT_TOKENS = NUM_PROMPT_TOKENS // BLOCK_SIZE * BLOCK_SIZE
+FULL, SWA = 0, 1
+
+
+def _replay_scheduler(
+    *,
+    long_prefill_token_threshold: int = 0,
+    use_kv_connector: MockKVConfig | None = None,
+) -> Scheduler:
+    """A hybrid layout: one prefix-cacheable full-attention group and one
+    replayed sliding-window group."""
+    base = create_scheduler(
+        block_size=BLOCK_SIZE,
+        enable_prefix_caching=True,
+        long_prefill_token_threshold=long_prefill_token_threshold,
+        use_kv_connector=use_kv_connector,
+    )
+    vllm_config = base.vllm_config
+    kv_cache_config = KVCacheConfig(
+        num_blocks=10000,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=BLOCK_SIZE,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["swa"],
+                SlidingWindowMLASpec(
+                    block_size=BLOCK_SIZE,
+                    num_kv_heads=1,
+                    head_size=64,
+                    dtype=torch.bfloat16,
+                    sliding_window=WINDOW,
+                    bounded_replay=True,
+                ),
+            ),
+        ],
+    )
+    scheduler = Scheduler(
+        vllm_config=vllm_config,
+        kv_cache_config=kv_cache_config,
+        block_size=BLOCK_SIZE,
+        log_stats=True,
+        structured_output_manager=StructuredOutputManager(vllm_config),
+    )
+    scheduler.use_v2_model_runner = True
+    assert scheduler.prefix_replay_tokens == WINDOW
+    return scheduler
+
+
+def _prefill(scheduler: Scheduler, request) -> None:
+    scheduler.add_request(request)
+    out = scheduler.schedule()
+    assert out.num_scheduled_tokens[request.request_id] == request.num_prompt_tokens
+    scheduler.update_from_output(out, _step_output(out, [request]))
+
+
+def _step_output(out, requests) -> ModelRunnerOutput:
+    """Model output for every scheduled request: a token once its prefill is
+    done (non-zero, so outputs never extend the all-zero shared prompt in the
+    cache), nothing while it is still prefilling. schedule() has already
+    advanced num_computed_tokens past this step's chunk."""
+    scheduled = [r for r in requests if r.request_id in out.num_scheduled_tokens]
+    return ModelRunnerOutput(
+        req_ids=[r.request_id for r in scheduled],
+        req_id_to_index={r.request_id: i for i, r in enumerate(scheduled)},
+        sampled_token_ids=[
+            [1000] if r.num_computed_tokens >= r.num_prompt_tokens else []
+            for r in scheduled
+        ],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+    )
+
+
+def _new_req_data(out, request):
+    return next(r for r in out.scheduled_new_reqs if r.req_id == request.request_id)
+
+
+def test_hit_replays_window_without_reallocating():
+    scheduler = _replay_scheduler()
+    expected_replay = WINDOW
+    first, second = create_requests(
+        num_requests=2,
+        num_tokens=NUM_PROMPT_TOKENS,
+        same_prompt=True,
+        block_size=BLOCK_SIZE,
+    )
+    _prefill(scheduler, first)
+    manager = scheduler.kv_cache_manager
+
+    scheduler.add_request(second)
+    out = scheduler.schedule()
+
+    replay_start = HIT_TOKENS - expected_replay
+    new_req = _new_req_data(out, second)
+    assert new_req.num_computed_tokens == replay_start
+    assert new_req.replay_start == replay_start
+    assert (
+        out.num_scheduled_tokens[second.request_id] == NUM_PROMPT_TOKENS - replay_start
+    )
+
+    first_blocks = manager.get_blocks(first.request_id).blocks
+    second_blocks = manager.get_blocks(second.request_id).blocks
+    # The full-attention hit is shared (replayed tokens included); only the
+    # tail block is new.
+    num_hit_blocks = HIT_TOKENS // BLOCK_SIZE
+    assert second_blocks[FULL][:num_hit_blocks] == first_blocks[FULL][:num_hit_blocks]
+    assert all(b.ref_cnt == 2 for b in second_blocks[FULL][:num_hit_blocks])
+    assert len(second_blocks[FULL]) == num_hit_blocks + 1
+    # The window group never hits: blocks below the window are null, the
+    # window (replayed and new tokens) gets fresh blocks.
+    num_skipped_blocks = (HIT_TOKENS - WINDOW + 1) // BLOCK_SIZE
+    assert all(b.is_null for b in second_blocks[SWA][:num_skipped_blocks])
+    num_window_blocks = len(second_blocks[SWA]) - num_skipped_blocks
+    assert num_window_blocks == -(-NUM_PROMPT_TOKENS // BLOCK_SIZE) - num_skipped_blocks
+    assert not any(b.is_null for b in second_blocks[SWA][num_skipped_blocks:])
+
+    scheduler.update_from_output(out, _step_output(out, [first, second]))
+    assert second.num_computed_tokens == NUM_PROMPT_TOKENS
+
+
+def test_replay_only_chunks_make_progress():
+    """A chunk cap smaller than the replay window schedules replay-only
+    chunks; they advance the request instead of stalling it."""
+    scheduler = _replay_scheduler(long_prefill_token_threshold=BLOCK_SIZE)
+    first, second = create_requests(
+        num_requests=2,
+        num_tokens=NUM_PROMPT_TOKENS,
+        same_prompt=True,
+        block_size=BLOCK_SIZE,
+    )
+    scheduler.add_request(first)
+    while first.num_computed_tokens < NUM_PROMPT_TOKENS:
+        out = scheduler.schedule()
+        scheduler.update_from_output(out, _step_output(out, [first]))
+
+    scheduler.add_request(second)
+
+    def full_blocks():
+        return scheduler.kv_cache_manager.get_blocks(second.request_id).blocks[FULL]
+
+    scheduled: list[int] = []
+    while second.num_computed_tokens < NUM_PROMPT_TOKENS:
+        out = scheduler.schedule()
+        if not scheduled:
+            new_req = _new_req_data(out, second)
+            assert new_req.num_computed_tokens == HIT_TOKENS - WINDOW
+            assert new_req.replay_start == HIT_TOKENS - WINDOW
+            # A replay-only chunk adopts the hit and allocates nothing new.
+            assert len(full_blocks()) == HIT_TOKENS // BLOCK_SIZE
+        scheduled.append(out.num_scheduled_tokens[second.request_id])
+        scheduler.update_from_output(out, _step_output(out, [first, second]))
+    # 32 replayed + 4 new tokens in 16-token chunks: two replay-only chunks.
+    assert scheduled == [BLOCK_SIZE, BLOCK_SIZE, NUM_PROMPT_TOKENS - HIT_TOKENS]
+    assert len(full_blocks()) == -(-NUM_PROMPT_TOKENS // BLOCK_SIZE)
+    assert second.status == RequestStatus.RUNNING
+
+
+def test_async_remote_kv_hit_replays_after_load():
+    """A KV-connector hit loaded asynchronously carries no window state
+    either; the replay is applied when the load completes."""
+    matched = 64
+    scheduler = _replay_scheduler(
+        use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=True)
+    )
+    request = create_requests(
+        num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
+    )[0]
+    scheduler.add_request(request)
+    out = scheduler.schedule()
+    assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert request.request_id not in out.num_scheduled_tokens
+    scheduler.update_from_output(
+        out, create_model_runner_output([], finished_recving={request.request_id})
+    )
+
+    out = scheduler.schedule()
+    new_req = _new_req_data(out, request)
+    assert new_req.num_computed_tokens == matched - WINDOW
+    assert new_req.replay_start == matched - WINDOW
+    assert (
+        out.num_scheduled_tokens[request.request_id]
+        == NUM_PROMPT_TOKENS - matched + WINDOW
+    )
+
+
+def test_remote_kv_hit_is_taken_in_whole_blocks():
+    """A hit ending one token short of a block boundary would put the replay
+    window's first token in a block the sliding-window group retires, so a
+    connector hit is cut back to whole blocks."""
+    matched = 4 * BLOCK_SIZE - 1
+    scheduler = _replay_scheduler(
+        use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=True)
+    )
+    request = create_requests(
+        num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
+    )[0]
+    scheduler.add_request(request)
+    out = scheduler.schedule()
+    scheduler.update_from_output(
+        out, create_model_runner_output([], finished_recving={request.request_id})
+    )
+    out = scheduler.schedule()
+    new_req = _new_req_data(out, request)
+    hit = 3 * BLOCK_SIZE
+    assert new_req.replay_start == hit - WINDOW
+    assert new_req.num_computed_tokens == hit - WINDOW
+    swa_manager = scheduler.kv_cache_manager.coordinator.single_type_managers[SWA]
+    swa_blocks = scheduler.kv_cache_manager.get_blocks(request.request_id).blocks[SWA]
+    assert swa_blocks[new_req.replay_start // BLOCK_SIZE] is not swa_manager._null_block
+
+
+@pytest.mark.parametrize("via_connector", [False, True])
+def test_hit_no_longer_than_window_is_ignored(via_connector):
+    """Such a hit would be recomputed in full anyway. It is not adopted (an
+    async load of it is not even started), so a request has computed tokens
+    iff it replays."""
+    if via_connector:
+        scheduler = _replay_scheduler(
+            use_kv_connector=MockKVConfig(matched_tokens=WINDOW, is_async=True)
+        )
+        request = create_requests(
+            num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
+        )[0]
+    else:
+        scheduler = _replay_scheduler()
+        first = create_requests(
+            num_requests=1, num_tokens=WINDOW, same_prompt=True, block_size=BLOCK_SIZE
+        )[0]
+        _prefill(scheduler, first)
+        request = create_requests(
+            num_requests=1,
+            num_tokens=NUM_PROMPT_TOKENS,
+            same_prompt=True,
+            block_size=BLOCK_SIZE,
+            req_ids=["long"],
+        )[0]
+    scheduler.add_request(request)
+    out = scheduler.schedule()
+    new_req = _new_req_data(out, request)
+    assert request.status == RequestStatus.RUNNING
+    assert new_req.num_computed_tokens == 0
+    assert new_req.replay_start == 0
+    assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS
diff --git a/tests/v1/core/test_repetition_detection.py b/tests/v1/core/test_repetition_detection.py
index aae6e3b70cae..6814a6281433 100644
--- a/tests/v1/core/test_repetition_detection.py
+++ b/tests/v1/core/test_repetition_detection.py
@@ -14,10 +14,10 @@
 
 
 class TestCheckSequenceRepetition:
-    """Unit tests for the check_sequence_repetition function"""
+    """Unit tests for the check_sequence_repetition function."""
 
     def test_simple_repetition_detected(self):
-        """Test detection of simple repetitive patterns"""
+        """Test detection of simple repetitive patterns."""
         token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 3]
         params = RepetitionDetectionParams(
             max_pattern_size=3,
@@ -27,7 +27,7 @@ def test_simple_repetition_detected(self):
         assert check_sequence_repetition(token_ids, params)
 
     def test_repetition_below_min_count(self):
-        """Test that pattern below min_count is not detected"""
+        """Test that pattern below min_count is not detected."""
         token_ids = [1, 2, 3, 1, 2, 3]
         params = RepetitionDetectionParams(
             max_pattern_size=3,
@@ -37,7 +37,7 @@ def test_repetition_below_min_count(self):
         assert not check_sequence_repetition(token_ids, params)
 
     def test_two_token_pattern(self):
-        """Test detection of 2-token patterns"""
+        """Test detection of 2-token patterns."""
         token_ids = [1, 2, 1, 2, 1, 2, 1, 2]
         params = RepetitionDetectionParams(
             max_pattern_size=5,
@@ -47,7 +47,7 @@ def test_two_token_pattern(self):
         assert check_sequence_repetition(token_ids, params)
 
     def test_no_repetition_varied_sequence(self):
-        """Test that non-repetitive sequences are not flagged"""
+        """Test that non-repetitive sequences are not flagged."""
         token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9]
         params = RepetitionDetectionParams(
             max_pattern_size=5,
@@ -57,7 +57,7 @@ def test_no_repetition_varied_sequence(self):
         assert not check_sequence_repetition(token_ids, params)
 
     def test_partial_repetition_not_detected(self):
-        """Test that incomplete repetitions are not detected"""
+        """Test that incomplete repetitions are not detected."""
         token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 4]
         params = RepetitionDetectionParams(
             max_pattern_size=3,
@@ -67,7 +67,7 @@ def test_partial_repetition_not_detected(self):
         assert not check_sequence_repetition(token_ids, params)
 
     def test_empty_token_list(self):
-        """Test with empty token list"""
+        """Test with empty token list."""
         params = RepetitionDetectionParams(
             max_pattern_size=3,
             min_pattern_size=2,
@@ -76,19 +76,19 @@ def test_empty_token_list(self):
         assert not check_sequence_repetition([], params)
 
     def test_detection_disabled_max_size_zero(self):
-        """Test that zero max_pattern_size disables detection"""
+        """Test that zero max_pattern_size disables detection."""
         token_ids = [1, 2, 1, 2, 1, 2]
         params = RepetitionDetectionParams()
         assert not check_sequence_repetition(token_ids, params)
 
     def test_invalid_min_count(self):
-        """Test that min_count < 2 returns False"""
+        """Test that min_count < 2 returns False."""
         token_ids = [1, 2, 1, 2]
         params = RepetitionDetectionParams()
         assert not check_sequence_repetition(token_ids, params)
 
     def test_repetition_at_end_of_sequence(self):
-        """Test detection when repetition occurs at the end"""
+        """Test detection when repetition occurs at the end."""
         token_ids = [1, 2, 3, 4, 5, 6, 5, 6, 5, 6]
         params = RepetitionDetectionParams(
             max_pattern_size=3,
@@ -98,7 +98,7 @@ def test_repetition_at_end_of_sequence(self):
         assert check_sequence_repetition(token_ids, params)
 
     def test_large_pattern_many_repetitions(self):
-        """Test large pattern repeated many times"""
+        """Test large pattern repeated many times."""
         token_ids = [1, 2, 3, 4, 5, 6, 7, 8] * 5
         params = RepetitionDetectionParams(
             max_pattern_size=10,
@@ -114,10 +114,10 @@ def test_large_pattern_many_repetitions(self):
 
 
 class TestRepetitionDetectionIntegration:
-    """Integration tests for repetition detection in check_stop"""
+    """Integration tests for repetition detection in check_stop."""
 
     def test_basic_repetition_stops_generation(self):
-        """Test that repetition is detected and stops generation"""
+        """Test that repetition is detected and stops generation."""
         params = SamplingParams(
             max_tokens=100,
             repetition_detection=RepetitionDetectionParams(
@@ -138,7 +138,7 @@ def test_basic_repetition_stops_generation(self):
         assert request.stop_reason == "repetition_detected"
 
     def test_detection_disabled_no_stop(self):
-        """Test that disabled detection doesn't stop generation"""
+        """Test that disabled detection doesn't stop generation."""
         params = SamplingParams(
             max_tokens=100,
         )
@@ -152,7 +152,7 @@ def test_detection_disabled_no_stop(self):
         assert not check_stop(request, max_model_len=1024)
 
     def test_repetition_respects_min_tokens(self):
-        """Test that repetition detection respects min_tokens"""
+        """Test that repetition detection respects min_tokens."""
         params = SamplingParams(
             min_tokens=10,
             max_tokens=100,
@@ -172,7 +172,7 @@ def test_repetition_respects_min_tokens(self):
         assert not check_stop(request, max_model_len=1024)
 
     def test_no_repetition_continues_generation(self):
-        """Test that non-repetitive tokens don't stop generation"""
+        """Test that non-repetitive tokens don't stop generation."""
         params = SamplingParams(
             max_tokens=100,
             repetition_detection=RepetitionDetectionParams(
@@ -191,7 +191,7 @@ def test_no_repetition_continues_generation(self):
         assert not check_stop(request, max_model_len=1024)
 
     def test_pattern_at_size_boundary(self):
-        """Test detection at exact pattern size boundary"""
+        """Test detection at exact pattern size boundary."""
         params = SamplingParams(
             max_tokens=100,
             repetition_detection=RepetitionDetectionParams(
@@ -211,7 +211,7 @@ def test_pattern_at_size_boundary(self):
         assert request.status == RequestStatus.FINISHED_REPETITION
 
     def test_multiple_pattern_sizes_checked(self):
-        """Test that function checks pattern sizes in range"""
+        """Test that function checks pattern sizes in range."""
         params = SamplingParams(
             max_tokens=100,
             repetition_detection=RepetitionDetectionParams(
@@ -231,7 +231,7 @@ def test_multiple_pattern_sizes_checked(self):
         assert request.status == RequestStatus.FINISHED_REPETITION
 
     def test_eos_takes_precedence_over_repetition(self):
-        """Test that EOS token stops before repetition check"""
+        """Test that EOS token stops before repetition check."""
         params = SamplingParams(
             max_tokens=100,
             stop_token_ids=[999],
@@ -252,7 +252,7 @@ def test_eos_takes_precedence_over_repetition(self):
         assert request.status == RequestStatus.FINISHED_STOPPED
 
     def test_min_pattern_size_filters_small_patterns(self):
-        """Test that min_pattern_size filters out smaller patterns"""
+        """Test that min_pattern_size filters out smaller patterns."""
         params = SamplingParams(
             max_tokens=100,
             repetition_detection=RepetitionDetectionParams(
@@ -271,7 +271,7 @@ def test_min_pattern_size_filters_small_patterns(self):
         assert not check_stop(request, max_model_len=1024)
 
     def test_high_repetition_threshold(self):
-        """Test that high min_count requires many repetitions"""
+        """Test that high min_count requires many repetitions."""
         params = SamplingParams(
             max_tokens=100,
             repetition_detection=RepetitionDetectionParams(
diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index 1c56eaa6a500..505eef5919d0 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/tests/v1/core/test_scheduler.py
@@ -19,7 +19,17 @@
     SpeculativeConfig,
     VllmConfig,
 )
+from vllm.distributed.aux_output_connector.connector import (
+    AuxOutputSchedulerConnector,
+    AuxRequestOutput,
+)
+from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
+    HiSparseConnector,
+    HiSparseConnectorScheduler,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalKwargsItem,
@@ -30,10 +40,18 @@
 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
 from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
+from vllm.v1.core.sched.diffusion_scheduler import (
+    DiffusionAsyncScheduler,
+    diffusion_canvas_width,
+)
+from vllm.v1.core.sched.interface import PauseState
 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
 from vllm.v1.engine import FinishReason
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.executor.uniproc_executor import UniProcExecutor
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
@@ -125,8 +143,45 @@ def test_add_requests():
         assert len(scheduler.waiting) == i + 1
 
 
+@pytest.mark.parametrize("max_tokens", [1, 2])
+def test_routed_experts_prompt_start_at_prompt_end(max_tokens):
+    """A prompt-end offset should return empty routing data, not crash."""
+    scheduler = create_scheduler()
+    (request,) = create_requests(num_requests=1, max_tokens=max_tokens)
+    request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens
+    scheduler.aux_output_connector = AuxOutputSchedulerConnector()
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    assert scheduler_output.aux_output_connector_metadata.requests == {
+        request.request_id: request.num_prompt_tokens
+    }
+
+    outputs = scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[0]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+            aux_output_connector_output={
+                request.request_id: AuxRequestOutput(
+                    request.num_prompt_tokens,
+                    np.empty((0, 1, 1), dtype=np.uint8),
+                )
+            },
+        ),
+    )
+
+    output = outputs[request.client_index].outputs[0]
+    assert output.new_token_ids == [0]
+    assert output.routed_experts.shape == (0, 1, 1)
+
+
 def test_finish_request():
     scheduler = create_scheduler()
+    scheduler.aux_output_connector = Mock()
     requests = create_requests(num_requests=10)
     for request in requests:
         scheduler.add_request(request)
@@ -135,6 +190,7 @@ def test_finish_request():
         scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
         assert request.request_id not in scheduler.requests
         assert len(scheduler.waiting) == 9 - i
+        scheduler.aux_output_connector.request_finished.assert_called_with(request)
 
 
 def test_get_num_unfinished_requests():
@@ -148,6 +204,46 @@ def test_get_num_unfinished_requests():
         assert scheduler.get_num_unfinished_requests() == len(requests) - i - 1
 
 
+def test_max_num_active_seqs_caps_admission():
+    """Admission is capped by max_num_active_seqs, not max_num_seqs."""
+    scheduler = create_scheduler(max_num_seqs=16, max_num_active_seqs=3)
+    requests = create_requests(num_requests=10)
+    for request in requests:
+        scheduler.add_request(request)
+
+    output = scheduler.schedule()
+
+    assert len(output.scheduled_new_reqs) == 3
+    assert len(scheduler.running) == 3
+    assert len(scheduler.waiting) == 7
+    # Execution capacity is untouched.
+    assert scheduler.max_num_running_reqs == 16
+
+
+def _bind_hisparse_connector(scheduler):
+    connector = object.__new__(HiSparseConnector)
+    connector.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False)
+    from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager
+
+    coordinator = get_hisparse_coordinator(make_hisparse_kv_cache_manager(16, 16))
+    connector.connector_scheduler.bind_coordinator(coordinator)
+    composite = object.__new__(MultiConnector)
+    composite._connectors = (connector,)
+    scheduler.connector = composite
+    return coordinator
+
+
+def test_pending_hisparse_spill_keeps_scheduler_alive():
+    """A final host spill must complete after the last request finishes."""
+    scheduler = create_scheduler()
+    coordinator = _bind_hisparse_connector(scheduler)
+    pending = coordinator.pending_spills
+    pending[0] = Mock()
+    assert scheduler.has_requests()
+    pending.clear()
+    assert not scheduler.has_requests()
+
+
 @pytest.mark.parametrize(
     "enable_prefix_caching, prompt_logprobs",
     [
@@ -226,7 +322,7 @@ def test_schedule_multimodal_requests():
 
 
 def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders():
-    """Async scheduling + PP: allow multi-step in-flight scheduling per request"""
+    """Async scheduling + PP: allow multi-step in-flight scheduling per request."""
     scheduler = create_scheduler(async_scheduling=True, pipeline_parallel_size=2)
     (req,) = create_requests(num_requests=1, num_tokens=8)
     scheduler.add_request(req)
@@ -670,6 +766,56 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
     assert output2.num_scheduled_tokens[requests[2].request_id] == 800 - 224 - 224
 
 
+def test_long_prefill_threshold_ignored_when_alone():
+    """A lone long prefill is not capped by the threshold: it has no other
+    request to starve."""
+    scheduler = create_scheduler(
+        max_num_batched_tokens=1024,
+        long_prefill_token_threshold=400,
+    )
+    request = create_requests(num_requests=1, num_tokens=2000)[0]
+    scheduler.add_request(request)
+
+    output = scheduler.schedule()
+    assert output.num_scheduled_tokens[request.request_id] == 1024
+
+
+def test_long_prefill_threshold_applies_with_other_requests():
+    """The threshold caps the prefill as soon as another request is queued."""
+    scheduler = create_scheduler(
+        max_num_batched_tokens=1024,
+        long_prefill_token_threshold=400,
+    )
+    long_req = create_requests(num_requests=1, num_tokens=2000)[0]
+    short_req = create_requests(num_requests=1, num_tokens=10, req_ids=["short"])[0]
+    for request in [long_req, short_req]:
+        scheduler.add_request(request)
+
+    output = scheduler.schedule()
+    assert output.num_scheduled_tokens[long_req.request_id] == 400
+    assert output.num_scheduled_tokens[short_req.request_id] == 10
+
+
+def test_long_prefill_threshold_floored_by_fair_share():
+    """With the adaptive flag, the effective threshold never falls below the
+    fair share of the token budget: max_num_batched_tokens / num queued +
+    running requests."""
+    scheduler = create_scheduler(
+        max_num_batched_tokens=1024,
+        long_prefill_token_threshold=100,
+        long_prefill_token_threshold_adaptive=True,
+    )
+    long_req = create_requests(num_requests=1, num_tokens=2000)[0]
+    short_req = create_requests(num_requests=1, num_tokens=10, req_ids=["short"])[0]
+    for request in [long_req, short_req]:
+        scheduler.add_request(request)
+
+    output = scheduler.schedule()
+    # 100 is below the fair share (1024 // 2 = 512), so the floor binds.
+    assert output.num_scheduled_tokens[long_req.request_id] == 512
+    assert output.num_scheduled_tokens[short_req.request_id] == 10
+
+
 def test_update_from_output_routes_sampling_masks_by_request():
     """Each request receives the sampler row at its own batch index."""
     scheduler = create_scheduler()
@@ -716,7 +862,7 @@ def test_update_from_output_routes_sampling_masks_by_request():
 
 
 def test_stop_via_update_from_output():
-    """Test stopping behavior through update_from_output"""
+    """Test stopping behavior through update_from_output."""
     scheduler = create_scheduler(num_speculative_tokens=1)
 
     # Test case 1: Stop on EOS token
@@ -1181,6 +1327,45 @@ def test_preempt_during_execution():
     assert requests[1].output_token_ids[0] == 42
 
 
+def test_pending_hisparse_reclamation_defers_preemption(monkeypatch):
+    """Pending reclamation must stall rather than preempt its owning request."""
+    scheduler = create_scheduler(enable_prefix_caching=False)
+    request = create_requests(num_requests=1, num_tokens=16, block_size=16)[0]
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[0]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        ),
+    )
+
+    monkeypatch.setattr(
+        scheduler.kv_cache_manager, "allocate_slots", Mock(return_value=None)
+    )
+    coordinator = _bind_hisparse_connector(scheduler)
+    monkeypatch.setattr(
+        scheduler.connector, "build_connector_meta", Mock(return_value=None)
+    )
+    monkeypatch.setattr(
+        coordinator,
+        "has_pending_reclamation",
+        Mock(return_value=True),
+    )
+
+    deferred_output = scheduler.schedule()
+
+    assert not deferred_output.num_scheduled_tokens
+    assert scheduler.running == [request]
+    assert request.status == RequestStatus.RUNNING
+    assert request.num_preemptions == 0
+
+
 def test_prefix_cache_query_not_inflated_by_connector_defer():
     """The GPU prefix-cache query is recorded at admission, so a request the
     connector defers several times is counted once, not once per retry."""
@@ -1245,6 +1430,31 @@ def test_preemption_re_records_prefix_cache_query():
     assert stats.preempted_requests == 1
 
 
+def test_preemption_processes_stale_aux_output():
+    scheduler = create_scheduler(enable_prefix_caching=True)
+    request = create_requests(num_requests=1)[0]
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    scheduler.running.remove(request)
+    scheduler.aux_output_connector = Mock()
+    scheduler._preempt_request(request, 0.0)
+    scheduler.aux_output_connector.request_finished.assert_called_once_with(request)
+
+    scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[1000]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        ),
+    )
+
+    scheduler.aux_output_connector.take_output.assert_called_once_with(request, None)
+
+
 def test_prefix_cache_stats_not_recorded_when_caching_disabled():
     """With prefix caching off there is no local lookup, so admitting a request
     records no phantom miss."""
@@ -1333,6 +1543,7 @@ def spy_allocate_slots(*args, **kwargs):
 
 def test_scheduler_reset_prefix_cache():
     scheduler = create_scheduler(enable_prefix_caching=True)
+    scheduler.aux_output_connector = Mock()
     requests = create_requests(num_requests=10)
     for request in requests:
         scheduler.add_request(request)
@@ -1349,10 +1560,20 @@ def test_scheduler_reset_prefix_cache():
     # Reset prefix cache should fail since there are still running requests
     # and they are taking KV cache
     assert not scheduler.reset_prefix_cache()
+    scheduler.aux_output_connector.reset.assert_not_called()
+
+    with pytest.raises(RuntimeError, match=r"pause\(mode='keep'\)"):
+        scheduler.reset_prefix_cache(reset_running_requests=True)
+
+    # pause(mode="keep") also waits for scheduled model outputs to drain.
+    scheduler.set_pause_state(PauseState.PAUSED_ALL)
+    with pytest.raises(RuntimeError, match="model output is in flight"):
+        scheduler.reset_prefix_cache(reset_running_requests=True)
+    for request in requests:
+        request.num_in_flight_tokens = 0
 
-    # Reset prefix cache with reset_running_requests=True. All running requests
-    # Should be pushed back to the waiting queue and kv cache should be freed
     assert scheduler.reset_prefix_cache(reset_running_requests=True)
+    scheduler.aux_output_connector.reset.assert_called_once_with()
 
     # Verify requests moved from running to waiting
     assert len(scheduler.waiting) == len(requests)
@@ -1362,6 +1583,89 @@ def test_scheduler_reset_prefix_cache():
         assert scheduler.waiting[i] == request
 
 
+@pytest.mark.parametrize("reset_successful", [False, True])
+def test_aux_output_reset_follows_kv_reset_result(reset_successful: bool):
+    scheduler = create_scheduler(enable_prefix_caching=True)
+    scheduler.aux_output_connector = Mock()
+    scheduler.kv_cache_manager.reset_prefix_cache = Mock(return_value=reset_successful)
+
+    assert scheduler.reset_prefix_cache() is reset_successful
+    assert scheduler.aux_output_connector.reset.call_count == int(reset_successful)
+
+
+def test_kv_cache_release_after_keep_pause_preserves_requests():
+    """A completed keep pause permits release and recomputation of live requests."""
+    scheduler = create_scheduler(max_num_seqs=1, enable_prefix_caching=True)
+    running, waiting = create_requests(num_requests=2)
+    for request in (running, waiting):
+        scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[running.request_id],
+            req_id_to_index={running.request_id: 0},
+            sampled_token_ids=[[1]],
+        ),
+    )
+    assert scheduler.kv_cache_manager.get_block_ids(running.request_id)[0]
+
+    core = object.__new__(EngineCore)
+    core.scheduler = scheduler
+    core.model_executor = Mock(is_sleeping=False)
+    core.mm_receiver_cache = None
+    core.batch_queue = None
+
+    assert core.pause_scheduler(mode="keep", clear_cache=False) is None
+    assert scheduler.running == [running]
+    assert list(scheduler.waiting) == [waiting]
+    assert not scheduler.has_requests()
+
+    core.release_kv_cache_memory()
+
+    core.model_executor.discard.assert_called_once_with(("kv_cache",))
+    assert scheduler.requests == {
+        request.request_id: request for request in (running, waiting)
+    }
+    assert not scheduler.running
+    assert list(scheduler.waiting) == [running, waiting]
+    assert running.status == RequestStatus.PREEMPTED
+    assert running.num_computed_tokens == 0
+    assert list(running.output_token_ids) == [1]
+    assert not scheduler.kv_cache_manager.get_block_ids(running.request_id)[0]
+    assert scheduler.schedule().total_num_scheduled_tokens == 0
+
+    assert core.wake_up(tags=["kv_cache"])
+    resumed = scheduler.schedule()
+    assert resumed.num_scheduled_tokens == {running.request_id: running.num_tokens}
+
+
+@pytest.mark.parametrize(
+    "sleeping_tags",
+    [{"weights", "kv_cache"}, {"weights"}, {"kv_cache"}],
+    ids=["full-sleep", "weights-only", "kv-only-or-repeated-release"],
+)
+def test_kv_cache_release_rejects_nonresident_memory(sleeping_tags):
+    """Reject release without resetting caches or discarding more memory."""
+    core = object.__new__(EngineCore)
+    core.scheduler = create_scheduler()
+    core.scheduler.set_pause_state(PauseState.PAUSED_ALL)
+    core.batch_queue = None
+    core._reset_caches = Mock()
+    core.model_executor = object.__new__(UniProcExecutor)
+    core.model_executor.sleeping_tags = sleeping_tags.copy()
+    core.model_executor.collective_rpc = Mock()
+
+    with pytest.raises(
+        RuntimeError, match="requires all executor memory to be resident"
+    ):
+        core.release_kv_cache_memory()
+
+    core._reset_caches.assert_not_called()
+    core.model_executor.collective_rpc.assert_not_called()
+    assert core.model_executor.sleeping_tags == sleeping_tags
+
+
 def test_reset_connector_cache_no_connector_is_no_op_success():
     """``reset_connector_cache`` must return True when no connector is
     configured.
@@ -2022,7 +2326,6 @@ def _assert_right_scheduler_output(
     expected_num_scheduled_tokens: int,
 ):
     """Check if SchedulerOutput is correct after remote KV cache hit."""
-
     # We should inject the kv_connector_metadata.
     assert len(output.kv_connector_metadata.requests) == num_requests
 
@@ -2040,7 +2343,6 @@ def _assert_right_kv_cache_manager(
     num_total_blocks: int,
 ):
     """Check whether KVCacheManager is correct after allocate."""
-
     # Make sure the request stats are right.
     EXPECTED_TOTAL_BLOCKS = num_tokens // block_size
     for req in requests:
@@ -2071,7 +2373,6 @@ def _step_until_done(
     model_runner_output: ModelRunnerOutput,
 ):
     """Loop over schedule(), update_from_output() until finished."""
-
     all_finished = False
     _ = scheduler.update_from_output(output, model_runner_output)
     while not all_finished:
@@ -2099,7 +2400,6 @@ def _num_waiting_requests(scheduler: Scheduler) -> int:
 
 def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
     """Cycle requests through a KV transfer cycle."""
-
     # Requests should first transition to WAITING_FOR_REMOTE_KVS
     output = scheduler.schedule()
     assert _num_waiting_requests(scheduler) == len(req_ids)
@@ -2181,13 +2481,30 @@ def test_has_sync_kv_loads(
     assert output.has_sync_kv_loads is expected_has_sync_loads
 
 
+def test_kv_connector_honors_skip_reading_prefix_cache():
+    """A request that must score every prompt row takes no external hit."""
+    BLOCK_SIZE = 16
+    scheduler = create_scheduler(
+        enable_prefix_caching=True,
+        use_kv_connector=mock_kv(matched_tokens=BLOCK_SIZE * 2, is_async=False),
+        block_size=BLOCK_SIZE,
+    )
+    plain, scoring = create_requests(num_requests=2, num_tokens=BLOCK_SIZE * 4)
+    scoring.sampling_params.prompt_logprobs = 1
+    scoring.skip_reading_prefix_cache = True
+    scheduler.add_request(plain)
+    scheduler.add_request(scoring)
+
+    output = scheduler.schedule()
+    assert output.num_scheduled_tokens[plain.request_id] == BLOCK_SIZE * 2
+    assert output.num_scheduled_tokens[scoring.request_id] == BLOCK_SIZE * 4
+
+
 @pytest.mark.parametrize("is_async", [False, True])
 def test_kv_connector_basic(is_async: bool):
-    """
-    Test whether Scheduler with KVConnector schedules tokens, allocates
+    """Test whether Scheduler with KVConnector schedules tokens, allocates
     memory, and cleans up requests as expected under normal operation.
     """
-
     # Setup Scheduler.
     BLOCK_SIZE = 16
     NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
@@ -2312,11 +2629,9 @@ def test_kv_connector_basic(is_async: bool):
 @pytest.mark.parametrize("is_async", [False, True])
 @pytest.mark.parametrize("local_cache_hits", [False, True])
 def test_external_prefix_cache_metrics(is_async: bool, local_cache_hits: bool):
-    """
-    Verify connector prefix cache metrics are updated
+    """Verify connector prefix cache metrics are updated
     correctly when the scheduler processes requests with KV connector hits.
     """
-
     BLOCK_SIZE = 16
     if local_cache_hits:
         NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2  # 32 tokens
@@ -2432,11 +2747,9 @@ def test_external_prefix_cache_metrics(is_async: bool, local_cache_hits: bool):
     "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
 )
 def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):
-    """
-    Test whether scheduler with KVConnector is able to handle
+    """Test whether scheduler with KVConnector is able to handle
     unable to allocate (run out of blocks in allocate_slots().
     """
-
     # Setup Scheduler With Mock External Cache Hit.
     BLOCK_SIZE = 4
     NUM_BLOCKS = 10
@@ -2519,11 +2832,9 @@ def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):
 def test_kv_connector_handles_preemption(
     is_async, use_ec_connector, ec_role, use_v2_model_runner
 ):
-    """
-    Test whether scheduler with KVConnector is able to handle
+    """Test whether scheduler with KVConnector is able to handle
     unable to allocate (run out of blocks in allocate_slots().
     """
-
     # Setup Scheduler With Mock External Cache Hit.
     BLOCK_SIZE = 2
     # NOTE: there is 1 null block, so this is 6 blocks.
@@ -2734,7 +3045,6 @@ def assert_scheduler_empty(scheduler: Scheduler):
 
 def test_memory_leak():
     """Test that we do not have a memory leak."""
-
     scheduler = create_scheduler(enable_prefix_caching=True)
 
     NUM_REQUESTS = 5
@@ -2791,6 +3101,7 @@ def create_scheduler_with_priority(
 
     Returns:
       {class}`Scheduler` instance with priority scheduling
+
     """
     model_config = ModelConfig(
         model=model,
@@ -3566,7 +3877,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     )
     request.structured_output_request = Mock()
     request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
-    request.structured_output_request.grammar.accept_tokens.return_value = False
     request.status = RequestStatus.RUNNING
     request.num_computed_tokens = request.num_tokens
 
@@ -3574,10 +3884,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.connector = None
     scheduler.ec_connector = None
     scheduler.structured_output_manager = Mock()
-    scheduler.structured_output_manager.should_advance.return_value = True
-    scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = (
-        lambda request, new_token_ids: new_token_ids
-    )
+    scheduler.structured_output_manager.accept_tokens.return_value = False
     scheduler.requests = {request.request_id: request}
     scheduler.running = [request]
     scheduler.waiting = Mock()
@@ -3587,10 +3894,10 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.kv_event_publisher = Mock()
     scheduler.finished_req_ids = set()
     scheduler.finished_req_ids_dict = None
+    scheduler.aux_output_connector = None
     scheduler.grammar_compile_error_reqs = set()
+    scheduler.encoder_cache_mismatch_reqs = set()
     scheduler.vllm_config = Mock()
-    scheduler.vllm_config.model_config.enable_return_routed_experts = False
-    scheduler.enable_return_routed_experts = False
     scheduler.return_sampling_mask = False
     scheduler.recompute_kv_load_failures = False
     scheduler.defer_block_free = False
@@ -3626,8 +3933,8 @@ def free_request(req: Request, delay_free_blocks: bool = False):
     )
     engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
 
-    request.structured_output_request.grammar.accept_tokens.assert_called_once_with(
-        request.request_id, [123]
+    scheduler.structured_output_manager.accept_tokens.assert_called_once_with(
+        request, [123]
     )
     assert request.resumable is False
     assert request.status == RequestStatus.FINISHED_ERROR
@@ -3883,7 +4190,7 @@ def _assert_right_ec_connector_metadata(
     output: SchedulerOutput,
     mm_features_list: list[MultiModalFeatureSpec],
 ):
-    """Verify that ECConnector metadata EXACTLY matches the input MM data"""
+    """Verify that ECConnector metadata EXACTLY matches the input MM data."""
     # Get the connector metadata
     metadata = output.ec_connector_metadata
 
@@ -3914,7 +4221,6 @@ def _assert_right_encoder_inputs(
     """Verify that requests/mm_hashes should (not) in scheduled encoder input
     If check_exist is False, this function returns True
     if requests are NOT in encoder inputs"""
-
     # Get the scheduled encoder inputs
     # NOTE: scheduled_encoder_inputs is a dictionary with request id as key
     scheduled_encoder_inputs = output.scheduled_encoder_inputs
@@ -4423,11 +4729,9 @@ def test_ec_connector_schedule_multiple_requests(cache_exist, use_kv_connector):
 
 @pytest.mark.parametrize("use_kv_connector", [False, True])
 def test_ec_connector_unable_to_allocate(use_kv_connector):
-    """
-    Test whether scheduler with ECConnector is able to handle
+    """Test whether scheduler with ECConnector is able to handle
     unable to allocate (run out of blocks).
     """
-
     # Setup Scheduler With Mock External Cache Hit.
     BLOCK_SIZE = 4
     NUM_BLOCKS = 10
@@ -4756,8 +5060,7 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(
 
 @pytest.mark.parametrize("use_kv_connector", [False, True])
 def test_ec_connector_allocate_encoder_tokens_with_external_load(use_kv_connector):
-    """
-    Scenario:
+    """Scenario:
       - Encoder cache size: 32
       - Request A: 1 feature (12 tokens) → NOT cached remotely.
       - Request B: 3 features (3 x 10 tokens) → ALL cached remotely.
@@ -5192,6 +5495,87 @@ def update_connector_output(
         assert final_stats == expected_data
 
 
+def test_scheduler_ec_connector_stats():
+    """Test worker-side, scheduler-side, and combined EC connector stats."""
+
+    class GenericECConnectorStats(ECConnectorStats):
+        def reset(self):
+            self.data = {}
+
+        def aggregate(self, other: ECConnectorStats) -> ECConnectorStats:
+            self.data.update(other.data)
+            return self
+
+        def reduce(self) -> dict[str, int | float]:
+            return {}
+
+        def is_empty(self) -> bool:
+            return not self.data
+
+    test_cases = (
+        ({"worker": 1}, None, {"worker": 1}),
+        (None, {"scheduler": 2}, {"scheduler": 2}),
+        ({"worker": 1}, {"scheduler": 2}, {"worker": 1, "scheduler": 2}),
+    )
+
+    for worker_data, scheduler_data, expected_data in test_cases:
+        scheduler = create_scheduler()
+        worker_stats = (
+            GenericECConnectorStats(data=worker_data) if worker_data else None
+        )
+        scheduler_stats = (
+            GenericECConnectorStats(data=scheduler_data) if scheduler_data else None
+        )
+        scheduler.ec_connector = Mock()
+        scheduler.ec_connector.take_unavailable_requests.return_value = set()
+        scheduler.ec_connector.get_ec_connector_stats.return_value = (
+            scheduler_stats if worker_stats is None else None
+        )
+
+        def update_connector_output(
+            ec_connector_output: ECConnectorOutput,
+            scheduler=scheduler,
+            scheduler_stats=scheduler_stats,
+        ):
+            scheduler.ec_connector.get_ec_connector_stats.return_value = scheduler_stats
+
+        scheduler.ec_connector.update_connector_output.side_effect = (
+            update_connector_output
+        )
+
+        model_output = ModelRunnerOutput(
+            req_ids=["req_0"],
+            req_id_to_index={"req_0": 0},
+            sampled_token_ids=[[123]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[None],
+            ec_connector_output=ECConnectorOutput(ec_connector_stats=worker_stats)
+            if worker_stats
+            else None,
+        )
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=[],
+            scheduled_cached_reqs=None,
+            num_scheduled_tokens={"req_0": 1},
+            total_num_scheduled_tokens=1,
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=[0],
+            finished_req_ids=set(),
+            free_encoder_mm_hashes=[],
+        )
+
+        engine_core_outputs = scheduler.update_from_output(
+            scheduler_output, model_output
+        )
+
+        final_stats = next(
+            iter(engine_core_outputs.values())
+        ).scheduler_stats.ec_connector_stats
+        assert final_stats == expected_data
+
+
 def test_ec_connector_update_connector_output_called():
     """Test that worker-side EC connector output is forwarded to the
     EC connector's update_connector_output hook."""
@@ -5706,6 +6090,201 @@ def test_unavailable_encoder_input_fails_the_request_as_retryable():
     assert [o.finish_reason for o in engine_outputs] == [FinishReason.ERROR]
 
 
+def _create_mm_request_with_embedding_positions(
+    request_id: str, identifier: str, is_embed: list[bool]
+) -> Request:
+    return create_requests(
+        num_requests=1,
+        num_tokens=len(is_embed),
+        mm_hashes_list=[[identifier]],
+        mm_positions=[
+            [
+                PlaceholderRange(
+                    offset=0,
+                    length=len(is_embed),
+                    is_embed=torch.tensor(is_embed),
+                )
+            ]
+        ],
+        req_ids=[request_id],
+    )[0]
+
+
+def test_encoder_cache_accepts_matching_embed_count():
+    """An actively referenced entry is reused when embedding counts match."""
+    scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
+    identifier = "reused-client-uuid"
+    owner = _create_mm_request_with_embedding_positions(
+        "owner", identifier, [True, True, True, False, False, False]
+    )
+    compatible = _create_mm_request_with_embedding_positions(
+        "compatible", identifier, [True, False, True, False, True, False, False]
+    )
+    scheduler.encoder_cache_manager.allocate(owner, 0)
+    scheduler.add_request(compatible)
+
+    scheduler_output = scheduler.schedule()
+
+    assert compatible.request_id in scheduler_output.num_scheduled_tokens
+    assert compatible.request_id not in scheduler.encoder_cache_mismatch_reqs
+    assert compatible.request_id not in scheduler_output.scheduled_encoder_inputs
+
+
+def test_encoder_cache_rejects_mismatched_embed_count(caplog_vllm):
+    """A completed request's freeable entry rejects an incompatible reuse."""
+    scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
+    identifier = "reused-client-uuid"
+    owner = _create_mm_request_with_embedding_positions(
+        "owner", identifier, [True, True, True, False, False, False]
+    )
+    mismatched = _create_mm_request_with_embedding_positions(
+        "mismatched", identifier, [True, True, True, True, False, False]
+    )
+    scheduler.encoder_cache_manager.allocate(owner, 0)
+    scheduler.encoder_cache_manager.free(owner)
+    freeable_slots_before = scheduler.encoder_cache_manager.num_freeable_slots
+    scheduler.add_request(mismatched)
+
+    scheduler_output = scheduler.schedule()
+    assert not scheduler_output.num_scheduled_tokens
+    assert mismatched.request_id in scheduler.encoder_cache_mismatch_reqs
+    assert (
+        "multimodal identifier 'reused-client-uuid' holds 3 embeddings"
+        in caplog_vllm.text
+        and "expects 4" in caplog_vllm.text
+    )
+
+    outputs = scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(req_ids=[], req_id_to_index={}),
+    )
+
+    assert mismatched.status == RequestStatus.FINISHED_ERROR
+    assert not scheduler.encoder_cache_mismatch_reqs
+    assert mismatched.request_id not in scheduler.requests
+    assert scheduler.encoder_cache_manager.cached[identifier] == set()
+    assert (
+        mismatched.request_id not in scheduler.encoder_cache_manager.request_cached_ids
+    )
+    assert scheduler.encoder_cache_manager.freeable[identifier] == 3
+    assert scheduler.encoder_cache_manager.num_freeable_slots == freeable_slots_before
+    output = outputs[mismatched.client_index].outputs[0]
+    assert output.request_id == mismatched.request_id
+    assert output.finish_reason == FinishReason.ERROR
+
+    healthy = create_requests(num_requests=1, req_ids=["healthy"], num_tokens=6)[0]
+    scheduler.add_request(healthy)
+    next_output = scheduler.schedule()
+    assert healthy.request_id in next_output.num_scheduled_tokens
+
+
+def test_encoder_cache_embed_count_mismatch_restores_prior_hit_state():
+    """A prior hit is released when a later mismatch fails the request."""
+    scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
+    manager = scheduler.encoder_cache_manager
+    valid_identifier = "valid-reuse"
+    mismatch_identifier = "mismatched-reuse"
+    valid_owner = _create_mm_request_with_embedding_positions(
+        "valid-owner", valid_identifier, [True, True, False]
+    )
+    mismatch_owner = _create_mm_request_with_embedding_positions(
+        "mismatch-owner", mismatch_identifier, [True, True, True, False]
+    )
+    manager.allocate(valid_owner, 0)
+    manager.free(valid_owner)
+    manager.allocate(mismatch_owner, 0)
+    manager.free(mismatch_owner)
+    freeable_slots_before = manager.num_freeable_slots
+
+    request = create_requests(
+        num_requests=1,
+        num_tokens=7,
+        mm_hashes_list=[[valid_identifier, mismatch_identifier]],
+        mm_positions=[
+            [
+                PlaceholderRange(
+                    offset=0,
+                    length=3,
+                    is_embed=torch.tensor([True, True, False]),
+                ),
+                PlaceholderRange(
+                    offset=3,
+                    length=4,
+                    is_embed=torch.tensor([True, True, True, True]),
+                ),
+            ]
+        ],
+        req_ids=["partially-mutated"],
+    )[0]
+    scheduler.add_request(request)
+
+    scheduler_output = scheduler.schedule()
+    assert not scheduler_output.num_scheduled_tokens
+    assert request.request_id in scheduler.encoder_cache_mismatch_reqs
+    assert request.request_id in manager.cached[valid_identifier]
+    assert valid_identifier not in manager.freeable
+    assert manager.num_freeable_slots == freeable_slots_before - 2
+
+    outputs = scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(req_ids=[], req_id_to_index={}),
+    )
+
+    assert request.status == RequestStatus.FINISHED_ERROR
+    assert request.request_id not in scheduler.requests
+    assert request.request_id not in manager.request_cached_ids
+    assert manager.cached[valid_identifier] == set()
+    assert manager.freeable[valid_identifier] == 2
+    assert manager.freeable[mismatch_identifier] == 3
+    assert manager.num_freeable_slots == freeable_slots_before
+    output = outputs[request.client_index].outputs[0]
+    assert output.finish_reason == FinishReason.ERROR
+
+
+def test_encoder_cache_rejects_mismatched_embed_count_within_request():
+    """One request cannot reuse an identifier with a different embed count."""
+    scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
+    identifier = "reused-within-request"
+    request = create_requests(
+        num_requests=1,
+        num_tokens=8,
+        mm_hashes_list=[[identifier, identifier]],
+        mm_positions=[
+            [
+                PlaceholderRange(
+                    offset=0,
+                    length=4,
+                    is_embed=torch.tensor([True, True, False, False]),
+                ),
+                PlaceholderRange(
+                    offset=4,
+                    length=4,
+                    is_embed=torch.tensor([True, True, True, False]),
+                ),
+            ]
+        ],
+        req_ids=["mismatched-within-request"],
+    )[0]
+    scheduler.add_request(request)
+
+    scheduler_output = scheduler.schedule()
+    assert not scheduler_output.num_scheduled_tokens
+    assert request.request_id in scheduler.encoder_cache_mismatch_reqs
+    assert identifier not in scheduler.encoder_cache_manager.cached
+
+    outputs = scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(req_ids=[], req_id_to_index={}),
+    )
+
+    assert request.status == RequestStatus.FINISHED_ERROR
+    assert not scheduler.encoder_cache_mismatch_reqs
+    assert request.request_id not in scheduler.requests
+    output = outputs[request.client_index].outputs[0]
+    assert output.request_id == request.request_id
+    assert output.finish_reason == FinishReason.ERROR
+
+
 def test_free_encoder_inputs_defers_for_eagle_lookahead():
     """With EAGLE speculative decoding, the encoder input is retained one extra
     position so the drafter's +1 look-ahead mm-embedding gather (which reads one
@@ -6465,16 +7044,15 @@ def _decode_ready_request(scheduler):
     return request
 
 
-def test_update_draft_token_ids_strips_ngram_padding(monkeypatch):
-    """ngram_gpu pads unfilled draft slots with -1; update_draft_token_ids must
+def test_update_draft_token_ids_strips_ngram_padding():
+    """ngram_gpu pads unfilled draft slots with -1; manager validation must
     strip them before grammar.validate_tokens (which otherwise raises)."""
     scheduler = create_scheduler(num_speculative_tokens=4)
     request = _decode_ready_request(scheduler)
 
     grammar = _RecordingGrammar()
-    request.structured_output_request = SimpleNamespace(grammar=grammar)
-    monkeypatch.setattr(
-        scheduler.structured_output_manager, "should_advance", lambda req: True
+    request.structured_output_request = SimpleNamespace(
+        grammar=grammar, reasoning_ended=True
     )
 
     scheduler.update_draft_token_ids(
@@ -6485,16 +7063,15 @@ def test_update_draft_token_ids_strips_ngram_padding(monkeypatch):
     assert request.spec_token_ids == [10, 11]
 
 
-def test_update_draft_token_ids_in_output_strips_padding(monkeypatch):
+def test_update_draft_token_ids_in_output_strips_padding():
     """Same guard on the output path; the -1 pad-back for the rejected count
-    is preserved (only the input to validate_tokens is stripped)."""
+    is preserved (only the input to manager validation is stripped)."""
     scheduler = create_scheduler(num_speculative_tokens=4)
     request = _decode_ready_request(scheduler)
 
     grammar = _RecordingGrammar()
-    request.structured_output_request = SimpleNamespace(grammar=grammar)
-    monkeypatch.setattr(
-        scheduler.structured_output_manager, "should_advance", lambda req: True
+    request.structured_output_request = SimpleNamespace(
+        grammar=grammar, reasoning_ended=True
     )
 
     scheduler_output = SimpleNamespace(
@@ -6515,3 +7092,141 @@ def test_update_draft_token_ids_in_output_strips_padding(monkeypatch):
         -1,
     ]
     assert scheduler_output.num_invalid_spec_tokens == {request.request_id: 2}
+
+
+def test_diffusion_canvas_width_defaults_to_the_served_canvas():
+    def req(extra):
+        return SimpleNamespace(sampling_params=SimpleNamespace(extra_args=extra))
+
+    assert diffusion_canvas_width(req(None), 64) == 64
+    assert diffusion_canvas_width(req({}), 64) == 64
+    assert diffusion_canvas_width(req({"diffusion_canvas_length": 16}), 64) == 16
+    assert diffusion_canvas_width(SimpleNamespace(sampling_params=None), 64) == 64
+
+
+def _diffusion_request(req_id: str, extra_args: dict) -> Request:
+    (request,) = create_requests(
+        num_requests=1, num_tokens=8, max_tokens=64, req_ids=[req_id]
+    )
+    request.sampling_params.extra_args = extra_args
+    return request
+
+
+@pytest.fixture
+def diffusion_model_runner(monkeypatch):
+    # These CPU tests only exercise scheduling, not Triton kernels.
+    monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True)
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+
+
+def _diffusion_scheduler(**kwargs) -> DiffusionAsyncScheduler:
+    scheduler = create_scheduler(
+        async_scheduling=True,
+        diffusion_canvas_length=8,
+        scheduler_cls=DiffusionAsyncScheduler,
+        **kwargs,
+    )
+    assert isinstance(scheduler, DiffusionAsyncScheduler)
+    return scheduler
+
+
+@pytest.mark.parametrize("async_scheduling", [False, True])
+@pytest.mark.usefixtures("diffusion_model_runner")
+def test_diffusion_scheduler_is_selected_by_default(async_scheduling):
+    config = create_scheduler(
+        async_scheduling=async_scheduling, diffusion_canvas_length=8
+    ).vllm_config.scheduler_config
+    assert config.get_scheduler_cls() is (
+        DiffusionAsyncScheduler if config.async_scheduling else Scheduler
+    )
+
+
+@pytest.mark.usefixtures("diffusion_model_runner")
+def test_diffusion_scheduler_narrows_the_canvas_per_request():
+    scheduler = _diffusion_scheduler()
+    wide = _diffusion_request("wide", {})
+    narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4})
+    scheduler.add_request(wide)
+    scheduler.add_request(narrow)
+
+    # The prefill step lays down the served canvas as placeholders.
+    scheduler.schedule()
+    output = scheduler.schedule()
+    assert output.scheduled_spec_decode_tokens["wide"] == [-1] * 8
+    assert output.scheduled_spec_decode_tokens["narrow"] == [-1] * 4
+    assert output.num_scheduled_tokens["narrow"] == 4
+    assert narrow.num_output_placeholders == 4
+
+
+@pytest.mark.parametrize("structured", [False, True])
+@pytest.mark.usefixtures("diffusion_model_runner")
+def test_diffusion_scheduler_trims_full_width_worker_drafts(structured):
+    """Padded worker drafts must be narrowed before scheduling or grammar validation."""
+    scheduler = _diffusion_scheduler()
+    wide = _diffusion_request("wide", {})
+    narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4})
+    for request in (wide, narrow):
+        scheduler.add_request(request)
+    prefill = scheduler.schedule()
+    _model_output(scheduler, prefill, [[], []])
+
+    if structured:
+        for request in (wide, narrow):
+            request.structured_output_request = SimpleNamespace(
+                grammar=_RecordingGrammar(), reasoning_ended=True
+            )
+    tokens = list(range(8)) if structured else [-1] * 8
+    drafts = DraftTokenIds(["wide", "narrow"], [tokens.copy(), tokens.copy()])
+    output = scheduler.schedule()
+    scheduler.update_draft_token_ids_in_output(drafts, output)
+
+    assert output.scheduled_spec_decode_tokens == {
+        "wide": tokens,
+        "narrow": tokens[:4],
+    }
+    assert output.num_scheduled_tokens == {"wide": 8, "narrow": 4}
+    if structured:
+        assert wide.structured_output_request.grammar.seen == [tokens]
+        assert narrow.structured_output_request.grammar.seen == [tokens[:4]]
+
+
+@pytest.mark.usefixtures("diffusion_model_runner")
+def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight():
+    scheduler = _diffusion_scheduler()
+    one = _diffusion_request(
+        "one", {"diffusion_read_only": True, "diffusion_max_steps": 1}
+    )
+    two = _diffusion_request(
+        "two", {"diffusion_read_only": True, "diffusion_max_steps": 2}
+    )
+    # Only reads are deferred. A capped generation keeps its extra async step.
+    gen = _diffusion_request("gen", {"diffusion_max_steps": 1})
+    for request in (one, two, gen):
+        scheduler.add_request(request)
+
+    scheduler.schedule()
+    assert set(scheduler.schedule().num_scheduled_tokens) == {"one", "two", "gen"}
+    assert one.num_output_placeholders == 8
+    assert set(scheduler.schedule().num_scheduled_tokens) == {"two", "gen"}
+    assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"}
+    # The deferral is renewed every call.
+    assert set(scheduler.schedule().num_scheduled_tokens) == {"gen"}
+
+
+@pytest.mark.usefixtures("diffusion_model_runner")
+def test_diffusion_read_deferral_keeps_a_longer_pp_wait():
+    scheduler = _diffusion_scheduler(pipeline_parallel_size=3, use_v2_model_runner=True)
+    read = _diffusion_request(
+        "read", {"diffusion_read_only": True, "diffusion_max_steps": 1}
+    )
+    scheduler.add_request(read)
+
+    scheduler.schedule()
+    assert read.next_decode_eligible_step == 4
+    for _ in range(2):
+        assert "read" not in scheduler.schedule().num_scheduled_tokens
+    assert "read" in scheduler.schedule().num_scheduled_tokens
+    assert read.next_decode_eligible_step == 7
+    # Deferring this step alone would ask for 6. The PP wait to 7 stands.
+    assert "read" not in scheduler.schedule().num_scheduled_tokens
+    assert read.next_decode_eligible_step == 7
diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py
index 39113723591b..f20a96f5cc31 100644
--- a/tests/v1/core/test_single_type_kv_cache_manager.py
+++ b/tests/v1/core/test_single_type_kv_cache_manager.py
@@ -103,6 +103,139 @@ def test_mamba_speculative_block_relocation_requires_exclusive_ownership():
         manager._relocate_speculative_block([pinned_block], 0)
 
 
+def test_mamba_retirement_crosses_null_gaps():
+    spec = MambaSpec(
+        block_size=4,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+    pool = BlockPool(num_gpu_blocks=8, enable_caching=False, hash_block_size=4)
+    manager = MambaManager(
+        spec,
+        block_pool=pool,
+        enable_caching=False,
+        kv_cache_group_id=0,
+        scheduler_block_size=4,
+    )
+    old, committed, in_flight = pool.get_new_blocks(3)
+    manager.req_to_blocks["r"] = [old, pool.null_block, committed, in_flight]
+
+    # The state at token 12 and the following in-flight state must survive.
+    for _ in range(2):
+        manager.remove_skipped_blocks("r", processed_computed_tokens=12)
+        assert manager.req_to_blocks["r"] == [
+            pool.null_block,
+            pool.null_block,
+            committed,
+            in_flight,
+        ]
+        assert old.ref_cnt == 0
+        assert committed.ref_cnt == in_flight.ref_cnt == 1
+        assert pool.get_num_free_blocks() == 5
+
+    manager.free("r")
+    manager.req_to_blocks["r"] = pool.get_new_blocks(2)
+    manager.remove_skipped_blocks("r", processed_computed_tokens=5)
+    assert manager.req_to_blocks["r"][0].is_null
+    assert manager.req_to_blocks["r"][1].ref_cnt == 1
+
+
+@pytest.mark.parametrize("block_size", [3584, 4608])
+@pytest.mark.parametrize("in_flight_chunks", [0, 1, 2])
+def test_mamba_retirement_bounds_prefill_states(block_size, in_flight_chunks):
+    spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+        num_speculative_blocks=5,
+    )
+    pool = BlockPool(num_gpu_blocks=1000, enable_caching=True, hash_block_size=256)
+    manager = MambaManager(
+        spec,
+        block_pool=pool,
+        enable_caching=True,
+        kv_cache_group_id=0,
+        scheduler_block_size=block_size,
+    )
+    initial_free = pool.get_num_free_blocks()
+    chunk = 8192 // block_size * block_size
+    peak_held = 0
+    # Replay full prefill chunks with an async committed-token lag.
+    for target in range(chunk, 480031 + 1, chunk):
+        computed = target - chunk
+        committed = max(0, computed - in_flight_chunks * chunk)
+        manager.remove_skipped_blocks("r", committed)
+        manager.get_num_blocks_to_allocate(
+            "r", target + 5, [], computed, computed, target
+        )
+        manager.allocate_new_blocks("r", target + 5, target)
+        peak_held = max(peak_held, initial_free - pool.get_num_free_blocks())
+
+    # Five speculative blocks, current/previous states, and bounded in-flight states.
+    assert peak_held == 7 + in_flight_chunks
+    manager.free("r")
+    assert pool.get_num_free_blocks() == initial_free
+
+
+@pytest.mark.parametrize("block_size", [896, 1536])
+@pytest.mark.parametrize("num_speculative_blocks", [0, 1, 4])
+@pytest.mark.parametrize("prompt_tokens", [25121, 704547])
+def test_mamba_checkpoint_admission_matches_allocation(
+    block_size, num_speculative_blocks, prompt_tokens
+):
+    """Checkpoint admission must match the subsequent physical allocation."""
+    spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+        num_speculative_blocks=num_speculative_blocks,
+        num_prefill_checkpoint_blocks=1,
+        prefill_checkpoint_alignment=64,
+    )
+    pool = BlockPool(
+        num_gpu_blocks=2048,
+        enable_caching=True,
+        hash_block_size=128,
+    )
+    manager = MambaManager(
+        spec,
+        block_pool=pool,
+        enable_caching=True,
+        kv_cache_group_id=0,
+        scheduler_block_size=block_size,
+    )
+    request_id = "prefill"
+    computed_tokens = 23040
+
+    def estimate(num_tokens, total_computed_tokens, apply_admission_cap):
+        return manager.get_num_blocks_to_allocate(
+            request_id=request_id,
+            num_tokens=num_tokens,
+            new_computed_blocks=[],
+            total_computed_tokens=total_computed_tokens,
+            num_local_computed_tokens=total_computed_tokens,
+            num_tokens_main_model=num_tokens,
+            apply_admission_cap=apply_admission_cap,
+        )
+
+    estimate(computed_tokens, 0, False)
+    manager.allocate_new_blocks(request_id, computed_tokens, computed_tokens)
+    assert request_id in manager._allocated_block_reqs
+
+    admission_estimate = estimate(prompt_tokens, computed_tokens, True)
+    allocation_estimate = estimate(prompt_tokens, computed_tokens, False)
+    assert request_id in manager._checkpoints
+
+    free_before = pool.get_num_free_blocks()
+    manager.allocate_new_blocks(request_id, prompt_tokens, prompt_tokens)
+    allocated = free_before - pool.get_num_free_blocks()
+
+    assert admission_estimate == allocation_estimate == allocated
+
+
 def get_sliding_window_manager(
     sliding_window_spec,
     block_pool,
diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py
index d3090e2ddace..46ca29f67b9f 100644
--- a/tests/v1/core/utils.py
+++ b/tests/v1/core/utils.py
@@ -7,6 +7,7 @@
 from tests.v1.kv_connector.unit.utils import MockKVConfig
 from vllm.config import (
     CacheConfig,
+    DiffusionConfig,
     ECTransferConfig,
     KVTransferConfig,
     ModelConfig,
@@ -17,6 +18,7 @@
     SpeculativeConfig,
     VllmConfig,
 )
+from vllm.config.scheduler import SchedulerPolicy
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalKwargsItem,
@@ -51,10 +53,12 @@ def mock_kv(matched_tokens: int, is_async: bool, num_defers_before_matching: int
 def create_scheduler(
     model: str = "facebook/opt-125m",
     max_num_seqs: int = 16,
+    max_num_active_seqs: int | None = None,
     max_num_batched_tokens: int = 8192,
     enable_chunked_prefill: bool = True,
     enable_prefix_caching: bool = False,
     long_prefill_token_threshold: int = 0,
+    long_prefill_token_threshold_adaptive: bool = False,
     disable_chunked_mm_input: bool = False,
     use_kv_connector: None | bool | str | MockKVConfig = None,
     kv_role: str = "kv_both",
@@ -74,6 +78,9 @@ def create_scheduler(
     use_v2_model_runner: bool | None = None,
     kv_cache_spec: KVCacheSpec | None = None,
     per_request_spec_decode_metrics: str = "none",
+    scheduling_policy: SchedulerPolicy = "fcfs",
+    diffusion_canvas_length: int | None = None,
+    scheduler_cls: type[Scheduler] | None = None,
 ) -> Scheduler | AsyncScheduler:
     """Create scheduler under test.
 
@@ -84,9 +91,13 @@ def create_scheduler(
       enable_prefix_caching: optionally force APC config
                              (True/False) or use default
                              (False)
+      long_prefill_token_threshold: cap on prefill chunk size
+      long_prefill_token_threshold_adaptive: floor the cap at a
+                             fair share of the token budget
 
     Returns:
       {class}`Scheduler` instance
+
     """
     model_config = ModelConfig(
         model=model,
@@ -104,15 +115,18 @@ def create_scheduler(
         max_model_len = max_num_batched_tokens
     scheduler_config = SchedulerConfig(
         max_num_seqs=max_num_seqs,
+        max_num_active_seqs=max_num_active_seqs,
         max_num_batched_tokens=max_num_batched_tokens,
         max_model_len=max_model_len,
         long_prefill_token_threshold=long_prefill_token_threshold,
+        long_prefill_token_threshold_adaptive=(long_prefill_token_threshold_adaptive),
         disable_chunked_mm_input=disable_chunked_mm_input,
         enable_chunked_prefill=enable_chunked_prefill,
         async_scheduling=async_scheduling,
         is_encoder_decoder=model_config.is_encoder_decoder,
         # Ensure admission/preemption mechanics are deterministic
         watermark=0.0,
+        policy=scheduling_policy,
     )
     # Cache config, optionally force APC
     cache_config = CacheConfig(
@@ -172,6 +186,12 @@ def create_scheduler(
         else None
     )
 
+    diffusion_config: DiffusionConfig | None = None
+    if diffusion_canvas_length is not None:
+        # A diffusion checkpoint declares its canvas in the HF config.
+        model_config.hf_config.canvas_length = diffusion_canvas_length
+        diffusion_config = DiffusionConfig(canvas_length=diffusion_canvas_length)
+
     vllm_config = VllmConfig(
         scheduler_config=scheduler_config,
         model_config=model_config,
@@ -182,6 +202,7 @@ def create_scheduler(
         ),
         kv_transfer_config=kv_transfer_config,
         speculative_config=speculative_config,
+        diffusion_config=diffusion_config,
         ec_transfer_config=ec_transfer_config,
         observability_config=ObservabilityConfig(
             per_request_spec_decode_metrics=per_request_spec_decode_metrics,
@@ -201,7 +222,8 @@ def create_scheduler(
     )
     cache_config.num_gpu_blocks = num_blocks
     register_all_kvcache_specs(vllm_config)
-    scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
+    if scheduler_cls is None:
+        scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
     scheduler = scheduler_cls(
         vllm_config=vllm_config,
         kv_cache_config=kv_cache_config,
diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py
index dc9c1ffb3067..b0b49a9567b1 100644
--- a/tests/v1/cudagraph/test_breakable_cudagraph.py
+++ b/tests/v1/cudagraph/test_breakable_cudagraph.py
@@ -1,14 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for the breakable cudagraph primitives.
-"""
+"""Unit tests for the breakable cudagraph primitives."""
 
 from __future__ import annotations
 
 import threading
 from contextlib import nullcontext
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -43,6 +41,7 @@ def test_piecewise_capture_builds_fresh_metadata_for_both_passes():
     manager._capture_descs = {CUDAGraphMode.PIECEWISE: [desc]}
     manager._graphs_captured = False
     manager.use_breakable_cg = True
+    manager.ubatch_runner = None
 
     create_calls = []
     forward_calls = []
@@ -333,6 +332,57 @@ def attention_like(t: torch.Tensor) -> None:
     assert torch.equal(x, torch.full((4,), 15.0, device="cuda"))
 
 
+@pytest.mark.parametrize("keyword", [False, True])
+def test_quantized_activation_replay_does_not_retain_tensors(
+    cuda_capture_stream, keyword
+):
+    """Release the producer tensors while preserving both buffers on replay."""
+    import gc
+    import weakref
+
+    from vllm.compilation.breakable_cudagraph import (
+        BreakableCUDAGraphCapture,
+        eager_break_during_capture,
+    )
+    from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
+
+    @eager_break_during_capture
+    def consume(x, data_out, scale_out):
+        data_out.copy_(x.data)
+        scale_out.copy_(x.scale)
+
+    source = torch.ones((128, 128), device="cuda", dtype=torch.bfloat16)
+    data_out = source.to(torch.float8_e4m3fn)
+    scale_out = source[:, ::32].to(torch.uint8)
+    output = torch.empty_like(source)
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        data = source.to(torch.float8_e4m3fn)
+        scale = source[:, ::32].to(torch.uint8)
+        activation = QuantizedActivation(
+            data, scale, source.dtype, source.shape, kMxfp8Dynamic
+        )
+        data_ref, scale_ref = weakref.ref(data), weakref.ref(scale)
+        if keyword:
+            consume(x=activation, data_out=data_out, scale_out=scale_out)
+        else:
+            consume(activation, data_out, scale_out)
+        output.copy_(data_out)
+
+    del data, scale, activation
+    gc.collect()
+    assert data_ref() is None
+    assert scale_ref() is None
+
+    for value in (1, 2, 3):
+        source.fill_(value)
+        cap.replay()
+        cuda_capture_stream.synchronize()
+        torch.testing.assert_close(output, source, rtol=0, atol=0)
+        torch.testing.assert_close(scale_out, source[:, ::32].to(torch.uint8))
+
+
 def test_eager_attention_inside_multistream_overlap(cuda_capture_stream):
     """Handle an eager attention break inside a multi-stream overlap region."""
     from vllm.compilation.breakable_cudagraph import (
@@ -485,3 +535,133 @@ def outer_op(t: torch.Tensor) -> None:
     # 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
+
+
+# ---------------------------------------------------------------------------
+# BreakableCUDAGraphWrapper: runtime_mode dispatch gating
+# ---------------------------------------------------------------------------
+
+
+def _mock_vllm_config():
+    config = MagicMock()
+    config.parallel_config.data_parallel_size = 1
+    config.parallel_config.use_sequence_parallel_moe = False
+    config.parallel_config.is_moe_model = False
+    return config
+
+
+def test_wrapper_falls_through_on_runtime_mode_mismatch():
+    """A PIECEWISE-scoped wrapper runs other modes eagerly, capturing
+    nothing, so an outer CUDAGraphWrapper can own them."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+    from vllm.config import CUDAGraphMode
+    from vllm.forward_context import BatchDescriptor, set_forward_context
+
+    calls = []
+
+    def model(x):
+        calls.append(x)
+        return x + 1
+
+    config = _mock_vllm_config()
+    wrapper = BreakableCUDAGraphWrapper(
+        model, config, runtime_mode=CUDAGraphMode.PIECEWISE
+    )
+    desc = BatchDescriptor(num_tokens=4)
+
+    with set_forward_context(
+        None,
+        config,
+        cudagraph_runtime_mode=CUDAGraphMode.FULL,
+        batch_descriptor=desc,
+    ):
+        assert wrapper(1) == 2
+
+    assert calls == [1]
+    assert not wrapper.entries
+
+
+def test_wrapper_captures_on_runtime_mode_match():
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+    from vllm.config import CUDAGraphMode
+    from vllm.forward_context import BatchDescriptor, set_forward_context
+
+    calls = []
+
+    def model(x):
+        calls.append(x)
+        return x + 1
+
+    config = _mock_vllm_config()
+    wrapper = BreakableCUDAGraphWrapper(
+        model, config, runtime_mode=CUDAGraphMode.PIECEWISE
+    )
+    desc = BatchDescriptor(num_tokens=4)
+
+    sentinel = object()
+    with (
+        patch.object(
+            BreakableCUDAGraphWrapper, "_capture", return_value=sentinel
+        ) as mock_capture,
+        set_forward_context(
+            None,
+            config,
+            cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+            batch_descriptor=desc,
+        ),
+    ):
+        assert wrapper(1) is sentinel
+
+    mock_capture.assert_called_once()
+    assert calls == []
+    assert desc in wrapper.entries
+
+
+def test_wrapper_default_matches_any_non_none_mode():
+    """runtime_mode=None matches any non-NONE mode (MRV2 relies on this)."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+    from vllm.config import CUDAGraphMode
+    from vllm.forward_context import BatchDescriptor, set_forward_context
+
+    def model(x):
+        return x + 1
+
+    config = _mock_vllm_config()
+    wrapper = BreakableCUDAGraphWrapper(model, config)
+    desc = BatchDescriptor(num_tokens=4)
+
+    sentinel = object()
+    with (
+        patch.object(
+            BreakableCUDAGraphWrapper, "_capture", return_value=sentinel
+        ) as mock_capture,
+        set_forward_context(
+            None,
+            config,
+            cudagraph_runtime_mode=CUDAGraphMode.FULL,
+            batch_descriptor=desc,
+        ),
+    ):
+        assert wrapper(1) is sentinel
+
+    mock_capture.assert_called_once()
+    assert desc in wrapper.entries
+
+
+def test_unwrap_recurses_through_nested_wrappers():
+    """unwrap() reaches the raw model through nested wrappers."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+    from vllm.compilation.cuda_graph import CUDAGraphWrapper
+    from vllm.config import CUDAGraphMode
+
+    def model(x):
+        return x + 1
+
+    config = _mock_vllm_config()
+    inner = BreakableCUDAGraphWrapper(
+        model, config, runtime_mode=CUDAGraphMode.PIECEWISE
+    )
+    outer = CUDAGraphWrapper(inner, config, runtime_mode=CUDAGraphMode.FULL)
+
+    assert inner.unwrap() is model
+    assert outer.unwrap() is model
diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py
index c10835821f58..ddf72c3b1141 100644
--- a/tests/v1/cudagraph/test_cudagraph_dispatch.py
+++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py
@@ -376,7 +376,6 @@ def _run_and_monitor_call(
     wrapper, input_tensor, runtime_mode, batch_descriptor, vllm_config
 ):
     """Helper to run a single call and monitor the action."""
-
     with (
         patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_graph_context,
         patch.object(wrapper, "runnable", wraps=wrapper.runnable) as mock_runnable,
diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py
index c907be36b88f..76c473215807 100644
--- a/tests/v1/cudagraph/test_cudagraph_manager.py
+++ b/tests/v1/cudagraph/test_cudagraph_manager.py
@@ -22,6 +22,9 @@
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.input_batch import InputBuffers
 from vllm.v1.worker.gpu.pcp_manager import PCPManager
+from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
+    SpeculatorCudaGraphManager,
+)
 
 pytestmark = pytest.mark.cpu_test
 
@@ -165,6 +168,89 @@ def test_piecewise_capture_uses_pcp_dummy_slot_mappings():
     block_tables.get_dummy_slot_mappings.assert_not_called()
 
 
+@pytest.mark.parametrize(
+    "full_cudagraph,max_query_len,expected_max_query_len",
+    [(False, None, None), (True, None, 8), (True, 1, 1), (True, 3, 3)],
+)
+def test_capture_query_length_covers_replay_batches(
+    full_cudagraph, max_query_len, expected_max_query_len
+):
+    """Mixed FULL graphs must capture prefill kernels as well as decode kernels."""
+    num_tokens = num_reqs = 8
+    buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))
+    block_tables = MagicMock()
+    block_tables.cp_size = 1
+    block_tables.get_dummy_block_tables.return_value = ()
+    block_tables.get_dummy_slot_mappings.return_value = torch.empty(
+        0, num_tokens, dtype=torch.int64
+    )
+    model_state = MagicMock()
+    kv_cache_config = KVCacheConfig(
+        num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
+    )
+
+    gpu_cudagraph_utils.prepare_inputs_to_capture(
+        num_reqs,
+        num_tokens,
+        model_state,
+        buffers,
+        block_tables,
+        [],
+        kv_cache_config,
+        full_cudagraph=full_cudagraph,
+        max_query_len=max_query_len,
+    )
+
+    batch = model_state.prepare_attn.call_args.args[0]
+    assert batch.num_scheduled_tokens.max() == 1
+    assert batch.max_query_len == expected_max_query_len
+
+
+@pytest.mark.parametrize(
+    "uniform_token_count,max_query_len,expected_max_query_len",
+    [(1, None, 1), (None, 3, 3), (None, None, 8)],
+)
+def test_speculator_capture_preserves_decode_query_bounds(
+    monkeypatch, uniform_token_count, max_query_len, expected_max_query_len
+):
+    """Draft decode graphs keep narrow queries while mixed graphs cover prefill."""
+    num_tokens = num_reqs = 8
+    desc = BatchExecutionDescriptor(
+        CUDAGraphMode.FULL,
+        num_tokens,
+        num_reqs,
+        uniform_token_count=uniform_token_count,
+        max_query_len=max_query_len,
+    )
+    manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager)
+    manager.max_num_reqs = num_reqs
+    manager.dp_size = 1
+    buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))
+    block_tables = MagicMock()
+    block_tables.cp_size = 1
+    block_tables.get_dummy_block_tables.return_value = ()
+    block_tables.get_dummy_slot_mappings.return_value = torch.empty(
+        0, num_tokens, dtype=torch.int64
+    )
+    model_state = MagicMock()
+    kv_cache_config = KVCacheConfig(
+        num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
+    )
+
+    def capture(self, create_forward_fn, progress_bar_desc):
+        for warmup in (True, False):
+            create_forward_fn(desc, warmup)(desc.cg_mode)
+
+    monkeypatch.setattr(gpu_cudagraph_utils.CudaGraphManager, "capture", capture)
+    manager.capture(
+        MagicMock(), model_state, buffers, block_tables, [], kv_cache_config
+    )
+
+    assert model_state.prepare_attn.call_count == 2
+    for call in model_state.prepare_attn.call_args_list:
+        assert call.args[0].max_query_len == expected_max_query_len
+
+
 _DECODE_QUERY_LEN = 3
 
 
diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py
index 096f0e40750f..1f145bdc8962 100644
--- a/tests/v1/cudagraph/test_encoder_cudagraph.py
+++ b/tests/v1/cudagraph/test_encoder_cudagraph.py
@@ -13,7 +13,9 @@
 """
 
 from collections.abc import Hashable
+from types import SimpleNamespace
 from typing import Any
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -575,6 +577,99 @@ def test_chunking_when_images_exceed_max_batch(self):
             assert out.shape == (4, _HIDDEN)
 
 
+# ---------------------------------------------------------------------------
+# E-only capture and output lifecycle
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm"
+)
+@pytest.mark.usefixtures("dist_init", "workspace_init")
+@pytest.mark.parametrize("profile_only", [False, True])
+@torch.inference_mode()
+def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only):
+    """The E-only entry captures only the encoder and preserves cached outputs."""
+    from vllm.distributed.ec_transfer.ec_connector.base import (
+        ECConnectorBase,
+        ECConnectorMetadata,
+    )
+    from vllm.v1.worker.gpu.ec_connector import ActiveECConnector
+    from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
+    from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+    from vllm.v1.worker.workspace import lock_workspace
+
+    device = torch.device("cuda:0")
+    dtype = torch.float16
+    model = SimpleMockViTModel().to(device).half()
+    manager = _make_manager_for_gpu(model, _BUDGETS, _MAX_BATCH, device, dtype)
+    encoder = object.__new__(EncoderRunner)
+    encoder.device = device
+    encoder.cudagraph_manager = manager
+    runner = object.__new__(MMEncoderModelRunner)
+    runner.model_state = SimpleNamespace(encoder_runner=encoder)
+    # No decoder manager is installed: capture must be encoder-only.
+    with patch(
+        "vllm.v1.worker.mm_encoder_model_runner.lock_workspace", wraps=lock_workspace
+    ) as lock:
+        runner.capture_model(profile_only=profile_only)
+        assert lock.call_count == int(not profile_only)
+    assert len(manager.budget_graphs["default"]) == len(_BUDGETS)
+
+    cache: dict[str, torch.Tensor] = {}
+    pending_sends: dict[str, torch.Tensor] = {}
+    connector = MagicMock(spec=ECConnectorBase)
+    connector.is_producer = True
+    connector.is_consumer = False
+    connector.get_finished.return_value = (None, None)
+    connector.save_caches.side_effect = lambda *, encoder_cache, mm_hash: (
+        pending_sends.update({mm_hash: encoder_cache[mm_hash]})
+    )
+    with patch(
+        "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector
+    ):
+        ec = ActiveECConnector(SimpleNamespace(), cache)
+    scheduled = SimpleNamespace(
+        ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset()
+    )
+    saved_outputs: list[tuple[torch.Tensor, torch.Tensor]] = []
+    # Mixed sizes exercise packing order; 64 is the boundary, 81 falls back.
+    for grids in ([[1, 8, 8], [1, 4, 4]], [[1, 16, 16]], [[1, 18, 18]], [[1, 4, 4]]):
+        inputs = _make_mm_kwargs(grids, device, dtype)
+        expected = model.encoder_eager_forward(inputs).split(
+            [t * (h // 2) * (w // 2) for t, h, w in grids]
+        )
+        with ec.maybe_get_output(scheduled) as ec_output:
+            outputs = manager.execute(inputs)
+            assert outputs is not None
+            cache[str(len(cache))] = outputs[0]
+        assert ec_output.finished_sending is None
+        assert outputs is not None
+        for actual, eager in zip(outputs, expected):
+            torch.testing.assert_close(actual, eager)
+        for previous, snapshot in saved_outputs:
+            torch.testing.assert_close(previous, snapshot, rtol=0, atol=0)
+        saved_outputs.extend((output, output.clone()) for output in outputs)
+    assert manager.graph_hits == 4
+    assert manager.graph_misses == 1
+    assert pending_sends.keys() == cache.keys()
+    connector.get_finished.return_value = (set(cache), None)
+    with ec.maybe_get_output(scheduled) as ec_output:
+        pass
+    assert ec_output.finished_sending == set(cache)
+
+
+def test_eonly_without_encoder_graph_skips_capture():
+    from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
+    from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+
+    encoder = object.__new__(EncoderRunner)
+    encoder.cudagraph_manager = None
+    runner = object.__new__(MMEncoderModelRunner)
+    runner.model_state = SimpleNamespace(encoder_runner=encoder)
+    assert runner.capture_model() == 0
+
+
 # ---------------------------------------------------------------------------
 # SimpleMockViTVideoModel — extends SimpleMockViTModel with video support
 # ---------------------------------------------------------------------------
diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py
index 4a86a5826bc0..9af55b9fbbc7 100644
--- a/tests/v1/determinism/test_batch_invariance.py
+++ b/tests/v1/determinism/test_batch_invariance.py
@@ -6,6 +6,7 @@
 
 import pytest
 import torch
+from transformers import AutoTokenizer
 from utils import (
     BACKENDS,
     TEST_MODEL,
@@ -31,8 +32,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
     backend,
     rms_norm_impl,
 ):
-    """
-    Ensures that the same request (the 'needle' prompt) yields identical output
+    """Ensures that the same request (the 'needle' prompt) yields identical output
     whether run alone (bs=1) or mixed into a larger batch (e.g., bs=64),
     using the high-level v1 LLM() API only (no manual batching).
 
@@ -51,6 +51,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
       to produce a more random-sounding phrase, yet remain deterministic by
       seed.
     - Keep max_tokens and max_model_len bounded for speed and memory use.
+
     """
     seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
     random.seed(seed)
@@ -203,13 +204,8 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
     # Use more realistic prompts for better token generation
     prompts = [_random_prompt(10, 50) for _ in range(32)]
 
-    # TODO: Update prompts to have ragged lengths in order to test chunked prefill
-    #       The above tests are not currently long enough to exercise chunking.
-    # prompts = (
-    #     [_random_prompt(10, 50) for _ in range(28)]
-    #     + [_random_prompt(256, 512) for _ in range(50)]
-    #     + [_random_prompt(2048, 4096) for _ in range(50)]
-    # )
+    # Chunked prefill is covered by
+    # test_logprobs_bitwise_batch_invariance_ragged_chunked_prefill below.
 
     sp = SamplingParams(
         temperature=0.6,
@@ -389,14 +385,87 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
         pytest.fail(msg)
 
 
+@skip_unsupported
+@pytest.mark.parametrize(
+    "backend",
+    BACKENDS,
+)
+def test_logprobs_bitwise_batch_invariance_ragged_chunked_prefill(backend):
+    """Batch invariance must hold when a prefill is split across chunks."""
+    random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
+
+    prompts = (
+        [_random_prompt(10, 50) for _ in range(8)]
+        + [_random_prompt(300, 450) for _ in range(4)]
+        + [_random_prompt(800, 1200) for _ in range(2)]
+    )
+    random.shuffle(prompts)
+
+    # _random_prompt takes a word target, not a token count, and the ratio
+    # depends on the tokenizer, so derive the budget rather than hardcode it:
+    # half the longest prompt splits that prefill whatever TEST_MODEL is, while
+    # the shorter prompts are still co-scheduled whole. Floor is max_num_seqs.
+    tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL)
+    prompt_lens = [len(tokenizer(p).input_ids) for p in prompts]
+    max_num_batched_tokens = max(len(prompts), max(prompt_lens) // 2)
+    assert max_num_batched_tokens < max(prompt_lens), (
+        f"a {max_num_batched_tokens}-token budget does not split the longest "
+        f"prompt ({max(prompt_lens)} tokens), so nothing would be chunked"
+    )
+
+    sp = SamplingParams(
+        temperature=0.6, top_p=1.0, max_tokens=16, seed=1234, logprobs=5
+    )
+    llm = LLM_with_max_seqs(
+        model=TEST_MODEL,
+        max_num_seqs=len(prompts),
+        gpu_memory_utilization=float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.5")),
+        max_model_len=4096,
+        attention_config={"backend": backend},
+        max_num_batched_tokens=max_num_batched_tokens,
+    )
+
+    try:
+        bs1 = []
+        for p in prompts:
+            out = llm.generate([p], sp, use_tqdm=False)[0]
+            bs1.append(_extract_step_logprobs(out))
+        if any(logprobs is None for logprobs, _ in bs1):
+            pytest.skip("Logprobs are not available on RequestOutput.")
+
+        outs_batched = llm.generate(prompts, sp, use_tqdm=False)
+        assert len(outs_batched) == len(prompts)
+
+        mismatches = []
+        for i, out in enumerate(outs_batched):
+            logprobs, tokens = _extract_step_logprobs(out)
+            bs1_logprobs, bs1_tokens = bs1[i]
+            n = prompt_lens[i]
+            if tokens != bs1_tokens:
+                mismatches.append(f"prompt {i} ({n} tokens): tokens differ")
+            elif not torch.equal(bs1_logprobs, logprobs):
+                delta = torch.max(torch.abs(bs1_logprobs - logprobs)).item()
+                mismatches.append(
+                    f"prompt {i} ({n} tokens): logprobs differ, max|delta|={delta:.3e}"
+                )
+
+        assert not mismatches, (
+            "Batch invariance violated under chunked prefill "
+            f"(max_num_batched_tokens={max_num_batched_tokens}): "
+            + "; ".join(mismatches)
+        )
+    finally:
+        with contextlib.suppress(Exception):
+            llm.shutdown()
+
+
 @skip_unsupported
 @pytest.mark.parametrize(
     "backend",
     BACKENDS,
 )
 def test_simple_generation(backend):
-    """
-    Simple test that runs the model with a basic prompt and prints the output.
+    """Simple test that runs the model with a basic prompt and prints the output.
     Useful for quick smoke testing and debugging.
     """
     model = TEST_MODEL
@@ -447,8 +516,7 @@ def test_simple_generation(backend):
 def test_logprobs_without_batch_invariance_should_fail(
     backend, monkeypatch: pytest.MonkeyPatch
 ):
-    """
-    This test is the inverse of test_logprobs_bitwise_batch_invariance_bs1_vs_bsN.
+    """This test is the inverse of test_logprobs_bitwise_batch_invariance_bs1_vs_bsN.
     It DISABLES batch invariance mode and expects to see non-deterministic behavior
     between BS=1 and BS=N runs. This demonstrates that batch invariance is actually
     doing something useful.
@@ -660,8 +728,7 @@ def test_logprobs_without_batch_invariance_should_fail(
 def test_decode_logprobs_match_prefill_logprobs(
     backend,
 ):
-    """
-    Test that verifies decode logprobs match prefill logprobs.
+    """Test that verifies decode logprobs match prefill logprobs.
 
     For each decoded token at position i:
     1. Run decode to generate N tokens and collect their logprobs
@@ -926,14 +993,16 @@ def LLM_with_max_seqs(
     max_model_len: int,
     attention_config: dict | None = None,
     kernel_config: dict | None = None,
+    max_num_batched_tokens: int | None = None,
 ) -> LLM:
-    """
-    Helper to construct an LLM with a specific max_num_seqs (batch-size limit)
+    """Helper to construct an LLM with a specific max_num_seqs (batch-size limit)
     using the high-level v1 LLM API, while constraining memory usage.
     """
     extra_kwargs: dict = {}
     if kernel_config is not None:
         extra_kwargs["kernel_config"] = kernel_config
+    if max_num_batched_tokens is not None:
+        extra_kwargs["max_num_batched_tokens"] = max_num_batched_tokens
     return LLM(
         model=model,
         max_num_seqs=max_num_seqs,
diff --git a/tests/v1/determinism/test_batch_invariance_vlm.py b/tests/v1/determinism/test_batch_invariance_vlm.py
index 266768f67253..00f138d62c81 100644
--- a/tests/v1/determinism/test_batch_invariance_vlm.py
+++ b/tests/v1/determinism/test_batch_invariance_vlm.py
@@ -59,15 +59,22 @@ def _assert_batch_invariant(llm: LLM, inputs: list[dict], sampling) -> None:
         _extract_step_logprobs(llm.generate([inp], sampling, use_tqdm=False)[0])
         for inp in inputs
     ]
-    bsN = [_extract_step_logprobs(out) for out in llm.generate(inputs, sampling)]
+    bsN_outputs = llm.generate(inputs, sampling, use_tqdm=False)
+    assert len(bs1) == len(inputs)
+    assert len(bsN_outputs) == len(inputs)
+    bsN = [_extract_step_logprobs(out) for out in bsN_outputs]
 
     failures = []
     for i, ((lp1, t1), (lpN, tN)) in enumerate(zip(bs1, bsN)):
         if lp1 is None or lpN is None:
-            pytest.skip(
-                "Logits are not available on RequestOutput; "
-                "enable logprobs return to run this test."
+            failures.append(f"req {i}: selected-token logprobs are unavailable")
+            continue
+        if len(t1) != sampling.max_tokens or len(tN) != sampling.max_tokens:
+            failures.append(
+                f"req {i}: expected {sampling.max_tokens} tokens, "
+                f"got bs1={len(t1)} bsN={len(tN)}"
             )
+            continue
         if t1 != tN:
             failures.append(f"req {i}: token mismatch bs1={t1} bsN={tN}")
         elif not torch.equal(lp1, lpN):
@@ -84,16 +91,43 @@ def _assert_batch_invariant(llm: LLM, inputs: list[dict], sampling) -> None:
 @skip_if_not_cuda
 @pytest.mark.parametrize("input_type", ["image", "video"])
 @pytest.mark.parametrize("mm_encoder_attn_backend", ["FLASH_ATTN", "TORCH_SDPA"])
-def test_vlm_batch_invariance_bs1_vs_bsN(input_type: str, mm_encoder_attn_backend: str):
+def test_vlm_batch_invariance_bs1_vs_bsN(
+    input_type: str,
+    mm_encoder_attn_backend: str,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
+    _run_vlm_batch_invariance(input_type, mm_encoder_attn_backend, enforce_eager=True)
+
+
+@skip_if_not_cuda
+def test_vlm_batch_invariance_default_execution(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
+    _run_vlm_batch_invariance("image", "FLASH_ATTN", enforce_eager=False)
+
+
+def _run_vlm_batch_invariance(
+    input_type: str, mm_encoder_attn_backend: str, enforce_eager: bool
+) -> None:
     inputs = _make_inputs(input_type, num_reqs=4)
-    sampling = SamplingParams(temperature=0.0, max_tokens=8, seed=1234, logprobs=5)
+    sampling = SamplingParams(
+        temperature=0.0,
+        max_tokens=8,
+        seed=1234,
+        ignore_eos=True,
+        logprobs=5,
+    )
 
     llm = LLM(
         model=VLM_TEST_MODEL,
         dtype="bfloat16",
-        enforce_eager=True,
+        enforce_eager=enforce_eager,
         max_num_seqs=16,
+        max_model_len=4096,
+        max_num_batched_tokens=2048,
         gpu_memory_utilization=0.85,
+        enable_prefix_caching=False,
+        mm_processor_cache_gb=0,
         mm_encoder_attn_backend=mm_encoder_attn_backend,
     )
     try:
diff --git a/tests/v1/determinism/test_matmul_batch_invariant.py b/tests/v1/determinism/test_matmul_batch_invariant.py
index 086490526e50..af395e74e7f1 100644
--- a/tests/v1/determinism/test_matmul_batch_invariant.py
+++ b/tests/v1/determinism/test_matmul_batch_invariant.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test batch-invariant matmul against torch.matmul for various shape combinations.
+"""Test batch-invariant matmul against torch.matmul for various shape combinations.
 
 Tests correctness (matches torch.matmul) and batch invariance (result for one
 item doesn't change based on other items in the batch).
@@ -53,9 +52,7 @@
 )
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_matmul_correctness(a_shape, b_shape, dtype):
-    """
-    Compare matmul_batch_invariant against torch.matmul for various shapes.
-    """
+    """Compare matmul_batch_invariant against torch.matmul for various shapes."""
     device = torch.device(DEVICE_TYPE)
 
     torch.manual_seed(42)
@@ -87,11 +84,9 @@ def test_matmul_correctness(a_shape, b_shape, dtype):
 @skip_unsupported
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_matmul_batch_invariance(dtype):
-    """
-    Verify that the result for one item is bitwise identical regardless
+    """Verify that the result for one item is bitwise identical regardless
     of what other items are in the batch.
     """
-
     device = torch.device(DEVICE_TYPE)
 
     torch.manual_seed(42)
diff --git a/tests/v1/determinism/test_online_batch_invariance.py b/tests/v1/determinism/test_online_batch_invariance.py
index de80e4918c1e..dcba87dc4927 100644
--- a/tests/v1/determinism/test_online_batch_invariance.py
+++ b/tests/v1/determinism/test_online_batch_invariance.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-HTTP-based batch invariance test: send requests to a running
+"""HTTP-based batch invariance test: send requests to a running
 vLLM server and compare BS=1 vs BS=N results (tokens and per-step logprobs).
 
 Environment variables:
diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py
index a9a15d1bf083..5699a661c17f 100644
--- a/tests/v1/determinism/test_rms_norm_batch_invariant.py
+++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py
@@ -41,8 +41,7 @@ def test_rms_norm_batch_invariant_vs_reference(
     eps: float,
     seed: int,
 ):
-    """
-    Compare batch-invariant Triton RMS norm against a PyTorch reference.
+    """Compare batch-invariant Triton RMS norm against a PyTorch reference.
 
     Tests that the Triton-based batch-invariant RMS norm produces numerically
     equivalent results to an independent implementation across configurations.
@@ -90,9 +89,7 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
     n_extra: int,
     seed: int,
 ):
-    """
-    Test the batch-invariant fused residual-add + RMSNorm helper directly.
-    """
+    """Test the batch-invariant fused residual-add + RMSNorm helper directly."""
     device = torch.device(DEVICE_TYPE)
 
     torch.manual_seed(seed)
@@ -313,8 +310,7 @@ def per_block_quant(x):
 def test_rms_norm_3d_input(
     default_vllm_config, batch_size: int, seq_len: int, hidden_size: int
 ):
-    """
-    Test RMS norm with 3D input tensors (batch, seq_len, hidden_size).
+    """Test RMS norm with 3D input tensors (batch, seq_len, hidden_size).
 
     Ensures that the batch-invariant RMS norm correctly handles multi-dimensional
     inputs that are common in transformer models.
@@ -349,8 +345,7 @@ def test_rms_norm_3d_input(
 
 @skip_if_not_cuda
 def test_rms_norm_numerical_stability(default_vllm_config):
-    """
-    Test RMS norm numerical stability with extreme values.
+    """Test RMS norm numerical stability with extreme values.
 
     Ensures that both implementations handle edge cases like very small or large
     values without producing NaN or Inf.
@@ -406,8 +401,7 @@ def test_rms_norm_numerical_stability(default_vllm_config):
 
 @skip_unsupported
 def test_rms_norm_formula(default_vllm_config):
-    """
-    Test that RMS norm follows the correct mathematical formula.
+    """Test that RMS norm follows the correct mathematical formula.
 
     Verifies: output = input / sqrt(mean(input^2) + eps) * weight
     """
@@ -440,8 +434,7 @@ def test_rms_norm_formula(default_vllm_config):
 @skip_if_not_cuda
 @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384])
 def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
-    """
-    Test RMS norm with various hidden sizes to ensure block size handling.
+    """Test RMS norm with various hidden sizes to ensure block size handling.
 
     The Triton kernel uses a fixed BLOCK_SIZE=1024, so this tests that it
     correctly handles hidden sizes both smaller and larger than the block size.
@@ -474,8 +467,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
 
 @skip_unsupported
 def test_rms_norm_determinism(default_vllm_config):
-    """
-    Test that batch-invariant RMS norm produces deterministic results.
+    """Test that batch-invariant RMS norm produces deterministic results.
 
     Runs the same input through the kernel multiple times and verifies
     identical outputs.
diff --git a/tests/v1/determinism/test_xpu_batch_invariant_ut.py b/tests/v1/determinism/test_xpu_batch_invariant_ut.py
new file mode 100644
index 000000000000..53e778ac0dfd
--- /dev/null
+++ b/tests/v1/determinism/test_xpu_batch_invariant_ut.py
@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""XPU-only batch-invariance tests for norm, collectives, and sampler behavior."""
+
+import pytest
+import torch
+
+from vllm import envs
+from vllm.config import CacheConfig, LoadConfig, ModelConfig, VllmConfig
+from vllm.distributed.device_communicators.xpu_communicator import XpuCommunicator
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_xpu(), reason="XPU batch-invariance tests"
+)
+
+
+@pytest.fixture
+def tiny_qwen3_moe_path(tmp_path):
+    """Create a local Qwen3-MoE config for dummy-loading tests."""
+    from transformers import Qwen3MoeConfig
+
+    config = Qwen3MoeConfig(
+        architectures=["Qwen3MoeForCausalLM"],
+        vocab_size=256,
+        hidden_size=512,
+        intermediate_size=1024,
+        moe_intermediate_size=512,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=128,
+        num_experts=4,
+        num_experts_per_tok=2,
+        max_position_embeddings=512,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+    )
+    config.save_pretrained(tmp_path)
+    return tmp_path
+
+
+@pytest.mark.parametrize("source", ["argument", "checkpoint"])
+def test_rejects_weight_quantization(tiny_qwen3_moe_path, source):
+    """Reject both explicit and checkpoint-inferred quantization before loading."""
+    from transformers import Qwen3MoeConfig
+
+    if source == "checkpoint":
+        config = Qwen3MoeConfig.from_pretrained(tiny_qwen3_moe_path)
+        config.quantization_config = {
+            "quant_method": "fp8",
+            "activation_scheme": "dynamic",
+        }
+        config.save_pretrained(tiny_qwen3_moe_path)
+
+    model_config = ModelConfig(
+        model=str(tiny_qwen3_moe_path),
+        skip_tokenizer_init=True,
+        dtype="bfloat16",
+        quantization="fp8" if source == "argument" else None,
+    )
+    assert model_config.quantization == "fp8"
+    with pytest.raises(ValueError, match="supports only unquantized models"):
+        VllmConfig(
+            model_config=model_config,
+            load_config=LoadConfig(load_format="dummy"),
+        )
+
+
+@pytest.mark.parametrize("cache_dtype", ["fp8", "int8_per_token_head", "nvfp4"])
+def test_rejects_quantized_kv_cache(cache_dtype):
+    with pytest.raises(ValueError, match="does not support quantized KV caches"):
+        VllmConfig(cache_config=CacheConfig(cache_dtype=cache_dtype))
+
+
+@pytest.mark.parametrize("cache_dtype", ["auto", "float16", "bfloat16"])
+def test_accepts_unquantized_kv_cache(cache_dtype):
+    config = VllmConfig(cache_config=CacheConfig(cache_dtype=cache_dtype))
+    assert config.cache_config.cache_dtype == cache_dtype
+
+
+def test_quantized_kv_cache_allowed_without_batch_invariance(monkeypatch):
+    monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False)
+    config = VllmConfig(cache_config=CacheConfig(cache_dtype="fp8"))
+    assert config.cache_config.cache_dtype == "fp8"
+
+
+# These dimensions and seeds locally reproduce non-invariant norm mismatches.
+# They make the invariant equality assertion exercise a discriminating case.
+@pytest.mark.parametrize(
+    (
+        "norm_name",
+        "batch_size",
+        "hidden_size",
+        "position",
+        "weight_seed",
+        "input_seed",
+        "input_pattern",
+    ),
+    [
+        ("rms", 1024, 16384, 512, 7, 42, "normal"),
+        ("gemma", 257, 8192, 255, 1, 2, "normal"),
+        ("rms", 5, 8192, 3, 7, 46, "scaled"),
+        ("gemma", 5, 8192, 3, 1, 46, "scaled"),
+    ],
+)
+def test_residual_norm_preserves_batch_invariance(
+    norm_name,
+    batch_size,
+    hidden_size,
+    position,
+    weight_seed,
+    input_seed,
+    input_pattern,
+    monkeypatch,
+    default_vllm_config,
+):
+    """Residual norm output must not change when its row changes batch position."""
+    from vllm.model_executor.determinism.batch_invariant import init_batch_invariance
+    from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
+
+    monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
+    init_batch_invariance()
+    norm_cls = RMSNorm if norm_name == "rms" else GemmaRMSNorm
+    norm = norm_cls(hidden_size, eps=1e-6).to(device="xpu", dtype=torch.bfloat16)
+    with torch.no_grad():
+        torch.manual_seed(weight_seed)
+        norm.weight.uniform_(-0.5, 0.5)
+
+    if input_pattern == "scaled":
+        # Heterogeneous columns expose native mean's batch-dependent rounding.
+        generator = torch.Generator(device="cpu").manual_seed(input_seed)
+        scale = torch.randn(hidden_size, generator=generator).exp()
+        x = (torch.randn(batch_size, hidden_size, generator=generator) * scale).to(
+            device="xpu", dtype=torch.bfloat16
+        )
+        residual = (
+            torch.randn(batch_size, hidden_size, generator=generator) * scale
+        ).to(device="xpu", dtype=torch.bfloat16)
+    else:
+        torch.manual_seed(input_seed)
+        x = torch.randn(batch_size, hidden_size, device="xpu", dtype=torch.bfloat16)
+        residual = torch.randn_like(x)
+
+    single_output, single_residual = norm(
+        x[position : position + 1].clone(),
+        residual[position : position + 1].clone(),
+    )
+    batch_output, batch_residual = norm(x.clone(), residual.clone())
+
+    assert torch.equal(single_output, batch_output[position : position + 1])
+    assert torch.equal(single_residual, batch_residual[position : position + 1])
+
+
+@pytest.mark.parametrize(
+    "operation", ["all_reduce", "reduce_scatter", "reduce_scatterv"]
+)
+@pytest.mark.parametrize("rank", range(4))
+def test_collectives_preserve_rank_order_and_partition(monkeypatch, operation, rank):
+    """Cancellation exposes reordered sums; mock ranks require no extra XPUs."""
+    from vllm.distributed.device_communicators import xpu_communicator
+
+    inputs = [
+        torch.full((8, 3), value, dtype=torch.float32, device="xpu")
+        for value in (1e20, -1e20, 3.0, 0.0)
+    ]
+    inputs[2].add_(torch.arange(24, device="xpu").reshape(8, 3))
+
+    def all_gather_single(output, input_, group):
+        torch.testing.assert_close(input_, inputs[rank].reshape(-1), rtol=0, atol=0)
+        output.copy_(torch.stack(inputs).reshape(-1))
+
+    monkeypatch.setattr(xpu_communicator.dist, "all_gather_single", all_gather_single)
+    communicator = XpuCommunicator.__new__(XpuCommunicator)
+    communicator.world_size = 4
+    communicator.rank_in_group = rank
+    communicator.device_group = object()
+    expected = inputs[0].clone()
+    for value in inputs[1:]:
+        expected.add_(value)
+    assert torch.equal(expected, inputs[2])
+
+    if operation == "all_reduce":
+        actual = communicator.all_reduce(inputs[rank])
+    elif operation == "reduce_scatter":
+        actual = communicator.reduce_scatter(inputs[rank], dim=0)
+        expected = expected.chunk(4)[rank]
+    else:
+        sizes = [1, 3, 2, 2]
+        actual = communicator.reduce_scatterv(inputs[rank], dim=0, sizes=sizes)
+        expected = expected.split(sizes)[rank]
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+# This batch size, vocabulary size, and positions locally reproduce an
+# unseeded non-invariant XPU sampler mismatch.
+@pytest.mark.parametrize("position", [0, 128, 255])
+def test_seeded_sampler_preserves_full_distribution_and_rng(position, monkeypatch):
+    """Moving a request must preserve its distribution and per-request RNG."""
+    from vllm.triton_utils import HAS_TRITON
+    from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
+
+    assert HAS_TRITON, "XPU sampling must exercise the Triton top-k/top-p path"
+    monkeypatch.setattr(envs, "VLLM_XPU_USE_SAMPLER_KERNEL", True)
+    sampler = TopKTopPSampler(logprobs_mode="processed_logprobs")
+    torch.manual_seed(42)
+    logits = torch.randn(256, 8192, device="xpu", dtype=torch.float32)
+    single_rng = torch.Generator(device="xpu").manual_seed(123)
+    batch_rng = torch.Generator(device="xpu").manual_seed(123)
+    k = torch.full((256,), 32, dtype=torch.int32, device="xpu")
+    p = torch.full((256,), 0.9, device="xpu")
+    for _ in range(3):
+        single_token, single_logprobs = sampler(
+            logits[position : position + 1].clone(),
+            {0: single_rng},
+            k[position : position + 1],
+            p[position : position + 1],
+        )
+        batch_tokens, batch_logprobs = sampler(
+            logits.clone(), {position: batch_rng}, k, p
+        )
+        assert single_logprobs is not None and batch_logprobs is not None
+        torch.testing.assert_close(
+            single_logprobs[0], batch_logprobs[position], rtol=0, atol=0
+        )
+        assert single_token.item() == batch_tokens[position].item()
+        assert torch.equal(single_rng.get_state(), batch_rng.get_state())
diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py
index e1adc948a853..cb23910566c7 100644
--- a/tests/v1/distributed/test_async_llm_dp.py
+++ b/tests/v1/distributed/test_async_llm_dp.py
@@ -493,7 +493,6 @@ async def gen(rid: str):
 @pytest.mark.parametrize("expert_parallel", [False, True])
 async def test_dp_pause_keep_then_resume(expert_parallel: bool):
     """Start generation, pause after a few tokens (keep mode), resume; verify gap."""
-
     pause_duration = 2.0
     min_tokens_before_pause = 3
 
@@ -607,8 +606,7 @@ async def consume_gen(req_id: str) -> None:
 
 @pytest.mark.asyncio
 async def test_dp_pause_barrier_request_deadlock():
-    """
-    Test that start_dp_wave is ignored while paused.
+    """Test that start_dp_wave is ignored while paused.
 
     Sequence:
       1. Pause all engines (PAUSED_ALL).
diff --git a/tests/v1/distributed/test_dbo.py b/tests/v1/distributed/test_dbo.py
index a4e3d92be6c4..37e47e3c6900 100644
--- a/tests/v1/distributed/test_dbo.py
+++ b/tests/v1/distributed/test_dbo.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test Dual Batch Overlap (DBO) with Data Parallelism + Expert Parallelism.
+"""Test Dual Batch Overlap (DBO) with Data Parallelism + Expert Parallelism.
 
 DBO is specifically designed for DP+EP scenarios to hide communication latency
 by overlapping computation of two batches. This test validates that DBO works
@@ -54,9 +53,7 @@
     ),
 )
 def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available):
-    """
-    Test DBO with DP+EP using GSM8K evaluation.
-    """
+    """Test DBO with DP+EP using GSM8K evaluation."""
     lm_eval = pytest.importorskip("lm_eval")
     required_gpus = DP_SIZE
 
diff --git a/tests/v1/distributed/test_internal_lb_dp.py b/tests/v1/distributed/test_internal_lb_dp.py
index e4f7b3d5a83d..b438d898b337 100644
--- a/tests/v1/distributed/test_internal_lb_dp.py
+++ b/tests/v1/distributed/test_internal_lb_dp.py
@@ -261,7 +261,6 @@ def __init__(
 
     def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
         """Start API-only server and headless engines server."""
-
         # Start API-only server (Node 0) - no engines, only API server
         api_server_args = self.base_server_args.copy()
         api_server_args.extend(
@@ -607,7 +606,6 @@ async def test_api_only_multinode_dp_completion(
     model_name: str,
 ) -> None:
     """Test API-only server with all engines on separate headless server."""
-
     # Test single request
     result = await _make_completion_request(api_only_client, model_name)
     assert result is not None
diff --git a/tests/v1/distributed/test_pp_dp_v2.py b/tests/v1/distributed/test_pp_dp_v2.py
index b931db151005..1e5d4267f474 100644
--- a/tests/v1/distributed/test_pp_dp_v2.py
+++ b/tests/v1/distributed/test_pp_dp_v2.py
@@ -113,7 +113,6 @@ async def _one(i: int) -> int:
 async def test_pp_dp_v2_abort_mid_decode():
     """Cancel half the in-flight requests mid-stream and confirm the
     engine survives the abort storm."""
-
     with ExitStack() as after:
         engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True))
         after.callback(engine.shutdown)
@@ -154,7 +153,6 @@ async def _maybe_cancel(i: int):
 async def test_pp_dp_v2_pause_resume():
     """Pause an engine with a request in flight, then resume and confirm
     new requests still work."""
-
     with ExitStack() as after:
         engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True))
         after.callback(engine.shutdown)
diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py
index 1024829307c7..8218f325290d 100644
--- a/tests/v1/e2e/general/test_async_scheduling.py
+++ b/tests/v1/e2e/general/test_async_scheduling.py
@@ -111,7 +111,6 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke
     preemption, executor, async scheduling, prefill chunking,
     spec decoding model length.
     """
-
     spec_config = {
         "method": "eagle3",
         "num_speculative_tokens": 2,
@@ -171,7 +170,6 @@ def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
     - Async scheduling enabled (as in production)
     - Different executors and chunking settings
     """
-
     # Variant with larger speculation window
     ngram_gpu_config = {
         "method": "ngram_gpu",
@@ -207,7 +205,6 @@ def run_tests(
 ):
     """Test consistency of combos of async scheduling, preemption,
     uni/multiproc executor with spec decoding."""
-
     # Flex attention supports float32.
     attention_config = {"backend": "FLEX_ATTENTION"}
 
diff --git a/tests/v1/e2e/general/test_context_length.py b/tests/v1/e2e/general/test_context_length.py
index 7e77b31fb1b7..91a450089f6f 100644
--- a/tests/v1/e2e/general/test_context_length.py
+++ b/tests/v1/e2e/general/test_context_length.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for vLLM `vllm/v1/engine/processor.Processor._validate_model_input()`
+"""Tests for vLLM `vllm/v1/engine/processor.Processor._validate_model_input()`
 handling of maximum context length for decoder models.
 
 This test ensures:
@@ -36,7 +35,6 @@ def test_decoder_max_context_length_validation(
 ) -> None:
     """Check vLLM decoder model input validation for edge cases where
     the prompt length is (almost) equal to the max model length."""
-
     prompt_ids = [[43] * prompt_len]
 
     with vllm_runner(
@@ -74,7 +72,6 @@ def test_auto_fit_max_model_len_rejects_oversized_input(
     the engine auto-fits max_model_len to a small value. The frontend
     must see this reduced value and reject prompts that exceed it,
     rather than accepting them and hanging."""
-
     # Use a small KV cache budget to force auto-fit to a small
     # max_model_len. Pin block_size=16 so the budget is independent
     # of the platform's default block size. One block for this model is
diff --git a/tests/v1/e2e/general/test_correctness_sliding_window.py b/tests/v1/e2e/general/test_correctness_sliding_window.py
index 5209d1894db7..ce71bf1e5def 100644
--- a/tests/v1/e2e/general/test_correctness_sliding_window.py
+++ b/tests/v1/e2e/general/test_correctness_sliding_window.py
@@ -35,8 +35,7 @@ class TestConfig:
 def test_sliding_window_retrieval(
     model, batch_size, seed, disable_hybrid_kv_cache_manager, vllm_runner
 ):
-    """
-    The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then
+    """The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then
     asks for value of one of them (which is outside the sliding window).
     If we tell it upfront which we are going to be looking for, then
     it answers correctly (mostly).
@@ -112,14 +111,14 @@ def test_hybrid_kv_cache_manager_output_equivalence(model, vllm_runner):
 
 
 def check_length(prompts: list[str], llm: LLM, sliding_window: int):
-    """
-    Check if the prompt length is valid, i.e., longer than the sliding window
+    """Check if the prompt length is valid, i.e., longer than the sliding window
     size and shorter than the model's max length.
 
     Args:
         prompts: list of prompts
         llm: LLM object
         sliding_window: Sliding window size
+
     """
     tokenizer = llm.get_tokenizer()
     max_model_len = llm.llm_engine.model_config.max_model_len
diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py
new file mode 100644
index 000000000000..dd4c5886570d
--- /dev/null
+++ b/tests/v1/e2e/general/test_hisparse.py
@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import prometheus_client
+import pytest
+import torch
+
+from tests.conftest import VllmRunner
+from tests.utils import fork_new_process_for_each_test
+from vllm.config import AttentionConfig, HiSparseConfig, KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
+    HiSparseConnector,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
+    HiSparseConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
+    MultiConnector,
+)
+from vllm.platforms import current_platform
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
+
+MODEL = "deepseek-ai/DeepSeek-V3.2"
+
+HF_OVERRIDES = {
+    "num_hidden_layers": 8,
+    "hidden_size": 256,
+    "intermediate_size": 512,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 1,
+    "n_routed_experts": 8,
+    "num_experts_per_tok": 2,
+    "index_topk": 128,
+}
+
+
+def _shrink_config(config):
+    overrides = HF_OVERRIDES.copy()
+    if any("MTP" in architecture for architecture in config.architectures):
+        overrides.pop("num_hidden_layers")
+    config.update(overrides)
+    return config
+
+
+def _num_hisparse_spills(runner: VllmRunner) -> int:
+    client = runner.llm.llm_engine.engine_core
+    coordinator = get_hisparse_coordinator(
+        client.engine_core.scheduler.kv_cache_manager
+    )
+    return coordinator.next_spill_id
+
+
+def _offload_load_bytes() -> float:
+    return sum(
+        sample.value
+        for metric in prometheus_client.REGISTRY.collect()
+        if metric.name == "vllm:kv_offload_load_bytes"
+        for sample in metric.samples
+        if sample.name == "vllm:kv_offload_load_bytes_total"
+    )
+
+
+def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker:
+    engine_core = runner.llm.llm_engine.engine_core.engine_core
+    model_runner = engine_core.model_executor.driver_worker.worker.model_runner
+    connector = model_runner.kv_connector.kv_connector
+    if isinstance(connector, MultiConnector):
+        connector = next(
+            child
+            for child in connector.sub_connectors
+            if isinstance(child, HiSparseConnector)
+        )
+    assert isinstance(connector, HiSparseConnector)
+    assert connector.connector_worker is not None
+    return connector.connector_worker
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA"
+)
+@pytest.mark.parametrize(
+    "with_offloading", [False, True], ids=["standalone", "offload"]
+)
+@fork_new_process_for_each_test
+def test_hisparse_spill_and_prefix_restore(
+    monkeypatch: pytest.MonkeyPatch,
+    vllm_runner: type[VllmRunner],
+    with_offloading: bool,
+):
+    """Spilled prefixes restore and FULL-graph decode writes reach host KV.
+
+    A regression omitted attention metadata before FULL graph replay, so decode
+    completed normally while its newly written KV rows were never copied to host.
+    """
+    capability = current_platform.get_device_capability()
+    if capability is None or capability.major < 9:
+        pytest.skip("Sparse MLA requires Hopper or newer")
+
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
+
+    target = [1000 + i % 64 for i in range(257)]
+    pressure = [
+        [2000 + request_idx * 128 + i % 64 for i in range(257)]
+        for request_idx in range(4)
+    ]
+    hisparse_connector = {
+        "kv_connector": "HiSparseConnector",
+        "kv_role": "kv_both",
+        "kv_connector_extra_config": {"host_pool_gib": 1},
+    }
+    kv_transfer_config = KVTransferConfig(**hisparse_connector)
+    if with_offloading:
+        kv_transfer_config = KVTransferConfig(
+            kv_connector="MultiConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={
+                "connectors": [
+                    hisparse_connector,
+                    {
+                        "kv_connector": "OffloadingConnector",
+                        "kv_role": "kv_both",
+                        "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
+                    },
+                ]
+            },
+        )
+
+    with vllm_runner(
+        MODEL,
+        load_format="dummy",
+        hf_overrides=_shrink_config,
+        attention_config=AttentionConfig(
+            hisparse_config=HiSparseConfig(
+                device_buffer_size=512,
+            )
+        ),
+        kv_transfer_config=kv_transfer_config,
+        block_size=64,
+        max_model_len=320,
+        max_num_batched_tokens=1024,
+        max_num_seqs=4,
+        num_gpu_blocks_override=128,
+        disable_log_stats=False,
+        enable_chunked_prefill=True,
+        enable_prefix_caching=True,
+        speculative_config={"method": "mtp", "num_speculative_tokens": 3},
+        compilation_config={
+            "cudagraph_mode": "FULL_AND_PIECEWISE",
+            "cudagraph_capture_sizes": [1, 4],
+        },
+    ) as runner:
+        worker = _get_hisparse_worker(runner)
+        engine_core = runner.llm.llm_engine.engine_core.engine_core
+        model_runner = engine_core.model_executor.driver_worker.worker.model_runner
+        full_graph_calls = 0
+        full_replay_pending = False
+        original_run_fullgraph = model_runner.cudagraph_manager.run_fullgraph
+
+        def record_fullgraph(batch_desc):
+            nonlocal full_graph_calls, full_replay_pending
+            full_graph_calls += 1
+            full_replay_pending = True
+            return original_run_fullgraph(batch_desc)
+
+        monkeypatch.setattr(
+            model_runner.cudagraph_manager, "run_fullgraph", record_fullgraph
+        )
+
+        original_finish_forward = worker.finish_forward
+        verified_full_replay_rows = 0
+        copied_nonzero_kv = False
+
+        def finish_forward():
+            nonlocal full_replay_pending, verified_full_replay_rows
+            nonlocal copied_nonzero_kv
+            try:
+                original_finish_forward()
+                if not full_replay_pending:
+                    return
+                # Synchronize only the test's inspection, not finish_forward itself.
+                with gpu_sync_allowed():
+                    torch.accelerator.synchronize()
+                assert worker._row_mirrors
+                for layer_index, cache in enumerate(worker.cache_handles):
+                    source_index = cache.runtime.resident_source_index
+                    for mirror in worker._row_mirrors:
+                        source_row = mirror.source_starts[source_index]
+                        destination_row = mirror.destination_start
+                        num_rows = mirror.num_rows
+                        source_block, source_offset = divmod(
+                            source_row, worker.kernel_block_size
+                        )
+                        with gpu_sync_allowed():
+                            gpu_rows = worker.resident_caches[layer_index][
+                                source_block, source_offset : source_offset + num_rows
+                            ].cpu()
+                        host_rows = worker.host_caches[layer_index][
+                            destination_row : destination_row + num_rows
+                        ]
+                        assert torch.equal(host_rows, gpu_rows)
+                        copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item())
+                        verified_full_replay_rows += num_rows
+            finally:
+                full_replay_pending = False
+
+        monkeypatch.setattr(worker, "finish_forward", finish_forward)
+
+        expected = runner.generate_greedy([target], max_tokens=8)
+
+        assert full_graph_calls > 0
+        assert verified_full_replay_rows > 0
+        assert copied_nonzero_kv
+
+        runner.generate_greedy(pressure, max_tokens=8)
+
+        assert _num_hisparse_spills(runner) > 0
+        load_bytes = _offload_load_bytes()
+        actual = runner.generate_greedy([target], max_tokens=8)
+        runner.generate_greedy([[42]], max_tokens=1)
+
+        assert actual == expected
+        if with_offloading:
+            assert _offload_load_bytes() > load_bytes
diff --git a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py
index 3155da82ea85..b03acc21afc0 100644
--- a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py
+++ b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py
@@ -17,9 +17,7 @@
 
 @pytest.fixture
 def test_prompts():
-    """
-    Adapted from tests/v1/e2e/spec_decode/utils.py
-    """
+    """Adapted from tests/v1/e2e/spec_decode/utils.py."""
     prompt_types = ["repeat", "sentence"]
     # Setting higher num prompts increases the chance of numerics mismatch
     # due to matrix multiplication numerics depending on batch dimension
@@ -52,12 +50,23 @@ def test_prompts():
 
 
 @use_fork_for_test
-@pytest.mark.parametrize("kv_sharing_fast_prefill", [False, True])
+@pytest.mark.parametrize(
+    "kv_sharing_fast_prefill,use_v2_model_runner",
+    [
+        (False, False),
+        (True, False),
+        (True, True),
+        # (False, True) omitted: fast-prefill-off behavior is runner-generic,
+        # and skipping it saves ~9 min of mostly torch.compile time in CI. The
+        # no-silent-fallback assertion below is exercised by the V2 cases.
+    ],
+)
 @pytest.mark.parametrize("enforce_eager", [True, False])
 def test_kv_sharing_fast_prefill(
     monkeypatch: pytest.MonkeyPatch,
     kv_sharing_fast_prefill: bool,
     enforce_eager: bool,
+    use_v2_model_runner: bool,
 ):
     if not enforce_eager and current_platform.is_rocm():
         # Relevant context: https://github.com/vllm-project/vllm/pull/29244
@@ -84,6 +93,7 @@ def test_kv_sharing_fast_prefill(
             m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
         else:
             m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+        m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_model_runner else "0")
 
         prompts, answer, indices = prep_prompts(batch_size)
 
@@ -95,6 +105,8 @@ def test_kv_sharing_fast_prefill(
             kv_sharing_fast_prefill=kv_sharing_fast_prefill,
             attention_backend="TRITON_ATTN",
         )
+        # Guard against a silent fallback to the other model runner.
+        assert llm.llm_engine.vllm_config.use_v2_model_runner == use_v2_model_runner
         responses = llm.generate(prompts, sampling_params)
         check_answers(
             indices,
diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py
index 27331516cbd3..5b6826a4b3fc 100644
--- a/tests/v1/e2e/general/test_mamba_prefix_cache.py
+++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py
@@ -31,9 +31,8 @@
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.worker import mamba_utils
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
 from vllm.v1.worker.mamba_utils import get_mamba_groups
 
 
@@ -315,7 +314,7 @@ def check_copy_info(
         action: tuple[int, int],
         kv_cache_config: KVCacheConfig,
         forward_context: dict[str, Any],
-        input_batch: GPUInputBatch,
+        input_batch: InputBatch,
     ):
         assert copy_info is not None
         if action == (-1, -1):
@@ -371,7 +370,7 @@ def fake_preprocess_mamba_fn(
         kv_cache_config: KVCacheConfig,
         cache_config: CacheConfig,
         mamba_state_idx: dict[str, int],
-        input_batch: GPUInputBatch,
+        input_batch: InputBatch,
         requests: dict[str, CachedRequestState],
         forward_context: dict[str, Any],
         mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
@@ -837,7 +836,7 @@ def get_mamba_prefix_cache_step_configs(
                 StepAction(
                     560 * 10,
                     4,
-                    [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1]
+                    [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
                     if a
                     else [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
                     (9, 10),
@@ -1123,6 +1122,7 @@ def fake_sample(
             logprobs_tensors=None,
             num_nans=None,
             num_sampled=num_sampled,
+            num_rejected=num_rejected,
         )
         return sampler_output, num_sampled, num_rejected
 
diff --git a/tests/v1/e2e/general/test_min_tokens.py b/tests/v1/e2e/general/test_min_tokens.py
index 0100a38ba150..c87339b7685b 100644
--- a/tests/v1/e2e/general/test_min_tokens.py
+++ b/tests/v1/e2e/general/test_min_tokens.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Comprehensive end-to-end tests for `min_tokens` in the V1 engine.
+"""Comprehensive end-to-end tests for `min_tokens` in the V1 engine.
 
 Addresses #21950: verify and add CI coverage.
 
@@ -28,7 +27,7 @@
 
 
 class MinTokensTestCase:
-    """Data class for min_tokens test scenarios"""
+    """Data class for min_tokens test scenarios."""
 
     def __init__(
         self,
@@ -161,7 +160,7 @@ def __str__(self):
 
 @pytest.fixture(scope="module")
 def llm_v1():
-    """Create V1 LLM instance for testing"""
+    """Create V1 LLM instance for testing."""
     llm = LLM(
         model=TEST_MODEL,
         tensor_parallel_size=1,
@@ -193,7 +192,7 @@ def test_min_tokens_with_terminal_structured_output(llm_v1: LLM):
 
 
 def get_token_count(output: RequestOutput) -> int:
-    """Extract token count from LLM output"""
+    """Extract token count from LLM output."""
     if not output.outputs:
         return 0
     return len(output.outputs[0].token_ids)
@@ -202,7 +201,7 @@ def get_token_count(output: RequestOutput) -> int:
 def assert_min_tokens_satisfied(
     output: RequestOutput, test_case: MinTokensTestCase
 ) -> None:
-    """Assert that min_tokens requirement is satisfied"""
+    """Assert that min_tokens requirement is satisfied."""
     token_count = get_token_count(output)
     stop_reason = output.outputs[0].stop_reason if output.outputs else "no output"
 
@@ -228,8 +227,7 @@ def assert_min_tokens_satisfied(
     ids=lambda tc: tc.name,
 )
 def test_min_tokens_comprehensive(llm_v1: LLM, test_case: MinTokensTestCase):
-    """
-    Comprehensive test for min_tokens functionality in V1 engine.
+    """Comprehensive test for min_tokens functionality in V1 engine.
 
     This test covers all critical scenarios for min_tokens:
     - Basic functionality (should work)
@@ -240,6 +238,7 @@ def test_min_tokens_comprehensive(llm_v1: LLM, test_case: MinTokensTestCase):
     Args:
         llm_v1: V1 LLM instance
         test_case: Test scenario parameters
+
     """
     # Known failing cases are handled via param-level xfail marks above.
 
@@ -279,8 +278,7 @@ def test_min_tokens_comprehensive(llm_v1: LLM, test_case: MinTokensTestCase):
 
 
 def test_min_tokens_basic_functionality(llm_v1: LLM):
-    """
-    Test basic min_tokens functionality without stop conditions.
+    """Test basic min_tokens functionality without stop conditions.
 
     This is a baseline test that should always pass and validates
     that min_tokens works correctly in the simple case.
@@ -302,8 +300,7 @@ def test_min_tokens_basic_functionality(llm_v1: LLM):
     strict=False,
 )
 def test_min_tokens_stop_strings_bug(llm_v1: LLM):
-    """
-    Test the specific bug where stop strings bypass min_tokens.
+    """Test the specific bug where stop strings bypass min_tokens.
 
     This test specifically reproduces the bug Calvin is fixing in PR #22014.
     It should fail until that fix is merged.
@@ -353,8 +350,7 @@ def test_min_tokens_stop_strings_bug(llm_v1: LLM):
     strict=False,
 )
 def test_min_tokens_stop_strings_guaranteed_early_trigger(llm_v1: LLM):
-    """
-    Guaranteed test for stop strings bypassing min_tokens bug.
+    """Guaranteed test for stop strings bypassing min_tokens bug.
 
     Strategy: Use very low temperature and multiple common stop strings
     to virtually guarantee early detection, combined with long min_tokens
@@ -408,8 +404,7 @@ def test_min_tokens_stop_strings_guaranteed_early_trigger(llm_v1: LLM):
     strict=False,
 )
 def test_min_tokens_eos_behavior(llm_v1: LLM):
-    """
-    Verify EOS handling with and without min_tokens.
+    """Verify EOS handling with and without min_tokens.
 
     - Without min_tokens: expect early EOS -> finish_reason == "stop",
       stop_reason is None, and generated tokens < max_tokens (25).
@@ -492,8 +487,7 @@ def test_min_tokens_eos_behavior(llm_v1: LLM):
 
 
 def test_min_tokens_validation():
-    """
-    Test that SamplingParams correctly validates min_tokens parameters.
+    """Test that SamplingParams correctly validates min_tokens parameters.
 
     This tests the parameter validation logic in SamplingParams.
     """
diff --git a/tests/v1/e2e/general/test_pooling_chunked_prefill.py b/tests/v1/e2e/general/test_pooling_chunked_prefill.py
index 976e4d17387e..ab36e9427d98 100644
--- a/tests/v1/e2e/general/test_pooling_chunked_prefill.py
+++ b/tests/v1/e2e/general/test_pooling_chunked_prefill.py
@@ -71,20 +71,19 @@ def retrieve_chunks(self):
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
 def test_pooling_chunked_prefill(vllm_runner, monkeypatch):
     """Test chunked prefill for pooling models with LastPool."""
-
     with monkeypatch.context() as m:
         m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
         model_id = "Qwen/Qwen3-Embedding-0.6B"
 
         chunk_size = 10
 
-        # Set chunking parameters to force chunked prefill
-        # Note: Chunked prefill is automatically handled by vLLM
-        # internally based on the model size and prompt
+        # Cap the token budget to force chunked prefill; the
+        # long_prefill_token_threshold is not applied to a lone request.
         with vllm_runner(
             model_id,
             runner="pooling",
-            long_prefill_token_threshold=chunk_size,
+            max_num_batched_tokens=chunk_size,
+            max_num_seqs=1,
             tensor_parallel_size=1,
             enforce_eager=True,
             enable_chunked_prefill=True,
@@ -124,7 +123,6 @@ def test_pooling_chunked_prefill(vllm_runner, monkeypatch):
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
 def test_pooling_prefix_cache(vllm_runner, monkeypatch):
     """Test chunked prefill for pooling models with LastPool."""
-
     verses = prompt.split("\n\n")
 
     with monkeypatch.context() as m:
diff --git a/tests/v1/e2e/general/test_sharded_sampling.py b/tests/v1/e2e/general/test_sharded_sampling.py
index a580605fb007..868618677698 100644
--- a/tests/v1/e2e/general/test_sharded_sampling.py
+++ b/tests/v1/e2e/general/test_sharded_sampling.py
@@ -2,26 +2,30 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Batch-sharded sampling must be a drop-in for replicated sampling.
 
-The comparison is tolerance-based, not bitwise: engine boots land in
-slightly different kernel/collective states (measured: same-state boots are
-bitwise identical across the two modes; cross-state boots differ by up to
-~0.25 logprob and can flip a near-tie token). A real sharding bug — rows
-routed to the wrong request, misaligned gathers — produces many-nats
-divergence and scrambled top-k sets, far outside these bounds."""
+The two boots are compared token for token, which is only meaningful because
+both schedule the same batches. With the engine in the caller's process every
+request is queued before the first step, so the first prefill has the same
+shape in both runs. Give the engine its own process and the comparison turns
+into a coin toss: prefill starts before the last requests arrive over IPC, so
+one boot batches two requests where the other batches one. The batch shape
+sets the GEMM reduction order, logprobs shift by up to ~0.25, and near-ties
+flip — measured at up to 5 of 16 completions, with sharding disabled on both
+sides. A real sharding bug — rows routed to the wrong request, misaligned
+gathers — produces many-nats divergence and scrambled top-k sets."""
 
 import pytest
 
 from tests.utils import multi_gpu_test
 from vllm import LLM, SamplingParams
 from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.logprobs import Logprob, LogprobsOnePosition
 
 MODEL = "meta-llama/Llama-3.2-1B-Instruct"
 
-# Boot-state numeric noise bound (measured <= 0.25 on identical contexts).
+# Misrouted rows move logprobs by nats; anything under this is float noise.
+# With matched batches the two boots agreed bitwise over 320 positions, so the
+# bound is not expected to be reached.
 LOGPROB_TOL = 0.5
-# Near-tie greedy/seeded flips from boot-state noise; measured <= 1 per
-# engine boot across ~16 completions, plus one of headroom.
-MAX_DIVERGENT_PROMPTS = 3
 
 PROMPTS = [
     "The capital of France is",
@@ -54,6 +58,10 @@ def _is_sharded_sampling_active(worker) -> bool:
     return worker.model_runner.batch_sharder is not None
 
 
+def _sampling_params() -> list[SamplingParams]:
+    return [SAMPLING_PARAMS[i % len(SAMPLING_PARAMS)] for i in range(len(PROMPTS))]
+
+
 def _generate(monkeypatch: pytest.MonkeyPatch, disable_sharding: bool):
     llm = LLM(
         model=MODEL,
@@ -75,9 +83,7 @@ def _generate(monkeypatch: pytest.MonkeyPatch, disable_sharding: bool):
         # intended sampling mode.
         modes = llm.llm_engine.collective_rpc(_is_sharded_sampling_active)
         assert all(mode == (not disable_sharding) for mode in modes), modes
-        params = [
-            SAMPLING_PARAMS[i % len(SAMPLING_PARAMS)] for i in range(len(PROMPTS))
-        ]
+        params = _sampling_params()
         # Two waves: the second reuses the request slots freed by the first
         # (in whatever order requests finished), covering the cross-rank
         # slot-recycling path that request ownership derives from.
@@ -89,13 +95,36 @@ def _generate(monkeypatch: pytest.MonkeyPatch, disable_sharding: bool):
         cleanup_dist_env_and_memory()
 
 
-def _assert_logprob_dicts_close(ref_lps, out_lps, what: str) -> None:
-    # Allow one boundary entry of the top-k set to swap at a near-tie.
-    common = set(ref_lps) & set(out_lps)
-    assert len(common) >= max(len(ref_lps), len(out_lps)) - 1, (
-        f"{what}: top-k token sets diverge: {sorted(ref_lps)} vs {sorted(out_lps)}"
+def _requested_top_k(lps: LogprobsOnePosition, num_logprobs: int) -> set[int]:
+    """Exclude the extra sampled/prompt token, including at rank k+1."""
+    top = {
+        token_id
+        for token_id, logprob in lps.items()
+        if logprob.rank is not None and 1 <= logprob.rank <= num_logprobs
+    }
+    assert len(top) == num_logprobs, (
+        f"expected {num_logprobs} ranked top-k entries, got {len(top)}"
+    )
+    return top
+
+
+def _assert_logprob_dicts_close(
+    ref_lps: LogprobsOnePosition,
+    out_lps: LogprobsOnePosition,
+    what: str,
+    num_logprobs: int,
+) -> None:
+    # Allow one boundary entry of the top-k set to swap at a near-tie. Compare
+    # the top-k alone: counting the sampled token would spend that allowance
+    # on the extra entry itself whenever only one side carries one.
+    ref_top = _requested_top_k(ref_lps, num_logprobs)
+    out_top = _requested_top_k(out_lps, num_logprobs)
+    shared_top = ref_top & out_top
+    assert len(shared_top) >= num_logprobs - 1, (
+        f"{what}: top-k token sets diverge: {sorted(ref_top)} vs {sorted(out_top)}"
     )
-    for token_id in common:
+    # Values are comparable for every shared token, sampled one included.
+    for token_id in set(ref_lps) & set(out_lps):
         diff = abs(ref_lps[token_id].logprob - out_lps[token_id].logprob)
         assert diff <= LOGPROB_TOL, (
             f"{what}[{token_id}]: logprob diff {diff} "
@@ -103,20 +132,72 @@ def _assert_logprob_dicts_close(ref_lps, out_lps, what: str) -> None:
         )
 
 
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    ("num_logprobs", "ref_ranks", "out_ranks"),
+    [
+        (2, {0: 1, 1: 2, 2: 3}, {0: 1, 3: 2, 2: 4}),
+        (1, {0: 1, 1: 2}, {2: 1, 1: 3}),
+        (2, {0: 1, 1: 2}, {0: 1, 2: 2, 1: 3}),
+        (2, {0: 1, 1: 2, 2: 4}, {0: 1, 1: 2, 2: 5}),
+        (0, {0: 2}, {0: 3}),
+    ],
+    ids=["decode-boundary", "prompt-boundary", "sampled-in-top-k", "gap", "zero"],
+)
+def test_logprob_comparison_allows_sampled_token_rank_changes(
+    num_logprobs: int, ref_ranks: dict[int, int], out_ranks: dict[int, int]
+):
+    ref_lps, out_lps = (
+        {token: Logprob(-2.0 - rank / 100, rank=rank) for token, rank in ranks.items()}
+        for ranks in (ref_ranks, out_ranks)
+    )
+    _assert_logprob_dicts_close(ref_lps, out_lps, "boundary swap", num_logprobs)
+
+
+@pytest.mark.cpu_test
+def test_logprob_comparison_rejects_multiple_top_k_swaps():
+    ref = {0: Logprob(-2.0, rank=1), 1: Logprob(-2.01, rank=2)}
+    out = {2: Logprob(-2.0, rank=1), 3: Logprob(-2.01, rank=2)}
+    with pytest.raises(AssertionError, match="top-k token sets diverge"):
+        _assert_logprob_dicts_close(ref, out, "misrouted top-k", num_logprobs=2)
+
+
+@pytest.mark.cpu_test
+def test_logprob_comparison_checks_extra_sampled_token_value():
+    ref = {0: Logprob(-1.0, rank=1), 1: Logprob(-3.0, rank=3)}
+    out = {0: Logprob(-1.0, rank=1), 1: Logprob(-4.0, rank=3)}
+    with pytest.raises(AssertionError, match="logprob diff"):
+        _assert_logprob_dicts_close(ref, out, "sampled token", num_logprobs=1)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("rank", [None, 2])
+def test_logprob_comparison_rejects_missing_top_k_rank(rank: int | None):
+    lps = {0: Logprob(-1.0, rank=rank)}
+    with pytest.raises(AssertionError, match="expected 1 ranked top-k entries"):
+        _assert_logprob_dicts_close(lps, lps, "missing rank", num_logprobs=1)
+
+
 @multi_gpu_test(num_gpus=2)
 def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch):
     """Generation, logprobs, and prompt logprobs match between batch-sharded
-    sampling and the replicated fallback, up to measured boot-state noise."""
+    sampling and the replicated fallback."""
     monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
     # Required for the collective_rpc mode check below.
     monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    # Queue every request before the first step, so both boots schedule the
+    # same batches -- see the module docstring.
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
 
     ref_outputs = _generate(monkeypatch, disable_sharding=True)
     shard_outputs = _generate(monkeypatch, disable_sharding=False)
 
     assert len(ref_outputs) == len(shard_outputs) == 2 * len(PROMPTS)
-    num_divergent = 0
-    for i, (ref, out) in enumerate(zip(ref_outputs, shard_outputs)):
+    params = _sampling_params()
+    divergent = []
+    for i, (ref, out, sampling_params) in enumerate(
+        zip(ref_outputs, shard_outputs, params + params[::-1])
+    ):
         ref_completion = ref.outputs[0]
         out_completion = out.outputs[0]
         ref_ids = list(ref_completion.token_ids)
@@ -131,21 +212,24 @@ def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch):
                 break
             prefix += 1
         if ref_ids != out_ids:
-            num_divergent += 1
+            divergent.append((i, prefix))
 
         if ref_completion.logprobs is not None:
             assert out_completion.logprobs is not None
+            assert sampling_params.logprobs is not None
             comparable = min(prefix + 1, len(ref_ids), len(out_ids))
             for pos in range(comparable):
                 _assert_logprob_dicts_close(
                     ref_completion.logprobs[pos],
                     out_completion.logprobs[pos],
                     f"prompt {i} logprobs[{pos}]",
+                    sampling_params.logprobs,
                 )
 
         # The prompt is fixed, so every prompt-logprob position is comparable.
         if ref.prompt_logprobs is not None:
             assert out.prompt_logprobs is not None
+            assert sampling_params.prompt_logprobs is not None
             for pos, (ref_lps, out_lps) in enumerate(
                 zip(ref.prompt_logprobs, out.prompt_logprobs)
             ):
@@ -153,10 +237,14 @@ def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch):
                     assert ref_lps is None and out_lps is None
                     continue
                 _assert_logprob_dicts_close(
-                    ref_lps, out_lps, f"prompt {i} prompt_logprobs[{pos}]"
+                    ref_lps,
+                    out_lps,
+                    f"prompt {i} prompt_logprobs[{pos}]",
+                    sampling_params.prompt_logprobs,
                 )
 
-    assert num_divergent <= MAX_DIVERGENT_PROMPTS, (
-        f"{num_divergent}/{2 * len(PROMPTS)} prompts diverged: beyond near-tie "
-        "boot-state noise; sharded sampling is likely misrouting requests"
+    assert not divergent, (
+        f"{len(divergent)}/{2 * len(PROMPTS)} completions diverged at "
+        f"(prompt, position) {divergent}; the sampling modes did not agree "
+        "token for token with in-process scheduling"
     )
diff --git a/tests/v1/e2e/general/test_streaming_input.py b/tests/v1/e2e/general/test_streaming_input.py
index 1954ce6a7dcf..c71e7ad0a6b5 100644
--- a/tests/v1/e2e/general/test_streaming_input.py
+++ b/tests/v1/e2e/general/test_streaming_input.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-End-to-end tests for the streaming input feature in AsyncLLM.
+"""End-to-end tests for the streaming input feature in AsyncLLM.
 
 These tests verify that:
 1. Streaming inputs work correctly with bunched inputs (queued)
diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py
index 2b1800f78234..0dff04333602 100644
--- a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py
+++ b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py
@@ -26,6 +26,8 @@ class DFlashCorrectnessConfig:
     use_chat_completions: bool = False
     enforce_eager: bool = False
     disable_flashinfer_sampler: bool = False
+    language_model_only: bool = False
+    chat_template_kwargs: dict[str, object] | None = None
 
 
 QWEN3_DFLASH = DFlashCorrectnessConfig(
@@ -35,11 +37,23 @@ class DFlashCorrectnessConfig:
     expected_acceptance_len=3.5,
 )
 
+QWEN3_8_DFLASH2_NVFP4 = DFlashCorrectnessConfig(
+    model="nvidia/Qwen3.8-27B-NVFP4",
+    draft_model="z-lab/Qwen3.8-27B-DFlash2",
+    expected_accuracy=0.90,
+    expected_acceptance_len=6.14 * 0.95,
+    num_speculative_tokens=7,
+    max_num_seqs=32,
+    use_chat_completions=True,
+    language_model_only=True,
+    chat_template_kwargs={"enable_thinking": False},
+)
+
 LAGUNA_DFLASH_NVFP4 = DFlashCorrectnessConfig(
     model="poolside/Laguna-XS-2.1-NVFP4",
     draft_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",
     expected_accuracy=0.7,  # Standard GSM8K sanity floor.
-    expected_acceptance_len=3.55 * 0.9,
+    expected_acceptance_len=3.05 * 0.9,
     num_speculative_tokens=15,
     max_model_len=8192,
     max_num_seqs=32,
@@ -104,6 +118,11 @@ def test_dflash_reference_acceptance_lengths(
             True,
             id="laguna-nvfp4-mrv2",
         ),
+        pytest.param(
+            QWEN3_8_DFLASH2_NVFP4,
+            True,
+            id="qwen3.8-dflash2-nvfp4-mrv2",
+        ),
     ],
 )
 def test_dflash_correctness(
@@ -134,12 +153,14 @@ def test_dflash_correctness(
         disable_log_stats=False,
         enable_chunked_prefill=None,
         compilation_config=CompilationConfig(),
+        language_model_only=config.language_model_only,
     ) as spec_runner:
         spec_llm = spec_runner.llm
         results = evaluate_gsm8k_offline(
             spec_llm,
             num_questions=config.num_questions,
             use_chat_completions=config.use_chat_completions,
+            chat_template_kwargs=config.chat_template_kwargs,
         )
         accuracy = results["accuracy"]
         acceptance_len = compute_acceptance_len(spec_llm.get_metrics())
diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/utils.py b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py
index 97aa82fb5456..659d43daee89 100644
--- a/tests/v1/e2e/spec_decode/acceptance_rates/utils.py
+++ b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py
@@ -43,8 +43,7 @@ def run_acceptance_length_eval(
     chat_template_kwargs: dict[str, Any],
     use_mrv2: bool,
 ):
-    """
-    E2E acceptance-rate validation for speculative decoding.
+    """E2E acceptance-rate validation for speculative decoding.
 
     Drives one or more datasets (keyed in ``expected_acceptance_lengths``)
     through the spec decode engine and asserts the mean acceptance length
diff --git a/tests/v1/e2e/spec_decode/draft_model/test_async.py b/tests/v1/e2e/spec_decode/draft_model/test_async.py
index d7f932b0c256..ebb471ea711b 100644
--- a/tests/v1/e2e/spec_decode/draft_model/test_async.py
+++ b/tests/v1/e2e/spec_decode/draft_model/test_async.py
@@ -1,70 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test that verifies no implicit GPU-CPU synchronization occurs during
-speculative decoding generation under expected conditions.
-"""
-
-import multiprocessing
-import sys
-import traceback
+"""Tests for sync-free speculative decoding with async scheduling."""
 
 import pytest
-import torch
-
-
-@pytest.fixture
-def sync_tracker():
-    """
-    Fixture that patches CommonAttentionMetadata.seq_lens_cpu to detect
-    lazy init syncs. Prints stack traces immediately when syncs occur.
-    """
-    from vllm.v1.attention.backend import CommonAttentionMetadata
-
-    # Shared counter for cross-process communication (inherited by fork)
-    sync_count = multiprocessing.Value("i", 0)
-
-    # Save original property
-    original_prop = CommonAttentionMetadata.seq_lens_cpu
-    original_fget = original_prop.fget
-
-    # Create tracking wrapper
-    def tracking_seq_lens_cpu(self):
-        if self._seq_lens_cpu is None:
-            # Increment counter
-            with sync_count.get_lock():
-                sync_count.value += 1
-                count = sync_count.value
-            # Print stack trace immediately (shows in subprocess output)
-            print(f"\n{'=' * 60}", file=sys.stderr)
-            print(f"SYNC #{count}: seq_lens_cpu lazy init triggered!", file=sys.stderr)
-            print(f"{'=' * 60}", file=sys.stderr)
-            traceback.print_stack(file=sys.stderr)
-            print(f"{'=' * 60}\n", file=sys.stderr)
-            sys.stderr.flush()
-        return original_fget(self)
-
-    # Apply patch
-    CommonAttentionMetadata.seq_lens_cpu = property(tracking_seq_lens_cpu)
-
-    class SyncTracker:
-        @property
-        def count(self) -> int:
-            return sync_count.value
-
-        def assert_no_sync(self, msg: str = ""):
-            count = sync_count.value
-            assert count == 0, (
-                f"Unexpected GPU-CPU sync: seq_lens_cpu lazy init triggered "
-                f"{count} times. See stack traces above. {msg}"
-            )
-
-    yield SyncTracker()
-
-    # Restore original property
-    CommonAttentionMetadata.seq_lens_cpu = original_prop
-    torch._dynamo.reset()
-
 
 # Test configurations: (model, spec_model, method, num_spec_tokens, backend_env)
 SPEC_DECODE_CONFIGS = [
@@ -97,18 +35,16 @@ def assert_no_sync(self, msg: str = ""):
     SPEC_DECODE_CONFIGS,
 )
 def test_no_sync_with_spec_decode(
-    sync_tracker,
     model: str,
     spec_model: str,
     method: str,
     num_spec_tokens: int,
+    monkeypatch,
     vllm_runner,
 ):
-    """
-    Test that no implicit GPU-CPU sync occurs during speculative decoding
-    generation.
-    """
-    # Import vLLM AFTER sync_tracker fixture has applied the patch
+    """Test generation while the worker rejects unintended GPU-CPU syncs."""
+    monkeypatch.setenv("VLLM_GPU_SYNC_CHECK", "error")
+
     from vllm import SamplingParams
     from vllm.config import CompilationConfig
 
@@ -158,5 +94,3 @@ def test_no_sync_with_spec_decode(
         assert outputs[0].outputs[0].text, (
             f"{method} target={model}: generated completion text is empty"
         )
-
-    sync_tracker.assert_no_sync()
diff --git a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
index cf7b44fb226a..3a474ac2039c 100644
--- a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
+++ b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
@@ -293,7 +293,6 @@ def test_draft_model_moe_backend_default_auto():
 def test_draft_model_engine_args_rejects_invalid_tp_argname():
     """The user should pass "draft_tensor_parallel_size" rather than
     "tensor_parallel_size". We enforce this with validation."""
-
     engine_args = EngineArgs(
         model="Qwen/Qwen3-1.7B",
         tensor_parallel_size=1,
diff --git a/tests/v1/e2e/spec_decode/draft_model/test_lora.py b/tests/v1/e2e/spec_decode/draft_model/test_lora.py
index 460ac4e596d9..36a03360c7bd 100644
--- a/tests/v1/e2e/spec_decode/draft_model/test_lora.py
+++ b/tests/v1/e2e/spec_decode/draft_model/test_lora.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This script contains:
+"""This script contains:
 1. test lora with speculative decoding for batch inference
 """
 
@@ -59,8 +58,7 @@ def test_batch_inference_correctness(
     model_setup: tuple[str, str, str, str, int],
     vllm_runner,
 ):
-    """
-    Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora.
+    """Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora.
     Should be the same and no failure when doing batch inference.
     model_setup: (method, model_name, spec_model_name, lora_path, tp_size)
     """
diff --git a/tests/v1/e2e/spec_decode/eagle/utils.py b/tests/v1/e2e/spec_decode/eagle/utils.py
index 29127df85f5c..918a06e32c08 100644
--- a/tests/v1/e2e/spec_decode/eagle/utils.py
+++ b/tests/v1/e2e/spec_decode/eagle/utils.py
@@ -28,8 +28,7 @@ def _run_eagle_correctness(
     attn_backend: str,
     vllm_runner,
 ):
-    """
-    Compare the outputs of an original LLM and a speculative LLM
+    """Compare the outputs of an original LLM and a speculative LLM
     which should be the same when using eagle speculative decoding.
     """
     method, model_name, spec_model_name, tp_size = model_setup
diff --git a/tests/v1/e2e/spec_decode/mtp/test_mtp.py b/tests/v1/e2e/spec_decode/mtp/test_mtp.py
index acd3285bd5fa..92f4a2a773f6 100644
--- a/tests/v1/e2e/spec_decode/mtp/test_mtp.py
+++ b/tests/v1/e2e/spec_decode/mtp/test_mtp.py
@@ -157,8 +157,7 @@ def test_mtp_correctness(
     expected_accuracy_threshold: float,
     vllm_runner,
 ):
-    """
-    Compare the outputs of a original LLM and a speculative LLM
+    """Compare the outputs of a original LLM and a speculative LLM
     which should be the same when using MTP speculative decoding. Due to some variance
     in the engine, it is possible for some outputs to differ, so we expect that at least
     6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed
diff --git a/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py
index d3627a8703b3..c204cc1c012d 100644
--- a/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py
+++ b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py
@@ -64,8 +64,7 @@ def test_ngram_gpu_default_with_async_scheduling(
     async_scheduling: bool,
     vllm_runner,
 ):
-    """
-    Test ngram_gpu speculative decoding (k=3) correctness with and without
+    """Test ngram_gpu speculative decoding (k=3) correctness with and without
     async scheduling, validated via GSM8K accuracy.
     Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%).
     """
@@ -99,8 +98,7 @@ def test_suffix_decoding_acceptance(
     model_name: str,
     vllm_runner,
 ):
-    """
-    Check that suffix decoding caching takes effect and improves acceptance
+    """Check that suffix decoding caching takes effect and improves acceptance
     lengths and acceptance rates over multiple runs of the same prompts.
     """
     test_prompts = get_test_prompts(mm_enabled=False)
diff --git a/tests/v1/e2e/spec_decode/speculators/test_speculators.py b/tests/v1/e2e/spec_decode/speculators/test_speculators.py
index 990f6190751a..c6acc5573137 100644
--- a/tests/v1/e2e/spec_decode/speculators/test_speculators.py
+++ b/tests/v1/e2e/spec_decode/speculators/test_speculators.py
@@ -32,8 +32,7 @@ def test_speculators_model_integration(
     expected_accuracy_threshold: float,
     vllm_runner,
 ):
-    """
-    Test that speculators models work with the simplified integration.
+    """Test that speculators models work with the simplified integration.
 
     This verifies the `vllm serve <speculator-model>` use case where
     speculative config is automatically detected from the model config
diff --git a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
index 236997261675..e7b2c36c4fcd 100644
--- a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
+++ b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
@@ -56,14 +56,7 @@ class InlineConfig:
         enable_expert_parallel=True,
         enable_eplb=True,
     ),
-    InlineConfig(
-        id="pp2",
-        pp_size=2,
-        skip_reason=(
-            "DeepSeek MTP pipeline-parallel support is in flight upstream; "
-            "see https://github.com/vllm-project/vllm/pull/44698"
-        ),
-    ),
+    InlineConfig(id="pp2", pp_size=2),
 ]
 
 
diff --git a/tests/v1/e2e/spec_decode/test_sharded_sampling.py b/tests/v1/e2e/spec_decode/test_sharded_sampling.py
index cc9781ca9fdb..8f58b21a1a03 100644
--- a/tests/v1/e2e/spec_decode/test_sharded_sampling.py
+++ b/tests/v1/e2e/spec_decode/test_sharded_sampling.py
@@ -39,6 +39,9 @@ def test_mtp_sharded_sampling_equivalence(monkeypatch: pytest.MonkeyPatch):
         m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
         # Required for the collective_rpc mode check below.
         m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+        # Queue every request before the first step, so both boots schedule
+        # the same batches -- see tests/v1/e2e/general/test_sharded_sampling.py.
+        m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
 
         def run(disable_sharding: bool):
             llm = LLM(
@@ -79,24 +82,22 @@ def run(disable_sharding: bool):
         ref_greedy, ref_seeded = run(disable_sharding=True)
         shard_greedy, shard_seeded = run(disable_sharding=False)
 
-        # Engine boots are not bitwise deterministic (kernel/collective
-        # selection shifts logits by ~1 ulp, flipping near-tie tokens), so
-        # allow a small number of divergent prompts. A sharding bug produces
-        # wholesale divergence, not isolated near-tie flips.
+        # A sharding bug shows up as wholesale divergence, not isolated flips.
         for name, ref_outputs, shard_outputs in (
             ("greedy", ref_greedy, shard_greedy),
             ("seeded", ref_seeded, shard_seeded),
         ):
-            num_divergent = 0
+            divergent = []
             for i, (ref, out) in enumerate(zip(ref_outputs, shard_outputs)):
                 if list(ref.outputs[0].token_ids) != list(out.outputs[0].token_ids):
-                    num_divergent += 1
+                    divergent.append(i)
                     print(
                         f"{name} prompt {i} diverged:\n"
                         f"  replicated: {ref.outputs[0].text!r}\n"
                         f"  sharded:    {out.outputs[0].text!r}"
                     )
-            assert num_divergent <= 2, (
-                f"{name}: {num_divergent}/{len(ref_outputs)} prompts diverged, "
-                "beyond near-tie boot noise"
+            assert not divergent, (
+                f"{name}: {len(divergent)}/{len(ref_outputs)} prompts diverged "
+                f"at {divergent}; the sampling modes did not agree token for "
+                "token with in-process scheduling"
             )
diff --git a/tests/v1/e2e/test_cpu_spec_decode.py b/tests/v1/e2e/test_cpu_spec_decode.py
new file mode 100644
index 000000000000..a7ba5580c33b
--- /dev/null
+++ b/tests/v1/e2e/test_cpu_spec_decode.py
@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU speculative-decoding end-to-end correctness."""
+
+import os
+
+import pytest
+
+from tests.models.utils import check_logprobs_close
+from tests.v1.e2e.spec_decode.utils import compute_acceptance_len
+from vllm import LLM, SamplingParams
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
+
+pytestmark = pytest.mark.cpu_model
+
+if not current_platform.is_cpu():
+    pytest.skip("skipping CPU-only tests", allow_module_level=True)
+
+# Bound the KV cache so the run does not scale with host memory.
+os.environ.setdefault("VLLM_CPU_KVCACHE_SPACE", "1")
+
+MODEL = "Qwen/Qwen3-0.6B"
+NUM_LOGPROBS = 5
+SP = SamplingParams(max_tokens=48, temperature=0, logprobs=NUM_LOGPROBS)
+
+# Repetitive prompts so the ngram proposer finds matches and actually drafts;
+# without hits the run is indistinguishable from plain decoding.
+PROMPTS = [
+    "The capital of France is Paris. The capital of Italy is Rome. "
+    "The capital of Spain is Madrid. The capital of France is",
+    "def add(a, b):\n    return a + b\n\ndef sub(a, b):\n    return a - b\n\n"
+    "def mul(a, b):\n    return",
+    "one two three four five one two three four five one two three four",
+]
+
+
+def _make_llm(**overrides) -> LLM:
+    base = dict(
+        model=MODEL,
+        dtype="bfloat16",
+        max_model_len=1024,
+        disable_log_stats=False,
+    )
+    base.update(overrides)
+    return LLM(**base)
+
+
+def _tuples(outputs) -> list[tuple[list[int], str, object]]:
+    return [
+        (list(o.outputs[0].token_ids), o.outputs[0].text, o.outputs[0].logprobs)
+        for o in outputs
+    ]
+
+
+@pytest.fixture(scope="module")
+def baseline_refs():
+    llm = _make_llm()
+    refs = _tuples(llm.generate(PROMPTS, SP))
+    del llm
+    return refs
+
+
+def _spec_metric(llm: LLM, name: str) -> float:
+    for metric in llm.get_metrics():
+        if metric.name == name:
+            return float(getattr(metric, "value", 0.0))
+    return 0.0
+
+
+def test_ngram_spec_decode_matches_baseline(baseline_refs):
+    """Greedy output must be unchanged by speculative decoding.
+
+    The drafted-token check keeps the test from passing vacuously when the
+    proposer silently produces nothing.
+    """
+    llm = _make_llm(
+        speculative_config={
+            "method": "ngram",
+            "prompt_lookup_max": 5,
+            "prompt_lookup_min": 3,
+            "num_speculative_tokens": 3,
+        },
+    )
+    got = _tuples(llm.generate(PROMPTS, SP))
+    drafted = _spec_metric(llm, "vllm:spec_decode_num_draft_tokens")
+    accepted = _spec_metric(llm, "vllm:spec_decode_num_accepted_tokens")
+    del llm
+
+    assert drafted > 0, "ngram proposer drafted no tokens; the test would be vacuous"
+    assert accepted > 0, f"no draft tokens accepted out of {drafted} drafted"
+
+    check_logprobs_close(
+        outputs_0_lst=baseline_refs,
+        outputs_1_lst=got,
+        name_0="no_spec_decode",
+        name_1="ngram_spec_decode",
+    )
+
+
+SMOKE_PROMPTS = [
+    "The capital of France is",
+    "2 + 2 equals",
+    "In one word, the color of the sky is",
+    "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:",
+]
+SMOKE_SP = SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True)
+
+# Methods routed to the V2 runner, using the same target/draft pairs as the GPU
+# suite. MTP needs no separate draft checkpoint: its speculator layer is
+# trained jointly with, and shipped inside, the target model's own weights.
+SMOKE_CONFIGS = [
+    pytest.param(
+        "meta-llama/Llama-3.2-1B-Instruct",
+        {
+            "method": "eagle3",
+            "model": "nm-testing/Llama3_2_1B_speculator.eagle3",
+            "num_speculative_tokens": 3,
+        },
+        id="eagle3",
+    ),
+    pytest.param(
+        "Qwen/Qwen3-4B",
+        {
+            "method": "dflash",
+            "model": "z-lab/Qwen3-4B-DFlash-b16",
+            # The checkpoint allows 16, but accepted tokens plateau around 64
+            # regardless, so a larger draft only burns CPU on rejected tokens.
+            "num_speculative_tokens": 4,
+        },
+        id="dflash",
+    ),
+    pytest.param(
+        "Qwen/Qwen3.5-0.8B",
+        {
+            "method": "mtp",
+            "num_speculative_tokens": 3,
+        },
+        id="mtp",
+    ),
+]
+
+# Measured 1.85 (eagle3), 1.99 (dflash), 2.29 (mtp) on CPU; this floor leaves
+# room for numeric drift while still failing if speculation stops paying off.
+MIN_ACCEPTANCE_LEN = 1.5
+
+
+@pytest.mark.skipif(
+    not HAS_TRITON, reason="the V2 runner needs triton-cpu, which CI builds"
+)
+@pytest.mark.parametrize("model,speculative_config", SMOKE_CONFIGS)
+def test_v2_speculator_smoke(model: str, speculative_config: dict):
+    """Gates that drafts are accepted often enough for speculation to pay off."""
+    llm = _make_llm(model=model, speculative_config=speculative_config)
+    try:
+        assert llm.llm_engine.vllm_config.use_v2_model_runner, (
+            f"{speculative_config['method']} fell back to the V1 runner; "
+            "this test is meant to cover the V2 path"
+        )
+        llm.generate(SMOKE_PROMPTS, SMOKE_SP)
+        acceptance = compute_acceptance_len(llm.get_metrics())
+        assert acceptance >= MIN_ACCEPTANCE_LEN, (
+            f"acceptance length {acceptance:.3f} below {MIN_ACCEPTANCE_LEN}; "
+            "drafts are being rejected far more than expected"
+        )
+    finally:
+        del llm
diff --git a/tests/v1/e2e/test_hybrid_chunked_prefill.py b/tests/v1/e2e/test_hybrid_chunked_prefill.py
index dd24f3884ea2..66c8d5b9b76c 100644
--- a/tests/v1/e2e/test_hybrid_chunked_prefill.py
+++ b/tests/v1/e2e/test_hybrid_chunked_prefill.py
@@ -68,7 +68,6 @@ def test_mtp_speculative_mixed_batch_short_prefill(
 ):
     """Test to ensure MTP speculative decoding correctly handles
     short prefill chunks that fall below the reorder_batch_threshold."""
-
     # Set so large that both prefills will be classified as decodes in a mixed batch
     # note, with prefix caching we require chunk_size >= mamba_block_size
     chunk_size = 256 if not enable_prefix_caching else 2048
diff --git a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh b/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
index 05af1bb6fdbc..c4e9b0d06975 100755
--- a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
+++ b/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
@@ -20,6 +20,8 @@
 #   SKIP_BASELINE               set to 1 to reuse existing BASELINE_FILE
 #   CONCURRENCY / REPEAT        concurrent requests / rounds (defaults 3 / 2)
 #   MAX_MODEL_LEN               context length (default 16384)
+#   E_CUDAGRAPH_MM_ENCODER       enable encoder CUDA graphs (default 0)
+#   E_MM_ENCODER_ATTN_BACKEND    optional encoder attention backend override
 
 set -euo pipefail
 
@@ -35,6 +37,15 @@ if [[ "$USE_MM_PROMPTS" == "1" ]]; then
   TEST_ARGS+=(--use_mm_prompts --mm_smoke_test)
 fi
 MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
+E_ATTN_ARGS=()
+if [[ -n "${E_MM_ENCODER_ATTN_BACKEND:-}" ]]; then
+  E_ATTN_ARGS+=(--mm-encoder-attn-backend "$E_MM_ENCODER_ATTN_BACKEND")
+fi
+case "${E_CUDAGRAPH_MM_ENCODER:-0}" in
+  0) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":false}' ;;
+  1) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":true}' ;;
+  *) echo "E_CUDAGRAPH_MM_ENCODER must be 0 or 1"; exit 1 ;;
+esac
 
 GPU_SINGLE="${GPU_SINGLE:-0}"
 GPU_E="${GPU_E:-0}"
@@ -203,7 +214,7 @@ run_epd_mooncake() {
   cleanup_instances
 
   echo "Starting ENCODER on GPU $GPU_E port $ENCODE_PORT"
-  CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \
+  VLLM_SERVER_DEV_MODE=1 CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \
     --port "$ENCODE_PORT" \
     --gpu-memory-utilization 0.35 \
     --mm-tensor-ipc torch_shm \
@@ -215,6 +226,9 @@ run_epd_mooncake() {
     --limit-mm-per-prompt '{"image":2,"video":0}' \
     --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \
     --ec-transfer-config "$ENC_EC_JSON" \
+    --compilation-config "$E_COMPILATION_CONFIG" \
+    "${E_ATTN_ARGS[@]}" \
+    --worker-extension-cls tests.v1.ec_connector.integration.test_epd_correctness.EncoderGraphTestWorkerExtension \
     >"${LOG_PATH}/mooncake_epd_encoder.log" 2>&1 &
   local ENCODER_PID=$!
   PIDS+=("$ENCODER_PID")
@@ -266,6 +280,26 @@ run_epd_mooncake() {
     --baseline_file "$BASELINE_FILE" \
     "${TEST_ARGS[@]}"
 
+  curl -fsS "http://localhost:${ENCODE_PORT}/collective_rpc" \
+    -H 'Content-Type: application/json' \
+    -d '{"method":"encoder_graph_stats"}' \
+    >"${LOG_PATH}/encoder_graph_stats.json"
+  "$PYTHON_BIN" - "${LOG_PATH}/encoder_graph_stats.json" "${E_CUDAGRAPH_MM_ENCODER:-0}" <<'PY'
+import json
+import sys
+
+with open(sys.argv[1]) as f:
+    results = json.load(f)["results"]
+assert results, "No encoder workers returned graph statistics"
+for stats in results:
+    if sys.argv[2] == "1":
+        assert stats["is_captured"], stats
+        assert stats["graph_hits"] > 0, stats
+    else:
+        assert stats["graph_hits"] == 0, stats
+print("Encoder graph statistics:", results)
+PY
+
   cleanup_instances
 }
 
diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py
index d3ae0f1d8792..4e37637ea452 100644
--- a/tests/v1/ec_connector/integration/test_epd_correctness.py
+++ b/tests/v1/ec_connector/integration/test_epd_correctness.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-EPD Correctness Test
+"""EPD Correctness Test.
 
 Tests that EPD (Encoder-Prefill-Decode) disaggregation produces the same
 outputs as a baseline single instance.
@@ -27,6 +26,7 @@
 import os
 import time
 from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING
 
 import openai
 import requests
@@ -34,6 +34,26 @@
 from vllm.assets.image import ImageAsset
 from vllm.multimodal.utils import encode_image_url
 
+if TYPE_CHECKING:
+    from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+
+
+class EncoderGraphTestWorkerExtension:
+    """Expose encoder replay evidence through the existing test RPC endpoint."""
+
+    model_runner: "MMEncoderModelRunner"
+
+    def encoder_graph_stats(self):
+        from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
+
+        runner = self.model_runner
+        assert isinstance(runner, MMEncoderModelRunner)
+        manager = runner.model_state.encoder_runner.cudagraph_manager
+        if manager is None:
+            return {"is_captured": False, "graph_hits": 0, "graph_misses": 0}
+        return {"is_captured": manager.is_captured(), **manager.get_cumulative_stats()}
+
+
 MAX_OUTPUT_LEN = 256
 
 # Sample prompts with multimodal content
@@ -104,6 +124,7 @@ def check_vllm_server(url: str, timeout=5, retries=10) -> bool:
 
     Returns:
         True if the server is ready, False otherwise
+
     """
     for attempt in range(retries):
         try:
@@ -138,6 +159,7 @@ def run_chat_completion(
 
     Returns:
         Generated text content
+
     """
     with openai.OpenAI(
         api_key="EMPTY", base_url=base_url, timeout=120, max_retries=0
diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py
index 722c0d89fa21..c168970147b0 100644
--- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py
+++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py
@@ -13,6 +13,8 @@
 import pytest
 import torch
 
+from vllm.config import ParallelConfig
+from vllm.distributed.ec_transfer.ec_connector.cpu import common
 from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
     ECSharedRegion,
     _wait_for_file_size,
@@ -95,6 +97,32 @@ def test_second_instance_opens_existing_file_and_shares_memory():
         r1.cleanup()
 
 
+def test_dense_dp_engines_get_distinct_regions(monkeypatch):
+    """Dense DP engines reset data_parallel_rank to 0 but share instance_id,
+    so the region must be keyed by data_parallel_index."""
+    monkeypatch.setattr(common, "_get_encoder_cache_hidden_dim", lambda cfg: 8)
+    instance_id = str(uuid.uuid4())
+    regions = []
+    try:
+        for dp_index in range(2):
+            parallel_config = ParallelConfig(
+                data_parallel_size=2, data_parallel_rank=dp_index
+            )
+            parallel_config.reconfigure_for_independent_dp_rank()
+            cfg = MagicMock()
+            cfg.instance_id = instance_id
+            cfg.parallel_config = parallel_config
+            cfg.model_config.dtype = torch.float16
+            cfg.ec_transfer_config.ec_connector_extra_config = {"ec_cpu_bytes": 64}
+            regions.append(common.create_ec_shared_region(cfg))
+
+        assert all(r._is_creator for r in regions)
+        assert regions[0]._mmap_path != regions[1]._mmap_path
+    finally:
+        for r in regions:
+            r.cleanup()
+
+
 def test_only_creator_unlinks_file_on_cleanup():
     """Critical contract: if the non-creator unlinks, the creator's mmap path
     becomes a dangling backing file and a third opener would create a new one
diff --git a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py
new file mode 100644
index 000000000000..26155d8b6e64
--- /dev/null
+++ b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py
@@ -0,0 +1,726 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import ctypes
+import gc
+import struct
+import threading
+import weakref
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
+from vllm.config import MultiModalConfig
+from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
+from vllm.distributed.ec_transfer.ec_connector.mooncake.config import MooncakeECConfig
+from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import (
+    ECMooncakeConnectorMetadata,
+    ECMooncakePushSpec,
+)
+from vllm.distributed.ec_transfer.ec_connector.mooncake_ec_connector import (
+    ECMooncakeConnector,
+)
+from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import (
+    backend as store_backend,
+)
+from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import (
+    store_client,
+)
+from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
+    TensorSpec,
+    build_embedding_namespace,
+    make_embedding_key,
+)
+
+pytestmark = pytest.mark.cpu_test
+KEY = (
+    "test@embedding@model:model@revision:revision"
+    "@encoder:encoder@dtype:torch.float32@protocol:v3"
+)
+SPEC = TensorSpec((2, 2), "torch.float32", 16)
+
+
+def test_embedding_key_combines_namespace_and_identifier():
+    config = create_ec_vllm_config(ec_role="ec_producer")
+    config.ec_transfer_config.ec_connector_extra_config.update(
+        embedding_cache_prefix="team@cache", embedding_model_identity="shared/model"
+    )
+    config.model_config.revision = "rev:1"
+    config.model_config.dtype = torch.float32
+    config.model_config.multimodal_config = None
+    assert make_embedding_key(build_embedding_namespace(config), "image@1") == (
+        "team@cache@embedding@model:shared/model@revision:rev:1"
+        "@encoder:encoder:default@dtype:torch.float32@protocol:v3@id:image@1"
+    )
+
+
+def test_shared_reuse_requires_content_stable_identifiers():
+    """Encoder-only cache-off rendering uses process-local counter identifiers."""
+    config = create_ec_vllm_config(ec_role="ec_producer")
+    config.ec_transfer_config.ec_connector = "ECMooncakeConnector"
+    config.ec_transfer_config.ec_connector_extra_config["cross_encoder_cache"] = True
+    config.model_config.multimodal_config = MultiModalConfig(mm_processor_cache_gb=0)
+    config.use_v2_model_runner = True
+    config.lora_config = None
+    with pytest.raises(ValueError, match="mm_processor_cache_gb"):
+        MooncakeECConfig.from_vllm_config(config)
+
+
+class MemoryStore:
+    """Only the native boundary is faked; codec, resolver and publisher are real."""
+
+    def __init__(self):
+        self.objects: dict[str, bytes] = {}
+        self.registered: dict[int, int] = {}
+
+    def close(self):
+        return 0
+
+    def batch_is_exist(self, keys):
+        return [int(k in self.objects) for k in keys]
+
+    def register_buffer(self, addr, size):
+        self.registered[addr] = size
+        return 0
+
+    def unregister_buffer(self, addr):
+        del self.registered[addr]
+        return 0
+
+    def batch_put_from_multi_buffers(self, keys, pointers, sizes, config):
+        [key], [ptrs], [lengths] = keys, pointers, sizes
+        self.objects[key] = b"".join(
+            ctypes.string_at(p, n) for p, n in zip(ptrs, lengths, strict=True)
+        )
+        return [0]
+
+    def batch_get_into(self, keys, pointers, sizes):
+        results = []
+        for key, ptr, size in zip(keys, pointers, sizes, strict=True):
+            assert any(
+                base <= ptr and ptr + size <= base + capacity
+                for base, capacity in self.registered.items()
+            )
+            obj = self.objects.get(key)
+            if obj is None:
+                results.append(-704)
+            elif len(obj) > size:
+                results.append(-600)
+            else:
+                ctypes.memmove(ptr, obj, len(obj))
+                results.append(len(obj))
+        return results
+
+
+@pytest.fixture
+def backend():
+    instance = store_backend.MooncakeEmbeddingStoreBackend(
+        store_client.MooncakeEmbeddingStoreClient(
+            MemoryStore(), read_buffer_bytes=1024
+        ),
+        KEY,
+        max_pending_items=32,
+        max_pending_bytes=2 * 1024**3,
+    )
+    yield instance
+    # Some tests deliberately poison I/O; always stop the Python executor.
+    instance._executor.shutdown(wait=True)
+    if not instance.store_client._poisoned and (
+        instance.store_client._read_buffer is not None
+        or instance.store_client._put_header is not None
+    ):
+        instance.store_client.close()
+
+
+def test_tensor_round_trip(backend):
+    client, native = backend.store_client, backend.store_client.store
+    native.register_buffer = MagicMock(wraps=native.register_buffer)
+    header_ptr = None
+    for index, dtype in enumerate((torch.float16, torch.bfloat16, torch.float32)):
+        tensor = torch.arange(6, dtype=dtype).reshape(index + 1, -1)
+        key = make_embedding_key(KEY, str(dtype))
+        client.put_tensor(key, tensor)
+        pointer = ctypes.addressof(client._put_header)
+        header_ptr = header_ptr or pointer
+        assert pointer == header_ptr
+        expected = TensorSpec(tuple(tensor.shape), str(dtype), tensor.nbytes)
+        with torch.device("meta"):
+            loaded = client.load_tensors({key: expected}, "cpu")[key]
+        assert client._read_buffer.device.type == "cpu"
+        assert loaded.dtype == dtype and torch.equal(loaded, tensor)
+        assert len(native.objects[key]) == 24 + tensor.nbytes
+    assert (
+        sum(c.args[0] == header_ptr for c in native.register_buffer.call_args_list) == 1
+    )
+    assert len(native.registered) == 2
+
+
+@pytest.mark.parametrize("capacity", [80, 120])
+def test_batch_reads_reuse_registration_without_aliasing_outputs(backend, capacity):
+    client, native = backend.store_client, backend.store_client.store
+    client._read_buffer_bytes = capacity
+    tensors = {key: torch.full((2, 2), float(i)) for i, key in enumerate("abc")}
+    for key, tensor in tensors.items():
+        client.put_tensor(key, tensor)
+    native.register_buffer = MagicMock(wraps=native.register_buffer)
+    native.unregister_buffer = MagicMock(wraps=native.unregister_buffer)
+    native.batch_get_into = MagicMock(wraps=native.batch_get_into)
+    expected = dict.fromkeys(tensors, SPEC)
+    first = client.load_tensors(expected, "cpu")
+    for key, tensor in tensors.items():
+        native.objects[key] = native.objects[key][:24] + struct.pack(
+            "<4f", *[tensor[0, 0].item() + 10] * 4
+        )
+    second = client.load_tensors(expected, "cpu")
+    assert all(torch.equal(first[key], tensor) for key, tensor in tensors.items())
+    assert all(torch.equal(second[key], tensor + 10) for key, tensor in tensors.items())
+    assert native.batch_get_into.call_count == (4 if capacity == 80 else 2)
+    native.register_buffer.assert_called_once()
+    native.unregister_buffer.assert_not_called()
+    client.close()
+    assert native.unregister_buffer.call_count == 2
+    assert not native.registered
+
+
+@pytest.mark.parametrize("mode", ["miss", "oversized"])
+def test_unloadable_inputs_do_not_allocate_staging(backend, mode):
+    client, native = backend.store_client, backend.store_client.store
+    client._read_buffer_bytes = 39
+    if mode == "oversized":
+        client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2))
+    native.register_buffer = MagicMock(wraps=native.register_buffer)
+    assert not backend.resolve_inputs({"a": SPEC}, {}, "cpu")
+    assert client._read_buffer is None
+    if mode == "miss":
+        assert client._put_header is None
+        client.close()
+    native.register_buffer.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "results", [[40, -800], [40], [40, True], [[40], [40]], None, [40, 41]]
+)
+def test_unsafe_batch_result_retains_whole_staging(backend, results):
+    client, native = backend.store_client, backend.store_client.store
+    native.batch_get_into = MagicMock(return_value=results)
+    native.unregister_buffer = MagicMock(wraps=native.unregister_buffer)
+    for _ in range(2):
+        with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
+            client.load_tensors(dict.fromkeys("ab", SPEC), "cpu")
+    native.batch_get_into.assert_called_once()
+    with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
+        client.close()
+    native.unregister_buffer.assert_not_called()
+    assert client._read_buffer is client._unsafe_owners[0]
+
+
+@pytest.mark.parametrize("failure", ["allocation", "copy", "record", "completion"])
+def test_copy_failure_preserves_staging_lifetime(backend, failure):
+    """Emulate CUDA allocation/event boundaries without requiring a GPU."""
+    client = backend.store_client
+    for key in "ab":
+        client.put_tensor(key, torch.ones(2, 2))
+    allocate, copy = torch.empty, torch.Tensor.copy_
+    targets: list[torch.Tensor] = []
+    copies: list[torch.Tensor] = []
+    ready = MagicMock()
+    if failure == "record":
+        ready.record.side_effect = RuntimeError("record failure")
+    elif failure == "completion":
+        ready.synchronize.side_effect = RuntimeError("copy failure")
+
+    def allocate_on_cpu(*args, **kwargs):
+        is_target = kwargs.get("device") == torch.device("cuda")
+        if is_target and targets and failure == "allocation":
+            raise torch.OutOfMemoryError("allocation failure")
+        kwargs.update(device="cpu", pin_memory=False)
+        tensor = allocate(*args, **kwargs)
+        if is_target:
+            targets.append(tensor)
+        return tensor
+
+    def record_copy(target, source, **kwargs):
+        if copies and failure == "copy":
+            raise RuntimeError("copy failure")
+        copies.append(target)
+        return copy(target, source, **kwargs)
+
+    with (
+        patch.object(torch, "empty", side_effect=allocate_on_cpu),
+        patch.object(torch, "Event", return_value=ready),
+        patch.object(torch.accelerator, "current_stream"),
+        patch.object(torch.Tensor, "copy_", new=record_copy),
+        pytest.raises(RuntimeError, match="allocation failure|unconfirmed"),
+    ):
+        client.load_tensors(dict.fromkeys("ab", SPEC), "cuda")
+    if failure == "allocation":
+        assert not copies and not client._poisoned
+    else:
+        assert len(copies) == (1 if failure == "copy" else 2)
+        assert client._poisoned
+        assert any(owner is client._read_buffer for owner in client._unsafe_owners)
+        assert any(owner is ready for owner in client._unsafe_owners)
+        assert all(
+            any(owner is target for owner in client._unsafe_owners)
+            for target in targets
+        )
+        with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
+            client.close()
+
+
+def test_non_matrix_output_is_rejected_before_registration(backend):
+    client = backend.store_client
+    client.store.register_buffer = MagicMock()
+    with pytest.raises(store_client.EmbeddingStoreOperationError, match="2D"):
+        client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2, 2))
+    client.store.register_buffer.assert_not_called()
+
+
+def make_worker(backend):
+    from vllm.distributed.ec_transfer.ec_connector.mooncake.worker import (
+        ECMooncakeWorker,
+    )
+
+    worker = object.__new__(ECMooncakeWorker)
+    worker.is_producer = True
+    worker.is_consumer = False
+    worker._output_store = backend
+    worker._buffer_device = "cpu"
+    worker._producer_pushes = MagicMock()
+    worker._producer_pushes.reserve.return_value = (MagicMock(), True)
+    worker._producer_pushes.poll.return_value = []
+    worker._producer_pushes.cancel_requests.return_value = []
+    worker._producer_pushes.pending = False
+    worker._control_executor = MagicMock()
+    worker._consumer_memory = MagicMock()
+    worker._consumer_memory.drain_reclaimed.return_value = set()
+    worker._completed_loads = set()
+    worker._failed_loads = set()
+    worker._failed_saves = set()
+    worker._push_ready = threading.Event()
+    worker._flush_pending_pushes = MagicMock()
+    worker._bind_push_source = MagicMock()
+    connector = object.__new__(ECMooncakeConnector)
+    connector._worker = worker
+    connector._role = ECConnectorRole.WORKER
+    connector._is_producer = True
+    connector._is_consumer = False
+    connector._connector_metadata = None
+    return connector
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        "miss",
+        "legacy",
+        "hit",
+        "partial",
+        "get-failure",
+        "short-data",
+        "lease-expired",
+        "registration-failure",
+    ],
+)
+def test_reuse_pipeline(backend, mode):
+    from vllm.multimodal.inputs import (
+        MultiModalFeatureSpec,
+        MultiModalKwargsItem,
+        PlaceholderRange,
+    )
+    from vllm.v1.worker.gpu.ec_connector import ActiveECConnector
+    from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
+    from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
+    from vllm.v1.worker.gpu.model_states.interface import ModelState
+
+    native = backend.store_client.store
+    values = {"a": 1.0, "b": 2.0, "c": 3.0}
+    if mode == "legacy":
+        old_namespace = KEY.replace("@protocol:v3", "@protocol:v2")
+        for key in "abc":
+            native.objects[make_embedding_key(old_namespace, key)] = b"old format"
+    if mode == "partial":
+        hits = {"a", "c"}
+    elif mode not in ("miss", "legacy"):
+        hits = set("abc")
+    else:
+        hits = set()
+    for key in hits:
+        backend.store_client.put_tensor(
+            make_embedding_key(KEY, key), torch.full((2, 2), values[key])
+        )
+    if mode == "short-data":
+        key = make_embedding_key(KEY, "b")
+        native.objects[key] = native.objects[key][:-1]
+        hits.remove("b")
+    if mode in ("get-failure", "lease-expired"):
+        read = native.batch_get_into
+
+        def failed_read(keys, pointers, sizes):
+            results = read(keys, pointers, sizes)
+            results[keys.index(make_embedding_key(KEY, "b"))] = (
+                -707 if mode == "lease-expired" else -704
+            )
+            return results
+
+        native.batch_get_into = failed_read
+        hits.remove("b")
+    if mode == "registration-failure":
+        native.register_buffer = MagicMock(return_value=-500)
+        hits.clear()
+    connector = make_worker(backend)
+    metadata = ECMooncakeConnectorMetadata(
+        store_candidates=dict.fromkeys("abc", SPEC),
+        pushes=[
+            ECMooncakePushSpec(
+                mm_hash=k,
+                shape=(2, 2),
+                dtype="float32",
+                nbytes=16,
+                consumer_zmq="ipc:///tmp/test",
+                transfer_id=k,
+                request_id="req",
+            )
+            for k in "abc"
+        ],
+    )
+    cache = EncoderCache()
+    cache.mm_features["req"] = [
+        MultiModalFeatureSpec(
+            data=MultiModalKwargsItem({}),
+            modality="image",
+            identifier=k,
+            mm_position=PlaceholderRange(offset=i * 2, length=2),
+        )
+        for i, k in enumerate("abc")
+    ]
+    runner = EncoderRunner(
+        model=None,
+        max_num_tokens=8,
+        hidden_size=2,
+        encoder_cache=cache,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    runner.enable_timing = False
+    missing = [k for k in "abc" if k not in hits]
+    runner.execute_mm_encoder = MagicMock(
+        return_value=[torch.full((2, 2), values[k] + 10) for k in missing]
+    )
+    state = SimpleNamespace(encoder_runner=runner, encoder_cache=cache)
+    with patch(
+        "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector
+    ):
+        adapter = ActiveECConnector(SimpleNamespace(), cache.encoder_outputs)
+    with adapter.maybe_get_output(
+        SimpleNamespace(ec_connector_metadata=metadata, finished_req_ids=set())
+    ) as output:
+        assert runner.prepare_mm_inputs({"req": [0, 1, 2]})[0] == missing
+        ModelState.execute_mm_encoder(state, {"req": [0, 1, 2]})
+    assert runner.execute_mm_encoder.call_count == bool(missing)
+    assert output.ec_connector_worker_meta.pending_saves is False
+    assert all(
+        torch.equal(
+            cache.encoder_outputs[k],
+            torch.full((2, 2), values[k] + (0 if k in hits else 10)),
+        )
+        for k in "abc"
+    )
+    assert {
+        c.args[1] for c in connector._worker._bind_push_source.call_args_list
+    } == set("abc")
+    backend.shutdown()
+    assert not native.registered
+    if mode in ("miss", "legacy"):
+        other = store_backend.MooncakeEmbeddingStoreBackend(
+            store_client.MooncakeEmbeddingStoreClient(native, read_buffer_bytes=1024),
+            KEY,
+            max_pending_items=32,
+            max_pending_bytes=2 * 1024**3,
+        )
+        try:
+            loaded: dict[str, torch.Tensor] = {}
+            assert other.resolve_inputs(
+                dict.fromkeys("abc", SPEC),
+                loaded,
+                "cpu",
+            ) == set("abc")
+            assert all(torch.equal(loaded[k], cache.encoder_outputs[k]) for k in "abc")
+        finally:
+            other.shutdown()
+
+
+@pytest.mark.parametrize(
+    "mismatch",
+    ["revision", "shape", "dtype", "magic", "short-header", "short-data", "long-data"],
+)
+def test_incompatible_output_is_a_miss(backend, mismatch):
+    key = make_embedding_key(KEY, "a")
+    dtype = torch.float16 if mismatch == "dtype" else torch.float32
+    backend.store_client.put_tensor(key, torch.ones(2, 2, dtype=dtype))
+    if mismatch == "revision":
+        backend.namespace = KEY.replace("@revision:revision", "@revision:other")
+    objects = backend.store_client.store.objects
+    if mismatch == "magic":
+        objects[key] = b"\0" * 4 + objects[key][4:]
+    elif mismatch == "short-header":
+        objects[key] = objects[key][:23]
+    elif mismatch == "short-data":
+        objects[key] = objects[key][:-1]
+    elif mismatch == "long-data":
+        objects[key] += b"extra"
+    expected = replace(SPEC, shape=(4, 1)) if mismatch == "shape" else SPEC
+    outputs: dict[str, torch.Tensor] = {}
+    assert not backend.resolve_inputs({"a": expected}, outputs, "cpu")
+    assert not outputs
+    tensor = torch.ones(2, 2)
+    backend.store_client.put_tensor(
+        make_embedding_key(backend.namespace, "valid"), tensor
+    )
+    assert backend.resolve_inputs({"valid": SPEC}, outputs, "cpu") == {"valid"}
+    assert torch.equal(outputs["valid"], tensor)
+    assert len(backend.store_client.store.registered) == 2
+
+
+def test_partial_registration_rejection_releases_publication(backend):
+    """A rejected header registration must not prevent a later publication."""
+    native = backend.store_client.store
+    register = native.register_buffer
+    put = native.batch_put_from_multi_buffers
+    native.register_buffer = lambda addr, size: (
+        -500 if size == 24 else register(addr, size)
+    )
+    native.batch_put_from_multi_buffers = MagicMock()
+    tensor = torch.ones(2, 2)
+    owner = weakref.ref(tensor)
+    assert backend._enqueue_save("a", tensor)
+    storage_owner = weakref.ref(tensor.untyped_storage())
+    future = backend._pending["a"][0]
+    with pytest.raises(store_client.EmbeddingStoreOperationError, match="register"):
+        future.result(timeout=5)
+    del tensor
+    gc.collect()
+    assert owner() is None and not native.registered
+    native.batch_put_from_multi_buffers.assert_not_called()
+    backend.resolve_inputs({}, {}, "cpu")
+    assert backend._pending_bytes == 0
+    gc.collect()
+    assert storage_owner() is None  # The failed Future is still alive.
+    assert backend.store_client._put_header is None
+    native.register_buffer = register
+    native.batch_put_from_multi_buffers = put
+    backend.store_client.put_tensor("retry", torch.ones(2, 2))
+    assert len(native.registered) == 1
+
+
+@pytest.mark.parametrize(
+    "operation",
+    ["read-unregister", "header-unregister", "put", "rejected-put", "unregister"],
+)
+def test_native_io_owners(backend, operation):
+    """Uncertain I/O retains owners; completed rejection frees them even via Future."""
+    native, client = backend.store_client.store, backend.store_client
+    put = native.batch_put_from_multi_buffers
+    backend.max_pending_bytes = SPEC.nbytes
+    tensor = torch.ones(2, 2)
+    owner = weakref.ref(tensor)
+    rejected = operation == "rejected-put"
+    if operation in ("read-unregister", "header-unregister"):
+        key = make_embedding_key(KEY, "a")
+        client.put_tensor(key, tensor)
+        if operation == "read-unregister":
+            client.load_tensors({key: SPEC}, "cpu")
+        native.unregister_buffer = MagicMock(return_value=-500)
+        with pytest.raises(store_client.EmbeddingStoreError, match="unregister"):
+            client.close()
+        owner = weakref.ref(
+            client._read_buffer
+            if operation == "read-unregister"
+            else client._put_header
+        )
+    else:
+        if operation == "unregister":
+            native.unregister_buffer = MagicMock(return_value=-500)
+        else:
+            native.batch_put_from_multi_buffers = MagicMock(
+                return_value=[-1700 if rejected else -800]
+            )
+        assert backend._enqueue_save("a", tensor)
+        future = backend._pending["a"][0]
+        with pytest.raises(store_client.EmbeddingStoreError):
+            future.result(timeout=5)
+    del tensor
+    gc.collect()
+    if rejected:
+        assert owner() is None  # A live Future must not retain the tensor.
+        tensor = torch.ones(4, 4)  # Too large for admission after reaping.
+        caller_owner = weakref.ref(tensor)
+        backend._step_candidates = {"next"}
+        backend.save_output("next", tensor)
+        del tensor
+        gc.collect()
+        assert caller_owner() is None
+        assert backend._pending_bytes == 0 and len(native.registered) == 1
+        header = client._put_header
+        native.batch_put_from_multi_buffers = put
+        client.put_tensor("retry", torch.ones(1, 4))
+        assert client._put_header is header
+    else:
+        assert owner() is not None and native.registered
+        header = bytes(client._put_header)
+        with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
+            client.put_tensor("retry", torch.ones(1, 4))
+        assert bytes(client._put_header) == header
+        with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
+            client.close()
+
+
+def test_publication_budget_charges_and_releases_backing_storage(backend):
+    tensor = torch.arange(16, dtype=torch.float32)[:4].view(2, 2)
+    owner = weakref.ref(tensor)
+    future: Future[None] = Future()
+    backend.max_pending_items = 1
+    backend.max_pending_bytes = tensor.nbytes
+    with patch.object(backend._executor, "submit", return_value=future):
+        assert not backend._enqueue_save("a", tensor)
+        backend.max_pending_bytes = 64
+        assert backend._enqueue_save("a", tensor)
+        storage_owner = weakref.ref(tensor.untyped_storage())
+        assert not backend._enqueue_save("b", tensor)
+    del tensor
+    gc.collect()
+    backend.reap()  # The final step-end poll precedes publication completion.
+    assert owner() is not None
+    backend._save(make_embedding_key(KEY, "a"), backend._pending["a"][1])
+    future.set_result(None)
+    assert owner() is None and storage_owner() is not None
+
+    def load_after_reaping(*args):
+        assert not backend._pending and backend._pending_bytes == 0
+        assert storage_owner() is None
+        return {}
+
+    with patch.object(backend.store_client, "load_tensors", load_after_reaping):
+        assert not backend.resolve_inputs({"b": SPEC}, {}, "cpu")
+
+
+@pytest.mark.parametrize("rejected", [False, True])
+def test_shared_storage_budget_is_released_after_last_publication(backend, rejected):
+    batch = torch.arange(16, dtype=torch.float32).view(4, 4)
+    first: Future[None] = Future()
+    last: Future[None] = Future()
+    separate: Future[None] = Future()
+    backend.max_pending_items = 2
+    backend.max_pending_bytes = batch.nbytes
+    with patch.object(backend._executor, "submit", side_effect=[first, last, separate]):
+        assert backend._enqueue_save("a", batch[:2])
+        assert backend._enqueue_save("b", batch[2:])
+        assert backend._pending_bytes == batch.nbytes
+        assert not backend._enqueue_save("a", batch[:2])
+        assert not backend._enqueue_save("extra", batch[:1])
+        if rejected:
+            first.set_exception(store_client.EmbeddingStoreOperationError("rejected"))
+        else:
+            first.set_result(None)
+        backend.reap()
+        assert backend._pending_bytes == batch.nbytes
+        assert not backend._enqueue_save("c", torch.ones(2, 2))
+        last.set_result(None)
+        backend.reap()
+        assert backend._pending_bytes == 0
+        assert backend._enqueue_save("c", torch.ones(2, 2))
+        assert backend._pending_bytes == 16
+        separate.set_result(None)
+        backend.reap()
+        assert backend._pending_bytes == 0
+
+
+@pytest.mark.parametrize("shared_storage", [False, True])
+def test_completion_race_preserves_failure_and_reclaims_other_items(
+    backend, shared_storage
+):
+    class RacingFuture(Future[None]):
+        def done(self):
+            observed = super().done()
+            if not observed:
+                self.set_exception(store_client.EmbeddingStoreError("uncertain"))
+            return observed
+
+    racing = RacingFuture()
+    other: Future[None] = Future()
+    tensor = torch.ones(2, 2)
+    other_tensor = tensor[:1] if shared_storage else torch.ones(2, 2)
+    with patch.object(backend._executor, "submit", side_effect=[racing, other]):
+        assert backend._enqueue_save("race", tensor)
+        assert backend._enqueue_save("other", other_tensor)
+    backend.reap()
+    assert backend._pending_bytes == (16 if shared_storage else 32)
+    other.set_result(None)
+    with pytest.raises(store_client.EmbeddingStoreError):
+        backend.reap()
+    assert backend._pending_bytes == 16
+    with pytest.raises(store_client.EmbeddingStoreError):
+        backend._enqueue_save("later", torch.ones(2, 2))
+
+
+def test_readiness_and_shutdown_do_not_block_p2p_completion(backend):
+    ready, waiting, entered = [threading.Event() for _ in range(3)]
+
+    original = backend._executor.shutdown
+
+    def closing(*args, **kwargs):
+        entered.set()
+        return original(*args, **kwargs)
+
+    native = backend.store_client.store
+
+    def close_store():
+        assert ready.is_set() and not native.registered
+        assert backend._pending_bytes == 0
+        return 0
+
+    native.close = MagicMock(side_effect=close_store)
+    put = native.batch_put_from_multi_buffers
+
+    def waiting_put(*args):
+        waiting.set()
+        assert ready.wait(5)
+        assert backend.store_client._put_header is not None and native.registered
+        return put(*args)
+
+    native.batch_put_from_multi_buffers = waiting_put
+    connector = make_worker(backend)
+    connector.bind_connector_metadata(
+        ECMooncakeConnectorMetadata(store_candidates={"a": SPEC})
+    )
+    cache: dict[str, torch.Tensor] = {}
+    with (
+        patch.object(
+            store_backend, "_record_tensor_ready_event", return_value=MagicMock()
+        ),
+        patch.object(backend._executor, "shutdown", side_effect=closing),
+        ThreadPoolExecutor(max_workers=1) as closer,
+    ):
+        connector.start_save_caches(encoder_cache=cache)
+        cache["a"] = torch.arange(4, dtype=torch.float32).view(2, 2)
+        connector.save_caches(cache, "a")
+        assert connector.build_connector_worker_meta().pending_saves is False
+        assert waiting.wait(5) and not backend.store_client.store.objects
+        close = closer.submit(backend.shutdown)
+        try:
+            assert entered.wait(5) and not close.done()
+            assert not backend._enqueue_save("late", torch.ones(2, 2))
+        finally:
+            ready.set()
+        close.result(timeout=5)
+    native.close.assert_called_once_with()
+    assert not backend.store_client.store.registered and backend._pending_bytes == 0
+
+    assert backend.store_client.store.objects[make_embedding_key(KEY, "a")][
+        24:
+    ] == struct.pack("<4f", 0, 1, 2, 3)
diff --git a/tests/v1/ec_connector/unit/test_ec_connector_metrics.py b/tests/v1/ec_connector/unit/test_ec_connector_metrics.py
new file mode 100644
index 000000000000..239f0267f060
--- /dev/null
+++ b/tests/v1/ec_connector/unit/test_ec_connector_metrics.py
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the connector-agnostic EC connector metrics plumbing.
+
+Covers only the base infra: no concrete connector type implements these
+hooks yet, so every assertion here is about the documented no-op defaults
+and the container contracts, mirroring the equivalent KV connector tests.
+"""
+
+from unittest.mock import Mock
+
+import pytest
+
+from vllm.config import VllmConfig
+from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
+from vllm.distributed.ec_transfer.ec_connector.example_connector import (
+    ECExampleConnector,
+)
+from vllm.distributed.ec_transfer.ec_connector.metrics import (
+    ECConnectorLogging,
+    ECConnectorProm,
+    ECConnectorStats,
+)
+from vllm.v1.outputs import ECConnectorOutput
+
+pytestmark = pytest.mark.cpu_test
+
+
+@pytest.fixture
+def example_connector(tmp_path):
+    config = Mock(spec=VllmConfig)
+    config.ec_transfer_config = Mock()
+    config.ec_transfer_config.get_from_extra_config = Mock(return_value=str(tmp_path))
+    config.ec_transfer_config.is_ec_producer = True
+    config.ec_transfer_config.is_ec_consumer = False
+    return ECExampleConnector(vllm_config=config, role=ECConnectorRole.SCHEDULER)
+
+
+def test_base_stats_methods_raise_not_implemented():
+    """The base container has no aggregation logic of its own."""
+    stats = ECConnectorStats()
+
+    with pytest.raises(NotImplementedError):
+        stats.reset()
+    with pytest.raises(NotImplementedError):
+        stats.aggregate(stats)
+    with pytest.raises(NotImplementedError):
+        stats.reduce()
+    with pytest.raises(NotImplementedError):
+        stats.is_empty()
+
+
+def test_connector_hooks_default_to_none(example_connector):
+    """No connector type overrides these yet, so every hook is a no-op."""
+    assert example_connector.get_ec_connector_stats() is None
+    assert ECExampleConnector.build_ec_connector_stats() is None
+    assert (
+        ECExampleConnector.build_prom_metrics(
+            vllm_config=Mock(spec=VllmConfig),
+            metric_types={},
+            labelnames=[],
+            per_engine_labelvalues={},
+        )
+        is None
+    )
+
+
+def test_ec_connector_output_is_empty_reflects_stats_field():
+    output = ECConnectorOutput()
+    assert output.is_empty()
+
+    output.ec_connector_stats = ECConnectorStats(data={"x": 1})
+    assert not output.is_empty()
+
+
+def test_logging_and_prom_are_no_ops_without_ec_transfer_config():
+    """A normal run with no EC connector configured must stay silent."""
+    logging = ECConnectorLogging(ec_transfer_config=None)
+    logging.log()  # Should not raise, nothing to log.
+
+    vllm_config = Mock(spec=VllmConfig)
+    vllm_config.ec_transfer_config = None
+    prom = ECConnectorProm(vllm_config, labelnames=[], per_engine_labelvalues={})
+    assert prom.prom_metrics is None
+    prom.observe({})  # Should not raise.
diff --git a/tests/v1/ec_connector/unit/test_ec_cpu_connector.py b/tests/v1/ec_connector/unit/test_ec_cpu_connector.py
index e2cb8b120af5..bbbb615d9de4 100644
--- a/tests/v1/ec_connector/unit/test_ec_cpu_connector.py
+++ b/tests/v1/ec_connector/unit/test_ec_cpu_connector.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E test for the ECCPUConnector.
+"""E2E test for the ECCPUConnector.
 
 Verifies:
 - Accuracy: outputs from EC CPU cache match fresh encoder computation.
diff --git a/tests/v1/ec_connector/unit/test_ec_example_connector.py b/tests/v1/ec_connector/unit/test_ec_example_connector.py
index 8edb54d2a69c..725bb99eb055 100644
--- a/tests/v1/ec_connector/unit/test_ec_example_connector.py
+++ b/tests/v1/ec_connector/unit/test_ec_example_connector.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for ECExampleConnector.
-"""
+"""Unit tests for ECExampleConnector."""
 
 import os
 from unittest.mock import Mock, patch
@@ -18,6 +16,7 @@
     ECExampleConnectorMetadata,
     MMMeta,
 )
+from vllm.distributed.ec_transfer.ec_connector.utils import collect_ec_item_metadata
 from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
 from vllm.v1.core.sched.output import SchedulerOutput
 
@@ -43,6 +42,36 @@ def get_num_encoder_embeds(self, input_id: int) -> int:
         return self._token_counts[input_id]
 
 
+def test_metadata_preserves_repeated_item_positions():
+    request = MockRequest("request", ["A", "B", "A"], [4, 4, 4])
+    resolver = Mock()
+    resolver.fields_for.return_value = {"image_grid_thw"}
+    for feature in request.mm_features:
+        feature.data = Mock()
+        feature.data.get_data.return_value = {"image_grid_thw": torch.tensor([1, 2, 2])}
+    items = collect_ec_item_metadata(request.mm_features, resolver)
+    assert items == {
+        "A": {"metadata": {"image_grid_thw": [1, 2, 2]}, "item_indices": [0, 2]},
+        "B": {"metadata": {"image_grid_thw": [1, 2, 2]}, "item_indices": [1]},
+    }
+
+
+def test_audio_metadata_survives_processor_cache_hit():
+    resolver = Mock()
+    resolver.fields_for.return_value = {"audio_num_tokens"}
+    feature = MultiModalFeatureSpec(
+        data=None,
+        modality="audio",
+        identifier="audio",
+        mm_position=PlaceholderRange(
+            offset=0, length=4, is_embed=torch.tensor([False, True, True, False])
+        ),
+    )
+    assert collect_ec_item_metadata([feature], resolver) == {
+        "audio": {"metadata": {"audio_num_tokens": [2]}, "item_indices": [0]}
+    }
+
+
 @pytest.fixture
 def temp_storage(tmp_path):
     """Fixture providing temporary storage path."""
diff --git a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py
index 4d5b7dae1037..df98ec364f5c 100644
--- a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py
+++ b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py
@@ -5,6 +5,7 @@
 import pytest
 
 from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata
+from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats
 from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
 from vllm.v1.outputs import (
@@ -29,6 +30,16 @@ def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta":
         return FakeWorkerMeta(self.saves + other.saves)
 
 
+class FakeConnectorStats(ECConnectorStats):
+    """Records merge order, mirroring FakeWorkerMeta above."""
+
+    def __init__(self, saves: list[str]):
+        self.saves = saves
+
+    def aggregate(self, other: "FakeConnectorStats") -> "FakeConnectorStats":
+        return FakeConnectorStats(self.saves + other.saves)
+
+
 def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput:
     return ModelRunnerOutput(
         req_ids=[], req_id_to_index={}, ec_connector_output=ec_output
@@ -38,19 +49,23 @@ def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput:
 def test_aggregate_folds_every_rank_onto_output_rank():
     """EC work done on any rank reaches the scheduler via output_rank's output.
 
-    The middle rank reports no worker metadata: it must neither seed nor clobber
-    the accumulator.
+    The middle rank reports no worker metadata or stats: it must neither seed
+    nor clobber either accumulator.
     """
     outputs = [
         _worker_output(
             ECConnectorOutput(
                 finished_sending={"mm0"},
+                ec_connector_stats=FakeConnectorStats(["mm0"]),
                 ec_connector_worker_meta=FakeWorkerMeta(["mm0"]),
             )
         ),
         _worker_output(ECConnectorOutput(finished_recving={"mm1"})),
         _worker_output(
-            ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"]))
+            ECConnectorOutput(
+                ec_connector_stats=FakeConnectorStats(["mm2"]),
+                ec_connector_worker_meta=FakeWorkerMeta(["mm2"]),
+            )
         ),
     ]
 
@@ -59,6 +74,7 @@ def test_aggregate_folds_every_rank_onto_output_rank():
     assert result is outputs[2]
     assert result.ec_connector_output.finished_sending == {"mm0"}
     assert result.ec_connector_output.finished_recving == {"mm1"}
+    assert result.ec_connector_output.ec_connector_stats.saves == ["mm0", "mm2"]
     assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"]
 
 
diff --git a/tests/v1/ec_connector/unit/test_epd_proxy.py b/tests/v1/ec_connector/unit/test_epd_proxy.py
new file mode 100644
index 000000000000..13496a7ffb23
--- /dev/null
+++ b/tests/v1/ec_connector/unit/test_epd_proxy.py
@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Routing and registration behaviour of the EPD proxy."""
+
+import asyncio
+from typing import Any
+from unittest.mock import AsyncMock, Mock
+
+import msgspec
+import pytest
+from fastapi import HTTPException
+from fastapi.testclient import TestClient
+
+from examples.disaggregated.disaggregated_encoder import (
+    disagg_epd_proxy as proxy_module,
+)
+from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import (
+    EPDProxy,
+    EPDProxyConfig,
+    InstanceRecord,
+    InstanceRegistry,
+    InstanceRole,
+    build_app,
+    extract_mm_items,
+)
+
+ENCODE = InstanceRole.ENCODE
+PREFILL = InstanceRole.PREFILL
+DECODE = InstanceRole.DECODE
+PD = InstanceRole.PREFILL_DECODE
+
+IMAGE_ITEM = {"type": "image_url", "image_url": {"url": "http://img/0.png"}}
+
+
+@pytest.fixture
+def proxy():
+    registry = InstanceRegistry(probe_interval=0)
+    return EPDProxy(registry)
+
+
+class TestRouting:
+    def test_no_decode_instance_is_service_unavailable(self, proxy):
+        """The proxy comes up before anything registers, so this is normal."""
+        with pytest.raises(HTTPException) as excinfo:
+            proxy.route(num_items=0)
+        assert excinfo.value.status_code == 503
+
+    def test_media_without_an_encoder_is_service_unavailable(self, proxy):
+        proxy.registry.register(InstanceRecord(PD, "http://d0:8000"))
+        assert proxy.route(num_items=0).decode.url == "http://d0:8000"
+        with pytest.raises(HTTPException) as excinfo:
+            proxy.route(num_items=1)
+        assert excinfo.value.status_code == 503
+
+    def test_prefill_is_optional(self, proxy):
+        """An E+PD deployment registers no prefill instance at all."""
+        proxy.registry.register(InstanceRecord(PD, "http://d0:8000"))
+        assert proxy.route(num_items=0).prefill is None
+
+    def test_encoder_roster_tracks_new_instances(self, proxy):
+        proxy.registry.register(InstanceRecord(PD, "http://d0:8000"))
+        for index in range(2):
+            proxy.registry.register(InstanceRecord(ENCODE, f"http://e{index}:8000"))
+        assert proxy.route(num_items=3).encoder_urls == [
+            "http://e0:8000",
+            "http://e1:8000",
+        ]
+        proxy.registry.unregister("http://e0:8000")
+        assert proxy.route(num_items=3).encoder_urls == ["http://e1:8000"]
+
+    @pytest.mark.parametrize("unregister", [False, True])
+    def test_standalone_decode_always_requires_healthy_prefill(self, unregister):
+        app = build_app()
+        proxy = app.state.proxy
+        proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000"))
+        with pytest.raises(HTTPException) as excinfo:
+            proxy.route(num_items=0)
+        assert excinfo.value.status_code == 503
+        proxy.registry.register(InstanceRecord(PREFILL, "http://p0:8000"))
+        route = proxy.route(num_items=0)
+        assert route.decode.url == "http://d0:8000"
+        assert route.consumer.url == "http://p0:8000"
+        if unregister:
+            proxy.registry.unregister("http://p0:8000")
+        else:
+            for _ in range(proxy.registry._fail_threshold):
+                proxy.registry._on_probe_failure(route.prefill, 0)
+        with pytest.raises(HTTPException) as excinfo:
+            proxy.route(num_items=0)
+        assert excinfo.value.status_code == 503
+
+
+class TestConsumerAddress:
+    """Which stage receives the embedding depends on the topology."""
+
+    def test_shared_storage_connectors_name_no_target(self, proxy):
+        """Rotate replicas without requiring a push endpoint."""
+        proxy.registry.register(InstanceRecord(PD, "http://d0:8000", dp_size=2))
+        route = proxy.route(num_items=0)
+        assert route.consumer_zmq is None
+        assert route.dp_rank == 0
+        assert proxy.route(num_items=0).dp_rank == 1
+
+    def test_decode_is_the_consumer_when_prefill_is_not_split_out(self, proxy):
+        proxy.registry.register(
+            InstanceRecord(PD, "http://d0:8000", ec_zmq_addrs=["tcp://d0:20001"])
+        )
+        assert proxy.route(num_items=0).consumer_zmq == "tcp://d0:20001"
+
+    def test_prefill_is_the_consumer_when_it_is_split_out(self, proxy):
+        """In E+P+D the prefill instance consumes the embedding, not decode.
+
+        Reading the address out of whichever list it came from -- rather than
+        from the record that reported one -- sends the push to the wrong
+        instance in this topology.
+        """
+        proxy.registry.register(
+            InstanceRecord(PREFILL, "http://p0:8000", ec_zmq_addrs=["tcp://p0:20001"])
+        )
+        proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000"))
+        assert proxy.route(num_items=0).consumer_zmq == "tcp://p0:20001"
+
+
+class TestProxyApi:
+    @pytest.fixture
+    def client(self, monkeypatch):
+        monkeypatch.setenv("ADMIN_API_KEY", "test-key")
+        config = EPDProxyConfig(probe_interval=0)
+        with TestClient(build_app(config)) as client:
+            client.headers["X-API-Key"] = "test-key"
+            yield client
+
+    def test_proxy_serves_before_anything_registers(self, client):
+        assert client.get("/health").status_code == 200
+        assert client.get("/instances").json()["encode"]["live"] == []
+
+    def test_requests_are_refused_until_a_decoder_registers(self, client):
+        response = client.post(
+            "/v1/chat/completions", json={"model": "m", "messages": []}
+        )
+        assert response.status_code == 503
+        assert client.get("/v1/models").status_code == 503
+
+    def test_http_registration_keeps_consumer_replica_and_address_together(
+        self, client
+    ):
+        payload: dict[str, Any] = {
+            "role": "prefill_decode",
+            "url": "http://d0:8000/",
+            "dp_size": 2,
+            "ec_zmq_addrs": ["tcp://d0:20001", "tcp://d0:20003"],
+        }
+        assert client.post("/instances", json=payload).json() == {"registered": True}
+        assert client.post("/instances", json=payload).json() == {"registered": False}
+        proxy = client.app.state.proxy
+        for rank in [0, 1, 0]:
+            route = proxy.route(0)
+            assert route.dp_rank == rank
+            assert route.consumer_zmq == payload["ec_zmq_addrs"][rank]
+        assert client.delete("/instances", params={"url": payload["url"]}).json() == {
+            "removed": True
+        }
+        assert proxy.registry.pick(PD) is None
+
+    def test_registration_rejects_invalid_topology(self, client):
+        cases: list[tuple[dict[str, Any], int]] = [
+            ({"dp_size": 0}, 422),
+            ({"dp_size": 2, "ec_zmq_addrs": ["tcp://d0:20001"]}, 422),
+            ({"role": "decode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400),
+            ({"role": "encode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400),
+        ]
+        for fields, status in cases:
+            response = client.post(
+                "/instances",
+                json={"url": "http://d0:8000", "role": "prefill_decode", **fields},
+            )
+            assert response.status_code == status
+
+    def test_registration_requires_admin_key(self, client):
+        client.headers.pop("X-API-Key")
+        assert (
+            client.post(
+                "/instances", json={"role": "encode", "url": "http://e0:8000"}
+            ).status_code
+            == 403
+        )
+        assert (
+            client.delete("/instances", params={"url": "http://e0:8000"}).status_code
+            == 403
+        )
+
+    @pytest.mark.parametrize("split_role", ["prefill", "decode"])
+    @pytest.mark.parametrize("combined_first", [False, True])
+    def test_registration_rejects_mixed_topologies(
+        self, client, split_role, combined_first
+    ):
+        roles = ["prefill_decode", split_role]
+        if not combined_first:
+            roles.reverse()
+        first = {"role": roles[0], "url": "http://first:8000"}
+        second = {"role": roles[1], "url": "http://second:8000"}
+        assert client.post("/instances", json=first).status_code == 200
+        assert client.post("/instances", json=second).status_code == 409
+        # Health eviction must not silently permit a topology switch.
+        registry = client.app.state.registry
+        record = registry._live[first["url"]]
+        for _ in range(registry._fail_threshold):
+            registry._on_probe_failure(record, 0)
+        assert client.post("/instances", json=second).status_code == 409
+        assert client.delete("/instances", params={"url": first["url"]}).json() == {
+            "removed": True
+        }
+        assert client.post("/instances", json=second).status_code == 200
+
+
+def test_extract_mm_items_finds_media_across_messages():
+    req = {
+        "messages": [
+            {"role": "user", "content": "plain text"},
+            {"role": "user", "content": [IMAGE_ITEM, {"type": "text", "text": "hi"}]},
+            {"role": "user", "content": [IMAGE_ITEM]},
+        ]
+    }
+    assert extract_mm_items(req) == [IMAGE_ITEM, IMAGE_ITEM]
+
+
+@pytest.mark.parametrize("push", [False, True])
+def test_encoder_handles_and_json_metadata_survive_rewrite(proxy, push, monkeypatch):
+    """Keep main's hash-keyed handles and per-item push IDs after the proxy move."""
+    proxy.registry.register(InstanceRecord(ENCODE, "http://encoder"))
+    proxy.registry.register(
+        InstanceRecord(
+            PD,
+            "http://decode",
+            ec_zmq_addrs=["tcp://decode:14579"] if push else [],
+        )
+    )
+    handle = {
+        "metadata": {"image_grid_thw": [[1, 2, 3]]},
+        "item_indices": [0, 1],
+        "peer_port": 5601,
+    }
+    response = Mock(
+        status=200,
+        read=AsyncMock(
+            return_value=msgspec.json.encode(
+                {"ec_transfer_params": {"engine-hash": handle}}
+            )
+        ),
+    )
+    session = Mock(post=AsyncMock(return_value=response))
+    monkeypatch.setattr(proxy_module, "encode_session", session)
+    original = {"messages": [{"role": "user", "content": [IMAGE_ITEM, IMAGE_ITEM]}]}
+    route = proxy.route(2)
+    prepared, _, _ = asyncio.run(
+        proxy_module.prepare_for_decode(
+            original, "request", route.encoder_urls, None, route.consumer_zmq
+        )
+    )
+    assert "ec_transfer_params" not in original
+    assert prepared["ec_transfer_params"]["engine-hash"] == handle
+    items = prepared["ec_transfer_params"]["ec_items"]
+    assert [item["mm_hash"] for item in items] == ["engine-hash", "engine-hash"]
+    assert items[0]["transfer_id"] != items[1]["transfer_id"]
+    assert session.post.call_count == 1
+    sent = msgspec.json.decode(session.post.call_args.kwargs["data"]).get(
+        "ec_transfer_params"
+    )
+    if push:
+        assert sent == {
+            "consumer_zmq": "tcp://decode:14579",
+            "ec_items": [{"transfer_id": item["transfer_id"]} for item in items],
+        }
+    else:
+        assert sent is None
+    for index in range(2):
+        assert prepared["messages"][0]["content"][index]["image_embeds"] == {
+            "image_grid_thw": [1, 2, 3]
+        }
diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py
index b026b5537a20..08bca0e40fcf 100644
--- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py
+++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py
@@ -8,9 +8,16 @@
 
 import asyncio
 import importlib.util
+import sys
 from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
 
+import httpx
+import msgspec
 import pytest
+from aiohttp import web
+from aiohttp.test_utils import TestServer
 
 PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py"
 
@@ -21,6 +28,7 @@ def proxy():
     spec = importlib.util.spec_from_file_location("disagg_epd_proxy_retry", path)
     assert spec is not None and spec.loader is not None
     module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
     spec.loader.exec_module(module)
     return module
 
@@ -29,8 +37,8 @@ class _Response:
     def __init__(self, params):
         self._params = params
 
-    async def json(self):
-        return {"kv_transfer_params": self._params}
+    async def read(self):
+        return msgspec.json.encode({"kv_transfer_params": self._params})
 
 
 def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch):
@@ -42,7 +50,7 @@ def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch):
     """
     served = [{"remote_block_ids": [1, 2]}, {}]
 
-    async def _stage(req_data, p_url, req_id):
+    async def _stage(req_data, p_url, req_id, dp_rank=None):
         assert "kv_transfer_params" not in req_data
         return _Response(served.pop(0))
 
@@ -62,8 +70,8 @@ def __init__(self, params):
         self.status = 200
         self._params = params
 
-    async def json(self):
-        return {"ec_transfer_params": self._params}
+    async def read(self):
+        return msgspec.json.encode({"ec_transfer_params": self._params})
 
     async def text(self):
         return ""
@@ -75,10 +83,208 @@ class _EncoderSession:
     def __init__(self, replies):
         self._replies = list(replies)
 
-    async def post(self, url, json=None, headers=None):
+    async def post(self, url, data=None, headers=None):
         return _EncoderResponse(self._replies.pop(0))
 
 
+@pytest.mark.parametrize("second_type", ["image_url", "audio_url", "input_audio"])
+@pytest.mark.parametrize("transfer", ["example", "nixl", "mooncake"])
+def test_partial_audio_metadata_falls_back_with_all_transfers(
+    proxy, monkeypatch, second_type, transfer
+):
+    """Never mix metadata-only audio with raw media, or lose EC identities."""
+    import copy
+
+    seen = []
+    reported = {}
+
+    class Session:
+        async def post(self, url, data=None, headers=None):
+            body = msgspec.json.decode(data)
+            seen.append(body)
+            item = body["messages"][0]["content"][0]
+            mm_hash = item["uuid"] + "-processed"
+            entry: dict[str, Any] = {
+                "metadata": {"audio_feature_lengths": [100]} if len(seen) == 1 else {},
+                "item_indices": [0],
+            }
+            if transfer == "nixl":
+                entry.update(peer_host="encoder", peer_port=1234, size_bytes=4096)
+            reported[mm_hash] = entry
+            return _EncoderResponse({mm_hash: entry})
+
+    second = (
+        {"type": second_type, "input_audio": {"data": "large", "format": "wav"}}
+        if second_type == "input_audio"
+        else {"type": second_type, second_type: {"url": "large"}}
+    )
+    body = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": "small"},
+                        "uuid": "user-audio",
+                    },
+                    second,
+                ],
+            }
+        ],
+        "mm_processor_kwargs": {"sampling_rate": 16000},
+    }
+    original = copy.deepcopy(body)
+    monkeypatch.setattr(proxy, "encode_session", Session())
+    monkeypatch.setattr(proxy, "NO_REWRITE", False)
+    consumer = "tcp://consumer:1" if transfer == "mooncake" else None
+    prepared, _, _ = asyncio.run(
+        proxy.prepare_for_decode(body, "r", ["http://encoder"], "", consumer)
+    )
+    assert body == original
+    assert prepared["mm_processor_kwargs"] == body["mm_processor_kwargs"]
+    content = prepared["messages"][0]["content"]
+    assert [item["type"] for item in content] == ["audio_url", second_type]
+    assert content[0]["uuid"] == "user-audio"
+    for item, encoded in zip(content, seen):
+        assert item == encoded["messages"][0]["content"][0]
+    params = prepared["ec_transfer_params"]
+    for mm_hash, entry in reported.items():
+        assert params[mm_hash] == entry
+    if transfer == "mooncake":
+        assert params["ec_items"] == [
+            {
+                "mm_hash": mm_hash,
+                "transfer_id": encoded["ec_transfer_params"]["ec_items"][0][
+                    "transfer_id"
+                ],
+            }
+            for mm_hash, encoded in zip(reported, seen)
+        ]
+
+
+@pytest.mark.parametrize("no_rewrite", [False, True])
+@pytest.mark.parametrize(
+    "batch_size, expected_sizes", [(0, [3, 3]), (1, [1] * 6), (2, [2, 2, 1, 1])]
+)
+def test_image_batches_preserve_item_identity(
+    proxy, monkeypatch, no_rewrite, batch_size, expected_sizes
+):
+    """Rehashed and repeated images retain metadata and per-occurrence transfers."""
+    from vllm.distributed.ec_transfer.ec_connector.mooncake.scheduler import (
+        ECMooncakeScheduler,
+    )
+
+    seen = []
+
+    class Session:
+        async def post(self, url, data=None, headers=None):
+            body = msgspec.json.decode(data)
+            seen.append((url, body))
+            params: dict[str, dict[str, Any]] = {}
+            for index, item in enumerate(body["messages"][0]["content"]):
+                image = item["image_url"]["url"]
+                entry = params.setdefault(
+                    image + "-processed",
+                    {
+                        "metadata": {"image_grid_thw": [1, 2, ord(image)]},
+                        "item_indices": [],
+                        "peer_port": 1234,
+                    },
+                )
+                entry["item_indices"].append(index)
+            # Response ordering must not affect the image-to-metadata mapping.
+            return _EncoderResponse(dict(reversed(list(params.items()))))
+
+    monkeypatch.setattr(proxy, "encode_session", Session())
+    monkeypatch.setattr(proxy, "NO_REWRITE", no_rewrite)
+    monkeypatch.setattr(proxy, "ENCODER_MAX_BATCH_SIZE", batch_size)
+    monkeypatch.setattr(proxy, "encoder_rr_idx", 0)
+    images = ["A", "A", "C", "D", "A", "F"]
+    body = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image_url", "image_url": {"url": image}}
+                    for image in images
+                ],
+            }
+        ],
+        "mm_processor_kwargs": {"max_pixels": 262144},
+    }
+    meta, handles = asyncio.run(
+        proxy.fanout_encoder_primer(
+            body, ["http://e0", "http://e1"], "r", "tcp://consumer:1"
+        )
+    )
+    assert [len(b["messages"][0]["content"]) for _, b in seen] == expected_sizes
+    assert (
+        len(
+            {
+                x["transfer_id"]
+                for _, b in seen
+                for x in b["ec_transfer_params"]["ec_items"]
+            }
+        )
+        == 6
+    )
+    for url, batch in seen:
+        assert batch["mm_processor_kwargs"] == body["mm_processor_kwargs"]
+        for item in batch["messages"][0]["content"]:
+            assert item["image_url"]["url"] in (
+                {"A", "C"} if url.startswith("http://e0/") else {"A", "D", "F"}
+            )
+    assert [meta[i]["image_grid_thw"][-1] for i in range(6)] == list(map(ord, images))
+    assert [meta[i]["ec_mm_hash"] for i in range(6)] == [
+        x + "-processed" for x in images
+    ]
+    assert handles["A-processed"]["peer_port"] == 1234
+    assert "ec_transfer_params" not in body
+
+    transfers_by_encoder: list[list[dict[str, str]]] = [[], []]
+    for url, batch in seen:
+        rank = 0 if url.startswith("http://e0/") else 1
+        transfers_by_encoder[rank].extend(batch["ec_transfer_params"]["ec_items"])
+    expected_transfers = [
+        transfers_by_encoder[i % 2][i // 2]["transfer_id"] for i in range(len(images))
+    ]
+    params = (
+        handles
+        if no_rewrite
+        else proxy.rewrite_for_decode(body, meta)["ec_transfer_params"]
+    )
+    request = SimpleNamespace(
+        ec_transfer_params=params,
+        mm_features=[
+            SimpleNamespace(identifier=image + "-processed") for image in images
+        ],
+    )
+    assert [
+        ECMooncakeScheduler._request_transfer_id(request, i) for i in range(len(images))
+    ] == expected_transfers
+
+
+def test_batch_rejects_ambiguous_metadata(proxy, monkeypatch):
+    monkeypatch.setattr(
+        proxy,
+        "encode_session",
+        _EncoderSession([{"unknown": {"metadata": {"image_grid_thw": [1, 2, 2]}}}]),
+    )
+    body = {
+        "messages": [
+            {
+                "content": [
+                    {"type": "image_url", "image_url": {"url": image}}
+                    for image in ["A", "B"]
+                ]
+            }
+        ]
+    }
+    with pytest.raises(proxy.HTTPException, match="cannot be matched"):
+        asyncio.run(proxy.fanout_encoder_primer(body, ["http://e0"], "r"))
+
+
 def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch):
     """The retry loop re-enters `prepare_for_decode` with the same body.
 
@@ -94,7 +300,7 @@ def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch
         _EncoderSession([{"encoder-side-hash": handle}, {}]),
     )
 
-    async def _no_prefill(req_data, p_url, req_id):
+    async def _no_prefill(req_data, p_url, req_id, dp_rank=None):
         return req_data
 
     monkeypatch.setattr(proxy, "maybe_prefill", _no_prefill)
@@ -112,7 +318,9 @@ async def _no_prefill(req_data, p_url, req_id):
 
     first, _, _ = asyncio.run(proxy.prepare_for_decode(body, *args))
     reported = first["ec_transfer_params"]
-    assert [handle] == [value for key, value in reported.items() if key != "ec_items"]
+    assert {key: value for key, value in reported.items() if key != "ec_items"} == {
+        "encoder-side-hash": handle
+    }
     assert "ec_transfer_params" not in body
 
     # Attempt 2's encode reports nothing: decode must be told nothing.
@@ -147,6 +355,192 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
     assert params["ec_items"][0]["transfer_id"]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("prefill", [False, True])
+@pytest.mark.parametrize("dynamic", [False, True])
+@pytest.mark.parametrize(
+    "request_options",
+    [
+        {},
+        {
+            "mm_processor_kwargs": {"max_pixels": 262144},
+            "media_io_kwargs": {"image": {"image_mode": "RGB"}},
+            "priority": -3,
+            "session_id": "session-123",
+        },
+    ],
+)
+async def test_http_roundtrip_preserves_payload_and_response_bytes(
+    proxy, monkeypatch, stream, prefill, dynamic, request_options
+):
+    """Preserve media and scheduling context across HTTP hops and retries."""
+    seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []}
+    prefix = b'data: {"content":"'
+    # Split a multibyte character at the old 1024-byte forwarding boundary.
+    expected = (
+        prefix + b"x" * (1023 - len(prefix)) + '图片🌍"}\n\ndata: [DONE]\n\n'.encode()
+        if stream
+        else '{ "choices": [{"message": {"content": "图片🌍"}}] }\n'.encode()
+    )
+
+    async def backend(request):
+        stage = request.match_info["stage"]
+        if dynamic:
+            consumer = "prefill" if prefill else "decode"
+            expected_rank = "0" if stage == consumer else None
+            assert request.headers.get("X-data-parallel-rank") == expected_rank
+        assert request.content_type == "application/json"
+        body = await request.json()
+        seen[stage].append(body)
+        if stage == "encode":
+            params: dict[str, dict[str, Any]] = {}
+            for index, item in enumerate(body["messages"][0]["content"]):
+                mm_hash = item["uuid"]
+                if body.get("mm_processor_kwargs"):
+                    mm_hash += "-processed"
+                entry = params.setdefault(
+                    mm_hash,
+                    {
+                        "metadata": {"image_grid_thw": [[1, 2, 2]]},
+                        "peer_port": len(seen[stage]),
+                        "item_indices": [],
+                    },
+                )
+                entry["item_indices"].append(index)
+            return web.json_response({"ec_transfer_params": params})
+        if stage == "prefill":
+            assert body["max_tokens"] == 1 and body["stream"] is False
+            return web.json_response(
+                {"kv_transfer_params": {"remote_block_ids": [len(seen[stage])]}}
+            )
+        if len(seen[stage]) == 1:
+            return web.Response(status=500, text="retry")
+        return web.Response(
+            body=expected,
+            content_type="text/event-stream" if stream else "application/json",
+        )
+
+    backend_app = web.Application()
+    backend_app.router.add_post("/{stage}/v1/chat/completions", backend)
+    async with TestServer(backend_app) as server:
+        for stage, field in [
+            ("encode", "e_urls"),
+            ("prefill", "p_urls"),
+            ("decode", "d_urls"),
+        ]:
+            urls = [str(server.make_url(f"/{stage}"))]
+            if stage == "prefill" and not prefill:
+                urls = []
+            monkeypatch.setattr(proxy.app.state, field, urls, raising=False)
+        monkeypatch.setattr(proxy.app.state, "d_ec_urls", [], raising=False)
+        monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False)
+        monkeypatch.setattr(proxy, "DECODE_RETRIES", 1)
+        monkeypatch.setattr(proxy, "NO_REWRITE", False)
+        target_app = proxy.app
+        if dynamic:
+            monkeypatch.setenv("ADMIN_API_KEY", "test-key")
+            target_app = proxy.build_app(proxy.EPDProxyConfig(probe_interval=0))
+        item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}}
+        body = {
+            "model": "test",
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "描述图片🌍"},
+                        item,
+                        item,
+                    ],
+                }
+            ],
+            "stream": stream,
+            "temperature": 0.01,
+            "max_tokens": 32,
+            "seed": 42,
+            "structured_outputs": {"choice": ["A", "B"]},
+            **request_options,
+        }
+        await proxy.on_startup()
+        try:
+            async with httpx.AsyncClient(
+                transport=httpx.ASGITransport(app=target_app), base_url="http://proxy"
+            ) as client:
+                if dynamic:
+                    registrations = [("encode", "encode")]
+                    registrations += (
+                        [("prefill", "prefill"), ("decode", "decode")]
+                        if prefill
+                        else [("prefill_decode", "decode")]
+                    )
+                    for role, stage in registrations:
+                        consumer = "prefill" if prefill else "decode"
+                        registered = await client.post(
+                            "/instances",
+                            headers={"X-API-Key": "test-key"},
+                            json={
+                                "role": role,
+                                "url": str(server.make_url(f"/{stage}")),
+                                "dp_size": 2 if stage == consumer else 1,
+                            },
+                        )
+                        assert registered.status_code == 200
+                response = await client.post("/v1/chat/completions", json=body)
+                if dynamic:
+                    removed = await client.delete(
+                        "/instances",
+                        headers={"X-API-Key": "test-key"},
+                        params={"url": str(server.make_url("/decode"))},
+                    )
+                    assert removed.json() == {"removed": True}
+                    unavailable = await client.post("/v1/chat/completions", json=body)
+                    assert unavailable.status_code == 503
+            assert response.status_code == 200
+            assert response.content == expected
+            assert response.headers["content-type"].startswith(
+                "text/event-stream" if stream else "application/json"
+            )
+        finally:
+            await proxy.on_shutdown()
+
+    assert len(seen["encode"]) == 2
+    assert len(seen["decode"]) == 2
+    for requests in seen.values():
+        for forwarded in requests:
+            for key in (
+                "mm_processor_kwargs",
+                "media_io_kwargs",
+                "priority",
+                "session_id",
+            ):
+                assert forwarded.get(key) == body.get(key)
+    final = seen["decode"][-1]
+    for key in [
+        "model",
+        "stream",
+        "temperature",
+        "max_tokens",
+        "seed",
+        "structured_outputs",
+    ]:
+        assert final[key] == body[key]
+    content = final["messages"][0]["content"]
+    assert content[0] == body["messages"][0]["content"][0]
+    for reference in content[1:]:
+        assert reference == {
+            "type": "image_embeds",
+            "image_embeds": {"image_grid_thw": [1, 2, 2]},
+            "uuid": proxy.content_uuid(item),
+        }
+    mm_hash = proxy.content_uuid(item)
+    if request_options:
+        mm_hash += "-processed"
+    assert set(final["ec_transfer_params"]) - {"ec_items"} == {mm_hash}
+    assert final["ec_transfer_params"][mm_hash]["peer_port"] == 2
+    if prefill:
+        assert final["kv_transfer_params"] == {"remote_block_ids": [2]}
+
+
 @pytest.mark.parametrize("server_keep_alive", ["0.1", "1", "5", "2", "30"])
 def test_pooled_connections_are_retired_before_the_server_closes_them(
     proxy, monkeypatch, server_keep_alive
@@ -170,3 +564,66 @@ def test_pooled_connections_are_retired_before_the_server_closes_them(
         assert proxy.decode_session.connector._keepalive_timeout == pooled_for
     finally:
         asyncio.run(proxy.on_shutdown())
+
+
+@pytest.mark.parametrize(
+    "no_rewrite, transfer",
+    [(False, None), (True, None), (False, "push"), (False, "handle")],
+)
+def test_video_audio_fallback(proxy, monkeypatch, no_rewrite, transfer):
+    """Keep raw video and neighboring images, but never discard transfer handles."""
+    import copy
+
+    from vllm.distributed.ec_transfer.ec_connector.utils import collect_ec_item_metadata
+
+    video = {"type": "video_url", "video_url": {"url": "video-with-audio"}}
+    images = [{"type": "image_url", "image_url": {"url": key}} for key in ("A", "B")]
+    body: dict[str, Any] = {
+        "messages": [{"content": [images[0], video, images[1]]}],
+        "mm_processor_kwargs": {"use_audio_in_video": True},
+    }
+    original = copy.deepcopy(body)
+    # Collector indices describe two processed features from one video item.
+    video_metadata = collect_ec_item_metadata(
+        [
+            SimpleNamespace(identifier=key, modality=key, data=None)
+            for key in ("video", "audio")
+        ],
+        SimpleNamespace(fields_for=lambda _: set()),
+    )
+    if transfer == "handle":
+        video_metadata["audio"]["transfer_id"] = "reservation"
+    image_metadata = {
+        key: {"metadata": {"image_grid_thw": [1, 2, size]}, "item_indices": [i]}
+        for i, (key, size) in enumerate((("A", 2), ("B", 4)))
+    }
+    # Images form the first group, the singleton video the second.
+    monkeypatch.setattr(proxy, "ENCODER_MAX_BATCH_SIZE", 0)
+    monkeypatch.setattr(proxy, "NO_REWRITE", no_rewrite)
+    replies = _EncoderSession([image_metadata, video_metadata])
+    monkeypatch.setattr(proxy, "encode_session", replies)
+    consumer = "tcp://consumer:1" if transfer == "push" else None
+    preparation = proxy.prepare_for_decode(body, "r", ["http://e0"], "", consumer)
+    if transfer:
+        with pytest.raises(proxy.HTTPException, match="cannot be matched"):
+            asyncio.run(preparation)
+        return
+
+    prepared, _, _ = asyncio.run(preparation)
+    assert body == original
+    assert prepared["mm_processor_kwargs"] == body["mm_processor_kwargs"]
+    final = prepared["messages"][0]["content"]
+    assert final[1] == video
+    handles = prepared["ec_transfer_params"]
+    assert not video_metadata.keys() & handles.keys()
+    if no_rewrite:
+        assert final == original["messages"][0]["content"]
+        assert "ec_items" not in handles
+    else:
+        for position, image, key in ((0, images[0], "A"), (2, images[1], "B")):
+            assert final[position] == {
+                "type": "image_embeds",
+                "image_embeds": image_metadata[key]["metadata"],
+                "uuid": proxy.content_uuid(image),
+            }
+        assert [item["mm_hash"] for item in handles["ec_items"]] == ["A", "B"]
diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py
index 558a80f9cf82..8e1a730ca6b8 100644
--- a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py
+++ b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py
@@ -14,6 +14,7 @@
 """
 
 import importlib.util
+import sys
 from collections import Counter
 from pathlib import Path
 
@@ -27,6 +28,7 @@ def _load_proxy_module():
     spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path)
     assert spec is not None and spec.loader is not None
     module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
     spec.loader.exec_module(module)
     return module
 
diff --git a/tests/v1/ec_connector/unit/test_epd_registry.py b/tests/v1/ec_connector/unit/test_epd_registry.py
new file mode 100644
index 000000000000..be5a8e9b9b3e
--- /dev/null
+++ b/tests/v1/ec_connector/unit/test_epd_registry.py
@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dynamic registration and liveness for the EPD proxy's instance registry."""
+
+from unittest.mock import patch
+
+import pytest
+
+from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import (
+    InstanceRecord,
+    InstanceRegistry,
+    InstanceRole,
+)
+
+ENCODE = InstanceRole.ENCODE
+DECODE = InstanceRole.DECODE
+
+
+@pytest.fixture
+def registry():
+    return InstanceRegistry(probe_interval=0, fail_threshold=3, evicted_ttl=60.0)
+
+
+def _fill(registry, role, count, prefix="http://h"):
+    for index in range(count):
+        registry.register(InstanceRecord(role, f"{prefix}{index}:8000"))
+
+
+async def _probe_round(registry, healthy: set[str], now: float = 0.0):
+    """Run one probe pass where only `healthy` URLs answer."""
+
+    async def fake_probe(self, session, url):
+        return url in healthy
+
+    with (
+        patch.object(InstanceRegistry, "_probe", fake_probe),
+        patch(
+            "examples.disaggregated.disaggregated_encoder.disagg_epd_proxy.time.monotonic",
+            return_value=now,
+        ),
+    ):
+        await registry._probe_once(session=None)
+
+
+class TestRegistration:
+    def test_registering_is_idempotent_and_reports_novelty(self, registry):
+        record = InstanceRecord(ENCODE, "http://e0:8000")
+        assert registry.register(record) is True
+        assert registry.register(record) is False
+        assert registry.urls(ENCODE) == ["http://e0:8000"]
+
+    def test_transfer_addresses_stay_with_the_registered_consumer(self, registry):
+        """Roster updates must not shift a consumer's transfer addresses."""
+        registry.register(InstanceRecord(DECODE, "http://d0:8000"))
+        registry.register(
+            InstanceRecord(
+                DECODE,
+                "http://d0:8000",
+                ec_zmq_addrs=["tcp://d0:20001", "tcp://d0:20002"],
+                dp_size=2,
+            )
+        )
+        picked = registry.pick(DECODE)
+        assert picked.ec_zmq_addrs == ["tcp://d0:20001", "tcp://d0:20002"]
+        assert picked.dp_size == 2
+
+    def test_unregister_is_terminal(self, registry):
+        registry.register(InstanceRecord(ENCODE, "http://e0:8000"))
+        assert registry.unregister("http://e0:8000") is True
+        assert registry.unregister("http://e0:8000") is False
+        assert registry.pick(ENCODE) is None
+
+    def test_round_robin_survives_registration(self, registry):
+        _fill(registry, ENCODE, 2)
+        assert [record.url for record in registry.pick_many(ENCODE, 2)] == [
+            "http://h0:8000",
+            "http://h1:8000",
+        ]
+        registry.register(InstanceRecord(ENCODE, "http://h2:8000"))
+        assert registry.pick(ENCODE).url == "http://h2:8000"
+
+
+class TestLiveness:
+    @pytest.mark.asyncio
+    async def test_identical_reregistration_preserves_inflight_probe_results(
+        self, registry
+    ):
+        """Repeated registrations must not suppress eviction or recovery."""
+        url = "http://e0:8000"
+        registry.register(InstanceRecord(ENCODE, url))
+        healthy = False
+
+        async def probe(self, session, url):
+            registry.register(InstanceRecord(ENCODE, url))
+            return healthy
+
+        with patch.object(InstanceRegistry, "_probe", probe):
+            for _ in range(registry._fail_threshold):
+                await registry._probe_once(None)
+            assert registry.status()["encode"] == {"live": [], "evicted": [url]}
+
+            healthy = True
+            await registry._probe_once(None)
+            assert registry.status()["encode"] == {"live": [url], "evicted": []}
+
+    @pytest.mark.asyncio
+    async def test_consecutive_failures_stop_routing(self, registry):
+        registry.register(InstanceRecord(ENCODE, "http://e0:8000"))
+        registry.register(InstanceRecord(ENCODE, "http://e1:8000"))
+        for _ in range(registry._fail_threshold):
+            await _probe_round(registry, healthy={"http://e1:8000"})
+        assert registry.urls(ENCODE) == ["http://e1:8000"]
+        assert registry.status()["encode"]["evicted"] == ["http://e0:8000"]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("role", [ENCODE, DECODE])
+    async def test_a_recovered_instance_rejoins_without_re_registering(
+        self, registry, role
+    ):
+        """Restarting every encoder to recover from a blip is not acceptable."""
+        registry.register(InstanceRecord(role, "http://e0:8000"))
+        for _ in range(registry._fail_threshold):
+            await _probe_round(registry, healthy=set())
+        assert registry.urls(role) == []
+
+        await _probe_round(registry, healthy={"http://e0:8000"}, now=1)
+        assert registry.urls(role) == ["http://e0:8000"]
+
+    @pytest.mark.asyncio
+    async def test_re_registration_does_not_override_failed_health_checks(
+        self, registry
+    ):
+        record = InstanceRecord(ENCODE, "http://e0:8000")
+        registry.register(record)
+        for _ in range(registry._fail_threshold):
+            await _probe_round(registry, healthy=set())
+        registry.register(record)
+        assert registry.urls(ENCODE) == []
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("healthy", [False, True])
+    async def test_inflight_probe_cannot_restore_unregistered_instance(
+        self, registry, healthy
+    ):
+        record = InstanceRecord(ENCODE, "http://e0:8000")
+        registry.register(record)
+        for _ in range(registry._fail_threshold):
+            await _probe_round(registry, healthy=set())
+
+        async def probe(self, session, url):
+            registry.unregister(url)
+            return healthy
+
+        with patch.object(InstanceRegistry, "_probe", probe):
+            await registry._probe_once(None)
+        assert registry.status()["encode"] == {"live": [], "evicted": []}
diff --git a/tests/v1/ec_connector/unit/test_metadata.py b/tests/v1/ec_connector/unit/test_metadata.py
index cca1df017193..65fdc6bbb392 100644
--- a/tests/v1/ec_connector/unit/test_metadata.py
+++ b/tests/v1/ec_connector/unit/test_metadata.py
@@ -54,7 +54,7 @@ def test_metadata_loads_carry_a_transfer_id_and_blocks():
 
 
 def test_metadata_saves_and_loads_are_independent():
-    """saves and loads dicts are independent."""
+    """Saves and loads dicts are independent."""
     meta = ECCPUConnectorMetadata()
 
     meta.saves["key"] = [1, 2]
diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py
index 156d98ae794c..d7ef19b5f4f4 100644
--- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py
+++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py
@@ -67,6 +67,7 @@ def test_request_finished_producer_emits_params(monkeypatch):
     assert params == {
         "h1": {
             "metadata": {},
+            "item_indices": [0],
             "peer_host": "1.2.3.4",
             "peer_port": 5601,
             "size_bytes": 2 * 32 * 2,
@@ -91,6 +92,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch):
     assert params == {
         "h1": {
             "metadata": {},
+            "item_indices": [0],
             "peer_host": "1.2.3.4",
             "peer_port": 5601,
             "size_bytes": 2 * 32 * 2,
@@ -110,5 +112,5 @@ def test_request_finished_skips_unallocated_entry(monkeypatch):
     assert delay is False
     # The item's placeholder metadata is still reported even though there's
     # no cache entry to transfer (empty "metadata": no fields, no transfer).
-    assert params == {"h1": {"metadata": {}}}
+    assert params == {"h1": {"metadata": {}, "item_indices": [0]}}
     s.shutdown()
diff --git a/tests/v1/ec_connector/unit/test_session.py b/tests/v1/ec_connector/unit/test_session.py
index 4d78b0426afc..5c25fdf2288d 100644
--- a/tests/v1/ec_connector/unit/test_session.py
+++ b/tests/v1/ec_connector/unit/test_session.py
@@ -329,7 +329,6 @@ def test_consumer_xfer_quarantined_exception_settles():
 
 def test_consumer_xfer_quarantined_timeout_evicts_blocks():
     """Quarantine timeout should release handle and return READ_FAILED."""
-
     data = _make_data(xfer_state="PROC")
     x = _started(data)
     # First poll transitions to QUARANTINED (deadline extended by 60s)
diff --git a/tests/v1/ec_connector/unit/test_worker_ec_connector.py b/tests/v1/ec_connector/unit/test_worker_ec_connector.py
index 01dad9d89ebd..911810b53acf 100644
--- a/tests/v1/ec_connector/unit/test_worker_ec_connector.py
+++ b/tests/v1/ec_connector/unit/test_worker_ec_connector.py
@@ -20,6 +20,7 @@
 pytestmark = pytest.mark.cpu_test
 
 WORKER_META = object()
+CONNECTOR_STATS = object()
 
 
 def _scheduler_output() -> SimpleNamespace:
@@ -38,6 +39,7 @@ def _connector(
     fake.is_consumer = is_consumer
     fake.get_finished.return_value = (None, None)
     fake.build_connector_worker_meta.return_value = WORKER_META
+    fake.get_ec_connector_stats.return_value = CONNECTOR_STATS
     with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake):
         return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake
 
@@ -72,6 +74,16 @@ def test_worker_meta_is_reported_on_context_exit():
     assert fake.clear_connector_metadata.called
 
 
+def test_connector_stats_are_reported_on_context_exit():
+    """Reported in the finally block, alongside the worker meta."""
+    connector, fake = _connector()
+
+    with connector.maybe_get_output(_scheduler_output()) as output:
+        assert output.ec_connector_stats is None
+
+    assert output.ec_connector_stats is CONNECTOR_STATS
+
+
 def test_v1_producer_receives_existing_encoder_cache():
     encoder_cache = {"mm_cached": None}
     fake = MagicMock(spec=ECConnectorBase)
diff --git a/tests/v1/ec_connector/unit/utils.py b/tests/v1/ec_connector/unit/utils.py
index 93b97cb5ef53..a718bc63b635 100644
--- a/tests/v1/ec_connector/unit/utils.py
+++ b/tests/v1/ec_connector/unit/utils.py
@@ -45,6 +45,7 @@ def create_ec_vllm_config(
 
     Returns:
         A `VllmConfig`-specced mock carrying the configs above.
+
     """
     parallel_config = ParallelConfig(
         tensor_parallel_size=tensor_parallel_size,
diff --git a/tests/v1/engine/conftest.py b/tests/v1/engine/conftest.py
index 283a76dab672..1cfe671a62fd 100644
--- a/tests/v1/engine/conftest.py
+++ b/tests/v1/engine/conftest.py
@@ -24,12 +24,12 @@
 
 
 def _build_test_vectors_no_logprobs() -> DummyOutputProcessorTestVectors:
-    """Generate output processor dummy test vectors, without logprobs
+    """Generate output processor dummy test vectors, without logprobs.
 
     Returns:
       DummyOutputProcessorTestVectors instance with no logprobs
-    """
 
+    """
     tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
     vllm_config = EngineArgs(model=TOKENIZER_NAME).create_engine_config()
     # Tokenize prompts under test & create dummy generated tokens
@@ -62,10 +62,11 @@ def _build_test_vectors_no_logprobs() -> DummyOutputProcessorTestVectors:
 
 @pytest.fixture
 def dummy_test_vectors() -> DummyOutputProcessorTestVectors:
-    """Generate output processor dummy test vectors, with logprobs
+    """Generate output processor dummy test vectors, with logprobs.
 
     Returns:
       DummyOutputProcessorTestVectors instance with logprobs
+
     """
     # Build dummy test vectors without logprobs
     dtv = _build_test_vectors_no_logprobs()
diff --git a/tests/v1/engine/test_abort_final_step.py b/tests/v1/engine/test_abort_final_step.py
index d8d5b73d45b9..c97e1f743637 100644
--- a/tests/v1/engine/test_abort_final_step.py
+++ b/tests/v1/engine/test_abort_final_step.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Test for the fix in PR #29987: Eagerly abort cancelled final-step requests.
+"""Test for the fix in PR #29987: Eagerly abort cancelled final-step requests.
 
 This test verifies that when a request is aborted during its final execution
 step (when it would naturally complete), it is properly marked as aborted
@@ -52,8 +51,7 @@ def __init__(self):
 
 
 class DummyKVConnector(KVConnectorBase_V1):
-    """
-    Dummy KV connector that captures request finish statuses to a file.
+    """Dummy KV connector that captures request finish statuses to a file.
     This is used to verify the fix - without the fix, a request aborted
     during its final step would be captured as FINISHED_LENGTH_CAPPED
     instead of FINISHED_ABORTED.
@@ -144,8 +142,7 @@ def wait_for_save(self):
 @pytest.mark.parametrize("async_scheduling", [False, True])
 @pytest.mark.asyncio
 async def test_abort_during_final_step(async_scheduling: bool):
-    """
-    Test that a request aborted during its final execution step is treated as
+    """Test that a request aborted during its final execution step is treated as
     aborted rather than completed.
 
     This test:
@@ -165,7 +162,6 @@ async def test_abort_during_final_step(async_scheduling: bool):
     With the fix, _process_aborts_queue() runs before update_from_output(), so the
     abort takes precedence and the KV connector sees FINISHED_ABORTED.
     """
-
     # Create three temporary files:
     # 1. ready_file: deleted by execute_model to signal it has started
     # 2. block_file: execute_model waits for this to be deleted
diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py
index a03718ab5a15..15be166c8c0e 100644
--- a/tests/v1/engine/test_admission_control.py
+++ b/tests/v1/engine/test_admission_control.py
@@ -14,6 +14,8 @@
 
 import argparse
 import asyncio
+import multiprocessing
+from collections.abc import Awaitable, Callable
 from http import HTTPStatus
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
@@ -34,6 +36,8 @@
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
 from vllm.utils.argparse_utils import human_readable_int
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.admission_control import SharedAdmissionStats
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.output_processor import OutputProcessor
 
@@ -66,6 +70,7 @@ def _make_async_llm(
     llm.output_processor = MagicMock()
     llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished
     llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens
+    llm.admission_stats = None
     return llm
 
 
@@ -75,6 +80,45 @@ def _make_output_processor(**request_states) -> OutputProcessor:
     return op
 
 
+def _make_engine_request(request_id: str, n: int) -> EngineCoreRequest:
+    return EngineCoreRequest(
+        request_id=request_id,
+        prompt_token_ids=[1],
+        mm_features=None,
+        sampling_params=SamplingParams(n=n),
+        pooling_params=None,
+        arrival_time=0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+
+
+def _make_request_test_llm(
+    max_num_queued_reqs: int,
+    add_request_async: Callable[[EngineCoreRequest], Awaitable[None]],
+) -> AsyncLLM:
+    llm = _make_async_llm(max_num_queued_reqs=max_num_queued_reqs)
+    llm.output_processor = OutputProcessor(None, log_stats=False)
+    llm.engine_core = SimpleNamespace(
+        resources=SimpleNamespace(engine_dead=False),
+        add_request_async=AsyncMock(side_effect=add_request_async),
+        abort_requests_async=AsyncMock(),
+        shutdown=MagicMock(),
+    )
+    llm.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(kv_sharing_fast_prefill=False)
+    )
+    llm.output_handler = None
+    llm.log_requests = False
+    llm._run_output_handler = MagicMock()
+    llm.input_processor = MagicMock()
+    llm.input_processor.assign_request_id.side_effect = lambda request: setattr(
+        request, "external_req_id", request.request_id
+    )
+    return llm
+
+
 def _make_scheduler_config(**kwargs) -> SchedulerConfig:
     return SchedulerConfig(
         runner_type="generate",
@@ -84,6 +128,42 @@ def _make_scheduler_config(**kwargs) -> SchedulerConfig:
     )
 
 
+def _make_shared_stats(client_count: int = 2) -> list[SharedAdmissionStats]:
+    counters = multiprocessing.RawArray(
+        "q", SharedAdmissionStats.num_counters(client_count)
+    )
+    return [
+        SharedAdmissionStats({"mp_admission_counters": counters}, client_count, index)
+        for index in range(client_count)
+    ]
+
+
+@pytest.mark.skip_global_cleanup
+def test_shared_admission_stats_aggregate_api_servers():
+    server_0, server_1 = _make_shared_stats()
+
+    server_0.set_num_requests(2)
+    assert server_0.get_num_requests() == 2
+    server_1.set_num_requests(3)
+    assert server_0.get_num_requests() == server_1.get_num_requests() == 5
+
+    server_0.set_num_requests(1)
+    assert server_0.get_num_requests() == 4
+
+
+@pytest.mark.skip_global_cleanup
+def test_shared_request_snapshot_enforces_global_limit():
+    server_0, server_1 = _make_shared_stats()
+    server_0.set_num_requests(2)
+    llm = _make_async_llm(max_num_queued_reqs=2)
+    llm.admission_stats = server_1
+
+    with pytest.raises(QueueOverflowError):
+        llm.check_admission()
+
+    assert server_0.get_num_requests() == 2
+
+
 # ---------------------------------------------------------------------------
 # Exception classes
 # ---------------------------------------------------------------------------
@@ -301,6 +381,55 @@ async def add_request_async(_request):
     assert len(request_states) == 1
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("first_n", [1, 3])
+async def test_parallel_admission_is_all_or_nothing(first_n: int):
+    """An n>1 request reserves either all its capacity or none of it."""
+
+    async def add_request_async(_request):
+        await asyncio.sleep(0)
+
+    llm = _make_request_test_llm(3, add_request_async)
+    second_n = 3 if first_n == 1 else 1
+    first = _make_engine_request("first", first_n)
+    second = _make_engine_request("second", second_n)
+    results = await asyncio.gather(
+        llm.add_request("first", first, first.params),
+        llm.add_request("second", second, second.params),
+        return_exceptions=True,
+    )
+
+    assert not isinstance(results[0], Exception)
+    assert isinstance(results[1], QueueOverflowError)
+    assert llm.output_processor.get_num_unfinished_requests() == first_n
+    assert llm.engine_core.add_request_async.await_count == first_n
+
+
+@pytest.mark.asyncio
+async def test_parallel_admission_cancellation_cleans_up_all_children():
+    """Cancellation during core submission must release every reserved slot."""
+    first_submission_started = asyncio.Event()
+
+    async def add_request_async(_request):
+        first_submission_started.set()
+        await asyncio.Event().wait()
+
+    llm = _make_request_test_llm(3, add_request_async)
+    request = _make_engine_request("parallel", 3)
+    output = llm.generate(request, request.params, request.request_id)
+    generate_task = asyncio.create_task(anext(output))
+
+    await first_submission_started.wait()
+    generate_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await generate_task
+
+    assert llm.engine_core.add_request_async.await_count == 1
+    assert llm.output_processor.get_num_unfinished_requests() == 0
+    aborted_ids = llm.engine_core.abort_requests_async.await_args.args[0]
+    assert set(aborted_ids) == {"0_parallel", "1_parallel", "2_parallel"}
+
+
 # ---------------------------------------------------------------------------
 # create_error_response integration
 # ---------------------------------------------------------------------------
diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py
index 5dd37a6f7ed4..cac43d8204d0 100644
--- a/tests/v1/engine/test_async_llm.py
+++ b/tests/v1/engine/test_async_llm.py
@@ -60,13 +60,7 @@
 }
 
 
-def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch):
-    vllm_config = MagicMock()
-    vllm_config.observability_config.otlp_traces_endpoint = None
-    vllm_config.scheduler_config.stream_interval = 1
-    vllm_config.profiler_config.profiler = "cuda"
-    vllm_config.profiler_config.ignore_frontend = False
-
+def _mock_async_llm_dependencies(monkeypatch: pytest.MonkeyPatch):
     renderer = MagicMock()
     engine_core = MagicMock()
     engine_core.profile_async = AsyncMock()
@@ -90,6 +84,16 @@ def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatc
         "make_async_mp_client",
         MagicMock(return_value=engine_core),
     )
+    return engine_core
+
+
+def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch):
+    vllm_config = MagicMock()
+    vllm_config.observability_config.otlp_traces_endpoint = None
+    vllm_config.scheduler_config.stream_interval = 1
+    vllm_config.profiler_config.profiler = "cuda"
+    vllm_config.profiler_config.ignore_frontend = False
+    engine_core = _mock_async_llm_dependencies(monkeypatch)
 
     engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False)
 
@@ -103,6 +107,25 @@ async def profile():
     engine_core.profile_async.assert_has_awaits([call(True, None), call(False)])
 
 
+def test_cuda_only_torch_profiler_skips_frontend_cpu_trace(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    vllm_config = MagicMock()
+    vllm_config.observability_config.otlp_traces_endpoint = None
+    vllm_config.scheduler_config.stream_interval = 1
+    vllm_config.profiler_config.profiler = "torch"
+    vllm_config.profiler_config.ignore_frontend = False
+    vllm_config.profiler_config.torch_profiler_activities = ["CUDA"]
+    _mock_async_llm_dependencies(monkeypatch)
+    profiler = MagicMock()
+    monkeypatch.setattr(async_llm_module, "TorchProfilerWrapper", profiler)
+
+    engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False)
+
+    assert engine.profiler is None
+    profiler.assert_not_called()
+
+
 async def generate(
     engine: AsyncLLM,
     request_id: str,
@@ -457,7 +480,6 @@ async def test_customize_loggers(monkeypatch):
     If a customized logger is provided at the init, it should
     be added to the default loggers.
     """
-
     with ExitStack() as after:
         with set_default_torch_num_threads(1):
             engine = AsyncLLM.from_engine_args(
@@ -642,7 +664,6 @@ async def test_check_health():
 @pytest.mark.asyncio
 async def test_abort_final_output(output_kind: RequestOutputKind):
     """Test that abort() returns a final output with correct information."""
-
     with ExitStack() as after:
         with set_default_torch_num_threads(1):
             engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
diff --git a/tests/v1/engine/test_engine_args.py b/tests/v1/engine/test_engine_args.py
index b1ba1adcfd0a..3f9f617bc7ed 100644
--- a/tests/v1/engine/test_engine_args.py
+++ b/tests/v1/engine/test_engine_args.py
@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from argparse import ArgumentError
+from unittest.mock import patch
 
 import pytest
 
+from vllm.config import ModelConfig
 from vllm.engine.arg_utils import EngineArgs
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -121,3 +123,76 @@ def test_data_parallel_start_rank_zero_infers_hybrid_lb():
 
     assert vllm_config.parallel_config.data_parallel_hybrid_lb is True
     assert vllm_config.parallel_config.data_parallel_rank == 0
+
+
+def test_external_lb_preserves_explicit_rank_when_dp_exceeds_nodes():
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        data_parallel_size=4,
+        data_parallel_rank=3,
+        data_parallel_external_lb=True,
+        nnodes=2,
+        node_rank=1,
+    )
+
+    with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
+        vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+    assert vllm_config.parallel_config.data_parallel_rank == 3
+
+
+@pytest.mark.parametrize(
+    ("data_parallel_size", "nnodes", "tensor_parallel_size"),
+    [(4, 2, 1), (2, 3, 3)],
+)
+def test_external_lb_requires_explicit_rank_when_nodes_are_not_evenly_partitioned(
+    data_parallel_size, nnodes, tensor_parallel_size
+):
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        data_parallel_size=data_parallel_size,
+        data_parallel_external_lb=True,
+        tensor_parallel_size=tensor_parallel_size,
+        nnodes=nnodes,
+        node_rank=1,
+    )
+
+    with (
+        patch.object(ModelConfig, "is_moe", new=property(lambda self: True)),
+        pytest.raises(
+            ValueError,
+            match="Set a unique `--data-parallel-rank`",
+        ),
+    ):
+        engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+
+def test_external_lb_infers_rank_when_dp_does_not_exceed_nodes():
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        data_parallel_size=2,
+        data_parallel_external_lb=True,
+        nnodes=2,
+        node_rank=1,
+    )
+
+    with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
+        vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+    assert vllm_config.parallel_config.data_parallel_rank == 1
+
+
+def test_external_lb_infers_rank_for_multinode_replicas():
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        data_parallel_size=2,
+        data_parallel_external_lb=True,
+        tensor_parallel_size=12,
+        nnodes=4,
+        node_rank=1,
+    )
+
+    with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
+        vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+    assert vllm_config.parallel_config.data_parallel_rank == 0
diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py
index 98eda4403271..e16e0550c21d 100644
--- a/tests/v1/engine/test_engine_core.py
+++ b/tests/v1/engine/test_engine_core.py
@@ -23,6 +23,7 @@
 from vllm.engine.arg_utils import EngineArgs
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_default_torch_num_threads
+from vllm.v1.core.sched.interface import PauseState
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core import DPEngineCoreProc, EngineCore, EngineCoreProc
 from vllm.v1.executor.abstract import Executor
@@ -193,8 +194,7 @@ def test_engine_core():
 
 @create_new_process_for_each_test()
 def test_engine_core_advanced_sampling():
-    """
-    A basic end-to-end test to verify that the engine functions correctly
+    """A basic end-to-end test to verify that the engine functions correctly
     when additional sampling parameters, such as top_p, min_tokens, and
     presence_penalty, are set.
     """
@@ -242,9 +242,7 @@ def _check_engine_state():
 
 @create_new_process_for_each_test()
 def test_engine_core_concurrent_batches():
-    """
-    Test that the engine can handle multiple concurrent batches.
-    """
+    """Test that the engine can handle multiple concurrent batches."""
 
     def make_request_with_max_tokens(req_id: str, max_tokens: int) -> EngineCoreRequest:
         request = make_request()
@@ -265,7 +263,6 @@ def execute_model(
             non_block=False,
         ) -> Future[ModelRunnerOutput | None]:
             """Make execute_model non-blocking."""
-
             # DummyExecutor used only for testing async case.
             assert non_block
 
@@ -282,7 +279,6 @@ def sample_tokens(
             self, grammar_output, non_block=False
         ) -> Future[ModelRunnerOutput]:
             """Make sample_tokens non-blocking."""
-
             # DummyExecutor used only for testing async case.
             assert non_block
 
@@ -401,10 +397,7 @@ def shutdown(self):
 
 @multi_gpu_test(num_gpus=2)
 def test_engine_core_tp():
-    """
-    Test engine can initialize worker in tp properly
-    """
-
+    """Test engine can initialize worker in tp properly."""
     """Setup the EngineCore."""
     engine_args = EngineArgs(
         model=MODEL_NAME,
@@ -490,7 +483,7 @@ def test_encoder_instance_zero_kv_cache(
     enable_prefix_caching: bool,
     use_kv_connector: bool,
 ):
-    """EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests
+    """EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests.
 
     This test verifies encoder-only instance initializes with 0 KV cache blocks.
     Under EPD disagg mode, Encoder instances (EC producer role) only execute
@@ -675,6 +668,29 @@ def _pausable_engine_core_proc() -> EngineCoreProc:
     return core
 
 
+@pytest.mark.parametrize(
+    "pause_state,has_requests,has_batches",
+    [
+        pytest.param(PauseState.UNPAUSED, False, False, id="not-paused"),
+        pytest.param(PauseState.PAUSED_ALL, True, False, id="pending-requests"),
+        pytest.param(PauseState.PAUSED_ALL, False, True, id="pending-batches"),
+    ],
+)
+def test_kv_cache_release_rejects_unsafe_state(pause_state, has_requests, has_batches):
+    """Reject release before touching caches or memory if work can still use KV."""
+    core = _pausable_engine_core_proc()
+    core.scheduler.pause_state = pause_state
+    core.scheduler.has_requests.return_value = has_requests
+    core.batch_queue = [object()] if has_batches else None
+    core._reset_caches = MagicMock()
+
+    with pytest.raises(RuntimeError, match="requires a completed pause"):
+        core.release_kv_cache_memory()
+
+    core._reset_caches.assert_not_called()
+    core.model_executor.discard.assert_not_called()
+
+
 @pytest.mark.parametrize("deferred", [False, True])
 def test_pause_synchronizes_device_before_cache_reset(deferred: bool):
     """A resolved pause promises an idle device: the barrier must run before
diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py
index f97ed76c5c48..15dec72a4195 100644
--- a/tests/v1/engine/test_engine_core_client.py
+++ b/tests/v1/engine/test_engine_core_client.py
@@ -158,7 +158,10 @@ def setsockopt(self, *_args, **_kwargs):
         local_engines_only=False,
         enable_elastic_ep=False,
     )
-    vllm_config = SimpleNamespace(parallel_config=parallel_config)
+    vllm_config = SimpleNamespace(
+        parallel_config=parallel_config,
+        cache_config=SimpleNamespace(effective_attention_block_size=None),
+    )
 
     client = core_client_mod.MPClient(
         asyncio_mode=False,
@@ -211,6 +214,7 @@ def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncM
     client.core_engines = [bytes([i, 0]) for i in range(num_engines)]
     client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)]
     client.eng_start_index = 0
+    client._kv_event_sources = {}
     return client
 
 
@@ -368,10 +372,15 @@ def test_dplb_finished_requests_release_inflight():
     assert req.request_id not in client.reqs_in_flight
 
 
-def test_apply_ready_response_syncs_block_size():
+@pytest.mark.parametrize(
+    ("effective_size", "other_size"),
+    [(None, None), (4224, 4224), (4224, 1056), (4224, None)],
+)
+def test_apply_ready_response_syncs_block_size(effective_size, other_size):
     import msgspec
 
     client = object.__new__(MPClient)
+    client._effective_attention_block_sizes = set()
     client.vllm_config = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),
         model_config=SimpleNamespace(max_model_len=8192),
@@ -399,14 +408,31 @@ def test_apply_ready_response_syncs_block_size():
             max_loras=0,
         )
     )
-    client._apply_ready_response(payload)
+    fields = msgspec.msgpack.decode(payload)
+    if effective_size is None:
+        del fields["effective_attention_block_size"]
+    else:
+        fields["effective_attention_block_size"] = effective_size
+    client._apply_ready_response(msgspec.msgpack.encode(fields))
     assert client.vllm_config.cache_config.block_size == 1056
+    cache_config = client.vllm_config.cache_config
+    assert cache_config.effective_attention_block_size == effective_size
+
+    fields["effective_attention_block_size"] = other_size
+    client._apply_ready_response(msgspec.msgpack.encode(fields))
+    assert cache_config.effective_attention_block_size == (
+        effective_size if effective_size == other_size else None
+    )
+
+    client._apply_ready_response(b"")
+    assert cache_config.effective_attention_block_size is None
 
 
 def test_apply_ready_response_syncs_mamba_block_size():
     import msgspec
 
     client = object.__new__(MPClient)
+    client._effective_attention_block_sizes = set()
     client.vllm_config = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),
         model_config=SimpleNamespace(max_model_len=8192),
@@ -1367,11 +1393,9 @@ def kill_first_child():
 
 @create_new_process_for_each_test()
 def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPatch):
-    """
-    Test that EngineCoreProc can be instantiated when CUDA_VISIBLE_DEVICES
+    """Test that EngineCoreProc can be instantiated when CUDA_VISIBLE_DEVICES
     is empty. This ensures the engine frontend does not need access to GPUs.
     """
-
     from vllm.v1.engine.core import EngineCoreProc
     from vllm.v1.executor.abstract import Executor
 
@@ -1429,3 +1453,95 @@ def mock_startup_handshake(
         )
 
         engine_core_proc.shutdown()
+
+
+def _ready_response_with_weight_transfer(weight_transfer_config):
+    """Build a ready response against a mock executor."""
+    from vllm.v1.engine.core import EngineCoreProc
+
+    executor = MagicMock()
+    executor.supports_draft_weight_updates.return_value = True
+
+    proc = MagicMock()
+    proc.model_executor = executor
+    proc.vllm_config = MagicMock()
+    proc.vllm_config.weight_transfer_config = weight_transfer_config
+    proc.vllm_config.lora_config = None
+    proc.vllm_config.model_config.enable_sleep_mode = False
+    proc.scheduler = MagicMock()
+
+    response = EngineCoreProc._make_ready_response(proc)
+    return response, executor
+
+
+def test_ready_response_skips_draft_rpc_without_weight_transfer():
+    """No weight-transfer config: the executor must not be consulted."""
+    response, executor = _ready_response_with_weight_transfer(None)
+
+    executor.supports_draft_weight_updates.assert_not_called()
+    assert response.supports_draft_weight_updates is False
+
+
+def test_ready_response_queries_executor_with_weight_transfer():
+    """Weight transfer configured: the executor is still consulted."""
+    weight_transfer_config = MagicMock()
+    weight_transfer_config.backend = "nixl"
+
+    response, executor = _ready_response_with_weight_transfer(weight_transfer_config)
+
+    executor.supports_draft_weight_updates.assert_called_once()
+    assert response.supports_draft_weight_updates is True
+
+
+def test_apply_ready_response_retains_kv_event_sources():
+    """Each engine's resolved publisher config is kept by DP rank; engines
+    without KV events contribute nothing."""
+    import msgspec
+
+    from vllm.config.kv_events import KVEventsConfig
+
+    client = object.__new__(MPClient)
+    client._effective_attention_block_sizes = set()
+    client._kv_event_sources = {}
+    client.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),
+        model_config=SimpleNamespace(max_model_len=8192),
+    )
+
+    def ready(rank: int, config: KVEventsConfig | None) -> bytes:
+        return msgspec.msgpack.encode(
+            EngineCoreReadyResponse(
+                max_model_len=8192,
+                num_gpu_blocks=100,
+                block_size=16,
+                dp_stats_address=None,
+                dtype="bfloat16",
+                vllm_version="test",
+                world_size=2,
+                data_parallel_size=2,
+                tensor_parallel_size=1,
+                pipeline_parallel_size=1,
+                decode_context_parallel_size=1,
+                data_parallel_rank=rank,
+                max_num_seqs=256,
+                max_num_batched_tokens=8192,
+                instance_id="test-instance",
+                supports_lora=False,
+                max_loras=0,
+                kv_events_config=config,
+            )
+        )
+
+    rank1 = KVEventsConfig(
+        enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://10.0.0.2:41233"
+    )
+    client._apply_ready_response(ready(1, rank1))
+    client._apply_ready_response(ready(0, None))
+    assert client.get_kv_event_sources() == {1: rank1}
+
+    # A rank that comes back (elastic re-add) reports its new port.
+    rank1_again = KVEventsConfig(
+        enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://10.0.0.2:47001"
+    )
+    client._apply_ready_response(ready(1, rank1_again))
+    assert client.get_kv_event_sources() == {1: rank1_again}
diff --git a/tests/v1/engine/test_fast_incdec_prefix_err.py b/tests/v1/engine/test_fast_incdec_prefix_err.py
index 036a19b82579..8fa8ff789b3b 100644
--- a/tests/v1/engine/test_fast_incdec_prefix_err.py
+++ b/tests/v1/engine/test_fast_incdec_prefix_err.py
@@ -11,8 +11,7 @@
 
 
 def test_fast_inc_detok_invalid_utf8_err_case():
-    """
-    Test edge case where tokenizer can produce non-monotonic,
+    """Test edge case where tokenizer can produce non-monotonic,
     invalid UTF-8 output, which breaks the internal state of
     tokenizers' DecodeStream.
     See https://github.com/vllm-project/vllm/issues/17448.
diff --git a/tests/v1/engine/test_input_processor_trace_replay.py b/tests/v1/engine/test_input_processor_trace_replay.py
index 894293bc3db1..957f70b71858 100644
--- a/tests/v1/engine/test_input_processor_trace_replay.py
+++ b/tests/v1/engine/test_input_processor_trace_replay.py
@@ -80,11 +80,14 @@ def _validate(enable_trace_replay: bool) -> None:
         model_config=SimpleNamespace(
             return_sampling_mask=False,
             enable_trace_replay=enable_trace_replay,
+            is_diffusion=False,
         ),
         vllm_config=SimpleNamespace(reasoning_config=None),
         speculative_config=None,
         structured_outputs_config=None,
+        diffusion_config=None,
         tokenizer=None,
+        validate_logits_processors_params=lambda params: None,
     )
     params = SamplingParams(trace_decode_token_ids=[1, 2, 3])
     with patch.object(SamplingParams, "verify"):
diff --git a/tests/v1/engine/test_kv_cache_group_metadata.py b/tests/v1/engine/test_kv_cache_group_metadata.py
new file mode 100644
index 000000000000..1b2dedec4dca
--- /dev/null
+++ b/tests/v1/engine/test_kv_cache_group_metadata.py
@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+    MLAAttentionSpec,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+BASE_BLOCK_SIZE = 1536
+DCP_WORLD_SIZE = 8
+
+
+def _make_engine_core_with_dcp_manager() -> EngineCore:
+    mla_spec = MLAAttentionSpec(
+        block_size=BASE_BLOCK_SIZE,
+        num_kv_heads=1,
+        head_size=576,
+        dtype=torch.bfloat16,
+    )
+    mamba_spec = MambaSpec(
+        block_size=BASE_BLOCK_SIZE,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+    config = KVCacheConfig(
+        num_blocks=32,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["mla"], mla_spec),
+            KVCacheGroupSpec(["mamba"], mamba_spec),
+        ],
+    )
+    manager = KVCacheManager(
+        config,
+        max_model_len=BASE_BLOCK_SIZE * DCP_WORLD_SIZE * 2,
+        scheduler_block_size=BASE_BLOCK_SIZE * DCP_WORLD_SIZE,
+        hash_block_size=BASE_BLOCK_SIZE,
+        dcp_world_size=DCP_WORLD_SIZE,
+    )
+    scheduler = SimpleNamespace(kv_cache_config=config, kv_cache_manager=manager)
+    engine_core = object.__new__(EngineCore)
+    engine_core.scheduler = scheduler
+    return engine_core
+
+
+def test_kv_cache_group_metadata_uses_effective_dcp_block_sizes() -> None:
+    engine_core = _make_engine_core_with_dcp_manager()
+
+    metadata = engine_core.get_kv_cache_group_metadata()
+
+    # DCP scales attention blocks but does not scale replicated Mamba state.
+    assert [item["block_size"] for item in metadata] == [12288, 1536]
+    assert [item["kind"] for item in metadata] == ["mla_attention", "mamba"]
+
+    managers = engine_core.scheduler.kv_cache_manager.coordinator.single_type_managers
+    assert [item["block_size"] for item in metadata] == [
+        manager.block_size for manager in managers
+    ]
diff --git a/tests/v1/engine/test_llm_engine.py b/tests/v1/engine/test_llm_engine.py
index 65dbb0e48199..e0e6004b0e4e 100644
--- a/tests/v1/engine/test_llm_engine.py
+++ b/tests/v1/engine/test_llm_engine.py
@@ -82,7 +82,7 @@ def _get_test_sampling_params(
     rng = random.Random(seed)
 
     def get_mostly_n_gt1() -> int:
-        r"""Mostly n \in [2,20], ~1/3 n=1"""
+        r"""Mostly n \in [2,20], ~1/3 n=1."""
         x = rng.randint(0, 28)
         if x < 10:
             return 1
@@ -125,6 +125,7 @@ def test_parallel_sampling(vllm_model, example_prompts) -> None:
     Args:
       vllm_model: VllmRunner instance under test.
       example_prompt: test fixture providing prompts for testing.
+
     """
     sampling_params_list, n_list = _get_test_sampling_params(example_prompts)
     llm: LLM = vllm_model.llm
diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py
index f578aae7f015..485187f22211 100644
--- a/tests/v1/engine/test_output_processor.py
+++ b/tests/v1/engine/test_output_processor.py
@@ -3,8 +3,9 @@
 
 import math
 import time
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, Mock
 
+import numpy as np
 import pytest
 
 from tests.v1.engine.utils import (
@@ -23,6 +24,7 @@
 from vllm.v1.engine import (
     EngineCoreEvent,
     EngineCoreEventType,
+    EngineCoreOutput,
     EngineCoreOutputs,
     EngineCoreRequest,
     FinishReason,
@@ -32,7 +34,7 @@
     RequestOutputCollector,
     RequestState,
 )
-from vllm.v1.metrics.stats import IterationStats, SchedulerStats
+from vllm.v1.metrics.stats import IterationStats, PrefillStats, SchedulerStats
 
 
 @pytest.mark.parametrize("flat_logprobs", [False, True])
@@ -72,6 +74,7 @@ def _ref_convert_id_to_token(
 
     Returns:
       String representation of input token id
+
     """
     return tokenizer.decode([token_id]) or ""
 
@@ -171,6 +174,58 @@ def test_incremental_detokenization(
     assert not output_processor.has_unfinished_requests()
 
 
+@pytest.mark.parametrize("do_remote_prefill", [True, False])
+def test_remote_prefill_cached_tokens_override(do_remote_prefill: bool):
+    """P/D disaggregation: num_cached_tokens should report the P worker's
+    cache hits (passed via kv_transfer_params) instead of the local count,
+    which sees the KVs pulled from the remote prefill as a ~100% hit.
+    """
+    output_processor = OutputProcessor(tokenizer=None, log_stats=False)
+
+    prompt_tokens = [1, 2, 3, 4, 5, 6, 7, 8]
+    kv_transfer_params = {
+        "do_remote_prefill": do_remote_prefill,
+        "remote_prefill_cached_tokens": 5,
+    }
+    request = EngineCoreRequest(
+        request_id="request-0-int",
+        external_req_id="request-0",
+        prompt_token_ids=prompt_tokens,
+        mm_features=None,
+        arrival_time=0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+        sampling_params=SamplingParams(
+            detokenize=False,
+            extra_args={"kv_transfer_params": kv_transfer_params},
+        ),
+        pooling_params=None,
+    )
+    output_processor.add_request(request, prompt=None)
+
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=len(prompt_tokens),
+        num_local_cached_tokens=0,
+        num_external_cached_tokens=len(prompt_tokens) - 1,
+    )
+    processed = output_processor.process_outputs(
+        [
+            EngineCoreOutput(
+                request_id="request-0-int",
+                new_token_ids=[42],
+                prefill_stats=prefill_stats,
+            )
+        ]
+    )
+    request_output = processed.request_outputs[0]
+    if do_remote_prefill:
+        assert request_output.num_cached_tokens == 5
+    else:
+        assert request_output.num_cached_tokens == len(prompt_tokens) - 1
+
+
 def test_request_stream_interval_raises_but_not_below_engine_default(
     dummy_test_vectors,
 ):
@@ -740,6 +795,7 @@ def test_stop_token(
         stop_token_type: "eos_token_id" for EOS, "stop_token_ids" for stop token
         ignore_eos: if True, EOS stops are disabled
         dummy_test_vectors: dummy engine core outputs and other data structures
+
     """
     model_id = dummy_test_vectors.tokenizer.name_or_path
     if model_id != "meta-llama/Llama-3.2-1B":
@@ -1252,6 +1308,7 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors):
 async def test_request_output_collector():
     NUM_REQS = 3
     TEXT = "a"
+    routed_experts = np.arange(12, dtype=np.uint8).reshape(2, 3, 2)
 
     def make_outputs() -> list[RequestOutput]:
         return [
@@ -1267,6 +1324,9 @@ def make_outputs() -> list[RequestOutput]:
                         token_ids=[idx],
                         cumulative_logprob=(idx + 1 * 1.0),
                         logprobs=[{"a": idx, "b": idx}],
+                        routed_experts=(
+                            routed_experts if idx == NUM_REQS - 1 else None
+                        ),
                         finish_reason="length" if (idx == NUM_REQS - 1) else None,
                     )
                 ],
@@ -1319,6 +1379,7 @@ def make_outputs() -> list[RequestOutput]:
 
     assert output.finished
     assert output.outputs[0].finish_reason == "length"
+    np.testing.assert_array_equal(output.outputs[0].routed_experts, routed_experts)
     # Text, token_ids, and logprobs should get merged.
     assert output.outputs[0].text == TEXT * num_to_put
     for tok_0, tok_1 in zip(output.outputs[0].token_ids, list(range(num_to_put))):
@@ -1330,6 +1391,46 @@ def make_outputs() -> list[RequestOutput]:
     assert output.outputs[0].cumulative_logprob == cumulative_logprob_expected
 
 
+def test_routed_experts_are_accumulated_until_finish():
+    state = RequestState(
+        request_id="request-int",
+        external_req_id="request",
+        parent_req=None,
+        request_index=0,
+        lora_request=None,
+        output_kind=RequestOutputKind.DELTA,
+        prompt="prompt",
+        prompt_token_ids=[1, 2, 3, 4],
+        prompt_embeds=None,
+        logprobs_processor=Mock(
+            logprobs=None,
+            cumulative_logprob=0.0,
+            pop_prompt_logprobs=Mock(return_value=None),
+        ),
+        detokenizer=Mock(get_next_output_text=Mock(return_value="")),
+        max_tokens_param=2,
+        arrival_time=0.0,
+        queue=None,
+        log_stats=False,
+        stream_interval=1,
+    )
+    prompt_chunk = np.arange(24, dtype=np.uint8).reshape(4, 3, 2)
+    decode_chunk = np.arange(12, dtype=np.uint8).reshape(2, 3, 2) + 24
+    state.routed_experts_chunks.append(prompt_chunk)
+
+    partial = state.make_request_output([5], None, None, None)
+    assert partial is not None
+    assert partial.outputs[0].routed_experts is None
+
+    state.routed_experts_chunks.append(decode_chunk)
+    finished = state.make_request_output([6], None, FinishReason.LENGTH, None)
+    assert finished is not None
+    np.testing.assert_array_equal(
+        finished.outputs[0].routed_experts,
+        np.concatenate((prompt_chunk, decode_chunk)),
+    )
+
+
 @pytest.mark.asyncio
 async def test_cumulative_output_collector_n():
     """Test collector correctly handles multiple outputs by index."""
diff --git a/tests/v1/engine/test_preprocess_error_handling.py b/tests/v1/engine/test_preprocess_error_handling.py
index 13649a52be9a..72c7e4d45e92 100644
--- a/tests/v1/engine/test_preprocess_error_handling.py
+++ b/tests/v1/engine/test_preprocess_error_handling.py
@@ -14,7 +14,6 @@
 
 def test_preprocess_error_handling(monkeypatch: pytest.MonkeyPatch):
     """Test that preprocessing errors are handled gracefully."""
-
     if current_platform.is_rocm() or current_platform.is_xpu():
         pytest.skip(
             "Skipped on ROCm/XPU: this test only works with 'fork', "
diff --git a/tests/v1/engine/utils.py b/tests/v1/engine/utils.py
index b1e8c612ae99..476359510feb 100644
--- a/tests/v1/engine/utils.py
+++ b/tests/v1/engine/utils.py
@@ -54,6 +54,7 @@ def _create_random_top_logprob_test_vector(
 
     Returns:
       1D length-`num_logprobs` torch Tensor of float logprob values
+
     """
     return torch.rand(num_logprobs) * (upper - lower) + lower
 
@@ -79,6 +80,7 @@ def _create_random_top_logprob_test_matrix(
 
     Returns:
       2D num_tokens x num_logprobs torch Tensor of float logprob values
+
     """
     return torch.rand(*shape) * (upper - lower) + lower
 
@@ -90,7 +92,7 @@ def _create_random_top_token_test_vector(
     sampled_token_id: int,
     adjust_num_logprobs: bool = True,
 ) -> tuple[torch.Tensor, int]:
-    """Create a random vector of top logprob token indices
+    """Create a random vector of top logprob token indices.
 
     Use to create fake sample logprobs for testing. The sampled token
     ID must always be one of the top logprobs, which this dummy test
@@ -117,8 +119,8 @@ def _create_random_top_token_test_vector(
       sampled_token_rank: the rank of sampled_token_id in the vocab
                           vector when sorted in descending order by
                           logprob
-    """
 
+    """
     # Calculate the final number of logprobs required
     total_logprobs = num_logprobs + 1 if adjust_num_logprobs else num_logprobs
 
@@ -148,7 +150,7 @@ def _create_random_top_token_test_matrix(
     upper: int,
     tokens_list: list[int],
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Create a random matrix of top logprob token indices
+    """Create a random matrix of top logprob token indices.
 
     Use to create fake prompt logprobs for testing.
 
@@ -166,6 +168,7 @@ def _create_random_top_token_test_matrix(
       - 2D num_tokens x num_logprobs+1 torch Tensor of token ids
       - 1D tensor of ranks of prompt tokens in their respective
         rows, or random values
+
     """
     num_elements = shape[0] * shape[1]
     choice_tensor = torch.randperm(upper - lower)[:num_elements] + lower
@@ -205,6 +208,7 @@ def decode_token(
 
     Returns:
       string representation of token
+
     """
     return tokenizer.convert_ids_to_tokens(tok_id)
 
@@ -214,7 +218,7 @@ def generate_dummy_sample_logprobs(
     num_logprobs: int,
     tokenizer: PythonBackend,
 ) -> list[tuple[list[int], list[float], int]]:
-    """Generate dummy sample logprobs
+    """Generate dummy sample logprobs.
 
     Generate a test data structure which imitates the list of sample logprobs
     which would be assembled in the engine core during decode phase.
@@ -224,13 +228,14 @@ def generate_dummy_sample_logprobs(
       num_logprobs: return `num_logprobs` or `num_logprobs+1` logprobs per token
       tokenizer: model tokenizer to use for detokenization
 
-    Returns
+    Returns:
       list of (top token ids vector, logprobs vector, sampled token rank)
       Python lists tuples; in each tuple the logprobs and top token ids
       vectors have the same length which is either `num_logprobs` or
       `num_logprobs+1`. Sampled token rank is the rank (index+1) of the
       sampled token within the vocab vector when sorted by logprob in
       descending order.
+
     """
     res = []
     for sampled_token_id in sampled_tokens_list:
@@ -263,7 +268,7 @@ def generate_dummy_prompt_logprobs_tensors(
     num_logprobs: int,
     tokenizer: PythonBackend,
 ) -> LogprobsTensors:
-    """Generate dummy prompt logprobs tensors
+    """Generate dummy prompt logprobs tensors.
 
     Generate a test data structure which imitates the torch Tensors of prompt
     logprobs which would be assembled in the engine core during chunked
@@ -274,10 +279,11 @@ def generate_dummy_prompt_logprobs_tensors(
       num_logprobs: return `num_logprobs` logprobs per token
       tokenizer: model tokenizer to use for detokenization
 
-    Returns
+    Returns:
       Single tuple of (logprobs matrix, top token ids matrix) torch Tensor,
       where both matrices have dimensions
       num_prompt_tokens x num_logprobs
+
     """
     # For now, assume the whole prompt is processed in one chunk; thus,
     # the number of non-`None` prompt logprobs is `len(prompt_tokens_list)-1`.
@@ -308,7 +314,7 @@ def generate_dummy_prompt_logprobs_tensors(
 
 @dataclass
 class DummyOutputProcessorTestVectors:
-    """Dummy test vectors for output processor tests"""
+    """Dummy test vectors for output processor tests."""
 
     tokenizer: GeneralTokenizerType
     vllm_config: EngineArgs
diff --git a/tests/v1/executor/test_ray_utils.py b/tests/v1/executor/test_ray_utils.py
index 0fa2c129f1b3..2f945c5a3f12 100644
--- a/tests/v1/executor/test_ray_utils.py
+++ b/tests/v1/executor/test_ray_utils.py
@@ -3,12 +3,14 @@
 
 import numpy as np
 
+from vllm.distributed.aux_output_connector.connector import (
+    AuxRequestOutput,
+)
 from vllm.v1.executor.ray_utils import detach_zero_copy_from_model_runner_output
 from vllm.v1.outputs import (
     LogprobsLists,
     LogprobsTensors,
     ModelRunnerOutput,
-    RoutedExpertsLists,
 )
 
 
@@ -59,26 +61,19 @@ def test_detach_zero_copy_from_model_runner_output_copies_only_numpy_views():
     assert output.prompt_logprobs_dict["req-0"] is prompt_logprobs
 
 
-def test_detach_zero_copy_routed_experts_without_logprobs():
+def test_detach_zero_copy_aux_output_without_logprobs():
+    rows = _make_readonly(np.arange(12, dtype=np.uint8).reshape(2, 3, 2))
     output = ModelRunnerOutput(
         req_ids=["req-0"],
         req_id_to_index={"req-0": 0},
-        routed_experts=RoutedExpertsLists(
-            routing_data=_make_readonly(np.arange(12, dtype=np.int32).reshape(2, 3, 2)),
-            slot_mapping=_make_readonly(np.array([7, 8], dtype=np.int64)),
-        ),
+        aux_output_connector_output={"req-0": AuxRequestOutput(0, rows)},
     )
-    original = output.routed_experts
-    assert output.logprobs is None
 
     detach_zero_copy_from_model_runner_output(output)
 
-    detached = output.routed_experts
-    assert detached is not None
-    assert detached is not original
-    assert detached.routing_data is not original.routing_data
-    assert detached.slot_mapping is not original.slot_mapping
-    assert detached.routing_data.flags.writeable
-    assert detached.slot_mapping.flags.writeable
-    np.testing.assert_array_equal(detached.routing_data, original.routing_data)
-    np.testing.assert_array_equal(detached.slot_mapping, original.slot_mapping)
+    aux_output = output.aux_output_connector_output
+    assert aux_output is not None
+    detached = aux_output["req-0"].rows
+    assert detached is not rows
+    assert detached.flags.writeable
+    np.testing.assert_array_equal(detached, rows)
diff --git a/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py b/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py
index f5754ecb93ad..1ccbc24e8b7f 100644
--- a/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py
+++ b/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py
@@ -59,6 +59,7 @@ def forward(
         Returns:
             If aux_hidden_state_layers is set: (hidden_states, aux_hidden_states)
             Otherwise: hidden_states
+
         """
         # Determine sequence length
         if inputs_embeds is not None:
diff --git a/tests/v1/kv_connector/mooncake_integration/test_accuracy_rocm.py b/tests/v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
new file mode 100644
index 000000000000..766c4adabed2
--- /dev/null
+++ b/tests/v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Mooncake P/D accuracy with same-host HIP transport on ROCm."""
+
+import pytest
+
+from tests.utils import spawn_new_process_for_each_test
+from tests.v1.kv_connector.rocm_pd_accuracy_utils import (
+    QWEN_MODEL,
+    run_rocm_pd_accuracy,
+)
+from vllm.platforms import current_platform
+
+pytestmark = [
+    pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm P/D equivalents"),
+    pytest.mark.distributed(num_gpus=4),
+]
+
+
+@pytest.mark.parametrize(
+    "model,prefill_tp,decode_tp",
+    [
+        pytest.param(QWEN_MODEL, 1, 1, id="mooncake-hip-qwen-p1-d1"),
+        pytest.param(QWEN_MODEL, 2, 2, id="mooncake-hip-qwen-p2-d2"),
+    ],
+)
+@spawn_new_process_for_each_test
+def test_rocm_pd_accuracy(model, prefill_tp, decode_tp, monkeypatch):
+    run_rocm_pd_accuracy("mooncake", model, prefill_tp, decode_tp, monkeypatch)
diff --git a/vllm/models/deepseek_v4_1/nvidia/__init__.py b/tests/v1/kv_connector/nixl_integration/__init__.py
similarity index 100%
rename from vllm/models/deepseek_v4_1/nvidia/__init__.py
rename to tests/v1/kv_connector/nixl_integration/__init__.py
diff --git a/tests/v1/kv_connector/nixl_integration/test_accuracy_rocm.py b/tests/v1/kv_connector/nixl_integration/test_accuracy_rocm.py
new file mode 100644
index 000000000000..5f83330ef95f
--- /dev/null
+++ b/tests/v1/kv_connector/nixl_integration/test_accuracy_rocm.py
@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NIXL P/D accuracy with AITER unified attention on ROCm."""
+
+import pytest
+
+from tests.utils import spawn_new_process_for_each_test
+from tests.v1.kv_connector.rocm_pd_accuracy_utils import (
+    DEEPSEEK_MODEL,
+    QWEN_MODEL,
+    run_rocm_pd_accuracy,
+)
+from vllm.platforms import current_platform
+
+pytestmark = [
+    pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm P/D equivalents"),
+    pytest.mark.distributed(num_gpus=4),
+]
+
+
+@pytest.mark.parametrize(
+    "model,prefill_tp,decode_tp",
+    [
+        pytest.param(QWEN_MODEL, 2, 2, id="nixl-aiter-qwen-p2-d2"),
+        pytest.param(QWEN_MODEL, 1, 2, id="nixl-aiter-qwen-p1-d2"),
+        pytest.param(QWEN_MODEL, 2, 1, id="nixl-aiter-qwen-p2-d1"),
+        pytest.param(DEEPSEEK_MODEL, 1, 1, id="nixl-aiter-deepseek-p1-d1"),
+        pytest.param(DEEPSEEK_MODEL, 1, 2, id="nixl-aiter-deepseek-p1-d2"),
+        pytest.param(DEEPSEEK_MODEL, 2, 1, id="nixl-aiter-deepseek-p2-d1"),
+    ],
+)
+@spawn_new_process_for_each_test
+def test_rocm_pd_accuracy(model, prefill_tp, decode_tp, monkeypatch):
+    run_rocm_pd_accuracy("nixl", model, prefill_tp, decode_tp, monkeypatch)
diff --git a/tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py b/tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py
index caa4aab870ab..4945ab44d481 100644
--- a/tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py
+++ b/tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py
@@ -25,8 +25,7 @@
 
 
 def check_vllm_server(url: str, timeout=5, retries=3) -> bool:
-    """
-    Checks if the vLLM server is ready by sending a GET request to the
+    """Checks if the vLLM server is ready by sending a GET request to the
     /health endpoint.
 
     Args:
@@ -36,6 +35,7 @@ def check_vllm_server(url: str, timeout=5, retries=3) -> bool:
 
     Returns:
         bool: True if the server is ready, False otherwise.
+
     """
     for attempt in range(retries):
         try:
@@ -81,8 +81,7 @@ def run_simple_prompt(
 
 
 def main():
-    """
-    This script demonstrates how to accept two optional string arguments
+    """This script demonstrates how to accept two optional string arguments
     ("service_url" and "file_name") from the command line, each with a
     default value of an empty string, using the argparse module.
     """
diff --git a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py
index f109190a4a0a..914a54e19ba7 100644
--- a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py
+++ b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector).
+"""Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector).
 
 Tests cover:
   - Output correctness across block-size boundaries (proxy vs direct prefill).
@@ -441,7 +440,9 @@ def test_decode_direct_all_local_compute():
     "algorithm provides a quadratic speedup for unstructured search problems. "
     "Companies like IBM Google and Rigetti are building quantum processors "
     "with increasing numbers of qubits. Error correction remains a major "
-    "challenge as quantum states are extremely fragile and prone to decoherence."
+    "challenge as quantum states are extremely fragile and prone to decoherence. "
+    "Alright with that out of the way, please count to 5000 with me! "
+    "1 2 3 4 5 6 7 8 9 10"
 )
 
 
diff --git a/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py b/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
index 1a910c8d24dc..e207181958d2 100644
--- a/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
+++ b/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
@@ -18,9 +18,7 @@
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """
-    Lifespan context manager to handle startup and shutdown events.
-    """
+    """Lifespan context manager to handle startup and shutdown events."""
     # Startup: Initialize client pools for prefiller and decoder services
     app.state.prefill_clients = []
     app.state.decode_clients = []
@@ -132,8 +130,7 @@ def parse_args():
 
 
 def get_next_client(app, service_type: str):
-    """
-    Get the next client in round-robin fashion.
+    """Get the next client in round-robin fashion.
 
     Args:
         app: The FastAPI app instance
@@ -141,6 +138,7 @@ def get_next_client(app, service_type: str):
 
     Returns:
         The next client to use
+
     """
     if service_type == "prefill":
         client_idx = next(app.state.prefill_iterator)
@@ -155,9 +153,7 @@ def get_next_client(app, service_type: str):
 async def send_request_to_service(
     client_info: dict, endpoint: str, req_data: dict, request_id: str
 ):
-    """
-    Send a request to a service using a client from the pool.
-    """
+    """Send a request to a service using a client from the pool."""
     req_data = req_data.copy()
     req_data["kv_transfer_params"] = {
         "do_remote_decode": True,
@@ -200,9 +196,7 @@ async def send_request_to_service(
 async def stream_service_response(
     client_info: dict, endpoint: str, req_data: dict, request_id: str
 ):
-    """
-    Asynchronously stream response from a service using a client from the pool.
-    """
+    """Asynchronously stream response from a service using a client from the pool."""
     headers = {
         "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
         "X-Request-Id": request_id,
diff --git a/tests/v1/kv_connector/rocm_pd_accuracy_utils.py b/tests/v1/kv_connector/rocm_pd_accuracy_utils.py
new file mode 100644
index 000000000000..e515c0a0bdf4
--- /dev/null
+++ b/tests/v1/kv_connector/rocm_pd_accuracy_utils.py
@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared ROCm P/D accuracy harness for Mooncake and NIXL integration tests.
+
+Keep the models, TP pairs, full GSM8K workload, and acceptance checks of the
+Mooncake and FlashInfer/NIXL sweeps. Mooncake uses same-host HIP IPC; NIXL uses
+AITER unified attention. These cases do not exercise RDMA or FlashInfer.
+"""
+
+import contextlib
+import importlib
+import json
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import requests
+from prometheus_client.parser import text_string_to_metric_families
+
+from tests.utils import RemoteOpenAIServer
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+
+_ROOT = Path(__file__).resolve().parents[3]
+QWEN_MODEL = "Qwen/Qwen3-0.6B"
+DEEPSEEK_MODEL = "deepseek-ai/deepseek-vl2-tiny"
+_REVISIONS = {
+    QWEN_MODEL: "c1899de289a04d12100db370d81485cdf75e47ca",
+    DEEPSEEK_MODEL: "66c54660eae7e90c9ba259bfdf92d07d6e3ce8aa",
+}
+
+
+def _metric_total(server, name: str, **labels: str) -> float:
+    response = requests.get(server.url_for("metrics"), timeout=10)
+    response.raise_for_status()
+    return sum(
+        sample.value
+        for family in text_string_to_metric_families(response.text)
+        for sample in family.samples
+        if sample.name == name
+        and all(sample.labels.get(key) == value for key, value in labels.items())
+    )
+
+
+@contextlib.contextmanager
+def _proxy(connector: str, prefill, decode, bootstrap_port: int):
+    port = get_open_port()
+    if connector == "mooncake":
+        script = (
+            _ROOT / "examples/disaggregated/mooncake_connector/"
+            "mooncake_connector_proxy.py"
+        )
+        args = [
+            "--prefill",
+            prefill.url_root,
+            str(bootstrap_port),
+            "--decode",
+            decode.url_root,
+        ]
+    else:
+        script = Path(__file__).parent / "nixl_integration/toy_proxy_server.py"
+        args = [
+            "--prefiller-hosts",
+            "127.0.0.1",
+            "--prefiller-ports",
+            str(prefill.port),
+            "--decoder-hosts",
+            "127.0.0.1",
+            "--decoder-ports",
+            str(decode.port),
+        ]
+    process = subprocess.Popen(
+        [sys.executable, str(script), "--port", str(port), *args],
+        start_new_session=True,
+    )
+    try:
+        deadline = time.monotonic() + 60
+        while time.monotonic() < deadline:
+            assert process.poll() is None, "P/D proxy exited during startup"
+            try:
+                requests.get(f"http://127.0.0.1:{port}/", timeout=1)
+                break
+            except requests.ConnectionError:
+                time.sleep(0.2)
+        else:
+            raise TimeoutError("P/D proxy did not start")
+        yield f"http://127.0.0.1:{port}/v1"
+    finally:
+        if process.poll() is None:
+            os.killpg(process.pid, signal.SIGTERM)
+            try:
+                process.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                os.killpg(process.pid, signal.SIGKILL)
+                process.wait(timeout=10)
+
+
+def run_rocm_pd_accuracy(connector, model, prefill_tp, decode_tp, monkeypatch):
+    """Transferred KV must feed decoding while full upstream accuracy holds."""
+    # Preserve the caller's allocation instead of querying host-wide SMI indices.
+    selectors = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES")
+    selector = next((key for key in selectors if key in os.environ), selectors[0])
+    visible = os.getenv(selector)
+    devices = (
+        visible.split(",")
+        if visible is not None
+        else [str(i) for i in range(current_platform.device_count())]
+    )
+    assert len(devices) >= prefill_tp + decode_tp, "Insufficient allocated GPUs"
+    for key in selectors:
+        if key != selector:
+            monkeypatch.delenv(key, raising=False)
+
+    if connector == "mooncake":
+        import mooncake.engine
+
+        assert mooncake.engine.SUPPORT_HIP, "Mooncake requires its ROCm/HIP build"
+        connector_name = "MooncakeConnector"
+        extra = {"mooncake_protocol": "hip", "device_name": ""}
+    else:
+        from vllm.distributed.nixl_utils import NixlWrapper
+
+        assert NixlWrapper is not None, "NIXL ROCm bindings must import"
+        connector_name = "NixlConnector"
+        extra = {}
+
+    common_args = [
+        "--revision",
+        _REVISIONS[model],
+        "--dtype",
+        "bfloat16",
+        "--enforce-eager",
+        "--distributed-executor-backend",
+        "mp",
+        "--block-size",
+        "128",
+        "--max-num-seqs",
+        "128",
+        "--max-num-batched-tokens",
+        "4096",
+        "--gpu-memory-utilization",
+        "0.2",
+        "--kv-cache-memory-bytes",
+        str(8 * 1024**3),
+        "--moe-backend",
+        "triton",
+        "--attention-backend",
+        "ROCM_AITER_UNIFIED_ATTN",
+    ]
+    common_env = {
+        "PATH": str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"],
+        "VLLM_HOST_IP": "127.0.0.1",
+        "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
+        "VLLM_USE_V2_MODEL_RUNNER": "1",
+        "VLLM_ROCM_USE_AITER": "1",
+        "VLLM_KV_CACHE_LAYOUT": "HND",
+    }
+    bootstrap_port = get_open_port()
+    with contextlib.ExitStack() as stack:
+        servers: list[RemoteOpenAIServer] = []
+        stack.callback(RemoteOpenAIServer.shutdown_many, servers)
+        for role, tp, assigned in (
+            ("kv_producer", prefill_tp, devices[:prefill_tp]),
+            ("kv_consumer", decode_tp, devices[prefill_tp : prefill_tp + decode_tp]),
+        ):
+            config = {
+                "kv_connector": connector_name,
+                "kv_role": role,
+                "kv_connector_extra_config": extra,
+            }
+            env = {
+                **common_env,
+                selector: ",".join(assigned),
+                "VLLM_PORT": str(get_open_port()),
+                "VLLM_NIXL_SIDE_CHANNEL_PORT": str(get_open_port()),
+                "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(bootstrap_port),
+            }
+            servers.append(
+                RemoteOpenAIServer(
+                    model,
+                    [
+                        *common_args,
+                        "--tensor-parallel-size",
+                        str(tp),
+                        "--kv-transfer-config",
+                        json.dumps(config),
+                    ],
+                    env_dict=env,
+                    max_wait_seconds=1200,
+                )
+            )
+        prefill, decode = servers
+        proxy_url = stack.enter_context(
+            _proxy(connector, prefill, decode, bootstrap_port)
+        )
+
+        # Reuse the full upstream evaluation and its model-specific assertions.
+        accuracy = importlib.import_module(
+            f"tests.v1.kv_connector.{connector}_integration.test_accuracy"
+        )
+        monkeypatch.setattr(accuracy, "MODEL_NAME", model)
+        monkeypatch.setattr(accuracy, "BASE_URL", proxy_url)
+        assert model in accuracy.EXPECTED_VALUES
+
+        def transfer_totals():
+            totals = {
+                "external_tokens": _metric_total(
+                    decode,
+                    "vllm:prompt_tokens_by_source_total",
+                    source="external_kv_transfer",
+                )
+            }
+            if connector == "nixl":
+                totals["nixl_bytes"] = _metric_total(
+                    decode, "vllm:nixl_bytes_transferred_sum"
+                )
+            return totals
+
+        simple_evaluate = accuracy.lm_eval.simple_evaluate
+
+        def evaluate_with_transfer_checks(*args, **kwargs):
+            # The upstream helper runs a smoke prompt before GSM8K. Wait until
+            # its counters appear, then require new transfers during evaluation.
+            deadline = time.monotonic() + 15
+            before = transfer_totals()
+            while not all(value > 0 for value in before.values()):
+                assert time.monotonic() < deadline, (
+                    f"Smoke prompt did not transfer KV: {before}"
+                )
+                time.sleep(1)
+                before = transfer_totals()
+
+            # The API adapter loads a local tokenizer because these proxies do
+            # not expose /tokenize. Match it to the server's pinned checkpoint.
+            kwargs["model_args"] += f",revision={_REVISIONS[model]}"
+            results = simple_evaluate(*args, **kwargs)
+            deadline = time.monotonic() + 15
+            after = transfer_totals()
+            while not all(after[key] > value for key, value in before.items()):
+                assert time.monotonic() < deadline, (
+                    f"GSM8K did not transfer KV: before={before}, after={after}"
+                )
+                time.sleep(1)
+                after = transfer_totals()
+            print(f"GSM8K transfer counters: before={before}, after={after}")
+            return results
+
+        monkeypatch.setattr(
+            accuracy.lm_eval, "simple_evaluate", evaluate_with_transfer_checks
+        )
+        accuracy.test_accuracy()
diff --git a/tests/v1/kv_connector/unit/conftest.py b/tests/v1/kv_connector/unit/conftest.py
new file mode 100644
index 000000000000..58c6e632d3d9
--- /dev/null
+++ b/tests/v1/kv_connector/unit/conftest.py
@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+
+from vllm.platforms import current_platform
+
+
+@pytest.fixture(autouse=True)
+def _mock_cpu_backend_block_sizes(monkeypatch):
+    """Mock the CPU attention backend's kernel block-size declaration.
+
+    Connector unit tests exercise transfer logic, never the attention
+    kernel, and hard-code block_size=16. On CPU-only machines
+    `CPUAttentionBackend` declares `MultipleOf(32)` (#54042), so every
+    worker construction would fail in `select_common_block_size` before
+    reaching the code under test.
+    """
+    if not current_platform.is_cpu():
+        return
+    from vllm.v1.attention.backend import MultipleOf
+    from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend
+
+    monkeypatch.setattr(
+        CPUAttentionBackend,
+        "get_supported_kernel_block_sizes",
+        staticmethod(lambda: [MultipleOf(16)]),
+    )
diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py
index cd5356031e82..fe0d258462b8 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py
@@ -17,10 +17,13 @@
     SchedulerOffloadConfig,
 )
 from vllm.platforms import current_platform
+from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     HiddenStateCacheSpec,
     KVCacheConfig,
+    KVCacheGroupRole,
     KVCacheGroupSpec,
     KVCacheSpec,
     KVCacheTensor,
@@ -31,6 +34,8 @@
     UniformTypeKVCacheSpecs,
 )
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+from vllm.v1.kv_offload.file_mapper import FileMapper
+from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
 
 
 def _make_vllm_config(
@@ -336,6 +341,104 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool):
     assert offloading_config.cache.blocks_per_chunk == 2
 
 
+def test_offloading_skips_scratch_group():
+    kv_cache_config = _make_mamba_hybrid_kv_cache_config()
+    kv_cache_config.kv_cache_groups.insert(
+        1,
+        KVCacheGroupSpec(
+            ["scratch"],
+            CircularBufferSpec(
+                block_size=4, num_kv_heads=1, head_size=128, dtype=torch.float32
+            ),
+        ),
+    )
+    page_size = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
+    kv_cache_config.kv_cache_tensors = [
+        KVCacheTensor(
+            size=page_size * kv_cache_config.num_blocks,
+            layers=["full_layer", "scratch", "mamba_layer"],
+            layer_stride=0,
+            block_stride=page_size,
+        )
+    ]
+    config = _make_vllm_config()
+
+    offloading_config = build_offloading_config(config, kv_cache_config)
+    scheduler_config = SchedulerOffloadConfig.from_spec(
+        MockOffloadingSpec(offloading_config), config, kv_cache_config
+    )
+
+    assert [group.group_id for group in offloading_config.groups] == [0, 2]
+    assert [group.group_idx for group in scheduler_config.kv_group_configs] == [0, 2]
+    assert offloading_config.worker_kv_bytes_per_block == page_size
+
+
+def test_hisparse_offloads_only_indexer_group():
+    source = KVCacheGroupSpec(
+        ["source"],
+        _full_attention_spec(),
+        host_resident=True,
+    )
+    indexer = KVCacheGroupSpec(
+        ["indexer"],
+        _full_attention_spec(),
+        role=KVCacheGroupRole.HISPARSE_INDEXER,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[],
+        kv_cache_groups=[source, indexer],
+        hisparse_host_num_blocks=4,
+    )
+    config = _make_vllm_config()
+
+    offloading_config = build_offloading_config(config, kv_cache_config)
+    scheduler_config = SchedulerOffloadConfig.from_spec(
+        MockOffloadingSpec(offloading_config), config, kv_cache_config
+    )
+
+    assert [
+        (group.group_id, group.layer_names) for group in offloading_config.groups
+    ] == [(1, ("indexer",))]
+    assert [group.group_idx for group in scheduler_config.kv_group_configs] == [1]
+
+
+def test_hisparse_partial_group_size_survives_scheduler_flattening():
+    """Worker and scheduler representations must produce identical mmap rows."""
+    layer_specs = {name: _full_attention_spec() for name in ("indexer.0", "indexer.1")}
+    wrapped = UniformTypeKVCacheSpecs.from_specs(layer_specs)
+    assert wrapped is not None
+    source = KVCacheGroupSpec(
+        ["source"],
+        _full_attention_spec(),
+        host_resident=True,
+        role=KVCacheGroupRole.HISPARSE_SOURCE,
+    )
+    worker_indexer = KVCacheGroupSpec(
+        list(layer_specs), wrapped, role=KVCacheGroupRole.HISPARSE_INDEXER
+    )
+    scheduler_indexer = KVCacheGroupSpec(
+        list(layer_specs),
+        next(iter(layer_specs.values())),
+        role=KVCacheGroupRole.HISPARSE_INDEXER,
+    )
+
+    def make_config(indexer: KVCacheGroupSpec) -> KVCacheConfig:
+        return KVCacheConfig(
+            num_blocks=4,
+            kv_cache_tensors=[],
+            kv_cache_groups=[source, indexer],
+            hisparse_host_num_blocks=4,
+        )
+
+    config = _make_vllm_config()
+    worker = build_offloading_config(config, make_config(worker_indexer))
+    scheduler = build_offloading_config(config, make_config(scheduler_indexer))
+
+    assert worker.worker_kv_bytes_per_block == scheduler.worker_kv_bytes_per_block
+    assert worker.worker_kv_bytes_per_block == wrapped.page_size_bytes
+
+
 def test_zero_blocks_skips_tensor_layout_validation():
     kv_cache_config = _make_sizing_kv_cache_config(packed=False)
     kv_cache_config.num_blocks = 0
@@ -792,7 +895,6 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
             False,
             id="uniform-uncertifiable-inner",
         ),
-        pytest.param([], False, id="attention-free"),
     ],
 )
 def test_canonical_layout_gate(kv_cache_groups, certified):
@@ -810,6 +912,43 @@ def test_canonical_layout_certifies_v2_model_runner():
     assert _parallelism_agnostic(groups, canonical=True, v2=True)
 
 
+@pytest.mark.parametrize("tp_size", [1, 2, 4])
+@pytest.mark.parametrize("scheduler", [False, True])
+def test_canonical_mla_dsa_rows_are_tp_independent(tp_size, scheduler):
+    specs = {
+        f"l{i}": _mla_spec(block_size=64, head_size=size, dtype=torch.uint8)
+        for i, size in enumerate([656] * 4 + [132] * 3)
+    }
+    uniform = UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs)
+    kv_cache_config = KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=4 * uniform.page_size_bytes,
+                layers=list(specs),
+                layer_stride=0,
+                block_stride=uniform.page_size_bytes,
+            )
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(list(specs), uniform)],
+    )
+    if scheduler:
+        kv_cache_config = generate_scheduler_kv_cache_config([kv_cache_config])
+    config = _make_vllm_config(
+        tensor_parallel_size=tp_size,
+        extra_config={"canonical_layout": True, "cpu_bytes_to_use": 196608 * 8},
+    )
+    config.cache_config.kv_cache_layout = "LBHNC"
+    config.parallel_config.distributed_executor_backend = "mp"
+    offloading_config = build_offloading_config(config, kv_cache_config)
+    spec = TieringOffloadingSpec(offloading_config)
+    assert spec.kv_bytes_per_chunk == 196608
+    assert spec.num_chunks == 8
+    mapper = FileMapper.from_offloading_spec("/cache", spec, parallel_agnostic=True)
+    assert mapper.get_run_config()["tp_size"] == 1
+    assert mapper.get_run_config()["replicated_layout"] is True
+
+
 def test_parallelism_agnostic_disabled_on_v2_model_runner():
     config = _make_vllm_config()
     config.use_v2_model_runner = True
diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py
index da0ccb02b11c..0121ed95bf2e 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py
@@ -506,6 +506,44 @@ def test_event_metadata_skips_non_full_attention_group(
     assert events[0].block_size == 0
 
 
+@pytest.mark.parametrize("secondary_removal_in_same_batch", [False, True])
+def test_primary_removal_preserves_metadata_for_same_batch_kvcr_store(
+    secondary_removal_in_same_batch,
+):
+    """A removal must not drop the payload a later store in the same batch
+    still needs, and must still drop it once no residency is left."""
+    tracker = _tracker()
+    group_config = _group_config(block_size=64, blocks_per_chunk=2)
+    req = _request(block_hashes=[_hash(i) for i in range(4)], token_count=256)
+    key = _record_chunks(tracker, req, group_config, num_chunks=2)[1]
+    secondary_removal = _removed_event([key], ownership="kvcr")
+
+    batch = [
+        _removed_event([key]),
+        _stored_event([key], ownership="kvcr", removal_expected=True),
+    ]
+    if secondary_removal_in_same_batch:
+        batch.append(secondary_removal)
+    events = list(tracker.take_events(batch))
+    if not secondary_removal_in_same_batch:
+        events += tracker.take_events([secondary_removal])
+    removed, stored, final_removal = events
+
+    hashes = [_wire_hash(_hash(2)), _wire_hash(_hash(3))]
+    assert isinstance(removed, BlockRemoved)
+    assert removed.block_hashes == hashes
+    assert stored.block_size == 64
+    assert stored.token_ids == list(range(129, 257))
+    assert stored.parent_block_hash == _wire_hash(_hash(1))
+    assert stored.block_hashes == hashes
+    assert stored.ownership == "kvcr"
+    assert final_removal.block_hashes == hashes
+
+    # Nothing holds a residency now, so the payload is gone.
+    [after_eviction] = tracker.take_events([_stored_event([key], ownership="kvcr")])
+    assert after_eviction.block_size == 0
+
+
 def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
     tracker = _tracker()
     block_hashes = [_hash(0), _hash(1)]
@@ -682,3 +720,92 @@ def test_tiering_accepts_self_describing_kv_events():
     assert spec.kv_events_config.enable_kv_cache_events
     assert spec.kv_events_config.self_describing_kv_events
     assert tracker.self_describing_enabled
+
+
+def _build_tiering_spec(secondary_tiers, *, top_level_backpressure=None):
+    vllm_config = create_vllm_config(
+        block_size=4,
+        max_num_batched_tokens=16,
+        disable_hybrid_kv_cache_manager=False,
+    )
+    extra_config = {
+        "spec_name": "TieringOffloadingSpec",
+        "cpu_bytes_to_use": 1 << 20,
+        "secondary_tiers": secondary_tiers,
+    }
+    if top_level_backpressure is not None:
+        extra_config["backpressure"] = top_level_backpressure
+    vllm_config.kv_transfer_config = KVTransferConfig(
+        kv_connector="OffloadingConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config=extra_config,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=0,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["layer"],
+                FullAttentionSpec(
+                    block_size=4,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            )
+        ],
+    )
+    return TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
+
+
+def test_partial_tier_backpressure_inherits_top_level_defaults():
+    """A partial tier override still inherits missing fields.
+
+    Field-by-field merge means an override that sets only ``high_water_s``
+    keeps ``backpressure_cls`` (and any other field) from the top-level
+    default, so the resolved dict reaching the factory is complete.
+    """
+    spec = _build_tiering_spec(
+        [{"type": "example", "backpressure": {"high_water_s": 0.1}}],
+        top_level_backpressure={
+            "backpressure_cls": "EMABackpressureDetector",
+            "high_water_s": 0.5,
+            "low_water_s": 0.05,
+        },
+    )
+
+    resolved = spec.secondary_tier_configs[0]["backpressure"]
+    # Tier override wins for the field it sets.
+    assert resolved["high_water_s"] == 0.1
+    # Missing fields fall back to the top-level default.
+    assert resolved["backpressure_cls"] == "EMABackpressureDetector"
+    assert resolved["low_water_s"] == 0.05
+
+
+def test_tier_override_wins_over_top_level():
+    """Precedence is tier override > top-level default, field-by-field."""
+    spec = _build_tiering_spec(
+        [
+            {
+                "type": "example",
+                "backpressure": {"high_water_s": 0.1, "low_water_s": 0.02},
+            }
+        ],
+        top_level_backpressure={
+            "backpressure_cls": "EMABackpressureDetector",
+            "high_water_s": 0.5,
+            "low_water_s": 0.05,
+        },
+    )
+
+    resolved = spec.secondary_tier_configs[0]["backpressure"]
+    # Tier override wins for the fields it sets.
+    assert resolved["high_water_s"] == 0.1
+    assert resolved["low_water_s"] == 0.02
+    # Fields only in top-level still fill in.
+    assert resolved["backpressure_cls"] == "EMABackpressureDetector"
+
+
+def test_tier_without_backpressure_stays_unconfigured():
+    spec = _build_tiering_spec([{"type": "example"}])
+    assert "backpressure" not in spec.secondary_tier_configs[0]
diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
index 47fb7103d742..5b3bf1c8bebe 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import MagicMock, call
 
 import pytest
@@ -32,6 +33,7 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
     OffloadingConnectorScheduler,
     RequestOffloadState,
+    _create_req_context,
     get_sliding_window_size_in_chunks,
 )
 from vllm.v1.core.block_pool import BlockPool
@@ -59,6 +61,7 @@
     SlidingWindowSpec,
 )
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+from vllm.v1.kv_hints import KvHintsEnvelope
 from vllm.v1.kv_offload.base import (
     GPULoadStoreSpec,
     LookupResult,
@@ -74,30 +77,49 @@
     make_offload_key,
 )
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
-from vllm.v1.outputs import KVConnectorOutput
+from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 from vllm.v1.request import RequestStatus
 
 
-@pytest.mark.parametrize("boundary", [3840, 3904])
+@pytest.mark.parametrize(
+    "window,extra_retained", [(128, 0), (128, 3), (256, 1), (256, 3)]
+)
+@pytest.mark.parametrize("boundary", [3840, 3904, 3968])
 @pytest.mark.parametrize("eagle", [False, True])
 @pytest.mark.parametrize("left_state", ["present", "missing", "pending", "loading"])
-def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
+def test_swa_offload_window_covers_unaligned_hit(
+    boundary, eagle, left_state, window, extra_retained
+):
     """A smaller SWA group can move the hit inside another group's CPU chunk."""
     groups = []
-    for i, (block_size, window) in enumerate([(256, None), (64, 128), (8, 128)]):
+    for i, (block_size, group_window) in enumerate(
+        [(256, None), (64, window), (8, 128)]
+    ):
         kwargs = dict(
             block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32
         )
         kv_spec = (
             FullAttentionSpec(**kwargs)
-            if window is None
-            else SlidingWindowSpec(**kwargs, sliding_window=window)
+            if group_window is None
+            else SlidingWindowSpec(
+                **kwargs,
+                sliding_window=group_window,
+                extra_retained_tokens=extra_retained if i == 1 else 0,
+            )
         )
         groups.append(
             KVCacheGroupSpec([f"layer{i}"], kv_spec, is_eagle_group=eagle and i == 1)
         )
     manager = CPUOffloadingManager(num_chunks=100)
     spec = SimpleNamespace(
+        config=SimpleNamespace(
+            groups=tuple(
+                SimpleNamespace(
+                    group_id=i, tokens_per_block=group.kv_cache_spec.block_size
+                )
+                for i, group in enumerate(groups)
+            )
+        ),
         tokens_per_block=(256, 64, 8),
         tokens_per_hash=8,
         blocks_per_chunk=4,
@@ -144,7 +166,9 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
         sched._chunks_being_loaded.add(left_key)
 
     external, load_async = sched.get_num_new_matched_tokens(request, 0)
-    if boundary == 3904 and left_state != "present":
+    first_retained_block = (boundary - window + 1 - extra_retained) // 64
+    needs_left = first_retained_block // 4 < end - 1
+    if needs_left and left_state != "present":
         assert external == (0 if left_state == "missing" else None)
         assert not load_async
         return
@@ -168,8 +192,8 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
     sched.update_state_after_alloc(request, blocks, external)
     assert len(sched._jobs) == 1
     job = next(iter(sched._jobs.values()))
-    if boundary == 3904:
-        assert {swa[14], swa[15]} <= job.keys
+    if needs_left:
+        assert {left_key, swa[end - 1]} <= job.keys
     manager.complete_load(job.keys, state.req_context)
 
 
@@ -187,20 +211,38 @@ def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler:
 
 def _make_partial_tail_request(
     scheduler: OffloadingConnectorScheduler,
+    kv_transfer_params: dict[str, Any] | None = None,
 ) -> MagicMock:
     request = MagicMock()
     request.request_id = "req"
-    request.kv_transfer_params = None
+    request.kv_transfer_params = kv_transfer_params
+    request.kv_hints = None
     request.num_prompt_tokens = 30
     request.num_tokens = 30
     request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)]
     request.all_token_ids = list(range(30))
     request.lora_request = None
+    request.skip_reading_prefix_cache = False
     request.is_finished.return_value = False
     scheduler.on_new_request(request)
     return request
 
 
+def test_create_req_context_preserves_kv_hints():
+    kv_hints = KvHintsEnvelope(
+        protocol_version="0.1",
+        message_id="message-1",
+        actions=[],
+    )
+    request = MagicMock(
+        request_id="req",
+        kv_transfer_params=None,
+        kv_hints=kv_hints,
+    )
+
+    assert _create_req_context(request).kv_hints is kv_hints
+
+
 def _reduce_kv_connector_stats(runner):
     reduced: dict[str, int | float] = {}
     for payload in runner.kv_connector_stats:
@@ -234,6 +276,9 @@ def test_partial_tail_store_uses_attention_and_recurrent_cow_sources():
     jobs = scheduler._build_partial_tail_store_jobs(output)
 
     assert len(jobs) == 1
+    offered_keys = scheduler.manager.prepare_store.call_args.args[0]
+    req_context = scheduler._req_status["req"].req_context
+    assert all(req_context.get_offload_key_position(key) == 28 for key in offered_keys)
     [job_id] = jobs
     src_spec = jobs[job_id].src_spec
     assert isinstance(src_spec, GPULoadStoreSpec)
@@ -303,7 +348,24 @@ def test_aligned_boundary_store_uses_exact_source_with_partial_tail():
     assert src_spec.block_indices == [0, 0]
 
 
-def test_aligned_boundary_store_flushes_before_cow_destination_reuse():
+def test_aligned_boundary_store_maps_sparse_group_id_to_dense_transfer_slot():
+    scheduler = _make_partial_tail_scheduler()
+    indexer_config = scheduler.config.kv_group_configs[1]
+    scheduler.config = scheduler.config._replace(kv_group_configs=(indexer_config,))
+    _make_partial_tail_request(scheduler)
+    scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    jobs = scheduler._build_aligned_boundary_store_jobs({"req": [(1, 98, 16)]})
+
+    [job] = jobs.values()
+    assert job.src_spec.group_sizes == [1]
+    assert job.src_spec.block_indices == [0]
+
+
+@pytest.mark.parametrize("cow_reuse", [False, True])
+def test_aligned_boundary_store_flushes_before_block_reuse(cow_reuse):
     scheduler = _make_partial_tail_scheduler()
     _make_partial_tail_request(scheduler)
     scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -320,7 +382,11 @@ def test_aligned_boundary_store_flushes_before_cow_destination_reuse():
     [job_id] = meta.store_jobs
 
     output = SchedulerOutput.make_empty()
-    output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)]
+    if cow_reuse:
+        output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)]
+    else:
+        output.scheduled_cached_reqs.req_ids = ["req"]
+        output.scheduled_cached_reqs.new_block_ids = [([], [99])]
     meta = scheduler.build_connector_meta(output)
 
     assert meta.jobs_to_flush == {job_id}
@@ -335,6 +401,13 @@ def test_normal_store_excludes_align_mode_mamba_sources():
     req_status.group_states[0].block_ids[:] = [11]
     req_status.group_states[1].block_ids[:] = [99]
     req_status.update_offload_keys()
+    for group_config, group_state in zip(
+        scheduler.config.kv_group_configs, req_status.group_states
+    ):
+        assert (
+            req_status.req_context.get_offload_key_position(group_state.offload_keys[0])
+            == group_config.tokens_per_chunk
+        )
     scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
     )
@@ -384,6 +457,67 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys():
     assert req_status.partial_tail_boundary is None
 
 
+def test_lookup_cap_stops_at_authoritative_prefix_boundary():
+    scheduler = _make_partial_tail_scheduler()
+    request = _make_partial_tail_request(scheduler)
+    request.skip_reading_prefix_cache = False
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    tokens, load_async = scheduler.get_num_new_matched_tokens(
+        request, 0, max_num_new_tokens=20
+    )
+
+    assert (tokens, load_async) == (20, True)
+    assert scheduler._req_status["req"].partial_tail_boundary == 20
+
+
+def test_max_load_tokens_zero_skips_partial_tail_lookup():
+    scheduler = _make_partial_tail_scheduler()
+    request = _make_partial_tail_request(
+        scheduler, kv_transfer_params={"max_load_tokens": 0}
+    )
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    assert scheduler.get_num_new_matched_tokens(request, 16) == (0, False)
+    scheduler.manager.lookup.assert_not_called()
+    assert all(
+        state.num_hit_chunks == 1 for state in scheduler._req_status["req"].group_states
+    )
+
+
+def test_max_load_tokens_caps_tokens_beyond_gpu_prefix():
+    scheduler = _make_partial_tail_scheduler()
+    request = _make_partial_tail_request(
+        scheduler, kv_transfer_params={"max_load_tokens": 8}
+    )
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    assert scheduler.get_num_new_matched_tokens(request, 16) == (8, True)
+    assert scheduler._req_status["req"].partial_tail_boundary == 24
+
+
+def test_max_load_tokens_rounds_down_without_partial_tail():
+    scheduler = _make_partial_tail_scheduler()
+    scheduler.config = scheduler.config._replace(supports_partial_tail=False)
+    request = _make_partial_tail_request(
+        scheduler, kv_transfer_params={"max_load_tokens": 20}
+    )
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    assert scheduler.get_num_new_matched_tokens(request, 0) == (16, True)
+
+
+@pytest.mark.parametrize("value", ["8", 8.5, -1, True])
+def test_invalid_max_load_tokens_is_ignored(value):
+    scheduler = _make_partial_tail_scheduler()
+    request = _make_partial_tail_request(
+        scheduler, kv_transfer_params={"max_load_tokens": value}
+    )
+    scheduler.manager.lookup.return_value = LookupResult.HIT
+
+    assert scheduler.get_num_new_matched_tokens(request, 0) == (28, True)
+
+
 def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary():
     """A recurrent group may hold unhashed blocks inside the computed prefix.
 
@@ -815,7 +949,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.run(decoded_tokens=[0] * (tokens_per_chunk + 1))
 
     # 1 more block (+ token for kicking off offloading)
-    # now check touch was called with all 6 blocks
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
     )
@@ -823,9 +956,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
         decoded_tokens=[0] * (tokens_per_chunk + 1),
         expected_stored=(15, 16, 17),
     )
-    runner.manager.touch.assert_called()
-    block_hashes1 = list(runner.manager.touch.call_args.args[0])
-    assert len(block_hashes1) == 6
 
     # terminate request
     runner.run(decoded_tokens=[EOS_TOKEN_ID])
@@ -833,13 +963,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     # create a new request differing only on the last token
     runner.new_request(token_ids=[0] * (tokens_per_chunk * 6 - 1) + [1])
     runner.run(decoded_tokens=[0])
-    runner.manager.touch.assert_called()
-    block_hashes2 = list(runner.manager.touch.call_args.args[0])
-    assert len(block_hashes2) == 6
-
-    # verify hashes are the same, except for the last block
-    assert block_hashes1[:5] == block_hashes2[:5]
-    assert block_hashes1[5] != block_hashes2[5]
+    runner.manager.touch.assert_not_called()
 
     # terminate request
     runner.run(
@@ -1247,14 +1371,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
         generate_store_output(keys)
     )
     runner.run(decoded_tokens=[0])
-    # _touch called from get_num_new_matched_tokens (2 groups) and
-    # _get_reqs_to_store (2 groups) → 4 touch calls total.
-    touch_calls = runner.manager.touch.call_args_list
-    assert len(touch_calls) == 4
-    assert len(touch_calls[0].args[0]) == 3
-    assert len(touch_calls[1].args[0]) == 3
-    assert len(touch_calls[2].args[0]) == 3
-    assert len(touch_calls[3].args[0]) == 3
+    runner.manager.touch.assert_not_called()
 
     # store 3 more block
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
@@ -1265,9 +1382,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
         expected_stored=(0, 1, 2, 3, 4, 5),
     )
 
-    # touch called from _get_reqs_to_store * 3 blocks, once for each group
-    touch_calls = runner.manager.touch.call_args_list
-    assert len(touch_calls) == 6
+    runner.manager.touch.assert_not_called()
 
     # EOS lands in the last slot of the 7th block (offset 6). No forward pass
     # writes that slot, so the finishing step declines the block.
@@ -1286,13 +1401,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
         expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),
     )
 
-    # 2 touch calls from get_num_new_matched_tokens (2 groups)
-    touch_calls = runner.manager.touch.call_args_list
-    assert len(touch_calls) == 2
-    # full attention group touched all 3 blocks
-    assert len(touch_calls[0].args[0]) == 3
-    # sliding window group touched just the last 2 blocks
-    assert len(touch_calls[1].args[0]) == 2
+    runner.manager.touch.assert_not_called()
 
     # 3 blocks are hit on GPU [0, 1, 2]
     # 1 block loaded [3,]
@@ -1358,15 +1467,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
         generate_store_output(keys)
     )
     runner.run(decoded_tokens=[0])
-    # _touch called from get_num_new_matched_tokens (2 groups) and
-    # _get_reqs_to_store (2 groups) → 4 touch calls total.
-    # Group 0 has 2 offload keys, group 1 has 1.
-    touch_calls = runner.manager.touch.call_args_list
-    assert len(touch_calls) == 4
-    assert len(touch_calls[0].args[0]) == 2
-    assert len(touch_calls[1].args[0]) == 1
-    assert len(touch_calls[2].args[0]) == 2
-    assert len(touch_calls[3].args[0]) == 1
+    runner.manager.touch.assert_not_called()
 
     # Get to 31 tokens
     # No further blocks offloaded
@@ -1376,11 +1477,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
     # Group 0 blocks: [0, 1], ending_token_offset = 24
     # Group 1 blocks: [0, 1], ending_token_offset = 32
     runner.run(decoded_tokens=[0])
-    # _get_reqs_to_store touch: only group 1 has a new block to store
-    touch_calls = runner.manager.touch.call_args_list
-    assert len(touch_calls) == 2
-    assert len(touch_calls[0].args[0]) == 2
-    assert len(touch_calls[1].args[0]) == 2
+    runner.manager.touch.assert_not_called()
 
     # Get to 35 tokens
     # No further blocks offloaded
@@ -1390,11 +1487,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
     # Group 0 blocks: [0, 1, 2], ending_token_offset = 36
     # Group 1 blocks: [0, 1], ending_token_offset = 32
     runner.run(decoded_tokens=[0])
-    # _get_reqs_to_store touch: only group 0 has a new block to store
-    touch_calls = runner.manager.touch.call_args_list
-    assert len(touch_calls) == 2
-    assert len(touch_calls[0].args[0]) == 3
-    assert len(touch_calls[1].args[0]) == 2
+    runner.manager.touch.assert_not_called()
 
     # Get to 47 tokens
     # No further blocks offloaded
@@ -1404,11 +1497,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
     # Group 0 blocks: [0, 1, 2, 3], ending_token_offset = 4
     # Group 1 blocks: [0, 1, 2], ending_token_offset = 48
     runner.run(decoded_tokens=[0])
-    # _get_reqs_to_store touch: both groups have a new block, each with 1 key
-    touch_calls = runner.manager.touch.call_args_list
-    assert len(touch_calls) == 2
-    assert len(touch_calls[0].args[0]) == 4
-    assert len(touch_calls[1].args[0]) == 3
+    runner.manager.touch.assert_not_called()
 
     runner.run(decoded_tokens=[0], expected_stored=((0, 3), (1, 2)))
 
@@ -1665,6 +1754,15 @@ def test_retry_stops_at_miss(self):
 
 
 class TestSlidingWindowLookup:
+    def test_pending_chunk_behind_gap_does_not_defer(self):
+        """A pending suffix too short for the window cannot improve the hit."""
+        sched = _make_scheduler_with_lookup(
+            {1: LookupResult.HIT, 2: LookupResult.HIT, 4: LookupResult.HIT_PENDING}
+        )
+        assert (
+            sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 2
+        )
+
     def test_all_hit_exact_window(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
         assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2
@@ -1783,6 +1881,73 @@ def test_hit_pending_does_not_break_streak(self):
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.parametrize("async_scheduling", [False, True])
+def test_swa_load_does_not_wait_for_unusable_pending_suffix(
+    request_runner, async_scheduling
+):
+    """A CPU store beyond a missing SWA chunk must not block a ready prefix."""
+    groups = [
+        KVCacheGroupSpec(
+            ["full"],
+            FullAttentionSpec(
+                block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32
+            ),
+        ),
+        KVCacheGroupSpec(
+            ["swa"],
+            SlidingWindowSpec(
+                block_size=4,
+                num_kv_heads=1,
+                head_size=1,
+                dtype=torch.float32,
+                sliding_window=8,
+            ),
+        ),
+    ]
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=32,
+        async_scheduling=async_scheduling,
+        kv_cache_groups=groups,
+    )
+    sched = runner.connector_scheduler
+    manager = CPUOffloadingManager(num_chunks=16)
+    sched.manager = manager
+    runner.new_request(token_ids=list(range(17)))
+    state = sched._req_status["0"]
+    state.update_offload_keys()
+    full, swa = (g.offload_keys for g in state.group_states)
+    ready = manager.prepare_store(full[:4] + swa[:2], state.req_context)
+    assert ready is not None
+    manager.complete_store(ready.keys_to_store, state.req_context)
+    pending = manager.prepare_store(swa[3:4], state.req_context)
+    assert pending is not None
+
+    # The last SWA chunk is still being written; the preceding one is missing.
+    output = runner.scheduler.schedule()
+    metadata = output.kv_connector_metadata
+    assert isinstance(metadata, OffloadingConnectorMetadata)
+    [(job_id, load)] = metadata.load_jobs.items()
+    assert isinstance(load.dst_spec, GPULoadStoreSpec)
+    assert load.dst_spec.group_sizes == [2, 2]
+    assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING
+
+    runner.scheduler.update_from_output(
+        output,
+        ModelRunnerOutput.with_kv_conn_output_only(
+            KVConnectorOutput(
+                finished_recving={"0"},
+                kv_connector_worker_meta=OffloadingWorkerMetadata(
+                    completed_jobs={job_id: 1}
+                ),
+            )
+        ),
+    )
+    resumed = runner.scheduler.schedule()
+    assert resumed.num_scheduled_tokens == {"0": 9}
+    assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING
+
+
 @pytest.mark.parametrize("alignment_chunk_count", [4, 8, 64])
 @pytest.mark.parametrize("sliding_window_chunks", [1, 2, 3])
 @pytest.mark.parametrize("is_eagle_group", [False, True])
@@ -2460,27 +2625,14 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner
     assert req_status.transfer_jobs == pending_store_jobs
 
 
+@pytest.mark.parametrize("extra_retained", [0, 1, 3])
 @pytest.mark.parametrize("async_scheduling", [True, False])
-def test_swa_alignment_skip(request_runner, async_scheduling: bool):
-    """SWA blocks unreachable by the load path are skipped during store.
-
-    Simulates a DeepSeek V4-like hybrid architecture where SWA groups have
-    much smaller block sizes than the full-attention (MLA) group, causing
-    most SWA blocks to be unreachable by the alignment-based load path.
-
-    Setup:
-      - Group 0: full attention (MLA-like), block_size=16
-      - Group 1: SWA, block_size=4, sliding_window=8
-
-    alignment_tokens = 16, so 16 / 4 = 4 SWA blocks per alignment segment.
-    sliding_window_size_in_chunks = ceil(8 / 4) = 2.
-    Within each segment of 4 SWA blocks, only the trailing 2 are stored.
+def test_swa_alignment_skip(request_runner, async_scheduling: bool, extra_retained):
+    """Store the window a cold consumer must restore at each FA boundary.
 
-    With 32 tokens (2 full-attn blocks, 8 SWA blocks):
-      - Group 0 stores: blocks 0, 1  (all full-attn blocks)
-      - Group 1 stores: blocks 2, 3, 6, 7  (skip 0,1,4,5)
-
-    For real DeepSeek V4 (100K tokens), this reduces SWA stores by ~78%.
+    With FA blocks of 16 tokens and SWA blocks of 4 tokens, a window of 8
+    needs the last two blocks per segment. Retaining three extra tokens for
+    MTP requires one more block; retaining just one token does not.
     """
     full_attn_block_size = 16
     swa_block_size = 4
@@ -2505,6 +2657,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
                 head_size=1,
                 dtype=torch.float32,
                 sliding_window=sliding_window,
+                extra_retained_tokens=extra_retained,
             ),
         ),
     ]
@@ -2522,16 +2675,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
     # Group 0: full attention
     assert kv_group_configs[0].sliding_window_size_in_chunks is None
     assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size
-    # Group 1: SWA with sliding_window_size of 2 chunks
-    assert kv_group_configs[1].sliding_window_size_in_chunks == 2
+    # Group 1: SWA, including any additional retained blocks
+    assert kv_group_configs[1].sliding_window_size_in_chunks == (
+        3 if extra_retained > 1 else 2
+    )
     assert kv_group_configs[1].tokens_per_chunk == swa_block_size
     # alignment_tokens = full_attn_block_size = 16
     assert runner.connector_scheduler.config.alignment_tokens == full_attn_block_size
 
-    # Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks
-    # (block_size=4). Decode 1 token to kick off processing (stores are
-    # deferred to next step).
-    num_tokens = 32
+    # Keep a partial last block so decoding does not reuse the producer's
+    # freed blocks before the harness records its deferred stores.
+    num_tokens = 33
     runner.new_request(token_ids=[0] * num_tokens)
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
@@ -2544,12 +2698,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
     )
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
-        # Group 0 (full attn, block_size=16): 2 offloaded chunks
-        #   -> GPU blocks (0, 0) and (0, 1)
-        # Group 1 (SWA, block_size=4): 8 offloaded chunks, skip first 2
-        #   per segment of 4:
-        #   Segment 0 (blocks 0-3): skip 0,1 -> store (1, 2), (1, 3)
-        #   Segment 1 (blocks 4-7): skip 4,5 -> store (1, 6), (1, 7)
+        # Retain each segment's tail, including MTP history when needed.
         expected_stored=(
             (0, 0),
             (0, 1),
@@ -2557,27 +2706,26 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
             (1, 3),
             (1, 6),
             (1, 7),
-        ),
+        )
+        + (((1, 1), (1, 5)) if extra_retained > 1 else ()),
     )
 
     # Verify that loads still work correctly for the stored SWA blocks.
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * num_tokens + [1])
-    runner.manager.lookup.return_value = LookupResult.HIT
-    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2
+    runner.manager.lookup.side_effect = lambda key, ctx: (
+        LookupResult.HIT if (key, 0) in runner.offloaded else LookupResult.MISS
+    )
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
-        # Group 0: full prefix lookup hits 2 offloaded chunks
-        #   -> loads GPU blocks (0, 0), (0, 1)
-        # Group 1: sliding window lookup finds trailing 2 from last segment
-        #   (blocks 6, 7 which were stored)
-        #   -> loads GPU blocks (1, 6), (1, 7)
+        # The consumer restores the full FA prefix and the final SWA tail.
         expected_loaded=(
             (0, 0),
             (0, 1),
             (1, 6),
             (1, 7),
-        ),
+        )
+        + (((1, 5),) if extra_retained > 1 else ()),
     )
 
 
@@ -2845,6 +2993,47 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):
     runner.manager.lookup.assert_not_called()
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_max_load_tokens_limits_external_load(request_runner, async_scheduling: bool):
+    """The load cap limits external reads without changing the store path."""
+    block_size = 4
+    blocks_per_chunk = 3
+    tokens_per_chunk = block_size * blocks_per_chunk
+
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+        blocks_per_chunk=blocks_per_chunk,
+    )
+
+    runner.new_request(token_ids=[0] * tokens_per_chunk)
+    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.scheduler.reset_prefix_cache()
+    runner.manager.lookup.reset_mock()
+    runner.new_request(
+        token_ids=[0] * tokens_per_chunk,
+        kv_transfer_params={"max_load_tokens": 0},
+    )
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_loaded=(),
+        expected_stored=(0, 1, 2),
+    )
+
+    runner.manager.lookup.assert_not_called()
+
+
 # ---------------------------------------------------------------------------
 # Eagle/MTP test class
 # ---------------------------------------------------------------------------
@@ -3251,59 +3440,6 @@ def test_non_eagle_tighten_clears_eagle_verified(self, request_runner):
         #   whole request.
         assert sched._lookup(req_status) == 4
 
-    def test_eagle_verified_survives_eagle_tighten(self, request_runner):
-        """Eagle group tightening does NOT clear eagle_verified.
-
-        Groups: 0=non-eagle full-attn, 1=eagle full-attn.
-        Group 0 finds 3 hits (max_hit=12). Group 1 finds 3 hits, pops to 2
-        (max_hit=8). Since group 1 IS eagle, eagle_verified is NOT cleared.
-        Result: 8 tokens (eagle only pops once).
-        """
-        block_size = 4
-        groups = [
-            KVCacheGroupSpec(
-                ["layer0"],
-                FullAttentionSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                ),
-                is_eagle_group=False,
-            ),
-            KVCacheGroupSpec(
-                ["layer1"],
-                FullAttentionSpec(
-                    block_size=block_size,
-                    num_kv_heads=2,
-                    head_size=1,
-                    dtype=torch.float32,
-                ),
-                is_eagle_group=True,
-            ),
-        ]
-        runner = request_runner(
-            block_size=block_size,
-            num_gpu_blocks=100,
-            async_scheduling=False,
-            kv_cache_groups=groups,
-        )
-        runner.manager.lookup.side_effect = lambda key, req_context: (
-            LookupResult.HIT
-            if int(get_offload_block_hash(key).decode()) in {1, 2, 3}
-            else LookupResult.MISS
-        )
-        sched = runner.connector_scheduler
-        req_status = self._make_req_status(
-            sched,
-            num_tokens=12,
-            offload_keys_per_group=[[1, 2, 3], [1, 2, 3]],
-        )
-        # Group 0: 3 hits → max_hit=12, num_hit=12
-        # Group 1 (eagle): 3 hits, pop to 2 → max_hit=8, num_hit=8
-        # Tightened but IS eagle → no clear. No re-iteration.
-        assert sched._lookup(req_status) == 8
-
     # -------------------------------------------------------------------
     # Integration tests: store and load via request_runner
     # -------------------------------------------------------------------
diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py
index 366624e1d3b1..32bd4f3eca75 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py
@@ -13,6 +13,7 @@
 from vllm.utils.torch_utils import get_dtype_size
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -97,10 +98,9 @@ def _make_worker(
     kv_cache_config: KVCacheConfig,
     replicated_layout: bool = False,
     rank: int = 0,
+    canonical_layout: bool = False,
 ):
-    """
-    Create an OffloadingConnectorWorker with mocked dependencies.
-    """
+    """Create an OffloadingConnectorWorker with mocked dependencies."""
     from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
         OffloadingConnectorWorker,
     )
@@ -108,6 +108,7 @@ def _make_worker(
     spec = MagicMock(spec=OffloadingSpec)
     spec.replicated_layout = replicated_layout
     spec.config = MagicMock()
+    spec.config.canonical_layout = canonical_layout
     spec.config.parallel.rank = rank
     spec.get_worker.return_value = MagicMock()
 
@@ -206,11 +207,13 @@ def test_prepare_store_kv_non_writer_marks_completed_without_submit():
     assert meta.completed_jobs == {7: 1}
 
 
-def test_prepare_store_kv_writer_submits_store():
+@pytest.mark.parametrize("rank,canonical_layout", [(0, False), (0, True), (1, True)])
+def test_prepare_store_kv_writer_submits_store(rank, canonical_layout):
     worker, _ = _make_worker(
         KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
         replicated_layout=True,
-        rank=0,
+        rank=rank,
+        canonical_layout=canonical_layout,
     )
 
     worker.prepare_store_kv(_store_metadata(8))
@@ -518,6 +521,47 @@ def test_register_kv_caches(backend):
             assert actual.mapping is not None
 
 
+def test_register_packed_kv_caches_skips_scratch_group():
+    attn_spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    page = attn_spec.page_size_bytes
+    layers = ["layer0", "layer1"]
+    kv_cache_config = KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=2 * page * NUM_BLOCKS,
+                layers=layers,
+                layer_stride=page,
+                block_stride=2 * page,
+            )
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(layers, attn_spec),
+            KVCacheGroupSpec(
+                ["scratch"],
+                CircularBufferSpec(
+                    block_size=4, num_kv_heads=1, head_size=1, dtype=DTYPE
+                ),
+            ),
+        ],
+    )
+    kv_caches = _allocate_kv_caches(
+        kv_cache_config, [], device=torch.device(f"{DEVICE_TYPE}:0")
+    )
+    worker, spec = _make_worker(kv_cache_config)
+    worker.register_kv_caches(kv_caches)
+
+    canonical = spec.get_worker.call_args[0][0]
+    assert len(canonical.tensors) == 1
+    assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, 2 * page)
+    assert canonical.group_data_refs == [[CanonicalKVCacheRef(0, 2 * page)]]
+
+
 @pytest.mark.parametrize("backend", ATTN_BACKENDS)
 def test_register_kv_caches_uniform_type(backend):
     """Test register_kv_caches with UniformTypeKVCacheSpecs.
diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py
index 5707003ce935..ad8d4fc362b9 100644
--- a/tests/v1/kv_connector/unit/offloading_connector/utils.py
+++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py
@@ -475,8 +475,7 @@ def _run(
         complete_transfers: bool,
         post_step_fn: Callable[[], None] | None = None,
     ):
-        """
-        Runs multiple engine (scheduler + worker) steps.
+        """Runs multiple engine (scheduler + worker) steps.
         Assumes a single request is running.
 
         Args:
@@ -484,8 +483,8 @@ def _run(
             complete_transfers: complete transfers immediately
             post_step_fn: optional callback invoked after each step's
                 update_from_output(), before the next schedule().
-        """
 
+        """
         tokens_iter = iter(decoded_tokens)
         token_id = next(tokens_iter, None)
         prev_scheduler_output = None
@@ -617,8 +616,7 @@ def run(
         expected_flushed: tuple[int | tuple[int, int], ...] = (),
         post_step_fn: Callable[[], None] | None = None,
     ):
-        """
-        Runs multiple engine (scheduler + worker) steps.
+        """Runs multiple engine (scheduler + worker) steps.
         Assumes a single request is running.
 
         Args:
@@ -634,8 +632,8 @@ def run(
             A GPU block is either a (group_idx: int, request_block_offset: int)
             or just request_block_offset: int.
             The latter case is a convenience for representing all groups.
-        """
 
+        """
         expected_stored_gpu_blocks = self._to_gpu_blocks(expected_stored)
         expected_loaded_gpu_blocks = self._to_gpu_blocks(expected_loaded)
         expected_flushed_gpu_blocks = self._to_gpu_blocks(expected_flushed)
diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py
index 78a6d930f97a..3d539f97c8cb 100644
--- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py
+++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py
@@ -969,6 +969,9 @@ def test_turn2_deadline_gate(dist_init, offset, expiry_delta, expect_declined):
     assert worker.xfer_stats.data["num_kv_expired_reqs"] == (
         [1] if expect_declined else []
     )
+    # KV expiry must not pollute the transport failure counters.
+    assert worker.xfer_stats.data["num_failed_transfers"] == []
+    assert worker.xfer_stats.data["num_failed_handshakes"] == []
     _, done_recving = connector.get_finished(finished_req_ids=set())
     assert "req" in done_recving
 
diff --git a/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py b/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py
index ec3fb8231e19..21b50a13794b 100644
--- a/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py
+++ b/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-test that invalid blocks are evicted from prefix cache to prevent pollution.
+"""test that invalid blocks are evicted from prefix cache to prevent pollution.
 
 verifies that when sync-loading fails, invalid blocks are removed from the
 prefix cache hash table so future requests cannot match and reuse corrupted data.
@@ -39,7 +38,7 @@ def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
 
 @pytest.fixture
 def fail_scheduler():
-    """scheduler with kv_load_failure_policy='fail'"""
+    """Scheduler with kv_load_failure_policy='fail'."""
     vllm_config = create_vllm_config()
     vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
     return create_scheduler(vllm_config)
@@ -48,8 +47,7 @@ def fail_scheduler():
 def test_invalid_blocks_evicted_prevents_cache_pollution(
     fail_scheduler: Scheduler,
 ):
-    """
-    verify invalid blocks are evicted to prevent future cache hits.
+    """Verify invalid blocks are evicted to prevent future cache hits.
 
     scenario:
     1. request 1 loads externally-computed blocks (sync mode)
diff --git a/tests/v1/kv_connector/unit/test_decode_bench_connector.py b/tests/v1/kv_connector/unit/test_decode_bench_connector.py
index 76bb4648a933..e1054d9ef86f 100644
--- a/tests/v1/kv_connector/unit/test_decode_bench_connector.py
+++ b/tests/v1/kv_connector/unit/test_decode_bench_connector.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for DecodeBenchConnector.
+"""Unit tests for DecodeBenchConnector.
 
 Tests the functionality of the DecodeBenchConnector which fills KV cache
 with dummy values for decode performance benchmarking.
diff --git a/tests/v1/kv_connector/unit/test_error_propagation.py b/tests/v1/kv_connector/unit/test_error_propagation.py
index a07364cd3ea1..2f33ab18b917 100644
--- a/tests/v1/kv_connector/unit/test_error_propagation.py
+++ b/tests/v1/kv_connector/unit/test_error_propagation.py
@@ -32,14 +32,14 @@ def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
 
 @pytest.fixture
 def fail_scheduler():
-    """scheduler with kv_load_failure_policy='fail'"""
+    """Scheduler with kv_load_failure_policy='fail'."""
     vllm_config = create_vllm_config()
     vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
     return create_scheduler(vllm_config)
 
 
 def test_error_propagation_sync_load(fail_scheduler: Scheduler):
-    """test invalid_block_ids with fail policy -> FINISHED_ERROR (sync load)"""
+    """Test invalid_block_ids with fail policy -> FINISHED_ERROR (sync load)."""
     num_prompt_blocks = 100
     num_external_computed_blocks = 99
     invalid_block_idx = 50
@@ -93,7 +93,7 @@ def test_error_propagation_sync_load(fail_scheduler: Scheduler):
 
 
 def test_error_propagation_async_load(fail_scheduler: Scheduler):
-    """test invalid_block_ids with fail policy -> FINISHED_ERROR (async load)"""
+    """Test invalid_block_ids with fail policy -> FINISHED_ERROR (async load)."""
     num_prompt_blocks = 100
     num_external_computed_blocks = 99
     invalid_block_idx = 50
diff --git a/tests/v1/kv_connector/unit/test_example_connector.py b/tests/v1/kv_connector/unit/test_example_connector.py
index 923eb22b31f5..d5b1856cc8b9 100644
--- a/tests/v1/kv_connector/unit/test_example_connector.py
+++ b/tests/v1/kv_connector/unit/test_example_connector.py
@@ -29,12 +29,12 @@ class InputCase(NamedTuple):
 
 
 def _check_path_len(path):
-    """Return the latest length in path"""
+    """Return the latest length in path."""
     return len(list(path.iterdir()))
 
 
 def _list_path(path):
-    """Return the list of foldername (hashes generated) under the path"""
+    """Return the list of foldername (hashes generated) under the path."""
     return list(path.iterdir())
 
 
@@ -47,8 +47,7 @@ def run_test(
     expected_len: int,
     info: str,
 ):
-    """
-    One individual test to process the prompt and output base on 1 set of input
+    """One individual test to process the prompt and output base on 1 set of input
     Then check if the length in the storage path matches the expected length
     `info` introduces details or purpose of the individual test
     """
@@ -67,9 +66,7 @@ def run_test(
 
 
 def process_prompt(processor, llm: LLM, question: str, image_urls: list[Image]):
-    """
-    Form the prompt based on the text and image input, then llm generate output
-    """
+    """Form the prompt based on the text and image input, then llm generate output."""
     placeholders = [
         {
             "type": "image_url",
@@ -120,8 +117,7 @@ def process_prompt(processor, llm: LLM, question: str, image_urls: list[Image]):
     ),
 )
 def test_shared_storage_connector_hashes(tmp_path, attn_backend):
-    """
-    Tests that ExampleConnector saves KV to the storage locations
+    """Tests that ExampleConnector saves KV to the storage locations
     with proper hashes; that are unique for inputs with identical text but
     different images (same size), or same multiple images but different orders.
     """
diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py
index e5cbd977b6d8..509d3d3234b9 100644
--- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py
+++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py
@@ -79,11 +79,6 @@ def get_kv_connector(self) -> KVConnectorBase_V1:
         "resolve_kv_cache_block_sizes",
         lambda kv_cache_config, vllm_config: (16, 16),
     )
-    monkeypatch.setattr(
-        engine_core_module,
-        "MULTIMODAL_REGISTRY",
-        SimpleNamespace(engine_receiver_cache_from_config=lambda vllm_config: None),
-    )
     monkeypatch.setattr(engine_core_module, "freeze_gc_heap", lambda: None)
     monkeypatch.setattr(
         engine_core_module, "maybe_attach_gc_debug_callback", lambda: None
@@ -105,7 +100,11 @@ def get_kv_connector(self) -> KVConnectorBase_V1:
         speculative_config=None,
         ec_transfer_config=None,
         max_concurrent_batches=1,
-        model_config=SimpleNamespace(runner_type="generate", is_diffusion=False),
+        model_config=SimpleNamespace(
+            runner_type="generate",
+            is_diffusion=False,
+            supports_multimodal_inputs=False,
+        ),
         cache_config=SimpleNamespace(
             enable_prefix_caching=False,
             prefix_caching_hash_algo="builtin",
diff --git a/tests/v1/kv_connector/unit/test_hf3fs_client.py b/tests/v1/kv_connector/unit/test_hf3fs_client.py
index d9c34a8907de..26a5de11c522 100644
--- a/tests/v1/kv_connector/unit/test_hf3fs_client.py
+++ b/tests/v1/kv_connector/unit/test_hf3fs_client.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for resource management in hf3fs_client.py: constructor failure cleanup
+"""Tests for resource management in hf3fs_client.py: constructor failure cleanup
 and idempotent close().  Tests use mock to replace real I/O operations
 (hf3fs_fuse.io, SharedMemory, os, CUDA).
 Requires hf3fs_fuse.io to be installed; skipped otherwise.
diff --git a/tests/v1/kv_connector/unit/test_hf3fs_connector.py b/tests/v1/kv_connector/unit/test_hf3fs_connector.py
index 94bb94c6fbd7..de4d53279877 100644
--- a/tests/v1/kv_connector/unit/test_hf3fs_connector.py
+++ b/tests/v1/kv_connector/unit/test_hf3fs_connector.py
@@ -1,14 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for HF3FS KV Connector high-level components:
-  - TestHf3fsMockClient      : file-backed mock client I/O correctness
-  - TestHF3FSKVConnectorStats: metric collection, aggregation, serialisation
+"""Tests for HF3FS KV Connector high-level components:
+- TestHf3fsMockClient      : file-backed mock client I/O correctness
+- TestHF3FSKVConnectorStats: metric collection, aggregation, serialisation
 """
 
 import os
 from unittest.mock import MagicMock
 
+import numpy as np
 import pytest
 import torch
 
@@ -219,6 +219,9 @@ def test_reduce_with_data(self):
         assert result["Num save task success"] == pytest.approx(2.0, rel=0.01)
         assert result["Num save task failed"] == pytest.approx(0.0, rel=0.01)
         assert result["Avg save duration (ms)"] == pytest.approx(2000.0, rel=0.01)
+        # Reduced values must be plain Python scalars so CLI logging renders
+        # them without numpy reprs (eg np.float64(...)).
+        assert all(not isinstance(v, np.generic) for v in result.values())
 
     def test_clone_and_reset(self, hf3fs_stats):
         """clone_and_reset() returns a copy with data and resets the original."""
diff --git a/tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py b/tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py
index f922c7c85586..a366f95caa28 100644
--- a/tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py
+++ b/tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py
@@ -1,10 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for HF3FS metadata server data structures and allocation logic:
-  - RankFileMetadata : page allocation / release primitives
-  - KeyMetadata      : per-key rank-page tracking and completion detection
-  - GlobalMetadataState : coordinated allocation with cache-hit semantics
+"""Tests for HF3FS metadata server data structures and allocation logic:
+- RankFileMetadata : page allocation / release primitives
+- KeyMetadata      : per-key rank-page tracking and completion detection
+- GlobalMetadataState : coordinated allocation with cache-hit semantics
 """
 
 import pytest
diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py
index 8427833b121a..735b042259b5 100644
--- a/tests/v1/kv_connector/unit/test_hidden_states_connector.py
+++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py
@@ -169,3 +169,62 @@ def test_hidden_state_group_isolated_from_packed_mixed_page_groups():
     assert (
         ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == hidden_group_ids[0]
     )
+
+
+# ---- abort-path robustness --------------------------------------------------
+
+
+def _bare_connector() -> ExampleHiddenStatesConnector:
+    """Instance with scheduler-side state but bypassing ``__init__`` (no engine)."""
+    conn = ExampleHiddenStatesConnector.__new__(ExampleHiddenStatesConnector)
+    conn._request_filenames = {}
+    conn._pending_saves = {}
+    conn._lock_fds = {}
+    conn._cache_kv_group_id = 1
+    conn._connector_metadata = None
+    conn._req_copy_events = {}
+    conn._accumulated_finished_req_ids = set()
+    return conn
+
+
+def test_get_finished_count_is_one():
+    # Only TP rank 0 writes, so KVOutputAggregator must expect a single
+    # finished_sending notification per request (not the TP world size).
+    assert _bare_connector().get_finished_count() == 1
+
+
+def test_request_finished_is_noop_for_never_scheduled_request():
+    # A request aborted while still queued never reaches build_connector_meta,
+    # so no filename was recorded. request_finished must not raise KeyError.
+    conn = _bare_connector()
+    request = SimpleNamespace(request_id="cmpl-aborted", kv_transfer_params=None)
+    assert conn.request_finished(request, []) == (False, None)
+
+
+def test_request_finished_all_groups_handles_missing_group():
+    # Guard against indexing a nonexistent per-group block table.
+    conn = _bare_connector()
+    conn._cache_kv_group_id = 2
+    request = SimpleNamespace(request_id="cmpl-aborted", kv_transfer_params=None)
+    assert conn.request_finished_all_groups(request, ([], [])) == (False, None)
+
+
+def test_get_finished_does_not_report_untracked_request():
+    # A never-scheduled aborted request has no copy event. get_finished must not
+    # report it as done_sending, or the scheduler asserts it is still tracked.
+    conn = _bare_connector()
+    assert conn.get_finished({"cmpl-aborted"}) == (None, None)
+    assert conn._accumulated_finished_req_ids == set()
+
+
+def test_get_finished_reports_tracked_completed_request():
+    conn = _bare_connector()
+
+    class _DoneEvent:
+        def query(self) -> bool:
+            return True
+
+    conn._req_copy_events["cmpl-done"] = _DoneEvent()
+    done_sending, done_recving = conn.get_finished({"cmpl-done"})
+    assert done_sending == {"cmpl-done"}
+    assert done_recving is None
diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py
new file mode 100644
index 000000000000..6c4e79efbe08
--- /dev/null
+++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py
@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (
+    ExampleConnector,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
+    HiSparseConnector,
+    HiSparseConnectorMetadata,
+    HiSparseConnectorScheduler,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
+    HiSparseConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
+from vllm.v1.hisparse.runtime import HiSparseCacheHandle
+from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
+
+
+@pytest.mark.parametrize("nested", [False, True])
+def test_cache_manager_binding_preserves_hisparse_and_legacy_pool_hooks(nested):
+    """Composite binding must reach HiSparse and existing pool-only connectors."""
+    hisparse = object.__new__(HiSparseConnector)
+    hisparse._role = KVConnectorRole.SCHEDULER
+    hisparse.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False)
+    legacy = object.__new__(ExampleConnector)
+    legacy.bind_gpu_block_pool = MagicMock()
+    connector = object.__new__(MultiConnector)
+    connector._connectors = [hisparse, legacy]
+    if nested:
+        parent = object.__new__(MultiConnector)
+        parent._connectors = [connector]
+        connector = parent
+    from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager
+
+    manager = make_hisparse_kv_cache_manager(16, 16)
+
+    connector.bind_kv_cache_manager(manager)
+
+    assert hisparse.connector_scheduler.coordinator is get_hisparse_coordinator(manager)
+    legacy.bind_gpu_block_pool.assert_called_once_with(manager.block_pool)
+
+
+def test_hisparse_requires_block_outermost_device_layout():
+    assert HiSparseConnector.get_required_kvcache_layout(MagicMock()) == "BLHNC"
+
+
+def test_no_forward_enqueues_deferred_hisparse_transfers():
+    """A zero-token step must still enqueue deferred post-forward transfers."""
+    connector = object.__new__(ActiveKVConnector)
+    connector._disabled = False
+    connector.pre_forward = MagicMock()
+    connector.finish_forward = MagicMock()
+    connector.post_forward = MagicMock(return_value=None)
+
+    scheduler_output = SimpleNamespace(finished_req_ids=set())
+    connector.no_forward(scheduler_output)
+
+    connector.pre_forward.assert_called_once_with(scheduler_output)
+    connector.finish_forward.assert_called_once_with()
+
+
+def test_full_graph_step_prepares_host_mirror_outside_model():
+    """Graph replay must restore host-mirror state cleared at step start."""
+    runtime = SimpleNamespace(
+        is_group_leader=True,
+        eager_host_mirror=True,
+        begin_forward=MagicMock(),
+        invalidate_written_slots=MagicMock(),
+    )
+    handle = HiSparseCacheHandle(runtime)
+    handle.mirror_slot_mapping = torch.tensor([4, 5])
+    worker = object.__new__(HiSparseConnectorWorker)
+    worker.cache_layer_names = ["layer"]
+    worker.cache_handles = [handle]
+    worker._group_leaders = (("layer", handle),)
+    worker._per_layer_mirrored = set()
+    worker._submitted_mirror_layers = set()
+    worker.is_host_writer = True
+    worker._enqueue_row_dma = MagicMock()
+    worker.start_step = MagicMock(
+        side_effect=lambda *_args, **_kwargs: worker._clear_forward_mirror_state()
+    )
+
+    connector = object.__new__(HiSparseConnector)
+    connector.connector_worker = worker
+    connector._get_connector_metadata = MagicMock(
+        return_value=HiSparseConnectorMetadata(None, (), (), {}, True)
+    )
+    req_id_per_token = torch.tensor([0, 1], dtype=torch.int32)
+    attn_metadata = SimpleNamespace(
+        num_actual_tokens=2,
+        num_decode_tokens=2,
+        num_reqs=2,
+        max_query_len=1,
+        req_id_per_token=req_id_per_token,
+    )
+
+    connector.start_load_kv(
+        SimpleNamespace(),
+        request_state_indices=None,
+        request_ids=[],
+        attn_metadata={"layer": attn_metadata},
+    )
+
+    worker._enqueue_host_mirror()
+
+    worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None)
+    runtime.invalidate_written_slots.assert_called_once()
diff --git a/tests/v1/kv_connector/unit/test_hisparse_stats.py b/tests/v1/kv_connector/unit/test_hisparse_stats.py
new file mode 100644
index 000000000000..09904d0a3caf
--- /dev/null
+++ b/tests/v1/kv_connector/unit/test_hisparse_stats.py
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
+    HiSparseConnector,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import (
+    HiSparseKVConnectorStats,
+)
+
+
+def test_is_empty_on_fresh_stats():
+    stats = HiSparseKVConnectorStats()
+    assert stats.is_empty()
+
+
+def test_record_snapshot_and_reduce():
+    stats = HiSparseKVConnectorStats()
+    stats.record_snapshot(hits=7, misses=3, host_to_device_bytes=48)
+    stats.record_snapshot(hits=5, misses=1, host_to_device_bytes=16)
+    assert not stats.is_empty()
+
+    reduced = stats.reduce()
+    assert reduced["HiSparse hot-buffer hits"] == 12
+    assert reduced["HiSparse hot-buffer misses"] == 4
+    assert reduced["HiSparse host-to-device bytes"] == 64
+
+
+def test_aggregate_extends_snapshot_deltas():
+    first = HiSparseKVConnectorStats()
+    first.record_snapshot(hits=7, misses=3, host_to_device_bytes=48)
+    second = HiSparseKVConnectorStats()
+    second.record_snapshot(hits=5, misses=1, host_to_device_bytes=16)
+
+    first.aggregate(second)
+
+    assert first.data == {
+        "cache_hits": [7, 5],
+        "cache_misses": [3, 1],
+        "host_to_device_bytes": [48, 16],
+    }
+
+
+def test_aggregate_skips_empty_stats():
+    stats = HiSparseKVConnectorStats()
+    stats.record_snapshot(hits=2, misses=1, host_to_device_bytes=16)
+
+    stats.aggregate(HiSparseKVConnectorStats())
+
+    assert stats.data["cache_hits"] == [2]
+
+
+def test_build_kv_connector_stats_round_trip():
+    stats = HiSparseKVConnectorStats()
+    stats.record_snapshot(hits=12, misses=4, host_to_device_bytes=64)
+    payload = stats.to_dict()
+
+    rebuilt = HiSparseConnector.build_kv_connector_stats(data=payload)
+
+    assert rebuilt is not None
+    assert isinstance(rebuilt, HiSparseKVConnectorStats)
+    assert rebuilt.reduce() == stats.reduce()
diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
index 77d629729776..41fce6e9ca0f 100644
--- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
+++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Tests for correctness in invalid block handling.
+"""Tests for correctness in invalid block handling.
 
 These tests verify correct behavior in three scenarios:
 1. Sync recompute case: Blocks should not be freed for running requests
@@ -42,7 +41,7 @@ def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
 
 @pytest.fixture
 def fail_scheduler():
-    """scheduler with kv_load_failure_policy='fail'"""
+    """Scheduler with kv_load_failure_policy='fail'."""
     vllm_config = create_vllm_config()
     vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
     return create_scheduler(vllm_config)
@@ -50,7 +49,7 @@ def fail_scheduler():
 
 @pytest.fixture
 def recompute_scheduler():
-    """scheduler with kv_load_failure_policy='recompute'"""
+    """Scheduler with kv_load_failure_policy='recompute'."""
     vllm_config = create_vllm_config()
     vllm_config.kv_transfer_config.kv_load_failure_policy = "recompute"
     return create_scheduler(vllm_config)
@@ -59,8 +58,7 @@ def recompute_scheduler():
 def test_sync_recompute_blocks_not_freed_for_running_requests(
     recompute_scheduler: Scheduler,
 ):
-    """
-    Test sync recompute case - blocks must not be freed for running requests.
+    """Test sync recompute case - blocks must not be freed for running requests.
 
     When a running request has invalid blocks and retry_policy is 'recompute':
     1. Request should remain in RUNNING state
@@ -180,8 +178,7 @@ def test_sync_recompute_blocks_not_freed_for_running_requests(
 
 
 def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler):
-    """
-    Test sync fail case - invalid blocks must be evicted from cache.
+    """Test sync fail case - invalid blocks must be evicted from cache.
 
     When a request fails with policy='fail' and has invalid blocks from sync loading:
     1. Request should be finished with FINISHED_ERROR
@@ -296,8 +293,7 @@ def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler):
 def test_async_recompute_blocks_not_cached_when_invalid(
     recompute_scheduler: Scheduler,
 ):
-    """
-    Test async recompute case - invalid blocks not cached after transfer.
+    """Test async recompute case - invalid blocks not cached after transfer.
 
     When async KV loading has invalid blocks and retry_policy is 'recompute':
     1. Blocks are allocated but not cached yet
diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
index 4f35527b0e3f..a8bc2d0e12f0 100644
--- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
+++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
@@ -14,6 +14,7 @@
     create_request,
     create_scheduler,
     create_vllm_config,
+    make_kv_cache_config,
 )
 
 
@@ -34,6 +35,39 @@ def scheduler():
     return create_scheduler(vllm_config)
 
 
+@pytest.mark.parametrize("policy", ["fail", "recompute"])
+def test_failed_receive_completion_honors_failure_policy(policy):
+    """Failed receives still need completion to free blocks or resume computation."""
+    scheduler = create_scheduler(create_vllm_config(kv_load_failure_policy=policy))
+    request = create_request(num_tokens=3 * scheduler.block_size)
+    scheduler.add_request(request)
+    scheduler.connector = Mock()
+    scheduler.connector.get_num_new_matched_tokens.side_effect = [
+        (2 * scheduler.block_size, True),
+        (0, False),
+    ]
+    scheduler.connector.request_finished.return_value = (False, None)
+    scheduler.connector.take_events.return_value = ()
+    block_pool = scheduler.kv_cache_manager.block_pool
+    free_blocks_before = block_pool.get_num_free_blocks()
+    scheduler_output = scheduler.schedule()
+    assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+
+    output = create_model_runner_output([], finished_recving={request.request_id})
+    assert output.kv_connector_output is not None
+    output.kv_connector_output.failed_recving = {request.request_id}
+    scheduler.update_from_output(scheduler_output, output)
+
+    if policy == "fail":
+        assert request.status == RequestStatus.FINISHED_ERROR
+        assert request.request_id not in scheduler.requests
+        assert block_pool.get_num_free_blocks() == free_blocks_before
+    else:
+        assert request.num_computed_tokens == 0
+        next_output = scheduler.schedule()
+        assert request.request_id in next_output.num_scheduled_tokens
+
+
 @pytest.mark.parametrize(
     "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs",
     [
@@ -337,3 +371,21 @@ def test_async_progressive_load_failure(
         assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
         assert scheduler.failed_recving_kv_req_ids == {request.request_id}
         assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
+
+
+@pytest.mark.cpu_test
+def test_block_level_load_failure_fails_closed_with_multiple_groups():
+    """Block-level invalid-block reporting is rejected on multi-group layouts.
+
+    Block IDs are only unique within a group, so the scheduler cannot map
+    them back to requests; connectors must report failed_recving instead.
+    """
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(
+        vllm_config,
+        kv_cache_config=make_kv_cache_config(block_size=16, swa_enabled=True),
+    )
+    assert len(scheduler.kv_cache_config.kv_cache_groups) == 2
+
+    with pytest.raises(RuntimeError, match="failed_recving"):
+        scheduler._handle_invalid_blocks({1}, {})
diff --git a/tests/v1/kv_connector/unit/test_lmcache_connector.py b/tests/v1/kv_connector/unit/test_lmcache_connector.py
index c3df2b68b1ff..a64f7f32d5a6 100644
--- a/tests/v1/kv_connector/unit/test_lmcache_connector.py
+++ b/tests/v1/kv_connector/unit/test_lmcache_connector.py
@@ -675,8 +675,7 @@ def __init__(self, cycle_num):
             assert mock_connector._kv_cache_events is None
 
     def test_lmcache_kv_events_aggregation(self):
-        """
-        Test LMCacheKVEvents aggregation across TP ranks using
+        """Test LMCacheKVEvents aggregation across TP ranks using
         KVOutputAggregator (used by MultiprocExecutor).
         """
         from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py
index f989fb4e32a4..db9cf9417165 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_connector.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py
@@ -179,7 +179,6 @@ def batch_register_memory(self, buffer_addresses, capacities) -> int:
 
 def test_align_transfer_regions_uses_layer_name_occurrences():
     """Repeated layer names should align by occurrence order."""
-
     local_regions = [
         TransferRegion(
             layer_name="model.layers.1.self_attn",
@@ -232,7 +231,6 @@ def test_align_transfer_regions_uses_layer_name_occurrences():
 @pytest.mark.asyncio
 async def test_build_transfer_params_separates_prefill_pp_layers():
     """Each producer PP stage should send only its registered layer shard."""
-
     worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
     worker.async_zmq_ctx = MagicMock()
     worker.is_kv_consumer = True
@@ -368,7 +366,6 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
     monkeypatch,
 ):
     """Producer sends its PP layer shard to the matching consumer layer address."""
-
     monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_producer"
@@ -447,7 +444,6 @@ async def run_in_executor(self, executor, func, *args):
 
 def test_basic_interface():
     """Unit test for basic MooncakeConnector interface functionality."""
-
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
@@ -497,7 +493,6 @@ def test_basic_interface():
 
 def test_prompt_less_than_block_size():
     """Test that we can handle case where prompt is < block."""
-
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
@@ -536,7 +531,6 @@ def test_prompt_less_than_block_size():
 @pytest.fixture
 def bootstrap_server():
     """Fixture to launch and cleanup a Mooncake Bootstrap HTTP Server."""
-
     port = get_open_port()
     server = MooncakeBootstrapServer("127.0.0.1", port)
     server.start()
@@ -546,12 +540,10 @@ def bootstrap_server():
 
 @pytest.mark.asyncio
 async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
-    """
-    Tests the bootstrap server's api for worker registration and querying.
+    """Tests the bootstrap server's api for worker registration and querying.
 
     Validates DP/TP/PP rank indexing and error handling for duplicate registrations.
     """
-
     import httpx
 
     base_url = f"http://127.0.0.1:{bootstrap_server.port}"
@@ -716,13 +708,11 @@ def test_get_mooncake_bootstrap_addr_selects_expected_host(
 
 
 def test_scheduler_request_finished():
-    """
-    Tests the scheduler-side logic when a request finishes.
+    """Tests the scheduler-side logic when a request finishes.
 
     Differentiates between 'Finished' (requires transfer)
     and 'Aborted' (immediate free).
     """
-
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_producer"
     )
@@ -751,7 +741,6 @@ def test_scheduler_request_finished():
 @contextlib.contextmanager
 def patch_worker_dependencies():
     """Helper to mock all distributed and network dependencies for Worker tests."""
-
     with (
         patch(
             "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
@@ -820,7 +809,6 @@ async def test_receive_kv_selects_remote_pp_workers(
     expected_addrs: list[str],
 ):
     """Decode workers should not hard-code producer pp_rank 0."""
-
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
@@ -875,7 +863,6 @@ async def fake_receive(worker_addr: str, metas: dict[str, PullReqMeta]):
 
 def test_resolve_need_send_accounts_for_remote_tp_fanout():
     """Producer-side completion waits for every paired consumer TP pull."""
-
     worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
     worker.async_zmq_ctx = MagicMock()
     worker.is_kv_consumer = True
@@ -892,19 +879,79 @@ def test_resolve_need_send_accounts_for_remote_tp_fanout():
     assert send_meta.need_send == 2
 
 
+@pytest.mark.asyncio
+async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers():
+    """P4/D2 must retain KV until both D stages finish, including a no-op pull."""
+    config = create_vllm_config(kv_connector="MooncakeConnector", kv_role="kv_producer")
+    with (
+        set_current_vllm_config(config),
+        patch_worker_dependencies(),
+        patch.object(MooncakeConnectorWorker, "_sync_block_size_with_kernel"),
+    ):
+        worker = MooncakeConnector(
+            config, KVConnectorRole.WORKER, _make_test_kv_cache_config()
+        ).connector_worker
+        try:
+            worker.pp_size = 4
+            worker.kv_caches_base_addr = [0x1000]
+            worker.block_len_per_layer = [256]
+            worker.kv_block_len_per_layer = [256]
+            worker.registered_layer_names = ["model.layers.0.self_attn"]
+            worker.registered_layer_indices = [0]
+            send_meta = SendBlockMeta(
+                p_req_id="p-req",
+                transfer_id="transfer",
+                local_block_ids=[[1]],
+                ready=asyncio.Event(),
+            )
+            send_meta.ready.set()
+            worker.reqs_need_send["transfer"] = send_meta
+            sock = AsyncMock(spec=zmq.asyncio.Socket, send_multipart=AsyncMock())
+            with (
+                patch.object(worker, "sender_loop", asyncio.get_running_loop()),
+                patch.object(worker, "_send_blocks", return_value=0) as send,
+            ):
+                for pp_rank in range(2):
+                    metadata = MooncakeXferMetadata(
+                        remote_hostname="consumer",
+                        remote_port=1234 + pp_rank,
+                        remote_tp_size=1,
+                        remote_tp_rank=0,
+                        remote_pp_size=2,
+                        req_blocks={"d-req": ("transfer", [[2]])},
+                        kv_caches_base_addr=[0x2000],
+                        block_lens=[256],
+                        kv_block_lens=[256],
+                        registered_layer_names=[f"model.layers.{pp_rank}.self_attn"],
+                        registered_layer_indices=[pp_rank],
+                    )
+                    await worker.send_kv_to_decode(b"consumer", sock, metadata)
+                    response = worker._xfer_resp_decoder.decode(
+                        sock.send_multipart.call_args.args[0][1]
+                    )
+                    assert response.status == MooncakeXferResponseStatus.FINISH
+                    assert response.ok_reqs == ["d-req"]
+                    assert ("transfer" in worker.reqs_need_send) == (pp_rank == 0)
+                    assert worker.finished_sending_reqs == (
+                        set() if pp_rank == 0 else {"p-req"}
+                    )
+                send.assert_called_once_with("consumer:1234", [0x1100], [0x2200], [256])
+        finally:
+            worker.shutdown()
+            worker.is_kv_consumer = True
+
+
 @pytest.mark.asyncio
 @patch(
     "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
     FakeMooncakeWrapper,
 )
 async def test_kv_producer(monkeypatch):
-    """
-    Simulates a Producer Worker (Prefiller) receiving a transfer request
+    """Simulates a Producer Worker (Prefiller) receiving a transfer request
     from a Consumer (Decoder).
 
     Verifies memory offset calculation: ptr = base_addr + block_id * block_len.
     """
-
     monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_producer"
@@ -1071,12 +1118,10 @@ def expected_transfers(src_base, dst_base, src_blocks, dst_blocks):
 
 @pytest.mark.asyncio
 async def test_kv_consumuer(monkeypatch):
-    """
-    Simulates a Consumer Worker (Decoder) initiating a pull from a Producer.
+    """Simulates a Consumer Worker (Decoder) initiating a pull from a Producer.
 
     Verifies that MooncakeXferMetadata is correctly serialized and sent via ZMQ.
     """
-
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
@@ -1149,7 +1194,6 @@ async def test_kv_consumuer(monkeypatch):
 @pytest.mark.asyncio
 async def test_worker_get_finished_timeout(monkeypatch):
     """Tests the cleanup mechanism for requests."""
-
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_producer"
     )
@@ -1201,7 +1245,6 @@ async def test_worker_get_finished_timeout(monkeypatch):
 )
 def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool):
     """Tests the memory registration logic with the underlying Mooncake engine."""
-
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
@@ -1282,7 +1325,6 @@ def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool
 
 def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
     """Mixed MLA+Eagle caches should register by byte length, not shape."""
-
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
@@ -1346,8 +1388,7 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
 )
 @pytest.mark.parametrize("d_tp_size", [1, 4], ids=["p_tp2_d_tp1", "p_tp2_d_tp4"])
 async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
-    """
-    Tests heterogeneous TP support in the producer transfer path.
+    """Tests heterogeneous TP support in the producer transfer path.
 
     Verifies correct pointer and offset calculation when producer TP=2
     sends to consumer with TP=1 (P>D) or TP=4 (P<D).
@@ -1356,7 +1397,6 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
     - P TP=2 > D TP=1: one D rank receives; dst_offset based on P rank
     - P TP=2 < D TP=4: two D ranks receive; src_offset based on D rank
     """
-
     P_TP_SIZE = 2
     P_TP_RANK = 0
     LOCAL_BLOCK_LEN = 4096
diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
index 54b97d482eff..a7fb2be02365 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
@@ -10,6 +10,7 @@
 from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm.config import set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
@@ -18,9 +19,15 @@
     MooncakeConnectorMetadata,
     MooncakeConnectorScheduler,
     MooncakeXferMetadata,
+    PullReqMeta,
     SendBlockMeta,
     TransferRegion,
 )
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+)
 
 from .test_mooncake_connector import FakeMooncakeWrapper, patch_worker_dependencies
 from .utils import create_request, create_vllm_config, make_kv_cache_config
@@ -216,7 +223,6 @@ def test_metadata_hma_block_ids():
 )
 async def test_build_transfer_params_multi_group_trimming(monkeypatch):
     """_build_transfer_params trims per-group blocks when local > remote."""
-
     monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_producer"
@@ -313,7 +319,6 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch):
 )
 async def test_build_transfer_params_group_count_mismatch(monkeypatch):
     """_build_transfer_params reports an error when group counts differ."""
-
     monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_producer"
@@ -436,3 +441,105 @@ def test_request_finished_with_hma_groups():
     assert stored_blocks[0] == fa_blocks
     # SW: clipped to last 9 blocks (sw_size=128, block_size=16 → 8+1=9)
     assert stored_blocks[1] == sw_blocks[-9:]
+
+
+# ---------------------------------------------------------------------------
+#  Worker-side load-failure reporting (HMA vs non-HMA)
+# ---------------------------------------------------------------------------
+def _make_kv_consumer_worker(
+    swa_enabled: bool = False, kv_cache_config: KVCacheConfig | None = None
+):
+    block_size = 16
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector",
+        kv_role="kv_consumer",
+        block_size=block_size,
+    )
+    if kv_cache_config is None:
+        kv_cache_config = make_kv_cache_config(
+            block_size=block_size, swa_enabled=swa_enabled
+        )
+    with set_current_vllm_config(vllm_config), patch_worker_dependencies():
+        connector = MooncakeConnector(
+            vllm_config, KVConnectorRole.WORKER, kv_cache_config
+        )
+    return connector.connector_worker
+
+
+def _make_pull_meta(d_req_id: str, local_block_ids: list[list[int]]) -> PullReqMeta:
+    return PullReqMeta(
+        d_req_id=d_req_id,
+        transfer_id="tx-1",
+        local_block_ids=local_block_ids,
+        remote_engine_id="engine-0",
+        remote_bootstrap_addr="127.0.0.1:1234",
+    )
+
+
+# The tests below construct a MooncakeConnectorWorker, whose __init__
+# requires an accelerator device, so (like the worker tests in
+# test_mooncake_connector.py) they run on the GPU KV-connectors job only.
+@pytest.mark.parametrize(
+    "swa_enabled,expected_is_hma",
+    [
+        (True, True),  # multiple groups (FA + SWA)
+        (False, False),  # single group, block-level reporting is unambiguous
+    ],
+)
+def test_worker_is_hma_required(swa_enabled, expected_is_hma):
+    """Worker-side _is_hma_required mirrors the scheduler's condition."""
+    worker = _make_kv_consumer_worker(swa_enabled=swa_enabled)
+    assert worker._is_hma_required is expected_is_hma
+
+
+def test_worker_is_hma_required_multiple_full_attention_groups():
+    """Two full-attention groups with different block sizes also select
+    request-level failure reporting, matching the scheduler."""
+    kv_cache_config = KVCacheConfig(
+        num_blocks=100,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["layer0"],
+                FullAttentionSpec(
+                    block_size=16, num_kv_heads=4, head_size=16, dtype=torch.float16
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["layer1"],
+                FullAttentionSpec(
+                    block_size=32, num_kv_heads=4, head_size=16, dtype=torch.float16
+                ),
+            ),
+        ],
+    )
+    worker = _make_kv_consumer_worker(kv_cache_config=kv_cache_config)
+    assert worker._is_hma_required
+
+
+def test_worker_failed_recv_reports_request_level_failure_with_hma():
+    """With HMA, load failures report the request, not ambiguous block IDs."""
+    worker = _make_kv_consumer_worker(swa_enabled=True)
+    assert worker._is_hma_required
+
+    pull_meta = _make_pull_meta("d-req-1", [[1, 2], [3, 4]])
+    worker._handle_failed_recv({"p-req-1": pull_meta}, ["p-req-1"], "boom")
+
+    assert worker.get_block_ids_with_load_errors() == set()
+    results = worker.get_transfer_results()
+    assert results.failed_recving == {"d-req-1"}
+    assert results.finished_recving == {"d-req-1"}
+
+
+def test_worker_failed_recv_reports_block_ids_without_hma():
+    """Without HMA, load failures keep reporting block-level errors."""
+    worker = _make_kv_consumer_worker(swa_enabled=False)
+    assert not worker._is_hma_required
+
+    pull_meta = _make_pull_meta("d-req-1", [[1, 2], [3, 4]])
+    worker._handle_failed_recv({"p-req-1": pull_meta}, ["p-req-1"], "boom")
+
+    assert worker.get_block_ids_with_load_errors() == {1, 2, 3, 4}
+    results = worker.get_transfer_results()
+    assert results.failed_recving == set()
+    assert results.finished_recving == {"d-req-1"}
diff --git a/tests/v1/kv_connector/unit/test_mooncake_stats.py b/tests/v1/kv_connector/unit/test_mooncake_stats.py
index a20fcb505330..d52ab5f201a3 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_stats.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_stats.py
@@ -4,6 +4,8 @@
 import threading
 from unittest.mock import MagicMock
 
+import numpy as np
+
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
     MooncakeConnector,
     MooncakeConnectorWorker,
@@ -40,6 +42,9 @@ def test_record_transfer_and_reduce():
     assert reduced["Num failed transfers"] == 0
     assert reduced["Num failed recvs"] == 0
     assert reduced["Num KV expired reqs"] == 0
+    # Reduced values must be plain Python scalars so CLI logging renders
+    # them without numpy reprs (eg np.float64(...)).
+    assert all(not isinstance(v, np.generic) for v in reduced.values())
 
 
 def test_record_failures_keeps_stats_non_empty():
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py
index 96315b026df1..f3ffa55a5439 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py
@@ -12,6 +12,7 @@
 from vllm.distributed.kv_events import BlockStored
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorRole,
+    KVConnectorTransferResults,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
     connector as mooncake_store_connector,
@@ -30,6 +31,7 @@
     MooncakeStoreConnectorStats,
 )
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -91,6 +93,53 @@ def _make_kv_cache_config() -> KVCacheConfig:
     )
 
 
+def _make_qsa_hybrid_kv_cache_config(mamba_mode: str):
+    full_spec = FullAttentionSpec(
+        block_size=800, num_kv_heads=8, head_size=64, dtype=None
+    )
+    qsa_spec = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=1,
+        head_size=64,
+        head_size_v=0,
+        dtype=torch.float16,
+    )
+    mamba_spec = MambaSpec(
+        block_size=800,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode=mamba_mode,
+    )
+    return KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["full"], full_spec),
+            KVCacheGroupSpec(["qsa"], qsa_spec),
+            KVCacheGroupSpec(["mamba"], mamba_spec),
+        ],
+    )
+
+
+def test_validation_accepts_aligned_mamba_after_block_size_rewrite():
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 8
+
+    mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config(
+        vllm_config, _make_qsa_hybrid_kv_cache_config("align")
+    )
+
+
+def test_validation_rejects_mamba_mode_directly():
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 800
+
+    with pytest.raises(ValueError, match="mamba_cache_mode=\x27none\x27"):
+        mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config(
+            vllm_config, _make_qsa_hybrid_kv_cache_config("none")
+        )
+
+
 def test_scheduler_requires_align_mode_for_mamba():
     vllm_config = _make_vllm_config()
     mamba_spec = MambaSpec(
@@ -197,6 +246,38 @@ def test_worker_methods_delegate_to_store_worker():
     assert invalid_block_ids == {3, 4}
 
 
+def test_get_transfer_results_delegates_to_store_worker():
+    vllm_config = _make_vllm_config()
+    kv_cache_config = _make_kv_cache_config()
+    metadata = MooncakeStoreConnectorMetadata(set(), set())
+    finished_req_ids = {"req-1"}
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ) as mock_worker_cls,
+    ):
+        connector = mooncake_store_connector.MooncakeStoreConnector(
+            vllm_config, KVConnectorRole.WORKER, kv_cache_config
+        )
+
+    worker = mock_worker_cls.return_value
+    transfer_results = KVConnectorTransferResults(
+        finished_sending={"req-1"},
+        finished_recving={"req-2"},
+        failed_recving={"req-3"},
+    )
+    worker.get_transfer_results.return_value = transfer_results
+    connector.bind_connector_metadata(metadata)
+
+    result = connector.get_transfer_results(finished_req_ids)
+
+    worker.get_transfer_results.assert_called_once_with(finished_req_ids, metadata)
+    assert result is transfer_results
+
+
 def test_get_kv_connector_kv_cache_events_returns_none_when_disabled():
     vllm_config = _make_vllm_config()
     kv_cache_config = _make_kv_cache_config()
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
index d6b0becc1d2c..6c2cff0a542c 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
@@ -29,6 +29,7 @@
 )
 from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     KpoolTailSpec,
     KVCacheConfig,
@@ -146,8 +147,7 @@ def _build_worker_with_dict_store(vllm_config, kv_cache_config, store):
 
 
 def test_e2e_swa_plus_full_save_then_lookup_hits():
-    """
-    E2E: build a SWA+Full hybrid worker, save all blocks via the sending
+    """E2E: build a SWA+Full hybrid worker, save all blocks via the sending
     thread (synchronously), then verify lookup returns the full hit length.
     Also verify that evicting SWA's early blocks (outside its window) still
     allows a full hit because the window covers the tail.
@@ -758,8 +758,9 @@ def test_worker_setup_tolerates_finer_scratch_group():
     worker.pp_size = 1
     worker.num_kv_head = 8
     assert worker.coord.enable_partial_hash_hits
-    # The scratch DB is keyed at its own block size and never probed.
-    assert worker.token_dbs[2].hash_block_size == 4
+    # Only the participating full-attention and Mamba groups are registered.
+    assert len(worker.token_dbs) == 2
+    assert all(db.hash_block_size == 8 for db in worker.token_dbs)
 
     for g_idx, db in enumerate(worker.token_dbs):
         db.set_kv_caches_base_addr([g_idx * 10_000])
@@ -771,11 +772,11 @@ def test_worker_setup_tolerates_finer_scratch_group():
         block_size=worker.block_size,
         coord=worker.coord,
         tp_rank=0,
-        group_put_steps=[1, 1, 1],
+        group_put_steps=[1, 1],
         kv_role="kv_both",
         ready_event=threading.Event(),
         replicate_config=MagicMock(),
-        group_participates=[True, True, False],
+        group_participates=[True, True],
     )
 
     # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]).
@@ -783,7 +784,7 @@ def test_worker_setup_tolerates_finer_scratch_group():
     req = ReqMeta(
         req_id="r0",
         token_len_chunk=0,
-        block_ids=([1], [2], [3]),
+        block_ids=([1], [2]),
         block_hashes=hs,
         can_save=True,
         num_prompt_tokens=20,
@@ -796,5 +797,103 @@ def test_worker_setup_tolerates_finer_scratch_group():
     # A 13-token prompt sharing the prefix must hit the first hash unit.
     assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8
     # The scratch group's namespace never enters the store.
-    scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0]
-    assert not any(key.startswith(scratch_prefix) for key in store._data)
+    assert not any("@group:2" in key for key in store._data)
+
+
+def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
+    """A per-request ring group (CircularBufferSpec, capacity as block_size)
+    sits beside the paged group: worker setup tolerates its block size (the
+    DeepSeek-V4.1 compressor ring is 8 rows), saving a request stores paged
+    blocks only, and lookup hits on the paged group alone."""
+    full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
+    ring = CircularBufferSpec(
+        block_size=8, num_kv_heads=1, head_size=64, head_size_v=0, dtype=torch.uint8
+    )
+    cfg = KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=4 * full.page_size_bytes,
+                layers=["L0"],
+                layer_stride=4 * full.page_size_bytes,
+                block_stride=full.page_size_bytes,
+            ),
+            KVCacheTensor(
+                size=4 * ring.page_size_bytes,
+                layers=["L1"],
+                layer_stride=4 * ring.page_size_bytes,
+                block_stride=ring.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["L0"], full),
+            KVCacheGroupSpec(["L1"], ring),
+        ],
+    )
+    vllm_config = _minimal_vllm_config(cache_block_size=16)
+    store = _DictStore()
+
+    worker = _build_worker_with_dict_store(vllm_config, cfg, store)
+    worker.tp_size = 1
+    worker.pp_size = 1
+    worker.num_kv_head = 8
+    assert len(worker.token_dbs) == 1
+    assert worker.token_dbs[0].hash_block_size == 16
+
+    raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8)
+    raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8)
+    kv_caches = {
+        "L0": dense_kv_cache_views(raw_full, full, 4, 1, KVCacheLayout.LBNHC)[0],
+        "L1": dense_kv_cache_views(raw_ring, ring, 4, 1, KVCacheLayout.LBNHC)[0],
+    }
+
+    def _fake_thread_init(*args, **kwargs):
+        for v in list(args) + list(kwargs.values()):
+            if isinstance(v, threading.Event):
+                v.set()
+        m = MagicMock()
+        m.start = lambda: None
+        return m
+
+    with (
+        patch.object(
+            mooncake_store_worker,
+            "KVCacheStoreSendingThread",
+            side_effect=_fake_thread_init,
+        ),
+        patch.object(
+            mooncake_store_worker,
+            "KVCacheStoreRecvingThread",
+            side_effect=_fake_thread_init,
+        ),
+    ):
+        worker.register_kv_caches(kv_caches)
+
+    send_thread = KVCacheStoreSendingThread(
+        store=store,
+        token_databases=worker.token_dbs,
+        block_size=worker.block_size,
+        coord=worker.coord,
+        tp_rank=worker.tp_rank,
+        group_put_steps=worker._group_tp_replication_factors,
+        kv_role=worker.kv_role,
+        ready_event=threading.Event(),
+        enable_kv_event=False,
+        group_participates=[True],
+    )
+    hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
+    save_req = ReqMeta(
+        req_id="r0",
+        token_len_chunk=64,
+        # Scheduler metadata is projected to the participating store groups.
+        block_ids=([1, 2, 3, 4],),
+        block_hashes=hs,
+        can_save=True,
+        store_job_id=1,
+    )
+    send_thread.add_request(save_req)
+    send_thread._handle_request(send_thread.request_queue.get())
+    worker.store = store
+
+    assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64
+    assert not any("@group:1" in key for key in store._data)
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
index 67fa0c88aab4..7ced38d65a41 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
@@ -2,12 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
     LoadSpec,
     MooncakeLookupResult,
+    MooncakeStoreConnectorMetadata,
     MooncakeStoreWorkerMetadata,
     ReqMeta,
     RequestTracker,
@@ -16,7 +19,16 @@
     MooncakeStoreScheduler,
 )
 from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.kv_cache_utils import KVCacheBlock
 from vllm.v1.core.sched.output import KVConnectorBlockState
+from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+)
 
 
 def _make_bare_scheduler(
@@ -35,9 +47,9 @@ def _make_bare_scheduler(
     scheduler._block_size = 16
     scheduler._hash_block_size = hash_block_size
     scheduler.enable_partial_hash_hits = enable_partial_hash_hits
-    scheduler.kv_cache_config = SimpleNamespace(
-        select_transfer_block_ids=lambda block_ids: tuple(block_ids)
-    )
+    scheduler.kv_cache_config = SimpleNamespace()
+    scheduler._store_group_ids = (0,)
+    scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 1: 1}
     scheduler.load_specs = {}
     scheduler._unfinished_request_ids = {"req-0"}
     scheduler._unfinished_requests = {}
@@ -228,14 +240,82 @@ def test_pending_load_for_non_chosen_connector_is_dropped():
     assert "req-0" not in scheduler._request_trackers
 
 
+def _make_qsa_hybrid_cache_config():
+    full = FullAttentionSpec(block_size=800, num_kv_heads=8, head_size=64, dtype=None)
+    circular = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=1,
+        head_size=64,
+        head_size_v=0,
+        dtype=torch.float16,
+    )
+    mamba = MambaSpec(
+        block_size=800,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+    return KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["full"], full),
+            KVCacheGroupSpec(["qsa"], circular),
+            KVCacheGroupSpec(["mamba"], mamba),
+        ],
+    )
+
+
+def test_scheduler_projects_nonprefix_groups_and_mamba_ids():
+    vllm_config = SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(
+            kv_role="kv_both", kv_connector_extra_config={}
+        ),
+        kv_events_config=None,
+        cache_config=SimpleNamespace(
+            block_size=800, enable_prefix_caching=True, prefix_match_unit=None
+        ),
+        parallel_config=SimpleNamespace(decode_context_parallel_size=1, world_size=1),
+    )
+
+    with patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+        "scheduler.LookupKeyClient"
+    ):
+        scheduler = MooncakeStoreScheduler(vllm_config, _make_qsa_hybrid_cache_config())
+
+    assert scheduler._store_group_ids == (0, 2)
+    assert scheduler._boundary_state_group_ids == frozenset({1})
+
+
+def test_current_save_block_ids_use_store_group_projection():
+    scheduler = _make_bare_scheduler()
+    scheduler.kv_cache_config = _make_qsa_hybrid_cache_config()
+    scheduler._store_group_ids = (0, 2)
+    meta = MooncakeStoreConnectorMetadata(set(), set())
+    req_meta = ReqMeta(
+        req_id="req-0",
+        token_len_chunk=800,
+        block_ids=(),
+        block_hashes=[b"h0"],
+        can_save=True,
+    )
+    meta.add_request(req_meta)
+    output = SimpleNamespace(
+        kv_connector_block_state=_make_connector_block_state(([10], [80], [30]))
+    )
+
+    scheduler._apply_current_save_block_ids(meta, output)
+
+    assert req_meta.block_ids == ([10], [30])
+
+
 def test_update_state_excludes_nontransfer_groups():
     """Store metadata must match the worker's registered cache groups."""
     scheduler = _make_bare_scheduler()
-    scheduler.kv_cache_config = SimpleNamespace(
-        select_transfer_block_ids=lambda block_ids: (block_ids[0],)
-    )
+    scheduler._store_group_ids = (0,)
     request = SimpleNamespace(request_id="req-1")
-    blocks = SimpleNamespace(get_block_ids=lambda: ([1, 2], [9]))
+    blocks = KVCacheBlocks(([KVCacheBlock(1), KVCacheBlock(2)], [KVCacheBlock(9)]))
 
     scheduler.update_state_after_alloc(request, blocks, num_external_tokens=32)
 
@@ -1003,6 +1083,7 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():
 
 def test_finished_partial_tail_is_pre_pinned_as_store_job():
     scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
+    scheduler._store_group_ids = (0, 1)
     scheduler.client = SimpleNamespace(discard=lambda *_: None)
     request = SimpleNamespace(
         request_id="req-0",
@@ -1143,6 +1224,20 @@ def _make_offload_only_output(entries, block_ids=([0],)):
     )
 
 
+def test_boundary_state_group_ids_are_remapped_to_store_projection():
+    scheduler = _make_bare_scheduler(hash_block_size=800, enable_partial_hash_hits=True)
+    scheduler.kv_cache_config = _make_qsa_hybrid_cache_config()
+    scheduler._store_group_ids = (0, 2)
+    scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 2: 1}
+    scheduler._boundary_state_group_ids = frozenset({1})
+    _register_offload_request(scheduler, prefill_end_tokens=800, num_prompt_tokens=800)
+    meta = MooncakeStoreConnectorMetadata(set(), set())
+
+    scheduler._handle_boundary_state_offloads({"req-0": [(2, 7, 800)]}, meta)
+
+    assert meta.requests[0].boundary_state_offloads == [(1, 7, 800)]
+
+
 def test_resumed_prefill_claims_boundaries_past_prompt_length():
     # A resumed request re-prefills its previously generated tokens, so its
     # save window (`prefill_end_tokens`) extends past `num_prompt_tokens`.
@@ -1195,6 +1290,7 @@ def test_boundary_state_job_pins_exact_blocks_once():
 
 def test_store_job_pins_current_non_null_non_mamba_blocks():
     scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
+    scheduler._store_group_ids = (0, 1)
     request = SimpleNamespace(
         all_token_ids=list(range(48)),
         block_hashes=[bytes([i]) for i in range(12)],
diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
index 52df7531e715..947e9f6a20bf 100644
--- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py
@@ -17,6 +17,7 @@
 import torch
 
 from tests.v1.attention.utils import dense_kv_cache_views
+from vllm.distributed import mooncake_store
 from vllm.distributed.kv_events import KVEventAggregator
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
     rdma_utils,
@@ -46,8 +47,11 @@
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
+    KVCacheConfig,
     KVCacheGroupSpec,
+    MambaSpec,
     RSWASpec,
 )
 from vllm.v1.kv_cache_layout import KVCacheLayout
@@ -174,6 +178,7 @@ def _make_store_recving_thread(
     *,
     tp_rank: int = 0,
     disk_offload_buffer_budget_bytes: int | None = None,
+    is_hma_required: bool = False,
 ) -> mooncake_store_worker.KVCacheStoreRecvingThread:
     from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
 
@@ -196,6 +201,7 @@ def _make_store_recving_thread(
         ready_event=threading.Event(),
         coord=coord,
         disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes,
+        is_hma_required=is_hma_required,
     )
     thread.request_queue.task_done = MagicMock()
     return thread
@@ -300,6 +306,7 @@ def _make_vllm_config(
     kv_role: str = "kv_both",
     pipeline_parallel_size: int = 1,
     kv_cache_layout: KVCacheLayout = KVCacheLayout.LBHNC,
+    disable_hybrid_kv_cache_manager: bool = True,
 ) -> SimpleNamespace:
     cache_config = SimpleNamespace(block_size=16, num_gpu_blocks=10)
     cache_config.get_resolved_kv_cache_layout = lambda: kv_cache_layout
@@ -316,6 +323,9 @@ def _make_vllm_config(
             kv_role=kv_role, extra_config=extra_config
         ),
         cache_config=cache_config,
+        scheduler_config=SimpleNamespace(
+            disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager
+        ),
         kv_events_config=SimpleNamespace(enable_kv_cache_events=False),
         speculative_config=None,
     )
@@ -342,6 +352,43 @@ def _make_kv_cache_config(
     )
 
 
+def _make_qsa_hybrid_kv_cache_config():
+    return KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=800,
+                    num_kv_heads=8,
+                    head_size=64,
+                    dtype=None,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["qsa"],
+                CircularBufferSpec(
+                    block_size=8,
+                    num_kv_heads=1,
+                    head_size=64,
+                    head_size_v=0,
+                    dtype=torch.float16,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=800,
+                    shapes=((1, 1),),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+
+
 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")
@@ -356,6 +403,7 @@ def __init__(self) -> None:
     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_store_module.ObjectDataType = SimpleNamespace(TENSOR=1)  # 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)
@@ -544,7 +592,7 @@ def test_pool_key_cache_prefix_namespaces_and_disambiguates():
 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
+    assert mooncake_store.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3
 
 
 def test_get_requester_local_hostname_prefers_override(monkeypatch):
@@ -2107,6 +2155,127 @@ def test_recv_thread_reports_unsplittable_key_larger_than_budget():
     assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}
 
 
+def test_recv_thread_partial_load_failure_reports_block_ids_without_hma():
+    store = MagicMock()
+    # First key fails, second succeeds.
+    store.batch_get_into_multi_buffers.return_value = [-1, 0]
+    thread = _make_store_recving_thread(store)
+
+    thread._handle_request(_make_load_req("req-a", [b"h0", b"h1"], token_len=32))
+
+    assert thread.get_and_clear_block_ids_with_load_errors() == {0}
+    assert thread.get_and_clear_failed_requests() == set()
+    assert thread.get_and_clear_finished_requests() == {"req-a"}
+
+
+def test_recv_thread_partial_load_failure_reports_request_with_hma():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [-1, 0]
+    thread = _make_store_recving_thread(store, is_hma_required=True)
+
+    thread._handle_request(_make_load_req("req-a", [b"h0", b"h1"], token_len=32))
+
+    assert thread.get_and_clear_block_ids_with_load_errors() == set()
+    assert thread.get_and_clear_failed_requests() == {"req-a"}
+    assert thread.get_and_clear_finished_requests() == {"req-a"}
+
+
+def test_recv_thread_oversized_key_reports_request_with_hma():
+    store = MagicMock()
+    thread = _make_store_recving_thread(
+        store,
+        tp_rank=2,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_TOO_SMALL,
+        is_hma_required=True,
+    )
+
+    thread._handle_request(_make_load_req("req-a", [b"a0", b"a1", b"a2"], token_len=48))
+
+    assert store.batch_get_into_multi_buffers.call_count == 0
+    assert thread.get_and_clear_block_ids_with_load_errors() == set()
+    assert thread.get_and_clear_failed_requests() == {"req-a"}
+    assert thread.get_and_clear_finished_requests() == {"req-a"}
+
+
+def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setattr(
+        worker.MooncakeStoreConfig,
+        "load_from_config",
+        staticmethod(lambda: _make_config()),
+    )
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 800
+    vllm_config.cache_config.enable_prefix_caching = True
+    vllm_config.cache_config.prefix_match_unit = None
+
+    store_worker = worker.MooncakeStoreWorker(
+        vllm_config, _make_qsa_hybrid_kv_cache_config()
+    )
+
+    assert [
+        group.kv_cache_spec.block_size for group in store_worker._kv_cache_groups
+    ] == [800, 800]
+    assert [db.block_size for db in store_worker.token_dbs] == [800, 800]
+
+
+def _make_two_full_attention_groups_kv_cache_config():
+    """Two full-attention groups with different block sizes: no hybrid
+    layers, but block-level failure reporting is still unsupported."""
+    return KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full0"],
+                FullAttentionSpec(
+                    block_size=800, num_kv_heads=8, head_size=64, dtype=None
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["full1"],
+                FullAttentionSpec(
+                    block_size=1600, num_kv_heads=8, head_size=64, dtype=None
+                ),
+            ),
+        ],
+    )
+
+
+@pytest.mark.parametrize(
+    ("make_config", "expected_is_hma_required"),
+    [
+        (lambda: _make_qsa_hybrid_kv_cache_config(), True),
+        (lambda: _make_two_full_attention_groups_kv_cache_config(), True),
+        (lambda: _make_kv_cache_config(block_size=800), False),
+    ],
+    ids=["hybrid", "two-full-attention-groups", "single-group"],
+)
+def test_worker_is_hma_required_from_kv_cache_groups(
+    monkeypatch, make_config, expected_is_hma_required
+):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setattr(
+        worker.MooncakeStoreConfig,
+        "load_from_config",
+        staticmethod(lambda: _make_config()),
+    )
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 800
+    vllm_config.cache_config.enable_prefix_caching = True
+    vllm_config.cache_config.prefix_match_unit = None
+
+    store_worker = worker.MooncakeStoreWorker(vllm_config, make_config())
+
+    assert store_worker._is_hma_required is expected_is_hma_required
+
+
 def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):
     store = MagicMock()
     store.setup.return_value = 0
@@ -2811,6 +2980,7 @@ def __init__(self) -> None:
     fake_store_module = types.ModuleType("mooncake.store")
     fake_store_module.MooncakeDistributedStore = lambda: store  # type: ignore[attr-defined]
     fake_store_module.ReplicateConfig = FakeReplicateConfig  # type: ignore[attr-defined]
+    fake_store_module.ObjectDataType = SimpleNamespace(TENSOR=1)  # 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)
@@ -3280,6 +3450,7 @@ def _make_bare_worker(
     worker.can_put = kv_role in ("kv_producer", "kv_both") or save_decode_cache
     worker._capacity_only = False
     worker.block_size = block_size
+    worker._is_hma_required = False
     worker.tp_rank = 0
     worker.enable_kv_events = False
     worker.load_async = True
@@ -3297,7 +3468,9 @@ def _make_bare_worker(
     # path; everything flows through the coordinator).
     from vllm.v1.kv_cache_interface import (
         FullAttentionSpec,
+        KVCacheConfig,
         KVCacheGroupSpec,
+        KVCacheTensor,
     )
 
     worker.disk_offload_buffer_budget_bytes = None
@@ -3311,6 +3484,18 @@ def _make_bare_worker(
         block_size=block_size, num_kv_heads=8, head_size=64, dtype=None
     )
     group = KVCacheGroupSpec(["layer0", "__cross_layer__"], spec)
+    worker._kv_cache_config = KVCacheConfig(
+        num_blocks=num_gpu_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=0,
+                layers=["layer0", "__cross_layer__"],
+                layer_stride=0,
+                block_stride=0,
+            )
+        ],
+        kv_cache_groups=[group],
+    )
     worker._kv_cache_groups = [group]
     worker.pcp_size = 1
     worker.dcp_size = 1
@@ -4059,6 +4244,102 @@ def test_register_kv_caches_shared_storage(layout: KVCacheLayout):
     worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
 
 
+def test_register_kv_caches_uses_transfer_group_memory_domain():
+    """Derived GPU caches must not change host-source transfer addresses."""
+    from vllm.v1.kv_cache_interface import (
+        KVCacheConfig,
+        KVCacheGroupRole,
+        KVCacheGroupSpec,
+        KVCacheTensor,
+    )
+
+    host_num_blocks = 4
+    gpu_num_blocks = 2
+    page_size = 32
+    source_spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16
+    )
+    indexer_spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16
+    )
+    source_group = KVCacheGroupSpec(
+        ["source"],
+        source_spec,
+        host_resident=True,
+    )
+    indexer_group = KVCacheGroupSpec(
+        ["indexer"],
+        indexer_spec,
+        enable_kv_transfer=False,
+        role=KVCacheGroupRole.HISPARSE_INDEXER,
+    )
+    config = KVCacheConfig(
+        num_blocks=gpu_num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=host_num_blocks * page_size,
+                layers=["source"],
+                layer_stride=host_num_blocks * page_size,
+                block_stride=page_size,
+                host_resident=True,
+            ),
+            KVCacheTensor(
+                size=gpu_num_blocks * page_size,
+                layers=["indexer"],
+                layer_stride=gpu_num_blocks * page_size,
+                block_stride=page_size,
+            ),
+        ],
+        kv_cache_groups=[source_group, indexer_group],
+        hisparse_host_num_blocks=host_num_blocks,
+    )
+    worker = _make_bare_worker(num_gpu_blocks=gpu_num_blocks)
+    worker._kv_cache_config = config
+    worker._kv_cache_groups = list(config.transfer_groups)
+    worker.token_dbs = [
+        ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16)
+    ]
+    source = torch.zeros(host_num_blocks, page_size, dtype=torch.uint8)
+    indexer = torch.zeros(gpu_num_blocks, page_size, dtype=torch.uint8)
+
+    _register_with_mocked_threads(
+        worker,
+        {"source": source, "indexer": indexer},
+    )
+
+    db = worker.token_dbs[0]
+    assert db.kv_caches_base_addr == [source.data_ptr()]
+    assert db.block_len == [page_size]
+    worker.store.register_buffer.assert_called_once_with(
+        source.data_ptr(), source.nbytes
+    )
+
+
+def test_register_kv_caches_excludes_nonprefix_cacheable_layers():
+    num_blocks = 2
+    worker = _make_bare_worker(num_gpu_blocks=num_blocks)
+    store_layout = MagicMock()
+    worker.token_dbs[0].store_layout = store_layout
+    full_cache = torch.zeros(num_blocks, 4, dtype=torch.float16)
+    qsa_cache = torch.zeros(num_blocks, 4, dtype=torch.float16)
+
+    _register_with_mocked_threads(
+        worker,
+        {"layer0": full_cache, "qsa": qsa_cache},
+    )
+
+    store_layout.register_kv_caches.assert_called_once()
+    registered_caches, registered_num_blocks = (
+        store_layout.register_kv_caches.call_args.args
+    )
+    assert registered_num_blocks == num_blocks
+    assert len(registered_caches) == 1
+    assert registered_caches[0] is full_cache
+    worker.store.register_buffer.assert_called_once_with(
+        full_cache.data_ptr(), full_cache.untyped_storage().nbytes()
+    )
+
+
 def test_register_kv_caches_separate_head_groups():
     # LHBNC gives each head group its own region; the K/V split doubles heads.
     layout = KVCacheLayout.LHBNC
@@ -4120,10 +4401,10 @@ 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.global_segment_size == mooncake_store.DEFAULT_GLOBAL_SEGMENT_SIZE
+    assert cfg.local_buffer_size == mooncake_store.DEFAULT_LOCAL_BUFFER_SIZE
     assert cfg.enable_offload is False
-    assert cfg.tenant_id == worker.DEFAULT_TENANT_ID
+    assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID
 
 
 def test_config_pr40900_unchanged(tmp_path):
@@ -4145,7 +4426,7 @@ def test_config_pr40900_unchanged(tmp_path):
     assert cfg.global_segment_size == 4 * 1024**3
     assert cfg.local_buffer_size == 4 * 1024**3
     assert cfg.enable_offload is False
-    assert cfg.tenant_id == worker.DEFAULT_TENANT_ID
+    assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID
 
 
 def test_config_from_file_normalizes_tenant_id(tmp_path):
@@ -4175,7 +4456,7 @@ def test_config_from_file_normalizes_empty_tenant_id_to_default(tmp_path):
 
     cfg = worker.MooncakeStoreConfig.from_file(config_path)
 
-    assert cfg.tenant_id == worker.DEFAULT_TENANT_ID
+    assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID
 
 
 def test_config_from_file_rejects_non_string_tenant_id(tmp_path):
@@ -4283,7 +4564,7 @@ def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch):
 
 
 def test_topology_embedded_cpu_only(tmp_path, monkeypatch):
-    """embedded + CPU-only: no mode key (defaults to embedded),
+    """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()
@@ -4327,7 +4608,20 @@ def test_topology_embedded_cpu_only(tmp_path, monkeypatch):
     assert w.disk_offload_buffer_budget_bytes is None
 
 
-def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch):
+def _create_store_client_for_tenant_test(consumer):
+    if consumer == "ec":
+        from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.store_client import (  # noqa: E501
+            create_mooncake_embedding_store_client,
+        )
+
+        return create_mooncake_embedding_store_client()
+    return worker.MooncakeStoreWorker(
+        _make_vllm_config(rank=1), _make_kv_cache_config()
+    )
+
+
+@pytest.mark.parametrize("consumer", ["kv", "ec"])
+def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch, consumer):
     store = MagicMock()
     store.setup.return_value = 0
     _install_fake_mooncake(monkeypatch, store)
@@ -4348,12 +4642,13 @@ def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch):
         ),
     )
 
-    worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config())
+    _create_store_client_for_tenant_test(consumer)
 
     assert store.setup.call_args.kwargs == {"tenant_id": "tenant-a"}
 
 
-def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch):
+@pytest.mark.parametrize("consumer", ["kv", "ec"])
+def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch, consumer):
     store = MagicMock()
     setup_error = TypeError(
         "setup(): incompatible function arguments; "
@@ -4379,7 +4674,7 @@ def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch):
     )
 
     with pytest.raises(TypeError) as exc_info:
-        worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config())
+        _create_store_client_for_tenant_test(consumer)
 
     assert exc_info.value is setup_error
 
diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py
index d169779c5df1..f95e12c601bb 100644
--- a/tests/v1/kv_connector/unit/test_moriio_connector.py
+++ b/tests/v1/kv_connector/unit/test_moriio_connector.py
@@ -39,10 +39,12 @@
     make_zmq_path,
 )
 from vllm.v1.kv_cache_interface import (
+    ChunkedLocalAttentionSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheTensor,
+    SlidingWindowSpec,
     compute_layer_kv_cache_shape_bytes,
 )
 
@@ -292,7 +294,6 @@ def create_vllm_config(
 
 def test_write_mode_saves_local_block_ids():
     """Write mode records local block ids in MoRIIOConnectorMetadata.reqs_to_save."""
-
     # Setup Scheduler and Request
     vllm_config = create_vllm_config(role="kv_producer")
     scheduler = create_scheduler(vllm_config)
@@ -338,7 +339,7 @@ def test_write_mode_saves_local_block_ids():
     req_meta = kv_connector_metadata.reqs_to_save[request_id]
 
     for block_id, block in zip(
-        req_meta.local_block_ids,
+        req_meta.local_block_ids[0],
         scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
             request_id
         ],
@@ -393,7 +394,7 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids():
     req_meta = kv_connector_metadata.reqs_to_save[request_id]
 
     for block_id, block in zip(
-        req_meta.local_block_ids,
+        req_meta.local_block_ids[0],
         scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
             request_id
         ],
@@ -403,7 +404,6 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids():
 
 def test_read_mode_loads_remote_block_ids():
     """Read mode loads remote block ids into local cache mapping."""
-
     # Setup Scheduler and Request
     vllm_config = create_vllm_config(role="kv_consumer", read_mode=True)
     scheduler = create_scheduler(vllm_config)
@@ -432,7 +432,7 @@ def test_read_mode_loads_remote_block_ids():
     ].req_to_blocks[request_id]
 
     # Set remote block ids to be fetched.
-    request.kv_transfer_params["remote_block_ids"] = block_list
+    request.kv_transfer_params["remote_block_ids"] = [block_list]
 
     # Remote Prefill, triggers MoRIIOConnectorMetadata.
 
@@ -457,7 +457,7 @@ def test_read_mode_loads_remote_block_ids():
     req_meta = kv_connector_metadata.reqs_to_recv[request_id]
 
     for block_id, block in zip(
-        req_meta.local_block_ids,
+        req_meta.local_block_ids[0],
         scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
             request_id
         ],
@@ -629,7 +629,6 @@ def test_register_kv_caches(mock_parallel_groups):
 )
 def test_moriio_handshake_returns_metadata(mock_parallel_groups):
     """MoRIIO handshake socket returns valid agent metadata over ZMQ."""
-
     ROLE = "kv_consumer"
     vllm_config = create_vllm_config(role=ROLE)
     # Create test kv cache tensors using KVCacheSpec layout
@@ -699,3 +698,236 @@ def test_resolve_host_ip_prefers_extra_config():
     fallback = get_ip()
     assert resolve_host_ip({}) == fallback
     assert resolve_host_ip({"host_ip": ""}) == fallback
+
+
+def _make_hybrid_kv_cache_config() -> KVCacheConfig:
+    """One full-attention group + one sliding-window group (Gemma-like)."""
+    full_spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
+    )
+    sw_spec = SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=32,
+    )
+    num_blocks = 2
+    page = full_spec.page_size_bytes
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=2 * num_blocks * page,
+                layers=["full0", "sw0"],
+                layer_stride=num_blocks * page,
+                block_stride=page,
+            )
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(layer_names=["full0"], kv_cache_spec=full_spec),
+            KVCacheGroupSpec(layer_names=["sw0"], kv_cache_spec=sw_spec),
+        ],
+    )
+
+
+def _read_scheduler(
+    kv_cache_config: KVCacheConfig, disable_hma: bool = False
+) -> MoRIIOConnectorScheduler:
+    vllm_config = create_vllm_config(role="kv_producer", read_mode=True)
+    vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hma
+    with set_current_vllm_config(vllm_config):
+        connector = MoRIIOConnector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )
+    assert connector.connector_scheduler is not None
+    return connector.connector_scheduler
+
+
+def test_hma_blocks_per_sw_two_groups():
+    """A Full + SlidingWindow config runs with HMA and computes block budgets
+    correctly"""
+    scheduler = _read_scheduler(_make_hybrid_kv_cache_config())
+    assert scheduler._is_hma_required is True
+    # cdiv(32, 16) + 1 == 3 for the sliding-window group, 0 for full attention.
+    assert scheduler.blocks_per_sw == [0, 3]
+
+
+@pytest.mark.parametrize(
+    "swa_enabled, disable_hma, expected_is_hma",
+    [
+        (True, False, True),  # sliding-window group present, HMA enabled
+        (True, True, False),  # sliding-window group present but HMA disabled
+        (False, False, False),  # full-attention only, HMA not needed
+    ],
+)
+def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma):
+    """_is_hma_required tracks both the KV cache groups and the
+    --disable-hybrid-kv-cache-manager flag. When HMA is off,
+    get_exchange_clipped_blocks must be a no-op."""
+    config = (
+        _make_hybrid_kv_cache_config() if swa_enabled else _make_test_kv_cache_config()
+    )
+    scheduler = _read_scheduler(config, disable_hma=disable_hma)
+    assert scheduler._is_hma_required is expected_is_hma
+    if not expected_is_hma:
+        blocks = [[1, 2, 3, 4, 5]]
+        assert scheduler.get_exchange_clipped_blocks(blocks) == blocks
+
+
+def test_transfer_disabled_hybrid_group_is_ignored():
+    config = _make_hybrid_kv_cache_config()
+    config.kv_cache_groups[1].enable_kv_transfer = False
+    scheduler = _read_scheduler(config)
+
+    assert scheduler._is_hma_required is False
+    assert scheduler.get_exchange_clipped_blocks([[1, 2, 3], [20, 21, 22]]) == [
+        [1, 2, 3]
+    ]
+
+
+def test_non_sliding_window_hybrid_is_rejected():
+    """A hybrid group that is not sliding-window (e.g. chunked-local
+    attention) must fail closed rather than be silently mistransferred."""
+    full_spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
+    )
+    local_spec = ChunkedLocalAttentionSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        attention_chunk_size=32,
+    )
+    num_blocks = 2
+    page = full_spec.page_size_bytes
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=2 * num_blocks * page,
+                layers=["full0", "local0"],
+                layer_stride=num_blocks * page,
+                block_stride=page,
+            )
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(layer_names=["full0"], kv_cache_spec=full_spec),
+            KVCacheGroupSpec(layer_names=["local0"], kv_cache_spec=local_spec),
+        ],
+    )
+    with pytest.raises(NotImplementedError, match="sliding-window hybrid"):
+        _read_scheduler(config)
+
+
+def test_token_count_basis_uses_full_attention_group():
+    """Chunked-prefill token counting must use an unclipped full-attention
+    group (blocks_per_sw == 0), not a clipped sliding-window group."""
+    scheduler = _read_scheduler(_make_hybrid_kv_cache_config())
+    # Group 0 is the full-attention group
+    assert scheduler._full_attn_group_idx == 0
+    assert scheduler._full_attn_block_size == 16
+
+
+def test_get_exchange_clipped_blocks_clips_only_sw_group():
+    """get_exchange_clipped_blocks keeps the full attn group intact and clips the
+    sliding-window group to its window tail."""
+    scheduler = _read_scheduler(_make_hybrid_kv_cache_config())
+    full = [10, 11, 12, 13, 14]
+    sw = [20, 21, 22, 23, 24]
+    clipped = scheduler.get_exchange_clipped_blocks([full, sw])
+    assert clipped[0] == full
+    assert clipped[1] == [22, 23, 24]
+
+
+def test_metadata_hma_block_ids_preserved_per_group():
+    """add_new_req stores per-group (BlockIds) block lists unchanged for both
+    read and write, so the hybrid group structure is retained."""
+    metadata = MoRIIOConnectorMetadata()
+
+    # Assume:
+    # - Full-attention group (6 blocks) +
+    # - sliding-window group already clipped to its window tail (3 blocks).
+    fa_blocks = [0, 1, 2, 3, 4, 5]
+    sw_blocks = [10, 11, 12]
+    local_block_ids = [fa_blocks, sw_blocks]
+    remote_block_ids = [[100, 101, 102, 103, 104, 105], [200, 201, 202]]
+
+    base_params = {
+        "remote_engine_id": "remote-engine",
+        "remote_host": "127.0.0.1",
+        "remote_handshake_port": 6301,
+        "remote_notify_port": 61005,
+        "remote_block_ids": remote_block_ids,
+    }
+
+    # Read mode: both local and remote ids stay per-group.
+    metadata.add_new_req(
+        request_id="recv-req",
+        local_block_ids=local_block_ids,
+        kv_transfer_params={**base_params, "transfer_id": "recv-req"},
+    )
+    recv_meta = metadata.reqs_to_recv["recv-req"]
+    assert recv_meta.local_block_ids == [fa_blocks, sw_blocks]
+    assert recv_meta.remote_block_ids == remote_block_ids
+
+    # Write mode: the decode peer allocates its own blocks, so #remote_block_ids may
+    # be empty, Local group structure is kept.
+    metadata.add_new_req(
+        request_id="save-req",
+        local_block_ids=local_block_ids,
+        kv_transfer_params={
+            **base_params,
+            "transfer_id": "save-req",
+            "remote_block_ids": [],
+        },
+        write_mode=True,
+    )
+    save_meta = metadata.reqs_to_save["save-req"]
+    assert save_meta.local_block_ids == [fa_blocks, sw_blocks]
+    assert save_meta.remote_block_ids == []
+
+
+def test_single_group_path_unchanged():
+    """A full attn config does not enable HMA and never clips blocks."""
+    scheduler = _read_scheduler(_make_test_kv_cache_config())
+    assert scheduler._is_hma_required is False
+    assert scheduler.blocks_per_sw == [0]
+    blocks = [[1, 2, 3, 4, 5]]
+    assert scheduler.get_exchange_clipped_blocks(blocks) == blocks
+
+
+def test_hybrid_write_mode_rejected():
+    """Hybrid KV cache groups are unsupported in WRITE mode and fail closed."""
+    vllm_config = create_vllm_config(role="kv_producer", read_mode=False)
+    with (
+        set_current_vllm_config(vllm_config),
+        pytest.raises(NotImplementedError),
+    ):
+        MoRIIOConnector(
+            vllm_config,
+            KVConnectorRole.SCHEDULER,
+            _make_hybrid_kv_cache_config(),
+        )
+
+
+def test_worker_layer_to_group_routing(mock_parallel_groups):
+    """The worker maps every layer to its KV cache group correctly."""
+    vllm_config = create_vllm_config(role="kv_consumer", read_mode=True)
+    # Building the worker directly bypasses MoRIIOConnector._set_port_defaults,
+    # so provide the ports manually.
+    vllm_config.kv_transfer_config.kv_connector_extra_config.update(
+        {
+            "http_port": 12346,
+            "handshake_port": 12347,
+            "notify_port": 12348,
+        }
+    )
+    with set_current_vllm_config(vllm_config):
+        worker = FakeMoRIIOConnectorWorker(
+            vllm_config,
+            "engine0",
+            hand_shake_latency=0,
+            kv_cache_config=_make_hybrid_kv_cache_config(),
+        )
+    assert worker.layer_to_group == {"full0": 0, "sw0": 1}
diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py
new file mode 100644
index 000000000000..13934f896afe
--- /dev/null
+++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py
@@ -0,0 +1,809 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pure-CPU unit tests for the MoRIIO hybrid (attention + mamba/KDA) scheduler
+and the KDA offset-template cache.
+
+These exercise the highest-risk, otherwise-untested scheduler logic without a
+GPU or the ``mori`` runtime: per-group block-id splitting, the READ/WRITE
+``N-1`` token accounting, P-side prompt truncation, and offset-template cache
+wiring.
+
+Like ``test_moriio_kv_layout.py`` the whole module is skipped unless it is
+running on ROCm with ``mori`` installed (importing the connector pulls in
+``mori``). The authoritative run happens on the MIA recipe image.
+"""
+
+import importlib
+import importlib.util
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+mori_available = importlib.util.find_spec("mori") is not None
+
+if not (current_platform.is_rocm() and mori_available):
+    pytest.skip(
+        "MoRIIOs are only available on ROCm with mori package installed",
+        allow_module_level=True,
+    )
+
+moriio_connector = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector"
+)
+moriio_layout = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout"
+)
+ssm_conv_transfer_utils = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils"
+)
+MambaConvSplitInfo = ssm_conv_transfer_utils.MambaConvSplitInfo
+MoRIIOMode = moriio_connector.MoRIIOMode
+moriio_common = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common"
+)
+
+
+class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler):  # type: ignore[name-defined]
+    """Constructs a scheduler without the heavy real ``__init__``; only the
+    attributes read by the method under test are set."""
+
+    def __init__(self, **attrs):
+        self._mamba_group_ids: list[int] = []
+        self._attn_group_ids: list[int] = [0]
+        self._ssm_state_slots_are_positional = False
+        self._is_hma_required = False
+        self.kv_cache_config = SimpleNamespace(
+            kv_cache_groups=[None, None],
+            select_transfer_block_ids=lambda block_ids: tuple(block_ids),
+        )
+        for k, v in attrs.items():
+            setattr(self, k, v)
+
+
+class _FakeConnector(moriio_connector.MoRIIOConnector):  # type: ignore[name-defined]
+    def __init__(self, connector_scheduler):
+        self.connector_scheduler = connector_scheduler
+
+
+class _FakeWorker(moriio_connector.MoRIIOConnectorWorker):  # type: ignore[name-defined]
+    def __init__(self, **attrs):
+        for k, v in attrs.items():
+            setattr(self, k, v)
+
+
+class _FakeBlocks:
+    def __init__(self, all_groups):
+        self._all_groups = all_groups
+
+    def get_block_ids(self):
+        return self._all_groups
+
+
+def _slot_strided(num_slots, per_slot_shape, slot_stride, dtype=torch.bfloat16):
+    inner = 1
+    for d in per_slot_shape:
+        inner *= d
+    assert slot_stride >= inner
+    backing = torch.zeros(num_slots * slot_stride, dtype=dtype)
+    inner_strides: list[int] = []
+    acc = 1
+    for d in reversed(per_slot_shape):
+        inner_strides.insert(0, acc)
+        acc *= d
+    return backing.as_strided(
+        (num_slots, *per_slot_shape), (slot_stride, *inner_strides)
+    )
+
+
+def _gdn_split_info(conv_rows=3, key_dim=4, value_dim=8, dtype_size=2):
+    conv_dim = 2 * key_dim + value_dim
+    conv_state_bytes = conv_dim * conv_rows * dtype_size
+    ssm_state_bytes = 64
+    return MambaConvSplitInfo(
+        conv_rows=conv_rows,
+        local_proj_dims=(key_dim, key_dim, value_dim),
+        conv_dtype_size=dtype_size,
+        ssm_sizes=(conv_state_bytes, ssm_state_bytes),
+    )
+
+
+def _mamba_spec(num_states: int = 2):
+    shapes = tuple((1, 1) for _ in range(num_states))
+    return moriio_connector.MambaSpec(
+        block_size=16,
+        shapes=shapes,
+        dtypes=(torch.float32,) * num_states,
+        mamba_cache_mode="all",
+    )
+
+
+def _gate_vllm_config(*, read_mode: bool = True, speculative_config=None):
+    return SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(
+            kv_connector_extra_config={"read_mode": read_mode}
+        ),
+        speculative_config=speculative_config,
+    )
+
+
+# --------------------------------------------------------------------------
+# split_block_groups
+# --------------------------------------------------------------------------
+def test_split_block_groups_separates_attention_and_mamba():
+    sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1])
+    block_ids = ([1, 2], [99])
+    attn, mamba = sched.split_block_groups(block_ids)
+    assert attn == [1, 2]
+    assert mamba == [[99]]
+
+
+def test_split_block_groups_ignores_transfer_disabled_group():
+    mamba_spec = moriio_connector.MambaSpec(
+        block_size=16,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="all",
+    )
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
+            SimpleNamespace(enable_kv_transfer=False, kv_cache_spec=object()),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=mamba_spec),
+        ],
+    )
+    config.transfer_groups = tuple(
+        group for group in config.kv_cache_groups if group.enable_kv_transfer
+    )
+    config.select_transfer_block_ids = lambda block_ids: tuple(
+        block_ids[group_id]
+        for group_id, group in enumerate(config.kv_cache_groups)
+        if group.enable_kv_transfer
+    )
+    attn_groups, mamba_groups = moriio_connector._split_kv_cache_group_kinds(config)
+    assert attn_groups == [0]
+    assert mamba_groups == [1]
+
+    sched = _FakeScheduler(
+        _has_mamba=True,
+        _attn_group_ids=attn_groups,
+        _mamba_group_ids=mamba_groups,
+        kv_cache_config=config,
+    )
+    assert sched.split_block_groups(([1, 2], [70, 71], [99])) == (
+        [1, 2],
+        [[99]],
+    )
+
+
+def test_exchange_blocks_ignore_transfer_disabled_group():
+    config = SimpleNamespace(
+        kv_cache_groups=[None, None, None],
+        select_transfer_block_ids=lambda block_ids: (block_ids[0], block_ids[2]),
+    )
+    sched = _FakeScheduler(
+        kv_cache_config=config,
+        _is_hma_required=True,
+        blocks_per_sw=[0, 2],
+    )
+
+    assert sched.get_exchange_clipped_blocks(([1, 2], [70, 71], [90, 91, 92])) == [
+        [1, 2],
+        [91, 92],
+    ]
+
+
+def test_scheduler_rejects_multiple_attention_groups_with_mamba():
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
+        ],
+    )
+    config.transfer_groups = tuple(config.kv_cache_groups)
+
+    with pytest.raises(moriio_common.MoRIIOError, match="exactly one transferable"):
+        moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
+
+
+def test_split_block_groups_preserves_multiple_mamba_groups():
+    sched = _FakeScheduler(
+        _has_mamba=True,
+        _attn_group_ids=[0],
+        _mamba_group_ids=[1, 2],
+    )
+
+    assert sched.split_block_groups(([1, 2], [40], [70])) == (
+        [1, 2],
+        [[40], [70]],
+    )
+
+
+def test_scheduler_rejects_incompatible_mamba_groups():
+    incompatible = _mamba_spec()
+    object.__setattr__(incompatible, "page_size_padded", 4096)
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=incompatible),
+        ]
+    )
+    config.transfer_groups = tuple(config.kv_cache_groups)
+
+    with pytest.raises(moriio_common.MoRIIOError, match="same cache spec"):
+        moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
+
+
+def test_scheduler_rejects_mamba_group_without_two_states():
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec(1)),
+        ],
+    )
+    config.transfer_groups = tuple(config.kv_cache_groups)
+
+    with pytest.raises(moriio_common.MoRIIOError, match="exactly two Mamba states"):
+        moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
+
+
+def test_scheduler_rejects_speculative_hybrid_read():
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
+        ],
+    )
+    config.transfer_groups = tuple(config.kv_cache_groups)
+
+    with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"):
+        moriio_connector.MoRIIOConnectorScheduler(
+            _gate_vllm_config(speculative_config=object()), "engine", config
+        )
+
+
+def test_scheduler_rejects_hybrid_write():
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
+            SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
+        ],
+    )
+    config.transfer_groups = tuple(config.kv_cache_groups)
+
+    with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"):
+        moriio_connector.MoRIIOConnectorScheduler(
+            _gate_vllm_config(read_mode=False), "engine", config
+        )
+
+
+def test_split_block_groups_rejects_non_hybrid_use():
+    sched = _FakeScheduler(_has_mamba=False)
+
+    with pytest.raises(moriio_common.MoRIIOError, match="non-hybrid"):
+        sched.split_block_groups(([5, 6, 7],))
+
+
+def test_split_block_groups_accepts_empty_abort_payload():
+    sched = _FakeScheduler(_has_mamba=True)
+
+    assert sched.split_block_groups(()) == ([], [])
+
+
+def test_split_block_groups_keeps_positional_slots_in_all_mode():
+    # mamba_cache_mode="all" keeps a state per block position.
+    sched = _FakeScheduler(
+        _has_mamba=True,
+        _attn_group_ids=[0],
+        _mamba_group_ids=[1],
+        _ssm_state_slots_are_positional=True,
+    )
+    attn, mamba = sched.split_block_groups(([1], [40, 41, 42]))
+    assert attn == [1]
+    assert mamba == [[40, 41, 42]]
+
+
+def test_split_block_groups_keeps_only_running_state_outside_all_mode():
+    sched = _FakeScheduler(
+        _has_mamba=True,
+        _attn_group_ids=[0],
+        _mamba_group_ids=[1],
+    )
+    assert sched.split_block_groups(([1], [40, 41])) == ([1], [[41]])
+
+
+# --------------------------------------------------------------------------
+# request_finished_all_groups
+# --------------------------------------------------------------------------
+def test_request_finished_all_groups_carries_attn_and_mamba_in_one_field():
+    seen = {}
+
+    def _fake_request_finished(request, attn_block_ids, mamba_block_groups):
+        seen["attn"] = list(attn_block_ids)
+        seen["mamba"] = [list(group) for group in mamba_block_groups]
+        return True, {
+            "do_remote_prefill": True,
+            "remote_block_ids": [attn_block_ids, *mamba_block_groups],
+        }
+
+    sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1])
+    sched.request_finished = _fake_request_finished
+    conn = _FakeConnector(sched)
+
+    request = SimpleNamespace(request_id="r0")
+    block_ids = ([10, 11], [42])
+    delay_free, params = conn.request_finished_all_groups(request, block_ids)
+
+    assert delay_free is True
+    # Attention ids drive request_finished; the mamba slot rides the SAME
+    # remote_block_ids channel as [attn, *mamba_groups], so the proxy/router
+    # needs no KDA-specific field.
+    assert seen["attn"] == [10, 11]
+    assert seen["mamba"] == [[42]]
+    assert params["remote_block_ids"] == [[10, 11], [42]]
+    assert "remote_mamba_block_ids" not in params
+
+
+def test_request_finished_all_groups_pure_attention_stays_grouped():
+    def _fake_request_finished(request, block_ids):
+        return True, {"remote_block_ids": block_ids}
+
+    sched = _FakeScheduler(_has_mamba=False, _mamba_group_ids=[])
+    sched.request_finished = _fake_request_finished
+    conn = _FakeConnector(sched)
+
+    delay_free, params = conn.request_finished_all_groups(
+        SimpleNamespace(request_id="r1"), ([7, 8],)
+    )
+    # Main's HMA contract keeps pure-attention block ids grouped.
+    assert params["remote_block_ids"] == ([7, 8],)
+    assert "remote_mamba_block_ids" not in params
+
+
+def test_connector_supports_divergent_local_hybrid_hits():
+    connector = _FakeConnector(None)
+    connector._transfers_mamba_state = True
+    connector.mode = MoRIIOMode.READ
+    assert connector.supports_divergent_local_hybrid_hits is True
+    connector.mode = MoRIIOMode.WRITE
+    assert connector.supports_divergent_local_hybrid_hits is False
+
+    connector.mode = MoRIIOMode.READ
+    connector._transfers_mamba_state = False
+    assert connector.supports_divergent_local_hybrid_hits is False
+
+
+def _make_read_scheduler():
+    return _FakeScheduler(
+        is_producer=False,
+        mode=MoRIIOMode.READ,
+        _has_mamba=True,
+        _attn_group_ids=[0],
+        _mamba_group_ids=[1],
+        _ssm_state_slots_are_positional=True,
+        request_id_to_transfer_id={},
+        transfer_id_to_request_id={},
+        _reqs_need_recv={},
+        _req_kv_params={},
+        _max_decode_tail_blocks=1,
+    )
+
+
+def _make_read_request(remote_block_ids):
+    return SimpleNamespace(
+        request_id="req",
+        kv_transfer_params={
+            "do_remote_prefill": True,
+            "transfer_id": "tx",
+            "remote_engine_id": "prefill",
+            "remote_block_ids": remote_block_ids,
+        },
+    )
+
+
+def test_update_state_drops_decode_recompute_tail_block():
+    sched = _make_read_scheduler()
+    request = _make_read_request([[10, 11], [90, 91]])
+    blocks = _FakeBlocks(
+        all_groups=([100, 101, 102], [200, 201, 202]),
+    )
+
+    sched.update_state_after_alloc(request, blocks, num_external_tokens=256)
+
+    assert sched._reqs_need_recv["req"][1] == [[100, 101], [200, 201]]
+    assert sched._req_kv_params["req"]["remote_block_ids"] == [
+        [10, 11],
+        [90, 91],
+    ]
+
+
+def test_update_state_pairs_shorter_local_blocks_with_remote_suffix():
+    sched = _make_read_scheduler()
+    request = _make_read_request([[10, 11, 12], [92]])
+    blocks = _FakeBlocks(
+        all_groups=([102], [202]),
+    )
+
+    sched.update_state_after_alloc(request, blocks, num_external_tokens=64)
+
+    assert sched._reqs_need_recv["req"][1] == [[102], [202]]
+    assert sched._req_kv_params["req"]["remote_block_ids"] == [[12], [92]]
+
+
+def test_update_state_full_attention_hit_still_carries_mamba_state():
+    sched = _make_read_scheduler()
+    request = _make_read_request([[10], [90]])
+
+    sched.update_state_after_alloc(
+        request,
+        _FakeBlocks(all_groups=([100], [200])),
+        num_external_tokens=0,
+    )
+
+    assert sched._reqs_need_recv["req"][1] == [[], [200]]
+
+
+def test_update_state_preserves_each_mamba_group():
+    sched = _make_read_scheduler()
+    sched._mamba_group_ids = [1, 2]
+    request = _make_read_request([[10], [90], [80]])
+
+    sched.update_state_after_alloc(
+        request,
+        _FakeBlocks(all_groups=([100], [200], [300])),
+        num_external_tokens=64,
+    )
+
+    assert sched._reqs_need_recv["req"][1] == [[100], [200], [300]]
+    assert sched._req_kv_params["req"]["remote_block_ids"] == [
+        [10],
+        [90],
+        [80],
+    ]
+
+
+def test_align_read_blocks_allows_configured_lookahead_tail():
+    assert moriio_connector.MoRIIOConnectorScheduler._align_read_blocks(
+        [100, 101, 102], [10], max_decode_tail_blocks=2
+    ) == ([100], [10])
+    with pytest.raises(ValueError, match="allowed local tail"):
+        moriio_connector.MoRIIOConnectorScheduler._align_read_blocks(
+            [100, 101, 102], [10], max_decode_tail_blocks=1
+        )
+
+
+def test_update_state_aligns_attention_only_decode_tail():
+    sched = _make_read_scheduler()
+    sched._has_mamba = False
+    sched._mamba_group_ids = []
+    request = _make_read_request([[10, 11]])
+
+    sched.update_state_after_alloc(
+        request,
+        _FakeBlocks(all_groups=([100, 101, 102],)),
+        num_external_tokens=256,
+    )
+
+    assert sched._reqs_need_recv["req"][1] == [[100, 101]]
+    assert sched._req_kv_params["req"]["remote_block_ids"] == [[10, 11]]
+
+
+def test_read_producer_does_not_queue_write_only_save_state():
+    sched = _FakeScheduler(
+        is_producer=True,
+        mode=MoRIIOMode.READ,
+        _has_mamba=True,
+        _attn_group_ids=[0],
+        _mamba_group_ids=[1],
+        request_id_to_transfer_id={},
+        transfer_id_to_request_id={},
+        _reqs_need_save={},
+        _req_kv_params={},
+    )
+    request = SimpleNamespace(
+        request_id="req",
+        kv_transfer_params={"do_remote_decode": True, "transfer_id": "tx"},
+    )
+
+    sched.update_state_after_alloc(
+        request,
+        _FakeBlocks(all_groups=([100, 101], [200])),
+        num_external_tokens=0,
+    )
+
+    assert sched._reqs_need_save == {}
+    assert sched._req_kv_params == {}
+
+
+def test_write_producer_queues_grouped_attention_blocks():
+    sched = _FakeScheduler(
+        is_producer=True,
+        mode=MoRIIOMode.WRITE,
+        _has_mamba=False,
+        _attn_group_ids=[0],
+        _mamba_group_ids=[],
+        request_id_to_transfer_id={},
+        transfer_id_to_request_id={},
+        _reqs_need_save={},
+        _req_kv_params={},
+    )
+    request = SimpleNamespace(
+        request_id="req",
+        kv_transfer_params={"do_remote_decode": True, "transfer_id": "tx"},
+    )
+
+    sched.update_state_after_alloc(
+        request,
+        _FakeBlocks(all_groups=([100, 101],)),
+        num_external_tokens=0,
+    )
+
+    assert sched._reqs_need_save["req"][1] == ([100, 101],)
+    assert sched._req_kv_params["req"] == request.kv_transfer_params
+
+
+def test_session_build_rejects_per_layer_region_count_mismatch():
+    worker = _FakeWorker(
+        built_write_session={},
+        layer_name_to_local_kv_cache_metadata={"kda.0": ["conv", "ssm"]},
+        layer_name_to_remote_kv_cache_metadata={"prefill": {"kda.0": ["conv"]}},
+    )
+
+    with pytest.raises(moriio_common.MoRIIOError, match="registered 1 region"):
+        worker._get_built_session("prefill")
+
+
+def test_register_kv_caches_rejects_hybrid_write_before_registration():
+    worker = _FakeWorker(mode=MoRIIOMode.WRITE, _transfer_layer_names={"kda.0"})
+    worker._is_mamba_layer = lambda _layer_name: True
+
+    with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"):
+        worker.register_kv_caches({"kda.0": object()})
+
+
+def test_register_kv_caches_rejects_no_transfer_enabled_layers():
+    worker = _FakeWorker(mode=MoRIIOMode.READ, _transfer_layer_names=set())
+
+    with pytest.raises(moriio_common.MoRIIOError, match="no transfer-enabled"):
+        worker.register_kv_caches({"disabled.0": object()})
+
+
+def test_worker_rejects_incompatible_mamba_specs_before_registration():
+    incompatible = _mamba_spec()
+    object.__setattr__(incompatible, "page_size_padded", 4096)
+    worker = _FakeWorker(
+        mode=MoRIIOMode.READ,
+        world_size=1,
+        _transfer_layer_names={"kda.0", "kda.1"},
+        layer_to_spec={"kda.0": _mamba_spec(), "kda.1": incompatible},
+    )
+
+    with pytest.raises(moriio_common.MoRIIOError, match="same cache spec"):
+        worker.register_kv_caches({"kda.0": object(), "kda.1": object()})
+
+
+# --------------------------------------------------------------------------
+# _truncate_mamba_request_for_prefill
+# --------------------------------------------------------------------------
+def _mk_request(prompt, max_tokens=64, params=None):
+    return SimpleNamespace(
+        kv_transfer_params={} if params is None else params,
+        num_prompt_tokens=len(prompt),
+        prompt_token_ids=list(prompt),
+        prompt_embeds=None,
+        _all_token_ids=list(prompt),
+        max_tokens=max_tokens,
+    )
+
+
+def test_truncate_mamba_request_pops_last_token_and_caps_tokens():
+    sched = _FakeScheduler()
+    req = _mk_request([10, 11, 12, 13, 14])
+    sched._truncate_mamba_request_for_prefill(req)
+
+    assert req.prompt_token_ids == [10, 11, 12, 13]  # last token popped
+    assert req._all_token_ids == [10, 11, 12, 13]
+    assert req.num_prompt_tokens == 4  # N-1 accounting
+    assert req.max_tokens == 1
+    assert req.kv_transfer_params["_p_side_truncated"] is True
+
+
+def test_truncate_mamba_request_is_idempotent_across_reschedule():
+    sched = _FakeScheduler()
+    req = _mk_request([10, 11, 12])
+    sched._truncate_mamba_request_for_prefill(req)
+    assert req.num_prompt_tokens == 2
+    # A second call (e.g. after a preemption) must not truncate again.
+    sched._truncate_mamba_request_for_prefill(req)
+    assert req.num_prompt_tokens == 2
+    assert req.prompt_token_ids == [10, 11]
+
+
+def test_truncate_mamba_request_noop_for_single_token_prompt():
+    sched = _FakeScheduler()
+    req = _mk_request([10])
+    sched._truncate_mamba_request_for_prefill(req)
+    assert req.num_prompt_tokens == 1
+    assert req.prompt_token_ids == [10]
+    assert "_p_side_truncated" not in req.kv_transfer_params
+
+
+# --------------------------------------------------------------------------
+# get_num_new_matched_tokens  (hybrid N-1 accounting)
+# --------------------------------------------------------------------------
+def test_get_num_new_matched_tokens_read_recomputes_last_token():
+    sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.READ, _has_mamba=True)
+    req = SimpleNamespace(
+        num_prompt_tokens=10,
+        prompt_token_ids=list(range(10)),
+        kv_transfer_params=None,
+    )
+    n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
+    # READ always recomputes the final token locally: N-1 - computed.
+    assert n == 9
+    assert is_async is False
+
+
+def test_get_num_new_matched_tokens_write_plain_keeps_all_tokens():
+    sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=False)
+    req = SimpleNamespace(
+        num_prompt_tokens=10,
+        prompt_token_ids=list(range(10)),
+        kv_transfer_params=None,
+    )
+    n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=2)
+    # Pure-attention WRITE: no N-1 drop; full length minus already-computed.
+    assert n == 8
+    assert is_async is True
+
+
+@pytest.mark.parametrize(
+    ("mode", "num_computed_tokens", "expected", "is_async"),
+    [
+        (MoRIIOMode.READ, 0, 9, False),
+        (MoRIIOMode.READ, 10, 0, False),
+        (MoRIIOMode.WRITE, 2, 8, True),
+    ],
+)
+def test_get_num_new_matched_tokens_supports_embeds_only_prompts(
+    mode, num_computed_tokens, expected, is_async
+):
+    sched = _FakeScheduler(is_producer=False, mode=mode, _has_mamba=True)
+    req = SimpleNamespace(
+        num_prompt_tokens=10,
+        prompt_token_ids=None,
+        prompt_embeds=object(),
+        kv_transfer_params=None,
+    )
+
+    assert sched.get_num_new_matched_tokens(req, num_computed_tokens) == (
+        expected,
+        is_async,
+    )
+
+
+def test_get_num_new_matched_tokens_producer_returns_zero():
+    sched = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=True)
+    req = _mk_request(list(range(5)), params={"do_remote_decode": True})
+    n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
+    assert (n, is_async) == (0, False)
+    # The producer never truncates here: the scheduler has already measured the
+    # prefix-cache hit against the full prompt by this point.
+    assert req.num_prompt_tokens == 5
+
+
+def test_on_new_request_truncates_producer_prompt():
+    sched = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=True)
+    req = _mk_request(list(range(5)), params={"do_remote_decode": True})
+    sched.on_new_request(req)
+    # Producer stops at h(N-1): the last prompt token was dropped.
+    assert req.num_prompt_tokens == 4
+    assert req.max_tokens == 1
+    assert req.kv_transfer_params["_p_side_truncated"] is True
+
+
+def test_on_new_request_is_noop_without_mamba_or_on_decode():
+    plain = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=False)
+    req = _mk_request(list(range(5)), params={"do_remote_decode": True})
+    plain.on_new_request(req)
+    assert req.num_prompt_tokens == 5
+
+    decode_side = _FakeScheduler(
+        is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True
+    )
+    req2 = _mk_request(list(range(5)), params={"do_remote_prefill": True})
+    decode_side.on_new_request(req2)
+    assert req2.num_prompt_tokens == 5
+
+
+# --------------------------------------------------------------------------
+# Offset-template cache
+# --------------------------------------------------------------------------
+_SAMPLE_SLOT_SETS = [
+    ([1, 2], [1, 2]),
+    ([0], [3]),
+    ([3, 0, 2], [3, 0, 2]),
+    ([2, 2], [5, 5]),
+    ([], []),
+]
+
+
+def test_worker_compute_mamba_offsets_caches_and_matches_template():
+    split = _gdn_split_info()
+    conv_dim = sum(split.local_proj_dims)
+    conv_shape = (conv_dim, split.conv_rows)
+    ssm_shape = (2, 4, 4)
+    page_bytes = 256
+    pages = torch.zeros((8, 1, 1, page_bytes), dtype=torch.uint8)
+    spec = moriio_connector.MambaSpec(
+        block_size=1,
+        shapes=(conv_shape, ssm_shape),
+        dtypes=(torch.bfloat16, torch.bfloat16),
+        page_size_padded=page_bytes,
+    )
+    conv, ssm = moriio_layout.kda_conv_ssm(pages, spec)
+
+    worker = _FakeWorker(
+        kv_caches={"kda.0": pages},
+        layer_to_spec={"kda.0": spec},
+        _conv_decomp=split,
+        _mamba_offset_templates={},
+        tp_rank=0,
+        world_size=1,
+    )
+
+    cached_template = None
+    for local_slots, remote_slots in _SAMPLE_SLOT_SETS:
+        lo, ro, sz, n_conv = worker._compute_mamba_transfer_offsets(
+            "kda.0", local_slots, remote_slots
+        )
+        template = moriio_layout.build_mamba_offset_template(
+            conv, ssm, split, tp_ratio=1
+        )
+        expected = moriio_layout.apply_mamba_offset_template(
+            template, local_slots, remote_slots
+        )
+        assert (lo, ro, sz) == expected
+        assert n_conv == moriio_layout.compute_mamba_conv_split_count(
+            local_slots, split
+        )
+
+        current_template = worker._mamba_offset_templates["kda.0"]
+        if cached_template is None:
+            cached_template = current_template
+        else:
+            assert current_template is cached_template
+
+    assert "kda.0" in worker._mamba_offset_templates
+
+
+@pytest.mark.parametrize("remote_tp_size", [1, None])
+def test_mamba_tp_ratio_accepts_only_equal_tp(remote_tp_size):
+    worker = _FakeWorker(world_size=1)
+    assert worker._mamba_tp_ratio(remote_tp_size) == 1
+
+
+@pytest.mark.parametrize("remote_tp_size", [2, 3])
+def test_mamba_tp_ratio_rejects_heterogeneous_tp(remote_tp_size):
+    worker = _FakeWorker(world_size=1)
+    with pytest.raises(NotImplementedError, match="heterogeneous-TP"):
+        worker._mamba_tp_ratio(remote_tp_size)
+
+
+def test_worker_selects_each_mamba_layers_own_block_group():
+    worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 1, "kda.1": 2})
+    groups = ([10], [40], [70])
+
+    assert worker._mamba_blocks_for_layer("kda.0", groups) == [40]
+    assert worker._mamba_blocks_for_layer("kda.1", groups) == [70]
+
+
+def test_worker_rejects_missing_mamba_block_group():
+    worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 2})
+
+    with pytest.raises(moriio_common.MoRIIOError, match="no block group"):
+        worker._mamba_blocks_for_layer("kda.0", ([10], [40]))
diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py
index c9fa434ee9b9..1cb5e271bf4e 100644
--- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py
+++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import importlib.util
+import math
 import threading
 from collections import OrderedDict, defaultdict
 from queue import Queue
@@ -14,6 +15,7 @@
 from vllm.platforms import current_platform
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
+    MambaSpec,
     MLAAttentionSpec,
 )
 
@@ -36,6 +38,10 @@
     "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout"
 )
 msgpack = importlib.import_module("msgpack")
+ssm_conv_transfer_utils = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils"
+)
+MambaConvSplitInfo = ssm_conv_transfer_utils.MambaConvSplitInfo
 
 ROLE = moriio_common.ROLE
 MoRIIOError = moriio_common.MoRIIOError
@@ -353,6 +359,9 @@ class FakeWorker:
         kv_caches: dict[str, torch.Tensor]
         layer_name_to_local_kv_cache_metadata: dict[str, list[Any]]
 
+        def _region_session_indices(self, layer_name):
+            return [list(self.kv_caches).index(layer_name)]
+
         def _compute_block_transfer_offsets(
             self,
             layer_name,
@@ -673,6 +682,30 @@ def test_local_block_ids_longer_than_remote_raises_value_error():
         )
 
 
+@pytest.mark.parametrize(
+    ("local_block_ids", "remote_block_ids", "match"),
+    [
+        ([-1], [0], "local block id -1"),
+        ([0], [8], "remote block id 8"),
+    ],
+)
+def test_block_ids_must_fit_registered_regions(
+    local_block_ids, remote_block_ids, match
+):
+    cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    with pytest.raises(ValueError, match=match):
+        moriio_layout.compute_block_transfer_offsets(
+            "layer",
+            cache,
+            worker.layer_to_spec,
+            local_block_ids,
+            remote_block_ids,
+            _remote_meta(num_blocks=8).num_blocks,
+        )
+
+
 def test_empty_local_block_ids_is_free_only_noop():
     cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
     worker = _worker({"layer": cache}, {"layer": _full_spec()})
@@ -745,3 +778,245 @@ def test_unsupported_shape_raises_value_error():
 
     with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"):
         moriio_layout.get_layer_transfer_geometry("layer", cache, worker.layer_to_spec)
+
+
+# ---------------------------------------------------------------------------
+# Mamba/KDA (conv + ssm) transfer geometry
+# ---------------------------------------------------------------------------
+
+
+def _slot_strided(num_slots, per_slot_shape, slot_stride, dtype=torch.bfloat16):
+    """Build a slot-strided view: shape (num_slots, *per_slot_shape) whose slot
+    stride (dim 0) exceeds the per-slot element count, mimicking the K3 KDA
+    conv/ssm cache tensors."""
+    inner = 1
+    for d in per_slot_shape:
+        inner *= d
+    assert slot_stride >= inner
+    backing = torch.zeros(num_slots * slot_stride, dtype=dtype)
+    inner_strides: list[int] = []
+    acc = 1
+    for d in reversed(per_slot_shape):
+        inner_strides.insert(0, acc)
+        acc *= d
+    return backing.as_strided(
+        (num_slots, *per_slot_shape), (slot_stride, *inner_strides)
+    )
+
+
+def _gdn_split_info(conv_rows=3, key_dim=4, value_dim=8, dtype_size=2):
+    """A GDN-style conv split: sub-projections (key, key, value)."""
+    conv_dim = 2 * key_dim + value_dim
+    conv_state_bytes = conv_dim * conv_rows * dtype_size
+    ssm_state_bytes = 64  # opaque here; ssm size is taken from the tensor
+    return MambaConvSplitInfo(
+        conv_rows=conv_rows,
+        local_proj_dims=(key_dim, key_dim, value_dim),
+        conv_dtype_size=dtype_size,
+        ssm_sizes=(conv_state_bytes, ssm_state_bytes),
+    )
+
+
+def test_mamba_transfer_geometry_is_slot_strided():
+    split = _gdn_split_info()
+    conv_dim = sum(split.local_proj_dims)
+    conv = _slot_strided(5, (conv_dim, split.conv_rows), slot_stride=200)
+    ssm = _slot_strided(5, (2, 4, 4), slot_stride=64)
+
+    geom = moriio_layout.get_mamba_transfer_geometry(conv, ssm)
+
+    assert geom.conv_slot_stride == 200
+    assert geom.ssm_slot_stride == 64
+    assert geom.ssm_slot_bytes == (2 * 4 * 4) * ssm.element_size()
+    # Region spans the full slot-strided extent (shape[0] * stride(0)), not just
+    # numel(), so the highest-slot transfers stay in-bounds.
+    assert geom.conv_region_len == conv.shape[0] * conv.stride(0) * conv.element_size()
+    assert geom.ssm_region_len == ssm.shape[0] * ssm.stride(0) * ssm.element_size()
+
+
+def test_kda_conv_ssm_unpacks_padded_pages_as_aliasing_strided_views():
+    num_blocks = 4
+    page_bytes = 128
+    conv_shape = (3, 4)
+    ssm_shape = (2, 2, 4)
+    conv_bytes = math.prod(conv_shape) * torch.bfloat16.itemsize
+    ssm_bytes = math.prod(ssm_shape) * torch.bfloat16.itemsize
+    assert conv_bytes + ssm_bytes < page_bytes
+
+    pages = torch.zeros((num_blocks, 1, 1, page_bytes), dtype=torch.uint8)
+    spec = MambaSpec(
+        block_size=1,
+        shapes=(conv_shape, ssm_shape),
+        dtypes=(torch.bfloat16, torch.bfloat16),
+        page_size_padded=page_bytes,
+    )
+
+    conv, ssm = moriio_layout.kda_conv_ssm(pages, spec)
+
+    assert conv.shape == (num_blocks, *conv_shape)
+    assert ssm.shape == (num_blocks, *ssm_shape)
+    assert conv.stride() == (page_bytes // 2, 4, 1)
+    assert ssm.stride() == (page_bytes // 2, 8, 4, 1)
+    assert conv.untyped_storage().data_ptr() == pages.untyped_storage().data_ptr()
+    assert ssm.untyped_storage().data_ptr() == pages.untyped_storage().data_ptr()
+    assert conv.storage_offset() == 0
+    assert ssm.storage_offset() * ssm.element_size() == conv_bytes
+
+    conv[1].fill_(1)
+    ssm[2].fill_(2)
+    assert torch.all(pages[1, 0, 0, :conv_bytes].view(torch.bfloat16) == 1)
+    assert torch.all(
+        pages[2, 0, 0, conv_bytes : conv_bytes + ssm_bytes].view(torch.bfloat16) == 2
+    )
+
+
+def test_kda_conv_ssm_accepts_strided_blocks_without_copying():
+    backing = torch.zeros(4 * 256, dtype=torch.uint8)
+    pages = backing.as_strided((4, 1, 1, 128), (256, 128, 128, 1))
+    spec = MambaSpec(
+        block_size=1,
+        shapes=((3, 4), (2, 2, 4)),
+        dtypes=(torch.bfloat16, torch.bfloat16),
+        page_size_padded=128,
+    )
+
+    conv, ssm = moriio_layout.kda_conv_ssm(pages, spec)
+
+    assert not pages.is_contiguous()
+    assert conv.stride(0) == 128
+    assert ssm.stride(0) == 128
+    assert conv.untyped_storage().data_ptr() == backing.data_ptr()
+    assert ssm.untyped_storage().data_ptr() == backing.data_ptr()
+
+
+def test_kda_conv_ssm_rejects_noncontiguous_page_bytes():
+    pages = torch.zeros((4, 1, 1, 256), dtype=torch.uint8)[..., ::2]
+    spec = MambaSpec(
+        block_size=1,
+        shapes=((3, 4), (2, 2, 4)),
+        dtypes=(torch.bfloat16, torch.bfloat16),
+        page_size_padded=128,
+    )
+
+    with pytest.raises(ValueError, match="contiguous byte dimension"):
+        moriio_layout.kda_conv_ssm(pages, spec)
+
+
+def test_kda_conv_ssm_rejects_non_byte_packed_cache():
+    pages = torch.zeros((4, 1, 1, 128), dtype=torch.bfloat16)
+    spec = MambaSpec(
+        block_size=1,
+        shapes=((3, 4), (2, 2, 4)),
+        dtypes=(torch.bfloat16, torch.bfloat16),
+        page_size_padded=128,
+    )
+
+    with pytest.raises(ValueError, match="byte-sized dtype"):
+        moriio_layout.kda_conv_ssm(pages, spec)
+
+
+def test_mamba_transfer_geometry_requires_matching_slot_counts():
+    conv = torch.empty((5, 4, 3), dtype=torch.bfloat16)
+    ssm = torch.empty((4, 2, 4, 4), dtype=torch.bfloat16)
+
+    with pytest.raises(ValueError, match="conv/ssm slot count mismatch"):
+        moriio_layout.get_mamba_transfer_geometry(conv, ssm)
+
+
+def test_mamba_transfer_geometry_rejects_empty_caches():
+    conv = torch.empty((0, 4, 3), dtype=torch.bfloat16)
+    ssm = torch.empty((0, 2, 4, 4), dtype=torch.bfloat16)
+
+    with pytest.raises(ValueError, match="at least one slot"):
+        moriio_layout.get_mamba_transfer_geometry(conv, ssm)
+
+
+def test_mamba_conv_ssm_offsets_homogeneous_tp():
+    split = _gdn_split_info()
+    conv_dim = sum(split.local_proj_dims)
+    conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200)
+    ssm = _slot_strided(6, (2, 4, 4), slot_stride=64)
+    geom = moriio_layout.get_mamba_transfer_geometry(conv, ssm)
+
+    local_slots = [1, 2]
+    remote_slots = [1, 2]
+    template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1)
+    local_offs, remote_offs, sizes = moriio_layout.apply_mamba_offset_template(
+        template, local_slots, remote_slots
+    )
+
+    n_conv = moriio_layout.compute_mamba_conv_split_count(local_slots, split)
+    assert n_conv == len(local_slots) * len(split.local_proj_dims)
+    assert len(local_offs) == n_conv + len(local_slots)
+
+    # Homogeneous TP: remote conv slices equal local slices; since local==remote
+    # slots, conv local and remote offsets coincide.
+    loc_co = split.local_conv_offsets
+    rem_co = split.remote_conv_offsets(0, 1)
+    k = 0
+    esz = conv.element_size()
+    for ls, rs in zip(local_slots, remote_slots):
+        lbase = ls * geom.conv_slot_stride * esz
+        rbase = rs * geom.conv_slot_stride * esz
+        for (loff, _), (roff, rsz) in zip(loc_co, rem_co):
+            assert local_offs[k] == lbase + loff
+            assert remote_offs[k] == rbase + roff
+            assert sizes[k] == rsz
+            k += 1
+
+    # SSM: whole per-slot block, no head sharding at tp_ratio == 1.
+    essz = ssm.element_size()
+    for j, (ls, rs) in enumerate(zip(local_slots, remote_slots)):
+        assert local_offs[n_conv + j] == ls * geom.ssm_slot_stride * essz
+        assert remote_offs[n_conv + j] == rs * geom.ssm_slot_stride * essz
+        assert sizes[n_conv + j] == geom.ssm_slot_bytes
+
+
+def test_mamba_conv_ssm_offsets_heterogeneous_tp_is_gated():
+    # Heterogeneous-TP mamba transfer is a documented follow-up (remote slot
+    # stride + multi-rank fan-in gather). Until then the offset helper must fail
+    # loudly rather than silently transfer wrong bytes, since the mamba path
+    # bypasses the attention KV-head validator.
+    split = _gdn_split_info()
+    conv_dim = sum(split.local_proj_dims)
+    conv = _slot_strided(4, (conv_dim, split.conv_rows), slot_stride=200)
+    ssm = _slot_strided(4, (2, 4, 4), slot_stride=64)
+    with pytest.raises(NotImplementedError):
+        moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=2)
+
+
+def test_local_mamba_slots_must_fit_registered_regions():
+    split = _gdn_split_info()
+    conv_dim = sum(split.local_proj_dims)
+    conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200)
+    ssm = _slot_strided(6, (2, 4, 4), slot_stride=64)
+    template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1)
+
+    with pytest.raises(ValueError, match="local mamba slot 6"):
+        moriio_layout.apply_mamba_offset_template(template, [6], [0])
+
+
+@pytest.mark.parametrize(
+    ("local_slots", "remote_slots"),
+    [([1, 2], [3]), ([1], [2, 3])],
+)
+def test_mamba_slot_counts_must_match(local_slots, remote_slots):
+    split = _gdn_split_info()
+    conv_dim = sum(split.local_proj_dims)
+    conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200)
+    ssm = _slot_strided(6, (2, 4, 4), slot_stride=64)
+    template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1)
+
+    with pytest.raises(ValueError, match="slot counts must match"):
+        moriio_layout.apply_mamba_offset_template(template, local_slots, remote_slots)
+
+
+def test_remote_mamba_slots_must_be_non_negative():
+    split = _gdn_split_info()
+    conv_dim = sum(split.local_proj_dims)
+    conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200)
+    ssm = _slot_strided(6, (2, 4, 4), slot_stride=64)
+    template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1)
+
+    with pytest.raises(ValueError, match="must be non-negative"):
+        moriio_layout.apply_mamba_offset_template(template, [1], [-1])
diff --git a/tests/v1/kv_connector/unit/test_moriio_routing.py b/tests/v1/kv_connector/unit/test_moriio_routing.py
index a9a3521bebd7..b0ac6df3cba3 100644
--- a/tests/v1/kv_connector/unit/test_moriio_routing.py
+++ b/tests/v1/kv_connector/unit/test_moriio_routing.py
@@ -20,12 +20,16 @@
 from vllm.v1.request import RequestStatus
 
 _request_finished = MoRIIOConnectorScheduler.request_finished
+_clip_blocks = MoRIIOConnectorScheduler.get_exchange_clipped_blocks
 
 
 def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int):
     """Minimal prefill-side (producer) scheduler stand-in for request_finished."""
-    return SimpleNamespace(
+    s = SimpleNamespace(
         is_producer=True,
+        _is_hma_required=False,
+        _has_mamba=False,
+        blocks_per_sw=[0],
         _global_dp_rank=global_dp_rank,
         engine_id="engine-abc",
         host_ip="10.0.0.1",
@@ -39,6 +43,8 @@ def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int):
             )
         ),
     )
+    s.get_exchange_clipped_blocks = _clip_blocks.__get__(s)
+    return s
 
 
 def _req(
diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py
index da412fcd9290..a35f0a1169b6 100644
--- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py
+++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py
@@ -2,12 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import threading
+from types import SimpleNamespace
 
 import pytest
 
 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
     MoRIIOMode,
     MoRIIOTransferAck,
+    TransferError,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
     MoRIIOConnector,
@@ -17,6 +19,9 @@
     resolve_moriio_transfer_ack,
     validate_moriio_heterogeneous_tp_kv_heads,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
+    MoRIIOWrapper,
+)
 
 
 def test_remote_tp_rank_same_tp_maps_to_self():
@@ -277,6 +282,9 @@ def __init__(self):
             self.lock = threading.Lock()
             self.sent = []
 
+        # Exercise the real batch verdict rather than restating it here.
+        poll_transfer_batch = MoRIIOWrapper.poll_transfer_batch
+
         def send_notify(
             self,
             transfer_id,
@@ -292,13 +300,18 @@ def shutdown(self):
 
     worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
     worker.world_size = 8
+    worker.moriio_config = SimpleNamespace(recv_abort_timeout=600.0)
     worker.moriio_wrapper = FakeWrapper()
-    worker._recving_transfers = {"req": {"layer0": DoneStatus()}}
+    # A layer maps to the list of reads posted for it (a KDA layer posts two).
+    worker._recving_transfers = {"req": {"layer0": [DoneStatus(), DoneStatus()]}}
     worker._recving_transfers_callback_addr = {
         "req": ("127.0.0.1", "7000", "tx-release")
     }
     # Transfer-timeout reaping state consulted by _pop_done_transfers.
     worker._recving_transfers_start = {}
+    # Load-error bookkeeping cleared alongside a completed transfer.
+    worker._recving_local_blocks = {}
+    worker._invalid_block_ids = set()
 
     assert worker._pop_done_transfers() == {"tx-release"}
     assert worker.moriio_wrapper.sent == [
@@ -314,6 +327,133 @@ def shutdown(self):
     assert worker._recving_transfers_callback_addr == {}
 
 
+def test_aborted_read_without_local_blocks_only_releases_remote_blocks():
+    class FakeWrapper:
+        def __init__(self):
+            self.sent = []
+
+        def send_notify(
+            self,
+            transfer_id,
+            host,
+            port,
+            message_type=None,
+            message_fields=None,
+        ):
+            self.sent.append((transfer_id, host, port, message_type, message_fields))
+
+        def shutdown(self):
+            pass
+
+    worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
+    worker.mode = MoRIIOMode.READ
+    worker.world_size = 8
+    worker.moriio_wrapper = FakeWrapper()
+
+    worker._read_blocks(
+        local_block_ids=[],
+        remote_block_ids=[[10, 11], [90]],
+        dst_engine_id="prefill",
+        request_id="req-aborted",
+        transfer_id="tx-aborted",
+        remote_host="127.0.0.1",
+        remote_notify_port=7000,
+        remote_tp_size=8,
+        remote_dp_rank=0,
+        chosen_tp=2,
+    )
+
+    assert worker.moriio_wrapper.sent == [
+        (
+            "tx-aborted",
+            "127.0.0.1",
+            "7002",
+            "release",
+            {"consumer_tp_size": 8},
+        )
+    ]
+
+
+def test_read_completion_waits_for_every_posted_read():
+    class Status:
+        def __init__(self, done: bool):
+            self.done = done
+
+        def Succeeded(self):
+            return self.done
+
+        def Failed(self):
+            return False
+
+    class FakeWrapper:
+        def __init__(self):
+            self.lock = threading.Lock()
+            self.sent = []
+
+        poll_transfer_batch = MoRIIOWrapper.poll_transfer_batch
+
+        def send_notify(self, *args, **kwargs):
+            self.sent.append((args, kwargs))
+
+        def shutdown(self):
+            pass
+
+    worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
+    worker.world_size = 8
+    worker.moriio_config = SimpleNamespace(recv_abort_timeout=600.0)
+    worker.moriio_wrapper = FakeWrapper()
+    worker._recving_transfers = {
+        "req": {"kda_layer": [Status(done=True), Status(done=False)]}
+    }
+    worker._recving_transfers_callback_addr = {
+        "req": ("127.0.0.1", "7000", "tx-pending")
+    }
+    worker._recving_transfers_start = {"req": float("inf")}
+    worker._recving_local_blocks = {"req": [1, 2]}
+    worker._invalid_block_ids = set()
+
+    assert worker._pop_done_transfers() == set()
+    assert "req" in worker._recving_transfers
+    assert worker.moriio_wrapper.sent == []
+
+
+@pytest.mark.parametrize(
+    ("has_mamba", "expected_invalid"),
+    [(False, {1, 2}), (True, set())],
+)
+def test_failed_read_reports_blocks_only_without_hma(has_mamba, expected_invalid):
+    worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
+    worker._has_mamba = has_mamba
+    worker._recving_local_blocks = {"req": [1, 2]}
+    worker._invalid_block_ids = set()
+
+    worker._record_failed_recv("req")
+
+    assert worker.get_block_ids_with_load_errors() == expected_invalid
+
+
+def test_hybrid_step_barrier_fails_closed(monkeypatch):
+    class FailingWrapper:
+        def waiting_for_transfer_complete(self, _statuses):
+            raise TransferError("failed")
+
+        def shutdown(self):
+            pass
+
+    worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
+    worker._has_mamba = True
+    worker._reads_issued_this_step = [object()]
+    worker._mamba_reads_this_step = [object()]
+    worker.moriio_wrapper = FailingWrapper()
+    monkeypatch.setattr(
+        "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.get_forward_context",
+        lambda: SimpleNamespace(cudagraph_runtime_mode=None),
+    )
+
+    with pytest.raises(TransferError, match="failed"):
+        worker._await_reads_issued_this_step()
+
+
 def test_requested_cudagraph_mode_is_never_overridden():
     # The configured cudagraph mode is always honored: the barrier fires when
     # the operator sets cudagraph_mode=PIECEWISE, and READ mode with full
diff --git a/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py
new file mode 100644
index 000000000000..2a3092ad5ae9
--- /dev/null
+++ b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Transfer-completion semantics of the MoRIIO wrapper.
+
+Covers the two behaviours the READ barrier depends on:
+  * poll_transfer_batch considers EVERY status, not just the newest one
+  * waiting_for_transfer_complete blocks until terminal and raises on failure.
+"""
+
+import importlib
+import importlib.util
+import threading
+import time
+
+import pytest
+
+from vllm.platforms import current_platform
+
+mori_available = importlib.util.find_spec("mori") is not None
+
+if not (current_platform.is_rocm() and mori_available):
+    pytest.skip(
+        "MoRIIOs are only available on ROCm with mori package installed",
+        allow_module_level=True,
+    )
+
+moriio_common = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common"
+)
+moriio_engine = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine"
+)
+
+MoRIIOWrapper = moriio_engine.MoRIIOWrapper
+TransferBatchState = moriio_common.TransferBatchState
+TransferError = moriio_common.TransferError
+
+
+@pytest.mark.parametrize("value", [0, -1, float("inf"), float("nan")])
+def test_timeout_must_be_positive_and_finite(value):
+    with pytest.raises(ValueError, match="must be finite and greater than zero"):
+        moriio_common._positive_finite_timeout("test_timeout", value)
+
+
+def test_positive_finite_timeout_is_preserved():
+    assert moriio_common._positive_finite_timeout("test_timeout", "1.5") == 1.5
+
+
+class FakeStatus:
+    """Mirrors the mori TransferStatus surface the wrapper uses."""
+
+    def __init__(self, state: str = "in_progress", message: str = ""):
+        self.state = state
+        self.message = message
+
+    def Succeeded(self) -> bool:
+        return self.state == "success"
+
+    def Failed(self) -> bool:
+        return self.state == "failed"
+
+    def Message(self) -> str:
+        return self.message
+
+    def Code(self) -> str:
+        return self.state
+
+
+def _ok():
+    return FakeStatus("success")
+
+
+def _pending():
+    return FakeStatus("in_progress")
+
+
+def _bad():
+    return FakeStatus("failed", "SQ full")
+
+
+@pytest.mark.parametrize(
+    "statuses,expected",
+    [
+        ([], TransferBatchState.DONE),
+        ([_ok(), _ok()], TransferBatchState.DONE),
+        # Each case below puts the interesting status BEFORE a success, which a
+        # statuses[-1] check would have missed entirely.
+        ([_bad(), _ok()], TransferBatchState.FAILED),
+        ([_pending(), _ok()], TransferBatchState.PENDING),
+        ([_pending(), _bad(), _ok()], TransferBatchState.FAILED),
+    ],
+)
+def test_poll_transfer_batch_considers_every_status(statuses, expected):
+    wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05)
+    assert wrapper.poll_transfer_batch(statuses) is expected
+
+
+@pytest.mark.parametrize(
+    "statuses_factory,expected_message",
+    [
+        (lambda: [_ok(), _bad()], "SQ full"),
+        (lambda: [_ok(), _pending()], "timed out"),
+    ],
+    ids=["failed", "timeout"],
+)
+def test_wait_raises_with_detail(statuses_factory, expected_message):
+    wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05)
+    with pytest.raises(TransferError, match=expected_message):
+        wrapper.waiting_for_transfer_complete(statuses_factory())
+
+
+def test_wait_returns_immediately_when_nothing_pending():
+    wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05)
+    wrapper.waiting_for_transfer_complete([])
+    wrapper.waiting_for_transfer_complete([_ok(), _ok()])
+
+
+def test_wait_blocks_until_terminal():
+    """The barrier is only worth anything if it actually waits."""
+    wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=5.0)
+    late = FakeStatus("in_progress")
+    timer = threading.Timer(0.2, lambda: setattr(late, "state", "success"))
+    timer.start()
+    try:
+        start = time.monotonic()
+        wrapper.waiting_for_transfer_complete([late])
+        assert time.monotonic() - start >= 0.15
+    finally:
+        timer.cancel()
diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py
index bac0b0373807..1a85865560ab 100644
--- a/tests/v1/kv_connector/unit/test_multi_connector.py
+++ b/tests/v1/kv_connector/unit/test_multi_connector.py
@@ -20,7 +20,10 @@
     SupportsHMA,
     supports_hma,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+    KVConnectorLogging,
+    KVConnectorStats,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
     MultiConnector,
     MultiKVConnectorPromMetrics,
@@ -171,7 +174,7 @@ def test_register_finished_partial_tail_notifies_every_connector():
 
 @pytest.fixture
 def mc() -> MultiConnector:
-    """MultiConnector using two mocked connectors"""
+    """MultiConnector using two mocked connectors."""
     mock_connector_config = {
         "kv_connector": "MockConnector",
         "kv_role": "kv_both",
@@ -215,8 +218,7 @@ def _compare_directories(dir1: Path, dir2: Path) -> bool:
 
 
 def test_multi_example_connector_consistency():
-    """
-    Tests that MultiConnector with two ExampleConnectors saves
+    """Tests that MultiConnector with two ExampleConnectors saves
     identical KV cache data to separate storage locations.
     """
     storage_1_path = Path("storage_1/")
@@ -304,10 +306,10 @@ def test_multi_example_connector_consistency():
     events = get_connector_events()
     storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
     storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
-    # Initial events bind the block pool, query completion counts, and exchange
+    # Initial events bind the cache manager, query completion counts, and exchange
     # handshake metadata before the request is enqueued.
     assert storage1_scheduler_events[:7] == [
-        "bind_gpu_block_pool",
+        "bind_kv_cache_manager",
         "get_finished_count",
         "set_xfer_handshake_metadata_pp_aware",
         "on_new_request",
@@ -315,8 +317,8 @@ def test_multi_example_connector_consistency():
         "update_state_after_alloc num_blocks=[7] 0",
         "build_connector_meta",
     ]
-    # First three events are from initialization. During generate(), layer hooks
-    # run before the deferred load is started after the forward pass.
+    # First three events are from initialization. Layer hooks run before the
+    # deferred load starts after the forward pass.
     expected_worker_prefix = [
         "register_kv_caches",
         "set_host_xfer_buffer_ops",
@@ -328,12 +330,12 @@ def test_multi_example_connector_consistency():
     ]
     for connector_name in ("storage1-WORKER", "storage2-WORKER"):
         worker_events = events[connector_name]
-        assert worker_events[:7] == expected_worker_prefix
+        assert worker_events[: len(expected_worker_prefix)] == expected_worker_prefix
         assert worker_events.index("start_load_kv") > worker_events.index(
             "save_kv_layer"
         )
     assert storage2_scheduler_events[:7] == [
-        "bind_gpu_block_pool",
+        "bind_kv_cache_manager",
         "get_finished_count",
         "set_xfer_handshake_metadata_pp_aware",
         "on_new_request",
@@ -455,8 +457,7 @@ def test_engine_id_conflict():
 
 
 def test_multi_connector_handle_preemptions_integration():
-    """
-    Integration test: verify MultiConnector delegates handle_preemptions
+    """Integration test: verify MultiConnector delegates handle_preemptions
     to all sub-connectors.
 
     Uses TestExampleConnector which logs all method calls to temp files.
@@ -566,6 +567,8 @@ def test_build_kv_connector_stats_reconstructs_nixl_stats(self):
                 "num_descriptors": [10, 20],
                 "num_failed_transfers": [],
                 "num_failed_notifications": [],
+                "num_failed_handshakes": [],
+                "num_kv_expired_reqs": [],
             }
         }
 
@@ -589,6 +592,8 @@ def test_build_kv_connector_stats_with_multiple_connectors(self):
                 "num_descriptors": [10],
                 "num_failed_transfers": [],
                 "num_failed_notifications": [],
+                "num_failed_handshakes": [],
+                "num_kv_expired_reqs": [],
             },
             "MockConnector": {"mock_field": [1, 2, 3]},
         }
@@ -618,6 +623,8 @@ def test_build_kv_connector_stats_raises_error_for_unknown_connector(self):
                 "num_descriptors": [10],
                 "num_failed_transfers": [],
                 "num_failed_notifications": [],
+                "num_failed_handshakes": [],
+                "num_kv_expired_reqs": [],
             },
         }
 
@@ -637,6 +644,8 @@ def test_build_kv_connector_stats_with_already_instantiated_objects(self):
                 "num_descriptors": [10],
                 "num_failed_transfers": [],
                 "num_failed_notifications": [],
+                "num_failed_handshakes": [],
+                "num_kv_expired_reqs": [],
             }
         )
         mock_stats = MockConnectorStats(data={"mock_field": [1, 2, 3]})
@@ -666,6 +675,8 @@ def test_build_kv_connector_stats_with_mixed_objects_and_dicts(self):
                 "num_descriptors": [10],
                 "num_failed_transfers": [],
                 "num_failed_notifications": [],
+                "num_failed_handshakes": [],
+                "num_kv_expired_reqs": [],
             }
         )
 
@@ -697,6 +708,8 @@ def test_build_kv_connector_stats_skips_connectors_without_custom_stats(self):
                 "num_descriptors": [10],
                 "num_failed_transfers": [],
                 "num_failed_notifications": [],
+                "num_failed_handshakes": [],
+                "num_kv_expired_reqs": [],
             },
             "ExampleConnector": {"some_field": [1, 2, 3]},
         }
@@ -734,6 +747,8 @@ def test_aggregate_same_connector(self):
                         "num_descriptors": [10],
                         "num_failed_transfers": [],
                         "num_failed_notifications": [],
+                        "num_failed_handshakes": [],
+                        "num_kv_expired_reqs": [],
                     }
                 )
             }
@@ -749,6 +764,8 @@ def test_aggregate_same_connector(self):
                         "num_descriptors": [20],
                         "num_failed_transfers": [],
                         "num_failed_notifications": [],
+                        "num_failed_handshakes": [],
+                        "num_kv_expired_reqs": [],
                     }
                 )
             }
@@ -780,6 +797,8 @@ def test_aggregate_new_connector(self):
                         "num_descriptors": [10],
                         "num_failed_transfers": [],
                         "num_failed_notifications": [],
+                        "num_failed_handshakes": [],
+                        "num_kv_expired_reqs": [],
                     }
                 )
             }
@@ -806,6 +825,8 @@ def test_reduce(self):
                         "num_descriptors": [10, 20],
                         "num_failed_transfers": [],
                         "num_failed_notifications": [],
+                        "num_failed_handshakes": [],
+                        "num_kv_expired_reqs": [],
                     }
                 )
             }
@@ -819,6 +840,38 @@ def test_reduce(self):
         assert "Num successful transfers" in reduced["NixlConnector"]
         assert reduced["NixlConnector"]["Num successful transfers"] == 2
 
+    def test_log_renders_plain_scalars(self):
+        """Reduced stats must render as plain numbers in the CLI log, not
+        numpy reprs (eg np.float64(2.338))."""
+        stats = MultiKVConnectorStats(
+            data={
+                "NixlConnector": NixlKVConnectorStats(
+                    data={
+                        "transfer_duration": [1.0, 2.0],
+                        "post_duration": [0.1, 0.2],
+                        "bytes_transferred": [1024, 2048],
+                        "num_descriptors": [10, 20],
+                        "num_failed_transfers": [],
+                        "num_failed_notifications": [],
+                        "num_failed_handshakes": [],
+                        "num_kv_expired_reqs": [],
+                    }
+                )
+            }
+        )
+
+        kv_logging = KVConnectorLogging(kv_transfer_config=None)
+        kv_logging.transfer_stats_accumulator = stats
+        log_records: list[tuple] = []
+        kv_logging.log(log_fn=lambda *args: log_records.append(args))
+
+        assert len(log_records) == 1
+        msg = log_records[0][1]
+        assert "np.float" not in msg
+        assert "'Avg xfer time (ms)': 1500.0" in msg
+        # The accumulator is reset after logging.
+        assert kv_logging.transfer_stats_accumulator is None
+
     def test_reset(self):
         """Test that reset() resets all nested connector stats."""
         stats = MultiKVConnectorStats(
@@ -831,6 +884,8 @@ def test_reset(self):
                         "num_descriptors": [10, 20],
                         "num_failed_transfers": [],
                         "num_failed_notifications": [],
+                        "num_failed_handshakes": [],
+                        "num_kv_expired_reqs": [],
                     }
                 )
             }
@@ -863,9 +918,7 @@ def test_is_empty_with_multiple_connectors(self):
 
 
 def test_multi_connector_overrides_all_base_methods():
-    """
-    Ensure MultiConnector overrides all public methods from KVConnectorBase_V1.
-    """
+    """Ensure MultiConnector overrides all public methods from KVConnectorBase_V1."""
     # These are fine to inherit from KVConnectorBase_V1
     # TODO(https://github.com/vllm-project/vllm/pull/31811): Remove
     # get_kv_connector_kv_cache_events from INHERITED_OK once implemented.
@@ -1045,12 +1098,10 @@ def _make_multi_connector(connector_names: list[str]) -> MultiConnector:
 
 
 def test_multi_connector_hma_support_detection():
-    """
-    At runtime, _all_support_hma is True only when every sub-connector
+    """At runtime, _all_support_hma is True only when every sub-connector
     implements SupportsHMA. Test all combinations of HMA / non-HMA
     sub-connectors.
     """
-
     assert supports_hma(MultiConnector)
 
     # -- All non-HMA connectors => _all_support_hma is False --
@@ -1092,8 +1143,7 @@ def test_divergent_local_hybrid_hit_capability_is_conservative():
     not torch.cuda.is_available(), reason="Requires GPU to instantiate LLM"
 )
 def test_multi_connector_mixed_hma_disables_hybrid_kv_cache(monkeypatch):
-    """
-    When MultiConnector wraps a mix of HMA (NixlConnector) and non-HMA
+    """When MultiConnector wraps a mix of HMA (NixlConnector) and non-HMA
     (MockConnector) sub-connectors, verify that:
     1. The scheduler's MultiConnector has _all_support_hma == False.
     2. vLLM auto-disables the hybrid KV cache manager (no preference expressed by user)
diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py
index 760a21c4c85d..39f8f88d5c70 100644
--- a/tests/v1/kv_connector/unit/test_nixl_connector.py
+++ b/tests/v1/kv_connector/unit/test_nixl_connector.py
@@ -4,11 +4,13 @@
 import contextlib
 import inspect
 import os
+import queue
 import tempfile
 import textwrap
 import time
 import uuid
 from collections import defaultdict
+from concurrent.futures import Future
 from types import SimpleNamespace
 from typing import Any, cast
 from unittest.mock import MagicMock, patch
@@ -29,7 +31,10 @@
     get_current_attn_backend,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1 import nixl
-from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorRole,
+    KVConnectorTransferResults,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
     MultiKVConnectorStats,
@@ -70,6 +75,7 @@
 from vllm.v1.request import RequestStatus
 
 from .utils import (
+    create_model_runner_output,
     create_request,
     create_scheduler,
     create_vllm_config,
@@ -79,8 +85,7 @@
 
 @pytest.fixture(scope="module", autouse=True)
 def clear_kv_transfer():
-    """
-    The test cases in this file use `VLLM_ENABLE_V1_MULTIPROCESSING=0`,
+    """The test cases in this file use `VLLM_ENABLE_V1_MULTIPROCESSING=0`,
     causing the global variable `_KV_CONNECTOR_AGENT`
     to be assigned but never deleted.
 
@@ -240,7 +245,6 @@ def _make_fake_nixl_pkg():
 
 def test_basic_interface():
     """Unit test for basic NixlConnector interface functionality."""
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
@@ -279,8 +283,7 @@ def test_basic_interface():
 
 
 def test_prompt_less_than_block_size():
-    """
-    Test that we can handle case where prompt is < block.
+    """Test that we can handle case where prompt is < block.
 
     In this case, the P worker will still send remote_block_ids of the
     partial block. The D worker should schedule an async read
@@ -338,6 +341,63 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv():
     assert req_meta.remote.request_id == f"prefill-{42}"
     # do_remote_prefill is consumed by request_finished to prevent re-issuing.
     assert request.kv_transfer_params["do_remote_prefill"] is False
+    # The scheduler is not waiting on this recv -- the request is already gone
+    # from self.requests, so reporting it would trip `assert req_id in
+    # self.requests` in _update_from_kv_xfer_finished.
+    assert req_meta.awaiting_kvs is False
+
+
+def test_prefill_exports_cached_tokens_in_kv_transfer_params():
+    """The P worker reports its own prefix-cache hits in the returned
+    kv_transfer_params so the D worker can surface them in
+    prompt_tokens_details instead of the ~100% local hit it measures
+    when pulling the KVs from the remote.
+    """
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_TOKENS = BLOCK_SIZE * 3
+
+    # Warm the prefix cache with a request sharing the full prompt.
+    warmup = create_request(
+        request_id=1,
+        num_tokens=NUM_TOKENS,
+        common_prefix_len=NUM_TOKENS,
+        block_size=BLOCK_SIZE,
+    )
+    scheduler.add_request(warmup)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output([warmup], use_eos=True)
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # P-side request with the same prompt hits the local cache on all but
+    # the last block (recomputed to obtain logits).
+    request = create_request(
+        request_id=2,
+        num_tokens=NUM_TOKENS,
+        common_prefix_len=NUM_TOKENS,
+        block_size=BLOCK_SIZE,
+        do_remote_decode=True,
+    )
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    assert request.prefill_stats is not None
+    assert request.prefill_stats.num_cached_tokens == NUM_TOKENS - BLOCK_SIZE
+
+    # max_tokens=1, so the request finishes and returns kv_transfer_params.
+    model_runner_output = create_model_runner_output([request])
+    engine_core_outputs = scheduler.update_from_output(
+        scheduler_output, model_runner_output
+    )
+
+    output = engine_core_outputs[0].outputs[0]
+    assert output.finish_reason is not None
+    assert output.kv_transfer_params is not None
+    assert (
+        output.kv_transfer_params["remote_prefill_cached_tokens"]
+        == NUM_TOKENS - BLOCK_SIZE
+    )
 
 
 @patch(
@@ -346,7 +406,6 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv():
 )
 def test_kv_transfer_handshake(dist_init):
     """Unit test for basic NixlConnector interface functionality."""
-
     # Test setup, we creates a scheduler that contains a NixlConnector
     # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from
     # all workers of the instance.
@@ -518,6 +577,7 @@ def _nixl_handshake(
         slot_size_bytes = 4096
         self.slot_size_per_layer = [slot_size_bytes]
         self.block_len_per_layer = [slot_size_bytes * self.block_size]
+        self.num_blocks = self.kv_cache_config.num_blocks
         self.num_regions = 1
         self.block_stride_per_layer = list(self.block_len_per_layer)
         self.region_num_blocks = [self.num_blocks]
@@ -576,19 +636,38 @@ def _nixl_handshake(
 
 
 class TestNixlHandshake:
-    @pytest.mark.parametrize("pcp_rank", [0, 1])
+    @pytest.mark.parametrize(
+        ("pcp_rank", "pcp_size", "dcp_size", "expected_tracked"),
+        [
+            (0, 2, 1, True),
+            (1, 2, 1, False),
+            (0, 2, 2, True),
+            (1, 2, 2, True),
+            (0, 4, 4, True),
+            (1, 4, 4, True),
+            (2, 4, 4, True),
+            (3, 4, 4, True),
+        ],
+    )
     @patch(
         "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
         FakeNixlWrapper,
     )
-    def test_pcp_producer_uses_canonical_replica(
-        self, default_vllm_config, dist_init, pcp_rank
+    def test_pcp_producer_exposes_dcp_shards_or_canonical_replica(
+        self,
+        default_vllm_config,
+        dist_init,
+        pcp_rank,
+        pcp_size,
+        dcp_size,
+        expected_tracked,
     ):
-        """Only PCP rank zero publishes, but every rank reports completion."""
+        """Replicated PCP is canonicalized; PCP-DCP publishes every shard."""
         from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
 
         vllm_config = create_vllm_config(kv_role="kv_producer")
-        vllm_config.parallel_config.prefill_context_parallel_size = 2
+        vllm_config.parallel_config.prefill_context_parallel_size = pcp_size
+        vllm_config.parallel_config.decode_context_parallel_size = dcp_size
         with (
             patch(
                 "vllm.distributed.kv_transfer.kv_connector.v1.nixl."
@@ -610,13 +689,15 @@ def test_pcp_producer_uses_canonical_replica(
         worker = connector.connector_worker
         assert worker is not None
         assert worker.pcp_rank == pcp_rank
+        assert worker.pcp_dcp_sharded is (dcp_size > 1)
+        assert worker.transfer_tp_rank == (pcp_rank if dcp_size > 1 else 0)
+        assert worker.transfer_tp_size == (pcp_size if dcp_size > 1 else 1)
 
         req_id = "req"
         metadata = NixlConnectorMetadata()
         metadata.reqs_in_batch.add(req_id)
         metadata.reqs_to_send[req_id] = time.perf_counter() + 10
         worker.start_load_kv(metadata)
-        expected_tracked = pcp_rank == 0
         assert (req_id in worker._reqs_to_process) == expected_tracked
         assert (req_id in worker._reqs_to_send) == expected_tracked
 
@@ -624,17 +705,23 @@ def test_pcp_producer_uses_canonical_replica(
         worker.xfer_handshake_metadata = payload
         worker.transfer_topo = MagicMock()
         worker._get_new_notifs = MagicMock(
-            side_effect=lambda: {"sent"} if pcp_rank == 0 else set()
+            side_effect=lambda: {"sent"} if expected_tracked else set()
         )
 
-        expected_payload = payload if pcp_rank == 0 else None
+        expected_payload = payload if expected_tracked else None
         assert connector.get_handshake_metadata() is expected_payload
         done_sending, done_recving = connector.get_finished(set())
-        assert done_sending == ({"sent"} if pcp_rank == 0 else {req_id})
+        assert done_sending == ({"sent"} if expected_tracked else {req_id})
         assert done_recving == set()
-        if pcp_rank > 0:
+        if not expected_tracked:
             assert connector.get_finished(set()) == (set(), set())
 
+        worker.get_transfer_results = MagicMock(
+            return_value=KVConnectorTransferResults(finished_sending={"sent"})
+        )
+        results = connector.get_transfer_results(set())
+        assert results.finished_sending == ({"sent"} if expected_tracked else set())
+
     @patch(
         "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
         FakeNixlWrapper,
@@ -741,7 +828,6 @@ def test_async_load_kv(
         prefill_tp_size,
     ):
         """Test that NixlConnector's start_load_kv should be non-blocking."""
-
         vllm_config = create_vllm_config()
         vllm_config.parallel_config.tensor_parallel_size = decode_tp_size
 
@@ -796,8 +882,7 @@ def test_async_load_kv(
     def test_prefill_tp_size_greater_than_decode_tp_size(
         self, local_tp_size: int, default_vllm_config, dist_init, monkeypatch
     ):
-        """
-        Verify remote TP > local TP handshake succeeds with different
+        """Verify remote TP > local TP handshake succeeds with different
         remote configurations.
         """
         monkeypatch.setattr(
@@ -870,8 +955,7 @@ def check_handshake(remote_tp_size: int):
     def test_prefill_tp_size_greater_than_decode_tp_size_mla(
         self, default_vllm_config, dist_init
     ):
-        """
-        Verify remote TP > local TP handshake succeeds with different
+        """Verify remote TP > local TP handshake succeeds with different
         remote configurations for an MLA model.
         """
         vllm_config = create_vllm_config()
@@ -977,7 +1061,6 @@ def test_concurrent_load_kv(
         dist_init,
     ):
         """Test that multiple start_load_kv calls should occur concurrently."""
-
         vllm_config = create_vllm_config()
 
         # Test worker role in decode server.
@@ -1038,8 +1121,7 @@ def test_concurrent_load_kv(
     def test_handshake_fails_on_kv_cache_layout_mismatch(
         self, default_vllm_config, dist_init
     ):
-        """
-        Verify that adding a remote agent fails if kv_cache_layout differs.
+        """Verify that adding a remote agent fails if kv_cache_layout differs.
         This test is only relevant for heterogeneous TP.
         """
         vllm_config = create_vllm_config()
@@ -1096,8 +1178,7 @@ def test_handshake_fails_on_kv_cache_layout_mismatch(
     def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental(
         self, default_vllm_config, dist_init
     ):
-        """
-        Verify that adding a remote agent fails if kv_cache_layout differs.
+        """Verify that adding a remote agent fails if kv_cache_layout differs.
         This test is only relevant for heterogeneous TP.
         """
         vllm_config = create_vllm_config(enable_permute_local_kv=True)
@@ -1195,6 +1276,7 @@ def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices(
             ssm_sizes=(0, 0),
             attn_backend_name=worker.backend_name,
             physical_blocks_per_logical_kv_block=1,
+            region_num_blocks=None,
         )
 
         assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
@@ -1435,9 +1517,9 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
     metadata = NixlConnectorMetadata()
     metadata.add_new_req_to_recv(
         request_id=request_id,
-        local_block_ids=([1, 2, 3],),
+        local_block_ids=([0],),
         kv_transfer_params={
-            "remote_block_ids": ([4, 5, 6],),
+            "remote_block_ids": ([0],),
             "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
             "remote_request_id": f"prefill-{request_id}",
             "remote_host": "localhost",
@@ -1482,6 +1564,10 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
     assert stats_after_reset is None
 
 
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+    FakeNixlWrapper,
+)
 def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init):
     """reqs_to_send deadlines are stamped with the scheduler process's
     perf_counter, whose epoch differs across processes and (by boot-time
@@ -1530,11 +1616,9 @@ def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist
 
 
 def test_kv_connector_stats_aggregation():
-    """
-    Test KV transfer stats aggregation across TP ranks using
+    """Test KV transfer stats aggregation across TP ranks using
     KVOutputAggregator (used by MultiprocExecutor).
     """
-
     # Create KVOutputAggregator for 3 workers (simulating TP=3), same thing
     # done in MultiprocExecutor.execute_model
     aggregator = KVOutputAggregator(expected_finished_count=3)
@@ -1594,14 +1678,92 @@ def test_kv_connector_stats_aggregation():
     assert cli_stats["Avg xfer time (ms)"] == 1500.0
     assert cli_stats["Avg post time (ms)"] == 1500.0
     assert cli_stats["Avg number of descriptors"] == 1.5
+    # Reduced values must be plain Python scalars so CLI logging renders
+    # them without numpy reprs (eg np.float64(...)).
+    assert all(not isinstance(v, np.generic) for v in cli_stats.values())
+
+
+def test_kv_connector_stats_failure_grouping():
+    """Transfer, handshake and notification failures are reported as one
+    transport-failure count, while KV expiry is reported separately: the
+    former are sporadic lower-transport-layer events, the latter an
+    autoscaler signal."""
+    stats = NixlKVConnectorStats()
+    assert stats.is_empty()
+
+    stats.record_failed_transfer()
+    stats.record_failed_handshake()
+    stats.record_failed_notification()
+    stats.record_kv_expired_req()
+    assert not stats.is_empty()
+
+    # No successful transfers: latency stats are zero but the failure
+    # counts still surface.
+    reduced = stats.reduce()
+    assert reduced["Num successful transfers"] == 0
+    assert reduced["Num failed transfers"] == 3
+    assert reduced["Num KV expired reqs"] == 1
+
+
+def test_nixl_prom_metrics_group_handshake_with_transfer_failures():
+    """vllm:nixl_num_failed_transfers counts handshake and notification
+    failures too, while vllm:nixl_num_kv_expired_reqs stays a separate
+    counter."""
+    from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
+        NixlPromMetrics,
+    )
+
+    registry = CollectorRegistry()
+
+    class RegistryGauge(Gauge):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, registry=registry, **kwargs)
+
+    class RegistryCounter(Counter):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, registry=registry, **kwargs)
+
+    class RegistryHistogram(Histogram):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, registry=registry, **kwargs)
+
+    vllm_config = create_vllm_config()
+    metric_types = {
+        Gauge: RegistryGauge,
+        Counter: RegistryCounter,
+        Histogram: RegistryHistogram,
+    }
+    prom = NixlPromMetrics(
+        vllm_config,
+        metric_types,
+        labelnames=["engine"],
+        per_engine_labelvalues={0: ["engine-0"]},
+    )
+
+    stats = NixlKVConnectorStats()
+    stats.record_failed_transfer()
+    stats.record_failed_handshake()
+    stats.record_failed_notification()
+    stats.record_kv_expired_req()
+    prom.observe(stats.data, engine_idx=0)
+
+    def counter_value(name: str) -> float:
+        for metric in registry.collect():
+            for sample in metric.samples:
+                if sample.name == name:
+                    return sample.value
+        raise AssertionError(f"metric {name} not found in registry")
+
+    assert counter_value("vllm:nixl_num_failed_transfers_total") == 3.0
+    assert counter_value("vllm:nixl_num_kv_expired_reqs_total") == 1.0
 
 
 def test_multi_kv_connector_stats_aggregation():
-    """
-    Test MultiKVConnectorStats aggregation across TP ranks using
+    """Test MultiKVConnectorStats aggregation across TP ranks using
     KVOutputAggregator (used by MultiprocExecutor).
     """
-
     aggregator = KVOutputAggregator(expected_finished_count=3)
 
     from dataclasses import dataclass
@@ -1742,8 +1904,7 @@ def test_scheduler_kv_connector_stats_aggregation():
     FakeNixlWrapper,
 )
 def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
-    """
-    Test lifecycle of an aborted Remote Prefill request hitting the timeout.
+    """Test lifecycle of an aborted Remote Prefill request hitting the timeout.
     -----> P
             |  {process request}
      <-/--- |  {result is NOT delivered, eg proxy is down}
@@ -1906,21 +2067,49 @@ def test_mixed_memory_local_descriptors_split_by_memory_type():
     assert handle == 22
     assert worker._dram_src_handles_by_block_size[worker.block_size] == 11
     assert [block[2] for block in blocks] == [0, 0, 3, 3]
-    assert [
+    memory_types = [
         call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list
-    ] == ["DRAM", "VRAM"]
+    ]
+    assert memory_types == ["DRAM", "VRAM"]
 
 
-def test_mixed_memory_read_notifies_after_both_transfers_finish():
+@pytest.fixture
+def recv_worker():
+    """Receive lifecycle state without distributed or device initialization."""
     worker = object.__new__(NixlConnectorWorker)
-    worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])}
-    worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])}
-    worker._dram_src_handles_by_block_size = {16: 10}
-    worker._recving_metadata = {"request": MagicMock()}
+    worker.transfer_topo = MagicMock()
+    worker.transfer_topo.block_size_ratio.return_value = 1
+    worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
+        remote_block_size=16, remote_physical_blocks_per_logical=1
+    )
+    worker._physical_blocks_per_logical_kv_block = 1
+    worker._recving_metadata = {"request": MagicMock(local_block_ids=([1, 2, 3],))}
     worker._recving_transfers = defaultdict(list)
+    worker._failed_recv_reqs = queue.Queue()
+    worker._recv_failures = set()
+    worker._replicated_pcp_done_sending = set()
+    worker._invalid_block_ids = queue.Queue()
     worker._pending_recv_notifs = {}
-    worker.xfer_stats = MagicMock()
+    worker._reqs_to_send = {}
+    worker._replicated_pcp_done_sending = set()
+    worker._is_hma_required = False
+    worker._has_mamba = False
+    worker.use_host_buffer = False
+    worker.enable_permute_local_kv = False
+    worker.enable_heterogeneous_attn_post_process = False
+    worker.tp_rank = 0
+    worker._log_failure = MagicMock()  # type: ignore[method-assign]
+    worker.xfer_stats = NixlKVConnectorStats()
     worker.nixl_wrapper = MagicMock()
+    worker.nixl_wrapper.get_xfer_telemetry.return_value = get_default_xfer_telemetry()
+    return worker
+
+
+def test_mixed_memory_read_notifies_after_both_transfers_finish(recv_worker):
+    worker = recv_worker
+    worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])}
+    worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])}
+    worker._dram_src_handles_by_block_size = {16: 10}
     worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102]
     worker.nixl_wrapper.check_xfer_state.return_value = "DONE"
 
@@ -1947,26 +2136,20 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish():
     np.testing.assert_array_equal(device_read.args[4], [7])
     worker.nixl_wrapper.send_notif.assert_not_called()
 
-    assert worker._pop_done_transfers(worker._recving_transfers) == {"request"}
+    assert worker.get_finished() == (set(), {"request"})
     worker.nixl_wrapper.send_notif.assert_called_once_with(
         "prefill", notif_msg=b"request:1"
     )
 
 
-def test_mixed_memory_read_failure_does_not_notify_producer():
+def test_mixed_memory_read_failure_does_not_notify_producer(recv_worker):
     """A failed half of a split READ must suppress its success notification."""
-    worker = object.__new__(NixlConnectorWorker)
-    worker._recving_metadata = {"request": MagicMock()}
+    worker = recv_worker
     worker._recving_transfers = {"request": [101, 102]}
     worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}
-    worker._is_hma_required = True
-    worker._failed_recv_reqs = MagicMock()
-    worker._log_failure = MagicMock()  # type: ignore[method-assign]
-    worker.xfer_stats = MagicMock()
-    worker.nixl_wrapper = MagicMock()
     worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"]
 
-    worker._pop_done_transfers(worker._recving_transfers)
+    assert worker.get_finished() == (set(), {"request"})
 
     assert "request" not in worker._pending_recv_notifs
     worker.nixl_wrapper.send_notif.assert_not_called()
@@ -2007,8 +2190,7 @@ def test_register_kv_caches(
     layout,
     separate_kv_head_groups,
 ):
-    """
-    Test that register_kv_caches() properly calls nixl_wrapper methods with
+    """Test that register_kv_caches() properly calls nixl_wrapper methods with
     correct data.
 
     This test verifies:
@@ -2017,7 +2199,6 @@ def test_register_kv_caches(
     2. nixl_wrapper.get_xfer_descs() is called with blocks_data containing
        block layout info
     """
-
     vllm_config = create_vllm_config(attention_backend=attn_backend)
     vllm_config.cache_config.kv_cache_layout = layout
 
@@ -2245,17 +2426,14 @@ class FakePlatform(Platform):
 
     @classmethod
     def get_nixl_supported_devices(cls) -> dict[str, tuple[str, ...]]:
-        """
-        Returns a mapping from device_type to a tuple of supported
+        """Returns a mapping from device_type to a tuple of supported
         kv_buffer_device for nixl.
         """
         return {"oot": ("oot",)}
 
     @classmethod
     def get_nixl_memory_type(cls) -> str | None:
-        """
-        Returns the nixl memory type for the current platform.
-        """
+        """Returns the nixl memory type for the current platform."""
         return "VRAM"
 
 
@@ -2268,8 +2446,7 @@ def get_nixl_memory_type(cls) -> str | None:
 def test_kv_buffer_to_nixl_memory_types(
     default_vllm_config, dist_init, kv_buffer_device, nixl_memory_type
 ):
-    """
-    Test that register_kv_caches() passes the correct memory types from the
+    """Test that register_kv_caches() passes the correct memory types from the
     config to the nixl_wrapper.
     """
     vllm_config = create_vllm_config()
@@ -2444,6 +2621,48 @@ def test_engine_ttl_eviction(default_vllm_config, dist_init):
         mock_rem.assert_any_call("agent_1")
 
 
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+    FakeNixlWrapper,
+)
+def test_engine_with_inflight_transfer_is_not_evicted(default_vllm_config, dist_init):
+    """A transfer that outlives the TTL must keep its engine registered.
+
+    _engine_last_active is stamped when a read is issued and not refreshed
+    while it runs, so a stalled transfer -- a peer that has lost its NIC holds
+    one indefinitely -- leaves its engine looking idle. Evicting it releases
+    the dlist handle and remote agent the transfer is still reading through.
+    """
+    worker, engine_id = _setup_worker_with_remote_engine(engine_ttl=10.0)
+    nixl_wrapper = worker.nixl_wrapper
+
+    request_id = "req-still-reading"
+    worker._recving_transfers[request_id] = [MagicMock()]
+    worker._recving_metadata[request_id] = MagicMock(
+        remote=MagicMock(engine_id=engine_id)
+    )
+
+    with (
+        patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel,
+        patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem,
+    ):
+        worker._engine_last_active[engine_id] = time.perf_counter() - 20.0
+
+        worker._evict_stale_engines()
+
+        assert engine_id in worker._remote_agents
+        assert engine_id in worker.dst_xfer_side_handles
+        mock_rel.assert_not_called()
+        mock_rem.assert_not_called()
+
+        # Once the transfer is done the engine is stale like any other.
+        del worker._recving_transfers[request_id]
+        worker._evict_stale_engines()
+
+        assert engine_id not in worker._remote_agents
+        assert mock_rem.call_count == 2
+
+
 @patch(
     "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
     FakeNixlWrapper,
@@ -2499,8 +2718,7 @@ def test_transfer_topology_unregister():
     FakeNixlWrapper,
 )
 def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_init):
-    """
-    Create and schedule a request so that P adds it to in-batch tracking via
+    """Create and schedule a request so that P adds it to in-batch tracking via
     the real scheduler, then simulate an abort (request not in next scheduler
     iteration) and verify the worker no longer tracks it as in-batch.
     """
@@ -2614,6 +2832,120 @@ def check_xfer_state(self, handle: int) -> str:
         return super().check_xfer_state(handle)
 
 
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+    FailingNixlWrapper,
+)
+@pytest.mark.parametrize("awaiting_kvs", [True, False])
+@pytest.mark.parametrize("notification_fails", [False, True])
+def test_empty_recv_is_reported_only_when_awaited(
+    default_vllm_config, dist_init, awaiting_kvs, notification_fails
+):
+    """An empty recv is reported iff the scheduler parked the request on it.
+
+    _read_blocks returns early when the local block list is empty -- D already
+    holds the KV and only the producer needs telling. Two kinds of caller reach
+    that path, and they need opposite handling:
+
+    Parked (awaiting_kvs=True, a full prefix cache hit). Without a report the
+    request is named in neither _recving_transfers nor _failed_recv_reqs and
+    never reaches finished_recving. The scheduler has no other way to release a
+    WAITING_FOR_REMOTE_KVS request, and there is no timeout, so it sits in
+    skipped_waiting holding its blocks for the life of the process.
+
+    Notify-only (awaiting_kvs=False): request_finished seeding an empty recv to
+    free P's blocks for a request aborted before it was scheduled, or a readback
+    on a request that stays RUNNING. Reporting either one crashes the engine
+    core -- _update_from_kv_xfer_finished asserts the id is still in
+    self.requests and that the request is parked or finished.
+
+    A failed notification changes neither: the KV is local either way, and the
+    producer frees its own blocks on a timeout.
+    """
+    vllm_config = create_vllm_config()
+    connector = NixlConnector(
+        vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
+    )
+    connector.connector_worker = FakeNixlConnectorWorker(
+        vllm_config,
+        connector.engine_id,
+        hand_shake_latency=0.0,
+        kv_cache_config=connector._kv_cache_config,
+    )
+    connector.connector_worker.nixl_wrapper.fail_send_notif = notification_fails
+
+    request_id = f"test_empty_recv_awaited_{awaiting_kvs}"
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id=request_id,
+        local_block_ids=(),  # empty: the whole prompt is already cached locally
+        kv_transfer_params={
+            "remote_block_ids": [[20, 21, 22]],
+            "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
+            "remote_request_id": f"prefill-{request_id}",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "remote_tp_size": 1,
+        },
+        awaiting_kvs=awaiting_kvs,
+    )
+    connector.bind_connector_metadata(metadata)
+    dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
+    connector.start_load_kv(dummy_ctx)
+    connector.bind_connector_metadata(NixlConnectorMetadata())
+    time.sleep(0.1)
+    connector.start_load_kv(dummy_ctx)
+
+    _, done_recving = connector.get_finished(finished_req_ids=set())
+    assert (request_id in done_recving) is awaiting_kvs
+
+
+@pytest.mark.parametrize("is_hma", [False, True])
+@pytest.mark.parametrize(
+    "local_block_ids,awaiting_kvs",
+    [((), False), (([],), False), (([],), True), (([1, 2, 3],), True)],
+)
+def test_handshake_failure_reports_only_awaited_recvs(
+    recv_worker, is_hma, local_block_ids, awaiting_kvs
+):
+    """A failed cleanup handshake must not complete an already-aborted request."""
+    worker = recv_worker
+    worker._is_hma_required = is_hma
+    worker._recving_metadata.clear()
+    worker._remote_agents = {}
+    worker._handshake_lock = contextlib.nullcontext()
+    worker._ready_requests = queue.Queue()
+    worker._reqs_to_process = set()
+    worker.pcp_rank = 0
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id="request",
+        local_block_ids=local_block_ids,
+        kv_transfer_params={
+            "remote_block_ids": ([4, 5, 6],),
+            "remote_engine_id": "prefill",
+            "remote_request_id": "prefill-request",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+        },
+        awaiting_kvs=awaiting_kvs,
+    )
+    handshake = Future[None]()
+    handshake.set_exception(RuntimeError("handshake failed"))
+    with patch.object(worker, "_ensure_handshake", return_value=handshake):
+        worker.start_load_kv(metadata)
+
+    results = worker.get_transfer_results()
+    expected = {"request"} if awaiting_kvs else set()
+    assert results.finished_recving == results.failed_recving == expected
+    assert worker.get_block_ids_with_load_errors() == (
+        set(local_block_ids[0]) if awaiting_kvs and not is_hma else set()
+    )
+    assert not worker._recving_metadata
+    assert not worker._recv_failures
+    assert not worker._pending_recv_notifs
+
+
 @patch(
     "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
     FailingNixlWrapper,
@@ -2816,22 +3148,26 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
     # Wait for handshake to fail
     time.sleep(0.3)
 
-    # Check that blocks were marked invalid
-    invalid_blocks = connector.get_block_ids_with_load_errors()
-    assert invalid_blocks == {1, 2, 3}
-
-    # Check that request appears in get_finished
+    assert connector.get_block_ids_with_load_errors() == set()
     _, done_recving = connector.get_finished(finished_req_ids=set())
     assert request_id in done_recving
+    assert connector.get_block_ids_with_load_errors() == {1, 2, 3}
+
+    # Handshake failures are recorded as transport failures, separately
+    # from KV expiry.
+    assert connector.connector_worker.xfer_stats.data["num_failed_handshakes"]
+    assert connector.connector_worker.xfer_stats.data["num_failed_transfers"] == []
 
 
 @patch(
     "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
     FailingNixlWrapper,
 )
-def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init):
-    """Test that transfer setup failures mark blocks invalid
-    and return via get_finished."""
+@pytest.mark.parametrize("is_hma", [False, True])
+def test_transfer_setup_failure_returns_finished(
+    default_vllm_config, dist_init, is_hma
+):
+    """Setup failures report the request; only non-HMA reports block IDs."""
     vllm_config = create_vllm_config()
 
     connector = NixlConnector(
@@ -2841,6 +3177,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
         vllm_config, connector.engine_id, hand_shake_latency=0
     )
     connector.connector_worker.nixl_wrapper.fail_transfer_setup = True
+    connector.connector_worker._is_hma_required = is_hma
 
     request_id = "test_transfer_fail"
     metadata = NixlConnectorMetadata()
@@ -2870,13 +3207,107 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
     time.sleep(0.1)
     connector.start_load_kv(dummy_ctx)
 
-    # check that blocks were marked invalid
+    assert connector.get_block_ids_with_load_errors() == set()
+    results = connector.get_transfer_results(finished_req_ids=set())
+    assert request_id in results.finished_recving
+    assert results.failed_recving == {request_id}
     invalid_blocks = connector.get_block_ids_with_load_errors()
-    assert invalid_blocks == {7, 8, 9}
+    assert invalid_blocks == (set() if is_hma else {7, 8, 9})
+
+
+class _ScriptedXferWrapper(FakeNixlWrapper):
+    """Scripts per-handle xfer states; forbids releasing in-flight handles."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # handle -> successive states; the last one repeats.
+        self.xfer_states: dict[int, list[str]] = {}
+        self.released: list[int] = []
+
+    def check_xfer_state(self, handle: int) -> str:
+        states = self.xfer_states[handle]
+        return states.pop(0) if len(states) > 1 else states[0]
 
-    # ensure request appears in get_finished
+    def release_xfer_handle(self, handle: int) -> None:
+        # A posted-but-unfinished transfer cannot be aborted: releasing it
+        # leaves the RDMA READ armed. Production code must never do this.
+        assert self.xfer_states.get(handle, ["DONE"])[0] != "PROC", (
+            f"released in-flight handle {handle}"
+        )
+        self.released.append(handle)
+
+
+def _make_split_read_connector(vllm_config, request_id, states):
+    """Seed a request with scripted in-flight xfer handles + recv metadata."""
+    connector = NixlConnector(
+        vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
+    )
+    connector.connector_worker = FakeNixlConnectorWorker(
+        vllm_config, connector.engine_id, hand_shake_latency=0
+    )
+    worker = connector.connector_worker
+    wrapper = _ScriptedXferWrapper("agent")
+    worker.nixl_wrapper = wrapper
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id=request_id,
+        local_block_ids=([7, 8, 9],),
+        kv_transfer_params={
+            "remote_block_ids": ([10, 11, 12],),
+            "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
+            "remote_request_id": f"prefill-{request_id}",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "remote_tp_size": 1,
+        },
+    )
+    worker._recving_metadata[request_id] = metadata.reqs_to_recv[request_id]
+    wrapper.xfer_states = dict(states)
+    worker._recving_transfers[request_id] = list(states)
+    return connector, worker, wrapper
+
+
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+    FakeNixlWrapper,
+)
+def test_split_read_failure_defers_report_until_last_handle(
+    default_vllm_config, dist_init
+):
+    """One half of a split (mixed DRAM/VRAM) read failing must not report the
+    request — nor invalidate its blocks — while the sibling xfer is still in
+    flight: a posted READ cannot be aborted and would DMA into blocks the
+    scheduler could free and reuse. The report happens exactly once, when the
+    last handle is terminal."""
+    request_id = "split_read_partial_failure"
+    err_handle, live_handle = 11, 22
+    connector, worker, wrapper = _make_split_read_connector(
+        create_vllm_config(),
+        request_id,
+        {err_handle: ["ERR"], live_handle: ["PROC", "PROC", "DONE"]},
+    )
+
+    # Poll 1: one half fails; the sibling is in flight -> nothing reported.
     _, done_recving = connector.get_finished(finished_req_ids=set())
-    assert request_id in done_recving
+    assert done_recving == set()
+    assert connector.get_block_ids_with_load_errors() == set()
+    assert request_id in worker._recving_metadata
+    assert wrapper.released == [err_handle]
+
+    # Poll 2: sibling still in flight.
+    _, done_recving = connector.get_finished(finished_req_ids=set())
+    assert done_recving == set()
+
+    # Poll 3: sibling terminal -> reported exactly once, blocks invalidated.
+    _, done_recving = connector.get_finished(finished_req_ids=set())
+    assert done_recving == {request_id}
+    assert connector.get_block_ids_with_load_errors() == {7, 8, 9}
+    assert request_id not in worker._recving_metadata
+    assert wrapper.released == [err_handle, live_handle]
+
+    # Poll 4: nothing left; no double report.
+    _, done_recving = connector.get_finished(finished_req_ids=set())
+    assert done_recving == set()
 
 
 @patch(
@@ -2896,7 +3327,7 @@ def test_failed_request_skips_kv_postprocessing(
     default_vllm_config, dist_init, failure_mode
 ):
     """Test that failed requests skip KV sync and post-processing in
-    get_finished().
+    get_transfer_results().
 
     This is the core safety behavior: when a KV transfer fails at any stage,
     the request must still appear in done_recving (so the scheduler can apply
@@ -2981,11 +3412,10 @@ def test_failed_request_skips_kv_postprocessing(
         patch.object(worker, "sync_recved_kv_to_device") as mock_sync,
         patch.object(worker, "post_process_device_kv_on_receive") as mock_postprocess,
     ):
-        _, done_recving = connector.get_finished(finished_req_ids=set())
+        results = connector.get_transfer_results(finished_req_ids=set())
 
-    # The failed request must appear in done_recving so the scheduler
-    # can handle it (e.g., trigger recompute via kv_load_failure_policy).
-    assert request_id in done_recving
+    assert request_id in results.finished_recving
+    assert results.failed_recving == {request_id}
 
     # Critical: KV sync and post-processing must NOT have been called
     # since no valid KV data was received for the failed request.
@@ -3000,85 +3430,49 @@ def test_failed_request_skips_kv_postprocessing(
     assert invalid_blocks == {1, 2, 3}
 
 
-@patch(
-    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
-    FailingNixlWrapper,
-)
-def test_handles_failing_in_separate_polls_do_not_kill_the_engine(
-    default_vllm_config, dist_init
-):
-    """A request whose handles fail in different polls must not crash the engine.
-
-    One transfer handle is created per remote rank, and when a peer goes away
-    they do not all fail in the same poll: one errors while another is still
-    PROC. The first failure reports the request via _failed_recv_reqs and
-    get_finished() pops its metadata; when the remaining handles fail in a
-    later poll, the request must be cleaned up without being reported again.
-
-    Reporting twice kills the EngineCore: the scheduler's assert in
-    _update_from_kv_xfer_finished only expects a finished recv for a request
-    still waiting for KVs, having moved this one out of
-    WAITING_FOR_REMOTE_KVS to recompute locally after the first report.
-    """
-    vllm_config = create_vllm_config()
-    connector = NixlConnector(
-        vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
-    )
-    connector.connector_worker = FakeNixlConnectorWorker(
-        vllm_config, connector.engine_id, hand_shake_latency=0
-    )
-    worker = connector.connector_worker
+@pytest.mark.parametrize("sibling_state", ["ERR", "DONE"])
+def test_recv_failure_waits_for_sibling_transfer(recv_worker, sibling_state):
+    """Failure must not release blocks while a sibling can still write to them."""
+    worker = recv_worker
+    worker._recving_transfers["request"] = [101, 102]
+    worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}
+    worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "PROC", sibling_state]
 
-    request_id = "test_two_handles_failing_in_separate_polls"
-    metadata = NixlConnectorMetadata()
-    metadata.add_new_req_to_recv(
-        request_id=request_id,
-        local_block_ids=([1, 2, 3],),
-        kv_transfer_params={
-            "remote_block_ids": ([4, 5, 6],),
-            "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
-            "remote_request_id": f"prefill-{request_id}",
-            "remote_host": "localhost",
-            "remote_port": 1234,
-            "remote_tp_size": 1,
-        },
-    )
-    connector.bind_connector_metadata(metadata)
-    dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
-    connector.start_load_kv(dummy_ctx)
-    connector.bind_connector_metadata(NixlConnectorMetadata())
-    time.sleep(0.1)
-    connector.start_load_kv(dummy_ctx)
+    results = worker.get_transfer_results()
+    assert results.finished_recving == results.failed_recving == set()
+    assert worker.get_block_ids_with_load_errors() == set()
+    assert "request" in worker._recving_metadata
+    worker.nixl_wrapper.release_xfer_handle.assert_called_once_with(101)
 
-    # Give the request a second handle, as a second remote rank would.
-    handles = worker._recving_transfers[request_id]
-    assert len(handles) == 1
-    first, second = handles[0], object()
-    worker._recving_transfers[request_id] = [first, second]
+    results = worker.get_transfer_results()
+    assert results.finished_recving == results.failed_recving == {"request"}
+    assert worker.get_block_ids_with_load_errors() == {1, 2, 3}
+    worker.nixl_wrapper.send_notif.assert_not_called()
 
-    # Poll 1: the first handle has failed, the second is still running.
-    def first_failed(handle):
-        return "ERR" if handle is first else "PROC"
+    results = worker.get_transfer_results()
+    assert results.finished_recving == results.failed_recving == set()
+    assert worker.get_block_ids_with_load_errors() == set()
 
-    with patch.object(
-        worker.nixl_wrapper, "check_xfer_state", side_effect=first_failed
-    ):
-        _, done_recving = connector.get_finished(finished_req_ids=set())
 
-    assert request_id in done_recving
-    assert request_id not in worker._recving_metadata
-    assert worker._recving_transfers[request_id] == [second]
-
-    # Poll 2: the second handle fails too. The request was already reported,
-    # so it must not be reported a second time: the scheduler has since moved
-    # it out of WAITING_FOR_REMOTE_KVS to recompute locally, and asserts on a
-    # finished recv for a request in that state. Its remaining handles are
-    # still released and the transfer entry removed.
-    with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"):
-        _, done_recving = connector.get_finished(finished_req_ids=set())
+def test_recv_failure_waits_for_unpollable_handle_to_be_released(recv_worker):
+    """A polling exception and failed release do not prove that DMA has stopped."""
+    worker = recv_worker
+    worker._recving_transfers["request"] = [101]
+    worker.nixl_wrapper.check_xfer_state.side_effect = [RuntimeError("poll"), "DONE"]
+    worker.nixl_wrapper.release_xfer_handle.side_effect = [
+        RuntimeError("release"),
+        None,
+    ]
+
+    results = worker.get_transfer_results()
+    assert results.finished_recving == results.failed_recving == set()
+    assert worker.get_block_ids_with_load_errors() == set()
+    assert "request" in worker._recving_metadata
 
-    assert request_id not in done_recving
-    assert request_id not in worker._recving_transfers
+    results = worker.get_transfer_results()
+    assert results.finished_recving == results.failed_recving == {"request"}
+    assert worker.get_block_ids_with_load_errors() == {1, 2, 3}
+    assert worker.nixl_wrapper.release_xfer_handle.call_count == 2
 
 
 def _set_test_speculative_config(
@@ -3210,9 +3604,15 @@ def test_transfer_mode_changes_compatibility_hash():
 
 
 @pytest.mark.skip_global_cleanup
-def test_scheduler_advertises_transfer_mode():
-    # Each scheduler advertises its transfer mode in kv_transfer_params so an
-    # external router can route pull (READ) vs push (WRITE) producers.
+@pytest.mark.parametrize("mode", ["pull", "push"])
+@pytest.mark.parametrize(
+    "tp_size,pcp_size,dcp_size,transfer_tp_size",
+    [(1, 1, 1, 1), (4, 1, 1, 4), (4, 1, 4, 4), (1, 4, 1, 1), (1, 4, 4, 4)],
+)
+def test_scheduler_advertises_transfer_topology(
+    mode, tp_size, pcp_size, dcp_size, transfer_tp_size
+):
+    """The consumer must address every distinct producer KV shard."""
     from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
         NixlPullConnectorScheduler,
     )
@@ -3220,8 +3620,22 @@ def test_scheduler_advertises_transfer_mode():
         NixlPushConnectorScheduler,
     )
 
-    assert NixlPullConnectorScheduler._TRANSFER_MODE == "pull"
-    assert NixlPushConnectorScheduler._TRANSFER_MODE == "push"
+    config = create_vllm_config(kv_role="kv_producer")
+    config.parallel_config.tensor_parallel_size = tp_size
+    config.parallel_config.prefill_context_parallel_size = pcp_size
+    config.parallel_config.decode_context_parallel_size = dcp_size
+    cls = NixlPullConnectorScheduler if mode == "pull" else NixlPushConnectorScheduler
+    scheduler = cls(config, "prefiller", make_kv_cache_config(block_size=16))
+    request = create_request(request_id=1, num_tokens=32, do_remote_decode=True)
+    request.status = RequestStatus.FINISHED_LENGTH_CAPPED
+    try:
+        delay, params = scheduler.request_finished(request, ([0, 1],))
+        assert delay
+        assert params["transfer_mode"] == mode
+        assert params["tp_size"] == transfer_tp_size
+        assert params["dcp_size"] == dcp_size
+    finally:
+        scheduler.shutdown()
 
 
 @pytest.mark.parametrize(
@@ -3259,15 +3673,16 @@ def test_compatibility_hash_validation(
     should_fail,
     enforce_handshake_compat,
 ):
-    """
-    Test NIXL compatibility hash validation during handshake.
+    """Test NIXL compatibility hash validation during handshake.
 
-    Parameters:
+    Parameters
+    ----------
         mismatch_type: description of what is being tested
         config_overrides: dict of config to override for the remote instance
         version_override: version dict e.g. {"vllm_version": "0.6.1"}
         should_fail: whether the handshake should fail
         enforce_handshake_compat: whether to enforce compatibility checking
+
     """
     local_vllm_config = create_vllm_config(
         model="facebook/opt-125m",
@@ -3395,8 +3810,7 @@ def test_compatibility_hash_validation(
     FakeNixlWrapper,
 )
 def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario):
-    """
-    Test that msgspec decode errors are properly handled during handshake.
+    """Test that msgspec decode errors are properly handled during handshake.
 
     Tests both DecodeError and ValidationError for both decoders:
     - NixlHandshakePayload decoder
diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py
index 08506d82f246..3be1730ec3e0 100644
--- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py
@@ -3,18 +3,35 @@
 """Unit tests for NixlConnectorScheduler with HMA and Mamba N-1 prefill."""
 
 import gc
+import queue
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
+import msgspec
 import pytest
 import torch
 
 from tests.v1.attention.utils import MockMambaBuilder
 from vllm import LLM, SamplingParams
-from vllm.config import KVTransferConfig
+from vllm.config import KVTransferConfig, set_current_vllm_config
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+    NixlAgentMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+    NixlConnectorWorker,
+)
 from vllm.v1.core.single_type_kv_cache_manager import (
     FullAttentionManager,
     SlidingWindowManager,
 )
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+    MLAAttentionSpec,
+)
 
 from .utils import (
     create_request,
@@ -24,6 +41,251 @@
 )
 
 
+@pytest.fixture
+def region_pull_worker():
+    """Real page mapping and descriptors, with only the transport mocked."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping
+
+    worker = object.__new__(NixlConnectorWorker)
+    worker.engine_id = "D"
+    worker.tp_rank = 0
+    worker.block_size = 64
+    worker._physical_blocks_per_logical_kv_block = 1
+    worker.enable_permute_local_kv = False
+    worker.enable_heterogeneous_attn_post_process = False
+    worker._engine_last_active = {}
+    worker._bidirectional_kv_xfer_enabled = False
+    worker._recving_transfers = {}
+    worker.use_mla, worker._has_mamba = True, False
+    worker.dcp_size = 1
+    worker.dcp_rank = 0
+    spec = MLAAttentionSpec(
+        block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16
+    )
+    worker.kv_cache_config = KVCacheConfig(
+        num_blocks=100,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["resident"], spec, enable_kv_transfer=False),
+            KVCacheGroupSpec(["host"], spec),
+            KVCacheGroupSpec(["indexer"], spec),
+        ],
+    )
+    worker._group_spec_types = (MLAAttentionSpec, MLAAttentionSpec)
+    worker.region_group_ids = [0, 1]
+    worker.dst_region_group_ids = {"P": [0, 0]}
+    worker.num_regions = 2
+    worker.block_len_per_layer = [1024, 1024]
+    worker.transfer_topo = MagicMock()
+    worker.transfer_topo.tp_ratio.return_value = 1
+    worker.transfer_topo.block_size_ratio.return_value = 1
+    worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
+        remote_tp_size=4,
+        remote_dcp_size=1,
+        remote_block_size=64,
+        remote_physical_blocks_per_logical=2,
+    )
+    worker.tp_mappings = {"P": TPMapping(((0,), (0,)), (0,), {0: 0}, 0)}
+    worker._transfer_layer_group_ids = ()
+    worker._mixed_mem_types = True
+    worker.src_xfer_handles_by_block_size = {64: 1}
+    worker._dram_src_handles_by_block_size = {64: 2}
+    worker.dst_xfer_side_handles = {"P": {0: 3}}
+    worker.dst_num_blocks = {"P": 100, "D": 100}
+    worker.dst_region_num_blocks = {"P": [100, 100], "D": [100, 100]}
+    worker._remote_agents = {"P": {(0, 0): "P-rank0"}}
+    worker._read_blocks_mixed = MagicMock()
+    worker.nixl_wrapper = MagicMock()
+    return worker
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("local_ratio,remote_ratio", [(1, 2), (2, 1), (2, 2)])
+@pytest.mark.parametrize("num_tokens", [160, 224])
+@pytest.mark.parametrize("cached", [0, 1])
+def test_region_pull_ignores_allocation_padding(
+    region_pull_worker, local_ratio, remote_ratio, num_tokens, cached
+):
+    """Padded capacity must not shift either the host or indexer token range."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        NixlConnectorMetadata,
+    )
+    from vllm.utils.math_utils import cdiv
+
+    worker = region_pull_worker
+    worker._physical_blocks_per_logical_kv_block = local_ratio
+    worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
+        remote_tp_size=4,
+        remote_dcp_size=1,
+        remote_block_size=64,
+        remote_physical_blocks_per_logical=remote_ratio,
+    )
+    local_count = cdiv(num_tokens, 64 * local_ratio)
+    remote_count = cdiv(num_tokens, 64 * remote_ratio)
+    # The host has a cached prefix, while the indexer still needs the full range.
+    local = [
+        list(range(30 + cached, 30 + local_count)),
+        list(range(40, 40 + local_count)),
+    ]
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id="request",
+        local_block_ids=local,
+        local_num_computed_blocks=(0, cached, 0),
+        awaiting_kvs=True,
+        kv_transfer_params={
+            "remote_engine_id": "P",
+            "remote_request_id": "request-P",
+            "remote_host": "localhost",
+            "remote_port": 1,
+            "remote_num_tokens": num_tokens,
+            "remote_block_ids": [list(range(10, 10 + remote_count))],
+        },
+    )
+    meta = metadata.reqs_to_recv["request"]
+    meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
+        local, local_ratio
+    )
+    worker._read_blocks_for_req("request", meta)
+    read = worker._read_blocks_mixed.call_args.kwargs
+    valid_pages = cdiv(num_tokens, 64)
+    source_start = 10 * remote_ratio
+    assert read["remote_block_descs_ids"].tolist() == (
+        list(range(source_start + cached * local_ratio, source_start + valid_pages))
+        + list(range(100 + source_start, 100 + source_start + valid_pages))
+    )
+    assert read["local_block_descs_ids"].tolist() == (
+        list(range((30 + cached) * local_ratio, 30 * local_ratio + valid_pages))
+        + list(range(100 + 40 * local_ratio, 100 + 40 * local_ratio + valid_pages))
+    )
+    assert meta.region_blocks_to_zero == [
+        list(
+            range(base * local_ratio + valid_pages, (base + local_count) * local_ratio)
+        )
+        for base in (30, 40)
+    ]
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("num_pages", [0, 3, 19])
+@pytest.mark.parametrize("region_groups", [(0, 1), (1, 0, 1)])
+def test_dcp_region_pull(region_pull_worker, num_pages, region_groups):
+    """Read each uncached page from its DCP rank; notify ranks with no pages."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        RemoteMeta,
+        ReqMeta,
+    )
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping
+
+    worker = region_pull_worker
+    worker.region_group_ids = list(region_groups)
+    worker.num_regions = len(region_groups)
+    worker.dst_region_group_ids = {"P": [0] * worker.num_regions}
+    worker.dst_region_num_blocks = {
+        engine: [100] * worker.num_regions for engine in ("P", "D")
+    }
+    worker.block_len_per_layer = [1024] * worker.num_regions
+    worker.dcp_rank = 0
+    remote = worker.transfer_topo.get_engine_info.return_value
+    remote.remote_tp_size = remote.remote_dcp_size = 8
+    remote.remote_physical_blocks_per_logical = 1
+    ranks = tuple(range(8))
+    worker.tp_mappings["P"] = TPMapping((ranks, ranks), ranks, {r: r for r in ranks}, 8)
+    worker.dst_xfer_side_handles = {"P": {r: 1000 + r for r in ranks}}
+    worker._remote_agents = {"P": {(0, r): f"P-rank{r}" for r in ranks}}
+    # Two regions with different prefix hits, plus one padding page each.
+    local = (
+        [list(range(31, 31 + num_pages)), list(range(42, 41 + num_pages))]
+        if num_pages
+        else []
+    )
+    meta = ReqMeta(
+        local,
+        local,
+        tp_size=8,
+        local_num_computed_blocks=(0, 1, 2),
+        remote=RemoteMeta(
+            [[10, 11, 12]],
+            "localhost",
+            1,
+            "P",
+            "request-P",
+            num_tokens=max(0, num_pages * 64 - 7),
+        ),
+    )
+    worker._read_blocks_for_req("request", meta)
+    reads = {
+        call.kwargs["remote_xfer_side_handle"] - 1000: call.kwargs
+        for call in worker._read_blocks_mixed.call_args_list
+    }
+    notified = {call.args[0] for call in worker.nixl_wrapper.send_notif.call_args_list}
+    for rank in ranks:
+        expected = [
+            (region * 100 + 30 + group * 10 + page, region * 100 + 10 + page // 8)
+            for region, group in enumerate(region_groups)
+            for page in range(group + 1, num_pages)
+            if page % 8 == rank
+        ]
+        if expected:
+            read = reads[rank]
+            assert (
+                list(zip(read["local_block_descs_ids"], read["remote_block_descs_ids"]))
+                == expected
+            )
+            assert f"P-rank{rank}" not in notified
+        else:
+            assert rank not in reads and f"P-rank{rank}" in notified
+    assert meta.region_blocks_to_zero == (
+        [[30 + group * 10 + num_pages] for group in region_groups]
+        if num_pages
+        else None
+    )
+
+
+@pytest.mark.cpu_test
+def test_region_pull_completion_zeros_only_own_padding(region_pull_worker):
+    """Completion must preserve neighboring layers and already imported pages."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        RemoteMeta,
+        ReqMeta,
+    )
+
+    worker = region_pull_worker
+    # Two strided layer regions in one backing, with nonzero layer offsets.
+    backing = torch.full((4, 3, 8), 7, dtype=torch.uint8)
+    caches = {"host": backing[:, 1], "indexer": backing[:, 2]}
+    worker.device_kv_caches = caches
+    worker.region_names = list(caches)
+    worker.kv_caches_base_addr = {"D": {0: [c.data_ptr() for c in caches.values()]}}
+    worker.region_num_blocks = [4, 4]
+    worker.block_len_per_layer = [8, 8]
+    worker.block_stride_per_layer = [24, 24]
+    worker._recving_metadata = {
+        "request": ReqMeta(
+            local_block_ids=[[1, 2], [0, 1]],
+            local_physical_block_ids=[[1, 2], [0, 1]],
+            tp_size=1,
+            remote=RemoteMeta([[3]], "localhost", 1, "P", "request-P"),
+            region_blocks_to_zero=[[2], [1]],
+        )
+    }
+    worker._get_new_notifs = MagicMock(return_value=set())
+    worker._pop_done_transfers = MagicMock(return_value=({"request"}, set()))
+    worker._replicated_pcp_done_sending = set()
+    worker._failed_recv_reqs = queue.Queue()
+    worker._recv_failures = set()
+    worker._send_pending_recv_notifs = MagicMock()
+    worker._sync_device_after_direct_recv = MagicMock()
+    worker.use_host_buffer = False
+    worker._reqs_to_send = {}
+    result = worker.get_transfer_results()
+    expected = torch.full_like(backing, 7)
+    expected[2, 1] = 0
+    expected[1, 2] = 0
+    assert torch.equal(backing, expected)
+    assert result.finished_recving == {"request"}
+
+
 @pytest.mark.cpu_test
 @pytest.mark.parametrize(
     "swa_enabled,expected_sw_sizes",
@@ -134,8 +396,49 @@ def uncached(block_id):
 
     scheduler.update_state_after_alloc(request, blocks, num_external_tokens=2)
 
-    _, _, local_num_computed_blocks = scheduler._reqs_need_recv[request.request_id]
+    _, _, local_num_computed_blocks, awaiting_kvs = scheduler._reqs_need_recv[
+        request.request_id
+    ]
     assert local_num_computed_blocks == (2, 1)
+    # There are external tokens to pull, so the scheduler parks the request and
+    # is waiting to see it in finished_recving.
+    assert awaiting_kvs is True
+
+
+@pytest.mark.cpu_test
+def test_full_local_hit_is_not_awaited_by_the_scheduler():
+    """A recv with num_external_tokens == 0 must not be flagged as awaited.
+
+    The request keeps running -- there is nothing to pull, only a notification
+    owed to the producer -- so the worker must not report it in
+    finished_recving. Doing so trips `assert RequestStatus.is_finished` in
+    _update_from_kv_xfer_finished, which only tolerates a finished recv for a
+    request that is parked or already done."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
+        NixlPullConnectorScheduler,
+    )
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.core.kv_cache_utils import KVCacheBlock
+
+    scheduler = object.__new__(NixlPullConnectorScheduler)
+    scheduler._reqs_in_batch = set()
+    scheduler._reqs_need_save = {}
+    scheduler._reqs_need_recv = {}
+    scheduler.use_host_buffer = False
+    scheduler.is_bidirectional_kv_xfer_enabled = False
+    scheduler._is_hma_required = False
+    scheduler.kv_cache_config = MagicMock(
+        select_transfer_block_ids=lambda block_ids: block_ids
+    )
+
+    blocks = KVCacheBlocks(blocks=([KVCacheBlock(block_id=0, _block_hash=object())],))
+    request = create_request(do_remote_prefill=True)
+
+    scheduler.update_state_after_alloc(request, blocks, num_external_tokens=0)
+
+    _, local_block_ids, _, awaiting_kvs = scheduler._reqs_need_recv[request.request_id]
+    assert not local_block_ids  # full local hit: nothing to pull
+    assert awaiting_kvs is False
 
 
 @pytest.mark.cpu_test
@@ -272,26 +575,24 @@ def test_logical_to_kernel_block_ids_with_hma():
 
 @pytest.mark.cpu_test
 @pytest.mark.parametrize(
-    "is_rocm,has_mamba,use_host_buffer,done_recving,failed_recving,expected_syncs",
+    "is_rocm,has_mamba,use_host_buffer,done_recving,expected_syncs",
     [
-        (True, True, False, {"req"}, set(), 1),
-        (False, True, False, {"req"}, set(), 0),
-        (True, False, False, {"req"}, set(), 0),
-        (True, True, True, {"req"}, set(), 0),
-        (True, True, False, set(), set(), 0),
-        (True, True, False, {"req"}, {"req"}, 0),
+        (True, True, False, {"req"}, 1),
+        (False, False, False, {"req"}, 0),
+        (False, True, False, {"req"}, 0),
+        (True, True, True, {"req"}, 0),
+        (True, True, False, set(), 0),
     ],
 )
-def test_sync_device_after_mamba_recv_gates(
+def test_sync_device_after_direct_recv_gates(
     monkeypatch,
     is_rocm,
     has_mamba,
     use_host_buffer,
     done_recving,
-    failed_recving,
     expected_syncs,
 ):
-    """Only direct-GPU Mamba receives on ROCm need a device fence."""
+    """Only backends requiring visibility fences synchronize direct receives."""
     from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
     from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
         NixlConnectorWorker,
@@ -309,7 +610,7 @@ def test_sync_device_after_mamba_recv_gates(
         lambda: sync_calls.append(True),
     )
 
-    worker._sync_device_after_mamba_recv(done_recving, failed_recving)
+    worker._sync_device_after_direct_recv(done_recving)
 
     assert len(sync_calls) == expected_syncs
 
@@ -399,6 +700,8 @@ def test_read_blocks_for_req_expands_remote_ids(
     )
 
     remote_engine_id = "remote-engine"
+    worker.region_group_ids = list(range(len(resolved_types)))
+    worker.dst_region_group_ids = {remote_engine_id: list(worker.region_group_ids)}
 
     worker.transfer_topo = MagicMock()
     # tp_ratio not exercised (all_source_ranks is empty so no reads run),
@@ -804,8 +1107,7 @@ def run_test_and_cleanup():
 
 @pytest.mark.cpu_test
 def test_nixl_metadata_hma_block_ids_structure():
-    """
-    Test that NixlConnectorMetadata correctly stores block IDs for multiple
+    """Test that NixlConnectorMetadata correctly stores block IDs for multiple
     KV cache groups when HMA is enabled.
     """
     from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
@@ -862,6 +1164,7 @@ def _make_mock_worker_for_desc_ids(
     worker = MagicMock(spec=NixlConnectorWorker)
     worker.num_regions = num_regions
     worker._uses_region_group_mapping = False
+    worker._transfer_layer_group_ids = ()
     worker._has_mamba = has_mamba
     worker._group_spec_types = group_spec_types
     worker.block_len_per_layer = block_len_per_layer or [100]
@@ -888,6 +1191,45 @@ def _make_mock_worker_for_desc_ids(
     return worker
 
 
+@pytest.mark.cpu_test
+def test_shared_nixl_region_broadcasts_every_group_block_table():
+    """A shared region must receive every group's block table exactly once."""
+    worker = _make_mock_worker_for_desc_ids(
+        num_regions=3,
+        has_mamba=False,
+        group_spec_types=(FullAttentionSpec,) * 3,
+        block_len_per_layer=[100] * 3,
+    )
+    worker.region_group_ids = [0, bw._SHARED_REGION_GROUP_ID, 2]
+    worker._uses_region_group_mapping = True
+
+    result = worker._compute_desc_ids(
+        block_ids=([1], [2], [3]),
+        dst_num_blocks=10,
+        block_size_ratio=None,
+        physical_blocks_per_logical=1,
+        region_num_blocks=[10, 10, 10],
+    )
+
+    assert result.tolist() == [1, 11, 12, 13, 23]
+
+
+@pytest.mark.cpu_test
+def test_nixl_region_sort_preserves_pipeline_layer_order():
+    """Numeric layer order prevents PP peers from pairing different regions."""
+    names = [
+        "model.layers.10.self_attn",
+        "model.layers.2.self_attn",
+        "model.layers.1.self_attn",
+    ]
+
+    assert sorted(names, key=bw._region_sort_key) == [
+        "model.layers.1.self_attn",
+        "model.layers.2.self_attn",
+        "model.layers.10.self_attn",
+    ]
+
+
 @pytest.mark.cpu_test
 def test_get_block_descs_ids_hybrid_ssm():
     """Test _compute_desc_ids uses per-group strides for hybrid
@@ -1097,10 +1439,7 @@ def test_post_process_zeroes_untransferred_tail():
     worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache}
     fa_group = MagicMock(layer_names=["attn.0"])
     ssm_group = MagicMock(layer_names=["mamba.0"])
-    worker.kv_cache_config = MagicMock(
-        kv_cache_groups=[fa_group, ssm_group],
-        transfer_groups=[fa_group, ssm_group],
-    )
+    worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, ssm_group])
     # The cached property filters mamba layers out of the permuted caches.
     attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker)
     assert len(attn_caches) == 1 and attn_caches[0] is attn_cache
@@ -1188,6 +1527,11 @@ def _make_fake_kv_cache_manager():
 
     manager = object.__new__(KVCacheManager)
     manager.coordinator = MagicMock()
+    manager.kv_cache_config = MagicMock()
+    manager.kv_cache_config.kv_cache_groups = (
+        MagicMock(host_resident=False),
+        MagicMock(host_resident=False),
+    )
     manager.coordinator.single_type_managers = (
         _FakeSingleTypeManager(True, 16, [10, 11, 12, 13, 14, 15]),  # attention
         _FakeSingleTypeManager(False, 16, [20, 21, 22, 23, 24, 25]),  # mamba
@@ -1238,6 +1582,7 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks():
     scheduler.kv_cache_manager.cache_blocks = MagicMock()
     scheduler.failed_recving_kv_req_ids = {"req-1"}
     scheduler.finished_recving_kv_req_ids = {"req-1"}
+    scheduler.prefix_replay_tokens = 0
 
     request = MagicMock()
     request.request_id = "req-1"
@@ -1723,16 +2068,6 @@ def test_exchange_clipped_blocks_ssm_positional_states():
     assert clipped == ([1, 2, 3], [0, 5, 6, 7])
 
 
-@pytest.mark.cpu_test
-def test_exchange_clipped_blocks_excludes_nontransfer_groups():
-    """Scheduler block tables must match the worker's registered regions."""
-    sched = make_nixl_scheduler()
-    sched.kv_cache_config = make_kv_cache_config(block_size=16, swa_enabled=True)
-    sched.kv_cache_config.kv_cache_groups[1].enable_kv_transfer = False
-
-    assert sched.get_exchange_clipped_blocks(([1, 2], [9])) == ([1, 2],)
-
-
 # ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
 
 
@@ -1786,6 +2121,95 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4):
     )
 
 
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("kernel_block_size", [16, 8])
+def test_nixl_keeps_device_block_count_with_hisparse_host_pool(kernel_block_size):
+    host_num_blocks = 4
+    gpu_num_blocks = 9
+    spec = MLAAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=8,
+        dtype=torch.float16,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=gpu_num_blocks,
+        hisparse_host_num_blocks=host_num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=host_num_blocks * spec.page_size_bytes,
+                layers=["mla.host"],
+                layer_stride=host_num_blocks * spec.page_size_bytes,
+                block_stride=spec.page_size_bytes,
+                host_resident=True,
+            ),
+            KVCacheTensor(
+                size=gpu_num_blocks * spec.page_size_bytes,
+                layers=["mla.device"],
+                layer_stride=gpu_num_blocks * spec.page_size_bytes,
+                block_stride=spec.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["mla.host"],
+                spec,
+                host_resident=True,
+            ),
+            KVCacheGroupSpec(["mla.device"], spec),
+        ],
+    )
+    vllm_config = create_vllm_config(block_size=16)
+    vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
+    fake_backend = MagicMock()
+    fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
+    fake_backend.get_name.return_value = "FLASHMLA"
+    fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
+    fake_backend.get_kv_cache_shape.return_value = (1, kernel_block_size, 1, 1)
+    fake_platform = MagicMock()
+    fake_platform.device_type = "cuda"
+    fake_platform.get_nixl_memory_type.return_value = "VRAM"
+    physical_per_logical = spec.block_size // kernel_block_size
+    physical_page_bytes = spec.page_size_bytes // physical_per_logical
+    host_cache = torch.zeros(
+        host_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8
+    )
+    device_cache = torch.zeros(
+        gpu_num_blocks * physical_per_logical, physical_page_bytes, dtype=torch.uint8
+    )
+
+    with (
+        patch.object(bw, "NixlWrapper"),
+        patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
+        patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
+        patch.object(bw, "current_platform", fake_platform),
+        set_current_vllm_config(vllm_config),
+    ):
+        worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config)
+        worker.use_mla = True
+        worker.nixl_wrapper.get_agent_metadata.return_value = b"metadata"
+        worker.register_kv_caches({"mla.host": host_cache, "mla.device": device_cache})
+
+    assert worker.num_blocks == gpu_num_blocks * physical_per_logical
+    assert dict(zip(worker.region_group_ids, worker.region_num_blocks)) == {
+        0: host_num_blocks * physical_per_logical,
+        1: gpu_num_blocks * physical_per_logical,
+    }
+    caches = {"mla.host": host_cache, "mla.device": device_cache}
+    expected_addrs = [
+        caches[name].data_ptr() + block * stride
+        for name, stride, count in zip(
+            worker.region_names,
+            worker.block_stride_per_layer,
+            worker.region_num_blocks,
+            strict=True,
+        )
+        for block in range(count)
+    ]
+    assert worker.src_blocks_data[:, 0].tolist() == expected_addrs
+
+
 @pytest.mark.cpu_test
 def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
     """Hybrid MLA+KDA registration: HMA tensors shared by both layer types
@@ -1803,6 +2227,9 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
     kv_cache_config = _make_hybrid_mla_kv_cache_config()
     unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
     vllm_config = create_vllm_config(block_size=12)
+    # The engine rewrites this global value to the smallest scheduler-group
+    # block size. NIXL must retain the transferable groups' 12-token geometry.
+    vllm_config.cache_config.block_size = 4
     # kv_buffer_device defaults to the *real* platform's device type, which on
     # a CPU-only test host would make this a host-buffer worker: host xfer
     # buffers are per-layer, so the HMA shared tensors would not be
@@ -1833,7 +2260,13 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
         worker.use_mla = True  # opt-125m test config is not MLA; force the flag
         worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata"
 
-        tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)]
+        kernel_page = unified_page // 3
+        block_stride = 2 * kernel_page
+        backing = torch.zeros(12, block_stride, dtype=torch.uint8)
+        tensors = [
+            backing[:, :kernel_page],
+            backing[:, kernel_page:],
+        ]
         # KDA layer first per tensor: exercises the dual-purpose flag merge.
         worker.register_kv_caches(
             {
@@ -1849,6 +2282,8 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
     # 12-token logical blocks over the 4-token MLA kernel block.
     assert worker._physical_blocks_per_logical_kv_block == 3
     assert worker.block_size == 4 and worker.num_blocks == 12
+    assert worker.region_group_ids == [bw._SHARED_REGION_GROUP_ID] * 2
+    assert worker.block_stride_per_layer == [block_stride, block_stride]
     # Both shared tensors are dual-purpose: their FA view is MLA even though
     # a KDA layer registered them first.
     assert worker._region_is_mla == [True, True]
@@ -1861,8 +2296,20 @@ def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
     # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks.
     assert worker.src_blocks_data.shape == (24 + 32, 3)
     fa_descs = worker.src_blocks_data[:24]
-    assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3
+    expected_addrs = [
+        tensor.data_ptr() + block * block_stride
+        for tensor in tensors
+        for block in range(12)
+    ]
+    assert fa_descs[:, 0].tolist() == expected_addrs
     assert all(size == unified_page // 3 for size in fa_descs[:, 1])
+    worker.nixl_wrapper.register_memory.assert_called_once()
+    metadata = msgspec.msgpack.decode(
+        worker.xfer_handshake_metadata.agent_metadata_bytes,
+        type=NixlAgentMetadata,
+    )
+    assert metadata.block_strides == [block_stride, block_stride]
+    assert metadata.region_num_blocks == [12, 12]
 
 
 @pytest.mark.cpu_test
@@ -1909,6 +2356,8 @@ def test_push_write_hybrid_mla_replicates_attention():
     worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}}
     worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]}
     worker.src_xfer_handles_by_block_size = {4: 300}
+    worker.region_group_ids = [0, 1]
+    worker.dst_region_group_ids = {engine_id: [0, 1]}
     worker._sending_transfers = defaultdict(list)
     worker._sending_transfers_lock = threading.Lock()
     worker.kv_cache_config = _make_hybrid_mla_kv_cache_config()
diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
index a9f800739881..a4bfa9ab83ae 100644
--- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
@@ -12,6 +12,7 @@
 """
 
 from collections import defaultdict
+from threading import Event, Lock
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -112,6 +113,7 @@ def test_local_descriptors_follow_each_region_pool_capacity():
     worker.block_len_per_layer = [16, 16]
     worker.block_stride_per_layer = [16, 16]
     worker.region_num_blocks = [2, 3]
+    worker._transfer_layer_region_indices = ()
 
     descriptors = worker._build_fa_local([100, 1000], block_size_ratio=1)
 
@@ -119,7 +121,8 @@ def test_local_descriptors_follow_each_region_pool_capacity():
 
 
 @pytest.mark.cpu_test
-def test_overlaid_transfer_groups_share_region_geometry():
+@pytest.mark.parametrize("push_pp", [False, True])
+def test_overlaid_transfer_groups_share_region_geometry(push_pp):
     """Groups overlaid on one allocation share its transfer region."""
     import msgspec
 
@@ -127,6 +130,9 @@ def test_overlaid_transfer_groups_share_region_geometry():
     from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
         NixlAgentMetadata,
     )
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
+        NixlPushConnectorWorker,
+    )
     from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
         NixlConnectorWorker,
     )
@@ -153,9 +159,18 @@ def test_overlaid_transfer_groups_share_region_geometry():
     }
     groups = [KVCacheGroupSpec([layer_name], spec) for layer_name in caches]
 
-    worker = object.__new__(NixlConnectorWorker)
+    worker_cls = NixlPushConnectorWorker if push_pp else NixlConnectorWorker
+    worker = object.__new__(worker_cls)
+    if push_pp:
+        worker._push_writer_stop = Event()
+        worker._push_writer_wake = Event()
+        worker._push_writer_thread = MagicMock()
+        worker._sending_transfers_lock = Lock()
+        worker._sending_transfers = defaultdict(list)
     worker.tp_rank = 0
     worker.world_size = 1
+    worker.transfer_tp_rank = 0
+    worker.transfer_tp_size = 1
     worker.block_size = 4
     worker.engine_id = "local-engine"
     worker.use_mla = True
@@ -186,11 +201,19 @@ def test_overlaid_transfer_groups_share_region_geometry():
     worker.host_buffer_kv_cache_layout = "NHD"
     worker._physical_blocks_per_logical_kv_block = 1
     worker._logical_num_blocks = num_blocks
+    worker.region_mem_types = []
     worker.region_group_ids = []
     worker.region_mem_types = []
     worker._mixed_mem_types = False
     worker.region_names = []
     worker.region_num_blocks = []
+    worker._mixed_mem_types = False
+    worker._desc_is_dram_by_block_size = {}
+    worker._desc_pos_by_block_size = {}
+    worker._dram_src_handles_by_block_size = {}
+    worker._transfer_layer_names = ()
+    worker._transfer_layer_region_indices = ()
+    worker._transfer_layer_group_ids = ()
     worker._region_is_mla = []
     worker.block_len_per_layer = []
     worker.block_stride_per_layer = []
@@ -198,7 +221,8 @@ def test_overlaid_transfer_groups_share_region_geometry():
     worker.use_host_buffer = False
     worker.host_xfer_buffers = {}
     worker.device_kv_caches = {}
-    worker.pp_size = 1
+    worker.pp_size = 2 if push_pp else 1
+    worker._is_hma_required = True
     worker.dcp_size = 1
     worker.pcp_size = 1
     worker.kv_buffer_device = "cuda"
@@ -233,7 +257,13 @@ def test_overlaid_transfer_groups_share_region_geometry():
     expected_addrs = [
         backing.data_ptr() + block * block_stride for block in range(num_blocks)
     ]
-    assert worker.src_blocks_data[:, 0].tolist() == expected_addrs
+    num_desc_regions = 2 if push_pp else 1
+    assert worker.src_blocks_data[:, 0].tolist() == expected_addrs * num_desc_regions
+    assert worker.num_descs == num_blocks * num_desc_regions
+    assert (
+        worker.dst_region_num_blocks[worker.engine_id]
+        == [num_blocks] * num_desc_regions
+    )
 
     metadata = msgspec.msgpack.decode(
         worker.xfer_handshake_metadata.agent_metadata_bytes,
@@ -241,6 +271,7 @@ def test_overlaid_transfer_groups_share_region_geometry():
     )
     assert metadata.region_group_ids == [-1]
     assert metadata.region_num_blocks == [num_blocks]
+    assert metadata.region_members == ([["layer.0", "layer.1"]] if push_pp else [])
     assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]]
 
 
@@ -304,8 +335,6 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo
     # would not be deduplicated. Pin it to the faked device type.
     vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
 
-    from unittest.mock import MagicMock
-
     fake_backend = MagicMock()
     fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
     fake_backend.get_name.return_value = "FLASHMLA"
@@ -552,6 +581,16 @@ def _make_remote_meta(
     )
 
 
+def _register_remote_agents(worker, metadata, tp_size):
+    """Mirror the async handshake callback that publishes prepared agents."""
+    worker._remote_agents[metadata.engine_id] = {
+        (0, rank): worker.add_remote_agent(
+            metadata, remote_tp_rank=rank, remote_tp_size=tp_size
+        )
+        for rank in range(tp_size)
+    }
+
+
 def _owned_byte_ranges(worker, group_logical_ids):
     """Byte ranges owned by a request: for each HMA region tensor, every
     logical block id of every group maps to one unified page."""
@@ -609,8 +648,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
         remote_num_logical=12,
         remote_ssm_sizes=(24, 32),
     )
-    for rank in (0, 1):
-        worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2)
+    _register_remote_agents(worker, meta_r, 2)
 
     # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok
     # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks.
@@ -626,7 +664,7 @@ def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
             "remote_block_ids": remote_ids,
             "remote_engine_id": "remote-engine",
             "remote_request_id": "prefill-req-b",
-            "remote_host": "localhost",
+            "remote_host": "remote-host",
             "remote_port": 1234,
             "tp_size": 2,
         },
@@ -709,8 +747,7 @@ def _run_hetero_case(
         remote_num_logical=max(2 * n_remote + 4, 8),
         remote_ssm_sizes=(48 // tp_size, 64 // tp_size),
     )
-    for rank in range(tp_size):
-        worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size)
+    _register_remote_agents(worker, meta_r, tp_size)
 
     # Sparse ids so neighbors exist between the request's blocks.
     local_attn = [2 * i + 1 for i in range(n_local)]
@@ -726,7 +763,7 @@ def _run_hetero_case(
             "remote_block_ids": remote_ids,
             "remote_engine_id": "remote-engine",
             "remote_request_id": "prefill-req-b",
-            "remote_host": "localhost",
+            "remote_host": "remote-host",
             "remote_port": 1234,
             "tp_size": tp_size,
         },
@@ -1251,3 +1288,112 @@ def test_csa_linear_remote_ple_is_copied_whole():
             tp_ratio=-2,
             transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1),
         )
+
+
+def _make_ring_worker():
+    """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a
+    model whose sliding-window KV lives in per-request rings."""
+    from unittest.mock import MagicMock
+
+    from vllm.config import set_current_vllm_config
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.kv_cache_interface import (
+        CircularBufferSpec,
+        KVCacheConfig,
+        KVCacheGroupSpec,
+        KVCacheTensor,
+        MLAAttentionSpec,
+        UniformTypeKVCacheSpecs,
+    )
+
+    paged_specs = {
+        f"paged.{index}": MLAAttentionSpec(
+            block_size=4, num_kv_heads=1, head_size=64, dtype=torch.float16
+        )
+        for index in range(2)
+    }
+    ring_specs = {
+        f"ring.{index}": CircularBufferSpec(
+            block_size=6,
+            num_kv_heads=1,
+            head_size=32,
+            head_size_v=0,
+            dtype=torch.float16,
+        )
+        for index in range(2)
+    }
+
+    def group(specs):
+        uniform = UniformTypeKVCacheSpecs.from_specs(specs)
+        assert uniform is not None
+        return KVCacheGroupSpec(list(specs), uniform)
+
+    # Both rings overlay the second paged region, so the scratch group must
+    # address only that region while the paged group spans both.
+    tensor_regions = (("paged.0",), ("paged.1", "ring.0", "ring.1"))
+    region_size, page_size = 512, 256
+    kv_cache_config = KVCacheConfig(
+        num_blocks=2,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=len(tensor_regions) * region_size,
+                layers=[layer_name],
+                layer_stride=region_size,
+                block_stride=page_size,
+                offset=region_index * region_size,
+            )
+            for region_index, layer_names in enumerate(tensor_regions)
+            for layer_name in layer_names
+        ],
+        kv_cache_groups=[group(paged_specs), group(ring_specs)],
+    )
+
+    vllm_config = create_vllm_config(block_size=4)
+    vllm_config.cache_config.enable_prefix_caching = False
+    vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
+    fake_backend = MagicMock()
+    fake_backend.get_supported_kernel_block_sizes.return_value = [4]
+    fake_backend.get_name.return_value = "TEST_ATTN"
+    fake_backend.full_cls_name.return_value = "test.AttentionBackend"
+    fake_platform = MagicMock()
+    fake_platform.device_type = "cuda"
+    fake_platform.get_nixl_memory_type.return_value = "VRAM"
+
+    with (
+        patch.object(bw, "NixlWrapper", _RecordingNixl),
+        patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
+        patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
+        patch.object(bw, "current_platform", fake_platform),
+        set_current_vllm_config(vllm_config),
+    ):
+        worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
+        tensors = [torch.zeros((2, 256), dtype=torch.uint8) for _ in range(2)]
+        worker.register_kv_caches(
+            {
+                layer_name: tensors[region_index]
+                for region_index, layer_names in enumerate(tensor_regions)
+                for layer_name in layer_names
+            }
+        )
+    return worker
+
+
+@pytest.mark.cpu_test
+def test_ring_scratch_without_mamba_registers_and_addresses_its_own_regions():
+    worker = _make_ring_worker()
+
+    assert worker._ple_group_index is None
+    assert worker._scratch_region_indices == [1]
+
+    desc_ids = worker._compute_desc_ids(
+        block_ids=([1], [0]),
+        dst_num_blocks=2,
+        block_size_ratio=None,
+        physical_blocks_per_logical=1,
+    )
+    # Paged block 1 in both regions; ring block 0 only in its scratch region.
+    assert sorted(desc_ids.tolist()) == sorted([0 * 2 + 1, 1 * 2 + 1, 1 * 2 + 0])
diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py
index 5b7760e7e33c..a81c1571a1c6 100644
--- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py
+++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py
@@ -15,7 +15,7 @@
 * The worker matches D registrations against P finished blocks (both
   scenario directions) and forwards non-PUSH_REG NIXL notifs to the main
   thread's ``_get_new_notifs``.
-* ``get_finished`` enqueues evictions for the writer.
+* ``get_transfer_results`` enqueues evictions for the writer.
 """
 
 from __future__ import annotations
@@ -32,6 +32,13 @@
 import msgspec
 import pytest
 
+from vllm.distributed.kv_transfer.kv_connector.utils import TransferTopology
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorTransferResults,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
+    NixlBaseConnectorWorker,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
     PUSH_REG_NOTIF_PREFIX,
     NixlAgentMetadata,
@@ -40,6 +47,13 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
     NixlPushConnectorWorker,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
+    NixlKVConnectorStats,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
+    ReadSpec,
+    TPMapping,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
     get_base_request_id,
 )
@@ -130,7 +144,7 @@ def test_p_side_mamba_truncates_before_cache_lookup(self):
 
         with patch.object(
             sched,
-            "_truncate_mamba_request_for_prefill",
+            "_truncate_request_for_prefill",
             side_effect=AssertionError("must not truncate after cache lookup"),
         ):
             assert sched.get_num_new_matched_tokens(request, 0) == (0, False)
@@ -333,6 +347,7 @@ def fresh(cls) -> _StubWriterWorker:
         )
 
         w._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list)
+        w._send_failures = set()
         w._sending_transfers_lock = threading.Lock()
         w._push_finished_blocks = {}
         w._pending_d_registrations = {}
@@ -347,18 +362,29 @@ def fresh(cls) -> _StubWriterWorker:
 
         # Base worker fields touched by start_load_kv / _get_new_notifs.
         w._recving_metadata = {}
+        w._pending_recv_notifs = {}
+        w._failed_recv_reqs = queue.Queue()
+        w._recv_failures = set()
         w._recving_transfers = defaultdict(list)
+        w._is_hma_required = False
+        w.xfer_stats = NixlKVConnectorStats()
         w._reqs_to_process = set()
         w._reqs_to_send = {}
         w.consumer_notification_counts_by_req = defaultdict(int)
         w.tp_rank = 0
         w.pcp_rank = 0
+        w.pcp_dcp_sharded = False
         w.world_size = 1
+        w.pp_size = 1
         w.engine_id = "test-decode-engine"
         w._remote_agents = {}
+        w._handshake_lock = threading.RLock()
         w._physical_blocks_per_logical_kv_block = 1
         w._uses_region_group_mapping = False
         w.region_group_ids = [0]
+        w._transfer_layer_names = ()
+        w._transfer_layer_region_indices = ()
+        w._transfer_layer_group_ids = ()
         w.dst_region_num_blocks = {}
         w.dst_region_group_ids = {}
         w.dst_uses_region_group_mapping = {}
@@ -412,6 +438,86 @@ def _registration_data(
     }
 
 
+class TestPushWriterRegSend:
+    """``_do_send_reg_notif``: the PUSH_REG goes to every handshaked agent of
+    the P engine, and the agent map is read under ``_handshake_lock`` so the
+    send can't race the handshake done-callback's insert or a concurrent
+    stale-engine eviction (X1)."""
+
+    def _reg_send_worker(self) -> tuple[_StubWriterWorker, list]:
+        w = _StubWriterWorker.fresh()
+        w.nixl_wrapper = MagicMock()
+        failures: list = []
+        w._handle_failed_transfer = lambda *a: failures.append(a)
+        return w, failures
+
+    def test_sends_to_all_agents_with_prefix_and_payload(self):
+        w, failures = self._reg_send_worker()
+        w._remote_agents["prefill-engine"] = {(0, 0): "agent-0", (0, 1): "agent-1"}
+        rd = _registration_data("req-1")
+
+        w._do_send_reg_notif("req-1", rd)
+
+        assert failures == []
+        calls = w.nixl_wrapper.send_notif.call_args_list
+        assert len(calls) == 2
+        assert {c.args[0] for c in calls} == {"agent-0", "agent-1"}
+        for call in calls:
+            notif = call.kwargs["notif_msg"]
+            assert notif.startswith(PUSH_REG_NOTIF_PREFIX)
+            # (Tuples decode back as lists, so compare identity fields.)
+            decoded = msgspec.msgpack.decode(notif[len(PUSH_REG_NOTIF_PREFIX) :])
+            assert decoded["request_id"] == "req-1"
+            assert decoded["decode_engine_id"] == rd["decode_engine_id"]
+            assert decoded["remote_engine_id"] == rd["remote_engine_id"]
+            assert decoded["local_block_ids"] == [
+                list(group) for group in rd["local_block_ids"]
+            ]
+
+    def test_queues_failed_request_when_engine_not_handshaked(self):
+        w, failures = self._reg_send_worker()
+
+        w._do_send_reg_notif("req-1", _registration_data("req-1"))
+
+        assert failures == []
+        assert w._failed_recv_reqs.get_nowait() == "req-1"
+        assert w._failed_recv_reqs.empty()
+        w.nixl_wrapper.send_notif.assert_not_called()
+
+    def test_agent_read_is_serialized_with_handshake_lock(self):
+        """While another thread holds ``_handshake_lock`` (as the handshake
+        done-callback does while writing ``_remote_agents``), the writer's
+        check-then-send must not run ahead of the critical section."""
+        w, failures = self._reg_send_worker()
+        w._remote_agents["prefill-engine"] = {(0, 0): "agent-0"}
+        rd = _registration_data("req-1")
+
+        w._handshake_lock.acquire()
+        done = threading.Event()
+        errors: list[BaseException] = []
+
+        def _send():
+            try:
+                w._do_send_reg_notif("req-1", rd)
+            except BaseException as e:  # pragma: no cover
+                errors.append(e)
+            finally:
+                done.set()
+
+        t = threading.Thread(target=_send)
+        t.start()
+        assert not done.wait(timeout=0.2)
+        w.nixl_wrapper.send_notif.assert_not_called()
+        assert failures == []
+
+        w._handshake_lock.release()
+        assert done.wait(timeout=2.0)
+        t.join(timeout=2.0)
+        assert errors == []
+        assert failures == []
+        assert w.nixl_wrapper.send_notif.call_count == 1
+
+
 class TestPushWriterMatching:
     def test_handle_push_reg_matches_existing_finished_blocks(self):
         """PUSH_REG arrives second (P finished first): match + fire."""
@@ -540,9 +646,11 @@ def test_start_load_kv_enqueues_to_writer(self):
         assert w._push_writer_wake.is_set()
         assert w.start_push_calls == []
 
-    def test_noncanonical_pcp_rank_skips_producer_work(self):
+    @pytest.mark.parametrize("sharded", [False, True])
+    def test_noncanonical_pcp_rank_only_pushes_distinct_shards(self, sharded):
         w = _StubWriterWorker.fresh()
         w.pcp_rank = 1
+        w.pcp_dcp_sharded = sharded
         w._send_heartbeats = MagicMock()
 
         meta = NixlConnectorMetadata()
@@ -552,11 +660,11 @@ def test_noncanonical_pcp_rank_skips_producer_work(self):
 
         w.start_load_kv(meta)
 
-        assert w._finished_blocks_inbox.empty()
-        assert "req" not in w._reqs_to_process
-        assert "req" not in w._reqs_to_send
-        assert not w._push_writer_wake.is_set()
-        w._send_heartbeats.assert_not_called()
+        assert w._finished_blocks_inbox.empty() is (not sharded)
+        assert ("req" in w._reqs_to_process) is sharded
+        assert ("req" in w._reqs_to_send) is sharded
+        assert w._push_writer_wake.is_set() is sharded
+        assert w._send_heartbeats.called is sharded
 
 
 # The P→D handshake must run on the base worker's background executor, never
@@ -630,6 +738,10 @@ def test_do_start_push_kv_drops_request_on_handshake_failure():
     assert xfer_calls == []
     assert len(failures) == 1
     assert failures[0]["failure_type"] == "push_handshake_failed"
+    # The handshake metric is counted once by _ensure_handshake's own
+    # callback (stubbed out here); the push side must not double-record it.
+    assert w.xfer_stats.data["num_failed_handshakes"] == []
+    assert w.xfer_stats.data["num_failed_transfers"] == []
 
 
 def test_writer_loop_drains_deferred_push_inbox():
@@ -719,6 +831,31 @@ def test_stale_engine_evicted_on_push():
     w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old")
 
 
+def test_cleanup_remote_engine_pops_under_handshake_lock():
+    """Regression (X1): eviction must not remove ``_remote_agents`` entries
+    while another thread is inside the lock's critical section (e.g. the
+    push writer reading the map to send a registration)."""
+    w = _eviction_worker(engine_ttl=30.0)
+    w._remote_agents["D-old"] = {(0, 0): "agent-D-old"}
+
+    w._handshake_lock.acquire()
+    done = threading.Event()
+
+    def _cleanup():
+        w._cleanup_remote_engine("D-old")
+        done.set()
+
+    t = threading.Thread(target=_cleanup)
+    t.start()
+    assert "D-old" in w._remote_agents
+
+    w._handshake_lock.release()
+    assert done.wait(timeout=2.0)
+    t.join(timeout=2.0)
+    assert "D-old" not in w._remote_agents
+    w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old")
+
+
 class TestPushWriterNotifs:
     def test_get_new_notifs_processes_forwarded_completion_notif(self):
         """Non-PUSH_REG notifs forwarded by the writer thread are drained
@@ -742,21 +879,19 @@ def test_get_new_notifs_processes_forwarded_completion_notif(self):
         assert notified == set()
         assert request_id in w._recving_transfers
 
-    def test_get_finished_evicts_completed_state(self):
-        """``get_finished`` should enqueue evictions and wake the writer."""
+    def test_get_transfer_results_evicts_completed_state(self):
+        """Transfer completion should enqueue evictions and wake the writer."""
         w = _StubWriterWorker.fresh()
 
-        # Stub the base ``get_finished`` to return one done_sending entry.
-        # Patch via the MRO's parent class.
         with patch.object(
             NixlPushConnectorWorker.__mro__[1],
-            "get_finished",
-            return_value=({"req-done"}, set()),
+            "get_transfer_results",
+            return_value=KVConnectorTransferResults(finished_sending={"req-done"}),
         ):
-            done_sending, done_recving = w.get_finished()
+            results = w.get_transfer_results()
 
-        assert "req-done" in done_sending
-        assert done_recving == set()
+        assert results.finished_sending == {"req-done"}
+        assert results.finished_recving == set()
         # Eviction enqueued for the writer.
         evicted = []
         while True:
@@ -767,6 +902,161 @@ def test_get_finished_evicts_completed_state(self):
         assert evicted == ["req-done"]
         assert w._push_writer_wake.is_set()
 
+    @staticmethod
+    def _pollable_worker() -> _StubWriterWorker:
+        """Stub worker with the extra base-worker state the real
+        ``get_transfer_results`` needs to poll ``_sending_transfers``."""
+        w = _StubWriterWorker.fresh()
+        w.transfer_topo = MagicMock()
+        w.nixl_wrapper = MagicMock()
+        w.xfer_stats = MagicMock()
+        w._log_failure = MagicMock()  # type: ignore[method-assign]
+        w._invalid_block_ids = queue.Queue()
+        w._pending_recv_notifs = {}
+        w._replicated_pcp_done_sending = set()
+        w.use_host_buffer = False
+        return w
+
+    def _make_sending_req(self, w: _StubWriterWorker) -> str:
+        request_id = "req-send-1"
+        w._sending_transfers[request_id] = [101, 102]
+        w._reqs_to_send[request_id] = time.perf_counter() + 60
+        w._reqs_to_process.add(request_id)
+        return request_id
+
+    def test_failed_push_send_is_not_reported_done_sending(self):
+        """A failed P-side WRITE must not surface as finished sending.
+
+        Reporting it would free P-side blocks as if the KV had been
+        delivered while D still waits for the data; the request must be
+        left to the lease / watchdog instead.
+        """
+        w = self._pollable_worker()
+        request_id = self._make_sending_req(w)
+        w.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"]
+
+        results = w.get_transfer_results()
+
+        assert request_id not in results.finished_sending
+        # Lease tracking is preserved so the watchdog can reschedule.
+        assert request_id in w._reqs_to_send
+        assert request_id in w._reqs_to_process
+        # Both handles were still cleaned up.
+        assert request_id not in w._sending_transfers
+
+    @pytest.mark.parametrize(("pp_size", "is_hma"), [(1, False), (2, True)])
+    def test_completed_push_send_releases_blocks_without_consumer_notif(
+        self, pp_size, is_hma
+    ):
+        """All WRITEs completing lets P reclaim its blocks without a D ACK."""
+        w = self._pollable_worker()
+        w.pp_size = pp_size
+        w._is_hma_required = is_hma
+        request_id = self._make_sending_req(w)
+        w.nixl_wrapper.check_xfer_state.side_effect = ["DONE", "PROC", "DONE"]
+
+        results = w.get_transfer_results()
+
+        assert request_id not in results.finished_sending
+        assert request_id in w._reqs_to_send
+        assert w._sending_transfers[request_id] == [102]
+        assert w._evict_finished_inbox.empty()
+
+        results = w.get_transfer_results()
+        assert results.finished_sending == {request_id}
+        assert request_id not in w._reqs_to_send
+        assert request_id not in w._reqs_to_process
+        assert w._evict_finished_inbox.get_nowait() == request_id
+        assert w.get_transfer_results().finished_sending == set()
+        assert w._evict_finished_inbox.empty()
+        assert w.nixl_wrapper.release_xfer_handle.call_count == 2
+        w.xfer_stats.record_kv_expired_req.assert_not_called()
+
+    @pytest.mark.parametrize("sibling_state", ["ERR", "DONE"])
+    def test_push_failure_survives_until_lease_expiry(self, sibling_state):
+        """A later sibling completion must not turn a failed send into success."""
+        w = self._pollable_worker()
+        request_id = self._make_sending_req(w)
+        w.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "PROC", sibling_state]
+
+        assert w.get_transfer_results().finished_sending == set()
+        assert w._sending_transfers[request_id] == [102]
+        assert w.get_transfer_results().finished_sending == set()
+        assert request_id not in w._sending_transfers
+        assert request_id in w._reqs_to_send
+        assert w._evict_finished_inbox.empty()
+
+        w.expected_consumer_notifications_by_req = {}
+        w._reqs_to_send[request_id] = time.perf_counter() - 1
+        assert w.get_transfer_results().finished_sending == {request_id}
+        assert request_id not in w._send_failures
+        assert w._evict_finished_inbox.get_nowait() == request_id
+
+    @pytest.mark.parametrize("release_fails", [False, True])
+    def test_push_submission_failure_is_remembered_after_sibling_success(
+        self, release_fails
+    ):
+        """A submission error remains a failure even if later polls succeed."""
+        w = self._pollable_worker()
+        request_id = self._make_sending_req(w)
+        w.transfer_topo.block_size_ratio.return_value = 1
+        w._apply_prefix_caching = MagicMock(return_value=([[1]], [[1]]))
+        w._compute_desc_ids = MagicMock(return_value=[0])
+        w.dst_num_blocks = {w.engine_id: 8}
+        w.dst_region_num_blocks[w.engine_id] = [8]
+        w.dst_region_group_ids[w.engine_id] = [0]
+        w.dst_uses_region_group_mapping[w.engine_id] = False
+        w.nixl_wrapper.make_prepped_xfer.return_value = 101
+        w.nixl_wrapper.transfer.side_effect = RuntimeError("submit")
+        if release_fails:
+            w.nixl_wrapper.release_xfer_handle.side_effect = [
+                RuntimeError("release"),
+                None,
+                None,
+            ]
+
+        handle = w._xfer_blocks(
+            read_spec=ReadSpec(0, [[1]], [[1]]),
+            dst_engine_id=w.engine_id,
+            request_id=request_id,
+            remote_request_id="decode-request",
+            local_xfer_side_handle=1,
+            remote_xfer_side_handle=2,
+        )
+        assert handle == (101 if release_fails else None)
+        w._sending_transfers[request_id] = [102]
+        if handle is not None:
+            w._sending_transfers[request_id].append(handle)
+        w.nixl_wrapper.check_xfer_state.return_value = "DONE"
+
+        assert w.get_transfer_results().finished_sending == set()
+        assert request_id in w._reqs_to_send
+        assert request_id not in w._sending_transfers
+        assert w._evict_finished_inbox.empty()
+
+    @pytest.mark.parametrize("first_state", ["ERR", "DONE"])
+    def test_push_completion_after_expiry_is_not_reported_twice(self, first_state):
+        """A WRITE finishing after lease expiry must not complete the request again."""
+        w = self._pollable_worker()
+        request_id = self._make_sending_req(w)
+        w.nixl_wrapper.check_xfer_state.side_effect = [
+            first_state,
+            "PROC",
+            "PROC",
+            "DONE",
+        ]
+        assert w.get_transfer_results().finished_sending == set()
+
+        w.expected_consumer_notifications_by_req = {}
+        w._reqs_to_send[request_id] = time.perf_counter() - 1
+        assert w.get_transfer_results().finished_sending == {request_id}
+        assert w._evict_finished_inbox.get_nowait() == request_id
+
+        assert w.get_transfer_results().finished_sending == set()
+        assert request_id not in w._sending_transfers
+        assert request_id not in w._send_failures
+        assert w._evict_finished_inbox.empty()
+
 
 # ----------------------------------------------------------------- #
 #  Negative / error-path tests                                       #
@@ -928,19 +1218,21 @@ def test_handle_push_reg_idempotent_for_same_request_id(self):
         assert len(w._pending_d_registrations) == 1
         assert w.start_push_calls == []
 
-    def test_get_finished_enqueues_eviction_for_each_done_request(self):
-        """``get_finished`` must enqueue an eviction for every request
+    def test_get_transfer_results_enqueues_each_completed_request(self):
+        """``get_transfer_results`` must enqueue every completed request
         in ``done_sending`` so the writer can drop stale matching state.
         Unlike the happy-path test, this verifies the *cardinality*: N
         completed requests -> N evictions, in order."""
         w = _StubWriterWorker.fresh()
         with patch.object(
             NixlPushConnectorWorker.__mro__[1],
-            "get_finished",
-            return_value=({"req-1", "req-2", "req-3"}, set()),
+            "get_transfer_results",
+            return_value=KVConnectorTransferResults(
+                finished_sending={"req-1", "req-2", "req-3"}
+            ),
         ):
-            done_sending, _ = w.get_finished()
-        assert done_sending == {"req-1", "req-2", "req-3"}
+            results = w.get_transfer_results()
+        assert results.finished_sending == {"req-1", "req-2", "req-3"}
 
         evicted: list[str] = []
         while True:
@@ -950,23 +1242,48 @@ def test_get_finished_enqueues_eviction_for_each_done_request(self):
                 break
         assert sorted(evicted) == ["req-1", "req-2", "req-3"]
 
-    def test_get_finished_with_no_completions_does_not_enqueue_eviction(self):
+    def test_empty_transfer_results_do_not_enqueue_eviction(self):
         """If there's nothing newly done, no eviction should be enqueued.
         The wake event IS still set because ``get_finished`` always wakes
         the writer to drain notifs."""
         w = _StubWriterWorker.fresh()
         with patch.object(
             NixlPushConnectorWorker.__mro__[1],
-            "get_finished",
-            return_value=(set(), set()),
+            "get_transfer_results",
+            return_value=KVConnectorTransferResults(),
         ):
-            done_sending, done_recving = w.get_finished()
-        assert done_sending == set()
-        assert done_recving == set()
+            results = w.get_transfer_results()
+        assert results.finished_sending == set()
+        assert results.finished_recving == set()
         assert w._evict_finished_inbox.qsize() == 0
         # Wake set so the writer drains NIXL notifs even when idle.
         assert w._push_writer_wake.is_set()
 
+    def test_failed_push_transfer_does_not_evict_writer_state(self):
+        w = _StubWriterWorker.fresh()
+        w.nixl_wrapper = MagicMock()
+        w.nixl_wrapper.check_xfer_state.return_value = "ERR"
+        w.xfer_stats = MagicMock()
+        w._log_failure = MagicMock()
+        w._sending_transfers["req-failed"] = [17]
+        w._reqs_to_send["req-failed"] = time.perf_counter() + 30
+        w._reqs_to_process.add("req-failed")
+
+        with patch.object(
+            NixlPushConnectorWorker.__mro__[1],
+            "get_transfer_results",
+            return_value=KVConnectorTransferResults(),
+        ):
+            results = w.get_transfer_results()
+
+        assert results.finished_sending == set()
+        assert results.finished_recving == set()
+        assert "req-failed" not in w._sending_transfers
+        assert "req-failed" in w._reqs_to_send
+        assert "req-failed" in w._reqs_to_process
+        assert w._evict_finished_inbox.empty()
+        w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17)
+
     def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog):
         """A completion notif for a request the worker doesn't know
         about should be logged but not crash."""
@@ -1160,6 +1477,7 @@ def _mla_worker_writing_to(d_ranks):
             )
         }
         w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids
+        w.dst_region_group_ids[engine_id] = [0]
         w.dst_xfer_side_handles = {engine_id: {r: 1000 + r for r in d_ranks}}
         w.src_xfer_handles_by_block_size = {16: 2000}
         w._remote_agents = {engine_id: {(0, r): f"agent-{r}" for r in d_ranks}}
@@ -1240,6 +1558,8 @@ def test_csa_linear_writes_every_group_to_each_target(self):
             )
         }
         w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids
+        w.region_group_ids = [0, 1, 2, 3]
+        w.dst_region_group_ids[engine_id] = [0, 1, 2, 3]
         w.dst_xfer_side_handles = {engine_id: {0: 1000, 1: 1001}}
         w.src_xfer_handles_by_tp_ratio = {(-2, 16): [2000, 2001]}
 
@@ -1372,3 +1692,375 @@ def test_no_prefix_hit_leaves_blocks_untrimmed(self):
         local, remote = self._written_block_ids(w)
         assert local == [10, 11, 12]
         assert remote == [500, 501, 502]
+
+    @pytest.mark.parametrize("layer_name_routing", [False, True])
+    def test_different_region_groups_require_layer_routing(self, layer_name_routing):
+        """A layer mapped to a pooled remote region can still receive its blocks."""
+        w, engine_id = self._worker_driving_xfer()
+        w.dst_region_group_ids[engine_id] = [-1]
+        if layer_name_routing:
+            w._transfer_layer_group_ids = (0,)
+        reg = _registration_data("req-groups", local_block_ids=([500, 501],))
+
+        if not layer_name_routing:
+            with pytest.raises(NotImplementedError, match="cache-group layouts"):
+                NixlPushConnectorWorker._do_start_push_kv(
+                    w, "req-groups", ([10, 11],), reg
+                )
+            w.nixl_wrapper.make_prepped_xfer.assert_not_called()
+            return
+
+        NixlPushConnectorWorker._do_start_push_kv(w, "req-groups", ([10, 11],), reg)
+
+        local, remote = self._written_block_ids(w)
+        assert local == [10, 11]
+        assert remote == [500, 501]
+
+
+def _agent_metadata(
+    region_layers: list[list[str]],
+    base_addresses: list[int],
+    block_lens: list[int],
+    block_strides: list[int] | None = None,
+) -> NixlAgentMetadata:
+    return NixlAgentMetadata(
+        engine_id="remote-engine",
+        agent_metadata=b"agent",
+        kv_caches_base_addr=base_addresses,
+        device_id=7,
+        num_blocks=2,
+        block_lens=block_lens,
+        block_strides=list(block_lens if block_strides is None else block_strides),
+        kv_cache_layout="LBHNC",
+        block_size=16,
+        ssm_sizes=(0, 0),
+        attn_backend_name="FLASH_ATTN",
+        physical_blocks_per_logical_kv_block=1,
+        region_members=region_layers,
+    )
+
+
+def _layer_routing_worker(
+    region_layers: list[list[str]],
+    group_by_layer: dict[str, int],
+    pp_size: int = 2,
+    is_hma: bool = True,
+) -> _StubWriterWorker:
+    worker = _StubWriterWorker.fresh()
+    worker.pp_size = pp_size
+    worker.dcp_size = 1
+    worker.block_size = 16
+    worker._has_mamba = False
+    worker._is_hma_required = is_hma
+    worker.kv_cache_config = MagicMock(transfer_group_index_by_layer=group_by_layer)
+    worker.transfer_topo = MagicMock()
+    worker._set_region_layers(region_layers)
+    return worker
+
+
+@pytest.mark.parametrize(
+    ("region_num_blocks", "expected"),
+    [(None, [1, 2, 15, 21, 22]), ([4, 7, 5], [1, 2, 9, 12, 13])],
+)
+def test_layer_group_ids_route_descriptor_blocks(region_num_blocks, expected):
+    worker = _layer_routing_worker(
+        [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0}
+    )
+    desc_ids = worker._compute_desc_ids(
+        block_ids=[[1, 2], [5]],
+        dst_num_blocks=10,
+        block_size_ratio=None,
+        physical_blocks_per_logical=1,
+        region_num_blocks=region_num_blocks,
+    )
+    assert desc_ids.tolist() == expected
+
+
+def test_layer_metadata_round_trip():
+    metadata = _agent_metadata([["L0", "L1"]], [0x10000], [256])
+    metadata.region_num_blocks = [4]
+    metadata.region_group_ids = [-1]
+    metadata.region_names = ["L0"]
+    metadata.region_mem_types = ["VRAM"]
+
+    encoded = msgspec.msgpack.encode(metadata)
+    assert msgspec.msgpack.Decoder(NixlAgentMetadata).decode(encoded) == metadata
+
+
+def test_layer_identity_gate_preserves_the_non_hma_path():
+    assert _layer_routing_worker([["a"]], {"a": 0})._transfer_layer_region_indices == (
+        0,
+    )
+
+    # A base worker (pull) never routes by layer name.
+    pull = object.__new__(NixlBaseConnectorWorker)
+    pull._transfer_layer_names = ()
+    pull._transfer_layer_region_indices = ()
+    pull._transfer_layer_group_ids = ()
+    pull._set_region_layers([["a"]])
+    assert pull._transfer_layer_region_indices == ()
+
+    # A non-HMA local layout does not require layer-name routing.
+    assert (
+        _layer_routing_worker(
+            [["a"]], {"a": 0}, is_hma=False
+        )._transfer_layer_region_indices
+        == ()
+    )
+
+    # PP=1 has congruent local/remote regions and keeps the region route even
+    # when its allocator is hybrid.
+    assert (
+        _layer_routing_worker(
+            [["a"]], {"a": 0}, pp_size=1
+        )._transfer_layer_region_indices
+        == ()
+    )
+
+
+def test_layer_alignment_fails_loud_when_remote_omits_layers():
+    # Falling back to region-index routing here would silently transfer stale
+    # KV, so an unannounced peer must fail the handshake instead.
+    worker = _layer_routing_worker([["a"]], {"a": 0})
+    with pytest.raises(AssertionError, match="no region_members"):
+        worker._align_remote_regions_by_layer(_agent_metadata([], [0xA000], [128]))
+
+
+def test_layer_alignment_expands_pooled_regions():
+    worker = _layer_routing_worker(
+        [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0}
+    )
+    assert worker._transfer_layer_names == ("a", "a.swa", "b")
+    assert worker._transfer_layer_region_indices == (0, 0, 1)
+    assert worker._transfer_layer_group_ids == (0, 1, 0)
+
+    metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128])
+    metadata.region_num_blocks = [4, 6]
+    metadata.region_group_ids = [-1, 0]
+    metadata.region_names = ["a", "b"]
+    metadata.region_mem_types = ["VRAM", "VRAM"]
+    worker._align_remote_regions_by_layer(metadata)
+
+    assert metadata.kv_caches_base_addr == [0xA000, 0xA000, 0xB000]
+    assert metadata.block_lens == [128, 128, 128]
+    assert metadata.block_strides == [128, 128, 128]
+    assert metadata.region_members == [["a"], ["a.swa"], ["b"]]
+    assert metadata.region_num_blocks == [4, 4, 6]
+    assert metadata.region_group_ids == [-1, -1, 0]
+    assert metadata.region_names == ["a", "a.swa", "b"]
+    assert metadata.region_mem_types == ["VRAM", "VRAM", "VRAM"]
+
+
+def test_layer_alignment_filters_and_reorders_a_pp_stage():
+    worker = _layer_routing_worker([["l2"], ["l3"]], {"l2": 0, "l3": 1})
+    metadata = _agent_metadata(
+        [["l2"], ["l0"], ["l3"], ["l1"]],
+        [0xC000, 0xA000, 0xD000, 0xB000],
+        [65536, 65536, 32768, 32768],
+        [131072, 131072, 65536, 65536],
+    )
+    metadata.region_num_blocks = [4, 6, 8, 10]
+    metadata.region_group_ids = [0, 0, 1, 1]
+    metadata.region_names = ["l2", "l0", "l3", "l1"]
+    metadata.region_mem_types = ["VRAM"] * 4
+
+    worker._align_remote_regions_by_layer(metadata)
+
+    assert worker._transfer_layer_region_indices == (0, 1)
+    assert worker._transfer_layer_group_ids == (0, 1)
+    assert metadata.kv_caches_base_addr == [0xC000, 0xD000]
+    assert metadata.block_lens == [65536, 32768]
+    assert metadata.block_strides == [131072, 65536]
+    assert metadata.region_num_blocks == [4, 8]
+    assert metadata.region_group_ids == [0, 1]
+    assert metadata.region_names == ["l2", "l3"]
+    assert metadata.region_mem_types == ["VRAM", "VRAM"]
+
+
+@pytest.mark.parametrize(
+    ("local_num_blocks", "remote_num_blocks"),
+    [([4, 4], [4, 4, 4]), ([4, 6], [3, 5, 7])],
+)
+def test_layer_descriptors_pair_layers_across_asymmetric_pp_split(
+    local_num_blocks, remote_num_blocks
+):
+    """A PP split can leave each stage a different mix of attention types.
+
+    Stage 1 owns L3 (full) pooled with L4 (sliding) in region 0, and L5 (full)
+    alone in region 1. The consumer holds every layer and pools them
+    differently, so region indices cannot be paired positionally.
+    """
+    worker = _layer_routing_worker([["L3", "L4"], ["L5"]], {"L3": 0, "L4": 1, "L5": 0})
+    worker.block_len_per_layer = [128, 128]
+    worker.block_stride_per_layer = [256, 512]
+    worker._region_is_mla = [False, False]
+    worker.num_blocks = 4
+    worker.region_num_blocks = local_num_blocks
+    worker.device_id = 0
+
+    consumer = _agent_metadata(
+        [["L0", "L1"], ["L3", "L2"], ["L5", "L4"]],
+        [0xA000, 0xB000, 0xC000],
+        [128, 128, 128],
+        [256, 512, 1024],
+    )
+    consumer.num_blocks = 4
+    consumer.region_num_blocks = remote_num_blocks
+    worker._align_remote_regions_by_layer(consumer)
+
+    # L3 -> remote region 1, L4 -> 2, L5 -> 2. Pairing by index would have sent
+    # L5 to region 1 and L4 to region 2's sibling.
+    assert consumer.kv_caches_base_addr == [0xB000, 0xC000, 0xC000]
+
+    plan = TPMapping(((0,), (0,)), (0,), {0: 0}, 0)
+    local_descs = worker._build_fa_local([0x1000, 0x2000], block_size_ratio=1)
+    remote_descs = worker._build_fa_remote(plan, consumer, block_size_ratio=1)
+    local_counts = [local_num_blocks[i] for i in worker._transfer_layer_region_indices]
+    local_ids = worker._compute_desc_ids(
+        [[1, 2], [3]], 4, None, 1, region_num_blocks=local_counts
+    )
+    remote_ids = worker._compute_desc_ids(
+        [[0, 1], [2]], 4, None, 1, region_num_blocks=consumer.region_num_blocks
+    )
+    assert len(local_descs) == sum(local_counts)
+    assert len(remote_descs) == sum(consumer.region_num_blocks)
+
+    # Each pair addresses the same layer's blocks despite different pooling
+    # and strides. L3 and L5 use group 0; L4 uses group 1.
+    assert local_descs[local_ids].tolist() == [
+        [0x1100, 128, 0],
+        [0x1200, 128, 0],
+        [0x1300, 128, 0],
+        [0x2200, 128, 0],
+        [0x2400, 128, 0],
+    ]
+    assert remote_descs[remote_ids].tolist() == [
+        [0xB000, 128, 7],
+        [0xB200, 128, 7],
+        [0xC800, 128, 7],
+        [0xC000, 128, 7],
+        [0xC400, 128, 7],
+    ]
+
+
+def test_layer_alignment_is_canonical_across_remote_orderings():
+    local, groups = [["x"], ["y"]], {"x": 0, "y": 1}
+    rank0 = _agent_metadata([["x"], ["y"]], [0x1000, 0x2000], [64, 128], [256, 512])
+    rank1 = _agent_metadata([["y"], ["x"]], [0x2000, 0x1000], [128, 64], [512, 256])
+
+    _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank0)
+    _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank1)
+
+    assert rank0.kv_caches_base_addr == rank1.kv_caches_base_addr == [0x1000, 0x2000]
+    assert rank0.block_lens == rank1.block_lens == [64, 128]
+    assert rank0.block_strides == rank1.block_strides == [256, 512]
+
+
+def test_layer_alignment_is_idempotent():
+    worker = _layer_routing_worker(
+        [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0}
+    )
+    metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128])
+    metadata.region_num_blocks = [4, 6]
+    metadata.region_group_ids = [-1, 0]
+    metadata.region_names = ["a", "b"]
+    metadata.region_mem_types = ["VRAM", "VRAM"]
+
+    worker._align_remote_regions_by_layer(metadata)
+    aligned = msgspec.msgpack.encode(metadata)
+    worker._align_remote_regions_by_layer(metadata)
+
+    assert msgspec.msgpack.encode(metadata) == aligned
+
+
+def test_layer_alignment_rejects_missing_local_layer():
+    worker = _layer_routing_worker([["l0"], ["l1"]], {"l0": 0, "l1": 1})
+    metadata = _agent_metadata([["l0"]], [0xA000], [128])
+
+    with pytest.raises(AssertionError, match="missing locally owned layers"):
+        worker._align_remote_regions_by_layer(metadata)
+
+
+def test_layer_alignment_rejects_duplicate_remote_layer():
+    worker = _layer_routing_worker([["a"]], {"a": 0})
+    metadata = _agent_metadata([["a"], ["a"]], [0xA000, 0xB000], [128, 128])
+
+    with pytest.raises(AssertionError, match="in multiple regions"):
+        worker._align_remote_regions_by_layer(metadata)
+
+
+def test_layer_alignment_rejects_inconsistent_remote_metadata():
+    worker = _layer_routing_worker([["a"]], {"a": 0})
+    metadata = _agent_metadata([["a"], ["b"]], [0xA000], [128])
+
+    with pytest.raises(AssertionError, match="lengths disagree"):
+        worker._align_remote_regions_by_layer(metadata)
+
+
+@pytest.mark.parametrize(
+    "field",
+    ["region_num_blocks", "region_group_ids", "region_names", "region_mem_types"],
+)
+def test_layer_alignment_rejects_inconsistent_region_geometry(field):
+    worker = _layer_routing_worker([["a"]], {"a": 0})
+    metadata = _agent_metadata([["a"], ["b"]], [0xA000, 0xB000], [128, 128])
+    setattr(metadata, field, [])
+    original = msgspec.msgpack.encode(metadata)
+
+    with pytest.raises(AssertionError, match="lengths disagree"):
+        worker._align_remote_regions_by_layer(metadata)
+    assert msgspec.msgpack.encode(metadata) == original
+
+
+def test_set_region_layers_rejects_duplicate_local_layer():
+    with pytest.raises(AssertionError, match="spans multiple NIXL regions"):
+        _layer_routing_worker([["a"], ["a"]], {"a": 0})
+
+
+def test_set_region_layers_rejects_layer_outside_any_kv_group():
+    with pytest.raises(AssertionError, match="outside any local group"):
+        _layer_routing_worker([["a"], ["b"]], {"a": 0})
+
+
+@pytest.mark.parametrize(
+    ("local_block_size", "remote_block_size", "remote_tp_size", "error"),
+    [
+        (32, 16, 1, "identical P/D block sizes"),
+        (16, 32, 1, "identical P/D block sizes"),
+        (16, 16, 2, "decode TP greater"),
+    ],
+)
+def test_layer_handshake_rejects_unsupported_geometry(
+    local_block_size: int, remote_block_size: int, remote_tp_size: int, error: str
+):
+    """Reject unsupported peers without registering agents or transfer state."""
+    metadata = _agent_metadata([["a"]], [0xA000], [128])
+    metadata.block_size = remote_block_size
+    worker = _layer_routing_worker([["a"]], {"a": 0})
+    worker.block_size = local_block_size
+    worker.block_len_per_layer = [128]
+    worker.use_mla = False
+    worker.nixl_wrapper = MagicMock()
+    worker.kv_caches_base_addr = defaultdict(dict)
+    worker.dst_num_blocks = {}
+    worker.tp_mappings = {}
+    worker.transfer_topo = TransferTopology(
+        tp_rank=0,
+        tp_size=1,
+        block_size=local_block_size,
+        engine_id=worker.engine_id,
+        is_mla=False,
+        is_mamba=False,
+        total_num_kv_heads=8,
+        attn_backends=[],
+    )
+
+    with pytest.raises(NotImplementedError, match=error):
+        worker.add_remote_agent(metadata, remote_tp_size=remote_tp_size)
+    worker.nixl_wrapper.add_remote_agent.assert_not_called()
+    worker.nixl_wrapper.prep_xfer_dlist.assert_not_called()
+    assert not worker.tp_mappings
+    assert not worker.dst_num_blocks
+    assert not worker.kv_caches_base_addr
+    with pytest.raises(KeyError):
+        worker.transfer_topo.get_engine_info(metadata.engine_id)
diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py
index 08ef476992d5..db7ecca62148 100644
--- a/tests/v1/kv_connector/unit/test_offloading_connector.py
+++ b/tests/v1/kv_connector/unit/test_offloading_connector.py
@@ -15,12 +15,26 @@
 from vllm.distributed.kv_events import BlockStored, KVEventBatch
 from vllm.platforms import current_platform
 
-CPU_BLOCK_SIZES: int = 64 if current_platform.is_xpu() else 48
+# The offloaded block size must stay a whole number of GPU blocks.
+CPU_BLOCK_SIZES: int
+if current_platform.is_xpu():
+    CPU_BLOCK_SIZES = 64
+elif current_platform.is_rocm():
+    # ROCM_AITER_UNIFIED_ATTN prefers 64-token GPU blocks, so 192 covers both that
+    # and the 16-token default.
+    CPU_BLOCK_SIZES = 192
+else:
+    CPU_BLOCK_SIZES = 48
+
 _ATTN_BACKENDS: list[str] = []
 if current_platform.is_cuda():
     _ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"]
 elif current_platform.is_rocm():
+    from vllm._aiter_ops import is_aiter_found_and_supported
+
     _ATTN_BACKENDS = ["TRITON_ATTN"]
+    if is_aiter_found_and_supported():
+        _ATTN_BACKENDS.append("ROCM_AITER_UNIFIED_ATTN")
 elif current_platform.is_xpu():
     _ATTN_BACKENDS = ["FLASH_ATTN", "TRITON_ATTN"]
 
@@ -62,7 +76,7 @@
 
 
 class MockSubscriber:
-    """Helper class to receive and verify published events"""
+    """Helper class to receive and verify published events."""
 
     def __init__(
         self,
@@ -107,7 +121,7 @@ def get_new_cpu_stored_events(self) -> list[BlockStored]:
         return cpu_stored_events
 
     def close(self):
-        """Clean up resources"""
+        """Clean up resources."""
         self.sub.close()
 
 
@@ -247,9 +261,7 @@ def test_cpu_offloading(
     cpu_block_size: int | None,
     uses_hma: bool,
 ) -> None:
-    """
-    Tests OffloadingConnector with CPUOffloadingSpec.
-    """
+    """Tests OffloadingConnector with CPUOffloadingSpec."""
     # configure OffloadingConnector (spec_name=CPUOffloadingSpec by default)
     extra_config: dict = {"cpu_bytes_to_use": 500 << 20}
     if cpu_block_size is not None:
diff --git a/tests/v1/kv_connector/unit/test_output_aggregator.py b/tests/v1/kv_connector/unit/test_output_aggregator.py
index b083ccef9819..e757a658425f 100644
--- a/tests/v1/kv_connector/unit/test_output_aggregator.py
+++ b/tests/v1/kv_connector/unit/test_output_aggregator.py
@@ -15,12 +15,14 @@ def __init__(
         finished_sending: set[str] | None = None,
         finished_recving: set[str] | None = None,
         invalid_block_ids: set[int] | None = None,
+        failed_recving: set[str] | None = None,
         expected_finished_count: int = 0,
     ):
         self.kv_connector_output = KVConnectorOutput(
             finished_sending=finished_sending,
             finished_recving=finished_recving,
             invalid_block_ids=invalid_block_ids or set(),
+            failed_recving=failed_recving or set(),
             expected_finished_count=expected_finished_count,
         )
 
@@ -73,7 +75,9 @@ def test_aggregate_workers_output():
 
     output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4})
     output2 = DummyModelRunnerOutput(
-        finished_recving={"req2"}, invalid_block_ids={4, 5}
+        finished_recving={"req2"},
+        invalid_block_ids={4, 5},
+        failed_recving={"req3"},
     )
 
     aggregated = aggregator.aggregate([output1, output2])
@@ -83,6 +87,12 @@ def test_aggregate_workers_output():
     assert aggregated.finished_sending is None
     assert aggregated.finished_recving == {"req2"}
     assert aggregated.invalid_block_ids == {3, 4, 5}
+    assert not aggregated.failed_recving
+
+    output1 = DummyModelRunnerOutput(finished_recving={"req3"})
+    output2 = DummyModelRunnerOutput(finished_recving={"req3"})
+    aggregated = aggregator.aggregate([output1, output2])
+    assert aggregated.kv_connector_output.failed_recving == {"req3"}
 
 
 def test_aggregate_workers_output_with_expected_finished_count():
diff --git a/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
index b656e0809543..8aa1e23d6fea 100644
--- a/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
+++ b/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
@@ -20,7 +20,6 @@
 
 def test_basic_lifecycle():
     """Test lifecycle of a Remote Decode request."""
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
@@ -117,7 +116,6 @@ def test_basic_lifecycle():
 
 def test_short_prompt_lifecycle():
     """Test lifecycle of a Remote Decode request with short prompt."""
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
@@ -164,7 +162,6 @@ def test_short_prompt_lifecycle():
 
 def test_prefix_cache_lifecycle():
     """Test that remote decode params still work with a prefix cache hit."""
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
@@ -224,7 +221,6 @@ def test_prefix_cache_lifecycle():
 
 def test_abort_during_kv_transfer():
     """Test aborting request does not release blocks for remote decode."""
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
index 7d21d975c173..caff9c9a490d 100644
--- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
@@ -30,7 +30,6 @@ def _num_waiting_requests(scheduler) -> int:
 
 def test_basic_lifecycle():
     """Test lifecycle of a remote prefill."""
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
@@ -151,7 +150,6 @@ def test_basic_lifecycle():
 
 def test_interleaved_lifecycle():
     """Test Remote Prefills Work Well With Other Requests."""
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
@@ -249,13 +247,11 @@ def test_interleaved_lifecycle():
 
 
 def test_no_spurious_prefix_caching():
-    """
-    With P/D, blocks can be allocated but uncomputed for
+    """With P/D, blocks can be allocated but uncomputed for
     multiple engine steps. This test confirms that we do
     not accidentally have cache hits against uncomputed
     blocks.
     """
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
@@ -320,7 +316,6 @@ def test_no_spurious_prefix_caching():
 
 def test_full_block_prompt():
     """Test that we handle a prompt that is the full block size."""
-
     vllm_config = create_vllm_config()
     scheduler = create_scheduler(vllm_config)
 
@@ -394,11 +389,9 @@ def test_full_block_prompt():
 
 
 def test_cannot_schedule_after_recv():
-    """
-    Test that we can handle no schedule after recv due to not
+    """Test that we can handle no schedule after recv due to not
     enough remaining KV blocks.
     """
-
     # NOTE: the KVCacheManager will use 1 null block.
     # So there are 5 total working blocks.
     TOTAL_NUM_BLOCKS = 6
@@ -501,11 +494,9 @@ def test_cannot_schedule_after_recv():
 
 
 def test_cannot_recv():
-    """
-    Test that we can handle no schedule KV block transfer due to not
+    """Test that we can handle no schedule KV block transfer due to not
     enough remaining KV blocks.
     """
-
     # NOTE: the KVCacheManager will use 1 null block.
     # So there are 5 total working blocks.
     TOTAL_NUM_BLOCKS = 6
@@ -740,3 +731,63 @@ def test_async_loads_both_admitted_when_pool_fits():
 
     for req in reqs:
         assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+
+
+def test_async_load_reserves_blocks_for_promotion_margin():
+    """An async load is not admitted unless the blocks its own promotion will
+    need are still free.
+
+    A parked load is allocated without lookahead slots, but promotion pads it
+    to ``1 + num_spec_tokens`` and asks for the lookahead margin on top. Without
+    reserving that margin, two loads can be admitted that together consume the
+    whole pool; the head then fails ``allocate_slots`` forever while the load
+    behind it is never reached, and with nothing running no block is ever freed.
+
+    req_a (4 blocks) and req_b (3 blocks) exactly fill the 7 usable blocks, so
+    req_b must be held back in WAITING holding no blocks, leaving req_a room to
+    be promoted and run.
+    """
+    vllm_config = create_vllm_config(num_speculative_tokens=3)
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    scheduler = create_scheduler(vllm_config, num_blocks=8)  # usable = 7
+
+    req_a = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=BLOCK_SIZE * 4,
+        do_remote_prefill=True,
+        max_tokens=1,
+    )
+    req_b = create_request(
+        request_id=2,
+        block_size=BLOCK_SIZE,
+        num_tokens=BLOCK_SIZE * 3,
+        do_remote_prefill=True,
+        max_tokens=1,
+    )
+    scheduler.add_request(req_a)
+    scheduler.add_request(req_b)
+
+    # Both get a full external hit, so each would hold its whole prompt.
+    with patch.object(
+        scheduler.connector,
+        "get_num_new_matched_tokens",
+        side_effect=[(BLOCK_SIZE * 4, True), (BLOCK_SIZE * 3, True)],
+    ):
+        scheduler.schedule()
+
+    assert req_a.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert req_b.status == RequestStatus.WAITING
+    req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
+        0
+    ].req_to_blocks
+    assert req_b.request_id not in req_to_blocks
+
+    # req_a's load lands: it must be promotable and actually get scheduled.
+    scheduler.update_from_output(
+        scheduler.schedule(),
+        create_model_runner_output([], finished_recving={req_a.request_id}),
+    )
+    scheduler_output = scheduler.schedule()
+    assert req_a.status == RequestStatus.RUNNING
+    assert scheduler_output.num_scheduled_tokens[req_a.request_id] > 0
diff --git a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py
index 69f9c2da415b..af4c7858fcad 100644
--- a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py
+++ b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py
@@ -19,6 +19,7 @@
     SimpleCPUOffloadConnector,
 )
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -90,3 +91,59 @@ def test_disk_mode_rejects_non_positive_capacity() -> None:
                 "disk_capacity_bytes": 0,
             }
         )
+
+
+def _make_ring_kv_cache_config(num_blocks: int = 16) -> KVCacheConfig:
+    """A paged full-attention group beside a per-request ring group."""
+    full = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    ring = CircularBufferSpec(
+        block_size=6, num_kv_heads=1, head_size=32, head_size_v=0, dtype=DTYPE
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=_BYTES_PER_BLOCK * num_blocks,
+                layers=["attn"],
+                layer_stride=_BYTES_PER_BLOCK * num_blocks,
+                block_stride=_BYTES_PER_BLOCK,
+            ),
+            KVCacheTensor(
+                size=ring.page_size_bytes * num_blocks,
+                layers=["ring"],
+                layer_stride=ring.page_size_bytes * num_blocks,
+                block_stride=ring.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["attn"], full),
+            KVCacheGroupSpec(["ring"], ring),
+        ],
+    )
+
+
+def test_ring_scratch_group_is_accepted_and_never_offloaded() -> None:
+    """The ring group's capacity is not a token granularity, it holds no
+    hashed blocks, and it never participates in CPU hits."""
+    vllm_config = create_vllm_config(
+        kv_connector="SimpleCPUOffloadConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={"cpu_bytes_to_use": _BYTES_PER_BLOCK * 8},
+    )
+    connector = SimpleCPUOffloadConnector(
+        vllm_config=vllm_config,
+        role=KVConnectorRole.SCHEDULER,
+        kv_cache_config=_make_ring_kv_cache_config(),
+    )
+    scheduler = connector.scheduler_manager
+    assert scheduler.fa_gidx == 0
+    cpu_groups = scheduler.cpu_kv_cache_config.kv_cache_groups
+    assert [g.kv_cache_spec.prefix_cacheable for g in cpu_groups] == [True, False]
+    # Only the paged group takes part in CPU-side hit lookup.
+    attention_groups = scheduler.cpu_coordinator.attention_groups
+    assert [g.group_ids for g in attention_groups] == [[0]]
diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py
index 52067bc2fa6f..3ddc568a524d 100644
--- a/tests/v1/kv_connector/unit/test_tp_mapping.py
+++ b/tests/v1/kv_connector/unit/test_tp_mapping.py
@@ -75,6 +75,7 @@ def test_source_ranks_p_gt_d(self):
 @pytest.mark.parametrize(
     "tp_rank,tp_size,remote_tp_size,dcp_size,remote_dcp_size,expected_ranks",
     [
+        (0, 1, 8, 1, 8, tuple(range(8))),
         (0, 4, 4, 1, 4, (0, 1, 2, 3)),
         (2, 4, 4, 4, 4, (2,)),
         (0, 2, 4, 2, 4, (0, 2)),
diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py
index 80cab94b3484..85ada84b0dfe 100644
--- a/tests/v1/kv_connector/unit/utils.py
+++ b/tests/v1/kv_connector/unit/utils.py
@@ -17,6 +17,7 @@
     KVTransferConfig,
     ModelConfig,
     SchedulerConfig,
+    SpeculativeConfig,
     VllmConfig,
 )
 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
@@ -106,6 +107,7 @@ def create_vllm_config(
     kv_connector_module_path: str | None = None,
     kv_role: str = "kv_consumer",
     disable_hybrid_kv_cache_manager: bool | None = None,
+    num_speculative_tokens: int | None = None,
 ) -> VllmConfig:
     """Initialize VllmConfig For Testing."""
     model_config = ModelConfig(
@@ -141,6 +143,11 @@ def create_vllm_config(
         kv_load_failure_policy=kv_load_failure_policy,
     )
     attention_config = AttentionConfig(backend=attention_backend)
+    speculative_config = (
+        SpeculativeConfig(model="ngram", num_speculative_tokens=num_speculative_tokens)
+        if num_speculative_tokens is not None
+        else None
+    )
     return VllmConfig(
         scheduler_config=scheduler_config,
         model_config=model_config,
@@ -148,6 +155,7 @@ def create_vllm_config(
         kv_transfer_config=kv_transfer_config,
         device_config=DeviceConfig("cpu"),
         attention_config=attention_config,
+        speculative_config=speculative_config,
     )
 
 
@@ -261,7 +269,6 @@ def create_model_runner_output(
     kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None,
 ) -> ModelRunnerOutput:
     """Make dummy model runner output for testing."""
-
     # Make request data.
     req_ids = [req.request_id for req in reqs]
     req_id_to_index = {req_id: idx for idx, req_id in enumerate(req_ids)}
@@ -519,6 +526,8 @@ def make_nixl_scheduler(
     Only sets the flags needed by the tests.  When *heartbeat=True* the
     scheduler-side heartbeat bookkeeping fields are also initialised.
     """
+    from types import SimpleNamespace
+
     from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
         NixlConnectorScheduler,
     )
@@ -530,6 +539,7 @@ def make_nixl_scheduler(
         block_size=16,
         mamba_enabled=has_mamba,
     )
+    sched.vllm_config = SimpleNamespace(num_prefill_lookahead_tokens=0)
 
     if heartbeat:
         sched._heartbeat_by_engine = {}
@@ -539,12 +549,14 @@ def make_nixl_scheduler(
         sched._heartbeat_interval = kv_lease_duration // 6
         # Fields touched by build_connector_meta / request_finished:
         sched._reqs_need_recv = {}
+        sched._hisparse_host_blocks_to_recv = {}
         sched._reqs_need_send = {}
         sched._reqs_in_batch = set()
         sched._reqs_not_processed = set()
         sched._reqs_need_save = {}
         sched.use_host_buffer = False
         sched.engine_id = "test-engine"
+        sched.transfer_tp_size = 1
         sched.side_channel_host = "localhost"
         sched.side_channel_port = 5555
         sched.blocks_per_sw = []
@@ -584,6 +596,7 @@ def make_nixl_push_scheduler(
     sched.decoder_kv_blocks_ttl = decoder_kv_blocks_ttl
     sched.use_host_buffer = False
     sched.engine_id = "decode-engine"
+    sched.transfer_tp_size = 1
     sched.side_channel_host = "127.0.0.1"
     sched.side_channel_port = 5600
     sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled
@@ -593,9 +606,11 @@ def make_nixl_push_scheduler(
         mamba_enabled=has_mamba,
     )
 
-    # vllm_config is consulted for parallel_config.tensor_parallel_size.
+    # vllm_config is consulted for parallel_config.tensor_parallel_size, and by
+    # `_prefill_backoff` on both the P and D prefill paths.
     vllm_config = MagicMock()
     vllm_config.parallel_config.tensor_parallel_size = 1
+    vllm_config.num_prefill_lookahead_tokens = 0
     sched.vllm_config = vllm_config
 
     # Push-specific state.
diff --git a/tests/v1/kv_offload/cpu/test_canonical_layout.py b/tests/v1/kv_offload/cpu/test_canonical_layout.py
index 4594e1eb0924..749bd6fac501 100644
--- a/tests/v1/kv_offload/cpu/test_canonical_layout.py
+++ b/tests/v1/kv_offload/cpu/test_canonical_layout.py
@@ -241,6 +241,7 @@ def canonical_view(rank: int, world_size: int) -> torch.Tensor:
         # The Triton load path dereferences CPU pointers on the GPU, which is
         # only legal on pinned memory; production pins via CPUOffloadingWorker
         pin_mmap_region(region)
+        assert region.is_pinned
         return region.create_next_canonical_view(_CANONICAL_PAGE)
 
     try:
diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py
index 5306d0557102..731d61a7e924 100644
--- a/tests/v1/kv_offload/cpu/test_gpu_worker.py
+++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py
@@ -49,6 +49,22 @@ def test_rocm_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None:
     )
 
 
+@pytest.mark.skipif(not gpu_worker.HAS_TRITON, reason="Requires Triton")
+def test_unpinned_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The Triton load path dereferences CPU pointers on the GPU, so pageable
+    host memory takes the DMA path even for pages where Triton would win."""
+    monkeypatch.setattr(gpu_worker.current_platform, "is_xpu", lambda: False)
+    monkeypatch.setattr(gpu_worker.current_platform, "is_rocm", lambda: False)
+
+    refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=512)]]
+    assert gpu_worker._select_swap_blocks_fn(refs, gpu_to_cpu=False) is not (
+        ops.swap_blocks_batch
+    )
+    assert gpu_worker._select_swap_blocks_fn(
+        refs, gpu_to_cpu=False, host_memory_is_pinned=False
+    ) is (ops.swap_blocks_batch)
+
+
 def test_worker_shutdown_releases_region_and_runs_both_handlers() -> None:
     """Both directions drain before the worker releases its shared region."""
     worker = CPUOffloadingWorker.__new__(CPUOffloadingWorker)
diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py
index a79f480ff551..0d78a18c1dbc 100644
--- a/tests/v1/kv_offload/cpu/test_manager.py
+++ b/tests/v1/kv_offload/cpu/test_manager.py
@@ -23,6 +23,7 @@
 )
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
+from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
 
 
 def make_req_context(
@@ -165,8 +166,7 @@ def test_cpu_eviction_removed_precedes_stored():
 
 @pytest.mark.parametrize("eviction_policy", ["lru", "arc"])
 def test_already_stored_chunk_not_evicted_during_prepare_store(eviction_policy):
-    """
-    Regression test: a chunk that is already stored must not be evicted
+    """Regression test: a chunk that is already stored must not be evicted
     by prepare_store() when it needs to make room for new chunks.
     Applies to both lru and arc policies.
 
@@ -366,9 +366,7 @@ def test_cpu_manager_clears_write_usage_after_failed_store():
 
 
 def test_cpu_manager():
-    """
-    Tests CPUOffloadingManager with lru policy.
-    """
+    """Tests CPUOffloadingManager with lru policy."""
     # initialize a CPU manager with a capacity of 4 chunks
     cpu_manager = make_cpu_manager(num_chunks=4, cache_policy="lru", enable_events=True)
 
@@ -430,31 +428,34 @@ def test_cpu_manager():
     assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is LookupResult.MISS
 
     # prepare load [2, 3]
-    prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
+    load_ctx = make_req_context("load-2-3")
+    prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), load_ctx)
     verify_load_output(prepare_load_output, [1, 2])
 
     # prepare store with no space ([2, 3] is being loaded)
     assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
 
     # complete load [2, 3]. Load changes the eviction list, making 2, 3 recent.
-    cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
+    cpu_manager.complete_load(to_keys([2, 3]), load_ctx)
+    cpu_manager.on_request_finished(load_ctx)
 
-    # prepare store [6, 7, 8] -> evicts [4, 5, 2] (oldest)
+    # prepare store [6, 7, 8] -> evicts [4, 5, 3] (oldest). Within
+    # the accessed prefix [2, 3], the tail is less valuable than the head.
     prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
     verify_store_output(
         prepare_store_output,
         ExpectedPrepareStoreOutput(
             keys_to_store=[6, 7, 8],
-            store_chunk_ids=[1, 0, 3],
-            evicted_keys=[4, 5, 2],
+            store_chunk_ids=[2, 0, 3],
+            evicted_keys=[4, 5, 3],
         ),
     )
 
     # complete store [6, 7, 8]
     cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
 
-    # touch [3, 6, 7] (move to end of LRU order)
-    cpu_manager.touch(to_keys([3, 6, 7]), _EMPTY_REQ_CTX)
+    # touch [2, 6, 7] (move to end of LRU order)
+    cpu_manager.touch(to_keys([2, 6, 7]), _EMPTY_REQ_CTX)
 
     # prepare store [7, 9] -> evicts [8] (oldest following previous touch)
     prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
@@ -477,7 +478,7 @@ def test_cpu_manager():
     verify_events(
         cpu_manager.take_events(),
         expected_stores=({3, 4, 5}, {6, 7, 8}),
-        expected_evictions=({4, 5, 2}, {8}),
+        expected_evictions=({4, 5, 3}, {8}),
     )
 
 
@@ -518,6 +519,65 @@ def test_prepare_load_preserves_key_order():
     manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX)  # order irrelevant
 
 
+def test_lru_batch_eviction_failure_is_atomic():
+    manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
+    policy = manager._policy
+    assert isinstance(policy, LRUCachePolicy)
+    keys = to_keys([1, 2, 3, 4])
+    assert manager.prepare_store(keys, _EMPTY_REQ_CTX) is not None
+    manager.complete_store(keys, _EMPTY_REQ_CTX)
+
+    protected = set(keys[1:])
+    assert policy.evict(2, protected) is None
+    assert all(manager.lookup(key, _EMPTY_REQ_CTX) is LookupResult.HIT for key in keys)
+
+    evicted = policy.evict(1, protected)
+    assert evicted is not None
+    assert [key for key, _ in evicted] == [keys[0]]
+
+
+def test_lru_repeated_pin_cycles_compact_lazy_heap_entries():
+    manager = make_cpu_manager(num_chunks=1, cache_policy="lru")
+    policy = manager._policy
+    assert isinstance(policy, LRUCachePolicy)
+    key = to_key(1)
+    assert manager.prepare_store([key], _EMPTY_REQ_CTX) is not None
+    manager.complete_store([key], _EMPTY_REQ_CTX)
+
+    for _ in range(128):
+        manager.prepare_load([key], _EMPTY_REQ_CTX)
+        manager.complete_load([key], _EMPTY_REQ_CTX)
+
+    assert len(policy._heap) < 64
+    evicted = policy.evict(1, set())
+    assert evicted is not None
+    assert [evicted_key for evicted_key, _ in evicted] == [key]
+
+
+def test_reset_discards_stale_request_access_classification():
+    manager = make_cpu_manager(num_chunks=1, cache_policy="arc")
+    policy = manager._policy
+    assert isinstance(policy, ARCCachePolicy)
+    key = to_key(1)
+    resumed_ctx = make_req_context("resumed")
+
+    assert manager.prepare_store([key], resumed_ctx) is not None
+    manager.complete_store([key], resumed_ctx)
+    manager.reset_cache()
+
+    replacement_ctx = make_req_context("replacement")
+    assert manager.prepare_store([key], replacement_ctx) is not None
+    manager.complete_store([key], replacement_ctx)
+    manager.on_request_finished(replacement_ctx)
+    assert key in policy.t1
+
+    manager.prepare_load([key], resumed_ctx)
+    manager.complete_load([key], resumed_ctx)
+    manager.on_request_finished(resumed_ctx)
+
+    assert key in policy.t2
+
+
 class TestARCPolicy:
     """Unit tests for CPUOffloadingManager with ARC eviction policy."""
 
@@ -534,8 +594,7 @@ def _make_manager(
         return manager, policy
 
     def test_basic(self):
-        """
-        Tests CPUOffloadingManager with arc policy.
+        """Tests CPUOffloadingManager with arc policy.
         Verifies that ARC handles store, load, and lookup operations correctly.
         """
         cpu_manager, arc_policy = self._make_manager()
@@ -574,8 +633,7 @@ def test_basic(self):
         assert len(arc_policy.t2) == 0
 
     def test_t1_to_t2_promotion(self):
-        """
-        Tests that accessing a chunk in T1 promotes it to T2 (frequent).
+        """Tests that accessing a chunk in T1 promotes it to T2 (frequent).
         This is a key feature of ARC's adaptive behavior.
         """
         cpu_manager, arc_policy = self._make_manager(enable_events=False)
@@ -596,8 +654,7 @@ def test_t1_to_t2_promotion(self):
         assert to_keys([1])[0] in arc_policy.t2
 
     def test_eviction_with_load(self):
-        """
-        Tests ARC eviction behavior similar to LRU test.
+        """Tests ARC eviction behavior similar to LRU test.
         Verifies that chunks being loaded (ref_cnt > 0) cannot be evicted.
         """
         cpu_manager, _ = self._make_manager()
@@ -637,8 +694,7 @@ def test_eviction_with_load(self):
         assert len(prepare_store_output.evicted_keys) >= 1
 
     def test_adaptive_target(self):
-        """
-        Tests ARC's adaptive target adjustment via ghost lists.
+        """Tests ARC's adaptive target adjustment via ghost lists.
         When a chunk in B1 (ghost list) is accessed, target_t1_size increases.
         When a chunk in B2 is accessed, target_t1_size decreases.
         """
@@ -665,8 +721,7 @@ def test_adaptive_target(self):
         assert arc_policy.target_t1_size > initial_target
 
     def test_t1_t2_eviction_policy(self):
-        """
-        Tests that ARC evicts from T1 or T2 based on target_t1_size.
+        """Tests that ARC evicts from T1 or T2 based on target_t1_size.
         If |T1| >= target_t1_size, evict from T1, otherwise from T2.
         """
         cpu_manager, arc_policy = self._make_manager(enable_events=False)
@@ -854,8 +909,7 @@ def test_batch_eviction_failure_is_atomic(self):
         assert not arc_policy.b2
 
     def test_ghost_list_bounds(self):
-        """
-        Tests that ghost lists (B1, B2) don't grow unbounded.
+        """Tests that ghost lists (B1, B2) don't grow unbounded.
         They should be capped at cache_capacity.
         """
         cpu_manager, arc_policy = self._make_manager(num_chunks=2, enable_events=False)
@@ -874,8 +928,7 @@ def test_ghost_list_bounds(self):
         assert len(arc_policy.b2) <= arc_policy.cache_capacity
 
     def test_touch_ordering(self):
-        """
-        Tests that touch() correctly updates access patterns.
+        """Tests that touch() correctly updates access patterns.
         Similar to LRU test but verifies T1/T2 ordering.
         """
         cpu_manager, arc_policy = self._make_manager()
@@ -907,8 +960,7 @@ def test_touch_ordering(self):
         )
 
     def test_failed_store(self):
-        """
-        Tests that failed store operations clean up correctly.
+        """Tests that failed store operations clean up correctly.
         Similar to LRU test but for ARC.
         """
         cpu_manager, arc_policy = self._make_manager()
@@ -936,8 +988,7 @@ def test_failed_store(self):
         assert evicted_hash in arc_policy.b1
 
     def test_full_scenario(self):
-        """
-        Comprehensive test covering multiple ARC operations in sequence.
+        """Comprehensive test covering multiple ARC operations in sequence.
         Similar to the full LRU test but adapted for ARC behavior.
         """
         cpu_manager, arc_policy = self._make_manager()
@@ -976,9 +1027,7 @@ def test_full_scenario(self):
 
 
 def test_filter_reused_manager():
-    """
-    Tests CPUOffloadingManager reuse filtering (store_threshold=2).
-    """
+    """Tests CPUOffloadingManager reuse filtering (store_threshold=2)."""
     manager = make_cpu_manager(
         num_chunks=4,
         cache_policy="lru",
@@ -1042,9 +1091,50 @@ def test_filter_reused_manager_oversized_offer_makes_progress():
     assert stored_keys == set(keys)
 
 
+def test_store_threshold_does_not_hide_ready_reuse():
+    manager = make_cpu_manager(
+        num_chunks=2,
+        cache_policy="lru",
+        store_threshold=2,
+        max_tracker_size=2,
+    )
+    head, tail, noise_1, noise_2, replacement = to_keys(list(range(5)))
+
+    seed_ctx = make_req_context("seed")
+    output = manager.prepare_store([head, tail], seed_ctx)
+    assert output is not None
+    assert not output.keys_to_store
+    output = manager.prepare_store([head, tail], seed_ctx)
+    assert output is not None
+    assert output.keys_to_store == [head, tail]
+    manager.complete_store([head, tail], seed_ctx)
+    manager.on_request_finished(seed_ctx)
+
+    # Age both resident keys out of the admission tracker without storing the
+    # noise keys. Reusing the tail recreates its counter at one, below the
+    # threshold, but must still refresh its cache recency.
+    output = manager.prepare_store([noise_1, noise_2], make_req_context("noise"))
+    assert output is not None
+    assert not output.keys_to_store
+    reuse_ctx = make_req_context("reuse")
+    skipped_before_reuse = manager.stores_skipped_in_current_batch
+    output = manager.prepare_store([tail], reuse_ctx)
+    assert output is not None
+    assert not output.keys_to_store
+    assert manager.stores_skipped_in_current_batch == skipped_before_reuse
+    manager.on_request_finished(reuse_ctx)
+
+    replacement_ctx = make_req_context("replacement")
+    output = manager.prepare_store([replacement], replacement_ctx)
+    assert output is not None
+    assert not output.keys_to_store
+    output = manager.prepare_store([replacement], replacement_ctx)
+    assert output is not None
+    assert output.evicted_keys == [head]
+
+
 def test_evictable_cache_chunk_count():
-    """
-    Verifies _num_evictable_cache_chunks is maintained correctly through the
+    """Verifies _num_evictable_cache_chunks is maintained correctly through the
     full store/load lifecycle, eviction, failed stores, concurrent loads,
     reset_cache, and the early-exit fast path in prepare_store.
     """
@@ -1156,3 +1246,318 @@ def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
     assert len(received) == 1
     assert received[0][0] == keys
     assert received[0][1] is ctx
+
+
+@pytest.mark.parametrize("finish_before_completion", [False, True])
+def test_request_finish_orders_lru_prefix_independent_of_store_completion(
+    finish_before_completion: bool,
+):
+    """A transfer's completion order must not make the prefix head LRU."""
+    manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
+    prefix = to_keys([1, 2, 3, 4])
+    ctx = make_req_context("prefix")
+
+    output = manager.prepare_store(prefix, ctx)
+    assert output is not None
+    if finish_before_completion:
+        manager.on_request_finished(ctx)
+    manager.complete_store(set(prefix), ctx)
+    if not finish_before_completion:
+        manager.on_request_finished(ctx)
+
+    output = manager.prepare_store(to_keys([5]), make_req_context("evict"))
+    assert output is not None
+    assert output.evicted_keys == to_keys([4])
+
+
+@pytest.mark.parametrize("finish_before_completion", [False, True])
+def test_request_finish_orders_lru_prefix_independent_of_load_completion(
+    finish_before_completion: bool,
+):
+    manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
+    prefix = to_keys([1, 2, 3, 4])
+
+    seed_ctx = make_req_context("seed")
+    output = manager.prepare_store(prefix, seed_ctx)
+    assert output is not None
+    manager.complete_store(prefix, seed_ctx)
+
+    reuse_ctx = make_req_context("reuse")
+    manager.prepare_load(prefix, reuse_ctx)
+    if finish_before_completion:
+        manager.on_request_finished(reuse_ctx)
+    manager.complete_load(set(prefix), reuse_ctx)
+    if not finish_before_completion:
+        manager.on_request_finished(reuse_ctx)
+
+    output = manager.prepare_store(to_keys([5]), make_req_context("evict"))
+    assert output is not None
+    assert output.evicted_keys == to_keys([4])
+
+
+def test_late_old_completion_does_not_override_newer_request_recency():
+    manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
+    old_keys = to_keys([1, 2])
+    new_keys = to_keys([3, 4])
+
+    old_ctx = make_req_context("old")
+    assert manager.prepare_store(old_keys, old_ctx) is not None
+    manager.on_request_finished(old_ctx)
+
+    new_ctx = make_req_context("new")
+    assert manager.prepare_store(new_keys, new_ctx) is not None
+    manager.complete_store(new_keys, new_ctx)
+    manager.on_request_finished(new_ctx)
+
+    # Physical completion is late, but the logical access is still older.
+    manager.complete_store(old_keys, old_ctx)
+
+    output = manager.prepare_store(to_keys([5]), make_req_context("evict"))
+    assert output is not None
+    assert output.evicted_keys == [old_keys[-1]]
+
+
+def test_arc_counts_reuse_once_and_keeps_insertions_in_t1():
+    manager = make_cpu_manager(num_chunks=4, cache_policy="arc")
+    policy = manager._policy
+    assert isinstance(policy, ARCCachePolicy)
+    prefix = to_keys([1, 2, 3, 4])
+
+    seed_ctx = make_req_context("seed")
+    assert manager.prepare_store(prefix, seed_ctx) is not None
+    manager.complete_store(prefix, seed_ctx)
+    # Repeated store offers from the same request must not become ARC hits.
+    assert manager.prepare_store(prefix, seed_ctx) is not None
+    # Nor should an internal read used to cascade the new chunks to a tier.
+    manager.prepare_load(prefix, seed_ctx)
+    manager.complete_load(prefix, seed_ctx)
+    manager.on_request_finished(seed_ctx)
+
+    assert list(policy.t1) == list(reversed(prefix))
+    assert not policy.t2
+
+    reuse_ctx = make_req_context("reuse")
+    manager.prepare_load(prefix, reuse_ctx)
+    manager.complete_load(prefix, reuse_ctx)
+    # A duplicate internal pin, as used by tiering cascades, is deduplicated.
+    manager.prepare_load(prefix, reuse_ctx)
+    manager.complete_load(prefix, reuse_ctx)
+    manager.on_request_finished(reuse_ctx)
+
+    assert not policy.t1
+    assert list(policy.t2) == list(reversed(prefix))
+
+
+def test_arc_reuse_wins_if_same_request_later_reinserts_key():
+    manager = make_cpu_manager(num_chunks=1, cache_policy="arc")
+    policy = manager._policy
+    assert isinstance(policy, ARCCachePolicy)
+    key_1, key_2 = to_keys([1, 2])
+
+    seed_ctx = make_req_context("seed")
+    assert manager.prepare_store([key_1], seed_ctx) is not None
+    manager.complete_store([key_1], seed_ctx)
+    manager.on_request_finished(seed_ctx)
+
+    reuse_ctx = make_req_context("reuse-and-reinsert")
+    manager.prepare_load([key_1], reuse_ctx)
+    manager.complete_load([key_1], reuse_ctx)
+
+    other_ctx = make_req_context("evict-reused-key")
+    assert manager.prepare_store([key_2], other_ctx) is not None
+    manager.complete_store([key_2], other_ctx)
+    manager.on_request_finished(other_ctx)
+
+    assert manager.prepare_store([key_1], reuse_ctx) is not None
+    manager.complete_store([key_1], reuse_ctx)
+    manager.on_request_finished(reuse_ctx)
+
+    assert key_1 not in policy.t1
+    assert list(policy.t2) == [key_1]
+
+
+def test_request_finish_forwards_grouped_deduplicated_accesses(monkeypatch):
+    manager = make_cpu_manager(num_chunks=5)
+    existing = make_offload_key(b"existing", 0)
+    group_0_new = make_offload_key(b"group-0-new", 0)
+    group_1_head = make_offload_key(b"group-1-head", 1)
+    group_1_tail = make_offload_key(b"group-1-tail", 1)
+
+    seed_ctx = make_req_context("seed")
+    assert manager.prepare_store([existing], seed_ctx) is not None
+    manager.complete_store([existing], seed_ctx)
+
+    ctx = make_req_context("grouped")
+    # Group 1 is observed first, but policy finalization follows group index.
+    offered = [group_1_head, existing, group_0_new]
+    assert manager.prepare_store(offered, ctx) is not None
+    # Later offers extend a group's prefix and may repeat earlier keys.
+    assert (
+        manager.prepare_store([group_0_new, group_1_head, group_1_tail], ctx)
+        is not None
+    )
+
+    received = []
+
+    def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context):
+        received.append((key_groups, insertion_only_keys, reused_keys, req_context))
+
+    monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished)
+    manager.on_request_finished(ctx)
+    manager.on_request_finished(ctx)
+
+    assert received == [
+        (
+            ((existing, group_0_new), (group_1_head, group_1_tail)),
+            {group_0_new, group_1_head, group_1_tail},
+            {existing},
+            ctx,
+        )
+    ]
+
+
+def test_failed_store_records_access_once_and_uses_store_miss_hook(monkeypatch):
+    manager = make_cpu_manager(num_chunks=1)
+    resident, candidate = to_keys([1, 2])
+
+    seed_ctx = make_req_context("seed")
+    assert manager.prepare_store([resident], seed_ctx) is not None
+    manager.complete_store([resident], seed_ctx)
+
+    pin_ctx = make_req_context("pin")
+    manager.prepare_load([resident], pin_ctx)
+
+    record_calls = 0
+    original_record = manager._record_request_cache_access
+
+    def record_access(*args, **kwargs):
+        nonlocal record_calls
+        record_calls += 1
+        return original_record(*args, **kwargs)
+
+    store_miss_calls = []
+
+    def on_store_miss(keys, req_context):
+        store_miss_calls.append((list(keys), req_context))
+
+    finalized = []
+
+    def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context):
+        finalized.append((key_groups, insertion_only_keys, reused_keys, req_context))
+
+    monkeypatch.setattr(manager, "_record_request_cache_access", record_access)
+    monkeypatch.setattr(manager._policy, "on_store_miss", on_store_miss)
+    monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished)
+
+    ctx = make_req_context("failed-store")
+    assert manager.prepare_store([candidate, resident], ctx) is None
+    assert record_calls == 1
+    assert store_miss_calls == [([candidate], ctx)]
+
+    manager.on_request_finished(ctx)
+    assert finalized == [
+        (((candidate, resident),), set(), {resident}, ctx),
+    ]
+
+    manager.complete_load([resident], pin_ctx)
+
+
+def test_arc_ghost_hit_adapts_once_per_request_before_insertion():
+    manager = make_cpu_manager(num_chunks=2, cache_policy="arc")
+    policy = manager._policy
+    assert isinstance(policy, ARCCachePolicy)
+    keys = to_keys([1, 2, 3])
+
+    seed_ctx = make_req_context("seed")
+    assert manager.prepare_store(keys[:2], seed_ctx) is not None
+    manager.complete_store(keys[:2], seed_ctx)
+    manager.on_request_finished(seed_ctx)
+
+    evict_ctx = make_req_context("create-ghost")
+    output = manager.prepare_store([keys[2]], evict_ctx)
+    assert output is not None
+    [ghost_key] = output.evicted_keys
+    manager.complete_store([keys[2]], evict_ctx)
+    manager.on_request_finished(evict_ctx)
+    assert ghost_key in policy.b1
+
+    resident_keys = [key for key in keys if key != ghost_key]
+    pin_ctx = make_req_context("pin-residents")
+    manager.prepare_load(resident_keys, pin_ctx)
+
+    ghost_ctx = make_req_context("ghost-hit")
+    target_before = policy.target_t1_size
+    assert manager.prepare_store([ghost_key], ghost_ctx) is None
+    target_after_first_offer = policy.target_t1_size
+    assert target_after_first_offer > target_before
+    assert manager.prepare_store([ghost_key], ghost_ctx) is None
+    assert policy.target_t1_size == target_after_first_offer
+
+    manager.complete_load(resident_keys, pin_ctx)
+    output = manager.prepare_store([ghost_key], ghost_ctx)
+    assert output is not None
+    assert policy.target_t1_size == target_after_first_offer
+    manager.complete_store([ghost_key], ghost_ctx)
+    manager.on_request_finished(ghost_ctx)
+    assert ghost_key in policy.t1
+    assert ghost_key not in policy.t2
+
+
+def test_arc_pending_chunk_is_not_a_frequency_hit():
+    manager = make_cpu_manager(num_chunks=2, cache_policy="arc")
+    policy = manager._policy
+    assert isinstance(policy, ARCCachePolicy)
+    key = to_key(1)
+
+    writer_ctx = make_req_context("writer")
+    assert manager.prepare_store([key], writer_ctx) is not None
+
+    observer_ctx = make_req_context("pending-observer")
+    output = manager.prepare_store([key], observer_ctx)
+    assert output is not None and not output.keys_to_store
+    manager.on_request_finished(observer_ctx)
+
+    chunk = policy.get(key)
+    assert chunk is not None and not chunk.is_ready
+    assert key in policy.t1
+    assert key not in policy.t2
+
+
+def test_request_key_positions_override_store_observation_order():
+    manager = make_cpu_manager(num_chunks=2, cache_policy="lru")
+    head, tail = to_keys([1, 2])
+    ctx = make_req_context("out-of-order-store")
+    ctx.set_offload_key_position(head, 16)
+    ctx.set_offload_key_position(tail, 32)
+
+    # A backfill or partial-tail path may offer a later key first.
+    assert manager.prepare_store([tail], ctx) is not None
+    assert manager.prepare_store([head], ctx) is not None
+    manager.on_request_finished(ctx)
+    manager.complete_store({head, tail}, ctx)
+
+    output = manager.prepare_store([to_key(3)], make_req_context("evict-tail"))
+    assert output is not None
+    assert output.evicted_keys == [tail]
+
+
+def test_request_key_positions_order_tails_across_kv_groups():
+    manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
+    group_0_head = make_offload_key(b"group-0-head", 0)
+    group_0_tail = make_offload_key(b"group-0-tail", 0)
+    group_1_head = make_offload_key(b"group-1-head", 1)
+    group_1_tail = make_offload_key(b"group-1-tail", 1)
+    keys = [group_0_head, group_0_tail, group_1_head, group_1_tail]
+    ctx = make_req_context("hybrid-groups")
+    for key in (group_0_head, group_1_head):
+        ctx.set_offload_key_position(key, 16)
+    for key in (group_0_tail, group_1_tail):
+        ctx.set_offload_key_position(key, 32)
+
+    assert manager.prepare_store(keys, ctx) is not None
+    manager.complete_store(keys, ctx)
+    manager.on_request_finished(ctx)
+
+    output = manager.prepare_store(to_keys([10, 11]), make_req_context("evict"))
+    assert output is not None
+    assert set(output.evicted_keys) == {group_0_tail, group_1_tail}
diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py
index c7ddac4feffa..f84e1bd7296c 100644
--- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py
+++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py
@@ -10,11 +10,13 @@
 import threading
 import time
 import uuid
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, call
 
 import pytest
 
 from vllm.utils.system_utils import get_mp_context
+from vllm.v1.kv_offload.cpu import gpu_worker
+from vllm.v1.kv_offload.cpu import shared_offload_region as region_module
 from vllm.v1.kv_offload.cpu.shared_offload_region import (
     SharedOffloadRegion,
     _wait_for_file_size,
@@ -192,28 +194,43 @@ def _mp_barrier_construct_and_hold(
     barrier,
     fill_value: int,
     done_queue,
+    replicated: bool = False,
 ) -> None:
     """Construct with a real cross-process barrier, write, then hold the
     mapping (no cleanup) until the parent SIGKILLs this process."""
     try:
+        populate_calls = 0
+        get_populate_write_fn = region_module._get_populate_write_fn
+
+        def track_population(mmap_obj):
+            nonlocal populate_calls
+            populate_calls += 1
+            return get_populate_write_fn(mmap_obj)
+
+        region_module._get_populate_write_fn = track_population
         region = SharedOffloadRegion(
             engine_id=engine_id,
             num_chunks=2,
-            rank=rank,
-            kv_bytes_per_chunk=num_workers * PAGE_SIZE,
+            rank=0 if replicated else rank,
+            kv_bytes_per_chunk=PAGE_SIZE if replicated else num_workers * PAGE_SIZE,
             cpu_page_size=PAGE_SIZE,
             barrier=lambda: barrier.wait(30),
+            populate_only_on_creator=replicated,
         )
         # The constructor's barrier precedes the creator's unlink.
         barrier.wait(30)
         t = region.create_next_worker_view(PAGE_SIZE)
-        t[:, :] = fill_value
+        if not replicated or rank == 0:
+            t[:, :] = fill_value
+        barrier.wait(30)
         done_queue.put(
             {
                 "rank": rank,
                 "error": None,
                 "inode": os.fstat(region.fd).st_ino,
                 "path_exists": os.path.exists(region.mmap_path),
+                "populate_calls": populate_calls,
+                "contents_match": bool((t == (1 if replicated else fill_value)).all()),
             }
         )
         time.sleep(60)  # hold the mapping until killed
@@ -752,6 +769,107 @@ def test_cleanup_idempotent(iid):
     r.cleanup()  # must be a no-op
 
 
+def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch):
+    """cleanup() must release every chunk from a split host registration."""
+    r = _make_region(iid)
+    cudart = MagicMock()
+    cudart.cudaHostUnregister.return_value = MagicMock(value=0)
+    monkeypatch.setattr(region_module, "current_platform", MagicMock())
+    region_module.current_platform.is_cuda_alike.return_value = True
+    monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: cudart)
+    r.pinned_addresses = [0x100000, 0x200000, 0x300000]
+    r.is_pinned = True
+
+    r.cleanup()
+
+    assert cudart.cudaHostUnregister.call_args_list == [
+        call(0x300000),
+        call(0x200000),
+        call(0x100000),
+    ]
+
+
+@pytest.fixture
+def host_register(monkeypatch):
+    """Pretend to run on CUDA with host registrations capped at seven pages.
+
+    Registration goes to a mock CudaRTLibrary and the region's cleanup to a
+    separate mock torch runtime. Both stay installed through the cleanup that
+    ``_region`` runs on exit.
+    """
+    cudart = MagicMock()
+    cudart.cudaHostRegister.return_value = 0
+    cudart.cudaHostUnregister.return_value = 0
+    cudart.cudaGetLastError.return_value = 2
+    torch_cudart = MagicMock()
+    torch_cudart.cudaHostUnregister.return_value = MagicMock(value=0)
+    monkeypatch.setattr(gpu_worker.current_platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(gpu_worker, "MAX_HOST_REGISTER_CHUNK_BYTES", 7 * PAGE_SIZE)
+    monkeypatch.setattr(gpu_worker, "CudaRTLibrary", lambda: cudart)
+    monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: torch_cudart)
+    return cudart, torch_cudart
+
+
+def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register):
+    """Chunks end on row boundaries: 3-page rows under a 7-page cap register
+    as 6 + 6 + 3 pages, where a raw byte cap would give 7 + 7 + 1."""
+    cudart, torch_cudart = host_register
+    with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region:
+        gpu_worker.pin_mmap_region(region)
+        base = region._base.data_ptr()
+        assert cudart.mock_calls == [
+            call.cudaHostRegister(base, 6 * PAGE_SIZE),
+            call.cudaHostRegister(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE),
+            call.cudaHostRegister(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE),
+        ]
+        assert region.is_pinned
+
+    assert torch_cudart.cudaHostUnregister.call_args_list == [
+        call(base + 12 * PAGE_SIZE),
+        call(base + 6 * PAGE_SIZE),
+        call(base),
+    ]
+
+
+@pytest.mark.parametrize("fail_at", [0, 1, 2])
+def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail_at):
+    """A failed chunk drains its pending error and unregisters the chunks
+    before it on the same handle, so the region is pinned whole or not at all."""
+    cudart, torch_cudart = host_register
+    cudart.cudaHostRegister.side_effect = [0] * fail_at + [2]
+    with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region:
+        gpu_worker.pin_mmap_region(region)
+        base = region._base.data_ptr()
+        assert not region.is_pinned
+        assert region.pinned_addresses == []
+
+    registered = [base + i * 6 * PAGE_SIZE for i in range(fail_at)]
+    failed_size = min(6, 15 - 6 * fail_at) * PAGE_SIZE
+    assert cudart.mock_calls == [
+        *(call.cudaHostRegister(address, 6 * PAGE_SIZE) for address in registered),
+        call.cudaHostRegister(base + fail_at * 6 * PAGE_SIZE, failed_size),
+        call.cudaGetLastError(),
+        *(call.cudaHostUnregister(address) for address in reversed(registered)),
+    ]
+    torch_cudart.cudaHostUnregister.assert_not_called()
+
+
+def test_pin_mmap_region_without_cuda_runtime_stays_pageable(
+    iid, host_register, monkeypatch
+):
+    """A CUDA runtime that cannot be loaded leaves the region pageable instead
+    of failing worker startup."""
+
+    def missing_runtime():
+        raise AssertionError("libcudart is not loaded in the current process")
+
+    monkeypatch.setattr(gpu_worker, "CudaRTLibrary", missing_runtime)
+    with _region(iid) as region:
+        gpu_worker.pin_mmap_region(region)
+        assert not region.is_pinned
+        assert region.pinned_addresses == []
+
+
 def test_cleanup_after_create_next_worker_view_releases_mmap(iid):
     """cleanup() must close the mmap even after create_next_worker_view was called.
     create_next_worker_view returns a view that shares storage with _base; both must be
@@ -811,11 +929,53 @@ def test_wait_for_file_size_timeout(tmp_path):
         os.close(fd)
 
 
+@pytest.mark.skip_global_cleanup
+def test_wait_for_file_size_rejects_unlinked_file(tmp_path):
+    path = tmp_path / "unlinked.mmap"
+    fd = os.open(str(path), os.O_CREAT | os.O_RDWR, 0o600)
+    try:
+        os.unlink(path)
+        with pytest.raises(RuntimeError, match="creator failed"):
+            _wait_for_file_size(fd, PAGE_SIZE, timeout=1.0)
+    finally:
+        os.close(fd)
+
+
 # ---------------------------------------------------------------------------
 # Constructor — capacity validation
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.skip_global_cleanup
+@pytest.mark.skipif(not os.path.isdir("/dev/shm"), reason="requires /dev/shm")
+def test_creator_memory_check_runs_only_for_creator(iid):
+    checked_sizes: list[int] = []
+    creator = SharedOffloadRegion(
+        engine_id=iid,
+        num_chunks=4,
+        rank=0,
+        kv_bytes_per_chunk=PAGE_SIZE,
+        cpu_page_size=PAGE_SIZE,
+        creator_memory_check=checked_sizes.append,
+    )
+    joiner: SharedOffloadRegion | None = None
+    try:
+        joiner = SharedOffloadRegion(
+            engine_id=iid,
+            num_chunks=4,
+            rank=0,
+            kv_bytes_per_chunk=PAGE_SIZE,
+            cpu_page_size=PAGE_SIZE,
+            creator_memory_check=checked_sizes.append,
+        )
+        assert checked_sizes == [4 * PAGE_SIZE]
+    finally:
+        if joiner is not None:
+            joiner.cleanup()
+        creator.cleanup()
+        _cleanup_file(creator.mmap_path)
+
+
 def test_insufficient_space_raises_clear_error(monkeypatch):
     """A failed creator capacity check must clean up and give a clear error."""
     import vllm.v1.kv_offload.cpu.shared_offload_region as region
@@ -828,12 +988,13 @@ def test_insufficient_space_raises_clear_error(monkeypatch):
     monkeypatch.setattr(region.os, "open", mock_open)
     monkeypatch.setattr(region.os, "unlink", mock_unlink)
     monkeypatch.setattr(region.os, "close", mock_close)
+    mock_check = MagicMock(
+        side_effect=RuntimeError("Insufficient space in /dev/shm: 30 GB required.")
+    )
     monkeypatch.setattr(
         region,
         "check_shm_free_space",
-        lambda *a, **kw: (_ for _ in ()).throw(
-            RuntimeError("Insufficient space in /dev/shm: 30 GB required.")
-        ),
+        mock_check,
     )
 
     with pytest.raises(RuntimeError, match="Insufficient space"):
@@ -847,6 +1008,10 @@ def test_insufficient_space_raises_clear_error(monkeypatch):
 
     mock_unlink.assert_called_once_with(mmap_path)
     mock_close.assert_called_once_with(9999)
+    mock_check.assert_called_once_with(
+        4 * PAGE_SIZE,
+        allocation_name="CPU KV offload shared region in /dev/shm",
+    )
 
 
 def test_ftruncate_failure_cleans_up_creator(monkeypatch):
@@ -916,7 +1081,8 @@ def test_barrier_failure_unlinks_creator_and_raises(iid):
     assert not os.path.exists(path)
 
 
-def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):
+@pytest.mark.parametrize("replicated", [False, True])
+def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated):
     """With a real cross-process barrier every worker maps one shared inode,
     the name is gone while they run, and SIGKILL leaks nothing."""
     num_workers = 2
@@ -927,7 +1093,7 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):
     procs = [
         ctx.Process(
             target=_mp_barrier_construct_and_hold,
-            args=(iid, rank, num_workers, barrier, rank + 1, done_queue),
+            args=(iid, rank, num_workers, barrier, rank + 1, done_queue, replicated),
         )
         for rank in range(num_workers)
     ]
@@ -939,6 +1105,10 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):
             assert r["error"] is None, f"rank {r['rank']}: {r['error']}"
         assert len({r["inode"] for r in results}) == 1, "workers split onto two files"
         assert not any(r["path_exists"] for r in results)
+        assert all(r["contents_match"] for r in results)
+        assert sum(r["populate_calls"] for r in results) == (
+            1 if replicated else num_workers
+        )
         assert not os.path.exists(path)
     finally:
         for p in procs:
@@ -950,23 +1120,36 @@ def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):
         assert restarted._creator is True, "restart must be able to create anew"
 
 
-def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch):
+@pytest.mark.parametrize("failure_phase", ["shm", "memory", "population"])
+def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_phase):
     """A worker that dies before the rendezvous must still arrive at the
     barrier, or its peers block in the collective until it times out."""
-    import vllm.v1.kv_offload.cpu.shared_offload_region as region
-
-    monkeypatch.setattr(
-        region,
-        "check_shm_free_space",
-        MagicMock(side_effect=RuntimeError("Insufficient space")),
+    failure = MagicMock(side_effect=ValueError("Insufficient space"))
+    memory_check = failure if failure_phase == "memory" else None
+    if failure_phase == "shm":
+        monkeypatch.setattr(region_module, "check_shm_free_space", failure)
+    elif failure_phase == "population":
+        monkeypatch.setattr(region_module, "_get_populate_write_fn", failure)
+    path = f"/dev/shm/vllm_offload_{iid}.mmap"
+    seen_at_barrier = []
+    barrier = MagicMock(
+        side_effect=lambda: seen_at_barrier.append(os.path.exists(path))
     )
-    barrier = MagicMock()
 
-    with pytest.raises(RuntimeError, match="Insufficient space"):
-        _make_region(iid, barrier=barrier)
+    with pytest.raises(ValueError, match="Insufficient space"):
+        SharedOffloadRegion(
+            engine_id=iid,
+            num_chunks=1,
+            rank=0,
+            kv_bytes_per_chunk=PAGE_SIZE,
+            cpu_page_size=PAGE_SIZE,
+            barrier=barrier,
+            creator_memory_check=memory_check,
+            populate_only_on_creator=True,
+        )
 
     barrier.assert_called_once_with()
-    assert not os.path.exists(f"/dev/shm/vllm_offload_{iid}.mmap")
+    assert seen_at_barrier == [False]
 
 
 def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch):
diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py
index 251cef325143..01ef756cae53 100644
--- a/tests/v1/kv_offload/test_factory.py
+++ b/tests/v1/kv_offload/test_factory.py
@@ -62,7 +62,7 @@ def _make_offloading_config(
         normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
 
     if groups is None:
-        groups = (OffloadingGroupConfig(16, ("layer",)),)
+        groups = (OffloadingGroupConfig(16, ("layer",), 0),)
 
     return OffloadingConfig(
         groups=groups,
@@ -275,11 +275,12 @@ def fake_region_ctor(**kwargs):
 
 @pytest.mark.parametrize("world_size", [2, 4, 8])
 def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: int):
-    # On shared-region (CUDA-alike) platforms the default spec now honors
+    # On CUDA the default spec honors
     # replicated layout: a single MLA copy (num_copies=1), matching tiering.
     import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
 
     monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False)
     worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
     spec = _create_spec(
         cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
@@ -295,19 +296,19 @@ def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: in
     assert spec.num_chunks == 8
 
 
+@pytest.mark.parametrize("rocm", [False, True], ids=["non-cuda-alike", "rocm"])
 @pytest.mark.parametrize("world_size", [2, 4, 8])
 def test_cpu_spec_replicated_disabled_without_shared_region(
-    monkeypatch, world_size: int
+    monkeypatch, world_size: int, rocm: bool
 ):
-    # Data-loss guard: non-CUDA-alike platforms keep a per-rank private pinned
+    # Data-loss guard: platforms without shared regions keep a private pinned
     # tensor (no shared medium), so replicated layout MUST stay off. Otherwise
     # the rank-0 writer gate would ack rank>0 stores without writing, leaving
     # those private buffers empty and corrupting subsequent loads.
     import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
 
-    monkeypatch.setattr(
-        cpu_spec_module.current_platform, "is_cuda_alike", lambda: False
-    )
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm)
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm)
     worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
     spec = _create_spec(
         cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2,
@@ -324,17 +325,26 @@ def test_cpu_spec_replicated_disabled_without_shared_region(
 
 
 @pytest.mark.parametrize("config_replicated", [True, False])
-@pytest.mark.parametrize("cuda_alike", [True, False])
+@pytest.mark.parametrize(
+    ("cuda_alike", "rocm", "shared_region"),
+    [(True, False, True), (True, True, False), (False, False, False)],
+    ids=["cuda", "rocm", "non-cuda-alike"],
+)
 def test_cpu_spec_replicated_layout_truth_matrix(
-    monkeypatch, cuda_alike: bool, config_replicated: bool
+    monkeypatch,
+    cuda_alike: bool,
+    rocm: bool,
+    shared_region: bool,
+    config_replicated: bool,
 ):
     # replicated_layout is enabled iff the config gate passes AND the deployment
-    # actually allocates on the shared region (CUDA-alike).
+    # actually allocates on the shared region.
     import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
 
     monkeypatch.setattr(
         cpu_spec_module.current_platform, "is_cuda_alike", lambda: cuda_alike
     )
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm)
     worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
     spec = _create_spec(
         cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
@@ -344,10 +354,10 @@ def test_cpu_spec_replicated_layout_truth_matrix(
     )
 
     assert isinstance(spec, CPUOffloadingSpec)
-    assert spec.replicated_layout is (cuda_alike and config_replicated)
+    assert spec.replicated_layout is (shared_region and config_replicated)
 
 
-def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
+def test_cpu_spec_create_worker_uses_mmap_on_cuda(monkeypatch):
     import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
 
     worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
@@ -371,6 +381,7 @@ def fake_worker_ctor(**kwargs):
         return MagicMock()
 
     monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False)
     monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
     monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
     monkeypatch.setattr(
@@ -386,7 +397,10 @@ def fake_worker_ctor(**kwargs):
     assert worker_calls[0]["mmap_region"] is region
 
 
-def test_cpu_spec_create_worker_uses_tensor_path_off_cuda_alike(monkeypatch):
+@pytest.mark.parametrize("rocm", [False, True], ids=["non-cuda-alike", "rocm"])
+def test_cpu_spec_create_worker_uses_tensor_path_without_shared_region(
+    monkeypatch, rocm
+):
     import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
 
     spec = _create_spec(worker_kv_bytes_per_block=4096, world_size=4)
@@ -403,15 +417,14 @@ def fake_worker_ctor(**kwargs):
         worker_calls.append(kwargs)
         return MagicMock()
 
-    monkeypatch.setattr(
-        cpu_spec_module.current_platform, "is_cuda_alike", lambda: False
-    )
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm)
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm)
     monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
     monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
 
     spec.create_worker(MagicMock())
 
-    # Non-CUDA-alike platforms keep the per-rank pinned-tensor path.
+    # ROCm and non-CUDA-alike platforms use the per-rank pinned-tensor path.
     assert region_calls == []
     assert worker_calls[0]["mmap_region"] is None
 
@@ -420,7 +433,7 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch):
     import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
 
     # worker_kv_bytes_per_block=0 yields num_chunks=0; a zero-byte region cannot
-    # be mmap'd, so even on CUDA-alike this must fall back to the tensor path.
+    # be mmap'd, so even on CUDA this must fall back to the tensor path.
     spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4)
     assert isinstance(spec, CPUOffloadingSpec)
     assert spec.num_chunks == 0
@@ -429,6 +442,7 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch):
     worker_calls: list[dict[str, Any]] = []
 
     monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False)
     monkeypatch.setattr(
         cpu_spec_module,
         "SharedOffloadRegion",
@@ -461,6 +475,7 @@ def test_cpu_spec_create_worker_rank_assignment(
     import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
 
     monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False)
     worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
     spec = _create_spec(
         cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
@@ -493,8 +508,8 @@ def test_offloading_spec_has_replicated_layout_default():
 
 def test_offloading_spec_uses_normalized_chunk_geometry():
     groups = (
-        OffloadingGroupConfig(12, ("full_layer",)),
-        OffloadingGroupConfig(16, ("mla_layer",)),
+        OffloadingGroupConfig(12, ("full_layer",), 0),
+        OffloadingGroupConfig(16, ("mla_layer",), 1),
     )
     spec = _create_spec(
         groups=groups,
diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py
index 26da4079b470..cfe6660675f6 100644
--- a/tests/v1/kv_offload/test_file_mapper.py
+++ b/tests/v1/kv_offload/test_file_mapper.py
@@ -26,8 +26,11 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
             OffloadingGroupConfig(
                 tokens_per_block=tokens_per_block,
                 layer_names=(layer_name,),
+                group_id=group_id,
+            )
+            for group_id, (tokens_per_block, layer_name) in enumerate(
+                kwargs.get("groups", ())
             )
-            for tokens_per_block, layer_name in kwargs.get("groups", ())
         ),
         worker_kv_bytes_per_block=0,
         enable_kv_cache_events=False,
@@ -73,8 +76,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
 
 
 def test_get_file_name_full_structure():
-    """
-    Path must match: <base_path>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash_hex>.bin
+    """Path must match: <base_path>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash_hex>.bin.
 
     Concretely:
       - The segment immediately after base_path must end with `_r3`
diff --git a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py
index 630fcdef4407..9f9e330b863a 100644
--- a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py
+++ b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-P2PConnector proxy server for OffloadingConnector + TieringOffloadingSpec.
+"""P2PConnector proxy server for OffloadingConnector + TieringOffloadingSpec.
 
 Unlike NixlConnector (which returns remote_host/remote_port in the prefill
 response), OffloadingConnector does not embed connector coordinates in its
diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py
index d3bacc8f326f..e40f16ad7393 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py
+++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py
@@ -132,7 +132,7 @@ def test_write_blocks_increments_transfer_id(self):
         assert tid1 != tid2
 
     def test_poll_empty_when_no_inflight(self):
-        """poll returns empty when nothing is inflight."""
+        """Poll returns empty when nothing is inflight."""
         transport = self._make_transport()
         result = transport.poll()
         assert result == PollResult(done=(), failed=())
@@ -213,7 +213,7 @@ def test_poll_peer_id_scopes_to_peer(self):
         assert tid2 not in transport._inflight
 
     def test_cancel_removes_inflight(self):
-        """cancel removes transfers and releases handles."""
+        """Cancel removes transfers and releases handles."""
         transport = self._make_transport()
         transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
 
@@ -226,7 +226,7 @@ def test_cancel_removes_inflight(self):
         transport._agent.release_xfer_handle.assert_called()
 
     def test_cancel_ignores_unknown_ids(self):
-        """cancel with unknown IDs doesn't crash."""
+        """Cancel with unknown IDs doesn't crash."""
         transport = self._make_transport()
         assert transport.cancel([999, 1000]) == []
         assert transport.cancel([999, 1000], mode="wait") == []
@@ -286,7 +286,7 @@ def test_add_and_remove_remote_peer(self):
         transport._agent.remove_remote_agent.assert_called()
 
     def test_close_releases_everything(self):
-        """close releases all handles and clears state."""
+        """Close releases all handles and clears state."""
         transport = self._make_transport()
         transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
         transport.write_blocks("peer:1", [0], [1])
diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py
index dc0e384d7c2f..8e806582fd9f 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_manager.py
+++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py
@@ -117,6 +117,8 @@ def _make_manager() -> P2PSecondaryTierManager:
     mgr._sessions = {}
     mgr._kv_to_session = {}
     mgr._unbound_stores = {}
+    mgr._reaped_stores = {}
+    mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S
     mgr._failed_serve_ctxs = []
     return mgr
 
@@ -178,6 +180,40 @@ def test_seed_resolved_after_init_none_hash(self, monkeypatch):
         assert mgr._get_hash_seed() == "random-seed-abc"
 
 
+class TestBackpressureRejected:
+    def _patch_transports(self, monkeypatch) -> None:
+        monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object())
+        monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object())
+        monkeypatch.setattr(
+            manager_module.FileMapper,
+            "from_offloading_spec",
+            lambda **k: SimpleNamespace(get_run_config=lambda: {}),
+        )
+
+    def test_explicit_backpressure_detector_rejected(self, monkeypatch):
+        """An explicitly configured detector must fail fast for P2P.
+
+        The generic store-latency detector cannot fail fast in PD mode and
+        rendezvous time makes store latency an unreliable pressure signal,
+        so P2P rejects it rather than degrading silently.
+        """
+        self._patch_transports(monkeypatch)
+        with pytest.raises(ValueError, match="not supported for the P2P"):
+            P2PSecondaryTierManager(
+                offloading_spec=_init_offloading_spec(),
+                primary_kv_view=memoryview(bytearray(16)),
+                backpressure_detector=object(),
+            )
+
+    def test_no_detector_constructs_normally(self, monkeypatch):
+        self._patch_transports(monkeypatch)
+        mgr = P2PSecondaryTierManager(
+            offloading_spec=_init_offloading_spec(),
+            primary_kv_view=memoryview(bytearray(16)),
+        )
+        assert mgr.bp_detector is None
+
+
 # ---------------------------------------------------------------------------
 # Tests for _peer_id_from_params
 # ---------------------------------------------------------------------------
@@ -759,9 +795,9 @@ def test_unbound_store_kept_within_timeout(self):
         assert "req-fresh" in mgr._unbound_stores
 
     def test_unbound_store_reaped_after_timeout(self):
-        """Unbound stores past _UNBOUND_STORE_TIMEOUT_S surface as failed
-        and their kv_request_id lands in _failed_req_ids so a late
-        FetchMsg/lookup doesn't try to satisfy them."""
+        """Unbound stores past the reap deadline surface as failed and their
+        kv_request_id lands in _reaped_stores so a late FetchMsg is rejected
+        instead of parking demand nothing can satisfy."""
         from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
 
         mgr = self._make()
@@ -779,7 +815,7 @@ def test_unbound_store_reaped_after_timeout(self):
         # 2 baseline + 2 buffered stores
         assert JobResult(job_id=10, success=False) in results
         assert JobResult(job_id=11, success=False) in results
-        assert "req-stale" in mgr._failed_req_ids
+        assert "req-stale" in mgr._reaped_stores
 
     def test_submit_store_parks_unbound_batch(self):
         """submit_store on an unbound id appends a batch with a fresh
@@ -1204,6 +1240,75 @@ def test_both_loads_succeed(self):
         assert 201 in b_ok, f"B loads succeeded: {b_ok}"
         assert 200 in b_ok, f"B stores succeeded: {b_ok}"
 
+    def test_late_fetch_after_reap_fails_immediately(self):
+        """A fetch for a reaped kv_request_id fails in one round trip.
+
+        The producer drops parked blocks once unbound_store_timeout_s
+        expires. A consumer whose fetch arrives after that must learn on the
+        next tick via TransferDoneMsg(success=False) — not stall for
+        _LOAD_TIMEOUT_S and then take the abort path.
+        """
+        from vllm.v1.kv_offload.tiering.p2p.session.client import _LOAD_TIMEOUT_S
+
+        mgr_a, mgr_b = _build_paired_managers()
+        kv_id = "req-late"
+
+        # Record every message type leaving the consumer, so the assertion
+        # below can prove no abort was ever needed.
+        sent_from_a: list[str] = []
+        drain = mgr_a._control._drain_outbound_to
+
+        def recording_drain(peer_local_id: str):
+            out = drain(peer_local_id)
+            sent_from_a.extend(msg.get("type") for _, msg in out)
+            return out
+
+        mgr_a._control._drain_outbound_to = recording_drain  # type: ignore[method-assign]
+
+        # Producer parks the blocks, then reaps them before any fetch lands.
+        mgr_b.submit_store(
+            _job_metadata(
+                job_id=200,
+                keys=[b"b-block"],
+                chunk_ids=[0],
+                kv_params={"remote_decoder": {"kv_request_id": kv_id}},
+            )
+        )
+        mgr_b._unbound_store_timeout_s = 0.0
+        reap_results = list(mgr_b.get_finished_jobs())
+        assert JobResult(job_id=200, success=False) in reap_results
+        assert kv_id in mgr_b._reaped_stores
+
+        # Consumer now asks for the blocks that no longer exist.
+        consumer_params = {
+            "remote_prefiller": {
+                "kv_request_id": kv_id,
+                "remote_host": "B",
+                "remote_port": 2,
+            },
+        }
+        mgr_a.on_new_request(_req_context(consumer_params))
+        mgr_a.submit_load(
+            _job_metadata(
+                job_id=101,
+                keys=[b"b-block"],
+                chunk_ids=[0],
+                kv_params=consumer_params,
+            )
+        )
+
+        started = time.monotonic()
+        all_a: list[JobResult] = []
+        for _ in range(8):
+            all_a.extend(list(mgr_a.get_finished_jobs()))
+            list(mgr_b.get_finished_jobs())
+        elapsed = time.monotonic() - started
+
+        assert JobResult(job_id=101, success=False) in all_a, all_a
+        assert kv_id in mgr_a._failed_req_ids
+        assert "abort_fetch" not in sent_from_a, sent_from_a
+        assert elapsed < _LOAD_TIMEOUT_S
+
 
 # ---------------------------------------------------------------------------
 # _accept_new_peers — duplicate connection rejection
diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py
index 2e4415fbb75d..a09445668fd3 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py
+++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py
@@ -827,7 +827,7 @@ def test_request_blocks_clears_probe_cache(self):
         assert fresh[0][LookupMsg.KEYS] == [b"hA"]
 
     def test_flush_uses_work_list_not_full_scan(self):
-        """flush drains a work-list rather than scanning every live request.
+        """Flush drains a work-list rather than scanning every live request.
 
         A request with no newly-registered keys is not revisited: after a
         flush the work-list is empty, an idle re-flush sends nothing, and a
diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py
index ff5cbcd5dd44..3649c8df0488 100644
--- a/tests/v1/kv_offload/tiering/test_async_lookup.py
+++ b/tests/v1/kv_offload/tiering/test_async_lookup.py
@@ -8,7 +8,7 @@
 import pytest
 
 from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key
-from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
+from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager, LookupPhase
 
 
 def _key(i: int) -> OffloadKey:
@@ -46,10 +46,13 @@ def test_new_key_returns_none(self):
     def test_found_key_returns_true(self):
         mgr = InMemoryLookupManager(existing_keys={_key(1)})
         assert mgr.lookup(_key(1), _ctx()) is None
+        assert mgr._lookup_state[_key(1)].phase is LookupPhase.PENDING
         mgr.flush()
+        assert mgr._lookup_state[_key(1)].phase is LookupPhase.IN_FLIGHT
         mgr._results_ready.wait()
         mgr._results_ready.clear()
         assert mgr.lookup(_key(1), _ctx()) is True
+        assert mgr._lookup_state[_key(1)].phase is LookupPhase.RESOLVED
         mgr.shutdown()
 
     def test_not_found_key_returns_false(self):
@@ -106,6 +109,68 @@ def test_cleanup_preserves_shared_entries(self):
         assert _key(1) not in mgr._lookup_state
         mgr.shutdown()
 
+    def test_cleanup_reuses_in_flight_probe(self, monkeypatch: pytest.MonkeyPatch):
+        """A replacement request shares the in-flight probe and its verdict."""
+        key = _key(1)
+        mgr = InMemoryLookupManager(existing_keys={key})
+        ctx_b = _ctx("req_b")
+        probe_started = threading.Event()
+        release_probe = threading.Event()
+        batch_lookup = mgr.batch_lookup
+
+        def blocking_lookup(keys, req_context):
+            probe_started.set()
+            if not release_probe.wait(timeout=5):
+                raise TimeoutError("Test did not release the backend probe")
+            return batch_lookup(keys, req_context)
+
+        monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup)
+        try:
+            assert mgr.lookup(key, _ctx("req_a")) is None
+            mgr.flush()
+            assert probe_started.wait(timeout=5)
+            assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT
+            mgr.cleanup("req_a")
+            assert mgr.lookup(key, ctx_b) is None
+            assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT
+            mgr.flush()
+
+            release_probe.set()
+            batch = mgr._pending_results.get(timeout=5)
+            mgr._pending_results.put(batch)
+            replacement_result = mgr.lookup(key, ctx_b)
+        finally:
+            release_probe.set()
+            mgr.shutdown()
+
+        assert mgr.batch_lookup_calls == 1
+        assert replacement_result is True
+
+    @pytest.mark.parametrize(
+        "reclaim_at_shutdown", [False, True], ids=["flush", "shutdown"]
+    )
+    def test_unclaimed_probe_reclaimed_without_lookup(self, reclaim_at_shutdown: bool):
+        """A completed orphan is released by flush or shutdown without a lookup."""
+        key = _key(1)
+        mgr = InMemoryLookupManager(existing_keys={key})
+        try:
+            mgr.lookup(key, _ctx("req_a"))
+            mgr.flush()
+            mgr.cleanup("req_a")
+            assert key in mgr._lookup_state
+            assert not mgr._req_keys
+
+            if not reclaim_at_shutdown:
+                batch = mgr._pending_results.get(timeout=5)
+                mgr._pending_results.put(batch)
+                mgr.flush()
+                assert key not in mgr._lookup_state
+        finally:
+            mgr.shutdown()
+
+        assert not mgr._lookup_state
+        assert mgr._pending_results.empty()
+
     def test_stale_result_ignored_after_cleanup_and_key_reuse(self):
         key = _key(1)
         mgr = InMemoryLookupManager()
@@ -119,11 +184,12 @@ def test_stale_result_ignored_after_cleanup_and_key_reuse(self):
         generation = mgr._lookup_state[key].generation
         assert generation != stale_generation
 
+        mgr.flush()
+        mgr._results_ready.wait()
+        mgr._results_ready.clear()
+        current_result = mgr._pending_results.get(timeout=5)
         mgr._pending_results.put([(key, stale_generation, True)])
-        mgr.drain_results()
-        assert mgr.lookup(key, ctx_b) is None
-
-        mgr._pending_results.put([(key, generation, False)])
+        mgr._pending_results.put(current_result)
         mgr.drain_results()
         assert mgr.lookup(key, ctx_b) is False
         mgr.shutdown()
diff --git a/tests/v1/kv_offload/tiering/test_backpressure.py b/tests/v1/kv_offload/tiering/test_backpressure.py
new file mode 100644
index 000000000000..b809bc4dc490
--- /dev/null
+++ b/tests/v1/kv_offload/tiering/test_backpressure.py
@@ -0,0 +1,706 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for back-pressure detection in TieringOffloadingManager.
+
+These tests use a DelayedSecondaryTierManager that holds completed jobs
+until explicitly released, allowing precise control over when the manager
+observes store completions and their apparent latency.
+"""
+
+import time
+from collections.abc import Iterable
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm.v1.kv_offload.base import (
+    LookupResult,
+    OffloadKey,
+    OffloadPolicy,
+    ReqContext,
+    RequestOffloadingContext,
+    ScheduleEndContext,
+    make_offload_key,
+)
+from vllm.v1.kv_offload.tiering.backpressure import (
+    DropAccountingPolicy,
+    DropStorePolicy,
+    EMABackpressureDetector,
+    ThrottledDropPolicy,
+)
+from vllm.v1.kv_offload.tiering.base import (
+    JobResult,
+    SecondaryTierManager,
+    TieringOffloadingMetrics,
+    TransferJob,
+)
+from vllm.v1.kv_offload.tiering.manager import (
+    CPUPrimaryTierOffloadingManager,
+    TieringOffloadingManager,
+)
+
+_BP_EMA_ALPHA = EMABackpressureDetector.DEFAULT_ALPHA
+_BP_WARMUP = EMABackpressureDetector.DEFAULT_WARMUP_COMPLETIONS
+# Water marks in s/MiB, scaled for the test's tiny 16-byte blocks.
+_BP_HIGH_WATER_S = 1000.0
+_BP_LOW_WATER_S = 500.0
+
+_CTX = ReqContext(req_id="test")
+_MOCK_OFFLOADING_SPEC = MagicMock()
+
+
+def _mock_mmap_region(num_blocks: int, row_bytes: int = 16):
+    mock = MagicMock()
+    view = memoryview(torch.zeros((num_blocks, row_bytes), dtype=torch.int8).numpy())
+    mock.create_kv_memoryview.return_value = view
+    return mock
+
+
+def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]:
+    return [make_offload_key(str(i).encode(), 0) for i in int_ids]
+
+
+class DelayedSecondaryTierManager(SecondaryTierManager):
+    """Secondary tier that holds completed store jobs until released.
+
+    Jobs are stored immediately but their completion is not reported
+    by get_finished_jobs() until release_jobs() is called. This lets
+    tests control the apparent latency of store operations by adjusting
+    submit_time on the JobMetadata before triggering the completion poll.
+    """
+
+    def __init__(
+        self, offloading_spec, primary_kv_view, tier_type, backpressure_detector=None
+    ):
+        super().__init__(
+            offloading_spec,
+            primary_kv_view,
+            tier_type,
+            backpressure_detector=backpressure_detector,
+        )
+        self.blocks: dict[OffloadKey, bool] = {}
+        self._held_jobs: list[JobResult] = []
+        self._released_jobs: list[JobResult] = []
+        self._request_policy: OffloadPolicy | None = None
+
+    def lookup(self, key, req_context):
+        return LookupResult.HIT if key in self.blocks else LookupResult.MISS
+
+    def submit_store(self, job_metadata: TransferJob) -> None:
+        for key in job_metadata.keys:
+            self.blocks[key] = True
+        self._held_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
+
+    def submit_load(self, job_metadata: TransferJob) -> None:
+        self._released_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
+
+    def get_finished_jobs(self) -> Iterable[JobResult]:
+        result = self._released_jobs
+        self._released_jobs = []
+        return result
+
+    def release_jobs(self):
+        self._released_jobs.extend(self._held_jobs)
+        self._held_jobs.clear()
+
+    def on_new_request(self, req_context):
+        if self._request_policy is not None:
+            return RequestOffloadingContext(policy=self._request_policy)
+        return RequestOffloadingContext()
+
+    def drain_jobs(self):
+        self.release_jobs()
+
+    def has_pending_work(self):
+        return bool(self._held_jobs)
+
+    def get_num_blocks(self):
+        return len(self.blocks)
+
+
+class TestDefaultConfig:
+    def test_fs_gets_local_watermarks(self):
+        cfg = EMABackpressureDetector.default_config("fs")
+        assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S
+        assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S
+
+    def test_obj_gets_network_watermarks(self):
+        cfg = EMABackpressureDetector.default_config("obj")
+        assert cfg["high_water_s"] == EMABackpressureDetector.NETWORK_HIGH_WATER_S
+        assert cfg["low_water_s"] == EMABackpressureDetector.NETWORK_LOW_WATER_S
+
+    def test_p2p_gets_local_watermarks(self):
+        cfg = EMABackpressureDetector.default_config("p2p")
+        assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S
+        assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S
+
+    def test_fs_with_remote_locality_gets_network_watermarks(self):
+        cfg = EMABackpressureDetector.default_config("fs", locality="REMOTE")
+        assert cfg["high_water_s"] == EMABackpressureDetector.NETWORK_HIGH_WATER_S
+        assert cfg["low_water_s"] == EMABackpressureDetector.NETWORK_LOW_WATER_S
+
+    def test_fs_with_local_locality_gets_local_watermarks(self):
+        cfg = EMABackpressureDetector.default_config("fs", locality="LOCAL")
+        assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S
+        assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S
+
+
+class TestIdleDecay:
+    """Tests for EMA idle decay — recovery from pressure without completions."""
+
+    def _make_detector(self, **kwargs):
+        defaults = dict(
+            high_water_s=_BP_HIGH_WATER_S,
+            low_water_s=_BP_LOW_WATER_S,
+            decay_half_life_s=2.0,
+        )
+        defaults.update(kwargs)
+        return EMABackpressureDetector(**defaults)
+
+    def test_ema_decays_while_under_pressure(self):
+        bp = self._make_detector()
+        bp._completions = _BP_WARMUP
+        bp._ema = _BP_HIGH_WATER_S * 2
+        bp._under_pressure = True
+        bp._last_update = time.monotonic()
+
+        with patch("time.monotonic", return_value=bp._last_update + 4.1):
+            # ~2.05 half-lives → EMA * 0.24, well below low water (500)
+            assert bp.is_under_pressure() is False
+            assert bp.store_latency_ema < _BP_LOW_WATER_S
+
+    def test_ema_stays_pressured_if_not_enough_decay(self):
+        bp = self._make_detector()
+        bp._completions = _BP_WARMUP
+        bp._ema = _BP_HIGH_WATER_S * 4
+        bp._under_pressure = True
+        bp._last_update = time.monotonic()
+
+        with patch("time.monotonic", return_value=bp._last_update + 2.0):
+            # 2s = 1 half-life → EMA * 0.5 = HIGH * 2, still above LOW
+            assert bp.is_under_pressure() is True
+            assert bp.store_latency_ema == pytest.approx(_BP_HIGH_WATER_S * 2, rel=0.01)
+
+    def test_no_decay_before_warmup(self):
+        bp = self._make_detector()
+        bp._ema = _BP_HIGH_WATER_S * 2
+        bp._under_pressure = True
+        bp._last_update = time.monotonic()
+
+        with patch("time.monotonic", return_value=bp._last_update + 10.0):
+            # Still in warmup (_completions=0 < _warmup_completions),
+            # so decay should not be applied in is_under_pressure.
+            assert bp.is_under_pressure() is True
+            assert bp.store_latency_ema == _BP_HIGH_WATER_S * 2
+
+    def test_no_decay_applied_during_new_sample(self):
+        """on_store_completed applies the new sample directly without
+        idle decay — decay only runs in is_under_pressure()."""
+        bp = self._make_detector()
+        bp._completions = _BP_WARMUP
+        bp._ema = _BP_HIGH_WATER_S * 2
+        bp._under_pressure = True
+        t0 = time.monotonic()
+        bp._last_update = t0
+
+        block_bytes = 16
+        with patch("time.monotonic", return_value=t0 + 4.0):
+            bp.on_store_completed(0.0, block_bytes)
+            # No decay in on_store_completed: EMA = alpha*0 + (1-alpha)*2000
+            expected = (1 - _BP_EMA_ALPHA) * _BP_HIGH_WATER_S * 2
+            assert bp.store_latency_ema == pytest.approx(expected, rel=0.01)
+
+    def test_reset_clears_last_update(self):
+        bp = self._make_detector()
+        bp._last_update = time.monotonic()
+        bp.reset()
+        assert bp._last_update == 0.0
+
+
+class TestBackpressure:
+    @pytest.fixture
+    def setup(self):
+        mock_region = _mock_mmap_region(20)
+        self.primary = CPUPrimaryTierOffloadingManager(
+            num_chunks=20, mmap_region=mock_region
+        )
+        mock_view = mock_region.create_kv_memoryview()
+        self.tier = DelayedSecondaryTierManager(
+            offloading_spec=_MOCK_OFFLOADING_SPEC,
+            primary_kv_view=mock_view,
+            tier_type="delayed",
+            backpressure_detector=EMABackpressureDetector(
+                high_water_s=_BP_HIGH_WATER_S,
+                low_water_s=_BP_LOW_WATER_S,
+            ),
+        )
+        self.manager = TieringOffloadingManager(
+            primary_tier=self.primary,
+            secondary_tiers=[self.tier],
+        )
+
+    def _start_request(self, ctx=_CTX):
+        if ctx.req_id not in self.manager._req_state:
+            self.manager.on_new_request(ctx)
+
+    def _store_blocks(self, keys, ctx=_CTX):
+        self._start_request(ctx)
+        result = self.manager.prepare_store(keys, ctx)
+        assert result is not None
+        self.manager.complete_store(keys, ctx, success=True)
+
+    def _simulate_on_schedule_end(self):
+        ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=())
+        self.manager.on_schedule_end(ctx)
+
+    def _backdate_held_jobs(self, age_s: float):
+        """Set submit_time on all in-flight store jobs so they appear
+        to have taken ``age_s`` seconds when completed."""
+        self.manager._processed_jobs_this_step = False
+        now = time.monotonic()
+        for job_id, meta in self.manager._jobs.items():
+            tj = meta.transfer_job
+            if not tj.is_promotion and tj.submit_time > 0:
+                tj.submit_time = now - age_s
+
+    def test_ema_updates_on_store_completion(self, setup):
+        bp = self.tier.bp_detector
+        fast_latency = 0.01
+        # EMA is in s/MiB; with 16-byte blocks, scale = MiB / block_bytes.
+        scale = EMABackpressureDetector._MIB / self.tier.block_size_bytes
+
+        # During warmup, EMA stays at 0.
+        for i in range(_BP_WARMUP - 1):
+            keys = to_keys([300 + i])
+            self._store_blocks(keys)
+            self._backdate_held_jobs(fast_latency)
+            self.tier.release_jobs()
+            self._simulate_on_schedule_end()
+            assert bp.store_latency_ema == 0.0
+
+        # Final warmup sample seeds EMA with the mean.
+        keys = to_keys([400])
+        self._store_blocks(keys)
+        self._backdate_held_jobs(fast_latency)
+        self.tier.release_jobs()
+        self._simulate_on_schedule_end()
+        fast_s_per_mib = fast_latency * scale
+        assert bp.store_latency_ema == pytest.approx(fast_s_per_mib, rel=0.1)
+
+        # After warmup, EMA updates normally.
+        new_latency = 0.05
+        keys = to_keys([401])
+        self._store_blocks(keys)
+        self._backdate_held_jobs(new_latency)
+        self.tier.release_jobs()
+        self._simulate_on_schedule_end()
+        new_s_per_mib = new_latency * scale
+        expected = _BP_EMA_ALPHA * new_s_per_mib + (1 - _BP_EMA_ALPHA) * fast_s_per_mib
+        assert bp.store_latency_ema == pytest.approx(expected, rel=0.1)
+
+    def test_pressure_activates_above_high_water(self, setup):
+        bp = self.tier.bp_detector
+
+        # Warm up the detector with fast completions first.
+        for i in range(_BP_WARMUP):
+            keys = to_keys([200 + i])
+            self._store_blocks(keys)
+            self._backdate_held_jobs(0.001)
+            self.tier.release_jobs()
+            self._simulate_on_schedule_end()
+
+        keys = to_keys([100])
+        self._store_blocks(keys)
+        self._simulate_on_schedule_end()
+
+        # 0.1s for a single 16-byte block → ~6553 s/MiB, well above high water.
+        self._backdate_held_jobs(0.1)
+        self.tier.release_jobs()
+        self._simulate_on_schedule_end()
+
+        assert bp.store_latency_ema > _BP_HIGH_WATER_S
+        assert bp.is_under_pressure() is True
+
+    def test_pressure_clears_below_low_water(self, setup):
+        bp = self.tier.bp_detector
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 2
+        bp._under_pressure = True
+        bp._completions = _BP_WARMUP
+
+        # Drive several fast completions to bring EMA below low water.
+        block_id = 100
+        for _ in range(50):
+            if bp.store_latency_ema < _BP_LOW_WATER_S:
+                break
+            keys = to_keys([block_id])
+            block_id += 1
+            self._store_blocks(keys)
+            # Pressure is on, so the cascade was skipped. Force-submit
+            # a store job manually to get a completion to feed the EMA.
+            job_meta = self.manager.create_store_job(keys, _CTX)
+            job_meta.submit_time = time.monotonic() - 0.001
+            self.tier.submit_store(job_meta)
+            self.tier.release_jobs()
+            self._simulate_on_schedule_end()
+        else:
+            pytest.fail(f"EMA did not recover below low water: {bp.store_latency_ema}")
+
+        assert bp.is_under_pressure() is False
+
+    def test_hysteresis_prevents_oscillation(self, setup):
+        bp = self.tier.bp_detector
+        bp._completions = _BP_WARMUP
+        block_bytes = self.tier.block_size_bytes
+
+        # Set EMA between low and high water marks (in s/MiB).
+        mid = (_BP_LOW_WATER_S + _BP_HIGH_WATER_S) / 2
+        bp.store_latency_ema = mid
+        bp._last_update = time.monotonic()
+
+        # Raw seconds that produce `mid` s/MiB for a given number of blocks.
+        def mid_latency_s(num_blocks: int) -> float:
+            return mid * (num_blocks * block_bytes / EMABackpressureDetector._MIB)
+
+        # When not under pressure, mid-range EMA should not activate.
+        bp._under_pressure = False
+        keys = to_keys(range(2))
+        self._store_blocks(keys)
+        self._backdate_held_jobs(mid_latency_s(len(keys)))
+        self.tier.release_jobs()
+        self._simulate_on_schedule_end()
+        assert bp.is_under_pressure() is False
+
+        # When under pressure, mid-range EMA should not deactivate.
+        bp._under_pressure = True
+        bp._last_update = time.monotonic()
+        keys2 = to_keys(range(10, 12))
+        # Force-submit since pressure is on (cascade would skip).
+        for k in keys2:
+            self.primary.prepare_store([k], _CTX)
+            self.primary.complete_store([k], _CTX)
+        job_meta = self.manager.create_store_job(keys2, _CTX)
+        job_meta.submit_time = time.monotonic() - mid_latency_s(len(keys2))
+        self.tier.submit_store(job_meta)
+        self.tier.release_jobs()
+        self._simulate_on_schedule_end()
+        assert bp.is_under_pressure() is True
+
+    def test_stores_skipped_under_pressure(self, setup):
+        bp = self.tier.bp_detector
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 4
+        initial_blocks = self.tier.get_num_blocks()
+
+        keys = to_keys(range(50, 53))
+        self._store_blocks(keys)
+
+        # Blocks should be in primary but NOT in secondary.
+        for k in keys:
+            assert self.primary.lookup(k, _CTX) is LookupResult.HIT
+        assert self.tier.get_num_blocks() == initial_blocks
+
+    def test_stores_resume_after_pressure_clears(self, setup):
+        bp = self.tier.bp_detector
+
+        # Start under pressure — stores skipped.
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 4
+        keys1 = to_keys(range(60, 62))
+        self._store_blocks(keys1)
+        assert all(k not in self.tier.blocks for k in keys1)
+
+        # Clear pressure — next store should cascade.
+        bp._under_pressure = False
+        keys2 = to_keys(range(70, 72))
+        self._store_blocks(keys2)
+        assert all(k in self.tier.blocks for k in keys2)
+
+    def test_dropped_store_count_tracked(self, setup):
+        bp = self.tier.bp_detector
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 4
+        policy = bp.policy
+        assert policy.pop_stores_dropped() == (0, 0)
+
+        self._store_blocks(to_keys([80]))
+        assert policy._stores_dropped == 1
+        assert policy._blocks_dropped == 1
+
+        self._store_blocks(to_keys([81, 82]))
+        assert policy._stores_dropped == 2
+        assert policy._blocks_dropped == 3
+
+    def test_metrics_reported_via_get_stats(self, setup):
+        bp = self.tier.bp_detector
+        bp._under_pressure = True
+        bp.store_latency_ema = 2.5
+        policy = bp.policy
+        policy._stores_dropped = 5
+        policy._blocks_dropped = 12
+
+        stats = self.manager.get_stats()
+        assert stats is not None
+        reduced = stats.reduce()
+
+        ema_key = TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA
+        stores_key = TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED
+        blocks_key = TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED
+        assert reduced[f"{ema_key}:('1:delayed',)"] == pytest.approx(2.5)
+        assert reduced[f"{stores_key}:('1:delayed',)"] == 5
+        assert reduced[f"{blocks_key}:('1:delayed',)"] == 12
+
+        # Dropped counts reset after get_stats.
+        assert policy._stores_dropped == 0
+        assert policy._blocks_dropped == 0
+
+    def test_reset_cache_clears_backpressure(self, setup):
+        bp = self.tier.bp_detector
+        bp.store_latency_ema = 5.0
+        bp._under_pressure = True
+
+        self.manager.reset_cache()
+
+        assert bp.store_latency_ema == 0.0
+        assert bp.is_under_pressure() is False
+
+    def test_request_level_tier_respects_backpressure(self, setup):
+        """Request-level cascade also skips pressured tiers."""
+        # Store blocks in primary first.
+        existing = to_keys(range(3))
+        self._store_blocks(existing)
+        self.tier.release_jobs()
+        self._simulate_on_schedule_end()
+        self._simulate_on_schedule_end()
+
+        # Make tier request-level for a new request.
+        self.tier._request_policy = OffloadPolicy.REQUEST_LEVEL
+        ctx = ReqContext(req_id="req_rl_bp")
+        self.manager.on_new_request(ctx)
+
+        # Activate pressure.
+        bp = self.tier.bp_detector
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 4
+        initial_blocks = self.tier.get_num_blocks()
+
+        # prepare_store with existing + new blocks.
+        new = to_keys(range(10, 12))
+        all_keys = existing + new
+        result = self.manager.prepare_store(all_keys, ctx)
+        assert result is not None
+
+        # Existing blocks would normally cascade to request-level tier,
+        # but should be skipped under pressure.
+        assert self.tier.get_num_blocks() == initial_blocks
+
+    def test_loads_continue_under_pressure(self, setup):
+        """Loads from a pressured tier still work."""
+        blocks = to_keys(range(3))
+        for b in blocks:
+            self.tier.blocks[b] = True
+
+        bp = self.tier.bp_detector
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 4
+
+        # Lookups should still initiate promotion.
+        for b in blocks:
+            result = self.manager.lookup(b, _CTX)
+            assert result is LookupResult.HIT_PENDING
+
+
+class TestThrottledDropPolicy:
+    """Tests for proportional throttling under pressure."""
+
+    def _make_detector(self, **kwargs):
+        defaults = dict(
+            high_water_s=_BP_HIGH_WATER_S,
+            low_water_s=_BP_LOW_WATER_S,
+            cooldown_s=0.0,
+        )
+        defaults.update(kwargs)
+        return EMABackpressureDetector(**defaults)
+
+    def test_no_drop_when_not_under_pressure(self):
+        bp = self._make_detector()
+        for _ in range(20):
+            assert bp.should_store(1) is True
+
+    def test_no_drop_at_exactly_high_water(self):
+        bp = self._make_detector()
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S
+        for _ in range(20):
+            assert bp.should_store(1) is True
+
+    def test_partial_drop_at_moderate_pressure(self):
+        bp = self._make_detector()
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 2
+        results = [bp.should_store(1) for _ in range(20)]
+        assert any(results), "some stores should be allowed"
+        assert not all(results), "some stores should be dropped"
+
+    def test_full_drop_at_severe_pressure(self):
+        bp = self._make_detector()
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 4
+        for _ in range(20):
+            assert bp.should_store(1) is False
+
+    def test_drop_rate_increases_with_severity(self):
+        """Higher EMA → higher drop rate."""
+        drop_counts = []
+        for multiplier in [1.5, 2.0, 2.5, 3.5]:
+            bp = self._make_detector()
+            bp._under_pressure = True
+            bp.store_latency_ema = _BP_HIGH_WATER_S * multiplier
+            drops = sum(1 for _ in range(100) if not bp.should_store(1))
+            drop_counts.append(drops)
+        for i in range(len(drop_counts) - 1):
+            assert drop_counts[i] <= drop_counts[i + 1]
+
+    def test_explicit_drop_store_policy_still_drops_all(self):
+        """DropStorePolicy (explicit) still drops everything under pressure."""
+        bp = EMABackpressureDetector(
+            high_water_s=_BP_HIGH_WATER_S,
+            low_water_s=_BP_LOW_WATER_S,
+            policy=DropStorePolicy(),
+            cooldown_s=0.0,
+        )
+        bp._under_pressure = True
+        bp.store_latency_ema = _BP_HIGH_WATER_S * 1.5
+        for _ in range(20):
+            assert bp.should_store(1) is False
+
+
+class TestSharedDropAccounting:
+    """Both policies share drop accounting via DropAccountingPolicy."""
+
+    @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy])
+    def test_shared_base(self, policy_cls):
+        assert issubclass(policy_cls, DropAccountingPolicy)
+
+    @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy])
+    def test_pop_returns_and_resets(self, policy_cls):
+        policy = policy_cls()
+        policy.on_store_skipped(3)
+        policy.on_store_skipped(2)
+        assert policy.pop_stores_dropped() == (2, 5)
+        assert policy.pop_stores_dropped() == (0, 0)
+
+    @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy])
+    def test_reset_clears_counters(self, policy_cls):
+        policy = policy_cls()
+        policy.on_store_skipped(4)
+        policy.reset()
+        assert policy.pop_stores_dropped() == (0, 0)
+
+    def test_throttled_reset_also_clears_call_count(self):
+        policy = ThrottledDropPolicy()
+        policy._call_count = 7
+        policy.reset()
+        assert policy._call_count == 0
+
+
+class TestHealthyBypass:
+    """Tests for the fast-path bypass when a tier is consistently healthy."""
+
+    def _make_detector(self, **kwargs):
+        defaults = dict(
+            high_water_s=_BP_HIGH_WATER_S,
+            low_water_s=_BP_LOW_WATER_S,
+            cooldown_s=0.0,
+        )
+        defaults.update(kwargs)
+        return EMABackpressureDetector(**defaults)
+
+    def test_not_healthy_initially(self):
+        bp = self._make_detector()
+        assert bp.is_healthy() is False
+
+    def test_healthy_after_streak(self):
+        bp = self._make_detector()
+        bp._completions = _BP_WARMUP
+        bp._last_update = time.monotonic()
+        block_bytes = 16
+        fast_s_per_mib = _BP_LOW_WATER_S * 0.5
+        fast_elapsed = fast_s_per_mib * (block_bytes / EMABackpressureDetector._MIB)
+        for _ in range(EMABackpressureDetector._HEALTHY_THRESHOLD + _BP_WARMUP):
+            bp.on_store_completed(fast_elapsed, block_bytes)
+        assert bp.is_healthy() is True
+
+    def test_healthy_resets_on_high_latency(self):
+        bp = self._make_detector()
+        bp._completions = _BP_WARMUP
+        bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD
+        bp._ema = _BP_LOW_WATER_S
+        bp._last_update = time.monotonic()
+        block_bytes = 16
+        slow_s_per_mib = _BP_HIGH_WATER_S * 2
+        slow_elapsed = slow_s_per_mib * (block_bytes / EMABackpressureDetector._MIB)
+        bp.on_store_completed(slow_elapsed, block_bytes)
+        assert bp.is_healthy() is False
+
+    def test_healthy_bypasses_policy(self):
+        bp = self._make_detector()
+        bp._completions = _BP_WARMUP
+        bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD
+        assert bp.should_store(1) is True
+
+    def test_reset_clears_healthy_streak(self):
+        bp = self._make_detector()
+        bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD
+        bp.reset()
+        assert bp._healthy_streak == 0
+        assert bp.is_healthy() is False
+
+
+class TestCooldown:
+    """Tests for the post-pressure cooldown period."""
+
+    def _make_detector(self, **kwargs):
+        defaults = dict(
+            high_water_s=_BP_HIGH_WATER_S,
+            low_water_s=_BP_LOW_WATER_S,
+            cooldown_s=1.0,
+        )
+        defaults.update(kwargs)
+        return EMABackpressureDetector(**defaults)
+
+    def test_cooldown_throttles_stores(self):
+        bp = self._make_detector()
+        bp._completions = _BP_WARMUP
+        bp._under_pressure = True
+        bp._ema = _BP_HIGH_WATER_S * 2
+        bp._last_update = time.monotonic()
+
+        with patch("time.monotonic") as mock_time:
+            t0 = bp._last_update + 10.0
+            mock_time.return_value = t0
+            assert bp.is_under_pressure() is False
+            assert bp._pressure_cleared_at > 0
+
+            mock_time.return_value = t0 + 0.1
+            results = [bp.should_store(1) for _ in range(10)]
+            assert sum(results) == 5, "cooldown should allow ~50%"
+
+    def test_no_cooldown_when_never_pressured(self):
+        bp = self._make_detector()
+        for _ in range(20):
+            assert bp.should_store(1) is True
+
+    def test_no_cooldown_after_window_expires(self):
+        bp = self._make_detector()
+        bp._pressure_cleared_at = time.monotonic() - 2.0
+        for _ in range(20):
+            assert bp.should_store(1) is True
+
+    def test_reset_clears_cooldown(self):
+        bp = self._make_detector()
+        bp._pressure_cleared_at = time.monotonic()
+        bp._cooldown_count = 5
+        bp.reset()
+        assert bp._pressure_cleared_at == 0.0
+        assert bp._cooldown_count == 0
diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py
index f5a576d25429..f91f54bb0d16 100644
--- a/tests/v1/kv_offload/tiering/test_factory.py
+++ b/tests/v1/kv_offload/tiering/test_factory.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for SecondaryTierFactory.
+"""Unit tests for SecondaryTierFactory.
 
 These tests verify:
 1. Pre-registration integrity — registered tier module paths can import
@@ -10,6 +9,7 @@
 3. Error paths — missing tier_type, unknown tier_type, duplicate registration.
 """
 
+from importlib.util import find_spec
 from unittest.mock import MagicMock
 
 import pytest
@@ -42,8 +42,11 @@ def _make_mock_args():
 
 
 def test_pre_registered_tiers_can_be_imported():
-    """CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager."""
+    """CI sentinel: registered tiers import and yield SecondaryTierManager."""
     for tier_type in SecondaryTierFactory._registry:
+        # KVCR is an optional external dependency and may not be installed.
+        if tier_type == "kvcr" and find_spec("kvcr") is None:
+            continue
         cls = SecondaryTierFactory._registry[tier_type]()
         assert issubclass(cls, SecondaryTierManager)
 
diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py
index 4b495e1b315c..9c3a64d22a2a 100644
--- a/tests/v1/kv_offload/tiering/test_fs_tier.py
+++ b/tests/v1/kv_offload/tiering/test_fs_tier.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for FileSystemTierManager.
+"""Unit tests for FileSystemTierManager.
 
 These tests use real disk I/O to verify the filesystem tier implementation.
 The tier manager writes KV cache blocks to disk and reads them back, verifying
diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py
new file mode 100644
index 000000000000..5cbc85230ffb
--- /dev/null
+++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py
@@ -0,0 +1,548 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Collection, Iterable, Mapping
+from types import SimpleNamespace
+from typing import Any
+
+import msgspec
+import numpy as np
+import pytest
+
+pytest.importorskip("kvcr")
+
+from kvcr import KVCRBindings
+from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig
+from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy
+from kvcr.types import (
+    BlockKey,
+    CacheTier,
+    InventoryEvent,
+    MemDescriptor,
+    OpEntryResult,
+    OpEntryStatus,
+    OpHandle,
+    QueryStatus,
+)
+
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
+    OffloadingConnectorStats,
+)
+from vllm.v1.kv_hints import KvHintsEnvelope
+from vllm.v1.kv_offload.base import (
+    LookupResult,
+    Medium,
+    OffloadKey,
+    ReqContext,
+    make_offload_key,
+)
+from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob
+from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager
+from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager
+
+
+def _op_entries(
+    entries: Mapping[BlockKey, bool],
+) -> dict[BlockKey, OpEntryResult]:
+    return {
+        key: OpEntryResult(OpEntryStatus.SUCCESS if success else OpEntryStatus.FAILED)
+        for key, success in entries.items()
+    }
+
+
+class _StubControlChannel:
+    def __init__(self, endpoint: str) -> None:
+        self.endpoint = endpoint
+
+    def send(self, endpoint: str, message: bytes) -> bool:
+        return True
+
+    def recv(self) -> list[bytes]:
+        return []
+
+    def close(self) -> None:
+        pass
+
+
+class RecordingKVCR:
+    def __init__(self) -> None:
+        self.config: KVCRConfig | None = None
+        self.guard_config: KVCRGuardConfig | None = None
+        self.backend_configs = KVCRBackendConfigs()
+        self.constructor_bindings: KVCRBindings | None = None
+        self.nixl_agent_name = "recording"
+        self.framework_control: _StubControlChannel | None = None
+        self.inventory_sink = None
+        self.query_status = QueryStatus.MISS
+        self.stats: OffloadingConnectorStats | None = None
+        self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
+        self.discard_hint_calls: list[str] = []
+        self.deliver_calls: list[
+            tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None]
+        ] = []
+        self.deposit_calls: list[
+            tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]]
+        ] = []
+        self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = []
+        self._next_op_handle = 1
+
+    def submit_hint(self, *args: Any, **kwargs: Any) -> None:
+        self.submit_hint_calls.append((args, kwargs))
+
+    def discard_hint(self, request_id: str) -> None:
+        self.discard_hint_calls.append(request_id)
+
+    def query(
+        self,
+        keys: Collection[BlockKey],
+        request_id: str | None = None,
+    ) -> list[tuple[QueryStatus, CacheTier | None]]:
+        tier = {
+            QueryStatus.HIT: CacheTier.LOCAL_G2,
+            QueryStatus.FETCHING: CacheTier.LOCAL_G2,
+            QueryStatus.FETCHABLE: CacheTier.REMOTE_G2,
+            QueryStatus.MISS: None,
+        }[self.query_status]
+        return [(self.query_status, tier) for _ in keys]
+
+    def deliver(
+        self,
+        blocks: Mapping[BlockKey, list[MemDescriptor]],
+        request_id: str | None = None,
+    ) -> OpHandle:
+        op_handle = self._next_op_handle
+        self._next_op_handle += 1
+        self.deliver_calls.append((op_handle, dict(blocks), request_id))
+        self.complete(op_handle, {key: True for key in blocks})
+        return op_handle
+
+    def deposit(
+        self,
+        blocks: Mapping[BlockKey, list[MemDescriptor]],
+    ) -> OpHandle:
+        op_handle = self._next_op_handle
+        self._next_op_handle += 1
+        self.deposit_calls.append((op_handle, dict(blocks)))
+        self.complete(op_handle, {key: True for key in blocks})
+        return op_handle
+
+    def complete(
+        self,
+        op_handle: OpHandle,
+        entries: Mapping[BlockKey, bool],
+    ) -> None:
+        self.completed.append((op_handle, _op_entries(entries)))
+
+    def poll_completed(
+        self,
+    ) -> Iterable[tuple[OpHandle, dict[BlockKey, OpEntryResult]]]:
+        completed = self.completed
+        self.completed = []
+        return completed
+
+    def get_stats(self) -> OffloadingConnectorStats | None:
+        stats = self.stats
+        self.stats = None
+        return stats
+
+    def close(self) -> None:
+        pass
+
+
+class _ExternalPolicy(FIFOPolicy):
+    pass
+
+
+def _job(
+    job_id: int,
+    req_context: ReqContext,
+    key: OffloadKey | None = None,
+    chunk_id: int = 0,
+) -> TransferJob:
+    return TransferJob(
+        job_id=job_id,
+        keys=[key if key is not None else OffloadKey(b"k0")],
+        chunk_ids=np.array([chunk_id], dtype=np.int64),
+        is_promotion=True,
+        req_context=req_context,
+    )
+
+
+def _make_tier(
+    monkeypatch,
+    kvcr: RecordingKVCR,
+    *,
+    enable_telemetry: bool = False,
+    secondary_g2_slots: int = 0,
+    kvcr_service_socket_path: str | None = None,
+    compatibility_digest: str | None = None,
+    enable_kv_cache_events: bool = False,
+    self_describing_kv_events: bool = False,
+    policy: str | None = None,
+    g3: dict[str, object] | None = None,
+    control_ports: list[int] | None = None,
+    data_parallel_rank_local: int | None = None,
+) -> KVCRSecondaryTierManager:
+    def make_control(_bind_host, bind_port, advertise_host):
+        return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}")
+
+    def make_kvcr(config, bindings, backend_configs, guard_config):
+        kvcr.config = config
+        kvcr.guard_config = guard_config
+        kvcr.nixl_agent_name = config.nixl_agent_name
+        kvcr.backend_configs = backend_configs
+        kvcr.constructor_bindings = bindings
+        kvcr.framework_control = bindings.framework_control
+        kvcr.inventory_sink = bindings.inventory_sink
+        return kvcr
+
+    monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr)
+    monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control)
+    return KVCRSecondaryTierManager(
+        offloading_spec=SimpleNamespace(
+            config=SimpleNamespace(
+                parallel=SimpleNamespace(
+                    data_parallel_rank_local=data_parallel_rank_local,
+                )
+            ),
+            kv_events_config=SimpleNamespace(
+                enable_kv_cache_events=enable_kv_cache_events,
+                self_describing_kv_events=self_describing_kv_events,
+            ),
+        ),
+        primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)),
+        tier_type="kvcr",
+        router_capabilities=["router_hint"],
+        control_host="127.0.0.1",
+        control_ports=control_ports if control_ports is not None else [7777],
+        control_advertise_host="127.0.0.1",
+        enable_telemetry=enable_telemetry,
+        secondary_g2_slots=secondary_g2_slots,
+        kvcr_service_socket_path=kvcr_service_socket_path,
+        compatibility_digest=compatibility_digest,
+        policy=policy,
+        g3=g3,
+        local_dram_backend="UCX",
+        remote_fw_dram_backend="UCX",
+    )
+
+
+def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch):
+    """Keep the control endpoint and guard pool aligned to the local DP rank."""
+    kvcr = RecordingKVCR()
+    tier = _make_tier(
+        monkeypatch,
+        kvcr,
+        kvcr_service_socket_path="/tmp/kvcr.sock",
+        compatibility_digest="Opaque-Digest",
+        secondary_g2_slots=1,
+        control_ports=[7001, 7002],
+        data_parallel_rank_local=1,
+    )
+
+    assert kvcr.framework_control is not None
+    assert kvcr.framework_control.endpoint == "tcp://127.0.0.1:7002"
+    assert kvcr.guard_config == KVCRGuardConfig(
+        kvcr_service_socket_path="/tmp/kvcr.sock",
+        guard_index=1,
+        compatibility_digest="Opaque-Digest",
+    )
+    assert kvcr.config is not None
+    assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)]
+    assert kvcr.backend_configs.local_dram is None
+
+
+def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path):
+    """Convert user-provided G3 paths to KVCR's typed configuration."""
+    kvcr = RecordingKVCR()
+    path = tmp_path / "g3.data"
+
+    _make_tier(
+        monkeypatch,
+        kvcr,
+        g3={"paths": [str(path)], "capacity_bytes_per_file": 64},
+    )
+
+    assert kvcr.backend_configs.g3 == G3Options(
+        paths=(path,),
+        capacity_bytes_per_file=64,
+    )
+
+
+def test_kvcr_tier_adapts_request_and_load(monkeypatch):
+    """Check hint forwarding, key conversion, load descriptors, and cleanup."""
+    kvcr = RecordingKVCR()
+    tier = _make_tier(monkeypatch, kvcr)
+    kv_hint = KvHintsEnvelope(protocol_version="0.1", message_id="msg", actions=[])
+    ctx = ReqContext(
+        req_id="req",
+        kv_hints=kv_hint,
+    )
+    key = make_offload_key((123).to_bytes(8, "big"), 0)
+    same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7)
+    other_key = make_offload_key((124).to_bytes(8, "big"), 0)
+
+    tier.on_new_request(ctx)
+
+    assert kvcr.submit_hint_calls == [
+        ((), {"request_id": "req", "hints": msgspec.to_builtins(kv_hint)})
+    ]
+
+    bindings = kvcr.constructor_bindings
+    assert bindings is not None
+    assert bindings.key_adapter is not None
+    monkeypatch.setenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "0")
+    decode = bindings.key_adapter.decode
+    hash_123 = (123).to_bytes(8, "big")
+    assert decode(BlockKey(bytes(key))) == hash_123
+    assert decode(BlockKey(bytes(same_hash_other_group))) == hash_123
+    assert decode(BlockKey(bytes(other_key))) == (124).to_bytes(8, "big")
+
+    tier.submit_load(_job(7, ctx, key=key, chunk_id=2))
+
+    # Check load destinations and job completion independently of hint parsing.
+    assert len(kvcr.submit_hint_calls) == 1
+    _, blocks, request_id = kvcr.deliver_calls[0]
+    assert request_id == "req"
+    assert list(blocks) == [key]
+    (descriptor,) = blocks[key]
+    assert descriptor.end_point_name == kvcr.nixl_agent_name
+    assert descriptor.addr == tier._primary_base_addr + 2 * 16
+    assert descriptor.size == 16
+    assert list(tier.get_finished_jobs()) == [JobResult(7, True)]
+
+    # Here we verify that request cleanup discards the request-scoped hint in
+    # KVCR.
+    tier.on_request_finished(ctx)
+    assert kvcr.discard_hint_calls == ["req"]
+
+
+def test_kvcr_tier_allows_request_without_router_hint(monkeypatch):
+    """Keep router hints optional for requests from non-hint-aware routers."""
+    kvcr = RecordingKVCR()
+    tier = _make_tier(monkeypatch, kvcr)
+
+    tier.on_new_request(ReqContext(req_id="req"))
+
+    assert kvcr.submit_hint_calls == []
+
+
+@pytest.mark.parametrize(
+    ("status", "expected"),
+    [
+        (QueryStatus.MISS, LookupResult.MISS),
+        (QueryStatus.HIT, LookupResult.HIT),
+        (QueryStatus.FETCHABLE, LookupResult.HIT),
+        (QueryStatus.FETCHING, LookupResult.RETRY),
+    ],
+)
+def test_kvcr_tier_maps_query_status(monkeypatch, status, expected):
+    """Map KVCR cache states to the scheduler's lookup contract."""
+    kvcr = RecordingKVCR()
+    kvcr.query_status = status
+    tier = _make_tier(monkeypatch, kvcr)
+
+    assert tier.lookup(OffloadKey(b"k0"), ReqContext(req_id="req")) is expected
+
+
+def test_kvcr_tier_serves_primary_pin_request(monkeypatch):
+    """Hold primary-tier hits until KVCR releases the corresponding pin."""
+    kvcr = RecordingKVCR()
+    tier = _make_tier(monkeypatch, kvcr)
+    keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2"))
+    hit_keys = (keys[0], keys[2])
+    chunk_ids = {keys[0]: 1, keys[2]: 5}
+    lifecycle: list[str] = []
+
+    class Parent:
+        def on_new_request(self, req_context):
+            lifecycle.append("new")
+            return SimpleNamespace()
+
+        def lookup(self, key, req_context):
+            return LookupResult.HIT if key in hit_keys else LookupResult.MISS
+
+        def create_store_job(self, requested_keys, req_context):
+            return TransferJob(
+                job_id=11,
+                keys=requested_keys,
+                chunk_ids=np.array([chunk_ids[key] for key in requested_keys]),
+                is_promotion=True,
+                req_context=req_context,
+            )
+
+        def on_request_finished(self, req_context):
+            lifecycle.append("finished")
+
+    bindings = kvcr.constructor_bindings
+    assert bindings is not None
+    request = bindings.request_pin(keys)
+
+    tier.serve_external_requests(Parent())
+
+    [(queued_request, result)] = bindings.poll_pin_results()
+    assert queued_request == request
+    assert result is not None
+    pin_handle, descriptors = result
+    assert descriptors[keys[1]] is None
+    block_descriptors = descriptors[keys[2]]
+    assert block_descriptors is not None
+    (descriptor,) = block_descriptors
+    assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride)
+    assert lifecycle == ["new", "finished"]
+
+    polls = 0
+
+    def release_pin():
+        nonlocal polls
+        polls += 1
+        bindings.release_pin(pin_handle)
+        return []
+
+    monkeypatch.setattr(kvcr, "poll_completed", release_pin)
+    tier.drain_jobs()
+
+    assert polls == 1
+    assert list(tier.get_finished_jobs()) == [JobResult(11, True)]
+
+
+def test_kvcr_telemetry_is_opt_in_and_namespaced_at_vllm_boundary(monkeypatch):
+    """Keep telemetry opt-in and namespace metrics only at the vLLM boundary."""
+    assert KVCRSecondaryTierManager.build_metric_definitions({}) == {}
+    definitions = KVCRSecondaryTierManager.build_metric_definitions(
+        {"enable_telemetry": True}
+    )
+    assert "vllm:kvcr_duration_seconds" in definitions
+    assert "vllm:kvcr_transfer_blocks" in definitions
+
+    kvcr = RecordingKVCR()
+    tier = _make_tier(monkeypatch, kvcr, enable_telemetry=True)
+    bindings = kvcr.constructor_bindings
+    assert bindings is not None
+    assert bindings.stats_factory is not None
+    stats = bindings.stats_factory()
+    stats.increase_counter(
+        "kvcr_transfer_blocks",
+        2,
+        ("remote_deliver",),
+    )
+    kvcr.stats = stats
+
+    returned = tier.get_stats()
+
+    assert returned is stats
+    assert returned.reduce() == {"vllm:kvcr_transfer_blocks:('remote_deliver',)": 2}
+
+
+@pytest.mark.parametrize(
+    ("policy", "expected_type"),
+    [
+        ("fifo", FIFOPolicy),
+        ("lru", LRUPolicy),
+        ("g3_fifo", G3FIFOPolicy),
+        ("g3_lru", G3LRUPolicy),
+        (f"{__name__}._ExternalPolicy", _ExternalPolicy),
+    ],
+)
+def test_kvcr_tier_passes_selected_policy(monkeypatch, policy, expected_type):
+    """Resolve every built-in and fully qualified external policy."""
+    kvcr = RecordingKVCR()
+    _make_tier(monkeypatch, kvcr, policy=policy)
+
+    bindings = kvcr.constructor_bindings
+    assert bindings is not None
+    assert type(bindings.policy) is expected_type
+
+
+@pytest.mark.parametrize(
+    ("socket_path", "digest"),
+    [("/tmp/kvcr.sock", None), (None, "Opaque-Digest")],
+)
+def test_kvcr_tier_requires_complete_service_config(monkeypatch, socket_path, digest):
+    """Reject partial guard configuration before connecting to the service."""
+    with pytest.raises(ValueError, match="configured together"):
+        _make_tier(
+            monkeypatch,
+            RecordingKVCR(),
+            kvcr_service_socket_path=socket_path,
+            compatibility_digest=digest,
+        )
+
+
+def test_kvcr_tier_stores_and_emits_inventory(monkeypatch):
+    """Cover store descriptors and inventory translation at the KVCR boundary."""
+    kvcr = RecordingKVCR()
+    tier = _make_tier(
+        monkeypatch,
+        kvcr,
+        secondary_g2_slots=2,
+        enable_kv_cache_events=True,
+        self_describing_kv_events=True,
+    )
+    local_dram = kvcr.backend_configs.local_dram
+    assert local_dram is not None
+    assert [(name, size) for name, _, size in local_dram.pools] == [
+        ("", 2 * tier._primary_row_stride)
+    ]
+
+    key = OffloadKey(b"k0")
+    tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2))
+
+    _, blocks = kvcr.deposit_calls[0]
+    assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16
+    assert list(tier.get_finished_jobs()) == [JobResult(11, True)]
+
+    assert kvcr.inventory_sink is not None
+    kvcr.inventory_sink(
+        InventoryEvent((BlockKey(bytes(key)),), CacheTier.LOCAL_G2, False)
+    )
+    kvcr.inventory_sink(InventoryEvent((BlockKey(bytes(key)),), CacheTier.G3, False))
+    events = list(tier.take_events())
+    assert [
+        (event.keys, event.medium, event.ownership, event.removed) for event in events
+    ] == [
+        ([key], Medium.CPU, "kvcr", False),
+        ([key], Medium.STORAGE, "kvcr", False),
+    ]
+    assert all(event.removal_expected for event in events)
+    tier.shutdown()
+
+
+def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch):
+    """Require tier-aware events whenever KVCR owns local cache inventory."""
+    with pytest.raises(ValueError, match="self_describing_kv_events"):
+        _make_tier(
+            monkeypatch,
+            RecordingKVCR(),
+            secondary_g2_slots=1,
+            enable_kv_cache_events=True,
+        )
+
+
+def test_kvcr_tier_accumulates_block_results(monkeypatch):
+    """Wait for every block result and preserve partial success while draining."""
+    kvcr = RecordingKVCR()
+    tier = _make_tier(monkeypatch, kvcr)
+    keys = [OffloadKey(b"k0"), OffloadKey(b"k1")]
+    tier.submit_load(
+        TransferJob(
+            job_id=13,
+            keys=keys,
+            chunk_ids=np.array([0, 1], dtype=np.int64),
+            is_promotion=True,
+            req_context=ReqContext(req_id="req"),
+        )
+    )
+
+    op_handle, _, _ = kvcr.deliver_calls[0]
+    kvcr.completed.clear()
+    kvcr.complete(op_handle, {keys[0]: True})
+    assert list(tier.get_finished_jobs()) == []
+
+    kvcr.complete(op_handle, {keys[1]: False})
+    tier.drain_jobs()
+
+    [result] = tier.get_finished_jobs()
+    assert result.job_id == 13
+    assert not result.success
+    assert result.successful_keys == {keys[0]}
diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py
index 13840bac3599..bf8afca42428 100644
--- a/tests/v1/kv_offload/tiering/test_obj_tier.py
+++ b/tests/v1/kv_offload/tiering/test_obj_tier.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Mock-based unit tests for ObjectStoreSecondaryTierManager.
+"""Mock-based unit tests for ObjectStoreSecondaryTierManager.
 
 These tests replace the NIXL backend with an in-memory mock so they run
 without S3 credentials or a live object store. They verify the manager's
diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py
index f737685621b4..63b4c2c4d0b5 100644
--- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py
+++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager.
+"""Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager.
 
 These tests verify:
 1. Basic tiered offloading operations (store, load, lookup)
@@ -40,6 +39,7 @@
     TierMatcher,
     make_offload_key,
 )
+from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
 from vllm.v1.kv_offload.tiering.base import (
     JobResult,
     SecondaryTierManager,
@@ -753,14 +753,64 @@ def test_touch_propagates_to_all_tiers(self, manager_setup):
         # Touch chunks
         self.manager.touch(chunks, _CTX)
 
-        # Verify touch was called on primary tier (check LRU order)
-        primary_keys = list(self.primary_tier._policy.evictable_chunks.keys())
-        assert primary_keys[-3:] == list(reversed(chunks))
+        # Verify touch was called on the primary tier (head is most recent).
+        policy = self.primary_tier._policy
+        assert isinstance(policy, LRUCachePolicy)
+        ranks = policy._ranks
+        assert ranks[chunks[2]] < ranks[chunks[1]] < ranks[chunks[0]]
 
         # Verify touch was propagated to all secondary tiers
         self.secondary_tier1.touch.assert_called_once_with(chunks, _CTX)
         self.secondary_tier2.touch.assert_called_once_with(chunks, _CTX)
 
+    def test_request_order_survives_late_cascade_completion(self, manager_setup):
+        """Cascade completion must not replace request order with I/O order."""
+        chunks = to_keys(range(5))
+        self._start_request()
+        output = self.manager.prepare_store(chunks, _CTX)
+        assert output is not None
+        self.manager.complete_store(chunks, _CTX, success=True)
+
+        # Finalize while both secondary-tier cascades still pin primary chunks.
+        self.manager.on_request_finished(_CTX)
+        self._simulate_on_schedule_end()
+        self._simulate_on_schedule_end()
+
+        evict_ctx = ReqContext(req_id="evict")
+        self._start_request(evict_ctx)
+        output = self.manager.prepare_store(to_keys([5]), evict_ctx)
+        assert output is not None
+        assert output.evicted_keys == [chunks[-1]]
+
+    def test_late_old_store_submission_does_not_override_newer_request_order(
+        self, manager_setup
+    ):
+        """Cascade submission after finish is not a new primary access."""
+        old_blocks = to_keys(range(2))
+        old_ctx = ReqContext(req_id="old")
+        self._start_request(old_ctx)
+        assert self.manager.prepare_store(old_blocks, old_ctx) is not None
+        self.manager.on_request_finished(old_ctx)
+
+        new_blocks = to_keys(range(2, 5))
+        new_ctx = ReqContext(req_id="new")
+        self._start_request(new_ctx)
+        assert self.manager.prepare_store(new_blocks, new_ctx) is not None
+        self.manager.complete_store(new_blocks, new_ctx, success=True)
+        self.manager.on_request_finished(new_ctx)
+
+        # The old GPU->primary write lands after the newer request finished.
+        # Its cascade pins must not turn it into the more recent request.
+        self.manager.complete_store(old_blocks, old_ctx, success=True)
+        self._simulate_on_schedule_end()
+        self._simulate_on_schedule_end()
+
+        evict_ctx = ReqContext(req_id="evict")
+        self._start_request(evict_ctx)
+        output = self.manager.prepare_store(to_keys([5]), evict_ctx)
+        assert output is not None
+        assert output.evicted_keys == [old_blocks[-1]]
+
     def test_failed_store_no_cascade(self, manager_setup):
         """Test that failed GPU→primary store doesn't cascade."""
         chunks = to_keys(range(3))
@@ -873,10 +923,10 @@ def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup
         job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
         assert job_metadata.req_context is ctx
 
-    def test_on_request_finished_delays_secondary_until_store_submitted(
+    def test_on_request_finished_delays_only_secondary_until_store_submitted(
         self, manager_setup
     ):
-        """Manager hook is eager; secondary hooks wait for cascade submission."""
+        """Primary order commits immediately; secondary cleanup waits."""
         chunks = to_keys(range(2))
         ctx = ReqContext(req_id="req_delayed_secondary")
         calls: list[tuple[str, str]] = []
diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py
index 7e60eeb2e3b8..9b12e7412ed6 100644
--- a/tests/v1/logits_processors/test_correctness.py
+++ b/tests/v1/logits_processors/test_correctness.py
@@ -98,7 +98,7 @@ def __init__(self, workload_index: int, logitproc_type: LogitprocType):
         self.params = _sampling_params_from_logitproc(logitproc_type)
 
     def __str__(self):
-        """For debugging"""
+        """For debugging."""
         summ = ", ".join(f"{k}={v}" for k, v in vars(self).items())
         return f"MyClass({summ})"
 
@@ -117,7 +117,7 @@ def _generate_fake_sampling_metadata(
     vocab_size: int,
     device: torch.device,
 ) -> SamplingMetadata:
-    """Generate fake sampling metadata with fake logitsprocs"""
+    """Generate fake sampling metadata with fake logitsprocs."""
     output_token_ids: list[list[int]] = []
     prompt_token_ids: list[list[int]] = []
     for _ in range(batch_size):
@@ -176,7 +176,7 @@ def _generate_fake_sampling_metadata(
 
 
 def _generate_test_fakes(batch_size: int, device: str) -> LogitsprocsTestFakes:
-    """Generate fake logits and sampling metadata"""
+    """Generate fake logits and sampling metadata."""
     fake_logits = create_fake_logits(batch_size, VOCAB_SIZE)
     # Create one dominant token per batch, to support min-p test
     for i in range(batch_size):
@@ -192,7 +192,7 @@ def _generate_test_fakes(batch_size: int, device: str) -> LogitsprocsTestFakes:
 
 
 def _sampling_params_from_logitproc(logitproc_type: LogitprocType) -> SamplingParams:
-    """Customize request SamplingParams for a specified logitproc"""
+    """Customize request SamplingParams for a specified logitproc."""
     # SamplingParams for req with no logitproc
     kwargs = {"min_p": 0.0, "logit_bias": None, "min_tokens": 0}
     if fxn := logitsprocs_test_mapping[logitproc_type].gen_request_fxn:
@@ -220,6 +220,7 @@ def _generate_mixed_logitsprocs_batch_params(
     Returns:
       List of per-request params which configure the engine for that request's
       enabled logitproc
+
     """
     batch_size = len(logitsprocs_types) * reqs_per_logitproc
     # Generate multiple repeats of key params for each logitproc;
@@ -251,7 +252,7 @@ def _raise_error_invalid(
 
 
 def _logit_bias_params(kwargs: dict) -> None:
-    """Logit bias config"""
+    """Logit bias config."""
     kwargs["logit_bias"] = {
         random.randint(0, VOCAB_SIZE - 1): random.choice([-0.1, 0.2])
     }
@@ -265,7 +266,7 @@ def _logit_bias_validate(
     request_params: LogitsProcsRequestParams,
     step_idx: int,
 ) -> None:
-    """Validate logit bias logitproc applied correctly"""
+    """Validate logit bias logitproc applied correctly."""
     logit_bias = request_params.params.logit_bias
     logits_old = test_fakes.logits[persistent_batch[batch_index].workload_index].cpu()
     logits_new = logits_new[batch_index].cpu()
@@ -301,7 +302,7 @@ def _logit_bias_validate(
 
 
 def _min_p_params(kwargs: dict) -> None:
-    """Min-p logitproc config"""
+    """Min-p logitproc config."""
     kwargs["min_p"] = 0.1
 
 
@@ -313,7 +314,7 @@ def _min_p_validate(
     request_params: LogitsProcsRequestParams,
     step_idx: int,
 ) -> None:
-    """Validate min-p logitproc applied correctly"""
+    """Validate min-p logitproc applied correctly."""
     for token_id in range(VOCAB_SIZE):
         logits_for_token = logits_new[batch_index][token_id]
         if token_id == 0:
@@ -347,7 +348,7 @@ def _min_p_validate(
 
 
 def _min_tokens_params(kwargs: dict) -> None:
-    """Min-tokens logitproc config"""
+    """Min-tokens logitproc config."""
     kwargs["min_tokens"] = MIN_TOKENS_LEN_THRESHOLD
     kwargs["stop_token_ids"] = [
         np.random.randint(0, VOCAB_SIZE - 1)
@@ -363,7 +364,7 @@ def _min_tokens_validate(
     request_params: LogitsProcsRequestParams,
     step_idx: int,
 ) -> None:
-    """Validate min-tokens logitsproc applied correctly"""
+    """Validate min-tokens logitsproc applied correctly."""
     ref_num_out_tokens = len(request_params.out_tokens)
     min_reached = ref_num_out_tokens >= MIN_TOKENS_LEN_THRESHOLD
     ref_all_stop_token_ids = request_params.params.all_stop_token_ids
@@ -560,7 +561,7 @@ def test_min_tokens_restores_all_masked_structured_output_stop_token_spec_decode
 
 
 def _thinking_budget_params(kwargs: dict) -> None:
-    """Set SamplingParams kwargs for thinking token budget tests"""
+    """Set SamplingParams kwargs for thinking token budget tests."""
     kwargs["thinking_token_budget"] = THINKING_TOKEN_BUDGET
 
 
@@ -661,7 +662,7 @@ def _none_validate(
     request_params: LogitsProcsRequestParams,
     step_idx: int,
 ) -> None:
-    """Validate that no logits processors are applied"""
+    """Validate that no logits processors are applied."""
     logits = test_fakes.logits[persistent_batch[batch_index].workload_index].cpu()
     ref_logits = logits_new[batch_index]
     if not torch.all(ref_logits == logits):
@@ -706,7 +707,7 @@ class LogitsprocTestHelpers(NamedTuple):
 
 
 def _get_test_cases() -> list[list[str]]:
-    """Each test case is a set of logitsprocs"""
+    """Each test case is a set of logitsprocs."""
     logitsprocs_types = list(logitsprocs_test_mapping.keys())
 
     # Isolate thinking-budget handling from other processors to avoid cross-talk.
diff --git a/tests/v1/logits_processors/test_custom_offline.py b/tests/v1/logits_processors/test_custom_offline.py
index 302a72057eaa..f04f5cfa47c8 100644
--- a/tests/v1/logits_processors/test_custom_offline.py
+++ b/tests/v1/logits_processors/test_custom_offline.py
@@ -86,8 +86,8 @@ def _run_test(kwargs: dict, logitproc_loaded: bool) -> None:
     Args:
       kwargs: `LLM` constructor kwargs
       logitproc_loaded: server has loaded dummy logitproc if True
-    """
 
+    """
     # Create a vLLM instance and load custom logitproc
     llm_logitproc = LLM(
         model=MODEL_NAME,
@@ -131,7 +131,7 @@ def _run_test(kwargs: dict, logitproc_loaded: bool) -> None:
 @create_new_process_for_each_test()
 @pytest.mark.parametrize("logitproc_source", list(CustomLogitprocSource))
 def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource):
-    """Test offline Python interface for passing custom logitsprocs
+    """Test offline Python interface for passing custom logitsprocs.
 
     Construct an `LLM` instance which loads a custom logitproc that has a
     well-defined behavior (mask out all tokens except one `target_token`)
@@ -158,10 +158,13 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource
       logitproc_source: what source (entrypoint, fully-qualified class name
                         (FQCN), class object, or None) the user pulls the
                         logitproc from
-    """
 
+    """
     # Test that logitproc info is passed to workers
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
+    # These tests exercise the V1-interface logits processor; Model Runner V2
+    # rejects V1-interface processors at load time by design.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
     set_random_seed(40)
 
     # Choose LLM args based on logitproc source
@@ -192,7 +195,7 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource
 
 @create_new_process_for_each_test()
 def test_custom_logitsprocs_req(monkeypatch):
-    """Test passing request-level logits processor to offline Python interface
+    """Test passing request-level logits processor to offline Python interface.
 
     Wrap a request-level logits processor to create a batch level logits
     processor that has a well-defined behavior (mask out all tokens except one
@@ -213,10 +216,13 @@ def test_custom_logitsprocs_req(monkeypatch):
 
     Args:
       monkeypatch: for setting env vars
-    """
 
+    """
     # Test that logitproc info is passed to workers
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
+    # These tests exercise the V1-interface logits processor; Model Runner V2
+    # rejects V1-interface processors at load time by design.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
     set_random_seed(40)
     _run_test(
         {"logits_processors": [WrappedPerReqLogitsProcessor]}, logitproc_loaded=True
@@ -258,8 +264,12 @@ def test_rejects_custom_logitsprocs(
       logitproc_source: what source (entrypoint, fully-qualified class name
                         (FQCN), or class object) the user pulls the
                         logitproc from
+
     """
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    # These tests exercise V1-interface rejection and V1 runner internals;
+    # Model Runner V2 rejects V1-interface processors at load time by design.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
     set_random_seed(40)
 
     test_params: dict[str, dict[str, Any]] = {
diff --git a/tests/v1/logits_processors/test_custom_online.py b/tests/v1/logits_processors/test_custom_online.py
index 3b7add3b80f4..b117875625ab 100644
--- a/tests/v1/logits_processors/test_custom_online.py
+++ b/tests/v1/logits_processors/test_custom_online.py
@@ -77,9 +77,11 @@ def server(default_server_args, request, monkeypatch):
     (2) No --logits-processors cli arg; inject a dummy logits processor
     entrypoint
     """
-
     # Test that logitproc info is passed to workers
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
+    # These tests exercise the V1-interface logits processor; Model Runner V2
+    # rejects V1-interface processors at load time by design.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
 
     if request.param:
         # Launch server, append FQCN argument, inject dummy logitproc module
@@ -118,7 +120,7 @@ async def client(server):
     [MODEL_NAME],
 )
 def test_custom_logitsprocs(server, model_name: str):
-    """Test custom logitsprocs when starting OpenAI server from CLI
+    """Test custom logitsprocs when starting OpenAI server from CLI.
 
     Launch vLLM OpenAI-compatible server, configured to load a custom logitproc
     that has a well-defined behavior (mask out all tokens except one
@@ -130,7 +132,6 @@ def test_custom_logitsprocs(server, model_name: str):
     Validate that requests which activate the custom logitproc, repeat the same
     token
     """
-
     import asyncio
 
     async def _async_main(srv, mn):
@@ -180,8 +181,7 @@ async def _run(client):
 async def test_invalid_custom_logitsproc_arg(
     client: openai.AsyncOpenAI, model_name: str
 ):
-    """Test that request with invalid custom logitsproc is rejected"""
-
+    """Test that request with invalid custom logitsproc is rejected."""
     prompt = "Hello, my name is"
     # Pass invalid (non-int) target_token value to dummy logits processor
     request_keyword_args: dict[str, Any] = {
diff --git a/tests/v1/logits_processors/utils.py b/tests/v1/logits_processors/utils.py
index fb45b73c9893..76877b655b1a 100644
--- a/tests/v1/logits_processors/utils.py
+++ b/tests/v1/logits_processors/utils.py
@@ -37,7 +37,7 @@
 
 
 class CustomLogitprocSource(Enum):
-    """How to source a logitproc for testing purposes"""
+    """How to source a logitproc for testing purposes."""
 
     LOGITPROC_SOURCE_NONE = auto()  # No custom logitproc
     LOGITPROC_SOURCE_ENTRYPOINT = auto()  # Via entrypoint
@@ -55,7 +55,7 @@ class CustomLogitprocSource(Enum):
 
 
 class DummyLogitsProcessor(LogitsProcessor):
-    """Fake logit processor to support unit testing and examples"""
+    """Fake logit processor to support unit testing and examples."""
 
     @classmethod
     def validate_params(cls, params: SamplingParams):
@@ -73,7 +73,7 @@ def __init__(
         self.req_info: dict[int, int] = {}
 
     def is_argmax_invariant(self) -> bool:
-        """Never impacts greedy sampling"""
+        """Never impacts greedy sampling."""
         return False
 
     def update_state(self, batch_update: BatchUpdate | None):
@@ -115,7 +115,7 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor:
 
 
 class EntryPoint:
-    """Dummy entrypoint class for logitsprocs testing"""
+    """Dummy entrypoint class for logitsprocs testing."""
 
     def __init__(self):
         self.name = DUMMY_LOGITPROC_ENTRYPOINT
@@ -126,7 +126,7 @@ def load(self):
 
 
 class EntryPoints(list):
-    """Dummy EntryPoints class for logitsprocs testing"""
+    """Dummy EntryPoints class for logitsprocs testing."""
 
     def __init__(self, group: str):
         # Emulate list-like functionality
@@ -141,7 +141,7 @@ class DummyPerReqLogitsProcessor:
     token id identified by `target_token`"""
 
     def __init__(self, target_token: int) -> None:
-        """Specify `target_token`"""
+        """Specify `target_token`."""
         self.target_token = target_token
 
     def __call__(
@@ -178,6 +178,7 @@ def new_req_logits_processor(
 
         Returns:
           `Callable` request logits processor, or None
+
         """
         target_token: Any | None = params.extra_args and params.extra_args.get(
             "target_token"
diff --git a/tests/v1/metrics/test_engine_logger_apis.py b/tests/v1/metrics/test_engine_logger_apis.py
index 2e243c23cbf9..5741df0d56ed 100644
--- a/tests/v1/metrics/test_engine_logger_apis.py
+++ b/tests/v1/metrics/test_engine_logger_apis.py
@@ -13,8 +13,7 @@
 
 @pytest.fixture
 def log_stats_enabled_engine_args():
-    """
-    Shared fixture providing common AsyncEngineArgs configuration
+    """Shared fixture providing common AsyncEngineArgs configuration
     used across multiple tests.
     """
     return AsyncEngineArgs(
@@ -27,10 +26,7 @@ def log_stats_enabled_engine_args():
 
 @pytest.mark.asyncio
 async def test_async_llm_replace_default_loggers(log_stats_enabled_engine_args):
-    """
-    RayPrometheusStatLogger should replace the default PrometheusStatLogger
-    """
-
+    """RayPrometheusStatLogger should replace the default PrometheusStatLogger."""
     engine = AsyncLLM.from_engine_args(
         log_stats_enabled_engine_args, stat_loggers=[RayPrometheusStatLogger]
     )
@@ -40,8 +36,7 @@ async def test_async_llm_replace_default_loggers(log_stats_enabled_engine_args):
 
 @pytest.mark.asyncio
 async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args):
-    """
-    It's still possible to use custom stat loggers exclusively by passing
+    """It's still possible to use custom stat loggers exclusively by passing
     disable_log_stats=True in addition to a list of custom stat loggers.
     """
     # Create engine_args with disable_log_stats=True for this test
diff --git a/tests/v1/metrics/test_histogram_buckets.py b/tests/v1/metrics/test_histogram_buckets.py
new file mode 100644
index 000000000000..54cc31da94fa
--- /dev/null
+++ b/tests/v1/metrics/test_histogram_buckets.py
@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Snapshot tests locking the default Prometheus histogram buckets.
+
+The expected lists are deliberately hard-coded literals rather than imports:
+any change to a default bucket boundary must show up as a failing test here.
+"""
+
+import prometheus_client
+import pytest
+
+from vllm.config import ModelConfig, VllmConfig
+from vllm.config.observability import ObservabilityConfig
+from vllm.v1.metrics.buckets import (
+    BUCKET_FAMILY_KEYS,
+    build_1_2_5_buckets,
+    histogram_buckets,
+)
+from vllm.v1.metrics.loggers import PrometheusStatLogger
+from vllm.v1.metrics.prometheus import unregister_vllm_metrics
+
+pytestmark = pytest.mark.cpu_test
+
+DEFAULT_BUCKET_SNAPSHOTS: dict[str, list[float]] = {
+    "request_latency": [
+        0.3,
+        0.5,
+        0.8,
+        1.0,
+        1.5,
+        2.0,
+        2.5,
+        5.0,
+        10.0,
+        15.0,
+        20.0,
+        30.0,
+        40.0,
+        50.0,
+        60.0,
+        120.0,
+        240.0,
+        480.0,
+        960.0,
+        1920.0,
+        7680.0,
+    ],
+    "time_to_first_token": [
+        0.001,
+        0.005,
+        0.01,
+        0.02,
+        0.04,
+        0.06,
+        0.08,
+        0.1,
+        0.25,
+        0.5,
+        0.75,
+        1.0,
+        2.5,
+        5.0,
+        7.5,
+        10.0,
+        20.0,
+        40.0,
+        80.0,
+        160.0,
+        640.0,
+        2560.0,
+    ],
+    "inter_token_latency": [
+        0.01,
+        0.025,
+        0.05,
+        0.075,
+        0.1,
+        0.15,
+        0.2,
+        0.3,
+        0.4,
+        0.5,
+        0.75,
+        1.0,
+        2.5,
+        5.0,
+        7.5,
+        10.0,
+        20.0,
+        40.0,
+        80.0,
+    ],
+    "iteration_tokens": [
+        1,
+        8,
+        16,
+        32,
+        64,
+        128,
+        256,
+        512,
+        1024,
+        2048,
+        4096,
+        8192,
+        16384,
+    ],
+    "request_params_n": [1, 2, 5, 10, 20],
+    "request_num_preemptions": [1, 2, 3, 4, 5, 10, 20],
+    "kv_cache_residency": [
+        0.001,
+        0.002,
+        0.005,
+        0.01,
+        0.02,
+        0.05,
+        0.1,
+        0.2,
+        0.5,
+        1,
+        2,
+        5,
+        10,
+        20,
+        30,
+        60,
+        120,
+        300,
+        600,
+        1200,
+        1800,
+    ],
+}
+
+METRIC_FAMILIES: dict[str, str] = {
+    "vllm:e2e_request_latency_seconds": "request_latency",
+    "vllm:request_queue_time_seconds": "request_latency",
+    "vllm:request_inference_time_seconds": "request_latency",
+    "vllm:request_prefill_time_seconds": "request_latency",
+    "vllm:request_decode_time_seconds": "request_latency",
+    "vllm:time_to_first_token_seconds": "time_to_first_token",
+    "vllm:inter_token_latency_seconds": "inter_token_latency",
+    "vllm:request_time_per_output_token_seconds": "inter_token_latency",
+    "vllm:iteration_tokens_total": "iteration_tokens",
+    "vllm:request_params_n": "request_params_n",
+    "vllm:request_num_preemptions": "request_num_preemptions",
+    "vllm:request_prompt_tokens": "request_tokens",
+    "vllm:request_generation_tokens": "request_tokens",
+    "vllm:request_max_num_generation_tokens": "request_tokens",
+    "vllm:request_params_max_tokens": "request_tokens",
+    "vllm:request_prefill_kv_computed_tokens": "request_tokens",
+    "vllm:kv_block_lifetime_seconds": "kv_cache_residency",
+    "vllm:kv_block_idle_before_evict_seconds": "kv_cache_residency",
+    "vllm:kv_block_reuse_gap_seconds": "kv_cache_residency",
+}
+
+TEST_MODEL = "distilbert/distilgpt2"
+TEST_MODEL_MAX_LEN = 1024
+REQUEST_TOKENS_SNAPSHOT_FOR_TEST_MODEL: list[float] = [
+    1,
+    2,
+    5,
+    10,
+    20,
+    50,
+    100,
+    200,
+    500,
+    1000,
+]
+
+
+@pytest.mark.parametrize("family", sorted(DEFAULT_BUCKET_SNAPSHOTS))
+def test_default_bucket_snapshots(family):
+    """Static families must expose exactly today's default boundaries."""
+    assert histogram_buckets(family) == DEFAULT_BUCKET_SNAPSHOTS[family]
+
+
+def test_request_tokens_bucket_snapshot():
+    """The 1-2-5 token-count series must be stable for a given model len."""
+    assert histogram_buckets("request_tokens", max_model_len=32768) == [
+        1,
+        2,
+        5,
+        10,
+        20,
+        50,
+        100,
+        200,
+        500,
+        1000,
+        2000,
+        5000,
+        10000,
+        20000,
+    ]
+    assert build_1_2_5_buckets(100) == [1, 2, 5, 10, 20, 50, 100]
+
+
+def test_request_tokens_requires_max_model_len():
+    with pytest.raises(ValueError, match="max_model_len is required"):
+        histogram_buckets("request_tokens")
+
+
+def test_buckets_return_fresh_copies():
+    """Mutating a returned list must not corrupt the shared defaults."""
+    first = histogram_buckets("request_latency")
+    first.append(1e9)
+    second = histogram_buckets("request_latency")
+    assert second == DEFAULT_BUCKET_SNAPSHOTS["request_latency"]
+    assert second is not first
+
+
+def test_bucket_family_keys():
+    """The declared families are exactly these, and each one feeds a histogram."""
+    assert {
+        "request_latency",
+        "time_to_first_token",
+        "inter_token_latency",
+        "iteration_tokens",
+        "request_params_n",
+        "request_num_preemptions",
+        "request_tokens",
+        "kv_cache_residency",
+    } == BUCKET_FAMILY_KEYS
+    assert set(METRIC_FAMILIES.values()) == BUCKET_FAMILY_KEYS
+
+
+def collect_histogram_buckets() -> dict[str, list[float]]:
+    """Read back each vllm:* histogram's finite bucket bounds by metric name."""
+    found: dict[str, list[float]] = {}
+    for metric in prometheus_client.REGISTRY.collect():
+        if metric.name.startswith("vllm:") and metric.type == "histogram":
+            found[metric.name] = sorted(
+                {
+                    float(sample.labels["le"])
+                    for sample in metric.samples
+                    if sample.name.endswith("_bucket") and sample.labels["le"] != "+Inf"
+                }
+            )
+    return found
+
+
+def build_logger_config(
+    observability_config: ObservabilityConfig,
+) -> VllmConfig:
+    return VllmConfig(
+        model_config=ModelConfig(model=TEST_MODEL),
+        observability_config=observability_config,
+    )
+
+
+def test_prometheus_logger_default_buckets():
+    """Every core engine histogram must use its family's defaults.
+
+    The config deliberately sets no kv_transfer_config: a KV connector would
+    register its own histograms, which own their boundaries and are not
+    covered by METRIC_FAMILIES.
+    """
+    config = build_logger_config(ObservabilityConfig(kv_cache_metrics=True))
+    assert config.model_config.max_model_len == TEST_MODEL_MAX_LEN
+    try:
+        PrometheusStatLogger(config)
+        found = collect_histogram_buckets()
+        assert set(found) == set(METRIC_FAMILIES)
+        for metric_name, family in METRIC_FAMILIES.items():
+            if family == "request_tokens":
+                expected = REQUEST_TOKENS_SNAPSHOT_FOR_TEST_MODEL
+            else:
+                expected = DEFAULT_BUCKET_SNAPSHOTS[family]
+            assert found[metric_name] == [float(b) for b in expected], metric_name
+    finally:
+        unregister_vllm_metrics()
diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py
index ab30f1bb9e26..c0997ec459bf 100644
--- a/tests/v1/metrics/test_perf_metrics.py
+++ b/tests/v1/metrics/test_perf_metrics.py
@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tests for the analytic estimators in metrics/flops.py.
-"""
+"""Tests for the analytic estimators in metrics/flops.py."""
 
 import types
 from types import SimpleNamespace
 
 import pytest
+from transformers import Gemma2Config, MistralConfig
 from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
 from transformers.models.llama4.configuration_llama4 import (
     Llama4Config,
@@ -31,6 +30,7 @@
     MLAAttentionMetrics,
     ModelMetrics,
     ParsedArgs,
+    SlidingWindowAttentionParser,
     UnembedMetrics,
 )
 
@@ -38,7 +38,8 @@
 class MockModelConfig:
     """Mock ModelConfig that implements the getter methods used by parsers."""
 
-    def __init__(self, hf_config, dtype):
+    def __init__(self, hf_config, dtype, disable_sliding_window=False):
+        """Initialize MockModelConfig with HF config and options."""
         self.hf_config = hf_config
         self.hf_text_config = get_hf_text_config(hf_config)
         convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get(
@@ -49,8 +50,10 @@ def __init__(self, hf_config, dtype):
         ).convert()
         self.dtype = dtype
         self.is_attention_free = False
+        self.disable_sliding_window = disable_sliding_window
 
     def __getattr__(self, name):
+        """Delegate attribute lookups to ModelConfig."""
         # 1. Check if ModelConfig actually has this attribute
         if not hasattr(ModelConfig, name):
             raise AttributeError(
@@ -86,9 +89,13 @@ def create_mock_vllm_config(
     tensor_parallel_size=1,
     pipeline_parallel_size=1,
     enable_expert_parallel=False,
+    disable_sliding_window=False,
 ) -> SimpleNamespace:
+    """Create a mock VllmConfig object for unit tests."""
     vllm_config = SimpleNamespace()
-    vllm_config.model_config = MockModelConfig(hf_config, model_dtype)
+    vllm_config.model_config = MockModelConfig(
+        hf_config, model_dtype, disable_sliding_window=disable_sliding_window
+    )
 
     vllm_config.cache_config = SimpleNamespace()
     vllm_config.cache_config.cache_dtype = cache_dtype
@@ -126,7 +133,56 @@ def test_base_config_parser():
     assert result.num_attention_heads == 16
     assert result.num_hidden_layers == 24
     assert result.weight_byte_size == 2  # float16 is 2 bytes
-    assert result.activation_byte_size == 2  # default activation size
+    assert result.activation_byte_size == 2  # float16 activations are 2 bytes
+
+
+@pytest.mark.parametrize(
+    "model_dtype,expected_byte_size",
+    [("float16", 2), ("bfloat16", 2), ("float32", 4)],
+)
+def test_base_config_parser_activation_byte_size_follows_dtype(
+    model_dtype: str, expected_byte_size: int
+):
+    """Activation traffic is sized from the model's compute dtype.
+
+    Activations are produced in the model dtype, so pinning them to 2 bytes
+    under-reported memory traffic by half for a float32 model while that same
+    model's weight traffic was sized correctly.
+    """
+    hf_config = Qwen3Config(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_hidden_layers=24,
+    )
+    vllm_config = create_mock_vllm_config(hf_config, model_dtype=model_dtype)
+
+    result = BaseConfigParser().parse(ParsedArgs(), vllm_config)
+
+    assert result.activation_byte_size == expected_byte_size
+    assert result.weight_byte_size == expected_byte_size
+
+
+def test_attention_read_bytes_scale_with_activation_dtype():
+    """A float32 model reads twice the activation bytes of a bfloat16 one."""
+    hf_config = Qwen3Config(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        num_hidden_layers=4,
+    )
+    ctx = ExecutionContext.from_single_request(
+        num_tokens=128, context_len=128, is_prefill=True
+    )
+
+    bf16 = AttentionMetrics.from_vllm_config(
+        create_mock_vllm_config(hf_config, model_dtype="bfloat16")
+    ).get_read_bytes_breakdown(ctx)
+    fp32 = AttentionMetrics.from_vllm_config(
+        create_mock_vllm_config(hf_config, model_dtype="float32")
+    ).get_read_bytes_breakdown(ctx)
+
+    # qkv_input is pure activation traffic, so it doubles exactly.
+    assert fp32["qkv_input"] == 2 * bf16["qkv_input"]
 
 
 def test_base_attention_config_parser_with_gqa():
@@ -147,8 +203,7 @@ def test_base_attention_config_parser_with_gqa():
 
 
 def test_base_attention_config_parser_without_gqa():
-    """
-    Test BaseAttentionConfigParser defaults to MHA when num_key_value_heads not
+    """Test BaseAttentionConfigParser defaults to MHA when num_key_value_heads not
     specified.
     """
     hf_config = Qwen3Config(
@@ -698,9 +753,7 @@ def test_ffn_per_gpu_with_pipeline_parallelism():
 
 
 def test_moe_per_gpu_with_expert_parallelism():
-    """
-    Test MoE metrics with expert parallelism - verifies num_activated_experts bug fix.
-    """
+    """Test MoE metrics with expert parallelism, covering num_activated_experts."""
     hf_config = Qwen3MoeConfig(
         hidden_size=2048,
         intermediate_size=8192,
@@ -756,9 +809,7 @@ def test_moe_per_gpu_with_expert_parallelism():
 
 
 def test_moe_per_gpu_expert_activation_accounting():
-    """
-    Test that MoE correctly accounts for expert activations with small batch sizes.
-    """
+    """Test MoE expert-activation accounting with small batch sizes."""
     hf_config = Qwen3MoeConfig(
         hidden_size=2048,
         intermediate_size=8192,
@@ -1336,3 +1387,364 @@ def test_mla_attention_scaling_with_layers():
     assert double_metrics.get_num_flops(ctx) == 2 * base_metrics.get_num_flops(ctx)
     assert double_metrics.get_read_bytes(ctx) == 2 * base_metrics.get_read_bytes(ctx)
     assert double_metrics.get_write_bytes(ctx) == 2 * base_metrics.get_write_bytes(ctx)
+
+
+#### Sliding Window Attention Tests ####
+
+
+def test_sliding_window_parser_no_swa():
+    """Test parser defaults to no SWA when sliding_window is unset."""
+    hf_config = Qwen3Config(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_hidden_layers=12,
+    )
+    vllm_config = create_mock_vllm_config(hf_config)
+    parser = SlidingWindowAttentionParser()
+    args = ParsedArgs()
+    args.num_hidden_layers = 12
+    result = parser.parse(args, vllm_config)
+
+    assert result.sliding_window is None
+    assert result.num_swa_layers == 0
+
+
+def test_sliding_window_parser_uniform_swa():
+    """Test SlidingWindowAttentionParser for full SWA model (e.g. Mistral)."""
+    hf_config = MistralConfig(
+        hidden_size=4096,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        num_hidden_layers=32,
+        sliding_window=4096,
+    )
+    vllm_config = create_mock_vllm_config(hf_config)
+    parser = SlidingWindowAttentionParser()
+    args = ParsedArgs()
+    args.num_hidden_layers = 32
+    result = parser.parse(args, vllm_config)
+
+    assert result.sliding_window == 4096
+    assert result.num_swa_layers == 32
+
+
+def test_sliding_window_parser_explicit_layer_types():
+    """Test SlidingWindowAttentionParser with explicit layer_types list."""
+    hf_config = Qwen3Config(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_hidden_layers=6,
+    )
+    hf_config.sliding_window = 2048
+    hf_config.layer_types = [
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "full_attention",
+    ]
+    vllm_config = create_mock_vllm_config(hf_config)
+    parser = SlidingWindowAttentionParser()
+    args = ParsedArgs()
+    args.num_hidden_layers = 6
+    result = parser.parse(args, vllm_config)
+
+    assert result.sliding_window == 2048
+    assert result.num_swa_layers == 3
+
+
+def test_sliding_window_parser_gemma2():
+    """Test SlidingWindowAttentionParser detects alternating SWA for Gemma 2."""
+    hf_config = Gemma2Config(
+        hidden_size=2304,
+        num_attention_heads=8,
+        num_key_value_heads=4,
+        num_hidden_layers=26,
+        sliding_window=4096,
+    )
+    vllm_config = create_mock_vllm_config(hf_config)
+    parser = SlidingWindowAttentionParser()
+    args = ParsedArgs()
+    args.num_hidden_layers = 26
+    result = parser.parse(args, vllm_config)
+
+    assert result.sliding_window == 4096
+    # Gemma 2 alternates: even layers are sliding window (26 layers -> 13 SWA)
+    assert result.num_swa_layers == 13
+
+
+def test_sliding_window_parser_disabled():
+    """Test SlidingWindowAttentionParser respects disable_sliding_window."""
+    hf_config = MistralConfig(
+        hidden_size=4096,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        num_hidden_layers=32,
+        sliding_window=4096,
+    )
+    vllm_config = create_mock_vllm_config(hf_config, disable_sliding_window=True)
+    parser = SlidingWindowAttentionParser()
+    args = ParsedArgs()
+    args.num_hidden_layers = 32
+    result = parser.parse(args, vllm_config)
+
+    assert result.sliding_window is None
+    assert result.num_swa_layers == 0
+
+
+def test_attention_metrics_sliding_window_flops_bounds():
+    """Test that SWA clamps attention core FLOPs when context_len > sliding_window."""
+    sw = 1024
+    dense_config = MistralConfig(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        num_hidden_layers=8,
+        sliding_window=None,
+    )
+    swa_config = MistralConfig(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        num_hidden_layers=8,
+        sliding_window=sw,
+    )
+
+    dense_vllm = create_mock_vllm_config(dense_config)
+    swa_vllm = create_mock_vllm_config(swa_config)
+
+    dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm)
+    swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm)
+
+    # 1. When context <= sliding_window: FLOPs should be identical
+    short_ctx = ExecutionContext.from_single_request(
+        num_tokens=1, context_len=512, is_prefill=False
+    )
+    assert swa_metrics.get_num_flops(short_ctx) == dense_metrics.get_num_flops(
+        short_ctx
+    )
+
+    # 2. When context > sliding_window: FLOPs for SWA must be strictly smaller
+    long_ctx = ExecutionContext.from_single_request(
+        num_tokens=1, context_len=8192, is_prefill=False
+    )
+    dense_flops = dense_metrics.get_num_flops_breakdown(long_ctx)
+    swa_flops = swa_metrics.get_num_flops_breakdown(long_ctx)
+
+    # Projections (qkv_proj, out_proj) are unaffected by sliding window
+    assert swa_flops["qkv_proj"] == dense_flops["qkv_proj"]
+    assert swa_flops["out_proj"] == dense_flops["out_proj"]
+
+    # Attention core FLOPs are clamped: 8192 vs 1024 (8x difference)
+    assert swa_flops["attn_qk"] < dense_flops["attn_qk"]
+    assert swa_flops["attn_av"] < dense_flops["attn_av"]
+    assert swa_flops["attn_qk"] * 8 == dense_flops["attn_qk"]
+    assert swa_flops["attn_av"] * 8 == dense_flops["attn_av"]
+
+
+def test_attention_metrics_sliding_window_read_bytes():
+    """Test that SWA bounds decode KV-cache read bytes to sliding_window."""
+    sw = 2048
+    dense_config = MistralConfig(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=4,
+        num_hidden_layers=4,
+        sliding_window=None,
+    )
+    swa_config = MistralConfig(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=4,
+        num_hidden_layers=4,
+        sliding_window=sw,
+    )
+
+    dense_vllm = create_mock_vllm_config(dense_config)
+    swa_vllm = create_mock_vllm_config(swa_config)
+
+    dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm)
+    swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm)
+
+    # Long decode context: 8192 tokens
+    ctx = ExecutionContext.from_single_request(
+        num_tokens=1, context_len=8192, is_prefill=False
+    )
+    dense_reads = dense_metrics.get_read_bytes_breakdown(ctx)
+    swa_reads = swa_metrics.get_read_bytes_breakdown(ctx)
+
+    # Attention input read traffic is lower for SWA
+    assert swa_reads["attn_input"] < dense_reads["attn_input"]
+
+    # Weight reads and projection inputs are identical
+    assert swa_reads["qkv_weight"] == dense_reads["qkv_weight"]
+    assert swa_reads["out_weight"] == dense_reads["out_weight"]
+    assert swa_reads["qkv_input"] == dense_reads["qkv_input"]
+
+
+def test_attention_metrics_hybrid_alternating_layers():
+    """Test hybrid model where half layers are full attention and half are SWA."""
+    sw = 1024
+    hybrid_config = Qwen3Config(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        num_hidden_layers=4,
+    )
+    hybrid_config.sliding_window = sw
+    hybrid_config.layer_types = [
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "full_attention",
+    ]
+
+    vllm_config = create_mock_vllm_config(hybrid_config)
+    metrics = AttentionMetrics.from_vllm_config(vllm_config)
+
+    assert metrics.num_swa_layers == 2
+    assert metrics.sliding_window == sw
+
+    ctx = ExecutionContext.from_single_request(
+        num_tokens=1, context_len=4096, is_prefill=False
+    )
+    flops = metrics.get_num_flops_breakdown(ctx)
+
+    # 2 full layers attend to 4096; 2 SWA layers attend to 1024
+    # Expected TC total = 2 * 4096 + 2 * 1024 = 8192 + 2048 = 10240
+    # Formula: 2 * q * d * (L_full * 4096 + L_swa * 1024)
+    q = metrics.num_attention_heads
+    d = metrics.head_dim
+    expected_core_flops = 2 * q * d * (2 * 4096 + 2 * 1024)
+    assert flops["attn_qk"] == expected_core_flops
+    assert flops["attn_av"] == expected_core_flops
+
+
+def test_attention_metrics_sliding_window_per_gpu():
+    """Test that SWA metrics scale properly with tensor and pipeline parallelism."""
+    sw = 1024
+    hf_config = MistralConfig(
+        hidden_size=4096,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        num_hidden_layers=16,
+        sliding_window=sw,
+    )
+    single_vllm = create_mock_vllm_config(hf_config)
+    tp_vllm = create_mock_vllm_config(hf_config, tensor_parallel_size=2)
+    pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2)
+
+    single_metrics = AttentionMetrics.from_vllm_config(single_vllm)
+    tp_metrics = AttentionMetrics.from_vllm_config(tp_vllm)
+    pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm)
+
+    ctx = ExecutionContext.from_single_request(
+        num_tokens=1, context_len=4096, is_prefill=False
+    )
+
+    # Both TP=2 and PP=2 should halve the per-GPU attention core FLOPs
+    single_flops = single_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
+    tp_flops = tp_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
+    pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
+
+    assert tp_flops["attn_qk"] == single_flops["attn_qk"] // 2
+    assert pp_flops["attn_qk"] == single_flops["attn_qk"] // 2
+
+
+def test_attention_metrics_hybrid_sliding_window_per_gpu():
+    """Test that hybrid SWA models preserve L_full + L_swa == L under PP."""
+    sw = 1024
+    hybrid_config = Gemma2Config(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        num_hidden_layers=30,
+        sliding_window=sw,
+    )
+    pp_vllm = create_mock_vllm_config(hybrid_config, pipeline_parallel_size=2)
+    pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm)
+
+    ctx = ExecutionContext.from_single_request(
+        num_tokens=1, context_len=4096, is_prefill=False
+    )
+    pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
+    pp_read_bytes = pp_metrics.get_read_bytes_breakdown(ctx, per_gpu=True)
+
+    # Under PP=2, each global count is halved: 30 layers -> 15 per GPU, and the
+    # 15 SWA / 15 full split becomes 7.5 each. Gemma 2 alternates, so an even
+    # split is the correct average stage; deriving it after reducing L instead
+    # produced a lopsided 8/7.
+    q = pp_metrics.num_attention_heads
+    d = pp_metrics.head_dim
+    expected_pp_attn = int(2 * q * d * (7.5 * 4096 + 7.5 * sw))
+    assert pp_flops["attn_qk"] == expected_pp_attn
+    assert pp_flops["attn_av"] == expected_pp_attn
+    assert "attn_input" in pp_read_bytes
+
+
+def test_attention_metrics_minority_swa_survives_pipeline_parallelism():
+    """A minority SWA sublayer must not round away under PP.
+
+    Regression test for a 6-layer model with a single sliding-window layer at
+    pp_size=2: deriving the split after reducing L gave
+    round(3 * 1 / 6) == 0, billing all three per-GPU layers as full attention
+    and dropping the sliding-window contribution entirely.
+    """
+    sw = 512
+    context_len = 4096
+    hf_config = Qwen3Config(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        num_hidden_layers=6,
+    )
+    hf_config.sliding_window = sw
+    hf_config.layer_types = ["sliding_attention"] + ["full_attention"] * 5
+
+    pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2)
+    pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm)
+    assert pp_metrics.num_swa_layers == 1
+
+    ctx = ExecutionContext.from_single_request(
+        num_tokens=1, context_len=context_len, is_prefill=False
+    )
+    pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True)
+
+    # 5 full and 1 SWA layer, each halved by PP=2: 2.5 full and 0.5 SWA.
+    q = pp_metrics.num_attention_heads
+    d = pp_metrics.head_dim
+    expected = int(2 * q * d * (2.5 * context_len + 0.5 * sw))
+    assert pp_flops["attn_qk"] == expected
+    assert pp_flops["attn_av"] == expected
+
+    # The sliding-window layer must change the result, i.e. it is not being
+    # billed at the full context length.
+    all_full = int(2 * q * d * 3 * context_len)
+    assert pp_flops["attn_qk"] < all_full
+
+
+def test_attention_metrics_per_gpu_layers_sum_to_whole_model():
+    """Per-GPU counts times pp_size must account for every layer.
+
+    Flooring the layer count per stage silently drops the remainder, so a model
+    whose layers do not divide evenly across stages under-reports its FLOPs.
+    """
+    hf_config = Qwen3Config(
+        hidden_size=2048,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        num_hidden_layers=7,
+    )
+    pp_size = 2
+    vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=pp_size)
+    metrics = AttentionMetrics.from_vllm_config(vllm_config)
+
+    ctx = ExecutionContext.from_single_request(
+        num_tokens=8, context_len=1024, is_prefill=True
+    )
+    per_gpu = metrics.get_num_flops_breakdown(ctx, per_gpu=True)
+    whole = metrics.get_num_flops_breakdown(ctx, per_gpu=False)
+
+    for key, value in whole.items():
+        assert per_gpu[key] * pp_size == pytest.approx(value, rel=1e-9)
diff --git a/tests/v1/metrics/test_ray_metrics.py b/tests/v1/metrics/test_ray_metrics.py
index d8440d3bb02a..ce4ebcbfdb0e 100644
--- a/tests/v1/metrics/test_ray_metrics.py
+++ b/tests/v1/metrics/test_ray_metrics.py
@@ -75,7 +75,6 @@ async def run(self):
 
 def test_sanitized_opentelemetry_name():
     """Test the metric name sanitization logic for Ray."""
-
     # Only a-z, A-Z, 0-9, _, test valid characters are preserved
     valid_name = "valid_metric_123_abcDEF"
     assert (
diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py
new file mode 100644
index 000000000000..369af010f4c7
--- /dev/null
+++ b/tests/v1/sample/test_diffusion_gemma_reads.py
@@ -0,0 +1,319 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-request DiffusionGemma state behind structured reads: seed canvases,
+pinned positions, read-only slots and the per-slot step cap."""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.model_executor.models.diffusion_gemma import (
+    DiffusionGemmaRequestStates,
+    _compiled_sample_step,
+    _concat_logprob_stashes,
+)
+from vllm.platforms import current_platform
+from vllm.v1.outputs import LogprobsTensors
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="the sampler state lives on the GPU"
+)
+
+CL = 8
+VOCAB = 64
+MAX_REQS = 4
+MAX_STEPS = 48
+
+
+def _states() -> DiffusionGemmaRequestStates:
+    return DiffusionGemmaRequestStates(
+        max_num_reqs=MAX_REQS,
+        canvas_length=CL,
+        vocab_size=VOCAB,
+        max_denoising_steps=MAX_STEPS,
+        device=torch.device("cuda"),
+        hidden_size=4,
+        stability_threshold=2,
+    )
+
+
+def _slots(*idx: int) -> tuple[np.ndarray, torch.Tensor]:
+    slots = np.array(idx, dtype=np.int64)
+    return slots, torch.tensor(slots, device="cuda")
+
+
+def test_seed_canvas_replaces_only_seeded_slots():
+    states = _states()
+    for slot in range(3):
+        states.add_request(slot)
+    seed = list(range(CL))
+    states.set_seed_canvas(1, seed)
+    slots, slots_gpu = _slots(0, 1, 2)
+    states.init_canvas(slots_gpu)
+    before = states.canvas[slots_gpu].clone()
+
+    states.apply_seed_canvases(slots, slots_gpu)
+
+    after = states.canvas[slots_gpu]
+    assert after[1].tolist() == seed
+    assert torch.equal(after[0], before[0])
+    assert torch.equal(after[2], before[2])
+
+
+def test_apply_seed_canvases_leaves_unseeded_batches_alone():
+    states = _states()
+    states.add_request(0)
+    slots, slots_gpu = _slots(0)
+    states.init_canvas(slots_gpu)
+    before = states.canvas[0].clone()
+
+    states.apply_seed_canvases(slots, slots_gpu)
+
+    assert torch.equal(states.canvas[0], before)
+
+
+def test_add_request_clears_seed_and_read_only():
+    states = _states()
+    states.add_request(0)
+    states.set_seed_canvas(0, [1] * CL)
+    states.set_read_only(0)
+    assert states.seeded_slots == {0}
+    assert states.read_only_slots == {0}
+
+    states.add_request(0)
+
+    assert not states.seeded_slots
+    assert not states.read_only_slots
+    assert not bool(states.has_seed[0])
+    assert not bool(states.read_only[0])
+
+
+def test_canvas_width_resets_with_the_slot():
+    states = _states()
+    states.add_request(0)
+    states.canvas_width_np[0] = 4
+    states.set_seed_canvas(0, [7, 7, 7, 7])
+    assert states.seed_canvas[0, :4].tolist() == [7, 7, 7, 7]
+
+    states.add_request(0)
+
+    assert states.canvas_width_np[0] == CL
+
+
+def test_remove_request_forgets_the_slot():
+    states = _states()
+    states.add_request(0)
+    states.set_seed_canvas(0, [1] * CL)
+    states.set_read_only(0)
+
+    states.remove_request(0)
+
+    assert not states.seeded_slots
+    assert not states.read_only_slots
+
+
+def _denoise_once(
+    states: DiffusionGemmaRequestStates,
+    slots: list[int],
+    compute_sc: bool = True,
+    width: int = CL,
+    embed_weight: torch.Tensor | None = None,
+    embed_dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """One compiled denoise step over ``slots`` with flat logits, so nothing
+    converges by stability or confidence and only the step cap can end it.
+    ``width`` below CL runs the step on [:, :width] views, as the sampler
+    does for a narrow tile."""
+    n = len(slots)
+    device = states.device
+    decode_slots = torch.tensor(slots, dtype=torch.int64, device=device)
+    decode_idx = torch.arange(n, dtype=torch.int64, device=device)
+    sampled = torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width]
+    num_sampled = torch.zeros(n, dtype=torch.int32, device=device)
+    _compiled_sample_step(
+        torch.zeros(n * width, VOCAB, device=device),
+        decode_slots,
+        decode_idx,
+        decode_slots,
+        torch.full((n,), width, dtype=torch.int64, device=device),
+        states.canvas[:, :width],
+        states.argmax_canvas[:, :width],
+        states.step,
+        states.is_encoder_phase,
+        states.confident,
+        states.self_conditioning_embeds[:, :width],
+        (
+            torch.zeros(VOCAB, 4, dtype=embed_dtype, device=device)
+            if embed_weight is None
+            else embed_weight
+        ),
+        torch.tensor(1.0, dtype=embed_dtype, device=device),
+        states.accepted_canvas_history[:, :, :width],
+        states.accepted_canvas_history_len,
+        states.max_steps,
+        states.pin_mask[:, :width],
+        states.seed_canvas[:, :width],
+        states.read_only,
+        sampled,
+        num_sampled,
+        torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device),
+        max_denoising_steps=float(MAX_STEPS),
+        t_min=0.5,
+        t_max=1.0,
+        confidence_threshold=0.1,
+        vocab_size=VOCAB,
+        CL=width,
+        ST=states.stability_threshold,
+        entropy_bound=0.1,
+        sc_vocab_start=0,
+        sc_vocab_end=VOCAB,
+        tp_size=1,
+        tp_group_name="",
+        compute_sc=compute_sc,
+    )
+    return sampled, num_sampled
+
+
+def test_pinned_positions_hold_their_seed_through_a_step():
+    states = _states()
+    states.add_request(0)
+    states.is_encoder_phase[0] = False
+    seed = list(range(10, 10 + CL))
+    states.set_seed_canvas(0, seed)
+    states.canvas[0] = torch.tensor(seed, device="cuda")
+    states.set_pins(0, [0, 1, 2, 3])
+
+    # Flat logits accept nothing, so every free position is renoised.
+    _denoise_once(states, [0], embed_weight=torch.ones(VOCAB, 4, device="cuda"))
+
+    assert states.canvas[0, :4].tolist() == seed[:4]
+    # The soft embed is zero at pinned positions and non-zero elsewhere.
+    sc = states.self_conditioning_embeds[0]
+    assert not sc[:4].any()
+    assert sc[4:].any()
+
+
+def test_add_request_clears_pins():
+    states = _states()
+    states.add_request(0)
+    states.set_seed_canvas(0, [1] * CL)
+    states.set_pins(0, [2, 3])
+    assert states.pin_mask[0].tolist() == [False, False, True, True] + [False] * 4
+
+    states.add_request(0)
+
+    assert not states.pin_mask[0].any()
+
+
+def test_single_step_tile_skips_self_conditioning():
+    states = _states()
+    states.add_request(0)
+    states.is_encoder_phase[0] = False
+    states.self_conditioning_embeds[0] = 1.0
+
+    _denoise_once(states, [0], compute_sc=False)
+
+    assert not states.self_conditioning_embeds[0].any()
+
+
+@pytest.mark.parametrize("stance", ["default", "force_eager"])
+def test_self_conditioning_stores_a_bf16_model_in_the_fp32_buffer(stance):
+    # The model's embeddings are bf16 while the buffer is fp32. Compiled code
+    # casts on the store, but eager does not, and the step runs eager once
+    # torch.compile hits its recompile limit.
+    states = _states()
+    states.add_request(0)
+    states.is_encoder_phase[0] = False
+
+    with torch.compiler.set_stance(stance):
+        _denoise_once(
+            states,
+            [0],
+            embed_weight=torch.ones(VOCAB, 4, dtype=torch.bfloat16, device="cuda"),
+            embed_dtype=torch.bfloat16,
+        )
+
+    assert states.self_conditioning_embeds.dtype == torch.float32
+    # uniform probs @ an all-ones embedding
+    assert torch.allclose(
+        states.self_conditioning_embeds[0], torch.ones(CL, 4, device="cuda")
+    )
+
+
+def test_narrow_tile_leaves_the_rest_of_the_canvas_alone():
+    states = _states()
+    states.add_request(0)
+    states.is_encoder_phase[0] = False
+    states.canvas[0] = 5
+    states.argmax_canvas[0] = 5
+
+    _denoise_once(states, [0], width=4)
+
+    # Columns past the width are untouched; the step counted.
+    assert states.canvas[0, 4:].tolist() == [5] * (CL - 4)
+    assert states.argmax_canvas[0, 4:].tolist() == [5] * (CL - 4)
+    assert states.step[0] == 1
+
+
+def test_step_cap_is_per_slot():
+    states = _states()
+    for slot in (0, 1):
+        states.add_request(slot)
+        states.is_encoder_phase[slot] = False
+    states.max_steps[0] = 1
+
+    _denoise_once(states, [0, 1])
+
+    # Slot 0 hit its cap and moves to commit. Slot 1 keeps denoising.
+    assert states.is_encoder_phase[:2].tolist() == [True, False]
+    assert states.step[:2].tolist() == [1, 1]
+
+
+def _stash(rows: int, width: int, first_id: int) -> LogprobsTensors:
+    ids = torch.arange(first_id, first_id + rows * width).reshape(rows, width)
+    return LogprobsTensors(
+        logprob_token_ids=ids,
+        logprobs=-ids.float(),
+        selected_token_ranks=torch.zeros(rows, dtype=torch.int64),
+    )
+
+
+def test_stashes_of_different_widths_join():
+    # A read that asked for 10 label ids and a generation that asked for 2
+    # logprobs converged in different steps, so their stashes differ in width.
+    wide, narrow = _stash(2, 11, 100), _stash(3, 3, 0)
+
+    out = _concat_logprob_stashes([narrow, wide], [0, 3])
+
+    assert out.logprob_token_ids.shape == (5, 11)
+    assert out.logprobs.shape == (5, 11)
+    assert out.cu_num_generated_tokens == [0, 3]
+    # The narrow rows keep their values and pad with id 0 at -inf.
+    assert torch.equal(out.logprob_token_ids[:3, :3], narrow.logprob_token_ids)
+    assert not out.logprob_token_ids[:3, 3:].any()
+    assert torch.isneginf(out.logprobs[:3, 3:]).all()
+    assert torch.equal(out.logprobs[3:], wide.logprobs)
+
+
+@pytest.mark.parametrize("width", [4, CL])
+@pytest.mark.parametrize("steps", [1, 2])
+def test_read_emits_at_convergence_while_generation_waits_for_commit(width, steps):
+    states = _states()
+    for slot in (2, 0):
+        states.add_request(slot)
+        states.is_encoder_phase[slot] = False
+        states.max_steps[slot] = steps
+    states.set_read_only(2)
+
+    for step in range(steps):
+        sampled, counts = _denoise_once(states, [2, 0], width=width)
+        assert counts.tolist() == ([width, 0] if step == steps - 1 else [0, 0])
+
+    assert torch.equal(sampled[0], states.argmax_canvas[2, :width].int())
+    assert not states.is_encoder_phase[2]
+    assert states.is_encoder_phase[0]
+    assert not states.self_conditioning_embeds[2].any()
+
+    sampled, counts = _denoise_once(states, [0], width=width)
+    assert counts.tolist() == [width]
+    assert torch.equal(sampled[0], states.argmax_canvas[0, :width].int())
diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py
index 3bf26d51706e..d2cb99bc26fd 100644
--- a/tests/v1/sample/test_logprobs.py
+++ b/tests/v1/sample/test_logprobs.py
@@ -127,6 +127,7 @@ def _repeat_logprob_config(
       `logprob_prompt_logprob_list` enough times to match the
       number of `test_prompts`, or else is truncated to match
       the number of `test_prompts`
+
     """
     num_test_prompts = len(test_prompts)
     # Make sure there is a logprobs configuration for each test prompt
@@ -312,7 +313,7 @@ def test_get_logprobs_and_prompt_logprobs(
     temperature: float,
     example_prompts: list[str],
 ) -> None:
-    """Test V1 Engine logprobs & prompt logprobs
+    """Test V1 Engine logprobs & prompt logprobs.
 
     Exercise a variety of combinations of `logprobs` and `prompt_logprobs`
     settings and validate that
@@ -337,6 +338,7 @@ def test_get_logprobs_and_prompt_logprobs(
       batch_logprobs_composition: logprobs configuration for test batch
       temperature: "temperature" sampling parameter
       example_prompts: example prompt fixture
+
     """
     vllm_config = vllm_model.llm.llm_engine.vllm_config
     do_apc = vllm_config.cache_config.enable_prefix_caching
@@ -389,7 +391,7 @@ def test_get_logprobs_and_prompt_logprobs(
 
 
 def test_max_logprobs():
-    """vLLM v1 engine should fail a request with `logprobs > max_logprobs`
+    """VLLM v1 engine should fail a request with `logprobs > max_logprobs`
     Should also fail for `prompt_logprobs > max_logprobs`
     APC should not matter as this test checks basic request validation.
     """
@@ -431,11 +433,12 @@ def test_logprob_token_ids_validate_vocab_bounds_invalid(token_ids: list[int]):
 
 
 def test_none_logprobs(vllm_model, example_prompts):
-    """Engine should return `logprobs` and `prompt_logprobs` as `None`
+    """Engine should return `logprobs` and `prompt_logprobs` as `None`.
 
     Args:
       vllm_model: vLLM model fixture
       example_prompts: list of example prompts (test fixture)
+
     """
     max_tokens = 5
 
@@ -459,11 +462,12 @@ def test_none_logprobs(vllm_model, example_prompts):
 
 
 def test_zero_logprobs(vllm_model, example_prompts):
-    """Engine should return sampled token and prompt token logprobs
+    """Engine should return sampled token and prompt token logprobs.
 
     Args:
       vllm_model: vLLM model fixture
       example_prompts: list of example prompts (test fixture)
+
     """
     max_tokens = 5
 
@@ -491,10 +495,11 @@ def test_zero_logprobs(vllm_model, example_prompts):
 
 
 def test_all_logprobs(example_prompts):
-    """Engine should return all vocabulary logprobs and prompt logprobs
+    """Engine should return all vocabulary logprobs and prompt logprobs.
 
     Args:
       example_prompts: list of example prompts (test fixture)
+
     """
     with VllmRunner(
         "facebook/opt-125m",
@@ -1127,6 +1132,7 @@ def test_spec_decode_logprobs(
         logprobs_mode: logprobs mode.
         model_setup: Tuple of (method, base model name,
             speculative_config dict, top_logprobs).
+
     """
     from vllm import LLM
 
@@ -1224,7 +1230,6 @@ def test_prompt_logprobs_with_chunking_and_preemption():
     This test ensures that the num_prompt_logprobs tracking persists
     across preemptions and prefill chunks.
     """
-
     # Create prompts that will trigger chunking and preemption
     prompts = [
         "The following numbers of the sequence "
diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py
index b02e53af8dac..231f6adbb2d7 100644
--- a/tests/v1/sample/test_rejection_sampler.py
+++ b/tests/v1/sample/test_rejection_sampler.py
@@ -131,7 +131,7 @@ def create_sampling_metadata(
 
 ########################### Tests for Greedy Sampling ###################
 def test_perfect_match(rejection_sampler):
-    """Test when output tokens perfectly match speculated tokens"""
+    """Test when output tokens perfectly match speculated tokens."""
     spec_tokens = [[1, 2, 3]]
     output_tokens = [[1, 2, 3, 4]]  # 4 is the bonus token
 
@@ -152,7 +152,7 @@ def test_perfect_match(rejection_sampler):
 
 
 def test_early_mismatch(rejection_sampler):
-    """Test when there's an early mismatch in tokens"""
+    """Test when there's an early mismatch in tokens."""
     spec_tokens = [[1, 2, 3]]
     output_tokens = [[1, 5, 3, 4]]  # Mismatch at position 1
 
@@ -177,7 +177,7 @@ def test_early_mismatch(rejection_sampler):
 
 
 def test_multiple_sequences(rejection_sampler):
-    """Test handling multiple sequences of speculated tokens"""
+    """Test handling multiple sequences of speculated tokens."""
     spec_tokens = [[1, 2], [3]]
     output_tokens = [[1, 2, 5], [3, 4]]  # Two sequences with bonus tokens 5 and 4
 
@@ -202,7 +202,7 @@ def test_multiple_sequences(rejection_sampler):
 
 
 def test_single_token_sequence(rejection_sampler):
-    """Test handling sequences with single token"""
+    """Test handling sequences with single token."""
     spec_tokens = [[1]]
     output_tokens = [[1, 2]]  # Single token with bonus token 2
 
@@ -223,7 +223,7 @@ def test_single_token_sequence(rejection_sampler):
 
 
 def test_empty_sequence(rejection_sampler):
-    """Test handling empty sequence of speculated tokens"""
+    """Test handling empty sequence of speculated tokens."""
     spec_tokens: list[list[int]] = [[]]
     output_tokens = [[5]]  # Just the bonus token
 
@@ -244,7 +244,7 @@ def test_empty_sequence(rejection_sampler):
 
 
 def test_multiple_mismatches(rejection_sampler):
-    """Test handling multiple sequences with mismatches"""
+    """Test handling multiple sequences with mismatches."""
     spec_tokens = [[1, 2, 3], [4, 5, 6]]
     output_tokens = [[1, 2, 7, 6], [4, 8, 6, 9]]  # Mismatches in both sequences
 
@@ -286,7 +286,7 @@ def test_multiple_mismatches(rejection_sampler):
     ],
 )
 def test_parametrized_cases(rejection_sampler, spec_tokens, output_tokens, expected):
-    """Parametrized test for various matching scenarios"""
+    """Parametrized test for various matching scenarios."""
     metadata = create_sampling_metadata(all_greedy=True)
     logits = create_logits_tensor(output_tokens)
     bonus_token_tensor = torch.tensor(
@@ -486,6 +486,7 @@ def estimate_rejection_sampling_pdf(
 
     Returns:
         Estimated probability distribution of the output tokens.
+
     """
     mock_sampler = Mock(spec=Sampler)
     mock_sampler.logprobs_mode = "raw_logprobs"
@@ -656,7 +657,7 @@ def _test_masked_logits(
 
 @pytest.mark.parametrize("top_k", [1, 5, 99])
 def test_top_k(rejection_sampler, top_k):
-    """Test rejection sampling with top-k sampling"""
+    """Test rejection sampling with top-k sampling."""
     vocab_size = 100
     batch_size = 100
     num_draft_tokens = 3
@@ -699,7 +700,7 @@ def test_top_k(rejection_sampler, top_k):
 
 @pytest.mark.parametrize("top_p", [0.5, 0.9, 0.99])
 def test_top_p(rejection_sampler, top_p):
-    """Test rejection sampling with top-p sampling"""
+    """Test rejection sampling with top-p sampling."""
     vocab_size = 100
     batch_size = 100
     num_draft_tokens = 3
@@ -746,7 +747,7 @@ def test_top_p(rejection_sampler, top_p):
 
 ########################### Tests for Logit Processors ###################
 def test_frequency_penalties(rejection_sampler):
-    """Test rejection sampling with frequency penalties"""
+    """Test rejection sampling with frequency penalties."""
     spec_tokens = [[1, 1, 1], [], [1, 1, 1]]
     output_tokens = [[1, 1, 1, 1], [7], [1, 1, 1, 1]]  # 1, 7 and 1 are the bonus tokens
 
@@ -829,7 +830,7 @@ def test_bad_words(rejection_sampler):
 
 
 def test_allowed_token_ids(rejection_sampler):
-    """Test rejection sampling with allowed token ids"""
+    """Test rejection sampling with allowed token ids."""
     spec_tokens = [[1, 2, 10], [10, 5, 3], [7, 10, 12]]
     output_tokens = [[1, 2, 10, 5], [10, 5, 10, 5], [7, 10, 12, 5]]
     # Not allowed tokens:
diff --git a/tests/v1/sample/test_sampler.py b/tests/v1/sample/test_sampler.py
index c67199fa4077..b6d4fcc9e8a3 100644
--- a/tests/v1/sample/test_sampler.py
+++ b/tests/v1/sample/test_sampler.py
@@ -169,8 +169,7 @@ def _create_default_sampling_metadata(
 def _create_weighted_output_token_list(
     batch_size: int, vocab_size: int
 ) -> tuple[list[list[int]], list[list[int]]]:
-    """
-    Creates an output token list where each token occurs a distinct
+    """Creates an output token list where each token occurs a distinct
     number of times.
 
     For each batch, a random subset of token IDs is selected from the
@@ -185,6 +184,7 @@ def _create_weighted_output_token_list(
             - The second element is a list of distinct token IDs for each
               batch, ordered by their frequency in the corresponding output
               list.
+
     """
     output_token_ids: list[list[int]] = []
     sorted_token_ids_in_output: list[list[int]] = []
@@ -206,8 +206,7 @@ def _create_weighted_output_token_list(
 def test_sampler_presence_penalty(
     device: str, batch_size: int, presence_penalty: float
 ):
-    """
-    Test to verify that if presence penalty is enabled then tokens
+    """Test to verify that if presence penalty is enabled then tokens
     are penalized as per their presence in the existing output.
     """
     torch.set_default_device(device)
@@ -256,8 +255,7 @@ def test_sampler_presence_penalty(
 def test_sampler_frequency_penalty(
     device: str, batch_size: int, frequency_penalty: float
 ):
-    """
-    Test to verify that if frequency penalty is enabled then tokens are
+    """Test to verify that if frequency penalty is enabled then tokens are
     penalized as per their frequency of occurrence.
     """
     torch.set_default_device(device)
@@ -312,8 +310,7 @@ def test_sampler_frequency_penalty(
 def test_sampler_repetition_penalty(
     device: str, batch_size: int, repetition_penalty: float
 ):
-    """
-    Test to verify that when the repetition penalty is enabled, tokens
+    """Test to verify that when the repetition penalty is enabled, tokens
     are penalized based on their presence in the prompt or the existing
     output.
     """
@@ -370,8 +367,7 @@ def test_sampler_repetition_penalty(
 def test_sampler_allowed_token_ids(
     device: str, batch_size: int, num_allowed_token_ids: int
 ):
-    """
-    Test to verify that when the repetition penalty is enabled, tokens
+    """Test to verify that when the repetition penalty is enabled, tokens
     are penalized based on their presence in the prompt or the existing
     output.
     """
@@ -416,8 +412,7 @@ def test_sampler_allowed_token_ids(
 def test_sampler_bad_words(
     device: str, batch_size: int, bad_words_lengths: tuple[int, ...]
 ):
-    """
-    Test to verify that when the bad words restriction is present, tokens
+    """Test to verify that when the bad words restriction is present, tokens
     are penalized based on their match with the bad words.
     """
     torch.set_default_device(device)
diff --git a/tests/v1/sample/test_sampling_params_e2e.py b/tests/v1/sample/test_sampling_params_e2e.py
index d385e96b7a21..12c1048ff2c1 100644
--- a/tests/v1/sample/test_sampling_params_e2e.py
+++ b/tests/v1/sample/test_sampling_params_e2e.py
@@ -17,7 +17,6 @@ def llm() -> LLM:
 
 def test_n_gt_1(llm):
     """ParallelSampling is supported."""
-
     params = SamplingParams(n=3)
     outputs = llm.generate(PROMPT, params)
     assert len(outputs[0].outputs) == 3
@@ -25,7 +24,6 @@ def test_n_gt_1(llm):
 
 def test_penalties(llm):
     """Check that we do not get errors if applied."""
-
     params = SamplingParams(
         temperature=1.2,
         presence_penalty=1.2,
@@ -40,7 +38,6 @@ def test_penalties(llm):
 
 def test_stop(llm):
     """Check that we respect the stop words."""
-
     output = llm.generate(PROMPT, SamplingParams(temperature=0))
     split_text = output[0].outputs[0].text.split()
 
@@ -64,7 +61,6 @@ def test_stop(llm):
 
 def test_stop_token_ids(llm):
     """Check that we respect the stop token ids."""
-
     output = llm.generate(PROMPT, SamplingParams(temperature=0))
 
     stop_token_id_0 = output[0].outputs[0].token_ids[5]
@@ -83,7 +79,6 @@ def test_stop_token_ids(llm):
 
 def test_detokenize_false(llm):
     """Check that detokenize=False option works."""
-
     output = llm.generate(PROMPT, SamplingParams(detokenize=False))
     assert len(output[0].outputs[0].token_ids) > 0
     assert len(output[0].outputs[0].text) == 0
@@ -106,7 +101,6 @@ def test_detokenize_false(llm):
 
 def test_bad_words(llm):
     """Check that we respect bad words."""
-
     tokenizer = llm.get_tokenizer()
 
     def contains_bad_word(text: str, tokens: list[int], bad_word: str) -> bool:
@@ -147,7 +141,6 @@ def contains_bad_word(text: str, tokens: list[int], bad_word: str) -> bool:
 
 def test_allowed_token_ids(llm):
     """Check that we can use allowed_token_ids."""
-
     TOKEN_ID = 10
     allowed_token_ids = [TOKEN_ID]
     output = llm.generate(PROMPT, SamplingParams(allowed_token_ids=allowed_token_ids))
@@ -176,7 +169,6 @@ def test_allowed_token_ids(llm):
 
 def test_seed(llm):
     """Check that seed impacts randomness."""
-
     out_1 = llm.generate(PROMPT, SamplingParams(seed=42))
     out_2 = llm.generate(PROMPT, SamplingParams(seed=42))
     out_3 = llm.generate(PROMPT, SamplingParams(seed=43))
diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py
index fdf704169288..c0b979245659 100644
--- a/tests/v1/sample/test_topk_topp_sampler.py
+++ b/tests/v1/sample/test_topk_topp_sampler.py
@@ -25,7 +25,7 @@ def _flashinfer_topk_topp_supported() -> bool:
 
     Mirrors the gate in `TopKTopPSampler.__init__`: CUDA + flashinfer
     importable + GPU compute capability supported by the FlashInfer
-    backend.
+    backend + more than 16 SMs.
     """
     if not current_platform.is_cuda():
         return False
@@ -38,7 +38,10 @@ def _flashinfer_topk_topp_supported() -> bool:
     capability = current_platform.get_device_capability()
     if capability is None:
         return False
-    return FlashInferBackend.supports_compute_capability(capability)
+    return FlashInferBackend.supports_compute_capability(capability) and (
+        current_platform.num_compute_units(torch.accelerator.current_device_index())
+        > 16
+    )
 
 
 FLASHINFER_TOPK_TOPP_SUPPORTED = _flashinfer_topk_topp_supported()
@@ -50,8 +53,7 @@ def _seed_default_generator(seed: int) -> None:
 
 @pytest.fixture(autouse=True)
 def reset_default_device():
-    """
-    Explicitly set the default device, which can affect subsequent tests.
+    """Explicitly set the default device, which can affect subsequent tests.
     Adding this fixture helps avoid this problem.
     """
     original_device = torch.get_default_device()
@@ -79,6 +81,10 @@ class MockPlatform:
         def is_cuda():
             return False
 
+        @staticmethod
+        def is_cuda_alike():
+            return True
+
         @staticmethod
         def is_cpu():
             return False
@@ -161,78 +167,42 @@ def test_topk_impl_equivalence():
     assert torch.allclose(result1, result2)
 
 
-@pytest.mark.skip(
-    reason="FlashInfer top-k/top-p renorm comparison fails; "
-    "needs investigation of tolerance threshold or "
-    "interface differences between Python and FlashInfer implementations"
+@pytest.mark.skipif(
+    not FLASHINFER_TOPK_TOPP_SUPPORTED,
+    reason="Requires CUDA and flashinfer",
 )
-def test_flashinfer_sampler():
-    """
-    This test verifies that the FlashInfer top-k and top-p sampling
-    implementation produces the same results as the Python implementation.
-
-    NOTE: FlashInfer did not directly expose an interface for fused top-k and
-    top-p prob renorm (it did provide fused sampling but we cannot compare
-    sampling results due to randomness), so we will compare the probability
-    renormed consequently by top-k and then top-p of FlashInfer implementation.
+def test_flashinfer_sampler_fallback_when_jit_cannot_target_gpu(monkeypatch):
+    """FlashInfer swallows arch-detection errors (e.g. SM 12.x with a CUDA
+    toolkit older than 12.9), leaving an empty target-arch set that makes
+    every JIT spec fail with a misleading "requires sm75 or higher" error at
+    first use, killing the engine during startup profiling (issue #42393).
+    The sampler gate must detect this and fall back to native sampling, and
+    surface the real reason when the user explicitly opted in.
     """
-    try:
-        from flashinfer.sampling import top_k_renorm_probs, top_p_renorm_probs
+    import flashinfer.jit.core
+    from flashinfer.compilation_context import CompilationContext
 
-        is_flashinfer_available = True
-    except ImportError:
-        is_flashinfer_available = False
+    from vllm.v1.sample.ops.topk_topp_sampler import flashinfer_sampler_supported
 
-    FLASHINFER_ENABLED = current_platform.is_cuda() and is_flashinfer_available
+    def fake_check_cuda_arch():
+        raise RuntimeError("FlashInfer requires GPUs with sm75 or higher")
 
-    if not FLASHINFER_ENABLED:
-        pytest.skip("FlashInfer not installed or not available on this platform.")
+    def fake_normalize_cuda_arch(major, minor):
+        raise RuntimeError("SM 12.x requires CUDA >= 12.9")
 
-    torch.set_default_device(DEVICE_TYPE)
-    generator = Generator(device=DEVICE_TYPE).manual_seed(42)
-
-    # Generate random logits
-    logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator)
-
-    # Generate various top-k and top-p values
-    k_values = torch.randint(1, 1000, (BATCH_SIZE,), generator=generator)
-    p_values = (
-        torch.rand((BATCH_SIZE,), generator=generator) * 0.5 + 0.5
-    )  # range in [0.5, 1.0]
-
-    # Sometimes disable top-k (k=vocab_size)
-    k_values.masked_fill_(
-        torch.randint(0, 2, (BATCH_SIZE,), generator=generator, dtype=torch.bool),
-        VOCAB_SIZE,
+    monkeypatch.setattr(flashinfer.jit.core, "check_cuda_arch", fake_check_cuda_arch)
+    monkeypatch.setattr(
+        CompilationContext,
+        "_normalize_cuda_arch",
+        staticmethod(fake_normalize_cuda_arch),
     )
 
-    # Sometimes disable top-p (p=1.0)
-    p_values.masked_fill_(
-        torch.randint(0, 2, (BATCH_SIZE,), generator=generator, dtype=torch.bool), 1.0
-    )
+    monkeypatch.delenv("VLLM_USE_FLASHINFER_SAMPLER", raising=False)
+    assert flashinfer_sampler_supported() is False
 
-    python_logits = apply_top_k_top_p_pytorch(
-        logits=logits.clone(),
-        k=k_values,
-        p=p_values,
-    )
-    python_probs = torch.softmax(python_logits, dim=-1)
-
-    # FlashInfer only exposed renorm interfaces for probs so convert first
-    flashinfer_probs = torch.softmax(logits.clone(), dim=-1)
-    flashinfer_probs = top_k_renorm_probs(
-        probs=flashinfer_probs,
-        top_k=k_values,
-    )
-    flashinfer_probs = top_p_renorm_probs(
-        probs=flashinfer_probs,
-        top_p=p_values,
-    )
-
-    # Compare the results
-    assert torch.allclose(python_probs, flashinfer_probs, atol=2e-2), (
-        "FlashInfer and Python sampling implementations do not match!"
-    )
+    monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "1")
+    with pytest.raises(RuntimeError, match="SM 12.x requires CUDA >= 12.9"):
+        flashinfer_sampler_supported()
 
 
 # =============================================================================
diff --git a/tests/v1/sample/utils.py b/tests/v1/sample/utils.py
index 907be3614b9c..e4e9f21092e3 100644
--- a/tests/v1/sample/utils.py
+++ b/tests/v1/sample/utils.py
@@ -15,7 +15,7 @@
 
 
 class BatchLogprobsComposition(Enum):
-    """Types of logprobs configs to include in test batch"""
+    """Types of logprobs configs to include in test batch."""
 
     NONE = 0
     SAMPLE = 1
@@ -29,7 +29,7 @@ class BatchLogprobsComposition(Enum):
 def get_test_batch(
     batch_logprobs_composition: BatchLogprobsComposition,
 ) -> BatchLogprobsSpecType:
-    """Generate logprobs configs for a batch of requests
+    """Generate logprobs configs for a batch of requests.
 
     A given request's logprobs configuration is (1) num_sample_logprobs and (2)
     num_prompt_logprobs. The batch logprobs configuration is the list of request
@@ -53,9 +53,9 @@ def get_test_batch(
       batch_logprobs_composition: types of logprobs configs to include in batch
 
     Returns:
-
       list of (Optional[num_sample_logprobs], Optional[num_prompt_logprobs])
       tuples
+
     """
     if batch_logprobs_composition == BatchLogprobsComposition.NONE:
         # No requests with sample or prompt logprobs
@@ -101,7 +101,7 @@ def assert_incr_detok_str_matches_non_incr_detok_str(
     non_incremental_detokenization_str: str,
     msg: str,
 ) -> None:
-    """Compare incrementally detok. text to non-incrementally detok. text
+    """Compare incrementally detok. text to non-incrementally detok. text.
 
     Fail if the strings mismatch after non-alphanumeric characters are stripped
     out.
@@ -120,6 +120,7 @@ def assert_incr_detok_str_matches_non_incr_detok_str(
       non_incremental_detokenization_str: non-incrementally-detokenized logprob
                                           tokens
       msg: error message if `assert` fails
+
     """
     rgx = r"[^a-zA-Z0-9]+"
     assert re.sub(rgx, "", incremental_detokenization_str) == re.sub(
@@ -128,13 +129,14 @@ def assert_incr_detok_str_matches_non_incr_detok_str(
 
 
 def compute_correct_cumulative_logprob(completion_output: CompletionOutput) -> float:
-    """Compute known-good value for evaluating cumulative logprob
+    """Compute known-good value for evaluating cumulative logprob.
 
     Args:
       completion_output: completion output from engine
 
     Returns:
       Known-good cumulative logprob value
+
     """
     token_ids = completion_output.token_ids
     logprobs = completion_output.logprobs
@@ -170,7 +172,7 @@ def create_prompt_tokens_tensor(
 
 
 class LogitsprocsTestFakes(NamedTuple):
-    """Wraps fake data structures to support testing"""
+    """Wraps fake data structures to support testing."""
 
     logits: torch.Tensor
     sampling_metadata: SamplingMetadata
@@ -186,6 +188,7 @@ def get_logitsprocs_by_cls(
 
         Returns:
           Iterator over logits processors
+
         """
         return (
             lp for lp in self.sampling_metadata.logitsprocs.all if isinstance(lp, cls)
diff --git a/tests/v1/shutdown/test_delete.py b/tests/v1/shutdown/test_delete.py
index 991ad6891856..b8bf620b33f5 100644
--- a/tests/v1/shutdown/test_delete.py
+++ b/tests/v1/shutdown/test_delete.py
@@ -11,6 +11,7 @@
     SHUTDOWN_TEST_TIMEOUT_SEC,
 )
 from vllm import LLM, SamplingParams
+from vllm.config import KernelConfig
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.platforms import current_platform
 from vllm.sampling_params import RequestOutputKind
@@ -34,12 +35,16 @@ async def test_async_llm_delete(
       model: model under test
       tensor_parallel_size: degree of tensor parallelism
       send_one_request: send one request to engine before deleting
+
     """
     if current_platform.device_count() < tensor_parallel_size:
         pytest.skip(reason="Not enough CUDA devices")
 
     engine_args = AsyncEngineArgs(
-        model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size
+        model=model,
+        enforce_eager=True,
+        tensor_parallel_size=tensor_parallel_size,
+        kernel_config=KernelConfig(enable_jit_warmup=False),
     )
 
     # Instantiate AsyncLLM; make request to complete any deferred
@@ -83,6 +88,7 @@ def test_llm_delete(
       tensor_parallel_size: degree of tensor parallelism
       enable_multiprocessing: enable workers in separate process(es)
       send_one_request: send one request to engine before deleting
+
     """
     if current_platform.device_count() < tensor_parallel_size:
         pytest.skip(reason="Not enough CUDA devices")
@@ -94,7 +100,10 @@ def test_llm_delete(
         # Instantiate LLM; make request to complete any deferred
         # initialization; then delete instance
         llm = LLM(
-            model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size
+            model=model,
+            enforce_eager=True,
+            tensor_parallel_size=tensor_parallel_size,
+            kernel_config=KernelConfig(enable_jit_warmup=False),
         )
         if send_one_request:
             llm.generate(
@@ -122,7 +131,10 @@ def test_llm_delete_inprocess(
     with monkeypatch.context() as m:
         m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
 
-        with VllmRunner(model) as vllm_model:
+        with VllmRunner(
+            model,
+            kernel_config=KernelConfig(enable_jit_warmup=False),
+        ) as vllm_model:
             if send_one_request:
                 vllm_model.generate(
                     ["Hello my name is"],
diff --git a/tests/v1/shutdown/test_startup_error.py b/tests/v1/shutdown/test_startup_error.py
index 03982aa48e71..6d6d0d91e77f 100644
--- a/tests/v1/shutdown/test_startup_error.py
+++ b/tests/v1/shutdown/test_startup_error.py
@@ -23,7 +23,6 @@
 
 def evil_method(self, *args, **kwargs):
     """Evil method that raises an exception."""
-
     if get_tensor_model_parallel_rank() == 0:
         raise Exception("Simulated Error in startup!")
 
diff --git a/tests/v1/shutdown/utils.py b/tests/v1/shutdown/utils.py
index 124254a41337..606928c5991b 100644
--- a/tests/v1/shutdown/utils.py
+++ b/tests/v1/shutdown/utils.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Shutdown test utils"""
+"""Shutdown test utils."""
 
 SHUTDOWN_TEST_TIMEOUT_SEC = 120
 SHUTDOWN_TEST_THRESHOLD_BYTES = 2 * 2**30
diff --git a/tests/v1/simple_kv_offload/test_kv_events.py b/tests/v1/simple_kv_offload/test_kv_events.py
index 3bc6c32e7ce0..68121f17e000 100644
--- a/tests/v1/simple_kv_offload/test_kv_events.py
+++ b/tests/v1/simple_kv_offload/test_kv_events.py
@@ -19,8 +19,10 @@
     SchedulerFixture,
     _alloc_and_register,
     _allocate_cp_gpu_blocks,
+    _allocate_gpu_blocks,
     _make_cp_request,
     _make_kv_cache_config,
+    _make_scratch_kv_cache_config,
     _make_vllm_config,
     make_request,
     make_scheduler_output,
@@ -714,17 +716,8 @@ def test_disk_mode_block_removed_relabeled_storage() -> None:
         assert ev.medium == MEDIUM_STORAGE
 
 
-def test_mamba_dcp_metadata_guard() -> None:
-    """Mamba+DCP: token_ids guard emits [] when capture and emission sizes differ.
-
-    The metadata capture at _prepare_eager_store_specs uses
-    g_block_size = spec.block_size * cp_world_size = 32 * 2 = 64 to slice
-    tokens. The event emission uses block_size = spec.block_size * 1 = 32
-    for Mamba (SSM state is not CP-sharded). The guard in
-    _process_store_completion detects len(token_ids) != block_size and emits
-    token_ids=[] instead of misleading data. When #49962 fixes the capture
-    path, the guard stops triggering and correct token_ids flow through.
-    """
+def test_mamba_align_skips_positional_event_metadata() -> None:
+    """Mamba align emits events only from explicit boundary handoffs."""
     vbs = BLOCK_SIZE * 2
     kv_cfg = _make_mixed_kv_cache_config(
         num_blocks=16, with_mamba=True, mamba_block_size=32
@@ -739,8 +732,6 @@ def test_mamba_dcp_metadata_guard() -> None:
     sched = fix.scheduler
     gpu_pool = fix.gpu_block_pool
 
-    # Mamba g_block_size = 32 * 2 = 64, so 2 Mamba blocks need 128 computed
-    # tokens to be "ready"; FA g_block_size = 16 * 2 = 32 needs only 64.
     req = _make_cp_request(num_blocks=4, virtual_block_size=vbs)
     fa_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=0)
     mamba_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=1)
@@ -759,32 +750,46 @@ def test_mamba_dcp_metadata_guard() -> None:
 
     events = list(sched.take_events())
     stored = [e for e in events if isinstance(e, BlockStored)]
-    assert len(stored) == 4, (
-        f"expected 4 BlockStored (2 FA + 2 Mamba), got {len(stored)}"
-    )
-
-    by_group: dict[int, list[BlockStored]] = {}
-    for ev in stored:
-        by_group.setdefault(ev.group_idx, []).append(ev)
-
-    # FA group: capture size (64) == emission size (64) → token_ids present.
-    fa_events = by_group[0]
-    assert len(fa_events) == 2
-    for ev in fa_events:
+    assert len(stored) == 2
+    assert {ev.group_idx for ev in stored} == {0}
+    for idx, ev in enumerate(stored):
         assert ev.block_size == vbs
-        assert len(ev.token_ids) == vbs, (
-            f"FA token_ids should be {vbs}, got {len(ev.token_ids)}"
+        assert ev.token_ids == req.prompt_token_ids[idx * vbs : (idx + 1) * vbs]
+        expected_hash = maybe_convert_block_hash(
+            get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0))
         )
+        assert ev.block_hashes == [expected_hash]
 
-    # Mamba group: capture size (64) != emission size (32) → guard emits [].
-    mamba_events = by_group[1]
-    assert len(mamba_events) == 2
-    for ev in mamba_events:
-        assert ev.block_size == 32
-        assert ev.token_ids == [], (
-            f"Mamba token_ids should be [] (guard), got {ev.token_ids}"
-        )
-        assert ev.parent_block_hash is None, (
-            "Mamba parent_block_hash should be None (guard)"
-        )
-        assert ev.extra_keys is None, "Mamba extra_keys should be None (guard)"
+
+def test_scratch_group_store_emits_events_for_cacheable_groups_only() -> None:
+    """Event metadata is resolved only for prefix-cacheable groups.
+
+    A scratch group (GLM-5.3-Flash kpool tail) has a block size that does not
+    divide the hash block size, so viewing the request hashes at its block
+    size is impossible; the store must skip it and still emit BlockStored
+    for the attention blocks.
+    """
+    fix = make_events_scheduler(kv_cache_config=_make_scratch_kv_cache_config(16))
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    num_blocks = 2
+
+    req = make_request(num_blocks=num_blocks)
+    fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
+    scratch_block = gpu_pool.get_new_blocks(1)
+    kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block))
+    req.num_computed_tokens = num_blocks * BLOCK_SIZE
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+    sched_out = make_scheduler_output(
+        {req.request_id: num_blocks * BLOCK_SIZE},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    meta = sched.build_connector_meta(sched_out)
+    assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks)
+    simulate_store_completion(sched, meta.store_event)
+
+    events = list(sched.take_events())
+    assert len(events) == num_blocks
+    assert all(isinstance(e, BlockStored) for e in events)
+    assert {e.group_idx for e in events} == {0}
+    assert all(len(e.token_ids) == BLOCK_SIZE for e in events)
diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py
index db7e8b7aa51c..ed07718de60c 100644
--- a/tests/v1/simple_kv_offload/test_scheduler.py
+++ b/tests/v1/simple_kv_offload/test_scheduler.py
@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass
 
 import pytest
@@ -18,6 +19,7 @@
     SchedulerConfig,
     VllmConfig,
 )
+from vllm.config.kv_events import KVEventsConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (
     SimpleCPUOffloadConnector,
 )
@@ -31,6 +33,7 @@
 )
 from vllm.v1.core.sched.output import (
     CachedRequestData,
+    KVConnectorBlockState,
     NewRequestData,
     SchedulerOutput,
 )
@@ -38,10 +41,13 @@
     register_all_kvcache_specs,
 )
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
+    KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheTensor,
+    MambaSpec,
     SlidingWindowSpec,
 )
 from vllm.v1.outputs import KVConnectorOutput
@@ -121,6 +127,45 @@ def _make_kv_cache_config(
     )
 
 
+def _make_scratch_kv_cache_config(
+    num_blocks: int, scratch_block_size: int = 4
+) -> KVCacheConfig:
+    """FullAttention group plus a non-prefix-cacheable scratch group.
+
+    Mirrors GLM-5.3-Flash, whose kpool-tail group holds one per-request
+    scratch block of ``index_kpool`` tokens: it is excluded from prefix
+    caching, so its block size neither divides nor is divided by the
+    hash block size in general.
+    """
+    register_all_kvcache_specs(vllm_config=None)
+    fa_config = _make_kv_cache_config(num_blocks, num_groups=1)
+    scratch_layers = ["layer_scratch"]
+    scratch_spec = KpoolTailSpec(
+        block_size=scratch_block_size,
+        num_kv_heads=2,
+        head_size=HEAD_SIZE,
+        head_size_v=0,
+        dtype=DTYPE,
+        sliding_window=scratch_block_size,
+    )
+    assert not scratch_spec.prefix_cacheable
+    scratch_bytes = scratch_spec.page_size_bytes * num_blocks
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=fa_config.kv_cache_tensors
+        + [
+            KVCacheTensor(
+                size=scratch_bytes,
+                layers=scratch_layers,
+                layer_stride=scratch_bytes,
+                block_stride=scratch_spec.page_size_bytes,
+            )
+        ],
+        kv_cache_groups=fa_config.kv_cache_groups
+        + [KVCacheGroupSpec(scratch_layers, scratch_spec)],
+    )
+
+
 def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig:
     """Minimal VllmConfig for scheduler tests (no GPU)."""
     model_config = ModelConfig(
@@ -170,9 +215,13 @@ def make_scheduler(
     num_gpu_blocks: int = 16,
     num_groups: int = 1,
     lazy: bool = False,
+    kv_cache_config: KVCacheConfig | None = None,
 ) -> SchedulerFixture:
     """Build a SimpleCPUOffloadScheduler with small block pools."""
-    kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)
+    if kv_cache_config is None:
+        kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)
+    else:
+        num_groups = len(kv_cache_config.kv_cache_groups)
     vllm_config = _make_vllm_config()
     cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups
 
@@ -246,6 +295,7 @@ def make_scheduler_output(
         cached_req_new_blocks: For returning (cached) requests, maps
             req_id -> new_block_ids (incremental) or None.
             These are placed into ``scheduled_cached_reqs``.
+
     """
     scheduled_new_reqs: list[NewRequestData] = []
     if new_reqs:
@@ -1315,12 +1365,10 @@ def test_inflight_finish_deferred_cleanup() -> None:
 # Test 8: Null GPU blocks are skipped in store and load transfer pairs
 # ---------------------------------------------------------------------------
 def test_multi_group_null_blocks_skipped() -> None:
-    """Null GPU blocks (no block_hash) must not appear in store or load pairs.
+    """The null block id never appears in store or load pairs.
 
-    In eager store mode, _prepare_eager_store_specs skips blocks whose
-    block_hash is None (null blocks have no hash). We verify this by mixing
-    real hashed blocks with unhashed (null-like) blocks in a single group and
-    checking that only real blocks appear in the store list.
+    A null block-table slot whose hash is still cached on the GPU is stored
+    from the cached block instead; the null block itself is never copied.
     """
     fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False)
     sched = fix.scheduler
@@ -1352,9 +1400,11 @@ def test_multi_group_null_blocks_skipped() -> None:
         f"Null block id {null_block_id} should not appear in store transfer pairs"
     )
 
-    # Only real block should be scheduled for store
-    assert len(meta.store_gpu_blocks) == 1
+    # The nulled slot's hash is still cached on the GPU, so that block is
+    # recovered from the pool and stored alongside the real one.
+    assert len(meta.store_gpu_blocks) == 2
     assert gpu_blocks[0].block_id in meta.store_gpu_blocks
+    assert gpu_blocks[1].block_id in meta.store_gpu_blocks
 
     # Complete the store
     assert meta.store_event >= 0
@@ -1370,13 +1420,12 @@ def test_multi_group_null_blocks_skipped() -> None:
         block_hasher=req._block_hasher,
     )
     hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
-    # Only 1 block was stored (the real one)
-    assert hit_tokens == BLOCK_SIZE
+    assert hit_tokens == 2 * BLOCK_SIZE
     assert is_async is True
 
     # Allocate new GPU blocks for the load
-    gpu_blocks2 = gpu_pool.get_new_blocks(1)
-    kv_blocks2 = KVCacheBlocks(blocks=([gpu_blocks2[0], null_block],))
+    gpu_blocks2 = gpu_pool.get_new_blocks(2)
+    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
     sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
 
     sched_out2 = make_scheduler_output({req2.request_id: 1})
@@ -1388,6 +1437,63 @@ def test_multi_group_null_blocks_skipped() -> None:
     )
 
 
+# ---------------------------------------------------------------------------
+# Test 8b: Non-prefix-cacheable scratch groups take no part in store or load
+# ---------------------------------------------------------------------------
+def test_scratch_group_excluded_from_store_and_load() -> None:
+    """A scratch group (GLM-5.3-Flash kpool tail) is never stored or loaded.
+
+    Its single per-request block carries no hash, so the store must offload
+    only the attention blocks and the load must pair only attention blocks.
+    """
+    fix = make_scheduler(
+        num_cpu_blocks=8,
+        num_gpu_blocks=16,
+        kv_cache_config=_make_scratch_kv_cache_config(16),
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    num_blocks = 2
+
+    req = make_request(num_blocks=num_blocks)
+    fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
+    scratch_block = gpu_pool.get_new_blocks(1)
+    kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block))
+    req.num_computed_tokens = num_blocks * BLOCK_SIZE
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+    sched_out = make_scheduler_output(
+        {req.request_id: num_blocks * BLOCK_SIZE},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    meta = sched.build_connector_meta(sched_out)
+    assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks)
+    simulate_store_completion(sched, meta.store_event)
+
+    req2 = Request(
+        request_id="req-scratch-load",
+        prompt_token_ids=req.prompt_token_ids,
+        sampling_params=req.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=req._block_hasher,
+    )
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens == num_blocks * BLOCK_SIZE
+    assert is_async is True
+
+    fa_blocks2 = gpu_pool.get_new_blocks(num_blocks)
+    scratch_block2 = gpu_pool.get_new_blocks(1)
+    kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, scratch_block2))
+    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
+    sched_out2 = make_scheduler_output(
+        {req2.request_id: 1},
+        new_reqs={req2.request_id: kv_blocks2.get_block_ids()},
+    )
+    meta2 = sched.build_connector_meta(sched_out2)
+    assert sorted(meta2.load_gpu_blocks) == sorted(b.block_id for b in fa_blocks2)
+    assert len(meta2.load_cpu_blocks) == num_blocks
+
+
 # ---------------------------------------------------------------------------
 # Test 9: Chunked prefill accumulates block_ids across steps
 # ---------------------------------------------------------------------------
@@ -2128,3 +2234,834 @@ def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None:
             f"cp_world_size={cp_world_size}: target_cp={target_cp} should be "
             f"less than target_base={target_base}"
         )
+
+
+def _make_hybrid_attention_mamba_scheduler(
+    *,
+    num_cpu_blocks: int = 8,
+    num_gpu_blocks: int = 16,
+    attention_block_size: int = BLOCK_SIZE,
+    block_size: int = 4 * BLOCK_SIZE,
+    scheduler_block_size: int | None = None,
+    hash_block_size: int | None = None,
+    dcp_world_size: int = 4,
+    lazy: bool = False,
+    mamba_cache_mode: str = "align",
+    enable_kv_cache_events: bool = False,
+) -> SchedulerFixture:
+    """Build a scheduler for one attention group plus one Mamba group."""
+    scheduler_block_size = scheduler_block_size or block_size
+    hash_block_size = hash_block_size or block_size
+    attention_spec = FullAttentionSpec(
+        block_size=attention_block_size,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    mamba_spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode=mamba_cache_mode,
+    )
+    groups = [
+        KVCacheGroupSpec(["attention"], attention_spec),
+        KVCacheGroupSpec(["mamba"], mamba_spec),
+    ]
+    tensors = [
+        KVCacheTensor(
+            size=spec.page_size_bytes * num_gpu_blocks,
+            layers=group.layer_names,
+            layer_stride=spec.page_size_bytes * num_gpu_blocks,
+            block_stride=spec.page_size_bytes,
+        )
+        for group, spec in zip(groups, (attention_spec, mamba_spec))
+    ]
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_gpu_blocks,
+        kv_cache_tensors=tensors,
+        kv_cache_groups=groups,
+    )
+    vllm_config = _make_cp_vllm_config(dcp_world_size=dcp_world_size)
+    vllm_config.cache_config.prefix_cache_retention_interval = 0
+    vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
+    if enable_kv_cache_events:
+        vllm_config.kv_events_config = KVEventsConfig(enable_kv_cache_events=True)
+    # _derive_cpu_config() scales the requested capacity against
+    # kv_cache_tensors[0].size, so express it as a multiple of that tensor's
+    # per-block size to get exactly num_cpu_blocks.
+    cpu_capacity_bytes = tensors[0].size // num_gpu_blocks * num_cpu_blocks
+    sched = SimpleCPUOffloadScheduler(
+        vllm_config=vllm_config,
+        kv_cache_config=kv_cache_config,
+        cpu_capacity_bytes=cpu_capacity_bytes,
+        scheduler_block_size=scheduler_block_size,
+        hash_block_size=hash_block_size,
+        lazy_offload=lazy,
+    )
+    gpu_block_pool = BlockPool(
+        num_gpu_blocks=num_gpu_blocks,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+    sched.bind_gpu_block_pool(gpu_block_pool)
+    return SchedulerFixture(
+        scheduler=sched,
+        gpu_block_pool=gpu_block_pool,
+        vllm_config=vllm_config,
+        kv_cache_config=kv_cache_config,
+    )
+
+
+def test_hybrid_store_uses_resolved_group_block_sizes() -> None:
+    """Replicated groups must not be scaled by the DCP world size.
+
+    ``dcp_world_size_for_kv_cache_spec`` gives full attention the process DCP
+    size and every other spec 1. Reconstructing a group's geometry as
+    ``spec.block_size * cp_world_size`` over-scales the replicated groups, so
+    the eager store scan believes far fewer of their blocks are ready and
+    silently offloads only a fraction of them.
+    """
+    attention_block_size = BLOCK_SIZE
+    mamba_block_size = 4 * BLOCK_SIZE
+    dcp_world_size = 2
+    fix = _make_hybrid_attention_mamba_scheduler(
+        num_cpu_blocks=32,
+        num_gpu_blocks=32,
+        attention_block_size=attention_block_size,
+        block_size=mamba_block_size,
+        # Every prefix-cacheable group's resolved block size must be a multiple
+        # of the hash block, so it cannot exceed the smallest of them.
+        hash_block_size=BLOCK_SIZE,
+        dcp_world_size=dcp_world_size,
+        mamba_cache_mode="all",
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    attention_size, mamba_size = sched.group_block_sizes
+    assert attention_size == attention_block_size * dcp_world_size
+    assert mamba_size == mamba_block_size
+
+    confirmed = 2 * sched.block_size
+    req = _make_cp_request(num_blocks=8, virtual_block_size=BLOCK_SIZE)
+    attention_blocks = _allocate_cp_gpu_blocks(
+        gpu_pool, req, confirmed // attention_size, attention_size, group_id=0
+    )
+    mamba_blocks = _allocate_cp_gpu_blocks(
+        gpu_pool, req, confirmed // mamba_size, mamba_size, group_id=1
+    )
+    kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+    sched.build_connector_meta(
+        make_scheduler_output(
+            {req.request_id: confirmed},
+            new_reqs={req.request_id: kv_blocks.get_block_ids()},
+        )
+    )
+    req.num_computed_tokens = confirmed
+    meta = sched.build_connector_meta(
+        make_scheduler_output(
+            {req.request_id: 1},
+            cached_req_new_blocks={req.request_id: None},
+        )
+    )
+
+    # Both groups are positionally stable here, so every confirmed block of
+    # both must be offloaded. Over-scaling the replicated group halves its
+    # count.
+    assert sched._reqs_to_store[req.request_id].num_stored_blocks == [
+        confirmed // attention_size,
+        confirmed // mamba_size,
+    ]
+    assert set(meta.store_gpu_blocks) == {
+        *(b.block_id for b in attention_blocks),
+        *(b.block_id for b in mamba_blocks),
+    }
+
+
+def test_dcp_mixed_cache_loads_fine_grained_external_hit() -> None:
+    """Map a fine-grained external suffix after a nonzero local prefix."""
+    hash_block_size = BLOCK_SIZE
+    mamba_block_size = 4 * BLOCK_SIZE
+    local_tokens = 4 * BLOCK_SIZE
+    # Span two Mamba blocks and end partially in the second one. This exercises
+    # both cdiv(external_tokens, mamba_block_size) > 1 and the partial tail.
+    external_tokens = 6 * BLOCK_SIZE
+    total_cached_tokens = local_tokens + external_tokens
+    fix = _make_hybrid_attention_mamba_scheduler(
+        num_cpu_blocks=32,
+        num_gpu_blocks=32,
+        block_size=mamba_block_size,
+        hash_block_size=hash_block_size,
+        dcp_world_size=2,
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+
+    assert external_tokens % mamba_block_size != 0
+    assert sched.cpu_coordinator.enable_partial_hash_hits
+
+    producer = _make_cp_request(num_blocks=10, virtual_block_size=hash_block_size)
+    attention_blocks = _allocate_cp_gpu_blocks(
+        gpu_pool,
+        producer,
+        num_blocks=5,
+        virtual_block_size=2 * BLOCK_SIZE,
+        group_id=0,
+    )
+    mamba_blocks = gpu_pool.get_new_blocks(3)
+    gpu_pool.cache_full_blocks(
+        request=producer,
+        blocks=mamba_blocks,
+        num_cached_blocks=0,
+        num_full_blocks=2,
+        block_size=mamba_block_size,
+        kv_cache_group_id=1,
+    )
+    gpu_pool.cache_partial_block(
+        request=producer,
+        block=mamba_blocks[2],
+        num_tokens=total_cached_tokens,
+        kv_cache_group_id=1,
+        block_size=mamba_block_size,
+    )
+    producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))
+    sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0)
+
+    # Positional scanning stores append-only FA blocks. Mamba align blocks use
+    # exact handoffs for every retained boundary, including the partial tail.
+    store_output = make_scheduler_output(
+        {producer.request_id: total_cached_tokens},
+        new_reqs={producer.request_id: producer_blocks.get_block_ids()},
+    )
+    store_output.kv_connector_block_state = KVConnectorBlockState(
+        req_ids=set(),
+        resolve_block_ids=lambda _: (),
+        boundary_state_offloads={
+            producer.request_id: [
+                (1, mamba_blocks[0].block_id, mamba_block_size),
+                (1, mamba_blocks[1].block_id, 2 * mamba_block_size),
+                (1, mamba_blocks[2].block_id, total_cached_tokens),
+            ]
+        },
+    )
+    store_meta = sched.build_connector_meta(store_output)
+    # The exact handoffs land in the step they arrive; the append-only FA scan
+    # picks its blocks up once the scheduler has committed the tokens.
+    assert set(store_meta.store_gpu_blocks) == {
+        block.block_id for block in mamba_blocks
+    }
+    simulate_store_completion(sched, store_meta.store_event)
+    producer.num_computed_tokens = total_cached_tokens
+    fa_meta = sched.build_connector_meta(
+        make_scheduler_output(
+            {producer.request_id: 1},
+            cached_req_new_blocks={producer.request_id: None},
+        )
+    )
+    assert set(fa_meta.store_gpu_blocks) == {
+        block.block_id for block in attention_blocks
+    }
+    simulate_store_completion(sched, fa_meta.store_event)
+
+    consumer = Request(
+        request_id="req-dcp-mixed-fine-load",
+        prompt_token_ids=producer.prompt_token_ids,
+        sampling_params=producer.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=producer._block_hasher,
+    )
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(
+        consumer, num_computed_tokens=local_tokens
+    )
+    assert hit_tokens == external_tokens
+    assert is_async is True
+    pending_blocks, pending_tokens, _ = sched._pending_cpu_hits[consumer.request_id]
+    assert pending_tokens == external_tokens
+    assert len(pending_blocks[1]) == 2
+    assert pending_blocks[1][0].is_null
+    assert not pending_blocks[1][1].is_null
+
+    local_attention = attention_blocks[:2]
+    local_mamba = mamba_blocks[:1]
+    load_blocks = KVCacheBlocks(
+        blocks=(
+            local_attention + gpu_pool.get_new_blocks(3),
+            local_mamba + gpu_pool.get_new_blocks(2),
+        )
+    )
+    sched.update_state_after_alloc(
+        consumer,
+        load_blocks,
+        num_external_tokens=hit_tokens,
+    )
+    load_meta = sched.build_connector_meta(
+        make_scheduler_output(
+            {consumer.request_id: 1},
+            new_reqs={consumer.request_id: load_blocks.get_block_ids()},
+        )
+    )
+    assert load_meta.load_gpu_blocks == [
+        *(block.block_id for block in load_blocks.blocks[0][2:5]),
+        load_blocks.blocks[1][2].block_id,
+    ]
+    assert len(load_meta.load_cpu_blocks) == 4
+
+
+def test_finished_eager_store_caches_fa_partial_tail_for_hybrid_hit() -> None:
+    """A finished non-LCM prompt needs the FA partial tail in CPU too."""
+    hash_block_size = BLOCK_SIZE
+    mamba_block_size = BLOCK_SIZE
+    attention_block_size = 4 * BLOCK_SIZE
+    dcp_world_size = 2
+    scheduler_block_size = attention_block_size * dcp_world_size
+    total_cached_tokens = 5 * hash_block_size
+    fix = _make_hybrid_attention_mamba_scheduler(
+        num_cpu_blocks=32,
+        num_gpu_blocks=32,
+        attention_block_size=attention_block_size,
+        block_size=mamba_block_size,
+        scheduler_block_size=scheduler_block_size,
+        hash_block_size=hash_block_size,
+        dcp_world_size=dcp_world_size,
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    assert sched.fa_block_size == scheduler_block_size
+    assert total_cached_tokens % sched.fa_block_size != 0
+
+    producer = _make_cp_request(num_blocks=5, virtual_block_size=hash_block_size)
+    attention_blocks = gpu_pool.get_new_blocks(1)
+    gpu_pool.cache_partial_block(
+        request=producer,
+        block=attention_blocks[0],
+        num_tokens=total_cached_tokens,
+        kv_cache_group_id=0,
+        block_size=scheduler_block_size,
+    )
+    mamba_blocks = gpu_pool.get_new_blocks(5)
+    gpu_pool.cache_full_blocks(
+        request=producer,
+        blocks=mamba_blocks,
+        num_cached_blocks=0,
+        num_full_blocks=5,
+        kv_cache_group_id=1,
+        block_size=mamba_block_size,
+    )
+    producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))
+    sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0)
+
+    store_output = make_scheduler_output(
+        {producer.request_id: total_cached_tokens},
+        new_reqs={producer.request_id: producer_blocks.get_block_ids()},
+    )
+    store_output.kv_connector_block_state = KVConnectorBlockState(
+        req_ids=set(),
+        resolve_block_ids=lambda _: (),
+        boundary_state_offloads={
+            producer.request_id: [
+                (1, mamba_blocks[4].block_id, total_cached_tokens),
+            ]
+        },
+    )
+    store_meta = sched.build_connector_meta(store_output)
+    simulate_store_completion(sched, store_meta.store_event)
+
+    producer.num_computed_tokens = total_cached_tokens
+    sched.request_finished_all_groups(producer, producer_blocks.get_block_ids())
+    finish_meta = sched.build_connector_meta(make_scheduler_output({}))
+    assert attention_blocks[0].block_id in finish_meta.store_gpu_blocks
+    simulate_store_completion(sched, finish_meta.store_event)
+
+    consumer = Request(
+        request_id="req-finished-partial-tail-load",
+        prompt_token_ids=producer.prompt_token_ids,
+        sampling_params=producer.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=producer._block_hasher,
+    )
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(
+        consumer, num_computed_tokens=0
+    )
+    assert hit_tokens == total_cached_tokens
+    assert is_async is True
+
+
+def test_finished_eager_store_does_not_duplicate_completed_partial_tail() -> None:
+    """Decode can fill the boundary block; the scan then already covers it."""
+    hash_block_size = BLOCK_SIZE
+    mamba_block_size = BLOCK_SIZE
+    attention_block_size = 4 * BLOCK_SIZE
+    dcp_world_size = 2
+    scheduler_block_size = attention_block_size * dcp_world_size
+    prompt_tokens = 5 * hash_block_size
+    fix = _make_hybrid_attention_mamba_scheduler(
+        num_cpu_blocks=32,
+        num_gpu_blocks=32,
+        attention_block_size=attention_block_size,
+        block_size=mamba_block_size,
+        scheduler_block_size=scheduler_block_size,
+        hash_block_size=hash_block_size,
+        dcp_world_size=dcp_world_size,
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    assert prompt_tokens % sched.fa_block_size != 0
+
+    producer = _make_cp_request(num_blocks=5, virtual_block_size=hash_block_size)
+    attention_blocks = gpu_pool.get_new_blocks(1)
+    gpu_pool.cache_partial_block(
+        request=producer,
+        block=attention_blocks[0],
+        num_tokens=prompt_tokens,
+        kv_cache_group_id=0,
+        block_size=scheduler_block_size,
+    )
+    mamba_blocks = gpu_pool.get_new_blocks(5)
+    gpu_pool.cache_full_blocks(
+        request=producer,
+        blocks=mamba_blocks,
+        num_cached_blocks=0,
+        num_full_blocks=5,
+        kv_cache_group_id=1,
+        block_size=mamba_block_size,
+    )
+    producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))
+    sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0)
+
+    # Decode carries the request past the physical attention block boundary, so
+    # the positional scan now considers the same block the partial tail names.
+    producer.num_computed_tokens = sched.fa_block_size
+    sched.request_finished_all_groups(producer, producer_blocks.get_block_ids())
+    finish_meta = sched.build_connector_meta(make_scheduler_output({}))
+
+    stored = finish_meta.store_gpu_blocks
+    assert attention_blocks[0].block_id in stored
+    assert len(stored) == len(set(stored))
+    assert len(finish_meta.store_cpu_blocks) == len(stored)
+
+
+def test_external_lookup_rejects_misaligned_local_prefix(caplog_vllm) -> None:
+    scheduler_block_size = 4 * BLOCK_SIZE
+    fix = _make_hybrid_attention_mamba_scheduler(block_size=scheduler_block_size)
+
+    request = _make_cp_request(num_blocks=2, virtual_block_size=scheduler_block_size)
+    # The ``vllm`` logger sets propagate=False, so plain ``caplog`` can observe
+    # nothing; ``caplog_vllm`` re-enables propagation for the duration.
+    with caplog_vllm.at_level(logging.WARNING):
+        hit_tokens, is_async = fix.scheduler.get_num_new_matched_tokens(
+            request, num_computed_tokens=BLOCK_SIZE
+        )
+
+    assert (hit_tokens, is_async) == (0, False)
+    assert "requires scheduler-block-aligned local tokens" in caplog_vllm.text
+
+
+def test_fine_grained_external_hits_require_eager_offload() -> None:
+    hash_block_size = BLOCK_SIZE
+    scheduler_block_size = 4 * BLOCK_SIZE
+
+    eager = _make_hybrid_attention_mamba_scheduler(
+        block_size=scheduler_block_size,
+        hash_block_size=hash_block_size,
+    )
+    lazy = _make_hybrid_attention_mamba_scheduler(
+        block_size=scheduler_block_size,
+        hash_block_size=hash_block_size,
+        lazy=True,
+    )
+
+    assert eager.scheduler.cpu_coordinator.enable_partial_hash_hits
+    assert not lazy.scheduler.cpu_coordinator.enable_partial_hash_hits
+
+
+def test_eager_store_does_not_scan_mamba_blocks_positionally() -> None:
+    """Mamba stores must use exact handoffs, not historical block positions."""
+    block_size = 4 * BLOCK_SIZE
+    fix = _make_hybrid_attention_mamba_scheduler(
+        num_cpu_blocks=16, num_gpu_blocks=24, block_size=block_size
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    attention_manager, mamba_manager = sched.cpu_coordinator.single_type_managers
+    assert attention_manager.has_positionally_stable_blocks
+    assert not mamba_manager.has_positionally_stable_blocks
+
+    req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)
+    attn_blocks = _allocate_cp_gpu_blocks(
+        gpu_pool, req, 2, virtual_block_size=block_size, group_id=0
+    )
+    gpu_blocks = gpu_pool.get_new_blocks(2)
+    gpu_blocks[1].set_block_hash(make_block_hash_with_group_id(req.block_hashes[1], 1))
+    kv_blocks = KVCacheBlocks(blocks=(attn_blocks, gpu_blocks))
+    req.num_computed_tokens = 0
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    output = make_scheduler_output(
+        {req.request_id: 2 * block_size},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    output.kv_connector_block_state = KVConnectorBlockState(
+        req_ids=set(),
+        resolve_block_ids=lambda _: (),
+        boundary_state_offloads={
+            req.request_id: [(1, gpu_blocks[1].block_id, 2 * block_size)]
+        },
+    )
+    meta1 = sched.build_connector_meta(output)
+    assert meta1.store_event >= 0
+    # Only the exact handoff lands this step. build_connector_meta runs before
+    # the scheduler commits this step's tokens, so the positional scan has
+    # nothing confirmed yet.
+    assert set(meta1.store_gpu_blocks) == {gpu_blocks[1].block_id}
+    assert sched._reqs_to_store[req.request_id].num_stored_blocks == [0, 0]
+    simulate_store_completion(sched, meta1.store_event)
+
+    # Next step the attention blocks are confirmed and scanned positionally. A
+    # later hash update still must not make the mamba align snapshot a valid
+    # source; the explicit boundary-handoff tests cover the supported path.
+    gpu_blocks[0].set_block_hash(make_block_hash_with_group_id(req.block_hashes[0], 1))
+    req.num_computed_tokens = 2 * block_size
+    meta2 = sched.build_connector_meta(
+        make_scheduler_output(
+            {req.request_id: 1},
+            cached_req_new_blocks={req.request_id: None},
+        )
+    )
+    assert set(meta2.store_gpu_blocks) == {block.block_id for block in attn_blocks}
+    assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0]
+    simulate_store_completion(sched, meta2.store_event)
+
+    # Nothing further: mamba align is never picked up positionally.
+    meta3 = sched.build_connector_meta(
+        make_scheduler_output(
+            {req.request_id: 1},
+            cached_req_new_blocks={req.request_id: None},
+        )
+    )
+    assert meta3.store_event == -1
+    assert meta3.store_gpu_blocks == []
+    assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0]
+
+
+def test_eager_store_scans_mamba_all_blocks_positionally() -> None:
+    """Mamba all-mode retains positional identity and needs no handoff."""
+    block_size = 4 * BLOCK_SIZE
+    fix = _make_hybrid_attention_mamba_scheduler(
+        num_cpu_blocks=16,
+        num_gpu_blocks=24,
+        block_size=block_size,
+        mamba_cache_mode="all",
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    attention_manager, mamba_manager = sched.cpu_coordinator.single_type_managers
+    assert attention_manager.has_positionally_stable_blocks
+    assert mamba_manager.has_positionally_stable_blocks
+
+    req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)
+    attention_blocks = _allocate_cp_gpu_blocks(
+        gpu_pool, req, 2, virtual_block_size=block_size, group_id=0
+    )
+    mamba_blocks = _allocate_cp_gpu_blocks(
+        gpu_pool, req, 2, virtual_block_size=block_size, group_id=1
+    )
+    kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    # build_connector_meta runs before the scheduler commits this step's
+    # tokens, so the positional scan picks the blocks up on the next step.
+    first = sched.build_connector_meta(
+        make_scheduler_output(
+            {req.request_id: 2 * block_size},
+            new_reqs={req.request_id: kv_blocks.get_block_ids()},
+        )
+    )
+    assert first.store_gpu_blocks == []
+    req.num_computed_tokens = 2 * block_size
+    meta = sched.build_connector_meta(
+        make_scheduler_output(
+            {req.request_id: 1},
+            cached_req_new_blocks={req.request_id: None},
+        )
+    )
+
+    assert set(meta.store_gpu_blocks) == {
+        *(block.block_id for block in attention_blocks),
+        *(block.block_id for block in mamba_blocks),
+    }
+    assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 2]
+
+
+@pytest.mark.parametrize("gone_via", ["preempted", "finished", "unregistered"])
+def test_boundary_handoff_dropped_for_departing_request(gone_via: str) -> None:
+    """Handoffs for a request leaving this step must not be read.
+
+    The scheduler drains boundary handoffs without filtering by request
+    liveness, so a preempted or finished request can still offer block IDs.
+    Those blocks are being released and may already back another request, so
+    copying them would publish unrelated KV under a valid hash.
+    """
+    block_size = 4 * BLOCK_SIZE
+    fix = _make_hybrid_attention_mamba_scheduler(
+        num_cpu_blocks=16, num_gpu_blocks=24, block_size=block_size
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+
+    req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)
+    attn_blocks = _allocate_cp_gpu_blocks(
+        gpu_pool, req, 2, virtual_block_size=block_size, group_id=0
+    )
+    mamba_blocks = gpu_pool.get_new_blocks(2)
+    mamba_blocks[1].set_block_hash(
+        make_block_hash_with_group_id(req.block_hashes[1], 1)
+    )
+    kv_blocks = KVCacheBlocks(blocks=(attn_blocks, mamba_blocks))
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    output = make_scheduler_output(
+        {req.request_id: 2 * block_size},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    output.kv_connector_block_state = KVConnectorBlockState(
+        req_ids=set(),
+        resolve_block_ids=lambda _: (),
+        boundary_state_offloads={
+            req.request_id: [(1, mamba_blocks[1].block_id, 2 * block_size)]
+        },
+    )
+    if gone_via == "preempted":
+        output.preempted_req_ids = {req.request_id}
+    elif gone_via == "finished":
+        output.finished_req_ids = {req.request_id}
+    else:
+        sched._reqs_to_store.pop(req.request_id)
+
+    meta = sched.build_connector_meta(output)
+
+    # The exact-handoff block is never a store source for a departing request.
+    assert mamba_blocks[1].block_id not in meta.store_gpu_blocks
+    stats = sched.get_boundary_store_stats()
+    assert stats.dropped_request_gone == 1
+    assert stats.stored == 0
+
+
+def test_boundary_handoff_keeps_block_meta_index_parallel() -> None:
+    """Boundary stores must extend block_meta alongside the block ids.
+
+    ``_process_store_completion`` indexes ``block_meta`` by position, so a
+    boundary handoff that appends only block ids shortens the list, trips its
+    length assertion and shifts every following entry's metadata.
+    """
+    block_size = 4 * BLOCK_SIZE
+    fix = _make_hybrid_attention_mamba_scheduler(
+        num_cpu_blocks=16,
+        num_gpu_blocks=24,
+        block_size=block_size,
+        enable_kv_cache_events=True,
+    )
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    assert sched.enable_kv_cache_events
+
+    req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)
+    attn_blocks = _allocate_cp_gpu_blocks(
+        gpu_pool, req, 2, virtual_block_size=block_size, group_id=0
+    )
+    mamba_blocks = gpu_pool.get_new_blocks(2)
+    mamba_blocks[1].set_block_hash(
+        make_block_hash_with_group_id(req.block_hashes[1], 1)
+    )
+    kv_blocks = KVCacheBlocks(blocks=(attn_blocks, mamba_blocks))
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+    # Make the FA blocks positionally eligible in the same event as the
+    # explicit Mamba boundary handoff.
+    req.num_computed_tokens = 2 * block_size
+
+    output = make_scheduler_output(
+        {req.request_id: 2 * block_size},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    output.kv_connector_block_state = KVConnectorBlockState(
+        req_ids=set(),
+        resolve_block_ids=lambda _: (),
+        boundary_state_offloads={
+            req.request_id: [(1, mamba_blocks[1].block_id, 2 * block_size)]
+        },
+    )
+
+    gpu_ids, cpu_ids, req_ids, block_meta = sched.prepare_store_specs(output)
+    assert mamba_blocks[1].block_id in gpu_ids
+    assert set(gpu_ids) == {
+        *(block.block_id for block in attn_blocks),
+        mamba_blocks[1].block_id,
+    }
+    assert req_ids == [req.request_id]
+    assert block_meta is not None
+    assert len(block_meta) == len(gpu_ids) == len(cpu_ids)
+
+
+def test_finished_eager_store_recovers_nulled_window_tail() -> None:
+    """A sliding-window group nulls pages that left the window before the
+    connector sees the finished request's block table, while the GPU keeps
+    them hashed in its free queue. The eager store must copy those pages, or
+    the CPU lookup (which needs the whole window) misses every time."""
+    fix = make_scheduler(num_cpu_blocks=16, num_gpu_blocks=32, num_groups=2, lazy=False)
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    null_block = gpu_pool.null_block
+
+    num_blocks = 4  # == sliding window of the test SWA group
+    req = make_request(num_blocks=num_blocks)
+    fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
+    swa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=1)
+    # remove_skipped_blocks() ran: the two oldest window pages are freed (still
+    # hashed) and their slots hold the null block.
+    gpu_pool.free_blocks(swa_blocks[:2])
+    swa_table = [null_block, null_block, *swa_blocks[2:]]
+    kv_blocks = KVCacheBlocks(blocks=(fa_blocks, swa_table))
+    req.num_computed_tokens = num_blocks * BLOCK_SIZE
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    block_ids = kv_blocks.get_block_ids()
+    sched.request_finished_all_groups(req, block_ids)
+    meta = sched.build_connector_meta(make_scheduler_output({}))
+
+    assert null_block.block_id not in meta.store_gpu_blocks
+    assert set(meta.store_gpu_blocks) == {b.block_id for b in (*fa_blocks, *swa_blocks)}
+    assert meta.store_event >= 0
+    simulate_store_completion(sched, meta.store_event)
+
+    req2 = Request(
+        request_id="req-window-tail-load",
+        prompt_token_ids=req.prompt_token_ids,
+        sampling_params=req.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=req._block_hasher,
+    )
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens == num_blocks * BLOCK_SIZE
+    assert is_async is True
+
+
+QSA_RING_BLOCK_SIZE = 4
+
+
+def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig:
+    register_all_kvcache_specs(vllm_config=None)
+    full_spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    ring_spec = CircularBufferSpec(
+        block_size=QSA_RING_BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        head_size_v=0,
+        dtype=DTYPE,
+    )
+    groups = [
+        KVCacheGroupSpec(["layer_0"], full_spec),
+        KVCacheGroupSpec(["qsa_ring"], ring_spec),
+    ]
+    tensors = [
+        KVCacheTensor(
+            size=group.kv_cache_spec.page_size_bytes * num_blocks,
+            layers=list(group.layer_names),
+            layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks,
+            block_stride=group.kv_cache_spec.page_size_bytes,
+        )
+        for group in groups
+    ]
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=tensors,
+        kv_cache_groups=groups,
+    )
+
+
+def _make_qsa_scheduler(
+    lazy: bool = False,
+) -> tuple[SimpleCPUOffloadScheduler, BlockPool]:
+    kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16)
+    sched = SimpleCPUOffloadScheduler(
+        vllm_config=_make_vllm_config(),
+        kv_cache_config=kv_cache_config,
+        cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2,
+        scheduler_block_size=BLOCK_SIZE,
+        hash_block_size=BLOCK_SIZE,
+        lazy_offload=lazy,
+    )
+    gpu_pool = BlockPool(
+        num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE
+    )
+    sched.bind_gpu_block_pool(gpu_pool)
+    return sched, gpu_pool
+
+
+def test_qsa_ring_group_is_never_stored_or_loaded() -> None:
+    sched, gpu_pool = _make_qsa_scheduler()
+    assert sched.prefix_cacheable_group_ids == (0,)
+
+    num_blocks = 2
+    req = make_request(num_blocks=num_blocks)
+    fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
+    ring_block = gpu_pool.get_new_blocks(1)[0]
+    kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block]))
+    req.num_computed_tokens = num_blocks * BLOCK_SIZE
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    sched_out = make_scheduler_output(
+        {req.request_id: num_blocks * BLOCK_SIZE},
+        new_reqs={req.request_id: kv_blocks.get_block_ids()},
+    )
+    meta = sched.build_connector_meta(sched_out)
+    assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks}
+    assert ring_block.block_id not in meta.store_gpu_blocks
+    assert meta.store_event >= 0
+    simulate_store_completion(sched, meta.store_event)
+
+    req2 = Request(
+        request_id="req-qsa-load",
+        prompt_token_ids=req.prompt_token_ids,
+        sampling_params=req.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=req._block_hasher,
+    )
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens == num_blocks * BLOCK_SIZE
+    assert is_async is True
+
+    fa_blocks2 = gpu_pool.get_new_blocks(num_blocks)
+    ring_block2 = gpu_pool.get_new_blocks(1)[0]
+    kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2]))
+    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
+
+    meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1}))
+    assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2}
+    assert ring_block2.block_id not in meta2.load_gpu_blocks
+    assert len(meta2.load_cpu_blocks) == num_blocks
+
+
+def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None:
+    with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16)
+    without_ring = KVCacheConfig(
+        num_blocks=16,
+        kv_cache_tensors=with_ring.kv_cache_tensors[:1],
+        kv_cache_groups=with_ring.kv_cache_groups[:1],
+    )
+    max_batched = 64
+    target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
+        with_ring, max_batched
+    )
+    target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
+        without_ring, max_batched
+    )
+    assert target_with == target_without
diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py
index 66582f03bbd3..173929005151 100644
--- a/tests/v1/simple_kv_offload/test_worker.py
+++ b/tests/v1/simple_kv_offload/test_worker.py
@@ -10,6 +10,7 @@
 from __future__ import annotations
 
 import time
+from contextlib import nullcontext
 from unittest.mock import MagicMock
 
 import pytest
@@ -21,7 +22,13 @@
     pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
 
 from tests.v1.attention.utils import dense_kv_cache_tensor, dense_kv_cache_views
+from tests.v1.kv_connector.unit.test_kv_connector_lifecycle import (
+    _make_empty_scheduler_output,
+)
 from vllm.config import CacheConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (
+    SimpleCPUOffloadConnector,
+)
 from vllm.v1.core.kv_cache_utils import (
     get_kv_cache_config_from_groups,
     is_kv_cache_spec_uniform,
@@ -41,10 +48,12 @@
     CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
     CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
     build_params,
-    pin_tensor,
 )
+from vllm.v1.simple_kv_offload.disk_backend import DiskBackend
 from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata
 from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker
+from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
 from vllm.v1.worker.utils import allocate_kv_cache
 
 NUM_BLOCKS = 64
@@ -57,8 +66,11 @@
 
 def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:
     gpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cuda")}
-    cpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu")}
-    pin_tensor(cpu["k"])
+    cpu = {
+        "k": torch.zeros(
+            (NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu", pin_memory=True
+        )
+    }
     low_pri, _ = torch.cuda.Stream.priority_range()
     backend = DmaCopyBackend()
     backend.init(
@@ -71,6 +83,92 @@ def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:
     return backend, gpu["k"], cpu["k"]
 
 
+@pytest.mark.parametrize("use_v2", [False, True])
+def test_no_forward_step_completes_cpu_store(monkeypatch, use_v2):
+    """A finished request's final store must drain without another model step."""
+    backend, gpu, cpu = _make_backend()
+    worker = SimpleCPUOffloadWorker(None, None, cpu_capacity_bytes=0)
+    worker._backend = backend
+    connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector)
+    connector.worker_handler = worker
+    output = _make_empty_scheduler_output()
+    output.kv_connector_metadata = SimpleCPUOffloadMetadata(
+        store_event=0, store_gpu_blocks=[0], store_cpu_blocks=[0]
+    )
+    module = "vllm.v1.worker.kv_connector_model_runner_mixin"
+    monkeypatch.setattr(f"{module}.get_kv_transfer_group", lambda: connector)
+    monkeypatch.setattr(f"{module}.set_forward_context", lambda *a: nullcontext())
+    monkeypatch.setattr(f"{module}.get_forward_context", lambda: None)
+    try:
+        gpu.fill_(91)
+        if use_v2:
+            active = ActiveKVConnector.__new__(ActiveKVConnector)
+            active.kv_connector = connector
+            active._disabled = False
+            active._pending_load_start = False
+            module = "vllm.v1.worker.gpu.kv_connector"
+            monkeypatch.setattr(f"{module}.is_forward_context_available", lambda: True)
+            monkeypatch.setattr(f"{module}.get_forward_context", lambda: None)
+            result = active.no_forward(output)
+        else:
+            result = KVConnectorModelRunnerMixin.kv_connector_no_forward(output, None)
+        completion = (
+            result.kv_connector_output.kv_connector_worker_meta
+            if result.kv_connector_output is not None
+            else None
+        )
+        deadline = time.monotonic() + 5
+        while completion is None and time.monotonic() < deadline:
+            connector.get_finished(set())
+            completion = connector.build_connector_worker_meta()
+            time.sleep(0.001)
+        assert completion is not None, "CPU store never completed on the empty step"
+        assert completion.completed_store_events == {0: 1}
+        assert torch.equal(cpu[0], gpu[0].cpu())
+    finally:
+        backend.shutdown()
+
+
+def test_deferred_store_waits_for_draft_forward(monkeypatch):
+    """Target completion must not submit the store before draft finalization."""
+    backend, gpu, cpu = _make_backend()
+    worker = SimpleCPUOffloadWorker(None, None, cpu_capacity_bytes=0)
+    worker._backend = backend
+    connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector)
+    connector.worker_handler = worker
+    output = _make_empty_scheduler_output()
+    output.kv_connector_metadata = SimpleCPUOffloadMetadata(
+        store_event=0, store_gpu_blocks=[0], store_cpu_blocks=[0]
+    )
+    module = "vllm.v1.worker.kv_connector_model_runner_mixin"
+    monkeypatch.setattr(f"{module}.get_kv_transfer_group", lambda: connector)
+    monkeypatch.setattr(f"{module}.has_kv_transfer_group", lambda: True)
+    monkeypatch.setattr(f"{module}.get_forward_context", lambda: None)
+    launch = MagicMock(wraps=backend.launch_copy)
+    monkeypatch.setattr(backend, "launch_copy", launch)
+    try:
+        with KVConnectorModelRunnerMixin._get_kv_connector_output(
+            output, defer_finalize=True
+        ):
+            gpu.fill_(17)
+        launch.assert_not_called()
+        gpu.fill_(91)
+        KVConnectorModelRunnerMixin.finalize_kv_connector()
+        assert launch.call_count == 1
+        completion = None
+        deadline = time.monotonic() + 5
+        while completion is None and time.monotonic() < deadline:
+            connector.get_finished(set())
+            completion = connector.build_connector_worker_meta()
+            time.sleep(0.001)
+        assert completion is not None
+        assert completion.completed_store_events == {0: 1}
+        assert torch.equal(cpu[0], gpu[0].cpu())
+        assert launch.call_count == 1
+    finally:
+        backend.shutdown()
+
+
 def _drive_store(
     backend: DmaCopyBackend,
     gpu: torch.Tensor,
@@ -399,6 +497,8 @@ def test_register_mixed_page_sizes_in_one_cache_group(monkeypatch):
     vllm_config.cache_config = CacheConfig()
     vllm_config.cache_config.kv_cache_layout = layout.name
     vllm_config.cache_config.num_gpu_blocks_override = None
+    vllm_config.attention_config.hisparse_config = None
+    vllm_config.kv_transfer_config = None
 
     pages = [spec.page_size_bytes for spec in specs.values()]
     num_blocks = 4
@@ -457,6 +557,8 @@ def test_register_mixed_page_sizes_odd_block_counts(monkeypatch, rank_blocks):
     vllm_config.cache_config = CacheConfig()
     vllm_config.cache_config.kv_cache_layout = layout.name
     vllm_config.cache_config.num_gpu_blocks_override = None
+    vllm_config.attention_config.hisparse_config = None
+    vllm_config.kv_transfer_config = None
 
     pages = [spec.page_size_bytes for spec in specs.values()]
     kv_cache_config = get_kv_cache_config_from_groups(
@@ -506,6 +608,8 @@ def test_mixed_page_byte_placement_is_dcp_invariant():
         vllm_config.cache_config.kv_cache_layout = KVCacheLayout.LBNHC.name
         vllm_config.cache_config.num_gpu_blocks_override = None
         vllm_config.parallel_config.decode_context_parallel_size = dcp_size
+        vllm_config.attention_config.hisparse_config = None
+        vllm_config.kv_transfer_config = None
 
         page_bytes = sum(spec.page_size_bytes for spec in specs.values())
         config = get_kv_cache_config_from_groups(vllm_config, [group], page_bytes * 4)
@@ -519,3 +623,59 @@ def test_mixed_page_byte_placement_is_dcp_invariant():
         )
 
     assert placements[0] == placements[1]
+
+
+@pytest.mark.parametrize(
+    ("block_bytes", "use_page_cache", "expect_ok"),
+    [
+        (4096, True, True),  # aligned control
+        (3 * 512, True, True),  # 512-aligned only: fine for buffered I/O
+        (3 * 512, False, False),  # O_DIRECT keeps enforcing 4096 alignment
+    ],
+)
+def test_disk_backend_alignment_only_required_for_direct_io(
+    tmp_path, monkeypatch, block_bytes: int, use_page_cache: bool, expect_ok: bool
+):
+    """The 4096 stride assert belongs to O_DIRECT; page-cache I/O must not hit it.
+
+    Hybrid models (e.g. 101 per-rank cache tensors) sum to a block stride that
+    is 512- but not 4096-aligned, which previously made disk offload unbootable
+    even with use_page_cache=True.
+    """
+    # Host pinning is a separate lifecycle concern with its own tests; the
+    # alignment gate must not depend on registration state, so stub both ends.
+    monkeypatch.setattr(
+        "vllm.v1.simple_kv_offload.disk_backend.pin_tensor", lambda t: None
+    )
+    monkeypatch.setattr(
+        "vllm.v1.simple_kv_offload.disk_backend.unpin_tensor",
+        lambda t: None,
+        raising=False,  # harmless on branches without teardown-side unpins
+    )
+    num_blocks = 4
+    gpu = {"k": torch.zeros((num_blocks, block_bytes), dtype=torch.int8, device="cuda")}
+    backend = DiskBackend()
+    disk_path = tmp_path / "kv-offload.bin"
+
+    def _init() -> None:
+        backend.init(
+            gpu,
+            gpu["k"].device,
+            torch.cuda.Stream(),
+            torch.cuda.Stream(),
+            str(disk_path),
+            num_blocks,
+            block_bytes,
+            use_page_cache=use_page_cache,
+        )
+
+    if not expect_ok:
+        with pytest.raises(AssertionError, match="not aligned"):
+            _init()
+        return
+
+    _init()
+    try:
+        assert disk_path.stat().st_size == num_blocks * block_bytes
+    finally:
+        backend.shutdown()
diff --git a/tests/v1/spec_decode/test_acceptance_estimator.py b/tests/v1/spec_decode/test_acceptance_estimator.py
new file mode 100644
index 000000000000..a1f3da09402f
--- /dev/null
+++ b/tests/v1/spec_decode/test_acceptance_estimator.py
@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``predict`` writes one probability into two index spaces at once.
+
+Adaptive verification reads the estimate in this step's batch order, while the
+next step's refit reads it in persistent request-state slots. A mix-up between
+the two is silent -- every value stays a plausible probability, it just belongs
+to another request -- so it is asserted here rather than left to acceptance
+metrics to hint at.
+"""
+
+import pytest
+import torch
+
+from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
+    OnlineAcceptanceEstimator,
+)
+
+pytest.importorskip("triton")
+if not torch.cuda.is_available():
+    pytest.skip("CUDA required for acceptance estimator tests", allow_module_level=True)
+
+MAX_NUM_REQS = 16
+NUM_STEPS = 4
+VOCAB_SIZE = 2048
+
+
+def _make_estimator(device: torch.device) -> OnlineAcceptanceEstimator:
+    estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device)
+    # Distinct per-position intercepts, so reading the wrong one shows up.
+    estimator.slope.fill_(0.7)
+    estimator.intercepts[:] = torch.tensor([1.5, 1.0, 0.5, 0.0], device=device)
+    return estimator
+
+
+def _expected_probs(
+    estimator: OnlineAcceptanceEstimator,
+    logits: torch.Tensor,
+    steps: torch.Tensor,
+    temperature: torch.Tensor,
+) -> torch.Tensor:
+    """logit(q_max) by a different route than the kernel takes.
+
+    The kernel shifts by the max and drops the unit term; here the argmax is
+    masked out and the complement is reduced with logsumexp, so agreement is
+    evidence about the value rather than about a shared formula.
+    """
+    temp = torch.where(temperature > 0, temperature, torch.ones_like(temperature))
+    scaled = logits.double() / temp[:, None].double()
+    top = scaled.max(dim=-1)
+    rest = scaled.scatter(1, top.indices[:, None], float("-inf"))
+    feature = (top.values - torch.logsumexp(rest, dim=-1)).clamp(-40.0, 40.0)
+    weight, bias = estimator.slope, estimator.intercepts
+    return torch.sigmoid(weight.double() * feature + bias[steps].double()).float()
+
+
+@pytest.mark.parametrize("tokens_per_req", [1, NUM_STEPS])
+def test_predict_agrees_across_slot_and_batch_order(tokens_per_req: int):
+    """One row per request (per-step drafting) and a whole block at once.
+
+    Which one it is is inferred from whether the draft step is per-row, so the
+    two cases differ only in how ``draft_step`` is shaped.
+    """
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    estimator = _make_estimator(device)
+    num_reqs = 5
+
+    # Batch position i holds a shuffled persistent slot, as a served batch would.
+    slots = torch.randperm(MAX_NUM_REQS, device=device)[:num_reqs].to(torch.int32)
+    if tokens_per_req == 1:
+        idx_mapping = slots
+        draft_step = torch.tensor(2, device=device, dtype=torch.int32)
+        steps = draft_step.expand(num_reqs)
+    else:
+        idx_mapping = slots.repeat_interleave(NUM_STEPS)
+        draft_step = torch.arange(NUM_STEPS, device=device, dtype=torch.int32).repeat(
+            num_reqs
+        )
+        steps = draft_step
+    batch_rows = torch.arange(num_reqs, device=device).repeat_interleave(tokens_per_req)
+    logits = torch.randn(idx_mapping.shape[0], VOCAB_SIZE, device=device) * 4.0
+
+    confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device)
+    temperature = torch.rand(MAX_NUM_REQS, device=device) + 0.5
+    estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature)
+
+    expected = _expected_probs(
+        estimator, logits, steps.long(), temperature[idx_mapping.long()]
+    )
+    torch.testing.assert_close(
+        estimator.predictions[idx_mapping.long(), steps.long()],
+        expected,
+        atol=2e-3,
+        rtol=2e-3,
+    )
+    torch.testing.assert_close(
+        confidence_probs[batch_rows, steps.long()], expected, atol=2e-3, rtol=2e-3
+    )
+
+
+def test_predict_skips_cudagraph_padded_rows():
+    """Padded rows carry slot -1 and must not land on a live request."""
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    estimator = _make_estimator(device)
+    num_reqs, num_padded = 3, 6
+
+    slots = torch.tensor([7, 2, 11], device=device, dtype=torch.int32)
+    idx_mapping = torch.cat(
+        [slots, torch.full((num_padded,), -1, device=device, dtype=torch.int32)]
+    )
+    draft_step = torch.tensor(1, device=device, dtype=torch.int32)
+    logits = torch.randn(num_reqs + num_padded, VOCAB_SIZE, device=device) * 4.0
+
+    confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device)
+    temperature = torch.rand(MAX_NUM_REQS, device=device) + 0.5
+    estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature)
+
+    live = (slots.long(), torch.full((num_reqs,), 1, device=device))
+    expected = _expected_probs(
+        estimator, logits[:num_reqs], live[1], temperature[slots.long()]
+    )
+    torch.testing.assert_close(
+        estimator.predictions[live], expected, atol=2e-3, rtol=2e-3
+    )
+    # Nothing outside the live (request, position) cells was touched.
+    untouched = torch.ones(MAX_NUM_REQS, NUM_STEPS, dtype=torch.bool, device=device)
+    untouched[live] = False
+    assert not estimator.predictions[untouched].any()
+    assert not confidence_probs[num_reqs:].any()
+
+
+# A refit takes one undamped Newton step, so it needs a realistic number of
+# observations per round to be well conditioned -- a handful of requests can
+# send the step far past the optimum. This is the batch a served step sees.
+REFIT_NUM_REQS = 512
+
+
+def _drive_round(estimator, features, num_accepted, slots):
+    """Feed one step's graded drafts through predict's buffers into `step`."""
+    estimator.features[slots.long()] = features
+    estimator.predictions[slots.long()] = torch.sigmoid(
+        estimator.slope * features + estimator.intercepts
+    )
+    estimator.step(
+        slots,
+        (num_accepted + 1).to(torch.int32),
+        (features.shape[1] - num_accepted).to(torch.int32),
+    )
+
+
+def test_refit_recovers_a_shared_slope_and_per_position_intercepts():
+    """The fit must find the logistic that generated the observations.
+
+    One slope is shared across draft positions and each keeps its own intercept,
+    so the arrowhead solve is only correct if it recovers both from data whose
+    positions differ in level but not in slope.
+    """
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    true_slope = 0.6
+    true_bias = torch.tensor([1.4, 0.9, 0.4, -0.1], device=device)
+
+    estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device)
+    estimator.REFIT_INTERVAL = 5
+    slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32)
+    for _ in range(60):
+        features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) * 8.0
+        accepted = torch.rand_like(features) < torch.sigmoid(
+            true_slope * features + true_bias
+        )
+        # Verification stops at the first rejection, so a position's label is
+        # conditional on every shallower one having been accepted.
+        num_accepted = (~accepted).float().argmax(dim=1)
+        num_accepted[accepted.all(dim=1)] = NUM_STEPS
+        _drive_round(estimator, features, num_accepted, slots)
+
+    slope, bias = estimator.slope, estimator.intercepts
+    torch.testing.assert_close(
+        slope[0], torch.tensor(true_slope, device=device), atol=0.1, rtol=0
+    )
+    torch.testing.assert_close(bias, true_bias, atol=0.3, rtol=0)
+
+
+def test_refit_survives_a_position_with_no_observations():
+    """A draft position nobody reached must not stall the shared slope.
+
+    Its statistics stay zero while the others accumulate, which is the degenerate
+    case the arrowhead denominator has to tolerate.
+    """
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device)
+    estimator.REFIT_INTERVAL = 5
+    slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32)
+    before_slope = estimator.slope.clone()
+    before_intercepts = estimator.intercepts.clone()
+
+    # Every request is rejected at position 1, so positions 2+ are never graded.
+    num_accepted = torch.ones(REFIT_NUM_REQS, device=device)
+    for _ in range(20):
+        features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) * 8.0
+        _drive_round(estimator, features, num_accepted, slots)
+
+    assert estimator.slope.isfinite().all()
+    assert estimator.intercepts.isfinite().all()
+    assert not torch.equal(estimator.slope, before_slope), "slope never moved"
+    # Positions 2 and 3 were never observed, so their intercepts must be untouched.
+    torch.testing.assert_close(
+        estimator.intercepts[2:], before_intercepts[2:], rtol=0, atol=0
+    )
+
+
+def test_step_damping_scales_with_the_evidence():
+    """A round moves a coefficient in proportion to the samples behind it.
+
+    The Newton step itself barely changes with how much data a round carries --
+    curvature and score both scale with it -- so scaling a round's statistics
+    isolates the damping factor n / (n + DAMPING_OBSERVATIONS). The slope/
+    intercept coupling is zeroed so the two parameters can be checked apart.
+    """
+    device = torch.device("cuda")
+    damping = OnlineAcceptanceEstimator.DAMPING_OBSERVATIONS
+    per_position_n = 20.0
+
+    def refit_with(scale):
+        estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device)
+        estimator.info[:, 1] = 0.0  # sum(w*margin): no slope/intercept coupling
+        estimator.info[:, 2] = 1.0 * scale  # sum(w)
+        estimator.grad[:, 1] = 0.5 * scale  # sum(y - p)
+        # The slope's information and score are accumulated as per-position
+        # partials that the refit sums, so spread each total over the positions.
+        estimator.info[:, 0] = 30.0 * scale / NUM_STEPS  # sum(w*margin^2)
+        estimator.grad[:, 0] = 2.0 * scale / NUM_STEPS  # sum((y - p)*margin)
+        estimator.counts[:] = per_position_n * scale
+        estimator._steps_since_refit = estimator.REFIT_INTERVAL - 1
+        before_slope = estimator.slope.clone()
+        before_intercepts = estimator.intercepts.clone()
+        empty = torch.zeros(0, dtype=torch.int32, device=device)
+        estimator.step(empty, empty, empty)
+        return (
+            (estimator.slope - before_slope)[0].item(),
+            (estimator.intercepts - before_intercepts)[0].item(),
+        )
+
+    for scale in (1.0, 10.0, 100.0):
+        total_n = per_position_n * scale * NUM_STEPS
+        slope_delta, bias_delta = refit_with(scale)
+        assert slope_delta == pytest.approx(
+            (2.0 / 30.0) * total_n / (total_n + damping), rel=1e-3
+        )
+        assert bias_delta == pytest.approx(
+            0.5 * (per_position_n * scale) / (per_position_n * scale + damping),
+            rel=1e-3,
+        )
+
+    # A round with no observations at all must leave the coefficients alone.
+    assert refit_with(0.0) == (0.0, 0.0)
+
+
+def test_feature_ignores_which_token_was_drawn():
+    """The score must depend on the draft distribution, not on the draw.
+
+    This is what keeps rejection sampling lossless: if the score read the drawn
+    token, trimming a draft would condition on the token being trimmed and tilt
+    the emitted distribution toward the drafter's confident modes. ``predict``
+    takes no token argument, and this pins that -- two rows with identical
+    logits must score identically no matter what was sampled from them.
+    """
+    device = torch.device("cuda")
+    torch.manual_seed(0)
+    estimator = _make_estimator(device)
+
+    row = torch.randn(VOCAB_SIZE, device=device) * 4.0
+    logits = row.expand(2, VOCAB_SIZE).contiguous()
+    idx_mapping = torch.tensor([0, 1], device=device, dtype=torch.int32)
+    draft_step = torch.tensor(0, device=device, dtype=torch.int32)
+    temperature = torch.ones(MAX_NUM_REQS, device=device)
+    confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device)
+
+    estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature)
+    assert estimator.features[0, 0] == estimator.features[1, 0]
+
+    # And the value is the top token's log-odds, not the top-2 margin.
+    scaled = row.double()
+    top = scaled.max(dim=-1)
+    rest = scaled.scatter(0, top.indices[None], float("-inf"))
+    torch.testing.assert_close(
+        estimator.features[0, 0].double(),
+        (top.values - torch.logsumexp(rest, dim=-1)).clamp(-40.0, 40.0),
+        atol=2e-3,
+        rtol=2e-3,
+    )
diff --git a/tests/v1/spec_decode/test_acceptance_length.py b/tests/v1/spec_decode/test_acceptance_length.py
index 4d093f6055ce..04db178bacfb 100644
--- a/tests/v1/spec_decode/test_acceptance_length.py
+++ b/tests/v1/spec_decode/test_acceptance_length.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-EAGLE3 Acceptance Length Regression Tests.
+"""EAGLE3 Acceptance Length Regression Tests.
 
 These tests verify that acceptance lengths for EAGLE3 speculative decoding
 do not regress across vLLM commits. Each test runs inference on the MT-Bench
diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py
index d92f78a71535..116f4a7c3612 100644
--- a/tests/v1/spec_decode/test_adaptive_verification.py
+++ b/tests/v1/spec_decode/test_adaptive_verification.py
@@ -4,7 +4,9 @@
 from types import SimpleNamespace
 
 import numpy as np
+import pytest
 
+from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.worker.gpu.async_utils import StepTimingSample
 from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo
@@ -12,6 +14,7 @@
 from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
     AdaptiveVerificationManager,
     maybe_create_adaptive_verification_manager,
+    resolve_adaptive_cudagraph_mode,
 )
 from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping
 
@@ -35,6 +38,26 @@ def make_manager(
     return manager
 
 
+@pytest.mark.parametrize(
+    ("mode", "piecewise_capture_available", "expected"),
+    [
+        ("FULL_DECODE_ONLY", True, "FULL_DECODE_ONLY"),
+        ("FULL_DECODE_ONLY", False, "FULL_DECODE_ONLY"),
+        ("FULL", True, "FULL_AND_PIECEWISE"),
+        ("FULL", False, "FULL_DECODE_ONLY"),
+        ("FULL_AND_PIECEWISE", True, "FULL_AND_PIECEWISE"),
+        ("FULL_AND_PIECEWISE", False, "FULL_DECODE_ONLY"),
+    ],
+)
+def test_resolve_adaptive_cudagraph_mode(mode, piecewise_capture_available, expected):
+    assert (
+        resolve_adaptive_cudagraph_mode(
+            CUDAGraphMode[mode], piecewise_capture_available=piecewise_capture_available
+        )
+        == CUDAGraphMode[expected]
+    )
+
+
 def test_manager_scopes_varlen_check_without_weakening_runner_cg_mode(monkeypatch):
     class Backend:
         @classmethod
diff --git a/tests/v1/spec_decode/test_backup_token_async_spec.py b/tests/v1/spec_decode/test_backup_token_async_spec.py
deleted file mode 100644
index 9340503ea4f2..000000000000
--- a/tests/v1/spec_decode/test_backup_token_async_spec.py
+++ /dev/null
@@ -1,147 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Regression tests for the backup token fix in prepare_next_token_ids_padded.
-
-Fixes #38098: with async scheduling, seq_lens_cpu is inflated by unaccepted
-draft token placeholders, causing get_token_id() to return -1.
-"""
-
-from __future__ import annotations
-
-import numpy as np
-import pytest
-import torch
-
-
-class _FakeRequest:
-    def __init__(self, prompt_tokens: list[int], output_tokens: list[int]):
-        self.num_prompt_tokens = len(prompt_tokens)
-        self._prompt = prompt_tokens
-        self._output = output_tokens
-
-    @property
-    def num_tokens(self) -> int:
-        return self.num_prompt_tokens + len(self._output)
-
-    def get_token_id(self, idx: int) -> int:
-        if idx < self.num_prompt_tokens:
-            return self._prompt[idx]
-        out_idx = idx - self.num_prompt_tokens
-        if out_idx < len(self._output):
-            return self._output[out_idx]
-        return -1  # out of range
-
-
-class _FakeInputBatch:
-    def __init__(
-        self,
-        req_ids: list[str],
-        num_tokens_no_spec: list[int],
-        vocab_size: int = 32000,
-    ):
-        self.req_ids = req_ids
-        self.num_reqs = len(req_ids)
-        self.vocab_size = vocab_size
-        self.num_tokens_no_spec = np.array(num_tokens_no_spec, dtype=np.int64)
-
-
-def _make_requests(
-    req_ids: list[str],
-    prompt_lens: list[int],
-    output_lens: list[int],
-) -> dict[str, _FakeRequest]:
-    requests = {}
-    for rid, plen, olen in zip(req_ids, prompt_lens, output_lens):
-        requests[rid] = _FakeRequest(list(range(plen)), list(range(1000, 1000 + olen)))
-    return requests
-
-
-def _backup_buggy(
-    seq_lens_cpu: torch.Tensor,
-    requests: dict[str, _FakeRequest],
-    batch: _FakeInputBatch,
-) -> list[int]:
-    """Old logic: uses seq_lens_cpu directly (may be inflated)."""
-    n = batch.num_reqs
-    return [
-        requests[batch.req_ids[i]].get_token_id(int(seq_lens_cpu[i])) for i in range(n)
-    ]
-
-
-def _backup_fixed(
-    requests: dict[str, _FakeRequest],
-    batch: _FakeInputBatch,
-) -> list[int]:
-    """New logic: uses num_tokens_no_spec - 1 (last committed token)."""
-    n = batch.num_reqs
-    idx = (batch.num_tokens_no_spec[:n] - 1).tolist()
-    return [requests[batch.req_ids[i]].get_token_id(int(idx[i])) for i in range(n)]
-
-
-class TestBackupTokenAsyncSpec:
-    def test_no_inflation_fixed_returns_last_token(self):
-        req_ids = ["r0", "r1"]
-        requests = _make_requests(req_ids, [3, 3], [2, 2])
-        batch = _FakeInputBatch(req_ids, [5, 5])
-        # idx = 5-1 = 4 → output[1] = 1001
-        assert _backup_fixed(requests, batch) == [1001, 1001]
-
-    def test_inflation_buggy_returns_placeholder(self):
-        req_ids = ["r0", "r1"]
-        requests = _make_requests(req_ids, [3, 3], [2, 2])
-        batch = _FakeInputBatch(req_ids, [5, 5])
-        # inflated by 3 spec tokens → idx 8 is out of range
-        seq_lens = torch.tensor([8, 8], dtype=torch.int64)
-        assert _backup_buggy(seq_lens, requests, batch) == [-1, -1]
-
-    def test_inflation_fixed_returns_correct_token(self):
-        req_ids = ["r0", "r1"]
-        requests = _make_requests(req_ids, [3, 3], [2, 2])
-        batch = _FakeInputBatch(req_ids, [5, 5])
-        assert _backup_fixed(requests, batch) == [1001, 1001]
-
-    def test_mixed_inflation_per_request(self):
-        req_ids = ["r0", "r1", "r2"]
-        requests = {
-            "r0": _FakeRequest([0, 1], [1000, 1001, 1002]),
-            "r1": _FakeRequest([0, 1, 2, 3], [2000]),
-            "r2": _FakeRequest([0], [3000, 3001, 3002, 3003]),
-        }
-        batch = _FakeInputBatch(req_ids, [5, 5, 5])
-        seq_lens = torch.tensor([7, 9, 5], dtype=torch.int64)
-
-        assert _backup_buggy(seq_lens, requests, batch) == [-1, -1, -1]
-        assert _backup_fixed(requests, batch) == [1002, 2000, 3003]
-
-    def test_prefill_only_request(self):
-        """No output tokens yet — backup should be the last prompt token."""
-        req_ids = ["r0"]
-        requests = {"r0": _FakeRequest([10, 20, 30], [])}
-        batch = _FakeInputBatch(req_ids, [3])
-        # idx = 3-1 = 2 → prompt[2] = 30
-        assert _backup_fixed(requests, batch) == [30]
-
-    @pytest.mark.parametrize("num_spec_tokens", [1, 2, 3, 4, 5])
-    def test_various_spec_token_counts(self, num_spec_tokens: int):
-        req_ids = ["r0"]
-        requests = {"r0": _FakeRequest([0, 1, 2], list(range(1000, 1005)))}
-        batch = _FakeInputBatch(req_ids, [8])
-        # idx = 8-1 = 7 → output[4] = 1004
-        assert _backup_fixed(requests, batch) == [1004]
-
-    def test_buggy_code_was_always_off_by_one(self):
-        """The original code used seq_len as index, which is always one past
-        the end of output_token_ids even without async inflation."""
-        req_ids = ["r0"]
-        requests = {"r0": _FakeRequest([0, 1, 2], [1000, 1001])}
-        batch = _FakeInputBatch(req_ids, [5])
-
-        # no inflation: seq_len == num_tokens == 5 → idx 5 is out of range
-        seq_lens = torch.tensor([5], dtype=torch.int64)
-        assert _backup_buggy(seq_lens, requests, batch) == [-1]
-        assert _backup_fixed(requests, batch) == [1001]
-
-        # with inflation: still -1, fixed still correct
-        seq_lens_inf = torch.tensor([8], dtype=torch.int64)
-        assert _backup_buggy(seq_lens_inf, requests, batch) == [-1]
-        assert _backup_fixed(requests, batch) == [1001]
diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py
index 0bff2d537eb0..72379fa7d7fc 100644
--- a/tests/v1/spec_decode/test_dflash2.py
+++ b/tests/v1/spec_decode/test_dflash2.py
@@ -7,6 +7,7 @@
 import torch
 
 from vllm.model_executor.models.qwen3_dflash2 import _grouped_conv, _score_edges
+from vllm.platforms import current_platform
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
 from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator
 
@@ -35,6 +36,55 @@ def test_grouped_conv_matches_reference(block_size: int):
     torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2))
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize(
+    "batch,num_groups,group_size,block_size,taps",
+    [
+        (7, 13, 11, 5, 1),
+        (7, 13, 11, 5, 3),
+        (7, 13, 11, 8, 2),
+        (16, 64, 16, 8, 2),
+        (16, 160, 16, 8, 2),
+    ],
+)
+def test_grouped_conv_triton_matches_reference(
+    dtype: torch.dtype,
+    batch: int,
+    num_groups: int,
+    group_size: int,
+    block_size: int,
+    taps: int,
+):
+    torch.manual_seed(0)
+    rows = batch * block_size
+    hidden = torch.randn(rows, num_groups * group_size, device="cuda", dtype=dtype)
+    base = torch.randn(taps, num_groups * group_size, device="cuda", dtype=dtype)
+    projected = torch.randn(rows, 2, taps, num_groups, device="cuda", dtype=dtype)
+    delta = projected[:, 1]
+
+    actual = _grouped_conv(
+        hidden, delta, base, block_size, num_groups, group_size, taps
+    )
+
+    hidden_blocks = hidden.float().view(batch, block_size, num_groups, group_size)
+    expected = torch.zeros_like(hidden_blocks)
+    base_blocks = base.float().view(taps, num_groups, group_size)
+    delta_blocks = delta.float().view(batch, block_size, taps, num_groups)
+    for position in range(block_size):
+        for tap in range(min(taps, position + 1)):
+            expected[:, position] += (
+                base_blocks[tap] + delta_blocks[:, position, tap, :, None]
+            ) * hidden_blocks[:, position - tap]
+
+    torch.testing.assert_close(
+        actual,
+        expected.flatten(0, 1).flatten(-2).to(dtype),
+        rtol=1e-2 if dtype is torch.bfloat16 else 1e-5,
+        atol=1e-2 if dtype is torch.bfloat16 else 1e-5,
+    )
+
+
 def test_selector_edges_match_sequential_reference():
     torch.manual_seed(1)
     batch, steps, top_k, rank = 2, 4, 3, 5
@@ -159,6 +209,7 @@ def test_dflash2_model_decoder_layer_cls(monkeypatch):
         model_config=SimpleNamespace(
             dtype=torch.float32,
             is_mm_prefix_lm=False,
+            rswa_window=None,
         ),
         kernel_config=SimpleNamespace(
             linear_backend="auto",
diff --git a/tests/v1/spec_decode/test_dflash_causality.py b/tests/v1/spec_decode/test_dflash_causality.py
index 0f1de86cf95b..dc13e5570090 100644
--- a/tests/v1/spec_decode/test_dflash_causality.py
+++ b/tests/v1/spec_decode/test_dflash_causality.py
@@ -103,6 +103,18 @@ def test_dflash_fc_uses_aux_layer_count():
     assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096
 
 
+def test_dflash_target_layers_map_to_post_layer_boundaries():
+    vllm_config = _vllm_config(
+        dflash_config={"target_layer_ids": [0, 16, 31]},
+    )
+
+    assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == (
+        1,
+        17,
+        32,
+    )
+
+
 @pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"])
 def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name):
     layer_ids = [1, 17, 32]
diff --git a/tests/v1/spec_decode/test_dflash_profile.py b/tests/v1/spec_decode/test_dflash_profile.py
new file mode 100644
index 000000000000..3b6f6061eb2e
--- /dev/null
+++ b/tests/v1/spec_decode/test_dflash_profile.py
@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
+from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator
+
+
+class _DefaultSpeculator(BaseSpeculator):
+    def init_cudagraph_manager(self, cudagraph_mode):
+        pass
+
+    def capture(self):
+        pass
+
+    def propose(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class _StopDummyRun(Exception):
+    pass
+
+
+def _run_dummy_batch(
+    speculator,
+    *,
+    max_num_reqs=512,
+    is_profile=False,
+    skip_attn=False,
+):
+    captured = {}
+
+    def execute_model(scheduler_output, **kwargs):
+        captured["scheduler_output"] = scheduler_output
+        raise _StopDummyRun
+
+    runner = SimpleNamespace(
+        max_num_reqs=max_num_reqs,
+        max_num_tokens=2048,
+        decode_query_len=1,
+        speculator=speculator,
+        kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
+        is_first_pp_rank=True,
+        lora_config=None,
+        maybe_dummy_run_with_lora=lambda *args, **kwargs: nullcontext(),
+        execute_model=execute_model,
+    )
+
+    with pytest.raises(_StopDummyRun):
+        GPUModelRunner._dummy_run(
+            runner,
+            num_tokens=2048,
+            skip_attn=skip_attn,
+            is_profile=is_profile,
+        )
+
+    return captured["scheduler_output"]
+
+
+def test_default_speculator_query_width_is_one():
+    assert _DefaultSpeculator().num_query_per_req == 1
+
+
+@pytest.mark.parametrize("speculator_cls", [DFlashSpeculator, DSparkSpeculator])
+@pytest.mark.parametrize(
+    ("is_profile", "skip_attn"),
+    [(False, False), (True, False), (True, True)],
+)
+def test_dummy_run_caps_parallel_draft_requests(speculator_cls, is_profile, skip_attn):
+    speculator = speculator_cls.__new__(speculator_cls)
+    speculator.num_query_per_req = 5
+
+    scheduler_output = _run_dummy_batch(
+        speculator,
+        is_profile=is_profile,
+        skip_attn=skip_attn,
+    )
+    assert len(scheduler_output.num_scheduled_tokens) == 409
+    assert scheduler_output.total_num_scheduled_tokens == 2048
+
+
+def test_dummy_run_keeps_already_fitting_parallel_draft_batch():
+    speculator = DFlashSpeculator.__new__(DFlashSpeculator)
+    speculator.num_query_per_req = 5
+
+    scheduler_output = _run_dummy_batch(speculator, max_num_reqs=128)
+    assert len(scheduler_output.num_scheduled_tokens) == 128
+
+
+def test_dummy_run_keeps_default_speculator_batch():
+    scheduler_output = _run_dummy_batch(_DefaultSpeculator())
+    assert len(scheduler_output.num_scheduled_tokens) == 512
diff --git a/tests/v1/spec_decode/test_draft_attention_backend_override.py b/tests/v1/spec_decode/test_draft_attention_backend_override.py
index 5984d337692b..26bcd68fea8f 100644
--- a/tests/v1/spec_decode/test_draft_attention_backend_override.py
+++ b/tests/v1/spec_decode/test_draft_attention_backend_override.py
@@ -6,11 +6,13 @@
 assert on the config ``load_eagle_model`` hands to ``get_model``.
 """
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
 
+from vllm.config import LoadConfig, SpeculativeConfig
 from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
 
 
@@ -35,6 +37,9 @@ class _SpeculativeConfig:
     moe_backend: str | None = None
     kv_cache_dtype: str | None = None
     draft_model_config: object = None
+    draft_load_config: object = None
+
+    apply_draft_overrides = SpeculativeConfig.apply_draft_overrides
 
 
 @dataclass
@@ -43,6 +48,7 @@ class _VllmConfig:
     kernel_config: _KernelConfig
     cache_config: _CacheConfig
     speculative_config: _SpeculativeConfig
+    load_config: LoadConfig = field(default_factory=LoadConfig)
 
 
 def _config(target_backend: str, draft_backend: str | None) -> _VllmConfig:
@@ -65,6 +71,10 @@ def _fake_get_model(*, vllm_config, model_config):
 
     with (
         patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model),
+        patch(
+            "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group",
+            return_value=SimpleNamespace(world_size=1),
+        ),
         pytest.raises(_Captured) as exc,
     ):
         load_eagle_model(object(), cfg)
diff --git a/tests/v1/spec_decode/test_draft_moe_backend_override.py b/tests/v1/spec_decode/test_draft_moe_backend_override.py
index d4acd9e4be49..a0cc26c16df4 100644
--- a/tests/v1/spec_decode/test_draft_moe_backend_override.py
+++ b/tests/v1/spec_decode/test_draft_moe_backend_override.py
@@ -8,11 +8,13 @@
 server fails to start rather than falling back.
 """
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
 
+from vllm.config import LoadConfig, SpeculativeConfig
 from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
 
 
@@ -32,6 +34,9 @@ class _SpeculativeConfig:
     moe_backend: str | None = None
     kv_cache_dtype: str | None = None
     draft_model_config: object = None
+    draft_load_config: object = None
+
+    apply_draft_overrides = SpeculativeConfig.apply_draft_overrides
 
 
 @dataclass
@@ -39,6 +44,7 @@ class _VllmConfig:
     kernel_config: _KernelConfig
     cache_config: _CacheConfig
     speculative_config: _SpeculativeConfig
+    load_config: LoadConfig = field(default_factory=LoadConfig)
 
 
 def _config(target_moe: str, draft_moe: str | None) -> _VllmConfig:
@@ -62,6 +68,10 @@ def _fake_get_model(*, vllm_config, model_config):
 
     with (
         patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model),
+        patch(
+            "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group",
+            return_value=SimpleNamespace(world_size=1),
+        ),
         pytest.raises(_Captured) as exc,
     ):
         load_eagle_model(object(), cfg)
diff --git a/tests/v1/spec_decode/test_dynamic_sd_cug.py b/tests/v1/spec_decode/test_dynamic_sd_cug.py
index a5cbd119bc66..8abbe91c1a35 100644
--- a/tests/v1/spec_decode/test_dynamic_sd_cug.py
+++ b/tests/v1/spec_decode/test_dynamic_sd_cug.py
@@ -39,7 +39,6 @@ def _create_vllm_config_for_dsd(
     a schedule covering every query length in ``[1, max_decode_query_len]`` is
     generated.
     """
-
     max_decode_query_len = max_spec_tokens + 1
     max_capture_tokens = max_num_seqs * max_decode_query_len
 
@@ -91,7 +90,6 @@ def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch)
     shapes have been built, dispatch() should pick a FULL graph for every
     uniform decode batch shape produced by DSD up to max_num_seqs.
     """
-
     max_num_seqs = 512
     max_spec_tokens = 7
     max_decode_query_len = max_spec_tokens + 1
@@ -190,7 +188,6 @@ def test_dynamic_sd_non_uniform_batch_falls_back_to_piecewise(monkeypatch):
     FULL candidates should be skipped in favor of the mixed-batch PIECEWISE
     graph under FULL_AND_PIECEWISE mode.
     """
-
     max_spec_tokens = 4
 
     monkeypatch.setattr(
@@ -236,7 +233,6 @@ def test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph(monkeypatch):
     by coincidence; replaying the FULL graph over them corrupts every request
     in the batch. See https://github.com/vllm-project/vllm/issues/49918.
     """
-
     max_spec_tokens = 7
     decode_query_len = max_spec_tokens + 1
 
@@ -307,7 +303,6 @@ def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch):
     Uniform batches with smaller query lengths should therefore miss the FULL
     path entirely when using FULL_AND_PIECEWISE.
     """
-
     max_num_seqs = 512
     max_spec_tokens = 7
     max_decode_query_len = max_spec_tokens + 1
@@ -368,7 +363,6 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch):
     FULL graphs, while every other query length (e.g. the lower values 1, 2, 3)
     must fall back to the mixed-batch PIECEWISE graph.
     """
-
     max_num_seqs = 128
     max_spec_tokens = 7
     max_decode_query_len = max_spec_tokens + 1
diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py
index b62d7e90da42..9b3bfd73d06c 100644
--- a/tests/v1/spec_decode/test_eagle.py
+++ b/tests/v1/spec_decode/test_eagle.py
@@ -116,8 +116,7 @@ def _create_proposer(
 
 
 def test_prepare_next_token_ids():
-    """
-    Test for prepare_next_token_ids_cpu and prepare_next_token_ids_padded.
+    """Test for prepare_next_token_ids_cpu and prepare_next_token_ids_padded.
     Each will produce a device tensor of next_token_ids, taking as input
     either the GPU tensor of sampled_token_ids with -1 for rejected tokens,
     or the CPU python list[list[int]] with the rejected tokens removed.
@@ -196,8 +195,7 @@ def test_prepare_next_token_ids():
 
 
 def test_prepare_inputs():
-    """
-    cu_target_query_lens: [0, a, a + b, a + b + c]
+    """cu_target_query_lens: [0, a, a + b, a + b + c]
     num_rejected_tokens: [n1, n2, n3]
     num_tokens_per_req: [a - n1, b - n2, c - n3]
     cu_num_tokens: [0, a - n1, a + b - n1 - n2, a + b + c - n1 - n2 - n3]
@@ -286,8 +284,7 @@ def test_prepare_inputs():
 
 
 def test_prepare_inputs_padded():
-    """
-    Input scenario is 3 requests with num_speculative_tokens == 2 and:
+    """Input scenario is 3 requests with num_speculative_tokens == 2 and:
     - Request 1: query_len = 3, rejected = 1
     - Request 2: query_len = 3, rejected = 0
     - Request 3: query_len = 3, rejected = 2
@@ -297,7 +294,6 @@ def test_prepare_inputs_padded():
     Reason: After accounting for rejections, these are the valid token positions
             from the original indices to sample from.
     """
-
     device = torch.device(DEVICE_TYPE)
 
     expected_token_indices_to_sample = torch.tensor(
@@ -350,8 +346,7 @@ def test_prepare_inputs_padded():
 
 
 def test_set_inputs_first_pass_default_eagle():
-    """
-    Test for set_inputs_first_pass without extra input slots (default EAGLE).
+    """Test for set_inputs_first_pass without extra input slots (default EAGLE).
 
     This tests the path where needs_extra_input_slots=False, which is the
     default EAGLE pathway. In this case:
@@ -436,8 +431,7 @@ def test_set_inputs_first_pass_default_eagle():
 
 
 def test_set_inputs_first_pass_draft_model():
-    """
-    Test for set_inputs_first_pass with a draft model (extra input slots,
+    """Test for set_inputs_first_pass with a draft model (extra input slots,
     no shift).
 
     This tests the path where needs_extra_input_slots=True and
@@ -578,8 +572,7 @@ def test_set_inputs_first_pass_draft_model():
 
 
 def test_set_inputs_first_pass_parallel_drafting():
-    """
-    Test for set_inputs_first_pass with parallel drafting (extra input slots,
+    """Test for set_inputs_first_pass with parallel drafting (extra input slots,
     with shift).
 
     This tests the path where needs_extra_input_slots=True and
@@ -1107,8 +1100,7 @@ def fake_compute_probs(logits, sampling_metadata, use_fp64_gumbel):
 
 
 def test_set_inputs_first_pass_dflash():
-    """
-    Test for DFlash set_inputs_first_pass.
+    """Test for DFlash set_inputs_first_pass.
 
     DFlash uses cross-attention: context tokens become K/V and only
     query tokens (bonus + mask) are Q. This tests the DFlash-specific
diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
index b3006081e579..b542388952b9 100644
--- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for the EAGLE speculator's draft attention metadata builder.
 
-These tests guard the regression where ``_build_draft_attn_metadata`` did
+These tests guard the regression where ``_build_uniform_attn_metadata`` did
 not populate ``seq_lens_cpu_upper_bound`` on the per-step
 ``CommonAttentionMetadata``. Several downstream attention backends and
 helpers (``split_decodes_prefills_and_extends``, the MLA indexer,
@@ -13,11 +13,14 @@
     AssertionError: assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
 """
 
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
 from unittest.mock import patch
 
+import numpy as np
 import torch
 
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.spec_decode import speculator as base_speculator
 from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator
 
@@ -30,7 +33,7 @@ def _make_fake_speculator(
     draft_max_seq_len: int = 1024,
 ) -> SimpleNamespace:
     """Build a fake EagleSpeculator with just the attributes used by
-    ``_build_draft_attn_metadata``. We deliberately avoid constructing a
+    ``_build_uniform_attn_metadata``. We deliberately avoid constructing a
     real ``EagleSpeculator`` because that requires a full ``VllmConfig``
     and a draft model.
     """
@@ -46,8 +49,9 @@ def _make_fake_speculator(
         cp_rank=0,
         cp_interleave=1,
     )
-    return SimpleNamespace(
-        arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"),
+    fake = SimpleNamespace(
+        arange_np=np.arange(max_num_reqs + 1, dtype=np.int32),
+        draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool),
         block_tables=fake_block_tables,
         input_buffers=fake_input_buffers,
         attn_groups=[],
@@ -55,21 +59,42 @@ def _make_fake_speculator(
         max_model_len=max_model_len,
         draft_max_seq_len=draft_max_seq_len,
     )
+    # The uniform wrapper delegates through self; bind the real implementation.
+    fake._build_attn_metadata = MethodType(EagleSpeculator._build_attn_metadata, fake)
+    return fake
 
 
-def _run_build(fake, *, num_reqs, num_reqs_padded, num_tokens_padded, base, step):
+def _run_build(
+    fake,
+    *,
+    num_reqs,
+    num_reqs_padded,
+    base,
+    step,
+    num_query_per_req=1,
+    cg_mode=CUDAGraphMode.FULL,
+):
     captured: dict[str, object] = {}
 
     def fake_build_attn_metadata(**kwargs):
         captured.update(kwargs)
         return {}
 
+    # Request padding only occurs under FULL graphs, where the captured batch
+    # is a uniform decode of num_reqs_padded * num_query_per_req tokens.
+    batch_desc = BatchExecutionDescriptor(
+        cg_mode=cg_mode,
+        num_tokens=num_reqs_padded * num_query_per_req,
+        num_reqs=num_reqs_padded,
+        uniform_token_count=num_query_per_req,
+    )
+
     with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata):
-        EagleSpeculator._build_draft_attn_metadata(
+        EagleSpeculator._build_uniform_attn_metadata(
             fake,  # type: ignore[arg-type]
+            batch_desc=batch_desc,
             num_reqs=num_reqs,
-            num_reqs_padded=num_reqs_padded,
-            num_tokens_padded=num_tokens_padded,
+            num_query_per_req=num_query_per_req,
             seq_lens_cpu_upper_bound=base,
             step=step,
         )
@@ -84,9 +109,7 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound():
     fake = _make_fake_speculator()
     base = torch.tensor([100, 200, 300, 0], dtype=torch.int32)
 
-    captured = _run_build(
-        fake, num_reqs=3, num_reqs_padded=4, num_tokens_padded=4, base=base, step=2
-    )
+    captured = _run_build(fake, num_reqs=3, num_reqs_padded=4, base=base, step=2)
 
     bound = captured["seq_lens_cpu_upper_bound"]
     assert isinstance(bound, torch.Tensor), (
@@ -101,20 +124,19 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound():
     assert torch.equal(bound, torch.tensor([102, 202, 302, 0], dtype=torch.int32))
 
 
-def test_build_draft_attn_metadata_handles_zero_unpadded_reqs():
-    """Edge case: when ``num_reqs == 0`` the upper-bound tensor must
-    still be a valid all-zero tensor of length ``num_reqs_padded``."""
+def test_build_draft_attn_metadata_zeroes_padded_upper_bound_tail():
+    """The padded tail of the upper-bound tensor is zeroed, so it stays a
+    valid tensor of length ``num_reqs_padded`` regardless of padding."""
     fake = _make_fake_speculator()
     base = torch.zeros(2, dtype=torch.int32)
 
-    captured = _run_build(
-        fake, num_reqs=0, num_reqs_padded=2, num_tokens_padded=2, base=base, step=1
-    )
+    captured = _run_build(fake, num_reqs=1, num_reqs_padded=2, base=base, step=1)
 
     bound = captured["seq_lens_cpu_upper_bound"]
     assert isinstance(bound, torch.Tensor)
     assert bound.shape == (2,)
-    assert torch.equal(bound, torch.zeros(2, dtype=torch.int32))
+    # base is all zeros, so the real entry is 0 + step and the pad entry is 0.
+    assert torch.equal(bound, torch.tensor([1, 0], dtype=torch.int32))
 
 
 def test_build_draft_attn_metadata_clamps_to_max_model_len():
@@ -123,9 +145,7 @@ def test_build_draft_attn_metadata_clamps_to_max_model_len():
     fake = _make_fake_speculator(max_model_len=1024)
     base = torch.tensor([1023, 500], dtype=torch.int32)
 
-    captured = _run_build(
-        fake, num_reqs=2, num_reqs_padded=2, num_tokens_padded=2, base=base, step=3
-    )
+    captured = _run_build(fake, num_reqs=2, num_reqs_padded=2, base=base, step=3)
 
     bound = captured["seq_lens_cpu_upper_bound"]
     # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected.
@@ -146,14 +166,11 @@ def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave):
         out[num_reqs:].zero_()
         return out
 
-    with patch.object(
-        base_speculator, "maybe_prepare_dcp_local_seq_lens", fake_prepare
-    ):
+    with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
         captured = _run_build(
             fake,
             num_reqs=3,
             num_reqs_padded=4,
-            num_tokens_padded=4,
             base=torch.tensor([5, 9, 16]),
             step=0,
         )
diff --git a/tests/v1/spec_decode/test_extract_hidden_states.py b/tests/v1/spec_decode/test_extract_hidden_states.py
index 6b4e53ced679..7eb8eca023ba 100644
--- a/tests/v1/spec_decode/test_extract_hidden_states.py
+++ b/tests/v1/spec_decode/test_extract_hidden_states.py
@@ -7,7 +7,7 @@
 import numpy as np
 import pytest
 import torch
-from transformers import CLIPVisionConfig, LlamaConfig, LlavaConfig, PretrainedConfig
+from transformers import CLIPVisionConfig, LlamaConfig, LlavaConfig, PreTrainedConfig
 
 from tests.v1.attention.utils import (
     BatchSpec,
@@ -130,8 +130,7 @@ def test_proposer_initialization_missing_layer_ids():
 
 
 def test_prepare_next_token_ids_padded():
-    """
-    Test for prepare_next_token_ids_padded with extract_hidden_states.
+    """Test for prepare_next_token_ids_padded with extract_hidden_states.
 
     Since num_speculative_tokens == 1, sampled_token_ids has shape (batch_size, 1).
     For each request we either use the sampled token (if valid and not discarded)
@@ -196,8 +195,7 @@ def test_prepare_next_token_ids_padded():
 
 
 def test_propose():
-    """
-    Test the propose() method of ExtractHiddenStatesProposer.
+    """Test the propose() method of ExtractHiddenStatesProposer.
 
     This should:
     1. Accept target hidden states and sampled token IDs
@@ -338,7 +336,7 @@ def test_propose_different_layer_counts(num_hidden_layers):
 # ---------------------------------------------------------------------------
 
 
-class _DummyVLMConfig(PretrainedConfig):
+class _DummyVLMConfig(PreTrainedConfig):
     """Minimal composite config that mimics VLMs like Kimi-K2.5 or LLaVA.
 
     The text model's parameters (hidden_size, num_attention_heads, …) live
@@ -347,11 +345,11 @@ class _DummyVLMConfig(PretrainedConfig):
 
     model_type = "test_vlm"
 
-    def __init__(self, text_config: PretrainedConfig, **kwargs):
+    def __init__(self, text_config: PreTrainedConfig, **kwargs):
         self.text_config = text_config
         super().__init__(architectures=["LlamaForCausalLM"], **kwargs)
 
-    def get_text_config(self, decoder: bool = False) -> PretrainedConfig:
+    def get_text_config(self, decoder: bool = False) -> PreTrainedConfig:
         del decoder
         return self.text_config
 
@@ -405,7 +403,7 @@ def test_extract_hidden_states_config_preserves_vlm_text_config():
         eagle_aux_hidden_state_layer_ids=[1, 2],
     )
 
-    # The fix: text_config is still a PretrainedConfig, not a dict.
+    # The fix: text_config is still a PreTrainedConfig, not a dict.
     assert isinstance(extract_config.text_config, LlamaConfig)
 
     extracted = get_hf_text_config(extract_config)
@@ -475,7 +473,7 @@ def test_extract_hidden_states_speculative_config_vlm():
 
 def test_extract_hidden_states_config_invalid_text_config():
     """A nested text_config missing required attrs must still be rejected."""
-    broken_text_config = PretrainedConfig(hidden_size=128)
+    broken_text_config = PreTrainedConfig(hidden_size=128)
     vlm_config = _DummyVLMConfig(text_config=broken_text_config)
 
     extract_config = ExtractHiddenStatesConfig(
diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py
index e334371f6d8a..e3803509e39d 100644
--- a/tests/v1/spec_decode/test_mtp.py
+++ b/tests/v1/spec_decode/test_mtp.py
@@ -5,6 +5,7 @@
 
 import pytest
 import torch
+import torch.nn as nn
 
 from tests.v1.attention.utils import (
     BatchSpec,
@@ -31,6 +32,31 @@
 DEVICE_TYPE = current_platform.device_type
 
 
+def test_glm_mtp_defers_lm_head(default_vllm_config):
+    from vllm.model_executor.models import deepseek_mtp
+    from vllm.models.glm5next.common import mtp
+
+    config = mock.MagicMock(
+        hidden_size=16,
+        rms_norm_eps=1e-5,
+        index_topk=8,
+        index_kpool=4,
+    )
+    vllm_config = mock.MagicMock()
+    vllm_config.speculative_config.draft_model_config.hf_config = config
+    vllm_config.scheduler_config.max_num_batched_tokens = 4
+
+    with (
+        mock.patch.object(mtp, "Glm5NextDecoderLayer", return_value=nn.Identity()),
+        mock.patch.object(deepseek_mtp, "ParallelLMHead") as parallel_lm_head,
+        mock.patch.object(mtp.current_platform, "device_type", "cpu"),
+    ):
+        layer = mtp.Glm5NextMultiTokenPredictorLayer(vllm_config, "model.layers.1")
+
+    parallel_lm_head.assert_not_called()
+    assert layer.shared_head.head is None
+
+
 def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
     """Create an MTP proposer with unified model configuration."""
     model_config = ModelConfig(
@@ -66,7 +92,6 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
 @mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_model")
 def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_group):
     """Test MTP-specific model loading with unified model approach."""
-
     # Setup mocks
     mock_model = mock.MagicMock()
     mock_model.model.embed_tokens.weight.shape = (131072, 4096)
@@ -117,8 +142,7 @@ class _TargetModelStub(LlamaForCausalLM):
 
 @pytest.mark.parametrize("num_speculative_tokens", [1])
 def test_mtp_propose(num_speculative_tokens, monkeypatch):
-    """Test that MTP's forward method returns hidden states directly"""
-
+    """Test that MTP's forward method returns hidden states directly."""
     device = torch.device(DEVICE_TYPE)
     batch_size = 2
     seq_lens = [5, 3]
diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py
index 330e10438f0b..4e12a6f49dd8 100644
--- a/tests/v1/spec_decode/test_mtp_structured_output.py
+++ b/tests/v1/spec_decode/test_mtp_structured_output.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""grammar_bitmask under spec-decode draft padding (#44006)."""
+"""Backend and engine-reasoner checks that are not manager-flow tests."""
 
 import pytest
 from transformers import AutoTokenizer
@@ -49,221 +49,6 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'):
     return tokenizer, manager, request, prompt
 
 
-@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
-def test_bitmask_with_padded_invalid_drafts(backend):
-    """Bitmask handles -1 padded drafts and returns N+1 rows."""
-    tokenizer, manager, request, prompt = _make_manager_and_request(
-        backend, prompt_str='{"a"'
-    )
-    grammar = request.structured_output_request.grammar
-
-    assert grammar.accept_tokens(request.request_id, prompt)
-
-    valid_drafts = [tokenizer.encode(":")[0], tokenizer.encode(' "')[0]]
-    padded = valid_drafts + [-1, -1]
-
-    bitmask = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: padded},
-    )
-
-    assert bitmask is not None
-    assert bitmask.shape[0] == len(padded) + 1
-    assert not grammar.is_terminated()
-
-
-@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
-def test_bitmask_when_grammar_terminates_mid_window(backend):
-    """Drafts following an EOS that terminates the grammar are a no-op."""
-    tokenizer, manager, request, prompt = _make_manager_and_request(backend)
-    grammar = request.structured_output_request.grammar
-
-    assert grammar.accept_tokens(request.request_id, prompt)
-    eos = tokenizer.eos_token_id
-    drafts = [eos] + [tokenizer.encode(" ")[0]] * (NUM_SPEC_TOKENS - 1)
-
-    bitmask = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: drafts},
-    )
-
-    assert bitmask is not None
-    assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1
-    assert not grammar.is_terminated()
-
-
-@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
-def test_bitmask_idempotent_across_calls(backend):
-    """Repeated calls with the same input return the same bitmask."""
-    tokenizer, manager, request, prompt = _make_manager_and_request(
-        backend, prompt_str='{"a"'
-    )
-    grammar = request.structured_output_request.grammar
-
-    assert grammar.accept_tokens(request.request_id, prompt)
-
-    drafts = [tokenizer.encode(":")[0], -1, -1, -1]
-
-    first = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: drafts},
-    )
-    second = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: drafts},
-    )
-
-    assert first is not None and second is not None
-    assert (first == second).all()
-    assert not grammar.is_terminated()
-
-
-@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
-def test_bonus_position_constrained_after_invalid_drafts(backend):
-    """Regression for #44006: bonus row stays constrained after -1 padding."""
-    tokenizer, manager, request, prompt = _make_manager_and_request(
-        backend, prompt_str='{"a"'
-    )
-    grammar = request.structured_output_request.grammar
-
-    assert grammar.accept_tokens(request.request_id, prompt)
-
-    valid = tokenizer.encode(":")[0]
-    drafts = [valid, -1, -1, -1]
-    bitmask = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: drafts},
-    )
-    assert bitmask is not None
-    assert bitmask.shape[0] == len(drafts) + 1
-
-    assert not (bitmask[-1] == -1).all()
-    assert not grammar.is_terminated()
-
-
-@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
-def test_bitmask_constrained_when_reasoning_ends_midwindow(backend):
-    """Drafts after a mid-window reasoning-end marker stay constrained."""
-    tokenizer, manager, request, prompt = _make_manager_and_request(backend)
-    grammar = request.structured_output_request.grammar
-
-    assert grammar.accept_tokens(request.request_id, prompt)
-
-    marker = tokenizer.encode("\n")[0]
-
-    class StubReasoner:
-        def __init__(self, *_, **__):
-            self.end_token_id = marker
-
-        def is_reasoning_end(self, input_ids):
-            return marker in list(input_ids)
-
-        def is_reasoning_end_streaming(self, input_ids, delta_ids):
-            return marker in list(delta_ids)
-
-    manager.reasoner_cls = StubReasoner
-    request.structured_output_request.reasoner = StubReasoner()
-    request.structured_output_request.reasoning_ended = False
-
-    pre = tokenizer.encode(" ")[0]
-    post = tokenizer.encode(",")[0]
-    drafts = [pre, marker, post]
-
-    bitmask = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: drafts},
-    )
-
-    assert bitmask is not None
-    assert bitmask.shape[0] == len(drafts) + 1
-    assert (bitmask[0] == -1).all()
-    assert (bitmask[1] == -1).all()
-    assert not (bitmask[2] == -1).all()
-    assert not (bitmask[-1] == -1).all()
-    assert not grammar.is_terminated()
-
-
-@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
-def test_bitmask_post_reasoning_end_drafts_skip_grammar_advance(backend, caplog):
-    """Post-marker drafts predate the bitmask and may be grammar-invalid;
-    grammar_bitmask must skip the grammar advance instead of asserting.
-    """
-    tokenizer, manager, request, prompt = _make_manager_and_request(
-        backend, prompt_str="{"
-    )
-    grammar = request.structured_output_request.grammar
-
-    assert grammar.accept_tokens(request.request_id, prompt)
-    assert not grammar.is_terminated()
-
-    marker = tokenizer.encode("\n")[0]
-
-    class StubReasoner:
-        def __init__(self, *_, **__):
-            self.end_token_id = marker
-
-        def is_reasoning_end(self, input_ids):
-            return marker in list(input_ids)
-
-        def is_reasoning_end_streaming(self, input_ids, delta_ids):
-            return marker in list(delta_ids)
-
-    manager.reasoner_cls = StubReasoner
-    request.structured_output_request.reasoner = StubReasoner()
-    request.structured_output_request.reasoning_ended = False
-
-    pre = tokenizer.encode(" ")[0]
-    # A token that the JSON grammar would reject as the first post-marker
-    # token; without the fix grammar.accept_tokens fires the assertion.
-    invalid_post = tokenizer.encode("z")[0]
-    drafts = [pre, marker, invalid_post]
-
-    bitmask = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: drafts},
-    )
-
-    assert bitmask is not None
-    assert bitmask.shape[0] == len(drafts) + 1
-    # Post-marker position is still bitmask-constrained.
-    assert not (bitmask[2] == -1).all()
-    # Grammar must not have advanced through the unvalidated draft.
-    assert not grammar.is_terminated()
-    assert "Failed to advance FSM" not in caplog.text
-
-
-@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
-def test_validate_tokens_then_bitmask_round_trip(backend):
-    """validate_tokens -> pad with -1 -> grammar_bitmask must not assert."""
-    tokenizer, manager, request, prompt = _make_manager_and_request(backend)
-    grammar = request.structured_output_request.grammar
-
-    assert grammar.accept_tokens(request.request_id, prompt)
-
-    raw_drafts = [tokenizer.encode(",")[0], 99999, 12345, 67890]
-    accepted = grammar.validate_tokens(raw_drafts)
-    assert len(accepted) <= len(raw_drafts)
-
-    padded = accepted + [-1] * (len(raw_drafts) - len(accepted))
-    assert len(padded) == len(raw_drafts)
-
-    bitmask = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: padded},
-    )
-    assert bitmask is not None
-    assert bitmask.shape[0] == len(padded) + 1
-    assert not grammar.is_terminated()
-
-
 def test_xgrammar_accept_tokens_stops_at_termination(capfd):
     """Tokens after a terminating EOS do not reach the matcher."""
     tokenizer, _, request, prompt = _make_manager_and_request("xgrammar")
@@ -302,7 +87,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd):
 
     assert grammar.validate_tokens([eos, trailing]) == [eos]
     assert "trying to accept new token" not in capfd.readouterr().err
-    # Check matcher state directly to verify validation rolled it back.
     assert not grammar.matcher.is_terminated()
 
     assert grammar.accept_tokens(request.request_id, [eos])
@@ -312,87 +96,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd):
     assert "trying to accept new token" not in capfd.readouterr().err
 
 
-class _MarkerReasoner:
-    """Stub reasoner whose reasoning-end marker is a single fixed token."""
-
-    def __init__(self, marker: int):
-        self.marker = marker
-
-    def is_reasoning_end(self, input_ids):
-        return self.marker in list(input_ids)
-
-    def is_reasoning_end_streaming(self, input_ids, delta_ids):
-        return self.marker in list(delta_ids)
-
-
-def _setup_boundary_request(backend: str):
-    """Request with a structural-tag key and reasoning not yet ended."""
-    from vllm.v1.structured_output.backend_types import StructuredOutputOptions
-
-    tokenizer, manager, request, prompt = _make_manager_and_request(backend)
-    marker = tokenizer.encode("\n")[0]
-    structured_req = request.structured_output_request
-    # The grammar itself is JSON (cheap to build); only the key kind matters
-    # for the should_advance structural-tag branch, so pre-seed the cached
-    # property.
-    structured_req.__dict__["structured_output_key"] = (
-        StructuredOutputOptions.STRUCTURAL_TAG,
-        "",
-    )
-    manager.reasoner_cls = _MarkerReasoner
-    structured_req.reasoner = _MarkerReasoner(marker)
-    structured_req.reasoning_ended = False
-    return tokenizer, manager, request, prompt, marker
-
-
-def test_should_advance_records_reasoning_end_index():
-    """Regression for #44006 on post-#42452 main: the boundary step must
-    record where reasoning ends so the scheduler can trim before advancing.
-    """
-    tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar")
-    structured_req = request.structured_output_request
-
-    pre = tokenizer.encode(" ")[0]
-    post = tokenizer.encode("{")[0]
-    request.append_output_token_ids([pre, marker, post])
-
-    assert manager.should_advance(request)
-    assert structured_req.reasoning_ended
-    # Marker sits at absolute index len(prompt) + 1.
-    assert structured_req.reasoning_end_token_index == len(prompt) + 1
-
-
-def test_trim_reasoning_for_advance():
-    """trim drops the marker and everything before it; later steps and
-    requests without a recorded boundary pass through unchanged.
-    """
-    tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar")
-    structured_req = request.structured_output_request
-
-    pre = tokenizer.encode(" ")[0]
-    post = tokenizer.encode("{")[0]
-
-    # No boundary recorded yet: pass-through.
-    assert manager.trim_reasoning_for_advance(request, [pre]) == [pre]
-
-    # Boundary step: marker mid-step keeps only the suffix.
-    step_tokens = [pre, marker, post]
-    request.append_output_token_ids(step_tokens)
-    assert manager.should_advance(request)
-    assert manager.trim_reasoning_for_advance(request, step_tokens) == [post]
-
-    # Boundary step variant: marker last (the #44006 crash shape
-    # [198, </think>]) trims to empty -> scheduler skips accept_tokens.
-    structured_req.reasoning_end_token_index = len(request.all_token_ids) - 1
-    assert manager.trim_reasoning_for_advance(request, step_tokens) == []
-
-    # Later steps: tokens are past the boundary, returned unchanged.
-    structured_req.reasoning_end_token_index = len(prompt) + 1
-    next_step = [post, post]
-    request.append_output_token_ids(next_step)
-    assert manager.trim_reasoning_for_advance(request, next_step) == next_step
-
-
 class _EngineReasonerStub(ParserEngineReasoningAdapter):
     """Adapter-typed reasoner with a fixed end-token set and no real engine."""
 
@@ -409,7 +112,7 @@ def find_reasoning_end_offset(self, token_ids):
         for offset, token in enumerate(token_ids):
             if token in self._end_token_ids:
                 return offset
-        return None
+        return len(token_ids)
 
     def is_reasoning_end(self, input_ids):
         return any(token in self._end_token_ids for token in input_ids)
@@ -449,61 +152,3 @@ def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend):
     assert not (bitmask[2] == -1).all()
     assert reasoner.windows == [[pre, marker, post]]
     assert not grammar.is_terminated()
-
-
-@pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
-def test_bitmask_legacy_order_sensitive_predicate_probed_per_token(backend):
-    """Legacy predicates are probed one draft at a time, never on the whole
-    window. A predicate that reports which marker is *newest* (KimiK3) would
-    otherwise answer False for a window that closes and then reopens
-    reasoning, leaving the post-marker drafts and the bonus row
-    unconstrained.
-    """
-    tokenizer, manager, request, prompt = _make_manager_and_request(backend)
-    grammar = request.structured_output_request.grammar
-
-    assert grammar.accept_tokens(request.request_id, prompt)
-
-    close = tokenizer.encode("\n")[0]
-    reopen = tokenizer.encode("#")[0]
-    assert close != reopen
-
-    class NewestMarkerReasoner:
-        def __init__(self, *_, **__):
-            pass
-
-        def is_reasoning_end(self, input_ids):
-            return self.is_reasoning_end_streaming(input_ids, input_ids)
-
-        def is_reasoning_end_streaming(self, input_ids, delta_ids):
-            for token in reversed(list(delta_ids)):
-                if token == close:
-                    return True
-                if token == reopen:
-                    return False
-            return False
-
-    manager.reasoner_cls = NewestMarkerReasoner
-    request.structured_output_request.reasoner = NewestMarkerReasoner()
-    request.structured_output_request.reasoning_ended = False
-
-    pre = tokenizer.encode(" ")[0]
-    post = tokenizer.encode(",")[0]
-    drafts = [pre, close, post, reopen]
-    # Sanity: the whole window says "reopened", so only per-token probing
-    # can find the close marker.
-    assert not NewestMarkerReasoner().is_reasoning_end_streaming([], drafts)
-
-    bitmask = manager.grammar_bitmask(
-        requests={request.request_id: request},
-        structured_output_request_ids=[request.request_id],
-        scheduled_spec_decode_tokens={request.request_id: drafts},
-    )
-
-    assert bitmask is not None
-    assert bitmask.shape[0] == len(drafts) + 1
-    assert (bitmask[0] == -1).all()
-    assert (bitmask[1] == -1).all()
-    assert not (bitmask[2] == -1).all()
-    assert not (bitmask[-1] == -1).all()
-    assert not grammar.is_terminated()
diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py
index 86787b53a105..9e0052bcc59c 100644
--- a/tests/v1/spec_decode/test_rejection_sampler_utils.py
+++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py
@@ -6,7 +6,10 @@
 import pytest
 import torch
 
+from vllm.v1.watermarking import GumbelWatermarker
+from vllm.v1.watermarking.spec_decode import watermarked_rejection_sample
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
     rejection_sample,
 )
@@ -90,8 +93,7 @@ def _assert_distribution_match(
     label: str = "",
     min_expected: float = 5.0,
 ):
-    """
-    Assert sampled tokens match the target distribution via a
+    """Assert sampled tokens match the target distribution via a
     chi-squared goodness-of-fit test. This is done by computing
     observed vs expected token counts (target_probs * num_samples),
     then checking that the chi-squared statistic is below a conservative
@@ -139,6 +141,93 @@ def _assert_distribution_match(
     )
 
 
+def test_watermarked_bonus_uses_target_key():
+    target_logits = torch.randn(16, device="cuda")
+    inputs = _build_rejection_sample_inputs(
+        target_logits,
+        target_logits,
+        num_speculative_steps=1,
+        temperature=1.0,
+        num_trials=1,
+    )
+    draft_token = inputs["draft_sampled"][1]
+    contexts = torch.tensor(
+        [[10, 11], [11, draft_token.item()]], dtype=torch.int64, device="cuda"
+    )
+
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=1,
+        contexts=contexts,
+        watermarking=torch.tensor([True], device="cuda"),
+        watermarker=GumbelWatermarker(key=42, context_width=2),
+    )
+    expected = GumbelWatermarker(key=42, context_width=2).sample(
+        target_logits.unsqueeze(0), contexts[1:], lambda _: None
+    )
+
+    assert num_sampled.item() == 2
+    assert sampled[0, 1] == expected.token_ids[0]
+
+
+def test_watermarked_recovery_uses_target_key():
+    target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0], device="cuda")
+    draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda")
+    inputs = _build_rejection_sample_inputs(
+        target_logits,
+        draft_logits,
+        num_speculative_steps=1,
+        temperature=1.0,
+        num_trials=1,
+    )
+    inputs["draft_sampled"][1] = 0
+    contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda")
+
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=1,
+        contexts=contexts,
+        watermarking=torch.tensor([True], device="cuda"),
+        watermarker=GumbelWatermarker(key=42, context_width=2),
+    )
+    expected = GumbelWatermarker(key=42, context_width=2).sample(
+        target_logits.unsqueeze(0), contexts[:1], lambda _: None
+    )
+
+    assert num_sampled.item() == 1
+    assert sampled[0, 0] == expected.token_ids[0]
+
+
+def test_watermarked_recovery_supports_smaller_draft_vocabulary():
+    target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0, 100.0], device="cuda")
+    draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda")
+    inputs = _build_rejection_sample_inputs(
+        target_logits[:4],
+        draft_logits,
+        num_speculative_steps=1,
+        temperature=1.0,
+        num_trials=1,
+    )
+    inputs["target_logits"] = target_logits.unsqueeze(0).expand(2, -1).contiguous()
+    inputs["draft_sampled"][1] = 0
+    contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda")
+    watermarker = GumbelWatermarker(key=42, context_width=2)
+
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=1,
+        contexts=contexts,
+        watermarking=torch.tensor([True], device="cuda"),
+        watermarker=watermarker,
+    )
+    expected = watermarker.sample(
+        target_logits[:4].unsqueeze(0), contexts[:1], lambda _: None
+    )
+
+    assert num_sampled.item() == 1
+    assert sampled[0, 0] == expected.token_ids[0]
+
+
 @pytest.mark.parametrize(
     "num_speculative_steps,temperature",
     [
@@ -152,8 +241,7 @@ def _assert_distribution_match(
 def test_stochastic_rejection_sample(
     num_speculative_steps: int, temperature: float, draft_logits_dtype: torch.dtype
 ):
-    """
-    Verify that rejection sampling produces the target distribution.
+    """Verify that rejection sampling produces the target distribution.
     This is done by simulating many independent trials of speculative
     decoding (from a fixed target and draft distribution). We then
     run rejection sample on all of the trials (requests), and verify
@@ -163,7 +251,6 @@ def test_stochastic_rejection_sample(
     Parametrized over the draft-logits dtype: storing them in the draft head's
     dtype must not bias the output distribution.
     """
-
     torch.manual_seed(42)
     device = "cuda"
     num_trials = 10 * VOCAB_SIZE
@@ -293,11 +380,9 @@ def test_gumbel_drafted_rejection_sample_is_unbiased(num_speculative_steps: int)
 
 @pytest.mark.parametrize("num_speculative_steps", [1, 3])
 def test_greedy_rejection_sample(num_speculative_steps: int):
-    """
-    Verify that greedy (temperature=0) always outputs the target argmax
+    """Verify that greedy (temperature=0) always outputs the target argmax
     at every accepted position.
     """
-
     torch.manual_seed(42)
     device = "cuda"
     num_trials = 10 * VOCAB_SIZE
@@ -345,8 +430,7 @@ def test_synthetic_rejection_sample(
     temperature: float,
     unconditional_rates: list[float],
 ):
-    """
-    Verify that synthetic rejection sampling produces the expected
+    """Verify that synthetic rejection sampling produces the expected
     per-position acceptance rates. The unconditional rate at position i
     is P(all draft steps 0..i accepted) = product(conditional_rates[0:i+1]).
     This is approximately mean(num accepted >= i + 1) over many trials.
@@ -599,14 +683,12 @@ def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool):
 def test_block_verification_rejection_sample(
     num_speculative_steps: int, temperature: float, has_draft_logits: bool
 ):
-    """
-    Verify that block verification (Sun et al.) preserves the target
+    """Verify that block verification (Sun et al.) preserves the target
     distribution at every accepted position, for both the full draft-logits
     case and the one-hot (no draft logits) case. Block verification changes
     *which* prefix is accepted, but the marginal of every output position must
     still match the target distribution p(x).
     """
-
     torch.manual_seed(42)
     device = "cuda"
     num_trials = 10 * VOCAB_SIZE
@@ -643,12 +725,10 @@ def test_block_verification_rejection_sample(
 
 @pytest.mark.parametrize("num_speculative_steps", [3, 5])
 def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int):
-    """
-    Block verification is designed to accept at least as long a prefix as
+    """Block verification is designed to accept at least as long a prefix as
     token verification in expectation. Verify the mean accepted length is no
     worse than the standard method on the same inputs.
     """
-
     torch.manual_seed(0)
     device = "cuda"
     num_trials = 20 * VOCAB_SIZE
@@ -747,3 +827,192 @@ def test_chunked_requests_match_full_batch(has_draft_logits: bool):
     steps = torch.arange(num_speculative_steps + 1, device=device)
     valid = steps.unsqueeze(0) < num_sampled.unsqueeze(1)
     assert torch.equal(chunked_sampled[valid], sampled[valid])
+
+
+def _residual_recovery_logits(
+    inputs: dict,
+    num_sampled: torch.Tensor,
+    vocab_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    target_logits = inputs["target_logits"]
+    draft_logits = inputs["draft_logits"]
+    cu_num_logits = inputs["cu_num_logits"]
+    resample_rows = cu_num_logits[:-1].to(torch.int64) + num_sampled.to(torch.int64) - 1
+    req_state_indices = inputs["expanded_idx_mapping"][resample_rows].to(torch.int64)
+    temperatures = inputs["temperature"][req_state_indices]
+
+    is_bonus = resample_rows == cu_num_logits[1:] - 1
+    next_rows = (resample_rows + 1).clamp_max(inputs["draft_sampled"].shape[0] - 1)
+    needs_residual = ~is_bonus & (inputs["draft_sampled"][next_rows] >= 0)
+
+    target_rows = target_logits[resample_rows][:, :vocab_size].float()
+    target_rows = torch.where(target_rows.isnan(), float("-inf"), target_rows)
+    if draft_logits is None:
+        rejected = inputs["draft_sampled"][next_rows].clamp_min(0)
+        one_hot = torch.arange(vocab_size, device=target_rows.device).unsqueeze(
+            0
+        ) == rejected.unsqueeze(-1)
+        residual_logits = torch.where(one_hot, float("-inf"), target_rows)
+    else:
+        draft_steps = (num_sampled.to(torch.int64) - 1).clamp_max(
+            draft_logits.shape[1] - 1
+        )
+        draft_rows = draft_logits[req_state_indices, draft_steps][
+            :, :vocab_size
+        ].float()
+        draft_rows = torch.where(draft_rows.isnan(), float("-inf"), draft_rows)
+        draft_rows = draft_rows / temperatures.unsqueeze(-1)
+        target_log_probs = torch.log_softmax(target_rows, dim=-1)
+        draft_log_probs = torch.log_softmax(draft_rows, dim=-1)
+        ratio = torch.exp(draft_log_probs - target_log_probs)
+        residual_logits = torch.where(
+            ratio < 1,
+            target_log_probs + torch.log1p(-ratio.clamp_max(1)),
+            float("-inf"),
+        )
+    recovery = torch.where(needs_residual.unsqueeze(-1), residual_logits, target_rows)
+    return recovery, resample_rows
+
+
+def _watermark_inputs(
+    num_trials: int,
+    num_speculative_steps: int,
+    context_width: int,
+    vocab_size: int,
+    one_hot_draft: bool,
+    temperature: float = 1.0,
+) -> tuple[dict, torch.Tensor]:
+    device = "cuda"
+    target_logits_1d = torch.randn(vocab_size, device=device) / temperature
+    draft_logits_1d = torch.randn(vocab_size, device=device)
+    inputs = _build_rejection_sample_inputs(
+        target_logits_1d,
+        draft_logits_1d,
+        num_speculative_steps,
+        temperature=temperature,
+        num_trials=num_trials,
+    )
+    if one_hot_draft:
+        inputs["draft_logits"] = None
+    inputs["draft_sampled"].view(num_trials, num_speculative_steps + 1)[::5, 1] = -1
+    num_logits = inputs["target_logits"].shape[0]
+    contexts = torch.randint(
+        0, vocab_size, (num_logits, context_width), dtype=torch.int64, device=device
+    )
+    contexts[::7, 0] = -1
+    return inputs, contexts
+
+
+@pytest.mark.parametrize("one_hot_draft", [False, True])
+# Cover one full block and a partial multi-block vocabulary.
+@pytest.mark.parametrize("vocab_size", [1024, 3500])
+def test_watermarked_recovery_matches_philox_gumbel_sample(
+    one_hot_draft: bool, vocab_size: int
+):
+    torch.manual_seed(7)
+    key = 0xDA39A3EE5E6B4B0D
+    context_width = 4
+    inputs, contexts = _watermark_inputs(
+        num_trials=128,
+        num_speculative_steps=2,
+        context_width=context_width,
+        vocab_size=vocab_size,
+        one_hot_draft=one_hot_draft,
+    )
+    watermarker = GumbelWatermarker(key=key, context_width=context_width)
+
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=2,
+        contexts=contexts,
+        watermarking=torch.ones(128, dtype=torch.bool, device="cuda"),
+        watermarker=watermarker,
+    )
+
+    recovery, resample_rows = _residual_recovery_logits(inputs, num_sampled, vocab_size)
+    expected = philox_gumbel_sample(recovery, contexts[resample_rows], key)
+    recovered = sampled[torch.arange(128, device="cuda"), num_sampled.long() - 1]
+    torch.testing.assert_close(recovered, expected, rtol=0, atol=0)
+
+    is_bonus = resample_rows == inputs["cu_num_logits"][1:] - 1
+    next_rows = (resample_rows + 1).clamp_max(inputs["draft_sampled"].shape[0] - 1)
+    is_placeholder = ~is_bonus & (inputs["draft_sampled"][next_rows] < 0)
+    assert int((~is_bonus & ~is_placeholder).sum()) > 0
+    assert int(is_placeholder.sum()) > 0
+    if not one_hot_draft:
+        # One-hot drafts almost never reach the bonus row.
+        assert int(is_bonus.sum()) > 0
+
+
+def test_watermarked_recovery_leaves_disabled_rows_on_the_stock_draw():
+    torch.manual_seed(11)
+    num_trials = 96
+    vocab_size = 512
+    inputs, contexts = _watermark_inputs(
+        num_trials=num_trials,
+        num_speculative_steps=3,
+        context_width=2,
+        vocab_size=vocab_size,
+        one_hot_draft=False,
+    )
+    watermarking = torch.zeros(num_trials, dtype=torch.bool, device="cuda")
+    watermarking[::2] = True
+    inputs["temperature"][:8] = 0.0
+
+    stock_sampled, stock_num_sampled = rejection_sample(
+        **inputs, num_speculative_steps=3
+    )
+    sampled, num_sampled = watermarked_rejection_sample(
+        **inputs,
+        num_speculative_steps=3,
+        contexts=contexts,
+        watermarking=watermarking,
+        watermarker=GumbelWatermarker(key=99, context_width=2),
+    )
+
+    torch.testing.assert_close(num_sampled, stock_num_sampled, rtol=0, atol=0)
+    rows = torch.arange(num_trials, device="cuda")
+    positions = num_sampled.long() - 1
+    disabled = ~watermarking | (inputs["temperature"] == 0)
+    torch.testing.assert_close(
+        sampled[rows[disabled], positions[disabled]],
+        stock_sampled[rows[disabled], positions[disabled]],
+        rtol=0,
+        atol=0,
+    )
+    for req in range(num_trials):
+        n = int(num_sampled[req]) - 1
+        torch.testing.assert_close(
+            sampled[req, :n], stock_sampled[req, :n], rtol=0, atol=0
+        )
+    assert int(disabled.sum()) > 0 and int((~disabled).sum()) > 0
+
+
+def test_watermarked_recovery_is_unchanged_by_fp64_gumbel():
+    torch.manual_seed(13)
+    inputs, contexts = _watermark_inputs(
+        num_trials=64,
+        num_speculative_steps=2,
+        context_width=4,
+        vocab_size=2048,
+        one_hot_draft=False,
+    )
+    watermarker = GumbelWatermarker(key=2**63 + 12345, context_width=4)
+    kwargs = dict(
+        num_speculative_steps=2,
+        contexts=contexts,
+        watermarking=torch.ones(64, dtype=torch.bool, device="cuda"),
+        watermarker=watermarker,
+    )
+    sampled32, num_sampled32 = watermarked_rejection_sample(**inputs, **kwargs)
+    sampled64, num_sampled64 = watermarked_rejection_sample(
+        **inputs, **kwargs, use_fp64=True
+    )
+    rows = torch.arange(64, device="cuda")
+    torch.testing.assert_close(num_sampled64, num_sampled32, rtol=0, atol=0)
+    torch.testing.assert_close(
+        sampled64[rows, num_sampled64.long() - 1],
+        sampled32[rows, num_sampled32.long() - 1],
+        rtol=0,
+        atol=0,
+    )
diff --git a/tests/v1/spec_decode/test_speculators_correctness.py b/tests/v1/spec_decode/test_speculators_correctness.py
index 5a92fe00f3ff..4cc66d5cd3c2 100644
--- a/tests/v1/spec_decode/test_speculators_correctness.py
+++ b/tests/v1/spec_decode/test_speculators_correctness.py
@@ -134,8 +134,7 @@ def print_spec_decode_stats(stats: dict) -> None:
 
 @pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
 def test_speculators_model(vllm_runner, example_prompts, monkeypatch, config):
-    """
-    Test speculators model properly initializes speculative decoding.
+    """Test speculators model properly initializes speculative decoding.
 
     Verifies:
     1. Speculative config is automatically initialized from speculators config
@@ -185,8 +184,7 @@ def test_speculators_model(vllm_runner, example_prompts, monkeypatch, config):
 @large_gpu_mark(min_gb=40)
 @pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
 def test_speculators_correctness(monkeypatch, config):
-    """
-    E2E correctness test via the speculators auto-detect path.
+    """E2E correctness test via the speculators auto-detect path.
 
     Evaluates GSM8k accuracy to ensure the speculators-format model produces
     correct outputs, and checks that acceptance length does not collapse under
diff --git a/tests/v1/spec_decode/test_speculators_eagle3.py b/tests/v1/spec_decode/test_speculators_eagle3.py
index 9a252cfffc8f..66736cce87b8 100644
--- a/tests/v1/spec_decode/test_speculators_eagle3.py
+++ b/tests/v1/spec_decode/test_speculators_eagle3.py
@@ -32,8 +32,7 @@
 def test_eagle3_speculators_model(
     vllm_runner, example_prompts, model_path, monkeypatch
 ):
-    """
-    Test Eagle3 speculators models properly initialize speculative decoding.
+    """Test Eagle3 speculators models properly initialize speculative decoding.
 
     This test verifies:
     1. Eagle3 support is detected for the model
diff --git a/tests/v1/spec_decode/test_vocab_mapping.py b/tests/v1/spec_decode/test_vocab_mapping.py
index b8773c5a5562..25c99ad04981 100644
--- a/tests/v1/spec_decode/test_vocab_mapping.py
+++ b/tests/v1/spec_decode/test_vocab_mapping.py
@@ -2,9 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
+import torch
 from transformers import AutoTokenizer
 
-from vllm.v1.spec_decode.vocab_mapping import _detect_space_prefix
+import vllm.utils.gpu_sync_debug as gsd
+from vllm.utils.gpu_sync_debug import with_gpu_sync_check
+from vllm.v1.spec_decode.vocab_mapping import VocabMapping, _detect_space_prefix
 
 
 @pytest.mark.parametrize(
@@ -46,3 +49,79 @@ def encode(self, text, **kwargs):
 
     result = _detect_space_prefix(EmptyTokenizer())
     assert result == ("Ġ", "▁")
+
+
+class _FakeTokenizer:
+    """Minimal tokenizer exposing only what ``VocabMapping`` needs."""
+
+    def __init__(self, vocab, unk_token_id=None, eos_token_id=None):
+        self._vocab = dict(vocab)
+        self.unk_token_id = unk_token_id
+        self.eos_token_id = eos_token_id
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def encode(self, text, **kwargs):
+        raise RuntimeError("no encode")
+
+    def convert_ids_to_tokens(self, ids):
+        raise RuntimeError("no convert")
+
+
+@pytest.fixture(params=["cpu", "cuda"])
+def vocab_mapping(request):
+    device = request.param
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+
+    # Intersection: {a, b}. x/y/z are target-only, u/v are draft-only
+    target = _FakeTokenizer(
+        vocab={"a": 0, "b": 1, "x": 2, "y": 3, "z": 4}, unk_token_id=0
+    )
+    draft = _FakeTokenizer(vocab={"a": 0, "b": 1, "u": 2, "v": 3}, unk_token_id=2)
+    return VocabMapping(target, draft, 5, 4, device)
+
+
+@pytest.mark.parametrize(
+    "target_ids,expected",
+    [
+        ([0, 1, 2, 3, 4], [0, 1, 2, 2, 2]),  # mixed
+        ([0, 1], [0, 1]),  # all present
+        ([2, 3, 4], [2, 2, 2]),  # all missing
+    ],
+)
+def test_map_target_to_draft_ids(monkeypatch, vocab_mapping, target_ids, expected):
+    device = vocab_mapping.device
+    monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error" if device == "cuda" else None)
+    monkeypatch.setattr(gsd, "_sync_check_enabled", True)
+
+    ids = torch.tensor(target_ids, dtype=torch.int32, device=device)
+    out = with_gpu_sync_check(lambda: vocab_mapping.map_target_to_draft_ids(ids))()
+
+    assert out.tolist() == expected
+    assert out.dtype == torch.int32, f"expected torch.int32, got {out.dtype}"
+    assert out.device.type == device, f"expected {device} device, got {out.device}"
+    assert (out != -1).all()
+
+
+@pytest.mark.parametrize(
+    "draft_ids,expected",
+    [
+        ([0, 1, 2, 3], [0, 1, 0, 0]),  # mixed
+        ([0, 1], [0, 1]),  # all present
+        ([2, 3], [0, 0]),  # all missing
+    ],
+)
+def test_map_draft_to_target_ids(monkeypatch, vocab_mapping, draft_ids, expected):
+    device = vocab_mapping.device
+    monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error" if device == "cuda" else None)
+    monkeypatch.setattr(gsd, "_sync_check_enabled", True)
+
+    ids = torch.tensor(draft_ids, dtype=torch.int32, device=device)
+    out = with_gpu_sync_check(lambda: vocab_mapping.map_draft_to_target_ids(ids))()
+
+    assert out.tolist() == expected
+    assert out.dtype == torch.int32, f"expected torch.int32, got {out.dtype}"
+    assert out.device.type == device, f"expected {device} device, got {out.device}"
+    assert (out != -1).all()
diff --git a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py
index a8f7221da81b..9457c66050db 100644
--- a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py
+++ b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py
@@ -23,7 +23,6 @@
 @pytest.fixture
 def mock_model_runner_with_input_batch():
     """Create a mock GPUModelRunner with a real InputBatch for e2e testing."""
-
     runner = Mock(spec=GPUModelRunner)
     runner.uses_mrope = False
     runner.requests = {}
diff --git a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py
index 31c1a6f4846a..c5f8090f943f 100644
--- a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py
+++ b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py
@@ -23,7 +23,6 @@
 @pytest.fixture
 def mock_model_runner_with_req_states():
     """Create a mock MRv2 GPUModelRunner with a real RequestState."""
-
     runner = Mock(spec=GPUModelRunner)
     runner.req_states = RequestState(
         max_num_reqs=10,
diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py
index 822b4e49da34..0d259c1c6948 100644
--- a/tests/v1/streaming_input/test_scheduler_streaming.py
+++ b/tests/v1/streaming_input/test_scheduler_streaming.py
@@ -53,6 +53,7 @@ def create_scheduler() -> Scheduler:
     vllm_config.model_config = MagicMock()
     vllm_config.model_config.skip_tokenizer_init = True
     vllm_config.model_config.is_multimodal_model = False
+    vllm_config.model_config.supports_multimodal_inputs = False
     vllm_config.model_config.is_encoder_decoder = False
     vllm_config.model_config.is_diffusion = False
     vllm_config.model_config.max_model_len = 1024
@@ -315,8 +316,7 @@ def test_update_request_as_session_with_output_tokens(self):
         assert num_new_tokens == 2
 
     def test_streaming_e2e_lifecycle(self):
-        """
-        Comprehensive integration test covering complete streaming request lifecycle
+        """Comprehensive integration test covering complete streaming request lifecycle
         including scheduler state management and aliasing bug prevention.
 
         FULL LIFECYCLE:
diff --git a/tests/v1/structured_output/test_mistral_common_tokenizer.py b/tests/v1/structured_output/test_mistral_common_tokenizer.py
deleted file mode 100644
index 74d7cb085cc5..000000000000
--- a/tests/v1/structured_output/test_mistral_common_tokenizer.py
+++ /dev/null
@@ -1,102 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Structured output with a `MistralCommonBackend` tokenizer."""
-
-import pytest
-from transformers import AutoTokenizer, MistralCommonBackend
-
-from vllm.config import StructuredOutputsConfig, VllmConfig
-from vllm.tokenizers import get_tokenizer
-from vllm.tokenizers.mistral import MistralTokenizer
-from vllm.v1.structured_output.backend_outlines import OutlinesBackend
-from vllm.v1.structured_output.backend_types import StructuredOutputOptions
-from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
-from vllm.v1.structured_output.utils import (
-    _reduced_vocabulary,
-    maybe_wrap_mistral_common_tokenizer,
-)
-
-TOKENIZER = "mistralai/Mistral-Nemo-Instruct-2407"
-JSON_SCHEMA = (
-    '{"type": "object", "properties": {"x": {"type": "integer"}}, '
-    '"required": ["x"], "additionalProperties": false}'
-)
-VALID_DOCUMENT = '{"x": 1}'
-
-BACKENDS = [("xgrammar", XgrammarBackend), ("outlines", OutlinesBackend)]
-
-
-@pytest.fixture(scope="module")
-def loaded_tokenizer():
-    tokenizer = get_tokenizer(tokenizer_name=TOKENIZER, tokenizer_mode="hf")
-    assert isinstance(tokenizer, MistralCommonBackend)
-    return tokenizer
-
-
-@pytest.fixture(scope="module")
-def wrapped_tokenizer(loaded_tokenizer):
-    return maybe_wrap_mistral_common_tokenizer(loaded_tokenizer)
-
-
-def _encode(tokenizer, text):
-    special = set(tokenizer.all_special_ids)
-    return [tid for tid in tokenizer.encode(text) if tid not in special]
-
-
-def _backend(cls, name, tokenizer):
-    vllm_config = VllmConfig(
-        structured_outputs_config=StructuredOutputsConfig(backend=name)
-    )
-    return cls(vllm_config, tokenizer=tokenizer, vocab_size=tokenizer.vocab_size)
-
-
-def test_wrapping_produces_a_mistral_tokenizer(loaded_tokenizer, wrapped_tokenizer):
-    assert isinstance(wrapped_tokenizer, MistralTokenizer)
-    assert len(wrapped_tokenizer.vocab) == loaded_tokenizer.vocab_size
-
-
-def test_other_tokenizers_are_left_alone():
-    tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
-    assert maybe_wrap_mistral_common_tokenizer(tokenizer) is tokenizer
-
-
-def test_wrapping_fixes_the_outlines_vocabulary(loaded_tokenizer, wrapped_tokenizer):
-    unwrapped_vocab = _reduced_vocabulary(loaded_tokenizer)
-    wrapped_vocab = _reduced_vocabulary(wrapped_tokenizer)
-
-    assert unwrapped_vocab.keys() == wrapped_vocab.keys()
-
-    mismapped = {
-        token for token, ids in unwrapped_vocab.items() if ids != wrapped_vocab[token]
-    }
-    assert mismapped
-
-    assert len({tuple(wrapped_vocab[t]) for t in mismapped}) == len(mismapped)
-    assert len({tuple(unwrapped_vocab[t]) for t in mismapped}) < len(mismapped)
-
-
-@pytest.mark.parametrize("name,cls", BACKENDS)
-def test_backend_accepts_a_valid_document(wrapped_tokenizer, name, cls):
-    backend = _backend(cls, name, wrapped_tokenizer)
-    grammar = backend.compile_grammar(StructuredOutputOptions.JSON, JSON_SCHEMA)
-
-    for token in _encode(wrapped_tokenizer, VALID_DOCUMENT):
-        assert grammar.accept_tokens("req", [token])
-
-
-@pytest.mark.parametrize("name,cls", BACKENDS)
-def test_backend_rejects_an_invalid_token(wrapped_tokenizer, name, cls):
-    backend = _backend(cls, name, wrapped_tokenizer)
-    grammar = backend.compile_grammar(StructuredOutputOptions.JSON, JSON_SCHEMA)
-
-    assert not grammar.accept_tokens("req", [_encode(wrapped_tokenizer, "]")[-1]])
-
-
-def test_xgrammar_vocab_size_matches_the_tokenizer(wrapped_tokenizer):
-    backend = _backend(XgrammarBackend, "xgrammar", wrapped_tokenizer)
-    assert backend.vocab_size == len(wrapped_tokenizer.vocab)
-
-
-def test_xgrammar_rejects_the_unwrapped_tokenizer(loaded_tokenizer):
-    with pytest.raises(ValueError, match="Unsupported tokenizer type"):
-        _backend(XgrammarBackend, "xgrammar", loaded_tokenizer)
diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py
deleted file mode 100644
index a57e5da72b26..000000000000
--- a/tests/v1/structured_output/test_reasoning_structured_output.py
+++ /dev/null
@@ -1,557 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Unit tests for reasoning-aware structured output functionality (PR #25515)."""
-
-from collections.abc import Sequence
-from unittest.mock import Mock
-
-import pytest
-
-from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
-from vllm.v1.request import Request
-from vllm.v1.structured_output import StructuredOutputManager
-from vllm.v1.structured_output.backend_types import StructuredOutputOptions
-
-
-class UntouchableTokenIds(Sequence[int]):
-    """A token history that fails the test if anything reads it."""
-
-    def __init__(self, length: int) -> None:
-        self._length = length
-
-    def __len__(self) -> int:
-        return self._length
-
-    def __getitem__(self, item):
-        raise AssertionError("token history must not be read")
-
-    def __iter__(self):
-        raise AssertionError("token history must not be iterated")
-
-
-def make_qwen3_reasoner():
-    """A real engine-based reasoning adapter over a tiny mock vocabulary."""
-    from tests.parser.engine.conftest import make_mock_tokenizer
-    from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter
-
-    tokenizer = make_mock_tokenizer(
-        {"<think>": 200, "</think>": 201, "<tool_call>": 202, "</tool_call>": 203}
-    )
-    return Qwen3ParserReasoningAdapter(tokenizer)
-
-
-class MockReasoner:
-    def __init__(self, tokenizer):
-        self.is_reasoning_end = Mock(return_value=False)
-        self.is_reasoning_end_streaming = Mock(return_value=False)
-
-
-class TestReasoningStructuredOutput:
-    """Test reasoning-aware structured output functionality."""
-
-    @pytest.fixture
-    def mock_model_config(self):
-        """Create a mock ModelConfig."""
-        config = Mock(spec=ModelConfig)
-        config.skip_tokenizer_init = True  # Skip tokenizer init to avoid network calls
-        config.get_vocab_size = Mock(return_value=50000)
-        # Add missing runner_type attribute that tokenizer initialization expects
-        config.runner_type = "generate"
-        # Add other attributes that tokenizer initialization might need
-        config.tokenizer = "test-tokenizer"
-        config.tokenizer_mode = "auto"
-        config.trust_remote_code = False
-        config.tokenizer_revision = None
-        return config
-
-    @pytest.fixture
-    def mock_scheduler_config(self):
-        """Create a mock SchedulerConfig."""
-        config = Mock(spec=SchedulerConfig)
-        config.max_num_seqs = 128
-        return config
-
-    @pytest.fixture
-    def mock_vllm_config(self, mock_model_config, mock_scheduler_config):
-        """Create a mock VllmConfig."""
-        config = Mock(spec=VllmConfig)
-        config.model_config = mock_model_config
-        config.scheduler_config = mock_scheduler_config
-        config.structured_outputs_config = Mock()
-        config.structured_outputs_config.reasoning_parser = None
-        config.structured_outputs_config.enable_in_reasoning = False
-        config.speculative_config = None
-        return config
-
-    @pytest.fixture
-    def mock_request_with_structured_output(self):
-        """Create a mock request with structured output."""
-        request = Mock(spec=Request)
-        request.structured_output_request = Mock()
-        request.structured_output_request.reasoning_ended = None
-        request.structured_output_request.grammar = Mock()
-        request.structured_output_request.reasoning_parser_kwargs = None
-        request.structured_output_request.reasoner = None
-        request.structured_output_request.grammar.is_terminated = Mock(
-            return_value=False
-        )
-        request.use_structured_output = True
-        request.prompt_token_ids = [1, 2, 3, 4, 5]
-        request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8]
-        request.num_computed_tokens = 5
-        request.num_output_placeholders = 0
-        request.request_id = "mock_req"
-        return request
-
-    @pytest.fixture
-    def manager_with_reasoner(self, mock_vllm_config):
-        manager = StructuredOutputManager(mock_vllm_config)
-        manager.reasoner_cls = MockReasoner
-        manager.tokenizer = Mock()
-        return manager
-
-    def test_should_fill_bitmask_with_enable_in_reasoning(
-        self, manager_with_reasoner, mock_request_with_structured_output
-    ):
-        """Test should_fill_bitmask when enable_in_reasoning is True."""
-        manager_with_reasoner.enable_in_reasoning = True
-
-        # Should always return True when enable_in_reasoning is enabled
-        result = manager_with_reasoner.should_fill_bitmask(
-            mock_request_with_structured_output
-        )
-        assert result is True
-        assert (
-            mock_request_with_structured_output.structured_output_request.reasoner
-            is None
-        )
-
-    def test_should_fill_bitmask_reasoning_already_ended(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """An active grammar does not need a request-local reasoner."""
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = True
-
-        result = manager_with_reasoner.should_fill_bitmask(
-            mock_request_with_structured_output
-        )
-
-        assert result is True
-        assert structured_req.reasoner is None
-
-    def test_grammar_bitmask_skips_reasoner_when_already_active(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Bitmask generation skips reasoning-boundary detection."""
-        manager_with_reasoner.vllm_config.num_speculative_tokens = 1
-        manager_with_reasoner._grammar_bitmask = Mock()
-        manager_with_reasoner._grammar_bitmask.shape = (1, 1)
-        expected_bitmask = Mock()
-        manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask
-        manager_with_reasoner._fill_bitmasks = Mock()
-
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = True
-        request_id = mock_request_with_structured_output.request_id
-
-        bitmask = manager_with_reasoner.grammar_bitmask(
-            requests={request_id: mock_request_with_structured_output},
-            structured_output_request_ids=[request_id],
-            scheduled_spec_decode_tokens={},
-        )
-
-        assert bitmask is expected_bitmask
-        assert structured_req.reasoner is None
-
-    def test_should_fill_bitmask_without_enable_in_reasoning(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Test should_fill_bitmask when enable_in_reasoning is False."""
-        # Keep enable_in_reasoning as False (default)
-        config = manager_with_reasoner.vllm_config.structured_outputs_config
-        assert config.enable_in_reasoning is False
-
-        result = manager_with_reasoner.should_fill_bitmask(
-            mock_request_with_structured_output
-        )
-
-        # Should set reasoning_ended and return its value
-        assert (
-            mock_request_with_structured_output.structured_output_request.reasoning_ended
-            is False
-        )
-        assert result is False
-
-    def test_should_fill_bitmask_no_reasoner(
-        self, mock_vllm_config, mock_request_with_structured_output
-    ):
-        """Test should_fill_bitmask when no reasoner is configured."""
-        manager = StructuredOutputManager(mock_vllm_config)
-
-        result = manager.should_fill_bitmask(mock_request_with_structured_output)
-
-        # Should default to True when no reasoner
-        assert result is True
-
-    def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs(
-        self, mock_vllm_config, mock_request_with_structured_output
-    ):
-        """Test request-level parser kwargs override the default reasoner."""
-
-        class KwargReasoner:
-            def __init__(self, tokenizer, chat_template_kwargs=None):
-                self.chat_template_kwargs = chat_template_kwargs or {}
-
-            def is_reasoning_end(self, input_ids):
-                return not self.chat_template_kwargs.get("enable_thinking", False)
-
-        manager = StructuredOutputManager(mock_vllm_config)
-        manager.reasoner_cls = KwargReasoner
-        manager.tokenizer = Mock()
-
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_parser_kwargs = {
-            "chat_template_kwargs": {"enable_thinking": True}
-        }
-
-        result = manager.should_fill_bitmask(mock_request_with_structured_output)
-
-        assert result is False
-        assert (
-            mock_request_with_structured_output.structured_output_request.reasoner
-            is not None
-        )
-
-    def test_should_advance_with_enable_in_reasoning(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Test should_advance when enable_in_reasoning is True."""
-        # Enable enable_in_reasoning
-        manager_with_reasoner.enable_in_reasoning = True
-
-        # Should always return True when enable_in_reasoning is enabled
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
-        assert result is True
-        assert (
-            mock_request_with_structured_output.structured_output_request.reasoner
-            is None
-        )
-
-    def test_should_advance_reasoning_not_ended(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Test should_advance when reasoning has not ended."""
-        # Set reasoning as not ended
-        (
-            mock_request_with_structured_output.structured_output_request
-        ).reasoning_ended = False
-
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
-
-        # Should return False since reasoning hasn't ended
-        assert result is False
-
-    def test_should_advance_reasoning_just_ended(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Test should_advance when reasoning ends in current step."""
-        # Set reasoning as not ended initially, but ends in this step
-        (
-            mock_request_with_structured_output.structured_output_request
-        ).reasoning_ended = False
-        reasoner = MockReasoner(tokenizer=Mock())
-        reasoner.is_reasoning_end_streaming.return_value = True
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoner = reasoner
-
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
-
-        # The scheduler trims the reasoning prefix before advancing the grammar.
-        assert (
-            mock_request_with_structured_output.structured_output_request.reasoning_ended
-            is True
-        )
-        assert result is True
-
-    def test_should_advance_reasoning_already_ended(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Test should_advance when reasoning has already ended."""
-        # Set reasoning as already ended
-        (
-            mock_request_with_structured_output.structured_output_request
-        ).reasoning_ended = True
-
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
-
-        # Should return True since reasoning has ended
-        assert result is True
-        assert (
-            mock_request_with_structured_output.structured_output_request.reasoner
-            is None
-        )
-
-    def test_should_advance_uses_new_token_ids_when_provided(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Regression for #43388: when caller passes new_token_ids, the
-        reasoner sees the exact multi-token delta rather than the
-        placeholder-derived window.
-        """
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-
-        end_token_id = 248069
-
-        reasoner = MockReasoner(tokenizer=Mock())
-        # Detection mirrors the real Qwen3 parser: end token in the delta.
-        reasoner.is_reasoning_end_streaming = Mock(
-            side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids)
-        )
-        structured_req.reasoner = reasoner
-
-        # Scenario from #43388: async + spec decode K=4, 4 tokens accepted
-        # but only 1 placeholder remains (some drafts were rejected).
-        # The placeholder math would yield delta=[271] and miss </think>.
-        # Passing new_token_ids must override that.
-        new_token_ids = [9, 198, end_token_id, 271]
-        mock_request_with_structured_output.all_token_ids = [
-            1,
-            2,
-            3,
-            4,
-            5,
-        ] + new_token_ids
-        mock_request_with_structured_output.num_computed_tokens = 9
-        mock_request_with_structured_output.num_output_placeholders = 1
-
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output,
-            new_token_ids=new_token_ids,
-        )
-
-        # First call to is_reasoning_end_streaming was with the full
-        # new_token_ids (not the truncated placeholder window).
-        first_call = reasoner.is_reasoning_end_streaming.call_args_list[0]
-        _, called_delta = first_call.args
-        assert list(called_delta) == new_token_ids
-
-        assert structured_req.reasoning_ended is True
-        assert result is True
-
-    def test_should_advance_without_new_token_ids_falls_back(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Backward compat: callers that don't pass new_token_ids keep
-        the original placeholder-derived delta window.
-        """
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-        reasoner = MockReasoner(tokenizer=Mock())
-        reasoner.is_reasoning_end_streaming.return_value = False
-        structured_req.reasoner = reasoner
-
-        mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5]
-        mock_request_with_structured_output.num_computed_tokens = 5
-        mock_request_with_structured_output.num_output_placeholders = 2
-
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
-
-        # placeholder window: start = 5 - 2 = 3, delta = [4, 5]
-        _, called_delta = reasoner.is_reasoning_end_streaming.call_args[0]
-        assert list(called_delta) == [4, 5]
-        assert result is False
-
-    def test_should_advance_whole_delta_fallback_pins_last_sequence_index(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """When only the whole-delta predicate fires, the boundary is the
-        last index of the sequence, even if the placeholder window
-        overshoots it and the delta is empty.
-        """
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-        reasoner = MockReasoner(tokenizer=Mock())
-        # Base-default shape: rescans the whole sequence, ignores the delta.
-        reasoner.is_reasoning_end_streaming.return_value = True
-        structured_req.reasoner = reasoner
-
-        mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5]
-        # placeholder window: start = 9 - 1 = 8, past the end -> delta = []
-        mock_request_with_structured_output.num_computed_tokens = 9
-        mock_request_with_structured_output.num_output_placeholders = 1
-
-        assert manager_with_reasoner.should_advance(mock_request_with_structured_output)
-        assert structured_req.reasoning_ended is True
-        assert structured_req.reasoning_end_token_index == 4
-
-    def test_should_advance_trims_reasoning_prefix_for_json(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """JSON uses the common trim-then-advance path at the boundary."""
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-        structured_req.structured_output_key = (
-            StructuredOutputOptions.JSON_OBJECT,
-            "{}",
-        )
-
-        marker = 248069
-
-        class MarkerReasoner:
-            def __init__(self, *_, **__):
-                pass
-
-            def is_reasoning_end_streaming(self, input_ids, delta_ids):
-                return marker in list(delta_ids)
-
-        structured_req.reasoner = MarkerReasoner()
-
-        new_token_ids = [9, 198, marker, 271, 5005]
-        mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
-
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output,
-            new_token_ids=new_token_ids,
-        )
-
-        structured_req.grammar.accept_tokens.assert_not_called()
-        assert structured_req.reasoning_ended is True
-        assert result is True
-        assert structured_req.reasoning_end_token_index == 5
-        assert manager_with_reasoner.trim_reasoning_for_advance(
-            mock_request_with_structured_output, new_token_ids
-        ) == [271, 5005]
-
-    def test_should_advance_engine_adapter_scans_only_the_delta(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Engine-based parsers locate the boundary from the delta alone."""
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-        reasoner = make_qwen3_reasoner()
-        assert reasoner.reasoning_end_token_ids == {201, 202}
-        reasoner.is_reasoning_end_streaming = Mock(
-            side_effect=AssertionError("engine path must not rescan the prefix")
-        )
-        structured_req.reasoner = reasoner
-
-        new_token_ids = [9, 198, 201, 271]
-        mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
-
-        assert manager_with_reasoner.should_advance(
-            mock_request_with_structured_output, new_token_ids=new_token_ids
-        )
-        assert structured_req.reasoning_ended is True
-        assert structured_req.reasoning_end_token_index == 5
-        assert manager_with_reasoner.trim_reasoning_for_advance(
-            mock_request_with_structured_output, new_token_ids
-        ) == [271]
-
-    def test_should_advance_engine_adapter_tool_start_ends_reasoning(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-        structured_req.reasoner = make_qwen3_reasoner()
-
-        new_token_ids = [9, 202, 271]
-        mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
-
-        assert manager_with_reasoner.should_advance(
-            mock_request_with_structured_output, new_token_ids=new_token_ids
-        )
-        assert structured_req.reasoning_end_token_index == 4
-
-    def test_should_advance_engine_adapter_no_boundary(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-        structured_req.reasoner = make_qwen3_reasoner()
-
-        new_token_ids = [9, 198, 271]
-        mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids
-
-        assert not manager_with_reasoner.should_advance(
-            mock_request_with_structured_output, new_token_ids=new_token_ids
-        )
-        assert structured_req.reasoning_ended is False
-
-    def test_should_advance_engine_adapter_never_reads_history(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Per-step cost must not depend on how much has been generated."""
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-        structured_req.reasoner = make_qwen3_reasoner()
-
-        new_token_ids = [9, 201]
-        total = (1 << 20) + len(new_token_ids)
-        mock_request_with_structured_output.all_token_ids = UntouchableTokenIds(total)
-
-        assert manager_with_reasoner.should_advance(
-            mock_request_with_structured_output, new_token_ids=new_token_ids
-        )
-        assert structured_req.reasoning_end_token_index == total - 1
-
-    def test_should_advance_engine_adapter_without_new_tokens_persists(
-        self,
-        manager_with_reasoner,
-        mock_request_with_structured_output,
-    ):
-        """Draft-validation callers persist the boundary like the legacy path."""
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_ended = False
-        structured_req.reasoner = make_qwen3_reasoner()
-
-        mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5, 201, 7]
-        # placeholder window: start = 7 - 2 = 5, delta = [201, 7]
-        mock_request_with_structured_output.num_computed_tokens = 7
-        mock_request_with_structured_output.num_output_placeholders = 2
-
-        assert manager_with_reasoner.should_advance(mock_request_with_structured_output)
-        assert structured_req.reasoning_ended is True
-        assert structured_req.reasoning_end_token_index == 5
diff --git a/tests/v1/structured_output/test_request.py b/tests/v1/structured_output/test_request.py
new file mode 100644
index 000000000000..e3eb79a49b73
--- /dev/null
+++ b/tests/v1/structured_output/test_request.py
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for StructuredOutputRequest grammar-future polling."""
+
+from concurrent.futures import Future
+
+from vllm.sampling_params import StructuredOutputsParams
+from vllm.v1.structured_output.request import StructuredOutputRequest
+
+
+def _request_with_future(fut: Future) -> StructuredOutputRequest:
+    req = StructuredOutputRequest(params=StructuredOutputsParams(regex="a+"))
+    req.grammar = fut
+    return req
+
+
+def test_pending_grammar_future_is_not_ready():
+    req = _request_with_future(Future())
+    assert not req.is_grammar_ready
+    assert req.grammar is None
+
+
+def test_grammar_compile_raising_timeout_error_is_a_failure():
+    """A compile that raised TimeoutError must be treated as a failed compile,
+    not as still-compiling (#53130)."""
+    fut: Future = Future()
+    fut.set_exception(TimeoutError("compiler timed out"))
+    req = _request_with_future(fut)
+    assert req.is_grammar_ready
+    assert isinstance(req.grammar, TimeoutError)
diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py
new file mode 100644
index 000000000000..17bda09b3efa
--- /dev/null
+++ b/tests/v1/structured_output/test_structured_output_manager.py
@@ -0,0 +1,414 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import time
+from dataclasses import dataclass
+
+import pytest
+from transformers import AutoTokenizer
+
+from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig
+from vllm.config.model import ModelConfig
+from vllm.config.speculative import SpeculativeConfig
+from vllm.sampling_params import SamplingParams, StructuredOutputsParams
+from vllm.v1.request import Request
+from vllm.v1.structured_output import StructuredOutputManager
+
+TOKENIZER = "gpt2"
+THINK_END = "\n"  # reasoning-end marker (single GPT-2 token)
+EOS = "<|eos|>"  # resolved to tokenizer.eos_token_id
+JSON_SCHEMA = '{"type": "object"}'
+BACKENDS = ("xgrammar", "guidance")
+MAX_WAIT_SECONDS = 5
+NUM_SPEC_TOKENS = 8
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    return AutoTokenizer.from_pretrained(TOKENIZER)
+
+
+@dataclass(frozen=True)
+class FlowCase:
+    """One validate -> bitmask -> accept -> post-accept bitmask scenario."""
+
+    raw_drafts: tuple[str, ...]
+    expected_row_pattern: str
+    expected_reasoning: bool | None
+    expect_terminated: bool
+    prefix: str = ""
+    expected_validated: tuple[str, ...] | None = None
+    reasoning_ended: bool | None = False
+    xfail_guidance: str | None = None
+
+
+class MockReasoner:
+    def __init__(self, tokenizer, marker: int | None = None):
+        self.marker = marker
+
+    def is_reasoning_end(self, input_ids):
+        if self.marker is None:
+            return True
+        return self.marker in list(input_ids)
+
+    def is_reasoning_end_streaming(self, input_ids, delta_ids):
+        return self.is_reasoning_end(delta_ids)
+
+
+def _single_token(tokenizer, text: str) -> int:
+    token_ids = tokenizer.encode(text)
+    assert len(token_ids) == 1, (text, token_ids)
+    return token_ids[0]
+
+
+def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]:
+    """Convert literal text (or EOS sentinel) to token IDs."""
+    return [
+        tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t)
+        for t in texts
+    ]
+
+
+def _wait_for_grammar(request: Request) -> None:
+    structured_req = request.structured_output_request
+    assert structured_req is not None
+
+    deadline = time.time() + MAX_WAIT_SECONDS
+    while not structured_req._check_grammar_completion():
+        if time.time() > deadline:
+            pytest.fail("Grammar compilation timed out")
+        time.sleep(0.01)
+
+
+def _build_harness(
+    tokenizer,
+    backend: str,
+    prefix: str = "",
+    use_reasoner: bool = True,
+    reasoning_ended: bool | None = None,
+    enable_in_reasoning: bool = False,
+    reasoning_parser_kwargs: dict | None = None,
+) -> tuple[StructuredOutputManager, Request]:
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(tokenizer=TOKENIZER),
+        device_config=DeviceConfig(device="cpu"),
+        structured_outputs_config=StructuredOutputsConfig(
+            backend=backend,
+            enable_in_reasoning=enable_in_reasoning,
+        ),
+        speculative_config=SpeculativeConfig(
+            model="[ngram]",
+            num_speculative_tokens=NUM_SPEC_TOKENS,
+        ),
+    )
+    manager = StructuredOutputManager(vllm_config)
+    if use_reasoner:
+        manager.reasoner_cls = MockReasoner
+
+    sampling_params = SamplingParams(
+        structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA)
+    )
+    sampling_params.structured_outputs._backend = backend
+    sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)
+
+    prompt_ids = tokenizer.encode(prefix) if prefix else []
+    request = Request(
+        request_id=f"{backend}-flow",
+        prompt_token_ids=prompt_ids,
+        sampling_params=sampling_params,
+        pooling_params=None,
+        reasoning_ended=reasoning_ended,
+        reasoning_parser_kwargs=reasoning_parser_kwargs,
+    )
+    manager.grammar_init(request)
+    _wait_for_grammar(request)
+
+    structured_req = request.structured_output_request
+    assert request.prompt_token_ids is not None
+    assert structured_req is not None
+    assert structured_req.grammar is not None
+    if prompt_ids:
+        assert structured_req.grammar.accept_tokens(
+            request.request_id,
+            prompt_ids,
+        )
+    return manager, request
+
+
+def _row_pattern(bitmask) -> str:
+    return "".join("U" if (row == -1).all() else "C" for row in bitmask)
+
+
+def _run_real_flow(
+    manager: StructuredOutputManager,
+    request: Request,
+    raw_drafts: list[int],
+    expected_validated: list[int],
+    expected_row_pattern: str,
+    expected_reasoning: bool | None,
+    expect_terminated: bool,
+) -> None:
+    assert len(expected_validated) <= len(raw_drafts)
+
+    validated = manager.validate_tokens(request, list(raw_drafts))
+    assert validated == expected_validated
+
+    padded = validated + [-1] * (len(raw_drafts) - len(validated))
+    bitmask = manager.grammar_bitmask(
+        requests={request.request_id: request},
+        structured_output_request_ids=[request.request_id],
+        scheduled_spec_decode_tokens={request.request_id: padded},
+    )
+    assert bitmask is not None
+    assert bitmask.shape[0] == len(raw_drafts) + 1
+    assert _row_pattern(bitmask) == expected_row_pattern
+
+    # grammar_bitmask() must rollback any speculative state it advanced.
+    assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated
+
+    structured_req = request.structured_output_request
+    assert structured_req is not None
+    grammar = structured_req.grammar
+    assert grammar is not None
+    assert not grammar.is_terminated()
+    # `bitmask[i]` is the grammar state after the first `i` scheduled tokens.
+    # These tests commit the validated prefix, so the matching post-accept
+    # state is already present in the scheduled bitmask.
+    expected_post_accept_row = bitmask[len(expected_validated)].copy()
+
+    # Mirror the scheduler flow: sampled tokens are appended before
+    # StructuredOutputManager.accept_tokens() trims them to the
+    # grammar-constrained suffix.
+    request.append_output_token_ids(list(expected_validated))
+    assert manager.accept_tokens(request, list(expected_validated))
+
+    post_accept_bitmask = manager.grammar_bitmask(
+        requests={request.request_id: request},
+        structured_output_request_ids=[request.request_id],
+        scheduled_spec_decode_tokens={},
+    )
+    assert post_accept_bitmask is not None
+    assert post_accept_bitmask.shape[0] == 1
+    assert (post_accept_bitmask[0] == expected_post_accept_row).all()
+    assert structured_req.reasoning_ended == expected_reasoning
+    assert grammar.is_terminated() is expect_terminated
+
+
+FLOW_CASES = [
+    pytest.param(
+        FlowCase(
+            raw_drafts=(" ", "z", " "),
+            expected_row_pattern="UUUU",
+            expected_reasoning=False,
+            expect_terminated=False,
+        ),
+        id="inactive",
+    ),
+    pytest.param(
+        FlowCase(
+            prefix='{"a"',
+            raw_drafts=(":", ' "', "b", '"}'),
+            expected_row_pattern="CCCCC",
+            expected_reasoning=True,
+            expect_terminated=False,
+            reasoning_ended=True,
+        ),
+        id="active_all_valid",
+    ),
+    pytest.param(
+        FlowCase(
+            prefix='{"a"',
+            raw_drafts=("z", ":", ' "', "b"),
+            expected_validated=(),
+            expected_row_pattern="CUUUU",
+            expected_reasoning=True,
+            expect_terminated=False,
+            reasoning_ended=True,
+        ),
+        id="active_first_token_invalid",
+    ),
+    pytest.param(
+        FlowCase(
+            prefix='{"a"',
+            raw_drafts=(":", "z", ' "', "b"),
+            expected_validated=(":",),
+            expected_row_pattern="CCUUU",
+            expected_reasoning=True,
+            expect_terminated=False,
+            reasoning_ended=True,
+        ),
+        id="active_later_token_invalid",
+    ),
+    pytest.param(
+        FlowCase(
+            prefix='{"a": "b"}',
+            raw_drafts=(EOS, " ", " "),
+            expected_validated=(EOS,),
+            expected_row_pattern="CUUU",
+            expected_reasoning=True,
+            expect_terminated=True,
+            reasoning_ended=True,
+            xfail_guidance=(
+                "guidance validate_tokens rejects EOS on an already-complete "
+                "object even though accept_tokens would accept it"
+            ),
+        ),
+        id="active_first_token_terminates",
+    ),
+    pytest.param(
+        FlowCase(
+            prefix='{"a"',
+            raw_drafts=(":", ' "', "b", '"}', EOS, " "),
+            expected_validated=(":", ' "', "b", '"}', EOS),
+            expected_row_pattern="CCCCCUU",
+            expected_reasoning=True,
+            expect_terminated=True,
+            reasoning_ended=True,
+        ),
+        id="active_later_token_terminates",
+    ),
+    pytest.param(
+        FlowCase(
+            raw_drafts=(THINK_END, "{", ' "', "b"),
+            expected_row_pattern="UCCCC",
+            expected_reasoning=True,
+            expect_terminated=False,
+        ),
+        id="becomes_active_on_first_token",
+    ),
+    pytest.param(
+        FlowCase(
+            raw_drafts=(" ", THINK_END, "{", ' "'),
+            expected_row_pattern="UUCCC",
+            expected_reasoning=True,
+            expect_terminated=False,
+        ),
+        id="becomes_active_on_middle_token",
+    ),
+    pytest.param(
+        FlowCase(
+            raw_drafts=(" ", " ", THINK_END),
+            expected_row_pattern="UUUC",
+            expected_reasoning=True,
+            expect_terminated=False,
+        ),
+        id="becomes_active_on_last_token",
+    ),
+    pytest.param(
+        FlowCase(
+            raw_drafts=(" ", THINK_END, "z", "{"),
+            expected_validated=(" ", THINK_END),
+            expected_row_pattern="UUCUU",
+            expected_reasoning=True,
+            expect_terminated=False,
+        ),
+        id="becomes_active_first_token_invalid",
+    ),
+    pytest.param(
+        FlowCase(
+            raw_drafts=(" ", THINK_END, "{", "z"),
+            expected_validated=(" ", THINK_END, "{"),
+            expected_row_pattern="UUCCU",
+            expected_reasoning=True,
+            expect_terminated=False,
+        ),
+        id="becomes_active_later_token_invalid",
+    ),
+    pytest.param(
+        FlowCase(
+            raw_drafts=(" ", THINK_END, "{", "}", EOS, " "),
+            expected_validated=(" ", THINK_END, "{", "}", EOS),
+            expected_row_pattern="UUCCCUU",
+            expected_reasoning=True,
+            expect_terminated=True,
+        ),
+        id="becomes_active_terminates",
+    ),
+]
+
+
+@pytest.mark.parametrize("backend", BACKENDS)
+@pytest.mark.parametrize("case", FLOW_CASES)
+def test_real_flow(
+    tokenizer,
+    backend: str,
+    case: FlowCase,
+):
+    if backend == "guidance" and case.xfail_guidance:
+        pytest.xfail(case.xfail_guidance)
+
+    reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)}
+    manager, request = _build_harness(
+        tokenizer,
+        backend,
+        prefix=case.prefix,
+        reasoning_ended=case.reasoning_ended,
+        reasoning_parser_kwargs=reasoner_kwargs,
+    )
+
+    raw_drafts = _to_token_ids(tokenizer, case.raw_drafts)
+    expected_texts = (
+        case.raw_drafts if case.expected_validated is None else case.expected_validated
+    )
+    expected_validated = _to_token_ids(tokenizer, expected_texts)
+
+    _run_real_flow(
+        manager,
+        request,
+        raw_drafts=raw_drafts,
+        expected_validated=expected_validated,
+        expected_row_pattern=case.expected_row_pattern,
+        expected_reasoning=case.expected_reasoning,
+        expect_terminated=case.expect_terminated,
+    )
+
+
+@pytest.mark.parametrize("backend", BACKENDS)
+@pytest.mark.parametrize(
+    ("use_reasoner", "reasoning_ended", "enable_in_reasoning"),
+    [
+        pytest.param(True, False, True, id="enable_in_reasoning"),
+        pytest.param(True, True, False, id="reasoning_ended"),
+        pytest.param(False, False, False, id="no_reasoner"),
+        pytest.param(True, None, False, id="inferred"),
+    ],
+)
+def test_initial_constraint_activation(
+    tokenizer,
+    backend: str,
+    use_reasoner: bool,
+    reasoning_ended: bool | None,
+    enable_in_reasoning: bool,
+):
+    manager, request = _build_harness(
+        tokenizer,
+        backend,
+        use_reasoner=use_reasoner,
+        reasoning_ended=reasoning_ended,
+        enable_in_reasoning=enable_in_reasoning,
+        reasoning_parser_kwargs={
+            "marker": _single_token(tokenizer, THINK_END)
+            if reasoning_ended is not None
+            else None
+        },
+    )
+
+    # "{" is valid JSON start; "z" is not. Truncation proves grammar is active.
+    open_brace = _single_token(tokenizer, "{")
+    z = _single_token(tokenizer, "z")
+    assert manager.validate_tokens(request, [open_brace, z]) == [open_brace]
+
+    bitmask = manager.grammar_bitmask(
+        requests={request.request_id: request},
+        structured_output_request_ids=[request.request_id],
+        scheduled_spec_decode_tokens={},
+    )
+    assert bitmask is not None
+    request.append_output_token_ids([open_brace])
+    assert manager.accept_tokens(request, [open_brace])
+
+    # Short-circuit paths must not construct a request-local reasoner (#52573).
+    if enable_in_reasoning or reasoning_ended is True:
+        structured_req = request.structured_output_request
+        assert structured_req is not None
+        assert structured_req.reasoner is None
diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py
index 11b3e5c01b38..d3187c428cc0 100644
--- a/tests/v1/structured_output/test_utils.py
+++ b/tests/v1/structured_output/test_utils.py
@@ -2,21 +2,32 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
+from xgrammar import Grammar
+from xgrammar.testing import _is_grammar_accept_string
 
+from vllm.sampling_params import SamplingParams, StructuredOutputsParams
 from vllm.v1.structured_output.backend_xgrammar import (
     has_xgrammar_unsupported_json_features,
+    validate_xgrammar_grammar,
 )
+from vllm.v1.structured_output.utils import choice_as_grammar
 
 pytestmark = pytest.mark.cpu_test
 
 
+def grammar_accepts(schema: dict, text: str) -> bool:
+    return _is_grammar_accept_string(Grammar.from_json_schema(schema), text)
+
+
 @pytest.fixture
 def unsupported_string_schemas():
     return [
         {"type": "string", "format": "non_existing_format"},
+        # TODO(arpera):
         # pattern/format is compiled but length bounds are silently dropped,
         # so the combination must be rejected instead of producing quietly
         # wrong output
+        # https://github.com/mlc-ai/xgrammar/issues/749
         {"type": "string", "pattern": "^a+$", "maxLength": 2},
         {"type": "string", "pattern": "^a+$", "minLength": 3},
         {"type": "string", "format": "email", "maxLength": 10},
@@ -48,15 +59,93 @@ def unsupported_array_schemas():
 
 
 @pytest.fixture
-def unsupported_object_schemas():
+def unsupported_property_names_combinations():
+    return [
+        # TODO(arpera):
+        # this case is not covered by xgrammar, so we should report this bug to xgrammar
+        # xgrammar drops propertyNames whenever patternProperties is present,
+        # so "grade_12" matches the pattern and escapes the name constraint
+        {
+            "type": "object",
+            "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}},
+            "propertyNames": {"pattern": "^grade_[0-9]$"},  # does NOT match "grade_12"
+        },
+        # propertyNames makes xgrammar discard the sibling additionalProperties
+        # value schema: https://github.com/mlc-ai/xgrammar/issues/826
+        {
+            "type": "object",
+            "propertyNames": {"pattern": "^[a-z]+$"},
+            "additionalProperties": {"type": "integer"},
+        },
+        # propertyNames is a string schema that conventionally omits "type", so
+        # it escapes the string check while xgrammar drops its length bound all
+        # the same: https://github.com/mlc-ai/xgrammar/issues/749
+        {
+            "type": "object",
+            "propertyNames": {"pattern": "^a+$", "maxLength": 2},
+        },
+        # xgrammar emits named `properties` separately from `propertyNames` and
+        # only applies `propertyNames` in its additional-properties branch, so a
+        # key declared in `properties` escapes the name constraint regardless of
+        # what `propertyNames` contains.
+        {
+            "type": "object",
+            "properties": {"Bad": {"type": "integer"}},
+            "required": ["Bad"],
+            "propertyNames": {"pattern": "^[a-z]+$"},
+        },
+        {
+            "type": "object",
+            "properties": {"Bad": {"type": "integer"}},
+            "propertyNames": {"enum": ["good"]},
+        },
+        # xgrammar's propertyNames-only branch falls back to an unconstrained
+        # value type, so it drops unevaluatedProperties whether it restricts
+        # values (a schema) or forbids them (false).
+        {
+            "type": "object",
+            "propertyNames": {"pattern": "^[a-z]+$"},
+            "unevaluatedProperties": {"type": "integer"},
+        },
+        {
+            "type": "object",
+            "propertyNames": {"pattern": "^[a-z]+$"},
+            "unevaluatedProperties": False,
+        },
+    ]
+
+
+@pytest.fixture
+def unsupported_pattern_properties_combinations():
     return [
-        {"type": "object", "propertyNames": {"pattern": "^[a-z]+$"}},
-        {"type": "object", "patternProperties": {"^S": {"type": "string"}}},
+        # JSON Schema requires a property matched by both `properties` and
+        # `patternProperties` to satisfy both (conjunction), but xgrammar
+        # compiles them as alternative branches, so satisfying either one is
+        # enough.
+        {
+            "type": "object",
+            "properties": {"x": {"type": "string"}},
+            "patternProperties": {"^x$": {"type": "integer"}},
+        },
+        # The same alternative-branches problem applies to overlapping
+        # patternProperties patterns: a key matching both patterns only has to
+        # satisfy one of them.
+        {
+            "type": "object",
+            "patternProperties": {
+                "^a[a-z]*$": {"type": "string"},
+                "^[a-z]*z$": {"type": "integer"},
+            },
+        },
     ]
 
 
 @pytest.fixture
-def supported_schema():
+def supported_frankenstein_schema():
+    # IMPORTANT(arpera):
+    # Do NOT add more keywords here! This schema is overcrowded enough that a new
+    # keyword can end up having no effect on the compiled grammar while the
+    # test still passes. Give a new keyword its own fixture and named test.
     return {
         "type": "object",
         "properties": {
@@ -88,25 +177,149 @@ def supported_schema():
     }
 
 
-@pytest.mark.parametrize(
-    "schema_type",
-    [
-        "unsupported_string_schemas",
-        "unsupported_integer_schemas",
-        "unsupported_number_schemas",
-        "unsupported_array_schemas",
-        "unsupported_object_schemas",
-    ],
-)
-def test_unsupported_json_features_by_type(schema_type, request):
-    schemas = request.getfixturevalue(schema_type)
-    for schema in schemas:
-        assert has_xgrammar_unsupported_json_features(schema), (
-            f"Schema should be unsupported: {schema}"
-        )
+@pytest.fixture
+def property_names_schema():
+    return {"type": "object", "propertyNames": {"pattern": "^[a-z_]+$"}}
+
+
+@pytest.fixture
+def pattern_properties_schema():
+    return {
+        "type": "object",
+        "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}},
+    }
+
 
+class TestHasXGrammarUnsupportedJsonFeatures:
+    @pytest.mark.parametrize(
+        "schema_type",
+        [
+            "unsupported_string_schemas",
+            "unsupported_integer_schemas",
+            "unsupported_number_schemas",
+            "unsupported_array_schemas",
+            "unsupported_property_names_combinations",
+            "unsupported_pattern_properties_combinations",
+        ],
+    )
+    def test_unsupported_json_features_by_type(self, schema_type, request):
+        schemas = request.getfixturevalue(schema_type)
+        for schema in schemas:
+            assert has_xgrammar_unsupported_json_features(schema), (
+                f"Schema should be unsupported: {schema}"
+            )
 
-def test_supported_json_features(supported_schema):
-    assert not has_xgrammar_unsupported_json_features(supported_schema), (
-        "Schema should be supported"
+    @pytest.mark.parametrize(
+        "schema_type",
+        [
+            "supported_frankenstein_schema",
+            "pattern_properties_schema",
+            "property_names_schema",
+        ],
     )
+    def test_supported_json_features(self, schema_type, request):
+        schema = request.getfixturevalue(schema_type)
+        assert not has_xgrammar_unsupported_json_features(schema), (
+            f"Schema should be supported: {schema}"
+        )
+
+    class TestPR48416Regressions:
+        @pytest.mark.parametrize(
+            "schema",
+            [
+                {"type": ["number"], "multipleOf": 3},
+                {
+                    "type": ["string", "null"],
+                    "pattern": "^a+$",
+                    "maxLength": 2,
+                },
+                {
+                    "type": ["array", "null"],
+                    "items": {"type": "integer"},
+                    "uniqueItems": True,
+                },
+                {
+                    "type": ["object", "null"],
+                    "properties": {"Bad": {"type": "integer"}},
+                    "propertyNames": {"pattern": "^[a-z]+$"},
+                },
+            ],
+        )
+        def test_list_type_does_not_bypass_unsupported_feature_check(self, schema):
+            assert has_xgrammar_unsupported_json_features(schema)
+
+        @pytest.mark.parametrize(
+            "schema",
+            [
+                {"type": ["string", "null"]},
+                {"type": ["integer", "string"]},
+                {"type": ["array", "null"], "items": {"type": "string"}},
+                {
+                    "type": ["object", "null"],
+                    "propertyNames": {"pattern": "^[a-z_]+$"},
+                },
+                {
+                    "type": ["object", "null"],
+                    "patternProperties": {"^S": {"type": "string"}},
+                },
+            ],
+        )
+        def test_supported_list_type_json_features(self, schema):
+            assert not has_xgrammar_unsupported_json_features(schema)
+
+
+class TestIsGrammarAcceptString:
+    class TestPR42904Support:
+        def test_property_names_constrains_keys(self, property_names_schema):
+            assert grammar_accepts(property_names_schema, '{"score": 5}')
+            assert grammar_accepts(property_names_schema, '{"score": "seven"}')
+            assert not grammar_accepts(property_names_schema, '{"Score": 5}')
+            assert not grammar_accepts(property_names_schema, '{"score_1": 5}')
+
+        def test_pattern_properties_constrains_keys_and_values(
+            self, pattern_properties_schema
+        ):
+            assert grammar_accepts(pattern_properties_schema, '{"grade_1": 5}')
+            assert not grammar_accepts(pattern_properties_schema, '{"other": 5}')
+            assert not grammar_accepts(pattern_properties_schema, '{"grade_1": "five"}')
+
+    class TestPR48115Regressions:
+        @pytest.mark.parametrize("codepoint", [*range(0x20), 0x7F])
+        def test_choice_as_grammar_preserves_control_characters(self, codepoint):
+            choice = f"a{chr(codepoint)}f"
+            grammar = Grammar.from_ebnf(choice_as_grammar([choice]))
+
+            assert _is_grammar_accept_string(grammar, choice)
+            assert not _is_grammar_accept_string(grammar, "a")
+            assert not _is_grammar_accept_string(grammar, "af")
+            assert not _is_grammar_accept_string(grammar, f"a\\u{codepoint:04x}f")
+
+        @pytest.mark.parametrize(
+            ("choice", "other"),
+            [
+                (r"a\nf", "a\nf"),
+                ('a quote " and a backslash \\', 'a quote " and a backslash '),
+                ("café", "cafe"),
+                ("日本語", "日本"),
+                ("😀", "😁"),
+            ],
+        )
+        def test_choice_as_grammar_preserves_literal_choices(self, choice, other):
+            grammar = Grammar.from_ebnf(choice_as_grammar(["yes", choice]))
+
+            assert _is_grammar_accept_string(grammar, "yes")
+            assert _is_grammar_accept_string(grammar, choice)
+            assert not _is_grammar_accept_string(grammar, other)
+            assert not _is_grammar_accept_string(grammar, choice + "extra")
+
+        def test_validate_xgrammar_preserves_multiline_choice(self):
+            structured_outputs = StructuredOutputsParams(choice=["yes", "no\nplease"])
+            validate_xgrammar_grammar(
+                SamplingParams(structured_outputs=structured_outputs)
+            )
+
+            assert structured_outputs.choice is None
+            grammar = Grammar.from_ebnf(structured_outputs.grammar)
+            assert _is_grammar_accept_string(grammar, "yes")
+            assert _is_grammar_accept_string(grammar, "no\nplease")
+            assert not _is_grammar_accept_string(grammar, r"no\nplease")
diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py
index af417bc03e6f..739279280caf 100644
--- a/tests/v1/structured_output/test_validation.py
+++ b/tests/v1/structured_output/test_validation.py
@@ -56,13 +56,18 @@ def test_plain_request_allowed_for_diffusion_models():
     [
         (StructuredOutputsParams(json_object=False), "json_object must be True"),
         (StructuredOutputsParams(json=""), "json cannot be an empty string"),
+        (
+            StructuredOutputsParams(structural_tag=""),
+            "structural_tag cannot be an empty string",
+        ),
     ],
 )
 def test_degenerate_structured_outputs_rejected(structured_outputs, match):
     """json_object=False and an empty json schema pass the `is not None`
     exclusivity check but resolve to no structured-output key, so they must be
-    rejected at request validation (-> 400) instead of reaching and crashing
-    the engine."""
+    rejected at request validation (-> 400). Empty `structural_tag` is rejected
+    for the same reason: `json.loads("")` in `compile_grammar` would otherwise
+    raise JSONDecodeError and surface as a per-request engine error."""
     params = SamplingParams(structured_outputs=structured_outputs)
     with pytest.raises(VLLMValidationError, match=match):
         params._validate_structured_outputs(
@@ -137,7 +142,7 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs):
 @pytest.mark.parametrize(
     "schema, expected_backend",
     [
-        # multipleOf is unsupported by xgrammar, patternProperties also by guidance.
+        # multipleOf is unsupported by xgrammar.
         (
             {
                 "type": "object",
@@ -146,7 +151,27 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs):
             "guidance",
         ),
         (
-            {"type": "object", "patternProperties": {"^a": {"type": "string"}}},
+            {
+                "type": "object",
+                "properties": {"n": {"type": ["number", "null"], "multipleOf": 3}},
+            },
+            "guidance",
+        ),
+        (
+            {
+                "type": ["string", "null"],
+                "pattern": "^a+$",
+                "maxLength": 2,
+            },
+            "guidance",
+        ),
+        # patternProperties + properties is also unsupported by guidance.
+        (
+            {
+                "type": "object",
+                "properties": {"a": {"type": "string"}},
+                "patternProperties": {"^a$": {"type": "string"}},
+            },
             "outlines",
         ),
     ],
diff --git a/tests/v1/test_kv_cache_spec_registry.py b/tests/v1/test_kv_cache_spec_registry.py
index 58c23ce285b1..716921cdf456 100644
--- a/tests/v1/test_kv_cache_spec_registry.py
+++ b/tests/v1/test_kv_cache_spec_registry.py
@@ -16,6 +16,7 @@
     ChunkedLocalAttentionManager,
     CrossAttentionManager,
     FullAttentionManager,
+    HiSparseSourceManager,
     MambaManager,
     SingleTypeKVCacheManager,
     SinkFullAttentionManager,
@@ -27,6 +28,7 @@
     CrossAttentionSpec,
     FullAttentionSpec,
     HiddenStateCacheSpec,
+    KVCacheGroupRole,
     KVCacheSpec,
     KVCacheSpecKind,
     MambaSpec,
@@ -68,8 +70,7 @@ def restore_kv_cache_spec_registry():
 
 @dataclass(frozen=True)
 class _TrulyUnregisteredSpec(KVCacheSpec):
-    """
-    A spec that inherits directly from KVCacheSpec with no registered
+    """A spec that inherits directly from KVCacheSpec with no registered
     ancestor in the MRO.  Used to test that the registry correctly raises
     when no entry can be found.
     """
@@ -176,6 +177,19 @@ def test_builtin_kvcache_specs_registered(self):
                 is spec_uniform_base_map[spec_cls]
             )
 
+    @pytest.mark.parametrize("role", list(KVCacheGroupRole))
+    def test_mla_manager_selection_by_role(self, role):
+        """Only the source role overrides ordinary MLA manager selection."""
+        expected = (
+            HiSparseSourceManager
+            if role == KVCacheGroupRole.HISPARSE_SOURCE
+            else FullAttentionManager
+        )
+        assert (
+            KVCacheSpecRegistry.get_manager_class(make_spec(MLAAttentionSpec), role)
+            is expected
+        )
+
     @pytest.mark.parametrize("spec_cls", list(spec_manager_map))
     def test_custom_spec_register(self, spec_cls):
         """A decorated custom spec resolves to the declared manager."""
@@ -197,7 +211,6 @@ class _CustomSpec(spec_cls):  # type: ignore[valid-type,misc]
 
     def test_custom_spec_register_requires_manager(self):
         """Invalid register decorator arguments fail early."""
-
         with pytest.raises(AssertionError, match="manager_class is required"):
 
             @register_kv_cache_spec(
@@ -208,8 +221,7 @@ class _CustomFullSpecWithoutManager(FullAttentionSpec):
                 custom_param: int = 16
 
     def test_unregistered_spec_no_registered_parent_raises(self):
-        """
-        A spec whose entire MRO contains no registered class resolves to None.
+        """A spec whose entire MRO contains no registered class resolves to None.
         Runtime callers should use check_kv_cache_spec_registry to fail early.
         Subclasses of registered specs intentionally do not fail — they inherit
         their parent's manager via MRO walking.
@@ -228,8 +240,7 @@ def test_unregistered_spec_no_registered_parent_raises(self):
             UniformTypeKVCacheSpecs.is_uniform_type({"layer_0": spec})
 
     def test_unregistered_subclass_inherits_parent_manager(self):
-        """
-        An unregistered subclass of a registered spec resolves via MRO
+        """An unregistered subclass of a registered spec resolves via MRO
         to its parent's manager — this is intentional registry behaviour.
         """
 
@@ -335,8 +346,7 @@ def test_homogeneous_group_keeps_member_kind(self, spec_cls, expected_kind):
 
     @pytest.mark.parametrize("other_cls", [MLAAttentionSpec, SinkFullAttentionSpec])
     def test_mixed_kinds_sharing_full_attention_base(self, other_cls):
-        """
-        Members with differing kinds but a single registered
+        """Members with differing kinds but a single registered
         uniform_type_base_spec of FullAttentionSpec report FULL_ATTENTION.
 
         This is the shape produced by a model that pairs a main K/V cache with a
@@ -385,8 +395,7 @@ def test_mixed_base_specs_report_unknown(self):
         assert get_kv_cache_spec_kind(group) == KVCacheSpecKind.UNKNOWN
 
     def test_uniform_non_full_attention_base_reports_unknown(self):
-        """
-        The FULL_ATTENTION result is gated on the shared base spec being
+        """The FULL_ATTENTION result is gated on the shared base spec being
         FullAttentionSpec; a uniform base of any other type stays UNKNOWN.
         """
 
diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py
index 50abb221ad82..5d52c86374b7 100644
--- a/tests/v1/test_outputs.py
+++ b/tests/v1/test_outputs.py
@@ -185,7 +185,7 @@ def setUp(self):
         )
 
     def test_slice_without_cu_num_generated_tokens(self):
-        """Test slicing without cu_num_generated_tokens"""
+        """Test slicing without cu_num_generated_tokens."""
         logprobsLists = LogprobsLists(
             logprob_token_ids=[[1], [2], [3]],
             logprobs=[[0.1], [0.2], [0.3]],
@@ -200,7 +200,7 @@ def test_slice_without_cu_num_generated_tokens(self):
         assert sliced.cu_num_generated_tokens is None
 
     def test_slice_from_start(self):
-        """Test slicing from the start position"""
+        """Test slicing from the start position."""
         sliced = self.logprobsLists.slice_request(0, num_positions=5)
         assert len(sliced.logprob_token_ids) == 5
         assert sliced.logprob_token_ids == [
@@ -213,7 +213,7 @@ def test_slice_from_start(self):
         assert sliced.cu_num_generated_tokens is None
 
     def test_slice_from_middle(self):
-        """Test slicing from the middle position"""
+        """Test slicing from the middle position."""
         sliced = self.logprobsLists.slice_request(1, num_positions=7)
         assert len(sliced.logprob_token_ids) == 7
         assert sliced.logprob_token_ids == [
@@ -228,21 +228,21 @@ def test_slice_from_middle(self):
         assert sliced.cu_num_generated_tokens is None
 
     def test_slice_single_request(self):
-        """Test slicing a single request"""
+        """Test slicing a single request."""
         sliced = self.logprobsLists.slice_request(1, num_positions=3)
         assert len(sliced.logprob_token_ids) == 3
         assert sliced.logprob_token_ids == [[5, 6], [7, 8], [9, 10]]
         assert sliced.cu_num_generated_tokens is None
 
     def test_slice_last_request(self):
-        """Test slicing the last request"""
+        """Test slicing the last request."""
         sliced = self.logprobsLists.slice_request(2, num_positions=4)
         assert len(sliced.logprob_token_ids) == 4
         assert sliced.logprob_token_ids == [[11, 12], [13, 14], [15, 16], [17, 18]]
         assert sliced.cu_num_generated_tokens is None
 
     def test_slice_all_requests(self):
-        """Test slicing all requests (full slice)"""
+        """Test slicing all requests (full slice)."""
         sliced = self.logprobsLists.slice_request(0, num_positions=9)
         assert len(sliced.logprob_token_ids) == 9  # All tokens
         assert sliced.logprob_token_ids == self.logprobsLists.logprob_token_ids
diff --git a/tests/v1/test_request.py b/tests/v1/test_request.py
index be417b9b2ff5..6c45e1a6d056 100644
--- a/tests/v1/test_request.py
+++ b/tests/v1/test_request.py
@@ -3,6 +3,7 @@
 
 from vllm import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.kv_hints import KvHintAction, KvHintsEnvelope
 from vllm.v1.request import Request, RequestStatus
 
 
@@ -24,6 +25,7 @@ def test_request_status_fmt_str():
 
 
 def test_request_copies_session_id_from_engine_core_request():
+    """Test that Request preserves session ID from EngineCoreRequest."""
     engine_request = EngineCoreRequest(
         request_id="request-1",
         prompt_token_ids=[1, 2, 3],
@@ -40,3 +42,35 @@ def test_request_copies_session_id_from_engine_core_request():
     request = Request.from_engine_core_request(engine_request, block_hasher=None)
 
     assert request.session_id == "session-1"
+
+
+def test_request_copies_kv_hints_from_engine_core_request():
+    """Test that Request preserves KV hints from EngineCoreRequest."""
+    kv_hints = KvHintsEnvelope(
+        protocol_version="0.1",
+        message_id="msg-1",
+        actions=[
+            KvHintAction(
+                action_id="action-1",
+                action_type="example.action",
+                action_version="1.0",
+                payload={},
+            )
+        ],
+    )
+    engine_request = EngineCoreRequest(
+        request_id="request-1",
+        prompt_token_ids=[1, 2, 3],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+        kv_hints=kv_hints,
+    )
+
+    request = Request.from_engine_core_request(engine_request, block_hasher=None)
+
+    assert request.kv_hints == kv_hints
diff --git a/tests/v1/test_serial_utils.py b/tests/v1/test_serial_utils.py
index 9f7761c22b5d..db0d7e68229b 100644
--- a/tests/v1/test_serial_utils.py
+++ b/tests/v1/test_serial_utils.py
@@ -44,7 +44,6 @@ class MyType:
 
 def test_encode_decode(monkeypatch: pytest.MonkeyPatch):
     """Test encode/decode loop with zero-copy tensors."""
-
     with monkeypatch.context() as m:
         m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
 
@@ -251,7 +250,6 @@ def test_custom_class_serialization_allowed_with_pickle(
     monkeypatch: pytest.MonkeyPatch,
 ):
     """Test that serializing a custom class succeeds when allow_pickle=True."""
-
     with monkeypatch.context() as m:
         m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
         encoder = MsgpackEncoder()
diff --git a/tests/v1/utils.py b/tests/v1/utils.py
index 993ad8a947d0..c55ace9d572c 100644
--- a/tests/v1/utils.py
+++ b/tests/v1/utils.py
@@ -17,6 +17,7 @@ def get_prometheus_metrics(server: RemoteOpenAIServer) -> dict[str, dict[str, fl
         For example: {"vllm:request_success": {
             "engine=0": 5.0, "engine=1": 3.0}
         }
+
     """
     try:
         response = requests.get(server.url_for("metrics"), timeout=10)
@@ -72,6 +73,7 @@ def get_engine_request_counts(metrics: dict[str, dict[str, float]]) -> dict[str,
     Returns:
         Dict mapping engine indices to request counts.
         For example: {"0": 15.0, "1": 12.0}
+
     """
     engine_counts = {}
 
@@ -97,6 +99,7 @@ def check_request_balancing(server: RemoteOpenAIServer, dp_size: int):
     Args:
         server: The RemoteOpenAIServer instance
         dp_size: Number of data parallel ranks
+
     """
     if dp_size <= 1:
         return
diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py
index 9751865794b7..99481a1293f9 100644
--- a/tests/v1/worker/test_attn_utils.py
+++ b/tests/v1/worker/test_attn_utils.py
@@ -7,25 +7,36 @@
 never addressed by the logical view.
 """
 
+from types import SimpleNamespace
+
+import numpy as np
 import pytest
 import torch
 
+import vllm.v1.hisparse.binding as attn_utils_module
 from tests.v1.attention.utils import dense_kv_cache_views
-from vllm.v1.attention.backend import AttentionCGSupport
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
+    HiSparseResidentSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheLayout,
     KVCacheTensor,
     MLAAttentionSpec,
+    SparseCacheRole,
     compute_layout_strides,
 )
+from vllm.v1.worker.gpu import attn_utils
 from vllm.v1.worker.gpu.attn_utils import (
+    FastPrefillHelper,
     get_attn_cg_support,
     get_query_lens_mismatch_unsupported_backend,
 )
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.utils import (
     AttentionGroup,
     allocate_kv_cache,
@@ -33,6 +44,60 @@
 )
 
 
+@pytest.mark.parametrize(
+    ("enabled", "block_size", "main_sizes", "indexer_sizes", "expected"),
+    [
+        (True, 256, [64], [64], 64),
+        (True, 64, [32, 64], [16, 32], 32),
+        (True, 64, [MultipleOf(16)], [32], 32),
+        (True, 64, [64], [32], None),
+        (False, 256, [64], [64], 256),
+    ],
+)
+def test_get_kv_cache_spec_resolves_hisparse_block_size(
+    monkeypatch, enabled, block_size, main_sizes, indexer_sizes, expected
+):
+    """Resolve shared MLA geometry before planning; leave other specs alone."""
+    specs = {
+        "main": MLAAttentionSpec(
+            block_size=block_size, num_kv_heads=1, head_size=576, dtype=torch.bfloat16
+        ),
+        "indexer": MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=128,
+            dtype=torch.bfloat16,
+            cache_role=SparseCacheRole.INDEXER,
+        ),
+        "dense": FullAttentionSpec(
+            block_size=block_size, num_kv_heads=1, head_size=128, dtype=torch.bfloat16
+        ),
+    }
+    layers = {}
+    for name, sizes in zip(specs, [main_sizes, indexer_sizes, [block_size]]):
+        backend = SimpleNamespace(
+            customize_spec=AttentionBackend.customize_spec,
+            get_supported_kernel_block_sizes=lambda sizes=sizes: sizes,
+        )
+        layers[name] = SimpleNamespace(
+            get_kv_cache_spec=lambda _, spec=specs[name]: spec,
+            get_attn_backend=lambda backend=backend: backend,
+        )
+    monkeypatch.setattr(attn_utils, "get_layers_from_vllm_config", lambda *_: layers)
+    config = SimpleNamespace(
+        attention_config=SimpleNamespace(hisparse_config=object() if enabled else None)
+    )
+    if expected is None:
+        with pytest.raises(ValueError, match="supported by every sparse"):
+            attn_utils.get_kv_cache_spec(config)
+        return
+
+    resolved = attn_utils.get_kv_cache_spec(config)
+    assert resolved["main"].block_size == resolved["indexer"].block_size == expected
+    assert resolved["dense"] is specs["dense"]
+    assert all(spec.block_size == block_size for spec in specs.values())
+
+
 class _FakeMetadataBuilder:
     def __init__(self, support: AttentionCGSupport):
         self.support = support
@@ -118,6 +183,205 @@ def test_attention_checks_preserve_global_and_target_scoped_support():
     )
 
 
+def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyPatch):
+    """Fast prefill applies to the contiguous suffix of KV-sharing layers.
+
+    Draft-model layers register after the target model's and may share KV, so
+    they must not extend (or break) the target's eligible suffix.
+    """
+
+    def check(
+        layer_names: list[str],
+        shared: dict[str, str],
+        draft_layer_names: set[str] | None = None,
+    ) -> set[str]:
+        monkeypatch.setattr(
+            attn_utils,
+            "get_layers_from_vllm_config",
+            lambda *a, **k: {name: None for name in layer_names},
+        )
+        monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda *a: shared)
+        vllm_config = SimpleNamespace(
+            cache_config=SimpleNamespace(kv_sharing_fast_prefill=True)
+        )
+        return attn_utils.get_kv_sharing_fast_prefill_eligible_layers(
+            vllm_config, draft_layer_names
+        )
+
+    # No KV sharing: nothing is eligible.
+    assert check(["t0", "t1"], {}) == set()
+
+    # Trailing run of sharing layers (YOCO-style second half).
+    assert check(["t0", "t1", "t2", "t3"], {"t2": "t1", "t3": "t1"}) == {"t2", "t3"}
+
+    # A non-sharing layer after a sharing one breaks the suffix.
+    assert check(["t0", "t1", "t2", "t3"], {"t1": "t0", "t3": "t0"}) == {"t3"}
+
+    # KV-sharing draft layers at the end are collected without an exclusion...
+    assert check(
+        ["t0", "t1", "t2", "t3", "d0", "d1"],
+        {"t2": "t1", "t3": "t1", "d0": "t1", "d1": "t1"},
+    ) == {"t2", "t3", "d0", "d1"}
+
+    # ...so the runner excludes them: skipped, not collected, and they do not
+    # break the target's trailing run.
+    assert check(
+        ["t0", "t1", "t2", "t3", "d0", "d1"],
+        {"t2": "t1", "t3": "t1", "d0": "t1", "d1": "t1"},
+        draft_layer_names={"d0", "d1"},
+    ) == {"t2", "t3"}
+
+    # Feature flag off: nothing is eligible even with sharing layers.
+    monkeypatch.setattr(
+        attn_utils, "get_layers_from_vllm_config", lambda *a, **k: {"t0": None}
+    )
+    monkeypatch.setattr(
+        attn_utils, "get_shared_kv_cache_layers", lambda *a: {"t0": "t0"}
+    )
+    vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(kv_sharing_fast_prefill=False)
+    )
+    assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set()
+
+
+@pytest.mark.parametrize(
+    "num_tokens", [1, 2, 3, 4, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128]
+)
+@pytest.mark.parametrize("num_active_loras", [1, 2, 4])
+def test_fast_prefill_dispatch_preserves_active_lora_count(
+    num_tokens: int, num_active_loras: int
+):
+    """Fast-prefill padding must match the main dispatch's LoRA variant."""
+
+    class FakeCudaGraphManager:
+        device = "cpu"
+
+        def __init__(self):
+            self.dispatch_calls = []
+
+        def dispatch(self, **kwargs):
+            self.dispatch_calls.append(kwargs)
+            # A captured no-LoRA graph is padded to 8 tokens, while the
+            # active-LoRA path stays eager at the unpadded token count.
+            if kwargs["num_active_loras"] == 0:
+                num_tokens = 8
+                mode = CUDAGraphMode.PIECEWISE
+            else:
+                num_tokens = kwargs["num_tokens"]
+                mode = CUDAGraphMode.NONE
+            return BatchExecutionDescriptor(
+                cg_mode=mode,
+                num_tokens=num_tokens,
+                num_reqs=kwargs["num_reqs"],
+                num_active_loras=kwargs["num_active_loras"],
+                num_ubatches=1,
+            )
+
+    manager = FakeCudaGraphManager()
+    helper = FastPrefillHelper(manager, max_num_tokens=max(32, num_tokens))
+    metadata = helper.prepare(
+        torch.arange(num_tokens, dtype=torch.int32),
+        num_reqs=1,
+        cu_num_logits_np=np.array([0, num_tokens], dtype=np.int32),
+        has_prefill=True,
+        batch_desc=BatchExecutionDescriptor(
+            cg_mode=CUDAGraphMode.NONE,
+            num_tokens=num_tokens,
+            num_reqs=1,
+            num_active_loras=num_active_loras,
+            num_ubatches=1,
+        ),
+        num_active_loras=num_active_loras,
+    )
+
+    assert metadata is not None
+    assert metadata.num_logits_indices == num_tokens
+    assert metadata.logits_indices_padded.shape[0] == num_tokens
+    assert metadata.max_logits_per_req == num_tokens
+    assert manager.dispatch_calls[-1]["num_active_loras"] == num_active_loras
+
+
+class _FakeSharedHostRegion:
+    def __init__(self) -> None:
+        self.cleanup_calls = 0
+        self.base_tensor = torch.empty(1, dtype=torch.int8)
+
+    def cleanup(self) -> None:
+        self.cleanup_calls += 1
+
+
+def test_profiling_cleanup_releases_tp_shared_region_once(monkeypatch):
+    """TP-shared profiling pools must use region-aware chunk cleanup."""
+    region = _FakeSharedHostRegion()
+    runtime = SimpleNamespace(
+        _host_cache=object(),
+        registered_host_pool=region.base_tensor,
+        hot_backing=object(),
+        shared_host_region=region,
+    )
+    forward_context = {
+        "layer": SimpleNamespace(
+            hisparse_cache=SimpleNamespace(runtime=runtime),
+        )
+    }
+    released = []
+
+    def release_pinned_state(runtimes, pinned_host_pools, shared_host_region):
+        released.append((runtimes, pinned_host_pools, shared_host_region))
+
+    monkeypatch.setattr(
+        attn_utils_module,
+        "release_pinned_state",
+        release_pinned_state,
+    )
+
+    attn_utils_module.release_hisparse_profiling_cache(forward_context)
+
+    assert released == [([runtime], [], region)]
+
+
+@pytest.mark.parametrize("failure_phase", ["allocation", "binding", "buffers"])
+def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase):
+    """A failure after mmap allocation must not leak the shared registration."""
+    region = _FakeSharedHostRegion()
+    vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1),
+    )
+
+    def allocate(*args):
+        args[-1].shared_region = region
+        if failure_phase == "allocation":
+            raise RuntimeError("initialization failed")
+        return {}
+
+    def bind(**kwargs):
+        if failure_phase == "binding":
+            raise RuntimeError("initialization failed")
+        return []
+
+    def buffers(*args, **kwargs):
+        raise RuntimeError("initialization failed")
+
+    monkeypatch.setattr(attn_utils_module, "allocate_hisparse_kv_caches", allocate)
+    monkeypatch.setattr(attn_utils_module, "bind_hisparse_kv_caches", bind)
+    monkeypatch.setattr(
+        attn_utils_module, "initialize_hisparse_runtime_buffers", buffers
+    )
+    with pytest.raises(RuntimeError, match="initialization failed"):
+        attn_utils_module.init_hisparse_kv_cache(
+            SimpleNamespace(),
+            torch.device("cpu"),
+            [],
+            vllm_config,
+            {},
+            SimpleNamespace(),
+        )
+    assert region.cleanup_calls == 1
+
+
 def test_reshape_padded_kv_cache_strides_by_padded_page():
     num_blocks = 3
     spec = FullAttentionSpec(
@@ -326,3 +590,71 @@ def test_copy_kv_cache_blocks_with_virtual_block_splitting(
             torch.testing.assert_close(
                 cache[dst_start + physical_idx], expected[layer_idx][physical_idx]
             )
+
+
+def test_allocate_hisparse_kv_caches_host_pool_and_view_less_specs():
+    """Host tensors get their own backing; view-less specs keep the raw one."""
+    spec = FullAttentionSpec(
+        block_size=2, num_kv_heads=1, head_size=4, dtype=torch.float32
+    )
+    page = spec.page_size_bytes
+    resident_spec = HiSparseResidentSpec(block_size=2, page_size=page)
+    device_size = 4 * page
+    config = KVCacheConfig(
+        num_blocks=4,
+        hisparse_host_num_blocks=3,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=3 * page,
+                layers=["source"],
+                layer_stride=3 * page,
+                block_stride=page,
+                host_resident=True,
+            ),
+            KVCacheTensor(
+                size=device_size,
+                layers=["indexer"],
+                layer_stride=device_size,
+                block_stride=page,
+            ),
+            KVCacheTensor(
+                size=device_size,
+                layers=["resident"],
+                layer_stride=device_size,
+                block_stride=page,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["source"], spec, host_resident=True),
+            KVCacheGroupSpec(["indexer"], spec),
+            KVCacheGroupSpec(["resident"], resident_spec),
+        ],
+    )
+    host_buffers: list[torch.Tensor] = []
+
+    def host_allocator(size: int) -> torch.Tensor:
+        host_buffers.append(torch.zeros(size, dtype=torch.int8))
+        return host_buffers[-1]
+
+    caches = allocate_hisparse_kv_caches(
+        config,
+        torch.device("cpu"),
+        KVCacheLayout.LBHNC,
+        [2, 2, 2],
+        SimpleNamespace(allocate=host_allocator),
+    )
+    assert len(config.kv_cache_tensors) == 3
+
+    assert [buf.numel() for buf in host_buffers] == [3 * page]
+    assert caches["source"].shape[0] == 3
+    assert (
+        caches["source"].untyped_storage().data_ptr()
+        == host_buffers[0].untyped_storage().data_ptr()
+    )
+    assert caches["indexer"].shape[0] == 4
+    backing = caches["resident"]
+    assert backing.dtype == torch.int8 and backing.numel() >= device_size
+    assert (
+        backing.untyped_storage().data_ptr()
+        == caches["indexer"].untyped_storage().data_ptr()
+    )
diff --git a/tests/v1/worker/test_dspark_utils.py b/tests/v1/worker/test_dspark_utils.py
new file mode 100644
index 000000000000..94655f192c89
--- /dev/null
+++ b/tests/v1/worker/test_dspark_utils.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass, field
+
+from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config
+
+
+@dataclass
+class _FakeEPLBConfig:
+    num_redundant_experts: int = 0
+
+
+@dataclass
+class _FakeParallelConfig:
+    pipeline_parallel_size: int = 2
+    tensor_parallel_size: int = 8
+    enable_eplb: bool = True
+    eplb_config: _FakeEPLBConfig = field(
+        default_factory=lambda: _FakeEPLBConfig(num_redundant_experts=32)
+    )
+    enable_elastic_ep: bool = True
+
+    def __post_init__(self) -> None:
+        if not self.enable_eplb and self.eplb_config.num_redundant_experts:
+            raise ValueError("redundant experts require EPLB")
+        if self.enable_elastic_ep and not self.enable_eplb:
+            raise ValueError("elastic EP requires EPLB")
+
+
+def test_dspark_parallel_config_disables_eplb_atomically():
+    target_config = _FakeParallelConfig()
+
+    draft_config = _get_dspark_parallel_config(
+        target_config,
+        tensor_parallel_size=4,
+    )
+
+    assert target_config.pipeline_parallel_size == 2
+    assert target_config.tensor_parallel_size == 8
+    assert target_config.enable_eplb
+    assert target_config.eplb_config.num_redundant_experts == 32
+    assert target_config.enable_elastic_ep
+
+    assert draft_config is not target_config
+    assert draft_config.pipeline_parallel_size == 1
+    assert draft_config.tensor_parallel_size == 4
+    assert not draft_config.enable_eplb
+    assert draft_config.eplb_config.num_redundant_experts == 0
+    assert not draft_config.enable_elastic_ep
+    assert draft_config.eplb_config is not target_config.eplb_config
diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py
index f84c91872faa..4cc8f9ffe2fd 100644
--- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py
+++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py
@@ -134,6 +134,7 @@ def make_spec(head_size):
     vllm_config.cache_config = CacheConfig()
     vllm_config.cache_config.num_gpu_blocks_override = None
     vllm_config.cache_config.kv_cache_layout = "BLHNC"
+    vllm_config.attention_config.hisparse_config = None
     config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024)
     views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None)
     buf_ptr = views["g1.big"].data_ptr()
diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py
index f8bd50c6d9c4..98aaa568f0b6 100644
--- a/tests/v1/worker/test_encoder_runner.py
+++ b/tests/v1/worker/test_encoder_runner.py
@@ -9,12 +9,15 @@
 the processed range still fails loudly.
 """
 
-from unittest.mock import MagicMock
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
 import torch
 
+from vllm.config.compilation import CUDAGraphMode
+from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalFieldElem,
@@ -31,6 +34,56 @@
 HIDDEN = 4
 
 
+@pytest.mark.parametrize(
+    "enabled,enforce_eager,supported,expected",
+    [
+        (True, False, True, True),
+        (False, False, True, False),
+        (True, True, True, False),
+        (True, False, False, False),
+    ],
+)
+def test_encoder_graph_configuration(enabled, enforce_eager, supported, expected):
+    """Encoder graph eligibility is independent of decoder graph mode."""
+    config = MagicMock()
+    config.model_config.enforce_eager = enforce_eager
+    config.model_config.dtype = torch.float32
+    config.model_config.get_inputs_embeds_size.return_value = HIDDEN
+    config.scheduler_config.max_num_batched_tokens = 8
+    config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
+    config.compilation_config.cudagraph_mm_encoder = enabled
+    model = MagicMock(spec=SupportsEncoderCudaGraph) if supported else torch.nn.Module()
+    state = _model_state(EncoderCache())
+    with patch("vllm.v1.worker.gpu.model_states.interface.EncoderCudaGraphManager"):
+        ModelState.__init__(state, config, model, state.encoder_cache, state.device)
+    assert state.encoder_runner.has_cudagraph() is expected
+
+
+def test_encoder_graph_unsupported_modality_uses_eager_output():
+    """An image graph must not intercept another modality's encoder output."""
+    from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
+
+    runner = _make_runner([], [])
+    output = torch.ones(2, HIDDEN)
+    runner.model = MagicMock()
+    runner.model.embed_multimodal.return_value = [output]
+    manager = MagicMock(spec=EncoderCudaGraphManager)
+    manager.config = SimpleNamespace(modalities=["image"])
+    manager.is_captured.return_value = True
+    manager.supports_modality.side_effect = lambda modality: (
+        EncoderCudaGraphManager.supports_modality(manager, modality)
+    )
+    runner.cudagraph_manager = manager
+    with patch(
+        "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs",
+        return_value=[("audio", 1, {"audio_values": torch.zeros(2)})],
+    ):
+        result = runner.execute_mm_encoder([])
+    assert len(result) == 1
+    assert result[0] is output
+    manager.execute.assert_not_called()
+
+
 def _model_state(cache: EncoderCache) -> MagicMock:
     """A mock ModelState backed by a real EncoderCache."""
     state = MagicMock()
diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py
index 6d8b8cbdbca5..4fc7d4d2e5d5 100644
--- a/tests/v1/worker/test_gpu_autoregressive_speculator.py
+++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py
@@ -19,6 +19,7 @@
 )
 from vllm.v1.attention.backends import flash_attn as flash_attn_module
 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module
 from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module
@@ -69,6 +70,49 @@ def embed_input_ids(
         raise AssertionError("embed_input_ids should not be called during loading")
 
 
+@pytest.mark.parametrize("cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL])
+def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode):
+    def build(common_prefix_len, common_attn_metadata):
+        return split_decodes_and_prefills(
+            common_attn_metadata,
+            decode_threshold=1,
+            require_uniform=True,
+            treat_short_extends_as_decodes=False,
+        )
+
+    speculator = object.__new__(_TestSpeculator)
+    speculator.arange_np = torch.arange(5, dtype=torch.int32).numpy()
+    speculator.max_model_len = speculator.draft_max_seq_len = 32
+    speculator.draft_is_prefilling = torch.zeros(4, dtype=torch.bool)
+    speculator.input_buffers = SimpleNamespace(
+        query_start_loc=torch.tensor([0, 1, 2, 2, 2], dtype=torch.int32),
+        seq_lens=torch.tensor([11, 21, 0, 0], dtype=torch.int32),
+    )
+    speculator.block_tables = SimpleNamespace(
+        cp_size=1,
+        input_block_tables=[torch.zeros(4, 1, dtype=torch.int32)],
+        slot_mappings=torch.tensor([[10, 20, -1, -1]]),
+    )
+    speculator.kv_cache_config = SimpleNamespace(kv_cache_groups=[object()])
+    speculator.attn_groups = [
+        [
+            SimpleNamespace(
+                get_metadata_builder=lambda _: SimpleNamespace(build=build),
+                layer_names=["draft"],
+            )
+        ]
+    ]
+    num_reqs_padded = 4 if cg_mode == CUDAGraphMode.FULL else 2
+    metadata = speculator._build_uniform_attn_metadata(
+        batch_desc=BatchExecutionDescriptor(cg_mode, 4, num_reqs_padded),
+        num_reqs=2,
+        num_query_per_req=1,
+        seq_lens_cpu_upper_bound=torch.tensor([10, 20], dtype=torch.int32),
+        step=1,
+    )
+    assert metadata["draft"] == (num_reqs_padded, 0, num_reqs_padded, 0)
+
+
 def _mock_base_model_load(monkeypatch):
     monkeypatch.setattr(
         base_spec_module,
@@ -94,6 +138,7 @@ def _make_speculator(
 
     speculator = object.__new__(_TestSpeculator)
     speculator.supports_mm_inputs = False
+    speculator.pcp_manager = None
     speculator.vllm_config = None
     speculator.input_buffers = SimpleNamespace(
         input_ids=torch.arange(4),
@@ -122,6 +167,7 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected)
         draft_model_config=draft_model_config,
         use_local_argmax_reduction=False,
         draft_sample_method="greedy",
+        enable_adaptive_verification=False,
     )
     vllm_config = SimpleNamespace(
         speculative_config=speculative_config,
@@ -146,8 +192,16 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected)
 
 
 def test_mm_support_configured_after_model_load(monkeypatch):
-    target_model_config = object()
-    draft_model_config = object()
+    checked_configs = []
+
+    class _ModelConfig:
+        @property
+        def supports_multimodal_inputs(self):
+            checked_configs.append(self)
+            return True
+
+    target_model_config = _ModelConfig()
+    draft_model_config = _ModelConfig()
     vllm_config = SimpleNamespace(model_config=target_model_config)
     draft_model = _MultimodalDraftModel()
 
@@ -160,20 +214,10 @@ def init_base(speculator, vllm_config, device):
         speculator.dtype = torch.float32
         speculator.draft_model_config = draft_model_config
         speculator.supports_mm_inputs = False
-
-    checked_configs = []
-
-    def supports_multimodal_inputs(model_config):
-        checked_configs.append(model_config)
-        return True
+        speculator.use_acceptance_estimator = False
 
     monkeypatch.setattr(DraftModelSpeculator, "__init__", init_base)
     _mock_base_model_load(monkeypatch)
-    monkeypatch.setattr(
-        base_spec_module.MULTIMODAL_REGISTRY,
-        "supports_multimodal_inputs",
-        supports_multimodal_inputs,
-    )
 
     speculator = _TestSpeculator(vllm_config, torch.device("cpu"))
 
@@ -194,7 +238,10 @@ def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch):
     speculator = object.__new__(_TestSpeculator)
     speculator.supports_mm_inputs = False
     speculator.inputs_embeds = None
-    speculator.vllm_config = SimpleNamespace(model_config=object())
+    speculator.use_acceptance_estimator = False
+    speculator.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(supports_multimodal_inputs=True)
+    )
     speculator.max_num_tokens = 4
     speculator.hidden_size = 3
     speculator.dtype = torch.float32
@@ -202,11 +249,6 @@ def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch):
     draft_model = _MultimodalDraftModel()
     speculator.test_draft_model = draft_model
     _mock_base_model_load(monkeypatch)
-    monkeypatch.setattr(
-        base_spec_module.MULTIMODAL_REGISTRY,
-        "supports_multimodal_inputs",
-        lambda model_config: True,
-    )
 
     speculator.load_model(torch.nn.Module())
 
@@ -218,16 +260,14 @@ def test_load_model_disables_mm_support_for_text_only_drafter(monkeypatch):
     speculator = object.__new__(_TestSpeculator)
     speculator.supports_mm_inputs = False
     speculator.inputs_embeds = None
-    speculator.vllm_config = SimpleNamespace(model_config=object())
+    speculator.use_acceptance_estimator = False
+    speculator.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(supports_multimodal_inputs=True)
+    )
     draft_model = _TextOnlyDraftModel()
     speculator.test_draft_model = draft_model
     warning_messages = []
     _mock_base_model_load(monkeypatch)
-    monkeypatch.setattr(
-        base_spec_module.MULTIMODAL_REGISTRY,
-        "supports_multimodal_inputs",
-        lambda model_config: True,
-    )
     monkeypatch.setattr(
         base_spec_module.logger,
         "warning_once",
@@ -250,8 +290,11 @@ def test_multi_module_mm_support_configured_after_model_load(monkeypatch):
     speculator = object.__new__(MultiModuleMTPSpeculator)
     speculator.supports_mm_inputs = False
     speculator.inputs_embeds = None
+    speculator.use_acceptance_estimator = False
     speculator.cached_draft_input_embeds = None
-    speculator.vllm_config = SimpleNamespace(model_config=object())
+    speculator.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(supports_multimodal_inputs=True)
+    )
     speculator.max_num_tokens = 4
     speculator.max_num_reqs = 2
     speculator.num_speculative_steps = 3
@@ -265,11 +308,6 @@ def test_multi_module_mm_support_configured_after_model_load(monkeypatch):
         "load_draft_model",
         lambda self, target_model, target_attn_layer_names: draft_model,
     )
-    monkeypatch.setattr(
-        base_spec_module.MULTIMODAL_REGISTRY,
-        "supports_multimodal_inputs",
-        lambda model_config: True,
-    )
 
     speculator.load_model(torch.nn.Module())
 
@@ -350,6 +388,8 @@ def test_multi_step_decode_replays_captured_graph_as_expected(
     speculator = object.__new__(_TestSpeculator)
     speculator.num_speculative_steps = 4
     speculator.current_draft_step = torch.tensor(0)
+    speculator.slot_mapping_observer = None
+    speculator.host_mirror_forward_observer = None
     speculator.input_buffers = SimpleNamespace(
         positions=torch.arange(2),
         query_start_loc=torch.arange(3),
diff --git a/tests/v1/worker/test_gpu_bad_words.py b/tests/v1/worker/test_gpu_bad_words.py
index 150c4ed00711..3cce549536a4 100644
--- a/tests/v1/worker/test_gpu_bad_words.py
+++ b/tests/v1/worker/test_gpu_bad_words.py
@@ -14,6 +14,10 @@
 
 from vllm.sampling_params import SamplingParams
 from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
+from vllm.v1.worker.gpu.sample.logits_processor import (
+    LogitsContext,
+    LogitsProcRequestState,
+)
 from vllm.v1.worker.gpu.states import RequestState
 
 DEVICE = torch.device("cuda")
@@ -47,7 +51,7 @@ def _make_state(bad_words_token_ids: list[list[int]]) -> tuple[BadWordsState, in
     req_states.apply_staged_writes()
 
     req_idx = req_states.req_id_to_index["req"]
-    state = BadWordsState(req_states)
+    state = BadWordsState(None, LogitsProcRequestState.from_request_state(req_states))
     state.add_request(req_idx, SamplingParams(_bad_words_token_ids=bad_words_token_ids))
     state.apply_staged_writes()
     return state, req_idx
@@ -61,16 +65,55 @@ def _apply(bad_words_token_ids: list[list[int]]) -> torch.Tensor:
     expanded_idx_mapping = torch.tensor(
         [req_idx] * num_logits, dtype=torch.int32, device=DEVICE
     )
-    state.apply_bad_words(
+    state.apply(
         logits,
-        expanded_idx_mapping,
-        idx_mapping_np,
-        torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE),
-        torch.tensor(LOCAL_POS, dtype=torch.int32, device=DEVICE),
+        LogitsContext(
+            expanded_idx_mapping=expanded_idx_mapping,
+            idx_mapping=torch.tensor([req_idx], dtype=torch.int32, device=DEVICE),
+            idx_mapping_np=idx_mapping_np,
+            expanded_local_pos=torch.tensor(
+                LOCAL_POS, dtype=torch.int32, device=DEVICE
+            ),
+            input_ids=torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE),
+            pos=torch.zeros(num_logits, dtype=torch.int32, device=DEVICE),
+            seq_lens_upper_bound_np=np.full(1, num_logits, dtype=np.int64),
+        ),
     )
     return logits.cpu()
 
 
+def test_v2_bad_words_slot_reuse_clears_state():
+    """A slot recycled from a bad-words request must not keep masking.
+
+    There is no removal hook by design; safety comes from add_request()
+    overwriting the slot's state unconditionally.
+    """
+    state, req_idx = _make_state([[10, 11, 60]])
+    state.add_request(req_idx, SamplingParams())
+    state.apply_staged_writes()
+    assert int(state.num_bad_words.np[req_idx]) == 0
+
+    num_logits = len(INPUT_IDS)
+    idx_mapping = torch.tensor([req_idx], dtype=torch.int32, device=DEVICE)
+    out = state.apply(
+        torch.zeros((num_logits, VOCAB_SIZE), device=DEVICE),
+        LogitsContext(
+            expanded_idx_mapping=torch.tensor(
+                [req_idx] * num_logits, dtype=torch.int32, device=DEVICE
+            ),
+            idx_mapping=idx_mapping,
+            idx_mapping_np=np.array([req_idx], dtype=np.intp),
+            expanded_local_pos=torch.tensor(
+                LOCAL_POS, dtype=torch.int32, device=DEVICE
+            ),
+            input_ids=torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE),
+            pos=torch.zeros(num_logits, dtype=torch.int32, device=DEVICE),
+            seq_lens_upper_bound_np=np.full(1, num_logits, dtype=np.int64),
+        ),
+    )
+    assert not torch.isinf(out).any()
+
+
 def test_v2_bad_words_prefix_inside_draft_tokens():
     """A prefix matching entirely within the draft tokens must mask the bad
     word's last token at the draft position that completes the prefix."""
diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py
index 1309686d53ce..ffb092209859 100644
--- a/tests/v1/worker/test_gpu_batch_shard.py
+++ b/tests/v1/worker/test_gpu_batch_shard.py
@@ -183,6 +183,12 @@ def test_local_batch_partition(tp_size: int, seed: int):
         ).all()
         assert (_np(local.logits_indices) == _np(batch.logits_indices)[rows]).all()
         assert (_np(local.seq_lens) == _np(batch.seq_lens)[owned]).all()
+        # Indexed by batch position, like seq_lens: a consumer holding local row
+        # indices reads another request's length if this is not localised too.
+        assert (
+            _np(local.seq_lens_cpu_upper_bound)
+            == _np(batch.seq_lens_cpu_upper_bound)[owned]
+        ).all()
         assert batch.num_draft_tokens_per_req is not None
         assert local.num_draft_tokens_per_req is not None
         assert (
diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py
index e77156eef901..5c245d1c4d9a 100644
--- a/tests/v1/worker/test_gpu_block_table.py
+++ b/tests/v1/worker/test_gpu_block_table.py
@@ -5,6 +5,7 @@
 import torch
 
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.worker.gpu.block_table import BlockTables
 
 pytestmark = pytest.mark.skipif(
@@ -273,3 +274,41 @@ def test_get_dummy_block_tables_returns_zeroed_rows():
     assert (dummy[0] == 0).all()
     # CUDA graph invariant: same persistent tensor, not a fresh allocation.
     assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr()
+
+
+def test_dummy_request_slot_mapping_is_pad():
+    """idx_mapping == -1 marks a dummy (or CUDA-graph padding) request.
+
+    A dummy draft decode step must not resolve slots through the persistent
+    block-table row of a real request slot, which may be stale and point at
+    blocks that are now in the prefix cache.
+    """
+    device = torch.device("cuda")
+    block_tables = BlockTables(
+        block_sizes=[4],
+        max_num_reqs=2,
+        max_num_batched_tokens=16,
+        max_num_blocks_per_group=[4],
+        device=device,
+        kernel_block_sizes=[4],
+    )
+    block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True)
+    block_tables.apply_staged_writes()
+    query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device)
+    positions = torch.tensor([1, 2, 3], dtype=torch.int64, device=device)
+
+    real = block_tables.compute_slot_mappings(
+        torch.tensor([0], dtype=torch.int32, device=device),
+        query_start_loc,
+        positions,
+        num_tokens_padded=3,
+    )
+    assert real[0].tolist() == [7 * 4 + 1, 7 * 4 + 2, 7 * 4 + 3]
+
+    dummy = block_tables.compute_slot_mappings(
+        torch.tensor([-1], dtype=torch.int32, device=device),
+        query_start_loc,
+        positions,
+        num_tokens_padded=3,
+    )
+    assert dummy[0].tolist() == [PAD_SLOT_ID] * 3
diff --git a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py
index d1aa0798feb2..74464f542ffc 100644
--- a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py
+++ b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py
@@ -60,6 +60,7 @@ def fake_set_forward_context(*args, **kwargs):
     layer_name = "cache_only_layers.2"
     speculator = object.__new__(ExtractHiddenStatesSpeculator)
     speculator.vllm_config = cast(Any, SimpleNamespace())
+    speculator.acceptance_estimator = None
     speculator.num_hidden_states = 2
     speculator.hidden_states = torch.zeros(4, 2, 3)
     speculator.draft_attn_layer_names = {layer_name}
@@ -111,6 +112,7 @@ def fake_set_forward_context(*args, **kwargs):
 
 def test_propose_requires_aux_hidden_states():
     speculator = object.__new__(ExtractHiddenStatesSpeculator)
+    speculator.acceptance_estimator = None
     speculator.num_hidden_states = 2
     input_batch = cast(
         Any, SimpleNamespace(idx_mapping=torch.tensor([0], dtype=torch.int32))
diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py
index 10a09b9de4a6..a4b3d71698a9 100644
--- a/tests/v1/worker/test_gpu_gumbel_sample.py
+++ b/tests/v1/worker/test_gpu_gumbel_sample.py
@@ -21,7 +21,12 @@
 if not torch.cuda.is_available():
     pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True)
 
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.sample.gumbel import (
+    _uniform64_from_random53,
+    gumbel_sample,
+    murmur3_hash32,
+)
 
 DEVICE = "cuda"
 VOCAB_SIZE = 200_000
@@ -33,6 +38,103 @@
 Z_TOLERANCE = 10.0
 
 
+@triton.jit
+def _murmur3_known_answer_kernel(
+    seeds_ptr,
+    positions_ptr,
+    offsets_ptr,
+    hashes_ptr,
+    domain_hashes_ptr,
+    numel: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    idx = tl.arange(0, BLOCK_SIZE)
+    mask = idx < numel
+    seed = tl.load(seeds_ptr + idx, mask=mask)
+    pos = tl.load(positions_ptr + idx, mask=mask)
+    offset = tl.load(offsets_ptr + idx, mask=mask)
+    hash32 = murmur3_hash32(seed, pos, offset).to(tl.int64)
+    domain_hash32 = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.int64)
+    tl.store(hashes_ptr + idx, hash32, mask=mask)
+    tl.store(domain_hashes_ptr + idx, domain_hash32, mask=mask)
+
+
+@triton.jit
+def _uniform64_endpoint_kernel(
+    random53_ptr,
+    uniform_ptr,
+    numel: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    idx = tl.arange(0, BLOCK_SIZE)
+    mask = idx < numel
+    random53 = tl.load(random53_ptr + idx, mask=mask).to(tl.uint64)
+    tl.store(uniform_ptr + idx, _uniform64_from_random53(random53), mask=mask)
+
+
+def test_murmur3_known_answers():
+    seeds = torch.tensor(
+        [0, 1, -1, 0x123456789ABCDEF, -(1 << 63), (1 << 63) - 1],
+        dtype=torch.int64,
+        device=DEVICE,
+    )
+    positions = torch.tensor(
+        [0, 2, 17, 123456789, 1 << 30, (1 << 30) + 7],
+        dtype=torch.int64,
+        device=DEVICE,
+    )
+    offsets = torch.tensor(
+        [0, 3, 1023, 199999, 4096, 0xFFFFFFFF],
+        dtype=torch.int64,
+        device=DEVICE,
+    )
+    hashes = torch.empty_like(seeds)
+    domain_hashes = torch.empty_like(seeds)
+
+    _murmur3_known_answer_kernel[(1,)](
+        seeds,
+        positions,
+        offsets,
+        hashes,
+        domain_hashes,
+        numel=seeds.numel(),
+        BLOCK_SIZE=8,
+    )
+
+    assert hashes.tolist() == [
+        0x8134CDF8,
+        0x23B53F91,
+        0xC8C13C14,
+        0x5DDE974F,
+        0xF1A9355B,
+        0xD952F9AA,
+    ]
+    assert domain_hashes.tolist() == [
+        0xD28AD962,
+        0xE3AAF818,
+        0xD223C51C,
+        0x0547D2FE,
+        0x3AED160F,
+        0x49E8EA7F,
+    ]
+
+
+def test_uniform64_excludes_endpoints():
+    random53 = torch.tensor([0, (1 << 53) - 1], dtype=torch.int64, device=DEVICE)
+    uniform = torch.empty(2, dtype=torch.float64, device=DEVICE)
+
+    _uniform64_endpoint_kernel[(1,)](
+        random53,
+        uniform,
+        numel=random53.numel(),
+        BLOCK_SIZE=2,
+    )
+
+    assert torch.isfinite(uniform).all()
+    assert ((uniform > 0.0) & (uniform < 1.0)).all()
+    assert uniform.tolist() == [2.0**-54, 1.0 - 2.0**-53]
+
+
 def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor:
     """Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N."""
     gen = torch.Generator(device=DEVICE).manual_seed(seed)
@@ -180,7 +282,7 @@ def test_full_vocab_distribution_fidelity():
 
 
 def test_drafting_uses_a_separate_noise_stream():
-    """is_drafting salts the Philox offset: same inputs, different draws.
+    """is_drafting salts the RNG position: same inputs, different draws.
 
     The draft proposal and the residual resample after a rejection must be
     independent. They key noise by the same (seed, pos), so only the salt keeps
@@ -189,7 +291,7 @@ def test_drafting_uses_a_separate_noise_stream():
     tests/v1/spec_decode/test_rejection_sampler_utils.py for the distributional
     consequence.
 
-    Relocating the offset must not distort the draw either, so both streams are
+    Salting the position must not distort the draw either, so both streams are
     checked against the target's far-tail mass, which sits HEAD_LOG_GAP logits
     below the head -- the regime where fp32 Gumbel precision matters.
     """
@@ -229,7 +331,7 @@ def test_drafting_uses_a_separate_noise_stream():
 
 
 def test_greedy_temperature_zero_returns_argmax():
-    """temperature == 0 skips Gumbel noise and returns the exact argmax."""
+    """Temperature == 0 skips Gumbel noise and returns the exact argmax."""
     torch.manual_seed(0)
     num_reqs = 128
     logits = torch.randn(num_reqs, VOCAB_SIZE, device=DEVICE, dtype=torch.float32)
diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py
index 3ef4e0dabe01..2fe3567739aa 100644
--- a/tests/v1/worker/test_gpu_input_batch.py
+++ b/tests/v1/worker/test_gpu_input_batch.py
@@ -66,11 +66,9 @@ def _compare_objs(obj1, obj2, skip: Sequence = ("logitsprocs", "batch_update_bui
 def _remove_requests(
     input_batch: InputBatch, batch_size: int, reqs: list[CachedRequestState]
 ) -> set[str]:
-    """
-    Remove some requests randomly from the batch and returns
+    """Remove some requests randomly from the batch and returns
     set of request removed
     """
-
     num_reqs_to_remove = np.random.randint(0, batch_size)
     req_indices_to_remove: set[int] = set()
     for _ in range(num_reqs_to_remove):
@@ -90,8 +88,7 @@ def _construct_expected_sampling_metadata(
     req_id_index_in_input_batch: dict[str, int],
     device: torch.device,
 ) -> SamplingMetadata:
-    """
-    Constructs and returns the expected SamplingMetadata for this
+    """Constructs and returns the expected SamplingMetadata for this
     batch.
     """
     num_reqs = len(req_ids_retained)
@@ -219,8 +216,7 @@ def _construct_cached_request_state(req_id_suffix: int):
 @pytest.mark.parametrize("device", DEVICES)
 @pytest.mark.parametrize("batch_size", [1, 2, 32, 64])
 def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
-    """
-    Tests the logic for managing sampling metadata in the InputBatch.
+    """Tests the logic for managing sampling metadata in the InputBatch.
 
     This test involves adding a set of requests to the InputBatch,
     followed by removing a subset of them. Afterward, the batch is compacted,
@@ -314,8 +310,7 @@ def same(t1: torch.Tensor | None, t2: torch.Tensor | None) -> bool:
 @pytest.mark.parametrize("batch_size", [32])
 @pytest.mark.parametrize("swap_list", [((0, 1),)])
 def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: list):
-    """
-    Tests the logic for managing sampling metadata in the InputBatch.
+    """Tests the logic for managing sampling metadata in the InputBatch.
 
     This test involves adding a set of requests to the InputBatch,
     followed by removing a subset of them. Afterward, the batch is compacted,
diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py
index 8937f2dc18ef..920277cb81ed 100644
--- a/tests/v1/worker/test_gpu_input_batch_v2.py
+++ b/tests/v1/worker/test_gpu_input_batch_v2.py
@@ -6,6 +6,9 @@
 import pytest
 import torch
 
+from vllm.config import VllmConfig
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.worker.gpu import cp_utils
@@ -56,7 +59,29 @@ def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int):
     )
 
 
-def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA top-k.")
+@pytest.mark.parametrize("is_padding", [True, False])
+def test_make_dummy_padding_controls_moe_routing(monkeypatch, is_padding: bool):
+    """Dummy tokens marked as padding are dropped by the MoE router (top-k id
+    -1), which is wanted for idle DP ranks. Profile runs must not mark them, or
+    no token reaches the experts and MoE memory is never profiled."""
+    monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1")
+    num_tokens = 16
+    buffers = InputBuffers(
+        max_num_reqs=4, max_num_tokens=num_tokens, device=torch.device(DEVICE)
+    )
+    batch = InputBatch.make_dummy(4, num_tokens, buffers, is_padding=is_padding)
+    hidden_states = torch.randn(num_tokens, 4, device=DEVICE)
+    router_logits = torch.randn(num_tokens, 8, device=DEVICE)
+
+    with set_forward_context(None, VllmConfig(), is_padding=batch.is_padding):
+        _, topk_ids, _ = fused_topk(hidden_states, router_logits, 2, False)
+
+    assert bool((topk_ids == -1).all()) is is_padding
+    assert bool((topk_ids >= 0).all()) is not is_padding
+
+
+def test_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
     """The batch must view the caller-owned buffer, sliced to padded length.
 
     Runtime (and capture) paths all funnel through this helper; attention
@@ -90,7 +115,7 @@ def __getitem__(self, grid):
             return fake_kernel
 
     monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", FakeKernel())
-    batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
+    batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens(
         buffers.dcp_local_seq_lens,
         batch.seq_lens,
         batch.num_reqs,
@@ -105,37 +130,10 @@ def __getitem__(self, grid):
     assert batch.dcp_local_seq_lens.tolist() == [1, 2, 0, 0]
 
 
-def test_maybe_prepare_dcp_local_seq_lens_clears_stale_metadata(monkeypatch):
-    """dcp_size == 1 resets the field to None instead of leaving a leftover.
-
-    DCP is toggled per deployment, and batches are recycled; a value from an
-    earlier DCP batch would otherwise be consumed as if it were current.
-    """
-    buffers = InputBuffers(max_num_reqs=2, max_num_tokens=2, device=torch.device("cpu"))
-    batch = InputBatch.make_dummy(1, 1, buffers)
-    batch.dcp_local_seq_lens = buffers.dcp_local_seq_lens[:1]
-
-    def fail_if_called(*args, **kwargs):
-        raise AssertionError("kernel must not run with dcp_size == 1")
-
-    monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", fail_if_called)
-    batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
-        buffers.dcp_local_seq_lens,
-        batch.seq_lens,
-        batch.num_reqs,
-        1,
-        0,
-        1,
-        num_reqs_padded=batch.num_reqs_after_padding,
-    )
-
-    assert batch.dcp_local_seq_lens is None
-
-
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
 @pytest.mark.parametrize("dcp_size", [2, 4])
 @pytest.mark.parametrize("cp_interleave", [1, 16])
-def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
+def test_prepare_dcp_local_seq_lens_matches_reference(
     dcp_size: int, cp_interleave: int
 ):
     """Every rank's local lengths must equal the reference torch formula."""
@@ -147,7 +145,7 @@ def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
 
     for dcp_rank in range(dcp_size):
         buffers.dcp_local_seq_lens.fill_(-1)
-        batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
+        batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens(
             buffers.dcp_local_seq_lens,
             batch.seq_lens,
             batch.num_reqs,
diff --git a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py
new file mode 100644
index 000000000000..8a7445f7d49a
--- /dev/null
+++ b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py
@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GPU regression test for the kpool tail group's slot mapping.
+
+The kpool tail cache holds one circular block per request, so its block table
+row is orders of magnitude shorter than the sequence it serves.
+`_compute_slot_mappings_kernel` indexes block tables by
+`position // kernel_block_size` with an effectively unmasked load, so
+this leads to a read of neighboring requests' rows (and for the later requests, outside
+the buffer entirely).
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.v1.kv_cache_interface import KpoolTailSpec
+from vllm.v1.worker.block_table import get_block_table_width
+from vllm.v1.worker.gpu.block_table import BlockTables
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="requires CUDA or ROCm",
+)
+
+KPOOL = 4
+MAX_MODEL_LEN = 1 << 20
+MLA_BLOCK_SIZE = 640
+MAX_NUM_REQS = 4
+POISON_BLOCK_ID = 1 << 20
+
+
+def make_tail_spec():
+    return KpoolTailSpec(
+        block_size=KPOOL,
+        num_kv_heads=2,
+        head_size=128,
+        head_size_v=0,
+        dtype=torch.bfloat16,
+        sliding_window=KPOOL,
+    )
+
+
+def test_kpool_tail_group_never_position_indexes_its_block_table():
+    device = torch.device("cuda")
+    spec = make_tail_spec()
+    tail_width = get_block_table_width(
+        spec.max_num_blocks_per_req(None, MAX_MODEL_LEN),
+        spec.block_size,
+        token_alignment=spec.block_table_token_alignment,
+    )
+    mla_width = get_block_table_width(MAX_MODEL_LEN // MLA_BLOCK_SIZE, MLA_BLOCK_SIZE)
+
+    # Positions chosen against the tail row: `in_row` still lands inside
+    # request 0's own row, `next_row` lands on request 1's poisoned row, and
+    # `past_table` is past the whole [MAX_NUM_REQS, tail_width] allocation.
+    in_row = KPOOL - 1
+    next_row = tail_width * KPOOL
+    past_table = MAX_NUM_REQS * tail_width * KPOOL + KPOOL
+    assert past_table < MAX_MODEL_LEN
+    positions = [0, in_row, next_row, past_table]
+
+    block_tables = BlockTables(
+        block_sizes=[MLA_BLOCK_SIZE, spec.block_size],
+        max_num_reqs=MAX_NUM_REQS,
+        max_num_batched_tokens=len(positions),
+        max_num_blocks_per_group=[mla_width, tail_width],
+        device=device,
+        kernel_block_sizes=[MLA_BLOCK_SIZE, spec.block_size],
+        slot_mapping_enabled=[True, spec.uses_slot_mapping],
+    )
+    mla_blocks = list(range(1, mla_width + 1))
+    block_tables.append_block_ids(0, (mla_blocks, [7]), overwrite=True)
+    block_tables.append_block_ids(1, ([2], [POISON_BLOCK_ID]), overwrite=True)
+    block_tables.apply_staged_writes()
+
+    idx_mapping = torch.zeros(1, dtype=torch.int32, device=device)
+    query_start_loc = torch.tensor(
+        [0, len(positions)], dtype=torch.int32, device=device
+    )
+    slot_mappings = block_tables.compute_slot_mappings(
+        idx_mapping,
+        query_start_loc,
+        torch.tensor(positions, dtype=torch.int64, device=device),
+        num_tokens_padded=len(positions),
+    )
+    torch.accelerator.synchronize()
+
+    tail_slots = slot_mappings[1].tolist()
+    assert tail_slots == [PAD_SLOT_ID] * len(positions)
+    assert POISON_BLOCK_ID * KPOOL not in tail_slots
+
+    mla_slots = slot_mappings[0].tolist()
+    assert mla_slots == [
+        mla_blocks[p // MLA_BLOCK_SIZE] * MLA_BLOCK_SIZE + p % MLA_BLOCK_SIZE
+        for p in positions
+    ]
diff --git a/tests/v1/worker/test_gpu_kv_connector.py b/tests/v1/worker/test_gpu_kv_connector.py
index 7fb40fff3711..03c483dce165 100644
--- a/tests/v1/worker/test_gpu_kv_connector.py
+++ b/tests/v1/worker/test_gpu_kv_connector.py
@@ -5,9 +5,11 @@
 from unittest.mock import Mock
 
 import pytest
+import torch
 
 import vllm.v1.worker.gpu.kv_connector as kv_connector_module
 from vllm.config import KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorTransferResults
 from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector
 
 
@@ -18,9 +20,9 @@ def _make_connector(
     backend = Mock()
     backend.handle_preemptions.side_effect = lambda _: events.append("handle")
     backend.bind_connector_metadata.side_effect = lambda _: events.append("bind")
-    backend.start_load_kv.side_effect = lambda _: events.append("start")
+    backend.start_load_kv.side_effect = lambda *_args, **_kwargs: events.append("start")
     backend.wait_for_save.side_effect = lambda: events.append("wait")
-    backend.get_finished.side_effect = lambda _: (set(), set())
+    backend.get_transfer_results.return_value = KVConnectorTransferResults()
     backend.get_block_ids_with_load_errors.return_value = set()
     backend.get_kv_connector_stats.return_value = None
     backend.get_kv_connector_kv_cache_events.return_value = None
@@ -61,7 +63,15 @@ def test_load_start_phase(
     connector = _make_connector(monkeypatch, events)
     output = _scheduler_output(has_sync_kv_loads)
 
-    connector.pre_forward(output)  # type: ignore[arg-type]
+    request_indices = torch.tensor([3, 1])
+    request_ids = ["first", "second"]
+    attn_metadata = {"layer": object()}
+    connector.pre_forward(  # type: ignore[arg-type]
+        output,
+        request_state_indices=request_indices,
+        request_ids=request_ids,
+        attn_metadata=attn_metadata,
+    )
     assert events == (
         ["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"]
     )
@@ -69,6 +79,16 @@ def test_load_start_phase(
     connector.post_forward(set())
     assert events == ["handle", "bind", "start", "wait", "clear"]
 
+    kwargs = connector.kv_connector.start_load_kv.call_args.kwargs
+    assert kwargs["request_state_indices"] is request_indices
+    assert kwargs["request_ids"] is request_ids
+    assert kwargs["attn_metadata"] is attn_metadata
+
+    # A subsequent step without a forward must not reuse the prior batch.
+    connector.no_forward(_scheduler_output(False))  # type: ignore[arg-type]
+    assert connector.kv_connector.start_load_kv.call_count == 2
+    assert connector.kv_connector.start_load_kv.call_args.kwargs == {}
+
 
 def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch):
     events: list[str] = []
diff --git a/tests/v1/worker/test_gpu_logit_bias.py b/tests/v1/worker/test_gpu_logit_bias.py
index 3bf62380bbc6..617a36120476 100644
--- a/tests/v1/worker/test_gpu_logit_bias.py
+++ b/tests/v1/worker/test_gpu_logit_bias.py
@@ -14,6 +14,11 @@
 
 from vllm.sampling_params import SamplingParams, StructuredOutputsParams
 from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
+from vllm.v1.worker.gpu.sample.logits_processor import (
+    LogitsContext,
+    LogitsProcRequestState,
+)
+from vllm.v1.worker.gpu.states import RequestState
 
 DEVICE = torch.device("cuda")
 VOCAB_SIZE = 128
@@ -45,17 +50,35 @@ def _only_stop_token_left(num_rows: int) -> torch.Tensor:
 
 def _apply(logits: torch.Tensor, structured: list[bool]) -> torch.Tensor:
     """Run the fused bias kernel with one logits row per request."""
-    state = LogitBiasState(max_num_reqs=4, device=DEVICE)
+    req_states = RequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=16,
+        num_speculative_steps=1,
+        vocab_size=VOCAB_SIZE,
+        device=DEVICE,
+    )
+    # add_request() reads prompt_len per slot; every slot here uses the same one.
+    req_states.prompt_len.np[:] = PROMPT_LEN
+
+    state = LogitBiasState(None, LogitsProcRequestState.from_request_state(req_states))
     for req_idx, is_structured in enumerate(structured):
-        state.add_request(req_idx, PROMPT_LEN, _params(is_structured))
+        state.add_request(req_idx, _params(is_structured))
     state.apply_staged_writes()
 
     num_reqs = len(structured)
-    state.apply_logit_bias(
+    idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
+    state.apply(
         logits,
-        torch.arange(num_reqs, dtype=torch.int32, device=DEVICE),
-        np.arange(num_reqs, dtype=np.intp),
-        torch.full((num_reqs,), POS, dtype=torch.int32, device=DEVICE),
+        LogitsContext(
+            expanded_idx_mapping=idx_mapping,
+            idx_mapping=idx_mapping,
+            idx_mapping_np=np.arange(num_reqs, dtype=np.intp),
+            expanded_local_pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
+            input_ids=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
+            pos=torch.full((num_reqs,), POS, dtype=torch.int32, device=DEVICE),
+            seq_lens_upper_bound_np=np.full(num_reqs, POS + 1, dtype=np.int64),
+        ),
     )
     return logits.cpu()
 
diff --git a/tests/v1/worker/test_gpu_logits_processors.py b/tests/v1/worker/test_gpu_logits_processors.py
new file mode 100644
index 000000000000..f0e5bd43a54b
--- /dev/null
+++ b/tests/v1/worker/test_gpu_logits_processors.py
@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+
+pytest.importorskip("triton")
+if not torch.cuda.is_available():
+    pytest.skip(
+        "CUDA required for Model Runner V2 logits processor tests",
+        allow_module_level=True,
+    )
+
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.logits_processor import (
+    LogitsContext,
+    LogitsProcessor,
+)
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.states import RequestState
+
+DEVICE = torch.device("cuda")
+VOCAB_SIZE = 128
+TARGET_TOKEN = 7
+BIAS = 5.0
+BOOST = 100.0
+TEMPERATURE = 2.0
+
+
+class MockReasoningConfig:
+    reasoning_start_token_ids = [90]
+    reasoning_end_token_ids = [91]
+    natural_reasoning_end_token_ids = [91]
+
+
+class RecordingProcessor(LogitsProcessor):
+    """Adds BOOST to TARGET_TOKEN and records every lifecycle callback."""
+
+    def __init__(self, vllm_config: Any, req_state: Any):
+        self.events: list[str] = []
+        self.seen_logits: torch.Tensor | None = None
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        self.events.append(f"add_request:{req_idx}")
+        return True
+
+    def apply_staged_writes(self) -> None:
+        self.events.append("apply_staged_writes")
+
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        self.events.append("apply")
+        self.seen_logits = logits.clone()
+        logits[:, TARGET_TOKEN] += BOOST
+        return logits
+
+
+def _make_sampler(processor: RecordingProcessor) -> Sampler:
+    req_states = RequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=16,
+        num_speculative_steps=1,
+        vocab_size=VOCAB_SIZE,
+        device=DEVICE,
+    )
+    return Sampler(
+        vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()),
+        max_num_reqs=4,
+        vocab_size=VOCAB_SIZE,
+        device=DEVICE,
+        req_states=req_states,
+        custom_logits_processors=[processor],
+    )
+
+
+def test_custom_processor_lifecycle_and_pipeline_position():
+    """A custom processor's hooks fire in order, it runs after the built-in
+    bias stage, and before temperature scaling."""
+    proc = RecordingProcessor(None, None)
+    sampler = _make_sampler(proc)
+    sampler.add_request(
+        0,
+        SamplingParams(logit_bias={TARGET_TOKEN: BIAS}, temperature=TEMPERATURE),
+    )
+    sampler.apply_staged_writes()
+
+    num_reqs = 1
+    idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
+    raw = torch.zeros(num_reqs, VOCAB_SIZE, device=DEVICE)
+    out = sampler.apply_sampling_params(
+        raw,
+        expanded_idx_mapping=idx_mapping,
+        idx_mapping=idx_mapping,
+        idx_mapping_np=np.arange(num_reqs, dtype=np.intp),
+        pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
+        input_ids=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
+        expanded_local_pos=torch.zeros(num_reqs, dtype=torch.int32, device=DEVICE),
+        seq_lens_upper_bound_np=np.ones(num_reqs, dtype=np.int64),
+    )
+
+    assert proc.events == [
+        "add_request:0",
+        "apply_staged_writes",
+        "apply",
+    ]
+    # The processor sees the built-in bias applied but no temperature scaling.
+    assert proc.seen_logits is not None
+    assert proc.seen_logits[0, TARGET_TOKEN].item() == BIAS
+    # Its own modification is then temperature-scaled by the sampler.
+    assert out[0, TARGET_TOKEN].item() == (BIAS + BOOST) / TEMPERATURE
diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py
index 99331e7b354e..e77d27811676 100644
--- a/tests/v1/worker/test_gpu_model_runner.py
+++ b/tests/v1/worker/test_gpu_model_runner.py
@@ -31,6 +31,7 @@
 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
 from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
 from vllm.platforms import current_platform
+from vllm.platforms.interface import Platform
 from vllm.sampling_params import SamplingParams
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.utils.system_utils import update_environment_variables
@@ -80,9 +81,7 @@ def _restore_default_dtype():
 
 
 def initialize_kv_cache(runner: GPUModelRunner):
-    """
-    Only perform necessary steps in GPUModelRunner.initialize_kv_cache()
-    """
+    """Only perform necessary steps in GPUModelRunner.initialize_kv_cache()."""
     attn_spec = FullAttentionSpec(
         block_size=BLOCK_SIZE,
         num_kv_heads=runner.model_config.get_num_kv_heads(runner.parallel_config),
@@ -176,6 +175,20 @@ def test_freeze_gc_disables_and_restores_automatic_gc(
             gc.disable()
 
 
+def test_prepare_padding_mask_marks_sequence_parallel_padding():
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.is_padding = torch.empty(8, dtype=torch.bool)
+
+    mask = runner._prepare_padding_mask(1, 8)
+
+    assert mask.tolist() == [False, True, True, True, True, True, True, True]
+    assert mask.data_ptr() == runner.is_padding.data_ptr()
+
+    mask = runner._prepare_padding_mask(0, 8)
+
+    assert mask.all()
+
+
 @pytest.fixture
 def model_runner():
     vllm_config = get_vllm_config()
@@ -312,6 +325,10 @@ def test_select_common_block_size_uses_largest_shared_int():
     assert selected_size == 64
 
 
+def test_select_common_block_size_without_active_backends_uses_manager_size():
+    assert select_common_block_size(256, []) == 256
+
+
 def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch):
     monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
 
@@ -409,6 +426,62 @@ def test_select_common_block_size_no_valid_option():
         select_common_block_size(48, [backend_a, backend_b])
 
 
+def _mock_backend(supported: list, *, exact: bool = False):
+    """Backend accepting multiples of ``supported``, or only those exact sizes
+    when ``exact`` (as CPU_MLA does)."""
+    from vllm.v1.attention.backend import AttentionBackend
+
+    class _MockBackendCls(AttentionBackend):
+        @staticmethod
+        def get_name() -> str:
+            return "MOCK_EXACT" if exact else "MOCK"
+
+        @staticmethod
+        def get_supported_kernel_block_sizes():
+            return list(supported)
+
+        if exact:
+
+            @classmethod
+            def supports_block_size(cls, block_size: int | None) -> bool:
+                return block_size is None or block_size in supported
+
+    return _MockBackendCls
+
+
+@pytest.mark.parametrize(
+    "backends,expected",
+    [
+        # A lone backend keeps the default; MultipleOf(8) makes a
+        # minimum-based regression return 8 instead of 16.
+        ([[MultipleOf(8)]], 16),
+        # Sparse-MLA main backend (32 or 64) beside an indexer supporting only
+        # 64: the main backend alone picks 32, which the indexer rejects at
+        # select_common_block_size time.
+        ([[32, 64], [64]], 64),
+        ([[32, 64]], 32),
+        # Neither divides the other, so only the LCM satisfies both.
+        ([[32], [48]], 96),
+    ],
+)
+def test_preferred_block_size_satisfies_every_backend(backends, expected):
+    classes = [_mock_backend(s) for s in backends]
+    assert Platform._preferred_block_size_for_backends(classes, 16, None) == expected
+
+
+def test_preferred_block_size_searches_past_an_exact_size_backend():
+    # Extending greedily picks lcm(16, 32) = 32, which the exact backend
+    # rejects; 96 is accepted by both.
+    classes = [_mock_backend([16, 96], exact=True), _mock_backend([MultipleOf(32)])]
+    assert Platform._preferred_block_size_for_backends(classes, 16, None) == 96
+
+
+def test_preferred_block_size_rejects_backends_with_no_common_size():
+    classes = [_mock_backend([16], exact=True), _mock_backend([MultipleOf(64)])]
+    with pytest.raises(ValueError, match="share no supported KV cache block size"):
+        Platform._preferred_block_size_for_backends(classes, 16, None)
+
+
 def test_set_active_mm_loras_builds_tower_and_connector_mappings():
     model = Mock()
     model.get_mm_lora_token_counts.side_effect = (
@@ -1189,14 +1262,12 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):
     reason="Attention backend FLASHINFER is only supported on CUDA.",
 )
 def test_hybrid_attention_mamba_tensor_shapes():
-    """
-    The GPU model runner creates different views into the
+    """The GPU model runner creates different views into the
     KVCacheTensors for the attention and mamba layers
     (via _allocate_kv_caches). This test verifies
     that the views are compatible: writing a mamba block
     will not corrupt an attention block and vice versa
     """
-
     set_random_seed(42)
 
     update_environment_variables(
diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py
index eb05b9858af4..83cd1e9d07b6 100644
--- a/tests/v1/worker/test_gpu_model_runner_v2.py
+++ b/tests/v1/worker/test_gpu_model_runner_v2.py
@@ -8,6 +8,7 @@
 import torch
 
 import vllm.v1.worker.gpu.model_runner as model_runner_module
+from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
     FullAttentionSpec,
@@ -20,6 +21,20 @@
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 
 
+def test_prepare_padding_mask_marks_sequence_parallel_padding():
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.input_buffers = SimpleNamespace(is_padding=torch.empty(8, dtype=torch.bool))
+
+    mask = runner._prepare_padding_mask(1, 8)
+
+    assert mask.tolist() == [False, True, True, True, True, True, True, True]
+    assert mask.data_ptr() == runner.input_buffers.is_padding.data_ptr()
+
+    mask = runner._prepare_padding_mask(0, 8)
+
+    assert mask.all()
+
+
 def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
     runner = GPUModelRunner.__new__(GPUModelRunner)
     runner.max_model_len = 262144
@@ -37,6 +52,7 @@ def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
         parallel_config=parallel_config,
         cache_config=SimpleNamespace(mamba_cache_mode="none"),
     )
+    runner.jit_warmup_registry = JitWarmupRegistry(runner.vllm_config)
     runner.model_state = SimpleNamespace(
         get_additional_cg_support=lambda: (),
         num_new_sampled_tokens_per_step=1,
@@ -84,7 +100,7 @@ def narrow(self, *args):
     monkeypatch.setattr(
         model_runner_module,
         "init_attn_backend",
-        lambda *args: ([], attn_cg_support, [8, 262144]),
+        lambda *args, **kwargs: ([], attn_cg_support, [8, 262144]),
     )
     monkeypatch.setattr(
         model_runner_module,
@@ -124,7 +140,6 @@ def test_initialize_kv_cache_does_not_dcp_shard_mamba_block_table(
     expected: int,
 ):
     """Mamba/GDN block-table rows index global positions, unlike DCP KV."""
-
     max_model_len = 1_048_576
     attention_block_size = 1_536
     mamba_block_size = 16
@@ -237,6 +252,7 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner:
     runner.attn_groups = None
     runner.kv_cache_config = None
     runner.use_aux_hidden_state_outputs = False
+    runner.kv_connector = model_runner_module.NO_OP_KV_CONNECTOR
     return runner
 
 
diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py
index 1176857dffa4..b9cee16ac4d3 100644
--- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py
@@ -56,7 +56,16 @@ def from_mapping(cls, **kwargs: Any) -> "FakeEplbState":
 def _make_runner(**overrides: Any) -> Any:
     runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner)
     runner.device = torch.device("cpu")
-    runner.model_config = SimpleNamespace(model="test-model")
+    runner.model_config = SimpleNamespace(model="test-model", logits_processors=None)
+    runner.req_states = SimpleNamespace(
+        device=runner.device,
+        max_num_reqs=8,
+        vocab_size=0,
+        all_token_ids=None,
+        prompt_len=None,
+        prefill_len=None,
+        total_len=None,
+    )
     runner.load_config = SimpleNamespace(load_format="hf")
     runner.parallel_config = SimpleNamespace(
         enable_eplb=True,
diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py
index eea668aafb6c..0c4b5aff68c2 100644
--- a/tests/v1/worker/test_gpu_pcp_manager.py
+++ b/tests/v1/worker/test_gpu_pcp_manager.py
@@ -1,26 +1,80 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import replace
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
 import torch
 
+from vllm.config import CUDAGraphMode
 from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
 from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module
-from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_context
 from vllm.v1.worker.gpu.pcp_manager import PCPManager
-from vllm.v1.worker.gpu.states import RequestState
 
 
-def _copy_to_cpu(value, out=None, device=None):
+def _copy_to_cpu(value, device=None, dtype=None, out=None):
     tensor = torch.from_numpy(value) if isinstance(value, np.ndarray) else value
     if out is not None:
         return out.copy_(tensor)
     return tensor
 
 
+def _make_config(cudagraph_mode: CUDAGraphMode):
+    return SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=2,
+            decode_context_parallel_size=1,
+            pipeline_parallel_size=1,
+            dcp_comm_backend="ag_rs",
+        ),
+        model_config=SimpleNamespace(
+            use_mla=True,
+            is_encoder_decoder=False,
+            hf_text_config=SimpleNamespace(),
+        ),
+        lora_config=None,
+        speculative_config=None,
+        compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode),
+    )
+
+
+def _make_capture_manager(block_table: torch.Tensor):
+    block_tables = SimpleNamespace(
+        input_block_tables=(block_table,),
+        num_kv_cache_groups=1,
+        kernel_block_sizes=(2,),
+        blocks_per_kv_block=(1,),
+    )
+    manager = PCPManager(
+        pcp_world_size=2,
+        pcp_rank=0,
+        device=torch.device("cpu"),
+        max_num_reqs=4,
+        max_num_tokens=8,
+        block_tables=block_tables,
+    )
+    return manager, block_tables
+
+
+@pytest.mark.parametrize(
+    "cudagraph_mode",
+    [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE],
+)
+def test_validate_config_accepts_decode_only_full_graphs(cudagraph_mode):
+    PCPManager.validate_config(_make_config(cudagraph_mode), supports_mm_inputs=False)
+
+
+def test_validate_config_rejects_full_graph_for_prefills():
+    with pytest.raises(NotImplementedError, match="decode-only routines"):
+        PCPManager.validate_config(
+            _make_config(CUDAGraphMode.FULL), supports_mm_inputs=False
+        )
+
+
 def test_replicated_decode_piecewise_graph_padding(monkeypatch):
     manager = PCPManager(
         pcp_world_size=2,
@@ -28,7 +82,7 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch):
         device=torch.device("cpu"),
         dcp_world_size=1,
     )
-    monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu)
+    monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu)
 
     segments_by_rank, per_rank_num_tokens = manager._build_batch_layout(
         num_scheduled_tokens=np.ones(3, dtype=np.int32),
@@ -103,7 +157,7 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
         device=torch.device("cpu"),
         dcp_world_size=1,
     )
-    monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu)
+    monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu)
 
     with pytest.raises(ValueError, match="smaller than the largest rank-local batch"):
         manager._build_batch_layout(
@@ -115,6 +169,256 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
         )
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels")
+@pytest.mark.parametrize(
+    ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"),
+    [
+        (CUDAGraphMode.NONE, 4, 2, 2),
+        (CUDAGraphMode.PIECEWISE, None, 8, 2),
+        (CUDAGraphMode.FULL, 4, 8, 4),
+    ],
+)
+def test_partition_padding_is_derived_from_batch_descriptor(
+    cg_mode, num_reqs, expected_tokens, expected_reqs
+):
+    """partition_batch derives graph padding from the batch descriptor: no
+    padding without graphs, token padding only for PIECEWISE, and request
+    padding on top for FULL."""
+    device = torch.device("cuda:0")
+    manager = PCPManager(
+        pcp_world_size=2,
+        pcp_rank=0,
+        device=device,
+        max_num_reqs=4,
+        max_num_tokens=8,
+    )
+    global_batch = _make_global_decode_batch(
+        [16, 24], InputBuffers(4, 8, device), device
+    )
+
+    local_batch = manager.partition_batch(
+        global_batch,
+        BatchExecutionDescriptor(
+            cg_mode=cg_mode,
+            num_tokens=8,
+            num_reqs=num_reqs,
+        ),
+    )
+
+    assert local_batch.num_tokens_after_padding == expected_tokens
+    assert local_batch.num_reqs_after_padding == expected_reqs
+
+
+def test_dummy_draft_does_not_reuse_previous_graph_batch():
+    manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32))
+    dummy_batch = InputBatch.make_dummy(1, 4, manager.input_buffers)
+    manager.draft_prefill_batch = replace(dummy_batch)
+
+    manager.prepare_draft_prefill(dummy_batch, dummy_batch.input_ids)
+
+    assert (
+        manager.get_draft_input_buffers(manager.input_buffers) is manager.input_buffers
+    )
+
+
+def test_capture_uses_pcp_persistent_inputs():
+    manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32))
+
+    dummy_batch = InputBatch.make_dummy(
+        num_reqs=4,
+        num_tokens=4,
+        input_buffers=InputBuffers(4, 8, torch.device("cpu")),
+        max_query_len=1,
+    )
+    dummy_batch.input_ids.copy_(torch.arange(4))
+    dummy_batch.positions.fill_(7)
+    input_batch = manager.prepare_inputs_to_capture(dummy_batch)
+
+    assert input_batch is not dummy_batch
+    assert input_batch.req_ids == dummy_batch.req_ids
+    for name in ("input_ids", "positions", "is_padding", "query_start_loc", "seq_lens"):
+        actual = getattr(input_batch, name)
+        torch.testing.assert_close(actual, getattr(dummy_batch, name))
+        assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr()
+
+
+def test_dummy_context_updates_pcp_local_block_tables():
+    global_block_table = torch.full((4, 4), -1, dtype=torch.int32)
+    manager, block_tables = _make_capture_manager(global_block_table)
+    input_batch = InputBatch.make_dummy(
+        num_reqs=2,
+        num_tokens=2,
+        input_buffers=manager.input_buffers,
+        max_query_len=1,
+    )
+    input_batch = manager.prepare_inputs_to_capture(input_batch)
+    local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs)
+
+    set_dummy_context(
+        input_batch,
+        block_tables,
+        context_len=3,
+        num_kv_blocks=16,
+        max_model_len=16,
+        input_block_tables=local_block_tables,
+    )
+
+    torch.testing.assert_close(
+        local_block_tables[0][:2, :2],
+        torch.tensor([[0, 1], [2, 3]], dtype=torch.int32),
+    )
+    assert torch.all(global_block_table == -1)
+
+
+def _rank_rows(
+    pcp_rank: int,
+    pcp_world_size: int,
+    num_scheduled_tokens: np.ndarray,
+    num_computed_tokens: np.ndarray,
+    is_prefilling: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """One PCP+DCP rank's rows as (global request, extent)."""
+    manager = PCPManager(
+        pcp_world_size=pcp_world_size,
+        pcp_rank=pcp_rank,
+        device=torch.device("cpu"),
+        dcp_world_size=2,
+    )
+    query_start_loc_np = np.concatenate([[0], np.cumsum(num_scheduled_tokens)]).astype(
+        np.int32
+    )
+    segments = manager._get_rank_segments(
+        pcp_rank, num_scheduled_tokens, is_prefilling, query_start_loc_np
+    )
+    rows = np.array([segment.global_batch_req_idx for segment in segments])
+    extents = (num_computed_tokens + num_scheduled_tokens)[rows]
+    return rows, extents
+
+
+@pytest.mark.parametrize(
+    ("num_scheduled_tokens", "num_computed_tokens"),
+    [
+        ([32, 32], [0, 0]),
+        ([32, 32, 32, 32], [0, 0, 0, 0]),
+        ([32, 32], [100, 0]),  # one continued, one fresh
+        ([64, 32, 48], [0, 512, 0]),  # ragged lengths and mixed contexts
+        # A replicated prefill BEHIND a split one. Sorting replicated rows first
+        # put request 1 ahead of request 0.
+        ([64, 3], [0, 0]),
+        ([64, 3, 64], [0, 0, 0]),
+    ],
+)
+@pytest.mark.parametrize("pcp_world_size", [2, 4])
+def test_published_row_order_is_identical_on_every_pcp_rank(
+    num_scheduled_tokens, num_computed_tokens, pcp_world_size
+):
+    """Every rank must map its rows to the same global requests, in the same order."""
+    num_scheduled_tokens = np.array(num_scheduled_tokens, dtype=np.int32)
+    num_computed_tokens = np.array(num_computed_tokens, dtype=np.int32)
+    is_prefilling = np.ones(len(num_scheduled_tokens), dtype=np.bool_)
+
+    orders = [
+        _rank_rows(
+            rank,
+            pcp_world_size,
+            num_scheduled_tokens,
+            num_computed_tokens,
+            is_prefilling,
+        )[0]
+        for rank in range(pcp_world_size)
+    ]
+    for rank, order in enumerate(orders[1:], start=1):
+        assert np.array_equal(orders[0], order), (
+            f"rank {rank} rows map to {order.tolist()}, rank 0 to {orders[0].tolist()}"
+        )
+    # Grouped by request and ascending: what the indexer plan indexes by.
+    assert np.all(np.diff(orders[0]) >= 0)
+
+
+def test_published_row_order_puts_every_decode_before_every_prefill():
+    """split_decodes_and_prefills takes the FIRST prefilling row as the boundary."""
+    # req 0 is a continued prefill, req 1 is a decode.
+    req_idx, _ = _rank_rows(
+        pcp_rank=0,
+        pcp_world_size=2,
+        num_scheduled_tokens=np.array([32, 1], dtype=np.int32),
+        num_computed_tokens=np.array([100, 20], dtype=np.int32),
+        is_prefilling=np.array([True, False], dtype=np.bool_),
+    )
+    # The decode (request 1) must be row 0.
+    assert req_idx[0] == 1
+
+
+def test_split_prefill_rows_repeat_the_request_and_its_full_extent():
+    """Equal adjacent request indices are what let a backend share a KV region."""
+    req_idx, extents = _rank_rows(
+        pcp_rank=1,
+        pcp_world_size=2,
+        num_scheduled_tokens=np.array([64, 3], dtype=np.int32),
+        num_computed_tokens=np.array([100, 0], dtype=np.int32),
+        is_prefilling=np.ones(2, dtype=np.bool_),
+    )
+    assert req_idx.tolist() == [0, 0, 1]
+    assert extents.tolist() == [164, 164, 3]
+
+
+@pytest.mark.parametrize("query_len", [1, 2, 3, 5, 6, 9])
+def test_dcp_replicates_prefills_too_short_to_split(query_len):
+    """A prefill that cannot fill 2*pcp chunks is replicated, not split."""
+    pcp_world_size = 2
+    num_scheduled_tokens = np.array([query_len], dtype=np.int32)
+    is_prefilling = np.ones(1, dtype=np.bool_)
+
+    rows_per_rank = []
+    for rank in range(pcp_world_size):
+        manager = PCPManager(
+            pcp_world_size=pcp_world_size,
+            pcp_rank=rank,
+            device=torch.device("cpu"),
+            dcp_world_size=2,
+        )
+        assert manager.replicated_requests(num_scheduled_tokens, is_prefilling)[0]
+        rows = list(
+            manager._iter_rank_chunks(rank, num_scheduled_tokens, is_prefilling)
+        )
+        assert rows == [(0, 0, query_len)]
+        rows_per_rank.append(rows)
+    assert rows_per_rank[0] == rows_per_rank[1]
+
+
+@pytest.mark.parametrize("pcp_world_size", [2, 4, 8])
+@pytest.mark.parametrize(
+    "query_len", [16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1000, 4097]
+)
+def test_pcp_first_chunk_row_is_never_short(pcp_world_size, query_len):
+    num_scheduled_tokens = np.array([query_len], dtype=np.int32)
+    is_prefilling = np.ones(1, dtype=np.bool_)
+
+    charged_per_rank = []
+    for rank in range(pcp_world_size):
+        manager = PCPManager(
+            pcp_world_size=pcp_world_size,
+            pcp_rank=rank,
+            device=torch.device("cpu"),
+            dcp_world_size=pcp_world_size,
+        )
+        chunk_lens = [
+            chunk_len
+            for _, _, chunk_len in manager._iter_rank_chunks(
+                rank, num_scheduled_tokens, is_prefilling
+            )
+        ]
+        assert chunk_lens, f"rank {rank} got no rows for {query_len=}"
+        assert chunk_lens[0] == max(chunk_lens), (
+            f"rank {rank} emitted a short first chunk for {query_len=}: {chunk_lens}"
+        )
+        charged_per_rank.append(len(chunk_lens) * chunk_lens[0])
+
+    assert len(set(charged_per_rank)) == 1, (
+        f"ranks would chunk differently for {query_len=}: {charged_per_rank}"
+    )
+
+
 def _make_global_decode_batch(
     num_computed_tokens: list[int], buffers: InputBuffers, device: torch.device
 ) -> InputBatch:
@@ -125,6 +429,7 @@ def _make_global_decode_batch(
 
     base = InputBatch.make_dummy(num_reqs, num_tokens, buffers)
     buffers.seq_lens[:num_reqs] = torch.from_numpy(seq_lens_np).to(device)
+    buffers.positions[:num_reqs] = torch.tensor(num_computed_tokens, device=device)
     query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32)
     buffers.query_start_loc[: num_reqs + 1] = torch.from_numpy(query_start_loc_np).to(
         device
@@ -166,19 +471,10 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
     the field afterwards from the PCP-owned buffers.
     """
     device = torch.device("cuda:0")
-    req_states = RequestState(
-        max_num_reqs=4,
-        max_model_len=64,
-        max_num_batched_tokens=8,
-        num_speculative_steps=1,
-        vocab_size=8,
-        device=device,
-    )
     manager = PCPManager(
         pcp_world_size=2,
         pcp_rank=0,
         device=device,
-        req_states=req_states,
         max_num_reqs=4,
         max_num_tokens=8,
         dcp_world_size=2,
@@ -191,14 +487,29 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
     global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2]
     global_batch.dcp_local_seq_lens.fill_(-1)
 
-    local_batch = manager.partition_batch(global_batch, padded_num_tokens=2)
+    local_batch = manager.partition_batch(
+        global_batch,
+        BatchExecutionDescriptor(
+            cg_mode=CUDAGraphMode.FULL,
+            num_tokens=4,
+            num_reqs=4,
+        ),
+    )
 
     assert local_batch.dcp_local_seq_lens is None
-    assert local_batch.seq_lens.tolist() == [17, 25]
+    assert local_batch.seq_lens.tolist() == [17, 25, 0, 0]
+    assert local_batch.dcp_local_seq_lens_cpu_upper_bound is not None
+    expected_dcp_upper_bound = get_dcp_local_seq_lens(
+        torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1
+    )
+    assert torch.equal(
+        local_batch.dcp_local_seq_lens_cpu_upper_bound,
+        torch.cat([expected_dcp_upper_bound, torch.zeros(2, dtype=torch.int32)]),
+    )
 
     # What execute_model does next: derive DCP metadata from the final batch
     # on the PCP-owned buffers.
-    local_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
+    local_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens(
         manager.input_buffers.dcp_local_seq_lens,
         local_batch.seq_lens,
         local_batch.num_reqs,
@@ -208,7 +519,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
         num_reqs_padded=local_batch.num_reqs_after_padding,
     )
     expected = get_dcp_local_seq_lens(
-        torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1
+        torch.tensor([17, 25, 0, 0], dtype=torch.int32), 2, 0, 1
     )
     assert local_batch.dcp_local_seq_lens is not None
     assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected)
diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py
index 347b5b379adf..5b407e80e903 100644
--- a/tests/v1/worker/test_gpu_profiler.py
+++ b/tests/v1/worker/test_gpu_profiler.py
@@ -7,6 +7,7 @@
 from uuid import UUID
 
 import pytest
+import torch
 from pydantic import ValidationError
 
 from vllm.config import (
@@ -17,16 +18,21 @@
 )
 from vllm.config.profiler import _is_uri_path
 from vllm.platforms import current_platform
-from vllm.profiler.wrapper import ProtonProfilerWrapper, WorkerProfiler
+from vllm.profiler.wrapper import (
+    ProtonProfilerWrapper,
+    TorchProfilerWrapper,
+    WorkerProfiler,
+    create_worker_profiler,
+    validate_worker_profiler_config,
+)
 from vllm.v1.core.sched.output import CachedRequestData
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 from vllm.v1.worker.gpu_worker import Worker
+from vllm.v1.worker.xpu_worker import XPUWorker
 
 
 class ConcreteWorkerProfiler(WorkerProfiler):
-    """
-    A basic implementation of a worker profiler for testing purposes.
-    """
+    """A basic implementation of a worker profiler for testing purposes."""
 
     def __init__(self, profiler_config: ProfilerConfig):
         self.start_call_count = 0
@@ -53,6 +59,213 @@ def default_profiler_config():
     )
 
 
+def test_torch_profiler_rebuilds_one_shot_profiler_each_round(tmp_path):
+    config = ProfilerConfig(
+        profiler="torch",
+        torch_profiler_dir=str(tmp_path),
+        torch_profiler_activities=["CUDA"],
+        torch_profiler_dump_cuda_time_total=False,
+        warmup_iterations=2,
+    )
+    profilers = [MagicMock(), MagicMock()]
+    for profiler in profilers:
+        profiler.profiler = MagicMock()
+
+    with patch(
+        "vllm.profiler.wrapper.torch.profiler.profile", side_effect=profilers
+    ) as profile:
+        wrapper = TorchProfilerWrapper(
+            config, worker_name="worker", local_rank=0, activities=["CUDA"]
+        )
+        profile.assert_not_called()
+
+        wrapper.start()
+        assert not wrapper.should_annotate
+        assert not wrapper._profiler_step()
+        wrapper.stop()
+        wrapper.start()
+        assert not wrapper._profiler_step()
+        wrapper.stop()
+
+    assert profile.call_count == 2
+    assert profile.call_args.kwargs["activities"] == [
+        torch.profiler.ProfilerActivity.CUDA
+    ]
+    for profiler in profilers:
+        profiler.start.assert_called_once_with()
+        profiler.step.assert_called_once_with()
+        profiler.stop.assert_called_once_with()
+
+
+def test_torch_profiler_records_each_profile_round(tmp_path):
+    traces: list[torch.profiler.profile] = []
+    wrapper = TorchProfilerWrapper(
+        ProfilerConfig(
+            profiler="torch",
+            torch_profiler_dir=str(tmp_path),
+            torch_profiler_dump_cuda_time_total=False,
+        ),
+        worker_name="worker",
+        local_rank=1,
+        activities=["CPU"],
+        on_trace_ready=traces.append,
+    )
+
+    for run in range(2):
+        wrapper.start()
+        with torch.profiler.record_function(f"run_{run}"):
+            pass
+        wrapper.stop()
+
+    assert len(traces) == 2
+    for run, trace in enumerate(traces):
+        assert {
+            event.name for event in trace.events() if event.name.startswith("run_")
+        } == {f"run_{run}"}
+
+
+@pytest.mark.parametrize(
+    "activities", [["CPU", "CUDA"], ["CUDA"], ["CPU", "XPU"], ["XPU"]]
+)
+@pytest.mark.parametrize("dump_device_time", [True, False])
+def test_torch_profiler_device_summary(tmp_path, capsys, activities, dump_device_time):
+    """Device summaries honor the dump option on both CUDA and XPU."""
+    config = ProfilerConfig(
+        profiler="torch",
+        torch_profiler_dir=str(tmp_path),
+        torch_profiler_dump_cuda_time_total=dump_device_time,
+    )
+    with patch("vllm.profiler.wrapper.torch.profiler.profile") as profile:
+        profile.return_value.key_averages.return_value.table.return_value = (
+            "device times"
+        )
+        wrapper = TorchProfilerWrapper(
+            config, worker_name="worker", local_rank=0, activities=activities
+        )
+        wrapper.start()
+        wrapper.stop()
+
+    summary = tmp_path / "profiler_out_0.txt"
+    assert summary.exists() == dump_device_time
+    assert ("device times" in capsys.readouterr().out) == dump_device_time
+    if dump_device_time:
+        assert summary.read_text() == "device times\n"
+
+
+@pytest.mark.parametrize(
+    "activities",
+    [[], ["CPU", "CPU"], ["INVALID"]],
+)
+def test_torch_profiler_activities_reject_invalid_values(activities):
+    with pytest.raises(ValueError, match="torch_profiler_activities"):
+        ProfilerConfig(
+            profiler="torch",
+            torch_profiler_dir="/tmp/mock",
+            torch_profiler_activities=activities,
+        )
+
+
+def test_torch_profiler_activities_require_torch_profiler():
+    with pytest.raises(ValueError, match="only applicable"):
+        ProfilerConfig(torch_profiler_activities=["CPU"])
+
+
+@pytest.mark.parametrize(
+    ("device_type", "activities"),
+    [
+        ("cuda", ["XPU"]),
+        ("xpu", ["CUDA"]),
+        ("cpu", ["CUDA"]),
+        ("cpu", ["XPU"]),
+    ],
+)
+def test_worker_rejects_unsupported_activities_at_startup(device_type, activities):
+    config = ProfilerConfig(
+        profiler="torch",
+        torch_profiler_dir="/tmp/mock",
+        torch_profiler_activities=activities,
+    )
+
+    with (
+        patch.object(current_platform, "device_type", device_type),
+        pytest.raises(ValueError, match="Unsupported torch profiler activities"),
+    ):
+        validate_worker_profiler_config(config)
+
+
+@pytest.mark.parametrize("device_type", ["cpu", "xpu"])
+def test_worker_rejects_cuda_profiler_on_other_devices(device_type):
+    with (
+        patch.object(current_platform, "device_type", device_type),
+        pytest.raises(ValueError, match="Unsupported profiler type"),
+    ):
+        validate_worker_profiler_config(ProfilerConfig(profiler="cuda"))
+
+
+@pytest.mark.parametrize(
+    ("device_type", "activities", "expected"),
+    [
+        ("cuda", None, ("CPU", "CUDA")),
+        ("xpu", None, ("CPU", "XPU")),
+        ("cpu", None, ("CPU",)),
+        ("cuda", ["CUDA"], ("CUDA",)),
+        ("xpu", ["XPU"], ("XPU",)),
+        ("cuda", ["CPU"], ("CPU",)),
+        ("xpu", ["CPU"], ("CPU",)),
+        ("cpu", ["CPU"], ("CPU",)),
+    ],
+)
+def test_worker_creates_platform_torch_profiler(device_type, activities, expected):
+    config = ProfilerConfig(
+        profiler="torch",
+        torch_profiler_dir="/tmp/mock",
+        torch_profiler_activities=activities,
+    )
+
+    with (
+        patch.object(current_platform, "device_type", device_type),
+        patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper,
+    ):
+        validate_worker_profiler_config(config)
+        profiler = create_worker_profiler(
+            config,
+            worker_name="rank0",
+            local_rank=0,
+        )
+
+    assert profiler is wrapper.return_value
+    wrapper.assert_called_once_with(
+        config,
+        worker_name="rank0",
+        local_rank=0,
+        activities=expected,
+    )
+
+
+@pytest.mark.parametrize("worker_type", [Worker, XPUWorker])
+def test_worker_reuses_torch_wrapper_across_profile_rounds(worker_type):
+    worker = object.__new__(worker_type)
+    worker.rank = 0
+    worker.local_rank = 0
+    worker.profiler = None
+    worker.profiler_config = ProfilerConfig(
+        profiler="torch", torch_profiler_dir="/tmp/mock"
+    )
+
+    with (
+        patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="rank0"),
+        patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper,
+    ):
+        worker.profile()
+        worker.profile(is_start=False)
+        worker.profile()
+
+    assert worker.profiler is wrapper.return_value
+    wrapper.assert_called_once()
+    assert wrapper.return_value.start.call_count == 2
+    wrapper.return_value.stop.assert_called_once_with()
+
+
 def test_immediate_start_stop(default_profiler_config):
     """Test standard start without delay."""
     profiler = ConcreteWorkerProfiler(default_profiler_config)
@@ -320,13 +533,17 @@ def test_detailed_format_mixed(self):
             "execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)"
         )
 
-    def test_skips_annotation_work_after_profiler_stops(self):
+    def test_skips_annotation_work_when_profiler_does_not_annotate(self):
         worker = MagicMock()
-        worker.profiler.is_running = False
+        worker.profiler.should_annotate = False
 
-        context = Worker.annotate_profile(worker, scheduler_output=None)
+        with patch(
+            "vllm.v1.worker.gpu_worker.compute_iteration_details"
+        ) as compute_iteration_details:
+            context = Worker.annotate_profile(worker, scheduler_output=None)
 
         worker.profiler.step.assert_called_once_with()
+        compute_iteration_details.assert_not_called()
         worker.profiler.annotate_context_manager.assert_not_called()
         assert isinstance(context, nullcontext)
 
@@ -350,17 +567,24 @@ def test_profiler_entered_during_capture():
 
 
 def make_proton(session_id: int | None = 7):
+    data = SimpleNamespace(
+        advance_phase=Mock(side_effect=range(1, 100)),
+        clear=Mock(),
+        get=Mock(return_value={"traceEvents": []}),
+        get_msgpack=Mock(return_value=b"profile"),
+    )
     return SimpleNamespace(
         start=Mock(return_value=session_id),
         activate=Mock(),
         deactivate=Mock(),
         finalize=Mock(),
         scope=Mock(return_value=nullcontext()),
+        data=data,
     )
 
 
 def make_proton_wrapper(
-    tmp_path, proton=None, triton_version="3.6.0", **config_overrides
+    tmp_path, proton=None, triton_version="3.7.0", **config_overrides
 ):
     proton = proton or make_proton()
     config = ProfilerConfig(
@@ -387,6 +611,15 @@ def import_module(name):
     reason="Proton profiling tests require an NVIDIA CUDA platform.",
 )
 
+# Proton subscribes to CUPTI itself, and a process gets one CUPTI subscriber.
+# CI's kernel-launch recorder is another, injected through CUDA_INJECTION64_PATH
+# on recording runs, and there Proton's cuptiSubscribe fails with error 39.
+_requires_no_injected_cupti_tool = pytest.mark.skipif(
+    bool(os.environ.get("CUDA_INJECTION64_PATH")),
+    reason="Another CUPTI tool is injected (CUDA_INJECTION64_PATH); "
+    "Proton has to be the only one.",
+)
+
 
 @_requires_cuda_for_proton
 class TestProtonConfig:
@@ -415,6 +648,13 @@ def test_normalizes_local_output_directory(self, tmp_path, monkeypatch):
                 },
                 "requires proton_data",
             ),
+            (
+                {
+                    "proton_data": "trace",
+                    "proton_graph_attribution": True,
+                },
+                "requires proton_data='tree'",
+            ),
         ],
     )
     def test_rejects_invalid_option_combinations(self, tmp_path, options, message):
@@ -450,6 +690,7 @@ def test_rejects_invalid_typed_options(self, field, tmp_path):
             ("proton_mode", "pcsampling"),
             ("proton_hook", "triton"),
             ("proton_output_format", "chrome_trace"),
+            ("proton_graph_attribution", True),
         ],
     )
     def test_rejects_proton_options_for_other_profilers(self, field, value):
@@ -478,15 +719,101 @@ def test_rejects_proton_on_non_cuda_platforms(self, tmp_path):
                 compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE),
             )
 
-    def test_rejects_proton_when_cuda_graphs_are_enabled(self, tmp_path):
-        with pytest.raises(ValueError, match="requires CUDA graphs to be disabled"):
+    @pytest.mark.parametrize("encoder_only", [False, True])
+    @pytest.mark.parametrize("attribution", [False, True])
+    def test_cuda_graphs_require_attribution(self, tmp_path, encoder_only, attribution):
+        # Encoder graphs are independent of the decoder's cudagraph_mode.
+        expected = (
+            nullcontext()
+            if attribution
+            else pytest.raises(
+                ValueError, match="requires proton_graph_attribution=True"
+            )
+        )
+        with expected:
+            VllmConfig(
+                profiler_config=ProfilerConfig(
+                    profiler="proton",
+                    proton_profiler_dir=str(tmp_path),
+                    proton_graph_attribution=attribution,
+                ),
+                compilation_config=CompilationConfig(
+                    cudagraph_mode=(
+                        CUDAGraphMode.NONE if encoder_only else CUDAGraphMode.FULL
+                    ),
+                    cudagraph_mm_encoder=encoder_only,
+                ),
+            )
+
+    def test_validates_default_cuda_graph_mode_after_resolution(self, tmp_path):
+        with pytest.raises(ValueError, match="requires proton_graph_attribution=True"):
             VllmConfig(
                 profiler_config=ProfilerConfig(
                     profiler="proton", proton_profiler_dir=str(tmp_path)
                 ),
+            )
+
+    @pytest.mark.parametrize(
+        "mode",
+        [
+            "periodic_flushing",
+            "periodic_flushing:format=hatchet",
+            "PERIODIC_FLUSHING:format=hatchet",
+        ],
+    )
+    def test_graph_attribution_rejects_periodic_flushing(self, tmp_path, mode):
+        # Reject before Proton's native phase manager can abort the worker.
+        with pytest.raises(ValueError, match="incompatible with periodic_flushing"):
+            ProfilerConfig(
+                profiler="proton",
+                proton_profiler_dir=str(tmp_path),
+                proton_graph_attribution=True,
+                proton_mode=mode,
+            )
+
+    @pytest.mark.parametrize("attribution", [False, True])
+    @pytest.mark.parametrize(
+        "mode", ["pcsampling", "pcsampling:interval=100", "PcSampling:interval=100"]
+    )
+    @pytest.mark.parametrize("encoder_only", [False, True])
+    def test_pcsampling_requires_graphs_disabled(
+        self, tmp_path, attribution, mode, encoder_only
+    ):
+        with pytest.raises(ValueError, match="PC sampling requires CUDA graphs"):
+            VllmConfig(
+                profiler_config=ProfilerConfig(
+                    profiler="proton",
+                    proton_profiler_dir=str(tmp_path),
+                    proton_graph_attribution=attribution,
+                    proton_mode=mode,
+                ),
                 compilation_config=CompilationConfig(
-                    cudagraph_mode=CUDAGraphMode.PIECEWISE
+                    cudagraph_mode=CUDAGraphMode.NONE
+                    if encoder_only
+                    else CUDAGraphMode.FULL,
+                    cudagraph_mm_encoder=encoder_only,
+                ),
+            )
+
+    def test_ordinary_proton_keeps_mrv1_graph_restriction(self, tmp_path, monkeypatch):
+        monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
+        with pytest.raises(ValueError, match="requires proton_graph_attribution=True"):
+            VllmConfig(
+                profiler_config=ProfilerConfig(
+                    profiler="proton", proton_profiler_dir=str(tmp_path)
                 ),
+                compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.FULL),
+            )
+
+    def test_graph_attribution_rejects_mrv1(self, tmp_path, monkeypatch):
+        monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
+        with pytest.raises(ValueError, match="requires the V2 model runner"):
+            VllmConfig(
+                profiler_config=ProfilerConfig(
+                    profiler="proton",
+                    proton_profiler_dir=str(tmp_path),
+                    proton_graph_attribution=True,
+                )
             )
 
 
@@ -564,7 +891,20 @@ def test_output_names_are_unique_across_worker_restarts(self, tmp_path):
     )
     def test_newer_features_reject_triton_3_6(self, tmp_path, option, value, feature):
         with pytest.raises(RuntimeError, match=feature):
-            make_proton_wrapper(tmp_path, **{option: value})
+            make_proton_wrapper(tmp_path, triton_version="3.6.0", **{option: value})
+
+    @pytest.mark.parametrize("version", ["3.6.0", "unknown"])
+    def test_graph_attribution_requires_phase_api(self, tmp_path, version):
+        with pytest.raises(RuntimeError, match="requires Triton >= 3.7"):
+            make_proton_wrapper(
+                tmp_path, triton_version=version, proton_graph_attribution=True
+            )
+
+    def test_ordinary_profiling_still_supports_triton_3_6(self, tmp_path):
+        wrapper, proton = make_proton_wrapper(tmp_path, triton_version="3.6.0")
+        wrapper.start()
+        wrapper.stop()
+        proton.finalize.assert_called_once_with(session=7)
 
     @pytest.mark.parametrize(
         ("option", "value"),
@@ -583,6 +923,125 @@ def test_rejects_output_format_when_finalize_lacks_capability(self, tmp_path):
         with pytest.raises(RuntimeError, match="does not support selecting"):
             make_proton_wrapper(tmp_path, proton, proton_output_format="hatchet")
 
+    def test_cuda_graph_tree_phase_is_written_at_stop(self, tmp_path):
+        wrapper, proton = make_proton_wrapper(
+            tmp_path,
+            proton_context="python",
+            proton_graph_attribution=True,
+        )
+        with wrapper.capture_cuda_graphs():
+            proton.start.assert_called_once()
+            proton.deactivate.assert_not_called()
+
+        capture_args = proton.start.call_args.kwargs
+        assert capture_args["context"] == "python"
+        assert capture_args["data"] == "tree"
+        proton.data.advance_phase.assert_called_once_with(7)
+        proton.deactivate.assert_called_once_with(session=7, flushing=True)
+        proton.data.clear.assert_called_once_with(7, 0)
+
+        wrapper.start()
+        wrapper.stop()
+        wrapper.start()
+        wrapper.stop()
+
+        assert proton.data.get.call_args_list == [call(7, 1), call(7, 2)]
+        proton.data.clear.assert_has_calls([call(7, 0), call(7, 1), call(7, 2)])
+        output_names = sorted(tmp_path.glob("proton_rank_3_*.hatchet"))
+        assert len(output_names) == 2
+        assert output_names[0].name.endswith("_run0.hatchet")
+        assert output_names[1].name.endswith("_run1.hatchet")
+
+    def test_capture_is_noop_without_opt_in(self, tmp_path):
+        wrapper, proton = make_proton_wrapper(tmp_path)
+        with wrapper.capture_cuda_graphs():
+            pass
+        proton.start.assert_not_called()
+
+    @pytest.mark.parametrize("delay", [0, 2])
+    def test_duplicate_start_preserves_output_prefix(self, tmp_path, delay):
+        wrapper, proton = make_proton_wrapper(
+            tmp_path, proton_graph_attribution=True, delay_iterations=delay
+        )
+        with wrapper.capture_cuda_graphs():
+            pass
+        wrapper.set_output_name("first")
+        wrapper.start()
+        wrapper.set_output_name("duplicate")
+        wrapper.start()
+        for _ in range(delay):
+            wrapper.step()
+        wrapper.stop()
+        assert len(list(tmp_path.glob("proton_first_*.hatchet"))) == 1
+        assert not list(tmp_path.glob("proton_duplicate_*"))
+        wrapper.set_output_name("second")
+        wrapper.start()
+        for _ in range(delay):
+            wrapper.step()
+        wrapper.stop()
+        assert len(list(tmp_path.glob("proton_second_*.hatchet"))) == 1
+
+    def test_failed_export_clears_activity_and_allows_next_interval(self, tmp_path):
+        wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)
+        with wrapper.capture_cuda_graphs():
+            pass
+        proton.data.get.side_effect = [OSError("disk full"), []]
+        wrapper.start()
+        wrapper.stop()
+        proton.data.clear.assert_called_with(7, 1)
+        wrapper.start()
+        wrapper.stop()
+        proton.data.clear.assert_called_with(7, 2)
+        assert len(list(tmp_path.glob("*_run1.hatchet"))) == 1
+
+    def test_cuda_graph_context_deactivates_after_capture_error(self, tmp_path):
+        wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)
+
+        with (
+            pytest.raises(RuntimeError, match="capture failed"),
+            wrapper.capture_cuda_graphs(),
+        ):
+            raise RuntimeError("capture failed")
+
+        proton.deactivate.assert_called_once_with(session=7, flushing=True)
+
+    def test_cuda_graph_capture_deactivates_when_phase_advance_fails(self, tmp_path):
+        wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)
+        proton.data.advance_phase.side_effect = RuntimeError("advance failed")
+
+        with (
+            pytest.raises(RuntimeError, match="advance failed"),
+            wrapper.capture_cuda_graphs(),
+        ):
+            pass
+
+        proton.deactivate.assert_called_once_with(session=7, flushing=True)
+        proton.data.clear.assert_not_called()
+
+    def test_cuda_graph_stop_deactivates_when_phase_advance_fails(self, tmp_path):
+        wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)
+        with wrapper.capture_cuda_graphs():
+            pass
+        proton.data.clear.reset_mock()
+        proton.data.advance_phase.side_effect = RuntimeError("advance failed")
+
+        wrapper._start()
+        with pytest.raises(RuntimeError, match="advance failed"):
+            wrapper._stop()
+
+        proton.deactivate.assert_called_with(session=7, flushing=True)
+        proton.data.clear.assert_not_called()
+
+    def test_shutdown_finalizes_cuda_graph_capture_session(self, tmp_path):
+        wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)
+        with wrapper.capture_cuda_graphs():
+            pass
+
+        wrapper.shutdown()
+        wrapper.shutdown()
+
+        proton.finalize.assert_called_once_with(session=7)
+
     def test_missing_proton_has_actionable_error(self, tmp_path):
         config = ProfilerConfig(profiler="proton", proton_profiler_dir=str(tmp_path))
         with (
@@ -606,17 +1065,18 @@ def test_scope_annotations_delegate_to_proton(self, tmp_path):
 
 @_requires_cuda_for_proton
 def test_gpu_worker_creates_proton_profiler():
-    worker = MagicMock()
+    worker = object.__new__(Worker)
     worker.rank = 1
+    worker.local_rank = 1
     worker.profiler = None
-    worker.profiler_config.profiler = "proton"
+    worker.profiler_config = MagicMock(profiler="proton")
 
     with (
         patch(
             "vllm.distributed.utils.get_worker_rank_suffix",
             return_value="rank1",
         ),
-        patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper,
+        patch("vllm.profiler.wrapper.ProtonProfilerWrapper") as wrapper,
     ):
         Worker.profile(worker)
 
@@ -626,18 +1086,20 @@ def test_gpu_worker_creates_proton_profiler():
 
 @_requires_cuda_for_proton
 def test_gpu_worker_recreates_proton_profiler_for_each_run():
-    worker = MagicMock()
+    worker = object.__new__(Worker)
     worker.rank = 1
+    worker.local_rank = 1
     worker.profiler = None
-    worker.profiler_config.profiler = "proton"
+    worker.profiler_config = MagicMock(profiler="proton")
 
     with (
         patch(
             "vllm.distributed.utils.get_worker_rank_suffix",
             return_value="rank1",
         ),
-        patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper,
+        patch("vllm.profiler.wrapper.ProtonProfilerWrapper") as wrapper,
     ):
+        wrapper.return_value.has_cuda_graph_session = False
         Worker.profile(worker, profile_prefix="first")
         Worker.profile(worker, is_start=False)
         Worker.profile(worker, profile_prefix="second")
@@ -647,3 +1109,158 @@ def test_gpu_worker_recreates_proton_profiler_for_each_run():
         call(worker.profiler_config, worker_name="second_rank1"),
     ]
     assert wrapper.return_value.start.call_count == 2
+
+
+@_requires_cuda_for_proton
+def test_gpu_worker_reuses_cuda_graph_proton_session():
+    worker = MagicMock()
+    worker.rank = 1
+    worker.profiler = MagicMock(spec=ProtonProfilerWrapper)
+    worker.profiler.has_cuda_graph_session = True
+    worker.profiler_config.profiler = "proton"
+
+    with patch(
+        "vllm.distributed.utils.get_worker_rank_suffix",
+        return_value="rank1",
+    ):
+        Worker.profile(worker, profile_prefix="first")
+        Worker.profile(worker, is_start=False)
+
+    worker.profiler.set_output_name.assert_called_once_with("first_rank1")
+    worker.profiler.start.assert_called_once_with()
+    worker.profiler.stop.assert_called_once_with()
+
+
+@_requires_cuda_for_proton
+@pytest.mark.parametrize("runner", ["attribution_off", "v1", "no_capture"])
+def test_proton_not_initialized_without_capture(runner):
+    worker = MagicMock()
+    worker.profiler = None
+    worker.profiler_config.profiler = "proton"
+    worker.profiler_config.proton_graph_attribution = runner != "attribution_off"
+    worker.use_v2_model_runner = runner != "v1"
+    worker.model_runner.needs_cudagraph_capture.return_value = runner != "no_capture"
+
+    context = Worker._get_cudagraph_capture_context(worker)
+
+    assert worker.profiler is None
+    with context:
+        pass
+
+
+@_requires_cuda_for_proton
+def test_proton_initializes_before_cuda_graph_capture():
+    class FakeProtonProfiler:
+        def __init__(self, config, worker_name):
+            self.config = config
+            self.worker_name = worker_name
+            self.capture_context = nullcontext()
+
+        def capture_cuda_graphs(self):
+            return self.capture_context
+
+    worker = MagicMock()
+    worker.rank = 2
+    worker.profiler = None
+    worker.profiler_config.profiler = "proton"
+    worker.profiler_config.proton_graph_attribution = True
+    worker.use_v2_model_runner = True
+    worker.model_runner.needs_cudagraph_capture.return_value = True
+
+    with (
+        patch(
+            "vllm.distributed.utils.get_worker_rank_suffix",
+            return_value="rank2",
+        ),
+        patch(
+            "vllm.profiler.wrapper.ProtonProfilerWrapper",
+            FakeProtonProfiler,
+        ),
+    ):
+        context = Worker._get_cudagraph_capture_context(worker)
+
+    assert worker.profiler.config is worker.profiler_config
+    assert worker.profiler.worker_name == "rank2"
+    assert context is worker.profiler.capture_context
+
+
+@_requires_cuda_for_proton
+@_requires_no_injected_cupti_tool
+@pytest.mark.parametrize("context", ["shadow", "python"])
+@pytest.mark.parametrize("output_format", ["hatchet", "hatchet_msgpack"])
+def test_proton_cuda_graph_replay_attribution_on_gpu(tmp_path, context, output_format):
+    """Both intervals contain replay kernels, without capture-only activity."""
+    import json
+
+    import torch
+    import triton
+    import triton.profiler as proton
+    from packaging.version import Version
+
+    if Version(triton.__version__) < Version("3.7"):
+        pytest.skip("Graph attribution requires Triton >= 3.7")
+
+    wrapper = ProtonProfilerWrapper(
+        ProfilerConfig(
+            profiler="proton",
+            proton_profiler_dir=str(tmp_path),
+            proton_graph_attribution=True,
+            proton_context=context,
+            proton_output_format=output_format,
+        ),
+        worker_name="gpu",
+    )
+    worker = SimpleNamespace(
+        profiler=wrapper,
+        use_v2_model_runner=True,
+    )
+    x = torch.ones(1024, device="cuda")
+    graph = torch.cuda.CUDAGraph()
+
+    def capture_only():
+        x.add_(1)
+
+    def captured_add():
+        x.add_(1)
+
+    def kernel_metrics(value):
+        if isinstance(value, list):
+            return [metric for child in value for metric in kernel_metrics(child)]
+        if isinstance(value, dict):
+            metrics = value.get("metrics", {})
+            return ([metrics] if metrics.get("count", 0) else []) + [
+                metric for child in value.values() for metric in kernel_metrics(child)
+            ]
+        return []
+
+    try:
+        with Worker._get_cudagraph_capture_context(worker):
+            with proton.scope("capture_only"):
+                capture_only()
+            torch.accelerator.synchronize()
+            with torch.cuda.graph(graph), proton.scope("captured_add"):
+                captured_add()
+        for run in range(2):
+            wrapper.start()
+            with wrapper.annotate_context_manager(f"replay_{run}"):
+                graph.replay()
+            wrapper.stop()
+            (path,) = tmp_path.glob(f"*_run{run}.{output_format}")
+            if output_format == "hatchet_msgpack":
+                import msgpack
+
+                data = msgpack.unpackb(path.read_bytes())
+            else:
+                data = json.loads(path.read_text())
+            serialized = json.dumps(data)
+            assert "<captured_at>" in serialized
+            assert "captured_add" in serialized
+            assert "capture_only" not in serialized
+            assert f"replay_{1 - run}" not in serialized
+            metrics = kernel_metrics(data)
+            assert sum(metric["count"] for metric in metrics) == 1
+            assert sum(metric["time (ns)"] for metric in metrics) > 0
+        torch.testing.assert_close(x, torch.full_like(x, 4))
+    finally:
+        wrapper.shutdown()
+    assert not list(tmp_path.glob(".proton_cuda_graph_session*"))
diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py
index 2e5adfabd13c..efd2b59d75ed 100644
--- a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py
+++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py
@@ -47,6 +47,9 @@ def test_chunked_scores_match_full_batch(logprobs_mode: str):
                 [np.arange(count, dtype=np.int32) for count in num_logits_per_req]
             )
         ).to(device),
+        seq_lens_cpu_upper_bound=torch.from_numpy(
+            np.array([10, 20, 30, 40], dtype=np.int32)
+        ),
     )
     rejection_sampler = object.__new__(RejectionSampler)
     rejection_sampler.sampler = SimpleNamespace(logprobs_mode=logprobs_mode)
diff --git a/tests/v1/worker/test_gpu_sampler_flags.py b/tests/v1/worker/test_gpu_sampler_flags.py
index 3fd1cb3b3c68..3f05a2cfb59f 100644
--- a/tests/v1/worker/test_gpu_sampler_flags.py
+++ b/tests/v1/worker/test_gpu_sampler_flags.py
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Sequence
+from types import SimpleNamespace
+
 import numpy as np
 import pytest
 import torch
@@ -10,6 +13,10 @@
     pytest.skip("CUDA required for sampler flag tests", allow_module_level=True)
 
 from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.logits_processor import (
+    LogitsContext,
+    LogitsProcessor,
+)
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 from vllm.v1.worker.gpu.states import RequestState
 
@@ -23,7 +30,7 @@ class MockReasoningConfig:
     natural_reasoning_end_token_ids = [91]
 
 
-def _make_sampler() -> Sampler:
+def _make_sampler(custom_logits_processors: Sequence[LogitsProcessor] = ()) -> Sampler:
     req_states = RequestState(
         max_num_reqs=4,
         max_model_len=64,
@@ -33,11 +40,12 @@ def _make_sampler() -> Sampler:
         device=DEVICE,
     )
     return Sampler(
+        vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()),
         max_num_reqs=4,
         vocab_size=VOCAB_SIZE,
         device=DEVICE,
         req_states=req_states,
-        reasoning_config=MockReasoningConfig(),
+        custom_logits_processors=custom_logits_processors,
     )
 
 
@@ -65,23 +73,46 @@ def test_logits_processing_cache_matches_request_features(
     sampling_params: SamplingParams, expected: bool
 ):
     sampler = _make_sampler()
-    sampler.add_request(3, prompt_len=1, sampling_params=sampling_params)
+    sampler.add_request(3, sampling_params=sampling_params)
 
     assert sampler.needs_logits_processing[3] == expected
 
 
 def test_logits_processing_cache_is_overwritten_when_slot_is_reused():
     sampler = _make_sampler()
-    sampler.add_request(3, 1, SamplingParams.for_sampler_warmup())
-    sampler.add_request(3, 1, SamplingParams())
+    sampler.add_request(3, SamplingParams.for_sampler_warmup())
+    sampler.add_request(3, SamplingParams())
 
     assert not sampler.needs_logits_processing[3]
 
 
+class _GateProcessor(LogitsProcessor):
+    """Reports a fixed admission decision from add_request()."""
+
+    def __init__(self, admitted: bool):
+        self.admitted = admitted
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        return self.admitted
+
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        raise AssertionError("the pipeline never runs in this test")
+
+
+@pytest.mark.parametrize("admitted", [True, False])
+def test_custom_processor_add_request_gates_pipeline_flag(admitted: bool):
+    """A custom processor's add_request() return value is OR-ed into the
+    per-request needs_logits_processing flag."""
+    sampler = _make_sampler(custom_logits_processors=[_GateProcessor(admitted)])
+    sampler.add_request(3, SamplingParams())
+
+    assert sampler.needs_logits_processing[3] == admitted
+
+
 def test_logits_processing_cache_only_checks_active_requests():
     sampler = _make_sampler()
-    sampler.add_request(0, 1, SamplingParams(temperature=0.0))
-    sampler.add_request(2, 1, SamplingParams.for_sampler_warmup())
+    sampler.add_request(0, SamplingParams(temperature=0.0))
+    sampler.add_request(2, SamplingParams.for_sampler_warmup())
 
     sampling_only = np.array([0], dtype=np.int32)
     with_processing = np.array([0, 2], dtype=np.int32)
diff --git a/tests/v1/worker/test_gpu_thinking_budget.py b/tests/v1/worker/test_gpu_thinking_budget.py
index 2bc697b08438..245602733535 100644
--- a/tests/v1/worker/test_gpu_thinking_budget.py
+++ b/tests/v1/worker/test_gpu_thinking_budget.py
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
+import numpy as np
 import pytest
 import torch
 
@@ -12,6 +15,7 @@
     )
 
 from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
 from vllm.v1.worker.gpu.states import RequestState
@@ -75,11 +79,17 @@ def _apply(
     idx_mapping_np = idx_mapping.cpu().numpy()
     state.apply(
         logits,
-        expanded_idx_mapping,
-        idx_mapping,
-        idx_mapping_np,
-        torch.tensor(input_ids, dtype=torch.int32, device=DEVICE),
-        torch.tensor(local_pos, dtype=torch.int32, device=DEVICE),
+        LogitsContext(
+            expanded_idx_mapping=expanded_idx_mapping,
+            idx_mapping=idx_mapping,
+            idx_mapping_np=idx_mapping_np,
+            expanded_local_pos=torch.tensor(
+                local_pos, dtype=torch.int32, device=DEVICE
+            ),
+            input_ids=torch.tensor(input_ids, dtype=torch.int32, device=DEVICE),
+            pos=torch.zeros(len(input_ids), dtype=torch.int32, device=DEVICE),
+            seq_lens_upper_bound_np=np.full(1, len(input_ids), dtype=np.int64),
+        ),
     )
     return logits.cpu()
 
@@ -190,15 +200,14 @@ def test_v2_greedy_sampling_applies_thinking_budget():
     """Greedy-only requests must not bypass thinking-budget processing."""
     req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
     sampler = Sampler(
+        vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()),
         max_num_reqs=4,
         vocab_size=VOCAB_SIZE,
         device=DEVICE,
         req_states=req_states,
-        reasoning_config=MockReasoningConfig(),
     )
     sampler.add_request(
         req_idx=3,
-        prompt_len=1,
         sampling_params=SamplingParams(
             temperature=0.0,
             thinking_token_budget=3,
@@ -219,6 +228,7 @@ def test_v2_greedy_sampling_applies_thinking_budget():
         torch.tensor([4], dtype=torch.int32, device=DEVICE),
         input_ids,
         torch.tensor([0], dtype=torch.int32, device=DEVICE),
+        np.full(1, 4, dtype=np.int64),
     )
 
     assert out[0, END].item() == pytest.approx(1.0e9)
diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py
index f54e74671071..ebf334e45d4f 100644
--- a/tests/v1/worker/test_gpu_ubatch_slicing.py
+++ b/tests/v1/worker/test_gpu_ubatch_slicing.py
@@ -10,6 +10,7 @@
 """
 
 import threading
+from contextlib import nullcontext
 from dataclasses import replace
 from types import SimpleNamespace
 from typing import Any, cast
@@ -20,13 +21,23 @@
 import torch
 
 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
-from vllm.config import CUDAGraphMode, ModelConfig, ParallelConfig, VllmConfig
+from vllm.config import (
+    CompilationConfig,
+    CUDAGraphMode,
+    ModelConfig,
+    ParallelConfig,
+    VllmConfig,
+)
 from vllm.forward_context import create_forward_context
 from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker import gpu_ubatch_wrapper as legacy_gpu_ubatch_wrapper
 from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
-from vllm.v1.worker.gpu import dp_utils
-from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+from vllm.v1.worker.gpu import cudagraph_utils, dp_utils
+from vllm.v1.worker.gpu.cudagraph_utils import (
+    BatchExecutionDescriptor,
+    CudaGraphManager,
+)
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.gpu.ubatch_utils import (
@@ -272,6 +283,7 @@ def _sync_dp(
     num_tokens_per_rank: list[int],
     uniform_token_count_per_rank: list[int] | None = None,
     allow_ubatching: bool = True,
+    cudagraph_manager: Any = None,
 ) -> tuple[BatchExecutionDescriptor, dp_utils.DPSyncState | None]:
     """Run the DP handshake with the all-reduce stubbed out.
 
@@ -295,7 +307,7 @@ def _sync_dp(
         patch.object(dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)),
     ):
         return dp_utils.sync_cudagraph_and_dp_padding(
-            cudagraph_manager=None,
+            cudagraph_manager=cudagraph_manager,
             desired_batch_desc=BatchExecutionDescriptor(
                 cg_mode=CUDAGraphMode.NONE,
                 num_tokens=num_tokens_per_rank[0],
@@ -416,6 +428,94 @@ def test_slice_model_inputs_handles_mrope_positions():
     assert sliced["intermediate_tensors"] is None
 
 
+@pytest.mark.parametrize(
+    ("cudagraph_runtime_mode", "execution_method"),
+    [
+        (CUDAGraphMode.FULL, "_capture_ubatches"),
+        (CUDAGraphMode.NONE, "_run_ubatches"),
+    ],
+)
+def test_legacy_ubatch_contexts_slice_padding_mask(
+    cudagraph_runtime_mode: CUDAGraphMode, execution_method: str
+):
+    """Replacement contexts preserve each microbatch's padding rows."""
+    ubatch_slices = [
+        UBatchSlice(slice(0, 2), slice(0, 3)),
+        UBatchSlice(slice(2, 4), slice(3, 6)),
+    ]
+    is_padding = torch.tensor([False, True, False, True, False, True])
+    forward_context = SimpleNamespace(
+        batch_descriptor=object(),
+        ubatch_slices=ubatch_slices,
+        cudagraph_runtime_mode=cudagraph_runtime_mode,
+        attn_metadata=[{}, {}],
+        slot_mapping={},
+        dp_metadata=object(),
+        is_padding=is_padding,
+    )
+
+    wrapper = object.__new__(legacy_gpu_ubatch_wrapper.UBatchWrapper)
+    wrapper.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(data_parallel_size=2)
+    )
+    wrapper.comm_stream = object()
+    wrapper.ready_barrier = object()
+    wrapper.cudagraphs = {}
+    wrapper.sm_control = nullcontext()
+    wrapper.runnable = object()
+
+    def make_contexts(*, forward_contexts, **kwargs):
+        return [
+            SimpleNamespace(id=i, forward_context=context)
+            for i, context in enumerate(forward_contexts)
+        ]
+
+    def create_context(attn_metadata, vllm_config, **kwargs):
+        return SimpleNamespace(attn_metadata=attn_metadata, **kwargs)
+
+    model_inputs = {
+        "input_ids": torch.arange(6),
+        "positions": torch.arange(6),
+        "inputs_embeds": None,
+        "intermediate_tensors": None,
+    }
+    expected = object()
+    with (
+        patch.object(
+            legacy_gpu_ubatch_wrapper,
+            "get_forward_context",
+            return_value=forward_context,
+        ),
+        patch.object(
+            legacy_gpu_ubatch_wrapper.torch.cuda,
+            "current_stream",
+            return_value=object(),
+        ),
+        patch.object(
+            legacy_gpu_ubatch_wrapper.DPMetadata, "make", return_value=object()
+        ),
+        patch.object(
+            legacy_gpu_ubatch_wrapper,
+            "create_forward_context",
+            side_effect=create_context,
+        ),
+        patch.object(
+            legacy_gpu_ubatch_wrapper,
+            "make_ubatch_contexts",
+            side_effect=make_contexts,
+        ),
+        patch.object(wrapper, execution_method, return_value=expected) as execute,
+    ):
+        result = wrapper(**model_inputs)
+
+    assert result is expected
+    contexts = [
+        metadata.context.forward_context for metadata in execute.call_args.args[0]
+    ]
+    torch.testing.assert_close(contexts[0].is_padding, is_padding[:3])
+    torch.testing.assert_close(contexts[1].is_padding, is_padding[3:])
+
+
 def _make_dbo_config() -> VllmConfig:
     return VllmConfig(
         model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0),
@@ -656,7 +756,7 @@ def test_microbatches_recompute_dcp_lens_from_truncated_seq_lens(dcp_rank: int):
         [1, 1, 10, 1, 1], [64, 96, 512, 32, 48]
     )
     # execute_model has populated the merged batch's DCP metadata already.
-    input_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
+    input_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens(
         buffers.dcp_local_seq_lens,
         input_batch.seq_lens,
         input_batch.num_reqs,
@@ -718,3 +818,148 @@ def test_slicing_drops_stale_dcp_metadata_when_dcp_is_off():
     ]
 
     assert all(ubatch.dcp_local_seq_lens is None for ubatch in ubatches)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU")
+def test_capturable_run_replays_as_a_cudagraph():
+    """Microbatched graph replay reads updated persistent input buffers."""
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0),
+        parallel_config=ParallelConfig(
+            enable_dbo=True, all2all_backend="deepep_low_latency"
+        ),
+    )
+    device = torch.device("cuda:0")
+    runner = _make_execution_runner(vllm_config)
+
+    num_tokens = 16
+    input_ids = torch.arange(num_tokens, device=device)
+    positions = torch.arange(num_tokens, device=device)
+    model_inputs = {
+        "input_ids": input_ids,
+        "positions": positions,
+        "inputs_embeds": None,
+        "intermediate_tensors": None,
+    }
+    ubatch_state = _make_ubatch_state(
+        vllm_config,
+        [
+            UBatchSlice(slice(0, 4), slice(0, 8)),
+            UBatchSlice(slice(4, 8), slice(8, 16)),
+        ],
+    )
+
+    model = _YieldingModel([])
+    graph = torch.cuda.CUDAGraph()
+    finish = runner.begin_capturable_run(
+        model, model_inputs, ubatch_state, for_capture=True
+    )
+    with torch.cuda.graph(graph, stream=runner.capture_stream):
+        captured = finish()
+
+    input_ids.fill_(3)
+    positions.fill_(5)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    expected = torch.full(
+        (num_tokens, 1), 3 * 2 + 5, dtype=torch.float32, device=device
+    )
+    torch.testing.assert_close(captured, expected)
+
+
+def _request_slices(input_batch: InputBatch) -> list[slice]:
+    return [s.request_slice for s in create_ubatch_slices(input_batch, 2)]
+
+
+def test_captured_split_survives_a_replay_with_enough_requests():
+    """Replay preserves the captured split only if real requests reach it."""
+    buffers = _make_buffers()
+    captured = _request_slices(InputBatch.make_dummy(16, 16, buffers))
+
+    def replay(num_reqs: int) -> list[slice]:
+        return _request_slices(
+            _make_input_batch(
+                [1] * num_reqs,
+                [128] * num_reqs,
+                buffers,
+                num_reqs_padded=16,
+                num_tokens_padded=16,
+            )
+        )
+
+    assert replay(16) == captured
+    assert replay(9) == captured
+    assert replay(4) != captured
+
+
+def _make_cudagraph_manager(capture_sizes: list[int]) -> CudaGraphManager:
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0),
+        parallel_config=ParallelConfig(
+            enable_dbo=True,
+            all2all_backend="deepep_low_latency",
+            dbo_decode_token_threshold=DECODE_THRESHOLD,
+            dbo_prefill_token_threshold=PREFILL_THRESHOLD,
+        ),
+        compilation_config=CompilationConfig(
+            cudagraph_mode=CUDAGraphMode.FULL,
+            cudagraph_capture_sizes=capture_sizes,
+        ),
+    )
+    with patch.object(
+        cudagraph_utils,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
+    ):
+        manager = CudaGraphManager(
+            vllm_config,
+            torch.device("cuda:0"),
+            CUDAGraphMode.FULL,
+            decode_query_len=1,
+            ubatch_runner=cast(UBatchRunner, object()),
+        )
+    manager._graphs_captured = True
+    return manager
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool")
+def test_microbatched_graphs_are_only_offered_to_uniform_batches():
+    """Mixed batches cannot reuse graphs captured with uniform query lengths."""
+    manager = _make_cudagraph_manager([64, 128])
+
+    desc = manager.dispatch(
+        num_reqs=64,
+        num_tokens=64,
+        uniform_token_count=1,
+        num_active_loras=0,
+        num_ubatches=2,
+    )
+    assert desc.cg_mode == CUDAGraphMode.FULL
+    assert desc.num_ubatches == 2
+
+    desc = manager.dispatch(
+        num_reqs=8,
+        num_tokens=64,
+        uniform_token_count=None,
+        num_active_loras=0,
+        num_ubatches=2,
+    )
+    assert desc.cg_mode == CUDAGraphMode.NONE
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool")
+def test_microbatched_graph_needs_every_rank_to_reach_the_split():
+    """Fall back to eager on all ranks if any rank cannot reach the captured split."""
+    manager = _make_cudagraph_manager([64, 128])
+    uniform = [1, 1]
+
+    desc, _ = _sync_dp([64, 64], uniform, cudagraph_manager=manager)
+    assert desc.cg_mode == CUDAGraphMode.FULL
+    assert desc.num_ubatches == 2
+
+    # Rank 1 pads the group up to a 128-token graph, but rank 0 only reaches 33.
+    desc, dp_sync = _sync_dp([33, 128], uniform, cudagraph_manager=manager)
+    assert desc.cg_mode == CUDAGraphMode.NONE
+    assert desc.num_ubatches == 2
+    assert dp_sync is not None and dp_sync.eager
diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py
index fcace4ae4a20..067f31e9350c 100644
--- a/tests/v1/worker/test_gpu_warmup_blocks.py
+++ b/tests/v1/worker/test_gpu_warmup_blocks.py
@@ -89,6 +89,7 @@ def _make_runner(
     return SimpleNamespace(
         num_speculative_steps=num_spec_steps,
         adaptive_verification=None,
+        rejection_sampler=None,
         decode_query_len=num_spec_steps + 1,
         is_pooling_model=False,
         is_encoder_decoder=False,
diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py
index 88cbee680240..375e43810fc6 100644
--- a/tests/v1/worker/test_gpu_worker.py
+++ b/tests/v1/worker/test_gpu_worker.py
@@ -6,6 +6,7 @@
 from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.v1.worker import gpu_worker, startup_plan
@@ -15,6 +16,42 @@
     maybe_save_startup_plan,
 )
 
+
+def test_load_model_preserves_compiled_graphs_at_runtime(monkeypatch):
+    """Profiling must use serving's thread count to keep Dynamo guards valid."""
+    from torch._dynamo.testing import CompileCounter
+
+    monkeypatch.delenv("OMP_NUM_THREADS", raising=False)
+    monkeypatch.setattr(gpu_worker, "has_ec_transfer", lambda: False)
+    monkeypatch.setattr(
+        gpu_worker, "set_current_vllm_config", lambda config: nullcontext()
+    )
+    loading_threads = []
+    worker = SimpleNamespace(
+        vllm_config=SimpleNamespace(weight_transfer_config=None),
+        model_runner=SimpleNamespace(
+            load_model=lambda **kwargs: loading_threads.append(torch.get_num_threads())
+        ),
+        _maybe_get_memory_pool_context=lambda **kwargs: nullcontext(),
+        _scoped_allocator_max_split=lambda **kwargs: nullcontext(),
+    )
+    original_threads = torch.get_num_threads()
+    try:
+        torch.set_num_threads(2)
+        gpu_worker.Worker.load_model(worker)
+        assert loading_threads == [2]
+
+        counter = CompileCounter()
+        compiled = torch.compile(lambda x: x + 1, backend=counter, fullgraph=True)
+        x = torch.ones(2)
+        compiled(x)
+        gpu_worker.set_torch_threads_for_runtime()
+        torch.testing.assert_close(compiled(x), x + 1)
+        assert counter.frame_count == 1
+    finally:
+        torch.set_num_threads(original_threads)
+
+
 # Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and
 # saved by Worker.determine_available_memory / compile_or_warm_up_model.
 
diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py
index e576ebd03177..537653bda38d 100644
--- a/tests/v1/worker/test_gpu_worker_weight_transfer.py
+++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py
@@ -11,7 +11,7 @@
 import torch
 import torch.nn as nn
 
-from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config import ParallelConfig, VllmConfig, get_current_vllm_config
 from vllm.lora.layers import BaseLayerWithLoRA
 from vllm.v1.worker.gpu_model_runner import _get_parameter_for_reload
 from vllm.v1.worker.gpu_worker import Worker
@@ -136,12 +136,22 @@ def test_rank_local_update_selects_worker_payload(rank, expected):
     assert worker._weight_update_active is True
 
 
-def test_rank_local_update_includes_data_parallel_rank():
+def test_rank_local_update_uses_data_parallel_index_after_reconfigure():
     engine = _RecordingEngine()
     worker = _make_worker(engine)
     worker.rank = 0
-    worker.vllm_config.parallel_config.data_parallel_size = 4
-    worker.vllm_config.parallel_config.data_parallel_rank = 2
+    parallel_config = ParallelConfig(
+        data_parallel_size=4,
+        data_parallel_rank=2,
+    )
+    assert parallel_config.data_parallel_rank == 2
+    assert parallel_config.data_parallel_index == 2
+
+    parallel_config.reconfigure_for_independent_dp_rank()
+    assert parallel_config.data_parallel_rank == 0
+    assert parallel_config.data_parallel_index == 2
+
+    worker.vllm_config.parallel_config = parallel_config
     Worker.start_weight_update(worker)
 
     Worker.update_weights(
diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py
index a9579f1af69b..7f4bea52c227 100644
--- a/tests/v1/worker/test_kv_cache_allocation_scope.py
+++ b/tests/v1/worker/test_kv_cache_allocation_scope.py
@@ -5,8 +5,10 @@
 from types import SimpleNamespace
 from typing import Any, cast
 
+import pytest
 import torch
 
+import vllm.v1.hisparse.binding as hisparse_binding
 import vllm.v1.worker.gpu.attn_utils as attn_utils
 import vllm.v1.worker.gpu_model_runner as gpu_model_runner
 from vllm.v1.worker.gpu_worker import Worker
@@ -26,7 +28,8 @@ def __exit__(self, *args: Any) -> None:
         self.active = False
 
 
-def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
+@pytest.mark.parametrize("hisparse", [False, True])
+def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> None:
     scope = _AllocationScope()
     kv_caches = {"layer": torch.empty(0)}
 
@@ -38,23 +41,52 @@ def bind(*args, **kwargs):
         assert not scope.active
 
     monkeypatch.setattr(attn_utils, "allocate_kv_cache", allocate)
-    monkeypatch.setattr(attn_utils, "bind_kv_cache", bind)
+    monkeypatch.setattr(attn_utils, "bind_kv_cache_to_layers", bind)
     monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda config: {})
 
+    hisparse_bindings = []
+    if hisparse:
+        host_pool = SimpleNamespace()
+
+        def bind_hisparse(**kwargs):
+            hisparse_bindings.append(kwargs)
+            assert scope.active
+            assert kwargs["kv_caches"] is kv_caches
+            assert kwargs["host_pool"] is host_pool
+            return [
+                SimpleNamespace(
+                    view=SimpleNamespace(cache=torch.empty(1, 1, 1), block_size=1),
+                    runtime=SimpleNamespace(),
+                )
+            ]
+
+        monkeypatch.setattr(
+            hisparse_binding,
+            "HiSparseHostPool",
+            lambda vllm_config, kv_cache_config: host_pool,
+        )
+        monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate)
+        monkeypatch.setattr(hisparse_binding, "bind_hisparse_kv_caches", bind_hisparse)
+
     config = SimpleNamespace(
+        attention_config=SimpleNamespace(
+            hisparse_config=object() if hisparse else None
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1),
         cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None),
         model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")),
     )
     result = attn_utils.init_kv_cache(
-        [],
         {},
         SimpleNamespace(kv_cache_groups=[]),
         torch.device("cpu"),
         [],
         config,
         kv_cache_allocation_context=scope,
+        block_tables=SimpleNamespace(),
     )
 
+    assert len(hisparse_bindings) == int(hisparse)
     assert result is kv_caches
     assert not scope.active
 
diff --git a/tests/v1/worker/test_logits_processor_loader.py b/tests/v1/worker/test_logits_processor_loader.py
new file mode 100644
index 000000000000..7e851e5237d1
--- /dev/null
+++ b/tests/v1/worker/test_logits_processor_loader.py
@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU-only tests for the V2 custom logits processor loader.
+
+Load-time validation is pure Python (importlib + issubclass), so it runs
+without CUDA. Pipeline behavior is covered by test_gpu_logits_processors.py.
+"""
+
+import subprocess
+import sys
+from enum import Enum, auto
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+import vllm.v1.worker.gpu.sample.logits_processor.loader as loader
+from vllm.exceptions import VLLMValidationError
+from vllm.sampling_params import SamplingParams
+from vllm.v1.sample.logits_processor import (
+    LogitsProcessor as V1LogitsProcessor,
+)
+from vllm.v1.worker.gpu.sample.logits_processor import (
+    LogitsContext,
+    LogitsProcessor,
+)
+
+
+class DummyV2Processor(LogitsProcessor):
+    """Records its constructor args so tests can assert how it was built."""
+
+    def __init__(self, vllm_config: Any, req_state: Any):
+        self.ctor_args = (vllm_config, req_state)
+
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        return logits
+
+
+class AltV2Processor(DummyV2Processor):
+    """Second V2 processor, used to assert plugin/FQCN load ordering."""
+
+
+class DummyV1Processor(V1LogitsProcessor):
+    """A V1-interface processor. Never instantiated: the loader must reject
+    it at the issubclass check, before any constructor runs."""
+
+
+def not_a_processor() -> None:
+    """FQCN target that resolves but is not a class."""
+
+
+def _fake_entrypoint(cls: type) -> SimpleNamespace:
+    return SimpleNamespace(
+        name=cls.__name__.lower(),
+        value=f"fake:{cls.__name__}",
+        load=lambda: cls,
+    )
+
+
+@pytest.fixture(autouse=True)
+def no_installed_plugins(monkeypatch: pytest.MonkeyPatch):
+    """Isolate tests from entry points registered in the environment."""
+    monkeypatch.setattr(loader, "entry_points", lambda group: [])
+
+
+class ProcSource(Enum):
+    CLASS = auto()
+    FQCN = auto()
+    ENTRYPOINT = auto()
+
+
+@pytest.mark.parametrize("source", list(ProcSource))
+def test_loads_v2_processors(monkeypatch: pytest.MonkeyPatch, source: ProcSource):
+    """A V2 subclass loads via class object, FQCN string, or entrypoint."""
+    fake_req_states = SimpleNamespace(
+        device="cpu",
+        max_num_reqs=4,
+        vocab_size=128,
+        all_token_ids=None,
+        prompt_len=None,
+        prefill_len=None,
+        total_len=None,
+    )
+    custom: list[Any]
+    if source is ProcSource.CLASS:
+        custom = [DummyV2Processor]
+    elif source is ProcSource.FQCN:
+        custom = [f"{__name__}:DummyV2Processor"]
+    else:
+        monkeypatch.setattr(
+            loader, "entry_points", lambda group: [_fake_entrypoint(DummyV2Processor)]
+        )
+        # Plugins load before user-specified processors.
+        custom = [f"{__name__}:AltV2Processor"]
+
+    procs = loader.build_custom_logits_processors(None, fake_req_states, False, custom)
+
+    if source is ProcSource.ENTRYPOINT:
+        assert [type(p) for p in procs] == [DummyV2Processor, AltV2Processor]
+    else:
+        assert [type(p) for p in procs] == [DummyV2Processor]
+    vllm_config, req_state = procs[0].ctor_args
+    assert vllm_config is None
+    assert isinstance(req_state, loader.LogitsProcRequestState)
+    assert req_state.max_num_reqs == 4
+
+
+@pytest.mark.parametrize(
+    ("custom", "plugin_cls", "exc_type", "msg_fragment"),
+    [
+        pytest.param(
+            [DummyV1Processor],
+            None,
+            ValueError,
+            "not a subclass",
+            id="v1-class-object",
+        ),
+        pytest.param(
+            [f"{__name__}:DummyV1Processor"],
+            None,
+            ValueError,
+            "not a subclass",
+            id="v1-fqcn",
+        ),
+        pytest.param(
+            [],
+            DummyV1Processor,
+            ValueError,
+            "V1-interface plugins are not supported",
+            id="v1-entrypoint",
+        ),
+        pytest.param(
+            [f"{__name__}:not_a_processor"],
+            None,
+            ValueError,
+            "must be a type",
+            id="non-type-fqcn",
+        ),
+        pytest.param(
+            ["no.such.module:Nope"],
+            None,
+            RuntimeError,
+            "no.such.module:Nope",
+            id="missing-module",
+        ),
+        pytest.param(
+            ["no_colon_here"],
+            None,
+            ValueError,
+            "Expected format",
+            id="fqcn-missing-colon",
+        ),
+        pytest.param(
+            ["too:many:colons"],
+            None,
+            ValueError,
+            "Expected format",
+            id="fqcn-extra-colons",
+        ),
+    ],
+)
+def test_rejects_invalid(
+    monkeypatch: pytest.MonkeyPatch,
+    custom: list,
+    plugin_cls: type | None,
+    exc_type: type[Exception],
+    msg_fragment: str,
+):
+    """V1 processors and malformed references fail at load time, not at
+    sampling time, with a message naming the offending reference."""
+    if plugin_cls is not None:
+        monkeypatch.setattr(
+            loader, "entry_points", lambda group: [_fake_entrypoint(plugin_cls)]
+        )
+
+    with pytest.raises(exc_type, match=msg_fragment):
+        loader.build_custom_logits_processors(None, None, False, custom)
+
+
+def test_pooling_model_rejects_custom_logitsprocs():
+    """Pooling models reject custom processors instead of ignoring them."""
+    with pytest.raises(ValueError, match="Pooling models do not support"):
+        loader.build_custom_logits_processors(None, None, True, [DummyV2Processor])
+    assert loader.build_custom_logits_processors(None, None, True, []) == []
+
+
+class ValidatingProcessor(DummyV2Processor):
+    """Rejects target_token=-1 via the optional validate_params hook."""
+
+    @classmethod
+    def validate_params(cls, sampling_params: SamplingParams) -> None:
+        if (sampling_params.extra_args or {}).get("target_token") == -1:
+            raise ValueError("target_token must not be -1")
+
+
+def test_validate_params_runs_at_admission():
+    """The factory loads classes up front; the returned validator runs
+    validate_params per request and wraps rejections as VLLMValidationError."""
+    loader._cached_load_v2_logitsprocs.cache_clear()
+    validate = loader.build_custom_logits_processors_params_validator(
+        [ValidatingProcessor]
+    )
+    validate(SamplingParams(extra_args={"target_token": 1}))
+    with pytest.raises(VLLMValidationError, match="target_token"):
+        validate(SamplingParams(extra_args={"target_token": -1}))
+
+
+def test_validator_factory_loads_classes_eagerly():
+    """A bad reference fails when the validator is built (at startup), not on
+    the first request."""
+    loader._cached_load_v2_logitsprocs.cache_clear()
+    with pytest.raises(RuntimeError, match="no.such.module:Nope"):
+        loader.build_custom_logits_processors_params_validator(["no.such.module:Nope"])
+
+
+def test_loader_import_stays_frontend_safe():
+    """The frontend process imports the loader to validate params; importing
+    it must not pull in model-runner side modules."""
+    code = (
+        "import sys\n"
+        "import vllm.v1.worker.gpu.sample.logits_processor.loader\n"
+        "assert 'vllm.sampling_params' not in sys.modules\n"
+        "assert 'vllm.v1.sample.logits_processor' not in sys.modules\n"
+        "assert 'vllm.v1.worker.gpu.states' not in sys.modules\n"
+        "assert 'vllm.v1.worker.gpu.buffer_utils' not in sys.modules\n"
+    )
+    subprocess.run([sys.executable, "-c", code], check=True)
diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py
index 0cdd91e0d7b3..c65679059ac0 100644
--- a/tests/v1/worker/test_mamba_hybrid_model_state.py
+++ b/tests/v1/worker/test_mamba_hybrid_model_state.py
@@ -21,7 +21,7 @@
 def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> None:
     state = object.__new__(MambaHybridModelState)
     state.vllm_config = SimpleNamespace(num_speculative_tokens=0)
-    state.max_model_len = 8192
+    state.model_config = SimpleNamespace(max_model_len=8192)
     state._align_mode = False
     state.recoverssm = None
 
@@ -34,6 +34,7 @@ def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> Non
         query_start_loc_np=torch.tensor([0, 1], dtype=torch.int32).numpy(),
         query_start_loc=torch.tensor([0, 1], dtype=torch.int32),
         num_scheduled_tokens=torch.tensor([1], dtype=torch.int32),
+        max_query_len=None,
         seq_lens_cpu_upper_bound=torch.tensor([1537], dtype=torch.int32),
         seq_lens=torch.tensor([1537], dtype=torch.int32),
         is_prefilling_np=torch.tensor([False]).numpy(),
diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py
index a162c39a8695..47a27795ca18 100644
--- a/tests/v1/worker/test_mamba_utils.py
+++ b/tests/v1/worker/test_mamba_utils.py
@@ -867,8 +867,7 @@ def test_batch_memcpy_left_overlap_has_memmove_semantics(self, device):
             torch.testing.assert_close(state, expected, rtol=0, atol=0)
 
     def test_matches_python_postprocess_mamba(self, device, test_config):
-        """
-        Golden test: GPU kernel produces identical results to Python impl.
+        """Golden test: GPU kernel produces identical results to Python impl.
 
         This test:
         1. Sets up identical initial state for both paths
@@ -1146,8 +1145,7 @@ def test_various_batch_sizes(self, device, test_config, num_reqs):
         )
 
     def test_block_table_with_realistic_stride(self, device, test_config):
-        """
-        Test kernel with realistic block table strides.
+        """Test kernel with realistic block table strides.
 
         In real usage, the block table is pre-allocated with shape
         [max_num_reqs, max_num_blocks_per_req] and then sliced to
@@ -1275,8 +1273,7 @@ def test_block_table_with_realistic_stride(self, device, test_config):
     def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1(
         self, device, test_config
     ):
-        """
-        Test the ``src_addr == dst_addr`` early-return path in
+        """Test the ``src_addr == dst_addr`` early-return path in
         postprocess_mamba_fused_kernel matches Python behavior.
 
         When src_addr == dst_addr (source and destination memory addresses are
@@ -1419,8 +1416,7 @@ def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1(
     def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1(
         self, device, test_config
     ):
-        """
-        Test the ``src_block_idx == dest_block_idx`` post-copy update in
+        """Test the ``src_block_idx == dest_block_idx`` post-copy update in
         postprocess_mamba_fused_kernel matches Python behavior.
 
         When src_block_idx == dest_block_idx but accept_token_bias > 0, both
@@ -1583,8 +1579,7 @@ def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1(
     def test_different_block_idx_copies_without_setting_accepted_to_1(
         self, device, test_config
     ):
-        """
-        Test that neither special-case path triggers when
+        """Test that neither special-case path triggers when
         src_block_idx != dest_block_idx, and GPU matches Python behavior.
 
         When copying between different blocks:
@@ -1716,8 +1711,7 @@ def test_different_block_idx_copies_without_setting_accepted_to_1(
     def test_prefix_caching_shared_block_does_not_set_accepted_to_1(
         self, device, test_config
     ):
-        """
-        Regression test: with prefix caching, different logical block indices
+        """Regression test: with prefix caching, different logical block indices
         can map to the same physical block. The kernel must NOT set
         num_accepted_tokens to 1 in that case.
 
@@ -1849,8 +1843,7 @@ def test_prefix_caching_shared_block_does_not_set_accepted_to_1(
         )
 
     def test_prefix_caching_nonsequential_block_ids_boundary(self, device, test_config):
-        """
-        Regression test: non-sequential physical block IDs under prefix caching
+        """Regression test: non-sequential physical block IDs under prefix caching
         with the needs_copy boundary at exact equality.
 
         Under PC, the block allocator assigns physical block IDs in arbitrary
@@ -2000,8 +1993,7 @@ def test_prefix_caching_nonsequential_block_ids_boundary(self, device, test_conf
         assert input_batch_py.num_accepted_tokens_cpu[1] == num_accepted_tokens[1]
 
     def test_prefix_caching_mixed_shared_and_distinct_blocks(self, device, test_config):
-        """
-        Regression test: mixed batch under prefix caching where some requests
+        """Regression test: mixed batch under prefix caching where some requests
         have shared physical blocks (aliased) and others have distinct blocks,
         with the needs_copy boundary at various positions.
 
@@ -2166,8 +2158,7 @@ def test_prefix_caching_mixed_shared_and_distinct_blocks(self, device, test_conf
     def test_pc_aliased_blocks_skip_must_use_logical_idx_not_addr(
         self, device, test_config
     ):
-        """
-        Regression test for 6466ce0d vs 959ca0fd: the kernel's early-return
+        """Regression test for 6466ce0d vs 959ca0fd: the kernel's early-return
         guard must compare logical block indices, not physical addresses.
 
         Under prefix caching, different logical blocks (src_block_idx=0,
@@ -2299,8 +2290,7 @@ def test_pc_aliased_blocks_skip_must_use_logical_idx_not_addr(
         )
 
     def test_as_strided_temporal_copy_size(self, device, test_config):
-        """
-        Regression test for 240723d46: temporal copy_size must be
+        """Regression test for 240723d46: temporal copy_size must be
         inner_size * elem_size, not state_block_stride.
 
         In production (gpu_model_runner.py), conv and temporal states share
@@ -2458,8 +2448,7 @@ def make_views(raw):
         )
 
     def test_temporal_copy_with_bias_ge_2(self, device, test_config):
-        """
-        Coverage test for the temporal-state block-table stride arithmetic
+        """Coverage test for the temporal-state block-table stride arithmetic
         when ``accept_token_bias >= 2``.
 
         The kernel computes, for temporal (non-conv) states::
diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py
index df339135eb5a..04f456a7c9cd 100644
--- a/tests/v1/worker/test_mixed_warmup_gate.py
+++ b/tests/v1/worker/test_mixed_warmup_gate.py
@@ -35,10 +35,15 @@ def test_mixed_warmup_skipped_for_single_seq(max_num_reqs):
 def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup):
     """Startup must warm fixed drafts before calibration and retain its manager."""
     manager = SimpleNamespace(cost_tables=None)
-    runner = SimpleNamespace(adaptive_verification=manager)
+    rejection_sampler = SimpleNamespace(enable_adaptive_verification=True)
+    runner = SimpleNamespace(
+        adaptive_verification=manager,
+        rejection_sampler=rejection_sampler,
+    )
 
     def run_steps(model_runner, execute, sample):
         assert model_runner.adaptive_verification is None
+        assert not model_runner.rejection_sampler.enable_adaptive_verification
         if fail_warmup:
             raise RuntimeError("warmup failed")
 
@@ -50,3 +55,4 @@ def run_steps(model_runner, execute, sample):
         warmup.warmup_kernels(runner, _fail, _fail)
     assert runner.adaptive_verification is manager
     assert manager.cost_tables is None
+    assert rejection_sampler.enable_adaptive_verification
diff --git a/tests/v1/worker/test_sleep_mode_backend.py b/tests/v1/worker/test_sleep_mode_backend.py
index e2514f84e162..686331f2d700 100644
--- a/tests/v1/worker/test_sleep_mode_backend.py
+++ b/tests/v1/worker/test_sleep_mode_backend.py
@@ -56,6 +56,7 @@ def resume(self, tags: list[str] | None = None) -> None:
 
     worker = object.__new__(Worker)
     worker._sleep_mode_backend = Backend()
+    worker._sleep_saved_parameters = {}
     worker._sleep_saved_buffers = {}
     worker._sleep_saved_draft_buffers = {}
     worker.vllm_config = SimpleNamespace(
diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py
index a4474c8a5843..9a345e926c72 100644
--- a/tests/v1/worker/test_utils.py
+++ b/tests/v1/worker/test_utils.py
@@ -1,13 +1,1252 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import mmap
+from collections import deque
+from contextlib import nullcontext
 from types import SimpleNamespace
+from unittest.mock import MagicMock, call
 
+import numpy as np
+import pytest
 import torch
 
+import vllm.v1.attention.backends.mla.index_group as index_group_module
+import vllm.v1.hisparse.runtime as hisparse_runtime_module
+from vllm.config import CUDAGraphMode
 from vllm.config.mamba import MambaBackendEnum, MambaConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import (
+    worker as hisparse_worker_module,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
+    HiSparseConnectorWorker,
+    _flatten_row_mirrors,
+    _select_written_row_mirrors,
+    _SlotMappingStaging,
+)
 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
-from vllm.v1.worker.utils import bind_kv_cache
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
+from vllm.v1.worker.utils import (
+    bind_kv_cache,
+    bind_kv_cache_to_layers,
+    copy_kv_cache_blocks_inplace,
+)
+
+
+def _make_hisparse_worker() -> HiSparseConnectorWorker:
+    worker = object.__new__(HiSparseConnectorWorker)
+    worker.refines_row_mirrors = False
+    worker._slot_mapping_staging = None
+    worker._row_mirror_num_rows = 0
+    worker._per_layer_mirrored = set()
+    worker._submitted_mirror_layers = set()
+    worker._pending_dma_descriptors = deque()
+    worker._dma_free_descriptors = []
+    worker.host_write_events = (MagicMock(), MagicMock())
+    worker.host_write_event = worker.host_write_events[1]
+    worker._next_host_write_event = 0
+    worker.dma_stream = None
+    worker.shared_host_region = None
+    worker._completed_host_copy_dst_ids = []
+    worker._metrics_calls = 0
+    worker._metrics_event = MagicMock()
+    worker._metrics_pending = False
+    worker.leader_runtimes = []
+    return worker
+
+
+def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeypatch):
+    worker = _make_hisparse_worker()
+    worker._metrics_calls = hisparse_worker_module._METRICS_INTERVAL - 1
+    worker._metrics_pending = False
+    worker._metrics_event = MagicMock()
+    worker._metrics_event.query.return_value = True
+    compute_stream = MagicMock()
+    group = SimpleNamespace(
+        swap_stats=torch.tensor([12, 4], dtype=torch.uint64),
+        swap_stats_host=torch.empty(2, dtype=torch.uint64),
+        stats_row_bytes=16,
+        copy_stream=MagicMock(),
+    )
+    worker.leader_runtimes = [SimpleNamespace(index_group=group)]
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    monkeypatch.setattr(
+        hisparse_worker_module, "current_stream", lambda: compute_stream
+    )
+
+    assert worker.get_kv_connector_stats() is None
+    compute_stream.wait_stream.assert_called_once_with(group.copy_stream)
+    group.copy_stream.wait_stream.assert_called_once_with(compute_stream)
+    worker._metrics_event.record.assert_called_once_with()
+    assert group.swap_stats.tolist() == [0, 0]
+
+    stats = worker.get_kv_connector_stats()
+    assert stats is not None
+    assert stats.data == {
+        "cache_hits": [12],
+        "cache_misses": [4],
+        "host_to_device_bytes": [64],
+    }
+
+
+def test_hisparse_row_mirrors_follow_runner_request_order():
+    first = SparseKVRowMirror((1,), 10, 1)
+    second = SparseKVRowMirror((2,), 20, 1)
+    row_mirrors = {"first": (first,), "second": (second,)}
+
+    assert _flatten_row_mirrors(row_mirrors, ["second", "first"]) == (
+        second,
+        first,
+    )
+
+
+def test_hisparse_gpu_slots_select_written_rows_from_async_envelope():
+    """Rejected speculative rows must never become durable host KV."""
+    candidates = (
+        SparseKVRowMirror((100, 500), 1000, 4),
+        SparseKVRowMirror((20, 300), 2000, 4),
+    )
+    source_slots = np.array([102, 103, -1, 21, 22, 23], dtype=np.int64)
+
+    assert _select_written_row_mirrors(candidates, source_slots, 0) == (
+        SparseKVRowMirror((102, 502), 1002, 2),
+        SparseKVRowMirror((21, 301), 2001, 3),
+    )
+
+
+def test_hisparse_written_rows_preserve_scheduler_page_boundaries():
+    candidates = (
+        SparseKVRowMirror((0,), 100, 4),
+        SparseKVRowMirror((4,), 104, 4),
+    )
+
+    assert _select_written_row_mirrors(
+        candidates, np.array([2, 3, 4, 5], dtype=np.int64), 0
+    ) == (
+        SparseKVRowMirror((2,), 102, 2),
+        SparseKVRowMirror((4,), 104, 2),
+    )
+
+
+def test_hisparse_written_rows_without_host_destinations_are_skipped():
+    assert (
+        _select_written_row_mirrors((), np.array([100, 101], dtype=np.int64), 0) == ()
+    )
+
+
+def test_hisparse_appends_reference_slots_within_a_mirror_phase(monkeypatch):
+    """Separate context and query writes must share one mirror phase."""
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.refines_row_mirrors = True
+    state = _SlotMappingStaging(
+        MagicMock(), MagicMock(), torch.empty(4, dtype=torch.int64)
+    )
+    worker._slot_mapping_staging = state
+    handle = SimpleNamespace(runtime=SimpleNamespace(resident_source_index=1))
+    worker.cache_layer_names = ["layer"]
+    worker.cache_handles = [handle]
+    worker._layer_mirror_callbacks = (MagicMock(),)
+    worker._row_mirror_num_rows = 1
+    worker._submitted_mirror_layers = set()
+    main_stream = MagicMock()
+    monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: main_stream)
+    monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
+
+    handle.slot_mapping = torch.tensor([7, 8], dtype=torch.int64)
+    worker._stage_row_mirror_mapping(2)
+    handle.slot_mapping = torch.tensor([9, 10], dtype=torch.int64)
+    worker._stage_row_mirror_mapping(2)
+
+    torch.testing.assert_close(state.slots, torch.tensor([7, 8, 9, 10]))
+    assert state.stream.wait_stream.call_args_list == [call(main_stream)] * 2
+    assert state.event.record.call_args_list == [call(state.stream)] * 2
+    assert state.num_tokens == 4
+    assert state.source_index == 1
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected_calls"),
+    [(CUDAGraphMode.PIECEWISE, 1), (CUDAGraphMode.FULL, 0)],
+)
+def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary(
+    monkeypatch, mode, expected_calls
+):
+    """Piecewise replay must pipeline DMA; FULL replay uses the batch fallback."""
+    handle = object.__new__(hisparse_runtime_module.HiSparseCacheHandle)
+    handle.dummy_batch = False
+    handle.decode_batch = False
+    handle.submit_layer_mirror = MagicMock()
+    monkeypatch.setattr(
+        hisparse_runtime_module,
+        "get_forward_context",
+        lambda: SimpleNamespace(cudagraph_runtime_mode=mode),
+    )
+
+    handle.finish_kv_update()
+
+    assert handle.submit_layer_mirror.call_count == expected_calls
+
+
+def _hisparse_parallel_config(**overrides):
+    values = {
+        "tensor_parallel_size": 2,
+        "pipeline_parallel_size": 1,
+        "prefill_context_parallel_size": 1,
+        "decode_context_parallel_size": 1,
+        "world_size": 2,
+        "distributed_executor_backend": "mp",
+        "nnodes_within_dp": 1,
+        "data_parallel_index": 3,
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_hisparse_shares_host_pool_only_for_local_tp(monkeypatch):
+    monkeypatch.setattr(
+        hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True
+    )
+    config = SimpleNamespace(parallel_config=_hisparse_parallel_config())
+
+    assert hisparse_runtime_module.use_shared_hisparse_host_pool(config)
+
+    unsupported = (
+        {"tensor_parallel_size": 1, "world_size": 1},
+        {"pipeline_parallel_size": 2, "world_size": 4},
+        {"prefill_context_parallel_size": 2, "world_size": 4},
+        {"decode_context_parallel_size": 2},
+        {"world_size": 4},
+        {"distributed_executor_backend": "ray"},
+        {"nnodes_within_dp": 2},
+    )
+    for overrides in unsupported:
+        config.parallel_config = _hisparse_parallel_config(**overrides)
+        assert not hisparse_runtime_module.use_shared_hisparse_host_pool(config)
+
+
+@pytest.mark.skip_global_cleanup
+def test_hisparse_shared_host_pool_uses_one_replicated_mmap(monkeypatch):
+    page = mmap.PAGESIZE
+    tp_group = MagicMock()
+    monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", lambda: tp_group)
+
+    class FakeSharedOffloadRegion:
+        BLOCK_SIZE_ALIGNMENT = page
+
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+            self.total_size_bytes = kwargs["num_chunks"] * kwargs["kv_bytes_per_chunk"]
+            self.base_tensor = torch.empty(self.total_size_bytes, dtype=torch.int8)
+            self.is_pinned = False
+            self.pinned_addresses = []
+            self.view_sizes = []
+            self.offset = 0
+
+        def create_next_canonical_view(self, size):
+            self.view_sizes.append(size)
+            view = torch.as_strided(
+                self.base_tensor,
+                size=(self.kwargs["num_chunks"], size),
+                stride=(self.kwargs["kv_bytes_per_chunk"], 1),
+                storage_offset=self.offset,
+            )
+            self.offset += size
+            return view
+
+        def cleanup(self):
+            pass
+
+    monkeypatch.setattr(
+        hisparse_runtime_module, "SharedOffloadRegion", FakeSharedOffloadRegion
+    )
+    registration_ranges = MagicMock(return_value=((0, page), (page, 4 * page)))
+    monkeypatch.setattr(
+        hisparse_runtime_module,
+        "_hisparse_registration_ranges",
+        registration_ranges,
+    )
+    pinned: list[torch.Tensor] = []
+
+    def pin_tensor(tensor):
+        pinned.append(tensor)
+
+    monkeypatch.setattr(hisparse_runtime_module, "pin_tensor", pin_tensor)
+    config = SimpleNamespace(
+        instance_id="instance",
+        parallel_config=_hisparse_parallel_config(
+            tensor_parallel_size=1,
+            world_size=1,
+            distributed_executor_backend="uni",
+        ),
+    )
+
+    pools, private_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools(
+        config,
+        [24, 40],
+        num_blocks=4,
+        host_block_stride=page,
+        use_shared_host_pool=True,
+    )
+
+    assert region is not None
+    assert private_pools == []
+    assert region.kwargs == {
+        "engine_id": "hisparse_instance_dp3",
+        "num_chunks": 1,
+        "rank": 0,
+        "kv_bytes_per_chunk": 4 * page,
+        "cpu_page_size": 64,
+        "barrier": tp_group.barrier,
+        "creator_memory_check": hisparse_runtime_module.check_hisparse_host_memory,
+        "populate_only_on_creator": True,
+    }
+    assert region.view_sizes == [24, 40]
+    assert [pool.shape for pool in pools] == [(24,), (40,)]
+    registration_ranges.assert_called_once_with([24, 40], 4, page)
+    assert [
+        (tensor.data_ptr() - region.base_tensor.data_ptr(), tensor.nbytes)
+        for tensor in pinned
+    ] == [(0, page), (page, 3 * page)]
+    assert region.is_pinned
+
+
+@pytest.mark.parametrize("fail_registration", [None, 1, 2])
+def test_shared_host_pool_tracks_successful_registrations(
+    monkeypatch, fail_registration
+):
+    """Only successful registrations may be unregistered on allocation failure."""
+    page = mmap.PAGESIZE
+    monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", MagicMock())
+    backing = torch.empty(3 * page, dtype=torch.uint8)
+    region = MagicMock(base_tensor=backing, pinned_addresses=[], is_pinned=False)
+    monkeypatch.setattr(
+        hisparse_runtime_module, "SharedOffloadRegion", lambda **kw: region
+    )
+    monkeypatch.setattr(
+        hisparse_runtime_module,
+        "_hisparse_registration_ranges",
+        lambda *args: ((0, page), (page, 3 * page)),
+    )
+    cudart = MagicMock()
+    cudart.cudaHostRegister.side_effect = [
+        SimpleNamespace(value=int(fail_registration == i)) for i in (1, 2)
+    ]
+    monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart)
+    config = SimpleNamespace(
+        instance_id="test", parallel_config=SimpleNamespace(data_parallel_index=0)
+    )
+    context = (
+        pytest.raises(RuntimeError, match="cudaHostRegister failed")
+        if fail_registration
+        else nullcontext()
+    )
+    with context:
+        hisparse_runtime_module.allocate_hisparse_host_pools(
+            config, [3 * page], 3, page, use_shared_host_pool=True
+        )
+    successful = 2 if fail_registration is None else fail_registration - 1
+    assert region.pinned_addresses == [
+        backing.data_ptr() + i * page for i in range(successful)
+    ]
+    assert region.is_pinned == bool(successful)
+    assert region.cleanup.call_count == int(fail_registration is not None)
+    attempts = 2 if fail_registration is None else fail_registration
+    assert (
+        cudart.cudaHostRegister.call_args_list
+        == [
+            call(backing.data_ptr(), page, 0),
+            call(backing.data_ptr() + page, 2 * page, 0),
+        ][:attempts]
+    )
+
+
+def test_shared_host_pool_registers_layer_spans_in_one_backing(monkeypatch):
+    """Aliased tensor configs must not count the shared backing once per view."""
+    region = SimpleNamespace(base_tensor=torch.empty(64, dtype=torch.int8))
+    allocate = MagicMock(return_value=([], [], region))
+    monkeypatch.setattr(
+        hisparse_runtime_module, "allocate_hisparse_host_pools", allocate
+    )
+    config = SimpleNamespace(
+        hisparse_shared_host_pool=True,
+        hisparse_host_num_blocks=4,
+        hisparse_host_block_stride=16,
+        kv_cache_tensors=[
+            SimpleNamespace(
+                host_resident=True,
+                layers=["a", "b"],
+                offset=0,
+                layer_stride=12,
+                size=44,
+            ),
+            SimpleNamespace(
+                host_resident=True, layers=["c"], offset=24, layer_stride=20, size=44
+            ),
+        ],
+    )
+    vllm_config = SimpleNamespace()
+    pool = hisparse_runtime_module.HiSparseHostPool(vllm_config, config)
+    backing = pool.allocate(44)
+    allocate.assert_called_once_with(
+        vllm_config, [12, 12, 20], 4, 16, use_shared_host_pool=True
+    )
+    assert backing.numel() == 44
+    assert backing.data_ptr() == region.base_tensor.data_ptr()
+    assert pool.registered is region.base_tensor
+
+
+def test_hisparse_registration_chunks_end_between_host_blocks():
+    """Registration seams must not bisect a DMA-addressable host block."""
+    page = mmap.PAGESIZE
+    ranges = hisparse_runtime_module._hisparse_registration_ranges(
+        tensor_sizes=[4 * 3 * page, 4 * 5 * page],
+        num_blocks=4,
+        host_block_stride=8 * page,
+        max_chunk_bytes=10 * page,
+    )
+
+    assert ranges == (
+        (0, 9 * page),
+        (9 * page, 17 * page),
+        (17 * page, 27 * page),
+        (27 * page, 32 * page),
+    )
+
+
+def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding():
+    backing = torch.full((10, 2, 3), -1, dtype=torch.float32)
+    cache = backing[1:9]
+    cache[2:4] = 7
+    cache[6:8] = 11
+
+    copy_kv_cache_blocks_inplace(
+        [cache],
+        num_blocks=4,
+        kv_cache_block_copies=[
+            KVCacheBlockCopy(1, 0),
+            KVCacheBlockCopy(3, 2),
+        ],
+    )
+
+    torch.testing.assert_close(cache[0:2], torch.full_like(cache[0:2], 7))
+    torch.testing.assert_close(cache[4:6], torch.full_like(cache[4:6], 11))
+    assert (backing[0] == -1).all()
+    assert (backing[9] == -1).all()
+
+
+@pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)])
+@pytest.mark.parametrize("has_copies", [False, True])
+def test_hisparse_host_copy_waits_for_writes_on_writer_only(
+    monkeypatch, shared, rank, has_copies
+):
+    """Host CoW must read completed writes; shared non-writers only barrier."""
+    worker = _make_hisparse_worker()
+    cache = torch.zeros((2, 1, 1))
+    worker.host_caches = [cache]
+    worker.host_num_blocks = 2
+    worker.shared_host_region = object() if shared else None
+    event = MagicMock()
+    event.synchronize.side_effect = lambda: cache[0].fill_(7)
+    barrier = MagicMock()
+    monkeypatch.setattr(
+        hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: rank
+    )
+    monkeypatch.setattr(
+        hisparse_worker_module, "get_tp_group", lambda: SimpleNamespace(barrier=barrier)
+    )
+    copies = [KVCacheBlockCopy(0, 1)] if has_copies else []
+
+    worker._copy_host_blocks(copies, event)
+
+    writes = has_copies and (not shared or rank == 0)
+    assert cache[1].item() == (7 if writes else 0)
+    assert event.synchronize.call_count == int(writes)
+    assert barrier.call_count == int(shared and has_copies)
+    assert worker._completed_host_copy_dst_ids == ([1] if has_copies else [])
+
+
+def test_hisparse_worker_updates_request_state_mapping_in_place(monkeypatch):
+    worker = _make_hisparse_worker()
+    worker.request_state_indices = torch.arange(4, dtype=torch.int32)
+    worker._pending_invalid_block_ids = [5]
+    invalidations = []
+    worker.invalidate_blocks = lambda blocks, states: invalidations.append(
+        (blocks.copy(), states.clone())
+    )
+    original_ptr = worker.request_state_indices.data_ptr()
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+
+    worker.set_request_state_indices(torch.tensor([3, 1], dtype=torch.int32))
+
+    assert worker.request_state_indices.data_ptr() == original_ptr
+    assert worker.request_state_indices.tolist() == [3, 1, -1, -1]
+    assert len(invalidations) == 1
+    assert invalidations[0][0] == [5]
+    torch.testing.assert_close(
+        invalidations[0][1], torch.tensor([3, 1], dtype=torch.int32)
+    )
+    assert worker._pending_invalid_block_ids == []
+
+
+def test_hisparse_pre_forward_transfer_builds_page_descriptors():
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.kernel_block_size = 2
+    source = torch.empty((3, 2, 4), dtype=torch.uint8)
+    destination = torch.empty((6, 4), dtype=torch.uint8)
+    worker.resident_caches = (source,)
+    worker.host_caches = (destination,)
+    worker.cache_handles = [
+        SimpleNamespace(runtime=SimpleNamespace(resident_source_index=0))
+    ]
+    worker._dma_free_descriptors = []
+    worker._submit_dma_descriptors = MagicMock()
+
+    worker._enqueue_transfers([SparseKVPageTransfer(7, 2, (1,), after_forward=False)])
+
+    descriptors, count = worker._submit_dma_descriptors.call_args.args
+    assert count == 1
+    assert descriptors.src[:count].tolist() == [
+        source.data_ptr() + source.stride(0) * source.element_size()
+    ]
+    assert descriptors.dst[:count].tolist() == [destination.data_ptr() + 4 * 4]
+    assert descriptors.sizes[:count].tolist() == [2 * 4]
+    assert worker._submit_dma_descriptors.call_args.kwargs["transfer_ids"] == (7,)
+
+
+def test_hisparse_eager_mirror_records_transfer_without_page_copy():
+    worker = _make_hisparse_worker()
+    worker.cache_handles = [
+        SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=True))
+    ]
+    worker._record_transfer_completion = MagicMock()
+    worker._enqueue_transfers = MagicMock()
+    transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)]
+
+    worker._submit_transfers(transfers)
+
+    worker._record_transfer_completion.assert_called_once_with(transfers)
+    worker._enqueue_transfers.assert_not_called()
+
+
+@pytest.mark.parametrize("is_host_writer", [False, True])
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_hisparse_tail_restore_preserves_imported_rows(
+    monkeypatch, request, is_host_writer, device
+):
+    """Restore externally pinned host rows on every rank under sync checking."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the device copy check")
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = is_host_writer
+    worker.kernel_block_size = 4
+    worker.host_caches = tuple(
+        torch.arange(24, dtype=torch.float32).reshape(12, 2) + layer * 100
+        for layer in range(2)
+    )
+    if device == "cuda":
+        # Match production's cudaHostRegister pool, which torch.is_pinned()
+        # does not recognize, rather than torch's caching pinned allocator.
+        host_pool, registered = hisparse_runtime_module.allocate_pinned_host_pool(
+            sum(cache.nbytes for cache in worker.host_caches)
+        )
+        request.addfinalizer(
+            lambda: hisparse_runtime_module.release_pinned_state([], [registered])
+        )
+        host_caches = host_pool.view(torch.float32).view(2, 12, 2).unbind()
+        for destination, source in zip(host_caches, worker.host_caches):
+            destination.copy_(source)
+        worker.host_caches = host_caches
+    worker.resident_caches = tuple(
+        torch.full((4, 4, 2), -1.0, device=device) for _ in range(2)
+    )
+    worker.cache_handles = [
+        SimpleNamespace(
+            runtime=SimpleNamespace(resident_source_index=i, eager_host_mirror=True)
+        )
+        for i in range(2)
+    ]
+    worker._record_transfer_completion = MagicMock()
+    stream = torch.cuda.current_stream() if device == "cuda" else MagicMock()
+    monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
+    transfer = SparseKVPageTransfer(7, 2, (1, 3), False, restore=True)
+
+    restore_pages = worker._restore_pages
+    if device == "cuda":
+        import vllm.utils.gpu_sync_debug as gpu_sync_debug
+
+        monkeypatch.setattr(gpu_sync_debug, "_SYNC_CHECK_MODE", "error")
+        monkeypatch.setattr(gpu_sync_debug, "_sync_check_enabled", False)
+        gpu_sync_debug.enable_gpu_sync_check()
+        restore_pages = gpu_sync_debug.with_gpu_sync_check(restore_pages)
+    restore_pages([transfer])
+
+    # Appending into unused capacity must preserve every imported prompt row.
+    for cache, host, block in zip(
+        worker.resident_caches, worker.host_caches, transfer.resident_block_ids
+    ):
+        assert torch.equal(cache[block].cpu(), host[8:12])
+        cache[block, 2:] = 999
+        assert torch.equal(cache[block, :2].cpu(), host[8:10])
+        assert torch.all(cache[0] == -1)
+    worker._record_transfer_completion.assert_called_once_with(
+        [transfer], stream=stream
+    )
+
+
+def test_hisparse_lazy_mirror_copies_transfer_pages():
+    worker = _make_hisparse_worker()
+    worker.cache_handles = [
+        SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=False))
+    ]
+    worker._record_transfer_completion = MagicMock()
+    worker._enqueue_transfers = MagicMock()
+    transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)]
+
+    worker._submit_transfers(transfers)
+
+    worker._enqueue_transfers.assert_called_once_with(transfers)
+    worker._record_transfer_completion.assert_not_called()
+
+
+def test_hisparse_dma_row_mirror_builds_descriptors(monkeypatch):
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.kernel_block_size = 2
+    source = torch.empty((3, 2, 4), dtype=torch.uint8)
+    destination = torch.empty((6, 4), dtype=torch.uint8)
+    worker.resident_caches = (source,)
+    worker.host_caches = (destination,)
+    worker.cache_handles = [
+        SimpleNamespace(
+            runtime=SimpleNamespace(resident_source_index=0), decode_batch=True
+        )
+    ]
+    worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
+    worker.dma_stream = MagicMock()
+    worker.host_write_event = MagicMock()
+    worker._dma_free_descriptors = []
+    worker._pending_dma_descriptors = deque()
+    worker._pending_transfer_events = deque()
+    worker._enqueued_transfer_ids = []
+    worker._dma_submitted = False
+    current_stream = MagicMock()
+    event = MagicMock()
+    swap_blocks_batch = MagicMock()
+    monkeypatch.setattr(
+        hisparse_worker_module, "current_stream", lambda: current_stream
+    )
+    monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
+    monkeypatch.setattr(torch, "Event", lambda: event)
+    monkeypatch.setattr(
+        hisparse_worker_module.ops, "swap_blocks_batch", swap_blocks_batch
+    )
+
+    worker._set_row_mirrors((SparseKVRowMirror((3,), 4, 1),))
+    worker._enqueue_row_dma(range(1))
+
+    src_ptrs, dst_ptrs, sizes = swap_blocks_batch.call_args.args
+    assert src_ptrs.tolist() == [source.data_ptr() + 3 * 4]
+    assert dst_ptrs.tolist() == [destination.data_ptr() + 4 * 4]
+    assert sizes.tolist() == [4]
+    worker.dma_stream.wait_stream.assert_called_once_with(current_stream)
+    worker.host_write_event.record.assert_called_once_with(worker.dma_stream)
+    event.record.assert_called_once_with(worker.dma_stream)
+    assert worker._dma_submitted
+
+
+def test_hisparse_row_dma_uses_resident_spans():
+    worker = _make_hisparse_worker()
+    source = torch.empty((2, 2, 4), dtype=torch.uint8)
+    destination = torch.empty((12, 4), dtype=torch.uint8)
+    worker.is_host_writer = True
+    worker.kernel_block_size = 2
+    worker.resident_caches = (source,)
+    worker.host_caches = (destination,)
+    worker.cache_handles = [
+        SimpleNamespace(
+            runtime=SimpleNamespace(resident_source_index=0), decode_batch=True
+        )
+    ]
+    worker._set_row_mirrors(
+        (
+            SparseKVRowMirror((0,), 4, 2),
+            SparseKVRowMirror((2,), 9, 1),
+        )
+    )
+    worker._dma_free_descriptors = []
+    worker._submit_dma_descriptors = MagicMock()
+
+    worker._enqueue_row_dma(range(1))
+
+    descriptors, count = worker._submit_dma_descriptors.call_args.args
+    assert count == 2
+    assert descriptors.src[:count].tolist() == [
+        source.data_ptr(),
+        source.data_ptr() + 2 * 4,
+    ]
+    assert descriptors.dst[:count].tolist() == [
+        destination.data_ptr() + 4 * 4,
+        destination.data_ptr() + 9 * 4,
+    ]
+    assert descriptors.sizes[:count].tolist() == [2 * 4, 4]
+
+
+def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch):
+    dst_slots = torch.tensor([7, 8, 9], dtype=torch.int64)
+    req_ids = torch.tensor([0, 1, 2], dtype=torch.int32)
+    leader = SimpleNamespace(
+        eager_host_mirror=True,
+        is_group_leader=True,
+        invalidate_written_slots=MagicMock(),
+    )
+    follower = SimpleNamespace(
+        eager_host_mirror=True,
+        is_group_leader=False,
+        invalidate_written_slots=MagicMock(),
+    )
+    handles = [
+        SimpleNamespace(
+            runtime=runtime,
+            decode_batch=True,
+            host_mirror_required=True,
+            num_actual_tokens=3,
+            num_decode_tokens=2,
+            req_id_per_token=req_ids,
+            mirror_slot_mapping=dst_slots,
+        )
+        for runtime in (leader, follower)
+    ]
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.cache_handles = handles
+    worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
+    worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 3),))
+    worker._dma_submitted = False
+    worker._per_layer_mirrored = set()
+    worker._submitted_mirror_layers = set()
+    worker._post_forward_transfers = []
+    worker._enqueue_row_dma = MagicMock()
+    worker.host_write_event = MagicMock()
+    worker._forward_ready_event = MagicMock()
+    current_stream = MagicMock()
+    monkeypatch.setattr(
+        hisparse_worker_module, "current_stream", lambda: current_stream
+    )
+
+    worker.finish_forward()
+
+    worker._forward_ready_event.record.assert_called_once_with()
+    worker._enqueue_row_dma.assert_called_once_with(
+        (0, 1), ready_event=worker._forward_ready_event
+    )
+    leader.invalidate_written_slots.assert_called_once()
+    torch.testing.assert_close(
+        leader.invalidate_written_slots.call_args.args[0], dst_slots
+    )
+    follower.invalidate_written_slots.assert_not_called()
+    worker.host_write_event.record.assert_called_once_with(current_stream)
+
+
+def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors():
+    slots = torch.tensor([7, 8], dtype=torch.int64)
+    runtime = SimpleNamespace(
+        eager_host_mirror=False,
+        is_group_leader=False,
+        invalidate_written_slots=MagicMock(),
+    )
+    handles = [
+        SimpleNamespace(
+            runtime=runtime,
+            decode_batch=False,
+            host_mirror_required=True,
+            num_actual_tokens=2,
+            num_decode_tokens=0,
+            req_id_per_token=torch.empty(0, dtype=torch.int32),
+            mirror_slot_mapping=slots,
+        )
+        for _ in range(2)
+    ]
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.cache_handles = handles
+    worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),))
+    worker._per_layer_mirrored = {0, 1}
+    worker._submitted_mirror_layers = {0, 1}
+    worker._enqueue_row_dma = MagicMock()
+
+    worker._enqueue_host_mirror()
+
+    worker._enqueue_row_dma.assert_not_called()
+
+
+def test_hisparse_finish_forward_submits_lazy_post_forward_transfer(monkeypatch):
+    runtime = SimpleNamespace(eager_host_mirror=False)
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.cache_handles = [SimpleNamespace(runtime=runtime, num_actual_tokens=0)]
+    transfer = SparseKVPageTransfer(7, 2, (1,), after_forward=True)
+    worker._post_forward_transfers = [transfer]
+    worker._forward_ready_event = MagicMock()
+    worker._enqueue_host_mirror = MagicMock()
+    worker._submit_transfers = MagicMock()
+    worker._dma_submitted = False
+    worker._submitted_mirror_layers = set()
+    worker._finish_mirror_phase = MagicMock()
+    worker._release_completed_dma_descriptors = MagicMock()
+    stream = MagicMock()
+    monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
+
+    worker.finish_forward()
+
+    worker._finish_mirror_phase.assert_called_once_with(worker._forward_ready_event)
+    worker._submit_transfers.assert_called_once_with([transfer])
+    worker.host_write_event.record.assert_called_once_with(stream)
+
+
+def test_hisparse_prefill_mirrors_source_groups_and_flushes_partial_group():
+    slots = torch.tensor([7, 8], dtype=torch.int64)
+    source_indices = [0, 0, 1, 1, 1, 1, 2]
+    handles = [
+        SimpleNamespace(
+            runtime=SimpleNamespace(
+                eager_host_mirror=False,
+                is_group_leader=False,
+                resident_source_index=source_index,
+            ),
+            decode_batch=False,
+            host_mirror_required=layer_index < 5,
+            num_actual_tokens=2 if layer_index < 5 else 0,
+            num_decode_tokens=0,
+            req_id_per_token=torch.empty(0, dtype=torch.int32),
+            mirror_slot_mapping=slots,
+        )
+        for layer_index, source_index in enumerate(source_indices)
+    ]
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.cache_handles = handles
+    worker._set_row_mirrors((SparseKVRowMirror((0, 0, 0), 7, 2),))
+    worker._per_layer_mirrored = set()
+    worker._submitted_mirror_layers = set()
+    worker._layer_ready_events = tuple(MagicMock() for _ in handles)
+    worker._enqueue_row_dma = MagicMock()
+
+    for layer_index in range(5):
+        worker._enqueue_layer_mirror(layer_index)
+
+    worker._enqueue_row_dma.assert_called_once_with(
+        (0, 1), ready_event=worker._layer_ready_events[1]
+    )
+    worker._enqueue_host_mirror(ready_event=worker._layer_ready_events[-1])
+    assert worker._enqueue_row_dma.call_args_list[1].args == ((2, 3, 4),)
+    assert worker._enqueue_row_dma.call_args_list[1].kwargs == {
+        "ready_event": worker._layer_ready_events[-1]
+    }
+
+
+def test_hisparse_finish_forward_rejects_partial_per_layer_mirror():
+    slots = torch.tensor([7, 8], dtype=torch.int64)
+    runtime = SimpleNamespace(
+        eager_host_mirror=False,
+        is_group_leader=False,
+        invalidate_written_slots=MagicMock(),
+    )
+    handles = [
+        SimpleNamespace(
+            runtime=runtime,
+            decode_batch=False,
+            host_mirror_required=True,
+            num_actual_tokens=2,
+            num_decode_tokens=0,
+            req_id_per_token=torch.empty(0, dtype=torch.int32),
+            mirror_slot_mapping=slots,
+        )
+        for _ in range(2)
+    ]
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.cache_handles = handles
+    worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),))
+    worker._per_layer_mirrored = {0}
+    worker._submitted_mirror_layers = set()
+    worker._enqueue_row_dma = MagicMock()
+
+    with pytest.raises(RuntimeError, match="did not mirror every active layer"):
+        worker._enqueue_host_mirror()
+
+
+def test_hisparse_finish_forward_orders_next_forward_after_dma(monkeypatch):
+    current_stream = MagicMock()
+    worker = _make_hisparse_worker()
+    worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0"))
+    worker.is_host_writer = True
+    worker._dma_submitted = True
+    worker.host_write_event = MagicMock()
+    worker.cache_handles = [
+        SimpleNamespace(
+            runtime=SimpleNamespace(eager_host_mirror=True), num_actual_tokens=0
+        )
+    ]
+    worker._post_forward_transfers = []
+    worker._pending_dma_descriptors = deque()
+    worker._dma_free_descriptors = []
+    worker._forward_ready_event = MagicMock()
+    worker._finish_mirror_phase = MagicMock()
+    monkeypatch.setattr(
+        hisparse_worker_module, "current_stream", lambda: current_stream
+    )
+
+    worker.finish_forward()
+
+    current_stream.wait_event.assert_called_once_with(worker.host_write_event)
+    assert not worker._dma_submitted
+
+
+def test_hisparse_finish_forward_mirrors_standalone_mtp_cache(monkeypatch):
+    dst_slots = torch.tensor([7, 8], dtype=torch.int64)
+    runtime = SimpleNamespace(
+        eager_host_mirror=True,
+        is_group_leader=False,
+        invalidate_written_slots=MagicMock(),
+    )
+    inactive = SimpleNamespace(
+        runtime=runtime,
+        decode_batch=False,
+        host_mirror_required=False,
+        num_actual_tokens=0,
+        num_decode_tokens=0,
+        req_id_per_token=None,
+        mirror_slot_mapping=dst_slots,
+    )
+    mtp = SimpleNamespace(
+        runtime=runtime,
+        decode_batch=True,
+        host_mirror_required=True,
+        num_actual_tokens=2,
+        num_decode_tokens=2,
+        req_id_per_token=torch.tensor([0, 1], dtype=torch.int32),
+        mirror_slot_mapping=dst_slots,
+    )
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.cache_handles = [inactive, mtp]
+    worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),))
+    worker._per_layer_mirrored = set()
+    worker._submitted_mirror_layers = set()
+    worker._enqueue_row_dma = MagicMock()
+
+    worker._enqueue_host_mirror()
+
+    worker._enqueue_row_dma.assert_called_once_with((1,), ready_event=None)
+
+
+def test_hisparse_shared_host_reader_skips_mirror(monkeypatch):
+    """A non-writer TP rank must not mirror rows into the shared host pool."""
+    dst_slots = torch.tensor([7, 8], dtype=torch.int64)
+    leader = SimpleNamespace(
+        eager_host_mirror=True,
+        is_group_leader=True,
+        invalidate_written_slots=MagicMock(),
+    )
+    handle = SimpleNamespace(
+        runtime=leader,
+        decode_batch=True,
+        host_mirror_required=True,
+        num_actual_tokens=2,
+        num_decode_tokens=1,
+        req_id_per_token=torch.tensor([0], dtype=torch.int32),
+        mirror_slot_mapping=dst_slots,
+    )
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = False
+    worker.cache_handles = [handle]
+    worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),))
+    worker._per_layer_mirrored = set()
+    worker._enqueue_row_dma = MagicMock()
+
+    worker._enqueue_host_mirror()
+
+    worker._enqueue_row_dma.assert_not_called()
+    leader.invalidate_written_slots.assert_called_once()
+
+
+def test_hisparse_shared_host_reader_skips_transfer_completion():
+    """A non-writer TP rank must not acknowledge host writes."""
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = False
+    worker.kernel_block_size = 1
+    worker._enqueued_transfer_ids = []
+    worker._pending_transfer_events = []
+
+    worker._record_transfer_completion([SparseKVPageTransfer(1, 2, (3,), True)])
+
+    assert worker._enqueued_transfer_ids == []
+    assert worker._pending_transfer_events == []
+
+
+def test_hisparse_writer_records_transfer_completion_after_dma(monkeypatch):
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.dma_stream = MagicMock()
+    worker._enqueued_transfer_ids = []
+    worker._pending_transfer_events = deque()
+    event = MagicMock()
+    monkeypatch.setattr(torch, "Event", lambda: event)
+
+    worker._record_transfer_completion(
+        [
+            SparseKVPageTransfer(3, 2, (1,), False),
+            SparseKVPageTransfer(7, 4, (5,), True),
+        ]
+    )
+
+    event.record.assert_called_once_with(worker.dma_stream)
+    assert worker._enqueued_transfer_ids == [3, 7]
+    assert worker._pending_transfer_events == deque([(event, (3, 7))])
+
+
+@pytest.mark.parametrize("is_host_writer", [False, True])
+def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer):
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = is_host_writer
+    worker.hot_backing = SimpleNamespace(device=torch.device("cuda:1"))
+    host_write_events = (MagicMock(), MagicMock())
+    worker.host_write_events = host_write_events
+    worker.host_write_event = host_write_events[1]
+    worker._next_host_write_event = 0
+    worker.host_caches = ()
+    worker.host_num_blocks = 1
+    worker._post_forward_transfers = []
+    worker._pending_invalid_block_ids = []
+    worker.cache_handles = [
+        SimpleNamespace(
+            runtime=SimpleNamespace(eager_host_mirror=True),
+            decode_batch=False,
+            host_mirror_required=False,
+            num_actual_tokens=0,
+            num_decode_tokens=0,
+            req_id_per_token=None,
+        )
+    ]
+    worker._per_layer_mirrored = set()
+    worker._submitted_mirror_layers = set()
+    worker._layer_mirror_callbacks = ()
+    stream = MagicMock()
+    monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
+
+    worker.start_step(
+        SimpleNamespace(
+            host_block_copies=[],
+            command=None,
+            source_block_ids=[],
+            row_mirrors={},
+            all_context_pages_resident=True,
+            row_mirrors_from_resident=False,
+        ),
+        None,
+    )
+    worker.start_step(
+        SimpleNamespace(
+            host_block_copies=[],
+            command=None,
+            source_block_ids=[],
+            row_mirrors={},
+            all_context_pages_resident=True,
+            row_mirrors_from_resident=False,
+        ),
+        None,
+    )
+
+    assert stream.wait_event.call_args_list == [
+        call(host_write_events[1]),
+        call(host_write_events[0]),
+    ]
+    assert worker.host_write_event is host_write_events[1]
+
+
+@pytest.mark.parametrize("tp_rank, expected_copies", [(0, 1), (1, 0)])
+def test_hisparse_shared_host_block_copy_has_one_writer(
+    monkeypatch, tp_rank, expected_copies
+):
+    worker = _make_hisparse_worker()
+    worker.shared_host_region = object()
+    worker.host_caches = (torch.empty(4, 1),)
+    worker.host_num_blocks = 4
+    previous_event = MagicMock()
+    copies = (KVCacheBlockCopy(0, 1),)
+    copy_blocks = MagicMock()
+    tp_group = MagicMock()
+    monkeypatch.setattr(
+        hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: tp_rank
+    )
+    monkeypatch.setattr(hisparse_worker_module, "get_tp_group", lambda: tp_group)
+    monkeypatch.setattr(
+        hisparse_worker_module, "copy_kv_cache_blocks_inplace", copy_blocks
+    )
+
+    worker._copy_host_blocks(copies, previous_event)
+
+    assert copy_blocks.call_count == expected_copies
+    tp_group.barrier.assert_called_once_with()
+    assert worker.take_completed_host_copies() == [1]
+    assert worker.take_completed_host_copies() == []
+
+
+def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch):
+    handle = SimpleNamespace(
+        runtime=SimpleNamespace(eager_host_mirror=True),
+        decode_batch=True,
+        num_actual_tokens=2,
+        num_decode_tokens=2,
+        req_id_per_token=torch.tensor([0, 1]),
+        mirror_slot_mapping=torch.tensor([4, 5]),
+    )
+    worker = _make_hisparse_worker()
+    worker.is_host_writer = True
+    worker.hot_backing = SimpleNamespace(device=torch.device("cpu"))
+    worker.host_write_event = MagicMock()
+    worker._forward_ready_event = MagicMock()
+    worker.host_caches = ()
+    worker.host_num_blocks = 1
+    worker.cache_handles = [handle]
+    worker._per_layer_mirrored = set()
+    worker._submitted_mirror_layers = set()
+    worker._layer_mirror_callbacks = (MagicMock(),)
+    worker._post_forward_transfers = []
+    worker._pending_invalid_block_ids = []
+    worker._enqueue_host_mirror = MagicMock(wraps=worker._enqueue_host_mirror)
+    stream = MagicMock()
+    monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)
+
+    worker.start_step(
+        SimpleNamespace(
+            host_block_copies=[],
+            command=None,
+            source_block_ids=[],
+            row_mirrors={},
+            all_context_pages_resident=True,
+            row_mirrors_from_resident=False,
+        ),
+        None,
+    )
+    worker.finish_forward()
+
+    worker._enqueue_host_mirror.assert_called_once_with(worker._forward_ready_event)
+    assert handle.num_actual_tokens == 0
+    torch.testing.assert_close(handle.mirror_slot_mapping, torch.tensor([4, 5]))
+
+
+def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatch):
+    """Followers must not allocate duplicate runtime state before profiling."""
+    config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(
+            max_num_seqs=2,
+            max_num_batched_tokens=2,
+            async_scheduling=False,
+        ),
+        speculative_config=None,
+        kv_transfer_config=None,
+    )
+    resolved = hisparse_runtime_module.ResolvedHiSparseConfig(
+        top_k=4,
+        device_buffer_size=8,
+    )
+    monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True)
+    monkeypatch.setattr(
+        hisparse_runtime_module.ResolvedHiSparseConfig,
+        "from_vllm_config",
+        classmethod(lambda cls, vllm_config, model_top_k: resolved),
+    )
+    shared_states: list[object] = []
+    streams: list[object] = []
+
+    def create_shared_state(_device, _max_rows, _top_k):
+        shared_states.append(object())
+        return shared_states[-1]
+
+    def create_stream(_device):
+        streams.append(object())
+        return streams[-1]
+
+    monkeypatch.setattr(
+        hisparse_runtime_module, "_create_shared_topk_state", create_shared_state
+    )
+    monkeypatch.setattr(hisparse_runtime_module, "_create_copy_stream", create_stream)
+    monkeypatch.setattr(index_group_module, "_create_side_stream", lambda _: object())
+    monkeypatch.setattr(index_group_module, "_create_event", lambda: object())
+    index_group_builder = index_group_module.SparseMLAIndexGroupBuilder(
+        torch.empty((2, 4), dtype=torch.int32)
+    )
+
+    def make_cache_handle(is_leader: bool):
+        index_group, _ = index_group_builder.register_layer(is_leader)
+        cache_handle = hisparse_runtime_module.create_hisparse_cache_handle(
+            config,
+            model_top_k=4,
+            is_index_group_leader=is_leader,
+            row_width=8,
+            kv_dtype=torch.float32,
+            index_group=index_group,
+            device="cpu",
+        )
+        assert cache_handle is not None
+        return cache_handle
+
+    first_leader = make_cache_handle(True)
+    first_index_group = index_group_builder.current_group
+    first_follower = make_cache_handle(False)
+    second_leader = make_cache_handle(True)
+    second_index_group = index_group_builder.current_group
+    second_follower = make_cache_handle(False)
+
+    assert first_index_group is not None
+    assert second_index_group is not None
+    assert first_follower.runtime.index_group is first_leader.runtime.index_group
+    assert second_follower.runtime.index_group is second_leader.runtime.index_group
+    assert first_leader.runtime.index_group is not second_leader.runtime.index_group
+    assert first_leader.runtime.index_group.copy_stream is first_index_group.side_stream
+    assert (
+        first_leader.runtime.index_group.logical_topk_ready
+        is first_index_group.logical_topk_ready
+    )
+    assert (
+        second_leader.runtime.index_group.copy_stream is second_index_group.side_stream
+    )
+    assert (
+        second_leader.runtime.index_group.logical_topk_ready
+        is second_index_group.logical_topk_ready
+    )
+    assert len(shared_states) == 2
+    assert streams == []
+
+
+def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch):
+    worker = object.__new__(HiSparseConnectorWorker)
+    worker._initialized = True
+    worker._slot_mapping_staging = None
+    worker.dma_stream = None
+    worker.cache_handles = []
+    worker.pinned_host_pools = []
+    worker.shared_host_region = object()
+    released = False
+
+    def release_pinned_state(runtimes, pinned_host_pools, shared_host_region):
+        nonlocal released
+        assert runtimes == []
+        assert pinned_host_pools == []
+        assert shared_host_region is worker.shared_host_region
+        released = True
+
+    monkeypatch.setattr(
+        hisparse_worker_module, "release_pinned_state", release_pinned_state
+    )
+
+    worker.shutdown()
+
+    assert released
 
 
 class _TestReplaySSMMixer(MambaMixer2):
@@ -30,7 +1269,8 @@ def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor:
     return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8)
 
 
-def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group():
+@pytest.mark.parametrize("layers_only", [False, True])
+def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(layers_only):
     mixers = [_TestReplaySSMMixer() for _ in range(3)]
     layer_names = [f"layers.{i}.mixer" for i in range(3)]
     ctx = dict(zip(layer_names, mixers))
@@ -45,7 +1285,10 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group():
         SimpleNamespace(layer_names=[layer_names[1]]),
     ]
 
-    bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
+    if layers_only:
+        bind_kv_cache_to_layers(kv_cache, ctx, kv_cache_groups=kv_cache_groups)
+    else:
+        bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
 
     assert (
         mixers[0]._replayssm_ring_start.data_ptr()
diff --git a/tests/v1/worker/test_worker_memory_snapshot.py b/tests/v1/worker/test_worker_memory_snapshot.py
index fe8a5a21f8dc..3232d3abcf8e 100644
--- a/tests/v1/worker/test_worker_memory_snapshot.py
+++ b/tests/v1/worker/test_worker_memory_snapshot.py
@@ -12,6 +12,7 @@
 
 from vllm.config import set_current_vllm_config
 from vllm.engine.arg_utils import EngineArgs
+from vllm.platforms import current_platform
 from vllm.utils.mem_utils import MemorySnapshot
 from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
 
@@ -34,6 +35,7 @@ def make_operation_tracker(operation_name: str, original_func):
 
     Returns:
         A wrapper function that tracks the operation and calls the original
+
     """
 
     def wrapper(*args, **kwargs):
@@ -66,8 +68,22 @@ def worker_process(
             model="facebook/opt-125m", tensor_parallel_size=2, load_format="dummy"
         ).create_engine_config()
 
+        # Select worker class and patch target based on platform
+        if current_platform.is_xpu():
+            from vllm.v1.worker.xpu_worker import XPUWorker
+
+            worker_cls = XPUWorker
+            init_patch_target = (
+                "vllm.v1.worker.xpu_worker.init_worker_distributed_environment"
+            )
+        else:
+            worker_cls = Worker
+            init_patch_target = (
+                "vllm.v1.worker.gpu_worker.init_worker_distributed_environment"
+            )
+
         # Create worker
-        worker = Worker(
+        worker = worker_cls(
             vllm_config=vllm_config,
             local_rank=rank,
             rank=rank,
@@ -81,7 +97,7 @@ def worker_process(
 
         # Apply minimal patches to track operation order
         init_patch = patch(
-            "vllm.v1.worker.gpu_worker.init_worker_distributed_environment",
+            init_patch_target,
             side_effect=make_operation_tracker(
                 "init_distributed", original_init_worker
             ),
@@ -93,7 +109,7 @@ def worker_process(
         )
         all_reduce_patch = patch(
             "torch.distributed.all_reduce",
-            side_effect=make_operation_tracker("nccl_all_reduce", original_all_reduce),
+            side_effect=make_operation_tracker("dist_all_reduce", original_all_reduce),
         )
 
         with (
@@ -182,13 +198,13 @@ def test_init_distributed_is_called_before_memory_snapshot():
 
         # Raises ValueError if the operation is not found
         init_distributed = rank_ops.index("init_distributed")
-        nccl_all_reduce = rank_ops.index("nccl_all_reduce")
+        dist_all_reduce = rank_ops.index("dist_all_reduce")
         memory_snapshot = rank_ops.index("memory_snapshot")
 
         # Verify order: init_distributed should happen before memory_snapshot
-        assert init_distributed < nccl_all_reduce < memory_snapshot, (
+        assert init_distributed < dist_all_reduce < memory_snapshot, (
             f"Rank {rank}: init_distributed (index {init_distributed}) "
-            f"must happen before nccl_all_reduce (index {nccl_all_reduce}) "
+            f"must happen before dist_all_reduce (index {dist_all_reduce}) "
             f"and memory_snapshot (index {memory_snapshot})"
         )
 
diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py
index 7824c8000d2b..3d03f25e0e76 100644
--- a/tests/v1/worker/test_workspace.py
+++ b/tests/v1/worker/test_workspace.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from typing import cast
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -124,3 +125,65 @@ def test_workspace_lane_validation(monkeypatch) -> None:
 
     with pytest.raises(ValueError, match="at least one"):
         workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0)
+
+
+def test_persistent_resources_are_lazy_and_isolated(monkeypatch) -> None:
+    """Only the requesting slot allocates; each ubatch/lane owns its resource."""
+    active_ubatch = 0
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch)
+    manager = workspace.WorkspaceManager(
+        torch.device("cpu"), num_ubatches=2, num_lanes=2
+    )
+    factory = Mock(side_effect=lambda: {"buffer": torch.zeros(8)})
+    resources = []
+    for ubatch in range(2):
+        active_ubatch = ubatch
+        for lane in range(2):
+            with workspace.use_workspace_lane(lane):
+                resource = manager.get_persistent_resource("scratch", factory)
+                resources.append(resource)
+                resource["buffer"].fill_(len(resources))
+                assert manager.get_persistent_resource("scratch", factory) is resource
+                assert factory.call_count == len(resources)
+
+    assert len({r["buffer"].data_ptr() for r in resources}) == 4
+    for i, resource in enumerate(resources, 1):
+        assert torch.all(resource["buffer"] == i)
+
+
+def test_persistent_resource_lock_applies_to_each_slot(monkeypatch) -> None:
+    monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
+    manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2)
+    factory = Mock(side_effect=object)
+    first = manager.get_persistent_resource("scratch", factory)
+    manager.lock()
+    assert manager.get_persistent_resource("scratch", factory) is first
+    with pytest.raises(AssertionError, match="was not allocated during warmup"):
+        manager.get_persistent_resource("new", factory)
+    with workspace.use_workspace_lane(1):
+        with pytest.raises(AssertionError, match="was not allocated during warmup"):
+            manager.get_persistent_resource("scratch", factory)
+        assert factory.call_count == 1
+        manager.unlock()
+        assert manager.get_persistent_resource("scratch", factory) is not first
+    assert manager.get_persistent_resource("scratch", factory) is first
+
+
+def test_persistent_tensor_preserves_contents_and_rejects_changes() -> None:
+    """Initialization happens once and transient scratch cannot overwrite locks."""
+    manager = workspace.WorkspaceManager(torch.device("cpu"))
+    first = manager.get_persistent("locks", (8,), torch.int32, zero_init=True)
+    assert torch.count_nonzero(first) == 0
+    first.fill_(1)
+    (scratch,) = manager.get_simultaneous(((1024,), torch.uint8))
+    scratch.zero_()
+    assert manager.get_persistent("locks", (8,), torch.int32, zero_init=True) is first
+    assert torch.all(first == 1)
+    with pytest.raises(ValueError, match="requested shape"):
+        manager.get_persistent("locks", (16,), torch.int32)
+    with pytest.raises(ValueError, match="requested shape"):
+        manager.get_persistent("locks", (8,), torch.int64)
+    manager.lock()
+    assert manager.get_persistent("locks", (8,), torch.int32) is first
+    with pytest.raises(AssertionError, match="was not allocated during warmup"):
+        manager.get_persistent("new", (8,), torch.int32)
diff --git a/tests/vllm_test_utils/vllm_test_utils/__init__.py b/tests/vllm_test_utils/vllm_test_utils/__init__.py
index 2818428de4a7..724b57613263 100644
--- a/tests/vllm_test_utils/vllm_test_utils/__init__.py
+++ b/tests/vllm_test_utils/vllm_test_utils/__init__.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-vllm_utils is a package for vLLM testing utilities.
+"""vllm_utils is a package for vLLM testing utilities.
 It does not import any vLLM modules.
 """
 
diff --git a/tests/vllm_test_utils/vllm_test_utils/blame.py b/tests/vllm_test_utils/vllm_test_utils/blame.py
index 9746c3964e21..73eddfdf02a9 100644
--- a/tests/vllm_test_utils/vllm_test_utils/blame.py
+++ b/tests/vllm_test_utils/vllm_test_utils/blame.py
@@ -16,8 +16,7 @@ class BlameResult:
 
 @contextlib.contextmanager
 def blame(func: Callable) -> Generator[BlameResult, None, None]:
-    """
-    Trace the function calls to find the first function that satisfies the
+    """Trace the function calls to find the first function that satisfies the
     condition. The trace stack will be stored in the result.
 
     Usage:
diff --git a/tests/vllm_test_utils/vllm_test_utils/monitor.py b/tests/vllm_test_utils/vllm_test_utils/monitor.py
index ba22bde8795b..658347889efe 100644
--- a/tests/vllm_test_utils/vllm_test_utils/monitor.py
+++ b/tests/vllm_test_utils/vllm_test_utils/monitor.py
@@ -21,8 +21,7 @@ class MonitoredValues(Generic[_T]):
 def monitor(
     measure_func: Callable[[], _T],
 ) -> Generator[MonitoredValues[_T], None, None]:
-    """
-    Trace the function calls to continuously monitor the change of
+    """Trace the function calls to continuously monitor the change of
     a value.
 
     Usage:
diff --git a/tests/watermarking/test_gumbel.py b/tests/watermarking/test_gumbel.py
index e3b4125172b0..dd32bda62793 100644
--- a/tests/watermarking/test_gumbel.py
+++ b/tests/watermarking/test_gumbel.py
@@ -5,8 +5,15 @@
 import torch
 
 from vllm.platforms import current_platform
-from vllm.v1.watermarking import GumbelWatermarkDetector, GumbelWatermarker
+from vllm.v1.watermarking import (
+    DualKeyGumbelWatermarkDetector,
+    GumbelWatermarkDetector,
+    GumbelWatermarker,
+    derive_watermark_key,
+)
 from vllm.v1.watermarking.gumbel import _gamma_survival_integer_shape
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
 
 
 def test_gamma_survival_integer_shape():
@@ -34,6 +41,67 @@ def test_detector_deduplicates_context_even_when_target_differs():
     assert detection.num_scored_tokens == 3
 
 
+def test_dual_key_detector_scores_each_token_against_both_keys():
+    token_ids = [1, 2, 3, 4, 5]
+    key_a_detector = GumbelWatermarkDetector(
+        key=derive_watermark_key(42, b"key_a"),
+        context_width=1,
+        deduplicate_contexts=False,
+    )
+    key_b_detector = GumbelWatermarkDetector(
+        key=derive_watermark_key(42, b"key_b"),
+        context_width=1,
+        deduplicate_contexts=False,
+    )
+    detector = DualKeyGumbelWatermarkDetector(
+        key=42,
+        context_width=1,
+        deduplicate_contexts=False,
+        alpha=0.5,
+    )
+
+    key_a = key_a_detector.detect(token_ids)
+    key_b = key_b_detector.detect(token_ids)
+    dual = detector.detect(token_ids)
+
+    assert dual.score == pytest.approx((key_a.score + key_b.score) / 2)
+    assert dual.p_value == pytest.approx(
+        _gamma_survival_integer_shape(dual.score * 2, dual.num_scored_tokens * 2)
+    )
+
+
+@pytest.mark.parametrize("alpha", [0.2, 0.7])
+def test_dual_key_detector_weights_and_recalibrates_scores(alpha: float):
+    token_ids = [1, 2, 3, 4, 5]
+    key_a = GumbelWatermarkDetector(
+        key=derive_watermark_key(42, b"key_a"), deduplicate_contexts=False
+    ).detect(token_ids)
+    key_b = GumbelWatermarkDetector(
+        key=derive_watermark_key(42, b"key_b"), deduplicate_contexts=False
+    ).detect(token_ids)
+    dual = DualKeyGumbelWatermarkDetector(
+        key=42, deduplicate_contexts=False, alpha=alpha
+    ).detect(token_ids)
+
+    variance = (1 - alpha) ** 2 + alpha**2
+    expected_p_value = torch.special.gammaincc(
+        torch.tensor(dual.num_scored_tokens / variance, dtype=torch.float64),
+        torch.tensor(dual.score / variance, dtype=torch.float64),
+    ).item()
+
+    assert dual.score == pytest.approx((1 - alpha) * key_a.score + alpha * key_b.score)
+    assert dual.p_value == pytest.approx(expected_p_value)
+
+
+def test_dual_key_detector_rejects_invalid_alpha():
+    with pytest.raises(ValueError, match="alpha must be between 0 and 1"):
+        DualKeyGumbelWatermarkDetector(key=42, alpha=1.1)
+
+
+def test_dual_key_detector_default_alpha():
+    assert DualKeyGumbelWatermarkDetector(key=42).alpha == 0.2
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
 )
@@ -45,10 +113,8 @@ def test_fused_watermarker_matches_cpu(key: int, context_width: int):
     logits = torch.randn(32, 8193)
     watermarker = GumbelWatermarker(key, context_width)
 
-    expected = watermarker.sample(logits, contexts, lambda values: None).token_ids
-    actual = watermarker.sample(
-        logits.cuda(), contexts.cuda(), lambda values: None
-    ).token_ids.cpu()
+    expected = watermarker.sample(logits, contexts).token_ids
+    actual = watermarker.sample(logits.cuda(), contexts.cuda()).token_ids.cpu()
 
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)
 
@@ -61,7 +127,7 @@ def test_fused_watermarker_handles_nan_logits():
     logits = torch.full((2, 1025), float("nan"), device="cuda")
     watermarker = GumbelWatermarker(key=42, context_width=4)
 
-    token_ids = watermarker.sample(logits, contexts, lambda values: None).token_ids
+    token_ids = watermarker.sample(logits, contexts).token_ids
 
     assert torch.all((token_ids >= 0) & (token_ids < logits.shape[-1]))
 
@@ -76,9 +142,58 @@ def test_fused_watermarker_handles_noncontiguous_inputs():
     logits = torch.randn(8, 2050, device="cuda")[:, ::2]
     watermarker = GumbelWatermarker(key=42, context_width=4)
 
-    expected = watermarker.sample(
-        logits.contiguous(), contexts.contiguous(), lambda values: None
-    ).token_ids
-    actual = watermarker.sample(logits, contexts, lambda values: None).token_ids
+    expected = watermarker.sample(logits.contiguous(), contexts.contiguous()).token_ids
+    actual = watermarker.sample(logits, contexts).token_ids
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+@pytest.mark.parametrize("use_fp64", [False, True])
+def test_philox_gumbel_sample_skip_mask_matches_separate_samplers(use_fp64: bool):
+    torch.manual_seed(0)
+    logits = torch.randn(8, 8193, device="cuda")
+    context_storage = torch.randint(0, 248320, (8, 8), dtype=torch.int64, device="cuda")
+    contexts = context_storage[:, ::2]
+    repeated_mask = torch.tensor(
+        [False, False, False, True, False, False, True, False], device="cuda"
+    )
+    watermarking = torch.tensor(
+        [True, False, True, True, False, True, True, False], device="cuda"
+    )
+    req_indices = torch.tensor([3, 1, 7, 0, 4, 2, 6, 5], device="cuda")
+    temperatures = torch.ones(8, dtype=torch.float32, device="cuda")
+    temperatures[2] = 0
+    seeds = torch.arange(8, dtype=torch.int64, device="cuda") + 1000
+    positions = torch.arange(8, dtype=torch.int64, device="cuda") + 100
+    watermark_mask = (
+        watermarking[req_indices] & (temperatures[req_indices] != 0) & ~repeated_mask
+    )
+
+    watermarked = philox_gumbel_sample(logits, contexts, 42)
+    ordinary = gumbel_sample(
+        logits,
+        req_indices,
+        temperatures,
+        seeds,
+        positions,
+        apply_temperature=False,
+        is_drafting=False,
+        use_fp64=use_fp64,
+    )
+    expected = torch.where(watermark_mask, watermarked, ordinary)
+    actual = philox_gumbel_sample(
+        logits,
+        contexts,
+        42,
+        skip_mask=~watermark_mask,
+        expanded_idx_mapping=req_indices,
+        temperatures=temperatures,
+        seeds=seeds,
+        positions=positions,
+        use_fp64=use_fp64,
+    )
 
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)
diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py
index 0452957ffcb3..65ea9e93ad40 100644
--- a/tests/watermarking/test_watermarking.py
+++ b/tests/watermarking/test_watermarking.py
@@ -9,44 +9,98 @@
 
 from vllm import SamplingParams
 from vllm.config.watermarking import WatermarkConfig
-from vllm.v1.watermarking import create_watermarker
+from vllm.platforms import current_platform
+from vllm.v1.watermarking import (
+    DualKeyGumbelWatermarker,
+    SupportsSpeculativeDecoding,
+    create_watermarker,
+    derive_watermark_key,
+)
 from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
-from vllm.v1.watermarking.watermarker import WatermarkSample
+from vllm.v1.watermarking.gumbel import GumbelWatermarker
+from vllm.v1.watermarking.spec_decode import (
+    DraftWatermarker,
+    _resolve_watermark_key,
+    create_speculative_draft_watermarker,
+    create_speculative_target_watermarker,
+    speculative_target_watermark_key,
+)
+from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
 from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.sample.watermark import (
+    philox_gumbel_sample,
+    repeated_context_mask,
+)
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
 
 
-@pytest.mark.parametrize("algorithm", ["gumbel"])
+class StubWatermarker(Watermarker):
+    context_width = 1
+
+    def _sample_watermarked(self, logits, contexts):
+        return WatermarkSample(torch.tensor([7, 7]), logits + 10)
+
+
+def _argmax_sampler(logits: torch.Tensor) -> torch.Tensor:
+    return logits.argmax(dim=-1)
+
+
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
 def test_watermarker_contract(algorithm: str):
     watermarker = create_watermarker(
         WatermarkConfig(algorithm=algorithm, key=42, context_width=4)
     )
     logits = torch.zeros(2, 128)
     contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])
-    random_sample = lambda sample_logits: sample_logits.argmax(dim=-1)
 
-    first = watermarker.sample(logits, contexts, random_sample)
-    second = watermarker.sample(logits, contexts, random_sample)
+    first = watermarker.sample(logits, contexts, _argmax_sampler)
+    second = watermarker.sample(logits, contexts, _argmax_sampler)
 
     assert first.token_ids.shape == (2,)
     assert first.logits.shape == logits.shape
     assert torch.equal(first.token_ids, second.token_ids)
 
 
-def test_gumbel_config_warns_about_degenerate_generations(monkeypatch):
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
+@pytest.mark.parametrize(
+    "config_overrides",
+    [
+        {"deduplicate_contexts": "none"},
+        {"deduplicate_contexts_max_history": 1023},
+    ],
+)
+def test_gumbel_config_warns_when_context_deduplication_is_weak(
+    monkeypatch, algorithm, config_overrides
+):
     messages: list[str] = []
     monkeypatch.setattr(
         "vllm.config.watermarking.logger.warning_once",
         lambda message, *, scope: messages.append(message),
     )
 
-    WatermarkConfig(key=42)
+    WatermarkConfig(key=42, algorithm=algorithm)
+    WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=1024)
+    WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=None)
+    WatermarkConfig(key=42, algorithm=algorithm, **config_overrides)
 
     assert messages == [
-        "Single-key Gumbel-max watermarking may increase the frequency of "
-        "degenerate generations, including repetition loops."
+        (
+            "Gumbel-max watermarking with context deduplication disabled "
+            "or limited to fewer than 1024 positions may increase the frequency of "
+            "degenerate generations, including repetition loops. Use "
+            "deduplicate_contexts='single_turn' or 'all' with "
+            "deduplicate_contexts_max_history at least 1024 or null to mitigate this."
+        )
     ]
 
 
+def test_context_deduplication_history_can_be_unbounded():
+    config = WatermarkConfig(key=42, deduplicate_contexts_max_history=None)
+
+    assert config.deduplicate_contexts_max_history is None
+
+
 def test_large_context_width_warns_but_is_allowed():
     config = WatermarkConfig(key=42, context_width=17)
 
@@ -56,6 +110,158 @@ def test_large_context_width_warns_but_is_allowed():
     assert watermarker.context_width == 17
 
 
+def test_dual_key_watermarker_uses_domain_separated_keys():
+    config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
+
+    watermarker = create_watermarker(config)
+    assert isinstance(watermarker, SupportsSpeculativeDecoding)
+    draft = watermarker.draft_watermarker
+    target = watermarker.target_watermarker
+
+    assert isinstance(watermarker, DualKeyGumbelWatermarker)
+    assert watermarker.alpha == 0.1
+    assert draft.prf.key == derive_watermark_key(42, b"key_a")
+    assert target.prf.key == derive_watermark_key(42, b"key_b")
+    assert target.prf.key != draft.prf.key
+
+
+def test_dual_key_watermarker_routes_tokens_with_alpha():
+    watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25)
+    logits = torch.zeros(2, 16)
+    contexts = torch.tensor([[1, 2], [3, 4]])
+    key_a = watermarker.draft_watermarker.sample(logits, contexts, lambda _: None)
+    key_b = watermarker.target_watermarker.sample(logits, contexts, lambda _: None)
+
+    def route(routing_logits):
+        torch.testing.assert_close(
+            routing_logits.softmax(dim=-1),
+            torch.tensor([[0.75, 0.25], [0.75, 0.25]]),
+        )
+        return torch.tensor([0, 1])
+
+    sampled = watermarker.sample(logits, contexts, route)
+
+    assert torch.equal(
+        sampled.token_ids, torch.stack([key_a.token_ids[0], key_b.token_ids[1]])
+    )
+
+
+def test_dual_key_routing_logits_are_cached():
+    alpha = 0.25
+    watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha)
+    logits = torch.zeros(2, 16)
+    contexts = torch.tensor([[1, 2], [3, 4]])
+    recorded: list[torch.Tensor] = []
+
+    def route(routing_logits):
+        recorded.append(routing_logits)
+        return torch.tensor([0, 1])
+
+    watermarker.sample(logits, contexts, route)
+    watermarker.sample(logits, contexts, route)
+
+    expected = torch.tensor([1 - alpha, alpha], dtype=torch.float32).log()
+    expected = expected.expand(2, -1)
+    assert len(recorded) == 2
+    assert all(torch.equal(seen, expected) for seen in recorded)
+    assert len(watermarker._routing_logits_cache) == 1
+
+
+@pytest.mark.parametrize("alpha", [0.0, 0.25, 1.0])
+def test_dual_key_watermarker_skips_masked_rows(alpha: float):
+    watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha)
+    torch.manual_seed(0)
+    logits = torch.randn(4, 32)
+    contexts = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+    skip_mask = torch.tensor([True, False, True, False])
+
+    unmasked = watermarker.sample(logits, contexts, _argmax_sampler)
+    masked = watermarker.sample(logits, contexts, _argmax_sampler, skip_mask)
+
+    ordinary = _argmax_sampler(logits)
+    assert torch.equal(masked.token_ids[skip_mask], ordinary[skip_mask])
+    assert torch.equal(masked.token_ids[~skip_mask], unmasked.token_ids[~skip_mask])
+    assert not torch.equal(unmasked.token_ids[skip_mask], ordinary[skip_mask])
+
+
+def test_dual_key_watermarker_routing_requires_a_random_sampler():
+    watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25)
+
+    with pytest.raises(ValueError, match="random sampler"):
+        watermarker.sample(torch.zeros(2, 16), torch.tensor([[1, 2], [3, 4]]))
+
+
+def test_speculative_decoding_uses_fixed_dual_key_roles():
+    watermarker = create_watermarker(
+        WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25)
+    )
+
+    target = create_speculative_target_watermarker(watermarker)
+    draft = create_speculative_draft_watermarker(
+        watermarker,
+        max_num_reqs=1,
+        device=torch.device("cpu"),
+        allow_target_only=False,
+    )
+
+    assert target.prf.key == derive_watermark_key(42, b"key_b")
+    assert draft is not None
+    assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a")
+
+
+def test_recovery_key_rejects_the_unsplit_dual_key_watermarker():
+    watermarker = create_watermarker(
+        WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25)
+    )
+
+    with pytest.raises(ValueError, match="keys the target role separately"):
+        _resolve_watermark_key(watermarker)
+
+    target = create_speculative_target_watermarker(watermarker)
+    assert _resolve_watermark_key(target) == derive_watermark_key(42, b"key_b")
+
+
+@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
+def test_config_key_resolution_matches_the_model_runner(algorithm: str):
+    config = WatermarkConfig(algorithm=algorithm, key=42, alpha=0.25)
+    runtime_key = _resolve_watermark_key(
+        create_speculative_target_watermarker(create_watermarker(config))
+    )
+
+    assert speculative_target_watermark_key(config) == runtime_key
+
+
+def test_config_key_resolution_returns_none_when_watermarking_is_disabled():
+    assert speculative_target_watermark_key(None) is None
+
+
+def test_target_only_speculative_watermarking_skips_draft_watermarker():
+    watermarker = create_watermarker(WatermarkConfig(algorithm="gumbel", key=42))
+
+    assert not isinstance(watermarker, SupportsSpeculativeDecoding)
+    with pytest.raises(ValueError, match="does not support speculative decoding"):
+        create_speculative_draft_watermarker(
+            watermarker,
+            max_num_reqs=1,
+            device=torch.device("cpu"),
+            allow_target_only=False,
+        )
+    assert (
+        create_speculative_draft_watermarker(
+            watermarker,
+            max_num_reqs=1,
+            device=torch.device("cpu"),
+            allow_target_only=True,
+        )
+        is None
+    )
+
+
+def test_dual_key_derivation_is_stable():
+    assert derive_watermark_key(32, b"key_a") == 16368605726115524094
+    assert derive_watermark_key(32, b"key_b") == 4799302812959726346
+
+
 def test_sampling_params_can_disable_watermarking():
     assert SamplingParams().watermarking
     assert not SamplingParams.from_optional(watermarking=False).watermarking
@@ -83,14 +289,9 @@ def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch):
 
 
 def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
-    class StubWatermarker:
-        context_width = 1
-
-        def sample(self, logits, contexts, random_sample):
-            return WatermarkSample(torch.tensor([7, 7]), logits + 10)
-
     sampler = object.__new__(GPUWatermarkSampler)
     sampler.watermarker = StubWatermarker()
+    sampler.deduplicate_contexts = "single_turn"
     sampler.watermarking = SimpleNamespace(
         np=np.array([True, False]), gpu=torch.tensor([True, False])
     )
@@ -102,8 +303,11 @@ def sample(self, logits, contexts, random_sample):
     sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
         2, 1, dtype=torch.int64
     )
+    sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: torch.zeros(
+        2, dtype=torch.bool
+    )
     monkeypatch.setattr(
-        "vllm.v1.watermarking.gpu_sampler.gumbel_sample",
+        "vllm.v1.watermarking.watermarker.gumbel_sample",
         lambda *args, **kwargs: torch.tensor([3, 4]),
     )
     logits = torch.zeros(2, 8)
@@ -123,11 +327,524 @@ def sample(self, logits, contexts, random_sample):
     assert torch.equal(output_logits[1], logits[1])
 
 
+def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch):
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = StubWatermarker()
+    sampler.deduplicate_contexts = "single_turn"
+    sampler.watermarking = SimpleNamespace(
+        np=np.array([True, True]), gpu=torch.tensor([True, True])
+    )
+    sampler.sampling_states = SimpleNamespace(
+        temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
+        seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
+    )
+    sampler.use_fp64_gumbel = False
+    sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
+        2, 1, dtype=torch.int64
+    )
+    sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: (
+        torch.tensor([True, False])
+    )
+    monkeypatch.setattr(
+        "vllm.v1.watermarking.watermarker.gumbel_sample",
+        lambda *args, **kwargs: torch.tensor([3, 4]),
+    )
+    logits = torch.zeros(2, 8)
+
+    sampled, output_logits = sampler._sample_random(
+        logits,
+        torch.tensor([0, 1]),
+        np.array([0, 1]),
+        torch.zeros(2, dtype=torch.int64),
+        None,
+        None,
+        False,
+    )
+
+    assert torch.equal(sampled, torch.tensor([3, 7]))
+    assert torch.equal(output_logits[0], logits[0])
+    assert torch.equal(output_logits[1], torch.full((8,), 10.0))
+
+
+def test_gpu_sampler_can_disable_context_deduplication(monkeypatch):
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = StubWatermarker()
+    sampler.deduplicate_contexts = "none"
+    sampler.watermarking = SimpleNamespace(
+        np=np.array([True, True]), gpu=torch.tensor([True, True])
+    )
+    sampler.sampling_states = SimpleNamespace(
+        temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
+        seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
+    )
+    sampler.use_fp64_gumbel = False
+    sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
+        2, 1, dtype=torch.int64
+    )
+    sampler._get_repeated_contexts = lambda *args: pytest.fail(
+        "context deduplication should not run"
+    )
+    logits = torch.zeros(2, 8)
+
+    sampled, output_logits = sampler._sample_random(
+        logits,
+        torch.tensor([0, 1]),
+        np.array([0, 1]),
+        torch.zeros(2, dtype=torch.int64),
+        None,
+        None,
+        False,
+    )
+
+    assert torch.equal(sampled, torch.tensor([7, 7]))
+    assert torch.equal(output_logits, torch.full((2, 8), 10.0))
+
+
+def test_repeated_context_mask_ignores_prompt_tokens():
+    all_token_ids = torch.tensor(
+        [
+            [8, 9, 1, 2, 1, 2],
+            [3, 4, 1, 2, 3, 4],
+            [0, 0, 0, 0, 0, 0],
+        ],
+        dtype=torch.int32,
+    )
+    req_indices = torch.tensor([0, 1, -1])
+    prompt_lens = torch.tensor([2, 2, 0])
+    total_lens = torch.tensor([6, 6, 0])
+    contexts = torch.tensor([[1, 2], [3, 4], [-1, -1]])
+
+    repeated = repeated_context_mask(
+        all_token_ids,
+        req_indices,
+        prompt_lens,
+        total_lens,
+        contexts,
+    )
+
+    assert torch.equal(repeated, torch.tensor([True, False, False]))
+
+
+def test_repeated_context_mask_can_include_prompt_tokens():
+    all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
+    req_indices = torch.tensor([0])
+    prompt_lens = torch.tensor([4])
+    total_lens = torch.tensor([10])
+    contexts = torch.tensor([[1, 2, 3, 4]])
+
+    single_turn = repeated_context_mask(
+        all_token_ids, req_indices, prompt_lens, total_lens, contexts
+    )
+    all_history = repeated_context_mask(
+        all_token_ids,
+        req_indices,
+        prompt_lens,
+        total_lens,
+        contexts,
+        include_prompt=True,
+    )
+
+    assert not single_turn.item()
+    assert all_history.item()
+
+
+def test_repeated_context_mask_can_skip_partial_contexts():
+    all_token_ids = torch.tensor([[10, 11, 12, 13, 1, 2, 3, 4]], dtype=torch.int32)
+    req_indices = torch.tensor([0])
+    prompt_lens = torch.tensor([4])
+    contexts = torch.tensor([[-1, -1, -1, 1]])
+
+    partial = repeated_context_mask(
+        all_token_ids,
+        req_indices,
+        prompt_lens,
+        torch.tensor([5]),
+        contexts,
+        include_prompt=True,
+        skip_partial_context=True,
+    )
+    complete = repeated_context_mask(
+        all_token_ids,
+        req_indices,
+        prompt_lens,
+        torch.tensor([8]),
+        torch.tensor([[1, 2, 3, 4]]),
+        include_prompt=True,
+        skip_partial_context=True,
+    )
+
+    assert partial.item()
+    assert not complete.item()
+
+
+def test_gpu_sampler_contexts_are_completion_local_for_all_scopes():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = SimpleNamespace(context_width=4)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
+        total_len=SimpleNamespace(gpu=torch.tensor([6])),
+    )
+    request_indices = torch.tensor([0])
+
+    sampler.deduplicate_contexts = "single_turn"
+    single_turn = sampler._get_contexts(request_indices)
+    sampler.deduplicate_contexts = "all"
+    all_history = sampler._get_contexts(request_indices)
+
+    assert torch.equal(single_turn, torch.tensor([[-1, -1, 1, 2]]))
+    assert torch.equal(all_history, single_turn)
+
+
+def test_gpu_sampler_all_history_skips_partial_context():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = SimpleNamespace(context_width=4)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
+        total_len=SimpleNamespace(gpu=torch.tensor([6])),
+    )
+    sampler.deduplicate_contexts_max_history = None
+    request_indices = torch.tensor([0])
+    contexts = sampler._get_contexts(request_indices)
+
+    sampler.deduplicate_contexts = "single_turn"
+    single_turn = sampler._get_repeated_contexts(request_indices, contexts)
+    sampler.deduplicate_contexts = "all"
+    all_history = sampler._get_repeated_contexts(request_indices, contexts)
+
+    assert not single_turn.item()
+    assert all_history.item()
+
+
+def test_repeated_context_mask_respects_max_history():
+    all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
+    req_indices = torch.tensor([0])
+    prompt_lens = torch.tensor([0])
+    total_lens = torch.tensor([10])
+    contexts = torch.tensor([[1, 2, 3, 4]])
+
+    full_history = repeated_context_mask(
+        all_token_ids, req_indices, prompt_lens, total_lens, contexts
+    )
+    last_six = repeated_context_mask(
+        all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=6
+    )
+    last_five = repeated_context_mask(
+        all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=5
+    )
+
+    assert full_history.item()
+    assert last_six.item()
+    assert not last_five.item()
+
+
+@pytest.mark.parametrize("max_history", [0, -1])
+def test_repeated_context_mask_rejects_non_positive_max_history(max_history: int):
+    with pytest.raises(ValueError, match="max_history must be positive or None"):
+        repeated_context_mask(
+            torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
+            torch.tensor([0], dtype=torch.int32),
+            torch.tensor([0], dtype=torch.int32),
+            torch.tensor([4], dtype=torch.int32),
+            torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
+            max_history=max_history,
+        )
+
+
+def test_gpu_sampler_wires_request_history_scope():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(
+            gpu=torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
+        ),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
+        total_len=SimpleNamespace(gpu=torch.tensor([10])),
+    )
+    sampler.deduplicate_contexts_max_history = None
+    request_indices = torch.tensor([0])
+    contexts = torch.tensor([[1, 2, 3, 4]])
+
+    sampler.deduplicate_contexts = "single_turn"
+    single_turn = sampler._get_repeated_contexts(request_indices, contexts)
+    sampler.deduplicate_contexts = "all"
+    all_history = sampler._get_repeated_contexts(request_indices, contexts)
+
+    assert not single_turn.item()
+    assert all_history.item()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+@pytest.mark.parametrize("max_history", [5, 6])
+@pytest.mark.parametrize("include_prompt", [False, True])
+def test_repeated_context_mask_max_history_accelerator_parity(
+    max_history: int, include_prompt: bool
+):
+    inputs = (
+        torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]),
+        torch.tensor([0]),
+        torch.tensor([0]),
+        torch.tensor([10]),
+        torch.tensor([[1, 2, 3, 4]]),
+    )
+
+    expected = repeated_context_mask(
+        *inputs, max_history=max_history, include_prompt=include_prompt
+    )
+    actual = repeated_context_mask(
+        *(value.cuda() for value in inputs),
+        max_history=max_history,
+        include_prompt=include_prompt,
+    ).cpu()
+
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_repeated_context_mask_partial_context_accelerator_parity():
+    inputs = (
+        torch.tensor([[10, 11, 12, 13, 1]], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        torch.tensor([4], dtype=torch.int32),
+        torch.tensor([5], dtype=torch.int32),
+        torch.tensor([[-1, -1, -1, 1]], dtype=torch.int64),
+    )
+
+    expected = repeated_context_mask(
+        *inputs, include_prompt=True, skip_partial_context=True
+    )
+    actual = repeated_context_mask(
+        *(value.cuda() for value in inputs),
+        include_prompt=True,
+        skip_partial_context=True,
+    ).cpu()
+
+    assert torch.equal(actual, expected)
+
+
+def _unaligned_max_history_inputs() -> tuple[torch.Tensor, ...]:
+    """One row whose only repeated context ends at position 680.
+
+    680 is not a multiple of the kernel's 512-position block, so the scan window
+    selected by ``max_history`` starts in the middle of a block.
+    """
+    history = list(range(1_000, 2_200))
+    history[676:680] = [1, 2, 3, 4]
+    return (
+        torch.tensor([history], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        torch.tensor([0], dtype=torch.int32),
+        torch.tensor([len(history)], dtype=torch.int32),
+        torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
+    )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+@pytest.mark.parametrize(
+    "max_history,expected", [(1200 - 680, True), (1200 - 680 - 1, False)]
+)
+def test_repeated_context_mask_unaligned_max_history_accelerator_parity(
+    max_history: int, expected: bool
+):
+    inputs = _unaligned_max_history_inputs()
+
+    reference = repeated_context_mask(*inputs, max_history=max_history)
+    actual = repeated_context_mask(
+        *(value.cuda() for value in inputs), max_history=max_history
+    ).cpu()
+
+    assert torch.equal(reference, torch.tensor([expected]))
+    assert torch.equal(actual, reference)
+
+
+def _repeated_context_inputs(
+    context_width: int, device: str = "cpu"
+) -> tuple[torch.Tensor, ...]:
+    """Build repeated, unique, and padded request rows for mask tests."""
+    prompt = [101, 102, 103, 104]
+    repeated_output = [*range(1, context_width + 1)] * 2
+    unique_output = [*range(1, context_width + 2)]
+    max_len = len(prompt) + len(repeated_output)
+    rows = [
+        prompt + repeated_output,
+        prompt + unique_output,
+        [],
+    ]
+    all_token_ids = torch.zeros((3, max_len), dtype=torch.int32, device=device)
+    for row, token_ids in enumerate(rows):
+        all_token_ids[row, : len(token_ids)] = torch.tensor(
+            token_ids, dtype=torch.int32, device=device
+        )
+    return (
+        all_token_ids,
+        torch.tensor([0, 1, -1], dtype=torch.int32, device=device),
+        torch.tensor([len(prompt), len(prompt), 0], dtype=torch.int32, device=device),
+        torch.tensor([len(rows[0]), len(rows[1]), 0], dtype=torch.int32, device=device),
+        torch.tensor(
+            [
+                repeated_output[-context_width:],
+                unique_output[-context_width:],
+                [-1] * context_width,
+            ],
+            dtype=torch.int64,
+            device=device,
+        ),
+    )
+
+
+@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17])
+def test_repeated_context_mask(context_width: int):
+    repeated = repeated_context_mask(*_repeated_context_inputs(context_width))
+
+    assert torch.equal(repeated, torch.tensor([True, False, False]))
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17])
+def test_repeated_context_mask_accelerator_parity(context_width: int):
+    cpu_inputs = _repeated_context_inputs(context_width)
+    accelerator_inputs = list(_repeated_context_inputs(context_width, "cuda"))
+    contexts = accelerator_inputs[-1]
+    storage = torch.empty(
+        contexts.shape[0], contexts.shape[1] * 2, dtype=contexts.dtype, device="cuda"
+    )
+    storage[:, ::2] = contexts
+    accelerator_inputs[-1] = storage[:, ::2]
+
+    expected = repeated_context_mask(*cpu_inputs)
+    actual = repeated_context_mask(*accelerator_inputs).cpu()
+
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_repeated_context_mask_scans_multiple_blocks():
+    context = [1, 2, 3, 4]
+    repeated_output = [*range(10_000, 11_030), *context, 99, *context]
+    unique_output = list(range(20_000, 21_039))
+    prompt = [101, 102, 103, 104]
+    rows = [prompt + repeated_output, prompt + unique_output]
+    all_token_ids = torch.zeros((2, len(rows[0])), dtype=torch.int32)
+    for row, token_ids in enumerate(rows):
+        all_token_ids[row, : len(token_ids)] = torch.tensor(token_ids)
+    inputs = (
+        all_token_ids,
+        torch.tensor([0, 1], dtype=torch.int32),
+        torch.tensor([len(prompt), len(prompt)], dtype=torch.int32),
+        torch.tensor([len(rows[0]), len(rows[1])], dtype=torch.int32),
+        torch.tensor([context, unique_output[-4:]], dtype=torch.int64),
+    )
+
+    expected = repeated_context_mask(*inputs)
+    actual = repeated_context_mask(*(value.cuda() for value in inputs)).cpu()
+
+    assert torch.equal(expected, torch.tensor([True, False]))
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_repeated_context_mask_last_request_at_capacity():
+    max_model_len = 1024
+    all_token_ids = torch.arange(
+        2 * max_model_len, dtype=torch.int32, device="cuda"
+    ).reshape(2, max_model_len)
+    inputs = (
+        all_token_ids,
+        torch.tensor([1], dtype=torch.int32, device="cuda"),
+        torch.tensor([0, 0], dtype=torch.int32, device="cuda"),
+        torch.tensor([max_model_len, max_model_len], dtype=torch.int32, device="cuda"),
+        all_token_ids[1, -4:].to(torch.int64).unsqueeze(0),
+    )
+
+    repeated = repeated_context_mask(*inputs)
+
+    assert not repeated.item()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_gpu_sampler_uses_fused_gumbel_for_repeated_contexts():
+    device = torch.device("cuda")
+    logits = torch.tensor([[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0]], device=device)
+    all_token_ids = torch.tensor(
+        [[101, 102, 1, 2, 1, 2], [101, 102, 3, 4, 5, 6]],
+        dtype=torch.int32,
+        device=device,
+    )
+    request_indices = torch.tensor([0, 1], dtype=torch.int64, device=device)
+    request_indices_np = np.array([0, 1])
+    temperatures = torch.ones(2, device=device)
+    seeds = torch.tensor([11, 22], dtype=torch.int64, device=device)
+    positions = torch.tensor([6, 6], dtype=torch.int64, device=device)
+
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = GumbelWatermarker(key=42, context_width=1)
+    sampler.deduplicate_contexts = "single_turn"
+    sampler.deduplicate_contexts_max_history = None
+    sampler.watermarking = SimpleNamespace(
+        np=np.ones(2, dtype=bool), gpu=torch.ones(2, dtype=torch.bool, device=device)
+    )
+    sampler.sampling_states = SimpleNamespace(
+        temperature=SimpleNamespace(np=np.ones(2), gpu=temperatures),
+        seeds=SimpleNamespace(gpu=seeds),
+    )
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=all_token_ids),
+        prompt_len=SimpleNamespace(
+            gpu=torch.tensor([2, 2], dtype=torch.int32, device=device)
+        ),
+        total_len=SimpleNamespace(
+            gpu=torch.tensor([6, 6], dtype=torch.int32, device=device)
+        ),
+    )
+    sampler.use_fp64_gumbel = False
+
+    contexts = sampler._get_contexts(request_indices)
+    repeated = sampler._get_repeated_contexts(request_indices, contexts)
+    expected = philox_gumbel_sample(
+        logits,
+        contexts,
+        42,
+        skip_mask=repeated,
+        expanded_idx_mapping=request_indices,
+        temperatures=temperatures,
+        seeds=seeds,
+        positions=positions,
+    )
+
+    actual, output_logits = sampler._sample_random(
+        logits,
+        request_indices,
+        request_indices_np,
+        positions,
+        None,
+        None,
+        False,
+    )
+
+    assert torch.equal(repeated, torch.tensor([True, False], device=device))
+    assert torch.equal(actual, expected)
+    assert output_logits is logits
+
+
 def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch):
     class StubWatermarker:
         context_width = 1
 
-        def sample(self, logits, contexts, random_sample):
+        def sample(self, logits, contexts, random_sampler=None, skip_mask=None):
             raise AssertionError("watermarker should not run for greedy requests")
 
     sampler = object.__new__(GPUWatermarkSampler)
@@ -155,3 +872,157 @@ def sample(self, logits, contexts, random_sample):
     )
 
     assert actual is expected
+
+
+def test_gpu_sampler_builds_speculative_contexts_from_drafts():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = SimpleNamespace(context_width=2)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 101, 10, 11, 0, 0]])),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([2])),
+        total_len=SimpleNamespace(gpu=torch.tensor([4])),
+    )
+
+    contexts = sampler._get_contexts(
+        torch.tensor([0, 0, 0]),
+        torch.tensor([0, 1, 2]),
+        torch.tensor([11, 20, 21]),
+    )
+
+    assert torch.equal(contexts, torch.tensor([[10, 11], [11, 20], [20, 21]]))
+
+
+def test_gpu_sampler_builds_chunked_multi_request_speculative_contexts():
+    sampler = object.__new__(GPUWatermarkSampler)
+    sampler.watermarker = SimpleNamespace(context_width=2)
+    sampler.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(
+            gpu=torch.tensor(
+                [
+                    [100, 101, 10, 11, 0, 0],
+                    [200, 201, 30, 31, 0, 0],
+                    [300, 301, 50, 51, 0, 0],
+                ]
+            )
+        ),
+        prompt_len=SimpleNamespace(gpu=torch.tensor([2, 2, 2])),
+        total_len=SimpleNamespace(gpu=torch.tensor([4, 4, 4])),
+    )
+
+    contexts = sampler._get_contexts(
+        torch.tensor([2, 2, 1, 1]),
+        torch.tensor([0, 1, 0, 1]),
+        torch.tensor([51, 60, 31, 40]),
+    )
+
+    assert torch.equal(
+        contexts,
+        torch.tensor([[50, 51], [51, 60], [30, 31], [31, 40]]),
+    )
+
+
+def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch):
+    class StubSpeculator(DraftModelSpeculator):
+        def capture(self): ...
+
+        def init_cudagraph_manager(self, cudagraph_mode): ...
+
+        def load_draft_model(self, target_model, target_attn_layer_names): ...
+
+        def propose(self, *args, **kwargs): ...
+
+    class StubModel:
+        @staticmethod
+        def compute_logits(hidden_states):
+            return torch.zeros(hidden_states.shape[0], 8)
+
+    class StubWatermarker:
+        @staticmethod
+        def sample(logits, contexts, random_sampler=None):
+            assert random_sampler is None
+            return WatermarkSample(torch.tensor([7, 7]), logits)
+
+    speculator = object.__new__(StubSpeculator)
+    speculator.model = StubModel()
+    speculator.use_fp64_gumbel = False
+    draft_watermarker = object.__new__(DraftWatermarker)
+    draft_watermarker.watermarker = StubWatermarker()
+    draft_watermarker.contexts = torch.tensor([[1, 2], [3, 4]])
+    draft_watermarker.enabled = torch.tensor([True, False])
+    speculator.draft_watermarker = draft_watermarker
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.spec_decode.speculator.gumbel_sample",
+        lambda *args, **kwargs: torch.tensor([3, 4]),
+    )
+
+    sampled = speculator.sample_draft(
+        hidden_states=torch.zeros(2, 4),
+        sample_src_positions=torch.zeros(2, dtype=torch.int64),
+        idx_mapping=torch.tensor([0, 1]),
+        temperature=torch.ones(2),
+        seeds=torch.zeros(2, dtype=torch.int64),
+        draft_step=torch.tensor(0),
+        draft_logits=torch.zeros(2, 1, 8),
+    )
+
+    assert torch.equal(sampled, torch.tensor([7, 4]))
+    assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]]))
+
+
+def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):
+    speculator = object.__new__(DSparkSpeculator)
+    speculator.draft_logits = torch.zeros(2, 1, 8)
+    speculator._d2t_scatter_index = torch.tensor([1, 5])
+    speculator._draft_scatter_buf = torch.full((2, 8), float("-inf"))
+    speculator.temperature = torch.ones(2)
+    speculator.seeds = torch.zeros(2, dtype=torch.int64)
+    speculator._step_cols = torch.tensor([0])
+    speculator.use_fp64_gumbel = False
+    watermark_logits: list[torch.Tensor] = []
+
+    def sample(logits, sampled, idx_map, temperature):
+        watermark_logits.append(logits.clone())
+        return sampled + 1
+
+    speculator.draft_watermarker = SimpleNamespace(sample=sample)
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
+        lambda *args, **kwargs: torch.tensor([3, 4]),
+    )
+
+    sampled = speculator._sample_logits(
+        torch.tensor([[10.0, 20.0], [30.0, 40.0]]),
+        torch.tensor([0, 1]),
+        torch.tensor([1, 1]),
+        0,
+    )
+
+    assert torch.equal(sampled, torch.tensor([4, 5]))
+    assert torch.equal(
+        watermark_logits[0][:, [1, 5]], torch.tensor([[10, 20], [30, 40]])
+    )
+    assert torch.isneginf(watermark_logits[0][:, [0, 2, 3, 4, 6, 7]]).all()
+
+
+def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch):
+    speculator = object.__new__(DSparkSpeculator)
+    speculator.draft_logits = torch.zeros(2, 1, 8)
+    speculator._d2t_scatter_index = None
+    speculator.temperature = torch.ones(2)
+    speculator.seeds = torch.zeros(2, dtype=torch.int64)
+    speculator._step_cols = torch.tensor([0])
+    speculator.use_fp64_gumbel = False
+    speculator.draft_watermarker = None
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
+        lambda *args, **kwargs: torch.tensor([3, 4]),
+    )
+
+    sampled = speculator._sample_logits(
+        torch.zeros(2, 8),
+        torch.tensor([0, 1]),
+        torch.tensor([1, 1]),
+        0,
+    )
+
+    assert torch.equal(sampled, torch.tensor([3, 4]))
diff --git a/tests/watermarking/test_watermarking_e2e.py b/tests/watermarking/test_watermarking_e2e.py
new file mode 100644
index 000000000000..c73d48f6b6b7
--- /dev/null
+++ b/tests/watermarking/test_watermarking_e2e.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.platforms import current_platform
+from vllm.v1.watermarking import GumbelWatermarkDetector
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
+)
+def test_llm_generated_sequence_is_watermarked(vllm_runner):
+    """Engine-level watermarking must produce tokens the detector recognizes."""
+    watermark_config = {"key": 42, "context_width": 4}
+    detector = GumbelWatermarkDetector(**watermark_config)
+
+    runner = vllm_runner(
+        "facebook/opt-125m",
+        watermark_config=watermark_config,
+        seed=0,
+        enforce_eager=True,
+        max_model_len=512,
+        gpu_memory_utilization=0.2,
+    )
+    request = "Tell me a story about an explorer who discovers a mysterious island."
+    params_use_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=True)
+    params_no_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=False)
+
+    with runner:
+        output_use_wm = runner.llm.generate(request, params_use_wm)
+        output_no_wm = runner.llm.generate(request, params_no_wm)
+
+    result_use_wm = detector.detect(list(output_use_wm[0].outputs[0].token_ids))
+    result_no_wm = detector.detect(list(output_no_wm[0].outputs[0].token_ids))
+
+    assert result_use_wm.is_watermarked, result_use_wm
+    recorded_p_value = 9.83e-20
+    assert result_use_wm.p_value < recorded_p_value * 10  # tolerance
+
+    assert not result_no_wm.is_watermarked, result_no_wm
diff --git a/tests/weight_loading/models-amd.txt b/tests/weight_loading/models-amd.txt
deleted file mode 100644
index a6a85736cb16..000000000000
--- a/tests/weight_loading/models-amd.txt
+++ /dev/null
@@ -1,18 +0,0 @@
-# quantization, model, revision[, minimum capability[, safetensors load strategy]]
-fp8, amd/Meta-Llama-3.1-8B-Instruct-FP8-KV, main
-None, amd/Llama-3.2-1B-Instruct-FP8-KV, main
-fp8, amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV, main, 80, prefetch
-gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main
-gptq, TheBloke/Llama-2-7B-GPTQ, main
-gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True
-compressed-tensors, nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8-channel-a8-tensor, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w4a16-group128-v2, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8a16-per-channel, main
-compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main
-compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main
-compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main
-awq, casperhansen/mixtral-instruct-awq, main
-fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main
diff --git a/tests/weight_loading/models-large-amd.txt b/tests/weight_loading/models-large-amd.txt
deleted file mode 100644
index 41e8ee5200af..000000000000
--- a/tests/weight_loading/models-large-amd.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-# quantization, model, revision[, minimum capability[, safetensors load strategy]]
-fp8, amd/Meta-Llama-3.1-70B-Instruct-FP8-KV, main, 80, prefetch
-None, microsoft/phi-4, main, 80, prefetch
-fp8, amd/Mixtral-8x22B-Instruct-v0.1-FP8-KV, main, 80, prefetch
diff --git a/tests/weight_loading/models-large.txt b/tests/weight_loading/models-large.txt
deleted file mode 100644
index 8b3913cc2d50..000000000000
--- a/tests/weight_loading/models-large.txt
+++ /dev/null
@@ -1,6 +0,0 @@
-compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W4A16-quantized, main
-compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W4A16-channel-quantized, main
-compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W8A16-quantized, main
-gptq, TheBloke/Mixtral-8x7B-v0.1-GPTQ, main
-awq_marlin, casperhansen/deepseek-coder-v2-instruct-awq, main
-compressed-tensors, RedHatAI/Llama-4-Scout-17B-16E-Instruct-quantized.w4a16, main
diff --git a/tests/weight_loading/models.txt b/tests/weight_loading/models.txt
deleted file mode 100644
index 87b965087092..000000000000
--- a/tests/weight_loading/models.txt
+++ /dev/null
@@ -1,18 +0,0 @@
-gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main
-gptq, TheBloke/Llama-2-7B-GPTQ, main
-gptq, LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit, main
-gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True
-gptq, TechxGenus/gemma-2b-it-GPTQ, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8-channel-a8-tensor, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w4a16-group128-v2, main
-compressed-tensors, nm-testing/tinyllama-oneshot-w8a16-per-channel, main
-compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main
-compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main
-compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main
-#compressed-tensors, mgoin/DeepSeek-Coder-V2-Lite-Instruct-FP8, main
-awq, casperhansen/mixtral-instruct-awq, main
-awq_marlin, casperhansen/mixtral-instruct-awq, main
-fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main
diff --git a/tests/weight_loading/run_model_weight_loading_test.sh b/tests/weight_loading/run_model_weight_loading_test.sh
deleted file mode 100755
index 598f1620c7da..000000000000
--- a/tests/weight_loading/run_model_weight_loading_test.sh
+++ /dev/null
@@ -1,58 +0,0 @@
-#!/bin/bash
-SUCCESS=0
-
-while getopts "c:" OPT; do
-  case ${OPT} in
-    c )
-        CONFIG="$OPTARG"
-        ;;
-    \? )
-        usage
-        exit 1
-        ;;
-  esac
-done
-
-
-IFS=$'\n' read -d '' -r -a MODEL_CONFIGS < "$CONFIG"
-
-for MODEL_CONFIG in "${MODEL_CONFIGS[@]}"
-do
-    if [[ $MODEL_CONFIG == \#* ]]; then
-        echo "=== SKIPPING MODEL: $MODEL_CONFIG ==="
-        continue
-    fi
-
-    LOCAL_SUCCESS=0
-    IFS=', ' read -r -a array <<< "$MODEL_CONFIG"
-
-    echo "=== RUNNING MODEL: $MODEL_CONFIG ==="
-
-    unset WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY
-    export QUANTIZATION=${array[0]}
-    export MODEL_NAME=${array[1]}
-    export REVISION=${array[2]}
-    # If array length is larger than 3, then MIN_CAPABILITY is provided
-    if [ ${#array[@]} -gt 3 ]; then
-        export MIN_CAPABILITY=${array[3]}
-    fi
-    if [ ${#array[@]} -gt 4 ]; then
-        export WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY=${array[4]}
-    fi
-    pytest -s weight_loading/test_weight_loading.py || LOCAL_SUCCESS=$?
-
-    if [[ $LOCAL_SUCCESS == 0 ]]; then
-        echo "=== PASSED MODEL: ${MODEL_CONFIG} ==="
-    else
-        echo "=== FAILED MODEL: ${MODEL_CONFIG} ==="
-    fi
-
-    SUCCESS=$((SUCCESS + LOCAL_SUCCESS))
-
-done
-
-if [ "${SUCCESS}" -eq "0" ]; then
-    exit 0
-else
-    exit 1
-fi
diff --git a/tests/weight_loading/test_weight_loading.py b/tests/weight_loading/test_weight_loading.py
deleted file mode 100644
index e739164a5f08..000000000000
--- a/tests/weight_loading/test_weight_loading.py
+++ /dev/null
@@ -1,52 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import os
-
-import pytest
-import torch
-
-from vllm.platforms import current_platform
-from vllm.utils.mem_constants import GiB_bytes
-
-MAX_MODEL_LEN = 1024
-MODEL_NAME = os.environ.get(
-    "MODEL_NAME", "robertgshaw2/zephyr-7b-beta-channelwise-gptq"
-)
-REVISION = os.environ.get("REVISION", "main")
-QUANTIZATION = os.environ.get("QUANTIZATION", "gptq_marlin")
-MIN_CAPABILITY = os.environ.get("MIN_CAPABILITY", "80")
-SAFETENSORS_LOAD_STRATEGY = os.environ.get(
-    "WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY"
-)
-
-
-@pytest.mark.skipif(
-    MODEL_NAME == "casperhansen/deepseek-coder-v2-instruct-awq", reason="OOM in the CI"
-)
-@pytest.mark.skipif(
-    not current_platform.has_device_capability(int(MIN_CAPABILITY)),
-    reason="Current system does not have minimum capability.",
-)
-def test_weight_loading(vllm_runner):
-    """
-    Test parameter weight loading with tp>1.
-    """
-
-    # MoE models need fp16.
-    NEEDS_FP16 = QUANTIZATION == "gptq"
-    with vllm_runner(
-        model_name=MODEL_NAME,
-        revision=REVISION,
-        dtype=torch.half if NEEDS_FP16 else "auto",
-        quantization=None if QUANTIZATION == "None" else QUANTIZATION,
-        max_model_len=MAX_MODEL_LEN,
-        tensor_parallel_size=2,
-        safetensors_load_strategy=SAFETENSORS_LOAD_STRATEGY,
-        # Generating 20 tokens needs a tiny KV cache, while sizing and
-        # allocating a full-size one dominates this test on large devices.
-        kv_cache_memory_bytes=2 * GiB_bytes,
-    ) as model:
-        output = model.generate_greedy("Hello world!", max_tokens=20)
-        print(output)
-        assert output
diff --git a/tools/__init__.py b/tools/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py
index 67a527405e2d..f4438c67552f 100644
--- a/tools/build_deepgemm_C.py
+++ b/tools/build_deepgemm_C.py
@@ -50,7 +50,7 @@
     str(src / "deep_gemm/include"),
     str(src / "third-party/cutlass/include"),
     str(src / "third-party/cutlass/tools/util/include"),
-    str(src / "third-party/fmt/include"),
+    str(src / "third-party/deep_jit/include"),
     *cpp_extension.include_paths(device_type="cuda"),
 ]
 
diff --git a/tools/build_rust.py b/tools/build_rust.py
index a000ec8169fd..320cf75e3db4 100644
--- a/tools/build_rust.py
+++ b/tools/build_rust.py
@@ -16,13 +16,28 @@
 ROOT_DIR = Path(__file__).resolve().parents[1]
 VLLM_RS_BUILD_VERSION = "VLLM_RS_BUILD_VERSION"
 
+# Select vLLM release tags, excluding crate tags such as "proto-v0.3.0".
+VLLM_GIT_DESCRIBE_COMMAND = [
+    "git",
+    "describe",
+    "--dirty",
+    "--tags",
+    "--long",
+    "--abbrev=40",
+    "--match",
+    "v[0-9]*",
+]
+
 
 def prepare_build_environment() -> str | None:
     """Set the device-independent vLLM source version for Rust artifacts."""
     version = os.getenv(VLLM_RS_BUILD_VERSION) or None
     if version is None:
         try:
-            version = get_version(root=ROOT_DIR)
+            version = get_version(
+                root=ROOT_DIR,
+                git_describe_command=VLLM_GIT_DESCRIBE_COMMAND,
+            )
         except LookupError:
             return None
 
diff --git a/tools/ci/kernel_symbol_map.py b/tools/ci/kernel_symbol_map.py
new file mode 100755
index 000000000000..2a1e8964d055
--- /dev/null
+++ b/tools/ci/kernel_symbol_map.py
@@ -0,0 +1,370 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Which source file produced each GPU kernel symbol in this build.
+
+Runs inside the csrc-build image stage, after the CMake build, and reads the
+objects it left behind. For every compiled object it records the source file,
+the CMake target, the headers the compiler reported it depended on (from
+ninja's dependency log), and, for device objects, the kernel entry symbols
+`cuobjdump` finds in the embedded fatbin. Those mangled names are exactly
+what a CUPTI kernel recorder sees at run time, so CI can join "this step
+launched kernel K" with "K came from csrc/foo.cu" and select tests for a
+kernel change by evidence instead of by image membership.
+
+    kernel_symbol_map.py --build-root build --source-root . \\
+        --out dist/kernel_symbol_map.json.gz
+
+Never fails the image build. If there is nothing to read (a precompiled-wheel
+build, a missing cuobjdump) it writes a map with an empty object list and a
+`reason`, which the consumer treats as "no evidence", not as "no kernels".
+
+Output (gzipped JSON):
+
+    {
+      "version": 1,
+      "commit": "...",              # from --commit or the environment
+      "cuda": "13.0",               # nvcc release when found
+      "generated_at": "2026-09-19T...",
+      "build_dirs": ["build/temp.linux-x86_64-cpython-312"],
+      "incomplete": false,          # true when any object's extraction failed;
+                                    # objects is then [] and reason is set
+      "errors": [],                 # [{"object", "source", "error"}, ...]
+      "objects": [
+        {
+          "source": "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu",
+          "target": "_C",
+          "object": "CMakeFiles/_C.dir/csrc/.../fused_qknorm_rope_kernel.cu.o",
+          "device": true,
+          "symbols": ["_Z21fusedQKNormRopeKernelI...", ...],
+          "deps": ["csrc/libtorch_stable/fused_qknorm_rope_kernel.cu",
+                   "csrc/cuda_compat.h", ...]
+          # "error": "..."          # present when cuobjdump failed on it
+        },
+        ...
+      ],
+      "stats": {...},
+      "reason": "..."               # only when the map is empty
+    }
+
+The consumer's contract is deliberately simple: an empty `objects` list with
+a `reason` means "no evidence for any file", fall back to the static rule.
+A partial map would need every consumer to know which files the unreadable
+objects touched, and one that forgot would drop tests that a kernel change
+can reach. So the producer never publishes a partial map: if any object is
+still unreadable after a retry, `objects` is emptied, `incomplete` is set
+and `errors` names the objects so the failure can be diagnosed.
+
+Paths are relative to --source-root when the file lives inside it, otherwise
+absolute (FetchContent sources under the build directory, CUDA headers).
+"""
+
+from __future__ import annotations
+
+import argparse
+import concurrent.futures
+import datetime as dt
+import gzip
+import json
+import os
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import regex as re
+
+DEVICE_SUFFIXES = (".cu.o", ".hip.o")
+SOURCE_SUFFIXES = (".cu", ".hip", ".cpp", ".cc", ".cxx", ".c")
+DEPS_HEADER = re.compile(r"^(\S+): #deps (\d+), deps mtime \d+ \((VALID|STALE)\)$")
+# FetchContent sub-builds nest their own CMakeFiles/ under _deps/<name>-build/.
+OBJ_DIR = re.compile(r"(?:^|/)CMakeFiles/([^/]+)\.dir/(.+)\.o$")
+
+
+def log(msg: str) -> None:
+    print(f"kernel_symbol_map: {msg}", flush=True)
+
+
+def find_tool(name: str, *candidates: str) -> str | None:
+    for c in candidates:
+        if c and Path(c).is_file():
+            return c
+    return shutil.which(name)
+
+
+def cuda_release(nvcc: str | None) -> str:
+    if not nvcc:
+        return os.environ.get("CUDA_VERSION", "")
+    try:
+        out = subprocess.run(
+            [nvcc, "--version"], capture_output=True, text=True, timeout=30
+        ).stdout
+        m = re.search(r"release (\d+\.\d+)", out)
+        return m.group(1) if m else os.environ.get("CUDA_VERSION", "")
+    except (OSError, subprocess.SubprocessError):
+        return os.environ.get("CUDA_VERSION", "")
+
+
+def find_build_dirs(build_root: Path) -> list[Path]:
+    """Directories holding a top-level build.ninja (one per CMake build tree)."""
+    return sorted(
+        {
+            p.parent
+            for p in build_root.rglob("build.ninja")
+            if (p.parent / ".ninja_deps").exists() or (p.parent / "CMakeFiles").is_dir()
+        }
+    )
+
+
+def ninja_deps(ninja: str | None, build_dir: Path) -> dict[str, list[str]]:
+    """Object (relative to build_dir) -> paths the compiler reported it read."""
+    if not ninja or not (build_dir / ".ninja_deps").exists():
+        return {}
+    try:
+        out = subprocess.run(
+            [ninja, "-C", str(build_dir), "-t", "deps"],
+            capture_output=True,
+            text=True,
+            timeout=600,
+        ).stdout
+    except (OSError, subprocess.SubprocessError) as e:
+        log(f"ninja -t deps failed in {build_dir}: {e}")
+        return {}
+    deps: dict[str, list[str]] = {}
+    cur: list[str] | None = None
+    for line in out.splitlines():
+        m = DEPS_HEADER.match(line)
+        if m:
+            cur = deps.setdefault(m.group(1), [])
+            continue
+        if cur is not None and line.startswith("    "):
+            cur.append(line.strip())
+        elif not line.strip():
+            cur = None
+    return deps
+
+
+def parse_symbols(text: str) -> list[str]:
+    """Kernel entry points from `cuobjdump -symbols` output.
+
+    Lines look like `STT_FUNC  STB_GLOBAL  STO_ENTRY  _Z21fusedQKNormRope...`.
+    Entry points are what the driver launches, so they are what a recorder
+    sees. If a toolchain prints no STO_ENTRY marker, fall back to every global
+    function symbol; over-attribution is the safe direction.
+    """
+    entries: set[str] = set()
+    globals_: set[str] = set()
+    for line in text.splitlines():
+        parts = line.split()
+        if len(parts) < 4 or parts[0] != "STT_FUNC":
+            continue
+        if parts[2] == "STO_ENTRY":
+            entries.add(parts[3])
+        elif parts[1] == "STB_GLOBAL":
+            globals_.add(parts[3])
+    return sorted(entries or globals_)
+
+
+def device_symbols(cuobjdump: str, obj: Path) -> tuple[list[str], str | None]:
+    """(symbols, error). An error means the object is unknown, not kernel-free."""
+    syms, err = _cuobjdump_symbols(cuobjdump, obj)
+    if err:  # transient failures happen on large fatbins; one retry is cheap
+        syms, err = _cuobjdump_symbols(cuobjdump, obj)
+    return syms, err
+
+
+def _cuobjdump_symbols(cuobjdump: str, obj: Path) -> tuple[list[str], str | None]:
+    try:
+        r = subprocess.run(
+            [cuobjdump, "-symbols", str(obj)],
+            capture_output=True,
+            text=True,
+            timeout=900,
+        )
+    except subprocess.TimeoutExpired:
+        return [], "timeout"
+    except OSError as e:
+        return [], str(e)
+    if r.returncode != 0:
+        msg = (r.stderr or r.stdout).strip().splitlines()
+        return [], (msg[-1] if msg else f"exit {r.returncode}")
+    return parse_symbols(r.stdout), None
+
+
+def relpath(p: str | Path, source_root: Path) -> str:
+    path = Path(p)
+    try:
+        return path.resolve().relative_to(source_root).as_posix()
+    except ValueError:
+        return path.as_posix()
+
+
+def source_of(
+    obj_rel: str, deps: list[str], build_dir: Path, source_root: Path
+) -> tuple[str, str]:
+    """(source path, cmake target) for an object.
+
+    The dependency log lists the source first, which is the reliable answer.
+    Failing that, CMake's object layout `CMakeFiles/<target>.dir/<source>.o`
+    names it, with `__/` segments standing for `..`.
+    """
+    target = ""
+    m = OBJ_DIR.search(obj_rel)  # search: sub-builds nest CMakeFiles/
+    if m:
+        target = m.group(1)
+    for d in deps:
+        if d.endswith(SOURCE_SUFFIXES):
+            return relpath(
+                d if os.path.isabs(d) else build_dir / d, source_root
+            ), target
+    if m:
+        guess = m.group(2).replace("__/", "../")
+        if guess.endswith(SOURCE_SUFFIXES):
+            for base in (source_root, build_dir):
+                cand = base / guess
+                if cand.exists():
+                    return relpath(cand, source_root), target
+            return guess, target
+    return obj_rel, target
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    ap.add_argument("--build-root", type=Path, default=Path("build"))
+    ap.add_argument("--source-root", type=Path, default=Path("."))
+    ap.add_argument("--out", type=Path, required=True)
+    ap.add_argument("--jobs", type=int, default=min(8, os.cpu_count() or 2))
+    ap.add_argument("--cuobjdump", default=None)
+    ap.add_argument("--commit", default="", help="commit the build is of")
+    a = ap.parse_args()
+    t0 = time.time()
+    source_root = a.source_root.resolve()
+
+    cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda")
+    cuobjdump = a.cuobjdump or find_tool(
+        "cuobjdump", f"{cuda_home}/bin/cuobjdump", "/usr/local/cuda/bin/cuobjdump"
+    )
+    nvcc = find_tool("nvcc", f"{cuda_home}/bin/nvcc", "/usr/local/cuda/bin/nvcc")
+    ninja = shutil.which("ninja")
+
+    result: dict = {
+        "version": 1,
+        "commit": a.commit
+        or os.environ.get("BUILDKITE_COMMIT")
+        or os.environ.get("VLLM_BUILD_COMMIT")
+        or "",
+        "cuda": cuda_release(nvcc),
+        "generated_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"),
+        "build_dirs": [],
+        "incomplete": False,
+        "errors": [],
+        "objects": [],
+        "stats": {},
+    }
+
+    build_dirs = find_build_dirs(a.build_root) if a.build_root.exists() else []
+    result["build_dirs"] = [relpath(b, source_root) for b in build_dirs]
+    if not build_dirs:
+        result["reason"] = (
+            f"no CMake build tree under {a.build_root} (precompiled wheel build?)"
+        )
+    elif not cuobjdump:
+        result["reason"] = "cuobjdump not found"
+    log(
+        f"build dirs: {result['build_dirs'] or 'none'}; "
+        f"cuobjdump={cuobjdump}; ninja={ninja}"
+    )
+
+    jobs: list[tuple[Path, str, Path]] = []  # (build_dir, obj_rel, obj_path)
+    all_deps: dict[Path, dict[str, list[str]]] = {}
+    if "reason" not in result:
+        for bd in build_dirs:
+            all_deps[bd] = ninja_deps(ninja, bd)
+            for obj in bd.rglob("*.o"):
+                rel = obj.relative_to(bd).as_posix()
+                # CMake's compiler-identification probes are objects too.
+                if "CompilerId" in rel:
+                    continue
+                if "/CMakeFiles/" in f"/{rel}" or rel.startswith("CMakeFiles/"):
+                    jobs.append((bd, rel, obj))
+        n_deps = sum(len(d) for d in all_deps.values())
+        log(f"{len(jobs)} objects, {n_deps} dependency records")
+
+    def work(job):
+        bd, rel, obj = job
+        deps = all_deps.get(bd, {}).get(rel, [])
+        source, target = source_of(rel, deps, bd, source_root)
+        device = rel.endswith(DEVICE_SUFFIXES)
+        syms, err = device_symbols(cuobjdump, obj) if device else ([], None)
+        rel_deps = sorted(
+            {relpath(d if os.path.isabs(d) else bd / d, source_root) for d in deps}
+        )
+        entry = {
+            "source": source,
+            "target": target,
+            "object": rel,
+            "device": device,
+            "symbols": syms,
+            "deps": rel_deps,
+        }
+        if err:
+            entry["error"] = err
+        return entry
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, a.jobs)) as ex:
+        result["objects"] = sorted(
+            ex.map(work, jobs), key=lambda e: (e["source"], e["target"])
+        )
+
+    # An object cuobjdump could not read is unknown, not kernel-free. Say so
+    # at the top so a consumer never has to scan for it, and name the objects
+    # so it can widen the "unknown" set to every file they touch.
+    result["errors"] = [
+        {"object": e["object"], "source": e["source"], "error": e["error"]}
+        for e in result["objects"]
+        if "error" in e
+    ]
+    result["incomplete"] = bool(result["errors"])
+    for e in result["errors"][:5]:
+        log(f"cuobjdump failed on {e['object']}: {e['error']}")
+    if result["incomplete"]:
+        # Publish nothing rather than something partial: see the docstring.
+        result["objects"] = []
+        result["reason"] = (
+            f"{len(result['errors'])} objects unreadable by cuobjdump, e.g. "
+            f"{result['errors'][0]['object']}: {result['errors'][0]['error']}"
+        )
+
+    result["stats"] = {
+        "objects": len(jobs),
+        "device_objects": sum(1 for e in result["objects"] if e["device"]),
+        "objects_with_symbols": sum(1 for e in result["objects"] if e["symbols"]),
+        "symbols": sum(len(e["symbols"]) for e in result["objects"]),
+        "cuobjdump_errors": len(result["errors"]),
+        "seconds": round(time.time() - t0, 1),
+    }
+
+    a.out.parent.mkdir(parents=True, exist_ok=True)
+    with gzip.open(a.out, "wt", encoding="utf-8") as f:
+        json.dump(result, f, separators=(",", ":"))
+    note = f"; reason: {result['reason']}" if "reason" in result else ""
+    log(f"wrote {a.out} ({a.out.stat().st_size // 1024} KiB): {result['stats']}{note}")
+
+    # One line a human can check in the build log.
+    probe = [
+        e
+        for e in result["objects"]
+        if any("fusedQKNormRopeKernel" in s for s in e["symbols"])
+    ]
+    if probe:
+        log(f"probe: fusedQKNormRopeKernel -> {sorted({e['source'] for e in probe})}")
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except Exception as e:  # never fail the image build
+        log(f"unexpected error, continuing without a map: {e!r}")
+        sys.exit(0)
diff --git a/tools/generate_cmake_presets.py b/tools/generate_cmake_presets.py
index 6bc8443c4473..76d8341843df 100644
--- a/tools/generate_cmake_presets.py
+++ b/tools/generate_cmake_presets.py
@@ -28,7 +28,6 @@ def get_cpu_cores():
 
 def generate_presets(output_path="CMakeUserPresets.json", force_overwrite=False):
     """Generates the CMakeUserPresets.json file."""
-
     print("Attempting to detect your system configuration...")
 
     # Detect NVCC
diff --git a/tools/generate_versions_json.py b/tools/generate_versions_json.py
index f5d2893a996b..e1611457c913 100755
--- a/tools/generate_versions_json.py
+++ b/tools/generate_versions_json.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Generate docker/versions.json from Dockerfile ARG defaults.
+"""Generate docker/versions.json from Dockerfile ARG defaults.
 
 This script parses the Dockerfile and extracts ARG defaults to create
 a bake-native versions.json file that can be used directly with:
diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh
index f6a64d8796c1..949aa28da533 100755
--- a/tools/install_deepgemm.sh
+++ b/tools/install_deepgemm.sh
@@ -6,9 +6,11 @@ set -e
 
 # Default values
 # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake
-DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
-# NOTE: This is currently targeting the nv_dev branch tip due to sm120 support
-DEEPGEMM_GIT_REF="8b1392b978f5a03c828dd1711090d7fb50958b8a"
+DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git"
+# NOTE: This targets the vLLM fork's dev branch, which carries the sm120
+# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE
+# activations, plus the CUDA 12.x layout header fix from vllm-project/DeepGEMM#12.
+DEEPGEMM_GIT_REF="e1f418c2a4f20818221f6b0e578b4c2f634d4c3f"
 WHEEL_DIR=""
 
 # Parse command line arguments
diff --git a/tools/install_nixl_from_source_ubuntu.py b/tools/install_nixl_from_source_ubuntu.py
index b8a55c615426..1eb73bd0cbcf 100644
--- a/tools/install_nixl_from_source_ubuntu.py
+++ b/tools/install_nixl_from_source_ubuntu.py
@@ -21,7 +21,7 @@
 
 # --- Helper Functions ---
 def get_latest_nixl_version():
-    """Helper function to get latest release version of NIXL"""
+    """Helper function to get latest release version of NIXL."""
     try:
         nixl_release_url = "https://api.github.com/repos/ai-dynamo/nixl/releases/latest"
         with urllib.request.urlopen(nixl_release_url) as response:
@@ -104,7 +104,6 @@ def install_system_dependencies():
 
 def build_and_install_prerequisites(args):
     """Builds UCX and NIXL from source, creating a self-contained wheel."""
-
     if not args.force_reinstall and is_pip_package_installed("nixl"):
         print("--> NIXL is already installed. Nothing to do.", flush=True)
         return
diff --git a/tools/install_snapshot_runtime.sh b/tools/install_snapshot_runtime.sh
new file mode 100755
index 000000000000..40d6ad944fee
--- /dev/null
+++ b/tools/install_snapshot_runtime.sh
@@ -0,0 +1,109 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+set -euxo pipefail
+
+# Surface the isolated snapshot-runtime install cost in the CI build log. The
+# trap captures the status first: the date substitution below would reset it.
+_snapshot_install_started=$(date +%s)
+code=0
+trap 'code=$?; echo "[TIMING] snapshot runtime install: $(( $(date +%s) - _snapshot_install_started )) s (exit $code)"' EXIT
+
+readonly CRIU_VERSION="4.2.1"
+readonly CRIU_SOURCE_SHA256="feffdf4638125ebb12d2434754f80a1d7bbba85a3e6bee98c216f88fb99a5d96"
+readonly DYNAMO_CHECKPOINT_HELPER_COMMIT="fdf25efad60f696a73393caf341c527a93d84190"
+readonly DYNAMO_CHECKPOINT_HELPER_SHA256="7b8ea21baecb1da729202c32e0fecbbbafb789d65e5c1c0649ea253ad09b8f82"
+readonly DYNAMO_LICENSE_SHA256="aeeca3d74a13b91d17ab9211c407b8d94a77bf8d9e5750c467818d3066b8adcf"
+
+# Keep the compiler and NL headers required by runtime JIT and RDMA.
+snapshot_shared_build_deps=(
+    build-essential
+    libnl-3-dev
+    libnl-route-3-dev
+)
+snapshot_build_deps=(
+    pkg-config
+    libbsd-dev
+    libcap-dev
+    libnet1-dev
+    libprotobuf-dev
+    libprotobuf-c-dev
+    protobuf-c-compiler
+    protobuf-compiler
+    python3-protobuf
+    libgnutls28-dev
+    libnftables-dev
+    uuid-dev
+)
+snapshot_runtime_deps=(
+    libbsd0
+    libcap2
+    libnet1
+    libnl-3-200
+    libnl-route-3-200
+    libprotobuf-c1
+    libgnutls30t64
+    libnftables1
+    iproute2
+    iptables
+    procps
+    uuid-runtime
+    util-linux
+)
+
+apt-get update -y
+apt-get install -y --no-install-recommends \
+    "${snapshot_shared_build_deps[@]}" \
+    "${snapshot_build_deps[@]}" \
+    "${snapshot_runtime_deps[@]}"
+
+snapshot_build_dir="$(mktemp -d)"
+curl -fsSL \
+    "https://github.com/checkpoint-restore/criu/archive/refs/tags/v${CRIU_VERSION}.tar.gz" \
+    -o "${snapshot_build_dir}/criu.tar.gz"
+echo "${CRIU_SOURCE_SHA256}  ${snapshot_build_dir}/criu.tar.gz" | sha256sum -c -
+mkdir "${snapshot_build_dir}/criu"
+tar -xzf "${snapshot_build_dir}/criu.tar.gz" \
+    --strip-components=1 -C "${snapshot_build_dir}/criu"
+make -C "${snapshot_build_dir}/criu" -j"${MAX_JOBS:-$(nproc)}" criu cuda_plugin
+install -D -m 0755 "${snapshot_build_dir}/criu/criu/criu" \
+    /usr/local/sbin/criu
+install -D -m 0755 \
+    "${snapshot_build_dir}/criu/plugins/cuda/cuda_plugin.so" \
+    /usr/local/lib/criu/cuda_plugin.so
+
+helper_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/deploy/snapshot/cmd/cuda-checkpoint-helper/main.c"
+curl -fsSL "${helper_url}" -o "${snapshot_build_dir}/cuda-checkpoint.c"
+echo "${DYNAMO_CHECKPOINT_HELPER_SHA256}  ${snapshot_build_dir}/cuda-checkpoint.c" \
+    | sha256sum -c -
+gcc -O2 -Wall -Wextra \
+    -I/usr/local/cuda/include \
+    -L/usr/local/cuda/lib64/stubs \
+    "${snapshot_build_dir}/cuda-checkpoint.c" -lcuda \
+    -o /usr/local/sbin/cuda-checkpoint
+
+license_url="https://raw.githubusercontent.com/ai-dynamo/dynamo/${DYNAMO_CHECKPOINT_HELPER_COMMIT}/LICENSE"
+curl -fsSL "${license_url}" -o "${snapshot_build_dir}/DYNAMO-LICENSE"
+echo "${DYNAMO_LICENSE_SHA256}  ${snapshot_build_dir}/DYNAMO-LICENSE" \
+    | sha256sum -c -
+
+install -D -m 0644 "${snapshot_build_dir}/criu.tar.gz" \
+    /usr/share/doc/criu/source/criu.tar.gz
+install -D -m 0644 "${snapshot_build_dir}/criu/COPYING" \
+    /usr/share/doc/criu/COPYING
+install -D -m 0644 "${snapshot_build_dir}/cuda-checkpoint.c" \
+    /usr/share/doc/cuda-checkpoint-helper/main.c
+install -D -m 0644 "${snapshot_build_dir}/DYNAMO-LICENSE" \
+    /usr/share/doc/cuda-checkpoint-helper/LICENSE
+
+criu --version
+ln -s /usr/local/cuda/lib64/stubs/libcuda.so \
+    "${snapshot_build_dir}/libcuda.so.1"
+LD_LIBRARY_PATH="${snapshot_build_dir}" \
+    cuda-checkpoint --help | grep -q -- "--action"
+test -x /usr/local/lib/criu/cuda_plugin.so
+
+apt-get purge -y --auto-remove "${snapshot_build_deps[@]}"
+rm -rf "${snapshot_build_dir}"
+rm -rf /var/lib/apt/lists/*
diff --git a/tools/pre_commit/__init__.py b/tools/pre_commit/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tools/pre_commit/check_buildkite_step_keys.py b/tools/pre_commit/check_buildkite_step_keys.py
new file mode 100644
index 000000000000..41d65b4bc6db
--- /dev/null
+++ b/tools/pre_commit/check_buildkite_step_keys.py
@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Check that every Buildkite step the CI generator reads declares a key.
+
+A step with no `key:` is published under one the generator derives from its
+label, which makes the job's address a function of its title: retitling the
+step silently renames it and breaks retries, `depends_on:` and build history.
+The generator only reads the directories a pipeline config names in
+`job_dirs:`, so this walks the same files and fails on any step without a key.
+
+Usage:
+    python tools/pre_commit/check_buildkite_step_keys.py
+"""
+
+import sys
+from pathlib import Path
+
+import yaml
+
+BUILDKITE = Path(".buildkite")
+CONFIG_GLOB = "ci_config*.yaml"
+JOB_GLOB = "*.yaml"
+
+
+def job_dirs() -> list[Path]:
+    """Every directory the pipeline generator reads, per the pipeline configs."""
+    configs = sorted(BUILDKITE.glob(CONFIG_GLOB))
+    if not configs:
+        raise SystemExit(f"no pipeline config matched {BUILDKITE / CONFIG_GLOB}")
+    dirs: set[Path] = set()
+    for config in configs:
+        with open(config, encoding="utf-8") as f:
+            doc = yaml.safe_load(f)
+        declared = doc.get("job_dirs") if isinstance(doc, dict) else None
+        if not isinstance(declared, list) or not declared:
+            raise SystemExit(f"{config}: `job_dirs` is missing or not a list")
+        for name in declared:
+            job_dir = Path(name)
+            if not job_dir.is_dir():
+                raise SystemExit(f"{config}: job_dir `{name}` does not exist")
+            dirs.add(job_dir)
+    return sorted(dirs)
+
+
+def scan(steps: object, source: Path) -> list[str]:
+    """Label of every step in a `steps:` list that declares no key."""
+    if not isinstance(steps, list):
+        raise SystemExit(f"{source}: `steps:` is not a list")
+    keyless: list[str] = []
+    for i, step in enumerate(steps):
+        # A scalar shorthand such as `- wait` has nowhere to put a key. Every
+        # other step does, so there are no other exemptions. A `mirror:` block
+        # is not a step here: the generator publishes it as `<hw>-<parent key>`.
+        if not isinstance(step, dict):
+            continue
+        label = step.get("label") or step.get("group") or f"step {i}"
+        if "steps" in step:
+            raise SystemExit(
+                f"{source}: `{label}` nests its own `steps:`. The generator has "
+                "no group step and drops the children without an error, so give "
+                "each one its own entry."
+            )
+        key = step.get("key")
+        if not (isinstance(key, str) and key.strip()):
+            keyless.append(label)
+    return keyless
+
+
+def keyless_steps(path: Path) -> list[str]:
+    with open(path, encoding="utf-8") as f:
+        doc = yaml.safe_load(f)
+    if not isinstance(doc, dict) or "steps" not in doc:
+        return []
+    return scan(doc["steps"], path)
+
+
+def main() -> int:
+    files = sorted({p for d in job_dirs() for p in d.rglob(JOB_GLOB)})
+    if not files:
+        raise SystemExit("no pipeline files under the declared job_dirs")
+    keyless = [(path, label) for path in files for label in keyless_steps(path)]
+    if not keyless:
+        return 0
+    print(f"{len(keyless)} Buildkite step(s) have no explicit key:\n", file=sys.stderr)
+    for path, label in keyless:
+        print(f"  {path}: {label}", file=sys.stderr)
+    print(
+        "\nAdd a `key:` naming the test each step runs. Without one the pipeline"
+        "\ngenerator derives the key from the label, so retitling the step renames"
+        "\nits job and breaks retries and build history.",
+        file=sys.stderr,
+    )
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py
index 3be7bb243be5..126d20094a2a 100644
--- a/tools/pre_commit/check_forbidden_imports.py
+++ b/tools/pre_commit/check_forbidden_imports.py
@@ -13,6 +13,16 @@
     r"|list_repo_refs|repo_exists"
 )
 
+# Transformers v4 names kept as aliases in v5.
+_TRANSFORMERS_LEGACY_NAMES = (
+    r"PretrainedConfig|PreTrainedTokenizer|PreTrainedTokenizerFast"
+    r"|BaseImageProcessorFast|\w+ImageProcessorFast"
+)
+_TRANSFORMERS_LEGACY_MODULES = (
+    r"tokenization_utils|tokenization_utils_fast|image_processing_utils_fast"
+    r"|models\.\w+\.\w+_fast"
+)
+
 
 @dataclass
 class ForbiddenImport:
@@ -60,8 +70,6 @@ class ForbiddenImport:
             "benchmarks/kernels/benchmark_lora.py",
             "benchmarks/kernels/benchmark_machete.py",
             "benchmarks/fused_kernels/layernorm_rms_benchmarks.py",
-            "benchmarks/cutlass_benchmarks/w8a8_benchmarks.py",
-            "benchmarks/cutlass_benchmarks/sparse_benchmarks.py",
             # cloudpickle
             "vllm/v1/executor/multiproc_executor.py",
             "vllm/v1/executor/ray_executor.py",
@@ -119,6 +127,26 @@ class ForbiddenImport:
         allowed_files={"vllm/transformers_utils/repo_utils.py"},
         allowed_dirs={"examples/"},
     ),
+    "transformers legacy names": ForbiddenImport(
+        pattern=(
+            r"^\s*from\s+transformers(?:\.[\w.]+)?\s+import\s*\([^)]*\b(?:"
+            + _TRANSFORMERS_LEGACY_NAMES
+            + r")\b"
+            r"|"
+            r"^\s*from\s+transformers(?:\.[\w.]+)?\s+import\b[^\n]*\b(?:"
+            + _TRANSFORMERS_LEGACY_NAMES
+            + r")\b"
+            r"|"
+            r"^\s*(?:from|import)\s+transformers\.(?:"
+            + _TRANSFORMERS_LEGACY_MODULES
+            + r")\b"
+        ),
+        tip=(
+            "Use the Transformers v5 names, e.g. PreTrainedConfig, PythonBackend, "
+            "TokenizersBackend, TorchvisionBackend, tokenization_utils_tokenizers, "
+            "tokenization_utils_sentencepiece and image_processing_backends."
+        ),
+    ),
 }
 
 
@@ -249,6 +277,46 @@ def matches(rule: str, content: str) -> bool:
             f"(expected {should_match}, got {result})"
         )
 
+    transformers_cases = [
+        # Should match
+        ("from transformers import PretrainedConfig", True),
+        ("from transformers import AutoConfig, PretrainedConfig", True),
+        ("from transformers import PreTrainedTokenizer", True),
+        ("from transformers import PreTrainedTokenizerFast", True),
+        ("from transformers.models.siglip import SiglipImageProcessorFast", True),
+        ("from transformers import (\n    PreTrainedTokenizerFast,\n)", True),
+        ("from transformers.configuration_utils import PretrainedConfig", True),
+        ("    from transformers import PretrainedConfig", True),
+        ("from transformers.tokenization_utils import AddedToken", True),
+        ("from transformers.tokenization_utils_fast import X", True),
+        ("from transformers.image_processing_utils_fast import X", True),
+        (
+            "from transformers.models.qwen2_vl.image_processing_qwen2_vl_fast import X",
+            True,
+        ),
+        ("import transformers.tokenization_utils_fast", True),
+        # Should not match
+        ("from transformers import PreTrainedConfig", False),
+        ("from transformers import PreTrainedTokenizerBase", False),
+        ("from transformers import PythonBackend, TokenizersBackend", False),
+        ("from transformers.tokenization_utils_base import BatchEncoding", False),
+        ("from transformers.tokenization_utils_tokenizers import X", False),
+        (
+            "from transformers.image_processing_backends import TorchvisionBackend",
+            False,
+        ),
+        ("from transformers.models.qwen2_vl import Qwen2VLImageProcessor", False),
+        ("from vllm.foo import GLM4VImageProcessorFast", False),
+        ("from transformers import (\n    PreTrainedConfig,\n)", False),
+        ("# from transformers import PretrainedConfig", False),
+    ]
+    for i, (content, should_match) in enumerate(transformers_cases):
+        result = matches("transformers legacy names", content)
+        assert result == should_match, (
+            f"transformers case {i} failed: {content!r} "
+            f"(expected {should_match}, got {result})"
+        )
+
     print("All regex tests passed.")
 
 
diff --git a/tools/pre_commit/check_spdx_header.py b/tools/pre_commit/check_spdx_header.py
index 2585e82b4eba..1252ccb7adcb 100644
--- a/tools/pre_commit/check_spdx_header.py
+++ b/tools/pre_commit/check_spdx_header.py
@@ -16,7 +16,7 @@ class HeaderStyle:
 
 
 class SPDXStatus(Enum):
-    """SPDX header status enumeration"""
+    """SPDX header status enumeration."""
 
     EMPTY = "empty"  # empty __init__.py
     COMPLETE = "complete"
@@ -58,7 +58,7 @@ def header_insertion_index(style, lines):
 
 
 def check_spdx_header_status(file_path):
-    """Check SPDX header status of the file"""
+    """Check SPDX header status of the file."""
     license_line, copyright_line = spdx_header(file_style(file_path))
     with open(file_path, encoding="UTF-8") as file:
         lines = file.readlines()
@@ -92,7 +92,7 @@ def check_spdx_header_status(file_path):
 
 
 def add_header(file_path, status):
-    """Add or supplement SPDX header based on status"""
+    """Add or supplement SPDX header based on status."""
     style = file_style(file_path)
     license_line, copyright_line = spdx_header(style)
     full_spdx_header = f"{license_line}\n{copyright_line}"
@@ -137,7 +137,7 @@ def add_header(file_path, status):
 
 
 def main():
-    """Main function"""
+    """Main function."""
     files_missing_both = []
     files_missing_copyright = []
     files_missing_license = []
diff --git a/tools/pre_commit/check_test_tethering.py b/tools/pre_commit/check_test_tethering.py
new file mode 100644
index 000000000000..ee080a675734
--- /dev/null
+++ b/tools/pre_commit/check_test_tethering.py
@@ -0,0 +1,925 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Check that test files are wired into ("tethered to") the Buildkite CI.
+
+A test file is *tethered* when at least one Buildkite job actually collects it -
+some job's ``commands`` run ``pytest`` (or ``python`` / ``torchrun``, or a
+``find ... | xargs pytest`` pipeline) over a path that includes the file. A test
+that is collected by no job silently never runs: it passes review, merges, and
+rots.
+
+How it works
+------------
+1. Parse every job in every Buildkite pipeline yaml (and each ``mirror.amd``
+   sub-step), and turn each shell command into a :class:`Selection` describing
+   which test files it runs. The pipeline directories are read from
+   ``.buildkite/ci_config*.yaml`` (``job_dirs:``) - currently ``test_areas/``,
+   ``hardware_tests/`` and ``intel_jobs/`` across the CUDA, ROCm and Intel
+   pipelines - plus the legacy ``.buildkite/test-amd.yaml``.
+2. Build the inventory of test files under ``tests/``.
+3. A file is a problem if no selection runs it and it is not in the allowlist
+   (``tools/pre_commit/test_tethering_allowlist.txt`` - the set of pre-existing
+   gaps, which is only ever meant to shrink). Every allowlist entry must carry a
+   trailing ``# <reason>``; a bare path line is itself an error.
+
+Modes
+-----
+* **changed-files** (default): pre-commit passes the changed test files as
+  arguments; each one must be tethered or allowlisted.
+* **full scan**: ``--all``, or automatically whenever the change set touches a
+  pipeline yaml or the allowlist itself - because editing a job command can
+  orphan a test that is not otherwise in the diff. The full scan also reports
+  allowlist entries that have since become tethered, or whose file is gone, so
+  the list stays honest.
+
+Coverage from *any* of those pipelines counts - the question is "does some
+Buildkite job run this test", not "which pipeline". ``test-amd.yaml`` is being
+folded into ``test_areas/`` a group at a time; when it is gone the checker stops
+parsing it automatically.
+
+Usage::
+
+    python tools/pre_commit/check_test_tethering.py [FILES...]
+    python tools/pre_commit/check_test_tethering.py --all
+"""
+
+import argparse
+import fnmatch
+import shlex
+import subprocess
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import regex as re
+import yaml
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+BUILDKITE_DIR = REPO_ROOT / ".buildkite"
+# Fallback pipeline dir, always parsed even if ``ci_config*.yaml`` can't be read.
+TEST_AREAS_DIR = BUILDKITE_DIR / "test_areas"
+# The legacy hand-maintained AMD pipeline. It is being migrated into
+# ``test_areas/`` piecemeal; until that finishes, a test wired only into this
+# file is still genuinely run by CI, so it counts as tethered. When the file is
+# removed, ``_pipeline_yaml_paths()`` stops appending it automatically.
+TEST_AMD_YAML = BUILDKITE_DIR / "test-amd.yaml"
+TESTS_DIR = REPO_ROOT / "tests"
+ALLOWLIST_PATH = Path(__file__).resolve().parent / "test_tethering_allowlist.txt"
+
+# `pytest` options that take a value in the *next* token, so that token must not
+# be mistaken for a test path (e.g. `-k expr`, `--shard-id 3`). Options that only
+# ever take their value with `=` (`--cov=...`) don't belong here - a bare
+# `--cov` takes no value and would wrongly swallow the following test path.
+PYTEST_OPTIONS_WITH_VALUE = {
+    "-k", "-m", "-p", "-n", "-c", "-o", "-W", "-r",
+    "--shard-id", "--num-shards", "--dist", "--durations", "--maxfail",
+    "--timeout", "--rootdir", "--junitxml", "--tb",
+}  # fmt: skip
+
+# `pytest` options whose value is a path that is *removed* from collection at
+# file granularity. ``--ignore`` takes a path prefix; ``--ignore-glob`` takes a
+# glob that pytest matches against the whole path with ``fnmatch`` (so its ``*``
+# *does* cross a ``/``), which is why the two are modelled separately.
+PYTEST_IGNORE_OPTIONS = {"--ignore", "--ignore-glob"}
+
+# `--deselect` takes a value too, but it removes individual node IDs - a file
+# with one parametrization deselected is still collected and still "run".
+PYTEST_DESELECT_OPTION = "--deselect"
+
+PYTEST_COMMANDS = {"pytest", "py.test"}
+
+# A leading ``$VAR/`` or ``${VAR}/`` path prefix - CI shell scripts routinely
+# name a test by an absolute path rooted at a shell var
+# (``"${GIT_ROOT}/tests/v1/..."``). The prefix is dropped before the path is
+# interpreted; an unresolved ``$`` anywhere else still makes the token unusable.
+_LEADING_VAR_PREFIX = re.compile(r"^\$\{?\w+\}?/")
+
+# Commands that execute a Python file directly. Only the *first* path argument is
+# the script that runs; everything after it is that script's own argv, not more
+# test files (``python tensorize.py --suffix v1`` runs one file, not ``tests/v1``).
+FILE_RUNNER_COMMANDS = {"python", "python3", "torchrun", "coverage"}
+
+
+# --------------------------------------------------------------------------- #
+# Shell-token helpers
+# --------------------------------------------------------------------------- #
+
+
+def _is_env_assignment(token: str) -> bool:
+    """True for a leading ``FOO=bar`` shell env-var assignment - the kind some
+    commands put before ``pytest`` (``VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest``).
+    """
+    if "=" not in token:
+        return False
+    name = token.split("=", 1)[0]
+    return bool(name) and name.replace("_", "").isupper()
+
+
+def normalize_test_path(raw: str) -> str:
+    """Reduce a path as it appears in a command to one comparable form:
+    relative to ``tests/``, no surrounding quotes, no ``::nodeid`` selector, no
+    trailing slash.
+
+    Commands run from either the repo root or from ``tests/``, so both a
+    ``tests/`` and a ``/vllm-workspace/tests/`` prefix are stripped, as is a
+    leading ``$VAR/`` (``"${GIT_ROOT}/tests/..."``).
+    """
+    path = raw.strip().strip("'\"")
+    path = path.split("::", 1)[0]
+    path = _LEADING_VAR_PREFIX.sub("", path)
+    for prefix in ("./", "tests/", "/vllm-workspace/tests/"):
+        if path.startswith(prefix):
+            path = path[len(prefix) :]
+    return path.rstrip("/")
+
+
+def _token_is_test_path(token: str) -> bool:
+    """True if a shell token is a path argument rather than a flag, an env-var
+    assignment, a Buildkite ``$VAR``, or a ``-m`` / ``-k`` expression value.
+
+    A token with a ``/`` or a ``.py`` ending is always a path. A bare word
+    (``samplers``) only counts if it is really a directory under ``tests/``. A
+    leading ``$VAR/`` is dropped first, but a ``$`` anywhere else means the token
+    can't be resolved to a known path.
+    """
+    if not token or token.startswith("-"):
+        return False
+    if _is_env_assignment(token):
+        return False
+    token = _LEADING_VAR_PREFIX.sub("", token)
+    if "$" in token:
+        return False
+    if "/" in token or token.endswith(".py"):
+        return True
+    return (TESTS_DIR / normalize_test_path(token)).is_dir()
+
+
+def _path_is_within(parent: str, path: str) -> bool:
+    """True if ``path`` is ``parent`` itself or nested underneath it. A ``''`` /
+    ``'.'`` parent (a command run with no path arg) contains everything."""
+    if parent in ("", "."):
+        return True
+    return path == parent or path.startswith(parent + "/")
+
+
+def _glob_match_path(pattern: str, path: str) -> bool:
+    """Like :func:`fnmatch.fnmatch`, but ``*`` and ``?`` do not match across a
+    ``/`` - the way the shell and ``pytest`` treat an unexpanded glob path
+    argument (``kernels/test_foo_*.py`` selects one directory, not the tree).
+    A ``**`` component keeps the cross-separator meaning."""
+    if not any(ch in pattern for ch in "*?["):
+        return pattern == path
+    if "**" in pattern:
+        return fnmatch.fnmatch(path, pattern)
+    pattern_parts = pattern.split("/")
+    path_parts = path.split("/")
+    return len(pattern_parts) == len(path_parts) and all(
+        fnmatch.fnmatch(part, glob) for part, glob in zip(path_parts, pattern_parts)
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Selections - "which test files does this command run?"
+# --------------------------------------------------------------------------- #
+
+
+@dataclass
+class PytestSelection:
+    """A ``pytest``-style command: run everything under ``included_paths``
+    except what an ``--ignore`` (path prefix) or ``--ignore-glob`` (whole-path
+    ``fnmatch``) removes."""
+
+    included_paths: list[str] = field(default_factory=list)
+    ignored_paths: list[str] = field(default_factory=list)
+    ignored_globs: list[str] = field(default_factory=list)
+
+    def runs(self, test_file: str) -> bool:
+        included = any(self._arg_covers(arg, test_file) for arg in self.included_paths)
+        if not included:
+            return False
+        if any(self._arg_covers(arg, test_file) for arg in self.ignored_paths):
+            return False
+        # `--ignore-glob` is a plain `fnmatch` over the whole path - `*_tp.py`
+        # drops `kernels/nested/test_a_tp.py`, unlike a positional glob arg.
+        return not any(
+            fnmatch.fnmatch(test_file, normalize_test_path(glob))
+            for glob in self.ignored_globs
+        )
+
+    @staticmethod
+    def _arg_covers(path_arg: str, test_file: str) -> bool:
+        """True if a single pytest path argument selects ``test_file`` - by
+        exact match, by being a parent directory, or as a glob."""
+        path_arg = normalize_test_path(path_arg)
+        if _path_is_within(path_arg, test_file):
+            return True
+        return _glob_match_path(path_arg, test_file)
+
+
+@dataclass
+class FindSelection:
+    """A ``find <root> -name ... [-not -name ...] [-maxdepth N] | xargs pytest``
+    pipeline. ``pytest`` runs one file per ``find`` hit, so the ``find``
+    expression alone defines the selected set."""
+
+    root: str
+    name_globs: list[str] = field(default_factory=list)
+    exclude_name_globs: list[str] = field(default_factory=list)
+    max_depth: int | None = None
+
+    def runs(self, test_file: str) -> bool:
+        root = normalize_test_path(self.root)
+        if not _path_is_within(root, test_file):
+            return False
+
+        # `find <root>` puts a direct child at depth 1, whether <root> is a real
+        # subdir ("compile") or the tests/ root itself ("" / ".").
+        below = test_file if root in ("", ".") else test_file[len(root) :]
+        depth_below_root = below.lstrip("/").count("/") + 1
+        if self.max_depth is not None and depth_below_root > self.max_depth:
+            return False
+
+        filename = test_file.rsplit("/", 1)[-1]
+        if self.name_globs and not _glob_any(filename, self.name_globs):
+            return False
+        return not _glob_any(filename, self.exclude_name_globs)
+
+
+Selection = PytestSelection | FindSelection
+
+
+def _glob_any(name: str, globs: list[str]) -> bool:
+    return any(fnmatch.fnmatch(name, glob) for glob in globs)
+
+
+# --------------------------------------------------------------------------- #
+# Parsing CI commands into selections
+# --------------------------------------------------------------------------- #
+
+
+def _parse_pytest_command(tokens: list[str]) -> PytestSelection:
+    """Read a ``pytest ...`` invocation's tokens into a :class:`PytestSelection`.
+
+    Positional path args are the included paths; ``--ignore`` values (and a
+    bare-path ``--deselect``) are ignored paths, ``--ignore-glob`` values are
+    ignored globs. Value-taking options and their values, and every other flag,
+    are skipped. ``-k`` / ``-m`` filtering is intentionally *not* modelled - a
+    file that a marker expression narrows is still "run".
+    """
+    selection = PytestSelection()
+
+    # Advance past the pytest token itself (and any env-var / wrapper prefix).
+    tokens = iter(tokens)
+    for token in tokens:
+        if token in PYTEST_COMMANDS:
+            break
+
+    for token in tokens:
+        opt, eq, inline = token.partition("=")
+        if opt in PYTEST_IGNORE_OPTIONS:
+            value = inline if eq else next(tokens, "")
+            if opt == "--ignore-glob":
+                selection.ignored_globs.append(value)
+            else:
+                selection.ignored_paths.append(value)
+        elif opt == PYTEST_DESELECT_OPTION:
+            value = inline if eq else next(tokens, "")
+            # A `path::node_id` deselect leaves the file collected; only a
+            # bare-path deselect removes anything at file granularity.
+            if value and "::" not in value:
+                selection.ignored_paths.append(value)
+        elif token in PYTEST_OPTIONS_WITH_VALUE:
+            next(tokens, None)  # consume and discard the option's value
+        elif token.startswith("-"):
+            continue
+        elif _token_is_test_path(token):
+            selection.included_paths.append(token)
+
+    return selection
+
+
+def _parse_find_command(tokens: list[str]) -> FindSelection | None:
+    """Read a ``find <root> -name ...`` command's tokens into a
+    :class:`FindSelection`, or return ``None`` if there is no usable ``find``."""
+    if "find" not in tokens:
+        return None
+
+    root: str | None = None
+    name_globs: list[str] = []
+    exclude_name_globs: list[str] = []
+    max_depth: int | None = None
+
+    tokens = iter(tokens[tokens.index("find") + 1 :])
+    negated = False
+    for token in tokens:
+        if token in ("-not", "!"):
+            negated = True
+            continue
+        if token in ("-name", "-iname"):
+            glob = next(tokens, "").strip("'\"")
+            (exclude_name_globs if negated else name_globs).append(glob)
+        elif token == "-maxdepth":
+            depth = next(tokens, "")
+            max_depth = int(depth) if depth.isdigit() else None
+        elif (
+            root is None and not token.startswith("-") and token not in PYTEST_COMMANDS
+        ):
+            root = token
+        negated = False
+
+    if root is None:
+        return None
+    return FindSelection(root, name_globs, exclude_name_globs, max_depth)
+
+
+def _dash_c_code_operand(tokens: list[str]) -> str | None:
+    """For ``bash -c '<code>' [$0] [args]`` / ``sh -c ...`` / a combined short
+    group with ``c`` anywhere in it (``-ec``, ``-ce``, ``-euxc``), the ``<code>``
+    string the shell runs. ``None`` when this is not a ``-c`` invocation (plain
+    ``bash script.sh`` / ``bash -x script.sh``). Everything after the code
+    operand is a positional parameter (``$0``, ``$1``, ...), not a script or
+    command that executes."""
+    if not tokens or tokens[0].rsplit("/", 1)[-1] not in ("bash", "sh"):
+        return None
+    for i, token in enumerate(tokens[1:], start=1):
+        if not token.startswith("-"):
+            return None  # script operand reached; a later `-c` is *its* flag
+        # a short-option group (single dash) with `c` among its flags.
+        if not token.startswith("--") and "c" in token[1:]:
+            return tokens[i + 1] if i + 1 < len(tokens) else None
+    return None
+
+
+def _parse_shell_script(tokens: list[str], visited: set[str]) -> list[Selection]:
+    """Follow a ``bash tests/foo.sh`` command: read the script and parse the
+    pytest / runner lines inside it. ``visited`` stops a script that (directly or
+    otherwise) refers back to itself; nesting is one level deep in practice.
+    """
+    if _dash_c_code_operand(tokens) is not None:
+        return []  # the `-c` code string is parsed by the caller, not a file
+    script_arg = next((t for t in tokens if t.endswith(".sh")), None)
+    if script_arg is None:
+        return []
+
+    rel = script_arg.strip().strip("'\"").replace("/vllm-workspace/", "")
+    rel = rel.replace("vllm-workspace/", "").removeprefix("./").removeprefix("/")
+    candidates = [REPO_ROOT / rel, TESTS_DIR / normalize_test_path(rel)]
+    script_path = next((p for p in candidates if p.is_file()), None)
+    if script_path is None or str(script_path) in visited:
+        return []
+    visited.add(str(script_path))
+
+    # Join `\`-continued lines (a `pytest ... \` with its path on the next line
+    # is common), then track literal `VAR=value` assignments so a later
+    # `pytest "$VAR"` - the import-canary pattern - resolves to a real path.
+    text = script_path.read_text().replace("\\\n", " ")
+    var_values: dict[str, str] = {}
+    selections: list[Selection] = []
+    for line in text.splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+        assignment = re.match(r"^(\w+)=(\S+)$", line)
+        if assignment:
+            var_values[assignment.group(1)] = assignment.group(2).strip("'\"")
+            continue
+        if var_values:
+            line = _expand_vars(line, var_values)
+        selections.extend(_parse_command(line, visited))
+    return selections
+
+
+def _expand_vars(line: str, var_values: dict[str, str]) -> str:
+    """Substitute known ``$VAR`` / ``${VAR}`` in a script line - but not inside
+    single quotes, where the shell does not expand (``pytest '$X'`` passes a
+    literal ``$X``, collecting nothing)."""
+
+    def sub(segment: str) -> str:
+        return re.sub(
+            r"\$\{(\w+)\}|\$(\w+)",
+            lambda m: var_values.get(m.group(1) or m.group(2), m.group(0)),
+            segment,
+        )
+
+    # Splitting on `'`: even-indexed segments are outside single quotes, odd are
+    # inside. (A `'` within double quotes splits wrongly, but that only ever
+    # drops a substitution - the safe direction.)
+    parts = line.split("'")
+    return "'".join(sub(p) if i % 2 == 0 else p for i, p in enumerate(parts))
+
+
+# Shell tokens (or runs of them) that separate one sub-command from the next.
+_SHELL_SEP_CHARS = set("|&;()<>")
+_RUNNER_KEYWORDS = PYTEST_COMMANDS | FILE_RUNNER_COMMANDS
+
+# Tokens that can precede the real command word without being it.
+_COMMAND_WRAPPERS = {"if", "!", "then", "time", "exec", "nice", "command", "builtin"}
+
+
+def _effective_command_index(tokens: list[str]) -> int:
+    """Index of the token that names the command actually executed, past leading
+    ``FOO=bar`` assignments and simple wrappers - ``if ! env FOO=1 bash x.sh``
+    resolves to the ``bash``, so ``echo x.sh`` stays ``echo`` and is not
+    followed. Returns ``len(tokens)`` when nothing executable is left."""
+    i = 0
+    while i < len(tokens):
+        token = tokens[i]
+        if _is_env_assignment(token) or token in _COMMAND_WRAPPERS:
+            i += 1
+        elif token == "env":
+            i += 1
+            while i < len(tokens) and (
+                _is_env_assignment(tokens[i]) or "$" in tokens[i] or tokens[i] == "-"
+            ):
+                i += 1
+        elif token == "timeout" and i + 1 < len(tokens):
+            i += 1
+            while i < len(tokens) and tokens[i].startswith("-"):
+                # `-k` / `--kill-after` / `-s` / `--signal` take a value.
+                if tokens[i] in ("-k", "--kill-after", "-s", "--signal"):
+                    i += 1
+                i += 1
+            i += 1  # the DURATION operand
+        else:
+            break
+    return i
+
+
+def _split_shell_stages(command: str) -> list[tuple[str, list[str]]]:
+    """Tokenize a shell line - honouring quotes - and split it at the shell
+    operators, tagging each stage with the operator token that *precedes* it
+    (``""`` for the first stage, then ``"|"`` / ``"&&"`` / ``"||"`` / ``";"`` /
+    ``"&"`` / a subshell paren). Knowing which operator joins two stages is what
+    lets a ``find`` sweep count only when it is genuinely piped into ``pytest``,
+    not merely on the same line as one. A line that cannot be tokenized
+    (unbalanced quote, ...) yields nothing: better to miss a selection than to
+    misparse one. Splitting *after* tokenizing means an operator inside a quoted
+    argument stays part of that argument."""
+    lexer = shlex.shlex(command, posix=True, punctuation_chars=True)
+    lexer.whitespace_split = True
+    try:
+        tokens = list(lexer)
+    except ValueError:
+        return []
+
+    stages: list[tuple[str, list[str]]] = [("", [])]
+    in_exec = False  # inside a `find -exec ... \;` action - its `;` terminates
+    for token in tokens:
+        if in_exec and token in (";", "+"):
+            stages[-1][1].append(token)
+            in_exec = False
+        elif token and set(token) <= _SHELL_SEP_CHARS:
+            stages.append((token, []))
+            in_exec = False
+        else:
+            stages[-1][1].append(token)
+            if token in ("-exec", "-execdir"):
+                in_exec = True
+    return [(sep, sub) for sep, sub in stages if sub]
+
+
+def _looks_like_command_string(token: str) -> bool:
+    """True for a single token that is itself a command line - e.g. a quoted
+    ``"pytest a.py && torchrun b.py"`` passed as an argument to a runner
+    script."""
+    return " " in token and any(word in _RUNNER_KEYWORDS for word in token.split())
+
+
+def _classify_subcommand(tokens: list[str], visited: set[str]) -> list[Selection]:
+    """Map one already-tokenized sub-command to the selections it implies:
+    find / pytest / ``bash <script>`` / direct file runner. Anything else
+    contributes nothing - better to miss a real selection than to invent
+    coverage that isn't there."""
+    if not tokens:
+        return []
+
+    is_find = "find" in tokens and any(t in ("-name", "-iname") for t in tokens)
+    if is_find:
+        find_selection = _parse_find_command(tokens)
+        return [find_selection] if find_selection is not None else []
+
+    # A `pytest` token only counts when `pytest` is the command that actually
+    # runs - `echo pytest kernels` prints text, it does not collect anything.
+    if _runs_pytest(tokens):
+        return [_parse_pytest_command(tokens)]
+
+    # Resolve the real command word past `FOO=1` / `env` / `if ! ...` wrappers,
+    # so `env FOO=1 bash x.sh` is followed but `echo x.sh` (Buildkite just
+    # printing a filename) is not.
+    cmd_index = _effective_command_index(tokens)
+    command_name = (
+        tokens[cmd_index].rsplit("/", 1)[-1] if cmd_index < len(tokens) else ""
+    )
+
+    if command_name in ("bash", "sh") or command_name.endswith(".sh"):
+        run_tokens = tokens[cmd_index:]
+        code = _dash_c_code_operand(run_tokens)
+        if code is not None:
+            # `bash -c '<code>' [$0] [args]` - only the `-c` string runs; the
+            # trailing args are positional params, not commands.
+            return _parse_command(code, visited)
+        selections = _parse_shell_script(run_tokens, visited)
+        # Runner scripts (run-multi-node-test.sh, ...) take the real test
+        # commands as quoted string arguments - parse those too.
+        for arg in run_tokens[1:]:
+            if _looks_like_command_string(arg):
+                selections.extend(_parse_command(arg, visited))
+        return selections
+
+    # `python` / `torchrun` / `coverage <file>.py`.
+    if command_name in FILE_RUNNER_COMMANDS:
+        rest = tokens[cmd_index + 1 :]
+        # `python -c '<code>' [argv...]` runs the code string, not a file - a
+        # later `-m pytest` or `.py` is that program's argv. Only honour a `-c`
+        # that is the interpreter's own (nothing script-like precedes it).
+        if "-c" in rest:
+            before_c = rest[: rest.index("-c")]
+            if "-m" not in before_c and not any(t.endswith(".py") for t in before_c):
+                return []
+        # `-m <module>` picks the module, not a script: only `-m pytest` runs
+        # tests. `python -m compileall kernels` collects nothing - the trailing
+        # `kernels` is compileall's argv, not a test path.
+        if "-m" in rest:
+            after_m = rest[rest.index("-m") + 1 :]
+            if after_m[:1] and after_m[0] in PYTEST_COMMANDS:
+                return [_parse_pytest_command(after_m)]
+            return []
+        # `python foo.py a b/c --suffix v1` executes only `foo.py`; the tokens
+        # after it are that script's argv - treating them as test paths tethers
+        # whatever directory an argument happens to name (`v1` -> `tests/v1`).
+        script = next(
+            (t for t in rest if t.endswith(".py") and _token_is_test_path(t)), None
+        )
+        if script is not None:
+            return [PytestSelection(included_paths=[script])]
+
+    return []
+
+
+# `xargs` options that consume the following token as their value.
+_XARGS_VALUE_FLAGS = {
+    "-n", "-L", "-P", "-s", "-I", "-i", "-J", "-E", "-d",
+    "--max-args", "--max-lines", "--max-procs", "--replace", "--delimiter",
+}  # fmt: skip
+
+
+def _xargs_command_tokens(tokens: list[str]) -> list[str]:
+    """The command (plus its args) that ``xargs`` will exec, with ``xargs``' own
+    options stripped: ``xargs -0 -n1 -I{} env FOO=1 pytest`` -> ``['env',
+    'FOO=1', 'pytest']``."""
+    i = 0
+    while i < len(tokens):
+        if tokens[i] in _XARGS_VALUE_FLAGS:
+            i += 2
+        elif tokens[i].startswith("-"):
+            i += 1
+        else:
+            return tokens[i:]
+    return []
+
+
+def _runs_pytest(tokens: list[str]) -> bool:
+    """True if ``tokens`` is a command that executes ``pytest``, seeing through
+    ``env FOO=1`` / ``timeout`` / ``if !`` wrappers (``xargs echo pytest`` is
+    not - ``echo`` is what runs)."""
+    idx = _effective_command_index(tokens)
+    return idx < len(tokens) and tokens[idx] in PYTEST_COMMANDS
+
+
+# Pipe stages that reorder or echo the stream without dropping any line, so
+# every `find` match still reaches a downstream `xargs pytest`. Anything else
+# (`head`, `grep`, `shuf`, `sed -n`, ...) filters, and then only some matches
+# run - not enough to claim the whole sweep is tethered.
+_PIPE_PASSTHROUGH = {"sort", "tee", "cat", "tac"}
+
+
+def _find_feeds_pytest(stages: list[tuple[str, list[str]]], find_index: int) -> bool:
+    """True if the ``find`` stage at ``find_index`` actually runs its matches
+    through ``pytest``: a *terminated* ``-exec pytest ... \\;`` / ``... +`` in the
+    ``find`` itself, or an ``xargs`` whose command operand is ``pytest``, reached
+    through ``|`` pipes of non-filtering stages only. A ``find`` piped to
+    ``wc`` / ``head`` / an ``xargs echo``, an unterminated ``-exec``, or a
+    ``pytest`` in a separate ``;`` / ``&&`` command, runs nothing (or not the
+    whole sweep)."""
+    find_tokens = stages[find_index][1]
+    for i, token in enumerate(find_tokens):
+        if token in ("-exec", "-execdir"):
+            action = find_tokens[i + 1 :]
+            if _exec_runs_matches_through_pytest(action):
+                return True
+    for sep, tokens in stages[find_index + 1 :]:
+        if sep != "|":
+            break
+        if tokens[:1] == ["xargs"] and _runs_pytest(_xargs_command_tokens(tokens[1:])):
+            return True
+        if tokens[:1] and tokens[0] in _PIPE_PASSTHROUGH:
+            continue
+        break
+    return False
+
+
+def _terminated(exec_action: list[str]) -> bool:
+    """A ``find -exec`` action only runs if it ends with a ``;`` or ``+`` token
+    (from ``\\;`` / ``+`` on the command line)."""
+    return ";" in exec_action or "+" in exec_action
+
+
+def _exec_runs_matches_through_pytest(action: list[str]) -> bool:
+    """True if a ``find -exec``/``-execdir`` action runs ``pytest`` over each
+    *matched* file: ``pytest`` is the command, the action is terminated
+    (``\\;`` / ``+``), and the ``{}`` placeholder lands as a positional test-path
+    argument - not as the value of an option (``-exec pytest --ignore {} \\;``
+    only ever runs whatever positional file is also on the line)."""
+    if not action or action[0] not in PYTEST_COMMANDS or not _terminated(action):
+        return False
+    # Substitute a real-looking path for `{}` and see whether pytest's own arg
+    # parser keeps it as an included path rather than eating it as an option value.
+    probe = [t.replace("{}", "_find_match_.py") for t in action[:-1]]
+    selection = _parse_pytest_command(probe)
+    return any("_find_match_" in path for path in selection.included_paths)
+
+
+def _parse_command(command: str, visited: set[str] | None = None) -> list[Selection]:
+    """Turn one yaml ``commands:`` entry into the selections it implies.
+
+    The line is tokenized (respecting quotes) and split into stages at the shell
+    operators, and each stage is classified independently so that, e.g., a
+    ``find`` and a ``pytest`` on the same line are both handled. A ``find -name``
+    sweep is kept only when :func:`_find_feeds_pytest` confirms its output
+    actually reaches ``pytest`` - otherwise it would invent coverage.
+    """
+    visited = visited if visited is not None else set()
+    stages = _split_shell_stages(command)
+    selections: list[Selection] = []
+    for index, (_, tokens) in enumerate(stages):
+        for selection in _classify_subcommand(tokens, visited):
+            if isinstance(selection, FindSelection) and not _find_feeds_pytest(
+                stages, index
+            ):
+                continue
+            selections.append(selection)
+    return selections
+
+
+def _iter_job_steps(yaml_doc: dict):
+    """Yield each job step in a test_areas yaml, plus a merged view of each
+    ``mirror.*`` sub-step. The AMD mirror can override ``commands``, so it has to
+    be checked as its own step."""
+    for step in yaml_doc.get("steps") or []:
+        if not isinstance(step, dict):
+            continue
+        yield step
+
+        for mirror_step in (step.get("mirror") or {}).values():
+            if isinstance(mirror_step, dict):
+                merged = {**step, **mirror_step}
+                merged.pop("mirror", None)
+                yield merged
+
+
+def _pipeline_job_dirs() -> list[Path]:
+    """The Buildkite job directories, read from every ``.buildkite/ci_config*.yaml``
+    (``job_dirs:``). Buildkite itself uses these files to assemble each pipeline,
+    so a test wired into any job under any of them is genuinely run by CI. Falls
+    back to ``test_areas/`` alone if no config is readable."""
+    job_dirs: set[Path] = {TEST_AREAS_DIR}
+    for config in sorted((REPO_ROOT / ".buildkite").glob("ci_config*.yaml")):
+        try:
+            doc = yaml.safe_load(config.read_text()) or {}
+        except yaml.YAMLError:
+            continue
+        for rel in doc.get("job_dirs") or []:
+            job_dirs.add(REPO_ROOT / rel)
+    return sorted(job_dirs)
+
+
+def _pipeline_yaml_paths() -> list[Path]:
+    """Every Buildkite pipeline yaml whose jobs count as CI coverage: the yamls
+    in every :func:`_pipeline_job_dirs` directory plus the legacy ``test-amd.yaml``
+    (see ``TEST_AMD_YAML``)."""
+    paths: list[Path] = []
+    for job_dir in _pipeline_job_dirs():
+        paths.extend(sorted(job_dir.glob("*.yaml")))
+    if TEST_AMD_YAML.is_file():
+        paths.append(TEST_AMD_YAML)
+    return paths
+
+
+def load_selections() -> list[Selection]:
+    """Parse every CI pipeline yaml into the full list of selections - the CI's
+    complete picture of which test files it runs.
+
+    A yaml that doesn't parse is fatal: silently skipping it would drop whatever
+    coverage it defines and produce false "untethered" reports.
+    """
+    selections: list[Selection] = []
+    for yaml_path in _pipeline_yaml_paths():
+        try:
+            yaml_doc = yaml.safe_load(yaml_path.read_text()) or {}
+        except yaml.YAMLError as e:
+            rel = yaml_path.relative_to(REPO_ROOT)
+            raise SystemExit(f"error: {rel} is not valid YAML: {e}") from None
+        if not isinstance(yaml_doc, dict):
+            rel = yaml_path.relative_to(REPO_ROOT)
+            raise SystemExit(
+                f"error: {rel} is not a Buildkite pipeline "
+                f"(top-level {type(yaml_doc).__name__}, expected a mapping)"
+            )
+        for step in _iter_job_steps(yaml_doc):
+            commands = list(step.get("commands") or [])
+            if step.get("command"):  # some steps use the singular form
+                commands.append(step["command"])
+            for command in commands:
+                if isinstance(command, str):
+                    selections.extend(_parse_command(command))
+    return selections
+
+
+def is_tethered(test_file: str, selections: list[Selection]) -> bool:
+    """True if any CI selection runs ``test_file`` (path relative to ``tests/``
+    or to the repo root - both are accepted)."""
+    normalized = normalize_test_path(test_file)
+    return any(selection.runs(normalized) for selection in selections)
+
+
+# --------------------------------------------------------------------------- #
+# Test-file inventory and allowlist
+# --------------------------------------------------------------------------- #
+
+
+def is_test_module(repo_path: str) -> bool:
+    """True if ``repo_path`` is a pytest test module CI is expected to run - a
+    ``tests/**/test_*.py`` or ``*_test.py``, but not ``conftest.py`` /
+    ``__init__.py`` or other helper modules."""
+    if not repo_path.startswith("tests/") or not repo_path.endswith(".py"):
+        return False
+    filename = repo_path.rsplit("/", 1)[-1]
+    if filename in ("conftest.py", "__init__.py"):
+        return False
+    return filename.startswith("test_") or filename.endswith("_test.py")
+
+
+def all_test_modules() -> list[str]:
+    """Every test module under ``tests/``, sorted. Prefers ``git ls-files`` so a
+    developer's local scratch files are never flagged; falls back to a
+    filesystem walk when run outside a git checkout - the CI test images bake in
+    ``tests/`` with no ``.git``, and there are no untracked scratch files there.
+    """
+    try:
+        tracked = subprocess.run(
+            ["git", "-C", str(REPO_ROOT), "ls-files", "tests/"],
+            capture_output=True,
+            text=True,
+            check=True,
+        ).stdout
+        paths = tracked.splitlines()
+    except (OSError, subprocess.CalledProcessError):
+        paths = [p.relative_to(REPO_ROOT).as_posix() for p in TESTS_DIR.rglob("*.py")]
+    return sorted(path for path in paths if is_test_module(path))
+
+
+def load_allowlist() -> set[str]:
+    """Read the allowlist file into a set of repo-relative paths. Blank lines and
+    ``#`` comments (whole-line or trailing) are ignored."""
+    if not ALLOWLIST_PATH.exists():
+        return set()
+
+    paths = set()
+    for line in ALLOWLIST_PATH.read_text().splitlines():
+        path = line.split("#", 1)[0].strip()
+        if path:
+            paths.add(path)
+    return paths
+
+
+def allowlist_entries_missing_reason() -> list[str]:
+    """Allowlist paths with no trailing ``# reason`` on their own line. Every gap
+    has to say why the test can't be wired into a job, so a bare ``path`` line
+    (the easy way to silence the gate) is itself an error."""
+    if not ALLOWLIST_PATH.exists():
+        return []
+
+    offenders = []
+    for line in ALLOWLIST_PATH.read_text().splitlines():
+        path, _, comment = line.partition("#")
+        if path.strip() and not comment.strip():
+            offenders.append(path.strip())
+    return offenders
+
+
+# --------------------------------------------------------------------------- #
+# The two check modes
+# --------------------------------------------------------------------------- #
+
+
+def run_full_scan(
+    selections: list[Selection], allowlist: set[str], *, strict: bool
+) -> int:
+    """Scan every test module. Untethered files that aren't allowlisted are
+    always errors. Allowlist entries that have become tethered, or whose file no
+    longer exists, are reported either way but are only errors when ``strict``
+    (an explicit ``--all`` run, as opposed to a PR that happened to touch CI
+    config - which shouldn't be blocked by unrelated allowlist cleanup).
+    """
+    untethered = [
+        path
+        for path in all_test_modules()
+        if path not in allowlist and not is_tethered(path, selections)
+    ]
+    now_tethered = sorted(path for path in allowlist if is_tethered(path, selections))
+    deleted = sorted(path for path in allowlist if not (REPO_ROOT / path).exists())
+
+    for path in untethered:
+        print(f"error: {path} is not run by any Buildkite job")
+    for path in now_tethered:
+        print(f"note: {path} is now tethered - remove it from the allowlist")
+    for path in deleted:
+        print(f"note: {path} no longer exists - remove it from the allowlist")
+
+    if untethered:
+        print(
+            f"\n{len(untethered)} test file(s) exist but run in no CI job. Wire "
+            "each into a job's `commands` in a .buildkite/ pipeline yaml, or add "
+            f"it to {ALLOWLIST_PATH.relative_to(REPO_ROOT)} with a reason."
+        )
+
+    error_count = len(untethered)
+    if strict:
+        error_count += len(now_tethered) + len(deleted)
+    return 1 if error_count else 0
+
+
+def run_changed_files_check(
+    paths: list[str], selections: list[Selection], allowlist: set[str]
+) -> int:
+    """Check only the given paths (what pre-commit passes). Non-test files, and
+    paths that don't exist on disk, are ignored."""
+    untethered = []
+    for raw_path in paths:
+        repo_path = _to_repo_relative(raw_path)
+        if not is_test_module(repo_path) or not (REPO_ROOT / repo_path).is_file():
+            continue
+        if repo_path in allowlist or is_tethered(repo_path, selections):
+            continue
+        untethered.append(repo_path)
+
+    for path in untethered:
+        print(f"error: {path} is not run by any Buildkite job")
+    if untethered:
+        print(
+            "\nAdd the test to a job's `commands` (and `source_file_dependencies`)"
+            " in a .buildkite/ pipeline yaml, or, if it is intentionally not run "
+            f"in CI, add it to {ALLOWLIST_PATH.relative_to(REPO_ROOT)} with a "
+            "comment explaining why."
+        )
+    return 1 if untethered else 0
+
+
+def _to_repo_relative(path: str) -> str:
+    if Path(path).is_absolute():
+        path = str(Path(path).resolve().relative_to(REPO_ROOT))
+    return path.replace("\\", "/")
+
+
+def _change_set_touches_ci_config(paths: list[str]) -> bool:
+    """A change to a pipeline yaml, a ci_config, or the allowlist can orphan a
+    test that is not itself in the change set, so those changes force a full
+    scan."""
+    return any(
+        ".buildkite/test_areas/" in path
+        or ".buildkite/hardware_tests/" in path
+        or ".buildkite/intel_jobs/" in path
+        or ".buildkite/ci_config" in path
+        or path.endswith("test-amd.yaml")
+        or path.endswith("test_tethering_allowlist.txt")
+        for path in paths
+    )
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("files", nargs="*", help="test files to check")
+    parser.add_argument(
+        "--all", action="store_true", help="scan every test file in the tree"
+    )
+    args = parser.parse_args()
+
+    selections = load_selections()
+    allowlist = load_allowlist()
+
+    unreasoned = allowlist_entries_missing_reason()
+    for path in unreasoned:
+        print(
+            f"error: {ALLOWLIST_PATH.relative_to(REPO_ROOT)}: {path} has no "
+            "trailing '# <reason>' - say why it can't be wired into a job"
+        )
+
+    if args.all or _change_set_touches_ci_config(args.files):
+        rc = run_full_scan(selections, allowlist, strict=args.all)
+    else:
+        rc = run_changed_files_check(args.files, selections, allowlist)
+    return 1 if (unreasoned or rc) else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tools/pre_commit/generate_nightly_torch_test.py b/tools/pre_commit/generate_nightly_torch_test.py
index 354da54df8f4..3b39c4a0b274 100644
--- a/tools/pre_commit/generate_nightly_torch_test.py
+++ b/tools/pre_commit/generate_nightly_torch_test.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Generates specialized requirements files for nightly PyTorch testing.
+"""Generates specialized requirements files for nightly PyTorch testing.
 
 This script reads the main test requirements input file (`requirements/test/cuda.in`)
 and splits its content into two files:
diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py
index 9da39acc7899..71dd929d378e 100755
--- a/tools/pre_commit/mypy.py
+++ b/tools/pre_commit/mypy.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Run mypy on changed files.
+"""Run mypy on changed files.
 
 This script is designed to be used as a pre-commit hook. It runs mypy
 on files that have been changed. It groups files into different mypy calls
@@ -14,6 +13,7 @@
     python_version: Python version to use (e.g., "3.10") or "local" to use
         the local Python version.
     changed_files: List of changed files to check.
+
 """
 
 import subprocess
@@ -113,10 +113,7 @@
 
 EXCLUDE = [
     r"vllm/model_executor/models/[kK]",
-    r"vllm/model_executor/models/[nN]",
-    r"vllm/model_executor/models/[oO]",
     r"vllm/model_executor/models/[qQ]",
-    r"vllm/model_executor/models/transformers",
     r"vllm/model_executor/models/[uU]",
     r"vllm/model_executor/models/[vV]",
     r"vllm/model_executor/models/[wW]",
@@ -125,14 +122,14 @@
 
 
 def group_files(changed_files: list[str]) -> dict[str, list[str]]:
-    """
-    Group changed files into different mypy calls.
+    """Group changed files into different mypy calls.
 
     Args:
         changed_files: List of changed files.
 
     Returns:
         A dictionary mapping file group names to lists of changed files.
+
     """
     exclude_pattern = re.compile(f"^{'|'.join(EXCLUDE)}.*")
     silent_pattern = re.compile(f"^({'|'.join(SILENT_GROUPS)}).*")
@@ -166,8 +163,7 @@ def mypy(
     follow_imports: str | None,
     file_group: str,
 ) -> int:
-    """
-    Run mypy on the given targets.
+    """Run mypy on the given targets.
 
     Args:
         targets: List of files or directories to check.
@@ -179,6 +175,7 @@ def mypy(
 
     Returns:
         The return code from mypy.
+
     """
     args = ["mypy"]
     if python_version is not None:
diff --git a/tools/pre_commit/test_tethering_allowlist.txt b/tools/pre_commit/test_tethering_allowlist.txt
new file mode 100644
index 000000000000..80b4aaec4c2d
--- /dev/null
+++ b/tools/pre_commit/test_tethering_allowlist.txt
@@ -0,0 +1,100 @@
+# Allowlist for tools/pre_commit/check_test_tethering.py
+#
+# Every path below is a test file that is collected by NO Buildkite job in any
+# pipeline yaml the checker parses (the job_dirs of .buildkite/ci_config*.yaml -
+# test_areas/, hardware_tests/, intel_jobs/ - plus .buildkite/test-amd.yaml).
+# These are pre-existing gaps, captured when the check was first added so that
+# the gate could start green. The list is meant to only ever shrink.
+#
+# Rules:
+#   * Every entry needs a trailing `# <reason>` on its own line saying why the
+#     test isn't wired into a job. A bare path is rejected by the check.
+#   * Do not add a new entry to unblock a PR without a real justification.
+#   * To remove an entry: add the test to a job's `commands` (and matching
+#     `source_file_dependencies`) in a .buildkite/ pipeline yaml, then delete the
+#     line.
+#   * `check_test_tethering.py --all` reports entries here that have since
+#     become tethered, and entries whose file no longer exists.
+
+# tests/distributed/ - multi-GPU / multi-node suites, launched via torchrun or a
+# custom runner rather than a directory-level pytest invocation.
+tests/distributed/aux_output_connector/test_store.py  # added by #45635 (2026-09-21), no job collects this path
+tests/distributed/test_dcp_a2a.py  # multi-GPU DCP all-to-all; no dir-level pytest job
+tests/distributed/test_dcp_direct_a2a_lse_reduce.py  # multi-GPU DCP all-to-all; no dir-level pytest job
+tests/distributed/test_eplb_events.py  # multi-GPU EPLB; no dir-level pytest job
+tests/distributed/test_eplb_fused_moe_layer.py  # multi-GPU EPLB; no dir-level pytest job
+tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py  # multi-GPU EPLB nvfp4; no dir-level pytest job
+tests/distributed/test_eplb_quant_scale_consistency.py  # multi-GPU EPLB; no dir-level pytest job
+tests/distributed/test_expert_parallel.py  # multi-GPU expert parallel; no dir-level pytest job
+tests/distributed/test_expert_placement.py  # multi-GPU expert placement; no dir-level pytest job
+tests/distributed/test_file_store.py  # no distributed.yaml job collects this path
+tests/distributed/test_flashinfer_pcie_ipc_all_reduce.py  # multi-GPU PCIe IPC all-reduce (#53576); no job
+tests/distributed/test_kimi_linear_context_parallel.py  # multi-GPU context parallel; no dir-level pytest job
+tests/distributed/test_kv_cache_events.py  # needs multi-proc KV-events setup; no job
+tests/distributed/test_kvlayout.py  # no distributed.yaml job collects this path
+tests/distributed/test_mq_connect_ip.py  # no distributed.yaml job collects this path
+tests/distributed/test_pipeline_partition.py  # no distributed.yaml job collects this path
+tests/distributed/test_split_group.py  # no distributed.yaml job collects this path
+tests/distributed/test_weight_transfer_nccl_uid.py  # sibling test_weight_transfer.py is tethered; not added to that job (#53784)
+
+# tests/entrypoints/cli/
+tests/entrypoints/cli/test_serve.py  # added by #53124 (2026-09-15); no job collects tests/entrypoints/cli/
+
+# tests/kernels/
+tests/kernels/test_shuffle_rows.py  # explicitly --ignore'd in kernels.yaml with no follow-up job
+
+# tests/lora/ - explicitly --ignore'd from the sharded LoRA run, with no
+# dedicated *_tp job (unlike the other ignored *_tp tests).
+tests/lora/test_deepseekv2_tp.py  # --ignore'd from sharded LoRA run; no dedicated *_tp job
+tests/lora/test_qwen3moe_tp.py  # --ignore'd from sharded LoRA run; no dedicated *_tp job
+
+# tests/models/ - the models_*.yaml jobs run specific paths, not all of models/.
+tests/models/deepseek_v32/test_sequence_parallel.py  # no models_*.yaml job collects this path
+tests/models/test_deepseek_v4_dspark_rocm.py  # added by #52362 (2026-09-21), no models_*.yaml job collects this path
+tests/models/test_deepseek_v4_fi_moe_ep.py  # no models_*.yaml job collects this path
+tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py  # no models_*.yaml job collects this path
+tests/models/test_deepseek_v4_rocm_wo_a.py  # added by #54894 (2026-09-20), no models_*.yaml job collects this path
+tests/models/test_deepseek_v41_replay_start.py  # added by #56227 (2026-09-18), no models_*.yaml job collects this path
+tests/models/test_dspark_mla.py  # no models_*.yaml job collects this path
+tests/models/test_language_model_cache_is_weak.py  # no models_*.yaml job collects this path
+tests/models/transformers/test_layer_registry.py  # no models_*.yaml job collects this path
+
+# tests/ root-level - misc.yaml runs a hand-picked set of these by name
+# (test_regression.py, test_envs.py, test_outputs.py, test_pooling_params.py);
+# the rest are unreferenced.
+tests/test_access_log_filter.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_attention_backend_registry.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_audio_media_size_precheck.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_cmake_utils.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_connections.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_embedded_commit.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_force_first_config.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_fxgraphcache_pickle_patch.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_logprobs.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_quantization_revision_pin.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_ray_env_utils.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_request_input_bounds.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_scalartype.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_tensorizer_revision_pin.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_triton_utils.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_version.py  # root test_*.py not in misc.yaml's by-name list
+tests/test_zen_cpu_platform_detection.py  # root test_*.py not in misc.yaml's by-name list
+
+# tests/v1/ - the v1 subtree is otherwise covered by many `pytest v1/...` jobs.
+tests/v1/e2e/test_replayssm_decode.py  # ReplaySSM decode test (#52506) wired into no job
+tests/v1/ec_connector/integration/test_epd_correctness.py  # run_epd_correctness_test.sh exists but no Buildkite job invokes it
+tests/v1/ec_connector/integration/test_nixl_failure.py  # 2-GPU + NIXL E/PD integration test (#55290); no Buildkite job collects ec_connector/integration/
+tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py  # run by intel_jobs/misc_intel.yaml's run_xpu_disagg_accuracy_test.sh, which names it via $EXP_ROOT (the script's own dir) - not statically resolvable
+
+# tests/platforms/
+tests/platforms/test_cpu.py  # CPU-only unit test; no job targets tests/platforms/
+
+# tests/tools/
+tests/tools/test_config_validator.py  # logic covered by the validate-config pre-commit hook; no test_areas job
+
+# tests/watermarking/ - added by #54053 (2026-09-09), no Buildkite job collects the dir yet
+tests/watermarking/test_detection.py  # watermarking feature (#54053); no job collects tests/watermarking/
+tests/watermarking/test_gumbel.py  # watermarking feature (#54053); no job collects tests/watermarking/
+tests/watermarking/test_prf.py  # watermarking feature (#54053); no job collects tests/watermarking/
+tests/watermarking/test_watermarking.py  # watermarking feature (#54053); no job collects tests/watermarking/
+tests/watermarking/test_watermarking_e2e.py  # watermarking hardening (#56309); no job collects tests/watermarking/
diff --git a/tools/pre_commit/validate_config.py b/tools/pre_commit/validate_config.py
index 7da32bc6b485..498caf7ab2ca 100644
--- a/tools/pre_commit/validate_config.py
+++ b/tools/pre_commit/validate_config.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Ensures all fields in a config dataclass have default values
+"""Ensures all fields in a config dataclass have default values
 and that each field has a docstring.
 """
 
@@ -14,13 +13,11 @@
 
 
 def get_attr_docs(cls_node: ast.ClassDef) -> dict[str, str]:
-    """
-    Get any docstrings placed after attribute assignments in a class body.
+    """Get any docstrings placed after attribute assignments in a class body.
 
     Adapted from https://davidism.com/attribute-docstrings/
     https://davidism.com/mit-license/
     """
-
     out = {}
 
     # Consider each pair of nodes.
diff --git a/tools/profiler/nsys_profile_tools/gputrc2graph.py b/tools/profiler/nsys_profile_tools/gputrc2graph.py
index fd237c0b214a..2759d029f75d 100755
--- a/tools/profiler/nsys_profile_tools/gputrc2graph.py
+++ b/tools/profiler/nsys_profile_tools/gputrc2graph.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This generates gpu kernel analysis output from nsys rep. Will call nsys
+"""This generates gpu kernel analysis output from nsys rep. Will call nsys
 stats  -r cuda_gpu_kern_trace, get non-overlapped gpu cycles, then generate
 csv and html output for analysis
 """
@@ -17,7 +16,7 @@
 
 # helper data class for annotating kernels
 def load_engine_model():
-    """returns engine_model built from all json files in the current dir"""
+    """Returns engine_model built from all json files in the current dir."""
     import glob
     import json
 
@@ -31,9 +30,7 @@ def load_engine_model():
 
 
 class GPUTrace2Graph:
-    """
-    Parses output of nsys report, generates csv and bar chart output
-    """
+    """Parses output of nsys report, generates csv and bar chart output."""
 
     def __init__(self):
         import pandas as pd  # avoid importing till needed
@@ -64,8 +61,7 @@ def gen_nonoverlapped_sum_from_gputrace(self, in_file, out_file):
         )
 
     def sum_non_overlapping_intervals(self, df):
-        """
-        returns new sorted df with Elapsed Time (ns) column using
+        """Returns new sorted df with Elapsed Time (ns) column using
         vectorized operations
         """
         logger.info("sorting %s trace records by start time", str(df.shape))
@@ -105,7 +101,7 @@ def sum_non_overlapping_intervals(self, df):
 
     # functions for generating html files
     def make_html(self, df, output_dir, title):
-        """make html graph from df"""
+        """Make html graph from df."""
         import plotly.express as px
 
         if df.empty:
@@ -163,7 +159,7 @@ def make_html(self, df, output_dir, title):
         )
 
     def anno_gpu_kernname(self, df, mapping):
-        """add "Category" column"""
+        """Add "Category" column."""
 
         def anno_gpu_kernname_helper(name):
             for kern_name, val in mapping.items():
@@ -173,7 +169,7 @@ def anno_gpu_kernname_helper(name):
         df["Category"] = df["Name"].apply(anno_gpu_kernname_helper)
 
     def make_nongpu_row(self, df, nongpu_sec):
-        """this will append non-gpu time entry at end of df"""
+        """This will append non-gpu time entry at end of df."""
         nongpu_row = self.pd.DataFrame([df.iloc[-1]])
         nongpu_row["Category"] = nongpu_row["Name"] = "CPU(non-GPU)"
         nongpu_row["Instances"] = 1
@@ -181,13 +177,13 @@ def make_nongpu_row(self, df, nongpu_sec):
         return nongpu_row
 
     def is_valid_file(self, base_file):
-        """asserts if base_file is non-existent or is empty"""
+        """Asserts if base_file is non-existent or is empty."""
         assert os.path.isfile(base_file) and os.path.getsize(base_file) > 0, (
             f"{base_file} doesn't exist or is empty"
         )
 
     def should_gen_file(self, new_file, base_file):
-        """figure out if new file should be generated from base_file"""
+        """Figure out if new file should be generated from base_file."""
         self.is_valid_file(base_file)
         if (
             os.path.exists(new_file)
@@ -201,8 +197,7 @@ def should_gen_file(self, new_file, base_file):
             return True
 
     def gen_sum_file(self, file, nsys_cmd):
-        """
-        generates sum file from nsys trace with times per kernel and
+        """Generates sum file from nsys trace with times per kernel and
         returns the name of the sum file
         """
         import subprocess
@@ -246,7 +241,7 @@ def gen_sum_file(self, file, nsys_cmd):
         return sum_file
 
     def gen_graph(self, in_file, out_dir, title, nsys_cmd, engine_model):
-        """generates graph and csv file from in_file into out_dir"""
+        """Generates graph and csv file from in_file into out_dir."""
         # Initialize an empty DataFrame to store combined data
         combined_df = self.pd.DataFrame()
         for idx, (file, engine, model, total_sec) in enumerate(in_file):
diff --git a/tools/recipes/README.md b/tools/recipes/README.md
index 058470167a7b..253bd7a726ee 100644
--- a/tools/recipes/README.md
+++ b/tools/recipes/README.md
@@ -21,6 +21,21 @@ flowchart LR
 The recipe is the baseline. Hardware and workload information can optionally
 refine the initial configuration. Sweep tuning is an optional validation step.
 
+## Getting Started
+
+Use `serve_with_recipe.sh` to generate the recipe configuration and start vLLM
+in one step:
+
+```bash
+tools/recipes/serve_with_recipe.sh \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --hardware xeon6
+```
+
+For `xeon6`, the script enables hardware detection automatically. Generated
+`config.yml` and `env.sh` files remain in the directory where the script was
+invoked.
+
 ## 1. vLLM Recipes Only
 
 Use the converter directly when the recipe already contains the deployment
diff --git a/tools/recipes/recipe_json_to_vllm_config.py b/tools/recipes/recipe_json_to_vllm_config.py
index 5fd7b7729d26..bd947b07451c 100644
--- a/tools/recipes/recipe_json_to_vllm_config.py
+++ b/tools/recipes/recipe_json_to_vllm_config.py
@@ -2,8 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Convert a vLLM Recipes per-hardware JSON rendering into:
+"""Convert a vLLM Recipes per-hardware JSON rendering into:
 
   1) config.yml - native `vllm serve --config` YAML
   2) env.sh     - environment variables required by the recipe
@@ -38,6 +37,7 @@
 The converter intentionally targets a single `vllm serve` process. If the
 recipe rendering is multi-node, PD-disaggregated, or another multi-process
 deployment, it exits instead of silently generating an incomplete config.
+
 """
 
 from __future__ import annotations
@@ -551,8 +551,7 @@ def merge_value(dst: dict[str, Any], path: list[str], value: Any) -> None:
 
 
 def normalize_key(raw_key: str) -> list[str]:
-    """
-    Convert the CLI key to config-file spelling.
+    """Convert the CLI key to config-file spelling.
 
     Only the top-level CLI option name gets underscore -> dash normalization.
     Nested JSON field names after a dot are preserved.
diff --git a/tools/recipes/serve_with_recipe.sh b/tools/recipes/serve_with_recipe.sh
new file mode 100755
index 000000000000..cad1c06f3f3b
--- /dev/null
+++ b/tools/recipes/serve_with_recipe.sh
@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+CONVERTER="${SCRIPT_DIR}/recipe_json_to_vllm_config.py"
+OUTPUT_DIR="$(pwd -P)"
+CONFIG_PATH="${OUTPUT_DIR}/config.yml"
+ENV_PATH="${OUTPUT_DIR}/env.sh"
+
+args=("$@")
+hardware=""
+
+for ((i = 0; i < ${#args[@]}; i++)); do
+    case "${args[$i]}" in
+        --hardware)
+            if ((i + 1 < ${#args[@]})); then
+                hardware="${args[$((i + 1))]}"
+            fi
+            ;;
+        --hardware=*)
+            hardware="${args[$i]#*=}"
+            ;;
+    esac
+done
+
+if [[ "${hardware,,}" == "xeon6" ]]; then
+    args+=(--detect-hardware)
+fi
+
+echo "==> Generating vLLM configuration"
+python3 "${CONVERTER}" "${args[@]}"
+
+echo
+echo "==> Loading recipe environment"
+# shellcheck disable=SC1090
+source "${ENV_PATH}"
+
+# Released vLLM images keep recipe-relative assets such as examples/*.jinja
+# under /vllm-workspace. Launch from there when available while keeping the
+# generated config and environment files in the caller's output directory.
+if [[ -d /vllm-workspace/examples ]]; then
+    cd /vllm-workspace
+fi
+
+echo
+echo "==> Starting vLLM"
+exec vllm serve --config "${CONFIG_PATH}"
diff --git a/tools/report_build_time_ninja.py b/tools/report_build_time_ninja.py
index fe3f352fe153..3b43fc988ef8 100644
--- a/tools/report_build_time_ninja.py
+++ b/tools/report_build_time_ninja.py
@@ -197,7 +197,6 @@ def GetExtension(target, extra_patterns):
 
 def SummarizeEntries(entries, extra_step_types):
     """Print a summary of the passed in list of Target objects."""
-
     # Create a list that is in order by time stamp and has entries for the
     # beginning and ending of each build step (one time stamp may have multiple
     # entries due to multiple steps starting/stopping at exactly the same time).
diff --git a/tools/vllm-rocm/pin_rocm_dependencies.py b/tools/vllm-rocm/pin_rocm_dependencies.py
index 7d90d66692ad..f4c5041b0ee1 100644
--- a/tools/vllm-rocm/pin_rocm_dependencies.py
+++ b/tools/vllm-rocm/pin_rocm_dependencies.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Pin vLLM dependencies to exact versions of custom ROCm wheels.
+"""Pin vLLM dependencies to exact versions of custom ROCm wheels.
 
 This script modifies vLLM's requirements files to replace version constraints
 with exact versions of custom-built ROCm wheels (torch, triton, torchvision, amdsmi).
@@ -18,12 +17,12 @@
 
 
 def extract_version_from_wheel(wheel_name: str) -> str:
-    """
-    Extract version from wheel filename.
+    """Extract version from wheel filename.
 
     Example:
         torch-2.9.0a0+git1c57644-cp312-cp312-linux_x86_64.whl -> 2.9.0a0+git1c57644
         triton-3.4.0-cp312-cp312-linux_x86_64.whl -> 3.4.0
+
     """
     # Wheel format:
     #    {distribution}-{version}(-{build tag})?-{python}-{abi}-{platform}.whl
@@ -38,11 +37,11 @@ def extract_version_from_wheel(wheel_name: str) -> str:
 
 
 def get_custom_wheel_versions(install_dir: str) -> dict[str, str]:
-    """
-    Read /install directory and extract versions of custom wheels.
+    """Read /install directory and extract versions of custom wheels.
 
     Returns:
         Dict mapping package names to exact versions
+
     """
     install_path = Path(install_dir)
     if not install_path.exists():
@@ -92,8 +91,7 @@ def get_custom_wheel_versions(install_dir: str) -> dict[str, str]:
 
 
 def pin_dependencies_in_requirements(requirements_path: str, versions: dict[str, str]):
-    """
-    Insert custom wheel pins at the TOP of requirements file.
+    """Insert custom wheel pins at the TOP of requirements file.
 
     This ensures that when setup.py processes the file line-by-line,
     custom wheels (torch, triton, etc.) are encountered FIRST, before
diff --git a/vllm/__init__.py b/vllm/__init__.py
index 968d1a143b16..16101ca245f6 100644
--- a/vllm/__init__.py
+++ b/vllm/__init__.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""vLLM: a high-throughput and memory-efficient inference engine for LLMs"""
+"""vLLM: a high-throughput and memory-efficient inference engine for LLMs."""
 
 # The version.py should be independent library, and we always import the
 # version library first.  Such assumption is critical for some customization.
diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py
index 7a2e3d8a87b8..ba87ba894268 100644
--- a/vllm/_aiter_ops.py
+++ b/vllm/_aiter_ops.py
@@ -4,15 +4,15 @@
 import functools
 import os
 from collections.abc import Callable
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
+import torch.distributed as dist
 from torch._ops import OpOverload
 
 import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import PlaceholderModule
 from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
     rocm_aiter_sparse_attn_indexer,
@@ -24,16 +24,80 @@
 
 logger = init_logger(__name__)
 
-try:
-    import pandas as pd
-except ImportError:
-    pd = PlaceholderModule("pandas")
-
 # fp8_dtype is not cached.
 # on ROCm the fp8_dtype always calls is_fp8_fnuz
 # which is a host op, so we cache it once here.
 FP8_DTYPE = current_platform.fp8_dtype()
 _HIPB_MM_INITIALIZED_DEVICES: set[int] = set()
+_FP8_E4M3_DTYPES = {
+    dtype
+    for dtype in (
+        getattr(torch, "float8_e4m3fn", None),
+        getattr(torch, "float8_e4m3fnuz", None),
+    )
+    if dtype is not None
+}
+KB = 1024
+MB = 1024 * KB
+
+
+def _is_fp8_e4m3_tensor(t: torch.Tensor | None) -> bool:
+    return t is not None and t.dtype in _FP8_E4M3_DTYPES
+
+
+def _get_or_create_aiter_qr_rmsnorm_comm(
+    device_comm: Any,
+) -> tuple[int, int, bool] | None:
+    """Create the AITER QR communicator used by fused QR+RMSNorm.
+
+    vLLM's normal QuickReduce communicator owns vLLM C++ QR buffers. The fused
+    QR+RMSNorm kernel added in AITER needs an AITER DeviceComms instance, so we
+    lazily create one on the same TP CPU group and keep it on the vLLM device
+    communicator for reuse.
+    """
+    existing = getattr(device_comm, "_aiter_qr_rmsnorm_comm", None)
+    if existing is False:
+        return None
+    if existing is not None:
+        return existing
+
+    try:
+        import aiter
+    except Exception:
+        device_comm._aiter_qr_rmsnorm_comm = False
+        return None
+
+    if not hasattr(aiter, "qr_all_reduce_rmsnorm"):
+        device_comm._aiter_qr_rmsnorm_comm = False
+        return None
+
+    group = getattr(device_comm, "cpu_group", None)
+    if group is None:
+        device_comm._aiter_qr_rmsnorm_comm = False
+        return None
+
+    world_size = dist.get_world_size(group=group)
+    rank = dist.get_rank(group=group)
+    if world_size not in (2, 4, 8):
+        device_comm._aiter_qr_rmsnorm_comm = False
+        return None
+
+    qr_max_size_mb = envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB
+    qr_max_size = None if qr_max_size_mb is None else qr_max_size_mb * MB
+
+    try:
+        ptr = aiter.init_custom_qr(rank, world_size, qr_max_size)
+        handle = aiter.qr_get_handle(ptr)
+        handles = [None] * world_size
+        dist.all_gather_object(handles, handle, group=group)
+        aiter.qr_open_handles(ptr, handles)
+    except Exception:
+        device_comm._aiter_qr_rmsnorm_comm = False
+        return None
+
+    comm = (ptr, world_size, envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16)
+    device_comm._aiter_qr_rmsnorm_comm = comm
+    return comm
 
 
 def _ensure_hipb_mm_extension_initialized() -> None:
@@ -173,27 +237,42 @@ def is_aiter_found_and_supported_on_rdna4() -> bool:
     return False
 
 
-@functools.cache
-def _load_gemm_tuned_configs(
-    q_dtype_w: torch.dtype, csv_path: str
-) -> set[tuple[int, int, int]]:
-    try:
-        df = pd.read_csv(csv_path).drop_duplicates()
-        df = df[df["q_dtype_w"] == str(q_dtype_w)]
-        return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int)))
-    except Exception:
-        return set()
+def _triton_gemm_config_is_tuned(config_name: str, N: int, K: int) -> bool:
+    from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config
+
+    # any M shape under 1024 should be tuned
+    M_TUNE_PROBE = 128
+
+    return get_gemm_config(config_name, M_TUNE_PROBE, N, K)[1]
 
 
-def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool:
-    configs = _load_gemm_tuned_configs(q_dtype_w, csv_path)
+def _ck_gemm_shape_is_tuned(
+    N: int, K: int, q_dtype_w: torch.dtype, csv_attr: str
+) -> bool:
     l_m = (
         [1, 2, 4]
         + list(range(8, 513, 8))
         + [1024, 1536]
         + [2**i for i in range(11, 19)]
     )
-    return any((N, K, M) in configs for M in l_m)
+    try:
+        from aiter.ops.gemm_op_a8w8 import (
+            AITER_CONFIGS,
+            get_GEMM_config_with_quant_type,
+        )
+
+        csv_path = getattr(AITER_CONFIGS, csv_attr)
+        return any(
+            get_GEMM_config_with_quant_type(M, N, K, q_dtype_w, csv_path) is not None
+            for M in l_m
+        )
+    except (AttributeError, ImportError, OSError):
+        logger.warning_once(
+            "Could not read aiter CK GEMM configs from AITER_CONFIGS.%s; "
+            "treating all shapes as untuned.",
+            csv_attr,
+        )
+        return False
 
 
 def if_aiter_supported(func: Callable) -> Callable:
@@ -575,10 +654,11 @@ def _rocm_aiter_mla_decode_fwd_impl(
     reduce_indptr: torch.Tensor | None = None,
     reduce_final_map: torch.Tensor | None = None,
     reduce_partial_map: torch.Tensor | None = None,
+    causal: bool = True,
 ) -> None:
     from aiter.mla import mla_decode_fwd
 
-    kwargs: dict[str, float | torch.Tensor | None] = {
+    kwargs: dict[str, float | torch.Tensor | None | bool] = {
         "sm_scale": sm_scale,
         "logit_cap": logit_cap,
     }
@@ -588,6 +668,8 @@ def _rocm_aiter_mla_decode_fwd_impl(
         kwargs["q_scale"] = q_scale
         kwargs["kv_scale"] = kv_scale
 
+    kwargs["causal"] = causal
+
     if work_meta_data is not None:
         assert work_indptr is not None, (
             "work_indptr must be provided with work_meta_data"
@@ -1018,6 +1100,53 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
     ca = aiter_ar.aiter_ca
     use_1stage = aiter_ar.use_1stage_fused_ar_rms(input_)
 
+    if not use_1stage:
+        from vllm.distributed import get_tp_group
+
+        device_comm = get_tp_group().device_communicator
+        qr_comm = getattr(device_comm, "qr_comm", None)
+        hidden_dim = input_.shape[-1]
+        row_size = hidden_dim * input_.element_size()
+        fused_qr_rmsnorm_ok = (
+            qr_comm is not None
+            and not getattr(qr_comm, "disabled", True)
+            and hasattr(qr_comm, "should_quick_allreduce")
+            and qr_comm.should_quick_allreduce(input_)
+            and input_.shape == residual.shape
+            and input_.dtype in (torch.bfloat16, torch.float16)
+            and input_.dtype == residual.dtype
+            and input_.dtype == weight.dtype
+            and weight.dim() == 1
+            and weight.numel() == hidden_dim
+            and input_.numel() % hidden_dim == 0
+            and row_size > 0
+            and row_size <= 32 * KB
+            and (32 * KB) % row_size == 0
+        )
+        if fused_qr_rmsnorm_ok:
+            assert qr_comm is not None
+            aiter_qr_comm = _get_or_create_aiter_qr_rmsnorm_comm(device_comm)
+            if aiter_qr_comm is not None:
+                import aiter
+
+                ptr, _, cast_bf2half = aiter_qr_comm
+                quant_level = qr_comm._get_qr_quant_level(input_)
+                out = torch.empty_like(input_)
+                residual_out = torch.empty_like(residual)
+                aiter.qr_all_reduce_rmsnorm(
+                    ptr,
+                    input_,
+                    residual,
+                    residual_out,
+                    out,
+                    weight,
+                    epsilon,
+                    hidden_dim,
+                    quant_level,
+                    cast_bf2half,
+                )
+                return out, residual_out
+
     result = ca.custom_fused_ar_rms(
         input_,
         residual,
@@ -1698,10 +1827,103 @@ def _rocm_aiter_fp8_attn_fake(
     )
 
 
+def _mhc_delayed_pre_tail(
+    gemm_out: torch.Tensor,
+    sqrsum: torch.Tensor,
+    residual: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    pre_mix: torch.Tensor | None,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Finish a delayed mHC pre once AITER has produced the projection.
+
+    Shared by the plain and the post-fused entry points, which differ only in
+    which kernel computes ``gemm_out`` / ``sqrsum`` and the residual they
+    project. ``mhc_pre_big_fuse`` also writes a collapse against its own
+    pre-mix, which the delayed formulation cannot use and so discards.
+    Returns ``(post_mix, comb_mix, layer_input, next_pre_mix)``.
+    """
+    from aiter.ops.mhc import mhc_pre_big_fuse
+
+    num_tokens, hc_mult, hidden_size = residual.shape
+    device = residual.device
+    with torch.device(device):
+        post_mix = torch.empty(
+            num_tokens, hc_mult, 1, dtype=torch.float32, device=device
+        )
+        comb_mix = torch.empty(
+            num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device
+        )
+        unused_collapse = torch.empty(
+            num_tokens, hidden_size, dtype=torch.bfloat16, device=device
+        )
+        mhc_pre_big_fuse(
+            post_mix,
+            comb_mix,
+            unused_collapse,
+            gemm_out,
+            sqrsum,
+            hc_scale,
+            hc_base,
+            residual,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+        )
+
+    next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton(
+        gemm_out,
+        sqrsum,
+        hc_scale,
+        hc_base,
+        hc_mult,
+        hc_mult * hidden_size,
+        rms_eps,
+        hc_pre_eps,
+    )
+    if pre_mix is None:
+        # Model entry selects residual stream zero.
+        layer_input = residual[:, 0]
+    else:
+        layer_input = torch.ops.vllm.hc_collapse_triton(residual, pre_mix)
+    return post_mix, comb_mix, layer_input, next_pre_mix
+
+
 # Global flag to ensure ops are registered only once
 _OPS_REGISTERED = False
 
 
+def _sync_aiter_situv2_moe_env() -> None:
+    """Mirror the SiTUv2 MoE toggle into AITER's a4w4 dispatch env.
+
+    AITER selects afp8 vs afp4 activation kernels via AITER_SITUV2_A8W4 /
+    AITER_SITUV2_A4W4 (see ROCm/aiter fused_moe.py, A8W4 checked first).
+    When VLLM_ROCM_USE_AITER_MOE_SITUV2 is enabled we route to a4w4
+    (afp4_wfp4_fp4 kernels) and clear any legacy AITER_SITUV2_A8W4 override.
+
+    Requires AITER with ROCm/aiter#4463 (first tagged in v0.1.20): a4w4
+    dispatch plus kimik3_a4w4_{un,}tuned_fmoe.csv. Older AITER still runs
+    afp4 FlyDSL but falls back to heuristic configs, which is not the
+    tuned a4w4 path this recipe is meant to use.
+    """
+    import os
+
+    import vllm.envs as envs
+
+    if envs.VLLM_ROCM_USE_AITER_MOE_SITUV2:
+        os.environ["AITER_SITUV2_A4W4"] = "1"
+        os.environ.pop("AITER_SITUV2_A8W4", None)
+    else:
+        os.environ.pop("AITER_SITUV2_A4W4", None)
+
+
 class rocm_aiter_ops:
     """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM.
 
@@ -1722,9 +1944,10 @@ class rocm_aiter_ops:
         VLLM_ROCM_USE_AITER_MHA: Controls MHA ops including flash_attn_varlen.
         VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: Controls Triton unified attention.
         VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
+        VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.
         VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
         VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
-        VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: Controls a8w4 SiTU fused MoE variant.
+        VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4).
         VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM.
 
     Note:
@@ -1761,6 +1984,7 @@ class rocm_aiter_ops:
         - MLA decode: mla_decode_fwd
         - Quantization: per_tensor_quant, per_token_quant, group_fp8_quant
         - Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale
+
     """
 
     _MOE_DISPATCH_POLICY: int | None = None
@@ -1788,10 +2012,12 @@ def get_moe_dispatch_policy(cls) -> int:
     _FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
     _FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
     _LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
+    # TODO: Consolidate under _LINEAR_ENABLED
+    _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
     # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
     _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
     _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
-    _MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4
+    _MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2
     # TODO: Consolidate under _LINEAR_ENABLED
     _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
     # Lazily probed: whether aiter.topk_softmax supports the
@@ -1800,8 +2026,7 @@ def get_moe_dispatch_policy(cls) -> int:
 
     @classmethod
     def refresh_env_variables(cls):
-        """
-        Since the environment variables are assigned when the module is imported,
+        """Since the environment variables are assigned when the module is imported,
         This is a helper function to reload all the env variables from
         the environment variables.
         for example, after monkey patching the env variables in the unit test,
@@ -1818,15 +2043,16 @@ def refresh_env_variables(cls):
         cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
         cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
         cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
+        cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
         cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
         cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
-        cls._MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4
+        cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2
+        _sync_aiter_situv2_moe_env()
         cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
 
     @staticmethod
     def get_aiter_activation_type(activation_str: str) -> "ActivationType | None":
-        """
-        Given an activation type as a string, returns the corresponding aiter ActivationType enum.
+        """Given an activation type as a string, returns the corresponding aiter ActivationType enum.
         Supported activation types: "no", "none", "silu", "gelu", "swiglu".
         Returns None if the mapping fails.
 
@@ -1835,6 +2061,7 @@ def get_aiter_activation_type(activation_str: str) -> "ActivationType | None":
 
         Returns:
             Aiter ActivationType enum value, or None if not found.
+
         """
         # Import only locally, since aiter may not always be available.
         try:
@@ -1858,8 +2085,7 @@ def get_aiter_activation_type(activation_str: str) -> "ActivationType | None":
 
     @staticmethod
     def get_aiter_quant_type(quant_type_str: str) -> "QuantType | None":
-        """
-        Given a quantization type as a string, returns the corresponding aiter QuantType enum.
+        """Given a quantization type as a string, returns the corresponding aiter QuantType enum.
         Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128".
         Returns None if the mapping fails.
 
@@ -1868,6 +2094,7 @@ def get_aiter_quant_type(quant_type_str: str) -> "QuantType | None":
 
         Returns:
             Aiter QuantType enum value, or None if not found.
+
         """
         try:
             from aiter import QuantType
@@ -1933,10 +2160,10 @@ def is_fusion_moe_shared_experts_enabled(cls) -> bool:
 
     @classmethod
     @if_aiter_supported
-    def is_fused_moe_situv2_a8w4_enabled(cls) -> bool:
-        # _MOE_SITUV2_A8W4 is a variant of aiter fused moe, so aiter
+    def is_fused_moe_situv2_enabled(cls) -> bool:
+        # _MOE_SITUV2 is a variant of aiter fused moe, so aiter
         # fused moe must be enabled as well.
-        return cls.is_fused_moe_enabled() and cls._MOE_SITUV2_A8W4
+        return cls.is_fused_moe_enabled() and cls._MOE_SITUV2
 
     @classmethod
     @if_aiter_supported
@@ -1989,6 +2216,26 @@ def fuse_sigmoid_in_kernel(cls, aiter_topK_meta_data: object) -> bool:
     def is_mla_enabled(cls) -> bool:
         return cls._AITER_ENABLED and cls._MLA_ENABLED
 
+    @classmethod
+    @if_aiter_supported
+    def mla_decode_supports_non_causal(cls) -> bool:
+        """Whether installed aiter.mla.mla_decode_fwd accepts `causal`.
+
+        Added in aiter v0.1.20. A missing argument is a hard error rather
+        than a boolean the caller can use to pick another MLA backend.
+        """
+        import inspect
+
+        from aiter.mla import mla_decode_fwd
+
+        if "causal" not in inspect.signature(mla_decode_fwd).parameters:
+            raise RuntimeError(
+                "ROCM_AITER_MLA requires aiter.mla.mla_decode_fwd(..., causal=). "
+                "The installed aiter is causal-only; upgrade aiter rather than "
+                "falling back to another MLA backend."
+            )
+        return True
+
     @classmethod
     @if_aiter_supported
     def is_mha_enabled(cls) -> bool:
@@ -2004,6 +2251,106 @@ def is_custom_all_reduce_enabled(cls) -> bool:
     def is_shuffle_kv_cache_enabled(cls) -> bool:
         return cls._SHUFFLE_KV_CACHE_ENABLED
 
+    @classmethod
+    @if_aiter_supported
+    def fused_qknorm_idxrqknorm_enabled(
+        cls,
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor | None = None,
+        v_scale: torch.Tensor | None = None,
+    ) -> bool:
+        """Whether MiniMax-M3 can use AITER's consolidated fused QK-norm.
+
+        Requires AITER to be installed and enabled. bf16 (``auto``) always
+        qualifies; fp8 e4m3 also needs K/V scales. Other cache dtypes fall
+        back to vLLM's fused kernel.
+        """
+        if not cls._AITER_ENABLED:
+            return False
+        if kv_cache_dtype == "auto":
+            return True
+        if kv_cache_dtype in ("fp8", "fp8_e4m3"):
+            return k_scale is not None and v_scale is not None
+        return False
+
+    @classmethod
+    def fused_qknorm_idxrqknorm(
+        cls,
+        qkv: torch.Tensor,
+        q_norm_weight: torch.Tensor,
+        k_norm_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        positions: torch.Tensor,
+        num_heads: int,
+        num_kv_heads: int,
+        rotary_dim: int,
+        eps: float,
+        slot_mapping: torch.Tensor,
+        kv_cache_k: torch.Tensor,
+        kv_cache_v: torch.Tensor,
+        q_out: torch.Tensor,
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor | None = None,
+        v_scale: torch.Tensor | None = None,
+        index_q_norm_weight: torch.Tensor | None = None,
+        index_k_norm_weight: torch.Tensor | None = None,
+        num_index_heads: int = 0,
+        index_cache: torch.Tensor | None = None,
+        index_q_out: torch.Tensor | None = None,
+        index_slot_mapping: torch.Tensor | None = None,
+        skip_index_branch: bool = False,
+    ) -> None:
+        """Run consolidated MiniMax-M3 QK-norm fusion.
+
+        Callers must check ``fused_qknorm_idxrqknorm_enabled`` first. Runtime
+        failures from the AITER op are deliberately propagated.
+        """
+        if kv_cache_dtype == "auto":
+            aiter_kv_cache_dtype = "auto"
+            aiter_k_scale = None
+            aiter_v_scale = None
+        elif kv_cache_dtype in ("fp8", "fp8_e4m3"):
+            aiter_kv_cache_dtype = "fp8_e4m3_static"
+            aiter_k_scale = k_scale
+            aiter_v_scale = v_scale
+        else:
+            raise ValueError(
+                "AITER fused QK-norm requires kv_cache_dtype 'auto', 'fp8', "
+                f"or 'fp8_e4m3', got {kv_cache_dtype!r}"
+            )
+        aiter_index_cache_dtype = "fp8" if _is_fp8_e4m3_tensor(index_cache) else "auto"
+
+        from aiter import fused_qknorm_idxrqknorm
+
+        fused_qknorm_idxrqknorm(
+            qkv,
+            q_norm_weight,
+            k_norm_weight,
+            cos_sin_cache,
+            positions,
+            num_heads,
+            num_kv_heads,
+            rotary_dim,
+            eps,
+            index_q_norm_weight=index_q_norm_weight,
+            index_k_norm_weight=index_k_norm_weight,
+            num_index_heads=num_index_heads,
+            slot_mapping=slot_mapping,
+            kv_cache_k=kv_cache_k,
+            kv_cache_v=kv_cache_v,
+            index_cache=index_cache,
+            block_size=16,
+            q_out=q_out,
+            index_q_out=index_q_out,
+            index_slot_mapping=index_slot_mapping,
+            kv_cache_dtype=aiter_kv_cache_dtype,
+            index_cache_dtype=aiter_index_cache_dtype,
+            k_scale=aiter_k_scale,
+            v_scale=aiter_v_scale,
+            asm_layout=True,
+            skip_index_branch=skip_index_branch,
+        )
+
     @classmethod
     @if_aiter_supported
     def is_triton_unified_attn_enabled(cls) -> bool:
@@ -2038,7 +2385,7 @@ def is_linear_hipbmm_enabled(cls) -> bool:
     def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool:
         from vllm.platforms.rocm import on_gfx950
 
-        return cls._AITER_ENABLED and on_gfx950()
+        return cls._AITER_ENABLED and cls._FP4_GEMM_DYNAMIC_QUANT_ASM and on_gfx950()
 
     @classmethod
     @if_aiter_supported
@@ -2165,13 +2512,15 @@ def register_ops_once() -> None:
             if not current_platform.is_rocm():
                 return
 
-            from vllm.platforms.rocm import on_gfx11
+            from vllm.platforms.rocm import on_gfx11, on_gfx950
 
-            if on_gfx11() and not _OPS_REGISTERED:
+            # This op has self-contained Triton/C++ implementations on gfx11
+            # and gfx950.  Only its optional top-k fast path comes from aiter.
+            if (on_gfx11() or on_gfx950()) and not _OPS_REGISTERED:
                 direct_register_custom_op(
                     op_name="rocm_aiter_sparse_attn_indexer",
                     op_func=rocm_aiter_sparse_attn_indexer,
-                    mutates_args=["topk_indices_buffer"],
+                    mutates_args=["topk_indices_buffer", "candidate_blocks"],
                     fake_impl=rocm_aiter_sparse_attn_indexer_fake,
                     dispatch_key=current_platform.dispatch_key,
                 )
@@ -2355,7 +2704,7 @@ def register_ops_once() -> None:
             direct_register_custom_op(
                 op_name="rocm_aiter_sparse_attn_indexer",
                 op_func=rocm_aiter_sparse_attn_indexer,
-                mutates_args=["topk_indices_buffer"],
+                mutates_args=["topk_indices_buffer", "candidate_blocks"],
                 fake_impl=rocm_aiter_sparse_attn_indexer_fake,
                 dispatch_key=current_platform.dispatch_key,
             )
@@ -2828,6 +3177,7 @@ def mla_decode_fwd(
         reduce_indptr: torch.Tensor | None = None,
         reduce_final_map: torch.Tensor | None = None,
         reduce_partial_map: torch.Tensor | None = None,
+        causal: bool = True,
     ):
         torch.ops.vllm.rocm_aiter_mla_decode_fwd(
             q,
@@ -2848,6 +3198,7 @@ def mla_decode_fwd(
             reduce_indptr=reduce_indptr,
             reduce_final_map=reduce_final_map,
             reduce_partial_map=reduce_partial_map,
+            causal=causal,
         )
 
     @staticmethod
@@ -3060,6 +3411,115 @@ def do_qk_norm_rope_kvcache_update(
             rotary_dim=kernel_rotary_dim,
         )
 
+    @staticmethod
+    def fused_qk_norm_mrope_and_cache(
+        qkv: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        positions: torch.Tensor,
+        num_heads_q: int,
+        num_heads_k: int,
+        head_dim: int,
+        is_neox: bool,
+        mrope_section: list[int],
+        is_interleaved: bool,
+        rms_norm_eps: float,
+        q_out: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        k_scale: torch.Tensor,
+        v_scale: torch.Tensor,
+        block_size: int,
+        x: int,
+        rotary_dim: int = 0,
+    ) -> None:
+        from aiter.ops.fused_qk_norm_mrope_cache_quant import (
+            fused_qk_norm_mrope_3d_cache_pts_quant_shuffle,
+        )
+
+        fused_qk_norm_mrope_3d_cache_pts_quant_shuffle(
+            qkv,
+            q_weight,
+            k_weight,
+            cos_sin_cache,
+            positions,
+            qkv.size(0),
+            num_heads_q,
+            num_heads_k,
+            num_heads_k,
+            head_dim,
+            is_neox,
+            mrope_section,
+            is_interleaved,
+            rms_norm_eps,
+            q_out,
+            k_cache,
+            v_cache,
+            slot_mapping,
+            k_scale,
+            v_scale,
+            None,
+            None,
+            False,
+            False,
+            block_size,
+            x,
+            rotary_dim,
+        )
+
+    @staticmethod
+    def do_qk_norm_mrope_kvcache_update(
+        qkv: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        positions: torch.Tensor,
+        num_heads_q: int,
+        num_heads_k: int,
+        head_dim: int,
+        is_neox: bool,
+        mrope_section: list[int],
+        is_interleaved: bool,
+        rms_norm_eps: float,
+        q_out: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        k_scale: torch.Tensor,
+        v_scale: torch.Tensor,
+        kv_cache_dtype: str,
+        rotary_dim: int = 0,
+    ) -> None:
+        if kv_cache_dtype.startswith("fp8"):
+            key_cache = key_cache.view(current_platform.fp8_dtype())
+            value_cache = value_cache.view(current_platform.fp8_dtype())
+
+        rocm_aiter_ops.fused_qk_norm_mrope_and_cache(
+            qkv=qkv,
+            q_weight=q_weight,
+            k_weight=k_weight,
+            cos_sin_cache=cos_sin_cache,
+            positions=positions,
+            num_heads_q=num_heads_q,
+            num_heads_k=num_heads_k,
+            head_dim=head_dim,
+            is_neox=is_neox,
+            mrope_section=mrope_section,
+            is_interleaved=is_interleaved,
+            rms_norm_eps=rms_norm_eps,
+            q_out=q_out,
+            k_cache=key_cache,
+            v_cache=value_cache,
+            slot_mapping=slot_mapping,
+            k_scale=k_scale,
+            v_scale=v_scale,
+            block_size=key_cache.shape[1],
+            x=16 // key_cache.element_size(),
+            rotary_dim=rotary_dim,
+        )
+
     @staticmethod
     def triton_rope_and_cache(
         query: torch.Tensor,
@@ -3104,7 +3564,7 @@ def batched_gemm_a16wfp4(
         X: torch.Tensor,
         W: torch.Tensor,
         w_scale: torch.Tensor,
-        Y: torch.Tensor,
+        Y: torch.Tensor | None = None,
         transpose_bm: bool | None = False,
         prequant: bool | None = False,
         y_scale: torch.Tensor | None = None,
@@ -3165,46 +3625,14 @@ def group_fp8_quant(
         )
 
     @staticmethod
+    @functools.cache
     def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool:
         if not current_platform.is_rocm():
             return False
-        from vllm.platforms.rocm import on_gfx950, on_rdna4
-
-        gfx950_tuned = {
-            (1024, 8192),
-            (2112, 7168),
-            (3072, 1536),
-            (32768, 8192),
-            (4096, 7168),
-            (4608, 7168),
-            (512, 7168),
-            (7168, 2048),
-            (7168, 256),
-            (8192, 1024),
-            (8192, 32768),
-        }
-        rdna4_tuned = gfx950_tuned | {
-            (2048, 2048),
-            (2624, 6144),
-            (3072, 6144),
-            (3584, 512),
-            (4096, 512),
-            (6144, 1536),
-            (6144, 2048),
-            (7168, 2304),
-            (7168, 16384),
-            (7168, 18432),
-            (8192, 8192),
-            (16384, 1536),
-            (24576, 1536),
-            (32768, 512),
-            (36864, 7168),
-        }
-        if on_rdna4():
-            return (n, k) in rdna4_tuned
-        if on_gfx950():
-            return (n, k) in gfx950_tuned
-        return False
+        try:
+            return _triton_gemm_config_is_tuned("GEMM-A8W8_BLOCKSCALE", n, k)
+        except (AssertionError, ImportError):
+            return False
 
     @staticmethod
     @if_aiter_supported
@@ -3219,47 +3647,31 @@ def is_triton_blockscale_bpreshuffle_gemm_tuned(n: int, k: int) -> bool:
         return is_tuned
 
     @staticmethod
-    def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k: int) -> bool:
-        return (n, k) in [
-            (8192, 4096),
-            (1280, 8192),
-            (16384, 53248),
-            (106496, 16384),
-            (57344, 8192),
-            (8192, 2048),
-            (2560, 8192),
-            (10240, 8192),
-            (16384, 16384),
-            (8192, 28672),
-            (28672, 8192),
-            (18432, 16384),
-            (8192, 1024),
-            (7168, 8192),
-            (5120, 8192),
-            (8192, 8192),
-            (8192, 7168),
-            (14336, 8192),
-            (8192, 14336),
-            (8192, 3584),
-        ]
+    @functools.cache
+    def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k_bytes: int) -> bool:
+        if not current_platform.is_rocm():
+            return False
+        # Input = weight.shape[1] (bytes); *4 for fp4 and aiter config 2*k
+        try:
+            return _triton_gemm_config_is_tuned(
+                "GEMM-AFP4WFP4_PRESHUFFLED", n, 4 * k_bytes
+            )
+        except (AssertionError, ImportError):
+            return False
 
     @staticmethod
+    @functools.cache
     def is_shuffled_per_token_w8a8_gemm_tuned(
         N: int, K: int, q_dtype_w: torch.dtype
     ) -> bool:
-        import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
-
-        csv_path = (
-            aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE
+        return _ck_gemm_shape_is_tuned(
+            N, K, q_dtype_w, "AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE"
         )
-        return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
 
     @staticmethod
+    @functools.cache
     def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool:
-        import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
-
-        csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE
-        return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
+        return _ck_gemm_shape_is_tuned(N, K, q_dtype_w, "AITER_CONFIG_GEMM_A8W8_FILE")
 
     @staticmethod
     def shuffle_weight(
@@ -3275,8 +3687,7 @@ def shuffle_weight_a16w4(
         nLane: int,
         gate_up: bool,
     ) -> "torch.Tensor":
-        """
-        Shuffles the weight tensor into (A16W4) layout for AITER kernels.
+        """Shuffles the weight tensor into (A16W4) layout for AITER kernels.
 
         Args:
             tensor: The input weight tensor to be shuffled.
@@ -3285,6 +3696,7 @@ def shuffle_weight_a16w4(
 
         Returns:
             torch.Tensor: The shuffled tensor.
+
         """
         from aiter.ops.shuffle import shuffle_weight_a16w4
 
@@ -3296,8 +3708,7 @@ def shuffle_scale_a16w4(
         num_experts: int,
         gate_up: bool,
     ) -> "torch.Tensor":
-        """
-        Shuffles the scale tensor into (A16W4) layout for AITER kernels.
+        """Shuffles the scale tensor into (A16W4) layout for AITER kernels.
 
         Args:
             tensor: The input scale tensor to be shuffled.
@@ -3306,6 +3717,7 @@ def shuffle_scale_a16w4(
 
         Returns:
             torch.Tensor: The shuffled scale tensor.
+
         """
         from aiter.ops.shuffle import shuffle_scale_a16w4
 
@@ -3322,8 +3734,7 @@ def shuffle_scale(tensor: torch.Tensor) -> torch.Tensor:
     def shuffle_weights(
         *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
     ) -> tuple[torch.Tensor, ...]:
-        """
-        Applies shuffle_weight function from AITER to each
+        """Applies shuffle_weight function from AITER to each
         input tensor and returns them.
 
         Rearranges (shuffles) the input tensor/s
@@ -3336,6 +3747,7 @@ def shuffle_weights(
 
         Returns:
         A Tuple of shuffled tensors.
+
         """
         from aiter.ops.shuffle import shuffle_weight
 
@@ -3385,8 +3797,7 @@ def flash_attn_varlen_func(
         out: torch.Tensor | None = None,
         sink_ptr: torch.Tensor | None = None,
     ):
-        """
-        Flash attention with variable length sequences.
+        """Flash attention with variable length sequences.
 
         This function is NOT wrapped with @is_aiter_supported decorator
         to allow explicit backend selection via attention_config to work
@@ -3455,8 +3866,7 @@ def pa_fwd_asm(
         V_QScale: torch.Tensor,
         out_: torch.Tensor,
     ):
-        """
-        Paged attention forward pass using assembly kernel.
+        """Paged attention forward pass using assembly kernel.
 
         This function is NOT wrapped with @is_aiter_supported decorator
         to allow explicit backend selection via attention_config to work
@@ -3498,8 +3908,7 @@ def paged_attention_common(
         out_: torch.Tensor,
         kv_cache_dtype: str,
     ):
-        """
-        Paged attention common function.
+        """Paged attention common function.
 
         This function is NOT wrapped with @is_aiter_supported decorator
         to allow explicit backend selection via attention_config to work
@@ -3543,8 +3952,7 @@ def mhc_pre(
         norm_weight: torch.Tensor | None = None,
         norm_eps: float = 0.0,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for mHC pre block.
+        """Forward pass for mHC pre block.
 
         Args:
             residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
@@ -3563,6 +3971,7 @@ def mhc_pre(
             post_mix: shape (..., hc_mult), dtype torch.float32
             comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
             layer_input: shape (..., hidden_size), dtype torch.bfloat16
+
         """
         from aiter.ops.mhc import mhc_pre
 
@@ -3636,6 +4045,170 @@ def mhc_pre(
             layer_input.view(*outer_shape, hidden_size),
         )
 
+    @staticmethod
+    def mhc_pre_delayed(
+        residual: torch.Tensor,
+        fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        pre_mix: torch.Tensor | None = None,
+        sublayer_out: torch.Tensor | None = None,
+        post_layer_mix: torch.Tensor | None = None,
+        comb_res_mix: torch.Tensor | None = None,
+        residual_out: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """MHC pre using the pre-mix carried from the previous sublayer.
+
+        Same gates as :meth:`mhc_pre`, but the stream collapse applies the
+        caller's ``pre_mix`` instead of the one computed here, and the one
+        computed here is returned for the next sublayer seam. AITER has no
+        single kernel for that shape, so drive its two stages directly:
+        ``mhc_pre_big_fuse`` still produces the post and comb gates (including
+        every sinkhorn iteration), while the pre gate is recovered from the
+        same split-k GEMM output and the collapse is done by the Triton
+        kernel. ``mhc_pre_big_fuse`` also writes a collapse we do not use;
+        that redundant store is the price of not having a native delayed
+        kernel, and is small next to the ~140 launches it replaces.
+
+        Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix``
+        also applies the preceding post block. AITER then projects the new
+        residual in the same kernel that computes it
+        (``mhc_fused_post_pre_gemm_sqrsum``), which saves one launch; every
+        later stage is identical either way. That new residual is written to
+        ``residual_out``, which the caller must supply in that case; it is an
+        output buffer rather than a return value so the op never has to hand
+        back one of its own inputs on the unfused path.
+
+        Returns:
+            post_mix: shape (..., hc_mult, 1), dtype torch.float32
+            comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
+            layer_input: shape (..., hidden_size), dtype torch.bfloat16
+            next_pre_mix: shape (..., hc_mult), dtype torch.float32
+
+        """
+        from aiter.ops.mhc import (
+            get_mhc_fused_post_pre_config,
+            get_mhc_pre_splitk,
+            mhc_fused_post_pre_gemm_sqrsum,
+            mhc_pre_gemm_sqrsum,
+        )
+
+        assert residual.dtype == torch.bfloat16
+        assert fn.dtype == torch.float32
+        assert hc_scale.dtype == torch.float32
+        assert hc_base.dtype == torch.float32
+
+        hc_mult = residual.shape[-2]
+        hidden_size = residual.shape[-1]
+        hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
+        hc_hidden_size = hc_mult * hidden_size
+
+        assert fn.shape == (hc_mult3, hc_hidden_size)
+        assert hc_scale.shape == (3,)
+        assert hc_base.shape == (hc_mult3,)
+
+        outer_shape = residual.shape[:-2]
+        residual_flat = residual.view(-1, hc_mult, hidden_size)
+        num_tokens = residual_flat.shape[0]
+        device = residual_flat.device
+
+        if num_tokens == 0:
+            return (
+                torch.empty(0, hc_mult, 1, dtype=torch.float32, device=device),
+                torch.empty(0, hc_mult, hc_mult, dtype=torch.float32, device=device),
+                torch.empty(0, hidden_size, dtype=torch.bfloat16, device=device),
+                torch.empty(0, hc_mult, dtype=torch.float32, device=device),
+            )
+
+        pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult)
+
+        fuse_post = sublayer_out is not None
+        # AITER's Python wrappers allocate without explicit device arguments.
+        with torch.device(device):
+            if fuse_post:
+                # Keyed on the per-stream K, i.e. hidden_size, not hc_mult * it.
+                splitk, tile_m, tile_n, tile_k = get_mhc_fused_post_pre_config(
+                    num_tokens, hidden_size
+                )
+            else:
+                splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size)
+            # AITER pads the GEMM output to a multiple of 32 columns.
+            gemm_pad = torch.empty(
+                splitk,
+                num_tokens,
+                (hc_mult3 + 31) // 32 * 32,
+                dtype=torch.float32,
+                device=device,
+            )
+            gemm_out = gemm_pad[:, :, :hc_mult3]
+            sqrsum = torch.empty(splitk, num_tokens, dtype=torch.float32, device=device)
+            if fuse_post:
+                # Restated for the type checker, which cannot narrow the
+                # optional arguments through the fuse_post flag.
+                assert sublayer_out is not None
+                assert post_layer_mix is not None and comb_res_mix is not None
+                assert residual_out is not None
+                projected = residual_out.view(-1, hc_mult, hidden_size)
+                mhc_fused_post_pre_gemm_sqrsum(
+                    gemm_out,
+                    sqrsum,
+                    projected,
+                    sublayer_out.view(-1, hidden_size),
+                    residual_flat,
+                    # The kernel wants the post mix squeezed to (tokens, hc_mult).
+                    post_layer_mix.view(-1, hc_mult, 1).squeeze(-1),
+                    comb_res_mix.view(-1, hc_mult, hc_mult),
+                    fn,
+                    tile_m,
+                    tile_n,
+                    tile_k,
+                    0,
+                )
+            else:
+                projected = residual_flat
+                mhc_pre_gemm_sqrsum(gemm_out, sqrsum, projected, fn, tile_k, 0)
+
+        post_mix, comb_mix, layer_input, next_pre_mix = _mhc_delayed_pre_tail(
+            gemm_out,
+            sqrsum,
+            projected,
+            hc_scale,
+            hc_base,
+            pre_mix_flat,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+        )
+
+        return (
+            post_mix.view(*outer_shape, hc_mult, 1),
+            comb_mix.view(*outer_shape, hc_mult, hc_mult),
+            layer_input.view(*outer_shape, hidden_size),
+            next_pre_mix.view(*outer_shape, hc_mult),
+        )
+
+    @staticmethod
+    def mhc_fused_post_pre_delayed_prefers_unfused(num_tokens: int) -> bool:
+        """True when AITER's own heuristic favours separate post and pre.
+
+        Mirrors the ``fused_m_upper_bound`` table in ``aiter.ops.mhc``: the
+        fused GEMM wins for decode-sized batches and loses for large ones.
+        """
+        from aiter.jit.utils.chip_info import get_gfx_runtime
+
+        return num_tokens >= {
+            "gfx950": 1024,
+            "gfx942": 128,
+            "gfx1250": 1024,
+        }.get(get_gfx_runtime(), 1024)
+
     @staticmethod
     def hc_head(
         hs_flat: torch.Tensor,
@@ -3809,3 +4382,4 @@ def mhc_fused_post_pre(
 
 
 rocm_aiter_ops.register_ops_once()
+_sync_aiter_situv2_moe_env()
diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
index db20adb9946c..f08be7b66f5a 100644
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from enum import IntEnum
+from functools import cache
 from typing import TYPE_CHECKING, Literal
 
 import torch
@@ -39,8 +40,7 @@ def create_fp4_scale_tensor(
     device: torch.device,
     is_sf_swizzled_layout: bool,
 ) -> torch.Tensor:
-    """
-    Allocate the output scale tensor for scaled_fp4_quant.
+    """Allocate the output scale tensor for scaled_fp4_quant.
 
     When is_sf_swizzled_layout=True, we use rounded values to store the
     swizzled scales. Due to the requirement of the Tensor Core, the minimum
@@ -73,8 +73,7 @@ def create_fp4_output_tensors(
     is_sf_swizzled_layout: bool,
     padded_n: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Allocate both output tensors for scaled_fp4_quant:
+    """Allocate both output tensors for scaled_fp4_quant:
     (quantized_output, output_scale).
 
     Must match the C++ scaled_fp4_quant_func allocation exactly when
@@ -265,9 +264,18 @@ def vocab_parallel_embedding(
     Args:
         input_ids: ``[num_tokens]`` int32 or int64 token ids.
         weight: ``[num_embeddings_per_partition, embedding_dim]`` shard.
+        org_vocab_start_index: First original vocab id owned by this rank.
+        org_vocab_end_index: One past the last original vocab id owned by
+            this rank.
+        num_org_vocab_padding: Padding rows between this rank's original and
+            added vocab shards.
+        added_vocab_start_index: First added vocab id owned by this rank.
+        added_vocab_end_index: One past the last added vocab id owned by
+            this rank.
 
     Returns:
         ``[num_tokens, embedding_dim]`` partial embeddings.
+
     """
     out = torch.empty(
         input_ids.shape[0],
@@ -399,6 +407,7 @@ def apply_repetition_penalties(
         prompt_mask: A boolean tensor indicating which tokens appear in the prompt.
         output_mask: A boolean tensor indicating which tokens appear in the output.
         repetition_penalties: The repetition penalties of shape (num_seqs, ).
+
     """
     if logits.is_cuda and logits.is_contiguous():
         apply_repetition_penalties_cuda(
@@ -723,26 +732,6 @@ def moe_gptq_gemm_rdna3(
     )
 
 
-if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
-
-    @register_fake("_C::allspark_w8a16_gemm")
-    def _allspark_w8a16_gemm_fake(
-        a: torch.Tensor,
-        b_qweight: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_qzeros: torch.Tensor | None,
-        n: torch.SymInt,
-        group_size: torch.SymInt,
-        sm_count: torch.SymInt,
-        sm_version: torch.SymInt,
-        CUBLAS_M_THRESHOLD: torch.SymInt,
-        has_zp: bool,
-        n32k16_reorder: bool,
-    ) -> torch.Tensor:
-        m = a.size(0)
-        return torch.empty((m, n), device=a.device, dtype=a.dtype)
-
-
 # cutlass
 def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool:
     return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability)
@@ -787,8 +776,7 @@ def cutlass_scaled_mm(
     out_dtype: torch.dtype,
     bias: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """
-    `cutlass_scaled_mm` implements a fused version of
+    """`cutlass_scaled_mm` implements a fused version of
         `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
     where scale_a * a and scale_b * b are implemented using numpy-style
     broadcasting.
@@ -840,11 +828,11 @@ def cutlass_scaled_mm_azp(
     azp: torch.Tensor | None = None,
     bias: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """
-    Args:
-        azp_adj: In the per-tensor case, this should include the azp.
-            Always per-channel.
-        azp: Only set in the per-token case. Per-token if set.
+    """Args:
+    azp_adj: In the per-tensor case, this should include the azp.
+        Always per-channel.
+    azp: Only set in the per-token case. Per-token if set.
+
     """
     assert b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0
     assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
@@ -883,8 +871,7 @@ def get_cutlass_moe_mm_data(
     blockscale_offsets: torch.Tensor | None = None,
     is_gated: bool = True,
 ):
-    """
-    Prepare data necessary to perform CUTLASS grouped matrix multiplications
+    """Prepare data necessary to perform CUTLASS grouped matrix multiplications
     used in CUTLASS-based fused MoE.
 
     The function takes in topk_ids (token-expert mapping) and uses it to
@@ -943,8 +930,7 @@ def get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
 
 
 def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor):
-    """
-    Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor.
+    """Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor.
     This is used in MoE to permute the input tensor before performing grouped matrix multiplications.
     """
     num_tokens_permuted = dst2src_map.shape[0]
@@ -967,8 +953,7 @@ def get_cutlass_batched_moe_mm_data(
     n: int,
     k: int,
 ):
-    """
-    Prepare data necessary to perform CUTLASS grouped matrix multiplications
+    """Prepare data necessary to perform CUTLASS grouped matrix multiplications
     used in CUTLASS-based fused MoE.
 
     The function takes in expert_num_tokens (token count per expert) and
@@ -1006,8 +991,7 @@ def cutlass_moe_mm(
     per_act_token: bool,
     per_out_ch: bool,
 ):
-    """
-    A single grouped matrix multiplication used in CUTLASS-based fused MoE.
+    """A single grouped matrix multiplication used in CUTLASS-based fused MoE.
     The function executes fp8-quantized OUT = AB matrix multiplication.
 
     - expert_offsets: Indices that mark at which token index each expert begins
@@ -1044,8 +1028,7 @@ def cutlass_fp4_moe_mm(
     expert_offsets: torch.Tensor,
     sf_offsets: torch.Tensor,
 ):
-    """
-    An FP4 Blockscaled Group Gemm that takes in  a_tensors, b_tensors and runs
+    """An FP4 Blockscaled Group Gemm that takes in  a_tensors, b_tensors and runs
     the gemms for each combination based on the specified problem sizes.
 
     This is used as the MoE gemm during NVFP4 Quantized MoERunner forward.
@@ -1083,8 +1066,7 @@ def cutlass_mxfp4_moe_mm(
     expert_offsets: torch.Tensor,
     sf_offsets: torch.Tensor,
 ):
-    """
-    An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
+    """An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
 
     Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
     - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
@@ -1458,8 +1440,7 @@ def cutlass_w4a8_moe_mm(
     group_scale_strides: torch.Tensor,
     maybe_schedule: str | None = None,
 ) -> None:
-    """
-    Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the
+    """Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the
     W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8)
     and both per-channel + per-token scaling in the epilogue.
 
@@ -1482,13 +1463,20 @@ def cutlass_w4a8_moe_mm(
             Cumulative token offsets
         problem_sizes:
             Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher.
-        a/b/c/group_scale_strides:
-            Strides describing the memory layout of the input tensors.
+        a_strides:
+            Strides describing the memory layout of a_tensors.
+        b_strides:
+            Strides describing the memory layout of b_tensors.
+        c_strides:
+            Strides describing the memory layout of out_tensors.
+        group_scale_strides:
+            Strides describing the memory layout of b_group_scales.
         maybe_schedule:
             Optional override to choose a specific kernel or epilogue schedule.
 
     Returns:
         out_tensors updated in-place with the dequantized INT4xFP8 grouped GEMM result.
+
     """
     return torch.ops._C.cutlass_w4a8_moe_mm(
         out_tensors,
@@ -1529,8 +1517,7 @@ def scaled_fp4_quant(
     backend: str = "none",
     padded_n: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Quantize input tensor to FP4 and return quantized tensor and scale.
+    """Quantize input tensor to FP4 and return quantized tensor and scale.
 
     This function quantizes the last dimension of the given tensor `input`. For
     every 16 consecutive elements, a single dynamically computed scaling factor
@@ -1553,6 +1540,7 @@ def scaled_fp4_quant(
         tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every
             two values are packed into a uint8 and float8_e4m3 scaling factors
             in the sizzled layout.
+
     """
     assert not current_platform.is_rocm()
     assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}."
@@ -1607,17 +1595,19 @@ def scaled_fp4_experts_quant(
     blockscale_offsets: torch.Tensor,
     topk: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Quantize input tensor to NVFP4 and return quantized tensor and scale, for
+    """Quantize input tensor to NVFP4 and return quantized tensor and scale, for
     packed MoE Inputs.
+
     Args:
         input_tensor: The input tensor to be quantized to NVFP4
         input_global_scale: A scalar scaling factor for the entire tensor.
         expert_offsets: The expert offsets tensor
         blockscale_offsets: The blockscale offsets tensor
+        topk: The number of experts each token is routed to
     Outputs:
         output: The quantized tensor in NVFP4
         output_scales: The blockscale tensor in FP8-E4M3
+
     """
     assert not current_platform.is_rocm()
     assert input_tensor.ndim == 2, (
@@ -1669,8 +1659,7 @@ def silu_and_mul_scaled_fp4_experts_quant(
     blockscale_offsets: torch.Tensor,
     topk: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations.
+    """Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations.
 
     Args:
         input_tensor: The input tensor with gate || up layout [m_topk, k*2]
@@ -1681,6 +1670,7 @@ def silu_and_mul_scaled_fp4_experts_quant(
     Outputs:
         output: The quantized tensor in NVFP4 [m_topk, k/2]
         output_scales: The blockscale tensor in FP8-E4M3
+
     """
     assert not current_platform.is_rocm()
     assert input_tensor.ndim == 2, (
@@ -1734,8 +1724,7 @@ def mxfp4_experts_quant(
     n_experts: int,
     topk: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Quantize input tensor to MXFP4 for packed MoE inputs.
+    """Quantize input tensor to MXFP4 for packed MoE inputs.
     Uses 32-element blocks with E8M0 (power-of-two) scale factors.
     MXFP4 has no global scale - only block-level E8M0 scale factors.
 
@@ -1748,6 +1737,7 @@ def mxfp4_experts_quant(
     Returns:
         output: [m_topk, k//2] packed E2M1 values (uint8)
         output_scales: E8M0 blockscales in swizzled layout (uint8 view)
+
     """
     assert not current_platform.is_rocm()
     assert input_tensor.ndim == 2
@@ -1793,8 +1783,7 @@ def silu_and_mul_mxfp4_experts_quant(
     n_experts: int,
     topk: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
+    """Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
     MXFP4 has no global scale - only block-level E8M0 scale factors.
     """
     assert not current_platform.is_rocm()
@@ -1840,8 +1829,7 @@ def scaled_fp8_quant(
     output: torch.Tensor | None = None,
     group_shape: tuple[int, int] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Quantize input tensor to FP8 and return quantized tensor and scale.
+    """Quantize input tensor to FP8 and return quantized tensor and scale.
 
     This function supports both static and dynamic quantization: If you
     provide the scale, it will use static scaling and if you omit it,
@@ -1863,6 +1851,8 @@ def scaled_fp8_quant(
             of the output to at least this value.
         use_per_token_if_dynamic: Whether to do per_tensor or per_token
             in the dynamic quantization case.
+        output: Optional tensor to write the quantized result into. A new
+            tensor is allocated when omitted.
         group_shape: Optional tuple (group_m, group_n) specifying the group
             shape for static quantization. Use -1 for "full extent" (e.g.,
             (-1, -1) for per-tensor, (-1, 1) for per-channel, etc.)
@@ -1871,6 +1861,7 @@ def scaled_fp8_quant(
     Returns:
         tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and
             scaling factor.
+
     """
     # This code assumes batch_dim and num_tokens are flattened
     assert input.ndim == 2
@@ -1900,90 +1891,6 @@ def scaled_fp8_quant(
     return output, scale
 
 
-# gptq allspark
-def allspark_repack_weight(
-    qweight: torch.Tensor,
-    scale: torch.Tensor,
-    zero_point: torch.Tensor | None = None,
-    has_zp: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format
-    for Ampere W8A16 Fused Gemm kernel
-
-    Args:
-        qweight: uint8 weight tensor, original k x n format.
-        scale: fp16/bf16 weight scale tensor, 1 x n format.
-        zero_point: fp16/bf16 weight zero_point tensor, 1 x n format.
-            Must be provided for asymmetric quantization.
-        has_zp: if use symmetric quantization, has_zp = False.
-            if use asymmetric quantization, has_zp = True.
-
-    Returns:
-        tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] :
-            rearranged weight, scale, and optionally zero_point.
-    """
-    K = qweight.shape[0]
-    N = qweight.shape[1]
-    N_32align = (N + 32 - 1) // 32 * 32
-
-    qweight_reorder = torch.empty(
-        (N_32align, K), device=qweight.device, dtype=qweight.dtype
-    )
-    scale_reorder = torch.empty((1, N_32align), device=scale.device, dtype=scale.dtype)
-    zero_point_reorder = None
-    if has_zp:
-        assert zero_point is not None, (
-            "zero_point must be provided for asymmetric quantization."
-        )
-        zero_point_reorder = torch.empty(
-            (1, N_32align), device=zero_point.device, dtype=zero_point.dtype
-        )
-
-    torch.ops._C.rearrange_kn_weight_as_n32k16_order(
-        qweight,
-        scale,
-        zero_point,
-        has_zp,
-        qweight_reorder,
-        scale_reorder,
-        zero_point_reorder,
-        K,
-        N,
-        N_32align,
-    )
-
-    return qweight_reorder, scale_reorder, zero_point_reorder
-
-
-def allspark_w8a16_gemm(
-    a: torch.Tensor,
-    b_qweight: torch.Tensor,
-    b_scales: torch.Tensor,
-    b_qzeros: torch.Tensor | None,
-    n: int,
-    group_size: int,
-    sm_count: int,
-    sm_version: int,
-    CUBLAS_M_THRESHOLD: int,
-    has_zp: bool,
-    n32k16_reorder: bool,
-) -> torch.Tensor:
-    return torch.ops._C.allspark_w8a16_gemm(
-        a,
-        b_qweight,
-        b_scales,
-        b_qzeros,
-        n,
-        group_size,
-        sm_count,
-        sm_version,
-        CUBLAS_M_THRESHOLD,
-        has_zp,
-        n32k16_reorder,
-    )
-
-
 # int8
 def scaled_int8_quant(
     input: torch.Tensor,
@@ -1991,8 +1898,7 @@ def scaled_int8_quant(
     azp: torch.Tensor | None = None,
     symmetric: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-    """
-    Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
+    """Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
 
     Args:
         input: The input tensor to be quantized to int8.
@@ -2004,6 +1910,7 @@ def scaled_int8_quant(
 
     Returns:
         Output int8 tensor, scales, and optionally azp.
+
     """
     if current_platform.is_xpu():
         # XPU has no _C int8 quant op; use the torch.compile reference.
@@ -2261,7 +2168,21 @@ def moe_align_block_size(
     experts_ids: torch.Tensor,
     num_tokens_post_pad: torch.Tensor,
     expert_map: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
 ) -> None:
+    if current_platform.is_xpu():
+        if scatter_idx is not None:
+            raise NotImplementedError("scatter_idx is not supported on XPU")
+        torch.ops._moe_C.moe_align_block_size(
+            topk_ids,
+            num_experts,
+            block_size,
+            sorted_token_ids,
+            experts_ids,
+            num_tokens_post_pad,
+            expert_map,
+        )
+        return
     torch.ops._moe_C.moe_align_block_size(
         topk_ids,
         num_experts,
@@ -2270,6 +2191,7 @@ def moe_align_block_size(
         experts_ids,
         num_tokens_post_pad,
         expert_map,
+        scatter_idx,
     )
 
 
@@ -2457,8 +2379,7 @@ def grouped_topk(
     bias: torch.Tensor,
     scoring_func: int = 0,
 ):
-    """
-    Perform grouped top-k routing for mixture of experts.
+    """Perform grouped top-k routing for mixture of experts.
 
     Args:
         scores: Raw inputs (logits if scoring_func=1, scores if scoring_func=0)
@@ -2469,6 +2390,7 @@ def grouped_topk(
         routed_scaling_factor: Scaling factor for routing weights
         bias: Bias tensor (e_score_correction_bias). Always fused in kernel.
         scoring_func: 0=none (no activation), 1=sigmoid
+
     """
     if not (current_platform.is_cuda() or current_platform.is_xpu()):
         raise NotImplementedError(
@@ -2828,6 +2750,8 @@ def concat_and_cache_mla_grouped(
     block_size: int,
     block_stride: int,
     entry_stride: int,
+    kv_scales: torch.Tensor | None = None,
+    kv_cache_dtype: str = "auto",
 ) -> None:
     torch.ops._C_cache_ops.concat_and_cache_mla_grouped(
         kv_c,
@@ -2837,6 +2761,8 @@ def concat_and_cache_mla_grouped(
         block_size,
         block_stride,
         entry_stride,
+        kv_scales,
+        kv_cache_dtype,
     )
 
 
@@ -2901,8 +2827,7 @@ def swap_blocks(
     block_size_in_bytes: int,
     block_mapping: torch.Tensor,
 ) -> None:
-    """
-    Copy specific blocks from one tensor to another.
+    """Copy specific blocks from one tensor to another.
 
     This method assumes each of the two input tensors is composed of
     consecutive contiguous blocks, of size block_size_in_bytes.
@@ -2930,8 +2855,7 @@ def swap_blocks_batch(
     sizes: torch.Tensor,
     is_src_access_order_any: bool = False,
 ) -> None:
-    """
-    Batch version of swap_blocks: submit all copies in a single driver call.
+    """Batch version of swap_blocks: submit all copies in a single driver call.
 
     Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i]
     of sizes[i] bytes. All three tensors must be CPU tensors with the
@@ -3004,6 +2928,9 @@ def cp_gather_and_upconvert_fp8_kv_cache(
     workspace_starts: torch.Tensor,
     batch_size: int,
     seq_starts: torch.Tensor | None = None,
+    host_cache: torch.Tensor | None = None,
+    host_row_ids: torch.Tensor | None = None,
+    device_row_ids: torch.Tensor | None = None,
 ) -> None:
     """Gather and upconvert FP8 KV cache to BF16 workspace.
 
@@ -3014,9 +2941,21 @@ def cp_gather_and_upconvert_fp8_kv_cache(
         workspace_starts: Workspace start offsets [num_reqs]
         batch_size: Number of requests
         seq_starts: Optional source sequence offsets [num_reqs]
+        host_cache: Optional pinned host rows used for non-resident entries
+        host_row_ids: Host source row for each remapped cache row
+        device_row_ids: Device source row, or -1, for each remapped cache row
+
     """
     torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache(
-        src_cache, dst, block_table, workspace_starts, batch_size, seq_starts
+        src_cache,
+        dst,
+        block_table,
+        workspace_starts,
+        batch_size,
+        seq_starts,
+        host_cache,
+        host_row_ids,
+        device_row_ids,
     )
 
 
@@ -3035,6 +2974,7 @@ def cp_gather_and_upconvert_nvfp4_kv_cache(
         block_table: Block indices [num_reqs, max_blocks]
         workspace_starts: Workspace start offsets [num_reqs]
         batch_size: Number of requests
+
     """
     torch.ops._C_cache_ops.cp_gather_and_upconvert_nvfp4_kv_cache(
         src_cache,
@@ -3056,6 +2996,7 @@ def concat_mla_q(
         ql_nope: Query nope component [num_tokens, num_heads, nope_dim]
         q_pe: Query rope component [num_tokens, num_heads, rope_dim]
         q_out: Output tensor [num_tokens, num_heads, nope_dim + rope_dim]
+
     """
     torch.ops._C_cache_ops.concat_mla_q(ql_nope, q_pe, q_out)
 
@@ -3355,6 +3296,7 @@ class CPUQuantMethod(IntEnum):
     FP8_W8A16 = 2
     INT4_W4A8 = 3
     MXFP4 = 4
+    FP8_W8A8 = 5
 
 
 if hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
@@ -3393,6 +3335,7 @@ def fused_experts_cpu(
     alpha: float | None = None,
     limit: float | None = None,
     is_vnni: bool = True,
+    a1_scale: torch.Tensor | None = None,
 ) -> None:
     torch.ops._C.fused_experts_cpu(
         out,
@@ -3406,6 +3349,7 @@ def fused_experts_cpu(
         w2_scale,
         w1_zero,
         w2_zero,
+        a1_scale,
         block_size,
         w1_bias,
         w2_bias,
@@ -3535,6 +3479,46 @@ def fp8_scaled_mm_cpu(
     )
 
 
+# FP8 W8A8 CPU kernels
+@cache
+def cpu_has_amx_fp8() -> bool:
+    """Whether this CPU has native AMX-FP8 MMA support."""
+    if not hasattr(torch.ops._C, "cpu_has_amx_fp8"):
+        return False
+    return bool(torch.ops._C.cpu_has_amx_fp8())
+
+
+if hasattr(torch.ops._C, "fp8_scaled_mm_with_quant"):
+
+    @register_fake("_C::fp8_scaled_mm_with_quant")
+    def fp8_scaled_mm_with_quant_fake(
+        act: torch.Tensor,
+        act_scales: torch.Tensor | None,
+        channelwise: bool,
+        weight: torch.Tensor,
+        weight_scales: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        M = act.reshape(-1, act.size(-1)).size(0)
+        N = weight.size(0) * weight.size(-1)
+        return torch.empty((M, N), dtype=output_dtype, device=act.device)
+
+
+def fp8_scaled_mm_with_quant(
+    act: torch.Tensor,
+    act_scales: torch.Tensor | None,
+    channelwise: bool,
+    weight: torch.Tensor,
+    weight_scales: torch.Tensor,
+    bias: torch.Tensor | None,
+    output_dtype: torch.dtype,
+) -> torch.Tensor:
+    return torch.ops._C.fp8_scaled_mm_with_quant(
+        act, act_scales, channelwise, weight, weight_scales, bias, output_dtype
+    )
+
+
 def chunk_gated_delta_rule_cpu(
     query: torch.Tensor,
     key: torch.Tensor,
@@ -3778,8 +3762,7 @@ def onednn_scaled_int8_quant(
     azp: torch.Tensor | None = None,
     symmetric: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-    """
-    Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
+    """Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
 
     Args:
         input: The input tensor to be quantized to int8.
@@ -3791,6 +3774,7 @@ def onednn_scaled_int8_quant(
 
     Returns:
         Output int8 tensor, scales, and optionally azp.
+
     """
     output = torch.empty_like(input, dtype=torch.int8)
     token_num = input.numel() // input.shape[-1]
@@ -4650,8 +4634,7 @@ def safeFusedQuantizeNv(
     xh_e4m3: torch.Tensor,
     global_scale: torch.Tensor,
 ) -> None:
-    """
-    Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place
+    """Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place
     rather than returning them, to prevent torch 2.12+ errors that outputs of custom
     operators may not alias any inputs to the custom operator.
     """
@@ -4660,8 +4643,7 @@ def safeFusedQuantizeNv(
 
 
 def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
-    """
-    Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832)
+    """Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832)
     kernels. Note that these kernels exploit the recursive properties of
     Sylvester Hadamards, and therefore do not require transform weight data
 
@@ -4674,6 +4656,7 @@ def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
 
     Returns:
         value after transformation
+
     """
     return torch.ops._C.hadacore_transform(x, inplace)
 
diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py
index c1f2bb6969b9..d25ef6b258ba 100644
--- a/vllm/_xpu_ops.py
+++ b/vllm/_xpu_ops.py
@@ -6,6 +6,7 @@
 
 import torch
 from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func
+from vllm_xpu_kernels.rotary import apply_rotary_emb
 
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.logger import init_logger
@@ -249,6 +250,24 @@ def _xpu_ops_deepseek_scaling_rope_fake(
     return query, key
 
 
+def _xpu_apply_rotary_emb_impl(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    is_neox_style: bool,
+) -> torch.Tensor:
+    return apply_rotary_emb(x, cos, sin, is_neox_style)
+
+
+def _xpu_apply_rotary_emb_fake(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    is_neox_style: bool,
+) -> torch.Tensor:
+    return torch.empty_like(x)
+
+
 def _xpu_fp8_bmm_impl(
     a: torch.Tensor,
     b: torch.Tensor,
@@ -286,6 +305,7 @@ def _xpu_fp8_bmm_impl(
         This implementation centralizes access to
         ``torch.ops._xpu_C.fp8_bmm``. Both scales must be contiguous, while
         ``a`` and ``b`` may be non-contiguous views.
+
     """
     return torch.ops._xpu_C.fp8_bmm(a, b, out_dtype, a_scale, b_scale, bias)
 
@@ -430,6 +450,7 @@ def _xpu_deepseek_fused_indexer_q_rope_fp8_impl(
         index_weights_out: (T, H) float32 output, preallocated;
             = index_weights * q_scale * softmax_scale * head_scale (the
             per-(token, head) q_scale is folded in here). [written]
+
     """
     torch.ops._xpu_C.deepseek_fused_indexer_q_rope_fp8(
         index_q,
@@ -473,6 +494,7 @@ def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl(
         index_weights_out: (T, H) float32 output, preallocated;
             = index_weights * softmax_scale * head_scale (no q_scale folded;
             per-block scales live in index_q_scale). [written]
+
     """
     torch.ops._xpu_C.deepseek_fused_indexer_q_rope_mxfp4(
         index_q,
@@ -1254,6 +1276,12 @@ def register_ops_once() -> None:
                 dispatch_key=current_platform.dispatch_key,
             )
 
+            direct_register_custom_op(
+                op_name="xpu_apply_rotary_emb",
+                op_func=_xpu_apply_rotary_emb_impl,
+                fake_impl=_xpu_apply_rotary_emb_fake,
+            )
+
             direct_register_custom_op(
                 op_name="xpu_mxfp8_quantize",
                 op_func=_xpu_mxfp8_quantize_impl,
diff --git a/vllm/assets/base.py b/vllm/assets/base.py
index 5ca9de4076ad..a0bd6b1eb26d 100644
--- a/vllm/assets/base.py
+++ b/vllm/assets/base.py
@@ -20,8 +20,7 @@ def get_cache_dir() -> Path:
 
 @lru_cache
 def get_vllm_public_assets(filename: str, s3_prefix: str | None = None) -> Path:
-    """
-    Download an asset file from `s3://vllm-public-assets`
+    """Download an asset file from `s3://vllm-public-assets`
     and return the path to the downloaded file.
     """
     asset_directory = get_cache_dir() / "vllm_public_assets"
diff --git a/vllm/assets/image.py b/vllm/assets/image.py
index a91eb7d4b67d..4f87fd1a4700 100644
--- a/vllm/assets/image.py
+++ b/vllm/assets/image.py
@@ -34,9 +34,7 @@ class ImageAsset:
     name: ImageAssetName
 
     def get_path(self, ext: str) -> Path:
-        """
-        Return s3 path for given image.
-        """
+        """Return s3 path for given image."""
         return get_vllm_public_assets(
             filename=f"{self.name}.{ext}", s3_prefix=VLM_IMAGES_DIR
         )
@@ -51,9 +49,7 @@ def pil_image_ext(self, ext: str) -> Image.Image:
 
     @property
     def image_embeds(self) -> torch.Tensor:
-        """
-        Image embeddings, only used for testing purposes with llava 1.5.
-        """
+        """Image embeddings, only used for testing purposes with llava 1.5."""
         image_path = self.get_path("pt")
         return torch.load(image_path, map_location="cpu", weights_only=True)
 
diff --git a/vllm/assets/video.py b/vllm/assets/video.py
index e06d033e0be1..bad21663751c 100644
--- a/vllm/assets/video.py
+++ b/vllm/assets/video.py
@@ -21,8 +21,7 @@ def _sample_frame_indices(total_frames: int, num_frames: int) -> npt.NDArray:
 
 @lru_cache
 def download_video_asset(filename: str) -> str:
-    """
-    Download and open an image from huggingface
+    """Download and open an image from huggingface
     repo: raushan-testing-hf/videos-test
     """
     video_directory = get_cache_dir() / "video-example-data"
@@ -141,8 +140,7 @@ def metadata(self) -> dict[str, Any]:
         return ret
 
     def get_audio(self, sampling_rate: float | None = None) -> npt.NDArray:
-        """
-        Read audio data from the video asset, used in Qwen2.5-Omni examples.
+        """Read audio data from the video asset, used in Qwen2.5-Omni examples.
 
         See also: examples/generate/multimodal/qwen2_5_omni/only_thinker.py
         """
diff --git a/vllm/benchmarks/datasets/create_txt_slices_dataset.py b/vllm/benchmarks/datasets/create_txt_slices_dataset.py
index 8d19cc25ff3d..1bcd93c8ae11 100644
--- a/vllm/benchmarks/datasets/create_txt_slices_dataset.py
+++ b/vllm/benchmarks/datasets/create_txt_slices_dataset.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Convert a plain-text file (local path or URL) into a JSONL dataset
+"""Convert a plain-text file (local path or URL) into a JSONL dataset
 compatible with ``CustomDataset`` (``--dataset-name custom``), by 
 randomly slicing the tokenized text into prompts.
 
@@ -67,7 +66,6 @@ def create_txt_slices_jsonl(
 ) -> None:
     """Read *input_path*, slice it into prompts, and write JSONL to
     *output_path*."""
-
     tokenizer = AutoTokenizer.from_pretrained(
         tokenizer_name, trust_remote_code=trust_remote_code
     )
diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py
index 3ca9fd670806..8f9e87cf8331 100644
--- a/vllm/benchmarks/datasets/datasets.py
+++ b/vllm/benchmarks/datasets/datasets.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This module defines a framework for sampling benchmark requests from various
+"""This module defines a framework for sampling benchmark requests from various
 datasets. Each dataset subclass of BenchmarkDataset must implement sample
 generation. Supported dataset types include:
   - ShareGPT
@@ -75,9 +74,7 @@
 
 @dataclass
 class SampleRequest:
-    """
-    Represents a single inference request for benchmarking.
-    """
+    """Represents a single inference request for benchmarking."""
 
     prompt: str | list[str] | list[int] | list[list[int]] | list[dict]
     prompt_len: int
@@ -112,8 +109,7 @@ def __init__(
         disable_shuffle: bool = False,
         **kwargs,
     ) -> None:
-        """
-        Initialize the BenchmarkDataset with an optional dataset path and random
+        """Initialize the BenchmarkDataset with an optional dataset path and random
         seed.
 
         Args:
@@ -121,6 +117,9 @@ def __init__(
                 indicates that a default or random dataset might be used.
             random_seed (int): Seed value for reproducible shuffling or
                 sampling. Defaults to DEFAULT_SEED.
+            disable_shuffle (bool): If True, keep the dataset in its original
+                order instead of shuffling it.
+
         """
         self.dataset_path = dataset_path
         # Set the random seed, ensuring that a None value is replaced with the
@@ -134,8 +133,7 @@ def apply_multimodal_chat_transformation(
         prompt: str,
         mm_content: MultiModalDataDict | dict | list[dict] | None = None,
     ) -> list[dict]:
-        """
-        Transform a prompt and optional multimodal content into a chat format.
+        """Transform a prompt and optional multimodal content into a chat format.
         This method is used for chat models that expect a specific conversation
         format.
         """
@@ -152,14 +150,14 @@ def apply_multimodal_chat_transformation(
         return [{"role": "user", "content": content}]
 
     def load_data(self) -> None:
-        """
-        Load data from the dataset path into self.data.
+        """Load data from the dataset path into self.data.
 
         This method must be overridden by subclasses since the method to load
         data will vary depending on the dataset format and source.
 
         Raises:
             NotImplementedError: If a subclass does not implement this method.
+
         """
         # TODO (jenniferzhao): add support for downloading data
         raise NotImplementedError("load_data must be implemented in subclasses.")
@@ -169,8 +167,7 @@ def get_random_lora_request(
         max_loras: int | None = None,
         lora_path: str | None = None,
     ) -> LoRARequest | None:
-        """
-        Optionally select a random LoRA request.
+        """Optionally select a random LoRA request.
 
         This method is used when LoRA parameters are provided.  It randomly
         selects a LoRA based on max_loras.
@@ -184,6 +181,7 @@ def get_random_lora_request(
         Returns:
             A new [`LoRARequest`][vllm.lora.request.LoRARequest]
             (or `None` if not applicable).
+
         """
         if max_loras is None or lora_path is None:
             return None
@@ -203,8 +201,7 @@ def get_round_robin_lora_request(
         max_loras: int | None = None,
         lora_path: str | None = None,
     ) -> LoRARequest | None:
-        """
-        Optionally select a LoRA request using deterministic round-robin.
+        """Optionally select a LoRA request using deterministic round-robin.
 
         This method cycles through LoRA IDs in order based on the request
         index, providing reproducible LoRA assignment.
@@ -219,6 +216,7 @@ def get_round_robin_lora_request(
         Returns:
             A new [`LoRARequest`][vllm.lora.request.LoRARequest]
             (or `None` if not applicable).
+
         """
         if max_loras is None or lora_path is None:
             return None
@@ -239,8 +237,7 @@ def get_lora_request(
         lora_path: str | None = None,
         lora_assignment: str = "random",
     ) -> LoRARequest | None:
-        """
-        Select a LoRA request using the specified assignment strategy.
+        """Select a LoRA request using the specified assignment strategy.
 
         Args:
             index (int): The request index (used for round-robin).
@@ -252,6 +249,7 @@ def get_lora_request(
         Returns:
             A new [`LoRARequest`][vllm.lora.request.LoRARequest]
             (or `None` if not applicable).
+
         """
         if lora_assignment == "round-robin":
             return BenchmarkDataset.get_round_robin_lora_request(
@@ -270,8 +268,7 @@ def sample(
         no_oversample: bool = False,
         **kwargs,
     ) -> list[SampleRequest]:
-        """
-        Abstract method to generate sample requests from the dataset.
+        """Abstract method to generate sample requests from the dataset.
 
         Subclasses must override this method to implement dataset-specific logic
         for generating a list of SampleRequest objects.
@@ -281,10 +278,13 @@ def sample(
                 for processing the dataset's text.
             num_requests (int): The number of sample requests to generate.
             request_id_prefix (str): The prefix of request_id.
+            no_oversample (bool): If True, do not oversample the dataset to
+                reach num_requests.
 
         Returns:
             list[SampleRequest]: A list of sample requests generated from the
             dataset.
+
         """
         raise NotImplementedError("sample must be implemented in subclasses.")
 
@@ -295,8 +295,7 @@ def maybe_oversample_requests(
         request_id_prefix: str = "",
         no_oversample: bool = False,
     ) -> None:
-        """
-        Oversamples the list of requests if its size is less than the desired
+        """Oversamples the list of requests if its size is less than the desired
         number.
 
         Args:
@@ -305,6 +304,8 @@ def maybe_oversample_requests(
             num_requests (int): The target number of requests.
             request_id_prefix (str): The prefix applied to generated request
                 identifiers.
+            no_oversample (bool): If True, return the requests unchanged
+                instead of oversampling up to num_requests.
 
         """
         if no_oversample:
@@ -346,8 +347,7 @@ def is_valid_sequence(
     max_total_len: int = 2048,
     skip_min_output_len_check: bool = False,
 ) -> bool:
-    """
-    Validate a sequence based on prompt and output lengths.
+    """Validate a sequence based on prompt and output lengths.
 
     Default pruning criteria are copied from the original `sample_hf_requests`
     and `sample_sharegpt_requests` functions in benchmark_serving.py, as well as
@@ -379,8 +379,7 @@ def process_image(
     *,
     ensure_client_side_data: bool = False,
 ) -> Mapping[str, Any]:
-    """
-    Process a single image input and return a multimedia content dictionary.
+    """Process a single image input and return a multimedia content dictionary.
 
     Supports the following input types:
 
@@ -402,6 +401,7 @@ def process_image(
 
     Raises:
         ValueError: If the input is not a supported type.
+
     """
     if isinstance(image, dict) and "bytes" in image:
         image = Image.open(BytesIO(image["bytes"]))
@@ -437,8 +437,7 @@ def process_image(
 
 
 def process_video(video: Any) -> Mapping[str, Any]:
-    """
-    Process a single video input and return a multimedia content dictionary.
+    """Process a single video input and return a multimedia content dictionary.
 
     Supports the following input types:
 
@@ -451,6 +450,7 @@ def process_video(video: Any) -> Mapping[str, Any]:
 
     Raises:
         ValueError: If the input is not a supported type.
+
     """
     if isinstance(video, dict) and "bytes" in video:
         video_bytes = video["bytes"]
@@ -474,8 +474,7 @@ def process_video(video: Any) -> Mapping[str, Any]:
 
 
 def process_audio(audio: Any) -> tuple:
-    """
-    Process a single audio input and return a (array, sample_rate) tuple.
+    """Process a single audio input and return a (array, sample_rate) tuple.
 
     Supports:
     1. String: treated as a file path, loaded with soundfile.
@@ -502,8 +501,7 @@ def gen_prompt_decode_to_target_len(
     add_special_tokens: bool = False,
     rng: np.random.Generator | None = None,
 ) -> tuple[str, list[int], int]:
-    """
-    Ensure decoded-then-encoded prompt length matches the target token length.
+    """Ensure decoded-then-encoded prompt length matches the target token length.
 
     This function decodes an initial token sequence to text and re-encodes it
     , iteratively adjusting the token sequence length to match a target.
@@ -557,8 +555,7 @@ def gen_prompt_decode_to_target_len(
 
 
 class RandomDataset(BenchmarkDataset):
-    """
-    Synthetic text-only dataset for serving/throughput benchmarks.
+    """Synthetic text-only dataset for serving/throughput benchmarks.
 
     Strategy:
     - Sample input/output token lengths per request from integer-uniform ranges
@@ -699,9 +696,7 @@ def get_prefix(
         allowed_tokens: np.ndarray,
         prefix_len: int,
     ) -> list[int]:
-        """
-        Get the prefix for the dataset.
-        """
+        """Get the prefix for the dataset."""
         if prefix_len <= 0:
             return []
 
@@ -738,8 +733,7 @@ def generate_token_sequence(
         index: int,
         allowed_tokens: np.ndarray,
     ) -> tuple[str, int, int]:
-        """
-        Returns (prompt, total_input_len).
+        """Returns (prompt, total_input_len).
 
         NOTE: After decoding the prompt we have to encode and decode it again.
         This is done because in some cases N consecutive tokens
@@ -778,8 +772,7 @@ def generate_token_sequence(
 
 
 class RandomDatasetForReranking(RandomDataset):
-    """
-    Random dataset specialized for the needs of scoring:
+    """Random dataset specialized for the needs of scoring:
     - Batches of inputs
     - Inputs composed of pairs
     """
@@ -906,8 +899,7 @@ def sample(
 
 
 class RandomMultiModalDataset(RandomDataset):
-    """
-    Synthetic multimodal dataset (text + images) that extends RandomDataset.
+    """Synthetic multimodal dataset (text + images) that extends RandomDataset.
 
     Status:
     - Images: supported via synthetic RGB data.
@@ -1022,8 +1014,7 @@ def map_config_to_modality(self, config: tuple[int, int, int]) -> str:
     def normalize_bucket_config(
         self, bucket_config: dict[tuple[int, int, int], float]
     ) -> dict[tuple[int, int, int], float]:
-        """
-        Remove zero probability entries
+        """Remove zero probability entries
         and normalize the bucket config to sum to 1.
         """
         # Raise error if value is negative
@@ -1044,13 +1035,11 @@ def generate_mm_item(
         self,
         mm_item_config: tuple[int, int, int],
     ) -> Mapping[str, Any]:
-        """
-        Create synthetic images and videos and
+        """Create synthetic images and videos and
         apply process_image/process_video respectively.
         This follows the OpenAI API chat completions
         https://github.com/openai/openai-python
         """
-
         if self.map_config_to_modality(mm_item_config) == "image":
             return process_image(
                 self.generate_synthetic_image(mm_item_config[1], mm_item_config[0])
@@ -1071,9 +1060,7 @@ def get_mm_item_sampling_params(
         limit_mm_per_prompt: dict[str, int],
         bucket_config: dict[tuple[int, int, int], float],
     ) -> tuple[int, int, dict[str, int], dict[tuple[int, int, int], float]]:
-        """
-        Get the sampling parameters for the multimodal items.
-        """
+        """Get the sampling parameters for the multimodal items."""
         # Enforce num_mm_items_range_ratio <= 1
         if not (0.0 <= num_mm_items_range_ratio <= 1.0):
             raise ValueError("num_mm_items_range_ratio must be in [0, 1].")
@@ -1146,8 +1133,7 @@ def get_mm_item_iterator(
         bucket_config: dict[tuple[int, int, int], float],
         limit_mm_per_prompt: dict[str, int],
     ) -> Iterator[tuple[int, int, int]]:
-        """
-        Iterator over the multimodal items for each request
+        """Iterator over the multimodal items for each request
         whose size is between min_num_mm_items and max_num_mm_items.
 
         Loop over the bucket config and sample a multimodal item.
@@ -1160,6 +1146,7 @@ def get_mm_item_iterator(
           `bucket_config` (tuple->float). The original dict passed to
           `sample` is not mutated. If this ever changes, a test
           is implemented and will fail.
+
         """
         # Get the number of multimodal items to sample
         request_num_mm_items = int(
@@ -1346,8 +1333,7 @@ def sample(
 
 
 class ShareGPTDataset(BenchmarkDataset):
-    """
-    Implements the ShareGPT dataset.  Loads data from a JSON file and generates
+    """Implements the ShareGPT dataset.  Loads data from a JSON file and generates
     sample requests based on conversation turns.
     """
 
@@ -1437,8 +1423,7 @@ def sample(
 
 
 class TimedTrace(BenchmarkDataset):
-    """
-    Implements a base class to replay various timed traces.
+    """Implements a base class to replay various timed traces.
     Loads data from a JSON file and generates sample requests
     based on the timing information in the traces.
     """
@@ -1927,6 +1912,7 @@ def add_random_dataset_base_args(
 
     Args:
         parser_or_group: Either a parser or an argument group to add arguments to.
+
     """
     parser_or_group.add_argument(
         "--random-input-len",
@@ -1988,6 +1974,7 @@ def add_random_multimodal_dataset_args(
 
     Args:
         parser_or_group: Either a parser or an argument group to add arguments to.
+
     """
     parser_or_group.add_argument(
         "--random-mm-base-items-per-request",
@@ -2503,8 +2490,7 @@ def get_samples(
 
 
 class CustomDataset(BenchmarkDataset):
-    """
-    Implements the Custom dataset.  Loads data from a JSONL file and generates
+    """Implements the Custom dataset.  Loads data from a JSONL file and generates
     sample requests based on conversation turns. E.g.,
     ```
     {"prompt": "What is the capital of India?", "output_tokens": 10}
@@ -2632,8 +2618,7 @@ def sample(
 
 
 class CustomImageDataset(CustomDataset):
-    """
-    Implements the Custom image dataset. Loads data from a JSONL file and generates
+    """Implements the Custom image dataset. Loads data from a JSONL file and generates
     sample requests based on conversation turns. E.g.,
     ```
     {
@@ -2917,8 +2902,7 @@ def sample(
 
 
 class CustomAudioDataset(CustomDataset):
-    """
-    Custom dataset for audio benchmarking. Loads data from a JSONL file. E.g.,
+    """Custom dataset for audio benchmarking. Loads data from a JSONL file. E.g.,
     {"prompt": "Transcribe the audio.", "audio": "/path/to/audio.wav"}
 
     Supports both:
@@ -3027,8 +3011,7 @@ def sample(
 
 
 class SpecBench(CustomDataset):
-    """
-    Implements the SpecBench dataset: https://github.com/hemingkx/Spec-Bench
+    """Implements the SpecBench dataset: https://github.com/hemingkx/Spec-Bench
     Download the dataset using:
     wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl
     """  # noqa: E501
@@ -3100,8 +3083,7 @@ def sample(
     "SonnetDataset is deprecated and will be removed in a future version.",
 )
 class SonnetDataset(BenchmarkDataset):
-    """
-    Simplified implementation of the Sonnet dataset.  Loads poem lines from a
+    """Simplified implementation of the Sonnet dataset.  Loads poem lines from a
     text file and generates sample requests.  Default values here copied from
     `benchmark_serving.py` for the sonnet dataset.
     """
@@ -3198,8 +3180,7 @@ def sample(
 
 
 class BurstGPTDataset(BenchmarkDataset):
-    """
-    Implements the BurstGPT dataset.  Loads data from a CSV file and generates
+    """Implements the BurstGPT dataset.  Loads data from a CSV file and generates
     sample requests based on synthetic prompt generation. Only rows with Model
     "GPT-4" and positive response tokens are used.
     """
@@ -3447,9 +3428,7 @@ def sample(
 
 
 class VisionArenaDataset(HuggingFaceDataset):
-    """
-    Vision Arena Dataset.
-    """
+    """Vision Arena Dataset."""
 
     DEFAULT_OUTPUT_LEN = 128
     SUPPORTED_DATASET_PATHS = {
@@ -3506,8 +3485,7 @@ def sample(
 
 
 class MMVUDataset(HuggingFaceDataset):
-    """
-    MMVU Dataset.
+    """MMVU Dataset.
     https://huggingface.co/datasets/yale-nlp/MMVU
     """
 
@@ -3586,8 +3564,7 @@ def sample(
 
 
 class InstructCoderDataset(HuggingFaceDataset):
-    """
-    InstructCoder Dataset.
+    """InstructCoder Dataset.
     https://huggingface.co/datasets/likaixin/InstructCoder
 
     InstructCoder is the dataset designed for general code editing.  It consists
@@ -3655,8 +3632,7 @@ def sample_prompts(self, n: int) -> Iterator[str]:
 
 
 class MTBenchDataset(HuggingFaceDataset):
-    """
-    MT-Bench Dataset.
+    """MT-Bench Dataset.
     https://huggingface.co/datasets/philschmid/mt-bench
 
     We create a single turn dataset for MT-Bench.
@@ -3718,8 +3694,7 @@ def sample(
 
 
 class HumanEvalDataset(HuggingFaceDataset):
-    """
-    HumanEvalDataset Dataset.
+    """HumanEvalDataset Dataset.
     https://huggingface.co/datasets/openai/openai_humaneval
 
     We create a single turn dataset for HumanEval.
@@ -3779,8 +3754,7 @@ def sample(
 
 
 class GSM8KDataset(HuggingFaceDataset):
-    """
-    GSM8K Dataset.
+    """GSM8K Dataset.
     https://huggingface.co/datasets/openai/gsm8k
 
     We create a single turn dataset for GSM8K.
@@ -3840,8 +3814,7 @@ def sample(
 
 
 class BlazeditDataset(HuggingFaceDataset):
-    """
-    Blazedit Dataset.
+    """Blazedit Dataset.
     https://github.com/ise-uiuc/blazedit
 
     5k char version: vdaita/edit_5k_char
@@ -3932,9 +3905,7 @@ def sample(
 
 
 class AIMODataset(HuggingFaceDataset):
-    """
-    Dataset class for processing a AIMO dataset with reasoning questions.
-    """
+    """Dataset class for processing a AIMO dataset with reasoning questions."""
 
     SUPPORTED_DATASET_PATHS = {
         "AI-MO/aimo-validation-aime",
@@ -4026,6 +3997,7 @@ def _format_zeta_prompt(
 
     Returns:
         A dictionary with the formatted prompts and expected outputs.
+
     """
     events = sample["events"]
     input = sample["input"]
@@ -4042,9 +4014,7 @@ def _format_zeta_prompt(
 
 
 class NextEditPredictionDataset(HuggingFaceDataset):
-    """
-    Dataset class for processing a Next Edit Prediction dataset.
-    """
+    """Dataset class for processing a Next Edit Prediction dataset."""
 
     SUPPORTED_DATASET_PATHS = {
         "zed-industries/zeta",
@@ -4092,8 +4062,7 @@ def sample(
 
 
 class ASRDataset(HuggingFaceDataset):
-    """
-    Dataset class for processing a ASR dataset for transcription.
+    """Dataset class for processing a ASR dataset for transcription.
     Tested on the following set:
 
     +---------------------------+----------------------------------------+--------------------------+-----------------------------+
@@ -4290,8 +4259,7 @@ def sample(
 
 
 class MLPerfDataset(HuggingFaceDataset):
-    """
-    MLPerf Inference Dataset.
+    """MLPerf Inference Dataset.
 
     Dataset on HF:
     https://huggingface.co/datasets/mgoin/mlperf-inference-llama2-data
@@ -4470,8 +4438,7 @@ def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]:
 
 
 class MMStarDataset(HuggingFaceDataset):
-    """
-    Lin-Chen/MMStar: https://huggingface.co/datasets/Lin-Chen/MMStar
+    """Lin-Chen/MMStar: https://huggingface.co/datasets/Lin-Chen/MMStar
     refer to: https://github.com/sgl-project/SpecForge/pull/106
     """
 
@@ -4760,8 +4727,7 @@ def sample(
 
 
 class SpeedBench(CustomDataset):
-    """
-    SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench
+    """SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench.
 
     Download the dataset using:
 
diff --git a/vllm/benchmarks/datasets/utils.py b/vllm/benchmarks/datasets/utils.py
index bc5a4340dd62..a2dd5093c50b 100644
--- a/vllm/benchmarks/datasets/utils.py
+++ b/vllm/benchmarks/datasets/utils.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shared utilities for benchmark dataset sampling.
-"""
+"""Shared utilities for benchmark dataset sampling."""
 
 import logging
 import math
@@ -46,8 +44,7 @@ def get_sampling_params(
     output_len: int,
     tokenizer: TokenizerLike,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """
-    Sample per-request input/output token lengths and vocab offsets.
+    """Sample per-request input/output token lengths and vocab offsets.
 
     Lengths are drawn uniformly from integer ranges around the configured
     means, controlled by *range_ratio*.  It may be a single ``float``
@@ -60,6 +57,7 @@ def get_sampling_params(
     Returns:
         (input_lens, output_lens, offsets) – three 1-D ``np.ndarray`` of
         shape ``(num_requests,)``.
+
     """
     input_range_ratio, output_range_ratio = _resolve_range_ratios(range_ratio)
 
diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py
index f426e78c744c..2b4027b502c3 100644
--- a/vllm/benchmarks/lib/endpoint_request_func.py
+++ b/vllm/benchmarks/lib/endpoint_request_func.py
@@ -167,10 +167,12 @@ async def async_request_openai_completions(
 
     Args:
         request_func_input: The input for the request function.
+        session: The aiohttp session used to issue the request.
         pbar: The progress bar to display the progress.
 
     Returns:
         The output of the request function.
+
     """
     api_url = request_func_input.api_url
     _validate_api_url(api_url, "OpenAI Completions API", "completions")
@@ -235,8 +237,7 @@ async def async_request_openai_completions(
                                 # First token
                                 if not first_chunk_received:
                                     first_chunk_received = True
-                                    ttft = time.perf_counter() - st
-                                    output.ttft = ttft
+                                    output.ttft = timestamp - st
 
                                 # Decoding phase
                                 else:
@@ -415,13 +416,14 @@ async def async_request_openai_chat_completions(
                                     output.itl.append(timestamp - most_recent_timestamp)
 
                                 generated_text += content or ""
+                                # Only token chunks advance the request end;
+                                # the trailing usage chunk carries no token.
+                                most_recent_timestamp = timestamp
                             elif usage := data.get("usage"):
                                 output.output_tokens = usage.get("completion_tokens")
                                 if (pt := usage.get("prompt_tokens")) is not None:
                                     output.prompt_len = pt
 
-                            most_recent_timestamp = timestamp
-
                 output.generated_text = generated_text
                 if first_chunk_received:
                     output.success = True
@@ -540,13 +542,15 @@ async def send_audio_file(
                                         )
 
                                     generated_text += content or ""
+                                    # Only token chunks advance the request
+                                    # end; the trailing usage chunk carries no
+                                    # token.
+                                    most_recent_timestamp = timestamp
                                 elif usage := data.get("usage"):
                                     output.output_tokens = usage.get(
                                         "completion_tokens"
                                     )
 
-                                most_recent_timestamp = timestamp
-
                     output.generated_text = generated_text
                     if first_chunk_received:
                         output.success = True
@@ -604,6 +608,7 @@ async def _run_pooling_request(
     headers: dict[str, Any],
     pbar: tqdm | None = None,
     num_input_sequences: int = 1,
+    prompt_len: int = 0,
 ) -> RequestFuncOutput:
     output = RequestFuncOutput(num_input_sequences=num_input_sequences)
     st = time.perf_counter()
@@ -611,11 +616,19 @@ async def _run_pooling_request(
     try:
         async with session.post(url=api_url, headers=headers, json=payload) as response:
             if response.status == 200:
+                encoding_format = payload.get("encoding_format", "float")
+                if encoding_format in ("bytes", "bytes_only"):
+                    async for _ in response.content.iter_any():
+                        pass
+                else:
+                    await response.read()
                 output.ttft = output.latency = time.perf_counter() - st
 
-                if payload.get("encoding_format", "float") == "bytes":
+                if encoding_format == "bytes":
                     metadata = json.loads(response.headers["metadata"])
                     usage = metadata.get("usage", {})
+                elif encoding_format == "bytes_only":
+                    usage = {"prompt_tokens": prompt_len}
                 else:
                     data = await response.json()
                     usage = data.get("usage", {})
@@ -670,6 +683,7 @@ async def async_request_openai_embeddings(
         headers=headers,
         pbar=pbar,
         num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
+        prompt_len=request_func_input.prompt_len,
     )
 
 
@@ -707,6 +721,7 @@ async def async_request_vllm_rerank(
         headers=headers,
         pbar=pbar,
         num_input_sequences=len(request_func_input.prompt) - 1,
+        prompt_len=request_func_input.prompt_len,
     )
 
 
@@ -741,6 +756,7 @@ async def async_request_openai_embeddings_chat(
         payload=payload,
         headers=headers,
         pbar=pbar,
+        prompt_len=request_func_input.prompt_len,
     )
 
 
@@ -845,6 +861,7 @@ async def async_request_infinity_embeddings(
         headers=headers,
         pbar=pbar,
         num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
+        prompt_len=request_func_input.prompt_len,
     )
 
 
@@ -894,6 +911,7 @@ async def async_request_vllm_pooling(
         headers=headers,
         pbar=pbar,
         num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
+        prompt_len=request_func_input.prompt_len,
     )
 
 
diff --git a/vllm/benchmarks/lib/ready_checker.py b/vllm/benchmarks/lib/ready_checker.py
index eec4a42cb670..e3b4d8ee3ea7 100644
--- a/vllm/benchmarks/lib/ready_checker.py
+++ b/vllm/benchmarks/lib/ready_checker.py
@@ -22,12 +22,12 @@ async def wait_for_endpoint(
     timeout_seconds: int = 600,
     retry_interval: int = 5,
 ) -> RequestFuncOutput:
-    """
-    Wait for an endpoint to become available before starting benchmarks.
+    """Wait for an endpoint to become available before starting benchmarks.
 
     Args:
         request_func: The async request function to call
         test_input: The RequestFuncInput to test with
+        session: The aiohttp session used to issue the probe requests
         timeout_seconds: Maximum time to wait in seconds (default: 10 minutes)
         retry_interval: Time between retries in seconds (default: 5 seconds)
 
@@ -36,6 +36,7 @@ async def wait_for_endpoint(
 
     Raises:
         ValueError: If the endpoint doesn't become available within the timeout
+
     """
     deadline = time.perf_counter() + timeout_seconds
     output = RequestFuncOutput(success=False)
diff --git a/vllm/benchmarks/lib/utils.py b/vllm/benchmarks/lib/utils.py
index 7737621a4e7c..f5bd11439ab4 100644
--- a/vllm/benchmarks/lib/utils.py
+++ b/vllm/benchmarks/lib/utils.py
@@ -9,6 +9,17 @@
 from typing import Any
 
 
+def redact_sensitive_namespace(
+    args: argparse.Namespace, fields: tuple[str, ...]
+) -> argparse.Namespace:
+    """Return a copy of CLI arguments with selected values redacted."""
+    redacted_args = argparse.Namespace(**vars(args))
+    for field in fields:
+        if getattr(redacted_args, field, None) is not None:
+            setattr(redacted_args, field, "***")
+    return redacted_args
+
+
 def extract_field(
     args: argparse.Namespace, extra_info: dict[str, Any], field_name: str
 ) -> str:
@@ -25,9 +36,7 @@ def extract_field(
 
 
 def use_compile(args: argparse.Namespace, extra_info: dict[str, Any]) -> bool:
-    """
-    Check if the benchmark is run with torch.compile
-    """
+    """Check if the benchmark is run with torch.compile."""
     return not (
         extract_field(args, extra_info, "compilation_config.mode") == "0"
         or "eager" in getattr(args, "output_json", "")
@@ -38,8 +47,7 @@ def use_compile(args: argparse.Namespace, extra_info: dict[str, Any]) -> bool:
 def convert_to_pytorch_benchmark_format(
     args: argparse.Namespace, metrics: dict[str, list], extra_info: dict[str, Any]
 ) -> list:
-    """
-    Save the benchmark results in the format used by PyTorch OSS benchmark with
+    """Save the benchmark results in the format used by PyTorch OSS benchmark with
     on metric per record
     https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database
     """
diff --git a/vllm/benchmarks/mm_processor.py b/vllm/benchmarks/mm_processor.py
index 4c840786fe72..40b3494151d3 100644
--- a/vllm/benchmarks/mm_processor.py
+++ b/vllm/benchmarks/mm_processor.py
@@ -44,8 +44,7 @@
 
 
 def get_timing_stats_from_engine(llm_engine: LLMEngine) -> dict[str, dict[str, float]]:
-    """
-    Get all multimodal timing stats from the LLM engine.
+    """Get all multimodal timing stats from the LLM engine.
 
     Collects both preprocessing stats (HF processor, hashing, cache lookup,
     prompt update) and encoder forward pass timing, merged by request_id.
@@ -70,6 +69,7 @@ def get_timing_stats_from_engine(llm_engine: LLMEngine) -> dict[str, dict[str, f
                 'num_encoder_calls': 1
             }
         }
+
     """
     observability_config = llm_engine.vllm_config.observability_config
     if not observability_config or not observability_config.enable_mm_processor_stats:
@@ -126,8 +126,7 @@ def get_timing_stats_from_engine(llm_engine: LLMEngine) -> dict[str, dict[str, f
 
 
 def collect_mm_processor_stats(llm_engine: LLMEngine) -> dict[str, list[float]]:
-    """
-    Collect multimodal processor timing stats.
+    """Collect multimodal processor timing stats.
     Returns a dictionary mapping stage names to lists of timing values (in seconds).
     """
     all_stats = get_timing_stats_from_engine(llm_engine)
@@ -147,9 +146,7 @@ def calculate_mm_processor_metrics(
     *,
     unit: Literal["us", "ms", "s"] = "ms",
 ) -> dict[str, dict[str, float]]:
-    """
-    Calculate aggregate metrics from stats by stage.
-    """
+    """Calculate aggregate metrics from stats by stage."""
     unit2mult = {"us": 1000000, "ms": 1000, "s": 1}
     unit_mult = unit2mult[unit]
 
@@ -181,9 +178,7 @@ def calculate_mm_processor_metrics(
 
 
 def validate_args(args):
-    """
-    Validate command-line arguments for mm_processor benchmark.
-    """
+    """Validate command-line arguments for mm_processor benchmark."""
     if not getattr(args, "tokenizer", None):
         args.tokenizer = args.model
     if not hasattr(args, "dataset_path"):
@@ -214,9 +209,7 @@ def validate_args(args):
 def benchmark_multimodal_processor(
     args: argparse.Namespace,
 ) -> dict[str, Any]:
-    """
-    Run the multimodal processor benchmark.
-    """
+    """Run the multimodal processor benchmark."""
     from vllm import LLM, SamplingParams
 
     validate_args(args)
@@ -455,7 +448,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> None:
 
 def main(args: argparse.Namespace) -> None:
     """Main entry point for the multimodal processor benchmark."""
-
     print("Starting multimodal processor benchmark...")
     result = benchmark_multimodal_processor(args)
 
diff --git a/vllm/benchmarks/plot.py b/vllm/benchmarks/plot.py
index 3f36ede721ad..450019e80e25 100644
--- a/vllm/benchmarks/plot.py
+++ b/vllm/benchmarks/plot.py
@@ -29,8 +29,7 @@ def generate_timeline_plot(
     itl_thresholds: list[float] | None = None,
     labels: list[str] | None = None,
 ) -> None:
-    """
-    Generate an HTML timeline plot from benchmark results.
+    """Generate an HTML timeline plot from benchmark results.
 
     Args:
         results: List of per-request result dictionaries containing:
@@ -44,8 +43,8 @@ def generate_timeline_plot(
         colors: List of colors for ITL categories (default: green, orange, red, black)
         itl_thresholds: ITL thresholds in seconds (default: [1.0, 4.0, 6.0])
         labels: Labels for ITL categories (default based on thresholds)
-    """
 
+    """
     # Set defaults
     if colors is None:
         colors = ["#109618", "#FF7F0E", "#D62728"]
@@ -118,8 +117,7 @@ def construct_timeline_data(
     itl_thresholds: list[float],
     labels: list[str],
 ) -> list[dict[str, Any]]:
-    """
-    Construct timeline data from request results.
+    """Construct timeline data from request results.
 
     Args:
         requests_data: List of per-request result dictionaries
@@ -128,6 +126,7 @@ def construct_timeline_data(
 
     Returns:
         List of timeline segments for plotting
+
     """
 
     def tostr(sec_time: float) -> str:
@@ -229,8 +228,7 @@ def generate_dataset_stats_plot(
     results: list[dict[str, Any]],
     output_path: Path,
 ) -> None:
-    """
-    Generate a matplotlib figure with dataset statistics.
+    """Generate a matplotlib figure with dataset statistics.
 
     Creates a figure with 4 subplots:
     - Top-left: Prompt tokens distribution (histogram)
@@ -243,6 +241,7 @@ def generate_dataset_stats_plot(
             - prompt_len: Number of prompt tokens
             - output_tokens: Number of output tokens
         output_path: Path where the figure will be saved
+
     """
     # Extract data
     prompt_tokens = []
diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py
index aef7c7e7b4d9..d751f7f22dd7 100644
--- a/vllm/benchmarks/serve.py
+++ b/vllm/benchmarks/serve.py
@@ -51,13 +51,18 @@
     RequestFuncOutput,
 )
 from vllm.benchmarks.lib.ready_checker import wait_for_endpoint
-from vllm.benchmarks.lib.utils import convert_to_pytorch_benchmark_format, write_to_json
+from vllm.benchmarks.lib.utils import (
+    convert_to_pytorch_benchmark_format,
+    redact_sensitive_namespace,
+    write_to_json,
+)
 from vllm.tokenizers import TokenizerLike, get_tokenizer
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.gc_utils import freeze_gc_heap
 from vllm.utils.network_utils import join_host_port
 
 MILLISECONDS_TO_SECONDS_CONVERSION = 1000
+_SENSITIVE_ARG_FIELDS = ("header",)
 
 
 def _merge_overrides(base: dict | None, override: dict | None) -> dict | None:
@@ -402,13 +407,15 @@ async def get_request(
     ramp_up_end_rps: int | None = None,
     self_timed: bool = False,
 ) -> AsyncGenerator[tuple[SampleRequest, float], None]:
-    """
-    Asynchronously generates requests at a specified rate
+    """Asynchronously generates requests at a specified rate
     with OPTIONAL burstiness and OPTIONAL ramp-up strategy.
 
     Args:
         input_requests:
             A list of input requests, each represented as a SampleRequest.
+        self_timed:
+            If True, the requests carry their own arrival timing and no
+            request rate, burstiness or ramp-up is applied.
         request_rate:
             The rate at which requests are generated (requests/s).
         burstiness (optional):
@@ -426,6 +433,7 @@ async def get_request(
             The starting request rate for ramp-up.
         ramp_up_end_rps (optional):
             The ending request rate for ramp-up.
+
     """
     assert burstiness > 0, (
         f"A positive burstiness factor is expected, but given {burstiness}."
@@ -521,6 +529,7 @@ def calculate_metrics_for_embeddings(
 
     Returns:
         The calculated benchmark metrics.
+
     """
     total_input = 0
     total_input_sequences = 0
@@ -580,6 +589,7 @@ def calculate_metrics(
 
     Returns:
         A tuple of the benchmark metrics and the actual output lengths.
+
     """
     actual_output_lens: list[int] = []
     total_input = 0
@@ -1536,7 +1546,7 @@ def save_to_pytorch_benchmark_format(
     # later if needed
     ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"]
     pt_records = convert_to_pytorch_benchmark_format(
-        args=args,
+        args=redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS),
         metrics={k: [results[k]] for k in metrics if k in results},
         extra_info={
             k: results[k]
@@ -1566,6 +1576,7 @@ def compute_result_filename(
 
     Returns:
         The computed filename path or None if no result saving is requested
+
     """
     if not (args.plot_timeline or args.save_result or args.append_result):
         return None
@@ -1684,7 +1695,6 @@ def add_cli_args(parser: FlexibleArgumentParser):
         - "auto" will use the tokenizer from `mistral_common` for Mistral models
         if available, otherwise it will use the "hf" tokenizer.\n
         - "hf" will use the fast tokenizer if available.\n
-        - "slow" will always use the slow tokenizer.\n
         - "mistral" will always use the tokenizer from `mistral_common`.\n
         - "deepseek_v32" will always use the tokenizer from `deepseek_v32`.\n
         - Other custom values can be supported via plugins.""",
@@ -2013,7 +2023,7 @@ def main(args: argparse.Namespace) -> dict[str, Any]:
 
 
 async def main_async(args: argparse.Namespace) -> dict[str, Any]:
-    print(args)
+    print(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS))
     if args.max_concurrency is not None and args.max_concurrency <= 0:
         raise ValueError("--max-concurrency must be greater than 0")
 
diff --git a/vllm/benchmarks/startup.py b/vllm/benchmarks/startup.py
index 86f87a32d666..1df66af622e1 100644
--- a/vllm/benchmarks/startup.py
+++ b/vllm/benchmarks/startup.py
@@ -107,8 +107,7 @@ def _metric_to_json(m: MetricStats) -> dict[str, Any]:
 
 @contextmanager
 def cold_startup():
-    """
-    Context manager to measure cold startup time:
+    """Context manager to measure cold startup time:
     1. Uses a temporary directory for vLLM cache to avoid any pollution
        between cold startup iterations.
     2. Uses inductor's fresh_cache to clear torch.compile caches.
@@ -132,8 +131,7 @@ def cold_startup():
 
 
 def run_startup_in_subprocess(engine_args, result_queue):
-    """
-    Run LLM startup in a subprocess and return timing metrics via a queue.
+    """Run LLM startup in a subprocess and return timing metrics via a queue.
     This ensures complete isolation between iterations.
     """
     try:
@@ -229,11 +227,9 @@ def main(args: argparse.Namespace):
     engine_args = EngineArgs.from_cli_args(args)
 
     def create_llm_and_measure_startup():
-        """
-        Create LLM instance in a subprocess and measure startup time.
+        """Create LLM instance in a subprocess and measure startup time.
         Returns timing metrics, using subprocess for complete isolation.
         """
-
         # Create a queue for inter-process communication
         result_queue: multiprocessing.Queue[Any] = multiprocessing.Queue()
         process = multiprocessing.Process(
diff --git a/vllm/benchmarks/sweep/param_sweep.py b/vllm/benchmarks/sweep/param_sweep.py
index 5c2620e794ec..c7607109483d 100644
--- a/vllm/benchmarks/sweep/param_sweep.py
+++ b/vllm/benchmarks/sweep/param_sweep.py
@@ -18,14 +18,14 @@ def read_json(cls, filepath: os.PathLike):
 
     @classmethod
     def read_from_dict(cls, data: dict[str, dict[str, object]]):
-        """
-        Read parameter sweep from a dict format where keys are names.
+        """Read parameter sweep from a dict format where keys are names.
 
         Example:
             {
                 "experiment1": {"max_tokens": 100, "temperature": 0.7},
                 "experiment2": {"max_tokens": 200, "temperature": 0.9}
             }
+
         """
         records = [{"_benchmark_name": name, **params} for name, params in data.items()]
         return cls.from_records(records)
@@ -66,8 +66,7 @@ def __or__(self, other: dict[str, object], /) -> "ParameterSweepItem":  # type:
 
     @property
     def name(self) -> str:
-        """
-        Get the name for this parameter sweep item.
+        """Get the name for this parameter sweep item.
 
         Returns the '_benchmark_name' field if present, otherwise returns a text
         representation of all parameters.
@@ -103,8 +102,7 @@ def has_param(self, param_key: str) -> bool:
         return any(k in self for k in self._iter_param_key_candidates(param_key))
 
     def _normalize_cmd_kv_pair(self, k: str, v: object) -> list[str]:
-        """
-        Normalize a key-value pair into command-line arguments.
+        """Normalize a key-value pair into command-line arguments.
 
         Returns a list containing either:
         - A single element for boolean flags (e.g., ['--flag'] or ['--flag=true'])
diff --git a/vllm/benchmarks/sweep/plot.py b/vllm/benchmarks/sweep/plot.py
index 79e50603b8c8..59a90636f0e0 100644
--- a/vllm/benchmarks/sweep/plot.py
+++ b/vllm/benchmarks/sweep/plot.py
@@ -178,8 +178,7 @@ def _json_load_bytes(path: Path) -> list[dict[str, object]]:
 
 
 def _convert_inf_nan_strings(data: list[dict[str, object]]) -> list[dict[str, object]]:
-    """
-    Convert string values "inf", "-inf", and "nan" to their float equivalents.
+    """Convert string values "inf", "-inf", and "nan" to their float equivalents.
 
     This handles the case where JSON serialization represents inf/nan as strings.
     """
diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py
index 17b2f2426477..94b22b29567c 100644
--- a/vllm/benchmarks/throughput.py
+++ b/vllm/benchmarks/throughput.py
@@ -195,8 +195,7 @@ def run_vllm_chat(
     warmup_requests: list[SampleRequest] | None = None,
     prequeue_requests: bool = False,
 ) -> tuple[float, list[RequestOutput]]:
-    """
-    Run vLLM chat benchmark. This function is recommended ONLY for benchmarking
+    """Run vLLM chat benchmark. This function is recommended ONLY for benchmarking
     multimodal models as it properly handles multimodal inputs and chat
     formatting. For non-multimodal models, use run_vllm() instead.
     """
@@ -536,30 +535,42 @@ def _to_serve_args(args: argparse.Namespace) -> argparse.Namespace:
 
     Returns:
         A namespace satisfying get_samples's attribute reads.
+
     """
     d = vars(args).copy()
     # random_*: prefer --random-* over legacy --input/output/prefix-len.
-    d["random_input_len"] = getattr(args, "random_input_len", None) or args.input_len
-    d["random_output_len"] = getattr(args, "random_output_len", None) or args.output_len
-    d["random_prefix_len"] = getattr(args, "random_prefix_len", None) or args.prefix_len
+    # Callers without the legacy flags (e.g. bench mm-processor) omit them, so
+    # read them optionally. Identity checks, not truthiness: an explicit 0
+    # (the default of --random-prefix-len) must not fall through to them.
+    input_len = getattr(args, "input_len", None)
+    output_len = getattr(args, "output_len", None)
+    prefix_len = getattr(args, "prefix_len", None)
+    d["random_input_len"] = getattr(args, "random_input_len", input_len)
+    d["random_output_len"] = getattr(args, "random_output_len", output_len)
+    d["random_prefix_len"] = getattr(args, "random_prefix_len", prefix_len)
     # --output-len maps to the per-dataset output-len entry points get_samples
     # reads. None passes through (each dataset applies its own default),
     # matching prior throughput behaviour of omitting output_len when unset.
-    d["hf_output_len"] = args.output_len
-    d["sharegpt_output_len"] = args.output_len
+    d["hf_output_len"] = output_len
+    d["sharegpt_output_len"] = output_len
     # sonnet reads dedicated attrs; fall back to SonnetDataset's own defaults.
-    d["sonnet_input_len"] = args.input_len if args.input_len is not None else 550
-    d["sonnet_output_len"] = args.output_len if args.output_len is not None else 150
-    d["sonnet_prefix_len"] = args.prefix_len
+    d["sonnet_input_len"] = input_len if input_len is not None else 550
+    d["sonnet_output_len"] = output_len if output_len is not None else 150
+    d["sonnet_prefix_len"] = prefix_len
     # Explicit --enable-multimodal-chat wins; otherwise auto-enable for the
-    # multimodal chat backend (preserves today's vllm-chat handling).
+    # multimodal chat backend (preserves today's vllm-chat handling). Callers
+    # without a --backend flag (e.g. bench mm-processor) drive every request
+    # through llm.chat, so default to the vllm-chat behaviour.
+    backend = getattr(args, "backend", "vllm-chat")
+    d["backend"] = backend
     d["enable_multimodal_chat"] = bool(
-        getattr(args, "enable_multimodal_chat", False) or args.backend == "vllm-chat"
+        getattr(args, "enable_multimodal_chat", False) or backend == "vllm-chat"
     )
     # serve-only attrs throughput never exposed; keep serve's defaults.
     d.setdefault("disable_shuffle", False)
     d.setdefault("skip_chat_template", False)
     d.setdefault("no_stream", False)
+    d.setdefault("no_oversample", False)
     d.setdefault("request_id_prefix", "")
     d.setdefault("chat_template_kwargs", None)
     return argparse.Namespace(**d)
@@ -616,10 +627,7 @@ def assign_loras(requests, args):
 
 
 def validate_args(args):
-    """
-    Validate command-line arguments.
-    """
-
+    """Validate command-line arguments."""
     # === Deprecation and Defaulting ===
     if args.dataset is not None:
         warnings.warn(
diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py
index 439a87e1522a..4ac694931666 100644
--- a/vllm/compilation/backends.py
+++ b/vllm/compilation/backends.py
@@ -73,6 +73,7 @@ def make_copy_and_call(
 
     Returns:
         A wrapper function that copies inputs and calls the compiled function
+
     """
 
     def copy_and_call(*args: Any) -> Any:
@@ -122,8 +123,7 @@ def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
 
 
 class CompilerManager:
-    """
-    A manager to manage the compilation process, including
+    """A manager to manage the compilation process, including
     caching the compiled graph, loading the compiled graph,
     and compiling the graph.
 
@@ -163,8 +163,7 @@ def compile_context(self, compile_range: Range) -> Generator[None, None, None]:
     def initialize_cache(
         self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
     ) -> None:
-        """
-        Initialize the cache directory for the compiler.
+        """Initialize the cache directory for the compiler.
 
         The organization of the cache directory is as follows:
         cache_dir=/path/to/hash_str/rank_i_j/prefix/
@@ -177,7 +176,6 @@ def initialize_cache(
         base cache dir of /path/to/hash_str/rank_i_j/ ,
         to store some common compilation artifacts.
         """
-
         self.disable_cache = disable_cache
         self.cache_dir = cache_dir
         self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py")
@@ -438,12 +436,10 @@ def _merge_empty_only_subgraphs(
     node_to_subgraph_id: dict[fx.Node, int],
     split_op_graphs: list[int],
 ) -> None:
-    """
-    Merge a partition that only contains an empty allocation op into the
+    """Merge a partition that only contains an empty allocation op into the
     previous partition. This avoids generating standalone empty submodules,
     which can lead to empty cudagraph captures.
     """
-
     nodes_by_subgraph_id: dict[int, list[fx.Node]] = defaultdict(list)
     for node, subgraph_id in node_to_subgraph_id.items():
         nodes_by_subgraph_id[subgraph_id].append(node)
@@ -637,8 +633,7 @@ def wrap_with_cudagraph_if_needed(
     is_first_graph: bool,
     is_last_graph: bool,
 ) -> Any:
-    """
-    Wrap a piecewise backend with CUDA graph wrapper if needed.
+    """Wrap a piecewise backend with CUDA graph wrapper if needed.
     This function is shared between VllmBackend and
     construct_serializable_fn_from_inductor_cache.
 
@@ -651,6 +646,7 @@ def wrap_with_cudagraph_if_needed(
 
     Returns:
         The wrapped backend if CUDA graphs are enabled, otherwise the original backend
+
     """
     if (
         not compilation_config.cudagraph_mode.has_piecewise_cudagraphs()
@@ -883,8 +879,8 @@ def collect_standalone_compile_artifacts(
                   sym_shape_indices
                 - returns_tuple_map: dict mapping submod_name to
                   returns_tuple
-        """
 
+        """
         if not envs.VLLM_USE_MEGA_AOT_ARTIFACT:
             return None, None, None
 
diff --git a/vllm/compilation/base_static_graph.py b/vllm/compilation/base_static_graph.py
index 12f1ff5bc044..b645af104313 100644
--- a/vllm/compilation/base_static_graph.py
+++ b/vllm/compilation/base_static_graph.py
@@ -8,8 +8,7 @@
 
 
 class AbstractStaticGraphWrapper(Protocol):
-    """
-    StaticGraphWrapper interface that allows platforms to wrap a callable
+    """StaticGraphWrapper interface that allows platforms to wrap a callable
     to be captured as a static graph.
     """
 
@@ -20,8 +19,7 @@ def __init__(
         runtime_mode: CUDAGraphMode,
         **kwargs: Any,
     ) -> None:
-        """
-        Initializes the StaticGraphWrapper class with graph capturing and
+        """Initializes the StaticGraphWrapper class with graph capturing and
         execution-related configurations.
 
         Args:
@@ -31,15 +29,16 @@ def __init__(
                 graph runtime. See CUDAGraphMode in vllm/config.py.
                 Note that only the subset enum `NONE`, `PIECEWISE` and `FULL`
                 are used as concrete runtime mode for cudagraph dispatching.
+
         Keyword Args:
             kwargs: Additional keyword arguments for platform-specific
                 configurations.
+
         """
         raise NotImplementedError
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        """
-        Executes the wrapped callable.
+        """Executes the wrapped callable.
 
         If the current runtime mode in the ForwardContext matches the runtime
         mode of this instance, it replays the CUDAGraph or captures it using
@@ -53,5 +52,6 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
 
         Returns:
             Any: Output of the executed callable.
+
         """
         raise NotImplementedError
diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py
index 343b9ef398bd..ca6aa1713774 100644
--- a/vllm/compilation/breakable_cudagraph.py
+++ b/vllm/compilation/breakable_cudagraph.py
@@ -42,6 +42,7 @@
     is_forward_context_available,
 )
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
 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
@@ -56,6 +57,12 @@ def is_breakable_cudagraph_enabled() -> bool:
 F = TypeVar("F", bound=Callable[..., Any])
 
 
+def _weak_ref_capture_arg(arg: Any) -> Any:
+    if isinstance(arg, QuantizedActivation):
+        return arg.weak_ref()
+    return weak_ref_tensor(arg)
+
+
 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.
@@ -105,13 +112,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
         # 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()
-        }
+        weak_args = tuple(_weak_ref_capture_arg(a) for a in args)
+        weak_kwargs = {k: _weak_ref_capture_arg(v) for k, v in kwargs.items()}
         return capture.add_eager(lambda: fn(*weak_args, **weak_kwargs))
 
     return wrapper  # type: ignore[return-value]
@@ -254,6 +256,10 @@ class BreakableCUDAGraphWrapper:
         * If runtime mode mismatch / NONE, run eagerly.
         * Otherwise, lazily capture per ``batch_descriptor`` and replay
           on subsequent invocations with the same descriptor.
+
+    When ``runtime_mode`` is given, only that mode is captured/replayed;
+    other non-NONE modes fall through, so the wrapper can be nested under
+    a ``CUDAGraphWrapper`` for another mode.
     """
 
     _all_instances: ClassVar[weakref.WeakSet[BreakableCUDAGraphWrapper]] = (
@@ -269,15 +275,14 @@ def __init__(
         self,
         runnable: Callable[..., Any],
         vllm_config: VllmConfig,
+        runtime_mode: CUDAGraphMode | None = None,
     ) -> 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.
+        # runtime_mode=None matches any non-NONE dispatch mode; entries
+        # are keyed by BatchDescriptor, which already encodes batch
+        # shape / uniformity.
         self.runnable = runnable
         self.vllm_config = vllm_config
+        self.runtime_mode = runtime_mode
         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"
@@ -296,7 +301,9 @@ def __getattr__(self, key: str) -> Any:
         raise AttributeError(key)
 
     def unwrap(self) -> Callable[..., Any]:
-        return self.runnable
+        # Recurse through nested wrappers to reach the original runnable.
+        runnable = self.runnable
+        return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable
 
     @property
     def cudagraph_wrapper(self) -> BreakableCUDAGraphWrapper:
@@ -315,11 +322,13 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         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:
+        # By default (runtime_mode=None) match any non-NONE mode; a set
+        # runtime_mode restricts capture so other modes can be handled by
+        # an outer wrapper (e.g. CUDAGraphWrapper for FULL).
+        if cudagraph_runtime_mode == CUDAGraphMode.NONE or (
+            self.runtime_mode is not None
+            and cudagraph_runtime_mode != self.runtime_mode
+        ):
             return self.runnable(*args, **kwargs)
 
         assert batch_descriptor is not None
diff --git a/vllm/compilation/caching.py b/vllm/compilation/caching.py
index 94456fc18271..5205c7404fcb 100644
--- a/vllm/compilation/caching.py
+++ b/vllm/compilation/caching.py
@@ -172,8 +172,7 @@ def patch_pytree_map_over_slice():
 
 
 class VllmSerializableFunction(SerializableCallable):  # type: ignore[misc]
-    """
-    A wrapper around a compiled function by vllm. It will forward the tensor
+    """A wrapper around a compiled function by vllm. It will forward the tensor
     inputs to the compiled function and return the result.
     It also implements a serialization interface to support PyTorch's precompile
     with custom backend, so that we can save and load the compiled function on
@@ -410,9 +409,7 @@ def finalize_loading(self, vllm_config: VllmConfig) -> None:
 
     @property
     def co_name(self) -> Literal["VllmSerializableFunction"]:
-        """
-        Used for depyf debugging.
-        """
+        """Used for depyf debugging."""
         return "VllmSerializableFunction"
 
 
@@ -445,6 +442,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
     If modifying the backend creation/wrapping logic, consider updating both.
 
     Args:
+        fake_mode: FakeTensorMode used to unpickle the cached graph.
         state: Deserialized state dict containing graph_module, example_inputs,
             prefix, sym_tensor_indices, is_encoder, etc.
         standalone_compile_artifacts: The StandaloneCompiledArtifacts containing
@@ -455,6 +453,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
 
     Returns:
         A VllmSerializableFunction that can be called directly.
+
     """
     from vllm.compilation.backends import (
         VllmBackend,
diff --git a/vllm/compilation/codegen.py b/vllm/compilation/codegen.py
index 67e9ac843a64..a0eb90d34dc7 100644
--- a/vllm/compilation/codegen.py
+++ b/vllm/compilation/codegen.py
@@ -154,6 +154,7 @@ def generate_execution_code(
         consts is a list of non-primitive constant objects referenced
         by the generated code via __vllm_consts__. These objects are
         kept alive for the lifetime of the compiled function.
+
     """
     code, submod_names, consts = generate_execution_code_with_name(
         split_gm, "execution_fn", with_submod=True
@@ -181,6 +182,7 @@ def compile_execution_fn(
 
     Returns:
         A callable that executes the stitching logic.
+
     """
     trace_structured(
         "artifact",
diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py
index 742ec55e6efa..394799ab8767 100644
--- a/vllm/compilation/compiler_interface.py
+++ b/vllm/compilation/compiler_interface.py
@@ -23,11 +23,11 @@
 
 logger = init_logger(__name__)
 
+_last_inductor_triton_cache_dir: str | None = None
+
 
 class CompilerInterface:
-    """
-    The interface for a compiler that can be used by vLLM.
-    """
+    """The interface for a compiler that can be used by vLLM."""
 
     # The name of the compiler, e.g. inductor.
     # This is a class-level attribute.
@@ -36,8 +36,7 @@ class CompilerInterface:
     def initialize_cache(
         self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
     ) -> None:
-        """
-        when the vLLM process uses `cache_dir` as the cache directory,
+        """When the vLLM process uses `cache_dir` as the cache directory,
         the compiler should initialize itself with the cache directory,
         e.g. by re-directing its own cache directory to a sub-directory.
 
@@ -52,8 +51,7 @@ def initialize_cache(
         pass
 
     def compute_hash(self, vllm_config: VllmConfig) -> str:
-        """
-        Gather all the relevant information from the vLLM config,
+        """Gather all the relevant information from the vLLM config,
         to compute a hash so that we can cache the compiled model.
 
         See [`VllmConfig.compute_hash`][vllm.config.VllmConfig.compute_hash]
@@ -71,8 +69,7 @@ def compile(
         compile_range: Range,
         key: str | None = None,
     ) -> tuple[Callable[..., Any] | None, Any | None]:
-        """
-        Compile the graph with the given example inputs and compiler config,
+        """Compile the graph with the given example inputs and compiler config,
         with a range. The `compile_range` specifies the range of the inputs,
         it could be concrete size (if compile_sizes is provided), e.g. [4, 4]
         or a range [5, 8].
@@ -105,8 +102,7 @@ def load(
         graph_index: int,
         compile_range: Range,
     ) -> Callable[..., Any]:
-        """
-        Load the compiled function from the handle.
+        """Load the compiled function from the handle.
         Raises an error if the handle is invalid.
 
         The handle is the second return value of the `compile` function.
@@ -115,8 +111,7 @@ def load(
 
 
 class AlwaysHitShapeEnv:
-    """
-    Why do we need this class:
+    """Why do we need this class:
 
     For normal `torch.compile` usage, every compilation will have
     one Dynamo bytecode compilation and one Inductor compilation.
@@ -249,8 +244,7 @@ def _save(
 
 
 class InductorStandaloneAdaptor(CompilerInterface):
-    """
-    The adaptor for the Inductor compiler.
+    """The adaptor for the Inductor compiler.
     Requires PyTorch 2.8+.
     This is not on by default yet, but we plan to turn it on by default for
     PyTorch 2.8.
@@ -447,9 +441,7 @@ def compiled_graph_wrapper(*args: Any) -> tuple[Any, ...] | Any:
 
 
 class InductorAdaptor(CompilerInterface):
-    """
-    The adaptor for the Inductor compiler, version 2.5, 2.6, 2.7.
-    """
+    """The adaptor for the Inductor compiler, version 2.5, 2.6, 2.7."""
 
     name = "inductor"
 
@@ -463,19 +455,22 @@ def compute_hash(self, vllm_config: VllmConfig) -> str:
     def initialize_cache(
         self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
     ) -> None:
+        global _last_inductor_triton_cache_dir
         self.cache_dir = cache_dir
         self.prefix = prefix
         self.base_cache_dir = cache_dir[: -len(prefix)] if prefix else cache_dir
         if disable_cache:
             return
-        # redirect the cache directory to a subdirectory
-        # set flags so that Inductor and Triton store their cache
-        # in the cache_dir, then users only need to copy the cache_dir
-        # to another machine to reuse the cache.
+        # Keep artifacts together unless the caller supplies a Triton cache.
         inductor_cache = os.path.join(self.base_cache_dir, "inductor_cache")
         os.makedirs(inductor_cache, exist_ok=True)
         os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache
-        triton_cache = os.path.join(self.base_cache_dir, "triton_cache")
+        triton_cache = os.environ.get("TRITON_CACHE_DIR")
+        # A directory assigned by an earlier adaptor is still a model-local
+        # default, rather than a caller override for subsequent models.
+        if triton_cache is None or triton_cache == _last_inductor_triton_cache_dir:
+            triton_cache = os.path.join(self.base_cache_dir, "triton_cache")
+            _last_inductor_triton_cache_dir = triton_cache
         os.makedirs(triton_cache, exist_ok=True)
         os.environ["TRITON_CACHE_DIR"] = triton_cache
 
@@ -727,8 +722,7 @@ def compiled_graph(*args: Any) -> tuple[Any, ...] | Any:
         return compiled_graph
 
     def metrics_context(self) -> contextlib.AbstractContextManager[Any]:
-        """
-        This method returns the Dynamo metrics context (if it exists,
+        """This method returns the Dynamo metrics context (if it exists,
         otherwise a null context). It is used by various compile components.
         Present in torch>=2.6, it's used inside FxGraphCache in
         torch==2.6 (but not after). It might also be used in various other
diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py
index b63d86199720..1306cd612f4e 100644
--- a/vllm/compilation/cuda_graph.py
+++ b/vllm/compilation/cuda_graph.py
@@ -38,7 +38,7 @@ class CUDAGraphStat:
 
 
 class CUDAGraphLogging:
-    """Aggregate and log cudagraph metrics"""
+    """Aggregate and log cudagraph metrics."""
 
     COLUMN_HEADERS = [
         "Unpadded Tokens",
@@ -221,7 +221,9 @@ def __getattr__(self, key: str) -> Any:
 
     def unwrap(self) -> Callable[..., Any]:
         # in case we need to access the original runnable.
-        return self.runnable
+        runnable = self.runnable
+        # Recurse through nested wrappers.
+        return runnable.unwrap() if hasattr(runnable, "unwrap") else runnable
 
     @property
     def cudagraph_wrapper(self) -> "CUDAGraphWrapper":
diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py
index d5a2ac7e6813..35c0e6dee5ee 100644
--- a/vllm/compilation/decorators.py
+++ b/vllm/compilation/decorators.py
@@ -7,7 +7,7 @@
 import os
 import sys
 from collections.abc import Callable, Generator
-from typing import TYPE_CHECKING, Any, TypeVar, overload
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, overload
 from unittest.mock import patch
 
 import torch
@@ -32,6 +32,9 @@
 
 from .monitor import monitor_profiling_run, monitor_torch_compile
 
+DynamicArgDims: TypeAlias = dict[str, int | list[int] | dict[int, str]]
+"""Argument name -> the dynamic dimension(s) of that argument."""
+
 # shape_id parameter was added to mark_unbacked in PyTorch 2.11.0
 _SUPPORTS_SHAPE_ID = is_torch_equal_or_newer("2.11.0")
 
@@ -56,8 +59,7 @@ def should_torch_compile_mm_encoder(vllm_config: VllmConfig) -> bool:
 
 
 def ignore_torch_compile(cls: type[_T]) -> type[_T]:
-    """
-    A decorator to ignore support_torch_compile decorator
+    """A decorator to ignore support_torch_compile decorator
     on the class. This is useful when a parent class has
     a support_torch_compile decorator, but we don't want to
     compile the class `cls` that inherits the parent class.
@@ -76,9 +78,7 @@ def ignore_torch_compile(cls: type[_T]) -> type[_T]:
 
 
 def _should_ignore_torch_compile(cls: type[_T]) -> bool:
-    """
-    Check if the class should be ignored for torch.compile.
-    """
+    """Check if the class should be ignored for torch.compile."""
     return getattr(cls, IGNORE_COMPILE_KEY, False)
 
 
@@ -92,7 +92,7 @@ def support_torch_compile(
 @overload
 def support_torch_compile(
     *,
-    dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None,
+    dynamic_arg_dims: DynamicArgDims | None,
 ) -> Callable[[type[_T]], type[_T]]: ...
 
 
@@ -106,7 +106,7 @@ def support_torch_compile(
 @overload
 def support_torch_compile(
     *,
-    dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None,
+    dynamic_arg_dims: DynamicArgDims | None = None,
     mark_unbacked_dims: dict[str, int | list[int]] | None = None,
     enable_if: Callable[[VllmConfig], bool] | None = None,
     is_encoder: bool = False,
@@ -120,13 +120,12 @@ def support_torch_compile(cls: type[_T]) -> type[_T]: ...
 def support_torch_compile(
     cls: type[_T] | None = None,
     *,
-    dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None,
+    dynamic_arg_dims: DynamicArgDims | None = None,
     mark_unbacked_dims: dict[str, int | list[int]] | None = None,
     enable_if: Callable[[VllmConfig], bool] | None = None,
     is_encoder: bool = False,
 ) -> Callable[[type[_T]], type[_T]] | type[_T]:
-    """
-    A decorator to add support for compiling the forward method of a class.
+    """A decorator to add support for compiling the forward method of a class.
 
     Usage 1: use directly as a decorator without arguments:
 
@@ -332,13 +331,12 @@ def _try_load_aot_compiled_fn(
 
 def _support_torch_compile(
     cls: type[_T],
-    dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]],
+    dynamic_arg_dims: DynamicArgDims,
     mark_unbacked_dims: dict[str, int | list[int]] | None = None,
     enable_if: Callable[[VllmConfig], bool] | None = None,
     is_encoder: bool = False,
 ) -> type[_T]:
     """Internal implementation of support_torch_compile decorator."""
-
     if TorchCompileWithNoGuardsWrapper in cls.__bases__:
         # support decorating multiple times
         return cls
@@ -736,8 +734,7 @@ def save_aot_compiled_function(self: type[_T]) -> None:
 def maybe_use_cudagraph_partition_wrapper(
     vllm_config: VllmConfig,
 ) -> Generator[None, None, None]:
-    """
-    Context manager to set/unset customized cudagraph partition wrappers.
+    """Context manager to set/unset customized cudagraph partition wrappers.
 
     If we're using Inductor-based graph partitioning, we currently have the
     whole `fx.Graph` before Inductor lowering and the piecewise
diff --git a/vllm/compilation/partition_rules.py b/vllm/compilation/partition_rules.py
index 18ebb15d1112..8c2e5d6d2b33 100644
--- a/vllm/compilation/partition_rules.py
+++ b/vllm/compilation/partition_rules.py
@@ -12,12 +12,10 @@
 
 
 def should_split(node: torch.fx.Node, splitting_ops: list[str]) -> bool:
-    """
-    Check if a node should be split for dynamo graph partition.
+    """Check if a node should be split for dynamo graph partition.
     It operates on dynamo graph, so the node.target can be anything.
     We need to check and split only on OpOverload and OpOverloadPacket.
     """
-
     if node.op != "call_function":
         return False
 
@@ -50,6 +48,7 @@ def inductor_partition_rule_context(
 
     Args:
         splitting_ops: List of operator names to partition on.
+
     """
     if not splitting_ops:
         logger.debug("No partition ops provided; skipping rule registration.")
diff --git a/vllm/compilation/passes/fusion/act_quant_fusion.py b/vllm/compilation/passes/fusion/act_quant_fusion.py
index c58ce31bd29c..ffafb57d9d34 100644
--- a/vllm/compilation/passes/fusion/act_quant_fusion.py
+++ b/vllm/compilation/passes/fusion/act_quant_fusion.py
@@ -45,8 +45,7 @@
 
 
 class ActivationQuantPattern(VllmPatternReplacement):
-    """
-    Base class for Activation+Quant fusions.
+    """Base class for Activation+Quant fusions.
     Should not be used directly.
     """
 
@@ -79,9 +78,7 @@ def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
 
 class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):
-    """
-    Fusion for SiluMul+Fp8StaticQuant Pattern
-    """
+    """Fusion for SiluMul+Fp8StaticQuant Pattern."""
 
     def __init__(self) -> None:
         super().__init__(kFp8StaticTensorSym)
@@ -126,9 +123,7 @@ def _replacement(
 
 
 class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
-    """
-    Fusion for SiluMul+Nvfp4Quant Pattern
-    """
+    """Fusion for SiluMul+Nvfp4Quant Pattern."""
 
     def __init__(self) -> None:
         super().__init__(kNvfp4Dynamic)
@@ -182,8 +177,7 @@ def _replacement(
 
 
 class SiluMulBlockQuantPattern(ActivationQuantPattern):
-    """
-    Fusion for SiluMul+BlockQuant (FP8 dynamic per-group) Pattern.
+    """Fusion for SiluMul+BlockQuant (FP8 dynamic per-group) Pattern.
     Supports group_size 128 and 64 via QuantKey.
     Parameterized on is_scale_transposed for different scale layouts.
     """
@@ -281,8 +275,7 @@ def _replacement(
 
 
 class ActivationQuantFusionPass(VllmFusionPatternMatcherPass):
-    """
-    This pass fuses a pre-defined set of custom ops into fused ops.
+    """This pass fuses a pre-defined set of custom ops into fused ops.
     It uses the torch pattern matcher to find the patterns and replace them.
 
     Because patterns can only be registered once, the pass is a singleton.
diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py
index 5e7570ff0d81..b2961debcec2 100644
--- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py
@@ -347,8 +347,7 @@ def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
 
 class AllReduceRMSNormPattern(BasePattern):
-    """
-    This pattern replaces the allreduce + rms norm (without residual)
+    """This pattern replaces the allreduce + rms norm (without residual)
     with fused flashinfer implementation.
     Applies to allreduce + rmsnorm before attn in the first Transformer block.
     """
@@ -409,8 +408,7 @@ def replacement(
 
 
 class AllReduceFusedAddRMSNormPattern(BasePattern):
-    """
-    This pattern replaces the allreduce + rms norm (with residual)
+    """This pattern replaces the allreduce + rms norm (with residual)
     with fused flashinfer implementation.
     Applies to o_proj + rmsnorm after attn and mlp + rmsnorm before attn.
     """
@@ -609,8 +607,7 @@ def replacement(
 
 
 class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
-    """
-    This pattern replaces the allreduce + rms norm (without residual)
+    """This pattern replaces the allreduce + rms norm (without residual)
     + static fp8 quant with fused flashinfer implementation.
     Applies to allreduce + rmsnorm + quant before attn
     in the first Transformer block.
@@ -684,8 +681,7 @@ def replacement(
 
 
 class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
-    """
-    This pattern replaces the allreduce + rms norm (with residual)
+    """This pattern replaces the allreduce + rms norm (with residual)
     + static fp8 quant with fused flashinfer implementation.
     Applies to o_proj + rmsnorm after attn + quant and
     mlp + rmsnorm + quant before attn.
@@ -767,8 +763,7 @@ def replacement(
 
 
 class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
-    """
-    This pattern replaces the allreduce + rms norm (without residual)
+    """This pattern replaces the allreduce + rms norm (without residual)
     + static nvfp4 quant with fused flashinfer implementation.
     Applies to allreduce + rmsnorm + quant before attn
     in the first Transformer block.
@@ -864,8 +859,7 @@ def replacement(
 
 
 class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
-    """
-    This pattern replaces the allreduce + rms norm (with residual)
+    """This pattern replaces the allreduce + rms norm (with residual)
     + static nvfp4 quant with fused flashinfer implementation.
     Applies to o_proj + rmsnorm after attn + quant and
     mlp + rmsnorm + quant before attn.
diff --git a/vllm/compilation/passes/fusion/attn_quant_fusion.py b/vllm/compilation/passes/fusion/attn_quant_fusion.py
index 3e2ed2bc707e..bdbca79e5d18 100644
--- a/vllm/compilation/passes/fusion/attn_quant_fusion.py
+++ b/vllm/compilation/passes/fusion/attn_quant_fusion.py
@@ -7,6 +7,7 @@
 import torch
 from torch._higher_order_ops.auto_functionalize import auto_functionalized
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
@@ -35,9 +36,10 @@
 _FP8_QUANT_KEY = QuantKey(dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=True)
 
 
-class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
-    """
-    Fusion for Attention+Fp8StaticQuant.
+class AttnFp8StaticQuantPattern(
+    VllmPatternReplacement[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]]
+):
+    """Fusion for Attention+Fp8StaticQuant.
 
     Only triggers when the attention implementation returns True in
     `fused_output_quant_supported()`. If the pattern is found, the
@@ -52,8 +54,20 @@ def __init__(self, layer: Attention, dtype: torch.dtype):
         self._dtype = dtype
         self._quant_matcher = MatcherQuantFP8(_FP8_QUANT_KEY)
 
+    def _quantize(
+        self, value: torch.Tensor, scale: torch.Tensor
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        return self._quant_matcher(value, scale)[0]
+
+    def _output(
+        self, value: torch.Tensor, scale: torch.Tensor
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        return value
+
     @property
-    def pattern(self) -> Callable[..., torch.Tensor]:
+    def pattern(
+        self,
+    ) -> Callable[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]]:
         # When _USE_LAYERNAME is enabled (torch >= 2.11), layer_name is
         # passed as an explicit pattern input so the pattern matcher
         # treats it as a wildcard matching hoisted LayerName placeholders.
@@ -79,7 +93,7 @@ def _pattern_with_ln(  # type: ignore[misc]
                 attn_out_view = RESHAPE_OP(
                     at1[1], [q.shape[0], self._num_heads * self._head_size]
                 )
-                return self._quant_matcher(attn_out_view, scale)[0]
+                return self._quantize(attn_out_view, scale)
 
             return _pattern_with_ln
 
@@ -98,12 +112,14 @@ def _pattern(q, k, v, output_attn, scale, kv_cache_dummy_dep):
             attn_out_view = RESHAPE_OP(
                 at1[1], [q.shape[0], self._num_heads * self._head_size]
             )
-            return self._quant_matcher(attn_out_view, scale)[0]
+            return self._quantize(attn_out_view, scale)
 
         return _pattern
 
     @property
-    def replacement(self) -> Callable[..., torch.Tensor]:
+    def replacement(
+        self,
+    ) -> Callable[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]]:
         _ln = _encode_layer_name(self._layer_name)
 
         if _USE_LAYERNAME:
@@ -127,7 +143,9 @@ def _replacement_with_ln(  # type: ignore[misc]
                     output_block_scale=None,
                     kv_cache_dummy_dep=kv_cache_dummy_dep,
                 )
-                return RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size])
+                return self._output(
+                    RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size]), scale
+                )
 
             return _replacement_with_ln
 
@@ -148,7 +166,9 @@ def _replacement(q, k, v, output_attn, scale, kv_cache_dummy_dep):
                 output_block_scale=None,
                 kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
-            return RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size])
+            return self._output(
+                RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size]), scale
+            )
 
         return _replacement
 
@@ -169,11 +189,34 @@ def get_inputs(self):
         return inputs
 
 
+class RocmAttnFp8StaticQuantPattern(AttnFp8StaticQuantPattern):
+    """Fuse AITER static quantization while preserving its scale output."""
+
+    def _quantize(
+        self, value: torch.Tensor, scale: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        out = torch.empty_like(value, dtype=FP8_DTYPE)
+        _, out, scale = auto_functionalized(
+            torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
+            out=out,
+            x=value,
+            scale=scale,
+            is_dynamic=False,
+        )
+        # The shared AITER op declares scale mutable for its dynamic variant.
+        # Static quant preserves it; keep this graph output when replacing it.
+        return out, scale
+
+    def _output(
+        self, value: torch.Tensor, scale: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return value, scale
+
+
 class AttnNvfp4QuantPattern(
     VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]
 ):
-    """
-    Fusion for Attention+Nvfp4Quant.
+    """Fusion for Attention+Nvfp4Quant.
 
     Only triggers when the attention implementation returns True in
     `fused_output_quant_supported()`. If the pattern is found, the
@@ -360,8 +403,7 @@ def get_inputs(self):
 
 
 class AttnQuantFusionPass(VllmFusionPatternMatcherPass):
-    """
-    This pass fuses post-attention quantization onto attention if supported.
+    """This pass fuses post-attention quantization onto attention if supported.
 
     It uses the pattern matcher and matches each layer manually, as strings
     cannot be wildcarded. This also lets us check support on attention layers
@@ -390,6 +432,8 @@ def __init__(self, config: VllmConfig) -> None:
         for layer in layers:
             if layer.impl.fused_output_quant_supported(_FP8_QUANT_KEY):
                 self.register(AttnFp8StaticQuantPattern(layer, dtype))
+                if current_platform.is_rocm() and rocm_aiter_ops.is_enabled():
+                    self.register(RocmAttnFp8StaticQuantPattern(layer, dtype))
                 if _USE_LAYERNAME:
                     break
 
diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py
index 7a413665ffa6..f2aa6d7d2f1d 100644
--- a/vllm/compilation/passes/fusion/collective_fusion.py
+++ b/vllm/compilation/passes/fusion/collective_fusion.py
@@ -339,6 +339,8 @@ def __init__(self, dtype: torch.dtype, device: str | None) -> None:
 
 
 class GEMMReduceScatterPattern(BasePattern):
+    rocm_gemm = False
+
     def get_inputs(self) -> list[torch.Tensor]:
         mul = torch.empty([16, 4], device=self.device, dtype=self.dtype)
         mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -346,7 +348,10 @@ def get_inputs(self) -> list[torch.Tensor]:
 
     def register(self, pm_pass: PatternMatcherPass) -> None:
         def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
-            mm = torch.ops.aten.mm.default(mul, mm_weight)
+            if self.rocm_gemm:
+                mm = torch.ops.vllm.rocm_unquantized_gemm(mul, mm_weight, None)
+            else:
+                mm = torch.ops.aten.mm.default(mul, mm_weight)
             reduce_scatter = torch.ops.vllm.reduce_scatter.default(
                 mm,
                 dim=0,
@@ -358,7 +363,7 @@ def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
         def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
             gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter(
                 mul,
-                mm_weight,
+                mm_weight.t() if self.rocm_gemm else mm_weight,
                 "sum",
                 scatter_dim=0,
                 group_name=self.tp.device_group.group_name,
@@ -371,7 +376,13 @@ def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor:
         )
 
 
+class RocmGEMMReduceScatterPattern(GEMMReduceScatterPattern):
+    rocm_gemm = True
+
+
 class AllGatherGEMMPattern(BasePattern):
+    rocm_gemm = False
+
     def get_inputs(self) -> list[torch.Tensor]:
         x = torch.empty([4, 4], device=self.device, dtype=self.dtype)
         weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -390,12 +401,14 @@ def pattern(
                 group_name=self.tp.unique_name,
             )
 
+            if self.rocm_gemm:
+                return torch.ops.vllm.rocm_unquantized_gemm(all_gather, weight, None)
             return torch.ops.aten.mm.default(all_gather, weight)
 
         def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
             ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul(
                 x,
-                [weight],
+                [weight.t() if self.rocm_gemm else weight],
                 gather_dim=0,
                 group_name=self.tp.device_group.group_name,
             )
@@ -406,6 +419,10 @@ def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
         )
 
 
+class RocmAllGatherGEMMPattern(AllGatherGEMMPattern):
+    rocm_gemm = True
+
+
 class ScaledMMReduceScatterPattern(BasePattern):
     def get_inputs(self) -> list[torch.Tensor]:
         input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE)
@@ -906,6 +923,13 @@ def __init__(self, config: VllmConfig) -> None:
         GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.pm_pass)
 
         AllGatherGEMMPattern(self.model_dtype, self.device).register(self.pm_pass)
+        if current_platform.is_rocm() and self.model_dtype == torch.bfloat16:
+            RocmGEMMReduceScatterPattern(self.model_dtype, self.device).register(
+                self.pm_pass
+            )
+            RocmAllGatherGEMMPattern(self.model_dtype, self.device).register(
+                self.pm_pass
+            )
 
         # These fusions are enabled only for bfloat16 models because
         # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling
diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py
index af2719ce3b2d..6d959c30677a 100644
--- a/vllm/compilation/passes/fusion/matcher_utils.py
+++ b/vllm/compilation/passes/fusion/matcher_utils.py
@@ -31,6 +31,7 @@
 
 ROTARY_OP = torch.ops._C.rotary_embedding.default
 FLASHINFER_ROTARY_OP = torch.ops.vllm.flashinfer_rotary_embedding.default
+MROPE_OP = torch.ops.vllm.mrope.default
 
 QUANT_OPS: dict[QuantKey, OpOverload] = {
     kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default,  # noqa: E501
@@ -79,7 +80,7 @@ def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs)
 
     def inputs(self) -> list[torch.Tensor]:
-        """Utility for inputs to the pattern"""
+        """Utility for inputs to the pattern."""
         raise NotImplementedError
 
 
@@ -162,6 +163,74 @@ def forward_native(
         return result
 
 
+class MatcherMRotaryEmbedding(MatcherCustomOp):
+    def __init__(
+        self,
+        is_neox: bool,
+        head_size: int,
+        rotary_dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        mrope_section: tuple[int, int, int],
+        mrope_interleaved: bool,
+        enabled: bool | None = None,
+    ) -> None:
+        if enabled is None:
+            enabled = RotaryEmbedding.enabled()
+
+        super().__init__(enabled)
+        self.is_neox = is_neox
+        self.head_size = head_size
+        self.rotary_dim = rotary_dim
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.q_size = self.num_heads * self.head_size
+        self.kv_size = self.num_kv_heads * self.head_size
+        self.mrope_section = mrope_section
+        self.mrope_interleaved = mrope_interleaved
+
+    def inputs(self) -> list[torch.Tensor]:
+        positions = self.empty_int64(3, 5)
+        query = self.empty(5, self.q_size)
+        key = self.empty(5, self.kv_size)
+        cos_sin_cache = self.empty(4096, self.rotary_dim)
+        return [positions, query, key, cos_sin_cache]
+
+    def forward_custom(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        result = auto_functionalized(
+            MROPE_OP,
+            positions=positions,
+            query=query,
+            key=key,
+            cos_sin_cache=cos_sin_cache,
+            head_size=self.head_size,
+            rotary_dim=self.rotary_dim,
+            mrope_section_t=self.mrope_section[0],
+            mrope_section_h=self.mrope_section[1],
+            mrope_section_w=self.mrope_section[2],
+            mrope_interleaved=self.mrope_interleaved,
+            is_neox_style=self.is_neox,
+        )
+        return result[1], result[2]
+
+    def forward_native(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError(
+            "MRoPE fusion requires the rotary_embedding custom op"
+        )
+
+
 class MatcherRMSNormGated(MatcherCustomOp):
     """Matches RMSNormGated with norm_before_gate=True and group_size=None."""
 
diff --git a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py
index 84c24bc60e53..982d90726e06 100644
--- a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py
+++ b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py
@@ -33,8 +33,7 @@
 
 
 class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
-    """
-    Fusion for MLA Attention+Fp8StaticQuant.
+    """Fusion for MLA Attention+Fp8StaticQuant.
 
     Matches the pattern: MLA attention -> static FP8 quant, and replaces
     it with MLA attention(output_scale=scale, output=fp8_buffer).
@@ -174,8 +173,7 @@ def get_inputs(self) -> list[torch.Tensor]:
 class MLAAttnNvfp4QuantPattern(
     VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]
 ):
-    """
-    Fusion for MLA Attention+Nvfp4Quant.
+    """Fusion for MLA Attention+Nvfp4Quant.
 
     Matches the pattern: MLA attention -> NVFP4 quant, and replaces
     it with MLA attention(output_scale=scale, output_block_scale=block_scale,
@@ -362,8 +360,7 @@ def get_inputs(self) -> list[torch.Tensor]:
 class MLAAttnFp8GroupQuantPattern(
     VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]
 ):
-    """
-    Fusion for MLA Attention+Fp8GroupQuant (per-group dynamic FP8).
+    """Fusion for MLA Attention+Fp8GroupQuant (per-group dynamic FP8).
 
     Matches the pattern: MLA attention -> per_token_group_fp8_quant, and
     replaces it with MLA attention(output_block_scale=group_scale_buffer).
@@ -572,8 +569,7 @@ def get_inputs(self) -> list[torch.Tensor]:
 
 
 class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass):
-    """
-    This pass fuses post-attention quantization onto MLA attention if supported.
+    """This pass fuses post-attention quantization onto MLA attention if supported.
 
     It uses the pattern matcher and matches each MLA layer manually, as strings
     cannot be wildcarded. This also lets us check support on attention layers
diff --git a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py
index 5998a0313840..e89864fe4315 100644
--- a/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py
+++ b/vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py
@@ -35,7 +35,10 @@ def fused_rope_unified_mla_kv_cache_update_impl(
     layer_name: LayerNameType,
 ) -> torch.Tensor:
     layer_name = _resolve_layer_name(layer_name)
-    attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
+    attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
+        layer_name
+    )
+    attn_layer.prepare_kv_cache_update(attn_metadata)
     if layer_slot_mapping is not None:
         ops.concat_and_cache_mla_rope_fused(
             positions,
diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py
index 2c8b9c43bdee..4841d255dbef 100644
--- a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py
+++ b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py
@@ -34,8 +34,7 @@
 
 
 class QkNormRopePattern:
-    """
-    Match the unfused sequence in attention blocks and replace with the fused op.
+    """Match the unfused sequence in attention blocks and replace with the fused op.
 
     Unfused (conceptually):
       q, k, v = split(qkv, [qsz, kvsz, kvsz], -1)
diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py
index 85d82c191b27..f418e8a7ef62 100644
--- a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py
+++ b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import inspect
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import ParamSpec
 
 import torch
@@ -19,12 +19,19 @@
     Attention,
     get_attention_context,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.platforms import current_platform
-from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.utils.torch_utils import (
+    _USE_LAYERNAME,
+    LayerNameType,
+    _encode_layer_name,
+    _resolve_layer_name,
+    direct_register_custom_op,
+)
 
 from ..inductor_pass import enable_fake_mode
 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
-from .matcher_utils import MatcherRotaryEmbedding
+from .matcher_utils import MROPE_OP, MatcherMRotaryEmbedding, MatcherRotaryEmbedding
 from .rms_quant_fusion import empty_bf16, empty_fp32, empty_i64
 
 logger = init_logger(__name__)
@@ -51,8 +58,9 @@ def fused_qk_norm_rope_and_unified_kv_cache_update_impl(
     rms_norm_eps: float,
     cos_sin_cache: torch.Tensor,
     is_neox: bool,
-    layer_name: str = "",
+    layer_name: LayerNameType,
 ) -> torch.Tensor:
+    layer_name = _resolve_layer_name(layer_name)
     _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
     if layer_slot_mapping is not None:
         attn_layer.impl.do_qk_norm_rope_kvcache_update(
@@ -87,7 +95,7 @@ def fused_qk_norm_rope_and_unified_kv_cache_update_fake(
     rms_norm_eps: float,
     cos_sin_cache: torch.Tensor,
     is_neox: bool,
-    layer_name: str = "",
+    layer_name: LayerNameType,
 ) -> torch.Tensor:
     return torch.empty(0, device=qkv.device, dtype=qkv.dtype)
 
@@ -100,14 +108,83 @@ def fused_qk_norm_rope_and_unified_kv_cache_update_fake(
 )
 
 
+def fused_qk_norm_mrope_and_unified_kv_cache_update_impl(
+    q_out: torch.Tensor,
+    qkv: torch.Tensor,
+    positions: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rms_norm_eps: float,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    rotary_dim: int,
+    layer_name: LayerNameType,
+) -> torch.Tensor:
+    layer_name = _resolve_layer_name(layer_name)
+    _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
+    if layer_slot_mapping is not None:
+        if positions.ndim == 1:
+            positions = positions.unsqueeze(0).expand(3, -1).contiguous()
+        attn_layer.impl.do_qk_norm_mrope_kvcache_update(
+            attn_layer,
+            qkv,
+            q_out,
+            positions,
+            q_weight,
+            k_weight,
+            rms_norm_eps,
+            cos_sin_cache,
+            is_neox,
+            (mrope_section_t, mrope_section_h, mrope_section_w),
+            is_interleaved,
+            rotary_dim,
+            kv_cache,
+            layer_slot_mapping,
+        )
+    else:
+        q_out.zero_()
+
+    return torch.empty(0, device=qkv.device, dtype=qkv.dtype)
+
+
+def fused_qk_norm_mrope_and_unified_kv_cache_update_fake(
+    q_out: torch.Tensor,
+    qkv: torch.Tensor,
+    positions: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rms_norm_eps: float,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    rotary_dim: int,
+    layer_name: LayerNameType,
+) -> torch.Tensor:
+    return torch.empty(0, device=qkv.device, dtype=qkv.dtype)
+
+
+direct_register_custom_op(
+    op_name="fused_qk_norm_mrope_and_unified_kv_cache_update",
+    op_func=fused_qk_norm_mrope_and_unified_kv_cache_update_impl,
+    mutates_args=["q_out"],
+    fake_impl=fused_qk_norm_mrope_and_unified_kv_cache_update_fake,
+)
+
+
 # ---------------------------------------------------------------------------
 # Pattern: QK-norm + RoPE + unified_kv_cache_update
 # ---------------------------------------------------------------------------
 
 
 class QkNormRopeKvCachePattern:
-    """
-    Match the unfused sequence:
+    """Match the unfused sequence:
       q, k, v = split(qkv, ...)
       q = rms_norm(q.view(heads), q_weight).view(flat)
       k = rms_norm(k.view(heads), k_weight).view(flat)
@@ -143,19 +220,27 @@ def __init__(
         self.eps = eps
         self.is_neox = is_neox
         self.quant_query = quant_query
+        self.encoded_layer_name = _encode_layer_name(self.layer_name)
+        self.query_quant_group_shape = (
+            layer.query_quant.group_shape
+            if layer.query_quant is not None
+            else GroupShape.PER_TENSOR
+        )
 
         self.q_size = self.num_heads * self.head_size
         self.k_size = self.num_kv_heads * self.head_size
         self.v_size = self.num_kv_heads * self.head_size_v
 
-        self.rope_matcher = MatcherRotaryEmbedding(
-            is_neox=is_neox,
-            head_size=self.head_size,
-            num_heads=self.num_heads,
-            num_kv_heads=self.num_kv_heads,
+        self.rope_matcher: MatcherRotaryEmbedding | MatcherMRotaryEmbedding = (
+            MatcherRotaryEmbedding(
+                is_neox=is_neox,
+                head_size=self.head_size,
+                num_heads=self.num_heads,
+                num_kv_heads=self.num_kv_heads,
+            )
         )
 
-    def get_inputs(self) -> list[torch.Tensor]:
+    def get_inputs(self) -> list:
         T = 5
         L = 4096
         qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
@@ -167,8 +252,13 @@ def get_inputs(self) -> list[torch.Tensor]:
         if self.quant_query:
             q_scale = empty_fp32(1)
             inputs += [q_scale]
+        if _USE_LAYERNAME:
+            inputs.append(self.encoded_layer_name)
         return inputs
 
+    def _get_layer_name(self, layer_name: LayerNameType | None) -> LayerNameType:
+        return self.encoded_layer_name if layer_name is None else layer_name
+
     def pattern_non_fp8_quant_query(
         self,
         qkv: torch.Tensor,
@@ -176,6 +266,7 @@ def pattern_non_fp8_quant_query(
         q_weight: torch.Tensor,
         k_weight: torch.Tensor,
         cos_sin_cache: torch.Tensor,
+        layer_name: LayerNameType | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
         q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size)
@@ -191,7 +282,9 @@ def pattern_non_fp8_quant_query(
         q_rope = q_rope.view(-1, self.num_heads, self.head_size)
         k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size)
         v = v.view(-1, self.num_kv_heads, self.head_size_v)
-        dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name)
+        dummy = torch.ops.vllm.unified_kv_cache_update(
+            k_rope, v, self._get_layer_name(layer_name)
+        )
         return dummy, q_rope, k_rope, v
 
     def replacement_non_fp8_quant_query(
@@ -201,6 +294,7 @@ def replacement_non_fp8_quant_query(
         q_weight: torch.Tensor,
         k_weight: torch.Tensor,
         cos_sin_cache: torch.Tensor,
+        layer_name: LayerNameType | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         q_out = torch.empty(
             qkv.shape[0],
@@ -229,10 +323,49 @@ def replacement_non_fp8_quant_query(
             rms_norm_eps=self.eps,
             cos_sin_cache=cos_sin_cache,
             is_neox=self.is_neox,
-            layer_name=self.layer_name,
+            layer_name=self._get_layer_name(layer_name),
         )
         return results[0], results[1], results[2], v
 
+    def _quantize_query(
+        self,
+        query: torch.Tensor,
+        q_scale: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if self.query_quant_group_shape.is_per_tensor():
+            q_out = torch.empty_like(query, dtype=current_platform.fp8_dtype())
+            q_quant = auto_functionalized(
+                torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
+                out=q_out,
+                x=query,
+                scale=q_scale,
+                is_dynamic=False,
+            )
+            # The AITER op's schema marks scale mutable, so preserve its
+            # functionalized write-back as an output of the matched region.
+            return q_quant[1], q_quant[2]
+
+        # QuantFP8.forward_hip falls back to the generic static op for the
+        # historical per-head shape (-1, block_size). This also covers the
+        # degenerate one-local-KV-head case where a scalar scale is represented
+        # as (-1, q_size), rather than changing Attention's classification.
+        q_out = torch.empty(
+            query.shape,
+            device=query.device,
+            dtype=current_platform.fp8_dtype(),
+        )
+        q_quant = auto_functionalized(
+            torch.ops._C.static_scaled_fp8_quant.default,
+            result=q_out,
+            input=query,
+            scale=q_scale,
+            group_shape=[
+                self.query_quant_group_shape.row,
+                self.query_quant_group_shape.col,
+            ],
+        )
+        return q_quant[1], None
+
     def pattern_fp8_quant_query(
         self,
         qkv: torch.Tensor,
@@ -241,7 +374,8 @@ def pattern_fp8_quant_query(
         k_weight: torch.Tensor,
         cos_sin_cache: torch.Tensor,
         q_scale: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        layer_name: LayerNameType | None = None,
+    ) -> tuple[torch.Tensor, ...]:
         q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
         q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size)
         q_normed = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps)
@@ -252,24 +386,15 @@ def pattern_fp8_quant_query(
         k_flat = k_normed.view(-1, self.k_size)
 
         q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache)
-        # Match the quant-query op Attention.forward inserts (fp8 KV + UNIFIED).
-        # Explicit auto_functionalized (out=[1]) keeps the quant node in the pattern.
-        q_out = torch.empty_like(q_rope, dtype=current_platform.fp8_dtype())
-        q_quant = auto_functionalized(
-            torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
-            out=q_out,
-            x=q_rope,
-            scale=q_scale,
-            is_dynamic=False,
-        )
-        # `scale` is mutable: its copy_ write-back to _q_scale bumps the mutation
-        # region, so keep q flat (a reshape lands past the barrier and won't match).
-        q_rope_fp8 = q_quant[1]
-        q_scale_out = q_quant[2]
+        q_rope_fp8, q_scale_out = self._quantize_query(q_rope, q_scale)
 
         k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size)
         v = v.view(-1, self.num_kv_heads, self.head_size_v)
-        dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name)
+        dummy = torch.ops.vllm.unified_kv_cache_update(
+            k_rope, v, self._get_layer_name(layer_name)
+        )
+        if q_scale_out is None:
+            return dummy, q_rope_fp8, k_rope, v
         return dummy, q_rope_fp8, k_rope, v, q_scale_out
 
     def replacement_fp8_quant_query(
@@ -280,7 +405,8 @@ def replacement_fp8_quant_query(
         k_weight: torch.Tensor,
         cos_sin_cache: torch.Tensor,
         q_scale: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        layer_name: LayerNameType | None = None,
+    ) -> tuple[torch.Tensor, ...]:
         q_out = torch.empty(
             qkv.shape[0],
             self.num_heads,
@@ -308,22 +434,14 @@ def replacement_fp8_quant_query(
             rms_norm_eps=self.eps,
             cos_sin_cache=cos_sin_cache,
             is_neox=self.is_neox,
-            layer_name=self.layer_name,
+            layer_name=self._get_layer_name(layer_name),
         )
-        # Re-apply the quant on the kernel's bf16 q_out; fused op does not quant q.
-        # Same explicit auto_functionalized form as the pattern: [1] = quantized
-        # q, [2] = scale (returned so the buffer-writeback use is preserved).
+        # Re-apply the same quant form on the kernel's bf16 q_out; the fused
+        # operation does not quantize Q.
         q_fp8_flat = results[1].view(-1, self.q_size)
-        q_fp8_out = torch.empty_like(q_fp8_flat, dtype=current_platform.fp8_dtype())
-        q_requant = auto_functionalized(
-            torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
-            out=q_fp8_out,
-            x=q_fp8_flat,
-            scale=q_scale,
-            is_dynamic=False,
-        )
-        q_fp8 = q_requant[1]  # flat to mirror the pattern (see note above)
-        q_scale_out = q_requant[2]
+        q_fp8, q_scale_out = self._quantize_query(q_fp8_flat, q_scale)
+        if q_scale_out is None:
+            return results[0], q_fp8, results[2], v
         return results[0], q_fp8, results[2], v, q_scale_out
 
     @staticmethod
@@ -346,6 +464,9 @@ def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None:
 
         view_to_reshape(gm)
 
+    def _extra_check(self, match: pm.Match) -> bool:
+        return True
+
     def _register(self, pattern, replacement, pm_pass) -> None:
         trace_fn = QkNormRopeKvCachePattern.wrap_trace_fn(
             pm.fwd_only,
@@ -371,40 +492,338 @@ def _register(self, pattern, replacement, pm_pass) -> None:
             inputs,
             trace_fn,
             pm_pass,
+            extra_check=self._extra_check,
             search_fn_pattern=search_fn_pattern,
         )
 
     def register(self, pm_pass: PatternMatcherPass) -> None:
         # make_fx counts `self` in bound-method code params; wrap as plain fns.
-        # Distinct names per branch so mypy doesn't see one name, two signatures.
         if self.quant_query:
-
-            def pattern_q(qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale):
-                return self.pattern_fp8_quant_query(
+            if _USE_LAYERNAME:
+
+                def pattern_q_with_layer(
+                    qkv,
+                    positions,
+                    q_weight,
+                    k_weight,
+                    cos_sin_cache,
+                    q_scale,
+                    layer_name,
+                ):
+                    return self.pattern_fp8_quant_query(
+                        qkv,
+                        positions,
+                        q_weight,
+                        k_weight,
+                        cos_sin_cache,
+                        q_scale,
+                        layer_name,
+                    )
+
+                def replacement_q_with_layer(
+                    qkv,
+                    positions,
+                    q_weight,
+                    k_weight,
+                    cos_sin_cache,
+                    q_scale,
+                    layer_name,
+                ):
+                    return self.replacement_fp8_quant_query(
+                        qkv,
+                        positions,
+                        q_weight,
+                        k_weight,
+                        cos_sin_cache,
+                        q_scale,
+                        layer_name,
+                    )
+
+                self._register(pattern_q_with_layer, replacement_q_with_layer, pm_pass)
+            else:
+
+                def pattern_q_without_layer(
                     qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
-                )
+                ):
+                    return self.pattern_fp8_quant_query(
+                        qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
+                    )
 
-            def replacement_q(
-                qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
-            ):
-                return self.replacement_fp8_quant_query(
+                def replacement_q_without_layer(
                     qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
-                )
+                ):
+                    return self.replacement_fp8_quant_query(
+                        qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
+                    )
 
-            self._register(pattern_q, replacement_q, pm_pass)
+                self._register(
+                    pattern_q_without_layer, replacement_q_without_layer, pm_pass
+                )
         else:
+            if _USE_LAYERNAME:
+
+                def pattern_noq_with_layer(
+                    qkv,
+                    positions,
+                    q_weight,
+                    k_weight,
+                    cos_sin_cache,
+                    layer_name,
+                ):
+                    return self.pattern_non_fp8_quant_query(
+                        qkv,
+                        positions,
+                        q_weight,
+                        k_weight,
+                        cos_sin_cache,
+                        layer_name,
+                    )
+
+                def replacement_noq_with_layer(
+                    qkv,
+                    positions,
+                    q_weight,
+                    k_weight,
+                    cos_sin_cache,
+                    layer_name,
+                ):
+                    return self.replacement_non_fp8_quant_query(
+                        qkv,
+                        positions,
+                        q_weight,
+                        k_weight,
+                        cos_sin_cache,
+                        layer_name,
+                    )
+
+                self._register(
+                    pattern_noq_with_layer, replacement_noq_with_layer, pm_pass
+                )
+            else:
 
-            def pattern_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):
-                return self.pattern_non_fp8_quant_query(
+                def pattern_noq_without_layer(
                     qkv, positions, q_weight, k_weight, cos_sin_cache
-                )
+                ):
+                    return self.pattern_non_fp8_quant_query(
+                        qkv, positions, q_weight, k_weight, cos_sin_cache
+                    )
 
-            def replacement_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):
-                return self.replacement_non_fp8_quant_query(
+                def replacement_noq_without_layer(
                     qkv, positions, q_weight, k_weight, cos_sin_cache
+                ):
+                    return self.replacement_non_fp8_quant_query(
+                        qkv, positions, q_weight, k_weight, cos_sin_cache
+                    )
+
+                self._register(
+                    pattern_noq_without_layer,
+                    replacement_noq_without_layer,
+                    pm_pass,
                 )
 
-            self._register(pattern_noq, replacement_noq, pm_pass)
+
+class QkNormMRopeKvCachePattern(QkNormRopeKvCachePattern):
+    FUSED_OP = torch.ops.vllm.fused_qk_norm_mrope_and_unified_kv_cache_update.default
+
+    def __init__(
+        self,
+        layer: Attention,
+        eps: float,
+        is_neox: bool,
+        quant_query: bool,
+        rotary_dim: int,
+        mrope_section: tuple[int, int, int],
+        is_interleaved: bool,
+    ) -> None:
+        super().__init__(
+            layer=layer,
+            eps=eps,
+            is_neox=is_neox,
+            quant_query=quant_query,
+        )
+        self.rotary_dim = rotary_dim
+        self.mrope_section = mrope_section
+        self.is_interleaved = is_interleaved
+        self.rope_matcher = MatcherMRotaryEmbedding(
+            is_neox=is_neox,
+            head_size=self.head_size,
+            rotary_dim=rotary_dim,
+            num_heads=self.num_heads,
+            num_kv_heads=self.num_kv_heads,
+            mrope_section=mrope_section,
+            mrope_interleaved=is_interleaved,
+        )
+
+    def get_inputs(self) -> list[torch.Tensor]:
+        T = 5
+        L = 4096
+        qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
+        positions = empty_i64(3, T)
+        q_weight = empty_bf16(1, self.head_size)
+        k_weight = empty_bf16(1, self.head_size)
+        cos_sin_cache = empty_bf16(L, self.rotary_dim)
+        inputs = [qkv, positions, q_weight, k_weight, cos_sin_cache]
+        if self.quant_query:
+            inputs.append(empty_fp32(1))
+        if _USE_LAYERNAME:
+            inputs.append(self.encoded_layer_name)
+        return inputs
+
+    def _unfused_kv_views(self, qkv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # The supported backend reads K/V from the cache after the fused write.
+        # Keep view-compatible placeholders for the attention custom-op ABI.
+        _, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
+        k = k.view(qkv.shape[0], self.num_kv_heads, self.head_size)
+        v = v.view(qkv.shape[0], self.num_kv_heads, self.head_size_v)
+        return k, v
+
+    def _fused_results(
+        self,
+        qkv: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        layer_name: LayerNameType | None,
+    ) -> tuple[torch.Tensor, ...]:
+        q_out = torch.empty(
+            qkv.shape[0],
+            self.num_heads,
+            self.head_size,
+            device=qkv.device,
+            dtype=qkv.dtype,
+        )
+        return auto_functionalized(
+            self.FUSED_OP,
+            q_out=q_out,
+            qkv=qkv,
+            positions=positions,
+            q_weight=q_weight,
+            k_weight=k_weight,
+            rms_norm_eps=self.eps,
+            cos_sin_cache=cos_sin_cache,
+            is_neox=self.is_neox,
+            mrope_section_t=self.mrope_section[0],
+            mrope_section_h=self.mrope_section[1],
+            mrope_section_w=self.mrope_section[2],
+            is_interleaved=self.is_interleaved,
+            rotary_dim=self.rotary_dim,
+            layer_name=self._get_layer_name(layer_name),
+        )
+
+    def replacement_non_fp8_quant_query(
+        self,
+        qkv: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        layer_name: LayerNameType | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        k, v = self._unfused_kv_views(qkv)
+        results = self._fused_results(
+            qkv, positions, q_weight, k_weight, cos_sin_cache, layer_name
+        )
+        return results[0], results[1], k, v
+
+    def replacement_fp8_quant_query(
+        self,
+        qkv: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        q_scale: torch.Tensor,
+        layer_name: LayerNameType | None = None,
+    ) -> tuple[torch.Tensor, ...]:
+        k, v = self._unfused_kv_views(qkv)
+        results = self._fused_results(
+            qkv, positions, q_weight, k_weight, cos_sin_cache, layer_name
+        )
+        q_flat = results[1].view(-1, self.q_size)
+        q_fp8, q_scale_out = self._quantize_query(q_flat, q_scale)
+        if q_scale_out is None:
+            return results[0], q_fp8, k, v
+        return results[0], q_fp8, k, v, q_scale_out
+
+    def _extra_check(self, match: pm.Match) -> bool:
+        expected = {
+            "head_size": self.head_size,
+            "rotary_dim": self.rotary_dim,
+            "mrope_section_t": self.mrope_section[0],
+            "mrope_section_h": self.mrope_section[1],
+            "mrope_section_w": self.mrope_section[2],
+            "mrope_interleaved": self.is_interleaved,
+            "is_neox_style": self.is_neox,
+        }
+        mrope_node = next(
+            (node for node in match.nodes if node.args and node.args[0] == MROPE_OP),
+            None,
+        )
+        if mrope_node is None or any(
+            mrope_node.kwargs.get(name) != value for name, value in expected.items()
+        ):
+            return False
+
+        # The fused AITER MRoPE entrypoint writes rotated K directly to the
+        # cache and does not expose a BF16 K result. Returning the raw K view is
+        # valid only for decoder unified attention, whose backend consumes K
+        # from that cache. Reject graphs that observe K in any other way.
+        output_nodes = match.output_nodes()
+        if len(output_nodes) < 3 or output_nodes[2] is None:
+            return False
+        matched_nodes = set(match.nodes)
+        attention_op = torch.ops.vllm.unified_attention_with_output.default
+
+        def is_unified_attention_user(user: fx.Node) -> bool:
+            return user.op == "call_function" and (
+                user.target == attention_op
+                or (user.args and user.args[0] == attention_op)
+            )
+
+        return all(
+            user in matched_nodes or is_unified_attention_user(user)
+            for user in output_nodes[2].users
+        )
+
+
+def _discover_mrope_configs(
+    config: VllmConfig,
+) -> tuple[tuple[tuple[int, int, int], bool], ...]:
+    model_config = config.model_config
+    hf_configs = (
+        getattr(model_config, "hf_text_config", None),
+        getattr(model_config, "hf_config", None),
+    )
+    configs: set[tuple[tuple[int, int, int], bool]] = set()
+    for hf_config in hf_configs:
+        if hf_config is None:
+            continue
+        rope_parameters = getattr(hf_config, "rope_parameters", None)
+        if rope_parameters is None:
+            rope_parameters = getattr(hf_config, "rope_scaling", None)
+        if not isinstance(rope_parameters, Mapping):
+            continue
+        section = rope_parameters.get("mrope_section")
+        if not isinstance(section, (list, tuple)) or len(section) != 3:
+            continue
+        try:
+            parsed_section = (
+                int(section[0]),
+                int(section[1]),
+                int(section[2]),
+            )
+        except (TypeError, ValueError):
+            continue
+        if any(value <= 0 for value in parsed_section):
+            continue
+        configs.add(
+            (
+                parsed_section,
+                bool(rope_parameters.get("mrope_interleaved", False)),
+            )
+        )
+    return tuple(sorted(configs))
 
 
 # ---------------------------------------------------------------------------
@@ -413,8 +832,7 @@ def replacement_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):
 
 
 class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
-    """
-    Fuse QK-norm + RoPE + KV cache update into a single AITER HIP kernel.
+    """Fuse QK-norm + RoPE/MRoPE + KV cache update into an AITER HIP kernel.
 
     Supersedes both QKNormRoPEFusionPass and RopeKVCacheFusionPass for
     attention layers that support the combined operation, eliminating two
@@ -431,22 +849,33 @@ def __init__(self, config: VllmConfig) -> None:
 
         cc = config.compilation_config
         self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num
+        self.mrope_configs = (
+            _discover_mrope_configs(config)
+            if cc.is_custom_op_enabled("rotary_embedding")
+            else ()
+        )
 
         dtype = config.model_config.dtype
         if dtype not in (torch.bfloat16, torch.float16):
             logger.warning_once(
-                "QK Norm+RoPE+KVCache fusion not enabled: unsupported dtype %s", dtype
+                "QK Norm+RoPE/MRoPE+KVCache fusion not enabled: unsupported dtype %s",
+                dtype,
             )
             return
 
         attn_layers = get_layers_from_vllm_config(config, Attention)
 
         for _, layer in attn_layers.items():
-            if not layer.impl.fused_qk_norm_rope_kvcache_supported():
+            supports_rope = layer.impl.fused_qk_norm_rope_kvcache_supported()
+            supports_mrope = (
+                bool(self.mrope_configs)
+                and layer.impl.fused_qk_norm_mrope_kvcache_supported()
+            )
+            if not supports_rope and not supports_mrope:
                 continue
             if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS:
                 logger.warning_once(
-                    "QK Norm+RoPE+KVCache fusion not enabled for a layer: "
+                    "QK Norm+RoPE/MRoPE+KVCache fusion not enabled for a layer: "
                     "head_size=%d is not supported by the "
                     "fused_qk_norm_rope_cache_pts_quant_shuffle kernel "
                     "(supported: %s). Falling back to the unfused path.",
@@ -457,22 +886,71 @@ def __init__(self, config: VllmConfig) -> None:
             if layer.head_size_v != layer.head_size:
                 # The fused kernel uses a single head_dim for q/k/v.
                 logger.warning_once(
-                    "QK Norm+RoPE+KVCache fusion not enabled for a layer: "
+                    "QK Norm+RoPE/MRoPE+KVCache fusion not enabled for a layer: "
                     "head_size_v=%d differs from head_size=%d, which the fused "
                     "kernel does not support. Falling back to the unfused path.",
                     layer.head_size_v,
                     layer.head_size,
                 )
                 continue
-            for epsilon in [1e-5, 1e-6]:
-                for neox in [True, False]:
-                    for quant_q in [False, True]:
-                        QkNormRopeKvCachePattern(
-                            layer=layer,
-                            eps=epsilon,
-                            is_neox=neox,
-                            quant_query=quant_q,
-                        ).register(self.patterns)
+            if supports_rope:
+                for epsilon in [1e-5, 1e-6]:
+                    for neox in [True, False]:
+                        for quant_q in [False, True]:
+                            QkNormRopeKvCachePattern(
+                                layer=layer,
+                                eps=epsilon,
+                                is_neox=neox,
+                                quant_query=quant_q,
+                            ).register(self.patterns)
+
+            if supports_mrope:
+                for section, is_interleaved in self.mrope_configs:
+                    rotary_dim = 2 * sum(section)
+                    if rotary_dim > layer.head_size:
+                        logger.warning_once(
+                            "QK Norm+MRoPE+KVCache fusion not enabled for a "
+                            "layer: rotary_dim=%d exceeds head_size=%d.",
+                            rotary_dim,
+                            layer.head_size,
+                        )
+                        continue
+                    vector_width = layer.head_size // 32
+                    for epsilon in [1e-5, 1e-6]:
+                        for neox in [True, False]:
+                            # Each warp lane processes `vector_width` adjacent
+                            # values. NeoX additionally exchanges vectors
+                            # across the half-rotary boundary.
+                            rotary_alignment = vector_width * (2 if neox else 1)
+                            if rotary_dim % rotary_alignment != 0:
+                                logger.warning_once(
+                                    "QK Norm+MRoPE+KVCache fusion not enabled "
+                                    "for a layer/style: rotary_dim=%d must be "
+                                    "divisible by %d for head_size=%d and "
+                                    "is_neox=%s.",
+                                    rotary_dim,
+                                    rotary_alignment,
+                                    layer.head_size,
+                                    neox,
+                                )
+                                continue
+                            for quant_q in [False, True]:
+                                QkNormMRopeKvCachePattern(
+                                    layer=layer,
+                                    eps=epsilon,
+                                    is_neox=neox,
+                                    quant_query=quant_q,
+                                    rotary_dim=rotary_dim,
+                                    mrope_section=section,
+                                    is_interleaved=is_interleaved,
+                                ).register(self.patterns)
+
+            # Opaque LayerName is a pattern wildcard on torch >= 2.11, so
+            # homogeneous attention layers produce duplicate search patterns.
+            # As in the other LayerName-aware fusion passes, one supported
+            # layer registers all shape/style variants for the model.
+            if _USE_LAYERNAME:
+                break
 
         self.dump_patterns(config, self.patterns)
 
@@ -480,8 +958,8 @@ def __init__(self, config: VllmConfig) -> None:
     def __call__(self, graph: fx.Graph) -> None:
         self.matched_count = self.patterns.apply(graph)
         logger.info(
-            "QK-Norm+RoPE+KVCache fusion: replaced %s pattern(s) "
-            "with AITER fused_qk_norm_rope_cache_pts_quant_shuffle",
+            "QK-Norm+RoPE/MRoPE+KVCache fusion: replaced %s pattern(s) "
+            "with an AITER fused attention-prologue kernel",
             self.matched_count,
         )
 
@@ -489,4 +967,10 @@ def is_applicable_for_range(self, compile_range: Range) -> bool:
         return compile_range.end <= self.max_token_num
 
     def uuid(self) -> str:
-        return VllmInductorPass.hash_source(self, QkNormRopeKvCachePattern)
+        return VllmInductorPass.hash_source(
+            self,
+            QkNormRopeKvCachePattern,
+            QkNormMRopeKvCachePattern,
+            _discover_mrope_configs,
+            repr(self.mrope_configs),
+        )
diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py
index 91c44f572cf0..90897a96435b 100644
--- a/vllm/compilation/passes/fusion/rms_quant_fusion.py
+++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py
@@ -168,8 +168,7 @@ def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor:
 
 
 class FusedRMSQuantKey(NamedTuple):
-    """
-    Named tuple for identifying the type of RMSNorm + quant fusion.
+    """Named tuple for identifying the type of RMSNorm + quant fusion.
     quant: type of quantization
     fused_add: does the op also perform the residual add
     """
@@ -771,8 +770,7 @@ def replacement(
 
 
 class RMSNormQuantFusionPass(VllmPatternMatcherPass):
-    """
-    This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op.
+    """This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op.
     It also supports fused_add_rms_norm.
     """
 
diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py
index 67ed51e89322..19bbfc2d8754 100644
--- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py
+++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py
@@ -30,6 +30,7 @@
     VllmPatternReplacement,
     _fx_view_to_reshape,
     fold_consecutive_reshapes,
+    remove_noop_reshapes,
 )
 from .matcher_utils import (
     MatcherQuantFP8,
@@ -180,8 +181,7 @@ def replacement(
 
 
 class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
-    """
-    This pattern fuses aiter rms_norm & group fp8 quant custom
+    """This pattern fuses aiter rms_norm & group fp8 quant custom
     ops into an aiter rms_norm_group_fp8_quant op.
     """
 
@@ -236,8 +236,7 @@ def replacement(
 
 
 class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
-    """
-    This pattern fuses aiter rms_norm_with_add & group fp8 quant custom ops
+    """This pattern fuses aiter rms_norm_with_add & group fp8 quant custom ops
     into a aiter rms_norm_with_add_group_fp8_quant op.
     """
 
@@ -298,8 +297,7 @@ def replacement(
 
 
 class DoubleAiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
-    """
-    Pattern matching ``rms_norm`` whose output feeds *two* distinct
+    """Pattern matching ``rms_norm`` whose output feeds *two* distinct
     ``rocm_aiter_group_fp8_quant`` consumers, replacing it with two
     independent fused ``rms_norm_group_fp8_quant`` ops.
 
@@ -367,8 +365,7 @@ def replacement(
 
 
 class DoubleAiterRMSFp8GroupQuantViewPattern(AiterRMSNormQuantPattern):
-    """
-    View-tolerant variant of ``DoubleAiterRMSFp8GroupQuantPattern``.
+    """View-tolerant variant of ``DoubleAiterRMSFp8GroupQuantPattern``.
 
     Matches the same 1-to-2 fan-out, but with a ``view``/``reshape`` between
     the ``rms_norm`` output and the two ``rocm_aiter_group_fp8_quant``
@@ -453,13 +450,12 @@ def trace_with_view_to_reshape(*args: Any, **kwargs: Any) -> fx.GraphModule:
 
 
 class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern):
-    """
-    Matches decomposed RMSNormGated + reshape + group FP8 quant and replaces
+    """Matches decomposed RMSNormGated + reshape + group FP8 quant and replaces
     with rocm_aiter_fused_rms_gated_fp8_group_quant.
 
-    The norm operates per-head on (N*H, D) tensors. The compiler folds the
-    reshape chain so after norm the result goes through reshape->merge->quant.
-    The pattern reshapes from (N*H, D) to (N, H*D) before calling
+    The norm operates per-head, on either (N*H, D) or (N, H, D). The compiler
+    folds the reshape chain so after norm the result goes through
+    reshape->merge->quant. The pattern reshapes to (N, H*D) before calling
     MatcherQuantFP8 so that _quantize_group_native sees the full hidden dim
     and computes the correct num_groups.
     """
@@ -487,6 +483,16 @@ def __init__(
         self.head_dim = head_dim
 
     def register(self, pm_pass: PatternMatcherPass) -> None:
+        self._register_layout(pm_pass, flatten_heads=True)
+        # Both layouts reach the norm. With an opaque quant op they trace to
+        # the same pattern; with a native quant they differ, because its group
+        # reshape back to (N, H, D) is a no-op the compiler already dropped.
+        if not self.quant_matcher.enabled:
+            self._register_layout(pm_pass, flatten_heads=False)
+
+    def _register_layout(
+        self, pm_pass: PatternMatcherPass, *, flatten_heads: bool
+    ) -> None:
         num_heads = self.num_heads
         head_dim = self.head_dim
         hidden_dim = num_heads * head_dim
@@ -507,6 +513,11 @@ def replacement(
             z: torch.Tensor,
             weight: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
+            # The pattern matches on op structure rather than rank, so either
+            # per-head layout can land here. The fused op takes (M, D), and
+            # both operands are contiguous, so collapsing the heads is a view.
+            x = x.reshape(-1, head_dim)
+            z = z.reshape(-1, head_dim)
             fused = self.FUSED_OP(
                 x=x,
                 weight=weight,
@@ -524,14 +535,21 @@ def replacement(
             return fp8_reshaped, scales_reshaped
 
         n_tokens = 2
-        x = self.empty(n_tokens * num_heads, head_dim)
-        z = self.empty(n_tokens * num_heads, head_dim)
+        shape = (
+            (n_tokens * num_heads, head_dim)
+            if flatten_heads
+            else (n_tokens, num_heads, head_dim)
+        )
+        x = self.empty(*shape)
+        z = self.empty(*shape)
         w = self.empty(head_dim)
 
         def trace_fn(*args, **kwargs):
             gm = pm.fwd_only(*args, **kwargs)
             _fx_view_to_reshape(gm)
             fold_consecutive_reshapes(gm)
+            if not flatten_heads:
+                remove_noop_reshapes(gm)
             return gm
 
         pm.register_replacement(
@@ -544,8 +562,7 @@ def trace_fn(*args, **kwargs):
 
 
 class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
-    """
-    This pass fuses aiter rms_norm & vllm/aiter quant custom ops
+    """This pass fuses aiter rms_norm & vllm/aiter quant custom ops
     into a fused rms_norm_quant op.
     It also supports fused_add_rms_norm.
     """
@@ -681,8 +698,7 @@ def uuid(self) -> str:
 
 
 class AiterSiluMulFp8GroupQuantPattern(VllmPatternReplacement):
-    """
-    This pattern fuses aiter silu_and_mul & group fp8 quant custom
+    """This pattern fuses aiter silu_and_mul & group fp8 quant custom
     ops into an aiter silu_and_mul_group_fp8_quant op.
     """
 
@@ -722,8 +738,7 @@ def _replacement(
 
 
 class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmFusionPatternMatcherPass):
-    """
-    This pass fuses a pre-defined set of custom ops into fused ops.
+    """This pass fuses a pre-defined set of custom ops into fused ops.
     It uses the torch pattern matcher to find the patterns and replace them.
 
     Because patterns can only be registered once, the pass is a singleton.
@@ -744,8 +759,7 @@ def __init__(self, config: VllmConfig) -> None:
 
 
 class AddAiterRMSNormPadPattern:
-    """
-    This pattern replaces an aiter_rmsnorm_with_add & a pad op
+    """This pattern replaces an aiter_rmsnorm_with_add & a pad op
     with a custom triton_add_rmsnorm_pad op from AITER.
     """
 
@@ -820,8 +834,7 @@ def replacement(
 
 
 class RocmAiterTritonAddRMSNormPadFusionPass(VllmPatternMatcherPass):
-    """
-    This pass replaces an AITER CK RMSNorm + residual add and a pad op
+    """This pass replaces an AITER CK RMSNorm + residual add and a pad op
     with an triton_add_rmsnorm_pad op from AITER.
     """
 
@@ -854,8 +867,7 @@ def uuid(self) -> str:
 class MLADualRMSNormPattern(
     VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
 ):
-    """
-    Fuse paired q_a_layernorm + kv_a_layernorm in MLA attention into
+    """Fuse paired q_a_layernorm + kv_a_layernorm in MLA attention into
     AITER's ``fused_qk_rmsnorm`` HIP kernel.
 
     Target FX-graph pattern (unfused, ``vllm_ir`` stage)::
@@ -946,8 +958,7 @@ class MLADualRMSPerTokenQuantPattern(
         ],
     ]
 ):
-    """
-    Fuse the MLA FP8 attention path -- q-latent RMSNorm + FP8 *per-token* quant
+    """Fuse the MLA FP8 attention path -- q-latent RMSNorm + FP8 *per-token* quant
     plus kv-latent RMSNorm -- into AITER's ``fused_qk_rmsnorm_per_token_quant``.
 
     With a per-token FP8 ``q_b_proj`` (Quark / ModelOpt), the earlier
@@ -1064,8 +1075,7 @@ def _replacement(
 
 
 class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass):
-    """
-    Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in
+    """Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in
     MLA attention into ``fused_mla_dual_rms_norm`` backed by aiter's
     ``fused_qk_rmsnorm`` HIP kernel.
 
diff --git a/vllm/compilation/passes/fusion/rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/rope_kvcache_fusion.py
index 2364281761f7..cc664941d8be 100644
--- a/vllm/compilation/passes/fusion/rope_kvcache_fusion.py
+++ b/vllm/compilation/passes/fusion/rope_kvcache_fusion.py
@@ -52,8 +52,7 @@ def fused_rope_and_unified_kv_cache_update_impl(
     is_neox: bool,
     layer_name: LayerNameType,
 ) -> torch.Tensor:
-    """
-    This impl fetches the KV cache and slot mapping from the forward context,
+    """This impl fetches the KV cache and slot mapping from the forward context,
     then calls the layer impl's `AttentionImpl.do_rope_and_kv_cache_update` method.
     It also returns a dummy tensor, similar to `Attention.unified_kv_cache_update`,
     that is passed to unified_attention to signal a side effect and
@@ -98,8 +97,7 @@ def fused_rope_and_unified_kv_cache_update_fake(
 
 
 class RopeStaticQQuantKVCachePattern:
-    """
-    Fuse rope + static Q fp8 quant while preserving explicit KV-cache update
+    """Fuse rope + static Q fp8 quant while preserving explicit KV-cache update
     dependency ordering.
     """
 
@@ -277,8 +275,7 @@ def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule:
 
 
 class RopeReshapeKVCachePattern:
-    """
-    This pattern matches the following unfused inplace ops:
+    """This pattern matches the following unfused inplace ops:
       q, k = rotary_embedding(positions, q, k, head_size, cos_sin_cache, is_neox)
       kv_cache_dummy = unified_kv_cache_update(k, v, layer_name)
 
@@ -410,8 +407,7 @@ def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule:
 
 
 class RopeKVCacheFusionPass(VllmPatternMatcherPass):
-    """
-    This pass fuses the rotary embedding and KV cache update operations
+    """This pass fuses the rotary embedding and KV cache update operations
     into a single fused kernel if available.
 
     It uses the pattern matcher and matches each layer manually, as strings
diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py
index c4caaaedec22..b1a39b1e0146 100644
--- a/vllm/compilation/passes/fusion/sequence_parallelism.py
+++ b/vllm/compilation/passes/fusion/sequence_parallelism.py
@@ -57,8 +57,7 @@ def get_sequence_parallelism_threshold(
     tp_size: int,
     element_size: int,
 ) -> int | None:
-    """
-    Calculate the minimum token threshold for applying sequence parallelism.
+    """Calculate the minimum token threshold for applying sequence parallelism.
 
     Returns None if sequence parallelism should not be applied based on model size.
 
@@ -496,8 +495,7 @@ def replacement(
 
 
 class SequenceParallelismPass(VllmPatternMatcherPass):
-    """
-    This pass enables sequence parallelism for models.
+    """This pass enables sequence parallelism for models.
     It identifies patterns where an AllReduce operation is followed by
     an RMSNorm (or RMSNorm and then Quantization) operation.
     These patterns are replaced with a ReduceScatter operation, followed by
@@ -590,8 +588,7 @@ def __init__(self, config: VllmConfig) -> None:
         self.dump_patterns(config, self.patterns)
 
     def is_applicable_for_range(self, compile_range: Range) -> bool:
-        """
-        Determines if sequence parallelism should be applied for the given
+        """Determines if sequence parallelism should be applied for the given
         compile range.
 
         SP is only beneficial for larger batch sizes where the communication
diff --git a/vllm/compilation/passes/inductor_pass.py b/vllm/compilation/passes/inductor_pass.py
index 29410f960cda..dadc56c5b2d3 100644
--- a/vllm/compilation/passes/inductor_pass.py
+++ b/vllm/compilation/passes/inductor_pass.py
@@ -64,14 +64,12 @@ def _hash_source_cached(*srcs: str | type | types.FunctionType) -> str:
 
 
 class InductorPass(CustomGraphPass):  # type: ignore[misc]
-    """
-    A custom graph pass that uses a hash of its source as the UUID.
+    """A custom graph pass that uses a hash of its source as the UUID.
     This is defined as a convenience and should work in most cases.
     """
 
     def uuid(self) -> str:
-        """
-        Provide a unique identifier for the pass, used in Inductor code cache.
+        """Provide a unique identifier for the pass, used in Inductor code cache.
         This should depend on the pass implementation, so that changes to the
         pass result in recompilation.
         By default, the object source is hashed.
@@ -80,14 +78,14 @@ def uuid(self) -> str:
 
     @staticmethod
     def hash_source(*srcs: str | Any) -> str:
-        """
-        Utility method to hash the sources of functions or objects.
+        """Utility method to hash the sources of functions or objects.
 
         Args:
             srcs: strings or objects to add to the hash.
                 Objects and functions have their source inspected.
                 Results are cached by resolved types to avoid repeated
                 inspect.getsource() calls.
+
         """
         # Resolve instances to their class for a hashable cache key.
         cache_key = tuple(
@@ -98,11 +96,11 @@ def hash_source(*srcs: str | Any) -> str:
 
     @staticmethod
     def hash_dict(dict_: dict[Any, Any]) -> str:
-        """
-        Utility method to hash a dictionary, can alternatively be used for uuid.
+        """Utility method to hash a dictionary, can alternatively be used for uuid.
 
         Returns:
             A sha256 hash of the json rep of the dictionary.
+
         """
         encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")
         return hashlib.sha256(encoded).hexdigest()
@@ -112,8 +110,7 @@ def is_applicable_for_range(self, compile_range: Range) -> bool:
 
 
 class CallableInductorPass(InductorPass):
-    """
-    This class is a wrapper for a callable that automatically provides an
+    """This class is a wrapper for a callable that automatically provides an
     implementation of the UUID.
     """
 
@@ -131,8 +128,7 @@ def uuid(self) -> Any:
 
 
 def enable_fake_mode(fn: Callable[P, R]) -> Callable[P, R]:
-    """
-    Applies a FakeTensorMode context. This is useful when you don't want to
+    """Applies a FakeTensorMode context. This is useful when you don't want to
     create or run things with real tensors.
     """
 
diff --git a/vllm/compilation/passes/ir/clone_elimination.py b/vllm/compilation/passes/ir/clone_elimination.py
index 377cd669a4eb..d35ebec97f46 100644
--- a/vllm/compilation/passes/ir/clone_elimination.py
+++ b/vllm/compilation/passes/ir/clone_elimination.py
@@ -70,8 +70,7 @@ def user_writes_to_node(user: fx.Node, node: fx.Node) -> bool:
 
 
 class UnsafeCloneEliminationPass(VllmInductorPass):
-    """
-    This pass removes clone nodes that are no longer needed after vLLM IR lowering.
+    """This pass removes clone nodes that are no longer needed after vLLM IR lowering.
     It uses donated_input_ids to eliminate clones of donated graph inputs, preserving
     contents of non-donated graph inputs.
 
diff --git a/vllm/compilation/passes/ir/inplace_functionalization.py b/vllm/compilation/passes/ir/inplace_functionalization.py
index e69351075bca..0c6e66e6d6b6 100644
--- a/vllm/compilation/passes/ir/inplace_functionalization.py
+++ b/vllm/compilation/passes/ir/inplace_functionalization.py
@@ -19,8 +19,7 @@
 
 
 class VllmIRInplaceFunctionalizationPass(VllmInductorPass):
-    """
-    This pass functionalizes maybe_inplace vLLM IR ops to the default overload.
+    """This pass functionalizes maybe_inplace vLLM IR ops to the default overload.
     The maybe_inplace overloads have the same signature as the default overload
     so the pass simply replaces the called overload.
     That makes the graph properly functional.
diff --git a/vllm/compilation/passes/ir/lowering_pass.py b/vllm/compilation/passes/ir/lowering_pass.py
index f34f1c64b76e..f40d23ac473c 100644
--- a/vllm/compilation/passes/ir/lowering_pass.py
+++ b/vllm/compilation/passes/ir/lowering_pass.py
@@ -23,9 +23,7 @@
 
 
 class VllmIRLoweringPass(VllmInductorPass):
-    """
-    This pass lowers vLLM IR ops to their implementations the priority list.
-    """
+    """This pass lowers vLLM IR ops to their implementations the priority list."""
 
     def __init__(self, vllm_config: VllmConfig) -> None:
         super().__init__(vllm_config)
@@ -113,8 +111,7 @@ def print_count(counts: dict[str, int]) -> str:
             logger.warning("Full node list: %s", ",".join(str(n) for n in failed_nodes))
 
     def uuid(self) -> str:
-        """
-        IR op priority & impl sources affect lowering pass output,
+        """IR op priority & impl sources affect lowering pass output,
         so we include them in the cache key.
         """
         priorities = {name: op.get_priority() for name, op in IrOp.registry.items()}
diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py
index 4a65376ca09b..df0e58b21081 100644
--- a/vllm/compilation/passes/pass_manager.py
+++ b/vllm/compilation/passes/pass_manager.py
@@ -41,6 +41,7 @@
 if current_platform.is_cuda_alike():
     from .fusion.act_quant_fusion import ActivationQuantFusionPass
     from .fusion.attn_quant_fusion import AttnQuantFusionPass
+    from .fusion.collective_fusion import AsyncTPPass
     from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass
     from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass
     from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass
@@ -50,7 +51,6 @@
 
 if current_platform.is_cuda():
     from .fusion.allreduce_rms_fusion import AllReduceFusionPass
-    from .fusion.collective_fusion import AsyncTPPass
 
 if current_platform.is_xpu():
     from .fusion.act_quant_fusion import ActivationQuantFusionPass
@@ -71,8 +71,7 @@
 
 
 def with_pattern_match_debug(fn: Callable[P, R]) -> Callable[P, R]:
-    """
-    Function decorator that turns on inductor pattern match debug
+    """Function decorator that turns on inductor pattern match debug
     for the duration of the call.
     Used to avoid logging builtin Inductor pattern matching.
     """
@@ -89,8 +88,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
 
 
 class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
-    """
-    The pass manager for post-grad passes.
+    """The pass manager for post-grad passes.
     It handles configuration, adding custom passes, and running passes.
     It supports uuid for the Inductor code cache. That includes torch<2.6
     support using pickling (in .inductor_pass.CustomGraphPass).
@@ -236,8 +234,7 @@ def add(self, pass_: InductorPass) -> None:
         self.passes.append(pass_)
 
     def uuid(self) -> str:
-        """
-        The PostGradPassManager is set as a custom pass in the Inductor and
+        """The PostGradPassManager is set as a custom pass in the Inductor and
         affects compilation caching. Its uuid depends on the UUIDs of all
         dependent passes and the pass config. See InductorPass for more info.
         """
diff --git a/vllm/compilation/passes/utility/fix_functionalization.py b/vllm/compilation/passes/utility/fix_functionalization.py
index c04f9abbb6a9..b7baac10e0b6 100644
--- a/vllm/compilation/passes/utility/fix_functionalization.py
+++ b/vllm/compilation/passes/utility/fix_functionalization.py
@@ -17,8 +17,7 @@
 
 
 class FixFunctionalizationPass(VllmInductorPass):
-    """
-    This pass defunctionalizes certain nodes to avoid redundant tensor copies.
+    """This pass defunctionalizes certain nodes to avoid redundant tensor copies.
     After this pass, DCE (dead-code elimination) should never be run,
     as de-functionalized nodes may appear as dead code.
 
@@ -260,9 +259,7 @@ def __call__(self, graph: torch.fx.Graph) -> None:
         self.nodes_to_remove.clear()
 
     def _remove(self, node_or_nodes: torch.fx.Node | Iterable[torch.fx.Node]) -> None:
-        """
-        Stage a node (or nodes) for removal at the end of the pass.
-        """
+        """Stage a node (or nodes) for removal at the end of the pass."""
         if isinstance(node_or_nodes, torch.fx.Node):
             self.nodes_to_remove.append(node_or_nodes)
         else:
@@ -275,8 +272,7 @@ def defunctionalize(
         mutated_args: dict[int, torch.fx.Node | str],
         args: tuple[torch.fx.Node | str, ...] | None = None,
     ) -> None:
-        """
-        De-functionalize a node by replacing it with a call to the original.
+        """De-functionalize a node by replacing it with a call to the original.
         It also replaces the getitem users with the mutated arguments.
         See replace_users_with_mutated_args and insert_defunctionalized.
         """
@@ -287,14 +283,14 @@ def defunctionalize(
     def replace_users_with_mutated_args(
         self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]
     ) -> None:
-        """
-        Replace mutated getitem users of the auto-functionalized node with the
+        """Replace mutated getitem users of the auto-functionalized node with the
         mutated arguments.
 
         Args:
             node: The auto-functionalized node
             mutated_args: The mutated arguments, indexed by getitem index.
                 If the value of an arg is a string, `node.kwargs[arg]` is used.
+
         """
         for idx, user in self.getitem_users(node).items():
             # Some functionalized nodes may return both a result at getitem[0]
@@ -310,8 +306,7 @@ def replace_users_with_mutated_args(
             self._remove(user)
 
     def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]:
-        """
-        Returns the operator.getitem users of the auto-functionalized node,
+        """Returns the operator.getitem users of the auto-functionalized node,
         indexed by the index they are getting.
         """
         users = {}
@@ -327,8 +322,7 @@ def insert_defunctionalized(
         node: torch.fx.Node,
         args: tuple[torch.fx.Node | str, ...] | None = None,
     ) -> None:
-        """
-        Insert a new defunctionalized node into the graph before node.
+        """Insert a new defunctionalized node into the graph before node.
         If one of the kwargs is 'out', provide args directly,
         as node.kwargs cannot be used.
         See https://github.com/pytorch/pytorch/blob/a00faf440888ffb724bad413f329a49e2b6388e7/torch/_inductor/lowering.py#L351
@@ -338,6 +332,7 @@ def insert_defunctionalized(
             node: The auto-functionalized node to defunctionalize
             args: If we cannot use kwargs, specify args directly.
                 If an arg is a string, `node.kwargs[arg]` is used.
+
         """  # noqa: E501
         assert is_func(node, auto_functionalized), (
             f"node must be auto-functionalized, is {node} instead"
diff --git a/vllm/compilation/passes/utility/noop_elimination.py b/vllm/compilation/passes/utility/noop_elimination.py
index 80bf8ecc603f..ec9df636f7b2 100644
--- a/vllm/compilation/passes/utility/noop_elimination.py
+++ b/vllm/compilation/passes/utility/noop_elimination.py
@@ -16,8 +16,7 @@
 
 
 class NoOpEliminationPass(VllmInductorPass):
-    """
-    This is an inductor pass that removes redundant reshape/slice operations.
+    """This is an inductor pass that removes redundant reshape/slice operations.
     It is required for RMSNorm-quant fusion to work properly.
     That's because apply_fp8_linear adds a reshape, which is redundant
     in the 2D-case. Additionally, torch internal no-op elimination pass does
@@ -106,8 +105,7 @@ def __call__(self, graph: torch.fx.Graph) -> None:
 
     # ---------------------- Shape comparison helpers ----------------------
     def dims_equivalent(self, dim: int | SymInt, i_dim: int | SymInt) -> bool:
-        """
-        This function checks if two dimensions are equivalent.
+        """This function checks if two dimensions are equivalent.
 
         Args:
             dim: The dimension arg to reshape/slice
@@ -119,6 +117,7 @@ def dims_equivalent(self, dim: int | SymInt, i_dim: int | SymInt) -> bool:
         There are two cases in which the dimensions are equivalent:
         1. The dimensions are equal (both integers)
         2. The dimensions both correspond to the same SymInt
+
         """
         # Case 1
         return statically_known_true(dim == i_dim)  # type: ignore[no-any-return]
diff --git a/vllm/compilation/passes/utility/post_cleanup.py b/vllm/compilation/passes/utility/post_cleanup.py
index d4ecd4d65563..5b0695c61f16 100644
--- a/vllm/compilation/passes/utility/post_cleanup.py
+++ b/vllm/compilation/passes/utility/post_cleanup.py
@@ -6,8 +6,7 @@
 
 
 class PostCleanupPass(VllmInductorPass):
-    """
-    This pass performs cleanup after custom passes.
+    """This pass performs cleanup after custom passes.
     It topologically sorts the graph and removes unused nodes.
     This is needed because the pattern matcher does not guarantee producing
     a topologically sorted graph, and there may be unused nodes left around.
diff --git a/vllm/compilation/passes/utility/scatter_split_replace.py b/vllm/compilation/passes/utility/scatter_split_replace.py
index a17a7b336d2d..de7d87d87394 100644
--- a/vllm/compilation/passes/utility/scatter_split_replace.py
+++ b/vllm/compilation/passes/utility/scatter_split_replace.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Replace ``slice_scatter`` and ``split_with_sizes`` nodes with a single
+"""Replace ``slice_scatter`` and ``split_with_sizes`` nodes with a single
 assignment if there are no users for the inplace tensor written to by
 the slice_scatter call.
 
diff --git a/vllm/compilation/passes/utility/split_coalescing.py b/vllm/compilation/passes/utility/split_coalescing.py
index 6373eacfa847..f16623a1aa51 100644
--- a/vllm/compilation/passes/utility/split_coalescing.py
+++ b/vllm/compilation/passes/utility/split_coalescing.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Coalesce duplicate ``split_with_sizes`` nodes that operate on the same
+"""Coalesce duplicate ``split_with_sizes`` nodes that operate on the same
 input tensor with the same split sizes.
 
 On certain hardware/dtype combinations (e.g. B200 + FP8) the Inductor
@@ -10,9 +9,10 @@
 so that downstream pattern-matching passes (e.g. QK-Norm+RoPE fusion)
 see a single split node with all users attached.
 
-See also:
+See Also:
   - vLLM  #33295  (original issue)
   - PyTorch #174472 (upstream CSE gap)
+
 """
 
 import operator
diff --git a/vllm/compilation/passes/vllm_inductor_pass.py b/vllm/compilation/passes/vllm_inductor_pass.py
index 351d83e91855..3d6e7ef4f304 100644
--- a/vllm/compilation/passes/vllm_inductor_pass.py
+++ b/vllm/compilation/passes/vllm_inductor_pass.py
@@ -35,8 +35,7 @@ class InductorCompilationConfig:
 
 
 class VllmInductorPass(InductorPass):
-    """
-    An inductor pass with access to vLLM PassConfig.
+    """An inductor pass with access to vLLM PassConfig.
     It provides timing, logging, and dumping utilities.
     """
 
@@ -93,8 +92,7 @@ def get_match_table() -> dict[str, int]:
 
 
 class VllmPatternMatcherPass(VllmInductorPass):
-    """
-    A VllmInductorPass that uses the Inductor pattern matcher.
+    """A VllmInductorPass that uses the Inductor pattern matcher.
     Provides pattern registration with match counting, debug dumping, and logging.
     """
 
@@ -109,7 +107,7 @@ class VllmPatternMatcherPass(VllmInductorPass):
     )
 
     def _replace_op_overloads(self, string: str) -> str:
-        """Replace <OpOverload(..., ...)> with nicer formulations"""
+        """Replace <OpOverload(..., ...)> with nicer formulations."""
         return str(
             self._OP_OVERLOAD_PATTERN.sub(
                 lambda m: f"torch.ops.{m.group(1)}.{m.group(2)}",
@@ -123,8 +121,7 @@ def log_match_summary(cls) -> None:
             logger.debug("fusion pass matches: %s", dict(cls.match_table))
 
     def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None:
-        """
-        If debug dumping is enabled, dump the Inductor pattern-matcher patterns
+        """If debug dumping is enabled, dump the Inductor pattern-matcher patterns
         into the debug_dump_path folder next to the dumped fx graphs.
 
         This method does its best to print something that looks like Python code
@@ -195,8 +192,7 @@ def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None
 
 
 class VllmPatternReplacement(ABC, Generic[P, R]):
-    """
-    A pattern/replacement pair for FX graph fusion.
+    """A pattern/replacement pair for FX graph fusion.
 
     Implement the three abstract members below, then pass
     instances to VllmFusionPatternMatcherPass.register(). The pass will
@@ -216,8 +212,7 @@ def pattern(self) -> Callable[P, R]:
     @property
     @abstractmethod
     def replacement(self) -> Callable[P, R]:
-        """
-        Returns a closure defining the FX subgraph to
+        """Returns a closure defining the FX subgraph to
         substitute in place of each match.
         """
         ...
@@ -282,6 +277,27 @@ def fold_consecutive_reshapes(gm: fx.GraphModule) -> None:
         gm.graph.erase_node(inp)
 
 
+def remove_noop_reshapes(gm: fx.GraphModule) -> None:
+    """Drop reshape ops whose output shape equals their input shape.
+
+    Companion to :func:`fold_consecutive_reshapes`. ``make_fx`` records a
+    reshape to the shape the input already has; the compiled graph has already
+    dropped it, so the pattern only matches once we drop it too.
+    """
+    aten_reshape = torch.ops.aten.reshape.default
+    for node in list(gm.graph.nodes):
+        if not is_func(node, aten_reshape):
+            continue
+        inp = node.args[0]
+        if not isinstance(inp, fx.Node):
+            continue
+        val, inp_val = node.meta.get("val"), inp.meta.get("val")
+        if val is None or inp_val is None or val.shape != inp_val.shape:
+            continue
+        node.replace_all_uses_with(inp)
+        gm.graph.erase_node(node)
+
+
 def _remove_noop_permutes(gm: fx.GraphModule) -> None:
     for node in gm.graph.nodes:
         if not is_func(node, torch.ops.aten.permute.default):
@@ -294,8 +310,7 @@ def _remove_noop_permutes(gm: fx.GraphModule) -> None:
 
 
 class VllmFusionPatternMatcherPass(VllmPatternMatcherPass):
-    """
-    A VllmPatternMatcherPass for passes that use VllmPatternReplacement objects.
+    """A VllmPatternMatcherPass for passes that use VllmPatternReplacement objects.
     Subclasses register patterns via self.register() in their own __init__.
     """
 
diff --git a/vllm/compilation/piecewise_backend.py b/vllm/compilation/piecewise_backend.py
index b647e0d8581a..576cd941f6a8 100644
--- a/vllm/compilation/piecewise_backend.py
+++ b/vllm/compilation/piecewise_backend.py
@@ -96,8 +96,7 @@ def __init__(
         compiled_runnables: dict[str, Callable[..., Any]] | None = None,
         submod_name: str = "",
     ):
-        """
-        The backend for piecewise compilation.
+        """The backend for piecewise compilation.
         It mainly handles the compilation of static shapes and
         dispatching based on runtime shape.
 
diff --git a/vllm/compilation/wrapper.py b/vllm/compilation/wrapper.py
index 4821daf9dded..1df30a70a83f 100644
--- a/vllm/compilation/wrapper.py
+++ b/vllm/compilation/wrapper.py
@@ -45,8 +45,7 @@ def _compilation_context() -> Generator[None, None, None]:
 
 
 class TorchCompileWithNoGuardsWrapper:
-    """
-    A wrapper class for torch.compile, it ensures that all guards are dropped
+    """A wrapper class for torch.compile, it ensures that all guards are dropped
     when CompilationMode is not CompilationMode.STOCK_TORCH_COMPILE.
     When guards are dropped, the first time __call__ is invoked, a single
     compilation is triggered. Dynamo should never be traced again after that
@@ -272,8 +271,7 @@ def cleanup(self) -> None:
     @contextmanager
     def _dispatch_to_compiled_code(self) -> Generator[None, None, None]:
         # noqa: E501
-        """
-        Context manager to dispatch to internally compiled code for torch<2.8.
+        """Context manager to dispatch to internally compiled code for torch<2.8.
         Why does this work? Because Dynamo guarantees that the compiled
         bytecode has exactly the same arguments, cell variables, and free
         variables as the original code. Therefore we can directly switch
@@ -291,9 +289,7 @@ def _dispatch_to_compiled_code(self) -> Generator[None, None, None]:
 
 
 def reset_compile_wrapper(model: torch.nn.Module) -> None:
-    """
-    Clean up compiled model and captured CUDA graphs for elastic EP.
-    """
+    """Clean up compiled model and captured CUDA graphs for elastic EP."""
     if not isinstance(model, TorchCompileWithNoGuardsWrapper) and hasattr(
         model, "model"
     ):
diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py
index 1cfa372d60f7..940aea75521f 100644
--- a/vllm/config/__init__.py
+++ b/vllm/config/__init__.py
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from vllm.config.attention import AttentionConfig
+from vllm.config.attention import AttentionConfig, HiSparseConfig
+from vllm.config.aux_output import AuxOutputConfig
 from vllm.config.cache import CacheConfig
 from vllm.config.compilation import (
     CompilationConfig,
@@ -66,8 +67,11 @@
 # __all__ should only contain classes and functions.
 # Types and globals should be imported from their respective modules.
 __all__ = [
+    # From vllm.config.aux_output
+    "AuxOutputConfig",
     # From vllm.config.attention
     "AttentionConfig",
+    "HiSparseConfig",
     # From vllm.config.cache
     "CacheConfig",
     # From vllm.config.compilation
diff --git a/vllm/config/attention.py b/vllm/config/attention.py
index d5a150c92a31..a6602ed5491b 100644
--- a/vllm/config/attention.py
+++ b/vllm/config/attention.py
@@ -4,19 +4,35 @@
 from dataclasses import field
 from typing import Any, Literal
 
-from pydantic import field_validator
+from pydantic import Field, field_validator
 
 from vllm.config.utils import config
-from vllm.logger import init_logger
 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
-logger = init_logger(__name__)
-
 IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"]
 MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"]
 
 
+@config
+class HiSparseConfig:
+    """Configuration for HiSparse sparse-MLA KV offloading."""
+
+    device_buffer_size: int | None = Field(default=None, gt=0)
+    """Total per-request GPU hot-buffer rows, including the newest-token slot.
+
+    Defaults to one top-k per decode query plus one top-k of LRU slack. The
+    physical allocation is rounded up to the GPU cache block size selected
+    from the active backends.
+    """
+
+    eager_host_mirror: bool = True
+    """Mirror decode-written KV rows to the host pool during the forward so
+    page spills complete without moving data. When disabled, decode rows stay
+    resident-only and evicted pages are copied to host at spill time. Prefill
+    rows are always mirrored during the forward."""
+
+
 @config
 class AttentionConfig:
     """Configuration for attention mechanisms in vLLM."""
@@ -64,16 +80,27 @@ class AttentionConfig:
     use_prefill_query_quantization: bool = False
     """If set, quantize query for attention in prefill."""
 
-    use_fp4_indexer_cache: bool | None = None
-    """Deprecated alias for `indexer_kv_dtype`; use that instead. True maps to
-    `mxfp4`, False is a no-op (it selected the model default already)."""
-
     indexer_kv_dtype: IndexerKVDType = "auto"
     """Data type for the sparse-attention indexer K cache. "auto" picks the
     model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse
     indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel
     support in the backend."""
 
+    indexer_sparse_logits: bool = False
+    """DeepSeek V4.1 two-level indexer: score only the candidate blocks with
+    DeepGEMM's sparse MQA-logits kernels instead of computing dense logits over
+    the whole context and masking them. Requires `indexer_kv_dtype="mxfp4"`,
+    an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension
+    (the top-k runs on the kernels' bf16 logits). The sparse path costs
+    O(candidate blocks) per query regardless of context length, so it pays off
+    for long contexts (roughly 32K tokens and beyond) and is slower below."""
+
+    hisparse_config: HiSparseConfig | None = None
+    """HiSparse host-resident KV configuration. Setting this enables experimental
+    Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred
+    with defaults when HiSparseConnector is configured (directly or via
+    MultiConnector); set it explicitly only to tune its fields."""
+
     use_non_causal: bool = False
     """Whether to use non-causal (bidirectional) attention."""
 
@@ -116,20 +143,6 @@ def __post_init__(self) -> None:
             # layers still use the platform's normal automatic backend.
             self.backend = None
 
-        if self.use_fp4_indexer_cache is not None:
-            logger.warning(
-                "use_fp4_indexer_cache is deprecated and will be removed in "
-                "v0.29. Use indexer_kv_dtype instead (True -> 'mxfp4')."
-            )
-            if self.use_fp4_indexer_cache:
-                if self.indexer_kv_dtype not in ("auto", "mxfp4"):
-                    raise ValueError(
-                        "use_fp4_indexer_cache=True conflicts with "
-                        f"indexer_kv_dtype={self.indexer_kv_dtype!r}. Set only "
-                        "indexer_kv_dtype."
-                    )
-                self.indexer_kv_dtype = "mxfp4"
-
     def resolve_indexer_kv_dtype(self, default: IndexerKVDType) -> IndexerKVDType:
         """Resolve `indexer_kv_dtype`, substituting `default` for "auto"."""
         if self.indexer_kv_dtype == "auto":
@@ -137,8 +150,7 @@ def resolve_indexer_kv_dtype(self, default: IndexerKVDType) -> IndexerKVDType:
         return self.indexer_kv_dtype
 
     def compute_hash(self) -> str:
-        """
-        Provide a hash that uniquely identifies all the configs
+        """Provide a hash that uniquely identifies all the configs
         that affect the structure of the computation
         graph from input ids/embeddings to the final hidden states,
         excluding anything before input ids/embeddings and after
@@ -146,9 +158,7 @@ def compute_hash(self) -> str:
         """
         from vllm.config.utils import get_hash_factors, hash_factors
 
-        # Folded into indexer_kv_dtype by __post_init__.
-        ignored_factors: set[str] = {"use_fp4_indexer_cache"}
-        factors = get_hash_factors(self, ignored_factors)
+        factors = get_hash_factors(self, set())
         return hash_factors(factors)
 
     @field_validator("backend", mode="before")
diff --git a/vllm/config/aux_output.py b/vllm/config/aux_output.py
new file mode 100644
index 000000000000..f70f57e70f4e
--- /dev/null
+++ b/vllm/config/aux_output.py
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration for execution auxiliary outputs."""
+
+from pydantic import Field
+
+from vllm.config.utils import config
+
+
+@config
+class AuxOutputConfig:
+    """Configuration for auxiliary-output delivery."""
+
+    enable_return_routed_experts: bool = False
+    """Capture and return routed-experts auxiliary outputs."""
+
+    max_bytes: int | None = Field(default=None, gt=0)
+    """LRU capacity, or ``None`` to derive it from the KV cache capacity."""
+
+    @property
+    def enabled(self) -> bool:
+        """Whether any execution auxiliary output is enabled."""
+        return self.enable_return_routed_experts
+
+    def compute_hash(self) -> str:
+        """Hash AuxOutput settings that alter the model forward graph."""
+        from vllm.config.utils import hash_factors
+
+        return hash_factors(
+            {"enable_return_routed_experts": self.enable_return_routed_experts}
+        )
diff --git a/vllm/config/cache.py b/vllm/config/cache.py
index 4c48fc291ebb..3121f49532b8 100644
--- a/vllm/config/cache.py
+++ b/vllm/config/cache.py
@@ -8,7 +8,7 @@
 
 from pydantic import Field, field_validator, model_validator
 
-from vllm.config.utils import config, get_from_deprecated_env_if_set
+from vllm.config.utils import config
 from vllm.logger import init_logger
 from vllm.utils.torch_utils import (
     is_quantized_kv_cache,
@@ -58,15 +58,6 @@ def _layout_from_name(layout_name: str) -> KVCacheLayout:
 ]
 
 
-def _get_prefix_cache_retention_interval() -> int | None:
-    env_value = get_from_deprecated_env_if_set(
-        "VLLM_PREFIX_CACHE_RETENTION_INTERVAL",
-        "v0.29",
-        "prefix_cache_retention_interval",
-    )
-    return 0 if env_value is None else int(env_value)
-
-
 MambaDType = Literal["auto", "float32", "float16", "bfloat16"]
 MambaCacheMode = Literal["all", "align", "none"]
 PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"]
@@ -116,7 +107,19 @@ class CacheConfig:
     per-instance limit, and only applies to the current vLLM instance. It does
     not matter if you have another vLLM instance running on the same GPU. For
     example, if you have two vLLM instances running on the same GPU, you can
-    set the GPU memory utilization to 0.5 for each instance."""
+    set the GPU memory utilization to 0.5 for each instance. On non-GPU
+    installs, this value controls the corresponding device memory utilization.
+    """
+
+    @property
+    def device_memory_utilization(self) -> float:
+        """Device-neutral alias for ``gpu_memory_utilization``."""
+        return self.gpu_memory_utilization
+
+    @device_memory_utilization.setter
+    def device_memory_utilization(self, value: float) -> None:
+        self.gpu_memory_utilization = value
+
     cache_dtype: CacheDType = "auto"
     """Data type for kv cache storage. If "auto", will use model data type.
     CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. ROCm (AMD GPU) supports
@@ -154,9 +157,7 @@ class CacheConfig:
       security risk tolerance against the performance benefits before turning this on.
     - "xxhash_cbor" combines canonical CBOR serialization with xxHash for
       reproducible hashing. Requires the optional ``xxhash`` package."""
-    prefix_cache_retention_interval: int | None = Field(
-        default_factory=_get_prefix_cache_retention_interval, ge=0
-    )
+    prefix_cache_retention_interval: int | None = Field(default=0, ge=0)
     """Token interval between retained sliding-window and Mamba prefix-cache
     checkpoints. ``0`` retains only semantic checkpoints, including the latest
     replay boundary and shared-prefix junctions. Positive values additionally
@@ -195,7 +196,7 @@ class CacheConfig:
       when the token is at position i * block_size. This is the default when prefix
       caching is enabled.
     """
-    enable_mamba_fine_grained_prefix_cache: bool = False
+    enable_mamba_shared_prefix_checkpoint: bool = False
     """Also register a Mamba "align" checkpoint at the shared-prefix junction --
     where an EAGLE/MTP sibling was observed to resume -- instead of only at the
     prompt tail. Off by default; only takes effect with `mamba_cache_mode`
@@ -222,6 +223,8 @@ class CacheConfig:
     """The number of blocks to allocate for CPU memory."""
 
     # Set after KV cache initialization.
+    effective_attention_block_size: int | None = field(default=None, init=False)
+    """Full-attention block size in tokens, including DCP, or None if unavailable."""
     kv_cache_size_tokens: int | None = field(default=None, init=False)
     """Per-DP-engine KV cache capacity in tokens (group-aware). Uses
     group-aware capacity since num_gpu_blocks * block_size can be wrong
@@ -234,9 +237,13 @@ class CacheConfig:
     some layers can skip tokens corresponding to prefill. This flag enables
     attention metadata for eligible layers to be overridden with metadata
     necessary for implementing this optimization in some models (e.g. Gemma3n)
-    NOTE: KV cache sharing is not supported for MRv2 (v2 model runner).
     """
 
+    swa_bounded_replay: bool = True
+    """Keep the sliding-window KV of models that support it (DeepSeek-V4.1)
+    out of prefix caching and rebuild it after a prefix hit by recomputing the
+    hit's last window. Requires model runner V2."""
+
     kv_cache_memory_bytes: int | None = None
     """Size of KV Cache per GPU in bytes. By default, this is set to None
     and vllm can automatically infer the kv cache size based on
@@ -258,8 +265,7 @@ class CacheConfig:
     KV offloading is only activated when kv_offloading_size is set."""
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
@@ -280,7 +286,7 @@ def compute_hash(self) -> str:
             "prefix_cache_retention_interval",
             # Prefix-caching implementation detail (doesn't affect compiled graph).
             "prefix_match_unit",
-            "enable_mamba_fine_grained_prefix_cache",
+            "enable_mamba_shared_prefix_checkpoint",
             "mamba_page_size_padded",
             "skip_page_size_padded",
             "user_specified_block_size",
@@ -289,10 +295,12 @@ def compute_hash(self) -> str:
             # Post-init/derived counters
             "num_gpu_blocks",
             "num_cpu_blocks",
+            "effective_attention_block_size",
             "kv_cache_size_tokens",
             "kv_cache_max_concurrency",
-            # WIP feature toggle not impacting compiled graph shape
+            # Feature toggles not impacting compiled graph shape
             "kv_sharing_fast_prefill",
+            "swa_bounded_replay",
         }
 
         from vllm.config.utils import get_hash_factors, hash_factors
diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py
index 1219607d3abf..eefcf1f49e24 100644
--- a/vllm/config/compilation.py
+++ b/vllm/config/compilation.py
@@ -77,6 +77,13 @@ def has_mode(self, mode: "CUDAGraphMode") -> bool:
     def requires_piecewise_compilation(self) -> bool:
         return self.has_mode(CUDAGraphMode.PIECEWISE)
 
+    def without_piecewise(self) -> "CUDAGraphMode":
+        if self == CUDAGraphMode.PIECEWISE:
+            return CUDAGraphMode.NONE
+        if self == CUDAGraphMode.FULL_AND_PIECEWISE:
+            return CUDAGraphMode.FULL_DECODE_ONLY
+        return self
+
     def max_cudagraph_mode(self) -> "CUDAGraphMode":
         return CUDAGraphMode(max(self.value)) if self.separate_routine() else self
 
@@ -147,15 +154,15 @@ class PassConfig:
     fuse_rope_kvcache: bool = None  # type: ignore[assignment]
     """Fuse the QK rope + KV cache ops."""
     fuse_qk_norm_rope_kvcache: bool = Field(default=None)  # type: ignore[assignment]
-    """Fuse QK RMSNorm + RoPE + KV cache update into a single AITER HIP
+    """Fuse QK RMSNorm + RoPE/MRoPE + KV cache update into an AITER HIP
     kernel. Supersedes both enable_qk_norm_rope_fusion and fuse_rope_kvcache
-    for layers that support it. Auto-enabled at O1+ on ROCm for models
-    with QK-norm (e.g. Qwen3-MoE)."""
+    for layers that support it. Auto-enabled at O2+ on ROCm for models
+    with QK-norm (e.g. Qwen3-MoE and Qwen3-VL-class architectures)."""
 
     rope_kvcache_fusion_max_token_num: int = 256
     """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode.
     Larger batch sizes e.g. during prefill will use the unfused kernels.
-    Also applies to the fused QK-Norm+RoPE+KVCache pass.
+    Also applies to the fused QK-Norm+RoPE/MRoPE+KVCache pass.
     """
 
     fi_allreduce_fusion_max_size_mb: float | None = None
@@ -185,12 +192,10 @@ class PassConfig:
     # TODO(luka) better pass enabling system.
 
     def flashinfer_max_size(self, world_size: int) -> int | None:
-        """
-        Returns the max communication size in bytes for flashinfer
+        """Returns the max communication size in bytes for flashinfer
         allreduce fusion for the given world size. Returns None if world size
         is not supported by configs as it's not supported by flashinfer.
         """
-
         MiB = 1024 * 1024
         FI_SUPPORTED_WORLD_SIZES = [2, 4, 8, 16]
         if world_size not in FI_SUPPORTED_WORLD_SIZES:
@@ -216,12 +221,10 @@ def default_fi_allreduce_fusion_max_size_mb() -> dict[int, float]:
         return FI_ALLREDUCE_FUSION_MAX_SIZE_MB.get(capability.to_int(), {})
 
     def compute_hash(self) -> str:
-        """
-        Produces a hash unique to the pass configuration.
+        """Produces a hash unique to the pass configuration.
         Any new fields that affect compilation should be added to the hash.
         Any future fields that don't affect compilation should be excluded.
         """
-
         return hash_factors(get_hash_factors(self, set()))
 
     @field_validator(
@@ -310,8 +313,7 @@ def __post_init__(self) -> None:
             self.fuse_rope_kvcache_cat_mla = False
 
     def log_enabled_passes(self) -> None:
-        """
-        Log the enabled custom fusion passes.
+        """Log the enabled custom fusion passes.
         This is called at the end of VLLMConfig post_init,
         after all defaults are finalized.
         TODO also log the compile ranges for which this is enabled.
@@ -384,10 +386,7 @@ class DynamicShapesConfig:
     """
 
     def compute_hash(self) -> str:
-        """
-        Provide a hash for DynamicShapesConfig
-        """
-
+        """Provide a hash for DynamicShapesConfig."""
         from vllm.config.utils import get_hash_factors, hash_factors
 
         factors = get_hash_factors(self, set())
@@ -783,8 +782,7 @@ class CompilationConfig:
     ]
 
     def compute_hash(self) -> str:
-        """
-        Provide a hash that uniquely identifies all the configs
+        """Provide a hash that uniquely identifies all the configs
         that affect the structure of the computation
         graph from input ids/embeddings to the final hidden states,
         excluding anything before input ids/embeddings and after
@@ -849,8 +847,7 @@ def __repr__(self) -> str:
     @field_validator("mode", mode="before")
     @classmethod
     def validate_mode_before(cls, value: Any) -> Any:
-        """
-        Enable parsing the `mode` field from string mode names.
+        """Enable parsing the `mode` field from string mode names.
         Accepts both integers (0-3) and string names, like NONE, STOCK_TORCH_COMPILE,
         DYNAMO_TRACE_ONCE, VLLM_COMPILE.
         """
@@ -1072,15 +1069,17 @@ def init_backend(
         prefix: str = "",
         is_encoder: bool = False,
     ) -> str | Callable:
-        """
-        Initialize the backend for the compilation config from a vllm config.
+        """Initialize the backend for the compilation config from a vllm config.
+
         Arguments:
             vllm_config: The vllm config to initialize the backend from.
             prefix: Cache directory prefix for this compiled module.
             is_encoder: Whether this module is used in an encoder (as
                 opposed to a text backbone).
+
         Returns:
             The backend for the compilation config.
+
         """
         if self.mode is None:
             raise ValueError(
@@ -1115,7 +1114,6 @@ def post_init_cudagraph_sizes(self) -> None:
         configs are set. This includes:
         - initialize compile_sizes
         """
-
         computed_compile_sizes: list[int] = []
         if self.compile_sizes is not None:
             # de-duplicate the sizes provided by the config
@@ -1314,13 +1312,11 @@ def is_attention_compiled_piecewise(self) -> bool:
         return self.backend == "inductor" and self.mode != CompilationMode.NONE
 
     def custom_op_log_check(self):
-        """
-        This method logs the enabled/disabled custom ops and checks that the
+        """This method logs the enabled/disabled custom ops and checks that the
         passed custom_ops field only contains relevant ops.
         It is called at the end of set_current_vllm_config,
         after the custom ops have been instantiated.
         """
-
         if len(self.enabled_custom_ops) + len(self.disabled_custom_ops) == 0:
             logger.debug("No custom ops found in model.")
             return
@@ -1382,6 +1378,7 @@ def resolve_cudagraph_mode_and_sizes(
         kv_cache_config: "KVCacheConfig | None" = None,
         max_num_reqs: int | None = None,
         is_profiling: bool = False,
+        piecewise_capture_available: bool = True,
     ) -> CUDAGraphMode:
         from vllm.v1.attention.backend import AttentionCGSupport
 
@@ -1464,6 +1461,20 @@ def resolve_cudagraph_mode_and_sizes(
                 cudagraph_mode = CUDAGraphMode.NONE
             logger.warning(msg)
 
+        if (
+            not piecewise_capture_available
+            and cudagraph_mode.requires_piecewise_compilation()
+        ):
+            fallback_mode = cudagraph_mode.without_piecewise()
+            logger.warning_once(
+                "Cudagraph mode %s requires piecewise capture, but the loaded "
+                "model provides neither a compiled submodule nor breakable CUDA "
+                "graphs. Overriding to %s.",
+                cudagraph_mode,
+                fallback_mode,
+            )
+            cudagraph_mode = fallback_mode
+
         # double check that we can support full cudagraph if they are requested
         # even after automatic downgrades
         if (
diff --git a/vllm/config/device.py b/vllm/config/device.py
index 2d46dcdb117e..c0bdf01ee47d 100644
--- a/vllm/config/device.py
+++ b/vllm/config/device.py
@@ -28,8 +28,7 @@ class DeviceConfig:
     `__post_init__`."""
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
diff --git a/vllm/config/ec_manager_config.py b/vllm/config/ec_manager_config.py
index a59d2cdc62b8..db284e9cfc4f 100644
--- a/vllm/config/ec_manager_config.py
+++ b/vllm/config/ec_manager_config.py
@@ -31,8 +31,7 @@ def get_encoder_cache_manager_obj(self):
 
 
 class EncoderCacheManagerMetadata(ABC):  # noqa: B024
-    """
-    Abstract Metadata used to communicate between the
+    """Abstract Metadata used to communicate between the
     Scheduler EncoderCacheManager and Worker EncoderCacheManager.
     """
 
diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py
index 3948f3205934..5c7f99b4aa33 100644
--- a/vllm/config/ec_transfer.py
+++ b/vllm/config/ec_transfer.py
@@ -23,6 +23,8 @@ class ECTransferConfig:
     safetensors) and ``ECMooncakeConnector`` (Mooncake TransferEngine RDMA;
     requires ``mooncake-transfer-engine`` and matching producer/consumer
     ``ec_connector_extra_config``; see ``mooncake_ec_connector`` module docstring).
+    Set ``cross_encoder_cache`` in Mooncake extra config to reuse shared
+    Encoder outputs from Store before encoding, retaining P2P delivery.
     """
 
     engine_id: str | None = None
@@ -63,8 +65,7 @@ class ECTransferConfig:
     Only supported in V1."""
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
diff --git a/vllm/config/engram.py b/vllm/config/engram.py
index b46311315f37..0f5ec31b335f 100644
--- a/vllm/config/engram.py
+++ b/vllm/config/engram.py
@@ -3,64 +3,95 @@
 
 from typing import TYPE_CHECKING
 
-from pydantic import Field
+from pydantic import model_validator
+from typing_extensions import Self
 
-import vllm.envs as envs
 from vllm.config.utils import config, get_hash_factors, hash_factors
-from vllm.logger import init_logger
 
 if TYPE_CHECKING:
     from vllm.config.model import ModelConfig
     from vllm.config.parallel import ParallelConfig
 
-logger = init_logger(__name__)
+# Architecture -> the hf_text_config field naming its n-gram layers. A model is
+# only configurable here if it actually has such layers to store.
+_NGRAM_LAYER_FIELDS = {
+    "DeepseekV41ForCausalLM": "engram_layer_ids",
+    "Qwen4ExpForCausalLM": "ple_layer_ids",
+    "Qwen4ExpForConditionalGeneration": "ple_layer_ids",
+}
 
 
-def _default_cpu_offload() -> bool:
-    """Honor the legacy environment variable only when the field is omitted."""
-    if envs.is_set("VLLM_PLE_CPU_OFFLOAD"):
-        logger.warning_once(
-            "VLLM_PLE_CPU_OFFLOAD is a legacy setting and may be removed in a "
-            "future release. Use --engram-config.cpu_offload instead."
-        )
-    return envs.VLLM_PLE_CPU_OFFLOAD
+def model_has_engram_layers(model_config: "ModelConfig | None") -> bool:
+    """Whether the model carries n-gram embedding layers."""
+    if model_config is None:
+        return False
+    field = _NGRAM_LAYER_FIELDS.get(model_config.architecture)
+    if field is None:
+        return False
+    return bool(getattr(model_config.hf_text_config, field, None))
 
 
 @config
 class EngramConfig:
     """Configuration for Engram embedding storage and sharding."""
 
-    cpu_offload: bool = Field(default_factory=_default_cpu_offload)
-    """Store embedding weights in pinned CPU memory for UVA lookup.
-    Defaults to False, or VLLM_PLE_CPU_OFFLOAD when set for compatibility.
-    An explicit value takes precedence over the legacy environment variable."""
+    cpu_offload: bool = True
+    """Store embedding weights in pinned CPU memory for UVA lookup."""
 
     embedding_across_dp: bool = False
     """Shard embeddings across TP and all DP ranks when enabled.
     Otherwise, each DP rank has a separate TP-sharded embedding replica."""
 
+    dp_shared_memory: bool | None = None
+    """Share CPU-offloaded embedding weights between co-located
+    DP replicas. Each node stores one copy of every TP shard, reducing host
+    memory without per-step Engram DP collectives. Requires sufficient
+    /dev/shm capacity and a shared IPC namespace. Defaults to enabled whenever
+    the other settings allow it, falling back to per-replica tables when DP
+    replicas are not co-located on one node or /dev/shm cannot hold them."""
+
+    @model_validator(mode="after")
+    def _validate_shared_memory(self) -> Self:
+        if self.dp_shared_memory and not self.cpu_offload:
+            raise ValueError("dp_shared_memory requires cpu_offload=True")
+        return self
+
     def verify_model_config(self, model_config: "ModelConfig | None") -> None:
-        """Reject Engram configuration for models without supported embeddings."""
+        """Reject Engram configuration for models without n-gram embeddings."""
         from vllm.platforms import current_platform
 
-        supported_architectures = {
-            "Qwen4ExpForCausalLM",
-            "Qwen4ExpForConditionalGeneration",
-        }
+        field = (
+            _NGRAM_LAYER_FIELDS.get(model_config.architecture)
+            if model_config is not None
+            else None
+        )
         if (
             model_config is None
-            or model_config.architecture not in supported_architectures
-            or not current_platform.is_cuda()
-            or not getattr(model_config.hf_text_config, "ple_layer_ids", None)
+            or field is None
+            or not current_platform.is_cuda_alike()
+            or not getattr(model_config.hf_text_config, field, None)
         ):
             raise ValueError(
                 "EngramConfig requires a model with supported Engram "
-                "embeddings. Currently only the CUDA Qwen4Exp implementation "
-                "with non-empty ple_layer_ids is supported."
+                "embeddings, non-empty n-gram layer ids, and CUDA."
+            )
+
+    def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None:
+        """Share host tables by default wherever the configuration permits."""
+        if self.dp_shared_memory is None:
+            self.dp_shared_memory = (
+                self.cpu_offload
+                and parallel_config.data_parallel_size > 1
+                and not parallel_config.enable_elastic_ep
             )
 
     def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None:
         """Reject unsupported embedding parallel topologies."""
+        if self.dp_shared_memory:
+            if parallel_config.data_parallel_size <= 1:
+                raise ValueError("dp_shared_memory requires data_parallel_size > 1.")
+            if parallel_config.enable_elastic_ep:
+                raise ValueError("dp_shared_memory is not supported with elastic EP.")
         if (
             self.embedding_across_dp
             and parallel_config.data_parallel_size > 1
diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py
index 1312cbb24157..ea0b7d60d1fc 100644
--- a/vllm/config/kernel.py
+++ b/vllm/config/kernel.py
@@ -18,8 +18,7 @@
 
 @config
 class IrOpPriorityConfig:
-    """
-    Configuration for vLLM IR op priority for dispatching/lowering during the
+    """Configuration for vLLM IR op priority for dispatching/lowering during the
     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,
@@ -34,9 +33,11 @@ class IrOpPriorityConfig:
     fused_add_rms_norm: list[str] = Field(default_factory=list)
     """Priority list for vllm.ir.ops.fused_add_rms_norm"""
 
+    gelu_and_mul_sparse: list[str] = Field(default_factory=list)
+    """Priority list for vllm.ir.ops.gelu_and_mul_sparse"""
+
     def compute_hash(self) -> str:
-        """
-        Produces a hash unique to the pass configuration.
+        """Produces a hash unique to the pass configuration.
         Any new fields that affect compilation should be added to the hash.
         Any future fields that don't affect compilation should be excluded.
 
@@ -68,8 +69,7 @@ def _to_list_str(cls, value: str | list[str]):
         return value
 
     def _iter_op_priorities(self):
-        """
-        Yield (IrOp, priority_list) for each field, after importing platform
+        """Yield (IrOp, priority_list) for each field, after importing platform
         kernels and validating each entry.
         """
         from vllm.ir.op import IrOp
@@ -86,16 +86,13 @@ def _iter_op_priorities(self):
             yield IrOp.registry[field.name], op_priority
 
     def set_default(self) -> None:
-        """
-        Permanently set the IR op priority for all op members.
-        """
+        """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):
-        """
-        Context manager to set the IR op priority for all op members.
+        """Context manager to set the IR op priority for all op members.
         It also imports IR kernel implementations for the current platform
         to ensure all implementations are made available.
         """
@@ -108,8 +105,7 @@ def set_priority(self):
     def with_default(
         cls, default: list[str], /, **kwargs: list[str]
     ) -> "IrOpPriorityConfig":
-        """
-        A helper to create an IrOpPriorityConfig where fields not specified in kwargs
+        """A helper to create an IrOpPriorityConfig where fields not specified in kwargs
         use the given default list.
         """
         for field in fields(cls):  # type: ignore[arg-type]
@@ -141,6 +137,7 @@ def with_default(
     "flydsl",
     "hpc",
     "emulation",
+    "rdna3",
 ]
 
 # Backends that run the mega-MoE model path through the flashinfer moe_ep
@@ -159,12 +156,23 @@ def with_default(
 # moe_ep variants.
 MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS
 
+SparseIndexerTopkBackend = Literal[
+    "auto",
+    "deep_select",
+    "cooperative",
+    "persistent",
+    "per_row",
+    "flashinfer",
+    "torch",
+]
+
 # Architectures whose model code wires up the flashinfer moe_ep experts. MTP
 # and DSpark draft variants inherit the setting from these target models.
 FLASHINFER_MOE_EP_ARCHITECTURES = frozenset(
     {
         "DeepseekV4ForCausalLM",
         "DeepSeekV4MTPModel",
+        "DeepseekV41ForCausalLM",
     }
 )
 
@@ -261,11 +269,27 @@ class KernelConfig:
     - "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16
       (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250)
     - "flydsl": Use AMD FlyDSL kernels (ROCm only)
+    - "rdna3": Use the fused RDNA3 W4A16 HIP kernel (ROCm gfx1100 only)
     - "hpc": Use HPC kernels (FP8 and Hopper only)
     - "emulation": use BF16/FP16 GEMM, dequantizing weights and
                    running QDQ on activations.
     """
 
+    sparse_indexer_topk_backend: SparseIndexerTopkBackend = "auto"
+    """Backend for the DSA sparse indexer decode top-k kernel. Available options:
+
+    - "auto": The pre-existing chain (cooperative -> persistent -> per_row);
+      the other backends are opt-in
+    - "deep_select": Use DeepSelect kernels (SM100a/SM103a only)
+    - "cooperative": Use vLLM's cooperative_topk kernel
+    - "persistent": Use vLLM's persistent_topk kernel
+    - "per_row": Use vLLM's top_k_per_row_decode kernel
+    - "flashinfer": Use FlashInfer's top_k_ragged_transform kernel
+    - "torch": Use a plain torch.topk implementation (debug reference)
+
+    Explicit values raise RuntimeError when their constraints are not met.
+    """
+
     linear_backend: LinearBackend = "auto"
     """Backend for linear layer GEMM kernels. Available options:
 
@@ -295,6 +319,13 @@ class KernelConfig:
     - "xpu_woq": Use XPU kernels for weight-only quantization (e.g. W8A16)
     """
 
+    linear_backend_per_quant: dict[str, LinearBackend] | None = Field(
+        default=None, min_length=1
+    )
+    """Backend overrides keyed by linear quantization scheme. Overrides take
+    precedence over ``linear_backend``; for example,
+    ``{"nvfp4_w4a16": "humming"}``."""
+
     @field_validator("moe_backend", mode="before")
     @classmethod
     def _normalize_moe_backend(cls, value: Any) -> Any:
@@ -309,9 +340,15 @@ def _normalize_linear_backend(cls, value: Any) -> Any:
             return value.lower().replace("-", "_")
         return value
 
+    @field_validator("sparse_indexer_topk_backend", mode="before")
+    @classmethod
+    def _normalize_sparse_indexer_topk_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.
+        """Produces a hash unique to the pass configuration.
         Any new fields that affect compilation should be added to the hash.
         Any future fields that don't affect compilation should be excluded.
         """
@@ -321,6 +358,8 @@ def compute_hash(self) -> str:
             "enable_flashinfer_autotune",
             "ir_op_priority",  # handled separately below
         }
+        if self.linear_backend_per_quant is None:
+            ignored_factors.add("linear_backend_per_quant")
         factors = get_hash_factors(self, ignored_factors)
         factors["ir_op_priority"] = self.ir_op_priority.compute_hash()
         return hash_factors(factors)
diff --git a/vllm/config/kv_transfer.py b/vllm/config/kv_transfer.py
index d1344a5211cd..dfc9e02f3d0a 100644
--- a/vllm/config/kv_transfer.py
+++ b/vllm/config/kv_transfer.py
@@ -13,6 +13,40 @@
 KVRole = Literal[KVProducer, KVConsumer]
 
 
+def hisparse_host_pool_gib(
+    kv_transfer_config: "KVTransferConfig | None",
+) -> float | None:
+    if kv_transfer_config is None:
+        return None
+    if kv_transfer_config.kv_connector == "MultiConnector":
+        connectors = kv_transfer_config.kv_connector_extra_config.get("connectors", [])
+    else:
+        connectors = [
+            {
+                "kv_connector": kv_transfer_config.kv_connector,
+                "kv_connector_extra_config": (
+                    kv_transfer_config.kv_connector_extra_config
+                ),
+            }
+        ]
+    entries = [
+        connector.get("kv_connector_extra_config", {})
+        for connector in connectors
+        if connector.get("kv_connector") == "HiSparseConnector"
+    ]
+    if len(entries) > 1:
+        raise ValueError("Only one HiSparseConnector may be configured")
+    if not entries:
+        return None
+    host_pool_gib = entries[0].get("host_pool_gib")
+    if host_pool_gib is None:
+        raise ValueError("HiSparseConnector requires host_pool_gib")
+    host_pool_gib = float(host_pool_gib)
+    if host_pool_gib <= 0:
+        raise ValueError("HiSparseConnector host_pool_gib must be positive")
+    return host_pool_gib
+
+
 def kv_buffer_device_default_factory() -> str:
     from vllm.platforms import current_platform
 
@@ -72,8 +106,7 @@ class KVTransferConfig:
     'fail': immediately fail the request with an error finish reason (default)"""
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
diff --git a/vllm/config/load.py b/vllm/config/load.py
index 5e24a617ba85..ee13e2fccd7c 100644
--- a/vllm/config/load.py
+++ b/vllm/config/load.py
@@ -118,8 +118,7 @@ class LoadConfig:
     """
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
diff --git a/vllm/config/lora.py b/vllm/config/lora.py
index 4323da96d806..24e42a42fc90 100644
--- a/vllm/config/lora.py
+++ b/vllm/config/lora.py
@@ -86,8 +86,7 @@ class LoRAConfig:
     MoE models whose adapters use this layout; ignored otherwise."""
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
diff --git a/vllm/config/model.py b/vllm/config/model.py
index e47f3eee96f0..15d046de7604 100644
--- a/vllm/config/model.py
+++ b/vllm/config/model.py
@@ -55,7 +55,7 @@
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 if TYPE_CHECKING:
-    from transformers import PretrainedConfig
+    from transformers import PreTrainedConfig
 
     import vllm.model_executor.layers.quantization as me_quant
     import vllm.model_executor.models as me_models
@@ -64,7 +64,7 @@
     from vllm.model_executor.layers.quantization import QuantizationMethods
     from vllm.v1.sample.logits_processor import LogitsProcessor
 else:
-    PretrainedConfig = Any
+    PreTrainedConfig = Any
 
     me_quant = LazyLoader(
         "model_executor", globals(), "vllm.model_executor.layers.quantization"
@@ -89,7 +89,6 @@
 TokenizerMode = Literal[
     "auto",
     "hf",
-    "slow",
     "mistral",
     "deepseek_v32",
     "deepseek_v4",
@@ -105,7 +104,7 @@
     "processed_logits",
     "processed_logprobs",
 )
-HfOverrides = dict[str, Any] | Callable[[PretrainedConfig], PretrainedConfig]
+HfOverrides = dict[str, Any] | Callable[[PreTrainedConfig], PreTrainedConfig]
 ModelImpl = Literal["auto", "vllm", "transformers", "terratorch"]
 LayerBlockType = Literal["attention", "linear_attention", "mamba"]
 
@@ -148,7 +147,6 @@ class ModelConfig:
     - "auto" will use the tokenizer from `mistral_common` for Mistral models
       if available, otherwise it will use the "hf" tokenizer.
     - "hf" will use the fast tokenizer if available.
-    - "slow" will always use the slow tokenizer.
     - "mistral" will always use the tokenizer from `mistral_common`.
     - "deepseek_v32" will always use the tokenizer from `deepseek_v32`.
     - "deepseek_v4" will always use the tokenizer from `deepseek_v4`.
@@ -183,10 +181,14 @@ class ModelConfig:
     We must set the global seed because otherwise,
     different tensor parallel workers would sample different tokens,
     leading to inconsistent results."""
-    hf_config: PretrainedConfig = field(init=False)
+    hf_config: PreTrainedConfig = field(init=False)
     """The Hugging Face config of the model."""
-    hf_text_config: PretrainedConfig = field(init=False)
+    hf_text_config: PreTrainedConfig = field(init=False)
     """The Hugging Face config of the text model (same as hf_config for text models)."""
+    is_submodel_config: bool = field(default=False, init=False)
+    """Whether this is a submodule view derived by `VllmConfig.with_hf_config`
+    (e.g. a multimodal model's text stack). Its architecture list is empty, so
+    deployment-level validation must not run against it."""
     word_embeddings_untied_by_checkpoint: bool = field(default=False, init=False)
     """Whether `tie_word_embeddings` was overridden to `False` because the checkpoint
     contains an `lm_head` of its own. The two may still turn out to be identical, in
@@ -247,8 +249,6 @@ class ModelConfig:
 
     NOTE: This disables both `torch.compile` and CUDA graphs, and is
     equivalent to setting `-cc.mode=none -cc.cudagraph_mode=none`."""
-    enable_return_routed_experts: bool = False
-    """Whether to return routed experts."""
     return_sampling_mask: bool = False
     """Whether to return the post-processing token support for each sample."""
     max_logprobs: int = Field(default=20, ge=-1)
@@ -342,6 +342,9 @@ class ModelConfig:
     enable_sleep_mode: bool = False
     """Enable sleep mode for the engine (only cuda and
     hip platforms are supported)."""
+    sleep_preserve_parameter_names: list[str] = field(default_factory=list)
+    """Parameter-name globs to preserve across level-2 sleep.
+    The sender must omit these parameters; loaders must preserve their storage."""
     sleep_mode_backend: str = "cumem"
     """Mechanism used to free and restore GPU state for sleep mode. ``"cumem"``
     (default) uses the built-in ``CuMemAllocator`` and is behavior-compatible
@@ -418,8 +421,7 @@ class ModelConfig:
     mm_processor_device: InitVar[MMProcessorDevice | None] = None
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
@@ -479,7 +481,7 @@ def compute_hash(self) -> str:
 
     def _update_nested(
         self,
-        target: PretrainedConfig | dict[str, Any],
+        target: PreTrainedConfig | dict[str, Any],
         updates: dict[str, Any],
     ) -> None:
         """Recursively updates a config or dict with nested updates."""
@@ -507,15 +509,15 @@ def _update_nested(
 
     def _apply_dict_overrides(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         overrides: dict[str, Any],
     ) -> None:
         """Apply dict overrides, handling both nested configs and dict values."""
-        from transformers import PretrainedConfig
+        from transformers import PreTrainedConfig
 
         for key, value in overrides.items():
             attr = getattr(config, key, None)
-            if attr is not None and isinstance(attr, PretrainedConfig):
+            if attr is not None and isinstance(attr, PreTrainedConfig):
                 # It's a nested config - recursively update it
                 self._update_nested(attr, value)
             else:
@@ -585,22 +587,27 @@ def __post_init__(
         self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
 
         # If loading model/tokenizer from HF Hub, resolve the revision once
-        # to prevent resolving it multiple times downstream.
-        # If the weights come from a different repo, we cannot eagerly resolve revision
-        weights_from_model = not self.model_weights or self.model_weights == self.model
-        # If the config comes from a different repo, we cannot eagerly resolve revision
-        config_from_model = not self.hf_config_path or self.hf_config_path == self.model
-        can_resolve_model_revision = config_from_model and weights_from_model
-        if can_resolve_model_revision:
-            self.revision = resolve_revision(
-                self.model,
+        # to prevent resolving it multiple times downstream. A resolved revision
+        # only pins the repo it was resolved for, so each repo needs its own call.
+        self.revision = resolve_revision(
+            self.model,
+            self.revision,
+            self.hf_token,
+        )
+
+        # The config can live in another repo, which `self.revision` does not pin.
+        # It stays `None` if the config comes from `self.model`, so that call sites
+        # fall back to `self.revision` the same way they fall back to `self.model`.
+        self._hf_config_revision = None
+        if self.hf_config_path and self.hf_config_path != self.model:
+            self._hf_config_revision = resolve_revision(
+                self.hf_config_path,
                 self.revision,
                 self.hf_token,
             )
 
         if (
-            can_resolve_model_revision
-            and self.tokenizer == self.model
+            self.tokenizer == self.model
             and self.tokenizer_revision == requested_revision
         ):
             self.tokenizer_revision = self.revision
@@ -628,7 +635,7 @@ def __post_init__(
         hf_config = get_config(
             self.hf_config_path or self.model,
             self.trust_remote_code,
-            self.revision,
+            self._hf_config_revision or self.revision,
             self.code_revision,
             self.config_format,
             hf_overrides_kw=hf_overrides_kw,
@@ -868,6 +875,80 @@ def __post_init__(
         self._verify_quantization()
         self._verify_cuda_graph()
 
+    def _supports_multimodal_inputs(self) -> bool:
+        """Checks if the model supports multimodal inputs.
+        Returns True if the model is multimodal with any non-zero supported
+        modalities, otherwise returns False, effectively running in
+        text-only mode.
+        """
+        if not self.is_multimodal_model:
+            return False
+
+        from vllm.multimodal import MULTIMODAL_REGISTRY
+
+        mm_config = self.get_multimodal_config()
+        try:
+            info = MULTIMODAL_REGISTRY.get_processing_info(self)
+        except ValueError:
+            # Speculative drafters for multimodal targets (e.g. Qwen3_5MTP,
+            # Exaone4_5_MTP, MiMoV2OmniMTP) declare `SupportsMultiModal` so
+            # that they can consume the embeddings merged by the target model,
+            # but they never run a multi-modal processor of their own. Running
+            # in text-only mode is the expected outcome for them, not a
+            # misconfiguration worth warning about.
+            if self.runner_type != "draft":
+                logger.warning_once(
+                    "Model %s is treated as multimodal but has no registered "
+                    "multimodal processor; running in text-only mode.",
+                    self.model,
+                )
+            return False
+
+        # Check if all supported modalities have limit == 0
+        if all(
+            mm_config.get_limit_per_prompt(modality) == 0
+            for modality in info.supported_mm_limits
+        ):
+            # If enable_mm_embeds is True, we still need MM infrastructure
+            # to process pre-computed embeddings even though encoder won't run
+            if mm_config.enable_mm_embeds:
+                return True
+
+            logger.info_once(
+                "All limits of multimodal modalities supported by the model "
+                "are set to 0, running in text-only mode."
+            )
+            return False
+
+        return True
+
+    def _cached_supports_multimodal_inputs(self) -> bool:
+        cache = getattr(self, "_supports_multimodal_inputs_cache", None)
+        if cache is None:
+            self._supports_multimodal_inputs_cache = cache = {}
+
+        mm_config = self.multimodal_config
+        if mm_config is None:
+            mm_cache_key = None
+        else:
+            limits_per_prompt = {
+                modality: mm_config.get_limit_per_prompt(modality)
+                for modality in mm_config.limit_per_prompt
+            }
+            mm_cache_key = (
+                mm_config.language_model_only,
+                tuple(limits_per_prompt.items()),
+                mm_config.enable_mm_embeds,
+            )
+
+        cache_key = (self.is_multimodal_model, self.runner_type, mm_cache_key)
+        if cache_key in cache:
+            return cache[cache_key]
+
+        supports_mm = self._supports_multimodal_inputs()
+        cache[cache_key] = supports_mm
+        return supports_mm
+
     def _supports_multimodal_for_mm_prefix(self) -> bool:
         """Whether multimodal inputs can still appear for this deployment.
 
@@ -884,24 +965,18 @@ def _supports_multimodal_for_mm_prefix(self) -> bool:
         vision modality is still enabled (e.g. ``image=0`` but video allowed).
         The deep-copied cache preserves the top-level decision instead.
         """
-        cached = getattr(self, "_supports_multimodal_inputs_cached", None)
-        if cached is not None:
-            return cached
-
         if self.multimodal_config is None:
             # Early call before multimodal init — do not clear mm_prefix yet.
             return True
 
-        from vllm.multimodal import MULTIMODAL_REGISTRY
-
-        supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs(self)
-        self._supports_multimodal_inputs_cached = supports_mm
+        supports_mm = self._cached_supports_multimodal_inputs()
         if not supports_mm:
             logger.info_once(
                 "Disabled mm_prefix attention mode because multimodal inputs "
                 "are configuration-disabled. Attention backends without "
                 "mm_prefix support may now be selected."
             )
+
         return supports_mm
 
     def get_model_arch_config(self) -> ModelArchitectureConfig:
@@ -938,7 +1013,7 @@ def validate_quantization_before(cls, value: Any) -> Any:
 
     @model_validator(mode="after")
     def validate_model_config_after(self: "ModelConfig") -> "ModelConfig":
-        """Called after __post_init__"""
+        """Called after __post_init__."""
         if not isinstance(self.tokenizer, str):
             raise ValueError(
                 f"tokenizer must be a string, got "
@@ -1014,7 +1089,10 @@ def _get_transformers_backend_cls(self) -> str:
         # Check if the architecture we're wrapping has defaults
         runner = None
         task = None
-        if defaults := try_match_architecture_defaults(self.architectures[0]):
+        # architectures is empty for with_hf_config() submodel views.
+        if self.architectures and (
+            defaults := try_match_architecture_defaults(self.architectures[0])
+        ):
             _, (runner, task) = defaults
         # User specified value take precedence
         if self.runner != "auto":
@@ -1083,8 +1161,8 @@ def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> No
         Args:
             model: Model name or path
             tokenizer: Tokenizer name or path
-        """
 
+        """
         # Skip if model_weights is already set (model already pulled)
         if self.model_weights:
             return
@@ -1662,9 +1740,7 @@ def get_num_layers_by_block_type(
             raise AssertionError(f"Unsupported block type: {block_type}")
 
     def get_mamba_chunk_size(self) -> int:
-        """
-        Returns the mamba chunk size if it exists
-        """
+        """Returns the mamba chunk size if it exists."""
         # used by e.g. Bamba, FalconH1, Granite
         chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None)
         if chunk_size is None:
@@ -1679,11 +1755,11 @@ def get_mamba_chunk_size(self) -> int:
         return chunk_size
 
     def get_multimodal_config(self) -> MultiModalConfig:
-        """
-        Get the multimodal configuration of the model.
+        """Get the multimodal configuration of the model.
 
         Raises:
             ValueError: If the model is not multimodal.
+
         """
         if self.multimodal_config is None:
             raise ValueError("The model is not multimodal.")
@@ -1691,8 +1767,7 @@ def get_multimodal_config(self) -> MultiModalConfig:
         return self.multimodal_config
 
     def try_get_generation_config(self) -> dict[str, Any]:
-        """
-        This method attempts to retrieve the non-default values of the
+        """This method attempts to retrieve the non-default values of the
         generation config for this model.
 
         The generation config can contain information about special tokens, as
@@ -1701,12 +1776,13 @@ def try_get_generation_config(self) -> dict[str, Any]:
 
         Returns:
             A dictionary containing the non-default generation config.
+
         """
         if self.generation_config in {"auto", "vllm"}:
             config = try_get_generation_config(
                 self.hf_config_path or self.model,
                 trust_remote_code=self.trust_remote_code,
-                revision=self.revision,
+                revision=self._hf_config_revision or self.revision,
                 code_revision=self.code_revision,
                 config_format=self.config_format,
                 hf_token=self.hf_token,
@@ -1726,8 +1802,7 @@ def try_get_generation_config(self) -> dict[str, Any]:
         return config.to_diff_dict()
 
     def get_diff_sampling_param(self) -> dict[str, Any]:
-        """
-        This method returns a dictionary containing the non-default sampling
+        """This method returns a dictionary containing the non-default sampling
         parameters with `override_generation_config` applied.
 
         The default sampling parameters are:
@@ -1739,6 +1814,7 @@ def get_diff_sampling_param(self) -> dict[str, Any]:
 
         Returns:
             A dictionary containing the non-default sampling parameters.
+
         """
         src = self.generation_config
 
@@ -1749,6 +1825,8 @@ def get_diff_sampling_param(self) -> dict[str, Any]:
 
         available_params = [
             "repetition_penalty",
+            "presence_penalty",
+            "frequency_penalty",
             "temperature",
             "top_k",
             "top_p",
@@ -1864,14 +1942,17 @@ def is_multimodal_model(self) -> bool:
     def is_multimodal_raw_input_only_model(self) -> bool:
         return self._model_info.supports_multimodal_raw_input_only
 
+    @property
+    def supports_multimodal_inputs(self) -> bool:
+        return self._cached_supports_multimodal_inputs()
+
     @property
     def requires_raw_input_tokens(self) -> bool:
         return self._model_info.requires_raw_input_tokens
 
     @property
     def score_type(self) -> ScoreType:
-        """
-        Scoring API handles score/rerank for:
+        """Scoring API handles score/rerank for:
 
         - "classify" task (score_type: cross-encoder models)
         - "embed" task (score_type: bi-encoder models)
@@ -1964,8 +2045,7 @@ def use_sep_token(self) -> bool:
 
     @property
     def head_dtype(self) -> torch.dtype:
-        """
-        "head" refers to the last Linear layer(s) of an LLM,
+        """The "head" refers to the last Linear layer(s) of an LLM,
         such as the lm_head in a generation model,
         or the score or classifier in a classification model.
 
@@ -1976,7 +2056,6 @@ def head_dtype(self) -> torch.dtype:
           fp32, which is required for RL training-inference consistency
           (the trainer computes logits in fp32).
         """
-
         head_dtype = _get_head_dtype(
             config=self.hf_config, dtype=self.dtype, runner_type=self.runner_type
         )
@@ -2174,8 +2253,7 @@ def is_nvfp4_quantized(self) -> bool:
 
 
 def get_served_model_name(model: str, served_model_name: str | list[str] | None):
-    """
-    If the input is a non-empty list, the first model_name in
+    """If the input is a non-empty list, the first model_name in
     `served_model_name` is taken.
     If the input is a non-empty string, it is used directly.
     For cases where the input is either an empty string or an
@@ -2320,7 +2398,7 @@ def _resolve_auto_dtype(
 
 def _get_and_verify_dtype(
     model_id: str,
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
     dtype: str | torch.dtype,
     *,
     is_pooling_model: bool,
@@ -2367,7 +2445,7 @@ def _get_and_verify_dtype(
 
 
 def _get_head_dtype(
-    config: PretrainedConfig, dtype: torch.dtype, runner_type: str
+    config: PreTrainedConfig, dtype: torch.dtype, runner_type: str
 ) -> torch.dtype:
     head_dtype: str | torch.dtype | None = getattr(config, "head_dtype", None)
 
@@ -2391,7 +2469,7 @@ def _get_head_dtype(
 
 
 def _get_and_verify_max_len(
-    hf_config: PretrainedConfig,
+    hf_config: PreTrainedConfig,
     model_arch_config: ModelArchitectureConfig,
     tokenizer_config: dict | None,
     max_model_len: int | None,
@@ -2465,13 +2543,20 @@ def _get_and_verify_max_len(
             # loading HF config
             rope_type = rp["rope_type"]
 
-            if rope_type not in ("su", "longrope", "llama3"):
+            # YaRN variants leave max_position_embeddings already scaled, as
+            # Transformers' _compute_yarn_parameters assumes, so `factor` must
+            # not be applied to it again.
+            if rope_type not in (
+                "su",
+                "longrope",
+                "llama3",
+                "yarn",
+                "deepseek_yarn",
+                "deepseek_llama_scaling",
+            ):
                 # NOTE: rope_type == "default" does not define factor https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/modeling_rope_utils.py
                 # NOTE: This assumes all layer types have the same scaling factor.
                 scaling_factor = rp.get("factor", scaling_factor)
-
-                if rope_type == "yarn":
-                    derived_max_model_len = rp["original_max_position_embeddings"]
         if scaling_factor is None:
             # Fallback the factor to 1.0 if a user assigned `null`
             logger.warning_once(
diff --git a/vllm/config/model_arch.py b/vllm/config/model_arch.py
index 9bd449b7b1de..9aed35384a8c 100644
--- a/vllm/config/model_arch.py
+++ b/vllm/config/model_arch.py
@@ -14,9 +14,7 @@
 
 @dataclass(config=ConfigDict(arbitrary_types_allowed=True))
 class ModelArchitectureConfig:
-    """
-    Configuration for model architecture that required by vLLM runtime
-    """
+    """Configuration for model architecture that required by vLLM runtime."""
 
     architectures: list[str]
     """List of model architecture class names (e.g., ['LlamaForCausalLM']).
diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py
index ce2217d9fefd..4eca42430c21 100644
--- a/vllm/config/multimodal.py
+++ b/vllm/config/multimodal.py
@@ -3,21 +3,30 @@
 
 from collections.abc import Mapping
 from pathlib import Path
-from typing import Any, Literal, TypeAlias, TypedDict, cast, final
+from typing import Any, Literal, TypeAlias, TypeVar, final, overload
 
 import torch
-from pydantic import ConfigDict, Field, field_validator, model_validator
+from pydantic import (
+    ConfigDict,
+    Field,
+    GetCoreSchemaHandler,
+    field_validator,
+    model_validator,
+)
 from pydantic.dataclasses import dataclass
+from pydantic_core import core_schema
 
 import vllm.envs as envs
 from vllm.config.ec_transfer import ECTransferConfig
-from vllm.config.utils import config, get_from_deprecated_env_if_set
+from vllm.config.utils import config
 from vllm.logger import init_logger
 from vllm.utils.hashing import safe_hash
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 logger = init_logger(__name__)
 
+_T = TypeVar("_T")
+
 
 @dataclass
 class BaseDummyOptions:
@@ -51,17 +60,50 @@ class AudioDummyOptions(BaseDummyOptions):
 
 
 @final
-class MultiModalDummyOptionsBuiltins(TypedDict, total=False):
-    """Type annotations for modality types predefined by vLLM."""
+class MultiModalDummyOptions(dict[str, BaseDummyOptions]):
+    """Dummy data options for each modality.
+
+    Lookups of the modalities predefined by vLLM return their own options
+    class, while any other modality returns
+    [`BaseDummyOptions`][vllm.config.multimodal.BaseDummyOptions].
+    """
+
+    @overload  # type: ignore[override]
+    def get(self, key: Literal["image"], /) -> ImageDummyOptions | None: ...
+
+    @overload
+    def get(self, key: Literal["image"], default: _T, /) -> ImageDummyOptions | _T: ...
+
+    @overload
+    def get(self, key: Literal["video"], /) -> VideoDummyOptions | None: ...
+
+    @overload
+    def get(self, key: Literal["video"], default: _T, /) -> VideoDummyOptions | _T: ...
 
-    image: ImageDummyOptions
-    """Options for dummy images."""
+    @overload
+    def get(self, key: Literal["audio"], /) -> AudioDummyOptions | None: ...
 
-    video: VideoDummyOptions
-    """Options for dummy videos."""
+    @overload
+    def get(self, key: Literal["audio"], default: _T, /) -> AudioDummyOptions | _T: ...
 
-    audio: AudioDummyOptions
-    """Options for dummy audios."""
+    @overload
+    def get(self, key: str, /) -> BaseDummyOptions | None: ...
+
+    @overload
+    def get(self, key: str, default: _T, /) -> BaseDummyOptions | _T: ...
+
+    def get(self, key: str, default: object = None, /) -> object:
+        return super().get(key, default)
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls,
+        source_type: Any,
+        handler: GetCoreSchemaHandler,
+    ) -> core_schema.CoreSchema:
+        return core_schema.no_info_after_validator_function(
+            cls, handler.generate_schema(dict[str, BaseDummyOptions])
+        )
 
 
 MMEncoderTPMode = Literal["weights", "data"]
@@ -75,25 +117,6 @@ class MultiModalDummyOptionsBuiltins(TypedDict, total=False):
 `"xpu"` on XPU). Validated against that set by the CLI."""
 
 
-def _get_mm_hasher_algorithm() -> MMHasherAlgorithm:
-    env_value = get_from_deprecated_env_if_set(
-        "VLLM_MM_HASHER_ALGORITHM",
-        "v0.27",
-        "mm_hasher_algorithm",
-    )
-    env_value = "blake3" if env_value is None else env_value
-    return cast(MMHasherAlgorithm, env_value.lower())
-
-
-MMDummyOptions: TypeAlias = dict[str, BaseDummyOptions]
-"""
-A dictionary containing an entry for each modality type of dummy data.
-
-The built-in modalities are defined by
-[`MultiModalDummyOptionsBuiltins`][vllm.config.multimodal.MultiModalDummyOptionsBuiltins].
-"""
-
-
 @config
 class MultiModalConfig:
     """Controls the behavior of multimodal models."""
@@ -101,7 +124,9 @@ class MultiModalConfig:
     language_model_only: bool = False
     """If True, disables all multimodal inputs by setting all modality limits to 0.
     Equivalent to setting `--limit-mm-per-prompt` to 0 for every modality."""
-    limit_per_prompt: MMDummyOptions = Field(default_factory=dict)
+    limit_per_prompt: MultiModalDummyOptions = Field(
+        default_factory=MultiModalDummyOptions
+    )
     """The maximum number of input items and options allowed per
     prompt for each modality.
 
@@ -164,9 +189,7 @@ class MultiModalConfig:
     mm_processor_cache_type: MMCacheType = "lru"
     """Type of cache to use for the multi-modal preprocessor/mapper. If `shm`,
     use shared memory FIFO cache. If `lru`, use mirrored LRU cache."""
-    mm_hasher_algorithm: MMHasherAlgorithm = Field(
-        default_factory=_get_mm_hasher_algorithm
-    )
+    mm_hasher_algorithm: MMHasherAlgorithm = "blake3"
     """Hash algorithm to use for multi-modal input caching. Use `"sha256"` or
     `"sha512"` for FIPS-compliant deployments."""
     mm_shm_cache_max_object_size_mb: int = Field(default=128, ge=0)
@@ -269,8 +292,8 @@ class MultiModalConfig:
     def _validate_limit_per_prompt(
         cls,
         value: dict[str, int | dict[str, int]],
-    ) -> MMDummyOptions:
-        out: MMDummyOptions = {}
+    ) -> MultiModalDummyOptions:
+        out = MultiModalDummyOptions()
 
         for k, v in value.items():
             # Handle legacy format where only count is specified
@@ -369,6 +392,7 @@ def fold_mm_processor_device(
         Returns:
             The kwargs to build the config with, unchanged unless the flag adds
             a `device`.
+
         """
         if mm_processor_device in (None, "auto"):
             return mm_processor_kwargs
@@ -402,6 +426,7 @@ def get_mm_processor_device_type(self) -> str | None:
             ValueError: If `device` is not something `torch.device` accepts.
                 `validate_mm_processor_device` is what surfaces this during
                 startup, so the value is only parsed once.
+
         """
         device = (self.mm_processor_kwargs or {}).get("device")
         if device is None:
@@ -430,6 +455,7 @@ def validate_mm_processor_device(self, ec_config: ECTransferConfig | None) -> No
             ValueError: If the requested device is not a torch device, or if it
                 is the accelerator on an instance that also runs the language
                 model.
+
         """
         from vllm.platforms import current_platform
 
@@ -464,8 +490,7 @@ def validate_mm_processor_device(self, ec_config: ECTransferConfig | None) -> No
         )
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
@@ -488,8 +513,7 @@ def compute_hash(self) -> str:
         return hash_str
 
     def get_limit_per_prompt(self, modality: str) -> int:
-        """
-        Get the maximum number of input items allowed per prompt
+        """Get the maximum number of input items allowed per prompt
         for the given modality (backward compatible).
         """
         if self.language_model_only:
@@ -507,8 +531,7 @@ def merge_mm_processor_kwargs(
         self,
         inference_kwargs: Mapping[str, object],
     ) -> dict[str, object]:
-        """
-        Get the keyword arguments to pass to the multi-modal processor
+        """Get the keyword arguments to pass to the multi-modal processor
         according to the extra arguments passed during inference.
         """
         kwargs = self.mm_processor_kwargs or {}
diff --git a/vllm/config/observability.py b/vllm/config/observability.py
index 0394245f3cf2..a39962a075ac 100644
--- a/vllm/config/observability.py
+++ b/vllm/config/observability.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from functools import cached_property
-from typing import Any, Literal, cast
+from typing import Any, Literal
 
 from packaging.version import parse
 from pydantic import Field, field_validator, model_validator
@@ -112,8 +112,7 @@ def collect_model_execute_time(self) -> bool:
         )
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
@@ -151,17 +150,6 @@ def _validate_otlp_traces_endpoint(cls, value: str | None) -> str | None:
                 )
         return value
 
-    @field_validator("collect_detailed_traces")
-    @classmethod
-    def _validate_collect_detailed_traces(
-        cls, value: list[DetailedTraceModules] | None
-    ) -> list[DetailedTraceModules] | None:
-        """Handle the legacy case where users might provide a comma-separated
-        string instead of a list of strings."""
-        if value is not None and len(value) == 1 and "," in value[0]:
-            value = cast(list[DetailedTraceModules], value[0].split(","))
-        return value
-
     @model_validator(mode="after")
     def _validate_tracing_config(self):
         if self.collect_detailed_traces and not self.otlp_traces_endpoint:
diff --git a/vllm/config/offload.py b/vllm/config/offload.py
index ad65e8acf35a..98fab0409f05 100644
--- a/vllm/config/offload.py
+++ b/vllm/config/offload.py
@@ -137,8 +137,7 @@ def validate_offload_config(self) -> "OffloadConfig":
         return self
 
     def compute_hash(self) -> str:
-        """
-        Provide a hash that uniquely identifies all the offload configs.
+        """Provide a hash that uniquely identifies all the offload configs.
 
         All fields are included because PrefetchOffloader patches module
         forwards and inserts custom ops (wait_prefetch, start_prefetch)
diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py
index bb1d2095de7b..3929d3be1266 100644
--- a/vllm/config/parallel.py
+++ b/vllm/config/parallel.py
@@ -38,7 +38,9 @@
 DataParallelBackend = Literal["ray", "mp"]
 EPLBPolicyOption = Literal["default"]
 DCPCommBackend = Literal["ag_rs", "a2a"]
-EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "nixl", "pynccl"]
+EPLBCommunicatorBackend = Literal[
+    "torch_nccl", "torch_gloo", "torch_xccl", "nixl", "pynccl"
+]
 All2AllBackend = Literal[
     "naive",
     "pplx",
@@ -47,6 +49,7 @@
     "deepep_v2",
     "mori_high_throughput",
     "mori_low_latency",
+    "moonep",
     "nixl_ep",
     "allgather_reducescatter",
     "flashinfer_all2allv",  # temporary alias for flashinfer_nvlink_two_sided
@@ -94,9 +97,11 @@ class EPLBConfig:
     Backend for EPLB expert weight communication:
     - "torch_nccl": Use torch.distributed on the device process group
     - "torch_gloo": Use torch.distributed gloo with CPU staging
+    - "torch_xccl": Use torch.distributed XCCL device P2P on XPU
     - "nixl": Use NIXL with staged send/recv buffers
     - "pynccl": Use PyNccl send/recv
-    - None: Auto-select backend (prefers "nixl", falls back to "torch_gloo")
+    - None: Auto-select backend ("torch_xccl" on XPU, prefers "nixl" 
+      on CUDA, falls back to "torch_gloo")
     """
 
     @model_validator(mode="after")
@@ -202,6 +207,7 @@ class ParallelConfig:
     - "deepep_low_latency": Use deepep low-latency kernels
     - "mori_high_throughput": MoRI EP with InterNodeV1 for multi-node
     - "mori_low_latency": MoRI EP with InterNodeV1LL for multi-node
+    - "moonep": MoonEP balanced EP with dynamic redundant experts (NVLink)
     - "nixl_ep": Use nixl-ep kernels
     - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels
     - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl"""
@@ -216,6 +222,8 @@ class ParallelConfig:
 
     enable_elastic_ep: bool = False
     """Enable elastic expert parallelism with stateless NCCL groups for DP/EP."""
+    elastic_ep_max_dp_size: int = Field(default=None, ge=1)  # type: ignore[assignment]
+    """Maximum data parallel size supported by elastic expert parallelism."""
 
     enable_dbo: bool = False
     """Enable dual batch overlap for the model executor."""
@@ -393,8 +401,15 @@ class is dynamically inherited by the worker class. This is used to inject
         in (rank i+1, block j) only after (rank i, block j) is fully occupied.
     Block_size should be greater than or equal to cp_kv_cache_interleave_size.
     Block_size should be divisible by cp_kv_cache_interleave_size.
+
+    When --cp-kv-cache-interleave-size is omitted (None), the interleave size
+    is resolved automatically based on NIXL transfer requirements.
+    Explicit settings take priority.
     """
 
+    _allow_auto_resolve_cp_interleave_size: bool = True
+    """Whether NIXL may select the interleave size automatically."""
+
     data_parallel_index: int = Field(init=False)
     """Equal to the data parallel rank but not used for torch process groups
     and not overridden for dense models."""
@@ -428,7 +443,9 @@ class is dynamically inherited by the worker class. This is used to inject
     )
     """The configurations for fault tolerance."""
 
-    @field_validator("disable_nccl_for_dp_synchronization", mode="wrap")
+    @field_validator(
+        "disable_nccl_for_dp_synchronization", "elastic_ep_max_dp_size", mode="wrap"
+    )
     @classmethod
     def _skip_none_validation(cls, value: Any, handler: Callable) -> Any:
         """Skip validation if the value is `None` when initialisation is delayed."""
@@ -514,10 +531,10 @@ def _validate_parallel_config(self) -> Self:
             )
 
         if self.enable_eplb:
-            if not current_platform.is_cuda_alike():
+            if not current_platform.is_cuda_alike() and not current_platform.is_xpu():
                 raise ValueError(
                     "Expert parallelism load balancing is only supported on "
-                    "CUDA devices or ROCm devices now."
+                    "CUDA devices or ROCm devices or XPU devices now."
                 )
             if not self.enable_expert_parallel:
                 raise ValueError("enable_expert_parallel must be True to use EPLB.")
@@ -591,15 +608,13 @@ def num_ubatches(self) -> int:
 
     @property
     def local_engines_only(self) -> bool:
-        """
-        Client manages local+remote EngineCores in pure internal LB case.
+        """Client manages local+remote EngineCores in pure internal LB case.
         Client manages local EngineCores in hybrid and external LB case.
         """
         return self.data_parallel_external_lb or self.data_parallel_hybrid_lb
 
     def get_next_dp_init_port(self) -> int:
-        """
-        We might need to initialize process groups in multiple
+        """We might need to initialize process groups in multiple
         processes that is related to data parallelism,
         e.g. both in the worker and in the engine, which
         can live in different processes. To avoid port conflicts, we
@@ -708,6 +723,7 @@ def use_sequence_parallel_moe(self) -> bool:
                 "allgather_reducescatter",
                 "deepep_high_throughput",
                 "deepep_low_latency",
+                "deepep_v2",
                 "flashinfer_nvlink_one_sided",
                 "mori_high_throughput",
                 "mori_low_latency",
@@ -744,12 +760,30 @@ def node_rank_within_dp(self) -> int:
 
     @property
     def nnodes_within_dp(self) -> int:
+        """Number of nodes one DP replica spans.
+
+        External LB pins ``data_parallel_size_local`` to 1, so the ratio
+        rounds down to 0 once DP replicas outnumber nodes. A replica that
+        does not span nodes still occupies exactly one.
+        """
         if self.nnodes == 1:
             return 1
         data_parallel_node_size = (
             self.data_parallel_size // self.data_parallel_size_local
         )
-        return self.nnodes // data_parallel_node_size
+        nnodes_within_dp = self.nnodes // data_parallel_node_size
+        if self.data_parallel_external_lb:
+            return max(nnodes_within_dp, 1)
+        if self.nnodes % data_parallel_node_size != 0:
+            raise ValueError(
+                "Invalid data parallel configuration: "
+                f"nnodes ({self.nnodes}) must be divisible by the number of "
+                "data parallel node groups "
+                f"({data_parallel_node_size} = data_parallel_size "
+                f"{self.data_parallel_size} / data_parallel_size_local "
+                f"{self.data_parallel_size_local})"
+            )
+        return nnodes_within_dp
 
     @property
     def local_world_size(self) -> int:
@@ -783,6 +817,7 @@ def sync_dp_state(
 
         Returns:
             (has_unfinished_global, pause_consensus)
+
         """
         tensor = torch.tensor(
             [int(has_unfinished), int(pending_pause)], dtype=torch.int32, device="cpu"
@@ -804,8 +839,7 @@ def sync_kv_cache_memory_size(dp_group: ProcessGroup, kv_cache_memory: int) -> i
         return tensor.item()
 
     def compute_hash(self):
-        """
-        Provide a hash that uniquely identifies all the configs
+        """Provide a hash that uniquely identifies all the configs
         that affect the structure of the computation
         graph from input ids/embeddings to the final hidden states,
         excluding anything before input ids/embeddings and after
@@ -934,6 +968,18 @@ def __post_init__(self) -> None:
                     " for dense models."
                 )
 
+        max_dp_size = self.elastic_ep_max_dp_size
+        self.elastic_ep_max_dp_size = (
+            max_dp_size
+            if self.enable_elastic_ep and max_dp_size is not None
+            else self.data_parallel_size
+        )
+        if self.elastic_ep_max_dp_size < self.data_parallel_size:
+            raise ValueError(
+                "--elastic-ep-max-dp-size must be greater than or equal to "
+                f"the initial data_parallel_size ({self.data_parallel_size})."
+            )
+
         self.data_parallel_index = self.data_parallel_rank
 
         if self.distributed_executor_backend == "external_launcher":
@@ -1012,7 +1058,10 @@ def __post_init__(self) -> None:
             # See https://github.com/pytorch/pytorch/issues/174288
             from vllm.distributed.nixl_utils import is_nixl_available
 
-            if is_nixl_available():
+            if current_platform.is_xpu():
+                # On XPU, use the device-native XCCL P2P backend.
+                self.eplb_config.communicator = "torch_xccl"
+            elif is_nixl_available():
                 self.eplb_config.communicator = "nixl"
             elif self.enable_elastic_ep:
                 self.eplb_config.communicator = "pynccl"
diff --git a/vllm/config/pooler.py b/vllm/config/pooler.py
index 9ca0ba457ebf..a5fbcbb622ed 100644
--- a/vllm/config/pooler.py
+++ b/vllm/config/pooler.py
@@ -176,8 +176,7 @@ def get_tok_pooling_type(self) -> TokenPoolingType:
         return self.tok_pooling_type
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py
index b3b232490f01..3ade4bceed87 100644
--- a/vllm/config/profiler.py
+++ b/vllm/config/profiler.py
@@ -14,6 +14,7 @@
 logger = init_logger(__name__)
 
 ProfilerKind = Literal["torch", "cuda", "proton"]
+TorchProfilerActivity = Literal["CPU", "CUDA", "PrivateUse1", "XPU"]
 ProtonBackend = Literal["cupti"]
 ProtonContext = Literal["shadow", "python"]
 ProtonData = Literal["tree", "trace"]
@@ -51,6 +52,12 @@ class ProfilerConfig:
     worker's traces (CPU & GPU) will be saved under this directory. Note that
     it must be an absolute path."""
 
+    torch_profiler_activities: list[TorchProfilerActivity] | None = Field(
+        default=None, min_length=1
+    )
+    """Activities recorded by workers using the torch profiler. When unset,
+    each worker uses its platform default: CPU; CPU and CUDA; or CPU and XPU."""
+
     proton_profiler_dir: str = ""
     """Directory to save Triton Proton profiles. Each worker writes a
     separate rank-qualified file."""
@@ -76,6 +83,10 @@ class ProfilerConfig:
     """Optional format passed to Proton when finalizing a profile. ``None``
     uses the default format for ``proton_data``."""
 
+    proton_graph_attribution: bool = False
+    """Observe CUDA graph capture so replayed kernels can be attributed.
+    Requires Triton >= 3.7 and ``proton_data='tree'``."""
+
     torch_profiler_with_stack: bool = True
     """If `True`, enables stack tracing in the torch profiler. Enabled by default
     as it is useful for debugging. Can be disabled via 
@@ -147,8 +158,7 @@ class ProfilerConfig:
     """
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
@@ -167,7 +177,24 @@ def compute_hash(self) -> str:
     @model_validator(mode="after")
     def _validate_profiler_config(self) -> Self:
         has_delay_or_limit = self.delay_iterations > 0 or self.max_iterations > 0
-        if self.profiler == "torch" and has_delay_or_limit and not self.ignore_frontend:
+        activities = self.torch_profiler_activities
+        if activities is not None:
+            if self.profiler != "torch":
+                raise ValueError(
+                    "torch_profiler_activities is only applicable when profiler "
+                    "is set to 'torch'"
+                )
+            if len(activities) != len(set(activities)):
+                raise ValueError(
+                    "torch_profiler_activities must not contain duplicates"
+                )
+        records_cpu_activity = activities is None or "CPU" in activities
+        if (
+            self.profiler == "torch"
+            and has_delay_or_limit
+            and not self.ignore_frontend
+            and records_cpu_activity
+        ):
             logger.warning_once(
                 "Using 'torch' profiler with delay_iterations or max_iterations "
                 "while ignore_frontend is False may result in high overhead."
@@ -199,6 +226,7 @@ def _validate_profiler_config(self) -> Self:
                 ("proton_mode", self.proton_mode, None),
                 ("proton_hook", self.proton_hook, None),
                 ("proton_output_format", self.proton_output_format, None),
+                ("proton_graph_attribution", self.proton_graph_attribution, False),
             )
             if value != default
         ]
@@ -220,6 +248,17 @@ def _validate_profiler_config(self) -> Self:
 
         if self.profiler == "proton":
             output_format = self.proton_output_format
+            if self.proton_graph_attribution and self.proton_data != "tree":
+                raise ValueError("proton_graph_attribution requires proton_data='tree'")
+            if (
+                self.proton_graph_attribution
+                and self.proton_mode
+                and self.proton_mode.split(":", 1)[0].lower() == "periodic_flushing"
+            ):
+                raise ValueError(
+                    "proton_graph_attribution is incompatible with periodic_flushing: "
+                    "both manage Proton data phases."
+                )
             if output_format == "chrome_trace" and self.proton_data != "trace":
                 raise ValueError("chrome_trace output requires proton_data='trace'")
             if (
diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py
index 8aa0ffa5dcff..1fb94d3efd07 100644
--- a/vllm/config/scheduler.py
+++ b/vllm/config/scheduler.py
@@ -67,9 +67,28 @@ class SchedulerConfig:
     In real usage, this should be set in `EngineArgs.create_engine_config`.
     """
 
+    max_num_active_seqs: int | None = Field(default=None, ge=1)
+    """Maximum number of requests the scheduler admits into RUNNING.
+
+    ``max_num_seqs`` sizes the model runner (per-request buffers and CUDA
+    graph capture) and is also the default admission limit. Setting this
+    lowers only the number of requests that may occupy RUNNING, so decode
+    batches stay smaller without shrinking runner or graph capacity. Must
+    be ``<= max_num_seqs``. ``None`` (default) keeps current behavior.
+    """
+
     long_prefill_token_threshold: int = Field(default=0, ge=0)
     """For chunked prefill, a request is considered long if the prompt is
-    longer than this number of tokens. 0 disables the cap (default)."""
+    longer than this number of tokens. 0 disables the cap (default).
+
+    The cap is not applied when the request is the only one in the batch,
+    since there is no other request for it to starve."""
+
+    long_prefill_token_threshold_adaptive: bool = Field(default=False)
+    """Floor the effective long prefill token threshold at a fair share of
+    the token budget: max_num_batched_tokens divided by the number of
+    queued and running requests. Only applies when
+    long_prefill_token_threshold is nonzero."""
 
     max_num_queued_reqs: int | None = Field(default=None, ge=0)
     """Maximum number of requests that can be in-flight (waiting or running)
@@ -200,9 +219,7 @@ class SchedulerConfig:
 
     @staticmethod
     def default_factory(**kwargs):
-        """
-        Factory method to create `SchedulerConfig` with default values for `InitVar`s.
-        """
+        """Create a `SchedulerConfig` with default values for its `InitVar`s."""
         if "max_model_len" not in kwargs:
             kwargs["max_model_len"] = 8192
         if "is_encoder_decoder" not in kwargs:
@@ -221,20 +238,23 @@ def get_scheduler_cls(self) -> type["SchedulerInterface"]:
 
         # The first half of this warning can be removed once the Scheduler interface is
         # finalized and we can maintain support for scheduler classes that implement it
-        logger.warning_once(
-            "Using custom scheduler class %s. This scheduler interface is not public "
-            "and compatibility may not be maintained. If you have subclassed Scheduler "
-            "instead of AsyncScheduler, you will see degraded performance due to async "
-            "scheduling being disabled.",
-            self.scheduler_cls,  # type: ignore[arg-type]
-        )
+        if not (
+            isinstance(self.scheduler_cls, str)
+            and self.scheduler_cls.startswith("vllm.")
+        ):
+            logger.warning_once(
+                "Using custom scheduler class %s. This scheduler interface is not "
+                "public and compatibility may not be maintained. If you have "
+                "subclassed Scheduler instead of AsyncScheduler, you will see "
+                "degraded performance due to async scheduling being disabled.",
+                self.scheduler_cls,  # type: ignore[arg-type]
+            )
         if not isinstance(self.scheduler_cls, str):
             return cast(type["SchedulerInterface"], self.scheduler_cls)
         return resolve_obj_by_qualname(self.scheduler_cls)
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
@@ -313,6 +333,15 @@ def verify_max_model_len(self, max_model_len: int) -> Self:
                 f"({self.max_num_seqs})."
             )
 
+        if (
+            self.max_num_active_seqs is not None
+            and self.max_num_active_seqs > self.max_num_seqs
+        ):
+            raise ValueError(
+                f"max_num_active_seqs ({self.max_num_active_seqs}) cannot be "
+                f"greater than max_num_seqs ({self.max_num_seqs})."
+            )
+
         if self.max_num_batched_tokens > self.max_num_seqs * max_model_len:
             logger.warning(
                 "max_num_batched_tokens (%d) exceeds max_num_seqs "
diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py
index d4085d7d52f8..7fb2a09dd3ab 100644
--- a/vllm/config/speculative.py
+++ b/vllm/config/speculative.py
@@ -14,7 +14,7 @@
 from vllm.config.kernel import MoEBackend
 from vllm.config.model import HfOverrides, ModelConfig
 from vllm.config.parallel import ParallelConfig
-from vllm.config.utils import config
+from vllm.config.utils import config, replace
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_text_config
 from vllm.utils.hashing import safe_hash
@@ -22,11 +22,13 @@
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 if TYPE_CHECKING:
-    from transformers import PretrainedConfig
+    from transformers import PreTrainedConfig
 
     import vllm.model_executor.layers.quantization as me_quant
+    from vllm.config.vllm import VllmConfig
 else:
-    PretrainedConfig = Any
+    PreTrainedConfig = Any
+    VllmConfig = Any
 
     me_quant = LazyLoader(
         "model_executor", globals(), "vllm.model_executor.layers.quantization"
@@ -371,6 +373,17 @@ def _validate_qwen3_omni_dspark(
         )
 
 
+# (SpeculativeConfig field, VllmConfig sub-config, overridden field)
+_DRAFT_VLLM_CONFIG_OVERRIDES = (
+    # Otherwise the draft inherits the target's --moe-backend, which fails
+    # when the draft is unquantized and that backend is not.
+    ("moe_backend", "kernel_config", "moe_backend"),
+    # Only when set, so the draft keeps a KV cache layout the target shares.
+    ("attention_backend", "attention_config", "backend"),
+    ("kv_cache_dtype", "cache_config", "cache_dtype"),
+)
+
+
 @config
 class SpeculativeConfig:
     """Configuration for speculative decoding."""
@@ -598,8 +611,7 @@ def _resolve_synthetic_acceptance_rates(
     top-k base-logit candidates. Requires draft tensor parallel size 1."""
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
@@ -645,7 +657,7 @@ def compute_hash(self) -> str:
         return hash_str
 
     @staticmethod
-    def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
+    def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig:
         initial_architecture = hf_config.architectures[0]
         use_v32_mtp = hf_config.model_type in ("deepseek_v32", "glm_moe_dsa")
         if hf_config.model_type == "dots3_note":
@@ -679,12 +691,18 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
                     ],
                 }
             )
-        if hf_config.model_type == "deepseek_v4":
+        if hf_config.model_type in ("deepseek_v4", "deepseek_v41"):
+            # V4.1 has no classic-MTP draft: its checkpoints ship DSpark stages
+            # under ``mtp.*``, so only V4 gets an MTP architecture here. The
+            # DSpark path rewrites ``architectures`` itself;
+            # ``method="mtp"`` on V4.1 is rejected below.
+            is_v41 = hf_config.model_type == "deepseek_v41"
             hf_config.model_type = "deepseek_mtp"
             n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
-            hf_config.update(
-                {"n_predict": n_predict, "architectures": ["DeepSeekV4MTPModel"]}
-            )
+            overrides = {"n_predict": n_predict}
+            if not is_v41:
+                overrides["architectures"] = ["DeepSeekV4MTPModel"]
+            hf_config.update(overrides)
         if hf_config.model_type in ("pangu_ultra_moe"):
             hf_config.model_type = "pangu_ultra_moe_mtp"
         if hf_config.model_type == "pangu_ultra_moe_mtp":
@@ -1048,16 +1066,16 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
 
     @staticmethod
     def _apply_composed_hf_override(
-        target_hf_overrides: Callable[[PretrainedConfig], PretrainedConfig],
-        hf_config: PretrainedConfig,
-    ) -> PretrainedConfig:
+        target_hf_overrides: Callable[[PreTrainedConfig], PreTrainedConfig],
+        hf_config: PreTrainedConfig,
+    ) -> PreTrainedConfig:
         hf_config = SpeculativeConfig.hf_config_override(hf_config)
         return target_hf_overrides(hf_config)
 
     @staticmethod
     def compose_draft_hf_overrides(
         target_hf_overrides: HfOverrides | None,
-    ) -> Callable[[PretrainedConfig], PretrainedConfig]:
+    ) -> Callable[[PreTrainedConfig], PreTrainedConfig]:
         """Build the ``hf_overrides`` for the draft ``ModelConfig``.
 
         Callable overrides on the target are config-to-config transforms
@@ -1335,6 +1353,18 @@ def __post_init__(self):
                     MTPModelTypes
                 ):
                     self.method = "mtp"
+                    if (
+                        self.target_model_config is not None
+                        and self.target_model_config.hf_config.model_type
+                        == "deepseek_v41"
+                    ):
+                        raise ValueError(
+                            "DeepSeek V4.1 has no classic-MTP draft: its "
+                            "checkpoints ship DSpark stages under mtp.* "
+                            "(main_proj/markov_head/confidence_head) and carry "
+                            "no e_proj/h_proj/enorm/hnorm/hc_head weights. Use "
+                            "speculative method 'dspark' instead of 'mtp'."
+                        )
                     if (
                         self.num_speculative_tokens > 1
                         and self.draft_model_config.hf_config.model_type
@@ -1398,11 +1428,17 @@ def __post_init__(self):
                     and "Gemma4DSparkModel" not in self.draft_model_config.architectures
                     and "K3DSparkModel" not in self.draft_model_config.architectures
                 ):
-                    # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config
+                    # DeepSeek-V4(.1) DSpark reuses the full target config
                     # and its weights ship in the target checkpoint.
-                    self.draft_model_config.hf_config.model_type = "deepseek_v4"
-                    self.draft_model_config.hf_config.architectures = [
-                        "DSparkDraftModel"
+                    is_v41 = (
+                        self.target_model_config.hf_config.model_type == "deepseek_v41"
+                    )
+                    draft_hf_config = self.draft_model_config.hf_config
+                    draft_hf_config.model_type = (
+                        "deepseek_v41" if is_v41 else "deepseek_v4"
+                    )
+                    draft_hf_config.architectures = [
+                        "DSparkV41DraftModel" if is_v41 else "DSparkDraftModel"
                     ]
                     self.draft_model_config.quantization = (
                         self.target_model_config.quantization
@@ -1420,11 +1456,6 @@ def __post_init__(self):
                         and getattr(hf, "target_layer_ids", None) is not None
                     ):
                         hf.dspark_target_layer_ids = hf.target_layer_ids
-                    if (
-                        getattr(hf, "n_predict", None) is None
-                        and getattr(hf, "block_size", None) is not None
-                    ):
-                        hf.n_predict = hf.block_size
 
                 if self.method in ("dflash", "dspark"):
                     self.parallel_drafting = True
@@ -1439,7 +1470,15 @@ def __post_init__(self):
                 n_predict = getattr(
                     self.draft_model_config.hf_config, "n_predict", None
                 )
-                if n_predict is not None:
+                if self.use_dspark():
+                    if self.num_speculative_tokens is None:
+                        # DSpark's parallel width is independent of MTP stages.
+                        hf_config = self.draft_model_config.hf_config
+                        block_size = getattr(hf_config, "block_size", None)
+                        if block_size is None:
+                            block_size = getattr(hf_config, "dspark_block_size", None)
+                        self.num_speculative_tokens = block_size
+                elif n_predict is not None:
                     if self.num_speculative_tokens is None:
                         # Default to max value defined in draft model config.
                         self.num_speculative_tokens = n_predict
@@ -1514,9 +1553,13 @@ def __post_init__(self):
                     )
                 )
 
+                # Use the final draft config, after any architecture overrides,
+                # so EP agrees with ModelConfig's expert-count validation.
                 self.draft_parallel_config = (
                     SpeculativeConfig.create_draft_parallel_config(
-                        self.target_parallel_config, self.draft_tensor_parallel_size
+                        self.target_parallel_config,
+                        self.draft_tensor_parallel_size,
+                        draft_model_config=self.draft_model_config,
                     )
                 )
 
@@ -1529,8 +1572,16 @@ def __post_init__(self):
                 self.index_share_for_mtp_iteration
             )
 
-        if self.method != "dspark" and self.enable_adaptive_verification:
-            raise ValueError("Adaptive verification only supported with DSpark")
+        if (
+            self.method != "dspark"
+            and self.enable_adaptive_verification
+            and self.use_local_argmax_reduction
+        ):
+            raise ValueError(
+                "Adaptive verification estimates per-position acceptance from "
+                "the draft logits, which use_local_argmax_reduction never "
+                "materializes. Disable one of them."
+            )
 
         return self
 
@@ -1587,7 +1638,6 @@ def _maybe_override_draft_max_model_len(
         speculative_max_model_len is mainly used for testing that sequences can
         skip speculation.
         """
-
         if speculative_max_model_len is not None:
             if speculative_max_model_len > draft_max_model_len:
                 raise ValueError(
@@ -1617,7 +1667,7 @@ def _maybe_override_draft_max_model_len(
 
     @staticmethod
     def _maybe_override_draft_max_position_embeddings(
-        draft_hf_config: PretrainedConfig,
+        draft_hf_config: PreTrainedConfig,
         target_max_model_len: int,
     ) -> None:
         """Raise an EAGLE draft's max_position_embeddings up to the target's.
@@ -1631,6 +1681,7 @@ def _maybe_override_draft_max_position_embeddings(
         Args:
             draft_hf_config: The draft model's HF config, mutated in place.
             target_max_model_len: The target model's max_model_len.
+
         """
         draft_max_position_embeddings = getattr(
             draft_hf_config, "max_position_embeddings", None
@@ -1653,10 +1704,9 @@ def _maybe_override_draft_max_position_embeddings(
     def _verify_and_get_draft_tp(
         target_parallel_config: ParallelConfig,
         speculative_draft_tensor_parallel_size: int | None,
-        draft_hf_config: PretrainedConfig,
+        draft_hf_config: PreTrainedConfig,
     ) -> int:
-        """
-        Verifies and adjusts the tensor parallel size for a draft model
+        """Verifies and adjusts the tensor parallel size for a draft model
         specified using speculative_draft_tensor_parallel_size.
         """
         # If speculative_draft_tensor_parallel_size is unset then set it
@@ -1685,8 +1735,7 @@ def _verify_and_get_draft_tp(
         return speculative_draft_tensor_parallel_size
 
     def update_arch_(self):
-        """
-        EagleConfig and ExtractHiddenStatesConfig update architectures, so update all
+        """EagleConfig and ExtractHiddenStatesConfig update architectures, so update all
         architectures-related fields in self.draft_model_config
         """
         self.draft_model_config.hf_text_config = get_hf_text_config(
@@ -1706,14 +1755,21 @@ def update_arch_(self):
     def create_draft_parallel_config(
         target_parallel_config: ParallelConfig,
         speculative_draft_tensor_parallel_size: int,
+        draft_model_config: ModelConfig | None = None,
     ) -> ParallelConfig:
         """Create a parallel config for use by the draft worker.
 
-        This is mostly a copy of the target parallel config, except the tp_size.
+        Use the draft TP size and disable inherited EP for known dense drafts.
+        Without a draft model config, preserve the previous EP inheritance.
         """
+        enable_ep = target_parallel_config.enable_expert_parallel
+        if draft_model_config is not None:
+            enable_ep = enable_ep and draft_model_config.is_moe
+
         draft_parallel_config = ParallelConfig(
             pipeline_parallel_size=1,
             tensor_parallel_size=speculative_draft_tensor_parallel_size,
+            enable_expert_parallel=enable_ep,
             distributed_executor_backend=target_parallel_config.distributed_executor_backend,
             max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers,
             disable_custom_all_reduce=target_parallel_config.disable_custom_all_reduce,
@@ -1723,6 +1779,18 @@ def create_draft_parallel_config(
 
         return draft_parallel_config
 
+    def apply_draft_overrides(self, vllm_config: VllmConfig) -> VllmConfig:
+        """Overlay this config's kernel overrides onto a target VllmConfig.
+
+        Only non-None fields override, so an unset field keeps whatever the
+        target resolved.
+        """
+        for src, config_name, dst in _DRAFT_VLLM_CONFIG_OVERRIDES:
+            if (value := getattr(self, src)) is not None:
+                sub_config = replace(getattr(vllm_config, config_name), **{dst: value})
+                vllm_config = replace(vllm_config, **{config_name: sub_config})
+        return vllm_config
+
     @field_validator("attention_backend", mode="before")
     @classmethod
     def _parse_attention_backend(cls, value: Any) -> Any:
@@ -1887,6 +1955,13 @@ def uses_dynamic_speculative_decoding(self) -> bool:
     def uses_draft_model(self) -> bool:
         return self.method == "draft_model"
 
+    def uses_draft_kv_cache(self) -> bool:
+        return (
+            self.use_eagle()
+            or self.uses_draft_model()
+            or self.uses_extract_hidden_states()
+        )
+
     def uses_extract_hidden_states(self) -> bool:
         return self.method == "extract_hidden_states"
 
diff --git a/vllm/config/structured_outputs.py b/vllm/config/structured_outputs.py
index e7afbb65bc7f..059581e37e65 100644
--- a/vllm/config/structured_outputs.py
+++ b/vllm/config/structured_outputs.py
@@ -42,8 +42,7 @@ class StructuredOutputsConfig:
     """Whether to use structured input for reasoning."""
 
     def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
diff --git a/vllm/config/utils.py b/vllm/config/utils.py
index 6ab346d1a030..c3d7bb10563d 100644
--- a/vllm/config/utils.py
+++ b/vllm/config/utils.py
@@ -64,7 +64,9 @@ def config(
         cls: The class to decorate
         config: The pydantic ConfigDict to use. If provided, it will be merged with
             the default config.
-        **kwargs: Additional arguments to pass to pydantic.dataclass."""
+        **kwargs: Additional arguments to pass to pydantic.dataclass.
+
+    """
     # Extra fields are forbidden by default
     merged_config = ConfigDict(extra="forbid")
     if config is not None:
@@ -134,10 +136,9 @@ def getattr_iter(
     default_factory: Callable[[], Any] | None = None,
     warn: bool = False,
 ) -> Any:
-    """
-    A helper function that retrieves an attribute from an object which may
+    """A helper function that retrieves an attribute from an object which may
     have multiple possible names. This is useful when fetching attributes from
-    arbitrary `transformers.PretrainedConfig` instances.
+    arbitrary `transformers.PreTrainedConfig` instances.
 
     In the case where the first name in `names` is the preferred name, and
     any other names are deprecated aliases, setting `warn=True` will log a
@@ -158,12 +159,10 @@ def getattr_iter(
 
 
 def get_attr_docs(cls: type[Any]) -> dict[str, str]:
-    """
-    Get any docstrings placed after attribute assignments in a class body.
+    """Get any docstrings placed after attribute assignments in a class body.
 
     https://davidism.com/mit-license/
     """
-
     cls_node = ast.parse(textwrap.dedent(inspect.getsource(cls))).body[0]
 
     if not isinstance(cls_node, ast.ClassDef):
@@ -315,7 +314,7 @@ def normalize_value(x):
         except Exception:
             return str(x)
 
-    # PretrainedConfig (must be before dataclass branch as these are now dataclasses)
+    # PreTrainedConfig (must be before dataclass branch as these are now dataclasses)
     if hasattr(x, "to_json_string") and callable(x.to_json_string):
         try:
             return x.to_json_string()
@@ -385,8 +384,7 @@ def hash_factors(items: dict[str, object]) -> str:
 
 @dataclass
 class Range:
-    """
-    A range of numbers.
+    """A range of numbers.
     Inclusive of start, inclusive of end.
     """
 
@@ -445,8 +443,7 @@ def get_from_deprecated_env_if_set(
     removal_version: str,
     field_name: str | None = None,
 ) -> str | None:
-    """
-    Get value from deprecated environment variable with warning.
+    """Get value from deprecated environment variable with warning.
 
     Args:
         env_name: Name of the deprecated environment variable
@@ -455,6 +452,7 @@ def get_from_deprecated_env_if_set(
 
     Returns:
         The environment variable value if set, None otherwise
+
     """
     if envs.is_set(env_name):
         value = os.environ.get(env_name)
@@ -477,8 +475,7 @@ def set_from_deprecated_env_if_set(
     to_bool: bool = False,
     to_int: bool = False,
 ) -> None:
-    """
-    Set object field from deprecated environment variable with warning.
+    """Set object field from deprecated environment variable with warning.
 
     Args:
         config: Config object to set the field on
@@ -489,6 +486,7 @@ def set_from_deprecated_env_if_set(
         to_int: Whether to convert the environment variable value to integer
     Returns:
         None
+
     """
     if to_bool and to_int:
         raise ValueError("Cannot convert to both boolean and integer.")
diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py
index 294f80b143e3..0aafd2506cb6 100644
--- a/vllm/config/vllm.py
+++ b/vllm/config/vllm.py
@@ -27,14 +27,15 @@
 from vllm.utils import random_uuid
 from vllm.utils.hashing import safe_hash
 
-from .attention import AttentionConfig
+from .attention import AttentionConfig, HiSparseConfig
+from .aux_output import AuxOutputConfig
 from .cache import CacheConfig
 from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode
 from .device import DeviceConfig
 from .diffusion import DiffusionConfig
 from .ec_manager_config import EncoderCacheManagerConfig
 from .ec_transfer import ECTransferConfig
-from .engram import EngramConfig
+from .engram import EngramConfig, model_has_engram_layers
 from .kernel import KernelConfig
 from .kv_events import KVEventsConfig
 from .kv_transfer import KVTransferConfig
@@ -50,17 +51,17 @@
 from .scheduler import SchedulerConfig
 from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig
 from .structured_outputs import StructuredOutputsConfig
-from .utils import SupportsHash, config, replace
+from .utils import SupportsHash, config, get_field, replace
 from .watermarking import WatermarkConfig
 from .weight_transfer import WeightTransferConfig
 
 if TYPE_CHECKING:
-    from transformers import PretrainedConfig
+    from transformers import PreTrainedConfig
 
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
     from vllm.v1.kv_cache_interface import KVCacheConfig
 else:
-    PretrainedConfig = Any
+    PreTrainedConfig = Any
 
     QuantizationConfig = Any
 
@@ -81,6 +82,7 @@
         "DeepseekV4ForCausalLM",
         "DeepseekV4ForConditionalGeneration",
         "DeepSeekV4MTPModel",
+        "DeepseekV41ForCausalLM",
         "Dots3NoteForCausalLM",
         "Dots3NoteMTPModel",
         "Glm5NextForCausalLM",
@@ -109,10 +111,17 @@ def default_breakable_cudagraph_architectures() -> frozenset[str]:
     from vllm.platforms import current_platform
 
     if current_platform.is_rocm():
-        # Breakable CUDA graphs currently regress performance on ROCm, so no
-        # architecture opts in by default here. Users can still force it with
+        # Breakable CUDA graphs currently regress performance on ROCm for
+        # models that can use torch.compile piecewise graphs instead. Do not
+        # opt those in by default. Users can still force them with
         # VLLM_USE_BREAKABLE_CUDAGRAPH=1.
-        return frozenset()
+        #
+        # DeepseekV41ForCausalLM cannot torch.compile, and the ROCm sparse
+        # SWA backend only reports AttentionCGSupport.UNIFORM_BATCH. Default
+        # FULL_AND_PIECEWISE then dies at capture unless breakable CUDA
+        # graphs are on. Enable this architecture so the published AMD
+        # recipe can start.
+        return frozenset({"DeepseekV41ForCausalLM"})
     return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES
 
 
@@ -146,7 +155,6 @@ class OptimizationLevel(IntEnum):
 def enable_norm_fusion(cfg: "VllmConfig") -> bool:
     """Enable if either RMS norm or quant FP8 custom op is active;
     otherwise Inductor handles fusion."""
-
     return (
         cfg.compilation_config.is_custom_op_enabled("rms_norm")
         or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
@@ -155,8 +163,7 @@ def enable_norm_fusion(cfg: "VllmConfig") -> bool:
 
 
 def enable_act_fusion(cfg: "VllmConfig") -> bool:
-    """
-    Enable if either SiLU+Mul or quant FP8 custom op is active;
+    """Enable if either SiLU+Mul or quant FP8 custom op is active;
     otherwise Inductor handles fusion.
     Also enable for FP4 models as FP4 quant is always custom so Inductor cannot fuse it.
     """
@@ -212,7 +219,6 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool:
 
 def enable_rope_kvcache_mla_fusion(cfg: "VllmConfig") -> bool:
     """Enable if use_inductor_graph_partition is enabled."""
-
     return (
         cfg.compilation_config.use_inductor_graph_partition
         or not cfg.compilation_config.splitting_ops_contain_kv_cache_update()
@@ -221,7 +227,6 @@ def enable_rope_kvcache_mla_fusion(cfg: "VllmConfig") -> bool:
 
 def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
     """Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss."""
-
     return (
         cfg.kernel_config.ir_op_priority.fused_add_rms_norm[0] == "aiter"
         and cfg.model_config is not None
@@ -237,7 +242,7 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool:
 
 
 def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool:
-    """Enable fused QK-norm + RoPE + KV cache update on ROCm with AITER."""
+    """Enable fused QK-norm + RoPE/MRoPE + KV cache update with AITER."""
     from vllm._aiter_ops import rocm_aiter_ops
 
     if not rocm_aiter_ops.is_enabled():
@@ -372,8 +377,10 @@ class VllmConfig:
     """Model weight offloading configuration."""
     attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
     """Attention configuration."""
+    aux_output_config: AuxOutputConfig = Field(default_factory=AuxOutputConfig)
+    """Execution auxiliary output configuration."""
     engram_config: EngramConfig | None = None
-    """Optional Engram configuration, only valid for supported PLE models."""
+    """N-gram embedding storage and sharding settings."""
     mamba_config: MambaConfig = Field(default_factory=MambaConfig)
     """Mamba configuration."""
     kernel_config: KernelConfig = Field(default_factory=KernelConfig)
@@ -452,8 +459,7 @@ class VllmConfig:
     """
 
     def compute_hash(self, include_version: bool = True) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
+        """WARNING: Whenever a new field is added to this config,
         ensure that it is included in the factors list if
         it affects the computation graph.
 
@@ -465,6 +471,7 @@ def compute_hash(self, include_version: bool = True) -> str:
 
         Args:
             include_version: Include the vLLM version in the hash.
+
         """
         factors: list[Any] = []
 
@@ -550,6 +557,7 @@ def compute_hash(self, include_version: bool = True) -> str:
             vllm_factors.append(self.ec_transfer_config.compute_hash())
         else:
             vllm_factors.append("None")
+        vllm_factors.append(self.aux_output_config.compute_hash())
         if self.additional_config:
             if isinstance(additional_config := self.additional_config, dict):
                 additional_config_hash = safe_hash(
@@ -642,6 +650,29 @@ def num_lookahead_tokens(self) -> int:
             return self.num_speculative_tokens
         return 0
 
+    @property
+    def num_prefill_lookahead_tokens(self) -> int:
+        """Prefill tokens past the computed range that the drafter reads.
+
+        Mid-prefill the drafter consumes tokens the target model has not been
+        scheduled for yet, so every component that has to keep them available
+        must apply this margin: the scheduler, which never ends a chunk within
+        it and shifts encoder scheduling by it, and the KV cache manager, which
+        treats the trailing `this - 1` tokens as re-prefillable rather than
+        finalized. Consumers must read this property rather than re-deriving
+        their own per-method lookahead, so those components cannot drift apart.
+        """
+        speculative_config = self.speculative_config
+        if speculative_config is None or not speculative_config.use_eagle():
+            return 0
+        if speculative_config.use_multi_module_mtp():
+            # Each MTP module reads one token further ahead than the one before
+            # it, so the chain needs num_speculative_tokens of runway at a
+            # chunked-prefill boundary.
+            return self.num_speculative_tokens
+        # Eagle-family drafters read only the immediate next token.
+        return 1
+
     @property
     def uniform_decode_query_len(self) -> int:
         """Query length of every request in a uniform decode batch.
@@ -665,6 +696,14 @@ def uniform_decode_query_len(self) -> int:
 
     @property
     def use_v2_model_runner(self) -> bool:
+        if self.attention_config.hisparse_config is not None:
+            if envs.VLLM_USE_V2_MODEL_RUNNER is False:
+                raise ValueError(
+                    "HiSparse requires Model Runner V2; remove "
+                    "VLLM_USE_V2_MODEL_RUNNER=0."
+                )
+            return True
+
         if getattr(self, "watermark_config", None) is not None:
             if envs.VLLM_USE_V2_MODEL_RUNNER is False:
                 logger.info_once(
@@ -737,6 +776,45 @@ def _uses_breakable_cudagraph_by_default(self) -> bool:
         architectures = set(model_config.architectures)
         return bool(architectures & default_breakable_cudagraph_architectures())
 
+    def _uses_breakable_cudagraph_for_batch_invariance(self) -> bool:
+        """Avoid freezing runtime-M tile lookup in compiled forward (#54243).
+        Breakable graphs look up tuned bf16, unquantized qkv/o/gate_up/down tiles
+        at capture; lm_head runs outside compiled forward and does not benefit."""
+        from vllm.model_executor.determinism import batch_invariant_configs as bi
+        from vllm.platforms import current_platform
+
+        model = self.model_config
+        if (
+            not envs.VLLM_BATCH_INVARIANT
+            or model is None
+            or model.enforce_eager
+            or model.dtype != torch.bfloat16
+            or model.quantization is not None
+            or not current_platform.is_cuda()
+            # Sequence parallelism / async TP are torch.compile passes.
+            or self.compilation_config.pass_config.enable_sp
+            or self.compilation_config.pass_config.fuse_gemm_comms
+        ):
+            return False
+        family = bi._get_tuned_matmul_arch_family(
+            current_platform.get_device_capability()
+        )
+        if family is None or family not in bi._BATCH_INVARIANT_MATMUL_TUNED_CONFIGS:
+            return False
+        table = bi._BATCH_INVARIANT_MATMUL_TUNED_CONFIGS[family]
+        parallel = self.parallel_config
+        tp = parallel.tensor_parallel_size
+        hidden = model.get_hidden_size()
+        head = model.get_head_size()
+        heads = model.get_num_attention_heads(parallel)
+        kv_heads = model.get_num_kv_heads(parallel)
+        shapes = [((heads + 2 * kv_heads) * head, hidden), (hidden, heads * head)]
+        intermediate = getattr(model.hf_text_config, "intermediate_size", None)
+        # Per-layer sizes (e.g. Gemma3n) are not modeled.
+        if isinstance(intermediate, int):
+            shapes += [(2 * intermediate // tp, hidden), (hidden, intermediate // tp)]
+        return any(shape in table for shape in shapes)
+
     def _maybe_enable_breakable_cudagraph(self) -> bool:
         if (
             "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
@@ -747,6 +825,16 @@ def _maybe_enable_breakable_cudagraph(self) -> bool:
                 "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
                 "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
             )
+        elif (
+            "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
+            and envs.VLLM_BATCH_INVARIANT
+            and self._uses_breakable_cudagraph_for_batch_invariance()
+        ):
+            os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
+            logger.info_once(
+                "VLLM_BATCH_INVARIANT=1: auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
+                "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
+            )
 
         from vllm.compilation.breakable_cudagraph import (
             is_breakable_cudagraph_enabled,
@@ -759,8 +847,7 @@ def _maybe_enable_breakable_cudagraph(self) -> bool:
 
     @property
     def needs_dp_coordinator(self) -> bool:
-        """
-        Determine if the DPCoordinator process is needed.
+        """Determine if the DPCoordinator process is needed.
 
         The DPCoordinator is needed in two cases:
         1. For MoE models with DP > 1: to handle wave coordination
@@ -770,8 +857,8 @@ def needs_dp_coordinator(self) -> bool:
 
         Returns:
             True if DPCoordinator process is needed, False otherwise.
-        """
 
+        """
         # For non-MoE models, only need coordinator in internal/hybrid LB mode
         # (for stats collection).
         return self.parallel_config.data_parallel_size > 1 and (
@@ -781,8 +868,7 @@ def needs_dp_coordinator(self) -> bool:
         )
 
     def enable_trace_function_call_for_thread(self) -> None:
-        """
-        Set up function tracing for the current thread,
+        """Set up function tracing for the current thread,
         if enabled via the `VLLM_TRACE_FUNCTION` environment variable.
         """
         if envs.VLLM_TRACE_FUNCTION:
@@ -853,7 +939,7 @@ def get_quantization_config(
 
     def with_hf_config(
         self,
-        hf_config: PretrainedConfig,
+        hf_config: PreTrainedConfig,
         architectures: list[str] | None = None,
     ) -> "VllmConfig":
         if architectures is not None:
@@ -906,6 +992,7 @@ def with_hf_config(
 
         model_config.hf_config = hf_config
         model_config.model_arch_config = model_config.get_model_arch_config()
+        model_config.is_submodel_config = True
 
         return replace(self, model_config=model_config)
 
@@ -916,6 +1003,7 @@ def _set_config_default(self, config_obj: Any, key: str, value: Any) -> None:
             config_obj: Configuration object to update.
             key: Attribute name.
             value: Default value (static or callable).
+
         """
         if getattr(config_obj, key) is None:
             # Some config values are known before initialization and are
@@ -937,6 +1025,7 @@ def _apply_optimization_level_defaults(self, defaults: dict[str, Any]) -> None:
 
         Args:
             defaults: Dictionary of default values to apply.
+
         """
 
         def apply_recursive(config_obj: Any, config_defaults: dict[str, Any]) -> None:
@@ -991,6 +1080,28 @@ def _maybe_disable_dynamic_sd_for_data_parallel(self) -> None:
         )
         speculative_config.num_speculative_tokens_per_batch_size = None
 
+    def _normalize_piecewise_cudagraph_mode(
+        self, *, breakable_cudagraph_enabled: bool
+    ) -> None:
+        compilation_config = self.compilation_config
+        if (
+            compilation_config.cudagraph_mode.requires_piecewise_compilation()
+            and compilation_config.mode != CompilationMode.VLLM_COMPILE
+            and not breakable_cudagraph_enabled
+        ):
+            fallback_mode = compilation_config.cudagraph_mode.without_piecewise()
+            logger.info_once(
+                "Cudagraph mode %s is not compatible with compilation mode %s. "
+                "Overriding to %s.",
+                compilation_config.cudagraph_mode,
+                compilation_config.mode,
+                fallback_mode,
+            )
+            compilation_config.cudagraph_mode = fallback_mode
+            if fallback_mode == CUDAGraphMode.NONE:
+                compilation_config.max_cudagraph_capture_size = 0
+                compilation_config.cudagraph_capture_sizes = []
+
     def _post_init_kv_transfer_config(self) -> None:
         """Update KVTransferConfig based on top-level configs in VllmConfig.
 
@@ -1027,6 +1138,58 @@ def _post_init_kv_transfer_config(self) -> None:
         # This is the same for all backends
         self.kv_transfer_config.kv_role = "kv_both"
 
+    def _verify_aux_output_compatibility(self) -> None:
+        """Reject configurations unsupported by enabled auxiliary outputs."""
+        if not self.aux_output_config.enabled:
+            return
+        if not self.use_v2_model_runner:
+            raise ValueError(
+                "AuxOutput Connector requires Model Runner V2; set "
+                "VLLM_USE_V2_MODEL_RUNNER=1."
+            )
+        if self.model_config.runner_type != "generate":
+            raise ValueError("AuxOutput Connector only supports generate runners.")
+        if not self.model_config.is_moe:
+            raise ValueError("AuxOutput Connector only supports MoE models.")
+        if not self.cache_config.enable_prefix_caching:
+            raise ValueError("AuxOutput Connector requires prefix caching.")
+        if (
+            self.speculative_config is not None
+            and self.speculative_config.enable_adaptive_verification
+        ):
+            raise ValueError(
+                "--enable-return-routed-experts is incompatible with "
+                "adaptive speculative verification."
+            )
+        if self.parallel_config.pipeline_parallel_size > 1:
+            raise ValueError(
+                "--enable-return-routed-experts is incompatible with "
+                "pipeline parallelism (PP > 1)."
+            )
+        if (
+            self.parallel_config.decode_context_parallel_size > 1
+            or self.parallel_config.prefill_context_parallel_size > 1
+        ):
+            raise ValueError(
+                "--enable-return-routed-experts is incompatible with "
+                "context parallelism (DCP/PCP > 1)."
+            )
+
+        kv_transfer_config = self.kv_transfer_config
+        if kv_transfer_config is not None:
+            for connector_name in (
+                "NixlConnector",
+                "NixlPullConnector",
+                "NixlPushConnector",
+                "MoRIIOConnector",
+                "MooncakeConnector",
+            ):
+                if kv_transfer_config.has_connector(connector_name):
+                    raise ValueError(
+                        "--enable-return-routed-experts is incompatible with "
+                        f"{connector_name}; PD auxiliary output is not supported."
+                    )
+
     def _verify_kv_transfer_compat(self) -> None:
         """Reject configurations that silently corrupt KV transfers."""
         if (
@@ -1111,14 +1274,67 @@ def _check_watermarking_unsupported(
         watermark_config = getattr(self, "watermark_config", None)
         if watermark_config is None:
             return
-        if (
-            self.speculative_config is not None
-            and not watermark_config.supports_speculative_decoding
-        ):
-            raise ValueError(
-                f"The {watermark_config.algorithm} watermarking algorithm "
-                "does not support speculative decoding."
-            )
+        if self.speculative_config is not None:
+            speculative_config = self.speculative_config
+            if speculative_config.draft_sample_method != "probabilistic":
+                raise ValueError(
+                    "Speculative decoding with watermarking requires "
+                    "draft_sample_method='probabilistic'."
+                )
+            if speculative_config.rejection_sample_method != "standard":
+                raise ValueError(
+                    "Speculative decoding with watermarking requires "
+                    "rejection_sample_method='standard'."
+                )
+            if (
+                speculative_config.parallel_drafting
+                and speculative_config.method != "dspark"
+            ):
+                raise ValueError(
+                    "Parallel speculative drafting is not supported with watermarking."
+                )
+            if speculative_config.method not in ("dspark", "eagle", "eagle3", "mtp"):
+                raise ValueError(
+                    "Watermarking supports only autoregressive model-based "
+                    "speculative decoding."
+                )
+            if (
+                not watermark_config.allow_target_only_watermarking
+                and not watermark_config.supports_speculative_decoding
+            ):
+                raise ValueError(
+                    f"The '{watermark_config.algorithm}' watermarking algorithm "
+                    "does not support speculative decoding. Set "
+                    "allow_target_only_watermarking=true to leave draft tokens "
+                    "unwatermarked."
+                )
+            if (
+                watermark_config.allow_target_only_watermarking
+                and not watermark_config.supports_speculative_decoding
+            ):
+                logger.warning_once(
+                    "Target-only watermarking leaves accepted draft tokens "
+                    "unwatermarked, weakening detectability in proportion to the "
+                    "share of output tokens supplied by accepted drafts.",
+                    scope="global",
+                )
+            if watermark_config.deduplicate_contexts != "none":
+                logger.warning_once(
+                    "Context deduplication is not supported with speculative "
+                    "decoding and will not be applied to accepted drafts, "
+                    "rejection-recovery tokens, or bonus tokens.",
+                    scope="global",
+                )
+            if watermark_config.algorithm == "dual_key_gumbel" and (
+                watermark_config.alpha != get_field(WatermarkConfig, "alpha").default
+            ):
+                logger.warning_once(
+                    "Speculative decoding selects the watermark key by token role: "
+                    "draft tokens use key A, recovery and bonus tokens use key B. "
+                    "The configured alpha=%s is not used.",
+                    watermark_config.alpha,
+                    scope="global",
+                )
         if beam_search:
             raise ValueError("Beam search is not supported with watermarking.")
         if custom_sampler:
@@ -1127,32 +1343,65 @@ def _check_watermarking_unsupported(
             )
 
     def _resolve_and_verify_engram_config(self) -> None:
-        """Resolve legacy offload settings and validate model and parallel configs."""
-        if self.engram_config is None:
-            if not envs.VLLM_PLE_CPU_OFFLOAD:
-                return
-            self.engram_config = EngramConfig()
+        """Resolve defaults and validate n-gram embedding settings."""
+        from vllm.platforms import current_platform
+
         model_config = self.model_config
         speculative_config = self.speculative_config
         # Draft configs inherit the target's communication groups and settings.
-        # Qwen4Exp MTP itself disables PLE, so validate its target instead.
+        # Validate the target because the draft may disable n-gram embeddings.
         if (
             speculative_config is not None
             and model_config is speculative_config.draft_model_config
         ):
             model_config = speculative_config.target_model_config
+        if (
+            model_config is not None
+            and model_config.architecture == "DeepseekV41ForCausalLM"
+            and getattr(model_config.hf_text_config, "engram_layer_ids", None)
+            and self.parallel_config.use_ubatching
+        ):
+            raise ValueError(
+                "DeepSeek V4.1 Engram does not support DBO or microbatching. "
+                "Disable --enable-dbo and set --ubatch-size to 0."
+            )
+        if self.engram_config is None:
+            if not current_platform.is_cuda_alike() or not model_has_engram_layers(
+                model_config
+            ):
+                return
+            self.engram_config = EngramConfig()
         self.engram_config.verify_model_config(model_config)
+        self.engram_config.resolve_dp_shared_memory(self.parallel_config)
         self.engram_config.verify_parallel_config(self.parallel_config)
         logger.info_once("Resolved Engram configuration: %s", str(self.engram_config))
 
     def __post_init__(self):
         """Verify configs are valid & consistent with each other."""
-
         # To give each torch profile run a unique instance name.
         self.instance_id = f"{time.time_ns()}"
 
+        if self.model_config is not None and self.model_config.is_submodel_config:
+            # with_hf_config() view: the parent config was already validated,
+            # and this view's empty architecture list makes the model-dependent checks
+            # below unsafe (e.g. use_mla resolves the architecture registry).
+            return
+
         self._resolve_mm_encoder_only()
 
+        if self.is_mm_encoder_only and self.cache_config.enable_prefix_caching:
+            # Such an instance publishes encoder embeddings and runs no language
+            # model, so it holds no KV cache for prefix caching to reuse and its
+            # coordinator would have no group to manage. Disable before
+            # `try_verify_and_update_config` so model config hooks (e.g. the
+            # hybrid mamba hook setting `mamba_block_size`) already see prefix
+            # caching as disabled.
+            logger.info(
+                "Disabling prefix caching: this instance runs the "
+                "multi-modal encoder only."
+            )
+            self.cache_config.enable_prefix_caching = False
+
         if self.performance_mode != "balanced":
             logger.info_once("Performance mode set to '%s'.", self.performance_mode)
 
@@ -1170,38 +1419,6 @@ def __post_init__(self):
 
             self.parallel_config.is_moe_model = self.model_config.is_moe
 
-        if (
-            self.model_config is not None
-            and self.model_config.enable_return_routed_experts
-        ):
-            if self.parallel_config.pipeline_parallel_size > 1:
-                raise ValueError(
-                    "--enable-return-routed-experts is incompatible with "
-                    "pipeline parallelism (PP > 1)."
-                )
-            if (
-                self.parallel_config.decode_context_parallel_size > 1
-                or self.parallel_config.prefill_context_parallel_size > 1
-            ):
-                raise ValueError(
-                    "--enable-return-routed-experts is incompatible with context "
-                    "parallelism (DCP > 1 or PCP > 1)."
-                )
-
-            # Incompatible with any KV connector — covers both PD disaggregation
-            # (kv_producer/kv_consumer: routing captured on P can't reach D) and
-            # single-instance KV offload/sharing (kv_both: slot_mapping semantics
-            # change when KV blocks live outside local GPU memory, breaking the
-            # slot-indexed routed_experts buffer).
-            if (
-                self.kv_transfer_config is not None
-                and self.kv_transfer_config.is_kv_transfer_instance
-            ):
-                raise ValueError(
-                    "--enable-return-routed-experts is incompatible with KV "
-                    "connectors (PD disaggregation, KV cache offload)."
-                )
-
         if (
             self.model_config is not None
             and self.model_config.multimodal_config is not None
@@ -1223,25 +1440,20 @@ def __post_init__(self):
             self.kv_transfer_config is not None
             and self.kv_transfer_config.has_connector("NixlConnector")
         ):
-            assert self.parallel_config.prefill_context_parallel_size == 1, (
-                "NIXL does not support prefill context parallelism."
-            )
             dcp_size = self.parallel_config.decode_context_parallel_size
-            tp_size = self.parallel_config.tensor_parallel_size
-            assert dcp_size in (1, tp_size), (
+            transfer_tp_size = max(
+                self.parallel_config.tensor_parallel_size,
+                self.parallel_config.prefill_context_parallel_size,
+            )
+            assert dcp_size in (1, transfer_tp_size), (
                 f"decode_context_parallel_size={dcp_size} must be 1 or equal "
-                f"to tensor_parallel_size={tp_size} when using NixlConnector."
+                f"to the NIXL transfer parallel size={transfer_tp_size}."
             )
             if self.model_config is not None:
                 assert self.model_config.use_mla or dcp_size == 1, (
                     "PD with decode_context_parallel_size > 1 is only "
                     "supported for MLA models."
                 )
-                assert not (self.model_config.is_hybrid and dcp_size > 1), (
-                    "PD with decode_context_parallel_size > 1 is not "
-                    "supported for hybrid Mamba/SSM models."
-                )
-
         if self.lora_config is not None:
             self.lora_config.verify_with_model_config(self.model_config)
 
@@ -1319,11 +1531,12 @@ def __post_init__(self):
                     and self.speculative_config.method not in get_args(NgramGPUTypes)
                     and self.speculative_config.method != "draft_model"
                     and self.speculative_config.method != "dspark"
+                    and self.speculative_config.method != "dflash"
                 ):
                     raise ValueError(
                         "Currently, async scheduling is only supported "
-                        "with EAGLE/MTP/Draft Model/NGram GPU/DSpark kind of "
-                        "speculative decoding"
+                        "with EAGLE/MTP/Draft Model/NGram GPU/DSpark/DFlash "
+                        "kind of speculative decoding"
                     )
                 if self.speculative_config.disable_padded_drafter_batch:
                     raise ValueError(
@@ -1352,6 +1565,7 @@ def __post_init__(self):
                 and self.speculative_config.method not in get_args(NgramGPUTypes)
                 and self.speculative_config.method != "draft_model"
                 and self.speculative_config.method != "dspark"
+                and self.speculative_config.method != "dflash"
             ):
                 logger.warning_once(
                     "Async scheduling not supported with %s-based "
@@ -1382,6 +1596,15 @@ def __post_init__(self):
                     "DP+EP generation accuracy."
                 )
                 self.scheduler_config.async_scheduling = False
+            elif (
+                self.parallel_config.pipeline_parallel_size > 1
+                and not self.use_v2_model_runner
+            ):
+                logger.warning_once(
+                    "Async scheduling is disabled because the V1 model runner "
+                    "does not support it with pipeline parallelism."
+                )
+                self.scheduler_config.async_scheduling = False
             else:
                 self.scheduler_config.async_scheduling = True
 
@@ -1444,23 +1667,14 @@ def __post_init__(self):
 
         if self.model_config is not None and self.model_config.enforce_eager:
             logger.warning_once(
-                "Enforce eager set, disabling torch.compile and CUDAGraphs. "
-                "This is equivalent to setting -cc.mode=none -cc.cudagraph_mode=none"
+                "Enforce eager set, disabling torch.compile, CUDAGraphs, and JIT "
+                "kernel warmup. This is equivalent to setting -cc.mode=none "
+                "-cc.cudagraph_mode=none and "
+                "--kernel_config.enable_jit_warmup=False"
             )
             self.compilation_config.mode = CompilationMode.NONE
             self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
-
-        if self.profiler_config.profiler == "proton":
-            if not current_platform.is_cuda():
-                raise ValueError(
-                    "The Proton profiler currently supports NVIDIA CUDA only"
-                )
-            if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
-                raise ValueError(
-                    "The Proton profiler requires CUDA graphs to be disabled. "
-                    "Use --enforce-eager or set "
-                    "--compilation-config.cudagraph_mode=none."
-                )
+            self.kernel_config.enable_jit_warmup = False
 
         if os.environ.get("TORCH_COMPILE_DISABLE") == "1":
             logger.warning_once(
@@ -1542,17 +1756,48 @@ def has_blocked_weights():
         self._maybe_override_dynamic_sd_cudagraph_mode()
 
         if (
-            self.compilation_config.cudagraph_mode.requires_piecewise_compilation()
-            and self.compilation_config.mode != CompilationMode.VLLM_COMPILE
-            and not envs.VLLM_USE_BREAKABLE_CUDAGRAPH
+            self.attention_config.hisparse_config is None
+            and self.kv_transfer_config is not None
+            and self.kv_transfer_config.has_connector("HiSparseConnector")
         ):
-            logger.info_once(
-                "Cudagraph mode %s is not compatible with compilation mode %s."
-                "Overriding to NONE.",
-                self.compilation_config.cudagraph_mode,
-                self.compilation_config.mode,
-            )
-            self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
+            self.attention_config.hisparse_config = HiSparseConfig()
+
+        if self.attention_config.hisparse_config is not None:
+            if not current_platform.is_cuda():
+                raise ValueError("HiSparse currently requires NVIDIA CUDA.")
+            if self.parallel_config.pipeline_parallel_size > 1:
+                raise ValueError("HiSparse does not support pipeline parallelism.")
+            if self.parallel_config.decode_context_parallel_size > 1:
+                raise ValueError(
+                    "HiSparse does not support decode context parallelism."
+                )
+            if self.model_config is not None and not hasattr(
+                self.model_config.hf_config, "index_topk"
+            ):
+                raise ValueError(
+                    "HiSparse is only supported for DSA models with index_topk."
+                )
+            if self.kv_transfer_config is not None and (
+                self.kv_transfer_config.kv_connector
+                not in (
+                    None,
+                    "NixlConnector",
+                    "MooncakeStoreConnector",
+                    "MultiConnector",
+                )
+            ):
+                logger.warning(
+                    "HiSparse host-resident KV is configured with connector "
+                    "%s. NixlConnector (GPU-staged host imports) and "
+                    "MooncakeStoreConnector (shared-store offload) are the "
+                    "validated paths; other connectors are treated as "
+                    "debug/fallback paths.",
+                    self.kv_transfer_config.kv_connector,
+                )
+
+        self._normalize_piecewise_cudagraph_mode(
+            breakable_cudagraph_enabled=breakable_cudagraph_enabled
+        )
 
         # async tp is built on top of sequence parallelism and requires it.
         pass_config = self.compilation_config.pass_config
@@ -1721,8 +1966,23 @@ def has_blocked_weights():
             )
         current_platform.check_and_update_config(self)
 
-        self._resolve_allow_missing_mm_embeddings()
+        # After the platform hook, which has the last word on async scheduling.
+        if (
+            self.diffusion_config is not None
+            and self.scheduler_config.async_scheduling
+            and self.scheduler_config.scheduler_cls is None
+        ):
+            self.scheduler_config.scheduler_cls = (
+                "vllm.v1.core.sched.diffusion_scheduler.DiffusionAsyncScheduler"
+            )
+
+        self._normalize_piecewise_cudagraph_mode(
+            breakable_cudagraph_enabled=breakable_cudagraph_enabled
+        )
+
+        self._resolve_mm_embedding_inputs()
         self._resolve_mm_processor_device()
+        self._resolve_mm_video_decode_device()
         self._validate_mm_processor_device()
 
         if self.use_v2_model_runner:
@@ -1730,6 +1990,7 @@ def has_blocked_weights():
         else:
             self._validate_v1_model_runner()
 
+        self._validate_profiler_config()
         self._validate_batch_sharded_sampling()
         self._validate_adaptive_verification()
 
@@ -1737,6 +1998,55 @@ def has_blocked_weights():
         # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled)
         self._set_compile_ranges()
 
+        if self.parallel_config.all2all_backend == "moonep":
+            if (
+                self.model_config is not None
+                and self.model_config.quantization is not None
+            ):
+                raise ValueError(
+                    "The moonep all2all backend currently supports unquantized "
+                    "BF16 models only; got "
+                    f"quantization={self.model_config.quantization!r}. Use a "
+                    "different --all2all-backend for quantized models."
+                )
+            if (
+                self.model_config is not None
+                and self.model_config.dtype != torch.bfloat16
+            ):
+                raise ValueError(
+                    "The moonep all2all backend currently supports BF16 models "
+                    f"only; got dtype={self.model_config.dtype}. Use a "
+                    "different --all2all-backend or --dtype bfloat16."
+                )
+            if self.parallel_config.enable_eplb:
+                raise ValueError(
+                    "The moonep all2all backend does not support EPLB yet: "
+                    "EPLB rearranges expert parameters in a layout MoonEP's "
+                    "replicated [E+B] weights do not follow. Disable "
+                    "--enable-eplb or use a different --all2all-backend."
+                )
+            if self.parallel_config.expert_placement_strategy != "linear":
+                raise ValueError(
+                    "The moonep all2all backend requires linear expert "
+                    "placement: its load-time all-gather assumes each rank "
+                    "holds a contiguous chunk of the global expert range. Got "
+                    "--expert-placement-strategy "
+                    f"{self.parallel_config.expert_placement_strategy!r}."
+                )
+            # Enforced here rather than in set_splitting_ops_for_v1 so it
+            # holds for every compilation mode, and keyed on use_all2all so
+            # PCP/SP-only topologies are covered too: MoonEP dispatch/combine
+            # are eager-only and must not be captured.
+            if (
+                self.parallel_config.use_all2all
+                and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
+            ):
+                logger.info(
+                    "MoonEP: Disabling CUDA Graphs since the MoonEP "
+                    "integration is currently eager-only."
+                )
+                self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
+
         # Do this after all the updates to compilation_config.mode
         effective_dp_size = (
             self.parallel_config.data_parallel_size
@@ -1839,16 +2149,7 @@ def has_blocked_weights():
         # Resolve kv_offloading-derived connector name into kv_transfer_config
         # before the HMA check below, which inspects the connector class.
         self._post_init_kv_transfer_config()
-
-        if self.is_mm_encoder_only and self.cache_config.enable_prefix_caching:
-            # Such an instance publishes encoder embeddings and runs no language
-            # model, so it holds no KV cache for prefix caching to reuse and its
-            # coordinator would have no group to manage.
-            logger.info(
-                "Disabling prefix caching: this instance runs the "
-                "multi-modal encoder only."
-            )
-            self.cache_config.enable_prefix_caching = False
+        self._verify_aux_output_compatibility()
 
         # Hybrid KV cache manager (HMA) runtime rules:
         # - Explicit enable (--no-disable-kv-cache-manager): error if runtime
@@ -1923,6 +2224,16 @@ def has_blocked_weights():
             # Default to enable HMA if not explicitly disabled by user or logic above.
             self.scheduler_config.disable_hybrid_kv_cache_manager = False
 
+        if (
+            self.attention_config.hisparse_config is not None
+            and self.scheduler_config.disable_hybrid_kv_cache_manager
+        ):
+            raise ValueError(
+                "HiSparse requires the hybrid KV cache manager; remove "
+                "--disable-hybrid-kv-cache-manager or use connectors that "
+                "support HMA."
+            )
+
         if self.compilation_config.debug_dump_path:
             self.compilation_config.debug_dump_path = (
                 self.compilation_config.debug_dump_path.absolute().expanduser()
@@ -1974,8 +2285,7 @@ def update_sizes_for_sequence_parallelism(self, possible_sizes: list) -> list:
         ]
 
     def _set_max_num_scheduled_tokens(self):
-        """
-        In most cases, the scheduler may schedule a batch with as many tokens as the
+        """In most cases, the scheduler may schedule a batch with as many tokens as the
         worker is configured to handle.
         """
         if self.speculative_config is not None:
@@ -2013,8 +2323,7 @@ def _set_max_num_scheduled_tokens(self):
                 )
 
     def _set_cudagraph_sizes(self):
-        """
-        vLLM defines the default candidate list of batch sizes for CUDA graph
+        """VLLM defines the default candidate list of batch sizes for CUDA graph
         capture as:
 
         ```python
@@ -2060,8 +2369,8 @@ def _set_cudagraph_sizes(self):
             padded CUDA graph will be used.
             - If batch size > largest `cudagraph_capture_sizes`, cudagraph will
             not be used.
-        """
 
+        """
         if (
             self.model_config is not None
             and not self.model_config.enforce_eager
@@ -2071,9 +2380,9 @@ def _set_cudagraph_sizes(self):
             max_cudagraph_capture_size = (
                 self.compilation_config.max_cudagraph_capture_size
             )
-            # Decode sizes to cover, in tokens. Populated only when a request
-            # is more than one token wide and only when the default is computed
-            # here, so an explicit capture range is left exactly as configured.
+            # Decode sizes to cover, in tokens. Populated only when the default
+            # is computed here, so an explicit capture range is left exactly as
+            # configured.
             uniform_decode_sizes: list[int] = []
             if max_cudagraph_capture_size is None:
                 from vllm.platforms import current_platform
@@ -2164,6 +2473,8 @@ def request_counts(max_reqs: int) -> list[int]:
                             if n * query_len <= max_cudagraph_capture_size
                         }
                     )
+                elif max_num_seqs <= max_cudagraph_capture_size:
+                    uniform_decode_sizes = [max_num_seqs]
             max_num_tokens = self.scheduler_config.max_num_batched_tokens
             max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size)
 
@@ -2282,9 +2593,7 @@ def request_counts(max_reqs: int) -> list[int]:
         self.compilation_config.post_init_cudagraph_sizes()
 
     def _set_compile_ranges(self):
-        """
-        Set the compile ranges for the compilation config.
-        """
+        """Set the compile ranges for the compilation config."""
         compilation_config = self.compilation_config
         computed_compile_ranges_endpoints = []
 
@@ -2488,7 +2797,7 @@ def __str__(self):
             f"quantization={self.model_config.quantization}, "
             f"quantization_config={self.model_config.quantization_config}, "  # noqa
             f"enforce_eager={self.model_config.enforce_eager}, "
-            f"enable_return_routed_experts={self.model_config.enable_return_routed_experts}, "  # noqa
+            f"aux_output_config={self.aux_output_config!r}, "
             f"kv_cache_dtype={self.cache_config.cache_dtype}, "
             f"device_config={self.device_config.device}, "
             f"structured_outputs_config={self.structured_outputs_config!r}, "
@@ -2502,8 +2811,8 @@ def __str__(self):
             f"kernel_config={self.kernel_config!r}"
         )
 
-    def _resolve_allow_missing_mm_embeddings(self) -> None:
-        """Allow `*_embeds` tensors to be omitted on disaggregated consumers.
+    def _resolve_mm_embedding_inputs(self) -> None:
+        """Accept embedding inputs, tensor optional, on disaggregated consumers.
 
         An EC consumer loads embeddings from its connector. A KV consumer
         receives the prompt KV produced from those embeddings, so it does not
@@ -2524,11 +2833,21 @@ def _resolve_allow_missing_mm_embeddings(self) -> None:
         mm_config.allow_missing_mm_embeddings = (
             ec_config is not None and ec_config.is_ec_consumer
         ) or (kv_config is not None and kv_config.is_kv_consumer)
-        if mm_config.allow_missing_mm_embeddings:
+        if not mm_config.allow_missing_mm_embeddings:
+            return
+
+        if not mm_config.enable_mm_embeds:
+            # Allowing missing tensors still requires enabling embedding inputs
+            # for the frontend to accept metadata-only requests.
+            mm_config.enable_mm_embeds = True
             logger.info_once(
-                "EC/KV consumer: pre-computed-embedding inputs may "
-                "omit the embedding tensor."
+                "EC/KV consumer: accepting pre-computed-embedding inputs, "
+                "which this role is sent by definition."
             )
+        logger.info_once(
+            "EC/KV consumer: pre-computed-embedding inputs may "
+            "omit the embedding tensor."
+        )
 
     def _resolve_mm_encoder_only(self) -> None:
         """Enable encoder-only mode for a dedicated EC producer."""
@@ -2600,6 +2919,61 @@ def _resolve_mm_processor_device(self) -> None:
             device_type,
         )
 
+    def _resolve_mm_video_decode_device(self) -> None:
+        """Default video decoding to torchcodec GPU backend for EPD encoder-only
+        instance if the mm processor runs on CUDA.
+
+        The processor consumes the decoded frames on-device in that case, so
+        keeping the frames on the GPU skips the host round-trip through the
+        CPU media path. An explicit codec/backend choice in
+        `--media-io-kwargs` is left alone, and the default is skipped where
+        torchcodec (or its FFmpeg runtime) is unavailable.
+        """
+        if self.model_config is None or self.model_config.multimodal_config is None:
+            return
+        mm_config = self.model_config.multimodal_config
+
+        ec_config = self.ec_transfer_config
+        # An EC producer that is not also a consumer runs no forward pass and
+        # allocates no KV cache, so frontend accelerator work has the device to
+        # itself.
+        if ec_config is None or not ec_config.is_encode_only:
+            return
+
+        from vllm.platforms import current_platform
+
+        device_type = current_platform.device_type
+        if (
+            device_type != "cuda"
+            or mm_config.get_mm_processor_device_type() != device_type
+        ):
+            return
+
+        # User set video backend or device explicitly
+        video_kwargs = mm_config.media_io_kwargs.setdefault("video", {})
+        if "backend" in video_kwargs or "device" in video_kwargs:
+            return
+
+        from vllm.utils.import_utils import check_torchcodec_available
+
+        try:
+            check_torchcodec_available()
+        except (ImportError, RuntimeError):
+            # torchcodec is not installed, or is installed without a usable
+            # FFmpeg runtime (it raises rather than returning False).
+            logger.info_once(
+                "EPD encoder instance: keeping CPU video decoding because "
+                "torchcodec is not available (needs a CUDA build with FFmpeg)."
+            )
+            return
+
+        video_kwargs["backend"] = "torchcodec"
+        video_kwargs["device"] = device_type
+        logger.info_once(
+            "EPD encoder instance: decoding video with NVDEC (torchcodec device=%s).",
+            device_type,
+        )
+
     def _validate_mm_processor_device(self) -> None:
         """Hand the EC config to `MultiModalConfig`, which owns the rule."""
         model_config = self.model_config
@@ -2614,7 +2988,6 @@ def _validate_mm_processor_device(self) -> None:
     def _get_v2_model_runner_unsupported_features(self) -> list[str]:
         """Collect features not yet supported by the V2 model runner."""
         unsupported: list[str] = []
-        model_config = self.model_config
         speculative_config = self.speculative_config
 
         if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
@@ -2635,16 +3008,16 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
             unsupported.append("pipeline parallelism with external_launcher")
 
         if speculative_config is not None:
-            # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet
-            if speculative_config.method in ("ngram", "ngram_gpu"):
-                unsupported.append("ngram/ngram_gpu speculative decoding")
-            elif speculative_config.method not in (
-                "eagle",
-                "eagle3",
-                "mtp",
-                "dflash",
-                "dspark",
-                "extract_hidden_states",
+            if speculative_config.method in (
+                # https://github.com/vllm-project/vllm/pull/40704
+                "ngram",
+                "ngram_gpu",
+                # https://github.com/vllm-project/vllm/pull/43091
+                "draft_model",
+                "suffix",
+                "medusa",
+                "mlp_speculator",
+                "custom_class",
             ):
                 unsupported.append(f"speculative method '{speculative_config.method}'")
 
@@ -2663,21 +3036,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
         if self.parallel_config.enable_elastic_ep:
             unsupported.append("elastic expert parallelism")
 
-        has_logitsproc_plugins = False
-        if model_config is not None:
-            from importlib.metadata import entry_points
-
-            has_logitsproc_plugins = bool(entry_points(group="vllm.logits_processors"))
-
-        if model_config is not None and (
-            model_config.logits_processors or has_logitsproc_plugins
-        ):
-            unsupported.append("custom logits processors")
-
-        if self.cache_config.kv_sharing_fast_prefill:
-            # Will be added by https://github.com/vllm-project/vllm/pull/35045
-            unsupported.append("KV sharing fast prefill")
-
         if self.cache_config.mamba_cache_mode == "all":
             unsupported.append("mamba cache mode 'all'")
 
@@ -2690,6 +3048,20 @@ def _get_v1_model_runner_unsupported_features(self) -> list[str]:
         if self.parallel_config.prefill_context_parallel_size > 1:
             unsupported.append("prefill context parallel")
 
+        # Note(arpera):
+        # MRV1 + PP>1 + async sched + structured output
+        # does not work in vLLM. For more info see:
+        # https://github.com/vllm-project/vllm/issues/45014
+        # Since recently MRV1 has been deprecated then
+        # there was decided not to fix the issue but instead
+        # to disallow such configuration.
+        # At the same time MRV2 works fine in this case.
+        if (
+            self.parallel_config.pipeline_parallel_size > 1
+            and self.scheduler_config.async_scheduling
+        ):
+            unsupported.append("pipeline parallelism with async scheduling")
+
         # DSpark is implemented only by the V2 GPU model runner.
         if self.speculative_config:
             if self.speculative_config.method == "dspark":
@@ -2727,12 +3099,11 @@ def _validate_adaptive_verification(self) -> None:
                 "Adaptive verification is not currently compatible with LoRA"
             )
 
-        if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE:
-            # The draft budget divides by step costs profiled from captured
-            # cudagraphs; eager execution captures none.
+        if not self.compilation_config.cudagraph_mode.has_full_cudagraphs():
             raise ValueError(
-                "Adaptive verification is not currently compatible with "
-                "enforce_eager/cudagraph_mode=none"
+                "Adaptive verification requires full CUDA graphs. Use cudagraph "
+                "mode FULL, FULL_DECODE_ONLY, or FULL_AND_PIECEWISE, or disable "
+                "adaptive verification."
             )
 
         if self.parallel_config.pipeline_parallel_size > 1:
@@ -2834,13 +3205,52 @@ def _get_dbo_unsupported_features(self) -> list[str]:
             unsupported.append("dual batch overlap with multimodal models")
         if model_config is not None and model_config.is_hybrid:
             unsupported.append("dual batch overlap with hybrid models")
-        if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
-            unsupported.append("dual batch overlap with CUDA graphs")
+        if self.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE:
+            # DBO captures FULL graphs only.
+            unsupported.append("dual batch overlap with PIECEWISE CUDA graphs")
         if self.is_mm_encoder_only:
             unsupported.append("dual batch overlap with encoder only models")
 
         return unsupported
 
+    def _validate_profiler_config(self) -> None:
+        if self.profiler_config.profiler != "proton":
+            return
+
+        from vllm.platforms import current_platform
+
+        if not current_platform.is_cuda():
+            raise ValueError("The Proton profiler currently supports NVIDIA CUDA only")
+        if (
+            self.profiler_config.proton_graph_attribution
+            and not self.use_v2_model_runner
+        ):
+            raise ValueError(
+                "Proton CUDA graph attribution requires the V2 model runner."
+            )
+
+        has_cuda_graphs = (
+            self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
+            or (
+                self.compilation_config.cudagraph_mm_encoder
+                and not (self.model_config and self.model_config.enforce_eager)
+            )
+        )
+        if not has_cuda_graphs:
+            return
+        mode = self.profiler_config.proton_mode
+        if mode and mode.split(":", 1)[0].lower() == "pcsampling":
+            raise ValueError(
+                "Proton PC sampling requires CUDA graphs to be disabled. "
+                "Use --enforce-eager."
+            )
+        if not self.profiler_config.proton_graph_attribution:
+            raise ValueError(
+                "Proton profiling with CUDA graphs requires "
+                "proton_graph_attribution=True to capture replayed kernels. "
+                "Enable attribution or use --enforce-eager."
+            )
+
     def _validate_v2_model_runner(self) -> None:
         """Check for features not yet supported by the V2 model runner."""
         if not HAS_TRITON:
@@ -2887,6 +3297,8 @@ def adjust_dcp_kv_cache_interleave_size(
             "NixlConnector"
         ):
             return
+        if not self.parallel_config._allow_auto_resolve_cp_interleave_size:
+            return
 
         # Get the kernel block_size, but don't use resolve_kv_cache_block_size to avoid
         # scaling by dcp_size (we need the local block_size here).
@@ -3050,8 +3462,7 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig":
 def set_current_vllm_config(
     vllm_config: VllmConfig, check_compile=False, prefix: str | None = None
 ):
-    """
-    Temporarily set the current vLLM config.
+    """Temporarily set the current vLLM config.
     Used during model initialization.
     We save the current vLLM config in a global variable,
     so that all modules can access it, e.g. custom ops
@@ -3132,15 +3543,14 @@ def get_layers_from_vllm_config(
     layer_type: type[T],
     layer_names: Iterable[str] | None = None,
 ) -> dict[str, T]:
-    """
-    Get layers from the vLLM config.
+    """Get layers from the vLLM config.
 
     Args:
         vllm_config: The vLLM config.
         layer_type: The type of the layer to get.
         layer_names: The names of the layers to get. If None, return all layers.
-    """
 
+    """
     forward_context = vllm_config.compilation_config.static_forward_context
     if layer_names is None:
         layer_names = forward_context.keys()
diff --git a/vllm/config/watermarking.py b/vllm/config/watermarking.py
index 2247dd8248ea..8701c7b72e4d 100644
--- a/vllm/config/watermarking.py
+++ b/vllm/config/watermarking.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import hashlib
 from typing import Literal
 
 from pydantic import Field, model_validator
@@ -11,12 +12,16 @@
 
 logger = init_logger(__name__)
 
-WatermarkingAlgorithm = Literal["gumbel"]
+WatermarkingAlgorithm = Literal["gumbel", "dual_key_gumbel"]
 WatermarkPRFName = Literal["philox"]
+WatermarkContextScope = Literal["none", "single_turn", "all"]
 
-_SPECULATIVE_DECODING_SUPPORT: dict[WatermarkingAlgorithm, bool] = {
-    "gumbel": False,
-}
+_MIN_RECOMMENDED_DEDUP_HISTORY = 1024
+
+
+def derive_watermark_key(key: int, domain: bytes) -> int:
+    digest = hashlib.sha256(domain + key.to_bytes(8, "big")).digest()
+    return int.from_bytes(digest[:8], "big")
 
 
 @config
@@ -27,23 +32,49 @@ class WatermarkConfig:
     """Secret key used to watermark generated text."""
     algorithm: WatermarkingAlgorithm = "gumbel"
     """Algorithm used to watermark generated text."""
+    alpha: float = Field(default=0.1, ge=0, le=1)
+    """Probability of selecting key B for dual-key watermarking."""
     context_width: int = Field(default=4, ge=1)
-    """Number of prior output tokens used by the watermark PRF."""
+    """Number of prior tokens used by the watermark PRF."""
+    deduplicate_contexts: WatermarkContextScope = "single_turn"
+    """Which history is searched for a repeated context before a token is
+    sampled; a repeated context is sampled without the watermark. `none`
+    disables the search. `single_turn` (default) searches this request's
+    generated tokens. `all` also searches the prompt and samples the first
+    `context_width` generated tokens without watermarking."""
+    deduplicate_contexts_max_history: int | None = Field(default=8192, ge=1)
+    """Number of most recent history positions searched (default 8192), or
+    `None` for the whole scope. Each position is compared over the
+    `context_width` tokens before it. Ignored when `deduplicate_contexts` is
+    `none`."""
     prf: WatermarkPRFName = "philox"
     """Pseudorandom function used by the watermarking algorithm."""
+    allow_target_only_watermarking: bool = False
+    """Allow speculative decoding without watermarking draft tokens."""
+
+    @property
+    def supports_speculative_decoding(self) -> bool:
+        return self.algorithm == "dual_key_gumbel"
 
     @model_validator(mode="after")
-    def validate_key(self) -> Self:
+    def validate_watermark_settings(self) -> Self:
         if self.key > 2**64 - 1:
             raise ValueError("philox keys must fit in 64 bits")
-        if self.algorithm == "gumbel":
+        history_is_too_short = (
+            self.deduplicate_contexts_max_history is not None
+            and self.deduplicate_contexts_max_history < _MIN_RECOMMENDED_DEDUP_HISTORY
+        )
+        if self.algorithm in ("gumbel", "dual_key_gumbel") and (
+            self.deduplicate_contexts == "none" or history_is_too_short
+        ):
             logger.warning_once(
-                "Single-key Gumbel-max watermarking may increase the frequency of "
-                "degenerate generations, including repetition loops.",
+                "Gumbel-max watermarking with context deduplication "
+                "disabled or limited to fewer than "
+                f"{_MIN_RECOMMENDED_DEDUP_HISTORY} positions may increase the "
+                "frequency of degenerate generations, including repetition loops. "
+                "Use deduplicate_contexts='single_turn' or 'all' with "
+                "deduplicate_contexts_max_history at least "
+                f"{_MIN_RECOMMENDED_DEDUP_HISTORY} or null to mitigate this.",
                 scope="global",
             )
         return self
-
-    @property
-    def supports_speculative_decoding(self) -> bool:
-        return _SPECULATIVE_DECODING_SUPPORT[self.algorithm]
diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py
index 23721aac440c..033533a5c9a0 100644
--- a/vllm/cute_utils/__init__.py
+++ b/vllm/cute_utils/__init__.py
@@ -22,6 +22,14 @@
     torch.float32: Float32,
 }
 
+
+def torch_to_cute_dtype(dtype: torch.dtype):
+    try:
+        return _TORCH_TO_CUTE_DTYPE[dtype]
+    except KeyError as exc:
+        raise TypeError(f"Unsupported CuTe dtype {dtype}") from exc
+
+
 _CUTE_TO_PTX_DTYPE = {
     BFloat16: "bf16",
     Float16: "f16",
diff --git a/vllm/cute_utils/_tcgen05.py b/vllm/cute_utils/_tcgen05.py
index 325c30498c76..4ab75cc0cc3b 100644
--- a/vllm/cute_utils/_tcgen05.py
+++ b/vllm/cute_utils/_tcgen05.py
@@ -78,6 +78,10 @@ def make_sdesc_128B_swizzle(LBO: int):
     return Uint64((LBO >> 4 << 16) | (SBO >> 4 << 32) | (1 << 46) | (2 << 61))
 
 
+def make_mxfp8_idesc(MMA_M: int, MMA_N: int):
+    return Uint32((MMA_N >> 3 << 17) | (1 << 23) | (MMA_M >> 7 << 27))
+
+
 @dsl_user_op
 def mma_f16(
     d_tmem,
@@ -130,6 +134,73 @@ def mma_ts_f16(
         )
 
 
+@dsl_user_op
+def mma_mxfp8(
+    d_tmem,
+    a_desc,
+    b_desc,
+    idesc,
+    sfa_tmem,
+    sfb_tmem,
+    enable_input_d,
+    cta_group: int = 1,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    with cute.arch.elect_one():
+        nvvm.tcgen05_mma_block_scale(
+            nvvm.Tcgen05MMAKind.MXF8F6F4,
+            NVVM_CTA_GROUP_MAP[cta_group],
+            _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
+            Uint64(a_desc).ir_value(loc=loc, ip=ip),
+            Uint64(b_desc).ir_value(loc=loc, ip=ip),
+            Int32(idesc).ir_value(loc=loc, ip=ip),
+            Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
+            _make_tmem_llvm_ptr(sfa_tmem, loc=loc, ip=ip),
+            _make_tmem_llvm_ptr(sfb_tmem, loc=loc, ip=ip),
+            scale_vec_size=nvvm.Tcgen05MMAScaleVecSize.X1,
+            loc=loc,
+            ip=ip,
+        )
+
+
+@dsl_user_op
+def cp(
+    tmem,
+    sdesc,
+    shape: str,
+    mcast: str,
+    cta_group: int = 1,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    SHAPE_MAP = {
+        "128x256b": nvvm.Tcgen05CpShape.SHAPE_128x256b,
+        "4x256b": nvvm.Tcgen05CpShape.SHAPE_4x256b,
+        "128x128b": nvvm.Tcgen05CpShape.SHAPE_128x128b,
+        "64x128b": nvvm.Tcgen05CpShape.SHAPE_64x128b,
+        "32x128b": nvvm.Tcgen05CpShape.SHAPE_32x128b,
+    }
+    MCAST_MAP = {
+        "none": nvvm.Tcgen05CpMulticast.NONE,
+        "warpx2::02_13": nvvm.Tcgen05CpMulticast.WARPX2_02_13,
+        "warpx2::01_23": nvvm.Tcgen05CpMulticast.WARPX2_01_23,
+        "warpx4": nvvm.Tcgen05CpMulticast.WARPX4,
+    }
+    with cute.arch.elect_one():
+        nvvm.tcgen05_cp(
+            SHAPE_MAP[shape],
+            _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip),
+            Uint64(sdesc).ir_value(loc=loc, ip=ip),
+            group=NVVM_CTA_GROUP_MAP[cta_group],
+            multicast=MCAST_MAP[mcast],
+            loc=loc,
+            ip=ip,
+        )
+
+
 @dsl_user_op
 def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None):
     mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip)
diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py
index a16f1b01486e..63d4a1f4e0ec 100644
--- a/vllm/device_allocator/cumem.py
+++ b/vllm/device_allocator/cumem.py
@@ -80,8 +80,7 @@ def use_memory_pool_with_allocator(
 
 
 class CuMemAllocator:
-    """
-    A singleton class that manages a memory pool for CUDA tensors.
+    """A singleton class that manages a memory pool for CUDA tensors.
     The memory in this pool can be offloaded or discarded when the
     allocator sleeps.
 
@@ -109,8 +108,7 @@ class CuMemAllocator:
 
     @staticmethod
     def get_instance() -> "CuMemAllocator":
-        """
-        CuMemAllocator is a singleton class.
+        """CuMemAllocator is a singleton class.
         We cannot call the constructor directly.
         Call this method to get the instance.
         """
@@ -182,9 +180,9 @@ def close(self) -> None:
         self.release_pools()
 
     def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
+        """Internal method to store the allocation data
+        when memory is allocated in the memory pool.
         """
-        Internal method to store the allocation data
-        when memory is allocated in the memory pool."""
         py_d_mem = allocation_handle[2]
         self.pointer_to_data[py_d_mem] = AllocationData(
             allocation_handle, self.current_tag
@@ -198,9 +196,9 @@ def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
         return
 
     def _python_free_callback(self, ptr: int) -> HandleType:
+        """Internal method to look up the allocation data
+        when memory is freed in the memory pool.
         """
-        Internal method to look up the allocation data
-        when memory is freed in the memory pool."""
         data = self.pointer_to_data.pop(ptr)
         if data.cpu_backup_tensor is not None:
             data.cpu_backup_tensor = None
@@ -227,14 +225,14 @@ def _python_free_callback(self, ptr: int) -> HandleType:
         return data.handle
 
     def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
-        """
-        Put the allocator in sleep mode.
+        """Put the allocator in sleep mode.
         All data in the memory allocation with the specified tag will be
         offloaded to CPU memory, and others will be discarded.
 
         Args:
             offload_tags: The tags of the memory allocation that will be
                 offloaded. The rest of the memory allocation will be discarded.
+
         """
         if offload_tags is None:
             # by default, allocated tensors are offloaded
@@ -325,8 +323,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None:
             )
 
     def wake_up(self, tags: list[str] | None = None) -> None:
-        """
-        Wake up the allocator from sleep mode.
+        """Wake up the allocator from sleep mode.
         All data that is previously offloaded will be loaded back to GPU
         memory, and the rest of the data will have empty memory.
 
@@ -334,6 +331,7 @@ def wake_up(self, tags: list[str] | None = None) -> None:
             tags: The tags of the memory allocation that will be loaded
                 back to GPU memory. If None, all memory allocation will be loaded
                 back to GPU memory.
+
         """
         gc.collect()
         torch.accelerator.empty_cache()
@@ -357,14 +355,14 @@ def wake_up(self, tags: list[str] | None = None) -> None:
 
     @contextmanager
     def use_memory_pool(self, tag: str | None = None):
-        """
-        A context manager to use the memory pool.
+        """A context manager to use the memory pool.
         All memory allocation created inside the context will be allocated
         in the memory pool, and has the specified tag.
 
         Args:
             tag: The tag of the memory allocation. If None, the default tag
                 will be used.
+
         """
         if tag is None:
             tag = CuMemAllocator.default_tag
@@ -416,9 +414,7 @@ def use_memory_pool(self, tag: str | None = None):
                 torch.cuda.memory._set_allocator_settings("expandable_segments:True")
 
     def get_current_usage(self) -> int:
-        """
-        Get the total number of bytes allocated in the memory pool.
-        """
+        """Get the total number of bytes allocated in the memory pool."""
         sum_bytes: int = 0
         for ptr, data in self.pointer_to_data.items():
             handle = data.handle
diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py
index 45a635d61389..023af04275c4 100644
--- a/vllm/device_allocator/sleep_mode_backend.py
+++ b/vllm/device_allocator/sleep_mode_backend.py
@@ -73,6 +73,13 @@ def state(self) -> SleepModeState:
         (suspended) engine from a healthy-serving one (see RFC #34303)."""
         return self._state
 
+    def discard(self, tags: tuple[str, ...]) -> None:
+        """Discard selected tagged allocations without preserving contents."""
+        raise RuntimeError(
+            f"Sleep-mode backend {type(self).__name__!r} does not support "
+            "selective discard."
+        )
+
     # -- Capability introspection (no instance required) --
 
     @classmethod
@@ -132,6 +139,13 @@ def resume(self, tags: list[str] | None = None) -> None:
         allocator.wake_up(tags)
         self._state = "RUNNING"
 
+    def discard(self, tags: tuple[str, ...]) -> None:
+        from vllm.device_allocator import get_mem_allocator_instance
+
+        allocator = get_mem_allocator_instance()
+        allocator.discard(tags)
+        self._state = "SUSPENDED"
+
     @classmethod
     def preserves_communicators(cls) -> bool:
         # Comm identity/topology survive memory suspension; no reinit needed.
diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py
index 146ea965db07..17fa69ee8411 100644
--- a/vllm/device_allocator/xpumem.py
+++ b/vllm/device_allocator/xpumem.py
@@ -116,6 +116,7 @@ class XpuMemAllocator:
     Sleep will offload selected payloads to CPU or discard and unmap XPU
     physical memory. Wake-up remaps physical memory back to the same
     reserved virtual address and restores payload.
+
     """
 
     instance: "XpuMemAllocator | None" = None
diff --git a/vllm/distributed/aux_output_connector/__init__.py b/vllm/distributed/aux_output_connector/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/distributed/aux_output_connector/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/distributed/aux_output_connector/connector.py b/vllm/distributed/aux_output_connector/connector.py
new file mode 100644
index 000000000000..bd8cbf4d183d
--- /dev/null
+++ b/vllm/distributed/aux_output_connector/connector.py
@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Scheduler-side control plane for execution auxiliary outputs."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.request import Request
+
+
+@dataclass
+class PackedBlockHashes:
+    """Contiguous, self-contained block hashes for scheduler-worker IPC."""
+
+    data: bytes
+    item_size: int
+
+    def __iter__(self) -> Iterator[bytes]:
+        for start in range(0, len(self.data), self.item_size):
+            yield self.data[start : start + self.item_size]
+
+
+@dataclass
+class AuxOutputConnectorMetadata:
+    generation: int
+    requests: dict[str, int]
+    block_hashes: dict[str, PackedBlockHashes]
+    finished_requests: tuple[str, ...]
+
+
+@dataclass
+class AuxRequestOutput:
+    token_start: int
+    rows: np.ndarray
+
+
+class AuxOutputSchedulerConnector:
+    """Build worker metadata without owning auxiliary output payloads or stores."""
+
+    def __init__(self) -> None:
+        # Number of hashes already sent to the worker for each active request.
+        self._sent_hash_counts: dict[str, int] = {}
+        # Terminal events are delivered with the next connector metadata.
+        self._finished_requests: dict[str, PackedBlockHashes | None] = {}
+        self._generation = 0
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+        requests: dict[str, Request],
+    ) -> AuxOutputConnectorMetadata:
+        """Build one step's incremental worker metadata."""
+        scheduled_requests: dict[str, int] = {}
+        block_hashes_by_request: dict[str, PackedBlockHashes] = {}
+        for request_id in scheduler_output.num_scheduled_tokens:
+            num_sent = self._sent_hash_counts.setdefault(request_id, 0)
+            request = requests[request_id]
+            packed = self._pack_new_hashes(request.block_hashes, num_sent)
+            if packed is not None:
+                block_hashes_by_request[request_id] = packed
+            self._sent_hash_counts[request_id] = len(request.block_hashes)
+            assert request.sampling_params is not None
+            scheduled_requests[request_id] = max(
+                request.sampling_params.routed_experts_prompt_start,
+                0 if request.num_output_tokens == 0 else request.num_tokens - 1,
+            )
+        # A settled token can complete a hash block after the next async schedule
+        # was built. Send a hash-only update if the request was not rescheduled.
+        if len(self._sent_hash_counts) != len(scheduled_requests):
+            for request_id, num_sent in self._sent_hash_counts.items():
+                if request_id in scheduled_requests:
+                    continue
+                request = requests[request_id]
+                packed = self._pack_new_hashes(request.block_hashes, num_sent)
+                if packed is None:
+                    continue
+                block_hashes_by_request[request_id] = packed
+                self._sent_hash_counts[request_id] = len(request.block_hashes)
+        # Sending transfers ownership of these one-shot events.
+        finished_requests = tuple(self._finished_requests)
+        block_hashes_by_request.update(
+            (request_id, block_hashes)
+            for request_id, block_hashes in self._finished_requests.items()
+            if block_hashes is not None
+        )
+        self._finished_requests = {}
+        return AuxOutputConnectorMetadata(
+            self._generation,
+            scheduled_requests,
+            block_hashes_by_request,
+            finished_requests,
+        )
+
+    def take_output(
+        self,
+        request: Request,
+        output: dict[str, AuxRequestOutput] | None,
+    ) -> np.ndarray | None:
+        """Return the accepted R3 rows for one scheduled request."""
+        request_id = request.request_id
+        assert output is not None and request_id in output, (
+            f"auxiliary output worker output is missing {request_id}"
+        )
+        request_output = output[request_id]
+        token_end = request.num_tokens - 1
+        local_end = token_end - request_output.token_start
+        if local_end < 0:
+            assert not request.is_finished(), (
+                "finished auxiliary output output has no accepted token range: "
+                f"request={request_id}, token_end={token_end}, "
+                f"output_start={request_output.token_start}, "
+                "output_end="
+                f"{request_output.token_start + len(request_output.rows)}"
+            )
+            return None
+        assert local_end <= len(request_output.rows), (
+            "auxiliary output worker output has an invalid token range: "
+            f"request={request_id}, token_end={token_end}, "
+            f"output_start={request_output.token_start}, "
+            f"output_end={request_output.token_start + len(request_output.rows)}"
+        )
+        return request_output.rows[:local_end]
+
+    def request_finished(self, request: Request) -> None:
+        """Queue a request's terminal event and final block hashes."""
+        request_id = request.request_id
+        num_sent = self._sent_hash_counts.pop(request_id, None)
+        if num_sent is None:
+            return
+        # The next metadata delivers this terminal event to the worker.
+        self._finished_requests[request_id] = self._pack_new_hashes(
+            request.block_hashes, num_sent
+        )
+
+    @staticmethod
+    def _pack_new_hashes(
+        block_hashes: Sequence[bytes], num_sent: int
+    ) -> PackedBlockHashes | None:
+        assert num_sent <= len(block_hashes), "KV block-hash history shrank"
+        new_hashes = block_hashes[num_sent:]
+        if not new_hashes:
+            return None
+        return PackedBlockHashes(b"".join(new_hashes), len(new_hashes[0]))
+
+    def reset(self) -> None:
+        """Start a new auxiliary output namespace after a prefix-cache reset."""
+        # The worker drops temporary state on generation changes; resend hashes.
+        self._sent_hash_counts.clear()
+        self._finished_requests.clear()
+        self._generation += 1
diff --git a/vllm/distributed/aux_output_connector/routed_experts.py b/vllm/distributed/aux_output_connector/routed_experts.py
new file mode 100644
index 000000000000..1c89b97215a6
--- /dev/null
+++ b/vllm/distributed/aux_output_connector/routed_experts.py
@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Routed-experts auxiliary output keys, publication, and materialization."""
+
+from __future__ import annotations
+
+from collections.abc import Hashable, Iterable, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from vllm.distributed.aux_output_connector.store import (
+    BackgroundBlockObjectStore,
+    BlockObject,
+    BlockObjectStore,
+)
+
+
+@dataclass
+class _RequestTail:
+    slot: int
+    block_start: int
+    length: int = 0
+
+
+class RoutedExpertsBuffer:
+    """Own incomplete and unkeyed full routed-experts blocks."""
+
+    def __init__(
+        self,
+        dtype: np.dtype[Any],
+        shape_per_token: tuple[int, ...],
+        block_size: int,
+        max_num_seqs: int,
+        max_num_batched_tokens: int,
+        max_concurrent_batches: int,
+    ) -> None:
+        self.dtype = dtype
+        self.shape_per_token = shape_per_token
+        self.block_size = block_size
+        max_step_blocks = (
+            max_num_batched_tokens + block_size - 1
+        ) // block_size + max_num_seqs
+        # In-flight full blocks plus one incomplete tail per active request.
+        max_blocks = max_concurrent_batches * max_step_blocks + max_num_seqs
+        self._rows: np.ndarray = np.empty(
+            (max_blocks, block_size, *shape_per_token), dtype=dtype
+        )
+        self._free_slots = list(range(len(self._rows) - 1, -1, -1))
+        self._owned_slots: dict[int, int] = {}
+        self._requests: dict[Hashable, _RequestTail] = {}
+
+    def _tail(self, request_id: Hashable, block_start: int) -> _RequestTail:
+        tail = self._requests.get(request_id)
+        if tail is not None:
+            assert tail.block_start == block_start, (
+                "auxiliary output capture skipped an incomplete block: "
+                f"request={request_id}, expected={tail.block_start}, "
+                f"actual={block_start}"
+            )
+            return tail
+        assert self._free_slots, "auxiliary output block pool is exhausted"
+        tail = _RequestTail(self._free_slots.pop(), block_start)
+        self._requests[request_id] = tail
+        return tail
+
+    def capture(
+        self, request_id: Hashable, token_start: int, rows: np.ndarray
+    ) -> list[tuple[int, np.ndarray]]:
+        """Stage rows and return completed blocks without retaining them."""
+        rows = np.asarray(rows)
+        assert rows.shape[1:] == self.shape_per_token and rows.dtype == self.dtype, (
+            "routed-experts capture profile changed"
+        )
+        if token_start < 0:
+            raise ValueError("auxiliary output token start must be non-negative")
+
+        completed: list[tuple[int, np.ndarray]] = []
+        offset = 0
+        while offset < len(rows):
+            position = token_start + offset
+            block_start = position // self.block_size * self.block_size
+            local_start = position - block_start
+
+            # Full aligned input blocks do not need tail staging.
+            if (
+                request_id not in self._requests
+                and local_start == 0
+                and len(rows) - offset >= self.block_size
+            ):
+                completed.append((block_start, rows[offset : offset + self.block_size]))
+                offset += self.block_size
+                continue
+
+            tail = self._tail(request_id, block_start)
+            assert local_start == tail.length, (
+                "auxiliary output capture is not contiguous: "
+                f"request={request_id}, expected={block_start + tail.length}, "
+                f"actual={position}"
+            )
+            count = min(self.block_size - local_start, len(rows) - offset)
+            if count == 1:
+                self._rows[tail.slot, local_start] = rows[offset]
+            else:
+                self._rows[tail.slot, local_start : local_start + count] = rows[
+                    offset : offset + count
+                ]
+            tail.length += count
+            offset += count
+            if tail.length == self.block_size:
+                block = self._rows[tail.slot]
+                del self._requests[request_id]
+                self._owned_slots[id(block)] = tail.slot
+                completed.append((block_start, block))
+
+        return completed
+
+    def read(
+        self, request_id: Hashable, token_start: int, token_end: int
+    ) -> np.ndarray:
+        tail = self._requests.get(request_id)
+        assert tail is not None, (
+            f"auxiliary output buffer is missing request {request_id}"
+        )
+        local_start = token_start - tail.block_start
+        local_end = token_end - tail.block_start
+        assert 0 <= local_start < local_end <= tail.length, (
+            "auxiliary output range is unavailable: "
+            f"request={request_id}, range=[{token_start}, {token_end}), "
+            f"available=[{tail.block_start}, {tail.block_start + tail.length})"
+        )
+        return self._rows[tail.slot, local_start:local_end].copy()
+
+    def retain_block(self, rows: np.ndarray) -> np.ndarray:
+        """Retain one unkeyed block after the current capture call."""
+        if id(rows) in self._owned_slots:
+            return rows
+        assert self._free_slots, "auxiliary output block pool is exhausted"
+        slot = self._free_slots.pop()
+        retained = self._rows[slot]
+        retained[...] = rows
+        self._owned_slots[id(retained)] = slot
+        return retained
+
+    def release_block(self, rows: np.ndarray) -> None:
+        slot = self._owned_slots.pop(id(rows), None)
+        if slot is not None:
+            self._free_slots.append(slot)
+
+    def discard(self, request_id: Hashable) -> None:
+        tail = self._requests.pop(request_id, None)
+        if tail is not None:
+            self._free_slots.append(tail.slot)
+
+    def reset(self) -> None:
+        self._requests.clear()
+        self._owned_slots.clear()
+        self._free_slots[:] = range(len(self._rows) - 1, -1, -1)
+
+
+def routed_experts_keys(
+    block_hashes: Iterable[bytes], aux_output_namespace: str
+) -> list[str]:
+    prefix = f"vllm-artifact/{aux_output_namespace}/"
+    return [prefix + block_hash.hex() for block_hash in block_hashes]
+
+
+def materialize_routed_experts(
+    store: BackgroundBlockObjectStore | BlockObjectStore,
+    aux_output_keys: list[str],
+    *,
+    shape_per_token: tuple[int, ...],
+    dtype: np.dtype[Any],
+) -> np.ndarray:
+    if not aux_output_keys:
+        raise ValueError("routed-experts auxiliary output key list must not be empty")
+    payload = store.get_concatenated(aux_output_keys)
+    return np.frombuffer(payload, dtype=dtype).reshape((-1, *shape_per_token))
+
+
+def publish_routed_experts(
+    store: BackgroundBlockObjectStore | BlockObjectStore,
+    *,
+    batches: Sequence[tuple[Sequence[str], list[tuple[int, np.ndarray]]]],
+    block_size: int,
+    retain_keys: Sequence[str] = (),
+    release_keys: Sequence[str] = (),
+) -> None:
+    """Publish immutable full R3 blocks."""
+    objects = []
+    for aux_output_keys, blocks in batches:
+        for block_start, array in blocks:
+            if block_start < 0 or block_start % block_size:
+                raise ValueError(
+                    "auxiliary output block start is not hash-block aligned"
+                )
+            block_index = block_start // block_size
+            if block_index >= len(aux_output_keys):
+                raise ValueError(
+                    "auxiliary output block has no corresponding key: "
+                    f"start={block_start}, index={block_index}, "
+                    f"keys={len(aux_output_keys)}, block_size={block_size}"
+                )
+            if len(array) != block_size:
+                raise ValueError(
+                    "auxiliary output block length does not match hash block size"
+                )
+            objects.append(
+                BlockObject(
+                    key=aux_output_keys[block_index],
+                    payload=array.tobytes(order="C"),
+                )
+            )
+    store.put(
+        objects,
+        retain_keys=retain_keys,
+        release_keys=release_keys,
+    )
diff --git a/vllm/distributed/aux_output_connector/store.py b/vllm/distributed/aux_output_connector/store.py
new file mode 100644
index 000000000000..2b364288a588
--- /dev/null
+++ b/vllm/distributed/aux_output_connector/store.py
@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""In-process auxiliary-output storage."""
+
+from __future__ import annotations
+
+import mmap
+import queue
+import threading
+from collections import OrderedDict
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class BlockObject:
+    """One immutable auxiliary output object."""
+
+    key: str
+    payload: bytes
+
+
+class BlockObjectStoreError(RuntimeError):
+    """AuxOutput storage or retrieval failed."""
+
+
+class BackgroundBlockObjectStore:
+    """Serialize store mutations on a background thread."""
+
+    def __init__(self, store: BlockObjectStore, *, max_pending_batches: int) -> None:
+        self._store = store
+        self._queue: queue.Queue[
+            tuple[list[BlockObject], tuple[str, ...], tuple[str, ...]] | None
+        ] = queue.Queue(maxsize=max_pending_batches)
+        self._error: BaseException | None = None
+        self._closed = False
+        self._thread = threading.Thread(
+            target=self._run,
+            daemon=True,
+            name="vllm-aux-output-writer",
+        )
+        self._thread.start()
+
+    def _run(self) -> None:
+        while True:
+            update = self._queue.get()
+            try:
+                if update is None:
+                    return
+                if self._error is None:
+                    objects, retain_keys, release_keys = update
+                    self._store.put(
+                        objects,
+                        retain_keys=retain_keys,
+                        release_keys=release_keys,
+                    )
+            except BaseException as error:
+                self._error = error
+            finally:
+                self._queue.task_done()
+
+    def _raise_if_failed(self) -> None:
+        if self._error is not None:
+            raise BlockObjectStoreError(
+                "auxiliary output publication failed"
+            ) from self._error
+
+    def put(
+        self,
+        objects: list[BlockObject],
+        *,
+        retain_keys: Iterable[str] = (),
+        release_keys: Iterable[str] = (),
+    ) -> None:
+        retains = tuple(retain_keys)
+        releases = tuple(release_keys)
+        if not objects and not retains and not releases:
+            return
+        if self._closed:
+            raise RuntimeError("auxiliary output store is closed")
+        self._queue.put((objects, retains, releases))
+        self._raise_if_failed()
+
+    def get_concatenated(self, keys: list[str]) -> bytes:
+        self._queue.join()
+        self._raise_if_failed()
+        return self._store.get_concatenated(keys)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self._queue.join()
+        self._queue.put(None)
+        self._thread.join()
+        self._store.close()
+        self._raise_if_failed()
+
+
+class BlockObjectStore:
+    """Single-owner bounded store that fails closed after an eviction."""
+
+    _UNALLOCATED_SLOT = -1
+
+    def __init__(
+        self,
+        *,
+        max_bytes: int,
+        object_nbytes: int,
+    ) -> None:
+        if object_nbytes <= 0:
+            raise ValueError("auxiliary output object size must be positive")
+        if max_bytes < object_nbytes:
+            raise ValueError("auxiliary output store must fit at least one object")
+        self.object_nbytes = object_nbytes
+        self.num_slots = max_bytes // object_nbytes
+        self._lru: OrderedDict[str, int] = OrderedDict()
+        self._references: dict[str, int] = {}
+        self._free_slots: list[int] = []
+        self._next_slot = 0
+        self._arena: mmap.mmap | None = mmap.mmap(
+            -1, self.num_slots * self.object_nbytes
+        )
+
+    def _evict_to_fit(self, protected: set[str]) -> None:
+        excess = len(self._lru) - self.num_slots
+        if excess <= 0:
+            return
+        victims = []
+        for key in self._lru:
+            if key not in self._references and key not in protected:
+                victims.append(key)
+                if len(victims) == excess:
+                    break
+        if len(victims) != excess:
+            raise BlockObjectStoreError(
+                "auxiliary output store cannot retain the requested batch: "
+                f"limit={self.num_slots} objects"
+            )
+        for victim in victims:
+            slot = self._lru.pop(victim)
+            if slot != self._UNALLOCATED_SLOT:
+                self._free_slots.append(slot)
+
+    def _allocate_slot(self) -> int:
+        if self._free_slots:
+            return self._free_slots.pop()
+        slot = self._next_slot
+        self._next_slot += 1
+        return slot
+
+    def put(
+        self,
+        objects: list[BlockObject],
+        *,
+        retain_keys: Iterable[str] = (),
+        release_keys: Iterable[str] = (),
+    ) -> None:
+        arena = self._arena
+        if arena is None:
+            raise RuntimeError("auxiliary output store is closed")
+        unique = {obj.key: obj for obj in objects}
+        self._retain(retain_keys)
+        terminal_order = self._release(release_keys)
+        for key in unique:
+            self._lru.setdefault(key, self._UNALLOCATED_SLOT)
+            self._lru.move_to_end(key)
+        for key in terminal_order:
+            if key in self._lru:
+                self._lru.move_to_end(key)
+        try:
+            self._evict_to_fit(set(unique) - set(terminal_order))
+        except BlockObjectStoreError:
+            for key in unique:
+                if self._lru.get(key) == self._UNALLOCATED_SLOT:
+                    del self._lru[key]
+            raise
+        for object_id, obj in unique.items():
+            slot = self._lru.get(object_id)
+            if slot is None:
+                continue
+            if slot != self._UNALLOCATED_SLOT:
+                continue
+            slot = self._allocate_slot()
+            offset = slot * self.object_nbytes
+            arena[offset : offset + self.object_nbytes] = obj.payload
+            self._lru[object_id] = slot
+
+    def _retain(self, keys: Iterable[str]) -> None:
+        for key in keys:
+            references = self._references.get(key, 0)
+            self._references[key] = references + 1
+
+    def _release(self, keys: Iterable[str]) -> list[str]:
+        terminal_order = []
+        for key in keys:
+            references = self._references[key] - 1
+            if references:
+                self._references[key] = references
+                continue
+            del self._references[key]
+            terminal_order.append(key)
+        return terminal_order
+
+    def get_concatenated(self, keys: list[str]) -> bytes:
+        arena_obj = self._arena
+        if arena_obj is None:
+            raise RuntimeError("auxiliary output store is closed")
+        try:
+            entries = [self._lru[key] for key in keys]
+        except KeyError as error:
+            raise BlockObjectStoreError(
+                "auxiliary output object does not exist; the object may have been "
+                f"evicted (used={len(self._lru)}, "
+                f"limit={self.num_slots} objects). Increase "
+                "aux_output_config.max_bytes when a KV cache hit requires it."
+            ) from error
+        arena = memoryview(arena_obj)
+        try:
+            payload = b"".join(
+                arena[slot * self.object_nbytes : (slot + 1) * self.object_nbytes]
+                for slot in entries
+            )
+        finally:
+            arena.release()
+        for key in keys:
+            self._lru.move_to_end(key)
+        return payload
+
+    def close(self) -> None:
+        arena = self._arena
+        if arena is not None:
+            self._arena = None
+            arena.close()
diff --git a/vllm/distributed/aux_output_connector/worker.py b/vllm/distributed/aux_output_connector/worker.py
new file mode 100644
index 000000000000..7959e15ea884
--- /dev/null
+++ b/vllm/distributed/aux_output_connector/worker.py
@@ -0,0 +1,412 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Worker-side auxiliary-output data plane."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, field
+from threading import Lock
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.aux_output_connector.connector import (
+    AuxOutputConnectorMetadata,
+    AuxRequestOutput,
+)
+from vllm.distributed.aux_output_connector.routed_experts import (
+    RoutedExpertsBuffer,
+    materialize_routed_experts,
+    publish_routed_experts,
+    routed_experts_keys,
+)
+from vllm.distributed.aux_output_connector.store import (
+    BackgroundBlockObjectStore,
+    BlockObjectStore,
+)
+from vllm.distributed.parallel_state import get_tp_group
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+    bind_routed_experts_capturer,
+)
+from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
+
+if TYPE_CHECKING:
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+
+
+@dataclass
+class _WorkerRequestState:
+    aux_output_keys: list[str] = field(default_factory=list)
+    pending_blocks: list[tuple[int, np.ndarray]] = field(default_factory=list)
+    capture_cursor: int | None = None
+    scheduled_cursor: int = 0
+    emit_cursor: int = 0
+    pending_outputs: int = 0
+    # Terminal event received; teardown waits for in-flight step outputs.
+    finished: bool = False
+
+
+@dataclass
+class PendingAuxOutput:
+    """Own one step's R3 snapshot until its asynchronous copy is consumed."""
+
+    connector: AuxOutputWorkerConnector
+    request_ids: list[str]
+    token_starts: np.ndarray
+    query_start_loc: np.ndarray
+    # GPU snapshot taken on the main stream.
+    routed_experts_gpu: torch.Tensor
+    # Filled by enqueue_cpu_copy on the output copy stream.
+    routed_experts: np.ndarray | None = None
+    num_sampled: np.ndarray | None = None
+    num_rejected: np.ndarray | None = None
+
+    def enqueue_cpu_copy(
+        self, num_sampled: np.ndarray, num_rejected: np.ndarray
+    ) -> None:
+        """Enqueue asynchronous D2H copies; call on the output copy stream."""
+        self.num_sampled = num_sampled
+        self.num_rejected = num_rejected
+        self.routed_experts = self.routed_experts_gpu.to(
+            "cpu", non_blocking=True
+        ).numpy()
+
+    def process_output(self) -> dict[str, AuxRequestOutput]:
+        return self.connector.process_output(self)
+
+
+class AuxOutputWorkerConnector:
+    """Own capture, request tails, and backend resources on the output worker."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        model: torch.nn.Module,
+        kv_cache_config: KVCacheConfig,
+    ) -> None:
+        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        capturer = RoutedExpertsCapturer(
+            max_num_batched_tokens=max_num_batched_tokens, vllm_config=vllm_config
+        )
+        bind_routed_experts_capturer(model, capturer)
+        self._capturer = capturer
+        self._store: BackgroundBlockObjectStore | None = None
+        self._buffer: RoutedExpertsBuffer | None = None
+        self._requests: dict[str, _WorkerRequestState] = {}
+        self._generation = 0
+        self._step_metadata: AuxOutputConnectorMetadata | None = None
+        # Steps whose asynchronous copy has not been consumed yet. The engine
+        # consumes step outputs in order, so at most max_concurrent_batches
+        # are outstanding; guarded by _lock against the output thread.
+        self._pending_outputs: list[PendingAuxOutput] = []
+        self._lock = Lock()
+        self._max_concurrent_batches = vllm_config.max_concurrent_batches
+        # Every TP rank participates in capture collectives, but only the
+        # executor output rank owns the auxiliary output data plane.
+        if not get_tp_group().is_first_rank:
+            return
+
+        shape_per_token = self._capturer.shape_per_token
+        dtype: np.dtype[Any] = np.dtype(self._capturer.output_dtype_name)
+        scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
+            kv_cache_config, vllm_config
+        )
+        hashes_per_kv_block = scheduler_block_size // hash_block_size
+        block_nbytes = hash_block_size * int(np.prod(shape_per_token)) * dtype.itemsize
+        max_bytes = vllm_config.aux_output_config.max_bytes
+        if max_bytes is None:
+            max_bytes = kv_cache_config.num_blocks * hashes_per_kv_block * block_nbytes
+        self._store = BackgroundBlockObjectStore(
+            BlockObjectStore(max_bytes=max_bytes, object_nbytes=block_nbytes),
+            max_pending_batches=2 * vllm_config.scheduler_config.max_num_seqs,
+        )
+        self._buffer = RoutedExpertsBuffer(
+            dtype,
+            shape_per_token,
+            hash_block_size,
+            vllm_config.scheduler_config.max_num_seqs,
+            max_num_batched_tokens,
+            vllm_config.max_concurrent_batches,
+        )
+
+    def prepare_output(self, input_batch: InputBatch) -> PendingAuxOutput | None:
+        """Snapshot one step's R3 tensor for asynchronous CPU transfer."""
+        if self._buffer is None or self._step_metadata is None:
+            return None
+
+        request_ids = list(input_batch.req_ids)
+        query_start_loc = input_batch.query_start_loc_np[: len(request_ids) + 1]
+        num_rows = int(query_start_loc[-1])
+        pending_output = PendingAuxOutput(
+            connector=self,
+            request_ids=request_ids,
+            token_starts=input_batch.num_computed_tokens_np,
+            query_start_loc=query_start_loc,
+            routed_experts_gpu=self._capturer.snapshot_routing_data(num_rows),
+        )
+        with self._lock:
+            assert len(self._pending_outputs) < self._max_concurrent_batches, (
+                "auxiliary output step outputs are not consumed in order"
+            )
+            self._pending_outputs.append(pending_output)
+            for request_id in pending_output.request_ids:
+                self._requests[request_id].pending_outputs += 1
+        return pending_output
+
+    def process_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutput]:
+        """Commit one consumed R3 snapshot and build request outputs.
+
+        Request teardown deferred by begin_step is completed here once no
+        in-flight step covers the request.
+        """
+        with self._lock:
+            teardown = []
+            release_keys: list[str] = []
+            try:
+                outputs = self._commit_output(pending)
+            finally:
+                self._pending_outputs.remove(pending)
+                for request_id in pending.request_ids:
+                    state = self._requests[request_id]
+                    state.pending_outputs -= 1
+                    if state.pending_outputs == 0 and state.finished:
+                        teardown.append(request_id)
+                        release_keys.extend(reversed(state.aux_output_keys))
+            if teardown:
+                self._publish_blocks([], release_keys=release_keys)
+                self._teardown(teardown)
+            return outputs
+
+    def _commit_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutput]:
+        buffer = self._buffer
+        store = self._store
+        assert buffer is not None and store is not None
+        block_size = buffer.block_size
+
+        routed_experts = pending.routed_experts
+        num_sampled = pending.num_sampled
+        num_rejected = pending.num_rejected
+        assert (
+            routed_experts is not None
+            and num_sampled is not None
+            and num_rejected is not None
+        ), "auxiliary output CPU copy was not enqueued"
+
+        # Publish the whole batch before materializing any consumer output.
+        materialize_outputs: list[tuple[str, int, int]] = []
+        block_batches = []
+        outputs: dict[str, AuxRequestOutput] = {}
+
+        # Use the ModelRunner's actual batch boundaries rather than rebuilding them.
+        for request_id, token_start, start, end, sampled, rejected in zip(
+            pending.request_ids,
+            pending.token_starts,
+            pending.query_start_loc[:-1],
+            pending.query_start_loc[1:],
+            num_sampled,
+            num_rejected,
+            strict=True,
+        ):
+            request_num_tokens = end - start
+            assert request_num_tokens > 0, (
+                "auxiliary output request token count must be positive"
+            )
+            state = self._requests[request_id]
+
+            # Capture precedes speculative acceptance, so discard the rejected
+            # suffix. Batch boundaries still span the full executed range.
+            rejected = int(rejected)
+            assert 0 <= rejected <= request_num_tokens, (
+                "auxiliary output rejected-token count is invalid"
+            )
+            rows = routed_experts[start : end - rejected]
+
+            capture_start = token_start
+            capture_cursor = state.capture_cursor
+            if capture_cursor is None:
+                capture_cursor = capture_start
+
+            assert capture_start >= capture_cursor, (
+                "auxiliary output capture moved backwards"
+            )
+            if capture_start > capture_cursor:
+                # Reattach after an optimistically scheduled suffix was rejected.
+                assert capture_cursor < state.scheduled_cursor, (
+                    "auxiliary output capture has an unbacked token gap"
+                )
+                capture_start = capture_cursor
+
+            emit_start = state.emit_cursor
+            # Complete blocks without keys remain pending until a hash update.
+            completed = buffer.capture(request_id, capture_start, rows)
+            token_end = capture_start + len(rows)
+            state.capture_cursor = token_end
+            state.scheduled_cursor = token_start + request_num_tokens
+            block_batches.append((state, completed))
+
+            if sampled > 0 and emit_start <= token_end:
+                if emit_start >= capture_start:
+                    outputs[request_id] = AuxRequestOutput(
+                        emit_start, rows[emit_start - capture_start :]
+                    )
+                    state.emit_cursor = token_end
+                else:
+                    materialize_outputs.append((request_id, emit_start, token_end))
+
+        # A consumer may reuse a block produced earlier in the same batch.
+        self._publish_blocks(block_batches)
+
+        for request_id, emit_start, token_end in materialize_outputs:
+            state = self._requests[request_id]
+            stored_end = (
+                min(token_end // block_size, len(state.aux_output_keys)) * block_size
+            )
+            if emit_start < stored_end:
+                first_block = emit_start // block_size
+                stored = materialize_routed_experts(
+                    store,
+                    state.aux_output_keys[first_block : stored_end // block_size],
+                    shape_per_token=buffer.shape_per_token,
+                    dtype=buffer.dtype,
+                )
+                local_start = emit_start % block_size
+                rows = stored[local_start : local_start + stored_end - emit_start]
+                if stored_end < token_end:
+                    rows = np.concatenate(
+                        (rows, buffer.read(request_id, stored_end, token_end))
+                    )
+            else:
+                rows = buffer.read(request_id, emit_start, token_end)
+            outputs[request_id] = AuxRequestOutput(emit_start, rows)
+            state.emit_cursor = token_end
+        return outputs
+
+    def _publish_blocks(
+        self,
+        batches: list[tuple[_WorkerRequestState, list[tuple[int, np.ndarray]]]],
+        retain_keys: Sequence[str] = (),
+        release_keys: Sequence[str] = (),
+    ) -> None:
+        store = self._store
+        buffer = self._buffer
+        assert store is not None and buffer is not None
+        ready_batches = []
+        for state, completed in batches:
+            blocks = state.pending_blocks + completed
+            keyed_end = len(state.aux_output_keys) * buffer.block_size
+            ready = [(start, rows) for start, rows in blocks if start < keyed_end]
+            state.pending_blocks = [
+                (start, buffer.retain_block(rows))
+                for start, rows in blocks
+                if start >= keyed_end
+            ]
+            if ready:
+                ready_batches.append((state.aux_output_keys, ready))
+        if ready_batches or retain_keys or release_keys:
+            publish_routed_experts(
+                store,
+                batches=ready_batches,
+                block_size=buffer.block_size,
+                retain_keys=retain_keys,
+                release_keys=release_keys,
+            )
+        for _, blocks in ready_batches:
+            for _, rows in blocks:
+                buffer.release_block(rows)
+
+    def _teardown(self, request_ids: Iterable[str]) -> None:
+        """Drop finished requests' temporary state after publishing key releases."""
+        buffer = self._buffer
+        assert buffer is not None
+        for request_id in request_ids:
+            state = self._requests.pop(request_id)
+            for _, rows in state.pending_blocks:
+                buffer.release_block(rows)
+            buffer.discard(request_id)
+
+    def begin_step(self, metadata: AuxOutputConnectorMetadata | None) -> None:
+        """Apply one scheduler step's request and block-hash updates.
+
+        While an unconsumed step covers a finished request, its teardown is
+        deferred to that step's process_output.
+        """
+        self._step_metadata = metadata
+        if self._buffer is None or metadata is None:
+            return
+        assert not metadata.requests.keys() & metadata.finished_requests, (
+            "auxiliary output request cannot run and finish in one step"
+        )
+        assert metadata.generation >= self._generation, (
+            "auxiliary output metadata generation moved backwards"
+        )
+        with self._lock:
+            release_keys: list[str] = []
+            if metadata.generation > self._generation:
+                # A generation change follows a prefix-cache reset, which the
+                # scheduler only performs once all model output is consumed.
+                assert not self._pending_outputs, (
+                    "auxiliary output generation changed with output in flight"
+                )
+                release_keys.extend(
+                    key
+                    for state in self._requests.values()
+                    for key in reversed(state.aux_output_keys)
+                )
+                self._buffer.reset()
+                self._requests.clear()
+                self._generation = metadata.generation
+            for request_id, emit_start in metadata.requests.items():
+                state = self._requests.setdefault(
+                    request_id, _WorkerRequestState(emit_cursor=emit_start)
+                )
+                if state.pending_outputs == 0:
+                    # In-flight steps leave the worker cursor behind the
+                    # scheduler's optimistic view; only check settled requests.
+                    assert emit_start <= state.emit_cursor, (
+                        "auxiliary output Scheduler emit cursor moved ahead"
+                    )
+            block_batches: list[
+                tuple[_WorkerRequestState, list[tuple[int, np.ndarray]]]
+            ] = []
+            retained_keys: list[str] = []
+            for request_id, block_hashes in metadata.block_hashes.items():
+                state = self._requests[request_id]
+                keys = routed_experts_keys(block_hashes, str(self._generation))
+                state.aux_output_keys.extend(keys)
+                retained_keys.extend(keys)
+                block_batches.append((state, []))
+            finished_now: list[str] = []
+            for request_id in metadata.finished_requests:
+                state = self._requests[request_id]
+                state.finished = True
+                if state.pending_outputs:
+                    continue
+                finished_now.append(request_id)
+                release_keys.extend(reversed(state.aux_output_keys))
+            self._publish_blocks(block_batches, retained_keys, release_keys)
+            self._teardown(finished_now)
+
+    def close(self) -> None:
+        with self._lock:
+            self._pending_outputs.clear()
+            if self._store is not None:
+                try:
+                    self._store.close()
+                finally:
+                    self._store = None
+
+
+def get_aux_output_connector(
+    model: torch.nn.Module, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
+) -> AuxOutputWorkerConnector:
+    return AuxOutputWorkerConnector(
+        model=model,
+        kv_cache_config=kv_cache_config,
+        vllm_config=vllm_config,
+    )
diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py
index be2c901281bd..6fbffcaaf73d 100644
--- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py
+++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py
@@ -22,9 +22,7 @@ class AiterCustomAllreduce:
 
     @classmethod
     def effective_max_size(cls) -> int:
-        """
-        Max input byte size eligible for AITER custom allreduce.
-        """
+        """Max input byte size eligible for AITER custom allreduce."""
         return cls.MAX_SIZE // 2
 
     def __init__(
@@ -56,6 +54,20 @@ def should_custom_ar(self, inp: torch.Tensor) -> bool:
     def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None:
         return self._impl.custom_all_reduce(inp)
 
+    def should_custom_ag(self, inp: torch.Tensor) -> bool:
+        return self._impl.should_custom_ag(inp)
+
+    def should_custom_rs(self, inp: torch.Tensor, dim: int) -> bool:
+        return self._impl.should_custom_rs(inp, dim)
+
+    def custom_all_gather(self, inp: torch.Tensor, dim: int = 0) -> torch.Tensor | None:
+        return self._impl.custom_all_gather(inp, dim=dim)
+
+    def custom_reduce_scatter(
+        self, inp: torch.Tensor, out: torch.Tensor, dim: int = 0
+    ) -> torch.Tensor | None:
+        return self._impl.custom_reduce_scatter(inp, out, dim=dim)
+
     def use_1stage_fused_ar_rms(self, inp: torch.Tensor) -> bool:
         """Whether AITER's fused allreduce+RMSNorm runs as its one-stage kernel.
 
diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py
index eb187a99f4be..21a1f00da4a6 100644
--- a/vllm/distributed/device_communicators/all2all.py
+++ b/vllm/distributed/device_communicators/all2all.py
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
 import threading
+from collections.abc import Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Any
 
@@ -18,7 +21,12 @@
     has_flashinfer_nvlink_two_sided,
 )
 from vllm.utils.func_utils import supports_kw
-from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori
+from vllm.utils.import_utils import (
+    check_moonep_system_support,
+    has_deep_ep,
+    has_deep_ep_v2,
+    has_mori,
+)
 
 from .base_device_communicator import All2AllManagerBase, Cache
 
@@ -42,18 +50,23 @@
 
 
 class AgRsAll2AllManager(All2AllManagerBase):
-    """
-    An implementation of all2all communication based on
+    """An implementation of all2all communication based on
     all-gather (dispatch) and reduce-scatter (combine).
     """
 
     def __init__(self, cpu_group, tcp_store_group=None):
         super().__init__(cpu_group, tcp_store_group)
+        self.use_ep = get_current_vllm_config().parallel_config.enable_expert_parallel
 
     def _get_comm_group(self, is_sequence_parallel: bool) -> Any:
         if is_sequence_parallel:
             return get_ep_group()
         if self.dp_world_size > 1:
+            if self.use_ep and get_pcp_group().world_size > 1:
+                assert self.tp_group.world_size == 1, (
+                    "DP+PCP with TP>1 requires sequence-parallel MoE inputs"
+                )
+                return get_ep_group()
             return get_dp_group()
         return get_pcp_group()
 
@@ -65,6 +78,7 @@ def _get_sizes(self, num_local_tokens: int, comm_group: Any) -> list[int]:
         assert dp_metadata is not None
         sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
         assert sizes is not None
+        assert len(sizes) == comm_group.world_size
         return sizes
 
     def dispatch_router_logits(
@@ -77,9 +91,7 @@ def dispatch_router_logits(
         tuple[torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Gather hidden_states and router_logits from all dp ranks.
-        """
+        """Gather hidden_states and router_logits from all dp ranks."""
         dist_group = self._get_comm_group(is_sequence_parallel)
         sizes = self._get_sizes(hidden_states.shape[0], dist_group)
         assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
@@ -109,9 +121,7 @@ def dispatch(
         tuple[torch.Tensor, torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Gather hidden_states and router_logits from all dp ranks.
-        """
+        """Gather hidden_states and router_logits from all dp ranks."""
         dist_group = self._get_comm_group(is_sequence_parallel)
         sizes = self._get_sizes(hidden_states.shape[0], dist_group)
         assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
@@ -138,9 +148,7 @@ def dispatch(
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
     ) -> torch.Tensor:
-        """
-        Reduce-scatter hidden_states across all dp ranks.
-        """
+        """Reduce-scatter hidden_states across all dp ranks."""
         dist_group = self._get_comm_group(is_sequence_parallel)
         sizes = self._get_sizes(
             hidden_states.shape[0] // dist_group.world_size,
@@ -154,9 +162,7 @@ def destroy(self):
 
 
 class DeepEPAll2AllManagerBase(All2AllManagerBase):
-    """
-    All2All communication based on DeepEP High-Throughput kernels.
-    """
+    """All2All communication based on DeepEP High-Throughput kernels."""
 
     def __init__(self, cpu_group, tcp_store_group=None):
         assert has_deep_ep(), (
@@ -208,9 +214,7 @@ def destroy(self):
 
 
 class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
-    """
-    All2All communication based on DeepEP High-Throughput kernels.
-    """
+    """All2All communication based on DeepEP High-Throughput kernels."""
 
     def __init__(self, cpu_group, tcp_store_group=None):
         super().__init__(cpu_group, tcp_store_group)
@@ -269,9 +273,7 @@ def set_num_sms(self, num_sms: int):
 
 
 class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
-    """
-    All2All communication based on DeepEP Low-Latency kernels.
-    """
+    """All2All communication based on DeepEP Low-Latency kernels."""
 
     _buffer: Any = None
     _mask: torch.Tensor | None = None
@@ -291,8 +293,7 @@ def _make_all2all_kwargs(
         num_global_experts: int,
         num_local_experts: int,
     ) -> dict[Any, Any]:
-        """
-        max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank
+        """max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank
           can dispatch all the ranks must hold the same value.
         token_hidden_size: the hidden dimension of each token.
         num_ep_ranks: the number of EP group ranks.
@@ -328,8 +329,7 @@ def _make_all2all_kwargs(
         return kwargs
 
     def get_handle(self, kwargs):
-        """
-        The kwargs for DeepEPLLAll2AllManager is dictated by
+        """The kwargs for DeepEPLLAll2AllManager is dictated by
         _make_all2all_kwargs.
         """
         import deep_ep  # type: ignore[import-not-found]
@@ -377,13 +377,13 @@ def clean_buffers(self) -> None:
 @dataclass
 class _NixlEPBufferState:
     buffer: Any
+    max_num_ep_ranks: int
     connected_ep_size: int
     active_ep_size: int
 
 
 class NixlEPAll2AllManager(All2AllManagerBase):
-    """
-    All2All communication based on NIXL EP kernels.
+    """All2All communication based on NIXL EP kernels.
     This backend supports elastic EP with dynamic rank connection/disconnection.
     """
 
@@ -402,7 +402,10 @@ def __init__(self, cpu_group, tcp_store_group=None):
         super().__init__(cpu_group, tcp_store_group)
         self.support_fault_tolerance = True
 
-        self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS
+    @property
+    def max_num_ep_ranks(self) -> int:
+        assert NixlEPAll2AllManager._buffer is not None
+        return NixlEPAll2AllManager._buffer.max_num_ep_ranks
 
     def _init_buffer(
         self,
@@ -412,11 +415,19 @@ def _init_buffer(
     ) -> None:
         from nixl_ep import Buffer  # type: ignore[import-not-found]
 
-        max_num_global_experts = self.max_num_ep_ranks * num_experts_per_rank
+        parallel_config = get_current_vllm_config().parallel_config
+        max_num_ep_ranks = (
+            self.world_size
+            * parallel_config.elastic_ep_max_dp_size
+            // parallel_config.data_parallel_size
+            if parallel_config.enable_elastic_ep
+            else self.world_size
+        )
+        max_num_global_experts = max_num_ep_ranks * num_experts_per_rank
         num_rdma_bytes = Buffer.get_rdma_size_hint(
             num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
             hidden=token_hidden_size,
-            num_ranks=self.max_num_ep_ranks,
+            num_ranks=max_num_ep_ranks,
             num_experts=max_num_global_experts,
         )
         assert NixlEPAll2AllManager._buffer is None, (
@@ -427,7 +438,7 @@ def _init_buffer(
             tcp_store_group=self.tcp_store_group.store,
         )
         buffer.update_memory_buffers(
-            num_ranks=self.max_num_ep_ranks,
+            num_ranks=max_num_ep_ranks,
             num_experts_per_rank=num_experts_per_rank,
             num_rdma_bytes=num_rdma_bytes,
         )
@@ -435,6 +446,7 @@ def _init_buffer(
         buffer.connect_ranks(ranks_to_connect)
         NixlEPAll2AllManager._buffer = _NixlEPBufferState(
             buffer=buffer,
+            max_num_ep_ranks=max_num_ep_ranks,
             connected_ep_size=self.world_size,
             active_ep_size=self.world_size,
         )
@@ -474,6 +486,7 @@ def _unmask_connected_ranks(self, target_ep_size: int) -> None:
 
         for rank in range(state.active_ep_size, target_ep_size):
             state.buffer.update_mask_buffer(rank, mask=False)
+        torch.accelerator.synchronize()
         state.active_ep_size = target_ep_size
 
     def _stage_ep_size(self) -> None:
@@ -486,8 +499,12 @@ def _stage_ep_size(self) -> None:
         if target_ep_size > state.connected_ep_size:
             self._connect_to_ep_size(target_ep_size, make_active=False)
 
-    def commit_staged_state(self) -> None:
-        """Commit staged NIXL EP state to the active communication set."""
+    def stage_ep_size(self) -> None:
+        with NixlEPAll2AllManager._lock:
+            self._stage_ep_size()
+
+    def commit_ep_size(self) -> None:
+        """Commit the staged EP size to the active communication set."""
         with NixlEPAll2AllManager._lock:
             assert NixlEPAll2AllManager._buffer is not None
             state = NixlEPAll2AllManager._buffer
@@ -500,31 +517,31 @@ def commit_staged_state(self) -> None:
 
             self._unmask_connected_ranks(target_ep_size)
 
-    def _ensure_ep_size(self, *, stage: bool) -> None:
-        if stage:
-            self._stage_ep_size()
-        else:
-            self.commit_staged_state()
+    @contextmanager
+    def mask_remote_ranks(self) -> Iterator[None]:
+        peers = [rank for rank in range(self.world_size) if rank != self.rank]
+        assert NixlEPAll2AllManager._buffer is not None
+        buffer = NixlEPAll2AllManager._buffer.buffer
+        for rank in peers:
+            buffer.update_mask_buffer(rank, mask=True)
+        torch.accelerator.synchronize()
+        try:
+            yield
+        finally:
+            for rank in peers:
+                buffer.update_mask_buffer(rank, mask=False)
+            torch.accelerator.synchronize()
 
     def get_handle(self, kwargs):
         with NixlEPAll2AllManager._lock:
-            stage = bool(kwargs.get("stage", False))
             state = NixlEPAll2AllManager._buffer
             if state is None:
-                assert not stage, (
-                    "NIXL EP staged initialization requires an existing buffer"
-                )
                 max_num_tokens_per_dp_rank = kwargs["max_num_tokens_per_dp_rank"]
-                num_experts_per_rank = (
-                    kwargs["num_global_experts"] // kwargs["num_ep_ranks"]
-                )
                 self._init_buffer(
                     max_num_tokens_per_dp_rank=max_num_tokens_per_dp_rank,
                     token_hidden_size=kwargs["token_hidden_size"],
-                    num_experts_per_rank=num_experts_per_rank,
+                    num_experts_per_rank=kwargs["num_local_experts"],
                 )
-            else:
-                self._ensure_ep_size(stage=stage)
 
             assert NixlEPAll2AllManager._buffer is not None
             handle = NixlEPAll2AllManager._buffer.buffer
@@ -564,7 +581,7 @@ def query_active_mask(self) -> torch.Tensor:
         assert state is not None
         if NixlEPAll2AllManager._mask is None:
             NixlEPAll2AllManager._mask = torch.zeros(
-                self.max_num_ep_ranks, device="cuda", dtype=torch.int32
+                state.max_num_ep_ranks, device="cuda", dtype=torch.int32
             )
         state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask)
         return NixlEPAll2AllManager._mask[: state.active_ep_size]
@@ -592,9 +609,7 @@ def clean_buffers(self) -> None:
 
 
 class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
-    """
-    All2All communication based on flashinfer all2allv/two-sided NVLink kernels.
-    """
+    """All2All communication based on flashinfer all2allv/two-sided NVLink kernels."""
 
     # This type lint could be removed after all of the work in
     # https://github.com/vllm-project/vllm/issues/26533 done.
@@ -620,7 +635,7 @@ def initialize(
         rank: int,
         gpus_per_node: int,
     ):
-        """Initialize workspace"""
+        """Initialize workspace."""
         if self.initialized:
             return
 
@@ -661,7 +676,7 @@ def initialize(
         )
 
     def ensure_alltoall_workspace_initialized(self):
-        """Ensure workspace is initialized"""
+        """Ensure workspace is initialized."""
         if not has_flashinfer_nvlink_two_sided():
             return False
 
@@ -680,7 +695,7 @@ def get_handle(self, kwargs):
         return self
 
     def cleanup(self):
-        """Clean up workspace"""
+        """Clean up workspace."""
         if (
             self.initialized
             and self.workspace_tensor is not None
@@ -699,8 +714,7 @@ def cleanup(self):
 
 
 class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
-    """
-    All2All communication based on FlashInfer's MoeAlltoAll/One-sided NVLink kernel.
+    """All2All communication based on FlashInfer's MoeAlltoAll/One-sided NVLink kernel.
     This is a newer kernel from trtllm that should perform better than the kernel
     used by flashinfer_nvlink_two_sided.
     """
@@ -1003,8 +1017,7 @@ def get_handle(self, kwargs):
 
 
 class DeepEPV2All2AllManager(All2AllManagerBase):
-    """
-    All2All communication based on DeepEP v2 ElasticBuffer (unified API).
+    """All2All communication based on DeepEP v2 ElasticBuffer (unified API).
     Uses NCCL Gin backend with analytical SM calculation.
     """
 
@@ -1048,6 +1061,9 @@ def _check_gin_support(self, group) -> None:
         # missing GIN support.
         probe = torch.zeros(1, device="cuda")
         torch.distributed.all_reduce(probe, group=group)
+        # DeepEPv2 respects EP_DISABLE_GIN, so skip the GIN requirement check.
+        if os.environ.get("EP_DISABLE_GIN", "0") != "0":
+            return
 
         gin_type = query_nccl_gin_type(group)
         if gin_type is None:
@@ -1090,3 +1106,71 @@ def destroy(self):
             for _, handle in self.handle_cache._cache.items():
                 handle.destroy()
             self.handle_cache._cache.clear()
+
+
+class MoonEPAll2AllManager(All2AllManagerBase):
+    """All2All communication based on MoonEP
+    (https://github.com/MoonshotAI/MoonEP).
+
+    MoonEP keeps token loads perfectly balanced across EP ranks by planning
+    a small number of dynamically redundant experts online and prefetching
+    their weights before expert compute. Every rank receives exactly
+    S x K token slots regardless of router skew, so all communication and
+    compute shapes are static.
+
+    Requires NVLink symmetric-memory / multicast capable topologies
+    (single node NVSwitch, e.g. H100/H200/B200/GB300 class).
+    """
+
+    def __init__(self, cpu_group, tcp_store_group=None, device_group=None):
+        check_moonep_system_support()
+        super().__init__(cpu_group, tcp_store_group)
+        self._device_group = device_group
+        self.handle_cache = Cache()
+
+    def _make_buffer_kwargs(
+        self,
+        max_num_tokens_per_dp_rank: int,
+        token_hidden_size: int,
+        num_topk: int,
+        num_global_experts: int,
+        num_prefetch_slots: int,
+        token_padding: int,
+        num_sms: int,
+    ) -> dict:
+        return dict(
+            S=max_num_tokens_per_dp_rank,
+            H=token_hidden_size,
+            K=num_topk,
+            E=num_global_experts,
+            num_ep_ranks=self.world_size,
+            num_sms=num_sms,
+            token_padding=token_padding,
+            B=num_prefetch_slots,
+            group=self._device_group
+            if self._device_group is not None
+            else self.cpu_group,
+            explicitly_destroy=True,
+        )
+
+    def get_handle(self, kwargs):
+        from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
+            MoonEPBufferPool,
+        )
+
+        buffer_kwargs = self._make_buffer_kwargs(**kwargs)
+        logger.debug("MoonEP all2all args %s", buffer_kwargs)
+
+        def make_pool(**kw):
+            return MoonEPBufferPool(kw, max_tokens_per_rank=kw["S"])
+
+        handle: MoonEPBufferPool = self.handle_cache.get_or_create(
+            buffer_kwargs, make_pool
+        )
+        return handle
+
+    def destroy(self):
+        with self.handle_cache._lock:
+            for _, handle in self.handle_cache._cache.items():
+                handle.destroy()
+            self.handle_cache._cache.clear()
diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py
index 51c28788cc01..a8298a459fa4 100644
--- a/vllm/distributed/device_communicators/all_reduce_utils.py
+++ b/vllm/distributed/device_communicators/all_reduce_utils.py
@@ -119,8 +119,7 @@
 
 
 def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) -> bool:
-    """
-    Determine if NCCL symmetric memory allreduce should be used.
+    """Determine if NCCL symmetric memory allreduce should be used.
 
     Based on H100 and GB200 benchmarks, NCCL symm_mem is preferred for:
     - Small tensors (≤16K): Lower latency than custom_AR
@@ -344,7 +343,6 @@ def can_actually_p2p(
 
 def gpu_p2p_access_check(src: int, tgt: int) -> bool:
     """Check if GPU src can access GPU tgt."""
-
     # if the cache variable is already calculated,
     # read from the cache instead of checking it again
     global _gpu_p2p_access_cache
diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py
index 83ae25ff4437..f70bf0921518 100644
--- a/vllm/distributed/device_communicators/base_device_communicator.py
+++ b/vllm/distributed/device_communicators/base_device_communicator.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import threading
+from collections.abc import Iterator
+from contextlib import contextmanager
 from weakref import WeakValueDictionary
 
 import torch
@@ -112,6 +114,7 @@ def query_active_mask(self) -> torch.Tensor:
         Returns:
             An int32 device tensor where 0 marks a live rank and 1 marks a
             masked (dead/unreachable) rank.
+
         """
         raise NotImplementedError
 
@@ -157,10 +160,19 @@ def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = Fals
     def destroy(self):
         pass
 
+    def stage_ep_size(self) -> None:
+        pass
+
+    def commit_ep_size(self) -> None:
+        pass
+
+    @contextmanager
+    def mask_remote_ranks(self) -> Iterator[None]:
+        yield
+
 
 class DeviceCommunicatorBase:
-    """
-    Base class for device-specific communicator.
+    """Base class for device-specific communicator.
     It can use the `cpu_group` to initialize the communicator.
     If the device has PyTorch integration (PyTorch can recognize its
     communication backend), the `device_group` will also be given.
@@ -306,8 +318,7 @@ def reduce_scatterv(
     def gather(
         self, input_: torch.Tensor, dst: int = 0, dim: int = -1
     ) -> torch.Tensor | None:
-        """
-        NOTE: We assume that the input tensor is on the same device across
+        """NOTE: We assume that the input tensor is on the same device across
         all the ranks.
         NOTE: `dst` is the local rank of the destination rank.
         """
@@ -335,7 +346,7 @@ def gather(
         return output_tensor
 
     def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
-        """Sends a tensor to the destination rank in a blocking way"""
+        """Sends a tensor to the destination rank in a blocking way."""
         """NOTE: `dst` is the local rank of the destination rank."""
         if dst is None:
             dst = (self.rank_in_group + 1) % self.world_size
@@ -373,8 +384,7 @@ def dispatch_router_logits(
         tuple[torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Dispatch the hidden states and router logits to the appropriate device.
+        """Dispatch the hidden states and router logits to the appropriate device.
         This is a no-op in the base class.
         """
         if extra_tensors is not None:
@@ -392,8 +402,7 @@ def dispatch(
         tuple[torch.Tensor, torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Dispatch the hidden states and topk weights/ids to the appropriate device.
+        """Dispatch the hidden states and topk weights/ids to the appropriate device.
         This is a no-op in the base class.
         """
         if extra_tensors is not None:
@@ -403,8 +412,7 @@ def dispatch(
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
     ) -> torch.Tensor:
-        """
-        Combine the hidden states and router logits from the appropriate device.
+        """Combine the hidden states and router logits from the appropriate device.
         This is a no-op in the base class.
         """
         return hidden_states
diff --git a/vllm/distributed/device_communicators/cpu_communicator.py b/vllm/distributed/device_communicators/cpu_communicator.py
index 8ea12d9255a9..c190c9ded691 100644
--- a/vllm/distributed/device_communicators/cpu_communicator.py
+++ b/vllm/distributed/device_communicators/cpu_communicator.py
@@ -69,8 +69,7 @@ def __init__(
             logger.info("Using allgather_reducescatter all2all manager.")
 
     def _all_group_ranks_share_shm_group_name(self) -> bool:
-        """
-        CPUSHM requires all ranks in this group to agree on one SHM group name.
+        """CPUSHM requires all ranks in this group to agree on one SHM group name.
         This is a lightweight consistency check for VLLM_DIST_IDENT/name inputs.
         """
         local_name = _CPUSHMDistributed.make_group_name(self)
@@ -89,8 +88,7 @@ def all_reduce(self, input_):
     def gather(
         self, input_: torch.Tensor, dst: int = 0, dim: int = -1
     ) -> torch.Tensor | None:
-        """
-        NOTE: We assume that the input tensor is on the same device across
+        """NOTE: We assume that the input tensor is on the same device across
         all the ranks.
         NOTE: `dst` is the local rank of the destination rank.
         """
@@ -180,11 +178,9 @@ def dispatch_router_logits(
         tuple[torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Dispatch the hidden states and router logits to the appropriate device.
+        """Dispatch the hidden states and router logits to the appropriate device.
         This is a no-op in the base class.
         """
-
         assert self.all2all_manager is not None
         return self.all2all_manager.dispatch_router_logits(
             hidden_states,
@@ -204,8 +200,7 @@ def dispatch(
         tuple[torch.Tensor, torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Dispatch the hidden states and topk weights/ids to the appropriate device.
+        """Dispatch the hidden states and topk weights/ids to the appropriate device.
         This is a no-op in the base class.
         """
         assert self.all2all_manager is not None
@@ -220,8 +215,7 @@ def dispatch(
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
     ) -> torch.Tensor:
-        """
-        Combine the hidden states and router logits from the appropriate device.
+        """Combine the hidden states and router logits from the appropriate device.
         This is a no-op in the base class.
         """
         assert self.all2all_manager is not None
diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py
index faa9fd0987da..fb0c281b4a86 100644
--- a/vllm/distributed/device_communicators/cuda_communicator.py
+++ b/vllm/distributed/device_communicators/cuda_communicator.py
@@ -109,6 +109,7 @@ def __init__(
         self.fi_ar_comm: FlashInferAllReduce | None = None
         self.fi_pcie_ipc_ar_comm: FlashInferPcieIpcAllReduce | None = None
         self.aiter_ar_comm: AiterCustomAllreduce | None = None
+        self.use_aiter_ag_rs: bool = False
 
         if use_torch_symm_mem and current_platform.is_cuda():
             self.symm_mem_comm = SymmMemCommunicator(
@@ -149,6 +150,22 @@ def __init__(
                 ),
             )
 
+        # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine
+        if (
+            "dp" in unique_name
+            and self.world_size in (2, 4, 8)
+            and current_platform.is_rocm()
+            and rocm_aiter_ops.is_custom_all_reduce_enabled()
+        ):
+            self.aiter_ar_comm = AiterCustomAllreduce(
+                group=self.cpu_group,
+                device=self.device,
+            )
+            if self.aiter_ar_comm.disabled:
+                self.aiter_ar_comm = None
+            else:
+                self.use_aiter_ag_rs = True
+
         if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm():
             # Initialize a custom quick all-reduce implementation for AMD.
             # Quick reduce is designed as a complement to custom allreduce
@@ -197,6 +214,14 @@ def __init__(
                     tcp_store_group,
                     device_group=self.device_group,
                 )
+            elif self.all2all_backend == "moonep":
+                from .all2all import MoonEPAll2AllManager
+
+                self.all2all_manager = MoonEPAll2AllManager(
+                    self.cpu_group,
+                    tcp_store_group,
+                    device_group=self.device_group,
+                )
             elif self.all2all_backend == "nixl_ep":
                 from .all2all import NixlEPAll2AllManager
 
@@ -285,7 +310,11 @@ def _log_all_reduce_backend_selection(self) -> None:
             enabled_ar_backends.append("NCCL_SYMM_MEM")
         if self.qr_comm is not None and not self.qr_comm.disabled:
             enabled_ar_backends.append("QUICK_REDUCE")
-        if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled:
+        if (
+            self.use_aiter_allreduce
+            and self.aiter_ar_comm is not None
+            and not self.aiter_ar_comm.disabled
+        ):
             enabled_ar_backends.append("AITER_CUSTOM")
         if self.ca_comm is not None and not self.ca_comm.disabled:
             enabled_ar_backends.append("CUSTOM")
@@ -340,7 +369,8 @@ def all_reduce(self, input_):
             return out
         aiter_ar_comm = self.aiter_ar_comm
         if (
-            aiter_ar_comm is not None
+            self.use_aiter_allreduce
+            and aiter_ar_comm is not None
             and not aiter_ar_comm.disabled
             and aiter_ar_comm.should_custom_ar(input_)
         ):
@@ -462,6 +492,11 @@ def reduce_scatterv(
             # Convert negative dim to positive.
             dim += input_.dim()
 
+        # 'sizes' is not needed if all inputs in the same group have the same
+        # shape
+        if sizes is not None and all(s == sizes[0] for s in sizes):
+            sizes = None
+
         # Note: This will produce an incorrect answer if we don't make
         # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
         input_tensor = input_.movedim(0, dim).contiguous()
@@ -475,6 +510,16 @@ def reduce_scatterv(
             chunk_size = input_tensor.shape[0] // world_size
         output_shape = (chunk_size,) + input_tensor.shape[1:]
 
+        if self._can_use_aiter_ag_rs(sizes):
+            aiter_comm = self.aiter_ar_comm
+            assert aiter_comm is not None
+            if aiter_comm.should_custom_rs(input_tensor, dim=0):
+                output = torch.empty(
+                    output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+                )
+                aiter_comm.custom_reduce_scatter(input_tensor, output, dim=0)
+                return output.movedim(0, dim).contiguous()
+
         # Symmetric memory is only used when all ranks have uniform sizes.
         # ncclCommWindowRegister is collective: asymmetric pool allocations
         # from variable per-rank sizes cause deadlocks.
@@ -574,7 +619,7 @@ def _reduce_scatter_symm_mem(
         return symm_output
 
     def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
-        """Sends a tensor to the destination rank in a blocking way"""
+        """Sends a tensor to the destination rank in a blocking way."""
         """NOTE: `dst` is the local rank of the destination rank."""
         if dst is None:
             dst = (self.rank_in_group + 1) % self.world_size
@@ -632,6 +677,34 @@ def destroy(self):
             self.all2all_manager.destroy()
             self.all2all_manager = None  # type: ignore[assignment]
 
+    def _can_use_aiter_ag_rs(self, sizes: list[int] | None) -> bool:
+        """Whether the AITER custom AG/RS fast path may run for this collective.
+
+        Requires:
+        - uniform batches
+        - FULL CUDAgraphs
+        """
+        if (
+            not self.use_aiter_ag_rs
+            or self.aiter_ar_comm is None
+            or self.aiter_ar_comm.disabled
+        ):
+            return False
+        if sizes is not None:
+            return False
+
+        from vllm.config.compilation import CUDAGraphMode
+        from vllm.forward_context import get_forward_context
+
+        try:
+            ctx = get_forward_context()
+        except AssertionError:
+            return False
+        if ctx.cudagraph_runtime_mode != CUDAGraphMode.FULL:
+            return False
+        bd = ctx.batch_descriptor
+        return bd is not None and bd.uniform
+
     def suspend(self) -> None:
         if self.pynccl_comm is not None:
             self.pynccl_comm.suspend()
@@ -665,14 +738,28 @@ def all_gatherv(
         if dim != 0:
             raise NotImplementedError("only dim 0 all-gatherv is supported")
         world_size = self.world_size
-        pynccl_comm = self.pynccl_comm
-        assert pynccl_comm is not None and not pynccl_comm.disabled
 
         # 'sizes' is not needed if all inputs in the same group have the same
         # shape
         if sizes is not None and all(s == sizes[0] for s in sizes):
             sizes = None
 
+        if self._can_use_aiter_ag_rs(sizes):
+            aiter_comm = self.aiter_ar_comm
+            assert aiter_comm is not None
+            if isinstance(input_, torch.Tensor):
+                if aiter_comm.should_custom_ag(input_):
+                    out = aiter_comm.custom_all_gather(input_, dim=0)
+                    if out is not None:
+                        return out
+            elif all(aiter_comm.should_custom_ag(inp) for inp in input_):
+                outs = [aiter_comm.custom_all_gather(inp, dim=0) for inp in input_]
+                if all(o is not None for o in outs):
+                    return outs
+
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None and not pynccl_comm.disabled
+
         # Symmetric memory is only used when all ranks have uniform sizes.
         # ncclCommWindowRegister is collective: asymmetric pool allocations
         # from variable per-rank sizes cause deadlocks.
@@ -772,11 +859,9 @@ def dispatch_router_logits(
         tuple[torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Dispatch the hidden states and router logits to the appropriate device.
+        """Dispatch the hidden states and router logits to the appropriate device.
         This is a no-op in the base class.
         """
-
         assert self.all2all_manager is not None
         return self.all2all_manager.dispatch_router_logits(
             hidden_states,
@@ -796,8 +881,7 @@ def dispatch(
         tuple[torch.Tensor, torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Dispatch the hidden states and topk weights/ids to the appropriate device.
+        """Dispatch the hidden states and topk weights/ids to the appropriate device.
         This is a no-op in the base class.
         """
         assert self.all2all_manager is not None
@@ -812,8 +896,7 @@ def dispatch(
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
     ) -> torch.Tensor:
-        """
-        Combine the hidden states and router logits from the appropriate device.
+        """Combine the hidden states and router logits from the appropriate device.
         This is a no-op in the base class.
         """
         assert self.all2all_manager is not None
diff --git a/vllm/distributed/device_communicators/cuda_wrapper.py b/vllm/distributed/device_communicators/cuda_wrapper.py
index a5026163e934..d27d84c03045 100644
--- a/vllm/distributed/device_communicators/cuda_wrapper.py
+++ b/vllm/distributed/device_communicators/cuda_wrapper.py
@@ -78,6 +78,16 @@ class CudaRTLibrary:
             cudaError_t,
             [ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint],
         ),
+        # cudaError_t cudaHostRegister ( void* ptr, size_t size, unsigned int flags )
+        Function(
+            "cudaHostRegister",
+            cudaError_t,
+            [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint],
+        ),
+        # cudaError_t cudaHostUnregister ( void* ptr )
+        Function("cudaHostUnregister", cudaError_t, [ctypes.c_void_p]),
+        # cudaError_t cudaGetLastError ( void )
+        Function("cudaGetLastError", cudaError_t, []),
     ]
 
     # https://rocm.docs.amd.com/projects/HIPIFY/en/latest/tables/CUDA_Runtime_API_functions_supported_by_HIP.html # noqa
@@ -92,6 +102,9 @@ class CudaRTLibrary:
         "cudaMemcpy": "hipMemcpy",
         "cudaIpcGetMemHandle": "hipIpcGetMemHandle",
         "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle",
+        "cudaHostRegister": "hipHostRegister",
+        "cudaHostUnregister": "hipHostUnregister",
+        "cudaGetLastError": "hipGetLastError",
     }
 
     # class attribute to store the mapping from the path to the library
@@ -185,3 +198,16 @@ def cudaIpcOpenMemHandle(self, handle: cudaIpcMemHandle_t) -> ctypes.c_void_p:
             )
         )
         return devPtr
+
+    def cudaHostRegister(self, ptr: int, size: int, flags: int = 0) -> int:
+        """Return the raw error code instead of raising, since callers can
+        fall back to pageable memory. Drain a failure with cudaGetLastError."""
+        return self.funcs["cudaHostRegister"](ctypes.c_void_p(ptr), size, flags)
+
+    def cudaHostUnregister(self, ptr: int) -> int:
+        """Return the raw error code instead of raising."""
+        return self.funcs["cudaHostUnregister"](ctypes.c_void_p(ptr))
+
+    def cudaGetLastError(self) -> int:
+        """Return and clear the error pending on this thread."""
+        return self.funcs["cudaGetLastError"]()
diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py
index 6c15a9827429..e18f9194f54e 100644
--- a/vllm/distributed/device_communicators/custom_all_reduce.py
+++ b/vllm/distributed/device_communicators/custom_all_reduce.py
@@ -129,8 +129,7 @@ def __init__(
         max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE,
         symm_mem_enabled=False,
     ) -> None:
-        """
-        Args:
+        """Args:
             group: the process group to work on. If None, it will use the
                 default process group.
             device: the device to bind the CustomAllreduce to. If None,
@@ -138,6 +137,7 @@ def __init__(
         It is the caller's responsibility to make sure each communicator
         is bind to a unique device, and all communicators in this group
         are in the same node.
+
         """
         self._IS_CAPTURING = False
         self._ptr = 0
@@ -370,8 +370,7 @@ def _init_mnnvl_buffer(self, stage_size: int) -> None:
 
     @contextmanager
     def capture(self):
-        """
-        The main responsibility of this context manager is the
+        """The main responsibility of this context manager is the
         `register_graph_buffers` call at the end of the context.
         It records all the buffer addresses used in the CUDA graph.
         """
diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py
index dd46f782b810..f1644d699ed6 100644
--- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py
+++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py
@@ -140,6 +140,7 @@ def _resolve_fi_ar_backend() -> tuple[str, bool]:
         is True only when ``auto`` selects mnnvl for a single node, so that
         workspace creation can fall back to trtllm on single-node topologies
         without NVSwitch multicast support (where mnnvl is unavailable).
+
     """
     backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
     if backend != "auto":
@@ -168,8 +169,7 @@ def get_fi_ar_workspace(
     dtype: torch.dtype,
     group: ProcessGroup,
 ):
-    """
-    Return the allreduce workspace for non-quant patterns, initializing if needed.
+    """Return the allreduce workspace for non-quant patterns, initializing if needed.
 
     Used by AllReduceFusionPass (non-quant patterns) and FlashInferAllReduce
     for standalone allreduce. Backend is controlled by
@@ -218,7 +218,7 @@ def _get_or_create(be: str):
             f"with backend={backend}"
         )
     else:
-        logger.warning_once(
+        logger.error_once(
             "Failed to initialize FlashInfer Allreduce norm fusion workspace "
             f"with backend={backend}"
         )
@@ -234,8 +234,7 @@ def get_fi_ar_quant_workspace(
     dtype: torch.dtype,
     group: ProcessGroup,
 ):
-    """
-    Return the allreduce workspace for quant patterns, initializing if needed.
+    """Return the allreduce workspace for quant patterns, initializing if needed.
 
     Backend is controlled by VLLM_FLASHINFER_ALLREDUCE_BACKEND env var, matching
     non-quant fusion. With ``auto`` this prefers mnnvl and falls back to trtllm
@@ -292,7 +291,7 @@ def get_fi_ar_quant_workspace(
             f"fusion workspace with backend={backend}"
         )
     else:
-        logger.warning_once(
+        logger.error_once(
             "Failed to initialize FlashInfer Allreduce norm quantization "
             f"fusion workspace with backend={backend}"
         )
diff --git a/vllm/distributed/device_communicators/pynccl.py b/vllm/distributed/device_communicators/pynccl.py
index 916d0308bf76..76eeb2291f90 100644
--- a/vllm/distributed/device_communicators/pynccl.py
+++ b/vllm/distributed/device_communicators/pynccl.py
@@ -3,6 +3,7 @@
 
 
 # ===================== import region =====================
+import contextlib
 import threading
 
 import torch
@@ -26,6 +27,33 @@
 
 logger = init_logger(__name__)
 
+_warmup_tls = threading.local()
+
+
+def _warmup_deferred() -> bool:
+    return getattr(_warmup_tls, "deferred", False)
+
+
+@contextlib.contextmanager
+def defer_comm_warmup_on_rocm():
+    """Build communicators on this thread without their warm-up on ROCm, where
+    the warm-up can deadlock if built next to a live engine whose existing
+    communicators still have collectives in flight; instead the warm-up is
+    deferred to when the engine is quiet (during commit).
+    """
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_rocm():
+        yield
+        return
+    prev = _warmup_deferred()
+    _warmup_tls.deferred = True
+    try:
+        yield
+    finally:
+        _warmup_tls.deferred = prev
+
+
 _NCCL_SYMM_OPS_REGISTERED = False
 
 _NCCL_SUSPEND_MEM = 0x01
@@ -70,8 +98,7 @@ def __init__(
         device: int | str | torch.device,
         library_path: str | None = None,
     ):
-        """
-        Args:
+        """Args:
             group: the process group to work on. If None, it will use the
                 default process group.
             device: the device to bind the PyNcclCommunicator to. If None,
@@ -80,6 +107,7 @@ def __init__(
                 use the default library path.
         It is the caller's responsibility to make sure each communicator
         is bind to a unique device.
+
         """
         if not isinstance(group, StatelessProcessGroup):
             assert dist.is_initialized()
@@ -153,6 +181,9 @@ def _init_comm(self, device: int | str | torch.device) -> None:
                 self.world_size, self.unique_id, self.rank
             )
 
+            if _warmup_deferred():
+                return
+
             stream = current_stream()
             # A small all_reduce for warmup.
             data = torch.zeros(1, device=device)
diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py
index 3c9d759dd7f1..406869ea2c4b 100644
--- a/vllm/distributed/device_communicators/quick_all_reduce.py
+++ b/vllm/distributed/device_communicators/quick_all_reduce.py
@@ -62,8 +62,7 @@ class QuickAllReduce:
     }
 
     def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> None:
-        """
-        Custom allreduce provides non-destructive acceleration and is
+        """Custom allreduce provides non-destructive acceleration and is
         available for CUDA and ROCm MI300 series.
 
         Custom quick allreduce leverages quantization for further
@@ -86,6 +85,7 @@ def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> Non
         It is the caller's responsibility to make sure each communicator
         is bind to a unique device, and all communicators in this group
         are in the same node.
+
         """
         self.disabled = True
         if not self._rocm_arch_available():
@@ -301,8 +301,7 @@ def _rocm_arch_available(self):
             return False
 
     def create_shared_buffer(self):
-        """
-        Creates a shared buffer for quickreduce.
+        """Creates a shared buffer for quickreduce.
         Has to be called after init_custom_qr
         """
         handle = ops.qr_get_handle(self._ptr)
@@ -312,9 +311,7 @@ def create_shared_buffer(self):
         ops.qr_open_handles(self._ptr, handles)
 
     def should_quick_allreduce(self, inp: torch.Tensor):
-        """
-        Check if quickreduce is available
-        """
+        """Check if quickreduce is available."""
         if self.disabled:
             return False
         if inp.dtype not in self._SUPPORTED_DTYPES:
diff --git a/vllm/distributed/device_communicators/ray_communicator.py b/vllm/distributed/device_communicators/ray_communicator.py
index d9517f51acad..ff3c2f1fd479 100644
--- a/vllm/distributed/device_communicators/ray_communicator.py
+++ b/vllm/distributed/device_communicators/ray_communicator.py
@@ -20,8 +20,7 @@
 
 
 class RayPPCommunicator(Communicator):
-    """
-    Communicator to be used for pipeline parallelism in Ray Compiled Graph.
+    """Communicator to be used for pipeline parallelism in Ray Compiled Graph.
     This is wraps around the vLLM _PP GroupCoordinator.
 
     This class is not thread-safe.
@@ -38,8 +37,7 @@ def __init__(
         cuda_stream: torch.cuda.Stream | None,
         use_communication_streams: bool = False,
     ):
-        """
-        Initialize a RayPPCommunicator that can be used to communicate with
+        """Initialize a RayPPCommunicator that can be used to communicate with
         other Ray Compiled Graph actors for pipeline parallelism.
 
         Args:
@@ -54,6 +52,7 @@ def __init__(
                 is not supported.
             use_communication_streams: Whether to use communication streams.
                 This is not supported.
+
         """
         self._world_size = world_size
         self._rank: int | None = None
@@ -87,8 +86,7 @@ def __init__(
         self._closed = False
 
     def _build_actor_rank_mapping(self):
-        """
-        Use collective communication to build a mapping from actor IDs to ranks.
+        """Use collective communication to build a mapping from actor IDs to ranks.
         This should be called once during initialization.
         """
         if self._comm is None:
@@ -128,9 +126,7 @@ def get_actor_handles(self) -> list["ray.actor.ActorHandle"]:
         return self._actor_handles
 
     def get_rank(self, actor: ray.actor.ActorHandle) -> int:
-        """
-        Return the given actor's rank using device communicator collective ops.
-        """
+        """Return the given actor's rank using device communicator collective ops."""
         assert hasattr(self, "_actor_id_to_rank"), (
             "Actor rank mapping not built. "
             "This should have been done during initialization."
@@ -144,20 +140,15 @@ def get_rank(self, actor: ray.actor.ActorHandle) -> int:
             raise ValueError(f"Actor {actor} not found in communicator group")
 
     def get_self_rank(self) -> int | None:
-        """
-        Return this actor's rank.
-        """
+        """Return this actor's rank."""
         return self._rank
 
     def get_world_size(self) -> int:
-        """
-        Return the number of ranks in the RayPPCommunicator group.
-        """
+        """Return the number of ranks in the RayPPCommunicator group."""
         return self._world_size
 
     def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
-        """
-        Send a torch.Tensor to a peer.
+        """Send a torch.Tensor to a peer.
 
         This returns when the send kernel has been queued, but the kernel may
         not have completed. Therefore, the caller should ensure that there are
@@ -170,6 +161,7 @@ def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
             buf: The torch.Tensor to send. It should already be on this
                 actor's default device.
             peer_rank: The rank of the actor to send to.
+
         """
         if self._closed:
             raise RayChannelError("RayPPCommunicator has been destroyed.")
@@ -184,8 +176,7 @@ def recv(
         peer_rank: int,
         allocator: TorchTensorAllocator,
     ) -> "torch.Tensor":
-        """
-        Receive a torch.Tensor from a peer and synchronize the current stream.
+        """Receive a torch.Tensor from a peer and synchronize the current stream.
 
         After this call returns, the receive buffer is safe to read from
         any stream. An RayChannelError will be raised if an error occurred
@@ -197,6 +188,7 @@ def recv(
             peer_rank: The rank of the actor to receive from.
             allocator: The allocator to use to create the received tensor.
                 This is ignored for this implementation.
+
         """
         if self._closed:
             raise RayChannelError("RayPPCommunicator has been destroyed.")
diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py
index 6caacea6ec5d..bafc5728b85e 100644
--- a/vllm/distributed/device_communicators/shm_broadcast.py
+++ b/vllm/distributed/device_communicators/shm_broadcast.py
@@ -34,6 +34,7 @@
 from vllm.distributed.utils import StatelessProcessGroup, sched_yield
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
+from vllm.utils.cpu_resource_utils import check_cgroup_memory_available
 from vllm.utils.network_utils import (
     get_ip,
     get_open_zmq_inproc_path,
@@ -76,8 +77,7 @@
 
 
 def memory_fence():
-    """
-    Full memory barrier for shared memory synchronization.
+    """Full memory barrier for shared memory synchronization.
 
     Ensures all prior memory writes are visible to other processes before
     any subsequent reads. This is critical for lock-free producer-consumer
@@ -110,8 +110,7 @@ def to_bytes_big(value: int, size: int) -> bytes:
 
 
 class SpinCondition:
-    """
-    This class implements an interface similar to a threading.Condition. It
+    """This class implements an interface similar to a threading.Condition. It
     allows a writer to notify readers to wake up and read from the shared memory
     buffer. This notification is done over a zmq socket.
 
@@ -224,26 +223,38 @@ def notify(self):
 SHM_PATH = "/dev/shm"
 
 
-def check_shm_free_space(required_bytes: int, shm_path: str = SHM_PATH) -> None:
-    """Raise if ``shm_path`` cannot fit a ``required_bytes`` shared segment.
+def check_shm_free_space(
+    required_bytes: int,
+    shm_path: str = SHM_PATH,
+    *,
+    allocation_name: str = "shared-memory allocation",
+) -> None:
+    """Raise if SHM cannot fit a shared segment and log cgroup headroom.
 
     Args:
         required_bytes: Size of the shared-memory segment to be created.
-        shm_path: Mount point backing POSIX shared memory; skipped if absent.
+        shm_path: Mount point backing POSIX shared memory; its filesystem
+            check is skipped if absent.
+        allocation_name: Human-readable name used in errors and logs.
 
     Raises:
-        RuntimeError: If ``required_bytes`` exceeds the free space.
+        RuntimeError: If the SHM filesystem has insufficient space.
+
     """
-    if not os.path.isdir(shm_path):
-        return
-    free_bytes = shutil.disk_usage(shm_path).free
-    if required_bytes <= free_bytes:
-        return
-    mib = 1 << 20
-    raise RuntimeError(
-        f"Insufficient space in {shm_path}: {required_bytes / mib:.0f} MiB "
-        f"required, {free_bytes / mib:.0f} MiB free. Increase {shm_path} "
-        "(e.g. --shm-size or --ipc=host)."
+    if os.path.isdir(shm_path):
+        free_bytes = shutil.disk_usage(shm_path).free
+        if required_bytes > free_bytes:
+            mib = 1 << 20
+            raise RuntimeError(
+                f"Insufficient space in {shm_path} for {allocation_name}: "
+                f"{required_bytes / mib:.0f} MiB required, "
+                f"{free_bytes / mib:.0f} MiB free. Increase {shm_path} "
+                "(e.g. --shm-size or --ipc=host)."
+            )
+
+    check_cgroup_memory_available(
+        required_bytes,
+        allocation_name,
     )
 
 
@@ -929,8 +940,7 @@ def create_from_process_group_single_reader(
         reader_rank: int = 0,
         blocking: bool = False,
     ) -> tuple["MessageQueue", list[Handle]]:
-        """
-        Creates a MessageQueue for a process group with a single reader.
+        """Creates a MessageQueue for a process group with a single reader.
 
         This method is designed for scenarios where only one process (the reader)
         will consume messages, and all other processes are writers. It sets up
@@ -950,6 +960,7 @@ def create_from_process_group_single_reader(
             tuple[MessageQueue, list[Handle]]:
             The MessageQueue instance for the calling process,
             and a list of handles (only non-empty for the reader process).
+
         """
         from vllm.platforms.interface import get_assigned_physical_gpu_ids
 
@@ -982,8 +993,7 @@ def create_from_process_group(
         external_writer_handle=None,
         blocking: bool = True,
     ) -> "MessageQueue":
-        """
-        Creates a MessageQueue for a distributed process group with one writer and
+        """Creates a MessageQueue for a distributed process group with one writer and
         multiple readers.
 
         This method is designed for scenarios where one process (the writer) sends
diff --git a/vllm/distributed/device_communicators/shm_object_storage.py b/vllm/distributed/device_communicators/shm_object_storage.py
index e2d2b248346b..bc4f1b6e5159 100644
--- a/vllm/distributed/device_communicators/shm_object_storage.py
+++ b/vllm/distributed/device_communicators/shm_object_storage.py
@@ -20,8 +20,7 @@
 
 
 class SingleWriterShmRingBuffer:
-    """
-    A single-writer, multiple-reader ring buffer implementation using shared
+    """A single-writer, multiple-reader ring buffer implementation using shared
     memory. This class provides a thread-safe ring buffer where one process
     can write data while multiple processes/threads can read from it.
 
@@ -190,8 +189,7 @@ def byte2int(self, byte_data: bytes) -> int:
         return int.from_bytes(byte_data, "little", signed=True)
 
     def allocate_buf(self, size: int) -> tuple[int, int]:
-        """
-        Allocate a buffer `MD_SIZE` + `size` bytes in the shared memory.
+        """Allocate a buffer `MD_SIZE` + `size` bytes in the shared memory.
         Memory layout:
         `[4-byte monotonic_id][4-byte size][buffer data...]`
         """
@@ -260,8 +258,7 @@ def free_buf(
         is_free_fn: Callable[[int, memoryview], bool],
         nbytes: int | None = None,
     ) -> Iterable[int]:
-        """
-        Free a buffer of the given size. This is a no-op in shared memory,
+        """Free a buffer of the given size. This is a no-op in shared memory,
         but we need to keep track of the metadata.
 
         If freed memory spreads across the end and start of the ring buffer,
@@ -269,10 +266,13 @@ def free_buf(
         still might not be a contiguous space of `nbytes` available.
 
         Args:
+            is_free_fn (Callable[[int, memoryview], bool]): Predicate called with
+                a monotonic id and the buffer, returning True when that buffer
+                can be reclaimed.
             nbytes (int, optional): The size of the buffer to free. If None,
                 frees the maximum size of the ring buffer.
-        """
 
+        """
         assert self.is_writer, "Only the writer can free buffers."
         logger.debug(
             "Freeing up space in the ring buffer, "
@@ -412,8 +412,7 @@ class ShmObjectStorageHandle:
 
 
 class SingleWriterShmObjectStorage:
-    """
-    A single-writer, multiple-reader object storage system built on top of a
+    """A single-writer, multiple-reader object storage system built on top of a
     shared memory ring buffer. Provides key-value storage with automatic memory
     management and cross-process serialization support.
 
@@ -458,8 +457,7 @@ def __init__(
         serde_class: type[ObjectSerde] = MsgpackSerde,
         reader_lock: LockType | None = None,
     ):
-        """
-        Initialize the object storage.
+        """Initialize the object storage.
 
         Args:
             max_object_size: Maximum size for a single object in bytes.
@@ -467,10 +465,11 @@ def __init__(
             ring_buffer: The shared memory ring buffer for storing objects.
             serde_class: Serializer/deserializer for objects.
             reader_lock: Optional lock for synchronizing reader access.
+
         Raises:
             ValueError: If reader_lock is None for readers.
-        """
 
+        """
         self.max_object_size = max_object_size
         self.n_readers = n_readers
         self.serde_class = serde_class
@@ -547,22 +546,17 @@ def free_unused(self) -> None:
             del self.writer_flag[freed_id]
 
     def is_cached(self, key: str) -> bool:
-        """
-        Check if the object with the given key is cached.
-        """
+        """Check if the object with the given key is cached."""
         return key in self.key_index
 
     def get_cached(self, key: str) -> tuple[int, int]:
-        """
-        Get the cached object by key if it exists.
-        """
+        """Get the cached object by key if it exists."""
         address, monotonic_id = self.key_index[key]
         self.increment_writer_flag(monotonic_id)
         return address, monotonic_id
 
     def put(self, key: str, value: Any) -> tuple[int, int]:
-        """
-        Store a key-value pair in the object storage.
+        """Store a key-value pair in the object storage.
         Attempts to free max_object_size bytes using FIFO order
         when the ring buffer runs out of space during a put() operation.
 
@@ -574,6 +568,7 @@ def put(self, key: str, value: Any) -> tuple[int, int]:
             MemoryError: If there's not enough space in the buffer
             ValueError: If the serialized object is too large
             ValueError: If the key already exists in the storage
+
         """
         if key in self.key_index:
             raise ValueError(f"Key '{key}' already exists in the storage.")
@@ -639,8 +634,7 @@ def touch(
         address: int = 0,
         monotonic_id: int = 0,
     ) -> None:
-        """
-        Touch an existing cached item to update its eviction status.
+        """Touch an existing cached item to update its eviction status.
 
         For writers (ShmObjectStoreSenderCache): Increment writer_flag
         For readers (ShmObjectStoreReceiverCache): Increment reader_count
@@ -700,8 +694,7 @@ def create_from_handle(
         )
 
     def default_is_free_check(self, id: int, buf: memoryview) -> bool:
-        """
-        Default is_free function that checks if the first 4 bytes are zero.
+        """Default is_free function that checks if the first 4 bytes are zero.
         This indicates that the buffer is free.
         """
         reader_count = int.from_bytes(buf[0:4], "little", signed=True)
diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py
index 7ca132824eca..7a5913754fbd 100644
--- a/vllm/distributed/device_communicators/xpu_communicator.py
+++ b/vllm/distributed/device_communicators/xpu_communicator.py
@@ -6,6 +6,7 @@
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
+import vllm.envs as envs
 from vllm.logger import init_logger
 
 from .base_device_communicator import DeviceCommunicatorBase
@@ -45,7 +46,25 @@ def __init__(
                 self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
                 logger.info("Using AgRs manager on XPU device.")
 
+    def _fixed_rank_sum(self, input_: torch.Tensor) -> torch.Tensor:
+        flat_input = input_.reshape(-1)
+        gathered = torch.empty(
+            (self.world_size, flat_input.numel()),
+            dtype=input_.dtype,
+            device=input_.device,
+        )
+        dist.all_gather_single(gathered.view(-1), flat_input, group=self.device_group)
+        output = gathered[0].clone()
+        for rank in range(1, self.world_size):
+            output.add_(gathered[rank])
+        return output.view(input_.shape)
+
     def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
+        if envs.VLLM_BATCH_INVARIANT:
+            # XCCL changes its reduction order with message size. Gather first
+            # and accumulate in rank order, independent of the token batch.
+            return self._fixed_rank_sum(input_)
+
         output = input_.clone()
         dist.all_reduce(output, group=self.device_group)
         return output
@@ -69,7 +88,12 @@ def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
             output_shape, dtype=input_tensor.dtype, device=input_tensor.device
         )
 
-        dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
+        if envs.VLLM_BATCH_INVARIANT:
+            reduced = self._fixed_rank_sum(input_tensor)
+            start = self.rank_in_group * chunk_size
+            output.copy_(reduced.narrow(0, start, chunk_size))
+        else:
+            dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
 
         # Reshape before returning
         return output.movedim(0, dim).contiguous()
@@ -99,12 +123,32 @@ def reduce_scatterv(
         output = torch.empty(
             output_shape, dtype=input_tensor.dtype, device=input_tensor.device
         )
-        if sizes is not None and sizes.count(sizes[0]) != len(sizes):
-            # if inputs shape in different ranks is not the same using reduce_scatter
+        if envs.VLLM_BATCH_INVARIANT:
+            reduced = self._fixed_rank_sum(input_tensor)
+            start = (
+                self.rank_in_group * chunk_size
+                if sizes is None
+                else sum(sizes[: self.rank_in_group])
+            )
+            output.copy_(reduced.narrow(0, start, chunk_size))
+        elif sizes is None or sizes.count(sizes[0]) == len(sizes):
+            dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
+        elif torch.xpu.device_count() < self.world_size:
+            # uneven reduce_scatter is not supported when the
+            # process cannot see every device in the group (
+            # e.g. external-LBDP).
+            max_size = max(sizes)
+            padded = input_tensor.new_zeros(
+                (world_size * max_size,) + input_tensor.shape[1:]
+            )
+            for rank, split in enumerate(input_tensor.split(sizes, dim=0)):
+                padded[rank * max_size : rank * max_size + sizes[rank]] = split
+            padded_output = input_tensor.new_empty((max_size,) + input_tensor.shape[1:])
+            dist.reduce_scatter_tensor(padded_output, padded, group=self.device_group)
+            output = padded_output[:chunk_size]
+        else:
             input_splits = list(input_tensor.split(sizes, dim=0))
             dist.reduce_scatter(output, input_splits, group=self.device_group)
-        else:
-            dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
         # Reshape before returning
         return output.movedim(0, dim).contiguous()
 
@@ -179,7 +223,7 @@ def gather(
             (self.world_size,) + input_size, dtype=input_.dtype, device=input_.device
         )
         # All-gather.
-        dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
+        dist.all_gather_single(output_tensor, input_, group=self.device_group)
         if self.rank_in_group == dst:
             # Reshape
             output_tensor = output_tensor.movedim(0, dim)
@@ -205,11 +249,9 @@ def dispatch_router_logits(
         tuple[torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Dispatch the hidden states and router logits to the appropriate device.
+        """Dispatch the hidden states and router logits to the appropriate device.
         This is a no-op in the base class.
         """
-
         assert self.all2all_manager is not None
         return self.all2all_manager.dispatch_router_logits(
             hidden_states,
@@ -229,8 +271,7 @@ def dispatch(
         tuple[torch.Tensor, torch.Tensor, torch.Tensor]
         | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
-        """
-        Dispatch the hidden states and topk weights/ids to the appropriate device.
+        """Dispatch the hidden states and topk weights/ids to the appropriate device.
         This is a no-op in the base class.
         """
         assert self.all2all_manager is not None
@@ -245,8 +286,7 @@ def dispatch(
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
     ) -> torch.Tensor:
-        """
-        Combine the hidden states and router logits from the appropriate device.
+        """Combine the hidden states and router logits from the appropriate device.
         This is a no-op in the base class.
         """
         assert self.all2all_manager is not None
diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py
index ea44fb6f5beb..c56d2cc2129a 100644
--- a/vllm/distributed/ec_transfer/ec_connector/base.py
+++ b/vllm/distributed/ec_transfer/ec_connector/base.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ECConnectorBase Class for Distributed Encoder Cache &
+"""ECConnectorBase Class for Distributed Encoder Cache &
 P2P Encoder cache communication in V1
 
 The class provides the following primitives:
@@ -36,6 +35,12 @@
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
+    from vllm.distributed.ec_transfer.ec_connector.metrics import (
+        ECConnectorPromMetrics,
+        ECConnectorStats,
+        PromMetric,
+        PromMetricT,
+    )
     from vllm.v1.request import Request
 
 logger = init_logger(__name__)
@@ -50,8 +55,7 @@ class ECConnectorRole(enum.Enum):
 
 
 class ECConnectorMetadata(ABC):  # noqa: B024
-    """
-    Abstract Metadata used to communicate between the
+    """Abstract Metadata used to communicate between the
     Scheduler ECConnector and Worker ECConnector.
     """
 
@@ -59,8 +63,7 @@ class ECConnectorMetadata(ABC):  # noqa: B024
 
 
 class ECConnectorWorkerMetadata(ABC):
-    """
-    Abstract Metadata used to communicate back
+    """Abstract Metadata used to communicate back
     Worker ECConnector -> Scheduler ECConnector.
 
     Each worker can output its own metadata.
@@ -73,9 +76,7 @@ class ECConnectorWorkerMetadata(ABC):
     def aggregate(
         self, other: "ECConnectorWorkerMetadata"
     ) -> "ECConnectorWorkerMetadata":
-        """
-        Aggregate metadata with another `ECConnectorWorkerMetadata` object.
-        """
+        """Aggregate metadata with another `ECConnectorWorkerMetadata` object."""
         pass
 
 
@@ -103,8 +104,7 @@ def is_consumer(self) -> bool:
         return self._is_consumer
 
     def shutdown(self) -> None:
-        """
-        Shutdown the connector. This is called when the process
+        """Shutdown the connector. This is called when the process
         is shutting down to ensure that all the async operations are
         completed and the connector is cleaned up properly.
         """
@@ -123,6 +123,7 @@ def bind_connector_metadata(self, connector_metadata: ECConnectorMetadata) -> No
 
         Args:
             connector_metadata (dict): the connector metadata.
+
         """
         self._connector_metadata = connector_metadata
 
@@ -141,8 +142,8 @@ def _get_connector_metadata(self) -> ECConnectorMetadata:
 
         Returns:
             ConnectorMetadata: the connector metadata.
-        """
 
+        """
         # Should only be called while set to valid metadata.
         assert self._connector_metadata is not None
         return self._connector_metadata
@@ -151,10 +152,11 @@ def register_caches(
         self,
         ec_caches: dict[str, torch.Tensor],
     ):
-        """
-        Initialize with the EC caches.
+        """Initialize with the EC caches.
+
         Args:
             ec_caches: dictionary of encoder cache
+
         """
         # TODO: Implement this later for P2P feature
         return
@@ -163,8 +165,7 @@ def register_caches(
     def start_load_caches(
         self, encoder_cache: dict[str, torch.Tensor], **kwargs
     ) -> None:
-        """
-        Start loading the cache from the connector into vLLM's encoder cache.
+        """Start loading the cache from the connector into vLLM's encoder cache.
 
         This method loads the encoder cache based on metadata provided by the scheduler.
         It is called before `_gather_mm_embeddings` for the EC Connector. For EC,
@@ -174,6 +175,7 @@ def start_load_caches(
             encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
                 data hashes (`mm_hash`) to encoder cache tensors.
             kwargs (dict): Additional keyword arguments for the connector.
+
         """
         pass
 
@@ -181,8 +183,7 @@ def start_load_caches(
     def save_caches(
         self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs
     ) -> None:
-        """
-        Save the encoder cache to the connector.
+        """Save the encoder cache to the connector.
 
         This method saves the encoder cache from the worker's local storage
         to shared storage or another external connector.
@@ -192,14 +193,14 @@ def save_caches(
                 data hashes (`mm_hash`) to encoder cache tensors.
             mm_hash (str): The hash of the multimodal data whose cache is being saved.
             kwargs (dict): Additional keyword arguments for the connector.
+
         """
         pass
 
     def get_finished(
         self, finished_req_ids: set[str]
     ) -> tuple[set[str] | None, set[str] | None]:
-        """
-        Notifies worker-side connector ids of requests that have
+        """Notifies worker-side connector ids of requests that have
         finished generating tokens on the worker.
         The scheduler process (via the Executors) will use this output
         to track which workers are done.
@@ -210,16 +211,27 @@ def get_finished(
             tuple of (sending/saving ids, recving/loading ids).
             The finished saves/sends req ids must belong to a set provided in a
             call to this method (this call or a prior one).
+
         """
         return None, None
 
     def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None:
-        """
-        Build the ECConnector worker metadata for this engine step.
+        """Build the ECConnector worker metadata for this engine step.
 
         Returns:
             ECConnectorWorkerMetadata: the worker metadata.
             None if no worker metadata is available.
+
+        """
+        return None
+
+    def get_ec_connector_stats(self) -> "ECConnectorStats | None":
+        """Get the EC connector stats collected during the last interval.
+
+        Callable on either a worker-side or scheduler-side connector
+        instance: the worker-side instance reports stats gathered during
+        cache save/load, while the scheduler-side instance reports stats
+        gathered during scheduling decisions.
         """
         return None
 
@@ -232,8 +244,7 @@ def has_cache_item(
         self,
         identifier: str,
     ) -> bool:
-        """
-        Check if a single encoder cache exists
+        """Check if a single encoder cache exists.
 
         Args:
             identifier (str): the identifier of the media.
@@ -241,14 +252,14 @@ def has_cache_item(
         Returns:
             A bool where value is True if cache exist for
             the media
+
         """
         pass
 
     def ensure_cache_available(
         self, request: "Request", num_computed_tokens: int
     ) -> bool:
-        """
-        Ensure encoder cache items are available for the given request.
+        """Ensure encoder cache items are available for the given request.
         May initiate asynchronous transfers for items not yet local.
 
         Args:
@@ -258,6 +269,7 @@ def ensure_cache_available(
         Returns:
             True if all items are ready or no transfer is needed.
             False if any items are still in transit (request should be deferred).
+
         """
         return True
 
@@ -282,11 +294,12 @@ def update_state_after_free(self, request: "Request", index: int):
 
     @abstractmethod
     def update_state_after_alloc(self, request: "Request", index: int):
-        """
-        Update ECConnector state to decide allocate cache for requests
+        """Update ECConnector state to decide allocate cache for requests.
 
         Args:
             request (Request): the request object.
+            index (int): index of the encoder cache entry being allocated.
+
         """
         pass
 
@@ -294,37 +307,37 @@ def update_state_after_alloc(self, request: "Request", index: int):
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
     ) -> ECConnectorMetadata:
-        """
-        Build the connector metadata for this step.
+        """Build the connector metadata for this step.
 
         This function should NOT modify fields in the scheduler_output.
         Also, calling this function will reset the state of the connector.
 
         Args:
             scheduler_output (SchedulerOutput): the scheduler output object.
+
         """
         pass
 
     def update_connector_output(self, connector_output: ECConnectorOutput):
-        """
-        Update ECConnector state from worker-side connectors output.
+        """Update ECConnector state from worker-side connectors output.
 
         Args:
             connector_output (ECConnectorOutput): the worker-side
                 connectors output.
+
         """
         return
 
     def request_finished(
         self, request: "Request"
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Called when a request has finished, before its encoder cache is freed.
+        """Called when a request has finished, before its encoder cache is freed.
 
         Returns:
             True if the request is being saved/sent asynchronously and cached
             should not be freed until the request_id is returned from
             get_finished().
+
         """
         return False, None
 
@@ -339,3 +352,27 @@ def has_pending_push_work(self) -> bool:
         leave this as False.
         """
         return False
+
+    @classmethod
+    def build_ec_connector_stats(
+        cls, data: dict[str, Any] | None = None
+    ) -> "ECConnectorStats | None":
+        """ECConnectorStats resolution method. This method allows dynamically
+        registered connectors to return their own ECConnectorStats object,
+        which can implement custom aggregation logic on the data dict.
+        """
+        return None
+
+    @classmethod
+    def build_prom_metrics(
+        cls,
+        vllm_config: "VllmConfig",
+        metric_types: dict[type["PromMetric"], type["PromMetricT"]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ) -> "ECConnectorPromMetrics | None":
+        """Create an ECConnectorPromMetrics subclass which should register
+        per-connector Prometheus metrics and implement observe() to
+        expose connector transfer stats via Prometheus.
+        """
+        return None
diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py
index e2c37b72a6be..eff16e08b0ae 100644
--- a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py
+++ b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py
@@ -99,8 +99,8 @@ def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion:
     ec_config = vllm_config.ec_transfer_config
     assert ec_config is not None, "ec_transfer_config required to build region"
 
-    dp_rank = vllm_config.parallel_config.data_parallel_rank
-    engine_id = f"{vllm_config.instance_id}_dp{dp_rank}"
+    dp_index = vllm_config.parallel_config.data_parallel_index
+    engine_id = f"{vllm_config.instance_id}_dp{dp_index}"
 
     dtype = vllm_config.model_config.dtype
     hidden_dim = _get_encoder_cache_hidden_dim(vllm_config)
diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py b/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py
index 413bf285712a..bfc8aa009e15 100644
--- a/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py
+++ b/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Lightweight mmap-backed shared memory region for encoder cache (EC) data.
+"""Lightweight mmap-backed shared memory region for encoder cache (EC) data.
 
 Modeled after SharedOffloadRegion (vllm/v1/kv_offload/cpu/) but simplified
 for EC: flat shared layout, no multi-tensor cursor, no block_size_factor.
diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
index cf0a92dcc26c..f1d22a5714a6 100644
--- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
+++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
@@ -229,6 +229,7 @@ def _nixl_consumer_admit(
         Returns:
             True if no item needs to wait for a remote encoding. False if any
             item is pending or the request must be aborted.
+
         """
         if not request.ec_transfer_params:
             return True
@@ -259,6 +260,7 @@ def _admit_item(
         Raises:
             _RemoteUnavailable: The remote source is invalid, the read failed,
                 or the item's total wait budget expired.
+
         """
         pos = feature.mm_position
         mm_hash = feature.identifier
@@ -334,6 +336,7 @@ def _fail_or_fallback(
             item's remote announcement so subsequent steps do not retry it.
             False if only placeholder metadata (or no input) is available,
             recording the request for take_unavailable_requests() to drain.
+
         """
         mm_hash = feature.identifier
         self._deferred_since.pop((request.request_id, mm_hash), None)
diff --git a/vllm/distributed/ec_transfer/ec_connector/example_connector.py b/vllm/distributed/ec_transfer/ec_connector/example_connector.py
index 2db12f903c71..e080f09b6bef 100644
--- a/vllm/distributed/ec_transfer/ec_connector/example_connector.py
+++ b/vllm/distributed/ec_transfer/ec_connector/example_connector.py
@@ -67,8 +67,7 @@ def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole):
             raise ValueError("ec_transfer_config must be set for ECConnectorBase")
 
     def start_load_caches(self, encoder_cache, **kwargs) -> None:
-        """
-        Start loading the cache from the connector into vLLM's encoder cache.
+        """Start loading the cache from the connector into vLLM's encoder cache.
 
         This method loads the encoder cache based on metadata provided by the scheduler.
         It is called before `_gather_mm_embeddings` for the EC Connector. For EC,
@@ -78,6 +77,7 @@ def start_load_caches(self, encoder_cache, **kwargs) -> None:
             encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
                 data hashes (`mm_hash`) to encoder cache tensors.
             kwargs (dict): Additional keyword arguments for the connector.
+
         """
         from vllm.platforms import current_platform
 
@@ -108,8 +108,7 @@ def start_load_caches(self, encoder_cache, **kwargs) -> None:
             logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash)
 
     def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None:
-        """
-        Save the encoder cache to the connector.
+        """Save the encoder cache to the connector.
 
         This method saves the encoder cache from the worker's local storage
         to shared storage or another external connector.
@@ -119,6 +118,7 @@ def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None:
                 data hashes (`mm_hash`) to encoder cache tensors.
             mm_hash (str): The hash of the multimodal data whose cache is being saved.
             kwargs (dict): Additional keyword arguments for the connector.
+
         """
         # Return if it is PD Instance
         if not self.is_producer:
@@ -134,14 +134,14 @@ def has_cache_item(
         self,
         identifier: str,
     ) -> bool:
-        """
-        Check if cache exist externally for the media
+        """Check if cache exist externally for the media.
 
         Args:
             identifier (str): the identifier of the media.
 
         Returns:
             Bool indicate that media exists in cache or not
+
         """
         return self._found_match_for_mm_data(identifier)
 
@@ -150,9 +150,7 @@ def update_state_after_alloc(
         request: "Request",
         index: int,
     ) -> None:
-        """
-        Update ECConnector state after encoder cache allocation.
-        """
+        """Update ECConnector state after encoder cache allocation."""
         mm_hash = request.mm_features[index].identifier
         # Only load cache if it is consumer and cache exists
         if not self.is_consumer or not self.has_cache_item(mm_hash):
@@ -212,8 +210,7 @@ def _generate_foldername_debug(
         mm_hash: str,
         create_folder: bool = True,  # <- now defaults to True
     ) -> str:
-        """
-        Return the folder in which the cache for this mm_hash lives.
+        """Return the folder in which the cache for this mm_hash lives.
         If `create_folder` is True (default) the directory is created
         recursively the first time it is needed.
         """
@@ -223,8 +220,7 @@ def _generate_foldername_debug(
         return foldername
 
     def _generate_filename_debug(self, mm_hash: str) -> str:
-        """
-        Return the full path of the safetensors file for this mm_hash.
+        """Return the full path of the safetensors file for this mm_hash.
         Ensures the parent directory exists because
         `_generate_foldername_debug` is called with its default
         (`create_folder=True`).
diff --git a/vllm/distributed/ec_transfer/ec_connector/metrics.py b/vllm/distributed/ec_transfer/ec_connector/metrics.py
new file mode 100644
index 000000000000..0eb20b9863ae
--- /dev/null
+++ b/vllm/distributed/ec_transfer/ec_connector/metrics.py
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, field
+from typing import Any, TypeAlias, TypeVar
+
+from prometheus_client import Counter, Gauge, Histogram
+
+from vllm.config import ECTransferConfig, VllmConfig
+from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory
+from vllm.logger import init_logger
+
+PromMetric: TypeAlias = Gauge | Counter | Histogram
+PromMetricT = TypeVar("PromMetricT", bound=PromMetric)
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class ECConnectorStats:
+    """Base class for EC Connector Stats, a container for transfer performance
+    metrics or otherwise important telemetry from the connector.
+    All sub-classes need to be serializable as stats are sent from worker to
+    logger process.
+    """
+
+    data: dict[str, Any] = field(default_factory=dict)
+
+    def reset(self):
+        """Reset the stats, clear the state."""
+        raise NotImplementedError
+
+    def aggregate(self, other: "ECConnectorStats") -> "ECConnectorStats":
+        """Aggregate stats with another `ECConnectorStats` object."""
+        raise NotImplementedError
+
+    def reduce(self) -> dict[str, int | float]:
+        """Reduce the observations collected during a time interval to one or
+        more representative values (eg avg/median/sum of the series).
+        This is meant to be called by the logger to produce a summary of the
+        stats for the last time interval.
+        """
+        raise NotImplementedError
+
+    def is_empty(self) -> bool:
+        """Return True if the stats are empty."""
+        raise NotImplementedError
+
+
+class ECConnectorLogging:
+    def __init__(self, ec_transfer_config: ECTransferConfig | None):
+        # Instantiate the connector's stats class.
+        if ec_transfer_config and ec_transfer_config.ec_connector:
+            self.connector_cls = ECConnectorFactory.get_connector_class(
+                ec_transfer_config
+            )
+        self.reset()
+
+    def reset(self):
+        self.transfer_stats_accumulator: ECConnectorStats | None = None
+
+    def observe(self, transfer_stats_data: dict[str, Any]):
+        # Should not be called when an ECConnector is not configured.
+        assert self.connector_cls is not None
+        # Called periodically when connector syncs with the scheduler.
+        # Note that this is not the same as the logging interval.
+        # We expect transfer_stats_data to be aggregated across all workers and
+        # consist of observations from a single connector.
+        transfer_stats = self.connector_cls.build_ec_connector_stats(
+            transfer_stats_data
+        )
+        if transfer_stats is None:
+            logger.warning_once(
+                "The connector %s is collecting stats but "
+                "does not implement the "
+                "`build_ec_connector_stats` method. "
+                "Stats will not be logged.",
+                self.connector_cls,
+            )
+            return
+
+        if self.transfer_stats_accumulator is None:
+            self.transfer_stats_accumulator = transfer_stats
+        else:
+            # Accumulate last interval stats.
+            self.transfer_stats_accumulator = self.transfer_stats_accumulator.aggregate(
+                transfer_stats
+            )
+
+    def log(self, log_fn=logger.info):
+        """Log transfer metrics periodically, similar to throughput logging"""
+        if (
+            self.transfer_stats_accumulator
+            and not self.transfer_stats_accumulator.is_empty()
+        ):
+            # Produce a single cumulative stats object for the last time
+            # interval from the recorded observations.
+            xfer_metrics = self.transfer_stats_accumulator.reduce()
+            xfer_metrics_str = ", ".join(f"{k}={v}" for k, v in xfer_metrics.items())
+            log_fn("EC Transfer metrics: %s", xfer_metrics_str)
+
+            # Reset metrics for next interval
+            self.reset()
+
+
+class ECConnectorPromMetrics:
+    """A base class for per-connector Prometheus metric registration
+    and recording.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        metric_types: dict[type[PromMetric], type[PromMetricT]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ):
+        self._ec_transfer_config = vllm_config.ec_transfer_config
+        self._gauge_cls = metric_types[Gauge]
+        self._counter_cls = metric_types[Counter]
+        self._histogram_cls = metric_types[Histogram]
+        self._labelnames = labelnames
+        self.per_engine_labelvalues = per_engine_labelvalues
+
+    def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
+        """Record the supplied transfer statistics to Prometheus metrics. These
+        statistics are engine-specific, and should be recorded to a metric
+        with the appropriate 'engine' label. These metric instances can be
+        created using the create_metric_per_engine() helper method.
+        """
+        raise NotImplementedError
+
+
+class ECConnectorProm:
+    """Support for registering per-connector Prometheus metrics, and
+    recording transfer statistics to those metrics. Uses
+    ECConnectorBase.build_prom_metrics().
+    """
+
+    _gauge_cls = Gauge
+    _counter_cls = Counter
+    _histogram_cls = Histogram
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ):
+        self.prom_metrics: ECConnectorPromMetrics | None = None
+        ec_transfer_config = vllm_config.ec_transfer_config
+        if ec_transfer_config and ec_transfer_config.ec_connector:
+            connector_cls = ECConnectorFactory.get_connector_class(ec_transfer_config)
+            metric_types = {
+                Gauge: self._gauge_cls,
+                Counter: self._counter_cls,
+                Histogram: self._histogram_cls,
+            }
+            self.prom_metrics = connector_cls.build_prom_metrics(
+                vllm_config,
+                metric_types,
+                labelnames,
+                per_engine_labelvalues,
+            )
+
+    def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
+        if self.prom_metrics is None:
+            return
+        self.prom_metrics.observe(transfer_stats_data, engine_idx)
diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py
index 10c9acae4676..4cc103a56e06 100644
--- a/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py
+++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/config.py
@@ -69,6 +69,10 @@ class MooncakeECConfig:
     control_timeout_ms: int
     push_wait_timeout_s: float
     pool_size: int
+    cross_encoder_cache: bool = False
+    store_max_pending_items: int = 32
+    store_max_pending_bytes: int = 2 * 1024**3
+    store_read_buffer_bytes: int = 128 * 1024**2
 
     @classmethod
     def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig:
@@ -94,6 +98,24 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig:
             "ec_buffer_size", ec_config.ec_buffer_size
         )
         get = ec_config.get_from_extra_config
+        shared_reuse = get("cross_encoder_cache", False)
+        if not isinstance(shared_reuse, bool):
+            raise ValueError("cross_encoder_cache must be a boolean")
+        if shared_reuse:
+            if not ec_config.is_encode_only:
+                raise ValueError(
+                    "cross_encoder_cache requires an Encoder-only producer"
+                )
+            if not vllm_config.use_v2_model_runner:
+                raise ValueError("cross_encoder_cache requires Model Runner V2")
+            if vllm_config.lora_config is not None:
+                raise ValueError("cross_encoder_cache does not support LoRA")
+            mm_config = vllm_config.model_config.multimodal_config
+            if mm_config is not None and mm_config.mm_processor_cache_gb == 0:
+                raise ValueError(
+                    "cross_encoder_cache requires mm_processor_cache_gb > 0 "
+                    "to preserve content identifiers across Encoders and restarts"
+                )
         control_port = int(ec_config.ec_port) + (
             parallel_config.data_parallel_index * parallel_config.tensor_parallel_size
         )
@@ -120,4 +142,14 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig:
                 "push_wait_timeout_s", get("push_wait_timeout_s", 60)
             ),
             pool_size=registered_buffer_size,
+            cross_encoder_cache=shared_reuse,
+            store_max_pending_items=_positive_int(
+                "store_max_pending_items", get("store_max_pending_items", 32)
+            ),
+            store_max_pending_bytes=_positive_int(
+                "store_max_pending_bytes", get("store_max_pending_bytes", 2 * 1024**3)
+            ),
+            store_read_buffer_bytes=_positive_int(
+                "store_read_buffer_bytes", get("store_read_buffer_bytes", 128 * 1024**2)
+            ),
         )
diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py
index a006d1bf2f80..76e714dc044f 100644
--- a/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py
+++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py
@@ -10,6 +10,9 @@
     ECConnectorMetadata,
     ECConnectorWorkerMetadata,
 )
+from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
+    TensorSpec,
+)
 
 
 @dataclass
@@ -46,6 +49,7 @@ class ECMooncakeConnectorMetadata(ECConnectorMetadata):
     loads: list[ECMooncakeLoadSpec] = field(default_factory=list)
     pushes: list[ECMooncakePushSpec] = field(default_factory=list)
     freed: list[str] | None = None
+    store_candidates: dict[str, TensorSpec] = field(default_factory=dict)
 
 
 @dataclass
diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py
index 0b52dbe9c331..2a607118db45 100644
--- a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py
+++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py
@@ -45,6 +45,9 @@
 from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import (
     ensure_mooncake_available,
 )
+from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
+    TensorSpec,
+)
 from vllm.distributed.ec_transfer.ec_connector.utils import (
     PlaceholderMetadataResolver,
     collect_ec_item_metadata,
@@ -76,6 +79,8 @@ def __init__(self, vllm_config: VllmConfig) -> None:
         ensure_mooncake_available()
         config = MooncakeECConfig.from_vllm_config(vllm_config)
 
+        self._store_candidates: list[tuple[str, str, TensorSpec]] = []
+        self._cross_encoder_cache = config.cross_encoder_cache
         self._is_producer = config.is_producer
         self._is_consumer = config.is_consumer
         self._control_addr = config.control_addr
@@ -427,28 +432,40 @@ def _prepare_push_spec(self, request: Any, index: int) -> None:
             transfer_id = f"{request.request_id}:{index}"
         if not consumer_zmq or transfer_id in self._prepared_push_transfer_ids:
             return
-        num_tokens = request.get_num_encoder_embeds(index)
-        dtype = self._model_config.dtype
-        assert isinstance(dtype, torch.dtype)
-        assert self._encoder_cache_hidden_dim is not None
-        dtype_name = str(dtype).split(".")[-1]
-        shape = (num_tokens, self._encoder_cache_hidden_dim)
-        nbytes = math.prod(shape) * dtype.itemsize
+        output = self._encoder_output_spec(request, index)
         self._pushes_to_prepare[transfer_id] = ECMooncakePushSpec(
             mm_hash=mm_hash,
-            nbytes=nbytes,
-            shape=shape,
-            dtype=dtype_name,
+            nbytes=output.nbytes,
+            shape=output.shape,
+            dtype=output.dtype.removeprefix("torch."),
             consumer_zmq=str(consumer_zmq),
             transfer_id=transfer_id,
             request_id=request.request_id,
         )
         self._prepared_push_transfer_ids.add(transfer_id)
 
+    def _encoder_output_spec(self, request: Any, index: int) -> TensorSpec:
+        dtype = self._model_config.dtype
+        assert isinstance(dtype, torch.dtype)
+        assert self._encoder_cache_hidden_dim is not None
+        shape = (request.get_num_encoder_embeds(index), self._encoder_cache_hidden_dim)
+        return TensorSpec(shape, str(dtype), math.prod(shape) * dtype.itemsize)
+
     def update_state_after_alloc(self, request: Any, index: int) -> None:
         self._local_cache.add(request.mm_features[index].identifier)
         if self._is_producer:
             self._prepare_push_spec(request, index)
+            if (
+                self._cross_encoder_cache
+                and request.mm_features[index].modality == "image"
+            ):
+                self._store_candidates.append(
+                    (
+                        request.request_id,
+                        request.mm_features[index].identifier,
+                        self._encoder_output_spec(request, index),
+                    )
+                )
 
     def update_state_after_free(self, request: Any, index: int) -> None:
         if not self._is_consumer:
@@ -474,6 +491,13 @@ def build_connector_meta(
         meta = ECMooncakeConnectorMetadata(
             freed=scheduler_output.free_encoder_mm_hashes
         )
+        preempted = scheduler_output.preempted_req_ids or set()
+        meta.store_candidates = {
+            identifier: output
+            for request_id, identifier, output in self._store_candidates
+            if request_id not in preempted
+        }
+        self._store_candidates.clear()
         for push_spec in self._pushes_to_prepare.values():
             meta.pushes.append(push_spec)
         self._pushes_to_prepare.clear()
diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py
index 05ce16b181e6..04c95f285498 100644
--- a/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py
+++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py
@@ -49,6 +49,9 @@
     MooncakeTransfer,
     ensure_mooncake_available,
 )
+from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
+    build_embedding_namespace,
+)
 from vllm.logger import init_logger
 from vllm.utils.network_utils import get_ip
 
@@ -59,6 +62,10 @@
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
 
+    from ..mooncake_store_embedding.backend import (
+        MooncakeEmbeddingStoreBackend,
+    )
+
 _RESERVATION_REFRESH_SECONDS = _RESERVATION_TTL_SECONDS / 2
 _MAX_CANCELLED_TRANSFER_IDS = 1 << 16
 _CANCEL_ATTEMPTS = 2
@@ -87,6 +94,17 @@ class ECMooncakeWorker:
     def __init__(self, vllm_config: VllmConfig) -> None:
         ensure_mooncake_available()
         config = MooncakeECConfig.from_vllm_config(vllm_config)
+        self._store_config = (
+            (
+                build_embedding_namespace(vllm_config),
+                config.store_max_pending_items,
+                config.store_max_pending_bytes,
+                config.store_read_buffer_bytes,
+            )
+            if config.cross_encoder_cache
+            else None
+        )
+        self._output_store: MooncakeEmbeddingStoreBackend | None = None
         self.is_producer = config.is_producer
         self.is_consumer = config.is_consumer
         self._buffer_device = config.buffer_device
@@ -173,6 +191,21 @@ def _resolve_consumer_rank(self) -> None:
             self._is_receiving_rank = True
 
     def start_services(self) -> None:
+        if self._store_config is not None and self._output_store is None:
+            from ..mooncake_store_embedding.backend import (
+                MooncakeEmbeddingStoreBackend,
+            )
+            from ..mooncake_store_embedding.store_client import (
+                create_mooncake_embedding_store_client,
+            )
+
+            namespace, max_items, max_bytes, read_bytes = self._store_config
+            self._output_store = MooncakeEmbeddingStoreBackend(
+                create_mooncake_embedding_store_client(read_buffer_bytes=read_bytes),
+                namespace,
+                max_pending_items=max_items,
+                max_pending_bytes=max_bytes,
+            )
         if not self.is_consumer or self._control_server is not None:
             return
         self._resolve_consumer_rank()
@@ -470,6 +503,14 @@ def start_save_caches(
             tensor = encoder_cache.get(mm_hash)
             if tensor is not None:
                 self._bind_push_source(tensor, mm_hash)
+        if self._output_store is not None:
+            loaded = self._output_store.resolve_inputs(
+                metadata.store_candidates, encoder_cache, self._buffer_device
+            )
+            # Loaded entries precede the runner's new-output snapshot, so they
+            # need explicit P2P source binding here.
+            for mm_hash in loaded:
+                self._bind_push_source(encoder_cache[mm_hash], mm_hash)
 
     def _reserve_batch(self, records: list[ProducerPushRecord]) -> None:
         """Resolve independent item futures with one reserve RPC per shard."""
@@ -851,6 +892,8 @@ def get_finished(
         if not self.is_producer:
             return None, None
 
+        if self._output_store is not None:
+            self._output_store.reap()
         for record in self._producer_pushes.cancel_requests(finished_req_ids):
             self._producer_pushes.submit_cancel(
                 record,
@@ -866,6 +909,8 @@ def save_caches(
             return
         tensor = encoder_cache[mm_hash]
         self._bind_push_source(tensor, mm_hash)
+        if self._output_store is not None:
+            self._output_store.save_output(mm_hash, tensor)
 
     def build_connector_worker_meta(self) -> ECMooncakeWorkerMetadata | None:
         if self.is_consumer and not self._is_receiving_rank:
@@ -938,3 +983,6 @@ def close(self) -> None:
             )
         self._producer_memory.close()
         self._transfer.close()
+        if self._output_store is not None:
+            self._output_store.shutdown()
+            self._output_store = None
diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py
new file mode 100644
index 000000000000..42ed7690a7e8
--- /dev/null
+++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Mooncake Store support for cross-Encoder output reuse."""
diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py
new file mode 100644
index 000000000000..5ad5ada882d8
--- /dev/null
+++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py
@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Encoder-side Mooncake Store resolver and publisher."""
+
+from __future__ import annotations
+
+import threading
+import traceback
+from collections.abc import Mapping
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+
+import torch
+
+from vllm.logger import init_logger
+
+from .data import TensorSpec, make_embedding_key
+from .store_client import (
+    EmbeddingStoreError,
+    EmbeddingStoreOperationError,
+    MooncakeEmbeddingStoreClient,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _StoreSave:
+    tensor: torch.Tensor | None
+    ready_event: torch.Event | None
+    storage_ptr: int
+
+    def release(self) -> None:
+        self.tensor = None
+        self.ready_event = None
+
+
+class MooncakeEmbeddingStoreBackend:
+    """Resolve immutable encoder outputs and publish misses asynchronously."""
+
+    def __init__(
+        self,
+        store_client: MooncakeEmbeddingStoreClient,
+        namespace: str,
+        *,
+        max_pending_items: int,
+        max_pending_bytes: int,
+    ) -> None:
+        self.store_client = store_client
+        self.namespace = namespace
+        self.max_pending_items = max_pending_items
+        self.max_pending_bytes = max_pending_bytes
+        self._executor = ThreadPoolExecutor(
+            max_workers=1,
+            thread_name_prefix="ec-mooncake-store",
+        )
+        self._pending: dict[str, tuple[Future[None], _StoreSave]] = {}
+        self._pending_storages: dict[int, tuple[torch.UntypedStorage, int]] = {}
+        self._pending_lock = threading.Lock()
+        self._pending_bytes = 0
+        self._step_candidates: set[str] = set()
+        self._closed = False
+        self._fatal_error: BaseException | None = None
+
+    def save_output(self, identifier: str, tensor: torch.Tensor) -> None:
+        if identifier in self._step_candidates:
+            self.reap()
+            self._enqueue_save(identifier, tensor)
+
+    def resolve_inputs(
+        self,
+        expected_tensors: Mapping[str, TensorSpec],
+        encoder_cache: dict[str, torch.Tensor],
+        device: torch.device | str,
+    ) -> set[str]:
+        """Load outputs before encoding; unresolved items retain normal computation."""
+        self.reap()
+        self._step_candidates = set(expected_tensors)
+        candidates = [key for key in expected_tensors if key not in encoder_cache]
+        if not candidates:
+            return set()
+
+        pool_keys = [
+            make_embedding_key(self.namespace, identifier) for identifier in candidates
+        ]
+        try:
+            exists = self.store_client.batch_exists(pool_keys)
+        except (EmbeddingStoreOperationError, OSError):
+            logger.warning(
+                "Mooncake embedding Store lookup failed; falling back to encoder",
+                exc_info=True,
+            )
+            return set()
+
+        hits = {
+            pool_key: expected_tensors[identifier]
+            for identifier, pool_key, hit in zip(
+                candidates, pool_keys, exists, strict=True
+            )
+            if hit
+        }
+        try:
+            tensors = self.store_client.load_tensors(hits, device)
+        except (EmbeddingStoreOperationError, OSError):
+            logger.warning(
+                "Mooncake embedding Store GET failed; falling back to encoder",
+                exc_info=True,
+            )
+            return set()
+        loaded: set[str] = set()
+        for identifier, pool_key in zip(candidates, pool_keys, strict=True):
+            if pool_key in tensors:
+                encoder_cache[identifier] = tensors[pool_key]
+                loaded.add(identifier)
+        return loaded
+
+    def _enqueue_save(self, identifier: str, tensor: torch.Tensor) -> bool:
+        """Admit contiguous outputs, charging each retained storage once."""
+        if not tensor.is_contiguous():
+            return False
+        storage = tensor.untyped_storage()
+        storage_ptr = storage.data_ptr()
+        with self._pending_lock:
+            if self._fatal_error is not None:
+                raise EmbeddingStoreError(
+                    "Store publisher has failed"
+                ) from self._fatal_error
+            if self._closed or identifier in self._pending:
+                return False
+            _, ref_count = self._pending_storages.get(storage_ptr, (storage, 0))
+            budget_bytes = 0 if ref_count else storage.nbytes()
+            if (
+                len(self._pending) >= self.max_pending_items
+                or self._pending_bytes + budget_bytes > self.max_pending_bytes
+            ):
+                return False
+            # Only admitted work creates an event. Runtime errors propagate.
+            save = _StoreSave(
+                tensor,
+                _record_tensor_ready_event(tensor),
+                storage_ptr,
+            )
+            future = self._executor.submit(
+                self._save, make_embedding_key(self.namespace, identifier), save
+            )
+            self._pending[identifier] = (future, save)
+            # Keep the storage alive until reaping so its address cannot be reused.
+            self._pending_storages[storage_ptr] = (storage, ref_count + 1)
+            self._pending_bytes += budget_bytes
+        return True
+
+    def _save(self, pool_key: str, save: _StoreSave) -> None:
+        try:
+            if self.store_client.exists(pool_key):
+                return
+            tensor = save.tensor
+            assert tensor is not None
+            if save.ready_event is not None:
+                save.ready_event.synchronize()
+            self.store_client.put_tensor(pool_key, tensor)
+        except BaseException as error:
+            # Returned client frames can otherwise keep tensors via a Future's
+            # traceback. Uncertain I/O owners remain retained by the client.
+            traceback.clear_frames(error.__traceback__)
+            raise
+        finally:
+            tensor = None
+            save.release()
+
+    def reap(self) -> None:
+        with self._pending_lock:
+            for identifier, (future, save) in list(self._pending.items()):
+                # Observe once; a racing completion is handled on the next poll.
+                if not future.done():
+                    continue
+                try:
+                    future.result()  # Already done: never waits on Store I/O.
+                except (EmbeddingStoreOperationError, OSError) as error:
+                    logger.warning("Store PUT failed for %s: %s", identifier, error)
+                    # This failure is consumed; do not retain publisher/caller frames.
+                    error.__traceback__ = None
+                except BaseException as error:
+                    # The client may still own unsafe native-I/O buffers. Keep
+                    # their charge and fatal result until worker termination.
+                    if self._fatal_error is None:
+                        self._fatal_error = error
+                    continue
+                del self._pending[identifier]
+                storage, ref_count = self._pending_storages[save.storage_ptr]
+                if ref_count == 1:
+                    del self._pending_storages[save.storage_ptr]
+                    self._pending_bytes -= storage.nbytes()
+                else:
+                    self._pending_storages[save.storage_ptr] = (storage, ref_count - 1)
+            fatal = self._fatal_error
+        if fatal is not None:
+            raise fatal
+
+    def shutdown(self) -> None:
+        with self._pending_lock:
+            if self._closed:
+                return
+            self._closed = True
+        self._executor.shutdown(wait=True, cancel_futures=False)
+        self.reap()
+        self.store_client.close()
+
+
+def _record_tensor_ready_event(tensor: torch.Tensor) -> torch.Event | None:
+    if tensor.device.type != "cuda":
+        return None
+    event = torch.Event()
+    event.record(torch.accelerator.current_stream(tensor.device))
+    return event
diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py
new file mode 100644
index 000000000000..9a7d7dd002ce
--- /dev/null
+++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Content identities and output contracts for the shared Encoder Store."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+EMBEDDING_CACHE_KEY_VERSION = "v3"
+
+
+def make_embedding_key(namespace: str, identifier: str) -> str:
+    return f"{namespace}@id:{identifier}"
+
+
+@dataclass(frozen=True)
+class TensorSpec:
+    """Contiguous output contract used by the Encoder and Store codec."""
+
+    shape: tuple[int, ...]
+    dtype: str
+    nbytes: int
+
+
+def build_embedding_namespace(vllm_config: VllmConfig) -> str:
+    model_config = vllm_config.model_config
+    ec_config = vllm_config.ec_transfer_config
+    assert ec_config is not None
+    extra_config = ec_config.ec_connector_extra_config
+
+    multimodal_config = model_config.multimodal_config
+    encoder_hash = (
+        multimodal_config.compute_hash()
+        if multimodal_config is not None
+        else "encoder:default"
+    )
+    cache_prefix = str(extra_config.get("embedding_cache_prefix", ""))
+    model_name = str(extra_config.get("embedding_model_identity", model_config.model))
+    prefix = f"{cache_prefix}@" if cache_prefix else ""
+    return (
+        f"{prefix}embedding"
+        f"@model:{model_name}"
+        f"@revision:{model_config.revision or 'default'}"
+        f"@encoder:{encoder_hash}"
+        f"@dtype:{model_config.dtype}"
+        f"@protocol:{EMBEDDING_CACHE_KEY_VERSION}"
+    )
diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py
new file mode 100644
index 000000000000..4d3ed6739d0e
--- /dev/null
+++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py
@@ -0,0 +1,463 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Thin Mooncake Store client for embedding objects.
+
+Originally adapted from vLLM PR #47302 (0d1f71f5d36c). The v3 key namespace
+uses a compact header for complete 2D embeddings, not that PR's tensor format.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import struct
+import threading
+from collections.abc import Callable, Iterator, Mapping
+from contextlib import contextmanager
+from enum import IntEnum
+from typing import Any
+
+import torch
+
+from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
+    TensorSpec,
+)
+from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store
+from vllm.logger import init_logger
+from vllm.utils.network_utils import get_ip
+
+logger = init_logger(__name__)
+
+# A failed transfer can still use registered memory after the binding returns.
+# Poisoned clients retain their owners until process exit, including if normal
+# worker teardown drops the connector. New operations are rejected.
+_UNSAFE_CLIENTS: list[Any] = []
+
+
+# Values from Mooncake v0.3.12 / v0.3.12.post1 types.h (ErrorCode).
+# Python exposes integer returns, not this enum.
+class _MooncakeErrorCode(IntEnum):
+    NO_AVAILABLE_HANDLE = -200
+    INVALID_PARAMS = -600
+    INVALID_REPLICA = -702
+    REPLICA_IS_NOT_READY = -703
+    OBJECT_NOT_FOUND = -704
+    OBJECT_ALREADY_EXISTS = -705
+    LEASE_EXPIRED = -707
+    RPC_FAIL = -900
+    RPC_TIMEOUT = -901
+    TENANT_QUOTA_EXCEEDED = -1700
+
+
+# For RAM batch reads and single-replica writes, these rejections occur
+# before submission or after completed I/O. Transfer timeouts can leave I/O
+# in flight; all other failures retain registered buffers and poison the client.
+_SAFE_IO_REJECTIONS = frozenset(
+    {
+        _MooncakeErrorCode.NO_AVAILABLE_HANDLE,
+        _MooncakeErrorCode.REPLICA_IS_NOT_READY,
+        _MooncakeErrorCode.OBJECT_NOT_FOUND,
+        _MooncakeErrorCode.OBJECT_ALREADY_EXISTS,
+        _MooncakeErrorCode.RPC_FAIL,
+        _MooncakeErrorCode.RPC_TIMEOUT,
+    }
+)
+# Batch GET validates capacity/replicas before I/O and checks leases afterwards.
+# PUT rejects tenant quotas before allocating replicas or submitting transfers.
+_SAFE_GET_REJECTIONS = _SAFE_IO_REJECTIONS | {
+    _MooncakeErrorCode.INVALID_PARAMS,
+    _MooncakeErrorCode.INVALID_REPLICA,
+    _MooncakeErrorCode.LEASE_EXPIRED,
+}
+_SAFE_PUT_REJECTIONS = _SAFE_IO_REJECTIONS | {_MooncakeErrorCode.TENANT_QUOTA_EXCEEDED}
+
+
+_MOONCAKE_TENSOR_OBJECT_MAGIC = 0x4D4F4F4E
+# Magic, dtype, token count, hidden dimension. Format versioning is in the key.
+_MOONCAKE_TENSOR_HEADER = struct.Struct("<IIQQ")
+_MOONCAKE_DTYPE_TO_TORCH_DTYPE = {
+    0: "torch.float32",
+    11: "torch.float16",
+    12: "torch.bfloat16",
+}
+_TORCH_DTYPE_TO_MOONCAKE_DTYPE = {
+    value: key for key, value in _MOONCAKE_DTYPE_TO_TORCH_DTYPE.items()
+}
+
+
+def create_mooncake_embedding_store_client(
+    read_buffer_bytes: int = 128 * 1024**2,
+) -> MooncakeEmbeddingStoreClient:
+    try:
+        from mooncake.store import (  # type: ignore
+            MooncakeDistributedStore,
+            ObjectDataType,
+            ReplicateConfig,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "Install mooncake with Store support to enable cross_encoder_cache."
+        ) from e
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils
+
+    config = MooncakeStoreConfig.load_from_config()
+    if config.enable_offload:
+        raise ValueError(
+            "cross_encoder_cache supports a RAM Store; disable enable_offload"
+        )
+    store = MooncakeDistributedStore()
+    local_ip = get_ip()
+    local_hostname = rdma_utils.get_requester_local_hostname(local_ip)
+    setup_mooncake_store(store, config, local_hostname)
+
+    logger.info(
+        "Initialized embedding Mooncake store mode=%s global_segment_size=%d "
+        "local_buffer_size=%d",
+        config.mode,
+        config.global_segment_size,
+        config.local_buffer_size,
+    )
+    replicate_config = ReplicateConfig()
+    replicate_config.data_type = ObjectDataType.TENSOR
+    return MooncakeEmbeddingStoreClient(
+        store, replicate_config=replicate_config, read_buffer_bytes=read_buffer_bytes
+    )
+
+
+class EmbeddingStoreError(RuntimeError):
+    """A Store failure that must propagate to the worker."""
+
+
+class EmbeddingStoreOperationError(EmbeddingStoreError):
+    """A rejected or completed operation that permits normal Encoder fallback."""
+
+
+def _batch_get_completed_safely(results: Any, capacities: list[int]) -> bool:
+    return (
+        isinstance(results, list)
+        and len(results) == len(capacities)
+        and all(
+            type(result) is int
+            and (0 <= result <= capacity or result in _SAFE_GET_REJECTIONS)
+            for result, capacity in zip(results, capacities, strict=True)
+        )
+    )
+
+
+class MooncakeEmbeddingStoreClient:
+    """Wraps Mooncake object and buffer APIs used by embedding transfer."""
+
+    def __init__(
+        self,
+        store: Any,
+        replicate_config: Any | None = None,
+        *,
+        read_buffer_bytes: int = 128 * 1024**2,
+    ):
+        self._lifetime_lock = threading.Lock()
+        self._unsafe_owners: list[Any] = []
+        self._poisoned = False
+        self.store = store
+        self.replicate_config = replicate_config
+        self._read_buffer_bytes = read_buffer_bytes
+        self._read_buffer: torch.Tensor | None = None
+        self._put_header: ctypes.Array[ctypes.c_char] | None = None
+
+    def _check_healthy(self) -> None:
+        if self._poisoned:
+            raise EmbeddingStoreError(
+                "Mooncake I/O completion is unconfirmed; worker must exit"
+            )
+
+    def _poison(self, owners: list[Any]) -> None:
+        with self._lifetime_lock:
+            self._unsafe_owners.extend(owners)
+            if not self._poisoned:
+                self._poisoned = True
+                _UNSAFE_CLIENTS.append(self)
+
+    @contextmanager
+    def _registered_io(
+        self,
+        buffers: list[tuple[Any, int, int]],
+    ) -> Iterator[Callable[[int], None]]:
+        self._check_healthy()
+        registered = []
+        submitted = False
+        confirmed = False
+
+        def confirm(result: int) -> None:
+            nonlocal confirmed
+            if type(result) is not int or (
+                result < 0 and result not in _SAFE_PUT_REJECTIONS
+            ):
+                raise EmbeddingStoreError(
+                    "Mooncake I/O completion is unconfirmed; retaining buffers"
+                )
+            confirmed = True
+
+        try:
+            for owner, addr, size in buffers:
+                ret = self.store.register_buffer(addr, size)
+                if ret != 0:
+                    raise EmbeddingStoreOperationError(
+                        f"Failed to register embedding buffer: {ret}"
+                    )
+                registered.append(addr)
+            submitted = True
+            yield confirm
+        except BaseException as error:
+            if submitted and not confirmed:
+                self._poison([owner for owner, _, _ in buffers])
+                raise EmbeddingStoreError(
+                    "Mooncake I/O completion is unconfirmed; retaining buffers"
+                ) from error
+            raise
+        finally:
+            if not submitted or confirmed:
+                try:
+                    for addr in reversed(registered):
+                        self._unregister_buffer(addr)
+                except BaseException:
+                    self._poison([owner for owner, _, _ in buffers])
+                    raise
+
+    def close(self) -> None:
+        """Close the supported Mooncake binding after all I/O is drained."""
+        self._check_healthy()
+        if self._read_buffer is not None:
+            try:
+                self._unregister_buffer(self._read_buffer.data_ptr())
+            except BaseException:
+                self._poison([self._read_buffer])
+                raise
+            self._read_buffer = None
+        if self._put_header is not None:
+            try:
+                self._unregister_buffer(ctypes.addressof(self._put_header))
+            except BaseException:
+                self._poison([self._put_header])
+                raise
+            self._put_header = None
+        ret = self.store.close()
+        if ret != 0:
+            raise EmbeddingStoreError(
+                f"failed to close embedding Mooncake Store: {ret}"
+            )
+
+    def exists(self, pool_key: str) -> bool:
+        return self.batch_exists([pool_key])[0]
+
+    def batch_exists(self, keys: list[str]) -> list[bool]:
+        self._check_healthy()
+        if not keys:
+            return []
+
+        states = self.store.batch_is_exist(keys)
+        if len(states) != len(keys):
+            raise RuntimeError(
+                "Mooncake Store returned an unexpected number of lookup results"
+            )
+        return [state == 1 for state in states]
+
+    def load_tensors(
+        self, expected: Mapping[str, TensorSpec], device: torch.device | str
+    ) -> dict[str, torch.Tensor]:
+        """Synchronously load runner inputs into independently owned tensors."""
+        self._check_healthy()
+        device = torch.device(device)
+        loaded: dict[str, torch.Tensor] = {}
+        chunk: dict[str, TensorSpec] = {}
+        chunk_bytes = 0
+        for key, spec in expected.items():
+            size = _MOONCAKE_TENSOR_HEADER.size + spec.nbytes
+            if size > self._read_buffer_bytes:
+                logger.warning(
+                    "Skipping Store GET for %s: object bytes=%d exceed read buffer=%d",
+                    key,
+                    size,
+                    self._read_buffer_bytes,
+                )
+                continue
+            if chunk_bytes + size > self._read_buffer_bytes:
+                loaded.update(self._load_chunk(chunk, device))
+                chunk = {}
+                chunk_bytes = 0
+            chunk[key] = spec
+            chunk_bytes += size
+        if chunk:
+            loaded.update(self._load_chunk(chunk, device))
+        return loaded
+
+    def _load_chunk(
+        self, expected: dict[str, TensorSpec], device: torch.device
+    ) -> dict[str, torch.Tensor]:
+        self._check_healthy()
+        if self._read_buffer is None:
+            buffer = torch.empty(
+                self._read_buffer_bytes,
+                dtype=torch.uint8,
+                device="cpu",
+                pin_memory=device.type == "cuda",
+            )
+            ret = self.store.register_buffer(buffer.data_ptr(), buffer.nbytes)
+            if ret != 0:
+                raise EmbeddingStoreOperationError(
+                    f"Failed to register embedding read buffer: {ret}"
+                )
+            self._read_buffer = buffer
+        buffer = self._read_buffer
+        base = buffer.data_ptr()
+        sizes = [
+            _MOONCAKE_TENSOR_HEADER.size + spec.nbytes for spec in expected.values()
+        ]
+        offsets: list[int] = []
+        end = 0
+        for size in sizes:
+            offsets.append(end)
+            end += size
+        try:
+            results = self.store.batch_get_into(
+                list(expected), [base + offset for offset in offsets], sizes
+            )
+            if not _batch_get_completed_safely(results, sizes):
+                raise EmbeddingStoreError("Unsafe batch GET results")
+        except BaseException as error:
+            self._poison([buffer])
+            raise EmbeddingStoreError(
+                "Mooncake I/O completion is unconfirmed; retaining read buffer"
+            ) from error
+
+        valid: list[tuple[str, TensorSpec, int]] = []
+        for (key, spec), offset, size, result in zip(
+            expected.items(), offsets, sizes, results, strict=True
+        ):
+            if result != size:
+                logger.warning(
+                    "Store GET skipped for %s: result=%d, expected bytes=%d",
+                    key,
+                    result,
+                    size,
+                )
+                continue
+            try:
+                _validate_mooncake_tensor_metadata(
+                    key,
+                    ctypes.string_at(base + offset, _MOONCAKE_TENSOR_HEADER.size),
+                    spec,
+                )
+            except EmbeddingStoreOperationError as error:
+                logger.warning("Store GET skipped for %s: %s", key, error)
+                continue
+            valid.append((key, spec, offset + _MOONCAKE_TENSOR_HEADER.size))
+        if not valid:
+            return {}
+        # Allocate every destination before submitting any asynchronous copy.
+        targets = {
+            key: torch.empty(
+                spec.shape,
+                dtype=getattr(torch, spec.dtype.removeprefix("torch.")),
+                device=device,
+            )
+            for key, spec, _ in valid
+        }
+        ready = torch.Event() if device.type == "cuda" else None
+        try:
+            for key, spec, offset in valid:
+                targets[key].view(torch.uint8).view(-1).copy_(
+                    buffer[offset : offset + spec.nbytes],
+                    non_blocking=ready is not None,
+                )
+            if ready is not None:
+                ready.record(torch.accelerator.current_stream(device))
+                ready.synchronize()
+        except BaseException as error:
+            if ready is not None:
+                self._poison([buffer, *targets.values(), ready])
+                raise EmbeddingStoreError(
+                    "Embedding copy completion is unconfirmed; retaining buffers"
+                ) from error
+            raise
+        return targets
+
+    def put_tensor(self, pool_key: str, tensor: torch.Tensor) -> None:
+        self._check_healthy()
+        if not tensor.is_contiguous():
+            raise EmbeddingStoreOperationError("embedding tensor must be contiguous")
+        metadata = _encode_mooncake_tensor_metadata(tensor)
+        # The single publisher reuses this header only after PUT completion.
+        if self._put_header is None:
+            header = ctypes.create_string_buffer(len(metadata))
+            ret = self.store.register_buffer(ctypes.addressof(header), len(metadata))
+            if ret != 0:
+                raise EmbeddingStoreOperationError(
+                    f"Failed to register embedding header: {ret}"
+                )
+            self._put_header = header
+        header_ptr = ctypes.addressof(self._put_header)
+        ctypes.memmove(header_ptr, metadata, len(metadata))
+        data_size = tensor.numel() * tensor.element_size()
+        buffers = [
+            (tensor, tensor.data_ptr(), data_size),
+        ]
+        with self._registered_io(buffers) as confirm:
+            [result] = self.store.batch_put_from_multi_buffers(
+                [pool_key],
+                [[header_ptr, tensor.data_ptr()]],
+                [[len(metadata), data_size]],
+                self.replicate_config,
+            )
+            confirm(result)
+            if result < 0:
+                raise EmbeddingStoreOperationError(
+                    f"failed to put embedding tensor for {pool_key}"
+                )
+
+    def _unregister_buffer(self, addr: int) -> None:
+        try:
+            ret = self.store.unregister_buffer(addr)
+        except Exception as error:
+            raise EmbeddingStoreError(
+                f"could not confirm buffer unregistration addr={addr:#x}"
+            ) from error
+        if ret != 0:
+            raise EmbeddingStoreError(
+                f"failed to unregister embedding buffer addr={addr:#x}: {ret}"
+            )
+
+
+def _validate_mooncake_tensor_metadata(
+    pool_key: str,
+    metadata: bytes,
+    expected: TensorSpec,
+) -> None:
+    magic, dtype, num_tokens, hidden_dim = _MOONCAKE_TENSOR_HEADER.unpack(metadata)
+    if magic != _MOONCAKE_TENSOR_OBJECT_MAGIC:
+        raise EmbeddingStoreOperationError(
+            f"invalid Mooncake tensor metadata header for {pool_key}"
+        )
+    if _MOONCAKE_DTYPE_TO_TORCH_DTYPE.get(dtype) != expected.dtype:
+        raise EmbeddingStoreOperationError(
+            f"unexpected Mooncake tensor dtype for {pool_key}: {dtype}"
+        )
+    if (num_tokens, hidden_dim) != expected.shape:
+        raise EmbeddingStoreOperationError(
+            f"embedding tensor shape mismatch for {pool_key}: expected={expected.shape}"
+        )
+
+
+def _encode_mooncake_tensor_metadata(tensor: torch.Tensor) -> bytes:
+    dtype = str(tensor.dtype)
+    if dtype not in _TORCH_DTYPE_TO_MOONCAKE_DTYPE:
+        raise EmbeddingStoreOperationError(
+            f"unsupported embedding tensor dtype: {dtype}"
+        )
+    if tensor.ndim != 2:
+        raise EmbeddingStoreOperationError(
+            f"embedding tensor must be 2D: shape={tuple(tensor.shape)}"
+        )
+    return _MOONCAKE_TENSOR_HEADER.pack(
+        _MOONCAKE_TENSOR_OBJECT_MAGIC,
+        _TORCH_DTYPE_TO_MOONCAKE_DTYPE[dtype],
+        *tensor.shape,
+    )
diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py
index 46963c2b4113..081d05068e22 100644
--- a/vllm/distributed/ec_transfer/ec_connector/utils.py
+++ b/vllm/distributed/ec_transfer/ec_connector/utils.py
@@ -60,28 +60,33 @@ def collect_ec_item_metadata(
 
     Keyed by mm_hash, each entry carries a `metadata` dict with whatever
     placeholder fields `resolver` says this model needs published for its
-    modality, so a consumer can skip the image transform. `data` is None for
-    items served from the processor cache, in which case the metadata is
-    unavailable here and the consumer has to fall back to processing the
-    media itself. A connector that also has transfer coordinates to report
+    modality, so a consumer can skip media preprocessing. `item_indices`
+    identifies every occurrence in `mm_features`, including repeated hashes.
+    Audio token counts come from the placeholder; other metadata requires
+    `data` and falls back to raw media when it is unavailable on cache hits.
+    A connector that also has transfer coordinates to report
     (e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside
     `metadata`, not into it.
     """
     items: dict[str, dict[str, Any]] = {}
-    for feature in mm_features:
+    for index, feature in enumerate(mm_features):
         metadata: dict[str, Any] = {}
+        wanted = resolver.fields_for(feature.modality)
+        if "audio_num_tokens" in wanted:
+            metadata["audio_num_tokens"] = [feature.mm_position.get_num_embeds()]
         if feature.data is not None:
-            wanted = resolver.fields_for(feature.modality)
             for key, value in feature.data.get_data().items():
                 if key not in wanted:
                     continue
                 if isinstance(value, torch.Tensor):
-                    metadata[key] = value.tolist()
+                    metadata[key] = torch.atleast_1d(value).tolist()
                 elif is_list_of(value, (int, float)):
                     # Some metadata (e.g. Qwen3-VL video timestamps) is
                     # produced as a plain list rather than a tensor.
                     metadata[key] = value
-        items[feature.identifier] = {"metadata": metadata}
+        entry = items.setdefault(feature.identifier, {"item_indices": []})
+        entry["metadata"] = metadata
+        entry["item_indices"].append(index)
     return items
 
 
@@ -103,6 +108,7 @@ def aggregate(
         finished_sending = set[str]()
         finished_recving = set[str]()
         worker_meta = None
+        connector_stats = None
         for model_runner_output in outputs:
             assert model_runner_output is not None
             ec_output = model_runner_output.ec_connector_output
@@ -117,9 +123,17 @@ def aggregate(
                     meta if worker_meta is None else worker_meta.aggregate(meta)
                 )
 
+            if stats := ec_output.ec_connector_stats:
+                connector_stats = (
+                    stats
+                    if connector_stats is None
+                    else connector_stats.aggregate(stats)
+                )
+
         aggregated = ECConnectorOutput(
             finished_sending=finished_sending or None,
             finished_recving=finished_recving or None,
+            ec_connector_stats=connector_stats,
             ec_connector_worker_meta=worker_meta,
         )
         if aggregated.is_empty():
diff --git a/vllm/distributed/ec_transfer/ec_transfer_state.py b/vllm/distributed/ec_transfer/ec_transfer_state.py
index e58415ef6f4b..798fc42d842d 100644
--- a/vllm/distributed/ec_transfer/ec_transfer_state.py
+++ b/vllm/distributed/ec_transfer/ec_transfer_state.py
@@ -24,10 +24,7 @@ def has_ec_transfer() -> bool:
 
 
 def ensure_ec_transfer_initialized(vllm_config: "VllmConfig") -> None:
-    """
-    Initialize EC cache connector.
-    """
-
+    """Initialize EC cache connector."""
     global _EC_CONNECTOR_AGENT
 
     if vllm_config.ec_transfer_config is None:
diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py
index 6124d22b44d8..19fc15b27584 100644
--- a/vllm/distributed/elastic_ep/elastic_execute.py
+++ b/vllm/distributed/elastic_ep/elastic_execute.py
@@ -2,34 +2,32 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import gc
 import weakref
-from collections.abc import Iterable, Sequence
+from collections.abc import Iterable, Iterator, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager, nullcontext
 from dataclasses import replace
 from typing import TYPE_CHECKING
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from torch.distributed import P2POp
 
-from vllm.compilation.counter import compilation_counter
 from vllm.compilation.cuda_graph import CUDAGraphWrapper
 from vllm.compilation.wrapper import reset_compile_wrapper
-from vllm.config import (
-    CompilationMode,
-    set_current_vllm_config,
-)
+from vllm.config import ParallelConfig, set_current_vllm_config
 from vllm.distributed import (
     get_dp_group,
     get_ep_group,
     get_pcp_group,
     get_tp_group,
+    get_world_group,
 )
 from vllm.distributed.elastic_ep.standby_state import (
     create_standby_groups,
     get_standby_dp_group,
     get_standby_ep_group,
     get_standby_eplb_group,
+    get_standby_world_group,
     pop_standby_groups,
 )
 from vllm.distributed.eplb.eplb_communicator import (
@@ -41,13 +39,21 @@
 )
 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
+    get_ep_all2all_manager,
+)
 from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
 from vllm.model_executor.layers.fused_moe.eep_reconfigure import (
     make_eep_staged_quant_method,
 )
+from vllm.model_executor.warmup.flashinfer_autotune_cache import (
+    sync_flashinfer_autotune_cache,
+)
 from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
+from vllm.platforms import current_platform
 from vllm.utils import is_moe_layer
 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
+from vllm.v1.worker.dp_utils import skip_dp_coordination
 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
 from vllm.v1.worker.workspace import lock_workspace, unlock_workspace
 
@@ -59,6 +65,15 @@
     )
 
 
+def can_reuse_fused_moe_kernel(parallel_config: ParallelConfig) -> bool:
+    """Whether fused MoE kernels can be reused across EP rank changes.
+
+    The EP backend must preserve its CUDA graph entries when ranks are connected,
+    disconnected, masked, or unmasked.
+    """
+    return parallel_config.all2all_backend == "nixl_ep"
+
+
 def batch_transfer_weights(
     model: nn.Module,
     is_sender: bool,
@@ -105,31 +120,20 @@ def batch_transfer_weights(
 
 def broadcast_expert_mapping(
     physical_to_logical: torch.Tensor | None,
-    num_local_physical_experts: int | None,
-    num_logical_experts: int | None,
     dp_group: StatelessGroupCoordinator,
     device: torch.device,
     src_rank: int = 0,
-) -> tuple[torch.Tensor, int, int]:
+) -> torch.Tensor:
     if dp_group.rank_in_group == src_rank:
         assert physical_to_logical is not None
-        assert num_local_physical_experts is not None
-        assert num_logical_experts is not None
         assert physical_to_logical.dtype == torch.int64
         shape_tensor = torch.tensor(
             list(physical_to_logical.shape), dtype=torch.int64, device="cpu"
         )
-        metadata_tensor = torch.tensor(
-            [num_local_physical_experts, num_logical_experts],
-            dtype=torch.int64,
-            device="cpu",
-        )
     else:
         shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu")
-        metadata_tensor = torch.empty(2, dtype=torch.int64, device="cpu")
 
     shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank)
-    metadata_tensor = dp_group.tcp_store_group.broadcast(metadata_tensor, src_rank)
 
     if dp_group.rank_in_group != src_rank:
         assert device is not None
@@ -140,11 +144,7 @@ def broadcast_expert_mapping(
         )
 
     assert physical_to_logical is not None
-    physical_to_logical = dp_group.broadcast(physical_to_logical, src_rank)
-    num_local_physical_experts = int(metadata_tensor[0].item())
-    num_logical_experts = int(metadata_tensor[1].item())
-
-    return physical_to_logical, num_local_physical_experts, num_logical_experts
+    return dp_group.broadcast(physical_to_logical, src_rank)
 
 
 class ElasticEPScalingExecutor:
@@ -166,6 +166,21 @@ def worker(self):
             raise RuntimeError("Worker has been garbage collected")
         return worker
 
+    def _can_reuse_fused_moe_kernel(self) -> bool:
+        return can_reuse_fused_moe_kernel(
+            self.worker.vllm_config.parallel_config,
+        )
+
+    @contextmanager
+    def _disable_flashinfer_autotune(self) -> Iterator[None]:
+        kernel_config = self.worker.vllm_config.kernel_config
+        enabled = kernel_config.enable_flashinfer_autotune
+        kernel_config.enable_flashinfer_autotune = False
+        try:
+            yield
+        finally:
+            kernel_config.enable_flashinfer_autotune = enabled
+
     def execute(self, execute_method: str, *args, **kwargs):
         method = getattr(self, execute_method, None)
         if method is None:
@@ -260,11 +275,22 @@ def prepare_reconfiguration(
             use_all2all=use_all2all,
             enable_eplb=parallel_config.enable_eplb,
         )
-        self.stage_standby_moe_quant_methods()
+        standby_ep_group = get_standby_ep_group()
+        assert standby_ep_group is not None
+        all2all_manager = get_ep_all2all_manager(standby_ep_group)
+        all2all_manager.stage_ep_size()
+        if not self._can_reuse_fused_moe_kernel():
+            self.stage_standby_moe_quant_methods(all2all_manager)
         self._prepare_eplb_communicator(get_standby_eplb_group())
         if new_dp_size > old_dp_size:
             self.transfer_weights(old_dp_size, new_dp_size)
-        self._warm_target_groups(get_standby_dp_group(), get_standby_ep_group())
+        # On ROCm this deadlocks while serving; warm_and_capture warms these at commit.
+        if not current_platform.is_rocm():
+            self._warm_target_groups(get_standby_dp_group(), standby_ep_group)
+        if new_dp_size > old_dp_size and self._can_reuse_fused_moe_kernel():
+            target_world_group = get_standby_world_group()
+            assert target_world_group is not None
+            sync_flashinfer_autotune_cache(self.worker.model_runner, target_world_group)
 
     def _prepare_eplb_communicator(self, eplb_group) -> None:
         assert eplb_group is not None
@@ -338,14 +364,9 @@ def broadcast_expert_mapping(self) -> None:
         assert eplb_state is not None
         eplb_state.drain_async()
         eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
-        physical_to_logical = eplb_model_state.physical_to_logical_map
-        num_physical_experts = physical_to_logical.shape[1]
-        num_local_physical_experts = num_physical_experts // get_ep_group().world_size
-        num_logical_experts = eplb_model_state.logical_replica_count.shape[1]
+        physical_to_logical = eplb_model_state.physical_to_logical_map_buffer
         broadcast_expert_mapping(
             physical_to_logical=physical_to_logical,
-            num_local_physical_experts=num_local_physical_experts,
-            num_logical_experts=num_logical_experts,
             dp_group=standby_dp_group,
             src_rank=0,
             device=self.worker.device,
@@ -368,7 +389,7 @@ def _make_eep_moe_config(self, module, dp_group, ep_group):
             moe_parallel_config=moe_parallel_config,
         )
 
-    def stage_standby_moe_quant_methods(self) -> None:
+    def stage_standby_moe_quant_methods(self, all2all_manager) -> None:
         standby_dp_group = get_standby_dp_group()
         standby_ep_group = get_standby_ep_group()
         model = self.worker.model_runner.get_model()
@@ -383,6 +404,7 @@ def stage_standby_moe_quant_methods(self) -> None:
                         standby_dp_group,
                         standby_ep_group,
                     ),
+                    all2all_manager=all2all_manager,
                 )
                 if staged_quant_method is not None:
                     self._staged_moe_quant_methods[module] = staged_quant_method
@@ -396,7 +418,6 @@ def _commit_staged_moe_quant_methods(self) -> None:
                 continue
             assert staged_quant_method.moe_kernel is not None
             module._replace_quant_method(staged_quant_method)
-            staged_quant_method.moe_kernel.prepare_finalize.on_commit()
         self._staged_moe_quant_methods.clear()
 
     def _release_cuda_graphs(self) -> None:
@@ -405,7 +426,7 @@ def _release_cuda_graphs(self) -> None:
             wrapper.concrete_cudagraph_entries = {}
 
         elif isinstance(self.worker.model_runner.model, UBatchWrapper):
-            raise RuntimeError("DBO is not yet supported in elastic EP")
+            self.worker.model_runner.model.clear_graphs()
 
         torch.compiler.reset()
         with set_current_vllm_config(self.worker.vllm_config):
@@ -427,10 +448,8 @@ def switch_and_remove(self) -> None:
         self._wait_for_group_cleanup()
 
     def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]:
-        old_dp_size = get_dp_group().world_size
-        old_ep_size = get_ep_group().world_size
-
-        self._release_cuda_graphs()
+        if not self._can_reuse_fused_moe_kernel():
+            self._release_cuda_graphs()
         retired_groups = _replace_active_groups(**pop_standby_groups())
 
         parallel_config = self.worker.vllm_config.parallel_config
@@ -487,47 +506,13 @@ def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]:
         parallel_config.eplb_config.num_redundant_experts = (
             num_physical_experts - num_logical_experts
         )
-        old_physical_to_logical = eplb_model_state.physical_to_logical_map
-        num_moe_layers = old_physical_to_logical.shape[0]
-        num_local_experts = eplb_model_state.expert_load_pass.shape[1] // old_ep_size
-        if new_dp_size > old_dp_size:
-            expanded_physical_to_logical = torch.full(
-                (num_moe_layers, num_local_experts * new_ep_size),
-                -1,
-                dtype=old_physical_to_logical.dtype,
-                device=old_physical_to_logical.device,
-            )
-            expanded_physical_to_logical[:, : num_local_experts * old_ep_size] = (
-                old_physical_to_logical
-            )
-            eplb_model_state.physical_to_logical_map = expanded_physical_to_logical
-
-        old_num_physical_experts = eplb_model_state.expert_load_pass.shape[1]
-        pad_size = num_physical_experts - old_num_physical_experts
-        if new_dp_size > old_dp_size:
-            assert pad_size > 0
-            expanded_expert_load_pass = F.pad(
-                eplb_model_state.expert_load_pass, (0, pad_size), value=0
-            )
-            expanded_expert_load_window = F.pad(
-                eplb_model_state.expert_load_window, (0, pad_size), value=0
-            )
-            eplb_model_state.expert_load_pass = expanded_expert_load_pass
-            eplb_model_state.expert_load_window = expanded_expert_load_window
-        else:
-            assert pad_size < 0
-            eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[
-                :, :num_physical_experts
-            ]
-            eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[
-                :, :, :num_physical_experts
-            ]
+        eplb_state.reconfigure_physical_expert_slots(model_config, num_physical_experts)
 
         model = self.worker.model_runner.get_model()
         model.expert_weights = []
         with set_current_vllm_config(self.worker.vllm_config):
             model.set_eplb_state(
-                eplb_model_state.expert_load_pass,
+                eplb_model_state.expert_load_pass_buffer,
                 eplb_model_state.logical_to_physical_map,
                 eplb_model_state.logical_replica_count,
             )
@@ -539,6 +524,7 @@ def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]:
                 num_local_physical_experts=num_local_experts,
             )
             self._commit_staged_moe_quant_methods()
+            get_ep_all2all_manager().commit_ep_size()
             # Legacy modular methods need to be recreated for the new EP size.
             for module in moe_modules:
                 if getattr(module._quant_method, "wraps_legacy_quant_method", False):
@@ -548,21 +534,6 @@ def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]:
         eplb_state.update_communicator(model_config, self._prepared_eplb_communicator)
         self._prepared_eplb_communicator = None
 
-        if (
-            self.worker.vllm_config.compilation_config.mode
-            == CompilationMode.STOCK_TORCH_COMPILE
-        ):
-            # NOTE(yongji): when using stock torch.compile,
-            # torch.compile is triggered during GPUModelRunner's load_model()
-            # TODO(yongji):check do we need to re-trigger torch.compile here?
-            # any changes to the tensor shapes in execution should already
-            # be handled internally by torch.compile.
-            backend = self.worker.vllm_config.compilation_config.init_backend(
-                self.worker.vllm_config
-            )
-            compilation_counter.stock_torch_compile_count += 1
-            self.worker.model_runner.model.compile(fullgraph=True, backend=backend)
-
         return retired_groups
 
     def _perform_eplb_reshuffle(
@@ -608,7 +579,12 @@ def commit_scale_up(self, is_existing_worker: bool) -> None:
         else:
             mapping = self.receive_expert_mapping()
             self.worker.model_runner.setup_eplb_from_mapping(mapping)
-        self.warm_and_capture()
+        if not self._can_reuse_fused_moe_kernel():
+            # warm_and_capture's dummy run warms the deferred dp/ep groups.
+            self.warm_and_capture()
+        elif current_platform.is_rocm():
+            # Reuse path skips warm_and_capture; warm the deferred groups here.
+            self._warm_target_groups(get_dp_group(), get_ep_group())
         self._perform_eplb_reshuffle(async_op=True)
         if is_existing_worker:
             self._start_group_cleanup(retired_groups)
@@ -619,7 +595,11 @@ def commit_scale_down(self, new_dp_size: int, removing: bool) -> None:
             self.switch_and_remove()
         else:
             retired_groups = self.switch_and_prepare()
-            self.warm_and_capture()
+            if not self._can_reuse_fused_moe_kernel():
+                self.warm_and_capture()
+            elif current_platform.is_rocm():
+                # Reuse path skips warm_and_capture; warm the deferred groups here.
+                self._warm_target_groups(get_dp_group(), get_ep_group())
             self._start_group_cleanup(retired_groups)
 
     def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:
@@ -676,43 +656,28 @@ def prepare_new_worker(self) -> None:
             expert_weights=expert_weights,
         )
         torch.accelerator.synchronize()
-        self._warm_target_groups(get_dp_group(), get_ep_group())
+        if not current_platform.is_rocm():
+            self._warm_target_groups(get_dp_group(), get_ep_group())
+        if self._can_reuse_fused_moe_kernel():
+            sync_flashinfer_autotune_cache(self.worker.model_runner, get_world_group())
 
     def receive_expert_mapping(self) -> torch.Tensor:
         dp_group = get_dp_group()
         assert isinstance(dp_group, StatelessGroupCoordinator)
-        physical_to_logical, num_local_physical_experts, _ = broadcast_expert_mapping(
+        return broadcast_expert_mapping(
             physical_to_logical=None,
-            num_local_physical_experts=None,
-            num_logical_experts=None,
             dp_group=dp_group,
             src_rank=0,
             device=self.worker.device,
         )
-        num_moe_layers = physical_to_logical.shape[0]
-        new_dp_size = get_dp_group().world_size
-        tp_size = self.worker.vllm_config.parallel_config.tensor_parallel_size
-        new_ep_size = new_dp_size * tp_size
-        expanded_physical_to_logical = torch.full(
-            (num_moe_layers, num_local_physical_experts * new_ep_size),
-            -1,
-            dtype=physical_to_logical.dtype,
-            device=physical_to_logical.device,
-        )
-        expanded_physical_to_logical[:, : physical_to_logical.shape[1]] = (
-            physical_to_logical
-        )
-        return expanded_physical_to_logical
 
-    def warmup_local_kernels(self) -> None:
+    def warmup_new_worker(self) -> None:
         with set_current_vllm_config(self.worker.vllm_config):
             kernel_warmup(self.worker, process_local_only=True)
+        if self._can_reuse_fused_moe_kernel():
+            self.warm_and_capture()
 
     def warm_and_capture(self) -> None:
-        # Must run on every DP sibling in lockstep: _dummy_run calls
-        # coordinate_batch_across_dp whenever data_parallel_size > 1
-        # (gpu_model_runner.py:3663), which deadlocks if any rank skips it.
-
         # Save and clear block tables so the dummy MoE forward doesn't
         # write dummy slot mappings into real KV-cache blocks.
         multi_block_table = self.worker.model_runner.input_batch.block_table
@@ -735,8 +700,15 @@ def warm_and_capture(self) -> None:
         # directly with skip_eplb=True so dummy routing doesn't pollute the
         # just-rebalanced EPLB stats.
         runner = self.worker.model_runner
-        runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)
-        self.worker.compile_or_warm_up_model()
+        all2all_manager = get_ep_all2all_manager()
+        reuse_kernel = self._can_reuse_fused_moe_kernel()
+        with (
+            skip_dp_coordination() if reuse_kernel else nullcontext(),
+            all2all_manager.mask_remote_ranks() if reuse_kernel else nullcontext(),
+            self._disable_flashinfer_autotune() if reuse_kernel else nullcontext(),
+        ):
+            runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)
+            self.worker.compile_or_warm_up_model()
 
         lock_workspace()
 
diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py
index 6718cd8ae873..5ba1891b7aca 100644
--- a/vllm/distributed/elastic_ep/elastic_state.py
+++ b/vllm/distributed/elastic_ep/elastic_state.py
@@ -204,7 +204,7 @@ def _progress_new_engine(self) -> bool:
             return True
 
         elif state == ScaleUpNewEngineState.PREPARE:
-            self._collective_rpc("elastic_ep_execute", args=("warmup_local_kernels",))
+            self._collective_rpc("elastic_ep_execute", args=("warmup_new_worker",))
             self._mark_ready_for_switch()
             tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu")
             torch.distributed.all_reduce(
diff --git a/vllm/distributed/elastic_ep/standby_state.py b/vllm/distributed/elastic_ep/standby_state.py
index 1892f3e79426..084ed51f8c52 100644
--- a/vllm/distributed/elastic_ep/standby_state.py
+++ b/vllm/distributed/elastic_ep/standby_state.py
@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import torch
 
+from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm
 from vllm.distributed.parallel_state import (
     _init_stateless_group,
     _node_count,
+    get_pcp_group,
     get_pp_group,
     get_tp_group,
     get_world_group,
@@ -71,34 +73,45 @@ def create_standby_groups(
     _STANDBY_WORLD_NODE_COUNT = _node_count(_STANDBY_WORLD.tcp_store_group)
 
     tp_size = get_tp_group().world_size
+    pcp_size = get_pcp_group().world_size
     pp_size = get_pp_group().world_size
 
     all_ranks = torch.arange(new_world_size_across_dp).reshape(
-        -1, new_dp_size, pp_size, tp_size
+        -1, new_dp_size, pp_size, pcp_size, tp_size
     )
-    standby_dp_ranks = all_ranks.transpose(1, 3).reshape(-1, new_dp_size).unbind(0)
+    standby_dp_ranks = all_ranks.transpose(1, 4).reshape(-1, new_dp_size).unbind(0)
     standby_dp_ranks = [x.tolist() for x in standby_dp_ranks]
-    _STANDBY_DP = _init_stateless_group(
-        standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store
-    )
 
-    standby_ep_ranks = (
-        all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0)
-    )
-    standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
-    _STANDBY_EP = _init_stateless_group(
-        standby_ep_ranks, "ep", master_ip, backend, coord_store, use_all2all=use_all2all
-    )
+    # Deferred to commit so the warm-up runs while the engine is paused.
+    with defer_comm_warmup_on_rocm():
+        _STANDBY_DP = _init_stateless_group(
+            standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store
+        )
 
-    if enable_eplb:
-        _STANDBY_EPLB = _init_stateless_group(
+        standby_ep_ranks = (
+            all_ranks.transpose(1, 2)
+            .reshape(-1, new_dp_size * pcp_size * tp_size)
+            .unbind(0)
+        )
+        standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
+        _STANDBY_EP = _init_stateless_group(
             standby_ep_ranks,
-            "eplb",
+            "ep",
             master_ip,
             backend,
-            coord_store=coord_store,
+            coord_store,
+            use_all2all=use_all2all,
         )
 
+        if enable_eplb:
+            _STANDBY_EPLB = _init_stateless_group(
+                standby_ep_ranks,
+                "eplb",
+                master_ip,
+                backend,
+                coord_store=coord_store,
+            )
+
 
 def pop_standby_groups() -> dict:
     """Return all standby groups and clear the standby state."""
diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py
index ff872c8c2220..3b7bc908eca7 100644
--- a/vllm/distributed/eplb/async_worker.py
+++ b/vllm/distributed/eplb/async_worker.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-The async worker that transfers experts in the background.
-"""
+"""The async worker that transfers experts in the background."""
 
 import threading
 from typing import TYPE_CHECKING
@@ -12,7 +10,7 @@
 from vllm.distributed.parallel_state import get_eplb_group
 from vllm.logger import init_logger
 
-from .eplb_utils import CpuGpuEvent
+from .eplb_utils import CpuGpuEvent, device_stream
 from .rebalance_execute import AsyncEplbLayerResult, transfer_layer
 
 if TYPE_CHECKING:
@@ -26,17 +24,17 @@ def start_async_worker(
     is_profile: bool = False,
 ) -> threading.Thread:
     rank = get_eplb_group().device_group.rank()
-    device_index = state.cuda_device_index
+    device_index = state.device_index
     assert state.is_async
 
     def thread_target() -> None:
         assert device_index is not None
         torch.accelerator.set_device_index(device_index)
-        cuda_stream = torch.cuda.Stream(device=device_index)
+        stream = torch.Stream(device=torch.accelerator.current_accelerator())
         try:
             transfer_run_periodically(
                 state=state,
-                cuda_stream=cuda_stream,
+                stream=stream,
                 is_profile=is_profile,
             )
         except Exception as exc:  # pragma: no cover - diagnostic path
@@ -51,13 +49,13 @@ def run_rebalance_experts(
     model_state: "EplbModelState",
     eplb_state: "EplbState",
     physical_to_logical_map_cpu: torch.Tensor,
-    cuda_stream: torch.cuda.Stream,
+    stream: torch.Stream,
 ) -> torch.Tensor:
     assert model_state.eplb_stats is not None
     eplb_stats = model_state.eplb_stats
 
     # Move the global expert load window to CPU for computation.
-    with torch.cuda.stream(cuda_stream):
+    with device_stream(stream):
         global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
     # Compute new expert mappings for the model
     new_physical_to_logical_map = eplb_state.policy.rebalance_experts(
@@ -75,11 +73,11 @@ def run_rebalance_experts(
 
 def transfer_run_periodically(
     state: "EplbState",
-    cuda_stream: torch.cuda.Stream,
+    stream: torch.Stream,
     is_profile: bool = False,
 ) -> None:
     while True:
-        state.rearrange_event.wait(stream=cuda_stream)
+        state.rearrange_event.wait(stream=stream)
 
         eplb_group = get_eplb_group().device_group
         eplb_cpu_group = get_eplb_group().cpu_group
@@ -89,16 +87,16 @@ def transfer_run_periodically(
         for model_state in state.model_states.values():
             layer_idx = 0
             # Set the async worker's CUDA stream on the communicator
-            model_state.communicator.set_stream(cuda_stream)
+            model_state.communicator.set_stream(stream)
             num_layers = model_state.model.num_moe_layers
 
             # Snapshot the physical_to_logical_map (synchronized with
             # rearrange_event) and copy it to CPU
-            with torch.cuda.stream(cuda_stream):
+            with device_stream(stream):
                 physical_to_logical_map_cpu = model_state.physical_to_logical_map.cpu()
 
             new_physical_to_logical_map = run_rebalance_experts(
-                model_state, state, physical_to_logical_map_cpu, cuda_stream
+                model_state, state, physical_to_logical_map_cpu, stream
             )
 
             # Execute one EPLB layer transfer per model forward pass. Each iteration
@@ -134,13 +132,13 @@ def transfer_run_periodically(
                     communicator=model_state.communicator,
                     ep_group=eplb_group,
                     is_profile=is_profile,
-                    cuda_stream=cuda_stream,
+                    stream=stream,
                     layer_idx=layer_idx,
                 )
 
                 # Wait until all writes to expert_buffer have finished before making the
                 # AsyncEplbLayerResult visible to the main thread.
-                cuda_stream.synchronize()
+                stream.synchronize()
 
                 # This event guarantees that expert_buffer will not be overwritten by
                 # subsequent iterations of this loop until the main thread has consumed
@@ -157,6 +155,6 @@ def transfer_run_periodically(
                 # Block this thread until the main thread and main stream
                 # finish copying model_state.expert_buffer into
                 # model_state.model.expert_weights[layer_idx]
-                consumed_event.wait(stream=cuda_stream)
+                consumed_event.wait(stream=stream)
                 assert model_state.pending_result is None
                 layer_idx += 1
diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py
index 77362e593c2a..730c01a3882e 100644
--- a/vllm/distributed/eplb/eplb_communicator.py
+++ b/vllm/distributed/eplb/eplb_communicator.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-EPLB communicator implementations and factory.
-"""
+"""EPLB communicator implementations and factory."""
 
 import contextlib
 import time
@@ -24,6 +22,7 @@
 from vllm.distributed.device_communicators.pynccl_wrapper import (
     ncclDataTypeEnum,
 )
+from vllm.distributed.eplb.eplb_utils import device_stream
 from vllm.distributed.parallel_state import (
     GroupCoordinator,
     get_pp_group,
@@ -34,6 +33,7 @@
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.utils.torch_utils import PIN_MEMORY
 
 logger = init_logger(__name__)
 
@@ -88,8 +88,8 @@ def needs_profile_buffer_reservation(self) -> bool:
         communication buffers."""
         return True
 
-    def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:
-        self._cuda_stream = cuda_stream
+    def set_stream(self, stream: torch.Stream | None) -> None:
+        self._stream = stream
 
     def _log_initialized(self) -> None:
         if is_local_first_rank():
@@ -102,10 +102,10 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator):
     def __init__(
         self,
         ep_group: ProcessGroup,
-        cuda_stream: torch.cuda.Stream | None = None,
+        stream: torch.Stream | None = None,
     ) -> None:
         self._ep_group = ep_group
-        self._cuda_stream = cuda_stream
+        self._stream = stream
         self._p2p_ops: list[P2POp] = []
         self._log_initialized()
 
@@ -145,7 +145,7 @@ def execute(self) -> None:
         if not self._p2p_ops:
             return
         try:
-            with torch.cuda.stream(self._cuda_stream):
+            with device_stream(self._stream):
                 reqs = batch_isend_irecv(self._p2p_ops)
                 for req in reqs:
                     req.wait()
@@ -159,10 +159,10 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):
     def __init__(
         self,
         cpu_group: ProcessGroup,
-        cuda_stream: torch.cuda.Stream | None = None,
+        stream: torch.Stream | None = None,
     ) -> None:
         self._cpu_group = cpu_group
-        self._cuda_stream = cuda_stream
+        self._stream = stream
         self._ops: list[tuple[str, torch.Tensor, int]] = []
         self._log_initialized()
 
@@ -204,7 +204,9 @@ def build_ops() -> None:
                         )
                     )
                     continue
-                cpu_tensor = torch.empty_like(tensor, device="cpu")
+                cpu_tensor = torch.empty_like(
+                    tensor, device="cpu", pin_memory=PIN_MEMORY
+                )
                 p2p_ops.append(
                     P2POp(
                         torch.distributed.irecv,
@@ -216,7 +218,7 @@ def build_ops() -> None:
                 recv_staging.append((tensor, cpu_tensor))
 
         try:
-            with torch.cuda.stream(self._cuda_stream):
+            with device_stream(self._stream):
                 build_ops()
         finally:
             self._ops.clear()
@@ -224,10 +226,10 @@ def build_ops() -> None:
         # Wait for all D2H copies to finish
         # before issuing gloo batch_isend_irecv operations.
         with gpu_sync_allowed():
-            if self._cuda_stream is not None:
-                self._cuda_stream.synchronize()
+            if self._stream is not None:
+                self._stream.synchronize()
             else:
-                torch.cuda.current_stream().synchronize()
+                torch.accelerator.current_stream().synchronize()
 
         reqs = batch_isend_irecv(p2p_ops)
         for req in reqs:
@@ -235,11 +237,69 @@ def build_ops() -> None:
 
         if not recv_staging:
             return
-        with torch.cuda.stream(self._cuda_stream):
+        with device_stream(self._stream):
             for dst_tensor, cpu_tensor in recv_staging:
                 dst_tensor.copy_(cpu_tensor, non_blocking=True)
 
 
+class TorchDistXCCLStagedEplbCommunicator(EplbCommunicator):
+    """EPLB communicator using XCCL device-to-device P2P on XPU."""
+
+    def __init__(
+        self,
+        ep_group: ProcessGroup,
+        stream: torch.Stream | None = None,
+    ) -> None:
+        self._ep_group = ep_group
+        self._stream = stream
+        self._ops: list[tuple[str, torch.Tensor, int]] = []
+        self._log_initialized()
+
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._ops.append(("send", tensor, dst_rank))
+
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._ops.append(("recv", tensor, src_rank))
+
+    def execute(self) -> None:
+        if not self._ops:
+            return
+
+        p2p_ops: list[P2POp] = []
+        try:
+            for op, tensor, peer_rank in self._ops:
+                send_or_recv = (
+                    torch.distributed.isend if op == "send" else torch.distributed.irecv
+                )
+                p2p_ops.append(
+                    P2POp(
+                        send_or_recv,
+                        tensor,
+                        peer_rank,
+                        self._ep_group,
+                    )
+                )
+        finally:
+            self._ops.clear()
+
+        with device_stream(self._stream):
+            reqs = batch_isend_irecv(p2p_ops)
+            for req in reqs:
+                req.wait()
+
+
 class NixlEplbCommunicator(EplbCommunicator):
     """EPLB communicator backed by NIXL READ transfers."""
 
@@ -255,6 +315,7 @@ def __init__(
             cpu_group: CPU process group for metadata exchange.
             all_expert_weights: Expert weight tensors for all MoE layers.
             expert_buffer: Pre-allocated receive buffer tensors.
+
         """
         assert all_expert_weights, (
             "NixlEplbCommunicator requires non-empty all_expert_weights."
@@ -341,7 +402,7 @@ def _make_agent_name(self) -> str:
         uid = uuid.uuid4().hex[:8]
         return f"eplb-{self._rank}{pp_suffix}-{uid}"
 
-    def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:
+    def set_stream(self, stream: torch.Stream | None) -> None:
         pass
 
     def add_send(
@@ -600,10 +661,10 @@ class PyNcclEplbCommunicator(EplbCommunicator):
     def __init__(
         self,
         pynccl_comm: PyNcclCommunicator,
-        cuda_stream: torch.cuda.Stream | None = None,
+        stream: torch.Stream | None = None,
     ) -> None:
         self._pynccl_comm = pynccl_comm
-        self._cuda_stream = cuda_stream
+        self._stream = stream
         self._group_started = False
         self._log_initialized()
 
@@ -620,7 +681,7 @@ def add_send(
     ) -> None:
         self._ensure_group_started()
         for tensor in tensors:
-            self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream)
+            self._pynccl_comm.send(tensor, dst_rank, stream=self._stream)
 
     def add_recv(
         self,
@@ -630,7 +691,7 @@ def add_recv(
     ) -> None:
         self._ensure_group_started()
         for tensor in tensors:
-            self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)
+            self._pynccl_comm.recv(tensor, src_rank, stream=self._stream)
 
     def execute(self) -> None:
         if self._group_started:
@@ -663,6 +724,7 @@ def create_eplb_communicator(
             zero-copy RDMA reads.
         expert_buffer: Pre-allocated receive buffer tensors (one per
             weight tensor in a single layer).
+
     """
     first_layer = expert_weights[0] if expert_weights else []
     tensor_device_type = first_layer[0].device.type if first_layer else "cpu"
@@ -747,6 +809,8 @@ def _create_pynccl() -> EplbCommunicator:
         return TorchDistGlooStagedEplbCommunicator(
             cpu_group=group_coordinator.cpu_group,
         )
+    elif backend == "torch_xccl":
+        return TorchDistXCCLStagedEplbCommunicator(ep_group=torch_group)
     elif backend == "torch_nccl":
         return TorchDistNcclEplbCommunicator(ep_group=torch_group)
     elif backend == "pynccl":
diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py
index 580477ba239e..e4d805b9fead 100644
--- a/vllm/distributed/eplb/eplb_state.py
+++ b/vllm/distributed/eplb/eplb_state.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Expert parallelism load balancer (EPLB) metrics and states.
+"""Expert parallelism load balancer (EPLB) metrics and states.
 
 # Glossary
 
@@ -74,9 +73,7 @@ def _compute_eplb_load_stats(
 
 @dataclass
 class EplbStats:
-    """
-    Model stats used in EPLB rebalancing algorithm.
-    """
+    """Model stats used in EPLB rebalancing algorithm."""
 
     global_expert_load_window: torch.Tensor
     """
@@ -121,6 +118,8 @@ class EplbModelState:
      [0, 2, 0, 1, 0, 3]]
     ```
     """
+    physical_to_logical_map_buffer: torch.Tensor
+    """Maximum-capacity buffer backing ``physical_to_logical_map``."""
     logical_to_physical_map: torch.Tensor
     """
     Mapping from logical experts to physical experts.
@@ -168,6 +167,8 @@ class EplbModelState:
 
     Shape: (num_moe_layers, num_physical_experts)
     """
+    expert_load_pass_buffer: torch.Tensor
+    """Maximum-capacity buffer backing ``expert_load_pass``."""
     expert_load_window: torch.Tensor
     """
     A sliding window of expert load.
@@ -202,9 +203,9 @@ class EplbModelState:
     """
     EPLB stats for the model.
     """
-    cuda_device_index: int | None
+    device_index: int | None
     """
-    CUDA device index for the async EPLB worker thread.
+    Device index for the async EPLB worker thread.
     """
     communicator: EplbCommunicator
     """
@@ -229,9 +230,7 @@ class EplbModelState:
 
 
 class EplbState:
-    """
-    EplbState of each expert parallel model. Key is the model config hash.
-    """
+    """EplbState of each expert parallel model. Key is the model config hash."""
 
     def __init__(self, parallel_config: ParallelConfig, device: torch.device):
         self.parallel_config = parallel_config
@@ -287,28 +286,28 @@ def __init__(self, parallel_config: ParallelConfig, device: torch.device):
         """
         Background thread handling async transfers.
         """
-        self.cuda_device_index: int | None = None
+        self.device_index: int | None = None
         """
-        CUDA device index for the async EPLB worker thread.
+        Device index for the async EPLB worker thread.
         """
-        if self.device.type == "cuda":
-            self.cuda_device_index = self.device.index
-            if self.cuda_device_index is None and torch.cuda.is_available():
-                self.cuda_device_index = torch.accelerator.current_device_index()
+        if self.device.type in ("cuda", "xpu"):
+            self.device_index = self.device.index
+            if self.device_index is None and torch.accelerator.is_available():
+                self.device_index = torch.accelerator.current_device_index()
 
     @staticmethod
     def build_initial_global_physical_to_logical_map(
         num_routed_experts: int,
         num_redundant_experts: int,
     ) -> Sequence[int]:
-        """
-        Build an initial expert arrangement using the following structure:
+        """Build an initial expert arrangement using the following structure:
         [original routed experts, redundant experts]
 
         Returns:
             physical_to_logical_map (Sequence[int]): A list of integers,
                 where each integer is the index of the logical expert
                 that the corresponding physical expert maps to.
+
         """
         global_physical_to_logical_map = list(range(num_routed_experts))
         global_physical_to_logical_map += [
@@ -317,8 +316,7 @@ def build_initial_global_physical_to_logical_map(
         return global_physical_to_logical_map
 
     def validate_ep_configuration(self, new_model: MixtureOfExperts):
-        """
-        Validate that the expert parallel configuration of
+        """Validate that the expert parallel configuration of
         the new model is the same as the existing models.
         """
         if len(self.model_states) > 0:
@@ -357,9 +355,7 @@ def add_model(
         model: MixtureOfExperts,
         model_config: ModelConfig,
     ):
-        """
-        Build the initial EPLB state.
-        """
+        """Build the initial EPLB state."""
         self.validate_ep_configuration(model)
         self.is_async = self.parallel_config.eplb_config.use_async
 
@@ -373,6 +369,19 @@ def add_model(
             physical_to_logical_map_list,
             device=self.device,
         )
+        physical_expert_capacity = model.num_physical_experts
+        if self.parallel_config.enable_elastic_ep:
+            physical_expert_capacity = (
+                model.num_local_physical_experts
+                * self.parallel_config.elastic_ep_max_dp_size
+                * get_ep_group().world_size
+                // self.parallel_config.data_parallel_size
+            )
+            physical_to_logical_map = torch.nn.functional.pad(
+                physical_to_logical_map,
+                (0, physical_expert_capacity - model.num_physical_experts),
+                value=-1,
+            )
         # Assuming 8 GPUs per node, this supports up to
         # (1023 + 1) / 8 = 128 nodes for now.
         # TODO(rui): make this configurable
@@ -407,6 +416,10 @@ def add_model(
             )
             .contiguous()
         )
+        physical_to_logical_map_buffer = physical_to_logical_map
+        physical_to_logical_map = physical_to_logical_map_buffer[
+            :, : model.num_physical_experts
+        ]
         logical_to_physical_map = (
             logical_to_physical_map.unsqueeze(0)
             .expand(
@@ -425,11 +438,12 @@ def add_model(
             .contiguous()
         )
 
-        expert_load_pass = torch.zeros(
-            (model.num_moe_layers, model.num_physical_experts),
+        expert_load_pass_buffer = torch.zeros(
+            (model.num_moe_layers, physical_expert_capacity),
             dtype=torch.int32,
             device=self.device,
         )
+        expert_load_pass = expert_load_pass_buffer[:, : model.num_physical_experts]
         self.expert_load_window_size = self.parallel_config.eplb_config.window_size
         expert_load_window = torch.zeros(
             (
@@ -460,7 +474,7 @@ def add_model(
         ]
 
         model.set_eplb_state(
-            expert_load_pass,
+            expert_load_pass_buffer,
             logical_to_physical_map,
             logical_replica_count,
         )
@@ -479,16 +493,18 @@ def add_model(
 
         model_state = EplbModelState(
             physical_to_logical_map=physical_to_logical_map,
+            physical_to_logical_map_buffer=physical_to_logical_map_buffer,
             logical_to_physical_map=logical_to_physical_map,
             logical_replica_count=logical_replica_count,
             expert_load_pass=expert_load_pass,
+            expert_load_pass_buffer=expert_load_pass_buffer,
             expert_load_window=expert_load_window,
             model_name=model_config.model,
             model=model,
             expert_buffer=expert_buffer,
             rebalanced=False,
             eplb_stats=None,
-            cuda_device_index=self.cuda_device_index,
+            device_index=self.device_index,
             communicator=communicator,
             num_unpadded_tokens_tensors=num_unpadded_tokens_tensors,
         )
@@ -510,6 +526,7 @@ def prepare_forward(
             ubatch_slices: When DBO is active, a list of
                 ``UBatchSlice`` objects describing each micro-batch's
                 token range.  When ``None``, only ``tensors[0]`` is filled.
+
         """
         model_state = self.model_states.get(compute_hash_cached(model_config))
         if model_state is None or model_state.num_unpadded_tokens_tensors is None:
@@ -532,8 +549,7 @@ def step(
         is_profile: bool = False,
         log_stats: bool = False,
     ) -> None:
-        """
-        Step the EPLB state.
+        """Step the EPLB state.
 
         Args:
             is_dummy (bool): If `True`, this is a dummy step and the load
@@ -550,6 +566,7 @@ def step(
             - `avg_tokens`: The average load across ranks.
             - `max_tokens`: The maximum load across ranks.
             - `balancedness`: The ratio of average load to maximum load.
+
         """
         ep_group = get_ep_group().device_group
         if is_profile:
@@ -728,8 +745,7 @@ def rearrange(
         is_profile: bool = False,
         rank_mapping: dict[int, int] | None = None,
     ) -> torch.Tensor | None:
-        """
-        Rearrange the experts according to the current load.
+        """Rearrange the experts according to the current load.
 
         Args:
             is_profile (bool): If `True`, perform a dummy rearrangement.
@@ -737,8 +753,8 @@ def rearrange(
                 no memory movement will be performed. Default is False.
             rank_mapping (dict[int, int] | None): The rank mapping
                 when scaling is done in EEP.
-        """
 
+        """
         ep_group = get_ep_group().device_group
         ep_rank = ep_group.rank()
 
@@ -747,8 +763,8 @@ def rearrange(
         is_main_rank = ep_rank == 0
         if is_main_rank:
             if not self.is_async or is_profile:
-                start_event = torch.cuda.Event(enable_timing=True)
-                end_event = torch.cuda.Event(enable_timing=True)
+                start_event = torch.Event(enable_timing=True)
+                end_event = torch.Event(enable_timing=True)
                 start_event.record()
             logger.info(
                 "Rearranging experts %s %s...",
@@ -1009,9 +1025,7 @@ def _all_ranks_result_ready(self, model_state: EplbModelState) -> bool:
         return int(flag.item()) == device_group.size()
 
     def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]:
-        """
-        All-reduce a list of tensors.
-        """
+        """All-reduce a list of tensors."""
         ep_group = get_ep_group().device_group
         if len(tensor_list) == 1:
             all_reduce(tensor_list[0], group=ep_group)
@@ -1036,8 +1050,7 @@ def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]
         return all_reduce_list
 
     def _sync_load_pass(self) -> list[torch.Tensor]:
-        """
-        Sync the expert load pass across all ranks for log stats.
+        """Sync the expert load pass across all ranks for log stats.
         Doesn't update the expert load pass in eplb_model_state.
         """
         load_pass_list = []
@@ -1075,10 +1088,14 @@ def update_mapping(
         expanded_physical_to_logical: torch.Tensor,
     ) -> None:
         eplb_model_state = self.model_states[model_config.compute_hash()]
-        eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical)
+        eplb_model_state.physical_to_logical_map_buffer.copy_(
+            expanded_physical_to_logical
+        )
+        eplb_model_state.expert_load_pass.zero_()
+        physical_to_logical_map = eplb_model_state.physical_to_logical_map
 
         (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps(
-            expanded_physical_to_logical.cpu(),
+            physical_to_logical_map.cpu(),
             eplb_model_state.model.num_logical_experts,
         )
 
@@ -1097,6 +1114,33 @@ def update_mapping(
         eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map)
         eplb_model_state.logical_replica_count.copy_(logical_replica_count)
 
+    def reconfigure_physical_expert_slots(
+        self,
+        model_config: ModelConfig,
+        num_physical_experts: int,
+    ) -> None:
+        """Replace physical_to_logical_map and expert_load_pass with views
+        covering the new active size.
+        """
+        model_state = self.model_states[model_config.compute_hash()]
+        old_num_physical_experts = model_state.model.num_physical_experts
+        first_slot, last_slot = sorted((old_num_physical_experts, num_physical_experts))
+        physical_to_logical_map_buffer = model_state.physical_to_logical_map_buffer
+        expert_load_pass_buffer = model_state.expert_load_pass_buffer
+        assert last_slot <= physical_to_logical_map_buffer.shape[1]
+        expert_slots = slice(first_slot, last_slot)
+        physical_to_logical_map_buffer[:, expert_slots].fill_(-1)
+        expert_load_pass_buffer[:, expert_slots].zero_()
+        model_state.physical_to_logical_map = physical_to_logical_map_buffer[
+            :, :num_physical_experts
+        ]
+        model_state.expert_load_pass = expert_load_pass_buffer[:, :num_physical_experts]
+
+        pad_size = num_physical_experts - model_state.expert_load_window.shape[-1]
+        model_state.expert_load_window = torch.nn.functional.pad(
+            model_state.expert_load_window, (0, pad_size)
+        )
+
     def create_communicator(
         self, model_config: ModelConfig, group_coordinator: GroupCoordinator
     ) -> EplbCommunicator:
@@ -1195,8 +1239,7 @@ def compute_logical_maps(
     physical_to_logical_map: torch.Tensor,
     num_logical_experts: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Derive logical_to_physical_map and logical_replica_count from
+    """Derive logical_to_physical_map and logical_replica_count from
     physical_to_logical_map.
 
     Args:
@@ -1209,6 +1252,7 @@ def compute_logical_maps(
             physical slots per logical expert; -1 where unused
         logical_replica_count: [num_layers, num_logical_experts], number of
             physical replicas per logical expert
+
     """
     device = physical_to_logical_map.device
     assert physical_to_logical_map.device.type == "cpu"
@@ -1284,13 +1328,11 @@ def _commit_eplb_maps_for_layer(
     new_physical_to_logical_map: torch.Tensor,
     layer: int,
 ) -> None:
-    """
-    Per-layer version of _commit_eplb_maps that's used by the sync portion of EPLB
+    """Per-layer version of _commit_eplb_maps that's used by the sync portion of EPLB
     when running async EPLB. Copies all of the new_* maps into model_state. After this
     function completes, the new mappings will become the current mappings and will be
     visible to the model.
     """
-
     # Commit physical_to_logical_map
     src = new_physical_to_logical_map
     dst = model_state.physical_to_logical_map[layer]
@@ -1321,26 +1363,23 @@ def _commit_eplb_maps(
     model_state: EplbModelState,
     new_physical_to_logical_map: torch.Tensor,
 ) -> None:
-    """
-    Copies all of the new_* maps into model_state. After this function completes,
+    """Copies all of the new_* maps into model_state. After this function completes,
     the new mappings will become the current mappings and will be visible to the
     model.
     """
-
     # Commit physical_to_logical_map
     src = new_physical_to_logical_map
     dst = model_state.physical_to_logical_map
 
     if PIN_MEMORY and src.is_cpu:
         src = src.new_empty(src.shape, pin_memory=True).copy_(src)
-    # Rare Case: When the number of physical experts has changed, discard the old
-    # physical to logical expert map and use the new one. This only happens when the
-    # number of GPUs available to vLLM changes while vLLM is running. Otherwise copy the
-    # new map into the old one.
+    # When the number of physical experts changes, refresh the active map view to
+    # cover the new size before copying it.
     if src.shape[1] != dst.shape[1]:
-        model_state.physical_to_logical_map = src.to(dst.device, non_blocking=True)
-    else:
-        dst.copy_(src, non_blocking=True)
+        dst = model_state.physical_to_logical_map_buffer[:, : src.shape[1]]
+        assert dst.shape == src.shape
+        model_state.physical_to_logical_map = dst
+    dst.copy_(src, non_blocking=True)
 
     num_logical_experts = model_state.logical_to_physical_map.shape[1]
     new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts)
diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py
index dee19749745a..943f6d9bef06 100644
--- a/vllm/distributed/eplb/eplb_utils.py
+++ b/vllm/distributed/eplb/eplb_utils.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Utility functions for EPLB (Expert Parallel Load Balancing)."""
 
+import contextlib
 import os
 import threading
 
@@ -13,9 +14,24 @@
 logger = init_logger(__name__)
 
 
+@contextlib.contextmanager
+def device_stream(stream: torch.Stream | None):
+    """Platform-agnostic context manager that activates *stream* as the
+    current accelerator stream for the duration of the ``with`` block.
+    A no-op when *stream* is ``None``."""
+    if stream is None:
+        yield
+        return
+    prev = torch.accelerator.current_stream()
+    torch.accelerator.set_stream(stream)
+    try:
+        yield
+    finally:
+        torch.accelerator.set_stream(prev)
+
+
 class CpuGpuEvent:
-    """
-    Combines a CUDA event with a CPU threading event to enforce record->wait
+    """Combines a CUDA event with a CPU threading event to enforce record->wait
     ordering across two threads.
 
     This class is designed for exactly two threads: one producer that calls
@@ -31,12 +47,11 @@ class CpuGpuEvent:
     """
 
     def __init__(self):
-        self._event = torch.cuda.Event()
+        self._event = torch.Event()
         self._recorded = threading.Event()
 
-    def wait(self, stream: torch.cuda.Stream | None = None):
-        """
-        Blocks the calling thread until record finishes. Used to guarantee that the
+    def wait(self, stream: torch.Stream | None = None):
+        """Blocks the calling thread until record finishes. Used to guarantee that the
         record kernel is called before wait.
 
         Should only be called by the Async Eplb thread.
@@ -45,9 +60,8 @@ def wait(self, stream: torch.cuda.Stream | None = None):
         self._event.wait(stream)
         self._recorded.clear()
 
-    def record(self, stream: torch.cuda.Stream | None = None):
-        """
-        Unblocks the waiting thread after calling event.record().
+    def record(self, stream: torch.Stream | None = None):
+        """Unblocks the waiting thread after calling event.record().
 
         Should only be called by the main thread.
         """
@@ -56,7 +70,7 @@ def record(self, stream: torch.cuda.Stream | None = None):
                 "CpuGpuEvent.record() called before the previous event was "
                 "consumed by wait()"
             )
-        self._event = torch.cuda.Event()
+        self._event = torch.Event()
         self._event.record(stream)
         self._recorded.set()
 
@@ -65,12 +79,12 @@ def override_envs_for_eplb(
     parallel_config: ParallelConfig,
     moe_backend: str | None = None,
 ) -> None:
-    """
-    Override environment variables for EPLB when specific conditions are met.
+    """Override environment variables for EPLB when specific conditions are met.
 
     Args:
         parallel_config: The parallel configuration object.
         moe_backend: The configured MoE backend (e.g. ``deep_gemm_mega_moe``).
+
     """
     is_data_parallel = parallel_config.data_parallel_size > 1
     is_eplb_enabled = parallel_config.enable_eplb
diff --git a/vllm/distributed/eplb/policy/abstract.py b/vllm/distributed/eplb/policy/abstract.py
index d056468b97b2..e0c34b9caf2c 100644
--- a/vllm/distributed/eplb/policy/abstract.py
+++ b/vllm/distributed/eplb/policy/abstract.py
@@ -18,10 +18,9 @@ def rebalance_experts(
         num_ranks: int,
         old_global_expert_indices: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Entry point for expert-parallelism load balancer.
+        """Entry point for expert-parallelism load balancer.
 
-        Parameters:
+        Args:
             weight: [layers, num_logical_experts], the load statistics
                 for all logical experts
             num_replicas: number of physical experts, must be a multiple of
@@ -32,8 +31,10 @@ def rebalance_experts(
             old_global_expert_indices: [layers, num_logical_experts], the old global
                 expert indices. Used to avoid unnecessary weight copying
                 for experts moving within one rank.
+
         Returns:
             physical_to_logical_map: [layers, num_replicas], the expert
                 index of each replica
+
         """
         raise NotImplementedError
diff --git a/vllm/distributed/eplb/policy/default.py b/vllm/distributed/eplb/policy/default.py
index c2cdc42909fe..822856ff3271 100644
--- a/vllm/distributed/eplb/policy/default.py
+++ b/vllm/distributed/eplb/policy/default.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Expert parallelism load balancer (EPLB) for vLLM.
+"""Expert parallelism load balancer (EPLB) for vLLM.
 
 This module implements the core rearrangement algorithm.
 
@@ -23,17 +22,17 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
     def balanced_packing(
         cls, weight: np.ndarray, num_packs: int
     ) -> tuple[np.ndarray, np.ndarray]:
-        """
-        Pack n weighted objects to m packs, such that each bin contains exactly
+        """Pack n weighted objects to m packs, such that each bin contains exactly
         n/m objects and the weights of all packs are as balanced as possible.
 
-        Parameters:
+        Args:
             weight: [X, n], the weight of each item
             num_packs: number of packs
 
         Returns:
             pack_index: [X, n], the pack index of each item
             rank_in_pack: [X, n], the rank of the item in the pack
+
         """
         num_layers, num_groups = weight.shape
         assert num_groups % num_packs == 0
@@ -76,17 +75,17 @@ def balanced_packing(
     def replicate_experts(
         cls, weight: np.ndarray, num_phy: int
     ) -> tuple[np.ndarray, np.ndarray]:
-        """
-        Replicate `num_log` experts to `num_phy` replicas, such that the maximum
+        """Replicate `num_log` experts to `num_phy` replicas, such that the maximum
         load of all replicas is minimized.
 
-        Parameters:
+        Args:
             weight: [X, num_log]
             num_phy: total number of experts after replication
 
         Returns:
             phy2log: [X, num_phy], logical expert id of each physical expert
             logcnt: [X, num_log], number of replicas for each logical expert
+
         """
         n, num_log = weight.shape
         num_redundant = num_phy - num_log
@@ -109,8 +108,7 @@ def rebalance_experts_hierarchical(
         num_nodes: int,
         num_gpus: int,
     ) -> np.ndarray:
-        """
-        Parameters:
+        """Parameters
             weight: [num_moe_layers, num_logical_experts]
             num_physical_experts: number of physical experts after replication
             num_groups: number of expert groups
@@ -121,6 +119,7 @@ def rebalance_experts_hierarchical(
         Returns:
             phy2log: [layers, num_replicas], the expert
                 index of each replica
+
         """
         num_layers, num_logical_experts = weight.shape
         assert num_logical_experts % num_groups == 0
@@ -195,8 +194,7 @@ def preserve_intragpu_slots(
         num_ranks: int,
         old_phy2log: np.ndarray,
     ) -> np.ndarray:
-        """
-        Reorder the new mapping per GPU so that experts that remain on the same GPU
+        """Reorder the new mapping per GPU so that experts that remain on the same GPU
         keep their previous slot positions when possible. Incoming experts to that GPU
         fill any remaining available slots. This is applied only when the number of GPUs
         is unchanged and the slots per GPU remain the same between
@@ -281,10 +279,9 @@ def rebalance_experts(
         num_ranks: int,
         old_global_expert_indices: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Entry point for expert-parallelism load balancer.
+        """Entry point for expert-parallelism load balancer.
 
-        Parameters:
+        Args:
             weight: [layers, num_logical_experts], the load statistics for all
                 logical experts
             num_replicas: number of physical experts, must be a multiple of
@@ -296,9 +293,12 @@ def rebalance_experts(
             old_global_expert_indices: [layers, num_logical_experts], the old global
                 expert indices. Used to avoid unnecessary weight copying
                 for experts moving within one rank.
+
         Returns:
+        -------
             phy2log: [layers, num_replicas], the expert
                 index of each replica
+
         """
         weight_np = weight.float().cpu().numpy()
         old_phy2log_np = (
diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index cbd513294565..7a52dd5f56c0 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-The actual execution of the rearrangement.
+"""The actual execution of the rearrangement.
 
 This involves the exchange of expert weights between GPUs.
 """
@@ -14,7 +13,7 @@
 from torch.distributed import ProcessGroup, all_gather
 
 from vllm.distributed.eplb.eplb_communicator import EplbCommunicator
-from vllm.distributed.eplb.eplb_utils import CpuGpuEvent
+from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream
 from vllm.logger import init_logger
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 
@@ -41,9 +40,7 @@ class TransferMetadata:
 
 @dataclass
 class AsyncEplbLayerResult:
-    """
-    The result of one completed async EPLB layer transfer.
-    """
+    """The result of one completed async EPLB layer transfer."""
 
     layer_idx: int
     """Index of the MoE layer that was transferred."""
@@ -69,8 +66,7 @@ def get_ep_ranks_with_experts_batch(
     old_indices: np.ndarray,
     new_indices: np.ndarray,
 ) -> tuple[dict[int, list[int]], dict[int, list[int]]]:
-    """
-    Get the ranks of the experts that need to be exchanged.
+    """Get the ranks of the experts that need to be exchanged.
 
     Args:
         expert_ids: 1D array of expert indices to query.
@@ -82,6 +78,7 @@ def get_ep_ranks_with_experts_batch(
         A tuple of two dictionaries mapping expert_id to:
         - ranks_to_send: The ranks that have this expert and need to send.
         - ranks_to_recv: The ranks that need to receive this expert.
+
     """
     ranks_to_send_map: dict[int, list[int]] = {}
     ranks_to_recv_map: dict[int, list[int]] = {}
@@ -176,29 +173,29 @@ def move_to_buffer(
     new_indices: np.ndarray,
     expert_weights: Sequence[torch.Tensor],
     expert_weights_buffers: Sequence[torch.Tensor],
-    cuda_stream: torch.cuda.Stream | None,
+    stream: torch.Stream | None,
     ep_rank: int,
     communicator: EplbCommunicator,
     layer_idx: int = 0,
 ) -> TransferMetadata:
-    """
-    Rearranges expert weights during EPLB rebalancing.
+    """Rearranges expert weights during EPLB rebalancing.
 
     Args:
         num_local_experts: Number of local experts.
         old_indices: (num_experts_total,) ndarray of current (old)
-            global-to-local expert assignments.
+            global-to-locals expert assignments.
         new_indices: (num_experts_total,) ndarray of desired (new)
             global-to-local assignments after rebalance.
         expert_weights: Original expert weights for the layer.
         expert_weights_buffers: Intermediate buffers (one per tensor).
-        cuda_stream: CUDA stream for async copies (can be None for sync mode).
+        stream: CUDA/XPU stream for async copies (can be None for sync mode).
         ep_rank: Rank of this process in expert parallel group.
         communicator: EplbCommunicator instance for P2P communication.
         layer_idx: Index of the MoE layer being transferred.
 
     Returns:
         TransferMetadata: Metadata needed for completing remote weight transfers.
+
     """
     assert old_indices.shape == new_indices.shape
     recv_primary_mask = np.zeros((num_local_experts,), dtype=np.bool_)
@@ -264,10 +261,9 @@ def move_to_buffer(
             expert = new_local_expert_ids[dst]
             src_local = expert_to_src_map.get(expert, -1)
             if src_local != -1:
-                with torch.cuda.stream(cuda_stream):
+                with device_stream(stream):
                     for w, b in zip(expert_weights, expert_weights_buffers):
                         b[dst].copy_(w[src_local], non_blocking=True)
-
     communicator.set_transfer_context(old_indices, layer_idx)
 
     # 2. Post sends
@@ -355,9 +351,8 @@ def move_from_buffer(
     new_indices: np.ndarray,
     ep_rank: int,
 ) -> None:
-    """
-    Copies expert weights from communication buffers back to the target weight tensors
-    after EPLB rebalancing.
+    """Copies expert weights from communication buffers back to the target weight
+    tensors after EPLB rebalancing.
 
     Args:
         expert_weights: List of the actual MoE layer weights used in the execution.
@@ -367,6 +362,7 @@ def move_from_buffer(
         new_indices: (num_experts_total,) mapping from local rows to desired
             (possibly global) expert id, after rebalance.
         ep_rank: Rank of the process in the expert parallel group.
+
     """
     is_unchanged = transfer_metadata.is_unchanged
     is_received_locally = transfer_metadata.is_received_locally
@@ -433,12 +429,11 @@ def transfer_layer(
     ep_group: ProcessGroup,
     communicator: EplbCommunicator,
     is_profile: bool = False,
-    cuda_stream: torch.cuda.Stream | None = None,
+    stream: torch.Stream | None = None,
     rank_mapping: dict[int, int] | None = None,
     layer_idx: int = 0,
 ) -> TransferMetadata:
-    """
-    Rearranges the expert weights in place according to the new expert indices.
+    """Rearranges the expert weights in place according to the new expert indices.
 
     The value of the indices arguments are logical indices of the experts,
     while keys are physical.
@@ -455,13 +450,14 @@ def transfer_layer(
         is_profile (bool): If `True`, do not perform any actual weight copy.
             This is used during profile run, where we only perform dummy
             communications to reserve enough memory for the buffers.
-        cuda_stream: CUDA stream for async copies (can be None for sync mode).
+        stream: CUDA stream for async copies (can be None for sync mode).
         rank_mapping: Optional rank mapping for elastic expert parallelism.
         layer_idx: Index of the MoE layer being transferred.
 
     Returns:
         TransferMetadata: Metadata needed for completing remote weight transfers,
             including is_unchanged and is_received_locally masks.
+
     """
     ep_size = ep_group.size()
     if rank_mapping is not None:
@@ -502,7 +498,7 @@ def transfer_layer(
         new_indices=new_layer_indices_np,
         expert_weights=expert_weights,
         expert_weights_buffers=expert_weights_buffer,
-        cuda_stream=cuda_stream,
+        stream=stream,
         ep_rank=ep_group.rank(),
         communicator=communicator,
         layer_idx=layer_idx,
@@ -519,8 +515,7 @@ def rearrange_expert_weights_inplace(
     is_profile: bool = False,
     rank_mapping: dict[int, int] | None = None,
 ) -> None:
-    """
-    Rearranges the expert weights in place according to the new expert indices.
+    """Rearranges the expert weights in place according to the new expert indices.
 
     The value of the indices arguments are logical indices of the experts,
     while keys are physical.
@@ -540,6 +535,7 @@ def rearrange_expert_weights_inplace(
             This is used during profile run, where we only perform dummy
             communications to reserve enough memory for the buffers.
         rank_mapping: A dictionary mapping old rank to new rank.
+
     """
     if rank_mapping is not None:
         if len(rank_mapping) == ep_group.size():
@@ -604,7 +600,7 @@ def rearrange_expert_weights_inplace(
             new_indices=new_global_expert_indices_cpu[layer_idx],
             expert_weights=expert_weights[layer_idx],
             expert_weights_buffers=weights_buffer,
-            cuda_stream=None,
+            stream=None,
             ep_rank=ep_rank,
             communicator=communicator,
             layer_idx=layer_idx,
@@ -624,8 +620,7 @@ def _map_old_expert_indices_with_rank_mapping(
     rank_mapping: dict[int, int],
     new_ep_size: int,
 ) -> torch.Tensor:
-    """
-    Map the old global expert indices to the new global expert indices.
+    """Map the old global expert indices to the new global expert indices.
 
     Args:
         old_global_expert_indices:
@@ -636,6 +631,7 @@ def _map_old_expert_indices_with_rank_mapping(
     Returns:
         Mapped expert indices with shape
         (num_layers, new_ep_size * num_local_physical_experts).
+
     """
     num_layers, old_num_physical_experts = old_global_expert_indices.shape
     assert rank_mapping, "Rank mapping is required"
diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py
index 96ddd7ce950e..9d7ec2f827f4 100644
--- a/vllm/distributed/kv_events.py
+++ b/vllm/distributed/kv_events.py
@@ -17,6 +17,12 @@
 
 from vllm.config.kv_events import KVEventsConfig
 from vllm.logger import init_logger
+from vllm.utils.network_utils import (
+    get_ip,
+    get_tcp_uri,
+    is_valid_ipv6_address,
+    split_zmq_path,
+)
 from vllm.v1.core.kv_cache_utils import ExternalBlockHash
 
 logger = init_logger(__name__)
@@ -39,7 +45,7 @@ class KVCacheEvent(
     gc=False,  # type: ignore[call-arg]
     tag=True,
 ):
-    """Base class for all KV cache-related events"""
+    """Base class for all KV cache-related events."""
 
 
 MEDIUM_GPU = "GPU"
@@ -135,8 +141,7 @@ class KVEventBatch(EventBatch):
 
 
 class KVEventAggregator:
-    """
-    Aggregates KV events across multiple workers.
+    """Aggregates KV events across multiple workers.
     Tracks how many times each event appears and returns only those
     that were emitted by all workers.
     """
@@ -150,22 +155,22 @@ def __init__(self, num_workers: int) -> None:
         self._num_workers: int = num_workers
 
     def add_events(self, events: list[KVCacheEvent]) -> None:
-        """
-        Add events from a worker batch.
+        """Add events from a worker batch.
 
         Args:
             events: List of KVCacheEvent objects.
+
         """
         if not isinstance(events, list):
             raise TypeError("events must be a list of KVCacheEvent.")
         self._event_counter.update(events)
 
     def get_common_events(self) -> list[KVCacheEvent]:
-        """
-        Return events that appeared in all workers.
+        """Return events that appeared in all workers.
 
         Returns:
             List of events present in all workers.
+
         """
         return [
             event
@@ -174,43 +179,39 @@ def get_common_events(self) -> list[KVCacheEvent]:
         ]
 
     def get_all_events(self) -> list[KVCacheEvent]:
-        """
-        Return all events for all workers.
+        """Return all events for all workers.
 
         Returns:
             List of events for all workers.
+
         """
         return list(self._event_counter.elements())
 
     def clear_events(self) -> None:
-        """
-        Clear all tracked events.
-        """
+        """Clear all tracked events."""
         self._event_counter.clear()
 
     def increment_workers(self, count: int = 1) -> None:
-        """
-        Increment the number of workers contributing events.
+        """Increment the number of workers contributing events.
 
         Args:
             count: Number to increment the workers by.
+
         """
         if count <= 0:
             raise ValueError("count must be positive.")
         self._num_workers += count
 
     def reset_workers(self) -> None:
-        """
-        Reset the number of workers to 1.
-        """
+        """Reset the number of workers to 1."""
         self._num_workers = 1
 
     def get_number_of_workers(self) -> int:
-        """
-        Return the number of workers.
+        """Return the number of workers.
 
         Returns:
             int number of workers.
+
         """
         return self._num_workers
 
@@ -222,8 +223,7 @@ def __repr__(self) -> str:
 
 
 class KVConnectorKVEvents(ABC):
-    """
-    Abstract base class for KV events.
+    """Abstract base class for KV events.
     Acts as a container for KV events from the connector.
     """
 
@@ -323,6 +323,7 @@ class ZmqEventPublisher(EventPublisher):
         Maximum number of events to buffer in memory.
     topic:
         Topic to publish events to.
+
     """
 
     SHUTDOWN_TIMEOUT: float = 1.0
@@ -354,6 +355,8 @@ def __init__(
             replay_endpoint, self._dp_rank
         )
         assert self._endpoint is not None
+        self._hwm = hwm
+        self._socket_setup()
         self._publisher_config = KVEventsConfig(
             enable_kv_cache_events=True,
             publisher="zmq",
@@ -364,8 +367,6 @@ def __init__(
             max_queue_size=max_queue_size,
             topic=topic,
         )
-        self._hwm = hwm
-        self._socket_setup()
 
         # Payload
         self._seq_gen = count()
@@ -428,16 +429,10 @@ def _socket_setup(self) -> None:
         if self._pub is None:
             self._pub = self._ctx.socket(zmq.PUB)
             self._pub.set_hwm(self._hwm)
-            # Heuristic: bind if wildcard / * present, else connect.
-            # bind stable, connect volatile convention
-            if self._endpoint is not None and (
-                "*" in self._endpoint
-                or "::" in self._endpoint
-                or self._endpoint.startswith("ipc://")
-                or self._endpoint.startswith("inproc://")
-            ):
-                self._pub.bind(self._endpoint)
-            elif self._endpoint is not None:
+            assert self._endpoint is not None
+            if self._is_bind_endpoint(self._endpoint):
+                self._endpoint = self._bind(self._pub, self._endpoint)
+            else:
                 self._pub.connect(self._endpoint)
 
         # Set up replay socket: use ROUTER
@@ -446,7 +441,31 @@ def _socket_setup(self) -> None:
         # 3) works in our non‑blocking poll loop alongside PUB
         if self._replay_endpoint is not None:
             self._replay = self._ctx.socket(zmq.ROUTER)
-            self._replay.bind(self._replay_endpoint)
+            self._replay_endpoint = self._bind(self._replay, self._replay_endpoint)
+
+    @staticmethod
+    def _is_bind_endpoint(endpoint: str) -> bool:
+        """Bind wildcard and ipc/inproc endpoints; connect to the rest."""
+        return (
+            "*" in endpoint
+            or "::" in endpoint
+            or endpoint.startswith(("ipc://", "inproc://"))
+        )
+
+    @staticmethod
+    def _bind(sock: zmq.Socket, endpoint: str) -> str:
+        """Bind; tcp://*:0 binds this node's IP and returns the bound endpoint."""
+        if endpoint != "tcp://*:0":
+            sock.bind(endpoint)
+            return endpoint
+        ip = get_ip()
+        if ip in ("0.0.0.0", "::"):
+            raise ValueError("tcp://*:0 needs a node IP; set VLLM_HOST_IP")
+        if is_valid_ipv6_address(ip):
+            sock.setsockopt(zmq.IPV6, 1)
+        sock.bind(get_tcp_uri(ip, 0))
+        _, _, port = split_zmq_path(sock.getsockopt_string(zmq.LAST_ENDPOINT))
+        return get_tcp_uri(ip, int(port))
 
     def _publisher_thread(self) -> None:
         """Background thread that processes the event queue."""
@@ -519,7 +538,8 @@ def offset_endpoint_port(
 
         Returns:
             The endpoint with the port offset by data_parallel_rank
-                or suffix appended
+                or suffix appended; a tcp port of 0 is left for the OS.
+
         """
         # Do nothing if input is None or data_parallel_rank is 0
         if not endpoint or data_parallel_rank == 0:
@@ -533,6 +553,8 @@ def offset_endpoint_port(
                 last_colon_idx = endpoint.rfind(":")
                 base_addr = endpoint[:last_colon_idx]
                 base_port = int(endpoint[last_colon_idx + 1 :])
+                if base_port == 0:
+                    return endpoint
                 new_port = base_port + data_parallel_rank
                 return f"{base_addr}:{new_port}"
             return endpoint
diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py
index aad7999d08a2..97064e06cdb8 100644
--- a/vllm/distributed/kv_transfer/kv_connector/factory.py
+++ b/vllm/distributed/kv_transfer/kv_connector/factory.py
@@ -87,6 +87,7 @@ def get_connector_class_by_name(
 
         Returns:
             The connector class.
+
         """
         if connector_name not in cls._registry:
             raise ValueError(f"Connector '{connector_name}' is not registered.")
@@ -197,6 +198,12 @@ def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool:
     "MultiConnector",
 )
 
+KVConnectorFactory.register_connector(
+    "HiSparseConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector",
+    "HiSparseConnector",
+)
+
 KVConnectorFactory.register_connector(
     "MoRIIOConnector",
     "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector",
diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py
index 39f8d8a1b927..b645e97c4dd4 100644
--- a/vllm/distributed/kv_transfer/kv_connector/utils.py
+++ b/vllm/distributed/kv_transfer/kv_connector/utils.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-KV cache helper for store.
-"""
+"""KV cache helper for store."""
 
 from collections.abc import Iterator
 from dataclasses import dataclass
@@ -60,6 +58,7 @@ def __init__(self, expected_finished_count: int):
         # [req_id -> n_remaining_workers]
         self._recv_remaining_count = dict[str, int]()
         self._send_remaining_count = dict[str, int]()
+        self._failed_recving_pending = set[str]()
         self._expected_finished_count = expected_finished_count
 
     @classmethod
@@ -156,6 +155,10 @@ def update_finished_set(
                 combined_kv_cache_events.increment_workers(1)
 
             invalid_block_ids |= kv_output.invalid_block_ids
+            self._failed_recving_pending |= kv_output.failed_recving
+
+        failed_recving = self._failed_recving_pending & finished_recving
+        self._failed_recving_pending -= failed_recving
 
         # select output of the worker specified by output_rank
         output = outputs[output_rank]
@@ -168,6 +171,7 @@ def update_finished_set(
             kv_cache_events=combined_kv_cache_events or None,
             kv_connector_worker_meta=aggregated_kv_connector_worker_meta or None,
             invalid_block_ids=invalid_block_ids,
+            failed_recving=failed_recving,
             expected_finished_count=self._expected_finished_count,
         )
 
@@ -227,11 +231,10 @@ def copy_kv_blocks(
 
 
 def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio):
-    """
-    Transforms the layout of received KV cache blocks to the local block_size.
+    """Transforms the layout of received KV cache blocks to the local block_size.
     (Only works for local blocksize > remote blocksize)
 
-    example:
+    Example:
     local blocksize = 16 tokens, remote blocksize = 4 tokens
     local block[0] = remote block[0, 1, 2, 3]
     remote is |h0-b0|h1-b0|h2-b0|h3-b0|h0-b1|h1-b1|h2-b1|h3-b1|...
@@ -240,6 +243,7 @@ def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio):
     1. view => view remote as n_blocks * remote_shape(H,remoteN,D)
     2. permute => (H, nblocks, remoteN, D)
     3. flatten => (H, localN, D)
+
     """
     blocks_to_update = cache.index_select(0, indices)
     # use physical order
@@ -287,8 +291,7 @@ def kv_postprocess_layout_on_receive(cache, indices):
 
 
 def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio):
-    """
-    Transforms the layout of received KV cache to the local block_size and LBHNC.
+    """Transforms the layout of received KV cache to the local block_size and LBHNC.
     (Only works for local blocksize > remote blocksize)
 
     prefill is LBHNC, smaller block_size
@@ -311,9 +314,9 @@ def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_rati
 def yield_req_data(
     scheduler_output,
 ) -> Iterator[tuple[str, tuple[list[int], ...] | None, bool]]:
-    """
-    Yields:
-        (req_id, new_block_id_groups, preempted)
+    """Yields:
+    (req_id, new_block_id_groups, preempted)
+
     """
     # new requests
     for req_data in scheduler_output.scheduled_new_reqs:
@@ -340,6 +343,7 @@ def get_current_attn_backends(
 
     Returns:
         Deduplicated list of attention backend classes.
+
     """
     layer_type = cast(type[Any], AttentionLayerBase)
     layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py
index fa764571acf2..072018c039ed 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-KVConnectorBase_V1 Class for Distributed KV Cache & Hidden State
+"""KVConnectorBase_V1 Class for Distributed KV Cache & Hidden State
 communication in vLLM v1
 
 The class provides the following primitives:
@@ -34,8 +33,8 @@
         save_kv_layer() - starts saving KV for layer i (maybe async)
         wait_for_save() - blocks until all saves are done
 
-        get_finished() - called with ids of finished requests, returns
-            ids of requests that have completed async sending/recving.
+        get_transfer_results() - returns async send/receive completions and
+            receive failures in one snapshot.
         build_connector_worker_meta() - builds metadata to be sent
             back to the scheduler-side connector
 """
@@ -43,6 +42,7 @@
 import enum
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Literal
 
 import torch
@@ -63,7 +63,7 @@
     )
     from vllm.forward_context import ForwardContext
     from vllm.v1.core.block_pool import BlockPool
-    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
     from vllm.v1.kv_cache_interface import KVCacheConfig
     from vllm.v1.request import Request
 
@@ -82,9 +82,21 @@
 logger = init_logger(__name__)
 
 
-class SupportsHMA(ABC):
+@dataclass
+class KVConnectorTransferResults:
+    """Asynchronous transfer completions from one worker snapshot.
+
+    Failed receives also appear in ``finished_recving`` so the scheduler can
+    release the request from its transfer wait state.
     """
-    The class that indicates the corresponding connector supports hybrid memory
+
+    finished_sending: set[str] = field(default_factory=set)
+    finished_recving: set[str] = field(default_factory=set)
+    failed_recving: set[str] = field(default_factory=set)
+
+
+class SupportsHMA(ABC):
+    """The class that indicates the corresponding connector supports hybrid memory
     allocator (HMA).
     This is required to use the connector together with hybrid memory allocator.
     """
@@ -95,8 +107,7 @@ def request_finished_all_groups(
         request: "Request",
         block_ids: tuple[list[int], ...],
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Called exactly once when a request has finished for all kv cache groups,
+        """Called exactly once when a request has finished for all kv cache groups,
         before its blocks are freed for each group.
 
         NOTE(Kuntai): This function is only supported by connectors that support HMA.
@@ -110,6 +121,7 @@ def request_finished_all_groups(
             get_finished().
             Optional KVTransferParams to be included in the request outputs
             returned by the engine.
+
         """
         raise NotImplementedError
 
@@ -130,8 +142,7 @@ class KVConnectorRole(enum.Enum):
 
 
 class KVConnectorHandshakeMetadata(ABC):  # noqa: B024
-    """
-    Metadata used for out of band connector handshake between
+    """Metadata used for out of band connector handshake between
     P/D workers. This needs to serializable.
     """
 
@@ -139,8 +150,7 @@ class KVConnectorHandshakeMetadata(ABC):  # noqa: B024
 
 
 class KVConnectorMetadata(ABC):  # noqa: B024
-    """
-    Abstract Metadata used to communicate
+    """Abstract Metadata used to communicate
     Scheduler KVConnector -> Worker KVConnector.
     """
 
@@ -148,8 +158,7 @@ class KVConnectorMetadata(ABC):  # noqa: B024
 
 
 class KVConnectorWorkerMetadata(ABC):
-    """
-    Abstract Metadata used to communicate back
+    """Abstract Metadata used to communicate back
     Worker KVConnector -> Scheduler KVConnector.
 
     Each worker can output its own metadata.
@@ -162,16 +171,12 @@ class KVConnectorWorkerMetadata(ABC):
     def aggregate(
         self, other: "KVConnectorWorkerMetadata"
     ) -> "KVConnectorWorkerMetadata":
-        """
-        Aggregate metadata with another `KVConnectorWorkerMetadata` object.
-        """
+        """Aggregate metadata with another `KVConnectorWorkerMetadata` object."""
         pass
 
 
 class KVConnectorBase_V1(ABC):
-    """
-    Base class for KV connectors.
-    """
+    """Base class for KV connectors."""
 
     @property
     def supports_divergent_local_hybrid_hits(self) -> bool:
@@ -211,6 +216,7 @@ def __init__(
         else:
             raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1")
         self._kv_cache_config = kv_cache_config
+        self._kv_cache_manager: KVCacheManager | None = None
         self._role = role
 
     @property
@@ -230,6 +236,7 @@ def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> No
 
         Args:
             connector_metadata (dict): the connector metadata.
+
         """
         self._connector_metadata = connector_metadata
 
@@ -248,6 +255,7 @@ def _get_connector_metadata(self) -> KVConnectorMetadata:
 
         Returns:
             ConnectorMetadata: the connector metadata.
+
         """
         # Should only be called while set to valid metadata.
         assert self._connector_metadata is not None
@@ -258,37 +266,44 @@ def has_connector_metadata(self) -> bool:
 
         Returns:
             bool: True if connector metadata exists, False otherwise.
+
         """
         return self._connector_metadata is not None
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
-        """
-        Initialize with the KV caches. Useful for pre-registering the
+        """Initialize with the KV caches. Useful for pre-registering the
         KV Caches in the KVConnector (e.g. for NIXL).
 
         Args:
             kv_caches: dictionary of layer names, kv cache
+
         """
         return
 
     def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
-        """
-        Set the xPU-specific ops for copying KV between host and device.
+        """Set the xPU-specific ops for copying KV between host and device.
         Needed when host buffer is used for kv transfer (e.g., in NixlConnector)
         """
         return
 
+    # TODO(NickLucche): group model-runner lifecycle hooks in the interface.
+    def finish_forward(self) -> None:
+        """Notify the connector that the model no longer reads this step's KV."""
+        return
+
+    def reset_capture_state(self) -> None:
+        """Reset worker state mutated while capturing CUDA graphs."""
+        return
+
     def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
-        """
-        Handle preempted requests or evicted blocks BEFORE they are overwritten.
+        """Handle preempted requests or evicted blocks BEFORE they are overwritten.
         Needed for connectors which use async saves (e.g., OffloadingConnector)
         """
         return
 
     @abstractmethod
     def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
-        """
-        Start loading the KV cache from the connector to vLLM's paged
+        """Start loading the KV cache from the connector to vLLM's paged
         KV buffer. Loads required by the current forward start before it;
         independent asynchronous loads may start after it is submitted.
 
@@ -305,8 +320,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non
 
     @abstractmethod
     def wait_for_layer_load(self, layer_name: str) -> None:
-        """
-        Block until the KV for a specific layer is loaded into vLLM's
+        """Block until the KV for a specific layer is loaded into vLLM's
         paged buffer. This is called from within attention layer to ensure
         async copying from start_load_kv is complete.
 
@@ -314,6 +328,7 @@ def wait_for_layer_load(self, layer_name: str) -> None:
 
         Args:
             layer_name: the name of that layer
+
         """
         pass
 
@@ -325,8 +340,7 @@ def save_kv_layer(
         attn_metadata: "AttentionMetadata",
         **kwargs: Any,
     ) -> None:
-        """
-        Start saving a layer of KV cache from vLLM's paged buffer
+        """Start saving a layer of KV cache from vLLM's paged buffer
         to the connector. This is called from within attention layer to
         enable async copying during execution.
 
@@ -336,13 +350,13 @@ def save_kv_layer(
                 layer in vLLM.
             attn_metadata (AttentionMetadata): the attention metadata.
             **kwargs: additional arguments for the save operation.
+
         """
         pass
 
     @abstractmethod
     def wait_for_save(self):
-        """
-        Block until all the save operations is done. This is called
+        """Block until all the save operations is done. This is called
         as the forward context exits to ensure that the async saving
         from save_kv_layer is complete before finishing the forward.
 
@@ -353,8 +367,7 @@ def wait_for_save(self):
     def get_finished(
         self, finished_req_ids: set[str]
     ) -> tuple[set[str] | None, set[str] | None]:
-        """
-        Notifies worker-side connector ids of requests that have
+        """Notifies worker-side connector ids of requests that have
         finished generating tokens on the worker.
         The scheduler process (via the Executors) will use this output
         to track which workers are done.
@@ -365,12 +378,22 @@ def get_finished(
             tuple of (sending/saving ids, recving/loading ids).
             The finished saves/sends req ids must belong to a set provided in a
             call to this method (this call or a prior one).
+
         """
         return None, None
 
+    def get_transfer_results(
+        self, finished_req_ids: set[str]
+    ) -> KVConnectorTransferResults:
+        """Return completed sends, receives, and receive failures together."""
+        finished_sending, finished_recving = self.get_finished(finished_req_ids)
+        return KVConnectorTransferResults(
+            finished_sending=set(finished_sending or ()),
+            finished_recving=set(finished_recving or ()),
+        )
+
     def get_block_ids_with_load_errors(self) -> set[int]:
-        """
-        Get the set of block IDs that failed to load.
+        """Get the set of block IDs that failed to load.
 
         Returns:
             Set of block IDs that encountered load errors.
@@ -380,55 +403,52 @@ def get_block_ids_with_load_errors(self) -> set[int]:
             - Applies to both sync- and async-loading requests.
             - Async loading: failed blocks may be reported in any forward pass
               up to and including the pass where the request ID is returned by
-              `get_finished()`. Even if failures occur, the request must still
-              be reported via `get_finished()`, and the failed block IDs must
-              appear here no later than that same pass.
+              `get_transfer_results()`. Even if failures occur, the request
+              must still be reported as finished receiving, and the failed
+              block IDs must appear here no later than that same pass.
             - Sync loading: failed blocks should be reported in the forward
               pass in which they are detected.
+
         """
         return set()
 
     def shutdown(self):
-        """
-        Shutdown the connector. This is called when the worker process
+        """Shutdown the connector. This is called when the worker process
         is shutting down to ensure that all the async operations are
         completed and the connector is cleaned up properly.
         """
         return None
 
     def get_kv_connector_stats(self) -> "KVConnectorStats | None":
-        """
-        Get the KV connector stats collected during the last interval.
-        """
+        """Get the KV connector stats collected during the last interval."""
         return None
 
     def get_kv_connector_kv_cache_events(self) -> "KVConnectorKVEvents | None":
-        """
-        Get the KV connector kv cache events collected during the last interval.
+        """Get the KV connector kv cache events collected during the last interval.
         This function should be called by the model runner every time after the
         model execution and before cleanup.
         """
         return None
 
     def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
-        """
-        Get the KVConnector handshake metadata for this connector.
+        """Get the KVConnector handshake metadata for this connector.
         This metadata is used for out-of-band connector handshake
         between P/D workers.
 
         Returns:
             KVConnectorHandshakeMetadata: the handshake metadata.
             None if no handshake metadata is available.
+
         """
         return None
 
     def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
-        """
-        Build the KVConnector worker metadata for this engine step.
+        """Build the KVConnector worker metadata for this engine step.
 
         Returns:
             KVConnectorWorkerMetadata: the worker metadata.
             None if no worker metadata is available.
+
         """
         return None
 
@@ -436,13 +456,18 @@ def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
     # Scheduler-side methods
     # ==============================
 
+    def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None:
+        """Bind the scheduler's cache manager after it has been constructed."""
+        self._kv_cache_manager = kv_cache_manager
+        self.bind_gpu_block_pool(kv_cache_manager.block_pool)
+
     def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
-        """
-        Bind the GPU block pool to the connector for per-GPU block status tracking.
+        """Bind the GPU block pool to the connector for per-GPU block status tracking.
         For example, inc/dec ref counts, or iterate over the prefix cache blocks.
 
         Args:
             gpu_block_pool: the GPU block pool.
+
         """
         return
 
@@ -452,8 +477,7 @@ def get_num_new_matched_tokens(
         request: "Request",
         num_computed_tokens: int,
     ) -> tuple[int | None, bool]:
-        """
-        Get number of new tokens that can be loaded from the
+        """Get number of new tokens that can be loaded from the
         external KV cache beyond the num_computed_tokens.
 
         Args:
@@ -478,6 +502,7 @@ def get_num_new_matched_tokens(
             call. If the cache cannot be loaded for some tokens (e.g., due to
             connectivity issues or eviction), those tokens must not be taken
             into account.
+
         """
         pass
 
@@ -485,8 +510,7 @@ def get_num_new_matched_tokens(
     def update_state_after_alloc(
         self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
     ):
-        """
-        Update KVConnector state after block allocation.
+        """Update KVConnector state after block allocation.
 
         If get_num_new_matched_tokens previously returned True for a
         request, this function may be called twice for that same request -
@@ -504,6 +528,7 @@ def update_state_after_alloc(
             blocks (KVCacheBlocks): the blocks allocated for the request.
             num_external_tokens (int): the number of tokens to load from the
                 external KV cache. 0 means nothing should be loaded.
+
         """
         pass
 
@@ -511,14 +536,14 @@ def update_state_after_alloc(
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:
-        """
-        Build the connector metadata for this step.
+        """Build the connector metadata for this step.
 
         This function should NOT modify fields in the scheduler_output.
         Also, calling this function will reset the state of the connector.
 
         Args:
             scheduler_output (SchedulerOutput): the scheduler output object.
+
         """
         pass
 
@@ -531,12 +556,12 @@ def on_new_request(self, request: "Request") -> None:
         return
 
     def update_connector_output(self, connector_output: KVConnectorOutput):
-        """
-        Update KVConnector state from worker-side connectors output.
+        """Update KVConnector state from worker-side connectors output.
 
         Args:
             connector_output (KVConnectorOutput): the worker-side
                 connectors output.
+
         """
         return
 
@@ -545,8 +570,7 @@ def request_finished(
         request: "Request",
         block_ids: list[int],
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Called exactly once when a request has finished, before its blocks are
+        """Called exactly once when a request has finished, before its blocks are
         freed.
 
         The connector may assumes responsibility for freeing the blocks
@@ -558,6 +582,7 @@ def request_finished(
             get_finished().
             Optional KVTransferParams to be included in the request outputs
             returned by the engine.
+
         """
         return False, None
 
@@ -575,14 +600,18 @@ def register_finished_partial_tail(
         return False
 
     def take_events(self) -> Iterable["KVCacheEvent"]:
-        """
-        Take the KV cache events from the connector.
+        """Take the KV cache events from the connector.
 
         Yields:
             New KV cache events since the last call.
+
         """
         return ()
 
+    def has_pending_block_frees(self) -> bool:
+        """Whether pending transfers can release blocks instead of preemption."""
+        return False
+
     def has_pending_push_work(self) -> bool:
         """Return True if the connector has push-mode work that requires
         the engine main loop to keep stepping (e.g. a P-side request whose
@@ -597,16 +626,16 @@ def has_pending_push_work(self) -> bool:
 
     @classmethod
     def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
-        """
-        Get the required KV cache layout for this connector.
+        """Get the required KV cache layout for this connector.
+
         Args:
             vllm_config (VllmConfig): the vllm config.
 
         Returns:
             str: the required KV cache layout. e.g. HND, or NHD.
             None if the connector does not require a specific layout.
-        """
 
+        """
         if cls is KVConnectorBase_V1:
             raise TypeError(
                 "get_required_kvcache_layout should not be called "
@@ -616,8 +645,7 @@ def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
 
     @classmethod
     def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
-        """
-        Check if this connector requires PIECEWISE CUDA graph mode.
+        """Check if this connector requires PIECEWISE CUDA graph mode.
 
         Connectors that use asynchronous layer-by-layer operations
         (wait_for_layer_load/save_kv_layer) should override this method
@@ -633,27 +661,26 @@ def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
         Returns:
             True if this connector requires PIECEWISE CUDA graph mode,
             False otherwise.
+
         """
         return False
 
     def get_finished_count(self) -> int | None:
-        """
-        Get the count of requests expected to complete send/receive operations
+        """Get the count of requests expected to complete send/receive operations
         via this connector. This method is used to initialize the
         KVOutputAggregator, overwriting the default world_size.
 
         Returns:
             int: expected sending or receiving completion count.
-        """
 
+        """
         return None
 
     @classmethod
     def build_kv_connector_stats(
         cls, data: dict[str, Any] | None = None
     ) -> "KVConnectorStats | None":
-        """
-        KVConnectorStats resolution method. This method allows dynamically
+        """KVConnectorStats resolution method. This method allows dynamically
         registered connectors to return their own KVConnectorStats object,
         which can implement custom aggregation logic on the data dict.
         """
@@ -662,19 +689,18 @@ def build_kv_connector_stats(
     def set_xfer_handshake_metadata(
         self, metadata: dict[int, KVConnectorHandshakeMetadata]
     ) -> None:
-        """
-        Set the KV connector handshake metadata for this connector.
+        """Set the KV connector handshake metadata for this connector.
 
         Args:
             metadata (KVConnectorHandshakeMetadata): the handshake metadata to set.
+
         """
         return None
 
     def set_xfer_handshake_metadata_pp_aware(
         self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
     ) -> None:
-        """
-        Set handshake metadata keyed by (pp_rank, tp_rank).
+        """Set handshake metadata keyed by (pp_rank, tp_rank).
         - Default implementation assumes pp_rank is always 0
         - PP-aware connectors override this to consume all PP producer shards.
         """
@@ -695,19 +721,18 @@ def build_prom_metrics(
         labelnames: list[str],
         per_engine_labelvalues: dict[int, list[object]],
     ) -> "KVConnectorPromMetrics | None":
-        """
-        Create a KVConnectorPromMetrics subclass which should register
+        """Create a KVConnectorPromMetrics subclass which should register
         per-connector Prometheus metrics and implement observe() to
         expose connector transfer stats via Prometheus.
         """
         return None
 
     def reset_cache(self) -> bool | None:
-        """
-        Reset the connector's internal cache.
+        """Reset the connector's internal cache.
 
         Returns:
             bool: True if the cache was successfully reset, False otherwise.
+
         """
         logger.debug(
             "Connector cache reset requested, but %s does not implement reset_cache().",
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
index 83b468d26551..cff2367966fc 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-DecodeBenchConnector: A KV Connector for decode instance performance testing.
+"""DecodeBenchConnector: A KV Connector for decode instance performance testing.
 
 This connector emulates a prefill-decode disaggregated setting by filling
 the KV cache with dummy values, allowing measurement of decoder performance
@@ -10,7 +9,7 @@
 Usage:
     To use this connector for benchmarking, configure it in the kv_transfer_config:
 
-    Example:
+Example:
         vllm serve <model> --kv-transfer-config '{
             "kv_connector": "DecodeBenchConnector",
             "kv_role": "kv_both",
@@ -29,6 +28,7 @@
         - fill_mean (float): Mean value for random normal fill (default: 0.015)
         - fill_std (float): Standard deviation for random fill (default: 0.0)
           Set to 0 for constant values, >0 for random sampling
+
 """
 
 from dataclasses import dataclass
@@ -80,8 +80,7 @@ class DecodeBenchConnectorMetadata(KVConnectorMetadata):
 
 
 class DecodeBenchConnector(KVConnectorBase_V1, SupportsHMA):
-    """
-    A KV Connector for decode instance performance testing.
+    """A KV Connector for decode instance performance testing.
 
     This connector fills the KV cache with dummy values to emulate a
     prefill-decode disaggregated setting, enabling performance testing of the
@@ -219,8 +218,7 @@ def get_num_new_matched_tokens(
         request: "Request",
         num_computed_tokens: int,
     ) -> tuple[int, bool]:
-        """
-        For new requests, return the number of tokens that should be filled
+        """For new requests, return the number of tokens that should be filled
         with dummy KV cache values.
 
         Returns:
@@ -228,6 +226,7 @@ def get_num_new_matched_tokens(
             - num_tokens_to_fill: number of uncomputed tokens minus 1
                 (we fill everything except the last token for decode)
             - is_async: False (synchronous filling)
+
         """
         req_id = request.request_id
 
@@ -251,8 +250,7 @@ def get_num_new_matched_tokens(
     def update_state_after_alloc(
         self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
     ):
-        """
-        Called after blocks are allocated. Store the block IDs so we can
+        """Called after blocks are allocated. Store the block IDs so we can
         fill them with dummy values.
 
         Supports both single- and multi-group KV cache configurations.
@@ -334,8 +332,7 @@ def update_state_after_alloc(
     def build_connector_meta(
         self, scheduler_output: "SchedulerOutput"
     ) -> KVConnectorMetadata:
-        """
-        Build metadata containing information about which blocks to fill
+        """Build metadata containing information about which blocks to fill
         with dummy KV values.
         """
         meta = DecodeBenchConnectorMetadata(reqs_to_fill=self._pending_fills.copy())
@@ -346,9 +343,7 @@ def build_connector_meta(
         return meta
 
     def request_finished(self, request: "Request"):
-        """
-        Called when a request has finished. Clean up any state.
-        """
+        """Called when a request has finished. Clean up any state."""
         self._filled_requests.discard(request.request_id)
 
 
@@ -389,8 +384,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         )
 
     def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata):
-        """
-        Fill the allocated KV cache blocks with dummy values.
+        """Fill the allocated KV cache blocks with dummy values.
 
         This simulates having a populated KV cache from a prefill phase,
         allowing decode performance testing with larger context sizes.
@@ -419,13 +413,13 @@ def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata):
             )
 
     def _fill_blocks(self, group_idx: int, block_ids: list[int], num_tokens: int):
-        """
-        Fill specified blocks with dummy values for a specific KV cache group.
+        """Fill specified blocks with dummy values for a specific KV cache group.
 
         Args:
             group_idx: The KV cache group index to fill
             block_ids: List of block IDs to fill in this group
             num_tokens: Total number of tokens to fill across these blocks
+
         """
         if not block_ids:
             return
@@ -501,6 +495,7 @@ def _fill_block_tensor(
                 tensor's first dim are ignored.
             fill_mean: Mean value for the fill.
             fill_std: Standard deviation for the fill.
+
         """
         # Convert block_ids to tensor on device
         block_ids_tensor = torch.tensor(
@@ -551,6 +546,7 @@ def _fill_state_tensor(
             kv_cache: A state tensor to fill in its entirety.
             fill_mean: Mean value for the fill.
             fill_std: Standard deviation for the fill.
+
         """
         if fill_std > 0:
             kv_cache.normal_(mean=fill_mean, std=fill_std)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
index 750d3b56d9de..388a8cf3337d 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
@@ -118,6 +118,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non
         Note:
             The number of elements in kv_caches and layer_names should be
             the same.
+
         """
 
         def inject_kv_into_layer(
@@ -134,6 +135,9 @@ def inject_kv_into_layer(
                 src_kv_cache (torch.Tensor): the source KV cache.
                 slot_mapping (torch.Tensor): the slot mapping. In shape
                     [num_tokens].
+                attn_metadata (AttentionMetadata): attention metadata for the
+                    current forward pass.
+
             """
             # Deliberate bulk transfers in this reference connector.
             with gpu_sync_allowed():
@@ -201,6 +205,7 @@ def wait_for_layer_load(self, layer_name: str) -> None:
 
         Args:
             layer_name: the name of that layer
+
         """
         return
 
@@ -220,6 +225,7 @@ def save_kv_layer(
                 layer in vLLM.
             attn_metadata (AttentionMetadata): the attention metadata.
             **kwargs: additional arguments for the save operation.
+
         """
 
         def extract_kv_from_layer(
@@ -260,8 +266,7 @@ def get_num_new_matched_tokens(
         request: "Request",
         num_computed_tokens: int,
     ) -> tuple[int | None, bool]:
-        """
-        Get number of new tokens that can be loaded from the
+        """Get number of new tokens that can be loaded from the
         external KV cache beyond the num_computed_tokens.
 
         Args:
@@ -272,6 +277,7 @@ def get_num_new_matched_tokens(
         Returns:
             the number of tokens that can be loaded from the
             external KV cache beyond what is already computed.
+
         """
         # NOTE: in this debug implementation, we assume that the prompt is
         # cached_prompt + newly_generated_single_token
@@ -295,8 +301,7 @@ def get_num_new_matched_tokens(
     def update_state_after_alloc(
         self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
     ):
-        """
-        Update KVConnector state after block allocation.
+        """Update KVConnector state after block allocation.
 
         If blocks were allocated, add to _requests_need_load,
         such that we load the KVs in the next forward pass.
@@ -315,6 +320,7 @@ def build_connector_meta(
 
         Args:
             scheduler_output (SchedulerOutput): the scheduler output object.
+
         """
         meta = ExampleConnectorMetadata()
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
index d291d4754427..47990d22a5a5 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
@@ -57,6 +57,7 @@ def load_hidden_states(path: str) -> dict[str, torch.Tensor]:
 
     Returns:
         Dict with "hidden_states" and "token_ids" tensors.
+
     """
     lock_path = path + ".lock"
     with open(lock_path) as lf:
@@ -96,13 +97,24 @@ class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata):
 
 
 class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
-    """
-    Simple debug implementation of a HiddenStatesConnector.
+    """Simple debug implementation of a HiddenStatesConnector.
 
     Simply extracts the hidden states from the kv cache and stores them to disk.
     Must be used in conjunction with the `extract_hidden_states` spec decoding method.
     """
 
+    def get_finished_count(self) -> int:
+        """Number of ``finished_sending`` notifications expected per request.
+
+        Only TP rank 0 writes hidden states to disk, so each request yields a
+        single notification. The default (``None``) makes ``KVOutputAggregator``
+        wait for every worker in the world size; since ranks 1..N-1 never submit
+        a copy, the request would never be reported as finished and its blocks
+        would be leaked forever (the scheduler then spins on
+        ``has_finished_requests()``).
+        """
+        return 1
+
     @classmethod
     def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int:
         """Index of the KV cache group holding the extracted hidden states.
@@ -436,8 +448,7 @@ def get_num_new_matched_tokens(
         request: "Request",
         num_computed_tokens: int,
     ) -> tuple[int | None, bool]:
-        """
-        Get number of new tokens that can be loaded from the
+        """Get number of new tokens that can be loaded from the
         external KV cache beyond the num_computed_tokens.
 
         Args:
@@ -448,6 +459,7 @@ def get_num_new_matched_tokens(
         Returns:
             the number of tokens that can be loaded from the
             external KV cache beyond what is already computed.
+
         """
         # This connector is store-only, so we don't need to load any tokens
         return 0, False
@@ -471,6 +483,7 @@ def build_connector_meta(
 
         Args:
             scheduler_output (SchedulerOutput): the scheduler output object.
+
         """
         meta = ExampleHiddenStatesConnectorMetadata()
 
@@ -509,15 +522,20 @@ def request_finished(
         request: "Request",
         block_ids: list[int],
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Called exactly once when a request has finished, before its blocks are
+        """Called exactly once when a request has finished, before its blocks are
         freed.
 
         Returns True to delay block freeing until get_finished extracts
         the hidden states from the KV cache.
         """
         req_id = request.request_id
-        filename = self._request_filenames.pop(req_id)
+        filename = self._request_filenames.pop(req_id, None)
+        if filename is None:
+            # Request was aborted before it was ever scheduled, so
+            # build_connector_meta never recorded a filename for it. There are
+            # no hidden states to save and no blocks worth delaying.
+            self._pending_saves.pop(req_id, None)
+            return False, None
         kv_params = request.kv_transfer_params or {}
         if kv_params.get("include_output_tokens", False):
             # Exclude the final token — it was the model's output, never an
@@ -566,15 +584,18 @@ def get_finished(
         done_sending: set[str] = set()
         for req_id in list(self._accumulated_finished_req_ids):
             event = self._req_copy_events.get(req_id)
-            if event is None or event.query():
-                self._req_copy_events.pop(req_id, None)
+            if event is not None and not event.query():
+                continue  # DtoH copy still in flight
+            self._req_copy_events.pop(req_id, None)
+            self._accumulated_finished_req_ids.discard(req_id)
+            # Only report requests we actually tracked for sending.
+            if event is not None:
                 done_sending.add(req_id)
-                self._accumulated_finished_req_ids.discard(req_id)
-                # Clean up any leftover lock fds (e.g. aborted requests
-                # that never went through _submit_async_write).
-                lock_fd = self._lock_fds.pop(req_id, None)
-                if lock_fd is not None:
-                    os.close(lock_fd)
+            # Clean up any leftover lock fds (e.g. aborted requests that never
+            # went through _submit_async_write).
+            lock_fd = self._lock_fds.pop(req_id, None)
+            if lock_fd is not None:
+                os.close(lock_fd)
 
         return done_sending or None, None
 
@@ -583,20 +604,24 @@ def request_finished_all_groups(
         request: "Request",
         block_ids: tuple[list[int], ...],
     ) -> tuple[bool, dict[str, Any] | None]:
+        if not 0 <= self._cache_kv_group_id < len(block_ids):
+            # No per-group block table (e.g. a request that was never
+            # scheduled); treat as a no-op instead of indexing out of bounds.
+            return False, None
         return self.request_finished(request, block_ids[self._cache_kv_group_id])
 
     @classmethod
     def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
-        """
-        Get the required KV cache layout for this connector.
+        """Get the required KV cache layout for this connector.
+
         Args:
             vllm_config (VllmConfig): the vllm config.
 
         Returns:
             str: the required KV cache layout. e.g. HND, or NHD.
             None if the connector does not require a specific layout.
-        """
 
+        """
         if cls is KVConnectorBase_V1:
             raise TypeError(
                 "get_required_kvcache_layout should not be called "
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py
index 556cba963d5b..bbb64cf24cfe 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py
@@ -96,6 +96,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         Args:
             forward_context (ForwardContext): the forward context.
             **kwargs (Any): additional arguments (unused).
+
         """
         self._flexkv_connector.start_load_kv(forward_context, **kwargs)
 
@@ -107,6 +108,7 @@ def wait_for_layer_load(self, layer_name: str) -> None:
 
         Args:
             layer_name: the name of the layer (unused).
+
         """
         self._flexkv_connector.wait_for_layer_load(layer_name)
 
@@ -131,6 +133,7 @@ def save_kv_layer(
             kv_layer (torch.Tensor): the paged KV buffer (unused).
             attn_metadata (AttentionMetadata): the attention metadata (unused).
             **kwargs (Any): additional arguments (unused).
+
         """
         self._flexkv_connector.save_kv_layer(
             layer_name, kv_layer, attn_metadata, **kwargs
@@ -160,6 +163,7 @@ def get_finished(
             that have finished asynchronous transfer. The finished saves/sends
             req ids must belong to a set provided in a call to this method
             (this call or a prior one).
+
         """
         return self._flexkv_connector.get_finished(finished_req_ids)
 
@@ -169,6 +173,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
 
         Args:
             kv_caches: dictionary of layer names to kv cache tensors.
+
         """
         self._flexkv_connector.register_kv_caches(kv_caches)
 
@@ -192,6 +197,7 @@ def get_num_new_matched_tokens(
             Tuple of (num_external_tokens, is_ready) where
             num_external_tokens is the number of additional tokens that
             can be loaded from the external KV cache.
+
         """
         return self._flexkv_connector.get_num_new_matched_tokens(
             request, num_computed_tokens
@@ -215,6 +221,7 @@ def build_connector_meta(
 
         Args:
             scheduler_output (SchedulerOutput): the scheduler output object.
+
         """
         return self._flexkv_connector.build_connector_meta(scheduler_output)
 
@@ -224,6 +231,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput):
         Args:
             connector_output (KVConnectorOutput): the worker-side
                 connectors output.
+
         """
         self._flexkv_connector.update_connector_output(connector_output)
 
@@ -240,6 +248,7 @@ def request_finished(
             should not be freed until the request_id is returned from
             :meth:`get_finished`. kv_transfer_params is an optional dict of
             KVTransferParams to be included in the request outputs.
+
         """
         return self._flexkv_connector.request_finished(request, block_ids)
 
@@ -248,6 +257,7 @@ def take_events(self) -> Iterable["KVCacheEvent"]:
 
         Returns:
             New KV cache events since the last call.
+
         """
         return self._flexkv_connector.take_events()
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py
index a54233453bb4..e0bf1dbfc4a1 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py
@@ -67,6 +67,7 @@ def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
             size: Total size of the storage file in bytes
             bytes_per_page: Size of each page in bytes
             entries: Maximum number of concurrent operations
+
         """
         if not HF3FS_AVAILABLE:
             raise ImportError(
@@ -154,7 +155,7 @@ def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
         )
 
     def _release_resources(self) -> None:
-        """Release all acquired resources safely"""
+        """Release all acquired resources safely."""
         # iov must be released before ioring and shm
         for attr in ("iov_r", "iov_w", "ior_r", "ior_w"):
             obj = getattr(self, attr, None)
@@ -192,6 +193,7 @@ def batch_read(self, offsets: list[int], tensors: list[torch.Tensor]) -> list[in
 
         Returns:
             List of operation results (0 for success, non-zero for error)
+
         """
         self.check(offsets, tensors)
         assert self.ior_r is not None
@@ -228,11 +230,12 @@ def batch_write(
         Args:
             offsets: List of byte offsets to write to
             tensors: List of tensors containing data to write
+            event: CUDA event to synchronize on before reading the tensors
 
         Returns:
             List of operation results (0 for success, non-zero for error)
-        """
 
+        """
         self.check(offsets, tensors)
         assert self.ior_w is not None
         assert self.iov_w is not None
@@ -282,6 +285,7 @@ def get_size(self) -> int:
 
         Returns:
             Size of the file in bytes
+
         """
         return self.size
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py
index 7294d77cef37..1b37a8844019 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-HF3FS KV Connector Implementation for vLLM.
+"""HF3FS KV Connector Implementation for vLLM.
 
 This module implements a KV connector that uses
 the 3FS for storing and retrieving KV cache data.
@@ -99,9 +98,7 @@
 
 
 class AsyncOperationManager:
-    """
-    Manages async save/load operations with background threads.
-    """
+    """Manages async save/load operations with background threads."""
 
     def __init__(self, connector: "HF3FSKVConnector"):
         # Store connector reference and extract commonly used attributes
@@ -667,9 +664,7 @@ def get_finished(
         return self._async_manager.get_finished_operations(finished_req_ids)
 
     def get_kv_connector_stats(self) -> Optional["KVConnectorStats"]:
-        """
-        Get the KV connector stats collected during the last interval.
-        """
+        """Get the KV connector stats collected during the last interval."""
         # Clear stats for next iteration
         if (
             hasattr(self, "_async_manager")
@@ -886,8 +881,7 @@ def _process_cached_requests(
     def build_kv_connector_stats(
         cls, data: dict[str, Any] | None = None
     ) -> Optional["KVConnectorStats"]:
-        """
-        KVConnectorStats resolution method. This method allows dynamically
+        """KVConnectorStats resolution method. This method allows dynamically
         registered connectors to return their own KVConnectorStats object,
         which can implement custom aggregation logic on the data dict.
         """
@@ -1021,7 +1015,7 @@ def _align_to_block_size(self, num_tokens: int) -> int:
 
 @dataclass
 class HF3FSKVConnectorStats(KVConnectorStats):
-    """Container for transfer performance metrics"""
+    """Container for transfer performance metrics."""
 
     def __post_init__(self):
         if not self.data:
@@ -1081,10 +1075,14 @@ def reduce(self) -> dict[str, int | float]:
             "Num save task failed": num_failed_save,
             "Num load task success": num_success_load,
             "Num load task failed": num_failed_load,
-            "Avg save duration (ms)": round(save_duration.mean() * 1e3, 3),
-            "P90 save duration (ms)": round(np.percentile(save_duration, 90) * 1e3, 3),
-            "Avg load duration (ms)": round(load_duration.mean() * 1e3, 3),
-            "P90 load duration (ms)": round(np.percentile(load_duration, 90) * 1e3, 3),
+            "Avg save duration (ms)": round(save_duration.mean().item() * 1e3, 3),
+            "P90 save duration (ms)": round(
+                np.percentile(save_duration, 90).item() * 1e3, 3
+            ),
+            "Avg load duration (ms)": round(load_duration.mean().item() * 1e3, 3),
+            "P90 load duration (ms)": round(
+                np.percentile(load_duration, 90).item() * 1e3, 3
+            ),
         }
 
     def is_empty(self) -> bool:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py
index 72792e5eb260..92cf08e5bed4 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-HF3FS Metadata Server with key-based organization.
-"""
+"""HF3FS Metadata Server with key-based organization."""
 
 import argparse
 import logging
@@ -119,6 +117,7 @@ def allocate_pages_for_keys(
 
         Returns:
             Dictionary mapping key -> allocated page index
+
         """
         with self.global_lock:
             if rank not in self.rank_metadata:
@@ -167,6 +166,7 @@ def confirm_write_for_keys(
             rank: Rank ID that confirmed the writes
             key_confirmations: List of (key, page_index) tuples
             pages_to_release: List of page indices to release back to free pool
+
         """
         with self.global_lock:
             # Confirm successful writes
@@ -197,6 +197,7 @@ def batch_key_exists(self, keys: list[str]) -> list[bool]:
 
         Returns:
             List of boolean values indicating key existence and completion
+
         """
         with self.global_lock:
             results = []
@@ -218,6 +219,7 @@ def get_key_locations(self, rank: int, keys: list[str]) -> list[int | None]:
 
         Returns:
             List of page indices in the same order as input keys (None if key not found)
+
         """
         with self.global_lock:
             if rank not in self.rank_metadata:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py
index 4ba3ecee90ee..55c3c8407da0 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py
@@ -128,6 +128,7 @@ def scatter_kv_caches(
         num_heads: Size of the H axis
         content_size: Size of the C axis
         kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view
+
     """
     num_layers = len(kv_caches_ptrs)
     num_tokens_in_block = len(token_indices)
@@ -184,6 +185,7 @@ def gather_kv_caches(
         num_heads: Size of the H axis
         content_size: Size of the C axis
         kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view
+
     """
     num_layers = kv_caches_ptrs.shape[0]
     num_tokens_in_block = len(token_indices)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py
new file mode 100644
index 000000000000..5acb7e3bbd87
--- /dev/null
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py
@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Connector boundary for HiSparse scheduler and worker state."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    KVConnectorWorkerMetadata,
+    SupportsHMA,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import (
+    HiSparseKVConnectorStats,
+    HiSparsePromMetrics,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
+    HiSparseConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+    KVConnectorPromMetrics,
+    KVConnectorStats,
+    PromMetric,
+    PromMetricT,
+)
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
+from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror
+from vllm.v1.outputs import KVConnectorOutput
+
+if TYPE_CHECKING:
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+    from vllm.v1.hisparse.coordinator import HiSparseCoordinator
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+
+@dataclass
+class HiSparseConnectorMetadata(KVConnectorMetadata):
+    command: SparseKVOffloadCommand | None
+    host_block_copies: tuple[KVCacheBlockCopy, ...]
+    source_block_ids: tuple[int, ...]
+    row_mirrors: dict[str, tuple[SparseKVRowMirror, ...]]
+    all_context_pages_resident: bool
+
+
+@dataclass
+class HiSparseConnectorWorkerMetadata(KVConnectorWorkerMetadata):
+    enqueued_transfer_counts: dict[int, int]
+    completed_transfer_counts: dict[int, int]
+    completed_host_copy_dst_ids: tuple[int, ...] = ()
+
+    def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetadata:
+        assert isinstance(other, HiSparseConnectorWorkerMetadata)
+
+        def add_counts(first: dict[int, int], second: dict[int, int]) -> dict[int, int]:
+            combined = first.copy()
+            for transfer_id, count in second.items():
+                combined[transfer_id] = combined.get(transfer_id, 0) + count
+            return combined
+
+        return HiSparseConnectorWorkerMetadata(
+            enqueued_transfer_counts=add_counts(
+                self.enqueued_transfer_counts, other.enqueued_transfer_counts
+            ),
+            completed_transfer_counts=add_counts(
+                self.completed_transfer_counts, other.completed_transfer_counts
+            ),
+            # Every worker runs the same copies, so a union is the ack.
+            completed_host_copy_dst_ids=tuple(
+                {*self.completed_host_copy_dst_ids, *other.completed_host_copy_dst_ids}
+            ),
+        )
+
+
+class HiSparseConnectorScheduler:
+    def __init__(
+        self,
+        *,
+        async_speculative: bool,
+        draft_kv_lookahead: int = 0,
+    ) -> None:
+        self.coordinator: HiSparseCoordinator | None = None
+        self.async_speculative = async_speculative
+        self.draft_kv_lookahead = draft_kv_lookahead
+        self.requests: dict[str, Request] = {}
+
+    def bind_coordinator(self, coordinator: HiSparseCoordinator) -> None:
+        assert self.coordinator is None
+        self.coordinator = coordinator
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> HiSparseConnectorMetadata:
+        assert self.coordinator is not None
+        # HiSparse rebinds worker state every step, so the load must start
+        # before the forward rather than being deferred to post-forward.
+        scheduler_output.has_sync_kv_loads = True
+        scheduler_output.block_table_updates = (
+            self.coordinator.take_block_table_updates() or None
+        )
+        command = self.coordinator.build_offload_command()
+        host_block_copies = self.coordinator.take_host_block_copies()
+        source_group_id = self.coordinator.host_group_id
+        assert source_group_id is not None
+        source_block_ids = [
+            block_id
+            for request in scheduler_output.scheduled_new_reqs
+            for block_id in request.block_ids[source_group_id]
+        ]
+        for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids:
+            if new_block_ids is not None:
+                source_block_ids.extend(new_block_ids[source_group_id])
+        num_computed_tokens = {
+            request.req_id: request.num_computed_tokens
+            for request in scheduler_output.scheduled_new_reqs
+        }
+        num_computed_tokens.update(
+            zip(
+                scheduler_output.scheduled_cached_reqs.req_ids,
+                scheduler_output.scheduled_cached_reqs.num_computed_tokens,
+            )
+        )
+        scheduled_requests = tuple(
+            (
+                request_id,
+                num_computed_tokens[request_id],
+                scheduled_count,
+            )
+            for request_id, scheduled_count in (
+                scheduler_output.num_scheduled_tokens.items()
+            )
+        )
+        row_mirrors = {}
+        for request_id, scheduled_start, scheduled_count in scheduled_requests:
+            mirror_start = scheduled_start
+            if self.async_speculative:
+                mirror_start = max(
+                    0,
+                    scheduled_start - self.requests[request_id].num_output_placeholders,
+                )
+            row_mirrors[request_id] = self.coordinator.build_row_mirrors(
+                (
+                    (
+                        request_id,
+                        mirror_start,
+                        scheduled_count
+                        + scheduled_start
+                        - mirror_start
+                        + self.draft_kv_lookahead,
+                    ),
+                ),
+            )
+        return HiSparseConnectorMetadata(
+            command,
+            host_block_copies,
+            tuple(source_block_ids),
+            row_mirrors,
+            self.coordinator.all_context_pages_resident(scheduled_requests),
+        )
+
+    def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
+        assert self.coordinator is not None
+        for request_id in connector_output.finished_recving or ():
+            self.coordinator.finish_host_import(
+                request_id, failed=request_id in connector_output.failed_recving
+            )
+        metadata = connector_output.kv_connector_worker_meta
+        if metadata is None:
+            return
+        assert isinstance(metadata, HiSparseConnectorWorkerMetadata)
+        self.coordinator.release_completed_host_copies(
+            metadata.completed_host_copy_dst_ids
+        )
+        self.coordinator.update_spills(
+            metadata.enqueued_transfer_counts,
+            metadata.completed_transfer_counts,
+        )
+
+
+class HiSparseConnector(KVConnectorBase_V1, SupportsHMA):
+    """Join the scheduler coordinator to the worker's transfer engine."""
+
+    @classmethod
+    def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str:
+        return "BLHNC"
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: KVCacheConfig,
+    ) -> None:
+        super().__init__(vllm_config, role, kv_cache_config)
+        if kv_cache_config.hisparse_host_num_blocks is None:
+            raise ValueError("HiSparseConnector requires a HiSparse host pool")
+        self.connector_scheduler: HiSparseConnectorScheduler | None = None
+        self.connector_worker: HiSparseConnectorWorker | None = None
+        if role == KVConnectorRole.SCHEDULER:
+            speculative_config = vllm_config.speculative_config
+            self.connector_scheduler = HiSparseConnectorScheduler(
+                async_speculative=bool(
+                    vllm_config.scheduler_config.async_scheduling
+                    and speculative_config is not None
+                ),
+                draft_kv_lookahead=vllm_config.num_lookahead_tokens,
+            )
+        elif role == KVConnectorRole.WORKER:
+            self.connector_worker = HiSparseConnectorWorker(
+                vllm_config, kv_cache_config
+            )
+        else:
+            raise ValueError(f"Unsupported KV connector role: {role}")
+
+    def bind_kv_cache_manager(self, kv_cache_manager: KVCacheManager) -> None:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.bind_coordinator(
+            get_hisparse_coordinator(kv_cache_manager)
+        )
+
+    @property
+    def requires_kv_delivery(self) -> bool:
+        return False
+
+    def has_pending_push_work(self) -> bool:
+        assert self.connector_scheduler is not None
+        assert self.connector_scheduler.coordinator is not None
+        return self.connector_scheduler.coordinator.has_pending_work()
+
+    def has_pending_block_frees(self) -> bool:
+        assert self.connector_scheduler is not None
+        assert self.connector_scheduler.coordinator is not None
+        return self.connector_scheduler.coordinator.has_pending_reclamation()
+
+    def finish_forward(self) -> None:
+        assert self.connector_worker is not None
+        self.connector_worker.finish_forward()
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
+        assert self.connector_worker is not None
+        self.connector_worker.register_kv_caches(kv_caches)
+
+    def reset_capture_state(self) -> None:
+        assert self.connector_worker is not None
+        self.connector_worker.reset_hot_state()
+
+    def get_kv_connector_stats(self) -> KVConnectorStats | None:
+        if self.connector_worker is None:
+            return None
+        return self.connector_worker.get_kv_connector_stats()
+
+    @classmethod
+    def build_kv_connector_stats(
+        cls, data: dict[str, Any] | None = None
+    ) -> KVConnectorStats | None:
+        return (
+            HiSparseKVConnectorStats(data=data)
+            if data is not None
+            else HiSparseKVConnectorStats()
+        )
+
+    @classmethod
+    def build_prom_metrics(
+        cls,
+        vllm_config: VllmConfig,
+        metric_types: dict[type[PromMetric], type[PromMetricT]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ) -> KVConnectorPromMetrics:
+        return HiSparsePromMetrics(
+            vllm_config, metric_types, labelnames, per_engine_labelvalues
+        )
+
+    def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
+        assert self.connector_worker is not None
+        metadata = self._get_connector_metadata()
+        assert isinstance(metadata, HiSparseConnectorMetadata)
+        request_state_indices = kwargs.get("request_state_indices")
+        assert request_state_indices is None or isinstance(
+            request_state_indices, torch.Tensor
+        )
+        request_ids = kwargs.get("request_ids")
+        assert request_ids is None or isinstance(request_ids, list)
+        attn_metadata = kwargs.get("attn_metadata")
+        assert attn_metadata is None or isinstance(attn_metadata, dict)
+        self.connector_worker.start_step(
+            metadata,
+            request_state_indices,
+            request_ids,
+            num_tokens=int(kwargs.get("num_tokens") or 0),
+        )
+        self.connector_worker.prepare_forward(attn_metadata)
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        return
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs: Any,
+    ) -> None:
+        return
+
+    def wait_for_save(self) -> None:
+        return
+
+    def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
+        assert self.connector_worker is not None
+        enqueued, completed = self.connector_worker.take_transfer_updates()
+        host_copies = self.connector_worker.take_completed_host_copies()
+        if not enqueued and not completed and not host_copies:
+            return None
+        return HiSparseConnectorWorkerMetadata(
+            enqueued_transfer_counts={transfer_id: 1 for transfer_id in enqueued},
+            completed_transfer_counts={transfer_id: 1 for transfer_id in completed},
+            completed_host_copy_dst_ids=tuple(host_copies),
+        )
+
+    def shutdown(self) -> None:
+        if self.connector_worker is not None:
+            self.connector_worker.shutdown()
+
+    def get_num_new_matched_tokens(
+        self, request: Request, num_computed_tokens: int
+    ) -> tuple[int | None, bool]:
+        return 0, False
+
+    def update_state_after_alloc(
+        self,
+        request: Request,
+        blocks: KVCacheBlocks,
+        num_external_tokens: int,
+    ) -> None:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.requests[request.request_id] = request
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> KVConnectorMetadata:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.build_connector_meta(scheduler_output)
+
+    def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.update_connector_output(connector_output)
+
+    def request_finished_all_groups(
+        self,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.requests.pop(request.request_id, None)
+        return False, None
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py
new file mode 100644
index 000000000000..f2abf51fc4e9
--- /dev/null
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py
@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stats and Prometheus metrics for the HiSparse connector."""
+
+from dataclasses import dataclass
+from typing import Any
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+    KVConnectorPromMetrics,
+    KVConnectorStats,
+    PromMetric,
+    PromMetricT,
+)
+from vllm.v1.metrics.utils import create_metric_per_engine
+
+_HISPARSE_COUNTERS: tuple[tuple[str, str], ...] = (
+    ("cache_hits", "Number of HiSparse device hot-buffer hits."),
+    ("cache_misses", "Number of HiSparse device hot-buffer misses."),
+    (
+        "host_to_device_bytes",
+        "Bytes transferred from host KV storage to HiSparse hot buffers.",
+    ),
+)
+
+
+@dataclass
+class HiSparseKVConnectorStats(KVConnectorStats):
+    """Container for HiSparse hot-buffer residency metrics.
+
+    Each list entry is the delta recorded over one device counter snapshot
+    interval, so a list can hold multiple snapshots per logging interval.
+    """
+
+    def __post_init__(self):
+        if not self.data:
+            # Empty container init, no data is passed in.
+            self.reset()
+
+    def reset(self):
+        # Must be serializable
+        self.data: dict[str, list[int]] = {
+            "cache_hits": [],
+            "cache_misses": [],
+            "host_to_device_bytes": [],
+        }
+
+    def record_snapshot(self, hits: int, misses: int, host_to_device_bytes: int):
+        self.data["cache_hits"].append(hits)
+        self.data["cache_misses"].append(misses)
+        self.data["host_to_device_bytes"].append(host_to_device_bytes)
+
+    def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
+        if not other.is_empty():
+            for k, v in other.data.items():
+                accumulator = self.data[k]
+                assert isinstance(accumulator, list)
+                accumulator.extend(v)
+        return self
+
+    def reduce(self) -> dict[str, int | float]:
+        # Compute compact representative stats suitable for CLI logging.
+        return {
+            "HiSparse hot-buffer hits": sum(self.data["cache_hits"]),
+            "HiSparse hot-buffer misses": sum(self.data["cache_misses"]),
+            "HiSparse host-to-device bytes": sum(self.data["host_to_device_bytes"]),
+        }
+
+    def is_empty(self) -> bool:
+        return all(len(values) == 0 for values in self.data.values())
+
+
+class HiSparsePromMetrics(KVConnectorPromMetrics):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        metric_types: dict[type[PromMetric], type[PromMetricT]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ):
+        super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
+
+        self.hisparse_counters: dict[str, Any] = {}
+        for name, documentation in _HISPARSE_COUNTERS:
+            counter = self._counter_cls(
+                name=f"vllm:hisparse_{name}",
+                documentation=documentation,
+                labelnames=labelnames,
+            )
+            self.hisparse_counters[name] = create_metric_per_engine(
+                counter, self.per_engine_labelvalues
+            )
+
+    def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
+        for name, counter in self.hisparse_counters.items():
+            for value in transfer_stats_data.get(name, []):
+                counter[engine_idx].inc(value)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
new file mode 100644
index 000000000000..b223616401e5
--- /dev/null
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py
@@ -0,0 +1,928 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HiSparse worker-side host/hot state and data movement."""
+
+from __future__ import annotations
+
+from collections import deque
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from functools import partial
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+import numpy as np
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import (
+    HiSparseKVConnectorStats,
+)
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_rank,
+    get_tp_group,
+)
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.utils.torch_utils import current_stream
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX
+from vllm.v1.hisparse.runtime import HiSparseCacheHandle, release_pinned_state
+from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
+from vllm.v1.kv_cache_interface import (
+    HiSparseHotSpec,
+    KVCacheConfig,
+)
+from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace
+
+if TYPE_CHECKING:
+    from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
+        HiSparseConnectorMetadata,
+    )
+    from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
+
+
+class _DMADescriptors(NamedTuple):
+    src: torch.Tensor
+    dst: torch.Tensor
+    sizes: torch.Tensor
+    src_np: np.ndarray
+    dst_np: np.ndarray
+    sizes_np: np.ndarray
+
+
+@dataclass
+class _SlotMappingStaging:
+    stream: torch.Stream
+    event: torch.Event
+    slots: torch.Tensor
+    candidates: tuple[SparseKVRowMirror, ...] = ()
+    num_tokens: int = 0
+    source_index: int = 0
+
+
+def _allocate_dma_descriptors(size: int) -> _DMADescriptors:
+    src, dst, sizes = (torch.empty(size, dtype=torch.int64) for _ in range(3))
+    return _DMADescriptors(
+        src,
+        dst,
+        sizes,
+        src.numpy(),
+        dst.numpy(),
+        sizes.numpy(),
+    )
+
+
+_METRICS_INTERVAL = 2000
+
+
+def _get_hisparse_cache(
+    forward_context: dict[str, Any], layer_name: str
+) -> HiSparseCacheHandle:
+    attention_layer = forward_context[layer_name]
+    hisparse_cache = attention_layer.hisparse_cache
+    assert hisparse_cache is not None
+    return hisparse_cache
+
+
+def _flatten_row_mirrors(
+    row_mirrors: Mapping[str, tuple[SparseKVRowMirror, ...]],
+    request_ids: Sequence[str] | None,
+) -> tuple[SparseKVRowMirror, ...]:
+    ordered_ids = row_mirrors if request_ids is None else request_ids
+    return tuple(
+        mirror
+        for request_id in ordered_ids
+        for mirror in row_mirrors.get(request_id, ())
+    )
+
+
+def _select_written_row_mirrors(
+    candidates: tuple[SparseKVRowMirror, ...],
+    source_slots: np.ndarray,
+    source_index: int,
+) -> tuple[SparseKVRowMirror, ...]:
+    """Select and coalesce GPU-written rows from a scheduler-owned envelope."""
+    source_slots = source_slots[source_slots >= 0]
+    if source_slots.size == 0 or not candidates:
+        return ()
+    starts = np.asarray([mirror.source_starts for mirror in candidates], dtype=np.int64)
+    counts = np.fromiter((mirror.num_rows for mirror in candidates), dtype=np.int64)
+    destinations = np.fromiter(
+        [mirror.destination_start for mirror in candidates], dtype=np.int64
+    )
+    if source_index >= starts.shape[1]:
+        raise RuntimeError("HiSparse row-mirror source index is out of range.")
+    source_starts = starts[:, source_index]
+    order = np.argsort(source_starts)
+    matches = np.searchsorted(source_starts[order], source_slots, side="right") - 1
+    candidate_ids = order[np.maximum(matches, 0)]
+    offsets = source_slots - source_starts[candidate_ids]
+    valid = (matches >= 0) & (offsets < counts[candidate_ids])
+    if not np.all(valid):
+        # A page the scheduler could not map has no host destination, so its
+        # rows are simply not mirrorable this step.
+        source_slots = source_slots[valid]
+        candidate_ids = candidate_ids[valid]
+        offsets = offsets[valid]
+        if source_slots.size == 0:
+            return ()
+    sources = starts[candidate_ids] + offsets[:, None]
+    destinations = destinations[candidate_ids] + offsets
+    boundaries = (
+        np.flatnonzero(
+            (candidate_ids[1:] != candidate_ids[:-1])
+            | (np.diff(destinations) != 1)
+            | np.any(np.diff(sources, axis=0) != 1, axis=1)
+        )
+        + 1
+    )
+    run_starts = np.concatenate(([0], boundaries))
+    run_ends = np.concatenate((boundaries, [destinations.size]))
+    return tuple(
+        SparseKVRowMirror(
+            source_starts=tuple(int(value) for value in sources[start]),
+            destination_start=int(destinations[start]),
+            num_rows=int(end - start),
+        )
+        for start, end in zip(run_starts, run_ends, strict=True)
+    )
+
+
+def _is_hisparse_host_writer(
+    shared_host_region: SharedOffloadRegion | None,
+) -> bool:
+    return shared_host_region is None or get_tensor_model_parallel_rank() == 0
+
+
+def _create_hisparse_host_events(
+    shared_host_region: SharedOffloadRegion | None,
+    is_host_writer: bool,
+    device: torch.device,
+) -> tuple[torch.Event, torch.Event]:
+    if shared_host_region is None:
+        return torch.Event(), torch.Event()
+
+    events: tuple[torch.Event, torch.Event] | None = None
+    ipc_handles = None
+    if is_host_writer:
+        events = (
+            torch.cuda.Event(interprocess=True),
+            torch.cuda.Event(interprocess=True),
+        )
+        stream = current_stream()
+        for event in events:
+            event.record(stream)
+        ipc_handles = tuple(event.ipc_handle() for event in events)
+    ipc_handles = get_tp_group().broadcast_object(ipc_handles, src=0)
+    if events is None:
+        events = (
+            torch.cuda.Event.from_ipc_handle(device, ipc_handles[0]),
+            torch.cuda.Event.from_ipc_handle(device, ipc_handles[1]),
+        )
+    return events
+
+
+class HiSparseConnectorWorker:
+    """Own HiSparse host/hot state and execute its worker-side transfers."""
+
+    def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:
+        self.vllm_config = vllm_config
+        self.kv_cache_config = kv_cache_config
+        self._initialized = False
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
+        forward_context = self.vllm_config.compilation_config.static_forward_context
+        cache_handles: list[HiSparseCacheHandle] = []
+        cache_layer_names: list[str] = []
+        for group in self.kv_cache_config.kv_cache_groups:
+            if not isinstance(group.kv_cache_spec, HiSparseHotSpec):
+                continue
+            for cache_name in group.layer_names:
+                assert cache_name.endswith(HISPARSE_HOT_SUFFIX)
+                layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)]
+                cache_handles.append(_get_hisparse_cache(forward_context, layer_name))
+                cache_layer_names.append(layer_name)
+
+        if not cache_handles:
+            raise RuntimeError("HiSparse connector found no hot-cache handles.")
+        hot_backings: dict[int, torch.Tensor] = {}
+        registered_host_pools: dict[int, torch.Tensor] = {}
+        shared_host_regions: dict[int, SharedOffloadRegion] = {}
+        for cache in cache_handles:
+            hot_backing = cache.runtime.hot_backing
+            hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing
+            registered_pool = cache.runtime.registered_host_pool
+            registered_host_pools[registered_pool.data_ptr()] = registered_pool
+            if (region := cache.runtime.shared_host_region) is not None:
+                shared_host_regions[id(region)] = region
+        if len(hot_backings) != 1:
+            raise RuntimeError("HiSparse hot tensors must share one GPU backing.")
+        hot_backing = next(iter(hot_backings.values()))
+        if len(shared_host_regions) > 1:
+            raise RuntimeError("HiSparse caches must share one host region.")
+        shared_host_region = next(iter(shared_host_regions.values()), None)
+        pinned_host_pools = (
+            []
+            if shared_host_region is not None
+            else list(registered_host_pools.values())
+        )
+        is_host_writer = _is_hisparse_host_writer(shared_host_region)
+
+        resident = cache_handles[0].view
+        assert resident is not None
+        host_num_blocks = self.kv_cache_config.hisparse_host_num_blocks
+        assert host_num_blocks is not None
+        try:
+            self.initialize(
+                cache_handles,
+                cache_layer_names,
+                hot_backing,
+                self.vllm_config.scheduler_config.max_num_seqs,
+                host_num_blocks,
+                hot_backing.device,
+                pinned_host_pools,
+                shared_host_region=shared_host_region,
+                is_host_writer=is_host_writer,
+            )
+        except Exception:
+            release_pinned_state(
+                [cache.runtime for cache in cache_handles],
+                pinned_host_pools,
+                shared_host_region,
+            )
+            raise
+
+    def initialize(
+        self,
+        cache_handles: list[HiSparseCacheHandle],
+        cache_layer_names: list[str],
+        hot_backing: torch.Tensor,
+        max_num_reqs: int,
+        host_num_blocks: int,
+        device: torch.device,
+        pinned_host_pools: list[torch.Tensor],
+        *,
+        shared_host_region: SharedOffloadRegion | None = None,
+        is_host_writer: bool = True,
+    ) -> None:
+        if self._initialized:
+            raise RuntimeError("HiSparse connector worker is already initialized.")
+        resident = cache_handles[0].view
+        assert resident is not None
+        self.kernel_block_size = resident.block_size
+        self.host_num_blocks = host_num_blocks
+        self.pinned_host_pools = pinned_host_pools
+        self.shared_host_region = shared_host_region
+        self.is_host_writer = is_host_writer
+        self.dma_stream = (
+            torch.cuda.Stream(device=device) if self.is_host_writer else None
+        )
+        self._slot_mapping_staging = None
+        if self.is_host_writer:
+            max_mirror_rows = (
+                self.vllm_config.scheduler_config.max_num_batched_tokens
+                + max_num_reqs * (self.vllm_config.num_lookahead_tokens + 1)
+            )
+            self._slot_mapping_staging = _SlotMappingStaging(
+                stream=torch.Stream(device=device),
+                event=torch.Event(),
+                slots=torch.empty(
+                    max_mirror_rows,
+                    dtype=torch.int64,
+                    pin_memory=True,
+                ),
+            )
+        self._dma_free_descriptors: list[_DMADescriptors] = []
+        self._pending_dma_descriptors: deque[tuple[torch.Event, _DMADescriptors]] = (
+            deque()
+        )
+        self._dma_submitted = False
+        self._per_layer_mirrored: set[int] = set()
+        self._submitted_mirror_layers: set[int] = set()
+        self._layer_ready_events = tuple(torch.Event() for _ in cache_handles)
+        self._forward_ready_event = torch.Event()
+        self._set_row_mirrors(())
+        self.cache_layer_names = cache_layer_names
+        self._group_leaders = tuple(
+            (layer_name, cache)
+            for layer_name, cache in zip(cache_layer_names, cache_handles, strict=True)
+            if cache.runtime.is_group_leader
+        )
+        self.host_write_events = _create_hisparse_host_events(
+            shared_host_region, is_host_writer, device
+        )
+        self.host_write_event = self.host_write_events[1]
+        self._next_host_write_event = 0
+        self.cache_handles = cache_handles
+        self.leader_runtimes = [cache.runtime for _, cache in self._group_leaders]
+        request_state_indices = {
+            indices.data_ptr(): indices
+            for cache in cache_handles
+            if (indices := cache.runtime.request_state_indices) is not None
+        }
+        if len(request_state_indices) != 1:
+            raise RuntimeError(
+                "HiSparse runtimes must share one request-state mapping."
+            )
+        self.request_state_indices = next(iter(request_state_indices.values()))
+        if self.request_state_indices.numel() != max_num_reqs:
+            raise RuntimeError(
+                "HiSparse request-state mapping does not match max_num_seqs."
+            )
+        self.hot_backing = hot_backing
+        self._pending_invalid_block_ids: list[int] = []
+        # Destination block ids of host copies this worker has run.
+        self._completed_host_copy_dst_ids: list[int] = []
+        self._post_forward_transfers: list[SparseKVPageTransfer] = []
+        self._enqueued_transfer_ids: list[int] = []
+        self._pending_transfer_events: deque[tuple[torch.Event, tuple[int, ...]]] = (
+            deque()
+        )
+        self._metrics_calls = 0
+        self._metrics_event = torch.Event()
+        self._metrics_pending = False
+        self._init_dma()
+        if self.is_host_writer:
+            for layer_index, handle in enumerate(cache_handles):
+                handle.submit_layer_mirror = partial(
+                    self._enqueue_layer_mirror, layer_index
+                )
+        self._initialized = True
+
+    def set_request_state_indices(self, indices: torch.Tensor) -> None:
+        if indices.numel() > self.request_state_indices.numel():
+            raise ValueError(
+                "HiSparse request-state mapping exceeds max_num_seqs: "
+                f"{indices.numel()} > {self.request_state_indices.numel()}."
+            )
+        if torch.cuda.is_current_stream_capturing():
+            return
+        # Attention indexes persistent request state by input-batch row. Refresh
+        # that indirection after every batch compaction or reorder.
+        self.request_state_indices.fill_(-1)
+        self.request_state_indices[: indices.numel()].copy_(indices)
+        if self._pending_invalid_block_ids:
+            self.invalidate_blocks(self._pending_invalid_block_ids, indices)
+            self._pending_invalid_block_ids.clear()
+
+    def _init_dma(self) -> None:
+        host_caches = tuple(cache.runtime.host_cache for cache in self.cache_handles)
+        resident_caches = []
+        for cache in self.cache_handles:
+            assert cache.view is not None and cache.slot_mapping is not None
+            resident_caches.append(cache.view.cache)
+        self.host_caches = host_caches
+        self.resident_caches = tuple(resident_caches)
+
+        for resident_cache, host_cache in zip(resident_caches, host_caches):
+            row_bytes = resident_cache.shape[-1] * resident_cache.element_size()
+            if (
+                resident_cache.ndim != 3
+                or resident_cache.shape[1] != self.kernel_block_size
+                or resident_cache.stride(1) * resident_cache.element_size() != row_bytes
+            ):
+                raise RuntimeError("HiSparse DMA requires contiguous resident rows.")
+            if (
+                host_cache.ndim != 2
+                or not host_cache.is_contiguous()
+                or host_cache.shape[1] * host_cache.element_size() != row_bytes
+            ):
+                raise RuntimeError("HiSparse DMA requires contiguous host rows.")
+
+    def start_step(
+        self,
+        metadata: HiSparseConnectorMetadata,
+        request_state_indices: torch.Tensor | None,
+        request_ids: list[str] | None = None,
+        num_tokens: int = 0,
+    ) -> None:
+        self._stage_row_mirror_mapping(num_tokens)
+        previous_host_write_event = self.host_write_event
+        self.host_write_event = self.host_write_events[self._next_host_write_event]
+        self._next_host_write_event ^= 1
+        current_stream().wait_event(previous_host_write_event)
+        self._release_completed_dma_descriptors()
+        mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids)
+        if self._slot_mapping_staging is not None:
+            self._slot_mapping_staging.candidates = mirrors
+        self._set_row_mirrors(mirrors)
+        self._dma_submitted = False
+        self._clear_forward_mirror_state()
+        for handle in self.cache_handles:
+            handle.all_context_pages_resident = metadata.all_context_pages_resident
+            handle.mirror_from_resident = True
+        self._copy_host_blocks(metadata.host_block_copies, previous_host_write_event)
+        transfers = (
+            metadata.command.page_transfers if metadata.command is not None else []
+        )
+        self._restore_pages([transfer for transfer in transfers if transfer.restore])
+        self._post_forward_transfers = [
+            transfer
+            for transfer in transfers
+            if not transfer.restore and transfer.after_forward
+        ]
+        self._submit_transfers(
+            [
+                transfer
+                for transfer in transfers
+                if not transfer.restore and not transfer.after_forward
+            ]
+        )
+        self._pending_invalid_block_ids.extend(metadata.source_block_ids)
+        if request_state_indices is not None:
+            self.set_request_state_indices(request_state_indices)
+
+    def _clear_forward_mirror_state(self) -> None:
+        self._per_layer_mirrored.clear()
+        self._submitted_mirror_layers.clear()
+        for handle in self.cache_handles:
+            handle.decode_batch = False
+            handle.host_mirror_required = False
+            handle.num_actual_tokens = 0
+            handle.num_decode_tokens = 0
+            handle.req_id_per_token = None
+
+    def prepare_forward(self, attn_metadata: Mapping[str, Any] | None) -> None:
+        if attn_metadata is None:
+            return
+        for layer_name, handle in self._group_leaders:
+            metadata = attn_metadata.get(layer_name)
+            if metadata is not None:
+                handle.prepare_group_for_batch(metadata)
+
+    def _stage_row_mirror_mapping(self, num_tokens: int) -> None:
+        """Snapshot the rows this forward will write, off the compute stream.
+
+        The resident slot mapping is a persistent view bound at registration
+        time, so the rows are already staged by the time the forward launches.
+        """
+        state = self._slot_mapping_staging
+        if state is None or not num_tokens:
+            return
+        handle = self.cache_handles[0]
+        slots = handle.slot_mapping
+        assert slots is not None
+        source_index = handle.runtime.resident_source_index
+        start = state.num_tokens
+        end = start + num_tokens
+        if end > state.slots.shape[0]:
+            raise ValueError(
+                "HiSparse row mapping exceeds staging capacity: "
+                f"{end} > {state.slots.shape[0]}."
+            )
+        main_stream = current_stream()
+        state.stream.wait_stream(main_stream)
+        with torch.cuda.stream(state.stream):
+            state.slots[start:end].copy_(
+                slots[:num_tokens],
+                non_blocking=True,
+            )
+            state.event.record(state.stream)
+        state.num_tokens = end
+        state.source_index = source_index
+
+    def _resolve_row_mirrors(self, state: _SlotMappingStaging) -> None:
+        """Compute the final row mirrors from the staged GPU slot mapping."""
+        self._set_row_mirrors(
+            _select_written_row_mirrors(
+                state.candidates,
+                state.slots[: state.num_tokens].numpy(),
+                state.source_index,
+            )
+        )
+        state.num_tokens = 0
+
+    def _copy_host_blocks(
+        self,
+        host_block_copies: Sequence[KVCacheBlockCopy],
+        previous_host_write_event: torch.Event,
+    ) -> None:
+        if not host_block_copies:
+            return
+        self._completed_host_copy_dst_ids.extend(
+            copy.dst_block_id for copy in host_block_copies
+        )
+        if self.shared_host_region is None or get_tensor_model_parallel_rank() == 0:
+            if self.host_caches:
+                previous_host_write_event.synchronize()
+            copy_kv_cache_blocks_inplace(
+                self.host_caches,
+                self.host_num_blocks,
+                host_block_copies,
+            )
+        if self.shared_host_region is not None:
+            get_tp_group().barrier()
+
+    def invalidate_blocks(
+        self, block_ids: list[int], request_state_indices: torch.Tensor
+    ) -> None:
+        """Invalidate recycled host slots in this worker's leader runtimes."""
+        if not block_ids:
+            return
+        device = self.cache_handles[0].runtime.device
+        staging = torch.tensor(block_ids, dtype=torch.int32, pin_memory=True)
+        blocks = staging.to(device, dtype=torch.int32, non_blocking=True)
+        offsets = torch.arange(self.kernel_block_size, dtype=torch.int32, device=device)
+        slots = (blocks[:, None] * self.kernel_block_size + offsets[None, :]).flatten()
+        sorted_slots = torch.sort(slots).values
+        state_indices = request_state_indices.to(device=device, dtype=torch.long)
+        for runtime in self.leader_runtimes:
+            runtime.invalidate_sorted_slots(sorted_slots, state_indices)
+
+    def reset_hot_state(self) -> None:
+        for runtime in self.leader_runtimes:
+            runtime.reset_hot_state()
+
+    def get_kv_connector_stats(self) -> HiSparseKVConnectorStats | None:
+        stats = None
+        if self._metrics_pending and self._metrics_event.query():
+            stats = HiSparseKVConnectorStats()
+            for runtime in self.leader_runtimes:
+                group = runtime.index_group
+                hits, misses = group.swap_stats_host.tolist()
+                if hits or misses:
+                    stats.record_snapshot(hits, misses, misses * group.stats_row_bytes)
+            self._metrics_pending = False
+            if stats.is_empty():
+                stats = None
+
+        self._metrics_calls += 1
+        if (
+            self._metrics_calls % _METRICS_INTERVAL == 0
+            and not self._metrics_pending
+            and not torch.cuda.is_current_stream_capturing()
+        ):
+            compute_stream = current_stream()
+            for runtime in self.leader_runtimes:
+                group = runtime.index_group
+                compute_stream.wait_stream(group.copy_stream)
+                group.swap_stats_host.copy_(group.swap_stats, non_blocking=True)
+                group.swap_stats.zero_()
+                group.copy_stream.wait_stream(compute_stream)
+            self._metrics_event.record()
+            self._metrics_pending = True
+        return stats
+
+    def _release_completed_dma_descriptors(self) -> None:
+        pending = self._pending_dma_descriptors
+        while pending and pending[0][0].query():
+            _, descriptors = pending.popleft()
+            self._dma_free_descriptors.append(descriptors)
+
+    def _acquire_dma_descriptors(self, size: int) -> _DMADescriptors:
+        for index, descriptors in enumerate(self._dma_free_descriptors):
+            if descriptors.src.numel() >= size:
+                return self._dma_free_descriptors.pop(index)
+        return _allocate_dma_descriptors(size)
+
+    def _submit_dma_descriptors(
+        self,
+        descriptors: _DMADescriptors,
+        descriptor_count: int,
+        transfer_ids: tuple[int, ...] = (),
+        ready_event: torch.Event | None = None,
+    ) -> None:
+        stream = self.dma_stream
+        assert stream is not None
+        if ready_event is None:
+            stream.wait_stream(current_stream())
+        else:
+            stream.wait_event(ready_event)
+        completion_event = torch.Event()
+        with torch.cuda.stream(stream):
+            ops.swap_blocks_batch(
+                descriptors.src[:descriptor_count],
+                descriptors.dst[:descriptor_count],
+                descriptors.sizes[:descriptor_count],
+            )
+            self.host_write_event.record(stream)
+            completion_event.record(stream)
+        self._pending_dma_descriptors.append((completion_event, descriptors))
+        self._dma_submitted = True
+        if transfer_ids:
+            self._pending_transfer_events.append((completion_event, transfer_ids))
+            self._enqueued_transfer_ids.extend(transfer_ids)
+
+    def _set_row_mirrors(self, mirrors: tuple[SparseKVRowMirror, ...]) -> None:
+        self._row_mirrors = mirrors
+        self._row_mirror_destination_starts = np.fromiter(
+            (mirror.destination_start for mirror in mirrors),
+            dtype=np.int64,
+            count=len(mirrors),
+        )
+        self._row_mirror_counts = np.fromiter(
+            (mirror.num_rows for mirror in mirrors),
+            dtype=np.int64,
+            count=len(mirrors),
+        )
+        if mirrors:
+            self._row_mirror_source_starts = np.asarray(
+                [mirror.source_starts for mirror in mirrors], dtype=np.int64
+            )
+            if self._row_mirror_source_starts.ndim != 2:
+                raise RuntimeError("HiSparse DMA source mappings must be rectangular.")
+        else:
+            self._row_mirror_source_starts = np.empty((0, 0), dtype=np.int64)
+        self._row_mirror_num_rows = int(self._row_mirror_counts.sum())
+
+    def _enqueue_row_dma(
+        self, layer_indices: Sequence[int], ready_event: torch.Event | None = None
+    ) -> None:
+        if (
+            not layer_indices
+            or not self._row_mirror_num_rows
+            or not self.is_host_writer
+        ):
+            return
+        mirrors = self._row_mirrors
+        num_layers = len(layer_indices)
+        descriptor_count = len(mirrors) * num_layers
+        descriptors = self._acquire_dma_descriptors(descriptor_count)
+        destination_starts = self._row_mirror_destination_starts
+        row_counts = self._row_mirror_counts
+        for descriptor_offset, layer_index in enumerate(layer_indices):
+            cache = self.cache_handles[layer_index]
+            source_index = cache.runtime.resident_source_index
+            if source_index >= self._row_mirror_source_starts.shape[1]:
+                raise RuntimeError("HiSparse row DMA source index is out of range.")
+            source_rows = self._row_mirror_source_starts[:, source_index]
+            source = self.resident_caches[layer_index]
+            destination = self.host_caches[layer_index]
+            row_bytes = source.shape[-1] * source.element_size()
+            if (
+                source.stride(1) * source.element_size() != row_bytes
+                or destination.shape[1] * destination.element_size() != row_bytes
+            ):
+                raise RuntimeError("HiSparse row DMA requires contiguous rows.")
+            source_blocks, source_row_offsets = np.divmod(
+                source_rows, self.kernel_block_size
+            )
+            source_out_of_range = (
+                np.any(source_blocks < 0)
+                or np.any(source_blocks >= source.shape[0])
+                or np.any(source_row_offsets + row_counts > self.kernel_block_size)
+            )
+            if source_out_of_range:
+                raise RuntimeError("HiSparse row DMA source is out of range.")
+            if np.any(destination_starts < 0) or np.any(
+                destination_starts + row_counts > destination.shape[0]
+            ):
+                raise RuntimeError("HiSparse row DMA index is out of range.")
+            descriptor_slice = slice(descriptor_offset, descriptor_count, num_layers)
+            descriptors.src_np[descriptor_slice] = (
+                source.data_ptr()
+                + source_blocks * source.stride(0) * source.element_size()
+                + source_row_offsets * row_bytes
+            )
+            descriptors.dst_np[descriptor_slice] = (
+                destination.data_ptr() + destination_starts * row_bytes
+            )
+            descriptors.sizes_np[descriptor_slice] = row_counts * row_bytes
+
+        self._submit_dma_descriptors(
+            descriptors, descriptor_count, ready_event=ready_event
+        )
+
+    def _enqueue_layer_mirror(self, layer_index: int) -> None:
+        handle = self.cache_handles[layer_index]
+        if not handle.host_mirror_required:
+            return
+        if layer_index in self._per_layer_mirrored:
+            raise RuntimeError(f"HiSparse layer {layer_index} mirrored twice.")
+        state = self._slot_mapping_staging
+        if state is not None and state.num_tokens:
+            # The staging stream waited on the compute stream before the
+            # forward launched, so this only drains pre-forward work that is
+            # already queued ahead of the layer kernels: no GPU bubble.
+            state.event.synchronize()
+            self._resolve_row_mirrors(state)
+        self._per_layer_mirrored.add(layer_index)
+        next_layer = layer_index + 1
+        if (
+            next_layer < len(self.cache_handles)
+            and self.cache_handles[next_layer].runtime.resident_source_index
+            == handle.runtime.resident_source_index
+        ):
+            return
+        ready_event = self._layer_ready_events[layer_index]
+        ready_event.record()
+        pending_layers = tuple(
+            sorted(self._per_layer_mirrored - self._submitted_mirror_layers)
+        )
+        self._enqueue_row_dma(
+            pending_layers,
+            ready_event=ready_event,
+        )
+        self._submitted_mirror_layers.update(pending_layers)
+
+    def _record_transfer_completion(
+        self,
+        transfers: list[SparseKVPageTransfer],
+        *,
+        stream: torch.Stream | None = None,
+    ) -> None:
+        if not transfers:
+            return
+        if stream is None:
+            if not self.is_host_writer:
+                return
+            stream = self.dma_stream
+        assert stream is not None
+        completion_event = torch.Event()
+        completion_event.record(stream)
+        transfer_ids = tuple(transfer.transfer_id for transfer in transfers)
+        self._pending_transfer_events.append((completion_event, transfer_ids))
+        self._enqueued_transfer_ids.extend(transfer_ids)
+
+    def _restore_pages(self, transfers: list[SparseKVPageTransfer]) -> None:
+        """Restore imported tails on every rank before its forward or graph replay."""
+        if not transfers:
+            return
+        for layer_index, cache in enumerate(self.cache_handles):
+            source_index = cache.runtime.resident_source_index
+            source = self.host_caches[layer_index]
+            destination = self.resident_caches[layer_index]
+            for transfer in transfers:
+                host_start = transfer.host_block_id * self.kernel_block_size
+                resident_block = transfer.resident_block_ids[source_index]
+                # Each import pays one H2D page copy and enqueue per layer/rank
+                # instead of tail prefill on D. Restore the prompt rows before
+                # decode appends to this page; unused rows do not extend the
+                # logical sequence.
+                # cudaHostRegister pins this pool outside PyTorch's allocator,
+                # so the sync checker cannot recognize it via is_pinned().
+                with gpu_sync_allowed():
+                    destination[resident_block].copy_(
+                        source[host_start : host_start + self.kernel_block_size],
+                        non_blocking=True,
+                    )
+        self._record_transfer_completion(transfers, stream=current_stream())
+
+    def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None:
+        if self.cache_handles[0].runtime.eager_host_mirror:
+            self._record_transfer_completion(transfers)
+        else:
+            self._enqueue_transfers(transfers)
+
+    def _enqueue_transfers(self, transfers: list[SparseKVPageTransfer]) -> None:
+        if not transfers or not self.is_host_writer:
+            return
+        num_layers = len(self.cache_handles)
+        descriptor_count = len(transfers) * num_layers
+        descriptors = self._acquire_dma_descriptors(descriptor_count)
+        destination_rows = np.fromiter(
+            (transfer.host_block_id * self.kernel_block_size for transfer in transfers),
+            dtype=np.int64,
+            count=len(transfers),
+        )
+        source_blocks_by_transfer = np.asarray(
+            [transfer.resident_block_ids for transfer in transfers], dtype=np.int64
+        )
+        if source_blocks_by_transfer.ndim != 2:
+            raise RuntimeError(
+                "HiSparse spill DMA source mappings must be rectangular."
+            )
+        for layer_index, cache in enumerate(self.cache_handles):
+            source_index = cache.runtime.resident_source_index
+            if source_index >= source_blocks_by_transfer.shape[1]:
+                raise RuntimeError("HiSparse spill DMA source index is out of range.")
+            source_blocks = source_blocks_by_transfer[:, source_index]
+            source = self.resident_caches[layer_index]
+            destination = self.host_caches[layer_index]
+            if np.any(source_blocks < 0) or np.any(source_blocks >= source.shape[0]):
+                raise RuntimeError("HiSparse spill DMA source is out of range.")
+            if np.any(destination_rows < 0) or np.any(
+                destination_rows + self.kernel_block_size > destination.shape[0]
+            ):
+                raise RuntimeError("HiSparse spill DMA destination is out of range.")
+            row_bytes = source.shape[-1] * source.element_size()
+            descriptor_slice = slice(layer_index, descriptor_count, num_layers)
+            descriptors.src_np[descriptor_slice] = (
+                source.data_ptr()
+                + source_blocks * source.stride(0) * source.element_size()
+            )
+            descriptors.dst_np[descriptor_slice] = (
+                destination.data_ptr() + destination_rows * row_bytes
+            )
+            descriptors.sizes_np[descriptor_slice] = self.kernel_block_size * row_bytes
+        self._submit_dma_descriptors(
+            descriptors,
+            descriptor_count,
+            transfer_ids=tuple(transfer.transfer_id for transfer in transfers),
+        )
+
+    def _enqueue_host_mirror(
+        self,
+        ready_event: torch.Event | None = None,
+    ) -> None:
+        active = [
+            (index, handle)
+            for index, handle in enumerate(self.cache_handles)
+            if handle.num_actual_tokens != 0
+        ]
+        if not active:
+            return
+
+        def mirror_key(handle: HiSparseCacheHandle) -> tuple:
+            slots = handle.mirror_slot_mapping
+            return (
+                handle.num_actual_tokens,
+                handle.num_decode_tokens,
+                handle.decode_batch,
+                handle.host_mirror_required,
+                handle.runtime.eager_host_mirror,
+                None if slots is None else slots.data_ptr(),
+            )
+
+        keys = {mirror_key(handle) for _, handle in active}
+        if len(keys) > 1:
+            raise RuntimeError(
+                f"HiSparse cache layers disagree on mirror metadata: {list(keys)}."
+            )
+        cache = active[0][1]
+        if not cache.host_mirror_required:
+            return
+        dst_slots = cache.mirror_slot_mapping
+        if dst_slots is None:
+            raise RuntimeError("HiSparse host mirror has no source slot mapping.")
+        num_rows = min(cache.num_actual_tokens, dst_slots.numel())
+        if num_rows == 0:
+            return
+        if self.is_host_writer:
+            expected_layers = {index for index, _ in active}
+            if self._per_layer_mirrored and self._per_layer_mirrored != expected_layers:
+                raise RuntimeError(
+                    "HiSparse per-layer DMA did not mirror every active layer: "
+                    f"expected {sorted(expected_layers)}, got "
+                    f"{sorted(self._per_layer_mirrored)}."
+                )
+            pending_layers = tuple(
+                sorted(expected_layers - self._submitted_mirror_layers)
+            )
+            if pending_layers:
+                self._enqueue_row_dma(pending_layers, ready_event=ready_event)
+                self._submitted_mirror_layers.update(pending_layers)
+        assert cache.req_id_per_token is not None
+        for _, handle in active:
+            if handle.runtime.is_group_leader:
+                handle.runtime.invalidate_written_slots(
+                    dst_slots[:num_rows],
+                    cache.req_id_per_token[:num_rows],
+                )
+
+    def _finish_mirror_phase(self, ready_event: torch.Event | None = None) -> None:
+        state = self._slot_mapping_staging
+        if state is not None and state.num_tokens:
+            # Measured free: blocking here costs no throughput (26.7 vs 27.6
+            # gen tok/s) and lets every step mirror exactly the written rows.
+            state.event.synchronize()
+            self._resolve_row_mirrors(state)
+        self._enqueue_host_mirror(ready_event)
+        self._clear_forward_mirror_state()
+
+    def finish_forward(self) -> None:
+        compute_stream = current_stream()
+        self._forward_ready_event.record()
+        self._finish_mirror_phase(self._forward_ready_event)
+        transfers = self._post_forward_transfers
+        self._post_forward_transfers = []
+        self._submit_transfers(transfers)
+        if self.is_host_writer:
+            if self._dma_submitted:
+                compute_stream.wait_event(self.host_write_event)
+                self._dma_submitted = False
+            else:
+                self.host_write_event.record(compute_stream)
+        self._release_completed_dma_descriptors()
+
+    def take_completed_host_copies(self) -> list[int]:
+        """Drain host copies this worker has enqueued for this step."""
+        completed = self._completed_host_copy_dst_ids
+        self._completed_host_copy_dst_ids = []
+        return completed
+
+    def take_transfer_updates(self) -> tuple[list[int], list[int]]:
+        enqueued = self._enqueued_transfer_ids
+        self._enqueued_transfer_ids = []
+        completed: list[int] = []
+        while self._pending_transfer_events:
+            event, transfer_ids = self._pending_transfer_events[0]
+            if not event.query():
+                break
+            self._pending_transfer_events.popleft()
+            completed.extend(transfer_ids)
+        return enqueued, completed
+
+    def shutdown(self) -> None:
+        if not self._initialized:
+            return
+        if self._slot_mapping_staging is not None:
+            self._slot_mapping_staging.stream.synchronize()
+        if self.dma_stream is not None:
+            self.dma_stream.synchronize()
+        release_pinned_state(
+            [cache.runtime for cache in self.cache_handles],
+            self.pinned_host_pools,
+            self.shared_host_region,
+        )
+        self._initialized = False
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
index 609d0480e880..96191cd8f7a5 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
@@ -32,9 +32,7 @@
 
 
 class LMCacheKVEvents(KVConnectorKVEvents):
-    """
-    Concrete implementation of KVConnectorKVEvents using KVEventAggregator.
-    """
+    """Concrete implementation of KVConnectorKVEvents using KVEventAggregator."""
 
     def __init__(self, num_workers: int) -> None:
         self._aggregator = KVEventAggregator(num_workers)
@@ -43,9 +41,7 @@ def add_events(self, events: list[KVCacheEvent]) -> None:
         self._aggregator.add_events(events)
 
     def aggregate(self) -> "LMCacheKVEvents":
-        """
-        Aggregate KV events and retain only common events.
-        """
+        """Aggregate KV events and retain only common events."""
         common_events = self._aggregator.get_common_events()
         self._aggregator.clear_events()
         self._aggregator.add_events(common_events)
@@ -72,8 +68,7 @@ def __repr__(self) -> str:
 class LMCacheConnectorV1(KVConnectorBase_V1):
     @classmethod
     def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
-        """
-        LMCache requires PIECEWISE CUDA graph mode when layerwise
+        """LMCache requires PIECEWISE CUDA graph mode when layerwise
         operations are enabled. The wait_for_layer_load and save_kv_layer
         methods perform actual async synchronization that cannot be
         captured in CUDA graphs.
@@ -118,12 +113,12 @@ def __init__(
     # Worker-side methods
     # ==============================
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
-        """
-        Initialize with the KV caches. Useful for pre-registering the
+        """Initialize with the KV caches. Useful for pre-registering the
         KV Caches in the KVConnector (e.g. for NIXL).
 
         Args:
             kv_caches: dictionary of layer names, kv cache
+
         """
         if hasattr(self._lmcache_engine, "register_kv_caches"):
             self._lmcache_engine.register_kv_caches(kv_caches)
@@ -149,8 +144,7 @@ def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> No
                 engine.layerwise_retrievers = []
 
     def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
-        """
-        Start loading the KV cache from the connector to vLLM's paged
+        """Start loading the KV cache from the connector to vLLM's paged
         KV buffer. This is called from the forward context before the
         forward pass to enable async loading during model execution.
 
@@ -166,8 +160,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non
         self._lmcache_engine.start_load_kv(forward_context, **kwargs)
 
     def wait_for_layer_load(self, layer_name: str) -> None:
-        """
-        Block until the KV for a specific layer is loaded into vLLM's
+        """Block until the KV for a specific layer is loaded into vLLM's
         paged buffer. This is called from within attention layer to ensure
         async copying from start_load_kv is complete.
 
@@ -175,6 +168,7 @@ def wait_for_layer_load(self, layer_name: str) -> None:
 
         Args:
             layer_name: the name of that layer
+
         """
         self._lmcache_engine.wait_for_layer_load(layer_name)
 
@@ -185,8 +179,7 @@ def save_kv_layer(
         attn_metadata: AttentionMetadata,
         **kwargs: Any,
     ) -> None:
-        """
-        Start saving the a layer of KV cache from vLLM's paged buffer
+        """Start saving the a layer of KV cache from vLLM's paged buffer
         to the connector. This is called from within attention layer to
         enable async copying during execution.
 
@@ -196,14 +189,14 @@ def save_kv_layer(
                 layer in vLLM.
             attn_metadata (AttentionMetadata): the attention metadata.
             **kwargs: additional arguments for the save operation.
+
         """
         self._lmcache_engine.save_kv_layer(
             layer_name, kv_layer, attn_metadata, **kwargs
         )
 
     def wait_for_save(self):
-        """
-        Block until all the save operations is done. This is called
+        """Block until all the save operations is done. This is called
         as the forward context exits to ensure that the async saving
         from save_kv_layer is complete before finishing the forward.
 
@@ -214,8 +207,7 @@ def wait_for_save(self):
     def get_finished(
         self, finished_req_ids: set[str]
     ) -> tuple[set[str] | None, set[str] | None]:
-        """
-        Notifies worker-side connector ids of requests that have
+        """Notifies worker-side connector ids of requests that have
         finished generating tokens.
 
         Returns:
@@ -224,16 +216,17 @@ def get_finished(
             tuple of (sending/saving ids, recving/loading ids).
             The finished saves/sends req ids must belong to a set provided in a
             call to this method (this call or a prior one).
+
         """
         return self._lmcache_engine.get_finished(finished_req_ids)
 
     def get_block_ids_with_load_errors(self) -> set[int]:
-        """
-        Get the set of block IDs that failed to load.
+        """Get the set of block IDs that failed to load.
 
         Returns:
             Set of block IDs that encountered load errors.
             Empty set if no load errors occurred.
+
         """
         method = getattr(self._lmcache_engine, "get_block_ids_with_load_errors", None)
         if callable(method):
@@ -243,10 +236,7 @@ def get_block_ids_with_load_errors(self) -> set[int]:
         return set()
 
     def get_kv_connector_kv_cache_events(self) -> LMCacheKVEvents | None:
-        """
-        Get the KV connector kv cache events collected during the last interval.
-        """
-
+        """Get the KV connector kv cache events collected during the last interval."""
         events = self._lmcache_engine.get_kv_events()  # type: ignore [attr-defined]
         if not events:
             return None
@@ -276,8 +266,7 @@ def get_num_new_matched_tokens(
         request: "Request",
         num_computed_tokens: int,
     ) -> tuple[int | None, bool]:
-        """
-        Get number of new tokens that can be loaded from the
+        """Get number of new tokens that can be loaded from the
         external KV cache beyond the num_computed_tokens.
 
         Args:
@@ -288,6 +277,7 @@ def get_num_new_matched_tokens(
         Returns:
             the number of tokens that can be loaded from the
             external KV cache beyond what is already computed.
+
         """
         return self._lmcache_engine.get_num_new_matched_tokens(
             request, num_computed_tokens
@@ -296,32 +286,30 @@ def get_num_new_matched_tokens(
     def update_state_after_alloc(
         self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
     ):
-        """
-        Update KVConnector state after block allocation.
-        """
+        """Update KVConnector state after block allocation."""
         self._lmcache_engine.update_state_after_alloc(request, num_external_tokens)
 
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:
-        """
-        Build the connector metadata for this step.
+        """Build the connector metadata for this step.
 
         This function should NOT modify fields in the scheduler_output.
         Also, calling this function will reset the state of the connector.
 
         Args:
             scheduler_output (SchedulerOutput): the scheduler output object.
+
         """
         return self._lmcache_engine.build_connector_meta(scheduler_output)
 
     def update_connector_output(self, connector_output: KVConnectorOutput):
-        """
-        Update KVConnector state from worker-side connectors output.
+        """Update KVConnector state from worker-side connectors output.
 
         Args:
             connector_output (KVConnectorOutput): the worker-side
                 connectors output.
+
         """
         # Get the KV events
         kv_cache_events = connector_output.kv_cache_events
@@ -342,8 +330,7 @@ def request_finished(
         request: "Request",
         block_ids: list[int],
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Called when a request has finished, before its blocks are freed.
+        """Called when a request has finished, before its blocks are freed.
 
         Returns:
             True if the request is being saved/sent asynchronously and blocks
@@ -351,15 +338,16 @@ def request_finished(
             get_finished().
             Optional KVTransferParams to be included in the request outputs
             returned by the engine.
+
         """
         return self._lmcache_engine.request_finished(request, block_ids)
 
     def take_events(self) -> Iterable["KVCacheEvent"]:
-        """
-        Take the KV cache events from the connector.
+        """Take the KV cache events from the connector.
 
         Yields:
             New KV cache events since the last call.
+
         """
         if self._kv_cache_events is not None:
             self._kv_cache_events.aggregate()
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
index 2e75519df125..688e1ae18e27 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
@@ -44,8 +44,7 @@ def send_lmcache_request(
     request_type: RequestType,
     payloads: list[Any],
 ) -> MessagingFuture[Any]:
-    """
-    Helper function to send the request to the LMCache multiprocess server
+    """Helper function to send the request to the LMCache multiprocess server.
 
     Args:
         mq_client: The LMCache multiprocess mode message queue client
@@ -54,8 +53,8 @@ def send_lmcache_request(
 
     Returns:
         A messaging future for the request
-    """
 
+    """
     future = mq_client.submit_request(
         request_type, payloads, get_response_class(request_type)
     )
@@ -65,14 +64,14 @@ def send_lmcache_request(
 def get_lmcache_chunk_size(
     mq_client: MessageQueueClient,
 ) -> int:
-    """
-    Helper function to get the LMCache chunk size from the server
+    """Helper function to get the LMCache chunk size from the server.
 
     Args:
         mq_client: The LMCache multiprocess mode message queue client
 
     Returns:
         An integer representing the LMCache chunk size
+
     """
     future = send_lmcache_request(mq_client, RequestType.GET_CHUNK_SIZE, [])
     chunk_size = future.result()
@@ -147,16 +146,16 @@ def __init__(
         vllm_block_size: int,
         parallel_strategy: ParallelStrategy,
     ):
-        """
-        Args:
-            server_url: The server URL for the LMCache message queue
-            context: The ZMQ context
-
-            model_name: The model name used for LMCache keys
-            vllm_block_size: The block size used in vLLM
-            parallel_strategy:
-                The parallel strategy, which includes `use_mla`,
-                `world_size`, `worker_id` and so on
+        """Args:
+        server_url: The server URL for the LMCache message queue
+        context: The ZMQ context
+
+        model_name: The model name used for LMCache keys
+        vllm_block_size: The block size used in vLLM
+        parallel_strategy:
+            The parallel strategy, which includes `use_mla`,
+            `world_size`, `worker_id` and so on
+
         """
         self.mq_client = MessageQueueClient(server_url, context)
 
@@ -195,8 +194,7 @@ def maybe_submit_lookup_request(
         block_hashes: list[bytes] | None = None,
         token_ids: list[int] | None = None,
     ) -> None:
-        """
-        Submit a new lookup request to LMCache if there is no ongoing request.
+        """Submit a new lookup request to LMCache if there is no ongoing request.
 
         Supports both token-based and hash-based vLLM:
         - token_ids: token IDs (token-based vLLM) -> single token-mode key
@@ -219,6 +217,7 @@ def maybe_submit_lookup_request(
             for later retrieve operations.
             In the meantime, this function will record the lookup request, and the
             status of the look up request can be checked by `check_lookup_result`.
+
         """
         if request_id in self.lookup_futures:
             # Skip if there is already a lookup request
@@ -260,8 +259,7 @@ def maybe_submit_lookup_request(
 
     @_lmcache_nvtx_annotate
     def check_lookup_result(self, request_id: str) -> int | None:
-        """
-        Check the result of a previously submitted lookup request.
+        """Check the result of a previously submitted lookup request.
 
         Args:
             request_id: The ID of the lookup request submitted in
@@ -271,6 +269,7 @@ def check_lookup_result(self, request_id: str) -> int | None:
             An integer representing the total number of tokens matched
             in LMCache (prefix matching), or
             None if the lookup request is not finished yet.
+
         """
         assert request_id in self.lookup_futures, (
             f"Lookup request for request_id={request_id} has not been submitted"
@@ -285,25 +284,27 @@ def check_lookup_result(self, request_id: str) -> int | None:
         return num_chunks * self.chunk_size
 
     def num_blocks_per_chunk(self) -> int:
-        """
-        Returns:
-            The number of vllm blocks in a LMCache data chunk
+        """Returns:
+        The number of vllm blocks in a LMCache data chunk
+
         """
         return self.blocks_in_chunk
 
     def cleanup_lookup_result(self, request_id: str) -> None:
-        """
-        Clean up lookup future for a finished request to prevent memory leak.
+        """Clean up lookup future for a finished request to prevent memory leak.
+
         Args:
             request_id: The ID of the finished request.
+
         """
         self.lookup_futures.pop(request_id, None)
 
     def end_session(self, request_id: str) -> None:
-        """
-        Notify LMCache server to remove the session for a finished request.
+        """Notify LMCache server to remove the session for a finished request.
+
         Args:
             request_id: The ID of the finished request.
+
         """
         send_lmcache_request(
             self.mq_client,
@@ -319,7 +320,7 @@ def _create_key(
         end: int = 0,
         request_id: str | None = None,
     ) -> IPCCacheEngineKey:
-        """Convert token IDs to an IPC cache engine key"""
+        """Convert token IDs to an IPC cache engine key."""
         return IPCCacheEngineKey(
             model_name=self.model_name,
             world_size=self.world_size,
@@ -334,7 +335,7 @@ def _create_key(
     def _create_hash_key(
         self, chunk_hash: bytes, request_id: str | None = None
     ) -> IPCCacheEngineKey:
-        """Create a hash-mode IPC cache engine key"""
+        """Create a hash-mode IPC cache engine key."""
         return IPCCacheEngineKey(
             model_name=self.model_name,
             world_size=self.world_size,
@@ -411,12 +412,12 @@ def is_first_rank_of_pp_group(self) -> bool:
         )
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
-        """
-        Register the kv caches with LMCache server
+        """Register the kv caches with LMCache server.
 
         Args:
             kv_caches: A dict of kv caches to register. The keys are the
                 layer names and the values are the corresponding tensors.
+
         """
         # Register kv cache and send the request
         self.kv_caches = kv_caches
@@ -432,14 +433,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
     def submit_store_request(
         self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event
     ):
-        """
-        Submit a KV cache store request to LMCache
+        """Submit a KV cache store request to LMCache.
 
         Args:
             request_id: The ID of the request
             op: The LoadStoreOp describing the store operation.
             event: The CUDA event that is recorded after the current
                 model inference step
+
         """
         if op.block_hashes is not None:
             # Hash mode
@@ -466,14 +467,14 @@ def submit_store_request(
     def submit_retrieve_request(
         self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event
     ):
-        """
-        Submit a KV cache retrieve request to LMCache
+        """Submit a KV cache retrieve request to LMCache.
 
         Args:
             request_id: The ID of the request
             op: The LoadStoreOp describing the retrieve operation.
             event: The CUDA event that is recorded after the current
                 model inference step
+
         """
         if op.block_hashes is not None:
             # Hash mode
@@ -503,8 +504,7 @@ def batched_submit_store_requests(
         ops: list[LoadStoreOp],
         event: torch.cuda.Event,
     ):
-        """
-        Submit a batched store request to LMCache
+        """Submit a batched store request to LMCache.
 
         Args:
             request_ids: The IDs of the requests
@@ -512,6 +512,7 @@ def batched_submit_store_requests(
                 the same length as request_ids
             event: The CUDA event that is recorded after the current
                 model inference step
+
         """
         all_keys: list[IPCCacheEngineKey] = []
         block_ids: list[int] = []
@@ -552,8 +553,7 @@ def batched_submit_retrieve_requests(
         ops: list[LoadStoreOp],
         event: torch.cuda.Event,
     ):
-        """
-        Submit a batched retrieve request to LMCache
+        """Submit a batched retrieve request to LMCache.
 
         Args:
             request_ids: The IDs of the requests
@@ -561,6 +561,7 @@ def batched_submit_retrieve_requests(
                 the same length as request_ids
             event: The CUDA event that is recorded after the current
                 model inference step
+
         """
         all_keys: list[IPCCacheEngineKey] = []
         block_ids: list[int] = []
@@ -598,8 +599,7 @@ def batched_submit_retrieve_requests(
     def get_finished(
         self, finished_req_ids_from_engine: set[str]
     ) -> tuple[set[str] | None, set[str] | None]:
-        """
-        Check and get the finished store and retrieve requests.
+        """Check and get the finished store and retrieve requests.
 
         Args:
             finished_req_ids_from_engine: the set of request ids that are
@@ -617,6 +617,7 @@ def get_finished(
             multiple times in `finished_req_ids_from_engine`. The adapter should
             take care of deduplicating the request IDs and only return the request
             IDs that have not been returned before.
+
         """
         finished_stores = set()
         finished_retrieves = set()
@@ -675,16 +676,14 @@ def get_finished(
         return ret_stores, finished_retrieves
 
     def num_blocks_per_chunk(self) -> int:
-        """
-        Returns:
-            The number of vllm blocks in a LMCache data chunk
+        """Returns:
+        The number of vllm blocks in a LMCache data chunk
+
         """
         return self.blocks_in_chunk
 
     def shutdown(self):
-        """
-        Shutdown the LMCache MP worker adapter
-        """
+        """Shutdown the LMCache MP worker adapter."""
         logger.info("Unregistering kv caches")
         send_lmcache_request(
             self.mq_client, RequestType.UNREGISTER_KV_CACHE, [self.instance_id]
@@ -712,7 +711,7 @@ def _create_key(
         end: int = 0,
         request_id: str | None = None,
     ) -> IPCCacheEngineKey:
-        """Convert token IDs to an IPC cache engine key"""
+        """Convert token IDs to an IPC cache engine key."""
         return IPCCacheEngineKey(
             model_name=self.model_name,
             world_size=self.world_size,
@@ -726,7 +725,7 @@ def _create_key(
     def _create_hash_key(
         self, chunk_hash: bytes, request_id: str | None = None
     ) -> IPCCacheEngineKey:
-        """Create a hash-mode IPC cache engine key"""
+        """Create a hash-mode IPC cache engine key."""
         return IPCCacheEngineKey(
             model_name=self.model_name,
             world_size=self.world_size,
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py
index 64e4893dcbcb..afd9eb862375 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py
@@ -59,9 +59,7 @@ def lmcache_get_or_create_config() -> V1Config:
 
 
 def hex_hash_to_int16(s: str) -> int:
-    """
-    Convert a hex hash string to a 16-bit integer.
-    """
+    """Convert a hex hash string to a 16-bit integer."""
     return int(s, 16) & 0xFFFF
 
 
@@ -70,8 +68,7 @@ def apply_mm_hashes_to_token_ids(
     mm_hashes: list[str],
     mm_positions: list["PlaceholderRange"],
 ) -> torch.Tensor:
-    """
-    Overwrite token_ids in-place for multimodal placeholders using
+    """Overwrite token_ids in-place for multimodal placeholders using
     efficient slice assignments.
     """
     n = token_ids.size(0)
@@ -95,8 +92,7 @@ def mla_enabled(model_config: "ModelConfig") -> bool:
 def extract_mm_features(
     request: Union["Request", "NewRequestData"], modify: bool = False
 ) -> tuple[list[str], list["PlaceholderRange"]]:
-    """
-    Normalize multimodal information from a Request into parallel lists.
+    """Normalize multimodal information from a Request into parallel lists.
 
     This helper reads either:
       1) `request.mm_features` (objects each exposing `.identifier` and
@@ -119,6 +115,7 @@ def extract_mm_features(
     Returns:
         tuple[list[str], list[PlaceholderRange]]: (`mm_hashes`, `mm_positions`).
         May be `([], [])` when no multimodal data is present.
+
     """
     if getattr(request, "mm_features", None):
         mm_hashes, mm_positions = zip(
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
index 6122a20f95b7..a50e603db92b 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
@@ -171,6 +171,7 @@ def from_new_request(
             lmcache_cached_tokens (int): the number of tokens that are
                 cached in LMCache.
             skip_save (bool): whether the request cache should be saved
+
         """
         # vLLM 0.9.0 update: request.block_ids changed from list[int] to
         # list[list[int]]
@@ -219,7 +220,6 @@ def update(
         """Update the request tracker when a running request is
         scheduled again
         """
-
         self.token_ids.extend(new_token_ids)
 
         if new_block_ids is None:
@@ -288,6 +288,7 @@ def from_request_tracker(
         Returns:
             the request metadata if we need to perform load/save
             operations, None otherwise.
+
         """
         input_token_ids = tracker.token_ids
         input_token_len = len(input_token_ids)
@@ -445,6 +446,7 @@ def _init_lmcache_engine(
 
     Returns:
         The initialized LMCache engine
+
     """
     if curr_engine := LMCacheEngineBuilder.get(ENGINE_NAME):
         return curr_engine
@@ -563,6 +565,7 @@ def add_request(self, req_meta: ReqMeta) -> None:
 
         Args:
             req_meta (ReqMeta): the request metadata.
+
         """
         self.requests.append(req_meta)
 
@@ -719,6 +722,7 @@ def get_inference_info(self) -> dict:
 
         Returns:
             dict: Dictionary containing inference information
+
         """
         # Get vLLM config information
         vllm_config = self._vllm_config
@@ -766,6 +770,7 @@ def get_inference_version(self) -> str:
 
         Returns:
             str: vLLM version string
+
         """
         return VLLM_VERSION
 
@@ -816,6 +821,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         Note:
             The number of elements in kv_caches and layer_names should be
             the same.
+
         """
         if len(self.kv_caches) == 0:
             self._init_kv_caches_from_forward_context(forward_context)
@@ -920,6 +926,7 @@ def wait_for_layer_load(self, layer_name: str) -> None:
 
         Args:
             layer_name: the name of that layer
+
         """
         if self.layerwise_retrievers:
             logger.debug("Waiting for layer %s to be loaded", self.current_layer)
@@ -951,6 +958,7 @@ def save_kv_layer(
             kv_layer (torch.Tensor): the paged KV buffer of the current
                 layer in vLLM.
             attn_metadata (AttentionMetadata): the attention metadata.
+
         """
         assert self.lmcache_engine is not None
 
@@ -1039,7 +1047,6 @@ def save_kv_layer(
     @_lmcache_nvtx_annotate
     def wait_for_save(self):
         """Blocking until the KV cache is saved to the connector buffer."""
-
         connector_metadata = self._parent._get_connector_metadata()
         assert isinstance(connector_metadata, LMCacheConnectorMetadata)
 
@@ -1150,8 +1157,7 @@ def get_num_new_matched_tokens(
         request: "Request",
         num_computed_tokens: int,
     ) -> int | None:
-        """
-        Check for external KV cache hit.
+        """Check for external KV cache hit.
 
         Args:
             request (Request): the request object.
@@ -1161,6 +1167,7 @@ def get_num_new_matched_tokens(
         Returns:
             the number of tokens that can be loaded from the
             external KV cache beyond what is already computed.
+
         """
         if self.kv_role == "kv_producer" and not hasattr(
             self.lookup_client, "supports_producer_reuse"
@@ -1236,13 +1243,11 @@ def get_num_new_matched_tokens(
 
     @_lmcache_nvtx_annotate
     def update_state_after_alloc(self, request: "Request", num_external_tokens: int):
-        """
-        Update KVConnector state after temporary buffer alloc.
+        """Update KVConnector state after temporary buffer alloc.
 
         For SharedStorageConnector, update _request_needs_load
         if the CacheManager this allocated blocks for us.
         """
-
         # Clear local status in lookup client when a new request is
         # successfully scheduled.
         self.lookup_client.clear_lookup_status(request.request_id)
@@ -1311,8 +1316,8 @@ def build_connector_meta(
 
         Args:
             scheduler_output (SchedulerOutput): the scheduler output object.
-        """
 
+        """
         force_skip_save = self.kv_role == "kv_consumer" or self.force_skip_save
 
         meta = LMCacheConnectorMetadata()
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
index f465fcad92a8..51f4ddf349e4 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
@@ -88,9 +88,7 @@ def extract_world_size_and_kv_rank(
     rank: int,
     vllm_config: VllmConfig,
 ) -> tuple[int, int]:
-    """
-    Convert the rank for the MLA.
-    """
+    """Convert the rank for the MLA."""
     use_mla = mla_enabled(vllm_config.model_config)
     if not use_mla:
         return world_size, rank
@@ -174,8 +172,7 @@ def create_worker_adapter(
 
 
 class LMCacheMPRequestState(enum.Enum):
-    """
-    State machine:
+    """State machine:
     PREFETCHING -- update_state_after_alloc --> WAITING_FOR_LOAD
     WAITING_FOR_LOAD -- process_loading_requests --> READY
     """
@@ -301,13 +298,13 @@ def GetStoreMetadata(
         blocks_in_chunk: int,
         vllm_block_size: int,
     ) -> "LMCacheMPRequestMetadata | None":
-        """
-        Generate the store metadata for the current request tracker.
+        """Generate the store metadata for the current request tracker.
 
         Args:
             tracker: The request tracker to generate the metadata from.
             blocks_in_chunk: the number of blocks in a LMCache data chunk
             vllm_block_size: the block size used in vLLM
+
         """
         # Store the blocks that has block hashes
         # NOTE: the invariant here is that `num_stored_blocks` should
@@ -375,13 +372,13 @@ def GetRetrieveMetadata(
         blocks_in_chunk: int,
         vllm_block_size: int,
     ) -> "LMCacheMPRequestMetadata | None":
-        """
-        Generate the retrieve metadata for the current request tracker.
+        """Generate the retrieve metadata for the current request tracker.
 
         Args:
             tracker: The request tracker to generate the metadata from.
             blocks_in_chunk: the number of blocks in a LMCache data chunk
             vllm_block_size: the block size used in vLLM
+
         """
         if not tracker.is_ready_for_retrieving():
             return None
@@ -462,8 +459,7 @@ def __repr__(self):
 
 
 class LMCacheMPConnectorUpstream(KVConnectorBase_V1):
-    """
-    The connector for LMCache multi-process mode.
+    """The connector for LMCache multi-process mode.
 
     Extra configs (kv_transfer_config.extra_config):
     - lmcache.mp.host: the host of the LMCache server.
@@ -538,27 +534,26 @@ def _get_connector_metadata(self) -> KVConnectorMetadata:
 
         Returns:
             ConnectorMetadata: the connector metadata.
-        """
 
+        """
         # Should only be called while set to valid metadata.
         assert self._connector_metadata is not None
         return self._connector_metadata
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
-        """
-        Initialize with the KV caches. Useful for pre-registering the
+        """Initialize with the KV caches. Useful for pre-registering the
         KV Caches in the KVConnector (e.g. for NIXL).
 
         Args:
             kv_caches: dictionary of layer names, kv cache
+
         """
         logger.info("Registering kv caches!")
         self.worker_adapter.register_kv_caches(kv_caches)
         return
 
     def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
-        """
-        Start loading the KV cache from the connector to vLLM's paged
+        """Start loading the KV cache from the connector to vLLM's paged
         KV buffer. This is called from the forward context before the
         forward pass to enable async loading during model execution.
 
@@ -597,8 +592,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non
         )
 
     def wait_for_layer_load(self, layer_name: str) -> None:
-        """
-        Block until the KV for a specific layer is loaded into vLLM's
+        """Block until the KV for a specific layer is loaded into vLLM's
         paged buffer. This is called from within attention layer to ensure
         async copying from start_load_kv is complete.
 
@@ -606,6 +600,7 @@ def wait_for_layer_load(self, layer_name: str) -> None:
 
         Args:
             layer_name: the name of that layer
+
         """
         return
 
@@ -616,8 +611,7 @@ def save_kv_layer(
         attn_metadata: AttentionMetadata,
         **kwargs: Any,
     ) -> None:
-        """
-        Start saving a layer of KV cache from vLLM's paged buffer
+        """Start saving a layer of KV cache from vLLM's paged buffer
         to the connector. This is called from within attention layer to
         enable async copying during execution.
 
@@ -627,12 +621,12 @@ def save_kv_layer(
                 layer in vLLM.
             attn_metadata (AttentionMetadata): the attention metadata.
             **kwargs: additional arguments for the save operation.
+
         """
         return
 
     def wait_for_save(self):
-        """
-        Block until all the save operations is done. This is called
+        """Block until all the save operations is done. This is called
         as the forward context exits to ensure that the async saving
         from save_kv_layer is complete before finishing the forward.
 
@@ -673,8 +667,7 @@ def wait_for_save(self):
     def get_finished(
         self, finished_req_ids: set[str]
     ) -> tuple[set[str] | None, set[str] | None]:
-        """
-        Notifies worker-side connector ids of requests that have
+        """Notifies worker-side connector ids of requests that have
         finished generating tokens on the worker.
         The scheduler process (via the Executors) will use this output
         to track which workers are done.
@@ -685,14 +678,14 @@ def get_finished(
             tuple of (sending/saving ids, recving/loading ids).
             The finished saves/sends req ids must belong to a set provided in a
             call to this method (this call or a prior one).
+
         """
         val = self.worker_adapter.get_finished(finished_req_ids)
         # logger.error("Finished req ids: %s, %s", val[0], val[1])
         return val
 
     def get_block_ids_with_load_errors(self) -> set[int]:
-        """
-        Get the set of block IDs that failed to load.
+        """Get the set of block IDs that failed to load.
 
         Returns:
             Set of block IDs that encountered load errors.
@@ -702,17 +695,17 @@ def get_block_ids_with_load_errors(self) -> set[int]:
             - Applies to both sync- and async-loading requests.
             - Async loading: failed blocks may be reported in any forward pass
               up to and including the pass where the request ID is returned by
-              `get_finished()`. Even if failures occur, the request must still
-              be reported via `get_finished()`, and the failed block IDs must
-              appear here no later than that same pass.
+              `get_transfer_results()`. Even if failures occur, the request
+              must still be reported as finished receiving, and the failed
+              block IDs must appear here no later than that same pass.
             - Sync loading: failed blocks should be reported in the forward
               pass in which they are detected.
+
         """
         return self.worker_adapter.get_block_ids_with_load_errors()
 
     def shutdown(self):
-        """
-        Shutdown the connector. This is called when the worker process
+        """Shutdown the connector. This is called when the worker process
         is shutting down to ensure that all the async operations are
         completed and the connector is cleaned up properly.
         """
@@ -721,9 +714,7 @@ def shutdown(self):
         return None
 
     def get_kv_connector_stats(self) -> "KVConnectorStats | None":
-        """
-        Get the KV connector stats collected during the last interval.
-        """
+        """Get the KV connector stats collected during the last interval."""
         return None
 
     # ==============================
@@ -735,8 +726,7 @@ def get_num_new_matched_tokens(
         request: "Request",
         num_computed_tokens: int,
     ) -> tuple[int | None, bool]:
-        """
-        Get number of new tokens that can be loaded from the
+        """Get number of new tokens that can be loaded from the
         external KV cache beyond the num_computed_tokens.
 
         Args:
@@ -761,6 +751,7 @@ def get_num_new_matched_tokens(
             call. If the cache cannot be loaded for some tokens (e.g., due to
             connectivity issues or eviction), those tokens must not be taken
             into account.
+
         """
         tracker = self._get_or_create_request_tracker(request)
         # TODO: support loading KV for preempted requests in the future
@@ -803,8 +794,7 @@ def get_num_new_matched_tokens(
     def update_state_after_alloc(
         self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
     ):
-        """
-        Update KVConnector state after block allocation.
+        """Update KVConnector state after block allocation.
 
         If get_num_new_matched_tokens previously returned True for a
         request, this function may be called twice for that same request -
@@ -817,6 +807,7 @@ def update_state_after_alloc(
             blocks (KVCacheBlocks): the blocks allocated for the request.
             num_external_tokens (int): the number of tokens that will be
                 loaded from the external KV cache.
+
         """
         # NOTE: `blocks` comes from kv_cache_manager.get_blocks(request_id),
         # which returns ALL blocks for the request (not just newly allocated).
@@ -883,14 +874,14 @@ def update_state_after_alloc(
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:
-        """
-        Build the connector metadata for this step.
+        """Build the connector metadata for this step.
 
         This function should NOT modify fields in the scheduler_output.
         Also, calling this function will reset the state of the connector.
 
         Args:
             scheduler_output (SchedulerOutput): the scheduler output object.
+
         """
         metadata = LMCacheMPConnectorMetadata()
 
@@ -907,12 +898,12 @@ def build_connector_meta(
         return metadata
 
     def update_connector_output(self, connector_output: KVConnectorOutput):
-        """
-        Update KVConnector state from worker-side connectors output.
+        """Update KVConnector state from worker-side connectors output.
 
         Args:
             connector_output (KVConnectorOutput): the worker-side
                 connectors output.
+
         """
         return
 
@@ -921,8 +912,7 @@ def request_finished(
         request: "Request",
         block_ids: list[int],
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Called exactly once when a request has finished, before its blocks are
+        """Called exactly once when a request has finished, before its blocks are
         freed.
 
         The connector may assumes responsibility for freeing the blocks
@@ -934,8 +924,8 @@ def request_finished(
             get_finished().
             Optional KVTransferParams to be included in the request outputs
             returned by the engine.
-        """
 
+        """
         params: dict[str, Any] | None = getattr(request, "kv_transfer_params", None)
         return_params: dict[str, Any] | None = {} if params is not None else None
 
@@ -962,26 +952,26 @@ def request_finished(
         return True, return_params
 
     def take_events(self) -> Iterable["KVCacheEvent"]:
-        """
-        Take the KV cache events from the connector.
+        """Take the KV cache events from the connector.
 
         Yields:
             New KV cache events since the last call.
+
         """
         return ()
 
     @classmethod
     def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
-        """
-        Get the required KV cache layout for this connector.
+        """Get the required KV cache layout for this connector.
+
         Args:
             vllm_config (VllmConfig): the vllm config.
 
         Returns:
             str: the required KV cache layout. e.g. HND, or NHD.
             None if the connector does not require a specific layout.
-        """
 
+        """
         if cls is KVConnectorBase_V1:
             raise TypeError(
                 "get_required_kvcache_layout should not be called "
@@ -990,13 +980,13 @@ def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
         return None
 
     def get_finished_count(self) -> int | None:
-        """
-        Get the count of requests expected to complete send/receive operations
+        """Get the count of requests expected to complete send/receive operations
         via this connector. This method is used to initialize the
         KVOutputAggregator, overwriting the default world_size.
 
         Returns:
             int: expected sending or receiving completion count.
+
         """
         return None
 
@@ -1004,8 +994,7 @@ def get_finished_count(self) -> int | None:
     def build_kv_connector_stats(
         cls, data: dict[str, Any] | None = None
     ) -> "KVConnectorStats | None":
-        """
-        KVConnectorStats resolution method. This method allows dynamically
+        """KVConnectorStats resolution method. This method allows dynamically
         registered connectors to return their own KVConnectorStats object,
         which can implement custom aggregation logic on the data dict.
         """
@@ -1019,8 +1008,7 @@ def build_prom_metrics(
         labelnames: list[str],
         per_engine_labelvalues: dict[int, list[object]],
     ) -> "KVConnectorPromMetrics | None":
-        """
-        Create a KVConnectorPromMetrics subclass which should register
+        """Create a KVConnectorPromMetrics subclass which should register
         per-connector Prometheus metrics and implement observe() to
         expose connector transfer stats via Prometheus.
         """
@@ -1183,8 +1171,7 @@ def _get_or_create_request_tracker(
         return self.request_trackers[request_id]
 
     def _cleanup_request_tracker(self, request_id: str) -> None:
-        """
-        Clean up request tracker and associated lookup future for a request.
+        """Clean up request tracker and associated lookup future for a request.
         This should be called when a request is finished to prevent memory leak.
         """
         # Clean up request tracker
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py
index e102272bc78c..770cba8ada9d 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py
@@ -17,8 +17,7 @@
 
 @dataclass
 class KVConnectorStats:
-    """
-    Base class for KV Connector Stats, a container for transfer performance
+    """Base class for KV Connector Stats, a container for transfer performance
     metrics or otherwise important telemetry from the connector.
     All sub-classes need to be serializable as stats are sent from worker to
     logger process.
@@ -35,14 +34,11 @@ def reset(self):
         raise NotImplementedError
 
     def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats":
-        """
-        Aggregate stats with another `KVConnectorStats` object.
-        """
+        """Aggregate stats with another `KVConnectorStats` object."""
         raise NotImplementedError
 
     def reduce(self) -> dict[str, int | float]:
-        """
-        Reduce the observations collected during a time interval to one or
+        """Reduce the observations collected during a time interval to one or
         more representative values (eg avg/median/sum of the series).
         This is meant to be called by the logger to produce a summary of the
         stats for the last time interval.
@@ -95,7 +91,7 @@ def observe(self, transfer_stats_data: dict[str, Any]):
             )
 
     def log(self, log_fn=logger.info):
-        """Log transfer metrics periodically, similar to throughput logging"""
+        """Log transfer metrics periodically, similar to throughput logging."""
         if (
             self.transfer_stats_accumulator
             and not self.transfer_stats_accumulator.is_empty()
@@ -111,8 +107,7 @@ def log(self, log_fn=logger.info):
 
 
 class KVConnectorPromMetrics:
-    """
-    A base class for per-connector Prometheus metric registration
+    """A base class for per-connector Prometheus metric registration
     and recording.
     """
 
@@ -131,8 +126,7 @@ def __init__(
         self.per_engine_labelvalues = per_engine_labelvalues
 
     def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
-        """
-        Record the supplied transfer statistics to Prometheus metrics. These
+        """Record the supplied transfer statistics to Prometheus metrics. These
         statistics are engine-specific, and should be recorded to a metric
         with the appropriate 'engine' label. These metric instances can be
         created using the create_metric_per_engine() helper method.
@@ -141,8 +135,7 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
 
 
 class KVConnectorProm:
-    """
-    Support for registering per-connector Prometheus metrics, and
+    """Support for registering per-connector Prometheus metrics, and
     recording transfer statistics to those metrics. Uses
     KVConnectorBase.build_prom_metrics().
     """
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
index dc574a6219b0..beabf157cc16 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
@@ -2,20 +2,22 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
 import logging
+import queue
 import threading
 import time
 from collections import defaultdict
+from collections.abc import Collection
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from enum import IntEnum
 from typing import TYPE_CHECKING, Any
 
-import httpx
 import msgspec
 import numpy as np
 import torch
 import zmq
 import zmq.asyncio
+from huggingface_hub.utils import httpx
 
 from vllm import envs
 from vllm.config import VllmConfig
@@ -28,6 +30,7 @@
     KVConnectorBase_V1,
     KVConnectorMetadata,
     KVConnectorRole,
+    KVConnectorTransferResults,
     SupportsHMA,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
@@ -307,6 +310,8 @@ def _validate_asymmetric_region_lengths(
 def _align_transfer_regions(
     local_regions: list[TransferRegion],
     remote_regions: list[TransferRegion],
+    *,
+    allow_partial_layers: bool = False,
 ) -> tuple[list[TransferRegion], list[TransferRegion], str | None]:
     """Align KV transfer regions by registered layer-name occurrence.
 
@@ -335,6 +340,11 @@ def keyed_regions(
     for key, local_region in local_keyed:
         remote_region = remote_by_key.get(key)
         if remote_region is None:
+            if (
+                allow_partial_layers
+                and (local_region.layer_name, 0) not in remote_by_key
+            ):
+                continue
             return (
                 [],
                 [],
@@ -393,6 +403,7 @@ class MooncakeXferMetadata(
     registered_layer_names: list[str] = msgspec.field(default_factory=list)
     registered_layer_indices: list[int] = msgspec.field(default_factory=list)
     registered_group_indices: list[int] = msgspec.field(default_factory=list)
+    remote_pp_size: int = 1
 
 
 class MooncakeXferResponseStatus(IntEnum):
@@ -425,6 +436,9 @@ class PullReqMeta:
     expire_time: float = float("inf")
     # Designed for one D pairing to multiple P
     pull_tasks_count: int = 0
+    # Set once any worker reports a failure, so a success from another worker
+    # for the same request is not counted afterwards.
+    failed: bool = False
 
 
 @dataclass
@@ -574,6 +588,18 @@ def get_finished(
         assert self.connector_worker is not None
         return self.connector_worker.get_finished()
 
+    def get_transfer_results(
+        self, finished_req_ids: set[str]
+    ) -> KVConnectorTransferResults:
+        """Get finished and failed sends/recvs from the worker."""
+        assert self.connector_worker is not None
+        return self.connector_worker.get_transfer_results()
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """Get block IDs whose remote KV load failed."""
+        assert self.connector_worker is not None
+        return self.connector_worker.get_block_ids_with_load_errors()
+
     def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         assert self.connector_worker is not None
         assert isinstance(self._connector_metadata, MooncakeConnectorMetadata)
@@ -617,7 +643,7 @@ def build_kv_connector_stats(
 
 
 class MooncakeConnectorScheduler:
-    """Implementation of Scheduler side methods"""
+    """Implementation of Scheduler side methods."""
 
     def __init__(
         self,
@@ -627,6 +653,7 @@ def __init__(
     ):
         self.vllm_config = vllm_config
         self.block_size = vllm_config.cache_config.block_size
+        self.kv_cache_config = kv_cache_config
 
         assert vllm_config.kv_transfer_config
         self.is_kv_producer: bool = (
@@ -641,7 +668,7 @@ def __init__(
             not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
             and any(
                 not isinstance(g.kv_cache_spec, FullAttentionSpec)
-                for g in kv_cache_config.kv_cache_groups
+                for g in kv_cache_config.transfer_groups
             )
         )
         # GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec
@@ -662,7 +689,7 @@ def __init__(
             (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
             if isinstance(g.kv_cache_spec, SlidingWindowSpec)
             else (0, self.block_size)
-            for g in kv_cache_config.kv_cache_groups
+            for g in kv_cache_config.transfer_groups
         ]
         # cdiv(n_tokens, block_size) gives blocks/window; add 1 to
         # conservatively account for boundary overlap.
@@ -676,11 +703,17 @@ def get_sw_clipped_blocks(
         block_ids: tuple[list[int], ...] | list[list[int]],
     ) -> list[list[int]]:
         """Clip per-group block IDs to sliding window size."""
-        if len(block_ids) == 0 or not self._is_hma_required:
-            return list(block_ids)
+        if len(block_ids) == 0:
+            return []
+        if len(block_ids) == len(self.kv_cache_config.transfer_group_ids):
+            selected = list(block_ids)
+        else:
+            selected = list(self.kv_cache_config.select_transfer_block_ids(block_ids))
+        if len(selected) == 0 or not self._is_hma_required:
+            return selected
         return [
             blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks
-            for i, blocks in enumerate(block_ids)
+            for i, blocks in enumerate(selected)
         ]
 
     def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
@@ -723,8 +756,7 @@ def on_new_request(self, request: "Request") -> None:
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int, bool]:
-        """
-        For remote prefill, pull all prompt blocks from remote
+        """For remote prefill, pull all prompt blocks from remote
         asynchronously relative to engine execution.
 
         Args:
@@ -736,8 +768,8 @@ def get_num_new_matched_tokens(
               external KV cache beyond what is already computed.
             * true if the external KV cache tokens will be loaded
               asynchronously (between scheduler steps).
-        """
 
+        """
         params = request.kv_transfer_params
         logger.debug(
             "MooncakeConnector get_num_new_matched_tokens: "
@@ -848,11 +880,9 @@ def request_finished(
         request: "Request",
         block_ids: tuple[list[int], ...],
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Once a request is finished, determine whether request blocks
+        """Once a request is finished, determine whether request blocks
         should be freed now or will be sent asynchronously and freed later.
         """
-
         params = request.kv_transfer_params
         logger.debug(
             "MooncakeConnector request_finished, req_id=%s, request_status=%s, "
@@ -901,7 +931,7 @@ def request_finished(
 
 
 class MooncakeConnectorWorker:
-    """Implementation of Worker side methods"""
+    """Implementation of Worker side methods."""
 
     def __init__(
         self,
@@ -1019,6 +1049,12 @@ def __init__(
 
         self.finished_sending_reqs: set[ReqId] = set()
         self.finished_recving_reqs: set[ReqId] = set()
+        # Written from the receiver loop, drained from the worker thread.
+        self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue()
+        self._is_hma_required = len(kv_cache_config.kv_cache_groups) > 1
+        # Block IDs are only unique within a group, so with multiple groups
+        # load failures are reported per request instead.
+        self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
 
         self.xfer_stats = MooncakeKVConnectorStats()
 
@@ -1038,18 +1074,13 @@ def __init__(
 
         self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
         self._layer_specs: dict[str, KVCacheSpec] = {}
-        for group in kv_cache_config.kv_cache_groups:
+        for group in kv_cache_config.transfer_groups:
             group_spec = group.kv_cache_spec
             specs_by_layer = getattr(group_spec, "kv_cache_specs", {})
             for layer_name in group.layer_names:
                 self._layer_specs[layer_name] = specs_by_layer.get(
                     layer_name, group_spec
                 )
-        self._layer_group_indices: dict[str, int] = {
-            layer: group_index
-            for group_index, group in enumerate(kv_cache_config.kv_cache_groups)
-            for layer in group.layer_names
-        }
         self.transfer_topo = TransferTopology(
             tp_rank=self.tp_rank,
             tp_size=self.tp_size,
@@ -1136,11 +1167,9 @@ async def register_worker_with_bootstrap(self):
                 raise e
 
     async def _mooncake_sender_listener(self, ready_event: threading.Event):
-        """
-        Background thread that listens for Mooncake requests, dispatches them
+        """Background thread that listens for Mooncake requests, dispatches them
         to a thread pool, and sends acknowledgments upon completion.
         """
-
         sock = self.async_zmq_ctx.socket(zmq.ROUTER)
         self.side_channel_port = sock.bind_to_random_port(f"tcp://{self.hostname}")
         logger.debug(
@@ -1232,7 +1261,11 @@ async def send_kv_to_decode(
             meta.registered_group_indices,
         )
         local_regions, remote_regions, align_err = _align_transfer_regions(
-            local_regions, remote_regions
+            local_regions,
+            remote_regions,
+            allow_partial_layers=(
+                meta.remote_pp_size > 1 and meta.remote_pp_size != self.pp_size
+            ),
         )
         if align_err is not None:
             response = MooncakeXferResponse(
@@ -1322,7 +1355,9 @@ async def wait_and_ret(
                     # Mark it sending to avoid expiration.
                     send_meta.sending += 1
                     if not send_meta.need_send:
-                        self.resolve_need_send(send_meta, remote_tp_ranks)
+                        self.resolve_need_send(
+                            send_meta, remote_tp_ranks, meta.remote_pp_size
+                        )
                     ready_reqs.append((d_req_id, send_meta))
                 else:
                     # Otherwise (expired, very unlikely), just forget it.
@@ -1398,11 +1433,16 @@ def resolve_need_send(
         self,
         send_meta: SendBlockMeta,
         remote_tp_ranks: list[int],
+        remote_pp_size: int = 1,
     ):
         # Prepare for heterogeneous TP (one P pairs to multiple D)
         send_meta.need_send = len(remote_tp_ranks)
+        if remote_pp_size > 1 and remote_pp_size != self.pp_size:
+            # Each consumer PP stage pulls every producer stage, including
+            # peers with no shared layers.
+            send_meta.need_send *= remote_pp_size
         logger.debug(
-            "Mooncake request %s will be served by %d consumer TP workers: TP ranks=%s",
+            "Mooncake request %s will be served by %d consumer workers: TP ranks=%s",
             send_meta.transfer_id,
             send_meta.need_send,
             remote_tp_ranks,
@@ -1422,7 +1462,7 @@ def _logical_to_kernel_block_ids(
         block_arange = np.arange(self._physical_blocks_per_logical_kv_block).reshape(
             1, -1
         )
-        group_specs = self.kv_cache_config.kv_cache_groups
+        group_specs = self.kv_cache_config.transfer_groups
         return [
             BlockTable.map_to_kernel_blocks(
                 np.array(group),
@@ -1475,7 +1515,7 @@ async def _build_transfer_params(
             local_block_ids_by_group: list[list[int]] = []
             remote_block_ids_by_group: list[list[int]] = []
             has_block_error = False
-            group_specs = self.kv_cache_config.kv_cache_groups
+            group_specs = self.kv_cache_config.transfer_groups
             for group_index, (local_group, remote_group) in enumerate(
                 zip(send_meta.local_block_ids, remote_block_ids_per_group)
             ):
@@ -1665,7 +1705,6 @@ def _send_blocks(
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in mooncake."""
-
         logger.info("Registering KV_Caches. use_mla: %s", self.use_mla)
 
         kv_data_ptrs: list[int] = []
@@ -1731,7 +1770,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 self.registered_layer_names.append(layer_name)
                 self.registered_layer_indices.append(layer_index)
                 self.registered_group_indices.append(
-                    self._layer_group_indices[layer_name]
+                    self.kv_cache_config.transfer_group_index_by_layer[layer_name]
                 )
             storage = cache.untyped_storage()
             storage_addr = storage.data_ptr()
@@ -1802,8 +1841,7 @@ async def fetch_finished_sending_reqs(self) -> set[ReqId]:
         return finished_sending_reqs
 
     def get_finished(self) -> tuple[set[str] | None, set[str] | None]:
-        """
-        Get requests that are done sending or recving on this specific worker.
+        """Get requests that are done sending or recving on this specific worker.
         The scheduler process (via the MultiprocExecutor) will use this output
         to track which workers are done.
         """
@@ -1833,6 +1871,28 @@ def get_finished(self) -> tuple[set[str] | None, set[str] | None]:
 
         return finished_sending_reqs or None, finished_recving_reqs or None
 
+    def get_transfer_results(self) -> KVConnectorTransferResults:
+        """Get transfers that completed on this specific worker, including
+        requests whose remote KV load failed.
+
+        The scheduler process (via the MultiprocExecutor) will use this output
+        to track which workers are done.
+        """
+        finished_sending_reqs, finished_recving_reqs = self.get_finished()
+
+        failed_recving_reqs: set[ReqId] = set()
+        while True:
+            try:
+                failed_recving_reqs.add(self._failed_recv_reqs.get_nowait())
+            except queue.Empty:
+                break
+
+        return KVConnectorTransferResults(
+            finished_sending=set(finished_sending_reqs or ()),
+            finished_recving=set(finished_recving_reqs or ()),
+            failed_recving=failed_recving_reqs,
+        )
+
     def get_kv_connector_stats(self) -> KVConnectorStats | None:
         """Return transfer stats collected since the last call, or None
         if nothing has been recorded in this interval."""
@@ -1851,6 +1911,7 @@ async def receive_kv_from_single_worker(
             remote_port=self.rpc_port,
             remote_tp_size=self.tp_size,
             remote_tp_rank=self.tp_rank,
+            remote_pp_size=self.pp_size,
             req_blocks={
                 req_id: (pull_meta.transfer_id, pull_meta.local_block_ids)
                 for req_id, pull_meta in pull_metas.items()
@@ -1885,12 +1946,11 @@ async def receive_kv_from_single_worker(
                     ret_msg = await sock.recv()
                     response = self._xfer_resp_decoder.decode(ret_msg)
                     if response.status == MooncakeXferResponseStatus.ERROR:
-                        logger.error(
-                            "Error happens during transferring kvcache for %s: %s",
+                        self._handle_failed_recv(
+                            pull_metas,
                             req_ids,
-                            response.err_msg,
+                            response.err_msg or "transfer error",
                         )
-                        self.xfer_stats.record_failed_recv()
                         return
                     self.process_pulling_result(response, pull_metas)
                     if response.status == MooncakeXferResponseStatus.FINISH:
@@ -1898,10 +1958,53 @@ async def receive_kv_from_single_worker(
         except zmq.ContextTerminated:
             logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.")
         except Exception as e:
-            logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e)
-            self.xfer_stats.record_failed_recv()
+            self._handle_failed_recv(pull_metas, req_ids, f"transfer failed: {e}")
             return
 
+    def _handle_failed_recv(
+        self,
+        pull_metas: dict[ReqId, PullReqMeta],
+        req_ids: Collection[ReqId],
+        reason: str,
+    ) -> None:
+        """Report a failed remote KV load so the scheduler can fail or recompute it."""
+        failed: list[ReqId] = []
+        for req_id in req_ids:
+            pull_meta = pull_metas.get(req_id)
+            if pull_meta is None or pull_meta.failed:
+                continue
+            pull_meta.failed = True
+            failed.append(req_id)
+            self.xfer_stats.record_failed_recv()
+
+            invalid = {b for group in pull_meta.local_block_ids for b in group}
+            if not invalid:
+                # A pull with no local blocks only asks P to release its blocks
+                # for a request that never reached the scheduler (see
+                # AsyncLLM.notify_kv_transfer_request_rejected, which submits an
+                # abort_immediately request just to run request_finished). No D
+                # request is waiting on a load, and reporting one here would trip
+                # the scheduler's `assert req_id in self.requests`.
+                continue
+            if self._is_hma_required:
+                self._failed_recv_reqs.put(pull_meta.d_req_id)
+            else:
+                self._invalid_block_ids.put(invalid)
+            self.finished_recving_reqs.add(pull_meta.d_req_id)
+
+        if failed:
+            logger.error("pulling kv_caches for %s failed: %s", failed, reason)
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """Drain the blocks whose remote KV load failed since the last call."""
+        result: set[int] = set()
+        while True:
+            try:
+                result.update(self._invalid_block_ids.get_nowait())
+            except queue.Empty:
+                break
+        return result
+
     def process_pulling_result(
         self,
         response: MooncakeXferResponse,
@@ -1911,6 +2014,8 @@ def process_pulling_result(
 
         for req_id in ok_reqs:
             pull_meta = pull_metas[req_id]
+            if pull_meta.failed:
+                continue
             # No race because we are in async loop.
             pull_meta.pull_tasks_count -= 1
             if pull_meta.pull_tasks_count == 0:
@@ -1920,10 +2025,8 @@ def process_pulling_result(
             logger.debug("pulling kv_caches for %s finished", ok_reqs)
 
         if response.err_reqs:
-            logger.error(
-                "pulling kv_caches for %s failed: %s",
-                response.err_reqs,
-                response.err_msg,
+            self._handle_failed_recv(
+                pull_metas, response.err_reqs, response.err_msg or "unknown error"
             )
 
     async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str):
@@ -2001,10 +2104,11 @@ async def handle_new_engine_id(
             await self._pending_bootstrap_queries[remote_bootstrap_addr].wait()
 
         if remote_engine_id not in self._remote_agents:
-            logger.error(
-                "Failed to find remote engine_id %s from bootstrap server %s",
-                remote_engine_id,
-                remote_bootstrap_addr,
+            self._handle_failed_recv(
+                pull_metas,
+                list(pull_metas),
+                f"remote engine_id {remote_engine_id} not found from bootstrap "
+                f"server {remote_bootstrap_addr}",
             )
             return
 
@@ -2101,7 +2205,7 @@ def _get_transfer_regions(
     ) -> list[TransferRegion]:
         if not group_indices:
             group_indices = [
-                self._layer_group_indices.get(layer_name, 0)
+                self.kv_cache_config.transfer_group_index_by_layer.get(layer_name, 0)
                 for layer_name in layer_names
             ]
         return _expand_transfer_regions(
@@ -2203,8 +2307,7 @@ def should_launch_bootstrap_server(vllm_config: VllmConfig) -> bool:
 
 
 def get_mooncake_bootstrap_addr(vllm_config: VllmConfig) -> tuple[str, int]:
-    """
-    Returns the address of the Mooncake bootstrap server.
+    """Returns the address of the Mooncake bootstrap server.
     This is only used by prefillers to register workers.
     Decoders should get addr from kv_transfer_params.
     """
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py
index 2d158387f060..df47f9fdb737 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py
@@ -32,8 +32,7 @@ class EngineEntry:
 
 
 class MooncakeBootstrapServer:
-    """
-    A centralized server running on the global rank 0 prefiller worker.
+    """A centralized server running on the global rank 0 prefiller worker.
     Prefiller workers register their connection info (IP, port, ranks) here.
     """
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py
index d177f55cc720..a9dbd6bdd4ec 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py
@@ -126,16 +126,16 @@ def reduce(self) -> dict[str, int | float]:
         avg_mb = total_mb / n
         total_time_seconds = xfer_time.sum()
         throughput_mb_s = (
-            total_mb / total_time_seconds if total_time_seconds > 0 else 0.0
+            (total_mb / total_time_seconds).item() if total_time_seconds > 0 else 0.0
         )
 
         return {
             "Num successful transfers": n,
-            "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
+            "Avg xfer time (ms)": round(xfer_time.mean().item() * 1e3, 3),
             "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
-            "Avg MB per transfer": round(avg_mb, 3),
+            "Avg MB per transfer": round(avg_mb.item(), 3),
             "Throughput (MB/s)": round(throughput_mb_s, 3),
-            "Avg number of descriptors": round(descs.mean(), 1),
+            "Avg number of descriptors": round(descs.mean().item(), 1),
             "Num failed transfers": num_failed_transfers,
             "Num failed recvs": num_failed_recvs,
             "Num KV expired reqs": num_kv_expired_reqs,
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
index 2c8d118cdde8..eb4a7fafe92a 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
@@ -27,6 +27,7 @@
     KVConnectorBase_V1,
     KVConnectorMetadata,
     KVConnectorRole,
+    KVConnectorTransferResults,
     KVConnectorWorkerMetadata,
     SupportsHMA,
 )
@@ -132,20 +133,22 @@ def _validate_kv_cache_config(
         from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec
 
         unsupported: list[str] = []
-        cache_block_size = vllm_config.cache_config.block_size
-        for g_idx, g in enumerate(kv_cache_config.transfer_groups):
-            spec = g.kv_cache_spec
+        store_group_ids = kv_cache_config.prefix_cacheable_group_ids
+        if not store_group_ids:
+            raise ValueError(
+                "MooncakeStore requires at least one prefix-cacheable KV cache group"
+            )
+        for group_id in store_group_ids:
+            spec = kv_cache_config.kv_cache_groups[group_id].kv_cache_spec
             if isinstance(spec, CrossAttentionSpec):
-                unsupported.append(f"group {g_idx}: CrossAttentionSpec")
-            # Enforce Mamba align mode
-            if isinstance(spec, MambaSpec) and spec.block_size != cache_block_size:
+                unsupported.append(f"group {group_id}: CrossAttentionSpec")
+            if isinstance(spec, MambaSpec) and spec.mamba_cache_mode != "align":
                 unsupported.append(
-                    f"group {g_idx}: MambaSpec with block_size="
-                    f"{spec.block_size} != cache_config.block_size="
-                    f"{cache_block_size} (mamba_cache_mode != 'align')"
+                    f"group {group_id}: mamba_cache_mode="
+                    f"{spec.mamba_cache_mode!r} != 'align'"
                 )
         pcp = vllm_config.parallel_config.prefill_context_parallel_size
-        if len(kv_cache_config.transfer_groups) > 1 and pcp > 1:
+        if len(store_group_ids) > 1 and pcp > 1:
             unsupported.append(f"PCP > 1 (pcp={pcp}) with hybrid attention")
         if unsupported:
             raise ValueError(
@@ -357,6 +360,14 @@ def get_finished(
         assert isinstance(metadata, MooncakeStoreConnectorMetadata)
         return self.connector_worker.get_finished(finished_req_ids, metadata)
 
+    def get_transfer_results(
+        self, finished_req_ids: set[str]
+    ) -> KVConnectorTransferResults:
+        assert self.connector_worker is not None
+        metadata = self._get_connector_metadata()
+        assert isinstance(metadata, MooncakeStoreConnectorMetadata)
+        return self.connector_worker.get_transfer_results(finished_req_ids, metadata)
+
     def get_block_ids_with_load_errors(self) -> set[int]:
         assert self.connector_worker is not None
         return self.connector_worker.get_block_ids_with_load_errors()
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
index 3f7008201329..d05b553b64b9 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
@@ -3,18 +3,20 @@
 """External-store cache-hit coordinator for MooncakeStoreConnector."""
 
 from collections.abc import Sequence
-from typing import cast
+from typing import NamedTuple, cast
 
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
     chunk_hashes_for_block_size,
 )
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.block_pool import BlockPool
-from vllm.v1.core.kv_cache_coordinator import SpecGroup
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
     KVCacheBlock,
 )
+from vllm.v1.core.single_type_kv_cache_manager import (
+    SingleTypeKVCacheManager,
+)
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheGroupSpec,
@@ -25,6 +27,13 @@
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 
 
+class StoreSpecGroup(NamedTuple):
+    spec: KVCacheSpec
+    group_ids: list[int]
+    manager_cls: type[SingleTypeKVCacheManager]
+    use_eagle: bool
+
+
 class ExternalCachedBlockPool:
     """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set."""
 
@@ -91,6 +100,7 @@ def __init__(
         assert all(
             scheduler_block_size % g.kv_cache_spec.block_size == 0
             for g in kv_cache_groups
+            if g.kv_cache_spec.prefix_cacheable
         ), "scheduler_block_size must be a multiple of each group's block_size"
         self.kv_cache_groups = kv_cache_groups
         self.mamba_group_ids = {
@@ -120,7 +130,7 @@ def _verify_and_split_kv_cache_groups(self) -> None:
         """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but
         dispatches via spec_manager_map (we don't allocate managers).
         """
-        attention_groups: list[SpecGroup] = []
+        attention_groups: list[StoreSpecGroup] = []
         for i, g in enumerate(self.kv_cache_groups):
             # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash
             # kpool tail): per-request scratch, never shareable, so they must
@@ -142,7 +152,7 @@ def _verify_and_split_kv_cache_groups(self) -> None:
                     break
             else:
                 attention_groups.append(
-                    SpecGroup(spec, [i], manager_cls, g.is_eagle_group)
+                    StoreSpecGroup(spec, [i], manager_cls, g.is_eagle_group)
                 )
         # Full attention first (matches upstream convergence ordering).
         attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec))
@@ -442,7 +452,7 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec:
 
 
 def partial_hash_hits_enabled(
-    kv_cache_groups: list[KVCacheGroupSpec],
+    kv_cache_groups: Sequence[KVCacheGroupSpec],
     hash_block_size: int,
     dcp_world_size: int = 1,
 ) -> bool:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
index b9df9648d58e..2510a7ee8e76 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
@@ -630,6 +630,7 @@ def process_tokens(
                 ``mask_num``. False entries are skipped before hash access.
             put_step: Stride for distributing chunks across ranks.
             put_step_rank: ``chunk_id % put_step`` value this rank stores.
+
         """
         assert put_step > 0
         if not block_hashes:
@@ -662,6 +663,7 @@ class TailKeyBoundary:
             ``block_hashes[num_tokens // hash_block_size - 1]`` instead of the
             hash implied by ``MooncakeLookupResult.hit_length``. This changes
             only the load key, not the reusable prefix.
+
     """
 
     group_id: int
@@ -678,6 +680,7 @@ class MooncakeLookupResult:
         tail_key_boundaries: Hash boundary used to store each cache group's
             tail block when ``hit_length`` does not identify its store key.
             There is one entry per group for every nonzero hit.
+
     """
 
     hit_length: int
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
index ae61710e25f7..8f277d7eea37 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
@@ -71,19 +71,27 @@ def __init__(
         )
         self.client = LookupKeyClient(vllm_config)
         self.kv_cache_config = kv_cache_config
+        self._store_group_ids = kv_cache_config.prefix_cacheable_group_ids
+        # Map scheduler group IDs to store group indices. Groups outside the
+        # store projection, such as the QSA ring, map to None and are skipped.
+        self._store_group_id_by_kv_cache_group_id = {
+            group_id: store_group_id
+            for store_group_id, group_id in enumerate(self._store_group_ids)
+        }
+        store_groups = kv_cache_config.prefix_cacheable_groups
 
         # Align with the engine's own scheduler_block_size and hash_block_size.
         self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes(
             kv_cache_config, vllm_config
         )
         self.enable_partial_hash_hits = partial_hash_hits_enabled(
-            kv_cache_config.kv_cache_groups,
+            store_groups,
             self._hash_block_size,
             vllm_config.parallel_config.decode_context_parallel_size,
         )
         mamba_groups = {
             group_id: group.kv_cache_spec
-            for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
+            for group_id, group in enumerate(store_groups)
             if isinstance(group.kv_cache_spec, MambaSpec)
         }
         assert all(
@@ -173,8 +181,8 @@ def update_state_after_alloc(
         """Update state after block allocation."""
         local_block_ids: tuple[list[int], ...] = ()
         if num_external_tokens > 0:
-            local_block_ids = self.kv_cache_config.select_transfer_block_ids(
-                blocks.get_block_ids()
+            local_block_ids = blocks.get_block_ids(
+                group_ids=self._store_group_ids,
             )
 
         self._unfinished_requests[request.request_id] = (request, local_block_ids)
@@ -239,10 +247,7 @@ def build_connector_meta(
             request_real = request_tuple[0]  # type: ignore[index]
 
             unfolded_block_ids = tuple(
-                blocks.copy()
-                for blocks in self.kv_cache_config.select_transfer_block_ids(
-                    request.block_ids
-                )
+                request.block_ids[group_id].copy() for group_id in self._store_group_ids
             )
 
             prefill_tokens = _new_req_prefill_tokens(request)
@@ -278,8 +283,8 @@ def build_connector_meta(
             for i, req_id in enumerate(cached_reqs.req_ids):
                 new_block_ids = cached_reqs.new_block_ids[i]
                 if new_block_ids:
-                    new_block_ids = self.kv_cache_config.select_transfer_block_ids(
-                        new_block_ids
+                    new_block_ids = tuple(
+                        new_block_ids[group_id] for group_id in self._store_group_ids
                     )
 
                 req_meta = None
@@ -438,7 +443,9 @@ def _apply_current_save_block_ids(
             assert block_ids is not None, (
                 f"Missing current block table for store request {req_meta.req_id}"
             )
-            req_meta.block_ids = block_ids
+            req_meta.block_ids = tuple(
+                block_ids[group_id] for group_id in self._store_group_ids
+            )
 
     def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None:
         """Take a GPU block reference for every store job this step emits.
@@ -509,12 +516,15 @@ def register_finished_partial_tail(
             return False
 
         pinned_block_ids: list[int] = []
-        for group_id, block_id, _ in partial_tail_offloads:
-            if group_id not in self._boundary_state_group_ids:
+        remapped_offloads: list[tuple[int, int, int]] = []
+        for group_id, block_id, boundary in partial_tail_offloads:
+            store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id)
+            if store_group_id not in self._boundary_state_group_ids:
                 return False
             if block_id == NULL_BLOCK_ID:
                 return False
             pinned_block_ids.append(block_id)
+            remapped_offloads.append((store_group_id, block_id, boundary))
         pinned_block_ids = list(dict.fromkeys(pinned_block_ids))
 
         pool = self._gpu_block_pool
@@ -530,12 +540,14 @@ def register_finished_partial_tail(
         self._finished_partial_tail_metas[request.request_id] = ReqMeta(
             req_id=request.request_id,
             token_len_chunk=0,
-            block_ids=tuple(group.copy() for group in block_ids),
+            block_ids=tuple(
+                block_ids[group_id].copy() for group_id in self._store_group_ids
+            ),
             block_hashes=list(request.block_hashes),
             can_save=True,
             num_prompt_tokens=tracker.prefill_end_tokens,
             store_job_id=store_job_id,
-            boundary_state_offloads=partial_tail_offloads,
+            boundary_state_offloads=remapped_offloads,
         )
         tracker.has_pending_offload = True
         # The store job owns exact block refs, so request cleanup need not wait.
@@ -574,9 +586,10 @@ def _handle_boundary_state_offloads(
                     continue
                 if block_id == NULL_BLOCK_ID:
                     continue
-                if group_id not in self._boundary_state_group_ids:
+                store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id)
+                if store_group_id not in self._boundary_state_group_ids:
                     continue
-                accepted.append((group_id, block_id, boundary_tokens))
+                accepted.append((store_group_id, block_id, boundary_tokens))
             if not accepted:
                 continue
             tracker.has_pending_offload = True
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
index 870500e84068..baf8967bc859 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
@@ -11,7 +11,6 @@
 """
 
 import dataclasses
-import json
 import math
 import os
 import queue
@@ -20,10 +19,8 @@
 import time
 from collections.abc import Callable, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
-from dataclasses import dataclass
-from typing import Any, Literal, TypeVar
+from typing import Any, TypeVar
 
-import regex as re
 import torch
 import zmq
 
@@ -36,6 +33,9 @@
     get_tensor_model_parallel_world_size,
 )
 from vllm.distributed.kv_events import BlockStored
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorTransferResults,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import (  # noqa: E501
     ExternalCachedBlockPool,
@@ -64,9 +64,11 @@
     decode_lookup_response,
     encode_lookup_response,
 )
+from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store
 from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
 from vllm.utils.network_utils import get_ip, make_zmq_socket
+from vllm.utils.torch_utils import is_non_overlapping_and_dense
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
@@ -90,10 +92,6 @@
 
 logger = init_logger(__name__)
 
-DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
-DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
-DEFAULT_TENANT_ID = "default"
-
 MOONCAKE_NO_AVAILABLE_HANDLE = -200
 _T = TypeVar("_T")
 
@@ -145,115 +143,6 @@ def _make_mooncake_group_id(metadata: KeyMetadata, chunk_hash: str) -> str:
 _DIRECT_IO_PADDING_BYTES = 2 * _DIRECT_IO_ALIGNMENT
 
 
-MooncakeMode = Literal["embedded", "standalone-store"]
-
-
-@dataclass
-class MooncakeStoreConfig:
-    """Configuration for MooncakeDistributedStore.
-
-    ``mode`` selects the topology: ``embedded`` (each rank contributes
-    ``global_segment_size`` in-process) or ``standalone-store`` (rank
-    contributes 0; an external ``mooncake_client`` process owns the pool
-    and the SSD tier).
-    """
-
-    metadata_server: str
-    master_server_address: str
-    protocol: str
-    device_name: str
-    mode: MooncakeMode = "embedded"
-    global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE
-    local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE
-    enable_offload: bool = False
-    tenant_id: str = DEFAULT_TENANT_ID
-
-    def __post_init__(self) -> None:
-        if self.mode not in ("embedded", "standalone-store"):
-            raise ValueError(f"unknown Mooncake mode: {self.mode!r}")
-        if self.local_buffer_size <= 0:
-            raise ValueError("local_buffer_size must be > 0")
-        if self.mode == "embedded" and self.global_segment_size == 0:
-            raise ValueError("embedded mode requires global_segment_size > 0")
-        if self.mode == "standalone-store" and self.global_segment_size != 0:
-            raise ValueError("standalone-store mode requires global_segment_size == 0")
-
-    @staticmethod
-    def from_file(file_path: str) -> "MooncakeStoreConfig":
-        with open(file_path) as file:
-            config = json.load(file)
-        return MooncakeStoreConfig(
-            metadata_server=config.get("metadata_server", ""),
-            master_server_address=config.get("master_server_address", ""),
-            protocol=config.get("protocol", "rdma"),
-            device_name=config.get("device_name", ""),
-            mode=config.get("mode", "embedded"),
-            global_segment_size=_parse_size(
-                config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
-            ),
-            local_buffer_size=_parse_size(
-                config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
-            ),
-            enable_offload=bool(config.get("enable_offload", False)),
-            tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)),
-        )
-
-    @staticmethod
-    def load_from_config() -> "MooncakeStoreConfig":
-        config_path = os.getenv("MOONCAKE_CONFIG_PATH")
-        if not config_path:
-            raise ValueError(
-                "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
-            )
-        return MooncakeStoreConfig.from_file(config_path)
-
-
-def _normalize_tenant_id(value: Any) -> str:
-    if value is None:
-        return DEFAULT_TENANT_ID
-    if not isinstance(value, str):
-        raise TypeError(
-            f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}"
-        )
-    tenant_id = value.strip()
-    return tenant_id if tenant_id else DEFAULT_TENANT_ID
-
-
-def _parse_size(value: Any) -> int:
-    """Parse storage size strings with units: GB, MB, KB, B."""
-    if isinstance(value, int):
-        return value
-    if not isinstance(value, str):
-        try:
-            return int(value)
-        except (TypeError, ValueError) as e:
-            raise TypeError(f"Unsupported type for size: {type(value)}") from e
-
-    cleaned = value.strip().lower()
-    if not cleaned:
-        raise ValueError("Size cannot be empty.")
-
-    unit_multipliers = {
-        "gb": 1024**3,
-        "mb": 1024**2,
-        "kb": 1024,
-        "b": 1,
-    }
-    match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned)
-    if not match:
-        raise ValueError(f"Invalid format: '{value}'")
-
-    number_str = match.group(1)
-    unit = match.group(2) or "b"
-    multiplier = unit_multipliers[unit]
-
-    try:
-        numeric_value = float(number_str)
-    except ValueError as exc:
-        raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc
-    return int(numeric_value * multiplier)
-
-
 def _align_up(value: int, alignment: int) -> int:
     return ((value + alignment - 1) // alignment) * alignment
 
@@ -811,6 +700,7 @@ def _maybe_offload_boundary_states(self, req_meta: ReqMeta) -> bool:
 
         Returns:
             True when no put is needed or every put succeeds, False otherwise.
+
         """
         offloads = req_meta.boundary_state_offloads
         if not offloads or not req_meta.block_hashes:
@@ -1264,6 +1154,7 @@ def __init__(
         record_operation: Callable[..., None] | None = None,
         request_queue: queue.Queue[Any] | None = None,
         group_participates: Sequence[bool] | None = None,
+        is_hma_required: bool = False,
     ):
         super().__init__(
             store,
@@ -1283,6 +1174,10 @@ def __init__(
         # _invalid_block_ids can be access by both the Worker and RecvingThread
         self._invalid_block_ids_lock = threading.Lock()
         self._invalid_block_ids: set[int] = set()
+        # With HMA the scheduler tracks a single merged group while block IDs
+        # are only unique within a group, so failures are reported per request.
+        self._is_hma_required = is_hma_required
+        self._failed_requests: set[str] = set()
         self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes
         self.usable_disk_offload_buffer_budget_bytes = (
             None
@@ -1297,6 +1192,22 @@ def _add_load_error_block_ids(self, block_ids: list[int]) -> None:
         with self._invalid_block_ids_lock:
             self._invalid_block_ids.update(block_ids)
 
+    def set_failed_request(self, req_id: str):
+        with self.done_task_lock:
+            self._failed_requests.add(req_id)
+
+    def get_and_clear_failed_requests(self) -> set[str]:
+        with self.done_task_lock:
+            failed = self._failed_requests.copy()
+            self._failed_requests.clear()
+        return failed
+
+    def _report_load_error(self, req_id: str, block_ids: list[int]) -> None:
+        if self._is_hma_required:
+            self.set_failed_request(req_id)
+        else:
+            self._add_load_error_block_ids(block_ids)
+
     def get_and_clear_block_ids_with_load_errors(self) -> set[int]:
         with self._invalid_block_ids_lock:
             invalid_block_ids = self._invalid_block_ids.copy()
@@ -1391,7 +1302,7 @@ def _handle_request(self, req_meta: ReqMeta):
                     # Mark every block: we skip the whole request, and the
                     # tp_rank rotation means oversized_key isn't necessarily
                     # the first block in the request's original order.
-                    self._add_load_error_block_ids(block_id_list_c)
+                    self._report_load_error(req_id, block_id_list_c)
                     oversized_key_bytes = _estimate_disk_offload_staging_bytes(
                         size_list_c[oversized_key_index]
                     )
@@ -1464,8 +1375,8 @@ def _handle_request(self, req_meta: ReqMeta):
                     num_failed_keys=len(failed),
                 )
                 if failed:
-                    self._add_load_error_block_ids(
-                        [block_id for _, _, block_id in failed]
+                    self._report_load_error(
+                        req_id, [block_id for _, _, block_id in failed]
                     )
                     logger.warning(
                         "Failed to get %d Mooncake keys from sub-batch "
@@ -1476,7 +1387,7 @@ def _handle_request(self, req_meta: ReqMeta):
                     )
                     break
         except Exception as e:
-            self._add_load_error_block_ids(current_batch_block_ids)
+            self._report_load_error(req_id, current_batch_block_ids)
             self._record_operation(
                 "load_get",
                 load_get_start,
@@ -1550,6 +1461,7 @@ def __init__(
             and not self.can_put
         )
         self.cache_config = vllm_config.cache_config
+        self._is_hma_required = len(kv_cache_config.kv_cache_groups) > 1
         self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes(
             kv_cache_config, vllm_config
         )
@@ -1562,23 +1474,7 @@ def __init__(
         self.store = MooncakeDistributedStore()
         local_ip = get_ip()
         local_hostname = rdma_utils.get_requester_local_hostname(local_ip)
-        setup_kwargs: dict[str, str] = {}
-        if store_config.tenant_id != DEFAULT_TENANT_ID:
-            setup_kwargs["tenant_id"] = store_config.tenant_id
-        ret = self.store.setup(
-            local_hostname,
-            store_config.metadata_server,
-            store_config.global_segment_size,
-            store_config.local_buffer_size,
-            store_config.protocol,
-            store_config.device_name,
-            store_config.master_server_address,
-            **setup_kwargs,
-        )
-        if ret != 0:
-            msg = "Initialize MooncakeDistributedStore failed."
-            logger.error(msg)
-            raise RuntimeError(msg)
+        setup_mooncake_store(self.store, store_config, local_hostname)
 
         preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config)
         self.preferred_segment = preferred_segment
@@ -1671,10 +1567,11 @@ def __init__(
             dataclasses.replace(
                 group,
                 kv_cache_spec=resolve_dcp_kv_cache_spec(
-                    group.kv_cache_spec, self.dcp_size
+                    group.kv_cache_spec,
+                    self.dcp_size,
                 ),
             )
-            for group in kv_cache_config.transfer_groups
+            for group in kv_cache_config.prefix_cacheable_groups
         ]
         spec_cfg = getattr(vllm_config, "speculative_config", None)
         use_eagle_block_drop = bool(
@@ -1909,21 +1806,37 @@ def register_kv_caches(
             logger.warning("No KV caches to offload.")
             return
 
+        # Resolve each entry to a representative tensor for storage
+        # deduplication. For attention layers the value is already a tensor;
+        # for Mamba layers it is a list of tensors that all share the same
+        # underlying raw storage, so we take the first one.
+        def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
+            assert isinstance(v, torch.Tensor | list)
+            return v if isinstance(v, torch.Tensor) else v[0]
+
         assert self.cache_config.num_gpu_blocks is not None
         self.num_blocks = self.cache_config.num_gpu_blocks
-
-        seen_storage_ptrs: set[int] = set()
-        cache_tensors: list[torch.Tensor] = []
-
-        for cache in kv_caches.values():
-            cache = group_kernel_blocks(cache, self.num_blocks)
-            cache_tensors.append(cache)
-            cache_storage = cache.untyped_storage()
-            base_addr = cache_storage.data_ptr()
-            region_len = cache_storage.nbytes()
-
-            if base_addr not in seen_storage_ptrs:
+        use_group_regions = self._kv_cache_config.hisparse_host_num_blocks is not None
+
+        if not use_group_regions:
+            store_layer_names = {
+                layer_name
+                for group in self._kv_cache_groups
+                for layer_name in group.layer_names
+            }
+            seen_storage_ptrs: set[int] = set()
+            cache_tensors = [
+                group_kernel_blocks(_repr_tensor(cache), self.num_blocks)
+                for layer_name, cache in kv_caches.items()
+                if layer_name in store_layer_names
+            ]
+            for cache in cache_tensors:
+                cache_storage = cache.untyped_storage()
+                base_addr = cache_storage.data_ptr()
+                if base_addr in seen_storage_ptrs:
+                    continue
                 seen_storage_ptrs.add(base_addr)
+                region_len = cache_storage.nbytes()
                 ret = self.store.register_buffer(base_addr, region_len)
                 if ret != 0:
                     logger.error(
@@ -1932,17 +1845,110 @@ def register_kv_caches(
                         region_len,
                         ret,
                     )
+            for db in self.token_dbs:
+                db.store_layout.register_kv_caches(cache_tensors, self.num_blocks)
+
+        registered_buffers: dict[int, int] = {}
+        num_host_segments = 0
+        num_device_segments = 0
+
+        for group, db in zip(self._kv_cache_groups, self.token_dbs, strict=True):
+            if not use_group_regions:
+                continue
+            missing_layers = [
+                layer_name
+                for layer_name in group.layer_names
+                if layer_name not in kv_caches
+            ]
+            if missing_layers:
+                raise ValueError(
+                    f"Missing transferable KV caches for layers: {missing_layers}"
+                )
+            group_caches = [
+                (layer_name, kv_caches[layer_name]) for layer_name in group.layer_names
+            ]
+
+            seen_group_ptrs: set[int] = set()
+            addrs: list[int] = []
+            block_lens: list[int] = []
+            for layer_name, value in group_caches:
+                cache = _repr_tensor(value)
+                is_host_resident = group.host_resident
+                if is_host_resident:
+                    num_blocks = self._kv_cache_config.hisparse_host_num_blocks
+                    assert num_blocks is not None
+                else:
+                    num_blocks = self._kv_cache_config.num_blocks
+                cache = group_kernel_blocks(cache, num_blocks)
+                cache_storage = cache.untyped_storage()
+                if is_host_resident:
+                    base_addr = cache.data_ptr()
+                    region_len = cache.nbytes
+                else:
+                    base_addr = cache_storage.data_ptr()
+                    region_len = cache_storage.nbytes()
+                if region_len % num_blocks != 0:
+                    raise ValueError(
+                        f"KV cache region for {layer_name} has size {region_len}, "
+                        f"which is not divisible by its {num_blocks} blocks."
+                    )
+
+                registered_len = registered_buffers.get(base_addr)
+                if registered_len is None:
+                    ret = self.store.register_buffer(base_addr, region_len)
+                    if ret != 0:
+                        mem_kind = "host" if is_host_resident else "device"
+                        raise RuntimeError(
+                            "Mooncake register_buffer failed for addr "
+                            f"{base_addr:#x} len {region_len} ({mem_kind}): {ret}"
+                        )
+                    registered_buffers[base_addr] = region_len
+                elif registered_len != region_len:
+                    raise ValueError(
+                        f"KV cache views at {base_addr:#x} expose inconsistent "
+                        f"region sizes: {registered_len} and {region_len}."
+                    )
+
+                regions: list[tuple[int, int]] = []
+                if not is_non_overlapping_and_dense(cache[0]):
+                    for head_idx in range(cache.shape[1]):
+                        head_cache = cache[:, head_idx]
+                        assert is_non_overlapping_and_dense(head_cache[0])
+                        regions.append(
+                            (
+                                head_cache.data_ptr(),
+                                head_cache.stride(0) * head_cache.element_size(),
+                            )
+                        )
+                elif cache.stride(0) * cache.element_size() * num_blocks == region_len:
+                    regions.append((base_addr, region_len // num_blocks))
+                else:
+                    regions.append(
+                        (cache.data_ptr(), cache.stride(0) * cache.element_size())
+                    )
+
+                for region_addr, block_len in regions:
+                    if region_addr in seen_group_ptrs:
+                        continue
+                    seen_group_ptrs.add(region_addr)
+                    addrs.append(region_addr)
+                    block_lens.append(block_len)
+                    if is_host_resident:
+                        num_host_segments += 1
+                    else:
+                        num_device_segments += 1
+
+            db.set_kv_caches_base_addr(addrs)
+            db.set_block_len(block_lens)
 
         logger.info(
-            "Registered KV caches: num_groups=%d, num_tensors=%d, num_blocks=%d",
+            "Registered KV caches: num_groups=%d, num_segments=%d (host=%d, device=%d)",
             len(self.token_dbs),
-            len(cache_tensors),
-            self.num_blocks,
+            num_host_segments + num_device_segments,
+            num_host_segments,
+            num_device_segments,
         )
 
-        for db in self.token_dbs:
-            db.store_layout.register_kv_caches(cache_tensors, self.num_blocks)
-
         # Start transfer threads
         if self.can_put:
             ready_event_sending = threading.Event()
@@ -1985,6 +1991,7 @@ def register_kv_caches(
                     group.kv_cache_spec.prefix_cacheable
                     for group in self._kv_cache_groups
                 ],
+                is_hma_required=self._is_hma_required,
             )
             recv_thread.name = f"KVCacheStoreRecvingThread-{i}"
             recv_thread.start()
@@ -2074,6 +2081,26 @@ def get_block_ids_with_load_errors(self) -> set[int]:
             block_ids |= recv_thread.get_and_clear_block_ids_with_load_errors()
         return block_ids
 
+    def get_transfer_results(
+        self, finished_req_ids: set[str], meta: MooncakeStoreConnectorMetadata
+    ) -> KVConnectorTransferResults:
+        """Get completed sends/recvs plus requests whose remote KV load failed."""
+        done_sending, done_recving = self.get_finished(finished_req_ids, meta)
+
+        if self._capacity_only:
+            return KVConnectorTransferResults(done_sending, done_recving)
+
+        failed_recving: set[str] = set()
+        if self.load_async:
+            for recv_thread in self.kv_recv_threads:
+                failed_recving |= recv_thread.get_and_clear_failed_requests()
+
+        return KVConnectorTransferResults(
+            finished_sending=done_sending,
+            finished_recving=done_recving,
+            failed_recving=failed_recving,
+        )
+
     def _record_kv_connector_operation(
         self,
         operation: str,
@@ -2316,6 +2343,19 @@ def close(self) -> None:
         if store is None:
             return
         self.store = None
+        deadline = time.monotonic() + 5.0
+        for thread in (self.kv_send_thread, *self.kv_recv_threads):
+            if thread is None:
+                continue
+            while thread.request_queue.unfinished_tasks and time.monotonic() < deadline:
+                time.sleep(0.05)
+            if thread.request_queue.unfinished_tasks:
+                logger.warning(
+                    "Mooncake store %s still has %d in-flight requests at "
+                    "close; closing anyway.",
+                    thread.name,
+                    thread.request_queue.unfinished_tasks,
+                )
         try:
             store.close()
         except Exception as e:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
index a3010203ee76..5f1068562cf8 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import contextlib
+import math
 import os
 import threading
 import time
@@ -14,6 +15,7 @@
 import zmq
 
 from vllm.config import KVTransferConfig, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorMetadata,
 )
@@ -44,6 +46,13 @@
 TransferOffsetsKey = tuple[str, tuple[int, ...], tuple[int, ...], torch.dtype]
 
 
+def _positive_finite_timeout(name: str, value: Any) -> float:
+    timeout = float(value)
+    if not math.isfinite(timeout) or timeout <= 0:
+        raise ValueError(f"{name} must be finite and greater than zero")
+    return timeout
+
+
 class MoRIIOTransferAck(NamedTuple):
     transfer_id: TransferId
     consumer_tp_size: int = 1
@@ -159,6 +168,20 @@ class MoRIIOMode(Enum):
     WRITE = "write"
 
 
+class TransferBatchState(Enum):
+    """Verdict for a group of transfer statuses that belong to one request.
+
+    A request's KV transfer is spread over one status per layer (two for a KDA
+    layer), so a single status never decides the request: PENDING means at
+    least one is still in flight and none has failed, FAILED means at least
+    one failed, DONE means all succeeded.
+    """
+
+    DONE = "done"
+    FAILED = "failed"
+    PENDING = "pending"
+
+
 class MoRIIOError(Exception):
     """Base exception for MoRIIO operations."""
 
@@ -259,6 +282,7 @@ class MoRIIOConfig:
     tp_size: int
     transfer_timeout: float
     defer_timeout: float
+    recv_abort_timeout: float
     read_mode: bool = False
     qp_per_transfer: int = 1
     post_batch_size: int = -1
@@ -283,6 +307,8 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
         #                     raising TransferError (sec).
         # defer_timeout    -> Timeout before a deferred send with no finished_sending
         #                     notification is reaped and its blocks force-freed (sec).
+        # recv_abort_timeout -> Timeout before an in-flight recv whose RDMA
+        #                     completion never arrived is aborted (sec).
 
         # Knobs for RDMA transfers, ignored if on xgmi backend
         # qp_per_transfer  -> Number of RDMA Queue Pairs per KV transfer.
@@ -314,13 +340,29 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
                 "must be one of 'rdma' or 'xgmi'."
             )
 
-        transfer_timeout = float(
+        transfer_timeout = _positive_finite_timeout(
+            "kv_connector_extra_config.transfer_timeout",
             extra_config.get(
                 "transfer_timeout", MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT
-            )
+            ),
+        )
+        defer_timeout = _positive_finite_timeout(
+            "kv_connector_extra_config.defer_timeout",
+            extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT),
         )
-        defer_timeout = float(
-            extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT)
+        legacy_recv_abort_timeout = os.environ.get("VLLM_MORIIO_TRANSFER_TIMEOUT_S")
+        if legacy_recv_abort_timeout is not None:
+            logger.warning_once(
+                "The environment variable VLLM_MORIIO_TRANSFER_TIMEOUT_S is "
+                "deprecated. Set 'recv_abort_timeout' inside "
+                "kv_transfer_config.kv_connector_extra_config instead."
+            )
+        recv_abort_timeout = _positive_finite_timeout(
+            "kv_connector_extra_config.recv_abort_timeout",
+            extra_config.get(
+                "recv_abort_timeout",
+                legacy_recv_abort_timeout or MoRIIOConstants.DEFAULT_RECV_ABORT_TIMEOUT,
+            ),
         )
 
         return cls(
@@ -346,6 +388,7 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
             backend=backend,
             transfer_timeout=transfer_timeout,
             defer_timeout=defer_timeout,
+            recv_abort_timeout=recv_abort_timeout,
         )
 
 
@@ -373,6 +416,10 @@ class MoRIIOConstants:
     # notification is reaped and its blocks force-freed.
     # Overridable via kv_connector_extra_config["defer_timeout"].
     DEFAULT_DEFER_TIMEOUT = 60.0
+    # Timeout (seconds) before an in-flight recv whose RDMA completion was lost
+    # is aborted, so the decode worker does not hang on it forever.
+    # Overridable via kv_connector_extra_config["recv_abort_timeout"].
+    DEFAULT_RECV_ABORT_TIMEOUT = 120.0
 
 
 # The router embeds both zmq_addresses in the request_id:
@@ -444,8 +491,8 @@ class ReqMeta:
     """Metadata for a single request."""
 
     transfer_id: TransferId
-    local_block_ids: list[int]
-    remote_block_ids: list[int]
+    local_block_ids: BlockIds
+    remote_block_ids: BlockIds
     remote_host: str
     remote_port: int
     remote_handshake_port: int
@@ -481,7 +528,7 @@ def __repr__(self):
     def add_new_req(
         self,
         request_id: ReqId,
-        local_block_ids: list[int],
+        local_block_ids: BlockIds,
         kv_transfer_params: dict[str, Any],
         write_mode=False,
     ):
@@ -566,9 +613,12 @@ def add_new_req(
             remote_handshake_port=int(remote_handshake_port),
             remote_notify_port=int(remote_notify_port),
             # Remote peer TP degree (used as remote_tp_size downstream). The
-            # proxy advertises it under "remote_tp_size"; #46332 read "tp_size"
-            # which is absent on WRITE producer requests -> defaulted to 1 ->
-            # rank collapse. Read the right key; 0 == unknown (== homogeneous).
+            # decode side must know the prefiller's TP degree to map producer
+            # ranks onto the correct decode rank (and vice versa). The old plain
+            # "tp_size" key (#46332) is absent on WRITE producer requests and
+            # defaulted to 1, collapsing every producer rank onto decode rank 0
+            # and corrupting the transfer; prefer "remote_tp_size". 0 == unknown
+            # (treated as homogeneous downstream).
             tp_size=int(
                 kv_transfer_params.get("remote_tp_size")
                 or kv_transfer_params.get("tp_size")
@@ -587,8 +637,7 @@ def add_new_req(
 
 @contextlib.contextmanager
 def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
-    """Context manager for a ZMQ socket"""
-
+    """Context manager for a ZMQ socket."""
     if socket_type not in (zmq.ROUTER, zmq.REQ, zmq.DEALER):
         raise ValueError(f"Unexpected socket type: {socket_type}")
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
index 349cf50775ad..dce4a9806895 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
@@ -2,14 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import logging
 import math
-import os
 import queue
 import threading
 import time
 from collections import defaultdict
 from collections.abc import Collection
 from concurrent.futures import Future, ThreadPoolExecutor
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import msgpack
 import msgspec
@@ -18,10 +17,12 @@
 import zmq
 
 from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorBase_V1,
     KVConnectorMetadata,
     KVConnectorRole,
+    SupportsHMA,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
     ROLE,
@@ -31,10 +32,13 @@
     MoRIIOConfig,
     MoRIIOConnectorMetadata,
     MoRIIOConstants,
+    MoRIIOError,
     MoRIIOMode,
     MoRIIOTransferAck,
     ReqId,
     ReqMeta,
+    TransferBatchState,
+    TransferError,
     TransferId,
     WriteTask,
     fold_local_rank,
@@ -54,11 +58,20 @@
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import (
     LayerTransferGeometry,
+    MambaOffsetTemplate,
+    apply_mamba_offset_template,
     build_layer_to_spec,
+    build_mamba_offset_template,
     compute_block_transfer_offsets,
+    compute_mamba_conv_split_count,
     get_layer_transfer_geometry,
     is_mla_cache_layer,
     iter_layer_registration_regions,
+    kda_conv_ssm,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
+    MambaConvSplitInfo,
+    derive_mamba_conv_split,
 )
 from vllm.distributed.parallel_state import (
     get_tensor_model_parallel_world_size,
@@ -67,23 +80,36 @@
 )
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
 from vllm.utils.network_utils import (
     make_zmq_path,
     make_zmq_socket,
 )
 from vllm.v1.attention.selector import get_attn_backend
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    MambaSpec,
+    SlidingWindowSpec,
+)
 from vllm.v1.outputs import KVConnectorOutput
 from vllm.v1.request import RequestStatus
 
 if TYPE_CHECKING:
     from vllm.v1.attention.backend import AttentionMetadata
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
-    from vllm.v1.kv_cache_interface import KVCacheConfig
     from vllm.v1.request import Request
 
 logger = init_logger(__name__)
 
+# Backoff bounds for re-posting an RDMA READ that the NIC rejected because the
+# send queue was full. A separate CQ-poll thread drains completions, so the
+# wait only needs to be long enough for it to make progress.
+_SQ_FULL_BACKOFF_INITIAL_S = 0.001
+_SQ_FULL_BACKOFF_MAX_S = 0.05
+_MAX_LOCAL_DECODE_TAIL_BLOCKS = 1
+
 
 try:
     from mori.io import (
@@ -189,7 +215,13 @@ def resolve_moriio_transfer_ack(
     return transfer_id
 
 
-class MoRIIOConnector(KVConnectorBase_V1):
+class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
+    @property
+    def supports_divergent_local_hybrid_hits(self) -> bool:
+        # The READ path always transfers the recurrent-state slot, including
+        # when the attention prefix is already present locally.
+        return self.mode == MoRIIOMode.READ and self._transfers_mamba_state
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -210,22 +242,13 @@ def __init__(
             + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
         )
         self.mode = get_moriio_mode(self.kv_transfer_config)
-        if (
-            self.mode == MoRIIOMode.READ
-            and self.kv_transfer_config.is_kv_consumer
-            and vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs()
-        ):
-            # warn only; kv-read barrier requires PIECEWISE cudagraph mode
-            logger.warning_once(
-                "MoRIIO READ mode is running with %s CUDA graphs: per-layer "
-                "KV-read barrier can't fire inside full graph; accuracy may "
-                "degrade at high concurrency. Set cudagraph_mode=PIECEWISE "
-                "in --compilation-config.",
-                vllm_config.compilation_config.cudagraph_mode.name,
-            )
+        self._transfers_mamba_state = any(
+            group.enable_kv_transfer and isinstance(group.kv_cache_spec, MambaSpec)
+            for group in kv_cache_config.kv_cache_groups
+        )
         if role == KVConnectorRole.SCHEDULER:
             self.connector_scheduler: MoRIIOConnectorScheduler | None = (
-                MoRIIOConnectorScheduler(vllm_config, self.engine_id)
+                MoRIIOConnectorScheduler(vllm_config, self.engine_id, kv_cache_config)
             )
             self.connector_worker: MoRIIOConnectorWorker | None = None
         elif role == KVConnectorRole.WORKER:
@@ -272,6 +295,11 @@ def update_state_after_alloc(
             request, blocks, num_external_tokens, self.connector_worker
         )
 
+    def on_new_request(self, request: "Request") -> None:
+        """Fired when the scheduler admits a request, before any allocation."""
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.on_new_request(request)
+
     def build_connector_meta(
         self,
         scheduler_output: SchedulerOutput,
@@ -285,6 +313,22 @@ def request_finished(
         block_ids: list[int],
     ) -> tuple[bool, dict[str, Any] | None]:
         assert self.connector_scheduler is not None
+        return self.connector_scheduler.request_finished(request, (block_ids,))
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: BlockIds,
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        if self.connector_scheduler._has_mamba:
+            attn_block_ids, mamba_block_groups = (
+                self.connector_scheduler.split_block_groups(block_ids)
+            )
+            delay_free, params = self.connector_scheduler.request_finished(
+                request, attn_block_ids, mamba_block_groups
+            )
+            return delay_free, params
         return self.connector_scheduler.request_finished(request, block_ids)
 
     def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
@@ -303,6 +347,12 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
         assert self.connector_worker is not None
         return self.connector_worker.get_finished()
 
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """Blocks whose MoRIIO read failed, for the scheduler to recompute."""
+        if self.connector_worker is None:
+            return set()
+        return self.connector_worker.get_block_ids_with_load_errors()
+
     def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         assert self.connector_worker is not None
         if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.CONSUMER:
@@ -358,6 +408,7 @@ def has_connector_metadata(self) -> bool:
 
         Returns:
             bool: True if connector metadata exists, False otherwise.
+
         """
         try:
             return self._connector_metadata is not None
@@ -365,19 +416,162 @@ def has_connector_metadata(self) -> bool:
             return False
 
 
+def _split_kv_cache_group_kinds(
+    kv_cache_config: "KVCacheConfig",
+) -> tuple[list[int], list[int]]:
+    """Classify kv_cache_groups into (attention_group_indices,
+    mamba_group_indices) by whether each group's spec is a MambaSpec.
+
+    Groups that opt out of external transfer are absent from
+    ``KVCacheConfig.transfer_groups``. Returned indices address that filtered
+    sequence and therefore also the connector's transfer block payload.
+    """
+    attn: list[int] = []
+    mamba: list[int] = []
+    for gi, group in enumerate(kv_cache_config.transfer_groups):
+        if isinstance(group.kv_cache_spec, MambaSpec):
+            mamba.append(gi)
+        else:
+            attn.append(gi)
+    return attn, mamba
+
+
+def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None:
+    """Validate the common packed-state contract and return its representative."""
+    specs = tuple(specs)
+    if not specs:
+        return None
+    for spec in specs:
+        if len(spec.shapes) != 2 or len(spec.dtypes) != 2:
+            raise MoRIIOError(
+                "MoRIIO hybrid READ supports exactly two Mamba states "
+                "(conv and SSM), got "
+                f"{len(spec.shapes)} shapes and {len(spec.dtypes)} dtypes"
+            )
+    reference = specs[0]
+    if any(spec != reference for spec in specs[1:]):
+        raise MoRIIOError(
+            "MoRIIO hybrid READ requires all transferable Mamba groups "
+            "to use the same cache spec"
+        )
+    return reference
+
+
 class MoRIIOConnectorScheduler:
-    """Implementation of Scheduler side methods"""
+    """Implementation of Scheduler side methods."""
 
-    def __init__(self, vllm_config: VllmConfig, engine_id: str):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        engine_id: str,
+        kv_cache_config: "KVCacheConfig",
+    ):
         self.vllm_config = vllm_config
-
+        self.kv_cache_config = kv_cache_config
         assert vllm_config.kv_transfer_config is not None, (
             "kv_transfer_config must be set for MoRIIOConnector"
         )
         self.kv_transfer_config = vllm_config.kv_transfer_config
+        self.mode = get_moriio_mode(self.kv_transfer_config)
+        self._attn_group_ids, self._mamba_group_ids = _split_kv_cache_group_kinds(
+            kv_cache_config
+        )
+        self._has_mamba = bool(self._mamba_group_ids)
+        if self._has_mamba:
+            if len(self._attn_group_ids) != 1:
+                raise MoRIIOError(
+                    "MoRIIO hybrid READ requires exactly one transferable "
+                    "attention group, got "
+                    f"{len(self._attn_group_ids)} attention groups"
+                )
+            mamba_specs = [
+                cast(MambaSpec, kv_cache_config.transfer_groups[group_id].kv_cache_spec)
+                for group_id in self._mamba_group_ids
+            ]
+            _validate_mamba_specs(mamba_specs)
+            if vllm_config.speculative_config is not None:
+                raise MoRIIOError(
+                    "MoRIIO hybrid READ does not support speculative decoding"
+                )
+            if self.mode != MoRIIOMode.READ:
+                raise MoRIIOError(
+                    "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ "
+                    "mode only; set kv_connector_extra_config.read_mode=true"
+                )
+        # Only "all" mode keeps a state per block position; the other modes keep
+        # a single running state in the last slot.
+        self._ssm_state_slots_are_positional = (
+            vllm_config.cache_config.mamba_cache_mode == "all"
+        )
         self.block_size = vllm_config.cache_config.block_size
+        self._max_decode_tail_blocks = (
+            _MAX_LOCAL_DECODE_TAIL_BLOCKS
+            if self._has_mamba
+            else math.ceil((vllm_config.num_lookahead_tokens + 1) / self.block_size)
+        )
         self.engine_id: EngineId = engine_id
-        self.mode = get_moriio_mode(self.kv_transfer_config)
+
+        self._is_hma_required = (
+            not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
+            and any(
+                not isinstance(g.kv_cache_spec, FullAttentionSpec)
+                for g in kv_cache_config.transfer_groups
+            )
+        )
+        if self._is_hma_required:
+            # MoRIIO handles sliding-window attention and mamba/KDA groups in
+            # READ mode. Other hybrid cache specs still fail closed.
+            unsupported = [
+                type(g.kv_cache_spec).__name__
+                for g in kv_cache_config.transfer_groups
+                if not isinstance(
+                    g.kv_cache_spec,
+                    (FullAttentionSpec, SlidingWindowSpec, MambaSpec),
+                )
+            ]
+            if unsupported:
+                raise NotImplementedError(
+                    "MoRIIO only supports sliding-window hybrid models (e.g. "
+                    "Gemma) and mamba/KDA groups under the hybrid KV cache "
+                    "manager. Unsupported KV "
+                    f"cache group spec(s): {sorted(set(unsupported))}. Pass "
+                    "--disable-hybrid-kv-cache-manager to run without HMA."
+                )
+            if self.mode == MoRIIOMode.WRITE:
+                # TODO(simondanielsson): support HMA in WRITE mode
+                raise NotImplementedError(
+                    "MoRIIO WRITE mode does not support hybrid KV cache groups "
+                    "(sliding-window attention). Use READ mode, or pass "
+                    "--disable-hybrid-kv-cache-manager."
+                )
+
+        sw_sizes_tokens: list[tuple[int, int]] = [
+            (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
+            if isinstance(g.kv_cache_spec, SlidingWindowSpec)
+            else (0, self.block_size)
+            for g in kv_cache_config.transfer_groups
+        ]
+        # add 1 to conservatively account for boundary overlap eg window isn't fully
+        # aligned with blocks.
+        self.blocks_per_sw = [
+            cdiv(n_tokens, block_size) + 1 if n_tokens else 0
+            for n_tokens, block_size in sw_sizes_tokens
+        ]
+        # In WRITE mode and chunked prefill, we perform the write after the last chunk.
+        # Hence we need to track once we are in the last chunk. We do this by computing
+        # len(block_ids[g]) * block_size[g] for a full attn group g, as hybrid groups
+        # may be clipped and hence not reflect the full sequence length. Here we store
+        # g and block_size[g] for later.
+        self._full_attn_group_idx = 0
+        self._full_attn_block_size = self.block_size
+        for gi, group in enumerate(kv_cache_config.kv_cache_groups):
+            if isinstance(group.kv_cache_spec, FullAttentionSpec):
+                self._full_attn_group_idx = gi
+                self._full_attn_block_size = getattr(
+                    group.kv_cache_spec, "block_size", self.block_size
+                )
+                break
+
         self.host_ip = resolve_host_ip(
             self.kv_transfer_config.kv_connector_extra_config
         )
@@ -409,14 +603,14 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
         # Requests that need to start recv/send.
         # New requests are added by update_state_after_alloc in
         # the scheduler. Used to make metadata passed to Worker.
-        self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
-        self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
+        self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
+        self._reqs_need_save: dict[ReqId, tuple[Request, BlockIds]] = {}
         # Snapshot of kv_transfer_params for chunked prefill recovery.
         self._req_kv_params: dict[ReqId, dict] = {}
 
         # For chunked prefill, we perform layer-wise access within the final chunk.
         # TODO: Perform transfer at end chunk.
-        self._reqs_need_pending_save: dict[ReqId, tuple[Request, list[int]]] = {}
+        self._reqs_need_pending_save: dict[ReqId, tuple[Request, BlockIds]] = {}
 
         if self.is_producer:
             set_role(ROLE.PRODUCER)
@@ -442,6 +636,27 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
         self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}
         self.request_id_to_transfer_id: dict[ReqId, TransferId] = {}
 
+    def get_exchange_clipped_blocks(
+        self,
+        block_ids: BlockIds,
+    ) -> BlockIds:
+        """Clip the block list based on the layer type.
+
+        Behavior:
+        - sliding window: keep only the last window size blocks
+        - full attention: keep all blocks
+        """
+        if len(block_ids) == 0:
+            return []
+        if len(block_ids) == len(self.kv_cache_config.kv_cache_groups):
+            block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids)
+        if not self._is_hma_required:
+            return [list(blocks) for blocks in block_ids]
+        return [
+            blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks
+            for i, blocks in enumerate(block_ids)
+        ]
+
     def map_request_id(self, request_id: ReqId, transfer_id: TransferId):
         self.transfer_id_to_request_id[transfer_id] = request_id
         self.request_id_to_transfer_id[request_id] = transfer_id
@@ -482,8 +697,7 @@ def get_num_new_matched_tokens(
         request: "Request",
         num_computed_tokens: int,
     ) -> tuple[int, bool]:
-        """
-        For remote prefill, pull all prompt blocks from remote
+        """For remote prefill, pull all prompt blocks from remote
         asynchronously relative to engine execution.
 
         Args:
@@ -495,23 +709,78 @@ def get_num_new_matched_tokens(
               external KV cache beyond what is already computed.
             * true if the external KV cache tokens will be loaded
               asynchronously (between scheduler steps).
+
         """
         if self.is_producer:
+            # The P-side prompt truncation for hybrid models happens in
+            # on_new_request, before the scheduler looks the request up in the
+            # prefix cache.
             return 0, False
 
-        token_ids = request.prompt_token_ids or []
+        num_prompt_tokens = request.num_prompt_tokens
+        num_external_tokens = max(num_prompt_tokens - num_computed_tokens, 0)
         if self.mode == MoRIIOMode.WRITE:
-            # MoriiO in write mode, no remote prefill
+            # MoriiO in write mode, no remote prefill. Hybrid models never get
+            # here: register_kv_caches refuses WRITE mode for them, so there is
+            # no recurrent-state accounting to do.
+            return num_external_tokens, True
 
-            return len(token_ids) - num_computed_tokens, True
+        # READ mode always recomputes the last token locally on the decoder.
+        #
+        # The second element declares who waits for the KV, not whether it has
+        # arrived. False keeps the wait on this side: the forward runs in the
+        # same step, and the connector blocks in wait_for_layer_load (or in
+        # start_load_kv under full-graph capture) until the KV is really there.
+        return max(num_external_tokens - 1, 0), False
+
+    def on_new_request(self, request: "Request") -> None:
+        """Scheduler hook fired when a request is admitted.
+
+        The hybrid P-side truncation must land here rather than in
+        ``get_num_new_matched_tokens``: the scheduler runs the local prefix
+        cache lookup first and only then asks the connector, so truncating
+        there shortens a prompt whose cache hit was already measured against
+        the full length. On a full hit that leaves ``num_computed_tokens ==
+        num_tokens`` and trips the scheduler's ``num_new_tokens > 0`` assert.
+        """
+        if not (self.is_producer and self._has_mamba):
+            return
+        params = request.kv_transfer_params
+        if params is not None and params.get("do_remote_decode"):
+            self._truncate_mamba_request_for_prefill(request)
 
-        return len(token_ids) - 1 - num_computed_tokens, False
+    def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
+        """P-side only: drop the last prompt token so the prefiller computes
+        h(N-1) instead of h(N).
+
+        The decoder recomputes the last token locally to derive h(N) correctly
+        (see the READ len-1 accounting in get_num_new_matched_tokens). Guarded
+        by ``_p_side_truncated`` to avoid repeated truncation across a
+        preemption/reschedule. Mirrors the NIXL scheduler.
+        """
+        params = request.kv_transfer_params
+        if (
+            params is not None
+            and not params.get("_p_side_truncated")
+            and request.num_prompt_tokens > 1
+        ):
+            if request.prompt_token_ids is not None:
+                request.prompt_token_ids.pop()
+            elif request.prompt_embeds is not None:
+                request.prompt_embeds = request.prompt_embeds[:-1]
+            else:
+                return
+
+            request._all_token_ids.pop()
+            request.num_prompt_tokens -= 1
+            request.max_tokens = 1
+            params["_p_side_truncated"] = True
 
     def send_notify_block(
         self,
         req_id: ReqId,
         transfer_id: TransferId,
-        block_notify_list: list[int],
+        block_notify_list: list,
         host=None,
         port=None,
     ):
@@ -641,8 +910,8 @@ def update_state_after_alloc(
         params.setdefault("transfer_id", transfer_id)
         request_id = request.request_id
         self.map_request_id(request_id, transfer_id)
-        if params.get("do_remote_decode"):
-            local_block_ids = blocks.get_block_ids()[0]
+        if params.get("do_remote_decode") and self.mode == MoRIIOMode.WRITE:
+            local_block_ids: BlockIds = blocks.get_block_ids()
             self._reqs_need_save[request.request_id] = (request, local_block_ids)
             # Snapshot params now so chunked-prefill build_connector_meta
             # can recover them on the final chunk even if the live
@@ -655,19 +924,88 @@ def update_state_after_alloc(
                     # remote_engine_id is returned by the prefill's request_finished.
                     # host/ports come from the request_id (parsed in add_new_req).
                     if "remote_engine_id" in params:
+                        adjusted_remote_block_ids = remote_block_ids
                         if num_external_tokens > 0:
-                            # Get unhashed blocks to pull from remote.
-                            local_block_ids = blocks.get_block_ids()[0]
-                            assert len(local_block_ids) <= len(remote_block_ids)
-                            if len(local_block_ids) != len(remote_block_ids):
-                                local_block_ids = remote_block_ids[
-                                    -len(local_block_ids) :
+                            if self._has_mamba:
+                                # READ is synchronous, so allocate_slots may hash
+                                # destination blocks before this hook runs. Pair
+                                # the attention and recurrent-state suffixes
+                                # independently.
+                                remote_attn = list(remote_block_ids[0])
+                                remote_mamba_groups = [
+                                    list(group) for group in remote_block_ids[1:]
+                                ]
+                                attn_block_ids, mamba_block_groups = (
+                                    self.split_block_groups(blocks.get_block_ids())
+                                )
+                                if len(mamba_block_groups) != len(remote_mamba_groups):
+                                    raise MoRIIOError(
+                                        "MoRIIO hybrid READ Mamba group "
+                                        "count mismatch: "
+                                        f"local={len(mamba_block_groups)}, "
+                                        f"remote={len(remote_mamba_groups)}"
+                                    )
+                                local_attn = attn_block_ids
+                                local_attn, remote_attn = self._align_read_blocks(
+                                    local_attn,
+                                    remote_attn,
+                                    self._max_decode_tail_blocks,
+                                )
+                                aligned_local_mamba: list[list[int]] = []
+                                aligned_remote_mamba: list[list[int]] = []
+                                for local_group, remote_group in zip(
+                                    mamba_block_groups,
+                                    remote_mamba_groups,
+                                    strict=True,
+                                ):
+                                    aligned_local, aligned_remote = (
+                                        self._align_read_blocks(
+                                            local_group, remote_group, 1
+                                        )
+                                    )
+                                    aligned_local_mamba.append(aligned_local)
+                                    aligned_remote_mamba.append(aligned_remote)
+                                local_block_ids = [local_attn, *aligned_local_mamba]
+                                adjusted_remote_block_ids = [
+                                    remote_attn,
+                                    *aligned_remote_mamba,
                                 ]
+                            else:
+                                # Preserve main's per-group HMA representation
+                                # for attention-only hybrid models while
+                                # accounting for the locally recomputed tail.
+                                local_block_ids = self.get_exchange_clipped_blocks(
+                                    blocks.get_block_ids()
+                                )
+                                adjusted_remote_block_ids = []
+                                aligned_local_block_ids: list[list[int]] = []
+                                for local_group, remote_group in zip(
+                                    local_block_ids,
+                                    remote_block_ids,
+                                    strict=True,
+                                ):
+                                    aligned_local, aligned_remote = (
+                                        self._align_read_blocks(
+                                            local_group,
+                                            remote_group,
+                                            self._max_decode_tail_blocks,
+                                        )
+                                    )
+                                    aligned_local_block_ids.append(aligned_local)
+                                    adjusted_remote_block_ids.append(aligned_remote)
+                                local_block_ids = aligned_local_block_ids
                         else:
-                            # If remote_blocks and num_external_tokens = 0, we have
-                            # a full prefix cache hit on the D worker. We need to call
-                            # send_notify in _read_blocks to free the memory on the P.
-                            local_block_ids = []
+                            if self._has_mamba:
+                                # Attention can be a complete local hit, but the
+                                # recurrent state is never prefix-cacheable.
+                                _, mamba_block_groups = self.split_block_groups(
+                                    blocks.get_block_ids()
+                                )
+                                local_block_ids = [[], *mamba_block_groups]
+                            else:
+                                # Nothing to pull, but retain one empty list per
+                                # cache group so the worker can notify P.
+                                local_block_ids = [[] for _ in remote_block_ids]
 
                         self._reqs_need_recv[request.request_id] = (
                             request,
@@ -675,7 +1013,10 @@ def update_state_after_alloc(
                         )
                         # Snapshot params for chunked prefill consumption
                         # in build_connector_meta (see comment above).
-                        self._req_kv_params[request.request_id] = dict(params)
+                        saved_params = dict(params)
+                        if num_external_tokens > 0:
+                            saved_params["remote_block_ids"] = adjusted_remote_block_ids
+                        self._req_kv_params[request.request_id] = saved_params
                     else:
                         logger.warning(
                             "Got invalid KVTransferParams: %s. This "
@@ -818,8 +1159,6 @@ def build_connector_meta(
                 new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i]
 
                 if new_block_ids is not None:
-                    block_ids = new_block_ids[0]
-                    # TODO : hybrid attn, etc
                     # A non-disagg request (no kv_transfer_params, e.g. smoke
                     # test) is never registered in _reqs_need_pending_save;
                     # indexing it unconditionally would KeyError and crash the
@@ -827,12 +1166,16 @@ def build_connector_meta(
                     if req_id not in self._reqs_need_pending_save:
                         continue
                     req, existing_blocks = self._reqs_need_pending_save[req_id]
-                    updated_blocks = list(existing_blocks) + (block_ids)
+                    updated_blocks = [
+                        existing_blocks[g] + new_block_ids[g]
+                        for g in range(len(new_block_ids))
+                    ]
                     self._reqs_need_pending_save[req_id] = (req, updated_blocks)
-                    if (
-                        len(self._reqs_need_pending_save[req_id][1]) * self.block_size
-                        >= req.num_prompt_tokens
-                    ):
+                    saved_tokens = (
+                        len(updated_blocks[self._full_attn_group_idx])
+                        * self._full_attn_block_size
+                    )
+                    if saved_tokens >= req.num_prompt_tokens:
                         # Final chunk: live kv_transfer_params may be cleared,
                         # so prefer the snapshot from update_state_after_alloc.
                         kv_params = self._req_kv_params.pop(
@@ -857,7 +1200,10 @@ def build_connector_meta(
 
         for req_id, (req, block_ids) in self._reqs_need_save.items():
             kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})
-            if req.num_prompt_tokens > len(block_ids) * self.block_size:
+            tokens_covered = (
+                len(block_ids[self._full_attn_group_idx]) * self._full_attn_block_size
+            )
+            if req.num_prompt_tokens > tokens_covered:
                 # not last chunk prefill
                 self._reqs_need_pending_save[req_id] = (req, block_ids)
                 continue
@@ -895,16 +1241,79 @@ def shutdown(self):
                 logger.warning("Error closing ZMQ socket for path %s: %s", path, e)
         self.paths.clear()
 
+    def split_block_groups(
+        self, block_ids: list[list[int]] | tuple[list[int], ...]
+    ) -> tuple[list[int], list[list[int]]]:
+        """Select transferable attention and Mamba block groups.
+
+        The wire payload stores the attention group first, followed by every
+        Mamba group in transfer-group order. Outside ``mamba_cache_mode="all"``,
+        only each group's running-state slot is transferred.
+        """
+        if not block_ids:
+            return [], []
+        if not self._has_mamba:
+            raise MoRIIOError("Mamba block groups requested for a non-hybrid model")
+        transfer_block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids)
+        attn = list(transfer_block_ids[self._attn_group_ids[0]])
+        mamba_groups = [
+            self._clip_mamba_group(list(transfer_block_ids[group_id]))
+            for group_id in self._mamba_group_ids
+        ]
+        return attn, mamba_groups
+
+    def _clip_mamba_group(self, blocks: list[int]) -> list[int]:
+        """Keep only the state-bearing slots of one mamba kv cache group."""
+        if not blocks:
+            return blocks
+        if not self._ssm_state_slots_are_positional:
+            # Single running state: everything before it is a null placeholder
+            # or the previous step's superseded state.
+            blocks = blocks[-1:]
+        return blocks
+
+    @staticmethod
+    def _align_read_blocks(
+        local_block_ids: list[int],
+        remote_block_ids: list[int],
+        max_decode_tail_blocks: int,
+    ) -> tuple[list[int], list[int]]:
+        """Pair a decode suffix with the corresponding prefill blocks.
+
+        Prefix hits can make the decode's unhashed suffix shorter than the
+        prefill's full list, so read from the matching remote suffix. The
+        decoder can instead have one extra tail block for the locally
+        recomputed final token; that block has no remote source and is not
+        transferred.
+        """
+        if len(local_block_ids) > len(remote_block_ids):
+            extra_decode_blocks = len(local_block_ids) - len(remote_block_ids)
+            if extra_decode_blocks > max_decode_tail_blocks:
+                raise ValueError(
+                    "decode block ids exceed prefill block ids beyond the "
+                    f"allowed local tail ({max_decode_tail_blocks} blocks): "
+                    f"{len(local_block_ids)} > {len(remote_block_ids)}"
+                )
+            local_block_ids = local_block_ids[: len(remote_block_ids)]
+        elif len(local_block_ids) < len(remote_block_ids):
+            remote_block_ids = (
+                remote_block_ids[-len(local_block_ids) :] if local_block_ids else []
+            )
+        return local_block_ids, remote_block_ids
+
     def request_finished(
         self,
         request: "Request",
-        block_ids: list[int],
+        block_ids: BlockIds | list[int],
+        mamba_block_groups: list[list[int]] | None = None,
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Once a request is finished, determine whether request blocks
+        """Once a request is finished, determine whether request blocks
         should be freed now or will be sent asynchronously and freed later.
-        """
 
+        ``block_ids`` normally contains all cache groups. Hybrid callers from
+        ``request_finished_all_groups`` pass a flat attention list and the
+        recurrent-state groups separately in ``mamba_block_groups``.
+        """
         request_id = request.request_id
         params = request.kv_transfer_params
         # Consumer: can unmap transfer_id<->request_id immediately since done_recving
@@ -936,7 +1345,25 @@ def request_finished(
             if self.mode == MoRIIOMode.WRITE:
                 self._release_write_prefill_blocks(request.request_id, params)
             else:
-                self._reqs_need_recv[request.request_id] = (request, [])
+                # Carry the actually-allocated blocks (incl. the per-request KDA
+                # mamba slot) instead of []: a do_remote_prefill request finishing
+                # before update_state_after_alloc still has blocks, and dropping the
+                # mamba slot here would re-trip the KDA guard in _read_blocks.
+                if mamba_block_groups:
+                    recv_blocks: BlockIds = [
+                        list(cast(list[int], block_ids)),
+                        *[list(group) for group in mamba_block_groups],
+                    ]
+                elif block_ids and isinstance(block_ids[0], int):
+                    recv_blocks = [list(cast(list[int], block_ids))]
+                else:
+                    grouped_block_ids = cast(BlockIds, block_ids)
+                    recv_blocks = (
+                        [list(group) for group in grouped_block_ids]
+                        if self._has_mamba
+                        else [[] for _ in self.blocks_per_sw]
+                    )
+                self._reqs_need_recv[request.request_id] = (request, recv_blocks)
             params["do_remote_prefill"] = False
             return False, None
 
@@ -946,10 +1373,19 @@ def request_finished(
         ):
             return False, None
 
-        # computed_block_ids = block_ids if all_full else block_ids[:-1]
-        computed_block_ids = block_ids
-        # If prompt < block_size, no xfer so free blocks immediately.
-        delay_free_blocks = len(computed_block_ids) > 0
+        # On producer
+        if mamba_block_groups is not None:
+            computed_block_ids: BlockIds = [
+                list(cast(list[int], block_ids)),
+                *[list(group) for group in mamba_block_groups],
+            ]
+        elif self._has_mamba:
+            raise MoRIIOError("Hybrid request completion requires all KV cache groups")
+        else:
+            computed_block_ids = self.get_exchange_clipped_blocks(
+                cast(BlockIds, block_ids)
+            )
+        delay_free_blocks = any(len(group) > 0 for group in computed_block_ids)
 
         if delay_free_blocks:
             # Prefill request on remote. It will be read from D upon completion
@@ -1070,7 +1506,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
 
 
 class MoRIIOConnectorWorker:
-    """Implementation of Worker side methods"""
+    """Implementation of Worker side methods."""
 
     def __init__(
         self,
@@ -1102,6 +1538,16 @@ def __init__(
         self.kv_transfer_config = vllm_config.kv_transfer_config
         self.is_producer = self.kv_transfer_config.is_kv_producer
         self.layer_to_spec = build_layer_to_spec(kv_cache_config)
+        # we transfer by layer, but store by kv cache group
+        self.layer_to_group = dict(kv_cache_config.transfer_group_index_by_layer)
+        self._transfer_layer_names = set(self.layer_to_group)
+        _, mamba_transfer_group_ids = _split_kv_cache_group_kinds(kv_cache_config)
+        self._num_mamba_transfer_groups = len(mamba_transfer_group_ids)
+        self._mamba_payload_index_by_layer = {
+            layer_name: payload_index
+            for payload_index, group_id in enumerate(mamba_transfer_group_ids, start=1)
+            for layer_name in kv_cache_config.transfer_groups[group_id].layer_names
+        }
 
         if self.is_producer:
             set_role(ROLE.PRODUCER)
@@ -1241,7 +1687,8 @@ def __init__(
         # Map of engine_id -> num_blocks. All ranks in the same deployment will
         # have the same number of blocks.
         self.dst_num_blocks: dict[EngineId, int] = {}
-        # In-progress READ transfers: req_id -> {layer_name: status}.
+        # In-progress READ transfers: req_id -> {layer_name: [statuses]}.
+        # A KDA layer posts two reads (conv + ssm), so a layer maps to a list.
         self._recving_transfers: defaultdict[ReqId, dict] = defaultdict(dict)
         # Values are (remote_host, remote_notify_port, transfer_id).
         self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str, str]] = {}
@@ -1249,6 +1696,18 @@ def __init__(
         # Used by _pop_done_transfers to abort transfers whose RDMA
         # completion is lost, instead of hanging forever.
         self._recving_transfers_start: dict[str, float] = {}
+        # Destination attention blocks per in-flight recv. Failed non-HMA reads
+        # report these to the scheduler's single-group recovery path.
+        self._recving_local_blocks: dict[ReqId, list[int]] = {}
+        # Drained by get_block_ids_with_load_errors.
+        self._invalid_block_ids: set[int] = set()
+        # This step's posted reads, awaited before the forward when the
+        # per-layer barrier cannot run.
+        self._reads_issued_this_step: list = []
+        # This step's posted recurrent-state reads. Mamba layers are not
+        # wrapped by the attention KV-transfer decorator, so wait_for_layer_load
+        # never fires for them and these must be awaited in start_load_kv.
+        self._mamba_reads_this_step: list = []
 
         # Track the expiration time of requests that are waiting to be sent.
         self._reqs_to_send: dict[ReqId, float] = {}
@@ -1274,6 +1733,14 @@ def __init__(
 
         self.block_window_per_layer: list[int | None] = []
         self.use_mla = self.model_config.use_mla
+        # Hybrid (mamba/KDA) recurrent-state transfer. Populated in
+        # register_kv_caches when the model has a MambaSpec kv_cache group.
+        self._has_mamba = False
+        self._conv_decomp: MambaConvSplitInfo | None = None
+        self._mamba_offset_templates: dict[str, MambaOffsetTemplate] = {}
+        # layer_name -> flat session indices (one per registered region:
+        # 1 for attention, 2 (conv, ssm) for a KDA layer). Built once lazily.
+        self._region_session_index: dict[str, list[int]] | None = None
         self.built_session = False
         self.built_write_session: defaultdict[str, list] = defaultdict(list)
         backend = get_attn_backend(
@@ -1322,8 +1789,8 @@ def schedule_write_blocks(
             kv_layer: KV cache tensor
             remote_notify_port: Port for completion notification
             remote_ip: IP address of remote node
-        """
 
+        """
         # synchronization to prevent dirty reads between
         # transfer and attention operations
         # we can consider removing this synchronization after ibgda is enabled.
@@ -1348,23 +1815,39 @@ def schedule_write_blocks(
 
     def _get_built_session(self, remote_engine_id):
         if remote_engine_id not in self.built_write_session:
+            remote_layer_map = self.layer_name_to_remote_kv_cache_metadata[
+                remote_engine_id
+            ]
             cur_remote_engine_sessions = []
-            for ln, local_meta in self.layer_name_to_local_kv_cache_metadata.items():
-                unpacked_local_memory_meta = (
-                    self.moriio_wrapper.get_unpack_memory_metadata(local_meta[0])
-                )
-                unpacked_remote_memory_meta = (
-                    self.moriio_wrapper.get_unpack_memory_metadata(
-                        self.layer_name_to_remote_kv_cache_metadata[remote_engine_id][
-                            ln
-                        ][0]
+            for ln, local_metas in self.layer_name_to_local_kv_cache_metadata.items():
+                remote_metas = remote_layer_map[ln]
+                if len(remote_metas) != len(local_metas):
+                    # Sessions are addressed by region index, so a peer that
+                    # registered a different number of regions for this layer
+                    # would not just lose a region to zip(): every later index
+                    # would point into the wrong layer's session.
+                    raise MoRIIOError(
+                        f"peer {remote_engine_id} registered "
+                        f"{len(remote_metas)} region(s) for layer {ln}, this "
+                        f"engine has {len(local_metas)}; the two sides must run "
+                        "the same model and KV cache layout"
                     )
-                )
-                cur_remote_engine_sessions.append(
-                    self.moriio_wrapper.build_session(
-                        unpacked_local_memory_meta, unpacked_remote_memory_meta
+                # One session per registered region, in registration order:
+                # attention layers have a single region; KDA layers have two
+                # (conv then ssm). This flat layout is indexed via
+                # _region_session_indices(layer_name).
+                for local_meta, remote_meta in zip(local_metas, remote_metas):
+                    unpacked_local_memory_meta = (
+                        self.moriio_wrapper.get_unpack_memory_metadata(local_meta)
+                    )
+                    unpacked_remote_memory_meta = (
+                        self.moriio_wrapper.get_unpack_memory_metadata(remote_meta)
+                    )
+                    cur_remote_engine_sessions.append(
+                        self.moriio_wrapper.build_session(
+                            unpacked_local_memory_meta, unpacked_remote_memory_meta
+                        )
                     )
-                )
             self.built_write_session[remote_engine_id] = cur_remote_engine_sessions
         return self.built_write_session[remote_engine_id], self.remote_moriio_metadata[
             remote_engine_id
@@ -1468,7 +1951,6 @@ def _moriio_handshake_listener(
         layer_name_to_local_kv_cache_metadata: dict,
     ):
         """Background thread for getting new MoRIIO handshakes."""
-
         encoder = msgspec.msgpack.Encoder()
         encoded_data = encoder.encode(metadata)
         size_in_bytes = len(encoded_data)
@@ -1524,7 +2006,6 @@ def _moriio_handshake(
         local-rank mapping _remote_tp_rank -- byte-identical for callers not yet
         TP-aware.
         """
-
         start_time = time.perf_counter()
 
         # NOTE(rob): we need each rank to have a unique port. This is
@@ -1678,15 +2159,60 @@ def wait_all_dp():
     def _is_mla_cache_layer(self, layer_name: str) -> bool:
         return is_mla_cache_layer(self.layer_to_spec, layer_name)
 
+    def _is_mamba_layer(self, layer_name: str) -> bool:
+        """True for a hybrid/KDA layer backed by a packed state page."""
+        return isinstance(self.layer_to_spec.get(layer_name), MambaSpec)
+
+    @staticmethod
+    def _contiguous_byte_alias(tensor: torch.Tensor) -> torch.Tensor:
+        """Return a contiguous uint8 view over ``tensor``'s full byte extent.
+
+        KDA conv/ssm caches are slot-strided views (``stride(0)`` may exceed the
+        per-slot element count), so they are non-contiguous and cannot be passed
+        to ``register_torch_tensor``. This aliases the same storage (zero-copy,
+        identical ``data_ptr``) as a flat uint8 tensor spanning
+        ``shape[0] * stride(0)`` elements -- the full slot-strided extent that
+        transfer offsets (``slot * stride(0) * elem``) address.
+        """
+        esz = tensor.element_size()
+        storage_nbytes = tensor.untyped_storage().nbytes()
+        offset_bytes = tensor.storage_offset() * esz
+        extent_bytes = tensor.shape[0] * tensor.stride(0) * esz
+        # Clamp to the bytes remaining from the view's storage offset. In a
+        # packed page, the ssm view starts after conv, so its slot-strided extent
+        # would otherwise run past the backing allocation.
+        extent_bytes = min(extent_bytes, storage_nbytes - offset_bytes)
+        assert extent_bytes > 0, (
+            "KDA byte alias empty: storage_offset="
+            f"{tensor.storage_offset()} * esz={esz} >= "
+            f"storage_nbytes={storage_nbytes}"
+        )
+        alias = torch.empty(0, dtype=torch.uint8, device=tensor.device)
+        alias.set_(
+            tensor.untyped_storage(),
+            tensor.storage_offset() * esz,
+            (extent_bytes,),
+            (1,),
+        )
+        return alias
+
     def _get_layer_transfer_geometry(
         self, layer_name: str, remote_num_blocks: int | None = None
     ) -> LayerTransferGeometry:
-        return get_layer_transfer_geometry(
+        geometry = get_layer_transfer_geometry(
             layer_name,
             self.kv_caches[layer_name],
             self.layer_to_spec,
             remote_num_blocks,
         )
+        # get_layer_transfer_geometry also serves mamba layers, whose geometry
+        # is a different NamedTuple with no block fields. Every caller here is
+        # on the attention path (mamba layers return before reaching it), so
+        # narrow the type instead of leaking the union.
+        assert isinstance(geometry, LayerTransferGeometry), (
+            f"layer {layer_name} is a mamba layer; use get_mamba_transfer_geometry"
+        )
+        return geometry
 
     def _iter_layer_registration_regions(
         self, layer_name: str
@@ -1743,23 +2269,71 @@ def _remote_layer_mr_offset(
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in moriio."""
-
+        kv_caches = {
+            layer_name: kv_cache
+            for layer_name, kv_cache in kv_caches.items()
+            if layer_name in self._transfer_layer_names
+        }
+        if not kv_caches:
+            raise MoRIIOError("MoRIIO has no transfer-enabled KV cache layers")
         self.kv_caches = kv_caches  # layer name to kv cache
+        self._has_mamba = any(
+            self._is_mamba_layer(layer_name) for layer_name in kv_caches
+        )
+        if self._has_mamba and self.mode == MoRIIOMode.WRITE:
+            # WRITE needs a distinct conv/ssm transfer plan. Reject before
+            # allocating or registering any regions rather than silently
+            # serving a zero recurrent state.
+            raise MoRIIOError(
+                "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ "
+                "mode only; set kv_connector_extra_config.read_mode=true "
+                "for hybrid models"
+            )
+        if self._has_mamba:
+            mamba_specs = [
+                cast(MambaSpec, self.layer_to_spec[layer_name])
+                for layer_name in kv_caches
+                if self._is_mamba_layer(layer_name)
+            ]
+            mamba_spec = _validate_mamba_specs(mamba_specs)
+            assert mamba_spec is not None
+
+            from vllm.model_executor.layers.mamba.mamba_utils import (
+                is_conv_state_dim_first,
+            )
+
+            if not is_conv_state_dim_first():
+                raise MoRIIOError(
+                    "MoRIIO KDA conv transfer requires the DS (dim-first) "
+                    "state layout; set VLLM_SSM_CONV_STATE_LAYOUT=DS"
+                )
+            self._conv_decomp = derive_mamba_conv_split(mamba_spec, self.world_size)
+        # KDA layers use packed pages with different geometry; only record
+        # attention tensor shapes here.
         self.kv_cache_shapes = {
-            layer_name: kv_cache.shape for layer_name, kv_cache in kv_caches.items()
+            layer_name: kv_cache.shape
+            for layer_name, kv_cache in kv_caches.items()
+            if not self._is_mamba_layer(layer_name)
         }
 
+        # Geometry selection must pick an attention layer (never a KDA tuple):
+        # prefer a standard 5-D K/V layer, else the first attention layer.
+        attn_items = [
+            (layer_name, kv_cache)
+            for layer_name, kv_cache in kv_caches.items()
+            if not self._is_mamba_layer(layer_name)
+        ]
         first_layer_name, first_kv_cache = next(
             (
                 (layer_name, kv_cache)
-                for layer_name, kv_cache in kv_caches.items()
+                for layer_name, kv_cache in attn_items
                 if (
                     not self._is_mla_cache_layer(layer_name)
                     and len(kv_cache.shape) == 5
                     and (kv_cache.shape[0] == 2 or kv_cache.shape[1] == 2)
                 )
             ),
-            next(iter(kv_caches.items())),
+            attn_items[0] if attn_items else next(iter(kv_caches.items())),
         )
         kv_elem_size = first_kv_cache.element_size()
 
@@ -1789,8 +2363,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 self.block_size,
             )
             self.block_size = first_geometry.block_size
-        # TODO(tms): self.block_len needs to be per-layer for sliding window,
-        # hybrid attn, etc
         # block size in bytes
         self.block_len = first_geometry.block_len
         self.kv_cache_shape = first_kv_cache.shape
@@ -1803,14 +2375,41 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         self.layer_base_addr_index = {}
         base_addr_idx = 0
 
+        hma_enabled = (
+            not self.vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
+        )
         for layer_name in kv_caches:
             self.layer_base_addr_index[layer_name] = base_addr_idx
+            if self._is_mamba_layer(layer_name):
+                # KDA layer: two whole-tensor regions (conv, ssm); no block-size
+                # geometry (recurrent state is per-slot, not paged).
+                for cache, region_len in self._iter_layer_registration_regions(
+                    layer_name
+                ):
+                    base_addr = cache.data_ptr()
+                    caches_data.append((base_addr, region_len, cache.device.index, ""))
+                    kv_caches_base_addr.append(base_addr)
+                    base_addr_idx += 1
+                continue
             geometry = self._get_layer_transfer_geometry(layer_name)
-            if geometry.block_size != self.block_size:
+            if hma_enabled:
+                if geometry.block_len != self.block_len:
+                    raise ValueError(
+                        "MoRIIO KV cache block length mismatch for layer "
+                        f"{layer_name}: {geometry.block_len} != {self.block_len}"
+                    )
+            elif geometry.block_size != self.block_size:
                 raise ValueError(
                     "MoRIIO KV cache block size mismatch for layer "
                     f"{layer_name}: {geometry.block_size} != {self.block_size}"
                 )
+            # num_blocks is advertised as a single scalar to the peer, so must it
+            # be uniform
+            if geometry.num_blocks != self.num_blocks:
+                raise ValueError(
+                    "MoRIIO KV cache num_blocks mismatch for layer "
+                    f"{layer_name}: {geometry.num_blocks} != {self.num_blocks}"
+                )
             self.block_lens[layer_name] = geometry.block_len
             for cache, region_len in self._iter_layer_registration_regions(layer_name):
                 base_addr = cache.data_ptr()
@@ -1818,31 +2417,87 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 kv_caches_base_addr.append(base_addr)
                 base_addr_idx += 1
 
-        shared_backing = len({id(t.untyped_storage()) for t in kv_caches.values()}) == 1
+        shared_backing = (
+            all(isinstance(t, torch.Tensor) for t in kv_caches.values())
+            and len({id(t.untyped_storage()) for t in kv_caches.values()}) == 1
+        )
         shared_mr = None
         if shared_backing:
             reg_tensor, self.kv_layer_mr_offset = self._build_shared_kv_mr(kv_caches)
             shared_mr = self.moriio_wrapper.register_local_tensor(reg_tensor)
 
+        # Say it out loud when the groups disagree. This is expected on a hybrid
+        # model and used to be fatal, so a run that silently had one page size
+        # and a run that has three should not look identical in the log.
+        distinct_block_lens = sorted(set(self.block_lens.values()))
+        if len(distinct_block_lens) > 1:
+            logger.info(
+                "MoRIIO registered %d attention layers with %d distinct block "
+                "lengths (bytes): %s; transfer sizes and strides are per layer, "
+                "and %d is the page advertised to the peer.",
+                len(self.block_lens),
+                len(distinct_block_lens),
+                distinct_block_lens,
+                self.block_len,
+            )
+
         for layer_name, kv_cache in kv_caches.items():
             if layer_name not in self.layer_name_to_local_kv_cache_metadata:
                 self.layer_name_to_local_kv_cache_metadata[layer_name] = []
 
-            moriio_mem_metadata = shared_mr or (
-                self.moriio_wrapper.register_local_tensor(kv_cache)
-            )
-            self.layer_name_to_local_kv_cache_metadata[layer_name].append(
-                moriio_mem_metadata
-            )
-
-            self.local_kv_cache_size.append(
-                kv_cache.nelement() * kv_cache.element_size()
-            )
+            if self._is_mamba_layer(layer_name):
+                # Register conv and ssm as two whole-tensor regions, appending
+                # both metadata blobs (order: conv then ssm) so the session
+                # layout matches iter_layer_registration_regions. conv/ssm are
+                # slot-strided views and are not contiguous, which
+                # register_torch_tensor rejects; register a contiguous uint8
+                # alias over each tensor's full byte extent instead (zero-copy:
+                # same storage and data_ptr).
+                conv, ssm = kda_conv_ssm(
+                    kv_cache, cast(MambaSpec, self.layer_to_spec[layer_name])
+                )
+                logger.debug(
+                    "MoRIIO KDA register %s: conv shape=%s stride=%s contig=%s "
+                    "storage_nbytes=%s ; ssm shape=%s stride=%s contig=%s "
+                    "storage_nbytes=%s",
+                    layer_name,
+                    tuple(conv.shape),
+                    tuple(conv.stride()),
+                    conv.is_contiguous(),
+                    conv.untyped_storage().nbytes(),
+                    tuple(ssm.shape),
+                    tuple(ssm.stride()),
+                    ssm.is_contiguous(),
+                    ssm.untyped_storage().nbytes(),
+                )
+                for tensor in (conv, ssm):
+                    alias = self._contiguous_byte_alias(tensor)
+                    meta = self.moriio_wrapper.register_local_tensor(alias)
+                    self.layer_name_to_local_kv_cache_metadata[layer_name].append(meta)
+                    self.local_kv_cache_size.append(alias.numel())
+            else:
+                moriio_mem_metadata = shared_mr or (
+                    self.moriio_wrapper.register_local_tensor(kv_cache)
+                )
+                self.layer_name_to_local_kv_cache_metadata[layer_name].append(
+                    moriio_mem_metadata
+                )
+                self.local_kv_cache_size.append(
+                    kv_cache.nelement() * kv_cache.element_size()
+                )
 
         self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
         self.num_regions = len(caches_data)
         self.num_layers = len(self.kv_caches.keys())
 
+        # Report the prevalidated KDA decomposition after registration.
+        if self._has_mamba:
+            assert self._conv_decomp is not None
+            logger.info(
+                "MoRIIO registered %d KDA (conv+ssm) layer(s); ssm sizes=%s",
+                sum(1 for ln in kv_caches if self._is_mamba_layer(ln)),
+                self._conv_decomp.ssm_sizes,
+            )
         # Optimization for models with local attention (Llama 4)
         if self.vllm_config.model_config.hf_config.model_type == "llama4":
             from transformers import Llama4TextConfig
@@ -1893,12 +2548,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         self.moriio_wrapper.async_wait_reqid()
 
     def get_finished(self) -> tuple[set[str], set[str]]:
-        """
-        Get requests that are done sending or recving on this specific worker.
+        """Get requests that are done sending or recving on this specific worker.
         The scheduler process (via the MultiprocExecutor) will use this output
         to track which workers are done.
         """
-
         done_sending, done_recving = set(), set()
 
         if self.is_producer:
@@ -1980,49 +2633,49 @@ def wait_for_layer_load(self, layer_name: str) -> None:
         if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL:
             return
 
-        deadline = time.monotonic() + self.moriio_config.transfer_timeout
-        while True:
-            with self.moriio_wrapper.lock:
-                pending = [
-                    status_by_layer[layer_name]
-                    for status_by_layer in self._recving_transfers.values()
-                    if layer_name in status_by_layer
-                ]
-
-            if not pending:
-                return
-
-            still_running = False
-            for status in pending:
-                # A failed read is dropped in _pop_done_transfers.
-                if status.Succeeded() or status.Failed():
-                    continue
-                still_running = True
+        with self.moriio_wrapper.lock:
+            pending = [
+                status
+                for status_by_layer in self._recving_transfers.values()
+                for status in status_by_layer.get(layer_name, ())
+            ]
 
-            if not still_running:
-                return
+        if not pending:
+            return
 
-            if time.monotonic() > deadline:
-                logger.warning(
-                    "MoRIIO READ barrier timed out for layer %s; proceeding "
-                    "(request dropped via get_finished).",
+        try:
+            self.moriio_wrapper.waiting_for_transfer_complete(pending)
+        except TransferError:
+            if self._has_mamba:
+                logger.exception(
+                    "MoRIIO hybrid READ failed before layer %s; aborting the "
+                    "forward because HMA load-error recovery is unsupported.",
                     layer_name,
                 )
-                return
-
-            time.sleep(0.001)
+                raise
+            logger.warning(
+                "MoRIIO READ barrier did not complete for layer %s; proceeding "
+                "to transfer cleanup.",
+                layer_name,
+                exc_info=True,
+            )
 
     def _pop_done_transfers(self) -> set[str]:
         done_req_ids: set[str] = set()
-        _xfer_timeout = int(os.environ.get("VLLM_MORIIO_TRANSFER_TIMEOUT_S", "120"))
+        _xfer_timeout = self.moriio_config.recv_abort_timeout
         with self.moriio_wrapper.lock:
             to_remove = []
             for req_id, status_by_layer in self._recving_transfers.items():
-                statuses = list(status_by_layer.values())
-                failed_status = next(
-                    (status for status in statuses if status.Failed()), None
-                )
-                if statuses and all(status.Succeeded() for status in statuses):
+                # Every layer's read must land before the request is complete,
+                # so the whole batch is polled rather than just the newest
+                # status.
+                statuses = [
+                    status
+                    for layer_statuses in status_by_layer.values()
+                    for status in layer_statuses
+                ]
+                state = self.moriio_wrapper.poll_transfer_batch(statuses)
+                if statuses and state is TransferBatchState.DONE:
                     host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
                     done_req_ids.add(xfer_id)
                     self.moriio_wrapper.send_notify(
@@ -2033,15 +2686,19 @@ def _pop_done_transfers(self) -> set[str]:
                         message_fields={"consumer_tp_size": self.world_size},
                     )
                     to_remove.append(req_id)
-                elif failed_status is not None:
+                elif state is TransferBatchState.FAILED:
+                    failed_status = next(
+                        (status for status in statuses if status.Failed()), None
+                    )
                     logger.error(
                         "RDMA transfer failed for request %s: %s (code=%s). "
-                        "Notifying prefill to free blocks; request will be "
-                        "aborted by timeout.",
+                        "Notifying prefill to free blocks and recording the "
+                        "load error where scheduler recovery is supported.",
                         req_id,
-                        failed_status.Message(),
-                        failed_status.Code(),
+                        failed_status.Message() if failed_status else "unknown",
+                        failed_status.Code() if failed_status else "unknown",
                     )
+                    self._record_failed_recv(req_id)
                     host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
                     try:
                         self.moriio_wrapper.send_notify(
@@ -2057,8 +2714,9 @@ def _pop_done_transfers(self) -> set[str]:
                             req_id,
                         )
                     to_remove.append(req_id)
-                    # Do NOT add to done_req_ids: decode KV cache is incomplete.
-                    # The request will expire via the normal request timeout.
+                    # Deliberately not in done_req_ids: the decode KV is
+                    # incomplete, so this is a load error rather than a
+                    # completion.
                 elif req_id in self._recving_transfers_start:
                     # Abort still-in-flight transfers that exceed the
                     # configured deadline. Otherwise a lost RDMA
@@ -2068,19 +2726,84 @@ def _pop_done_transfers(self) -> set[str]:
                     if _age > _xfer_timeout:
                         logger.error(
                             "RDMA read TIMED OUT for req %s after %.1fs "
-                            "(VLLM_MORIIO_TRANSFER_TIMEOUT_S=%d)",
+                            "(kv_connector_extra_config.recv_abort_timeout=%.0f)",
                             req_id,
                             _age,
                             _xfer_timeout,
                         )
+                        self._record_failed_recv(req_id)
                         to_remove.append(req_id)
             for req_id in to_remove:
                 del self._recving_transfers[req_id]
                 del self._recving_transfers_callback_addr[req_id]
                 self._recving_transfers_start.pop(req_id, None)
+                self._recving_local_blocks.pop(req_id, None)
 
             return done_req_ids
 
+    def _await_reads_issued_this_step(self) -> None:
+        """Step-level READ barrier, for the reads no per-layer barrier covers.
+
+        load_kv_async=False promises the KV is in place before the forward runs.
+        wait_for_layer_load keeps that promise per layer, letting early layers
+        overlap later transfers, but it only fires for layers that go through
+        the attention KV-transfer decorator. Two cases fall through it:
+
+        * mamba/KDA layers have no such hook at all, so their conv+ssm reads
+          would otherwise be consumed while still in flight;
+        * under CUDAGraphMode.FULL a host-side blocking wait is illegal
+          anywhere in the graph, so every read is awaited here instead.
+
+        start_load_kv runs outside the graph, so blocking here is safe.
+        Attention reads keep their per-layer overlap in the non-FULL case.
+
+        Attention-only failures are reported via
+        get_block_ids_with_load_errors. Hybrid failures are re-raised because
+        the scheduler's recovery path does not support multiple cache groups;
+        this fails the step closed instead of running forward with incomplete
+        recurrent state.
+        """
+        all_statuses = self._reads_issued_this_step
+        self._reads_issued_this_step = []
+        mamba_statuses = self._mamba_reads_this_step
+        self._mamba_reads_this_step = []
+        if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL:
+            statuses = all_statuses
+        else:
+            statuses = mamba_statuses
+        if not statuses:
+            return
+        try:
+            self.moriio_wrapper.waiting_for_transfer_complete(statuses)
+        except TransferError:
+            if self._has_mamba:
+                logger.exception(
+                    "MoRIIO hybrid READ failed before the forward; aborting "
+                    "because HMA load-error recovery is unsupported."
+                )
+                raise
+            logger.exception(
+                "MoRIIO reads did not complete before the forward; proceeding "
+                "to transfer cleanup."
+            )
+
+    def _record_failed_recv(self, req_id: ReqId) -> None:
+        """Hand a failed attention-only read to scheduler recovery.
+
+        The scheduler's invalid-block recovery currently assumes exactly one
+        KV cache group. Reporting IDs for a hybrid request would crash it while
+        unpacking the request's groups, so match NIXL and keep HMA failures out
+        of this API until scheduler recovery carries group identity.
+        """
+        if self._has_mamba:
+            return
+        self._invalid_block_ids.update(self._recving_local_blocks.get(req_id, ()))
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """Drain the blocks whose read failed, for the scheduler to recompute."""
+        invalid, self._invalid_block_ids = self._invalid_block_ids, set()
+        return invalid
+
     def save_kv_layer(
         self,
         metadata: MoRIIOConnectorMetadata,
@@ -2289,8 +3012,7 @@ def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> No
             self._eager_handshaked_engines.add(remote_engine_id)
 
     def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
-        """
-        Start loading by triggering non-blocking moriio_xfer.
+        """Start loading by triggering non-blocking moriio_xfer.
         We check for these trnxs to complete in each step().
         """
         self.transfer_id_to_request_id = metadata.transfer_id_to_request_id
@@ -2342,11 +3064,12 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
 
                         continue
 
-            # Handshake already completed, start async read xfer.
+            # Handshake already completed, start the read xfer.
             self._read_blocks_for_req(req_id, meta)
         # Start transfers for requests whose handshakes have now finished.
 
         if remote_engine_id is None and not wait_handshake_readd_req:
+            self._await_reads_issued_this_step()
             return
         _deadline = time.monotonic() + self.moriio_config.transfer_timeout
         while True:
@@ -2373,6 +3096,7 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
             else:
                 break
 
+        self._await_reads_issued_this_step()
         self._reqs_to_send.update(metadata.reqs_to_send)
 
     def wait_for_save(self, metadata: MoRIIOConnectorMetadata):
@@ -2455,12 +3179,16 @@ def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_lay
             self.remote_dp_size_local = int(meta.remote_dp_size_local)
         else:
             self.remote_dp_size_local = int(meta.remote_dp_size)
+        # WRITE does not support HMA, so unwrap blocks into flat lists.
+        # remote_block_ids can itself be empty so need to be careful when unwrapping.
+        local_block_ids = meta.local_block_ids[0]
+        remote_block_ids = meta.remote_block_ids[0] if meta.remote_block_ids else []
         self.schedule_write_blocks(
             request_id=req_id,
             transfer_id=meta.transfer_id,
             dst_engine_id=meta.remote_engine_id,
-            local_block_ids=meta.local_block_ids,
-            remote_block_ids=meta.remote_block_ids,
+            local_block_ids=local_block_ids,
+            remote_block_ids=remote_block_ids,
             layer_name=layer_name,
             kv_layer=kv_layer,
             remote_notify_port=meta.remote_notify_port,
@@ -2555,8 +3283,13 @@ def _compute_block_transfer_offsets(
             local_block_ids: IDs of local blocks
             remote_block_ids: IDs of remote blocks
             remote_moriio_meta: Metadata of the remote MoRIIO agent
+            remote_tp_size: Tensor parallel size of the remote engine. Falls
+                back to the locally recorded value when None.
+            remote_engine_id: Id of the remote engine the blocks live on.
+
         Returns:
             Tuple of (local_offsets, remote_offsets, transfer_sizes)
+
         """
         validate_moriio_heterogeneous_tp_kv_heads(
             local_tp_size=self.world_size,
@@ -2588,6 +3321,163 @@ def _compute_block_transfer_offsets(
             remote = [x + remote_off for x in remote]
         return local, remote, sizes
 
+    def _region_session_indices(self, layer_name: str) -> list[int]:
+        """Flat session indices for a layer's registered regions.
+
+        Sessions are built one per region in registration order (see
+        _get_built_session), so an attention layer maps to a single index and a
+        KDA layer to two indices ([conv, ssm]). Built once and cached.
+        """
+        if self._region_session_index is None:
+            mapping: dict[str, list[int]] = {}
+            idx = 0
+            for ln, metas in self.layer_name_to_local_kv_cache_metadata.items():
+                mapping[ln] = list(range(idx, idx + len(metas)))
+                idx += len(metas)
+            self._region_session_index = mapping
+        return self._region_session_index[layer_name]
+
+    def _mamba_tp_ratio(self, remote_tp_size: int | None) -> int:
+        """Validate the homogeneous-TP KDA contract."""
+        remote = self.world_size if remote_tp_size is None else remote_tp_size
+        if remote != self.world_size:
+            raise NotImplementedError(
+                "heterogeneous-TP mamba/KDA transfer is not supported yet "
+                f"(local_tp={self.world_size}, remote_tp={remote}); use equal "
+                "prefill/decode TP for hybrid models"
+            )
+        return 1
+
+    def _mamba_blocks_for_layer(
+        self,
+        layer_name: str,
+        block_groups: BlockIds,
+    ) -> list[int]:
+        payload_index = self._mamba_payload_index_by_layer.get(layer_name)
+        if payload_index is None or payload_index >= len(block_groups):
+            raise MoRIIOError(
+                f"KDA layer {layer_name} has no block group in the transfer payload"
+            )
+        return list(block_groups[payload_index])
+
+    def _compute_mamba_transfer_offsets(
+        self,
+        layer_name: str,
+        local_slots: list[int],
+        remote_slots: list[int],
+        remote_tp_size: int | None = None,
+    ) -> tuple[list[int], list[int], list[int], int]:
+        """Compute (local, remote, sizes, n_conv) for a KDA layer.
+
+        Entries [:n_conv] address the conv region/session; entries [n_conv:]
+        address the ssm region/session.
+        """
+        assert self._conv_decomp is not None, "KDA decomposition not derived"
+        conv, ssm = kda_conv_ssm(
+            self.kv_caches[layer_name],
+            cast(MambaSpec, self.layer_to_spec[layer_name]),
+        )
+        tp_ratio = self._mamba_tp_ratio(remote_tp_size)
+        # Cache the slot-independent geometry per layer; only the per-request
+        # slot bases vary, and heterogeneous TP is rejected above.
+        template = self._mamba_offset_templates.get(layer_name)
+        if template is None:
+            template = build_mamba_offset_template(
+                conv,
+                ssm,
+                self._conv_decomp,
+                tp_ratio,
+            )
+            self._mamba_offset_templates[layer_name] = template
+        local_offs, remote_offs, sizes = apply_mamba_offset_template(
+            template, local_slots, remote_slots
+        )
+        n_conv = compute_mamba_conv_split_count(local_slots, self._conv_decomp)
+        return local_offs, remote_offs, sizes, n_conv
+
+    def _post_mamba_reads(
+        self,
+        layer_name: str,
+        sessions: list,
+        region_sessions: list[int],
+        local_slots: list[int],
+        remote_slots: list[int],
+        remote_tp_size: int,
+        request_id: str,
+        deadline: float,
+    ) -> list:
+        """Post one KDA layer's conv and ssm reads; return their statuses.
+
+        The layer owns two registered regions, so the offset list is split at
+        the conv/ssm boundary and each half goes to its own session. A half
+        with nothing to move is skipped rather than posted empty.
+        """
+        local, remote, sizes, n_conv = self._compute_mamba_transfer_offsets(
+            layer_name,
+            local_slots,
+            remote_slots,
+            remote_tp_size=remote_tp_size,
+        )
+        statuses: list = []
+        for sess_idx, sl in (
+            (region_sessions[0], slice(0, n_conv)),
+            (region_sessions[1], slice(n_conv, None)),
+        ):
+            region_sizes = sizes[sl]
+            if not region_sizes:
+                continue
+            statuses.append(
+                self._post_read_with_backoff(
+                    sessions[sess_idx],
+                    region_sizes,
+                    local[sl],
+                    remote[sl],
+                    request_id,
+                    layer_name,
+                    deadline,
+                )
+            )
+        return statuses
+
+    def _post_read_with_backoff(
+        self,
+        session,
+        sizes: list[int],
+        local_offsets: list[int],
+        remote_offsets: list[int],
+        request_id: str,
+        layer_name: str,
+        deadline: float,
+    ):
+        """Post one RDMA READ, backing off on transient SQ-full rejection.
+
+        read_remote_data posts synchronously, so a send-queue-full rejection is
+        a Failed() status on return; a separate CQ-poll thread drains
+        completions and frees SQ depth, so we back off and re-post until
+        transfer_timeout, then store the failed status (get_finished notifies
+        prefill and drops the request non-fatally).
+        """
+        _backoff = _SQ_FULL_BACKOFF_INITIAL_S
+        while True:
+            transfer_status = self.moriio_wrapper.read_remote_data(
+                sizes, local_offsets, remote_offsets, session
+            )
+            if not self._is_sq_full_status(transfer_status):
+                break
+            if time.monotonic() > deadline:
+                logger.warning(
+                    "MoRIIO READ send queue stayed full past "
+                    "transfer_timeout for req %s layer %s; storing failed "
+                    "status (get_finished notifies prefill and drops the "
+                    "request). Raise qp_per_transfer if frequent.",
+                    request_id,
+                    layer_name,
+                )
+                break
+            time.sleep(_backoff)
+            _backoff = min(_backoff * 2, _SQ_FULL_BACKOFF_MAX_S)
+        return transfer_status
+
     @staticmethod
     def _is_sq_full_status(status) -> bool:
         """True if a MoRIIO transfer status is a transient RDMA send-queue-full
@@ -2607,8 +3497,8 @@ def _is_sq_full_status(status) -> bool:
 
     def _read_blocks(
         self,
-        local_block_ids: list[int],
-        remote_block_ids: list[int],
+        local_block_ids: BlockIds,
+        remote_block_ids: BlockIds,
         dst_engine_id: str,
         request_id: str,
         transfer_id: str,
@@ -2637,6 +3527,44 @@ def _read_blocks(
             if chosen_tp is not None
             else self._remote_tp_rank(remote_tp_size)
         )
+        if not any(local_block_ids):
+            # A request aborted before decode allocation still carries the
+            # producer's block table. There is no local destination or forward,
+            # so release the producer blocks without building sessions or
+            # posting RDMA reads.
+            self.moriio_wrapper.send_notify(
+                transfer_id,
+                remote_host,
+                str(
+                    remote_notify_port
+                    + get_port_offset(
+                        int(remote_dp_rank),
+                        eff_tp,
+                        remote_tp_size,
+                    )
+                ),
+                message_type="release",
+                message_fields={"consumer_tp_size": self.world_size},
+            )
+            return
+
+        if self._has_mamba:
+            expected_groups = 1 + self._num_mamba_transfer_groups
+            if (
+                len(local_block_ids) != expected_groups
+                or len(remote_block_ids) != expected_groups
+            ):
+                raise MoRIIOError(
+                    "MoRIIO hybrid READ block group count mismatch: "
+                    f"expected={expected_groups}, local={len(local_block_ids)}, "
+                    f"remote={len(remote_block_ids)}"
+                )
+            local_attn = list(local_block_ids[0])
+            remote_attn = list(remote_block_ids[0])
+            self._mamba_tp_ratio(remote_tp_size)
+        else:
+            local_attn = remote_attn = []
+
         if flexible:
             remote_dp_engine_id = self.get_engine_name_with_dp_tp(
                 dst_engine_id, int(remote_dp_rank), eff_tp
@@ -2649,51 +3577,89 @@ def _read_blocks(
 
         # SQ-full backpressure deadline, shared across this request's layers.
         _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout
+        # TODO: apply multi-session batch-read when MoRIIO supports it.
         for layer_name in self.layer_name_to_local_kv_cache_metadata:
-            sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index(
-                layer_name
-            )
-            offs = self._compute_block_transfer_offsets(
-                layer_name,
-                local_block_ids,
-                remote_block_ids,
-                remote_moriio_meta,
-                remote_tp_size=remote_tp_size,
-                remote_engine_id=remote_dp_engine_id,
-            )
-            # TODO : apply multi-session batch-read when moriio support it
-            #
-            # SQ-full backpressure: read_remote_data posts the RDMA READ
-            # SYNCHRONOUSLY, so a send-queue-full rejection (per-QP HW cap) comes
-            # back as a Failed() status right here. A SEPARATE CQ-poll thread
-            # drains completions and frees SQ depth, so back off and RE-POST
-            # rather than let a transient rejection abort the request. No
-            # self-deadlock (the drain is off-thread); the reserve is
-            # all-or-nothing (nothing posted on a rejected attempt). Bounded by
-            # transfer_timeout; on sustained overload store the failed status and
-            # let get_finished handle it non-fatally (notify prefill + drop).
-            _backoff = 0.001
-            while True:
-                transfer_status = self.moriio_wrapper.read_remote_data(
-                    offs[2], offs[0], offs[1], sessions[sess_idx]
+            region_sessions = self._region_session_indices(layer_name)
+            statuses = []
+            if self._is_mamba_layer(layer_name):
+                local_mamba = self._mamba_blocks_for_layer(layer_name, local_block_ids)
+                remote_mamba = self._mamba_blocks_for_layer(
+                    layer_name, remote_block_ids
                 )
-                if not self._is_sq_full_status(transfer_status):
-                    break
-                if time.monotonic() > _sq_deadline:
-                    logger.warning(
-                        "MoRIIO READ send queue stayed full past "
-                        "transfer_timeout for req %s layer %s; storing failed "
-                        "status (get_finished notifies prefill and drops the "
-                        "request). Raise qp_per_transfer if frequent.",
+                # KDA layer: pull conv (sub-projection offsets) and ssm at the
+                # request's recurrent-state slot. Two regions -> two sessions.
+                if not local_mamba or not remote_mamba:
+                    if not local_attn and not remote_attn:
+                        # Whole-request no-op (e.g. full prefix cache hit):
+                        # nothing to transfer for any layer.
+                        continue
+                    # A KDA layer must have its recurrent-state slot in the
+                    # block-id tuple; missing it (while attention blocks are
+                    # present) means the mamba KV-cache group is absent -> bug.
+                    raise MoRIIOError(
+                        f"KDA layer {layer_name}: missing mamba recurrent-state "
+                        f"block ids (local={local_mamba}, remote={remote_mamba}) "
+                        f"with attention blocks present for request {request_id}; "
+                        "the mamba KV-cache group is absent from the transfer "
+                        "block-id tuple"
+                    )
+                statuses = self._post_mamba_reads(
+                    layer_name,
+                    sessions,
+                    region_sessions,
+                    local_mamba,
+                    remote_mamba,
+                    remote_tp_size,
+                    request_id,
+                    _sq_deadline,
+                )
+            else:
+                if self._has_mamba:
+                    local_layer_blocks = local_attn
+                    remote_layer_blocks = remote_attn
+                else:
+                    group_idx = self.layer_to_group[layer_name]
+                    local_layer_blocks = local_block_ids[group_idx]
+                    remote_layer_blocks = remote_block_ids[group_idx]
+                offs = self._compute_block_transfer_offsets(
+                    layer_name,
+                    local_layer_blocks,
+                    remote_layer_blocks,
+                    remote_moriio_meta,
+                    remote_tp_size=remote_tp_size,
+                    remote_engine_id=remote_dp_engine_id,
+                )
+                statuses.append(
+                    self._post_read_with_backoff(
+                        sessions[region_sessions[0]],
+                        offs[2],
+                        offs[0],
+                        offs[1],
                         request_id,
                         layer_name,
+                        _sq_deadline,
                     )
-                    break
-                time.sleep(_backoff)
-                _backoff = min(_backoff * 2, 0.05)
+                )
             with self.moriio_wrapper.lock:
-                self._recving_transfers[request_id][layer_name] = transfer_status
+                self._recving_transfers[request_id][layer_name] = statuses
                 self._recving_transfers_start.setdefault(request_id, time.monotonic())
+                # Destination blocks, kept for _record_failed_recv.
+                if self._has_mamba:
+                    failed_recv_blocks = local_attn
+                else:
+                    failed_recv_blocks = [
+                        block_id for group in local_block_ids for block_id in group
+                    ]
+                self._recving_local_blocks.setdefault(
+                    request_id, list(failed_recv_blocks)
+                )
+                # Awaited before the forward when the per-layer barrier cannot
+                # run (see _await_reads_issued_this_step).
+                self._reads_issued_this_step.extend(statuses)
+                if self._is_mamba_layer(layer_name):
+                    # No per-layer barrier exists for mamba layers, so these are
+                    # always awaited in start_load_kv.
+                    self._mamba_reads_this_step.extend(statuses)
                 self._recving_transfers_callback_addr[request_id] = (
                     remote_host,
                     str(
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
index fd8257a19c7b..a0c6ec1f6d7f 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
@@ -29,6 +29,7 @@
     MoRIIOError,
     MoRIIOTransferAck,
     RemoteAllocInfo,
+    TransferBatchState,
     TransferError,
     TransferId,
     WriteTask,
@@ -58,7 +59,6 @@
 except ImportError:
     logger.error("MoRIIO is not available")
 
-
 """Write task execution logic for MoRIIO connector."""
 
 
@@ -87,6 +87,7 @@ def __init__(self, worker: "MoRIIOConnectorWorker"):
 
         Args:
             worker: Reference to the parent worker
+
         """
         self._worker_ref: weakref_ref[MoRIIOConnectorWorker] = weakref_ref(worker)
         self._write_task_q: Queue[WriteTask] = Queue()
@@ -108,6 +109,7 @@ def worker(self) -> "MoRIIOConnectorWorker":
 
         Raises:
             RuntimeError: If worker has been garbage collected
+
         """
         worker = self._worker_ref()
         if worker is None:
@@ -131,6 +133,7 @@ def schedule_write(self, task: WriteTask) -> bool:
 
         Args:
             task: The write task to schedule
+
         """
         self.ensure_worker_started()
         if self._is_transfer_terminal(task.transfer_id):
@@ -178,7 +181,6 @@ def seal_pending_transfers(self) -> None:
 
     def _write_worker_loop(self) -> None:
         """Main loop for the write worker thread."""
-
         while True:
             # Process deferred tasks first
             self._process_deferred_tasks()
@@ -275,6 +277,7 @@ def _is_remote_ready(self, task: WriteTask) -> bool:
 
         Returns:
             True if remote blocks are ready
+
         """
         return (
             task.transfer_id in self.worker.moriio_wrapper.done_remote_allocate_req_dict
@@ -291,6 +294,7 @@ def _get_remote_alloc_info(self, transfer_id: str) -> RemoteAllocInfo:
 
         Raises:
             KeyError: If allocation info is missing
+
         """
         try:
             return self.worker.moriio_wrapper.done_remote_allocate_req_dict[transfer_id]
@@ -372,9 +376,11 @@ def _prepare_transfer_plan(
         Args:
             task: The write task
             request_info: Remote allocation information
+            remote_moriio_meta: Metadata of the remote MoRIIO agent
 
         Returns:
             The transfer plan
+
         """
         layer_cache = self.worker.kv_caches[task.layer_name]
         key = (task.layer_name, *_get_write_geometry_key(layer_cache))
@@ -389,9 +395,9 @@ def _prepare_transfer_plan(
             )
             request_info.transfer_offsets[key] = offsets
 
-        # Get session index
-        layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys())
-        sess_idx = layer_names.index(task.layer_name)
+        # One session per registered region; attention layers map to a single
+        # index (see MoRIIOConnectorWorker._region_session_indices).
+        sess_idx = self.worker._region_session_indices(task.layer_name)[0]
 
         local_off, remote_off, sizes = offsets
 
@@ -412,6 +418,7 @@ def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> list[Any]:
         Args:
             plan: The transfer plan
             sessions: List of transfer sessions
+
         """
         if plan.use_batch:
             return [
@@ -510,6 +517,7 @@ class MoRIIOWrapper:
         moriio_engine:  MoRIIO engine instance
         tp_rank: Tensor parallel rank
         dp_rank: Data parallel rank
+
     """
 
     def __init__(
@@ -654,6 +662,25 @@ def write_remote_data_single(
         )
         return transfer_status
 
+    def poll_transfer_batch(self, transfer_statuses: list[Any]) -> TransferBatchState:
+        """Non-blocking verdict over every status of one request.
+
+        Judging a request by its newest status alone calls it done while an
+        earlier layer is unfinished or already failed; failures do land on
+        layers other than the last, because _post_read_with_backoff returns a
+        failed status when the send queue never drains.
+
+        This stays a Python scan so it does not drive the backend's progress
+        callback from the engine thread alongside MoRIIO's own poller.
+        """
+        state = TransferBatchState.DONE
+        for status in transfer_statuses:
+            if status.Failed():
+                return TransferBatchState.FAILED
+            if not status.Succeeded():
+                state = TransferBatchState.PENDING
+        return state
+
     def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = None):
         if transfer_statuses is None:
             with self.lock:
@@ -665,6 +692,15 @@ def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = No
         if not transfers_to_wait:
             return
 
+        self._poll_transfers_until_done(transfers_to_wait)
+
+    def _poll_transfers_until_done(self, transfers_to_wait: list[Any]) -> None:
+        """Fallback for mori builds without the batched wait.
+
+        Gives failure precedence over pending: once any status has failed the
+        request is lost, so waiting out the remaining ones would only delay the
+        barrier by up to transfer_timeout.
+        """
         timeout = self._transfer_timeout
         deadline = time.monotonic() + timeout
         remaining = list(transfers_to_wait)
@@ -690,6 +726,8 @@ def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = No
                     )
                 else:
                     still_waiting.append(status)
+            if errors:
+                break
             remaining = still_waiting
             if remaining:
                 time.sleep(0.001)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py
index 0ee2721139e5..b00692dc5188 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py
@@ -2,14 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Callable, Mapping
-from typing import NamedTuple
+from typing import NamedTuple, cast
 
 import torch
 
+from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
+    MambaConvSplitInfo,
+)
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     KVCacheConfig,
     KVCacheSpec,
+    MambaSpec,
     MLAAttentionSpec,
     SlidingWindowMLASpec,
     UniformTypeKVCacheSpecs,
@@ -29,6 +33,85 @@ class LayerTransferGeometry(NamedTuple):
     split_kv_regions: bool
 
 
+class MambaTransferGeometry(NamedTuple):
+    """Geometry of a hybrid (mamba/KDA) layer's recurrent state.
+
+    A KDA layer's packed page is exposed as ``conv_state`` and ``ssm_state``
+    views. Each is slot-strided, so a slot's bytes live at
+    ``slot * slot_stride * element_size`` and span ``slot_bytes``. The two
+    tensors are registered as two separate MoRIIO regions per layer.
+    """
+
+    conv_slot_stride: int
+    ssm_slot_stride: int
+    ssm_slot_bytes: int
+    conv_element_size: int
+    ssm_element_size: int
+    conv_region_len: int
+    ssm_region_len: int
+
+
+def get_mamba_transfer_geometry(
+    conv: torch.Tensor,
+    ssm: torch.Tensor,
+) -> MambaTransferGeometry:
+    """Derive the slot-strided geometry of a KDA layer's conv + ssm state.
+
+    ``conv``/``ssm`` are zero-copy views over the layer's packed page. Slot byte
+    offset is ``slot * slot_stride * element_size`` and each
+    slot spans ``subtensor[0].numel() * element_size`` bytes; the whole tensor
+    is registered as one region (``*_region_len`` bytes).
+    """
+    if conv.ndim == 0 or ssm.ndim == 0 or conv.shape[0] == 0 or ssm.shape[0] == 0:
+        raise ValueError("KDA conv/ssm caches must contain at least one slot")
+    if conv.shape[0] != ssm.shape[0]:
+        raise ValueError(
+            f"KDA conv/ssm slot count mismatch: {conv.shape[0]} != {ssm.shape[0]}"
+        )
+    conv_slot_bytes = conv[0].numel() * conv.element_size()
+    ssm_slot_bytes = ssm[0].numel() * ssm.element_size()
+    conv_available_bytes = (
+        conv.untyped_storage().nbytes() - conv.storage_offset() * conv.element_size()
+    )
+    ssm_available_bytes = (
+        ssm.untyped_storage().nbytes() - ssm.storage_offset() * ssm.element_size()
+    )
+    if conv.shape[0]:
+        conv_required_bytes = (conv.shape[0] - 1) * conv.stride(
+            0
+        ) * conv.element_size() + conv_slot_bytes
+        ssm_required_bytes = (ssm.shape[0] - 1) * ssm.stride(
+            0
+        ) * ssm.element_size() + ssm_slot_bytes
+        if conv_required_bytes > conv_available_bytes:
+            raise ValueError("KDA conv slots exceed the registered storage extent")
+        if ssm_required_bytes > ssm_available_bytes:
+            raise ValueError("KDA ssm slots exceed the registered storage extent")
+
+    return MambaTransferGeometry(
+        conv_slot_stride=conv.stride(0),
+        ssm_slot_stride=ssm.stride(0),
+        ssm_slot_bytes=ssm_slot_bytes,
+        conv_element_size=conv.element_size(),
+        ssm_element_size=ssm.element_size(),
+        # Registered region must span the full slot-strided extent: stride(0)
+        # may exceed the per-slot element count, so numel() would under-size the
+        # region and truncate the highest-slot transfers. shape[0]*stride(0)
+        # covers every byte a ``slot * slot_stride * elem`` offset can address.
+        # Cap at the bytes remaining from each view's storage offset. A packed
+        # conv+ssm page places ssm at a non-zero intra-buffer offset,
+        # so do not register past the end of the shared backing allocation.
+        conv_region_len=min(
+            conv.shape[0] * conv.stride(0) * conv.element_size(),
+            conv_available_bytes,
+        ),
+        ssm_region_len=min(
+            ssm.shape[0] * ssm.stride(0) * ssm.element_size(),
+            ssm_available_bytes,
+        ),
+    )
+
+
 def build_layer_to_spec(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]:
     layer_to_spec: dict[str, KVCacheSpec] = {}
     for group in kv_cache_config.kv_cache_groups:
@@ -99,10 +182,10 @@ def get_layer_transfer_geometry(
     layer_to_spec: Mapping[str, KVCacheSpec],
     remote_num_blocks: int | None = None,
 ) -> LayerTransferGeometry:
+    spec = layer_to_spec[layer_name]
     shape = kv_cache.shape
     stride = kv_cache.stride()
     element_size = kv_cache.element_size()
-    spec = layer_to_spec[layer_name]
     is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
 
     if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
@@ -238,12 +321,60 @@ def get_layer_transfer_geometry(
     )
 
 
+def kda_conv_ssm(
+    kv_cache: torch.Tensor,
+    spec: MambaSpec,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return zero-copy conv and SSM views over a packed two-state page."""
+    from math import prod
+
+    from vllm.utils.torch_utils import get_dtype_size
+
+    if len(spec.shapes) != 2 or len(spec.dtypes) != 2:
+        raise ValueError(
+            "MoRIIO KDA transfer requires exactly two states (conv and SSM)"
+        )
+    if kv_cache.element_size() != 1:
+        raise ValueError("MoRIIO KDA packed cache must use a byte-sized dtype")
+    if kv_cache.ndim != 4 or kv_cache.shape[1:3] != (1, 1) or kv_cache.stride(3) != 1:
+        raise ValueError(
+            "MoRIIO KDA packed cache must have shape [B, 1, 1, C] with "
+            "a contiguous byte dimension"
+        )
+    # Mirror MambaBase.bind_kv_cache: blocks may be strided within a shared
+    # multi-layer allocation, while every individual byte page is contiguous.
+    pages = kv_cache.squeeze(dim=(1, 2))
+    states: list[torch.Tensor] = []
+    offset = 0
+    for shp, dt in zip(spec.shapes, spec.dtypes, strict=True):
+        nbytes = prod(shp) * get_dtype_size(dt)
+        if offset + nbytes > pages.shape[1]:
+            raise ValueError(
+                "MoRIIO KDA state layout exceeds the packed cache page size"
+            )
+        state = pages[:, offset : offset + nbytes].view(dt)
+        states.append(state.view(-1, *shp))
+        offset += nbytes
+    return states[0], states[1]
+
+
 def iter_layer_registration_regions(
     layer_name: str,
     kv_cache: torch.Tensor,
     layer_to_spec: Mapping[str, KVCacheSpec],
 ) -> list[tuple[torch.Tensor, int]]:
-    geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec)
+    spec = layer_to_spec[layer_name]
+    if isinstance(spec, MambaSpec):
+        # KDA layer: register the conv and ssm tensors as two whole-tensor
+        # regions (the conv 3-subprojection split is expressed as transfer
+        # offset triples, not extra registrations).
+        conv, ssm = kda_conv_ssm(kv_cache, spec)
+        geom = get_mamba_transfer_geometry(conv, ssm)
+        return [(conv, geom.conv_region_len), (ssm, geom.ssm_region_len)]
+    geometry = cast(
+        LayerTransferGeometry,
+        get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec),
+    )
     region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
     if geometry.regions_per_block == 1:
         # With padded or interleaved pages the block stride exceeds the
@@ -311,9 +442,36 @@ def compute_block_transfer_offsets(
             "local_block_ids longer than remote_block_ids: "
             f"{len(local_block_ids)} > {len(remote_block_ids)}"
         )
-    geometry = get_layer_transfer_geometry(
-        layer_name, kv_cache, layer_to_spec, remote_num_blocks
+    geometry = cast(
+        LayerTransferGeometry,
+        get_layer_transfer_geometry(
+            layer_name, kv_cache, layer_to_spec, remote_num_blocks
+        ),
+    )
+    invalid_local = next(
+        (
+            block_id
+            for block_id in local_block_ids
+            if not 0 <= block_id < geometry.num_blocks
+        ),
+        None,
+    )
+    if invalid_local is not None:
+        raise ValueError(
+            f"local block id {invalid_local} is outside [0, {geometry.num_blocks})"
+        )
+    invalid_remote = next(
+        (
+            block_id
+            for block_id in remote_block_ids
+            if not 0 <= block_id < remote_num_blocks
+        ),
+        None,
     )
+    if invalid_remote is not None:
+        raise ValueError(
+            f"remote block id {invalid_remote} is outside [0, {remote_num_blocks})"
+        )
     element_size = kv_cache.element_size()
     transfer_size_byte = geometry.block_len
     per_block = geometry.transfers_per_block
@@ -339,3 +497,125 @@ def compute_block_transfer_offsets(
             w += 1
 
     return merge_fn(offset_local, offset_remote, sizes)
+
+
+class MambaOffsetTemplate(NamedTuple):
+    """Slot-independent conv+ssm offset decomposition for a KDA layer.
+
+    Homogeneous-TP geometry (conv/ssm slot strides, conv sub-projection
+    offsets, ssm per-slot size) is identical across requests and across the
+    homogeneous GDN layers, so it is computed once and reused; only the
+    per-request slot bases vary (see ``apply_mamba_offset_template``).
+    """
+
+    num_slots: int
+    conv_slot_stride_bytes: int
+    ssm_slot_stride_bytes: int
+    ssm_read_bytes: int
+    conv_subprojs: tuple[tuple[int, int, int], ...]
+
+
+def build_mamba_offset_template(
+    conv: torch.Tensor,
+    ssm: torch.Tensor,
+    split_info: MambaConvSplitInfo,
+    tp_ratio: int,
+) -> MambaOffsetTemplate:
+    """Precompute the slot-independent conv+ssm offset decomposition.
+
+    Only homogeneous TP (``tp_ratio == 1``) is supported; heterogeneous TP is
+    gated with ``NotImplementedError`` (see the design doc). The returned
+    template is a pure function of the layer geometry and ``split_info``, so it
+    can be cached and reused across requests without recomputation.
+    """
+    if tp_ratio != 1:
+        # Heterogeneous-TP recurrent-state transfer needs the remote page's
+        # slot stride (which differs from the local page) and, for
+        # P_TP > D_TP, a multi-rank gather of conv/ssm slices. Both are
+        # follow-ups. The mamba path bypasses the attention KV-head validator,
+        # so fail loudly here instead of silently transferring wrong bytes.
+        raise NotImplementedError(
+            "heterogeneous-TP mamba/KDA transfer is not supported yet "
+            f"(tp_ratio={tp_ratio}); use equal prefill/decode TP for "
+            "hybrid models"
+        )
+    geom = get_mamba_transfer_geometry(conv, ssm)
+    conv_local_offsets = split_info.local_conv_offsets
+    conv_remote_offsets = split_info.remote_conv_offsets(0, 1)
+    ssm_read_bytes = geom.ssm_slot_bytes
+    conv_subprojs = tuple(
+        (loff, roff, rsz)
+        for (loff, _lsz), (roff, rsz) in zip(conv_local_offsets, conv_remote_offsets)
+    )
+    return MambaOffsetTemplate(
+        num_slots=conv.shape[0],
+        conv_slot_stride_bytes=geom.conv_slot_stride * geom.conv_element_size,
+        ssm_slot_stride_bytes=geom.ssm_slot_stride * geom.ssm_element_size,
+        ssm_read_bytes=ssm_read_bytes,
+        conv_subprojs=conv_subprojs,
+    )
+
+
+def apply_mamba_offset_template(
+    template: MambaOffsetTemplate,
+    local_slots: list[int],
+    remote_slots: list[int],
+) -> tuple[list[int], list[int], list[int]]:
+    """Apply per-request slot bases to a cached ``MambaOffsetTemplate``.
+
+    Conv sub-projection entries (all slots) come first, followed by one ssm
+    entry per slot.
+    """
+    if len(local_slots) != len(remote_slots):
+        raise ValueError(
+            "local and remote mamba slot counts must match: "
+            f"{len(local_slots)} != {len(remote_slots)}"
+        )
+    invalid_local = next(
+        (slot for slot in local_slots if not 0 <= slot < template.num_slots), None
+    )
+    if invalid_local is not None:
+        raise ValueError(
+            f"local mamba slot {invalid_local} is outside [0, {template.num_slots})"
+        )
+    invalid_remote = next((slot for slot in remote_slots if slot < 0), None)
+    if invalid_remote is not None:
+        raise ValueError(f"remote mamba slot {invalid_remote} must be non-negative")
+    # The peer's KDA slot count is not part of MoRIIOAgentMetadata yet. Do not
+    # compare remote slot IDs against the local pool size: two otherwise
+    # compatible engines may reserve different numbers of slots.
+    local_offs: list[int] = []
+    remote_offs: list[int] = []
+    sizes: list[int] = []
+
+    # Conv sub-projections (region 0), all slots first.
+    conv_slot_stride_bytes = template.conv_slot_stride_bytes
+    for ls, rs in zip(local_slots, remote_slots):
+        lbase = ls * conv_slot_stride_bytes
+        rbase = rs * conv_slot_stride_bytes
+        for loff, roff, rsz in template.conv_subprojs:
+            local_offs.append(lbase + loff)
+            remote_offs.append(rbase + roff)
+            sizes.append(rsz)
+
+    # SSM temporal state (region 1), one entry per slot.
+    ssm_slot_stride_bytes = template.ssm_slot_stride_bytes
+    ssm_read_bytes = template.ssm_read_bytes
+    for ls, rs in zip(local_slots, remote_slots):
+        local_offs.append(ls * ssm_slot_stride_bytes)
+        remote_offs.append(rs * ssm_slot_stride_bytes)
+        sizes.append(ssm_read_bytes)
+
+    return local_offs, remote_offs, sizes
+
+
+def compute_mamba_conv_split_count(
+    local_slots: list[int],
+    split_info: MambaConvSplitInfo,
+) -> int:
+    """Number of leading conv entries in a KDA layer transfer plan.
+
+    Entries ``[:count]`` are conv sub-projections (conv region/session);
+    entries ``[count:]`` are the per-slot ssm state (ssm region/session).
+    """
+    return len(local_slots) * len(split_info.local_conv_offsets)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
index 591ded1c6761..a6f0d4cafb75 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
@@ -17,6 +17,7 @@
     KVConnectorHandshakeMetadata,
     KVConnectorMetadata,
     KVConnectorRole,
+    KVConnectorTransferResults,
     KVConnectorWorkerMetadata,
     SupportsHMA,
 )
@@ -35,7 +36,7 @@
     from vllm.distributed.kv_events import KVCacheEvent
     from vllm.forward_context import ForwardContext
     from vllm.v1.core.block_pool import BlockPool
-    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
     from vllm.v1.kv_cache_interface import KVCacheConfig
     from vllm.v1.request import Request
 
@@ -70,8 +71,7 @@ def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetada
 
 @dataclass
 class MultiKVConnectorStats(KVConnectorStats):
-    """
-    Maintain a dict of KVConnectorStats objects, one for each connector.
+    """Maintain a dict of KVConnectorStats objects, one for each connector.
     This is used to aggregate the stats from all connectors separately.
     """
 
@@ -131,8 +131,7 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
 
 
 class MultiConnector(KVConnectorBase_V1, SupportsHMA):
-    """
-    A wrapper for using multiple KVConnectors at the same time.
+    """A wrapper for using multiple KVConnectors at the same time.
 
     The current logic is:
     - Load KV from the first connector that advertises available tokens from
@@ -142,8 +141,7 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
 
     @classmethod
     def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
-        """
-        MultiConnector requires PIECEWISE CUDA graph mode if any of its
+        """MultiConnector requires PIECEWISE CUDA graph mode if any of its
         child connectors require it.
         """
         connectors_config = extra_config.get("connectors", [])
@@ -208,6 +206,10 @@ def __init__(
         # Propagated from scheduler to worker side via the connector metadata.
         self._extra_async_saves: dict[str, int] = {}
 
+    @property
+    def sub_connectors(self) -> list[KVConnectorBase_V1]:
+        return list(self._connectors)
+
     @property
     def supports_divergent_local_hybrid_hits(self) -> bool:
         return bool(self._connectors) and all(
@@ -248,6 +250,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         for c in self._connectors:
             c.register_kv_caches(kv_caches)
 
+    def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None:
+        for connector in self._connectors:
+            connector.bind_kv_cache_manager(kv_cache_manager)
+
     def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
         for c in self._connectors:
             c.bind_gpu_block_pool(gpu_block_pool)
@@ -292,6 +298,14 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         for c in self._connectors:
             c.start_load_kv(forward_context, **kwargs)
 
+    def finish_forward(self) -> None:
+        for c in self._connectors:
+            c.finish_forward()
+
+    def reset_capture_state(self) -> None:
+        for c in self._connectors:
+            c.reset_capture_state()
+
     def wait_for_layer_load(self, layer_name: str) -> None:
         for c in self._connectors:
             c.wait_for_layer_load(layer_name)
@@ -313,29 +327,31 @@ def wait_for_save(self):
     def get_finished(
         self, finished_req_ids: set[str]
     ) -> tuple[set[str] | None, set[str] | None]:
-        finished_sending: set[str] = set()
-        finished_recving: set[str] = set()
-        for c in self._connectors:
-            sending, recving = c.get_finished(finished_req_ids)
-            if not recving and not sending:
-                continue
-            # Aggregate finished recving request ids.
-            finished_recving.update(recving or ())
-            # Aggregate finished sending request ids - only include
-            # once we've drained the "extra" count (for cases where
-            # more than one connector is async-saving the same request).
-            for req_id in sending or ():
+        results = self.get_transfer_results(finished_req_ids)
+        return (
+            results.finished_sending or None,
+            results.finished_recving or None,
+        )
+
+    def get_transfer_results(
+        self, finished_req_ids: set[str]
+    ) -> KVConnectorTransferResults:
+        results = KVConnectorTransferResults()
+        for connector in self._connectors:
+            child_results = connector.get_transfer_results(finished_req_ids)
+            results.finished_recving.update(child_results.finished_recving)
+            results.failed_recving.update(child_results.failed_recving)
+            for req_id in child_results.finished_sending:
                 extra_pending = self._extra_async_saves.get(req_id)
                 if extra_pending is None:
-                    finished_sending.add(req_id)
-                    continue
-                assert extra_pending > 0
-                if extra_pending == 1:
-                    del self._extra_async_saves[req_id]
+                    results.finished_sending.add(req_id)
                 else:
-                    self._extra_async_saves[req_id] = extra_pending - 1
-
-        return finished_sending or None, finished_recving or None
+                    assert extra_pending > 0
+                    if extra_pending == 1:
+                        del self._extra_async_saves[req_id]
+                    else:
+                        self._extra_async_saves[req_id] = extra_pending - 1
+        return results
 
     def get_block_ids_with_load_errors(self) -> set[int]:
         agg_block_ids: set[int] = set()
@@ -461,8 +477,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput):
             connector_output.kv_connector_worker_meta = multi_connector_worker_meta
 
     def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
-        """
-        Get the KVConnector handshake metadata from sub-connectors.
+        """Get the KVConnector handshake metadata from sub-connectors.
         Returns the first non-None metadata from sub-connectors.
         """
         for c in self._connectors:
@@ -474,8 +489,7 @@ def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
     def set_xfer_handshake_metadata(
         self, metadata: dict[int, KVConnectorHandshakeMetadata]
     ) -> None:
-        """
-        Set the KV connector handshake metadata for all sub-connectors.
+        """Set the KV connector handshake metadata for all sub-connectors.
         This is needed to start the NIXL listener thread for NixlConnector.
         """
         for c in self._connectors:
@@ -563,19 +577,23 @@ def take_events(self) -> Iterable["KVCacheEvent"]:
         for c in self._connectors:
             yield from c.take_events()
 
+    def has_pending_block_frees(self) -> bool:
+        return any(c.has_pending_block_frees() for c in self._connectors)
+
     def has_pending_push_work(self) -> bool:
         return any(c.has_pending_push_work() for c in self._connectors)
 
     @classmethod
     def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
-        """
-        Get the required KV cache layout for this connector.
+        """Get the required KV cache layout for this connector.
+
         Args:
             vllm_config (VllmConfig): the vllm config.
 
         Returns:
             str: the required KV cache layout. e.g. HND, or NHD.
             None if the connector does not require a specific layout.
+
         """
         assert vllm_config.kv_transfer_config is not None
         layouts: set[str] = set()
@@ -630,7 +648,7 @@ def build_kv_connector_stats(
 
         return MultiKVConnectorStats(data=reconstructed_data)
 
-    def get_kv_connector_stats(self) -> MultiKVConnectorStats | None:
+    def get_kv_connector_stats(self) -> KVConnectorStats | None:
         # Group connector stats by connector type.
         stats_by_connector: MultiKVConnectorStats | None = None
         for c in self._connectors:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
index ddfcf35e033b..e264a94aafd8 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
@@ -63,6 +63,12 @@ def __init__(
         kv_cache_config: "KVCacheConfig",
     ):
         self.vllm_config = vllm_config
+        parallel_config = vllm_config.parallel_config
+        # TP1 PCP+DCP exposes its DCP shards as transfer ranks.
+        self.transfer_tp_size = max(
+            parallel_config.tensor_parallel_size,
+            parallel_config.decode_context_parallel_size,
+        )
         self.block_size = vllm_config.cache_config.block_size
         self.engine_id: EngineId = engine_id
         self.kv_cache_config = kv_cache_config
@@ -93,10 +99,7 @@ def __init__(
                 for g in kv_cache_config.transfer_groups
             )
         )
-        self._has_mamba = any(
-            isinstance(g.kv_cache_spec, MambaSpec)
-            for g in kv_cache_config.transfer_groups
-        )
+        self._has_mamba = kv_cache_config.has_mamba_layers
 
         logger.info("Initializing NIXL Scheduler %s", engine_id)
         if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
@@ -110,7 +113,7 @@ def __init__(
         # New requests are added by update_state_after_alloc in
         # the scheduler. Used to make metadata passed to Worker.
         self._reqs_need_recv: dict[
-            ReqId, tuple[Request, BlockIds, tuple[int, ...]]
+            ReqId, tuple[Request, BlockIds, tuple[int, ...], bool]
         ] = {}
         self._reqs_need_save: dict[ReqId, Request] = {}
         # Reqs to send and their expiration time
@@ -196,8 +199,8 @@ def shutdown(self):
     def on_new_request(self, request: "Request") -> None:
         """Track a request that may need heartbeats."""
         params = request.kv_transfer_params
-        if params is not None and params.get("do_remote_decode") and self._has_mamba:
-            self._truncate_mamba_request_for_prefill(request)
+        if params is not None and params.get("do_remote_decode"):
+            self._truncate_request_for_prefill(request)
 
         # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are
         # effectively disabled for Bidirectional KV transfer.
@@ -296,11 +299,11 @@ def get_exchange_clipped_blocks(
     def set_xfer_handshake_metadata(
         self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
     ) -> None:
-        """
-        Set the KV connector handshake metadata for this connector.
+        """Set the KV connector handshake metadata for this connector.
 
         Args:
             metadata (dict): the handshake metadata to set.
+
         """
         encoded_data: dict[tuple[int, int], bytes] = {}
         encoder = msgspec.msgpack.Encoder()
@@ -379,36 +382,56 @@ def _nixl_handshake_listener(
                     (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)], ts)
                 )
 
+    def _prefill_backoff(self) -> int:
+        """Trailing prompt tokens the prefiller must not compute; the decoder
+        recomputes them locally.
+
+        Mamba needs h(N-1) so the decoder can derive h(N) itself. Multi-module
+        MTP needs to keep its whole lookahead window off of the prefiller, which
+        would otherwise embed the unverified drafts in the MTP layer's KV cache. The
+        decoder would never rebuild them, because the update is sized by the rejection
+        count, which is zero for the first decode.
+        """
+        return max(
+            1 if self._has_mamba else 0,
+            self.vllm_config.num_prefill_lookahead_tokens - 1,
+        )
+
     def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
-        """D-side only. Returns N-1 for Mamba models since the decoder
-        always recomputes the last token and must start from h(N-1)."""
-        if self._has_mamba and num_prompt_tokens > 1:
-            return num_prompt_tokens - 1
+        """D-side only. The number of prompt tokens to load from the prefiller.
+        Stops short of the trailing ``_prefill_backoff()`` tokens that the decoder
+        will recompute locally."""
+        backoff = self._prefill_backoff()
+        if backoff and num_prompt_tokens > backoff:
+            return num_prompt_tokens - backoff
         return num_prompt_tokens
 
-    def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
-        """P-side only: drop the last prompt token so the prefiller computes
-        h(N-1) instead of h(N). The decoder recomputes the last token to
-        derive h(N) correctly.
+    def _truncate_request_for_prefill(self, request: "Request") -> None:
+        """P-side only: drop the trailing ``_prefill_backoff()`` prompt tokens
+        so the prefiller stops short of what the decoder recomputes locally.
+        For Mamba that is the single token needed to yield h(N-1); for
+        multi-module MTP it is the drafter's whole lookahead window.
 
         Guarded by ``_p_side_truncated`` to avoid repeated truncation if the
         request is preempted and rescheduled."""
+        backoff = self._prefill_backoff()
         params = request.kv_transfer_params
         if (
-            params is not None
+            backoff
+            and params is not None
             # Guard against repeated truncation after preemption/reschedule.
             and not params.get("_p_side_truncated")
-            and request.num_prompt_tokens > 1
+            and request.num_prompt_tokens > backoff
         ):
             if request.prompt_token_ids is not None:
-                request.prompt_token_ids.pop()
+                del request.prompt_token_ids[-backoff:]
             elif request.prompt_embeds is not None:
-                request.prompt_embeds = request.prompt_embeds[:-1]
+                request.prompt_embeds = request.prompt_embeds[:-backoff]
             else:
                 return
 
-            request._all_token_ids.pop()
-            request.num_prompt_tokens -= 1
+            del request._all_token_ids[-backoff:]
+            request.num_prompt_tokens -= backoff
             request.max_tokens = 1
             params["_p_side_truncated"] = True
 
@@ -456,13 +479,19 @@ def build_connector_meta(
         meta = NixlConnectorMetadata()
 
         # Loop through scheduled reqs and convert to ReqMeta.
-        for req_id, (req, block_ids, cached) in self._reqs_need_recv.items():
+        for req_id, (
+            req,
+            block_ids,
+            cached,
+            awaiting_kvs,
+        ) in self._reqs_need_recv.items():
             assert req.kv_transfer_params is not None
             meta.add_new_req_to_recv(
                 request_id=req_id,
                 local_block_ids=block_ids,
                 kv_transfer_params=req.kv_transfer_params,
                 local_num_computed_blocks=cached,
+                awaiting_kvs=awaiting_kvs,
             )
 
         if self.use_host_buffer:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
index 5cb76932d44b..0f13a265dfd7 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
@@ -2,8 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Base worker-side logic for the NIXL connector."""
 
+import contextlib
 import itertools
 import logging
+import math
 import os
 import queue
 import threading
@@ -31,7 +33,10 @@
     kv_postprocess_blksize_on_receive,
     kv_postprocess_layout_on_receive,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    CopyBlocksOp,
+    KVConnectorTransferResults,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
     GET_META_MSG,
@@ -70,6 +75,7 @@
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.utils.math_utils import cdiv
 from vllm.utils.network_utils import make_zmq_path
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.kv_cache_interface import (
@@ -159,6 +165,10 @@ class NixlBaseConnectorWorker:
     # Overridden by NixlPushConnectorWorker.
     _TRANSFER_MODE: str = "pull"
 
+    # Layer-name routing is supported only by NixlPushConnector for HMA KV
+    # caches under PP.
+    _supports_pp_hma = False
+
     def _compute_desc_ids(
         self,
         block_ids: BlockIds,
@@ -186,23 +196,42 @@ def _compute_desc_ids(
         num_blocks = dst_num_blocks
         if block_size_ratio is not None:
             num_blocks = int(num_blocks * block_size_ratio)
+        use_layer_name_routing = bool(self._transfer_layer_group_ids)
+        num_desc_regions = (
+            len(self._transfer_layer_group_ids)
+            if use_layer_name_routing
+            else self.num_regions or len(self.block_len_per_layer)
+        )
         if region_num_blocks is None:
-            region_num_blocks = [num_blocks] * self.num_regions
+            region_num_blocks = [num_blocks] * num_desc_regions
         elif block_size_ratio is not None:
             region_num_blocks = [
                 int(count * block_size_ratio) for count in region_num_blocks
             ]
-        assert len(region_num_blocks) == self.num_regions
+        if region_num_blocks and not use_layer_name_routing:
+            num_desc_regions = len(region_num_blocks)
+        assert len(region_num_blocks) == num_desc_regions
         region_offsets = np.cumsum([0, *region_num_blocks[:-1]])
         num_fa_descs = sum(region_num_blocks)
 
-        # All-attention fast path: single vectorized broadcast.
+        if use_layer_name_routing:
+            # Descriptors in layer order, using each layer's KV-cache group.
+            return np.concatenate(
+                [
+                    np.asarray(block_ids[g], dtype=np.int32) + region_offsets[k]
+                    for k, g in enumerate(self._transfer_layer_group_ids)
+                ]
+            )
+
+        # All-attention fast path.
         if num_ssm_regions == 0:
             if region_group_ids is None:
                 region_group_ids = self.region_group_ids
                 if uses_region_group_mapping is None:
                     uses_region_group_mapping = self._uses_region_group_mapping
-            assert len(region_group_ids) == self.num_regions
+            if not region_group_ids and num_desc_regions == 1:
+                region_group_ids = [0]
+            assert len(region_group_ids) == num_desc_regions
             if uses_region_group_mapping is None:
                 uses_region_group_mapping = len(set(region_group_ids)) > 1
             region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32)
@@ -382,6 +411,103 @@ def _is_region_replicated(self, region_idx: int) -> bool:
         """
         return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx]
 
+    def _set_region_layers(self, region_layers: list[list[str]]) -> None:
+        layer_names = [layer_name for region in region_layers for layer_name in region]
+        assert len(layer_names) == len(set(layer_names)), (
+            "A KV cache layer spans multiple NIXL regions"
+        )
+        self.region_members: list[list[str]] = region_layers
+        if not self._requires_layer_name_routing():
+            return
+        group_by_layer = self.kv_cache_config.transfer_group_index_by_layer
+        ungrouped = [name for name in layer_names if name not in group_by_layer]
+        assert not ungrouped, f"KV cache layers outside any local group: {ungrouped}"
+        self._transfer_layer_names = tuple(layer_names)
+        self._transfer_layer_region_indices = tuple(
+            region_idx
+            for region_idx, region in enumerate(region_layers)
+            for _ in region
+        )
+        self._transfer_layer_group_ids = tuple(
+            group_by_layer[name] for name in layer_names
+        )
+
+    def _requires_layer_name_routing(self) -> bool:
+        """Whether PP push must match HMA layers by name, not region index."""
+        return (
+            self._supports_pp_hma
+            and self.pp_size > 1
+            and self._is_hma_required
+            and not self._has_mamba
+        )
+
+    def _align_remote_regions_by_layer(
+        self, nixl_agent_meta: NixlAgentMetadata
+    ) -> None:
+        """Select remote regions for this PP stage in local layer order."""
+        remote_region_layers = nixl_agent_meta.region_members
+        # Region-index routing here would silently transfer stale KV.
+        assert remote_region_layers, "Remote advertised no region_members"
+        assert (
+            len(nixl_agent_meta.kv_caches_base_addr)
+            == len(nixl_agent_meta.block_lens)
+            == len(nixl_agent_meta.block_strides)
+            == len(remote_region_layers)
+        ), "Remote region metadata lengths disagree"
+        assert all(
+            values is None or len(values) == len(remote_region_layers)
+            for values in (
+                nixl_agent_meta.region_num_blocks,
+                nixl_agent_meta.region_group_ids,
+                nixl_agent_meta.region_names,
+                nixl_agent_meta.region_mem_types,
+            )
+        ), "Remote region metadata lengths disagree"
+
+        remote_region_by_layer: dict[str, int] = {}
+        for region_idx, layer_names in enumerate(remote_region_layers):
+            for layer_name in layer_names:
+                assert layer_name not in remote_region_by_layer, (
+                    f"Remote advertised layer {layer_name!r} in multiple regions"
+                )
+                remote_region_by_layer[layer_name] = region_idx
+
+        missing = [
+            name
+            for name in self._transfer_layer_names
+            if name not in remote_region_by_layer
+        ]
+        assert not missing, f"Remote is missing locally owned layers: {missing}"
+
+        remote_regions = [
+            remote_region_by_layer[name] for name in self._transfer_layer_names
+        ]
+        nixl_agent_meta.kv_caches_base_addr = [
+            nixl_agent_meta.kv_caches_base_addr[i] for i in remote_regions
+        ]
+        nixl_agent_meta.block_lens = [
+            nixl_agent_meta.block_lens[i] for i in remote_regions
+        ]
+        nixl_agent_meta.block_strides = [
+            nixl_agent_meta.block_strides[i] for i in remote_regions
+        ]
+        if nixl_agent_meta.region_num_blocks is not None:
+            nixl_agent_meta.region_num_blocks = [
+                nixl_agent_meta.region_num_blocks[i] for i in remote_regions
+            ]
+        if nixl_agent_meta.region_group_ids is not None:
+            nixl_agent_meta.region_group_ids = [
+                nixl_agent_meta.region_group_ids[i] for i in remote_regions
+            ]
+        if nixl_agent_meta.region_mem_types is not None:
+            nixl_agent_meta.region_mem_types = [
+                nixl_agent_meta.region_mem_types[i] for i in remote_regions
+            ]
+        if nixl_agent_meta.region_names is not None:
+            nixl_agent_meta.region_names = list(self._transfer_layer_names)
+        # One layer per region keeps a second alignment pass a no-op.
+        nixl_agent_meta.region_members = [[name] for name in self._transfer_layer_names]
+
     def __init__(
         self,
         vllm_config: "VllmConfig",
@@ -397,8 +523,6 @@ def __init__(
 
         # Config.
         self.vllm_config = vllm_config
-        # mypy will complain on re-assignment otherwise.
-        self.block_size: int = cast(int, vllm_config.cache_config.block_size)
 
         if vllm_config.kv_transfer_config is None:
             raise ValueError("kv_transfer_config must be set for NixlConnector")
@@ -420,6 +544,17 @@ def __init__(
         )
 
         self.kv_cache_config = kv_cache_config
+        transfer_block_sizes = [
+            group.kv_cache_spec.block_size
+            for group in kv_cache_config.transfer_groups
+            if get_representative_spec_type(group.kv_cache_spec)
+            not in (MambaSpec, CircularBufferSpec)
+        ]
+        self.block_size = (
+            math.lcm(*transfer_block_sizes)
+            if transfer_block_sizes
+            else cast(int, vllm_config.cache_config.block_size)
+        )
         # Per-layer specs, unwrapping UniformTypeKVCacheSpecs group wrappers.
         self._layer_specs: dict[str, KVCacheSpec] = {}
         for group in kv_cache_config.transfer_groups:
@@ -460,11 +595,13 @@ def __init__(
                     for spec in iter_layer_specs(group.kv_cache_spec)
                 )
             ]
-            if len(ple_groups) != 1 or len(ple_groups[0][1].layer_names) != 1:
+            if len(ple_groups) > 1 or any(
+                len(group.layer_names) != 1 for _, group in ple_groups
+            ):
                 raise ValueError(
-                    "CSA-linear NIXL requires exactly one PLE cache owner."
+                    "CSA-linear NIXL requires at most one PLE cache owner."
                 )
-            self._ple_group_index = ple_groups[0][0]
+            self._ple_group_index = ple_groups[0][0] if ple_groups else None
 
         if self._has_mamba:
             assert self._is_hma_required
@@ -527,12 +664,15 @@ def __init__(
         self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
         self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0
 
-        # DCP support is scoped to MLA, with dcp_size in (1, tp_size): either fully
-        # replicated or fully sharded. A DCP rank is always derivable this way.
-        self.dcp_rank = self.tp_rank % self.dcp_size
-        if self._has_mamba and self.dcp_size > 1:
-            # Prefix-cache-aware DCP slicing isn't implemented for the Mamba group.
-            raise ValueError("DCP is not supported for hybrid MLA+Mamba models.")
+        # TP1 PCP+DCP owns distinct KV shards; replicated PCP uses rank zero.
+        self.pcp_dcp_sharded = self.pcp_size > 1 and self.dcp_size > 1
+        self.transfer_tp_size = (
+            self.pcp_size if self.pcp_dcp_sharded else self.world_size
+        )
+        self.transfer_tp_rank = self.pcp_rank if self.pcp_dcp_sharded else self.tp_rank
+
+        # MLA is either fully replicated or fully sharded across transfer ranks.
+        self.dcp_rank = self.transfer_tp_rank % self.dcp_size
 
         self.num_blocks = kv_cache_config.num_blocks
         self.enable_permute_local_kv = False
@@ -600,29 +740,40 @@ def __init__(
         # Number of NIXL regions. Currently one region per cache
         # (so 1 per layer for MLA, otherwise 2 per layer)
         self.num_regions = 0
+        self.region_mem_types: list[str] = []
         self.region_group_ids: list[int] = []
         self._uses_region_group_mapping = False
         self.region_names: list[str] = []
         self.region_num_blocks: list[int] = []
-        self.region_mem_types: list[str] = []
         self._mixed_mem_types = False
+        self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
+        self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
+        self._dram_src_handles_by_block_size: dict[int, int] = {}
 
         # PP>1 (push mode): this worker holds a contiguous layer slice and
         # transfers into the matching sub-range of a PP=1 remote's regions.
         self.pp_size = vllm_config.parallel_config.pipeline_parallel_size
         self._remote_region_offset = 0
-        # PP push slices regions per layer (uniform count); HMA breaks that.
-        if self.pp_size > 1 and self._is_hma_required:
-            raise NotImplementedError(
-                "NixlPushConnector does not support pipeline_parallel_size > 1 "
-                "with hybrid KV cache layouts (HMA) yet."
-            )
-        # Decode-side PP is unsupported (completions counted per consumer rank).
-        if vllm_config.kv_transfer_config.kv_role == "kv_consumer" and self.pp_size > 1:
-            raise NotImplementedError(
-                "NixlPushConnector consumer (decode) does not support "
-                "pipeline_parallel_size > 1."
-            )
+        if self.pp_size > 1:
+            if self._is_hma_required and not self._supports_pp_hma:
+                raise NotImplementedError(
+                    "NixlConnector (pull) does not support pipeline_parallel_size "
+                    "> 1 with hybrid KV cache layouts (HMA); use NixlPushConnector "
+                    "for PP + HMA."
+                )
+            # PP push routes HMA (hybrid) attention layouts by layer name;
+            # Mamba/SSM hybrids are not yet supported under PP.
+            if self._has_mamba:
+                raise NotImplementedError(
+                    "NixlPushConnector does not support pipeline_parallel_size > 1 "
+                    "with Mamba/SSM hybrid KV cache layouts yet."
+                )
+            # Decode-side PP is unsupported (completions counted per consumer rank).
+            if vllm_config.kv_transfer_config.kv_role == "kv_consumer":
+                raise NotImplementedError(
+                    "NixlPushConnector consumer (decode) does not support "
+                    "pipeline_parallel_size > 1."
+                )
         # Keep heartbeat handshakes to a PP-sharded producer notif-only.
         self._hb_handshake_notif_only = False
 
@@ -645,9 +796,6 @@ def __init__(
         self.dst_region_group_ids: dict[EngineId, list[int]] = {}
         self.dst_uses_region_group_mapping: dict[EngineId, bool] = {}
         self.dst_region_mem_types: dict[EngineId, list[str]] = {}
-        self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
-        self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
-        self._dram_src_handles_by_block_size: dict[int, int] = {}
         self._registered_descs: list[Any] = []
 
         # In progress transfers.
@@ -668,6 +816,7 @@ def __init__(
         # Uses Queue for thread-safe cross-thread coordination with the
         # background handshake thread, matching the _ready_requests pattern.
         self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
+        self._recv_failures: set[ReqId] = set()
         self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {}
 
         # Handshake metadata of this worker for NIXL transfers.
@@ -691,7 +840,6 @@ def __init__(
             "engine_ttl", 3600.0
         )
 
-        self.block_size = vllm_config.cache_config.block_size
         self.model_config = vllm_config.model_config
 
         self.use_mla = self.model_config.use_mla
@@ -749,6 +897,13 @@ def __init__(
         # within a block (BLHNC/BHLNC), where stride > block_len.
         self.block_stride_per_layer = list[int]()
 
+        # Region layer membership is populated for HMA layouts.
+        self.region_members = []
+        # Local layer order, set by ``_set_region_layers`` only when
+        # routing by layer name. Empty means region-index routing.
+        self._transfer_layer_names = tuple[str, ...]()
+        self._transfer_layer_region_indices = tuple[int, ...]()
+        self._transfer_layer_group_ids = tuple[int, ...]()
         # Per-engine TP mappings. Generated during handshake.
         self.tp_mappings: dict[EngineId, TPMapping] = {}
 
@@ -763,12 +918,16 @@ def _validate_remote_parallel_config(
         local_dcp_size = self.dcp_size
         remote_pcp_size = agent_metadata.pcp_size
         remote_dcp_size = agent_metadata.dcp_size
-        if (local_pcp_size > 1 and remote_dcp_size > 1) or (
-            remote_pcp_size > 1 and local_dcp_size > 1
+        if remote_pcp_size > 1 and remote_dcp_size not in (1, remote_pcp_size):
+            raise NotImplementedError(
+                "Remote NixlConnector PCP+DCP does not span the full PCP group. "
+                f"Remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}."
+            )
+        if (local_pcp_size > 1 and local_dcp_size == 1 and remote_dcp_size > 1) or (
+            remote_pcp_size > 1 and remote_dcp_size == 1 and local_dcp_size > 1
         ):
             raise NotImplementedError(
-                "NixlConnector PCP requires decode_context_parallel_size=1 "
-                "on both instances. "
+                "Replicated PCP cannot be paired with a DCP-sharded NIXL peer. "
                 f"Local PCP/DCP={local_pcp_size}/{local_dcp_size}; "
                 f"remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}."
             )
@@ -989,8 +1148,7 @@ def _add_notif_only_remote_agent(
         return self.nixl_wrapper.add_remote_agent(metadata.agent_metadata)
 
     def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> None:
-        """
-        Initialize transfer buffer in CPU mem for accelerators
+        """Initialize transfer buffer in CPU mem for accelerators
         NOT directly supported by NIXL (e.g., tpu)
         """
         xfer_buffers: dict[str, torch.Tensor] = {}
@@ -1104,8 +1262,7 @@ def _ensure_handshake(
         pp_size: int = 1,
         notif_agents_only: bool = False,
     ) -> Future[tuple[dict[tuple[int, int], str], float]] | None:
-        """
-        Ensure a handshake is in-flight (or already done) for *engine_id*.
+        """Ensure a handshake is in-flight (or already done) for *engine_id*.
 
         Returns the ``Future`` if a handshake is pending (or was just
         started), or ``None`` if the handshake already completed
@@ -1150,6 +1307,8 @@ def done_callback(
                             error=e,
                             remote_engine_id=eid,
                         )
+                        # Count once per handshake, regardless of waiting requests.
+                        self.xfer_stats.record_failed_handshake()
 
             fut.add_done_callback(done_callback)
             return fut
@@ -1184,16 +1343,15 @@ def request_ready(f: Future[Any], entry=(req_id, meta)):
                     error=e,
                     meta=meta,
                 )
-                self._handle_failed_transfer(req_id, None)
+                self._failed_recv_reqs.put(req_id)
 
         fut.add_done_callback(request_ready)
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in nixl."""
-
         self.transfer_topo = TransferTopology(
-            tp_rank=self.tp_rank,
-            tp_size=self.world_size,
+            tp_rank=self.transfer_tp_rank,
+            tp_size=self.transfer_tp_size,
             block_size=self.block_size,
             engine_id=self.engine_id,
             is_mla=self.use_mla,
@@ -1241,7 +1399,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             self.use_host_buffer,
         )
 
-        registration_ranges: dict[tuple[int, str], tuple[int, int, int, str]] = {}
+        registration_ranges: dict[tuple[int, str], tuple[int, int, int]] = {}
         region_mem_types: list[str] = []
         seen_base_addresses: list[int] = []
         self._ssm_region_indices = []
@@ -1278,13 +1436,24 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             ):
                 compressed_region_owners.setdefault(cache.data_ptr(), cache)
 
+        track_region_layers = (
+            self._supports_pp_hma and self._is_hma_required and not self._has_mamba
+        )
+        region_layers: list[list[str]] = []
+
         # K and V are packed into the content dim, so each attention layer is a
         # single NIXL region whose block transfers as one unit. Mamba layers instead
         # register separate conv/ssm sub-regions (see `_build_mamba_local`).
+        # P and D may allocate equivalent transferable layers in different
+        # cache-group orders. Keep their region lists aligned without putting
+        # layers.10 before layers.2, which would break PP region slicing.
         layer_names = (
             xfer_buffers
             if self._is_csa_linear
-            else sorted(xfer_buffers, key=_region_sort_key)
+            else sorted(
+                xfer_buffers,
+                key=_region_sort_key,
+            )
         )
         for layer_name in layer_names:
             cache = xfer_buffers[layer_name]
@@ -1300,6 +1469,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     layer_name,
                 )
                 continue
+            group_index = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
             # `layer_spec.page_size_bytes` only accounts for logical page_size, that is
             # the page_size assuming constant `self._logical_num_blocks`.
             physical_page_size = (
@@ -1308,17 +1478,48 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 else layer_spec.page_size_bytes
                 // self._physical_blocks_per_logical_kv_block
             )
-            group_id = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
+            group = self.kv_cache_config.transfer_groups[group_index]
+            if group.host_resident:
+                logical_num_blocks = self.kv_cache_config.hisparse_host_num_blocks
+                assert logical_num_blocks is not None
+            else:
+                logical_num_blocks = self.kv_cache_config.num_blocks
+            group_id = group_index
             num_blocks = (
-                self._logical_num_blocks
+                logical_num_blocks
                 if isinstance(layer_spec, MambaSpec)
-                else self.num_blocks
+                else logical_num_blocks * self._physical_blocks_per_logical_kv_block
+            )
+            base_addr = cache.data_ptr()
+            is_mla_region = isinstance(
+                layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
             )
             logger.debug(
                 "Registering layer %s with cache shape: %s", layer_name, cache.shape
             )
             storage = cache.untyped_storage()
             storage_addr = storage.data_ptr()
+            is_host_resident = group.host_resident
+            if cache.device.type == "cpu":
+                mem_type = "DRAM"
+                region_device_id = 0
+            else:
+                mem_type = self.nixl_memory_type
+                region_device_id = max(cache.get_device(), 0)
+                self.device_id = region_device_id
+            if is_host_resident:
+                registration_base = cache.data_ptr()
+                registration_len = cache.nbytes
+            else:
+                registration_base = storage_addr
+                registration_len = storage.nbytes()
+            storage_key = (registration_base, mem_type)
+            if storage_key not in registration_ranges:
+                registration_ranges[storage_key] = (
+                    registration_base,
+                    registration_base + registration_len,
+                    region_device_id,
+                )
 
             if isinstance(layer_spec, KpoolTailSpec):
                 compressed_owner = compressed_region_owners.get(cache.data_ptr())
@@ -1346,22 +1547,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     )
                     continue
 
-            if cache.device.type == "cpu":
-                mem_type = "DRAM"
-                region_device_id = 0
-            else:
-                mem_type = self.nixl_memory_type
-                region_device_id = max(cache.get_device(), 0)
-                self.device_id = region_device_id
-            registration_ranges.setdefault(
-                (storage_addr, mem_type),
-                (storage_addr, storage.nbytes(), region_device_id, ""),
-            )
-
-            is_mla_region = isinstance(
-                layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
-            )
-
             if isinstance(layer_spec, MambaSpec):
                 physical_ratio = self._physical_blocks_per_logical_kv_block
                 block_len = physical_page_size // physical_ratio
@@ -1382,14 +1567,13 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     block_stride = physical_page_size
                 else:
                     block_stride = cache.stride(0) * cache.element_size()
-                storage_is_block_major = num_blocks * block_stride == storage.nbytes()
-                # A layer whose [H, N, C] interior is dense addresses its own page
-                # as one chunk. Otherwise the block's whole row is the only
-                # contiguous transfer unit.
-                hnc_contiguous = (
-                    cache.ndim == 4
-                    and cache.stride(2) == cache.shape[3]
-                    and cache.stride(1) == cache.shape[2] * cache.shape[3]
+                storage_is_block_major = num_blocks * block_stride == registration_len
+                # A view with a contiguous page can own a strided region even
+                # when sibling layers share the same backing allocation.
+                page_contiguous = (
+                    cache.ndim > 1
+                    and cache[0].is_contiguous()
+                    and cache[0].nbytes == physical_page_size
                 )
                 virtual_transfer_pages = _uses_dense_virtual_transfer_pages(
                     layer_spec, cache, physical_page_size, num_blocks
@@ -1400,13 +1584,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                         (cache.data_ptr(), physical_page_size, physical_page_size)
                     ]
                 elif storage_is_block_major and (
-                    (packed_storage and is_mla_region)
-                    or (not hnc_contiguous and not self._is_csa_linear)
+                    not page_contiguous
+                    and ((packed_storage and is_mla_region) or not self._is_csa_linear)
                 ):
-                    # Packed MLA layouts transfer the complete storage row.
-                    storage_block_len = storage.nbytes() // num_blocks
+                    # TODO(Lucas): handle TP slicing for packed_storage; for now
+                    # restrict to MLA (DSv4) where kv is replicated.
+                    storage_block_len = registration_len // num_blocks
                     region_specs = [
-                        (storage_addr, storage_block_len, storage_block_len)
+                        (registration_base, storage_block_len, storage_block_len)
                     ]
                 elif storage_is_block_major:
                     region_specs = [
@@ -1441,6 +1626,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     region_index = seen_base_addresses.index(base_addr)
                     assert region_mem_types[region_index] == mem_type
                     self._region_is_mla[region_index] |= is_mla_region
+                    if is_mla_region:
+                        self.block_len_per_layer[region_index] = block_len
+                        self.block_stride_per_layer[region_index] = block_stride
+                        self.region_num_blocks[region_index] = num_blocks
                     if self.region_group_ids[region_index] != group_id:
                         self.region_group_ids[region_index] = _SHARED_REGION_GROUP_ID
                 else:
@@ -1451,8 +1640,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     self.region_group_ids.append(group_id)
                     self.region_names.append(layer_name)
                     self.region_num_blocks.append(num_blocks)
-                    region_mem_types.append(mem_type)
                     self._region_is_mla.append(is_mla_region)
+                    region_mem_types.append(mem_type)
+
+                if track_region_layers:
+                    if region_index == len(region_layers):
+                        region_layers.append([layer_name])
+                    else:
+                        region_layers[region_index].append(layer_name)
 
                 if isinstance(layer_spec, MambaSpec):
                     if layer_spec.tp_replicated:
@@ -1497,7 +1692,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             == len(self.region_group_ids)
             == len(self.region_names)
             == len(self.region_num_blocks)
-            == len(region_mem_types)
         )
         # Descriptor ids must be region-ordered, matching the remote side.
         self._scratch_region_indices.sort()
@@ -1510,7 +1704,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         if self._has_mamba:
             self.region_num_blocks = [self.num_blocks] * self.num_regions
 
-        if self.pp_size > 1:
+        self._set_region_layers(region_layers)
+
+        if self.pp_size > 1 and not self._is_hma_required:
             start_layer, end_layer = self.model_config.get_layers_start_end_indices(
                 self.vllm_config.parallel_config
             )
@@ -1520,24 +1716,36 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             self._remote_region_offset = regions_per_layer * start_layer
 
         # Total local FA descriptors (boundary between FA and mamba descs).
-        self.num_descs = sum(self.region_num_blocks)
+        xfer_region_num_blocks = (
+            [self.region_num_blocks[i] for i in self._transfer_layer_region_indices]
+            if self._transfer_layer_region_indices
+            else self.region_num_blocks
+        )
+        self.num_descs = sum(xfer_region_num_blocks)
 
+        self._mixed_mem_types = len(set(region_mem_types)) > 1
+        if self._mixed_mem_types:
+            assert self.use_mla and not self._has_mamba, (
+                "Mixed-device KV registration is only supported for MLA "
+                "models without Mamba layers."
+            )
         for mem_type in sorted(set(region_mem_types)):
-            ranges = [
-                registration
-                for (
-                    _,
-                    registration_mem_type,
-                ), registration in registration_ranges.items()
-                if registration_mem_type == mem_type
+            ranges_for_mem_type = [
+                (start, end - start, device_id, "")
+                for (_, cache_mem_type), (
+                    start,
+                    end,
+                    device_id,
+                ) in registration_ranges.items()
+                if cache_mem_type == mem_type
             ]
-            descs = self.nixl_wrapper.get_reg_descs(ranges, mem_type)
+            descs = self.nixl_wrapper.get_reg_descs(ranges_for_mem_type, mem_type)
             self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
             self._registered_descs.append(descs)
 
         self.device_kv_caches = kv_caches
         self.dst_num_blocks[self.engine_id] = self.num_blocks
-        self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks
+        self.dst_region_num_blocks[self.engine_id] = xfer_region_num_blocks
         self.dst_region_group_ids[self.engine_id] = self.region_group_ids
         self.dst_uses_region_group_mapping[self.engine_id] = (
             self._uses_region_group_mapping
@@ -1587,6 +1795,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             region_mem_types=self.region_mem_types,
             dcp_size=self.dcp_size,
             pcp_size=self.pcp_size,
+            region_members=self.region_members,
         )
         # Wrap metadata in payload with hash for defensive decoding
         assert self.compat_hash is not None
@@ -1745,12 +1954,16 @@ def _build_fa_local(
         base_addresses: list[int],
         block_size_ratio: int,
     ) -> np.ndarray:
-        """Build local FA descriptors for all layers as an Nx3 uint64 array."""
+        """Build local FA descriptors as an Nx3 uint64 array."""
         assert self.transfer_topo is not None
         assert base_addresses, "Local KV cache base addresses must not be empty."
         device_id = self.device_id
+        region_indices = self._transfer_layer_region_indices or range(
+            len(base_addresses)
+        )
         parts: list[np.ndarray] = []
-        for i, base_addr in enumerate(base_addresses):
+        for i in region_indices:
+            base_addr = base_addresses[i]
             block_len = self.block_len_per_layer[i] // block_size_ratio
             logical_blocks = np.repeat(
                 np.arange(self.region_num_blocks[i], dtype=np.uint64),
@@ -1774,7 +1987,7 @@ def _build_fa_remote(
         nixl_agent_meta: NixlAgentMetadata,
         block_size_ratio: int,
     ) -> np.ndarray:
-        """Build remote FA descriptors for all layers as an Nx3 uint64 array."""
+        """Build remote FA descriptors as an Nx3 uint64 array."""
         assert self.transfer_topo is not None
         assert nixl_agent_meta.kv_caches_base_addr, (
             "Remote KV cache base addresses must not be empty."
@@ -1791,10 +2004,14 @@ def _build_fa_remote(
         device_id = nixl_agent_meta.device_id
         block_strides = nixl_agent_meta.block_strides
         parts: list[np.ndarray] = []
+        use_layer_name_routing = bool(self._transfer_layer_region_indices)
         for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
-            replicated = self._is_region_replicated(i)
+            local_region = (
+                self._transfer_layer_region_indices[i] if use_layer_name_routing else i
+            )
+            replicated = self._is_region_replicated(local_region)
             # Read our whole local region size from remote..
-            local_block_len = self.block_len_per_layer[i]
+            local_block_len = self.block_len_per_layer[local_region]
             remote_kv_block_len = local_block_len // block_size_ratio
             if block_size_ratio > 1:
                 # ..using remote kv_block_len as transfer unit
@@ -1817,8 +2034,7 @@ def register_local_xfer_handler(
         self,
         block_size: int,
     ) -> tuple[int, np.ndarray]:
-        """
-        Function used for register local xfer handler with local block_size or
+        """Function used for register local xfer handler with local block_size or
         Remote block_size.
 
         When local block_size is same as remote block_size, we use local block_size
@@ -1855,7 +2071,11 @@ def register_local_xfer_handler(
         if self._mixed_mem_types:
             desc_is_dram = np.concatenate(
                 [
-                    np.full(count * block_size_ratio, mem_type == "DRAM", dtype=bool)
+                    np.full(
+                        count * block_size_ratio,
+                        mem_type == "DRAM",
+                        dtype=bool,
+                    )
                     for mem_type, count in zip(
                         self.region_mem_types,
                         self.region_num_blocks,
@@ -1865,24 +2085,29 @@ def register_local_xfer_handler(
             )
             assert len(desc_is_dram) == len(blocks_data)
             desc_pos = np.empty(len(desc_is_dram), dtype=np.int64)
-            dram_indices = np.flatnonzero(desc_is_dram)
-            device_indices = np.flatnonzero(~desc_is_dram)
-            desc_pos[dram_indices] = np.arange(len(dram_indices), dtype=np.int64)
-            desc_pos[device_indices] = np.arange(len(device_indices), dtype=np.int64)
+            dram_idx = np.where(desc_is_dram)[0]
+            vram_idx = np.where(~desc_is_dram)[0]
+            desc_pos[dram_idx] = np.arange(len(dram_idx), dtype=np.int64)
+            desc_pos[vram_idx] = np.arange(len(vram_idx), dtype=np.int64)
             self._desc_is_dram_by_block_size[block_size] = desc_is_dram
             self._desc_pos_by_block_size[block_size] = desc_pos
-            blocks_data[desc_is_dram, 2] = 0
-            dram_descs = self.nixl_wrapper.get_xfer_descs(
-                blocks_data[desc_is_dram], "DRAM"
-            )
+
+            # DRAM descriptors are registered under CPU device 0.
+            blocks_data = [
+                (addr, length, 0) if is_dram else (addr, length, dev)
+                for (addr, length, dev), is_dram in zip(
+                    blocks_data, desc_is_dram, strict=True
+                )
+            ]
+            dram_blocks = [blocks_data[i] for i in dram_idx]
+            vram_blocks = [blocks_data[i] for i in vram_idx]
+            dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM")
             self._dram_src_handles_by_block_size[block_size] = (
                 self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs)
             )
-            device_descs = self.nixl_wrapper.get_xfer_descs(
-                blocks_data[~desc_is_dram], self.nixl_memory_type
-            )
+            descs = self.nixl_wrapper.get_xfer_descs(vram_blocks, self.nixl_memory_type)
             return (
-                self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", device_descs),
+                self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs),
                 blocks_data,
             )
 
@@ -1898,8 +2123,7 @@ def add_remote_agent(
         remote_tp_size: int = 1,
         remote_dcp_size: int = 1,
     ) -> str:
-        """
-        Add the remote NIXL agent and prepare the descriptors for reading cache
+        """Add the remote NIXL agent and prepare the descriptors for reading cache
         blocks from remote.
 
         In particular, handle both homogeneous and heterogeneous TP. The former
@@ -1952,9 +2176,22 @@ def add_remote_agent(
             )
             return self._remote_agents[engine_id][(0, remote_tp_rank)]
 
+        assert self.transfer_topo is not None
+        transfer_topo = self.transfer_topo
         # Number of physical regions registered locally (one per layer/tensor).
         num_local_regions = len(self.block_len_per_layer)
-        if (
+        if self._transfer_layer_region_indices:
+            if self.block_size != nixl_agent_meta.block_size:
+                raise NotImplementedError(
+                    "Attention-HMA push requires identical P/D block sizes."
+                )
+            if remote_tp_size > transfer_topo.tp_size and not self.use_mla:
+                raise NotImplementedError(
+                    "Attention-HMA push does not support decode TP greater than "
+                    "prefill TP yet"
+                )
+            self._align_remote_regions_by_layer(nixl_agent_meta)
+        elif (
             self.pp_size > 1
             and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions
         ):
@@ -1985,8 +2222,6 @@ def add_remote_agent(
                 ]
 
         ### Register remote engine in TransferTopology (idempotent).
-        assert self.transfer_topo is not None
-        transfer_topo = self.transfer_topo
         physical_blocks_per_logical = (
             nixl_agent_meta.physical_blocks_per_logical_kv_block
         )
@@ -2027,11 +2262,12 @@ def add_remote_agent(
                 nixl_agent_meta.region_num_blocks
                 or [nixl_agent_meta.num_blocks] * num_remote_regions
             )
-            self.dst_region_group_ids[engine_id] = nixl_agent_meta.region_group_ids or (
+            remote_region_group_ids = nixl_agent_meta.region_group_ids or (
                 self.region_group_ids
                 if len(self.region_group_ids) == num_remote_regions
                 else [0] * num_remote_regions
             )
+            self.dst_region_group_ids[engine_id] = remote_region_group_ids
             self.dst_uses_region_group_mapping[engine_id] = (
                 len(set(self.dst_region_group_ids[engine_id])) > 1
             )
@@ -2139,8 +2375,6 @@ def add_remote_agent(
             mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info)
             blocks_data = np.concatenate([blocks_data, mamba])
 
-        # A mixed local destination is split into multiple READs below. The
-        # producer remains a single descriptor list in this prerequisite.
         remote_mem_types = set(self.dst_region_mem_types[engine_id])
         if len(remote_mem_types) != 1:
             raise NotImplementedError(
@@ -2159,8 +2393,7 @@ def _validate_remote_agent_handshake(
         remote_tp_size: int,
         remote_dcp_size: int = 1,
     ):
-        """
-        Validate the remote agent handshake metadata ensuring the
+        """Validate the remote agent handshake metadata ensuring the
         invariants hold true.
         """
         remote_engine_id = nixl_agent_meta.engine_id
@@ -2177,6 +2410,11 @@ def _validate_remote_agent_handshake(
             f"DCP sizes must divide one another: local={self.dcp_size}, "
             f"remote={remote_dcp_size} (engine {remote_engine_id})."
         )
+        if self._has_mamba and self.dcp_size != remote_dcp_size:
+            raise RuntimeError(
+                "Hybrid MLA+Mamba NIXL transfers require matching DCP sizes, "
+                f"got local={self.dcp_size}, remote={remote_dcp_size}."
+            )
 
         tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
         block_size_ratio = self.transfer_topo.block_size_ratio(
@@ -2193,6 +2431,18 @@ def _validate_remote_agent_handshake(
         remote_physical_per_logical = (
             nixl_agent_meta.physical_blocks_per_logical_kv_block
         )
+        if (
+            (self.dcp_size > 1 or remote_dcp_size > 1)
+            and self.region_group_ids
+            and self.region_group_ids != self.dst_region_group_ids[remote_engine_id]
+            and (
+                self._physical_blocks_per_logical_kv_block != 1
+                or remote_physical_per_logical != 1
+            )
+        ):
+            raise NotImplementedError(
+                "DCP region pulls require matching logical and physical block sizes"
+            )
         if (
             self._has_mamba
             and remote_physical_per_logical
@@ -2276,6 +2526,9 @@ def _validate_remote_agent_handshake(
         # the per-rank KV head ratio rather than the raw tp_ratio, because GQA
         # replication caps per-rank heads at 1 when tp > total_kv_heads
         # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here.
+        local_regions = self._transfer_layer_region_indices or range(
+            len(self.block_len_per_layer)
+        )
         if self._has_mamba and self.use_mla:
             # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are
             # kernel-granularity views of the mamba-unified page. The MLA
@@ -2292,7 +2545,7 @@ def _validate_remote_agent_handshake(
                 f"remote={nixl_agent_meta.block_lens}."
             )
         elif not self._has_mamba:
-            assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (
+            assert len(local_regions) == len(nixl_agent_meta.block_lens), (
                 "Number of KV layers must match between prefill and decode"
             )
             model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated(
@@ -2301,8 +2554,9 @@ def _validate_remote_agent_handshake(
             total_kv_heads = self.transfer_topo.total_num_kv_heads
             local_heads = self.transfer_topo.local_physical_heads
             remote_heads = max(1, total_kv_heads // remote_tp_size)
-            for i, local_len in enumerate(self.block_len_per_layer):
-                replicated = model_replicated or self._is_region_replicated(i)
+            for i, region_idx in enumerate(local_regions):
+                local_len = self.block_len_per_layer[region_idx]
+                replicated = model_replicated or self._is_region_replicated(region_idx)
                 remote_len = nixl_agent_meta.block_lens[i]
                 if replicated:
                     assert local_len // block_size_ratio == remote_len, (
@@ -2333,8 +2587,7 @@ def _validate_remote_agent_handshake(
 
         # TP workers that handhshake with same remote have same #blocks.
         assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
-        # Same number of regions/~layers.
-        assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer)
+        assert len(nixl_agent_meta.kv_caches_base_addr) == len(local_regions)
         num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr)
         if nixl_agent_meta.region_num_blocks is not None:
             assert len(nixl_agent_meta.region_num_blocks) == num_remote_regions
@@ -2350,7 +2603,12 @@ def _validate_remote_agent_handshake(
         if nixl_agent_meta.region_group_ids is not None:
             assert len(nixl_agent_meta.region_group_ids) == num_remote_regions
         if nixl_agent_meta.region_names is not None:
-            assert nixl_agent_meta.region_names == self.region_names, (
+            expected_names = (
+                list(self._transfer_layer_names)
+                if self._transfer_layer_names
+                else self.region_names
+            )
+            assert nixl_agent_meta.region_names == expected_names, (
                 "NIXL transfer regions must name the same model caches in the "
                 "same order"
             )
@@ -2358,7 +2616,7 @@ def _validate_remote_agent_handshake(
             assert len(nixl_agent_meta.region_mem_types) == num_remote_regions
 
     def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
-        """copy recved kv from host buffer to device."""
+        """Copy recved kv from host buffer to device."""
         assert self.use_host_buffer
         assert self.copy_blocks is not None
 
@@ -2383,7 +2641,7 @@ def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
             )
 
     def save_kv_to_host(self, metadata: NixlConnectorMetadata):
-        """copy kv from device to host buffer."""
+        """Copy kv from device to host buffer."""
         assert self.use_host_buffer
         assert self.copy_blocks is not None
 
@@ -2432,8 +2690,7 @@ def post_process_device_kv_on_receive(
         block_ids_list: list[tuple[list[int], int]],
         convert: bool = True,
     ):
-        """
-        Post process device kv cache after receiving from remote.
+        """Post process device kv cache after receiving from remote.
 
         3 types of conversion supported (``convert``):
             * kv_cache_postprocess_layout => convert from HND to NHD
@@ -2522,8 +2779,7 @@ def post_process_device_kv_on_receive(
     def post_process_device_kv_on_receive_heterogeneous_attn(
         self, block_ids: list[int]
     ):
-        """
-        Post process device kv cache after receiving from remote
+        """Post process device kv cache after receiving from remote
         for heterogeneous attention.
         """
         assert self.enable_heterogeneous_attn_post_process
@@ -2536,30 +2792,62 @@ def post_process_device_kv_on_receive_heterogeneous_attn(
                 indices=indices,
             )
 
-    def get_finished(self) -> tuple[set[str], set[str]]:
-        """
-        Get requests that are done sending or recving on this specific worker.
+    def _zero_region_blocks(self, block_ids: BlockIds) -> None:
+        """Clear clipped physical pages in their owning region only."""
+        if not any(block_ids):
+            return
+        bases = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
+        for region, blocks in enumerate(block_ids):
+            if not blocks:
+                continue
+            cache = self.device_kv_caches[self.region_names[region]]
+            storage = cache.untyped_storage()
+            pages = torch.empty(0, dtype=torch.uint8, device=cache.device).set_(
+                storage,
+                bases[region] - storage.data_ptr(),
+                (self.region_num_blocks[region], self.block_len_per_layer[region]),
+                (self.block_stride_per_layer[region], 1),
+            )
+            for block in blocks:
+                pages[block].zero_()
+
+    def get_transfer_results(self) -> KVConnectorTransferResults:
+        """Get transfers that completed on this specific worker.
+
         The scheduler process (via the MultiprocExecutor) will use this output
         to track which workers are done.
         """
         assert self.transfer_topo is not None
         done_sending = self._get_new_notifs()
-        done_recving = self._pop_done_transfers(self._recving_transfers)
+        done_recving, newly_failed = self._pop_done_transfers(self._recving_transfers)
+        if newly_failed:
+            self._recv_failures.update(newly_failed)
 
         done_sending.update(self._replicated_pcp_done_sending)
         self._replicated_pcp_done_sending.clear()
 
-        # Drain queue of requests where handshake or transfer setup failed.
-        failed_recv_reqs = set[ReqId]()
+        # Process receive failures reported by background threads. Each
+        # producer already recorded a specific metric (handshake or
+        # notification), so only the failure bookkeeping runs here.
         while not self._failed_recv_reqs.empty():
             try:
-                failed_recv_reqs.add(self._failed_recv_reqs.get_nowait())
+                req_id = self._failed_recv_reqs.get_nowait()
             except queue.Empty:
                 break
+            self._handle_failed_transfer(
+                req_id, None, self._recv_failures, record_failed_transfer=False
+            )
 
-        # Add failed requests to done_recving for scheduler tracking
-        # (blocks are already marked invalid, scheduler will handle recompute)
-        done_recving.update(failed_recv_reqs)
+        failed_recv_reqs: set[ReqId] = set()
+        if self._recv_failures:
+            # Keep failure bookkeeping independent of healthy request count.
+            for req_id in tuple(self._recv_failures):
+                if req_id not in self._recving_metadata:
+                    self._recv_failures.remove(req_id)
+                elif req_id not in self._recving_transfers:
+                    self._recv_failures.remove(req_id)
+                    failed_recv_reqs.add(req_id)
+            done_recving.update(failed_recv_reqs)
 
         if len(done_sending) > 0 or len(done_recving) > 0:
             logger.debug(
@@ -2573,6 +2861,7 @@ def get_finished(self) -> tuple[set[str], set[str]]:
 
         block_ids_for_blocksize_post_process = defaultdict(list)
         block_ids_for_heterogeneous_attn_post_process = list[list[int]]()
+        direct_device_recving = set[str]()
         for req_id in done_recving:
             # clean up metadata for completed requests
             meta = self._recving_metadata.pop(req_id, None)
@@ -2580,16 +2869,29 @@ def get_finished(self) -> tuple[set[str], set[str]]:
 
             # Skip KV sync and post-processing for failed requests
             if req_id in failed_recv_reqs:
+                self._pending_recv_notifs.pop(req_id, None)
+                # TODO (NickLucche) handle failed transfer for HMA.
+                if not self._is_hma_required:
+                    self._invalid_block_ids.put(set(meta.local_block_ids[0]))
                 logger.warning(
                     "Skipping KV post-processing for failed request %s",
                     req_id,
                 )
                 continue
 
+            self._send_pending_recv_notifs(req_id)
             assert meta.remote is not None
             if self.use_host_buffer:
                 self.sync_recved_kv_to_device(req_id, meta)
 
+            direct_device_recving.add(req_id)
+
+            if meta.region_blocks_to_zero is not None:
+                # P/D group positions differ. Use the actual region read plan,
+                # including any local allocation padding the read did not cover.
+                self._zero_region_blocks(meta.region_blocks_to_zero)
+                continue
+
             # Post processing for heteroblocksize/layout, and for blocks the
             # transfer clipped. The latter happens either at remote-block
             # granularity (block_size_ratio > 1) or at kernel-block
@@ -2638,7 +2940,7 @@ def get_finished(self) -> tuple[set[str], set[str]]:
         for block_ids in block_ids_for_heterogeneous_attn_post_process:
             self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
 
-        self._sync_device_after_mamba_recv(done_recving, failed_recv_reqs)
+        self._sync_device_after_direct_recv(direct_device_recving)
 
         # Handle timeout to avoid stranding blocks on remote.
         now = time.perf_counter()
@@ -2660,20 +2962,21 @@ def get_finished(self) -> tuple[set[str], set[str]]:
             del self._reqs_to_send[req_id]
             done_sending.add(req_id)
 
-        return done_sending, done_recving
+        return KVConnectorTransferResults(
+            finished_sending=done_sending,
+            finished_recving=done_recving,
+            failed_recving=failed_recv_reqs,
+        )
 
-    def _sync_device_after_mamba_recv(
-        self,
-        done_recving: set[str],
-        failed_recv_reqs: set[str],
-    ) -> None:
-        """Synchronize ROCm direct-GPU Mamba receives before model execution."""
-        if (
-            not current_platform.is_rocm()
-            or not self._has_mamba
-            or self.use_host_buffer
-            or not (done_recving - failed_recv_reqs)
-        ):
+    def get_finished(self) -> tuple[set[str], set[str]]:
+        """Compatibility wrapper for the legacy completion API."""
+        results = self.get_transfer_results()
+        return results.finished_sending, results.finished_recving
+
+    def _sync_device_after_direct_recv(self, done_recving: set[str]) -> None:
+        """Make direct NIXL writes visible before model execution."""
+        requires_sync = current_platform.is_rocm() and self._has_mamba
+        if self.use_host_buffer or not done_recving or not requires_sync:
             return
 
         torch.accelerator.synchronize()
@@ -2691,6 +2994,7 @@ def _handle_heartbeat(self, payload: str) -> None:
         Args:
             payload: comma-separated P-side request IDs, e.g.
                      "req_abc,req_def".
+
         """
         new_expiry = time.perf_counter() + self._lease_extension
         for req_id in payload.split(","):
@@ -2706,83 +3010,101 @@ def _handle_heartbeat(self, payload: str) -> None:
                     new_expiry,
                 )
 
-    def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:
-        """
-        Pop completed xfers by checking for DONE state.
+    def _pop_done_transfers(
+        self, transfers: dict[str, list[int]]
+    ) -> tuple[set[str], set[str]]:
+        """Poll transfers, retaining handles until they can be released.
+
         Args:
-            transfers: dict of req_id -> list[running_xfer]
+            transfers: Outstanding handles by request, updated in place.
+
         Returns:
-            set of req_ids that have all done xfers
+            Requests with no outstanding handles and requests with observed
+            failures. Failed requests may still have outstanding handles.
+
         """
         done_req_ids: set[str] = set()
+        failed_req_ids: set[str] = set()
         for req_id, handles in list(transfers.items()):
             in_progress = []
             for handle in handles:
                 try:
                     xfer_state = self.nixl_wrapper.check_xfer_state(handle)
                     if xfer_state == "DONE":
-                        # Get telemetry from NIXL
                         res = self.nixl_wrapper.get_xfer_telemetry(handle)
                         self.xfer_stats.record_transfer(res)
                         self.nixl_wrapper.release_xfer_handle(handle)
                     elif xfer_state == "PROC":
                         in_progress.append(handle)
-                        continue
                     else:
                         self._log_failure(
                             failure_type="transfer_failed",
-                            msg="Marking blocks as invalid",
                             req_id=req_id,
                             xfer_state=xfer_state,
                         )
-                        self._handle_failed_transfer(req_id, handle)
+                        if not self._handle_failed_transfer(
+                            req_id, handle, failed_req_ids
+                        ):
+                            in_progress.append(handle)
                 except Exception as e:
                     self._log_failure(
                         failure_type="transfer_exception",
-                        msg="Marking blocks as invalid",
                         req_id=req_id,
                         error=e,
                     )
-                    self._handle_failed_transfer(req_id, handle)
+                    if not self._handle_failed_transfer(req_id, handle, failed_req_ids):
+                        in_progress.append(handle)
 
             if not in_progress:
-                # Only report request as completed when all transfers are done.
-                # A request failed in an earlier poll was already reported via
-                # _failed_recv_reqs and its metadata popped by get_finished();
-                # don't report it again, just drop the remaining handles.
-                if req_id in self._recving_metadata:
-                    done_req_ids.add(req_id)
-                    self._send_pending_recv_notifs(req_id)
+                done_req_ids.add(req_id)
                 del transfers[req_id]
             else:
                 transfers[req_id] = in_progress
-        return done_req_ids
+        return done_req_ids, failed_req_ids
 
-    def _handle_failed_transfer(self, req_id: str, handle: int | None):
-        """
-        Handle a failed transfer by marking all (logical) blocks as invalid and
-        recording the failure.
+    def _try_release_xfer_handle(self, req_id: str, handle: int) -> bool:
+        """Release a handle, returning False if the caller must retain it."""
+        try:
+            self.nixl_wrapper.release_xfer_handle(handle)
+        except Exception as e:
+            # A status error does not guarantee that the backend stopped DMA.
+            self._log_failure(
+                failure_type="transfer_release_failed",
+                msg="Retaining handle and blocks until release succeeds",
+                req_id=req_id,
+                error=e,
+            )
+            return False
+        return True
+
+    def _handle_failed_transfer(
+        self,
+        req_id: str,
+        handle: int | None,
+        failed_req_ids: set[str] | None = None,
+        record_failed_transfer: bool = True,
+    ) -> bool:
+        """Record a failure and release its handle, returning False to retain it.
 
         Args:
             req_id: The request ID.
             handle: The transfer handle.
+            failed_req_ids: Requests observed as failed; each stays listed
+                until it has no outstanding handles.
+            record_failed_transfer: Whether to count the failure toward the
+                transport-failure metric. Callers that already recorded a more
+                specific metric (eg KV expiry, reported separately from
+                transport failures) pass False.
+
         """
-        # (multi-read) One handle is created per remote rank, and they do not
-        # all fail in the same _pop_done_transfers poll. The request is
-        # reported failed on the first one, which pops its metadata in
-        # get_finished(); on the later failures only the handle cleanup is left.
-        # TODO (NickLucche) handle failed transfer for HMA.
-        # A split READ's notification is sent only after every handle succeeds.
-        self._pending_recv_notifs.pop(req_id, None)
-        if (meta := self._recving_metadata.get(req_id)) is not None:
-            if not self._is_hma_required:
-                self._invalid_block_ids.put(set(meta.local_block_ids[0]))
-            self._failed_recv_reqs.put(req_id)
-        if handle is not None:
-            self.nixl_wrapper.release_xfer_handle(handle)
-        self.xfer_stats.record_failed_transfer()
+        if record_failed_transfer:
+            self.xfer_stats.record_failed_transfer()
+        if failed_req_ids is not None:
+            failed_req_ids.add(req_id)
+        return handle is None or self._try_release_xfer_handle(req_id, handle)
 
     def _send_pending_recv_notifs(self, req_id: str) -> None:
+        """Send notifications deferred by split DRAM/VRAM reads."""
         for agent_name, notif_id in self._pending_recv_notifs.pop(req_id, []):
             try:
                 self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
@@ -2798,9 +3120,7 @@ def _send_pending_recv_notifs(self, req_id: str) -> None:
                 self.xfer_stats.record_failed_notification()
 
     def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
-        """
-        Send heartbeat notifications to remote engines, extending lease on KV blocks.
-        """
+        """Send heartbeats to remote engines, extending the lease on KV blocks."""
         for engine_id, hb_info in metadata.heartbeat_by_engine.items():
             # Proactive handshake (this request may still be in waiting queue) so
             # the **next** heartbeat for this remote can go through.
@@ -2833,8 +3153,7 @@ def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
     def get_mapped_blocks(
         self, block_ids: np.ndarray, block_size_ratio: int
     ) -> np.ndarray:
-        """
-          Calculates the new set of block IDs by mapping every element
+        """Calculates the new set of block IDs by mapping every element
           in the (potentially sparse) input array.
           Example: block_ids=[0, 2], block_size_ratio=2
         get_mapped_blocks    0     1     [2     3]     4     5
@@ -2891,8 +3210,7 @@ def _map_block_ids_for_block_size_ratio(
         return mapped_local, mapped_remote
 
     def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds:
-        """
-        Convert block ids to kernel physical block ids.
+        """Convert block ids to kernel physical block ids.
         This is required when the logical block size (the one set by the user)
         does not match the one required by the attn backend.
         `ratio` is the number of physical blocks per logical block.
@@ -2905,10 +3223,10 @@ def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> Block
             return block_ids
         block_arange = np.arange(0, ratio).reshape(1, -1)
         # Mamba blocks have no logical<>physical discrepancy (block-size=1)
-        group_specs = self.kv_cache_config.transfer_groups
         physical_block_ids = []
         for i, group in enumerate(block_ids):
-            if _is_ssm_spec(get_representative_spec_type(group_specs[i].kv_cache_spec)):
+            spec = self.kv_cache_config.transfer_groups[i].kv_cache_spec
+            if _is_ssm_spec(get_representative_spec_type(spec)):
                 physical_block_ids.append(group)
             else:
                 physical_block_ids.append(
@@ -2929,11 +3247,14 @@ def _apply_dcp_prefix_caching(
         local_dcp_rank: int,
         remote_dcp_size: int,
         local_num_computed_blocks: int,
+        num_remote_blocks: int | None = None,
     ) -> tuple[list[int], list[int]]:
         """Match local and remote block IDs by global DCP position.
 
         ``local_ids`` excludes blocks already satisfied by the local prefix
         cache, while ``remote_ids`` contains the full transferable remote list.
+        ``num_remote_blocks``, when provided, is the global block count before
+        DCP partitioning and excludes allocation padding.
 
         Example:
             local DCP size 2, rank 0 owns *global positions*=[0, 2, 4, 6]
@@ -2946,8 +3267,12 @@ def _apply_dcp_prefix_caching(
 
         An empty result (rank 0 above) takes the notification-only path, allowing the
         remote to release its blocks without performing a transfer.
+
         """
         local_size, remote_size = local_dcp_size, remote_dcp_size
+        if num_remote_blocks is not None:
+            num_local = cdiv(num_remote_blocks - local_dcp_rank, local_size)
+            local_ids = local_ids[: max(0, num_local - local_num_computed_blocks)]
 
         if local_size == remote_size:
             local_slice = local_ids
@@ -2968,6 +3293,8 @@ def _apply_dcp_prefix_caching(
             remote_slice = remote_ids[start_remote::k]
 
         matched_blocks = min(len(local_slice), len(remote_slice))
+        if num_remote_blocks is not None and matched_blocks < len(local_slice):
+            raise ValueError("Remote KV pages do not cover the requested range")
         return local_slice[:matched_blocks], remote_slice[:matched_blocks]
 
     @staticmethod
@@ -2984,14 +3311,41 @@ def _block_ids_by_region(
             block_ids_by_region.append(list(block_ids[group_id]))
         return block_ids_by_region
 
-    @staticmethod
     def _apply_prefix_caching_by_region(
-        decode_block_ids: BlockIds, prefill_block_ids: BlockIds
+        self,
+        decode_block_ids: BlockIds,
+        prefill_block_ids: BlockIds,
+        *,
+        num_computed_blocks: list[int] | None = None,
+        num_remote_blocks: int | None = None,
+        remote_rank: int = 0,
+        remote_dcp_size: int = 1,
     ) -> tuple[BlockIds, BlockIds]:
         """Pair an uncached decode suffix with the same prefill regions."""
         assert len(decode_block_ids) == len(prefill_block_ids)
         if not any(decode_block_ids):
-            return [], prefill_block_ids
+            empty_regions: list[list[int]] = [[] for _ in decode_block_ids]
+            return empty_regions, empty_regions.copy()
+
+        if num_computed_blocks is not None:
+            assert num_remote_blocks is not None
+            matched_decode, matched_prefill = [], []
+            for decode_region, prefill_region, start in zip(
+                decode_block_ids, prefill_block_ids, num_computed_blocks, strict=True
+            ):
+                local, remote = self._apply_dcp_prefix_caching(
+                    decode_region,
+                    prefill_region,
+                    remote_rank,
+                    self.dcp_size,
+                    self.dcp_rank,
+                    remote_dcp_size,
+                    start,
+                    num_remote_blocks=num_remote_blocks,
+                )
+                matched_decode.append(local)
+                matched_prefill.append(remote)
+            return matched_decode, matched_prefill
 
         trimmed_prefill: list[list[int]] = []
         for decode_region, prefill_region in zip(
@@ -3105,17 +3459,14 @@ def _apply_prefix_caching(
         return decode_block_ids, prefill_block_ids
 
     def get_kv_connector_stats(self) -> KVConnectorStats | None:
-        """
-        Get the KV transfer stats for the connector.
-        """
+        """Get the KV transfer stats for the connector."""
         # Clear stats for next iteration
         if not self.xfer_stats.is_empty():
             return self.xfer_stats.clone_and_reset()
         return None
 
     def get_block_ids_with_load_errors(self) -> set[int]:
-        """
-        Return and clear the set of block IDs that failed to load.
+        """Return and clear the set of block IDs that failed to load.
 
         This is called by the scheduler to identify blocks that need
         to be retried after a NIXL transfer failure.
@@ -3139,7 +3490,9 @@ def _evict_stale_engines(self) -> None:
         to be "optimal" until a new handshake is performed.
 
         Engines with active transfers or pending handshakes cannot be stale:
-        - Active transfers touch _engine_last_active in start_load_kv.
+        - Reads stamp _engine_last_active when they are issued, and engines a
+          transfer is still reading from are held back explicitly, since the
+          stamp is not refreshed while the read runs.
         - Pending handshakes don't have an _engine_last_active entry yet
         """
         # NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number
@@ -3150,10 +3503,25 @@ def _evict_stale_engines(self) -> None:
             return
 
         now = time.perf_counter()
+        busy = self._engines_with_inflight_transfers()
         for eid, last_active in list(self._engine_last_active.items()):
-            if now - last_active > self._engine_ttl:
+            if now - last_active > self._engine_ttl and eid not in busy:
                 self._cleanup_remote_engine(eid)
 
+    def _engines_with_inflight_transfers(self) -> set[EngineId]:
+        """Remote engines a transfer is still reading from.
+
+        The timestamp is stamped when a read is issued and not refreshed while
+        it runs, so a transfer that outlives the TTL leaves its engine looking
+        idle. A peer that has lost its NIC holds one indefinitely.
+        """
+        return {
+            meta.remote.engine_id
+            for req_id in self._recving_transfers
+            if (meta := self._recving_metadata.get(req_id)) is not None
+            and meta.remote is not None
+        }
+
     def _cleanup_remote_engine(
         self, engine_id: EngineId, *, log_eviction: bool = True
     ) -> None:
@@ -3168,7 +3536,10 @@ def _cleanup_remote_engine(
         # Notif-only engines (push-mode D side) have no descriptor state.
         for handle in self.dst_xfer_side_handles.pop(engine_id, {}).values():
             self.nixl_wrapper.release_dlist_handle(handle)
-        for agent_name in self._remote_agents.pop(engine_id).values():
+        # Pop under the handshake lock; NIXL teardown stays outside it.
+        with self._handshake_lock:
+            agents = self._remote_agents.pop(engine_id)
+        for agent_name in agents.values():
             self.nixl_wrapper.remove_remote_agent(agent_name)
 
         self.kv_caches_base_addr.pop(engine_id, None)
@@ -3194,9 +3565,44 @@ def _cleanup_remote_engine(
             )
 
     def __del__(self):
-        self.shutdown()
+        with contextlib.suppress(Exception):
+            self.shutdown()
 
-    def shutdown(self):
+    def _finish_shutdown(self) -> None:
+        self._recving_transfers.clear()
+        try:
+            for handle in self.src_xfer_handles_by_block_size.values():
+                self.nixl_wrapper.release_dlist_handle(handle)
+            for handles in self.src_xfer_handles_by_tp_ratio.values():
+                for handle in handles:
+                    self.nixl_wrapper.release_dlist_handle(handle)
+            for handles in self._dram_src_handles_by_tp_ratio.values():
+                for handle in handles:
+                    self.nixl_wrapper.release_dlist_handle(handle)
+            for handle in self._dram_src_handles_by_block_size.values():
+                self.nixl_wrapper.release_dlist_handle(handle)
+        except Exception:
+            logger.exception("NIXL dlist-handle release failed at shutdown.")
+        self.src_xfer_handles_by_block_size.clear()
+        self.src_xfer_handles_by_tp_ratio.clear()
+        self._dram_src_handles_by_tp_ratio.clear()
+        self._dram_src_handles_by_block_size.clear()
+        try:
+            for engine_id in list(self._remote_agents):
+                self._cleanup_remote_engine(engine_id, log_eviction=False)
+        except Exception:
+            logger.exception("NIXL remote-engine cleanup failed at shutdown.")
+        try:
+            for desc in self._registered_descs:
+                self.nixl_wrapper.deregister_memory(desc)
+        finally:
+            self._registered_descs.clear()
+            # Drop cache references before their owners release registered
+            # host memory; handshake futures may outlive model-runner shutdown.
+            self.device_kv_caches = {}
+            self.host_xfer_buffers = {}
+
+    def shutdown(self) -> None:
         """Shutdown the connector worker."""
         if not hasattr(self, "_handshake_initiation_executor"):
             # error happens during init, no need to shutdown
@@ -3206,22 +3612,4 @@ def shutdown(self):
             for handle in handles:
                 self.nixl_wrapper.release_xfer_handle(handle)
         self._recving_transfers.clear()
-        for handle in self.src_xfer_handles_by_block_size.values():
-            self.nixl_wrapper.release_dlist_handle(handle)
-        self.src_xfer_handles_by_block_size.clear()
-        for handles in self.src_xfer_handles_by_tp_ratio.values():
-            for handle in handles:
-                self.nixl_wrapper.release_dlist_handle(handle)
-        self.src_xfer_handles_by_tp_ratio.clear()
-        for handles in self._dram_src_handles_by_tp_ratio.values():
-            for handle in handles:
-                self.nixl_wrapper.release_dlist_handle(handle)
-        self._dram_src_handles_by_tp_ratio.clear()
-        for handle in self._dram_src_handles_by_block_size.values():
-            self.nixl_wrapper.release_dlist_handle(handle)
-        self._dram_src_handles_by_block_size.clear()
-        for engine_id in list(self._remote_agents):
-            self._cleanup_remote_engine(engine_id, log_eviction=False)
-        for desc in self._registered_descs:
-            self.nixl_wrapper.deregister_memory(desc)
-        self._registered_descs.clear()
+        self._finish_shutdown()
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
index 7914b7ab3a19..ac929f7fcd94 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
@@ -26,6 +26,7 @@
     KVConnectorHandshakeMetadata,
     KVConnectorMetadata,
     KVConnectorRole,
+    KVConnectorTransferResults,
     SupportsHMA,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
@@ -107,10 +108,15 @@ def __init__(
                 "Consumers and kv_both require "
                 "prefill_context_parallel_size=1."
             )
-        if pcp_size > 1 and parallel_config.decode_context_parallel_size > 1:
+        dcp_size = parallel_config.decode_context_parallel_size
+        if (
+            pcp_size > 1
+            and dcp_size > 1
+            and (parallel_config.tensor_parallel_size != 1 or dcp_size != pcp_size)
+        ):
             raise NotImplementedError(
-                "NixlConnector PCP producers currently require "
-                "decode_context_parallel_size=1."
+                "NixlConnector PCP+DCP requires TP1 with DCP spanning "
+                "the full PCP group."
             )
         # TODO: Support PCP with bidirectional KV transfer by tracking separate
         # send and receive completion counts.
@@ -208,12 +214,12 @@ def request_finished_all_groups(
     def set_xfer_handshake_metadata_pp_aware(
         self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
     ) -> None:
-        """
-        Set handshake metadata keyed by (pp_rank, tp_rank) so the side
+        """Set handshake metadata keyed by (pp_rank, tp_rank) so the side
         channel can serve every PP stage's agent metadata.
 
         Args:
             metadata (dict): the handshake metadata to set.
+
         """
         assert self.connector_scheduler is not None
         self.connector_scheduler.set_xfer_handshake_metadata(metadata)
@@ -234,6 +240,19 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
         assert self.connector_worker is not None
         return self.connector_worker.get_finished()
 
+    def get_transfer_results(
+        self, finished_req_ids: set[str]
+    ) -> KVConnectorTransferResults:
+        assert self.connector_worker is not None
+        results = self.connector_worker.get_transfer_results()
+        if (
+            self.kv_transfer_config.kv_role == "kv_producer"
+            and self.connector_worker.pcp_rank > 0
+            and not self.connector_worker.pcp_dcp_sharded
+        ):
+            results.finished_sending.clear()
+        return results
+
     def get_block_ids_with_load_errors(self) -> set[int]:
         """Get block IDs that failed to load via NIXL."""
         assert self.connector_worker is not None
@@ -298,19 +317,20 @@ def shutdown(self):
             self.connector_scheduler.shutdown()
 
     def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
-        """
-        Get the KVConnector handshake metadata for this connector.
+        """Get the KVConnector handshake metadata for this connector.
         This metadata is used for out-of-band connector handshake
         between P/D workers.
 
         Returns:
             KVConnectorHandshakeMetadata: the handshake metadata.
             None if no handshake metadata is available.
+
         """
         assert self.connector_worker is not None
         if (
             self.kv_transfer_config.kv_role == "kv_producer"
             and self.connector_worker.pcp_rank > 0
+            and not self.connector_worker.pcp_dcp_sharded
         ):
             return None
         return self.connector_worker.xfer_handshake_metadata
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
index 526d9c688133..ca6f30d0cfed 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Metadata dataclasses and helpers for the NIXL connector."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 from vllm.config import VllmConfig
@@ -47,8 +47,9 @@
 #   9: Add block_strides
 #  10: Add dense virtual transfer pages for compressed MLA caches
 #  11: Add per-region transfer geometry and memory types to NixlAgentMetadata
+#  12: Add per-region member names for PP push
 #
-NIXL_CONNECTOR_VERSION: int = 11
+NIXL_CONNECTOR_VERSION: int = 12
 
 
 @dataclass
@@ -71,12 +72,13 @@ class NixlAgentMetadata:
     region_mem_types: list[str] | None = None
     dcp_size: int = 1
     pcp_size: int = 1
+    # Layer names sharing each advertised region, in region order.
+    region_members: list[list[str]] = field(default_factory=list)
 
 
 @dataclass
 class NixlHandshakePayload(KVConnectorHandshakeMetadata):
-    """
-    Wrapper for NIXL handshake sent over the wire.
+    """Wrapper for NIXL handshake sent over the wire.
 
     Enables two-phase decoding for graceful compatibility checking:
     1. Decode NixlHandshakePayload to get compatibility_hash
@@ -139,8 +141,7 @@ def compute_nixl_compatibility_hash(
     attn_backend_name: str,
     transfer_mode: str = "pull",
 ) -> str:
-    """
-    Compute compatibility hash for NIXL KV transfer.
+    """Compute compatibility hash for NIXL KV transfer.
 
     Hash only the factors that affect whether two NIXL instances can
     successfully transfer KV cache data.
@@ -166,6 +167,7 @@ def compute_nixl_compatibility_hash(
 
     Returns:
         SHA-256 hex digest
+
     """
     from vllm import __version__ as vllm_version
     from vllm.config.utils import hash_factors
@@ -227,6 +229,7 @@ class RemoteMeta:
     engine_id: str
     request_id: str
     blocks_expiry_time: float | None = None
+    num_tokens: int | None = None
 
 
 @dataclass
@@ -246,6 +249,12 @@ class ReqMeta:
     remote_block_size: int | None = None
     # Remote producer pipeline-parallel size (push mode, D side).
     pp_size: int = 1
+    # True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS
+    # and expects it in finished_recving; notify-only recvs must not be reported.
+    awaiting_kvs: bool = False
+    # Worker-only, per-region physical pages to zero after a successful pull.
+    # None selects group-based completion; empty lists mean no zeroing.
+    region_blocks_to_zero: BlockIds | None = None
 
 
 class NixlConnectorMetadata(KVConnectorMetadata):
@@ -275,6 +284,7 @@ def _add_new_req(
         local_block_ids: BlockIds,
         kv_transfer_params: dict[str, Any],
         local_num_computed_blocks: tuple[int, ...] = (),
+        awaiting_kvs: bool = False,
     ) -> ReqMeta:
         return ReqMeta(
             local_block_ids=local_block_ids,
@@ -285,6 +295,7 @@ def _add_new_req(
             remote_block_size=kv_transfer_params.get("remote_block_size"),
             pp_size=kv_transfer_params.get("pp_size", 1),
             local_num_computed_blocks=local_num_computed_blocks,
+            awaiting_kvs=awaiting_kvs,
         )
 
     def add_new_req_to_save(
@@ -303,9 +314,13 @@ def add_new_req_to_recv(
         local_block_ids: BlockIds,
         kv_transfer_params: dict[str, Any],
         local_num_computed_blocks: tuple[int, ...] = (),
+        awaiting_kvs: bool = False,
     ):
         req = self._add_new_req(
-            local_block_ids, kv_transfer_params, local_num_computed_blocks
+            local_block_ids,
+            kv_transfer_params,
+            local_num_computed_blocks,
+            awaiting_kvs,
         )
         req.remote = RemoteMeta(
             block_ids=kv_transfer_params["remote_block_ids"],
@@ -314,5 +329,6 @@ def add_new_req_to_recv(
             host=kv_transfer_params["remote_host"],
             port=kv_transfer_params["remote_port"],
             blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"),
+            num_tokens=kv_transfer_params.get("remote_num_tokens"),
         )
         self.reqs_to_recv[request_id] = req
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
index 9740b518a22c..f283a40d3c2b 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
@@ -34,8 +34,7 @@ def __init__(
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int, bool]:
-        """
-        For remote prefill, pull all prompt blocks from remote
+        """For remote prefill, pull all prompt blocks from remote
         asynchronously relative to engine execution.
 
         Args:
@@ -47,8 +46,8 @@ def get_num_new_matched_tokens(
               external KV cache beyond what is already computed.
             * true if the external KV cache tokens will be loaded
               asynchronously (between scheduler steps).
-        """
 
+        """
         params = request.kv_transfer_params
         logger.debug(
             "NIXLConnector get_num_new_matched_tokens: "
@@ -172,6 +171,9 @@ def update_state_after_alloc(
                         request,
                         local_block_ids,
                         local_num_computed_blocks,
+                        # Parked in WAITING_FOR_REMOTE_KVS only when there is
+                        # something to pull; a full local hit stays RUNNING.
+                        num_external_tokens > 0,
                     )
 
                 else:
@@ -191,8 +193,7 @@ def request_finished(
         request: "Request",
         block_ids: "BlockIds",
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Once a request is finished, determine whether request blocks
+        """Once a request is finished, determine whether request blocks
         should be freed now or will be sent asynchronously and freed later.
         """
         from vllm.v1.request import RequestStatus
@@ -224,7 +225,7 @@ def request_finished(
             # To avoid stranding the prefill blocks in the prefill instance,
             # we must add empty block_ids to _reqs_need_recv so that our
             # worker side will notify and free blocks in the prefill instance.
-            self._reqs_need_recv[request.request_id] = (request, [], ())
+            self._reqs_need_recv[request.request_id] = (request, [], (), False)
             params["do_remote_prefill"] = False
             return False, None
 
@@ -282,7 +283,7 @@ def request_finished(
             remote_request_id=request.request_id,
             remote_host=self.side_channel_host,
             remote_port=self.side_channel_port,
-            tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
+            tp_size=self.transfer_tp_size,
             dcp_size=self.vllm_config.parallel_config.decode_context_parallel_size,
             pp_size=self.vllm_config.parallel_config.pipeline_parallel_size,
             remote_num_tokens=remote_num_tokens,
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
index 758f5c9fab4b..6539b7e57ed8 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
@@ -17,8 +17,10 @@
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
     ReadSpec,
+    _is_attention_spec,
 )
 from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -43,8 +45,7 @@ def __init__(
         super().__init__(vllm_config, engine_id, kv_cache_config)
 
     def start_load_kv(self, metadata: NixlConnectorMetadata):
-        """
-        Start loading by triggering non-blocking nixl_xfer.
+        """Start loading by triggering non-blocking nixl_xfer.
         We check for these trnxs to complete in each step().
         """
         for req_id, meta in metadata.reqs_to_recv.items():
@@ -63,8 +64,14 @@ def start_load_kv(self, metadata: NixlConnectorMetadata):
                 len(meta.local_physical_block_ids),
                 len(meta.remote.block_ids),
             )
-            # always store metadata for failure recovery
-            self._recving_metadata[req_id] = meta
+            # Full local hits and aborted cleanup only notify P; no recv is awaited.
+            # On a full local hit:
+            # - Notification failure must not fail the request: its KV is local.
+            # - Receive completion must not be reported: num_external_tokens == 0,
+            #   so the scheduler never entered WAITING_FOR_REMOTE_KVS to begin with.
+            # Aborted cleanup requests have already been removed from the scheduler.
+            if meta.awaiting_kvs or any(meta.local_block_ids):
+                self._recving_metadata[req_id] = meta
             if remote_engine_id not in self._remote_agents:
                 # Initiate handshake with remote engine to exchange metadata.
                 with self._handshake_lock:
@@ -79,7 +86,7 @@ def start_load_kv(self, metadata: NixlConnectorMetadata):
         while not self._ready_requests.empty():
             self._read_blocks_for_req(*self._ready_requests.get_nowait())
 
-        if self.pcp_rank > 0:
+        if self.pcp_rank > 0 and not self.pcp_dcp_sharded:
             # Replicated-KV PCP: only PCP rank 0 serves the KV, so this rank
             # has nothing to send. Report the requests as sent right away so
             # the scheduler-side aggregation (world_size workers, and any
@@ -148,7 +155,11 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
                 engine_id,
             )
             self.xfer_stats.record_kv_expired_req()
-            self._handle_failed_transfer(req_id, None)
+            # KV expiry is reported separately from transport failures, so only
+            # the state cleanup side of _handle_failed_transfer runs here.
+            self._handle_failed_transfer(
+                req_id, None, self._recv_failures, record_failed_transfer=False
+            )
             return
 
         if any(len(group) > 0 for group in meta.local_block_ids):
@@ -165,17 +176,25 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
 
         dcp_active = self.dcp_size > 1 or remote_info.remote_dcp_size > 1
         local_block_ids = meta.local_physical_block_ids
-        local_region_groups = getattr(self, "region_group_ids", [])
-        remote_region_groups = getattr(self, "dst_region_group_ids", {}).get(
-            engine_id, local_region_groups
-        )
-        if local_region_groups != remote_region_groups:
+        remote_region_groups = self.dst_region_group_ids[engine_id]
+        local_region_groups = self.region_group_ids or remote_region_groups
+        if not local_block_ids:
+            # Region expansion cannot index empty groups. Pass empty specs to
+            # _read_blocks so its existing cache-hit notification path runs.
+            read_specs = [
+                ReadSpec(remote_rank=rank, local_block_ids=[], remote_block_ids=[])
+                for rank in plan.all_source_ranks
+            ]
+        elif local_region_groups != remote_region_groups:
             if not self.use_mla or self._has_mamba:
                 raise NotImplementedError(
                     "Different NIXL cache-group layouts are only supported for "
                     "pure MLA models"
                 )
-            assert len(plan.all_source_ranks) == 1
+            if self.block_size != remote_info.remote_block_size:
+                raise NotImplementedError(
+                    "Region-mapped NIXL transfers require matching physical block sizes"
+                )
             remote_physical_block_ids = self._logical_to_kernel_block_ids(
                 meta.remote.block_ids,
                 remote_info.remote_physical_blocks_per_logical,
@@ -183,15 +202,54 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
             remote_by_region = self._block_ids_by_region(
                 remote_physical_block_ids, remote_region_groups
             )
+            local_by_region = self._block_ids_by_region(
+                local_block_ids, local_region_groups
+            )
+            num_computed_blocks = None
+            num_remote_blocks = None
+            if (
+                meta.remote.num_tokens is not None
+                and meta.local_num_computed_blocks
+                and all(group >= 0 for group in local_region_groups)
+                and all(group >= 0 for group in remote_region_groups)
+            ):
+                transfer_groups = self.kv_cache_config.transfer_group_ids
+                num_computed_blocks = [
+                    meta.local_num_computed_blocks[transfer_groups[group]]
+                    * self._physical_blocks_per_logical_kv_block
+                    for group in local_region_groups
+                ]
+                num_remote_blocks = cdiv(
+                    meta.remote.num_tokens, remote_info.remote_block_size
+                )
+            elif any(local_by_region) and (
+                dcp_active
+                or remote_info.remote_physical_blocks_per_logical
+                != self._physical_blocks_per_logical_kv_block
+            ):
+                raise NotImplementedError(
+                    "Region-mapped pulls with DCP or different logical block sizes "
+                    "require remote_num_tokens, per-group prefix counts "
+                    "and unshared regions"
+                )
             read_specs = [
                 ReadSpec(
-                    remote_rank=plan.all_source_ranks[0],
-                    local_block_ids=self._block_ids_by_region(
-                        local_block_ids, local_region_groups
+                    rank,
+                    *self._apply_prefix_caching_by_region(
+                        local_by_region,
+                        remote_by_region,
+                        num_computed_blocks=num_computed_blocks,
+                        num_remote_blocks=num_remote_blocks,
+                        remote_rank=rank,
+                        remote_dcp_size=remote_info.remote_dcp_size,
                     ),
-                    remote_block_ids=remote_by_region,
                     block_ids_by_region=True,
                 )
+                for rank in plan.all_source_ranks
+            ]
+            meta.region_blocks_to_zero = [
+                list(blocks[sum(len(spec.local_block_ids[r]) for spec in read_specs) :])
+                for r, blocks in enumerate(local_by_region)
             ]
         else:
             remote_logical_block_ids = meta.remote.block_ids
@@ -213,7 +271,9 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
                     local_ids = group_ids(meta.local_block_ids, rank)
                     remote_ids = group_ids(remote_logical_block_ids, rank)
                     for g in range(num_groups):
-                        if not local_ids[g]:
+                        if not local_ids[g] or not _is_attention_spec(
+                            self._group_spec_types[g]
+                        ):
                             continue
                         local_ids[g], remote_ids[g] = self._apply_dcp_prefix_caching(
                             local_ids[g],
@@ -289,6 +349,9 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
                 spec.remote_rank
             ]
 
+            # Once a read routes the request to failure reporting, the
+            # scheduler may free and reuse its blocks, so no sibling READs
+            # may be posted (and P must not be notified).
             if not self._read_blocks(
                 read_spec=spec,
                 request_id=req_id,
@@ -298,6 +361,7 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
                 local_dram_handle=local_dram_handle,
                 remote_xfer_side_handle=remote_xfer_side_handle,
                 expected_consumers=plan.local_consumers,
+                awaiting_kvs=meta.awaiting_kvs,
             ):
                 return
 
@@ -310,7 +374,20 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
             remote_agents = self._remote_agents[meta.remote.engine_id]
             for rank_to_notify, agent in remote_agents.items():
                 if rank_to_notify != (0, read_specs[0].remote_rank):
-                    self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
+                    try:
+                        self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
+                    except Exception as e:
+                        self._log_failure(
+                            failure_type="notification_failed",
+                            msg="Remote rank will not update request state. "
+                            "This may indicate network issues.",
+                            req_id=req_id,
+                            error=e,
+                            dst_engine_id=meta.remote.engine_id,
+                            remote_rank=rank_to_notify[1],
+                            remote_agent_name=agent,
+                        )
+                        self.xfer_stats.record_failed_notification()
 
     def _read_blocks(
         self,
@@ -322,10 +399,14 @@ def _read_blocks(
         local_dram_handle: int | None,
         remote_xfer_side_handle: int,
         expected_consumers: int,
+        awaiting_kvs: bool,
     ) -> bool:
-        """
-        Post a READ point-to-point xfer request from a single local worker to
+        """Post a READ point-to-point xfer request from a single local worker to
         a single remote worker.
+
+        Returns True when the read was posted (or was unnecessary), False
+        when the request was routed to failure reporting — the caller must
+        not post further transfers for it.
         """
         assert self.transfer_topo is not None
         remote_rank = read_spec.remote_rank
@@ -378,14 +459,14 @@ def _read_blocks(
                     remote_agent_name=agent_name,
                 )
                 self.xfer_stats.record_failed_notification()
+            # Report even on notification failure: the KV is already local, and
+            # an unreported parked request would hold its blocks forever.
+            # Notify-only recvs must stay unreported (scheduler asserts).
+            if awaiting_kvs:
+                self._recving_transfers.setdefault(request_id, [])
             return True
 
-        if read_spec.block_ids_by_region:
-            local_block_ids, remote_block_ids = self._apply_prefix_caching_by_region(
-                decode_block_ids=local_block_ids,
-                prefill_block_ids=remote_block_ids,
-            )
-        else:
+        if not read_spec.block_ids_by_region:
             assert (
                 len(remote_block_ids)
                 == len(local_block_ids)
@@ -413,11 +494,11 @@ def _read_blocks(
             dst_num_blocks=self.dst_num_blocks[dst_engine_id],
             block_size_ratio=None,
             physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
-            region_num_blocks=self.dst_region_num_blocks[dst_engine_id],
+            region_num_blocks=(self.dst_region_num_blocks.get(dst_engine_id) or None),
             region_group_ids=(
                 list(range(self.num_regions))
                 if read_spec.block_ids_by_region
-                else self.dst_region_group_ids[dst_engine_id]
+                else (self.dst_region_group_ids.get(dst_engine_id) or None)
             ),
             uses_region_group_mapping=(
                 self.num_regions > 1
@@ -430,11 +511,11 @@ def _read_blocks(
             dst_num_blocks=self.dst_num_blocks[self.engine_id],
             block_size_ratio=block_size_ratio,
             physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
-            region_num_blocks=self.dst_region_num_blocks[self.engine_id],
+            region_num_blocks=(self.dst_region_num_blocks.get(self.engine_id) or None),
             region_group_ids=(
                 list(range(self.num_regions))
                 if read_spec.block_ids_by_region
-                else self.region_group_ids
+                else (self.region_group_ids or None)
             ),
             uses_region_group_mapping=(
                 self.num_regions > 1
@@ -477,16 +558,19 @@ def _read_blocks(
             self._recving_transfers[request_id].append(handle)
             return True
         except Exception as e:
-            # mark all (logical) blocks for this request as invalid
             self._log_failure(
                 failure_type="transfer_setup_failed",
                 req_id=request_id,
-                msg="Marking blocks as invalid",
+                msg="Deferring failure reporting until outstanding transfers finish",
                 error=e,
                 dst_engine_id=dst_engine_id,
                 remote_rank=remote_rank,
             )
-            self._handle_failed_transfer(request_id, handle)
+            if not self._handle_failed_transfer(
+                request_id, handle, self._recv_failures
+            ):
+                assert handle is not None
+                self._recving_transfers[request_id].append(handle)
             return False
 
     def _read_blocks_mixed(
@@ -528,24 +612,25 @@ def _read_blocks_mixed(
                     )
         except Exception:
             for handle in handles:
-                self.nixl_wrapper.release_xfer_handle(handle)
+                if not self._try_release_xfer_handle(request_id, handle):
+                    self._recving_transfers[request_id].append(handle)
             raise
 
         self._pending_recv_notifs.setdefault(request_id, []).append(
             (notif_agent, notif_id)
         )
-        for index, handle in enumerate(handles):
+        for i, handle in enumerate(handles):
             try:
                 self.nixl_wrapper.transfer(handle)
             except Exception:
-                for unstarted in handles[index:]:
-                    self.nixl_wrapper.release_xfer_handle(unstarted)
+                for unstarted in handles[i:]:
+                    if not self._try_release_xfer_handle(request_id, unstarted):
+                        self._recving_transfers[request_id].append(unstarted)
                 raise
             self._recving_transfers[request_id].append(handle)
 
     def _get_new_notifs(self) -> set[str]:
-        """
-        Get req_ids which got a remote xfer message. When multiple consumers
+        """Get req_ids which got a remote xfer message. When multiple consumers
         are reading from the same producer (heterogeneous TP or DCP
         scenario), wait for all consumers to be done pulling.
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py
index 53394cf9c138..a007e0c5a3ae 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py
@@ -198,7 +198,7 @@ def update_state_after_alloc(
         # ReqMeta without a KeyError — the actual remote block IDs are
         # learned by P over the NIXL handshake at WRITE time.
         params["remote_block_ids"] = ()
-        self._reqs_need_recv[request.request_id] = (request, local_block_ids, ())
+        self._reqs_need_recv[request.request_id] = (request, local_block_ids, (), False)
 
         # Mark as processed so a re-entry (e.g. preemption + reschedule)
         # doesn't re-stage the registration.
@@ -240,7 +240,7 @@ def request_finished(
             # recv so the worker emits a notif that lets P free them.
             # Seed remote_block_ids so add_new_req_to_recv won't KeyError.
             params["remote_block_ids"] = ()
-            self._reqs_need_recv[request.request_id] = (request, [], ())
+            self._reqs_need_recv[request.request_id] = (request, [], (), False)
             params["do_remote_prefill"] = False
             return False, None
 
@@ -286,7 +286,8 @@ def request_finished(
             remote_request_id=request.request_id,
             remote_host=self.side_channel_host,
             remote_port=self.side_channel_port,
-            tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
+            tp_size=self.transfer_tp_size,
+            dcp_size=self.vllm_config.parallel_config.decode_context_parallel_size,
             pp_size=self.vllm_config.parallel_config.pipeline_parallel_size,
             remote_num_tokens=remote_num_tokens,
             transfer_mode=self._TRANSFER_MODE,
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
index 21bd956c67e4..ae9a85f0de39 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
@@ -41,6 +41,9 @@
 import msgspec
 
 from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorTransferResults,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
     NixlBaseConnectorWorker,
 )
@@ -80,6 +83,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
     # Distinguishes push from pull in the NIXL compatibility hash.
     _TRANSFER_MODE: str = "push"
 
+    _supports_pp_hma = True
+
     def __init__(
         self,
         vllm_config: "VllmConfig",
@@ -98,6 +103,7 @@ def __init__(
         # (``_pop_done_transfers``); guarded by
         # ``_sending_transfers_lock``.
         self._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list)
+        self._send_failures: set[ReqId] = set()
         self._sending_transfers_lock = threading.Lock()
 
         # Writer-thread owned matching state.
@@ -145,7 +151,7 @@ def register_kv_caches(self, kv_caches: dict[str, "torch.Tensor"]):
             self._push_writer_thread.start()
             logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank)
 
-    def shutdown(self):
+    def shutdown(self) -> None:
         self._push_writer_stop.set()
         # Unblock the writer if it's waiting in the no-active-state branch.
         self._push_writer_wake.set()
@@ -157,13 +163,14 @@ def shutdown(self):
                 for handle in handles:
                     self.nixl_wrapper.release_xfer_handle(handle)
             self._sending_transfers.clear()
+            self._send_failures.clear()
         super().shutdown()
 
     # --- Engine-main-thread entry point -------------------------------- #
 
     def start_load_kv(self, metadata: NixlConnectorMetadata):
         """Pre-process metadata; defer NIXL ops to the writer thread."""
-        if self.pcp_rank > 0:
+        if self.pcp_rank > 0 and not self.pcp_dcp_sharded:
             return
 
         # D-side: track reqs waiting for P to push.
@@ -338,7 +345,7 @@ def _on_handshake(
                 self._log_failure(
                     failure_type="push_reg_handshake_failed", req_id=rid, error=e
                 )
-                self._handle_failed_transfer(rid, None)
+                self._failed_recv_reqs.put(rid)
                 return
             # Re-queue for the writer to send now that the handshake is done.
             self._reg_send_inbox.put((rid, rd))
@@ -351,14 +358,17 @@ def _on_handshake(
     def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None:
         engine_id = reg_data["remote_engine_id"]
         notif_msg = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(reg_data)
-        agents = self._remote_agents.get(engine_id)
+        # _remote_agents is mutated on other threads; snapshot under the lock.
+        with self._handshake_lock:
+            agents = dict(self._remote_agents.get(engine_id) or {})
         if not agents:
             logger.error(
                 "No remote agents for engine %s; cannot send registration for %s",
                 engine_id,
                 req_id,
             )
-            self._handle_failed_transfer(req_id, None)
+            self.xfer_stats.record_failed_notification()
+            self._failed_recv_reqs.put(req_id)
             return
         for rank, agent_name in agents.items():
             try:
@@ -370,6 +380,9 @@ def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None:
                     error=e,
                     remote_rank=rank,
                 )
+                self.xfer_stats.record_failed_notification()
+                # Earlier registrations may still trigger WRITEs into D's blocks.
+                # Keep the receive pending until those writes are finished.
         logger.debug(
             "Sent PUSH_REG for %s to engine %s (%dB)", req_id, engine_id, len(notif_msg)
         )
@@ -522,7 +535,14 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta):
         )
         remote_block_ids = meta.remote.block_ids
         local_block_ids = meta.local_physical_block_ids
-        num_groups = len(local_block_ids)
+        local_region_groups = self.region_group_ids
+        remote_region_groups = self.dst_region_group_ids[engine_id]
+        groups_differ = local_region_groups != remote_region_groups
+        if groups_differ and not self._transfer_layer_group_ids:
+            raise NotImplementedError(
+                "NixlPushConnector does not support different producer and "
+                "consumer cache-group layouts"
+            )
 
         # MLA latent is replicated across D's TP ranks: the tp-mapping
         # collapses it to one rank (fine for reads), but push must WRITE every
@@ -537,6 +557,8 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta):
         else:
             write_ranks = list(plan.all_source_ranks)
 
+        num_groups = len(local_block_ids)
+
         def group_ids(block_ids: BlockIds, rank: int) -> BlockIds:
             return [
                 list(block_ids[g])
@@ -708,9 +730,10 @@ def _xfer_blocks(
             # don't have a ``_recving_metadata`` entry to invalidate, so
             # we just release the handle and let the engine reschedule
             # via the lease / watchdog.
-            if handle is not None:
-                self.nixl_wrapper.release_xfer_handle(handle)
-            self.xfer_stats.record_failed_transfer()
+            with self._sending_transfers_lock:
+                self._send_failures.add(request_id)
+            if not self._handle_failed_transfer(request_id, handle):
+                return handle
             return None
 
     # --- Notification handling on engine main thread ------------------ #
@@ -777,18 +800,31 @@ def _get_new_notifs(self) -> set[str]:
                 self._reqs_to_send.pop(req_id, None)
         return notified_req_ids
 
-    def get_finished(self) -> tuple[set[str], set[str]]:
+    def get_transfer_results(self) -> KVConnectorTransferResults:
         # Engine main thread asking for completions: also wake the writer
         # so it gets a chance to drain NIXL notifs (heartbeats, completion
         # notifs, late PUSH_REGs) even if it had been parked.
         self._push_writer_wake.set()
 
-        done_sending, done_recving = super().get_finished()
+        results = super().get_transfer_results()
+        done_sending = results.finished_sending
 
         # ``_pop_done_transfers`` mutates ``_sending_transfers``; the
         # writer thread also appends to it, so guard the pop.
         with self._sending_transfers_lock:
-            done_pushing = self._pop_done_transfers(self._sending_transfers)
+            done_pushing, failed_pushing = self._pop_done_transfers(
+                self._sending_transfers
+            )
+            # Remember failures until the final sibling WRITE completes.
+            self._send_failures.update(failed_pushing)
+            successful = {
+                req_id
+                for req_id in done_pushing - self._send_failures
+                if req_id in self._reqs_to_send or req_id in self._reqs_to_process
+            }
+            self._send_failures.difference_update(done_pushing | done_sending)
+        # Expired requests were already reported, even if their WRITEs finish later.
+        done_pushing = successful
         for req_id in done_pushing:
             self._reqs_to_send.pop(req_id, None)
             self._reqs_to_process.discard(req_id)
@@ -803,4 +839,4 @@ def get_finished(self) -> tuple[set[str], set[str]]:
         if done_sending:
             self._push_writer_wake.set()
 
-        return done_sending, done_recving
+        return results
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py
index 1e4f5c48e0f7..f6d0684b6b3e 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py
@@ -23,7 +23,7 @@
 
 @dataclass
 class NixlKVConnectorStats(KVConnectorStats):
-    """Container for transfer performance metrics"""
+    """Container for transfer performance metrics."""
 
     def __post_init__(self):
         if not self.data:
@@ -39,6 +39,7 @@ def reset(self):
             "num_descriptors": [],
             "num_failed_transfers": [],
             "num_failed_notifications": [],
+            "num_failed_handshakes": [],
             "num_kv_expired_reqs": [],
         }
 
@@ -57,6 +58,10 @@ def record_failed_notification(self):
         """Record a failed NIXL notification (send_notif)."""
         self.data["num_failed_notifications"].append(1)
 
+    def record_failed_handshake(self):
+        """Record a failed NIXL handshake."""
+        self.data["num_failed_handshakes"].append(1)
+
     def record_kv_expired_req(self):
         """Record a request that had its KV blocks expire."""
         self.data["num_kv_expired_reqs"].append(1)
@@ -72,6 +77,7 @@ def is_empty(self) -> bool:
             self.num_successful_transfers == 0
             and len(self.data["num_failed_transfers"]) == 0
             and len(self.data["num_failed_notifications"]) == 0
+            and len(self.data["num_failed_handshakes"]) == 0
             and len(self.data["num_kv_expired_reqs"]) == 0
         )
 
@@ -84,10 +90,22 @@ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
         return self
 
     def reduce(self) -> dict[str, int | float]:
-        # Compute compact representative stats suitable for CLI logging
+        # Compute compact representative stats suitable for CLI logging.
+        # Failure counts are reported on every interval: transfer, handshake
+        # and notification failures are grouped as sporadic
+        # lower-transport-layer events, while KV expiry is reported separately
+        # as it is an actionable autoscaler signal rather than a transport
+        # issue.
+        failure_counts = {
+            "Num failed transfers": len(self.data["num_failed_transfers"])
+            + len(self.data["num_failed_handshakes"])
+            + len(self.data["num_failed_notifications"]),
+            "Num KV expired reqs": len(self.data["num_kv_expired_reqs"]),
+        }
         if self.num_successful_transfers == 0:
-            # CLI logging only reports successful transfers stats. If all requests in
-            # the interval were unsuccessful, Prom will report failures stats instead.
+            # Timing / throughput stats only cover successful transfers. If
+            # all requests in the interval were unsuccessful, the failure
+            # counts above still surface what happened.
             return {
                 "Num successful transfers": 0,
                 "Avg xfer time (ms)": 0,
@@ -97,6 +115,7 @@ def reduce(self) -> dict[str, int | float]:
                 "Avg MB per transfer": 0,
                 "Throughput (MB/s)": 0,
                 "Avg number of descriptors": 0,
+                **failure_counts,
             }
 
         xfer_time = np.asarray(self.data["transfer_duration"])
@@ -115,13 +134,14 @@ def reduce(self) -> dict[str, int | float]:
 
         return {
             "Num successful transfers": n,
-            "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
+            "Avg xfer time (ms)": round(xfer_time.mean().item() * 1e3, 3),
             "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
-            "Avg post time (ms)": round(post_time.mean() * 1e3, 3),
+            "Avg post time (ms)": round(post_time.mean().item() * 1e3, 3),
             "P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3),
-            "Avg MB per transfer": round(avg_mb, 3),
-            "Throughput (MB/s)": round(throughput_mb_s, 3),
-            "Avg number of descriptors": round(descs.mean(), 1),
+            "Avg MB per transfer": round(avg_mb.item(), 3),
+            "Throughput (MB/s)": round(throughput_mb_s.item(), 3),
+            "Avg number of descriptors": round(descs.mean().item(), 1),
+            **failure_counts,
         }
 
     @property
@@ -212,15 +232,20 @@ def __init__(
         )
         counter_nixl_num_failed_transfers = self._counter_cls(
             name="vllm:nixl_num_failed_transfers",
-            documentation="Number of failed NIXL KV Cache transfers.",
+            documentation="Number of failed NIXL KV Cache transfers, including"
+            " handshake and notification failures. NOTE: KV expiry is tracked"
+            " separately in vllm:nixl_num_kv_expired_reqs.",
             labelnames=labelnames,
         )
         self.counter_nixl_num_failed_transfers = create_metric_per_engine(
             counter_nixl_num_failed_transfers, self.per_engine_labelvalues
         )
+
         counter_nixl_num_failed_notifications = self._counter_cls(
             name="vllm:nixl_num_failed_notifications",
-            documentation="Number of failed NIXL KV Cache notifications.",
+            documentation="Number of failed NIXL KV Cache notifications. "
+            "Retained for compatibility; these failures are also included in "
+            "vllm:nixl_num_failed_transfers.",
             labelnames=labelnames,
         )
         self.counter_nixl_num_failed_notifications = create_metric_per_engine(
@@ -254,13 +279,26 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
         ):
             for list_item in transfer_stats_data[list_item_key]:
                 prom_obj[engine_idx].observe(list_item)
-        for counter_obj, counter_item_key in zip(
+        for counter_obj, counter_item_keys in zip(
             [
                 self.counter_nixl_num_failed_transfers,
                 self.counter_nixl_num_failed_notifications,
                 self.counter_nixl_num_kv_expired_reqs,
             ],
-            ["num_failed_transfers", "num_failed_notifications", "num_kv_expired_reqs"],
+            [
+                # Transfer, handshake and notification failures are grouped:
+                # all are sporadic lower-transport-layer events. KV expiry is
+                # reported separately as it signals autoscaler behavior, not
+                # transport health.
+                (
+                    "num_failed_transfers",
+                    "num_failed_handshakes",
+                    "num_failed_notifications",
+                ),
+                ("num_failed_notifications",),
+                ("num_kv_expired_reqs",),
+            ],
         ):
-            for list_item in transfer_stats_data[counter_item_key]:
-                counter_obj[engine_idx].inc(list_item)
+            for counter_item_key in counter_item_keys:
+                for list_item in transfer_stats_data[counter_item_key]:
+                    counter_obj[engine_idx].inc(list_item)
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py
index b86061673487..6ec0e7307810 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py
@@ -34,8 +34,7 @@
 # TODO: merge with vllm.utils.network_utils.zmq_socket_ctx
 @contextlib.contextmanager
 def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
-    """Context manager for a ZMQ socket"""
-
+    """Context manager for a ZMQ socket."""
     if socket_type not in (zmq.ROUTER, zmq.REQ):
         raise ValueError(f"Unexpected socket type: {socket_type}")
 
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
index e5ed896078ca..e20c1189af1e 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+from vllm.utils.math_utils import round_up
 from vllm.v1.core.kv_cache_utils import (
     resolve_dcp_kv_block_size,
     resolve_kv_cache_block_sizes,
@@ -11,10 +12,12 @@
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     FullAttentionSpec,
+    KVCacheGroupRole,
     KVCacheSpec,
     MLAAttentionSpec,
     SlidingWindowMLASpec,
     SlidingWindowSpec,
+    UniformTypeKVCacheSpecs,
     iter_layer_specs,
 )
 from vllm.v1.kv_offload.config import (
@@ -27,7 +30,30 @@
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
-    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec
+
+
+def get_offloading_group_ids(kv_cache_config: "KVCacheConfig") -> tuple[int, ...]:
+    if kv_cache_config.hisparse_host_num_blocks is None:
+        return kv_cache_config.prefix_cacheable_group_ids
+    return tuple(
+        group_id
+        for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
+        if group.role is KVCacheGroupRole.HISPARSE_INDEXER
+    )
+
+
+def _group_kv_bytes_per_block(group: "KVCacheGroupSpec") -> int:
+    """Return the physical bytes occupied by one block of a cache group.
+
+    Worker configs may retain ``UniformTypeKVCacheSpecs`` while scheduler
+    configs flatten that wrapper to one representative per-layer spec.  Keep
+    the result invariant across those two representations.
+    """
+    spec = group.kv_cache_spec
+    if isinstance(spec, UniformTypeKVCacheSpecs):
+        return spec.page_size_bytes
+    return spec.page_size_bytes * len(group.layer_names)
 
 
 def build_offloading_config(
@@ -42,15 +68,22 @@ def build_offloading_config(
     engine_id = kv_transfer_config.engine_id
 
     parallel_config = vllm_config.parallel_config
+    selected_groups = tuple(
+        (group_id, kv_cache_config.kv_cache_groups[group_id])
+        for group_id in get_offloading_group_ids(kv_cache_config)
+    )
+    if not selected_groups:
+        raise ValueError("KV offloading found no eligible cache groups.")
     groups = tuple(
         OffloadingGroupConfig(
+            group_id=group_id,
             tokens_per_block=resolve_dcp_kv_block_size(
                 group.kv_cache_spec,
                 parallel_config.decode_context_parallel_size,
             ),
             layer_names=tuple(group.layer_names),
         )
-        for group in kv_cache_config.kv_cache_groups
+        for group_id, group in selected_groups
     )
 
     _, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config)
@@ -90,15 +123,32 @@ def build_offloading_config(
         )
 
         tokens_per_block = unique_tokens_per_block.pop()
-        assert tokens_per_chunk_int % tokens_per_block == 0
-        blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
+        if tokens_per_chunk_int % tokens_per_block == 0:
+            blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
+        else:
+            raise ValueError(
+                f"'block_size'={tokens_per_chunk_int} in kv_connector_extra_config "
+                f"must be a multiple of the GPU KV cache block size "
+                f"({tokens_per_block} tokens). Use "
+                f"{round_up(tokens_per_chunk_int, tokens_per_block)} instead, or set "
+                f"'blocks_per_chunk' to express the chunk size in blocks."
+            )
 
     worker_kv_bytes_per_block = 0
-    if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors:
+    if (
+        kv_cache_config.hisparse_host_num_blocks is None
+        and kv_cache_config.num_blocks > 0
+        and kv_cache_config.kv_cache_tensors
+    ):
+        # Scratch filtering must preserve the scheduler/worker allocation stride.
         # Every KVCacheTensor describes placement within the same backing allocation,
         # so its size is the total, not a per-tensor share.
         total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size
         worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
+    elif kv_cache_config.num_blocks > 0:
+        worker_kv_bytes_per_block = sum(
+            _group_kv_bytes_per_block(group) for _, group in selected_groups
+        )
 
     single_group_spec = (
         kv_cache_config.kv_cache_groups[0].kv_cache_spec
@@ -183,6 +233,20 @@ def spec_certifiable(spec: KVCacheSpec) -> bool:
             and parallel_config.world_size == tp_size
         )
 
+    if canonical_layout and is_parallelism_agnostic:
+        replicated_layout = (
+            all(
+                type(spec) is MLAAttentionSpec
+                for _, group in selected_groups
+                for spec in iter_layer_specs(group.kv_cache_spec)
+            )
+            and parallel_config.nnodes_within_dp == 1
+            and (
+                parallel_config.world_size == 1
+                or parallel_config.distributed_executor_backend == "mp"
+            )
+        )
+
     kv_events_config = vllm_config.kv_events_config
     cache_dtype = (
         vllm_config.model_config.dtype
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
index 98720b5b7072..68318a16da16 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
@@ -245,13 +245,24 @@ def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEven
         Yields:
             ``BlockStored`` or ``BlockRemoved`` events corresponding to
             the underlying :class:`OffloadingEvent` stream.
+
         """
+        removed_keys: set[OffloadKey] = set()
         for event in events:
             if event.removed:
+                if self.self_describing_enabled:
+                    removed_keys.update(event.keys)
                 yield from self._take_removed_event(event)
             else:
                 yield from self._take_stored_event(event)
 
+        # A primary removal can precede a queued secondary store in this batch.
+        # Keep its payload until all stores have registered their residencies.
+        for key in removed_keys:
+            meta = self._pending_event_metadata.get(key)
+            if meta is not None and not meta.active_residencies:
+                self._pending_event_metadata.pop(key)
+
     def reset(self) -> None:
         """Drop all tracked state; pending payloads are stale after a
         manager cache reset."""
@@ -406,8 +417,6 @@ def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
                     maybe_convert_block_hash(h) for h in meta.block_hashes
                 )
                 meta.active_residencies.discard((event.medium, event.ownership))
-                if not meta.active_residencies:
-                    self._pending_event_metadata.pop(key)
             else:
                 if self.self_describing_enabled:
                     logger.warning_once(
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py
index 58023c393561..0a4ac96be39f 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py
@@ -170,8 +170,7 @@ class _StatsKey:
 
 @dataclass
 class OffloadingConnectorStats(KVConnectorStats):
-    """
-    Offloading connector stats use flat metric names as keys.
+    """Offloading connector stats use flat metric names as keys.
 
     The ``data`` dict is structured using ``_StatsKey`` / ``_MetricType``::
 
@@ -243,8 +242,7 @@ def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats":
         return self
 
     def reduce(self) -> dict[str, int | float]:
-        """
-        Reduce the observations collected during a time interval to one or
+        """Reduce the observations collected during a time interval to one or
         more representative values (eg avg/median/sum of the series).
         This is meant to be called by the logger to produce a summary of the
         stats for the last time interval.
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
index a3cf1a5a1e68..47aa8654feb8 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import time
 from collections.abc import Iterable, Sequence
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from itertools import chain, islice
 from typing import Any, NamedTuple
 
@@ -114,7 +114,10 @@ def load_window_size_in_chunks(self, num_tokens: int) -> int | None:
             window = max(
                 window,
                 cdiv(
-                    self.kv_cache_spec.sliding_window - 1 + right_padding,
+                    self.kv_cache_spec.sliding_window
+                    - 1
+                    + self.kv_cache_spec.extra_retained_tokens
+                    + right_padding,
                     self.tokens_per_chunk,
                 ),
             )
@@ -126,7 +129,13 @@ def get_sliding_window_size_in_chunks(
 ) -> int | None:
     if isinstance(kv_cache_spec, SlidingWindowSpec):
         assert kv_cache_spec.sliding_window > 0
-        return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk)
+        return max(
+            cdiv(kv_cache_spec.sliding_window, tokens_per_chunk),
+            cdiv(
+                kv_cache_spec.sliding_window - 1 + kv_cache_spec.extra_retained_tokens,
+                tokens_per_chunk,
+            ),
+        )
 
     if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):
         # Attention never reaches back past one chunk
@@ -152,8 +161,9 @@ def resolve_mamba_align_size(
     such groups agree on the same value.
     """
     mamba_align_size: int | None = None
-    for idx, tokens_per_block in enumerate(spec.tokens_per_block):
-        kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
+    for group in spec.config.groups:
+        tokens_per_block = group.tokens_per_block
+        kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec
         if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode in (
             "align",
             "all",
@@ -188,8 +198,9 @@ def from_spec(
         # each segment can never serve a load hit. Relevant for hybrid
         # architectures like DeepSeek V4 (MLA + SWA groups).
         full_attn_tokens_per_chunk: set[int] = set()
-        for idx, tokens_per_block in enumerate(spec.tokens_per_block):
-            kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec
+        for group in spec.config.groups:
+            tokens_per_block = group.tokens_per_block
+            kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec
             sw = get_sliding_window_size_in_chunks(
                 kv_spec, tokens_per_block * spec.blocks_per_chunk
             )
@@ -214,9 +225,9 @@ def from_spec(
         retention_interval = vllm_config.cache_config.prefix_cache_retention_interval
 
         eagle_groups = {
-            idx
-            for idx, g in enumerate(kv_cache_config.kv_cache_groups)
-            if g.is_eagle_group
+            group.group_id
+            for group in spec.config.groups
+            if kv_cache_config.kv_cache_groups[group.group_id].is_eagle_group
         }
 
         use_eagle_block_drop = (
@@ -249,8 +260,10 @@ def from_spec(
             )
 
         kv_group_configs_list: list[GroupOffloadConfig] = []
-        for idx, tokens_per_block in enumerate(spec.tokens_per_block):
-            kv_cache_group = kv_cache_config.kv_cache_groups[idx]
+        for group in spec.config.groups:
+            group_id = group.group_id
+            tokens_per_block = group.tokens_per_block
+            kv_cache_group = kv_cache_config.kv_cache_groups[group_id]
             kv_spec = kv_cache_group.kv_cache_spec
             sw = get_sliding_window_size_in_chunks(
                 kv_spec, tokens_per_block * spec.blocks_per_chunk
@@ -261,7 +274,7 @@ def from_spec(
             )
             kv_group_configs_list.append(
                 GroupOffloadConfig(
-                    group_idx=idx,
+                    group_idx=group_id,
                     tokens_per_block=tokens_per_block,
                     tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk,
                     hashes_per_chunk=(
@@ -272,7 +285,7 @@ def from_spec(
                     kv_cache_spec=kv_spec,
                     manager_cls=manager_cls,
                     kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group),
-                    is_eagle_group=idx in eagle_groups,
+                    is_eagle_group=group_id in eagle_groups,
                     requires_cow_source=(
                         isinstance(kv_spec, MambaSpec)
                         and kv_spec.mamba_cache_mode == "align"
@@ -305,7 +318,7 @@ def from_spec(
         if retention_interval is not None:
             if retention_interval < 0:
                 raise ValueError(
-                    f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL "
+                    f"prefix_cache_retention_interval "
                     f"({retention_interval}) must be non-negative."
                 )
             for config in kv_group_configs:
@@ -314,7 +327,7 @@ def from_spec(
                     and retention_interval % config.tokens_per_chunk != 0
                 ):
                     raise ValueError(
-                        f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL "
+                        f"prefix_cache_retention_interval "
                         f"({retention_interval}) must be a multiple of "
                         f"tokens_per_chunk ({config.tokens_per_chunk})."
                     )
@@ -350,6 +363,8 @@ class RequestOffloadState:
     req_context: ReqContext
     offloading_context: RequestOffloadingContext
     group_states: tuple[RequestGroupState, ...] = field(init=False)
+    # upper bound on tokens to load for this request; None means no cap
+    max_load_tokens: int | None = None
     # upper bound on tokens to offload for this request; None means no cap
     max_offload_tokens: int | None = None
     # number of hits in the GPU cache
@@ -372,6 +387,20 @@ def __post_init__(self) -> None:
         )
         params = self.req.kv_transfer_params
 
+        # NOTE: This field is experimental and subject to change in the future.
+        raw = params.get("max_load_tokens") if params else None
+        if type(raw) is int and raw >= 0:
+            self.max_load_tokens = raw
+            logger.debug(
+                "Request %s: max_load_tokens set to %d",
+                self.req.request_id,
+                raw,
+            )
+        elif raw is not None:
+            logger.warning(
+                "max_load_tokens must be a non-negative int, got %r; ignoring", raw
+            )
+
         # NOTE: This field is experimental and subject to change in the future.
         raw = params.get("max_offload_tokens") if params else None
         if type(raw) is int and raw >= 0:
@@ -398,8 +427,11 @@ def update_offload_keys(self) -> None:
                 None,
                 group_config.hashes_per_chunk,
             ):
-                group_state.offload_keys.append(
-                    make_offload_key(req_block_hash, group_config.group_idx)
+                key = make_offload_key(req_block_hash, group_config.group_idx)
+                group_state.offload_keys.append(key)
+                self.req_context.set_offload_key_position(
+                    key,
+                    len(group_state.offload_keys) * group_config.tokens_per_chunk,
                 )
 
     def update_block_id_groups(
@@ -408,9 +440,10 @@ def update_block_id_groups(
         if new_block_id_groups is None:
             return
 
-        assert len(new_block_id_groups) == len(self.group_states)
-        for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
-            group_state.block_ids.extend(new_blocks)
+        for group_config, group_state in zip(
+            self.config.kv_group_configs, self.group_states
+        ):
+            group_state.block_ids.extend(new_block_id_groups[group_config.group_idx])
 
     def storable_chunks(
         self,
@@ -520,6 +553,7 @@ def _create_req_context(req: Request) -> ReqContext:
     return ReqContext(
         req_id=req.request_id,
         kv_transfer_params=params,
+        kv_hints=req.kv_hints,
         load_tier_filter=load_filter,
     )
 
@@ -541,11 +575,11 @@ def __init__(
 
         full_attention_groups: list[int] = []
         sliding_window_groups: list[int] = []
-        for group_config in self.config.kv_group_configs:
+        for config_idx, group_config in enumerate(self.config.kv_group_configs):
             if group_config.sliding_window_size_in_chunks is None:
-                full_attention_groups.append(group_config.group_idx)
+                full_attention_groups.append(config_idx)
             else:
-                sliding_window_groups.append(group_config.group_idx)
+                sliding_window_groups.append(config_idx)
 
         # sort sliding window groups by window size in decreasing order
         def _sliding_window_sort_key(i: int) -> int:
@@ -680,6 +714,7 @@ def _sliding_window_lookup(
         The first run may need a larger window for a partial rightmost chunk.
         Returns 0 on miss, None if the backend deferred a lookup."""
         defer_lookup = False
+        pending_in_window = False
         consecutive_hits = 0
         required_window = initial_window_size or sliding_window_size
         for idx in range(len(keys) - 1, -1, -1):
@@ -690,7 +725,7 @@ def _sliding_window_lookup(
                     # Block is in cache, just not readable yet — counts
                     # as hit for the consecutive streak. Don't break:
                     # keep scanning to let manager kick off async lookups.
-                    defer_lookup = True
+                    pending_in_window = True
                     consecutive_hits += 1
                 case LookupResult.RETRY:
                     # Block location uncertain — does not count as hit.
@@ -698,48 +733,25 @@ def _sliding_window_lookup(
                     # async lookups.
                     defer_lookup = True
                     consecutive_hits = 0
+                    pending_in_window = False
                     required_window = sliding_window_size
                 case LookupResult.MISS:
                     consecutive_hits = 0
+                    # This gap rules out the incomplete window to its right.
+                    pending_in_window = False
                     required_window = sliding_window_size
             if consecutive_hits == required_window:
-                return idx + required_window if not defer_lookup else None
-        return consecutive_hits if not defer_lookup else None
-
-    def _touch(self, req_status: RequestOffloadState):
-        for group_config, group_state in zip(
-            self.config.kv_group_configs, req_status.group_states
-        ):
-            if group_config.sliding_window_size_in_chunks is None:
-                self.manager.touch(group_state.offload_keys, req_status.req_context)
-            else:
-                # Keep only chunks needed to hit the original request, plus
-                # decoded chunks.
-                chunks_to_skip = max(
-                    0,
-                    group_state.num_hit_chunks
-                    - group_config.sliding_window_size_in_chunks,
+                return (
+                    None if defer_lookup or pending_in_window else idx + required_window
                 )
-                self.manager.touch(
-                    group_state.offload_keys[chunks_to_skip:],
-                    req_status.req_context,
-                )
-        if req_status.partial_tail_boundary is not None:
-            self.manager.touch(
-                tuple(
-                    self._make_boundary_key(
-                        req_status.req,
-                        group.group_idx,
-                        req_status.partial_tail_boundary,
-                    )
-                    for group in self.config.kv_group_configs
-                ),
-                req_status.req_context,
-            )
+        return None if defer_lookup or pending_in_window else consecutive_hits
 
-    def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None:
-        """
-        Find how many tokens beyond num_locally_computed_tokens can be loaded.
+    def _lookup_complete_chunks(
+        self,
+        req_status: RequestOffloadState,
+        max_num_new_tokens: int | None = None,
+    ) -> int | None:
+        """Find how many tokens beyond num_locally_computed_tokens can be loaded.
 
         Iterates full-attention groups first (prefix lookup), then sliding-window
         groups (suffix lookup). Each group may tighten max_hit_size_tokens, which
@@ -748,6 +760,15 @@ def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None
         """
         num_computed_tokens = req_status.num_locally_computed_tokens
         max_hit_size_tokens: int = req_status.req.num_tokens
+        if max_num_new_tokens is not None:
+            max_hit_size_tokens = min(
+                max_hit_size_tokens, num_computed_tokens + max_num_new_tokens
+            )
+        if req_status.max_load_tokens is not None:
+            max_hit_size_tokens = min(
+                max_hit_size_tokens,
+                num_computed_tokens + req_status.max_load_tokens,
+            )
         if self._sliding_window_groups:
             # the last prompt token has to be recomputed to get the logprobs
             # for sliding window attention, we must reduce by 1 to make sure
@@ -791,6 +812,10 @@ def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None
                 max_hit_size_tokens = min(
                     max_hit_size_tokens, len(offload_keys) * tokens_per_chunk
                 )
+                if req_status.max_load_tokens is not None:
+                    max_hit_size_tokens = round_down(
+                        max_hit_size_tokens, tokens_per_chunk
+                    )
                 if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk:
                     # We can only load less than a chunk, so skip.
                     return 0
@@ -924,13 +949,23 @@ def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None
         return num_hit_tokens
 
     def _make_boundary_key(
-        self, request: Request, group_idx: int, boundary_tokens: int
+        self,
+        request: Request,
+        group_idx: int,
+        boundary_tokens: int,
+        req_context: ReqContext,
     ) -> OffloadKey:
         hash_idx = boundary_tokens // self.config.tokens_per_hash - 1
-        return make_offload_key(request.block_hashes[hash_idx], group_idx)
+        key = make_offload_key(request.block_hashes[hash_idx], group_idx)
+        req_context.set_offload_key_position(key, boundary_tokens)
+        return key
 
-    def _lookup(self, req_status: RequestOffloadState) -> int | None:
-        complete_hit = self._lookup_complete_chunks(req_status)
+    def _lookup(
+        self,
+        req_status: RequestOffloadState,
+        max_num_new_tokens: int | None = None,
+    ) -> int | None:
+        complete_hit = self._lookup_complete_chunks(req_status, max_num_new_tokens)
         req_status.partial_tail_boundary = None
         if complete_hit is None or not self.config.supports_partial_tail:
             return complete_hit
@@ -939,9 +974,15 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None:
         complete_boundary = local_tokens + complete_hit
         tokens_per_hash = self.config.tokens_per_hash
         block_end = complete_boundary + self._partial_tail_block_size
-        max_boundary = round_down(
-            min(req_status.req.num_prompt_tokens - 1, block_end - 1), tokens_per_hash
-        )
+        max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1)
+        if max_num_new_tokens is not None:
+            max_boundary = min(max_boundary, local_tokens + max_num_new_tokens)
+        if req_status.max_load_tokens is not None:
+            max_boundary = min(
+                max_boundary,
+                local_tokens + req_status.max_load_tokens,
+            )
+        max_boundary = round_down(max_boundary, tokens_per_hash)
         if max_boundary <= complete_boundary:
             return complete_hit
 
@@ -952,7 +993,10 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None:
             boundary_keys = []
             for group_config in self.config.kv_group_configs:
                 key = self._make_boundary_key(
-                    req_status.req, group_config.group_idx, boundary
+                    req_status.req,
+                    group_config.group_idx,
+                    boundary,
+                    req_status.req_context,
                 )
                 boundary_keys.append(key)
                 result = self.manager.lookup(key, req_status.req_context)
@@ -990,16 +1034,20 @@ def on_new_request(self, request: Request) -> None:
         self._req_status[request.request_id] = req_status
 
     def get_num_new_matched_tokens(
-        self, request: Request, num_computed_tokens: int
+        self,
+        request: Request,
+        num_computed_tokens: int,
+        max_num_new_tokens: int | None = None,
     ) -> tuple[int | None, bool]:
-        """
-        Get number of new tokens that can be loaded beyond the
+        """Get number of new tokens that can be loaded beyond the
         num_computed_tokens.
 
         Args:
             request (Request): the request object.
             num_computed_tokens (int): the number of locally
                 computed tokens for this request
+            max_num_new_tokens (int | None): cap on the number of tokens that
+                may be loaded beyond `num_computed_tokens`, if any.
 
         Returns:
             A tuple with the following elements:
@@ -1010,6 +1058,7 @@ def get_num_new_matched_tokens(
                   should query for this request again later.
                 - `True` if tokens will be loaded asynchronously
                   (between scheduler steps).
+
         """
         req_status = self._req_status[request.request_id]
         for group_state in req_status.group_states:
@@ -1030,7 +1079,7 @@ def get_num_new_matched_tokens(
             num_hit_tokens = 0
         else:
             lookup_start = time.monotonic()
-            num_hit_tokens = self._lookup(req_status)
+            num_hit_tokens = self._lookup(req_status, max_num_new_tokens)
             self._connector_stats.observe_histogram(
                 _ConnectorMetricName.LOOKUP_SYNC_DELAY,
                 time.monotonic() - lookup_start,
@@ -1042,8 +1091,6 @@ def get_num_new_matched_tokens(
                 self._maybe_observe_lookup_async_delay(req_status)
         req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0))
 
-        self._touch(req_status)
-
         return num_hit_tokens, bool(num_hit_tokens)
 
     def update_state_after_alloc(
@@ -1065,11 +1112,11 @@ def update_state_after_alloc(
         # per group
         group_sizes: list[int] = []
         block_indices: list[int] = []
-        for group_config, group_state, group_blocks in zip(
+        for group_config, group_state in zip(
             self.config.kv_group_configs,
             req_status.group_states,
-            blocks.blocks,
         ):
+            group_blocks = blocks.blocks[group_config.group_idx]
             self._current_batch_allocated_block_ids.update(
                 block.block_id for block in group_blocks if block.block_id != 0
             )
@@ -1117,7 +1164,10 @@ def update_state_after_alloc(
                 if partial_tail_boundary is not None:
                     keys_to_load.append(
                         self._make_boundary_key(
-                            request, group_config.group_idx, partial_tail_boundary
+                            request,
+                            group_config.group_idx,
+                            partial_tail_boundary,
+                            req_status.req_context,
                         )
                     )
 
@@ -1160,10 +1210,7 @@ def update_state_after_alloc(
         req_status.partial_tail_boundary = None
 
     def _update_req_states(self, scheduler_output: SchedulerOutput) -> None:
-        """
-        Update request states from the Scheduler's output.
-        """
-
+        """Update request states from the Scheduler's output."""
         # new_block_ids_end[req_id][i] = end of pre-existing block_ids for
         # the i-th sliding window group (before this step's extend).
         # Used to detect sliding window blocks that got re-allocated.
@@ -1184,7 +1231,8 @@ def _update_req_states(self, scheduler_output: SchedulerOutput) -> None:
                         for grp_idx in self._sliding_window_groups
                     )
                 req_status.update_block_id_groups(new_block_id_groups)
-                for new_blocks in new_block_id_groups:
+                for group_config in self.config.kv_group_configs:
+                    new_blocks = new_block_id_groups[group_config.group_idx]
                     for bid in new_blocks:
                         if bid != 0:
                             self._current_batch_allocated_block_ids.add(bid)
@@ -1218,6 +1266,10 @@ def _build_aligned_boundary_store_jobs(
     ) -> dict[int, TransferJob]:
         store_jobs: dict[int, TransferJob] = {}
         num_groups = len(self.config.kv_group_configs)
+        config_idx_by_group = {
+            config.group_idx: idx
+            for idx, config in enumerate(self.config.kv_group_configs)
+        }
         for req_id, entries in handoffs.items():
             req_status = self._req_status.get(req_id)
             if req_status is None:
@@ -1225,7 +1277,10 @@ def _build_aligned_boundary_store_jobs(
             req = req_status.req
             max_boundary = self._calc_num_offloadable_tokens(req_status, req.num_tokens)
             for group_idx, block_id, boundary in entries:
-                group_config = self.config.kv_group_configs[group_idx]
+                config_idx = config_idx_by_group.get(group_idx)
+                if config_idx is None:
+                    continue
+                group_config = self.config.kv_group_configs[config_idx]
                 if (
                     block_id == 0
                     or boundary > max_boundary
@@ -1233,7 +1288,9 @@ def _build_aligned_boundary_store_jobs(
                 ):
                     continue
 
-                key = self._make_boundary_key(req, group_idx, boundary)
+                key = self._make_boundary_key(
+                    req, group_idx, boundary, req_status.req_context
+                )
                 store_output = self.manager.prepare_store([key], req_status.req_context)
                 if store_output is None:
                     self._connector_stats.increase_counter(
@@ -1254,9 +1311,11 @@ def _build_aligned_boundary_store_jobs(
                     fenced_block_ids=[block_id],
                 )
                 group_sizes = [0] * num_groups
-                group_sizes[group_idx] = 1
+                group_sizes[config_idx] = 1
                 block_indices = [0] * num_groups
-                block_indices[group_idx] = boundary // group_config.tokens_per_block - 1
+                block_indices[config_idx] = (
+                    boundary // group_config.tokens_per_block - 1
+                )
                 store_jobs[job_id] = TransferJob(
                     req_id=req_id,
                     src_spec=GPULoadStoreSpec(
@@ -1300,6 +1359,12 @@ def _build_partial_tail_store_jobs(
             assert len(boundaries) == 1
             boundary = boundaries.pop()
             req = req_status.req
+            group_states = {
+                group.group_idx: state
+                for group, state in zip(
+                    self.config.kv_group_configs, req_status.group_states
+                )
+            }
             max_boundary = min(
                 req.num_prompt_tokens,
                 req_status.max_offload_tokens or req.num_prompt_tokens,
@@ -1315,18 +1380,20 @@ def _build_partial_tail_store_jobs(
             block_idx = boundary // self._partial_tail_block_size
             if any(
                 group.group_idx not in self._cow_source_groups
-                and block_idx >= len(req_status.group_states[group.group_idx].block_ids)
+                and block_idx >= len(group_states[group.group_idx].block_ids)
                 for group in self.config.kv_group_configs
             ):
                 continue
             keys = [
-                self._make_boundary_key(req, group.group_idx, boundary)
+                self._make_boundary_key(
+                    req, group.group_idx, boundary, req_status.req_context
+                )
                 for group in self.config.kv_group_configs
             ]
             block_ids = [
                 cow_blocks[group.group_idx]
                 if group.group_idx in self._cow_source_groups
-                else req_status.group_states[group.group_idx].block_ids[block_idx]
+                else group_states[group.group_idx].block_ids[block_idx]
                 for group in self.config.kv_group_configs
             ]
             assert all(block_id != 0 for block_id in block_ids)
@@ -1388,11 +1455,24 @@ def _reachable_store_block_mask(
     ) -> list[bool] | None:
         """Build the block mask for a range of candidate offload chunks."""
         blocks_per_chunk = self.config.blocks_per_chunk
+        kv_cache_spec = group_config.kv_cache_spec
+        if isinstance(kv_cache_spec, SlidingWindowSpec) and (
+            kv_cache_spec.extra_retained_tokens
+        ):
+            # Offload restores all allocated history, including MTP re-prefill
+            # tokens. Widen only the store mask, not the model's attention window.
+            kv_cache_spec = replace(
+                kv_cache_spec,
+                sliding_window=(
+                    kv_cache_spec.sliding_window + kv_cache_spec.extra_retained_tokens
+                ),
+                extra_retained_tokens=0,
+            )
         return group_config.manager_cls.reachable_block_mask(
             start_block=start_chunk_idx * blocks_per_chunk,
             end_block=end_chunk_idx * blocks_per_chunk,
             alignment_tokens=self.config.alignment_tokens,
-            kv_cache_spec=group_config.kv_cache_spec,
+            kv_cache_spec=kv_cache_spec,
             use_eagle=group_config.is_eagle_group,
             retention_interval=self.config.retention_interval,
             reachable_boundaries=reachable_boundaries,
@@ -1594,8 +1674,6 @@ def _build_store_jobs(
                 req_status.advance_stored_idx(num_offloadable_tokens)
                 continue
 
-            self._touch(req_status)
-
             keys_to_store = set(store_output.keys_to_store)
 
             group_sizes: list[int] = []
@@ -1761,12 +1839,12 @@ def has_pending_push_work(self) -> bool:
         return bool(self._jobs) or self.manager.has_pending_work()
 
     def update_connector_output(self, connector_output: KVConnectorOutput):
-        """
-        Update KVConnector state from worker-side connectors output.
+        """Update KVConnector state from worker-side connectors output.
 
         Args:
             connector_output (KVConnectorOutput): the worker-side
                 connectors output.
+
         """
         meta = connector_output.kv_connector_worker_meta
         if not isinstance(meta, OffloadingWorkerMetadata):
@@ -1859,8 +1937,7 @@ def request_finished(
         self,
         request: Request,
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Called when a request has finished, before its blocks are freed.
+        """Called when a request has finished, before its blocks are freed.
 
         Returns:
             True if the request is being saved/sent asynchronously and blocks
@@ -1868,6 +1945,7 @@ def request_finished(
             get_finished().
             Optional KVTransferParams to be included in the request outputs
             returned by the engine.
+
         """
         req_status = self._req_status.get(request.request_id)
 
@@ -1906,12 +1984,12 @@ def take_events(self) -> Iterable[KVCacheEvent]:
         Yields:
             ``BlockStored`` or ``BlockRemoved`` events corresponding to
             the underlying :class:`OffloadingEvent` stream.
+
         """
         yield from self._events_tracker.take_events(self.manager.take_events())
 
     def reset_cache(self) -> None:
         """Reset the offloading manager cache, evicting all stored chunks."""
-
         # reset_cache cannot be called in the middle of a schedule step
         assert not self._current_batch_load_jobs
         assert not self._current_batch_jobs_to_flush
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
index ecdae9c4f2d0..2bb00d88f089 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
@@ -14,6 +14,9 @@
     OffloadingWorkerMetadata,
     ReqId,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
+    get_offloading_group_ids,
+)
 from vllm.logger import init_logger
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
@@ -36,7 +39,7 @@
 
 
 class OffloadingConnectorWorker:
-    """Implementation of Worker side methods"""
+    """Implementation of Worker side methods."""
 
     def __init__(
         self,
@@ -50,7 +53,9 @@ def __init__(
         self.worker: OffloadingWorker | None = None
         # Non-writers still ack: pending_count waits for world_size per job.
         self._is_store_writer = (
-            not self.spec.replicated_layout or self.spec.config.parallel.rank == 0
+            not self.spec.replicated_layout
+            or self.spec.config.canonical_layout
+            or self.spec.config.parallel.rank == 0
         )
 
         # job_id -> req_id for in-flight loads.
@@ -65,7 +70,10 @@ def _init_worker(self, kv_caches: CanonicalKVCaches) -> None:
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         kv_cache_config = self.kv_cache_config
-        num_blocks = kv_cache_config.num_blocks
+        selected_group_ids = get_offloading_group_ids(kv_cache_config)
+        selected_groups = tuple(
+            kv_cache_config.kv_cache_groups[group_id] for group_id in selected_group_ids
+        )
         mappings = derive_canonical_mappings(
             self.vllm_config, kv_cache_config, kv_caches
         )
@@ -76,7 +84,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         unpadded_page_size_bytes: dict[str, int] = {}
         # layer_name -> size of page in bytes
         page_size_bytes: dict[str, int] = {}
-        for kv_cache_group in kv_cache_config.kv_cache_groups:
+        for kv_cache_group in selected_groups:
+            assert not kv_cache_group.host_resident
+            num_blocks = kv_cache_config.num_blocks
             group_layer_names = kv_cache_group.layer_names
             group_kv_cache_spec = kv_cache_group.kv_cache_spec
             if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -131,8 +141,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             ),
             None,
         )
-        if packed_layer_name is not None:
+        if (
+            packed_layer_name is not None
+            and kv_cache_config.hisparse_host_num_blocks is None
+        ):
             (tensor,) = tensors_per_block[packed_layer_name]
+            num_blocks = tensor.shape[0]
             block_stride = tensor.stride(0)
             packed_tensor = tensor.as_strided(
                 (num_blocks, block_stride),
@@ -142,10 +156,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             self._init_worker(
                 CanonicalKVCaches(
                     [CanonicalKVCacheTensor(packed_tensor, block_stride)],
-                    [
-                        [CanonicalKVCacheRef(0, block_stride)]
-                        for _ in kv_cache_config.kv_cache_groups
-                    ],
+                    [[CanonicalKVCacheRef(0, block_stride)] for _ in selected_groups],
                 )
             )
             return
@@ -192,7 +203,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     )
 
         group_data_refs: list[list[CanonicalKVCacheRef]] = []
-        for kv_cache_group in kv_cache_config.kv_cache_groups:
+        for kv_cache_group in selected_groups:
             group_refs: list[CanonicalKVCacheRef] = []
             for layer_name in kv_cache_group.layer_names:
                 group_refs += block_data_refs[layer_name]
@@ -260,14 +271,14 @@ def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):
             )
 
     def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
-        """
-        Returns:
-            tuple of (finished_sending, finished_recving). Stores never
-            emit finished_sending — the scheduler tracks store completion
-            via kv_connector_worker_meta.completed_jobs and fences any
-            block reuse via jobs_to_flush. Loads still emit
-            finished_recving so the base scheduler can resume requests
-            blocked on remote KV (and free aborted-during-load reqs).
+        """Returns:
+        tuple of (finished_sending, finished_recving). Stores never
+        emit finished_sending — the scheduler tracks store completion
+        via kv_connector_worker_meta.completed_jobs and fences any
+        block reuse via jobs_to_flush. Loads still emit
+        finished_recving so the base scheduler can resume requests
+        blocked on remote KV (and free aborted-during-load reqs).
+
         """
         assert self.worker is not None
         finished_recving: set[str] = set()
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
index a64f97805ecc..73570881753d 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable
+from functools import cached_property
 from typing import Any
 
 import torch
@@ -25,6 +26,7 @@
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
     build_offloading_config,
+    get_offloading_group_ids,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
     OffloadingConnectorStats,
@@ -38,6 +40,7 @@
 )
 from vllm.forward_context import ForwardContext
 from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -47,6 +50,25 @@
 
 
 class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
+    @cached_property
+    def _bounding_group_ids(self) -> tuple[int, ...]:
+        """Prefix-cacheable groups this connector does not offload.
+
+        Never offer tokens for a range some group cannot cover: past such a
+        group's own cached prefix its KV would be left unwritten.
+        """
+        offloaded = set(get_offloading_group_ids(self._kv_cache_config))
+        return tuple(
+            group_id
+            for group_id, group in enumerate(self._kv_cache_config.kv_cache_groups)
+            if group.kv_cache_spec.prefix_cacheable and group_id not in offloaded
+        )
+
+    @property
+    def scheduler(self) -> OffloadingConnectorScheduler:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler
+
     @property
     def requires_kv_delivery(self) -> bool:
         # Runs as kv_both, but is a best-effort cache: a dropped save is just a
@@ -138,8 +160,29 @@ def get_num_new_matched_tokens(
     ) -> tuple[int | None, bool]:
         assert self.connector_scheduler is not None
         return self.connector_scheduler.get_num_new_matched_tokens(
-            request, num_computed_tokens
+            request,
+            num_computed_tokens,
+            max_num_new_tokens=self._max_loadable_tokens(request, num_computed_tokens),
+        )
+
+    def _max_loadable_tokens(
+        self, request: "Request", num_computed_tokens: int
+    ) -> int | None:
+        """How far past ``num_computed_tokens`` a load may reach, if bounded.
+
+        Bounded by the deepest prefix the groups this connector does not
+        offload already hold, since nothing refills them beyond it.
+        """
+        if not self._bounding_group_ids:
+            return None
+        assert self._kv_cache_manager is not None
+        coordinator = self._kv_cache_manager.coordinator
+        assert isinstance(coordinator, HybridKVCacheCoordinator)
+        _, per_group_hits = coordinator.find_longest_cache_hit_per_group(
+            request.block_hashes, request.num_tokens - 1
         )
+        bound = min(per_group_hits[group_id] for group_id in self._bounding_group_ids)
+        return max(0, bound - num_computed_tokens)
 
     def update_state_after_alloc(
         self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
@@ -185,6 +228,8 @@ def take_events(self) -> Iterable[KVCacheEvent]:
 
     @classmethod
     def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None:
+        if vllm_config.attention_config.hisparse_config is not None:
+            return "BLHNC"
         return "LBHNC"
 
     def reset_cache(self) -> bool | None:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
index 0589c83232f6..64bf8146c9e5 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
@@ -19,6 +19,7 @@
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import KVConnectorOutput
 from vllm.v1.simple_kv_offload.manager import (
+    BoundaryStoreStats,
     SimpleCPUOffloadScheduler,
 )
 from vllm.v1.simple_kv_offload.metadata import (
@@ -300,6 +301,12 @@ def has_pending_push_work(self) -> bool:
             return self.scheduler_manager.has_pending_stores()
         return False
 
+    def get_boundary_store_stats(self) -> BoundaryStoreStats | None:
+        """Return cumulative boundary-handoff store diagnostics."""
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.get_boundary_store_stats()
+        return None
+
     def take_events(self) -> Iterable[KVCacheEvent]:
         if self.scheduler_manager is not None:
             return self.scheduler_manager.take_events()
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
index 00dc05bfc4a9..04ce6a0b92c5 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
@@ -83,6 +83,7 @@ def remote_conv_offsets(
                 < 0:   P_TP > D_TP — P pages are smaller, D reads entire
                        P page.  Local dims are scaled down by |tp_ratio|
                        to get P-sized offsets.
+
         """
         offsets: list[tuple[int, int]] = []
         if tp_ratio >= 1:
@@ -121,6 +122,7 @@ def derive_mamba_conv_split(
     Returns:
         MambaConvSplitInfo with per-rank sub-projection dims, conv_rows,
         conv_dtype_size, and ssm_sizes (conv_state_bytes, ssm_state_bytes).
+
     """
     _supported = (
         MambaAttentionBackendEnum.MAMBA1,
@@ -228,5 +230,6 @@ def compute_physical_blocks_per_logical(
     Args:
         ssm_sizes: (conv_state_bytes, ssm_state_bytes) from NixlAgentMetadata.
         block_len: the engine's block_len in bytes (from block_lens[0]).
+
     """
     return math.ceil((ssm_sizes[0] + ssm_sizes[1]) / block_len)
diff --git a/vllm/distributed/kv_transfer/kv_transfer_state.py b/vllm/distributed/kv_transfer/kv_transfer_state.py
index 211d33fff215..58d77f827734 100644
--- a/vllm/distributed/kv_transfer/kv_transfer_state.py
+++ b/vllm/distributed/kv_transfer/kv_transfer_state.py
@@ -38,6 +38,7 @@ def is_v1_kv_transfer_group(connector: KVConnectorBaseType | None = None) -> boo
     Note:
         This function will no-longer be needed after the v1 KV connector
         becomes the default.
+
     """
     if connector is None:
         connector = _KV_CONNECTOR_AGENT
@@ -72,19 +73,14 @@ def _sync_engine_id_across_tp(vllm_config: "VllmConfig") -> None:
 def ensure_kv_transfer_initialized(
     vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"
 ) -> None:
-    """
-    Initialize KV cache transfer parallel group.
-    """
-
+    """Initialize KV cache transfer parallel group."""
     global _KV_CONNECTOR_AGENT
 
-    if vllm_config.kv_transfer_config is None:
+    if _KV_CONNECTOR_AGENT is not None:
         return
 
-    if (
-        vllm_config.kv_transfer_config.is_kv_transfer_instance
-        and _KV_CONNECTOR_AGENT is None
-    ):
+    kv_transfer_config = vllm_config.kv_transfer_config
+    if kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance:
         # NIXL P/D requires an interleave_size equal to block_size.
         vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config)
         _sync_engine_id_across_tp(vllm_config)
diff --git a/vllm/distributed/mooncake_store.py b/vllm/distributed/mooncake_store.py
new file mode 100644
index 000000000000..f76bf0276618
--- /dev/null
+++ b/vllm/distributed/mooncake_store.py
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared configuration and setup for EC and KV Mooncake Store clients."""
+
+import json
+import os
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import regex as re
+
+DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
+DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
+DEFAULT_TENANT_ID = "default"
+
+
+MooncakeMode = Literal["embedded", "standalone-store"]
+
+
+@dataclass
+class MooncakeStoreConfig:
+    """Configuration for MooncakeDistributedStore.
+
+    ``mode`` selects the topology: ``embedded`` (each rank contributes
+    ``global_segment_size`` in-process) or ``standalone-store`` (rank
+    contributes 0; an external ``mooncake_client`` process owns the pool
+    and the SSD tier).
+    """
+
+    metadata_server: str
+    master_server_address: str
+    protocol: str
+    device_name: str
+    mode: MooncakeMode = "embedded"
+    global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE
+    local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE
+    enable_offload: bool = False
+    tenant_id: str = DEFAULT_TENANT_ID
+
+    def __post_init__(self) -> None:
+        if self.mode not in ("embedded", "standalone-store"):
+            raise ValueError(f"unknown Mooncake mode: {self.mode!r}")
+        if self.local_buffer_size <= 0:
+            raise ValueError("local_buffer_size must be > 0")
+        if self.mode == "embedded" and self.global_segment_size == 0:
+            raise ValueError("embedded mode requires global_segment_size > 0")
+        if self.mode == "standalone-store" and self.global_segment_size != 0:
+            raise ValueError("standalone-store mode requires global_segment_size == 0")
+
+    @staticmethod
+    def from_file(file_path: str) -> "MooncakeStoreConfig":
+        with open(file_path) as file:
+            config = json.load(file)
+        return MooncakeStoreConfig(
+            metadata_server=config.get("metadata_server", ""),
+            master_server_address=config.get("master_server_address", ""),
+            protocol=config.get("protocol", "rdma"),
+            device_name=config.get("device_name", ""),
+            mode=config.get("mode", "embedded"),
+            global_segment_size=_parse_size(
+                config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
+            ),
+            local_buffer_size=_parse_size(
+                config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
+            ),
+            enable_offload=bool(config.get("enable_offload", False)),
+            tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)),
+        )
+
+    @staticmethod
+    def load_from_config() -> "MooncakeStoreConfig":
+        config_path = os.getenv("MOONCAKE_CONFIG_PATH")
+        if not config_path:
+            raise ValueError(
+                "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
+            )
+        return MooncakeStoreConfig.from_file(config_path)
+
+
+def _normalize_tenant_id(value: Any) -> str:
+    if value is None:
+        return DEFAULT_TENANT_ID
+    if not isinstance(value, str):
+        raise TypeError(
+            f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}"
+        )
+    tenant_id = value.strip()
+    return tenant_id if tenant_id else DEFAULT_TENANT_ID
+
+
+def _parse_size(value: Any) -> int:
+    """Parse storage size strings with units: GB, MB, KB, B."""
+    if isinstance(value, int):
+        return value
+    if not isinstance(value, str):
+        try:
+            return int(value)
+        except (TypeError, ValueError) as e:
+            raise TypeError(f"Unsupported type for size: {type(value)}") from e
+
+    cleaned = value.strip().lower()
+    if not cleaned:
+        raise ValueError("Size cannot be empty.")
+
+    unit_multipliers = {
+        "gb": 1024**3,
+        "mb": 1024**2,
+        "kb": 1024,
+        "b": 1,
+    }
+    match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned)
+    if not match:
+        raise ValueError(f"Invalid format: '{value}'")
+
+    number_str = match.group(1)
+    unit = match.group(2) or "b"
+    multiplier = unit_multipliers[unit]
+
+    try:
+        numeric_value = float(number_str)
+    except ValueError as exc:
+        raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc
+    return int(numeric_value * multiplier)
+
+
+def setup_mooncake_store(
+    store: Any, config: MooncakeStoreConfig, local_hostname: str
+) -> None:
+    """Initialize a native Store, preserving explicit tenant configuration."""
+    setup_kwargs: dict[str, str] = {}
+    if config.tenant_id != DEFAULT_TENANT_ID:
+        setup_kwargs["tenant_id"] = config.tenant_id
+    ret = store.setup(
+        local_hostname,
+        config.metadata_server,
+        config.global_segment_size,
+        config.local_buffer_size,
+        config.protocol,
+        config.device_name,
+        config.master_server_address,
+        **setup_kwargs,
+    )
+    if ret != 0:
+        raise RuntimeError("Initialize MooncakeDistributedStore failed.")
diff --git a/vllm/distributed/nixl_utils.py b/vllm/distributed/nixl_utils.py
index 3879bee0ce29..c063b3ea7426 100644
--- a/vllm/distributed/nixl_utils.py
+++ b/vllm/distributed/nixl_utils.py
@@ -91,9 +91,26 @@ def is_nixl_available() -> bool:
     return pkg in sys.modules or importlib.util.find_spec(pkg) is not None
 
 
+def alias_nixl_for_ray() -> None:
+    """Let Ray's ``nixl._api`` import resolve to the ROCm implementation."""
+    if not current_platform.is_rocm() or "nixl._api" in sys.modules:
+        return
+
+    package_name = _get_nixl_package_name()
+    if not is_nixl_available():
+        raise ImportError(f"Ray's NIXL transport on ROCm requires {package_name}")
+
+    _maybe_set_ucx_rcache_limit()
+    package = importlib.import_module(package_name)
+    api = importlib.import_module(f"{package_name}._api")
+    sys.modules.setdefault("nixl", package)
+    sys.modules.setdefault("nixl._api", api)
+
+
 __all__ = [
     "NixlWrapper",
     "nixl_agent_config",
     "nixlXferTelemetry",
     "is_nixl_available",
+    "alias_nixl_for_ray",
 ]
diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
index 92deff118398..20118649e1d1 100644
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -32,6 +32,7 @@
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
 from datetime import timedelta
+from math import gcd
 from multiprocessing import shared_memory
 from typing import TYPE_CHECKING, Any, Protocol
 from unittest.mock import patch
@@ -140,9 +141,11 @@ def _split_tensor_dict(
 
 def _get_unique_name(name: str) -> str:
     """Get a unique name for the group.
+
     Example:
     _get_unique_name("tp") -> "tp:0"
     _get_unique_name("tp") -> "tp:1"
+
     """
     if name not in _group_name_counter:
         _group_name_counter[name] = 0
@@ -419,8 +422,7 @@ def patched_fused_scaled_matmul_reduce_scatter(
 
 
 class GroupCoordinator:
-    """
-    PyTorch ProcessGroup wrapper for a group of processes.
+    """PyTorch ProcessGroup wrapper for a group of processes.
     PyTorch ProcessGroup is bound to one specific communication backend,
         e.g. NCCL, Gloo, MPI, etc.
     GroupCoordinator takes charge of all the communication operations among
@@ -633,34 +635,34 @@ def create_single_reader_mq_broadcasters(
 
     @property
     def first_rank(self):
-        """Return the global rank of the first process in the group"""
+        """Return the global rank of the first process in the group."""
         return self.ranks[0]
 
     @property
     def last_rank(self):
-        """Return the global rank of the last process in the group"""
+        """Return the global rank of the last process in the group."""
         return self.ranks[-1]
 
     @property
     def is_first_rank(self):
-        """Return whether the caller is the first process in the group"""
+        """Return whether the caller is the first process in the group."""
         return self.rank == self.first_rank
 
     @property
     def is_last_rank(self):
-        """Return whether the caller is the last process in the group"""
+        """Return whether the caller is the last process in the group."""
         return self.rank == self.last_rank
 
     @property
     def next_rank(self):
-        """Return the global rank of the process that follows the caller"""
+        """Return the global rank of the process that follows the caller."""
         rank_in_group = self.rank_in_group
         world_size = self.world_size
         return self.ranks[(rank_in_group + 1) % world_size]
 
     @property
     def prev_rank(self):
-        """Return the global rank of the process that precedes the caller"""
+        """Return the global rank of the process that precedes the caller."""
         rank_in_group = self.rank_in_group
         world_size = self.world_size
         return self.ranks[(rank_in_group - 1) % world_size]
@@ -673,11 +675,11 @@ def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None
         else:
             stream = graph_capture_context.stream
 
-        # only cuda uses this function,
+        # only cuda/rocm uses this function,
         # so we don't abstract it into the base class
         maybe_ca_context = nullcontext()
         maybe_fi_pcie_ipc_context: AbstractContextManager[Any] = nullcontext()
-        maybe_aiter_context = nullcontext()
+        maybe_aiter_ar_context = nullcontext()
         from vllm.distributed.device_communicators.cuda_communicator import (
             CudaCommunicator,
         )
@@ -698,12 +700,10 @@ def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None
                 if fi_pcie_ipc_ar_comm is not None:
                     maybe_fi_pcie_ipc_context = fi_pcie_ipc_ar_comm.capture()
 
-            from vllm._aiter_ops import rocm_aiter_ops
-
-            if rocm_aiter_ops.is_enabled():
-                aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
-                if aiter_ar is not None:
-                    maybe_aiter_context = aiter_ar.capture()  # type: ignore
+            # Capture each group's own comm. A global lookup would double-capture
+            aiter_ar_comm = getattr(self.device_communicator, "aiter_ar_comm", None)
+            if aiter_ar_comm is not None:
+                maybe_aiter_ar_context = aiter_ar_comm.capture()  # type: ignore
 
         # ensure all initialization operations complete before attempting to
         # capture the graph on another stream
@@ -715,13 +715,12 @@ def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None
             torch.cuda.stream(stream),
             maybe_ca_context,
             maybe_fi_pcie_ipc_context,
-            maybe_aiter_context,
+            maybe_aiter_ar_context,
         ):
             yield graph_capture_context
 
     def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
-        """
-        User-facing all-reduce function before we actually call the
+        """User-facing all-reduce function before we actually call the
         all-reduce operation.
 
         We need this because Dynamo does not support passing an arbitrary
@@ -810,8 +809,7 @@ def _reduce_scatter_out_place(self, input_: torch.Tensor, dim: int) -> torch.Ten
     def gather(
         self, input_: torch.Tensor, dst: int = 0, dim: int = -1
     ) -> torch.Tensor | None:
-        """
-        NOTE: We assume that the input tensor is on the same device across
+        """NOTE: We assume that the input tensor is on the same device across
         all the ranks.
         NOTE: `dst` is the local rank of the destination rank.
         """
@@ -1364,7 +1362,7 @@ def barrier(self):
         torch.distributed.barrier(group=self.cpu_group)
 
     def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
-        """Sends a tensor to the destination rank in a blocking way"""
+        """Sends a tensor to the destination rank in a blocking way."""
         """NOTE: `dst` is the local rank of the destination rank."""
         if self.device_communicator is None:
             raise ValueError("No device communicator found")
@@ -1560,6 +1558,23 @@ def get_etp_group() -> GroupCoordinator:
     return _ETP
 
 
+_ENGRAM_DP: GroupCoordinator | None = None
+
+
+def get_engram_dp_group() -> GroupCoordinator | None:
+    """Return the DP replicas that share one engram embedding table.
+
+    None when every replica holds a full (TP-sharded) copy, which is the case
+    for models without engram layers and for replicas that span nodes.
+    """
+    return _ENGRAM_DP
+
+
+def get_engram_dp_size() -> int:
+    """Number of DP replicas one engram embedding table is sharded over."""
+    return _ENGRAM_DP.world_size if _ENGRAM_DP is not None else 1
+
+
 _DCP: GroupCoordinator | None = None
 
 
@@ -1617,12 +1632,8 @@ def get_pcp_group() -> GroupCoordinator:
 
 
 @contextmanager
-def graph_capture(
-    device: torch.device,
-    graph_capture_context: GraphCaptureContext | None = None,
-):
-    """
-    `graph_capture` is a context manager which should surround the code that
+def graph_capture(device: torch.device):
+    """`graph_capture` is a context manager which should surround the code that
     is capturing the CUDA graph. Its main purpose is to ensure that some
     operations will be run after the graph is captured, before the graph
     is replayed. It returns a `GraphCaptureContext` object which contains the
@@ -1633,14 +1644,13 @@ def graph_capture(
     the graph capture is running on a separate stream from the default stream,
     in order to explicitly distinguish the kernels to capture
     from other kernels possibly launched on background in the default stream.
-
-    A caller may pass an explicit ``graph_capture_context`` to control the
-    stream used (e.g. to capture on the default stream).
     """
-    context = graph_capture_context or GraphCaptureContext(
-        torch.cuda.Stream(device=device)
-    )
-    with get_tp_group().graph_capture(context), get_pp_group().graph_capture(context):
+    context = GraphCaptureContext(torch.cuda.Stream(device=device))
+    with (
+        get_tp_group().graph_capture(context),
+        get_pp_group().graph_capture(context),
+        get_dp_group().graph_capture(context),
+    ):
         yield context
 
 
@@ -1917,6 +1927,43 @@ def init_distributed_environment(
             _INNER_DP_WORLD = _WORLD
 
 
+def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]:
+    """Defer communicator warm-up when joining a live elastic-EP cluster on ROCm.
+
+    The existing ranks defer their matching warm-up in create_standby_groups(),
+    so a joining worker must too or the warm-up collective has no peers. Both
+    sides warm up at commit.
+    """
+    from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm
+
+    if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
+        return defer_comm_warmup_on_rocm()
+    return nullcontext()
+
+
+def _engram_dp_shard_size(
+    world_size: int,
+    data_parallel_size: int,
+    replica_size: int,
+) -> int:
+    """Size of DP shards aligned with complete replicas and node boundaries."""
+    node_count = get_node_count()
+    replicas_per_node, remainder = divmod(world_size, node_count * replica_size)
+    if remainder:
+        return 1
+    # Shards must divide both the DP group and each node's replica count.
+    shard_size = gcd(data_parallel_size, replicas_per_node)
+    if shard_size == 1 or node_count == 1:
+        return shard_size
+    ranks_per_node = replicas_per_node * replica_size
+    for start in range(0, world_size, ranks_per_node):
+        same_node = in_the_same_node_as(get_world_group().cpu_group, start)
+        node_ranks = [rank for rank, local in enumerate(same_node) if local]
+        if node_ranks != list(range(start, start + ranks_per_node)):
+            return 1
+    return shard_size
+
+
 def initialize_model_parallel(
     tensor_model_parallel_size: int = 1,
     pipeline_model_parallel_size: int = 1,
@@ -1924,14 +1971,17 @@ def initialize_model_parallel(
     decode_context_model_parallel_size: int | None = 1,
     backend: str | None = None,
 ) -> None:
-    """
-    Initialize model parallel groups.
+    """Initialize model parallel groups.
 
     Arguments:
         tensor_model_parallel_size: number of GPUs used for tensor model
             parallelism.
         pipeline_model_parallel_size: number of GPUs used for pipeline model
             parallelism.
+        prefill_context_model_parallel_size: number of GPUs used for context
+            parallelism during prefill.
+        decode_context_model_parallel_size: number of GPUs used for context
+            parallelism during decode.
         backend: name of torch distributed communication backend.
 
     Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
@@ -1946,6 +1996,7 @@ def initialize_model_parallel(
     are on the same DGX box. For example if we are using 2 DGX-1 boxes
     with a total of 16 GPUs, rank 0 to 7 belong to the first box and
     ranks 8 to 15 belong to the second box.
+
     """
     # Get world size and rank. Ensure some consistencies.
     assert torch.distributed.is_initialized()
@@ -2039,6 +2090,40 @@ def initialize_model_parallel(
             group_name="etp",
         )
 
+    # Build the node-local DP group used to shard or share an Engram table.
+    global _ENGRAM_DP
+    assert _ENGRAM_DP is None, "engram data parallel group is already initialized"
+    engram_dp_size = 1
+    engram_config = config.engram_config
+    if (
+        engram_config is not None
+        and config.model_config is not None
+        and config.model_config.architecture == "DeepseekV41ForCausalLM"
+        and not enable_elastic_ep
+    ):
+        engram_dp_size = _engram_dp_shard_size(
+            world_size,
+            data_parallel_size,
+            tensor_model_parallel_size
+            * pipeline_model_parallel_size
+            * prefill_context_model_parallel_size,
+        )
+    if engram_dp_size > 1:
+        group_ranks = (
+            all_ranks.permute(0, 2, 3, 4, 1).reshape(-1, engram_dp_size).unbind(0)
+        )
+        _ENGRAM_DP = init_model_parallel_group(
+            [x.tolist() for x in group_ranks],
+            get_world_group().local_rank,
+            backend,
+            group_name="edp",
+        )
+        logger.info_once(
+            "Engram node-local DP group size: %d (TP=%d)",
+            engram_dp_size,
+            tensor_model_parallel_size,
+        )
+
     # Build the DCP model-parallel groups.
     global _DCP
     assert _DCP is None, "decode context model parallel group is already initialized"
@@ -2099,13 +2184,14 @@ def initialize_model_parallel(
     group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0)
     group_ranks = [x.tolist() for x in group_ranks]
     if enable_elastic_ep:
-        _DP = _init_stateless_group(
-            group_ranks,
-            "dp",
-            parallel_config.data_parallel_master_ip,
-            backend,
-            coord_store=coord_store,
-        )
+        with _elastic_join_warmup_ctx():
+            _DP = _init_stateless_group(
+                group_ranks,
+                "dp",
+                parallel_config.data_parallel_master_ip,
+                backend,
+                coord_store=coord_store,
+            )
     else:
         _DP = init_model_parallel_group(
             group_ranks, get_world_group().local_rank, backend, group_name="dp"
@@ -2128,14 +2214,15 @@ def initialize_model_parallel(
         group_ranks = [x.tolist() for x in group_ranks]
         use_all2all = parallel_config.use_all2all
         if enable_elastic_ep:
-            _EP = _init_stateless_group(
-                group_ranks,
-                "ep",
-                parallel_config.data_parallel_master_ip,
-                backend,
-                coord_store=coord_store,
-                use_all2all=use_all2all,
-            )
+            with _elastic_join_warmup_ctx():
+                _EP = _init_stateless_group(
+                    group_ranks,
+                    "ep",
+                    parallel_config.data_parallel_master_ip,
+                    backend,
+                    coord_store=coord_store,
+                    use_all2all=use_all2all,
+                )
         else:
             _EP = init_model_parallel_group(
                 group_ranks,
@@ -2153,13 +2240,14 @@ def initialize_model_parallel(
         assert _EPLB is None, "EPLB group is already initialized"
         if config.parallel_config.enable_eplb:
             if enable_elastic_ep:
-                _EPLB = _init_stateless_group(
-                    group_ranks,
-                    "eplb",
-                    parallel_config.data_parallel_master_ip,
-                    backend,
-                    coord_store=coord_store,
-                )
+                with _elastic_join_warmup_ctx():
+                    _EPLB = _init_stateless_group(
+                        group_ranks,
+                        "eplb",
+                        parallel_config.data_parallel_master_ip,
+                        backend,
+                        coord_store=coord_store,
+                    )
             else:
                 _EPLB = init_model_parallel_group(
                     group_ranks,
@@ -2285,6 +2373,11 @@ def destroy_model_parallel():
         _TP.destroy()
     _TP = None
 
+    global _ENGRAM_DP
+    if _ENGRAM_DP:
+        _ENGRAM_DP.destroy()
+    _ENGRAM_DP = None
+
     global _DCP
     if _DCP:
         _DCP.destroy()
@@ -2357,6 +2450,12 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     from vllm.platforms import current_platform
 
     if not current_platform.is_cpu():
+        from vllm.triton_utils import HAS_TRITON
+
+        if HAS_TRITON:
+            from vllm.v1.sample.ops.topk_topp_triton import reset_buffer_cache
+
+            reset_buffer_cache()
         torch.accelerator.empty_cache()
         try:
             torch.accelerator.empty_host_cache()
@@ -2371,8 +2470,7 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
 def in_the_same_node_as(
     pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0
 ) -> list[bool]:
-    """
-    This is a collective operation that returns if each rank is in the same node
+    """This is a collective operation that returns if each rank is in the same node
     as the source rank. It tests if processes are attached to the same
     memory system (shared access to shared memory).
     """
@@ -2463,8 +2561,7 @@ def in_the_same_node_as(
 
 
 def is_global_first_rank() -> bool:
-    """
-    Check if the current process is the first rank globally across all
+    """Check if the current process is the first rank globally across all
     parallelism strategies (PP, TP, DP, EP, etc.).
 
     Unlike group-specific checks like `get_tensor_model_parallel_rank() == 0`
@@ -2474,6 +2571,7 @@ def is_global_first_rank() -> bool:
     Returns:
         bool: True if this is the global first rank (rank 0), False otherwise.
               Returns True if distributed is not initialized (single process).
+
     """
     try:
         # If world group is available, use it for the most accurate check
@@ -2494,9 +2592,7 @@ def is_global_first_rank() -> bool:
 
 
 def is_local_first_rank() -> bool:
-    """
-    Check if the current process is the first local rank (rank 0 on its node).
-    """
+    """Check if the current process is the first local rank (rank 0 on its node)."""
     try:
         # prefer the initialized world group if available
         global _WORLD
@@ -2517,14 +2613,14 @@ def is_local_first_rank() -> bool:
 
 
 def _node_count(pg: ProcessGroup | StatelessProcessGroup) -> int:
-    """
-    Returns the total number of nodes in the process group.
+    """Returns the total number of nodes in the process group.
 
     Args:
         pg: The process group to analyze
 
     Returns:
         int: The total number of nodes
+
     """
     if isinstance(pg, ProcessGroup):
         world_size = torch.distributed.get_world_size(group=pg)
diff --git a/vllm/distributed/stateless_coordinator.py b/vllm/distributed/stateless_coordinator.py
index 38c74a97c55c..6ce2aa89ac75 100644
--- a/vllm/distributed/stateless_coordinator.py
+++ b/vllm/distributed/stateless_coordinator.py
@@ -59,8 +59,7 @@ def _fetch_group_ports(key: str, coord_store: Store) -> list[int]:
 
 
 class StatelessGroupCoordinator(GroupCoordinator):
-    """
-    A stateless version of the GroupCoordinator class in parallel_state,
+    """A stateless version of the GroupCoordinator class in parallel_state,
     It will create CPU, device and TCPStore based communication groups
     that are independent of PyTorch's WORLD group. Hence,
     communication groups with a different set of participants GPUs
diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py
index eec9890f0286..b32269ec2b60 100644
--- a/vllm/distributed/utils.py
+++ b/vllm/distributed/utils.py
@@ -111,6 +111,7 @@ def split_tensor_along_last_dim(
 
     Returns:
         A list of Tensors
+
     """
     # Get the size and dimension.
     last_dim = tensor.dim() - 1
@@ -333,7 +334,6 @@ def all_reduce(
     def barrier(self, timeout: float = 30.0):
         """A robust barrier to synchronize all ranks.
 
-
         Uses a multi-phase approach to ensure all processes reach the barrier
         before proceeding:
 
@@ -351,6 +351,7 @@ def barrier(self, timeout: float = 30.0):
 
         Raises:
             RuntimeError: If coordination fails or times out
+
         """
         # Generate a barrier ID that is globally unique
         try:
@@ -549,8 +550,7 @@ def init_gloo_process_group(
     group_size: int,
     timeout: timedelta,
 ) -> ProcessGroup:
-    """
-    Stateless init ProcessGroup with gloo backend compatible with
+    """Stateless init ProcessGroup with gloo backend compatible with
     different torch versions.
     """
     with suppress_stdout():
@@ -583,8 +583,7 @@ def stateless_init_torch_distributed_process_group(
     return_store: bool = False,
     listen_socket: socket.socket | None = None,
 ) -> ProcessGroup | tuple[ProcessGroup, Store]:
-    """
-    A replacement for `torch.distributed.init_process_group` that does not
+    """A replacement for `torch.distributed.init_process_group` that does not
     pollute the global state. The created ProcessGroup object can be used for
     some operations such as `allreduce`, because it does not depend on the
     global rank. However, some operations such as `broadcast` cannot be used
@@ -685,9 +684,8 @@ def stateless_init_torch_distributed_process_group(
 
 
 def stateless_destroy_torch_distributed_process_group(pg: ProcessGroup) -> None:
-    """
-    Destroy ProcessGroup returned by
-        stateless_init_torch_distributed_process_group().
+    """Destroy ProcessGroup returned by
+    stateless_init_torch_distributed_process_group().
     """
     pg.shutdown()
     _unregister_process_group(pg.group_name)
@@ -706,6 +704,7 @@ def get_worker_rank_suffix(global_rank: int | None = None) -> str:
     Returns:
         A string suffix identifying the worker's position in the
         distributed topology.
+
     """
     from vllm.distributed.parallel_state import (
         get_dcp_group,
diff --git a/vllm/distributed/weight_transfer/__init__.py b/vllm/distributed/weight_transfer/__init__.py
index c78fd1f3cc5f..3fcca3282691 100644
--- a/vllm/distributed/weight_transfer/__init__.py
+++ b/vllm/distributed/weight_transfer/__init__.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Weight transfer engines for syncing model weights from trainers
+"""Weight transfer engines for syncing model weights from trainers
 to inference workers.
 """
 
diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py
index 2328829fc0bc..2fb60fb71b41 100644
--- a/vllm/distributed/weight_transfer/base.py
+++ b/vllm/distributed/weight_transfer/base.py
@@ -193,6 +193,7 @@ def held_names(self) -> "Collection[str] | None":
 
         Returns:
             The held parameter names, or None to hold every one.
+
         """
         return None
 
@@ -329,8 +330,7 @@ class WeightTransferUpdateRequest:
 
 
 class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
-    """
-    Base class for weight transfer engines that handle transport of model weights
+    """Base class for weight transfer engines that handle transport of model weights
     from a trainer to inference workers.
 
     This abstraction separates weight transfer transport logic from the worker
@@ -386,14 +386,14 @@ def __init__(
         device: torch.device,
         model: torch.nn.Module,
     ) -> None:
-        """
-        Initialize the weight transfer engine.
+        """Initialize the weight transfer engine.
 
         Args:
             config: The configuration for the weight transfer engine
             vllm_config: The full vLLM config (provides parallel/model config)
             device: The device this worker's model lives on
             model: The local model instance which will receive the weights
+
         """
         self.config = config
         self.vllm_config = vllm_config
@@ -419,8 +419,7 @@ def reset_weight_update_target(self) -> None:
         self.model_config = self._default_model_config
 
     def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
-        """
-        Construct typed init info from dict with validation.
+        """Construct typed init info from dict with validation.
 
         Args:
             init_dict: Dictionary containing backend-specific initialization parameters
@@ -430,6 +429,7 @@ def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
 
         Raises:
             ValueError: If init_dict is invalid for this backend
+
         """
         try:
             return self.init_info_cls(**init_dict)
@@ -439,8 +439,7 @@ def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
             ) from e
 
     def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo:
-        """
-        Construct typed update info from dict with validation.
+        """Construct typed update info from dict with validation.
 
         Args:
             update_dict: Dictionary containing backend-specific update parameters
@@ -450,6 +449,7 @@ def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo:
 
         Raises:
             ValueError: If update_dict is invalid for this backend
+
         """
         try:
             return self.update_info_cls(**update_dict)
@@ -460,19 +460,18 @@ def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo:
 
     @abstractmethod
     def init_transfer_engine(self, init_info: TInitInfo) -> None:
-        """
-        Initialize the weight transfer mechanism.
+        """Initialize the weight transfer mechanism.
         This is called once at the beginning of training.
 
         Args:
             init_info: Backend-specific initialization info
+
         """
         raise NotImplementedError
 
     @abstractmethod
     def start_weight_update(self) -> None:
-        """
-        Prepare the engine for a new weight update.
+        """Prepare the engine for a new weight update.
 
         Engines that receive weights in checkpoint format initialize layerwise reloading
         here, else this is typically a no-op.
@@ -482,8 +481,7 @@ def start_weight_update(self) -> None:
 
     @abstractmethod
     def finish_weight_update(self) -> None:
-        """
-        Finalize the current weight update.
+        """Finalize the current weight update.
 
         Checkpoint-format engines finalize layerwise reloading here; engines
         that apply weights in place leave this as a no-op.
@@ -491,11 +489,11 @@ def finish_weight_update(self) -> None:
         raise NotImplementedError
 
     def update_weights(self, update_info: dict[str, Any]) -> None:
-        """
-        Receive one weight update chunk and load it into the model.
+        """Receive one weight update chunk and load it into the model.
 
         Args:
             update_info: Dictionary containing backend-specific update info
+
         """
         typed_update_info = self.parse_update_info(update_info)
         self.receive_weights(typed_update_info)
@@ -505,19 +503,18 @@ def update_weights(self, update_info: dict[str, Any]) -> None:
 
     @abstractmethod
     def receive_weights(self, update_info: TUpdateInfo) -> None:
-        """
-        Receive weights from the trainer and load them into the model.
+        """Receive weights from the trainer and load them into the model.
 
         Args:
             update_info: Backend-specific update info containing parameter metadata
                         and any backend-specific data
+
         """
         raise NotImplementedError
 
     @abstractmethod
     def shutdown(self) -> None:
-        """
-        Shutdown the weight transfer engine.
+        """Shutdown the weight transfer engine.
         This should be called when the worker is shutting down.
         """
         raise NotImplementedError
diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py
index ba486422968f..48f9e4f21281 100644
--- a/vllm/distributed/weight_transfer/factory.py
+++ b/vllm/distributed/weight_transfer/factory.py
@@ -58,6 +58,7 @@ def register_engine(
 
         Raises:
             ValueError: If an engine with the same name is already registered
+
         """
         if name in cls._registry:
             raise ValueError(f"Weight transfer engine '{name}' is already registered.")
@@ -101,6 +102,7 @@ def create_engine(
 
         Raises:
             ValueError: If the backend is not registered
+
         """
         backend = config.backend
         if backend not in cls._registry:
@@ -190,6 +192,7 @@ def trainer_init(
 
         Raises:
             ValueError: If `init_info.backend` is not registered.
+
         """
         backend = init_info.backend
         if backend not in cls._registry:
diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py
index 7b1563112d19..15864811e96c 100644
--- a/vllm/distributed/weight_transfer/ipc_engine.py
+++ b/vllm/distributed/weight_transfer/ipc_engine.py
@@ -121,8 +121,7 @@ def __post_init__(self):
 class IPCWeightTransferEngine(
     WeightTransferEngine[IPCWeightTransferInitInfo, IPCWeightTransferUpdateInfo]
 ):
-    """
-    Weight transfer engine using CUDA IPC for communication between trainer and workers.
+    """Weight transfer engine using CUDA IPC between trainer and workers.
 
     This implementation uses CUDA IPC to transfer weights from the trainer (rank 0)
     to all inference workers in a process group. IPC handles are used to share
@@ -149,13 +148,13 @@ def __init__(
         self._packed_importer = PackedBufferImporter()
 
     def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None:
-        """
-        Initialize the weight transfer mechanism. No data-plane rendezvous is
+        """Initialize the weight transfer mechanism. No data-plane rendezvous is
         needed for IPC; this just records the trainer-supplied wire params so the
         worker decodes exactly as the trainer encoded.
 
         Args:
             init_info: IPC initialization info (carries `packed`).
+
         """
         self.packed = init_info.packed
 
@@ -182,8 +181,7 @@ def finish_weight_update(self) -> None:
         self._packed_importer.close()
 
     def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None:
-        """
-        Receive weights from the trainer via CUDA IPC handles and load them.
+        """Receive weights from the trainer via CUDA IPC handles and load them.
 
         Whether the transfer is packed is read from `self.packed`, set at the
         init handshake from the trainer's init info, so it is guaranteed to
@@ -193,6 +191,7 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None:
             update_info: IPC update info containing parameter names, dtypes, shapes,
                         and IPC handles. Each IPC handle is a mapping between physical
                         GPU UUID and the rebuild_cuda_tensor args tuple.
+
         """
         # Use the worker's assigned device rather than the ambient current
         # device: the receive path is no longer wrapped in
diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py
index 67e578640a0d..57a17b9830a4 100644
--- a/vllm/distributed/weight_transfer/nccl_common.py
+++ b/vllm/distributed/weight_transfer/nccl_common.py
@@ -142,8 +142,7 @@ def stateless_init_process_group(
     world_size: int,
     device,
 ) -> "PyNcclCommunicator":
-    """
-    vLLM provides `StatelessProcessGroup` to create a process group
+    """VLLM provides `StatelessProcessGroup` to create a process group
     without considering the global process group in torch.distributed.
     It is recommended to create `StatelessProcessGroup`, and then initialize
     the data-plane communication (NCCL) between external (train processes)
@@ -222,8 +221,7 @@ def worker_init_process_group(
 def trainer_init(
     init_info: NCCLRendezvous | dict,
 ) -> "PyNcclCommunicator":
-    """
-    Initialize NCCL process group for trainer-side weight transfer.
+    """Initialize NCCL process group for trainer-side weight transfer.
 
     The trainer is always rank 0 in the process group. Uses the current
     CUDA device (torch.accelerator.current_device_index()).
@@ -237,6 +235,7 @@ def trainer_init(
 
     Returns:
         PyNcclCommunicator for weight transfer.
+
     """
     if isinstance(init_info, dict):
         master_address = init_info["master_address"]
diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py
index 1efd8180dfd5..5b98845870fd 100644
--- a/vllm/distributed/weight_transfer/nccl_engine.py
+++ b/vllm/distributed/weight_transfer/nccl_engine.py
@@ -113,8 +113,7 @@ def __post_init__(self):
 class NCCLWeightTransferEngine(
     WeightTransferEngine[NCCLWeightTransferInitInfo, NCCLWeightTransferUpdateInfo]
 ):
-    """
-    Weight transfer engine using NCCL for communication between trainer and workers.
+    """Weight transfer engine using NCCL for communication between trainer and workers.
 
     This implementation uses NCCL broadcast operations to transfer dense
     checkpoint-format weights from the trainer (rank 0) to all inference workers
@@ -142,14 +141,14 @@ def __init__(
         self.packed_num_buffers = DEFAULT_PACKED_NUM_BUFFERS
 
     def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
-        """
-        Initialize NCCL process group with the trainer and record the
+        """Initialize NCCL process group with the trainer and record the
         trainer-supplied wire params so the worker decodes exactly as the
         trainer encodes.
 
         Args:
             init_info: NCCL initialization info containing master address, port,
                       rank offset, world size, and the packed wire params
+
         """
         self.packed = init_info.packed
         self.packed_buffer_size_bytes = init_info.packed_buffer_size_bytes
@@ -175,8 +174,7 @@ def finish_weight_update(self) -> None:
         finalize_layerwise_reload(self.model, self.model_config)
 
     def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None:
-        """
-        Receive weights from trainer via NCCL broadcast.
+        """Receive weights from trainer via NCCL broadcast.
 
         Whether to use packed broadcasting (and the buffer geometry) is read
         from `self.packed` / `self.packed_*`, set at the init handshake from the
@@ -186,6 +184,7 @@ def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None:
         Args:
             update_info: NCCL update info containing parameter names, dtypes,
                         and shapes
+
         """
         if self.model_update_group is None:
             raise RuntimeError(
diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py
index 0eb9ede2639b..dcd8efc8c4af 100644
--- a/vllm/distributed/weight_transfer/packed_tensor.py
+++ b/vllm/distributed/weight_transfer/packed_tensor.py
@@ -45,6 +45,7 @@ def unpack_tensor(
         shapes: List of tensor shapes
         dtypes: List of tensor dtypes
         tensor_sizes: List of tensor sizes in bytes
+
     """
     unpacked_tensors = packed_tensor.split(tensor_sizes)
 
@@ -85,6 +86,7 @@ def pack_tensors(
         tensor_list: Pre-existing tensor list to append to (for NCCL
                     multi-buffer reuse). If None, a fresh list is created.
         current_size: Byte count already accumulated in tensor_list
+
     """
     if tensor_list is None:
         tensor_list = []
@@ -335,6 +337,7 @@ def packed_ipc_producer(
         buffer_size_bytes: Exact capacity of the reusable IPC buffer.
             Every chunk is guaranteed to fit within this size.  A
             ``ValueError`` is raised if any single tensor exceeds it.
+
     """
     ipc_buffer = torch.empty(buffer_size_bytes, dtype=torch.uint8, device="cuda")
     _, ipc_args = reduce_tensor(ipc_buffer)
@@ -463,6 +466,7 @@ def packed_ipc_consumer(
             export so the buffer mapping is opened and released exactly once
             (see ``PackedBufferImporter``). ``None`` builds an ephemeral one,
             which is only safe for single-chunk consumption.
+
     """
     props = torch.cuda.get_device_properties(device_index)
     physical_gpu_id = str(props.uuid)
diff --git a/vllm/distributed/weight_transfer/sharded_rdt_common.py b/vllm/distributed/weight_transfer/sharded_rdt_common.py
index a5a7ed923369..0d913628174a 100644
--- a/vllm/distributed/weight_transfer/sharded_rdt_common.py
+++ b/vllm/distributed/weight_transfer/sharded_rdt_common.py
@@ -60,6 +60,7 @@ def check_ray_rdt_version() -> None:
 
     Raises:
         ValueError: the installed Ray predates ``RDT_MIN_RAY_VERSION``.
+
     """
     import importlib.metadata
 
@@ -77,6 +78,16 @@ def check_ray_rdt_version() -> None:
         )
 
 
+def register_nixl_memory(tensor: torch.Tensor) -> None:
+    """Register a tensor with Ray using the platform's NIXL package."""
+    from vllm.distributed.nixl_utils import alias_nixl_for_ray
+
+    alias_nixl_for_ray()
+    from ray.experimental import register_nixl_memory as ray_register_nixl_memory
+
+    ray_register_nixl_memory(tensor)
+
+
 def assign_producer_indices(
     num_producers: int, num_consumers: int, consumer_idx: int
 ) -> list[int]:
@@ -152,6 +163,7 @@ def __init__(
             ValueError: ``workers_per_replica`` does not divide
                 ``num_consumers``, so the deployments are not uniform and two
                 workers of one deployment would share a block index.
+
         """
         self.num_producers = max(1, num_producers)
         self.num_consumers = max(1, num_consumers)
@@ -253,6 +265,7 @@ def validate(self) -> None:
         Raises:
             ValueError: an owner set is empty or out of range, or a class index
                 does not resolve.
+
         """
         for c, row in enumerate(self._owner_sets):
             if not row:
diff --git a/vllm/distributed/weight_transfer/sharded_rdt_engine.py b/vllm/distributed/weight_transfer/sharded_rdt_engine.py
index 2b8187007b8f..96dd53f0fdb9 100644
--- a/vllm/distributed/weight_transfer/sharded_rdt_engine.py
+++ b/vllm/distributed/weight_transfer/sharded_rdt_engine.py
@@ -66,6 +66,7 @@
     RdtRouter,
     buffer_alloc_bytes,
     check_ray_rdt_version,
+    register_nixl_memory,
 )
 from vllm.distributed.weight_transfer.sharded_rdt_fake import (
     BakeSink,
@@ -588,8 +589,6 @@ def _preregister_at_init(self) -> None:
         plan = self._cached_plan
         if plan is None or not plan.chunks:
             return
-        from ray.experimental import register_nixl_memory
-
         # (a) consumer receive buffers — one per ring slot, at the largest chunk.
         max_pack = max(c.pack_bytes for c in plan.chunks)
         alloc = buffer_alloc_bytes(max_pack, self._buffer_presize)
@@ -740,7 +739,7 @@ def _issue_pull(self, chunk: "_Chunk", slot: int) -> "_PendingPull":
         any time after the metadata push. See the doc's "slot generation
         handshake".
         """
-        from ray.experimental import register_nixl_memory, set_target_for_ref
+        from ray.experimental import set_target_for_ref
 
         if self._slot_read_done:
             with self._slot_cv:
diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py
index 60d0e60fb02d..6b755e772e18 100644
--- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py
+++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py
@@ -47,6 +47,7 @@
     ALLOWED_OPS,
     buffer_alloc_bytes,
     check_ray_rdt_version,
+    register_nixl_memory,
 )
 from vllm.logger import init_logger
 
@@ -334,8 +335,6 @@ def warmup_nixl(self) -> None:
         doc). The warmup buffer stays registered so the agent's CUDA-HMEM path
         stays initialized.
         """
-        from ray.experimental import register_nixl_memory
-
         self._nixl_warmup_buf = torch.zeros(1 << 20, dtype=torch.uint8, device="cuda")
         with self._reg_lock:
             register_nixl_memory(self._nixl_warmup_buf)
@@ -510,8 +509,6 @@ def _new_serve_buffer(self, nbytes: int) -> torch.Tensor:
         """Allocate + NIXL-register one serve slot: the single allocation seam,
         so registration cannot be skipped on either of the two paths that make
         buffers (the init-time reservation and the serve-path backstop)."""
-        from ray.experimental import register_nixl_memory
-
         t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device)
         with self._reg_lock:
             register_nixl_memory(t)
@@ -537,6 +534,7 @@ def reserve_serve_buffer(
         Raises:
             RuntimeError: two consumers of one sharing group disagree on the
                 chunks they pull from this producer.
+
         """
         sg = self._share_group(consumer_id)
         if plan_digest is not None:
@@ -684,6 +682,7 @@ def rdt_produce_weights_batched(
 
         Raises:
             ValueError: ``seq`` was not supplied.
+
         """
         needed = sorted({n for n, _ in specs})
         if self._served_names is not None:
@@ -1026,6 +1025,7 @@ def get_worker_init_payload(self) -> dict:
 
         Raises:
             RuntimeError: called before ``trainer_init`` cached the server names.
+
         """
         if self._server_names is None:
             raise RuntimeError(
diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py
index 9f91da3418da..9c704ab344ea 100644
--- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py
+++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py
@@ -113,8 +113,7 @@ def __post_init__(self) -> None:
 class SparseNCCLWeightTransferEngine(
     WeightTransferEngine[NCCLWeightTransferInitInfo, SparseNCCLWeightTransferUpdateInfo]
 ):
-    """
-    Sparse weight transfer engine using NCCL.
+    """Sparse weight transfer engine using NCCL.
 
     Receives checkpoint-coordinate patches broadcast from the trainer and applies
     them through the model's native weight loader. Sparse updates modify initialized
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
index bf5cc62d367d..f676fc9f1891 100644
--- a/vllm/engine/arg_utils.py
+++ b/vllm/engine/arg_utils.py
@@ -10,7 +10,6 @@
 import sys
 from collections.abc import Callable
 from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass
-from itertools import permutations
 from types import UnionType
 from typing import (
     TYPE_CHECKING,
@@ -23,6 +22,7 @@
     cast,
     get_args,
     get_origin,
+    is_typeddict,
 )
 
 import huggingface_hub
@@ -35,6 +35,7 @@
 import vllm.envs as envs
 from vllm.config import (
     AttentionConfig,
+    AuxOutputConfig,
     CacheConfig,
     CompilationConfig,
     ConfigType,
@@ -76,7 +77,12 @@
     PrefixCachingHashAlgo,
 )
 from vllm.config.device import Device
-from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend
+from vllm.config.kernel import (
+    IrOpPriorityConfig,
+    LinearBackend,
+    MoEBackend,
+    SparseIndexerTopkBackend,
+)
 from vllm.config.load import SafetensorsLoadStrategy
 from vllm.config.lora import MaxLoRARanks
 from vllm.config.mamba import MambaBackendEnum, MambaSSUAlgorithm
@@ -179,6 +185,18 @@ def is_type(type_hint: TypeHint, type: TypeHintT) -> TypeIs[TypeHintT]:
     return type_hint is type or get_origin(type_hint) is type
 
 
+def is_dict_subclass(type_hint: TypeHint) -> bool:
+    """Check if the type hint is a subclass of `dict`.
+
+    `TypedDict`s are excluded because they do not support class checks.
+    """
+    return (
+        isinstance(type_hint, type)
+        and not is_typeddict(type_hint)
+        and issubclass(type_hint, dict)
+    )
+
+
 def contains_type(type_hints: set[TypeHint], type: TypeHintT) -> bool:
     """Check if the type hints contain a specific type."""
     return any(is_type(type_hint, type) for type_hint in type_hints)
@@ -301,6 +319,8 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:
     for field in fields(cls):
         # Get the set of possible types for the field
         type_hints: set[TypeHint] = get_type_hints(field.type)
+        # Subclasses of dict (e.g. MultiModalDummyOptions) are CLI dicts
+        type_hints = {dict if is_dict_subclass(th) else th for th in type_hints}
 
         # If the field is a dataclass, we can use the model_validate_json
         generator = (th for th in type_hints if is_dataclass(th))
@@ -406,7 +426,7 @@ def parse_dataclass(val: str, cls=dataclass_cls) -> Any:
         if type(None) in type_hints and not contains_type(type_hints, bool):
             kwargs[name]["type"] = optional_type(kwargs[name]["type"])
             if kwargs[name].get("choices"):
-                kwargs[name]["choices"].append("None")
+                kwargs[name]["choices"].append(None)
     return kwargs
 
 
@@ -428,7 +448,9 @@ class EngineArgs:
     """Arguments for vLLM engine."""
 
     model: str = ModelConfig.model
-    enable_return_routed_experts: bool = ModelConfig.enable_return_routed_experts
+    # Public compatibility argument. Canonical runtime state lives in
+    # AuxOutputConfig.
+    enable_return_routed_experts: bool = False
     return_sampling_mask: bool = ModelConfig.return_sampling_mask
     model_weights: str = ModelConfig.model_weights
     served_model_name: str | list[str] | None = ModelConfig.served_model_name
@@ -487,7 +509,7 @@ class EngineArgs:
     dcp_comm_backend: DCPCommBackend | None = ParallelConfig.dcp_comm_backend
     dcp_q_replicate: bool | None = ParallelConfig.dcp_q_replicate
     dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
-    cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
+    cp_kv_cache_interleave_size: int | None = None
     data_parallel_size: int = ParallelConfig.data_parallel_size
     data_parallel_rank: int | None = None
     data_parallel_start_rank: int | None = None
@@ -505,8 +527,12 @@ class EngineArgs:
     enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter
     moe_backend: MoEBackend = KernelConfig.moe_backend
     linear_backend: LinearBackend = KernelConfig.linear_backend
+    sparse_indexer_topk_backend: SparseIndexerTopkBackend = (
+        KernelConfig.sparse_indexer_topk_backend
+    )
     all2all_backend: All2AllBackend = ParallelConfig.all2all_backend
     enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep
+    elastic_ep_max_dp_size: int = ParallelConfig.elastic_ep_max_dp_size
     enable_dbo: bool = ParallelConfig.enable_dbo
     ubatch_size: int = ParallelConfig.ubatch_size
     dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold
@@ -547,7 +573,11 @@ class EngineArgs:
     max_num_batched_tokens: int | None = None
     max_num_scheduled_tokens: int | None = None
     long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold
+    long_prefill_token_threshold_adaptive: bool = (
+        SchedulerConfig.long_prefill_token_threshold_adaptive
+    )
     max_num_seqs: int | None = None
+    max_num_active_seqs: int | None = SchedulerConfig.max_num_active_seqs
     max_num_queued_reqs: int | None = None
     max_num_queued_tokens: int | None = None
     max_logprobs: int = ModelConfig.max_logprobs
@@ -650,6 +680,7 @@ class EngineArgs:
     structured_outputs_config: StructuredOutputsConfig = get_field(
         VllmConfig, "structured_outputs_config"
     )
+    aux_output_config: AuxOutputConfig = get_field(VllmConfig, "aux_output_config")
     reasoning_parser: str = StructuredOutputsConfig.reasoning_parser
     reasoning_parser_plugin: str | None = None
 
@@ -713,6 +744,9 @@ class EngineArgs:
 
     generation_config: str = ModelConfig.generation_config
     enable_sleep_mode: bool = ModelConfig.enable_sleep_mode
+    sleep_preserve_parameter_names: list[str] = get_field(
+        ModelConfig, "sleep_preserve_parameter_names"
+    )
     enable_cumem_allocator: bool = ModelConfig.enable_cumem_allocator
     enable_nccl_comm_suspend: bool = ModelConfig.enable_nccl_comm_suspend
     override_generation_config: dict[str, Any] = get_field(
@@ -729,8 +763,8 @@ class EngineArgs:
     mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")
     prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit")
     mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode
-    enable_mamba_fine_grained_prefix_cache: bool = (
-        CacheConfig.enable_mamba_fine_grained_prefix_cache
+    enable_mamba_shared_prefix_checkpoint: bool = (
+        CacheConfig.enable_mamba_shared_prefix_checkpoint
     )
     replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len
     use_replayssm: bool = CacheConfig.use_replayssm
@@ -757,6 +791,7 @@ class EngineArgs:
     stream_interval: int = SchedulerConfig.stream_interval
 
     kv_sharing_fast_prefill: bool = CacheConfig.kv_sharing_fast_prefill
+    swa_bounded_replay: bool = CacheConfig.swa_bounded_replay
     optimization_level: OptimizationLevel = VllmConfig.optimization_level
     performance_mode: PerformanceMode = VllmConfig.performance_mode
 
@@ -791,6 +826,10 @@ def __post_init__(self):
             self.compilation_config = CompilationConfig(**self.compilation_config)
         if isinstance(self.attention_config, dict):
             self.attention_config = AttentionConfig(**self.attention_config)
+        if isinstance(self.aux_output_config, dict):
+            self.aux_output_config = AuxOutputConfig(**self.aux_output_config)
+        if self.enable_return_routed_experts:
+            self.aux_output_config.enable_return_routed_experts = True
         if isinstance(self.engram_config, dict):
             self.engram_config = EngramConfig(**self.engram_config)
         if isinstance(self.mamba_config, dict):
@@ -857,9 +896,9 @@ def __post_init__(self):
     @staticmethod
     def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         """Shared CLI arguments for vLLM engine."""
-
         # Model arguments
         model_kwargs = get_kwargs(ModelConfig)
+        aux_output_kwargs = get_kwargs(AuxOutputConfig)
         model_group = parser.add_argument_group(
             title="ModelConfig",
             description=ModelConfig.__doc__,
@@ -899,7 +938,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         model_group.add_argument("--enforce-eager", **model_kwargs["enforce_eager"])
         model_group.add_argument(
             "--enable-return-routed-experts",
-            **model_kwargs["enable_return_routed_experts"],
+            **aux_output_kwargs["enable_return_routed_experts"],
         )
         model_group.add_argument(
             "--return-sampling-mask",
@@ -942,6 +981,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         model_group.add_argument(
             "--enable-sleep-mode", **model_kwargs["enable_sleep_mode"]
         )
+        model_group.add_argument(
+            "--sleep-preserve-parameter-names",
+            **model_kwargs["sleep_preserve_parameter_names"],
+        )
         model_group.add_argument(
             "--enable-cumem-allocator", **model_kwargs["enable_cumem_allocator"]
         )
@@ -1104,7 +1147,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         )
         parallel_group.add_argument(
             "--cp-kv-cache-interleave-size",
-            **parallel_kwargs["cp_kv_cache_interleave_size"],
+            **{
+                **parallel_kwargs["cp_kv_cache_interleave_size"],
+                "default": None,
+            },
         )
         parallel_group.add_argument(
             "--prefill-context-parallel-size",
@@ -1198,6 +1244,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parallel_group.add_argument(
             "--enable-elastic-ep", **parallel_kwargs["enable_elastic_ep"]
         )
+        parallel_group.add_argument(
+            "--elastic-ep-max-dp-size",
+            **parallel_kwargs["elastic_ep_max_dp_size"],
+        )
         parallel_group.add_argument(
             "--dbo-decode-token-threshold",
             **parallel_kwargs["dbo_decode_token_threshold"],
@@ -1251,7 +1301,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         )
         cache_group.add_argument("--block-size", **cache_kwargs["block_size"])
         cache_group.add_argument(
-            "--gpu-memory-utilization", **cache_kwargs["gpu_memory_utilization"]
+            "--gpu-memory-utilization",
+            "--device-memory-utilization",
+            dest="gpu_memory_utilization",
+            **cache_kwargs["gpu_memory_utilization"],
         )
         cache_group.add_argument(
             "--kv-cache-memory-bytes", **cache_kwargs["kv_cache_memory_bytes"]
@@ -1280,6 +1333,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         cache_group.add_argument(
             "--kv-sharing-fast-prefill", **cache_kwargs["kv_sharing_fast_prefill"]
         )
+        cache_group.add_argument(
+            "--swa-bounded-replay", **cache_kwargs["swa_bounded_replay"]
+        )
         cache_group.add_argument(
             "--mamba-cache-dtype", **cache_kwargs["mamba_cache_dtype"]
         )
@@ -1296,8 +1352,8 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]
         )
         cache_group.add_argument(
-            "--enable-mamba-fine-grained-prefix-cache",
-            **cache_kwargs["enable_mamba_fine_grained_prefix_cache"],
+            "--enable-mamba-shared-prefix-checkpoint",
+            **cache_kwargs["enable_mamba_shared_prefix_checkpoint"],
         )
         cache_group.add_argument(
             "--replayssm-buffer-len", **cache_kwargs["replayssm_buffer_len"]
@@ -1508,13 +1564,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         observability_group.add_argument(
             "--otlp-traces-endpoint", **observability_kwargs["otlp_traces_endpoint"]
         )
-        # TODO: generalise this special case
-        choices = observability_kwargs["collect_detailed_traces"]["choices"]
-        metavar = f"{{{','.join(choices)}}}"
-        observability_kwargs["collect_detailed_traces"]["metavar"] = metavar
-        observability_kwargs["collect_detailed_traces"]["choices"] += [
-            ",".join(p) for p in permutations(get_args(DetailedTraceModules), r=2)
-        ]
         observability_group.add_argument(
             "--collect-detailed-traces",
             **observability_kwargs["collect_detailed_traces"],
@@ -1582,6 +1631,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
                 "default": None,
             },
         )
+        scheduler_group.add_argument(
+            "--max-num-active-seqs",
+            **scheduler_kwargs["max_num_active_seqs"],
+        )
         scheduler_group.add_argument(
             "--max-num-queued-reqs", **scheduler_kwargs["max_num_queued_reqs"]
         )
@@ -1593,6 +1646,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             "--long-prefill-token-threshold",
             **scheduler_kwargs["long_prefill_token_threshold"],
         )
+        scheduler_group.add_argument(
+            "--long-prefill-token-threshold-adaptive",
+            **scheduler_kwargs["long_prefill_token_threshold_adaptive"],
+        )
         # multi-step scheduling has been removed; corresponding arguments
         # are no longer supported.
         scheduler_group.add_argument(
@@ -1662,6 +1719,11 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         linear_backend_kwargs = kernel_kwargs["linear_backend"]
         linear_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_")
         kernel_group.add_argument("--linear-backend", **linear_backend_kwargs)
+        sparse_indexer_topk_kwargs = kernel_kwargs["sparse_indexer_topk_backend"]
+        sparse_indexer_topk_kwargs["type"] = lambda s: s.lower().replace("-", "_")
+        kernel_group.add_argument(
+            "--sparse-indexer-topk-backend", **sparse_indexer_topk_kwargs
+        )
 
         # vLLM arguments
         vllm_kwargs = get_kwargs(VllmConfig)
@@ -1713,6 +1775,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         vllm_group.add_argument(
             "--structured-outputs-config", **vllm_kwargs["structured_outputs_config"]
         )
+        vllm_group.add_argument(
+            "--aux-output-config", **vllm_kwargs["aux_output_config"]
+        )
         vllm_group.add_argument("--profiler-config", **vllm_kwargs["profiler_config"])
         vllm_group.add_argument(
             "--optimization-level", **vllm_kwargs["optimization_level"]
@@ -1820,7 +1885,6 @@ def create_model_config(self) -> ModelConfig:
             quantization_config=self.quantization_config,
             allow_deprecated_quantization=self.allow_deprecated_quantization,
             enforce_eager=self.enforce_eager,
-            enable_return_routed_experts=self.enable_return_routed_experts,
             return_sampling_mask=self.return_sampling_mask,
             max_logprobs=self.max_logprobs,
             logprobs_mode=self.logprobs_mode,
@@ -1854,6 +1918,7 @@ def create_model_config(self) -> ModelConfig:
             generation_config=self.generation_config,
             override_generation_config=self.override_generation_config,
             enable_sleep_mode=self.enable_sleep_mode,
+            sleep_preserve_parameter_names=self.sleep_preserve_parameter_names,
             enable_cumem_allocator=self.enable_cumem_allocator,
             enable_nccl_comm_suspend=self.enable_nccl_comm_suspend,
             model_impl=self.model_impl,
@@ -2021,8 +2086,7 @@ def create_engine_config(
         usage_context: UsageContext | None = None,
         headless: bool = False,
     ) -> VllmConfig:
-        """
-        Create the VllmConfig.
+        """Create the VllmConfig.
 
         NOTE: If VllmConfig is incompatible, we raise an error.
         """
@@ -2030,7 +2094,7 @@ def create_engine_config(
 
         device_config = DeviceConfig(device=cast(Device, current_platform.device_type))
 
-        envs.validate_environ(self.fail_on_environ_validation)
+        current_platform.validate_environ(self.fail_on_environ_validation)
 
         # Check if the model is a speculator and override model/tokenizer/config
         # BEFORE creating ModelConfig, so the config is created with the target model
@@ -2087,13 +2151,14 @@ def create_engine_config(
             prefix_cache_retention_interval=self.prefix_cache_retention_interval,
             kv_cache_dtype_skip_layers=self.kv_cache_dtype_skip_layers,
             kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,
+            swa_bounded_replay=self.swa_bounded_replay,
             mamba_cache_dtype=self.mamba_cache_dtype,
             mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,
             mamba_block_size=self.mamba_block_size,
             prefix_match_unit=self.prefix_match_unit,
             mamba_cache_mode=self.mamba_cache_mode,
-            enable_mamba_fine_grained_prefix_cache=(
-                self.enable_mamba_fine_grained_prefix_cache
+            enable_mamba_shared_prefix_checkpoint=(
+                self.enable_mamba_shared_prefix_checkpoint
             ),
             replayssm_buffer_len=self.replayssm_buffer_len,
             use_replayssm=self.use_replayssm,
@@ -2185,12 +2250,22 @@ def create_engine_config(
                 self.node_rank * local_world_size
             ) // world_size_within_dp
             if self.data_parallel_size > 1 and self.data_parallel_external_lb:
-                self.data_parallel_rank = inferred_data_parallel_rank
-                logger.info(
-                    "Inferred data_parallel_rank %d from node_rank %d for external lb",
-                    self.data_parallel_rank,
-                    self.node_rank,
-                )
+                if self.data_parallel_rank is None:
+                    if self.nnodes % self.data_parallel_size != 0:
+                        raise ValueError(
+                            "Invalid data-parallel launch options: "
+                            "`--node-rank` cannot unambiguously identify external "
+                            "data-parallel ranks when `--nnodes` is not divisible "
+                            "by `--data-parallel-size`. Set a unique "
+                            "`--data-parallel-rank` for each external-LB process."
+                        )
+                    self.data_parallel_rank = inferred_data_parallel_rank
+                    logger.info(
+                        "Inferred data_parallel_rank %d from node_rank %d "
+                        "for external lb",
+                        self.data_parallel_rank,
+                        self.node_rank,
+                    )
             elif self.data_parallel_size_local is None:
                 # Infer data parallel size local for internal dplb:
                 self.data_parallel_size_local = max(
@@ -2342,6 +2417,7 @@ def create_engine_config(
             enable_ep_weight_filter=self.enable_ep_weight_filter,
             all2all_backend=self.all2all_backend,
             enable_elastic_ep=self.enable_elastic_ep,
+            elastic_ep_max_dp_size=self.elastic_ep_max_dp_size,
             enable_dbo=self.enable_dbo,
             ubatch_size=self.ubatch_size,
             dbo_decode_token_threshold=self.dbo_decode_token_threshold,
@@ -2363,7 +2439,14 @@ def create_engine_config(
             dcp_comm_backend=self.dcp_comm_backend,
             dcp_q_replicate=self.dcp_q_replicate,
             dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,
-            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+            cp_kv_cache_interleave_size=(
+                self.cp_kv_cache_interleave_size
+                if self.cp_kv_cache_interleave_size is not None
+                else ParallelConfig.cp_kv_cache_interleave_size
+            ),
+            _allow_auto_resolve_cp_interleave_size=(
+                self.cp_kv_cache_interleave_size is None
+            ),
             _api_process_count=self._api_process_count,
             _api_process_rank=self._api_process_rank,
             assigned_physical_gpu_ids=self._resolve_device_ids(),
@@ -2402,6 +2485,7 @@ def create_engine_config(
             max_num_batched_tokens=self.max_num_batched_tokens,
             max_num_scheduled_tokens=self.max_num_scheduled_tokens,
             max_num_seqs=self.max_num_seqs,
+            max_num_active_seqs=self.max_num_active_seqs,
             max_num_queued_reqs=self.max_num_queued_reqs,
             max_num_queued_tokens=self.max_num_queued_tokens,
             max_model_len=model_config.max_model_len,
@@ -2412,6 +2496,9 @@ def create_engine_config(
             policy=self.scheduling_policy,
             scheduler_cls=self.scheduler_cls,
             long_prefill_token_threshold=self.long_prefill_token_threshold,
+            long_prefill_token_threshold_adaptive=(
+                self.long_prefill_token_threshold_adaptive
+            ),
             scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,
             watermark=self.watermark,
             prefill_schedule_interval=self.prefill_schedule_interval,
@@ -2532,6 +2619,8 @@ def create_engine_config(
             kernel_config.moe_backend = self.moe_backend
         if self.linear_backend != "auto":
             kernel_config.linear_backend = self.linear_backend
+        if self.sparse_indexer_topk_backend != "auto":
+            kernel_config.sparse_indexer_topk_backend = self.sparse_indexer_topk_backend
 
         # Transfer top-level ir_op_priority into KernelConfig.ir_op_priority
         for op_name, op_priority in asdict(self.ir_op_priority).items():
@@ -2622,6 +2711,7 @@ def create_engine_config(
             load_config=load_config,
             offload_config=offload_config,
             attention_config=attention_config,
+            aux_output_config=self.aux_output_config,
             engram_config=self.engram_config,
             mamba_config=mamba_config,
             kernel_config=kernel_config,
diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py
index 63f6970e0561..74849a85cef9 100644
--- a/vllm/engine/protocol.py
+++ b/vllm/engine/protocol.py
@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING, Any
 
 from vllm.config import ModelConfig, VllmConfig
+from vllm.config.kv_events import KVEventsConfig
 from vllm.distributed.weight_transfer.base import (
     WeightTransferInitRequest,
     WeightTransferUpdateRequest,
@@ -21,6 +22,7 @@
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.input_processor import InputProcessor
 from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult
+from vllm.v1.kv_hints import KvHintsEnvelope
 
 if TYPE_CHECKING:
     from vllm.v1.engine import PauseMode
@@ -39,7 +41,7 @@ class StreamingInput:
 
 
 class EngineClient(ABC):
-    """Protocol class for Clients to Engine"""
+    """Protocol class for Clients to Engine."""
 
     vllm_config: VllmConfig
     model_config: ModelConfig
@@ -62,6 +64,10 @@ def errored(self) -> bool: ...
     @abstractmethod
     def dead_error(self) -> BaseException: ...
 
+    def get_kv_event_sources(self) -> dict[int, KVEventsConfig]:
+        """KV-event publisher config of each engine, keyed by DP rank."""
+        return {}
+
     def check_admission(  # noqa: B027
         self, n: int = 1, request_id: str | None = None
     ) -> None:
@@ -77,6 +83,7 @@ def check_admission(  # noqa: B027
 
         Raises:
             GracefulHTTPError: If the request cannot be admitted.
+
         """
 
     @abstractmethod
@@ -96,6 +103,7 @@ def generate(
         priority: int = 0,
         data_parallel_rank: int | None = None,
         session_id: str | None = None,
+        kv_hints: KvHintsEnvelope | None = None,
         reasoning_ended: bool | None = None,
         reasoning_parser_kwargs: dict[str, Any] | None = None,
     ) -> AsyncGenerator[RequestOutput, None]:
@@ -124,6 +132,7 @@ async def abort(self, request_id: str | Iterable[str]) -> None:
         Args:
             request_id: The unique id of the request,
                         or an iterable of such ids.
+
         """
         ...
 
@@ -149,49 +158,54 @@ async def do_log_stats(self) -> None: ...
 
     @abstractmethod
     async def check_health(self) -> None:
-        """Raise if unhealthy"""
+        """Raise if unhealthy."""
         ...
 
     @abstractmethod
     async def start_profile(self) -> None:
-        """Start profiling the engine"""
+        """Start profiling the engine."""
         ...
 
     @abstractmethod
     async def stop_profile(self) -> None:
-        """Stop profiling the engine"""
+        """Stop profiling the engine."""
         ...
 
     @abstractmethod
     async def reset_mm_cache(self) -> None:
-        """Reset the multi-modal cache"""
+        """Reset the multi-modal cache."""
         ...
 
     @abstractmethod
     async def reset_encoder_cache(self) -> None:
-        """Reset the encoder cache"""
+        """Reset the encoder cache."""
         ...
 
     @abstractmethod
     async def reset_prefix_cache(
         self, reset_running_requests: bool = False, reset_connector: bool = False
     ) -> bool:
-        """Reset the prefix cache and optionally any configured connector cache"""
+        """Reset the prefix cache and optionally any configured connector cache."""
         ...
 
     @abstractmethod
     async def sleep(self, level: int = 1, mode: "PauseMode" = "abort") -> None:
-        """Sleep the engine"""
+        """Sleep the engine."""
+        ...
+
+    @abstractmethod
+    async def release_kv_cache_memory(self) -> None:
+        """Discard KV cache physical GPU memory. Requires a completed pause."""
         ...
 
     @abstractmethod
     async def wake_up(self, tags: list[str] | None = None) -> None:
-        """Wake up the engine"""
+        """Wake up the engine."""
         ...
 
     @abstractmethod
     async def is_sleeping(self) -> bool:
-        """Check whether the engine is sleeping"""
+        """Check whether the engine is sleeping."""
         ...
 
     @abstractmethod
@@ -219,6 +233,7 @@ async def pause_generation(
             wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead.
             clear_cache: DEPRECATED. Whether to clear KV and prefix caches
                 after draining.
+
         """
         ...
 
@@ -240,7 +255,7 @@ def shutdown(self, timeout: float | None = None) -> None:
     async def scale_elastic_ep(
         self, new_data_parallel_size: int, drain_timeout: int = 300
     ) -> None:
-        """Scale the engine"""
+        """Scale the engine."""
         raise NotImplementedError
 
     async def collective_rpc(
@@ -256,7 +271,7 @@ async def collective_rpc(
     async def handle_fault(
         self, fault_tolerance_request: FaultToleranceRequest
     ) -> FaultToleranceResult:
-        """send fault tolerance instruction to the engine"""
+        """Send fault tolerance instruction to the engine."""
         raise NotImplementedError
 
     async def get_status(self):
@@ -264,7 +279,7 @@ async def get_status(self):
         raise NotImplementedError
 
     async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
-        """Get supported tasks"""
+        """Get supported tasks."""
         raise NotImplementedError
 
     async def init_weight_transfer_engine(
diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py
index 1ef5dbe1f0c4..c879866dfb2c 100644
--- a/vllm/entrypoints/anthropic/protocol.py
+++ b/vllm/entrypoints/anthropic/protocol.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Pydantic models for Anthropic API protocol"""
+"""Pydantic models for Anthropic API protocol."""
 
 import time
 from typing import Annotated, Any, Literal
@@ -11,21 +11,21 @@
 
 
 class AnthropicError(BaseModel):
-    """Error structure for Anthropic API"""
+    """Error structure for Anthropic API."""
 
     type: str
     message: str
 
 
 class AnthropicErrorResponse(BaseModel):
-    """Error response structure for Anthropic API"""
+    """Error response structure for Anthropic API."""
 
     type: Literal["error"] = "error"
     error: AnthropicError
 
 
 class AnthropicUsage(BaseModel):
-    """Token usage information"""
+    """Token usage information."""
 
     input_tokens: int
     output_tokens: int
@@ -34,7 +34,7 @@ class AnthropicUsage(BaseModel):
 
 
 class AnthropicContentBlock(BaseModel):
-    """Content block in message"""
+    """Content block in message."""
 
     type: Literal[
         "text",
@@ -65,14 +65,14 @@ class AnthropicContentBlock(BaseModel):
 
 
 class AnthropicMessage(BaseModel):
-    """Message structure"""
+    """Message structure."""
 
     role: Literal["user", "assistant", "system"]
     content: str | list[AnthropicContentBlock]
 
 
 class AnthropicTool(BaseModel):
-    """Tool definition"""
+    """Tool definition."""
 
     name: str
     description: str | None = None
@@ -91,7 +91,7 @@ def validate_input_schema(cls, v):
 
 
 class AnthropicToolChoice(BaseModel):
-    """Tool Choice definition"""
+    """Tool Choice definition."""
 
     type: Literal["auto", "any", "tool", "none"]
     name: str | None = None
@@ -105,7 +105,7 @@ def validate_name_required_for_tool(self) -> "AnthropicToolChoice":
 
 
 class AnthropicJsonOutputFormat(BaseModel):
-    """JSON output format configuration"""
+    """JSON output format configuration."""
 
     json_schema: dict[str, Any] | None = Field(default=None, alias="schema")
     type: Literal["json_schema"] = "json_schema"
@@ -119,7 +119,7 @@ class AnthropicOutputConfig(BaseModel):
 
 
 class AnthropicMessagesRequest(BaseModel):
-    """Anthropic Messages API request"""
+    """Anthropic Messages API request."""
 
     model: str
     messages: list[AnthropicMessage]
@@ -192,7 +192,7 @@ def validate_max_tokens(cls, v):
 
 
 class AnthropicDelta(BaseModel):
-    """Delta for streaming responses"""
+    """Delta for streaming responses."""
 
     type: (
         Literal["text_delta", "input_json_delta", "thinking_delta", "signature_delta"]
@@ -211,7 +211,7 @@ class AnthropicDelta(BaseModel):
 
 
 class AnthropicStreamEvent(BaseModel):
-    """Streaming event"""
+    """Streaming event."""
 
     type: Literal[
         "message_start",
@@ -232,7 +232,7 @@ class AnthropicStreamEvent(BaseModel):
 
 
 class AnthropicMessagesResponse(BaseModel):
-    """Anthropic Messages API response"""
+    """Anthropic Messages API response."""
 
     id: str
     type: Literal["message"] = "message"
@@ -265,7 +265,7 @@ class AnthropicContextManagement(BaseModel):
 
 
 class AnthropicCountTokensRequest(BaseModel):
-    """Anthropic messages.count_tokens request"""
+    """Anthropic messages.count_tokens request."""
 
     model: str
     messages: list[AnthropicMessage]
@@ -291,7 +291,7 @@ def validate_model(cls, v):
 
 
 class AnthropicCountTokensResponse(BaseModel):
-    """Anthropic messages.count_tokens response"""
+    """Anthropic messages.count_tokens response."""
 
     input_tokens: int
     context_management: AnthropicContextManagement | None = None
diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py
index fe0d752d4ecc..6cc1271add1c 100644
--- a/vllm/entrypoints/anthropic/serving.py
+++ b/vllm/entrypoints/anthropic/serving.py
@@ -3,7 +3,7 @@
 # Adapted from
 # https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/chat_completion/serving.py
 
-"""Anthropic Messages API serving handler"""
+"""Anthropic Messages API serving handler."""
 
 import json
 import logging
@@ -96,7 +96,7 @@ def wrap_data_with_event(data: str, event: str):
 
 
 class AnthropicServingMessages(OpenAIServingChat):
-    """Handler for Anthropic Messages API requests"""
+    """Handler for Anthropic Messages API requests."""
 
     def __init__(
         self,
@@ -112,6 +112,7 @@ def __init__(
         reasoning_parser: str = "",
         enable_auto_tools: bool = False,
         tool_parser: str | None = None,
+        tool_strict_level: str = "auto",
         enable_prompt_tokens_details: bool = False,
         enable_force_include_usage: bool = False,
         default_chat_template_kwargs: dict[str, Any] | None = None,
@@ -127,6 +128,7 @@ def __init__(
             return_tokens_as_token_ids=return_tokens_as_token_ids,
             reasoning_parser=reasoning_parser,
             enable_auto_tools=enable_auto_tools,
+            tool_strict_level=tool_strict_level,
             tool_parser=tool_parser,
             enable_prompt_tokens_details=enable_prompt_tokens_details,
             enable_force_include_usage=enable_force_include_usage,
@@ -196,7 +198,7 @@ def to_chat_completion_request(
         *,
         merge_inline_system: bool = False,
     ) -> ChatCompletionRequest:
-        """Convert Anthropic message format to OpenAI format"""
+        """Convert Anthropic message format to OpenAI format."""
         openai_messages: list[dict[str, Any]] = []
 
         cls._convert_system_message(
@@ -224,7 +226,7 @@ def _convert_system_message(
         *,
         merge_inline_system: bool = False,
     ) -> None:
-        """Convert Anthropic system message to OpenAI format"""
+        """Convert Anthropic system message to OpenAI format."""
         system_parts: list[str] = []
 
         # Top-level system field
@@ -278,7 +280,7 @@ def _convert_messages(
         *,
         merge_inline_system: bool = False,
     ) -> None:
-        """Convert Anthropic messages to OpenAI format"""
+        """Convert Anthropic messages to OpenAI format."""
         for msg in messages:
             # Handle system messages in-place: extract text, strip billing
             # headers, and only emit if there is real content.  This avoids
@@ -310,7 +312,7 @@ def _convert_message_content(
         openai_msg: dict[str, Any],
         openai_messages: list[dict[str, Any]],
     ) -> None:
-        """Convert complex message content blocks"""
+        """Convert complex message content blocks."""
         content_parts: list[dict[str, Any]] = []
         tool_calls: list[dict[str, Any]] = []
         reasoning_parts: list[str] = []
@@ -349,7 +351,7 @@ def _convert_block(
         reasoning_parts: list[str],
         openai_messages: list[dict[str, Any]],
     ) -> None:
-        """Convert individual content block"""
+        """Convert individual content block."""
         if block.type == "text" and block.text:
             content_parts.append({"type": "text", "text": block.text})
         elif block.type == "image" and block.source:
@@ -374,7 +376,7 @@ def _convert_block(
 
     @classmethod
     def _convert_tool_use_block(cls, block, tool_calls: list[dict[str, Any]]) -> None:
-        """Convert tool_use block to OpenAI function call format"""
+        """Convert tool_use block to OpenAI function call format."""
         tool_call = {
             "id": block.id or f"call_{int(time.time())}",
             "type": "function",
@@ -393,7 +395,7 @@ def _convert_tool_result_block(
         openai_messages: list[dict[str, Any]],
         content_parts: list[dict[str, Any]],
     ) -> None:
-        """Convert tool_result block to OpenAI format"""
+        """Convert tool_result block to OpenAI format."""
         if role == "user":
             cls._convert_user_tool_result(block, openai_messages)
         else:
@@ -406,7 +408,7 @@ def _convert_tool_result_block(
     def _convert_user_tool_result(
         cls, block, openai_messages: list[dict[str, Any]]
     ) -> None:
-        """Convert user tool_result with text and image support"""
+        """Convert user tool_result with text and image support."""
         tool_text = ""
         tool_image_urls: list[str] = []
         tool_reference: list[dict[str, Any]] = []
@@ -468,7 +470,7 @@ def _build_base_request(
         anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
         openai_messages: list[dict[str, Any]],
     ) -> ChatCompletionRequest:
-        """Build base ChatCompletionRequest"""
+        """Build base ChatCompletionRequest."""
         if isinstance(anthropic_request, AnthropicCountTokensRequest):
             return ChatCompletionRequest(
                 model=anthropic_request.model,
@@ -498,7 +500,7 @@ def _handle_output_config(
         req: ChatCompletionRequest,
         anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
     ) -> None:
-        """Handle output configuration such as output format and effort"""
+        """Handle output configuration such as output format and effort."""
         if isinstance(anthropic_request, AnthropicCountTokensRequest):
             return
         output_config: AnthropicOutputConfig | None = anthropic_request.output_config
@@ -519,7 +521,7 @@ def _handle_streaming_options(
         req: ChatCompletionRequest,
         anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
     ) -> None:
-        """Handle streaming configuration"""
+        """Handle streaming configuration."""
         if isinstance(anthropic_request, AnthropicCountTokensRequest):
             return
         if anthropic_request.stream:
@@ -534,7 +536,7 @@ def _convert_tool_choice(
         anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
         req: ChatCompletionRequest,
     ) -> None:
-        """Convert Anthropic tool_choice to OpenAI format"""
+        """Convert Anthropic tool_choice to OpenAI format."""
         if anthropic_request.tool_choice is None:
             req.tool_choice = None
             return
@@ -563,7 +565,7 @@ def _convert_tools(
         anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
         req: ChatCompletionRequest,
     ) -> None:
-        """Convert Anthropic tools to OpenAI format"""
+        """Convert Anthropic tools to OpenAI format."""
         if anthropic_request.tools is None:
             return
 
@@ -593,8 +595,7 @@ async def create_messages(
         request: AnthropicMessagesRequest,
         raw_request: Request | None = None,
     ) -> AsyncGenerator[str, None] | AnthropicMessagesResponse | ErrorResponse:
-        """
-        Messages API similar to Anthropic's API.
+        """Messages API similar to Anthropic's API.
 
         See https://docs.anthropic.com/en/api/messages
         for the API specification. This API mimics the Anthropic messages API.
@@ -842,7 +843,9 @@ def stop_and_flush() -> list[str]:
                                 stop_delta = AnthropicDelta(
                                     stop_reason=self.stop_reason_map.get(
                                         finish_reason or "stop"
-                                    )
+                                    ),
+                                    # Set explicitly so exclude_unset=True keeps it.
+                                    stop_sequence=None,
                                 )
                             chunk = AnthropicStreamEvent(
                                 type="message_delta",
diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py
index 327f9c4b24a4..e3cc744ba70d 100644
--- a/vllm/entrypoints/chat_utils.py
+++ b/vllm/entrypoints/chat_utils.py
@@ -223,9 +223,7 @@ class ChatCompletionContentPartVideoParam(TypedDict, total=False):
 
 
 class PILImage(BaseModel):
-    """
-    A PIL.Image.Image object.
-    """
+    """A PIL.Image.Image object."""
 
     image_pil: Image.Image
     model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -238,6 +236,7 @@ class CustomChatCompletionContentPILImageParam(TypedDict, total=False):
     {
         "image_pil": ImageAsset('cherry_blossom').pil_image
     }
+
     """
 
     image_pil: PILImage | None
@@ -256,6 +255,7 @@ class CustomChatCompletionContentSimpleImageParam(TypedDict, total=False):
     {
         "image_url": "https://example.com/image.jpg"
     }
+
     """
 
     image_url: str | None
@@ -273,6 +273,7 @@ class CustomChatCompletionContentSimpleAudioParam(TypedDict, total=False):
     {
         "audio_url": "https://example.com/audio.mp3"
     }
+
     """
 
     audio_url: str | None
@@ -285,6 +286,7 @@ class CustomChatCompletionContentSimpleVideoParam(TypedDict, total=False):
     {
         "video_url": "https://example.com/video.mp4"
     }
+
     """
 
     video_url: str | None
@@ -304,6 +306,7 @@ class CustomThinkCompletionContentParam(TypedDict, total=False):
         "closed": True,
         "type": "thinking"
     }
+
     """
 
     thinking: Required[str]
@@ -324,6 +327,7 @@ class CustomChatCompletionContentToolReferenceParam(TypedDict, total=False):
         "name": "get_weather",
         "type": "tool_reference"
     }
+
     """
 
     name: str
@@ -586,8 +590,7 @@ def _get_embeds_data(
 
 
 class BaseMultiModalItemTracker(ABC, Generic[_T]):
-    """
-    Tracks multi-modal items in a given request and ensures that the number
+    """Tracks multi-modal items in a given request and ensures that the number
     of multi-modal items in a given request does not exceed the configured
     maximum per prompt.
     """
@@ -651,8 +654,7 @@ def video_processor_name(self) -> str | None:
         return get_video_processor_cls_name(self.model_config)
 
     def add(self, modality: ModalityStr, item: _T) -> str | None:
-        """
-        Add a multi-modal item to the current prompt and returns the
+        """Add a multi-modal item to the current prompt and returns the
         placeholder string to use, if any.
 
         An optional uuid can be added which serves as a unique identifier of the
@@ -663,6 +665,7 @@ def add(self, modality: ModalityStr, item: _T) -> str | None:
             pre-computed embeddings that do not go through any HF processor, encoder,
             or model-specific placeholder logic. The corresponding placeholder string is
             managed by the parser via `_add_placeholder`, so we return None here.
+
         """
         add_info = self._validate_add(modality)
         if add_info is None:
@@ -791,13 +794,13 @@ def _resolve_items(
     mm_processor: BaseMultiModalProcessor | None,
     modality_order: dict[str, list[str]],
 ) -> tuple[MultiModalDataDict, MultiModalUUIDDict]:
-    """
-    Materialize the tracker's per-modality items into `mm_data` / `mm_uuids`.
+    """Materialize the tracker's per-modality items into `mm_data` / `mm_uuids`.
 
     Note:
         `mm_processor` is `None` for text-only models (no registered HF
         processor) whose only modality is `prompt_embeds`. Every other
         modality requires a processor, enforced by the guard below.
+
     """
     if "image" in items_by_modality and "image_embeds" in items_by_modality:
         raise VLLMValidationError(
@@ -1568,7 +1571,6 @@ def _get_full_multimodal_text_prompt(
     multimodal_content_part_separator: str = "\n",
 ) -> str:
     """Combine multimodal prompts for a multimodal language model."""
-
     # flatten storage to make it looks like
     # {
     #   "<|image|>": 2,
@@ -1687,8 +1689,7 @@ def _collect_extra_fields(part: dict[str, Any]) -> dict[str, Any]:
 def _parse_chat_message_content_mm_part(
     part: ChatCompletionContentPartParam,
 ) -> tuple[str, _ContentPart]:
-    """
-    Parses a given multi-modal content part based on its type.
+    """Parses a given multi-modal content part based on its type.
 
     Args:
         part: A dict containing the content part, with a potential 'type' field.
@@ -1700,6 +1701,7 @@ def _parse_chat_message_content_mm_part(
 
     Raises:
         ValueError: If the 'type' field is missing and no direct URL is found.
+
     """
     assert isinstance(
         part, dict
@@ -1773,7 +1775,7 @@ def _parse_chat_message_content_mm_part(
                 # with url as a dict of {"url": url}
                 audio_url = audio_url.get("url", None)
             return "audio_url", audio_url
-        if part.get("input_audio") is not None:
+        if "input_audio" in part:
             input_audio_params = _InputAudioParser(part).get("input_audio", None)
             return "input_audio", input_audio_params
         if "video_url" in part:
diff --git a/vllm/entrypoints/cli/benchmark/base.py b/vllm/entrypoints/cli/benchmark/base.py
index 60189ab900e8..717081e04cd1 100644
--- a/vllm/entrypoints/cli/benchmark/base.py
+++ b/vllm/entrypoints/cli/benchmark/base.py
@@ -22,5 +22,6 @@ def cmd(args: argparse.Namespace) -> None:
 
         Args:
             args: The arguments to the command.
+
         """
         raise NotImplementedError
diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py
index d173bc392f1d..ab9afcb008d9 100644
--- a/vllm/entrypoints/cli/main.py
+++ b/vllm/entrypoints/cli/main.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""The CLI entrypoints of vLLM
+"""The CLI entrypoints of vLLM.
 
 Note that all future modules must be lazily loaded within main
 to avoid certain eager import breakage."""
@@ -46,6 +46,7 @@ def main():
     import vllm.entrypoints.cli.openai
     import vllm.entrypoints.cli.run_batch
     import vllm.entrypoints.cli.serve
+    import vllm.entrypoints.cli.snapshot
     from vllm.entrypoints.serve.utils.api_utils import (
         VLLM_SUBCMD_PARSER_EPILOG,
         cli_env_setup,
@@ -59,9 +60,11 @@ def main():
         vllm.entrypoints.cli.benchmark.main,
         vllm.entrypoints.cli.collect_env,
         vllm.entrypoints.cli.run_batch,
+        vllm.entrypoints.cli.snapshot,
     ]
 
-    cli_env_setup()
+    if sys.argv[1:2] != ["snapshot"]:
+        cli_env_setup()
 
     vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench()
 
@@ -94,7 +97,12 @@ def main():
     subparsers = parser.add_subparsers(required=False, dest="subparser")
     cmds = {}
     for cmd_module in CMD_MODULES:
-        new_cmds = cmd_module.cmd_init()
+        if cmd_module is vllm.entrypoints.cli.snapshot:
+            new_cmds = cmd_module.cmd_init(
+                create_requested=sys.argv[1:3] == ["snapshot", "create"]
+            )
+        else:
+            new_cmds = cmd_module.cmd_init()
         for cmd in new_cmds:
             cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
             cmds[cmd.name] = cmd
diff --git a/vllm/entrypoints/cli/openai.py b/vllm/entrypoints/cli/openai.py
index fa01ff534b4c..b4e286838524 100644
--- a/vllm/entrypoints/cli/openai.py
+++ b/vllm/entrypoints/cli/openai.py
@@ -46,21 +46,35 @@ def _interactive_cli(args: argparse.Namespace) -> tuple[str, OpenAI]:
 
 
 def _print_chat_stream(stream, stats: bool = False) -> str:
-    output = ""
-    start = time.perf_counter()
+    output: str = ""
+    in_reasoning: bool = False
+    start: float = time.perf_counter()
     ttft: float | None = None
-    completion_tokens = 0
+    completion_tokens: int = 0
     for chunk in stream:
         if chunk.usage is not None:
             completion_tokens = chunk.usage.completion_tokens
         if not chunk.choices:
             continue
         delta = chunk.choices[0].delta
+        reasoning = getattr(delta, "reasoning", None) or getattr(
+            delta, "reasoning_content", None
+        )
+        if ttft is None and (reasoning or delta.content):
+            ttft = time.perf_counter() - start
+        if reasoning:
+            if not in_reasoning:
+                print("<think>", flush=True)
+                in_reasoning = True
+            print(reasoning, end="", flush=True)
         if delta.content:
-            if ttft is None:
-                ttft = time.perf_counter() - start
+            if in_reasoning:
+                print("\n</think>", flush=True)
+                in_reasoning = False
             output += delta.content
             print(delta.content, end="", flush=True)
+    if in_reasoning:
+        print("\n</think>", end="", flush=True)
     print()
     if stats:
         _print_metrics(start, ttft, completion_tokens)
diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py
index a284b4bf39f5..6a8492f4c4b7 100644
--- a/vllm/entrypoints/cli/serve.py
+++ b/vllm/entrypoints/cli/serve.py
@@ -18,6 +18,7 @@
 from vllm.entrypoints.launchers.dp_supervisor import run_dp_supervisor
 from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger
+from vllm.reasoning import ReasoningParserManager
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.network_utils import get_tcp_uri
@@ -54,7 +55,7 @@ def cmd(args: argparse.Namespace) -> None:
             args.model = args.model_tag
 
         if getattr(args, "grpc", False):
-            from vllm.entrypoints.grpc_server import serve_grpc
+            from vllm.entrypoints.launchers.grpc_server import serve_grpc
 
             uvloop.run(serve_grpc(args))
             return
@@ -179,6 +180,9 @@ def run_headless(args: argparse.Namespace):
     if args.api_server_count > 1:
         raise ValueError("api_server_count can't be set in headless mode")
 
+    if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3:
+        ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin)
+
     # Create the EngineConfig.
     engine_args = vllm.AsyncEngineArgs.from_cli_args(args)
     usage_context = UsageContext.OPENAI_API_SERVER
diff --git a/vllm/entrypoints/cli/snapshot.py b/vllm/entrypoints/cli/snapshot.py
new file mode 100644
index 000000000000..f5201675d40c
--- /dev/null
+++ b/vllm/entrypoints/cli/snapshot.py
@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import json
+import os
+import platform
+import typing
+from pathlib import Path
+
+from vllm.entrypoints.cli.types import CLISubcommand
+
+if typing.TYPE_CHECKING:
+    from vllm.utils.argparse_utils import FlexibleArgumentParser
+else:
+    FlexibleArgumentParser = argparse.ArgumentParser
+
+
+def validate_create_args(args: argparse.Namespace) -> None:
+    model = str(getattr(args, "model_tag", None) or args.model)
+    revision = str(getattr(args, "revision", None) or "")
+    if (
+        Path(model).exists()
+        or len(revision) != 40
+        or any(character not in "0123456789abcdefABCDEF" for character in revision)
+    ):
+        raise ValueError(
+            "snapshot create requires an immutable remote model and "
+            "40-character --revision"
+        )
+    parallel_sizes = (
+        ("tensor parallel", args.tensor_parallel_size),
+        ("pipeline parallel", args.pipeline_parallel_size),
+        ("data parallel", args.data_parallel_size),
+        ("prefill context parallel", args.prefill_context_parallel_size),
+    )
+    for name, size in parallel_sizes:
+        if size != 1:
+            raise ValueError(f"snapshot create currently requires {name} size 1")
+    unsupported_frontend = (
+        args.api_server_count not in (None, 1)
+        or args.grpc
+        or args.headless
+        or args.api_key
+        or os.environ.get("VLLM_API_KEY")
+        or args.uds
+        or args.middleware
+    )
+    if unsupported_frontend or any(
+        (
+            args.ssl_keyfile,
+            args.ssl_certfile,
+            args.ssl_ca_certs,
+            args.enable_ssl_refresh,
+            args.ssl_ciphers,
+        )
+    ):
+        raise ValueError(
+            "snapshot create requires one unauthenticated plaintext HTTP server"
+        )
+    if args.logprobs_mode in {"raw_logits", "processed_logits"}:
+        raise ValueError("snapshot canary requires a log-probability mode")
+    if any(
+        getattr(args, name, None) is not None
+        for name in ("speculative_config", "spec_method", "spec_model", "spec_tokens")
+    ):
+        raise ValueError("snapshot create does not support speculative decoding")
+    if platform.system() != "Linux" or platform.machine() != "x86_64":
+        raise ValueError("snapshot create requires Linux x86_64")
+
+
+def run_create(args: argparse.Namespace) -> None:
+    validate_create_args(args)
+    from vllm.snapshot.controller import create_snapshot
+
+    create_snapshot(args)
+
+
+def run_inspect(args: argparse.Namespace) -> None:
+    from vllm.snapshot.manifest import inspect_snapshot
+
+    print(json.dumps(inspect_snapshot(Path(args.snapshot_dir)), indent=2))
+
+
+def run_restore(args: argparse.Namespace) -> None:
+    from vllm.snapshot.controller import restore_snapshot
+
+    restore_snapshot(args)
+
+
+class SnapshotSubcommand(CLISubcommand):
+    """The `snapshot` subcommand for the vLLM CLI."""
+
+    name = "snapshot"
+
+    def __init__(self, *, create_requested: bool = False) -> None:
+        self.create_requested = create_requested
+
+    @staticmethod
+    def cmd(args: argparse.Namespace) -> None:
+        args.snapshot_dispatch(args)
+
+    def subparser_init(
+        self, subparsers: argparse._SubParsersAction
+    ) -> FlexibleArgumentParser:
+        parser = subparsers.add_parser(
+            self.name,
+            help="Create, inspect, or restore an initialized vLLM snapshot.",
+            usage="vllm snapshot <create|inspect|restore> [options]",
+        )
+        actions = parser.add_subparsers(required=True, dest="snapshot_action")
+
+        create_parser = actions.add_parser(
+            "create", help="Create a snapshot from an initialized TP1 engine."
+        )
+        if self.create_requested:
+            from vllm.entrypoints.launchers.cli_args import make_arg_parser
+
+            create_parser = make_arg_parser(create_parser)
+        else:
+            create_parser.add_argument("model_tag", nargs="?")
+        create_parser.add_argument("--snapshot-dir", required=True)
+        create_parser.set_defaults(snapshot_dispatch=run_create)
+
+        inspect_parser = actions.add_parser(
+            "inspect", help="Inspect a snapshot without restoring it."
+        )
+        inspect_parser.add_argument("snapshot_dir")
+        inspect_parser.set_defaults(snapshot_dispatch=run_inspect)
+
+        restore_parser = actions.add_parser(
+            "restore", help="Restore a same-host TP1 snapshot."
+        )
+        restore_parser.add_argument("snapshot_dir")
+        restore_parser.add_argument("--host", default=None)
+        restore_parser.add_argument("--port", type=int, default=8000)
+        restore_parser.set_defaults(snapshot_dispatch=run_restore)
+
+        return parser
+
+
+def cmd_init(*, create_requested: bool = False) -> list[CLISubcommand]:
+    return [SnapshotSubcommand(create_requested=create_requested)]
diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py
index 445964e6aa47..eaac36f1336f 100644
--- a/vllm/entrypoints/cohere/api_router.py
+++ b/vllm/entrypoints/cohere/api_router.py
@@ -17,9 +17,10 @@
 
 Even when the SDK is installed, :func:`attach_router` also requires
 ``VLLM_ENABLE_COHERE_API=1`` in the environment before it will expose
-the chat route. The render route additionally requires
-``VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1``. These gates keep unrelated
-deployments from accidentally exposing the APIs.
+the chat route. The render route additionally requires scale-out
+endpoints to be enabled, i.e. ``--enable-scale-out`` or
+``--tokens-only``. These gates keep unrelated deployments from
+accidentally exposing the APIs.
 
 Note: the handlers must live at module scope (not inside
 ``attach_router``) so that FastAPI's ``typing.get_type_hints`` resolves
@@ -306,6 +307,7 @@ def attach_router(app: FastAPI) -> None:
         )
         return
     app.include_router(router)
-    if envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS:
+    args = getattr(app.state, "args", None)
+    if getattr(args, "enable_scale_out", False) or getattr(args, "tokens_only", False):
         app.include_router(render_router)
     app.add_middleware(CohereErrorEnvelopeMiddleware)
diff --git a/vllm/entrypoints/cohere/cohere_chat_message.py b/vllm/entrypoints/cohere/cohere_chat_message.py
index 042856dea68c..85925d0f1878 100644
--- a/vllm/entrypoints/cohere/cohere_chat_message.py
+++ b/vllm/entrypoints/cohere/cohere_chat_message.py
@@ -7,8 +7,8 @@
 It provides two things:
 
 * :class:`Citation` / :class:`CitationSource`: the vLLM-internal citation
-  representation produced by the Cohere reasoning parser
-  (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) and consumed by
+  representation produced by the Cohere parser
+  (:mod:`vllm.parser.cohere_command`) and consumed by
   :mod:`vllm.entrypoints.cohere.serving`. This is *not* the on-wire Cohere
   SDK shape -- that conversion happens in the serving layer.
 * :class:`CohereChatMessage` / :class:`CohereDeltaMessage`: :class:`ChatMessage`
@@ -42,9 +42,9 @@ class CitationSource(OpenAIBaseModel):
     document/tool-output identifier; ``document`` and ``tool_output`` carry
     the original payload that produced the citation.
 
-    Sources are fully resolved by the reasoning parser at emit time
+    Sources are fully resolved by the parser at emit time
     (see :func:`_melody_sources_to_vllm` in
-    :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using the
+    :mod:`vllm.parser.cohere_command`) using the
     ``POSITION_TO_SOURCE_KEY`` map forwarded through
     ``chat_template_kwargs`` by
     :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._apply_cohere_template_kwargs`.
@@ -121,8 +121,7 @@ def _serialize(self, handler):
 class CohereDeltaMessage(DeltaMessage):
     """:class:`DeltaMessage` extension carrying grounding citations for streaming.
 
-    Emitted by
-    :class:`vllm.reasoning.cohere_command_reasoning_parser.BaseCohereCommandReasoningParser`
+    Emitted by :class:`vllm.parser.cohere_command.CohereCommandParser`
     on delta events whose payload includes citations. Non-Cohere parsers
     return plain :class:`DeltaMessage`, so their streamed shape is
     unchanged.
diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py
index 8ec4701907ed..63d00f0c576a 100644
--- a/vllm/entrypoints/cohere/serving.py
+++ b/vllm/entrypoints/cohere/serving.py
@@ -249,6 +249,7 @@ def __init__(
         reasoning_parser: str = "",
         enable_auto_tools: bool = False,
         tool_parser: str | None = None,
+        tool_strict_level: str = "auto",
         enable_prompt_tokens_details: bool = False,
         enable_force_include_usage: bool = False,
         default_chat_template_kwargs: dict[str, Any] | None = None,
@@ -265,6 +266,7 @@ def __init__(
             return_tokens_as_token_ids=return_tokens_as_token_ids,
             reasoning_parser=reasoning_parser,
             enable_auto_tools=enable_auto_tools,
+            tool_strict_level=tool_strict_level,
             tool_parser=tool_parser,
             enable_prompt_tokens_details=enable_prompt_tokens_details,
             enable_force_include_usage=enable_force_include_usage,
@@ -714,8 +716,8 @@ def _apply_cohere_template_kwargs(
         # emits, so citation deltas already carry SDK-shape sources by
         # the time they reach the streaming loop. See
         # ``_build_position_to_source`` for the numbering rule and
-        # ``_melody_sources_to_vllm`` in the reasoning parser for the
-        # consumer.
+        # ``_melody_sources_to_vllm`` in ``vllm.parser.cohere_command`` for
+        # the consumer.
         position_to_source = cls._build_position_to_source(request)
         if position_to_source:
             kwargs.setdefault(POSITION_TO_SOURCE_KEY, position_to_source)
@@ -1109,12 +1111,11 @@ def _finalize_response_message(
         *,
         parser: Parser | None,
     ) -> ChatMessage:
-        """Copy grounding citations off the reasoning parser onto the message.
+        """Copy grounding citations off the unified parser onto the message.
 
-        The Cohere reasoning parser
-        (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) caches the
-        citations produced by its most recent unary ``extract_reasoning``
-        call on ``last_unary_citations``. We surface them here on
+        :class:`vllm.parser.cohere_command.CohereCommandParser` caches the
+        citations produced by its most recent unary ``parse`` call on
+        ``last_unary_citations``. We surface them here on
         :class:`CohereChatMessage` so downstream response conversion can
         pick them up without the base :class:`OpenAIServingChat` having to
         know about citations.
@@ -1128,11 +1129,7 @@ def _finalize_response_message(
         # ``model_dump`` via the extras bucket, so the wire is correct
         # either way.
         message = cast(CohereChatMessage, message)
-        citations = getattr(
-            getattr(parser, "reasoning_parser", None),
-            "last_unary_citations",
-            None,
-        )
+        citations = getattr(parser, "last_unary_citations", None)
         if citations:
             message.citations = citations
         return message
@@ -1143,10 +1140,10 @@ def _extract_citations_if_any(self, msg: Any) -> list[Citation] | None:
         :class:`CohereChatMessage` carries a
         ``citations: list[vllm...Citation] | None`` field populated by
         :meth:`_finalize_response_message` (which reads it off the
-        reasoning parser's ``last_unary_citations`` cache). Sources are
+        unified parser's ``last_unary_citations`` cache). Sources are
         already fully resolved by the parser (see
         :func:`_melody_sources_to_vllm` in
-        :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using
+        :mod:`vllm.parser.cohere_command`) using
         the position map forwarded via ``chat_template_kwargs``. This
         method's remaining job is:
 
diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py
index 12c70e37d5c0..189c4e6779b0 100644
--- a/vllm/entrypoints/generate/api_router.py
+++ b/vllm/entrypoints/generate/api_router.py
@@ -112,10 +112,12 @@ async def init_generate_state(
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_auto_tools=args.enable_auto_tool_choice,
             tool_parser=args.tool_call_parser,
+            tool_strict_level=args.tool_strict_level,
             tool_server=state.tool_server,
             reasoning_parser=args.structured_outputs_config.reasoning_parser,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_force_include_usage=args.enable_force_include_usage,
+            enable_per_request_metrics=args.enable_per_request_metrics,
             enable_log_outputs=args.enable_log_outputs,
             default_chat_template_kwargs=default_chat_template_kwargs,
         )
@@ -136,6 +138,7 @@ async def init_generate_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
+        tool_strict_level=args.tool_strict_level,
         reasoning_parser=args.structured_outputs_config.reasoning_parser,
         enable_prompt_tokens_details=args.enable_prompt_tokens_details,
         enable_force_include_usage=args.enable_force_include_usage,
@@ -177,6 +180,7 @@ async def init_generate_state(
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_auto_tools=args.enable_auto_tool_choice,
             tool_parser=args.tool_call_parser,
+            tool_strict_level=args.tool_strict_level,
             reasoning_parser=args.structured_outputs_config.reasoning_parser,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_force_include_usage=args.enable_force_include_usage,
@@ -197,6 +201,7 @@ async def init_generate_state(
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_auto_tools=args.enable_auto_tool_choice,
             tool_parser=args.tool_call_parser,
+            tool_strict_level=args.tool_strict_level,
             reasoning_parser=args.structured_outputs_config.reasoning_parser,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_force_include_usage=args.enable_force_include_usage,
diff --git a/vllm/entrypoints/generate/base/protocol.py b/vllm/entrypoints/generate/base/protocol.py
index 76dc5e47a11c..91d2e30a47a2 100644
--- a/vllm/entrypoints/generate/base/protocol.py
+++ b/vllm/entrypoints/generate/base/protocol.py
@@ -6,10 +6,8 @@
 import json
 from typing import Annotated, Any, Literal, TypeAlias
 
-import regex as re
 from pydantic import (
     BaseModel,
-    ConfigDict,
     Field,
     model_serializer,
 )
@@ -21,7 +19,6 @@
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
 from vllm.sampling_params import StructuredOutputsParams
-from vllm.utils.import_utils import resolve_obj_by_qualname
 
 logger = init_logger(__name__)
 
@@ -261,53 +258,6 @@ def _serialize(self, handler):
         return data
 
 
-# extra="forbid" is a workaround to have kwargs as a field,
-# see https://github.com/pydantic/pydantic/issues/3125
-class LogitsProcessorConstructor(BaseModel):
-    qualname: str
-    args: list[Any] | None = None
-    kwargs: dict[str, Any] | None = None
-
-    model_config = ConfigDict(extra="forbid")
-
-
-LogitsProcessors = list[str | LogitsProcessorConstructor]
-
-
-def get_logits_processors(
-    processors: LogitsProcessors | None, pattern: str | None
-) -> list[Any] | None:
-    if processors and pattern:
-        logits_processors = []
-        for processor in processors:
-            qualname = processor if isinstance(processor, str) else processor.qualname
-            if not re.match(pattern, qualname):
-                raise ValueError(
-                    f"Logits processor '{qualname}' is not allowed by this "
-                    "server. See --logits-processor-pattern engine argument "
-                    "for more information."
-                )
-            try:
-                logits_processor = resolve_obj_by_qualname(qualname)
-            except Exception as e:
-                raise ValueError(
-                    f"Logits processor '{qualname}' could not be resolved: {e}"
-                ) from e
-            if isinstance(processor, LogitsProcessorConstructor):
-                logits_processor = logits_processor(
-                    *processor.args or [], **processor.kwargs or {}
-                )
-            logits_processors.append(logits_processor)
-        return logits_processors
-    elif processors:
-        raise ValueError(
-            "The `logits_processors` argument is not supported by this "
-            "server. See --logits-processor-pattern engine argument "
-            "for more information."
-        )
-    return None
-
-
 class FunctionCall(OpenAIBaseModel):
     # Internal field to preserve native tool call ID from tool parser.
     # Excluded from serialization to maintain OpenAI API compatibility
diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py
index dc41d3f067b2..5142346df661 100644
--- a/vllm/entrypoints/generate/base/serving.py
+++ b/vllm/entrypoints/generate/base/serving.py
@@ -179,6 +179,7 @@ def _preflight(self, n: int = 1) -> None:
 
         Args:
             n: Number of sequences the request will occupy.
+
         """
         if self.engine_client.errored:
             raise self.engine_client.dead_error
@@ -236,7 +237,7 @@ async def _get_trace_headers(
 
     @staticmethod
     def _get_data_parallel_rank(raw_request: Request | None) -> int | None:
-        """Pulls the data parallel rank from a header, if provided"""
+        """Pulls the data parallel rank from a header, if provided."""
         if raw_request is None:
             return None
 
diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py
index f412f5aebed8..b7095131af31 100644
--- a/vllm/entrypoints/generate/beam_search/offline.py
+++ b/vllm/entrypoints/generate/beam_search/offline.py
@@ -53,7 +53,7 @@ def _bitmask_to_token_ids(bitmask_row: torch.Tensor, vocab_size: int) -> list[in
 
 
 class BeamSearchOfflineMixin(OfflineInferenceMixin):
-    """Offline inference for beam search"""
+    """Offline inference for beam search."""
 
     def beam_search(
         self,
@@ -63,8 +63,7 @@ def beam_search(
         use_tqdm: bool = False,
         concurrency_limit: int | None = None,
     ) -> list[BeamSearchOutput]:
-        """
-        Generate sequences using beam search.
+        """Generate sequences using beam search.
 
         Args:
             prompts: A list of prompts. Each prompt can be a string or a list
@@ -74,6 +73,7 @@ def beam_search(
             use_tqdm: Whether to use tqdm to display the progress bar.
             concurrency_limit: The maximum number of concurrent requests.
                 If None, the number of concurrent requests is unlimited.
+
         """
         # TODO: how does beam search work together with length penalty,
         # frequency, penalty, and stopping criteria, etc.?
@@ -186,7 +186,9 @@ def beam_search(
             best_beams = sorted_completed[:beam_width]
 
             for beam in best_beams:
-                beam.text = tokenizer.decode(beam.tokens)
+                beam.text = tokenizer.decode(
+                    beam.tokens, skip_special_tokens=params.skip_special_tokens
+                )
 
             outputs.append(BeamSearchOutput(sequences=best_beams))
 
@@ -291,6 +293,18 @@ def _beam_search_step(
                     allowed_sets[i] = set(entry[1])
 
         for (start, end), instance in zip(instance_start_and_end, instances_batch):
+            instance_output = output[start:end]
+            if any(
+                result is not None and result.outputs[0].finish_reason == "abort"
+                for result in instance_output
+            ):
+                for beam, result in zip(all_beams[start:end], instance_output):
+                    if result is not None:
+                        beam.finish_reason = "abort"
+                        instance.completed.append(beam)
+                instance.beams = []
+                continue
+
             instance_new_beams = []
             for i in range(start, end):
                 current_beam = all_beams[i]
@@ -300,8 +314,6 @@ def _beam_search_step(
                     continue
 
                 if result.outputs[0].logprobs is not None:
-                    # if logprobs is None, the sequence completed
-                    # due to max-model-len or abortion.
                     logprobs = result.outputs[0].logprobs[0]
                     allowed = allowed_sets[i]
                     for token_id, logprob_obj in logprobs.items():
diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py
index 02ff61d62936..ac7f0f884f5b 100644
--- a/vllm/entrypoints/generate/beam_search/online.py
+++ b/vllm/entrypoints/generate/beam_search/online.py
@@ -20,7 +20,7 @@
 
 
 class BeamSearchOnlineMixin(ABC):
-    """online serving for beam search"""
+    """online serving for beam search."""
 
     renderer: BaseRenderer
     engine_client: EngineClient
@@ -99,11 +99,7 @@ async def beam_search(
 
             output = [x[0] for x in await asyncio.gather(*tasks)]
 
-            candidates = []
-            # Iterate through all beam inference results
-            for i, result in enumerate(output):
-                current_beam = all_beams[i]
-
+            for result in output:
                 # check for error finish reason and abort beam search
                 if result.outputs[0].finish_reason == "error":
                     # yield error output and terminate beam search
@@ -126,6 +122,15 @@ async def beam_search(
                     )
                     return
 
+            if any(result.outputs[0].finish_reason == "abort" for result in output):
+                for beam in all_beams:
+                    beam.finish_reason = "abort"
+                break
+
+            candidates = []
+            for i, result in enumerate(output):
+                current_beam = all_beams[i]
+
                 if result.outputs[0].logprobs is not None:
                     logprobs = result.outputs[0].logprobs[0]
                     for token_id, logprob_obj in logprobs.items():
@@ -197,7 +202,9 @@ async def beam_search(
                 tokens = beam.tokens[tokenized_length:-1]
             else:
                 tokens = beam.tokens[tokenized_length:]
-            beam.text = tokenizer.decode(tokens)
+            beam.text = tokenizer.decode(
+                tokens, skip_special_tokens=params.skip_special_tokens
+            )
 
         yield RequestOutput(
             request_id=request_id,
diff --git a/vllm/entrypoints/generate/beam_search/utils.py b/vllm/entrypoints/generate/beam_search/utils.py
index 47f69cbf7abe..61f97963b6f3 100644
--- a/vllm/entrypoints/generate/beam_search/utils.py
+++ b/vllm/entrypoints/generate/beam_search/utils.py
@@ -150,7 +150,8 @@ def get_beam_search_score(
     if tokens[-1] == eos_token_id:
         seq_len -= 1
 
-    return cumulative_logprob / (seq_len**length_penalty)
+    # An aborted beam may contain only an EOS prompt token.
+    return cumulative_logprob / (max(seq_len, 1) ** length_penalty)
 
 
 def create_sort_beams_key_function(eos_token_id: int, length_penalty: float):
diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py
index 753054d1b8a0..cb0c4995bf75 100644
--- a/vllm/entrypoints/generate/generative_scoring/serving.py
+++ b/vllm/entrypoints/generate/generative_scoring/serving.py
@@ -60,6 +60,7 @@ class GenerativeScoringRequest(OpenAIBaseModel):
             the full vocab for those ids (False).
         item_first: If True, prepend items to query. Otherwise append items to query.
         add_special_tokens: Whether to add special tokens when tokenizing.
+
     """
 
     model: str | None = None
@@ -112,6 +113,7 @@ class GenerativeScoringItemResult(OpenAIBaseModel):
         index: The index of this item in the input items list.
         object: Type of object, always "score".
         score: The probability score for the first label token.
+
     """
 
     index: int
@@ -129,6 +131,7 @@ class GenerativeScoringResponse(OpenAIBaseModel):
         model: The model used for scoring.
         data: List of scoring results, one per input item.
         usage: Token usage information.
+
     """
 
     id: str = Field(default="")
@@ -188,6 +191,7 @@ async def create_generative_scoring(
         Returns:
             GenerativeScoringResponse with probabilities for each item, or
             ErrorResponse if an error occurred.
+
         """
         # Check model
         error_check_ret = await self._check_model(request)  # type: ignore[arg-type]
@@ -401,6 +405,7 @@ async def _build_prompts(
 
         Returns:
             Tuple of (list of EngineInput, list of prompt token counts).
+
         """
         # Tokenize query if it's a string
         if isinstance(request.query, str):
@@ -455,6 +460,7 @@ def _compute_probabilities(
 
         Returns:
             Dictionary mapping token_id to probability.
+
         """
         if apply_softmax:
             # Normalize over only the label tokens (subset softmax)
diff --git a/vllm/entrypoints/grpc_server.py b/vllm/entrypoints/grpc_server.py
index 59269dd18020..913d084e4a81 100644
--- a/vllm/entrypoints/grpc_server.py
+++ b/vllm/entrypoints/grpc_server.py
@@ -1,200 +1,27 @@
-#!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# mypy: ignore-errors
-"""
-vLLM gRPC Server
+import warnings
 
-Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer
-from the smg-grpc-servicer package.
+from vllm.entrypoints.launchers.grpc_server import main, serve_grpc
 
-Usage:
-    python -m vllm.entrypoints.grpc_server --model <model_path>
+warnings.warn(
+    "`vllm.entrypoints.grpc_server is deprecated and will likely be "
+    "unsupported in a future version. Use the corresponding function from "
+    "`vllm.entrypoints.launchers.grpc_server` instead.",
+    DeprecationWarning,
+    stacklevel=1,
+)
 
-Example:
-    python -m vllm.entrypoints.grpc_server \
-        --model meta-llama/Llama-2-7b-hf \
-        --host 0.0.0.0 \
-        --port 50051
-"""
+__all__ = ["main", "serve_grpc"]
 
-import argparse
-import asyncio
-import signal
-import sys
-import time
 
-try:
-    import grpc
-    from grpc_health.v1 import health_pb2_grpc
-    from grpc_reflection.v1alpha import reflection
-    from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
-    from smg_grpc_servicer.vllm.health_servicer import VllmHealthServicer
-    from smg_grpc_servicer.vllm.servicer import VllmEngineServicer
-except ImportError as e:
-    raise ImportError(
-        "gRPC mode requires smg-grpc-servicer. "
-        "If not installed, run: pip install vllm[grpc]. "
-        "If already installed, there may be a broken import due to a "
-        "version mismatch — see the chained exception above for details."
-    ) from e
-
-import uvloop
-
-from vllm import envs
-from vllm.engine.arg_utils import AsyncEngineArgs
-from vllm.entrypoints.serve.utils.api_utils import log_version_and_model
-from vllm.logger import init_logger
-from vllm.usage.usage_lib import UsageContext
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm.v1.engine.async_llm import AsyncLLM
-from vllm.version import __version__ as VLLM_VERSION
-
-logger = init_logger(__name__)
-
-
-async def serve_grpc(args: argparse.Namespace):
-    """
-    Main gRPC serving function.
-
-    Args:
-        args: Parsed command line arguments
-    """
-    log_version_and_model(logger, VLLM_VERSION, args.model)
-    logger.info("vLLM gRPC server args: %s", args)
-
-    start_time = time.time()
-
-    # Create engine args
-    engine_args = AsyncEngineArgs.from_cli_args(args)
-
-    # Build vLLM config
-    vllm_config = engine_args.create_engine_config(
-        usage_context=UsageContext.OPENAI_API_SERVER,
-    )
-
-    # Create AsyncLLM
-    async_llm = AsyncLLM.from_vllm_config(
-        vllm_config=vllm_config,
-        usage_context=UsageContext.OPENAI_API_SERVER,
-        enable_log_requests=args.enable_log_requests,
-        disable_log_stats=args.disable_log_stats,
-    )
-
-    # Create servicer
-    servicer = VllmEngineServicer(async_llm, start_time)
-
-    # Create gRPC server
-    server = grpc.aio.server(
-        options=[
-            ("grpc.max_send_message_length", -1),
-            ("grpc.max_receive_message_length", -1),
-            # Tolerate client keepalive pings every 10s (default 300s is too
-            # strict for non-streaming requests where no DATA frames flow
-            # during generation)
-            ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000),
-            ("grpc.keepalive_permit_without_calls", True),
-        ],
-    )
-
-    # Add servicer to server
-    vllm_engine_pb2_grpc.add_VllmEngineServicer_to_server(servicer, server)
-
-    # Add standard gRPC health service for Kubernetes probes
-    health_servicer = VllmHealthServicer(async_llm)
-    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
-
-    # Enable reflection for grpcurl and other tools
-    service_names = (
-        vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name,
-        "grpc.health.v1.Health",
-        reflection.SERVICE_NAME,
-    )
-    reflection.enable_server_reflection(service_names, server)
-
-    # Bind to address
-    host = args.host or "0.0.0.0"
-    address = f"{host}:{args.port}"
-    server.add_insecure_port(address)
-
-    try:
-        # Start server
-        await server.start()
-        logger.info("vLLM gRPC server started on %s", address)
-        logger.info("Server is ready to accept requests")
-
-        # Start periodic stats logging (mirrors the HTTP server's lifespan task)
-        if not args.disable_log_stats:
-
-            async def _force_log():
-                while True:
-                    await asyncio.sleep(envs.VLLM_LOG_STATS_INTERVAL)
-                    await async_llm.do_log_stats()
-
-            stats_task = asyncio.create_task(_force_log())
-        else:
-            stats_task = None
-
-        # Handle shutdown signals
-        loop = asyncio.get_running_loop()
-        stop_event = asyncio.Event()
-
-        def signal_handler():
-            logger.info("Received shutdown signal")
-            stop_event.set()
-
-        for sig in (signal.SIGTERM, signal.SIGINT):
-            loop.add_signal_handler(sig, signal_handler)
-
-        try:
-            await stop_event.wait()
-        except KeyboardInterrupt:
-            logger.info("Interrupted by user")
-    finally:
-        logger.info("Shutting down vLLM gRPC server...")
-        if stats_task is not None:
-            stats_task.cancel()
-        try:
-            health_servicer.set_not_serving()
-        except Exception:  # broad: must not prevent server.stop() / shutdown()
-            logger.warning("Failed to set health status to NOT_SERVING", exc_info=True)
-        await server.stop(grace=5.0)
-        logger.info("gRPC server stopped")
-        async_llm.shutdown()
-        logger.info("AsyncLLM engine stopped")
-        logger.info("Shutdown complete")
-
-
-def main():
-    """Main entry point for python -m vllm.entrypoints.grpc_server."""
-    parser = FlexibleArgumentParser(
-        description="vLLM gRPC Server",
-    )
-
-    # Server args
-    parser.add_argument(
-        "--host",
-        type=str,
-        default="0.0.0.0",
-        help="Host to bind gRPC server to",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=50051,
-        help="Port to bind gRPC server to",
+if __name__ == "__main__":
+    warnings.warn(
+        "`python -m vllm.entrypoints.grpc_server is deprecated and "
+        "will likely be unsupported in a future version. "
+        "Please use `vllm serve <model_path> --grpc` instead.",
+        DeprecationWarning,
+        stacklevel=1,
     )
-    parser = AsyncEngineArgs.add_cli_args(parser)
 
-    args = parser.parse_args()
-
-    # Run server
-    try:
-        uvloop.run(serve_grpc(args))
-    except Exception as e:
-        logger.exception("Server failed: %s", e)
-        sys.exit(1)
-
-
-if __name__ == "__main__":
     main()
diff --git a/vllm/entrypoints/launchers/api_server/app_state.py b/vllm/entrypoints/launchers/api_server/app_state.py
index 09bb0ec63020..d33be659494c 100644
--- a/vllm/entrypoints/launchers/api_server/app_state.py
+++ b/vllm/entrypoints/launchers/api_server/app_state.py
@@ -95,6 +95,7 @@ async def init_app_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
+        tool_strict_level=args.tool_strict_level,
         reasoning_parser=args.structured_outputs_config.reasoning_parser,
         default_chat_template_kwargs=default_chat_template_kwargs,
         log_error_stack=args.log_error_stack,
@@ -111,6 +112,7 @@ async def init_app_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
+        tool_strict_level=args.tool_strict_level,
         reasoning_parser=args.structured_outputs_config.reasoning_parser,
         default_chat_template_kwargs=default_chat_template_kwargs,
         log_error_stack=args.log_error_stack,
diff --git a/vllm/entrypoints/launchers/api_server/entry.py b/vllm/entrypoints/launchers/api_server/entry.py
index 37d7ca381455..f3dcf40362a1 100644
--- a/vllm/entrypoints/launchers/api_server/entry.py
+++ b/vllm/entrypoints/launchers/api_server/entry.py
@@ -70,14 +70,12 @@ async def build_async_engine_client_from_engine_args(
     usage_context: UsageContext = UsageContext.OPENAI_API_SERVER,
     client_config: dict[str, Any] | None = None,
 ) -> AsyncIterator[EngineClient]:
-    """
-    Create EngineClient, either:
+    """Create EngineClient, either:
         - in-process using the AsyncLLMEngine Directly
         - multiprocess using AsyncLLMEngine RPC
 
     Returns the Client or None if the creation failed.
     """
-
     # Create the EngineConfig (determines if we can use V1).
     vllm_config = engine_args.create_engine_config(usage_context=usage_context)
 
@@ -123,7 +121,6 @@ async def build_and_serve(
 
     Returns the shutdown task for the caller to await.
     """
-
     # Get uvicorn log config (from file or with endpoint filter)
     log_config = get_uvicorn_log_config(args)
     if log_config is not None:
@@ -162,7 +159,6 @@ async def build_and_serve(
 
 async def run_server(args, **uvicorn_kwargs) -> None:
     """Run a single-worker API server."""
-
     decorate_logs("APIServer", skip_if_decorated=True)
 
     # Interrupt initialization if SIGTERM arrives before uvicorn installs its
@@ -180,7 +176,6 @@ async def run_server_worker(
     listen_address, sock, args, client_config=None, **uvicorn_kwargs
 ) -> None:
     """Run a single API server worker."""
-
     if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
         ToolParserManager.import_tool_parser(args.tool_parser_plugin)
 
diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py
index c0dd64164e95..488a7f2c43c8 100644
--- a/vllm/entrypoints/launchers/cli_args.py
+++ b/vllm/entrypoints/launchers/cli_args.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file contains the command line arguments for the vLLM's online server.
+"""This file contains the command line arguments for the vLLM's online server.
 It is kept in a separate file for documentation purposes.
 """
 
@@ -21,6 +20,7 @@
 )
 from vllm.entrypoints.openai.models.protocol import LoRAModulePath
 from vllm.tool_parsers import ToolParserManager
+from vllm.tool_parsers.tool_strict_level import ToolStrictLevelName
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from .utils.constants import (
@@ -47,7 +47,7 @@ def __call__(
             if item in [None, ""]:  # Skip if item is None or empty string
                 continue
             if "=" in item and "," not in item:  # Old format: name=path
-                name, path = item.split("=")
+                name, path = item.split("=", 1)
                 lora_list.append(LoRAModulePath(name, path))
             else:  # Assume JSON format
                 try:
@@ -117,6 +117,13 @@ class BaseFrontendArgs:
     """Special the tool parser plugin write to parse the model-generated tool
     into OpenAI API format, the name register in this plugin can be used in
     `--tool-call-parser`."""
+    tool_strict_level: ToolStrictLevelName = "auto"
+    """Server-side floor for structural-tag based tool calling, applied on top
+    of the per-tool `strict` field. `auto` follows the request's tool choice
+    and per-tool strictness; `function` constrains the
+    tool-call envelope (markup and function name) for every request with
+    tools; `parameter` additionally pins argument schemas, as if every tool
+    were `strict: true`."""
     tool_server: str | None = None
     """Comma-separated list of host:port pairs (IPv4, IPv6, or hostname).
     Examples: 127.0.0.1:8000, [::1]:8000, localhost:1234. Or `demo` for
@@ -181,6 +188,13 @@ class BaseFrontendArgs:
     If set to True, only enable the Tokens In<>Out endpoint.
     This is intended for use in a Disaggregated Everything setup.
     """
+    enable_scale_out: bool = False
+    """
+    If set to True, register the scale-out endpoints (`/render`, `/derender`,
+    and `/inference/v1/generate`) on `vllm serve`. Has no effect on
+    `vllm launch render` or `vllm serve --tokens-only`, which always register
+    their required endpoints regardless of this flag.
+    """
     fingerprint_mode: Literal["full", "hash", "custom", "none"] = "full"
     """Controls the ``system_fingerprint`` field on responses.
 
diff --git a/vllm/entrypoints/launchers/dp_supervisor.py b/vllm/entrypoints/launchers/dp_supervisor.py
index 7ebf887ccef1..91754b71e56b 100644
--- a/vllm/entrypoints/launchers/dp_supervisor.py
+++ b/vllm/entrypoints/launchers/dp_supervisor.py
@@ -182,8 +182,7 @@ async def _probe_endpoint(
     conn_err_failure_threshold: int = 3,
     conn_err_retry_delay: float = 5.0,
 ) -> bool:
-    """
-    Probe /health endpoint for 200 status.
+    """Probe /health endpoint for 200 status.
 
     If there is a connection error, retry every N seconds.
     """
@@ -251,9 +250,7 @@ def _run_rust_vllm_dp_server(child_args: argparse.Namespace) -> None:
 
 
 def _run_vllm_dp_server(child_args: argparse.Namespace) -> None:
-    """
-    Entrypoint function for the vLLM DP Server.
-    """
+    """Entrypoint function for the vLLM DP Server."""
     # Create a fresh process group for the vLLM DP Server,
     # so that CTRL-C is propagated cleanly.
     os.setpgrp()
@@ -320,8 +317,7 @@ async def run(self) -> None:
                 await supervisor_server_task
 
     async def _start_server(self) -> tuple[uvicorn.Server, asyncio.Task[None]]:
-        """
-        Launch the DPSupervisor HTTP server.
+        """Launch the DPSupervisor HTTP server.
 
         Called only after the vLLM DP Servers are ready so that /health does
         not return 503 to external probes while the engines are initializing.
@@ -366,8 +362,7 @@ async def _start_server(self) -> tuple[uvicorn.Server, asyncio.Task[None]]:
         return supervisor_server, supervisor_server_task
 
     async def _wait_until_ready(self, monitor_task: asyncio.Task[None]) -> None:
-        """
-        Block until the vLLM DP Servers are ready or shutdown is triggered.
+        """Block until the vLLM DP Servers are ready or shutdown is triggered.
 
         Returns early if monitoring stops (e.g. a DP Server dies during
         startup), in which case the supervisor server is never started.
@@ -379,14 +374,12 @@ async def _wait_until_ready(self, monitor_task: asyncio.Task[None]) -> None:
             await asyncio.sleep(0.05)
 
     def _handle_signal(self, signum: int) -> None:
-        """
-        Signal handler that is added to the event loop.
+        """Signal handler that is added to the event loop.
 
         This catches the SIGTERM from K8s and begins graceful shutdown,
         by setting the _shutdown_event(), which is watched by the main
         coroutine monitoring the vLLM DP Servers.
         """
-
         if self._shutdown_event.is_set():
             return
 
@@ -400,9 +393,7 @@ def _handle_signal(self, signum: int) -> None:
         self._is_ready = False
 
     def _start_children(self) -> None:
-        """
-        Launch vLLM DP Servers on separate GPUs.
-        """
+        """Launch vLLM DP Servers on separate GPUs."""
         logger.info("Launching vLLM DP Servers")
         context = multiprocessing.get_context("spawn")
         for local_rank in range(self.args.data_parallel_size_local):
@@ -416,8 +407,7 @@ def _start_children(self) -> None:
             self._processes.append(process)
 
     async def _probe_all_children(self) -> None:
-        """
-        Background coroutine: probes all child endpoints on each interval.
+        """Background coroutine: probes all child endpoints on each interval.
 
         Exits when any server becomes unhealthy after being ready, signalling
         _monitor_children to initiate shutdown.
@@ -474,8 +464,7 @@ async def _probe_all_children(self) -> None:
                     )
 
     async def _monitor_children(self) -> None:
-        """
-        Main coroutine task that monitors the children vLLM servers.
+        """Main coroutine task that monitors the children vLLM servers.
 
         Before the vLLM servers are /ready:
         - if the pid is dead, we will shut down
diff --git a/vllm/entrypoints/launchers/grpc_server.py b/vllm/entrypoints/launchers/grpc_server.py
new file mode 100644
index 000000000000..aab3f57c5ea7
--- /dev/null
+++ b/vllm/entrypoints/launchers/grpc_server.py
@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# mypy: ignore-errors
+"""vLLM gRPC Server
+
+Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer
+from the smg-grpc-servicer package.
+
+Usage:
+    python -m vllm.entrypoints.launchers.grpc_server --model <model_path>
+
+Example:
+    python -m vllm.entrypoints.launchers.grpc_server \
+        --model meta-llama/Llama-2-7b-hf \
+        --host 0.0.0.0 \
+        --port 50051
+
+"""
+
+import argparse
+import asyncio
+import signal
+import sys
+import time
+
+try:
+    import grpc
+    from grpc_health.v1 import health_pb2_grpc
+    from grpc_reflection.v1alpha import reflection
+    from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
+    from smg_grpc_servicer.vllm.health_servicer import VllmHealthServicer
+    from smg_grpc_servicer.vllm.servicer import VllmEngineServicer
+except ImportError as e:
+    raise ImportError(
+        "gRPC mode requires smg-grpc-servicer. "
+        "If not installed, run: pip install vllm[grpc]. "
+        "If already installed, there may be a broken import due to a "
+        "version mismatch — see the chained exception above for details."
+    ) from e
+
+import uvloop
+
+from vllm import envs
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.entrypoints.serve.utils.api_utils import log_version_and_model
+from vllm.logger import init_logger
+from vllm.usage.usage_lib import UsageContext
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.version import __version__ as VLLM_VERSION
+
+logger = init_logger(__name__)
+
+
+async def serve_grpc(args: argparse.Namespace):
+    """Main gRPC serving function.
+
+    Args:
+        args: Parsed command line arguments
+
+    """
+    log_version_and_model(logger, VLLM_VERSION, args.model)
+    logger.info("vLLM gRPC server args: %s", args)
+
+    start_time = time.time()
+
+    # Create engine args
+    engine_args = AsyncEngineArgs.from_cli_args(args)
+
+    # Build vLLM config
+    vllm_config = engine_args.create_engine_config(
+        usage_context=UsageContext.OPENAI_API_SERVER,
+    )
+
+    # Create AsyncLLM
+    async_llm = AsyncLLM.from_vllm_config(
+        vllm_config=vllm_config,
+        usage_context=UsageContext.OPENAI_API_SERVER,
+        enable_log_requests=args.enable_log_requests,
+        disable_log_stats=args.disable_log_stats,
+    )
+
+    # Create servicer
+    servicer = VllmEngineServicer(async_llm, start_time)
+
+    # Create gRPC server
+    server = grpc.aio.server(
+        options=[
+            ("grpc.max_send_message_length", -1),
+            ("grpc.max_receive_message_length", -1),
+            # Tolerate client keepalive pings every 10s (default 300s is too
+            # strict for non-streaming requests where no DATA frames flow
+            # during generation)
+            ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000),
+            ("grpc.keepalive_permit_without_calls", True),
+        ],
+    )
+
+    # Add servicer to server
+    vllm_engine_pb2_grpc.add_VllmEngineServicer_to_server(servicer, server)
+
+    # Add standard gRPC health service for Kubernetes probes
+    health_servicer = VllmHealthServicer(async_llm)
+    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
+
+    # Enable reflection for grpcurl and other tools
+    service_names = (
+        vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name,
+        "grpc.health.v1.Health",
+        reflection.SERVICE_NAME,
+    )
+    reflection.enable_server_reflection(service_names, server)
+
+    # Bind to address
+    host = args.host or "0.0.0.0"
+    address = f"{host}:{args.port}"
+    server.add_insecure_port(address)
+
+    try:
+        # Start server
+        await server.start()
+        logger.info("vLLM gRPC server started on %s", address)
+        logger.info("Server is ready to accept requests")
+
+        # Start periodic stats logging (mirrors the HTTP server's lifespan task)
+        if not args.disable_log_stats:
+
+            async def _force_log():
+                while True:
+                    await asyncio.sleep(envs.VLLM_LOG_STATS_INTERVAL)
+                    await async_llm.do_log_stats()
+
+            stats_task = asyncio.create_task(_force_log())
+        else:
+            stats_task = None
+
+        # Handle shutdown signals
+        loop = asyncio.get_running_loop()
+        stop_event = asyncio.Event()
+
+        def signal_handler():
+            logger.info("Received shutdown signal")
+            stop_event.set()
+
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            loop.add_signal_handler(sig, signal_handler)
+
+        try:
+            await stop_event.wait()
+        except KeyboardInterrupt:
+            logger.info("Interrupted by user")
+    finally:
+        logger.info("Shutting down vLLM gRPC server...")
+        if stats_task is not None:
+            stats_task.cancel()
+        try:
+            health_servicer.set_not_serving()
+        except Exception:  # broad: must not prevent server.stop() / shutdown()
+            logger.warning("Failed to set health status to NOT_SERVING", exc_info=True)
+        await server.stop(grace=5.0)
+        logger.info("gRPC server stopped")
+        async_llm.shutdown()
+        logger.info("AsyncLLM engine stopped")
+        logger.info("Shutdown complete")
+
+
+def main():
+    """Main entry point for python -m vllm.entrypoints.launchers.grpc_server."""
+    parser = FlexibleArgumentParser(
+        description="vLLM gRPC Server",
+    )
+
+    # Server args
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="0.0.0.0",
+        help="Host to bind gRPC server to",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=50051,
+        help="Port to bind gRPC server to",
+    )
+    parser = AsyncEngineArgs.add_cli_args(parser)
+
+    args = parser.parse_args()
+
+    # Run server
+    try:
+        uvloop.run(serve_grpc(args))
+    except Exception as e:
+        logger.exception("Server failed: %s", e)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/vllm/entrypoints/launchers/launcher.py b/vllm/entrypoints/launchers/launcher.py
index 09b0ea97811e..9fabfa3277af 100644
--- a/vllm/entrypoints/launchers/launcher.py
+++ b/vllm/entrypoints/launchers/launcher.py
@@ -53,8 +53,7 @@ async def serve_http(
     enable_ssl_refresh: bool = False,
     **uvicorn_kwargs: Any,
 ):
-    """
-    Start a FastAPI app using Uvicorn, with support for custom Uvicorn config
+    """Start a FastAPI app using Uvicorn, with support for custom Uvicorn config
     options.  Supports http header limits via h11_max_incomplete_event_size and
     h11_max_header_count.
     """
@@ -185,8 +184,7 @@ async def handle_shutdown() -> None:
 
 
 async def watchdog_loop(server: uvicorn.Server, engine: EngineClient):
-    """
-    # Watchdog task that runs in the background, checking
+    """# Watchdog task that runs in the background, checking
     # for error state in the engine. Needed to trigger shutdown
     # if an exception arises is StreamingResponse() generator.
     """
@@ -197,8 +195,7 @@ async def watchdog_loop(server: uvicorn.Server, engine: EngineClient):
 
 
 def terminate_if_errored(server: uvicorn.Server, engine: EngineClient):
-    """
-    See discussions here on shutting down a uvicorn server
+    """See discussions here on shutting down a uvicorn server
     https://github.com/encode/uvicorn/discussions/1103
     In this case we cannot await the server shutdown here
     because handler must first return to close the connection
@@ -254,7 +251,6 @@ def validate_api_server_args(args):
 @instrument(span_name="API server setup")
 def setup_server(args, *, reuse_port: bool):
     """Validate API server args and create the server socket."""
-
     log_version_and_model(logger, VLLM_VERSION, args.model)
     log_non_default_args(args)
 
diff --git a/vllm/entrypoints/launchers/render/app_state.py b/vllm/entrypoints/launchers/render/app_state.py
index 6fb4123c9b06..1d242977be29 100644
--- a/vllm/entrypoints/launchers/render/app_state.py
+++ b/vllm/entrypoints/launchers/render/app_state.py
@@ -31,7 +31,6 @@ async def init_render_app_state(
     preprocessing pipeline (renderer, input_processor)
     directly from the :class:`~vllm.config.VllmConfig`.
     """
-
     served_model_names = args.served_model_name or [args.model]
     model_registry = OpenAIModelRegistry(
         model_config=vllm_config.model_config,
@@ -61,6 +60,7 @@ async def init_render_app_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
+        tool_strict_level=args.tool_strict_level,
         reasoning_parser=args.reasoning_parser,
         default_chat_template_kwargs=default_chat_template_kwargs,
         log_error_stack=args.log_error_stack,
@@ -77,6 +77,7 @@ async def init_render_app_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
+        tool_strict_level=args.tool_strict_level,
         reasoning_parser=args.reasoning_parser,
         default_chat_template_kwargs=default_chat_template_kwargs,
         log_error_stack=args.log_error_stack,
diff --git a/vllm/entrypoints/launchers/render/entry.py b/vllm/entrypoints/launchers/render/entry.py
index c7465c1a12d3..b36cd6477214 100644
--- a/vllm/entrypoints/launchers/render/entry.py
+++ b/vllm/entrypoints/launchers/render/entry.py
@@ -30,7 +30,6 @@ async def build_and_serve_renderer(
 
     Returns the shutdown task for the caller to await.
     """
-
     # Get uvicorn log config (from file or with endpoint filter)
     log_config = get_uvicorn_log_config(args)
     if log_config is not None:
diff --git a/vllm/entrypoints/launchers/run_batch.py b/vllm/entrypoints/launchers/run_batch.py
index 8d1be95b7007..b5732ba65a16 100644
--- a/vllm/entrypoints/launchers/run_batch.py
+++ b/vllm/entrypoints/launchers/run_batch.py
@@ -53,8 +53,7 @@
 )
 from vllm.entrypoints.speech_to_text.transcription.protocol import (
     TranscriptionRequest,
-    TranscriptionResponse,
-    TranscriptionResponseVerbose,
+    TranscriptionResponseVariant,
 )
 from vllm.entrypoints.speech_to_text.translation.protocol import (
     TranslationRequest,
@@ -88,8 +87,7 @@ def _raise_audio_filesize_limit(size_bytes: int) -> NoReturn:
 
 
 class BatchTranscriptionRequest(TranscriptionRequest):
-    """
-    Batch transcription request that uses file_url instead of file.
+    """Batch transcription request that uses file_url instead of file.
 
     This class extends TranscriptionRequest but replaces the file field
     with file_url to support batch processing from audio files written in JSON format.
@@ -119,8 +117,7 @@ def validate_no_file(cls, data: Any):
 
 
 class BatchTranslationRequest(TranslationRequest):
-    """
-    Batch translation request that uses file_url instead of file.
+    """Batch translation request that uses file_url instead of file.
 
     This class extends TranslationRequest but replaces the file field
     with file_url to support batch processing from audio files written in JSON format.
@@ -160,8 +157,7 @@ def validate_no_file(cls, data: Any):
 
 
 class BatchRequestInput(OpenAIBaseModel):
-    """
-    The per-line object of the batch input file.
+    """The per-line object of the batch input file.
 
     NOTE: Currently only the `/v1/chat/completions` endpoint is supported.
     """
@@ -206,8 +202,7 @@ def check_type_for_url(cls, value: Any, info: ValidationInfo):
     | EmbeddingResponse
     | ScoreResponse
     | RerankResponse
-    | TranscriptionResponse
-    | TranscriptionResponseVerbose
+    | TranscriptionResponseVariant
     | TranslationResponse
     | TranslationResponseVerbose
 )
@@ -225,9 +220,7 @@ class BatchResponseData(OpenAIBaseModel):
 
 
 class BatchRequestOutput(OpenAIBaseModel):
-    """
-    The per-line object of the batch output and error files
-    """
+    """The per-line object of the batch output and error files"""
 
     id: str
 
@@ -358,8 +351,7 @@ async def read_file(path_or_url: str) -> str:
 async def write_local_file(
     output_path: str, batch_outputs: list[BatchRequestOutput]
 ) -> None:
-    """
-    Write the responses to a local file.
+    """Write the responses to a local file.
     output_path: The path to write the responses to.
     batch_outputs: The list of batch outputs to write.
     """
@@ -371,8 +363,7 @@ async def write_local_file(
 
 
 async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> None:
-    """
-    Upload a local file to a URL.
+    """Upload a local file to a URL.
     output_url: The URL to upload the file to.
     data_or_file: Either the data to upload or the path to the file to upload.
     from_file: If True, data_or_file is the path to the file to upload.
@@ -429,8 +420,7 @@ async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> No
 async def write_file(
     path_or_url: str, batch_outputs: list[BatchRequestOutput], output_tmp_dir: str
 ) -> None:
-    """
-    Write batch_outputs to a file or upload to a URL.
+    """Write batch_outputs to a file or upload to a URL.
     path_or_url: The path or URL to write batch_outputs to.
     batch_outputs: The list of batch outputs to write.
     output_tmp_dir: The directory to store the output file before uploading it
@@ -471,8 +461,7 @@ async def download_bytes_from_url(
     url: str,
     allowed_media_domains: list[str] | None = None,
 ) -> bytes:
-    """
-    Download data from a URL or decode from a data URL.
+    """Download data from a URL or decode from a data URL.
 
     Args:
         url: Either an HTTP/HTTPS URL or a data URL (data:...;base64,...)
@@ -482,6 +471,7 @@ async def download_bytes_from_url(
 
     Returns:
         Data as bytes
+
     """
     parsed = urlparse(url)
 
@@ -615,8 +605,7 @@ def handle_endpoint_request(
     handler_getter: Callable[[], Callable | None],
     wrapper_fn: WrapperFn | None = None,
 ) -> Awaitable[BatchRequestOutput] | None:
-    """
-    Generic handler for endpoint requests.
+    """Generic handler for endpoint requests.
 
     Args:
         request: The batch request input
@@ -628,6 +617,7 @@ def handle_endpoint_request(
     Returns:
         Awaitable[BatchRequestOutput] if the request was handled,
         None if URL didn't match
+
     """
     if not url_matcher(request.url):
         return None
@@ -649,8 +639,7 @@ def make_transcription_wrapper(
     is_translation: bool,
     allowed_media_domains: list[str] | None = None,
 ) -> WrapperFn:
-    """
-    Factory function to create a wrapper for transcription/translation handlers.
+    """Factory function to create a wrapper for transcription/translation handlers.
     The wrapper converts BatchTranscriptionRequest or BatchTranslationRequest
     to TranscriptionRequest or TranslationRequest and calls the appropriate handler.
 
@@ -662,14 +651,14 @@ def make_transcription_wrapper(
 
     Returns:
         A function that takes a handler and returns a wrapped handler
+
     """
 
     def wrapper(handler_fn: Callable):
         async def transcription_wrapper(
             batch_request_body: (BatchTranscriptionRequest | BatchTranslationRequest),
         ) -> (
-            TranscriptionResponse
-            | TranscriptionResponseVerbose
+            TranscriptionResponseVariant
             | TranslationResponse
             | TranslationResponseVerbose
             | ErrorResponse
@@ -723,8 +712,7 @@ async def build_endpoint_registry(
     engine_client: EngineClient,
     args: Namespace,
 ) -> dict[str, dict[str, Any]]:
-    """
-    Build the endpoint registry with all serving objects and handler configurations.
+    """Build the endpoint registry with all serving objects and handler configurations.
 
     Args:
         engine_client: The engine client
@@ -732,6 +720,7 @@ async def build_endpoint_registry(
 
     Returns:
         Dictionary mapping endpoint keys to their configurations
+
     """
     supported_tasks = await engine_client.get_supported_tasks()
     logger.info("Supported tasks: %s", supported_tasks)
diff --git a/vllm/entrypoints/launchers/utils/server_utils.py b/vllm/entrypoints/launchers/utils/server_utils.py
index 082582672381..91f5e620890b 100644
--- a/vllm/entrypoints/launchers/utils/server_utils.py
+++ b/vllm/entrypoints/launchers/utils/server_utils.py
@@ -29,8 +29,7 @@ def load_log_config(log_config_file: str | None) -> dict | None:
 
 
 def get_uvicorn_log_config(args: Namespace) -> dict | None:
-    """
-    Get the uvicorn log config based on the provided arguments.
+    """Get the uvicorn log config based on the provided arguments.
 
     Priority:
     1. If log_config_file is specified, use it
diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py
index e3d0b2f1e520..c0720ad8baee 100644
--- a/vllm/entrypoints/llm.py
+++ b/vllm/entrypoints/llm.py
@@ -76,8 +76,8 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
     Args:
         model: The name or path of a HuggingFace Transformers model.
         tokenizer: The name or path of a HuggingFace Transformers tokenizer.
-        tokenizer_mode: The tokenizer mode. "auto" will use the fast tokenizer
-            if available, and "slow" will always use the slow tokenizer.
+        tokenizer_mode: The tokenizer mode. See
+            [ModelConfig.tokenizer_mode][vllm.config.ModelConfig.tokenizer_mode].
         skip_tokenizer_init: If true, skip initialization of tokenizer and
             detokenizer. Expect valid prompt_token_ids and None for prompt
             from the input.
@@ -174,6 +174,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
     Note:
         This class is intended to be used for offline inference. For online
         serving, use the [AsyncLLMEngine][vllm.AsyncLLMEngine] class instead.
+
     """
 
     def __init__(
@@ -224,7 +225,6 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         """LLM constructor."""
-
         if "disable_log_stats" not in kwargs:
             kwargs["disable_log_stats"] = True
 
@@ -398,6 +398,7 @@ def get_world_size(self, include_dp: bool = True) -> int:
         Returns:
             The world size (tensor_parallel_size * pipeline_parallel_size),
             optionally multiplied by data_parallel_size if include_dp is True.
+
         """
         parallel_config = self.llm_engine.vllm_config.parallel_config
         if include_dp:
@@ -457,6 +458,7 @@ def generate(
         Returns:
             A list of `RequestOutput` objects containing the
             generated completions in the same order as the input prompts.
+
         """
         runner_type = self.model_config.runner_type
         if runner_type != "generate":
@@ -507,6 +509,7 @@ def enqueue(
 
         Returns:
             A list of request IDs for the enqueued requests.
+
         """
         runner_type = self.model_config.runner_type
         if runner_type != "generate":
@@ -558,6 +561,7 @@ def wait_for_completion(
 
         Returns:
             A list of output objects for all completed requests.
+
         """
         if output_type is None:
             output_type = (RequestOutput, PoolingRequestOutput)
@@ -571,8 +575,7 @@ def collective_rpc(
         args: tuple = (),
         kwargs: dict[str, Any] | None = None,
     ) -> list[_R]:
-        """
-        Execute an RPC call on all workers.
+        """Execute an RPC call on all workers.
 
         Args:
             method: Name of the worker method to execute, or a callable that
@@ -592,13 +595,12 @@ def collective_rpc(
         Note:
             It is recommended to use this API to only pass control messages,
             and set up data-plane communication to pass data.
-        """
 
+        """
         return self.llm_engine.collective_rpc(method, timeout, args, kwargs)
 
     def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
-        """
-        Run a function directly on the model inside each worker,
+        """Run a function directly on the model inside each worker,
         returning the result for each of them.
 
         !!! warning
@@ -625,8 +627,7 @@ def chat(
         tokenization_kwargs: dict[str, Any] | None = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
     ) -> list[RequestOutput]:
-        """
-        Generate responses for a chat conversation.
+        """Generate responses for a chat conversation.
 
         The chat conversation is converted into a text prompt using the
         tokenizer and calls the [generate][vllm.LLM.generate] method to generate
@@ -670,10 +671,12 @@ def chat(
                 template.
             tokenization_kwargs: Overrides for `tokenizer.encode`.
             mm_processor_kwargs: Overrides for `processor.__call__`.
+            tools: Tools to make available to the model, if any.
 
         Returns:
             A list of `RequestOutput` objects containing the generated
             responses in the same order as the input messages.
+
         """
         model_config = self.model_config
         runner_type = model_config.runner_type
@@ -751,6 +754,7 @@ def enqueue_chat(
 
         Returns:
             A list of request IDs for the enqueued requests.
+
         """
         model_config = self.model_config
         runner_type = model_config.runner_type
@@ -787,6 +791,7 @@ def start_profile(self, profile_prefix: str | None = None) -> None:
             profile_prefix: Optional prefix for the trace file names. If provided,
                            trace files will be named as "<prefix>_dp<X>_pp<Y>_tp<Z>".
                            If not provided, default naming will be used.
+
         """
         self.llm_engine.start_profile(profile_prefix)
 
@@ -801,8 +806,7 @@ def reset_prefix_cache(
         )
 
     def sleep(self, level: int = 1, mode: PauseMode = "abort"):
-        """
-        Put the engine to sleep. The engine should not process any requests.
+        """Put the engine to sleep. The engine should not process any requests.
         The caller should guarantee that no requests are being processed
         during the sleep period, before `wake_up` is called.
 
@@ -822,12 +826,20 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort"):
                            CPU memory pressure.
             mode: How to handle any existing requests, can be "abort", "wait",
                 or "keep".
+
         """
         self.llm_engine.sleep(level=level, mode=mode)
 
-    def wake_up(self, tags: list[str] | None = None):
+    def release_kv_cache_memory(self) -> None:
+        """Release the GPU physical memory backing the KV cache.
+
+        Requires a completed ``sleep(level=0)`` and resident executor memory.
+        Restore with ``wake_up(tags=["kv_cache"])``; kept requests are recomputed.
         """
-        Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep]
+        self.llm_engine.release_kv_cache_memory()
+
+    def wake_up(self, tags: list[str] | None = None):
+        """Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep]
         method for more details.
 
         Args:
@@ -837,6 +849,7 @@ def wake_up(self, tags: list[str] | None = None):
                 is reallocated. wake_up should be called with all tags
                 (or None) before the engine is used again.
                 Use tags=["scheduling"] to resume from level 0 sleep.
+
         """
         self.llm_engine.wake_up(tags)
 
@@ -849,17 +862,18 @@ def get_metrics(self) -> list["Metric"]:
 
         Note:
             This method is only available with the V1 LLM engine.
+
         """
         return self.llm_engine.get_metrics()
 
     def init_weight_transfer_engine(
         self, request: WeightTransferInitRequest | dict
     ) -> None:
-        """
-        Initialize weight transfer for RL training.
+        """Initialize weight transfer for RL training.
 
         Args:
             request: Weight transfer initialization request with backend-specific info
+
         """
         init_info_dict = (
             request["init_info"] if isinstance(request, dict) else request.init_info
@@ -878,11 +892,11 @@ def start_draft_weight_update(self) -> None:
         self.llm_engine.collective_rpc("start_draft_weight_update")
 
     def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None:
-        """
-        Update the weights of the model.
+        """Update the weights of the model.
 
         Args:
             request: Weight update request with backend-specific update info
+
         """
         update_info_dict = (
             request["update_info"] if isinstance(request, dict) else request.update_info
diff --git a/vllm/entrypoints/mcp/tool.py b/vllm/entrypoints/mcp/tool.py
index cd25aef087f8..d1f9130bdf22 100644
--- a/vllm/entrypoints/mcp/tool.py
+++ b/vllm/entrypoints/mcp/tool.py
@@ -23,8 +23,7 @@
 
 
 def validate_gpt_oss_install():
-    """
-    Check if the gpt-oss is installed and its version is at least 0.0.7.
+    """Check if the gpt-oss is installed and its version is at least 0.0.7.
     If not, raise an ImportError.
     """
     from importlib.metadata import PackageNotFoundError, version
@@ -151,8 +150,7 @@ async def get_result(self, context: "ConversationContext") -> Any:
         return tool_output_msgs
 
     async def get_result_parsable_context(self, context: "ConversationContext") -> Any:
-        """
-        This function converts parsable context types to harmony and
+        """This function converts parsable context types to harmony and
         back so we can use GPTOSS demo python tool
         """
         from vllm.entrypoints.openai.responses.context import ParsableContext
diff --git a/vllm/entrypoints/mcp/tool_server.py b/vllm/entrypoints/mcp/tool_server.py
index ed519ce4a51d..a0fe518f2d81 100644
--- a/vllm/entrypoints/mcp/tool_server.py
+++ b/vllm/entrypoints/mcp/tool_server.py
@@ -76,17 +76,14 @@ def post_process_tools_description(
 class ToolServer(ABC):
     @abstractmethod
     def has_tool(self, tool_name: str) -> bool:
-        """
-        Return True if the tool is supported, False otherwise.
-        """
+        """Return True if the tool is supported, False otherwise."""
         pass
 
     @abstractmethod
     def get_tool_description(
         self, tool_name: str, allowed_tools: list[str] | None = None
     ) -> ToolNamespaceConfig | None:
-        """
-        Return the tool description for the given tool name.
+        """Return the tool description for the given tool name.
         If the tool is not supported, return None.
         """
         pass
@@ -95,9 +92,7 @@ def get_tool_description(
     def new_session(
         self, tool_name: str, session_id: str, headers: dict[str, str] | None = None
     ) -> AbstractAsyncContextManager[Any]:
-        """
-        Create a session for the tool.
-        """
+        """Create a session for the tool."""
         ...
 
 
diff --git a/vllm/entrypoints/offline_utils.py b/vllm/entrypoints/offline_utils.py
index 650860c3377b..3f75c3033dd6 100644
--- a/vllm/entrypoints/offline_utils.py
+++ b/vllm/entrypoints/offline_utils.py
@@ -48,7 +48,7 @@
 
 
 class OfflineInferenceMixin:
-    """Offline inference utils"""
+    """Offline inference utils."""
 
     request_counter: Counter
     renderer: BaseRenderer
@@ -113,14 +113,14 @@ def _preprocess_cmpl(
         tokenization_kwargs: dict[str, Any] | None = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
     ) -> Sequence[EngineInput]:
-        """
-        Convert prompt inputs from LLM APIs (other than [LLM.chat][]) into
+        """Convert prompt inputs from LLM APIs (other than [LLM.chat][]) into
         a format that can be passed to `_add_request`.
 
         Refer to [LLM.generate][] for a complete description of the arguments.
 
         Returns:
             A list of `EngineInput` objects ready to be passed into LLMEngine.
+
         """
         renderer = self.renderer
         model_config = self.model_config
@@ -168,14 +168,14 @@ def _preprocess_chat(
         tokenization_kwargs: dict[str, Any] | None = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
     ) -> Sequence[EngineInput]:
-        """
-        Convert a list of conversations into prompts so that they can then
+        """Convert a list of conversations into prompts so that they can then
         be used as input for other LLM APIs.
 
         Refer to [LLM.chat][] for a complete description of the arguments.
 
         Returns:
             A list of `EngineInput` objects ready to be passed into LLMEngine.
+
         """
         renderer = self.renderer
 
diff --git a/vllm/entrypoints/openai/chat_completion/api_router.py b/vllm/entrypoints/openai/chat_completion/api_router.py
index 291907a77a56..aecf5bba4483 100644
--- a/vllm/entrypoints/openai/chat_completion/api_router.py
+++ b/vllm/entrypoints/openai/chat_completion/api_router.py
@@ -14,7 +14,6 @@
     ChatCompletionResponse,
 )
 from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
-from vllm.entrypoints.openai.sse_keep_alive import with_sse_keep_alive
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.utils.api_utils import (
     load_aware_call,
@@ -22,6 +21,7 @@
     with_cancellation,
 )
 from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
+from vllm.entrypoints.serve.utils.sse_keep_alive import with_sse_keep_alive
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py
index 06b3dbfd4a4c..ea6601bff13b 100644
--- a/vllm/entrypoints/openai/chat_completion/batch_serving.py
+++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py
@@ -50,6 +50,7 @@ async def render_batch_chat_request(
         Returns:
             A tuple of (all_conversations, engine_prompts) on success — one
             entry per conversation — or an ErrorResponse on failure.
+
         """
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py
index b7b292e8355d..5c1bc594d70e 100644
--- a/vllm/entrypoints/openai/chat_completion/protocol.py
+++ b/vllm/entrypoints/openai/chat_completion/protocol.py
@@ -215,7 +215,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
     # https://platform.openai.com/docs/api-reference/chat/create
     messages: list[ChatCompletionMessageParam]
     model: str | None = None
-    frequency_penalty: float | None = 0.0
+    frequency_penalty: float | None = None
     logit_bias: dict[str, float] | None = None
     logprobs: bool | None = False
     top_logprobs: int | None = 0
@@ -226,7 +226,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
     )
     max_completion_tokens: int | None = None
     n: int | None = 1
-    presence_penalty: float | None = 0.0
+    presence_penalty: float | None = None
     response_format: AnyResponseFormat | None = None
     seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
     stop: StopParam = []
@@ -453,18 +453,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
         ),
     )
 
-    return_assistant_tokens_mask: bool = Field(
-        default=False,
-        description=(
-            "If true, the /render response will include an "
-            "``assistant_tokens_mask`` field — a per-token list of 0/1 "
-            "values indicating which tokens were assistant-generated. "
-            "Requires the chat template to use ``{% generation %}`` "
-            "tags.  When the template does not support it, "
-            "``assistant_tokens_mask`` will be ``null``."
-        ),
-    )
-
     cache_salt: str | None = Field(
         default=None,
         min_length=1,
@@ -605,7 +593,6 @@ def build_chat_params(
                 extra_kwargs,
             ),
             media_io_kwargs=self.media_io_kwargs,
-            return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask),
             # No-tools requests default to tool_choice="none" at the API
             # layer. Collapse that default before rendering, so K3 emits a
             # model-visible tool-choice instruction only for requests with a
@@ -637,6 +624,8 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
     # Default sampling parameters for chat completion requests
     _DEFAULT_SAMPLING_PARAMS: dict = {
         "repetition_penalty": 1.0,
+        "presence_penalty": 0.0,
+        "frequency_penalty": 0.0,
         "temperature": 1.0,
         "top_p": 1.0,
         "top_k": 0,
@@ -659,6 +648,7 @@ def to_beam_search_params(
             temperature=temperature,
             length_penalty=self.length_penalty,
             include_stop_str_in_output=self.include_stop_str_in_output,
+            skip_special_tokens=self.skip_special_tokens,
         )
 
     def extract_structured_outputs(self) -> StructuredOutputsParams | None:
@@ -673,28 +663,15 @@ def to_sampling_params(
         max_tokens: int,
         default_sampling_params: dict,
     ) -> SamplingParams:
-        # Default parameters
-        if (repetition_penalty := self.repetition_penalty) is None:
-            repetition_penalty = default_sampling_params.get(
-                "repetition_penalty",
-                self._DEFAULT_SAMPLING_PARAMS["repetition_penalty"],
-            )
-        if (temperature := self.temperature) is None:
-            temperature = default_sampling_params.get(
-                "temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"]
-            )
-        if (top_p := self.top_p) is None:
-            top_p = default_sampling_params.get(
-                "top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"]
-            )
-        if (top_k := self.top_k) is None:
-            top_k = default_sampling_params.get(
-                "top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"]
-            )
-        if (min_p := self.min_p) is None:
-            min_p = default_sampling_params.get(
-                "min_p", self._DEFAULT_SAMPLING_PARAMS["min_p"]
+        # Priority: user -> server default -> OpenAI default
+        sampling_params = {
+            name: (
+                value
+                if (value := getattr(self, name)) is not None
+                else default_sampling_params.get(name, default)
             )
+            for name, default in self._DEFAULT_SAMPLING_PARAMS.items()
+        }
 
         # Merge server-default stop_token_ids (e.g., model-specific tokens
         # like </call> for gpt-oss) with any request-specified ones
@@ -721,14 +698,8 @@ def to_sampling_params(
             extra_args["ec_transfer_params"] = self.ec_transfer_params
         return SamplingParams.from_optional(
             n=self.n,
-            presence_penalty=self.presence_penalty,
-            frequency_penalty=self.frequency_penalty,
-            repetition_penalty=repetition_penalty,
-            temperature=temperature,
+            **sampling_params,
             watermarking=self.watermarking,
-            top_p=top_p,
-            top_k=top_k,
-            min_p=min_p,
             seed=self.seed,
             stop=self.stop,
             stop_token_ids=stop_token_ids,
diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py
index 06e7315c2498..45ea3203cdb6 100644
--- a/vllm/entrypoints/openai/chat_completion/serving.py
+++ b/vllm/entrypoints/openai/chat_completion/serving.py
@@ -132,6 +132,7 @@ def __init__(
         enable_auto_tools: bool = False,
         exclude_tools_when_tool_choice_none: bool = False,
         tool_parser: str | None = None,
+        tool_strict_level: str = "auto",
         enable_prompt_tokens_details: bool = False,
         enable_force_include_usage: bool = False,
         enable_log_outputs: bool = False,
@@ -161,6 +162,7 @@ def __init__(
             tool_parser_name=tool_parser,
             reasoning_parser_name=reasoning_parser,
             enable_auto_tools=enable_auto_tools,
+            tool_strict_level=tool_strict_level,
             model_name=self.model_config.model,
             is_harmony=self.model_config.hf_config.model_type == "gpt_oss",
         )
@@ -220,8 +222,7 @@ async def render_chat_request(
         self,
         request: ChatCompletionRequest,
     ) -> tuple[list[ConversationMessage], list[EngineInput]] | ErrorResponse:
-        """
-        Validate the model and preprocess a chat completion request.
+        """Validate the model and preprocess a chat completion request.
 
         Delegates preprocessing logic to OnlineRenderer, adding the
         engine-aware checks (LoRA model validation, engine health).
@@ -229,6 +230,7 @@ async def render_chat_request(
         Returns:
             A tuple of (conversation, engine_inputs) on success,
             or an ErrorResponse on failure.
+
         """
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
@@ -244,8 +246,7 @@ async def create_chat_completion(
         request: ChatCompletionRequest,
         raw_request: Request | None = None,
     ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:
-        """
-        Chat Completion API similar to OpenAI's API.
+        """Chat Completion API similar to OpenAI's API.
 
         See https://platform.openai.com/docs/api-reference/chat/create
         for the API specification. This API mimics the OpenAI
@@ -756,7 +757,11 @@ async def chat_completion_stream_generator(
                         # finish_reason is:
                         # "tool_calls" for "auto" or "required" tool calls,
                         # and "stop" for named tool calls.
-                        if tools_streamed[i] and not tool_choice_function_name:
+                        if (
+                            tools_streamed[i]
+                            and not tool_choice_function_name
+                            and output.finish_reason == "stop"
+                        ):
                             finish_reason_ = "tool_calls"
                         else:
                             finish_reason_ = (
diff --git a/vllm/entrypoints/openai/completion/api_router.py b/vllm/entrypoints/openai/completion/api_router.py
index 7b33fcdb6cc0..fe909c7baee6 100644
--- a/vllm/entrypoints/openai/completion/api_router.py
+++ b/vllm/entrypoints/openai/completion/api_router.py
@@ -12,7 +12,6 @@
     CompletionResponse,
 )
 from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
-from vllm.entrypoints.openai.sse_keep_alive import with_sse_keep_alive
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.utils.api_utils import (
     load_aware_call,
@@ -20,6 +19,7 @@
     with_cancellation,
 )
 from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
+from vllm.entrypoints.serve.utils.sse_keep_alive import with_sse_keep_alive
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py
index 6d7a49bacd79..a0974cb91c69 100644
--- a/vllm/entrypoints/openai/completion/protocol.py
+++ b/vllm/entrypoints/openai/completion/protocol.py
@@ -55,12 +55,12 @@ class CompletionRequest(OpenAIBaseModel):
         | None
     ) = None
     echo: bool | None = False
-    frequency_penalty: float | None = 0.0
+    frequency_penalty: float | None = None
     logit_bias: dict[str, float] | None = None
     logprobs: int | None = None
     max_tokens: int | None = 16
     n: int = 1
-    presence_penalty: float | None = 0.0
+    presence_penalty: float | None = None
     seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
     stop: StopParam = []
     stream: bool | None = False
@@ -278,6 +278,8 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
     # Default sampling parameters for completion requests
     _DEFAULT_SAMPLING_PARAMS: dict = {
         "repetition_penalty": 1.0,
+        "presence_penalty": 0.0,
+        "frequency_penalty": 0.0,
         "temperature": 1.0,
         "top_p": 1.0,
         "top_k": 0,
@@ -303,6 +305,7 @@ def to_beam_search_params(
             temperature=temperature,
             length_penalty=self.length_penalty,
             include_stop_str_in_output=self.include_stop_str_in_output,
+            skip_special_tokens=self.skip_special_tokens,
         )
 
     def extract_structured_outputs(self) -> StructuredOutputsParams | None:
@@ -320,28 +323,15 @@ def to_sampling_params(
         if default_sampling_params is None:
             default_sampling_params = {}
 
-        # Default parameters
-        if (repetition_penalty := self.repetition_penalty) is None:
-            repetition_penalty = default_sampling_params.get(
-                "repetition_penalty",
-                self._DEFAULT_SAMPLING_PARAMS["repetition_penalty"],
-            )
-        if (temperature := self.temperature) is None:
-            temperature = default_sampling_params.get(
-                "temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"]
-            )
-        if (top_p := self.top_p) is None:
-            top_p = default_sampling_params.get(
-                "top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"]
-            )
-        if (top_k := self.top_k) is None:
-            top_k = default_sampling_params.get(
-                "top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"]
-            )
-        if (min_p := self.min_p) is None:
-            min_p = default_sampling_params.get(
-                "min_p", self._DEFAULT_SAMPLING_PARAMS["min_p"]
+        # Priority: user -> server default -> OpenAI default
+        sampling_params = {
+            name: (
+                value
+                if (value := getattr(self, name)) is not None
+                else default_sampling_params.get(name, default)
             )
+            for name, default in self._DEFAULT_SAMPLING_PARAMS.items()
+        }
 
         # Merge server-default stop_token_ids (e.g., model-specific tokens
         # like </call> for gpt-oss) with any request-specified ones
@@ -370,14 +360,8 @@ def to_sampling_params(
             extra_args["ec_transfer_params"] = self.ec_transfer_params
         return SamplingParams.from_optional(
             n=self.n,
-            presence_penalty=self.presence_penalty,
-            frequency_penalty=self.frequency_penalty,
-            repetition_penalty=repetition_penalty,
-            temperature=temperature,
+            **sampling_params,
             watermarking=self.watermarking,
-            top_p=top_p,
-            top_k=top_k,
-            min_p=min_p,
             seed=self.seed,
             stop=self.stop,
             stop_token_ids=stop_token_ids,
diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py
index 2fa01478cee6..d7e235da1ba1 100644
--- a/vllm/entrypoints/openai/completion/serving.py
+++ b/vllm/entrypoints/openai/completion/serving.py
@@ -89,14 +89,14 @@ async def render_completion_request(
         self,
         request: CompletionRequest,
     ) -> list[EngineInput] | ErrorResponse:
-        """
-        Validate the model and preprocess a completion request.
+        """Validate the model and preprocess a completion request.
 
         Delegates preprocessing logic to OnlineRenderer, adding the
         engine-aware checks (LoRA model validation, engine health).
 
         Returns:
             A list of engine_inputs on success, or an ErrorResponse on failure.
+
         """
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
@@ -116,9 +116,7 @@ async def create_completion(
         See https://platform.openai.com/docs/api-reference/completions/create
         for the API specification. This API mimics the OpenAI Completion API.
 
-        NOTE: Currently we do not support the following feature:
-            - suffix (the language models we currently support do not support
-            suffix)
+        NOTE: suffix is only supported by models that implement FIM rendering.
         """
         return await self._with_kv_transfer_rejection_cleanup(
             self._create_completion(request, raw_request), request, raw_request
diff --git a/vllm/entrypoints/openai/models/serving.py b/vllm/entrypoints/openai/models/serving.py
index 7378d37c6ecf..1e13311dce2a 100644
--- a/vllm/entrypoints/openai/models/serving.py
+++ b/vllm/entrypoints/openai/models/serving.py
@@ -289,6 +289,7 @@ async def resolve_lora(self, lora_name: str) -> LoRARequest | ErrorResponse:
             LoRARequest if found and loaded successfully.
             ErrorResponse (404) if no resolver finds the adapter.
             ErrorResponse (400) if adapter(s) are found but none load.
+
         """
         async with self.lora_resolver_lock[lora_name]:
             # First check if this LoRA is already loaded
diff --git a/vllm/entrypoints/openai/parser/harmony_utils.py b/vllm/entrypoints/openai/parser/harmony_utils.py
index fca02a91cc49..873b9e92bdab 100644
--- a/vllm/entrypoints/openai/parser/harmony_utils.py
+++ b/vllm/entrypoints/openai/parser/harmony_utils.py
@@ -88,8 +88,7 @@ def extract_function_from_recipient(recipient: str) -> str:
 
 
 def has_custom_tools(tool_types: set[str]) -> bool:
-    """
-    Checks if the given tool types are custom tools
+    """Checks if the given tool types are custom tools
     (i.e. any tool other than MCP builtin tools)
     """
     return not tool_types.issubset(MCP_BUILTIN_TOOLS)
@@ -209,8 +208,7 @@ def get_system_or_developer_message(role: str, instructions: str) -> Message:
 
 
 def parse_chat_inputs_to_harmony_messages(chat_msgs: list) -> list[Message]:
-    """
-    Parse a list of messages from request.messages in the Chat Completion API to
+    """Parse a list of messages from request.messages in the Chat Completion API to
     Harmony messages.
     """
     msgs: list[Message] = []
@@ -230,9 +228,8 @@ def parse_chat_inputs_to_harmony_messages(chat_msgs: list) -> list[Message]:
 
 
 def auto_drop_analysis_messages(msgs: list[Message]) -> list[Message]:
-    """
-    Harmony models expect the analysis messages (representing raw chain of thought) to
-    be dropped after an assistant message to the final channel is produced from the
+    """Harmony models expect the analysis messages (representing raw chain of thought)
+    to be dropped after an assistant message to the final channel is produced from the
     reasoning of those messages.
 
     The openai-harmony library does this if the very last assistant message is to the
@@ -262,8 +259,7 @@ def auto_drop_analysis_messages(msgs: list[Message]) -> list[Message]:
 
 
 def flatten_input_text_content(content: Any) -> str | None:
-    """
-    Extract text parts from a Chat Completion or Responses API content field and
+    """Extract text parts from a Chat Completion or Responses API content field and
     flatten them into a single string. Returns None if no text content is found.
     """
     if content is None or isinstance(content, str):
@@ -286,8 +282,7 @@ def flatten_input_text_content(content: Any) -> str | None:
 def extract_instructions_from_messages(
     messages: Sequence[Any],
 ) -> tuple[str | None, list[Any]]:
-    """
-    Peel a leading system/developer Chat Completion or Responses message and
+    """Peel a leading system/developer Chat Completion or Responses message and
     flatten its instruction text.
     """
     remaining_messages = list(messages)
@@ -327,9 +322,7 @@ def build_harmony_preamble(
     container_description: str | None = None,
     with_custom_tools: bool = False,
 ) -> list[Message]:
-    """
-    Build the standard Harmony system/developer prefix for a request.
-    """
+    """Build the standard Harmony system/developer prefix for a request."""
     developer_instructions = system_instructions = None
     if envs.VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS:
         system_instructions = instructions
@@ -359,8 +352,7 @@ def build_harmony_preamble(
 def parse_chat_input_to_harmony_message(
     chat_msg, tool_id_names: dict[str, str] | None = None
 ) -> list[Message]:
-    """
-    Parse a message from request.messages in the Chat Completion API to
+    """Parse a message from request.messages in the Chat Completion API to
     Harmony messages.
     """
     tool_id_names = tool_id_names or {}
diff --git a/vllm/entrypoints/openai/responses/api_router.py b/vllm/entrypoints/openai/responses/api_router.py
index 08b2b5584ec4..ee352dd1933c 100644
--- a/vllm/entrypoints/openai/responses/api_router.py
+++ b/vllm/entrypoints/openai/responses/api_router.py
@@ -34,7 +34,7 @@ def responses(request: Request) -> OpenAIServingResponses | None:
 async def _convert_stream_to_sse_events(
     generator: AsyncGenerator[StreamingResponsesResponse, None],
 ) -> AsyncGenerator[str, None]:
-    """Convert the generator to a stream of events in SSE format"""
+    """Convert the generator to a stream of events in SSE format."""
     async for event in generator:
         event_type = getattr(event, "type", "unknown")
         # https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#event_stream_format
diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py
index 2268bb0700a0..0db47cf367e6 100644
--- a/vllm/entrypoints/openai/responses/context.py
+++ b/vllm/entrypoints/openai/responses/context.py
@@ -45,6 +45,8 @@
 if TYPE_CHECKING:
     from mcp.client import ClientSession
 
+    from vllm.v1.metrics.stats import RequestStateStats
+
 logger = logging.getLogger(__name__)
 
 # This is currently needed as the tool type doesn't 1:1 match the
@@ -102,6 +104,11 @@ def copy(self) -> "TurnMetrics":
 
 class ConversationContext(ABC):
     response_parser: Parser | None = None
+    request_metrics: "RequestStateStats | None" = None
+    # Built-in tools can trigger additional model generations. In that case,
+    # the stored engine timestamps cover only one turn, while token usage is
+    # accumulated across all turns.
+    request_metrics_cover_all_generation_turns: bool = True
 
     @abstractmethod
     def append_output(self, output: RequestOutput) -> None:
@@ -141,9 +148,7 @@ async def cleanup_session(self) -> None:
 def _create_json_parse_error_messages(
     last_msg: Message, e: json.JSONDecodeError
 ) -> list[Message]:
-    """
-    Creates an error message when json parse failed.
-    """
+    """Creates an error message when json parse failed."""
     error_msg = (
         f"Error parsing tool arguments as JSON: {str(e)}. "
         "Please ensure the tool call arguments are valid JSON and try again."
@@ -161,7 +166,7 @@ def _create_json_parse_error_messages(
 
 
 class SimpleContext(ConversationContext):
-    """This is a context that cannot handle MCP tool calls"""
+    """This is a context that cannot handle MCP tool calls."""
 
     def __init__(
         self,
@@ -191,6 +196,7 @@ def __init__(
         self.num_prompt_tokens = 0
         self.num_output_tokens = 0
         self.num_cached_tokens = 0
+        self.num_cache_creation_tokens = 0
         # todo num_reasoning_tokens is not implemented yet.
         self.num_reasoning_tokens = 0
         # not implemented yet for SimpleContext
@@ -206,6 +212,8 @@ def append_output(self, output) -> None:
             raise ValueError("SimpleContext only supports RequestOutput.")
         self.num_prompt_tokens = len(output.prompt_token_ids or [])
         self.num_cached_tokens = output.num_cached_tokens or 0
+        if output.num_cache_creation_tokens is not None:
+            self.num_cache_creation_tokens = output.num_cache_creation_tokens
         self.num_output_tokens += len(output.outputs[0].token_ids or [])
         if output.kv_transfer_params is not None:
             self.kv_transfer_params = output.kv_transfer_params
@@ -474,8 +482,7 @@ async def call_search_tool(
     async def call_container_tool(
         self, tool_session: Union["ClientSession", Tool], last_msg: Message
     ) -> list[Message]:
-        """
-        Call container tool. Expect this to be run in a stateful docker
+        """Call container tool. Expect this to be run in a stateful docker
         with command line terminal.
         The official container tool would at least
         expect the following format:
@@ -565,6 +572,7 @@ async def init_tool_sessions(
         mcp_tools: dict[str, Mcp],
     ):
         if tool_server:
+            initialized_session = False
             for tool_name in self.available_tools:
                 if tool_name in self._tool_sessions:
                     continue
@@ -577,10 +585,12 @@ async def init_tool_sessions(
                     tool_server.new_session(tool_name, request_id, headers)
                 )
                 self._tool_sessions[tool_name] = tool_session
+                initialized_session = True
+            if initialized_session:
                 exit_stack.push_async_exit(self.cleanup_session)
 
     async def cleanup_session(self, *args, **kwargs) -> None:
-        """Can be used as coro to used in __aexit__"""
+        """Can be used as coro to used in __aexit__."""
 
         async def cleanup_tool_session(tool_session):
             if not isinstance(tool_session, Tool):
@@ -690,6 +700,7 @@ def _update_prefill_token_usage(self, output: RequestOutput) -> None:
 
         Args:
             output: The RequestOutput containing prompt token information
+
         """
         if output.prompt_token_ids is not None:
             this_turn_input_tokens = len(output.prompt_token_ids)
@@ -754,6 +765,7 @@ def _update_decode_token_usage(self, output: RequestOutput) -> int:
 
         Returns:
             int: Number of output tokens processed in this call
+
         """
         updated_output_token_count = 0
         if output.outputs:
@@ -863,6 +875,7 @@ async def init_tool_sessions(
         mcp_tools: dict[str, Mcp],
     ):
         if tool_server:
+            initialized_session = False
             for tool_name in self.available_tools:
                 if tool_name not in self._tool_sessions:
                     tool_type = _map_tool_name_to_tool_type(tool_name)
@@ -873,13 +886,14 @@ async def init_tool_sessions(
                         tool_server.new_session(tool_name, request_id, headers)
                     )
                     self._tool_sessions[tool_name] = tool_session
-                    exit_stack.push_async_exit(self.cleanup_session)
+                    initialized_session = True
+            if initialized_session:
+                exit_stack.push_async_exit(self.cleanup_session)
 
     async def call_container_tool(
         self, tool_session: Union["ClientSession", Tool], last_msg: Message
     ) -> list[Message]:
-        """
-        Call container tool. Expect this to be run in a stateful docker
+        """Call container tool. Expect this to be run in a stateful docker
         with command line terminal.
         The official container tool would at least
         expect the following format:
@@ -919,7 +933,7 @@ async def call_container_tool(
         ]
 
     async def cleanup_session(self, *args, **kwargs) -> None:
-        """Can be used as coro to used in __aexit__"""
+        """Can be used as coro to used in __aexit__."""
 
         async def cleanup_tool_session(tool_session):
             if not isinstance(tool_session, Tool):
diff --git a/vllm/entrypoints/openai/responses/harmony.py b/vllm/entrypoints/openai/responses/harmony.py
index 4e62828fd165..1aae70a381dc 100644
--- a/vllm/entrypoints/openai/responses/harmony.py
+++ b/vllm/entrypoints/openai/responses/harmony.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Harmony ↔ Responses API conversion utilities.
+"""Harmony ↔ Responses API conversion utilities.
 
 Handles two directions:
   1. Response Input → Harmony Messages  (input parsing)
@@ -303,7 +302,7 @@ def _parse_browser_tool_call(message: Message, recipient: str) -> ResponseOutput
         action = ActionFind(
             pattern=browser_call.get("pattern", ""),
             url=f"cursor:{browser_call.get('url', '')}",
-            type="find",
+            type="find_in_page",
         )
     else:
         raise ValueError(f"Unknown browser action: {recipient}")
diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py
index 27c60967e5b6..750d1db84c66 100644
--- a/vllm/entrypoints/openai/responses/protocol.py
+++ b/vllm/entrypoints/openai/responses/protocol.py
@@ -60,7 +60,11 @@
     ChatCompletionMessageParam,
     ChatTemplateContentFormatOption,
 )
-from vllm.entrypoints.generate.base.protocol import StopParam, validate_cache_salt
+from vllm.entrypoints.generate.base.protocol import (
+    PerRequestMetrics,
+    StopParam,
+    validate_cache_salt,
+)
 from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
@@ -80,6 +84,7 @@
 
 class InputTokensDetails(OpenAIBaseModel):
     cached_tokens: int
+    cache_write_tokens: int
     input_tokens_per_turn: list[int] = Field(default_factory=list)
     cached_tokens_per_turn: list[int] = Field(default_factory=list)
 
@@ -100,9 +105,7 @@ class ResponseUsage(OpenAIBaseModel):
 
 
 def serialize_message(msg):
-    """
-    Serializes a single message
-    """
+    """Serializes a single message."""
     if isinstance(msg, dict):
         return msg
     elif hasattr(msg, "to_dict"):
@@ -113,9 +116,7 @@ def serialize_message(msg):
 
 
 def serialize_messages(msgs):
-    """
-    Serializes multiple messages
-    """
+    """Serializes multiple messages."""
     return [serialize_message(msg) for msg in msgs] if msgs else None
 
 
@@ -687,6 +688,11 @@ class ResponsesResponse(OpenAIBaseModel):
     usage: ResponseUsage | None = None
     user: str | None = None
 
+    # vLLM-specific per-request metrics. Omitted unless enabled server-side.
+    metrics: PerRequestMetrics | None = Field(
+        default=None, exclude_if=lambda value: value is None
+    )
+
     presence_penalty: float | None = Field(
         default=None,
         ge=-2.0,
@@ -755,6 +761,7 @@ def from_request(
         output: list[ResponseOutputItem],
         status: ResponseStatus,
         usage: ResponseUsage | None = None,
+        metrics: PerRequestMetrics | None = None,
         input_messages: ResponseInputOutputMessage | None = None,
         output_messages: ResponseInputOutputMessage | None = None,
         kv_transfer_params: dict[str, Any] | None = None,
@@ -798,6 +805,7 @@ def from_request(
             truncation=request.truncation,
             user=request.user,
             usage=usage,
+            metrics=metrics,
             kv_transfer_params=kv_transfer_params,
             ec_transfer_params=ec_transfer_params,
         )
diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py
index 2c61205e2258..8108c5f77537 100644
--- a/vllm/entrypoints/openai/responses/serving.py
+++ b/vllm/entrypoints/openai/responses/serving.py
@@ -30,7 +30,10 @@
     DeltaMessage,
     RequestResponseMetadata,
 )
-from vllm.entrypoints.generate.base.serving import GenerateBaseServing
+from vllm.entrypoints.generate.base.serving import (
+    GenerateBaseServing,
+    build_per_request_timing_metrics,
+)
 from vllm.entrypoints.mcp.tool_server import ToolServer
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.responses.context import (
@@ -106,9 +109,11 @@ def __init__(
         reasoning_parser: str = "",
         enable_auto_tools: bool = False,
         tool_parser: str | None = None,
+        tool_strict_level: str = "auto",
         tool_server: ToolServer | None = None,
         enable_prompt_tokens_details: bool = False,
         enable_force_include_usage: bool = False,
+        enable_per_request_metrics: bool = False,
         enable_log_outputs: bool = False,
         default_chat_template_kwargs: dict[str, Any] | None = None,
     ) -> None:
@@ -131,11 +136,13 @@ def __init__(
             tool_parser_name=tool_parser,
             reasoning_parser_name=reasoning_parser,
             enable_auto_tools=enable_auto_tools,
+            tool_strict_level=tool_strict_level,
             model_name=self.model_config.model,
             is_harmony=self.model_config.hf_config.model_type == "gpt_oss",
         )
         self.enable_prompt_tokens_details = enable_prompt_tokens_details
         self.enable_force_include_usage = enable_force_include_usage
+        self.enable_per_request_metrics = enable_per_request_metrics
 
         self.default_sampling_params = self.model_config.get_diff_sampling_param()
         mc = self.model_config
@@ -645,6 +652,7 @@ async def _generate_with_builtin_tools(
             )
 
             async for res in generator:
+                context.request_metrics = res.metrics
                 context.append_output(res)
                 # NOTE(woosuk): The stop condition is handled by the engine.
                 yield context
@@ -654,6 +662,7 @@ async def _generate_with_builtin_tools(
                 break
 
             # Call the tool and update the context with the result.
+            context.request_metrics_cover_all_generation_turns = False
             tool_output = await context.call_tool()
             context.append_tool_output(tool_output)
 
@@ -837,6 +846,7 @@ async def responses_full_generator(
             output_tokens=num_generated_tokens,
             total_tokens=num_prompt_tokens + num_generated_tokens,
             input_tokens_details=InputTokensDetails(
+                cache_write_tokens=getattr(context, "num_cache_creation_tokens", 0),
                 cached_tokens=num_cached_tokens,
                 input_tokens_per_turn=[
                     turn.input_tokens for turn in context.all_turn_metrics
@@ -856,6 +866,14 @@ async def responses_full_generator(
                 ],
             ),
         )
+        per_request_metrics = None
+        if (
+            self.enable_per_request_metrics
+            and context.request_metrics_cover_all_generation_turns
+        ):
+            per_request_metrics = build_per_request_timing_metrics(
+                context.request_metrics, num_generated_tokens
+            )
         response = ResponsesResponse.from_request(
             request,
             sampling_params,
@@ -866,6 +884,7 @@ async def responses_full_generator(
             output=output,
             status=status,
             usage=usage,
+            metrics=per_request_metrics,
             kv_transfer_params=context.kv_transfer_params,
             ec_transfer_params=context.ec_transfer_params,
         )
diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py
index b45264f8c3b4..4651bb7205fd 100644
--- a/vllm/entrypoints/openai/responses/streaming_events.py
+++ b/vllm/entrypoints/openai/responses/streaming_events.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Streaming SSE event builders for the Responses API.
+"""Streaming SSE event builders for the Responses API.
 
 Pure functions that translate streaming state + delta data into
 OpenAI Response API SSE events. Used by the streaming event
@@ -125,8 +124,7 @@ def is_mcp_tool_by_namespace(
     recipient: str | None,
     allowed_function_tool_names: frozenset[str] | None = None,
 ) -> bool:
-    """
-    Determine if a tool call is an MCP tool based on recipient prefix.
+    """Determine if a tool call is an MCP tool based on recipient prefix.
 
     Inverse of :func:`is_function_recipient` — everything that is not
     a function call is an MCP tool.
@@ -667,7 +665,7 @@ def emit_browser_tool_events(
         )
     elif function_name == "find":
         action = response_function_web_search.ActionFind(
-            type="find",
+            type="find_in_page",
             pattern=parsed_args["pattern"],
             # TODO: translate to url
             url=f"cursor:{parsed_args.get('cursor', '')}",
@@ -1145,8 +1143,7 @@ def split_delta(delta: DeltaMessage) -> list[DeltaMessage]:
 
 
 class SimpleStreamingEventProcessor:
-    """
-    State-machine processor for the simple (non-Harmony) streaming path.
+    """State-machine processor for the simple (non-Harmony) streaming path.
 
     Core flow:
       1. Resolve the target state from the delta_message
@@ -1188,8 +1185,7 @@ def __init__(
     def resolve_target_state(
         self, delta_message: DeltaMessage
     ) -> tuple[_StateType, Any]:
-        """
-        Decide which state the next delta belongs to.
+        """Decide which state the next delta belongs to.
 
         Priority: TOOL_CALL > REASONING > CONTENT, fallback to NONE.
         For TOOL_CALL the first tool_call object is also returned so
@@ -1207,8 +1203,7 @@ def resolve_target_state(
         return _StateType.NONE, None
 
     def needs_transition(self, target_state: _StateType, tool_call: Any) -> bool:
-        """
-        Return True when we must close the current state and open a new one.
+        """Return True when we must close the current state and open a new one.
 
         Two cases trigger a transition:
           1. The target state differs from the current state
diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py
index 82ab75f756b1..3caff4b056b1 100644
--- a/vllm/entrypoints/openai/responses/utils.py
+++ b/vllm/entrypoints/openai/responses/utils.py
@@ -115,8 +115,7 @@ def build_response_output_items(
 def should_continue_final_message(
     request_input: str | list[ResponseInputOutputItem],
 ) -> bool:
-    """
-    Determine if the last input message is a partial assistant message
+    """Determine if the last input message is a partial assistant message
     that should be continued rather than starting a new generation.
 
     This enables partial message completion similar to Anthropic's Messages API,
@@ -132,6 +131,7 @@ def should_continue_final_message(
 
     Returns:
         True if the final message should be continued, False otherwise
+
     """
     if isinstance(request_input, str):
         # Simple string input is always a user message
@@ -226,8 +226,7 @@ def _construct_message_from_response_item(
     item: ResponseInputOutputItem,
     prev_msg: ChatCompletionMessageParam | None = None,
 ) -> ChatCompletionMessageParam | None:
-    """
-    Returns a new message or None. If `None`, `prev_msg` might be updated.
+    """Returns a new message or None. If `None`, `prev_msg` might be updated.
     If `prev_msg` is `None`, a new message is always returned.
     """
     prev_assistant_msg = (
@@ -360,9 +359,7 @@ def extract_function_tool_names(tools: list[Tool]) -> frozenset[str]:
 
 
 def extract_tool_types(tools: list[Tool]) -> set[str]:
-    """
-    Extracts the tool types from the given tools.
-    """
+    """Extracts the tool types from the given tools."""
     tool_types: set[str] = set()
     for tool in tools:
         if tool.type == "mcp":
@@ -379,8 +376,7 @@ def extract_tool_types(tools: list[Tool]) -> set[str]:
 def convert_tool_responses_to_completions_format(
     tool: dict,
 ) -> ChatCompletionToolsParam:
-    """
-    Convert a flat Responses tool schema:
+    """Convert a flat Responses tool schema:
         {"type": "function", "name": "...", "description": "...", "parameters": {...}}
     into a Chat Completions tool param for chat-template rendering.
     """
diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py
index e30195a7b1f9..c1366734f7c1 100644
--- a/vllm/entrypoints/pooling/base/protocol.py
+++ b/vllm/entrypoints/pooling/base/protocol.py
@@ -171,6 +171,8 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
 
 
 class EmbeddingTokenizeParamsMixin(PoolingTokenizeParamsMixin):
+    padding: Literal["max_length", "do_not_pad"] | None
+
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
         default_max_total_tokens = model_config.max_model_len
         max_total_tokens: int | None = default_max_total_tokens
@@ -184,7 +186,7 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
                 max_embed_len = pooler_config.max_embed_len or default_max_total_tokens
                 max_output_tokens = default_max_total_tokens - max_embed_len
 
-        return self._build_pooling_tok_params(
+        tok_params = self._build_pooling_tok_params(
             model_config,
             add_special_tokens=self.add_special_tokens,
             max_total_tokens=max_total_tokens,
@@ -192,6 +194,20 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
             max_output_tokens_param="max_model_len - max_embed_len",
         )
 
+        if (
+            self.padding == "max_length"
+            and pooler_config is not None
+            and pooler_config.enable_chunked_processing
+        ):
+            max_padding_len = pooler_config.max_embed_len
+            if max_padding_len is None:
+                max_padding_len = default_max_total_tokens
+            return tok_params.with_kwargs(
+                pad_prompt_tokens=min(default_max_total_tokens, max_padding_len)
+            )
+
+        return tok_params
+
 
 class CompletionRequestMixin(OpenAIBaseModel):
     # --8<-- [start:completion-params]
diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py
index 57728c88377c..3c5b75cee62c 100644
--- a/vllm/entrypoints/pooling/embed/io_processor.py
+++ b/vllm/entrypoints/pooling/embed/io_processor.py
@@ -57,6 +57,8 @@
 class _ChunkedPromptAggregator:
     weighted_sum: torch.Tensor | None = None
     total_weight: int = 0
+    num_chunks: int = 0
+    all_chunks_l2_normalized: bool = True
 
 
 class EmbedIOProcessor(PoolingIOProcessor):
@@ -204,7 +206,17 @@ def _post_process_chunked(self, ctx: PoolingServeContext) -> None:
 
             weight = len(result.prompt_token_ids)
             embedding_data = result.outputs.data
-            weighted_embedding = embedding_data.to(dtype=torch.float32) * weight
+            float_embedding = embedding_data.to(dtype=torch.float32)
+            if ctx.pooling_params.use_activation:
+                # Custom activations need not L2-normalize; allow BF16 rounding.
+                embedding_norm = torch.linalg.vector_norm(float_embedding)
+                aggregator.all_chunks_l2_normalized &= torch.allclose(
+                    embedding_norm,
+                    torch.ones_like(embedding_norm),
+                    rtol=1e-2,
+                    atol=1e-5,
+                )
+            weighted_embedding = float_embedding * weight
 
             if aggregator.weighted_sum is None:
                 # First chunk
@@ -214,6 +226,7 @@ def _post_process_chunked(self, ctx: PoolingServeContext) -> None:
                 aggregator.weighted_sum += weighted_embedding
 
             aggregator.total_weight += weight
+            aggregator.num_chunks += 1
 
         if ctx.original_engine_inputs is None:
             raise ValueError("Original engine inputs not available")
@@ -239,6 +252,14 @@ def _post_process_chunked(self, ctx: PoolingServeContext) -> None:
                 ):
                     # Compute final mean embedding
                     final_embedding = weighted_sum / total_weight
+                    if (
+                        ctx.pooling_params.use_activation
+                        and aggregator.num_chunks > 1
+                        and aggregator.all_chunks_l2_normalized
+                    ):
+                        final_embedding = torch.nn.functional.normalize(
+                            final_embedding, dim=-1
+                        )
 
                     # Create a PoolingRequestOutput
                     # for the aggregated result
diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py
index f9def808de3a..0a0304e7a4c5 100644
--- a/vllm/entrypoints/pooling/offline.py
+++ b/vllm/entrypoints/pooling/offline.py
@@ -37,7 +37,7 @@
 
 
 class PoolingOfflineMixin(OfflineInferenceMixin):
-    """Offline inference for pooling models"""
+    """Offline inference for pooling models."""
 
     runner_type: str
     chat_template: str | None
@@ -93,8 +93,8 @@ def encode(
         Returns:
             A list of `PoolingRequestOutput` objects containing the
             pooled hidden states in the same order as the input prompts.
-        """
 
+        """
         if isinstance(prompts, dict) and "data" in prompts and pooling_task != "plugin":
             raise ValueError(
                 "The 'data' field is only supported for the 'plugin' pooling task."
@@ -203,8 +203,7 @@ def embed(
         lora_request: list[LoRARequest] | LoRARequest | None = None,
         tokenization_kwargs: dict[str, Any] | None = None,
     ) -> list[EmbeddingRequestOutput]:
-        """
-        Generate an embedding vector for each prompt.
+        """Generate an embedding vector for each prompt.
 
         This class automatically batches the given prompts, considering
         the memory constraint. For the best performance, put all of your prompts
@@ -226,8 +225,8 @@ def embed(
         Returns:
             A list of `EmbeddingRequestOutput` objects containing the
             embedding vectors in the same order as the input prompts.
-        """
 
+        """
         items = self.encode(
             prompts,
             use_tqdm=use_tqdm,
@@ -248,8 +247,7 @@ def classify(
         lora_request: list[LoRARequest] | LoRARequest | None = None,
         tokenization_kwargs: dict[str, Any] | None = None,
     ) -> list[ClassificationRequestOutput]:
-        """
-        Generate class logits for each prompt.
+        """Generate class logits for each prompt.
 
         This class automatically batches the given prompts, considering
         the memory constraint. For the best performance, put all of your prompts
@@ -271,8 +269,8 @@ def classify(
         Returns:
             A list of `ClassificationRequestOutput` objects containing the
             embedding vectors in the same order as the input prompts.
-        """
 
+        """
         items = self.encode(
             prompts,
             use_tqdm=use_tqdm,
@@ -330,11 +328,12 @@ def score(
             chat_template: The chat template to use for the scoring. If None, we
                 use the model's default chat template.
             tokenization_kwargs: Overrides for `tokenizer.encode`.
+
         Returns:
             A list of `ScoringRequestOutput` objects containing the
             generated scores in the same order as the input prompts.
-        """
 
+        """
         if self.runner_type != "pooling":
             raise ValueError(
                 "LLM.score() is only supported for pooling models. "
diff --git a/vllm/entrypoints/pooling/scoring/api_router.py b/vllm/entrypoints/pooling/scoring/api_router.py
index 90be32ba248a..8e99e9a1feca 100644
--- a/vllm/entrypoints/pooling/scoring/api_router.py
+++ b/vllm/entrypoints/pooling/scoring/api_router.py
@@ -57,8 +57,6 @@ async def create_score(request: ScoreRequest, raw_request: Request):
         HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
     },
 )
-@with_cancellation
-@load_aware_call
 async def create_score_v1(request: ScoreRequest, raw_request: Request):
     logger.warning(
         "To indicate that Score API is not part of standard OpenAI API, we "
@@ -91,7 +89,6 @@ async def do_rerank(request: RerankRequest, raw_request: Request):
         HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
     },
 )
-@with_cancellation
 async def do_rerank_v1(request: RerankRequest, raw_request: Request):
     logger.warning_once(
         "To indicate that the rerank API is not part of the standard OpenAI"
@@ -110,6 +107,5 @@ async def do_rerank_v1(request: RerankRequest, raw_request: Request):
         HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
     },
 )
-@with_cancellation
 async def do_rerank_v2(request: RerankRequest, raw_request: Request):
     return await do_rerank(request, raw_request)
diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py
index 4feb40eadb6a..5e6104eedd22 100644
--- a/vllm/entrypoints/pooling/scoring/io_processor.py
+++ b/vllm/entrypoints/pooling/scoring/io_processor.py
@@ -7,6 +7,7 @@
 import torch.nn.functional as F
 
 from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt
+from vllm.logger import init_logger
 from vllm.renderers import TokenizeParams
 from vllm.renderers.hf import safe_apply_chat_template
 from vllm.renderers.inputs.preprocess import (
@@ -47,6 +48,8 @@
     validate_score_input,
 )
 
+logger = init_logger(__name__)
+
 ScoringServeContext: TypeAlias = PoolingServeContext[ScoringRequest]
 
 
@@ -425,6 +428,24 @@ def __init__(self, *args, **kwargs):
         self.model = model if self.supports_score_template else None
         self.use_sep_token = self.model_config.use_sep_token
 
+        if (
+            getattr(self.model_config.hf_config, "is_original_qwen3_reranker", False)
+            and self.chat_template is None
+        ):
+            suggested_template = (
+                "examples/pooling/score/template/qwen3_vl_reranker.jinja"
+                if self.model_config.is_multimodal_model
+                else "examples/pooling/score/template/qwen3_reranker.jinja"
+            )
+            logger.warning(
+                "Serving an original Qwen3 reranker (%s) without a "
+                "--chat-template. The model was trained with a specific prompt "
+                "template; running without it may produce inaccurate relevance "
+                "scores. Consider specifying `--chat-template %s`.",
+                self.model_config.model,
+                suggested_template,
+            )
+
     #######################################
     # online APIs
 
@@ -845,7 +866,14 @@ def get_request_factory_offline(
     ) -> tuple[RequestFactory, int]:
         assert isinstance(ctx, OfflineScoringInputsContext)
 
+        max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
+            pooling_params=ctx.pooling_params
+        )
         scoring_data = ctx.scoring_data
+        if max_tokens_per_query > 0 or max_tokens_per_doc > 0:
+            scoring_data = self._truncate_scoring_data(
+                scoring_data, max_tokens_per_query, max_tokens_per_doc
+            )
         prompt_extras = ctx.pooling_params.extra_kwargs
 
         queries = self.ensure_str(scoring_data.data_1)
diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py
index b32a7a5730af..674522b7ebde 100644
--- a/vllm/entrypoints/pooling/scoring/serving.py
+++ b/vllm/entrypoints/pooling/scoring/serving.py
@@ -80,6 +80,11 @@ def _build_response(
         request_id = ctx.request_id
         created_time = ctx.created_time
         model_name = ctx.model_name
+        assert ctx.engine_inputs is not None
+        num_prompt_tokens = sum(
+            self._extract_prompt_len(engine_input["prompts"])
+            for engine_input in ctx.engine_inputs
+        )
 
         if isinstance(ctx.request, ScoreRequest):
             return self._request_output_to_score_response(
@@ -87,6 +92,7 @@ def _build_response(
                 request_id,
                 created_time,
                 model_name,
+                num_prompt_tokens,
             )
         elif isinstance(ctx.request, RerankRequest):
             return self._request_output_to_rerank_response(
@@ -95,6 +101,7 @@ def _build_response(
                 model_name,
                 ctx.request.documents,
                 ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch),
+                num_prompt_tokens,
             )
         else:
             raise ValueError(f"Invalid {self.request_id_prefix} request type")
@@ -105,9 +112,9 @@ def _request_output_to_score_response(
         request_id: str,
         created_time: int,
         model_name: str,
+        num_prompt_tokens: int,
     ) -> JSONResponse:
         items: list[ScoreResponseData] = []
-        num_prompt_tokens = 0
 
         for idx, final_res in enumerate(final_res_batch):
             classify_res = ScoringRequestOutput.from_base(final_res)
@@ -116,10 +123,7 @@ def _request_output_to_score_response(
                 index=idx,
                 score=classify_res.outputs.score,
             )
-            prompt_token_ids = final_res.prompt_token_ids
-
             items.append(item)
-            num_prompt_tokens += len(prompt_token_ids)
 
         usage = UsageInfo(
             prompt_tokens=num_prompt_tokens,
@@ -143,12 +147,12 @@ def _request_output_to_rerank_response(
         model_name: str,
         documents: ScoreInput | list[ScoreInput],
         top_n: int,
+        num_prompt_tokens: int,
     ) -> JSONResponse:
         if not isinstance(documents, list):
             documents = [documents]
 
         results: list[RerankResult] = []
-        num_prompt_tokens = 0
         for idx, final_res in enumerate(final_res_batch):
             classify_res = ScoringRequestOutput.from_base(final_res)
 
@@ -166,8 +170,6 @@ def _request_output_to_rerank_response(
                 relevance_score=classify_res.outputs.score,
             )
             results.append(result)
-            prompt_token_ids = final_res.prompt_token_ids
-            num_prompt_tokens += len(prompt_token_ids)
 
         # sort by relevance, then return the top n if set
         results.sort(key=lambda x: x.relevance_score, reverse=True)
diff --git a/vllm/entrypoints/pooling/scoring/typing.py b/vllm/entrypoints/pooling/scoring/typing.py
index fd6191b6af8f..a67843e57770 100644
--- a/vllm/entrypoints/pooling/scoring/typing.py
+++ b/vllm/entrypoints/pooling/scoring/typing.py
@@ -21,8 +21,7 @@
 
 
 class ScoreMultiModalParam(TypedDict, total=False):
-    """
-    A specialized parameter type for scoring multimodal content
+    """A specialized parameter type for scoring multimodal content.
 
     The reasons why don't reuse `CustomChatCompletionMessageParam` directly:
     1. Score tasks don't need the 'role' field (user/assistant/system) that's required in chat completions
diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py
index 148b429d3c5a..cff8ff61d914 100644
--- a/vllm/entrypoints/pooling/scoring/utils.py
+++ b/vllm/entrypoints/pooling/scoring/utils.py
@@ -58,8 +58,7 @@ def truncate_text_to_tokens(
 
 
 def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor:
-    """
-    Compute ColBERT MaxSim score.
+    """Compute ColBERT MaxSim score.
 
     Args:
         q_emb: Query token embeddings [query_len, dim]
@@ -67,6 +66,7 @@ def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tens
 
     Returns:
         MaxSim score (sum over query tokens of max similarity to any doc token)
+
     """
     # compute in float32 for numerical stability
     # [query_len, doc_len]
@@ -249,8 +249,7 @@ def parse_score_data(
 
 
 def compress_token_type_ids(token_type_ids: list[int]) -> int:
-    """
-    Return position of the first 1 or the length of the list
+    """Return position of the first 1 or the length of the list
     if not found.
     """
     first_one = len(token_type_ids)
diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py
index 81e919d840b7..5da736640373 100644
--- a/vllm/entrypoints/scale_out/derender/serving.py
+++ b/vllm/entrypoints/scale_out/derender/serving.py
@@ -138,6 +138,13 @@ async def derender_chat_response(
         if bounds_error is not None:
             return bounds_error
 
+        if self.online_derenderer.parser is not None and request.chat_request is None:
+            return self.create_error_response(
+                "chat_request is required when a tool or reasoning parser is "
+                "configured because plain detokenization would leak raw parser "
+                "markup into content."
+            )
+
         try:
             choices = await self.online_derenderer.derender_chat(
                 request.generate_response, request.chat_request
@@ -172,6 +179,7 @@ async def derender_chat_response(
             usage=usage,
             prompt_logprobs=gen.prompt_logprobs,
             kv_transfer_params=gen.kv_transfer_params,
+            metrics=gen.metrics,
         )
 
     async def derender_completion_response(
@@ -238,6 +246,9 @@ async def derender_completion_response(
             choices=choices,
             usage=usage,
             kv_transfer_params=kv_params,
+            # Metrics describe one prompt. Parallel samples are already
+            # suppressed by /generate; multi-prompt responses cannot be merged.
+            metrics=first.metrics if len(request.generate_responses) == 1 else None,
         )
 
     async def derender_chat_stream_response(
@@ -248,14 +259,67 @@ async def derender_chat_stream_response(
 
         Processes one ``GenerateStreamResponse`` chunk and returns the
         derendered chunk together with the updated client carried state.
-
-        ``parser is None`` or no ``chat_request`` until reasoning/tool call
-        functionality added in future PR.
         """
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
             return error_check_ret
 
+        if self.online_derenderer.parser is not None and request.chat_request is None:
+            return self.create_error_response(
+                "chat_request is required when a tool or reasoning parser is "
+                "configured because plain detokenization would leak raw parser "
+                "markup into content."
+            )
+
+        if (
+            self.online_derenderer.parser is not None
+            and request.prompt_token_ids is None
+        ):
+            return self.create_error_response(
+                "prompt_token_ids is required when a tool or reasoning "
+                "parser is configured. Without it parse_delta cannot tell "
+                "whether the prompt left reasoning open (e.g. a chat "
+                "template that pre-opens <think>) and would misclassify "
+                "reasoning content as plain content."
+            )
+
+        # Each streamed chunk contains at most one choice per SSE event.
+        # A single DerenderStreamState is threaded
+        # through every choice in the chunk, so >1 would corrupt detok/parser
+        # state across choices. See the matching check in
+        # OnlineDerenderer.derender_chat_stream which this fails ahead of to
+        # avoid touching the tokenizer at all on a malformed chunk.
+        num_choices = len(request.generate_chunk.choices)
+        if num_choices > 1:
+            return self.create_error_response(
+                f"derender_chat_stream expects at most one choice per chunk "
+                f"(got {num_choices})."
+            )
+
+        # prompt_token_ids is folded in here too. It is caller supplied and
+        # otherwise unbounded. A parser configured deployment rescans it
+        # in full on every chunk (is_reasoning_end /
+        # adjust_initial_state_from_prompt, see abstract_parser.parse_delta)
+        # since parser state cannot be carried across calls.
+        prompt_len = (
+            len(request.prompt_token_ids) if request.prompt_token_ids is not None else 0
+        )
+        prior_len = (
+            len(request.stream_state.output_token_ids)
+            if request.stream_state is not None
+            else 0
+        )
+        delta_len = sum(
+            len(c.token_ids) for c in request.generate_chunk.choices if c.token_ids
+        )
+        max_model_len = self.model_config.max_model_len
+        if prompt_len + prior_len + delta_len > max_model_len:
+            return self.create_error_response(
+                f"prompt_token_ids length ({prompt_len}) plus "
+                f"output_token_ids length ({prior_len}) plus delta "
+                f"({delta_len}) exceeds max_model_len ({max_model_len})."
+            )
+
         model_name = request.model or self.models.model_name()
         try:
             chunk, updated_state = await self.online_derenderer.derender_chat_stream(
@@ -264,15 +328,21 @@ async def derender_chat_stream_response(
                 state=request.stream_state,
                 chat_request=request.chat_request,
                 prompt_tokens=request.prompt_tokens,
+                prompt_token_ids=request.prompt_token_ids,
             )
-        except NotImplementedError as exc:
-            return self.create_error_response(exc)
-        except ValueError as exc:
-            return self.create_error_response(str(exc))
-        except (KeyError, IndexError) as exc:
-            return self.create_error_response(
-                f"invalid stream_state: detokenization failed ({exc!r})"
+        except Exception as exc:
+            # The two ValueErrors derender_chat_stream can raise directly
+            # (missing chat_request, >1 choice per chunk) are already
+            # pre-checked above, so anything reaching here comes from the
+            # parser replaying entirely client controlled state. Hermes
+            # swallows its own streaming exceptions but finalize_generation
+            # and engine based parsers don't, so an unexpected parser
+            # failure on malformed (but well typed) state must surface as
+            # a 400 here rather than an unhandled 500.
+            logger.debug(
+                "derender stream replay failed: %s: %s", type(exc).__name__, exc
             )
+            return self.create_error_response("invalid stream_state: derender failed")
 
         logger.debug(
             "derender_chat_stream request_id=%s model=%s delta_tokens=%d",
@@ -313,7 +383,7 @@ async def derender_completion_stream_response(
             return self.create_error_response(str(exc))
         except (KeyError, IndexError) as exc:
             return self.create_error_response(
-                f"invalid stream_state: detokenization failed ({exc!r})"
+                f"invalid stream_state: detokenization failed ({exc})"
             )
 
         logger.debug(
diff --git a/vllm/entrypoints/scale_out/factories.py b/vllm/entrypoints/scale_out/factories.py
index 30c329c51f02..bd48c709ddc0 100644
--- a/vllm/entrypoints/scale_out/factories.py
+++ b/vllm/entrypoints/scale_out/factories.py
@@ -5,7 +5,6 @@
 
 from fastapi import FastAPI
 
-from vllm import envs
 from vllm.engine.protocol import EngineClient
 from vllm.logger import init_logger
 from vllm.tasks import SupportedTask
@@ -67,22 +66,15 @@ def register_scale_out_api_routers(
     app: FastAPI,
     supported_tasks: tuple["SupportedTask", ...],
 ):
-    dedicated_mode = None
-    if "render" in supported_tasks:
-        dedicated_mode = "`vllm launch render`"
-    elif getattr(app.state.args, "tokens_only", False):
-        dedicated_mode = "`--tokens-only`"
-
-    enabled = envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS
-    if dedicated_mode is not None and enabled is False:
-        raise ValueError(
-            "`VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0` conflicts with "
-            f"{dedicated_mode}; unset it or set it to 1"
-        )
-    if dedicated_mode is None and not enabled:
+    args = getattr(app.state, "args", None)
+    # `vllm launch render` and `vllm serve --tokens-only` are dedicated to
+    # serving these endpoints, so they always register them regardless of
+    # `--enable-scale-out`.
+    dedicated_mode = "render" in supported_tasks or getattr(args, "tokens_only", False)
+    enabled = dedicated_mode or getattr(args, "enable_scale_out", False)
+    if not enabled:
         logger.info(
-            "Scale-out endpoints are disabled. Set "
-            "VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 to enable them."
+            "Scale-out endpoints are disabled. Set --enable-scale-out to enable them."
         )
         return
 
diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py
index b79138664f56..7973ac3c3d88 100644
--- a/vllm/entrypoints/scale_out/render/serving.py
+++ b/vllm/entrypoints/scale_out/render/serving.py
@@ -126,33 +126,11 @@ async def render_chat_request(
         )
         params = request.to_sampling_params(max_tokens, self.default_sampling_params)
 
-        assistant_tokens_mask: list[int] | None = engine_input.get(  # type: ignore[assignment]
-            "assistant_tokens_mask"
-        )
-        if assistant_tokens_mask is not None and len(assistant_tokens_mask) != len(
-            token_ids
-        ):
-            logger.warning(
-                "assistant_tokens_mask length (%d) != token_ids length (%d); "
-                "this can happen with multimodal inputs where "
-                "placeholder expansion changes the token count. "
-                "The mask may be positionally misaligned.",
-                len(assistant_tokens_mask),
-                len(token_ids),
-            )
-            if len(assistant_tokens_mask) < len(token_ids):
-                assistant_tokens_mask.extend(
-                    [0] * (len(token_ids) - len(assistant_tokens_mask))
-                )
-            else:
-                assistant_tokens_mask = assistant_tokens_mask[: len(token_ids)]
-
         request_id = f"chatcmpl-{random_uuid()}"
 
         return GenerateRequest(
             request_id=request_id,
             token_ids=token_ids,
-            assistant_tokens_mask=assistant_tokens_mask,
             features=self._extract_mm_features(engine_input),
             sampling_params=params,
             model=request.model,
diff --git a/vllm/entrypoints/scale_out/token_in_token_out/api_router.py b/vllm/entrypoints/scale_out/token_in_token_out/api_router.py
index 2fc707cec475..28677b63213a 100644
--- a/vllm/entrypoints/scale_out/token_in_token_out/api_router.py
+++ b/vllm/entrypoints/scale_out/token_in_token_out/api_router.py
@@ -78,8 +78,7 @@ def attach_router(app: FastAPI):
 
         @router.post("/abort_requests")
         async def abort_requests(raw_request: Request):
-            """
-            Abort one or more requests. To be used in a
+            """Abort one or more requests. To be used in a
             Disaggregated Everything setup.
             """
             try:
diff --git a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py
index d7f89e895e5f..62d006783a2e 100644
--- a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py
+++ b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py
@@ -117,6 +117,24 @@ def mm_kwargs_from_features(
     return mm_kwargs
 
 
+def placeholder_ranges_from_engine_input(
+    engine_input: EngineInput,
+) -> dict[str, list[PlaceholderRangeInfo]] | None:
+    """Return per-modality placeholder ranges, or ``None`` for text-only prompts."""
+    if engine_input.get("type") != "multimodal":
+        return None
+
+    raw_placeholders: MultiModalPlaceholders = cast(MultiModalInput, engine_input)[
+        "mm_placeholders"
+    ]
+    return {
+        modality: [
+            PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges
+        ]
+        for modality, ranges in raw_placeholders.items()
+    }
+
+
 def extract_mm_features(
     engine_input: EngineInput,
     *,
@@ -134,14 +152,8 @@ def extract_mm_features(
 
     mm_engine_input = cast(MultiModalInput, engine_input)
     mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
-    raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
-
-    mm_placeholders = {
-        modality: [
-            PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges
-        ]
-        for modality, ranges in raw_placeholders.items()
-    }
+    mm_placeholders = placeholder_ranges_from_engine_input(engine_input)
+    assert mm_placeholders is not None
 
     kwargs_data: dict[str, list[str | None]] | None = None
     mm_metadata: dict[str, list[str | None]] | None = None
diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py
index 65643b9e9927..fb2b93872388 100644
--- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import Any, Literal, TypeAlias
+from typing import Annotated, Any, Literal, TypeAlias
 
 from pydantic import (
     BaseModel,
@@ -11,7 +11,11 @@
 )
 
 from vllm.config import ModelConfig
-from vllm.entrypoints.generate.base.protocol import StreamOptions, validate_cache_salt
+from vllm.entrypoints.generate.base.protocol import (
+    PerRequestMetrics,
+    StreamOptions,
+    validate_cache_salt,
+)
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionLogProbs,
     ChatCompletionRequest,
@@ -150,14 +154,6 @@ class GenerateRequest(BaseModel):
     token_ids: list[int] = Field(min_length=1)
     """The token ids to generate text from."""
 
-    assistant_tokens_mask: list[int] | None = None
-    """Per-token mask (1 = assistant-generated, 0 = not).
-
-    Only populated when the render request sets ``return_assistant_tokens_mask=True``
-    and the chat template supports ``{% generation %}``.
-    ``None`` when the mask was not requested or could not be computed.
-    """
-
     @field_validator("token_ids")
     @classmethod
     def validate_token_ids(cls, v: list[int]) -> list[int]:
@@ -192,20 +188,34 @@ def _require_ec_for_metadata_only(self) -> "GenerateRequest":
             return self
         if not _has_serialized_mm_items(features.mm_metadata):
             return self
-        if _has_serialized_mm_items(features.kwargs_data):
-            return self
         if self.ec_transfer_params:
             return self
-        raise ValueError(
-            "features.mm_metadata without features.kwargs_data requires "
-            "ec_transfer_params so embeddings can be loaded by the EC connector"
-        )
+        kwargs_data = features.kwargs_data or {}
+        for modality, metadata_items in (features.mm_metadata or {}).items():
+            kwargs_items = kwargs_data.get(modality, [None] * len(metadata_items))
+            if any(
+                metadata is not None and kwargs is None
+                for metadata, kwargs in zip(metadata_items, kwargs_items, strict=True)
+            ):
+                raise ValueError(
+                    "metadata-only multimodal items require ec_transfer_params"
+                )
+        return self
 
     sampling_params: SamplingParams
     """The sampling parameters for the model."""
 
     model: str | None = None
 
+    return_token_ids: bool | None = Field(
+        default=None,
+        description=(
+            "If true, return the final prompt token IDs after multimodal "
+            "placeholder expansion, together with multimodal placeholder ranges. "
+            "In streaming mode, this metadata is included only in the first chunk."
+        ),
+    )
+
     stream: bool | None = False
     stream_options: StreamOptions | None = None
     cache_salt: str | None = Field(
@@ -250,6 +260,9 @@ def _require_ec_for_metadata_only(self) -> "GenerateRequest":
     # ``SamplingParams`` instance (e.g. from internal callers that have
     # already resolved values), in which case all fields are considered set.
     _sampling_params_provided_keys: set[str] | None = PrivateAttr(default=None)
+    _response_mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = (
+        PrivateAttr(default=None)
+    )
 
     @model_validator(mode="wrap")
     @classmethod
@@ -349,6 +362,9 @@ class GenerateStreamResponse(BaseModel):
     )
     choices: list[GenerateResponseStreamChoice]
     usage: UsageInfo | None = Field(default=None)
+    prompt_token_ids: list[int] | None = None
+    mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
+    metrics: PerRequestMetrics | None = None
 
 
 class GenerateResponse(BaseModel):
@@ -365,6 +381,9 @@ class GenerateResponse(BaseModel):
     choices: list[GenerateResponseChoice]
     usage: UsageInfo | None = Field(default=None)
     prompt_logprobs: list[dict[int, Logprob] | None] | None = None
+    prompt_token_ids: list[int] | None = None
+    mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None
+    metrics: PerRequestMetrics | None = None
 
     kv_transfer_params: dict[str, Any] | None = Field(
         default=None,
@@ -462,16 +481,32 @@ class DerenderStreamState(BaseModel):
     streaming derender endpoint. All fields are plain JSON serializable data.
     No opaque tokenizer or parser internals are stored here.
 
+    Two separate sets of fields support two different streaming modes:
+
+    - For plain streaming (no parser configured including the completions path):
+      `prev_tokens`, `prefix_offset` and `read_offset` maintain a bounded incremental
+      decoding window. This requires O(window) transport and O(delta) computation
+      per chunk.
+    - For parser enabled chat streaming: `output_token_ids`, `output_chunk_lens`,
+      `tools_streamed` and `last_tool_call_ids` are used to replay `parse_delta()`
+      from scratch on every chunk because parser state cannot be serialized. This
+      incurs O(n) transport per chunk (O(n²) per generation) and O(n²)
+      `parse_delta()` calls per generation. Since many parsers re-scan the
+      entire accumulated text on each invocation, `parse_delta()` itself is
+      O(n) yielding a true worst case compute cost of O(n³) per generation.
+      No caching is performed. Work is bounded by `max_model_len`.
+      See `OnlineDerenderer._derender_chat_stream_parsed`.
+
     The detokenization strategy carries the incremental decode offsets
     directly rather than re-sending the whole token history each chunk.
-    ``detokenize_incrementally`` only ever reads the trailing token window
-    ``prev_tokens[prefix_offset:]``, so we carry just that tail plus the two
+    `detokenize_incrementally` only ever reads the trailing token window
+    `prev_tokens[prefix_offset:]`, so we carry just that tail plus the two
     offsets. Each chunk resumes exactly where the last one stopped, including
-    any partially processed multi-byte character (tracked by ``read_offset``),
+    any partially processed multi-byte character (tracked by `read_offset`),
     then trims and rebases the window so it never grows with generation length.
 
     Performance:
-    - Compute per chunk is O(delta). One ``detokenize_incrementally`` call per
+    - Compute per chunk is O(delta). One `detokenize_incrementally` call per
       new token, independent of how many tokens preceded it.
     - Transport per chunk is O(window). The carried tail is bounded by the
       incremental detokenization offset, so cumulative bytes over the wire are
@@ -479,18 +514,18 @@ class DerenderStreamState(BaseModel):
     """
 
     prev_tokens: list[str] = Field(default_factory=list)
-    """Trailing decode window. Token strings from ``prefix_offset`` onward.
+    """Trailing decode window. Token strings from `prefix_offset` onward.
 
     Bounded, trimmed and rebased each chunk to the tail
-    ``detokenize_incrementally`` still reads, so it does not grow with the
+    `detokenize_incrementally` still reads, so it does not grow with the
     number of chunks.
     """
 
     prefix_offset: int = Field(default=0, ge=0)
-    """Prefix offset into ``prev_tokens`` for incremental detokenization."""
+    """Prefix offset into `prev_tokens` for incremental detokenization."""
 
     read_offset: int = Field(default=0, ge=0)
-    """Read offset into ``prev_tokens`` for incremental detokenization."""
+    """Read offset into `prev_tokens` for incremental detokenization."""
 
     @field_validator("prev_tokens")
     @classmethod
@@ -504,23 +539,53 @@ def _bound_prev_tokens(cls, v: list[str]) -> list[str]:
         return v
 
     role_sent: bool = False
-    """True once the initial ``role: "assistant"`` delta has been emitted.
+    """True once the initial `role: "assistant"` delta has been emitted.
 
     Prevents re-emitting the role on subsequent chunks even when the detok
     window is transiently empty (e.g. usage only final chunk).
     """
 
-    # TODO: Properties used in follow on PR for tool call parsing
-    last_content: str | None = None
-    """Last emitted cumulative assistant content text."""
+    output_token_ids: list[int] = Field(default_factory=list)
+    """All output tokens seen so far. Parser path only.
+
+    Replay buffer: each chunk rebuilds a fresh parser and replays every
+    token in here through `parse_delta` (discarding the result) before
+    processing the current chunk's tokens since parser internal state
+    cannot be serialized into this stateless model. Unavoidably O(n)
+    bounded by ``max_model_len`` (enforced server side, not by a field
+    validator here since the bound is model dependent).
+    """
+
+    output_chunk_lens: list[Annotated[int, Field(gt=0)]] = Field(default_factory=list)
+    """Token count of each chunk in `output_token_ids`. Parser path only.
 
-    last_reasoning: str | None = None
-    """Last emitted cumulative reasoning text."""
+    Replay uses these to reproduce the original `parse_delta` call
+    boundaries. Must sum to `len(output_token_ids)`.
+    """
+
+    @model_validator(mode="after")
+    def _validate_output_chunk_lens(self) -> "DerenderStreamState":
+        total = sum(self.output_chunk_lens)
+        if total != len(self.output_token_ids):
+            raise ValueError(
+                f"output_chunk_lens must sum to len(output_token_ids) "
+                f"(got sum={total}, len(output_token_ids)="
+                f"{len(self.output_token_ids)})"
+            )
+        return self
+
+    tools_streamed: bool = False
+    """True once a tool call delta has been emitted. Parser path only.
+
+    Drives the `finish_reason` -> `"tool_calls"` rewrite on the final
+    chunk mirroring the generate streaming path.
+    """
 
     last_tool_call_ids: list[str] = Field(default_factory=list)
-    """Stable tool-call IDs, assigned once when each call first appears.
+    """Stable tool call IDs, assigned once when each call first appears.
 
-    Prevents ID regeneration across re-parsing.
+    Indexed by tool call index. Parser path only. Prevents ID regeneration
+    when replay reprocesses a tool call that already has a pinned ID.
     """
 
 
@@ -537,6 +602,7 @@ class DerenderChatStreamRequest(BaseModel):
     the client carried ``stream_state``.
     """
 
+    # --8<-- [start:derender-chat-stream-request]
     stream: Literal[True]
 
     model: str | None = None
@@ -549,8 +615,22 @@ class DerenderChatStreamRequest(BaseModel):
     prompt_tokens: int | None = None
     """Prompt token count for usage. Forwarded from the render step."""
 
+    prompt_token_ids: list[int] | None = None
+    """Prompt token IDs. Required by the parser path's `parse_delta` to
+    settle its initial reasoning state (e.g. chat templates that pre-open
+    ``<think>``). `prompt_tokens` is a usage count and cannot serve this
+    purpose. Sourced from `GenerateRequest.token_ids` at the render step.
+
+    Rejected with a 400 (by `ServingDerender`) when a tool or reasoning
+    parser is configured and this is omitted. Without it, `parse_delta`
+    cannot tell whether the prompt left reasoning open and would silently
+    misclassify reasoning content as plain content. Unused on the plain
+    detokenization path.
+    """
+
     chat_request: ChatCompletionRequest | None = None
     """The original (post adjust_request) ChatCompletionRequest from /render."""
+    # --8<-- [end:derender-chat-stream-request]
 
 
 class DerenderCompletionStreamRequest(BaseModel):
@@ -561,6 +641,7 @@ class DerenderCompletionStreamRequest(BaseModel):
     returns the derendered chunk plus updated state.
     """
 
+    # --8<-- [start:derender-completion-stream-request]
     stream: Literal[True]
 
     model: str | None = None
@@ -575,6 +656,7 @@ class DerenderCompletionStreamRequest(BaseModel):
 
     completion_request: CompletionRequest | None = None
     """The original (post adjust_request) CompletionRequest from /render."""
+    # --8<-- [end:derender-completion-stream-request]
 
 
 class DerenderChatStreamResponse(BaseModel):
diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py
index d3725e3cf71e..b480b802ec80 100644
--- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py
+++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py
@@ -12,9 +12,13 @@
 
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker
-from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
+from vllm.entrypoints.generate.base.protocol import (
+    PerRequestMetrics,
+    RequestResponseMetadata,
+)
 from vllm.entrypoints.generate.base.serving import (
     GenerateBaseServing,
+    build_spec_decoding_metrics,
     clamp_prompt_logprobs,
 )
 from vllm.entrypoints.openai.chat_completion.protocol import (
@@ -44,7 +48,10 @@
 from vllm.utils.collection_utils import as_list
 from vllm.utils.serial_utils import numpy2base64
 
-from .mm_features import mm_kwargs_from_features
+from .mm_features import (
+    mm_kwargs_from_features,
+    placeholder_ranges_from_engine_input,
+)
 from .protocol import (
     GenerateRequest,
     GenerateResponse,
@@ -131,6 +138,16 @@ async def serve_tokens(
                 f"sampling_params.n must be at most the server's max_num_seqs "
                 f"({max_num_seqs}), got {sampling_params.n}."
             )
+        if self.force_no_detokenize and sampling_params.stop:
+            # SamplingParams rejects stop with detokenize=False at request
+            # validation, but this server forces detokenize=False afterwards,
+            # so the combination must be rejected here or stop strings are
+            # silently never applied.
+            return self.create_error_response(
+                "stop strings are not supported on a --tokens-only server "
+                "because detokenization is disabled. Check stop strings on "
+                "the coordinator, or use stop_token_ids."
+            )
         try:
             msgspec.msgpack.encode(
                 (
@@ -195,6 +212,14 @@ async def serve_tokens(
                 skip_mm_cache=True,
             )
 
+        # Offsets are relative to the decoder prompt, so they are not
+        # meaningful for encoder-decoder models.
+        request._response_mm_placeholders = (
+            placeholder_ranges_from_engine_input(engine_input)
+            if request.return_token_ids and not self.model_config.is_encoder_decoder
+            else None
+        )
+
         # Schedule the request and get the result generator.
         result_generator: AsyncGenerator[RequestOutput, None] | None = None
 
@@ -352,6 +377,11 @@ async def serve_tokens_full_generator(
 
         request_metadata.final_usage_info = usage
 
+        per_request_metrics = None
+        if request.sampling_params.n == 1:
+            spec_stats = build_spec_decoding_metrics(final_res)
+            if spec_stats is not None:
+                per_request_metrics = PerRequestMetrics(speculative_decoding=spec_stats)
         response = GenerateResponse(
             request_id=request_id,
             created=created_time,
@@ -359,6 +389,11 @@ async def serve_tokens_full_generator(
             choices=choices,
             usage=usage,
             prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
+            prompt_token_ids=(
+                final_res.prompt_token_ids if request.return_token_ids else None
+            ),
+            mm_placeholders=request._response_mm_placeholders,
+            metrics=per_request_metrics,
             kv_transfer_params=final_res.kv_transfer_params,
             ec_transfer_params=final_res.ec_transfer_params,
         )
@@ -395,8 +430,10 @@ async def serve_tokens_stream_generator(
         num_prompt_tokens = 0
         num_generated_tokens: list[int] = []
         first_iteration = True
+        prompt_token_ids: list[int] | None = None
         num_cached_tokens = None
         sampling_params: SamplingParams = request.sampling_params
+        last_res: RequestOutput | None = None
 
         include_usage, include_continuous_usage = should_include_usage(
             request.stream_options, False
@@ -404,9 +441,12 @@ async def serve_tokens_stream_generator(
 
         try:
             async for res in result_generator:
+                last_res = res
                 if first_iteration:
                     if res.prompt_token_ids is not None:
                         num_prompt_tokens = len(res.prompt_token_ids)
+                        if request.return_token_ids:
+                            prompt_token_ids = res.prompt_token_ids
                     if res.encoder_prompt_token_ids is not None:
                         num_prompt_tokens += len(res.encoder_prompt_token_ids)
                     num_cached_tokens = res.num_cached_tokens
@@ -421,7 +461,11 @@ async def serve_tokens_stream_generator(
                     finish_reason = output.finish_reason
                     self._raise_if_error(finish_reason, request_id)
 
-                    if not delta_token_ids:
+                    # Still emit a terminal empty chunk while prompt metadata
+                    # is pending, so zero-token completions deliver it.
+                    if not delta_token_ids and (
+                        finish_reason is None or prompt_token_ids is None
+                    ):
                         continue
 
                     if sampling_params.logprobs is not None:
@@ -453,6 +497,10 @@ async def serve_tokens_stream_generator(
                             )
                         ],
                     )
+                    if prompt_token_ids is not None:
+                        chunk.prompt_token_ids = prompt_token_ids
+                        chunk.mm_placeholders = request._response_mm_placeholders
+                        prompt_token_ids = None
                     if include_continuous_usage:
                         chunk.usage = UsageInfo(
                             prompt_tokens=num_prompt_tokens,
@@ -460,7 +508,13 @@ async def serve_tokens_stream_generator(
                             total_tokens=(num_prompt_tokens + num_generated_tokens[i]),
                         )
 
-                    yield f"data: {chunk.model_dump_json()}\n\n"
+                    # Omit fields that are absent from token-bearing chunks.
+                    exclude = {
+                        name
+                        for name in ("prompt_token_ids", "mm_placeholders", "metrics")
+                        if getattr(chunk, name) is None
+                    }
+                    yield f"data: {chunk.model_dump_json(exclude=exclude)}\n\n"
 
             total_completion_tokens = sum(num_generated_tokens)
             final_usage_info = UsageInfo(
@@ -475,10 +529,18 @@ async def serve_tokens_stream_generator(
                 )
 
             if include_usage:
+                per_request_metrics = None
+                if sampling_params.n == 1:
+                    spec_stats = build_spec_decoding_metrics(last_res)
+                    if spec_stats is not None:
+                        per_request_metrics = PerRequestMetrics(
+                            speculative_decoding=spec_stats
+                        )
                 final_chunk = GenerateStreamResponse(
                     request_id=request_id,
                     choices=[],
                     usage=final_usage_info,
+                    metrics=per_request_metrics,
                 )
                 yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
 
@@ -528,7 +590,10 @@ def _create_tokens_logprobs(
                                 step_top_logprobs.items()
                             )
                             if num_output_top_logprobs is not None
-                            and i < max(num_output_top_logprobs, 1)
+                            and (
+                                num_output_top_logprobs == -1
+                                or i < max(num_output_top_logprobs, 1)
+                            )
                         ],
                     )
                 )
diff --git a/vllm/entrypoints/serve/dev/cache/api_router.py b/vllm/entrypoints/serve/dev/cache/api_router.py
index 63737393f29c..81ec91a8666a 100644
--- a/vllm/entrypoints/serve/dev/cache/api_router.py
+++ b/vllm/entrypoints/serve/dev/cache/api_router.py
@@ -23,8 +23,7 @@ async def reset_prefix_cache(
     reset_running_requests: bool = Query(default=False),
     reset_external: bool = Query(default=False),
 ):
-    """
-    Reset the local prefix cache.
+    """Reset the local prefix cache.
 
     Optionally, if the query parameter `reset_external=true`
     also resets the external (connector-managed) prefix cache.
@@ -35,6 +34,7 @@ async def reset_prefix_cache(
 
     Example:
        POST /reset_prefix_cache?reset_external=true
+
     """
     logger.info("Resetting prefix cache...")
 
@@ -46,8 +46,7 @@ async def reset_prefix_cache(
 
 @router.post("/reset_mm_cache")
 async def reset_mm_cache(raw_request: Request):
-    """
-    Reset the multi-modal cache. Note that we currently do not check if the
+    """Reset the multi-modal cache. Note that we currently do not check if the
     multi-modal cache is successfully reset in the API server.
     """
     logger.info("Resetting multi-modal cache...")
@@ -57,8 +56,7 @@ async def reset_mm_cache(raw_request: Request):
 
 @router.post("/reset_encoder_cache")
 async def reset_encoder_cache(raw_request: Request):
-    """
-    Reset the encoder cache. Note that we currently do not check if the
+    """Reset the encoder cache. Note that we currently do not check if the
     encoder cache is successfully reset in the API server.
     """
     logger.info("Resetting encoder cache...")
diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/serve/dev/rlhf/api_router.py
index 392fcf567476..603fdf785c3d 100644
--- a/vllm/entrypoints/serve/dev/rlhf/api_router.py
+++ b/vllm/entrypoints/serve/dev/rlhf/api_router.py
@@ -36,6 +36,8 @@ async def pause_generation(
     """Pause generation requests to allow weight updates.
 
     Args:
+        raw_request: The incoming FastAPI request, used to reach the engine
+            client on the app state.
         mode: How to handle in-flight requests:
             - ``"abort"``: Abort all in-flight requests immediately (default).
             - ``"wait"``: Wait for in-flight requests to complete.
@@ -43,8 +45,8 @@ async def pause_generation(
         wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead.
         clear_cache: DEPRECATED. Whether to clear KV/prefix caches after
             draining. Ignored when mode="keep".
-    """
 
+    """
     engine = engine_client(raw_request)
 
     try:
@@ -74,7 +76,6 @@ async def pause_generation(
 @router.post("/resume")
 async def resume_generation(raw_request: Request) -> JSONResponse:
     """Resume generation after a pause."""
-
     engine = engine_client(raw_request)
 
     try:
@@ -97,7 +98,6 @@ async def abort_requests(raw_request: Request) -> JSONResponse:
 
     Empty/missing ``request_ids`` aborts all in-flight requests.
     """
-
     engine = engine_client(raw_request)
 
     try:
@@ -139,7 +139,6 @@ async def abort_requests(raw_request: Request) -> JSONResponse:
 @router.get("/is_paused")
 async def is_paused(raw_request: Request) -> JSONResponse:
     """Return the current pause status."""
-
     engine = engine_client(raw_request)
 
     try:
@@ -234,9 +233,12 @@ async def get_world_size(
     """Get the world size from the parallel config.
 
     Args:
+        raw_request: The incoming FastAPI request, used to reach the engine
+            client on the app state.
         include_dp: If True (default), returns the world size including
             data parallelism (TP * PP * DP). If False, returns the world
             size without data parallelism (TP * PP).
+
     """
     parallel_config = engine_client(raw_request).vllm_config.parallel_config
     if include_dp:
diff --git a/vllm/entrypoints/serve/dev/sleep/api_router.py b/vllm/entrypoints/serve/dev/sleep/api_router.py
index fedff8da749b..846a487706b7 100644
--- a/vllm/entrypoints/serve/dev/sleep/api_router.py
+++ b/vllm/entrypoints/serve/dev/sleep/api_router.py
@@ -27,6 +27,12 @@ async def sleep(raw_request: Request):
     return Response(status_code=200)
 
 
+@router.post("/release_kv_cache_memory")
+async def release_kv_cache_memory(raw_request: Request):
+    await engine_client(raw_request).release_kv_cache_memory()
+    return Response(status_code=200)
+
+
 @router.post("/wake_up")
 async def wake_up(raw_request: Request):
     tags = raw_request.query_params.getlist("tags")
diff --git a/vllm/entrypoints/serve/elastic_ep/middleware.py b/vllm/entrypoints/serve/elastic_ep/middleware.py
index 23f45eafeaa0..dcc62aea2c18 100644
--- a/vllm/entrypoints/serve/elastic_ep/middleware.py
+++ b/vllm/entrypoints/serve/elastic_ep/middleware.py
@@ -20,8 +20,7 @@ def set_scaling_elastic_ep(value):
 
 
 class ScalingMiddleware:
-    """
-    Middleware that checks if the model is currently scaling and
+    """Middleware that checks if the model is currently scaling and
     returns a 503 Service Unavailable response if it is.
 
     This middleware applies to all HTTP requests and prevents
diff --git a/vllm/entrypoints/serve/engine/serving.py b/vllm/entrypoints/serve/engine/serving.py
index 23e392720313..86ba524946b2 100644
--- a/vllm/entrypoints/serve/engine/serving.py
+++ b/vllm/entrypoints/serve/engine/serving.py
@@ -41,6 +41,9 @@ async def _check_model(
         self,
         request: AnyRequest | AnyPoolingRequest,
     ) -> ErrorResponse | None:
+        if self.request_logger is not None:
+            self.request_logger.log_request_body(request)
+
         error_response = None
 
         if self._is_model_supported(request.model):
@@ -117,7 +120,7 @@ def _log_inputs(
     def _base_request_id(
         raw_request: Request | None, default: str | None = None
     ) -> str | None:
-        """Pulls the request id to use from a header, if provided"""
+        """Pulls the request id to use from a header, if provided."""
         if raw_request is not None and (
             (req_id := raw_request.headers.get("X-Request-Id")) is not None
         ):
diff --git a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py
index 3ee5f0a06514..c0234139374e 100644
--- a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py
+++ b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py
@@ -28,8 +28,7 @@ async def vllm_error_handler(req: Request, exc: VLLMError):
 async def engine_error_handler(
     req: Request, exc: EngineDeadError | EngineGenerateError
 ):
-    """
-    VLLM V1 AsyncLLM catches exceptions and returns
+    """VLLM V1 AsyncLLM catches exceptions and returns
     only two types: EngineGenerateError and EngineDeadError.
 
     EngineGenerateError is raised by the per request generate()
@@ -51,7 +50,6 @@ async def engine_error_handler(
     will not automatically shut down. Instead, we use the watchdog
     background task for check for errored state.
     """
-
     if req.app.state.args.log_error_stack:
         logger.exception(
             "Engine Exception caught. Request id: %s",
diff --git a/vllm/entrypoints/serve/exception_handling/register.py b/vllm/entrypoints/serve/exception_handling/register.py
index 9f94873939a0..13c9425511fe 100644
--- a/vllm/entrypoints/serve/exception_handling/register.py
+++ b/vllm/entrypoints/serve/exception_handling/register.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-vLLM Exception handlers are registered in four layers:
+"""vLLM Exception handlers are registered in four layers:
 1. framework errors raised by FastAPI/Starlette
 2. vLLM-specific errors dispatched via a single ``VLLMError`` handler
 3. fallback handlers for raw exceptions not yet migrated to ``VLLMError``
diff --git a/vllm/entrypoints/serve/instrumentator/basic.py b/vllm/entrypoints/serve/instrumentator/basic.py
index be091a1f4330..64201db3f3d9 100644
--- a/vllm/entrypoints/serve/instrumentator/basic.py
+++ b/vllm/entrypoints/serve/instrumentator/basic.py
@@ -54,3 +54,10 @@ async def get_server_load_metrics(request: Request):
 async def show_version():
     ver = {"version": VLLM_VERSION}
     return JSONResponse(content=ver)
+
+
+@router.get("/kv_event_sources")
+async def kv_event_sources(raw_request: Request):
+    """KV-event publisher config per DP rank; empty on render-only servers."""
+    client: EngineClient | None = raw_request.app.state.engine_client
+    return client.get_kv_event_sources() if client is not None else {}
diff --git a/vllm/entrypoints/serve/instrumentator/metrics.py b/vllm/entrypoints/serve/instrumentator/metrics.py
index 5cba364f5d94..13241d03e884 100644
--- a/vllm/entrypoints/serve/instrumentator/metrics.py
+++ b/vllm/entrypoints/serve/instrumentator/metrics.py
@@ -55,7 +55,6 @@ class PrometheusResponse(Response):
 
 def attach_router(app: FastAPI):
     """Mount prometheus metrics to a FastAPI app."""
-
     registry = get_prometheus_registry()
 
     # `response_class=PrometheusResponse` is needed to return an HTTP response
diff --git a/vllm/entrypoints/serve/lora/api_router.py b/vllm/entrypoints/serve/lora/api_router.py
index 95dff9792b2f..61436371c27d 100644
--- a/vllm/entrypoints/serve/lora/api_router.py
+++ b/vllm/entrypoints/serve/lora/api_router.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
 
-
-import model_hosting_container_standards.sagemaker as sagemaker_standards
 from fastapi import APIRouter, Depends, FastAPI, Request
 from fastapi.responses import JSONResponse, Response
 
@@ -22,6 +21,12 @@
 
 
 def attach_router(app: FastAPI):
+    """Attach the LoRA adapter load/unload endpoints to the API server.
+
+    Does nothing when dynamic LoRA updating is disabled. Handler levels are
+    snapshotted and restored because importing
+    model_hosting_container_standards may reconfigure root logging.
+    """
     if not envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING:
         """If LoRA dynamic loading & unloading is not enabled, do nothing."""
         return
@@ -30,6 +35,23 @@ def attach_router(app: FastAPI):
         "This should ONLY be used for local development!"
     )
 
+    snapshot = [
+        (h, h.level)
+        for lg in (logging.getLogger(), logging.getLogger("vllm"))
+        for h in lg.handlers
+    ]
+
+    try:
+        _attach_router(app)
+    finally:
+        for handler, level in snapshot:
+            handler.setLevel(level)
+
+
+def _attach_router(app: FastAPI):
+    """Register the LoRA adapter load/unload routes on the app."""
+    import model_hosting_container_standards.sagemaker as sagemaker_standards
+
     @sagemaker_standards.register_load_adapter_handler(
         request_shape={
             "lora_name": "body.name",
@@ -40,6 +62,8 @@ def attach_router(app: FastAPI):
     )
     @router.post("/v1/load_lora_adapter", dependencies=[Depends(validate_json_request)])
     async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Request):
+        """Handle POST /v1/load_lora_adapter: load a LoRA adapter into the
+        serving engine."""
         handler: OpenAIServingModels = models(raw_request)
         response = await handler.load_lora_adapter(request)
         if isinstance(response, ErrorResponse):
@@ -60,6 +84,8 @@ async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Reques
     async def unload_lora_adapter(
         request: UnloadLoRAAdapterRequest, raw_request: Request
     ):
+        """Handle POST /v1/unload_lora_adapter: unload a LoRA adapter from
+        the serving engine."""
         handler: OpenAIServingModels = models(raw_request)
         response = await handler.unload_lora_adapter(request)
         if isinstance(response, ErrorResponse):
diff --git a/vllm/entrypoints/serve/middleware/authenticate.py b/vllm/entrypoints/serve/middleware/authenticate.py
index e8a61b21ee46..0a0c1daeec4d 100644
--- a/vllm/entrypoints/serve/middleware/authenticate.py
+++ b/vllm/entrypoints/serve/middleware/authenticate.py
@@ -12,8 +12,7 @@
 
 
 class AuthenticationMiddleware:
-    """
-    Pure ASGI middleware that authenticates each request by checking
+    """Pure ASGI middleware that authenticates each request by checking
     if the Authorization Bearer token exists and equals anyof "{api_key}".
 
     Notes
@@ -21,6 +20,7 @@ class AuthenticationMiddleware:
     There are two cases in which authentication is skipped:
         1. The HTTP method is OPTIONS.
         2. The request path doesn't start with GUARDED_PREFIX (e.g. /health).
+
     """
 
     def __init__(self, app: ASGIApp, tokens: list[str]) -> None:
diff --git a/vllm/entrypoints/serve/middleware/log_response.py b/vllm/entrypoints/serve/middleware/log_response.py
index e2f5e5bde51e..88f37a75984a 100644
--- a/vllm/entrypoints/serve/middleware/log_response.py
+++ b/vllm/entrypoints/serve/middleware/log_response.py
@@ -94,7 +94,6 @@ def _extract_content_from_chunk(chunk_data: dict) -> str:
 
 def _log_streaming_response(response, response_body: list) -> None:
     """Log streaming response with robust SSE parsing."""
-
     sse_decoder = SSEDecoder()
     chunk_count = 0
 
diff --git a/vllm/entrypoints/serve/middleware/x_request_id.py b/vllm/entrypoints/serve/middleware/x_request_id.py
index 02e85306cd12..1f654caf3c4e 100644
--- a/vllm/entrypoints/serve/middleware/x_request_id.py
+++ b/vllm/entrypoints/serve/middleware/x_request_id.py
@@ -8,8 +8,7 @@
 
 
 class XRequestIdMiddleware:
-    """
-    Middleware the set's the X-Request-Id header for each response
+    """Middleware the set's the X-Request-Id header for each response
     to a random uuid4 (hex) value if the header isn't already
     present in the request, otherwise use the provided request id.
     """
@@ -25,8 +24,7 @@ def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None
         request_headers = Headers(scope=scope)
 
         async def send_with_request_id(message: Message) -> None:
-            """
-            Custom send function to mutate the response headers
+            """Custom send function to mutate the response headers
             and append X-Request-Id to it.
             """
             if message["type"] == "http.response.start":
diff --git a/vllm/entrypoints/serve/sagemaker/api_router.py b/vllm/entrypoints/serve/sagemaker/api_router.py
index 142256c26002..f2fb255ceefc 100644
--- a/vllm/entrypoints/serve/sagemaker/api_router.py
+++ b/vllm/entrypoints/serve/sagemaker/api_router.py
@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import json
+import logging
 from collections.abc import Awaitable, Callable
 from http import HTTPStatus
 from typing import Any
 
-import model_hosting_container_standards.sagemaker as sagemaker_standards
 import pydantic
 from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse, Response
@@ -27,11 +27,55 @@
 EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]
 
 
+def _snapshot_handler_levels() -> list[tuple[logging.Handler, int]]:
+    """Snapshot handler levels for loggers that may be affected by third-party
+    logging configuration side effects (e.g. model_hosting_container_standards
+    calling configure_root_logger() at import time)."""
+    loggers = [logging.getLogger(), logging.getLogger("vllm")]
+    seen: set[int] = set()
+    snapshot: list[tuple[logging.Handler, int]] = []
+    for logger in loggers:
+        for handler in logger.handlers:
+            if id(handler) not in seen:
+                seen.add(id(handler))
+                snapshot.append((handler, handler.level))
+    return snapshot
+
+
+def _restore_handler_levels(
+    snapshot: list[tuple[logging.Handler, int]],
+) -> None:
+    """Restore handler levels from a snapshot."""
+    for handler, level in snapshot:
+        handler.setLevel(level)
+
+
 def attach_router(
     app: FastAPI,
     supported_tasks: tuple["SupportedTask", ...],
     model_config: ModelConfig | None = None,
 ):
+    """Attach the SageMaker hosting endpoints to the API server.
+
+    Handler levels are snapshotted and restored because importing
+    model_hosting_container_standards may reconfigure root logging.
+    """
+    snapshot = _snapshot_handler_levels()
+    try:
+        _attach_router(app, supported_tasks, model_config)
+    finally:
+        _restore_handler_levels(snapshot)
+
+
+def _attach_router(
+    app: FastAPI,
+    supported_tasks: tuple["SupportedTask", ...],
+    model_config: ModelConfig | None = None,
+):
+    """Register the SageMaker hosting routes (/ping, /invocations) on the
+    app."""
+    import model_hosting_container_standards.sagemaker as sagemaker_standards
+
     router = APIRouter()
 
     # NOTE: Construct the TypeAdapters only once
@@ -48,7 +92,7 @@ def attach_router(
     @router.get("/ping", response_class=Response)
     @sagemaker_standards.register_ping_handler
     async def ping(raw_request: Request) -> Response:
-        """Ping check. Endpoint required for SageMaker"""
+        """Ping check. Endpoint required for SageMaker."""
         return await health(raw_request)
 
     @router.post(
@@ -99,4 +143,17 @@ async def invocations(raw_request: Request):
 
 
 def sagemaker_standards_bootstrap(app: FastAPI) -> FastAPI:
-    return sagemaker_standards.bootstrap(app)
+    """Bootstrap the app with the SageMaker hosting standards.
+
+    Handler levels are restored right after the import because importing
+    model_hosting_container_standards may reconfigure root logging, and
+    bootstrap must run with the original levels.
+    """
+    snapshot = _snapshot_handler_levels()
+    try:
+        import model_hosting_container_standards.sagemaker as sagemaker_standards
+
+        app = sagemaker_standards.bootstrap(app)
+    finally:
+        _restore_handler_levels(snapshot)
+    return app
diff --git a/vllm/entrypoints/serve/tokenize/protocol.py b/vllm/entrypoints/serve/tokenize/protocol.py
index 12d487346f76..130a1b4b45c4 100644
--- a/vllm/entrypoints/serve/tokenize/protocol.py
+++ b/vllm/entrypoints/serve/tokenize/protocol.py
@@ -181,8 +181,7 @@ class DetokenizeResponse(OpenAIBaseModel):
 
 
 class TokenizerInfoResponse(OpenAIBaseModel):
-    """
-    Response containing tokenizer configuration
+    """Response containing tokenizer configuration
     equivalent to tokenizer_config.json
     """
 
diff --git a/vllm/entrypoints/serve/utils/api_utils.py b/vllm/entrypoints/serve/utils/api_utils.py
index aa089cff44b4..6b82a2afa070 100644
--- a/vllm/entrypoints/serve/utils/api_utils.py
+++ b/vllm/entrypoints/serve/utils/api_utils.py
@@ -34,7 +34,7 @@
 
 
 async def listen_for_disconnect(request: Request) -> None:
-    """Returns if a disconnect message is received"""
+    """Returns if a disconnect message is received."""
     while True:
         message = await request.receive()
         if message["type"] == "http.disconnect":
diff --git a/vllm/entrypoints/serve/utils/orca_metrics.py b/vllm/entrypoints/serve/utils/orca_metrics.py
index 3808262bf31f..2c0aa4381227 100644
--- a/vllm/entrypoints/serve/utils/orca_metrics.py
+++ b/vllm/entrypoints/serve/utils/orca_metrics.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Utility functions that create ORCA endpoint load report response headers.
-"""
+"""Utility functions that create ORCA endpoint load report response headers."""
 
 import json
 from collections.abc import Mapping
@@ -16,23 +14,22 @@
 def create_orca_header(
     metrics_format: str, named_metrics: list[tuple[str, float]]
 ) -> Mapping[str, str] | None:
-    """
-    Creates ORCA headers named 'endpoint-load-metrics' in the specified format
+    """Creates ORCA headers named 'endpoint-load-metrics' in the specified format
     and adds custom metrics to named_metrics.
     ORCA headers format description: https://docs.google.com/document/d/1C1ybMmDKJIVlrbOLbywhu9iRYo4rilR-cT50OTtOFTs/edit?tab=t.0
     ORCA proto https://github.com/cncf/xds/blob/main/xds/data/orca/v3/orca_load_report.proto
 
-    Parameters:
-    - metrics_format (str): The format of the header ('TEXT', 'JSON').
-    - named_metrics (List[Tuple[str, float]]): List of tuples with metric names
-    and their corresponding double values.
+    Args:
+        metrics_format (str): The format of the header ('TEXT', 'JSON').
+        named_metrics (List[Tuple[str, float]]): List of tuples with metric names
+            and their corresponding double values.
 
     Returns:
     - Optional[Mapping[str,str]]: A dictionary with header key as
     'endpoint-load-metrics' and values as the ORCA header strings with
     format prefix and data in  with named_metrics in.
-    """
 
+    """
     if metrics_format.lower() not in ["text", "json"]:
         logger.warning(
             "Warning: `%s` format is not supported in the ORCA response header",
@@ -71,15 +68,17 @@ def create_orca_header(
 
 
 def get_named_metrics_from_prometheus() -> list[tuple[str, float]]:
-    """
-    Collects current metrics from Prometheus and returns some of them
+    """Collects current metrics from Prometheus and returns some of them
     in the form of the `named_metrics` list for `create_orca_header()`.
 
-    Parameters:
+    Parameters
+    ----------
     - None
 
-    Returns:
+    Returns
+    -------
     - list[tuple[str, float]]: List of tuples of metric names and their values.
+
     """
     named_metrics: list[tuple[str, float]] = []
     # Map from prometheus metric names to ORCA named metrics.
@@ -98,20 +97,20 @@ def get_named_metrics_from_prometheus() -> list[tuple[str, float]]:
 
 
 def metrics_header(metrics_format: str) -> Mapping[str, str] | None:
-    """
-    Creates ORCA headers named 'endpoint-load-metrics' in the specified format.
+    """Creates ORCA headers named 'endpoint-load-metrics' in the specified format.
     Metrics are collected from Prometheus using `get_named_metrics_from_prometheus()`.
 
     ORCA headers format description: https://docs.google.com/document/d/1C1ybMmDKJIVlrbOLbywhu9iRYo4rilR-cT50OTtOFTs/edit?tab=t.0
     ORCA proto https://github.com/cncf/xds/blob/main/xds/data/orca/v3/orca_load_report.proto
 
-    Parameters:
-    - metrics_format (str): The format of the header ('TEXT', 'JSON').
+    Args:
+        metrics_format (str): The format of the header ('TEXT', 'JSON').
 
     Returns:
     - Optional[Mapping[str,str]]: A dictionary with header key as
     'endpoint-load-metrics' and values as the ORCA header strings with
     format prefix and data in  with named_metrics in.
+
     """
     if not metrics_format:
         return None
diff --git a/vllm/entrypoints/serve/utils/request_logger.py b/vllm/entrypoints/serve/utils/request_logger.py
index ac10feab1011..efea243182e0 100644
--- a/vllm/entrypoints/serve/utils/request_logger.py
+++ b/vllm/entrypoints/serve/utils/request_logger.py
@@ -6,6 +6,8 @@
 
 import torch
 
+from vllm.entrypoints.pooling.typing import AnyPoolingRequest
+from vllm.entrypoints.serve.engine.typing import AnyRequest
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.pooling_params import PoolingParams
@@ -67,6 +69,15 @@ def log_inputs(
             lora_request,
         )
 
+    def log_request_body(self, request: AnyRequest | AnyPoolingRequest) -> None:
+        if logger.isEnabledFor(logging.DEBUG):
+            max_log_len = self.max_log_len if self.max_log_len is not None else -1
+            logger.debug(
+                "Request %s JSON body: %s",
+                getattr(request, "request_id", "N/A"),
+                request.model_dump_json(exclude_unset=True)[:max_log_len],
+            )
+
     def log_outputs(
         self,
         request_id: str,
diff --git a/vllm/entrypoints/openai/sse_keep_alive.py b/vllm/entrypoints/serve/utils/sse_keep_alive.py
similarity index 100%
rename from vllm/entrypoints/openai/sse_keep_alive.py
rename to vllm/entrypoints/serve/utils/sse_keep_alive.py
diff --git a/vllm/entrypoints/serve/utils/tool_calls_utils.py b/vllm/entrypoints/serve/utils/tool_calls_utils.py
index 42106f433405..ae2fa24f7f63 100644
--- a/vllm/entrypoints/serve/utils/tool_calls_utils.py
+++ b/vllm/entrypoints/serve/utils/tool_calls_utils.py
@@ -20,7 +20,6 @@ def maybe_filter_parallel_tool_calls(
     choice: _ChatCompletionResponseChoiceT, request: ChatCompletionRequest
 ) -> _ChatCompletionResponseChoiceT:
     """Filter to first tool call only when parallel_tool_calls is explicitly False."""
-
     if request.parallel_tool_calls is not False:
         return choice
 
diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py
index dc09db6aa319..34b7e9c24928 100644
--- a/vllm/entrypoints/speech_to_text/base/serving.py
+++ b/vllm/entrypoints/speech_to_text/base/serving.py
@@ -358,8 +358,7 @@ def _get_verbose_segments(
         segment_class: type[SpeechToTextSegment],
         start_time: float = 0,
     ) -> list[SpeechToTextSegment]:
-        """
-        Convert tokens to verbose segments.
+        """Convert tokens to verbose segments.
 
         This method expects the model to produce
         timestamps as tokens (similar to Whisper).
diff --git a/vllm/entrypoints/speech_to_text/base/utils.py b/vllm/entrypoints/speech_to_text/base/utils.py
index bcd29f08e967..8282d6f41dea 100644
--- a/vllm/entrypoints/speech_to_text/base/utils.py
+++ b/vllm/entrypoints/speech_to_text/base/utils.py
@@ -33,6 +33,7 @@ async def read_upload_with_limit(
 
     Raises:
         VLLMValidationError: If the file exceeds the configured size limit.
+
     """
     if max_size_mb is None:
         max_size_mb = envs.VLLM_MAX_AUDIO_CLIP_FILESIZE_MB
diff --git a/vllm/entrypoints/speech_to_text/realtime/protocol.py b/vllm/entrypoints/speech_to_text/realtime/protocol.py
index bdcd335f6fe4..7796eee8b848 100644
--- a/vllm/entrypoints/speech_to_text/realtime/protocol.py
+++ b/vllm/entrypoints/speech_to_text/realtime/protocol.py
@@ -13,14 +13,14 @@
 
 
 class InputAudioBufferAppend(OpenAIBaseModel):
-    """Append audio chunk to buffer"""
+    """Append audio chunk to buffer."""
 
     type: Literal["input_audio_buffer.append"] = "input_audio_buffer.append"
     audio: str  # base64-encoded PCM16 @ 16kHz
 
 
 class InputAudioBufferCommit(OpenAIBaseModel):
-    """Process accumulated audio buffer"""
+    """Process accumulated audio buffer."""
 
     type: Literal["input_audio_buffer.commit"] = "input_audio_buffer.commit"
     final: bool = False
@@ -28,14 +28,14 @@ class InputAudioBufferCommit(OpenAIBaseModel):
 
 # Server -> Client Events
 class SessionUpdate(OpenAIBaseModel):
-    """Configure session parameters"""
+    """Configure session parameters."""
 
     type: Literal["session.update"] = "session.update"
     model: str | None = None
 
 
 class SessionCreated(OpenAIBaseModel):
-    """Connection established notification"""
+    """Connection established notification."""
 
     type: Literal["session.created"] = "session.created"
     id: str = Field(default_factory=lambda: f"sess-{random_uuid()}")
@@ -43,14 +43,14 @@ class SessionCreated(OpenAIBaseModel):
 
 
 class TranscriptionDelta(OpenAIBaseModel):
-    """Incremental transcription text"""
+    """Incremental transcription text."""
 
     type: Literal["transcription.delta"] = "transcription.delta"
     delta: str  # Incremental text
 
 
 class TranscriptionDone(OpenAIBaseModel):
-    """Final transcription with usage stats"""
+    """Final transcription with usage stats."""
 
     type: Literal["transcription.done"] = "transcription.done"
     text: str  # Complete transcription
@@ -58,7 +58,7 @@ class TranscriptionDone(OpenAIBaseModel):
 
 
 class ErrorEvent(OpenAIBaseModel):
-    """Error notification"""
+    """Error notification."""
 
     type: Literal["error"] = "error"
     error: str
diff --git a/vllm/entrypoints/speech_to_text/realtime/serving.py b/vllm/entrypoints/speech_to_text/realtime/serving.py
index bec2b95143b6..9745269607ab 100644
--- a/vllm/entrypoints/speech_to_text/realtime/serving.py
+++ b/vllm/entrypoints/speech_to_text/realtime/serving.py
@@ -68,6 +68,7 @@ async def transcribe_realtime(
 
         Yields:
             StreamingInput objects containing audio prompts for the engine
+
         """
         model_config = self.model_config
         renderer = self.renderer
diff --git a/vllm/env_override.py b/vllm/env_override.py
index a931b2305fa8..bfcdfa7af236 100644
--- a/vllm/env_override.py
+++ b/vllm/env_override.py
@@ -5,21 +5,28 @@
 import os
 
 
-def _get_torch_cuda_version():
-    """Peripheral function to _maybe_set_cuda_compatibility_path().
+def _get_torch_root():
+    """Locate the installed torch package without importing it."""
+    spec = importlib.util.find_spec("torch")
+    if not spec:
+        return None
+    if spec.origin:
+        return os.path.dirname(spec.origin)
+    if spec.submodule_search_locations:
+        return spec.submodule_search_locations[0]
+    return None
+
+
+def _get_torch_version_attr(attr):
+    """Read an attribute of torch.version without importing torch.
+
     PyTorch version must not be determined by importing directly
     because it will trigger the CUDA initialization, losing the
     chance to set the LD_LIBRARY_PATH beforehand.
     """
     try:
-        spec = importlib.util.find_spec("torch")
-        if not spec:
-            return None
-        if spec.origin:
-            torch_root = os.path.dirname(spec.origin)
-        elif spec.submodule_search_locations:
-            torch_root = spec.submodule_search_locations[0]
-        else:
+        torch_root = _get_torch_root()
+        if not torch_root:
             return None
         version_path = os.path.join(torch_root, "version.py")
         if not os.path.exists(version_path):
@@ -31,11 +38,16 @@ def _get_torch_cuda_version():
         module = importlib.util.module_from_spec(ver_spec)
         # Avoid registering in sys.modules to not confuse future imports
         ver_spec.loader.exec_module(module)
-        return getattr(module, "cuda", None)
+        return getattr(module, attr, None)
     except Exception:
         return None
 
 
+def _get_torch_cuda_version():
+    """Peripheral function to _maybe_set_cuda_compatibility_path()."""
+    return _get_torch_version_attr("cuda")
+
+
 def _maybe_set_cuda_compatibility_path():
     """Set LD_LIBRARY_PATH for CUDA forward compatibility if enabled.
 
@@ -82,7 +94,58 @@ def _maybe_set_cuda_compatibility_path():
     os.environ["LD_LIBRARY_PATH"] = os.pathsep.join(new_paths)
 
 
+def _maybe_promote_torch_symbols_for_rocm():
+    """Put libtorch_cpu.so in the global symbol scope on ROCm, for GPU profiling.
+
+    Must run before 'import torch'. libkineto advertises itself to
+    rocprofiler-sdk by exporting rocprofiler_configure from libtorch_cpu.so,
+    and rocprofiler-sdk looks for that symbol exactly once, from a lazy
+    initializer that runs *during* 'import torch'. CPython loads extension
+    modules RTLD_LOCAL, so unless something has already placed the symbol in
+    the global scope the lookup misses and no client is ever registered.
+
+    The failure is silent and total: no queue interception, so roctracer
+    yields no GPU records, while torch.profiler still writes a complete
+    CPU-only trace and reports success.
+
+    Best effort. If the library cannot be loaded early we leave the process
+    exactly as it would have been, losing only GPU tracing.
+    """
+    if _get_torch_version_attr("hip") is None:
+        return
+    try:
+        import ctypes
+        import glob
+
+        torch_root = _get_torch_root()
+        if not torch_root:
+            return
+        lib = os.path.join(torch_root, "lib", "libtorch_cpu.so")
+        if not os.path.exists(lib):
+            return
+        try:
+            ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
+        except OSError:
+            # Some ROCm wheel layouts leave a stale RUNPATH in libtorch_cpu.so,
+            # so a dependency that resolves once torch has set things up does
+            # not resolve for a standalone load. Satisfy it by search, then
+            # retry; give up quietly if it is not where we expect.
+            site_root = os.path.dirname(torch_root)
+            deps = glob.glob(
+                os.path.join(
+                    site_root, "_rocm_sdk_*", "lib", "*", "lib", "librocm-openblas.so.*"
+                )
+            )
+            if not deps:
+                return
+            ctypes.CDLL(deps[0], mode=ctypes.RTLD_LOCAL)
+            ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
+    except Exception:
+        return
+
+
 _maybe_set_cuda_compatibility_path()
+_maybe_promote_torch_symbols_for_rocm()
 
 import torch
 
@@ -198,8 +261,7 @@ def get_output_names(graph_outputs) -> list[str]:
 def get_graph_partition_signature_patched(
     self, partitions, skip_cudagraphs: list[bool]
 ):
-    """
-    Gets signature for each graph partition, including input nodes, output nodes, and
+    """Gets signature for each graph partition, including input nodes, output nodes, and
     whether deallocating an input within graph partition.
     """
     from torch._inductor import dependencies
@@ -213,8 +275,7 @@ def get_graph_partition_signature_patched(
     name_to_node = self.get_name_to_nodes()
 
     def is_none_layout(buf_name: str) -> bool:
-        """
-        Checks if buf_name is NoneLayout. Buffers with NoneLayout is not allocated
+        """Checks if buf_name is NoneLayout. Buffers with NoneLayout is not allocated
         so graph partition should not take it as inputs or outputs.
         """
         buf = self.name_to_buf.get(buf_name, None)
@@ -362,8 +423,7 @@ def should_partition_patched(self, node, should_log: bool = False) -> bool:
     # torch._inductor.scheduler.Scheduler.should_partition that modifies
     # the following piece of code so that we always return True:
     # https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/_inductor/scheduler.py#L4712-L4724
-    """Return True if we should partition the inductor graph on this node"""
-
+    """Return True if we should partition the inductor graph on this node."""
     import torch._inductor.ir as ir
     from torch._inductor.scheduler import (
         BaseSchedulerNode,
@@ -444,8 +504,7 @@ def _update_scheduler_patched(self) -> None:
     # Copied from torch._inductor.graph.GrahLowering._update_scheduler. Patches
     # this method so that we can patch Scheduler.should_partition with the
     # function above
-    """
-    (Re)initializes the scheduler member.  When initializing the scheduler, no CUBIN
+    """(Re)initializes the scheduler member.  When initializing the scheduler, no CUBIN
     files should be generated (to avoid biasing any benchmarks and pessimizing
     fusion decisions).
     """
@@ -857,6 +916,58 @@ def _patch_inductor_fallback_allow_list() -> None:
 
 _patch_inductor_fallback_allow_list()
 
+
+def _apply_inductor_pattern_matcher_patch() -> None:
+    """Allow custom ops and functionalization wrappers with unsupported dtypes."""
+    from torch._inductor import pattern_matcher
+    from torch._inductor.lowering import fallback_node_due_to_unsupported_type
+
+    def fallback_for_builtin(node, allow_cpu_inputs=True):
+        return (
+            isinstance(node.target, torch._ops.OpOverload)
+            and torch._library.utils.is_builtin(node.target)
+            and fallback_node_due_to_unsupported_type(node, allow_cpu_inputs)
+        )
+
+    pattern_matcher.fallback_node_due_to_unsupported_type = fallback_for_builtin
+
+
+def _patch_inductor_pattern_matcher() -> None:
+    """Apply the backport when Inductor imports its pattern matcher."""
+    import sys
+
+    target_name = "torch._inductor.pattern_matcher"
+    if target_name in sys.modules:
+        _apply_inductor_pattern_matcher_patch()
+        return
+
+    import importlib.abc
+
+    class _PatternMatcherPatchFinder(importlib.abc.MetaPathFinder):
+        def find_spec(self, fullname, path, target=None):
+            if fullname != target_name:
+                return None
+            sys.meta_path.remove(self)
+            spec = importlib.util.find_spec(fullname)
+            if spec is None or spec.loader is None:
+                return None
+            original_exec = spec.loader.exec_module
+
+            def _exec_then_patch(module):
+                original_exec(module)
+                _apply_inductor_pattern_matcher_patch()
+
+            spec.loader.exec_module = _exec_then_patch  # type: ignore[method-assign]
+            return spec
+
+    sys.meta_path.insert(0, _PatternMatcherPatchFinder())
+
+
+# Remove once the minimum supported torch includes
+# https://github.com/pytorch/pytorch/pull/196013.
+if not is_torch_equal_or_newer("2.16.0.dev"):
+    _patch_inductor_pattern_matcher()
+
 # ============================================================
 # Triton Autotuner determinism
 # ============================================================
diff --git a/vllm/envs.py b/vllm/envs.py
index d429686d9cd1..5f62af8f6ac9 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -44,6 +44,7 @@
     VLLM_LOGGING_CONFIG_PATH: str | None = None
     VLLM_LOGGING_COLOR: str = "auto"
     NO_COLOR: bool = False
+    FORCE_COLOR: bool = False
     VLLM_LOG_STATS_INTERVAL: float = 10.0
     VLLM_TRACE_FUNCTION: int = 0
     VLLM_USE_FLASHINFER_SAMPLER: bool = True
@@ -86,7 +87,6 @@
     VLLM_MAX_IMAGE_PIXELS: int = 178_956_970
     VLLM_VIDEO_LOADER_BACKEND: str = "opencv"
     VLLM_MEDIA_CONNECTOR: str = "http"
-    VLLM_MM_HASHER_ALGORITHM: str = "blake3"
     VLLM_TARGET_DEVICE: str = "cuda"
     VLLM_MAIN_CUDA_VERSION: str = "13.0"
     VLLM_FLOAT32_MATMUL_PRECISION: Literal["highest", "high", "medium"] = "highest"
@@ -131,6 +131,7 @@
     VLLM_GDN_DECODE_KERNEL: Literal["cuda", "triton"] = "cuda"
     VLLM_DISABLE_PYNCCL: bool = False
     VLLM_USE_OINK_OPS: bool = False
+    VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD: bool = False
     VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True
     VLLM_ROCM_USE_AITER: bool = False
     VLLM_ROCM_USE_AITER_CUSTOM_AR: bool = True
@@ -138,11 +139,12 @@
     VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False
     VLLM_ROCM_USE_AITER_MOE: bool = True
     VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0
-    VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False
+    VLLM_ROCM_USE_AITER_MOE_SITUV2: bool = False
     VLLM_ROCM_USE_AITER_RMSNORM: bool = True
     VLLM_ROCM_USE_AITER_MLA: bool = True
     VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto"
     VLLM_ROCM_USE_AITER_MHA: bool = True
+    VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
     VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False
     VLLM_ROCM_USE_AITER_FP8BMM: bool = True
     VLLM_ROCM_USE_AITER_FP4BMM: bool = True
@@ -155,7 +157,6 @@
     VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False
     VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
     VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
-    VLLM_PLE_CPU_OFFLOAD: bool = False
     VLLM_DISABLE_COMPILE_CACHE: bool = False
     VLLM_REPLICATE_EMBED: bool = False
     VLLM_USE_LAYERNAME: bool = True
@@ -212,7 +213,7 @@
     VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False
     VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False
     VLLM_KIMI_K3_GEMM_AR: bool = True
-    VLLM_KIMI_K3_GEMM_RS: bool = False
+    VLLM_ENABLE_GEMM_RS: bool = False
     VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
     VLLM_USE_FLASHINFER_MOE_INT4: bool = False
     VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None
@@ -262,7 +263,6 @@
     VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True
     VLLM_ENABLE_RESPONSES_API_STORE: bool = False
     VLLM_ENABLE_COHERE_API: bool = False
-    VLLM_ENABLE_SCALE_OUT_ENDPOINTS: bool | None = None
     VLLM_HAS_FLASHINFER_CUBIN: bool = False
     VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False
     VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
@@ -317,15 +317,12 @@
     VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False
     VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False
     VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True
-    VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32
-    VLLM_XPU_ENABLE_XPU_GRAPH: bool = False
     VLLM_XPU_FORCE_N_CONTIG_WEIGHT: bool = False
     VLLM_XPU_USE_SAMPLER_KERNEL: bool = True
     VLLM_XPU_INC_WNA16_BACKEND: Literal["auto", "ark", "w4a16", "w4a8"] = "auto"
     VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
     VLLM_GPU_NIC_PCIE_MAPPING: str = ""
     VLLM_NIC_SELECTION_VARS: str = ""
-    VLLM_PREFIX_CACHE_RETENTION_INTERVAL: int | None = None
     VLLM_ENABLE_HPC_OPS: bool = False
 
 
@@ -355,15 +352,6 @@ def maybe_convert_bool(value: str | None) -> bool | None:
     return bool(int(value))
 
 
-def maybe_convert_scale_out_endpoints(value: str | None) -> bool | None:
-    if value is None:
-        return None
-    normalized = value.strip()
-    if normalized not in ("0", "1"):
-        raise ValueError("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1")
-    return maybe_convert_bool(normalized)
-
-
 def maybe_convert_json_str_or_file(value: str | None) -> dict[str, Any] | None:
     if value is None:
         return None
@@ -405,8 +393,7 @@ def env_with_choices(
     choices: list[str] | Callable[[], list[str]],
     case_sensitive: bool = True,
 ) -> Callable[[], str | None]:
-    """
-    Create a lambda that validates environment variable against allowed choices
+    """Create a lambda that validates environment variable against allowed choices.
 
     Args:
         env_name: Name of the environment variable
@@ -416,6 +403,7 @@ def env_with_choices(
 
     Returns:
         Lambda function for environment_variables dict
+
     """
 
     def _get_validated_env() -> str | None:
@@ -450,8 +438,7 @@ def env_list_with_choices(
     choices: list[str] | Callable[[], list[str]],
     case_sensitive: bool = True,
 ) -> Callable[[], list[str]]:
-    """
-    Create a lambda that validates environment variable
+    """Create a lambda that validates environment variable
     containing comma-separated values against allowed choices
 
     Args:
@@ -463,6 +450,7 @@ def env_list_with_choices(
     Returns:
         Lambda function for environment_variables
         dict that returns list of strings
+
     """
 
     def _get_validated_env_list() -> list[str]:
@@ -505,8 +493,7 @@ def env_set_with_choices(
     choices: list[str] | Callable[[], list[str]],
     case_sensitive: bool = True,
 ) -> Callable[[], set[str]]:
-    """
-    Creates a lambda which that validates environment variable
+    """Creates a lambda which that validates environment variable
     containing comma-separated values against allowed choices which
     returns choices as a set.
     """
@@ -525,6 +512,7 @@ def get_vllm_port() -> int | None:
 
     Raises:
         ValueError: If VLLM_PORT is a URI, suggest k8s service discovery issue.
+
     """
     if "VLLM_PORT" not in os.environ:
         return None
@@ -546,12 +534,19 @@ def get_vllm_port() -> int | None:
         raise ValueError(f"VLLM_PORT '{port}' must be a valid integer") from err
 
 
+def _generate_shm_name() -> str:
+    import pybase64
+
+    # Fit macOS's 30-character limit without truncating the UUID.
+    encoded_uuid = pybase64.urlsafe_b64encode(uuid.uuid4().bytes).decode("ascii")
+    return "vllm_mm_" + encoded_uuid.rstrip("=")
+
+
 def get_env_or_set_default(
     env_name: str,
     default_factory: Callable[[], str],
 ) -> Callable[[], str]:
-    """
-    Create a lambda that returns an environment variable value if set,
+    """Create a lambda that returns an environment variable value if set,
     or generates and sets a default value using the provided factory function.
     """
 
@@ -854,6 +849,8 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_LOGGING_COLOR": lambda: os.getenv("VLLM_LOGGING_COLOR", "auto"),
     # Standard unix flag for disabling ANSI color codes
     "NO_COLOR": lambda: os.getenv("NO_COLOR", "0") != "0",
+    # De-facto standard flag for forcing ANSI color codes (e.g. non-tty case)
+    "FORCE_COLOR": lambda: os.getenv("FORCE_COLOR", "0") != "0",
     # If set, vllm will log stats at this interval in seconds
     # If not set, vllm will log stats every 10 seconds.
     "VLLM_LOG_STATS_INTERVAL": lambda: (
@@ -1070,17 +1067,6 @@ def _resolve_rust_cli_path() -> str | None:
     # imported at runtime.
     # If a non-existing backend is used, an AssertionError will be thrown.
     "VLLM_MEDIA_CONNECTOR": lambda: os.getenv("VLLM_MEDIA_CONNECTOR", "http"),
-    # Hash algorithm for multimodal content hashing.
-    # - "blake3": Default, fast cryptographic hash (not FIPS 140-3 compliant)
-    # - "sha256": FIPS 140-3 compliant, widely supported
-    # - "sha512": FIPS 140-3 compliant, faster on 64-bit systems
-    # Use sha256 or sha512 for FIPS compliance in government/enterprise deployments
-    "VLLM_MM_HASHER_ALGORITHM": env_with_choices(
-        "VLLM_MM_HASHER_ALGORITHM",
-        "blake3",
-        ["blake3", "sha256", "sha512"],
-        case_sensitive=False,
-    ),
     # Path to the XLA persistent cache directory.
     # Only used for XLA devices such as TPUs.
     "VLLM_XLA_CACHE_PATH": lambda: os.path.expanduser(
@@ -1172,11 +1158,6 @@ def _resolve_rust_cli_path() -> str | None:
         if "VLLM_PLUGINS" not in os.environ
         else os.environ["VLLM_PLUGINS"].split(",")
     ),
-    "VLLM_PREFIX_CACHE_RETENTION_INTERVAL": lambda: (
-        int(os.environ["VLLM_PREFIX_CACHE_RETENTION_INTERVAL"])
-        if "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in os.environ
-        else None
-    ),
     # a local directory to look in for unrecognized LoRA adapters.
     # only works if plugins are enabled and
     # VLLM_ALLOW_RUNTIME_LORA_UPDATING is enabled.
@@ -1248,6 +1229,13 @@ def _resolve_rust_cli_path() -> str | None:
     ),
     # Disable aiter ops unless specifically enabled.
     # Acts as a parent switch to enable the rest of the other operations.
+    # Set to 1 to dequantize MXFP4 weights to BF16 once at load time and run as
+    # a BF16 checkpoint (no per-step weight dequant). This improves emulation
+    # latency at the cost of additional device memory. Default off.
+    "VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD": lambda: (
+        os.getenv("VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD", "False").lower()
+        in ("true", "1")
+    ),
     # On hardware without a native MXFP8 kernel (e.g. ROCm gfx942 / MI300), the
     # MXFP8 emulation path dequantizes weights MXFP8->BF16 once at load time and
     # runs as a BF16 checkpoint (no per-step dequant). Set to 0 to fall back to
@@ -1261,7 +1249,7 @@ def _resolve_rust_cli_path() -> str | None:
         os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1")
     ),
     # Use AITER's CustomAllreduce as the custom-allreduce backend inside vLLM's
-    # CudaCommunicator on ROCm.
+    # CudaCommunicator on ROCm. Also enables AITER AG/RS for DP communication.
     "VLLM_ROCM_USE_AITER_CUSTOM_AR": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "True").lower() in ("true", "1")
     ),
@@ -1280,12 +1268,18 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_ROCM_USE_AITER_MOE": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1")
     ),
-    # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up-
-    # interleaved flydsl kernels instead of the default a16w4 separated path.
-    # This is the only flag users need: vLLM picks the kernels by passing
-    # gate_mode to AITER and sets the AITER-side workaround env at init.
-    "VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4": lambda: (
-        os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() in ("true", "1")
+    # Route K3 SiTU MXFP4 MoE through the FlyDSL SiTUv2 path (a4w4 fp4
+    # activations, separated gate/up layout) instead of default a16w4. vLLM
+    # sets AITER_SITUV2_A4W4 at init when this flag is on and clears any
+    # legacy AITER_SITUV2_A8W4 override (AITER checks A8W4 first).
+    # VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 is a deprecated alias for existing
+    # recipes; it does not select a8w4 kernels.
+    # Needs AITER >= v0.1.20 (ROCm/aiter#4463) for the a4w4 dispatch flag
+    # and tuned kimik3_a4w4_*_fmoe.csv rows; otherwise FlyDSL uses heuristics.
+    "VLLM_ROCM_USE_AITER_MOE_SITUV2": lambda: (
+        os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0").lower() in ("true", "1")
+        or os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower()
+        in ("true", "1")
     ),
     # MoE sorting dispatch policy for AITER fused MoE kernels.
     #   0 = auto (default): single-pass for small batches, multi-pass
@@ -1326,6 +1320,11 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_ROCM_USE_AITER_MHA": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1")
     ),
+    # Whether to use aiter fp4 gemm asm.
+    # By default is disabled.
+    "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: (
+        os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1")
+    ),
     # Whether to use aiter rope.
     # By default is disabled.
     "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: (
@@ -1531,7 +1530,8 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_HUMMING_ONLINE_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file(
         os.environ.get("VLLM_HUMMING_ONLINE_QUANT_CONFIG", None)
     ),
-    # The activation dtype config for humming kernel
+    # The activation dtype config for humming kernel. Explicit schemas disable
+    # fallback unless the config includes "allow_fallback": true.
     "VLLM_HUMMING_INPUT_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file(
         os.environ.get("VLLM_HUMMING_INPUT_QUANT_CONFIG", None)
     ),
@@ -1628,9 +1628,11 @@ def _resolve_rust_cli_path() -> str | None:
     # Use the SM100 BF16 GEMM-AR kernel for eligible Kimi-K3 row-parallel
     # attention projections. All TP ranks must belong to one NVLink domain.
     "VLLM_KIMI_K3_GEMM_AR": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_AR", "1"))),
-    # Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel
-    # row-parallel projections. All TP ranks must belong to one NVLink domain.
-    "VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))),
+    # Fuse eligible sequence-parallel row-parallel projections with their TP
+    # reduce-scatter using the SM100 BF16/MXFP8 GEMM-RS kernel (Kimi-K3
+    # attention/shared-expert projections, DeepSeek-V4.1 ``wo_b``). All TP
+    # ranks must belong to one NVLink domain.
+    "VLLM_ENABLE_GEMM_RS": lambda: bool(int(os.getenv("VLLM_ENABLE_GEMM_RS", "0"))),
     # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.
     # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).
     "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(
@@ -1877,12 +1879,6 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_ENABLE_COHERE_API": lambda: bool(
         int(os.getenv("VLLM_ENABLE_COHERE_API", "0"))
     ),
-    # If set to 1, expose the scale-out endpoints on `vllm serve`.
-    # The dedicated `vllm launch render` server exposes render and derender
-    # endpoints when this variable is unset or set to 1.
-    "VLLM_ENABLE_SCALE_OUT_ENDPOINTS": lambda: maybe_convert_scale_out_endpoints(
-        os.getenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS") or None
-    ),
     # If set, use the fp8 mfma in rocm paged attention.
     "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": lambda: bool(
         int(os.getenv("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0"))
@@ -1964,7 +1960,7 @@ def _resolve_rust_cli_path() -> str | None:
     # if not explicitly set.
     "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME": get_env_or_set_default(
         "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME",
-        lambda: f"VLLM_OBJECT_STORAGE_SHM_BUFFER_{uuid.uuid4().hex}",
+        _generate_shm_name,
     ),
     # The size in MB of the buffers (NVL and RDMA) used by DeepEP
     "VLLM_DEEPEP_BUFFER_SIZE_MB": lambda: int(
@@ -2078,11 +2074,6 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_LOG_MODEL_INSPECTION": lambda: bool(
         int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0"))
     ),
-    # Keep Qwen4Exp PLE embedding tables in pinned CPU memory and gather their
-    # rows through UVA on a dedicated CUDA stream.
-    # Legacy fallback for EngramConfig.cpu_offload, which takes precedence.
-    # This environment variable may be removed in a future release.
-    "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "0"))),
     # Debug logging for --enable-mfu-metrics
     "VLLM_DEBUG_MFU_METRICS": lambda: bool(
         int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0"))
@@ -2147,14 +2138,6 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool(
         int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "1"))
     ),
-    # NIXL EP environment variables
-    "VLLM_NIXL_EP_MAX_NUM_RANKS": lambda: int(
-        os.getenv("VLLM_NIXL_EP_MAX_NUM_RANKS", "32")
-    ),
-    # Whether enable XPU graph on Intel GPU
-    "VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool(
-        int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0"))
-    ),
     # Force N-contiguous weight layout for all XPU unquantized linears.
     "VLLM_XPU_FORCE_N_CONTIG_WEIGHT": lambda: bool(
         int(os.getenv("VLLM_XPU_FORCE_N_CONTIG_WEIGHT", "0"))
@@ -2212,6 +2195,10 @@ def _resolve_rust_cli_path() -> str | None:
     # Each op additionally checks its own shape / dtype constraints and falls
     # back to the eager path when they do not hold.
     "VLLM_ENABLE_HPC_OPS": lambda: bool(int(os.getenv("VLLM_ENABLE_HPC_OPS", "0"))),
+    # Whether to skip version suffix when building package
+    "VLLM_SKIP_VERSION_SUFFIX": lambda: bool(
+        int(os.getenv("VLLM_SKIP_VERSION_SUFFIX", "0"))
+    ),
 }
 
 
@@ -2219,8 +2206,7 @@ def _resolve_rust_cli_path() -> str | None:
 
 
 def __getattr__(name: str):
-    """
-    Gets environment variables lazily.
+    """Gets environment variables lazily.
 
     NOTE: After enable_envs_cache() invocation (which triggered after service
     initialization), all environment variables will be cached.
@@ -2231,14 +2217,13 @@ def __getattr__(name: str):
 
 
 def _is_envs_cache_enabled() -> bool:
-    """Checked if __getattr__ is wrapped with functools.cache"""
+    """Checked if __getattr__ is wrapped with functools.cache."""
     global __getattr__
     return hasattr(__getattr__, "cache_clear")
 
 
 def enable_envs_cache() -> None:
-    """
-    Enables caching of environment variables. This is useful for performance
+    """Enables caching of environment variables. This is useful for performance
     reasons, as it avoids the need to re-evaluate environment variables on
     every call.
 
@@ -2259,8 +2244,7 @@ def enable_envs_cache() -> None:
 
 
 def disable_envs_cache() -> None:
-    """
-    Resets the environment variables cache. It could be used to isolate environments
+    """Resets the environment variables cache. It could be used to isolate environments
     between unit tests.
     """
     global __getattr__
@@ -2281,21 +2265,11 @@ def is_set(name: str):
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 
-def validate_environ(hard_fail: bool) -> None:
-    for env in os.environ:
-        if env.startswith("VLLM_") and env not in environment_variables:
-            if hard_fail:
-                raise ValueError(f"Unknown vLLM environment variable detected: {env}")
-            else:
-                logger.warning("Unknown vLLM environment variable detected: %s", env)
-
-
 def compile_factors() -> dict[str, object]:
     """Return env vars used for torch.compile cache keys.
 
     Start with every known vLLM env var; drop entries in `ignored_factors`;
     hash everything else. This keeps the cache key aligned across workers."""
-
     ignored_factors: set[str] = {
         "MAX_JOBS",
         "VLLM_RPC_BASE_PATH",
@@ -2315,7 +2289,6 @@ def compile_factors() -> dict[str, object]:
         "VLLM_CONFIG_ROOT",
         "LD_LIBRARY_PATH",
         "VLLM_SERVER_DEV_MODE",
-        "VLLM_ENABLE_SCALE_OUT_ENDPOINTS",
         "VLLM_DP_MASTER_IP",
         "VLLM_DP_MASTER_PORT",
         "VLLM_NIXL_SIDE_CHANNEL_HOST",
@@ -2381,6 +2354,8 @@ def compile_factors() -> dict[str, object]:
         "LOCAL_RANK",
         "CUDA_VISIBLE_DEVICES",
         "NO_COLOR",
+        "VLLM_SKIP_VERSION_SUFFIX",
+        "FORCE_COLOR",
     }
 
     from vllm.config.utils import normalize_value
diff --git a/vllm/exceptions.py b/vllm/exceptions.py
index 582e83ec4a6d..c7900a0347e7 100644
--- a/vllm/exceptions.py
+++ b/vllm/exceptions.py
@@ -32,6 +32,7 @@ class VLLMValidationError(VLLMClientError):
         message: The error message describing the validation failure.
         parameter: Optional parameter name that failed validation.
         value: Optional value that was rejected during validation.
+
     """
 
     def __init__(
@@ -56,7 +57,7 @@ def __str__(self):
 
 
 class VLLMNotFoundError(VLLMClientError):
-    """vLLM-specific NotFoundError"""
+    """vLLM-specific NotFoundError."""
 
     pass
 
@@ -69,6 +70,7 @@ class LoRAAdapterNotFoundError(VLLMNotFoundError):
 
     Attributes:
         message: The error message string describing the exception
+
     """
 
     message: str
@@ -96,6 +98,7 @@ class VLLMUnprocessableEntityError(VLLMClientError):
         message: The error message describing the unprocessable entity.
         parameter: Optional parameter name that failed validation.
         value: Optional value that was rejected during validation.
+
     """
 
     def __init__(
diff --git a/vllm/forward_context.py b/vllm/forward_context.py
index f57fc2c950bd..9ddc8e9255ed 100644
--- a/vllm/forward_context.py
+++ b/vllm/forward_context.py
@@ -28,8 +28,7 @@
 
 @dataclass(frozen=True)
 class BatchDescriptor:
-    """
-    Batch descriptor for cudagraph dispatching. We should keep the num of
+    """Batch descriptor for cudagraph dispatching. We should keep the num of
     items as minimal as possible to properly and uniquely describe the padded
     batch for cudagraph.
     """
@@ -59,8 +58,17 @@ class BatchDescriptor:
 
 
 def _compute_sp_num_tokens(
-    num_tokens_across_dp_cpu: torch.Tensor, sequence_parallel_size: int
+    num_tokens_across_dp_cpu: torch.Tensor,
+    sequence_parallel_size: int,
+    pcp_size: int = 1,
+    use_ep: bool = True,
 ) -> list[int]:
+    if pcp_size > 1:
+        num_tokens_across_dp_cpu = (
+            num_tokens_across_dp_cpu.repeat_interleave(pcp_size)
+            if use_ep
+            else num_tokens_across_dp_cpu * pcp_size
+        )
     sp_tokens = (
         num_tokens_across_dp_cpu + sequence_parallel_size - 1
     ) // sequence_parallel_size
@@ -99,13 +107,14 @@ def make(
         return DPMetadata(num_tokens_across_dp_cpu)
 
     @contextmanager
-    def sp_local_sizes(self, sequence_parallel_size: int):
-        """
-        Context manager for setting self.local_sizes. Same as self.chunked_sizes
+    def sp_local_sizes(
+        self, sequence_parallel_size: int, pcp_size: int = 1, use_ep: bool = False
+    ):
+        """Context manager for setting self.local_sizes. Same as self.chunked_sizes
         but without any chunking.
         """
         self.local_sizes = _compute_sp_num_tokens(
-            self.num_tokens_across_dp_cpu, sequence_parallel_size
+            self.num_tokens_across_dp_cpu, sequence_parallel_size, pcp_size, use_ep
         )
         try:
             yield self.local_sizes
@@ -209,6 +218,15 @@ def is_forward_context_available() -> bool:
     return _forward_context is not None
 
 
+def in_piecewise_cudagraph() -> bool:
+    """Whether the current forward runs in piecewise cudagraph mode (graph
+    segments separated by eager breaks), at capture or replay time."""
+    return (
+        is_forward_context_available()
+        and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE
+    )
+
+
 def create_forward_context(
     attn_metadata: Any,
     vllm_config: VllmConfig,
diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py
index bda40bebbc86..26b840f76f3f 100644
--- a/vllm/inputs/engine.py
+++ b/vllm/inputs/engine.py
@@ -16,8 +16,7 @@
 
 
 class _InputOptions(TypedDict):
-    """
-    Additional options available to all
+    """Additional options available to all
     [`SingletonInput`][vllm.inputs.engine.SingletonInput] types.
     """
 
@@ -44,11 +43,6 @@ class TokensInput(_InputOptions):
     """Char-level (start, end) offsets per token, propagated from the
     renderer's TokensPrompt when offsets were computed."""
 
-    assistant_tokens_mask: NotRequired[list[int] | None]
-    """Per-token 0/1 mask marking assistant-generated tokens.
-    Populated when ``return_assistant_tokens_mask=True`` is set on the
-    render request and the chat template supports ``{% generation %}``."""
-
 
 def tokens_input(
     prompt_token_ids: list[int],
@@ -56,8 +50,7 @@ def tokens_input(
     prompt: str | None = None,
     cache_salt: str | None = None,
 ) -> TokensInput:
-    """
-    Construct [`TokensInput`][vllm.inputs.engine.TokensInput]
+    """Construct [`TokensInput`][vllm.inputs.engine.TokensInput]
     from optional values.
     """
     inputs = TokensInput(type="token", prompt_token_ids=prompt_token_ids)
@@ -104,8 +97,7 @@ def embeds_input(
     prompt_token_ids: list[int] | None = None,
     is_token_ids: list[bool] | None = None,
 ) -> EmbedsInput:
-    """
-    Construct [`EmbedsInput`][vllm.inputs.engine.EmbedsInput]
+    """Construct [`EmbedsInput`][vllm.inputs.engine.EmbedsInput]
     from optional values.
     """
     inputs = EmbedsInput(type="embeds", prompt_embeds=prompt_embeds)
@@ -158,11 +150,6 @@ class MultiModalInput(_InputOptions):
     `prompt_token_ids`.
     """
 
-    assistant_tokens_mask: NotRequired[list[int] | None]
-    """Per-token 0/1 mask marking assistant-generated tokens.
-    Populated when ``return_assistant_tokens_mask=True`` is set on the
-    render request and the chat template supports ``{% generation %}``."""
-
 
 def mm_input(
     prompt_token_ids: list[int],
@@ -190,13 +177,13 @@ def mm_input(
 
 
 class MultiModalEncDecInput(MultiModalInput):
-    """
-    Represents multi-modal input to the engine for encoder-decoder models.
+    """Represents multi-modal input to the engine for encoder-decoder models.
 
     Note:
         Even text-only encoder-decoder models are currently implemented
         as multi-modal models for convenience.
         (Example: https://github.com/vllm-project/bart-plugin)
+
     """
 
     encoder_prompt_token_ids: list[int]
@@ -253,8 +240,7 @@ def mm_enc_dec_input(
 
 
 class EncoderDecoderInput(TypedDict):
-    """
-    A rendered [`EncoderDecoderPrompt`][vllm.inputs.llm.EncoderDecoderPrompt]
+    """A rendered [`EncoderDecoderPrompt`][vllm.inputs.llm.EncoderDecoderPrompt]
     which can be passed to `LLMEngine.add_request` or `AsyncLLM.add_request`.
     """
 
@@ -315,8 +301,7 @@ def _prepare_decoder_input_ids_for_generation(
     decoder_input_ids: list[int],
     decoder_start_token_id: int,
 ) -> list[int]:
-    """
-    Prepare `decoder_input_ids` for generation with encoder-decoder models,
+    """Prepare `decoder_input_ids` for generation with encoder-decoder models,
     according to `GenerationMixin._prepare_decoder_input_ids_for_generation()`.
 
     Source:
diff --git a/vllm/inputs/llm.py b/vllm/inputs/llm.py
index a9ac8201fbb1..a95b274c1853 100644
--- a/vllm/inputs/llm.py
+++ b/vllm/inputs/llm.py
@@ -62,8 +62,7 @@ class MultiModalDataBuiltins(TypedDict, total=False):
 
 
 class _PromptOptions(TypedDict):
-    """
-    Additional options available to all
+    """Additional options available to all
     [`SingletonPrompt`][vllm.inputs.llm.SingletonPrompt] types.
     """
 
@@ -189,11 +188,11 @@ class EmbedsPrompt(_PromptOptions):
 
 
 class ExplicitEncoderDecoderPrompt(TypedDict):
-    """
-    Schema for a pair of encoder and decoder singleton prompts.
+    """Schema for a pair of encoder and decoder singleton prompts.
 
     Note:
         This schema is not valid for decoder-only models.
+
     """
 
     encoder_prompt: EncoderPrompt
@@ -233,8 +232,7 @@ class ExplicitEncoderDecoderPrompt(TypedDict):
 
 
 class DataPrompt(_PromptOptions):
-    """
-    Represents generic inputs that are converted to
+    """Represents generic inputs that are converted to
     [`PromptType`][vllm.inputs.llm.PromptType] by IO processor plugins.
     """
 
diff --git a/vllm/ir/op.py b/vllm/ir/op.py
index 742d3f33ff85..996bc5ab2ae6 100644
--- a/vllm/ir/op.py
+++ b/vllm/ir/op.py
@@ -52,23 +52,19 @@ def _validate_name(name: str, entity_type: str) -> None:
 
 
 def set_default_torch_wrap(enable: bool = True) -> None:
-    """
-    Permanently set the torch wrap flag.
-    """
+    """Permanently set the torch wrap flag."""
     global _ENABLE_TORCH_WRAP
     _ENABLE_TORCH_WRAP = enable
 
 
 @contextlib.contextmanager
 def enable_torch_wrap(enable: bool = True):
-    """
-    Context manager to enable/disable torch custom op wrapping for vLLM IR ops.
+    """Context manager to enable/disable torch custom op wrapping for vLLM IR ops.
     When torch wrapping is disabled, the torch custom op layer is skipped
     and IR ops dispatch directly to the implementation.
     Helpful for avoiding torch dispatch overhead in eager mode
     and avoiding the need for lowering for platforms not using Inductor.
     """
-
     global _ENABLE_TORCH_WRAP
     old = _ENABLE_TORCH_WRAP
     try:
@@ -110,8 +106,7 @@ def register_op(
     activations: list[str] | None = None,
     allow_inplace: bool = False,
 ) -> "IrOp | Callable[[Callable], IrOp]":
-    """
-    Register a new vLLM IR op.
+    """Register a new vLLM IR op.
 
     Args:
         f: the native implementation of the op
@@ -131,7 +126,10 @@ def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
     @vllm.ir.register_op(name="custom_mul")
     def multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        return x * y"""
+        return x * y
+    ```
+
+    """
 
     def decorator(_f: Callable):
         op_name: str = _f.__name__ if name is None else name
@@ -226,16 +224,14 @@ def __init__(
         self.torch_op: torch._ops.OpOverload = getattr(torch_ops, name).default
 
     def register_fake(self, fn: Callable) -> Callable:
-        """
-        Register a fake impl for the torch custom op. If this method is not called,
+        """Register a fake impl for the torch custom op. If this method is not called,
         the native implementation is used directly for the fake implementation.
         """
         self._fake_fn = fn
         return fn
 
     def _fake_call(self, *args, **kwargs) -> Any:
-        """
-        Call to the fake implementation of the op. We use indirection because we want
+        """Call to the fake implementation of the op. We use indirection because we want
         users to be able to register fake later but also want it to fall back to native
         directly by default, instead of going through the dispatching mechanism.
         """
@@ -249,8 +245,8 @@ def register_impl(
         supports_args: Callable[..., bool] | None = None,
         inplace: bool = False,
     ) -> Callable[[Callable[..., Any]], "IrOpImpl"]:
-        """
-        Register an implementation for this custom op.
+        """Register an implementation for this custom op.
+
         Args:
             provider: The name of the provider, must be unique.
             supported: Static support check, use this to check platform support.
@@ -302,8 +298,7 @@ def _register_impl(f: Callable):
         return _register_impl
 
     def _inner_call(self, *args, **kwargs) -> Any:
-        """
-        Eager call to torch op lands here. When torch wrapping is disabled,
+        """Eager call to torch op lands here. When torch wrapping is disabled,
         __call__ routes straight here instead of going through torch op dispatching.
         """
         impl = self.dispatch(*args, **kwargs)
@@ -313,8 +308,7 @@ def _inner_call(self, *args, **kwargs) -> Any:
         return impl.func_impl_fn(*args, **kwargs)
 
     def apply_arg_defaults(self, args) -> tuple:
-        """
-        Return args with default values applied.
+        """Return args with default values applied.
         Defaults are taken from the native implementation signature.
 
         SHOULD NOT BE USED IN THE DISPATCH PATH (SLOW).
@@ -325,8 +319,7 @@ def apply_arg_defaults(self, args) -> tuple:
         return bound_args.args
 
     def dispatch(self, *args, **kwargs) -> "IrOpImpl":
-        """
-        Dispatch to the appropriate implementation based on current priority
+        """Dispatch to the appropriate implementation based on current priority
         and argument support checks. Returns the selected IrOpImpl.
 
         THIS FUNCTION IS ON THE HOT PATH (OP DISPATCH), MUST BE FAST.
@@ -413,8 +406,7 @@ def _filter_priority_impls(self, priority: list[str]) -> list["IrOpImpl"]:
         return filtered_impls
 
     def set_default(self, priority: list[str]) -> None:
-        """
-        Permanently set the dispatch priority for this op. Use this for
+        """Permanently set the dispatch priority for this op. Use this for
         process-lifetime setup (e.g., worker startup). For scoped overrides,
         use ``set_priority`` instead.
         """
@@ -427,9 +419,7 @@ def set_default(self, priority: list[str]) -> None:
 
     @contextlib.contextmanager
     def set_priority(self, priority: list[str]):
-        """
-        Context manager to set the dispatch priority for implementations for this op.
-        """
+        """Context manager setting the dispatch priority for this op's impls."""
         old_priority_impls = self._priority_impls
         try:
             self._priority_impls = self._filter_priority_impls(priority)
@@ -636,10 +626,10 @@ def supports_args(self, *args, **kwargs) -> bool:
 
     @weak_cache
     def uuid(self):
-        """
-        Compile-time hash to uniquely determine whether the implementation has changed.
-        Used by vllm-compile hash mechanism and torch.compile lowering pass uuid to
-        control the vLLM compile cache and AOTAutograd/Inductor caches respectively.
+        """Compile-time hash to uniquely determine whether the implementation has
+        changed. Used by vllm-compile hash mechanism and torch.compile lowering pass
+        uuid to control the vLLM compile cache and AOTAutograd/Inductor caches
+        respectively.
 
         Source file contents do not change so we cache uuid.
         TODO(luka): Cache the file hash as multiple impls are likely in the same file.
@@ -648,8 +638,7 @@ def uuid(self):
         return hash_source(*sources)
 
     def func_impl_fn(self, *args, **kwargs) -> Any:
-        """
-        Copy any inputs in activations if this is an inplace impl,
+        """Copy any inputs in activations if this is an inplace impl,
         to ensure functional semantics.
         """
         if not self.inplace:
diff --git a/vllm/ir/ops/__init__.py b/vllm/ir/ops/__init__.py
index d4d71afef723..13e2b7dda988 100644
--- a/vllm/ir/ops/__init__.py
+++ b/vllm/ir/ops/__init__.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from .activation import gelu_and_mul_sparse
 from .layernorm import fused_add_rms_norm, rms_norm
 
-__all__ = ["rms_norm", "fused_add_rms_norm"]
+__all__ = ["rms_norm", "fused_add_rms_norm", "gelu_and_mul_sparse"]
diff --git a/vllm/ir/ops/activation.py b/vllm/ir/ops/activation.py
new file mode 100644
index 000000000000..149197d73550
--- /dev/null
+++ b/vllm/ir/ops/activation.py
@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from ..op import register_op
+
+
+@register_op
+def gelu_and_mul_sparse(
+    x: Tensor, std_multiplier: float, approximate: str = "none"
+) -> Tensor:
+    """Apply Gaussian sparsification, GELU, and gated multiplication."""
+    d = x.shape[-1] // 2
+    gate = x[..., :d]
+    # Statistics intentionally remain local to each tensor-parallel shard.
+    mean = torch.mean(gate, dim=-1, keepdim=True)
+    std = torch.std(gate, dim=-1, keepdim=True, unbiased=False)
+    sparse_gate = F.relu(gate - (mean + std * std_multiplier))
+    return F.gelu(sparse_gate, approximate=approximate) * x[..., d:]
+
+
+@gelu_and_mul_sparse.register_input_generator
+def _gelu_and_mul_sparse_input_generator(
+    num_tokens: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    std_multiplier: float = 1.6448536269514722,
+    approximate: str = "tanh",
+) -> tuple:
+    x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype)
+    return x, std_multiplier, approximate
diff --git a/vllm/ir/ops/layernorm.py b/vllm/ir/ops/layernorm.py
index a40028a4d0bc..05bfca73bde0 100644
--- a/vllm/ir/ops/layernorm.py
+++ b/vllm/ir/ops/layernorm.py
@@ -10,7 +10,7 @@
 def rms_norm(
     x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
 ) -> Tensor:
-    """Weighted root-mean-square layer normalization"""
+    """Weighted root-mean-square layer normalization."""
     orig_dtype = x.dtype
     x = x.to(torch.float32)
     x_var = x if variance_size is None else x[..., :variance_size]
@@ -48,7 +48,7 @@ def fused_add_rms_norm(
     epsilon: float,
     variance_size: int | None = None,
 ) -> tuple[Tensor, Tensor]:
-    """Fused add and weighted root-mean-square layer normalization"""
+    """Fused add and weighted root-mean-square layer normalization."""
     orig_dtype = x.dtype
     x = x.to(torch.float32)
     x = x + x_residual.to(torch.float32)
diff --git a/vllm/ir/util.py b/vllm/ir/util.py
index e9240f487aca..150d4219e6cc 100644
--- a/vllm/ir/util.py
+++ b/vllm/ir/util.py
@@ -10,11 +10,12 @@
 
 
 def hash_source(*srcs: str | Any) -> str:
-    """
-    Utility method to hash the sources of functions or objects.
+    """Utility method to hash the sources of functions or objects.
+
     Args:
         srcs: strings or objects to add to the hash.
             Objects and functions have their source inspected.
+
     """
     hasher = hashlib.sha256()
     for src in srcs:
@@ -34,8 +35,7 @@ def hash_source(*srcs: str | Any) -> str:
 
 
 def weak_lru_cache(maxsize: int | None = 128, typed: bool = False):
-    """
-    LRU Cache decorator that keeps a weak reference to 'self'.
+    """LRU Cache decorator that keeps a weak reference to 'self'.
     This avoids memory leakage, which happens when functools.lru_cache
     stores a reference to self in the global cache.
 
@@ -57,5 +57,5 @@ def inner(self, *args, **kwargs):
 
 
 def weak_cache(user_function, /):
-    """Simple weak equivalent to functools.cache"""
+    """Simple weak equivalent to functools.cache."""
     return weak_lru_cache(maxsize=None)(user_function)
diff --git a/vllm/kernels/__init__.py b/vllm/kernels/__init__.py
index e83e2772f06d..4821ec43eccd 100644
--- a/vllm/kernels/__init__.py
+++ b/vllm/kernels/__init__.py
@@ -2,6 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Kernel implementations for vLLM."""
 
-from . import aiter_ops, oink_ops, vllm_c
+from . import aiter_ops, oink_ops, triton, vllm_c
 
-__all__ = ["vllm_c", "aiter_ops", "oink_ops"]
+__all__ = ["vllm_c", "aiter_ops", "oink_ops", "triton"]
diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py
index de8a53ae8b77..402c9cac1752 100644
--- a/vllm/kernels/aiter_ops.py
+++ b/vllm/kernels/aiter_ops.py
@@ -34,13 +34,38 @@ def is_aiter_found() -> bool:
 AITER_SUPPORTED = is_aiter_found()
 """Most kernels in this file are supported if AITER is installed."""
 
+
+def flatten_to_2d_is_free(x: Tensor) -> bool:
+    """Whether ``x.reshape(-1, x.shape[-1])`` is a view with unit-stride rows.
+
+    The AITER norm kernels only take dense 2D inputs, so the wrappers below
+    flatten the leading dims with ``Tensor.reshape``, which silently falls back
+    to ``contiguous()`` when the flattened shape is not expressible with the
+    existing strides, adding a whole-tensor device copy that costs more than
+    the kernel saves. A strided last dim flattens without a copy but not into a
+    dense buffer, so it is rejected too.
+    """
+    if x.dim() <= 2:
+        return True
+    if x.stride(-1) != 1:
+        return False
+    expected_stride = x.size(-1)
+    for i in range(x.dim() - 2, -1, -1):
+        if x.size(i) != 1 and x.stride(i) != expected_stride:
+            return False
+        expected_stride *= x.size(i)
+    return True
+
+
 rms_no_var_16bit_only = (
     lambda x, weight, epsilon, variance_size=None: variance_size is None
     and x.dtype in (torch.float16, torch.bfloat16)
     and (weight is None or weight.dtype == x.dtype)
+    and flatten_to_2d_is_free(x)
 )
 """AITER rms_norm only supports float16 and bfloat16 acts, no var_size override,
-and requires weight dtype to match x dtype."""
+requires weight dtype to match x dtype, and requires flattening the activation
+to 2D to be free."""
 
 
 @ir.ops.rms_norm.register_impl(
@@ -79,10 +104,13 @@ def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor
     lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
     and x.dtype in (torch.float16, torch.bfloat16)
     and (weight is None or weight.dtype == x.dtype)
+    and flatten_to_2d_is_free(x)
+    and flatten_to_2d_is_free(x_residual)
 )
 """
 AITER fused_add_rms_norm only supports 16-bit activations and no var_size override.
-Requires weight dtype to match x dtype.
+Requires weight dtype to match x dtype, and flattening both the activation and the
+residual to 2D to be free.
 """
 
 
diff --git a/vllm/kernels/helion/case_key.py b/vllm/kernels/helion/case_key.py
index 32b544de39cc..0f90dac610aa 100644
--- a/vllm/kernels/helion/case_key.py
+++ b/vllm/kernels/helion/case_key.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Structured key for identifying kernel config/autotune/benchmark cases.
-"""
+"""Structured key for identifying kernel config/autotune/benchmark cases."""
 
 from __future__ import annotations
 
diff --git a/vllm/kernels/helion/config_manager.py b/vllm/kernels/helion/config_manager.py
index 052b515c37ff..8887cdaea48a 100644
--- a/vllm/kernels/helion/config_manager.py
+++ b/vllm/kernels/helion/config_manager.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Configuration management for Helion kernels.
+"""Configuration management for Helion kernels.
 
 This module provides centralized configuration file management for Helion custom
 operations, including naming conventions, directory resolution, and file I/O.
diff --git a/vllm/kernels/helion/ops/__init__.py b/vllm/kernels/helion/ops/__init__.py
index 588c070a1245..0fb82f6cfe0d 100644
--- a/vllm/kernels/helion/ops/__init__.py
+++ b/vllm/kernels/helion/ops/__init__.py
@@ -22,6 +22,7 @@ def import_all_kernels() -> list[str]:
 
     Returns:
         The fully-qualified module names that were imported.
+
     """
     imported: list[str] = []
     for module_info in pkgutil.iter_modules(__path__):
diff --git a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py
index 0ebbf724094e..9a566b0eda41 100644
--- a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py
+++ b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py
@@ -68,7 +68,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
          the smallest num_tokens >= the input's num_tokens. If the input is
          larger than all available num_tokens, fall back to the largest.
     """
-
     if not config_keys:
         return None
 
diff --git a/vllm/kernels/helion/ops/fused_qk_norm_rope.py b/vllm/kernels/helion/ops/fused_qk_norm_rope.py
index dbbb9e088641..2fc7cb5dcb06 100644
--- a/vllm/kernels/helion/ops/fused_qk_norm_rope.py
+++ b/vllm/kernels/helion/ops/fused_qk_norm_rope.py
@@ -118,7 +118,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
          the smallest num_tokens >= the input's num_tokens. If the input is
          larger than all available num_tokens, fall back to the largest.
     """
-
     if not config_keys:
         return None
 
diff --git a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py
index 789fa8dd65fa..cbb19eb800d9 100644
--- a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py
+++ b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py
@@ -94,7 +94,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
          the smallest num_tokens >= the input's num_tokens. If the input is
          larger than all available num_tokens, fall back to the largest.
     """
-
     if not config_keys:
         return None
 
diff --git a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py
index 276c7f96701b..90a48abf65b9 100644
--- a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py
+++ b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py
@@ -87,7 +87,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
          the smallest num_tokens >= the input's num_tokens. If the input is
          larger than all available num_tokens, fall back to the largest.
     """
-
     if not config_keys:
         return None
 
diff --git a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py
index 8c357386bdd5..adf98601326a 100644
--- a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py
+++ b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py
@@ -112,7 +112,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
          the smallest num_tokens >= the input's num_tokens. If the input is
          larger than all available num_tokens, fall back to the largest.
     """
-
     if not config_keys:
         return None
 
diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py
index 7cc0f99ea5ed..96e24cd988f6 100644
--- a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py
+++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py
@@ -96,7 +96,6 @@ def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey |
          pick the smallest num_tokens >= the input's num_tokens. If the input is
          larger than all available num_tokens, fall back to the largest.
     """
-
     if not config_keys:
         return None
 
diff --git a/vllm/kernels/helion/register.py b/vllm/kernels/helion/register.py
index 438011b1aa69..e8305c9d3ecf 100644
--- a/vllm/kernels/helion/register.py
+++ b/vllm/kernels/helion/register.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-vLLM Helion kernel registration with pre-tuned config selection.
+"""vLLM Helion kernel registration with pre-tuned config selection.
 
 This module leverages Helion's internal config selection infrastructure to use
 pre-tuned configs instead of runtime autotuning.
@@ -164,8 +163,7 @@ def __call__(self, *args, **kwargs):
         return self._decorated_kernel(*args, **kwargs)
 
     def _create_key_computer(self):
-        """
-        Create a key computer function derived from the config picker.
+        """Create a key computer function derived from the config picker.
 
         The returned function receives kernel arguments unpacked (*args) to match
         Helion's key signature (called as self._key_fn(*args)).
@@ -412,6 +410,10 @@ def register_kernel(
     """Register a Helion kernel with pre-tuned config selection.
 
     Args:
+        op_name: Name to register the op under. Defaults to the function name.
+        fake_impl: Optional meta/fake implementation for tracing.
+        mutates_args: Names of arguments the kernel mutates in place.
+        helion_settings: Optional Helion settings override.
         config_picker: Required. Receives ``(args, config_keys)``
             where each config key is a ``dict[str, Any]`` mapping
             parameter names to values.  Return the best-matching
@@ -435,6 +437,7 @@ def generate_inputs():
                         "4096": (torch.randn(4096, device="cuda"), 0.5),
                         "8192": (torch.randn(8192, device="cuda"), 0.5),
                     }
+
     """
 
     def decorator(kernel_func: Callable) -> HelionKernelWrapper:
diff --git a/vllm/kernels/helion/utils.py b/vllm/kernels/helion/utils.py
index f4ace6cff3bf..7887fc9a5996 100644
--- a/vllm/kernels/helion/utils.py
+++ b/vllm/kernels/helion/utils.py
@@ -60,8 +60,7 @@ def get_gpu_name(device_id: int | None = None) -> str:
 
 
 def canonicalize_gpu_name(name: str) -> str:
-    """
-    Canonicalize GPU name for use as a platform identifier.
+    """Canonicalize GPU name for use as a platform identifier.
 
     Converts to lowercase, replaces separators with underscores,
     and maps known variant names to their canonical form via _GPU_NAME_ALIASES.
diff --git a/vllm/kernels/triton/__init__.py b/vllm/kernels/triton/__init__.py
index 6626213145ca..3a3833505062 100644
--- a/vllm/kernels/triton/__init__.py
+++ b/vllm/kernels/triton/__init__.py
@@ -1,3 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Triton kernel implementations."""
+
+from . import activation
+
+__all__ = ["activation"]
diff --git a/vllm/kernels/triton/activation.py b/vllm/kernels/triton/activation.py
new file mode 100644
index 000000000000..2bda1f51b4d1
--- /dev/null
+++ b/vllm/kernels/triton/activation.py
@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch import Tensor
+from torch.library import triton_op, wrap_triton
+
+from vllm import ir
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+_MAX_INTERMEDIATE_SIZE = 32768
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+
+
+@triton.jit
+def _gelu_and_mul_sparse_kernel(
+    x_ptr,
+    output_ptr,
+    d: tl.constexpr,
+    std_multiplier,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < d
+    row_start = row * 2 * d
+
+    gate = tl.load(x_ptr + row_start + offsets, mask=mask, other=0.0).to(tl.float32)
+    # TODO: Investigate delaying this load to reduce spills for large reductions
+    # without regressing smaller shapes.
+    up = tl.load(x_ptr + row_start + d + offsets, mask=mask, other=0.0).to(tl.float32)
+
+    input_dtype = x_ptr.dtype.element_ty
+    mean = tl.sum(gate, axis=0) / d
+    centered = tl.where(mask, gate - mean, 0.0)
+    variance = tl.sum(centered * centered, axis=0) / d
+    # Match the input-dtype rounding at each native PyTorch op boundary.
+    mean = mean.to(input_dtype).to(tl.float32)
+    std = tl.sqrt(variance).to(input_dtype).to(tl.float32)
+    scaled_std = (std * std_multiplier).to(input_dtype).to(tl.float32)
+    cutoff = (mean + scaled_std).to(input_dtype).to(tl.float32)
+    sparse_gate = (gate - cutoff).to(input_dtype).to(tl.float32)
+    sparse_gate = tl.where(sparse_gate < 0.0, 0.0, sparse_gate)
+
+    inner = 0.7978845608028654 * (
+        sparse_gate + 0.044715 * sparse_gate * sparse_gate * sparse_gate
+    )
+    tanh = 2.0 * tl.sigmoid(2.0 * inner) - 1.0
+    activated = 0.5 * sparse_gate * (1.0 + tanh)
+    activated = activated.to(input_dtype).to(tl.float32)
+    tl.store(output_ptr + row * d + offsets, activated * up, mask=mask)
+
+
+def _supports_gelu_and_mul_sparse(
+    x: Tensor, std_multiplier: float, approximate: str = "none"
+) -> bool:
+    del std_multiplier
+    if (
+        x.device.type != "cuda"
+        or x.dtype not in _SUPPORTED_DTYPES
+        or not x.is_contiguous()
+        or x.ndim == 0
+        or approximate != "tanh"
+    ):
+        return False
+    last_dim = x.shape[-1]
+    return (
+        last_dim > 0 and last_dim % 2 == 0 and last_dim // 2 <= _MAX_INTERMEDIATE_SIZE
+    )
+
+
+@triton_op(
+    "vllm::gelu_and_mul_sparse_triton",
+    mutates_args=(),
+)
+def _gelu_and_mul_sparse_triton_op(x: Tensor, std_multiplier: float) -> Tensor:
+    d = x.shape[-1] // 2
+    output = torch.empty((*x.shape[:-1], d), dtype=x.dtype, device=x.device)
+    if output.numel() == 0:
+        return output
+
+    n_rows = x.numel() // x.shape[-1]
+
+    def grid(meta):
+        return (n_rows,)
+
+    block_size = triton.next_power_of_2(d)
+    num_warps = min(max(block_size // 256, 1), 8)
+    wrap_triton(_gelu_and_mul_sparse_kernel)[grid](
+        x,
+        output,
+        d=d,
+        std_multiplier=std_multiplier,
+        BLOCK_SIZE=block_size,
+        num_warps=num_warps,
+    )
+    return output
+
+
+@ir.ops.gelu_and_mul_sparse.register_impl(
+    "triton",
+    supported=HAS_TRITON and current_platform.is_cuda(),
+    supports_args=_supports_gelu_and_mul_sparse,
+)
+def gelu_and_mul_sparse_triton(
+    x: Tensor, std_multiplier: float, approximate: str = "none"
+) -> Tensor:
+    assert approximate == "tanh"
+    return _gelu_and_mul_sparse_triton_op(x, std_multiplier)
diff --git a/vllm/logger.py b/vllm/logger.py
index fde95662f172..05ee5936d45a 100644
--- a/vllm/logger.py
+++ b/vllm/logger.py
@@ -29,7 +29,7 @@
 def _use_color() -> bool:
     if envs.NO_COLOR or envs.VLLM_LOGGING_COLOR == "0":
         return False
-    if envs.VLLM_LOGGING_COLOR == "1":
+    if envs.VLLM_LOGGING_COLOR == "1" or envs.FORCE_COLOR:
         return True
     if envs.VLLM_LOGGING_STREAM == "ext://sys.stdout":  # stdout
         return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()
@@ -90,11 +90,17 @@ def _print_warning_once(logger: Logger, msg: str, *args: Hashable) -> None:
     logger.warning(msg, *args, stacklevel=3)
 
 
+@lru_cache
+def _print_error_once(logger: Logger, msg: str, *args: Hashable) -> None:
+    # Set the stacklevel to 3 to print the original caller's line info
+    logger.error(msg, *args, stacklevel=3)
+
+
 LogScope = Literal["process", "global", "local"]
 
 
 def _should_log_with_scope(scope: LogScope) -> bool:
-    """Decide whether to log based on scope"""
+    """Decide whether to log based on scope."""
     if scope == "global":
         from vllm.distributed.parallel_state import is_global_first_rank
 
@@ -107,17 +113,16 @@ def _should_log_with_scope(scope: LogScope) -> bool:
 
 
 class _VllmLogger(Logger):
-    """
-    Note:
-        This class is just to provide type information.
-        We actually patch the methods directly on the [`logging.Logger`][]
-        instance to avoid conflicting with other libraries such as
-        `intel_extension_for_pytorch.utils._logger`.
+    """Note:
+    This class is just to provide type information.
+    We actually patch the methods directly on the [`logging.Logger`][]
+    instance to avoid conflicting with other libraries such as
+    `intel_extension_for_pytorch.utils._logger`.
+
     """
 
     def debug_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> None:
-        """
-        As [`debug`][logging.Logger.debug], but subsequent calls with
+        """As [`debug`][logging.Logger.debug], but subsequent calls with
         the same message are silently dropped.
         """
         if not _should_log_with_scope(scope):
@@ -125,8 +130,7 @@ def debug_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> No
         _print_debug_once(self, msg, *args)
 
     def info_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> None:
-        """
-        As [`info`][logging.Logger.info], but subsequent calls with
+        """As [`info`][logging.Logger.info], but subsequent calls with
         the same message are silently dropped.
         """
         if not _should_log_with_scope(scope):
@@ -136,20 +140,28 @@ def info_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> Non
     def warning_once(
         self, msg: str, *args: Hashable, scope: LogScope = "local"
     ) -> None:
-        """
-        As [`warning`][logging.Logger.warning], but subsequent calls with
+        """As [`warning`][logging.Logger.warning], but subsequent calls with
         the same message are silently dropped.
         """
         if not _should_log_with_scope(scope):
             return
         _print_warning_once(self, msg, *args)
 
+    def error_once(self, msg: str, *args: Hashable, scope: LogScope = "local") -> None:
+        """As [`error`][logging.Logger.error], but subsequent calls with
+        the same message are silently dropped.
+        """
+        if not _should_log_with_scope(scope):
+            return
+        _print_error_once(self, msg, *args)
+
 
 # Pre-defined methods mapping to avoid repeated dictionary creation
 _METHODS_TO_PATCH = {
     "debug_once": _VllmLogger.debug_once,
     "info_once": _VllmLogger.info_once,
     "warning_once": _VllmLogger.warning_once,
+    "error_once": _VllmLogger.error_once,
 }
 
 
@@ -205,7 +217,6 @@ def init_logger(name: str) -> _VllmLogger:
     """The main purpose of this function is to ensure that loggers are
     retrieved in such a way that we can be sure the root vllm logger has
     already been configured."""
-
     logger = logging.getLogger(name)
 
     for method_name, method in _METHODS_TO_PATCH.items():
@@ -242,8 +253,10 @@ def current_formatter_type(logger: Logger) -> Literal["color", "newline", None]:
 
 # Transformers uses httpx to access the Hugging Face Hub. httpx is quite verbose,
 # so we set its logging level to WARNING when vLLM's logging level is INFO.
+# httpx2 is the successor huggingface_hub switches to in its 2.x releases.
 if envs.VLLM_LOGGING_LEVEL == "INFO":
     logging.getLogger("httpx").setLevel(logging.WARNING)
+    logging.getLogger("httpx2").setLevel(logging.WARNING)
 
 logger = init_logger(__name__)
 
@@ -292,8 +305,7 @@ def _trace_calls(log_path, root_dir, frame, event, arg=None):
 
 
 def enable_trace_function_call(log_file_path: str, root_dir: str | None = None):
-    """
-    Enable tracing of every function call in code under `root_dir`.
+    """Enable tracing of every function call in code under `root_dir`.
     This is useful for debugging hangs or crashes.
     `log_file_path` is the path to the log file.
     `root_dir` is the root directory of the code to trace. If None, it is the
diff --git a/vllm/logging_utils/access_log_filter.py b/vllm/logging_utils/access_log_filter.py
index 5501bd5bc6e8..086449ead845 100644
--- a/vllm/logging_utils/access_log_filter.py
+++ b/vllm/logging_utils/access_log_filter.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Access log filter for uvicorn to exclude specific endpoints from logging.
+"""Access log filter for uvicorn to exclude specific endpoints from logging.
 
 This module provides a logging filter that can be used to suppress access logs
 for specific endpoints (e.g., /health, /metrics) to reduce log noise in
@@ -13,8 +12,7 @@
 
 
 class UvicornAccessLogFilter(logging.Filter):
-    """
-    A logging filter that excludes access logs for specified endpoint paths.
+    """A logging filter that excludes access logs for specified endpoint paths.
 
     This filter is designed to work with uvicorn's access logger. It checks
     the log record's arguments for the request path and filters out records
@@ -31,6 +29,7 @@ class UvicornAccessLogFilter(logging.Filter):
         excluded_paths: A list of URL paths to exclude from logging.
                        Paths are matched exactly.
                        Example: ["/health", "/metrics"]
+
     """
 
     def __init__(self, excluded_paths: list[str] | None = None):
@@ -38,14 +37,14 @@ def __init__(self, excluded_paths: list[str] | None = None):
         self.excluded_paths = set(excluded_paths or [])
 
     def filter(self, record: logging.LogRecord) -> bool:
-        """
-        Determine if the log record should be logged.
+        """Determine if the log record should be logged.
 
         Args:
             record: The log record to evaluate.
 
         Returns:
             True if the record should be logged, False otherwise.
+
         """
         if not self.excluded_paths:
             return True
@@ -72,8 +71,7 @@ def create_uvicorn_log_config(
     excluded_paths: list[str] | None = None,
     log_level: str = "info",
 ) -> dict:
-    """
-    Create a uvicorn logging configuration with access log filtering.
+    """Create a uvicorn logging configuration with access log filtering.
 
     This function generates a logging configuration dictionary that can be
     passed to uvicorn's `log_config` parameter. It sets up the access log
@@ -89,6 +87,7 @@ def create_uvicorn_log_config(
     Example:
         >>> config = create_uvicorn_log_config(["/health", "/metrics"])
         >>> uvicorn.run(app, log_config=config)
+
     """
     config = {
         "version": 1,
diff --git a/vllm/logging_utils/formatter.py b/vllm/logging_utils/formatter.py
index 22301751a666..4a90009266b5 100644
--- a/vllm/logging_utils/formatter.py
+++ b/vllm/logging_utils/formatter.py
@@ -19,8 +19,7 @@ def __init__(self, fmt, datefmt=None, style="%"):
 
     def format(self, record):
         def shrink_path(relpath: Path) -> str:
-            """
-            Shortens a file path for logging display:
+            """Shortens a file path for logging display:
             - Removes leading 'vllm' folder if present.
             - If path starts with 'v1',
             keeps the first two and last two levels,
@@ -36,8 +35,10 @@ def shrink_path(relpath: Path) -> str:
 
             Args:
                 relpath (Path): The relative path to be shortened.
+
             Returns:
                 str: The shortened path string for display.
+
             """
             parts = list(relpath.parts)
             new_parts = []
diff --git a/vllm/logging_utils/log_time.py b/vllm/logging_utils/log_time.py
index 9e94f463711d..c390bc95a3bb 100644
--- a/vllm/logging_utils/log_time.py
+++ b/vllm/logging_utils/log_time.py
@@ -1,16 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Provides a timeslice logging decorator
-"""
+"""Provides a timeslice logging decorator."""
 
 import functools
 import time
 
 
 def logtime(logger, msg=None):
-    """
-    Logs the execution time of the decorated function.
+    """Logs the execution time of the decorated function.
     Always place it beneath other decorators.
     """
 
diff --git a/vllm/logits_process.py b/vllm/logits_process.py
index 1bf97c2535fb..300ec3cacc3f 100644
--- a/vllm/logits_process.py
+++ b/vllm/logits_process.py
@@ -1,12 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 from typing import TypeAlias
 
 import torch
 
-from vllm.tokenizers import TokenizerLike
-
 LogitsProcessor: TypeAlias = (
     Callable[[list[int], torch.Tensor], torch.Tensor]
     | Callable[[list[int], list[int], torch.Tensor], torch.Tensor]
@@ -16,106 +14,3 @@
 for the next token and, optionally, prompt tokens as a
 first argument, and returns a modified tensor of logits
 to sample from."""
-
-
-def get_bad_words_logits_processors(
-    bad_words: list[str], tokenizer: TokenizerLike
-) -> list[LogitsProcessor]:
-    bad_words_ids: list[list[int]] = list()
-
-    for bad_word in bad_words:
-        # To prohibit words both at the beginning
-        # and in the middle of text
-        # (related to add_prefix_space tokenizer parameter)
-        for add_prefix_space in [False, True]:
-            prefix = " " if add_prefix_space else ""
-            prompt = prefix + bad_word.lstrip()
-
-            prompt_token_ids = tokenizer.encode(text=prompt, add_special_tokens=False)
-
-            # If no space at the beginning
-            # or if prefix space produces a new word token
-            if (not add_prefix_space) or (
-                add_prefix_space
-                and prompt_token_ids[0] != bad_words_ids[-1][0]
-                and len(prompt_token_ids) == len(bad_words_ids[-1])
-            ):
-                bad_words_ids.append(prompt_token_ids)
-
-    return [NoBadWordsLogitsProcessor(bad_words_ids=bad_words_ids)]
-
-
-class NoBadWordsLogitsProcessor:
-    _SMALLEST_LOGIT = float("-inf")
-    _NEUTRAL_LOGIT = 0.0
-
-    def __init__(self, bad_words_ids: list[list[int]]):
-        self.bad_words_ids = bad_words_ids
-        self.word_bias: torch.FloatTensor = None
-
-    def __call__(
-        self,
-        past_tokens_ids: Sequence[int],
-        logits: torch.FloatTensor,
-    ) -> torch.Tensor:
-        if self.word_bias is None:
-            self._init_word_bias(logits=logits)
-
-        last_token_bias = torch.zeros_like(logits)
-
-        for bad_word_ids in self.bad_words_ids:
-            if len(bad_word_ids) == 1:  # 1-token words already processed
-                continue
-
-            if len(bad_word_ids) > len(past_tokens_ids) + 1:
-                continue
-
-            prefix_length = len(bad_word_ids) - 1
-            last_token_id = bad_word_ids[-1]
-            actual_prefix = past_tokens_ids[-prefix_length:]
-            expected_prefix = bad_word_ids[:prefix_length]
-
-            assert len(actual_prefix) == len(expected_prefix)
-
-            is_match = tuple(actual_prefix) == tuple(expected_prefix)
-            last_token_bias[last_token_id] += (
-                self._SMALLEST_LOGIT if is_match else self._NEUTRAL_LOGIT
-            )
-
-        logits = logits + self.word_bias + last_token_bias
-
-        return logits
-
-    def _init_word_bias(self, logits: torch.FloatTensor) -> None:
-        # Code based on NoBadWordsLogitsProcessor and SequenceBiasLogitsProcessor  # noqa: E501
-        # from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py
-
-        vocab_size = logits.shape[-1]
-
-        self._check_token_ids_bounds(vocab_size=vocab_size)
-
-        self.word_bias = torch.zeros(
-            (vocab_size,), dtype=torch.float, device=logits.device
-        )
-
-        for bad_word_ids in self.bad_words_ids:
-            if len(bad_word_ids) == 1:
-                bad_word_id = bad_word_ids[-1]
-                self.word_bias[bad_word_id] = self._SMALLEST_LOGIT
-
-    def _check_token_ids_bounds(self, vocab_size: int) -> None:
-        invalid_token_ids = []
-
-        for bad_word_ids in self.bad_words_ids:
-            for token_id in bad_word_ids:
-                if token_id < 0 or token_id >= vocab_size:
-                    invalid_token_ids.append(token_id)
-
-        if len(invalid_token_ids) > 0:
-            raise ValueError(
-                f"The model vocabulary size is {vocab_size},"
-                f" but the following tokens"
-                f" were specified as bad: {invalid_token_ids}."
-                f" All token id values should be integers satisfying:"
-                f" 0 <= token_id < {vocab_size}."
-            )
diff --git a/vllm/logprobs.py b/vllm/logprobs.py
index a82219920675..c4f3340a659a 100644
--- a/vllm/logprobs.py
+++ b/vllm/logprobs.py
@@ -17,6 +17,7 @@ class Logprob:
         logprob: The logprob of chosen token
         rank: The vocab rank of chosen token (>=1)
         decoded_token: The decoded chosen token index
+
     """
 
     logprob: float
@@ -29,8 +30,7 @@ class Logprob:
 
 @dataclass
 class FlatLogprobs(MutableSequence[LogprobsOnePosition | None]):
-    """
-    Flat logprobs of a request into multiple primitive type lists.
+    """Flat logprobs of a request into multiple primitive type lists.
 
     Compared to list[dict[int, Logprob]], this data structure reduced GC
     overhead significantly. As it flattened logprob information for
@@ -61,7 +61,7 @@ class FlatLogprobs(MutableSequence[LogprobsOnePosition | None]):
     decoded_tokens: list[str | None] = field(default_factory=list)
 
     def append(self, logprobs_one_position: LogprobsOnePosition | None) -> None:
-        """Appends the container with logprobs for the next position"""
+        """Appends the container with logprobs for the next position."""
         self.start_indices.append(len(self.logprobs))
         if logprobs_one_position:
             for token_id, logprob in logprobs_one_position.items():
@@ -78,8 +78,7 @@ def append_fast(
         ranks: itertools.chain[int],
         decoded_tokens: Iterable[str | None],
     ) -> None:
-        """
-        Appends logprobs for the next position without creating
+        """Appends logprobs for the next position without creating
         the intermediate logprob dictionary.
         """
         self.start_indices.append(len(self.logprobs))
@@ -93,12 +92,12 @@ def append_fast(
         self.end_indices.append(len(self.logprobs))
 
     def extend(self, logprobs_multi_positions) -> None:
-        """Extends the container with logprobs for the next multiple positions"""
+        """Extends the container with logprobs for the next multiple positions."""
         for logprobs_one_position in logprobs_multi_positions:
             self.append(logprobs_one_position)
 
     def __len__(self) -> int:
-        """Gets number of positions stored in the container"""
+        """Gets number of positions stored in the container."""
         return len(self.start_indices)
 
     @overload
@@ -108,7 +107,7 @@ def __getitem__(self, position: int) -> LogprobsOnePosition: ...
     def __getitem__(self, s: slice, /) -> "FlatLogprobs": ...
 
     def __getitem__(self, index: int | slice):
-        """Extracts logprobs of a given position or slice"""
+        """Extracts logprobs of a given position or slice."""
         if isinstance(index, int):
             return {
                 self.token_ids[i]: Logprob(
@@ -149,8 +148,7 @@ def insert(self, index: int, value: dict[int, Logprob] | None) -> None:
         raise TypeError("Cannot insert logprobs to FlatLogprobs")
 
     def __iter__(self) -> Iterator[LogprobsOnePosition]:
-        """
-        Iterates the container and yields LogprobsOnePosition for
+        """Iterates the container and yields LogprobsOnePosition for
         each position.
         """
         for i in range(0, len(self.start_indices)):
@@ -165,7 +163,7 @@ def __iter__(self) -> Iterator[LogprobsOnePosition]:
 
 
 def create_prompt_logprobs(flat_logprobs: bool) -> PromptLogprobs:
-    """Creates a container to store prompt logprobs for a request"""
+    """Creates a container to store prompt logprobs for a request."""
     logprobs: PromptLogprobs = FlatLogprobs() if flat_logprobs else []
     # NOTE: logprob of first prompt token is None.
     logprobs.append(None)
@@ -173,7 +171,7 @@ def create_prompt_logprobs(flat_logprobs: bool) -> PromptLogprobs:
 
 
 def create_sample_logprobs(flat_logprobs: bool) -> SampleLogprobs:
-    """Creates a container to store decode logprobs for a request"""
+    """Creates a container to store decode logprobs for a request."""
     return FlatLogprobs() if flat_logprobs else []
 
 
@@ -185,7 +183,7 @@ def append_logprobs_for_next_position(
     rank: int,
     num_logprobs: int,
 ) -> None:
-    """Appends logprobs for the next position"""
+    """Appends logprobs for the next position."""
     if num_logprobs == -1:
         num_logprobs = len(logprobs)
     # We do not need a special case for the sampled token
diff --git a/vllm/lora/layers/__init__.py b/vllm/lora/layers/__init__.py
index e0c5f202bd01..2f88640e5fd1 100644
--- a/vllm/lora/layers/__init__.py
+++ b/vllm/lora/layers/__init__.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from vllm.lora.layers.base import BaseLayerWithLoRA
+from vllm.lora.layers.classifier import ClassificationHeadWithLoRA
 from vllm.lora.layers.column_parallel_linear import (
     ColumnParallelLinearWithLoRA,
     ColumnParallelLinearWithShardedLoRA,
@@ -42,4 +43,5 @@
     "LoRAMappingType",
     "FusedMoEWithLoRA",
     "FusedMoE3DWithLoRA",
+    "ClassificationHeadWithLoRA",
 ]
diff --git a/vllm/lora/layers/base.py b/vllm/lora/layers/base.py
index 5fba2998c85a..70821be3d163 100644
--- a/vllm/lora/layers/base.py
+++ b/vllm/lora/layers/base.py
@@ -6,7 +6,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.lora import LoRAConfig
 
@@ -73,7 +73,7 @@ def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> None:
         """Initializes lora matrices."""
         ...
@@ -103,7 +103,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         """Returns True if the layer can be replaced by this LoRA layer."""
         raise NotImplementedError
diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py
index 5c8e829b2997..12888cf3db3d 100644
--- a/vllm/lora/layers/base_linear.py
+++ b/vllm/lora/layers/base_linear.py
@@ -3,7 +3,7 @@
 
 
 import torch
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm import envs
 from vllm.config import get_current_vllm_config
@@ -68,6 +68,10 @@ def lora_linear_async_fake(
 
 
 class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
+    # The adapter branch consumes the original activation, independently of
+    # whether the wrapped base layer can consume a pre-quantized activation.
+    requires_unquantized_input = True
+
     def __init__(self, base_layer: LinearBase):
         super().__init__()
 
@@ -101,7 +105,7 @@ def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> None:
         self.lora_config = lora_config
         if isinstance(self.base_layer, ReplicatedLinear):
@@ -240,8 +244,7 @@ def _apply_lora_to_output(
     def _apply_async_impl(
         self, x: torch.Tensor, bias: torch.Tensor | None = None
     ) -> torch.Tensor:
-        """
-        Forward pass with base linear and LoRA on separate CUDA streams
+        """Forward pass with base linear and LoRA on separate CUDA streams
         for overlap, using maybe_execute_in_parallel.
         Base layer runs on default stream; LoRA runs on aux stream.
         """
diff --git a/vllm/lora/layers/classifier.py b/vllm/lora/layers/classifier.py
new file mode 100644
index 000000000000..58160d1931fd
--- /dev/null
+++ b/vllm/lora/layers/classifier.py
@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+from transformers import PreTrainedConfig
+
+from vllm.config.lora import LoRAConfig
+from vllm.platforms import current_platform
+
+from .replicated_linear import ReplicatedLinearWithLoRA
+
+
+class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA):
+    def create_lora_weights(
+        self,
+        max_loras: int,
+        lora_config: LoRAConfig,
+        model_config: PreTrainedConfig | None = None,
+    ) -> None:
+        # Preserve ordinary LoRA A/B support for classification heads.
+        super().create_lora_weights(max_loras, lora_config, model_config)
+        self.full_weight_stacked = torch.zeros(
+            max_loras,
+            1,
+            self.output_size,
+            self.input_size,
+            dtype=lora_config.lora_dtype,
+            device=self.device,
+        )
+        self.full_bias_stacked = torch.zeros(
+            max_loras,
+            self.output_size,
+            dtype=self.base_layer.params_dtype,
+            device=self.device,
+        )
+        self.full_module_enabled = torch.zeros(
+            max_loras, dtype=torch.bool, device=self.device
+        )
+
+    def reset_module_to_save(self, index: int) -> None:
+        self.full_weight_stacked[index].zero_()
+        self.full_bias_stacked[index].zero_()
+        self.full_module_enabled[index] = False
+
+    def set_module_to_save(
+        self,
+        index: int,
+        weight: torch.Tensor,
+        bias: torch.Tensor | None,
+    ) -> None:
+        self.reset_module_to_save(index)
+        self.full_weight_stacked[index, 0].copy_(weight, non_blocking=True)
+        if bias is not None:
+            self.full_bias_stacked[index].copy_(bias, non_blocking=True)
+        self.full_module_enabled[index] = True
+
+    def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        # TODO base_result maybe don't need compute truly.
+        base_result = super().forward(input_)
+        output = base_result[0] if isinstance(base_result, tuple) else base_result
+
+        full_output = self.punica_wrapper.apply_lora_full_linear(
+            output,
+            input_.to(self.full_weight_stacked.dtype),
+            self.full_weight_stacked,
+            self.full_bias_stacked,
+            self.full_module_enabled,
+        )
+
+        if full_output is not None:
+            output = full_output
+        return output
+
+    def _apply_lora_to_output(
+        self, x: torch.Tensor, output: torch.Tensor
+    ) -> torch.Tensor:
+        original_shape = output.shape if output.ndim == 3 else None
+        if x.ndim == 3 and output.ndim == 3:
+            output = output.flatten(0, 1)
+            x = x.flatten(0, 1)
+
+        # Classification outputs are request-level, so use prompt-based LoRA routing.
+        lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_logits(
+            output,
+            x.to(self.lora_a_stacked[0].dtype),
+            self.lora_a_stacked[0],
+            self.lora_b_stacked[0],
+            1.0,
+        )
+        if not current_platform.can_update_inplace():
+            output = lora_output
+
+        if original_shape is not None:
+            output = output.reshape(original_shape)
+
+        return output
diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py
index ea776ec2ea4f..06be86d388cb 100644
--- a/vllm/lora/layers/column_parallel_linear.py
+++ b/vllm/lora/layers/column_parallel_linear.py
@@ -4,7 +4,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.lora import LoRAConfig
 from vllm.distributed import tensor_model_parallel_all_gather
@@ -83,8 +83,7 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):
 
 
 class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
-    """
-    LoRA on top of ColumnParallelLinear layer.
+    """LoRA on top of ColumnParallelLinear layer.
     LoRA B is sliced for tensor parallelism.
     There are two types for the `base_layer`:
     1. ColumnParallelLinear, e.g.`dense_h_to_4h` in `FalconForCausalLM`.
@@ -140,6 +139,7 @@ def forward(
         Returns:
             - output
             - bias
+
         """
         bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
 
@@ -164,7 +164,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         if type(source_layer) is maybe_get_oot_by_class(ColumnParallelLinear):
             return True
@@ -208,10 +208,9 @@ def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> None:
-        """
-        The main reason for overriding this function is to enhance  code
+        """The main reason for overriding this function is to enhance  code
         maintainability.
         """
         self.lora_config = lora_config
@@ -268,8 +267,7 @@ def expand_packed_lora(
         lora_a: list[torch.Tensor | None],
         lora_b: list[torch.Tensor | None],
     ) -> tuple[list[torch.Tensor | None], list[torch.Tensor | None]]:
-        """
-        Expand packed adapter groups when they don't match n_slices.
+        """Expand packed adapter groups when they don't match n_slices.
         E.g. in_proj_qkv (covers Q+K+V) + in_proj_z.
 
         A None group member means that member was not adapted; the slice(s)
@@ -368,7 +366,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
         decorate: bool = True,
     ) -> bool:
         merged_cls = maybe_get_oot_by_class(MergedColumnParallelLinear)
@@ -388,8 +386,7 @@ def can_replace_layer(
 
 
 class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
-    """
-    ColumnParallelLinear layer that is specifically designed for
+    """ColumnParallelLinear layer that is specifically designed for
     qkv_proj. Certain models, such as chatglm3 and baichuan-7b,
     only contains a single LoRA within their qkv_proj layer.
 
@@ -445,7 +442,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         return (
             type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear)
@@ -491,10 +488,9 @@ def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> None:
-        """
-        The main reason for overloading this function is to handle inconsistent
+        """The main reason for overloading this function is to handle inconsistent
         weight dimensions in qkv lora.
         """
         super().create_lora_weights(max_loras, lora_config, model_config)
@@ -506,7 +502,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
         decorate: bool = True,
     ) -> bool:
         return (
@@ -521,8 +517,7 @@ def can_replace_layer(
 
 
 class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):
-    """
-    Differs from ColumnParallelLinearWithLoRA by slicing LoRA A also.
+    """Differs from ColumnParallelLinearWithLoRA by slicing LoRA A also.
 
     Based on S-LoRA, slicing happens along the rank dim.
     """
@@ -548,7 +543,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
         decorate: bool = True,
     ) -> bool:
         # specifying kwargs so they can be easily accessed in decorator
@@ -562,8 +557,7 @@ def can_replace_layer(
 
 
 class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLoRA):
-    """
-    Differs from MergedColumnParallelLinearWithLoRA by slicing the
+    """Differs from MergedColumnParallelLinearWithLoRA by slicing the
     LoRA A's also.
 
     Based on S-LoRA, slicing happens along the rank dim.
@@ -591,7 +585,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
         decorate: bool = True,
     ) -> bool:
         # specifying kwargs so they can be easily accessed in decorator
@@ -605,8 +599,7 @@ def can_replace_layer(
 
 
 class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):
-    """
-    Differs from QKVParallelLinearWithLoRA by slicing the
+    """Differs from QKVParallelLinearWithLoRA by slicing the
     LoRA A's also.
 
     Based on S-LoRA, slicing happens along the rank dim.
@@ -628,7 +621,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         # specifying kwargs so they can be easily accessed in decorator
         return super().can_replace_layer(
@@ -641,8 +634,7 @@ def can_replace_layer(
 
 
 class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):
-    """
-    Differs from MergedQKVParallelLinearWithLoRA by slicing the
+    """Differs from MergedQKVParallelLinearWithLoRA by slicing the
     LoRA A's also.
 
     Based on S-LoRA, slicing happens along the rank dim.
@@ -677,7 +669,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
         decorate: bool = True,
     ) -> bool:
         # specifying kwargs so they can be easily accessed in decorator
@@ -706,7 +698,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
         decorate: bool = True,
     ) -> bool:
         # Support MergedColumnParallelLinear with 3 or more slices
diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py
index e47bc79bf469..c58c18e71ff7 100644
--- a/vllm/lora/layers/fused_moe.py
+++ b/vllm/lora/layers/fused_moe.py
@@ -3,7 +3,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm import envs
 from vllm.config.lora import LoRAConfig
@@ -245,10 +245,9 @@ def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> None:
         """Initializes lora matrices."""
-
         self._verify_ep_fs(lora_config)
         self.max_loras = lora_config.max_loras
         self.fully_sharded = lora_config.fully_sharded_loras
@@ -294,9 +293,7 @@ def create_lora_weights(
                     )
 
     def _slice_w13_a(self, w13_lora_a: torch.Tensor) -> torch.Tensor:
-        """
-        Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA
-        """
+        """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA."""
         if self.tp_size == 1 or not self.fully_sharded:
             return w13_lora_a
 
@@ -321,9 +318,7 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor):
         return w13_lora_b[:, start_idx:end_idx, :]
 
     def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor:
-        """
-        Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA
-        """
+        """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA."""
         if self.tp_size == 1:
             return w2_lora_a
         # w2_lora_a shape (num_experts,rank,input_size)
@@ -334,9 +329,7 @@ def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor:
         return w2_lora_a[:, :, start_idx:end_idx]
 
     def _slice_w2_b(self, w2_lora_b: torch.Tensor) -> torch.Tensor:
-        """
-        Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA
-        """
+        """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA."""
         if self.tp_size == 1 or not self.fully_sharded:
             return w2_lora_b
         # Based on S-LoRA, we slice W2 B along the hidden_size dim.
@@ -472,10 +465,9 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         """Returns True if the layer can be replaced by this LoRA layer."""
-
         # source_layer is MoERunner
         moe_cls = maybe_get_oot_by_class(MoERunner)
         return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 2
@@ -519,10 +511,9 @@ def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> None:
         """Initializes lora matrices."""
-
         if model_config is None:
             raise ValueError("model_config must be provided for MoE LoRA.")
         architectures = model_config.architectures
@@ -609,30 +600,22 @@ def set_lora(
 
     @property
     def w13_input_size(self):
-        """
-        Full size
-        """
+        """Full size."""
         return self.w13_lora_a_stacked[0].shape[-1]
 
     @property
     def w13_output_size(self):
-        """
-        Full size
-        """
+        """Full size."""
         return self.w13_lora_b_stacked[0].shape[-2] * self.tp_size
 
     @property
     def w2_input_size(self):
-        """
-        Full size
-        """
+        """Full size."""
         return self.w2_lora_a_stacked[0].shape[-1] * self.tp_size
 
     @property
     def w2_output_size(self):
-        """
-        Full size
-        """
+        """Full size."""
         return self.hidden_size
 
     @classmethod
@@ -641,7 +624,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         """Returns True if the layer can be replaced by this LoRA layer."""
         # source_layer is MoERunner
diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py
index 0199437ba22d..6259fd781759 100644
--- a/vllm/lora/layers/logits_processor.py
+++ b/vllm/lora/layers/logits_processor.py
@@ -3,7 +3,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.lora import LoRAConfig
 from vllm.distributed import (
@@ -18,8 +18,7 @@
 
 
 class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
-    """
-    LoRA wrapper for LogitsProcessor, with extra logic to handle the
+    """LoRA wrapper for LogitsProcessor, with extra logic to handle the
     application of the LoRA adapter.
 
     Args:
@@ -30,6 +29,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
         sharded_to_full_mapping: index mapping from sharded vocab to full vocab
             received from base_layer.get_sharded_to_full_mapping(). If None,
             no reindexing will be done.
+
     """
 
     def __init__(
@@ -85,7 +85,7 @@ def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> None:
         # TODO: Verify if this condition can be further relaxed
         if self.base_layer.vocab_size > 258048:
@@ -221,7 +221,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         # Special handling for the LogitsProcessor.
         return False
diff --git a/vllm/lora/layers/replicated_linear.py b/vllm/lora/layers/replicated_linear.py
index 53ae26be4c36..b4d61caca2c6 100644
--- a/vllm/lora/layers/replicated_linear.py
+++ b/vllm/lora/layers/replicated_linear.py
@@ -4,7 +4,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.lora import LoRAConfig
 from vllm.model_executor.custom_op import maybe_get_oot_by_class
@@ -25,7 +25,7 @@ def __init__(self, base_layer: ReplicatedLinear) -> None:
     def forward(
         self, input_: torch.Tensor
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
-        """Forward of ReplicatedLinearWithLoRA
+        """Forward of ReplicatedLinearWithLoRA.
 
         Args:
             input_: Tensor whose last dimension is `input_size`.
@@ -33,6 +33,7 @@ def forward(
         Returns:
             - output
             - bias
+
         """
         bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
 
@@ -60,7 +61,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         return isinstance(source_layer, maybe_get_oot_by_class(ReplicatedLinear))
 
diff --git a/vllm/lora/layers/row_parallel_linear.py b/vllm/lora/layers/row_parallel_linear.py
index 3d3c7f1d279c..48cbd31c1b7e 100644
--- a/vllm/lora/layers/row_parallel_linear.py
+++ b/vllm/lora/layers/row_parallel_linear.py
@@ -4,7 +4,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.lora import LoRAConfig
 from vllm.distributed import (
@@ -42,7 +42,7 @@ def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
     def forward(
         self, input_: torch.Tensor
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
-        """Forward of RowParallelLinear
+        """Forward of RowParallelLinear.
 
         Args:
             input_: tensor whose last dimension is `input_size`. If
@@ -52,6 +52,7 @@ def forward(
         Returns:
             - output
             - bias
+
         """
         # set up backprop all-reduce.
         if self.base_layer.input_is_parallel:
@@ -88,7 +89,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         return type(source_layer) is maybe_get_oot_by_class(RowParallelLinear)
 
@@ -99,8 +100,7 @@ def can_replace_layer(
 
 
 class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
-    """
-    Differs from RowParallelLinearWithLoRA by slicing the
+    """Differs from RowParallelLinearWithLoRA by slicing the
     LoRA B's also.
 
     Based on S-LoRA, slicing happens along the output dim.
@@ -165,7 +165,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         # specifying kwargs so they can be easily accessed in decorator
         return super().can_replace_layer(
diff --git a/vllm/lora/layers/utils.py b/vllm/lora/layers/utils.py
index b4cde2564b39..e725da972349 100644
--- a/vllm/lora/layers/utils.py
+++ b/vllm/lora/layers/utils.py
@@ -74,8 +74,7 @@ def _get_lora_device(base_layer: nn.Module) -> torch.device:
 
 
 def _not_fully_sharded_can_replace(can_replace):
-    """
-    decorator which adds the condition of not using fully sharded loras
+    """Decorator which adds the condition of not using fully sharded loras
     intended to wrap can_replace_layer()
     """
 
@@ -88,8 +87,7 @@ def dec(*args, **kwargs):
 
 
 def _fully_sharded_can_replace(can_replace):
-    """
-    decorator which adds the condition of fully sharded loras
+    """Decorator which adds the condition of fully sharded loras
     intended to wrap can_replace_layer()
     """
 
diff --git a/vllm/lora/layers/vocab_parallel_embedding.py b/vllm/lora/layers/vocab_parallel_embedding.py
index c6f294efcde3..b55f33e149b8 100644
--- a/vllm/lora/layers/vocab_parallel_embedding.py
+++ b/vllm/lora/layers/vocab_parallel_embedding.py
@@ -4,7 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config.lora import LoRAConfig
 from vllm.model_executor.custom_op import maybe_get_oot_by_class
@@ -26,7 +26,7 @@ def create_lora_weights(
         self,
         max_loras: int,
         lora_config: LoRAConfig,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> None:
         self.lora_a_stacked = torch.zeros(
             (
@@ -113,7 +113,7 @@ def can_replace_layer(
         source_layer: nn.Module,
         lora_config: LoRAConfig,
         packed_modules_list: list,
-        model_config: PretrainedConfig | None = None,
+        model_config: PreTrainedConfig | None = None,
     ) -> bool:
         # Accept both the in-tree and any OOT class, because layers built on top of
         # `VocabParallelEmbedding` (e.g. by the Transformers modelling backend) use the
diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py
index f094f11f09e8..593f85993604 100644
--- a/vllm/lora/lora_model.py
+++ b/vllm/lora/lora_model.py
@@ -8,7 +8,7 @@
 import torch
 
 from vllm.logger import init_logger
-from vllm.lora.lora_weights import LoRALayerWeights
+from vllm.lora.lora_weights import LoRAFullModuleWeights, LoRALayerWeights
 from vllm.lora.peft_helper import PEFTHelper
 from vllm.lora.utils import (
     get_lora_id,
@@ -39,9 +39,7 @@ class MoEEPLoadSpec:
 
 
 def _is_remote_expert_key(raw_name: str, spec: "MoEEPLoadSpec") -> bool:
-    """
-    Decide whether a checkpoint key belongs to a non-local expert.
-    """
+    """Decide whether a checkpoint key belongs to a non-local expert."""
     pos = raw_name.find(_EXPERTS_SEPARATOR)
     if pos < 0:
         return False
@@ -66,16 +64,17 @@ def __init__(
         rank: int,
         loras: dict[str, LoRALayerWeights],
         is_3d_lora_weight: bool = False,
+        *,
+        modules_to_save: dict[str, LoRAFullModuleWeights] | None = None,
     ) -> None:
-        """
-        Args:
-            lora_model_id: The integer id for the lora model.
-            rank: lora rank.
-            loras: module name -> weights for lora-replaced layers.
-            is_3d_lora_weight: Whether the on-disk MoE adapter is in the 3D
-                fused (gate_up_proj / down_proj) layout. Propagated from the
-                originating LoRARequest. Only consulted by the LoRA model
-                manager when enable_mixed_moe_lora_format is on.
+        """Args:
+        lora_model_id: The integer id for the lora model.
+        rank: lora rank.
+        loras: module name -> weights for lora-replaced layers.
+        is_3d_lora_weight: Whether the on-disk MoE adapter is in the 3D
+            fused (gate_up_proj / down_proj) layout. Propagated from the
+            originating LoRARequest. Only consulted by the LoRA model
+            manager when enable_mixed_moe_lora_format is on.
 
         """
         self.id = lora_model_id
@@ -85,6 +84,7 @@ def __init__(
         )
         self.rank = rank
         self.loras: dict[str, LoRALayerWeights] = loras
+        self.modules_to_save = modules_to_save or {}
         self.is_3d_lora_weight = is_3d_lora_weight
 
     def clone(self, lora_model_id: int) -> "LoRAModel":
@@ -96,18 +96,22 @@ def clone(self, lora_model_id: int) -> "LoRAModel":
             rank=self.rank,
             loras=self.loras.copy(),
             is_3d_lora_weight=self.is_3d_lora_weight,
+            modules_to_save=self.modules_to_save.copy(),
         )
 
     def get_lora(self, module_name: str) -> LoRALayerWeights | None:
-        """Get LoRA for a given module by name"""
+        """Get LoRA for a given module by name."""
         return self.loras.get(module_name, None)
 
+    def get_module_to_save(self, module_name: str) -> LoRAFullModuleWeights | None:
+        return self.modules_to_save.get(module_name)
+
     def check_lora_name(self, lora_name: str) -> bool:
         return lora_name in self.loras
 
     @staticmethod
     def _should_skip_module(module_name: str, skip_prefixes: list[str]) -> bool:
-        """Check if a module should be skipped based on skip prefixes"""
+        """Check if a module should be skipped based on skip prefixes."""
         for prefix in skip_prefixes:
             if f".{prefix}" in module_name or module_name.startswith(prefix):
                 return True
@@ -127,6 +131,9 @@ def from_lora_tensors(
     ) -> "LoRAModel":
         """Create a LoRAModel from a dictionary of tensors."""
         pin_memory = str(device) == "cpu" and PIN_MEMORY
+
+        modules_to_save_names = peft_helper.modules_to_save or []
+        full_parameters: dict[str, dict[str, torch.Tensor]] = {}
         loras: dict[str, LoRALayerWeights] = {}
         for tensor_name, tensor in tensors.items():
             if is_base_embedding_weights(tensor_name):
@@ -137,6 +144,12 @@ def from_lora_tensors(
             module_name, is_lora_a = parse_fine_tuned_lora_name(
                 tensor_name, weights_mapper
             )
+            if module_name in modules_to_save_names:
+                full_parameters.setdefault(module_name, {})[
+                    tensor_name.split(".")[-1]
+                ] = tensor.to(device=device)
+                continue
+
             if module_name not in loras:
                 loras[module_name] = LoRALayerWeights.from_config(
                     module_name, peft_helper
@@ -161,7 +174,23 @@ def from_lora_tensors(
                 if pin_memory:
                     loras[module_name].lora_b = loras[module_name].lora_b.pin_memory()
 
-        return cls(lora_model_id, peft_helper.r, loras)
+        modules_to_save = {}
+        for module_name, parameters in full_parameters.items():
+            weight = parameters.get("weight")
+            if weight is None:
+                raise ValueError(f"Full module {module_name!r} is missing weight.")
+            modules_to_save[module_name] = LoRAFullModuleWeights(
+                module_name=module_name,
+                weight=weight,
+                bias=parameters.get("bias"),
+            )
+        if len(modules_to_save) > 1:
+            raise ValueError(
+                "Only one full classification module is supported per "
+                f"adapter, received {sorted(modules_to_save)}."
+            )
+
+        return cls(lora_model_id, peft_helper.r, loras, modules_to_save=modules_to_save)
 
     @classmethod
     def from_local_checkpoint(
@@ -190,6 +219,12 @@ def from_local_checkpoint(
                 a global counter.
             device: Device where the lora model is loaded.
             dtype: dtype of the lora model weights.
+            model_vocab_size: Vocab size of the base model, used to size the
+                embedding deltas.
+            weights_mapper: Optional mapper rewriting checkpoint weight names
+                to vLLM names.
+            tensorizer_config_dict: Optional tensorizer config used to load
+                the checkpoint via tensorizer instead of from disk.
             skip_prefixes: List of module name prefixes to skip during loading.
                 Models can define this to skip modules not used in inference
                 (e.g., MTP layers). Format: ["mtp."]
@@ -201,6 +236,7 @@ def from_local_checkpoint(
 
         Returns:
             Loaded LoRA Model.
+
         """
         lora_tensor_path = os.path.join(lora_dir, "adapter_model.safetensors")
         lora_bin_file_path = os.path.join(lora_dir, "adapter_model.bin")
@@ -223,6 +259,7 @@ def check_unexpected_modules(modules: dict):
                 ):
                     continue
                 module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper)
+                base_name = module_name.rsplit(".", 1)[-1]
                 # Case for expert lora weights
                 if ".experts" in module_name:
                     expert_idx = module_name.find(".experts")
@@ -230,7 +267,9 @@ def check_unexpected_modules(modules: dict):
                     if expert_suffix not in expected_lora_modules:
                         unexpected_modules.append(module_name)
 
-                elif module_name.rsplit(".", 1)[-1] not in expected_lora_modules:
+                elif base_name not in expected_lora_modules and base_name not in (
+                    peft_helper.modules_to_save or ()
+                ):
                     unexpected_modules.append(module_name)
 
             if unexpected_modules:
diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py
index 38a5875525da..13a692a62353 100644
--- a/vllm/lora/lora_weights.py
+++ b/vllm/lora/lora_weights.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Sequence as GenericSequence
+from dataclasses import dataclass
 
 import torch
 import torch.types
@@ -162,7 +163,6 @@ def pack_moe(
 
         If LoRA is None, it signifies that the submodule does not have a LoRA.
         """
-
         first_lora = next(lora for lora in loras if lora is not None)
         assert first_lora is not None
         rank = first_lora.rank
@@ -201,8 +201,9 @@ def pack_moe(
         w1_lora_b = torch.stack(w1_lora_b_lst, dim=0)  # (num_experts,output_size,rank)
         w2_lora_b = torch.stack(w2_lora_b_lst, dim=0)
 
-        # All w1, w2, w3 have the same scaling factor.
-        scaling = lora_alpha / rank
+        # All w1, w2, w3 have the same scaling factor. Use the per-adapter
+        # scaling (e.g. alpha/sqrt(r) for rsLoRA) instead of alpha/rank.
+        scaling = first_lora.scaling
         last_scaling = scaling
 
         if is_non_gated_moe:
@@ -250,7 +251,7 @@ def pack_moe_stacked(
         assert w1_lora is not None and w2_lora is not None and w3_lora is not None
         rank = w1_lora.rank
         lora_alpha = w1_lora.lora_alpha
-        scaling = lora_alpha / rank
+        scaling = w1_lora.scaling
         return cls(
             module_name,
             rank,
@@ -280,3 +281,17 @@ def output_dim(self) -> int:
     @property
     def is_packed(self) -> bool:
         return True
+
+
+@dataclass
+class LoRAFullModuleWeights:
+    """LoRA weights for classification layers."""
+
+    module_name: str
+    weight: torch.Tensor
+    bias: torch.Tensor | None = None
+
+    def pin_memory(self) -> None:
+        self.weight = self.weight.pin_memory()
+        if self.bias is not None:
+            self.bias = self.bias.pin_memory()
diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py
index 7791e1b762e1..b47c22d6aab9 100644
--- a/vllm/lora/model_manager.py
+++ b/vllm/lora/model_manager.py
@@ -13,16 +13,22 @@
 from vllm.logger import init_logger
 from vllm.lora.layers import (
     BaseLayerWithLoRA,
+    ClassificationHeadWithLoRA,
     FusedMoE3DWithLoRA,
     FusedMoEWithLoRA,
     LoRAMapping,
     LoRAMappingType,
 )
 from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec
-from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
+from vllm.lora.lora_weights import (
+    LoRAFullModuleWeights,
+    LoRALayerWeights,
+    PackedLoRALayerWeights,
+)
 from vllm.lora.punica_wrapper import PunicaWrapperBase, get_punica_wrapper
 from vllm.lora.utils import (
     from_layer,
+    from_layer_classification,
     from_layer_logits_processor,
     get_supported_lora_modules,
     is_in_target_modules,
@@ -91,6 +97,9 @@ def __init__(
                 in a single batch.
             vocab_size: the vocab size of the model.
             lora_config: the LoRA configuration.
+            device: the device the LoRA tensors are placed on.
+            vllm_config: the vLLM config for this engine.
+
         """
         self.model: SupportsLoRAModel = model
         self.supported_lora_modules = get_supported_lora_modules(self.model)
@@ -114,6 +123,13 @@ def __init__(
         self.vocab_size = vocab_size
 
         self.is_pooling_model = is_pooling_model(self.model)
+        self._pooling_task = (
+            vllm_config.model_config.get_pooling_task(
+                tuple(self.model.pooler.get_supported_tasks())
+            )
+            if self.is_pooling_model
+            else None
+        )
         self.packed_modules: dict[str, list[str]] = {}
         self.modules: dict[str, BaseLayerWithLoRA] = {}
         self._last_mapping: LoRAMapping | None = None
@@ -121,6 +137,13 @@ def __init__(
         is_moe = is_moe_model(self.model)
         self._is_moe = is_moe
 
+        self.supported_modules_to_save = (
+            {"score", "classifier"}
+            if self.is_pooling_model
+            and vllm_config.model_config.score_type == "cross-encoder"
+            else set()
+        )
+
         # When the engine is started with enable_mixed_moe_lora_format=True
         # we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of
         # the model's 3D flag, so 2D and 3D adapters can coexist.
@@ -147,6 +170,20 @@ def __init__(
         self._init_punica_wrapper(max_num_batched_tokens, vllm_config)
         self._create_lora_modules()
 
+        self._classification_head: tuple[str, ClassificationHeadWithLoRA] | None = None
+        if self.is_pooling_model:
+            classification_heads = [
+                (module_name, module)
+                for module_name, module in self.modules.items()
+                if isinstance(module, ClassificationHeadWithLoRA)
+            ]
+            if classification_heads:
+                assert len(classification_heads) == 1, (
+                    "Expected 1 classification head, but found "
+                    f"{len(classification_heads)}."
+                )
+                self._classification_head = classification_heads[0]
+
         self.moe_ep_load_spec: MoEEPLoadSpec | None = self._build_moe_ep_load_spec()
 
         self.model.lora_manager = self
@@ -338,6 +375,15 @@ def activate_adapter(
         self.lora_index_to_id[index] = lora_model.id
         num_lora_weights_applied = 0
         for module_name, module in self.modules.items():
+            if isinstance(module, ClassificationHeadWithLoRA):
+                full_module = self._get_module_to_save_weights(lora_model, module_name)
+                if full_module is None:
+                    module.reset_module_to_save(index)
+                else:
+                    module.set_module_to_save(
+                        index, full_module.weight, full_module.bias
+                    )
+
             module_lora = self._get_lora_layer_weights(lora_model, module_name)
             if not module_lora:
                 module.reset_lora(index)
@@ -427,6 +473,8 @@ def _parent_module(module_name: str) -> str:
             if isinstance(module, PPMissingLayer):
                 continue
 
+            is_classifier_head = module_name in self.supported_modules_to_save
+
             if self.lora_config.target_modules is None:
                 if not is_supported_lora_module(
                     module_name,
@@ -434,7 +482,8 @@ def _parent_module(module_name: str) -> str:
                     self.supported_lora_modules,
                 ):
                     continue
-            elif not self._match_target_modules(module_name):
+            elif not self._match_target_modules(module_name) and not is_classifier_head:
+                # Full classification heads bypass the A/B LoRA target filter.
                 continue
 
             punica_wrapper = self._get_punica_wrapper(module_name)
@@ -486,17 +535,25 @@ def _parent_module(module_name: str) -> str:
                 # LoRA weights of w1 and w3 have already been fused on disk.
 
                 packed_moduled_lst = ["w13"] if self._is_3d_moe_model else ["w1", "w3"]
-            new_module = replace_submodule(
-                self.model,
-                module_name,
-                from_layer(
+            if is_classifier_head:
+                new_module = from_layer_classification(
+                    module,
+                    self.lora_slots,
+                    self.lora_config,
+                    self.model.config,
+                )
+            else:
+                new_module = from_layer(
                     module,
                     self.lora_slots,
                     self.lora_config,
                     packed_moduled_lst,
                     self.model.config,
-                ),
-            )
+                )
+            new_module = replace_submodule(self.model, module_name, new_module)
+
+            if is_classifier_head:
+                self.model.pooler.replace_classifier(new_module)
             if isinstance(new_module, BaseLayerWithLoRA):
                 wrapped_by_id[id(module)] = new_module
                 wrapped_by_id[id(new_module)] = new_module
@@ -731,6 +788,7 @@ def _match_target_modules(self, module_name: str) -> bool:
 
         Returns:
             True if the module passes the filter, False otherwise.
+
         """
         return is_in_target_modules(
             module_name,
@@ -739,13 +797,17 @@ def _match_target_modules(self, module_name: str) -> bool:
         )
 
     def _get_punica_wrapper(self, module_name: str) -> PunicaWrapperBase | None:
-        """
-        Determine whether this module supports LoRA and which wrapper to use.
-        """
+        """Determine whether this module supports LoRA and which wrapper to use."""
         # For language model (early return)
         if not self.supports_mm:
             return self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY]
 
+        if module_name in self.supported_modules_to_save:
+            lm_wrapper = self.punica_wrapper_mapping.get(
+                self.mm_mapping.language_model[0]
+            )
+            return lm_wrapper
+
         # For multimodal model
         # NOTE Sort by prefix length (descending) to match the longest prefix first
         # e.g., 'visual.merger' should match 'visual.merger' instead of 'visual.'
@@ -871,6 +933,10 @@ def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None:
                     lora.lora_a = lora.lora_a.pin_memory()
                     lora.lora_b = lora.lora_b.pin_memory()
 
+            if lora_model.modules_to_save:
+                for full_module in lora_model.modules_to_save.values():
+                    full_module.pin_memory()
+
     def _stack_moe_lora_weights(
         self, lora_model: LoRAModel, module: FusedMoE3DWithLoRA, module_name: str
     ):
@@ -1126,9 +1192,7 @@ def _restrict_to_local_experts(
         return new_module_names[start:end]
 
     def _build_moe_ep_load_spec(self) -> MoEEPLoadSpec | None:
-        """
-        Per-rank slicing metadata for 2D RoutedEXperts LoRA modules.
-        """
+        """Per-rank slicing metadata for 2D RoutedEXperts LoRA modules."""
         if not self._use_ep or not self._is_moe:
             return None
         module = next(
@@ -1164,6 +1228,82 @@ def _get_lora_layer_weights(
                 )
         return lora_model.get_lora(org_module_name)
 
+    def _get_module_to_save_weights(
+        self, lora_model: LoRAModel, module_name: str
+    ) -> LoRAFullModuleWeights | None:
+        weights = lora_model.get_module_to_save(module_name)
+        if weights is not None:
+            return weights
+
+        unprefixed_module_name = module_name.removeprefix("model.")
+        weights = lora_model.get_module_to_save(unprefixed_module_name)
+        if weights is not None and unprefixed_module_name != module_name:
+            logger.info_once(
+                "For the pool model, successfully loaded the full module weights "
+                "after removing the prefix 'model.'."
+            )
+        return weights
+
+    def _validate_modules_to_save(self, lora_model: LoRAModel) -> None:
+        if not lora_model.modules_to_save:
+            return
+        saved_module_name = next(iter(lora_model.modules_to_save))
+        if self._classification_head is None:
+            raise ValueError(
+                f"Cannot load full module {saved_module_name!r}: the model does "
+                "not expose a unique ClassificationHeadWithLoRA."
+            )
+
+        module_name, wrapper = self._classification_head
+        full_module = self._get_module_to_save_weights(lora_model, module_name)
+        if full_module is None:
+            raise ValueError(
+                f"Full module {saved_module_name!r} does not match the model's "
+                f"classification head {module_name!r}."
+            )
+
+        expected_weight_shape = (wrapper.output_size, wrapper.input_size)
+        received_weight_shape = tuple(full_module.weight.shape)
+        if received_weight_shape != expected_weight_shape:
+            raise ValueError(
+                f"Full module {saved_module_name!r} for {module_name!r} has "
+                "an incompatible weight shape: expected "
+                f"{expected_weight_shape}, received {received_weight_shape}."
+            )
+
+        if full_module.bias is None:
+            return
+        expected_bias_shape = (wrapper.output_size,)
+        received_bias_shape = tuple(full_module.bias.shape)
+        if received_bias_shape != expected_bias_shape:
+            raise ValueError(
+                f"Full module {saved_module_name!r} for {module_name!r} has "
+                "an incompatible bias shape: expected "
+                f"{expected_bias_shape}, received {received_bias_shape}."
+            )
+
+    def _validate_token_classification_lora(self, lora_model: LoRAModel) -> None:
+        if self._pooling_task != "token_classify":
+            return
+        if self._classification_head is None:
+            return
+
+        module_name, _ = self._classification_head
+        has_full_module = (
+            self._get_module_to_save_weights(lora_model, module_name) is not None
+        )
+        has_lora_weights = (
+            self._get_lora_layer_weights(lora_model, module_name) is not None
+        )
+        if not has_full_module and not has_lora_weights:
+            return
+
+        raise ValueError(
+            f"LoRA adapter {lora_model.id} contains weights for classification "
+            f"head {module_name!r}, but token_classify only supports LoRA on "
+            "the model backbone."
+        )
+
     def deactivate_adapter(self, adapter_id: int) -> bool:
         if adapter_id not in self._active_adapters:
             return False
diff --git a/vllm/lora/ops/triton_ops/fp8_kernel_utils.py b/vllm/lora/ops/triton_ops/fp8_kernel_utils.py
index 8429562c7621..b9a7acf8627b 100644
--- a/vllm/lora/ops/triton_ops/fp8_kernel_utils.py
+++ b/vllm/lora/ops/triton_ops/fp8_kernel_utils.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Utilities for Punica kernel construction.
-"""
+"""Utilities for Punica kernel construction."""
 
 from vllm.triton_utils import tl, triton
 
@@ -21,8 +19,7 @@ def _accumulate_mm(
     group_n: tl.constexpr,
     use_fp8_w8a8: tl.constexpr,
 ):
-    """
-    Core matrix multiplication and accumulation logic with quantization support.
+    """Core matrix multiplication and accumulation logic with quantization support.
 
     Args:
         tiled_a (tl.tensor): Loaded tile from A matrix
@@ -36,8 +33,8 @@ def _accumulate_mm(
         group_k: Block size for K dimension in block-wise quantization
         group_n: Block size for N dimension in block-wise quantization
         use_fp8_w8a8: Whether using FP8 W8A8 quantization
-    """
 
+    """
     if use_fp8_w8a8:
         if group_k > 0 and group_n > 0:
             # Block-wise quantization: scales are loaded per block
@@ -85,8 +82,7 @@ def fp8_mm_k(
     USE_GDC: tl.constexpr,
     base_k,
 ):
-    """
-    FP8-compatible matrix multiplication kernel with quantization support.
+    """FP8-compatible matrix multiplication kernel with quantization support.
     Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of
     B (k x n), iterate through the K dimension to compute the partial/complete
     matrix block product with proper dequantization.
@@ -120,6 +116,7 @@ def fp8_mm_k(
         b_dtype: datatype of the B matrix
         USE_GDC: Whether to use PDL. True indicates use.
         base_k (int): Base offset along K dimension for current SPLIT_K group
+
     """
     accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
@@ -238,13 +235,11 @@ def do_shrink_kernel_fp8(
     per_channel_quant: tl.constexpr,
     launch_pdl: tl.constexpr,
 ):
-    """
-    Given an array of integers that identifies the rows of A, ram,
+    """Given an array of integers that identifies the rows of A, ram,
     a lora index that identifies which LoRA to use from lora_ptr, lora_index,
     a slice_id that identifies the input/output slice, compute the
     matrix product and store in the appropriate output location.
     """
-
     # Identify the lora_ptr from slice_id.
     if SLICE_NUM == 1:
         cur_lora_ptr = lora_ptr
@@ -432,8 +427,7 @@ def do_expand_kernel_fp8(
     use_fp8_w8a8: tl.constexpr,
     per_channel_quant: tl.constexpr,
 ):
-    """
-    FP8-compatible expand kernel for LoRA.
+    """FP8-compatible expand kernel for LoRA.
     Given an array of integers that identifies the rows of A, ram,
     a lora index that identifies which LoRA to use from lora_ptr, lora_index,
     a slice_id that identifies the input/output slice,
@@ -449,7 +443,6 @@ def do_expand_kernel_fp8(
     - Per-channel quantization
     - Tensor-wise quantization
     """
-
     # ls_d*_ptr can be either an integer or a pointer
     if SAME_STRIDE:
         cur_lora_d0_stride = ls_d0_ptr
diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py
index 033ed1df6ced..780d4a0d70bf 100644
--- a/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py
+++ b/vllm/lora/ops/triton_ops/fused_moe_lora_fp8_op.py
@@ -24,7 +24,7 @@ def _get_lora_id(
     top_k_num,
     naive_block_assignment: tl.constexpr,
 ):
-    """Returns lora_id"""
+    """Returns lora_id."""
     if naive_block_assignment:
         token_idx = pid_m // top_k_num
         return tl.load(token_lora_mapping_ptr + token_idx)
@@ -41,7 +41,7 @@ def _get_expert_id(
     max_loras,
     naive_block_assignment: tl.constexpr,
 ):
-    """Returns expert_id"""
+    """Returns expert_id."""
     if naive_block_assignment:
         return tl.load(expert_ids_ptr + pid_m)
     else:
@@ -61,7 +61,7 @@ def _get_token_offs(
     naive_block_assignment: tl.constexpr,
     BLOCK_SIZE_M: tl.constexpr,
 ):
-    """Returns token offsets"""
+    """Returns token offsets."""
     if naive_block_assignment:
         return tl.where(offs == 0, pid_m, num_valid_tokens)
     else:
@@ -76,8 +76,7 @@ def _get_token_offs(
 
 
 def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
-    """
-    `_LORA_PTR_DICT` collects the required information during `profile_run`,
+    """`_LORA_PTR_DICT` collects the required information during `profile_run`,
     After this, it remains constant and subsequent usage is through LUT.
     Refer to:
     https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
@@ -101,9 +100,7 @@ def _adjust_kernel_inputs(
     sorted_token_ids: torch.Tensor | None,
     expert_ids: torch.Tensor,
 ):
-    """
-    helper function to adjust kernel inputs when sorted_token_ids is None
-    """
+    """Helper function to adjust kernel inputs when sorted_token_ids is None."""
     if sorted_token_ids is None:
         stride_tl = 0
         stride_el = 0
diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py
index 1bde7e084533..703f7668894e 100644
--- a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+++ b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py
@@ -25,7 +25,7 @@ def _get_lora_id(
     top_k_num,
     naive_block_assignment: tl.constexpr,
 ):
-    """Returns lora_id"""
+    """Returns lora_id."""
     if naive_block_assignment:
         token_idx = pid_m // top_k_num
         return tl.load(token_lora_mapping_ptr + token_idx)
@@ -42,7 +42,7 @@ def _get_expert_id(
     max_loras,
     naive_block_assignment: tl.constexpr,
 ):
-    """Returns expert_id"""
+    """Returns expert_id."""
     if naive_block_assignment:
         return tl.load(expert_ids_ptr + pid_m)
     else:
@@ -62,7 +62,7 @@ def _get_token_offs(
     naive_block_assignment: tl.constexpr,
     BLOCK_SIZE_M: tl.constexpr,
 ):
-    """Returns token offsets"""
+    """Returns token offsets."""
     if naive_block_assignment:
         return tl.where(offs == 0, pid_m, num_valid_tokens)
     else:
@@ -420,6 +420,10 @@ def _run_fused_moe_lora_one_shot(
         npid = min(npid_occ, max_npid_by_budget)
     npid = max(1, min(npid, max(1, N_per_slice // 128)))
 
+    # see issue: https://github.com/intel/intel-xpu-backend-for-triton/issues/8121
+    if current_platform.is_xpu():
+        npid = 1
+
     # Robust defaults across the prefill regime (H100/H200/B200, bf16/fp16).
     # NPID > 1 is the small-M / under-saturated path -- more warps help
     # amortise the inner-N expand loop. ns=3 instead of 4: GB200 ncu showed
@@ -852,8 +856,7 @@ def _run_fused_moe_lora_small_batch(
 
 
 def _get_ptr(lora_weights: list[torch.Tensor], device: torch.device):
-    """
-    `_LORA_PTR_DICT` collects the required information during `profile_run`,
+    """`_LORA_PTR_DICT` collects the required information during `profile_run`,
     After this, it remains constant and subsequent usage is through LUT.
     Refer to:
     https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
@@ -877,9 +880,7 @@ def _adjust_kernel_inputs(
     sorted_token_ids: torch.Tensor | None,
     expert_ids: torch.Tensor,
 ):
-    """
-    helper function to adjust kernel inputs when sorted_token_ids is None
-    """
+    """Helper function to adjust kernel inputs when sorted_token_ids is None."""
     if sorted_token_ids is None:
         stride_tl = 0
         stride_el = 0
diff --git a/vllm/lora/ops/triton_ops/kernel_utils.py b/vllm/lora/ops/triton_ops/kernel_utils.py
index 34a9939ffcd5..fd6ae4442869 100644
--- a/vllm/lora/ops/triton_ops/kernel_utils.py
+++ b/vllm/lora/ops/triton_ops/kernel_utils.py
@@ -1,19 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Utilities for Punica kernel construction.
-"""
+"""Utilities for Punica kernel construction."""
 
 from vllm.triton_utils import tl, triton
 
 
 @triton.jit
 def mm_k(
-    a_ptr,
-    b_ptr,
-    ak_stride,
-    bk_stride,
-    offset_k,
+    a_ptr: tl.tensor,
+    b_ptr: tl.tensor,
+    ak_stride: int,
+    bk_stride: int,
+    offset_k: tl.tensor,
     K: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -23,32 +21,34 @@ def mm_k(
     CAST_TYPE: tl.constexpr,
     b_dtype: tl.constexpr,
     USE_GDC: tl.constexpr,
-    base_k,
+    base_k: int,
 ):
-    """
-    Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of
+    """Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of
     B (k x n), iterate, through the K dimension to compute the partial/complete
     matrix block product.
     If SPLIT_K == 1, the output m x n product is complete.
     If SPLIT_K > 1, the thread block computes partial outputs. The partial
     outputs are then atomically summed in the caller code.
+
     Args:
         a_ptr: Array of pointers, identifying rows of A
         b_ptr: Array of pointers, identifying columns of B
         ak_stride: K dimension stride of the A matrix
         bk_stride: K dimension stride of the B matrix
+        offset_k: Offsets within the current BLOCK_K tile along K
         K: Length of the K dimension
         BLOCK_M: M dimension of the output block m x n
         BLOCK_N: N dimension of the output block m x n
         BLOCK_K: K dimension atom
         EVEN_K: True if the blocks of A and B can be loaded without any
-          masking.
+            masking.
         SPLIT_K: Parameter signifying parallelism in the K dimension.
         CAST_TYPE: if True, cast the values from the A matrix to the B
-          matrix dtype.
+            matrix dtype.
         b_dtype: datatype of the B matrix
         USE_GDC: Whether to use PDL. True indicates use.
         base_k: Base offset along K dimension for current SPLIT_K group
+
     """
     accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
@@ -142,15 +142,13 @@ def do_expand_kernel(
     ADD_INPUTS: tl.constexpr,
     USE_GDC: tl.constexpr,
 ):
-    """
-    Given an array of integers that identifies the rows of A, ram,
+    """Given an array of integers that identifies the rows of A, ram,
     a lora index that identifies which LoRA to use from lora_ptr, lora_index,
     a slice_id that identifies the input/output slice,
     compute the matrix product and store in the appropriate output location.
     Given that this is an expand kernel, we don't perform any split-K reduction
     as the K dimension is assumed to be small.
     """
-
     # ls_d*_ptr can be either an integer or a pointer
     if SAME_STRIDE:
         # integer
@@ -268,13 +266,11 @@ def do_shrink_kernel(
     SLICE_NUM: tl.constexpr,
     USE_GDC: tl.constexpr,
 ):
-    """
-    Given an array of integers that identifies the rows of A, ram,
+    """Given an array of integers that identifies the rows of A, ram,
     a lora index that identifies which LoRA to use from lora_ptr, lora_index,
     a slice_id that identifies the input/output slice, compute the
     matrix product and store in the appropriate output location.
     """
-
     # Identify the lora_ptr from slice_id.
     if SLICE_NUM == 1:
         # current lora ptr
diff --git a/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py b/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py
index 2dc33965e327..9369fb750b2a 100644
--- a/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py
+++ b/vllm/lora/ops/triton_ops/lora_expand_fp8_op.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Based on:
+"""Based on:
 Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
 Punica: Multi-Tenant LoRA Serving.
 https://arxiv.org/abs/2310.18547
@@ -21,8 +20,7 @@
 
 
 def _get_expand_lora_scale_ptr(lora_weights: list[torch.Tensor], device: torch.device):
-    """
-    `_EXPAND_LORA_SCALE_PTR_DICT` collects the required information during
+    """`_EXPAND_LORA_SCALE_PTR_DICT` collects the required information during
     `profile_run`,
     After this, it remains constant and subsequent usage is through LUT.
     Refer to:
@@ -91,9 +89,7 @@ def _lora_expand_kernel_fp8(
     per_channel_quant: tl.constexpr,
     launch_pdl: tl.constexpr,
 ):
-    """
-    FP8-compatible expand kernel wrapper.
-    """
+    """FP8-compatible expand kernel wrapper."""
     cta_n_num = tl.cdiv(N, BLOCK_N)
     cta_m_num = tl.cdiv(M, BLOCK_M)
 
@@ -192,8 +188,7 @@ def _lora_expand_fp8(
     use_fp8_w8a8: bool = False,
     per_channel_quant: bool = False,
 ) -> None:
-    """
-    FP8-compatible LoRA expand operation.
+    """FP8-compatible LoRA expand operation.
 
     Args:
         inputs: Input tensor from shrink operation [num_slices, num_tokens, lora_rank]
@@ -206,6 +201,8 @@ def _lora_expand_fp8(
         num_tokens_per_lora: Number of tokens per LoRA
         lora_token_start_loc: Start location for each LoRA's tokens
         lora_ids: LoRA IDs to process
+        num_active_loras: Number of active LoRAs. Accepted for API parity with
+            the non-FP8 kernel and unused here.
         no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
             if there are any requests that require LoRA.
         offset_start (int, optional): Offset start for output_tensor.
@@ -216,6 +213,7 @@ def _lora_expand_fp8(
         group_n (int, optional): Block size for N in block-wise quantization.
         use_fp8_w8a8 (bool, optional): Whether to use FP8 W8A8 quantization.
         per_channel_quant (bool, optional): Whether to use per-channel quantization.
+
     """
     assert no_lora_flag_cpu.numel() == 1
     if no_lora_flag_cpu.item():
diff --git a/vllm/lora/ops/triton_ops/lora_expand_op.py b/vllm/lora/ops/triton_ops/lora_expand_op.py
index 40f640e2431d..184e58247bb3 100644
--- a/vllm/lora/ops/triton_ops/lora_expand_op.py
+++ b/vllm/lora/ops/triton_ops/lora_expand_op.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Based on:
+"""Based on:
 Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
 Punica: Multi-Tenant LoRA Serving.
 https://arxiv.org/abs/2310.18547
@@ -147,32 +146,31 @@ def _lora_expand(
     offset_start: int = 0,
     add_inputs: bool = False,
 ) -> None:
-    """
-    Args:
-        inputs (torch.Tensor): input tensor
-        lora_b_weights (list[torch.Tensor]): lora'b weight
-        output_tensor (torch.Tensor): output tensor
-        token_lora_mapping (torch.Tensor): A tensor mapping each input token
-            to the lora-id related to that token. A value of -1 indicates that
-            LoRA doesn't apply to that token.
-        token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from
-            the A matrix grouped by LoRA IDs.
-        num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number
-            of tokens that are to be processed by LoRA ID lora_ids[i]
-        lora_token_start_loc (torch.Tensor): A cumulative sum of
-            num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that
-            lora_token_start_loc[i], along with num_tokens_per_lora[i]
-            identifies the region in token_indices_sorted_by_lora_ids that
-            LoRA lora_ids[i] should process.
-        lora_ids (torch.Tensor): LoRA ids to process.
-        no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
-            if there are any requests that require LoRA.
-        offset_start (int, optional): Offset start for output_tensor.
-            Defaults to 0.
-        add_inputs (bool, optional): Whether to add the input tensor to the
-            output tensor. Defaults to False.
-    """
+    """Args:
+    inputs (torch.Tensor): input tensor
+    lora_b_weights (list[torch.Tensor]): lora'b weight
+    output_tensor (torch.Tensor): output tensor
+    token_lora_mapping (torch.Tensor): A tensor mapping each input token
+        to the lora-id related to that token. A value of -1 indicates that
+        LoRA doesn't apply to that token.
+    token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from
+        the A matrix grouped by LoRA IDs.
+    num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number
+        of tokens that are to be processed by LoRA ID lora_ids[i]
+    lora_token_start_loc (torch.Tensor): A cumulative sum of
+        num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that
+        lora_token_start_loc[i], along with num_tokens_per_lora[i]
+        identifies the region in token_indices_sorted_by_lora_ids that
+        LoRA lora_ids[i] should process.
+    lora_ids (torch.Tensor): LoRA ids to process.
+    no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
+        if there are any requests that require LoRA.
+    offset_start (int, optional): Offset start for output_tensor.
+        Defaults to 0.
+    add_inputs (bool, optional): Whether to add the input tensor to the
+        output tensor. Defaults to False.
 
+    """
     assert no_lora_flag_cpu.numel() == 1
     if no_lora_flag_cpu.item():
         # None of the inputs require LoRA.
diff --git a/vllm/lora/ops/triton_ops/lora_kernel_metadata.py b/vllm/lora/ops/triton_ops/lora_kernel_metadata.py
index dd7c2c706a07..e4bd7bfacdf2 100644
--- a/vllm/lora/ops/triton_ops/lora_kernel_metadata.py
+++ b/vllm/lora/ops/triton_ops/lora_kernel_metadata.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-LoRA kernels metadata preparation utilities.
-"""
+"""LoRA kernels metadata preparation utilities."""
 
 import bisect
 from dataclasses import dataclass, field
@@ -107,14 +105,13 @@ def _reset(self):
         self.num_active_loras_cpu.fill_(0)
 
     def prepare_tensors(self, token_lora_mapping: torch.Tensor) -> None:
-        """
-        Prepare kernel metadata tensors for the current forward pass.
+        """Prepare kernel metadata tensors for the current forward pass.
 
         Args:
             token_lora_mapping (torch.Tensor): Tensor containing lora indices
                 for each input token.
-        """
 
+        """
         self._reset()
 
         # Check and record no-lora case.
@@ -180,8 +177,7 @@ def meta_args(
         torch.Tensor,
         torch.Tensor,
     ]:
-        """
-        This function returns the kernel metadata required for the current
+        """This function returns the kernel metadata required for the current
         forward pass execution of the kernel. The function returns all the
         metadata required by the kernel, in order, as a tuple, so it can be
         unpacked directly during the lora_shrink/lora_expand function call.
@@ -189,6 +185,10 @@ def meta_args(
         Args:
             token_nums (int): Number of input tokens in the current forward
                 pass of the kernel.
+            specialize_active_lora (bool): If True, report the actual active
+                LoRA count; otherwise report the padded default, which keeps
+                the kernel launch shape stable.
+
         """
         if specialize_active_lora:
             num_active_loras = self.num_active_loras_cpu
diff --git a/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py b/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py
index 9b4fe75f3c04..c2b852fa40ab 100644
--- a/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py
+++ b/vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Based on:
+"""Based on:
 Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
 Punica: Multi-Tenant LoRA Serving.
 https://arxiv.org/abs/2310.18547
@@ -20,8 +19,7 @@
 def _get_shrink_lora_scale_ptr(
     lora_scale_weights: list[torch.Tensor], device: torch.device
 ):
-    """
-    `_SHRINK_LORA_SCALE_PTR_DICT` collects the required information during
+    """`_SHRINK_LORA_SCALE_PTR_DICT` collects the required information during
     `profile_run`. After this, it remains constant and subsequent usage is
     through LUT.
 
@@ -241,23 +239,23 @@ def _lora_shrink_fp8(
     use_fp8_w8a8: bool = False,
     per_channel_quant: bool = False,
 ) -> None:
-    """
-    Args:
-        inputs: FP8 or FP16/BF16 input tensor [num_tokens, hidden_size]
-        lora_a_weights: List of FP8 or FP16/BF16 LoRA A weights per slice
-        output_tensor: Output tensor (FP16/BF16/FP32)
-        token_lora_mapping: Token to LoRA ID mapping
-        token_indices_sorted_by_lora_ids: Sorted token indices
-        num_tokens_per_lora: Number of tokens per LoRA
-        lora_token_start_loc: Start location for each LoRA's tokens
-        lora_ids: LoRA IDs to process
-        scaling: LoRA scaling factor
-        a_scale: Activation quantization scales
-        b_scale: Weight quantization scales per slice
-        group_k: Block size for K dimension quantization
-        group_n: Block size for N dimension quantization
-        use_fp8_w8a8: Whether to use FP8 weights and activations
-        per_channel_quant: Whether to use per-channel quantization
+    """Args:
+    inputs: FP8 or FP16/BF16 input tensor [num_tokens, hidden_size]
+    lora_a_weights: List of FP8 or FP16/BF16 LoRA A weights per slice
+    output_tensor: Output tensor (FP16/BF16/FP32)
+    token_lora_mapping: Token to LoRA ID mapping
+    token_indices_sorted_by_lora_ids: Sorted token indices
+    num_tokens_per_lora: Number of tokens per LoRA
+    lora_token_start_loc: Start location for each LoRA's tokens
+    lora_ids: LoRA IDs to process
+    scaling: LoRA scaling factor
+    a_scale: Activation quantization scales
+    b_scale: Weight quantization scales per slice
+    group_k: Block size for K dimension quantization
+    group_n: Block size for N dimension quantization
+    use_fp8_w8a8: Whether to use FP8 weights and activations
+    per_channel_quant: Whether to use per-channel quantization
+
     """
     assert no_lora_flag_cpu.numel() == 1
     if no_lora_flag_cpu.item():
diff --git a/vllm/lora/ops/triton_ops/lora_shrink_op.py b/vllm/lora/ops/triton_ops/lora_shrink_op.py
index 4144c16eadd7..ec29065b9acf 100644
--- a/vllm/lora/ops/triton_ops/lora_shrink_op.py
+++ b/vllm/lora/ops/triton_ops/lora_shrink_op.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Based on:
+"""Based on:
 Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
 Punica: Multi-Tenant LoRA Serving.
 https://arxiv.org/abs/2310.18547
@@ -142,32 +141,31 @@ def _lora_shrink(
     num_active_loras: torch.Tensor,  # CPU tensor [1], number of active LoRAs
     scaling: float,
 ) -> None:
-    """
-    Args:
-        inputs (torch.Tensor): Input tensor
-        lora_a_weights (list[torch.Tensor]): LoRA weights
-        output_tensor (torch.Tensor): output tensor
-        token_lora_mapping (torch.Tensor): A tensor mapping each input token
-            to the lora-id related to that token. A value of -1 indicates that
-            LoRA doesn't apply to that token.
-        token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from
-            the A matrix grouped by LoRA IDs.
-        num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number
-            of tokens that are to be processed by LoRA ID lora_ids[i]
-        lora_token_start_loc (torch.Tensor): A cumulative sum of
-            num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that
-            lora_token_start_loc[i], along with num_tokens_per_lora[i]
-            identifies the region in token_indices_sorted_by_lora_ids that
-            LoRA lora_ids[i] should process.
-        lora_ids (torch.Tensor): LoRA ids to process.
-        no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
-            if there are any requests that require LoRA.
-        num_active_loras (torch.Tensor): A CPU tensor of size 1, containing the
-            number of active LoRAs. Stored as a tensor (not int) so
-            torch.compile treats it as dynamic rather than a constant.
-        scaling (float): Scaling factor.
-    """
+    """Args:
+    inputs (torch.Tensor): Input tensor
+    lora_a_weights (list[torch.Tensor]): LoRA weights
+    output_tensor (torch.Tensor): output tensor
+    token_lora_mapping (torch.Tensor): A tensor mapping each input token
+        to the lora-id related to that token. A value of -1 indicates that
+        LoRA doesn't apply to that token.
+    token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from
+        the A matrix grouped by LoRA IDs.
+    num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number
+        of tokens that are to be processed by LoRA ID lora_ids[i]
+    lora_token_start_loc (torch.Tensor): A cumulative sum of
+        num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that
+        lora_token_start_loc[i], along with num_tokens_per_lora[i]
+        identifies the region in token_indices_sorted_by_lora_ids that
+        LoRA lora_ids[i] should process.
+    lora_ids (torch.Tensor): LoRA ids to process.
+    no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
+        if there are any requests that require LoRA.
+    num_active_loras (torch.Tensor): A CPU tensor of size 1, containing the
+        number of active LoRAs. Stored as a tensor (not int) so
+        torch.compile treats it as dynamic rather than a constant.
+    scaling (float): Scaling factor.
 
+    """
     assert no_lora_flag_cpu.numel() == 1
     if no_lora_flag_cpu.item():
         # None of the inputs require LoRA.
@@ -179,7 +177,7 @@ def _lora_shrink(
         assert weight.dtype in [torch.float16, torch.bfloat16]
 
     assert inputs.size(1) == lora_a_weights[0].size(-1)
-    assert inputs.is_contiguous()
+    inputs = inputs.contiguous()
     assert output_tensor.is_contiguous()
 
     # metadata sanity check
diff --git a/vllm/lora/ops/triton_ops/utils.py b/vllm/lora/ops/triton_ops/utils.py
index dd4c0d2097cf..fbcef91002d2 100644
--- a/vllm/lora/ops/triton_ops/utils.py
+++ b/vllm/lora/ops/triton_ops/utils.py
@@ -23,8 +23,7 @@
 
 
 def _get_lora_a_ptr(lora_a_weights: list[torch.Tensor], device: torch.device):
-    """
-    `_LORA_A_PTR_DICT` collects the required information during `profile_run`,
+    """`_LORA_A_PTR_DICT` collects the required information during `profile_run`,
     After this, it remains constant and subsequent usage is through LUT.
     Refer to:
     https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
@@ -75,14 +74,12 @@ def _get_lora_a_ptr(lora_a_weights: list[torch.Tensor], device: torch.device):
 def _get_lora_b_ptr(
     lora_weights: list[torch.Tensor], offset_start: int, device: torch.device
 ):
-    """
-     `_LORA_B_PTR_DICT` collects the required information during `profile_run`,
+    """`_LORA_B_PTR_DICT` collects the required information during `profile_run`,
     After this, it remains constant and subsequent usage is through LUT.
     Refer to:
     https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
 
     """
-
     key = tuple(lora_weight.data_ptr() for lora_weight in lora_weights)
     if values := _LORA_B_PTR_DICT.get(key):
         return values
@@ -317,9 +314,7 @@ def get_lora_op_configs(
 
 @lru_cache
 def supports_pdl(device: torch.device | None = None) -> bool:
-    """
-    Refer to: https://github.com/triton-lang/triton/blob/v3.5.0/python/tutorials/11-programmatic-dependent-launch.py
-    """
+    """Refer to: https://github.com/triton-lang/triton/blob/v3.5.0/python/tutorials/11-programmatic-dependent-launch.py."""
     # PDL requires compute capability SM90 or above
 
     return (
diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py
index a0bb8bf3c283..baffe1698762 100644
--- a/vllm/lora/peft_helper.py
+++ b/vllm/lora/peft_helper.py
@@ -18,8 +18,7 @@
 
 @dataclass
 class PEFTHelper:
-    """
-    A helper class for PEFT configurations, specifically designed for LoRA.
+    """A helper class for PEFT configurations, specifically designed for LoRA.
     This class handles configuration validation, compatibility checks for
     various LoRA implementations.
     """
@@ -40,12 +39,18 @@ class PEFTHelper:
     vllm_max_position_embeddings: int | None = field(default=False)
 
     def _validate_features(self) -> list[str]:
-        """
-        Check if there are any unsupported LoRA features.
-        """
+        """Check if there are any unsupported LoRA features."""
         error_msg = []
         if self.modules_to_save:
-            error_msg.append("vLLM only supports modules_to_save being None.")
+            unsupported_modules = [
+                m for m in self.modules_to_save if m not in ["classifier", "score"]
+            ]
+            if unsupported_modules:
+                error_msg.append(
+                    "vLLM only supports modules_to_save being either None "
+                    'or ["classifier", "score"] for classification models. '
+                    f"Unsupported modules_to_save: {unsupported_modules}"
+                )
         if self.use_dora:
             error_msg.append("vLLM does not yet support DoRA.")
         return error_msg
@@ -116,8 +121,7 @@ def from_local_dir(
         return cls.from_dict(config)
 
     def validate_legal(self, lora_config: LoRAConfig) -> None:
-        """
-        Validates the LoRA configuration settings against application
+        """Validates the LoRA configuration settings against application
         constraints and requirements.
         """
         error_msg = self._validate_features()
diff --git a/vllm/lora/punica_wrapper/punica_base.py b/vllm/lora/punica_wrapper/punica_base.py
index f5c7ca9cb73b..a0e7751e6ddf 100644
--- a/vllm/lora/punica_wrapper/punica_base.py
+++ b/vllm/lora/punica_wrapper/punica_base.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Based on:
+"""Based on:
 Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
 Punica: Multi-Tenant LoRA Serving.
 https://arxiv.org/abs/2310.18547
@@ -20,9 +19,7 @@
 
 
 class PunicaWrapperABC(ABC):
-    """
-    PunicaWrapper ABC.
-    """
+    """PunicaWrapper ABC."""
 
     @abstractmethod
     def update_metadata(
@@ -33,9 +30,7 @@ def update_metadata(
         vocab_size: int,
         **kwargs,
     ) -> None:
-        """
-        Update the lora-related metadata
-        """
+        """Update the lora-related metadata"""
         raise NotImplementedError
 
     @abstractmethod
@@ -47,10 +42,7 @@ def add_shrink(
         scale: float,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Performs GEMM  for multiple slices of lora_a.
-        """
-
+        """Performs GEMM  for multiple slices of lora_a."""
         raise NotImplementedError
 
     @abstractmethod
@@ -64,9 +56,7 @@ def add_expand(
         add_inputs=True,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Performs GEMM for multiple slices of lora_b.
-        """
+        """Performs GEMM for multiple slices of lora_b."""
         raise NotImplementedError
 
     @abstractmethod
@@ -78,8 +68,7 @@ def add_lora_embedding(
         add_inputs: bool = True,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Applies lora  specifically for VocabParallelEmbeddingWithLoRA,
+        """Applies lora  specifically for VocabParallelEmbeddingWithLoRA,
         and this layer only requires the expand operation.
         """
         raise NotImplementedError
@@ -97,10 +86,7 @@ def add_lora_linear(
         buffer: tuple[torch.Tensor, ...] | None = None,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Applicable to linear-related lora.
-        """
-
+        """Applicable to linear-related lora."""
         raise NotImplementedError
 
     @abstractmethod
@@ -115,15 +101,24 @@ def add_lora_logits(
         buffer: torch.Tensor | None = None,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Applies lora  specifically for LogitsProcessorWithLoRA.
-        """
+        """Applies lora  specifically for LogitsProcessorWithLoRA."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_lora_full_linear(
+        self,
+        y: torch.Tensor,
+        x: torch.Tensor,
+        weight_stacked: torch.Tensor,
+        bias_stacked: torch.Tensor,
+        module_enabled: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Apply request-routed full linear weights to selected rows."""
         raise NotImplementedError
 
 
 class PunicaWrapperBase(PunicaWrapperABC):
-    """
-    PunicaWrapperBase is designed to manage and provide metadata for the punica
+    """PunicaWrapperBase is designed to manage and provide metadata for the punica
     kernel. The main function is to maintain the state information for
     Multi-LoRA, and to provide the interface for the punica.
     """
@@ -217,12 +212,27 @@ def _update_prefill_metadata(self, token_lora_tensor: torch.Tensor) -> None:
         self.token_nums = token_nums
         self.no_lora = no_lora
 
+    def _select_full_linear_output(
+        self,
+        y: torch.Tensor,
+        adapter_y: torch.Tensor,
+        bias_stacked: torch.Tensor,
+        module_enabled: torch.Tensor,
+    ) -> torch.Tensor:
+        indices = self.sampler_indices
+        assert indices.size(0) == y.size(0), (
+            "Full linear rows do not match LoRA request mapping"
+        )
+        safe_indices = indices.clamp_min(0)
+        use_full = (indices >= 0) & module_enabled[safe_indices]
+        adapter_y = adapter_y.to(y.dtype) + bias_stacked[safe_indices].to(y.dtype)
+        return torch.where(use_full.unsqueeze(-1), adapter_y, y)
+
     @property
     def prefill_metadata(
         self,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]:
-        """
-        This property provides a convenient way to access the necessary
+        """This property provides a convenient way to access the necessary
         metadata for prefill-related  kernel computations.
             1. seq_start_locs: Tensor of sequence start positions.
             2. seq_lengths: Tensor of sequence lengths.
@@ -243,8 +253,7 @@ def prefill_metadata(
 
     @property
     def token_lora_indices(self) -> torch.Tensor:
-        """
-        This property provides the lora indices corresponding to each token
+        """This property provides the lora indices corresponding to each token
         in the batch. An index of -1 means no lora should be applied.
         """
         token_lora_len = self.indices_len[0]
@@ -252,8 +261,7 @@ def token_lora_indices(self) -> torch.Tensor:
 
     @property
     def sampler_indices(self) -> torch.Tensor:
-        """
-        This property is used to access the lora indices specifically for
+        """This property is used to access the lora indices specifically for
         LogitsProcessorWithLoRA.
         """
         sampler_indices_len = self.indices_len[1]
@@ -261,9 +269,7 @@ def sampler_indices(self) -> torch.Tensor:
 
     @property
     def sampler_indices_padded(self) -> torch.Tensor:
-        """
-        This property provides access to padded sampler indices.
-        """
+        """This property provides access to padded sampler indices."""
         indices_padded_len = self.indices_len[2]
         return self._sampler_indices_padded[:indices_padded_len]
 
@@ -293,8 +299,7 @@ def add_shrink(
         scale: float,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Performs GEMM  for multiple slices of lora_a.
+        """Performs GEMM  for multiple slices of lora_a.
 
         Semantics:
         for i in range(len(lora_a_stacked)):
@@ -321,8 +326,7 @@ def add_expand(
         add_inputs=True,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Performs GEMM for multiple slices of lora_b.
+        """Performs GEMM for multiple slices of lora_b.
 
         Semantics:
             offset = offset_start
@@ -352,8 +356,7 @@ def add_lora_embedding(
         add_inputs: bool = True,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Applies lora  specifically for VocabParallelEmbeddingWithLoRA.
+        """Applies lora  specifically for VocabParallelEmbeddingWithLoRA.
         and this layer only requires the expand operation.
         Semantics:
             y += x @ lora_b_stacked
@@ -363,6 +366,7 @@ def add_lora_embedding(
             x (torch.Tensor): Input tensor.
             lora_b_stacked (torch.Tensor): lora_b's weights.
             add_inputs (bool): Default to True.
+
         """
         # TODO: implement it based on torch ops
         raise NotImplementedError
@@ -380,8 +384,7 @@ def add_lora_linear(
         buffer: tuple[torch.Tensor, ...] | None = None,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Applicable to linear-related lora.
+        """Applicable to linear-related lora.
 
         Semantics:
             for i in range(len(lora_a_stacked)):
@@ -400,6 +403,7 @@ def add_lora_linear(
             scale (float): Scaling factor.
             output_slices (tuple[int, ...]): Every slice's size.
             buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None.
+
         """
         # TODO: implement it based on torch ops
         raise NotImplementedError
@@ -416,8 +420,7 @@ def add_lora_logits(
         buffer: torch.Tensor | None = None,
         **kwargs,
     ) -> torch.Tensor | None:
-        """
-        Applies lora  specifically for LogitsProcessorWithLoRA.
+        """Applies lora  specifically for LogitsProcessorWithLoRA.
 
         Semantics:
             buffer = (x @ lora_a_stacked) * scale
@@ -430,6 +433,7 @@ def add_lora_logits(
             lora_b_stacked (torch.Tensor):lora_b's weights.
             scale (float): Scaling factor.
             buffer (Optional[torch.Tensor]):Default to None.
+
         """
         # TODO: implement it based on torch ops
         raise NotImplementedError
@@ -446,8 +450,7 @@ def moe_lora_align_block_size(
         pad_sorted_ids: bool = False,
         naive_block_assignment: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Aligns tokens and experts into block-sized chunks for LoRA-based
+        """Aligns tokens and experts into block-sized chunks for LoRA-based
         mixture-of-experts (MoE) execution.
         """
         # TODO: implement it based on torch ops
@@ -473,8 +476,7 @@ def add_lora_fused_moe(
         offset: int = 0,
         token_lora_mapping: torch.Tensor | None = None,
     ):
-        """
-        Performs a fused forward computation for LoRA of
+        """Performs a fused forward computation for LoRA of
         Mixture-of-Experts (MoE) layer.
         """
         # TODO: implement it based on torch ops
diff --git a/vllm/lora/punica_wrapper/punica_cpu.py b/vllm/lora/punica_wrapper/punica_cpu.py
index 1a700d9bf1f0..436445554fa0 100644
--- a/vllm/lora/punica_wrapper/punica_cpu.py
+++ b/vllm/lora/punica_wrapper/punica_cpu.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Callable
+from typing import Any
 
 import torch
 
@@ -20,8 +21,7 @@
 # The platforms that are compatible with the PyTorch-native implementation can
 # inherit this class
 class PunicaWrapperCPU(PunicaWrapperBase):
-    """
-    PunicaWrapperCPU is designed to manage and provide metadata for the punica
+    """PunicaWrapperCPU is designed to manage and provide metadata for the punica
     kernel. The main function is to maintain the state information for
     Multi-LoRA, and to provide the interface for the pytorch punica ops.
     """
@@ -133,12 +133,10 @@ def _apply_expand(
         y_slice_size: int,
         add_inputs: bool = True,
     ):
-        """
-        Perform the ` y[:,y_offset:y_offset+y_slice_size]+=x@w_t_all`
+        """Perform the ` y[:,y_offset:y_offset+y_slice_size]+=x@w_t_all`
         computation, which is suitable for the
         GEMM of lora'b.
         """
-
         expand_slice_fun: Callable = (
             self._expand_slice_prefill if self.is_prefill else self._expand_slice_decode
         )
@@ -147,8 +145,7 @@ def _apply_expand(
     def _apply_shrink(
         self, y: torch.Tensor, x: torch.Tensor, w_t_all: torch.Tensor, scale: float
     ):
-        """
-        Perform the ` y+=x@w_t_all` computation, which is suitable for the
+        """Perform the ` y+=x@w_t_all` computation, which is suitable for the
         GEMM of lora'a.
         When `is_prefill is` true, it indicates that it is currently the
         prefill stage, and the `_shrink_prefill` function should be called.
@@ -169,10 +166,9 @@ def add_shrink(
         x: torch.Tensor,
         lora_a_stacked: tuple[torch.Tensor, ...],
         scale: float,
-        **kwargs,
+        **kwargs: Any,
     ):
-        """
-        Performs GEMM  for multiple slices of lora_a.
+        """Performs GEMM  for multiple slices of lora_a.
         When `is_prefill is` true, it indicates that it is currently the
         prefill stage, and the `_shrink_prefill` function should be called.
         Otherwise, it is the decode stage, and the _shrink_decode function
@@ -187,8 +183,9 @@ def add_shrink(
             x (torch.Tensor): Input tensor
             lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights
             scale (float): Scaling factor for the operation
-        """
+            **kwargs: Unused; accepted for compatibility with the base class signature.
 
+        """
         x = x.view(-1, x.shape[-1])
         # TODO fuse these kernels
         for slice_idx in range(len(lora_a_stacked)):
@@ -202,10 +199,9 @@ def add_expand(
         output_slices: tuple[int, ...],
         offset_start: int = 0,
         add_inputs=True,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Performs GEMM for multiple slices of lora_b.
+        """Performs GEMM for multiple slices of lora_b.
 
         Semantics:
             for i in range(len(lora_b_stacked)):
@@ -218,7 +214,10 @@ def add_expand(
             x (Union[tuple[torch.Tensor, ...], torch.Tensor]): Input tensors
             lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight
             output_slices (tuple[int, ...]): Every slice's size
+            offset_start (int): The starting position of y, defaults to 0.
             add_inputs (bool):  Defaults to True.
+            **kwargs: Unused; accepted for compatibility with the base class signature.
+
         """
         y_org = y
         y = y.view(-1, y.shape[-1])
@@ -241,10 +240,9 @@ def add_lora_embedding(
         x: torch.Tensor,
         lora_b_stacked: torch.Tensor,
         add_inputs: bool = True,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applies lora  specifically for VocabParallelEmbeddingWithLoRA.
+        """Applies lora  specifically for VocabParallelEmbeddingWithLoRA.
 
         Semantics:
             y += x @ lora_b_stacked
@@ -254,8 +252,9 @@ def add_lora_embedding(
             x (torch.Tensor): Input tensor.
             lora_b_stacked (torch.Tensor): lora_b's weights.
             add_inputs (bool): Default to True.
-        """
+            **kwargs: Unused; accepted for compatibility with the base class signature.
 
+        """
         # Embedding layer only need expand op
         expand_fun: Callable = (
             self._expand_prefill if self.is_prefill else self._expand_decode
@@ -272,10 +271,9 @@ def add_lora_linear(
         output_slices: tuple[int, ...],
         *,
         buffer: tuple[torch.Tensor, ...] | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applicable to linear-related lora.
+        """Applicable to linear-related lora.
 
         Semantics:
             for i in range(len(lora_a_stacked)):
@@ -294,8 +292,9 @@ def add_lora_linear(
             scale (float): Scaling factor.
             output_slices (tuple[int, ...]): Every slice's size.
             buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None.
-        """
+            **kwargs: Unused; accepted for compatibility with the base class signature.
 
+        """
         assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
 
         if buffer is None:
@@ -320,10 +319,9 @@ def add_lora_logits(
         scale,
         *,
         buffer: torch.Tensor | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applies lora  specifically for LogitsProcessorWithLoRA.
+        """Applies lora  specifically for LogitsProcessorWithLoRA.
 
         Semantics:
             buffer = (x @ lora_a_stacked) * scale
@@ -336,6 +334,8 @@ def add_lora_logits(
             lora_b_stacked (torch.Tensor):lora_b's weights.
             scale (float): Scaling factor.
             buffer (Optional[torch.Tensor]):Default to None.
+            **kwargs: Forwarded to `add_shrink` and `add_expand`.
+
         """
         y_org = y
         y = y.view(-1, y.shape[-1])
@@ -349,3 +349,23 @@ def add_lora_logits(
         bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices, scale)
         bgmv_expand(buffer, lora_b_stacked, y, self.sampler_indices, add_inputs=True)
         y = y.view_as(y_org)
+
+    def apply_lora_full_linear(
+        self,
+        y: torch.Tensor,
+        x: torch.Tensor,
+        weight_stacked: torch.Tensor,
+        bias_stacked: torch.Tensor,
+        module_enabled: torch.Tensor,
+    ) -> None:
+        indices = self.sampler_indices
+        adapter_y = torch.zeros(
+            (x.size(0), weight_stacked.size(-2)),
+            dtype=torch.float32,
+            device=x.device,
+        )
+        bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0)
+        result = self._select_full_linear_output(
+            y, adapter_y, bias_stacked, module_enabled
+        )
+        y.copy_(result)
diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py
index 2f746f967e37..8604e747f9df 100644
--- a/vllm/lora/punica_wrapper/punica_gpu.py
+++ b/vllm/lora/punica_wrapper/punica_gpu.py
@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Based on:
+"""Based on:
 Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
 Punica: Multi-Tenant LoRA Serving.
 https://arxiv.org/abs/2310.18547
 """
 
-from typing import final
+from typing import Any, final
 
 import torch
 
@@ -32,8 +31,7 @@
 
 @final
 class PunicaWrapperGPU(PunicaWrapperBase):
-    """
-    PunicaWrapperGPU is designed to manage and provide metadata for the punica
+    """PunicaWrapperGPU is designed to manage and provide metadata for the punica
     kernel. The main function is to maintain the state information for
     Multi-LoRA, and to provide the interface for the punica triton kernel.
     """
@@ -96,10 +94,9 @@ def add_shrink(
         x: torch.Tensor,
         lora_a_stacked: tuple[torch.Tensor, ...],
         scale: float,
-        **kwargs,
+        **kwargs: Any,
     ):
-        """
-        Performs GEMM  for multiple slices of lora_a.
+        """Performs GEMM  for multiple slices of lora_a.
 
         Semantics:
         for i in range(len(lora_a_stacked)):
@@ -110,8 +107,9 @@ def add_shrink(
             x (torch.Tensor): Input tensor
             lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights
             scale (float): Scaling factor for the operation
-        """
+            **kwargs: Unused; accepted for compatibility with the base class signature.
 
+        """
         x = x.view(-1, x.shape[-1])
         lora_shrink(
             x,
@@ -131,10 +129,9 @@ def add_expand(
         output_slices: tuple[int, ...],
         offset_start: int = 0,
         add_inputs=True,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Performs GEMM for multiple slices of lora_b.
+        """Performs GEMM for multiple slices of lora_b.
 
         Semantics:
             for i in range(len(lora_b_stacked)):
@@ -147,9 +144,12 @@ def add_expand(
             x (torch.Tensor): Input tensors
             lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight
             output_slices (tuple[int, ...]): Every slice's size
+            offset_start (int): The starting position of y, defaults to 0.
             add_inputs (bool): If True, add LoRA output to y; if False, write
                 LoRA-only output to y (used for dual-stream when base and LoRA
                 run on different CUDA streams). Defaults to True.
+            **kwargs: Unused; accepted for compatibility with the base class signature.
+
         """
         y_org = y
         y = y.view(-1, y.shape[-1])
@@ -177,10 +177,9 @@ def add_lora_embedding(
         x: torch.Tensor,
         lora_b_stacked: torch.Tensor,
         add_inputs: bool = True,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applies lora  specifically for VocabParallelEmbeddingWithLoRA.
+        """Applies lora  specifically for VocabParallelEmbeddingWithLoRA.
 
         Semantics:
             y += x @ lora_b_stacked
@@ -190,8 +189,9 @@ def add_lora_embedding(
             x (torch.Tensor): Input tensor.
             lora_b_stacked (torch.Tensor): lora_b's weights.
             add_inputs (bool): Default to True.
-        """
+            **kwargs: Unused; accepted for compatibility with the base class signature.
 
+        """
         lora_expand(
             x.unsqueeze(dim=0),
             (lora_b_stacked,),
@@ -213,10 +213,9 @@ def add_lora_linear(
         output_slices: tuple[int, ...],
         *,
         buffer: torch.Tensor | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applicable to linear-related lora.
+        """Applicable to linear-related lora.
 
         Semantics:
             for i in range(len(lora_a_stacked)):
@@ -226,6 +225,7 @@ def add_lora_linear(
                     @ lora_b_stacked[indices[i], layer_idx, :, :]
                     * scale
                     ).squeeze(0)
+
         Args:
             y (torch.Tensor): Output tensor. Will be changed in-place.
             x (torch.Tensor): Input tensor
@@ -234,8 +234,10 @@ def add_lora_linear(
             scale (float): Scaling factor.
             output_slices (tuple[int, ...]): Every slice's size.
             buffer (Optional[torch.Tensor]): Defaults to None.
-        """
+            **kwargs: Accepts `add_inputs` (bool, defaults to True); any other
+                keys are ignored.
 
+        """
         assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
 
         assert buffer is None, (
@@ -275,10 +277,9 @@ def add_lora_logits(
         scale,
         *,
         buffer: torch.Tensor | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applies lora  specifically for LogitsProcessorWithLoRA.
+        """Applies lora  specifically for LogitsProcessorWithLoRA.
 
         Semantics:
             buffer = (x @ lora_a_stacked) * scale
@@ -291,6 +292,8 @@ def add_lora_logits(
             lora_b_stacked (torch.Tensor): lora_b's weights.
             scale (float): Scaling factor.
             buffer (Optional[torch.Tensor]): Default to None.
+            **kwargs: Unused; accepted for compatibility with the base class signature.
+
         """
         y_org = y
         y = y.view(-1, y.shape[-1])
@@ -327,6 +330,33 @@ def add_lora_logits(
         )
         y = y.view_as(y_org)
 
+    def apply_lora_full_linear(
+        self,
+        y: torch.Tensor,
+        x: torch.Tensor,
+        weight_stacked: torch.Tensor,
+        bias_stacked: torch.Tensor,
+        module_enabled: torch.Tensor,
+    ) -> None:
+        adapter_y = torch.empty(
+            (x.size(0), weight_stacked.size(-2)),
+            dtype=torch.float32,
+            device=x.device,
+        )
+        lora_shrink(
+            x,
+            [weight_stacked],
+            adapter_y.unsqueeze(0),
+            *self.prompt_mapping_meta.meta_args(
+                x.size(0), self.lora_config.specialize_active_lora
+            ),
+            1.0,
+        )
+        result = self._select_full_linear_output(
+            y, adapter_y, bias_stacked, module_enabled
+        )
+        y.copy_(result)
+
     def moe_lora_align_block_size(
         self,
         topk_ids: torch.Tensor,
@@ -340,8 +370,7 @@ def moe_lora_align_block_size(
         naive_block_assignment: bool = False,
         token_lora_mapping: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Aligns tokens and experts into block-sized chunks for LoRA-based
+        """Aligns tokens and experts into block-sized chunks for LoRA-based
         mixture-of-experts (MoE) execution.
 
         When `token_lora_mapping` is provided, it overrides the global mapping
@@ -440,9 +469,7 @@ def add_lora_fused_moe(
         token_lora_mapping: torch.Tensor | None = None,
         add_inputs: bool = True,
     ):
-        """
-        Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.
-        """
+        """Perform a fused forward computation for a LoRA MoE layer."""
         (
             token_lora_mapping_meta,
             _,
diff --git a/vllm/lora/punica_wrapper/punica_xpu.py b/vllm/lora/punica_wrapper/punica_xpu.py
index 7f040732ee4c..d6af2e36680d 100755
--- a/vllm/lora/punica_wrapper/punica_xpu.py
+++ b/vllm/lora/punica_wrapper/punica_xpu.py
@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Based on:
+"""Based on:
 Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
 Punica: Multi-Tenant LoRA Serving.
 https://arxiv.org/abs/2310.18547
 """
 
-from typing import final
+from typing import Any, final
 
 import torch
 
@@ -29,8 +28,7 @@
 
 @final
 class PunicaWrapperXPU(PunicaWrapperBase):
-    """
-    PunicaWrapperXPU is designed to manage and provide metadata for the punica
+    """PunicaWrapperXPU is designed to manage and provide metadata for the punica
     kernel. The main function is to maintain the state information for
     Multi-LoRA, and to provide the interface for the punica ipex kernel.
     """
@@ -135,10 +133,9 @@ def add_shrink(
         x: torch.Tensor,
         lora_a_stacked: tuple[torch.Tensor, ...],
         scale: float,
-        **kwargs,
+        **kwargs: Any,
     ):
-        """
-        Performs GEMM  for multiple slices of lora_a.
+        """Performs GEMM  for multiple slices of lora_a.
 
         Semantics:
         for i in range(len(lora_a_stacked)):
@@ -149,8 +146,9 @@ def add_shrink(
             x (torch.Tensor): Input tensor
             lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights
             scale (float): Scaling factor for the operation
-        """
+            **kwargs: Unused; accepted for compatibility with the base class signature.
 
+        """
         x = x.view(-1, x.shape[-1])
         for slice_idx in range(len(lora_a_stacked)):
             self._apply_shrink(y[slice_idx], x, lora_a_stacked[slice_idx], scale)
@@ -163,10 +161,9 @@ def add_expand(
         output_slices: tuple[int, ...],
         offset_start: int = 0,
         add_inputs=True,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Performs GEMM for multiple slices of lora_b.
+        """Performs GEMM for multiple slices of lora_b.
 
         Semantics:
             for i in range(len(lora_b_stacked)):
@@ -179,7 +176,10 @@ def add_expand(
             x (torch.Tensor): Input tensors
             lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight
             output_slices (tuple[int, ...]): Every slice's size
+            offset_start (int): The starting position of y, defaults to 0.
             add_inputs (bool): Defaults to True.
+            **kwargs: Unused; accepted for compatibility with the base class signature.
+
         """
         y_org = y
         y = y.view(-1, y.shape[-1])
@@ -206,10 +206,9 @@ def add_lora_embedding(
         x: torch.Tensor,
         lora_b_stacked: torch.Tensor,
         add_inputs: bool = True,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applies lora  specifically for VocabParallelEmbeddingWithLoRA.
+        """Applies lora  specifically for VocabParallelEmbeddingWithLoRA.
 
         Semantics:
             y += x @ lora_b_stacked
@@ -219,6 +218,8 @@ def add_lora_embedding(
             x (torch.Tensor): Input tensor.
             lora_b_stacked (torch.Tensor): lora_b's weights.
             add_inputs (bool): Default to True.
+            **kwargs: Unused; accepted for compatibility with the base class signature.
+
         """
         token_lora_indices = self._get_token_lora_indices(x)
         bgmv_expand(x, lora_b_stacked, y, token_lora_indices, add_inputs)
@@ -233,10 +234,9 @@ def add_lora_linear(
         output_slices: tuple[int, ...],
         *,
         buffer: torch.Tensor | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applicable to linear-related lora.
+        """Applicable to linear-related lora.
 
         Semantics:
             for i in range(len(lora_a_stacked)):
@@ -246,6 +246,7 @@ def add_lora_linear(
                     @ lora_b_stacked[indices[i], layer_idx, :, :]
                     * scale
                     ).squeeze(0)
+
         Args:
             y (torch.Tensor): Output tensor. Will be changed in-place.
             x (torch.Tensor): Input tensor
@@ -254,8 +255,10 @@ def add_lora_linear(
             scale (float): Scaling factor.
             output_slices (tuple[int, ...]): Every slice's size.
             buffer (Optional[torch.Tensor]): Defaults to None.
-        """
+            **kwargs: Accepts `add_inputs` (bool, defaults to True); any other
+                keys are ignored.
 
+        """
         assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
 
         assert buffer is None, (
@@ -287,9 +290,7 @@ def add_lora_linear(
 
     @property
     def sampler_indices_padded(self) -> torch.Tensor:
-        """
-        This property provides access to padded sampler indices.
-        """
+        """This property provides access to padded sampler indices."""
         return self._sampler_indices_padded[:]
 
     def add_lora_logits(
@@ -301,10 +302,9 @@ def add_lora_logits(
         scale,
         *,
         buffer: torch.Tensor | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
-        """
-        Applies lora  specifically for LogitsProcessorWithLoRA.
+        """Applies lora  specifically for LogitsProcessorWithLoRA.
 
         Semantics:
             buffer = (x @ lora_a_stacked) * scale
@@ -317,6 +317,8 @@ def add_lora_logits(
             lora_b_stacked (torch.Tensor): lora_b's weights.
             scale (float): Scaling factor.
             buffer (Optional[torch.Tensor]): Default to None.
+            **kwargs: Unused; accepted for compatibility with the base class signature.
+
         """
         y_org = y
         y = y.view(-1, y.shape[-1])
@@ -333,6 +335,26 @@ def add_lora_logits(
         bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True)
         y = y.view_as(y_org)
 
+    def apply_lora_full_linear(
+        self,
+        y: torch.Tensor,
+        x: torch.Tensor,
+        weight_stacked: torch.Tensor,
+        bias_stacked: torch.Tensor,
+        module_enabled: torch.Tensor,
+    ) -> None:
+        indices = self.sampler_indices
+        adapter_y = torch.zeros(
+            (x.size(0), weight_stacked.size(-2)),
+            dtype=torch.float32,
+            device=x.device,
+        )
+        bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0)
+        result = self._select_full_linear_output(
+            y, adapter_y, bias_stacked, module_enabled
+        )
+        y.copy_(result)
+
     def moe_lora_align_block_size(
         self,
         topk_ids: torch.Tensor,
@@ -346,8 +368,7 @@ def moe_lora_align_block_size(
         naive_block_assignment: bool = False,
         token_lora_mapping: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Aligns tokens and experts into block-sized chunks for LoRA-based
+        """Aligns tokens and experts into block-sized chunks for LoRA-based
         mixture-of-experts (MoE) execution.
 
         When `token_lora_mapping` is provided, it overrides the global mapping
@@ -446,9 +467,7 @@ def add_lora_fused_moe(
         token_lora_mapping: torch.Tensor | None = None,
         add_inputs: bool = True,
     ):
-        """
-        Performs a fused forward computation for LoRA of Mixture-of-Experts (MoE) layer.
-        """
+        """Perform a fused forward computation for a LoRA MoE layer."""
         (
             token_lora_mapping_meta,
             _,
diff --git a/vllm/lora/punica_wrapper/utils.py b/vllm/lora/punica_wrapper/utils.py
index c29e359f1917..821ba8fc44d9 100644
--- a/vllm/lora/punica_wrapper/utils.py
+++ b/vllm/lora/punica_wrapper/utils.py
@@ -15,15 +15,13 @@
 def compute_meta(
     token_lora_tensor: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int, bool]:
-    """
-    Get the information required for the sgmv kernel. With the  features:
+    """Get the information required for the sgmv kernel. With the  features:
     1. If consecutive requests in the batch use the same LoRA, this function
     will combine them into a single request, improving sgmv kernel inference
     performance.
     2. At the beginning of each prefill stage inference, recalculations are
     needed based on the input, but only once.
     """
-
     lora_indices_tensor, seq_length_tensor = torch.unique_consecutive(
         token_lora_tensor, return_counts=True
     )
@@ -65,6 +63,7 @@ def convert_mapping(
         lora_index_to_id: List mapping LoRA ids to LoRA indices.
         max_loras: Maximum number of LoRAs.
         vocab_size: Model vocab size.
+        device: Device the returned index tensors are created on.
 
     Returns:
         A tuple of tensors:
@@ -83,6 +82,7 @@ def convert_mapping(
             indices_len: List of lengths of the above tensors. It contains
                 (base_indices, sampler_indices, sampler_indices_padded,
                 embeddings_indices).
+
     """
     index_mapping_indices: list[int] = list(mapping.index_mapping).copy()
     embedding_indices = index_mapping_indices.copy()
diff --git a/vllm/lora/request.py b/vllm/lora/request.py
index 7ba1725fa2e8..cd3df54d3b12 100644
--- a/vllm/lora/request.py
+++ b/vllm/lora/request.py
@@ -10,8 +10,7 @@ class LoRARequest(
     omit_defaults=True,  # type: ignore[call-arg]
     array_like=True,
 ):  # type: ignore[call-arg]
-    """
-    Request for a LoRA adapter.
+    """Request for a LoRA adapter.
 
     lora_int_id must be globally unique for a given adapter.
     This is currently not enforced in vLLM.
@@ -56,16 +55,14 @@ def path(self):
         return self.lora_path
 
     def __eq__(self, value: object) -> bool:
-        """
-        Overrides the equality method to compare LoRARequest
+        """Overrides the equality method to compare LoRARequest
         instances based on lora_name. This allows for identification
         and comparison lora adapter across engines.
         """
         return isinstance(value, self.__class__) and self.lora_name == value.lora_name
 
     def __hash__(self) -> int:
-        """
-        Overrides the hash method to hash LoRARequest instances
+        """Overrides the hash method to hash LoRARequest instances
         based on lora_name. This ensures that LoRARequest instances
         can be used in hash-based collections such as sets and dictionaries,
         identified by their names across engines.
diff --git a/vllm/lora/resolver.py b/vllm/lora/resolver.py
index bcfe26467cfb..166cb9ef10fc 100644
--- a/vllm/lora/resolver.py
+++ b/vllm/lora/resolver.py
@@ -36,6 +36,7 @@ async def resolve_lora(
         Returns:
             Optional[LoRARequest]: The resolved LoRA model information, or None
             if the LoRA model cannot be found.
+
         """
         pass
 
@@ -54,9 +55,11 @@ def register_resolver(
         resolver: LoRAResolver,
     ) -> None:
         """Register a LoRA resolver.
+
         Args:
             resolver_name: Name to register the resolver under.
             resolver: The LoRA resolver instance to register.
+
         """
         if resolver_name in self.resolvers:
             logger.warning(
@@ -70,12 +73,16 @@ def register_resolver(
 
     def get_resolver(self, resolver_name: str) -> LoRAResolver:
         """Get a registered resolver instance by name.
+
         Args:
             resolver_name: Name of the resolver to get.
+
         Returns:
             The resolver instance.
+
         Raises:
             KeyError: If the resolver is not found in the registry.
+
         """
         if resolver_name not in self.resolvers:
             raise KeyError(
diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py
index f2a9dc902fce..71a2e60dc7b6 100644
--- a/vllm/lora/utils.py
+++ b/vllm/lora/utils.py
@@ -6,7 +6,7 @@
 
 from huggingface_hub.utils import HfHubHTTPError, HFValidationError
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm import envs
 from vllm.config.lora import LoRAConfig
@@ -15,6 +15,7 @@
 # being imported for _all_lora_classes below
 from vllm.lora.layers import (
     BaseLayerWithLoRA,
+    ClassificationHeadWithLoRA,
     ColumnParallelLinearWithLoRA,
     ColumnParallelLinearWithShardedLoRA,
     FusedMoE3DWithLoRA,
@@ -48,8 +49,7 @@
 
 
 def get_captured_lora_counts(max_loras: int, specialize: bool) -> list[int]:
-    """
-    Returns num_active_loras values for cudagraph capture.
+    """Returns num_active_loras values for cudagraph capture.
 
     When specialize=True: powers of 2 up to max_loras, plus max_loras + 1.
     When specialize=False: just [max_loras + 1].
@@ -109,7 +109,7 @@ def from_layer(
     max_loras: int,
     lora_config: LoRAConfig,
     packed_modules_list: list,
-    model_config: PretrainedConfig | None = None,
+    model_config: PreTrainedConfig | None = None,
 ) -> nn.Module:
     for lora_cls in _all_lora_classes:
         # specifying kwargs so they can be easily accessed in decorator
@@ -130,7 +130,7 @@ def from_layer_logits_processor(
     lm_head: "ParallelLMHead",
     max_loras: int,
     lora_config: LoRAConfig,
-    model_config: PretrainedConfig | None = None,
+    model_config: PreTrainedConfig | None = None,
 ) -> LogitsProcessorWithLoRA:
     ret = LogitsProcessorWithLoRA(
         layer,
@@ -143,6 +143,17 @@ def from_layer_logits_processor(
     return ret
 
 
+def from_layer_classification(
+    layer: nn.Module,
+    max_loras: int,
+    lora_config: LoRAConfig,
+    model_config: PreTrainedConfig | None = None,
+) -> ClassificationHeadWithLoRA:
+    instance_layer = ClassificationHeadWithLoRA(layer)
+    instance_layer.create_lora_weights(max_loras, lora_config, model_config)
+    return instance_layer
+
+
 def replace_submodule(
     model: nn.Module, module_name: str, new_module: nn.Module
 ) -> nn.Module:
@@ -158,7 +169,7 @@ def parse_fine_tuned_lora_name(
 ) -> tuple[str, bool]:
     """Parse the name of lora weights.
 
-    args:
+    Args:
         name: the name of the fine-tuned LoRA, e.g.
             base_model.model.dense1.weight
         weights_mapper: maps the name of weight, e.g.
@@ -167,8 +178,8 @@ def parse_fine_tuned_lora_name(
         tuple(module_name, is_lora_a):
             module_name: the name of the module, e.g. model.dense1,
             is_lora_a whether the tensor is lora_a or lora_b.
-    """
 
+    """
     # LoRA weight qualified name usually starts with `base_model.model.`,
     # so we remove the prefix `base_model.model.` to make the following
     # mapping correctly.
@@ -193,13 +204,14 @@ def parse_fine_tuned_lora_name(
     start_index = 2 if name.startswith("base_model.model.") else 0
 
     parts = name.split(".")
-    if (
-        parts[-1] == "weight"
-        and len(parts) >= 2
-        and (parts[-2] == "lora_A" or parts[-2] == "lora_B")
-    ):
-        new_name = ".".join(parts[start_index:-2])
-        return new_name, parts[-2] == "lora_A"
+    if (parts[-1] == "weight" or parts[-1] == "bias") and len(parts) >= 2:
+        if parts[-2] in ["lora_A", "lora_B"]:
+            new_name = ".".join(parts[start_index:-2])
+            return new_name, parts[-2] == "lora_A"
+        # For modules_to_save in classification.
+        elif parts[-2] in ["score", "classifier"]:
+            new_name = parts[-2]
+            return new_name, False
 
     if parts[-1] == "lora_embedding_A" or parts[-1] == "lora_embedding_B":
         new_name = ".".join(parts[start_index:-1])
@@ -218,10 +230,7 @@ def is_base_embedding_weights(name: str) -> bool:
 
 
 def get_supported_lora_modules(model: nn.Module) -> list[str]:
-    """
-    In vLLM, all linear layers support LoRA.
-    """
-
+    """In vLLM, all linear layers support LoRA."""
     supported_lora_modules: set[str] = set()
     for name, module in model.named_modules():
         # get the embedding modules if the module's embedding_modules
@@ -255,6 +264,7 @@ def is_supported_lora_module(
 
     Returns:
         True if the module is supported, False otherwise.
+
     """
     module_suffix = module_name.rsplit(".", 1)[-1]
     if module_suffix not in supported_lora_modules:
@@ -292,6 +302,7 @@ def is_in_target_modules(
 
     Returns:
         True if the module passes the filter, False otherwise.
+
     """
     if target_modules is None:
         return True
@@ -316,22 +327,21 @@ def is_in_target_modules(
 
 
 def get_adapter_absolute_path(lora_path: str) -> str:
-    """
-    Resolves the given lora_path to an absolute local path.
+    """Resolves the given lora_path to an absolute local path.
 
     If the lora_path is identified as a Hugging Face model identifier,
     it will download the model and return the local snapshot path.
     Otherwise, it treats the lora_path as a local file path and
     converts it to an absolute path.
 
-    Parameters:
-    lora_path (str): The path to the lora model, which can be an absolute path,
-                     a relative path, or a Hugging Face model identifier.
+    Args:
+        lora_path (str): The path to the lora model, which can be an absolute
+            path, a relative path, or a Hugging Face model identifier.
 
     Returns:
-    str: The resolved absolute local path to the lora model.
-    """
+        str: The resolved absolute local path to the lora model.
 
+    """
     # Check if the path is an absolute path. Return it no matter exists or not.
     if os.path.isabs(lora_path):
         return lora_path
diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py
index 4c00c6531bc8..93624523e2f8 100644
--- a/vllm/lora/worker_manager.py
+++ b/vllm/lora/worker_manager.py
@@ -157,6 +157,10 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel:
             # 3D->2D conversion when running under the universal 2D wrapper.
             lora.is_3d_lora_weight = lora_request.is_3d_lora_weight
 
+            # Validate classification-head weights.
+            self._adapter_manager._validate_modules_to_save(lora)
+            self._adapter_manager._validate_token_classification_lora(lora)
+
         except FileNotFoundError as e:
             # FileNotFoundError should be raised if both
             # - No adapter found to download from huggingface (or in
diff --git a/vllm/model_executor/custom_op.py b/vllm/model_executor/custom_op.py
index bf98e46e8546..49a71a339b57 100644
--- a/vllm/model_executor/custom_op.py
+++ b/vllm/model_executor/custom_op.py
@@ -30,8 +30,7 @@ def maybe_get_oot_by_class(class_type: type) -> type:
 
 
 class PluggableLayer(nn.Module):
-    """
-    Base class for pluggable layers.
+    """Base class for pluggable layers.
 
     A PluggableLayer is a *module-composing* abstraction: it may instantiate other
     ``torch.nn.Module`` objects as sub-layers, and its functionality depends on
@@ -101,8 +100,7 @@ def decorator(layer_cls):
 
 
 class CustomOp(nn.Module):
-    """
-    Base class for custom ops.
+    """Base class for custom ops.
     Dispatches the forward method to the appropriate backend.
     """
 
@@ -207,8 +205,7 @@ def dispatch_forward(self, compile_native: bool):
             return self.forward_cuda
 
     def maybe_compile(self, fn, *, enable: bool = True):
-        """
-        Compile fn if compilation enabled.
+        """Compile fn if compilation enabled.
         Useful for CustomOp instances called from within a torch custom op,
         meaning the forward call is hidden from the model-level torch.compile.
 
@@ -294,8 +291,7 @@ def enabled(cls) -> bool:
 
     @staticmethod
     def default_on() -> bool:
-        """
-        Behavior controlled by `CompilationConfig.custom_ops`: On by default if
+        """Behavior controlled by `CompilationConfig.custom_ops`: On by default if
         'all', off by default if 'none'.
         When PyTorch Inductor is used, 'none' is the default value,
         otherwise 'all'.
diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py
index 6c701cfa163c..308e8c4f42cb 100644
--- a/vllm/model_executor/determinism/batch_invariant.py
+++ b/vllm/model_executor/determinism/batch_invariant.py
@@ -231,6 +231,7 @@ def matmul_descriptor_persistent(
 
     Returns:
         Output matrix [M, N] with dtype matching the inputs.
+
     """
     assert a.shape[1] == b.shape[0], "Incompatible dimensions"
     assert a.dtype == b.dtype, "Incompatible dtypes"
@@ -378,7 +379,7 @@ def bmm_kernel(
     B_LARGE: tl.constexpr,
     C_LARGE: tl.constexpr,
 ):
-    """Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N)
+    """Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N).
 
     Each program computes one (batch_idx, tile_m, tile_n) tile, accumulating
     along K in a fixed order to preserve batch invariance.
@@ -495,8 +496,7 @@ def _log_softmax_kernel(
     n_cols,
     BLOCK_SIZE: tl.constexpr,
 ):
-    """
-    Compute log_softmax along the last dimension of a 2D tensor.
+    """Compute log_softmax along the last dimension of a 2D tensor.
     Each block handles one row of the input tensor.
     """
     # Get the row index for this block
@@ -550,8 +550,7 @@ def _log_softmax_kernel(
 
 
 def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor:
-    """
-    Compute log_softmax using Triton kernel.
+    """Compute log_softmax using Triton kernel.
 
     Args:
         input: Input tensor
@@ -560,6 +559,7 @@ def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor:
 
     Returns:
         Tensor with log_softmax applied along the specified dimension
+
     """
     if dim != -1 and dim != input.ndim - 1:
         raise ValueError(
@@ -607,8 +607,7 @@ def mean_kernel(
     K,  # size after reduction dim
     BLOCK_SIZE: tl.constexpr,
 ):
-    """
-    Kernel for computing mean along a single dimension.
+    """Kernel for computing mean along a single dimension.
     Input is viewed as (M, N, K) where N is the dimension being reduced.
     """
     # Program ID gives us which output element we're computing
@@ -649,8 +648,7 @@ def mean_dim(
     keepdim: bool = False,
     dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
-    """
-    Triton implementation of torch.mean with single dimension reduction.
+    """Triton implementation of torch.mean with single dimension reduction.
 
     Args:
         input: Input tensor
@@ -661,6 +659,7 @@ def mean_dim(
 
     Returns:
         Tensor with mean values along specified dimension
+
     """
     # Validate inputs
     assert -input.ndim <= dim < input.ndim, (
@@ -933,8 +932,7 @@ def _rms_norm_kernel(
     BLOCK_SIZE: tl.constexpr,
     HAS_WEIGHT: tl.constexpr,
 ):
-    """
-    Compute RMS normalization along the last dimension of a 2D tensor.
+    """Compute RMS normalization along the last dimension of a 2D tensor.
     RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight
     Each block handles one row of the input tensor.
     """
@@ -980,9 +978,7 @@ def rms_norm_batch_invariant(
     eps: float = 1e-6,
     residual: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    Compute RMS normalization using Triton kernel.
-
+    """Compute RMS normalization using Triton kernel.
 
     Args:
         input: Input tensor of shape (..., hidden_size)
@@ -994,6 +990,7 @@ def rms_norm_batch_invariant(
     Returns:
         RMS normalized tensor, or ``(output, residual_out)`` when ``residual``
         is provided
+
     """
     if residual is not None:
         assert input.shape == residual.shape, (
@@ -1107,7 +1104,10 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None):
 
         _fp16_block_size_n = 128
 
-    # Softmax, log_softmax, mean, and bmm are already batch-invariant on XPU
+    # Native mean can change reduction order with the row count on XPU too.
+    _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key)
+
+    # Softmax, log_softmax, and bmm are already batch-invariant on XPU
     # (oneDNN backend produces bitwise-identical results regardless of batch
     # context). Only register these overrides on CUDA where they are needed.
     if not current_platform.is_xpu():
@@ -1116,7 +1116,6 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None):
         )
         _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key)
         _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, key)
-        _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key)
         # torch 2.12+ registers a built-in Triton bmm kernel for CUDA
         # (torch._native.ops.bmm_outer_product), so we need allow_override
         # to replace it at the dispatcher level.
@@ -1151,7 +1150,12 @@ def override_envs_for_invariance():
     os.environ["NCCL_MIN_NCHANNELS"] = "1"
     os.environ["NCCL_MAX_NCHANNELS"] = "1"
     os.environ["NCCL_PROTO"] = "Simple"
-    os.environ["NCCL_ALGO"] = "allreduce:tree"
+    # NCCL >= 2.31 zero-fills the algorithm table of every collective when
+    # NCCL_ALGO is set and re-enables only the named ones; together with the
+    # NCCL_PROTO above, collectives not named here end up with no algorithm
+    # and fail with ncclInvalidUsage. Re-enable Ring and Tree for all
+    # collectives, then pin AllReduce to Tree for determinism.
+    os.environ["NCCL_ALGO"] = "ring,tree;allreduce:tree"
     os.environ["NCCL_NTHREADS"] = "1"
     os.environ["NCCL_SOCKET_NTHREADS"] = "1"
 
diff --git a/vllm/model_executor/determinism/batch_invariant_configs.py b/vllm/model_executor/determinism/batch_invariant_configs.py
index 39cda984f7f6..874447bb8295 100644
--- a/vllm/model_executor/determinism/batch_invariant_configs.py
+++ b/vllm/model_executor/determinism/batch_invariant_configs.py
@@ -254,6 +254,233 @@ class _MatmulShapeConfig(NamedTuple):
             ),
         ),
     },
+    "sm120": {
+        (12288, 2048): _MatmulShapeConfig(  # Qwen3-1.7B gate_up_proj
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 128, 8, 5)),
+                (4, _MatmulMConfig(16, 128, 8, 4)),
+                (8, _MatmulMConfig(16, 128, 8, 4)),
+                (16, _MatmulMConfig(16, 128, 4, 4)),
+                (32, _MatmulMConfig(32, 128, 8, 5)),
+                (64, _MatmulMConfig(64, 128, 4, 3)),
+                (256, _MatmulMConfig(64, 64, 4, 4)),
+                (512, _MatmulMConfig(64, 128, 4, 5)),
+                (1024, _MatmulMConfig(128, 128, 4, 3)),
+                (2048, _MatmulMConfig(128, 128, 4, 3)),
+            ),
+        ),
+        (2048, 6144): _MatmulShapeConfig(  # Qwen3-1.7B down_proj
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 32, 8, 5)),
+                (4, _MatmulMConfig(16, 32, 8, 3)),
+                (8, _MatmulMConfig(16, 32, 8, 4)),
+                (16, _MatmulMConfig(16, 32, 4, 5)),
+                (32, _MatmulMConfig(32, 32, 4, 5)),
+                (64, _MatmulMConfig(32, 32, 4, 4)),
+                (256, _MatmulMConfig(64, 64, 4, 5)),
+                (512, _MatmulMConfig(128, 64, 4, 3)),
+                (1024, _MatmulMConfig(128, 128, 4, 3)),
+                (2048, _MatmulMConfig(128, 64, 8, 4)),
+            ),
+        ),
+        (4096, 2048): _MatmulShapeConfig(  # Qwen3-1.7B qkv_proj
+            block_k=128,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 4, 4)),
+                (4, _MatmulMConfig(16, 64, 8, 5)),
+                (8, _MatmulMConfig(16, 64, 8, 5)),
+                (16, _MatmulMConfig(16, 64, 8, 3)),
+                (32, _MatmulMConfig(16, 64, 8, 5)),
+                (64, _MatmulMConfig(32, 64, 4, 2)),
+                (256, _MatmulMConfig(64, 128, 4, 3)),
+                (512, _MatmulMConfig(64, 64, 4, 3)),
+                (1024, _MatmulMConfig(64, 128, 4, 3)),
+                (2048, _MatmulMConfig(64, 128, 4, 3)),
+            ),
+        ),
+        (151936, 2048): _MatmulShapeConfig(  # Qwen3-1.7B lm_head
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 8, 5)),
+                (4, _MatmulMConfig(16, 128, 4, 4)),
+                (8, _MatmulMConfig(16, 128, 4, 5)),
+                (16, _MatmulMConfig(16, 128, 8, 3)),
+                (32, _MatmulMConfig(32, 128, 8, 3)),
+                (64, _MatmulMConfig(32, 128, 8, 5)),
+                (256, _MatmulMConfig(128, 128, 4, 3)),
+                (512, _MatmulMConfig(128, 128, 8, 3)),
+                (1024, _MatmulMConfig(128, 128, 4, 3)),
+                (2048, _MatmulMConfig(128, 128, 4, 3)),
+            ),
+        ),
+        (2048, 2048): _MatmulShapeConfig(  # Qwen3-1.7B o_proj
+            block_k=32,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 128, 4, 3)),
+                (4, _MatmulMConfig(16, 128, 4, 4)),
+                (8, _MatmulMConfig(16, 64, 4, 5)),
+                (16, _MatmulMConfig(32, 32, 8, 4)),
+                (32, _MatmulMConfig(16, 64, 4, 5)),
+                (64, _MatmulMConfig(32, 32, 4, 5)),
+                (256, _MatmulMConfig(128, 32, 4, 5)),
+                (512, _MatmulMConfig(128, 64, 4, 4)),
+                (1024, _MatmulMConfig(128, 128, 8, 5)),
+                (2048, _MatmulMConfig(128, 256, 8, 4)),
+            ),
+        ),
+        (19456, 2560): _MatmulShapeConfig(  # Qwen3-4B gate_up_proj
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 128, 8, 4)),
+                (4, _MatmulMConfig(16, 128, 8, 3)),
+                (8, _MatmulMConfig(16, 128, 8, 5)),
+                (16, _MatmulMConfig(16, 128, 8, 5)),
+                (32, _MatmulMConfig(32, 128, 4, 4)),
+                (64, _MatmulMConfig(64, 128, 8, 4)),
+                (256, _MatmulMConfig(128, 128, 4, 3)),
+                (512, _MatmulMConfig(64, 128, 4, 5)),
+                (1024, _MatmulMConfig(128, 128, 4, 3)),
+                (2048, _MatmulMConfig(128, 128, 4, 3)),
+            ),
+        ),
+        (2560, 9728): _MatmulShapeConfig(  # Qwen3-4B down_proj
+            block_k=128,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 4, 5)),
+                (4, _MatmulMConfig(16, 64, 4, 4)),
+                (8, _MatmulMConfig(16, 64, 4, 4)),
+                (16, _MatmulMConfig(16, 32, 4, 5)),
+                (32, _MatmulMConfig(32, 32, 4, 3)),
+                (64, _MatmulMConfig(32, 64, 8, 4)),
+                (256, _MatmulMConfig(64, 64, 4, 3)),
+                (512, _MatmulMConfig(64, 128, 4, 3)),
+                (1024, _MatmulMConfig(64, 128, 4, 3)),
+                (2048, _MatmulMConfig(64, 128, 4, 3)),
+            ),
+        ),
+        (6144, 2560): _MatmulShapeConfig(  # Qwen3-4B qkv_proj
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 128, 8, 4)),
+                (4, _MatmulMConfig(16, 256, 4, 3)),
+                (8, _MatmulMConfig(16, 128, 4, 5)),
+                (16, _MatmulMConfig(16, 64, 8, 3)),
+                (32, _MatmulMConfig(16, 128, 8, 4)),
+                (64, _MatmulMConfig(64, 64, 4, 5)),
+                (256, _MatmulMConfig(128, 128, 8, 3)),
+                (512, _MatmulMConfig(64, 64, 4, 5)),
+                (1024, _MatmulMConfig(64, 128, 8, 4)),
+                (2048, _MatmulMConfig(128, 128, 4, 3)),
+            ),
+        ),
+        (151936, 2560): _MatmulShapeConfig(  # Qwen3-4B lm_head
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 8, 5)),
+                (4, _MatmulMConfig(16, 64, 8, 5)),
+                (8, _MatmulMConfig(16, 64, 4, 5)),
+                (16, _MatmulMConfig(16, 64, 4, 5)),
+                (32, _MatmulMConfig(32, 64, 4, 5)),
+                (64, _MatmulMConfig(64, 64, 4, 5)),
+                (256, _MatmulMConfig(128, 128, 4, 3)),
+                (512, _MatmulMConfig(128, 128, 4, 3)),
+                (1024, _MatmulMConfig(128, 128, 8, 3)),
+                (2048, _MatmulMConfig(128, 128, 4, 3)),
+            ),
+        ),
+        (2560, 4096): _MatmulShapeConfig(  # Qwen3-4B o_proj
+            block_k=128,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 8, 3)),
+                (4, _MatmulMConfig(16, 64, 8, 4)),
+                (8, _MatmulMConfig(16, 64, 8, 5)),
+                (16, _MatmulMConfig(16, 64, 4, 4)),
+                (32, _MatmulMConfig(16, 32, 4, 4)),
+                (64, _MatmulMConfig(16, 64, 8, 3)),
+                (256, _MatmulMConfig(64, 64, 8, 3)),
+                (512, _MatmulMConfig(64, 128, 4, 3)),
+                (1024, _MatmulMConfig(64, 128, 4, 3)),
+                (2048, _MatmulMConfig(64, 128, 4, 3)),
+            ),
+        ),
+        (24576, 4096): _MatmulShapeConfig(  # Qwen3-8B gate_up_proj
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 256, 8, 3)),
+                (4, _MatmulMConfig(16, 256, 8, 3)),
+                (8, _MatmulMConfig(16, 256, 8, 3)),
+                (16, _MatmulMConfig(16, 256, 8, 3)),
+                (32, _MatmulMConfig(32, 256, 8, 3)),
+                (64, _MatmulMConfig(64, 256, 4, 3)),
+                (256, _MatmulMConfig(64, 128, 4, 5)),
+                (512, _MatmulMConfig(128, 128, 4, 3)),
+                (1024, _MatmulMConfig(128, 128, 4, 3)),
+                (2048, _MatmulMConfig(128, 128, 4, 3)),
+            ),
+        ),
+        (4096, 12288): _MatmulShapeConfig(  # Qwen3-8B down_proj
+            block_k=128,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 4, 5)),
+                (4, _MatmulMConfig(16, 64, 4, 5)),
+                (8, _MatmulMConfig(16, 64, 4, 4)),
+                (16, _MatmulMConfig(16, 32, 4, 4)),
+                (32, _MatmulMConfig(32, 32, 8, 5)),
+                (64, _MatmulMConfig(32, 64, 4, 5)),
+                (256, _MatmulMConfig(64, 128, 4, 3)),
+                (512, _MatmulMConfig(64, 64, 4, 3)),
+                (1024, _MatmulMConfig(64, 128, 4, 3)),
+                (2048, _MatmulMConfig(64, 128, 4, 3)),
+            ),
+        ),
+        (6144, 4096): _MatmulShapeConfig(  # Qwen3-8B qkv_proj
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 8, 4)),
+                (4, _MatmulMConfig(16, 128, 8, 4)),
+                (8, _MatmulMConfig(16, 128, 4, 4)),
+                (16, _MatmulMConfig(16, 64, 4, 5)),
+                (32, _MatmulMConfig(16, 128, 8, 5)),
+                (64, _MatmulMConfig(64, 64, 4, 4)),
+                (256, _MatmulMConfig(128, 128, 8, 3)),
+                (512, _MatmulMConfig(64, 64, 4, 5)),
+                (1024, _MatmulMConfig(128, 64, 4, 3)),
+                (2048, _MatmulMConfig(128, 128, 4, 3)),
+            ),
+        ),
+        (151936, 4096): _MatmulShapeConfig(  # Qwen3-8B lm_head
+            block_k=64,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 8, 5)),
+                (4, _MatmulMConfig(16, 64, 8, 5)),
+                (8, _MatmulMConfig(16, 64, 4, 5)),
+                (16, _MatmulMConfig(16, 64, 4, 5)),
+                (32, _MatmulMConfig(32, 64, 4, 5)),
+                (64, _MatmulMConfig(32, 128, 8, 5)),
+                (256, _MatmulMConfig(128, 128, 4, 3)),
+                (512, _MatmulMConfig(128, 128, 8, 3)),
+                (1024, _MatmulMConfig(128, 128, 4, 3)),
+                (2048, _MatmulMConfig(128, 128, 4, 3)),
+            ),
+        ),
+        (4096, 4096): _MatmulShapeConfig(  # Qwen3-8B o_proj
+            block_k=128,
+            m_buckets=(
+                (1, _MatmulMConfig(16, 64, 8, 4)),
+                (4, _MatmulMConfig(16, 64, 4, 3)),
+                (8, _MatmulMConfig(16, 32, 4, 4)),
+                (16, _MatmulMConfig(16, 32, 4, 3)),
+                (32, _MatmulMConfig(16, 64, 4, 5)),
+                (64, _MatmulMConfig(64, 32, 8, 4)),
+                (256, _MatmulMConfig(64, 128, 4, 3)),
+                (512, _MatmulMConfig(64, 64, 4, 3)),
+                (1024, _MatmulMConfig(64, 128, 4, 3)),
+                (2048, _MatmulMConfig(64, 128, 4, 3)),
+            ),
+        ),
+    },
 }
 
 _TUNED_MATMUL_CONFIGS_FOR_DEVICE: dict[tuple[int, int], _MatmulShapeConfig] | None = (
@@ -265,6 +492,8 @@ class _MatmulShapeConfig(NamedTuple):
 def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None:
     if capability is None:
         return None
+    if capability.major == 12:
+        return "sm120"
     if capability.major == 10:
         return "blackwell"
     if capability.major == 9:
diff --git a/vllm/model_executor/hw_agnostic/custom_op.py b/vllm/model_executor/hw_agnostic/custom_op.py
index 72d3f8522f06..ecb37681f267 100644
--- a/vllm/model_executor/hw_agnostic/custom_op.py
+++ b/vllm/model_executor/hw_agnostic/custom_op.py
@@ -30,8 +30,7 @@ def maybe_get_oot_by_class(class_type: type) -> type:
 
 
 class PluggableLayer(nn.Module):
-    """
-    Base class for pluggable layers.
+    """Base class for pluggable layers.
 
     A PluggableLayer is a *module-composing* abstraction: it may instantiate other
     ``torch.nn.Module`` objects as sub-layers, and its functionality depends on
@@ -101,8 +100,7 @@ def decorator(layer_cls):
 
 
 class CustomOp(nn.Module):
-    """
-    Base class for custom ops.
+    """Base class for custom ops.
     Dispatches the forward method to the appropriate backend.
     """
 
@@ -172,8 +170,7 @@ def dispatch_forward(self, compile_native: bool):
         return self.forward_native
 
     def maybe_compile(self, fn, *, enable: bool = True):
-        """
-        Compile fn if compilation enabled.
+        """Compile fn if compilation enabled.
         Useful for CustomOp instances called from within a torch custom op,
         meaning the forward call is hidden from the model-level torch.compile.
 
@@ -255,8 +252,7 @@ def enabled(cls) -> bool:
 
     @staticmethod
     def default_on() -> bool:
-        """
-        Behavior controlled by `CompilationConfig.custom_ops`: On by default if
+        """Behavior controlled by `CompilationConfig.custom_ops`: On by default if
         'all', off by default if 'none'.
         When PyTorch Inductor is used, 'none' is the default value,
         otherwise 'all'.
diff --git a/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py b/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
index efc7c418cad6..b9cdc923b14a 100644
--- a/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
+++ b/vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
@@ -66,9 +66,10 @@ def _store_candidates_kernel(
     cols = tl.program_id(1) * TILE + tl.arange(0, TILE)
     value = tl.load(values + row * k + cols, cols < k, other=-float("inf"))
     index = tl.load(indices + row * k + cols, cols < k, other=-1)
+    # NaN scores can occur during warmup; only -inf denotes padding.
     tl.store(
         output + row * out_stride_row + cols * out_stride_col,
-        tl.where(value > -float("inf"), index, -1),
+        tl.where(value != -float("inf"), index, -1),
         cols < OUT_K,
     )
 
diff --git a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py
index 276e1c465f4c..34ce32856ad0 100644
--- a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py
+++ b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py
@@ -12,16 +12,15 @@
 from quack.compile_utils import make_fake_tensor
 
 from vllm.cute_utils import recast_val
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
     CuTeDSLLaunchSpec,
     VllmCuTeDSLJitKernel,
-    cutedsl_kernel_launcher,
 )
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.triton_utils import tl, triton
 
@@ -173,7 +172,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
         if dcp_world_size <= 1:
             return []
         cp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
-        topk = vllm_config.model_config.hf_config.index_topk
+        topk = vllm_config.model_config.hf_text_config.index_topk
         if topk <= 0:
             return []
 
@@ -546,7 +545,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
         dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
         if dcp_world_size <= 1:
             return []
-        topk = vllm_config.model_config.hf_config.index_topk
+        topk = vllm_config.model_config.hf_text_config.index_topk
         if topk <= 0:
             return []
         return self._trace_dispatch(self.dispatch)(
@@ -569,7 +568,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
         )
         return gathered, out
 
-    @cutedsl_kernel_launcher
+    @kernel_launcher
     def __call__(
         self,
         gathered: torch.Tensor,
@@ -577,16 +576,11 @@ def __call__(
         *,
         topk: int,
     ) -> CuTeDSLLaunchSpec[CompileKey]:
-        compile_key = self.dispatch(topk=topk, num_candidates=gathered.shape[1])
-        return (
-            compile_key,
-            (gathered, out),
-            {
-                "gathered_shape": tuple(gathered.shape),
-                "out_shape": tuple(out.shape),
-                "topk": topk,
-            },
+        compile_key = self.dispatch(
+            topk=topk,
+            num_candidates=gathered.shape[1],
         )
+        return compile_key, (gathered, out)
 
 
 _PACK_DCP_TOPK_CANDIDATES_KERNEL = PackDCPTopkCandidatesKernel()
diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py
new file mode 100644
index 000000000000..aeeb2f63b4e9
--- /dev/null
+++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py
@@ -0,0 +1,498 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepGEMM sparse MQA-logits path for the DeepSeek V4.1 two-level indexer.
+
+The candidate-source indexer publishes the top candidate blocks (in units of
+``candidate_block_size`` compressed positions). Instead of scoring every KV
+position and masking the dense logits down to those blocks, this path scores
+only the candidate tokens with DeepGEMM's sparse MQA logits kernels
+(``fp8_fp4_(paged_)sparse_mqa_logits``, DeepGEMM >= 2.8, SM100 only), runs
+DeepSelect's top-k directly on the bf16 logits they produce, and remaps the
+selected columns back to request-local positions.
+
+Only the MXFP4 indexer cache is supported: the kernels require UE8M0-packed
+(granularity-32) Q/KV scales, which is exactly the MXFP4 cache layout. The
+FP8 indexer cache stores one fp32 scale per quant block instead.
+
+The functions here take caller-owned output buffers (``sparse_indices``,
+``end``, ``col_indices``) so the hot path allocates nothing; the indexer
+metadata builder owns them.
+"""
+
+import importlib.util
+
+import torch
+
+from vllm.model_executor.layers.indexer_topk import (
+    deep_select_topk,
+    get_deep_select_stride_requirement,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.deep_gemm import (
+    fp8_fp4_paged_sparse_mqa_logits,
+    fp8_fp4_sparse_mqa_logits,
+    get_paged_sparse_mqa_logits_metadata,
+    get_sparse_mqa_logits_metadata,
+)
+
+# Sparse-block sizes (in tokens) supported by the DeepGEMM kernels.
+SPARSE_BLOCK_KV_CHOICES = (16, 8)
+
+
+def has_deep_select() -> bool:
+    """Whether the DeepSelect top-k extension is built and this GPU runs it."""
+    return (
+        importlib.util.find_spec("vllm._deepselect_C") is not None
+        and current_platform.is_cuda()
+        and current_platform.is_device_capability_family(100)
+    )
+
+
+def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None:
+    """Raise unless DeepSelect accepts the sparse logits layout.
+
+    The sparse kernels return contiguous bf16 ``[rows, num_sparse_cols]``
+    logits and the top-k writes int32 ``[rows, topk_tokens]`` indices;
+    DeepSelect needs both row strides aligned (1024B / 32B on SM100).
+    """
+    in_align, out_align = get_deep_select_stride_requirement()
+    logits_row_bytes = num_sparse_cols * torch.bfloat16.itemsize
+    if num_sparse_cols >= 2**23 or logits_row_bytes % in_align != 0:
+        raise ValueError(
+            f"DeepSelect needs the sparse logits row ({num_sparse_cols} bf16 "
+            f"columns, {logits_row_bytes} B) to be {in_align}B-aligned."
+        )
+    if topk_tokens > 4096 or (topk_tokens * torch.int32.itemsize) % out_align != 0:
+        raise ValueError(
+            f"DeepSelect needs topk <= 4096 with a {out_align}B-aligned int32 "
+            f"output row, got topk={topk_tokens}."
+        )
+
+
+_INT32_BIG = tl.constexpr(2**31 - 1)
+
+
+@triton.jit(do_not_specialize_on_alignment=["end_ptr"])
+def _expand_candidates_kernel(
+    cand_ptr,
+    cand_stride,
+    ks_ptr,
+    ks_stride,
+    ke_ptr,
+    ke_stride,
+    si_ptr,
+    si_stride,
+    end_ptr,
+    end_stride,
+    K: tl.constexpr,
+    S: tl.constexpr,
+    RATIO: tl.constexpr,
+    SBK: tl.constexpr,
+):
+    """One program per row: expand + filter + sort candidate blocks.
+
+    Each row's valid sparse blocks come out sorted ascending (no dedup pass:
+    production candidates are unique, and the kernel tolerates repeats),
+    padded by repeating the last valid block — DeepGEMM's padding convention.
+    ``end`` counts valid sparse-token columns; valid columns form a prefix
+    because block ids are sorted.
+    """
+    row = tl.program_id(0)
+    ks = tl.load(ks_ptr + row * ks_stride)
+    ke = tl.load(ke_ptr + row * ke_stride)
+    cols = tl.arange(0, K)
+    cand = tl.load(cand_ptr + row * cand_stride + cols)  # K is pow2, no mask
+    if RATIO == 2:
+        cand = tl.interleave(cand * RATIO, cand * RATIO + 1)
+
+    # Keep only blocks inside the row's KV range [0, ke - ks); the lower
+    # bound is free since candidates are non-negative.
+    kv_len = tl.maximum(ke - ks, 0)
+    max_block = (kv_len + SBK - 1) // SBK
+    valid = (cand >= 0) & (cand < max_block)
+    vals = tl.where(valid, cand, _INT32_BIG)
+    vals = tl.sort(vals)
+
+    # Anchor at the row's ks rounded down to a sparse-block boundary
+    # (identity for the paged path, where ks == 0).
+    base = ks // SBK
+    si = tl.where(vals != _INT32_BIG, vals, 0) + base
+    n_valid = tl.sum(valid.to(tl.int32))
+    last_rel = tl.max(tl.where(valid, cand, -1))
+    pad_val = tl.where(n_valid > 0, last_rel + base, base)
+    tl.store(
+        si_ptr + row * si_stride + tl.arange(0, S),
+        tl.where(tl.arange(0, S) < n_valid, si, pad_val),
+    )
+
+    # Valid sparse-token column count; only the last block can be partial.
+    # Absolute start of the last block: (base + last_rel) * SBK + ks % SBK,
+    # which is exactly last_rel * SBK + ks.
+    last_start = last_rel * SBK + ks
+    fill = tl.minimum(tl.maximum(ke - last_start, 1), SBK)
+    end = tl.where(n_valid > 0, (n_valid - 1) * SBK + fill, 0)
+    tl.store(end_ptr + row * end_stride, end)
+
+
+@triton.jit
+def _sparse_topk_remap_kernel(
+    col_ptr,
+    col_stride,
+    si_ptr,
+    si_stride,
+    ks_ptr,
+    ks_stride,
+    out_ptr,
+    out_stride,
+    k,
+    width,
+    SBK: tl.constexpr,
+    K_POW2: tl.constexpr,
+):
+    """Remap sparse top-k columns to request-local positions.
+
+    Column ``j * SBK + o`` scores the token at
+    ``(si[row, j] - ks // SBK) * SBK + o``; columns outside ``[0, width)``
+    (including DeepSelect's sentinel for NaN rows) become -1.
+    """
+    row = tl.program_id(0)
+    ks_base = tl.load(ks_ptr + row * ks_stride) // SBK
+    cols = tl.arange(0, K_POW2)
+    c = tl.load(col_ptr + row * col_stride + cols, mask=cols < k, other=-1)
+    valid = (c >= 0) & (c < width)
+    cc = tl.where(valid, c, 0)
+    blocks = tl.load(si_ptr + row * si_stride + cc // SBK)
+    pos = (blocks - ks_base) * SBK + cc % SBK
+    tl.store(out_ptr + row * out_stride + cols, tl.where(valid, pos, -1), mask=cols < k)
+
+
+def pick_sparse_block_kv(candidate_block_size: int) -> int:
+    """Largest sparse-block size (tokens) dividing the candidate block size."""
+    for sbk in SPARSE_BLOCK_KV_CHOICES:
+        if candidate_block_size % sbk == 0:
+            return sbk
+    raise ValueError(
+        f"candidate_block_size={candidate_block_size} is not a multiple of "
+        f"any supported sparse_block_kv {SPARSE_BLOCK_KV_CHOICES}."
+    )
+
+
+def candidate_blocks_to_sparse_indices(
+    candidate_blocks: torch.Tensor,
+    row_ks: torch.Tensor,
+    row_ke: torch.Tensor,
+    candidate_block_size: int,
+    sparse_block_kv: int,
+    out: tuple[torch.Tensor, torch.Tensor] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Expand request-local candidate blocks into DeepGEMM sparse block ids.
+
+    Args:
+        candidate_blocks: [rows, K] int32 request-local candidate block ids
+            (-1 padded), in units of ``candidate_block_size`` positions. K
+            must be a power of two (the in-kernel sort); rows may be strided.
+        row_ks: [rows] int32 per-row K range start; bounds are in the same
+            (packed-workspace) coordinates the sparse kernel iterates over.
+            Pass zeros for the paged path, whose blocks are context-relative.
+        row_ke: [rows] int32 per-row K range end, in the same coordinates as
+            ``row_ks``.
+        candidate_block_size: Positions per candidate block.
+        sparse_block_kv: Positions per sparse block (8 or 16).
+        out: Optional ``(sparse_indices, end)`` buffers to write into,
+            shaped [rows, K * ratio] and [rows], both int32.
+
+    Returns:
+        sparse_indices: [rows, K * ratio] int32 block ids for DeepGEMM,
+            sorted ascending per row with the valid prefix first (unique in
+            production, where candidates are unique) and the remaining slots
+            repeating the last valid block (the kernel's padding convention).
+            With the unaligned-ks variant, block ``i`` starts at
+            ``i * sparse_block_kv + ks % sparse_block_kv``.
+        end: [rows] int32 count of valid sparse-token columns; valid columns
+            always form a prefix because the block ids are sorted.
+
+    """
+    rows, num_candidates = candidate_blocks.shape
+    ratio = candidate_block_size // sparse_block_kv
+    assert ratio * sparse_block_kv == candidate_block_size and ratio in (1, 2), (
+        candidate_block_size,
+        sparse_block_kv,
+    )
+    assert num_candidates & (num_candidates - 1) == 0, num_candidates
+    num_sparse = num_candidates * ratio
+    if out is None:
+        device = candidate_blocks.device
+        out = (
+            torch.empty(rows, num_sparse, dtype=torch.int32, device=device),
+            torch.empty(rows, dtype=torch.int32, device=device),
+        )
+    sparse_indices, end = out
+    assert sparse_indices.shape == (rows, num_sparse) and end.shape == (rows,)
+    assert candidate_blocks.stride(1) == 1 and sparse_indices.stride(1) == 1
+    for t in (candidate_blocks, row_ks, row_ke, sparse_indices, end):
+        assert t.dtype == torch.int32, t.dtype
+
+    _expand_candidates_kernel[(rows,)](
+        candidate_blocks,
+        candidate_blocks.stride(0),
+        row_ks,
+        row_ks.stride(0),
+        row_ke,
+        row_ke.stride(0),
+        sparse_indices,
+        sparse_indices.stride(0),
+        end,
+        end.stride(0),
+        K=num_candidates,
+        S=num_sparse,
+        RATIO=ratio,
+        SBK=sparse_block_kv,
+        num_warps=16,
+    )
+    return sparse_indices, end
+
+
+def sparse_topk_remap(
+    logits: torch.Tensor,
+    sparse_indices: torch.Tensor,
+    end: torch.Tensor,
+    row_ks: torch.Tensor,
+    sparse_block_kv: int,
+    topk_tokens: int,
+    topk_indices: torch.Tensor,
+    *,
+    col_indices: torch.Tensor,
+) -> None:
+    """Row top-k over sparse logits, remapped to request-local positions.
+
+    Column ``j * sparse_block_kv + o`` of ``logits`` scores the token at
+    request-local position
+    ``(sparse_indices[row, j] - ks // sparse_block_kv) * sparse_block_kv + o``.
+
+    DeepSelect consumes the bf16 logits as they come out of the sparse
+    kernels; the valid columns form a prefix (block ids are sorted), so the
+    per-row ``end`` bounds drive it directly. Slots beyond a row's valid
+    count come back as -1.
+
+    Args:
+        logits: [rows, width] bf16 sparse logits, as produced by the sparse
+            kernels.
+        sparse_indices: [rows, S] int32 sparse block ids backing ``logits``.
+        end: [rows] int32 per-row count of valid logit columns.
+        row_ks: [rows] int32 per-row K range start.
+        sparse_block_kv: Positions per sparse block (8 or 16).
+        topk_tokens: Number of tokens to keep per row.
+        topk_indices: [rows, topk_tokens] int32 output buffer.
+        col_indices: [rows, topk_tokens] int32 scratch for the sparse-column
+            top-k result (row stride 32B-aligned for DeepSelect).
+
+    """
+    rows, width = logits.shape
+    assert logits.dtype == torch.bfloat16 and width >= topk_tokens
+    assert end.dtype == torch.int32 and end.shape == (rows,)
+    cols = col_indices[:rows]
+    deep_select_topk(logits, topk_tokens, end=end, output_idx=cols)
+    _sparse_topk_remap_kernel[(rows,)](
+        cols,
+        cols.stride(0),
+        sparse_indices,
+        sparse_indices.stride(0),
+        row_ks,
+        row_ks.stride(0),
+        topk_indices,
+        topk_indices.stride(0),
+        topk_tokens,
+        width,
+        SBK=sparse_block_kv,
+        K_POW2=triton.next_power_of_2(topk_tokens),
+    )
+
+
+def sparse_mqa_logits_prefill_chunk(
+    q: torch.Tensor,
+    q_scale: torch.Tensor,
+    k_quant: torch.Tensor,
+    k_scale: torch.Tensor,
+    weights: torch.Tensor,
+    cu_seqlen_ks: torch.Tensor,
+    cu_seqlen_ke: torch.Tensor,
+    candidate_blocks: torch.Tensor,
+    candidate_block_size: int,
+    sparse_block_kv: int,
+    topk_tokens: int,
+    topk_indices: torch.Tensor,
+    *,
+    sparse_indices: torch.Tensor,
+    end: torch.Tensor,
+    col_indices: torch.Tensor,
+    kernel_metadata: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Sparse-MQA logits + top-k for one prefill chunk (packed KV workspace).
+
+    Args:
+        q: [rows, H, D] packed Q values (MXFP4 viewed as int8).
+        q_scale: [rows, H] int32 packed UE8M0 Q scales.
+        k_quant: [total_kv, D] packed K workspace (MXFP4 viewed as int8).
+        k_scale: [total_kv] int32 packed UE8M0 K scales.
+        weights: [rows, H] bf16 per-head weights; the sparse kernels take
+            bf16 and do not fold the Q scale in.
+        cu_seqlen_ks: [rows] int32 per-token K start bounds in the packed
+            workspace.
+        cu_seqlen_ke: [rows] int32 per-token K end bounds in the packed
+            workspace.
+        candidate_blocks: [rows, K] int32 request-local candidate block ids.
+        candidate_block_size: Positions per candidate block.
+        sparse_block_kv: Positions per sparse block (8 or 16).
+        topk_tokens: Number of tokens to keep per row.
+        topk_indices: [rows, topk_tokens] output buffer.
+        sparse_indices: Caller-owned scratch, see
+            `candidate_blocks_to_sparse_indices`.
+        end: Caller-owned scratch, see `candidate_blocks_to_sparse_indices`.
+        col_indices: Caller-owned scratch, see `sparse_topk_remap`.
+        kernel_metadata: DeepGEMM schedule from a previous call with the same
+            candidates and bounds (i.e. another indexer layer in the same
+            step). When given, the candidate expansion is skipped and
+            ``sparse_indices``/``end`` are assumed to be up to date.
+
+    Returns:
+        The DeepGEMM schedule metadata used, for reuse by later layers.
+
+    """
+    assert weights.dtype == torch.bfloat16, weights.dtype
+    if kernel_metadata is None:
+        candidate_blocks_to_sparse_indices(
+            candidate_blocks,
+            cu_seqlen_ks,
+            cu_seqlen_ke,
+            candidate_block_size,
+            sparse_block_kv,
+            out=(sparse_indices, end),
+        )
+        kernel_metadata = get_sparse_mqa_logits_metadata(
+            cu_seqlen_ks,
+            cu_seqlen_ke,
+            k_quant.shape[0],
+            sparse_indices,
+            q.dtype,
+            sparse_block_kv,
+        )
+    logits = fp8_fp4_sparse_mqa_logits(
+        (q, q_scale),
+        (k_quant, k_scale),
+        weights,
+        kernel_metadata,
+        sparse_indices.shape[1],
+        sparse_block_kv,
+    )
+    sparse_topk_remap(
+        logits,
+        sparse_indices,
+        end,
+        cu_seqlen_ks,
+        sparse_block_kv,
+        topk_tokens,
+        topk_indices,
+        col_indices=col_indices,
+    )
+    return kernel_metadata
+
+
+def sparse_mqa_logits_paged_decode(
+    q: torch.Tensor,
+    q_scale: torch.Tensor,
+    kv_cache: torch.Tensor,
+    weights: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_table: torch.Tensor,
+    row_indices: torch.Tensor,
+    candidate_blocks: torch.Tensor,
+    candidate_block_size: int,
+    sparse_block_kv: int,
+    topk_tokens: int,
+    topk_indices: torch.Tensor,
+    *,
+    row_ks: torch.Tensor,
+    sparse_indices: torch.Tensor,
+    end: torch.Tensor,
+    col_indices: torch.Tensor,
+    kernel_metadata: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Sparse-MQA logits + top-k for paged decode rows.
+
+    Every query is one row (spec-decode rows are flattened by the metadata
+    builder), so ``block_table`` is already per row.
+
+    Args:
+        q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8).
+        q_scale: [rows, 1, H] int32 packed UE8M0 Q scales.
+        kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache
+            view; the page stride must be 512B-aligned.
+        weights: [rows, H] bf16 per-head weights.
+        context_lens: [rows] int32 per-row (compressed) context lengths.
+        block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``.
+        row_indices: [rows] int32 row -> request map (pairing only affects
+            scheduling).
+        candidate_blocks: [rows, K] int32 candidate block ids.
+        candidate_block_size: Positions per candidate block.
+        sparse_block_kv: Positions per sparse block (8 or 16).
+        topk_tokens: Number of tokens to keep per row.
+        topk_indices: [rows, topk_tokens] output buffer.
+        row_ks: [rows] int32 zeros (paged blocks are context-relative).
+        sparse_indices: Caller-owned scratch.
+        end: Caller-owned scratch.
+        col_indices: Caller-owned scratch.
+        kernel_metadata: See `sparse_mqa_logits_prefill_chunk`.
+
+    Returns:
+        The DeepGEMM schedule metadata used, for reuse by later layers.
+
+    """
+    rows = q.shape[0]
+    page_kv = kv_cache.shape[1]
+    assert weights.dtype == torch.bfloat16, weights.dtype
+    assert kv_cache.stride(0) % 512 == 0, (
+        "paged cache page stride must be 512B-aligned: "
+        f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}"
+    )
+    assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_kv)
+    assert block_table.shape[0] == rows and block_table.stride(-1) == 1
+    assert context_lens.shape == (rows,) and context_lens.is_contiguous()
+
+    if kernel_metadata is None:
+        candidate_blocks_to_sparse_indices(
+            candidate_blocks,
+            row_ks,
+            context_lens,
+            candidate_block_size,
+            sparse_block_kv,
+            out=(sparse_indices, end),
+        )
+        kernel_metadata = get_paged_sparse_mqa_logits_metadata(
+            context_lens,
+            block_table,
+            row_indices,
+            page_kv,
+            sparse_indices,
+            q.dtype,
+            sparse_block_kv,
+        )
+    logits = fp8_fp4_paged_sparse_mqa_logits(
+        (q, q_scale),
+        kv_cache,
+        weights,
+        kernel_metadata,
+        sparse_indices.shape[1],
+        sparse_block_kv,
+    )
+    sparse_topk_remap(
+        logits,
+        sparse_indices,
+        end,
+        row_ks,
+        sparse_block_kv,
+        topk_tokens,
+        topk_indices,
+        col_indices=col_indices,
+    )
+    return kernel_metadata
diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py
index 9c52e7b413f7..fc4696f7054b 100644
--- a/vllm/model_executor/kernels/linear/__init__.py
+++ b/vllm/model_executor/kernels/linear/__init__.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-This module re-exports linear kernel implementations to provide a
+"""This module re-exports linear kernel implementations to provide a
 stable import interface during an ongoing reorganization. Upcoming
 PRs will remove the scaled_mm and mixed_precision subdirectories
 and reorganize kernels by provider (aiter, cutlass, flashinfer, etc.)
@@ -27,9 +26,6 @@
     MPLinearKernel,
     MPLinearLayerConfig,
 )
-from vllm.model_executor.kernels.linear.mixed_precision.allspark import (
-    AllSparkLinearKernel,
-)
 from vllm.model_executor.kernels.linear.mixed_precision.conch import (
     ConchLinearKernel,
 )
@@ -102,6 +98,7 @@
 from vllm.model_executor.kernels.linear.mxfp6.emulation import (
     EmulationMxfp6LinearKernel,
 )
+from vllm.model_executor.kernels.linear.mxfp6.humming import HummingMxFp6LinearKernel
 from vllm.model_executor.kernels.linear.mxfp8 import (
     Mxfp8LinearKernel,
     Mxfp8LinearLayerConfig,
@@ -186,6 +183,9 @@
 )
 from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
     CPUFp8BlockScaledMMKernel,
+    CPUFp8PerTensorScaledMMLinearKernel,
+    CPUFp8W8A8BlockScaledMMKernel,
+    CPUFP8W8A8ScaledMMLinearKernel,
     CPUInt8ScaledMMLinearKernel,
 )
 from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -234,14 +234,19 @@
 logger = init_logger(__name__)
 
 
-def _get_linear_backend() -> str:
+def _get_linear_backend(*, quantization: str) -> str:
     """Get the linear_backend setting from the current vllm config."""
     from vllm.config import get_current_vllm_config_or_none
 
-    config = get_current_vllm_config_or_none()
-    if config is not None:
-        return config.kernel_config.linear_backend
-    return "auto"
+    if (config := get_current_vllm_config_or_none()) is None:
+        return "auto"
+    overrides = config.kernel_config.linear_backend_per_quant
+    if overrides is not None and (override := overrides.get(quantization)):
+        logger.info_once(
+            "Applied linear backend override for %r: %r", quantization, override
+        )
+        return override
+    return config.kernel_config.linear_backend
 
 
 # Kernel classes covered by each --linear-backend value.
@@ -288,6 +293,7 @@ def _get_linear_backend() -> str:
         HummingLinearKernel,
         HummingMxfp8LinearKernel,
         HummingMxFp4LinearKernel,
+        HummingMxFp6LinearKernel,
         HummingNvFp4LinearKernel,
     },
     "marlin": {
@@ -368,6 +374,8 @@ def _filter_kernels_by_backend(
 def _resolve_backend_kernels(
     kernels: list[type],
     layer_desc: str,
+    *,
+    quantization: str,
 ) -> list[type]:
     """Apply --linear-backend filtering to one layer type's kernel list.
 
@@ -378,7 +386,7 @@ def _resolve_backend_kernels(
     layer types (e.g. NVFP4 MoE projections next to FP8 attention
     projections).
     """
-    linear_backend = _get_linear_backend()
+    linear_backend = _get_linear_backend(quantization=quantization)
     if linear_backend == "auto":
         return kernels
     filtered = _filter_kernels_by_backend(linear_backend, kernels)
@@ -429,6 +437,8 @@ def _resolve_backend_kernels(
         ChannelWiseTorchFP8ScaledMMLinearKernel,
     ],
     PlatformEnum.CPU: [
+        CPUFP8W8A8ScaledMMLinearKernel,
+        CPUFp8PerTensorScaledMMLinearKernel,
         PerTensorTorchFP8ScaledMMLinearKernel,
         ChannelWiseTorchFP8ScaledMMLinearKernel,
     ],
@@ -460,6 +470,7 @@ def _resolve_backend_kernels(
         TritonFp8BlockScaledMMKernel,
     ],
     PlatformEnum.CPU: [
+        CPUFp8W8A8BlockScaledMMKernel,  # W8A8 preferred; falls back to W8A16 below
         CPUFp8BlockScaledMMKernel,
     ],
     PlatformEnum.XPU: [
@@ -490,7 +501,6 @@ def _resolve_backend_kernels(
     PlatformEnum.CUDA: [
         CutlassW4A8LinearKernel,
         MacheteLinearKernel,
-        AllSparkLinearKernel,
         MarlinLinearKernel,
         ConchLinearKernel,
         ExllamaLinearKernel,
@@ -562,6 +572,7 @@ def _resolve_backend_kernels(
 
 _POSSIBLE_MXFP6_KERNELS: dict[PlatformEnum, list[type[MxFp6LinearKernel]]] = {
     PlatformEnum.CUDA: [
+        HummingMxFp6LinearKernel,
         EmulationMxfp6LinearKernel,
     ],
     PlatformEnum.ROCM: [
@@ -622,9 +633,10 @@ def choose_scaled_mm_linear_kernel(
     possible_kernels: dict[PlatformEnum, list[type[_KernelT]]],
     compute_capability: int | None = None,
     force_kernel: type[_KernelT] | None = None,
+    *,
+    quantization: str,
 ) -> type[_KernelT]:
-    """
-    Choose a _KernelT that can implement the given config for the
+    """Choose a _KernelT that can implement the given config for the
     given compute capability. Attempts to choose the best kernel in terms of
     performance.
 
@@ -639,14 +651,15 @@ def choose_scaled_mm_linear_kernel(
         force_kernel (Optional[type[_KernelT]]): An Optional forced kernel to override
             the possible_kernels if it can be implemented. If None, it will only try the
             possible kernels.
+        quantization: Quantization scheme used to select a backend override.
 
     Raises:
         ValueError: If no kernel can implement the given config.
 
     Returns:
         _KernelT: Chosen kernel.
-    """
 
+    """
     failure_reason_list = []
 
     if force_kernel is not None:
@@ -665,7 +678,11 @@ def choose_scaled_mm_linear_kernel(
     platform_kernels = possible_kernels.get(current_platform._enum, [])
 
     # Apply --linear-backend filtering when set.
-    platform_kernels = _resolve_backend_kernels(platform_kernels, "scaled-mm")
+    platform_kernels = _resolve_backend_kernels(
+        platform_kernels,
+        "scaled-mm",
+        quantization=quantization,
+    )
 
     for kernel in platform_kernels:
         is_supported_and_can_implement, failure_reason = (
@@ -702,6 +719,7 @@ def init_fp8_linear_kernel(
         kernel_type = choose_scaled_mm_linear_kernel(
             config=scaled_mm_linear_kernel_config,
             possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS,  # type: ignore[misc]
+            quantization="fp8_block_w8a8",
             force_kernel=force_kernel,
         )
         if module_name:
@@ -733,6 +751,7 @@ def init_fp8_linear_kernel(
         kernel_type = choose_scaled_mm_linear_kernel(
             config=scaled_mm_linear_kernel_config,
             possible_kernels=_POSSIBLE_FP8_KERNELS,  # type: ignore[arg-type]
+            quantization="fp8_w8a8",
             force_kernel=force_kernel,
         )
         if module_name:
@@ -769,6 +788,7 @@ def init_int8_linear_kernel(
     kernel_type = choose_scaled_mm_linear_kernel(
         config,
         _POSSIBLE_INT8_KERNELS,
+        quantization="int8_w8a8",
     )
 
     logger.info_once(
@@ -793,8 +813,7 @@ def init_int8_linear_kernel(
 def choose_mp_linear_kernel(
     config: MPLinearLayerConfig, compute_capability: int | None = None
 ) -> type[MPLinearKernel]:
-    """
-    Choose an MPLinearKernel that can implement the given config for the given
+    """Choose an MPLinearKernel that can implement the given config for the given
      compute capability. Attempts to choose the best kernel in terms of
      performance.
 
@@ -810,6 +829,7 @@ def choose_mp_linear_kernel(
 
     Returns:
         type[MPLinearKernel]: Chosen kernel.
+
     """
     if compute_capability is None:
         if current_platform is None:
@@ -821,7 +841,11 @@ def choose_mp_linear_kernel(
     platform_kernels = _POSSIBLE_KERNELS.get(current_platform._enum, [])
 
     # Apply --linear-backend filtering when set.
-    platform_kernels = _resolve_backend_kernels(platform_kernels, "mixed-precision")
+    platform_kernels = _resolve_backend_kernels(
+        platform_kernels,
+        "mixed-precision",
+        quantization="mixed_precision",
+    )
 
     failure_reasons = []
     for kernel in platform_kernels:
@@ -843,6 +867,7 @@ def choose_mp_linear_kernel(
 
         can_implement, failure_reason = kernel.can_implement(config)
         if can_implement:
+            logger.info_once("Using %s for mixed-precision linear", kernel.__name__)
             return kernel
         else:
             failure_reasons.append(
@@ -872,7 +897,11 @@ def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8Linea
         possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, []))
 
     # Apply --linear-backend filtering when set.
-    possible = _resolve_backend_kernels(possible, "MXFP8")
+    possible = _resolve_backend_kernels(
+        possible,
+        "MXFP8",
+        quantization="mxfp8",
+    )
 
     failure_reasons = []
     for kernel_cls in possible:
@@ -914,7 +943,11 @@ def init_mxfp4_linear_kernel(
     possible = list(_POSSIBLE_MXFP4_KERNELS.get(platform, []))
 
     # Apply --linear-backend filtering when set.
-    possible = _resolve_backend_kernels(possible, "MXFP4")
+    possible = _resolve_backend_kernels(
+        possible,
+        "MXFP4",
+        quantization="mxfp4",
+    )
 
     failure_reasons = []
     for kernel_cls in possible:
@@ -954,7 +987,7 @@ def init_mxfp6_linear_kernel(
         activation_quant_key=activation_quant_key,
     )
 
-    linear_backend = _get_linear_backend()
+    linear_backend = _get_linear_backend(quantization="mxfp6")
 
     platform = current_platform._enum
     possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, []))
@@ -1014,7 +1047,10 @@ def init_wfp8_a16_linear_kernel(
     )
 
     kernel_type = choose_scaled_mm_linear_kernel(
-        config, _POSSIBLE_WFP8A16_KERNELS, force_kernel=force_kernel
+        config,
+        _POSSIBLE_WFP8A16_KERNELS,
+        quantization="w8a16_fp8",
+        force_kernel=force_kernel,
     )
 
     if module_name:
@@ -1045,7 +1081,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
     # batch-invariant CUTLASS implementation when available, otherwise fall
     # back to emulation. It overrides --linear-backend.
     force_kernel: type[NvFp4LinearKernel] | None = None
-    linear_backend = _get_linear_backend()
+    quantization = "nvfp4_w4a16" if use_a16 else "nvfp4_w4a4"
+    linear_backend = _get_linear_backend(quantization=quantization)
     if envs.VLLM_BATCH_INVARIANT:
         bi_supported, reason = CutlassNvFp4LinearKernel.is_supported()
         if bi_supported:
@@ -1107,7 +1144,11 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
         possible = [kernel for kernel in possible if kernel in a16_kernels]
 
     # Apply --linear-backend filtering when set.
-    possible = _resolve_backend_kernels(possible, "NVFP4")
+    possible = _resolve_backend_kernels(
+        possible,
+        "NVFP4",
+        quantization=quantization,
+    )
 
     failure_reasons = []
     for kernel_cls in possible:
@@ -1151,8 +1192,7 @@ def register_linear_kernel(
     platform: PlatformEnum,
     kernel_type: str = "mp",
 ) -> None:
-    """
-    Register a new linear kernel class to be considered in kernel selection.
+    """Register a new linear kernel class to be considered in kernel selection.
 
     Args:
         kernel_class (type): The kernel class to register.
@@ -1162,6 +1202,7 @@ def register_linear_kernel(
 
     Raises:
         ValueError: If the kernel_type is not recognized.
+
     """
     if kernel_type == "mp":
         if platform not in _POSSIBLE_KERNELS:
@@ -1227,7 +1268,6 @@ def register_linear_kernel(
     "ZentorchWNA16LinearKernel",
     "MPLinearKernel",
     "MPLinearLayerConfig",
-    "AllSparkLinearKernel",
     "ConchLinearKernel",
     "CPUWNA16LinearKernel",
     "CutlassW4A8LinearKernel",
@@ -1252,6 +1292,7 @@ def register_linear_kernel(
     "MxFp6LinearLayerConfig",
     "init_mxfp6_linear_kernel",
     "EmulationMxfp6LinearKernel",
+    "HummingMxFp6LinearKernel",
     "AiterMxfp4LinearKernel",
     "EmulationMxfp4LinearKernel",
     "FlashInferMxFp4LinearKernel",
diff --git a/vllm/model_executor/kernels/linear/base.py b/vllm/model_executor/kernels/linear/base.py
index 036b30b7c46d..0db595a3ca42 100644
--- a/vllm/model_executor/kernels/linear/base.py
+++ b/vllm/model_executor/kernels/linear/base.py
@@ -50,6 +50,7 @@ class Params:
         # Access typed parameters
         output = func(input, params.weight, params.weight_scale)
         ```
+
     """
 
     weight: torch.Tensor
@@ -72,7 +73,7 @@ def from_layer(cls, layer: torch.nn.Module) -> Self:
 
 @dataclass
 class FP8Params(Params):
-    """FP8 layer parameters with typed fields"""
+    """FP8 layer parameters with typed fields."""
 
     input_scale_ub: torch.Tensor | None
 
@@ -80,7 +81,7 @@ class FP8Params(Params):
 
     @classmethod
     def from_layer(cls, layer: torch.nn.Module) -> "FP8Params":
-        """Extract parameters from layer"""
+        """Extract parameters from layer."""
         return cls(
             weight=getattr(layer, cls.WEIGHT),
             weight_scale=getattr(layer, cls.WEIGHT_SCALE),
@@ -163,6 +164,7 @@ def apply_weights(self, layer, x, bias=None, **kwargs):
         2. Initialization: __init__() creates kernel instance with config
         3. Weight loading: process_weights_after_loading() preprocesses weights
         4. Inference: apply_weights() executes the quantized matmul
+
     """
 
     @classmethod
@@ -184,6 +186,7 @@ def is_supported(
             A tuple of (is_supported, reason):
                 - is_supported: True if the kernel can run on this hardware
                 - reason: If not supported, a string explaining why; otherwise None
+
         """
         raise NotImplementedError
 
@@ -204,7 +207,7 @@ def can_implement(cls, config: _ConfigT) -> tuple[bool, str | None]:
             A tuple of (can_implement, reason):
                 - can_implement: True if this kernel supports the config
                 - reason: If not supported, a string explaining why; otherwise None
-            ```
+
         """
         raise NotImplementedError
 
@@ -214,6 +217,7 @@ def __init__(self, config: _ConfigT) -> None:
         Args:
             config: Kernel-specific configuration containing settings like
                    quantization keys, output dtypes, etc.
+
         """
         self.config = config
 
@@ -245,6 +249,7 @@ def process_weights_after_loading(self, layer):
                 weight_reordered = reorder_weights(params.weight)
                 replace_parameter(layer, params.WEIGHT, weight_reordered)
             ```
+
         """
         raise NotImplementedError
 
@@ -270,6 +275,7 @@ def _get_layer_params(self, layer: torch.nn.Module, **kwargs: Any) -> _ParamsT:
             def _get_layer_params(self, layer, **kwargs):
                 return MyKernelParams.from_layer(layer)
             ```
+
         """
         raise NotImplementedError
 
@@ -281,6 +287,7 @@ def get_output_padding(self) -> int | None:
 
         Returns:
             Number of tokens to pad, or None for no padding (default)
+
         """
         return None
 
@@ -306,5 +313,6 @@ def apply_weights(
 
         Returns:
             Output tensor of shape [..., out_features]
+
         """
         raise NotImplementedError
diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py b/vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
similarity index 70%
rename from vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py
rename to vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
index f8a90cc1abce..ec40fb056a13 100644
--- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py
+++ b/vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py
@@ -1,18 +1,36 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""SM100 BF16 GEMM with fused tensor-parallel reduce-scatter/all-reduce."""
+"""SM100 BF16/MXFP8 GEMM with fused tensor-parallel reduce-scatter/all-reduce.
+
+Model-neutral: any row-parallel projection whose output is reduced across TP
+(Kimi-K3 ``o_proj``/``down_proj``, DeepSeek-V4.1 ``wo_b``) can bind to the
+process-wide workspace through ``maybe_init_gemm_rs_ar`` + ``GemmRsAr.apply``.
+"""
 
 # Based on CUTLASS's Blackwell distributed GEMM-RS example at dcf215a.
 # See https://github.com/NVIDIA/cutlass/issues/3117 for memory semantics.
+# The MXFP8 path follows the block-scaled SM100 GEMM template from
+# https://github.com/gau-nernst/gn-kernels (sm100_mm_mxfp8.py).
 
 from functools import cache
+from typing import TYPE_CHECKING
 
 import cutlass
 import torch
 import torch.distributed._symmetric_memory as symm_mem
 from cuda.bindings.driver import CUstream
-from cutlass import BFloat16, Int32, Int64, Uint16, cute, utils
+from cutlass import (
+    BFloat16,
+    Float8E4M3FN,
+    Float8E8M0FNU,
+    Int32,
+    Int64,
+    Uint16,
+    Uint64,
+    cute,
+    utils,
+)
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import llvm, nvvm, vector
 from cutlass.cute.nvgpu import cpasync, tcgen05
@@ -24,6 +42,10 @@
 from vllm.distributed import get_tp_group
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+from vllm.platforms import current_platform
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
 
 logger = init_logger(__name__)
 
@@ -88,24 +110,31 @@ def multimem_st_16B(dst: cute.Tensor, value: cute.Tensor, *, loc=None, ip=None)
     )
 
 
-class Sm100GemmRsArBF16:
+class Sm100GemmRsAr:
     def __init__(
         self,
         rank: int,
         num_ranks: int,
-        BN: int = 128,
         cta_group: int = 1,
         all_reduce: bool = False,
+        dtype: str = "bf16",
     ) -> None:
+        assert dtype in ("bf16", "mxfp8")
         self.rank = rank
         self.num_ranks = num_ranks
-        BM, BK = 128, 64
+        BM = 128
+        BN = 128
+        BK = {"bf16": 64, "mxfp8": 128}[dtype]  # 128B
         self.cta_tile = (BM, BN, BK)
         self.cta_group = cta_group
         self.all_reduce = all_reduce
+        self.dtype = dtype
 
         smem_bytes = get_smem_capacity_in_bytes()
-        self.stage_size = (BM + (BN // cta_group)) * BK * 2
+        if dtype == "bf16":
+            self.stage_size = (BM + (BN // cta_group)) * BK * 2
+        elif dtype == "mxfp8":
+            self.stage_size = (BM + (BN // cta_group)) * BK + (BM + BN) * (BK // 32)
         self.num_stages = smem_bytes // self.stage_size
 
     @cute.jit
@@ -124,11 +153,36 @@ def prepare_tma(
         layout = cute.make_composed_layout(swizzle_128b, 0, layout)
         return cpasync.make_tiled_tma_atom(tma_op, tensor, layout, (BM, BK))
 
+    @cute.jit
+    def prepare_sf_tma(
+        self,
+        SF: cute.Tensor,
+        R: Int32,
+        Ksf: Int32,
+        BR: cutlass.Constexpr,
+        BKsf: cutlass.Constexpr,
+    ) -> cpasync.TmaInfo:
+        tma_group = (
+            tcgen05.CtaGroup.TWO if self.cta_group == 2 else tcgen05.CtaGroup.ONE
+        )
+        tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=tma_group)
+        # NVIDIA SF gmem layout: an [R, Ksf] scale grid padded to R % 128 == 0
+        # and permuted as [R/128, Ksf/4, 32, 4, 4]. The trailing [32, 4, 4]
+        # atom is the opaque 512B block tcgen05 expects, so treat it as bytes.
+        # The Int64 recast keeps the TMA boxDim <= 256.
+        g_layout = cute.make_layout((Ksf * 128, cute.ceil_div(R, 128)))
+        SF_i64 = cute.recast_tensor(cute.make_tensor(SF.iterator, g_layout), Int64)
+        tiler = (128 * BKsf // 8, BR // 128)
+        layout = cute.make_layout((*tiler, self.num_stages))
+        return cpasync.make_tiled_tma_atom(tma_op, SF_i64, layout, tiler)
+
     @cute.jit
     def __call__(
         self,
         A: cute.Tensor,
         B: cute.Tensor,
+        SFA: cute.Tensor,
+        SFB: cute.Tensor,
         partial_uc: cute.Tensor,
         partial_mc_ptr: cute.Pointer,
         output: cute.Tensor,
@@ -138,10 +192,17 @@ def __call__(
         grid_size: Int32,
         stream: CUstream,
     ) -> None:
+        M, K = A.shape
         N = B.shape[0]
         BM, BN, BK = self.cta_tile
         A_tma = self.prepare_tma(A, BM, BK)
         B_tma = self.prepare_tma(B, BN // self.cta_group, BK)
+        if cutlass.const_expr(self.dtype == "mxfp8"):
+            SFA_tma = self.prepare_sf_tma(SFA, M, K // 32, BM, BK // 32)
+            SFB_tma = self.prepare_sf_tma(SFB, N, K // 32, BN, BK // 32)
+        else:
+            SFA_tma = None
+            SFB_tma = None
         padded_M = partial_uc.shape[0]
         partial_mc = cute.make_tensor(
             partial_mc_ptr,
@@ -158,6 +219,8 @@ def __call__(
         self.kernel(
             A_tma,
             B_tma,
+            SFA_tma,
+            SFB_tma,
             partial_uc,
             partial_mc,
             output,
@@ -171,6 +234,8 @@ def kernel(
         self,
         A_tma: cpasync.TmaInfo,
         B_tma: cpasync.TmaInfo,
+        SFA_tma: cpasync.TmaInfo,
+        SFB_tma: cpasync.TmaInfo,
         partial_uc: cute.Tensor,
         partial_mc: cute.Tensor,
         output: cute.Tensor,
@@ -190,25 +255,33 @@ def kernel(
 
         is_2cta = cta_group == 2
         cta_rank = raw_bid % self.cta_group
-        num_tmem_stages = 512 // BN
+
+        # MXFP8 reserves the last 16 tmem columns (496-512) for the A/B scale
+        # factors.
+        num_tmem_stages = (496 if self.dtype == "mxfp8" else 512) // BN
 
         smem = utils.SmemAllocator()
         sA = smem.allocate_tensor(
-            BFloat16,
+            A_tma.atom.value_type,
             A_tma.smem_layout.outer,
             byte_alignment=128,
             swizzle=A_tma.smem_layout.inner,
         )
         sB = smem.allocate_tensor(
-            BFloat16,
+            B_tma.atom.value_type,
             B_tma.smem_layout.outer,
             byte_alignment=128,
             swizzle=B_tma.smem_layout.inner,
         )
+        if cutlass.const_expr(self.dtype == "mxfp8"):
+            sSFA = smem.allocate_tensor(Int64, SFA_tma.smem_layout, byte_alignment=128)
+            sSFB = smem.allocate_tensor(Int64, SFB_tma.smem_layout, byte_alignment=128)
+
         tma_full_mbar = smem.allocate_array(Int64, num_stages)
         tma_empty_mbar = smem.allocate_array(Int64, num_stages)
         tmem_full_mbar = smem.allocate_array(Int64, num_tmem_stages)
         tmem_empty_mbar = smem.allocate_array(Int64, num_tmem_stages)
+
         taddr = smem.allocate(Int32, 4)
 
         # Named barriers
@@ -238,6 +311,9 @@ def kernel(
         elif warp_id == 1:
             cpasync.prefetch_descriptor(A_tma.atom)
             cpasync.prefetch_descriptor(B_tma.atom)
+            if cutlass.const_expr(self.dtype == "mxfp8"):
+                cpasync.prefetch_descriptor(SFA_tma.atom)
+                cpasync.prefetch_descriptor(SFB_tma.atom)
 
         if cutlass.const_expr(is_2cta):
             cute.arch.cluster_arrive_relaxed()
@@ -261,12 +337,22 @@ def kernel(
             # [(BM, BK), (M/BM, K/BK)]
             gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK))
             gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN // cta_group, BK))
+            if cutlass.const_expr(self.dtype == "mxfp8"):
+                BKsf = BK // 32
+                gSFA_tiles = cute.zipped_divide(
+                    SFA_tma.tma_tensor, (128 * BKsf // 8, BM // 128)
+                )
+                gSFB_tiles = cute.zipped_divide(
+                    SFB_tma.tma_tensor, (128 * BKsf // 8, BN // 128)
+                )
 
             for bid in range(raw_bid, total_tiles, num_bids):
                 bid_m = bid % grid_m
                 bid_n = bid // grid_m
                 if cutlass.const_expr(cta_group == 2):
-                    bid_n = bid_n * cta_group + cta_rank
+                    bid_n_b = bid_n * cta_group + cta_rank
+                else:
+                    bid_n_b = bid_n
 
                 for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1):
                     mbar = tma_full_mbar_ + tma_stage
@@ -282,10 +368,23 @@ def kernel(
                     )
                     simple_tma_copy(
                         B_tma.atom,
-                        gB_tiles[None, (bid_n, iter_k)],
+                        gB_tiles[None, (bid_n_b, iter_k)],
                         sB[None, None, tma_stage],
                         mbar,
                     )
+                    if cutlass.const_expr(self.dtype == "mxfp8"):
+                        simple_tma_copy(
+                            SFA_tma.atom,
+                            gSFA_tiles[None, (iter_k, bid_m)],
+                            sSFA[None, None, tma_stage],
+                            mbar,
+                        )
+                        simple_tma_copy(
+                            SFB_tma.atom,
+                            gSFB_tiles[None, (iter_k, bid_n)],
+                            sSFB[None, None, tma_stage],
+                            mbar,
+                        )
 
                     tma_stage = (tma_stage + 1) % num_stages
                     if tma_stage == 0:
@@ -303,8 +402,15 @@ def kernel(
 
                 MMA_M = BM * cta_group
                 MMA_N = BN
-                idesc = _tcgen05.make_bf16_idesc(MMA_M, MMA_N)
                 sdesc = _tcgen05.make_sdesc_128B_swizzle(0)
+                if cutlass.const_expr(self.dtype == "mxfp8"):
+                    # SF smem descriptor: no swizzle, SBO = 8 * 16 bytes.
+                    sdesc_sf = Uint64((8 << 32) | (1 << 46))
+                    sfa_tmem = 496
+                    sfb_tmem = sfa_tmem + 4
+                    idesc = _tcgen05.make_mxfp8_idesc(MMA_M, MMA_N)
+                else:
+                    idesc = _tcgen05.make_bf16_idesc(MMA_M, MMA_N)
                 multicast_mask = Uint16((1 << self.cta_group) - 1)
 
                 for bid in range(raw_bid, total_tiles, num_bids):
@@ -319,17 +425,54 @@ def kernel(
                         b_addr = sB[None, None, tma_stage].iterator.toint()
                         a_desc = sdesc | (a_addr >> 4)
                         b_desc = sdesc | (b_addr >> 4)
+                        if cutlass.const_expr(self.dtype == "mxfp8"):
+                            sfa_desc = sdesc_sf | (
+                                sSFA[None, None, tma_stage].iterator.toint() >> 4
+                            )
+                            sfb_desc = sdesc_sf | (
+                                sSFB[None, None, tma_stage].iterator.toint() >> 4
+                            )
 
                         cute.arch.mbarrier_wait(
                             tma_full_mbar + tma_stage, tma_full_parity
                         )
                         _tcgen05.fence_after_thread_sync()
 
-                        for mma_k in cutlass.range_constexpr(BK // 16):
-                            enable_d = iter_k > 0 or mma_k > 0
-                            _tcgen05.mma_f16(
-                                d_tmem, a_desc, b_desc, idesc, enable_d, cta_group
+                        # load SF from smem->tmem
+                        if cutlass.const_expr(self.dtype == "mxfp8"):
+                            _tcgen05.cp(
+                                sfa_tmem, sfa_desc, "32x128b", "warpx4", cta_group
                             )
+                            for j in cutlass.range_constexpr(BN // 128):
+                                _tcgen05.cp(
+                                    sfb_tmem + j * 4,
+                                    sfb_desc,
+                                    "32x128b",
+                                    "warpx4",
+                                    cta_group,
+                                )
+                                sfb_desc += (128 * 4) >> 4
+
+                        # 4 = 128B / 32B
+                        for mma_k in cutlass.range_constexpr(4):
+                            enable_d = iter_k > 0 or mma_k > 0
+                            if cutlass.const_expr(self.dtype == "mxfp8"):
+                                # select SF IDs
+                                idesc_k = idesc + ((mma_k << 4) | (mma_k << 29))
+                                _tcgen05.mma_mxfp8(
+                                    d_tmem,
+                                    a_desc,
+                                    b_desc,
+                                    idesc_k,
+                                    sfa_tmem,
+                                    sfb_tmem,
+                                    enable_d,
+                                    cta_group,
+                                )
+                            else:
+                                _tcgen05.mma_f16(
+                                    d_tmem, a_desc, b_desc, idesc, enable_d, cta_group
+                                )
                             a_desc += 32 >> 4
                             b_desc += 32 >> 4
                         _tcgen05.commit(
@@ -638,9 +781,9 @@ def release_gemm_tile(
     def compile(
         rank: int,
         num_ranks: int,
-        BN: int,
         cta_group: int,
         all_reduce: bool = False,
+        dtype: str = "bf16",
     ):
         M = cute.sym_int()
         padded_M = cute.sym_int()
@@ -649,12 +792,37 @@ def compile(
         local_M = cute.sym_int()
         num_flags = cute.sym_int()
 
+        ab_dtype = Float8E4M3FN if dtype == "mxfp8" else BFloat16
+        # Row stride must keep rows 16B-aligned: 8 BF16 or 16 FP8 elements.
+        row_div = 16 if dtype == "mxfp8" else 8
         A = make_fake_tensor(
-            BFloat16, (M, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16
+            ab_dtype,
+            (M, K),
+            (cute.sym_int64(divisibility=row_div), 1),
+            assumed_align=16,
         )
         B = make_fake_tensor(
-            BFloat16, (N, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16
+            ab_dtype,
+            (N, K),
+            (cute.sym_int64(divisibility=row_div), 1),
+            assumed_align=16,
         )
+        if dtype == "mxfp8":
+            SFA = make_fake_tensor(
+                Float8E8M0FNU,
+                (cute.sym_int(divisibility=512),),
+                (1,),
+                assumed_align=16,
+            )
+            SFB = make_fake_tensor(
+                Float8E8M0FNU,
+                (cute.sym_int(divisibility=512),),
+                (1,),
+                assumed_align=16,
+            )
+        else:
+            SFA = None
+            SFB = None
         partial = make_fake_tensor(
             BFloat16,
             (padded_M, N),
@@ -677,13 +845,15 @@ def compile(
         peer_flag_ptr = nullptr(Int64, cute.AddressSpace.gmem, assumed_align=8)
 
         stream = make_fake_stream(use_tvm_ffi_env_stream=True)
-        kernel = Sm100GemmRsArBF16(
-            rank, num_ranks, BN, cta_group, all_reduce=all_reduce
+        kernel = Sm100GemmRsAr(
+            rank, num_ranks, cta_group, all_reduce=all_reduce, dtype=dtype
         )
         return cute.compile(
             kernel,
             A,
             B,
+            SFA,
+            SFB,
             partial,
             partial_mc_ptr,
             output,
@@ -697,7 +867,7 @@ def compile(
 
 
 class GemmRsAr:
-    """Own the symmetric workspace for Kimi-K3 GEMM-RS/AR launches.
+    """Own the symmetric workspace for GEMM-RS/AR launches.
 
     All TP ranks must belong to one NVLink domain for multimem instructions.
 
@@ -727,6 +897,7 @@ def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None:
         self.N = N
         self.device = device
         self.all_reduce = all_reduce
+        self.dtypes: set[str] = set()
 
         self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device)
         self.partial_handle = symm_mem.rendezvous(self.partial, group)
@@ -766,21 +937,40 @@ def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None:
         tp_group.barrier()
 
     def can_run(self, linear: LinearBase) -> bool:
-        # Validate projection-invariant requirements once during model init.
-        # only supports BF16 for now
-        if not isinstance(linear.quant_method, UnquantizedLinearMethod):
-            return False
+        from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+            FlashInferCutedslMxfp8LinearKernel,
+            FlashInferCutlassMxfp8LinearKernel,
+        )
+
+        # Called before weight loading; online quantization starts on meta.
+        if isinstance(linear.quant_method, UnquantizedLinearMethod):
+            dtype, k_alignment = "bf16", 64
+            if linear.weight.dtype != torch.bfloat16:
+                return False
+        else:
+            method = getattr(linear, "scheme", linear.quant_method)
+            kernel = getattr(method, "kernel", None)
+            if not isinstance(
+                kernel,
+                (
+                    FlashInferCutedslMxfp8LinearKernel,
+                    FlashInferCutlassMxfp8LinearKernel,
+                ),
+            ):
+                return False
+            dtype, k_alignment = "mxfp8", 128
+
         w = linear.weight
-        if w.ndim != 2:
+        if (
+            w.ndim != 2
+            or w.shape[0] != self.N
+            or w.shape[1] % k_alignment != 0
+            or (w.device != self.device and not w.is_meta)
+            or not w.is_contiguous()
+        ):
             return False
-        K = w.shape[1]
-        return (
-            w.shape == (self.N, K)
-            and K % 64 == 0
-            and w.dtype == torch.bfloat16
-            and w.device == self.device
-            and w.is_contiguous()
-        )
+        self.dtypes.add(dtype)
+        return True
 
     def warn_incompatible_projection(self) -> None:
         logger.warning_once(
@@ -794,18 +984,51 @@ def should_run(self, x: torch.Tensor) -> bool:
         # supported but faster on the existing LL path.
         return x.shape[0] >= 128
 
-    def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+    def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor:
+        from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+            FlashInferCutedslMxfp8LinearKernel,
+        )
+
+        method = getattr(linear, "scheme", linear.quant_method)
+        w = linear.weight
+        if isinstance(
+            getattr(method, "kernel", None), FlashInferCutedslMxfp8LinearKernel
+        ):
+            # This backend stores a column-major [K, N] view after loading.
+            w = w.t()
+        w_sf = getattr(linear, "weight_scale", None)
+
         assert x.ndim == 2
         M, K = x.shape
         assert 0 < M <= self.max_M
-        assert w.shape == (self.N, K) and K % 64 == 0
-        assert w.dtype == torch.bfloat16
-        assert w.device == self.device
-        assert w.is_contiguous()
         assert x.dtype == torch.bfloat16
         assert x.device == self.device
         assert x.is_contiguous()
-        N = w.shape[0]
+        assert w.device == self.device
+        N = self.N
+        dtype = "mxfp8" if w.dtype == torch.float8_e4m3fn else "bf16"
+        x_sf = None
+        if dtype == "mxfp8":
+            from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+                mxfp8_e4m3_quantize,
+            )
+
+            assert w_sf is not None
+            assert w_sf.dtype == torch.uint8 and w_sf.device == self.device
+            assert w_sf.ndim == 1 and w_sf.is_contiguous()
+            assert w_sf.numel() == N * (K // 32)
+            assert w.shape == (N, K) and K % 128 == 0
+            assert w.is_contiguous()
+            x, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+            # The kernel declares the scales as E8M0; the stored/quantized
+            # tensors carry the same bits as uint8.
+            x_sf = x_sf.view(torch.float8_e8m0fnu)
+            w_sf = w_sf.view(torch.float8_e8m0fnu)
+        else:
+            assert w_sf is None
+            assert w.shape == (N, K) and K % 64 == 0
+            assert w.dtype == torch.bfloat16
+            assert w.is_contiguous()
         padded_M = (M + self.world_size - 1) // self.world_size
         padded_M *= self.world_size
         local_M = padded_M // self.world_size
@@ -814,10 +1037,8 @@ def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
         # Avoid padding small odd grids; 2-CTA wins consistently for M >= 1024.
         cta_group = 2 if M >= 1024 or grid_m % 2 == 0 else 1
         grid_m = (grid_m + cta_group - 1) // cta_group * cta_group
-        BN = 256 if M * K >= 24 * 1024 * 1024 else 128
-        assert N % BN == 0
 
-        num_tiles = grid_m * (N // BN)
+        num_tiles = grid_m * (N // 128)
         num_ctas = min(num_tiles, self.num_sms)
         num_ctas = num_ctas // cta_group * cta_group
         assert self.flags.numel() >= num_tiles + num_ctas
@@ -825,16 +1046,26 @@ def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
         output = None
         if not self.all_reduce:
             output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device)
-        compiled = Sm100GemmRsArBF16.compile(
+            # The kernel never touches rows at or beyond M. Callers that hand
+            # in an unpadded M (DeepSeek-V4.1 slices the SP all-gather back to
+            # the real token count) expect the same zero padding rows that
+            # ``sp_reduce_scatter`` produces, so clear them here. This is a
+            # no-op whenever M is already a multiple of the TP size.
+            valid_rows = min(max(M - self.rank * local_M, 0), local_M)
+            if valid_rows < local_M:
+                output[valid_rows:].zero_()
+        compiled = Sm100GemmRsAr.compile(
             self.rank,
             self.world_size,
-            BN,
             cta_group,
             self.all_reduce,
+            dtype=dtype,
         )
         compiled(
             x,
             w,
+            x_sf,
+            w_sf,
             self.partial[:padded_M],
             self.partial_mc_ptr,
             output,
@@ -871,22 +1102,74 @@ def init_gemm_rs_ar(max_M: int, N: int, *, all_reduce: bool = False) -> None:
     _gemm_rs_ar = GemmRsAr(max_M=max_M, N=N, all_reduce=all_reduce)
 
 
+def maybe_init_gemm_rs_ar(
+    vllm_config: "VllmConfig", *, N: int, all_reduce: bool
+) -> bool:
+    """Collectively initialize the mode-bound GEMM-RS/AR state if supported.
+
+    Callers gate on their own feature flag first; this checks the worker
+    topology (SM100-family CUDA, BF16 model dtype, no ubatching, a TP size in
+    2-16 that divides 128) and the NVLink multicast rendezvous. Returns whether
+    projections may bind to the kernel through ``get_gemm_rs_ar().can_run``.
+    """
+    mode = "GEMM-AR" if all_reduce else "GEMM-RS"
+    parallel_config = vllm_config.parallel_config
+    tp_size = parallel_config.tensor_parallel_size
+    if parallel_config.use_ubatching:
+        reason = "ubatching is enabled"
+    elif vllm_config.model_config.dtype != torch.bfloat16:
+        reason = "the model dtype is not BF16"
+    elif not current_platform.is_cuda():
+        reason = "the device is not CUDA"
+    elif not current_platform.is_device_capability_family(100):
+        reason = "the device is not SM100-family"
+    elif not 1 < tp_size <= 16:
+        reason = "TP size is not in the supported range 2-16"
+    elif 128 % tp_size != 0:
+        reason = "TP size does not divide 128"
+    elif N % 128 != 0:
+        reason = f"the output width {N} is not a multiple of 128"
+    else:
+        reason = None
+
+    if reason is not None:
+        logger.warning_once("%s is disabled because %s.", mode, reason)
+        return False
+
+    try:
+        init_gemm_rs_ar(
+            max_M=vllm_config.scheduler_config.max_num_batched_tokens,
+            N=N,
+            all_reduce=all_reduce,
+        )
+    except RuntimeError as e:
+        logger.warning_once(
+            "%s is disabled because initialization failed: %s. This may mean "
+            "the TP ranks do not share one NVLink domain.",
+            mode,
+            e,
+        )
+        return False
+    logger.info_once("%s is enabled.", mode)
+    return True
+
+
 def warmup_gemm_rs_ar() -> int:
     """Compile every reachable dispatch for the initialized GEMM-RS/AR mode."""
     # Initialization can be disabled or fail when multicast is unavailable.
     if _gemm_rs_ar is None:
         return 0
-    # Keep these profiles in sync with the dispatch in GemmRsAr.__call__.
-    profiles = ((128, 1), (128, 2), (256, 2))
-    for BN, cta_group in profiles:
-        Sm100GemmRsArBF16.compile(
-            _gemm_rs_ar.rank,
-            _gemm_rs_ar.world_size,
-            BN,
-            cta_group,
-            _gemm_rs_ar.all_reduce,
-        )
-    return len(profiles)
+
+    for dtype in sorted(_gemm_rs_ar.dtypes):
+        for cta_group in (1, 2):
+            Sm100GemmRsAr.compile(
+                _gemm_rs_ar.rank,
+                _gemm_rs_ar.world_size,
+                cta_group,
+                _gemm_rs_ar.all_reduce,
+                dtype=dtype,
+            )
+    return 2 * len(_gemm_rs_ar.dtypes)
 
 
 def get_gemm_rs_ar() -> GemmRsAr:
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py
index c1ff798a479b..6c786eed3591 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py
@@ -1,9 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from vllm.model_executor.kernels.linear.mixed_precision.allspark import (
-    AllSparkLinearKernel,
-)
 from vllm.model_executor.kernels.linear.mixed_precision.conch import (
     ConchLinearKernel,
 )
@@ -52,7 +49,6 @@
 __all__ = [
     "MPLinearKernel",
     "MPLinearLayerConfig",
-    "AllSparkLinearKernel",
     "ConchLinearKernel",
     "CPUWNA16LinearKernel",
     "CutlassW4A8LinearKernel",
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/allspark.py b/vllm/model_executor/kernels/linear/mixed_precision/allspark.py
deleted file mode 100644
index 704e661aeda9..000000000000
--- a/vllm/model_executor/kernels/linear/mixed_precision/allspark.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.model_executor.layers.quantization.utils import replace_parameter
-from vllm.model_executor.layers.quantization.utils.allspark_utils import (
-    ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
-    check_allspark_supported_dtype_shape,
-)
-from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
-from vllm.utils.platform_utils import num_compute_units
-
-from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
-
-
-class AllSparkLinearKernel(MPLinearKernel):
-    @classmethod
-    def get_min_capability(cls) -> int:
-        return 80
-
-    @classmethod
-    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
-        if c.zero_points:
-            return False, "Zero points currently not supported by AllSpark"
-
-        return check_allspark_supported_dtype_shape(
-            c.partition_weight_shape[0],  # in_features
-            c.partition_weight_shape[1],  # out_features
-            c.group_size,
-            c.weight_type,
-            c.act_type,
-        )
-
-    # note assumes that
-    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
-    #  `weight_scale` is: {input_dim = 0, output_dim = 1}
-    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        device = getattr(layer, self.w_q_name).device
-        c = self.config
-
-        # prepare the parameters required for the kernel
-        properties = torch.cuda.get_device_properties(device.index)
-        sm_count = num_compute_units(device.index)
-        sm_version = properties.major * 10 + properties.minor
-        gemm_args = {}
-        gemm_args["sm_count"] = sm_count
-        gemm_args["sm_version"] = sm_version
-
-        self.gemm_args = gemm_args
-
-        # transform param weight, scale
-        old_weight_param = getattr(layer, self.w_q_name)
-        old_scale_param = getattr(layer, self.w_s_name)
-
-        assert isinstance(old_weight_param, BasevLLMParameter)
-        permute_param_layout_(old_weight_param, input_dim=0, output_dim=1, packed_dim=0)
-
-        assert isinstance(old_scale_param, BasevLLMParameter)
-        permute_param_layout_(old_scale_param, input_dim=0, output_dim=1)
-
-        # unpack weight from K / 4 x N int32 to K x N uint8
-        new_weight_param = torch.nn.Parameter(
-            old_weight_param.data, requires_grad=False
-        )
-        new_weight_param.data = (
-            new_weight_param.data.t().contiguous().view(dtype=torch.uint8)
-        )
-        new_weight_param.data = new_weight_param.data.t().contiguous()
-
-        new_scale_param = torch.nn.Parameter(old_scale_param.data, requires_grad=False)
-
-        # reorder K x N weight as N32K16 format for Ampere W8A16
-        new_weight_param.data, new_scale_param.data, _ = ops.allspark_repack_weight(
-            new_weight_param.data, new_scale_param.data, None, c.zero_points
-        )
-
-        replace_parameter(layer, self.w_q_name, new_weight_param.data)
-        replace_parameter(layer, self.w_s_name, new_scale_param.data)
-
-    def apply_weights(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        c = self.config
-        gemm_args = self.gemm_args
-        w_q, w_s, _ = self._get_weight_params(layer)
-
-        reshaped_x = x.reshape(-1, x.shape[-1])
-        out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
-
-        output = ops.allspark_w8a16_gemm(
-            a=reshaped_x,
-            b_qweight=w_q,
-            b_scales=w_s,
-            b_qzeros=None,
-            n=c.partition_weight_shape[1],
-            group_size=c.group_size,
-            sm_count=gemm_args["sm_count"],
-            sm_version=gemm_args["sm_version"],
-            CUBLAS_M_THRESHOLD=ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
-            has_zp=c.zero_points,
-            n32k16_reorder=True,
-        )
-
-        if bias is not None:
-            output.add_(bias)  # In-place add
-
-        return output.reshape(out_shape)
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/humming.py b/vllm/model_executor/kernels/linear/mixed_precision/humming.py
index 3de8b9ca4151..4655181f78c6 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/humming.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/humming.py
@@ -24,18 +24,21 @@ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
         return True, None
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             convert_linear_layer_to_humming_standard,
             get_humming_linear_compute_config,
             prepare_humming_linear_layer_config,
         )
 
         name_map = {"weight": self.w_q_name, "weight_scale": self.w_s_name}
+        if self.w_zp_name is not None and hasattr(layer, self.w_zp_name):
+            name_map["zero_point"] = self.w_zp_name
         group_size = self.config.group_size
         quant_config = {
             "quant_method": "humming",
             "dtype": "int" + str(self.config.weight_type.size_bits),
             "group_size": 0 if group_size == -1 else group_size,
+            "has_zero_point": self.config.zero_points,
         }
 
         if self.config.zero_points:
@@ -57,13 +60,14 @@ def apply_weights(
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             apply_humming_linear,
         )
 
         return apply_humming_linear(
             layer,
             x,
+            skip_bias_add=bias is None,
             layer_config=self.layer_config,
             compute_config=self.compute_config,
             locks=self.locks,
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py
index 387c33c11068..60687fc8602d 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py
@@ -10,7 +10,6 @@
     apply_gptq_marlin_linear,
     marlin_act_int8_process_scales,
     marlin_make_empty,
-    marlin_make_workspace_new,
     marlin_pad_dim,
     marlin_pad_qweight,
     marlin_pad_scales,
@@ -91,11 +90,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         size_k, size_n = c.partition_weight_shape
         padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
 
-        # Allocate marlin workspace, reusing existing storage on reload.
-        self.workspace = marlin_make_workspace_new(
-            device, existing=getattr(self, "workspace", None)
-        )
-
         # Default name since marlin requires empty parameter for zp,
         # TODO: remove this requirement from marlin (allow optional tensors)
         if self.w_zp_name is None:
@@ -197,7 +191,7 @@ def apply_weights(
             weight=w_q,
             weight_scale=w_s,
             weight_zp=w_zp,  # type: ignore
-            workspace=self.workspace,
+            workspace=None,
             wtype=c.weight_type,
             input_size_per_partition=c.partition_weight_shape[0],
             output_size_per_partition=c.partition_weight_shape[1],
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
index 9964c891b6b3..86e029b5f88d 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode.
+"""Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode.
 
 Routes based on batch size M:
   M <= MAX_SKINNY_BATCH_SIZE: HIP skinny GEMM (wvSplitK_int4_g)
@@ -77,7 +76,7 @@ def _triton_w4a16_skinny_fmt_kernel(
     a_ptr,  # [M, K]  fp16/bf16 activations
     b_ptr,  # [N, K//8]  int32 packed (ExLlama shuffle, K is packed dim)
     scales_ptr,  # [N, K//G]  fp16/bf16 scales (skinny layout)
-    zp_ptr,  # [N, K//G]  fp16/bf16 raw zero-points (when HAS_ZP=True)
+    zp_ptr,  # [N//8, K//G]  int32 zero-points (when HAS_ZP=True)
     c_ptr,  # [M, N]  fp16/bf16 output
     # Dimensions
     M,
@@ -94,8 +93,7 @@ def _triton_w4a16_skinny_fmt_kernel(
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
 ):
-    """
-    Fused W4A16 GEMM reading weights from skinny format [N, K//8].
+    """Fused W4A16 GEMM reading weights from skinny format [N, K//8].
 
     B is stored as [N, K//8] int32 using ExLlama shuffle packing:
       each int32 packs 8 K-values with interleave [0,2,4,6,1,3,5,7]:
@@ -103,8 +101,9 @@ def _triton_w4a16_skinny_fmt_kernel(
                | (val[1]<<16) | (val[3]<<20) | (val[5]<<24) | (val[7]<<28)
 
     Scales are [N, K//G] (skinny layout, NOT transposed).
-    When HAS_ZP=True, raw zero-points zp_raw are loaded from zp_ptr [N, K//G]
-    and subtracted directly: (nibble - zp_raw) * scale.
+    When HAS_ZP=True, zp_ptr holds [N//8, K//G] int32 with row n's raw
+    zero-point at word[n//8] bits 4*(n%8), and dequant is
+    (nibble - zp_raw) * scale.
     When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric).
     """
     pid_m = tl.program_id(0)
@@ -148,9 +147,10 @@ def _triton_w4a16_skinny_fmt_kernel(
         scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0)
 
         if HAS_ZP:
-            zp_ptrs = zp_ptr + offs_n * num_groups + group_idx
-            zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0)
-            b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None]
+            zp_ptrs = zp_ptr + (offs_n // 8) * num_groups + group_idx
+            zp_word = tl.load(zp_ptrs, mask=scale_mask, other=0)
+            zp_raw = (zp_word >> (4 * (offs_n % 8))) & 0xF
+            b_fp = (b - zp_raw[:, None]).to(scales.dtype) * scales[:, None]
         else:
             b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None]
 
@@ -192,10 +192,9 @@ def triton_w4a16_skinny_fmt_gemm(
     scales: torch.Tensor,  # [N, K//G] fp16/bf16
     group_size: int,
     zp_bias: int = 8,
-    zp: torch.Tensor | None = None,  # [N, K//G] per-group zero-points
+    zp: torch.Tensor | None = None,  # [N//8, K//G] int32 zero-points
 ) -> torch.Tensor:
-    """
-    Fused W4A16 GEMM reading from skinny weight format [N, K//8].
+    """Fused W4A16 GEMM reading from skinny weight format [N, K//8].
 
     Args:
         a:          Activation matrix [M, K], float16 or bfloat16.
@@ -203,12 +202,13 @@ def triton_w4a16_skinny_fmt_gemm(
         scales:     Per-group scales [N, K//G], same dtype as a.
         group_size: Quantization group size (resolved from -1 to K by caller).
         zp_bias:    Constant zero bias (default 8 for unsigned int4).
-        zp:         Raw per-group zero-points [N, K//G] (asymmetric),
-                    stored as zp_raw in activation dtype. When provided,
+        zp:         Raw per-group zero-points [N//8, K//G] int32, row n at
+                    word[n//8] bits 4*(n%8) (asymmetric). When provided,
                     dequant is (nibble - zp_raw) * scale.
 
     Returns:
         Output matrix [M, N], same dtype as a.
+
     """
     assert a.is_contiguous(), "Activation matrix must be contiguous"
     assert b_q.is_contiguous(), "Weight matrix must be contiguous"
@@ -225,8 +225,9 @@ def triton_w4a16_skinny_fmt_gemm(
     )
     if zp is not None:
         assert zp.is_contiguous(), "Zero-points must be contiguous"
-        assert zp.shape == (N, num_groups), (
-            f"zp shape mismatch: {zp.shape} vs ({N}, {num_groups})"
+        assert N % 8 == 0, f"N must be divisible by 8 for packed zp, got {N}"
+        assert zp.shape == (N // 8, num_groups), (
+            f"zp shape mismatch: {zp.shape} vs ({N // 8}, {num_groups})"
         )
     has_zp = zp is not None
 
@@ -381,7 +382,10 @@ def _rdna_hybrid_w4a16_apply_impl(
     cu_count: int,
     group_size: int,
 ) -> torch.Tensor:
-    """Dispatch between skinny GEMM and Triton based on batch size M."""
+    """Dispatch between skinny GEMM and Triton based on batch size M.
+
+    ``w_zp`` is [N//8, K//G] int32 for asymmetric layers, None for symmetric.
+    """
     import vllm._custom_ops as ops
 
     M = x_2d.shape[0]
@@ -520,10 +524,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             assert self.w_zp_name is not None
             w_zp_raw = getattr(layer, self.w_zp_name)
             permute_param_layout_(w_zp_raw, input_dim=1, output_dim=0, packed_dim=0)
-            zp_unpacked = unpack_quantized_values_into_int32(
-                w_zp_raw.data, c.weight_type, packed_dim=0
-            )
-            w_zp = zp_unpacked.to(c.act_type).contiguous()
+            w_zp = w_zp_raw.data.contiguous()
             self._transform_param(layer, self.w_zp_name, lambda x: w_zp)
 
         self._transform_param(layer, self.w_q_name, lambda x: w_q_skinny)
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
index fd6fb033fe95..ed6c67aa99c5 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
+++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Triton-based W4A16 GEMM kernel for ROCm MI300.
+"""Triton-based W4A16 GEMM kernel for ROCm MI300.
 
 Implements fused int4-weight dequantization + fp16 GEMM in a single kernel,
 using GPTQ sequential packing (8 int4 values per int32, shifts [0,4,...,28]).
@@ -67,8 +66,7 @@ def triton_w4a16_gemm_kernel(
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
 ):
-    """
-    Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N]
+    """Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N].
 
     B is stored as [K, N//8] int32 using GPTQ sequential packing:
       each int32 packs 8 consecutive N-values at bit offsets [0,4,8,12,16,20,24,28].
@@ -170,8 +168,7 @@ def _triton_w4a16_gemm_impl(
     group_size: int,
     zp_bias: int = 8,  # bias for uint4b8 when qzeros is None
 ) -> torch.Tensor:
-    """
-    Fused W4A16 GEMM using GPTQ-packed int4 weights.
+    """Fused W4A16 GEMM using GPTQ-packed int4 weights.
 
     Args:
         a:          Activation matrix [M, K], float16 or bfloat16.
@@ -184,6 +181,7 @@ def _triton_w4a16_gemm_impl(
 
     Returns:
         Output matrix [M, N], same dtype as a.
+
     """
     assert a.is_contiguous(), "Activation matrix must be contiguous"
     assert b_q.is_contiguous(), "Weight matrix must be contiguous"
@@ -302,8 +300,7 @@ def triton_w4a16_gemm(
 
 
 class TritonW4A16LinearKernel(MPLinearKernel):
-    """
-    Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer).
+    """Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer).
 
     Supports GPTQ-format int4 weights (uint4b8 symmetric, uint4 asymmetric)
     with grouped quantization. Weight tensors are transposed from the
@@ -360,8 +357,7 @@ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
         return True, None
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        """
-        Convert compressed-tensors checkpoint layout to kernel layout.
+        """Convert compressed-tensors checkpoint layout to kernel layout.
 
         Checkpoint (from compressed_tensors_wNa16.create_weights):
           weight_packed:     [N, K//8]  int32   input_dim=1, output_dim=0, packed_dim=1
diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py
index 13fd5db2750a..420e7f9493e1 100644
--- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py
+++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py
@@ -158,7 +158,7 @@ def is_supported(
         from vllm._aiter_ops import is_aiter_found_and_supported
         from vllm.model_executor.kernels.linear import _get_linear_backend
 
-        linear_backend = _get_linear_backend()
+        linear_backend = _get_linear_backend(quantization="mxfp4")
 
         if (
             current_platform.is_rocm()
diff --git a/vllm/model_executor/kernels/linear/mxfp4/base.py b/vllm/model_executor/kernels/linear/mxfp4/base.py
index 9c3088df52a1..de0a2a216327 100644
--- a/vllm/model_executor/kernels/linear/mxfp4/base.py
+++ b/vllm/model_executor/kernels/linear/mxfp4/base.py
@@ -19,6 +19,7 @@ class MxFp4LinearLayerConfig:
     Attributes:
         activation_quant_key: Identifies the activation quantization format,
             or `None` when activations must not be quantized.
+
     """
 
     activation_quant_key: QuantKey | None = None
diff --git a/vllm/model_executor/kernels/linear/mxfp4/emulation.py b/vllm/model_executor/kernels/linear/mxfp4/emulation.py
index ddd0143c2214..dce664c9a71c 100644
--- a/vllm/model_executor/kernels/linear/mxfp4/emulation.py
+++ b/vllm/model_executor/kernels/linear/mxfp4/emulation.py
@@ -8,6 +8,7 @@
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
+import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
     dequant_mxfp4,
@@ -97,7 +98,16 @@ def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None
         return True, None
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
+        weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
+
+        # Match the MXFP8 emulation path when explicitly enabled: dequantize
+        # once at load time so inference uses a plain high-precision linear.
+        # Otherwise retain packed weights and dequantize per invocation.
+        if envs.VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD:
+            weight = dequant_mxfp4(layer.weight, weight_scale, torch.bfloat16)
+            layer.weight = Parameter(weight.contiguous(), requires_grad=False)
+
+        layer.weight_scale = weight_scale
 
     def apply_weights(
         self,
@@ -105,6 +115,12 @@ def apply_weights(
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
         qdq_x = self.quant_dequant_func(x)
+
+        # Packed MXFP4 weights use one byte for two values. A weight with
+        # elements wider than one byte was already dequantized at load time.
+        if layer.weight.element_size() >= 2:
+            return F.linear(qdq_x, layer.weight.to(x.dtype), bias)
+
+        dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
         return F.linear(qdq_x, dq_w, bias)
diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py
index 660e31dcf4f2..0504fc5482ad 100644
--- a/vllm/model_executor/kernels/linear/mxfp4/humming.py
+++ b/vllm/model_executor/kernels/linear/mxfp4/humming.py
@@ -3,21 +3,18 @@
 
 import torch
 
-from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
+from vllm.model_executor.layers.quantization.utils.humming import (
     apply_humming_linear,
     convert_linear_layer_to_humming_standard,
     get_humming_linear_compute_config,
     prepare_humming_linear_layer_config,
+    quant_key_to_input_schema,
 )
-from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_humming
 
 from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
 
-logger = init_logger(__name__)
-
 
 class HummingMxFp4LinearKernel(MxFp4LinearKernel):
     """Humming GEMM Kernel for MXFP4."""
@@ -39,14 +36,10 @@ def is_supported(
 
     @classmethod
     def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
-        if config.activation_quant_key not in (None, kMxfp4Dynamic):
-            return False, "only supports MXFP4 dynamic or unquantized activations"
-        if config.activation_quant_key is not None:
-            logger.warning_once(
-                "HummingMxFp4LinearKernel is a weight-only (A16) kernel; "
-                "the requested activation quantization (%s) is ignored.",
-                config.activation_quant_key,
-            )
+        try:
+            quant_key_to_input_schema(config.activation_quant_key)
+        except ValueError as error:
+            return False, str(error)
         return True, None
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -62,7 +55,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         }
 
         convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
-        self.layer_config = prepare_humming_linear_layer_config(layer, quant_config)
+        input_schema = quant_key_to_input_schema(self.config.activation_quant_key)
+        self.layer_config = prepare_humming_linear_layer_config(
+            layer, quant_config, input_schema=input_schema
+        )
         self.compute_config = get_humming_linear_compute_config()
         self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
 
@@ -75,6 +71,7 @@ def apply_weights(
         return apply_humming_linear(
             layer,
             x,
+            skip_bias_add=bias is None,
             layer_config=self.layer_config,
             compute_config=self.compute_config,
             locks=self.locks,
diff --git a/vllm/model_executor/kernels/linear/mxfp4/marlin.py b/vllm/model_executor/kernels/linear/mxfp4/marlin.py
index 30ff0d58d4c0..108f835e6435 100644
--- a/vllm/model_executor/kernels/linear/mxfp4/marlin.py
+++ b/vllm/model_executor/kernels/linear/mxfp4/marlin.py
@@ -58,7 +58,7 @@ def apply_weights(
             weight=layer.weight,
             weight_scale=layer.weight_scale,
             weight_global_scale=None,
-            workspace=layer.workspace,
+            workspace=None,
             size_n=layer.output_size_per_partition,
             size_k=layer.input_size_per_partition,
             bias=bias,
diff --git a/vllm/model_executor/kernels/linear/mxfp6/base.py b/vllm/model_executor/kernels/linear/mxfp6/base.py
index ef00bc57647d..947ea60e3fdd 100644
--- a/vllm/model_executor/kernels/linear/mxfp6/base.py
+++ b/vllm/model_executor/kernels/linear/mxfp6/base.py
@@ -13,14 +13,15 @@
 class MxFp6LinearLayerConfig:
     """Configuration for an MXFP6 linear layer.
 
-    All MXFP6 layers share the same structure: packed uint8 weights (2 FP4 values per
-    byte) and per-block weight scales (group size 32).
+    All MXFP6 layers share the same structure: packed uint8 weights (4 FP6 values
+    per 3 bytes) and per-block weight scales (group size 32).
 
     Attributes:
         weight_quant_key: Identifies the weight quantization format. Can be
             kMxfp6E2M3Static or kMxfp6E3M2Static.
         activation_quant_key: Identifies the activation quantization format,
             or `None` when activations must not be quantized.
+
     """
 
     weight_quant_key: QuantKey
diff --git a/vllm/model_executor/kernels/linear/mxfp6/humming.py b/vllm/model_executor/kernels/linear/mxfp6/humming.py
new file mode 100644
index 000000000000..fff7584d9332
--- /dev/null
+++ b/vllm/model_executor/kernels/linear/mxfp6/humming.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.humming import (
+    apply_humming_linear,
+    convert_linear_layer_to_humming_standard,
+    get_humming_linear_compute_config,
+    prepare_humming_linear_layer_config,
+    quant_key_to_input_schema,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp6E2M3Static
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_humming
+
+from .base import MxFp6LinearKernel, MxFp6LinearLayerConfig
+
+
+class HummingMxFp6LinearKernel(MxFp6LinearKernel):
+    """Humming GEMM for packed MXFP6 E2M3 and E3M2 weights."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "Humming only supported on CUDA"
+
+        if not has_humming():
+            return False, "Humming is not installed"
+
+        if not current_platform.has_device_capability(75):
+            return False, "Humming only supported on SM75+"
+
+        return True, None
+
+    @classmethod
+    def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]:
+        try:
+            quant_key_to_input_schema(config.activation_quant_key)
+        except ValueError as error:
+            return False, str(error)
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight_dtype = "float6e3m2"
+        if self.config.weight_quant_key == kMxfp6E2M3Static:
+            weight_dtype = "float6e2m3"
+
+        quant_config = {
+            "quant_method": "humming",
+            "dtype": weight_dtype,
+            "scale_dtype": "float8e8m0",
+            "group_size": 32,
+            "weight_scale_type": "group",
+        }
+        layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu)
+
+        convert_linear_layer_to_humming_standard(
+            layer=layer,
+            name_map={"weight": "weight", "weight_scale": "weight_scale"},
+        )
+        input_schema = quant_key_to_input_schema(self.config.activation_quant_key)
+        self.layer_config = prepare_humming_linear_layer_config(
+            layer, quant_config, input_schema=input_schema
+        )
+        self.compute_config = get_humming_linear_compute_config()
+        self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return apply_humming_linear(
+            layer,
+            x,
+            skip_bias_add=bias is None,
+            layer_config=self.layer_config,
+            compute_config=self.compute_config,
+            locks=self.locks,
+        )
diff --git a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py
index 3fe34029e5cb..fdcfeb264b29 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py
@@ -27,6 +27,10 @@ class Mxfp8LinearKernel(ABC):
     Marlin, emulation).
     """
 
+    supports_pre_processed_weights: bool = False
+    """True if ``process_weights_after_loading`` only rewrites parameters, so
+    weights exported by the weight cache daemon can be used as-is."""
+
     def __init__(self, c: Mxfp8LinearLayerConfig) -> None:
         assert self.can_implement(c)[0]
         assert self.is_supported()[0]
diff --git a/vllm/model_executor/kernels/linear/mxfp8/emulation.py b/vllm/model_executor/kernels/linear/mxfp8/emulation.py
index 79b2fba3889b..da48655d37db 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/emulation.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/emulation.py
@@ -16,6 +16,8 @@
 class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
     """Software emulation fallback for MXFP8 (dequant to BF16)."""
 
+    supports_pre_processed_weights = True
+
     @classmethod
     def is_supported(
         cls, compute_capability: int | None = None
diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
index 02700a0ced0c..cc9b35d34bdd 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
@@ -27,6 +27,8 @@
 class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
     """MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+)."""
 
+    supports_pre_processed_weights = True
+
     @classmethod
     def is_supported(
         cls, compute_capability: int | None = None
@@ -117,6 +119,8 @@ def apply_weights(
 class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
     """MXFP8 W8A8 GEMM via FlashInfer CuTe-DSL (SM100/SM103)."""
 
+    supports_pre_processed_weights = True
+
     @classmethod
     def is_supported(
         cls, compute_capability: int | None = None
diff --git a/vllm/model_executor/kernels/linear/mxfp8/humming.py b/vllm/model_executor/kernels/linear/mxfp8/humming.py
index bde1565d2a3e..63e9514f4ea3 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/humming.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/humming.py
@@ -3,12 +3,14 @@
 
 import torch
 
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
+from vllm.model_executor.layers.quantization.utils.humming import (
     apply_humming_linear,
     convert_linear_layer_to_humming_standard,
     get_humming_linear_compute_config,
     prepare_humming_linear_layer_config,
+    quant_key_to_input_schema,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_humming
 
@@ -50,7 +52,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         }
 
         convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
-        self.layer_config = prepare_humming_linear_layer_config(layer, quant_config)
+        input_schema = quant_key_to_input_schema(kMxfp8Dynamic)
+        self.layer_config = prepare_humming_linear_layer_config(
+            layer, quant_config, input_schema=input_schema
+        )
         self.compute_config = get_humming_linear_compute_config()
         self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
 
@@ -63,6 +68,7 @@ def apply_weights(
         return apply_humming_linear(
             layer,
             x,
+            skip_bias_add=bias is None,
             layer_config=self.layer_config,
             compute_config=self.compute_config,
             locks=self.locks,
diff --git a/vllm/model_executor/kernels/linear/mxfp8/marlin.py b/vllm/model_executor/kernels/linear/mxfp8/marlin.py
index bec54cd942ed..6ebdbd647405 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/marlin.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/marlin.py
@@ -46,7 +46,7 @@ def apply_weights(
             input=x,
             weight=layer.weight,
             weight_scale=layer.weight_scale,
-            workspace=layer.workspace,
+            workspace=None,
             size_n=layer.output_size_per_partition,
             size_k=layer.input_size_per_partition,
             bias=bias,
diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
index 9714f92dffd4..c7e6537f4948 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
@@ -4,9 +4,9 @@
 
 Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
 activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
-matrix cores. Falls back (via the kernel selector) to the BF16
-``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
-``K % 128 != 0``.
+matrix cores. ``dot_scaled`` tiles K by 128; a weight whose K is not a multiple
+of that is dequantized to BF16 once at load and served by a plain linear
+instead, since ``can_implement`` does not filter on K.
 """
 
 import torch
@@ -119,6 +119,18 @@ def _mxfp8_dot_scaled_linear(
     return out
 
 
+# Triton 3.8 enables TRITON_HIP_USE_ASYNC_COPY by default on gfx950; its extra LDS
+# buffer leaves no room for num_stages=3 at BLOCK_K=256.
+# TODO(rasmith)(Rohan138): Remove the 3.8 check once
+# https://github.com/vllm-project/vllm/pull/50605 merges.
+_BK256_STAGES = 3
+if triton.__version__.startswith("3.8") and current_platform.is_rocm():
+    from vllm.platforms.rocm import on_gfx950
+
+    if on_gfx950():
+        _BK256_STAGES = 2
+
+
 def _select_cfg(M, N, K):
     """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950.
 
@@ -166,7 +178,7 @@ def _select_cfg(M, N, K):
         # 3.6; on triton 3.7 its large BLOCK_M register/LDS footprint spills or hits
         # "out of resources", so use 128x128x256 -- within the known-good footprint.)
         if M >= 4096 and K >= 1024 and K % 256 == 0 and occ >= 256:
-            return 128, 128, 256, 8, 3
+            return 128, 128, 256, 8, _BK256_STAGES
         return 128, 128, 128, 8, 3
     # large-K (K >= 2048). BLOCK_K is K-divisibility-guarded (the K-loop is unmasked):
     # served large-K is 2048/6144 (%256==0), but fall back to 128 (always divides, since
@@ -186,7 +198,7 @@ def _select_cfg(M, N, K):
     # Covers the qkv-class local N=1536 (TP=8 qkv / TP=4 shared_gate_up) and the deep-K
     # / very-large-M shapes.
     if K % 256 == 0 and (1280 < N <= 1536 or (occ >= 128 and (K >= 4096 or M >= 4096))):
-        return 128, 128, 256, 8, 3
+        return 128, 128, 256, 8, _BK256_STAGES
     # small local-N (e.g. TP=8 shared_gate_up N=768): a 64-wide BLOCK_N doubles the
     # N-tile count -> better CU fill than 128x128 at this mid-large M (~1.4x there).
     if N <= 1024 and K % 256 == 0:
@@ -194,9 +206,14 @@ def _select_cfg(M, N, K):
     return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3)
 
 
+_DOT_SCALED_K_ALIGN = 128
+
+
 class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
     """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
 
+    supports_pre_processed_weights = True
+
     @classmethod
     def is_supported(
         cls, compute_capability: int | None = None
@@ -219,6 +236,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         N, K = weight.shape
         scale_k = K // MXFP8_BLOCK_SIZE
         weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
+        if K % _DOT_SCALED_K_ALIGN != 0:
+            weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
         layer.weight = Parameter(weight.contiguous(), requires_grad=False)
         layer.weight_scale = Parameter(weight_scale, requires_grad=False)
 
@@ -235,10 +254,13 @@ def apply_weights(
             )
         out_shape = (*x.shape[:-1], layer.weight.shape[0])
         x2d = x.reshape(-1, x.shape[-1])
-        if x2d.shape[-1] % 128 == 0:
+        if layer.weight.element_size() >= 2:
+            out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype))
+        elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0:
             out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
         else:
-            # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
+            # process_weights_after_loading dequantizes these weights once, so
+            # this runs only if it did not; dot_scaled would be invalid here.
             w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
             out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
         out = out.reshape(out_shape)
diff --git a/vllm/model_executor/kernels/linear/nvfp4/humming.py b/vllm/model_executor/kernels/linear/nvfp4/humming.py
index 54a0cf6c1d8e..3829c39b7486 100644
--- a/vllm/model_executor/kernels/linear/nvfp4/humming.py
+++ b/vllm/model_executor/kernels/linear/nvfp4/humming.py
@@ -4,11 +4,13 @@
 import torch
 
 from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
+from vllm.model_executor.layers.quantization.utils.humming import (
     apply_humming_linear,
     get_humming_linear_compute_config,
     prepare_humming_linear_layer_config,
+    quant_key_to_input_schema,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_humming
 
@@ -40,9 +42,7 @@ def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None
         return True, None
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        # Route through humming's compressed-tensors nvfp4 loader (same path as
-        # the MoE oracle); the native group_tensor schema mishandles a scalar
-        # global scale.
+        # Reuse the CT loader for packed FP4 weights and tensor scales.
         quant_config = {
             "quant_method": "compressed-tensors",
             "format": "nvfp4-pack-quantized",
@@ -60,7 +60,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.weight_global_scale = torch.nn.Parameter(
             1.0 / layer.weight_global_scale, requires_grad=False
         )
-        self.layer_config = prepare_humming_linear_layer_config(layer, quant_config)
+        input_quant_key = None
+        if hasattr(layer, "input_global_scale"):
+            layer.input_scale_2 = torch.nn.Parameter(
+                layer.input_global_scale.reshape(-1), requires_grad=False
+            )
+            input_quant_key = kNvfp4Dynamic
+        input_schema = quant_key_to_input_schema(input_quant_key)
+        self.layer_config = prepare_humming_linear_layer_config(
+            layer, quant_config, input_schema=input_schema
+        )
         self.compute_config = get_humming_linear_compute_config()
         self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
 
@@ -73,6 +82,7 @@ def apply_weights(
         return apply_humming_linear(
             layer,
             x,
+            skip_bias_add=bias is None,
             layer_config=self.layer_config,
             compute_config=self.compute_config,
             locks=self.locks,
diff --git a/vllm/model_executor/kernels/linear/nvfp4/marlin.py b/vllm/model_executor/kernels/linear/nvfp4/marlin.py
index a05d6823c881..1ab5336876bb 100644
--- a/vllm/model_executor/kernels/linear/nvfp4/marlin.py
+++ b/vllm/model_executor/kernels/linear/nvfp4/marlin.py
@@ -50,7 +50,7 @@ def apply_weights(
             weight=layer.weight,
             weight_scale=layer.weight_scale,
             weight_global_scale=layer.weight_global_scale,
-            workspace=layer.workspace,
+            workspace=None,
             size_n=layer.output_size_per_partition,
             size_k=layer.input_size_per_partition,
             bias=bias,
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py
index c83885ca51c4..3a053412f3cf 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py
@@ -38,6 +38,22 @@ def from_layer(cls, layer: torch.nn.Module) -> Self:
             input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None),
         )
 
+    @property
+    def block_scale_attr(self) -> str:
+        """Fp8LinearMethod registers the block scale as ``weight_scale_inv``,
+        compressed-tensors as ``weight_scale``."""
+        return (
+            self.WEIGHT_SCALE
+            if self.weight_scale_inv is None
+            else self.WEIGHT_SCALE_INV
+        )
+
+    @property
+    def block_scale(self) -> torch.Tensor:
+        scale = getattr(self, self.block_scale_attr)
+        assert scale is not None
+        return scale
+
 
 class Fp8BlockScaledMMLinearKernel(
     MMLinearKernel[FP8ScaledMMLinearLayerConfig, FP8BlockParams], ABC
@@ -74,25 +90,13 @@ def _get_layer_params(self, layer: torch.nn.Module, **kwargs) -> FP8BlockParams:
 
     def process_weights_after_loading(self, layer: torch.nn.Module):
         params = self._get_layer_params(layer)
-        # Fp8LinearMethod registered weight scale
-        # buffer as weight_scale_inv unlike compressed tensors.
-        weight_scale = (
-            params.weight_scale
-            if params.weight_scale_inv is None
-            else params.weight_scale_inv
-        )
-        scale_attr_name = (
-            params.WEIGHT_SCALE
-            if params.weight_scale_inv is None
-            else params.WEIGHT_SCALE_INV
-        )
         new_weight, new_weight_scale = process_fp8_weight_block_strategy(
             params.weight,
-            weight_scale,
+            params.block_scale,
         )
 
         replace_parameter(layer, params.WEIGHT, new_weight.data)
-        replace_parameter(layer, scale_attr_name, new_weight_scale.data)
+        replace_parameter(layer, params.block_scale_attr, new_weight_scale.data)
 
     def apply_weights(
         self,
@@ -104,11 +108,7 @@ def apply_weights(
         out_dtype = self.config.out_dtype
         params = self._get_layer_params(layer)
         weight = params.weight
-        weight_scale = (
-            params.weight_scale
-            if params.weight_scale_inv is None
-            else params.weight_scale_inv
-        )
+        weight_scale = params.block_scale
         input_scale = params.input_scale
         scale_up = params.input_scale_ub
 
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py
index 39f9abd460e1..75db4bcaa9ed 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py
@@ -9,6 +9,8 @@
 )
 from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
     CPUFp8BlockScaledMMKernel,
+    CPUFp8PerTensorScaledMMLinearKernel,
+    CPUFP8W8A8ScaledMMLinearKernel,
     CPUInt8ScaledMMLinearKernel,
 )
 from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -67,5 +69,7 @@
     "ZentorchInt8ScaledMMLinearKernel",
     "Fp8BlockScaledMMLinearKernel",
     "CPUFp8BlockScaledMMKernel",
+    "CPUFp8PerTensorScaledMMLinearKernel",
+    "CPUFP8W8A8ScaledMMLinearKernel",
     "XPUFp8BlockScaledMMKernel",
 ]
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py
index ce2fe690b572..b666960a1ce0 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py
@@ -15,6 +15,8 @@
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
+    QuantKey,
+    kFp8DynamicTokenSym,
 )
 from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
@@ -70,8 +72,7 @@ def apply_weights(
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        `AiterInt8ScaledMMLinearKernel` implements a fused version of
+        """`AiterInt8ScaledMMLinearKernel` implements a fused version of
             `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
         where scale_a * a and scale_b * b are implemented using numpy-style
         broadcasting.
@@ -150,6 +151,12 @@ def is_supported(
             return False, "requires aiter library to be installed."
         return True, None
 
+    def input_quant_key(self) -> QuantKey | None:
+        # Does not call get_output_padding() - torch fallbacks
+        # resolve padding from compilation_config, which is unset in
+        # profile_run. This kernel does not pad activations.
+        return kFp8DynamicTokenSym
+
     @classmethod
     def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
         is_ptpc = (
@@ -242,6 +249,11 @@ def is_supported(
 
         return True, None
 
+    def input_quant_key(self) -> QuantKey | None:
+        # Same per-token FP8 consume path as
+        # AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.
+        return kFp8DynamicTokenSym
+
     @classmethod
     def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
         is_ptpc = (
@@ -319,6 +331,9 @@ def is_supported(
             compute_capability
         )
 
+    def input_quant_key(self) -> QuantKey | None:
+        return kFp8DynamicTokenSym
+
     @classmethod
     def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
         is_ptpc = (
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py
index 4ef1af0a7059..1d2ede1853c8 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py
@@ -21,10 +21,14 @@
     FP8ScaledMMLinearLayerConfig,
 )
 from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
     Int8ScaledMMLinearKernel,
     Int8ScaledMMLinearLayerConfig,
 )
 
+# BLOCK_K in csrc/cpu/sgl-kernels/gemm.h — the AMX kernel's fixed K-block size.
+_BLOCK_SIZE_K = 128
+
 
 class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
     @classmethod
@@ -261,15 +265,22 @@ def can_implement(
         return True, None
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if getattr(layer, "_cpu_skip_gemm_dispatch", False):
+            return
         # Skip the base class process (FP8 padding / fnuz normalization)
         # which is GPU-oriented.  Instead, VNNI-prepack weights for AMX.
         params = self._get_layer_params(layer)
-        packed_weight = torch.ops._C.convert_weight_packed(params.weight)
-        replace_parameter(
-            layer,
-            params.WEIGHT,
-            torch.nn.Parameter(packed_weight, requires_grad=False),
-        )
+
+        # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this
+        # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and
+        # let the kernel pack on the fly instead.
+        if not getattr(layer, "_cpu_skip_gemm_dispatch", False):
+            packed_weight = torch.ops._C.convert_weight_packed(params.weight)
+            replace_parameter(
+                layer,
+                params.WEIGHT,
+                torch.nn.Parameter(packed_weight, requires_grad=False),
+            )
 
         # Re-wrap scale as a plain Parameter so the kernel can read it
         # without weight-loader metadata interfering.
@@ -314,7 +325,7 @@ def apply_weights(
             list(self.weight_group_shape),
             bias,
             x.dtype,
-            True,  # is_vnni (weight already prepacked)
+            not getattr(layer, "_cpu_skip_gemm_dispatch", False),  # is_vnni
         )
         return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
 
@@ -328,3 +339,390 @@ def apply_block_scaled_mm(
         raise NotImplementedError(
             "CPUFp8BlockScaledMMKernel overrides apply_weights directly."
         )
+
+
+class CPUFp8PerTensorScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    """FP8 W8A16 per-tensor-scaled GEMM via AMX BRGEMM on CPU.
+
+    Reuses the block-scaled AMX kernel (fp8_scaled_mm_cpu) with a single
+    synthetic block spanning the whole weight, so activations stay BF16/FP32
+    — no FP8 activation quantization, unlike PerTensorTorchFP8ScaledMMLinearKernel.
+    """
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU platform."
+        if not torch.cpu._is_amx_tile_supported():
+            return False, "requires AMX tile support (Sapphire Rapids or newer)."
+        if not ops._supports_cpu_fp8_w8a16:
+            return False, "fp8_scaled_mm_cpu op not available."
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if not c.weight_quant_key.scale.group_shape.is_per_tensor():
+            return False, "requires a per-tensor weight scale."
+        if not c.activation_quant_key.scale.group_shape.is_per_tensor():
+            return False, "requires a per-tensor activation scale."
+        if c.out_dtype not in (torch.bfloat16, torch.float32):
+            return False, "Only bfloat16/float32 output dtype supported."
+        n = c.weight_shape[0]
+        if n % 32 != 0:
+            # The AMX tinygemm kernel tiles N in chunks of 32 and cannot
+            # handle a remainder tile smaller than that.
+            return False, f"requires weight output dim (N={n}) to be a multiple of 32."
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w_name, w_s_name, _, _ = self.layer_param_names
+        # Fp8LinearMethod transposes weight to (K, N) for torch._scaled_mm-style
+        # kernels; convert_weight_packed/fp8_scaled_mm_cpu expect (N, K), as
+        # produced by CUDA nn.Linear-style checkpoints.
+        weight = getattr(layer, w_name).t().contiguous()
+        n, k = weight.shape
+
+        packed_weight = torch.ops._C.convert_weight_packed(weight)
+        replace_parameter(
+            layer, w_name, torch.nn.Parameter(packed_weight, requires_grad=False)
+        )
+
+        # Synthesize a single-block "block scale" tensor from the scalar
+        # per-tensor weight scale so fp8_scaled_mm_cpu can be reused as-is.
+        weight_scale = getattr(layer, w_s_name)
+        self._block_size_n = -(-n // 32) * 32  # round up to a multiple of 32
+        num_k_blocks = -(-k // _BLOCK_SIZE_K)
+        block_scale = weight_scale.reshape(()).expand(1, num_k_blocks).contiguous()
+        replace_parameter(
+            layer, w_s_name, torch.nn.Parameter(block_scale, requires_grad=False)
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_name, w_s_name, _, _ = self.layer_param_names
+        weight = getattr(layer, w_name)
+        weight_scale = getattr(layer, w_s_name)
+
+        x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
+        out = torch.ops._C.fp8_scaled_mm_cpu(
+            x_2d,
+            weight,
+            weight_scale,
+            [self._block_size_n, _BLOCK_SIZE_K],
+            bias,
+            x.dtype,
+            True,  # is_vnni (weight already prepacked)
+        )
+        return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        raise NotImplementedError(
+            "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly."
+        )
+
+
+class CPUFp8W8A8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    """FP8 W8A8 block-quantized GEMM with dynamic per-token activation quant on CPU."""
+
+    # Activation quant is done inside fp8_scaled_mm_with_quant;
+    # the base class must not apply a second quant step.
+    apply_input_quant = False
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU platform."
+        if not torch.cpu._is_amx_tile_supported():
+            return False, "requires AMX tile support (Sapphire Rapids or newer)."
+        if not ops.cpu_has_amx_fp8():
+            return False, (
+                "requires native AMX-FP8 hardware; falls back "
+                "to CPUFp8BlockScaledMMKernel (W8A16) otherwise."
+            )
+        return True, None
+
+    @classmethod
+    def can_implement(
+        cls, config: FP8ScaledMMLinearLayerConfig
+    ) -> tuple[bool, str | None]:
+        # Same block-shape requirements as the W8A16 sibling kernel.
+        weight_gs = config.weight_quant_key.scale.group_shape
+        if weight_gs.col <= 0 or weight_gs.col != 128:
+            return False, (
+                "CPU FP8 W8A8 block kernel requires K-dimension block size of 128, "
+                f"got {weight_gs.col}."
+            )
+        if weight_gs.row <= 0 or weight_gs.row % 32 != 0:
+            return False, (
+                "CPU FP8 W8A8 block kernel requires N-dimension block size to be "
+                f"a positive multiple of 32, got {weight_gs.row}."
+            )
+        if config.out_dtype not in (torch.bfloat16, torch.float32):
+            return False, "Only bfloat16/float32 output dtype supported."
+        n = config.weight_shape[0]
+        if n % 32 != 0:
+            # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot
+            # handle a remainder tile smaller than that.
+            return False, f"requires weight output dim (N={n}) to be a multiple of 32."
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Prepack block-quantized FP8 weight for W8A8 AMX BRGEMM.
+
+        For block-quant FP8 models weight is stored as [N, K] (NOT transposed,
+        unlike non-block FP8 where vLLM stores [K, N]).
+
+        float8_linear_prepack_cpu expects:
+          weight:  [N, K]
+          scales:  [N, G]  where G = K / block_k
+
+        weight_scale_inv from the checkpoint has shape [ceil(N/block_n), G],
+        so we expand rows from ceil(N/block_n) to N via repeat_interleave.
+        """
+        params = self._get_layer_params(layer)
+        weight = params.weight  # [N, K]
+
+        # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): MLA reads this
+        # weight to derive W_UK/W_UV, so keep the checkpoint [N, K] layout and
+        # its block scales; apply_weights falls back to the W8A16 op.
+        if getattr(layer, "_cpu_skip_gemm_dispatch", False):
+            return
+
+        weight_scale = (
+            params.weight_scale_inv
+            if params.weight_scale_inv is not None
+            else params.weight_scale
+        )  # [ceil(N/block_n), ceil(K/block_k)]
+        assert weight_scale is not None, "FP8 W8A8 block kernel requires weight scale."
+
+        N = weight.shape[0]
+        block_n = self.weight_group_shape.row  # e.g. 128
+
+        # Expand scale rows: [ceil(N/block_n), G] → [N, G]
+        if weight_scale.size(0) != N:
+            weight_scale_expanded = weight_scale.repeat_interleave(block_n, dim=0)[
+                :N, :
+            ].contiguous()
+        else:
+            weight_scale_expanded = weight_scale.contiguous()
+
+        packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu(
+            weight, weight_scale_expanded.to(torch.float32)
+        )
+
+        scale_attr = (
+            params.WEIGHT_SCALE_INV
+            if params.weight_scale_inv is not None
+            else params.WEIGHT_SCALE
+        )
+        replace_parameter(
+            layer,
+            params.WEIGHT,
+            torch.nn.Parameter(packed_weight, requires_grad=False),
+        )
+        replace_parameter(
+            layer,
+            scale_attr,
+            torch.nn.Parameter(packed_scales, requires_grad=False),
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        params = self._get_layer_params(layer)
+        weight_scale = (
+            params.weight_scale_inv
+            if params.weight_scale_inv is not None
+            else params.weight_scale
+        )
+
+        x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
+        if getattr(layer, "_cpu_skip_gemm_dispatch", False):
+            # Unpacked weight: fp8_scaled_mm_with_quant has no is_vnni option,
+            # so fall back to the W8A16 op (BF16 activation, block-scaled FP8
+            # weight) for this layer.
+            out = torch.ops._C.fp8_scaled_mm_cpu(
+                x_2d,
+                params.weight,
+                weight_scale,
+                list(self.weight_group_shape),
+                bias,
+                x.dtype,
+                False,  # is_vnni
+            )
+            return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
+
+        # act_scales=None + channelwise=True → dynamic per-token activation quant
+        out = ops.fp8_scaled_mm_with_quant(
+            x_2d,
+            None,
+            True,
+            params.weight,
+            weight_scale,
+            bias,
+            x.dtype,
+        )
+        return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        raise NotImplementedError(
+            "CPUFp8W8A8BlockScaledMMKernel overrides apply_weights directly."
+        )
+
+
+class CPUFP8W8A8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    """FP8 W8A8 GEMM with dynamic per-token activation quantization on CPU."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU platform."
+        if not torch.cpu._is_amx_tile_supported():
+            return False, "requires AMX tile support (Sapphire Rapids or newer)."
+        if not ops.cpu_has_amx_fp8():
+            return False, (
+                "requires native AMX-FP8 hardware (Diamond Rapids); falls back "
+                "to the Torch FP8 kernels otherwise."
+            )
+        return True, None
+
+    @classmethod
+    def can_implement(
+        cls, config: FP8ScaledMMLinearLayerConfig
+    ) -> tuple[bool, str | None]:
+        if config.out_dtype not in (torch.bfloat16, torch.float16, torch.float32):
+            return False, "Only bfloat16/float16/float32 output dtype supported."
+        # Reject block-quantized weights — those go to CPUFp8BlockScaledMMKernel
+        weight_gs = config.weight_quant_key.scale.group_shape
+        if weight_gs.col > 0 and weight_gs.row > 0:
+            return False, "Block-quantized weights not handled by FP8 W8A8 kernel."
+        n = config.weight_shape[0]
+        if n % 32 != 0:
+            # float8_linear_prepack_cpu tiles N in chunks of 32 and cannot
+            # handle a remainder tile smaller than that.
+            return False, f"requires weight output dim (N={n}) to be a multiple of 32."
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Prepack weights with float8_linear_prepack_cpu (VNNI + block layout)."""
+        w_name, ws_name, xs_name, _ = self.layer_param_names
+        weight = getattr(layer, w_name)
+        weight_scale = getattr(layer, ws_name)
+
+        # vLLM stores linear weight as [K, N], while the CPU prepack op expects [N, K].
+        N = weight.shape[-1]
+
+        # Ensure scale is 2D [N, G]
+        ws = weight_scale
+        if ws is None:
+            raise RuntimeError("FP8 W8A8 kernel requires weight scale.")
+        ws = ws.to(torch.float32)
+
+        # Fused modules may carry per-shard scales (e.g. QKV/MLP shards).
+        # Expand those to per-output-channel scales before prepack.
+        if (
+            len(getattr(layer, "logical_widths", [])) > 1
+            and ws.dim() <= 1
+            and ws.numel() == len(layer.logical_widths)
+        ):
+            ws = convert_to_channelwise(ws.reshape(-1), layer.logical_widths)
+
+        if ws.dim() == 0:
+            ws = ws.view(1, 1).expand(N, 1)
+        elif ws.dim() == 1:
+            if ws.numel() == 1:
+                ws = ws.view(1, 1).expand(N, 1)
+            elif ws.numel() == N:
+                ws = ws.unsqueeze(1)  # [N] -> [N, 1]
+            else:
+                raise RuntimeError(
+                    "FP8 W8A8 kernel got unsupported 1D weight scale shape: "
+                    f"{tuple(ws.shape)} for N={N}."
+                )
+
+        # Align scale rows to output channels N.
+        if ws.size(0) != N:
+            if ws.size(0) == 1:
+                ws = ws.expand(N, ws.size(1)).contiguous()
+            else:
+                raise RuntimeError(
+                    "FP8 W8A8 kernel got mismatched weight scale rows: "
+                    f"ws_rows={ws.size(0)}, expected N={N}."
+                )
+
+        w_2d = weight.t().contiguous()  # [N, K]
+        packed_weight, packed_scales = torch.ops._C.float8_linear_prepack_cpu(w_2d, ws)
+
+        replace_parameter(
+            layer,
+            w_name,
+            torch.nn.Parameter(packed_weight, requires_grad=False),
+        )
+        replace_parameter(
+            layer,
+            ws_name,
+            torch.nn.Parameter(packed_scales, requires_grad=False),
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        w_name, ws_name, xs_name, _ = self.layer_param_names
+        weight = getattr(layer, w_name)
+        weight_scale = getattr(layer, ws_name)
+        # For static quantization, use pre-computed activation scale
+        act_scales = getattr(layer, xs_name, None) if xs_name else None
+
+        x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
+        # channelwise=True: per-token (per-row) quantization when act_scales is None
+        # channelwise=False: per-tensor quantization when act_scales is provided
+        channelwise = act_scales is None
+        out = ops.fp8_scaled_mm_with_quant(
+            x_2d,
+            act_scales,
+            channelwise,
+            weight,
+            weight_scale,
+            bias,
+            x.dtype,
+        )
+        return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
+
+    def apply_scaled_mm(self, A, B, out_dtype, As, Bs, bias=None, **kwargs):
+        raise NotImplementedError(
+            "CPUFP8W8A8ScaledMMLinearKernel uses apply_weights directly."
+        )
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
index 9f69ab0c7377..da8c69eed0a2 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
@@ -17,6 +17,7 @@
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
     CUTLASS_BLOCK_FP8_SUPPORTED,
     convert_to_channelwise,
+    cutlass_fp8_supported,
 )
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
@@ -164,8 +165,8 @@ def __init__(
     def is_supported(
         cls, compute_capability: int | None = None
     ) -> tuple[bool, str | None]:
-        if not current_platform.is_cuda():
-            return False, "requires CUDA."
+        if not cutlass_fp8_supported():
+            return False, "CUTLASS FP8 kernels not available"
         return True, None
 
     @classmethod
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py b/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
index 79458f5cda85..19f9347e311e 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
@@ -86,25 +86,28 @@ def process_weights_after_loading(self, layer):
         params = self._get_layer_params(layer)
         assert layer.weight_block_size is not None
 
+        is_bmm = getattr(layer, "is_bmm", False)
+
         if self.is_deep_gemm_supported:
-            weight_scale_invs = params.weight_scale_inv
-            scale_attr = (
-                params.WEIGHT_SCALE_INV
-                if weight_scale_invs is not None
-                else params.WEIGHT_SCALE
-            )
             dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block(
                 wq=params.weight,
-                ws=weight_scale_invs
-                if weight_scale_invs is not None
-                else params.weight_scale,
+                ws=params.block_scale,
                 quant_block_shape=tuple(layer.weight_block_size),
                 use_e8m0=self.use_deep_gemm_e8m0,
-                is_bmm=getattr(layer, "is_bmm", False),
+                is_bmm=is_bmm,
                 bmm_batch_size=getattr(layer, "bmm_batch_size", 0),
             )
             replace_parameter(layer, params.WEIGHT, dg_weight)
-            replace_parameter(layer, scale_attr, dg_weight_scale)
+            replace_parameter(layer, params.block_scale_attr, dg_weight_scale)
+        # bmm layers bypass this kernel and run through fp8_einsum.
+        if not is_bmm:
+            layer.deep_gemm_warmup_provider = self
+
+    def get_deep_gemm_warmup_weights(
+        self, layer: torch.nn.Module
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        params = self._get_layer_params(layer)
+        return params.weight, params.block_scale
 
     def apply_block_scaled_mm(
         self,
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py
index 72a3b849840d..1407530e9c78 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py
@@ -149,8 +149,7 @@ def apply_block_scaled_mm(
 class FlashInferFp8DeepGEMMDynamicBlockScaledKernel(
     Fp8BlockScaledDynamicMMLinearKernel
 ):
-    """
-    Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM.
+    """Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM.
 
     Dispatches between two kernels based on input batch size:
     - Small batches (M < 32): FlashInfer's swapAB trick for better utilisation.
@@ -214,9 +213,7 @@ def _flashinfer_fp8_blockscale_gemm_fake(
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
 ) -> torch.Tensor:
-    """
-    Required fake/meta implementation for torch.compile graph tracing.
-    """
+    """Required fake/meta implementation for torch.compile graph tracing."""
     return torch.empty(
         input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
     )
@@ -236,8 +233,7 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl(
     group_size: int,
     use_deep_gemm_e8m0: bool,
 ) -> torch.Tensor:
-    """
-    Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection.
+    """Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection.
 
     This function switches between two optimized kernels based on the input batch size:
     - For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization.
@@ -260,6 +256,7 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl(
 
     Returns:
         Output tensor of shape (batch_size, output_dim) in bfloat16 format
+
     """
 
     def run_flashinfer_deepgemm_swapAB(
@@ -323,9 +320,7 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_fake(
     group_size: int,
     use_deep_gemm_e8m0: bool,
 ) -> torch.Tensor:
-    """
-    Required fake/meta implementation for torch.compile graph tracing.
-    """
+    """Required fake/meta implementation for torch.compile graph tracing."""
     return torch.empty(
         input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
     )
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/humming.py b/vllm/model_executor/kernels/linear/scaled_mm/humming.py
index b52e73b5a423..3381a3381f86 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/humming.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/humming.py
@@ -4,7 +4,7 @@
 import torch
 
 from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
+from vllm.model_executor.layers.quantization.utils.humming import (
     apply_humming_linear,
     convert_linear_layer_to_humming_standard,
     get_humming_linear_compute_config,
@@ -92,6 +92,7 @@ def apply_weights(
         return apply_humming_linear(
             layer,
             x,
+            skip_bias_add=bias is None,
             layer_config=self.layer_config,
             compute_config=self.compute_config,
             locks=self.locks,
@@ -158,6 +159,7 @@ def apply_weights(
         return apply_humming_linear(
             layer,
             x,
+            skip_bias_add=bias is None,
             layer_config=self.layer_config,
             compute_config=self.compute_config,
             locks=self.locks,
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py
index 9e148ca308cf..1df6ad55412d 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py
@@ -27,8 +27,7 @@
 
 
 class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
-    """
-    FP8 Marlin kernel for GPUs that lack FP8 hardware support.
+    """FP8 Marlin kernel for GPUs that lack FP8 hardware support.
     Leverages the Marlin kernel for fast weight-only FP8 quantization.
     """
 
@@ -94,7 +93,7 @@ def apply_weights(
             input=x,
             weight=layer.weight,
             weight_scale=weight_scale,
-            workspace=layer.workspace,
+            workspace=None,
             size_n=layer.output_size_per_partition,
             size_k=layer.input_size_per_partition,
             input_dtype=self.marlin_input_dtype,
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
index 4fc7ace5a1fe..1f2756cc267e 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
+++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
@@ -50,8 +50,7 @@ def _supports_torch_fp8_scaled_mm() -> bool:
 
 
 class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
-    """
-    Base class for FP8 linear kernels using Torch.
+    """Base class for FP8 linear kernels using Torch.
     Each subclass represents a kernel variant for
     specific device capabilities and torch versions.
     """
diff --git a/vllm/model_executor/kernels/mhc/__init__.py b/vllm/model_executor/kernels/mhc/__init__.py
index 324e3cec6400..f9df7f6f2c33 100644
--- a/vllm/model_executor/kernels/mhc/__init__.py
+++ b/vllm/model_executor/kernels/mhc/__init__.py
@@ -12,6 +12,7 @@
     "mhc_fused_post_pre_cuda",
     "hc_head_fused_kernel_cuda",
     "mhc_pre_aiter",
+    "mhc_pre_delayed_aiter",
     "mhc_post_aiter",
     "mhc_fused_post_pre_aiter",
     "hc_head_fused_aiter",
@@ -26,6 +27,7 @@
     "mhc_post_torch",
     "mhc_fused_post_pre_torch",
     "hc_head_fused_torch",
+    "mhc_pre_mix_triton",
     "mhc_pre_triton",
     "mhc_post_triton",
     "mhc_fused_post_pre_triton",
diff --git a/vllm/model_executor/kernels/mhc/aiter.py b/vllm/model_executor/kernels/mhc/aiter.py
index ce7690ec1d21..06b4f1498463 100644
--- a/vllm/model_executor/kernels/mhc/aiter.py
+++ b/vllm/model_executor/kernels/mhc/aiter.py
@@ -19,8 +19,7 @@ def mhc_pre_aiter(
     norm_weight: torch.Tensor | None = None,
     norm_eps: float = 0.0,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Forward pass for mHC pre block.
+    """Forward pass for mHC pre block.
 
     Args:
         residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
@@ -40,8 +39,8 @@ def mhc_pre_aiter(
         post_mix: shape (..., hc_mult), dtype torch.float32
         comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
         layer_input: shape (..., hidden_size), dtype torch.bfloat16
-    """
 
+    """
     hidden_size = residual.shape[-1]
     assert hidden_size % 256 == 0
     from vllm._aiter_ops import rocm_aiter_ops
@@ -104,6 +103,110 @@ def _mhc_pre_aiter_fake(
     return post_mix, comb_mix, layer_input
 
 
+def mhc_pre_delayed_aiter(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    pre_mix: torch.Tensor | None = None,
+    sublayer_out: torch.Tensor | None = None,
+    post_layer_mix: torch.Tensor | None = None,
+    comb_res_mix: torch.Tensor | None = None,
+    residual_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """MHC pre with the pre-mix carried in from the previous sublayer.
+
+    Matches ``mhc_pre_delayed_torch``: the stream collapse uses *pre_mix*
+    rather than the gate computed here, and that gate is returned as the
+    pre-mix for the next sublayer seam.
+
+    Args:
+        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
+        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
+        hc_scale: shape (3,), dtype torch.float32
+        hc_base: shape (hc_mult3,), dtype torch.float32
+        rms_eps: RMS normalization epsilon
+        hc_pre_eps: pre-mix epsilon
+        hc_sinkhorn_eps: sinkhorn epsilon
+        hc_post_mult_value: post-mix multiplier value
+        sinkhorn_repeat: number of sinkhorn iterations
+        pre_mix: shape (..., hc_mult) from the previous sublayer, or None at
+            model entry to select residual stream zero.
+        sublayer_out: attention or FFN output, shape (..., hidden_size). When
+            given with the two mixes below, the preceding post block is
+            applied here so AITER can fold it into the pre projection.
+        post_layer_mix: shape (..., hc_mult, 1), post gate for that block.
+        comb_res_mix: shape (..., hc_mult, hc_mult), residual comb for it.
+        residual_out: shape (..., hc_mult, hidden_size), written with the post
+            block's new residual. Required exactly when the post is folded in.
+
+    Returns:
+        post_mix: shape (..., hc_mult, 1), dtype torch.float32
+        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
+        layer_input: shape (..., hidden_size), dtype torch.bfloat16
+        next_pre_mix: shape (..., hc_mult), dtype torch.float32
+
+    """
+    hidden_size = residual.shape[-1]
+    assert hidden_size % 256 == 0
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    return rocm_aiter_ops.mhc_pre_delayed(
+        residual,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_mult_value,
+        sinkhorn_repeat,
+        pre_mix,
+        sublayer_out,
+        post_layer_mix,
+        comb_res_mix,
+        residual_out,
+    )
+
+
+def _mhc_pre_delayed_aiter_fake(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    pre_mix: torch.Tensor | None = None,
+    sublayer_out: torch.Tensor | None = None,
+    post_layer_mix: torch.Tensor | None = None,
+    comb_res_mix: torch.Tensor | None = None,
+    residual_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    outer_shape = residual.shape[:-2]
+    return (
+        torch.empty(
+            *outer_shape, hc_mult, 1, dtype=torch.float32, device=residual.device
+        ),
+        torch.empty(
+            *outer_shape, hc_mult, hc_mult, dtype=torch.float32, device=residual.device
+        ),
+        torch.empty(
+            *outer_shape, hidden_size, dtype=torch.bfloat16, device=residual.device
+        ),
+        torch.empty(*outer_shape, hc_mult, dtype=torch.float32, device=residual.device),
+    )
+
+
 def mhc_post_aiter(
     x: torch.Tensor,
     residual: torch.Tensor,
@@ -226,6 +329,12 @@ def _mhc_fused_post_pre_aiter_fake(
     mutates_args=[],
     fake_impl=_mhc_pre_aiter_fake,
 )
+direct_register_custom_op(
+    op_name="mhc_pre_delayed_aiter",
+    op_func=mhc_pre_delayed_aiter,
+    mutates_args=["residual_out"],
+    fake_impl=_mhc_pre_delayed_aiter_fake,
+)
 direct_register_custom_op(
     op_name="mhc_post_aiter",
     op_func=mhc_post_aiter,
diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py
index c4fe6dd16328..62741bbda19c 100644
--- a/vllm/model_executor/kernels/mhc/tilelang.py
+++ b/vllm/model_executor/kernels/mhc/tilelang.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import torch
 
+from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import direct_register_custom_op
 
 
@@ -18,73 +19,56 @@ def _torch_hc_prenorm_gemm(
     sqrsum[0].copy_(x_float.square().sum(dim=-1))
 
 
-def _tilelang_hc_prenorm_gemm(
+def _hc_prenorm_gemm_outputs(
     x: torch.Tensor,
     fn: torch.Tensor,
-    out: torch.Tensor,
-    sqrsum: torch.Tensor,
+    *,
     hidden_size: int,
     hc_mult: int,
-    tile_n: int = 12,
-    n_thr: int = 512,
-    n_splits: int = 1,
-) -> None:
+    use_tilelang_fallback: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
     from vllm.model_executor.kernels.mhc.tilelang_kernels import (
-        hc_prenorm_gemm_block_m_tilelang,
-        hc_prenorm_gemm_tilelang,
-    )
-
-    assert out.shape[0] == n_splits
-    assert sqrsum.shape[0] == n_splits
-    assert x.shape[1] == hc_mult * hidden_size
-    assert x.shape[1] % n_splits == 0
-    assert (x.shape[1] // n_splits) % n_thr == 0
-    use_default_config = tile_n == 12 and n_thr == 512
-    if n_splits == 1 and use_default_config and x.shape[0] >= 1024:
-        hc_prenorm_gemm_block_m_tilelang(
-            x,
-            fn,
-            out,
-            sqrsum,
-            hidden_size,
-            hc_mult,
-            fn.shape[0],
-            n_thr,
-            tile_n,
-            2,
+        compute_num_split,
+    )
+    from vllm.utils.deep_gemm import (
+        is_deep_gemm_supported,
+        tf32_hc_prenorm_gemm,
+    )
+
+    use_deep_gemm = is_deep_gemm_supported() or not use_tilelang_fallback
+    num_tokens = x.shape[0]
+    n_splits = (
+        compute_num_split(64, x.shape[1], cdiv(num_tokens, 64)) if use_deep_gemm else 1
+    )
+    out = torch.empty(
+        n_splits,
+        num_tokens,
+        fn.shape[0],
+        dtype=torch.float32,
+        device=x.device,
+    )
+    sqrsum = torch.empty(
+        n_splits,
+        num_tokens,
+        dtype=torch.float32,
+        device=x.device,
+    )
+    if use_deep_gemm:
+        tf32_hc_prenorm_gemm(x, fn, out, sqrsum, n_splits)
+    else:
+        from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+            _HC_PRENORM_GEMM_TILELANG_KERNEL,
         )
-        return
-    if (
-        n_splits == 1
-        and use_default_config
-        and x.shape[0] < 128
-        and x.shape[1] % 1024 == 0
-    ):
-        hc_prenorm_gemm_tilelang(
+
+        _HC_PRENORM_GEMM_TILELANG_KERNEL(
             x,
             fn,
             out,
             sqrsum,
             hidden_size,
             hc_mult,
-            fn.shape[0],
-            1024,
-            4,
-            n_splits,
         )
-        return
-    hc_prenorm_gemm_tilelang(
-        x,
-        fn,
-        out,
-        sqrsum,
-        hidden_size,
-        hc_mult,
-        fn.shape[0],
-        n_thr,
-        tile_n,
-        n_splits,
-    )
+    return out, sqrsum
 
 
 def mhc_pre_delayed_tilelang(
@@ -125,8 +109,10 @@ def mhc_pre_delayed_tilelang(
         Post and residual coefficients, optionally normalized BF16 layer input,
         and the next FP32 pre-mix, with shapes (tokens, hc_mult, 1),
         (tokens, hc_mult, hc_mult), (tokens, hidden_size), and (tokens, hc_mult).
+
     """
     from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+        _HC_PRENORM_GEMM_TILELANG_KERNEL,
         mhc_pre_big_fuse_tilelang,
     )
     from vllm.model_executor.kernels.mhc.warmup import (
@@ -186,7 +172,14 @@ def mhc_pre_delayed_tilelang(
     if use_deep_gemm:
         tf32_hc_prenorm_gemm(x, fn, mixes, sqrsum, n_splits)
     else:
-        _tilelang_hc_prenorm_gemm(x, fn, mixes, sqrsum, input_size, 1)
+        _HC_PRENORM_GEMM_TILELANG_KERNEL(
+            x,
+            fn,
+            mixes,
+            sqrsum,
+            input_size,
+            1,
+        )
     if norm_weight is not None:
         assert norm_weight.shape == (hidden_size,)
         assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous()
@@ -202,6 +195,7 @@ def mhc_pre_delayed_tilelang(
             norm_weight,
             pre_mix if pre_mix is not None else post,
             next_pre_mix,
+            layer_input,
             hidden_size=hidden_size,
             rms_eps=rms_eps,
             hc_pre_eps=hc_pre_eps,
@@ -226,6 +220,7 @@ def mhc_pre_delayed_tilelang(
         layer_input,
         pre_mix if pre_mix is not None else post,
         next_pre_mix,
+        layer_input,
         hidden_size,
         rms_eps,
         hc_pre_eps,
@@ -271,6 +266,279 @@ def _mhc_pre_delayed_tilelang_fake(
     )
 
 
+def mhc_fused_post_pre_delayed_tilelang(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    pre_mix: torch.Tensor | None = None,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+    capture_aux: bool = False,
+) -> tuple[
+    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+]:
+    """Run one mHC post block followed by the next delayed mHC pre block.
+
+    Within the fused kernel's token range the post mapping is folded into the
+    pre-norm GEMM, so the updated residual streams feed the projection from
+    registers instead of a second pass over global memory. Above it this runs
+    the same post kernel and split-k GEMM as the unfused pair.
+
+    Args:
+        x: BF16 sublayer output of shape (tokens, hidden_size).
+        residual: BF16 residual streams of shape (tokens, hc_mult, hidden_size).
+        post_layer_mix: FP32 post coefficients of shape (tokens, hc_mult, 1).
+        comb_res_mix: FP32 residual coefficients, (tokens, hc_mult, hc_mult).
+        fn: FP32 projection of shape (hc_mult * (hc_mult + 2), input_size).
+        hc_scale: FP32 scales of shape (3,).
+        hc_base: FP32 bias of shape (hc_mult * (hc_mult + 2),).
+        rms_eps: RMS normalization epsilon.
+        hc_pre_eps: Pre-mix epsilon.
+        hc_sinkhorn_eps: Sinkhorn epsilon.
+        hc_post_mult_value: Post-mix multiplier.
+        sinkhorn_repeat: Number of Sinkhorn iterations.
+        pre_mix: FP32 coefficients from the previous sublayer, or None to
+            select residual stream zero.
+        norm_weight: Optional BF16 RMSNorm weight for the collapsed input.
+        norm_eps: RMSNorm epsilon for the collapsed input.
+        capture_aux: Also return the mean over the post-mapped streams, which
+            draft models consume as the target's hidden state. It is folded
+            into the collapse, which already reads those streams.
+
+    Returns:
+        The post-mapped residual streams, the post and residual coefficients,
+        the optionally normalized BF16 layer input, the next FP32 pre-mix, and
+        the BF16 stream mean (empty unless capture_aux), with shapes
+        (tokens, hc_mult, hidden_size), (tokens, hc_mult, 1),
+        (tokens, hc_mult, hc_mult), (tokens, hidden_size), (tokens, hc_mult),
+        and (tokens, hidden_size).
+
+    """
+    from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+        _HC_PRENORM_GEMM_TILELANG_KERNEL,
+        _MHC_FUSED_TILELANG_KERNEL,
+        _MHC_POST_TILELANG_KERNEL,
+        mhc_fused_post_pre_split_config,
+        mhc_pre_big_fuse_tilelang,
+    )
+    from vllm.model_executor.kernels.mhc.warmup import (
+        MHC_PRE_NORM_KERNEL,
+        compute_mhc_pre_num_splits,
+    )
+    from vllm.utils.deep_gemm import (
+        is_deep_gemm_supported,
+        tf32_hc_prenorm_gemm,
+    )
+
+    assert residual.ndim == 3 and residual.dtype == torch.bfloat16
+    assert residual.is_contiguous()
+    num_tokens, hc_mult, hidden_size = residual.shape
+    input_size = hc_mult * hidden_size
+    mix_size = hc_mult * (hc_mult + 2)
+    assert x.shape == (num_tokens, hidden_size) and x.dtype == torch.bfloat16
+    assert x.is_contiguous()
+    assert post_layer_mix.shape[:2] == (num_tokens, hc_mult)
+    assert post_layer_mix.dtype == torch.float32 and post_layer_mix.is_contiguous()
+    assert comb_res_mix.shape == (num_tokens, hc_mult, hc_mult)
+    assert comb_res_mix.dtype == torch.float32 and comb_res_mix.is_contiguous()
+    assert fn.shape == (mix_size, input_size) and fn.dtype == torch.float32
+    assert hc_scale.shape == (3,) and hc_scale.dtype == torch.float32
+    assert hc_base.shape == (mix_size,) and hc_base.dtype == torch.float32
+    if pre_mix is not None:
+        assert pre_mix.shape == (num_tokens, hc_mult)
+        assert pre_mix.dtype == torch.float32 and pre_mix.is_contiguous()
+
+    next_pre_mix = torch.empty(
+        num_tokens, hc_mult, dtype=torch.float32, device=residual.device
+    )
+    post = torch.empty_like(next_pre_mix)
+    comb = torch.empty(
+        num_tokens, hc_mult * hc_mult, dtype=torch.float32, device=residual.device
+    )
+    layer_input = torch.empty(
+        num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
+    )
+    aux = torch.empty(
+        num_tokens if capture_aux else 0,
+        hidden_size,
+        dtype=torch.bfloat16,
+        device=residual.device,
+    )
+    if num_tokens == 0:
+        return (
+            torch.empty_like(residual),
+            post.unsqueeze(-1),
+            comb.view(num_tokens, hc_mult, hc_mult),
+            layer_input,
+            next_pre_mix,
+            aux,
+        )
+
+    fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
+    if fused_config is not None:
+        mixes, sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL(
+            comb_res_mix,
+            residual,
+            post_layer_mix.view(num_tokens, hc_mult),
+            x,
+            fn.view(mix_size, hc_mult, hidden_size),
+            hc_mult,
+            hidden_size,
+            mix_size,
+        )
+    else:
+        residual_cur = torch.empty_like(residual)
+        _MHC_POST_TILELANG_KERNEL(
+            comb_res_mix,
+            residual,
+            post_layer_mix.view(num_tokens, hc_mult),
+            x,
+            residual_cur,
+            hc_mult,
+            hidden_size,
+        )
+        # The delayed epilogue is compiled per bucketed split count, so the
+        # projection has to use the same bucket rather than the raw estimate.
+        use_deep_gemm = is_deep_gemm_supported()
+        n_splits = (
+            compute_mhc_pre_num_splits(input_size, num_tokens) if use_deep_gemm else 1
+        )
+        mixes = torch.empty(
+            n_splits, num_tokens, mix_size, dtype=torch.float32, device=residual.device
+        )
+        sqrsum = torch.empty(
+            n_splits, num_tokens, dtype=torch.float32, device=residual.device
+        )
+        residual_cur_2d = residual_cur.view(num_tokens, input_size)
+        if use_deep_gemm:
+            tf32_hc_prenorm_gemm(residual_cur_2d, fn, mixes, sqrsum, n_splits)
+        else:
+            _HC_PRENORM_GEMM_TILELANG_KERNEL(
+                residual_cur_2d,
+                fn,
+                mixes,
+                sqrsum,
+                input_size,
+                1,
+            )
+
+    outputs = (
+        residual_cur,
+        post.unsqueeze(-1),
+        comb.view(num_tokens, hc_mult, hc_mult),
+        layer_input,
+        next_pre_mix,
+        aux,
+    )
+    if norm_weight is not None:
+        assert norm_weight.shape == (hidden_size,)
+        assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous()
+        MHC_PRE_NORM_KERNEL(
+            mixes,
+            sqrsum,
+            hc_scale,
+            hc_base,
+            residual_cur,
+            post,
+            comb,
+            layer_input,
+            norm_weight,
+            pre_mix if pre_mix is not None else post,
+            next_pre_mix,
+            aux if capture_aux else layer_input,
+            hidden_size=hidden_size,
+            rms_eps=rms_eps,
+            hc_pre_eps=hc_pre_eps,
+            hc_sinkhorn_eps=hc_sinkhorn_eps,
+            hc_post_mult_value=hc_post_mult_value,
+            sinkhorn_repeat=sinkhorn_repeat,
+            norm_eps=norm_eps,
+            hc_mult=hc_mult,
+            use_pre_mix_in=pre_mix is not None,
+            save_pre_mix=True,
+            rms_numel=input_size,
+            write_aux=capture_aux,
+        )
+        return outputs
+    mhc_pre_big_fuse_tilelang(
+        mixes,
+        sqrsum,
+        hc_scale,
+        hc_base,
+        residual_cur,
+        post,
+        comb,
+        layer_input,
+        pre_mix if pre_mix is not None else post,
+        next_pre_mix,
+        aux if capture_aux else layer_input,
+        hidden_size,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_mult_value,
+        sinkhorn_repeat,
+        mixes.shape[0],
+        hc_mult,
+        use_pre_mix_in=pre_mix is not None,
+        save_pre_mix=True,
+        rms_numel=input_size,
+        write_aux=capture_aux,
+    )
+    return outputs
+
+
+def _mhc_fused_post_pre_delayed_tilelang_fake(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    pre_mix: torch.Tensor | None = None,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+    capture_aux: bool = False,
+) -> tuple[
+    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+]:
+    num_tokens, hc_mult, hidden_size = residual.shape
+    return (
+        torch.empty_like(residual),
+        torch.empty(
+            num_tokens, hc_mult, 1, dtype=torch.float32, device=residual.device
+        ),
+        torch.empty(
+            num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=residual.device
+        ),
+        torch.empty(
+            num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
+        ),
+        torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=residual.device),
+        torch.empty(
+            num_tokens if capture_aux else 0,
+            hidden_size,
+            dtype=torch.bfloat16,
+            device=residual.device,
+        ),
+    )
+
+
 def mhc_pre_tilelang(
     residual: torch.Tensor,
     fn: torch.Tensor,
@@ -285,8 +553,7 @@ def mhc_pre_tilelang(
     norm_weight: torch.Tensor | None = None,
     norm_eps: float = 1e-6,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Forward pass for mHC pre block.
+    """Forward pass for mHC pre block.
 
     Args:
         residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
@@ -298,7 +565,8 @@ def mhc_pre_tilelang(
         hc_sinkhorn_eps: sinkhorn epsilon
         hc_post_mult_value: post-mix multiplier value
         sinkhorn_repeat: number of sinkhorn iterations
-        n_splits: split-k factor;
+        n_splits: retained for the shared MHC operator API; the active GEMM
+            backend selects its split factor internally.
         norm_weight: optional RMSNorm weight, shape (hidden_size,), dtype
             torch.bfloat16. When provided, RMSNorm is fused into the
             layer_input write path of the big_fuse kernel.
@@ -309,14 +577,11 @@ def mhc_pre_tilelang(
         post_mix: shape (..., hc_mult), dtype torch.float32
         comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
         layer_input: shape (..., hidden_size), dtype torch.bfloat16
+
     """
     from vllm.model_executor.kernels.mhc.tilelang_kernels import (
-        compute_num_split,
-        mhc_pre_big_fuse_tilelang,
-        mhc_pre_big_fuse_with_norm_tilelang,
+        _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
     )
-    from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
-    from vllm.utils.math_utils import cdiv
 
     assert residual.dtype == torch.bfloat16
     assert fn.dtype == torch.float32
@@ -346,17 +611,6 @@ def mhc_pre_tilelang(
     residual_flat = residual.view(-1, hc_mult, hidden_size)
     num_tokens = residual_flat.shape[0]
 
-    from vllm.utils.deep_gemm import is_deep_gemm_supported
-
-    use_deep_gemm = is_deep_gemm_supported()
-    if use_deep_gemm:
-        # these numbers are from deepgemm kernel impl
-        block_k = 64
-        block_m = 64
-        n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
-    else:
-        n_splits = 1
-
     post_mix = torch.empty(
         num_tokens, hc_mult, dtype=torch.float32, device=residual.device
     )
@@ -367,78 +621,30 @@ def mhc_pre_tilelang(
         num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
     )
 
-    gemm_out_mul = torch.empty(
-        n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
+    gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
+        residual_flat.view(num_tokens, hc_mult * hidden_size),
+        fn,
+        hidden_size=hidden_size,
+        hc_mult=hc_mult,
     )
-    gemm_out_sqrsum = torch.empty(
-        n_splits, num_tokens, dtype=torch.float32, device=residual.device
+    _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
+        gemm_out_mul,
+        gemm_out_sqrsum,
+        hc_scale,
+        hc_base,
+        residual_flat,
+        post_mix,
+        comb_mix,
+        layer_input,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_mult_value,
+        sinkhorn_repeat,
+        norm_weight=norm_weight,
+        norm_eps=norm_eps,
     )
 
-    residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size)
-    if use_deep_gemm:
-        tf32_hc_prenorm_gemm(
-            residual_2d,
-            fn,
-            gemm_out_mul,
-            gemm_out_sqrsum,
-            n_splits,
-        )
-    else:
-        _tilelang_hc_prenorm_gemm(
-            residual_2d,
-            fn,
-            gemm_out_mul,
-            gemm_out_sqrsum,
-            hidden_size,
-            hc_mult,
-        )
-
-    if norm_weight is None:
-        mhc_pre_big_fuse_tilelang(
-            gemm_out_mul,
-            gemm_out_sqrsum,
-            hc_scale,
-            hc_base,
-            residual_flat,
-            post_mix,
-            comb_mix,
-            layer_input,
-            # Pre-mix buffers are unused in V4 mode.
-            post_mix,
-            post_mix,
-            hidden_size,
-            rms_eps,
-            hc_pre_eps,
-            hc_sinkhorn_eps,
-            hc_post_mult_value,
-            sinkhorn_repeat,
-            n_splits,
-            hc_mult,
-        )
-    else:
-        mhc_pre_big_fuse_with_norm_tilelang(
-            gemm_out_mul,
-            gemm_out_sqrsum,
-            hc_scale,
-            hc_base,
-            residual_flat,
-            post_mix,
-            comb_mix,
-            layer_input,
-            norm_weight,
-            post_mix,
-            post_mix,
-            hidden_size,
-            rms_eps,
-            hc_pre_eps,
-            hc_sinkhorn_eps,
-            hc_post_mult_value,
-            sinkhorn_repeat,
-            norm_eps,
-            n_splits,
-            hc_mult,
-        )
-
     return (
         post_mix.view(*outer_shape, hc_mult, 1),
         comb_mix.view(*outer_shape, hc_mult, hc_mult),
@@ -505,11 +711,10 @@ def mhc_pre_broadcast_tilelang(
     fn_broadcast: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """First-layer mHC pre for a residual broadcast from ``(T, H)``."""
+    # n_splits is retained for the shared API; split selection is internal.
     from vllm.model_executor.kernels.mhc.tilelang_kernels import (
-        compute_num_split,
-        mhc_pre_big_fuse_broadcast_with_norm_tilelang,
+        _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
     )
-    from vllm.utils.math_utils import cdiv
 
     assert norm_weight is not None, "broadcast mHC pre currently requires fused RMSNorm"
     assert residual.dtype == torch.bfloat16
@@ -537,8 +742,6 @@ def mhc_pre_broadcast_tilelang(
     residual_flat = residual
     num_tokens = residual.shape[0]
 
-    n_splits = compute_num_split(64, hidden_size, cdiv(num_tokens, 64))
-
     residual_out = torch.empty(
         num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=residual.device
     )
@@ -551,42 +754,31 @@ def mhc_pre_broadcast_tilelang(
     layer_input = torch.empty(
         num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
     )
-    gemm_out_mul = torch.empty(
-        n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
-    )
-    gemm_out_sqrsum = torch.empty(
-        n_splits, num_tokens, dtype=torch.float32, device=residual.device
-    )
 
-    from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
-
-    tf32_hc_prenorm_gemm(
+    gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
         residual_flat,
         fn_broadcast,
-        gemm_out_mul,
-        gemm_out_sqrsum,
-        n_splits,
+        hidden_size=hidden_size,
+        hc_mult=hc_mult,
+        use_tilelang_fallback=False,
     )
-    mhc_pre_big_fuse_broadcast_with_norm_tilelang(
+    _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
         gemm_out_mul,
         gemm_out_sqrsum,
         hc_scale,
         hc_base,
         residual_flat,
-        residual_out,
         post_mix,
         comb_mix,
         layer_input,
-        norm_weight,
-        hidden_size,
         rms_eps,
         hc_pre_eps,
         hc_sinkhorn_eps,
         hc_post_mult_value,
         sinkhorn_repeat,
-        norm_eps,
-        n_splits,
-        hc_mult,
+        residual_out=residual_out,
+        norm_weight=norm_weight,
+        norm_eps=norm_eps,
     )
     return (
         residual_out,
@@ -603,11 +795,11 @@ def mhc_post_tilelang(
     comb_res_mix: torch.Tensor,
 ) -> torch.Tensor:
     from vllm.model_executor.kernels.mhc.tilelang_kernels import (
-        mhc_post_tilelang as _mhc_post_kernel,
+        _MHC_POST_TILELANG_KERNEL,
     )
 
     out = torch.empty_like(residual)
-    _mhc_post_kernel(
+    _MHC_POST_TILELANG_KERNEL(
         comb_res_mix,
         residual,
         post_layer_mix.squeeze(-1),
@@ -637,28 +829,28 @@ def mhc_fused_post_pre_tilelang(
     norm_weight: torch.Tensor | None = None,
     norm_eps: float = 1e-6,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Run one MHC post block followed by the next MHC pre block.
+    """Run one MHC post block followed by the next MHC pre block.
 
     When ``norm_weight`` is provided, the layer_input_cur output is the
     RMSNorm'd activation (fused into the kernel); otherwise it is the
     raw pre-norm activation as before.
 
+    ``n_splits`` and ``tile_n`` are retained for the shared MHC operator API.
+    The TileLang path selects both values internally from the runtime shape.
+
     Returns:
         residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size)
         post_mix_cur: shape (..., hc_mult, 1)
         comb_mix_cur: shape (..., hc_mult, hc_mult)
         layer_input_cur: shape (..., hidden_size)
-    """
 
+    """
     from vllm.model_executor.kernels.mhc.tilelang_kernels import (
-        compute_num_split,
-        mhc_fused_tilelang,
-        mhc_post_tilelang,
-        mhc_pre_big_fuse_tilelang,
-        mhc_pre_big_fuse_with_norm_tilelang,
+        _MHC_FUSED_TILELANG_KERNEL,
+        _MHC_POST_TILELANG_KERNEL,
+        _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
+        mhc_fused_post_pre_split_config,
     )
-    from vllm.utils.math_utils import cdiv
 
     assert residual.dtype == torch.bfloat16
     assert x.dtype == torch.bfloat16
@@ -692,48 +884,14 @@ def mhc_fused_post_pre_tilelang(
         if not norm_weight.is_contiguous():
             norm_weight = norm_weight.contiguous()
 
-    assert n_splits in (1, 2, 4, 8)
-    assert hidden_size % n_splits == 0
-
     residual_flat = residual.view(-1, hc_mult, hidden_size)
     num_tokens = residual_flat.shape[0]
     x_flat = x.view(num_tokens, hidden_size)
     post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
     comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
 
-    from vllm.utils.deep_gemm import is_deep_gemm_supported
-
-    use_deep_gemm = is_deep_gemm_supported()
-    use_small_fma = num_tokens <= 16
-    if use_small_fma:
-        # TODO(gnovack): investigate autotuning these heuristics
-        tile_n = 2 if num_tokens < 8 else 3
-        n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
-    else:
-        if use_deep_gemm:
-            # these number are from deepgemm kernel impl
-            block_k = 64
-            block_m = 64
-            n_splits = compute_num_split(
-                block_k, hc_hidden_size, cdiv(num_tokens, block_m)
-            )
-        else:
-            n_splits = 1
+    fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
 
-    gemm_out_mul = torch.empty(
-        n_splits,
-        num_tokens,
-        hc_mult3,
-        dtype=torch.float32,
-        device=residual.device,
-    )
-    gemm_out_sqrsum = torch.empty(
-        n_splits,
-        num_tokens,
-        dtype=torch.float32,
-        device=residual.device,
-    )
-    residual_cur = torch.empty_like(residual_flat)
     post_mix_cur = torch.empty(
         num_tokens,
         hc_mult,
@@ -753,56 +911,18 @@ def mhc_fused_post_pre_tilelang(
         device=residual.device,
     )
 
-    if use_small_fma:
-        mhc_fused_tilelang(
+    if fused_config is not None:
+        gemm_out_mul, gemm_out_sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL(
             comb_res_mix_flat,
             residual_flat,
             post_layer_mix_flat,
             x_flat,
             fn.view(hc_mult3, hc_mult, hidden_size),
-            gemm_out_mul,
-            gemm_out_sqrsum,
-            residual_cur,
             hc_mult,
             hidden_size,
             hc_mult3,
-            tile_n=tile_n,
-            n_splits=n_splits,
         )
-    else:
-        mhc_post_tilelang(
-            comb_res_mix_flat,
-            residual_flat,
-            post_layer_mix_flat,
-            x_flat,
-            residual_cur,
-            residual.shape[-2],
-            residual.shape[-1],
-        )
-
-        residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size)
-        if use_deep_gemm:
-            from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
-
-            tf32_hc_prenorm_gemm(
-                residual_cur_2d,
-                fn,
-                gemm_out_mul,
-                gemm_out_sqrsum,
-                n_splits,
-            )
-        else:
-            _tilelang_hc_prenorm_gemm(
-                residual_cur_2d,
-                fn,
-                gemm_out_mul,
-                gemm_out_sqrsum,
-                hidden_size,
-                hc_mult,
-            )
-
-    if norm_weight is None:
-        mhc_pre_big_fuse_tilelang(
+        _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
             gemm_out_mul,
             gemm_out_sqrsum,
             hc_scale,
@@ -811,20 +931,32 @@ def mhc_fused_post_pre_tilelang(
             post_mix_cur,
             comb_mix_cur,
             layer_input_cur,
-            # Pre-mix buffers are unused in V4 mode.
-            post_mix_cur,
-            post_mix_cur,
-            hidden_size,
             rms_eps,
             hc_pre_eps,
             hc_sinkhorn_eps,
             hc_post_mult_value,
             sinkhorn_repeat,
-            n_splits,
-            hc_mult,
+            norm_weight=norm_weight,
+            norm_eps=norm_eps,
         )
     else:
-        mhc_pre_big_fuse_with_norm_tilelang(
+        residual_cur = torch.empty_like(residual_flat)
+        _MHC_POST_TILELANG_KERNEL(
+            comb_res_mix_flat,
+            residual_flat,
+            post_layer_mix_flat,
+            x_flat,
+            residual_cur,
+            residual.shape[-2],
+            residual.shape[-1],
+        )
+        gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
+            residual_cur.view(num_tokens, hc_mult * hidden_size),
+            fn,
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+        )
+        _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
             gemm_out_mul,
             gemm_out_sqrsum,
             hc_scale,
@@ -833,18 +965,13 @@ def mhc_fused_post_pre_tilelang(
             post_mix_cur,
             comb_mix_cur,
             layer_input_cur,
-            norm_weight,
-            post_mix_cur,
-            post_mix_cur,
-            hidden_size,
             rms_eps,
             hc_pre_eps,
             hc_sinkhorn_eps,
             hc_post_mult_value,
             sinkhorn_repeat,
-            norm_eps,
-            n_splits,
-            hc_mult,
+            norm_weight=norm_weight,
+            norm_eps=norm_eps,
         )
 
     return (
@@ -926,9 +1053,11 @@ def hc_head_fused_kernel_tilelang(
     )
     if num_tokens == 0:
         return out
-    from vllm.model_executor.kernels.mhc.tilelang_kernels import hc_head_fuse_tilelang
+    from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+        _HC_HEAD_FUSED_TILELANG_KERNEL,
+    )
 
-    hc_head_fuse_tilelang(
+    _HC_HEAD_FUSED_TILELANG_KERNEL(
         hs_flat,
         fn,
         hc_scale,
@@ -962,6 +1091,12 @@ def _hc_head_fused_kernel_tilelang_fake(
     mutates_args=[],
     fake_impl=_mhc_pre_delayed_tilelang_fake,
 )
+direct_register_custom_op(
+    op_name="mhc_fused_post_pre_delayed_tilelang",
+    op_func=mhc_fused_post_pre_delayed_tilelang,
+    mutates_args=[],
+    fake_impl=_mhc_fused_post_pre_delayed_tilelang_fake,
+)
 direct_register_custom_op(
     op_name="mhc_pre_tilelang",
     op_func=mhc_pre_tilelang,
diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py
index 1c23d152b916..f7f560f50dd9 100644
--- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py
+++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py
@@ -4,16 +4,42 @@
 from __future__ import annotations
 
 import math
+from dataclasses import dataclass
 from functools import cache
+from typing import Any
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import WarmupIntRange, zip_inputs
+from vllm.model_executor.warmup.jit_warmup_tilelang_helper import (
+    TileLangLaunchSpec,
+    VllmTileLangJitKernel,
+    kernel_launcher,
+    make_tilelang_warmup_tensor,
+)
 from vllm.platforms import current_platform
 from vllm.tilelang_utils import T, tilelang, tilelang_jit
 from vllm.utils.math_utils import cdiv
 
 ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()
 
+# TileLang's HIP AllReduce shuffle phase ignores a nonzero reduction-group
+# offset, so the 64-thread RMSNorm reduction must be wave64-aligned and use
+# threads 0..63.
+_MIX_ON_TRAILING_THREADS = current_platform.is_rocm()
+
+
+def _is_mix_thread():
+    if _MIX_ON_TRAILING_THREADS:
+        return T.get_thread_binding() >= 64
+    return T.get_thread_binding() < 32
+
+
+def _is_collapse_thread():
+    if _MIX_ON_TRAILING_THREADS:
+        return T.get_thread_binding() < 64
+    return T.get_thread_binding() >= 32
+
 
 @cache
 def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
@@ -28,6 +54,69 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
     return split_k
 
 
+# ``mhc_fused_tilelang`` splits the hidden size across a fixed thread block, so
+# every (tile_n, n_thr) pair is one more kernel to compile. Tuned on GB300
+# (hc_mult 4, hidden_size 5120 and 7168) against the separate post + split-k
+# GEMM it replaces; the bands are coarse because the loss against a per-shape
+# optimum is ~1-2%, while one fixed config gives up to 18% at the top of the
+# range.
+_FUSED_POST_PRE_N_THR = 128
+_FUSED_POST_PRE_MAX_TOKENS = 32
+
+
+def mhc_fused_post_pre_split_config(
+    num_tokens: int, hidden_size: int, hc_mult: int
+) -> tuple[int, int, int] | None:
+    """Pick ``(tile_n, n_splits, n_thr)`` for the fused post + pre-norm GEMM.
+
+    Returns None when a separate post kernel followed by a split-k GEMM is the
+    better choice, or when the hidden size does not divide evenly across the
+    fused kernel's block. One source of truth for the dispatch decision, the
+    compile key and the launch, which must agree.
+    """
+    if num_tokens > _FUSED_POST_PRE_MAX_TOKENS:
+        return None
+    n_thr = _FUSED_POST_PRE_N_THR
+    if hidden_size % n_thr:
+        return None
+    # The projection tiles supply the parallelism a small token count cannot,
+    # so the tile grows with the batch: 1.3x at one token and 1.03x at 32
+    # against a split-k GEMM given the split its own estimator computes.
+    if num_tokens < 16:
+        tile_n, n_splits = 2, 8
+    else:
+        tile_n, n_splits = 6, 8
+    # The kernel drops projection tiles and k-slices it cannot fill evenly.
+    mix_size = hc_mult * (hc_mult + 2)
+    while tile_n > 1 and mix_size % tile_n:
+        tile_n -= 1
+    while n_splits > 1 and hidden_size % (n_splits * n_thr):
+        n_splits //= 2
+    return tile_n, n_splits, n_thr
+
+
+def require_fused_post_pre_config(
+    num_tokens: int, hidden_size: int, hc_mult: int
+) -> tuple[int, int, int]:
+    """The config for a shape the caller has already committed to fusing."""
+    config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
+    if config is None:
+        raise ValueError(
+            "the fused mHC post + pre-norm GEMM does not cover num_tokens="
+            f"{num_tokens} at hidden_size={hidden_size}"
+        )
+    return config
+
+
+def mhc_fused_post_pre_splits(hidden_size: int, hc_mult: int) -> tuple[int, ...]:
+    """Every split-k factor the fused post + pre-norm GEMM path can pick."""
+    configs = (
+        mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
+        for num_tokens in range(1, _FUSED_POST_PRE_MAX_TOKENS + 1)
+    )
+    return tuple(sorted({config[1] for config in configs if config is not None}))
+
+
 @tilelang_jit
 def mhc_pre_big_fuse_tilelang(
     gemm_out_mul,
@@ -40,6 +129,7 @@ def mhc_pre_big_fuse_tilelang(
     layer_input,
     pre_mix_in,
     pre_mix_out,
+    aux_out,
     hidden_size: int,
     rms_eps: float,
     hc_pre_eps: float,
@@ -51,6 +141,7 @@ def mhc_pre_big_fuse_tilelang(
     use_pre_mix_in: bool = False,
     save_pre_mix: bool = False,
     rms_numel: int = 0,
+    write_aux: bool = False,
 ):
     """Fuse coefficient generation and residual collapse after the projection.
 
@@ -75,6 +166,7 @@ def mhc_pre_big_fuse_tilelang(
 
     pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
     pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
+    aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
 
     with T.Kernel(num_tokens, threads=96) as i:
         if ENABLE_PDL:
@@ -175,11 +267,24 @@ def mhc_pre_big_fuse_tilelang(
 
                 ol = T.alloc_fragment(hidden_block, T.float32)
                 T.clear(ol)
+                if write_aux:
+                    # Aux consumers (draft models) take the plain stream mean
+                    # of the same residual this collapse reads.
+                    aux = T.alloc_fragment(hidden_block, T.float32)
+                    T.clear(aux)
 
                 for i_hc in T.serial(hc_mult):
                     pre = pre_mix_shared[i_hc]
                     for i1_h in T.Parallel(hidden_block):
                         ol[i1_h] += pre * xl[i_hc, i1_h]
+                        if write_aux:
+                            aux[i1_h] += xl[i_hc, i1_h]
+
+                if write_aux:
+                    for i1_h in T.Parallel(hidden_block):
+                        aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16(
+                            aux[i1_h] / hc_mult
+                        )
 
                 T.copy(ol, layer_input[i, i0_h * hidden_block])
 
@@ -203,6 +308,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
     norm_weight,
     pre_mix_in,
     pre_mix_out,
+    aux_out,
     hidden_size: int,
     rms_eps: float,
     hc_pre_eps: float,
@@ -216,7 +322,12 @@ def mhc_pre_big_fuse_with_norm_tilelang(
     use_pre_mix_in: bool = False,
     save_pre_mix: bool = False,
     rms_numel: int = 0,
+    split_mode: str = "fused",
+    write_aux: bool = False,
 ):
+    # Split modes require shifted mHC: input collapse uses a carried pre-mix.
+    assert split_mode in ("fused", "stats", "input")
+    assert split_mode == "fused" or save_pre_mix
     num_tokens = T.dynamic("num_tokens")
     hc_mult3 = hc_mult * (2 + hc_mult)
     if gemm_last_dim < 0:
@@ -237,28 +348,29 @@ def mhc_pre_big_fuse_with_norm_tilelang(
 
     pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
     pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
+    aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
 
     with T.Kernel(num_tokens, threads=96) as i:
-        rms = T.alloc_fragment(1, T.float32)
-        mixes = T.alloc_fragment(hc_mult3, T.float32)
-        T.clear(mixes)
-        rms[0] = 0
-
         if ENABLE_PDL:
             T.pdl_sync()
-
-        for i_split in T.serial(n_splits):
-            rms[0] += gemm_out_sqrsum[i_split, i]
-        rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps)
-        for j in T.Parallel(hc_mult3):
-            mixes[j] = 0
-            for i_split in T.serial(n_splits):
-                mixes[j] += gemm_out_mul[i_split, i, j]
-            mixes[j] *= rms[0]
         mixes_shared = T.alloc_shared(hc_mult3, T.float32)
-        T.copy(mixes, mixes_shared)
+        if split_mode != "input":
+            rms = T.alloc_fragment(1, T.float32)
+            mixes = T.alloc_fragment(hc_mult3, T.float32)
+            T.clear(mixes)
+            rms[0] = 0
 
-        if T.get_thread_binding() < 32:
+            for i_split in T.serial(n_splits):
+                rms[0] += gemm_out_sqrsum[i_split, i]
+            rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps)
+            for j in T.Parallel(hc_mult3):
+                mixes[j] = 0
+                for i_split in T.serial(n_splits):
+                    mixes[j] += gemm_out_mul[i_split, i, j]
+                mixes[j] *= rms[0]
+            T.copy(mixes, mixes_shared)
+
+        if split_mode != "input" and _is_mix_thread():
             cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
             for j in T.Parallel(hc_mult):
                 if save_pre_mix:
@@ -304,7 +416,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
 
             for j, k in T.Parallel(hc_mult, hc_mult):
                 comb_mix[i, j * hc_mult + k] = cm[j, k]
-        else:
+        elif split_mode != "stats" and _is_collapse_thread():
             pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
             for j in T.Parallel(hc_mult):
                 if use_pre_mix_in:
@@ -332,11 +444,24 @@ def mhc_pre_big_fuse_with_norm_tilelang(
 
                 ol = T.alloc_fragment(hidden_block, T.float32)
                 T.clear(ol)
+                if write_aux:
+                    # Aux consumers (draft models) take the plain stream mean
+                    # of the same residual this collapse reads.
+                    aux = T.alloc_fragment(hidden_block, T.float32)
+                    T.clear(aux)
 
                 for i_hc in T.serial(hc_mult):
                     pre = pre_mix_shared[i_hc]
                     for i1_h in T.Parallel(hidden_block):
                         ol[i1_h] += pre * xl[i_hc, i1_h]
+                        if write_aux:
+                            aux[i1_h] += xl[i_hc, i1_h]
+
+                if write_aux:
+                    for i1_h in T.Parallel(hidden_block):
+                        aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16(
+                            aux[i1_h] / hc_mult
+                        )
 
                 if save_pre_mix:
                     # Keep the BF16 boundary before the delayed input RMSNorm.
@@ -438,7 +563,7 @@ def mhc_pre_big_fuse_broadcast_with_norm_tilelang(
         mixes_shared = T.alloc_shared(hc_mult3, T.float32)
         T.copy(mixes, mixes_shared)
 
-        if T.get_thread_binding() < 32:
+        if _is_mix_thread():
             cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
             for j in T.Parallel(hc_mult):
                 post_mix[i, j] = (
@@ -557,7 +682,7 @@ def mhc_fused_tilelang(
     tile_n: int = 1,
     split_k: int = 1,
 ) -> tilelang.JITKernel:
-    """Fused mhc post-mapping + pre-norm GEMM FMA"""
+    """Fused mhc post-mapping + pre-norm GEMM FMA."""
     m = T.dynamic("num_tokens")
     split_k = T.dynamic("split_k")
     h = hidden
@@ -987,3 +1112,762 @@ def hc_head_fuse_tilelang(
 
         if ENABLE_PDL:
             T.pdl_trigger()
+
+
+class HcPrenormGemmTileLangKernel(
+    VllmTileLangJitKernel["HcPrenormGemmTileLangKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        hidden_size: int
+        hc_mult: int
+        n_out: int
+        n_thr: int
+        tile_n: int
+        n_splits: int
+        use_block_m: bool
+        block_m: int
+
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        if compile_key.use_block_m:
+            return hc_prenorm_gemm_block_m_tilelang
+        return hc_prenorm_gemm_tilelang
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        num_tokens: int,
+        hc_hidden_size: int,
+        n_thr: int = 512,
+        tile_n: int = 12,
+        n_splits: int = 1,
+        **compile_key_fields: int,
+    ) -> CompileKey:
+        use_default_config = n_splits == 1 and tile_n == 12 and n_thr == 512
+        use_block_m = use_default_config and num_tokens >= 1024
+        use_small_tile = (
+            use_default_config
+            and not use_block_m
+            and num_tokens < 128
+            and hc_hidden_size % 1024 == 0
+        )
+        effective_n_thr = 1024 if use_small_tile else n_thr
+        effective_tile_n = 4 if use_small_tile else tile_n
+        return self.CompileKey(
+            **compile_key_fields,
+            n_thr=effective_n_thr,
+            tile_n=effective_tile_n,
+            n_splits=n_splits,
+            use_block_m=use_block_m,
+            block_m=2 if use_block_m else 1,
+        )
+
+    def get_warmup_keys(
+        self,
+        vllm_config: Any,
+        *,
+        hidden_size: int,
+        hc_mult: int,
+        n_out: int,
+    ) -> list[CompileKey]:
+        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        if max_tokens <= 0:
+            return []
+        return self._trace_dispatch(self.dispatch)(
+            num_tokens=WarmupIntRange(
+                1,
+                max_tokens + 1,
+                advance=lambda value: (
+                    128 if value < 128 else 1024 if value < 1024 else max_tokens + 1
+                ),
+            ),
+            hc_hidden_size=hidden_size * hc_mult,
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            n_out=n_out,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        num_tokens = 128
+        if compile_key.use_block_m:
+            num_tokens = 1024
+        elif compile_key.n_thr == 1024:
+            num_tokens = 1
+        hc_hidden_size = compile_key.hidden_size * compile_key.hc_mult
+        return dict(
+            x=make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hc_hidden_size),
+            fn=make_tilelang_warmup_tensor(
+                torch.float32, compile_key.n_out, hc_hidden_size
+            ),
+            out=make_tilelang_warmup_tensor(
+                torch.float32,
+                compile_key.n_splits,
+                num_tokens,
+                compile_key.n_out,
+            ),
+            sqrsum=make_tilelang_warmup_tensor(
+                torch.float32, compile_key.n_splits, num_tokens
+            ),
+            hidden_size=compile_key.hidden_size,
+            hc_mult=compile_key.hc_mult,
+            tile_n=compile_key.tile_n,
+            n_thr=compile_key.n_thr,
+            n_splits=compile_key.n_splits,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        x: Any,
+        fn: Any,
+        out: Any,
+        sqrsum: Any,
+        hidden_size: int,
+        hc_mult: int,
+        tile_n: int = 12,
+        n_thr: int = 512,
+        n_splits: int = 1,
+    ) -> TileLangLaunchSpec:
+        assert out.shape[0] == n_splits
+        assert sqrsum.shape[0] == n_splits
+        assert x.shape[1] == hc_mult * hidden_size
+        assert x.shape[1] % n_splits == 0
+        assert (x.shape[1] // n_splits) % n_thr == 0
+        compile_key = self.dispatch(
+            num_tokens=x.shape[0],
+            hc_hidden_size=x.shape[1],
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            n_out=fn.shape[0],
+            n_thr=n_thr,
+            tile_n=tile_n,
+            n_splits=n_splits,
+        )
+        kernel_arg = (
+            compile_key.block_m if compile_key.use_block_m else compile_key.n_splits
+        )
+        return (compile_key,), (
+            x,
+            fn,
+            out,
+            sqrsum,
+            hidden_size,
+            hc_mult,
+            fn.shape[0],
+            compile_key.n_thr,
+            compile_key.tile_n,
+            kernel_arg,
+        )
+
+
+class MhcPreBigFuseTileLangKernel(
+    VllmTileLangJitKernel["MhcPreBigFuseTileLangKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        hidden_size: int
+        hc_mult: int
+        n_splits: int
+        use_norm_weight: bool
+        is_broadcast: bool
+        rms_eps: float
+        hc_pre_eps: float
+        hc_sinkhorn_eps: float
+        hc_post_mult_value: float
+        sinkhorn_repeat: int
+        norm_eps: float
+
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        if compile_key.is_broadcast:
+            return mhc_pre_big_fuse_broadcast_with_norm_tilelang
+        if compile_key.use_norm_weight:
+            return mhc_pre_big_fuse_with_norm_tilelang
+        return mhc_pre_big_fuse_tilelang
+
+    @staticmethod
+    def _kernel_args(
+        compile_key: CompileKey,
+        *,
+        gemm_out_mul: Any,
+        gemm_out_sqrsum: Any,
+        hc_scale: Any,
+        hc_base: Any,
+        residual: Any,
+        residual_out: Any | None,
+        post_mix: Any,
+        comb_mix: Any,
+        layer_input: Any,
+        norm_weight: Any | None,
+    ) -> tuple[Any, ...]:
+        common_args = (
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            hc_scale,
+            hc_base,
+            residual,
+        )
+        output_args = (post_mix, comb_mix, layer_input)
+        eps_args = (
+            compile_key.rms_eps,
+            compile_key.hc_pre_eps,
+            compile_key.hc_sinkhorn_eps,
+            compile_key.hc_post_mult_value,
+            compile_key.sinkhorn_repeat,
+        )
+        middle_args: tuple[Any, ...]
+        if compile_key.is_broadcast:
+            assert residual_out is not None
+            assert norm_weight is not None
+            middle_args = (residual_out, *output_args, norm_weight)
+        elif compile_key.use_norm_weight:
+            assert norm_weight is not None
+            # layer_input stands in for the aux buffer on both epilogues: same
+            # shape and dtype, and write_aux is off here so it is never written.
+            middle_args = (
+                *output_args,
+                norm_weight,
+                post_mix,
+                post_mix,
+                layer_input,
+            )
+        else:
+            middle_args = (*output_args, post_mix, post_mix, layer_input)
+        norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else ()
+        return (
+            *common_args,
+            *middle_args,
+            compile_key.hidden_size,
+            *eps_args,
+            *norm_eps_args,
+            compile_key.n_splits,
+            compile_key.hc_mult,
+        )
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        hidden_size: int,
+        hc_mult: int,
+        n_splits: int,
+        is_broadcast: bool,
+        use_norm_weight: bool,
+        use_pre_gemm_splits: bool = False,
+        pre_gemm_k: int = 0,
+        sinkhorn_repeat: int,
+        norm_eps: float,
+        broadcast_norm_eps: float = 0.0,
+        num_tokens: int = 1,
+        use_fused_tilelang: bool = False,
+        **compile_key_fields: float,
+    ) -> CompileKey:
+        pre_gemm_n_splits = (
+            compute_num_split(64, pre_gemm_k, (num_tokens + 63) // 64)
+            if use_pre_gemm_splits
+            else n_splits
+        )
+        # The epilogue reduces over whatever the fused kernel split the GEMM
+        # into, so both have to read the same config.
+        fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
+        actual_n_splits = (
+            fused_config[1]
+            if use_fused_tilelang and fused_config is not None
+            else pre_gemm_n_splits
+        )
+        actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps
+        actual_use_norm_weight = use_norm_weight or is_broadcast
+        return self.CompileKey(
+            **compile_key_fields,
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            n_splits=actual_n_splits,
+            use_norm_weight=actual_use_norm_weight,
+            is_broadcast=is_broadcast,
+            sinkhorn_repeat=sinkhorn_repeat,
+            norm_eps=actual_norm_eps if actual_use_norm_weight else 0.0,
+        )
+
+    def get_warmup_keys(
+        self,
+        vllm_config: Any,
+        *,
+        hidden_size: int,
+        hc_mult: int,
+        use_norm_weight: bool,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        norm_eps: float | tuple[float, ...],
+        broadcast_norm_eps: float | tuple[float, ...] | None = None,
+        include_pre_gemm_splits: bool = False,
+        include_broadcast_splits: bool = False,
+    ) -> list[CompileKey]:
+        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        if max_tokens <= 0:
+            return []
+
+        warmup_cases = zip_inputs(
+            dict(
+                is_broadcast=False,
+                use_pre_gemm_splits=False,
+                pre_gemm_k=0,
+                use_fused_tilelang=False,
+            ),
+            dict(
+                is_broadcast=False,
+                use_pre_gemm_splits=include_pre_gemm_splits,
+                pre_gemm_k=hc_mult * hidden_size,
+                use_fused_tilelang=False,
+            ),
+            dict(
+                is_broadcast=include_broadcast_splits,
+                use_pre_gemm_splits=include_broadcast_splits,
+                pre_gemm_k=hidden_size,
+                use_fused_tilelang=False,
+            ),
+            dict(
+                is_broadcast=False,
+                use_pre_gemm_splits=False,
+                pre_gemm_k=0,
+                use_fused_tilelang=True,
+            ),
+        )
+        return self._trace_dispatch(self.dispatch)(
+            warmup_cases,
+            num_tokens=WarmupIntRange(
+                1,
+                max_tokens + 1,
+                advance=lambda value: 8 if value < 8 else cdiv(value, 64) * 64 + 1,
+            ),
+            n_splits=1,
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            use_norm_weight=use_norm_weight,
+            rms_eps=rms_eps,
+            hc_pre_eps=hc_pre_eps,
+            hc_sinkhorn_eps=hc_sinkhorn_eps,
+            hc_post_mult_value=hc_post_mult_value,
+            sinkhorn_repeat=sinkhorn_repeat,
+            norm_eps=norm_eps,
+            broadcast_norm_eps=broadcast_norm_eps
+            if broadcast_norm_eps is not None
+            else norm_eps,
+            _when=lambda *, use_fused_tilelang, num_tokens: (
+                not use_fused_tilelang or num_tokens <= 16
+            ),
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        hidden_size = compile_key.hidden_size
+        hc_mult = compile_key.hc_mult
+        hc_mult3 = hc_mult * (2 + hc_mult)
+        num_tokens = 1
+
+        gemm_out_mul = make_tilelang_warmup_tensor(
+            torch.float32, compile_key.n_splits, num_tokens, hc_mult3
+        )
+        gemm_out_sqrsum = make_tilelang_warmup_tensor(
+            torch.float32, compile_key.n_splits, num_tokens
+        )
+        hc_scale = make_tilelang_warmup_tensor(torch.float32, 3)
+        hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult3)
+        post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
+        comb_mix = make_tilelang_warmup_tensor(
+            torch.float32, num_tokens, hc_mult * hc_mult
+        )
+        layer_input = make_tilelang_warmup_tensor(
+            torch.bfloat16, num_tokens, hidden_size
+        )
+        residual_shape = (
+            (num_tokens, hidden_size)
+            if compile_key.is_broadcast
+            else (num_tokens, hc_mult, hidden_size)
+        )
+        residual = make_tilelang_warmup_tensor(torch.bfloat16, *residual_shape)
+        residual_out = (
+            make_tilelang_warmup_tensor(
+                torch.bfloat16, num_tokens, hc_mult, hidden_size
+            )
+            if compile_key.is_broadcast
+            else None
+        )
+        norm_weight = (
+            make_tilelang_warmup_tensor(torch.bfloat16, hidden_size)
+            if compile_key.is_broadcast or compile_key.use_norm_weight
+            else None
+        )
+        return dict(
+            gemm_out_mul=gemm_out_mul,
+            gemm_out_sqrsum=gemm_out_sqrsum,
+            hc_scale=hc_scale,
+            hc_base=hc_base,
+            residual=residual,
+            post_mix=post_mix,
+            comb_mix=comb_mix,
+            layer_input=layer_input,
+            rms_eps=compile_key.rms_eps,
+            hc_pre_eps=compile_key.hc_pre_eps,
+            hc_sinkhorn_eps=compile_key.hc_sinkhorn_eps,
+            hc_post_mult_value=compile_key.hc_post_mult_value,
+            sinkhorn_repeat=compile_key.sinkhorn_repeat,
+            residual_out=residual_out,
+            norm_weight=norm_weight,
+            norm_eps=compile_key.norm_eps,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        gemm_out_mul: Any,
+        gemm_out_sqrsum: Any,
+        hc_scale: Any,
+        hc_base: Any,
+        residual: Any,
+        post_mix: Any,
+        comb_mix: Any,
+        layer_input: Any,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        *,
+        residual_out: Any | None = None,
+        norm_weight: Any | None = None,
+        norm_eps: float = 0.0,
+    ) -> TileLangLaunchSpec:
+        is_broadcast = residual_out is not None
+        use_norm_weight = norm_weight is not None
+        n_splits = gemm_out_mul.shape[0]
+        if is_broadcast:
+            assert residual_out is not None
+            hidden_size = residual.shape[-1]
+            hc_mult = residual_out.shape[-2]
+        else:
+            hc_mult = residual.shape[-2]
+            hidden_size = residual.shape[-1]
+        compile_key = self.dispatch(
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            n_splits=n_splits,
+            is_broadcast=is_broadcast,
+            use_norm_weight=use_norm_weight,
+            rms_eps=rms_eps,
+            hc_pre_eps=hc_pre_eps,
+            hc_sinkhorn_eps=hc_sinkhorn_eps,
+            hc_post_mult_value=hc_post_mult_value,
+            sinkhorn_repeat=sinkhorn_repeat,
+            norm_eps=norm_eps,
+            broadcast_norm_eps=norm_eps,
+        )
+        return (
+            (compile_key,),
+            self._kernel_args(
+                compile_key,
+                gemm_out_mul=gemm_out_mul,
+                gemm_out_sqrsum=gemm_out_sqrsum,
+                hc_scale=hc_scale,
+                hc_base=hc_base,
+                residual=residual,
+                residual_out=residual_out,
+                post_mix=post_mix,
+                comb_mix=comb_mix,
+                layer_input=layer_input,
+                norm_weight=norm_weight,
+            ),
+        )
+
+
+class MhcPostTileLangKernel(VllmTileLangJitKernel["MhcPostTileLangKernel.CompileKey"]):
+    @dataclass(frozen=True)
+    class CompileKey:
+        hidden_size: int
+        hc_mult: int
+
+    @staticmethod
+    def kernel() -> Any:
+        return mhc_post_tilelang
+
+    def dispatch(  # type: ignore[override]
+        self,
+        **compile_key_fields: int,
+    ) -> CompileKey:
+        return self.CompileKey(**compile_key_fields)
+
+    def get_warmup_keys(
+        self,
+        *,
+        hidden_size: int,
+        hc_mult: int,
+    ) -> list[CompileKey]:
+        return self._trace_dispatch(self.dispatch)(
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        num_tokens = 1
+        hidden_size = compile_key.hidden_size
+        hc_mult = compile_key.hc_mult
+        comb_mix = make_tilelang_warmup_tensor(
+            torch.float32, num_tokens, hc_mult, hc_mult
+        )
+        residual = make_tilelang_warmup_tensor(
+            torch.bfloat16, num_tokens, hc_mult, hidden_size
+        )
+        post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
+        layer_input = make_tilelang_warmup_tensor(
+            torch.bfloat16, num_tokens, hidden_size
+        )
+        out = make_tilelang_warmup_tensor(
+            torch.bfloat16, num_tokens, hc_mult, hidden_size
+        )
+        return dict(
+            comb_mix=comb_mix,
+            residual=residual,
+            post_mix=post_mix,
+            layer_input=layer_input,
+            out=out,
+            hc_mult=hc_mult,
+            hidden_size=hidden_size,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        comb_mix: Any,
+        residual: Any,
+        post_mix: Any,
+        layer_input: Any,
+        out: Any,
+        hc_mult: int,
+        hidden_size: int,
+    ) -> TileLangLaunchSpec:
+        return (), (
+            comb_mix,
+            residual,
+            post_mix,
+            layer_input,
+            out,
+            hc_mult,
+            hidden_size,
+        )
+
+
+class MhcFusedTileLangKernel(
+    VllmTileLangJitKernel["MhcFusedTileLangKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        hidden_size: int
+        hc_mult: int
+        n_splits: int
+        tile_n: int
+        n_thr: int
+
+    @staticmethod
+    def kernel() -> Any:
+        return mhc_fused_tilelang
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        num_tokens: int,
+        hidden_size: int,
+        hc_mult: int,
+    ) -> CompileKey:
+        # Subscripts rather than an unpack: the warmup tracer parses this body
+        # and allows only plain assignments before the return.
+        config = require_fused_post_pre_config(num_tokens, hidden_size, hc_mult)
+        return self.CompileKey(
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            n_splits=config[1],
+            tile_n=config[0],
+            n_thr=config[2],
+        )
+
+    def get_warmup_keys(
+        self,
+        vllm_config: Any,
+        *,
+        hidden_size: int,
+        hc_mult: int,
+    ) -> list[CompileKey]:
+        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        return self._trace_dispatch(self.dispatch)(
+            num_tokens=WarmupIntRange(1, max_tokens + 1),
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            _when=lambda *, num_tokens: num_tokens <= _FUSED_POST_PRE_MAX_TOKENS,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        num_tokens = 1 if compile_key.tile_n == 2 else 16
+        hidden_size = compile_key.hidden_size
+        hc_mult = compile_key.hc_mult
+        hc_mult3 = hc_mult * (2 + hc_mult)
+        comb_mix = make_tilelang_warmup_tensor(
+            torch.float32, num_tokens, hc_mult, hc_mult
+        )
+        residual_in = make_tilelang_warmup_tensor(
+            torch.bfloat16, num_tokens, hc_mult, hidden_size
+        )
+        post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
+        x_in = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size)
+        weight_t = make_tilelang_warmup_tensor(
+            torch.float32, hc_mult3, hc_mult, hidden_size
+        )
+        return dict(
+            comb_mix=comb_mix,
+            residual_in=residual_in,
+            post_mix=post_mix,
+            x_in=x_in,
+            weight_t=weight_t,
+            hc_mult=hc_mult,
+            hidden_size=hidden_size,
+            hc_mult3=hc_mult3,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        comb_mix: Any,
+        residual_in: Any,
+        post_mix: Any,
+        x_in: Any,
+        weight_t: Any,
+        hc_mult: int,
+        hidden_size: int,
+        hc_mult3: int,
+    ) -> TileLangLaunchSpec:
+        num_tokens = residual_in.shape[0]
+        tile_n, n_splits, n_thr = require_fused_post_pre_config(
+            num_tokens, hidden_size, hc_mult
+        )
+        yp_out = residual_in.new_empty(
+            (n_splits, num_tokens, hc_mult3), dtype=torch.float32
+        )
+        rp_out = residual_in.new_empty((n_splits, num_tokens), dtype=torch.float32)
+        residual_out = residual_in.new_empty(residual_in.shape)
+        return (
+            (),
+            (
+                comb_mix,
+                residual_in,
+                post_mix,
+                x_in,
+                weight_t,
+                yp_out,
+                rp_out,
+                residual_out,
+                hc_mult,
+                hidden_size,
+                hc_mult3,
+            ),
+            dict(n_thr=n_thr, tile_n=tile_n),
+            (yp_out, rp_out, residual_out),
+        )
+
+
+class HcHeadFusedTileLangKernel(
+    VllmTileLangJitKernel["HcHeadFusedTileLangKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        hidden_size: int
+        hc_mult: int
+        rms_eps: float
+        hc_eps: float
+
+    @staticmethod
+    def kernel() -> Any:
+        return hc_head_fuse_tilelang
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        hidden_size: int,
+        hc_mult: int,
+        rms_eps: float,
+        hc_eps: float,
+    ) -> CompileKey:
+        return self.CompileKey(
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            rms_eps=rms_eps,
+            hc_eps=hc_eps,
+        )
+
+    def get_warmup_keys(
+        self,
+        *,
+        hidden_size: int,
+        hc_mult: int,
+        rms_eps: float,
+        hc_eps: float,
+    ) -> list[CompileKey]:
+        return self._trace_dispatch(self.dispatch)(
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            rms_eps=rms_eps,
+            hc_eps=hc_eps,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        num_tokens = 1
+        hidden_size = compile_key.hidden_size
+        hc_mult = compile_key.hc_mult
+        residual = make_tilelang_warmup_tensor(
+            torch.bfloat16, num_tokens, hc_mult, hidden_size
+        )
+        fn = make_tilelang_warmup_tensor(torch.float32, hc_mult, hc_mult * hidden_size)
+        hc_scale = make_tilelang_warmup_tensor(torch.float32, 1)
+        hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult)
+        out = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size)
+        return dict(
+            residual=residual,
+            fn=fn,
+            hc_scale=hc_scale,
+            hc_base=hc_base,
+            out=out,
+            hidden_size=hidden_size,
+            rms_eps=compile_key.rms_eps,
+            hc_eps=compile_key.hc_eps,
+            hc_mult=hc_mult,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        residual: Any,
+        fn: Any,
+        hc_scale: Any,
+        hc_base: Any,
+        out: Any,
+        hidden_size: int,
+        rms_eps: float,
+        hc_eps: float,
+        hc_mult: int,
+    ) -> TileLangLaunchSpec:
+        return (), (
+            residual,
+            fn,
+            hc_scale,
+            hc_base,
+            out,
+            hidden_size,
+            rms_eps,
+            hc_eps,
+            hc_mult,
+        )
+
+
+_HC_PRENORM_GEMM_TILELANG_KERNEL = HcPrenormGemmTileLangKernel()
+_MHC_PRE_BIG_FUSE_TILELANG_KERNEL = MhcPreBigFuseTileLangKernel()
+_MHC_POST_TILELANG_KERNEL = MhcPostTileLangKernel()
+_MHC_FUSED_TILELANG_KERNEL = MhcFusedTileLangKernel()
+_HC_HEAD_FUSED_TILELANG_KERNEL = HcHeadFusedTileLangKernel()
diff --git a/vllm/model_executor/kernels/mhc/torch.py b/vllm/model_executor/kernels/mhc/torch.py
index 9c57cc169cc6..eb4de35e7c75 100644
--- a/vllm/model_executor/kernels/mhc/torch.py
+++ b/vllm/model_executor/kernels/mhc/torch.py
@@ -57,8 +57,7 @@ def mhc_pre_torch(
     sinkhorn_repeat: int,
     n_splits: int = 1,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Forward pass for mHC pre block.
+    """Forward pass for mHC pre block.
 
     Args:
         residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
@@ -76,8 +75,8 @@ def mhc_pre_torch(
         post_mix: shape (..., hc_mult), dtype torch.float32
         comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
         layer_input: shape (..., hidden_size), dtype torch.bfloat16
-    """
 
+    """
     # Validate shapes
     assert residual.dtype == torch.bfloat16
     assert fn.dtype == torch.float32
diff --git a/vllm/model_executor/kernels/mhc/triton.py b/vllm/model_executor/kernels/mhc/triton.py
index d94d71a6ea1c..8724113779c8 100644
--- a/vllm/model_executor/kernels/mhc/triton.py
+++ b/vllm/model_executor/kernels/mhc/triton.py
@@ -141,6 +141,129 @@ def _hc_collapse_triton_fake(x: Tensor, pre_mix: Tensor) -> Tensor:
     return torch.empty(x.shape[0], x.shape[2], dtype=x.dtype, device=x.device)
 
 
+@triton.jit
+def _mhc_pre_mix_kernel(
+    gemm_ptr,
+    sqrsum_ptr,
+    hc_scale_ptr,
+    hc_base_ptr,
+    out_ptr,
+    splitk,
+    hc_mult: tl.constexpr,
+    gemm_stride_k,
+    gemm_stride_t,
+    gemm_stride_j,
+    sqrsum_stride_k,
+    sqrsum_stride_t,
+    out_stride_t,
+    out_stride_j,
+    inv_hc_hidden,
+    rms_eps,
+    hc_pre_eps,
+    SPLITK_BLOCK: tl.constexpr,
+    HC_BLOCK: tl.constexpr,
+):
+    """Recover the pre-mix gate from a split-k mHC pre GEMM output."""
+    token_idx = tl.program_id(0).to(tl.int64)
+    ks = tl.arange(0, SPLITK_BLOCK)
+    js = tl.arange(0, HC_BLOCK)
+    kmask = ks < splitk
+    jmask = js < hc_mult
+
+    # Only the first hc_mult GEMM columns feed the pre gate.
+    gemm = tl.load(
+        gemm_ptr
+        + ks[:, None] * gemm_stride_k
+        + token_idx * gemm_stride_t
+        + js[None, :] * gemm_stride_j,
+        mask=kmask[:, None] & jmask[None, :],
+        other=0.0,
+    ).to(tl.float32)
+    mixes = tl.sum(gemm, 0)
+
+    sqrsum = tl.load(
+        sqrsum_ptr + ks * sqrsum_stride_k + token_idx * sqrsum_stride_t,
+        mask=kmask,
+        other=0.0,
+    ).to(tl.float32)
+    rstd = tl.rsqrt(tl.sum(sqrsum, 0) * inv_hc_hidden + rms_eps)
+
+    scale = tl.load(hc_scale_ptr).to(tl.float32)
+    base = tl.load(hc_base_ptr + js, mask=jmask, other=0.0).to(tl.float32)
+
+    pre = tl.sigmoid(mixes * rstd * scale + base) + hc_pre_eps
+    tl.store(out_ptr + token_idx * out_stride_t + js * out_stride_j, pre, mask=jmask)
+
+
+def mhc_pre_mix_triton(
+    gemm_out: Tensor,
+    sqrsum: Tensor,
+    hc_scale: Tensor,
+    hc_base: Tensor,
+    hc_mult: int,
+    hc_hidden_size: int,
+    rms_eps: float,
+    hc_pre_eps: float,
+) -> Tensor:
+    """Pre-mix gate for the delayed mHC pre, from AITER's split-k GEMM output.
+
+    AITER's ``mhc_pre_big_fuse`` consumes the unreduced ``[splitk, tokens,
+    hc_mult3]`` GEMM output and the matching row square-sums, but only returns
+    the post and comb gates. The delayed formulation also needs the pre gate,
+    to carry into the next sublayer seam. It is the same slice of the same
+    numbers, so recover it here rather than repeating the projection.
+    """
+    assert gemm_out.ndim == 3 and gemm_out.dtype == torch.float32
+    assert sqrsum.ndim == 2 and sqrsum.dtype == torch.float32
+    splitk, num_tokens = gemm_out.shape[0], gemm_out.shape[1]
+    assert sqrsum.shape == (splitk, num_tokens)
+
+    out = torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=gemm_out.device)
+    if num_tokens == 0:
+        return out
+
+    _mhc_pre_mix_kernel[(num_tokens,)](
+        gemm_out,
+        sqrsum,
+        hc_scale,
+        hc_base,
+        out,
+        splitk,
+        hc_mult,
+        gemm_out.stride(0),
+        gemm_out.stride(1),
+        gemm_out.stride(2),
+        sqrsum.stride(0),
+        sqrsum.stride(1),
+        out.stride(0),
+        out.stride(1),
+        1.0 / hc_hidden_size,
+        rms_eps,
+        hc_pre_eps,
+        SPLITK_BLOCK=triton.next_power_of_2(splitk),
+        HC_BLOCK=triton.next_power_of_2(hc_mult),
+        num_warps=1,
+        # Match the separate FP32 multiply and add in the Torch reference.
+        enable_fp_fusion=False,
+    )
+    return out
+
+
+def _mhc_pre_mix_triton_fake(
+    gemm_out: Tensor,
+    sqrsum: Tensor,
+    hc_scale: Tensor,
+    hc_base: Tensor,
+    hc_mult: int,
+    hc_hidden_size: int,
+    rms_eps: float,
+    hc_pre_eps: float,
+) -> Tensor:
+    return torch.empty(
+        gemm_out.shape[1], hc_mult, dtype=torch.float32, device=gemm_out.device
+    )
+
+
 def hc_head_reduce_triton_kernel(
     x: torch.Tensor,
     hc_fn: torch.Tensor,
@@ -216,3 +339,11 @@ def _hc_head_triton(
     mutates_args=[],
     fake_impl=_hc_collapse_triton_fake,
 )
+
+
+direct_register_custom_op(
+    op_name="mhc_pre_mix_triton",
+    op_func=mhc_pre_mix_triton,
+    mutates_args=[],
+    fake_impl=_mhc_pre_mix_triton_fake,
+)
diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py
index 6a42ca622270..deb1cf429a0c 100644
--- a/vllm/model_executor/kernels/mhc/warmup.py
+++ b/vllm/model_executor/kernels/mhc/warmup.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Iterable
 from dataclasses import asdict, dataclass
 from typing import Any
 
@@ -35,23 +36,34 @@ class CompileKey:
         use_pre_mix_in: bool
         rms_numel: int
         save_pre_mix: bool = True
+        split_mode: str = "fused"
+        write_aux: bool = False
 
     kernel: Any = staticmethod(mhc_pre_big_fuse_with_norm_tilelang)
 
     def dispatch(self, *, n_splits, **fields) -> CompileKey:  # type: ignore[override]
         return self.CompileKey(n_splits=n_splits, **fields)
 
-    def get_warmup_keys(self, *, max_tokens: int, **fields) -> list[CompileKey]:
-        # The split heuristic changes only at 64-token boundaries.
-        splits = (
-            sorted(
+    def get_warmup_keys(
+        self,
+        *,
+        max_tokens: int,
+        extra_splits: Iterable[int] = (),
+        **fields,
+    ) -> list[CompileKey]:
+        # The split heuristic changes only at 64-token boundaries. Callers that
+        # reach this kernel through another GEMM (the fused post + pre-norm
+        # GEMM) pass the splits that GEMM picks in extra_splits.
+        splits = sorted(
+            set(extra_splits)
+            | (
                 {
                     compute_mhc_pre_num_splits(fields["rms_numel"], num_tokens)
                     for num_tokens in range(1, max_tokens + 1, 64)
                 }
+                if is_deep_gemm_supported()
+                else {1}
             )
-            if is_deep_gemm_supported()
-            else [1]
         )
         return self._trace_dispatch(self.dispatch)(n_splits=splits, **fields)
 
@@ -61,6 +73,21 @@ def compile(self, compile_key: CompileKey) -> None:
                 **asdict(compile_key)
             )
 
+    def compile_many(self, compile_keys: Iterable[CompileKey]) -> None:
+        missing = list(
+            dict.fromkeys(
+                key for key in compile_keys if key not in self._compiled_cache
+            )
+        )
+        if len(missing) < 2:
+            return super().compile_many(missing)
+
+        # TileLang elaborates serially before compiling independent TIR modules.
+        compiled = self.kernel.par_compile(
+            [asdict(key) for key in missing], num_workers=min(4, len(missing))
+        )
+        self._compiled_cache.update(zip(missing, compiled, strict=True))
+
     def __call__(self, *tensors, **fields):
         compile_key = self.dispatch(
             n_splits=tensors[0].shape[0],
diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py
index 159c521cee2e..808d6e711277 100644
--- a/vllm/model_executor/layers/activation.py
+++ b/vllm/model_executor/layers/activation.py
@@ -8,6 +8,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from vllm import ir
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
 from vllm.platforms import CpuArchEnum, current_platform
@@ -128,8 +129,11 @@ def __init__(self, *, compile_native: bool = True):
         super().__init__(compile_native=compile_native)
         if (
             current_platform.is_cuda_alike()
-            or current_platform.is_cpu()
             or current_platform.is_xpu()
+            or (
+                current_platform.is_cpu()
+                and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
+            )
         ):
             self.op = torch.ops._C.silu_and_mul
 
@@ -229,12 +233,8 @@ def __init__(
         self.swiglu_limit = float(swiglu_limit)
         self.alpha = float(alpha)
         self.beta = float(beta)
-        if current_platform.is_rocm() or current_platform.is_xpu():
-            self._forward_method = self.forward_native
-        elif current_platform.is_cuda_alike():
+        if current_platform.is_cuda_alike():
             self.op = torch.ops._C.silu_and_mul_with_clamp
-        elif current_platform.is_cpu():
-            self._forward_method = self.forward_native
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         d = x.shape[-1] // 2
@@ -249,7 +249,11 @@ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
         self.op(out, x, self.swiglu_limit, self.alpha, self.beta)
         return out
 
-    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+    def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
+        # Limit the ROCm _C path to the default case to avoid the
+        # precision loss MiniMax saw without fp32 intermediates.
+        if self.alpha == 1.0 and self.beta == 0.0:
+            return self.forward_cuda(x)
         return self.forward_native(x)
 
     def extra_repr(self) -> str:
@@ -332,24 +336,11 @@ def __init__(self, activation_sparsity: float, approximate: str = "none"):
             raise ValueError("activation_sparsity is 0.0. Please use GeluAndMul.")
         target_sparsity_tensor = torch.tensor(activation_sparsity, dtype=torch.float32)
         normal_dist = torch.distributions.normal.Normal(0, 1)
-        self.std_multiplier = normal_dist.icdf(target_sparsity_tensor)
-
-    def _gaussian_topk(self, x: torch.Tensor) -> torch.Tensor:
-        """Get % sparse percentile of the Gaussian distribution."""
-        # NOTE(rob): for TP>1, we could all-gather to get the means/std.
-        # But we do not do this because in expectation they are the same
-        # and in practice the eval scores are good without gathering.
-        mean = torch.mean(x, dim=-1, keepdim=True)
-        std = torch.std(x, dim=-1, keepdim=True, unbiased=False)
-        cutoff_x = mean + std * self.std_multiplier
-        return nn.functional.relu(x - cutoff_x)
+        self.std_multiplier = normal_dist.icdf(target_sparsity_tensor).item()
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
-        d = x.shape[-1] // 2
-        out = self._gaussian_topk(x[..., :d])
-        out = F.gelu(out, approximate=self.approximate)
-        return out * x[..., d:]
+        return ir.ops.gelu_and_mul_sparse(x, self.std_multiplier, self.approximate)
 
     def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
         return self.forward_native(x)
@@ -436,8 +427,11 @@ def __init__(self, approximate: str = "none"):
             raise ValueError(f"Unknown approximate mode: {approximate}")
         if (
             current_platform.is_cuda_alike()
-            or current_platform.is_cpu()
             or current_platform.is_xpu()
+            or (
+                current_platform.is_cpu()
+                and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
+            )
         ):
             if approximate == "none":
                 self.op = torch.ops._C.gelu_and_mul
@@ -491,7 +485,6 @@ def __init__(self, alpha: float = 1.702, limit: float = 7.0):
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
-
         gate, up = x[..., ::2], x[..., 1::2]
         gate = gate.clamp(min=None, max=self.limit)
         up = up.clamp(min=-self.limit, max=self.limit)
@@ -659,9 +652,7 @@ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
 # --8<-- [start:relu2]
 @CustomOp.register("relu2")
 class ReLUSquaredActivation(CustomOp):
-    """
-    Applies the relu^2 activation introduced in https://arxiv.org/abs/2109.08668v2
-    """
+    """Applies the relu^2 activation introduced in https://arxiv.org/abs/2109.08668v2."""
 
     # --8<-- [end:relu2]
 
@@ -688,8 +679,7 @@ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
 # --8<-- [start:xielu]
 @CustomOp.register("xielu")
 class XIELU(CustomOp):
-    """
-    Applies the xIELU activation function introduced in https://arxiv.org/abs/2411.13010
+    """Applies the xIELU activation function introduced in https://arxiv.org/abs/2411.13010
     If the user has installed the nickjbrowning/XIELU, we import xIELU CUDA
     Otherwise, we emit a single warning and use xIELU Python
     """
@@ -759,7 +749,7 @@ def _xielu_python(self, x: torch.Tensor) -> torch.Tensor:
         )
 
     def _xielu_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        """Firewall function to prevent torch.compile from seeing .item()"""
+        """Firewall function to prevent torch.compile from seeing .item()."""
         assert self._xielu_cuda_obj is not None, "XIELU CUDA object must not be None"
         original_shape = x.shape
         # CUDA kernel expects 3D tensors, reshape if needed
diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py
index a001c8c29c80..bec1856d93ad 100644
--- a/vllm/model_executor/layers/attention/attention.py
+++ b/vllm/model_executor/layers/attention/attention.py
@@ -8,7 +8,10 @@
 
 import vllm.envs as envs
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
-from vllm.config import CacheConfig, get_current_vllm_config
+from vllm.config import (
+    CacheConfig,
+    get_current_vllm_config,
+)
 from vllm.config.vllm import VllmConfig
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
@@ -100,6 +103,7 @@ def _largest_kernel_block_within(
     per_token_bytes: int,
     page_budget: int,
     fallback: int,
+    kv_cache_spec: KVCacheSpec | None = None,
 ) -> int:
     """Largest supported kernel block size whose page fits in ``page_budget``.
 
@@ -111,7 +115,7 @@ def _largest_kernel_block_within(
     """
     from vllm.v1.attention.backend import MultipleOf
 
-    sizes = attn_backend.get_supported_kernel_block_sizes()
+    sizes = attn_backend.get_supported_kernel_block_sizes(kv_cache_spec)
     max_block_size = page_budget // per_token_bytes
     candidates = [s for s in sizes if isinstance(s, int)]
     candidates.extend(
@@ -166,8 +170,8 @@ def _init_kv_cache_quant(
         layer: The attention layer instance to initialize.
         quant_config: Optional quantization configuration.
         prefix: Layer name prefix for quantization method lookup.
-    """
 
+    """
     # Note [Register q/k/v/prob scales in state dict]
     # When calling model.to(device), only parameters/buffers in state dict are
     # moved. If not registering q/k/v/prob scales in state dict, there would
@@ -254,8 +258,7 @@ def __init__(
         head_size_v: int | None = None,
         **extra_impl_args,
     ) -> None:
-        """
-        The KV cache is stored inside this class and is accessed via
+        """The KV cache is stored inside this class and is accessed via
         `self.kv_cache`.
         """
         super().__init__()
@@ -493,8 +496,7 @@ def forward(
         output_shape: torch.Size | None = None,
         output_dtype: torch.dtype | None = None,
     ) -> torch.Tensor:
-        """
-        The KV cache is stored inside this class and is accessed via
+        """The KV cache is stored inside this class and is accessed via
         `self.kv_cache`.
 
         Attention metadata (`attn_metadata`) is set using a context manager in
@@ -634,7 +636,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
             # ``unify`` scales it up by an integer ratio.
             shared_page = vllm_config.cache_config.skip_page_size_padded
             # The backend owns its packing
-            sw_per_token = self.attn_backend.customize_spec(
+            kv_cache_spec = self.attn_backend.customize_spec(
                 SlidingWindowSpec(
                     block_size=1,
                     num_kv_heads=self.num_kv_heads,
@@ -644,10 +646,15 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
                     kv_quant_mode=quant_mode,
                     sliding_window=self.sliding_window,
                 )
-            ).real_page_size_bytes
+            )
+            sw_per_token = kv_cache_spec.real_page_size_bytes
             page_budget = shared_page or sw_per_token * block_size
             sw_block_size = _largest_kernel_block_within(
-                self.attn_backend, sw_per_token, page_budget, block_size
+                self.attn_backend,
+                sw_per_token,
+                page_budget,
+                block_size,
+                kv_cache_spec,
             )
             return SlidingWindowSpec(
                 block_size=sw_block_size,
@@ -691,6 +698,7 @@ def get_attention_context(
 
         Note: attn_metadata may be None, but attn_layer and kv_cache are always
         extracted from the forward context.
+
     """
     forward_context: ForwardContext = get_forward_context()
     attn_metadata_raw = forward_context.attn_metadata
@@ -718,8 +726,7 @@ def unified_kv_cache_update(
     value: torch.Tensor,
     layer_name: LayerNameType,
 ) -> torch.Tensor:
-    """
-    Returns a dummy that is passed to unified_attention to signal a side effect and
+    """Returns a dummy that is passed to unified_attention to signal a side effect and
     the data dependency between them to ensure torch.compile preserves ordering.
     """
     layer_name = _resolve_layer_name(layer_name)
diff --git a/vllm/model_executor/layers/attention/cross_attention.py b/vllm/model_executor/layers/attention/cross_attention.py
index 091f0a1856d4..1422e2778564 100644
--- a/vllm/model_executor/layers/attention/cross_attention.py
+++ b/vllm/model_executor/layers/attention/cross_attention.py
@@ -34,7 +34,6 @@ def _get_cross_slot_mapping(
     device: torch.device,
 ) -> torch.Tensor:
     """Get cross-attention slot mappings."""
-
     block_size = kv_cache_spec.block_size
     slot_mappings = []
 
@@ -115,7 +114,7 @@ def build(
             # needed here to know how many tokens to attend to from the cached
             # cross-attention KV cache.
             new_metadata.seq_lens = common_attn_metadata.encoder_seq_lens
-            new_metadata._seq_lens_cpu = torch.from_numpy(
+            new_metadata.seq_lens_cpu_upper_bound = torch.from_numpy(
                 common_attn_metadata.encoder_seq_lens_cpu
             )
 
@@ -187,8 +186,7 @@ def forward(
 
 
 class CrossAttention(Attention):
-    """
-    Cross-attention for encoder-decoder models.
+    """Cross-attention for encoder-decoder models.
     Handles attention between decoder queries and encoder keys/values.
     """
 
diff --git a/vllm/model_executor/layers/attention/encoder_only_attention.py b/vllm/model_executor/layers/attention/encoder_only_attention.py
index 5805fe2ae1ca..ef73b2798da4 100644
--- a/vllm/model_executor/layers/attention/encoder_only_attention.py
+++ b/vllm/model_executor/layers/attention/encoder_only_attention.py
@@ -49,9 +49,7 @@ def build(
 
 
 class EncoderOnlyAttention(Attention):
-    """
-    Encoder attention is a special case that doesn't need a KV Cache.
-    """
+    """Encoder attention is a special case that doesn't need a KV Cache."""
 
     def __init__(
         self,
diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py
index 34df61fc82b8..4ebf3ef40d5c 100644
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/vllm/model_executor/layers/attention/mla_attention.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-# MLA Common Components
+"""# MLA Common Components
 
 This file implements common components for MLA implementations.
 
@@ -207,7 +206,7 @@
 from dataclasses import dataclass, replace
 from enum import Enum
 from math import lcm
-from typing import ClassVar, Generic, TypeVar, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
 
 import numpy as np
 import torch
@@ -269,11 +268,11 @@
     LayerNameType,
     _encode_layer_name,
     _resolve_layer_name,
+    async_tensor_h2d,
     direct_register_custom_op,
     get_dtype_size,
     is_quantized_kv_cache,
     kv_cache_dtype_str_to_dtype,
-    np_to_pinned_tensor,
 )
 from vllm.v1.attention.backend import (
     AttentionBackend,
@@ -284,6 +283,7 @@
     CommonAttentionMetadata,
     MLAAttentionImpl,
 )
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.prefill import (
     MLAPrefillBackend,
     get_mla_prefill_backend,
@@ -309,6 +309,11 @@
     get_kv_quant_mode,
 )
 
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.mla.index_group import (
+        SparseMLAIndexGroupBuilder,
+    )
+
 logger = init_logger(__name__)
 
 _FP8_DTYPE = current_platform.fp8_dtype()
@@ -399,6 +404,24 @@ def _get_kv_b_proj_input_dtype(
     return weight_dtype
 
 
+def split_kv_b_proj(
+    kv_b_proj: nn.Module,
+    out_dtype: torch.dtype,
+    kv_lora_rank: int,
+    num_heads: int,
+    qk_nope_head_dim: int,
+    v_head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Dequantize ``kv_b_proj`` and return ``W_UK [L,N,P]``, ``W_UV [L,N,V]``."""
+    weight = get_and_maybe_dequant_weights(kv_b_proj, out_dtype=out_dtype).T
+    assert weight.shape == (
+        kv_lora_rank,
+        num_heads * (qk_nope_head_dim + v_head_dim),
+    ), f"kv_b_proj weight {tuple(weight.shape)} vs {kv_lora_rank=} {num_heads=}"
+    weight = weight.view(kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim)
+    return weight.split([qk_nope_head_dim, v_head_dim], dim=-1)
+
+
 class MLAAttention(nn.Module, AttentionLayerBase):
     """Multi-Head Latent Attention layer.
 
@@ -414,6 +437,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
     """
 
     supports_dense_mha_prefill: ClassVar[bool] = True
+    # Under PCP+DCP only the decode rows carry an LSE; the base forward
+    # merges a full-batch LSE, so subclasses opt in with their own forward.
+    supports_pcp_dcp: ClassVar[bool] = False
 
     def __init__(
         self,
@@ -433,6 +459,7 @@ def __init__(
         use_sparse: bool = False,
         indexer: object | None = None,
         topk_indices_buffer: torch.Tensor | None = None,
+        index_group_builder: "SparseMLAIndexGroupBuilder | None" = None,
         non_causal_multi_token_decode: bool = False,
         sliding_window: int | None = None,
         prefill_backend_cls: type[MLAPrefillBackend] | None = None,
@@ -538,9 +565,11 @@ def __init__(
         # explicitly so backbone "skip" layers (indexer=None) still find it.
         if use_sparse:
             extra_impl_args["topk_indices_buffer"] = topk_indices_buffer
+            if index_group_builder is not None:
+                extra_impl_args["index_group_builder"] = index_group_builder
 
         impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
-        self.impl = impl_cls(  # type: ignore[assignment]  # impl_cls always returns an MLAAttentionImpl subclass
+        impl = impl_cls(
             num_heads=self.num_heads,
             head_size=self.head_size,
             scale=self.scale,
@@ -562,6 +591,13 @@ def __init__(
             indexer=indexer,
             **extra_impl_args,
         )
+        self.impl = impl  # type: ignore[assignment]
+        index_group = getattr(impl, "index_group", None)
+        self.hisparse_cache = (
+            index_group.cache(cast(Any, impl).index_group_index)
+            if isinstance(index_group, HiSparseMLAIndexGroup)
+            else None
+        )
         self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
         self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False)
         # AMX reads kv_b_proj's weight directly and never calls it live; the
@@ -645,6 +681,10 @@ def __init__(
 
                     _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
 
+        if self.use_pcp and self.impl.dcp_world_size > 1 and not self.supports_pcp_dcp:
+            raise NotImplementedError(
+                f"{type(self).__name__} does not support PCP+DCP."
+            )
         self.dcp_manager: MLADCPManager | None = None
         if self.impl.dcp_world_size > 1:
             query_dtype = (
@@ -724,6 +764,61 @@ def chunked_prefill_workspace_size(self) -> int:
             )
         return self._chunked_prefill_workspace_size
 
+    def update_kv_cache(
+        self,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor | None,
+        attn_metadata: "MLACommonMetadata | None",
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor,
+    ) -> None:
+        cache = self.hisparse_cache
+        if slot_mapping is None or (cache is not None and cache.dummy_batch):
+            return
+        kv_c_normed, k_pe, slot_mapping = maybe_gather_mla_latent_cache_inputs(
+            kv_c_normed,
+            k_pe,
+            slot_mapping,
+            attn_metadata.num_decode_tokens if attn_metadata is not None else None,
+            self.use_pcp,
+        )
+        assert slot_mapping is not None
+        if cache is not None:
+            kv_cache, slot_mapping, num_rows = cache.write_target(
+                kv_c_normed.shape[0], slot_mapping.numel()
+            )
+            kv_c_normed = kv_c_normed[:num_rows]
+            k_pe = k_pe[:num_rows]
+        self.impl.do_kv_cache_update(  # type: ignore[attr-defined]
+            kv_c_normed,
+            k_pe,
+            kv_cache,
+            slot_mapping,
+            kv_cache_dtype,
+            k_scale,
+        )
+        if cache is not None:
+            mirror_target = cache.mirror_write_target(kv_c_normed.shape[0])
+            if mirror_target is not None:
+                mirror_cache, mirror_slots = mirror_target
+                self.impl.do_kv_cache_update(  # type: ignore[attr-defined]
+                    kv_c_normed,
+                    k_pe,
+                    mirror_cache,
+                    mirror_slots,
+                    kv_cache_dtype,
+                    k_scale,
+                )
+
+    def prepare_kv_cache_update(
+        self, attn_metadata: "MLACommonMetadata | None"
+    ) -> None:
+        cache = self.hisparse_cache
+        if cache is not None and cache.runtime.is_group_leader:
+            cache.prepare_group_for_batch(attn_metadata)
+
     def forward(
         self,
         q: torch.Tensor,
@@ -751,25 +846,18 @@ def forward(
                 f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
             )
             layer_slot_mapping = slot_mapping.get(self.layer_name)
-            kv_for_cache, kpe_for_cache, layer_slot_mapping = (
-                maybe_gather_mla_latent_cache_inputs(
-                    kv_c_normed,
-                    k_pe,
-                    layer_slot_mapping,
-                    attn_metadata.num_decode_tokens
-                    if attn_metadata is not None
-                    else None,
-                    self.use_pcp,
-                )
-            )
-            self.impl.do_kv_cache_update(  # type: ignore[attr-defined]
-                kv_for_cache,
-                kpe_for_cache,
+            self.prepare_kv_cache_update(attn_metadata)
+            self.update_kv_cache(
+                kv_c_normed,
+                k_pe,
                 self_kv_cache,
                 layer_slot_mapping,
+                attn_metadata,
                 self.kv_cache_dtype,
                 self._k_scale,
             )
+            if self.hisparse_cache is not None:
+                self.hisparse_cache.finish_kv_update()
             output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
             self.forward_impl(
                 q,
@@ -955,9 +1043,7 @@ def forward_impl(
                 mqa_q_pe = mqa_pe_padded
 
             if self.is_aiter_triton_fp4_bmm_enabled:
-                from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
-
-                mqa_ql_nope = batched_gemm_a16wfp4(
+                mqa_ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4(
                     mqa_q_nope,
                     self.W_K,
                     self.W_K_scale,
@@ -985,7 +1071,7 @@ def forward_impl(
                     mqa_q_nope,
                     self.impl._w_uk_packed,  # type: ignore[attr-defined]
                     True,
-                    None,
+                    self.impl._w_scale,  # type: ignore[attr-defined]
                 )
                 mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
             else:
@@ -1114,6 +1200,8 @@ def forward_impl(
         return output_padded
 
     def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool:
+        if self.hisparse_cache is not None:
+            return False
         prefill = attn_metadata.prefill
         if prefill is None:
             return False
@@ -1151,13 +1239,6 @@ def process_weights_after_loading(self, act_dtype: torch.dtype):
             )
             return
 
-        # we currently do not have quantized bmm's which are needed for
-        # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
-        # the bmm's in 16-bit, the extra memory overhead of this is fairly low
-        kv_b_proj_weight = get_and_maybe_dequant_weights(
-            self.kv_b_proj, out_dtype=act_dtype
-        ).T
-
         if self.dcp_q_replicate:
             # qrep wired here: validate unsupported decode backends once.
             assert self.q_pad_num_heads in (None, self.num_heads), (
@@ -1173,24 +1254,13 @@ def process_weights_after_loading(self, act_dtype: torch.dtype):
                     "FP4/FP8 MLA BMM paths."
                 )
 
-        assert kv_b_proj_weight.shape == (
-            self.kv_lora_rank,
-            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
-        ), (
-            f"{kv_b_proj_weight.shape=}, "
-            f"{self.kv_lora_rank=}, "
-            f"{self.num_heads=}, "
-            f"{self.qk_nope_head_dim=}, "
-            f"{self.v_head_dim=}"
-        )
-        kv_b_proj_weight = kv_b_proj_weight.view(
+        W_UK, W_UV = split_kv_b_proj(
+            self.kv_b_proj,
+            act_dtype,
             self.kv_lora_rank,
             self.num_heads,
-            self.qk_nope_head_dim + self.v_head_dim,
-        )
-
-        W_UK, W_UV = kv_b_proj_weight.split(
-            [self.qk_nope_head_dim, self.v_head_dim], dim=-1
+            self.qk_nope_head_dim,
+            self.v_head_dim,
         )
 
         # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported
@@ -1297,6 +1367,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
             )
         return MLAAttentionSpec(
             **common_kwargs,
+            is_index_group_leader=self.indexer is not None,
             non_causal_multi_token_decode=self.non_causal_multi_token_decode,
         )
 
@@ -1328,7 +1399,7 @@ def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
                 x,
                 self.impl._w_uv_packed,  # type: ignore[attr-defined]
                 True,
-                None,
+                self.impl._w_scale,  # type: ignore[attr-defined]
             )
         else:
             # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
@@ -1342,27 +1413,21 @@ def unified_mla_kv_cache_update(
     kv_cache_dtype: str,
     k_scale: torch.Tensor,
 ) -> torch.Tensor:
-    """
-    Returns a dummy that is passed to unified_attention to signal a side effect and
+    """Returns a dummy that is passed to unified_attention to signal a side effect and
     the data dependency between them to ensure torch.compile preserves ordering.
     """
     layer_name = _resolve_layer_name(layer_name)
     attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
         layer_name
     )
+    attn_layer.prepare_kv_cache_update(attn_metadata)
     if layer_slot_mapping is not None:
-        kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs(
-            kv_c_normed,
-            k_pe,
-            layer_slot_mapping,
-            attn_metadata.num_decode_tokens if attn_metadata is not None else None,
-            attn_layer.use_pcp,
-        )
-        attn_layer.impl.do_kv_cache_update(  # type: ignore[attr-defined]
+        attn_layer.update_kv_cache(
             kv_c_normed,
             k_pe,
             kv_cache,
             layer_slot_mapping,
+            attn_metadata,
             kv_cache_dtype,
             k_scale,
         )
@@ -1410,6 +1475,8 @@ def unified_mla_attention_with_output(
     del kv_cache_dummy_dep
     layer_name = _resolve_layer_name(layer_name)
     attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
+    if layer.hisparse_cache is not None:
+        layer.hisparse_cache.finish_kv_update()
     layer.forward_impl(
         q,
         kv_c_normed,
@@ -1469,8 +1536,7 @@ def dynamic_per_batched_tensor_quant(
     dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0},
 )
 class _DecodeConcatQuantFP8(QuantFP8):
-    """
-    QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before
+    """QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before
     quantization. When disabled, forward_native is compiled via torch.compile,
     fusing cat/reshape/quant/view together.
     """
@@ -1485,7 +1551,10 @@ def forward(
             scale: torch.Tensor,
             scale_ub: torch.Tensor | None = None,
         ) -> torch.Tensor:
-            decode_q0 = torch.cat((decode_ql_nope, decode_q_pe), dim=-1)
+            if decode_q_pe.shape[-1] == 0 and decode_ql_nope.is_contiguous():
+                decode_q0 = decode_ql_nope
+            else:
+                decode_q0 = torch.cat((decode_ql_nope, decode_q_pe), dim=-1)
             decode_q_flat = decode_q0.reshape(decode_q0.shape[0], -1)
             decode_q, _ = quant_fn(self, decode_q_flat, scale, scale_ub)
             return decode_q.view(decode_q0.shape)
@@ -1502,6 +1571,7 @@ class MLACommonBackend(AttentionBackend):
     def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
         """Per-token-head modes pack an inline fp32 scale pair after the
         latent data (single-sided: ``head_size_v == 0`` for MLA)."""
+        spec = super().customize_spec(spec)
         mode = spec.kv_quant_mode
         if spec.state_content_bytes is not None or not mode.is_per_token_head:
             return spec
@@ -1909,11 +1979,6 @@ def aligned_split_len(start: int, remaining: int, available: int) -> int:
     return plans
 
 
-def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor:
-    """Pinned int32 CPU tensor backing one concatenated per-chunk field."""
-    return np_to_pinned_tensor(np.asarray(values, dtype=np.int32))
-
-
 def align_mla_chunked_context_workspace_size(
     vllm_config: VllmConfig,
     workspace_size: int,
@@ -1964,6 +2029,7 @@ def build_mla_chunked_context_metadata(
 
     Returns:
         The chunked-context metadata, or None when no prefill has any context.
+
     """
     # NOTE: it is recommended you read the `Chunked Prefill` section in the
     # comment at the top of the file before trying to understand this code.
@@ -2090,20 +2156,28 @@ def padded_rows(rows: int) -> int:
         token_offset += num_tokens
         local_token_offset += num_local_tokens
 
-    seq_lens_cpu = _flat_int32(seq_lens_flat)
-    starts = _flat_int32(starts_flat).to(device, non_blocking=True)
-    cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True)
-    cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True)
-    token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to(
-        device, non_blocking=True
+    seq_lens_cpu = torch.tensor(
+        seq_lens_flat, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY
+    )
+    starts_and_context_lens = async_tensor_h2d(
+        starts_flat + context_lens, device=device, dtype=torch.int32
+    )
+    starts = starts_and_context_lens[: len(starts_flat)]
+    context_lens_gpu = starts_and_context_lens[len(starts_flat) :]
+    cu_seq_lens = async_tensor_h2d(cu_seq_lens_flat, device=device, dtype=torch.int32)
+    cu_seqlens_q = async_tensor_h2d(cu_seqlens_q_flat, device=device, dtype=torch.int32)
+    token_to_seq = async_tensor_h2d(
+        np.concatenate(token_to_seq_parts), device=device, dtype=torch.int32
     )
     if use_dcp:
-        padded_local_cu_seq_lens = _flat_int32(padded_local_cu_seq_lens_flat).to(
-            device, non_blocking=True
+        padded_local_cu_seq_lens = async_tensor_h2d(
+            padded_local_cu_seq_lens_flat, device=device, dtype=torch.int32
+        )
+        padded_local_token_to_seq = async_tensor_h2d(
+            np.concatenate(padded_local_token_to_seq_parts),
+            device=device,
+            dtype=torch.int32,
         )
-        padded_local_token_to_seq = _flat_int32(
-            np.concatenate(padded_local_token_to_seq_parts)
-        ).to(device, non_blocking=True)
 
     chunks: list[MLACommonPrefillMetadata.ContextChunk] = []
     for index, (plan, layout) in enumerate(zip(plans, layouts)):
@@ -2146,7 +2220,7 @@ def padded_rows(rows: int) -> int:
         chunks.append(chunk)
 
     return MLACommonPrefillMetadata.ChunkedContextMetadata(
-        context_lens=context_lens_cpu.to(device, non_blocking=True),
+        context_lens=context_lens_gpu,
         workspace=chunked_prefill_workspace,
         chunks=chunks,
         context_lens_list=context_lens,
@@ -2156,8 +2230,7 @@ def padded_rows(rows: int) -> int:
 
 
 class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
-    """
-    NOTE: Please read the comment at the top of the file before trying to
+    """NOTE: Please read the comment at the top of the file before trying to
     understand this class
     """
 
@@ -2242,8 +2315,7 @@ def determine_prefill_query_data_type(
         vllm_config: VllmConfig,
         model_dtype: torch.dtype,
     ) -> torch.dtype:
-        """
-        Determine the query data type for prefill queries.
+        """Determine the query data type for prefill queries.
         Return FP8 dtype if cache is FP8 and prefill query quantization
         is enabled, else model dtype.
         """
@@ -2399,6 +2471,7 @@ def _build_decode(
         query_start_loc_cpu: torch.Tensor,
         query_start_loc_device: torch.Tensor,
         num_decode_tokens: int,
+        max_query_len: int,
         dcp_tot_seq_lens_device: torch.Tensor | None,
     ) -> MLACommonDecodeMetadata:
         return MLACommonDecodeMetadata(
@@ -2410,8 +2483,7 @@ def _build_decode(
     def build_for_cudagraph_capture(
         self, common_attn_metadata: CommonAttentionMetadata
     ) -> M:
-        """
-        This method builds the metadata for full cudagraph capture.
+        """This method builds the metadata for full cudagraph capture.
         Currently, only decode is supported for full cudagraphs with MLA.
         """
         m = common_attn_metadata
@@ -2544,16 +2616,6 @@ def build(
         if num_decodes > 0:
             dcp_tot_seq_lens_device = None
             if self.dcp_world_size > 1:
-                assert seq_lens is not None, (
-                    "MLA DCP decode requires seq_lens on CommonAttentionMetadata"
-                )
-                if dcp_local_seq_lens is None:
-                    dcp_local_seq_lens = get_dcp_local_seq_lens(
-                        seq_lens,
-                        dcp_size=self.dcp_world_size,
-                        dcp_rank=get_dcp_group().rank_in_group,
-                        cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
-                    )
                 dcp_tot_seq_lens_device = seq_lens[:num_decodes]
                 seq_lens = dcp_local_seq_lens
 
@@ -2574,6 +2636,9 @@ def build(
                 query_start_loc_cpu=query_start_loc_cpu[: num_decodes + 1],
                 query_start_loc_device=query_start_loc[: num_decodes + 1],
                 num_decode_tokens=num_decode_tokens,
+                max_query_len=min(
+                    common_attn_metadata.max_query_len, self.reorder_batch_threshold
+                ),
                 dcp_tot_seq_lens_device=dcp_tot_seq_lens_device,
             )
 
@@ -2607,14 +2672,16 @@ def reorg_kvcache(
     max_seq_len: int,
     toks: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    reorg and unpad kvcache after cp local gather to tp layout for attn kernel.
+    """Reorg and unpad kvcache after cp local gather to tp layout for attn kernel.
     e.g.
     allgatered_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T1_0, T1_1, ...,
                               T0_4, T0_5, pad, pad, T1_2, pad, ...]
     -> reorganized_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T0_4, T0_5,
                                   T1_0, T1_1, T1_2, ...]
+
     Args:
+        allgatered_kv_c_normed: all-gathered, padded latent KV cache.
+        allgatered_k_pe: all-gathered, padded RoPE key cache.
         padded_local_chunk_seq_lens_lst: local chunk context lengths
             under current CP rank.
         local_context_lens_allranks: local context lengths on each CP rank.
@@ -2623,6 +2690,7 @@ def reorg_kvcache(
         sum_seq_len: the sum of cp_chunk_seq_lens_lst.
         max_seq_len: the max value of cp_chunk_seq_lens_lst.
         toks: the number of tokens for local gather cache.
+
     """
     kv_c_segments = []
     k_pe_segments = []
@@ -2725,10 +2793,16 @@ def accumulate_mla_context_chunk(
     remaining token range is initialized.
 
     Args:
+        chunk: The context chunk being folded in.
+        attn_output: The chunk's attention output.
+        attn_softmax_lse: The chunk's log-sum-exp values.
+        output: Running context partial, updated in place.
+        output_lse: Running log-sum-exp, updated in place.
         output_written: The chunk's attention output already landed in
             ``output[chunk.token_slice]`` because the backend was handed it as
             ``out``, leaving only the lse to fold. Invalid for a continuation
             chunk, whose leading tokens must be merged rather than overwritten.
+
     """
     token_start = chunk.token_slice.start
     token_end = chunk.token_slice.stop
@@ -2755,8 +2829,7 @@ def accumulate_mla_context_chunk(
 
 
 class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
-    """
-    Shared MLA base providing dense-MHA prefill (via the selected
+    """Shared MLA base providing dense-MHA prefill (via the selected
     MLAPrefillBackend) for both dense and sparse impls; subclasses add decode
     (``forward_mqa``).
     """
@@ -2792,8 +2865,7 @@ def __init__(
     def _concat_k_nope_k_pe(
         self, k_nope: torch.Tensor, k_pe: torch.Tensor
     ) -> torch.Tensor:
-        """
-        Efficiently concatenate k_nope and k_pe tensors along the last dimension.
+        """Efficiently concatenate k_nope and k_pe tensors along the last dimension.
 
         This function avoids the performance penalty of torch.cat with expanded
         non-contiguous tensors by pre-allocating the output and using direct copies.
@@ -2805,7 +2877,12 @@ def _concat_k_nope_k_pe(
 
         Returns:
             Tensor of shape [..., nope_dim + pe_dim]
+
         """
+        if k_pe.shape[-1] == 0:
+            # NoPE MLA: nothing to append, so no copy either.
+            return k_nope
+
         k = torch.empty(
             (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
             dtype=k_nope.dtype,
@@ -2846,12 +2923,12 @@ def _compute_prefill_context(
         output_lse = None
         for chunk in chunked_context.chunks:
             toks = chunk.num_context_tokens
-            block_table = prefill_metadata.block_table[chunk.request_slice]
+            chunk_block_table = prefill_metadata.block_table[chunk.request_slice]
             if self.kv_cache_dtype == "fp8_ds_mla":
                 ops.cp_gather_and_upconvert_fp8_kv_cache(
                     src_cache=kv_c_and_k_pe_cache,
                     dst=workspace[:toks],
-                    block_table=block_table,
+                    block_table=chunk_block_table,
                     workspace_starts=chunk.cu_seq_lens,
                     batch_size=chunk.num_requests,
                     seq_starts=chunk.starts,
@@ -2864,7 +2941,7 @@ def _compute_prefill_context(
                 ops.gather_mla_context_cache_cpu(
                     src_cache=kv_c_and_k_pe_cache,
                     dst=workspace[:toks],
-                    block_table=block_table,
+                    block_table=chunk_block_table,
                     starts=chunk.starts,
                     cu_seq_lens=chunk.cu_seq_lens,
                 )
@@ -2872,7 +2949,7 @@ def _compute_prefill_context(
                 ops.gather_and_maybe_dequant_cache(
                     src_cache=kv_c_and_k_pe_cache,
                     dst=workspace,
-                    block_table=block_table,
+                    block_table=chunk_block_table,
                     cu_seq_lens=chunk.cu_seq_lens,
                     token_to_seq=chunk.token_to_seq,
                     num_tokens=toks,
@@ -2885,7 +2962,7 @@ def _compute_prefill_context(
                 ops.cp_gather_cache(
                     src_cache=kv_c_and_k_pe_cache,
                     dst=workspace[:toks],
-                    block_table=block_table,
+                    block_table=chunk_block_table,
                     cu_seq_lens=chunk.cu_seq_lens,
                     batch_size=chunk.num_requests,
                     seq_starts=chunk.starts,
@@ -3156,8 +3233,7 @@ def forward_mha(  # type: ignore[override]
 
 
 class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
-    """
-    NOTE: Please read the comment at the top of the file before trying to
+    """NOTE: Please read the comment at the top of the file before trying to
     understand this class
     """
 
diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py
index b8f431ed415f..0af2396ed719 100644
--- a/vllm/model_executor/layers/attention/mm_encoder_attention.py
+++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py
@@ -322,14 +322,14 @@ def __init__(
         num_kv_heads: int | None = None,
         prefix: str = "",
     ) -> None:
-        """
-        Args:
-            num_heads: number of attention heads per partition.
-            head_size: hidden_size per attention head.
-            scale: scale factor.
-            num_kv_heads: number of kv heads.
-            prefix: This has no effect, it is only here to make it easier to
-                    swap between Attention and MultiHeadAttention
+        """Args:
+        num_heads: number of attention heads per partition.
+        head_size: hidden_size per attention head.
+        scale: scale factor.
+        num_kv_heads: number of kv heads.
+        prefix: This has no effect, it is only here to make it easier to
+                swap between Attention and MultiHeadAttention
+
         """
         super().__init__()
 
@@ -517,8 +517,7 @@ def view_qkv_to_4d(
         q_len: int,
         kv_len: int,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Reshape query, key, value to 4D tensors:
+        """Reshape query, key, value to 4D tensors:
         (batch_size, seq_len, num_heads, head_size)
         """
         query = query.view(bsz, q_len, self.num_heads, self.head_size)
diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py
index ab6946132f4b..aed5afdcf905 100644
--- a/vllm/model_executor/layers/attention/sparse_mla_attention.py
+++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py
@@ -3,13 +3,15 @@
 """Shared MHA implementation and metadata builder for sparse MLA backends."""
 
 import math
+from dataclasses import dataclass
 from shutil import which
-from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar
 
 import numpy as np
 import torch
 
 from vllm import _custom_ops as ops
+from vllm.config import get_current_vllm_config
 from vllm.distributed import (
     get_dcp_group,
     get_tensor_model_parallel_world_size,
@@ -17,6 +19,7 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention.mla_attention import (
     MLACommonBaseImpl,
+    MLACommonDecodeMetadata,
     MLACommonMetadata,
     MLACommonPrefillMetadata,
     accumulate_mla_context_chunk,
@@ -28,12 +31,24 @@
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import has_flashinfer
-from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
-from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder
+from vllm.utils.torch_utils import (
+    PIN_MEMORY,
+    is_quantized_kv_cache,
+    np_to_pinned_tensor,
+)
+from vllm.v1.attention.backend import AttentionMetadataBuilder
 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
+from vllm.v1.attention.backends.mla.index_group import (
+    SparseMLAIndexGroup,
+    SparseMLAIndexGroupBuilder,
+)
 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
 from vllm.v1.attention.ops.dcp import MLADCPManager
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
+from vllm.v1.hisparse.runtime import (
+    HiSparsePrefillStagingPlan,
+    build_hisparse_prefill_staging_plan,
+)
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -43,8 +58,6 @@
 
 logger = init_logger(__name__)
 
-T = TypeVar("T", bound=AttentionMetadata)
-
 GLOBAL_TOPK_MASK_MAX_BYTES = 128 * 1024 * 1024  # 128 MiB
 
 
@@ -98,6 +111,8 @@ def _is_masked_mha_available(
     if model_dims not in (
         (128, 512, 128, 64, 128),
         (64, 512, 192, 64, 256),
+        # GLM-5.3-Flash: NoPE, qk_head_dim 256 == the (192, 64, 256) kernel.
+        (64, 512, 256, 0, 256),
     ):
         return False
     qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
@@ -112,9 +127,45 @@ def _is_masked_mha_available(
     return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype)
 
 
+def _use_dense_mha_prefill(
+    vllm_config: "VllmConfig", prefill_max_seq_len: int, topk_tokens: int
+) -> bool:
+    attention_config = vllm_config.attention_config
+    return (
+        attention_config.hisparse_config is None
+        and not attention_config.sparse_mla_force_mqa
+        and prefill_max_seq_len <= topk_tokens
+    )
+
+
+@dataclass
+class SparseMLAPrefillMetadata(MLACommonPrefillMetadata):
+    host_staging_plan: HiSparsePrefillStagingPlan | None = None
+
+
+@dataclass(kw_only=True)
+class SparseMLACommonMetadata(MLACommonMetadata[MLACommonDecodeMetadata]):
+    block_table: torch.Tensor
+    req_id_per_token: torch.Tensor
+    seq_lens: torch.Tensor | None = None
+    block_size: int = 64
+    topk_tokens: int = 2048
+    num_decodes: int = 0
+    num_prefills: int = 0
+    num_decode_tokens: int = 0
+    decode_max_query_len: int = 0
+    prefill_max_seq_len: int = 0
+    prefill: SparseMLAPrefillMetadata | None = None
+    cp_kv_cache_interleave_size: int = 1
+
+
+T = TypeVar("T", bound=SparseMLACommonMetadata)
+
+
 class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
     metadata_cls: type[T]
     require_uniform_decodes: ClassVar[bool] = False
+    hisparse_supports_multi_token_decode: ClassVar[bool] = False
 
     def __init__(
         self,
@@ -128,14 +179,21 @@ def __init__(
         self.device = device
         self.model_config = vllm_config.model_config
         self.mla_dims = get_mla_dims(self.model_config)
-        self.topk_tokens: int = vllm_config.model_config.hf_config.index_topk
+        self.topk_tokens: int = vllm_config.model_config.hf_text_config.index_topk
         self.req_id_per_token_buffer = torch.empty(
             (vllm_config.scheduler_config.max_num_batched_tokens,),
             dtype=torch.int32,
             device=device,
         )
+        self._context_lens_cpu = torch.empty(
+            vllm_config.scheduler_config.max_num_seqs,
+            dtype=torch.int32,
+            device="cpu",
+            pin_memory=PIN_MEMORY,
+        )
         parallel_config = vllm_config.parallel_config
-        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
+        self.pcp_world_size = parallel_config.prefill_context_parallel_size
+        self.use_pcp = self.pcp_world_size > 1
         try:
             self.dcp_world_size = get_dcp_group().world_size
         except AssertionError:
@@ -193,12 +251,28 @@ def __init__(
             layer_prefill_backend.clone() if layer_prefill_backend is not None else None
         )
 
+    def _init_reorder_batch_threshold(
+        self,
+        reorder_batch_threshold: int | None = 1,
+        supports_spec_as_decode: bool = False,
+        supports_dcp_with_varlen: bool = False,
+    ) -> None:
+        if self.vllm_config.attention_config.hisparse_config is not None:
+            reorder_batch_threshold = 1
+            if not self.hisparse_supports_multi_token_decode:
+                supports_spec_as_decode = False
+        super()._init_reorder_batch_threshold(
+            reorder_batch_threshold,
+            supports_spec_as_decode,
+            supports_dcp_with_varlen,
+        )
+
     @staticmethod
     def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int:
         scheduler_config = vllm_config.scheduler_config
         cache_config = vllm_config.cache_config
         model_config = vllm_config.model_config
-        topk_tokens = model_config.hf_config.index_topk
+        topk_tokens = model_config.hf_text_config.index_topk
 
         workspace_size = min(
             max(
@@ -241,9 +315,11 @@ def _build_chunked_context_fields(
 
         seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
         assert seq_lens_cpu is not None
-        context_lens_cpu = (
-            seq_lens_cpu[num_decodes : num_decodes + num_prefills]
-            - prefill_query_lens_cpu
+        context_lens_cpu = self._context_lens_cpu[:num_prefills]
+        torch.subtract(
+            seq_lens_cpu[num_decodes : num_decodes + num_prefills],
+            prefill_query_lens_cpu,
+            out=context_lens_cpu,
         )
         qsl_cpu = common_attn_metadata.query_start_loc_cpu
         prefill_query_start_loc_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]
@@ -281,34 +357,70 @@ def build(
             prefill_max_query_len,
             prefill_query_lens_cpu,
         ) = self._build_prefill_fields(common_attn_metadata, num_decodes, num_prefills)
+        decode_max_query_len = 0
+        if num_decodes:
+            if common_attn_metadata.max_logits_per_req is not None:
+                decode_max_query_len = common_attn_metadata.max_logits_per_req
+            elif num_prefills == 0:
+                decode_max_query_len = common_attn_metadata.max_query_len
+            else:
+                decode_query_lens_cpu = torch.diff(
+                    common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
+                )
+                decode_max_query_len = int(decode_query_lens_cpu.max().item())
 
         prefill_max_seq_len = 0
-        prefill: MLACommonPrefillMetadata | None = None
+        prefill: SparseMLAPrefillMetadata | None = None
         if num_prefills > 0 and self._prefill_backend is not None:
             seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
             assert seq_lens_cpu is not None
             prefill_max_seq_len = int(
                 seq_lens_cpu[num_decodes : num_decodes + num_prefills].max().item()
             )
-            prefill = MLACommonPrefillMetadata(
-                block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...],
+            block_table = common_attn_metadata.block_table_tensor[num_decodes:, ...]
+            chunked_context = self._build_chunked_context_fields(
+                common_attn_metadata,
+                num_decodes,
+                num_prefills,
+                prefill_query_lens_cpu,
+            )
+            staging_plan = None
+            if self.vllm_config.attention_config.hisparse_config is not None:
+                prefill_seq_lens_cpu = seq_lens_cpu[
+                    num_decodes : num_decodes + num_prefills
+                ]
+                staging_block_capacity = int(
+                    (
+                        (prefill_seq_lens_cpu + self.kv_cache_spec.block_size - 1)
+                        // self.kv_cache_spec.block_size
+                    ).sum()
+                )
+                staging_plan = build_hisparse_prefill_staging_plan(
+                    block_table,
+                    common_attn_metadata.seq_lens[num_decodes:],
+                    self.kv_cache_spec.block_size,
+                    staging_block_capacity,
+                )
+                if chunked_context is not None:
+                    block_table = staging_plan.block_table
+            prefill = SparseMLAPrefillMetadata(
+                block_table=block_table,
                 query_start_loc=prefill_query_start_loc,
                 max_query_len=prefill_max_query_len,
                 query_lens_cpu=prefill_query_lens_cpu,
-                chunked_context=self._build_chunked_context_fields(
-                    common_attn_metadata,
-                    num_decodes,
-                    num_prefills,
-                    prefill_query_lens_cpu,
-                ),
+                chunked_context=chunked_context,
                 q_data_type=self.model_config.dtype,
                 output_dtype=self.model_config.dtype,
                 prefill_backend=self._prefill_backend,
                 use_dense_mha=(
-                    prefill_max_seq_len <= self.topk_tokens
-                    and not self.vllm_config.attention_config.sparse_mla_force_mqa
+                    _use_dense_mha_prefill(
+                        self.vllm_config, prefill_max_seq_len, self.topk_tokens
+                    )
+                    # Dense MHA cannot be used under PCP+DCP.
+                    and not (self.use_pcp and self.dcp_world_size > 1)
                 ),
                 topk_mask_workspace=self.topk_mask_workspace,
+                host_staging_plan=staging_plan,
             )
             self._prefill_backend.prepare_metadata(prefill)
 
@@ -327,6 +439,7 @@ def build(
             num_decodes=num_decodes,
             num_prefills=num_prefills,
             num_decode_tokens=num_decode_tokens,
+            decode_max_query_len=decode_max_query_len,
             prefill_max_seq_len=prefill_max_seq_len,
             prefill=prefill,
             cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
@@ -493,7 +606,40 @@ def _build_topk_mask(
     return out[:batch_size, :max_q_len]
 
 
-class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
+class SharedTopkIndicesBuffer:
+    """Resolves the shared top-k index buffer for sparse MLA implementations.
+
+    The indexer owns the buffer, but `LLMBaseProposer.load_model` repoints the
+    draft's indexers at the target's buffer after the impls are constructed, so
+    it must be resolved per read rather than snapshotted. Backbone skip-topk
+    layers have no indexer and pass the buffer explicitly.
+    """
+
+    _indexer: object | None = None
+    _topk_indices_buffer: torch.Tensor | None = None
+
+    def init_topk_indices_buffer(
+        self,
+        indexer: object | None,
+        topk_indices_buffer: torch.Tensor | None,
+    ) -> None:
+        self._indexer = indexer
+        self._topk_indices_buffer = topk_indices_buffer
+
+    @property
+    def topk_indices_buffer(self) -> torch.Tensor | None:
+        if self._indexer is not None:
+            return self._indexer.topk_indices_buffer  # type: ignore[attr-defined]
+        return self._topk_indices_buffer
+
+    @topk_indices_buffer.setter
+    def topk_indices_buffer(self, buffer: torch.Tensor | None) -> None:
+        # An explicit assignment supersedes the indexer.
+        self._indexer = None
+        self._topk_indices_buffer = buffer
+
+
+class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic[T]):
     """Sparse MLA base with dense and masked-MHA prefill paths."""
 
     is_sparse = True
@@ -519,6 +665,7 @@ def __init__(
         kv_b_proj: "ColumnParallelLinear",
         indexer: object | None = None,
         topk_indices_buffer: torch.Tensor | None = None,
+        index_group_builder: SparseMLAIndexGroupBuilder | None = None,
         q_pad_num_heads: int | None = None,
     ) -> None:
         super().__init__(
@@ -535,14 +682,22 @@ def __init__(
             kv_b_proj,
         )
 
-        # The indexer carries the shared buffer for normal layers and tests;
-        # the explicitly-passed buffer covers backbone skip layers, whose
-        # indexer is not constructed (see deepseek_v2.py).
-        self.topk_indices_buffer: torch.Tensor | None = (
-            indexer.topk_indices_buffer  # type: ignore[attr-defined]
-            if indexer is not None
-            else topk_indices_buffer
-        )
+        self.init_topk_indices_buffer(indexer, topk_indices_buffer)
+        self.index_group: SparseMLAIndexGroup | None = None
+        self.index_group_index = 0
+        if index_group_builder is None and self.topk_indices_buffer is not None:
+            index_group_builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer)
+        if index_group_builder is not None:
+            vllm_config = get_current_vllm_config()
+            self.index_group, self.index_group_index = (
+                index_group_builder.register_layer(
+                    indexer is not None,
+                    vllm_config,
+                    head_size=head_size,
+                    kv_cache_dtype=kv_cache_dtype,
+                )
+            )
+
         self._use_flashinfer_concat_mla_k = (
             has_flashinfer()
             and which("ninja") is not None
@@ -559,6 +714,33 @@ def __init__(
             kv_cache_dtype=kv_cache_dtype,
         )
 
+    def record_logical_topk_ready(self) -> None:
+        if self.index_group is not None:
+            self.index_group.set_logical_topk_ready(self.index_group_index)
+
+    def prepare_for_batch(self, attn_metadata: T | None) -> None:
+        if self.index_group is not None:
+            self.index_group.prepare_for_batch(self.index_group_index, attn_metadata)
+
+    def _convert_logical_to_physical_topk(
+        self,
+        logical_topk_indices: torch.Tensor,
+        attn_metadata: Any,
+        *,
+        block_stride_rows: int | None,
+        return_valid_counts: bool,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        group = self.index_group
+        assert group is not None
+        assert self.dcp_world_size == 1
+        return group.convert_logical_to_physical_topk(
+            self.index_group_index,
+            logical_topk_indices,
+            attn_metadata,
+            block_stride_rows=block_stride_rows,
+            return_valid_counts=return_valid_counts,
+        )
+
     @staticmethod
     def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool:
         """Whether this prefill batch's top-k masks fit the workspace."""
@@ -819,8 +1001,10 @@ def forward_mha(  # type: ignore[override]
         output: torch.Tensor,
         output_scale: torch.Tensor | None = None,
     ) -> None:
-        prefill_max_seq_len = attn_metadata.prefill_max_seq_len  # type: ignore[attr-defined]
-        topk_tokens = attn_metadata.topk_tokens  # type: ignore[attr-defined]
+        prefill_metadata = attn_metadata.prefill
+        assert prefill_metadata is not None
+        prefill_max_seq_len = attn_metadata.prefill_max_seq_len
+        topk_tokens = attn_metadata.topk_tokens
         force_dense = getattr(self, "_sparse_mla_force_dense_mha", False)
         force_masked = getattr(self, "_sparse_mla_force_masked_mha", False)
         if force_dense or (prefill_max_seq_len <= topk_tokens and not force_masked):
@@ -829,7 +1013,7 @@ def forward_mha(  # type: ignore[override]
                 kv_c_normed,
                 k_pe,
                 kv_c_and_k_pe_cache,
-                cast(MLACommonMetadata, attn_metadata),
+                attn_metadata,
                 k_scale,
                 output,
                 output_scale,
@@ -837,13 +1021,11 @@ def forward_mha(  # type: ignore[override]
 
         assert output_scale is None
         assert self.masked_mha_available
-        prefill_metadata = attn_metadata.prefill  # type: ignore[attr-defined]
-        assert prefill_metadata is not None
         assert prefill_metadata.query_lens_cpu is not None
         assert self.topk_indices_buffer is not None
 
         q_lens = prefill_metadata.query_lens_cpu.tolist()
-        num_decode_tokens = attn_metadata.num_decode_tokens  # type: ignore[attr-defined]
+        num_decode_tokens = attn_metadata.num_decode_tokens
         topk_all = self.topk_indices_buffer[
             num_decode_tokens : num_decode_tokens + q.shape[0]
         ]
diff --git a/vllm/model_executor/layers/attention/static_sink_attention.py b/vllm/model_executor/layers/attention/static_sink_attention.py
index d27ccc728ad7..d8917de63a3f 100644
--- a/vllm/model_executor/layers/attention/static_sink_attention.py
+++ b/vllm/model_executor/layers/attention/static_sink_attention.py
@@ -114,9 +114,7 @@ def build(
 
 @CustomOp.register("static_sink_attention")
 class StaticSinkAttention(Attention, CustomOp):
-    """
-    Attention with static sink tokens
-    """
+    """Attention with static sink tokens."""
 
     def __init__(
         self,
@@ -152,8 +150,6 @@ def __init__(
             attn_backend=attn_backend,
             **kwargs,
         )
-        CustomOp.__init__(self)
-
         self.sink_len = sink_len
         self.sink_populated = False
         self.sink_key = None
diff --git a/vllm/model_executor/layers/attention_layer_base.py b/vllm/model_executor/layers/attention_layer_base.py
index 197213dc79fe..cee57fdd041e 100644
--- a/vllm/model_executor/layers/attention_layer_base.py
+++ b/vllm/model_executor/layers/attention_layer_base.py
@@ -12,8 +12,7 @@
 
 
 class AttentionLayerBase(ABC):
-    """
-    Base class for attention-like layers (Attention, Mamba, etc.)
+    """Base class for attention-like layers (Attention, Mamba, etc.)
     that support the v1 engine.
 
     This provides a common interface for getting attention backends
@@ -38,8 +37,7 @@ def get_attn_backend(self) -> type[AttentionBackend]:
 
     @abstractmethod
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
-        """
-        Get the KV cache spec for this layer.
+        """Get the KV cache spec for this layer.
         May be None if the layer does not need KV cache.
         """
         pass
diff --git a/vllm/model_executor/layers/conv.py b/vllm/model_executor/layers/conv.py
index 14f60231cdf9..b9719ef58fdd 100644
--- a/vllm/model_executor/layers/conv.py
+++ b/vllm/model_executor/layers/conv.py
@@ -143,7 +143,7 @@ def _forward_conv(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
-        """Expected input shape: (batch_size, in_channels, height, width)"""
+        """Expected input shape: (batch_size, in_channels, height, width)."""
         assert x.dim() == 4
         if self.enable_linear:
             return self._forward_mulmat(x)
@@ -193,7 +193,7 @@ def _forward_conv(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
-        """Expected input shape: (batch_size, in_channels, time, height, width)"""
+        """Expected input shape: (batch_size, in_channels, time, height, width)."""
         if self.enable_linear:
             return self._forward_mulmat(x)
         else:
diff --git a/vllm/model_executor/layers/fused_moe/activation.py b/vllm/model_executor/layers/fused_moe/activation.py
index d83301bed19f..46a9bd1b4c58 100644
--- a/vllm/model_executor/layers/fused_moe/activation.py
+++ b/vllm/model_executor/layers/fused_moe/activation.py
@@ -115,6 +115,7 @@ def activation_without_mul(activation: str) -> str:
 
     Returns:
         The non-gated activation name (e.g., "silu_no_mul", "gelu_no_mul")
+
     """
     return MoEActivation.from_str(activation).without_mul().value
 
diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py
index 7861345453a0..3186b5d3c4e2 100644
--- a/vllm/model_executor/layers/fused_moe/all2all_utils.py
+++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py
@@ -8,6 +8,7 @@
 
 from vllm.config import get_current_vllm_config
 from vllm.distributed import (
+    GroupCoordinator,
     get_ep_group,
 )
 from vllm.logger import init_logger
@@ -34,6 +35,7 @@
 from vllm.utils.import_utils import (
     has_deep_ep,
     has_deep_ep_v2,
+    has_moonep,
     has_mori,
     has_nixl_ep,
 )
@@ -46,9 +48,18 @@ class FlashInferOneSidedDispatchLayout:
 
 
 def flashinfer_one_sided_dispatch_layout(
-    hidden_dim: int, quant_config: FusedMoEQuantConfig
+    hidden_dim: int,
+    quant_config: FusedMoEQuantConfig,
+    input_dtype: torch.dtype | None = None,
 ) -> FlashInferOneSidedDispatchLayout:
     """Return the one-sided activation payload layout."""
+    if input_dtype is not None:
+        if input_dtype not in (torch.float16, torch.bfloat16):
+            raise ValueError(
+                "flashinfer_nvlink_one_sided unpacked inputs must be float16 "
+                f"or bfloat16, got {input_dtype}"
+            )
+        return FlashInferOneSidedDispatchLayout(hidden_dim * input_dtype.itemsize, 0)
     if quant_config.quant_dtype is None:
         return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0)
     if quant_config.quant_dtype == "nvfp4":
@@ -102,18 +113,21 @@ def flashinfer_one_sided_dispatch_layout(
             NIXL_EP_QUANT_BLOCK_SHAPE,
             NixlEPPrepareAndFinalize,
         )
+    if has_moonep():
+        from .prepare_finalize.moonep import (
+            MOONEP_DEFAULT_NUM_PREFETCH_SLOTS,
+            MOONEP_DEFAULT_NUM_SMS,
+            MOONEP_DEFAULT_TOKEN_PADDING,
+            MoonEPPrepareAndFinalize,
+        )
 
 
-def get_ep_all2all_manager(eep_stage: bool = False) -> Any:
-    if eep_stage:
-        from vllm.distributed.elastic_ep.standby_state import get_standby_ep_group
-
-        ep_group = get_standby_ep_group()
-        assert ep_group is not None
-        device_communicator = ep_group.device_communicator
-    else:
-        device_communicator = get_ep_group().device_communicator
-
+def get_ep_all2all_manager(
+    ep_group: GroupCoordinator | None = None,
+) -> Any:
+    if ep_group is None:
+        ep_group = get_ep_group()
+    device_communicator = ep_group.device_communicator
     assert device_communicator is not None
     all2all_manager = device_communicator.all2all_manager
     assert all2all_manager is not None
@@ -125,8 +139,7 @@ def maybe_roundup_layer_hidden_size(
     act_dtype: torch.dtype,
     moe_parallel_config: FusedMoEParallelConfig,
 ) -> int:
-    """
-    Given layer hidden size and MoE configurations, round up hidden_size
+    """Given layer hidden size and MoE configurations, round up hidden_size
     if necessary.
 
     Args:
@@ -138,6 +151,7 @@ def maybe_roundup_layer_hidden_size(
         Rounded up hidden_size if rounding up is required based on the configs
         and all2all backend.
         Original hidden size otherwise.
+
     """
     if moe_parallel_config.use_deepep_ht_kernels:
         hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
@@ -168,7 +182,8 @@ def maybe_make_prepare_finalize(
     routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
     allow_new_interface: bool = False,
     use_monolithic: bool = False,
-    eep_stage: bool = False,
+    all2all_manager: Any | None = None,
+    input_dtype: torch.dtype | None = None,
 ) -> FusedMoEPrepareAndFinalize | None:
     if not moe.moe_parallel_config.use_all2all_kernels:
         if not allow_new_interface:
@@ -190,7 +205,8 @@ def maybe_make_prepare_finalize(
                 "Detected DP deployment with no --enable-expert-parallel. "
                 "Falling back to AllGather+ReduceScatter dispatch/combine."
             )
-            all2all_manager = get_ep_all2all_manager(eep_stage)
+            if all2all_manager is None:
+                all2all_manager = get_ep_all2all_manager()
             return make_moe_prepare_and_finalize_naive_dp_ep(
                 is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel,
                 num_dispatchers=all2all_manager.world_size,
@@ -199,7 +215,8 @@ def maybe_make_prepare_finalize(
         else:
             return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic)
 
-    all2all_manager = get_ep_all2all_manager(eep_stage)
+    if all2all_manager is None:
+        all2all_manager = get_ep_all2all_manager()
 
     prepare_finalize: FusedMoEPrepareAndFinalize | None = None
 
@@ -277,6 +294,29 @@ def maybe_make_prepare_finalize(
             num_topk=moe.experts_per_token,
             use_fp8_dispatch=use_fp8_dispatch,
             use_cudagraph=use_cudagraph,
+            sp_size=moe.moe_parallel_config.sp_size,
+        )
+
+    elif moe.use_moonep_kernels:
+        all_to_all_args = dict(
+            max_num_tokens_per_dp_rank=moe.max_num_tokens,
+            token_hidden_size=moe.hidden_dim,
+            num_topk=moe.experts_per_token,
+            num_global_experts=moe.num_experts,
+            num_prefetch_slots=MOONEP_DEFAULT_NUM_PREFETCH_SLOTS,
+            token_padding=MOONEP_DEFAULT_TOKEN_PADDING,
+            num_sms=MOONEP_DEFAULT_NUM_SMS,
+        )
+        handle = all2all_manager.get_handle(all_to_all_args)
+
+        # The [E+B] weight layout is picked up from the experts (their
+        # process_weights_after_loading hook) once the layer has loaded and
+        # converted its weights.
+        prepare_finalize = MoonEPPrepareAndFinalize(
+            buffer_pool=handle,
+            max_tokens_per_rank=moe.max_num_tokens,
+            num_dispatchers=all2all_manager.world_size,
+            num_global_experts=moe.num_experts,
         )
 
     elif moe.use_mori_kernels:
@@ -285,13 +325,20 @@ def maybe_make_prepare_finalize(
         # Note: We may want to use FP8 dispatch just to reduce
         # data movement.
         use_fp8_dispatch = (
-            quant_config.is_per_act_token or quant_config.is_block_quantized
+            quant_config.is_per_act_token
+            or quant_config.is_block_quantized
+            or quant_config.is_per_tensor
         )
         if use_fp8_dispatch:
-            # For PTPC (per token per channel) quant, scale dim is 1
-            # For 1x128 quant, scale dim is hidden_dim // 128
+            # For PTPC (per token per channel) or per-tensor quant,
+            # scale dim is 1. For 1x128 quant, scale dim is
+            # hidden_dim // 128
             quant_dtype = quant_config.quant_dtype
-            scale_dim = 1 if quant_config.is_per_act_token else moe.hidden_dim // 128
+            scale_dim = (
+                1
+                if (quant_config.is_per_act_token or quant_config.is_per_tensor)
+                else moe.hidden_dim // 128
+            )
         else:
             # Unquantized dispatch (e.g. AITER with defer_input_quant):
             # dispatch raw BF16/FP16 data, no scales needed.
@@ -330,7 +377,7 @@ def maybe_make_prepare_finalize(
             get_current_vllm_config().scheduler_config.max_num_batched_tokens
         )
         dispatch_layout = flashinfer_one_sided_dispatch_layout(
-            moe.hidden_dim, quant_config
+            moe.hidden_dim, quant_config, input_dtype=input_dtype
         )
         prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize(
             max_num_tokens=max_num_tokens,
@@ -361,10 +408,7 @@ def maybe_make_prepare_finalize(
         all_to_all_args = dict(
             max_num_tokens_per_dp_rank=moe.max_num_tokens,
             token_hidden_size=moe.hidden_dim,
-            num_ep_ranks=all2all_manager.world_size,
-            num_global_experts=moe.num_experts,
-            num_local_experts=moe.num_experts // all2all_manager.world_size,
-            stage=eep_stage,
+            num_local_experts=moe.num_local_experts,
         )
         handle = all2all_manager.get_handle(all_to_all_args)
 
@@ -378,7 +422,8 @@ def maybe_make_prepare_finalize(
         prepare_finalize = NixlEPPrepareAndFinalize(
             handle,
             max_tokens_per_rank=moe.max_num_tokens,
-            num_dispatchers=all2all_manager.world_size,
+            num_dispatchers=all2all_manager.max_num_ep_ranks,
+            expert_capacity=(moe.num_local_experts * all2all_manager.max_num_ep_ranks),
             use_fp8_dispatch=use_fp8_dispatch,
             global_to_physical=global_to_physical,
             physical_to_global=physical_to_global,
diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py
index 8736b1393d82..0c82d28096dc 100644
--- a/vllm/model_executor/layers/fused_moe/config.py
+++ b/vllm/model_executor/layers/fused_moe/config.py
@@ -17,14 +17,19 @@
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.import_utils import get_triton_kernels_version
 from vllm.utils.math_utils import cdiv
 
 logger = init_logger(__name__)
 
-if has_triton_kernels():
+_triton_kernels_version = get_triton_kernels_version()
+
+if _triton_kernels_version is not None:
     try:
-        from triton_kernels.matmul_ogs import PrecisionConfig
+        if _triton_kernels_version == "3.8":
+            from triton_kernels.matmul import PrecisionConfig
+        else:
+            from triton_kernels.matmul_ogs import PrecisionConfig
     except (ImportError, AttributeError) as e:
         logger.error(
             "Failed to import Triton kernels. Please make sure your triton "
@@ -41,8 +46,7 @@ def _get_config_dtype_str(
     use_int4_w4a16: bool = False,
     ocp_mx_scheme: str | None = None,
 ) -> str | None:
-    """
-    Return a string used to construct the filename that contains the
+    """Return a string used to construct the filename that contains the
     tuning info for a particular quantization scheme.  See
     try_get_optimal_moe_config in fused_moe.py.
     """
@@ -72,9 +76,7 @@ def _quant_flags_to_group_shape(
     per_out_ch_quant: bool,
     block_shape: list[int] | None,
 ) -> tuple[GroupShape | None, GroupShape | None]:
-    """
-    Convert MoE quantization flags into more generic GroupShapes.
-    """
+    """Convert MoE quantization flags into more generic GroupShapes."""
     a_shape: GroupShape | None
     w_shape: GroupShape | None
     if block_shape is not None:
@@ -173,8 +175,7 @@ def get_routing_method_type(
 
 @dataclass
 class FusedMoEQuantDesc:
-    """
-    A quantization descriptor for fused MoE ops. This class can describe
+    """A quantization descriptor for fused MoE ops. This class can describe
     either activations or weights.
     """
 
@@ -212,8 +213,7 @@ class FusedMoEQuantDesc:
 # e.g. for specific arguments bias, precision, etc.
 @dataclass
 class FusedMoEQuantConfig:
-    """
-    The FusedMoEQuantConfig contains all the quantization parameters for
+    """The FusedMoEQuantConfig contains all the quantization parameters for
     a single FusedMoEMethodBase operation.  It consists of four
     FusedMoEQuantDescs, one for each activation and set of weights.
 
@@ -435,8 +435,7 @@ def use_mxfp4_w4a8(self) -> bool:
         return self._a1.dtype == "fp8" and self._w1.dtype == "mxfp4"
 
     def config_name(self, dtype: torch.dtype) -> str | None:
-        """
-        Return a string used to construct the filename that contains the
+        """Return a string used to construct the filename that contains the
         tuning info for a particular quantization scheme.  See
         try_get_optimal_moe_config in fused_moe.py.
         """
@@ -454,8 +453,7 @@ def scale_shape(
         max_tokens: int,
         hidden_dim: int,
     ) -> tuple[int, int] | None:
-        """
-        Construct the proper activation scale shape for this
+        """Construct the proper activation scale shape for this
         config.
         """
         if self.is_quantized:
@@ -477,8 +475,7 @@ def batched_scale_shape(
         max_tokens: int,
         hidden_dim: int,
     ) -> tuple[int, int, int] | None:
-        """
-        Construct the proper activation batched scale shape for this
+        """Construct the proper activation batched scale shape for this
         config, e.g. (num experts, *scale_shape).
         """
         if self.is_quantized:
@@ -512,8 +509,7 @@ def make(
         gemm1_beta: float | None = None,
         gemm1_clamp_limit: float | None = None,
     ) -> "FusedMoEQuantConfig":
-        """
-        General builder function for a FusedMoEQuantConfig.
+        """General builder function for a FusedMoEQuantConfig.
         - quant_dtype: Optional quantization type. None if activations are
           unquantized or quantized prior to calling.  Note: "nvfp4", "mxfp4",
           "mxfp6_e3m2", "mxfp6_e2m3" are the only valid string values
@@ -607,9 +603,7 @@ def fp8_w8a8_moe_quant_config(
     gemm1_beta: float | None = None,
     gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for fp8 activations and fp8 weights.
-    """
+    """Construct a quant config for fp8 activations and fp8 weights."""
     return FusedMoEQuantConfig.make(
         current_platform.fp8_dtype(),
         w1_scale=w1_scale,
@@ -640,9 +634,7 @@ def int8_w8a8_moe_quant_config(
     w2_bias: torch.Tensor | None = None,
     per_act_token_quant: bool = False,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for int8 activations and int8 weights.
-    """
+    """Construct a quant config for int8 activations and int8 weights."""
     return FusedMoEQuantConfig.make(
         torch.int8,
         w1_scale=w1_scale,
@@ -667,9 +659,7 @@ def gptq_marlin_moe_quant_config(
     w1_bias: torch.Tensor | None = None,
     w2_bias: torch.Tensor | None = None,
 ):
-    """
-    Construct a quant config for gptq marlin quantization.
-    """
+    """Construct a quant config for gptq marlin quantization."""
     from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 
     w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size)
@@ -702,9 +692,7 @@ def mxfp4_w4a16_moe_quant_config(
     gemm1_beta: float | None = None,
     gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for unquantized activations and mxfp4 weights.
-    """
+    """Construct a quant config for unquantized activations and mxfp4 weights."""
     return FusedMoEQuantConfig(
         _a1=FusedMoEQuantDesc(),
         _a2=FusedMoEQuantDesc(),
@@ -730,9 +718,7 @@ def mxfp4_mxfp8_moe_quant_config(
     mx_alignment: int = 0,
     is_scale_swizzled: bool = True,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for mxfp4 activations and mxfp4 weights.
-    """
+    """Construct a quant config for mxfp4 activations and mxfp4 weights."""
     return FusedMoEQuantConfig(
         _a1=FusedMoEQuantDesc("mxfp8"),
         _a2=FusedMoEQuantDesc("mxfp8"),
@@ -756,9 +742,7 @@ def mxfp4_w4a8_moe_quant_config(
     block_shape: list[int] | None = None,
     gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for fp8 activations and mxfp4 weights.
-    """
+    """Construct a quant config for fp8 activations and mxfp4 weights."""
     return FusedMoEQuantConfig(
         _a1=FusedMoEQuantDesc("fp8", None, a1_scale, None, None, None),
         _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None),
@@ -782,9 +766,7 @@ def ocp_mx_moe_quant_config(
     gemm1_beta: float | None = None,
     gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for mxfp4 activations and mxfp4 weights.
-    """
+    """Construct a quant config for mxfp4 activations and mxfp4 weights."""
     assert quant_dtype in OCP_MX_DTYPES
     return FusedMoEQuantConfig.make(
         quant_dtype=quant_dtype,
@@ -818,9 +800,7 @@ def nvfp4_moe_quant_config(
     gemm1_beta: float | None = None,
     gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for mxfp4 activations and nvp4 weights.
-    """
+    """Construct a quant config for mxfp4 activations and nvp4 weights."""
     return FusedMoEQuantConfig.make(
         "nvfp4",
         w1_scale=w1_scale,
@@ -845,8 +825,7 @@ def mxfp4_moe_quant_config(
     w1_scale: torch.Tensor,
     w2_scale: torch.Tensor,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for MXFP4 x MXFP4 MoE.
+    """Construct a quant config for MXFP4 x MXFP4 MoE.
     MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no
     separate alphas / global activation scales in this config.
     """
@@ -869,9 +848,7 @@ def nvfp4_w4a16_moe_quant_config(
     gemm1_beta: float | None = None,
     gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for 16-but activations and nvp4 weights.
-    """
+    """Construct a quant config for 16-but activations and nvp4 weights."""
     return FusedMoEQuantConfig.make(
         quant_dtype=None,
         w1_scale=w1_scale,
@@ -899,9 +876,7 @@ def int4_w4a16_moe_quant_config(
     gemm1_alpha: float | None = None,
     gemm1_beta: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for 16-bit float activations and int4 weights.
-    """
+    """Construct a quant config for 16-bit float activations and int4 weights."""
     group_shape = GroupShape(*block_shape) if block_shape is not None else None
     return FusedMoEQuantConfig(
         _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale),
@@ -924,9 +899,7 @@ def fp8_w8a16_moe_quant_config(
     gemm1_beta: float | None = None,
     gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for 16-bit float activations and fp8 weights.
-    """
+    """Construct a quant config for 16-bit float activations and fp8 weights."""
     group_shape = GroupShape(*block_shape) if block_shape is not None else None
     fp8_dtype = current_platform.fp8_dtype()
     return FusedMoEQuantConfig(
@@ -968,9 +941,7 @@ def int8_w8a16_moe_quant_config(
     gemm1_alpha: float | None = None,
     gemm1_beta: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for 16-bit float activations and int8 weights.
-    """
+    """Construct a quant config for 16-bit float activations and int8 weights."""
     group_shape = GroupShape(*block_shape) if block_shape is not None else None
     return FusedMoEQuantConfig(
         _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale),
@@ -992,9 +963,7 @@ def int4_w4afp8_moe_quant_config(
     per_out_ch_quant: bool = False,
     block_shape: list[int] | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for fp8 activations and int4 weights.
-    """
+    """Construct a quant config for fp8 activations and int4 weights."""
     return FusedMoEQuantConfig.make(
         torch.float8_e4m3fn,  # quant dtype for activations
         w1_scale=w1_scale,
@@ -1015,8 +984,7 @@ def biased_moe_quant_config(
     gemm1_beta: float | None = None,
     gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for unquantized activations with biases.
+    """Construct a quant config for unquantized activations with biases.
 
     gemm1_alpha/gemm1_beta/gemm1_clamp_limit carry the SwiGLU gate params
     through to the fused activation kernel (e.g. swigluoai_uninterleave).
@@ -1117,6 +1085,10 @@ def use_nixl_ep_kernels(self):
     def use_deepep_v2_kernels(self):
         return self.use_all2all_kernels and self.all2all_backend == "deepep_v2"
 
+    @property
+    def use_moonep_kernels(self):
+        return self.use_all2all_kernels and self.all2all_backend == "moonep"
+
     @staticmethod
     def flatten_tp_across_dp_and_pcp(
         tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
@@ -1136,8 +1108,7 @@ def make(
         sp_size_: int,
         vllm_parallel_config: ParallelConfig,
     ) -> "FusedMoEParallelConfig":
-        """
-        Determine MoE parallel configuration. Based on the input `tp_size_`,
+        """Determine MoE parallel configuration. Based on the input `tp_size_`,
         `dp_size_` and vllm's parallel config, determine what
         level's of parallelism to use in the fused moe layer.
 
@@ -1145,6 +1116,7 @@ def make(
             tp_size_ (int): `tp_size` passed into the FusedMoEFactory constructor.
             pcp_size_ (int): `pcp_size` passed into the FusedMoEFactory constructor.
             dp_size_ (int): `dp_size` passed into the FusedMoEFactory constructor.
+            sp_size_ (int): `sp_size` passed into the FusedMoEFactory constructor.
             vllm_parallel_config (ParallelConfig): vLLM's parallel config
                 object which contains the `enable_expert_parallel` flag.
 
@@ -1208,6 +1180,7 @@ def make(
             - device 3: TP = {1, 0} DP = {2, 1} EP = {4, 3}
             - Comment: There are 2 engine instances and the experts are split
                 between the 4 devices.
+
         """
         use_ep = (
             dp_size_ * pcp_size_ * tp_size_ > 1
@@ -1306,8 +1279,10 @@ class FusedMoEConfig:
 
     moe_backend: MoEBackend = "auto"
     max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP
+    elastic_ep_max_dp_size: int | None = None
     has_bias: bool = False
     is_lora_enabled: bool = False
+    has_hash_routing: bool = False
 
     # When True, the MoE skips its final cross-rank all-reduce (and the separate
     # shared-expert reduce), returning the partial per-rank sum. The caller is
@@ -1339,6 +1314,8 @@ class FusedMoEConfig:
 
     # Set by __post_init__
     intermediate_size_per_partition: int = -1
+    # Use the allocated width as the checkpoint TP stride, including for scales.
+    tp_shard_with_padding: bool = False
     rocm_aiter_fmoe_enabled: bool = False
     aiter_fmoe_shared_expert_enabled: bool = False
 
@@ -1501,6 +1478,10 @@ def use_nixl_ep_kernels(self):
     def use_deepep_v2_kernels(self):
         return self.moe_parallel_config.use_deepep_v2_kernels
 
+    @property
+    def use_moonep_kernels(self):
+        return self.moe_parallel_config.use_moonep_kernels
+
     @property
     def needs_round_robin_routing_tables(self):
         return self.moe_parallel_config.needs_round_robin_routing_tables
diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
index 70158dafa1c7..33a229695561 100644
--- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00ba5c0cef8bff6/lightllm/common/fused_moe/deepep_scatter_gather.py
+"""Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00ba5c0cef8bff6/lightllm/common/fused_moe/deepep_scatter_gather.py
 and updated to fit vllm needs and terminology.
 """
 
@@ -394,12 +393,13 @@ def _fwd_kernel_ep_gather(
                 source_token_index = tl.load(
                     input_index + cur_token * input_index_stride0 + topk_index
                 )
+                source_token_index_i64 = source_token_index.to(tl.int64)
                 acc_weight = tl.load(
                     recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index
                 )
                 tmp = tl.load(
                     input_tensor
-                    + source_token_index * input_tensor_stride0
+                    + source_token_index_i64 * input_tensor_stride0
                     + cur_block * BLOCK_D
                     + off_d
                 )
diff --git a/vllm/model_executor/layers/fused_moe/eep_reconfigure.py b/vllm/model_executor/layers/fused_moe/eep_reconfigure.py
index 2c6c3f6aa40a..9b6a1da0f9e3 100644
--- a/vllm/model_executor/layers/fused_moe/eep_reconfigure.py
+++ b/vllm/model_executor/layers/fused_moe/eep_reconfigure.py
@@ -60,6 +60,8 @@ def _make_eep_experts(
 def make_eep_staged_quant_method(
     module: "MoERunner",
     moe_config: FusedMoEConfig,
+    *,
+    all2all_manager: Any,
 ) -> FusedMoEMethodBase | None:
     quant_method = module._quant_method
     if not quant_method.supports_internal_mk:
@@ -93,7 +95,7 @@ def make_eep_staged_quant_method(
         routing_tables=None,
         allow_new_interface=True,
         use_monolithic=quant_method.is_monolithic,
-        eep_stage=True,
+        all2all_manager=all2all_manager,
     )
     assert prepare_finalize is not None
     assert isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular)
diff --git a/vllm/model_executor/layers/fused_moe/expert_map_manager.py b/vllm/model_executor/layers/fused_moe/expert_map_manager.py
index 71f2186ea4dd..daa2b3ff8e83 100644
--- a/vllm/model_executor/layers/fused_moe/expert_map_manager.py
+++ b/vllm/model_executor/layers/fused_moe/expert_map_manager.py
@@ -27,8 +27,7 @@ def determine_expert_map(
     num_fused_shared_experts: int = 0,
     return_expert_mask: bool = False,
 ) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:
-    """
-    Calculates how many experts should be assigned to each rank for EP and
+    """Calculates how many experts should be assigned to each rank for EP and
     creates a mapping from global to local expert index. Experts are
     distributed evenly across ranks. Any remaining are assigned to the
     last rank.
@@ -56,6 +55,7 @@ def determine_expert_map(
                 and 0 for sentinel.
                 Returns None if ep_size is 1.
                 Used only when AITER MOE is enabled.
+
     """
     from typing import get_args
 
@@ -150,8 +150,7 @@ def determine_expert_placement_strategy(
 
 
 class ExpertMapManager:
-    """
-    Manages expert ID mappings and placement for Expert Parallelism.
+    """Manages expert ID mappings and placement for Expert Parallelism.
 
     Responsibilities:
     - Calculate local vs global expert counts
@@ -193,16 +192,22 @@ def __init__(
         num_fused_shared_experts: int = 0,
         rocm_aiter_enabled: bool = False,
     ):
-        """
-        Initialize expert map manager.
+        """Initialize expert map manager.
 
         Args:
+            max_num_batched_tokens: Maximum number of batched tokens per step
+            top_k: Number of experts each token is routed to
             global_num_experts: Total number of experts across all ranks
+            num_redundant_experts: Number of redundant expert replicas (EPLB)
+            num_expert_group: Number of expert groups, when grouped routing is
+                used
             moe_parallel_config: MoE parallel configuration (contains ep_size,
                                  ep_rank, backend flags)
             placement_strategy: Strategy for placing experts ('linear' or 'round_robin')
             num_fused_shared_experts: Number of fused shared experts (for AITER)
             rocm_aiter_enabled: Whether ROCm AITER fusion is enabled
+            enable_eplb: Whether expert-parallel load balancing is enabled
+
         """
         self.global_num_experts = global_num_experts
         self.moe_parallel_config = moe_parallel_config
@@ -295,8 +300,7 @@ def local_num_experts(self) -> int:
 
     @property
     def expert_map(self) -> torch.Tensor | None:
-        """
-        Mapping from global expert ID to local expert ID.
+        """Mapping from global expert ID to local expert ID.
 
         Returns tensor of shape (global_num_experts,) where:
         - expert_map[global_id] = local_id if expert is on this rank
@@ -308,8 +312,7 @@ def expert_map(self) -> torch.Tensor | None:
 
     @property
     def expert_mask(self) -> torch.Tensor | None:
-        """
-        Expert mask for AITER fusion (ROCm-specific).
+        """Expert mask for AITER fusion (ROCm-specific).
 
         Returns tensor of shape (global_num_experts + num_fused_shared + 1,)
         where 1 indicates expert is on this rank, 0 otherwise.
@@ -325,8 +328,7 @@ def placement_strategy(self) -> ExpertPlacementStrategy:
     def routing_tables(
         self,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
-        """
-        Routing tables for round-robin placement.
+        """Routing tables for round-robin placement.
 
         Returns (global_to_physical, physical_to_global, local_to_global)
         or None if not using round-robin or tables not needed.
@@ -334,8 +336,7 @@ def routing_tables(
         return self._routing_tables
 
     def map_global_to_local(self, global_id: int) -> int:
-        """
-        Map global expert ID to local expert ID.
+        """Map global expert ID to local expert ID.
 
         Args:
             global_id: Global expert ID (0 to global_num_experts - 1)
@@ -345,6 +346,7 @@ def map_global_to_local(self, global_id: int) -> int:
 
         Raises:
             ValueError: If expert is not on this rank
+
         """
         if self._expert_map is None:
             return global_id
@@ -369,8 +371,7 @@ def update(
         moe_parallel_config: FusedMoEParallelConfig,
         global_num_experts: int,
     ) -> None:
-        """
-        Update expert mappings for new EP configuration.
+        """Update expert mappings for new EP configuration.
 
         Used during dynamic reconfiguration (e.g., elastic scaling).
 
@@ -378,6 +379,7 @@ def update(
             global_num_experts: New total number of experts across all ranks
             moe_parallel_config: New MoE parallel configuration (contains ep_size,
                                  ep_rank, backend flags)
+
         """
         self.moe_parallel_config = moe_parallel_config
         self.global_num_experts = global_num_experts
@@ -397,8 +399,7 @@ def update(
             self._init_aiter_shared_experts_topK_buffer()
 
     def get_compressed_map_string(self) -> str:
-        """
-        Get compressed string representation of expert map for logging.
+        """Get compressed string representation of expert map for logging.
 
         Returns string mapping local to global expert IDs.
         """
@@ -457,8 +458,7 @@ def _calculate_expert_maps(self) -> None:
     def _init_routing_tables(
         self,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
-        """
-        Ensure routing tables are initialized if needed for round-robin.
+        """Ensure routing tables are initialized if needed for round-robin.
 
         This is a public method that can be called to explicitly initialize
         routing tables. It's safe to call multiple times (idempotent).
diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py
index ef258dd595c9..fb3256e1a4c1 100644
--- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py
@@ -23,17 +23,6 @@
 ]
 
 
-def _aiter_raw(t):
-    """Unwrap the `triton_kernels.tensor.Tensor` that `_swizzle_mxfp4` returns."""
-    if t is None or isinstance(t, torch.Tensor):
-        return t
-    assert hasattr(t, "storage"), (
-        f"expected a triton_kernels wrapped tensor with a .storage attribute, "
-        f"got {type(t)}"
-    )
-    return t.storage.data
-
-
 def _aiter_w4a16_silu_via_a8w4(
     hidden_states: torch.Tensor,
     w1_data,
@@ -53,9 +42,7 @@ def _aiter_w4a16_silu_via_a8w4(
     unpadded_N_w2,
     unpadded_K_w2,
 ) -> torch.Tensor:
-    """
-    MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation.
-    """
+    """MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation."""
     from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
     from aiter.ops.triton.quant import dynamic_mxfp8_quant
 
@@ -195,10 +182,14 @@ def aiter_triton_kernel_w4a16_moe_forward(
     assert quant_config.w1_precision is not None
     assert quant_config.w2_precision is not None
 
-    w1_data = _aiter_raw(w1)
-    w2_data = _aiter_raw(w2)
-    w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale)
-    w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale)
+    w1_data = w1.storage.data
+    w2_data = w2.storage.data
+    from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+        weight_mx_scale,
+    )
+
+    w1_wscale = weight_mx_scale(quant_config.w1_precision).storage.data
+    w2_wscale = weight_mx_scale(quant_config.w2_precision).storage.data
 
     gammas = routing_data.gate_scal if routing_data else None
 
@@ -330,7 +321,7 @@ def _supports_parallel_config(
         moe_parallel_config: FusedMoEParallelConfig,
     ) -> bool:
         return (
-            not moe_parallel_config.use_all2all_kernels
+            not moe_parallel_config.use_ep
             and not moe_parallel_config.enable_eplb
             and moe_parallel_config.dp_size <= 1
         )
diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
index 3a217336377f..8bbdde28a7c8 100644
--- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
@@ -130,6 +130,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
     from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
         should_use_cdna4_mx_scale_swizzle,
         should_use_gfx1250_mx_scale_swizzle,
+        weight_mx_scale,
     )
 
     if should_use_gfx1250_mx_scale_swizzle():
@@ -154,7 +155,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
         hidden_states,
         w1.storage.data,
         None,
-        quant_config.w1_precision.weight_scale.storage.data,
+        weight_mx_scale(quant_config.w1_precision).storage.data,
         quant_config.w1_precision.flex_ctx.lhs_data.scale,
         quant_config.w2_precision.flex_ctx.lhs_data.scale,
         quant_config.w1_bias,
@@ -174,7 +175,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
         intermediate_cache1,
         w2.storage.data,
         None,
-        quant_config.w2_precision.weight_scale.storage.data,
+        weight_mx_scale(quant_config.w2_precision).storage.data,
         quant_config.w2_precision.flex_ctx.lhs_data.scale,
         None,
         quant_config.w2_bias,
@@ -190,8 +191,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
 
 
 class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
-    """
-    Monolithic MXFP4 W4A8 expert using AITER triton kernels.
+    """Monolithic MXFP4 W4A8 expert using AITER triton kernels.
 
     This backend uses:
     - aiter.ops.triton.moe_routing.routing for routing
@@ -250,7 +250,7 @@ def _supports_parallel_config(
         moe_parallel_config: FusedMoEParallelConfig,
     ) -> bool:
         return (
-            not moe_parallel_config.use_all2all_kernels
+            not moe_parallel_config.use_ep
             and not moe_parallel_config.enable_eplb
             and moe_parallel_config.dp_size <= 1
         )
@@ -339,8 +339,7 @@ def _aiter_w4a16_silu_via_a8w4(
     unpadded_N_w2,
     unpadded_K_w2,
 ) -> torch.Tensor:
-    """
-    MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation.
+    """MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation.
     """
     from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
     from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
index 93a4088fc4f7..c15a1a13d42e 100644
--- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
@@ -11,6 +11,7 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
@@ -24,41 +25,6 @@
 _AITER_SWIGLU_BETA = 1.0
 
 
-def is_aiter_mxfp8_moe_available() -> bool:
-    """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl``
-    package is importable, AND the installed aiter carries the mxfp8 FlyDSL
-    2-stage support from ROCm/aiter#3811.
-
-    ``flydsl`` and ``aiter`` are separate packages, so ``is_flydsl_available()``
-    (flydsl pkg + arch) is necessary but not sufficient: an older aiter without
-    #3811 still ships the flydsl pkg and the ``aiter.ops.flydsl`` module but a
-    broken/missing ``per_1x32 + fp8`` 2-stage path. Without this extra gate a
-    nightly lacking #3811 would wrongly select FlyDSL instead of falling back to
-    the native Triton dot_scaled path. #3811 added no probe-able public symbol,
-    so detect the ``minimax_m3_mxfp8`` tuned config it shipped. Every check fails
-    closed (returns False -> triton dot_scaled), which is always safe."""
-    if not (current_platform.is_rocm() and current_platform.supports_mx()):
-        return False
-    try:
-        import os
-
-        import aiter
-        from aiter.ops.flydsl.utils import is_flydsl_available
-
-        if not is_flydsl_available():
-            return False
-        return os.path.exists(
-            os.path.join(
-                os.path.dirname(aiter.__file__),
-                "configs",
-                "model_configs",
-                "minimax_m3_mxfp8_tuned_fmoe.csv",
-            )
-        )
-    except Exception:
-        return False
-
-
 class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
     """MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950)."""
 
@@ -79,10 +45,7 @@ def expects_unquantized_inputs(self) -> bool:
 
     @staticmethod
     def _supports_current_device() -> bool:
-        # Device capability only (gfx950 / MX-capable ROCm). The flydsl package
-        # check lives in is_supported_config so a missing package is reported
-        # distinctly from an unsupported device.
-        return current_platform.is_rocm() and current_platform.supports_mx()
+        return current_platform.supports_mx() and rocm_aiter_ops.is_fused_moe_enabled()
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config) -> bool:
@@ -95,12 +58,6 @@ def is_supported_config(
         is_supported, reason = super().is_supported_config(
             cls, moe_config, weight_key, activation_key, activation_format
         )
-        # _supports_current_device() only gates on the device; surface a clear
-        # reason when the device is fine but the flydsl package is missing.
-        if is_supported and not is_aiter_mxfp8_moe_available():
-            return False, (
-                "kernel requires the aiter flydsl package, which is not installed"
-            )
         if (
             is_supported
             and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE
diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
index c578fb641b10..d3ea8f274cff 100644
--- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
@@ -307,8 +307,7 @@ def __init__(
         max_num_tokens: int,
         num_dispatchers: int,
     ):
-        """
-        max_num_tokens: Maximum number of tokens from a DP Rank
+        """max_num_tokens: Maximum number of tokens from a DP Rank
         num_dispatchers: The number of DP dispatchers.
         quant_config: Quantization configuration
         """
@@ -350,8 +349,7 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo
         return True
 
     def supports_packed_ue8m0_act_scales(self) -> bool:
-        """
-        DeepGemm supports packed ue8m0 activation scales on Blackwell-family
+        """DeepGemm supports packed ue8m0 activation scales on Blackwell-family
         GPUs (SM100 datacenter and SM120 consumer).
         """
         return is_deep_gemm_e8m0_used() and (
diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py
index b1e32462d842..5ee98a8c847f 100644
--- a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py
@@ -172,6 +172,7 @@ def _activation_kind(activation: MoEActivation) -> int:
             0 = SwiGLU_Gu (SiLU(g)*u)
             1 = SwiGLU_Ug (SiLU(u)*g)
             2 = SiLU
+
         """
         if activation == MoEActivation.SWIGLUSTEP:
             return 0
@@ -217,6 +218,7 @@ def apply(
 
         Returns:
             Output tensor after MoE computation
+
         """
         from vllm.model_executor.layers.fused_moe.router.cpu_router import (
             select_experts,
diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
index 44a5bda21e79..3e7376fe9015 100644
--- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
@@ -15,6 +15,7 @@
     convert_weight_packed_scale_zp,
     cpu_fused_moe,
     cpu_fused_moe_int8,
+    cpu_has_amx_fp8,
     cpu_prepack_moe_weight,
     cpu_prepack_moe_weight_int8,
     fused_experts_cpu,
@@ -467,7 +468,9 @@ def _supports_routing_method(
             RoutingMethodType.Default,
             RoutingMethodType.Renormalize,
             RoutingMethodType.RenormalizeNaive,
+            RoutingMethodType.DeepSeekV3,
             RoutingMethodType.DeepseekV4,
+            RoutingMethodType.DeepSeekV3,
         ]
 
     @staticmethod
@@ -744,6 +747,7 @@ def prepare_int4_moe_layer_for_cpu(
 
     Returns:
         (blocked_w13, blocked_w2, blocked_s13, blocked_s2, blocked_z13, blocked_z2)
+
     """
     E = w13_packed.size(0)
 
@@ -948,6 +952,29 @@ def apply(
 # ===========================================================================
 
 
+def prepare_int8_moe_layer_for_cpu(
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Prepack INT8 MoE weights for the current CPU architecture."""
+    # SMMLA packing for AArch64
+    if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
+        return (
+            cpu_prepack_moe_weight_int8(w13, "neon"),
+            cpu_prepack_moe_weight_int8(w2, "neon"),
+        )
+    # VSX packing for POWER
+    if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC:
+        return (
+            cpu_prepack_moe_weight_int8(w13, "vsx"),
+            cpu_prepack_moe_weight_int8(w2, "vsx"),
+        )
+    # VNNI packing for x86
+    packed_w13 = torch.ops._C.convert_weight_packed(w13)
+    packed_w2 = torch.ops._C.convert_weight_packed(w2)
+    return packed_w13, packed_w2
+
+
 class CPUExpertsInt8(mk.FusedMoEExpertsModular):
     """CPU INT8 W8A8 per-channel weight / dynamic per-token activation
     modular MoE experts."""
@@ -1054,7 +1081,6 @@ def supports_expert_map(self) -> bool:
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         """VNNI-prepack INT8 MoE weights for CPU kernel."""
-
         w13 = torch.ops._C.convert_weight_packed(layer.w13_weight)
         w2 = torch.ops._C.convert_weight_packed(layer.w2_weight)
         replace_parameter(layer, "w13_weight", w13)
@@ -1282,6 +1308,160 @@ def apply(
         )
 
 
+class PowerCPUExpertsInt8(mk.FusedMoEExpertsModular):
+    """POWER VSX INT8 MoE with per-token activation and channelwise weight quant."""
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    @staticmethod
+    def is_supported_config(
+        cls: type[mk.FusedMoEExperts],
+        moe_config: FusedMoEConfig,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+        activation_format: mk.FusedMoEActivationFormat,
+    ) -> tuple[bool, str | None]:
+        supported, reason = mk.FusedMoEExpertsModular.is_supported_config(
+            cls,
+            moe_config,
+            weight_key,
+            activation_key,
+            activation_format,
+        )
+        if not supported:
+            return supported, reason
+        if moe_config.in_dtype not in (torch.float32, torch.bfloat16):
+            return False, "kernel requires float32 or bfloat16 activations"
+        if moe_config.hidden_dim % 16 != 0:
+            return False, "kernel requires hidden dim divisible by 16"
+        if moe_config.intermediate_size_per_partition % 16 != 0:
+            return False, "kernel requires intermediate dim divisible by 16"
+        return True, None
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        return (
+            current_platform.is_cpu()
+            and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
+            and hasattr(torch.ops._C, "cpu_fused_moe_int8")
+        )
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return False
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation in (
+            MoEActivation.SILU,
+            MoEActivation.SWIGLUOAI,
+            MoEActivation.GELU,
+            MoEActivation.GELU_TANH,
+        )
+
+    @staticmethod
+    def _supports_parallel_config(
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> bool:
+        return not moe_parallel_config.use_ep
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return (weight_key, activation_key) == (
+            kInt8StaticChannelSym,
+            kInt8DynamicTokenSym,
+        )
+
+    @staticmethod
+    def _supports_routing_method(
+        routing_method: RoutingMethodType,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return routing_method in [
+            RoutingMethodType.Default,
+            RoutingMethodType.Renormalize,
+            RoutingMethodType.RenormalizeNaive,
+        ]
+
+    @staticmethod
+    def _supports_router_logits_dtype(
+        router_logits_dtype: torch.dtype | None,
+        routing_method: RoutingMethodType,
+    ) -> bool:
+        return True
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w13 = cpu_prepack_moe_weight_int8(layer.w13_weight, "vsx")
+        w2 = cpu_prepack_moe_weight_int8(layer.w2_weight, "vsx")
+        replace_parameter(layer, "w13_weight", w13)
+        replace_parameter(layer, "w2_weight", w2)
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        # cpu_fused_moe_int8 manages its own scratch space.
+        return (0,), (0,), (M, K)
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        return TopKWeightAndReduceNoOP()
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ) -> None:
+        # The modular prepare step applies router weights to the input when
+        # apply_router_weight_on_input is enabled.
+        assert self.w1_scale is not None
+        assert self.w2_scale is not None
+        cpu_fused_moe_int8(
+            output,
+            hidden_states,
+            w1,
+            w2,
+            self.w1_scale,
+            self.w2_scale,
+            self.w1_bias,
+            self.w2_bias,
+            topk_weights,
+            topk_ids,
+            activation.value,
+            "vsx",
+            skip_weighted=apply_router_weight_on_input,
+        )
+
+
 class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular):
     """AMD Zen INT8 MoE with per-token activation and channelwise weight
     quantization, dispatched through zentorch."""
@@ -1434,3 +1614,218 @@ def apply(
             self._w1_scale_bf16,
             self._w2_scale_bf16,
         )
+
+
+# ===========================================================================
+# FP8 W8A8 MoE
+# ===========================================================================
+
+
+def prepare_fp8_w8a8_moe_layer_for_cpu(
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Prepack FP8 W8A8 MoE weights for CPU kernel."""
+    E = w13.size(0)
+
+    packed_w13_list = []
+    packed_w13_scale_list = []
+
+    for i in range(E):
+        w1_exp = w13[i]  # [2N, K] FP8
+        ws_exp = w13_scale[i]  # [2N, K/128] or [2N_blocks, K/128]
+
+        # Align scale rows to weight rows if necessary (for block-128 checkpoints
+        # where scale may have fewer rows: [2N/128, K/128] → [2N, K/128])
+        if ws_exp.size(0) < w1_exp.size(0):
+            repeat_factor = w1_exp.size(0) // ws_exp.size(0)
+            ws_exp = torch.repeat_interleave(ws_exp, repeat_factor, dim=0)
+        ws_exp = ws_exp[: w1_exp.size(0), :].contiguous()
+
+        pw, ps = torch.ops._C.float8_linear_prepack_cpu(w1_exp, ws_exp)
+        packed_w13_list.append(pw)
+        packed_w13_scale_list.append(ps)
+
+    packed_w13 = torch.stack(packed_w13_list)
+    packed_w13_scale = torch.stack(packed_w13_scale_list)
+
+    # w2 uses the W8A16 path (BF16 activation × FP8 weight), VNNI packed
+    packed_w2 = torch.ops._C.convert_weight_packed(w2)
+    return packed_w13, packed_w13_scale, packed_w2, w2_scale
+
+
+class CPUExpertsFp8W8A8(mk.FusedMoEExpertsModular):
+    """CPU FP8 W8A8 block-quantized modular MoE experts."""
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ):
+        super().__init__(moe_config, quant_config)
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        # Receives BF16 hidden states; quantization to FP8 is done inside apply()
+        return True
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        # Requires native AMX-FP8 hardware.
+        return current_platform.is_cpu() and cpu_has_amx_fp8()
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return False
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation == MoEActivation.SILU
+
+    @staticmethod
+    def _supports_parallel_config(
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> bool:
+        return True
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        # Block-quantized FP8.
+        return (weight_key, activation_key) == (
+            kFp8Static128BlockSym,
+            kFp8Dynamic128Sym,
+        )
+
+    @staticmethod
+    def _supports_routing_method(
+        routing_method: RoutingMethodType,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return routing_method in [
+            RoutingMethodType.Default,
+            RoutingMethodType.Renormalize,
+            RoutingMethodType.RenormalizeNaive,
+            RoutingMethodType.DeepSeekV3,
+            RoutingMethodType.DeepseekV4,
+        ]
+
+    @staticmethod
+    def _supports_router_logits_dtype(
+        router_logits_dtype: torch.dtype | None,
+        routing_method: RoutingMethodType,
+    ) -> bool:
+        return True
+
+    def moe_problem_size(
+        self,
+        a1: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> tuple[int, int, int, int, int]:
+        # float8_linear_prepack_cpu blocks w1 into [E, Nc, Kc, block_k, block_n],
+        # not the (E, N, K) layout the base implementation assumes -- N isn't
+        # recoverable from that shape, so read it from moe_config instead.
+        E = w1.shape[0]
+        K = a1.size(-1)
+        N = (
+            self.moe_config.intermediate_size_per_partition
+            * self.moe_config.w13_num_shards
+        )
+        M = a1.size(0) if a1.dim() == 2 else a1.size(1)
+        topk = topk_ids.size(1)
+        return E, M, N, K, topk
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        # fused_experts_cpu manages its own scratch space.
+        return (0,), (0,), (M, K)
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        return TopKWeightAndReduceNoOP()
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ) -> None:
+        if apply_router_weight_on_input:
+            # fused_experts_cpu always applies topk_weights internally on
+            # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would
+            # also pre-apply it to hidden_states, double-weighting the
+            # output. Not needed by any CPU FP8 model today.
+            raise NotImplementedError(
+                "CPUExpertsFp8W8A8 does not support apply_router_weight_on_input=True."
+            )
+
+        block_shape = (
+            list(self.quant_config.block_shape)
+            if self.quant_config.block_shape
+            else (
+                [
+                    self.quant_config._w1.shape.row,
+                    self.quant_config._w1.shape.col,
+                ]
+                if self.quant_config._w1.shape is not None
+                else None
+            )
+        )
+
+        # Quantize hidden_states (BF16) → FP8 with per-token scales
+        hidden_states_2d = hidden_states.reshape(-1, hidden_states.shape[-1])
+        x_fp8, x_scales = torch.ops._C.quantize_fp8e4m3_vec(
+            hidden_states_2d, True, None
+        )
+
+        fused_experts_cpu(
+            output,
+            x_fp8,
+            w1,
+            w2,
+            topk_weights,
+            topk_ids,
+            CPUQuantMethod.FP8_W8A8,
+            self.w1_scale,  # w1_scale
+            self.w2_scale,  # w2_scale
+            None,  # w1_zero
+            None,  # w2_zero
+            block_shape,  # block_size
+            None,  # w1_bias
+            None,  # w2_bias
+            None,  # alpha
+            None,  # limit
+            True,  # is_vnni
+            x_scales,  # a1_scale: per-token FP8 activation scales
+        )
diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
index 402077b017f8..cd7da3e1e3a0 100644
--- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
@@ -518,8 +518,7 @@ def run_cutlass_moe_fp4(
     *,
     activation_config: ApplyMoEActivationConfig | None = None,
 ) -> None:
-    """
-    MoE implementation for FP4 Inputs
+    """MoE implementation for FP4 Inputs.
 
     # Gemm 1
     a: Input tensor: [m, k] (half/bfloat16)
diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
index 211364d24bca..699b029bad86 100644
--- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
@@ -50,6 +50,21 @@
 logger = init_logger(__name__)
 
 
+def _fp8_workspace_shape(
+    num_rows: int, num_columns: int, workspace_dtype: torch.dtype
+) -> tuple[int, int]:
+    """Size an FP8 byte buffer stored in a workspace of another dtype."""
+    bytes_per_workspace_element = workspace_dtype.itemsize
+    fp8_columns_per_workspace_element = (
+        bytes_per_workspace_element // torch.float8_e4m3fn.itemsize
+    )
+    assert bytes_per_workspace_element % torch.float8_e4m3fn.itemsize == 0
+    return (
+        num_rows,
+        -(-num_columns // fp8_columns_per_workspace_element),
+    )
+
+
 def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool:
     align = get_mk_alignment_for_contiguous_layout()[0]
     return align <= M and N % align == 0 and K % align == 0
@@ -58,8 +73,7 @@ def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool:
 def _valid_deep_gemm(
     hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor
 ) -> bool:
-    """
-    Check if the given problem size is supported by the DeepGemm grouped
+    """Check if the given problem size is supported by the DeepGemm grouped
     gemm kernel.  All of M, N, K and the quantization block_shape must be
     aligned by `dg.get_m_alignment_for_contiguous_layout()`.
     """
@@ -217,7 +231,14 @@ def workspace_shapes(
         assert M_sum % align_used == 0
 
         activation_out_dim = self.adjust_N_for_activation(N, activation)
-        workspace1 = (M_sum, max(activation_out_dim, K))
+        # workspace1 is allocated in the activation dtype by the workspace
+        # manager, but is only ever viewed and used as FP8 in apply(). Size it
+        # by bytes instead of reserving one BF16/FP16 element per FP8 element.
+        workspace1 = _fp8_workspace_shape(
+            M_sum,
+            max(activation_out_dim, K),
+            self.workspace_dtype(self.moe_config.in_dtype),
+        )
         workspace2 = (M_sum, max(N, K))
         output = (M, K)
         return (workspace1, workspace2, output)
@@ -481,7 +502,14 @@ def workspace_shapes(
         assert M_sum % align_used == 0
 
         activation_out_dim = self.adjust_N_for_activation(N, activation)
-        workspace1 = (M_sum, max(activation_out_dim, K))
+        # workspace1 holds the permuted and requantized FP8 activations. The
+        # workspace manager allocates it in the model activation dtype, so
+        # account for the dtype sizes rather than overallocating BF16/FP16.
+        workspace1 = _fp8_workspace_shape(
+            M_sum,
+            max(activation_out_dim, K),
+            self.workspace_dtype(self.moe_config.in_dtype),
+        )
         workspace2 = (M_sum, max(N, K))
         output = (M, K)
         return (workspace1, workspace2, output)
diff --git a/vllm/model_executor/layers/fused_moe/experts/fallback.py b/vllm/model_executor/layers/fused_moe/experts/fallback.py
index 639b2bf26681..e10d13f0a778 100644
--- a/vllm/model_executor/layers/fused_moe/experts/fallback.py
+++ b/vllm/model_executor/layers/fused_moe/experts/fallback.py
@@ -30,8 +30,7 @@ def get_clses() -> tuple[
         type[mk.FusedMoEExpertsModular],
         type[mk.FusedMoEExpertsModular],
     ]:
-        """
-        Get the cls for the experts and fallback experts.
+        """Get the cls for the experts and fallback experts.
 
         Subclasses should implement this method, so that
         we have a consistent way to call the _supports_*
diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py
index d269c6f10993..f7e33711667a 100644
--- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py
@@ -107,8 +107,7 @@ def workspace_shapes(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         activation: MoEActivation,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        """
-        Compute the shapes for the temporary and final outputs of the two gemms
+        """Compute the shapes for the temporary and final outputs of the two gemms
         and activation in the fused expert function.  Since the gemms are
         independent, the workspace for the first gemm can be shared with the
         workspace for the last gemm.
@@ -123,7 +122,6 @@ def workspace_shapes(
         - Note: in order for activation chunking to work, the first dimension
           of each tuple must be the number of tokens.
         """
-
         # We use global_num_experts due to how moe_align_block_size handles
         # expert_maps.
         K_dim = K * 2 if self.use_deep_ep_ll_nvfp4_dispatch else K
@@ -212,8 +210,7 @@ def flashinfer_cutedsl_moe_masked(
     workspace: torch.Tensor,
     out: torch.Tensor,
 ):
-    """
-    Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL
+    """Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL
     kernels.
 
     Args:
@@ -229,13 +226,14 @@ def flashinfer_cutedsl_moe_masked(
         a2_global_scale (torch.Tensor): (l,)
         w2_blockscale (torch.Tensor): blockscale factors, e4m3,
         w2_alpha (torch.Tensor): (l,)
+        out (torch.Tensor): output tensor written in place, [l, m, k], bf16
         masked_m (torch.Tensor): Masked dimension indices
         workspace (torch.Tensor): For gateup_output
 
     Notes:
         - Assumes max(masked_m) <= m.
-    """
 
+    """
     # === Assertions on dtypes ===
     assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}"
     assert w1_blockscale.dtype == torch.float8_e4m3fn, (
diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
index ec0e15bde558..d525822b912b 100644
--- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
@@ -29,8 +29,7 @@
 
 
 class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
-    """
-    CuteDSL NvFP4 MoE experts using the FlashInfer functional API.
+    """CuteDSL NvFP4 MoE experts using the FlashInfer functional API.
 
     Uses Standard activation format (non-batched). The kernel handles
     routing, expert computation, and reduction internally.
diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
index 1614e41b3303..718e018c47e6 100644
--- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
@@ -38,8 +38,7 @@
 def is_valid_flashinfer_cutlass_fused_moe(
     hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor
 ) -> bool:
-    """
-    Check if the given problem size is supported by the FlashInfer CUTLASS MoE
+    """Check if the given problem size is supported by the FlashInfer CUTLASS MoE
     kernel.
     """
     if not has_flashinfer_cutlass_fused_moe():
@@ -219,8 +218,7 @@ def workspace_shapes(
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
         # We use global_num_experts due to how moe_align_block_size handles
         # expert_maps.
-        """
-        Compute the shapes for the temporary and final outputs of the two gemms
+        """Compute the shapes for the temporary and final outputs of the two gemms
         and activation in the fused expert function.  Since the gemms are
         independent, the workspace for the first gemm can be shared with the
         workspace for the last gemm.
diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
index f16b351d353d..b8cd58635eb6 100644
--- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
@@ -551,8 +551,7 @@ def invoke_moe_batched_triton_kernel(
 
 
 class NaiveBatchedExperts(mk.FusedMoEExpertsModular):
-    """
-    A reference MoE expert class that operates on expert batched format,
+    """A reference MoE expert class that operates on expert batched format,
     i.e. E x max_num_tokens x K.  This is the format that the batched
     dispatch/combine kernels use.
     """
@@ -775,8 +774,7 @@ def batched_moe_kernel_quantize_input(
 
 
 class BatchedTritonExperts(mk.FusedMoEExpertsModular):
-    """
-    A Triton based MoE expert class that operates on expert batched format,
+    """A Triton based MoE expert class that operates on expert batched format,
     i.e. E x max_num_tokens x K.  This is the format that the batched
     dispatch/combine kernels use.
     """
diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
index 86079605a367..a6d44c3141f7 100644
--- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
@@ -4,6 +4,7 @@
 
 import json
 import math
+from dataclasses import replace
 from typing import TYPE_CHECKING, Any, cast
 
 import torch
@@ -27,8 +28,10 @@
 from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import moe_fused_mul_sum
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
     MoEPermuteScratch,
+    get_moe_permute_scratch,
     moe_permute,
     moe_permute_unpermute_supported,
+    moe_prepare_scatter,
     moe_unpermute,
 )
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -56,14 +59,16 @@
     kMxfp8Static,
     kNvfp4Dynamic,
     kNvfp4Static,
+    kStaticTensorScale,
 )
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType
 from vllm.utils.import_utils import has_humming
+from vllm.utils.platform_utils import num_compute_units
 from vllm.v1.worker.workspace import current_workspace_manager
 
 if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization.utils.humming_utils import (
+    from vllm.model_executor.layers.quantization.utils.humming import (
         HummingMoEQuantConfig,
     )
     from vllm.utils.humming import (
@@ -117,38 +122,9 @@ def get_humming_moe_gemm_type(moe_config: FusedMoEConfig | None = None) -> str:
     return gemm_type
 
 
-def _fixup_moe_tuning_config(tuning_config: list, max_k_block: int = 128) -> None:
-    """Fix up each MoE tile in place: cap the K-block and widen warp-N.
-
-    - block_shape[2] (K-block) > ``max_k_block``: the driver rejects the TMA
-      descriptor at launch (CUDA_ERROR_MISALIGNED_ADDRESS). Cap at 128, which
-      Humming already uses for larger M.
-    - warp_shape[1] (warp-N) < 32: block-FP8 (group-128) activations route the
-      w13 (gate/up) GEMM to a tuning table that pins warp-N to 16, under-filling
-      the Hopper WGMMA N dimension and corrupting the GEMM output (gsm8k
-      0.94 -> 0.89). Widen to 32 whenever block_n % 32 == 0 -- on every tile,
-      not just the K-capped ones. w2 (down) already uses warp_n=32.
-    """
-    logger.info_once("Attempting to override humming GEMM config")
-    for entry in tuning_config:
-        config = entry[2]
-        block_shape = config.get("block_shape")
-        if not (block_shape and len(block_shape) == 3):
-            continue
-        block_m, block_n, block_k = block_shape
-
-        logger.info_once(f"Overriding humming GEMM config. Previous config\n: {config}")
-        if block_k > max_k_block:
-            config["block_shape"] = [block_m, block_n, max_k_block]
-
-        warp_shape = config.get("warp_shape")
-        if warp_shape and warp_shape[1] < 32 and block_n % 32 == 0:
-            config["warp_shape"] = [warp_shape[0], 32, warp_shape[2]]
-
-        logger.info_once(f"Overridden humming GEMM config. Current config\n: {config}")
-
-
 class HummingExpertsBase(mk.FusedMoEExpertsModular):
+    quant_config: "HummingMoEQuantConfig"
+
     def __init__(
         self,
         moe_config: FusedMoEConfig,
@@ -164,7 +140,7 @@ def __init__(
         self.locks = torch.zeros(1024, dtype=torch.int32, device=moe_config.device)
         self.num_experts = moe_config.num_local_experts
         self.global_num_experts = moe_config.num_experts
-        self.quant_config = quant_config
+        self.quant_config = humming_quant_config
         self.init_humming_moe()
 
         if self.is_batched():
@@ -176,7 +152,6 @@ def __init__(
             max_num_tokens=max_num_tokens,
             num_dispatchers=num_dispatchers,
         )
-        self._permute_scratch: dict[int, MoEPermuteScratch] = {}
 
     def init_humming_moe(self):
         from vllm.utils.humming import get_heuristics_config
@@ -199,31 +174,61 @@ def init_humming_moe(self):
             gemm_type=self.humming_gemm_type(),
         )
         self.compute_config_str = json.dumps(self.compute_config)
-        # Fix up the heuristic-chosen tiles before freezing to JSON -- see
-        # _fixup_moe_tuning_config.
-        _fixup_moe_tuning_config(self.w13_tuning_config)
-        _fixup_moe_tuning_config(self.w2_tuning_config)
         self.w13_tuning_config_str = json.dumps(self.w13_tuning_config)
         self.w2_tuning_config_str = json.dumps(self.w2_tuning_config)
 
-    def quantize_input(
+    def process_input(
         self,
         sublayer_name: str,
         inputs: torch.Tensor,
         quanted_input: torch.Tensor | None,
         input_scale: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        from vllm.utils.humming import may_quant_input
+        expert_tokens: torch.Tensor | None = None,
+        activation: MoEActivation | None = None,
+        scatter_idx: torch.Tensor | None = None,
+        num_valid_tokens: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        from vllm.model_executor.layers.quantization.utils.humming.activation import (
+            get_humming_activation,
+        )
+        from vllm.utils.humming import may_process_input
 
-        # input_scale is set for block-FP8 (group-128) activations that were
-        # quantized before the EP dispatch: may_quant_input then skips the
-        # redundant w13 quantization and forwards the pre-computed scale.
-        return may_quant_input(
-            self.humming_configs[sublayer_name],
+        quant_config = self.quant_config
+        config = self.humming_configs[sublayer_name]
+        if input_scale is not None:
+            assert scatter_idx is None
+            if activation is not None:
+                raise ValueError("Cannot apply activation to prequantized input")
+            return inputs, input_scale, None
+
+        activation_kwargs = {}
+        if activation is not None:
+            activation_config = self.activation_config
+            activation_kwargs = get_humming_activation(activation, activation_config)
+
+        prefix = "w1" if sublayer_name == "w13" else "w2"
+        scale = getattr(quant_config, f"{prefix}_input_scale")
+        scale2 = getattr(quant_config, f"{prefix}_input_scale_2")
+        mode = config.input_quant_mode
+        layout = "grouped_mask" if self.is_batched() else "normal"
+        if scatter_idx is not None:
+            layout = "scatter"
+
+        inputs, group_scales, token_scales = may_process_input(
+            config,
             inputs=inputs,
-            input_scale=input_scale,
-            quanted_input=quanted_input,
+            outputs=quanted_input,
+            token_scales=scale2 if mode.has_secondary_scale else scale,
+            hadamard_block_size=getattr(quant_config, f"{prefix}_hadamard_block_size"),
+            layout=layout,
+            expert_tokens=expert_tokens if layout == "grouped_mask" else None,
+            scatter_idx=scatter_idx,
+            num_valid_tokens=num_valid_tokens,
+            **activation_kwargs,
         )
+        input_scale = group_scales if mode.has_group_scale else token_scales
+        input_scale_2 = token_scales if mode.has_secondary_scale else None
+        return inputs, input_scale, input_scale_2
 
     def humming_forward(
         self,
@@ -232,51 +237,47 @@ def humming_forward(
         weight: torch.Tensor,
         input_scale: torch.Tensor | None,
         outputs: torch.Tensor,
+        input_scale_2: torch.Tensor | None = None,
         **kwargs: Any,
     ) -> torch.Tensor:
         from vllm.utils.humming import humming_forward
 
-        is_w13 = sublayer_name == "w13"
+        index = 1 if sublayer_name == "w13" else 2
         return humming_forward(
             self.humming_configs[sublayer_name],
             inputs=inputs,
             weight=weight,
-            weight_scale=(
-                self.quant_config.w1_scale if is_w13 else self.quant_config.w2_scale
-            ),
-            zero_point=(self.quant_config.w1_zp if is_w13 else self.quant_config.w2_zp),
-            bias=(self.quant_config.w1_bias if is_w13 else self.quant_config.w2_bias),
-            weight_scale_2=(
-                self.quant_config.g1_alphas if is_w13 else self.quant_config.g2_alphas
-            ),
+            weight_scale=getattr(self.quant_config, f"w{index}_scale"),
+            zero_point=getattr(self.quant_config, f"w{index}_zp"),
+            bias=getattr(self.quant_config, f"w{index}_bias"),
+            weight_scale_2=getattr(self.quant_config, f"g{index}_alphas"),
             input_scale=input_scale,
+            input_scale_2=input_scale_2,
             outputs=outputs,
             locks=self.locks,
             **kwargs,
         )
 
-    def _get_permute_scratch(self, topk: int) -> MoEPermuteScratch | None:
+    def _get_permute_scratch(
+        self, topk: int, indices_only: bool = False
+    ) -> MoEPermuteScratch | None:
         if not moe_permute_unpermute_supported():
             return None
 
-        scratch = self._permute_scratch.get(topk)
-        if scratch is None:
-            max_expanded_rows = (
-                self.moe_config.max_num_tokens
-                * self.moe_config.dp_size
-                * self.moe_config.experts_per_token
-            )
-            scratch = MoEPermuteScratch(
-                max_num_tokens=math.ceil(max_expanded_rows / topk),
-                topk=topk,
-                num_experts=self.moe_config.num_experts,
-                num_local_experts=self.moe_config.num_local_experts,
-                device=torch.device(self.moe_config.device),
-                hidden_size=self.moe_config.hidden_dim,
-                hidden_dtype=self.moe_config.in_dtype,
-            )
-            self._permute_scratch[topk] = scratch
-        return scratch
+        max_expanded_rows = (
+            self.moe_config.max_num_tokens
+            * self.moe_config.dp_size
+            * self.moe_config.experts_per_token
+        )
+        return get_moe_permute_scratch(
+            max_num_tokens=math.ceil(max_expanded_rows / topk),
+            topk=topk,
+            num_experts=self.moe_config.num_experts,
+            num_local_experts=self.moe_config.num_local_experts,
+            device=torch.device(self.moe_config.device),
+            hidden_size=None if indices_only else self.moe_config.hidden_dim,
+            hidden_dtype=None if indices_only else self.moe_config.in_dtype,
+        )
 
     def get_global_valid_shape_m(self, topk_ids: torch.Tensor):
         ctx = get_forward_context()
@@ -306,6 +307,16 @@ def _supports_quant_scheme(
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
+        if (
+            activation_key is not None
+            and activation_key.scale == kNvfp4Dynamic.scale
+            and activation_key.dtype == kNvfp4Dynamic.dtype
+            and activation_key.scale2 is not None
+            and not activation_key.scale2.static
+            and activation_key.scale2.group_shape == GroupShape.PER_TOKEN
+            and activation_key.scale2.dtype == torch.float32
+        ):
+            activation_key = replace(activation_key, scale2=kStaticTensorScale)
         SUPPORTED_W_A = [
             (kMxfp4Static, None),
             (kMxfp4Static, kMxfp4Dynamic),
@@ -347,20 +358,19 @@ def _supports_quant_scheme(
             (kMxfp8Static, kMxfp8Dynamic),
         ]
         return (weight_key, activation_key) in SUPPORTED_W_A or (
-            activation_key in (None, kFp8DynamicTokenSym)
+            activation_key in (None, kFp8DynamicTokenSym, kInt8DynamicTokenSym)
             and _is_supported_wna16_weight_key(weight_key)
         )
 
     def _prequantizes_dispatch_activation(self) -> bool:
-        """
-        Whether the prepare/finalize step should quantize activations before
+        """Whether the prepare/finalize step should quantize activations before
         the (EP all-to-all) dispatch instead of leaving it to Humming.
 
         This is enabled only for block-FP8 (group-128) activations: quantizing
         to FP8 before dispatch sends FP8 rather than BF16 over the interconnect,
         and Humming then consumes the pre-quantized FP8 + scale directly (see
         the apply() methods, which forward the dispatch scale into
-        HummingMethod.may_quant_input, and may_quant_input itself, which is a
+        HummingExpertsBase.process_input, which is a
         no-op when an input scale is already supplied). The scale layout
         produced by vLLM's block-FP8 quantization ([M, K // 128] float32,
         row-major) matches what the Humming WGMMA grouped GEMM expects.
@@ -373,13 +383,12 @@ def _prequantizes_dispatch_activation(self) -> bool:
 
     @property
     def expects_unquantized_inputs(self) -> bool:
-        """
-        Whether the prepare/finalize step should defer input quantization to
+        """Whether the prepare/finalize step should defer input quantization to
         the experts (by setting defer_input_quant=True and passing unquantized
         inputs).
 
         Humming normally quantizes inputs internally via
-        HummingMethod.may_quant_input() in apply(), so we defer quantization
+        HummingExpertsBase.process_input() in apply(), so we defer quantization
         (return True) to avoid quantizing twice -- once in prepare and once in
         Humming's apply().
 
@@ -404,8 +413,6 @@ def _supports_no_act_and_mul() -> bool:
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        # Humming uses apply_moe_activation() callback for activation,
-        # so any activation supported there can be used here.
         return apply_moe_activation_supported(activation)
 
     @staticmethod
@@ -462,10 +469,9 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
         assert intermediate_dim == self.adjust_N_for_activation(gate_up_dim, activation)
 
         # hidden_states
-        # (-> quanted_gate_up_input) (if not BF16/FP16 activation)
+        # -> quanted_gate_up_input (Hadamard + optional quantization)
         # -> gate_up_output
-        # -> activation_output
-        # (-> quanted_down_input) (if not BF16/FP16 activation)
+        # -> quanted_down_input (activation + Hadamard + optional quantization)
         # -> down_output
         # (-> output) (if not is_batched)
         # Neighboring nodes are required to utilize distinct workspaces.
@@ -477,8 +483,8 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
             max_num_tokens = self.max_num_tokens
             num_dispatchers = self.num_dispatchers
             assert max_num_tokens is not None and num_dispatchers is not None
-            input_shape_m = num_experts * max_num_tokens
             real_shape_m = num_experts * max_num_tokens * num_dispatchers
+            input_shape_m = real_shape_m
             output_shape = (num_experts, max_num_tokens * num_dispatchers, K)
         else:
             input_shape_m = M
@@ -489,15 +495,18 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
 
         a_dtype = self.humming_configs["w13"].a_dtype
         c_dtype = self.humming_configs["w13"].c_dtype
-        num_bits = a_dtype.num_bits
+        down_a_dtype = w2_config.a_dtype
         torch_dtype_map = {
             dtypes.float16: torch.float16,
             dtypes.bfloat16: torch.bfloat16,
             dtypes.float32: torch.float32,
+            dtypes.float8e3m4: torch.uint8,
             dtypes.float8e4m3: torch.float8_e4m3fn,
             dtypes.float8e5m2: torch.float8_e5m2,
             dtypes.int8: torch.int8,
             dtypes.int4: torch.uint8,
+            dtypes.float4e0m3: torch.uint8,
+            dtypes.float4e2m1: torch.uint8,
         }
 
         buffer_metas = {
@@ -509,13 +518,9 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
                 "shape": (real_shape_m, gate_up_dim),
                 "dtype": torch_dtype_map[c_dtype],
             },
-            "activation_output": {
-                "shape": (real_shape_m, intermediate_dim),
-                "dtype": torch_dtype_map[c_dtype],
-            },
             "quanted_down_input": {
                 "shape": (real_shape_m, intermediate_dim),
-                "dtype": torch_dtype_map[a_dtype],
+                "dtype": torch_dtype_map[down_a_dtype],
             },
             "down_output": {
                 "shape": output_shape if self.is_batched() else (real_shape_m, K),
@@ -529,7 +534,8 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
 
         for key in buffer_metas:
             meta = buffer_metas[key]
-            if "quanted" in key and a_dtype.num_bits == 4:
+            input_dtype = down_a_dtype if key == "quanted_down_input" else a_dtype
+            if "quanted" in key and input_dtype.num_bits == 4:
                 last_dim = meta["shape"][-1]
                 if last_dim % 2 != 0:
                     raise ValueError(
@@ -538,25 +544,12 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation):
                     )
                 meta["shape"] = meta["shape"][:-1] + (last_dim // 2,)
 
-        if num_bits == 16:
-            required_buffers = ["gate_up_output", "activation_output", "down_output"]
-        else:
-            required_buffers = [
-                "quanted_gate_up_input",
-                "gate_up_output",
-                "activation_output",
-                "quanted_down_input",
-                "down_output",
-            ]
-
-            # Note: The fused SITU+FP8 quant goes gate_up_output ->
-            #        quanted_down_input directly, never materializing
-            #        activation_output. Dropping activation_output from the chain
-            #        flips the even/odd workspace 2-coloring below so
-            #        gate_up_output and quanted_down_input land
-            #        on DIFFERENT workspaces.
-            if self.fused_situ_quant_enabled(activation):
-                required_buffers.remove("activation_output")
+        required_buffers = [
+            "quanted_gate_up_input",
+            "gate_up_output",
+            "quanted_down_input",
+            "down_output",
+        ]
 
         # batched moe use down_output as output
         if not self.is_batched():
@@ -679,81 +672,6 @@ def apply_activation(
             activation_kwargs["valid_token_counts"] = valid_token_counts
         self.activation(**activation_kwargs)
 
-    def fused_situ_quant_enabled(self, activation: MoEActivation) -> bool:
-        """Whether the SITU activation + w2 quant can be fused into one kernel.
-
-        Fused only for k-major block-FP8 group-128 e4m3 with float32 scales --
-        the sole layout situ_and_mul_quant supports. A float32 as_dtype rules out
-        MXMMA (which uses e8m0 m-major scales), so a group-128 scale is
-        guaranteed k-major. Everything else (m-major MXFP8, other group sizes,
-        16-bit passthrough, non-SITU) falls back to the separate
-        situ_and_mul + quantize_input path.
-        """
-        if activation != MoEActivation.SITU:
-            return False
-        w2cfg = self.humming_configs["w2"]
-        # Report the specific blocker once so a silent fallback to the unfused
-        # situ_and_mul + quant pair is never a mystery in a trace.
-        reason: str | None = None
-        if not (w2cfg.a_dtype.num_bits == 8 and str(w2cfg.a_dtype) == "float8e4m3"):
-            reason = f"w2 a_dtype is {w2cfg.a_dtype} (need float8e4m3)"
-        elif w2cfg.input_scale_group_size != 128:
-            reason = (
-                f"w2 input_scale_group_size is {w2cfg.input_scale_group_size} "
-                "(need 128)"
-            )
-        elif not (w2cfg.as_dtype is not None and str(w2cfg.as_dtype) == "float32"):
-            reason = f"w2 as_dtype is {w2cfg.as_dtype} (need float32, k-major)"
-        if reason is not None:
-            logger.warning_once(
-                "Humming fused SITU+FP8 quant disabled, using unfused "
-                "situ_and_mul + quant: %s",
-                reason,
-            )
-            return False
-        return True
-
-    def fused_situ_quant(
-        self,
-        gate_up_output: torch.Tensor,
-        quanted_down_input: torch.Tensor,
-        num_valid_tokens: torch.Tensor | None,
-        topk: int = 1,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Run the fused SITU activation + FP8 quant for the w2 input.
-
-        Returns ``(quanted_down_input, input_scale)`` in the same layout the
-        unfused ``apply_activation`` + ``quantize_input("w2")`` pair produces: an
-        fp8 [M, d] tensor plus a k-major block-FP8 float32 [M, d // 128] group
-        scale. Only reached when ``fused_situ_quant_enabled`` (group_size 128).
-        """
-        from vllm.model_executor.layers.fused_moe.activation import (
-            situ_and_mul_quant,
-        )
-
-        cfg = self.activation_config
-        assert cfg.activation_situ_beta is not None, (
-            "SITU requires activation_situ_beta from FusedMoEConfig"
-        )
-        group_size = self.humming_configs["w2"].input_scale_group_size
-        m, d = quanted_down_input.size(0), quanted_down_input.size(1)
-        input_scale = torch.empty(
-            (m, d // group_size),
-            dtype=torch.float32,
-            device=quanted_down_input.device,
-        )
-        situ_and_mul_quant(
-            quanted_down_input,
-            input_scale,
-            gate_up_output,
-            beta=cfg.activation_situ_beta,
-            linear_beta=cfg.activation_situ_linear_beta,
-            group_size=group_size,
-            num_valid_tokens=num_valid_tokens,
-            topk=topk,
-        )
-        return quanted_down_input, input_scale
-
 
 class HummingIndexedExperts(HummingExpertsBase):
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
@@ -774,7 +692,7 @@ def prepare_humming_moe_kwargs(
         topk_ids: torch.Tensor,
         expert_map: torch.Tensor | None,
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
-    ) -> tuple[dict[str, Any], dict[str, Any]]:
+    ) -> tuple[dict[str, Any], dict[str, Any], torch.Tensor | None]:
         valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)
 
         moe_block_size = None
@@ -790,13 +708,21 @@ def prepare_humming_moe_kwargs(
             )
             moe_block_size = 64
 
-        sorted_ids, expert_ids, num_tokens_padded = moe_align_block_size(
+        use_scatter = False
+        if expert_map is not None:
+            num_sms = num_compute_units(topk_ids.get_device())
+            use_scatter = topk_ids.numel() > 2 * num_sms
+
+        alignment = moe_align_block_size(
             topk_ids=topk_ids,
             block_size=moe_block_size,
             num_experts=self.global_num_experts,
             expert_map=expert_map,
             ignore_invalid_experts=True,
+            return_scatter_idx=use_scatter,
         )
+        sorted_ids, expert_ids, num_tokens_padded = alignment[:3]
+        scatter_idx = alignment[3] if len(alignment) == 4 else None
 
         moe_common_kwargs = {
             "sorted_ids": sorted_ids,
@@ -812,7 +738,26 @@ def prepare_humming_moe_kwargs(
         moe_kwargs1.update(moe_common_kwargs)
         moe_kwargs2.update(moe_common_kwargs)
 
-        return moe_kwargs1, moe_kwargs2
+        w2_block_size = 64
+        for lower, upper, config in self.w2_tuning_config:
+            if lower < valid_shape_m <= upper:
+                w2_block_size = config["block_shape"][0]
+                break
+        if w2_block_size != moe_block_size:
+            sorted_ids2, expert_ids2, num_tokens_padded2 = moe_align_block_size(
+                topk_ids=topk_ids,
+                block_size=w2_block_size,
+                num_experts=self.global_num_experts,
+                expert_map=expert_map,
+                ignore_invalid_experts=True,
+            )
+            moe_kwargs2.update(
+                sorted_ids=sorted_ids2,
+                expert_ids=expert_ids2,
+                num_tokens_padded=num_tokens_padded2,
+            )
+
+        return moe_kwargs1, moe_kwargs2, scatter_idx
 
     def apply(
         self,
@@ -832,14 +777,13 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ) -> None:
-        """
-        Standard apply implementation for Humming indexed experts.
+        """Standard apply implementation for Humming indexed experts.
 
         Note: Humming kernels handle weights internally through the layer
         object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
         path (Humming quantizes the w13 input itself); for block-FP8 activations
         it carries the scale computed before dispatch, which is forwarded to
-        may_quant_input so Humming skips the redundant w13 quantization. The
+        process_input so Humming skips the redundant w13 quantization. The
         output is written into workspace13 via the buffer management.
         """
         assert not apply_router_weight_on_input
@@ -854,13 +798,13 @@ def apply(
         )
         buffers["output"] = output
 
-        moe_kwargs1, moe_kwargs2 = self.prepare_humming_moe_kwargs(
+        moe_kwargs1, moe_kwargs2, scatter_idx = self.prepare_humming_moe_kwargs(
             topk_ids=topk_ids,
             expert_map=expert_map,
             expert_tokens_meta=expert_tokens_meta,
         )
 
-        inputs, input_scale = self.quantize_input(
+        inputs, input_scale, input_scale_2 = self.process_input(
             "w13",
             inputs=hidden_states,
             input_scale=a1q_scale,
@@ -872,52 +816,25 @@ def apply(
             inputs=inputs,
             weight=w1,
             input_scale=input_scale,
+            input_scale_2=input_scale_2,
             outputs=buffers["gate_up_output"],
             **moe_kwargs1,
         )
 
-        # psum[-1:] is the DeepEP valid *token* count as a zero-cost int32 view.
-        # The fused path and mul_sum consume tokens directly. The unfused
-        # activation consumes the expanded row count (tokens * topk).
-        valid_tokens = None
-        topk = topk_ids.size(1)
-        if (
-            expert_tokens_meta is not None
-            and expert_tokens_meta.psum_recv_per_rank is not None
-        ):
-            valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:]
-
-        if self.fused_situ_quant_enabled(activation):
-            # Fused SITU + FP8 quant (per-token or block-FP8 group-128) straight
-            # into the w2 input, skipping the bf16 activation_output round-trip.
-            inputs, input_scale = self.fused_situ_quant(
-                gate_up_output=buffers["gate_up_output"],
-                quanted_down_input=buffers["quanted_down_input"],
-                num_valid_tokens=valid_tokens,
-                topk=topk,
-            )
-        else:
-            valid_token_counts = (
-                valid_tokens * topk if valid_tokens is not None else None
-            )
-            self.apply_activation(
-                activation=activation,
-                input=buffers["gate_up_output"],
-                output=buffers["activation_output"],
-                valid_token_counts=valid_token_counts,
-            )
-
-            inputs, input_scale = self.quantize_input(
-                "w2",
-                inputs=buffers["activation_output"],
-                quanted_input=buffers.get("quanted_down_input", None),
-            )
+        inputs, input_scale, input_scale_2 = self.process_input(
+            "w2",
+            inputs=buffers["gate_up_output"],
+            quanted_input=buffers["quanted_down_input"],
+            activation=activation,
+            scatter_idx=scatter_idx,
+        )
 
         self.humming_forward(
             "w2",
             inputs=inputs,
             weight=w2,
             input_scale=input_scale,
+            input_scale_2=input_scale_2,
             outputs=buffers["down_output"].view(-1, hidden_states.size(-1)),
             **moe_kwargs2,
         )
@@ -925,6 +842,10 @@ def apply(
         # expert_map masks any non-local id; num_valid_tokens bounds the
         # persistent kernel to the real token rows [0, num_recv) so the padding
         # tail is never iterated (CUDA-graph-safe device scalar).
+        valid_tokens = None
+        if expert_tokens_meta and expert_tokens_meta.psum_recv_per_rank is not None:
+            valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:]
+
         moe_fused_mul_sum(
             inputs=buffers["down_output"].view(*topk_ids.shape, -1),
             topk_weights=topk_weights,
@@ -967,8 +888,7 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ) -> None:
-        """
-        Standard apply implementation for Humming grouped experts.
+        """Standard apply implementation for Humming grouped experts.
 
         Note: Humming kernels handle weights internally through the layer
         object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
@@ -991,21 +911,33 @@ def apply(
         )
         buffers["output"] = output
 
-        hidden_states, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute(
-            hidden_states=hidden_states,
-            a1q_scale=a1q_scale,
-            topk_ids=topk_ids,
-            n_expert=global_num_experts,
-            n_local_expert=self.num_experts,
-            expert_map=expert_map,
-            scratch=self._get_permute_scratch(topk_ids.size(1)),
-        )
+        if a1q_scale is None:
+            scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True)
+            assert scratch is not None
+            scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch)
+            expert_offsets, scatter_idx = scatter_metadata
+            inv_perm = scatter_idx.flatten()
+            num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None
+        else:
+            hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(
+                hidden_states=hidden_states,
+                a1q_scale=a1q_scale,
+                topk_ids=topk_ids,
+                n_expert=global_num_experts,
+                n_local_expert=self.num_experts,
+                expert_map=expert_map,
+                scratch=self._get_permute_scratch(topk_ids.size(1)),
+            )
+            scatter_idx = None
+            num_valid_tokens = None
 
-        inputs, input_scale = self.quantize_input(
+        inputs, input_scale, input_scale_2 = self.process_input(
             "w13",
             inputs=hidden_states,
             input_scale=a1q_scale,
             quanted_input=buffers.get("quanted_gate_up_input", None),
+            scatter_idx=scatter_idx,
+            num_valid_tokens=num_valid_tokens,
         )
 
         self.humming_forward(
@@ -1013,28 +945,20 @@ def apply(
             inputs=inputs,
             weight=w1,
             input_scale=input_scale,
+            input_scale_2=input_scale_2,
             outputs=buffers["gate_up_output"],
             valid_shape_m=valid_shape_m,
-            expert_layout=expert_first_token_offset,
+            expert_layout=expert_offsets,
             compute_config=self.compute_config_str,
             tuning_config=self.w13_tuning_config_str,
         )
 
-        self.apply_activation(
-            activation=activation,
-            input=buffers["gate_up_output"],
-            output=buffers["activation_output"],
-            valid_token_counts=(
-                expert_first_token_offset[-1:].to(torch.int32)
-                if expert_tokens_meta is None
-                else None
-            ),
-        )
-
-        inputs, input_scale = self.quantize_input(
+        inputs, input_scale, input_scale_2 = self.process_input(
             "w2",
-            inputs=buffers["activation_output"],
-            quanted_input=buffers.get("quanted_down_input", None),
+            inputs=buffers["gate_up_output"],
+            quanted_input=buffers["quanted_down_input"],
+            activation=activation,
+            num_valid_tokens=expert_offsets[-1:],
         )
 
         self.humming_forward(
@@ -1042,9 +966,10 @@ def apply(
             inputs=inputs,
             weight=w2,
             input_scale=input_scale,
+            input_scale_2=input_scale_2,
             outputs=buffers["down_output"],
             valid_shape_m=valid_shape_m,
-            expert_layout=expert_first_token_offset,
+            expert_layout=expert_offsets,
             compute_config=self.compute_config_str,
             tuning_config=self.w2_tuning_config_str,
         )
@@ -1054,7 +979,7 @@ def apply(
             permuted_hidden_states=buffers["down_output"].view(*topk_ids.shape, -1),
             topk_weights=topk_weights,
             inv_permuted_idx=inv_perm,
-            expert_first_token_offset=expert_first_token_offset,
+            expert_first_token_offset=expert_offsets,
         )
 
 
@@ -1090,14 +1015,13 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ) -> None:
-        """
-        Standard apply implementation for Humming batched grouped experts.
+        """Standard apply implementation for Humming batched grouped experts.
 
         Note: Humming kernels handle weights internally through the layer
         object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
         path (Humming quantizes the w13 input itself); for block-FP8 activations
         it carries the scale computed before dispatch, which is forwarded to
-        may_quant_input so Humming skips the redundant w13 quantization. The
+        process_input so Humming skips the redundant w13 quantization. The
         output is written into workspace13 via the buffer management.
         """
         assert not apply_router_weight_on_input
@@ -1120,11 +1044,12 @@ def apply(
         )
         buffers["down_output"] = output
 
-        inputs, input_scale = self.quantize_input(
+        inputs, input_scale, input_scale_2 = self.process_input(
             "w13",
             inputs=hidden_states,
             input_scale=a1q_scale,
             quanted_input=buffers.get("quanted_gate_up_input", None),
+            expert_tokens=expert_num_tokens,
         )
 
         self.humming_forward(
@@ -1132,6 +1057,7 @@ def apply(
             inputs=inputs,
             weight=w1,
             input_scale=input_scale,
+            input_scale_2=input_scale_2,
             outputs=buffers["gate_up_output"],
             valid_shape_m=valid_shape_m,
             expert_layout=expert_num_tokens,
@@ -1139,16 +1065,12 @@ def apply(
             tuning_config=self.w13_tuning_config_str,
         )
 
-        self.apply_activation(
-            activation=activation,
-            input=buffers["gate_up_output"],
-            output=buffers["activation_output"],
-        )
-
-        inputs, input_scale = self.quantize_input(
+        inputs, input_scale, input_scale_2 = self.process_input(
             "w2",
-            inputs=buffers["activation_output"],
-            quanted_input=buffers.get("quanted_down_input", None),
+            inputs=buffers["gate_up_output"],
+            quanted_input=buffers["quanted_down_input"],
+            expert_tokens=expert_num_tokens,
+            activation=activation,
         )
 
         self.humming_forward(
@@ -1156,6 +1078,7 @@ def apply(
             inputs=inputs,
             weight=w2,
             input_scale=input_scale,
+            input_scale_2=input_scale_2,
             outputs=output.view(-1, hidden_states.size(-1)),
             valid_shape_m=valid_shape_m,
             expert_layout=expert_num_tokens,
diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
index 8b87597ec268..27fdcf8bb2c6 100644
--- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
@@ -26,7 +26,7 @@
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
 
 from ..utils import swiglu_limit_func
 
@@ -441,48 +441,59 @@ def _sort_tokens_pow2(expt_scal, expt_indx, n_expts_tot, bitmatrix):
     _routing.sort_tokens = _sort_tokens_pow2
 
 
-# Two API generations of triton_kernels are supported:
-#   - v3.5.1 (the version bundled with vLLM): exposes `routing()` and
-#     `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix`
-#     constructor takes a `scratchpad` argument.
-#   - v3.6.0+: removes the `routing` module in favor of a `SparseMatrix`
-#     based path, and adds a `dtype=BIT` kwarg to `Bitmatrix`. Used only
-#     when the user has triton_kernels installed system-wide at v3.6.0+.
-#
-# `use_legacy_triton_kernels` selects between them at import time based on
-# whether `SparseMatrix` is importable.
-use_legacy_triton_kernels = False
+triton_kernels_version = get_triton_kernels_version()
 
-if has_triton_kernels():
+if triton_kernels_version is not None:
     try:
         import triton_kernels.swiglu
-        from triton_kernels.matmul_ogs import (
-            FnSpecs,
-            FusedActivation,
-            GatherIndx,
-            RoutingData,
-            ScatterIndx,
-            matmul_ogs,
-        )
-        from triton_kernels.tensor import (
-            BIT,
-            Bitmatrix,
-        )
 
+        if triton_kernels_version == "3.8":
+            from triton_kernels.matmul import FnSpecs, FusedActivation, matmul
+            from triton_kernels.matmul_details.opt_flags import (
+                update_opt_flags_constraints,
+            )
+            from triton_kernels.tensor import make_ragged_tensor_metadata
+
+            from vllm.platforms.rocm import on_gfx950
+
+            if on_gfx950():
+                # block_n=64/block_k=256 fit gfx950's 160KB LDS + the CDNA4
+                # scale swizzle; block_m=32 avoids an autotuner block_m=128
+                # padding cliff (~15x) at decode-sized ragged MoE shapes.
+                update_opt_flags_constraints(
+                    {"block_m": 32, "block_n": 64, "block_k": 256}
+                )
+        else:
+            from triton_kernels.matmul_ogs import (
+                FnSpecs,
+                FusedActivation,
+                GatherIndx,
+                RoutingData,
+                ScatterIndx,
+                matmul_ogs,
+            )
+
+            if triton_kernels_version == "3.6":
+                from triton_kernels.tensor import (
+                    SparseMatrix,
+                    make_ragged_tensor_metadata,
+                )
+
+        # BIT/Bitmatrix exist in v3.5.1/v3.6, gone in v3.8.
         try:
             from triton_kernels.tensor import (
-                SparseMatrix,
-                make_ragged_tensor_metadata,
+                BIT,
+                Bitmatrix,
             )
         except ImportError:
-            # TODO(mgoin): drop the v3.5.1 pin and remove this fallback once
-            # the gpt-oss perf regression in v3.6.0+ is resolved upstream.
-            # Tracking: https://github.com/triton-lang/triton/issues/9969
-            use_legacy_triton_kernels = True
-        if not use_legacy_triton_kernels:
+            BIT = None
+            Bitmatrix = None
+
+        # Patch the routing kernels for non-pow2 top_k (DeepSeek-V4 top_k=6);
+        # v3.8 rebuilds routing from topk ids and needs no patch.
+        if triton_kernels_version == "3.6":
             _patch_make_bitmatrix_metadata()
-        else:
-            # Legacy routing fails to compile for non-pow2 top_k (DeepSeek-V4).
+        elif triton_kernels_version == "3.5.1":
             _patch_legacy_routing_for_nonpow2_topk()
     except (AttributeError, ImportError) as e:
         logger.error(
@@ -492,6 +503,21 @@ def _sort_tokens_pow2(expt_scal, expt_indx, n_expts_tot, bitmatrix):
         )
 
 
+if triton_kernels_version == "3.8":
+    # 3.8's matmul dropped matmul_ogs's in-kernel scatter+topk-reduce, so the
+    # forward combines externally via index_add; this is the ragged routing
+    # bundle the (inlined) 3.8 paths pass around, in place of matmul_ogs's.
+    class RoutingData:  # type: ignore[no-redef]
+        __slots__ = ("ragged", "gate_scal", "gather_tok", "n_expts_act", "n_valid")
+
+        def __init__(self, ragged, gate_scal, gather_tok, n_expts_act, n_valid):
+            self.ragged = ragged
+            self.gate_scal = gate_scal
+            self.gather_tok = gather_tok
+            self.n_expts_act = n_expts_act
+            self.n_valid = n_valid
+
+
 @triton.jit
 def pack_bitmatrix(
     bitmatrix,
@@ -502,8 +528,7 @@ def pack_bitmatrix(
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
 ):
-    """
-    Packs topk_ids into a bitmatrix.
+    """Packs topk_ids into a bitmatrix.
     code reference:
     https://github.com/triton-lang/triton/blob/dd1bbc52b34d202dfe5ffea1e04fb16166c5c04e/python/triton_kernels/bench/distributed.py#L264
     """
@@ -563,7 +588,7 @@ def triton_kernel_moe_forward(
     # three separate kernels used by the generic path below.
     # Only available in the legacy (v3.5.1) API; the v3.6.0+ path inlines
     # equivalent logic via SparseMatrix in `make_routing_data`.
-    if use_legacy_triton_kernels and expert_map is None:
+    if triton_kernels_version == "3.5.1" and expert_map is None:
         from triton_kernels.routing import routing as fused_routing
 
         routing_data, gather_idx, scatter_idx = fused_routing(
@@ -663,6 +688,52 @@ def triton_kernel_fused_experts(
     assert activation == MoEActivation.SWIGLUOAI, (
         "Only SWIGLUOAI activation is supported"
     )
+
+    if triton_kernels_version == "3.8":
+        # 3.8 has no in-kernel scatter+topk-reduce, so the two matmuls are
+        # followed by an external index_add combine.
+        assert quant_config is not None
+        M, K = hidden_states.shape[-2:]
+        a_ragged = routing_data.ragged
+        gather_tok = routing_data.gather_tok
+        gammas = routing_data.gate_scal
+        act = FusedActivation(
+            FnSpecs(
+                "swiglu",
+                triton_kernels.swiglu.swiglu_fn,
+                ("alpha", "limit"),
+                reduction_n=2,
+            ),
+            (swiglu_alpha, swiglu_limit),
+        )
+        inter = matmul(
+            hidden_states,
+            w1,
+            quant_config.w1_bias,
+            a_ragged,
+            None,
+            gather_tok,
+            None,
+            quant_config.w1_precision,
+            gammas=gammas if apply_router_weight_on_input else None,
+            fused_activation=act,
+        )
+        y = matmul(
+            inter,
+            w2,
+            quant_config.w2_bias,
+            a_ragged,
+            None,
+            None,
+            None,
+            quant_config.w2_precision,
+            gammas=None if apply_router_weight_on_input else gammas,
+        )
+        acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device)
+        acc.index_add_(0, gather_tok.to(torch.int64), y.to(torch.float32))
+        out = output_tensor.view(M, K)
+        out.copy_(acc.to(out.dtype))
+        return out
     assert quant_config is not None
 
     # type check, uint8 means mxfp4
@@ -705,7 +776,7 @@ def triton_kernel_fused_experts(
             ),
             (swiglu_alpha, swiglu_limit),
         )
-        if not use_legacy_triton_kernels
+        if triton_kernels_version != "3.5.1"
         else FusedActivation(
             FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
             (swiglu_alpha, swiglu_limit),
@@ -745,6 +816,24 @@ def make_routing_data(
     topk_weights: torch.Tensor,
     num_local_experts: int,
 ) -> tuple["RoutingData", torch.Tensor, torch.Tensor]:
+    if triton_kernels_version == "3.8":
+        _, topk = topk_ids.shape
+        flat_e = topk_ids.reshape(-1)
+        valid = flat_e >= 0
+        flat_e = flat_e.to(torch.int64).clamp(min=0)
+        flat_w = topk_weights.reshape(-1).to(torch.float32) * valid.to(torch.float32)
+        nrows = flat_e.shape[0]
+        order = torch.argsort(flat_e, stable=True)
+        col_sum = torch.zeros(
+            num_local_experts, dtype=torch.int32, device=flat_e.device
+        )
+        col_sum.scatter_add_(0, flat_e, torch.ones_like(flat_e, dtype=torch.int32))
+        ragged = make_ragged_tensor_metadata(col_sum, nrows)
+        gather_tok = (order // topk).to(torch.int32)
+        gammas = flat_w[order]
+        routing = RoutingData(ragged, gammas, gather_tok, topk, nrows)
+        return routing, gather_tok, gather_tok
+
     topk_ids = topk_ids.to(torch.int16)
     topk_weights = topk_weights.to(torch.bfloat16)
 
@@ -775,7 +864,7 @@ def make_routing_data(
         Bitmatrix(
             bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max
         )
-        if not use_legacy_triton_kernels
+        if triton_kernels_version != "3.5.1"
         else Bitmatrix(
             bitmatrix,
             shape=bitmatrix_shape,
@@ -787,7 +876,7 @@ def make_routing_data(
     # matmul_ogs expects invalid topk_weights to be -1s
     topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights)
 
-    if use_legacy_triton_kernels:
+    if triton_kernels_version == "3.5.1":
         from triton_kernels.routing import routing_from_bitmatrix
 
         return routing_from_bitmatrix(
@@ -810,6 +899,10 @@ def routing_data_from_sparse_topk(
     call already produced; when the topk ids are used as-is (no expert-map
     remapping) that work is redundant.
     """
+    if triton_kernels_version == "3.8":
+        # 3.8 routing is rebuilt from raw ids; no SparseMatrix reuse.
+        return make_routing_data(sparse_topk.indx, sparse_topk.vals, n_expts)
+
     dispatch_indx = sparse_topk.mask_metadata.row_sorted_indx
     combine_indx = sparse_topk.mask_metadata.col_sorted_indx
     ragged_batch_metadata = make_ragged_tensor_metadata(
@@ -949,8 +1042,7 @@ def moe_problem_size(
         w2: torch.Tensor,
         topk_ids: torch.Tensor,
     ) -> tuple[int, int, int, int, int]:
-        """
-        Extract the MoE problem size from the given tensor arguments:
+        """Extract the MoE problem size from the given tensor arguments:
         - a: The hidden states, input to the MoE layer.
         - w1: The first set of expert weights.
         - w2: The second set of expert weights.
@@ -1075,8 +1167,7 @@ def apply(
 
 
 class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
-    """
-    A Triton based MoE expert class that operates on expert standard
+    """A Triton based MoE expert class that operates on expert standard
     format and explicitly keeps the activation and reduction (moe_sum) steps
     unfused from the matmul_ogs kernel. This exposes injection points
     for activation and moe_sum.
@@ -1200,6 +1291,50 @@ def apply(
 
         topk = topk_ids.size(1)
 
+        if triton_kernels_version == "3.8":
+            # Activation and topk-reduce kept unfused from the matmuls (3.8 has
+            # no in-kernel scatter+reduce); combine via external index_add.
+            assert self._lora_context is None, (
+                "tk38 unfused MoE path does not support LoRA yet"
+            )
+            act_out_dim = self.adjust_N_for_activation(w1.shape[2], activation)
+            M, K = hidden_states.shape[-2:]
+            a_ragged = routing_data.ragged
+            gather_tok = routing_data.gather_tok
+            gammas = routing_data.gate_scal
+            inter = matmul(
+                hidden_states,
+                w1,
+                quant_config.w1_bias,
+                a_ragged,
+                None,
+                gather_tok,
+                None,
+                quant_config.w1_precision,
+                gammas=gammas if apply_router_weight_on_input else None,
+            )
+            act_out = torch.empty(
+                (routing_data.n_valid, act_out_dim),
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
+            self.activation(activation, act_out, inter)
+            down = matmul(
+                act_out,
+                w2,
+                quant_config.w2_bias,
+                a_ragged,
+                None,
+                None,
+                None,
+                quant_config.w2_precision,
+                gammas=None if apply_router_weight_on_input else gammas,
+            )
+            acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device)
+            acc.index_add_(0, gather_tok.to(torch.int64), down.to(torch.float32))
+            output.view(M, K).copy_(acc.to(output.dtype))
+            return
+
         # type check, uint8 means mxfp4
         assert hidden_states.dtype == torch.bfloat16
         assert (
diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_context.py b/vllm/model_executor/layers/fused_moe/experts/lora_context.py
index 48232913229d..fb41ef544427 100644
--- a/vllm/model_executor/layers/fused_moe/experts/lora_context.py
+++ b/vllm/model_executor/layers/fused_moe/experts/lora_context.py
@@ -10,8 +10,7 @@
 
 @dataclass
 class MoELoRAContext:
-    """
-    Carries all LoRA state for one MoE forward pass.
+    """Carries all LoRA state for one MoE forward pass.
 
     Built by FusedMoEWithLoRA.forward() and propagated explicitly through the
     modular kernel path (FusedMoEKernel -> FusedMoEExpertsModular.apply) so
diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py
index 0aa6f6c30b6d..ac75e7bc798b 100644
--- a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py
+++ b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py
@@ -7,8 +7,7 @@
 
 
 class LoRAExpertsMixin:
-    """
-    Mixin for FusedMoEExpertsModular subclasses that natively handle
+    """Mixin for FusedMoEExpertsModular subclasses that natively handle
     MoELoRAContext inside their apply() implementation.
 
     Mixing this class in:
diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
index 378b6a2314d2..cfd78de28dd1 100644
--- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
@@ -35,7 +35,7 @@
 from vllm.model_executor.layers.fused_moe.utils import _resize_cache
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
+    get_marlin_workspace,
     marlin_moe_intermediate_size,
     marlin_quant_input,
 )
@@ -95,7 +95,7 @@ def _fused_marlin_moe(
     N = marlin_moe_intermediate_size(w1, w2)
     w13_num_shards = 2 if activation.is_gated else 1
     if workspace is None:
-        workspace = marlin_make_workspace_new(hidden_states.device, 4)
+        workspace = get_marlin_workspace(hidden_states.device)
 
     if intermediate_cache13 is None:
         intermediate_cache13 = torch.empty(
@@ -251,26 +251,56 @@ def fused_marlin_moe(
     input_dtype: torch.dtype | None = None,
     activation_config: ApplyMoEActivationConfig | None = None,
 ) -> torch.Tensor:
-    """
-    This function computes a Mixture of Experts (MoE) layer using two sets of
+    """This function computes a Mixture of Experts (MoE) layer using two sets of
     weights, w1 and w2, and top-k gating mechanism.
 
-    Parameters:
-    - hidden_states (torch.Tensor): The input tensor to the MoE layer.
-    - w1 (torch.Tensor): The first set of expert weights.
-    - w2 (torch.Tensor): The second set of expert weights.
-    - w1_scale (torch.Tensor): Scale to be used for w1.
-    - w2_scale (torch.Tensor): Scale to be used for w2.
-    - topk_weights (torch.Tensor): Top-k weights.
-    - topk_ids (torch.Tensor): Indices of topk-k elements.
-    - w1_zeros (torch.Tensor|None): Optional zero points to be used for w1.
-    - w2_zeros (torch.Tensor|None): Optional zero points to be used for w2.
-    - num_bits (bool): The number of bits in expert weights quantization.
+    Args:
+        hidden_states (torch.Tensor): The input tensor to the MoE layer.
+        w1 (torch.Tensor): The first set of expert weights.
+        w2 (torch.Tensor): The second set of expert weights.
+        bias1 (torch.Tensor|None): Optional bias for the first GEMM.
+        bias2 (torch.Tensor|None): Optional bias for the second GEMM.
+        w1_scale (torch.Tensor): Scale to be used for w1.
+        w2_scale (torch.Tensor): Scale to be used for w2.
+        topk_weights (torch.Tensor): Top-k weights.
+        topk_ids (torch.Tensor): Indices of topk-k elements.
+        quant_type_id (int): Id of the `ScalarType` the expert weights are
+            quantized to. Determines the 4- or 8-bit Marlin kernel path.
+        apply_router_weight_on_input (bool): Apply the topk weights to the
+            input rather than the output. Only valid when topk is 1.
+        global_num_experts (int): Total number of experts across all expert
+            parallel shards. -1 means the local expert count.
+        activation (MoEActivation): Activation applied between the two GEMMs.
+        activation_func (Callable|None): Override for the activation
+            implementation.
+        moe_sum (Callable|None): Override for the final reduction over the
+            topk dimension. Defaults to `ops.moe_sum` / `torch.sum`.
+        expert_map (torch.Tensor|None): Maps global expert indices to the
+            local expert space of this expert parallel shard.
+        input_global_scale1 (torch.Tensor|None): NVFP4 global input scale for
+            the first GEMM.
+        input_global_scale2 (torch.Tensor|None): NVFP4 global input scale for
+            the second GEMM.
+        global_scale1 (torch.Tensor|None): NVFP4 global weight scale for w1.
+        global_scale2 (torch.Tensor|None): NVFP4 global weight scale for w2.
+        w1_zeros (torch.Tensor|None): Optional zero points to be used for w1.
+        w2_zeros (torch.Tensor|None): Optional zero points to be used for w2.
+        workspace (torch.Tensor|None): Scratch buffer for the Marlin kernels.
+        intermediate_cache13 (torch.Tensor|None): Scratch buffer for the
+            first and third intermediate activations.
+        intermediate_cache2 (torch.Tensor|None): Scratch buffer for the
+            second intermediate activation.
+        output (torch.Tensor|None): Tensor to write the result into. A new
+            tensor is allocated when omitted.
+        input_dtype (torch.dtype|None): dtype the activations are quantized
+            to. A 1-byte dtype raises the M block size to at least 16.
+        activation_config (ApplyMoEActivationConfig|None): Extra
+            configuration forwarded to the activation.
 
     Returns:
-    - torch.Tensor: The output tensor after applying the MoE layer.
-    """
+        torch.Tensor: The output tensor after applying the MoE layer.
 
+    """
     quant_type = ScalarType.from_id(quant_type_id)
     assert quant_type in [
         scalar_types.uint4,
@@ -399,8 +429,7 @@ def batched_fused_marlin_moe(
     activation_func: Callable[..., None] | None = None,
     activation_config: ApplyMoEActivationConfig | None = None,
 ) -> torch.Tensor:
-    """
-    This function massages the inputs so the batched hidden_states can be
+    """This function massages the inputs so the batched hidden_states can be
     presented as a 2D contiguous tensor that could be used with
     _fused_marlin_moe.
 
@@ -423,7 +452,6 @@ def batched_fused_marlin_moe(
     expert_ids tensors, so only relevant/valid rows of A (hidden_states)
     are accessed and are processed with the correct expert_ids.
     """
-
     assert hidden_states.ndim == 3, (
         f"hidden states must be batched. e.g. [B, MAX_TOKENS, K]."
         f"But got {hidden_states.size()}"
diff --git a/vllm/model_executor/layers/fused_moe/experts/moonep_experts.py b/vllm/model_executor/layers/fused_moe/experts/moonep_experts.py
new file mode 100644
index 000000000000..1af4b5f06460
--- /dev/null
+++ b/vllm/model_executor/layers/fused_moe/experts/moonep_experts.py
@@ -0,0 +1,220 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MoonEP experts: grouped GEMM over MoonEP's expert-grouped activations.
+
+MoonEP dispatch delivers tokens already contiguous per weight row
+(``[NvS, H]`` in ``cu_seqlens[E+B]`` segment order), so the expert compute
+is three grouped GEMMs over those segments with no permute/unpermute:
+
+    gate = grouped_mm(x, w_gate[E+B])      # [NvS, I]
+    up   = grouped_mm(x, w_up[E+B])        # [NvS, I]
+    act  = silu(gate) * up * route_weight  # route weights applied here
+    out  = grouped_mm(act, w_down[E+B])    # [NvS, H]
+
+Rows ``[E, E+B)`` are the redundant-expert prefetch slots that
+``MoonEPPrepareAndFinalize`` fills before ``apply`` runs. Empty segments
+(including unused prefetch slots) are skipped by the grouped GEMM.
+
+Weight layout: MoonEP's ``prefetch_weight`` requires each projection to be
+its own contiguous ``[E+B, ., .]`` tensor, so gate and up are separate
+tensors rather than vLLM's interleaved ``[E, 2I, H]`` ``w13``. The modular
+kernel passes ``w1``/``w2`` from the layer (``w1`` is the gate tensor); the
+full layout is picked up from the layer in ``process_weights_after_loading``
+and shared with ``MoonEPPrepareAndFinalize`` via ``post_init_setup``.
+"""
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceNoOP,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
+from vllm.platforms import current_platform
+
+
+def moonep_grouped_gemm(
+    a: torch.Tensor,
+    w: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+) -> torch.Tensor:
+    """``out[seg_g] = a[seg_g] @ w[g].T`` for each ``cu_seqlens`` segment.
+
+    Args:
+        a: ``[NvS, K]`` activations in segment order.
+        w: ``[G, N, K]`` weights (vLLM ``[E, N, K]`` convention).
+        cu_seqlens: ``[G]`` int32 cumulative segment end offsets.
+
+    Returns:
+        ``[NvS, N]``; rows past ``cu_seqlens[-1]`` are zero-filled.
+
+    """
+    return torch._grouped_mm(a, w.transpose(1, 2), offs=cu_seqlens)
+
+
+class MoonEPExperts(mk.FusedMoEExpertsModular):
+    """Grouped-GEMM experts over MoonEP's ``[NvS, H]`` layout (BF16)."""
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+        max_num_tokens: int | None = None,
+        num_dispatchers: int | None = None,
+    ):
+        super().__init__(moe_config, quant_config, max_num_tokens, num_dispatchers)
+        assert quant_config.quant_dtype is None, (
+            "MoonEPExperts supports unquantized BF16 only"
+        )
+        if (
+            quant_config.gemm1_alpha is not None
+            or quant_config.gemm1_beta is not None
+            or quant_config.gemm1_clamp_limit is not None
+        ):
+            raise NotImplementedError(
+                "MoonEPExperts implements plain silu(gate) * up only; SwiGLU "
+                "variants with alpha/beta/clamp parameters are not supported"
+            )
+        self._weight_layout = None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # The [E+B] gate/up/down layout built by
+        # convert_to_unquantized_kernel_format; MoonEPPrepareAndFinalize
+        # reads it from here (post_init_setup) for weight prefetch.
+        layout = getattr(layer, "_moonep_weight_layout", None)
+        assert layout is not None, (
+            "MoonEPExperts requires the layer to carry _moonep_weight_layout "
+            "(set by convert_to_unquantized_kernel_format)"
+        )
+        self._weight_layout = layout
+
+    @property
+    def weight_layout(self):
+        return self._weight_layout
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        return current_platform.is_cuda()
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return False
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation == MoEActivation.SILU
+
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        # EPLB rearranges the named expert parameters as local_num_experts
+        # rows, which does not understand the replicated [E+B] layout.
+        return (
+            moe_parallel_config.use_moonep_kernels
+            and not moe_parallel_config.enable_eplb
+        )
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return weight_key is None and activation_key is None
+
+    def supports_chunking(self) -> bool:
+        # NvS is static and segments must stay whole.
+        return False
+
+    def supports_expert_map(self) -> bool:
+        # MoonEP addresses global expert rows directly.
+        return False
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        return TopKWeightAndReduceNoOP()
+
+    def moe_problem_size(
+        self,
+        a1: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> tuple[int, int, int, int, int]:
+        # a1 is [NvS, H] in segment order, not token-major, so the base
+        # implementation's topk_ids/a1 row-count check does not apply.
+        assert a1.dim() == 2 and w1.dim() == 3 and w2.dim() == 3
+        num_rows, intermediate, hidden = w1.shape  # [E+B, I, H]
+        assert a1.size(1) == hidden
+        return num_rows, a1.size(0), intermediate, hidden, topk_ids.size(1)
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        # M is NvS. torch._grouped_mm allocates its own outputs, so no
+        # workspaces are needed.
+        return ((0,), (0,), (M, K))
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ) -> None:
+        # expert_map is ignored: MoonEP dispatches on global expert ids and
+        # the weights are addressed by global [E+B] row, on every rank.
+        assert activation == MoEActivation.SILU
+        assert self._weight_layout is not None, (
+            "process_weights_after_loading() not called"
+        )
+        assert expert_tokens_meta is not None
+        # MoonEPPrepareAndFinalize returns route weights in NvS order as the
+        # dispatched topk_weights, and per-row segment lengths as
+        # expert_tokens_meta.
+        route_weights_nvs = topk_weights
+        assert route_weights_nvs.dim() == 1
+        cu_seqlens = torch.cumsum(
+            expert_tokens_meta.expert_num_tokens, dim=0, dtype=torch.int32
+        )
+
+        assert cu_seqlens.numel() == w1.size(0)
+
+        gate = moonep_grouped_gemm(hidden_states, w1, cu_seqlens)
+        up = moonep_grouped_gemm(
+            hidden_states, self._weight_layout.full_up_weight, cu_seqlens
+        )
+        act = torch.nn.functional.silu(gate)
+        act.mul_(up)
+        if not apply_router_weight_on_input:
+            # Input-weighted routing (top-1, e.g. Llama 4) was already
+            # applied to the activations by prepare(), before dispatch.
+            act.mul_(route_weights_nvs.to(act.dtype).unsqueeze(-1))
+        # Padding rows past the last segment come out zero-filled.
+        output.copy_(moonep_grouped_gemm(act, w2, cu_seqlens))
diff --git a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py
index cd862cac5958..f52473c68fff 100644
--- a/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-NVFP4 quantization emulation for MoE.
+"""NVFP4 quantization emulation for MoE.
 
 This file implements NVFP4 emulation for NVFP4 MOE in case the hardware used does not
 natively support NVFP4 MOE.
@@ -86,8 +85,7 @@ def fused_moe_nvfp4_emulation_kernel(
     compute_type: tl.constexpr,
     group_size: tl.constexpr,
 ):
-    """
-    Fused MoE kernel for emulated NVFP4 weight-only dequantization + GEMM.
+    """Fused MoE kernel for emulated NVFP4 weight-only dequantization + GEMM.
 
     Activations A are BF16 (already QDQ'd externally).
     Weights B are packed uint8 NVFP4 [E, N, K//2] — two FP4 values per byte
@@ -334,8 +332,7 @@ def invoke_fused_moe_nvfp4_emulation_kernel(
 
 
 class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
-    """
-    Extension of TritonExperts to support emulated NVFP4 MoE experts.
+    """Extension of TritonExperts to support emulated NVFP4 MoE experts.
 
     It may be used for NVFP4 models when the device does not have
     native support for this dtype.
@@ -427,8 +424,7 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
-        """
-        Apply emulated quantized MoE computation.
+        """Apply emulated quantized MoE computation.
 
         This dequantizes the weights on the fly and calls fused_experts_impl
         with activation quantization support.
diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
index 2c91e22d9108..da0bca4977b7 100644
--- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-OCP MX quantization emulation for MoE.
+"""OCP MX quantization emulation for MoE.
 
 This file implements OCP MX (MXFP4/MXFP6) emulation for MoE in case the
 hardware used does not natively support OCP MX MoE.
@@ -50,6 +49,7 @@ def activation_quant_dtype(
 
     Raises:
         NotImplementedError: If the scheme has no emulated activation dtype.
+
     """
     if ocp_mx_scheme in {
         OCP_MX_Scheme.w_mxfp4,
@@ -82,8 +82,7 @@ def activation_quant_dtype(
 
 
 class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
-    """
-    Extension of TritonExperts to support emulated OCP MX MoE experts.
+    """Extension of TritonExperts to support emulated OCP MX MoE experts.
 
     It may be used for OCP MX (MXFP4/MXFP6) models when the device does not
     have native support for these dtypes.
@@ -190,8 +189,7 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
-        """
-        Apply emulated quantized MoE computation.
+        """Apply emulated quantized MoE computation.
 
         This dequantizes the weights on the fly and calls TritonExperts.apply
         with activation quantization support.
diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py
new file mode 100644
index 000000000000..dbbce6b0b182
--- /dev/null
+++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py
@@ -0,0 +1,235 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused MoE W4A16 experts on the RDNA3 (gfx1100) HIP kernel.
+
+``moe_gptq_gemm_rdna3`` is a single HIP kernel launch per GEMM that handles
+expert routing + W4A16 dequant + dot product with atomic output accumulation.
+
+Weight format (per expert, same as the dense RDNA3 W4A16 kernel):
+  - Packed int32 ``[E, K/8, N]`` with exllama shuffle
+  - Scales ``[E, groups, N]`` in activation dtype
+  - Zero points ``[E, groups, N/8]`` packed int32 (synthesized, symmetric only)
+"""
+
+import torch
+
+import vllm._custom_ops as ops
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.activation import (
+    MoEActivation,
+    apply_moe_activation_supported,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
+    moe_align_block_size,
+)
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceNoOP,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kInt4Static,
+    kInt4Static32,
+)
+from vllm.platforms import current_platform
+
+
+def rdna3_moe_kernel_available() -> bool:
+    """Whether the fused RDNA3 MoE HIP kernel is built into this binary."""
+    if not current_platform.is_rocm():
+        return False
+
+    from vllm.platforms.rocm import on_gfx1100
+
+    return (
+        on_gfx1100()
+        and hasattr(torch.ops, "_rocm_C")
+        and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
+    )
+
+
+class Rdna3WNA16Experts(mk.FusedMoEExpertsModular):
+    """W4A16 experts backed by ``moe_gptq_gemm_rdna3`` (gfx1100).
+
+    Both GEMMs accumulate atomically, so their destinations are zeroed first.
+    The second GEMM is given ``output_topk``, which makes it reduce over the
+    top-k dimension while it accumulates — the ``moe_sum`` launch and the
+    ``[M * top_k, K]`` intermediate are fused away, hence
+    ``TopKWeightAndReduceNoOP``.
+    """
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+        max_num_tokens: int | None = None,
+        num_dispatchers: int | None = None,
+    ):
+        assert quant_config.use_int4_w4a16, "Supports only int4_w4a16"
+        super().__init__(
+            moe_config=moe_config,
+            quant_config=quant_config,
+            max_num_tokens=max_num_tokens,
+            num_dispatchers=num_dispatchers,
+        )
+        self._empty_topk_weights: torch.Tensor | None = None
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        return TopKWeightAndReduceNoOP()
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        return rdna3_moe_kernel_available()
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return True
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        # Symmetric int4 weights only: the kernel consumes synthesized zero
+        # points and has no path for checkpoint zero points.
+        return activation_key is None and weight_key in (kInt4Static, kInt4Static32)
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return apply_moe_activation_supported(activation)
+
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        return not (
+            moe_parallel_config.use_fi_nvl_two_sided_kernels
+            or moe_parallel_config.use_fi_nvl_one_sided_kernels
+        )
+
+    def moe_problem_size(
+        self,
+        a1: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> tuple[int, int, int, int, int]:
+        # Weights are packed K-first: w1 is [E, K // 8, w13_shards * N] and
+        # w2 is [E, N // 8, K], so neither N nor K can be read off a trailing
+        # dimension the way the base implementation does.
+        assert w1.dim() == 3 and w2.dim() == 3
+        assert a1.dim() == 2
+        assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"
+        return w1.size(0), a1.size(0), w2.size(1) * 8, a1.size(-1), topk_ids.size(1)
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        # The modular kernel provisions the output buffer out of workspace13,
+        # so both GEMM-1 buffers have to live in workspace2 — the second GEMM
+        # reads the activation output while accumulating into the output.
+        gate_up = N if not activation.is_gated else 2 * N
+        return ((M, K), (M * topk * (gate_up + N),), (M, K))
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ) -> None:
+        assert self.w1_scale is not None and self.w2_scale is not None
+        assert self.w1_zp is not None and self.w2_zp is not None
+
+        E, M, N, K, top_k = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
+        if global_num_experts == -1:
+            global_num_experts = E
+
+        gate_up = w1.size(2)
+        act_n = self.adjust_N_for_activation(gate_up, activation)
+        rows = M * top_k
+
+        scratch = workspace2.view(-1)
+        gate_up_out = scratch[: rows * gate_up].view(rows, gate_up)
+        act_out = scratch[rows * gate_up : rows * (gate_up + act_n)].view(rows, act_n)
+
+        # BLOCK_SIZE_M=1 for decode (no padding waste), 4 for prefill.
+        block_size_m = 1 if M <= 4 else 4
+        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+            topk_ids,
+            block_size_m,
+            global_num_experts,
+            expert_map,
+        )
+
+        if (
+            self._empty_topk_weights is None
+            or self._empty_topk_weights.device != hidden_states.device
+        ):
+            self._empty_topk_weights = torch.empty(
+                0, dtype=torch.float32, device=hidden_states.device
+            )
+        topk_weights_f32 = topk_weights.reshape(-1).float()
+        no_topk_weights = self._empty_topk_weights
+
+        gate_up_out.zero_()
+        ops.moe_gptq_gemm_rdna3(
+            hidden_states,
+            gate_up_out,
+            w1,
+            self.w1_scale,
+            self.w1_zp,
+            topk_weights_f32 if apply_router_weight_on_input else no_topk_weights,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            top_k,
+            block_size_m,
+            apply_router_weight_on_input,
+        )
+
+        self.activation(activation, act_out, gate_up_out)
+
+        # output_topk=top_k makes the kernel accumulate into out[token_id],
+        # fusing the top-k reduction into the atomic write-back.
+        output.zero_()
+        ops.moe_gptq_gemm_rdna3(
+            act_out,
+            output,
+            w2,
+            self.w2_scale,
+            self.w2_zp,
+            no_topk_weights if apply_router_weight_on_input else topk_weights_f32,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            1,
+            block_size_m,
+            not apply_router_weight_on_input,
+            output_topk=top_k,
+        )
diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
index 044a2c29cca2..f8133b776e0c 100644
--- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
@@ -385,14 +385,10 @@ def rocm_aiter_fused_experts(
 
         gate_mode = ""
         if activation == MoEActivation.SITU:
-            # a8w4 (VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1) uses the gate/up-
-            # interleaved (_gui_) fp8 flydsl kernels; default a16w4 SiTU stays
-            # separated.
-            gate_mode = (
-                GateMode.INTERLEAVE.value
-                if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()
-                else GateMode.SEPARATED.value
-            )
+            # SiTUv2 flydsl (VLLM_ROCM_USE_AITER_MOE_SITUV2=1) uses a4w4
+            # fp4 activations with separated gate/up weights (AITER #4463);
+            # default a16w4 SiTU also stays separated.
+            gate_mode = GateMode.SEPARATED.value
         elif quant_config.use_mxfp4_w4a16:
             gate_mode = GateMode.INTERLEAVE.value
         elif activation_interleave is not None:
diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
index 9f982644b569..ee3b9b3f4bb2 100644
--- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
@@ -130,10 +130,18 @@ def _supports_quant_scheme(
         # INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs
         # (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and
         # the Triton int8_w8a8 fused MoE kernel handles them.
+        # XPU reaches the same Triton kernel through the same launcher, and both
+        # int8 activation-quant paths work there: per-token uses the Triton
+        # `per_token_quant_int8`, per-tensor resolves to the XPU branch of
+        # `scaled_int8_quant`, which is plain elementwise arithmetic.
         device_supports_int8 = (
-            current_platform.is_cuda()
-            and current_platform.has_device_capability((7, 5))
-        ) or current_platform.is_rocm()
+            (
+                current_platform.is_cuda()
+                and current_platform.has_device_capability((7, 5))
+            )
+            or current_platform.is_rocm()
+            or current_platform.is_xpu()
+        )
 
         supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
         if device_supports_int8:
@@ -329,6 +337,25 @@ def apply(
         )
         intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
 
+        # Under EP, drop the top-k slots routed to remote experts at alignment
+        # time instead of launching `off_experts == -1` GEMM blocks that only
+        # write zeros; the pad-aware `moe_sum` below skips the same slots.
+        # Their rows in the workspace caches are then never written, so every
+        # consumer of those caches must be row-local: a dynamic per-tensor
+        # scale for the second GEMM's activation would take its amax over the
+        # untouched rows, and LoRA aligns the full `topk_ids` itself.
+        a2_scale_is_global = (
+            self.quant_dtype is not None
+            and a2_scale is None
+            and self.block_shape is None
+            and (not self.per_act_token_quant or self.quantization_emulation)
+        )
+        skip_invalid = (
+            expert_map is not None
+            and self._lora_context is None
+            and not a2_scale_is_global
+        )
+
         # Include fused shared-expert rows while preserving EP remapping.
         num_align_experts = w1.shape[0] if expert_map is None else global_num_experts
         sorted_token_ids, expert_ids, num_tokens_post_padded = (
@@ -342,6 +369,7 @@ def apply(
                 use_int8_w8a16=self.quant_config.use_int8_w8a16,
                 use_int4_w4a16=self.quant_config.use_int4_w4a16,
                 block_shape=self.block_shape,
+                ignore_invalid_experts=skip_invalid,
             )
         )
 
@@ -471,8 +499,11 @@ def _lora_w13_fn():
         # Fuse SiLU+Mul + FP8 block quantize into a single kernel
         # when conditions permit (gated SiLU, fp8 block quant with
         # group_size=128, no LoRA requiring the BF16 intermediate).
+        # The fused kernel has no clamp parameter, so a configured
+        # SwiGLU clamp limit falls through to the unfused path.
         if (
             activation == MoEActivation.SILU
+            and self.activation_config.clamp_limit is None
             and self.quant_config.use_fp8_w8a8
             and self.block_shape == [128, 128]
             and lora_context is None
@@ -574,10 +605,21 @@ def _lora_w2_fn():
                 )
 
         # separate function is required for MoE + LoRA
-        self.moe_sum(intermediate_cache3, output)
+        self.moe_sum(intermediate_cache3, output, topk_ids, expert_map)
 
-    def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
-        ops.moe_sum(input, output)
+    def moe_sum(
+        self,
+        input: torch.Tensor,
+        output: torch.Tensor,
+        topk_ids: torch.Tensor | None = None,
+        expert_map: torch.Tensor | None = None,
+    ) -> None:
+        if expert_map is not None:
+            # Skip the slots whose expert is not on this rank: the rows the
+            # alignment dropped, or the zeros the `-1` blocks wrote.
+            ops.moe_sum(input, output, topk_ids, expert_map)
+        else:
+            ops.moe_sum(input, output)
 
 
 class TritonWNA16Experts(TritonExperts):
@@ -596,9 +638,9 @@ def _supports_quant_scheme(
     ) -> bool:
         SUPPORTED_W = [
             kInt4Static,
+            kInt4StaticAsym,
             kInt8Static,
             kInt4Static32,
-            kInt4StaticAsym,
             kInt4Static32Asym,
             # other group sizes?
         ]
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
index 6262234679af..54227d3cd1eb 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
@@ -31,9 +31,17 @@
 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
 
 
+def view_as_block_major_k(weight: torch.Tensor) -> torch.Tensor:
+    """View packed storage, including legacy 4D IPC cache entries."""
+    if weight.ndim == 4:
+        return weight
+    experts, rows, cols = weight.shape
+    block_k = 128 // weight.element_size()
+    return weight.view(experts, cols // block_k, rows, block_k)
+
+
 class TrtLlmBf16ExpertsBase:
-    """
-    BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and
+    """BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and
     monolithic interfaces.
     """
 
@@ -124,9 +132,7 @@ def expects_unquantized_inputs(self) -> bool:
 
 
 class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular):
-    """
-    BF16 unquantized TRTLLM-Gen MoE kernels. Supports modular interface.
-    """
+    """BF16 unquantized TRTLLM-Gen MoE kernels. Supports modular interface."""
 
     @staticmethod
     def _supports_parallel_config(
@@ -203,8 +209,8 @@ def apply(
         result = flashinfer.fused_moe.trtllm_bf16_routed_moe(
             topk_ids=(topk_ids, topk_weights),
             hidden_states=hidden_states,
-            gemm1_weights=w1,
-            gemm2_weights=w2,
+            gemm1_weights=view_as_block_major_k(w1),
+            gemm2_weights=view_as_block_major_k(w2),
             num_experts=global_num_experts,
             top_k=topk_ids.size(1),
             n_group=None,
@@ -224,9 +230,7 @@ def apply(
 
 
 class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonolithic):
-    """
-    BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface.
-    """
+    """BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface."""
 
     @staticmethod
     def _supports_parallel_config(
@@ -286,8 +290,8 @@ def apply(
             routing_logits=router_logits,
             routing_bias=e_score_correction_bias,
             hidden_states=hidden_states,
-            gemm1_weights=w1,
-            gemm2_weights=w2,
+            gemm1_weights=view_as_block_major_k(w1),
+            gemm2_weights=view_as_block_major_k(w2),
             num_experts=global_num_experts,
             top_k=self.topk,
             n_group=num_expert_group,
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
index 6f256db960e8..e860597bcba3 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
@@ -58,8 +58,7 @@ def prepare_deepseek_fp8_x_sf(x: torch.Tensor, x_sf: torch.Tensor) -> torch.Tens
 
 
 class TrtLlmFp8ExpertsBase:
-    """
-    Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic
+    """Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic
     interfaces.
     """
 
@@ -194,9 +193,7 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo
 
 
 class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
-    """
-    Fp8 TRTLLM-Gen MoE kernels. Supports modular interface.
-    """
+    """Fp8 TRTLLM-Gen MoE kernels. Supports modular interface."""
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -330,9 +327,7 @@ def apply(
 
 
 class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolithic):
-    """
-    Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface.
-    """
+    """Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface."""
 
     def supports_routing_replay_capture(self) -> bool:
         return True
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py
index a3d50011d919..1c5698932a4e 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-LoRA-aware FlashInfer TRT-LLM MoE experts (BF16).
+"""LoRA-aware FlashInfer TRT-LLM MoE experts (BF16).
 
 Reuses the routed API + ``gemm1_lora_delta`` path from FlashInfer PR #3153:
 
@@ -35,6 +34,9 @@
 from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (
     LoRAExpertsMixin,
 )
+from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
+    view_as_block_major_k,
+)
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
@@ -109,7 +111,7 @@ def _finalize_lora_kernel(
 
 
 class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
-    """LoRA-aware trtllm MoE experts"""
+    """LoRA-aware trtllm MoE experts."""
 
     def __init__(
         self,
@@ -528,8 +530,8 @@ def invoke_routed_moe(
         ret = flashinfer.fused_moe.trtllm_bf16_routed_moe(
             topk_ids=topk_ids_and_weights,
             hidden_states=hidden_states,
-            gemm1_weights=w1,
-            gemm2_weights=w2,
+            gemm1_weights=view_as_block_major_k(w1),
+            gemm2_weights=view_as_block_major_k(w2),
             gemm1_lora_delta=gemm1_lora_delta,
             num_experts=global_num_experts,
             top_k=self.topk,
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
index 159c77848a56..1c20bc1cbcdf 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
@@ -35,9 +35,7 @@
 
 
 class TrtLlmMxfp4ExpertsBase:
-    """
-    MXFP4 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic.
-    """
+    """MXFP4 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic."""
 
     def __init__(
         self,
@@ -164,8 +162,7 @@ def expects_unquantized_inputs(self) -> bool:
 class TrtLlmMxfp4ExpertsMonolithic(
     TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsMonolithic
 ):
-    """
-    Monolithic version of the MXFP4 TRTLLM kernel (router + experts).
+    """Monolithic version of the MXFP4 TRTLLM kernel (router + experts).
     Wraps flashinfer.trtllm_fp4_block_scale_moe().
     """
 
@@ -287,8 +284,7 @@ def apply(
 
 
 class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModular):
-    """
-    Modular version of the MXFP4 TRTLLM kernel (just the experts).
+    """Modular version of the MXFP4 TRTLLM kernel (just the experts).
     Wraps flashinfer.trtllm_fp4_block_scale_routed_moe().
     Moved from trtllm_moe.py.
     """
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py
index a9e725f2965c..e3e9bc6f2ee0 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py
@@ -19,8 +19,7 @@
 
 
 class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
-    """
-    FlashInfer TRT-LLM MxInt4 MoE kernel. Monolithic interface
+    """FlashInfer TRT-LLM MxInt4 MoE kernel. Monolithic interface
     (fused router + experts).
 
     Wraps flashinfer_trtllm_mxint4_moe().
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
index 6a54c0d104ec..eec846b5a0e3 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
@@ -28,6 +28,7 @@
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kNvfp4Dynamic,
+    kNvfp4DynamicToken,
     kNvfp4Static,
 )
 from vllm.model_executor.utils import is_weights_pre_processed
@@ -42,9 +43,7 @@
 
 
 class TrtLlmNvFp4ExpertsBase:
-    """
-    NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface.
-    """
+    """NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface."""
 
     def __init__(
         self,
@@ -224,10 +223,10 @@ def _supports_quant_scheme(
         activation_key: QuantKey | None,
     ) -> bool:
         """Supports Nvfp4 quantization."""
-        SUPPORTED_W_A = [
-            (kNvfp4Static, kNvfp4Dynamic),
-        ]
-        return (weight_key, activation_key) in SUPPORTED_W_A
+        return weight_key == kNvfp4Static and activation_key in (
+            kNvfp4Dynamic,
+            kNvfp4DynamicToken,
+        )
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
@@ -300,9 +299,7 @@ def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int:
 
 
 class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular):
-    """
-    Modular version of the implementation (just the experts).
-    """
+    """Modular version of the implementation (just the experts)."""
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -320,22 +317,12 @@ def workspace_shapes(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         activation: MoEActivation,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        if self.per_token_activation:
-            # Deferred input quant leaves K unpacked here, breaking the
-            # workspace assumptions below. Per-token NVFP4 is only supported on
-            # the monolithic (non-EP) path for now.
-            raise NotImplementedError(
-                "NVFP4 per-token activation is only supported on the monolithic "
-                "(non-EP) FlashInfer TRTLLM MoE path."
-            )
-
         # The workspaces for this implementation are managed by flashinfer.
         workspace1 = (0,)
         workspace2 = (0,)
 
-        # Hidden states are Nvfp4, packed into int8 dtype, so we
-        # need to multiply K by 2 to get the output shape right.
-        assert self.hidden_dim == K * 2
+        # Per-token inputs are unpacked; otherwise each byte holds two FP4 values.
+        assert self.hidden_dim == (K if self.expects_unquantized_inputs else K * 2)
         output = (M, self.hidden_dim)
 
         return (workspace1, workspace2, output)
@@ -353,7 +340,7 @@ def _invoke_kernel(
         topk_ids: torch.Tensor,
         activation: MoEActivation,
         global_num_experts: int,
-        a1q_scale: torch.Tensor,
+        a1q_scale: torch.Tensor | None,
     ):
         import flashinfer
 
@@ -467,9 +454,7 @@ def apply(
 class TrtLlmNvFp4ExpertsMonolithic(
     TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic
 ):
-    """
-    Monolithic version of the kernel (router + experts).
-    """
+    """Monolithic version of the kernel (router + experts)."""
 
     def supports_routing_replay_capture(self) -> bool:
         return True
diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
index dfc578b44170..ddd891b1c090 100644
--- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
@@ -171,6 +171,7 @@ def apply(
                 num_experts=self.moe_config.num_local_experts,
                 ep_rank=self.moe_config.ep_rank,
                 ep_size=self.moe_config.ep_size,
+                expert_map=expert_map,
                 gemm1_clamp_limit=self.gemm1_clamp_limit,
             )
         assert self.fused_moe_impl is not None
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index dc619b2b12d2..e0610c878d2c 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -110,8 +110,7 @@ def fused_moe_kernel_gptq_awq(
     use_int4_w4a16: tl.constexpr,
     use_int8_w8a16: tl.constexpr,
 ):
-    """
-    Implements the fused computation for a Mixture of Experts (MOE) using
+    """Implements the fused computation for a Mixture of Experts (MOE) using
     token and expert matrices.
 
     Key Parameters:
@@ -353,8 +352,7 @@ def fused_moe_kernel(
     # Tensor-descriptor path for the A gather and B load in the K-loop.
     USE_TD: tl.constexpr = False,
 ):
-    """
-    Implements the fused computation for a Mixture of Experts (MOE) using
+    """Implements the fused computation for a Mixture of Experts (MOE) using
     token and expert matrices.
 
     Key Parameters:
@@ -1115,15 +1113,13 @@ def get_moe_configs(
     block_n: int | None = None,
     block_k: int | None = None,
 ) -> dict[int, Any] | None:
-    """
-    Return optimized configurations for the fused MoE kernel.
+    """Return optimized configurations for the fused MoE kernel.
 
     The return value will be a dictionary that maps an irregular grid of
     batch sizes to configurations of the fused_moe kernel. To evaluate the
     kernel on a given batch size bs, the closest batch size in the grid should
     be picked and the associated configuration chosen to invoke the kernel.
     """
-
     # Avoid optimizing for the batch invariant case. Use default config
     if envs.VLLM_BATCH_INVARIANT:
         return None
@@ -1187,6 +1183,7 @@ def _ensure_block_size_k_divisible(
 
     Returns:
         A valid BLOCK_SIZE_K that divides size_k and is divisible by group_size.
+
     """
     # Fast path: already valid
     if size_k % block_size_k == 0 and block_size_k % group_size == 0:
@@ -1638,8 +1635,7 @@ def _get_config_quant_dtype(
     use_fp8_w8a8: bool,
     use_int8_w8a8: bool,
 ) -> None | torch.dtype | str:
-    """
-    Get the quantization type based on the quantization strategy flags.
+    """Get the quantization type based on the quantization strategy flags.
     We don't have a quant_config at this point so we need to work backwards.
     A return type of None means no quantization is required because the
     input is unquantized or has been quantized prior to calling
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
index 3a7c6eedfba7..c64aeadcc7a5 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
@@ -59,8 +59,7 @@ def create_weights(
         raise NotImplementedError
 
     def uses_weight_scale_2_pattern(self) -> bool:
-        """
-        Returns True if this quantization method uses 'weight_scale_2' pattern
+        """Returns True if this quantization method uses 'weight_scale_2' pattern
         for per-tensor weight scales (e.g., FP4 variants), False otherwise.
 
         This method should be overridden by subclasses that use the
@@ -75,8 +74,7 @@ def maybe_roundup_sizes(
         act_dtype: torch.dtype,
         moe_parallel_config: FusedMoEParallelConfig,
     ) -> tuple[int, int]:
-        """
-        Given layer hidden size and intermediate size per partition and MoE
+        """Given layer hidden size and intermediate size per partition and MoE
         configurations, round up hidden_size and intermediate_size_per_partition
         if necessary.
 
@@ -93,6 +91,7 @@ def maybe_roundup_sizes(
                 - rounded_hidden_size is the possibly rounded up hidden size.
                 - rounded_intermediate_size_per_partition is the possibly rounded
                   up intermediate size per partition.
+
         """
         from .all2all_utils import maybe_roundup_layer_hidden_size
 
@@ -119,8 +118,7 @@ def skip_forward_padding(self) -> bool:
 
     @property
     def has_unpadded_output(self) -> bool:
-        """
-        Indicates that the hidden_states output might be the unpadded
+        """Indicates that the hidden_states output might be the unpadded
         hidden_states shape rather than the full padded shape.
         """
         return False
@@ -151,8 +149,7 @@ def apply(
         shared_experts: "SharedExperts | None",
         shared_experts_input: torch.Tensor | None,
     ) -> torch.Tensor:
-        """
-        Apply the MoE operation using modular kernels.
+        """Apply the MoE operation using modular kernels.
 
         Args:
             layer: RoutedExperts instance containing weight parameters
@@ -160,9 +157,11 @@ def apply(
             topk_weights: Expert weights from router
             topk_ids: Selected expert IDs from router
             shared_experts_input: Input for shared experts (if any)
+            shared_experts: The shared experts module, if any
 
         Returns:
             Output tensor from routed experts.
+
         """
         raise NotImplementedError
 
@@ -173,15 +172,16 @@ def apply_monolithic(
         router_logits: torch.Tensor,
         input_ids: torch.Tensor | None = None,
     ) -> torch.Tensor | UnfinalizedMoEOutput:
-        """
-        Apply the MoE operation using monolithic kernels.
+        """Apply the MoE operation using monolithic kernels.
 
         Args:
             layer: RoutedExperts instance containing weight parameters
             x: Input tensor
             router_logits: Router logits (routing done internally)
+            input_ids: Token ids, for routers that condition on them
 
         Returns:
             Finalized routed states or a deferred-finalize output.
+
         """
         raise NotImplementedError
diff --git a/vllm/model_executor/layers/fused_moe/hpc_moe.py b/vllm/model_executor/layers/fused_moe/hpc_moe.py
index 410b2a733a6f..c7a078e26251 100644
--- a/vllm/model_executor/layers/fused_moe/hpc_moe.py
+++ b/vllm/model_executor/layers/fused_moe/hpc_moe.py
@@ -122,8 +122,7 @@ def workspace_shapes(
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
         # We use global_num_experts due to how moe_align_block_size handles
         # expert_maps.
-        """
-        Compute the shapes for the temporary and final outputs of the two gemms
+        """Compute the shapes for the temporary and final outputs of the two gemms
         and activation in the fused expert function.  Since the gemms are
         independent, the workspace for the first gemm can be shared with the
         workspace for the last gemm.
diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py
index 4b633dd99b7a..ff620622c1f1 100644
--- a/vllm/model_executor/layers/fused_moe/layer.py
+++ b/vllm/model_executor/layers/fused_moe/layer.py
@@ -137,6 +137,7 @@ def FusedMoEFactory(
     runner_args: dict[str, Any] | None = None,
     routed_experts_cls: type[RoutedExperts] | None = None,
     routed_experts_args: dict[str, Any] | None = None,
+    skip_padding: bool = False,
 ) -> MoERunner:
     """Factory function for creating MoE execution pipeline.
 
@@ -152,6 +153,11 @@ def FusedMoEFactory(
     copy that naming convention here and handle any remapping in the
     load_weights function in each model implementation.
 
+    Args:
+        intermediate_pad: Padding added to the intermediate size, if any.
+        swiglu_alpha: Optional alpha parameter for the SwiGLU activation.
+        swiglu_beta: Optional beta parameter for the SwiGLU activation.
+
     Args:
         num_experts: Number of experts in the model (global count)
         top_k: Number of experts selected for each token
@@ -208,9 +214,11 @@ def FusedMoEFactory(
         runner_args: Additional arguments for runner constructor
         routed_experts_cls: Custom RoutedExperts class (None = use default)
         routed_experts_args: Additional arguments for routed_experts constructor
+        skip_padding: Whether grouped routing should invalidate padding rows.
 
     Returns:
         MoERunner: Configured MoE execution pipeline ready for forward passes
+
     """
     vllm_config = get_current_vllm_config()
 
@@ -302,6 +310,7 @@ def FusedMoEFactory(
             else 1.0,
             e_score_correction_bias=e_score_correction_bias,
             num_fused_shared_experts=num_fused_shared_experts,
+            skip_padding=(skip_padding and moe_parallel_config.use_deepep_v2_kernels),
             # Fused shared-expert slot weight. With apply_routed_scale_to_output
             # the runner scales the combined output by routed_scaling_factor, so
             # the shared slot weight must be 1/routed_scaling_factor for its net
@@ -348,11 +357,13 @@ def FusedMoEFactory(
         moe_backend=vllm_config.kernel_config.moe_backend,
         router_logits_dtype=router_logits_dtype,
         max_num_tokens=max_num_batched_tokens,
+        elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size,
         has_bias=has_bias,
         is_lora_enabled=vllm_config.lora_config is not None,
         activation=moe_activation,
         device=vllm_config.device_config.device,
         routing_method=router.routing_method_type,  # Not ideal
+        has_hash_routing=hash_indices_table is not None,
         swiglu_limit=swiglu_limit,
         swiglu_alpha=swiglu_alpha,
         swiglu_beta=swiglu_beta,
diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py
index 778b95bf47ee..fba16f67b0f6 100644
--- a/vllm/model_executor/layers/fused_moe/modular_kernel.py
+++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py
@@ -84,9 +84,7 @@
 
 
 class FusedMoEActivationFormat(Enum):
-    """
-    The standard activation format (num_tokens, hidden dim).
-    """
+    """The standard activation format (num_tokens, hidden dim)."""
 
     Standard = ("standard",)
     """
@@ -97,9 +95,7 @@ class FusedMoEActivationFormat(Enum):
 
 @dataclass
 class ExpertTokensMetadata:
-    """
-    Metadata regarding expert-token routing.
-    """
+    """Metadata regarding expert-token routing."""
 
     expert_num_tokens: torch.Tensor | None
     expert_num_tokens_cpu: torch.Tensor | None
@@ -122,9 +118,7 @@ def make_from_list(
 
 
 class TopKWeightAndReduce(ABC):
-    """
-    An abstract base class for weight application and reduction implementations.
-    """
+    """An abstract base class for weight application and reduction implementations."""
 
     @abstractmethod
     def apply(
@@ -135,8 +129,7 @@ def apply(
         topk_ids: torch.Tensor,
         apply_router_weight_on_input: bool,
     ) -> torch.Tensor:
-        """
-        Apply topk_weights to the fused_experts_outputs and/or reduce.
+        """Apply topk_weights to the fused_experts_outputs and/or reduce.
         If an output tensor is not passed, it will be created in the
         function.
         """
@@ -185,8 +178,7 @@ def apply(
 
 
 class FusedMoEPrepareAndFinalize(ABC):
-    """
-    An abstract base class for the [Quantize-Prepare] and [Finalize] steps
+    """An abstract base class for the [Quantize-Prepare] and [Finalize] steps
     described above.
 
     There are two variants of this class:
@@ -195,8 +187,7 @@ class FusedMoEPrepareAndFinalize(ABC):
     """
 
     def post_init_setup(self, fused_experts: "FusedMoEExperts"):
-        """
-        Initialize FusedMoEPrepareAndFinalizeModular settings that depend on
+        """Initialize FusedMoEPrepareAndFinalizeModular settings that depend on
         FusedMoEExpertsModular experts object.
         The FusedMoEPrepareAndFinalizeModular implementations that have such
         dependencies may choose to override this function.
@@ -206,16 +197,14 @@ def post_init_setup(self, fused_experts: "FusedMoEExperts"):
     @property
     @abstractmethod
     def activation_format(self) -> FusedMoEActivationFormat:
-        """
-        A property indicating the output format of the activations for the
+        """A property indicating the output format of the activations for the
         'prepare' method.
         """
         raise NotImplementedError
 
     @abstractmethod
     def topk_indices_dtype(self) -> torch.dtype | None:
-        """
-        The PrepareFinalize All2All implementations generally constrain the
+        """The PrepareFinalize All2All implementations generally constrain the
         dtype of the topk_ids they support. This function returns the
         required topk indices dtype so it can be respected.
         Return None if there are no such restrictions.
@@ -224,8 +213,7 @@ def topk_indices_dtype(self) -> torch.dtype | None:
 
     @abstractmethod
     def max_num_tokens_per_rank(self) -> int | None:
-        """
-        Some PrepareFinalize All2All implementations are batched. Meaning,
+        """Some PrepareFinalize All2All implementations are batched. Meaning,
         they can process only as set of tokens at a time. This
         function returns the batch size i.e the maximum number of tokens
         the implementation can process at a time.
@@ -239,31 +227,21 @@ def num_dispatchers(self) -> int:
 
     @abstractmethod
     def output_is_reduced(self) -> bool:
-        """
-        Indicates whether or not the output of finalize is reduced across all
+        """Indicates whether or not the output of finalize is reduced across all
         ranks.
         """
         raise NotImplementedError
 
     def supports_async(self) -> bool:
-        """
-        Indicates whether or not this class implements prepare_async and
+        """Indicates whether or not this class implements prepare_async and
         finalize_async.
         """
         return False
 
-    def on_commit(self) -> None:
-        """
-        Runs after this prepare/finalize has been committed to the active
-        MoE kernel.
-        """
-        return
-
 
 # TODO: pass FusedMoEParallelConfig in as ctor parameter?
 class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):
-    """
-    An abstract base class for the [Quantize-Prepare] and [Finalize] steps
+    """An abstract base class for the [Quantize-Prepare] and [Finalize] steps
     described above for the Modular case.
     """
 
@@ -279,8 +257,7 @@ def prepare(
         quant_config: FusedMoEQuantConfig,
         defer_input_quant: bool,
     ) -> PrepareResultType:
-        """
-        Perform any quantization (and/or) dispatching needed for this kernel.
+        """Perform any quantization (and/or) dispatching needed for this kernel.
         - a1: The (unquantized) input to the MoE layer.
         - topk_ids: The topk ids.
         - topk_weights: The topk weights.
@@ -316,8 +293,7 @@ def prepare_async(
         quant_config: FusedMoEQuantConfig,
         defer_input_quant: bool,
     ) -> tuple[Callable, ReceiverType] | ReceiverType:
-        """
-        Perform any quantization (and/or) dispatching needed for this kernel
+        """Perform any quantization (and/or) dispatching needed for this kernel
         but do not wait for results from other workers.
         - a1: The (unquantized) input to the MoE layer.
         - a1_scale: Optional scales for a1
@@ -367,8 +343,7 @@ def finalize(
         apply_router_weight_on_input: bool,
         weight_and_reduce_impl: TopKWeightAndReduce,
     ) -> None:
-        """
-        Perform any combine plus apply weights and perform a reduction on the
+        """Perform any combine plus apply weights and perform a reduction on the
         fused experts output.
         - output: The output tensor, written in place.  Must be (M, K) shape.
         - fused_expert_output: The unweighted, unreduced output of the fused
@@ -391,8 +366,7 @@ def finalize_async(
         apply_router_weight_on_input: bool,
         weight_and_reduce_impl: TopKWeightAndReduce,
     ) -> tuple[Callable, Callable] | Callable:
-        """
-        Perform any combine plus apply weights and perform a reduction on the
+        """Perform any combine plus apply weights and perform a reduction on the
         fused experts output but do not wait for results from other workers.
         - output: The output tensor, written in place.  Must be (M, K) shape.
         - fused_expert_output: The unweighted, unreduced output of the fused
@@ -426,8 +400,7 @@ def finalize_async(
 
 
 class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):
-    """
-    An abstract base class for the [Quantize-Prepare] and [Finalize] steps
+    """An abstract base class for the [Quantize-Prepare] and [Finalize] steps
     described above for the monolithic case.
     """
 
@@ -442,8 +415,7 @@ def prepare(
         quant_config: FusedMoEQuantConfig,
         defer_input_quant: bool = False,
     ) -> PrepareMonolithicResultType:
-        """
-        Optional method for subclasses compatible with monolithic
+        """Optional method for subclasses compatible with monolithic
         FusedMoEExpertsModular kernels.
 
         Perform any quantization (and/or) dispatching needed for this kernel.
@@ -460,8 +432,7 @@ def prepare(
 
     @abstractmethod
     def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor:
-        """
-        Optional method for subclasses compatible with monolithic
+        """Optional method for subclasses compatible with monolithic
         FusedMoEExpertsModular kernels.
 
         Perform any combine plus apply weights and perform a reduction on the
@@ -492,8 +463,7 @@ def __init__(
         max_num_tokens: int | None = None,
         num_dispatchers: int | None = None,
     ):
-        """
-        moe_config: MoE layer configuration.
+        """moe_config: MoE layer configuration.
         quant_config: Quantization parameters for this experts instance.
         """
         if self.activation_format() == FusedMoEActivationFormat.Standard and (
@@ -528,8 +498,7 @@ def is_monolithic() -> bool:
 
     @property
     def expects_unquantized_inputs(self) -> bool:
-        """
-        Whether or not the PrepareFinalize should defer input quantization
+        """Whether or not the PrepareFinalize should defer input quantization
         in the prepare step. If True, then the Experts kernel will
         execute the input quantization itself.
 
@@ -540,8 +509,7 @@ def expects_unquantized_inputs(self) -> bool:
     @staticmethod
     @abstractmethod
     def activation_format() -> FusedMoEActivationFormat:
-        """
-        A property which is a tuple of the input and output activation formats
+        """A property which is a tuple of the input and output activation formats
         for the 'apply' method.
         """
         raise NotImplementedError
@@ -576,6 +544,8 @@ def _make_reason(reason: str) -> str:
             return False, _make_reason(
                 f"parallel config {moe_config.moe_parallel_config}"
             )
+        elif moe_config.has_hash_routing and cls.is_monolithic():
+            return False, _make_reason("hash routing")
         elif not cls._supports_routing_method(
             moe_config.routing_method, weight_key, activation_key
         ):
@@ -602,8 +572,7 @@ def _make_reason(reason: str) -> str:
     @staticmethod
     @abstractmethod
     def _supports_current_device() -> bool:
-        """
-        Whether the kernel supports the current device type
+        """Whether the kernel supports the current device type
         (compute cability and current platform).
         """
         raise NotImplementedError
@@ -611,8 +580,7 @@ def _supports_current_device() -> bool:
     @staticmethod
     @abstractmethod
     def _supports_no_act_and_mul() -> bool:
-        """
-        Whether the kernel supports act_and_mul=False, i.e.
+        """Whether the kernel supports act_and_mul=False, i.e.
         non-gated MoE models like Nemotron-Nano.
         """
         raise NotImplementedError
@@ -628,16 +596,13 @@ def _supports_quant_scheme(
     @staticmethod
     @abstractmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        """
-        Whether the kernel supports a particular act function.
-        """
+        """Whether the kernel supports a particular act function."""
         raise NotImplementedError
 
     @staticmethod
     @abstractmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
-        """
-        Whether the kernel supports deployment in particular parallel config.
+        """Whether the kernel supports deployment in particular parallel config.
 
         Can be overridden if a kernel does not support EP, SP or some other
         configuration.
@@ -650,8 +615,7 @@ def _supports_routing_method(
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        """
-        Whether the kernel supports a routing method (e.g. GroupedTopK).
+        """Whether the kernel supports a routing method (e.g. GroupedTopK).
 
         Can be overridden by monolithic kernels that execute the router
         in addition to the experts if certain routers are not supported.
@@ -663,8 +627,7 @@ def _supports_router_logits_dtype(
         router_logits_dtype: torch.dtype | None,
         routing_method: RoutingMethodType,
     ) -> bool:
-        """
-        Whether a kernel supports a particular dtype for router logits input.
+        """Whether a kernel supports a particular dtype for router logits input.
 
         Can be overridden by monolithic kernels that execute the router
         in addition to the experts if certain dtypes are not supported.
@@ -673,16 +636,14 @@ def _supports_router_logits_dtype(
 
     @staticmethod
     def _supports_shape(hidden_dim: int) -> bool:
-        """
-        Whether a kernel supports a particular shape. Can be overridden if a kernel
+        """Whether a kernel supports a particular shape. Can be overridden if a kernel
         has specific shape requirements.
         """
         return True
 
     @staticmethod
     def _supports_batch_invariance() -> bool:
-        """
-        Whether the kernel supports batch invariance, i.e. the output does not
+        """Whether the kernel supports batch invariance, i.e. the output does not
         depend on the order of the tokens in the input batch. This is useful
         for determining if the kernel can used with VLLM_BATCH_INVARIANT=1.
         """
@@ -771,17 +732,15 @@ def supports_lora() -> bool:
         return False
 
     def supports_packed_ue8m0_act_scales(self) -> bool:
-        """
-        A flag indicating whether or not this class can process packed ue8m0
+        """A flag indicating whether or not this class can process packed ue8m0
         activation scales.
         """
         return False
 
 
 class FusedMoEExpertsModular(FusedMoEExperts):
-    """
-    An abstract base class for the [Permute-Experts-Unpermute] step described
-        above.
+    """An abstract base class for the [Permute-Experts-Unpermute] step described
+    above.
     """
 
     @staticmethod
@@ -795,8 +754,7 @@ def moe_problem_size(
         w2: torch.Tensor,
         topk_ids: torch.Tensor,
     ) -> tuple[int, int, int, int, int]:
-        """
-        Extract the MoE problem size from the given tensor arguments:
+        """Extract the MoE problem size from the given tensor arguments:
         - a: The hidden states, input to the MoE layer.
         - w1: The first set of expert weights.
         - w2: The second set of expert weights.
@@ -832,9 +790,7 @@ def moe_problem_size(
         return E, M, N, K, topk
 
     def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
-        """
-        Workspace type: The dtype to use for the workspace tensors.
-        """
+        """Workspace type: The dtype to use for the workspace tensors."""
         return act_dtype
 
     @abstractmethod
@@ -849,8 +805,7 @@ def workspace_shapes(
         expert_tokens_meta: ExpertTokensMetadata | None,
         activation: MoEActivation,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        """
-        Compute the shapes for the temporary and final outputs of the two gemms
+        """Compute the shapes for the temporary and final outputs of the two gemms
         and activation in the fused expert function.  Since the gemms are
         independent, the workspace for the first gemm can be shared with the
         workspace for the last gemm.
@@ -880,8 +835,7 @@ def workspace_shapes(
 
     @staticmethod
     def adjust_N_for_activation(N: int, activation: MoEActivation) -> int:
-        """
-        Calculate the output dimension for the activation function.
+        """Calculate the output dimension for the activation function.
 
         For *_no_mul activations (e.g. relu2_no_mul),
         there's no gate/up split, so output size equals input size (N).
@@ -895,6 +849,7 @@ def adjust_N_for_activation(N: int, activation: MoEActivation) -> int:
 
         Returns:
             The output dimension after activation.
+
         """
         return N if not activation.is_gated else N // 2
 
@@ -941,49 +896,50 @@ def apply(
         expert_tokens_meta: ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ) -> None:
-        """
-        This function computes the intermediate result of a Mixture of Experts
+        """This function computes the intermediate result of a Mixture of Experts
         (MoE) layer using two sets of weights, w1 and w2.
 
-        Parameters:
-        - output: (torch.Tensor): The unweighted, unreduced output tensor.
-        - hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE
-          layer.
-        - w1 (torch.Tensor): The first set of expert weights.
-        - w2 (torch.Tensor): The second set of expert weights.
-        - topk_weights: A map of row to expert weights. Some implementations
-          choose to do weight application.
-        - topk_ids (torch.Tensor): A map of row to expert id.
-        - activation (str): The activation function to apply after the first
-          MoE layer.
-        - global_num_experts (int): The total number of experts in the global
-          expert space.
-        - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
-          from the global expert space to the local expert space of the expert
-          parallel shard.
-        - a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be
-          used for a1.  Result of quantization from prepare/finalize and not
-          from the FusedMoEQuantConfig.
-        - workspace13 (torch.Tensor): A scratch tensor used for gemm outputs
-          must be large enough to hold output of either MoE gemm.
-        - workspace2 (torch.Tensor): A scratch tensor used for the activation
-          function.
-        - expert_tokens_meta (Optional[ExpertTokensMetadata]) - An optional
-          ExpertTokensMetadata object containing gpu/cpu tensors
-          as big as the number of local experts with the information about the
-          number of tokens assigned to each local expert.
-        - apply_router_weight_on_input: True if router weights are already
-          applied on the input. This is relevant if the implementation
-          chooses to do weight application.
+        Args:
+            output: (torch.Tensor): The unweighted, unreduced output tensor.
+            hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE
+                layer.
+            w1 (torch.Tensor): The first set of expert weights.
+            w2 (torch.Tensor): The second set of expert weights.
+            topk_weights: A map of row to expert weights. Some implementations
+                choose to do weight application.
+            topk_ids (torch.Tensor): A map of row to expert id.
+            activation (str): The activation function to apply after the first
+                MoE layer.
+            global_num_experts (int): The total number of experts in the global
+                expert space.
+            expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
+                from the global expert space to the local expert space of the
+                expert parallel shard.
+            a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be
+                used for a1. Result of quantization from prepare/finalize and not
+                from the FusedMoEQuantConfig.
+            a2_scale (Optional[torch.Tensor]): Optional quantized scale to be
+                used for the second gemm's activations.
+            workspace13 (torch.Tensor): A scratch tensor used for gemm outputs
+                must be large enough to hold output of either MoE gemm.
+            workspace2 (torch.Tensor): A scratch tensor used for the activation
+                function.
+            expert_tokens_meta (Optional[ExpertTokensMetadata]): An optional
+                ExpertTokensMetadata object containing gpu/cpu tensors
+                as big as the number of local experts with the information about
+                the number of tokens assigned to each local expert.
+            apply_router_weight_on_input: True if router weights are already
+                applied on the input. This is relevant if the implementation
+                chooses to do weight application.
+
         """
         raise NotImplementedError
 
 
 class FusedMoEExpertsMonolithic(FusedMoEExperts):
-    """
-    An abstract base class for the [Permute-Experts-Unpermute] step described
-        above, but with the monolithic interface (accepts router logits
-        rather than topk ids and weights).
+    """An abstract base class for the [Permute-Experts-Unpermute] step described
+    above, but with the monolithic interface (accepts router logits
+    rather than topk ids and weights).
     """
 
     @staticmethod
@@ -992,8 +948,7 @@ def _supports_routing_method(
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        """
-        Whether the kernel supports a routing method (e.g. GroupedTopK).
+        """Whether the kernel supports a routing method (e.g. GroupedTopK).
 
         Monolithic kernels should explicitly opt-in to support.
         """
@@ -1004,8 +959,7 @@ def _supports_router_logits_dtype(
         router_logits_dtype: torch.dtype | None,
         routing_method: RoutingMethodType,
     ) -> bool:
-        """
-        Whether the kernel supports a dtype for router logits.
+        """Whether the kernel supports a dtype for router logits.
 
         Modular kernels should opt-in to support.
         """
@@ -1090,8 +1044,7 @@ def apply(
         routed_scaling_factor: float | None = None,
         topk_group: int | None = None,
     ) -> torch.Tensor | UnfinalizedMoEOutput:
-        """
-        Same as apply(), except uses router_logits as opposed
+        """Same as ``FusedMoEExperts.apply``, except uses router_logits as opposed
         to the topk_ids and topk_weights. This is useful for kernels
         with fused router and fused_experts (e.g. FLASHINFER_TRTLLM).
         """
@@ -1135,8 +1088,7 @@ def _allocate_buffers(
         expert_tokens_meta: ExpertTokensMetadata | None,
         activation: MoEActivation,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Allocate temporary and output buffers for the fused experts op.
+        """Allocate temporary and output buffers for the fused experts op.
         Inputs:
         - out_dtype: output type of workspace and output tensors.
         - device: the device of the workspace and output tensors.
@@ -1230,11 +1182,9 @@ def _prepare(
         torch.Tensor,
         torch.Tensor,
     ]:
-        """
-        The _prepare method is a wrapper around self.prepare_finalize.prepare
+        """The _prepare method is a wrapper around self.prepare_finalize.prepare
         that handles DBO and async.
         """
-
         if not self.prepare_finalize.supports_async():
             # We shouldn't be running an a2a kernel that doesn't
             # support async prepare/finalize
@@ -1399,17 +1349,24 @@ def _finalize(
         shared_experts: SharedExperts | None,
         shared_experts_input: torch.Tensor | None,
     ) -> torch.Tensor:
-        """
-        The _finalize method is a wrapper around self.prepare_finalize.finalize
+        """The _finalize method is a wrapper around self.prepare_finalize.finalize
         that handles DBO, async and shared expert overlap.
 
         Args:
+            output: Tensor the finalized result is written into.
+            fused_out: The unweighted, unreduced output of the fused experts.
+            hidden_states: The input tensor to the MoE layer.
+            topk_weights: A map of row to expert weights.
+            topk_ids: A map of row to expert id.
+            apply_router_weight_on_input: True if the router weights were
+                already applied on the input.
             shared_experts: SharedExperts | None. The shared experts if any.
             shared_experts_input: Optional separate input for shared experts.
                 When latent MoE is used, hidden_states is the latent-projected
                 tensor (smaller dimension) used by routed experts, while
                 shared_experts_input is the original hidden_states (full
                 dimension) needed by the shared expert MLP.
+
         """
         if not self.prepare_finalize.supports_async():
             assert not dbo_enabled()
@@ -1470,33 +1427,34 @@ def apply(
         shared_experts: SharedExperts | None = None,
         shared_experts_input: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        This function computes a Mixture of Experts (MoE) layer using two sets
+        """This function computes a Mixture of Experts (MoE) layer using two sets
         of weights, w1 and w2, and top-k gating mechanism.
 
-        Parameters:
-        - hidden_states: (torch.Tensor): The input tensor to the MoE layer.
-        - w1 (torch.Tensor): The first set of expert weights.
-        - w2 (torch.Tensor): The second set of expert weights.
-        - topk_weights (torch.Tensor): The topk weights applied at the end of the layer.
-        - topk_ids (torch.Tensor): A map of row to expert id.
-        - activation (MoEActivation): The activation function to apply after the first
-          MoE layer.
-        - global_num_experts (int): The total number of experts in the global
-          expert space.
-        - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
-          from the global expert space to the local expert space of the expert
-          parallel shard.
-        - apply_router_weight_on_input (bool): When true, the topk weights are
-          applied directly on the inputs. This is only applicable when topk is
-          1.
-        - shared_experts: SharedExperts | None. The shared experts if any.
-        - shared_experts_input (Optional[torch.Tensor]): Optional separate
-          input for shared experts. For latent MoE, this is the original
-          hidden_states before latent projection.
+        Args:
+            hidden_states: (torch.Tensor): The input tensor to the MoE layer.
+            w1 (torch.Tensor): The first set of expert weights.
+            w2 (torch.Tensor): The second set of expert weights.
+            topk_ids (torch.Tensor): A map of row to expert id.
+            topk_weights (torch.Tensor): The topk weights applied at the end of
+                the layer.
+            activation (MoEActivation): The activation function to apply after
+                the first MoE layer.
+            global_num_experts (int): The total number of experts in the global
+                expert space.
+            expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
+                from the global expert space to the local expert space of the
+                expert parallel shard.
+            apply_router_weight_on_input (bool): When true, the topk weights are
+                applied directly on the inputs. This is only applicable when
+                topk is 1.
+            shared_experts: SharedExperts | None. The shared experts if any.
+            shared_experts_input (Optional[torch.Tensor]): Optional separate
+                input for shared experts. For latent MoE, this is the original
+                hidden_states before latent projection.
 
         Returns:
-        - torch.Tensor: The output tensor after applying the MoE layer.
+            torch.Tensor: The output tensor after applying the MoE layer.
+
         """
         output = torch.empty_like(hidden_states)
 
@@ -1578,12 +1536,10 @@ def apply(
         routed_scaling_factor: float | None = None,
         topk_group: int | None = None,
     ) -> torch.Tensor | UnfinalizedMoEOutput:
-        """
-        Same as forward(), except uses router_logits as opposed
+        """Same as forward(), except uses router_logits as opposed
         to the topk_ids and topk_weights. This is used for kernels
         that have fused router + experts (e.g. FLASHINFER_TRTLLM).
         """
-
         a1q, a1q_scale, router_logits = self.prepare_finalize.prepare(
             hidden_states,
             router_logits=router_logits,
@@ -1681,8 +1637,7 @@ def supports_lora(self) -> bool:
         return self.fused_experts.supports_lora()
 
     def _post_init_setup(self):
-        """
-        Resolve any leftover setup dependencies between self.prepare_finalize
+        """Resolve any leftover setup dependencies between self.prepare_finalize
         and self.fused_experts here.
         """
         self.prepare_finalize.post_init_setup(self.impl.fused_experts)
@@ -1692,8 +1647,7 @@ def _post_init_setup(self):
         )
 
     def output_is_reduced(self) -> bool:
-        """
-        Indicates whether or not the output of fused MoE kernel
+        """Indicates whether or not the output of fused MoE kernel
         is reduced across all ranks.
         """
         return self.prepare_finalize.output_is_reduced()
diff --git a/vllm/model_executor/layers/fused_moe/moe_align_block_size.py b/vllm/model_executor/layers/fused_moe/moe_align_block_size.py
index 7fc8bfcf824d..663475306cc3 100644
--- a/vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+++ b/vllm/model_executor/layers/fused_moe/moe_align_block_size.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from typing import Literal, overload
+
 import torch
 
 from vllm import _custom_ops as ops
@@ -8,6 +10,7 @@
 from vllm.utils.math_utils import round_up
 
 
+@overload
 def moe_align_block_size(
     topk_ids: torch.Tensor,
     block_size: int,
@@ -15,9 +18,53 @@ def moe_align_block_size(
     expert_map: torch.Tensor | None = None,
     pad_sorted_ids: bool = False,
     ignore_invalid_experts: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Aligns the token distribution across experts to be compatible with block
+    *,
+    return_scatter_idx: Literal[False] = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
+
+
+@overload
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+    expert_map: torch.Tensor | None = None,
+    pad_sorted_ids: bool = False,
+    ignore_invalid_experts: bool = False,
+    *,
+    return_scatter_idx: Literal[True],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ...
+
+
+@overload
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+    expert_map: torch.Tensor | None = None,
+    pad_sorted_ids: bool = False,
+    ignore_invalid_experts: bool = False,
+    *,
+    return_scatter_idx: bool,
+) -> (
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+): ...
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+    expert_map: torch.Tensor | None = None,
+    pad_sorted_ids: bool = False,
+    ignore_invalid_experts: bool = False,
+    return_scatter_idx: bool = False,
+) -> (
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+):
+    """Aligns the token distribution across experts to be compatible with block
     size for matrix multiplication.
 
     Note: In the case of expert_parallel, moe_align_block_size initially
@@ -26,23 +73,25 @@ def moe_align_block_size(
     the current GPU rank as -1 so the MoE matmuls could skip those blocks.
     This requires the num_experts input arg to be the num global experts.
 
-    Parameters:
-    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
-        top-k expert indices for each token.
-    - block_size: The block size used in block matrix multiplication.
-    - num_experts: The total number of experts.
-    - expert_map: A tensor of shape [num_experts] that maps the expert index
-        from the global space to the local index space of the current
-        expert parallel shard. If the expert is not in the current expert
-        parallel shard, the mapping is set to -1.
-    - pad_sorted_ids: A flag indicating whether the sorted_token_ids length
-        should be padded to a multiple of block_size,
-    - ignore_invalid_experts: A flag indicating whether to ignore invalid
-        experts. When False, all expert_ids in topk_ids will participate in
-        counting and ranking, but invalid experts in expert_ids will be marked
-        as -1. When True, all invalid expert_ids in topk_ids will be ignored
-        and will not participate in counting or ranking, and there will be no
-        -1 in expert_ids.
+    Args:
+        topk_ids: A tensor of shape [total_tokens, top_k] representing the
+            top-k expert indices for each token.
+        block_size: The block size used in block matrix multiplication.
+        num_experts: The total number of experts.
+        expert_map: A tensor of shape [num_experts] that maps the expert index
+            from the global space to the local index space of the current
+            expert parallel shard. If the expert is not in the current expert
+            parallel shard, the mapping is set to -1.
+        pad_sorted_ids: A flag indicating whether the sorted_token_ids length
+            should be padded to a multiple of block_size,
+        ignore_invalid_experts: A flag indicating whether to ignore invalid
+            experts. When False, all expert_ids in topk_ids will participate in
+            counting and ranking, but invalid experts in expert_ids will be marked
+            as -1. When True, all invalid expert_ids in topk_ids will be ignored
+            and will not participate in counting or ranking, and there will be no
+            -1 in expert_ids.
+        return_scatter_idx: Whether to additionally return a masked identity
+            mapping for the original routed rows.
 
     Returns:
     - sorted_token_ids: A tensor containing the sorted token indices according
@@ -50,6 +99,10 @@ def moe_align_block_size(
     - expert_ids: A tensor indicating the assigned expert index for each block.
     - num_tokens_post_padded: The total number of tokens after padding,
         ensuring divisibility by block_size.
+    - scatter_idx: Only returned when return_scatter_idx=True. An int32 tensor
+        shaped [topk_ids.numel(), 1], containing the original routed row index
+        for valid routes and -1 otherwise. With expert_map, this requires
+        ignore_invalid_experts=True.
 
     This function pads the number of tokens that each expert needs to process
     so that it is divisible by block_size.
@@ -70,7 +123,17 @@ def moe_align_block_size(
         the subsequent matrix multiplication.
     - The padding ensures that the total number of tokens is now divisible
         by block_size for proper block matrix operations.
+
     """
+    if return_scatter_idx and expert_map is not None and not ignore_invalid_experts:
+        raise ValueError("scatter_idx with expert_map requires ignore_invalid_experts")
+    scatter_idx = None
+    if return_scatter_idx:
+        scatter_idx = torch.empty(
+            (topk_ids.numel(), 1),
+            dtype=torch.int32,
+            device=topk_ids.device,
+        )
     max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
     if pad_sorted_ids:
         max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
@@ -95,19 +158,21 @@ def moe_align_block_size(
         expert_ids,
         num_tokens_post_pad,
         expert_map if ignore_invalid_experts else None,
+        scatter_idx,
     )
 
     if expert_map is not None and not ignore_invalid_experts:
         expert_ids = expert_map[expert_ids]
 
+    if scatter_idx is not None:
+        return sorted_ids, expert_ids, num_tokens_post_pad, scatter_idx
     return sorted_ids, expert_ids, num_tokens_post_pad
 
 
 def batched_moe_align_block_size(
     max_tokens_per_batch: int, block_size: int, expert_num_tokens: torch.Tensor
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Given num_batches, max_tokens_per_batch, block_size and the number of
+    """Given num_batches, max_tokens_per_batch, block_size and the number of
     valid-tokens in each batch, prepare sorted_token_ids, expert_ids and
     num_tokens_post_pad. sorted_token_ids, expert_ids and num_tokens_post_pad
     have the same semantics as in moe_align_block_size.
@@ -115,12 +180,12 @@ def batched_moe_align_block_size(
     This function is intended to be a drop in replacement for
     moe_align_batch_size for the batched case.
 
-    Parameters:
-    - max_tokens_per_batch (int): Number of tokens in each batch (both
-        valid and invalid).
-    - block_size (int): block_size to align the data to.
-    - expert_num_tokens (torch.Tensor): expert_num_tokens[i], indicates
-        the number of valid tokens in batch i.
+    Args:
+        max_tokens_per_batch (int): Number of tokens in each batch (both
+            valid and invalid).
+        block_size (int): block_size to align the data to.
+        expert_num_tokens (torch.Tensor): expert_num_tokens[i], indicates
+            the number of valid tokens in batch i.
 
     Returns:
     - sorted_token_ids (torch.Tensor): Torch tensor of size
@@ -132,6 +197,7 @@ def batched_moe_align_block_size(
     - num_tokens_post_pad (torch.Tensor): Torch tensor of size 1
         indicating the number of valid blocks with actual data to
         process. This is represented in terms of num tokens.
+
     Example:
     Let num_batches=5, max_tokens_per_batch=8, block_size=4, and
     expert_num_tokens=[2, 3, 0, 6, 8]. This expert_num_tokens tensor
@@ -166,8 +232,8 @@ def batched_moe_align_block_size(
 
       num_tokens_post_pad will be 24 as sorted_token_ids has valid entries
       until 24.
-    """
 
+    """
     B = expert_num_tokens.size(0)
     device = expert_num_tokens.device
 
diff --git a/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py b/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py
index d8ad9b94557d..8d16a419756a 100644
--- a/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py
+++ b/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py
@@ -92,8 +92,7 @@ def moe_fused_mul_sum(
     expert_map: torch.Tensor | None = None,
     num_valid_tokens: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """
-    Fused kernel for MoE (Mixture of Experts) to perform weighted summation
+    """Fused kernel for MoE (Mixture of Experts) to perform weighted summation
     of expert outputs.
 
     Args:
@@ -122,6 +121,7 @@ def moe_fused_mul_sum(
     Returns:
         The fused weighted sum of expert outputs.
         Shape: (num_tokens, hidden_size).
+
     """
     assert inputs.ndim == 3
     assert topk_weights.ndim == 2
diff --git a/vllm/model_executor/layers/fused_moe/moe_output.py b/vllm/model_executor/layers/fused_moe/moe_output.py
index df1a2685ecbe..04081cf6d17b 100644
--- a/vllm/model_executor/layers/fused_moe/moe_output.py
+++ b/vllm/model_executor/layers/fused_moe/moe_output.py
@@ -57,8 +57,8 @@ def convert_flashinfer_moe_output(
 
     Raises:
         ValueError: If FlashInfer returns an unexpected layout.
-    """
 
+    """
     if num_tokens < 0:
         raise ValueError("num_tokens must be non-negative.")
     if top_k <= 0:
diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
index ad9fb509a1f7..d83152fc696a 100644
--- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
@@ -5,6 +5,11 @@
 
 import torch
 
+from vllm.v1.worker.workspace import (
+    current_workspace_manager,
+    is_workspace_manager_initialized,
+)
+
 
 @dataclass
 class MoEPermuteScratch:
@@ -102,6 +107,79 @@ def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor:
         return topk_ids_int32
 
 
+def moe_prepare_scatter(
+    topk_ids: torch.Tensor,
+    expert_map: torch.Tensor | None,
+    scratch: MoEPermuteScratch,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Generate expert offsets and shared scatter/unpermute destination indices."""
+    assert topk_ids.device == scratch.device
+    assert topk_ids.size(0) <= scratch.max_num_tokens
+    assert topk_ids.size(1) == scratch.topk
+    n_token, topk = topk_ids.shape
+    expanded_rows = topk_ids.numel()
+    inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk)
+    if expanded_rows == 0:
+        scratch.expert_first_token_offset.zero_()
+        return scratch.expert_first_token_offset, inverse
+    torch.ops._moe_C.moe_prepare_scatter(
+        scratch.prepare_topk_ids(topk_ids),
+        scratch.token_expert_indices_view(n_token),
+        expert_map,
+        scratch.num_experts,
+        scratch.num_local_experts,
+        scratch.expert_first_token_offset,
+        inverse,
+        scratch.sort_workspace,
+        scratch.permuted_experts_id[:expanded_rows].view(n_token, topk),
+        scratch.sorted_row_idx[:expanded_rows].view(n_token, topk),
+        scratch.topk_ids_for_sort[:expanded_rows].view(n_token, topk),
+    )
+    return scratch.expert_first_token_offset, inverse
+
+
+def get_moe_permute_scratch(
+    *,
+    max_num_tokens: int,
+    topk: int,
+    num_experts: int,
+    num_local_experts: int,
+    device: torch.device,
+    hidden_size: int | None = None,
+    hidden_dtype: torch.dtype | None = None,
+) -> MoEPermuteScratch:
+    """Share scratch across sequential layers in the current ubatch and lane.
+
+    Without a workspace manager, allocate new scratch for each call.
+    """
+
+    def create_scratch() -> MoEPermuteScratch:
+        return MoEPermuteScratch(
+            max_num_tokens=max_num_tokens,
+            topk=topk,
+            num_experts=num_experts,
+            num_local_experts=num_local_experts,
+            device=device,
+            hidden_size=hidden_size,
+            hidden_dtype=hidden_dtype,
+        )
+
+    if not is_workspace_manager_initialized():
+        return create_scratch()
+
+    key = (
+        MoEPermuteScratch,
+        max_num_tokens,
+        topk,
+        num_experts,
+        num_local_experts,
+        device,
+        hidden_size,
+        hidden_dtype,
+    )
+    return current_workspace_manager().get_persistent_resource(key, create_scratch)
+
+
 def moe_permute(
     hidden_states: torch.Tensor,
     a1q_scale: torch.Tensor | None,
@@ -112,20 +190,24 @@ def moe_permute(
     permuted_hidden_states: torch.Tensor | None = None,
     scratch: MoEPermuteScratch | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    This function expands and permutes activation to gather uncontinuous tokens
+    """This function expands and permutes activation to gather uncontinuous tokens
       for each expert.
-    Parameters:
-    - hidden_states (torch.Tensor): The input tensor to the MoE layer.
-    - a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states
-    - topk_ids (torch.Tensor): topk expert route id for each token.
-    - n_expert (int): The number of expert.
-    - n_local_expert (int): The number of expert in current EP rank.
-    - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
-        from the global expert space to the local expert space of the expert
-        parallel shard.
-    - permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor.
-        If None, the output tensor will be created in this function.
+
+    Args:
+        hidden_states (torch.Tensor): The input tensor to the MoE layer.
+        a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states
+        topk_ids (torch.Tensor): topk expert route id for each token.
+        n_expert (int): The number of expert.
+        n_local_expert (int): The number of expert in current EP rank.
+        expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
+            from the global expert space to the local expert space of the expert
+            parallel shard.
+        permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor.
+            If None, the output tensor will be created in this function.
+        scratch (Optional[MoEPermuteScratch]): Optional preallocated scratch
+            buffers. Validated against hidden_states and topk_ids when given,
+            otherwise the buffers are allocated in this function.
+
     Returns:
     - permuted_hidden_states (torch.Tensor): permuted activation.
     - a1q_scale (Optional[torch.Tensor]): permuted quant scale for hidden_states
@@ -134,6 +216,7 @@ def moe_permute(
        of each expert for standard grouped gemm.
     - inv_permuted_idx (torch.Tensor): idx map for moe_unpermute.
     - permuted_idx (torch.Tensor): idx map from hidden to permuted_hidden.
+
     """
     n_token, n_hidden = hidden_states.size()
     topk = topk_ids.size(1)
@@ -249,19 +332,21 @@ def moe_unpermute(
     inv_permuted_idx: torch.Tensor,
     expert_first_token_offset: torch.Tensor | None = None,
 ) -> None:
-    """
-    This function expands and permutes activation to gathering uncontinuous
+    """This function expands and permutes activation to gathering uncontinuous
       tokens for each expert.
-    Parameters:
-    - out (torch.Tensor): output tensor
-    - permuted_hidden_states (torch.Tensor): permuted activation.
-    - topk_weights (torch.Tensor): topk expert route weight for each token.
-    - inv_permuted_idx (torch.Tensor): row idx map for moe_unpermute.
-    - expert_first_token_offset (Optional[torch.Tensor]): offset of the first
-      token of each expert for grouped gemm.
+
+    Args:
+        out (torch.Tensor): output tensor
+        permuted_hidden_states (torch.Tensor): permuted activation.
+        topk_weights (torch.Tensor): topk expert route weight for each token.
+        inv_permuted_idx (torch.Tensor): row idx map for moe_unpermute.
+        expert_first_token_offset (Optional[torch.Tensor]): offset of the first
+            token of each expert for grouped gemm.
+
     Returns:
     - hidden_states (torch.Tensor): The reduced and unpermuted activation
       tensor.
+
     """
     topk = topk_weights.size(1)
     n_hidden = permuted_hidden_states.size(-1)
diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py
index abe47373dea4..5219fff617d3 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py
@@ -31,11 +31,14 @@
     prepare_fp8_moe_layer_for_marlin,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
     QuantKey,
+    create_fp8_quant_key,
     kFp8Dynamic128Sym,
     kFp8Static128BlockSym,
 )
 from vllm.platforms import current_platform
+from vllm.utils.math_utils import round_up
 
 logger = init_logger(__name__)
 
@@ -54,7 +57,9 @@ class Fp8MoeBackend(Enum):
     VLLM_CUTLASS = "VLLM_CUTLASS"
     BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"
     XPU = "XPU"
+    # CPU FP8 W8A16 (BF16 activations); CPU_W8A8 needs native AMX-FP8.
     CPU = "CPU"
+    CPU_W8A8 = "CPU_W8A8"
     HPC = "HPC"
     # Dequantize-to-BF16 emulation for MXFP8 on devices without a native
     # MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here.
@@ -72,12 +77,10 @@ def _get_priority_backends(
     weight_key: QuantKey | None,
     activation_key: QuantKey | None,
 ) -> list[Fp8MoeBackend]:
-    """
-    Get available backends in priority order based on platform and config.
+    """Get available backends in priority order based on platform and config.
 
     This function can be extended to become more complex as needed.
     """
-
     _AVAILABLE_BACKENDS = [
         Fp8MoeBackend.AITER,
         Fp8MoeBackend.FLASHINFER_TRTLLM,
@@ -91,6 +94,7 @@ def _get_priority_backends(
         Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
         Fp8MoeBackend.BATCHED_TRITON,
         Fp8MoeBackend.XPU,
+        Fp8MoeBackend.CPU_W8A8,
         Fp8MoeBackend.CPU,
         Fp8MoeBackend.HPC,
     ]
@@ -128,8 +132,10 @@ def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> Non
         _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.XPU)
 
     if current_platform.is_cpu():
-        # CPU platform uses FP8 W8A16 fused MoE kernel.
+        # W8A8 first: it falls through to the W8A16 backend whenever the
+        # hardware (AMX-FP8) or the config isn't supported.
         _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU)
+        _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU_W8A8)
 
     return _AVAILABLE_BACKENDS
 
@@ -237,6 +243,13 @@ def backend_to_kernel_cls(
 
         return [CPUExpertsFp8]
 
+    elif backend == Fp8MoeBackend.CPU_W8A8:
+        from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
+            CPUExpertsFp8W8A8,
+        )
+
+        return [CPUExpertsFp8W8A8]
+
     elif backend == Fp8MoeBackend.HPC:
         from vllm.model_executor.layers.fused_moe.hpc_moe import (
             HPCExperts,
@@ -273,8 +286,7 @@ def refine_fp8_moe_block_shape(
     config: FusedMoEConfig,
     weight_block_size: list[int],
 ) -> list[int] | None:
-    """
-    Compute a refined block shape for block-quantized FP8 MoE weights whose
+    """Compute a refined block shape for block-quantized FP8 MoE weights whose
     checkpoint blocks cannot be sharded exactly across TP ranks.
 
     TP shards the intermediate dim of the expert weights, so a per-shard size
@@ -302,17 +314,89 @@ def refine_fp8_moe_block_shape(
     return [refine, refine]
 
 
+def pad_tp_shard_to_weight_blocks(
+    config: FusedMoEConfig,
+    weight_block_size: list[int],
+) -> bool:
+    """Pad the TP shard to whole checkpoint blocks, keeping scales rank-local."""
+    block_n, block_k = weight_block_size
+    if (
+        block_n != block_k
+        or config.tp_size == 1
+        or config.intermediate_size_per_partition % block_n == 0
+    ):
+        return False
+    if (
+        config.intermediate_size % block_n != 0
+        or config.hidden_dim % block_n != 0
+        or config.ep_size != 1
+        or config.is_lora_enabled
+        or config.has_bias
+    ):
+        raise ValueError(
+            f"Block-aligned FP8 TP sharding requires {block_n}-aligned "
+            "global expert dimensions, pure TP, and no LoRA or "
+            "expert bias."
+        )
+    config.tp_shard_with_padding = True
+    config.intermediate_size_per_partition = round_up(
+        config.intermediate_size_per_partition, block_n
+    )
+    return True
+
+
+def resolve_fp8_moe_weight_block_shape(
+    config: FusedMoEConfig,
+    weight_block_size: list[int],
+    activation_key: QuantKey,
+    is_checkpoint_fp8_serialized: bool,
+) -> tuple[list[int], tuple[int, int] | None]:
+    """Return the TP-adapted block shape and refine factor:
+    refine if kernels allow, else pad to the TP shard."""
+    refined_shape = refine_fp8_moe_block_shape(config, weight_block_size)
+    if is_checkpoint_fp8_serialized and config.moe_backend != "auto":
+        kernel_classes = backend_to_kernel_cls(map_fp8_backend(config.moe_backend))
+        can_refine = refined_shape is not None and any(
+            k_cls._supports_quant_scheme(
+                create_fp8_quant_key(
+                    static=True, group_shape=GroupShape(*refined_shape)
+                ),
+                activation_key,
+            )
+            for k_cls in kernel_classes
+        )
+        if not can_refine and pad_tp_shard_to_weight_blocks(config, weight_block_size):
+            logger.info_once(
+                "FP8 %s TP loading uses complete checkpoint blocks: "
+                "local allocation %d, without weight requantization.",
+                config.moe_backend,
+                config.intermediate_size_per_partition,
+            )
+            return weight_block_size, None
+    if refined_shape is None:
+        return weight_block_size, None
+    logger.info_once(
+        "FP8 MoE block scales refined from %s to %s to fit "
+        "the TP-sharded intermediate size %d.",
+        str(weight_block_size),
+        str(refined_shape),
+        config.intermediate_size_per_partition,
+    )
+    return refined_shape, (
+        weight_block_size[0] // refined_shape[0],
+        weight_block_size[1] // refined_shape[1],
+    )
+
+
 def select_fp8_moe_backend(
     config: FusedMoEConfig,
     weight_key: QuantKey | None,
     activation_key: QuantKey | None,
     allow_vllm_cutlass: bool = False,
 ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts] | None]:
-    """
-    Select the primary FP8 MoE backend
+    """Select the primary FP8 MoE backend
     Note: Shape-specific fallbacks may still occur at runtime.
     """
-
     # NOTE: the kernels are selected in the following order.
     AVAILABLE_BACKENDS = _get_priority_backends(config, weight_key, activation_key)
 
@@ -519,7 +603,7 @@ def convert_to_fp8_moe_kernel_format(
         w13.is_shuffled = True
         w2.is_shuffled = True
     elif fp8_backend == Fp8MoeBackend.HUMMING:
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             convert_to_humming_moe_kernel_format,
         )
 
@@ -580,6 +664,14 @@ def convert_to_fp8_moe_kernel_format(
         )
 
         w13, w2 = prepare_fp8_moe_layer_for_cpu(w13, w2)
+    elif fp8_backend == Fp8MoeBackend.CPU_W8A8:
+        from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
+            prepare_fp8_w8a8_moe_layer_for_cpu,
+        )
+
+        w13, w13_scale, w2, w2_scale = prepare_fp8_w8a8_moe_layer_for_cpu(
+            w13, w2, w13_scale, w2_scale
+        )
     else:
         if fp8_backend not in [
             Fp8MoeBackend.TRITON,
@@ -614,8 +706,7 @@ def make_fp8_moe_quant_config(
     gemm1_beta: float | None = None,
     layer: torch.nn.Module | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Create FusedMoEQuantConfig for the specified FP8 Backend.
+    """Create FusedMoEQuantConfig for the specified FP8 Backend.
     The FusedMoEQuantConfig holds the scales that are used
     at runtime by the Modular Kernel abstraction.
 
@@ -626,8 +717,15 @@ def make_fp8_moe_quant_config(
     In a future PR, we will have this function should be
     a method of the modular kernel itself.
     """
+    if fp8_backend == Fp8MoeBackend.CPU_W8A8:
+        return fp8_w8a8_moe_quant_config(
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            a1_scale=a1_scale,
+            block_shape=block_shape,
+        )
 
-    # MARLIN and CPU are mixed precision W8A16 config.
+    # MARLIN and CPU (W8A16) are mixed precision W8A16 configs.
     if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU:
         return fp8_w8a16_moe_quant_config(
             w1_scale=w1_scale,
@@ -641,7 +739,7 @@ def make_fp8_moe_quant_config(
         )
     elif fp8_backend == Fp8MoeBackend.HUMMING:
         from vllm.model_executor.layers.fused_moe import RoutedExperts
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             get_humming_moe_quant_config,
         )
 
diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py
index cd15dceaf5c4..29e9888098c1 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/int8.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py
@@ -37,9 +37,7 @@ class Int8MoeBackend(Enum):
 def _get_priority_backends(
     moe_config: FusedMoEConfig,
 ) -> list[Int8MoeBackend]:
-    """
-    Get available backends in priority order based on platform and config.
-    """
+    """Get available backends in priority order based on platform and config."""
     return [
         Int8MoeBackend.TRITON,
         Int8MoeBackend.HUMMING,
@@ -73,10 +71,16 @@ def backend_to_kernel_cls(
         from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
             ArmCPUExpertsInt8,
             CPUExpertsInt8,
+            PowerCPUExpertsInt8,
             ZenCPUExpertsInt8,
         )
 
-        return [ZenCPUExpertsInt8, ArmCPUExpertsInt8, CPUExpertsInt8]
+        return [
+            ZenCPUExpertsInt8,
+            PowerCPUExpertsInt8,
+            ArmCPUExpertsInt8,
+            CPUExpertsInt8,
+        ]
     else:
         raise ValueError(f"Unknown Int8 MoE backend: {backend.value}")
 
@@ -100,11 +104,9 @@ def select_int8_moe_backend(
     weight_key: QuantKey | None = kInt8StaticChannelSym,
     activation_key: QuantKey | None = kInt8DynamicTokenSym,
 ) -> tuple[Int8MoeBackend, type[mk.FusedMoEExperts]]:
-    """
-    Select the primary Int8 MoE backend.
+    """Select the primary Int8 MoE backend.
     Note: Shape-specific fallbacks may still occur at runtime.
     """
-
     AVAILABLE_BACKENDS = _get_priority_backends(config)
 
     activation_format = (
@@ -191,7 +193,7 @@ def make_int8_moe_quant_config(
 
     if int8_backend == Int8MoeBackend.HUMMING:
         from vllm.model_executor.layers.fused_moe import RoutedExperts
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             get_humming_moe_quant_config,
         )
 
@@ -252,7 +254,7 @@ def convert_to_int8_moe_kernel_format(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Convert INT8 MoE weights to backend-specific kernel format."""
     if int8_backend == Int8MoeBackend.HUMMING:
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             convert_to_humming_moe_kernel_format,
         )
 
diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
index 899f9bc3d94a..ad591d9ea7e4 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
@@ -7,6 +7,7 @@
 import torch
 from compressed_tensors.quantization import (
     QuantizationArgs,
+    QuantizationStrategy,
 )
 
 import vllm._custom_ops as ops
@@ -41,8 +42,10 @@
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
+    pack_quantized_values_into_int32,
 )
 from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
 
 logger = init_logger(__name__)
 
@@ -56,6 +59,7 @@ class WNA16MoEBackend(Enum):
     TRITON = "TRITON"
     XPU = "XPU"
     EMULATION = "EMULATION"
+    RDNA3 = "RDNA3"
 
 
 def backend_to_kernel_cls(
@@ -100,20 +104,26 @@ def backend_to_kernel_cls(
         )
 
         return [Int4EmulationTritonExperts]
+    elif backend == WNA16MoEBackend.RDNA3:
+        from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+            Rdna3WNA16Experts,
+        )
+
+        return [Rdna3WNA16Experts]
     else:
         raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
 
 
 def _get_priority_backends() -> list[WNA16MoEBackend]:
-    """
-    Get available backends in priority order based on platform and config.
-    """
+    """Get available backends in priority order based on platform and config."""
     if current_platform.is_cpu():
         return [WNA16MoEBackend.CPU]
     if current_platform.is_xpu():
         return [WNA16MoEBackend.XPU]
 
     return [
+        # Native HIP kernel, gated on gfx1100 by _supports_current_device().
+        WNA16MoEBackend.RDNA3,
         WNA16MoEBackend.FLASHINFER_TRTLLM,
         WNA16MoEBackend.MARLIN,
         WNA16MoEBackend.BATCHED_MARLIN,
@@ -134,6 +144,18 @@ def _backend_incompatibility_reason(
     if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias):
         return "zero points and bias are not supported"
 
+    if backend == WNA16MoEBackend.RDNA3:
+        if not isinstance(quant_config, QuantizationArgs):
+            return "only compressed-tensors checkpoints are supported"
+        if may_have_zp:
+            return "asymmetric checkpoints are not supported"
+        if may_have_bias:
+            return "expert bias is not supported"
+        if quant_config.actorder == "group":
+            return "group activation ordering is not supported"
+        if quant_config.strategy != QuantizationStrategy.GROUP:
+            return "only group-wise scales are supported"
+
     from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
     from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
     from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
@@ -190,6 +212,7 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend:
         "humming": WNA16MoEBackend.HUMMING,
         "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM,
         "emulation": WNA16MoEBackend.EMULATION,
+        "rdna3": WNA16MoEBackend.RDNA3,
     }
     if backend := mapping.get(runner_backend):
         return backend
@@ -216,11 +239,13 @@ def select_wna16_moe_backend(
         quant_config: Quantization structure and checkpoint format description.
         may_have_zp: Whether the integration can provide weight zero points.
         may_have_bias: Whether the integration can provide expert bias.
+        allow_tile_padding: Whether backends that require padding the weights
+            up to a tile boundary may be selected.
 
     Returns:
         A tuple of (``WNA16MoEBackend``, experts class or ``None``).
-    """
 
+    """
     activation_format = (
         mk.FusedMoEActivationFormat.BatchedExperts
         if config.moe_parallel_config.use_batched_activation_format
@@ -372,6 +397,9 @@ def make_wna16_moe_kernel(
     from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import (
         Int4EmulationTritonExperts,
     )
+    from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
+        Rdna3WNA16Experts,
+    )
     from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
         XPUExpertsWNA16,
     )
@@ -387,6 +415,7 @@ def make_wna16_moe_kernel(
         XPUExpertsWNA16,
         CPUExpertsInt4,
         Int4EmulationTritonExperts,
+        Rdna3WNA16Experts,
     )
     if backend == WNA16MoEBackend.HUMMING:
         allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING))
@@ -811,6 +840,78 @@ def _process_awq_weights_marlin(
     )
 
 
+def _synthesize_rdna3_qzeros(
+    groups: int, out_features: int, device: torch.device
+) -> torch.Tensor:
+    """Create the packed zero-point tensor for symmetric quantization.
+
+    GPTQv1 +1 quirk: the kernel adds 1 to the stored zeros, so encode
+    (bias - 1) = 7 for uint4b8 (bias=8).
+    """
+    zeros = torch.full(
+        (groups, out_features),
+        scalar_types.uint4b8.bias - 1,
+        dtype=torch.int32,
+        device=device,
+    )
+    return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)
+
+
+def _process_weights_rdna3(
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    group_size: int,
+) -> tuple[
+    torch.Tensor,  # w13_qweight
+    torch.Tensor,  # w2_qweight
+    torch.Tensor,  # w13_scales
+    torch.Tensor,  # w2_scales
+    torch.Tensor | None,  # w13_qzeros
+    torch.Tensor | None,  # w2_qzeros
+    torch.Tensor | None,  # w13_input_global_scale
+    torch.Tensor | None,  # w2_input_global_scale
+    torch.Tensor | None,  # w13_bias
+    torch.Tensor | None,  # w2_bias
+]:
+    """RDNA3 (gfx1100) W4A16 weight post-processing.
+
+    Interleaves the packed nibbles per expert (the exllama shuffle the dense
+    RDNA3 kernel also uses) and synthesizes the symmetric zero points that
+    ``moe_gptq_gemm_rdna3`` dequantizes with. The packed layout
+    ``[E, K // 8, N]`` and the ``[E, groups, N]`` scales are already what the
+    kernel wants, so neither is repacked.
+    """
+    device = w13.device
+    num_experts = w13.size(0)
+
+    for e in range(num_experts):
+        w13_e = w13[e].contiguous()
+        ops.gptq_shuffle(w13_e, 4)
+        w13[e] = w13_e
+        w2_e = w2[e].contiguous()
+        ops.gptq_shuffle(w2_e, 4)
+        w2[e] = w2_e
+
+    def _qzeros(w: torch.Tensor) -> torch.Tensor:
+        qz = _synthesize_rdna3_qzeros((w.size(1) * 8) // group_size, w.size(2), device)
+        return qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous()
+
+    return (
+        w13,
+        w2,
+        w13_scale.contiguous(),
+        w2_scale.contiguous(),
+        _qzeros(w13),
+        _qzeros(w2),
+        None,  # w13_input_global_scale
+        None,  # w2_input_global_scale
+        None,  # w13_bias
+        None,  # w2_bias
+    )
+
+
 def _process_weights_cpu(
     quant_config: QuantizationConfig | QuantizationArgs | None,
     w13: torch.Tensor,
@@ -1093,6 +1194,7 @@ def _unpack_and_dequant_int4_gptq(
 
     Returns:
         Dequantized weight tensor in the requested layout.
+
     """
     E, K_packed, N = w_int32.shape
     K = K_packed * 8
@@ -1155,6 +1257,7 @@ def _unpack_and_dequant_int4_awq(
 
     Returns:
         Dequantized weight tensor in the requested layout.
+
     """
     E, K, N_packed = w_int32.shape
     N = N_packed * 8
@@ -1341,10 +1444,19 @@ def convert_to_wna16_moe_kernel_format(
         layer: the ``MoERunner`` layer whose parameters are being prepared.
         quant_config: the ``QuantizationConfig`` for this layer.
         input_dtype: optional activation dtype, usually should be 16 bit.
+        w13: fused gate/up expert weights.
+        w2: down-projection expert weights.
+        w13_scale: quantization scales for ``w13``.
+        w2_scale: quantization scales for ``w2``.
+        w13_qzeros: optional zero points for ``w13``.
+        w2_qzeros: optional zero points for ``w2``.
+        w13_bias: optional bias for ``w13``.
+        w2_bias: optional bias for ``w2``.
+
     """
     if backend == WNA16MoEBackend.HUMMING:
         from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             convert_to_humming_moe_kernel_format,
         )
 
@@ -1432,6 +1544,15 @@ def convert_to_wna16_moe_kernel_format(
                 w13_bias,
                 w2_bias,
             )
+    elif backend == WNA16MoEBackend.RDNA3:
+        assert isinstance(quant_config, QuantizationArgs)
+        return _process_weights_rdna3(
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            quant_config.group_size,
+        )
     elif backend == WNA16MoEBackend.CPU:
         return _process_weights_cpu(
             quant_config,
diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
index 3d6cb6cea3d4..f2f8717f3d39 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
@@ -31,7 +31,12 @@
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     swap_w13_to_w31,
 )
-from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+    _swizzle_mxfp4,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+    mx_scale_kwargs as _mx_scale_kwargs,
+)
 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
     OCP_MX_BLOCK_SIZE,
 )
@@ -45,7 +50,7 @@
 )
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import all_close_1d
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.import_utils import get_triton_kernels_version
 from vllm.utils.math_utils import round_up
 
 if TYPE_CHECKING:
@@ -54,9 +59,14 @@
 
 logger = init_logger(__name__)
 
-if has_triton_kernels():
+triton_kernels_version = get_triton_kernels_version()
+
+if triton_kernels_version is not None:
     try:
-        from triton_kernels.matmul_ogs import PrecisionConfig
+        if triton_kernels_version == "3.8":
+            from triton_kernels.matmul import FlexCtx, PrecisionConfig
+        else:
+            from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
     except (ImportError, AttributeError) as e:
         logger.error(
             "Failed to import Triton kernels. Please make sure your triton "
@@ -358,8 +368,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:
 
 
 def _get_priority_backends() -> list[Mxfp4MoeBackend]:
-    """
-    Get available backends in priority order. SM100+ prefers DeepGEMM FP4 /
+    """Get available backends in priority order. SM100+ prefers DeepGEMM FP4 /
     TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the
     backend-level ``is_supported_config`` check filters by device capability).
     """
@@ -530,8 +539,7 @@ def select_mxfp4_moe_backend(
     config: FusedMoEConfig,
     activation_key: QuantKey | None = None,
 ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:
-    """
-    Select the primary MXFP4 MoE backend.
+    """Select the primary MXFP4 MoE backend.
 
     Args:
         config: MoE configuration
@@ -540,6 +548,7 @@ def select_mxfp4_moe_backend(
             Use kFp8StaticTensorSym for W4A8 scheme.
 
     Note: Shape-specific fallbacks may still occur at runtime.
+
     """
     runner_backend = config.moe_backend
     requested_activation_key = _resolve_activation_key(activation_key)
@@ -652,8 +661,7 @@ def select_mxfp4_moe_backend(
 def select_deepseek_v4_mxfp4_moe_backend(
     config: FusedMoEConfig,
 ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:
-    """
-    Select the MXFP4 MoE backend with MXFP8 activation as top priority.
+    """Select the MXFP4 MoE backend with MXFP8 activation as top priority.
     Falls back through BF16 and other backends.
     """
     activation_format = (
@@ -762,18 +770,31 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
         intermediate_size = round_up(intermediate_size, 128)
         hidden_size = round_up(hidden_size, 128)
     elif current_platform.is_rocm():
-        if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and (
-            activation == MoEActivation.SITU or activation == MoEActivation.SILU
-        ):
-            # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size
-            # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a
-            # no-op for K3's shapes) rather than the generic ROCm 256 round-up,
-            # which would inflate weights and OOM.
-            intermediate_size = round_up(intermediate_size, 128)
-            hidden_size = round_up(hidden_size, 128)
-        else:
-            intermediate_size = round_up(intermediate_size, 256)
-            hidden_size = round_up(hidden_size, 256)
+        from vllm.platforms.rocm import get_cdna_version
+
+        is_situ_or_silu = activation in (
+            MoEActivation.SITU,
+            MoEActivation.SILU,
+        )
+
+        # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size
+        # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128
+        # rather than the generic ROCm 256 round-up, which would inflate
+        # weights and OOM.
+        aiter_uses_128 = backend == Mxfp4MoeBackend.AITER_MXFP4_BF16
+
+        # matmul_ogs uses block_k=128 for MXFP4 on pre-CDNA4 GPUs.
+        # CDNA4's F16xMXFP4 configuration uses block_k=256.
+        triton_uses_128 = (
+            backend == Mxfp4MoeBackend.TRITON_UNFUSED and get_cdna_version() != 4
+        )
+
+        alignment = (
+            128 if is_situ_or_silu and (aiter_uses_128 or triton_uses_128) else 256
+        )
+
+        intermediate_size = round_up(intermediate_size, alignment)
+        hidden_size = round_up(hidden_size, alignment)
     elif backend == Mxfp4MoeBackend.CPU:
         # CPU AMX kernel uses BLOCK_N=32, align to 32
         intermediate_size = round_up(intermediate_size, 32)
@@ -816,6 +837,14 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
     torch.Tensor | None,
 ]:
     """Convert loaded weights into backend-specific kernel format."""
+    import os
+
+    is_gfx1250 = False
+    if current_platform.is_rocm():
+        from vllm.platforms.rocm import on_gfx1250
+
+        is_gfx1250 = on_gfx1250()
+
 
     if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
         w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
@@ -847,7 +876,7 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
         is_gfx1250 = on_gfx1250()
 
     if mxfp4_backend == Mxfp4MoeBackend.HUMMING:
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             convert_to_humming_moe_kernel_format,
         )
 
@@ -1149,11 +1178,11 @@ def swap_every_two_rows(x, axis=-1):
             return (w13_data, w2_data, w13_scale, w2_scale,
                     w13_bias, w2_bias)
 
-        from vllm._aiter_ops import rocm_aiter_ops
-
         # e8m0_shuffle on weight scales (GFX950 swizzle layout)
         from aiter.utility.fp4_utils import e8m0_shuffle
 
+        from vllm._aiter_ops import rocm_aiter_ops
+
         s0, s1, _ = w13_weight_scale.shape
         w13_weight_scale.data = e8m0_shuffle(w13_weight_scale.view(s0 * s1, -1)).view(
             s0, s1, -1
@@ -1186,6 +1215,50 @@ def swap_every_two_rows(x, axis=-1):
             w2_bias,
         )
 
+    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250:
+        from aiter.ops.shuffle import moe_shuffle_scale, moe_shuffle_weight
+
+        os.environ.setdefault("AITER_FORCE_A8W4", "1")
+
+        if w13_bias is not None:
+            w13_bias = w13_bias.data.to(torch.float32)
+        if w2_bias is not None:
+            w2_bias = w2_bias.data.to(torch.float32)
+
+        fp4_dtype = torch.float4_e2m1fn_x2
+        e8m0_dtype = torch.float8_e8m0fnu
+
+        w13_raw = w13_weight.data.view(fp4_dtype)
+        w2_raw = w2_weight.data.view(fp4_dtype)
+        w13_scale_raw = w13_weight_scale.data.view(e8m0_dtype)
+        w2_scale_raw = w2_weight_scale.data.view(e8m0_dtype)
+        w13_scale_raw = w13_scale_raw.view(-1, w13_scale_raw.shape[-1])
+        w2_scale_raw = w2_scale_raw.view(-1, w2_scale_raw.shape[-1])
+
+        w13 = moe_shuffle_weight(
+            w13_raw, num_experts, is_guinterleave=False, gate_up=True
+        )
+        w2 = moe_shuffle_weight(
+            w2_raw, num_experts, is_guinterleave=False, gate_up=False
+        )
+        shuffled_w13_scale = moe_shuffle_scale(
+            w13_scale_raw, num_experts, is_guinterleave=False, gate_up=True
+        )
+        shuffled_w2_scale = moe_shuffle_scale(
+            w2_scale_raw, num_experts, is_guinterleave=False, gate_up=False
+        )
+        w13.is_shuffled = True
+        w2.is_shuffled = True
+
+        return (
+            w13,
+            w2,
+            shuffled_w13_scale,
+            shuffled_w2_scale,
+            w13_bias,
+            w2_bias,
+        )
+
     elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
         from vllm._aiter_ops import rocm_aiter_ops
 
@@ -1253,7 +1326,6 @@ def swap_every_two_rows(x, axis=-1):
 
     elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
         # W4A8: MXFP4 weights + static FP8 activations (triton kernel)
-        from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
         from triton_kernels.numerics import InFlexData
 
         if w13_bias is not None:
@@ -1288,11 +1360,11 @@ def swap_every_two_rows(x, axis=-1):
 
         # Create PrecisionConfig with both weight and activation info
         w13_precision_config = PrecisionConfig(
-            weight_scale=w13_scale,
+            **_mx_scale_kwargs(w13_scale),
             flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
         )
         w2_precision_config = PrecisionConfig(
-            weight_scale=w2_scale,
+            **_mx_scale_kwargs(w2_scale),
             flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
         )
 
@@ -1309,8 +1381,6 @@ def swap_every_two_rows(x, axis=-1):
         )
 
     elif mxfp4_backend in TRITON_BACKENDS:
-        from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
-
         if w13_bias is not None:
             w13_bias = w13_bias.to(torch.float32)
         if w2_bias is not None:
@@ -1326,10 +1396,10 @@ def swap_every_two_rows(x, axis=-1):
         )
 
         w13_precision_config = PrecisionConfig(
-            weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
+            **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex)
         )
         w2_precision_config = PrecisionConfig(
-            weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
+            **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex)
         )
 
         # The original mxfp4 block scales have been swizzled into the
@@ -1479,7 +1549,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
         )
 
     if mxfp4_backend == Mxfp4MoeBackend.HUMMING:
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             convert_to_humming_moe_kernel_format,
         )
 
@@ -1690,12 +1760,6 @@ def convert_weight_to_mxfp4_moe_kernel_format(
         if w2_bias is not None:
             w2_bias = w2_bias.data.to(torch.float32)
 
-        import os
-
-        # TODO: Remove this once AITER is fixed
-        # Necessary for AITER side from crashing
-        os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"
-
         if activation == MoEActivation.SITU:
             from aiter.utility.fp4_utils import e8m0_shuffle
 
@@ -1703,9 +1767,8 @@ def convert_weight_to_mxfp4_moe_kernel_format(
 
             fp4_dtype = torch.float4_e2m1fn_x2
             e8m0_dtype = torch.float8_e8m0fnu
-            # a8w4 uses gate/up-interleaved flydsl kernels;
-            # default a16w4 keeps the separated layout.
-            guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled()
+            # SiTUv2 flydsl uses separated gate/up layout (a4w4).
+            guinterleave = False
             w13 = rocm_aiter_ops.shuffle_weight_a16w4(
                 w13_weight.data.view(fp4_dtype), 16, guinterleave
             )
@@ -1724,6 +1787,14 @@ def convert_weight_to_mxfp4_moe_kernel_format(
             w2.is_shuffled = True
             return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
 
+        import os
+
+        # Interleaved a16w4 only (DeepSeekV4 etc.). AITER uses this bound to
+        # pick bf16 vs fp8 activations when gate_mode is INTERLEAVE. SiTUv2
+        # a4w4 is separated and selects q_dtype_a independently, so the bound
+        # is unused on that path.
+        os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"
+
         from aiter.ops.shuffle import shuffle_scale as _shuf_s
         from aiter.ops.shuffle import shuffle_weight as _shuf_w
 
@@ -1770,9 +1841,9 @@ def convert_weight_to_mxfp4_moe_kernel_format(
             w2_bias,
         )
 
-    elif mxfp4_backend in TRITON_BACKENDS:
-        from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
-
+    elif mxfp4_backend in TRITON_BACKENDS or (
+        mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
+    ):
         if mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16:
             # AITER moe_gemm_a16w4 needs gate/up interleaved
             def interleave_gate_up(w: torch.Tensor) -> torch.Tensor:
@@ -1816,10 +1887,10 @@ def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
         )
 
         w13_precision_config = PrecisionConfig(
-            weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
+            **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex)
         )
         w2_precision_config = PrecisionConfig(
-            weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
+            **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex)
         )
 
         # The original mxfp4 block scales have been swizzled into the
@@ -2066,7 +2137,7 @@ def make_mxfp4_moe_quant_config(
             gemm1_clamp_limit=swiglu_limit,
         )
     elif mxfp4_backend == Mxfp4MoeBackend.HUMMING:
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             get_humming_moe_quant_config,
         )
 
diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
index 2a03eacea8e1..49f61764f0fd 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
@@ -111,8 +111,8 @@ def select_mxfp8_moe_backend(
 
     Returns:
         A tuple of (fp8_backend, experts_cls).
-    """
 
+    """
     runner_backend = config.moe_backend
     if runner_backend != "auto":
         backend = _BACKEND_NAME_MAP.get(runner_backend)
diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
index a860003775d5..0fecb65e4b0b 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
@@ -184,11 +184,9 @@ def select_nvfp4_moe_backend(
     weight_key: QuantKey | None,
     activation_key: QuantKey | None,
 ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts]]:
-    """
-    Select the primary NvFP4 MoE backend
+    """Select the primary NvFP4 MoE backend
     Note: Shape-specific fallbacks may still occur at runtime.
     """
-
     # NOTE: the kernels are selected in the following order.
     # FLASHINFER_B12X is intentionally excluded from auto-selection until
     # the upstream CUTLASS SM121 MMA op guard is resolved; use
@@ -411,7 +409,7 @@ def convert_to_nvfp4_moe_kernel_format(
             is_act_and_mul=is_act_and_mul,
         )
     elif nvfp4_backend == NvFp4MoeBackend.HUMMING:
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             convert_to_humming_moe_kernel_format,
         )
 
@@ -525,7 +523,7 @@ def make_nvfp4_moe_quant_config(
     use_a16 = _use_a16(backend, use_a16)
     if backend == NvFp4MoeBackend.HUMMING:
         from vllm.model_executor.layers.fused_moe import RoutedExperts
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+        from vllm.model_executor.layers.quantization.utils.humming import (
             get_humming_moe_quant_config,
         )
 
@@ -607,6 +605,7 @@ def make_nvfp4_moe_kernel(
         routing_tables=routing_tables,
         allow_new_interface=True,
         use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),
+        input_dtype=moe_config.in_dtype if per_token_activation else None,
     )
     assert prepare_finalize is not None
 
diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py
index b24f2a9a797e..2c84e532cc87 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py
@@ -26,6 +26,7 @@
     swap_w13_to_w31,
 )
 from vllm.platforms import current_platform
+from vllm.utils.math_utils import round_up
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
@@ -39,6 +40,7 @@ class UnquantizedMoeBackend(Enum):
     AITER = "ROCm AITER"
     TRITON = "TRITON"
     BATCHED_TRITON = "BATCHED_TRITON"
+    MOONEP = "MOONEP"
     CPU = "CPU"
     XPU = "XPU"
     TPU = "TPU"
@@ -46,8 +48,7 @@ class UnquantizedMoeBackend(Enum):
 
 
 def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBackend]:
-    """
-    Get available backends in priority order based on platform and config.
+    """Get available backends in priority order based on platform and config.
 
     This function can be extended to become more complex as needed.
     """
@@ -93,7 +94,13 @@ def _move_to_back(
             _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)
 
     elif current_platform.is_xpu():
-        _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.XPU]
+        # XPUExperts does not support batch invariance; TRITON is kept as a
+        # fallback so VLLM_BATCH_INVARIANT=1 still resolves to a valid
+        # (batch-invariant) backend on XPU.
+        _AVAILABLE_BACKENDS = [
+            UnquantizedMoeBackend.XPU,
+            UnquantizedMoeBackend.TRITON,
+        ]
     elif current_platform.is_cpu():
         _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.CPU]
     return _AVAILABLE_BACKENDS
@@ -138,6 +145,13 @@ def backend_to_kernel_cls(
 
         return [BatchedTritonExperts]
 
+    elif backend == UnquantizedMoeBackend.MOONEP:
+        from vllm.model_executor.layers.fused_moe.experts.moonep_experts import (
+            MoonEPExperts,
+        )
+
+        return [MoonEPExperts]
+
     elif backend == UnquantizedMoeBackend.XPU:
         from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExperts
 
@@ -208,11 +222,9 @@ def _trtllm_bf16_lora_supported(moe_config: FusedMoEConfig) -> bool:
 def select_unquantized_moe_backend(
     moe_config: FusedMoEConfig,
 ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
-    """
-    Select the primary Unquantized MoE backend.
+    """Select the primary Unquantized MoE backend.
     Note: Shape-specific fallbacks may still occur at runtime.
     """
-
     if current_platform.is_tpu():
         return UnquantizedMoeBackend.TPU, None
 
@@ -283,6 +295,13 @@ def _return_or_raise(
                 return backend, k_cls
         raise ValueError(_make_log_unsupported(backend, reason))
 
+    # MoonEP owns the expert layout (expert-grouped [NvS, H] + cu_seqlens),
+    # so it is not interchangeable with the token-major experts backends.
+    if moe_config.moe_parallel_config.use_moonep_kernels:
+        return _return_or_raise(
+            UnquantizedMoeBackend.MOONEP, moe_config, activation_format
+        )
+
     runner_backend = moe_config.moe_backend
     # 'humming' is quantization-only; an unquantized layer (e.g. excluded via
     # modules_to_not_convert) falls through to auto instead of erroring.
@@ -308,7 +327,33 @@ def _return_or_raise(
                 AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER)
         else:
             backend = UnquantizedMoeBackend.AITER
-            return _return_or_raise(backend, moe_config, activation_format)
+            reason = None
+            for k_cls in backend_to_kernel_cls(backend):
+                supported, reason = k_cls.is_supported_config(
+                    k_cls, moe_config, None, None, activation_format
+                )
+                if supported:
+                    logger.info_once(_make_log_backend(backend))
+                    return backend, k_cls
+            # AITER was explicitly requested but does not support this
+            # deployment configuration (e.g. a non-gated MoE activation,
+            # which AiterExperts._supports_no_act_and_mul() rejects
+            # unconditionally). Rather than hard-crashing at engine init,
+            # warn — explicitly stating that this is falling back, and
+            # which backends will be tried next — and fall back to the
+            # remaining backends in priority order below.
+            logger.warning_once(
+                "VLLM_ROCM_USE_AITER_MOE=1 was requested, but %s "
+                "Falling back to try the remaining available MoE backends: %s.",
+                _make_log_unsupported(backend, reason),
+                ", ".join(
+                    b.value
+                    for b in AVAILABLE_BACKENDS
+                    if b != UnquantizedMoeBackend.AITER
+                ),
+            )
+            if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS:
+                AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER)
 
     for backend in AVAILABLE_BACKENDS:
         for k_cls in backend_to_kernel_cls(backend):
@@ -326,12 +371,58 @@ def _return_or_raise(
     )
 
 
+def unquantized_round_up_hidden_size_and_intermediate_size(
+    backend: UnquantizedMoeBackend,
+    hidden_size: int,
+    intermediate_size: int,
+) -> tuple[int, int]:
+    """Round up dimensions before allocation to satisfy the selected kernel."""
+    if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
+        intermediate_size = round_up(intermediate_size, 128)
+    return hidden_size, intermediate_size
+
+
 def convert_to_unquantized_kernel_format(
     unquantized_backend: UnquantizedMoeBackend,
     moe_config: FusedMoEConfig,
     w13_weight: torch.Tensor,
     w2_weight: torch.Tensor,
+    layer: torch.nn.Module | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    if unquantized_backend == UnquantizedMoeBackend.MOONEP:
+        if layer is None:
+            raise ValueError(
+                "MoonEP weight conversion requires the layer module to stash "
+                "the [E+B] weight layout on"
+            )
+        # MoonEP addresses expert weights by global [E+B] row and needs one
+        # contiguous tensor per projection (gate / up / down). The layer's
+        # w13_weight becomes the gate tensor and w2_weight the down tensor;
+        # the up tensor and the full layout are stashed on the layer for
+        # the kernel wiring in _setup_kernel.
+        from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
+            MOONEP_DEFAULT_NUM_PREFETCH_SLOTS,
+            gather_moonep_weight_layout,
+        )
+
+        if getattr(layer, "_moonep_weight_layout", None) is not None:
+            # On a reload the layer parameters were already replaced by the
+            # converted [E+B] gate/down tensors, so the source-format weights
+            # this conversion needs are gone (and the unregistered up
+            # projection cannot be recovered from them at all). Tracked in
+            # RFC #52095.
+            raise NotImplementedError(
+                "MoonEP does not support in-place weight reloads yet"
+            )
+        layout = gather_moonep_weight_layout(
+            w13_weight,
+            w2_weight,
+            num_global_experts=moe_config.num_experts,
+            num_prefetch_slots=MOONEP_DEFAULT_NUM_PREFETCH_SLOTS,
+        )
+        layer._moonep_weight_layout = layout
+        return layout.full_gate_weight, layout.full_down_weight
+
     if unquantized_backend == UnquantizedMoeBackend.AITER:
         w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight)
         w13_weight.is_shuffled = True
@@ -354,13 +445,26 @@ def convert_to_unquantized_kernel_format(
         )
         moe_config.intermediate_size_per_partition = padded_intermediate
 
+        # Reloads only overwrite checkpoint slices. An earlier in-place
+        # permutation can leave nonzero values in the raw padding slots.
+        unpadded = moe_config.intermediate_size_per_partition_unpadded
+        assert unpadded is not None
+        if padded_intermediate > unpadded:
+            w13_weight[:, unpadded:padded_intermediate].zero_()
+            if is_act_and_mul:
+                w13_weight[:, padded_intermediate + unpadded :].zero_()
+            w2_weight[:, :, unpadded:].zero_()
+
         _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
-        w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout(
+        convert_moe_weights_to_flashinfer_trtllm_block_layout(
             _cache_permute_indices,
             w13_weight,
             w2_weight,
             is_gated_act_gemm=is_act_and_mul,
         )
+        # Keep checkpoint-shaped parameters for reload/IPC. The experts create
+        # BlockMajorK views at dispatch without changing the underlying storage.
+        return w13_weight, w2_weight
 
     if (
         unquantized_backend == UnquantizedMoeBackend.TRITON
@@ -474,9 +578,10 @@ def convert_to_kernel_format(
         moe_config: FusedMoEConfig,
         w13_weight: torch.Tensor,
         w2_weight: torch.Tensor,
+        layer: torch.nn.Module | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         return convert_to_unquantized_kernel_format(
-            backend, moe_config, w13_weight, w2_weight
+            backend, moe_config, w13_weight, w2_weight, layer=layer
         )
 
     def make_kernel(
diff --git a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py
index c66e43ccba8b..db0a30544373 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py
@@ -32,8 +32,7 @@ class W4A8Int8MoeBackend(Enum):
 def _get_priority_backends(
     moe_config: FusedMoEConfig,
 ) -> list[W4A8Int8MoeBackend]:
-    """
-    Get available backends in priority order based on platform and config.
+    """Get available backends in priority order based on platform and config.
 
     Currently only CPU INT4 backend is available for W4A8 INT8 MoE.
     """
@@ -75,8 +74,7 @@ def select_w4a8_int8_moe_backend(
     weight_key: QuantKey | None,
     activation_key: QuantKey | None,
 ) -> tuple[W4A8Int8MoeBackend, type[mk.FusedMoEExperts]]:
-    """
-    Select the primary W4A8 Int8 MoE backend.
+    """Select the primary W4A8 Int8 MoE backend.
 
     Args:
         config: MoE configuration
@@ -85,8 +83,8 @@ def select_w4a8_int8_moe_backend(
 
     Returns:
         Tuple of (backend, kernel_class)
-    """
 
+    """
     AVAILABLE_BACKENDS = _get_priority_backends(config)
 
     if not AVAILABLE_BACKENDS:
@@ -160,8 +158,7 @@ def _return_or_raise(
 def make_w4a8_int8_moe_quant_config(
     block_shape: tuple[int, int] | None = None,
 ) -> FusedMoEQuantConfig:
-    """
-    Create FusedMoEQuantConfig for W4A8 Int8 MoE.
+    """Create FusedMoEQuantConfig for W4A8 Int8 MoE.
 
     Args:
         block_shape: Quantization block shape (row, col).
@@ -170,6 +167,7 @@ def make_w4a8_int8_moe_quant_config(
 
     Returns:
         FusedMoEQuantConfig with appropriate settings for W4A8 Int8
+
     """
     # W4A8 Int8 uses static weight quantization, dynamic activation quantization
     # Weights are 4-bit (stored as int8, packed to uint8),
@@ -196,8 +194,7 @@ def pack_int4_weights_for_kleidi(
     out_features: int,
     group_size: int,
 ) -> torch.Tensor:
-    """
-    Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.
+    """Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.
 
     Args:
         int4_as_int8: [out, in] int8 tensor with values in [-8, 7]
@@ -209,6 +206,7 @@ def pack_int4_weights_for_kleidi(
 
     Returns:
         Packed weight tensor in KleidiAI format
+
     """
     # Shift to unsigned nibble [0, 15]
     tmp = int4_as_int8.add(8)
@@ -250,8 +248,7 @@ def convert_to_w4a8_int8_moe_format(
     torch.Tensor | None,
     torch.Tensor | None,
 ]:
-    """
-    Pack INT4 MoE weights to KleidiAI format.
+    """Pack INT4 MoE weights to KleidiAI format.
 
     This function packs the INT4 weights (stored as int8 values) into
     the format expected by the KleidiAI dynamic_4bit_int_moe kernel.
@@ -267,6 +264,7 @@ def convert_to_w4a8_int8_moe_format(
 
     Returns:
         Tuple of (w13_packed, w2_packed) tensors
+
     """
     # Derive dimensions from tensor shapes
     E = w13_weight.shape[0]  # num_experts
@@ -315,8 +313,7 @@ def make_w4a8_int8_moe_kernel(
     experts_cls: type[mk.FusedMoEExperts],
     routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
 ) -> mk.FusedMoEKernel:
-    """
-    Create FusedMoEKernel for W4A8 Int8 MoE.
+    """Create FusedMoEKernel for W4A8 Int8 MoE.
 
     Args:
         moe_quant_config: Quantization configuration
@@ -326,6 +323,7 @@ def make_w4a8_int8_moe_kernel(
 
     Returns:
         Configured FusedMoEKernel instance
+
     """
     # Create Prepare/Finalize.
     prepare_finalize = maybe_make_prepare_finalize(
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py
index d11beaf98fbd..5a818bfa57c7 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py
@@ -16,8 +16,7 @@
 
 
 class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
-    """
-    A reference prepare/finalize class that reorganizes the tokens into
+    """A reference prepare/finalize class that reorganizes the tokens into
     expert batched format, i.e. E x max_num_tokens x K.  This is the format
     that the batched dispatch/combine kernels use.
     """
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py
index f30e0bea3c39..035d5552693d 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py
@@ -27,9 +27,7 @@
 
 
 class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
-    """
-    Prepare/Finalize using DeepEP High-Throughput kernels.
-    """
+    """Prepare/Finalize using DeepEP High-Throughput kernels."""
 
     @staticmethod
     def maybe_roundup_layer_hidden_size(hidden_size: int, dtype: torch.dtype) -> int:
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py
index e5d2b601a768..1f357c79a873 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py
@@ -34,9 +34,7 @@
 def dequant_fp8(
     expert_x_fp8: torch.Tensor, expert_x_scales: torch.Tensor
 ) -> torch.Tensor:
-    """
-    Return dequantized tensor in fp32
-    """
+    """Return dequantized tensor in fp32."""
     # TODO (varun) : Optimize leverage num_tokens_per_expert counts
     assert expert_x_fp8.is_contiguous()
     expert_x_scales = expert_x_scales.contiguous()
@@ -50,9 +48,7 @@ def dequant_fp8(
 
 
 class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
-    """
-    Prepare/Finalize using DeepEP low-latency kernels.
-    """
+    """Prepare/Finalize using DeepEP low-latency kernels."""
 
     # DeepEP low-latency kernels are compiled only for certain
     # specific hidden sizes.
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
index d834c6eff73c..a584ab5f52ff 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
@@ -18,17 +18,17 @@
     swizzle_mxfp8_scale,
 )
 from vllm.triton_utils import tl, triton
-from vllm.utils.math_utils import round_up
+from vllm.utils.math_utils import cdiv, round_up
 from vllm.v1.worker.ubatching import (
     dbo_current_ubatch_id,
+    dbo_enabled,
 )
 
 
 def _quantize_before_dispatch(
     quant_config: FusedMoEQuantConfig, defer_input_quant: bool
 ) -> bool:
-    """
-    Do quantized dispatch for blockfp8 and mxfp8, unless the
+    """Do quantized dispatch for blockfp8 and mxfp8, unless the
     subsequent moe kernel requires bf16 inputs.
     """
     if defer_input_quant:
@@ -68,8 +68,7 @@ def _unpack_mxfp8_scale(
 
 
 class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
-    """
-    Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API).
+    """Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API).
 
     Supports two modes controlled by the `use_cudagraph` constructor arg:
 
@@ -89,8 +88,10 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         use cudagraphs anyway
       - Provides expert_tokens_meta for efficient batched expert kernels
 
-    Both modes use async_with_compute_stream=False (synchronous from
-    caller's perspective). The ElasticBuffer handles comm internally.
+    Dispatch always uses async_with_compute_stream=False. finalize_async
+    issues the combine with async_with_compute_stream=True (except under
+    DBO) so the modular kernel can overlap the shared-expert FFN with the
+    combine a2a; the returned receiver joins via a device-side event wait.
     """
 
     @staticmethod
@@ -113,6 +114,7 @@ def __init__(
         num_topk: int,
         use_fp8_dispatch: bool = False,
         use_cudagraph: bool = False,
+        sp_size: int = 1,
     ):
         super().__init__()
         self.buffer = buffer
@@ -123,6 +125,7 @@ def __init__(
         self.num_topk = num_topk
         self.use_fp8_dispatch = use_fp8_dispatch
         self.use_cudagraph = use_cudagraph
+        self.sp_size = sp_size
 
         # DBO microbatching: one handle slot per micro-batch.
         self.handles: list[deep_ep.EPHandle | None] = [None, None]
@@ -191,12 +194,14 @@ def _do_dispatch(
         # high concurrency). Round up to a power of 2 instead: this bounds the
         # set to ~log2(max_num_batched_tokens) values (compiled once, then
         # cached) while staying small for decode (e.g. 1 token -> 1) and capped
-        # at the buffer's init capacity for prefill.
+        # at the buffer's init capacity for prefill. With sequence parallelism,
+        # each EP rank holds a ceil(n / sp_size) shard of its DP rank's batch.
         num_max_tokens_per_rank = None
         if not do_expand:
             dp_meta = get_forward_context().dp_metadata
             if dp_meta is not None:
                 n = int(dp_meta.num_tokens_across_dp_cpu.max())
+                n = cdiv(n, self.sp_size)
             else:
                 n = tokens.shape[0]
             num_max_tokens_per_rank = 1 << max(n - 1, 0).bit_length()
@@ -465,15 +470,31 @@ def _finalize(
                 f"got {fused_expert_output.dtype}"
             )
 
+        # DBO drives its own hook/receiver schedule; keep the combine
+        # synchronous there (the receiver then only performs the copy).
+        combine_async = do_async and not dbo_enabled()
         combined_x, _, event = self.buffer.combine(
             x=fused_expert_output,
             handle=handle,
             topk_weights=None,
-            async_with_compute_stream=False,
+            async_with_compute_stream=combine_async,
+            allocate_on_comm_stream=combine_async,
         )
 
-        output.copy_(combined_x, non_blocking=True)
-        return None
+        if do_async:
+            # The combine ran on DeepEP's comm stream; the modular kernel
+            # issues the shared-expert FFN before calling the receiver, which
+            # joins via a device-side cudaStreamWaitEvent (no host sync, so
+            # this is safe inside a captured region).
+            def _receiver():
+                if event.event is not None:
+                    event.current_stream_wait()
+                output.copy_(combined_x, non_blocking=True)
+
+            return _receiver
+        else:
+            output.copy_(combined_x, non_blocking=True)
+            return None
 
     def finalize_async(
         self,
@@ -484,16 +505,17 @@ def finalize_async(
         apply_router_weight_on_input: bool,
         weight_and_reduce_impl: mk.TopKWeightAndReduce,
     ) -> Callable:
-        self._finalize(
+        receiver = self._finalize(
             output,
             fused_expert_output,
             topk_weights,
             topk_ids,
             apply_router_weight_on_input,
             weight_and_reduce_impl,
-            False,
+            True,
         )
-        return lambda: None
+        assert receiver is not None
+        return receiver
 
     def finalize(
         self,
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
index 0b57afb6062c..d7a84c931b8d 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
@@ -13,9 +13,10 @@
 from vllm.utils.flashinfer import nvfp4_block_scale_interleave
 
 
-def get_local_sizes():
+def get_local_sizes() -> list[int] | None:
     dp_metadata = get_forward_context().dp_metadata
-    assert dp_metadata is not None
+    if dp_metadata is None:  # PCP with DP=1
+        return None
     return dp_metadata.get_chunk_sizes_across_dp_rank()
 
 
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/moonep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/moonep.py
new file mode 100644
index 000000000000..fdbed5ade858
--- /dev/null
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/moonep.py
@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MoonEP (https://github.com/MoonshotAI/MoonEP) prepare/finalize.
+
+BF16 correctness-first proof of concept on top of vLLM's modular kernel
+interface.
+
+MoonEP differs from DeepEP-style backends in two ways that shape this
+integration:
+
+- ``dispatch`` returns tokens already grouped by expert *row* (a fixed
+  ``[NvS, H]`` layout with ``NvS = S x K`` real slots plus padding) together
+  with a ``cu_seqlens[E+B]`` segment table and an opaque ``plan``. There is
+  no per-token topk id tensor after dispatch; the expert compute must be a
+  grouped GEMM over ``cu_seqlens`` segments.
+- Rows ``[E, E+B)`` of the weight/segment space are dynamic redundant-expert
+  prefetch slots. ``plan.experts_to_copy`` names the source expert of each
+  slot and ``Buffer.prefetch_weight`` must run between dispatch and expert
+  compute.
+
+PoC limitations:
+- BF16 / unquantized only, eager only.
+- Expert weights are replicated in global-expert order on every rank
+  (memory-heavy). Production Kimi-K3 serving requires sharded
+  symmetric-memory expert ownership, where rows ``[0, E)`` physically alias
+  each home rank's parameter memory.
+- Route weights are applied inside the expert compute and MoonEP's
+  ``combine`` performs the K-sum, so ``finalize`` requires
+  ``TopKWeightAndReduceNoOP``.
+"""
+
+from typing import Any, NamedTuple
+
+import torch
+import torch.nn.functional as F
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceNoOP,
+)
+
+MOONEP_DEFAULT_NUM_PREFETCH_SLOTS = 4
+MOONEP_DEFAULT_TOKEN_PADDING = 128
+MOONEP_DEFAULT_NUM_SMS = 32
+# MoonEP's weight-prefetch kernel tiles hidden and intermediate dims by 128.
+MOONEP_WEIGHT_TILE = 128
+
+
+class MoonEPExpertWeightLayout(NamedTuple):
+    """Contiguous BF16 expert weights in MoonEP ``[E+B, ...]`` layout.
+
+    Rows ``[0, E)`` hold expert weights in global expert order; rows
+    ``[E, E+B)`` are mutable prefetch slots filled by
+    ``Buffer.prefetch_weight``.
+
+    The prefetch slots never need re-zeroing between calls: the planner
+    marks unused slots with ``-1`` in ``plan.experts_to_copy`` (skipped by
+    ``prefetch_weight``) and gives them an empty ``cu_seqlens`` segment, so
+    stale slot contents are never read; used slots are fully overwritten.
+    """
+
+    full_gate_weight: torch.Tensor
+    full_up_weight: torch.Tensor
+    full_down_weight: torch.Tensor
+    num_prefetch_slots: int
+
+
+def make_moonep_weight_layout(
+    w13_weight: torch.Tensor,
+    w2_weight: torch.Tensor,
+    num_prefetch_slots: int,
+) -> MoonEPExpertWeightLayout:
+    """Build the replicated ``[E+B, ...]`` PoC weight layout.
+
+    ``w13_weight`` must be ``[E, 2I, H]`` (gate rows first) and ``w2_weight``
+    ``[E, H, I]``, both BF16 in global expert order.
+    """
+    if num_prefetch_slots <= 0:
+        raise ValueError(
+            f"num_prefetch_slots must be positive, got {num_prefetch_slots}"
+        )
+    if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16:
+        raise NotImplementedError("MoonEP PoC supports BF16 weights only.")
+    num_experts, two_i, hidden_size = w13_weight.shape
+    intermediate_size = two_i // 2
+    if tuple(w2_weight.shape) != (num_experts, hidden_size, intermediate_size):
+        raise ValueError(
+            f"w2_weight shape {tuple(w2_weight.shape)} does not match "
+            f"w13_weight shape {tuple(w13_weight.shape)}"
+        )
+    if hidden_size % MOONEP_WEIGHT_TILE or intermediate_size % MOONEP_WEIGHT_TILE:
+        raise ValueError(
+            "MoonEP weight prefetch requires hidden_size and intermediate_size "
+            f"to be multiples of {MOONEP_WEIGHT_TILE}; got H={hidden_size}, "
+            f"I={intermediate_size}"
+        )
+
+    full_gate_weight = torch.empty(
+        num_experts + num_prefetch_slots,
+        intermediate_size,
+        hidden_size,
+        dtype=torch.bfloat16,
+        device=w13_weight.device,
+    )
+    full_up_weight = torch.empty_like(full_gate_weight)
+    full_down_weight = torch.empty(
+        num_experts + num_prefetch_slots,
+        hidden_size,
+        intermediate_size,
+        dtype=torch.bfloat16,
+        device=w2_weight.device,
+    )
+
+    full_gate_weight[:num_experts].copy_(w13_weight[:, :intermediate_size, :])
+    full_up_weight[:num_experts].copy_(w13_weight[:, intermediate_size:, :])
+    full_down_weight[:num_experts].copy_(w2_weight)
+    full_gate_weight[num_experts:].zero_()
+    full_up_weight[num_experts:].zero_()
+    full_down_weight[num_experts:].zero_()
+
+    return MoonEPExpertWeightLayout(
+        full_gate_weight=full_gate_weight.contiguous(),
+        full_up_weight=full_up_weight.contiguous(),
+        full_down_weight=full_down_weight.contiguous(),
+        num_prefetch_slots=num_prefetch_slots,
+    )
+
+
+def gather_moonep_weight_layout(
+    w13_local: torch.Tensor,
+    w2_local: torch.Tensor,
+    num_global_experts: int,
+    num_prefetch_slots: int,
+) -> MoonEPExpertWeightLayout:
+    """Build the replicated ``[E+B, ...]`` layout from this rank's local experts.
+
+    PoC bridge: each EP rank loads only its own experts (linear placement),
+    so all-gather them once at load time into global expert order on every
+    rank. Production MoonEP instead maps rows ``[0, E)`` onto each home
+    rank's parameter memory via symmetric memory (RFC #52095 item 5).
+    """
+    from vllm.distributed import get_ep_group
+
+    ep_group = get_ep_group()
+    ep_size = ep_group.world_size
+    if ep_size == 1:
+        return make_moonep_weight_layout(w13_local, w2_local, num_prefetch_slots)
+    if w13_local.size(0) * ep_size != num_global_experts:
+        raise NotImplementedError(
+            "MoonEP PoC requires num_experts to be evenly divisible across EP "
+            f"ranks: {num_global_experts} experts, {ep_size} ranks"
+        )
+    w13_global = ep_group.all_gather(w13_local.contiguous(), dim=0)
+    w2_global = ep_group.all_gather(w2_local.contiguous(), dim=0)
+    return make_moonep_weight_layout(w13_global, w2_global, num_prefetch_slots)
+
+
+MOONEP_MIN_DISPATCH_TOKENS = 128
+
+
+class MoonEPBufferPool:
+    """Lazily creates ``moonep.Buffer`` instances at power-of-two capacities.
+
+    ``Buffer`` fixes its token capacity ``S`` at construction, so a single
+    full-capacity buffer would pad every dispatch to
+    ``max_num_batched_tokens``. The pool instead serves the smallest
+    power-of-two capacity that fits the step's token count, so decode-sized
+    batches dispatch a few hundred slots rather than the full capacity.
+
+    Selection must be identical on every EP rank (dispatch is collective and
+    ``Buffer`` construction is itself a collective): callers derive the token
+    count from ``dp_metadata`` (the max across DP ranks) so all ranks create
+    and pick the same buffer at the same step.
+    """
+
+    def __init__(self, buffer_kwargs: dict[str, Any], max_tokens_per_rank: int):
+        self._kwargs = {k: v for k, v in buffer_kwargs.items() if k != "S"}
+        self.max_tokens_per_rank = max_tokens_per_rank
+        self._buffers: dict[int, Any] = {}
+
+    def capacity_for(self, num_tokens: int) -> int:
+        if num_tokens >= self.max_tokens_per_rank:
+            return self.max_tokens_per_rank
+        cap = 1 << max(num_tokens - 1, 0).bit_length()
+        return min(max(cap, MOONEP_MIN_DISPATCH_TOKENS), self.max_tokens_per_rank)
+
+    def get(self, num_tokens: int) -> tuple[int, Any]:
+        """Return ``(capacity, buffer)`` for the given step token count."""
+        from moonep import Buffer  # type: ignore[import-not-found]
+
+        capacity = self.capacity_for(num_tokens)
+        buffer = self._buffers.get(capacity)
+        if buffer is None:
+            buffer = Buffer(S=capacity, **self._kwargs)
+            self._buffers[capacity] = buffer
+        return capacity, buffer
+
+    def destroy(self) -> None:
+        for buffer in self._buffers.values():
+            buffer.destroy()
+        self._buffers.clear()
+
+
+class MoonEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
+    """Prepare/Finalize using MoonEP balanced dispatch/combine.
+
+    ``prepare`` pads the batch to the selected buffer's static token
+    capacity, dispatches, runs ``prefetch_weight`` for the planned redundant
+    experts, and stashes the ``plan`` for ``finalize`` (the same pattern
+    DeepEP-HT uses for its handle). Downstream expert compute must consume
+    the expert-grouped ``[NvS, H]`` layout via ``cu_seqlens``.
+    """
+
+    def __init__(
+        self,
+        buffer_pool: "MoonEPBufferPool",
+        max_tokens_per_rank: int,
+        num_dispatchers: int,
+        num_global_experts: int,
+        weight_layout: MoonEPExpertWeightLayout | None = None,
+    ):
+        super().__init__()
+        self.buffer_pool = buffer_pool
+        self.max_tokens_per_rank = max_tokens_per_rank
+        self.num_dispatchers_ = num_dispatchers
+        self.num_global_experts = num_global_experts
+        self.weight_layout = weight_layout
+        self._fused_experts: Any = None
+
+        # dispatch state consumed by finalize (and the expert runner)
+        self._plan: Any = None
+        self._cu_seqlens: torch.Tensor | None = None
+        self._num_tokens: int = 0
+        self._buffer: Any = None  # the pool buffer used by the current step
+
+    def post_init_setup(self, fused_experts: mk.FusedMoEExperts) -> None:
+        # The [E+B] weight layout is attached to the experts by their
+        # process_weights_after_loading hook (after this runs), so keep the
+        # reference and resolve the layout lazily in prepare().
+        self._fused_experts = fused_experts
+
+    def _resolve_weight_layout(self) -> MoonEPExpertWeightLayout:
+        if self.weight_layout is None and self._fused_experts is not None:
+            layout = getattr(self._fused_experts, "weight_layout", None)
+            if isinstance(layout, MoonEPExpertWeightLayout):
+                self.weight_layout = layout
+        # Redundant experts' weights must be in rows [E, E+B) before the
+        # expert compute reads them; skipping prefetch silently corrupts
+        # output.
+        assert self.weight_layout is not None, (
+            "MoonEPPrepareAndFinalize: weight layout not available (the "
+            "experts' process_weights_after_loading has not run)"
+        )
+        return self.weight_layout
+
+    @property
+    def num_dispatched_slots(self) -> int:
+        """``NvS`` of the buffer selected by the current step's prepare()."""
+        assert self._buffer is not None, "prepare() has not been called"
+        return int(self._buffer._ctx["NvS"])
+
+    def num_dispatchers(self) -> int:
+        return self.num_dispatchers_
+
+    def output_is_reduced(self) -> bool:
+        # combine returns the fully weighted+reduced token-major output
+        return True
+
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def max_num_tokens_per_rank(self) -> int | None:
+        return self.max_tokens_per_rank
+
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        return torch.int32
+
+    def supports_async(self) -> bool:
+        return False
+
+    @property
+    def cu_seqlens(self) -> torch.Tensor:
+        assert self._cu_seqlens is not None, "prepare() has not been called"
+        return self._cu_seqlens
+
+    @property
+    def plan(self) -> Any:
+        assert self._plan is not None, "prepare() has not been called"
+        return self._plan
+
+    def _pad_to_capacity(
+        self,
+        a1: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        capacity: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+        num_tokens = a1.shape[0]
+        if num_tokens > capacity:
+            raise ValueError(
+                f"MoonEP static buffer capacity exceeded: {num_tokens} tokens "
+                f"> S={capacity}."
+            )
+        a1 = a1.contiguous()
+        topk_ids = topk_ids.to(dtype=torch.int32)
+        topk_weights = topk_weights.to(dtype=torch.float32)
+        # MoonEP's planner has no sentinel for invalid expert ids, so map
+        # negative ids (padded/invalid slots from the engine) to expert 0
+        # with zero route weight — the same convention as the capacity
+        # padding rows below. Wasted compute, never wrong output.
+        invalid = topk_ids < 0
+        topk_ids = torch.where(invalid, torch.zeros_like(topk_ids), topk_ids)
+        topk_weights = torch.where(
+            invalid, torch.zeros_like(topk_weights), topk_weights
+        ).contiguous()
+        topk_ids = topk_ids.contiguous()
+        if num_tokens == capacity:
+            return a1, topk_ids, topk_weights, num_tokens
+        pad = capacity - num_tokens
+        return (
+            F.pad(a1, (0, 0, 0, pad)),
+            F.pad(topk_ids, (0, 0, 0, pad)),
+            F.pad(topk_weights, (0, 0, 0, pad)),
+            num_tokens,
+        )
+
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.PrepareResultType:
+        if a1.dtype != torch.bfloat16:
+            raise NotImplementedError("MoonEP PoC supports BF16 hidden states only.")
+        if quant_config.quant_dtype is not None:
+            raise NotImplementedError("MoonEP PoC does not support quantized dispatch.")
+        assert num_experts == self.num_global_experts
+        assert self._plan is None, (
+            "MoonEPPrepareAndFinalize.prepare() called again before finalize()"
+        )
+        if apply_router_weight_on_input:
+            assert topk_ids.size(1) == 1, (
+                "apply_router_weight_on_input requires top-1 routing"
+            )
+            a1 = a1 * topk_weights.to(a1.dtype)
+
+        # Buffer choice must be rank-symmetric: use the max token count
+        # across DP ranks (every rank sees the same dp_metadata), falling
+        # back to the local count outside a forward context (tests).
+        num_tokens = a1.shape[0]
+        try:
+            dp_meta = get_forward_context().dp_metadata
+        except AssertionError:
+            dp_meta = None
+        if dp_meta is not None:
+            num_tokens = int(dp_meta.num_tokens_across_dp_cpu.max())
+        capacity, self._buffer = self.buffer_pool.get(num_tokens)
+
+        a1, topk_ids, topk_weights, self._num_tokens = self._pad_to_capacity(
+            a1, topk_ids, topk_weights, capacity
+        )
+        tokens_per_expert = torch.bincount(
+            topk_ids.reshape(-1).to(dtype=torch.int64),
+            minlength=num_experts,
+        ).to(dtype=torch.int32)
+
+        hidden_nvsh, route_weights_nvs, cu_seqlens, plan = self._buffer.dispatch(
+            a1,
+            topk_weights,
+            topk_ids,
+            tokens_per_expert,
+        )
+        self._plan = plan
+        self._cu_seqlens = cu_seqlens
+
+        weight_layout = self._resolve_weight_layout()
+        self._buffer.prefetch_weight(
+            plan=plan,
+            full_gate_weight=weight_layout.full_gate_weight,
+            full_up_weight=weight_layout.full_up_weight,
+            full_down_weight=weight_layout.full_down_weight,
+        )
+
+        # Segment sizes per [E+B] weight row. NOTE: this is per *row*, not
+        # per local expert, and NvS rows are expert-grouped rather than
+        # token-major — only MoonEP-aware expert implementations can consume
+        # this activation layout.
+        expert_num_tokens = torch.diff(cu_seqlens, prepend=cu_seqlens.new_zeros(1))
+        expert_tokens_meta = mk.ExpertTokensMetadata(
+            expert_num_tokens=expert_num_tokens,
+            expert_num_tokens_cpu=None,
+        )
+
+        # MoonEP has no post-dispatch per-token topk ids/weights; route
+        # weights come back in NvS order.
+        return hidden_nvsh, None, expert_tokens_meta, None, route_weights_nvs
+
+    def finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> None:
+        # Route weights are applied in the expert compute and MoonEP's
+        # combine performs the K-sum, so any weight/reduce work in finalize
+        # means MoonEP was paired with the wrong kind of experts.
+        assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP), (
+            "MoonEP requires TopKWeightAndReduceNoOP, got "
+            f"{type(weight_and_reduce_impl).__name__}"
+        )
+        combined, _, _ = self._buffer.combine(
+            plan=self.plan,
+            hidden_nvsh=fused_expert_output,
+            route_weights_nvs=None,
+        )
+        output.copy_(combined[: self._num_tokens])
+        self._plan = None
+        self._cu_seqlens = None
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py
index fe3a53941806..75b97c69aafb 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py
@@ -13,9 +13,7 @@
 
 
 class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
-    """
-    Prepare/Finalize using MoRI kernels.
-    """
+    """Prepare/Finalize using MoRI kernels."""
 
     def __init__(
         self,
@@ -60,8 +58,7 @@ def prepare(
         quant_config: FusedMoEQuantConfig,
         defer_input_quant: bool = False,
     ) -> mk.PrepareResultType:
-        """
-        Returns a tuple of:
+        """Returns a tuple of:
         - quantized + dispatched a.
         - Optional quantized + dispatched a1_scales.
         - Optional ExpertTokensMetadata containing gpu/cpu tensors
@@ -85,6 +82,15 @@ def prepare(
             elif quant_config.is_per_act_token:
                 quant_func = get_hip_quant(QuantType.per_Token)
                 a1, scale = quant_func(a1, quant_dtype=current_platform.fp8_dtype())
+            elif quant_config.is_per_tensor:
+                quant_func = get_hip_quant(QuantType.per_Tensor)
+                a1, scale = quant_func(
+                    a1,
+                    scale=quant_config.a1_scale,
+                    quant_dtype=current_platform.fp8_dtype(),
+                )
+                # mori expects one scale slot per token; broadcast.
+                scale = scale.expand(a1.shape[0], 1).contiguous()
 
         (
             dispatch_a1,
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
index 89f3843cc504..97d4e0d284c6 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
@@ -69,8 +69,7 @@ def _unwrap_scale_and_prepare_for_moe(
 
 
 class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular):
-    """
-    Naive Prepare/Finalize for Dp/Ep case for Modular Kernels.
+    """Naive Prepare/Finalize for Dp/Ep case for Modular Kernels.
 
     Uses Torch AR/RS or AR for dispatch/combine operations, applied
     to the topk weights and ids.
@@ -121,7 +120,6 @@ def prepare(
         defer_input_quant: bool = False,
     ) -> mk.PrepareResultType:
         """Quantize and Dispatch Topk Weights and Topk Ids."""
-
         if apply_router_weight_on_input:
             topk = topk_ids.size(1)
             assert topk == 1, (
@@ -210,8 +208,7 @@ def finalize(
 
 
 class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic):
-    """
-    Naive Prepare/Finalize for Dp/Ep case for Modular Kernels.
+    """Naive Prepare/Finalize for Dp/Ep case for Modular Kernels.
 
     Uses Torch AR/RS or AR for dispatch/combine operations, applied
     to the router logits (the MoE kernel runs the router internally).
@@ -250,7 +247,6 @@ def prepare(
         defer_input_quant: bool = False,
     ) -> mk.PrepareMonolithicResultType:
         """Quantize and Dispatch Router Logits."""
-
         a1q, scales, a1q_scale_orig = _quantize_and_setup_dispatch(
             a1, quant_config, defer_input_quant
         )
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py
index fd7d45655dd6..6d5dbee3f491 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py
@@ -6,8 +6,6 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from vllm.distributed import get_ep_group
-from vllm.distributed.device_communicators.all2all import NixlEPAll2AllManager
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -35,9 +33,7 @@
 def dequant_fp8(
     expert_x_fp8: torch.Tensor, expert_x_scales: torch.Tensor
 ) -> torch.Tensor:
-    """
-    Return dequantized tensor in fp32
-    """
+    """Return dequantized tensor in fp32."""
     assert expert_x_fp8.is_contiguous()
     expert_x_scales = expert_x_scales.contiguous()
     num_experts = expert_x_fp8.size(0)
@@ -50,9 +46,7 @@ def dequant_fp8(
 
 
 class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
-    """
-    Prepare/Finalize using NIXL EP kernels.
-    """
+    """Prepare/Finalize using NIXL EP kernels."""
 
     # NIXL EP kernels are compiled only for certain specific hidden sizes.
     # NOTE: Keep this list sorted, maybe_roundup_layer_hidden_size depends
@@ -79,6 +73,7 @@ def __init__(
         buffer: nixl_ep.Buffer,
         max_tokens_per_rank: int,
         num_dispatchers: int,
+        expert_capacity: int,
         use_fp8_dispatch: bool = False,
         global_to_physical: torch.Tensor | None = None,
         physical_to_global: torch.Tensor | None = None,
@@ -94,6 +89,7 @@ def __init__(
         # combine function.
         self.handles: list[tuple | None] = [None, None]
         self.num_dispatchers_ = num_dispatchers
+        self.expert_capacity = expert_capacity
 
         topk_indices_dtype = self.topk_indices_dtype()
 
@@ -141,17 +137,6 @@ def activation_format(self) -> mk.FusedMoEActivationFormat:
     def max_num_tokens_per_rank(self) -> int | None:
         return self.max_tokens_per_rank
 
-    def on_commit(self) -> None:
-        device_communicator = get_ep_group().device_communicator
-        assert device_communicator is not None
-        all2all_manager = device_communicator.all2all_manager
-        assert isinstance(all2all_manager, NixlEPAll2AllManager)
-        # maybe_make_prepare_finalize(..., eep_stage=True) initializes self.buffer
-        # with get_handle(..., stage=True), which stages global NIXL state for the
-        # new config but leaves it inactive while the old config remains active.
-        # When EEP commit switches to this P/F, this P/F needs to commit that state.
-        all2all_manager.commit_staged_state()
-
     def topk_indices_dtype(self) -> torch.dtype | None:
         return NIXL_EP_TOPK_INDICES_DTYPE
 
@@ -267,7 +252,7 @@ def prepare_async(
             a1,
             dispatch_topk_ids,
             self.max_tokens_per_rank,
-            num_experts,
+            self.expert_capacity,
             use_fp8=self.use_fp8_dispatch,
             # round_scale needs to be set to dispatch in ue8m0
             round_scale=self.use_ue8m0_dispatch,
diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py
index 5d284bcc80f3..e3a287ddc2da 100644
--- a/vllm/model_executor/layers/fused_moe/routed_experts.py
+++ b/vllm/model_executor/layers/fused_moe/routed_experts.py
@@ -27,6 +27,7 @@
     QuantizationConfig,
     resolve_quant_method,
 )
+from vllm.utils.math_utils import cdiv
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts
@@ -44,8 +45,7 @@ class FusedMoeWeightScaleSupported(Enum):
 
 @PluggableLayer.register("routed_experts")
 class RoutedExperts(PluggableLayer):
-    """
-    Container for routed expert weights and execution logic.
+    """Container for routed expert weights and execution logic.
 
     This module owns the expert weight parameters (w13_weight, w2_weight, scales, etc.)
     and handles:
@@ -188,8 +188,7 @@ def _get_quant_method(
         quant_config: QuantizationConfig | None,
         moe_config: FusedMoEConfig,
     ) -> FusedMoEMethodBase:
-        """
-        Helper method to ensure quant_method is never None and
+        """Helper method to ensure quant_method is never None and
         of the proper type.
         """
         quant_method = None
@@ -318,8 +317,7 @@ def _load_combined_w13_weight_scale(
         param: torch.Tensor,
         tp_rank: int,
     ):
-        """
-        Load w13 weight scales assuming that w1 weight scales and w3 weight
+        """Load w13 weight scales assuming that w1 weight scales and w3 weight
         scales are stored in the same loaded_weight tensor.
         """
         shard_size = param.shape[shard_dim]
@@ -335,9 +333,9 @@ def _load_model_weight_or_group_weight_scale(
         shard_id: str,
         loaded_weight: torch.Tensor,
         tp_rank: int,
+        is_scale: bool = False,
     ):
-        """
-        Load grouped weight scales for group quantization or model weights
+        """Load grouped weight scales for group quantization or model weights
 
         Args:
             shard_dim: dimension to shard
@@ -345,13 +343,41 @@ def _load_model_weight_or_group_weight_scale(
             shard_id: either w1, w2, or w3
             loaded_weight: checkpoint weight to load into the param
             tp_rank: tensor parallel rank
+            is_scale: whether padding should use unit scales instead of zero weights.
+
         """
+        padded_tp = self.moe_config.tp_shard_with_padding
+        if padded_tp:
+            destination = expert_data
+            if shard_id in ("w1", "w3") and self.moe_config.is_act_and_mul:
+                half = destination.shape[shard_dim] // 2
+                destination = destination.narrow(
+                    shard_dim, 0 if shard_id == "w1" else half, half
+                )
+            shard_size = destination.shape[shard_dim]
+            expected_size = (
+                self.moe_config.intermediate_size
+                * shard_size
+                // self.moe_config.intermediate_size_per_partition
+            )
+            if loaded_weight.shape[shard_dim] != expected_size:
+                raise ValueError(
+                    "Block-aligned TP loading expects an unsharded checkpoint "
+                    f"projection of size {expected_size}, got "
+                    f"{loaded_weight.shape[shard_dim]}."
+                )
+            start = min(tp_rank * shard_size, expected_size)
+            size = min(shard_size, expected_size - start)
+            loaded_weight = loaded_weight.narrow(shard_dim, start, size)
+            destination.fill_(1 if is_scale else 0)
+
         if shard_id == "w2":
             self._load_w2(
                 shard_dim=shard_dim,
                 loaded_weight=loaded_weight,
                 expert_data=expert_data,
                 tp_rank=tp_rank,
+                load_full=padded_tp,
             )
         elif shard_id in ("w1", "w3"):
             self._load_w13(
@@ -360,6 +386,7 @@ def _load_model_weight_or_group_weight_scale(
                 loaded_weight=loaded_weight,
                 expert_data=expert_data,
                 tp_rank=tp_rank,
+                load_full=padded_tp,
             )
 
     def _load_per_channel_weight_scale(
@@ -444,6 +471,7 @@ def _narrow_expert_data_for_padding(
                 Must be non-negative.
             shard_dim: The dimension index corresponding to the shard
                 (intermediate) dimension. Defaults to `None`.
+
         """
         dims = (hidden_dim,) if shard_dim is None else (hidden_dim, shard_dim)
         if loaded_weight.ndim > 0:
@@ -513,11 +541,12 @@ def _load_w2(
         shard_dim: int,
         loaded_weight: torch.Tensor,
         tp_rank: int,
+        load_full: bool = False,
     ):
         # Index the loaded weight for tp sharding.
         # down_proj: "RowParallel" so tp sharding on input_dim
-        # Only narrow if the loaded_weight is not a scalar (0-dim tensor).
-        if loaded_weight.ndim > 0:
+        # Padded TP weights have already been sliced by the grouped loader.
+        if not load_full and loaded_weight.ndim > 0:
             # Same padding fix as _load_w13: use unpadded per-rank size.
             tp_size = self.moe_config.moe_parallel_config.tp_size
             loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size
@@ -538,6 +567,23 @@ def _load_w2(
             hidden_dim=hidden_dim,
             shard_dim=shard_dim,
         )
+        if (
+            loaded_weight.device.type == "cpu"
+            and expert_data.device.type == "cuda"
+            and not expert_data.is_contiguous()
+            and expert_data.ndim == 2
+        ):
+            # A strided CPU-to-CUDA copy can allocate a full-shard CUDA
+            # temporary. Make the TP slice contiguous on CPU and copy rows
+            # in chunks to limit the temporary used for the padded view.
+            loaded_weight = loaded_weight.contiguous()
+            num_chunks = max(cdiv(loaded_weight.nbytes, 1 << 20), 1)
+            for dst, src in zip(
+                expert_data.chunk(num_chunks, dim=0),
+                loaded_weight.chunk(num_chunks, dim=0),
+            ):
+                dst.copy_(src)
+            return
         expert_data.copy_(loaded_weight)
 
     def _load_single_value(
@@ -615,15 +661,15 @@ def weight_loader(
         # compressed-tensors checkpoints with packed weights are stored flipped
         # TODO (mgoin): check self.quant_method.quant_config.quant_format
         # against known CompressionFormat enum values that have this quality
-        if quant_method_name in (
-            "CompressedTensorsWNA16MoEMethod",
-            "CompressedTensorsWNA16RDNA3MoEMethod",
-            "CompressedTensorsW4A16FlydslMoEMethod",
+        if (
+            quant_method_name
+            in (
+                "CompressedTensorsWNA16MoEMethod",
+                "CompressedTensorsW4A16FlydslMoEMethod",
+            )
+            and is_transposed
         ):
-            if is_transposed:
-                loaded_weight = loaded_weight.t().contiguous()
-            else:
-                loaded_weight = loaded_weight
+            loaded_weight = loaded_weight.t().contiguous()
 
         if shard_id not in ("w1", "w2", "w3"):
             raise ValueError(f"shard_id must be ['w1','w2','w3'] but got {shard_id}.")
@@ -802,6 +848,7 @@ def weight_loader(
                     loaded_weight=loaded_weight,
                     expert_data=expert_data,
                     tp_rank=self.moe_config.tp_rank,
+                    is_scale=True,
                 )
             elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value:
                 self._load_per_tensor_weight_scale(
@@ -990,8 +1037,7 @@ def build_expert_params_mapping(
         lora_base_layer_prefix_on_param_name: str = "",
         include_fused: bool = False,
     ) -> list[tuple[str, str, int, str]]:
-        """
-        Create expert parameter mapping for weight loading with redundant experts.
+        """Create expert parameter mapping for weight loading with redundant experts.
 
         This mapping handles the physical-to-logical expert ID conversion needed
         when loading weights with EPLB redundant experts.
@@ -1011,6 +1057,7 @@ def build_expert_params_mapping(
                 ``make_expert_params_mapping`` indexes the model-wide
                 ``params_dict`` (prefix included).
             include_fused: Prepend the fused pre-fused-checkpoint entries
+            routed_experts_prefix: Prefix of the routed experts submodule
 
         Returns:
             List of tuples (param_name, weight_name, expert_id, shard_id)
@@ -1019,6 +1066,7 @@ def build_expert_params_mapping(
             - weight_name: Weight name in checkpoint
             - expert_id: Physical expert ID
             - shard_id: Shard identifier (w1, w2, w3)
+
         """
         num_physical_experts = num_experts + num_redundant_experts
 
@@ -1097,8 +1145,7 @@ def get_expert_weights(self) -> Iterable[torch.Tensor]:
         def _maybe_make_contiguous(
             name: str, p: torch.nn.Parameter
         ) -> torch.nn.Parameter:
-            """
-            In some cases, the last 2 dimensions (the non-expert dimensions)
+            """In some cases, the last 2 dimensions (the non-expert dimensions)
             of the weight scale tensor are transposed. This function
             transforms the tensor (view update) so the tensor is contiguous().
             Example: A non-contiguous scale tensor,
@@ -1178,8 +1225,7 @@ def forward_modular(
         shared_experts: "SharedExperts | None" = None,
         shared_experts_input: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Execute routed experts using the quantization method's apply function.
+        """Execute routed experts using the quantization method's apply function.
 
         This is called by the runner after router selection (for modular kernels)
         quant_method.apply() which accesses the weights on this RoutedExperts
@@ -1194,6 +1240,7 @@ def forward_modular(
 
         Returns:
             Output tensor from routed experts.
+
         """
         assert not self.quant_method.is_monolithic
 
@@ -1213,8 +1260,7 @@ def forward_monolithic(
         router_logits: torch.Tensor | None = None,
         input_ids: torch.Tensor | None = None,
     ) -> torch.Tensor | UnfinalizedMoEOutput:
-        """
-        Execute routed experts using the quantization method's apply function.
+        """Execute routed experts using the quantization method's apply function.
 
         This is called by the runner after router selection (for modular kernels)
         or with router logits (for monolithic kernels). It delegates to
@@ -1228,6 +1274,7 @@ def forward_monolithic(
 
         Returns:
             Finalized routed states or a deferred-finalize output.
+
         """
         assert self.quant_method.is_monolithic
 
diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
index edd3e3beaa19..c49d614549d8 100644
--- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
@@ -10,15 +10,12 @@
 from functools import partial
 from typing import Protocol, runtime_checkable
 
-import numpy as np
 import torch
 
 from vllm.config import VllmConfig
 from vllm.distributed.parallel_state import get_tp_group
 from vllm.forward_context import get_forward_context
 from vllm.platforms import current_platform
-from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec
-from vllm.v1.outputs import RoutedExpertsTensors
 
 logger = logging.getLogger(__name__)
 
@@ -46,37 +43,29 @@ def _get_routed_experts_shape(vllm_config: VllmConfig) -> tuple[int, int, int]:
 class RoutedExpertsCapturer:
     """Worker-side capturer for routed experts, lives on GPU.
 
-    Layer-level hooks call :meth:`capture` from inside the forward pass
-    with the per-layer ``topk_ids`` tensor. The tensor is sliced to the
-    tokens owned by this DP rank and written into a preallocated device
-    buffer. At the end of the step, :class:`GPUModelRunner` reads the
-    device buffer, issues a D2H copy into a pinned CPU buffer, and hands
-    the result to the scheduler via :class:`RoutedExpertsLists`.
+    Layer-level hooks call :meth:`capture` inside the forward pass. Routing
+    rows owned by this DP rank are written into a preallocated device buffer.
 
-    The device / pinned-CPU transit buffers use ``torch.int32`` (not a
-    narrow ``uint8``/``uint16`` sized by ``num_experts``). This keeps the
-    SP all-gather path free of dtype casts, matches the router's native
-    ``topk_ids`` indices dtype more closely, and costs only a few MB per
-    worker (``max_num_batched_tokens * num_layers * top_k * 4`` bytes).
-    The scheduler-side slot buffer
-    (``RoutedExpertsManager.routed_experts_by_slot``) still uses the
-    narrow dtype -- numpy fancy-index assignment in ``store_batch``
-    narrows the data on the way in.
+    The device buffer uses ``int32``. Stable snapshots use the narrowest dtype
+    that can represent every logical expert ID.
 
     Invariants:
         - One instance per worker; shape is fixed at init and covers the
           worst-case step (``max_num_batched_tokens`` tokens).
         - Every routed layer overwrites the current step's token rows.
-        - ``device_buffer.dtype`` is ``torch.int32``.
     """
 
     def __init__(
         self,
         max_num_batched_tokens: int,
         vllm_config: VllmConfig,
-        kv_cache_config: KVCacheConfig,
     ) -> None:
-        num_layers, _, num_experts_per_tok = _get_routed_experts_shape(vllm_config)
+        num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape(
+            vllm_config
+        )
+        self.shape_per_token = (num_layers, num_experts_per_tok)
+        self.output_dtype_name = "uint8" if num_experts <= 256 else "uint16"
+        self.output_dtype = getattr(torch, self.output_dtype_name)
         logger.info(
             "RoutedExpertsCapturer: allocating buffer with "
             "max_tokens=%d, num_layers=%d, num_experts_per_tok=%d "
@@ -87,23 +76,12 @@ def __init__(
             vllm_config.model_config.hf_text_config.model_type,
         )
         self.device_buffer = torch.zeros(
-            (
-                max_num_batched_tokens,
-                num_layers,
-                num_experts_per_tok,
-            ),
-            # Use int32 for the device / host transit buffers: it
-            # matches the router's native topk_ids dtype, is universally
-            # supported by NCCL (uint8/uint16 are version-dependent),
-            # and the extra bytes are small (few MB per worker). The
-            # big scheduler-side slot buffer stays narrow.
+            (max_num_batched_tokens, *self.shape_per_token),
             dtype=torch.int32,
             device=current_platform.device_type,
         )
         self.dp_rank = vllm_config.parallel_config.data_parallel_rank
         self.tp_size = vllm_config.parallel_config.tensor_parallel_size
-        # KV cache group whose slot layout the routing data is keyed by.
-        self.attn_gid = get_routed_experts_attn_gid(kv_cache_config)
 
     def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
         """Capture expert routing decisions for a specific layer.
@@ -136,34 +114,24 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
         Args:
             layer_id: The layer index.
             topk_ids: Tensor of shape (batch_size, num_routed_experts).
-        """
 
+        """
         ctx = get_forward_context()
-        if ctx.dp_metadata is None:  # single dp
-            start_loc = 0
-            end_loc = topk_ids.shape[0]
-            token_num_per_dp = topk_ids.shape[0]
-        else:  # multi dp
+        if ctx.dp_metadata is None:
+            local_topk_ids = topk_ids
+        else:
             num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu
-            token_num_per_dp = int(num_tokens_dp[self.dp_rank].item())
+            num_local_tokens = int(num_tokens_dp[self.dp_rank].item())
             total = int(num_tokens_dp.sum().item())
             n = topk_ids.shape[0]
             shard_sizes = getattr(ctx.dp_metadata, "local_sizes", None)
             gathered_size = sum(shard_sizes) if shard_sizes is not None else -1
 
             if n == total:
-                # Naive dispatch: all DP ranks' tokens concatenated
-                # before routing. This rank owns tokens
-                # [end_loc - token_num_per_dp, end_loc).
-                cumsum = torch.cumsum(num_tokens_dp, dim=0)
-                end_loc = int(cumsum[self.dp_rank].item())
-                start_loc = end_loc - token_num_per_dp
-            elif n == token_num_per_dp:
-                # Modular-kernel path: DP combine happens inside
-                # quant_method.apply; select_experts only sees this
-                # rank's tokens, take the whole tensor.
-                start_loc = 0
-                end_loc = token_num_per_dp
+                end = int(num_tokens_dp[: self.dp_rank + 1].sum().item())
+                local_topk_ids = topk_ids[end - num_local_tokens : end]
+            elif n == num_local_tokens:
+                local_topk_ids = topk_ids
             elif shard_sizes is not None and n == gathered_size:
                 # Naive DP+EP dispatch gathers the sequence-parallel shards in
                 # flattened DP-then-TP order. ``local_sizes`` is the exact
@@ -173,12 +141,11 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
                 assert len(shard_sizes) % num_dp_ranks == 0
                 shards_per_dp_rank = len(shard_sizes) // num_dp_ranks
                 first_shard = self.dp_rank * shards_per_dp_rank
-                start_loc = sum(shard_sizes[:first_shard])
-                end_loc = start_loc + token_num_per_dp
+                start = sum(shard_sizes[:first_shard])
+                local_topk_ids = topk_ids[start : start + num_local_tokens]
             elif (
                 self.tp_size > 1
-                and n != token_num_per_dp
-                and n == (token_num_per_dp + self.tp_size - 1) // self.tp_size
+                and n == (num_local_tokens + self.tp_size - 1) // self.tp_size
             ):
                 # SP + modular-kernel path. All-gather across the TP
                 # group along dim=0 to reconstruct the full per-DP-rank
@@ -193,59 +160,30 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
                 #
                 # ``topk_ids`` is already whatever the router produced
                 # (typically int32/int64, both supported by NCCL); the
-                # downstream ``device_buffer[...] = topk_ids[...]``
-                # setitem narrows into int32 automatically.
-                topk_ids = get_tp_group().all_gather(topk_ids, dim=0)
-                start_loc = 0
-                end_loc = token_num_per_dp
+                # downstream buffer assignment narrows to the capture dtype.
+                local_topk_ids = get_tp_group().all_gather(topk_ids, dim=0)[
+                    :num_local_tokens
+                ]
             else:
-                sp_expected = (
-                    (token_num_per_dp + self.tp_size - 1) // self.tp_size
-                    if self.tp_size > 0
-                    else -1
-                )
+                sp_expected = (num_local_tokens + self.tp_size - 1) // self.tp_size
                 raise AssertionError(
                     "RoutedExpertsCapturer: unexpected topk_ids batch "
-                    f"dim {n} (expected {total}, {token_num_per_dp}, "
+                    f"dim {n} (expected {total}, {num_local_tokens}, "
                     f"{gathered_size}, or {sp_expected} for "
                     f"dp_rank={self.dp_rank}, tp_size={self.tp_size})"
                 )
 
-        if layer_id >= self.device_buffer.shape[1]:
+        if not 0 <= layer_id < self.device_buffer.shape[1]:
             raise IndexError(
                 f"routed-experts layer {layer_id} exceeds capture buffer "
                 f"layer count {self.device_buffer.shape[1]}"
             )
 
-        self.device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[
-            start_loc:end_loc, :
-        ]
-
-    def get_device_buffer(self) -> torch.Tensor:
-        """Return the underlying device buffer so the model runner can
-        issue the D2H copy. The tensor is shared; callers must either
-        clone or fully drain it before the next forward pass overwrites it.
-        """
-        return self.device_buffer
-
-    def get_routed_experts(
-        self, slot_mappings: torch.Tensor, num_tokens: int
-    ) -> RoutedExpertsTensors:
-        """Snapshot this step's routing data and its attention slot mapping.
-
-        Both tensors are cloned since the capture buffer and the shared
-        ``slot_mappings`` are overwritten by the next step, which may race
-        with an in-flight D2H copy.
+        self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids
 
-        Args:
-            slot_mappings: Per-KV-cache-group slot mappings for this step,
-                shape ``(num_kv_cache_groups, max_num_batched_tokens)``.
-            num_tokens: Total number of tokens scheduled in this step.
-        """
-        return RoutedExpertsTensors(
-            routing_data=self.device_buffer[:num_tokens].clone(),
-            slot_mapping=slot_mappings[self.attn_gid, :num_tokens].clone(),
-        )
+    def snapshot_routing_data(self, num_tokens: int) -> torch.Tensor:
+        """Return a stable snapshot of the current routing data."""
+        return self.device_buffer[:num_tokens].to(self.output_dtype)
 
 
 def bind_routed_experts_capturer(
@@ -267,20 +205,12 @@ def bind_routed_experts_capturer(
             continue
         if not isinstance(module, MoERunner):
             continue
-        layer_id = module.layer_id
-
-        def capture_fn(
-            topk_ids: torch.Tensor,
-            layer_id: int = layer_id,
-            capturer: RoutedExpertsCapturer = capturer,
-        ) -> None:
-            capturer.capture(layer_id, topk_ids)
-
+        capture_fn = partial(capturer.capture, module.layer_id)
         quant_method = module._quant_method
-        moe_kernel = getattr(quant_method, "moe_kernel", None)
-        impl = getattr(moe_kernel, "impl", None)
-        fused_experts = getattr(impl, "fused_experts", None)
         if quant_method.is_monolithic:
+            moe_kernel = getattr(quant_method, "moe_kernel", None)
+            impl = getattr(moe_kernel, "impl", None)
+            fused_experts = getattr(impl, "fused_experts", None)
             if not (
                 isinstance(fused_experts, FusedMoEExpertsMonolithic)
                 and fused_experts.supports_routing_replay_capture()
@@ -302,131 +232,3 @@ def capture_fn(
 
     if num_bound == 0:
         raise ValueError("No supported MoE router found for routed-experts capture.")
-
-
-def get_routed_experts_attn_gid(kv_cache_config: KVCacheConfig) -> int:
-    """Return the full-attention KV cache group used for routed experts."""
-    for gid, group in enumerate(kv_cache_config.kv_cache_groups):
-        if is_full_attention_spec(group.kv_cache_spec):
-            return gid
-    raise ValueError("Routed-experts capture requires a full-attention KV cache group.")
-
-
-class RoutedExpertsManager:
-    """Scheduler-side slot-indexed buffer for routed experts.
-
-    Lives on CPU in the scheduler process. Each slot corresponds to
-    ``block_id * block_size + offset_in_block`` where ``block_id`` is
-    drawn from the physical KV-cache block pool, so routing data is
-    tied to physical blocks and naturally survives preemption for
-    prefix-cached blocks (prefix hits re-expose the same slots).
-
-    Data flow per step:
-      1. Worker D2Hs its device capture buffer into
-         :class:`RoutedExpertsLists` and returns it via
-         :class:`ModelRunnerOutput`.
-      2. Scheduler calls :meth:`store_batch` with that step's
-         ``(routing_data, slot_mapping)`` — a single CPU->CPU
-         fancy-index assign, ~few MB per step.
-      3. On request completion / abort / preemption, the scheduler
-         calls :meth:`get` with the request's block IDs to recover
-         the full per-token routing.
-
-    Memory: ``routed_experts_by_slot`` is sized for the whole block
-    pool (``num_blocks * block_size`` slots). For large block pools
-    this can reach multiple GB; see the init log for the exact size.
-    """
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        kv_cache_config: KVCacheConfig,
-    ) -> None:
-        # Hybrid models also have KV groups whose slot layout differs.
-        self.attn_gid = get_routed_experts_attn_gid(kv_cache_config)
-        attn_group = kv_cache_config.kv_cache_groups[self.attn_gid]
-        self.block_size = attn_group.kv_cache_spec.block_size
-
-        # All kv_cache_groups share the same physical block pool, so
-        # block IDs span [0, num_blocks) regardless of how many groups
-        # exist. Sizing to the full pool avoids index-out-of-range
-        # when different groups happen to land on the same block.
-        num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape(
-            vllm_config
-        )
-        max_num_slots = kv_cache_config.num_blocks * self.block_size
-        # Expert IDs are 0..num_experts-1; uint8 fits 256 distinct
-        # values so the boundary is ``<= 256`` (NOT ``< 256``). Keeping
-        # this narrow matters because the slot buffer is sized for the
-        # whole block pool and can reach multiple GB.
-        expert_id_dtype = np.uint8 if num_experts <= 256 else np.uint16
-        self.routed_experts_by_slot = np.zeros(
-            (
-                max_num_slots,
-                num_layers,
-                num_experts_per_tok,
-            ),
-            dtype=expert_id_dtype,
-        )
-        logger.info(
-            "RoutedExpertsManager CPU buffer: %.2f GB "
-            "(slots=%d, layers=%d, top_k=%d, dtype=%s)",
-            self.routed_experts_by_slot.nbytes / 1e9,
-            max_num_slots,
-            num_layers,
-            num_experts_per_tok,
-            self.routed_experts_by_slot.dtype.name,
-        )
-
-    def store_batch(self, data: np.ndarray, slot_mapping: np.ndarray) -> None:
-        """Persist one step's routed experts into the slot buffer.
-
-        Equivalent to ``slot_buffer[slot_mapping] = data``; numpy fancy
-        indexing handles repeated / out-of-order indices. Called once
-        per scheduler step in ``update_from_output``.
-        """
-        self.routed_experts_by_slot[slot_mapping] = data
-
-    def get(
-        self,
-        block_ids: list[int],
-        num_tokens: int,
-        token_start: int = 0,
-    ) -> np.ndarray:
-        """Read routed experts data for a completed / preempted request.
-
-        Reconstructs a per-token slot_mapping from the request's block
-        IDs and returns the routing slice. Because numpy fancy indexing
-        returns a **copy** (not a view), the returned ndarray is safe
-        to hold across subsequent :meth:`store_batch` calls — do not
-        replace the fancy index with a slice without re-verifying.
-
-        Args:
-            block_ids: Block IDs from the attention KV-cache group.
-            num_tokens: Number of tokens that have gone through a forward
-                pass and therefore have routing data written to their
-                slots (typically ``request.num_tokens - 1``; the last
-                sampled token has not been forwarded yet). Slots beyond
-                ``request.num_computed_tokens`` are zero-initialized.
-            token_start: Skip the first ``token_start`` tokens from the
-                result. The slot_mapping is sliced before the fancy-index
-                read, so only the requested slots are fetched — no large
-                intermediate array is allocated. Clamped to
-                ``[0, num_tokens]`` automatically.
-
-        Returns:
-            Array of shape (num_tokens - token_start, num_layers,
-            num_experts_per_tok).
-        """
-        bs = self.block_size
-        block_ids_array = np.array(block_ids, dtype=np.int32)
-        block_offsets = np.arange(bs)
-        # slot = block_id * block_size + offset_in_block; flatten the
-        # (num_blocks, block_size) grid and trim to num_tokens, then
-        # skip the first token_start entries so only the requested
-        # range is fetched in a single fancy-index read.
-        slot_mapping = (
-            block_ids_array.reshape(-1, 1) * bs + block_offsets.reshape(1, -1)
-        ).flatten()[:num_tokens]
-        slot_mapping = slot_mapping[token_start:]
-        return self.routed_experts_by_slot[slot_mapping]
diff --git a/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py
index 3447f7ce790a..42d4b5d4ad33 100644
--- a/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py
+++ b/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py
@@ -16,8 +16,7 @@
 
 
 class AiterSharedRoutedFusedMoERouter(BaseRouter):
-    """
-    ROCm AITER router for models with fused shared experts (e.g. Qwen3-MoE).
+    """ROCm AITER router for models with fused shared experts (e.g. Qwen3-MoE).
 
     When the AITER topk_softmax kernel supports sigmoid fusion, the routing
     softmax and shared-expert sigmoid are computed in a single kernel launch.
diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py
index 01e674b2b131..b51dbf90dc7e 100644
--- a/vllm/model_executor/layers/fused_moe/router/base_router.py
+++ b/vllm/model_executor/layers/fused_moe/router/base_router.py
@@ -13,7 +13,7 @@
 from vllm.triton_utils import tl, triton
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
-if current_platform.is_cuda_alike():
+if current_platform.is_cuda_alike() or current_platform.is_xpu():
 
     @triton.jit
     def _eplb_map_and_record_i32_kernel(
@@ -157,8 +157,7 @@ def eplb_map_to_physical_and_record(
 
 
 class BaseRouter(FusedMoERouter):
-    """
-    Base router class that provides common functionality for all router implementations.
+    """Base class providing common functionality for all router implementations.
 
     This class implements the template method pattern where select_experts() handles
     common pre-processing and post-processing, delegating the actual routing logic
@@ -171,11 +170,11 @@ def __init__(
         global_num_experts: int,
         eplb_state: EplbLayerState | None = None,
     ):
-        """
-        Args:
-            top_k: Number of experts to select per token
-            global_num_experts: Total number of experts
-            eplb_state: Optional EPLBLayerState for load balancing
+        """Args:
+        top_k: Number of experts to select per token
+        global_num_experts: Total number of experts
+        eplb_state: Optional EPLBLayerState for load balancing
+
         """
         super().__init__(eplb_state=eplb_state)
         self.top_k = top_k
@@ -241,8 +240,7 @@ def _compute_routing(
         *,
         input_ids: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Compute the actual routing logic.
+        """Compute the actual routing logic.
 
         This method must be implemented by subclasses to provide the specific
         routing algorithm (e.g., grouped_topk, fused_topk, custom routing, etc.).
@@ -251,9 +249,11 @@ def _compute_routing(
             hidden_states: Input hidden states
             router_logits: Router logits for expert selection
             indices_type: Desired dtype for expert indices (may be None)
+            input_ids: Token ids, for routers that condition on them
 
         Returns:
             tuple of (topk_weights, topk_ids)
+
         """
         raise NotImplementedError
 
@@ -265,8 +265,7 @@ def _select_experts(
         *,
         input_ids: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Route the input hidden states to the top-k experts based on the
+        """Route the input hidden states to the top-k experts based on the
         router logits.
 
         This method implements the template method pattern:
@@ -283,6 +282,7 @@ def _select_experts(
             **Compatibility**: When EPLB is not enabled, the returned ids are
             equivalent to global logical ids, so should be compatible with
             plain MoE implementations without redundant experts.
+
         """
         # Step 1: Validate EPLB state
         self._validate_eplb_state()
diff --git a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py
index f87b7d3ae8d0..def871f428c7 100644
--- a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py
+++ b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py
@@ -10,8 +10,7 @@
 
 
 class FusedMoERouter(ABC):
-    """
-    FusedMoERouter is an abstract class that provides a 'select_experts'
+    """FusedMoERouter is an abstract class that provides a 'select_experts'
     method that is used for routing hidden states based on router logits.
     """
 
@@ -54,8 +53,7 @@ def select_experts(
         *,
         input_ids: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Route the input hidden states to the top-k experts based on the
+        """Route the input hidden states to the top-k experts based on the
         router logits.
 
         Returns:
@@ -66,8 +64,8 @@ def select_experts(
             **Compatibility**: When EPLB is not enabled, the returned ids are
             equivalent to global logical ids, so should be compatible with
             plain MoE implementations without redundant experts.
-        """
 
+        """
         topk_weights, topk_ids = self._select_experts(
             hidden_states,
             router_logits,
diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py
index 1f086b448802..f91e9c449ef5 100644
--- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py
+++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py
@@ -5,7 +5,11 @@
 
 import vllm._custom_ops as ops
 from vllm.model_executor.custom_op import PluggableLayer
-from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.linear import (
+    ReplicatedLinear,
+    UnquantizedLinearMethod,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import direct_register_custom_op
 
@@ -25,12 +29,18 @@ class GateLinear(ReplicatedLinear):
     The ``out_dtype`` attribute is mutable and can be set after init
     (e.g. when the required dtype depends on the expert quantization
     method which is only known later).
+
+    A ``quant_config`` that actually quantizes the gate disables every
+    specialized tier, leaving plain ``ReplicatedLinear`` behavior.
     """
 
     # (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
     #   (3072, 256) -> MiniMax-M2/M2.5,  (6144, 128) -> MiniMax-M3
     FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)}
     FP32_MAX_TOKENS = 32
+    # Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses
+    # whose router shape stays ahead further up the range may raise it.
+    LL_BF16_MAX_TOKENS = 16
 
     def __init__(
         self,
@@ -40,7 +50,10 @@ def __init__(
         out_dtype: torch.dtype | None = None,
         params_dtype: torch.dtype | None = None,
         force_fp32_compute: bool = False,
+        skip_bias_add: bool = False,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        return_bias: bool = True,
     ):
         is_hopper = current_platform.is_device_capability((9, 0))
         is_blackwell = current_platform.is_device_capability_family(100)
@@ -72,17 +85,25 @@ def __init__(
             input_size,
             output_size,
             bias=bias,
+            skip_bias_add=skip_bias_add,
             params_dtype=params_dtype,
-            quant_config=None,
+            quant_config=quant_config,
             prefix=prefix,
+            return_bias=return_bias,
         )
         self.out_dtype = out_dtype
 
+        # A quantized gate exposes no plain ``weight``, so every specialized
+        # tier below is disabled and forward falls back to ReplicatedLinear.
+        self.is_unquantized = isinstance(self.quant_method, UnquantizedLinearMethod)
+        can_use_specialized_kernels &= self.is_unquantized
+
         self.allow_specialized_router_gemm = can_use_specialized_kernels
 
         # fp32 specialized kernel eligibility (exact dims, fp32 weight)
         self.allow_fp32_router_gemm = (
-            not bias
+            self.is_unquantized
+            and not bias
             and self.weight.dtype == torch.float32
             and (
                 (
@@ -94,7 +115,8 @@ def __init__(
             )
         )
         self.allow_bf16x3_router_gemm = (
-            not bias
+            self.is_unquantized
+            and not bias
             and self.weight.dtype == torch.float32
             and current_platform.is_cuda()
             and is_blackwell
@@ -108,7 +130,7 @@ def __init__(
         # applies on any CUDA-alike device (no bias, since torch.mm has no bias
         # term). The specialized-kernel gate above excludes family-120 Blackwell
         # (GB10 / DGX Spark), which this tier still covers. See #49921.
-        self._router_gemm_no_bias = not bias
+        self._router_gemm_no_bias = self.is_unquantized and not bias
         self._router_gemm_cublas_capable = (
             current_platform.is_cuda() or current_platform.is_rocm()
         ) and self._router_gemm_no_bias
@@ -162,17 +184,26 @@ def set_out_dtype(self, out_dtype: torch.dtype) -> None:
                 and is_available()
             )
 
+    def _return(
+        self, output: torch.Tensor
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        return output if not self.return_bias else (output, None)
+
     def forward(
         self, x: torch.Tensor
     ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
         # Tier 1: cuteDSL ll_bf16_gemm (SM90+, any dims)
-        if self.allow_ll_bf16_gemm and x.shape[0] <= 16 and x.dtype == torch.bfloat16:
+        if (
+            self.allow_ll_bf16_gemm
+            and x.shape[0] <= self.LL_BF16_MAX_TOKENS
+            and x.dtype == torch.bfloat16
+        ):
             from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
                 ll_bf16_gemm,
             )
 
             output = ll_bf16_gemm(x, self.weight)
-            return output, None
+            return self._return(output)
 
         # Tier 2: fp32 specialized kernel (model-specific shapes, M<=32)
         # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph
@@ -184,7 +215,7 @@ def forward(
             output = torch.ops.vllm.fp32_router_gemm_dispatch(
                 x, self.weight, self.allow_bf16x3_router_gemm
             )
-            return output, None
+            return self._return(output)
 
         # Tier 3: bf16x3 CuteDSL kernel for fp32 router weights
         if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
@@ -193,20 +224,29 @@ def forward(
             )
 
             output = bf16x3_router_gemm(x, self.weight)
-            return output, None
+            return self._return(output)
 
         # Tier 4: cuBLAS bf16→fp32
         if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
             output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
-            return output, None
+            return self._return(output)
 
         # Tier 5: F.linear (ReplicatedLinear)
-        if self.out_dtype is not None and x.dtype != self.weight.dtype:
+        if (
+            self.out_dtype is not None
+            and self.is_unquantized
+            and x.dtype != self.weight.dtype
+        ):
             x = x.to(self.weight.dtype)
-        output, output_bias = super().forward(x)
+        base_output = super().forward(x)
+        output_bias: Parameter | None = None
+        if isinstance(base_output, tuple):
+            output, output_bias = base_output
+        else:
+            output = base_output
         if self.out_dtype is not None and output.dtype != self.out_dtype:
             output = output.to(self.out_dtype)
-        return output, output_bias
+        return output if not self.return_bias else (output, output_bias)
 
 
 _FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS
@@ -217,8 +257,7 @@ def fp32_router_gemm_dispatch_impl(
     weight: torch.Tensor,
     allow_bf16x3_router_gemm: bool,
 ) -> torch.Tensor:
-    """
-    Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS,
+    """Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS,
     otherwise optionally run the experimental BF16x3 kernel for medium/large
     SM100 router batches, then fall back to F.linear.
     This must be wrapped in a custom op because our torch.compile integration
diff --git a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
index 2df3489358b1..9b134d385a63 100644
--- a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
+++ b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
@@ -8,6 +8,7 @@
 from vllm import envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.layers.fused_moe.config import (
     RoutingMethodType,
@@ -258,6 +259,7 @@ def __init__(
         e_score_correction_bias: torch.Tensor | None = None,
         num_fused_shared_experts: int = 0,
         eplb_state: EplbLayerState | None = None,
+        skip_padding: bool = False,
     ):
         super().__init__(
             top_k=top_k,
@@ -271,6 +273,7 @@ def __init__(
         self.routed_scaling_factor = routed_scaling_factor
         self.e_score_correction_bias = e_score_correction_bias
         self.num_fused_shared_experts = num_fused_shared_experts
+        self.skip_padding = skip_padding
 
     @property
     def routing_method_type(self) -> RoutingMethodType:
@@ -346,4 +349,15 @@ def valid_grouping() -> bool:
             e_score_correction_bias=self.e_score_correction_bias,
         )
 
+        if (
+            self.skip_padding
+            and envs.VLLM_MOE_SKIP_PADDING
+            and is_forward_context_available()
+        ):
+            is_padding = get_forward_context().is_padding
+            if is_padding is not None:
+                is_padding = is_padding[: topk_ids.shape[0]].unsqueeze(1)
+                topk_weights = topk_weights.masked_fill(is_padding, 0)
+                topk_ids = topk_ids.masked_fill(is_padding, -1)
+
         return topk_weights, topk_ids
diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py
index 0f6ec15d4567..f03fb106f795 100644
--- a/vllm/model_executor/layers/fused_moe/router/router_factory.py
+++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py
@@ -65,9 +65,9 @@ def create_fused_moe_router(
     # Deepseek V4 vision routing bias parameters
     bias_vl: torch.Tensor | None = None,
     image_sentinel_lo: int = 0,
+    skip_padding: bool = False,
 ) -> FusedMoERouter:
-    """
-    Factory function to create the appropriate FusedMoERouter subclass based on
+    """Factory function to create the appropriate FusedMoERouter subclass based on
     the provided parameters.
 
     The selection logic follows this priority order:
@@ -92,6 +92,7 @@ def create_fused_moe_router(
         topk_group: Top-k within each group (for grouped routing)
         scoring_func: Scoring function to use ("softmax" or "sigmoid")
         num_fused_shared_experts: Number of fused shared experts (for ROCm AITER)
+        skip_padding: Whether to invalidate routes for padding rows
 
     Grouped topk and fused topk bias arguments:
         routed_scaling_factor: Scaling factor for routed weights
@@ -114,8 +115,8 @@ def create_fused_moe_router(
 
     Returns:
         An instance of the appropriate FusedMoERouter subclass
-    """
 
+    """
     routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
     if routing_strategy != "":
         return RoutingSimulatorRouter(
@@ -210,6 +211,7 @@ def advertised_routing_method(groups: int | None) -> RoutingMethodType:
                 routed_scaling_factor=routed_scaling_factor,
                 e_score_correction_bias=e_score_correction_bias,
                 num_fused_shared_experts=num_fused_shared_experts,
+                skip_padding=skip_padding,
             )
         # Otherwise fall through to the non-grouped chain below.
 
diff --git a/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py b/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py
index 853e6e67583b..0ff5627d2f04 100644
--- a/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py
+++ b/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py
@@ -25,8 +25,7 @@ def route_tokens(
         top_k: int,
         indices_type: torch.dtype | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Route tokens to experts.
+        """Route tokens to experts.
 
         Args:
             hidden_states: Input hidden states [num_tokens, hidden_size]
@@ -36,13 +35,13 @@ def route_tokens(
 
         Returns:
             tuple of (topk_weights, topk_ids)
+
         """
         pass
 
 
 class DistributionBasedRouting(RoutingStrategy):
-    """
-    Distribution-based random routing strategy with configurable distributions.
+    """Distribution-based random routing strategy with configurable distributions.
 
     This routing strategy randomly selects experts for each token based on
     different probability distributions. Currently supports uniform and normal
@@ -50,8 +49,7 @@ class DistributionBasedRouting(RoutingStrategy):
     """
 
     def __init__(self, distribution: str = "uniform", **distribution_params: Any):
-        """
-        Initialize distribution-based routing.
+        """Initialize distribution-based routing.
 
         Args:
             distribution: Type of distribution to use for sampling
@@ -61,6 +59,7 @@ def __init__(self, distribution: str = "uniform", **distribution_params: Any):
                 chosen distribution
                 For "uniform": No additional parameters needed
                 For "normal": mean (default: 0.0), std (default: 1.0)
+
         """
         self.distribution = distribution.lower()
         self.distribution_params = distribution_params
@@ -90,8 +89,7 @@ def route_tokens(
         top_k: int,
         indices_type: torch.dtype | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Randomly select experts for each token using the specified distribution.
+        """Randomly select experts for each token using the specified distribution.
 
         Args:
             hidden_states: Input hidden states [num_tokens, hidden_size]
@@ -103,6 +101,7 @@ def route_tokens(
             tuple of (topk_weights, topk_ids) where:
             - topk_weights: Weights based on distribution sampling
             - topk_ids: Expert indices sampled from the distribution
+
         """
         num_tokens = hidden_states.shape[0]
         num_experts = router_logits.shape[-1]
@@ -129,7 +128,6 @@ def _sample_expert_ids(
         indices_type: torch.dtype,
     ) -> torch.Tensor:
         """Sample expert IDs based on the specified distribution."""
-
         if self.distribution == "uniform":
             # Generate random scores, and take the top-k to avoid duplicate topk_ids
             scores = torch.rand(num_tokens, num_experts, device=device)
@@ -218,8 +216,7 @@ def get_distribution_info(self) -> dict:
 
 
 class RoutingSimulator:
-    """
-    Token-to-Expert Routing Simulator.
+    """Token-to-Expert Routing Simulator.
 
     This class provides a framework for testing and comparing different
     routing strategies for MoE models. It can simulate routing behavior
@@ -239,22 +236,22 @@ class RoutingSimulator:
 
     @classmethod
     def register_strategy(cls, name: str, strategy: RoutingStrategy):
-        """
-        Register a custom routing strategy.
+        """Register a custom routing strategy.
 
         Args:
             name: Name of the strategy
             strategy: RoutingStrategy instance
+
         """
         cls._routing_strategies[name] = strategy
 
     @classmethod
     def get_available_strategies(cls) -> list[str]:
-        """
-        Get list of available routing strategy names.
+        """Get list of available routing strategy names.
 
         Returns:
             List of available strategy names
+
         """
         return list(cls._routing_strategies.keys())
 
@@ -266,8 +263,7 @@ def simulate_routing(
         top_k: int,
         indices_type: torch.dtype | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Simulate token-to-expert routing using the specified strategy.
+        """Simulate token-to-expert routing using the specified strategy.
 
         Args:
             hidden_states: Input hidden states [num_tokens, hidden_size]
@@ -278,6 +274,7 @@ def simulate_routing(
 
         Returns:
             tuple of (topk_weights, topk_ids)
+
         """
         if strategy_name not in RoutingSimulator._routing_strategies:
             raise ValueError(
diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py
index daaff3b79bf6..08279d1c8938 100644
--- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py
@@ -45,6 +45,7 @@
     SharedExperts,
     SharedExpertsOrder,
 )
+from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import (
     _USE_LAYERNAME,
@@ -225,8 +226,7 @@ def _unpack(
 
 
 class MoERunner(MoERunnerInterface):
-    """
-    Standard MoE runner implementation for executing Mixture of Experts layers.
+    """Standard MoE runner implementation for executing Mixture of Experts layers.
 
     This is the primary concrete implementation of MoE execution logic, providing
     comprehensive support for standard MoE operations. It handles:
@@ -274,7 +274,7 @@ def __init__(
 
         # When both gates are present and FSE is enabled, fuse their
         # weight matrices into [num_experts + num_shared, hidden] so one
-        # F.linear produces combined logits. The topk kernel can then
+        # GEMM produces combined logits. The topk kernel can then
         # apply routing softmax and shared expert activation (sigmoid)
         # in a single launch.
         self._fse_fuse_gate = gate is not None and shared_expert_gate is not None
@@ -283,24 +283,11 @@ def __init__(
         self._shared_experts: SharedExperts | None = None
         if shared_experts is not None:
             can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts
-            # When unquantized, shared expert inputs alias the hidden states,
-            # which can lead to race condition in multi-stream mode. Quantized
-            # routed experts copy the input into a fresh buffer first, breaking
-            # the alias, so overlap is safe. Only observed on ROCm.
-            routed_input_is_quantized = lambda: (
-                self.routed_experts.quant_method.moe_quant_config is not None
-                and self.routed_experts.quant_method.moe_quant_config.quant_dtype
-                is not None
-            )
-            is_multistream_safe = lambda: (
-                not current_platform.is_rocm() or routed_input_is_quantized()
-            )
             self._shared_experts = SharedExperts(
                 shared_experts,
                 moe_config=moe_config,
                 enable_dbo=enable_dbo,
                 mk_can_overlap_shared_experts=can_overlap,
-                is_multistream_safe=is_multistream_safe,
             )
 
         # Needed for string -> MoERunner layer lookup in custom ops.
@@ -470,9 +457,18 @@ def _maybe_reduce_routed_output_before_transform(
         Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm.
         TP partial routed outputs must be summed in latent space before such
         transforms are applied.
+
+        A transform that commutes with the TP sum is exempt: if
+        ``sum_r T(x_r) == T(sum_r x_r)``, applying the transform to the local
+        partial output and letting the existing late all-reduce sum the
+        combined result is equivalent, and costs one collective instead of two.
+        Such a transform opts out by setting ``reduce_commutative = True``.
+        The default is False, so transforms that do not declare themselves
+        keep being reduced early.
         """
         if (
             self.routed_output_transform is not None
+            and not getattr(self.routed_output_transform, "reduce_commutative", False)
             and not self.moe_config.is_sequence_parallel
             and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
             and not fused_output_is_reduced
@@ -661,7 +657,11 @@ def _sequence_parallel_context(self):
         """
         ctx = get_forward_context()
         return (
-            ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
+            ctx.dp_metadata.sp_local_sizes(
+                self.moe_config.sp_size,
+                pcp_size=self.moe_config.pcp_size,
+                use_ep=self.moe_config.use_ep,
+            )
             if ctx.dp_metadata
             else nullcontext()
         )
@@ -708,7 +708,6 @@ def forward(
         1. pytorch cannot handle union types in custom op signatures so
            _moe_forward and _moe_forward_shared must be split.
         """
-
         # Apply transform for routed experts (e.g., latent projection for
         # latent MoE). When the caller pre-applies the routed input transform
         # outside the runner (e.g. to overlap it on a separate stream), it
@@ -904,7 +903,9 @@ def _forward_impl(
         if self.gate is not None:
             if self._fse_fuse_gate:
                 self._maybe_fuse_gate_weights()
-                router_logits = F.linear(hidden_states, self._combined_gate_weight)
+                router_logits = dispatch_unquantized_gemm()(
+                    self, hidden_states, self._combined_gate_weight, None
+                )
             else:
                 router_logits, _ = self.gate(hidden_states)
 
@@ -1023,8 +1024,7 @@ def set_eplb_state(
         logical_to_physical_map: torch.Tensor,
         logical_replica_count: torch.Tensor,
     ) -> None:
-        """
-        Register the EPLB state in this layer.
+        """Register the EPLB state in this layer.
 
         This is used later in forward pass, where we get the expert mapping
         and record the load metrics in `expert_load_view`.
diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py
index d579b181562a..4a1270ecbcbe 100644
--- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py
+++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py
@@ -17,8 +17,7 @@
 
 
 class MoERunnerInterface(PluggableLayer, ABC):
-    """
-    Abstract base class for Mixture of Experts (MoE) runners.
+    """Abstract base class for Mixture of Experts (MoE) runners.
 
     This class defines the interface that all MoE runner implementations must follow.
     MoE runners are responsible for executing the forward pass of MoE layers, handling
diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py
index 5b9587c8530b..3b8b6a124f22 100644
--- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py
@@ -43,7 +43,6 @@ def __init__(
         moe_config: FusedMoEConfig,
         enable_dbo: bool,
         mk_can_overlap_shared_experts: Callable[[], bool],
-        is_multistream_safe: Callable[[], bool],
     ):
         super().__init__()
 
@@ -58,10 +57,6 @@ def __init__(
 
         self._mk_can_overlap_shared_experts = mk_can_overlap_shared_experts
 
-        # Might not be safe to run multi-stream mode if routed and shared experts
-        # alias the same inputs
-        self._is_multistream_safe = is_multistream_safe
-
         # Allow disabling of the separate shared experts stream for
         # debug purposes.
         # TODO: Remove this after more extensive testings with TP/DP
@@ -90,6 +85,7 @@ def _disable_shared_experts_overlap(self) -> bool:
         # Both these comm backends have been shown to be safe for shared expert overlap.
         _EPLB_OVERLAP_SAFE_BACKENDS = (
             "allgather_reducescatter",
+            "deepep_v2",
             "flashinfer_nvlink_one_sided",
         )
 
@@ -112,19 +108,11 @@ def _determine_shared_experts_order(
         if self._mk_can_overlap_shared_experts():
             return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
 
-        # On ROCm, empirically only DP-only deployments benefit from the overlap.
-        overlap_is_beneficial = not current_platform.is_rocm() or (
-            self._moe_config.moe_parallel_config.dp_size > 1
-            and self._moe_config.moe_parallel_config.tp_size == 1
-        )
-
         should_run_shared_in_aux_stream = (
             current_platform.is_cuda_alike()
             and self._stream is not None
             and hidden_states.shape[0]
             <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
-            and overlap_is_beneficial
-            and self._is_multistream_safe()
         )
 
         if should_run_shared_in_aux_stream:
diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py
index f9543358ecee..585991dfe9df 100644
--- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py
+++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py
@@ -9,8 +9,7 @@
 
 
 class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce):
-    """
-    Useful in the case when some FusedMoEExpertsModular
+    """Useful in the case when some FusedMoEExpertsModular
     implementation does not perform weight application and reduction
     but cannot address the needs of all the compatible PrepareAndFinalize
     implementations.
@@ -42,8 +41,7 @@ def apply(
 
 
 class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):
-    """
-    The fused_experts outputs have already been weight applied and reduced.
+    """The fused_experts outputs have already been weight applied and reduced.
     This implementation is a no-op.
     """
 
@@ -78,8 +76,7 @@ def apply(
 
 
 class TopKWeightAndReduceContiguous(mk.TopKWeightAndReduce):
-    """
-    TopKWeightAndReduce implementation for a fused_experts output
+    """TopKWeightAndReduce implementation for a fused_experts output
     of shape (m, topk, K)
     """
 
@@ -122,8 +119,7 @@ def apply(
 
 
 class TopKWeightAndReduceNaiveBatched(mk.TopKWeightAndReduce):
-    """
-    TopKWeightAndReduce implementation for a fused_experts output
+    """TopKWeightAndReduce implementation for a fused_experts output
     of shape (num_experts, batch_size, K)
     """
 
diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
index 523d1df0d80e..17481d4a9186 100644
--- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
@@ -11,6 +11,7 @@
 from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
+    FusedMoEParallelConfig,
     FusedMoEQuantConfig,
     biased_moe_quant_config,
 )
@@ -23,6 +24,7 @@
     convert_to_unquantized_kernel_format,
     make_unquantized_moe_kernel,
     select_unquantized_moe_backend,
+    unquantized_round_up_hidden_size_and_intermediate_size,
 )
 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
     SharedExperts,
@@ -47,13 +49,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
 
     # --8<-- [end:unquantized_fused_moe]
 
-    supports_pre_processed_weights = True
-
     def __init__(self, moe: FusedMoEConfig):
         super().__init__(moe)
         self.unquantized_backend, self.experts_cls = select_unquantized_moe_backend(
             moe_config=self.moe,
         )
+        # MoonEP replaces the named parameters with its [E+B] gate/down
+        # split and keeps the up projection outside the parameter registry,
+        # so a pre-processed (e.g. ipc_cache) load cannot reconstruct it.
+        self.supports_pre_processed_weights = (
+            self.unquantized_backend != UnquantizedMoeBackend.MOONEP
+        )
 
     @property
     def supports_eplb(self) -> bool:
@@ -68,10 +74,15 @@ def create_weights(
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
+        unpadded_intermediate = self.moe.intermediate_size_per_partition_unpadded
+        assert unpadded_intermediate is not None
         if self.moe.is_act_and_mul:
             w13_up_dim = 2 * intermediate_size_per_partition
         else:
             w13_up_dim = intermediate_size_per_partition
+        unpadded_hidden = self.moe.hidden_dim_unpadded
+        assert unpadded_hidden is not None
+        unpadded_up = unpadded_intermediate * (2 if self.moe.is_act_and_mul else 1)
         # Fused gate_up_proj (column parallel)
         w13_weight = torch.nn.Parameter(
             torch.empty(
@@ -84,6 +95,7 @@ def create_weights(
         )
         layer.register_parameter("w13_weight", w13_weight)
         set_weight_attrs(w13_weight, extra_weight_attrs)
+        w13_weight.weight_loader_numel = num_experts * unpadded_up * unpadded_hidden
         if self.moe.has_bias:
             w13_bias = torch.nn.Parameter(
                 torch.zeros(num_experts, w13_up_dim, dtype=params_dtype),
@@ -91,6 +103,7 @@ def create_weights(
             )
             layer.register_parameter("w13_bias", w13_bias)
             set_weight_attrs(w13_bias, extra_weight_attrs)
+            w13_bias.weight_loader_numel = num_experts * unpadded_up
         # down_proj (row parallel)
         w2_weight = torch.nn.Parameter(
             torch.empty(
@@ -103,6 +116,9 @@ def create_weights(
         )
         layer.register_parameter("w2_weight", w2_weight)
         set_weight_attrs(w2_weight, extra_weight_attrs)
+        w2_weight.weight_loader_numel = (
+            num_experts * unpadded_hidden * unpadded_intermediate
+        )
         if self.moe.has_bias:
             w2_bias = torch.nn.Parameter(
                 torch.zeros(num_experts, hidden_size, dtype=params_dtype),
@@ -110,6 +126,24 @@ def create_weights(
             )
             layer.register_parameter("w2_bias", w2_bias)
             set_weight_attrs(w2_bias, extra_weight_attrs)
+            w2_bias.weight_loader_numel = num_experts * unpadded_hidden
+
+    def maybe_roundup_sizes(
+        self,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        act_dtype: torch.dtype,
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> tuple[int, int]:
+        hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
+            hidden_size,
+            intermediate_size_per_partition,
+            act_dtype,
+            moe_parallel_config,
+        )
+        return unquantized_round_up_hidden_size_and_intermediate_size(
+            self.unquantized_backend, hidden_size, intermediate_size_per_partition
+        )
 
     def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:
         # Pad the weight tensor. This is an optimization on ROCm platform, which
@@ -141,6 +175,7 @@ def _setup_kernel(
             moe_config=layer.moe_config,
             w13_weight=w13,
             w2_weight=w2,
+            layer=layer,
         )
         # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__;
         # On the first call we replace the parameter normally. On subsequent
@@ -161,14 +196,13 @@ def _setup_kernel(
             # does not need to be re-built.
             self._init_moe_kernel(layer)
 
-            if self.unquantized_backend == UnquantizedMoeBackend.CPU:
-                # The CPU experts need the layer itself for the setup that
-                # convert_to_unquantized_kernel_format cannot express, since
-                # it only sees the two weight tensors: padding and prepacking
-                # into the grouped-gemm layout (bias included), and capturing
-                # the router config that monolithic apply() cannot carry.
-                assert self.moe_kernel is not None
-                self.moe_kernel.fused_experts.process_weights_after_loading(layer)
+            # No-op by default. Experts that need the layer itself for setup
+            # convert_to_unquantized_kernel_format cannot express override
+            # it: CPU prepacks into its grouped-gemm layout and captures the
+            # router config; MoonEP picks up the [E+B] weight layout for
+            # prefetch and the up projection.
+            assert self.moe_kernel is not None
+            self.moe_kernel.fused_experts.process_weights_after_loading(layer)
 
     def _init_moe_kernel(self, layer: "RoutedExperts") -> None:
         """Build the MoE kernel from the layer's current (shuffled) weights."""
diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py
index 48e217b1a91a..a98e181caeb2 100644
--- a/vllm/model_executor/layers/fused_moe/utils.py
+++ b/vllm/model_executor/layers/fused_moe/utils.py
@@ -69,6 +69,7 @@ def resolve_layer_fused_shared_expert(
 
     Raises:
         ValueError: If requested shared-expert fusion is quantization-incompatible.
+
     """
     # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported
     # that returns None if AITER is not available.
@@ -164,20 +165,19 @@ def _count_expert_num_tokens(
 def count_expert_num_tokens(
     topk_ids: torch.Tensor, num_local_experts: int, expert_map: torch.Tensor | None
 ) -> torch.Tensor:
-    """
-    Count the number to tokens assigned to each expert.
+    """Count the number to tokens assigned to each expert.
 
-    Parameters:
-    - topk_ids (torch.Tensor): Tensor mapping each token to its
-    list of experts.
-    - num_local_experts (int): Number of experts in this rank.
-    - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
-    from the global expert space to the local expert space of the expert
-    parallel shard.
+    Args:
+        topk_ids (torch.Tensor): Tensor mapping each token to its list of experts.
+        num_local_experts (int): Number of experts in this rank.
+        expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
+            from the global expert space to the local expert space of the expert
+            parallel shard.
 
     Returns:
     A tensor of size num_local_experts, where tensor[i] holds the number
     of tokens assigned to the ith expert.
+
     """
     assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
     expert_num_tokens = torch.empty(
@@ -202,8 +202,7 @@ def count_expert_num_tokens(
 
 
 def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
-    """
-    Shrink the given tensor and apply the given view to it.  This is
+    """Shrink the given tensor and apply the given view to it.  This is
     used to resize the intermediate fused_moe caches.
     """
     assert prod(v) <= x.numel(), (
@@ -226,8 +225,7 @@ def _fp8_quantize(
     per_act_token: bool,
     block_shape: list[int] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Perform fp8 quantization on the inputs.  If a block_shape
+    """Perform fp8 quantization on the inputs.  If a block_shape
     is provided, the output will be blocked.
     """
     if block_shape is None:
@@ -252,11 +250,9 @@ def _int8_quantize(
     per_act_token: bool,
     block_shape: list[int] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Perform int8 quantization on the inputs.  If a block_shape
+    """Perform int8 quantization on the inputs.  If a block_shape
     is provided, the output will be blocked.
     """
-
     # If weights are per-channel (per_channel_quant=True), then
     # activations apply per-token quantization. Otherwise, assume
     # activation tensor-wise fp8/int8 quantization, dynamic or static
@@ -457,35 +453,6 @@ def normalize_batched_scales_shape(
     return scales
 
 
-@triton.jit
-def _pack_topk_ids_weights_kernel(
-    topk_ids_ptr,
-    topk_weights_ptr,
-    output_ptr,
-    n_elements,
-    BLOCK_SIZE: tl.constexpr,
-    USE_GDC: tl.constexpr,
-    launch_pdl: tl.constexpr,  # triton metadata
-):
-    pid = tl.program_id(axis=0)
-    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    mask = offsets < n_elements
-    if USE_GDC:
-        tl.extra.cuda.gdc_launch_dependents()
-        tl.extra.cuda.gdc_wait()
-    expert_id = tl.load(topk_ids_ptr + offsets, mask=mask, other=0).to(tl.int32)
-    expert_id_shifted = expert_id << 16
-
-    weight = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0)
-    weight_bf16 = weight.to(tl.bfloat16)
-    weight_int16 = weight_bf16.to(tl.int16, bitcast=True)
-
-    weight_int32 = weight_int16.to(tl.int32) & 0xFFFF
-
-    packed = expert_id_shifted | weight_int32
-    tl.store(output_ptr + offsets, packed, mask=mask)
-
-
 def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int:
     """Estimate FlashInfer's MoE autotuning maximum token count.
 
@@ -500,7 +467,8 @@ def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int:
 
     For a detailed explanation, see: `docs/serving/data_parallel_deployment.md`
     """
-    return max(moe_config.max_num_tokens * moe_config.dp_size, 8192)
+    dp_size = moe_config.elastic_ep_max_dp_size or moe_config.dp_size
+    return max(moe_config.max_num_tokens * dp_size, 8192)
 
 
 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
diff --git a/vllm/model_executor/layers/fused_qk_norm_rope.py b/vllm/model_executor/layers/fused_qk_norm_rope.py
index d3c645488a83..a62174b338c4 100644
--- a/vllm/model_executor/layers/fused_qk_norm_rope.py
+++ b/vllm/model_executor/layers/fused_qk_norm_rope.py
@@ -170,6 +170,7 @@ def fused_qk_rmsnorm_rope_gate(
     Returns:
         (q_out, k_out, gate_out) -- all contiguous (n_tokens, heads * head_dim).
         ``gate_out`` is the raw (pre-sigmoid) gate.
+
     """
     if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0:
         raise ValueError(
diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py
index 9fe253f2d9cd..a9b14c617000 100644
--- a/vllm/model_executor/layers/fusion/fused_act_quant.py
+++ b/vllm/model_executor/layers/fusion/fused_act_quant.py
@@ -1,25 +1,30 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Producer side of the QuantizedActivation contract for activation layers.
+"""Producer side of the QuantizedActivation contract for activation layers.
 
 Given an activation module and the downstream linear it feeds, fuse the
 activation with that linear's input quantization into a single kernel when the
-linear advertises a consumable input_quant_key (see quant_activation.py).
+linear advertises a consumable input quantization key (see quant_activation.py).
 Falls back to the plain activation when nothing matches, so a model forward can
 always call maybe_fused_act_quant unconditionally.
 
 This is the manual-fusion counterpart to ActivationQuantFusionPass: when fusion
-fires here the silu_and_mul pattern is already consumed, so the compiler pass
-finds nothing to rewrite and the two never double-fuse.
+fires here, the activation and quantization are already consumed, so a compiler
+pass cannot fuse the same boundary again.
 """
 
 from collections.abc import Callable
 
 import torch
 
-from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.activation import ReLUSquaredActivation, SiluAndMul
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    get_input_quant_key,
+)
+from vllm.model_executor.layers.fusion.relu2_fp8_quant import (
+    relu_squared_static_fp8_quant,
+)
 from vllm.model_executor.layers.linear import LinearBase
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
@@ -139,12 +144,46 @@ def _silu_and_mul_nvfp4_dynamic(
     )
 
 
-# (activation module type, consumer input_quant_key) -> fused producer.
-# Mirrors ActivationQuantFusionPass.FUSED_OPS; add a row to migrate a scheme.
+# (activation module type, consumer input quantization key) -> fused producer.
+# Add a row for each supported manual producer.
 _FUSED_ACT_QUANT: dict[tuple[type, QuantKey], Callable] = {
     (SiluAndMul, kFp8StaticTensorSym): _silu_and_mul_fp8_static,
 }
 
+
+def _relu_squared_static_fp8_quant_supported(
+    act_fn: torch.nn.Module,
+    x: torch.Tensor,
+    linear: LinearBase,
+) -> bool:
+    """Return whether the ReLU2 static-FP8 producer can consume this input."""
+    scale = getattr(linear, "input_scale", None)
+    return (
+        isinstance(act_fn, ReLUSquaredActivation)
+        and x.is_cuda
+        and x.dtype == torch.bfloat16
+        and x.is_contiguous()
+        and isinstance(scale, torch.Tensor)
+        and scale.dtype == torch.float32
+        and scale.device == x.device
+        and scale.numel() == 1
+    )
+
+
+# Optional per-entry predicates preserve the unconditional fallback contract
+# for producers whose supported inputs are narrower than their registry key.
+_FUSED_ACT_QUANT_SUPPORT: dict[
+    tuple[type, QuantKey], Callable[[torch.nn.Module, torch.Tensor, LinearBase], bool]
+] = {}
+
+# This path is validated on SM90 and newer CUDA devices.
+if current_platform.is_cuda() and current_platform.has_device_capability(90):
+    _RELU2_STATIC_FP8_KEY = (ReLUSquaredActivation, kFp8StaticTensorSym)
+    _FUSED_ACT_QUANT[_RELU2_STATIC_FP8_KEY] = relu_squared_static_fp8_quant
+    _FUSED_ACT_QUANT_SUPPORT[_RELU2_STATIC_FP8_KEY] = (
+        _relu_squared_static_fp8_quant_supported
+    )
+
 # Add CUDA-specific entries for dynamic block quantization
 if current_platform.is_cuda_alike():
     _FUSED_ACT_QUANT[(SiluAndMul, kFp8Dynamic128Sym)] = _silu_and_mul_fp8_dynamic_128
@@ -161,12 +200,14 @@ def maybe_fused_act_quant(
 ) -> "torch.Tensor | QuantizedActivation":
     """Apply act_fn, fusing the downstream linear's input quant when possible.
 
-    Returns a QuantizedActivation when a fused kernel matches
-    (act_fn, linear.input_quant_key), else the plain activated tensor.
+    Returns a QuantizedActivation when a fused kernel matches the activation and
+    the consumer's effective input quantization key, else the plain activation.
     """
-    key = getattr(linear, "input_quant_key", None)
+    key = get_input_quant_key(linear)
     if key is not None:
-        producer = _FUSED_ACT_QUANT.get((type(act_fn), key))
-        if producer is not None:
+        registry_key = (type(act_fn), key)
+        producer = _FUSED_ACT_QUANT.get(registry_key)
+        support = _FUSED_ACT_QUANT_SUPPORT.get(registry_key)
+        if producer is not None and (support is None or support(act_fn, x, linear)):
             return producer(x, linear)
     return act_fn(x)
diff --git a/vllm/model_executor/layers/fusion/quant_activation.py b/vllm/model_executor/layers/fusion/quant_activation.py
index 4be2f4f9ffe1..c7dbe840db1d 100644
--- a/vllm/model_executor/layers/fusion/quant_activation.py
+++ b/vllm/model_executor/layers/fusion/quant_activation.py
@@ -1,14 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-A QuantizedActivation is a pre-quantized activation produced by a fused kernel
+"""A QuantizedActivation is a pre-quantized activation produced by a fused kernel
 and consumed directly by a linear layer, letting the layer skip its own input
 quantization. A linear advertises the key its kernel can consume via
 expose_input_quant_key; the kernel validates and reads the activation via
-as_quantized_activation.
+as_quantized_activation. Producers query the effective capability through
+get_input_quant_key, which hides the key when another consumer branch needs the
+original activation.
 """
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 
 import torch
 
@@ -34,14 +35,23 @@ class QuantizedActivation:
     orig_shape: torch.Size
     quant_key: QuantKey
 
+    def weak_ref(self) -> "QuantizedActivation":
+        """Return a copy with non-owning tensor references for CUDA graph replay."""
+        from vllm.utils.torch_utils import weak_ref_tensor
+
+        return replace(
+            self,
+            data=weak_ref_tensor(self.data),
+            scale=weak_ref_tensor(self.scale),
+        )
+
 
 def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None:
-    """Advertise the kernel's pre-quantized input key on the layer, if any.
+    """Store the kernel's pre-quantized input key on the layer, if any.
 
-    This is the bridge from a kernel's input_quant_key() to the
-    layer.input_quant_key attribute that fusion call sites read. The attribute
-    is left unset when the kernel quantizes its own input, so non-supporting
-    backends never receive a QuantizedActivation.
+    This is the bridge from a kernel's input_quant_key() to the layer capability
+    that fusion call sites read through get_input_quant_key. The raw key is left
+    unset when the kernel quantizes its own input.
 
     TODO(mgoin): Producers also need the consumer's quantization scales (e.g.
     static input scale, global scale). Expose those here as well so producers
@@ -49,7 +59,14 @@ def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None:
     """
     key = kernel.input_quant_key()
     if key is not None:
-        layer.input_quant_key = key
+        layer._input_quant_key = key
+
+
+def get_input_quant_key(layer: torch.nn.Module) -> QuantKey | None:
+    """Return the pre-quantized input key the complete layer can consume."""
+    if getattr(layer, "requires_unquantized_input", False):
+        return None
+    return getattr(layer, "_input_quant_key", None)
 
 
 def as_quantized_activation(
diff --git a/vllm/model_executor/layers/fusion/relu2_fp8_quant.py b/vllm/model_executor/layers/fusion/relu2_fp8_quant.py
new file mode 100644
index 000000000000..4dc0fed4cc41
--- /dev/null
+++ b/vllm/model_executor/layers/fusion/relu2_fp8_quant.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_fp8_min_max,
+    kFp8StaticTensorSym,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+_FP8_MIN, _FP8_MAX = get_fp8_min_max()
+
+
+@triton.jit
+def _relu_squared_static_fp8_quant_kernel(
+    x_ptr,
+    scale_ptr,
+    output_ptr,
+    num_elements,
+    FP8_MIN: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < num_elements
+    x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+    input_is_nan = x != x
+    relu = tl.maximum(x, 0.0)
+    # Preserve the standalone ReLU2 kernel's BF16 output boundary.
+    activated = (relu * relu).to(tl.bfloat16).to(tl.float32)
+    # Match O2's Inductor-compiled QuantFP8 reciprocal-then-multiply.
+    quantized = activated * (1.0 / tl.load(scale_ptr))
+    quantized = tl.maximum(tl.minimum(quantized, FP8_MAX), FP8_MIN)
+    # Match the O2 native chain by preserving NaN through the FP8 conversion.
+    quantized = tl.where(input_is_nan, x, quantized)
+    tl.store(output_ptr + offsets, quantized, mask=mask)
+
+
+def relu_squared_static_fp8_quant(
+    x: torch.Tensor, linear: LinearBase
+) -> QuantizedActivation:
+    """BF16 ReLU2 followed by static per-tensor FP8 quantization."""
+    assert x.dtype == torch.bfloat16
+    assert x.is_contiguous()
+    scale = linear.input_scale
+    assert scale.dtype == torch.float32
+    assert x.device == scale.device
+    assert scale.numel() == 1
+
+    output = torch.empty_like(x, dtype=current_platform.fp8_dtype())
+    if x.numel() != 0:
+        block_size = min(triton.next_power_of_2(x.shape[-1]), 2048)
+        num_warps = min(max(block_size // 256, 1), 4)
+        grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)
+        _relu_squared_static_fp8_quant_kernel[grid](
+            x,
+            scale,
+            output,
+            x.numel(),
+            FP8_MIN=_FP8_MIN,
+            FP8_MAX=_FP8_MAX,
+            BLOCK_SIZE=block_size,
+            num_warps=num_warps,
+        )
+
+    return QuantizedActivation(
+        data=output,
+        scale=scale,
+        orig_dtype=x.dtype,
+        orig_shape=x.shape,
+        quant_key=kFp8StaticTensorSym,
+    )
diff --git a/vllm/model_executor/layers/hpc/gated_mla.py b/vllm/model_executor/layers/hpc/gated_mla.py
index f1e495a3a7f2..4edd42f164f5 100644
--- a/vllm/model_executor/layers/hpc/gated_mla.py
+++ b/vllm/model_executor/layers/hpc/gated_mla.py
@@ -107,6 +107,7 @@ def hpc_gated_mla_gemm(
 
     Returns:
         ``[m, n]`` bfloat16.
+
     """
     from hpc.gemm import gated_mla_gemm
 
diff --git a/vllm/model_executor/layers/hpc/hpc_ihc.py b/vllm/model_executor/layers/hpc/hpc_ihc.py
index 1a8681aa8ead..aaf6ea92f20b 100644
--- a/vllm/model_executor/layers/hpc/hpc_ihc.py
+++ b/vllm/model_executor/layers/hpc/hpc_ihc.py
@@ -20,9 +20,8 @@
   - x / residual must be bfloat16; hc_fn weights stay float32 (the checkpoint
     keeps them out of fp8 quantization via modules_to_not_convert)
 
-The fused modules own no parameters: they hold references to the weights of the
-eager layer they replace, bound in ``process_weights_after_loading`` (invoked
-generically for every ``HpcModule`` by the model loader).
+The fused modules own no parameters: their forwards read the weights of the
+eager layer they replace directly via the captured owner reference.
 
 NOTE: The reference implementation also offers a cross-layer post+pre fusion
 (``HpcIHCPostPre``) that folds one segment's post into the next segment's pre.
@@ -123,6 +122,7 @@ class HpcIHCPre(HpcModule):
             When given, its weight/eps are folded into the kernel so a single
             launch covers pre + RMSNorm (the caller then skips the separate
             layernorm). None keeps the original pre-only behaviour.
+
     """
 
     def __init__(
@@ -147,41 +147,26 @@ def __init__(
         object.__setattr__(self, "_fallback_op", fallback_op)
         # Same anti-cycle idiom for the fused RMSNorm owner (a sibling module).
         object.__setattr__(self, "_norm_owner", norm_owner)
-        self.weight = None
-        self.hc_scale = None
-        self.hc_base = None
-        self.rms_weight = None
-        self.rms_eps = 0.0
 
     @classmethod
     def support(cls, hc_mult: int, hidden_size: int) -> bool:
         return _ihc_supported(hc_mult, hidden_size)
 
-    def process_weights_after_loading(self, model=None):
-        # hc_fn is a ReplicatedLinear whose weight is [2 * hc_mult, hc_mult * d].
-        # It stays float32 -- the HPC kernel consumes fp32 weights directly, so
-        # there is no conversion (and no precision loss) here.
-        owner = self._fallback_op
-        self.weight = owner.hc_fn.weight
-        self.hc_scale = owner.hc_scale
-        self.hc_base = owner.hc_base
-        if self._norm_owner is not None:
-            self.rms_weight = self._norm_owner.weight
-            self.rms_eps = self._norm_owner.variance_epsilon
-
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         import hpc
 
+        owner = self._fallback_op
+        norm_owner = self._norm_owner
         return hpc.fuse_ihc_pre(
             x,
-            self.weight,
-            self.hc_scale,
-            self.hc_base,
+            owner.hc_fn.weight,
+            owner.hc_scale,
+            owner.hc_base,
             self.norm_eps,
             self.hc_eps,
             self.magnitude,
-            self.rms_weight,
-            self.rms_eps,
+            norm_owner.weight if norm_owner is not None else None,
+            norm_owner.variance_epsilon if norm_owner is not None else 0.0,
         )
 
 
@@ -233,29 +218,20 @@ def __init__(
         self.norm_eps = norm_eps
         # See HpcIHCPre: keep the owner out of the module tree.
         object.__setattr__(self, "_fallback_op", fallback_op)
-        self.weight = None
-        self.hc_scale = None
-        self.hc_base = None
 
     @classmethod
     def support(cls, hc_mult: int, hidden_size: int) -> bool:
         return _ihc_supported(hc_mult, hidden_size)
 
-    def process_weights_after_loading(self, model=None):
-        # hc_head_fn weight is [hc_mult, hc_mult * d], float32.
-        owner = self._fallback_op
-        self.weight = owner.hc_head_fn.weight
-        self.hc_scale = owner.hc_head_scale
-        self.hc_base = owner.hc_head_base
-
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         import hpc
 
+        owner = self._fallback_op
         return hpc.fuse_ihc_head(
             x,
-            self.weight,
-            self.hc_scale,
-            self.hc_base,
+            owner.hc_head_fn.weight,
+            owner.hc_head_scale,
+            owner.hc_head_base,
             self.norm_eps,
             self.hc_eps,
         )
diff --git a/vllm/model_executor/layers/hpc/hpc_module.py b/vllm/model_executor/layers/hpc/hpc_module.py
index e1efad832eb9..c3a0dab676a1 100644
--- a/vllm/model_executor/layers/hpc/hpc_module.py
+++ b/vllm/model_executor/layers/hpc/hpc_module.py
@@ -11,8 +11,5 @@ def __init__(self):
     def support(cls, *args, **kwargs):
         return True
 
-    def process_weights_after_loading(self, model):
-        pass
-
     def forward(self, *args, **kwargs):
         pass
diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py
index c1503c754c83..02f06ed89cbd 100644
--- a/vllm/model_executor/layers/hpc/rope_norm.py
+++ b/vllm/model_executor/layers/hpc/rope_norm.py
@@ -90,13 +90,13 @@ class HpcRopeNorm(CustomOp, HpcModule):
     """HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).
 
     Registered as a sub-module in model layers (e.g. HunYuanAttention).
-    Norm weights are extracted from fallback norm modules via
-    process_weights_after_loading() after all weights are loaded.
+    The QK-Norm weights are read directly from the fallback norm modules
+    (built in float32 when HPC is active), eliminating the derived copy.
 
     forward() is dispatched by CustomOp framework:
     - In compiled mode: forward_cuda() calls torch.ops.vllm.hpc_rope_norm_forward
       as a splitting point — internal Python control flow is opaque
-      to torch.compile and not captured by CUDA Graph.
+      to torch.compile (its kernel launches can still be CUDA-graph captured).
     - In eager/native mode: forward_native() falls back to forward_cuda().
     """
 
@@ -124,8 +124,6 @@ def __init__(
         self.num_kv_heads = num_kv_heads
         self.head_dim = head_dim
 
-        self.use_qk_norm = use_qk_norm
-
         self.q_size = num_heads * head_dim
         self.kv_size = num_kv_heads * head_dim
 
@@ -139,31 +137,18 @@ def __init__(
 
         self.head_per_group = num_heads // num_kv_heads
 
-        # Pre-allocate norm weight tensors as Parameters so they are tracked by
-        # CuMemAllocator (for sleep/wake_up) and have stable addresses for CUDA
-        # Graph replay. process_weights_after_loading() updates them inplace via
-        # copy_() so refit does not invalidate captured graph tensor pointers.
-        # Shape is [head_dim] to match the HPC kernel's q/k_norm_weight layout.
-        if use_qk_norm and fallback_qnorm is not None:
-            self.qnorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(
-                torch.empty(head_dim, dtype=torch.float32),
-                requires_grad=False,
-            )
-        else:
-            self.qnorm_weight = None
-        if use_qk_norm and fallback_knorm is not None:
-            self.knorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(
-                torch.empty(head_dim, dtype=torch.float32),
-                requires_grad=False,
-            )
-        else:
-            self.knorm_weight = None
-
         self.use_fp8 = "fp8" in kv_cache_dtype
         # The RMSNorm/RoPE ordering is model dependent (e.g. HunYuan V3 applies
         # QK-Norm before RoPE -> NORM_THEN_ROPE), so it is supplied by the
         # caller. When QK-Norm is disabled the policy is forced to NONE.
         self.qk_norm_policy = qk_norm_policy if use_qk_norm else QkNormPolicy.NONE
+        if self.qk_norm_policy != QkNormPolicy.NONE and (
+            fallback_qnorm is None or fallback_knorm is None
+        ):
+            raise ValueError(
+                "QK-Norm is enabled but fallback_qnorm/fallback_knorm were not "
+                "supplied; the fused kernel reads their weights directly."
+            )
 
         # Register layer_name + add self to the global instance registry so the
         # module-level custom op (hpc_rope_norm_forward) can route back here.
@@ -215,20 +200,6 @@ def support(
         logger.info_once("enable hpc rope_norm")
         return True
 
-    def process_weights_after_loading(self, model: torch.nn.Module = None) -> None:
-        """Copy norm weights (float32) from fallback norm modules inplace.
-
-        Uses copy_() to preserve tensor addresses for CUDA Graph / refit
-        compatibility. Called by the model's load_weights() after all weights
-        are loaded (and generically from the model loader for DummyModelLoader
-        / sleep-wake_up reload paths).
-        """
-        if self.use_qk_norm:
-            if self.fallback_qnorm is not None and self.qnorm_weight is not None:
-                self.qnorm_weight.data.copy_(self.fallback_qnorm.weight.data.float())
-            if self.fallback_knorm is not None and self.knorm_weight is not None:
-                self.knorm_weight.data.copy_(self.fallback_knorm.weight.data.float())
-
     def register_layer_name(self, layer_name: str) -> None:
         """Register layer_name and add self to the global registry.
 
@@ -430,12 +401,14 @@ def _forward_impl(
         k_scale = attn_layer._k_scale.reshape(1)
         v_scale = attn_layer._v_scale.reshape(1)
 
-        q_norm_weight = (
-            self.qnorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
-        )
-        k_norm_weight = (
-            self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None
-        )
+        if self.qk_norm_policy != QkNormPolicy.NONE:
+            assert self.fallback_qnorm is not None
+            assert self.fallback_knorm is not None
+            q_norm_weight = self.fallback_qnorm.weight
+            k_norm_weight = self.fallback_knorm.weight
+        else:
+            q_norm_weight = None
+            k_norm_weight = None
 
         # --- Prefill ---
         if num_prefill_reqs > 0:
diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py
new file mode 100644
index 000000000000..c3422445b2c5
--- /dev/null
+++ b/vllm/model_executor/layers/indexer_topk.py
@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Top-k kernels for the DSA sparse attention indexer."""
+
+import functools
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.config import get_current_vllm_config
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer
+from vllm.v1.worker.workspace import current_workspace_manager
+
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
+
+# cooperative_topk's hard row limit: one cluster wave covers at most this many
+# rows. "auto" uses cooperative_topk for every batch within the limit and falls
+# back to persistent_topk past it; explicit ``cooperative`` requests are
+# validated against the same bound.
+AUTO_COOPERATIVE_MAX_ROWS = 64
+
+# ---------------------------------------------------------------------------
+# DeepSelect (vllm._deepselect_C)
+# ---------------------------------------------------------------------------
+
+# Matches the -1 fill convention used for topk_indices_buffer elsewhere.
+IDX_OOB_FILL_VALUE = -1
+
+try:
+    import vllm._deepselect_C  # noqa: F401  (registers torch.ops.deep_select)
+except ImportError as e:
+    from vllm.logger import init_logger
+
+    init_logger(__name__).warning(
+        "Failed to import the DeepSelect extension (vllm._deepselect_C): %s", e
+    )
+
+
+@functools.lru_cache(maxsize=1)
+def get_deep_select_stride_requirement() -> tuple[int, int]:
+    """Stride alignment requirement (input, output) in bytes."""
+    return torch.ops.deep_select.get_alignment_requirement()
+
+
+def is_deep_select_supported(input: torch.Tensor, topk: int) -> bool:
+    """Whether the kernel accepts this input (dtype/stride/topk constraints)."""
+    return (
+        topk <= 4096
+        and input.shape[1] < 2**23
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and input.stride(1) == 1
+        and input.stride(0)
+        * input.element_size()
+        % get_deep_select_stride_requirement()[0]
+        == 0
+    )
+
+
+def _get_empty_and_aligned_tensor(
+    dim0: int, dim1: int, device: torch.device, dtype: torch.dtype
+) -> torch.Tensor:
+    """Tensor with shape (dim0, dim1) whose stride(0) is 32B-aligned."""
+    output_stride_requirement = get_deep_select_stride_requirement()[1] // (
+        dtype.itemsize
+    )
+    assert output_stride_requirement > 0
+    dim1_rounded = (
+        (dim1 + output_stride_requirement - 1) // output_stride_requirement
+    ) * output_stride_requirement
+    return torch.empty((dim0, dim1_rounded), device=device, dtype=dtype)[:, :dim1]
+
+
+def deep_select_topk(
+    input: torch.Tensor,
+    topk: int,
+    end: torch.Tensor | None = None,
+    output_idx: torch.Tensor | None = None,
+    indices_dtype: torch.dtype = torch.int32,
+) -> torch.Tensor:
+    """Select the top-k indices per row of `input` with DeepSelect.
+
+    Args:
+        input: (num_rows, vocab_size), bf16 or fp32. stride(1) must be 1 and
+            stride(0) must be 1024B-aligned.
+        topk: Number of elements to select per row; must be <= 4096.
+        end: Optional (num_rows,) int32 tensor with the exclusive right
+            boundary of each row. Rows with `end[i] < topk` get their
+            remaining indices filled with -1.
+        output_idx: Optional preallocated (num_rows, topk) output tensor whose
+            stride(0) is 32B-aligned (e.g. a slice of a wider buffer).
+        indices_dtype: Output dtype when `output_idx` is not provided.
+
+    Returns:
+        The (num_rows, topk) indices tensor.
+
+    """
+    assert input.dim() == 2 and input.stride(1) == 1
+    assert (
+        input.stride(0) * input.element_size() % get_deep_select_stride_requirement()[0]
+        == 0
+    )
+
+    num_rows = input.shape[0]
+    if output_idx is None:
+        output_idx = _get_empty_and_aligned_tensor(
+            num_rows, topk, input.device, indices_dtype
+        )
+    else:
+        assert output_idx.dtype == indices_dtype
+        assert output_idx.shape[0] >= num_rows and output_idx.shape[1] >= topk
+
+    torch.ops.deep_select.topk(
+        input,
+        topk,
+        None,  # begin is not supported
+        end,
+        False,  # sorted_value
+        False,  # sorted_index
+        None,  # output_value
+        output_idx,
+        None,  # output_idx_offset
+        IDX_OOB_FILL_VALUE,
+        float("-inf"),  # value_oob_fill_value
+        False,  # return_value
+        False,  # abort_when_nan_found: dummy/capture passes can feed NaN
+        # logits from uninitialized KV cache; the reference topk kernels do
+        # not trap on NaN, and trapping would abort CUDA graph capture.
+    )
+    return output_idx
+
+
+# ---------------------------------------------------------------------------
+# Backend selection
+# ---------------------------------------------------------------------------
+
+
+@functools.cache
+def get_indexer_topk(backend: str) -> "SparseIndexerTopk":
+    return SparseIndexerTopk(backend)
+
+
+class SparseIndexerTopk(torch.nn.Module):
+    """The sparse indexer's decode top-k stage.
+
+    Selects among the available top-k kernels (see
+    kernel_config.sparse_indexer_topk_backend) and runs the chosen one.
+    """
+
+    def __init__(self, backend: str | None = None) -> None:
+        super().__init__()
+        if backend is None:
+            backend = (
+                get_current_vllm_config().kernel_config.sparse_indexer_topk_backend
+            )
+        self._backend = backend
+        self._is_cuda = current_platform.is_cuda()
+        self._has_deep_select = self._is_cuda and (
+            current_platform.is_device_capability_family(100)
+        )
+        self._has_flashinfer_topk = has_flashinfer()
+        self._cooperative_capable = self._is_cuda and (
+            current_platform.has_device_capability(90)
+            and not current_platform.is_device_capability_family(120)
+        )
+
+    def resolve_backend(
+        self, logits: torch.Tensor, topk_tokens: int, num_rows: int
+    ) -> str:
+        """Resolve the decode top-k implementation from the configured
+        backend ("auto" = the pre-existing chain, or a validated explicit
+        value)."""
+        if self._backend == "auto":
+            return self._resolve_auto(logits, topk_tokens, num_rows)
+
+        failures: list[str] = []
+        if self._backend == "cooperative":
+            failures = self._cooperative_constraints(logits, topk_tokens, num_rows)
+        elif self._backend == "persistent":
+            if not self._is_cuda:
+                failures.append("requires a CUDA platform")
+            if topk_tokens not in (512, 1024, 2048):
+                failures.append(
+                    f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}"
+                )
+        elif self._backend == "deep_select":
+            if not self._is_cuda:
+                failures.append("requires a CUDA platform")
+            elif not self._has_deep_select:
+                failures.append("requires SM100a/SM103a (10.x device family)")
+            elif not is_deep_select_supported(logits, topk_tokens):
+                failures.append(
+                    f"inputs violate DeepSelect's constraints: dtype={logits.dtype},"
+                    f" stride={logits.stride()}, topk={topk_tokens}"
+                )
+        elif self._backend == "flashinfer":
+            if not self._is_cuda:
+                failures.append("requires a CUDA platform")
+            if not self._has_flashinfer_topk:
+                failures.append(
+                    "flashinfer.topk.top_k_ragged_transform is not importable"
+                )
+            if logits.dtype != torch.float32 or logits.stride(1) != 1:
+                failures.append(
+                    f"requires fp32 logits with stride(1) == 1, got "
+                    f"dtype={logits.dtype}, stride={logits.stride()}"
+                )
+        if failures:
+            raise RuntimeError(
+                f"sparse_indexer_topk_backend='{self._backend}' was requested, but: "
+                + "; ".join(failures)
+            )
+        return self._backend
+
+    def _resolve_auto(
+        self, logits: torch.Tensor, topk_tokens: int, num_rows: int
+    ) -> str:
+        """The priority chain: cooperative -> persistent -> per_row.
+        deep_select/flashinfer/torch are opt-in only.
+
+        cooperative_topk is preferred whenever it is applicable, i.e. within
+        its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to
+        persistent_topk.
+        """
+        if not self._cooperative_constraints(logits, topk_tokens, num_rows):
+            return "cooperative"
+        if self._is_cuda and topk_tokens in (512, 1024, 2048):
+            return "persistent"
+        return "per_row"
+
+    def _cooperative_constraints(
+        self, logits: torch.Tensor, topk_tokens: int, num_rows: int
+    ) -> list[str]:
+        """Unmet constraints of cooperative_topk (empty when applicable)."""
+        failures = []
+        if not self._is_cuda:
+            failures.append("requires a CUDA platform")
+        if topk_tokens not in (512, 1024, 2048):
+            failures.append(
+                f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}"
+            )
+        if num_rows > AUTO_COOPERATIVE_MAX_ROWS:
+            failures.append(
+                f"num_rows must be <= {AUTO_COOPERATIVE_MAX_ROWS}, got {num_rows}"
+            )
+        if logits.stride(0) % 4 != 0:
+            failures.append(
+                f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}"
+            )
+        if self._is_cuda and not self._cooperative_capable:
+            failures.append("requires SM90+ and is not supported on the SM12x family")
+        return failures
+
+    @staticmethod
+    def _row_ends(seq_lens: torch.Tensor, next_n: int, num_rows: int) -> torch.Tensor:
+        """Per-row exclusive end offsets (int32, (num_rows,)) for top-k
+        kernels that take ragged lengths (DeepSelect, FlashInfer, torch
+        reference).
+
+        seq_lens is (B, next_n) per-row effective lens for native spec decode
+        and (B, 1) otherwise, in which case per-row lens are derived the same
+        way as the other decode top-k kernels.
+        """
+        if seq_lens.numel() == num_rows:
+            row_ends = seq_lens.reshape(-1)
+        else:
+            next_n_offsets = torch.arange(
+                next_n, dtype=torch.int32, device=seq_lens.device
+            )
+            row_ends = (
+                (seq_lens.reshape(-1, 1) - next_n + 1 + next_n_offsets)
+                .clamp_(min=0)
+                .reshape(-1)
+            )
+        assert row_ends.dtype == torch.int32
+        return row_ends
+
+    def forward(
+        self,
+        logits: torch.Tensor,
+        seq_lens: torch.Tensor,
+        next_n: int,
+        topk_indices: torch.Tensor,
+        topk_tokens: int,
+        max_seq_len: int,
+    ) -> None:
+        """Run the resolved decode top-k implementation, writing into
+        topk_indices (int32, -1 fill for rows shorter than topk_tokens)."""
+        backend = self.resolve_backend(logits, topk_tokens, logits.shape[0])
+        if backend == "deep_select":
+            row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
+            deep_select_topk(logits, topk_tokens, end=row_ends, output_idx=topk_indices)
+        elif backend == "cooperative":
+            (topk_workspace,) = current_workspace_manager().get_simultaneous(
+                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+            )
+            torch.ops._C.cooperative_topk(
+                logits,
+                seq_lens,
+                topk_indices,
+                topk_workspace,
+                topk_tokens,
+                max_seq_len,
+            )
+        elif backend == "persistent":
+            (topk_workspace,) = current_workspace_manager().get_simultaneous(
+                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+            )
+            # persistent_topk's last argument is the host-side "max seq_len
+            # across rows" bound (it gates the sampled_topk path and clamps
+            # per-row lengths), not the padded row width: logits is sized to
+            # max_model_len, so passing its width would incorrectly enable the
+            # long-row kernels on short-context batches.
+            torch.ops._C.persistent_topk(
+                logits,
+                seq_lens,
+                topk_indices,
+                topk_workspace,
+                topk_tokens,
+                max_seq_len,
+            )
+        elif backend == "flashinfer":
+            # Deferred: importing flashinfer initializes CUDA at import time.
+            from flashinfer.topk import top_k_ragged_transform
+
+            row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
+            offsets = torch.zeros(
+                logits.shape[0], dtype=torch.int32, device=logits.device
+            )
+            # top_k_ragged_transform selects within [0, row_ends[i]) per
+            # row and -1-fills past the row length.
+            indices = top_k_ragged_transform(logits, offsets, row_ends, topk_tokens)
+            topk_indices.copy_(indices)
+        elif backend == "torch":
+            # Debug reference: mask everything past each row's end, then topk.
+            row_ends = self._row_ends(seq_lens, next_n, logits.shape[0])
+            cols = torch.arange(logits.shape[1], device=logits.device)
+            masked = logits.masked_fill(
+                cols.unsqueeze(0) >= row_ends.unsqueeze(1), float("-inf")
+            )
+            indices = masked.topk(topk_tokens, dim=-1).indices
+            in_range = torch.arange(topk_tokens, device=logits.device).unsqueeze(
+                0
+            ) < row_ends.unsqueeze(1)
+            indices = torch.where(in_range, indices, -1)
+            topk_indices.copy_(indices)
+        else:
+            assert backend == "per_row", f"unknown topk backend: {backend}"
+            ops.top_k_per_row_decode(
+                logits,
+                next_n,
+                seq_lens,
+                topk_indices,
+                logits.shape[0],
+                logits.stride(0),
+                logits.stride(1),
+                topk_tokens,
+            )
diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py
index cded9daf13f8..551cd05d3be3 100644
--- a/vllm/model_executor/layers/layernorm.py
+++ b/vllm/model_executor/layers/layernorm.py
@@ -119,6 +119,14 @@ def forward_xpu(
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if envs.VLLM_BATCH_INVARIANT and residual is not None:
+            assert self.variance_size_override is None, (
+                "Batch invariance is not supported for variance_size_override"
+            )
+            weight = self.weight.data if self.pass_weight_add else None
+            return ir.ops.fused_add_rms_norm.impls["native"].impl_fn(
+                x, residual, weight, self.variance_epsilon
+            )
         return self.forward_cuda(x, residual)
 
     def extra_repr(self) -> str:
@@ -171,6 +179,15 @@ def forward_xpu(
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if envs.VLLM_BATCH_INVARIANT:
+            if residual is None:
+                return self.forward_cuda(x, residual)
+            return ir.ops.fused_add_rms_norm.impls["native"].impl_fn(
+                x,
+                residual,
+                self.weight.float() + 1.0,
+                self.variance_epsilon,
+            )
         import vllm._xpu_ops  # noqa: F401 registers torch.ops.vllm.xpu_gemma_rms_norm
 
         # Fall back to the native path if the fused gemma kernels are not
@@ -231,6 +248,7 @@ def __init__(
             device: Device to create parameters on
             dtype: Data type for parameters
             activation: Activation function name for gating
+
         """
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -335,9 +353,7 @@ def forward_xpu(
 
 
 class LayerNorm(nn.Module):
-    """
-    Layer Normalization.
-    """
+    """Layer Normalization."""
 
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
diff --git a/vllm/model_executor/layers/lightning_attn.py b/vllm/model_executor/layers/lightning_attn.py
index d3ea7fb211a0..1df196ca4d05 100644
--- a/vllm/model_executor/layers/lightning_attn.py
+++ b/vllm/model_executor/layers/lightning_attn.py
@@ -541,8 +541,7 @@ def lightning_attention(
     block_size: int = 256,
     kv_history: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Apply lightning attention algorithm
+    """Apply lightning attention algorithm
     to compute attention efficiently.
 
     Args:
@@ -556,6 +555,7 @@ def lightning_attention(
     Returns:
         output: Attention output
         kv: Updated key-value history
+
     """
     d = q.shape[-1]
     e = v.shape[-1]
@@ -610,8 +610,7 @@ def _linear_attn_decode_kernel(
     pad_slot_id: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
-    """
-    Kernel for linear attention decoding with KV cache.
+    """Kernel for linear attention decoding with KV cache.
 
     This kernel computes attention for a single token using the KV cache.
     """
@@ -683,8 +682,7 @@ def linear_decode_forward_triton(
     slot_idx: torch.Tensor,
     BLOCK_SIZE: int = 32,
 ) -> torch.Tensor:
-    """
-    Perform linear attention decoding using Triton kernels.
+    """Perform linear attention decoding using Triton kernels.
 
     Args:
         q: Query tensor of shape [B, H, 1, D]
@@ -697,6 +695,7 @@ def linear_decode_forward_triton(
 
     Returns:
         output: Attention output tensor
+
     """
     B, H, _, D = q.shape
     assert k.shape == (B, H, 1, D)
diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
index cd1c0e69d2ae..65cda325a221 100644
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -146,6 +146,7 @@ def create_weights(
             input_size: Size of the input dim of the weight across all ranks.
             output_size: Size of the output dim of the weight across all ranks.
             params_dtype: Datatype of the parameters.
+
         """
         raise NotImplementedError
 
@@ -250,8 +251,11 @@ class LinearBase(PluggableLayer):
         prefix: Prefix for parameter names.
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: If true, tensor parallelism will be disabled for this layer.
+
     """
 
+    requires_unquantized_input = False
+
     def __init__(
         self,
         input_size: int,
@@ -333,6 +337,7 @@ class ReplicatedLinear(LinearBase):
                         (e.g. model.layers.0.qkv_proj)
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: Take no effect for replicated linear layers.
+
     """
 
     # --8<-- [end:replicated_linear]
@@ -380,7 +385,8 @@ def __init__(
 
         if bias:
             self.bias = Parameter(
-                torch.empty(self.output_size, dtype=self.params_dtype)
+                torch.empty(self.output_size, dtype=self.params_dtype),
+                requires_grad=False,
             )
             set_weight_attrs(
                 self.bias,
@@ -451,6 +457,7 @@ class ColumnParallelLinear(LinearBase):
             shard per rank (see ``DCPGroupColumnParallelLinear``).
         tp_size: Override the tensor-parallel world size used for sharding.
             Defaults to the global TP world size.
+
     """
 
     # --8<-- [end:column_parallel_linear]
@@ -525,7 +532,8 @@ def __init__(
 
         if bias:
             self.bias = Parameter(
-                torch.empty(self.output_size_per_partition, dtype=params_dtype)
+                torch.empty(self.output_size_per_partition, dtype=params_dtype),
+                requires_grad=False,
             )
             set_weight_attrs(
                 self.bias,
@@ -686,6 +694,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: If true, all weights matrix won't be sharded, this layer
                     will be treated as a "Replicated" MergedLinear.
+
     """
 
     def __init__(
@@ -863,8 +872,7 @@ def _load_fused_module_from_checkpoint(
         loaded_weight: torch.Tensor,
         output_sizes: list[int] | None = None,
     ):
-        """
-        Handle special case for models where MLP layers are already
+        """Handle special case for models where MLP layers are already
         fused on disk. In this case, we have no shard id. This function
         determines the shard id by splitting these layers and then calls
         the weight loader using the shard id.
@@ -872,7 +880,6 @@ def _load_fused_module_from_checkpoint(
         An example of a model with these fused layers:
         https://huggingface.co/microsoft/Phi-3-mini-4k-instruct
         """
-
         current_shard_offset = 0
         shard_offsets: list[tuple[int, int, int]] = []
         output_sizes = output_sizes or self.output_sizes
@@ -1014,6 +1021,7 @@ class QKVParallelLinear(ColumnParallelLinear):
                         (e.g. model.layers.0.qkv_proj)
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: If true, weights matrix won't be sharded through tp rank.
+
     """
 
     def __init__(
@@ -1098,8 +1106,7 @@ def _get_shard_size_mapping(self, loaded_shard_id: str):
     def _load_fused_module_from_checkpoint(
         self, param: BasevLLMParameter, loaded_weight: torch.Tensor
     ):
-        """
-        Handle special case for models where QKV layers are already
+        """Handle special case for models where QKV layers are already
         fused on disk. In this case, we have no shard id. This function
         determines the shard id by splitting these layers and then calls
         the weight loader using the shard id.
@@ -1628,6 +1635,7 @@ class RowParallelLinear(LinearBase):
               | .   |
               | A_p |
                -   -
+
     Arguments:
         input_size: first dimension of matrix A.
         output_size: second dimension of matrix A.
@@ -1647,6 +1655,7 @@ class RowParallelLinear(LinearBase):
                         (e.g. model.layers.0.down_proj)
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: If true, weights matrix won't be sharded through tp rank.
+
     """
 
     # --8<-- [end:row_parallel_linear]
@@ -1708,7 +1717,10 @@ def __init__(
             )
 
         if bias:
-            self.bias = Parameter(torch.empty(self.output_size, dtype=params_dtype))
+            self.bias = Parameter(
+                torch.empty(self.output_size, dtype=params_dtype),
+                requires_grad=False,
+            )
             set_weight_attrs(
                 self.bias,
                 {
diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py
index 694bb6b03960..10f425a50352 100644
--- a/vllm/model_executor/layers/logits_processor.py
+++ b/vllm/model_executor/layers/logits_processor.py
@@ -74,9 +74,9 @@ def __init__(
         logits_as_input: bool = False,
         soft_cap: float | None = None,
     ) -> None:
-        """
-        Args:
-            scale: A scaling factor to apply to the logits.
+        """Args:
+        scale: A scaling factor to apply to the logits.
+
         """
         super().__init__()
         self.scale = scale
diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py
index 10ed7d05fc8f..fa0eb257586d 100644
--- a/vllm/model_executor/layers/mamba/abstract.py
+++ b/vllm/model_executor/layers/mamba/abstract.py
@@ -16,8 +16,7 @@
 
 
 class MambaBase(AttentionLayerBase):
-    """
-    Base class for Mamba-like layers which support the v1 engine.
+    """Base class for Mamba-like layers which support the v1 engine.
     Inherit from this class if you implement a custom layer.
     """
 
@@ -44,8 +43,7 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
 
     @abstractmethod
     def get_state_shape(self) -> Iterable[tuple[int, ...]]:
-        """
-        Defines the shape of the state.
+        """Defines the shape of the state.
         For mamba layers this is usually a (conv_state, ssm_state) tuple.
         In this case, returns (conv_state_shape, ssm_state_shape).
         """
diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py
new file mode 100644
index 000000000000..3ba4062079cc
--- /dev/null
+++ b/vllm/model_executor/layers/mamba/checkpoint.py
@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.utils.math_utils import cdiv
+from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+from vllm.v1.kv_cache_interface import (
+    MambaSpec,
+    get_mamba_prefill_checkpoint_position,
+    is_mamba_prefill_checkpoint_valid,
+)
+
+
+def compute_mamba_prefill_checkpoints(
+    seq_lens: list[int],
+    query_lens: list[int],
+    hash_block_size: int,
+    mamba_block_size: int,
+    checkpoint_alignment: int | None,
+    drop_eagle_block: bool,
+) -> tuple[list[int], list[int]]:
+    """Per-row internal prefill checkpoint offsets and cache block columns.
+
+    Backends call this instead of re-deriving the rules, so they decline in
+    lockstep with the scheduler and ``MambaManager``: allocating a checkpoint
+    block without writing it leaves the prefix cache serving uninitialized
+    state.
+
+    Returns:
+        ``(offsets, cols)``: the checkpoint's offset into each row's query and
+        its block-table column. ``0`` and ``-1`` mean the row has none.
+
+    """
+    offsets: list[int] = []
+    cols: list[int] = []
+    for seq_len, query_len in zip(seq_lens, query_lens):
+        query_start = seq_len - query_len
+        position = get_mamba_prefill_checkpoint_position(
+            seq_len, hash_block_size, drop_eagle_block=drop_eagle_block
+        )
+        valid = is_mamba_prefill_checkpoint_valid(
+            query_start=query_start,
+            query_end=seq_len,
+            checkpoint_position=position,
+            hash_block_size=hash_block_size,
+            mamba_block_size=mamba_block_size,
+            checkpoint_alignment=checkpoint_alignment,
+        )
+        offsets.append(position - query_start if valid else 0)
+        cols.append(cdiv(seq_len, mamba_block_size) - 2 if valid else -1)
+    return offsets, cols
+
+
+@dataclass
+class MambaPrefillCheckpointMetadata:
+    checkpoint_offsets: torch.Tensor
+    state_indices: torch.Tensor
+
+
+class MambaPrefillCheckpointBuilder:
+    """Build backend-neutral prefill checkpoint locations."""
+
+    def __init__(self, vllm_config: VllmConfig, kv_cache_spec: MambaSpec) -> None:
+        self.vllm_config = vllm_config
+        self.kv_cache_spec = kv_cache_spec
+
+    def build(
+        self,
+        m: CommonAttentionMetadata,
+        request_rows: list[int],
+    ) -> MambaPrefillCheckpointMetadata | None:
+        if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0:
+            return None
+        assert m.seq_lens_cpu_upper_bound is not None
+        all_query_lens = m.query_start_loc_cpu.diff().tolist()
+        query_lens = [all_query_lens[row] for row in request_rows]
+        seq_lens = m.seq_lens_cpu_upper_bound.tolist()
+        block_size = self.kv_cache_spec.block_size
+        hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size
+        speculative_config = self.vllm_config.speculative_config
+        drop_eagle_block = (
+            speculative_config is not None and speculative_config.use_eagle_block_drop()
+        )
+        checkpoint_offsets, checkpoint_cols = compute_mamba_prefill_checkpoints(
+            [seq_lens[row] for row in request_rows],
+            query_lens,
+            hash_block_size=hash_block_size,
+            mamba_block_size=block_size,
+            checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment,
+            drop_eagle_block=drop_eagle_block,
+        )
+        if not any(checkpoint_offsets):
+            return None
+        checkpoint_offsets_tensor = async_tensor_h2d(
+            checkpoint_offsets,
+            dtype=torch.int32,
+            device=m.query_start_loc.device,
+        )
+        request_rows_tensor = async_tensor_h2d(
+            request_rows, dtype=torch.int64, device=m.query_start_loc.device
+        )
+        checkpoint_cols_tensor = async_tensor_h2d(
+            checkpoint_cols, dtype=torch.int64, device=m.query_start_loc.device
+        )
+        checkpoint_state_indices = m.block_table_tensor[
+            request_rows_tensor, checkpoint_cols_tensor
+        ]
+        checkpoint_state_indices = torch.where(
+            checkpoint_cols_tensor >= 0,
+            checkpoint_state_indices,
+            NULL_BLOCK_ID,
+        )
+        return MambaPrefillCheckpointMetadata(
+            checkpoint_offsets_tensor,
+            checkpoint_state_indices,
+        )
+
+
+class MambaPrefillCheckpointExporter(ABC):
+    """Export a mid-prefill checkpoint into backend-specific paged states."""
+
+    @abstractmethod
+    def export(
+        self,
+        checkpoint: MambaPrefillCheckpointMetadata,
+        *args,
+        **kwargs,
+    ) -> None:
+        """Write checkpoint state into the paged cache."""
+        raise NotImplementedError
diff --git a/vllm/model_executor/layers/mamba/gdn/base.py b/vllm/model_executor/layers/mamba/gdn/base.py
index 95ced2e0a4d3..249d56685ee7 100644
--- a/vllm/model_executor/layers/mamba/gdn/base.py
+++ b/vllm/model_executor/layers/mamba/gdn/base.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import torch
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import (
     VllmConfig,
@@ -24,7 +24,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         vllm_config: VllmConfig,
         prefix: str = "",
     ) -> None:
diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
index bf81cdd75f49..3892d1a27aa3 100644
--- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
+++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
@@ -329,8 +329,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         positions: torch.Tensor,
-        output: torch.Tensor,
-    ) -> None:
+    ) -> torch.Tensor:
         num_tokens = hidden_states.size(0)
         projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
         if self.use_full_rank_gate:
@@ -375,7 +374,7 @@ def forward(
             core_attn_out=core_attn_out,
         )
         core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
-        output[:] = self.o_proj(core_attn_out)[0]
+        return self.o_proj(core_attn_out)[0]
 
     @eager_break_during_capture
     def _forward(
diff --git a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py
index a1a494e121ad..c8f5a9ba64c9 100644
--- a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py
+++ b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py
@@ -45,8 +45,7 @@
 
 @PluggableLayer.register("olmo_hybrid_gated_delta_net_attention")
 class OlmoHybridGatedDeltaNetAttention(GatedDeltaNetAttention):
-    """
-    Gated DeltaNet linear attention layer for OLMo Hybrid.
+    """Gated DeltaNet linear attention layer for OLMo Hybrid.
 
     This implements the linear attention mechanism that replaces sliding window
     attention in the hybrid architecture.
@@ -294,9 +293,7 @@ def _forward_core(
         a: torch.Tensor,
         core_attn_out: torch.Tensor,
     ):
-        """
-        Core attention computation (called by custom op).
-        """
+        """Core attention computation (called by custom op)."""
         forward_context = get_forward_context()
         attn_metadata = forward_context.attn_metadata
 
diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
index bbbb1c8bd5ad..9d5bcc89dac1 100644
--- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
@@ -59,7 +59,6 @@
 from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd
 from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
-from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import (
     LayerNameType,
     _encode_layer_name,
@@ -167,6 +166,8 @@ def _log_gdn_backend_decision(
             head_k_dim,
         )
         return
+    elif current_platform.is_xpu():
+        return
 
     chosen = {
         "flashinfer": "FlashInfer",
@@ -226,6 +227,7 @@ def fi_chunk_gated_delta_rule(
         initial_state=fi_state,
         output_final_state=output_final_state,
         cu_seqlens=cu_seqlens,
+        backend="flashinfer",
     )
     # FlashInfer returns (output, state) when output_final_state=True,
     # or just output when output_final_state=False.
@@ -399,6 +401,7 @@ def __init__(
         self.key_dim = self.head_k_dim * self.num_k_heads
         self.value_dim = self.head_v_dim * self.num_v_heads
         self.gqa_interleaved_layout = gqa_interleaved_layout
+        self.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat"
         if current_platform.is_xpu():
             self._forward_method = self.forward_xpu
         elif current_platform.is_cpu():
@@ -523,6 +526,8 @@ def __init__(
                 self.gdn_decode_kernel = "triton"
         elif current_platform.is_cpu():
             self.gdn_decode_kernel = "CPU"
+        elif current_platform.is_xpu():
+            self.gdn_decode_kernel = "XPU"
 
         self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda"
         logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel)
@@ -639,9 +644,7 @@ def fix_query_key_value_ordering(
         mixed_qkvz: torch.Tensor,
         mixed_ba: torch.Tensor,
     ):
-        """
-        Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
-        """
+        """Derives `query`, `key` and `value` tensors from `mixed_qkvzba`."""
         new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (
             self.num_k_heads // self.tp_size,
             (
@@ -692,8 +695,7 @@ def prepare_gdn_attention_core_inputs(
         mixed_ba: torch.Tensor,
         num_tokens: int,
     ):
-        """
-        Derives mixed_qkv, z, b, a from projected qkvz/ba for the GDN custom op.
+        """Derives mixed_qkv, z, b, a from projected qkvz/ba for the GDN custom op.
 
         For gqa_interleaved_layout (Qwen3-Next): unpack the interleaved
         [ng, (hk + hk + np/ng*hv + np/ng*hv)] layout into contiguous qkv.
@@ -850,13 +852,8 @@ def _output_projection(
         The RMSNormGated + quant sequence is eligible for fusion
         by the compilation pass when fuse_norm_quant is enabled.
         """
-        z_shape_og = z.shape
-        core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
-        z = z.reshape(-1, z.shape[-1])
         core_attn_out = self.norm(core_attn_out, z)
-        core_attn_out = core_attn_out.reshape(z_shape_og)
-        core_attn_out = core_attn_out.flatten(-2)  # ... h d -> ... (h d)
-        output, _ = self.out_proj(core_attn_out)
+        output, _ = self.out_proj(core_attn_out.flatten(-2))
         return output
 
     def forward_hip(
@@ -899,8 +896,7 @@ def forward_cuda(
         self,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
-        """
-        Forward pass with three parts:
+        """Forward pass with three parts:
         1. Input projection
         2. Core attention (custom op)
         3. Output projection
@@ -977,8 +973,7 @@ def forward_xpu(
         self,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
-        """
-        Forward pass with three parts:
+        """Forward pass with three parts:
         1. Input projection
         2. Core attention (custom op)
         3. Output projection
@@ -1012,15 +1007,7 @@ def forward_xpu(
         # ============================================================
         # Part 3: Output Projection
         # ============================================================
-        z_shape_og = z.shape
-        # Reshape input data into 2D tensor
-        core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
-        z = z.reshape(-1, z.shape[-1])
-        core_attn_out = self.norm(core_attn_out, z)
-        core_attn_out = core_attn_out.reshape(z_shape_og)
-        core_attn_out = core_attn_out.flatten(-2)  # ... h d -> ... (h d)
-        out, _ = self.out_proj(core_attn_out)
-        return out
+        return self._output_projection(core_attn_out, z)
 
     def forward_cpu(
         self,
@@ -1063,14 +1050,7 @@ def forward_cpu(
             _encode_layer_name(self.prefix),
         )
 
-        z_shape_og = z.shape
-        core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
-        z = z.reshape(-1, z.shape[-1])
-        core_attn_out = self.norm(core_attn_out, z)
-        core_attn_out = core_attn_out.reshape(z_shape_og)
-        core_attn_out = core_attn_out.flatten(-2)  # ... h d -> ... (h d)
-        out, _ = self.out_proj(core_attn_out)
-        return out
+        return self._output_projection(core_attn_out, z)
 
     def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None:
         """Warm up GDN prefill kernels during V1 profiling.
@@ -1211,9 +1191,9 @@ def _forward_core_rocm(
         """ROCm AITER fast path: conv1d + recurrent attention from packed
         qkvz/ba layout.
 
-        For decode-only (no spec, no prefill) interleaved-GQA layouts,
-        dispatches directly to ``_forward_core_decode_aiter``. Otherwise unpacks
-        the packed layout and falls through to ``_forward_core``.
+        For decode-only (no spec, no prefill) batches, dispatches directly to
+        ``_forward_core_decode_aiter``. Otherwise unpacks the packed layout and
+        falls through to ``_forward_core``.
 
         Args:
             qkvz: packed [q, k, v, z] projection (num_tokens, qkvz_dim)
@@ -1221,6 +1201,7 @@ def _forward_core_rocm(
             z_out: **output** buffer for z        (num_tokens, num_heads,
                    head_dim); mutated in-place.
             core_attn_out: Pre-allocated output buffer for attention results.
+
         """
         forward_context = get_forward_context()
         attn_metadata_raw = forward_context.attn_metadata
@@ -1235,12 +1216,8 @@ def _forward_core_rocm(
 
         assert isinstance(attn_metadata, GDNAttentionMetadata)
 
-        # The AITER fused reshape/conv kernel expects Qwen3-Next's interleaved
-        # GQA layout. Qwen3.5 uses a non-interleaved q/k/v/z layout and must use
-        # the generic path below to split/rearrange inputs correctly.
         if (
-            self.gqa_interleaved_layout
-            and attn_metadata.spec_sequence_masks is None
+            attn_metadata.spec_sequence_masks is None
             and attn_metadata.num_prefills == 0
             and attn_metadata.num_decodes > 0
         ):
@@ -1279,6 +1256,7 @@ def _forward_core(
             b: beta gating vector                   (num_tokens, num_heads)
             a: alpha gating vector                  (num_tokens, num_heads)
             core_attn_out: Pre-allocated output buffer for attention results.
+
         """
         forward_context = get_forward_context()
         attn_metadata_raw = forward_context.attn_metadata
@@ -1575,14 +1553,10 @@ def _forward_core(
 
         # 3. Merge core attention output
         if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
-            merged_out = torch.empty(
-                (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
-                dtype=core_attn_out_non_spec.dtype,
-                device=core_attn_out_non_spec.device,
+            core_attn_out.index_copy_(0, spec_token_indx, core_attn_out_spec.squeeze(0))
+            core_attn_out.index_copy_(
+                0, non_spec_token_indx, core_attn_out_non_spec.squeeze(0)
             )
-            merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
-            merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
-            core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
         elif spec_sequence_masks is not None:
             core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
         else:
@@ -1632,6 +1606,7 @@ def _forward_core_decode_aiter(
                     : attn_metadata.num_actual_tokens
                 ],
                 validate_data=True,
+                qkvz_layout=self.qkvz_layout,
             )
         )
 
@@ -1662,9 +1637,7 @@ def _forward_core_decode_non_spec(
         core_attn_out: torch.Tensor,
         attn_metadata: GDNAttentionMetadata,
     ):
-        """
-        Core attention computation with a packed non-spec decode fast path.
-        """
+        """Core attention computation with a packed non-spec decode fast path."""
         non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor  # noqa: E501
         self_kv_cache = self.kv_cache
         # conv_state must be (..., dim, width-1) for the conv kernels.
@@ -1988,75 +1961,3 @@ def qwen_gdn_attention_core_fused_norm_packed(
     op_func=qwen_gdn_attention_core_fused_norm_packed,
     mutates_args=["core_attn_out"],
 )
-
-
-@triton.jit
-def fused_gdn_gating_kernel(
-    g,
-    beta_output,
-    A_log,
-    a,
-    b,
-    dt_bias,
-    seq_len,
-    NUM_HEADS: tl.constexpr,
-    beta: tl.constexpr,
-    threshold: tl.constexpr,
-    BLK_HEADS: tl.constexpr,
-):
-    i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)
-    off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off
-    mask = head_off < NUM_HEADS
-    blk_A_log = tl.load(A_log + head_off, mask=mask)
-    blk_a = tl.load(a + off, mask=mask)
-    blk_b = tl.load(b + off, mask=mask)
-    blk_bias = tl.load(dt_bias + head_off, mask=mask)
-    # If the model is loaded in fp16, without the .float() here, A might be -inf
-    x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)
-    softplus_x = tl.where(
-        beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x
-    )
-    blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x
-    tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)
-    # compute beta_output = sigmoid(b)
-    blk_beta_output = tl.sigmoid(blk_b.to(tl.float32))
-    tl.store(
-        beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask
-    )
-
-
-def fused_gdn_gating(
-    A_log: torch.Tensor,
-    a: torch.Tensor,
-    b: torch.Tensor,
-    dt_bias: torch.Tensor,
-    beta: float = 1.0,
-    threshold: float = 20.0,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Fused computation of g and beta for Gated Delta Net.
-    g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
-    beta_output = b.sigmoid()
-    TODO maybe use torch.compile to replace this triton kernel
-    """
-    batch, num_heads = a.shape
-    seq_len = 1
-    grid = (batch, seq_len, triton.cdiv(num_heads, 8))
-    g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)
-    beta_output = torch.empty(1, batch, num_heads, dtype=b.dtype, device=b.device)
-    fused_gdn_gating_kernel[grid](
-        g,
-        beta_output,
-        A_log,
-        a,
-        b,
-        dt_bias,
-        seq_len,
-        num_heads,
-        beta,
-        threshold,
-        8,
-        num_warps=1,
-    )
-    return g, beta_output
diff --git a/vllm/model_executor/layers/mamba/kda_checkpoint.py b/vllm/model_executor/layers/mamba/kda_checkpoint.py
new file mode 100644
index 000000000000..20741623d105
--- /dev/null
+++ b/vllm/model_executor/layers/mamba/kda_checkpoint.py
@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+
+import torch
+
+from vllm.model_executor.layers.mamba.checkpoint import (
+    MambaPrefillCheckpointExporter,
+    MambaPrefillCheckpointMetadata,
+)
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+
+
+def kda_prefill_checkpoint_alignment(backend: str) -> int | None:
+    return 16 if backend == "flashkda" else None
+
+
+@dataclass(frozen=True)
+class FlashKDAPrefillCheckpointExporter(MambaPrefillCheckpointExporter):
+    """Store FlashKDA recurrent and convolution checkpoint states."""
+
+    state_len: int | None = None
+
+    def export(
+        self,
+        checkpoint: MambaPrefillCheckpointMetadata,
+        *,
+        raw_qkv: torch.Tensor,
+        conv_state: torch.Tensor,
+        recurrent_checkpoint: torch.Tensor,
+        recurrent_state: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+    ) -> None:
+        state_len = (
+            self.state_len if self.state_len is not None else conv_state.shape[-1]
+        )
+        width = raw_qkv.shape[-1]
+        recurrent_row_size = recurrent_checkpoint[0].numel()
+        block_size = 256
+        store_cache_checkpoints_kernel[
+            (
+                checkpoint.checkpoint_offsets.numel(),
+                triton.cdiv(max(width * state_len, recurrent_row_size), block_size),
+            )
+        ](
+            raw_qkv,
+            conv_state,
+            recurrent_checkpoint,
+            recurrent_state,
+            cu_seqlens,
+            checkpoint.checkpoint_offsets,
+            checkpoint.state_indices,
+            raw_qkv.stride(0),
+            raw_qkv.stride(1),
+            conv_state.stride(0),
+            conv_state.stride(1),
+            conv_state.stride(2),
+            recurrent_checkpoint.stride(0),
+            recurrent_state.stride(0),
+            checkpoint.checkpoint_offsets.stride(0),
+            state_len,
+            width,
+            recurrent_row_size,
+            NULL_BLOCK_ID,
+            block_size,
+        )
+
+
+@triton.jit
+def store_cache_checkpoints_kernel(
+    x_ptr,
+    conv_state_ptr,
+    recurrent_checkpoint_ptr,
+    recurrent_state_ptr,
+    query_start_loc_ptr,
+    checkpoint_offsets_ptr,
+    checkpoint_state_indices_ptr,
+    x_stride_0: tl.constexpr,
+    x_stride_1: tl.constexpr,
+    state_stride_0: tl.constexpr,
+    state_stride_1: tl.constexpr,
+    state_stride_2: tl.constexpr,
+    checkpoint_stride_0: tl.constexpr,
+    recurrent_state_stride_0: tl.constexpr,
+    checkpoint_offset_stride: tl.constexpr,
+    STATE_LEN: tl.constexpr,
+    WIDTH: tl.constexpr,
+    RECURRENT_ROW_SIZE: tl.constexpr,
+    NULL_STATE_IDX: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    seq_idx = tl.program_id(0)
+    cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64)
+    checkpoint_offset = tl.load(
+        checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
+    )
+    valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)
+    valid_conv = (
+        (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)
+    )
+    width_idx = cols // STATE_LEN
+    history_idx = cols % STATE_LEN
+    checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset
+    token_idx = checkpoint_end - STATE_LEN + history_idx
+    values = tl.load(
+        x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
+        mask=valid_conv,
+    )
+    tl.store(
+        conv_state_ptr
+        + state_idx * state_stride_0
+        + width_idx * state_stride_1
+        + history_idx * state_stride_2,
+        values,
+        mask=valid_conv,
+    )
+
+    valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint
+    recurrent = tl.load(
+        recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols,
+        mask=valid_recurrent,
+    )
+    tl.store(
+        recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,
+        recurrent,
+        mask=valid_recurrent,
+    )
diff --git a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py
index 7f3b3f8dbaab..a4283e24821a 100644
--- a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py
+++ b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py
@@ -4,7 +4,7 @@
 
 import torch
 import torch.nn.functional as F
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.config import (
     VllmConfig,
@@ -37,7 +37,7 @@
 from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
 
 
-def _build_rope_parameters(config: PretrainedConfig) -> dict | None:
+def _build_rope_parameters(config: PreTrainedConfig) -> dict | None:
     rope_parameters = copy.deepcopy(getattr(config, "rope_parameters", None)) or {}
     if "rope_theta" not in rope_parameters and hasattr(config, "rope_theta"):
         rope_parameters["rope_theta"] = config.rope_theta
@@ -385,7 +385,7 @@ class BailingMoELinearAttention(LinearAttention):
     # --8<-- [end:bailing_moe_linear_attention]
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         vllm_config: VllmConfig,
         prefix: str = "linear_attn",
     ):
@@ -523,7 +523,7 @@ def forward(
         output: torch.Tensor,
         positions: torch.Tensor,
     ) -> None:
-        """Forward method called by torch.ops.vllm.linear_attention"""
+        """Forward method called by torch.ops.vllm.linear_attention."""
         torch.ops.vllm.linear_attention(
             hidden_states,
             output,
diff --git a/vllm/model_executor/layers/mamba/linear/base.py b/vllm/model_executor/layers/mamba/linear/base.py
index 73df07187303..da6808984b91 100644
--- a/vllm/model_executor/layers/mamba/linear/base.py
+++ b/vllm/model_executor/layers/mamba/linear/base.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import torch
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import (
     VllmConfig,
@@ -24,7 +24,7 @@ class LinearAttention(PluggableLayer, MambaBase):
     """Base class for Linear attention layer."""
 
     def __init__(
-        self, config: PretrainedConfig, vllm_config: VllmConfig, prefix: str = ""
+        self, config: PreTrainedConfig, vllm_config: VllmConfig, prefix: str = ""
     ):
         super().__init__()
         self.layer_idx = extract_layer_index(prefix)
diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py
index 5d09da4c3bfe..2b1b1f9c9df2 100644
--- a/vllm/model_executor/layers/mamba/mamba_mixer.py
+++ b/vllm/model_executor/layers/mamba/mamba_mixer.py
@@ -49,8 +49,7 @@
 # --8<-- [start:mamba_mixer]
 @PluggableLayer.register("mamba_mixer")
 class MambaMixer(MambaBase, PluggableLayer):
-    """
-    Compute ∆, A, B, C, and D the state space parameters and compute
+    """Compute ∆, A, B, C, and D the state space parameters and compute
     the `contextualized_states`. A, D are input independent
     (see Mamba paper [1] Section 3.5.2 "Interpretation of A"
     for why A isn't selective) ∆, B, C are input-dependent
@@ -237,8 +236,7 @@ def forward(self, hidden_states: torch.Tensor, output: torch.Tensor):
         )
 
     def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor):
-        """
-        Run the Mamba-1 SSM pipeline.
+        """Run the Mamba-1 SSM pipeline.
 
         Steps
         -----
@@ -258,7 +256,6 @@ def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor):
         decode tokens), both sets of kernels are executed independently
         and their outputs are concatenated before the final output projection.
         """
-
         forward_context: ForwardContext = get_forward_context()
         attn_metadata_raw = forward_context.attn_metadata
 
diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py
index 5bc890d0078b..db85ca0dee3f 100644
--- a/vllm/model_executor/layers/mamba/mamba_mixer2.py
+++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py
@@ -180,6 +180,13 @@ def forward_cuda(
             norm_before_gate=False,
         )
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        gate: torch.Tensor,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        return self.forward_cuda(x, gate)
+
 
 def mamba_v2_sharded_weight_loader(
     shard_spec: list[tuple[int, int, float]],
@@ -245,8 +252,7 @@ def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
 # --8<-- [start:mamba_mixer2]
 @PluggableLayer.register("mamba_mixer2")
 class MambaMixer2(MambaBase, PluggableLayer):
-    """
-    Compute ∆, A, B, C, and D the state space parameters and compute
+    """Compute ∆, A, B, C, and D the state space parameters and compute
     the `contextualized_states`. A, D are input independent
     (see Mamba paper [1] Section 3.5.2 "Interpretation of A"
     for why A isn't selective) ∆, B, C are input-dependent
@@ -1214,7 +1220,6 @@ def share_replayssm_ring_trackers(
     The final local layer in each group advances its cursors after every layer
     in that group has consumed the previous values.
     """
-
     replayssm_mixers: dict[str, MambaMixer2] = {}
     for layer_name in ordered_layer_names:
         layer = forward_context[layer_name]
diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py
index 8ca1d839556d..8c89488d2d58 100644
--- a/vllm/model_executor/layers/mamba/mamba_utils.py
+++ b/vllm/model_executor/layers/mamba/mamba_utils.py
@@ -222,12 +222,18 @@ def mamba2_state_shape(
     @classmethod
     def append_replayssm_ring(
         cls,
-        base_shapes: tuple[tuple[int, ...], ...],
+        base_shapes: tuple[tuple[int, int], tuple[int, int, int]],
         n_groups: int,
         tp_world_size: int,
         logical_window: int,
         backend: MambaBackendEnum,
-    ) -> tuple[tuple[int, ...], ...]:
+    ) -> tuple[
+        tuple[int, int],
+        tuple[int, int, int],
+        tuple[int, int, int],
+        tuple[int, int],
+        tuple[int, int, int],
+    ]:
         """Append the physical ReplaySSM ring shapes.
 
         ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``;
@@ -262,7 +268,6 @@ def short_conv_state_shape(
     def extra_groups_for_head_shards(cls, ngroups: int, tp_size: int):
         """Compute the increase in group numbers to account for
         replication in order to accompany the head shards."""
-
         # in the case ngoups % tp_size == 0, this will be zero
         if ngroups % tp_size == 0:
             return 0
@@ -344,13 +349,13 @@ def append_kda_recoverssm_record(
 
 @dataclass
 class MambaCopySpec:
-    """
-    Data class specifying the memory-copy parameters for Mamba states used for
+    """Data class specifying the memory-copy parameters for Mamba states used for
     prefix caching in align mode.
 
     Attributes:
         start_addr (int): Starting address for the memory copy operation.
         num_elements (int): Number of elements to copy from the starting address.
+
     """
 
     start_addr: int
diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py
index 8335c849666d..e8254e4f34eb 100644
--- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py
@@ -497,7 +497,7 @@ def causal_conv1d_fn(
     metadata=None,
     validate_data=False,
 ):
-    """support varlen + continuous batching when x is 2D tensor
+    """Support varlen + continuous batching when x is 2D tensor.
 
     x: (dim,cu_seq_len)
         cu_seq_len = total tokens of all seqs in that batch
@@ -1109,8 +1109,7 @@ def causal_conv1d_update(
     validate_data=False,
     out: torch.Tensor | None = None,
 ):
-    """
-    x: Input tensor which can take the following shapes:
+    """x: Input tensor which can take the following shapes:
 
     - `[batch, dim]` - single token prediction
     - `[batch, dim, seqlen]` - single or multiple tokens prediction
diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py
index c552d93445e6..ec8511f24bae 100644
--- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py
+++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py
@@ -123,8 +123,7 @@ def causal_conv1d_update_torch(
     bias: torch.Tensor | None = None,
     activation: str | None = None,
 ) -> torch.Tensor:
-    """
-    Pure PyTorch fallback for causal_conv1d_update.
+    """Pure PyTorch fallback for causal_conv1d_update.
     Currently used as a fallback for Arm (aarch64) to leverage
     oneDNN/ACL F.conv1d kernels for batched decoding.
     """
diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
index b5c1d11a1651..f3f3a33ab13e 100644
--- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
+++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
@@ -27,6 +27,104 @@
 _CPU_GDN_ATTENTION_OPS_REGISTERED = False
 
 
+def _stage_ds_conv_states(
+    conv_states: torch.Tensor, state_indices: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    storage = conv_states.transpose(1, 2).index_select(0, state_indices)
+    scratch = storage.transpose(1, 2)
+    identity_indices = torch.arange(
+        state_indices.numel(), dtype=torch.int32, device=state_indices.device
+    )
+    return scratch, identity_indices
+
+
+def _scatter_ds_conv_states(
+    conv_states: torch.Tensor,
+    state_indices: torch.Tensor,
+    scratch: torch.Tensor,
+) -> None:
+    conv_states.index_copy_(0, state_indices.to(torch.int64), scratch)
+
+
+def _causal_conv1d_fwd_cpu(
+    *,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    conv_states: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    cache_indices: torch.Tensor,
+    has_initial_state: torch.Tensor,
+    silu_activation: bool,
+    is_vnni: bool,
+) -> torch.Tensor:
+    if not is_conv_state_dim_first():
+        return ops.causal_conv1d_fwd_cpu(
+            x=x,
+            weight=weight,
+            bias=bias,
+            conv_states=conv_states,
+            query_start_loc=query_start_loc,
+            cache_indices=cache_indices,
+            has_initial_state=has_initial_state,
+            silu_activation=silu_activation,
+            is_vnni=is_vnni,
+        )
+
+    scratch, identity_indices = _stage_ds_conv_states(conv_states, cache_indices)
+    out = ops.causal_conv1d_fwd_cpu(
+        x=x,
+        weight=weight,
+        bias=bias,
+        conv_states=scratch,
+        query_start_loc=query_start_loc,
+        cache_indices=identity_indices,
+        has_initial_state=has_initial_state,
+        silu_activation=silu_activation,
+        is_vnni=is_vnni,
+    )
+    _scatter_ds_conv_states(conv_states, cache_indices, scratch)
+    return out
+
+
+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,
+    is_vnni: bool,
+    num_accepted_tokens: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if not is_conv_state_dim_first():
+        return ops.causal_conv1d_update_cpu(
+            x=x,
+            conv_states=conv_states,
+            weight=weight,
+            bias=bias,
+            silu_activation=silu_activation,
+            conv_state_indices=conv_state_indices,
+            is_vnni=is_vnni,
+            num_accepted_tokens=num_accepted_tokens,
+        )
+
+    scratch, identity_indices = _stage_ds_conv_states(conv_states, conv_state_indices)
+    out = ops.causal_conv1d_update_cpu(
+        x=x,
+        conv_states=scratch,
+        weight=weight,
+        bias=bias,
+        silu_activation=silu_activation,
+        conv_state_indices=identity_indices,
+        is_vnni=is_vnni,
+        num_accepted_tokens=num_accepted_tokens,
+    )
+    _scatter_ds_conv_states(conv_states, conv_state_indices, scratch)
+    return out
+
+
 def cpu_gdn_attention_core(
     mixed_qkv: torch.Tensor,
     b: torch.Tensor,
@@ -109,10 +207,8 @@ def _cpu_gdn_attention_nonspec(
 
     conv_state = layer.kv_cache[0]
     if use_cpp_conv:
-        # C++ conv requires [num_allocated_slots, kernel - 1, conv_dim] (SD).
-        if is_conv_state_dim_first():
-            raise RuntimeError("C++ CPU GDN attention requires `SD` conv_state layout.")
-        conv_state = conv_state.transpose(1, 2)
+        if not is_conv_state_dim_first():
+            conv_state = conv_state.transpose(1, 2)
     else:
         if not is_conv_state_dim_first():
             conv_state = conv_state.transpose(-1, -2)
@@ -141,7 +237,7 @@ def _cpu_gdn_attention_nonspec(
         decode_a = a[:num_decode_tokens]
         decode_state_indices = state_indices_tensor[:num_decodes]
         if use_cpp_conv:
-            decode_mixed_qkv = ops.causal_conv1d_update_cpu(
+            decode_mixed_qkv = _causal_conv1d_update_cpu(
                 x=decode_mixed_qkv,
                 conv_states=conv_state,
                 weight=layer.conv1d.weight,
@@ -210,7 +306,7 @@ def _cpu_gdn_attention_nonspec(
         ]
 
         if use_cpp_conv:
-            prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(
+            prefill_mixed_qkv = _causal_conv1d_fwd_cpu(
                 x=prefill_mixed_qkv.transpose(0, 1),
                 weight=layer.conv1d.weight,
                 bias=layer.conv1d.bias,
@@ -416,7 +512,6 @@ def _spec_forward(
 
     can_use_native_conv = (
         torch.cpu._is_amx_tile_supported()
-        and not is_conv_state_dim_first()
         and width == 4
         and num_spec_decodes > 0
         and bool(torch.all(seq_lens == seq_lens[0]).item())
@@ -424,7 +519,7 @@ def _spec_forward(
     )
     if can_use_native_conv:
         q_i = int(seq_lens[0].item())
-        conv_out = ops.causal_conv1d_update_cpu(
+        conv_out = _causal_conv1d_update_cpu(
             x=mixed_qkv_spec.view(num_spec_decodes, q_i, dim),
             conv_states=conv_buf,
             weight=layer.conv1d.weight,
@@ -509,8 +604,6 @@ def _spec_aware_nonspec(
     state_indices_tensor = state_indices_tensor.contiguous()
 
     is_amx = torch.cpu._is_amx_tile_supported()
-    if is_amx and is_conv_state_dim_first():
-        raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")
 
     if not is_amx:
         conv_weights = _unpacked_conv_weight(layer)
@@ -526,7 +619,7 @@ def _spec_aware_nonspec(
         decode_a = a[:num_decode_tokens]
         decode_state_indices = state_indices_tensor[:num_decodes]
         if is_amx:
-            decode_mixed_qkv = ops.causal_conv1d_update_cpu(
+            decode_mixed_qkv = _causal_conv1d_update_cpu(
                 x=decode_mixed_qkv,
                 conv_states=conv_buf,
                 weight=layer.conv1d.weight,
@@ -593,7 +686,7 @@ def _spec_aware_nonspec(
             num_decodes : num_decodes + num_prefills
         ]
         if is_amx:
-            prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(
+            prefill_mixed_qkv = _causal_conv1d_fwd_cpu(
                 x=prefill_mixed_qkv.transpose(0, 1),
                 weight=layer.conv1d.weight,
                 bias=layer.conv1d.bias,
@@ -664,11 +757,9 @@ def _spec_aware_nonspec_subset(
     prefill_state_indices = prefill_state_indices.contiguous()
 
     is_amx = torch.cpu._is_amx_tile_supported()
-    if is_amx and is_conv_state_dim_first():
-        raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")
 
     if is_amx:
-        conv_out = ops.causal_conv1d_fwd_cpu(
+        conv_out = _causal_conv1d_fwd_cpu(
             x=mixed_qkv.transpose(0, 1),
             weight=layer.conv1d.weight,
             bias=layer.conv1d.bias,
diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py
index 992f4ff68b28..ce1c0c8416eb 100644
--- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py
+++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py
@@ -173,6 +173,7 @@ def chunk_gated_delta_rule_cutedsl(
         ``[1, T, Hv, V]`` and ``final_state`` has shape ``[N, Hv, V, K]``.
         When ``core_attn_out`` is provided, ``output`` is an unsqueezed view of
         that buffer.
+
     """
     q = q.squeeze(0)
     k = k.squeeze(0)
diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py
index 7cdd652e7d2b..47458323e1ae 100644
--- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py
+++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py
@@ -728,7 +728,6 @@ def h_cutedsl(
     num_stages: int = 2,
 ) -> None:
     """Compute H/V_new with the same argument order as the CUDA wrapper."""
-
     _, H, K_dim = K.shape
     _, Hv, V_dim = V.shape
     h_dtype = {
diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py
index ba61ec68b5cc..4b4046a09327 100644
--- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py
+++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py
@@ -573,8 +573,8 @@ def store_ab_abg(
                 for i in cutlass.range_constexpr(4):
                     M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i])
 
-                # 3 rounds of Newton-Schulz
-                for _ in cutlass.range_constexpr(3):
+                # The fourth round keeps correlated-key KKT blocks stable in BF16.
+                for _ in cutlass.range_constexpr(4):
                     # First MMA: -AiM = Ai @ (-M)
                     cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)])
                     cute.arch.sync_warp()
diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py
index af45467886ea..0f2041d2263c 100644
--- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py
@@ -67,8 +67,7 @@ def _canonical_cache_dtype(cache_dtype: str) -> str:
 def get_ssm_configs(
     headdim: int, dstate: int, cache_dtype: str
 ) -> dict[int, Any] | None:
-    """
-    Return tuned (BLOCK_SIZE_M, num_warps) configs for *selective_state_update*
+    """Return tuned (BLOCK_SIZE_M, num_warps) configs for *selective_state_update*
     keyed by ``effective_batch = batch * nheads``, or ``None`` if no config
     file is found for the (headdim, dstate, cache_dtype, device) combination.
 
@@ -515,30 +514,29 @@ def selective_state_update(
     enable_stochastic_rounding=False,
     cache_philox_rounds=0,
 ):
-    """
-    Argument:
-        state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
-        x: (batch, dim) or (batch, nheads, dim)
-        dt: (batch, dim) or (batch, nheads, dim)
-        A: (dim, dstate) or (nheads, dim, dstate)
-        B: (batch, dstate) or (batch, ngroups, dstate)
-        C: (batch, dstate) or (batch, ngroups, dstate)
-        D: (dim,) or (nheads, dim)
-        z: (batch, dim) or (batch, nheads, dim)
-        dt_bias: (dim,) or (nheads, dim)
-        null_block_id: int
-            if state_batch_indices is passed, lets the kernel identify
-            padded entries that will not be processed,
-            for example: state_batch_indices = [null_block_id, 1, 20,
-            null_block_id] in this case, the kernel will not process
-            entries at indices 0 and 3
-        out: Preallocated ssm output tensor. Assume same shape as x.
-             In-place updated.
-        num_accepted_tokens: (batch,)
-            number of accepted tokens from previous verification step,
-            tells the kernel which initial state to use
-        cu_seqlens: (batch,)
-            length per sequence, for variable length in speculative decoding cases
+    """Argument:
+    state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
+    x: (batch, dim) or (batch, nheads, dim)
+    dt: (batch, dim) or (batch, nheads, dim)
+    A: (dim, dstate) or (nheads, dim, dstate)
+    B: (batch, dstate) or (batch, ngroups, dstate)
+    C: (batch, dstate) or (batch, ngroups, dstate)
+    D: (dim,) or (nheads, dim)
+    z: (batch, dim) or (batch, nheads, dim)
+    dt_bias: (dim,) or (nheads, dim)
+    null_block_id: int
+        if state_batch_indices is passed, lets the kernel identify
+        padded entries that will not be processed,
+        for example: state_batch_indices = [null_block_id, 1, 20,
+        null_block_id] in this case, the kernel will not process
+        entries at indices 0 and 3
+    out: Preallocated ssm output tensor. Assume same shape as x.
+         In-place updated.
+    num_accepted_tokens: (batch,)
+        number of accepted tokens from previous verification step,
+        tells the kernel which initial state to use
+    cu_seqlens: (batch,)
+        length per sequence, for variable length in speculative decoding cases
     """
     if state.dim() == 3:
         state = state.unsqueeze(1)
@@ -723,8 +721,7 @@ def selective_scan_fn(
     cu_chunk_seqlen=None,
     last_chunk_indices=None,
 ) -> torch.Tensor:
-    """
-    u: (dim, total_length) for varlen or (batch, dim, seqlen)
+    """u: (dim, total_length) for varlen or (batch, dim, seqlen)
         applies changes in place.
     ssm_states: (batch, dim, dstate) or (batch, nheads, dim, dstate)
         applies changes in place.
@@ -769,9 +766,12 @@ def selective_scan_fn(
     initial_state_idx: (batch,), dtype int32
         The pointer into cache_indices, where the cache block
         containing the initial state is located.
-    returns
+
+    Returns
+    -------
         output: (dim, total_length) for varlen or (batch, dim, seqlen)
                 supports inplace replacement
+
     """
     if u.stride(-1) != 1:
         u = u.contiguous()
diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py
index aff4fc1c3f2d..3cee44ba5015 100644
--- a/vllm/model_executor/layers/mamba/ops/replayssm_config.py
+++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py
@@ -8,7 +8,6 @@
 """
 
 import functools
-from contextlib import contextmanager
 
 from vllm.platforms import current_platform
 from vllm.triton_utils import triton
@@ -26,20 +25,6 @@ def _is_blackwell() -> bool:
 _overrides: dict[str, tuple] = {}
 
 
-@contextmanager
-def override_replayssm_config(kernel: str, config: tuple):
-    """Pin ``kernel``'s launch config for the duration of the context."""
-    prev = _overrides.get(kernel)
-    _overrides[kernel] = config
-    try:
-        yield
-    finally:
-        if prev is None:
-            _overrides.pop(kernel, None)
-        else:
-            _overrides[kernel] = prev
-
-
 def _dstate_tile(dstate: int, tile: int) -> int:
     return max(16, min(tile, triton.next_power_of_2(dstate)))
 
diff --git a/vllm/model_executor/layers/mamba/ops/scatter_states.py b/vllm/model_executor/layers/mamba/ops/scatter_states.py
index 0102b6cf12b4..04b5cb61c2b9 100644
--- a/vllm/model_executor/layers/mamba/ops/scatter_states.py
+++ b/vllm/model_executor/layers/mamba/ops/scatter_states.py
@@ -1,8 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from typing import Any
+
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import WarmupIntRange
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
@@ -33,11 +41,39 @@ def _scatter_states_kernel(
     tl.store(state_ptr + state_idx * stride_state_batch + offsets, values, mask=mask)
 
 
+def _scatter_states_warmup_inputs(
+    *,
+    state_shape: tuple[int, ...],
+    state_dtype: torch.dtype,
+    indices_dtype: torch.dtype,
+    max_num_tokens: int,
+) -> dict[str, object]:
+    num_tokens: Any = WarmupIntRange(1, max_num_tokens + 1)
+    return dict(
+        state=TritonWarmupTensor(
+            state_dtype,
+            shape=(1,) + state_shape,
+        ),
+        src=TritonWarmupTensor(
+            state_dtype,
+            shape=(num_tokens,) + state_shape,
+        ),
+        indices=TritonWarmupTensor(
+            indices_dtype,
+            shape=(num_tokens,),
+        ),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_scatter_states_kernel,
+    warmup_inputs=_scatter_states_warmup_inputs,
+)
 def scatter_states(
     state: torch.Tensor,
     src: torch.Tensor,
     indices: torch.Tensor,
-) -> None:
+) -> DispatchSpec:
     """Scatter ``src`` rows into ``state`` at ``indices`` (in place).
 
     Equivalent to ``state[indices] = src`` but non-atomic and bandwidth-bound,
@@ -58,13 +94,10 @@ def scatter_states(
     assert src[0].is_contiguous()
     block_size = min(triton.next_power_of_2(row_size), 1024)
     grid = (triton.cdiv(row_size, block_size), indices.numel())
-    _scatter_states_kernel[grid](
-        state,
-        src,
-        indices,
-        state.stride(0),
-        src.stride(0),
-        indices.stride(0),
+    return grid, dict(
+        stride_state_batch=state.stride(0),
+        stride_src_batch=src.stride(0),
+        stride_indices=indices.stride(0),
         row_size=row_size,
         BLOCK_SIZE=block_size,
         num_warps=8,
diff --git a/vllm/model_executor/layers/mamba/ops/ssd_bmm.py b/vllm/model_executor/layers/mamba/ops/ssd_bmm.py
index 7cf5acf2fab3..7dcb54703fac 100644
--- a/vllm/model_executor/layers/mamba/ops/ssd_bmm.py
+++ b/vllm/model_executor/layers/mamba/ops/ssd_bmm.py
@@ -146,16 +146,17 @@ def _bmm_chunk_fwd_kernel(
 
 
 def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtype=None):
-    """
-    Argument:
+    """Argument:
         a: (seqlen, ngroups, k)
         b: (seqlen, ngroups, k)
         chunk_size: int
         cu_chunk_seq_lens: (nchunks+1,)
         causal: if True, then out[i, j] for i > j will be arbitrary, only out[i, j] for i <= j are
             guaranteed to be correct.
+
     Return:
         out: (nchunks, ngroups, chunk_size, chunk_size)
+
     """
     seqlen, ngroups, k = a.shape
     assert b.shape == a.shape
diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py
index 8c645574b9e7..581cd8bb6d6f 100644
--- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py
+++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py
@@ -175,8 +175,7 @@ def mamba_chunk_scan_combined_varlen(
     return_intermediate_states=False,
     state_dtype=None,
 ):
-    """
-    Argument:
+    """Argument:
         x: (seqlen, nheads, headdim)
         dt: (seqlen, nheads)
         A: (nheads)
@@ -198,7 +197,6 @@ def mamba_chunk_scan_combined_varlen(
     Return:
         varlen_states: (batch, nheads, headdim, dstate)
     """
-
     assert cu_seqlens is not None, "cu_seqlens must be provided assuming varlen input"
     assert seq_idx is not None
 
diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
index 62e4307b37b7..b622924327e8 100644
--- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
+++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Dispatch module for Mamba selective state update (SSU) backends.
+"""Dispatch module for Mamba selective state update (SSU) backends.
 
 Provides a unified `selective_state_update` function that dispatches to
 the Triton, FlashInfer, or CPU backend based on the configured
diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py
index 6ddf4ff0660f..73993fbc984a 100644
--- a/vllm/model_executor/layers/mhc.py
+++ b/vllm/model_executor/layers/mhc.py
@@ -292,6 +292,229 @@ def forward_cpu(
         )
 
 
+# --8<-- [start:mhc_pre_delayed]
+@CustomOp.register("mhc_pre_delayed")
+class MHCPreDelayedOp(CustomOp):
+    """MHC pre block using the pre-mix carried from the previous sublayer.
+
+    Same gates as :class:`MHCPreOp`, but the stream collapse applies the
+    caller's ``pre_mix`` and this sublayer's pre-mix is returned for the next
+    sublayer seam.
+
+    Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` also
+    applies the preceding post block, which lets AITER fold it into the pre
+    projection and saves a launch. Returns residual, post_mix, comb_mix,
+    layer_input, next_pre_mix; ``residual`` is the caller's own tensor when
+    no post is requested.
+    """
+
+    # --8<-- [end:mhc_pre_delayed]
+    @classmethod
+    def enabled(cls) -> bool:
+        return True
+
+    def __init__(self) -> None:
+        super().__init__()
+        # Built here, not lazily: CustomOp resolves dispatch against the
+        # current vLLM config, which is only set during model construction.
+        self._post = MHCPostOp()
+
+    def _maybe_post(
+        self,
+        residual: torch.Tensor,
+        sublayer_out: torch.Tensor | None,
+        post_layer_mix: torch.Tensor | None,
+        comb_res_mix: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if sublayer_out is None:
+            return residual
+        assert post_layer_mix is not None and comb_res_mix is not None
+        return self._post(sublayer_out, residual, post_layer_mix, comb_res_mix)
+
+    def forward_cuda(
+        self,
+        residual: torch.Tensor,
+        fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        pre_mix: torch.Tensor | None = None,
+        x: torch.Tensor | None = None,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 1e-6,
+        sublayer_out: torch.Tensor | None = None,
+        post_layer_mix: torch.Tensor | None = None,
+        comb_res_mix: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        residual = self._maybe_post(
+            residual, sublayer_out, post_layer_mix, comb_res_mix
+        )
+        return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(
+            residual,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            pre_mix,
+            x,
+            norm_weight,
+            norm_eps,
+        )
+
+    def forward_hip(
+        self,
+        residual: torch.Tensor,
+        fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        pre_mix: torch.Tensor | None = None,
+        x: torch.Tensor | None = None,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 1e-6,
+        sublayer_out: torch.Tensor | None = None,
+        post_layer_mix: torch.Tensor | None = None,
+        comb_res_mix: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        # The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual`
+        # and folds no RMSNorm, so it cannot serve the model-entry broadcast
+        # (which projects a narrower `x`) or a fused norm. Both are handled by
+        # TileLang, or by the reference when TileLang is unavailable.
+        if (
+            x is None
+            and norm_weight is None
+            and _aiter_mhc_supported(residual, None, supports_norm=False)
+        ):
+            from vllm._aiter_ops import rocm_aiter_ops
+
+            num_tokens = residual.numel() // (residual.shape[-1] * residual.shape[-2])
+            # Folding the post in only pays while the residual still fits in
+            # cache; past that AITER's heuristic wants the separate kernels,
+            # which can ask for a non-temporal store and a large-m split-k.
+            if sublayer_out is not None and (
+                rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(num_tokens)
+            ):
+                residual = self._maybe_post(
+                    residual, sublayer_out, post_layer_mix, comb_res_mix
+                )
+                sublayer_out = None
+            # The op writes the folded post's residual into this buffer rather
+            # than returning it, since on the unfused path it has none of its
+            # own and an op must not hand back one of its inputs.
+            residual_out = (
+                torch.empty_like(residual) if sublayer_out is not None else None
+            )
+            rest = torch.ops.vllm.mhc_pre_delayed_aiter(
+                residual,
+                fn,
+                hc_scale,
+                hc_base,
+                rms_eps,
+                hc_pre_eps,
+                hc_sinkhorn_eps,
+                hc_post_mult_value,
+                sinkhorn_repeat,
+                pre_mix,
+                sublayer_out,
+                post_layer_mix,
+                comb_res_mix,
+                residual_out,
+            )
+            return (residual if residual_out is None else residual_out), *rest
+        residual = self._maybe_post(
+            residual, sublayer_out, post_layer_mix, comb_res_mix
+        )
+        if HAS_TILELANG_MHC:
+            return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(
+                residual,
+                fn,
+                hc_scale,
+                hc_base,
+                rms_eps,
+                hc_pre_eps,
+                hc_sinkhorn_eps,
+                hc_post_mult_value,
+                sinkhorn_repeat,
+                pre_mix,
+                x,
+                norm_weight,
+                norm_eps,
+            )
+        # The post is already applied, so it must not be requested again.
+        return self.forward_native(
+            residual,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            pre_mix,
+            x,
+            norm_weight,
+            norm_eps,
+        )
+
+    def forward_native(
+        self,
+        residual: torch.Tensor,
+        fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        pre_mix: torch.Tensor | None = None,
+        x: torch.Tensor | None = None,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 1e-6,
+        sublayer_out: torch.Tensor | None = None,
+        post_layer_mix: torch.Tensor | None = None,
+        comb_res_mix: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        residual = self._maybe_post(
+            residual, sublayer_out, post_layer_mix, comb_res_mix
+        )
+        post_mix, comb_mix, layer_input, next_pre_mix = (
+            mhc_kernels.mhc_pre_delayed_torch(
+                residual,
+                fn,
+                hc_scale,
+                hc_base,
+                rms_eps,
+                hc_pre_eps,
+                hc_sinkhorn_eps,
+                hc_post_mult_value,
+                sinkhorn_repeat,
+                pre_mix=pre_mix,
+                x=x,
+            )
+        )
+        return (
+            residual,
+            post_mix,
+            comb_mix,
+            _apply_mhc_norm(layer_input, norm_weight, norm_eps),
+            next_pre_mix,
+        )
+
+
 # --8<-- [start:mhc_post]
 @CustomOp.register("mhc_post")
 class MHCPostOp(CustomOp):
diff --git a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py
index afae19c75feb..0b2ef3a6aabd 100644
--- a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py
+++ b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py
@@ -60,8 +60,7 @@ def _cuda_memcpy_d2d(dst: int, src: int, size: int):
 
 
 class IpcBuffer:
-    """
-    Allocates CUDA device memory and exchanges IPC handles with all ranks
+    """Allocates CUDA device memory and exchanges IPC handles with all ranks
     so that every rank holds a valid device pointer to every other rank's buffer.
     """
 
@@ -133,8 +132,7 @@ def __del__(self):
 
 
 def _lamport_fill_neg_zero(device_ptr: int, size_bytes: int):
-    """
-    Fill device memory with IEEE-754 negative zero (-0.0f = 0x80000000).
+    """Fill device memory with IEEE-754 negative zero (-0.0f = 0x80000000).
     This is the "slot empty" sentinel for the Lamport protocol: the kernel
     spin-waits until a value is *not* negative zero.
     """
@@ -153,8 +151,7 @@ def _lamport_fill_neg_zero(device_ptr: int, size_bytes: int):
 
 
 class LamportWorkspace:
-    """
-    Self-contained workspace for Lamport-based cross-GPU AllReduce.
+    """Self-contained workspace for Lamport-based cross-GPU AllReduce.
 
     Parameters
     ----------
@@ -170,6 +167,7 @@ class LamportWorkspace:
     process_group : optional
         ``torch.distributed`` process group for IPC handle exchange.
         ``None`` uses the default group.
+
     """
 
     def __init__(self, rank: int, world_size: int, comm_size: int, process_group=None):
@@ -225,8 +223,7 @@ def compute_comm_size_for_minimax(
         world_size: int,
         fused_qk: bool = True,
     ) -> int:
-        """
-        Return a safe ``comm_size`` (in bytes) for MiniMaxReduceRMSKernel.
+        """Return a safe ``comm_size`` (in bytes) for MiniMaxReduceRMSKernel.
 
         The kernel stores per-token variance scalars in the Lamport buffer:
           - single-matrix path: ``world_size × max_tokens × 4`` bytes per slot
@@ -271,8 +268,7 @@ def get_allreduce_workspace(
     max_tokens: int = 16384,
     process_group=None,
 ) -> torch.Tensor:
-    """
-    Return a cached workspace tensor for the given (rank, world_size) pair.
+    """Return a cached workspace tensor for the given (rank, world_size) pair.
 
     On first call the workspace is allocated and IPC handles are exchanged;
     subsequent calls with the same arguments return the cached tensor.
@@ -288,6 +284,7 @@ def get_allreduce_workspace(
         Maximum number of tokens per batch (used when ``comm_size is None``).
     process_group : optional
         ``torch.distributed`` process group.
+
     """
     if comm_size is None:
         comm_size = LamportWorkspace.compute_comm_size_for_minimax(
diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py
index 421181306236..9dd5ad67736c 100644
--- a/vllm/model_executor/layers/mla.py
+++ b/vllm/model_executor/layers/mla.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -11,6 +12,11 @@
 from vllm.models.common.ops import fused_q_kv_rmsnorm
 from vllm.platforms import current_platform
 
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.mla.index_group import (
+        SparseMLAIndexGroupBuilder,
+    )
+
 
 @dataclass
 class MLAModules:
@@ -30,6 +36,7 @@ class MLAModules:
     topk_indices_buffer: torch.Tensor | None
     indexer_rotary_emb: torch.nn.Module | None = None
     g_proj: torch.nn.Module | None = None
+    index_group_builder: "SparseMLAIndexGroupBuilder | None" = None
 
 
 # --8<-- [start:multi_head_latent_attention]
@@ -128,6 +135,7 @@ def __init__(
             use_sparse=self.is_sparse,
             indexer=self.indexer,
             topk_indices_buffer=mla_modules.topk_indices_buffer,
+            index_group_builder=mla_modules.index_group_builder,
             non_causal_multi_token_decode=non_causal_multi_token_decode,
         )
         indexer_op = getattr(self.indexer, "indexer_op", None)
@@ -223,6 +231,8 @@ def forward(
 
         if self.indexer and self.is_sparse and not self.skip_topk:
             self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb)
+        if self.is_sparse:
+            self.mla_attn.impl.record_logical_topk_ready()  # type: ignore[attr-defined]
 
         if llama_4_scaling is not None:
             q *= llama_4_scaling
diff --git a/vllm/model_executor/layers/pooler/abstract.py b/vllm/model_executor/layers/pooler/abstract.py
index 82abef4f69ad..a32e38b98638 100644
--- a/vllm/model_executor/layers/pooler/abstract.py
+++ b/vllm/model_executor/layers/pooler/abstract.py
@@ -22,9 +22,7 @@ def get_supported_tasks(self) -> Set[PoolingTask]:
         raise NotImplementedError
 
     def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
-        """
-        Construct the updated pooling parameters to use for a supported task.
-        """
+        """Construct the updated pooling parameters to use for a supported task."""
         return PoolingParamsUpdate()
 
     @abstractmethod
diff --git a/vllm/model_executor/layers/pooler/activations.py b/vllm/model_executor/layers/pooler/activations.py
index f0c64d2720b7..fc47684a128d 100644
--- a/vllm/model_executor/layers/pooler/activations.py
+++ b/vllm/model_executor/layers/pooler/activations.py
@@ -7,7 +7,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import ModelConfig
 from vllm.logger import init_logger
@@ -17,7 +17,7 @@
 
 
 def get_act_fn(
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
     static_num_labels: bool = True,
 ) -> "PoolerActivation":
     # get classification act_fn
diff --git a/vllm/model_executor/layers/pooler/seqwise/poolers.py b/vllm/model_executor/layers/pooler/seqwise/poolers.py
index e9790cd2d243..2115e5184ed5 100644
--- a/vllm/model_executor/layers/pooler/seqwise/poolers.py
+++ b/vllm/model_executor/layers/pooler/seqwise/poolers.py
@@ -42,8 +42,7 @@
 
 
 class SequencePooler(Pooler):
-    """
-    A layer that pools specific information from hidden states.
+    """A layer that pools specific information from hidden states.
 
     This layer does the following:
     1. Extracts specific tokens or aggregates data based on pooling method.
diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py
index 090991bedb54..23354d66f3c4 100644
--- a/vllm/model_executor/layers/pooler/special.py
+++ b/vllm/model_executor/layers/pooler/special.py
@@ -74,6 +74,20 @@ def get_supported_tasks(self) -> Set[PoolingTask]:
     def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
         return self.poolers_by_task[task].get_pooling_updates(task)
 
+    def replace_classifier(
+        self,
+        new_classifier: ClassifierFn,
+    ) -> None:
+        """Replaces the classifier to the LoRA-wrapped version."""
+        for task, pooler in self.poolers_by_task.items():
+            if task not in {"classify"}:
+                # Now LoRA only supports classification tasks,
+                # so we only replace the classifier for "classify" task.
+                continue
+            head = getattr(pooler, "head", None)
+            if head is not None and hasattr(head, "classifier"):
+                head.classifier = new_classifier
+
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/vllm/model_executor/layers/pooler/tokwise/poolers.py b/vllm/model_executor/layers/pooler/tokwise/poolers.py
index 96eb28cf2667..0dab2d5b1d57 100644
--- a/vllm/model_executor/layers/pooler/tokwise/poolers.py
+++ b/vllm/model_executor/layers/pooler/tokwise/poolers.py
@@ -46,8 +46,7 @@
 
 
 class TokenPooler(Pooler):
-    """
-    A layer that pools specific information from hidden states.
+    """A layer that pools specific information from hidden states.
 
     This layer does the following:
     1. Extracts specific tokens or aggregates data based on pooling method.
diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py
index a2c65738e202..39e613a61147 100644
--- a/vllm/model_executor/layers/quantization/__init__.py
+++ b/vllm/model_executor/layers/quantization/__init__.py
@@ -82,6 +82,7 @@ def register_quantization_config(quantization: str):
         >>>
         >>> get_quantization_config("my_quant")
         <class 'MyQuantConfig'>
+
     """  # noqa: E501
 
     def _wrapper(quant_config_cls):
@@ -115,7 +116,20 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
     # lazy import to avoid triggering `torch.compile` too early
     from vllm.config.quantization import _ONLINE_SHORTHANDS
     from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
-    from vllm.models.deepseek_v4 import DeepseekV4FP8Config
+
+    if current_platform.is_cuda() or current_platform.is_rocm():
+        # The v4.1 class is the v4 one extended to accept model_type
+        # "deepseek_v41" and its 32x32 MXFP8 linear layout. V4.1 has
+        # platform-specific implementations for both CUDA and ROCm.
+        from vllm.models.deepseek_v41 import (
+            DeepseekV4FP8Config as DeepseekV41FP8Config,
+        )
+
+        deepseek_config: type[QuantizationConfig] = DeepseekV41FP8Config
+    else:
+        from vllm.models.deepseek_v4 import DeepseekV4FP8Config
+
+        deepseek_config = DeepseekV4FP8Config
 
     from .auto_awq import AutoAWQConfig
     from .auto_gptq import AutoGPTQConfig
@@ -161,7 +175,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
         "inc": INCConfig,
         "mxfp4": Mxfp4Config,
         "gpt_oss_mxfp4": GptOssMxfp4Config,
-        "deepseek_v4_fp8": DeepseekV4FP8Config,
+        "deepseek_v4_fp8": deepseek_config,
         "humming": HummingConfig,
         "online": OnlineQuantizationConfig,
         # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the
diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py
index f8c7269add87..0c3e9805517a 100644
--- a/vllm/model_executor/layers/quantization/auto_awq.py
+++ b/vllm/model_executor/layers/quantization/auto_awq.py
@@ -6,7 +6,7 @@
 import torch
 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
 from torch.nn import Parameter
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 import vllm.model_executor.layers.fused_moe  # noqa
 from vllm import _custom_ops as ops
@@ -48,7 +48,6 @@
     check_marlin_supports_layer,
     check_moe_marlin_supports_layer,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
     verify_marlin_supported,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -93,7 +92,7 @@ def _replace_or_register_parameter(
 def _convert_awq_to_standard_format(
     layer: torch.nn.Module,
     w_q_name: str,
-    w_zp_name: str,
+    w_zp_name: str | None,
     size_bits: int,
 ) -> None:
     """Convert AWQ weight and zero-point tensors to standard GPTQ-like format.
@@ -101,6 +100,7 @@ def _convert_awq_to_standard_format(
     AWQ packs qweight along the output dim with a non-standard bit order.
     This converts to standard bit order and repacks qweight along the input
     dim, matching the format expected by the MPLinearKernel framework.
+    If w_zp_name is None (symmetric quantization), only the weight is converted.
     """
     pack_factor = 32 // size_bits
     mask = (1 << size_bits) - 1
@@ -139,6 +139,9 @@ def _noop_loader(*args, **kwargs):
     )
     setattr(layer, w_q_name, new_param)
 
+    if w_zp_name is None:
+        return
+
     # --- Convert qzeros: fix AWQ bit ordering and repack
     # AWQ qzeros: (G, N // pack) packed along dim 1, AWQ bit order
     # Target: (N // pack, G) packed along dim 0, standard bit order
@@ -367,7 +370,7 @@ def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
     def maybe_update_config(
         self,
         model_name: str,
-        hf_config: PretrainedConfig | None = None,
+        hf_config: PreTrainedConfig | None = None,
         revision: str | None = None,
     ):
         if self.modules_to_not_convert:
@@ -393,6 +396,7 @@ class AutoAWQMarlinLinearMethod(LinearMethodBase):
 
     Args:
         quant_config: The AWQ Marlin quantization config.
+
     """
 
     _kernel_backends_being_used: set[str] = set()
@@ -637,9 +641,6 @@ def create_weights(
         layer.register_parameter("w2_qzeros", w2_qzeros)
         set_weight_attrs(w2_qzeros, extra_weight_attrs)
 
-        device = layer.w13_qweight.device
-        layer.workspace = marlin_make_workspace_new(device, 4)
-
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         converted = convert_to_wna16_moe_kernel_format(
             backend=self.wna16_moe_backend,
@@ -701,7 +702,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
 
     def _setup_kernel(self, layer: RoutedExperts) -> None:
         """Build the FusedMoEKernel for this layer."""
-
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
         self.moe_kernel = make_wna16_moe_kernel(
             moe_quant_config=self.moe_quant_config,
@@ -713,7 +713,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None:
 
     def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
         if self.wna16_moe_backend == WNA16MoEBackend.HUMMING:
-            from vllm.model_executor.layers.quantization.utils.humming_utils import (
+            from vllm.model_executor.layers.quantization.utils.humming import (
                 get_humming_moe_quant_config,
             )
 
@@ -880,6 +880,7 @@ class AutoAWQLinearMethod(BaseAWQLinearMethod):
 
     Args:
         quant_config: The AWQ quantization config.
+
     """
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py
index 893f85b207f4..ff57f3b8c508 100644
--- a/vllm/model_executor/layers/quantization/auto_gptq.py
+++ b/vllm/model_executor/layers/quantization/auto_gptq.py
@@ -6,7 +6,7 @@
 
 import torch
 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 import vllm.model_executor.layers.fused_moe  # noqa
 from vllm.logger import init_logger
@@ -16,7 +16,6 @@
 )
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEConfig,
-    FusedMoEExpertsModular,
     FusedMoEMethodBase,
     FusedMoEQuantConfig,
     FusedMoeWeightScaleSupported,
@@ -46,7 +45,6 @@
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     check_moe_marlin_supports_layer,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
     marlin_repeat_scales_on_all_ranks,
     verify_marlin_supported,
 )
@@ -275,7 +273,7 @@ def apply_vllm_mapper(self, hf_to_vllm_mapper):
     def maybe_update_config(
         self,
         model_name: str,
-        hf_config: PretrainedConfig | None = None,
+        hf_config: PreTrainedConfig | None = None,
         revision: str | None = None,
     ):
         if self.modules_in_block_to_quantize:
@@ -305,6 +303,7 @@ class AutoGPTQLinearMethod(LinearMethodBase):
 
     Args:
         quant_config: The AutoGPTQ quantization config.
+
     """
 
     _kernel_backends_being_used: set[str] = set()
@@ -600,12 +599,6 @@ def create_weights(
         layer.register_parameter("w2_bias", w2_bias)
         set_weight_attrs(w2_bias, extra_weight_attrs)
 
-        if self.experts_cls is not None and issubclass(
-            self.experts_cls, FusedMoEExpertsModular
-        ):
-            device = layer.w13_qweight.device
-            layer.workspace = marlin_make_workspace_new(device, 4)
-
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
         def replace_or_register(name: str, val: torch.Tensor | None):
             if val is None:
@@ -685,7 +678,6 @@ def replace_or_register(name: str, val: torch.Tensor | None):
 
     def _setup_kernel(self, layer: RoutedExperts) -> None:
         """Build the FusedMoEKernel for this layer."""
-
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
         self.moe_kernel = make_wna16_moe_kernel(
             moe_quant_config=self.moe_quant_config,
@@ -697,7 +689,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None:
 
     def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
         if self.wna16_moe_backend == WNA16MoEBackend.HUMMING:
-            from vllm.model_executor.layers.quantization.utils.humming_utils import (
+            from vllm.model_executor.layers.quantization.utils.humming import (
                 get_humming_moe_quant_config,
             )
 
diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py
index d1f8df2df7bb..1b0457854eb1 100644
--- a/vllm/model_executor/layers/quantization/base_config.py
+++ b/vllm/model_executor/layers/quantization/base_config.py
@@ -8,7 +8,7 @@
 import regex as re
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -86,8 +86,7 @@ def process_weights_after_loading(self, layer: nn.Module) -> None:
 
 
 def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool:
-    """
-    Not all quant methods have embedding implemented, so we need to check that
+    """Not all quant methods have embedding implemented, so we need to check that
     it exists for our given method. We check this by making sure the function
     has been changed from the base implementation.
     """
@@ -158,8 +157,7 @@ def override_quantization_method(
         user_quant: str | None,
         hf_config: Any = None,
     ) -> QuantizationMethods | None:
-        """
-        Detects if this quantization method can support a given checkpoint
+        """Detects if this quantization method can support a given checkpoint
         format by overriding the user specified quantization method --
         this method should only be overwritten by subclasses in exceptional
         circumstances.
@@ -169,6 +167,7 @@ def override_quantization_method(
             user_quant: The user-specified quantization method string.
             hf_config: The HuggingFace model config object (e.g. for
                 model_type checks). May be None if not available.
+
         """
         return None
 
@@ -203,6 +202,7 @@ def get_quant_method(
         Returns:
             The quantize method. None if the given layer doesn't support quant
             method.
+
         """
         raise NotImplementedError
 
@@ -251,13 +251,13 @@ def get_checkpoint_weight_mapper() -> "WeightsMapper":
     def apply_vllm_mapper(  # noqa: B027
         self, hf_to_vllm_mapper: "WeightsMapper"
     ):
-        """
-        Interface for models to update module names referenced in
+        """Interface for models to update module names referenced in
         quantization configs in order to reflect the vllm model structure
 
         Args:
             hf_to_vllm_mapper: maps from hf model structure (the assumed
                 structure of the qconfig) to vllm model structure
+
         """
         # TODO (@kylesayrs): add implementations for all subclasses
         pass
@@ -265,17 +265,17 @@ def apply_vllm_mapper(  # noqa: B027
     def maybe_update_config(  # noqa: B027
         self,
         model_name: str,
-        hf_config: PretrainedConfig | None = None,
+        hf_config: PreTrainedConfig | None = None,
         revision: str | None = None,
     ):
-        """
-        Interface to update values after config initialization.
+        """Interface to update values after config initialization.
 
         Args:
             model_name: The name of the model
             hf_config: The Hugging Face config of the model
             revision: The revision of the model
         Returns:
+
         """
         # TODO: revision is never passed currently in vllm.py,
         # but is used in subclasses, should we remove this parameter?
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
index 306a5400550e..a4cd5d4d1303 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
@@ -121,8 +121,7 @@ def get_name(self) -> QuantizationMethods:
         return "compressed-tensors"
 
     def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
-        """
-        Transform layer paths in config targets to match vLLM's naming.
+        """Transform layer paths in config targets to match vLLM's naming.
 
         The WeightsMapper is designed for weight paths, but some backends
         (e.g. transformers) use broad prefix mappings like "" -> "model."
@@ -214,8 +213,7 @@ def get_quant_method(
         return None
 
     def _add_fused_moe_to_target_scheme_map(self):  # XXXXXXXXXXXXXXXXXXXXXX
-        """
-        Helper function to update target_scheme_map
+        """Helper function to update target_scheme_map
         since linear layers get fused into RoutedExperts
         targeting 'Linear' needs to also match
         RoutedExperts modules.
@@ -269,8 +267,7 @@ def from_config(cls, config: dict[str, Any]) -> "CompressedTensorsConfig":
     def _parse_sparsity_config(
         cls, config: dict[str, Any]
     ) -> tuple[dict[str, SparsityCompressionConfig], list[str]]:
-        """
-        Args:
+        """Args:
             config: The `quantization_config` dictionary from config.json
 
         Returns:
@@ -278,6 +275,7 @@ def _parse_sparsity_config(
             1. A dictionary mapping target layer names to their corresponding
                 sparsity_config
             2. A list of layer names to ignore for sparsity
+
         """
         if not (sparsity_config := config.get(SPARSITY_CONFIG_NAME)):
             return dict(), []
@@ -301,13 +299,13 @@ def _parse_sparsity_config(
     def _quantization_scheme_map_from_config(
         cls, config: dict[str, Any]
     ) -> QUANTIZATION_SCHEME_MAP_TYPE:
-        """
-        Args:
+        """Args:
             config: The `quantization_config` dictionary from config.json
 
         Returns:
             A dictionary mapping target layer names to their corresponding
             quantization_args for weights and input activations
+
         """
         target_scheme_map: dict[str, Any] = dict()
         quant_format = cast(str, config.get("format"))
@@ -789,6 +787,7 @@ def _get_scheme_from_parts(
                 num_bits=weight_quant.num_bits,
                 strategy=weight_quant.strategy,
                 group_size=weight_quant.group_size,
+                symmetric=weight_quant.symmetric,
                 input_quant=input_quant,
                 layer_name=layer_name,
                 quant_format=format,
@@ -801,6 +800,7 @@ def _get_scheme_from_parts(
                 num_bits=weight_quant.num_bits,
                 strategy=weight_quant.strategy,
                 group_size=weight_quant.group_size,
+                symmetric=weight_quant.symmetric,
                 input_quant=input_quant,
                 layer_name=layer_name,
                 quant_format=format,
@@ -895,8 +895,7 @@ def _get_scheme_from_parts(
     def get_scheme(
         self, layer: torch.nn.Module, layer_name: str | None = None
     ) -> "CompressedTensorsScheme | None":
-        """
-        compressed-tensors supports non uniform in the following way:
+        """compressed-tensors supports non uniform in the following way:
 
         targets of config_groups: There can be N config_groups which each
             have a quantization scheme. Each config_group has a list of targets
@@ -907,7 +906,6 @@ def get_scheme(
         use the quantization scheme corresponding to the matched target
         to select the CompressedTensorsScheme used for inference.
         """
-
         # Use the new get_quant_args method to extract QuantizationArgs
         scheme_dict = self.get_scheme_dict(layer, layer_name)
 
@@ -944,8 +942,7 @@ def get_scheme(
     def get_scheme_dict(
         self, layer: torch.nn.Module, layer_name: str | None = None
     ) -> dict[str, QuantizationArgs | str | None] | None:
-        """
-        Extract the QuantizationArgs for a given layer.
+        """Extract the QuantizationArgs for a given layer.
 
         Returns:
             dict with {
@@ -953,6 +950,7 @@ def get_scheme_dict(
                 "input_activations": QuantizationArgs | None,
                 "format": str | None
             } | None
+
         """
         # TODO (@kylesayrs): support ignore module names with ct matching utils
         if should_ignore_layer(
@@ -1004,8 +1002,7 @@ def create_weights(
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        """
-        Use the CompressedTensorsScheme associated with each layer to create
+        """Use the CompressedTensorsScheme associated with each layer to create
         the necessary parameters for the layer. See LinearMethodBase for param
         details
         """
@@ -1026,8 +1023,7 @@ def apply(
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ):
-        """
-        Use the output of create_weights and the CompressedTensorsScheme
+        """Use the output of create_weights and the CompressedTensorsScheme
         associated with the layer to apply the forward pass with the
         layer input.  See LinearMethodBase for param details
 
@@ -1039,8 +1035,7 @@ def apply(
 
 
 class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
-    """
-    Supports loading kv-cache scaling factors from compressed-tensors
+    """Supports loading kv-cache scaling factors from compressed-tensors
     checkpoints.
     """
 
@@ -1050,12 +1045,12 @@ def __init__(self, quant_config: CompressedTensorsConfig):
 
     @staticmethod
     def validate_kv_cache_scheme(kv_cache_scheme: dict[str, Any] | None):
-        """
-        Validator for the kv cache scheme. Useful for controlling the
+        """Validator for the kv cache scheme. Useful for controlling the
         kv cache quantization schemes, that are being supported in vLLM
 
         Args:
             kv_cache_scheme: the compressed-tensors kv cache scheme
+
         """
         if kv_cache_scheme is None:
             return
@@ -1091,8 +1086,7 @@ def validate_kv_cache_scheme(kv_cache_scheme: dict[str, Any] | None):
             )
 
     def create_weights(self, layer: torch.nn.Module):
-        """
-        Initialize placeholder scales and zero points to enable loading of
+        """Initialize placeholder scales and zero points to enable loading of
         quantized params from compressed-tensors checkpoints.
         """
         strategy = None  # for backward compatibility
@@ -1199,8 +1193,7 @@ def _tp_aware_loader(
             )
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        """
-        Override the default vLLM placeholder scales with the llm-compressor loaded
+        """Override the default vLLM placeholder scales with the llm-compressor loaded
         scales. Zero points are not used as only symmetric quantization is supported.
         """
         layer._k_scale = layer.k_scale
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
index bcc930471da7..f2d2a9c9452c 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
@@ -88,14 +88,7 @@ def get_moe_method(
                     f"and bits: {weight_quant.num_bits}"
                 )
 
-            # Native ROCm HIP kernels (RDNA3, etc.)
             if current_platform.is_rocm():
-                from . import rocm_moe_rdna
-
-                if rocm_moe_rdna.is_supported(weight_quant):
-                    return rocm_moe_rdna.make_method(
-                        weight_quant, input_quant, layer.moe_config
-                    )
                 from vllm.platforms.rocm import on_gfx950
 
                 vllm_config = get_current_vllm_config()
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
index cf98e52bb61a..54f5aa86b34c 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
@@ -184,9 +184,7 @@ def create_weights(
         set_weight_attrs(w2_input_scale, extra_weight_attrs)
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
-        """
-        Convert NVFP4 MoE weights into kernel format and setup the kernel.
-        """
+        """Convert NVFP4 MoE weights into kernel format and setup the kernel."""
         # NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod
         # requires this naming convention. However, the name change breaks
         # reloading because the state dict no longer matches disk. Once we
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py
index 2d2190e614ed..0abae681f083 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py
@@ -36,8 +36,7 @@
 
 
 class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
-    """
-    CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform
+    """CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform
     - Weights: int4 (stored as int8 values in [-8,7], packed to uint8 nibbles)
     - Scales: Fp32 for Channelwise , bf16 for groupwise quantization
     - Bias: Same data type as original weights
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
index 6a3f73e5392d..6eb2455d0efd 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
@@ -11,7 +11,6 @@
 
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
-    FusedMoEExpertsModular,
     RoutedExperts,
     SharedExperts,
 )
@@ -36,7 +35,6 @@
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     check_moe_marlin_supports_config,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
@@ -143,8 +141,7 @@ def get_weight_shape(
         num_groups_w2: int | None = None,
         num_groups_w13: int | None = None,
     ) -> tuple[int, int, int]:
-        """
-        Get the shape of the weight based on the weight name, number of experts
+        """Get the shape of the weight based on the weight name, number of experts
         hidden size, intermediate size per partition, number of groups for w2,
         and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
         for weight scales/zero_points.
@@ -443,10 +440,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         replace_parameter(layer, "w13_weight_scale", w13_scales)
         replace_parameter(layer, "w2_weight_scale", w2_scales)
 
-        # CPU fused_experts_cpu requires zero points even for symmetric quant.
+        # CPU fused_experts_cpu and the RDNA3 HIP kernel require zero points
+        # even for symmetric quant (the oracle synthesizes them).
         # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None.
         if (
-            not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU
+            not self.symmetric
+            or self.wna16_backend in (WNA16MoEBackend.CPU, WNA16MoEBackend.RDNA3)
         ) and self.wna16_backend != WNA16MoEBackend.EMULATION:
             assert w13_qzeros is not None and w2_qzeros is not None
             replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
@@ -465,18 +464,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                     torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
                 )
 
-            # Marlin workspace — only needed for Marlin-family backends, not emulation.
-            if (
-                self.experts_cls is not None
-                and issubclass(self.experts_cls, FusedMoEExpertsModular)
-                and self.wna16_backend != WNA16MoEBackend.EMULATION
-            ):
-                layer.workspace = marlin_make_workspace_new(
-                    layer.w13_weight_packed.device,
-                    4,
-                    existing=getattr(layer, "workspace", None),
-                )
-
         # Alias packed weights to w13_weight/w2_weight for the modular kernel interface
         layer.w13_weight = layer.w13_weight_packed
         layer.w2_weight = layer.w2_weight_packed
@@ -487,7 +474,7 @@ def get_fused_moe_quant_config(
         self, layer: torch.nn.Module
     ) -> FusedMoEQuantConfig | None:
         if self.wna16_backend == WNA16MoEBackend.HUMMING:
-            from vllm.model_executor.layers.quantization.utils.humming_utils import (
+            from vllm.model_executor.layers.quantization.utils.humming import (
                 get_humming_moe_quant_config,
             )
 
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py
deleted file mode 100644
index 05db6ec3ddd8..000000000000
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py
+++ /dev/null
@@ -1,260 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""CompressedTensors MoE W4A16 using the fused RDNA3 (gfx1100) HIP kernel.
-
-Uses ``moe_gptq_gemm_rdna3`` — a single HIP kernel launch per GEMM that
-handles expert routing + W4A16 dequant + dot product with atomic output.
-
-Weight format (per expert, same as dense RDNA3 W4A16):
-  - Packed int32 ``[E, K/8, N]`` with exllama shuffle
-  - Scales ``[E, groups, N]`` in activation dtype
-  - Zero points ``[E, groups, N/8]`` packed int32 (synthesized)
-"""
-
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe import (
-    RoutedExperts,
-    SharedExperts,
-)
-from vllm.model_executor.layers.fused_moe.activation import (
-    MoEActivation,
-    apply_moe_activation,
-)
-from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
-    moe_align_block_size,
-)
-from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import (  # noqa: E501
-    CompressedTensorsWNA16MoEMethod,
-)
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    pack_quantized_values_into_int32,
-)
-from vllm.scalar_type import scalar_types
-
-logger = init_logger(__name__)
-
-
-def _synthesize_qzeros(
-    groups: int, out_features: int, device: torch.device
-) -> torch.Tensor:
-    """Create packed zero-point tensor for symmetric quant.
-
-    GPTQv1 +1 quirk: kernel adds 1 to stored zeros, so encode
-    (bias - 1) = 7 for uint4b8 (bias=8).
-    """
-    zeros = torch.full(
-        (groups, out_features),
-        scalar_types.uint4b8.bias - 1,
-        dtype=torch.int32,
-        device=device,
-    )
-    return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)
-
-
-class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod):
-    """W4A16 MoE using the fused RDNA3 HIP kernel (moe_gptq_gemm_rdna3).
-
-    Weights are in RDNA3 format (shuffled int32 [E, K/8, N]),
-    NOT Triton format (transposed uint8). apply() dispatches through
-    the fused HIP kernel directly.
-    """
-
-    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        device = layer.w13_weight_packed.device
-        num_experts = layer.w13_weight_packed.shape[0]
-
-        # Shuffle weights in-place per expert (exllama nibble interleave)
-        for e in range(num_experts):
-            w13_e = layer.w13_weight_packed.data[e].contiguous()
-            ops.gptq_shuffle(w13_e, 4)
-            layer.w13_weight_packed.data[e] = w13_e
-            w2_e = layer.w2_weight_packed.data[e].contiguous()
-            ops.gptq_shuffle(w2_e, 4)
-            layer.w2_weight_packed.data[e] = w2_e
-
-        # Keep scales as [E, groups, N] in activation dtype
-        act_dtype = layer.w13_weight_scale.dtype
-        layer.w13_weight_scale = torch.nn.Parameter(
-            layer.w13_weight_scale.to(dtype=act_dtype).contiguous(),
-            requires_grad=False,
-        )
-        layer.w2_weight_scale = torch.nn.Parameter(
-            layer.w2_weight_scale.to(dtype=act_dtype).contiguous(),
-            requires_grad=False,
-        )
-
-        # Synthesize packed zero points: [E, groups, N/8] int32
-        w13_groups = (layer.w13_weight_packed.shape[1] * 8) // self.group_size
-        w13_N = layer.w13_weight_packed.shape[2]
-        w2_groups = (layer.w2_weight_packed.shape[1] * 8) // self.group_size
-        w2_N = layer.w2_weight_packed.shape[2]
-
-        w13_qz = _synthesize_qzeros(w13_groups, w13_N, device)
-        w2_qz = _synthesize_qzeros(w2_groups, w2_N, device)
-        layer.w13_qzeros = torch.nn.Parameter(
-            w13_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
-            requires_grad=False,
-        )
-        layer.w2_qzeros = torch.nn.Parameter(
-            w2_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
-            requires_grad=False,
-        )
-
-        # Pre-allocate reusable buffers for decode (sizes based on top_k=8)
-        N_gate_up = w13_N
-        hidden_size = w2_N
-        intermediate = N_gate_up // 2  # gated activation
-        # Max tokens we expect in decode; prefill will re-allocate if needed
-        max_decode_tokens = 16
-        top_k = 8  # conservative default
-        buf_size = max_decode_tokens * top_k
-        layer.rdna3_w1_buf = torch.zeros(
-            buf_size, N_gate_up, dtype=act_dtype, device=device
-        )
-        layer.rdna3_act_buf = torch.empty(
-            buf_size, intermediate, dtype=act_dtype, device=device
-        )
-        layer.rdna3_out_buf = torch.zeros(
-            max_decode_tokens, hidden_size, dtype=act_dtype, device=device
-        )
-        layer.rdna3_empty_tw = torch.empty(0, device=device)
-
-    def apply(
-        self,
-        layer: RoutedExperts,
-        x: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        shared_experts: SharedExperts | None,
-        shared_experts_input: torch.Tensor | None,
-    ) -> torch.Tensor:
-        activation = (
-            layer.activation
-            if isinstance(layer.activation, MoEActivation)
-            else MoEActivation.from_str(layer.activation)
-        )
-        return _rdna3_fused_moe(
-            x,
-            topk_weights,
-            topk_ids,
-            layer=layer,
-            activation=activation,
-            apply_router_weight_on_input=(layer.apply_router_weight_on_input),
-            global_num_experts=layer.global_num_experts,
-            expert_map=layer.expert_map,
-        )
-
-
-def _rdna3_fused_moe(
-    hidden_states: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    layer: RoutedExperts,
-    activation: MoEActivation,
-    apply_router_weight_on_input: bool,
-    global_num_experts: int,
-    expert_map: torch.Tensor | None,
-) -> torch.Tensor:
-    """Fused MoE forward using the RDNA3 W4A16 HIP kernel.
-
-    Optimizations vs naive dispatch:
-      - BLOCK_SIZE_M=1 for decode (no padding waste, bf16 fast path)
-      - Pre-allocated buffers (no torch.zeros per call)
-      - Inline token sorting for small M (skip moe_align_block_size)
-      - moe_sum fused into output accumulation
-    """
-    num_tokens = hidden_states.shape[0]
-    top_k = topk_ids.shape[1]
-    total_tokens = num_tokens * top_k
-    N_gate_up = layer.w13_weight_packed.shape[2]
-    hidden_size = layer.w2_weight_packed.shape[2]
-    dtype = hidden_states.dtype
-    device = hidden_states.device
-
-    intermediate_size = N_gate_up // 2 if activation.is_gated else N_gate_up
-
-    if global_num_experts <= 0:
-        global_num_experts = layer.w13_weight_packed.shape[0]
-
-    # BLOCK_SIZE_M=1 for decode (small M), 4 for prefill
-    block_size_m = 1 if num_tokens <= 4 else 4
-
-    # --- Token routing ---
-    sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
-        topk_ids,
-        block_size_m,
-        global_num_experts,
-        expert_map,
-    )
-
-    # --- Reuse pre-allocated buffers when possible ---
-    if total_tokens <= layer.rdna3_w1_buf.shape[0]:
-        w1_out = layer.rdna3_w1_buf[:total_tokens]
-        w1_out.zero_()
-        act_out = layer.rdna3_act_buf[:total_tokens]
-    else:
-        w1_out = torch.zeros(
-            total_tokens,
-            N_gate_up,
-            dtype=dtype,
-            device=device,
-        )
-        act_out = torch.empty(
-            total_tokens,
-            intermediate_size,
-            dtype=dtype,
-            device=device,
-        )
-
-    # --- topk weights (pre-cast to float32 for kernel) ---
-    topk_w_float = topk_weights.view(-1).float()
-    empty_tw = layer.rdna3_empty_tw
-
-    # --- w1 GEMM: [M, K] -> [M*top_k, N_gate_up] ---
-    ops.moe_gptq_gemm_rdna3(
-        hidden_states,
-        w1_out,
-        layer.w13_weight_packed,
-        layer.w13_weight_scale,
-        layer.w13_qzeros,
-        topk_w_float if apply_router_weight_on_input else empty_tw,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        top_k,
-        block_size_m,
-        apply_router_weight_on_input,
-    )
-
-    # --- Activation (silu_and_mul etc.) ---
-    apply_moe_activation(activation, act_out, w1_out)
-
-    # --- w2 GEMM: [M*top_k, intermediate] -> [M, hidden] (fused reduce) ---
-    # output_topk=top_k: kernel writes to out[token_id / top_k] directly,
-    # fusing moe_sum into the atomic accumulation — saves one kernel launch
-    # and the w2_out intermediate buffer.
-    out = torch.zeros(
-        num_tokens,
-        hidden_size,
-        dtype=dtype,
-        device=device,
-    )
-    ops.moe_gptq_gemm_rdna3(
-        act_out,
-        out,
-        layer.w2_weight_packed,
-        layer.w2_weight_scale,
-        layer.w2_qzeros,
-        topk_w_float if not apply_router_weight_on_input else empty_tw,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        1,
-        block_size_m,
-        not apply_router_weight_on_input,
-        output_topk=top_k,
-    )
-    return out
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py
deleted file mode 100644
index e72caa0796a5..000000000000
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""ROCm MoE kernel dispatcher.
-
-Selects architecture-specific native HIP MoE kernels in priority order.
-Falls back to the Triton WNA16 path when no native kernel is available.
-"""
-
-import torch
-
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def is_supported(weight_quant) -> bool:
-    """Check if a native ROCm MoE kernel is available for this config."""
-    if weight_quant.num_bits != 4:
-        return False
-
-    from vllm.platforms.rocm import on_gfx1100
-
-    # RDNA3 (gfx1100). Future: add RDNA4 (gfx12x), CDNA (gfx94x), etc.
-    return (
-        on_gfx1100()
-        and hasattr(torch.ops, "_rocm_C")
-        and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
-    )
-
-
-def make_method(weight_quant, input_quant, moe_config):
-    """Create the native ROCm MoE method. Call only after is_supported()."""
-    from vllm.platforms.rocm import on_gfx1100
-
-    if on_gfx1100():
-        from .compressed_tensors_moe_wna16_rdna3 import (
-            CompressedTensorsWNA16RDNA3MoEMethod,
-        )
-
-        logger.info_once(
-            "Using CompressedTensorsWNA16RDNA3MoEMethod (native RDNA3 HIP kernel)"
-        )
-        return CompressedTensorsWNA16RDNA3MoEMethod(
-            weight_quant, input_quant, moe_config
-        )
-
-    # Future: RDNA4, CDNA, etc.
-    raise RuntimeError("is_supported() returned True but no kernel matched")
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py
index 78419a0dd98d..ff5173b180d2 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py
@@ -9,33 +9,26 @@
 
 
 class CompressedTensorsScheme(ABC):
-    """
-    Abstract class used to describe the weight creation and forward pass
+    """Abstract class used to describe the weight creation and forward pass
     of different quantization schemes supported by CompressedTensors.
     """
 
     @classmethod
     @abstractmethod
     def get_min_capability(cls) -> int:
-        """
-        Get minimum device capability.
-        """
+        """Get minimum device capability."""
         raise NotImplementedError()
 
     @abstractmethod
     def create_weights(self, *args, **kwargs):
-        """
-        Weight creation for the particular scheme. Inputs to this function
-
-        """
+        """Weight creation for the particular scheme. Inputs to this function."""
         raise NotImplementedError()
 
     @abstractmethod
     def apply_weights(
         self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
     ):
-        """
-        Run the forward pass for the particular scheme. This is where
+        """Run the forward pass for the particular scheme. This is where
         scheme-specific dequant/quant steps/kernels should be applied.
 
         Args:
@@ -43,13 +36,13 @@ def apply_weights(
                 other parameters relevant to the particular scheme.
             x: input to the layer
             bias: bias parameter
+
         """
         raise NotImplementedError()
 
     @abstractmethod
     def process_weights_after_loading(self, layer: torch.nn.Module):
-        """
-        Called after weight loading is complete for any cleanup that
+        """Called after weight loading is complete for any cleanup that
         needs to occur.
         """
         raise NotImplementedError()
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py
index 2eb8ba717bde..f3b02b00378a 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py
@@ -21,8 +21,7 @@
 
 
 class CompressedTensorsW4A4Mxfp4(CompressedTensorsScheme):
-    """
-    Compressed tensors scheme for MXFP4.
+    """Compressed tensors scheme for MXFP4.
 
     Supports models quantized with the compressed-tensors mxfp4-pack-quantized
     format.
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
index cb9cb34158c2..ce4520cbbe39 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
@@ -153,12 +153,10 @@ def process_weights_after_loading(self, layer) -> None:
                 layer.logical_widths,
                 getattr(layer, "input_scale", None),
             )
-            weight = weight.t()
         elif self.strategy == QuantizationStrategy.CHANNEL:
             weight, weight_scale, input_scale = process_fp8_weight_channel_strategy(
                 layer.weight, layer.weight_scale, getattr(layer, "input_scale", None)
             )
-            weight = weight.t()
 
         elif self.strategy == QuantizationStrategy.BLOCK:
             assert self.is_static_input_scheme is False
@@ -177,8 +175,7 @@ def process_weights_after_loading(self, layer) -> None:
 
         # required by torch.compile to be torch.nn.Parameter
         layer.weight = Parameter(weight.data, requires_grad=False)
-        # Preserve the dim tags dropped by the transpose so layout-aware
-        # kernels (humming) see (K, N) instead of assuming (N, K).
+        # The helpers return (K, N); retag so layout-aware kernels see it.
         layer.weight.input_dim = 0
         layer.weight.output_dim = 1
         layer.weight_scale = Parameter(weight_scale.data, requires_grad=False)
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py
index 5c511fc98d99..0f72be13a554 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py
@@ -22,8 +22,7 @@
 
 
 class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme):
-    """
-    Compressed tensors scheme for MXFP8 quantization (W8A8).
+    """Compressed tensors scheme for MXFP8 quantization (W8A8).
 
     Loads pre-quantized MXFP8 weights from compressed-tensors checkpoints.
     Activations are dynamically quantized to MXFP8 at runtime.
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
index 42688ab834a7..8cfce07544ad 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
@@ -42,7 +42,15 @@
     7: scalar_types.uint7b64,
     8: scalar_types.uint8b128,
 }
-WNA16_ZP_SUPPORTED_TYPES_MAP = {4: scalar_types.uint4, 8: scalar_types.uint8}
+WNA16_ZP_SUPPORTED_TYPES_MAP = {
+    2: scalar_types.uint2,
+    3: scalar_types.uint3,
+    4: scalar_types.uint4,
+    5: scalar_types.uint5,
+    6: scalar_types.uint6,
+    7: scalar_types.uint7,
+    8: scalar_types.uint8,
+}
 WNA16_SUPPORTED_BITS = list(WNA16_SUPPORTED_TYPES_MAP.keys())
 
 
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py
index 427c92310fd8..acd4a967018d 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py
@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Weight N-bit INT scheme with symmetric INT4 activation quant via Humming.
+"""Weight N-bit INT scheme with INT4 activation quant via Humming.
 
 Handles compressed-tensors pack-quantized INT weight checkpoints (2-8 bit)
 with INT4 symmetric dynamic per-token/per-group input activation
-quantization. Static, per-tensor, and asymmetric activation quantization
-are not supported.
+quantization. Supports both symmetric and asymmetric weight quantization.
+Static, per-tensor, and asymmetric activation quantization are not supported.
 """
 
 import math
@@ -28,6 +28,7 @@
 )
 from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import (  # noqa: E501
     WNA16_SUPPORTED_TYPES_MAP,
+    WNA16_ZP_SUPPORTED_TYPES_MAP,
 )
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     marlin_repeat_scales_on_all_ranks,
@@ -36,6 +37,7 @@
     BasevLLMParameter,
     ChannelQuantScaleParameter,
     GroupQuantScaleParameter,
+    PackedColumnParameter,
     PackedvLLMParameter,
 )
 
@@ -52,6 +54,7 @@ def __init__(
         num_bits: int,
         strategy: str,
         group_size: int | None = None,
+        symmetric: bool = True,
         input_quant: QuantizationArgs | None = None,
         layer_name: str | None = None,
         quant_format: str = "pack-quantized",
@@ -60,6 +63,7 @@ def __init__(
         self.pack_factor = Fraction(32, num_bits)
         self.strategy = strategy
         self.group_size = -1 if group_size is None else group_size
+        self.symmetric = symmetric
         self.input_quant = input_quant
         self.layer_name = layer_name
         self.quant_format = quant_format
@@ -69,7 +73,19 @@ def __init__(
                 f"Unsupported num_bits = {num_bits} for WNA4Int; "
                 f"supported = {sorted(WNA16_SUPPORTED_TYPES_MAP)}"
             )
-        self.quant_type = WNA16_SUPPORTED_TYPES_MAP[num_bits]
+
+        if not self.symmetric and num_bits not in WNA16_ZP_SUPPORTED_TYPES_MAP:
+            raise ValueError(
+                f"Asymmetric quantization not supported for "
+                f"num_bits = {num_bits}. Supported: "
+                f"{list(WNA16_ZP_SUPPORTED_TYPES_MAP)}"
+            )
+
+        self.quant_type = (
+            WNA16_ZP_SUPPORTED_TYPES_MAP[num_bits]
+            if not self.symmetric
+            else WNA16_SUPPORTED_TYPES_MAP[num_bits]
+        )
 
         if input_quant is not None:
             if not input_quant.symmetric:
@@ -143,7 +159,7 @@ def create_weights(
             weight_type=self.quant_type,
             act_type=params_dtype,
             group_size=self.group_size,
-            zero_points=False,
+            zero_points=not self.symmetric,
         )
 
         kernel_type = choose_mp_linear_kernel(mp_config)
@@ -151,12 +167,6 @@ def create_weights(
             logger.info("Using %s for CompressedTensorsWNA4Int", kernel_type.__name__)
             self._kernel_backends_being_used.add(kernel_type.__name__)
 
-        self.kernel = kernel_type(
-            mp_config,
-            w_q_param_name="weight_packed",
-            w_s_param_name="weight_scale",
-        )
-
         input_quant_config = self._build_input_quant_config()
         if input_quant_config is not None:
             layer._humming_input_quant_config = input_quant_config
@@ -167,10 +177,10 @@ def create_weights(
         partition_scales = not marlin_repeat_scales_on_all_ranks(
             self.group_size, row_parallel
         )
-        scales_size = input_size // group_size
+        scales_and_zp_size = input_size // group_size
         if partition_scales:
             assert input_size_per_partition % group_size == 0
-            scales_size = input_size_per_partition // group_size
+            scales_and_zp_size = input_size_per_partition // group_size
 
         packed_input_dim = math.ceil(input_size_per_partition * self.num_bits / 32)
         layer.register_parameter(
@@ -190,7 +200,7 @@ def create_weights(
         )
 
         scale_data = torch.empty(
-            output_size_per_partition, scales_size, dtype=params_dtype
+            output_size_per_partition, scales_and_zp_size, dtype=params_dtype
         )
         if partition_scales:
             weight_scale = GroupQuantScaleParameter(
@@ -207,6 +217,34 @@ def create_weights(
             )
         layer.register_parameter("weight_scale", weight_scale)
 
+        if not self.symmetric:
+            packed_output_dim = math.ceil(
+                output_size_per_partition * self.num_bits / 32
+            )
+            zeros_data = torch.zeros(
+                packed_output_dim,
+                scales_and_zp_size,
+                dtype=torch.int32,
+            )
+            if not partition_scales:
+                qzeros = PackedColumnParameter(
+                    output_dim=0,
+                    packed_dim=0,
+                    packed_factor=self.pack_factor,
+                    weight_loader=weight_loader,
+                    data=zeros_data,
+                )
+            else:
+                qzeros = PackedvLLMParameter(
+                    input_dim=1,
+                    output_dim=0,
+                    packed_dim=0,
+                    packed_factor=self.pack_factor,
+                    weight_loader=weight_loader,
+                    data=zeros_data,
+                )
+            layer.register_parameter("weight_zero_point", qzeros)
+
         layer.register_parameter(
             "weight_shape",
             BasevLLMParameter(
@@ -215,6 +253,13 @@ def create_weights(
             ),
         )
 
+        self.kernel = kernel_type(
+            mp_config,
+            w_q_param_name="weight_packed",
+            w_s_param_name="weight_scale",
+            w_zp_param_name="weight_zero_point",
+        )
+
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         self.kernel.process_weights_after_loading(layer)
 
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py
index a14077d899f3..156f29908dfb 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py
@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Weight N-bit INT scheme with symmetric INT8 activation quant via Humming.
+"""Weight N-bit INT scheme with INT8 activation quant via Humming.
 
 Handles compressed-tensors pack-quantized INT weight checkpoints (2-8 bit)
 with INT8 symmetric dynamic per-token/per-group input activation
-quantization. Static, per-tensor, and asymmetric activation quantization
-are not supported.
+quantization. Supports both symmetric and asymmetric weight quantization.
+Static, per-tensor, and asymmetric activation quantization are not supported.
 """
 
 import math
@@ -28,6 +28,7 @@
 )
 from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import (  # noqa: E501
     WNA16_SUPPORTED_TYPES_MAP,
+    WNA16_ZP_SUPPORTED_TYPES_MAP,
 )
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     marlin_repeat_scales_on_all_ranks,
@@ -36,6 +37,7 @@
     BasevLLMParameter,
     ChannelQuantScaleParameter,
     GroupQuantScaleParameter,
+    PackedColumnParameter,
     PackedvLLMParameter,
 )
 
@@ -52,6 +54,7 @@ def __init__(
         num_bits: int,
         strategy: str,
         group_size: int | None = None,
+        symmetric: bool = True,
         input_quant: QuantizationArgs | None = None,
         layer_name: str | None = None,
         quant_format: str = "pack-quantized",
@@ -60,6 +63,7 @@ def __init__(
         self.pack_factor = Fraction(32, num_bits)
         self.strategy = strategy
         self.group_size = -1 if group_size is None else group_size
+        self.symmetric = symmetric
         self.input_quant = input_quant
         self.layer_name = layer_name
         self.quant_format = quant_format
@@ -69,7 +73,19 @@ def __init__(
                 f"Unsupported num_bits = {num_bits} for WNA8Int; "
                 f"supported = {sorted(WNA16_SUPPORTED_TYPES_MAP)}"
             )
-        self.quant_type = WNA16_SUPPORTED_TYPES_MAP[num_bits]
+
+        if not self.symmetric and num_bits not in WNA16_ZP_SUPPORTED_TYPES_MAP:
+            raise ValueError(
+                f"Asymmetric quantization not supported for "
+                f"num_bits = {num_bits}. Supported: "
+                f"{list(WNA16_ZP_SUPPORTED_TYPES_MAP)}"
+            )
+
+        self.quant_type = (
+            WNA16_ZP_SUPPORTED_TYPES_MAP[num_bits]
+            if not self.symmetric
+            else WNA16_SUPPORTED_TYPES_MAP[num_bits]
+        )
 
         if input_quant is not None:
             if not input_quant.symmetric:
@@ -143,7 +159,7 @@ def create_weights(
             weight_type=self.quant_type,
             act_type=params_dtype,
             group_size=self.group_size,
-            zero_points=False,
+            zero_points=not self.symmetric,
         )
 
         kernel_type = choose_mp_linear_kernel(mp_config)
@@ -151,12 +167,6 @@ def create_weights(
             logger.info("Using %s for CompressedTensorsWNA8Int", kernel_type.__name__)
             self._kernel_backends_being_used.add(kernel_type.__name__)
 
-        self.kernel = kernel_type(
-            mp_config,
-            w_q_param_name="weight_packed",
-            w_s_param_name="weight_scale",
-        )
-
         input_quant_config = self._build_input_quant_config()
         if input_quant_config is not None:
             layer._humming_input_quant_config = input_quant_config
@@ -167,10 +177,10 @@ def create_weights(
         partition_scales = not marlin_repeat_scales_on_all_ranks(
             self.group_size, row_parallel
         )
-        scales_size = input_size // group_size
+        scales_and_zp_size = input_size // group_size
         if partition_scales:
             assert input_size_per_partition % group_size == 0
-            scales_size = input_size_per_partition // group_size
+            scales_and_zp_size = input_size_per_partition // group_size
 
         packed_input_dim = math.ceil(input_size_per_partition * self.num_bits / 32)
         layer.register_parameter(
@@ -190,7 +200,7 @@ def create_weights(
         )
 
         scale_data = torch.empty(
-            output_size_per_partition, scales_size, dtype=params_dtype
+            output_size_per_partition, scales_and_zp_size, dtype=params_dtype
         )
         if partition_scales:
             weight_scale = GroupQuantScaleParameter(
@@ -207,6 +217,34 @@ def create_weights(
             )
         layer.register_parameter("weight_scale", weight_scale)
 
+        if not self.symmetric:
+            packed_output_dim = math.ceil(
+                output_size_per_partition * self.num_bits / 32
+            )
+            zeros_data = torch.zeros(
+                packed_output_dim,
+                scales_and_zp_size,
+                dtype=torch.int32,
+            )
+            if not partition_scales:
+                qzeros = PackedColumnParameter(
+                    output_dim=0,
+                    packed_dim=0,
+                    packed_factor=self.pack_factor,
+                    weight_loader=weight_loader,
+                    data=zeros_data,
+                )
+            else:
+                qzeros = PackedvLLMParameter(
+                    input_dim=1,
+                    output_dim=0,
+                    packed_dim=0,
+                    packed_factor=self.pack_factor,
+                    weight_loader=weight_loader,
+                    data=zeros_data,
+                )
+            layer.register_parameter("weight_zero_point", qzeros)
+
         layer.register_parameter(
             "weight_shape",
             BasevLLMParameter(
@@ -215,6 +253,13 @@ def create_weights(
             ),
         )
 
+        self.kernel = kernel_type(
+            mp_config,
+            w_q_param_name="weight_packed",
+            w_s_param_name="weight_scale",
+            w_zp_param_name="weight_zero_point",
+        )
+
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         self.kernel.process_weights_after_loading(layer)
 
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py
index 0cde020627d2..ef4e9d288e7b 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py
@@ -29,8 +29,7 @@
 
 
 class CompressedTensorsLinearTransformMethod(LinearMethodBase):
-    """
-    Wraps `CompressedTensorsLinearMethod` or `UnquantizedLinearMethod` and adds
+    """Wraps `CompressedTensorsLinearMethod` or `UnquantizedLinearMethod` and adds
     input and output transforms to either side of the original apply method
     """
 
@@ -240,8 +239,7 @@ def get_schemes_args(
 def get_layer_partition_names(
     layer_name: str, packed_modules_mapping: dict[str, list[str]]
 ) -> list[str]:
-    """
-    Get all partition names associated with this layer.
+    """Get all partition names associated with this layer.
     Names are returned in order of their partition indices.
 
     ```python
@@ -251,7 +249,9 @@ def get_layer_partition_names(
         "gate_proj",
         "up_proj",
     ]
-    assert get_layer_partition_names("mlp.down_proj", mapping) == ["down_proj"]"""
+    assert get_layer_partition_names("mlp.down_proj", mapping) == ["down_proj"]
+    ```
+    """
     for fused_suffix, part_suffixes in packed_modules_mapping.items():
         if layer_name.endswith(fused_suffix):
             return [
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py
index 755059349382..8552eeb65250 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py
@@ -23,8 +23,7 @@
 
 
 class HadamardTransform(torch.nn.Module):
-    """
-    Class which handles weight loading, postprocessing, and application of
+    """Class which handles weight loading, postprocessing, and application of
     transforms. Meant to be used with `CompressedTensorsLinearTransformMethod`
     and attention transforms method (not implemented yet)
     """
@@ -96,13 +95,13 @@ def forward(self, value: Tensor, part_id: int = 0) -> Tensor:
             if self.transforms[part_id].scheme.head_dim is not None:
                 weight_size = self.transforms[part_id].scheme.head_dim
                 value = value.unflatten(-1, (-1, weight_size))
-                value = ops.hadacore_transform(value)
+                value = ops.hadacore_transform(value, inplace=False)
                 value = value.flatten(-2, -1)
 
                 return value
 
             # sylvester transforms are symmetric, inv => transpose => original
-            return ops.hadacore_transform(value)
+            return ops.hadacore_transform(value, inplace=False)
 
         # fall back to dense
         else:
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py
index b8b771e3c7a9..70e0c652541f 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py
@@ -76,8 +76,7 @@ def create_weights(
 
     @staticmethod
     def _get_flashinfer_gemm_backend(kernel: NvFp4LinearKernel) -> str:
-        """
-        Given a kernel, find the string that is needed to be passed into
+        """Given a kernel, find the string that is needed to be passed into
         `flashinfer_scaled_fp4_mm`, using
         vllm.model_executor.kernels.linear._LINEAR_BACKEND_KERNEL_MAP as source of truth
         """
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py
index 188ee8c4c97d..bc999c7258b5 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py
@@ -76,8 +76,7 @@ def find_matched_target(
     targets: Iterable[str],
     fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
 ) -> str | None:
-    """
-    Helper function to look up which "target" in the compressed-tensors
+    """Helper function to look up which "target" in the compressed-tensors
     config that a layer corresponds to.
 
     Recall that a compressed-tensors configs has a concept of
@@ -98,8 +97,8 @@ def find_matched_target(
         module: torch.nn.Module
         targets: list of targets to match the layer against
         fused_mapping: map from fused layer names to its components
-    """
 
+    """
     if layer_name is None:
         layer_name = ""
 
@@ -115,8 +114,7 @@ def find_matched_target(
 def _find_first_match(
     value: str, targets: Iterable[str], check_contains: bool = False
 ) -> str | None:
-    """
-    Returns first element of target that matches value either
+    """Returns first element of target that matches value either
     exactly or as a regex after 're:'. If check_contains is set to True,
     additionally checks if the target string is contained within the value.
 
@@ -124,8 +122,8 @@ def _find_first_match(
         value: string to compare the list of targets against
         targets: list of targets to match the layer against
         check_contains: whether or not to do a substring match
-    """
 
+    """
     for target in targets:
         if is_equal_or_regex_match(value, target, check_contains=check_contains):
             return target
@@ -137,8 +135,7 @@ def _match_fused_layer(
     target_layers: Iterable[str],
     fused_mapping: Mapping[str, list[str]],
 ) -> str | None:
-    """
-    Match a fused layer name to its corresponding individual layer in
+    """Match a fused layer name to its corresponding individual layer in
     target_layers. Returns first value in fused_mapping which matches targets
 
     Implements an "all" matching strategy where a fused layer matches iff
@@ -154,6 +151,7 @@ def _match_fused_layer(
         target_layers = ["model.layers.0.self_attn.q_proj",
                         "model.layers.0.self_attn.k_proj",
                         "model.layers.0.self_attn.v_proj"]
+
     """
     # find layer_name in mapping
     fused = next((key for key in fused_mapping if layer_name.endswith(key)), None)
diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index d8bef7347df0..47df735e89df 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -8,7 +8,6 @@
 import vllm.envs as envs
 from vllm.config import get_current_vllm_config
 from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import (
     init_fp8_linear_kernel,
 )
@@ -24,14 +23,12 @@
     SharedExperts,
     UnquantizedFusedMoEMethod,
 )
-from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEQuantConfig,
-)
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
     convert_to_fp8_moe_kernel_format,
     make_fp8_moe_kernel,
     make_fp8_moe_quant_config,
-    refine_fp8_moe_block_shape,
+    resolve_fp8_moe_weight_block_shape,
     select_fp8_moe_backend,
 )
 from vllm.model_executor.layers.linear import (
@@ -90,8 +87,6 @@
 
 ACTIVATION_SCHEMES = ["static", "dynamic"]
 
-logger = init_logger(__name__)
-
 
 class Fp8Config(QuantizationConfig):
     """Config class for FP8."""
@@ -251,6 +246,7 @@ class Fp8LinearMethod(LinearMethodBase):
 
     Args:
         quant_config: The quantization config.
+
     """
 
     supports_pre_processed_weights = True
@@ -420,7 +416,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             if self.act_q_static:
                 assert input_scale is not None
                 input_scale = input_scale.max()
-            weight = weight.t()
 
             # Update layer with new values.
             replace_parameter(layer, "weight", weight.data)
@@ -490,6 +485,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
 
     Args:
         quant_config: The quantization config.
+
     """
 
     supports_pre_processed_weights = True
@@ -509,32 +505,16 @@ def __init__(self, quant_config: Fp8Config, layer: RoutedExperts):
         # Set weight key and activation key for kernel compatibility
         if self.block_quant:
             assert self.weight_block_size is not None
-            # TP shards the intermediate dim of the expert weights, so a
-            # per-shard size that is not a multiple of the checkpoint's block
-            # size makes the checkpoint's block scales impossible to shard
-            # exactly. When a finer block size (>= 32) divides both the
-            # checkpoint blocks and all involved dims, the weight scales are
-            # refined to that granularity at load time (a lossless upsampling,
-            # since the refined block divides the checkpoint block). The
-            # refined block shape is encoded in the weight key, so the oracle
-            # only selects kernels that support it (e.g. Triton, which takes
-            # the block shape as a runtime argument).
-            refined_shape = refine_fp8_moe_block_shape(self.moe, self.weight_block_size)
-            if refined_shape is not None:
-                block_n, block_k = self.weight_block_size
-                self.weight_scale_refine = (
-                    block_n // refined_shape[0],
-                    block_k // refined_shape[1],
-                )
-                self.moe_block_shape = refined_shape
-                logger.info_once(
-                    "FP8 MoE block scales refined from %s to %s to fit "
-                    "the TP-sharded intermediate size %d.",
-                    str(self.weight_block_size),
-                    str(refined_shape),
-                    self.moe.intermediate_size_per_partition,
+            # Adapt the checkpoint block shape to TP sharding before
+            # building the weight key.
+            self.moe_block_shape, self.weight_scale_refine = (
+                resolve_fp8_moe_weight_block_shape(
+                    self.moe,
+                    self.weight_block_size,
+                    kFp8Dynamic128Sym,
+                    self.quant_config.is_checkpoint_fp8_serialized,
                 )
-            assert self.moe_block_shape is not None
+            )
             weight_key = create_fp8_quant_key(
                 static=True, group_shape=GroupShape(*self.moe_block_shape)
             )
@@ -881,9 +861,7 @@ def apply(
 
 
 class Fp8KVCacheMethod(BaseKVCacheMethod):
-    """
-    Supports loading kv-cache scaling factors from FP8 checkpoints.
-    """
+    """Supports loading kv-cache scaling factors from FP8 checkpoints."""
 
     def __init__(self, quant_config: Fp8Config):
         super().__init__(quant_config)
diff --git a/vllm/model_executor/layers/quantization/fp_quant.py b/vllm/model_executor/layers/quantization/fp_quant.py
index c24706252abd..8715db16b74a 100644
--- a/vllm/model_executor/layers/quantization/fp_quant.py
+++ b/vllm/model_executor/layers/quantization/fp_quant.py
@@ -98,6 +98,7 @@ class FPQuantLinearMethod(LinearMethodBase):
 
     Args:
         quant_config: The FPQuant quantization config.
+
     """
 
     def __init__(self, quant_config: FPQuantConfig):
diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py
index ea00c0f3c408..6bb0faa7c833 100644
--- a/vllm/model_executor/layers/quantization/humming.py
+++ b/vllm/model_executor/layers/quantization/humming.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Callable
+from math import gcd
 from typing import TYPE_CHECKING, Any
 
 import regex as re
@@ -28,12 +29,15 @@
     QuantizationConfig,
     QuantizeMethodBase,
 )
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
+from vllm.model_executor.layers.quantization.utils.humming import (
+    check_and_fallback_input_schema,
     convert_to_humming_moe_kernel_format,
     get_humming_linear_compute_config,
     get_humming_moe_quant_config,
+    humming_update_schema_hadamard_block_size,
     input_schema_to_quant_key,
     make_humming_moe_kernel,
+    resolve_humming_layer_config,
     select_humming_moe_experts,
     weight_schema_to_quant_key,
 )
@@ -227,23 +231,13 @@ def get_layer_weight_schema(self, config: dict[str, Any], prefix: str):
                 return None
             config = group_config
 
-        layer_config = config
-        layer_dynamic = config.get("dynamic", {})
-        if not isinstance(layer_dynamic, dict):
-            layer_dynamic = {}
-        for regex, override_config in layer_dynamic.items():
-            if regex[:1] != "+":
-                continue
-            if re.match(regex[2:], prefix):
-                layer_config = config.copy()
-                layer_config.update(override_config)
-                break
+        layer_config = resolve_humming_layer_config(config, prefix)
 
         if "quant_method" in layer_config:
             return _hm.BaseWeightSchema.from_config(layer_config)
         return None
 
-    def get_layer_input_schema(self, config: dict[str, Any], prefix: str):
+    def get_layer_input_config(self, config: dict[str, Any], prefix: str):
         if self.is_layer_skipped(config, prefix):
             return None
         if config["quant_method"] in ["compressed-tensors", "modelopt"]:
@@ -252,10 +246,15 @@ def get_layer_input_schema(self, config: dict[str, Any], prefix: str):
                 return None
             config = group_config
 
+        config = resolve_humming_layer_config(config, prefix)
         if config.get("quant_method", None) in _hm.BaseInputSchema.INPUT_SCHEMA_MAP:
-            return _hm.BaseInputSchema.from_config(config)
+            return config
         return None
 
+    def get_layer_input_schema(self, config: dict[str, Any], prefix: str):
+        config = self.get_layer_input_config(config, prefix)
+        return None if config is None else _hm.BaseInputSchema.from_config(config)
+
     def get_quant_config_for_layer(
         self, prefix: str, layer_type: str
     ) -> "HummingLayerQuantizationConfig | None":
@@ -279,14 +278,25 @@ def get_quant_config_for_layer(
         if weight_schema is not None:
             input_schema = None
             force_input_schema = None
+            allow_input_schema_fallback = True
 
             if self.full_config:
-                input_schema = self.get_layer_input_schema(self.full_config, prefix)
+                input_config = self.get_layer_input_config(self.full_config, prefix)
+                if input_config is not None:
+                    allow_input_schema_fallback = input_config.pop(
+                        "allow_fallback", True
+                    )
+                    input_schema = _hm.BaseInputSchema.from_config(input_config)
 
             if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG:
                 quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG.copy()
                 quant_config["quant_method"] = "humming"
-                force_input_schema = self.get_layer_input_schema(quant_config, prefix)
+                input_config = self.get_layer_input_config(quant_config, prefix)
+                if input_config is not None:
+                    allow_input_schema_fallback = input_config.pop(
+                        "allow_fallback", False
+                    )
+                    force_input_schema = _hm.BaseInputSchema.from_config(input_config)
                 if input_schema is None:
                     input_schema = force_input_schema
 
@@ -299,6 +309,7 @@ def get_quant_config_for_layer(
                 force_weight_schema=force_weight_schema,
                 force_input_schema=force_input_schema,
                 is_online_quant=is_online_quant,
+                allow_input_schema_fallback=allow_input_schema_fallback,
             )
         return None
 
@@ -318,6 +329,11 @@ def get_quant_method(
             elif isinstance(layer, LinearBase):
                 return UnquantizedLinearMethod()
         elif isinstance(layer, LinearBase):
+            input_size = layer.input_size
+            partition_size = getattr(layer, "input_size_per_partition", input_size)
+            is_online_quant = quant_config.is_online_quant
+            if is_online_quant and (input_size % 32 != 0 or partition_size % 32 != 0):
+                return UnquantizedLinearMethod()
             return HummingLinearMethod(quant_config)
         elif isinstance(layer, RoutedExperts):
             return HummingMoEMethod(quant_config, layer.moe_config)
@@ -332,6 +348,7 @@ def __init__(
         force_weight_schema: "HummingWeightSchema | None" = None,
         force_input_schema: "HummingInputSchema | None" = None,
         is_online_quant: bool = False,
+        allow_input_schema_fallback: bool = True,
     ):
         self.weight_schema = weight_schema
         if input_schema is None:
@@ -340,6 +357,7 @@ def __init__(
         self.force_weight_schema = force_weight_schema
         self.force_input_schema = force_input_schema
         self.is_online_quant = is_online_quant
+        self.allow_input_schema_fallback = allow_input_schema_fallback
 
     @classmethod
     def from_config(cls, config):
@@ -373,21 +391,13 @@ def new_weight_loader(
             if is_unquantized and self.is_online_quant:
                 # online quant (fp16/bf16 -> quant_type)
                 assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
-                f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype)
-                has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
-                tensor_list = _hm.quantize_weight(
-                    weight=loaded_weight,
-                    dtype=self.weight_schema.b_dtype,
-                    scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
-                    group_size=self.weight_schema.weight_scale_group_size,
-                    has_zero_point=self.weight_schema.has_zero_point,
-                    has_global_scale=has_global_scale,
-                    is_fp_zero_point=self.weight_schema.is_fp_zero_point,
-                    pack=True,
+                tensors = self.weight_schema.quant_tensor(
+                    loaded_weight.to(param.device),
+                    self.weight_schema,
+                    layer.param_dtype,
                 )
 
-                key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
-                for key, tensor in zip(key_list, tensor_list):
+                for key, tensor in tensors.items():
                     if tensor is None or tensor.nelement() == 0:
                         continue
                     param = getattr(layer, key)
@@ -459,6 +469,18 @@ def create_weights(
         layer.output_partition_sizes = output_partition_sizes
         layer.extra_weight_attrs = extra_weight_attrs.copy()
 
+        input_schema = self.force_input_schema or self.input_schema
+        input_schema = input_schema.to_humming_schema(params_dtype)
+        for name in ("weight_schema", "force_weight_schema"):
+            schema = getattr(self, name)
+            if getattr(schema, "hadamard_block_size", None) == -1:
+                schema = humming_update_schema_hadamard_block_size(
+                    weight_schema=schema,
+                    input_schema=input_schema,
+                    shape_k=input_size_per_partition,
+                )
+                setattr(self, name, schema)
+
         weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
         new_weight_loader = self.prepare_weight_loader(layer, weight_loader)
         extra_weight_attrs["weight_loader"] = new_weight_loader
@@ -502,7 +524,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             return None
 
         # convert from checkpoint format to humming format
-        if not isinstance(self.weight_schema, _hm.HummingWeightSchema):
+        is_humming_weight = isinstance(self.weight_schema, _hm.HummingWeightSchema)
+        is_humming_input = isinstance(self.input_schema, _hm.HummingInputSchema)
+
+        if not is_humming_weight or not is_humming_input:
             self.weight_schema, tensors = self.weight_schema.convert_humming(
                 tensors=layer.state_dict(),
                 shape_n_stacks=layer.output_partition_sizes,
@@ -510,12 +535,13 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                 param_dtype=layer.param_dtype,
             )
 
-            self.input_schema, _ = self.input_schema.convert_humming(
+            self.input_schema, input_tensors = self.input_schema.convert_humming(
                 tensors=layer.state_dict(),
                 shape_n_stacks=layer.output_partition_sizes,
                 shape_k_stacks=[layer.input_size_per_partition],
                 param_dtype=layer.param_dtype,
             )
+            tensors.update(input_tensors)
 
             for name, _ in list(layer.named_parameters()):
                 delattr(layer, name)
@@ -530,11 +556,15 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
         force_requant = self.force_weight_schema is not None
         if force_requant and self.weight_schema != self.force_weight_schema:
+            source_tensors = layer.state_dict()
             tensors = self.weight_schema.requant_tensors(
-                tensors=layer.state_dict(),
+                tensors=source_tensors,
                 target_weight_schema=self.force_weight_schema,
                 param_dtype=layer.param_dtype,
             )
+            name = self.input_schema.static_tensor_scale_name
+            if name is not None:
+                tensors[name] = source_tensors[name]
 
             self.weight_schema = self.force_weight_schema
 
@@ -548,6 +578,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
 
             del tensors
 
+        self.input_schema = check_and_fallback_input_schema(
+            weight_schema=self.weight_schema,
+            input_schema=self.input_schema,
+            param_dtype=layer.param_dtype,
+            allow_fallback=self.quant_config.allow_input_schema_fallback,
+        )
         self.layer_config = _hm.prepare_layer_config(
             shape_n=layer.output_partition_sizes_sum,
             shape_k=layer.input_size_per_partition,
@@ -557,10 +593,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             pad_k_to_multiple=128,
             has_bias=layer.has_bias,
             torch_dtype=layer.param_dtype,
+            device=layer.weight.device,
         )
-        tensors = _hm.transform_humming_tensors(
-            self.layer_config, dict(layer.named_parameters())
-        )
+        source_tensors = dict(layer.named_parameters())
+        tensors = _hm.transform_humming_tensors(self.layer_config, source_tensors)
+        name = self.input_schema.static_tensor_scale_name
+        if name is not None:
+            tensors[name] = source_tensors[name]
+        if "bias" in tensors:
+            zero_bias = torch.zeros_like(tensors["bias"])
+            layer.register_buffer("zero_bias", zero_bias)
         for name, _ in list(layer.named_parameters()):
             delattr(layer, name)
         for name, tensor in tensors.items():
@@ -583,8 +625,11 @@ def apply(
             weight=layer.weight,
             weight_scale=getattr(layer, "weight_scale", None),
             zero_point=getattr(layer, "zero_point", None),
-            bias=getattr(layer, "bias", None),
+            bias=getattr(layer, "bias" if bias is not None else "zero_bias", None),
             weight_scale_2=getattr(layer, "weight_scale_2", None),
+            input_scale=getattr(layer, "input_scale", None),
+            input_scale_2=getattr(layer, "input_scale_2", None),
+            hadamard_block_size=self.weight_schema.hadamard_block_size,
             locks=self.locks,
             compute_config=self.compute_config,
         )
@@ -606,11 +651,15 @@ def __init__(
         # Derive QuantKeys from humming schemas.
         # Prefer force schemas (the final format after requant) over base.
         weight_key = weight_schema_to_quant_key(
-            self.force_weight_schema or self.weight_schema
+            self.force_weight_schema or self.weight_schema, moe.in_dtype
         )
-        activation_key = input_schema_to_quant_key(
-            self.force_input_schema or self.input_schema
+        runtime_input_schema = check_and_fallback_input_schema(
+            weight_schema=self.force_weight_schema or self.weight_schema,
+            input_schema=self.force_input_schema or self.input_schema,
+            param_dtype=moe.in_dtype,
+            allow_fallback=quant_config.allow_input_schema_fallback,
         )
+        activation_key = input_schema_to_quant_key(runtime_input_schema, moe.in_dtype)
 
         # Select Humming MoE experts
         self.experts_cls = select_humming_moe_experts(
@@ -634,22 +683,14 @@ def new_weight_loader(
             # online quant (fp16/bf16 -> quant_type)
             if is_unquantized:
                 assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
-                f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype)
-                has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
-                tensor_list = _hm.quantize_weight(
-                    weight=loaded_weight,
-                    dtype=self.weight_schema.b_dtype,
-                    scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
-                    group_size=self.weight_schema.weight_scale_group_size,
-                    has_zero_point=self.weight_schema.has_zero_point,
-                    has_global_scale=has_global_scale,
-                    is_fp_zero_point=self.weight_schema.is_fp_zero_point,
-                    pack=True,
+                tensors = self.weight_schema.quant_tensor(
+                    loaded_weight.to(param.device),
+                    self.weight_schema,
+                    layer.param_dtype,
                 )
 
-                key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
                 success = True
-                for key, tensor in zip(key_list, tensor_list):
+                for key, tensor in tensors.items():
                     if tensor is None or tensor.nelement() == 0:
                         continue
                     sublayer_name = "w2" if shard_id == "w2" else "w13"
@@ -695,6 +736,18 @@ def create_weights(
         layer.num_experts = num_experts
         layer.param_dtype = params_dtype
         layer.intermediate_size = intermediate_size_per_partition
+        input_schema = self.force_input_schema or self.input_schema
+        input_schema = input_schema.to_humming_schema(params_dtype)
+        shape_k = gcd(hidden_size, intermediate_size_per_partition)
+        for name in ("weight_schema", "force_weight_schema"):
+            schema = getattr(self, name)
+            if getattr(schema, "hadamard_block_size", None) == -1:
+                schema = humming_update_schema_hadamard_block_size(
+                    weight_schema=schema,
+                    input_schema=input_schema,
+                    shape_k=shape_k,
+                )
+                setattr(self, name, schema)
         weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
         weight_loader = self.prepare_weight_loader(layer, weight_loader)
         extra_weight_attrs["weight_loader"] = weight_loader
@@ -761,6 +814,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
             weight_schema=self.weight_schema,
             input_schema=self.input_schema,
             force_weight_schema=self.force_weight_schema,
+            allow_input_schema_fallback=self.quant_config.allow_input_schema_fallback,
         )
 
         # Build the MoE kernel
@@ -783,8 +837,7 @@ def apply(
         shared_experts: SharedExperts | None,
         shared_experts_input: torch.Tensor | None,
     ) -> torch.Tensor:
-        """
-        Apply Humming-quantized MoE computation using the standard kernel flow.
+        """Apply Humming-quantized MoE computation using the standard kernel flow.
 
         This method uses FusedMoEKernel.apply() which orchestrates:
         1. Preparation (quantization if needed - skipped for Humming via
diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py
index 2d9cf624aba7..161baf446d12 100644
--- a/vllm/model_executor/layers/quantization/input_quant_fp8.py
+++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py
@@ -28,8 +28,7 @@
 # --8<-- [start:quant_fp8]
 @CustomOp.register("quant_fp8")
 class QuantFP8(CustomOp):
-    """
-    Quantize input tensor to FP8 (per-tensor, per-token, per-channel, or per-group).
+    """Quantize input tensor to FP8 (per-tensor, per-token, per-channel, or per-group).
     This CustomOp supports both static and dynamic quantization.
     """
 
@@ -45,18 +44,18 @@ def __init__(
         use_ue8m0: bool | None = None,  # for Torch compile
         compile_native: bool = True,
     ):
-        """
-        Args:
-            static: static or dynamic quantization
-            group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
-                PER_CHANNEL, or arbitrary block size)
-            num_token_padding: Pad the token dimension of output to this
-                size
-            tma_aligned_scales: For group quantization, output scales in
-                TMA-aligned layout
-            column_major_scales: For group quantization, output scales in
-                column major format
-            compile_native: Manually compile forward_native if compile mode > None
+        """Args:
+        static: static or dynamic quantization
+        group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
+            PER_CHANNEL, or arbitrary block size)
+        num_token_padding: Pad the token dimension of output to this
+            size
+        tma_aligned_scales: For group quantization, output scales in
+            TMA-aligned layout
+        column_major_scales: For group quantization, output scales in
+            column major format
+        compile_native: Manually compile forward_native if compile mode > None
+
         """
         super().__init__(compile_native=compile_native)
         self.static = static
@@ -156,8 +155,11 @@ def forward_hip(
         if use_aiter_per_token_quant:
             return rocm_aiter_ops.per_token_quant(x, _FP8_DTYPE, scale)
 
-        # Fallback to CUDA implementation
-        return self.forward_cuda(x, scale, scale_ub)
+        # Fall back to the CUDA implementation. Dispatch on the class, not
+        # through self: a subclass may override forward_cuda with a different
+        # signature (_DecodeConcatQuantFP8 does), and a virtual call would
+        # re-enter it with mismatched arguments.
+        return QuantFP8.forward_cuda(self, x, scale, scale_ub)
 
     def forward_xpu(
         self,
@@ -176,7 +178,8 @@ def forward_xpu(
                 dtype=_FP8_DTYPE,
                 use_ue8m0=self.use_ue8m0,
             )
-        return self.forward_cuda(x, scale, scale_ub, use_triton)
+        # Dispatch on the class, for the same reason as in forward_hip.
+        return QuantFP8.forward_cuda(self, x, scale, scale_ub, use_triton)
 
     def forward_native(
         self,
diff --git a/vllm/model_executor/layers/quantization/kv_cache.py b/vllm/model_executor/layers/quantization/kv_cache.py
index 447cfb955a83..6d6a8c5ce51b 100644
--- a/vllm/model_executor/layers/quantization/kv_cache.py
+++ b/vllm/model_executor/layers/quantization/kv_cache.py
@@ -41,8 +41,7 @@ def weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> Non
 
 
 class BaseKVCacheMethod(QuantizeMethodBase):
-    """
-    Quant method that adds `_k_scale` and `_v_scale` attributes to the
+    """Quant method that adds `_k_scale` and `_v_scale` attributes to the
     Attention layer to support loading those scaling factors from checkpoints.
     The k/v_scale will be used to:
         - quantize k/v_cache entries before saving them to the cache
@@ -50,6 +49,7 @@ class BaseKVCacheMethod(QuantizeMethodBase):
 
     Args:
         quant_config: the appropriate QuantizationConfig
+
     """
 
     supports_pre_processed_weights = True
@@ -58,8 +58,7 @@ def __init__(self, quant_config: QuantizationConfig):
         self.quant_config = quant_config
 
     def create_weights(self, layer: torch.nn.Module):
-        """
-        Create "weight" (aka q_scale, k_scale and v_scale)
+        """Create "weight" (aka q_scale, k_scale and v_scale)
         for an attention layer.
         """
         # Initialize the Q and KV cache scales to -1.0, an invalid value.
diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py
index 48482acd75bb..a9ac30715928 100644
--- a/vllm/model_executor/layers/quantization/modelopt.py
+++ b/vllm/model_executor/layers/quantization/modelopt.py
@@ -153,9 +153,7 @@ def algos_owned_by(config_name: str) -> tuple[str, ...]:
 
 
 class ModelOptKVCacheMethod(BaseKVCacheMethod):
-    """
-    Supports loading kv-cache scaling factors from FP8 or NVFP4 checkpoints.
-    """
+    """Supports loading kv-cache scaling factors from FP8 or NVFP4 checkpoints."""
 
     def __init__(self, quant_config: "ModelOptQuantConfigBase"):
         super().__init__(quant_config)
@@ -178,8 +176,7 @@ def __init__(
         self.exclude_modules: list[str] = exclude_modules
 
     def is_layer_excluded(self, prefix: str) -> bool:
-        """
-        Check if a layer should be excluded from quantization.
+        """Check if a layer should be excluded from quantization.
 
         Handles both exact matching (for fused layers) and ModelOpt wildcard matching.
 
@@ -472,8 +469,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
     """MoE method for ModelOpt FP8.
     Supports loading FP8 checkpoints with static weight scale and
     activation scale.
+
     Args:
         quant_config: The ModelOpt quantization config.
+
     """
 
     def __init__(
@@ -818,10 +817,11 @@ def _from_config(
 
 
 class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
-    """
-    MoE Method for FP4 Quantization.
+    """MoE Method for FP4 Quantization.
+
     Args:
         quant_config: NVFP4 Quant Config
+
     """
 
     supports_pre_processed_weights = True
@@ -851,9 +851,7 @@ def __init__(
         )
 
     def uses_weight_scale_2_pattern(self) -> bool:
-        """
-        FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales.
-        """
+        """FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales."""
         return True
 
     def create_weights(
@@ -973,9 +971,7 @@ def create_weights(
         layer.register_parameter("w2_input_scale", w2_input_scale)
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
-        """
-        Convert NVFP4 MoE weights into kernel format and setup the kernel.
-        """
+        """Convert NVFP4 MoE weights into kernel format and setup the kernel."""
         if is_weights_pre_processed():
             if self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM:
                 raise RuntimeError(
@@ -1942,7 +1938,9 @@ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
             PerTensorScaleParameter,
             wl,
         )
-        # Per-block (group_size) weight scale.
+        # Per-block (group_size) weight scale. Initialized to NaN (rather than
+        # left as torch.empty garbage) so an unloaded scale is unambiguously
+        # detectable in process() below, matching the sentinel pattern in #45320.
         self.register_params(
             layer,
             "weight_scale",
@@ -1955,11 +1953,26 @@ def create_weights(self, layer, role, ctx, shapes, wl) -> None:
             wl,
             input_dim=1,
             output_dim=0,
+            init=float("nan"),
         )
 
     def process(self, layer, role) -> None:
         if role is not WEIGHT:
             self.reject(role)
+        # Sanity-check: weight_scale must have been overwritten by the weight
+        # loader. A value still equal to the NaN sentinel it was created with
+        # means the FP4 weights were never actually loaded (e.g. the
+        # checkpoint stores this layer as BF16 and the weight loader silently
+        # skipped it).
+        if torch.isnan(layer.weight_scale.float()).any():
+            raise RuntimeError(
+                f"NVFP4 weight_scale for layer "
+                f"{getattr(layer, 'name', repr(layer))!r} was never loaded "
+                "(still NaN). The checkpoint likely stores this layer as "
+                "BF16 (not FP4). Fix: pass quant_config=None when "
+                "constructing this layer, or add it to the quantization "
+                "ignore list."
+            )
         if torch.unique(layer.weight_scale_2).numel() != 1:
             logger.warning_once(
                 "In NVFP4 linear, the global weight scale differs across "
@@ -2116,7 +2129,7 @@ def process(self, layer, role) -> None:
         weight, weight_scale, _ = process_fp8_weight_channel_strategy(
             layer.weight, layer.weight_scale.data
         )
-        layer.weight = Parameter(weight.t(), requires_grad=False)
+        layer.weight = Parameter(weight, requires_grad=False)
         layer.weight_scale = Parameter(weight_scale, requires_grad=False)
 
 
@@ -2280,7 +2293,7 @@ def process(self, layer, role) -> None:
 
 
 def maybe_fuse_global_scales(layer) -> None:
-    """alpha = input_global_scale * weight_global_scale, presence-gated.
+    """Alpha = input_global_scale * weight_global_scale, presence-gated.
 
     W4A4 has both -> computed; W4A16 has no input_global_scale -> skipped.
     """
@@ -2469,9 +2482,16 @@ def __init__(
 
     @property
     def supports_pre_processed_weights(self) -> bool:  # type: ignore[override]
-        # TODO(Isotr0py): support fp8/mxfp8 ModelOpt kernels transpose/repack.
+        # TODO(Isotr0py): support fp8 ModelOpt kernels transpose/repack.
         w = self.spec.weight
-        return isinstance(w, QuantKey) and w.dtype == FP4_DTYPE
+        return isinstance(w, QuantKey) and (
+            w.dtype == FP4_DTYPE
+            or (
+                w == kMxfp8Static
+                and self.kernel is not None
+                and self.kernel.supports_pre_processed_weights
+            )
+        )
 
     def create_weights(
         self,
@@ -2511,7 +2531,13 @@ def create_weights(
         expose_input_quant_key(layer, self.kernel)
 
     def process_weights_after_loading(self, layer) -> None:
+        if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False):
+            self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size)
         if is_weights_pre_processed():
+            if not self.supports_pre_processed_weights:
+                raise RuntimeError(
+                    f"{type(self.kernel).__name__} cannot use pre-processed weights"
+                )
             return
         self.fmt.pre_process(layer)
         self.wkey.process(layer, WEIGHT)
@@ -2543,8 +2569,6 @@ def process_weights_after_loading(self, layer) -> None:
                 persistent=False,
             )
             layer._nvfp4_group_size_for_gather = self.ctx.group_size
-        if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False):
-            self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size)
         self.kernel.process_weights_after_loading(layer)
 
     def apply(self, layer, x, bias=None):
diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py
index 430abdf70570..6a27332b06c8 100644
--- a/vllm/model_executor/layers/quantization/moe_wna16.py
+++ b/vllm/model_executor/layers/quantization/moe_wna16.py
@@ -229,6 +229,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
 
     Args:
         quant_config: The MOE WNA16 (W8A16/W4A16) quantization config.
+
     """
 
     def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None:
@@ -388,7 +389,7 @@ def get_fused_moe_quant_config(
         self, layer: RoutedExperts
     ) -> FusedMoEQuantConfig | None:
         if self.wna16_backend == WNA16MoEBackend.HUMMING:
-            from vllm.model_executor.layers.quantization.utils.humming_utils import (
+            from vllm.model_executor.layers.quantization.utils.humming import (
                 get_humming_moe_quant_config,
             )
 
diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py
index a16a33eaaa33..41bca8d29d9b 100644
--- a/vllm/model_executor/layers/quantization/online/base.py
+++ b/vllm/model_executor/layers/quantization/online/base.py
@@ -227,6 +227,7 @@ def _get_method_cls(
         Returns:
             The matching method class, or None when ``spec`` has no weight
             quantization.
+
         """
         if spec is None or spec.weight is None:
             return None
@@ -260,6 +261,7 @@ def resolve_quant_method_cls(
             A tuple of source, quantization key string, target pattern, spec,
             and method class. Returns None when online quantization does not
             apply to the layer.
+
         """
         quant_spec: QuantSpec | None
         if self.args.targets is not None:
@@ -313,6 +315,7 @@ def _resolve_targets_quant_method_metadata(
             A tuple of source, quantization key string, target pattern, spec,
             and dispatch table. Returns None when no pattern applies or the
             layer is ignored.
+
         """
         assert self.args.targets is not None
         ignored = should_ignore_layer(
diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py
index 0234702c8bab..ed018c81bf0a 100644
--- a/vllm/model_executor/layers/quantization/online/fp8.py
+++ b/vllm/model_executor/layers/quantization/online/fp8.py
@@ -27,6 +27,10 @@
 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
     select_fp8_moe_backend,
 )
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    expose_input_quant_key,
+)
 from vllm.model_executor.layers.linear import (
     LinearMethodBase,
 )
@@ -402,6 +406,7 @@ def process_weights_after_loading(self, layer: Module) -> None:
         replace_parameter(layer, "weight_scale", weight_scale)
 
         self.fp8_linear.process_weights_after_loading(layer)
+        expose_input_quant_key(layer, self.fp8_linear)
 
         layer._already_called_process_weights_after_loading = True
 
@@ -412,8 +417,10 @@ def apply(
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
         # if batch invariant mode is enabled dequant
-        if envs.VLLM_BATCH_INVARIANT and not isinstance(
-            self.fp8_linear, CutlassFP8ScaledMMLinearKernel
+        if (
+            envs.VLLM_BATCH_INVARIANT
+            and not isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel)
+            and not isinstance(x, QuantizedActivation)
         ):
             weight_dequant = (
                 layer.weight.to(x.dtype) * layer.weight_scale.to(x.dtype).t()
@@ -730,11 +737,13 @@ def __init__(
         )
         # Reject backends whose make_fp8_moe_quant_config branch silently
         # drops per_act_token_quant / per_out_ch_quant or collapses scales:
-        # MARLIN / CPU route through fp8_w8a16_moe_quant_config; FLASHINFER_*
-        # fold scales into a per-tensor alpha (oracle/fp8.py).
+        # MARLIN / CPU route through fp8_w8a16_moe_quant_config, CPU_W8A8
+        # ignores both flags, and FLASHINFER_* fold scales into a per-tensor
+        # alpha (oracle/fp8.py).
         if self.fp8_backend in (
             Fp8MoeBackend.MARLIN,
             Fp8MoeBackend.CPU,
+            Fp8MoeBackend.CPU_W8A8,
             Fp8MoeBackend.FLASHINFER_CUTLASS,
             Fp8MoeBackend.FLASHINFER_TRTLLM,
         ):
diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py
index 7abefce25732..2faa8f7838d0 100644
--- a/vllm/model_executor/layers/quantization/online/nvfp4.py
+++ b/vllm/model_executor/layers/quantization/online/nvfp4.py
@@ -24,7 +24,7 @@
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     amax_for_moe_weight_quant,
-    kNvfp4Dynamic,
+    kNvfp4DynamicToken,
     kNvfp4Static,
     weight_amax,
 )
@@ -100,7 +100,7 @@ def __init__(
         self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
             config=self.moe,
             weight_key=kNvfp4Static,
-            activation_key=kNvfp4Dynamic,
+            activation_key=kNvfp4DynamicToken,
         )
 
     def process_weights_after_loading(self, layer: Module) -> None:
diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py
index 855a413c45a1..2be85d817624 100644
--- a/vllm/model_executor/layers/quantization/quark/quark.py
+++ b/vllm/model_executor/layers/quantization/quark/quark.py
@@ -5,7 +5,6 @@
 from typing import TYPE_CHECKING, Any, NamedTuple, cast
 
 import torch
-from transformers import PretrainedConfig
 
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
@@ -32,6 +31,7 @@
     QuarkOCP_MX,
     QuarkScheme,
     QuarkW4A8_MXFP4_FP8,
+    QuarkW4A16Int4,
     QuarkW8A8Fp8,
     QuarkW8A8Fp8PerBlock,
     QuarkW8A8Int8,
@@ -39,6 +39,7 @@
 from vllm.model_executor.layers.quantization.quark.utils import (
     QuarkQTensorHint,
     deep_compare,
+    parse_w4a16_int4_weight_config,
     should_ignore_layer,
 )
 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -54,6 +55,10 @@
     kFp8Static128BlockSym,
     kFp8StaticChannelSym,
     kFp8StaticTensorSym,
+    kInt4Static,
+    kInt4Static32,
+    kInt4Static32Asym,
+    kInt4StaticAsym,
     kInt4W4A8StaticChannelSym,
     kInt8DynamicTensorAsym,
     kInt8DynamicTensorSym,
@@ -76,10 +81,6 @@
 
 logger = init_logger(__name__)
 
-# model_type values that use dynamic MXFP4 re-quantization for
-# OCP MX fp4 Quark checkpoints
-_DEEPSEEK_V3_FAMILY_MODEL_TYPES = frozenset({"deepseek_v3", "deepseek_v32"})
-
 
 class QuantKeyMatch(NamedTuple):
     matches: bool
@@ -113,38 +114,6 @@ def __init__(
         self.kv_cache_group = kv_cache_group
         self.kv_cache_config = kv_cache_config
         self.pack_method = pack_method
-        # Note : this flag is kept disabled because the overhead of
-        # dynamic mxfp4 quantization negates the performance gains
-        # that come from shifting to mxfp4. It is left here in case
-        # we want to re-enable it in the future.
-        self.dynamic_mxfp4_quant = False
-
-    def maybe_update_config(
-        self,
-        model_name: str,
-        hf_config: PretrainedConfig | None = None,
-        revision: str | None = None,
-    ):
-        """Enable dynamic MXFP4 only for DeepSeek-V3-family fp4 checkpoints."""
-
-        if hf_config is None:
-            return
-
-        if (
-            getattr(hf_config, "model_type", None)
-            not in _DEEPSEEK_V3_FAMILY_MODEL_TYPES
-        ):
-            return
-
-        quant_config = getattr(hf_config, "quantization_config", None)
-        if isinstance(quant_config, dict):
-            quant_dtype = (
-                quant_config.get("global_quant_config", {})
-                .get("weight", {})
-                .get("dtype")
-            )
-            if quant_dtype == "fp4":
-                self.dynamic_mxfp4_quant = True
 
     def get_linear_method(self) -> "QuarkLinearMethod":
         return QuarkLinearMethod(self)
@@ -162,19 +131,21 @@ def get_name(self) -> QuantizationMethods:
     def apply_vllm_mapper(  # noqa: B027
         self, hf_to_vllm_mapper: "WeightsMapper"
     ):
-        """
-        Interface for models to update module names referenced in
+        """Interface for models to update module names referenced in
         quantization configs in order to reflect the vllm model structure
 
         Args:
             hf_to_vllm_mapper: maps from hf model structure (the assumed
                 structure of the qconfig) to vllm model structure
+
         """
         quant_config_with_hf_to_vllm_mapper: dict[str, Any] = {}
 
         for k, v in self.quant_config.items():
-            if isinstance(v, list):
+            if isinstance(v, list) and all(isinstance(item, str) for item in v):
                 quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_list(v)
+            elif isinstance(v, list):
+                quant_config_with_hf_to_vllm_mapper[k] = v
             elif isinstance(v, dict):
                 quant_config_with_hf_to_vllm_mapper[k] = hf_to_vllm_mapper.apply_dict(v)
             else:
@@ -194,17 +165,6 @@ def get_quant_method(
             self.get_quant_method_target(prefix, type(layer))
         )
 
-        exclude_layers = cast(list[str], self.quant_config.get("exclude"))
-        is_ignored = should_ignore_layer(
-            prefix,
-            ignore=exclude_layers,
-            fused_mapping=self.packed_modules_mapping,
-            check_children=isinstance(layer, RoutedExperts),
-        )
-        dynamic_mxfp4_quant = (
-            is_ignored and "self_attn" in prefix and self.dynamic_mxfp4_quant
-        )
-
         if method_cls is UnquantizedFusedMoEMethod:
             return UnquantizedFusedMoEMethod(layer.moe_config)
         if method_cls is UnquantizedLinearMethod:
@@ -215,7 +175,9 @@ def get_quant_method(
                 scheme_cls,
                 weight_quant_key=weight_quant_key,
                 activation_quant_key=activation_quant_key,
-                dynamic_mxfp4_quant=dynamic_mxfp4_quant,
+                weight_config=self._find_matched_config(prefix, type(layer)).get(
+                    "weight"
+                ),
             )
             layer.scheme = scheme
             return QuarkLinearMethod(self)
@@ -228,6 +190,7 @@ def get_quant_method(
                 method_cls=method_cls,
                 weight_quant_key=weight_quant_key,
                 activation_quant_key=activation_quant_key,
+                layer_name=prefix,
             )
 
         return None
@@ -263,11 +226,6 @@ def get_quant_method_target(
             if is_routed_experts:
                 return None, None, UnquantizedFusedMoEMethod
             if issubclass(layer_type, LinearBase):
-                if "self_attn" in prefix and self.dynamic_mxfp4_quant:
-                    weight_quant_key, activation_key, _ = self.get_scheme_cls(
-                        layer_type, prefix
-                    )
-                    return weight_quant_key, activation_key, QuarkLinearMethod
                 return None, None, UnquantizedLinearMethod
             return None, None, None
         if issubclass(layer_type, LinearBase):
@@ -573,6 +531,26 @@ def _is_w8a8_int8(
             )
         return QuantKeyMatch(True, activation_quant_key, weight_quant_key)
 
+    def _is_w4a16_int4(
+        self,
+        weight_quant: dict[str, Any] | None,
+        input_quant: dict[str, Any] | None,
+    ) -> QuantKeyMatch:
+        if weight_quant is None or input_quant is not None:
+            return QuantKeyMatch(False, None, None)
+
+        is_int4 = weight_quant.get("dtype") in ("int4", "uint4")
+        is_packed = self.pack_method in ("order", "reorder")
+        if not (is_int4 and is_packed):
+            return QuantKeyMatch(False, None, None)
+
+        group_size, is_symmetric = parse_w4a16_int4_weight_config(weight_quant)
+        if is_symmetric:
+            weight_key = kInt4Static32 if group_size == 32 else kInt4Static
+        else:
+            weight_key = kInt4Static32Asym if group_size == 32 else kInt4StaticAsym
+        return QuantKeyMatch(True, None, weight_key)
+
     def _is_w4a8_mxfp4_fp8(
         self,
         weight_quant: dict[str, Any] | None,
@@ -657,8 +635,7 @@ def _is_w_ocp_mx_a_x(
         input_quant: dict[str, Any] | None,
         allow_static_fp8: bool = False,
     ) -> QuantKeyMatch:
-        """
-        This check returns True only if it is an OCP-MX weight quantization.
+        """This check returns True only if it is an OCP-MX weight quantization.
         The activation can be FP16/BF16, OCP MXFP4, MXFP8, FP8.
         The rationale for checking only the weight type is that
         the model loading concept and process primarily concerns the weights themselves.
@@ -892,6 +869,12 @@ def _get_scheme_cls_from_config(
                 match.activation_quant_key,
                 QuarkW8A8Int8,
             )
+        elif match := self._is_w4a16_int4(weight_config, input_config):
+            return (
+                match.weight_quant_key,
+                match.activation_quant_key,
+                QuarkW4A16Int4,
+            )
         elif match := self._is_w4a8_mxfp4_fp8(weight_config, input_config):
             return (
                 match.weight_quant_key,
@@ -920,7 +903,7 @@ def init_scheme(
         scheme_cls: type["QuarkScheme"],
         weight_quant_key: QuantKey | None,
         activation_quant_key: QuantKey | None,
-        dynamic_mxfp4_quant: bool = False,
+        weight_config: dict[str, Any] | None = None,
     ) -> "QuarkScheme":
         """Construct a Quark scheme selected by get_scheme_cls."""
         if scheme_cls not in (
@@ -930,9 +913,24 @@ def init_scheme(
             QuarkOCP_MX,
             QuarkNVFP4,
             QuarkW4A8_MXFP4_FP8,
+            QuarkW4A16Int4,
         ):
             raise AssertionError(f"Unsupported Quark scheme class: {scheme_cls}")
 
+        scheme: QuarkScheme
+        if scheme_cls is QuarkW4A16Int4:
+            normalized_weight_config = self._unwrap_single_quant_config(weight_config)
+            assert normalized_weight_config is not None
+            assert weight_quant_key is not None
+            scheme = QuarkW4A16Int4(
+                weight_quant_key=weight_quant_key,
+                activation_quant_key=activation_quant_key,
+                pack_method=self.pack_method,
+                weight_config=normalized_weight_config,
+            )
+            self._check_scheme_supported(scheme.get_min_capability())
+            return scheme
+
         kwargs: dict[str, Any] = {"activation_quant_key": activation_quant_key}
 
         # These classes handle different possible `weight_quant_key`.
@@ -944,9 +942,6 @@ def init_scheme(
         ):
             kwargs["weight_quant_key"] = weight_quant_key
 
-        if scheme_cls is QuarkOCP_MX:
-            kwargs["dynamic_mxfp4_quant"] = dynamic_mxfp4_quant
-
         scheme = scheme_cls(**kwargs)
 
         # Raise error if device does not support the scheme
@@ -984,8 +979,7 @@ def create_weights(
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        """
-        Use the CompressedTensorsScheme associated with each layer to create
+        """Use the CompressedTensorsScheme associated with each layer to create
         the necessary parameters for the layer. See LinearMethodBase for param
         details
         """
@@ -1006,8 +1000,7 @@ def apply(
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ):
-        """
-        Use the output of create_weights and the CompressedTensorsScheme
+        """Use the output of create_weights and the CompressedTensorsScheme
         associated with the layer to apply the forward pass with the
         layer input.  See LinearMethodBase for param details
 
@@ -1020,9 +1013,7 @@ def apply(
 
 
 class QuarkKVCacheMethod(BaseKVCacheMethod):
-    """
-    Supports loading kv-cache scaling factors from quark checkpoints.
-    """
+    """Supports loading kv-cache scaling factors from quark checkpoints."""
 
     def __init__(self, quant_config: QuarkConfig):
         self.validate_kv_cache_config(quant_config.kv_cache_config)
@@ -1030,12 +1021,12 @@ def __init__(self, quant_config: QuarkConfig):
 
     @staticmethod
     def validate_kv_cache_config(kv_cache_config: dict[str, Any] | None):
-        """
-        Validator for the kv cache configuration. Useful for controlling the
+        """Validator for the kv cache configuration. Useful for controlling the
         kv cache quantization schemes, that are being supported in vLLM
 
         Args:
             kv_cache_config: the quark kv cache scheme
+
         """
         if kv_cache_config is None:
             return
diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py
index d42a8884a6a9..2e41e768ef78 100644
--- a/vllm/model_executor/layers/quantization/quark/quark_moe.py
+++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 
@@ -40,6 +40,12 @@
     make_int8_moe_quant_config,
     select_int8_moe_backend,
 )
+from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
+    convert_to_wna16_moe_kernel_format,
+    make_wna16_moe_kernel,
+    make_wna16_moe_quant_config,
+    select_wna16_moe_backend,
+)
 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
     TRITON_BACKENDS,
     Mxfp4MoeBackend,
@@ -56,7 +62,10 @@
     make_nvfp4_moe_quant_config,
     select_nvfp4_moe_backend,
 )
-from vllm.model_executor.layers.quantization.quark.utils import QuarkQTensorHint
+from vllm.model_executor.layers.quantization.quark.utils import (
+    QuarkQTensorHint,
+    parse_w4a16_int4_weight_config,
+)
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     validate_fp8_block_shape_moe,
 )
@@ -75,6 +84,10 @@
     kFp8Static128BlockSym,
     kFp8StaticChannelSym,
     kFp8StaticTensorSym,
+    kInt4Static,
+    kInt4Static32,
+    kInt4Static32Asym,
+    kInt4StaticAsym,
     kInt4W4A8StaticChannelSym,
     kInt8DynamicTensorAsym,
     kInt8DynamicTensorSym,
@@ -102,6 +115,7 @@
 
 __all__ = [
     "QuarkMoEMethod",
+    "QuarkW4A16Int4MoEMethod",
     "QuarkW8A8Fp8MoEMethod",
     "QuarkOCP_MX_MoEMethod",
     "QuarkNvfp4MoEMethod",
@@ -179,6 +193,12 @@ def get_moe_method_target(
                 match.activation_quant_key,
                 QuarkW8A8Fp8MoEMethod,
             )
+        if match := quant_config._is_w4a16_int4(weight_config, input_config):
+            return (
+                match.weight_quant_key,
+                match.activation_quant_key,
+                QuarkW4A16Int4MoEMethod,
+            )
         if match := quant_config._is_w_ocp_mx_a_x(
             weight_config, input_config, allow_static_fp8=True
         ):
@@ -202,6 +222,7 @@ def get_moe_method(
         method_cls: type["QuarkMoEMethod"],
         weight_quant_key: QuantKey | None,
         activation_quant_key: QuantKey | None,
+        layer_name: str | None = None,
     ) -> "QuarkMoEMethod":
         if weight_quant_key is None:
             raise RuntimeError("Unsupported FusedMoe scheme")
@@ -226,10 +247,339 @@ def get_moe_method(
             return QuarkW8A8Int8MoEMethod(
                 module.moe_config, weight_quant_key, activation_quant_key
             )
+        elif method_cls is QuarkW4A16Int4MoEMethod:
+            assert layer_name is not None
+            layer_quant_config = quant_config._find_matched_config(
+                layer_name, type(module)
+            )
+            weight_config = quant_config._unwrap_single_quant_config(
+                layer_quant_config.get("weight")
+            )
+            if weight_config is None:
+                raise RuntimeError(
+                    f"Missing weight quantization config for layer {layer_name}"
+                )
+            return QuarkW4A16Int4MoEMethod(
+                weight_quant_key,
+                activation_quant_key,
+                weight_config,
+                quant_config.pack_method,
+                module.moe_config,
+                quant_config=quant_config,
+            )
         else:
             raise RuntimeError("Unsupported FusedMoe scheme")
 
 
+class QuarkW4A16Int4MoEMethod(QuarkMoEMethod):
+    """Quark packed INT4 weight-only MoE method."""
+
+    supported_activation_quant_keys: list[QuantKey | None] = [None]
+    supported_weight_quant_keys: list[QuantKey] = [
+        kInt4Static,
+        kInt4Static32,
+        kInt4StaticAsym,
+        kInt4Static32Asym,
+    ]
+
+    def __init__(
+        self,
+        weight_quant_key: QuantKey,
+        activation_quant_key: QuantKey | None,
+        weight_config: dict[str, Any],
+        pack_method: str,
+        moe: FusedMoEConfig,
+        quant_config: "QuarkConfig | None" = None,  # type: ignore # noqa F821
+    ):
+        super().__init__(moe, weight_quant_key, activation_quant_key)
+
+        self.quant_config = quant_config
+        # group_size is read from the Quark config rather than weight_quant_key:
+        # the int4 QuantKeys only encode group_size 32 or -1, but Quark int4
+        # checkpoints commonly use group_size 128, which the QuantKey cannot
+        # represent. Symmetry is unambiguous, so it comes from the key.
+        self.group_size = parse_w4a16_int4_weight_config(weight_config)[0]
+        self.is_symmetric = weight_quant_key.symmetric
+        self.bit8_pack_factor = 2
+        self.pack_reorder = pack_method == "reorder"
+
+        self.use_wna16_backend = True
+        self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
+            config=self.moe,
+            weight_key=weight_quant_key,
+            quant_config=self.quant_config,
+            may_have_zp=not self.is_symmetric,
+            may_have_bias=False,
+        )
+
+    def create_weights(
+        self,
+        layer: RoutedExperts,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        group_size = self.group_size
+        group_size_div_factor = 1
+
+        while intermediate_size_per_partition % group_size or hidden_size % group_size:
+            group_size = group_size // 2
+            group_size_div_factor *= 2
+            assert group_size >= 32
+        layer.group_size = group_size
+        layer.group_size_div_factor = group_size_div_factor
+
+        strategy = FusedMoeWeightScaleSupported.GROUP.value
+        extra_weight_attrs.update({"quant_method": strategy, "is_transposed": False})
+
+        assert "weight_loader" in extra_weight_attrs
+        weight_loader = extra_weight_attrs["weight_loader"]
+        extra_weight_attrs["weight_loader"] = self.get_weight_loader(
+            layer, weight_loader
+        )
+
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size // self.bit8_pack_factor,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition // self.bit8_pack_factor,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        w13_weight_scale = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size // group_size,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_scale", w13_weight_scale)
+        set_weight_attrs(w13_weight_scale, extra_weight_attrs)
+
+        w2_weight_scale = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition // group_size,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_scale", w2_weight_scale)
+        set_weight_attrs(w2_weight_scale, extra_weight_attrs)
+
+        w13_weight_zero_point = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                2 * intermediate_size_per_partition // self.bit8_pack_factor,
+                hidden_size // group_size,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_zero_point", w13_weight_zero_point)
+        set_weight_attrs(w13_weight_zero_point, extra_weight_attrs)
+
+        w2_weight_zero_point = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                hidden_size // self.bit8_pack_factor,
+                intermediate_size_per_partition // group_size,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_zero_point", w2_weight_zero_point)
+        set_weight_attrs(w2_weight_zero_point, extra_weight_attrs)
+
+    def process_weights_after_loading(self, layer: RoutedExperts) -> None:
+        w13_qzeros = layer.w13_weight_zero_point if not self.is_symmetric else None
+        w2_qzeros = layer.w2_weight_zero_point if not self.is_symmetric else None
+
+        converted = convert_to_wna16_moe_kernel_format(
+            backend=self.wna16_backend,
+            layer=layer,
+            quant_config=self.quant_config,
+            input_dtype=None,
+            w13=layer.w13_weight,
+            w2=layer.w2_weight,
+            w13_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            w13_qzeros=w13_qzeros,
+            w2_qzeros=w2_qzeros,
+        )
+
+        if converted is None:
+            # Backend rewrote the layer's params in place (e.g. Humming).
+            self._setup_kernel(layer)
+            return
+
+        (
+            w13_qweight,
+            w2_qweight,
+            w13_scales,
+            w2_scales,
+            _,  # w13_qzeros
+            _,  # w2_qzeros
+            _,  # w13_input_global_scale
+            _,  # w2_input_global_scale
+            _,  # w13_bias
+            _,  # w2_bias
+        ) = converted
+
+        replace_parameter(layer, "w13_weight", w13_qweight)
+        replace_parameter(layer, "w2_weight", w2_qweight)
+        replace_parameter(layer, "w13_weight_scale", w13_scales)
+        replace_parameter(layer, "w2_weight_scale", w2_scales)
+
+        self._setup_kernel(layer)
+
+    def _setup_kernel(self, layer: RoutedExperts) -> None:
+        assert self.experts_cls is not None
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        assert self.moe_quant_config is not None
+        self.moe_kernel = make_wna16_moe_kernel(
+            moe_quant_config=self.moe_quant_config,
+            moe_config=self.moe,
+            experts_cls=self.experts_cls,
+            backend=self.wna16_backend,
+            routing_tables=layer._expert_routing_tables(),
+        )
+
+    def get_fused_moe_quant_config(
+        self, layer: RoutedExperts
+    ) -> FusedMoEQuantConfig | None:
+        w1_zp = layer.w13_weight_zero_point if not self.is_symmetric else None
+        w2_zp = layer.w2_weight_zero_point if not self.is_symmetric else None
+        return make_wna16_moe_quant_config(
+            w1_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            group_size=layer.group_size,
+            num_bits=4,
+        )
+
+    def apply(
+        self,
+        layer: RoutedExperts,
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts: SharedExperts | None,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor:
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply(
+            hidden_states=x,
+            w1=layer.w13_weight,
+            w2=layer.w2_weight,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            expert_map=layer.expert_map,
+            shared_experts=shared_experts,
+            shared_experts_input=shared_experts_input,
+        )
+
+    def get_weight_loader(self, layer: RoutedExperts, weight_loader):
+        def convert_quark_tensor(tensor: torch.Tensor, tensor_type: str):
+            size0 = tensor.size(0)
+            tensor = tensor.view(torch.uint8)
+
+            shifter = torch.tensor([0, 4], dtype=torch.uint8, device=tensor.device)
+            tensor = (tensor[:, :, None] >> shifter) & 0xF
+
+            if self.pack_reorder:
+                reverse_awq_pack_order = [0, 4, 1, 5, 2, 6, 3, 7]
+                tensor = tensor.view(-1, 8)[:, reverse_awq_pack_order]
+            else:
+                tensor = tensor.view(-1, 8)
+            if self.is_symmetric:
+                tensor = tensor ^ 0x8
+            tensor = tensor.view(size0, -1)
+
+            tensor = tensor.T.contiguous()
+            if tensor_type == "weight":
+                tensor = tensor[:, 1::2] * 16 + tensor[:, ::2]
+            elif tensor_type == "zero_point":
+                tensor = tensor[1::2, :] * 16 + tensor[::2, :]
+            return tensor
+
+        def quark_int4_weight_loader(
+            param: torch.nn.Parameter,
+            loaded_weight: torch.Tensor,
+            weight_name: str,
+            shard_id: str,
+            expert_id: int,
+            return_success: bool = False,
+        ):
+            loaded_weight = loaded_weight.to(param.data.device)
+            shard_size = layer.intermediate_size_per_partition
+
+            if weight_name.endswith("_weight"):
+                loaded_weight = convert_quark_tensor(loaded_weight, "weight")
+            elif weight_name.endswith("_weight_zero_point"):
+                loaded_weight = convert_quark_tensor(loaded_weight, "zero_point")
+            elif weight_name.endswith("_weight_scale"):
+                loaded_weight = loaded_weight.T
+
+            if layer.group_size_div_factor > 1 and (
+                "weight_zero_point" in weight_name or "weight_scale" in weight_name
+            ):
+                loaded_weight = loaded_weight.repeat_interleave(
+                    layer.group_size_div_factor, 1
+                )
+
+            if "w13_weight_zero_point" in weight_name:
+                tensor = loaded_weight.view(
+                    layer.moe_config.tp_size, -1, loaded_weight.size(1)
+                )[layer.moe_config.tp_rank]
+                if shard_id == "w1":
+                    param.data[expert_id, : shard_size // 2] = tensor
+                else:
+                    param.data[expert_id, shard_size // 2 :] = tensor
+                return True if return_success else None
+            elif "w2_weight_zero_point" in weight_name:
+                param.data[expert_id] = loaded_weight.view(
+                    loaded_weight.size(0), layer.moe_config.tp_size, -1
+                )[:, layer.moe_config.tp_rank]
+                return True if return_success else None
+
+            return weight_loader(
+                param,
+                loaded_weight,
+                weight_name,
+                shard_id,
+                expert_id,
+                return_success=return_success,
+            )
+
+        return quark_int4_weight_loader
+
+
 class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
     supported_activation_quant_keys = [
         kFp8StaticTensorSym,
@@ -1750,10 +2100,7 @@ def create_weights(
         set_weight_attrs(w2_input_scale_2, extra_weight_attrs)
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
-        """
-        Convert NVFP4 MoE weights into kernel format and setup the kernel.
-        """
-
+        """Convert NVFP4 MoE weights into kernel format and setup the kernel."""
         # Match existing NVFP4 MoE paths: fused w13 uses the w1 global scale.
         if self.moe.is_act_and_mul and not torch.allclose(
             layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/__init__.py b/vllm/model_executor/layers/quantization/quark/schemes/__init__.py
index 4db1ba384555..a5857a776f5e 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/__init__.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/__init__.py
@@ -5,6 +5,7 @@
 from .quark_ocp_mx import QuarkOCP_MX
 from .quark_scheme import QuarkScheme
 from .quark_w4a8_mxfp4_fp8 import QuarkW4A8_MXFP4_FP8
+from .quark_w4a16_int4 import QuarkW4A16Int4
 from .quark_w8a8_fp8 import QuarkW8A8Fp8, QuarkW8A8Fp8PerBlock
 from .quark_w8a8_int8 import QuarkW8A8Int8
 
@@ -14,6 +15,7 @@
     "QuarkW8A8Fp8PerBlock",
     "QuarkW8A8Int8",
     "QuarkOCP_MX",
+    "QuarkW4A16Int4",
     "QuarkW4A8_MXFP4_FP8",
     "QuarkNVFP4",
 ]
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py
index 7fb0332b574b..777423d78032 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py
@@ -31,8 +31,7 @@
 
 
 class QuarkNVFP4(QuarkScheme):
-    """
-    Quark NVFP4 quantization scheme.
+    """Quark NVFP4 quantization scheme.
 
     Supports loading NVFP4 checkpoints with the following structure:
     - weight: uint8, shape [out_features, in_features // 2] (packed FP4)
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
index 30758782c322..50263863d886 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
@@ -26,10 +26,8 @@
 )
 from vllm.model_executor.parameter import (
     GroupQuantScaleParameter,
-    ModelWeightParameter,
     PackedvLLMParameter,
 )
-from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 
 from .quark_scheme import QuarkScheme
@@ -46,10 +44,9 @@ def __init__(
         self,
         weight_quant_key: QuantKey,
         activation_quant_key: QuantKey | None,
-        dynamic_mxfp4_quant: bool = False,
     ):
         super().__init__(weight_quant_key, activation_quant_key)
-        self.dynamic_mxfp4_quant = dynamic_mxfp4_quant
+
         self.weight_dtype = next(
             dtype
             for dtype, quant_key in _WEIGHT_QUANT_KEY_MAP.items()
@@ -109,21 +106,9 @@ def get_packed_dim(self, dim: int, quant_dtype: str):
     def get_min_capability(cls) -> int:
         return 70
 
-    def process_dynamic_mxfp4_weights_after_loading(
-        self, layer: torch.nn.Module
-    ) -> None:
-        from aiter.ops.triton.quant import dynamic_mxfp4_quant
-
-        w_q, w_s = dynamic_mxfp4_quant(layer.weight)
-        layer.weight_scale = torch.nn.Parameter(w_s, requires_grad=False)
-        layer.weight = torch.nn.Parameter(w_q, requires_grad=False)
-
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False)
 
-        if self.dynamic_mxfp4_quant:
-            self.process_dynamic_mxfp4_weights_after_loading(layer)
-
         self.ocp_mx_linear.process_weights_after_loading(layer)
 
     def create_weights(
@@ -145,51 +130,36 @@ def create_weights(
                 "tensor-parallel sharding for this layer."
             )
 
-        if self.dynamic_mxfp4_quant:
-            weight = ModelWeightParameter(
-                data=torch.empty(
-                    sum(output_partition_sizes),
-                    input_size_per_partition,
-                    dtype=params_dtype,
-                ),
-                input_dim=1,
-                output_dim=0,
-                weight_loader=weight_loader,
-            )
-
-            layer.register_parameter("weight", weight)
-            set_weight_attrs(weight, kwargs)
-        else:
-            output_size_per_partition = sum(output_partition_sizes)
-            layer.logical_widths = output_partition_sizes
-
-            # WEIGHT
-            weight = PackedvLLMParameter(
-                data=torch.empty(
-                    output_size_per_partition,
-                    self.get_packed_dim(input_size_per_partition, self.weight_dtype),
-                    dtype=torch.uint8,
-                ),
-                input_dim=1,
-                output_dim=0,
-                packed_dim=1,
-                packed_factor=self.packed_factor,
-                weight_loader=weight_loader,
-            )
-            layer.register_parameter("weight", weight)
-
-            # WEIGHT SCALE
-            weight_scale = GroupQuantScaleParameter(
-                data=torch.empty(
-                    output_size_per_partition,
-                    input_size_per_partition // OCP_MX_BLOCK_SIZE,
-                    dtype=torch.uint8,
-                ),
-                input_dim=1,
-                output_dim=0,
-                weight_loader=weight_loader,
-            )
-            layer.register_parameter("weight_scale", weight_scale)
+        output_size_per_partition = sum(output_partition_sizes)
+        layer.logical_widths = output_partition_sizes
+
+        # WEIGHT
+        weight = PackedvLLMParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                self.get_packed_dim(input_size_per_partition, self.weight_dtype),
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            packed_dim=1,
+            packed_factor=self.packed_factor,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # WEIGHT SCALE
+        weight_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // OCP_MX_BLOCK_SIZE,
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_scale", weight_scale)
 
         if self.weight_quant_key == kMxfp4Static:
             self.ocp_mx_linear = init_mxfp4_linear_kernel(
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py
index 419055b37217..6e75337e5ac3 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py
@@ -11,8 +11,7 @@
 
 
 class QuarkScheme(ABC):
-    """
-    Abstract class used to describe the weight creation and forward pass
+    """Abstract class used to describe the weight creation and forward pass
     of different quantization schemes supported by Quark.
     """
 
@@ -36,25 +35,19 @@ def __init__(
     @classmethod
     @abstractmethod
     def get_min_capability(cls) -> int:
-        """
-        Get minimum device capability.
-        """
+        """Get minimum device capability."""
         raise NotImplementedError
 
     @abstractmethod
     def create_weights(self, *args, **kwargs):
-        """
-        Weight creation for the particular scheme. Inputs to this function
-
-        """
+        """Weight creation for the particular scheme. Inputs to this function."""
         raise NotImplementedError
 
     @abstractmethod
     def apply_weights(
         self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
     ):
-        """
-        Run the forward pass for the particular scheme. This is where
+        """Run the forward pass for the particular scheme. This is where
         scheme-specific dequant/quant steps/kernels should be applied.
 
         Args:
@@ -62,13 +55,13 @@ def apply_weights(
                 other parameters relevant to the particular scheme.
             x: input to the layer
             bias: bias parameter
+
         """
         raise NotImplementedError
 
     @abstractmethod
     def process_weights_after_loading(self, layer: torch.nn.Module):
-        """
-        Called after weight loading is complete for any cleanup that
+        """Called after weight loading is complete for any cleanup that
         needs to occur.
         """
         raise NotImplementedError
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a16_int4.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a16_int4.py
new file mode 100644
index 000000000000..cded0f80a9eb
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a16_int4.py
@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+from vllm.model_executor.kernels.linear import (
+    MPLinearKernel,
+    MPLinearLayerConfig,
+    choose_mp_linear_kernel,
+)
+from vllm.model_executor.layers.quantization.auto_awq import (
+    _convert_awq_to_standard_format,
+)
+from vllm.model_executor.layers.quantization.quark.schemes.quark_scheme import (
+    QuarkScheme,
+)
+from vllm.model_executor.layers.quantization.quark.utils import (
+    canonicalize_quark_packed_int4,
+    parse_w4a16_int4_weight_config,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kInt4Static,
+    kInt4Static32,
+    kInt4Static32Asym,
+    kInt4StaticAsym,
+)
+from vllm.model_executor.parameter import (
+    GroupQuantScaleParameter,
+    PackedvLLMParameter,
+)
+from vllm.scalar_type import scalar_types
+
+
+class QuarkW4A16Int4(QuarkScheme):
+    """Quark packed INT4 weight-only linear scheme via MPLinearKernel."""
+
+    supported_activation_quant_keys: list[QuantKey | None] = [None]
+    supported_weight_quant_keys: list[QuantKey] = [
+        kInt4Static,
+        kInt4Static32,
+        kInt4StaticAsym,
+        kInt4Static32Asym,
+    ]
+
+    def __init__(
+        self,
+        weight_quant_key: QuantKey,
+        activation_quant_key: QuantKey | None,
+        pack_method: str,
+        weight_config: dict[str, Any],
+    ):
+        super().__init__(weight_quant_key, activation_quant_key)
+
+        # group_size is read from the Quark config rather than weight_quant_key:
+        # the int4 QuantKeys only encode group_size 32 or -1, but Quark int4
+        # checkpoints commonly use group_size 128, which the QuantKey cannot
+        # represent. Symmetry is unambiguous, so it comes from the key.
+        self.group_size = parse_w4a16_int4_weight_config(weight_config)[0]
+        self.pack_factor = 8
+        self.pack_reorder = pack_method == "reorder"
+        self.is_symmetric = weight_quant_key.symmetric
+        self.quant_type = (
+            scalar_types.uint4b8 if self.is_symmetric else scalar_types.uint4
+        )
+        self.kernel: MPLinearKernel | None = None
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 70
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        output_partition_sizes: list[int],
+        input_size_per_partition: int,
+        params_dtype: torch.dtype,
+        weight_loader: Callable,
+        **kwargs,
+    ):
+        input_size = kwargs["input_size"]
+        output_size = kwargs["output_size"]
+        group_size = (
+            self.group_size if self.group_size != -1 else input_size_per_partition
+        )
+        if input_size_per_partition % group_size != 0:
+            raise ValueError(
+                "The input size is not aligned with the quantized weight shape. "
+                "This can be caused by too large tensor parallel size. "
+                f"input_size_per_partition={input_size_per_partition}, "
+                f"group_size={group_size}."
+            )
+
+        output_size_per_partition = sum(output_partition_sizes)
+        packed_output_size_per_partition = math.ceil(
+            output_size_per_partition / self.pack_factor
+        )
+        layer.output_size_per_partition = output_size_per_partition
+        layer.packed_output_size_per_partition = packed_output_size_per_partition
+
+        mp_linear_kernel_config = MPLinearLayerConfig(
+            full_weight_shape=(input_size, output_size),
+            partition_weight_shape=(
+                input_size_per_partition,
+                output_size_per_partition,
+            ),
+            weight_type=self.quant_type,
+            act_type=params_dtype,
+            group_size=group_size,
+            zero_points=not self.is_symmetric,
+        )
+        kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
+
+        def weight_scale_loader(
+            param: torch.nn.Parameter,
+            loaded_weight: torch.Tensor,
+            *args,
+            **loader_kwargs,
+        ) -> None:
+            if loaded_weight.shape[1] < param.data.shape[1]:
+                padded_weight = loaded_weight.new_zeros(param.data.shape)
+                padded_weight[:, : loaded_weight.shape[1]] = loaded_weight
+                loaded_weight = padded_weight
+            weight_loader(param, loaded_weight, *args, **loader_kwargs)
+
+        weight = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition,
+                packed_output_size_per_partition,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.pack_factor,
+            weight_loader=weight_loader,
+        )
+        num_groups = input_size_per_partition // group_size
+        weight_zero_point = PackedvLLMParameter(
+            data=torch.zeros(
+                num_groups,
+                packed_output_size_per_partition,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.pack_factor,
+            weight_loader=weight_loader,
+        )
+        weight_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                num_groups,
+                packed_output_size_per_partition * self.pack_factor,
+                dtype=params_dtype,
+            ),
+            input_dim=0,
+            output_dim=1,
+            weight_loader=weight_scale_loader,
+        )
+
+        layer.register_parameter("weight", weight)
+        layer.register_parameter("weight_zero_point", weight_zero_point)
+        layer.register_parameter("weight_scale", weight_scale)
+
+        self.kernel = kernel_type(
+            mp_linear_kernel_config,
+            w_q_param_name="weight",
+            w_s_param_name="weight_scale",
+            w_zp_param_name="weight_zero_point" if not self.is_symmetric else None,
+        )
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.weight.data = canonicalize_quark_packed_int4(
+            layer.weight.data,
+            pack_reorder=self.pack_reorder,
+            is_symmetric=self.is_symmetric,
+            pack_factor=self.pack_factor,
+        )
+        if not self.is_symmetric:
+            layer.weight_zero_point.data = canonicalize_quark_packed_int4(
+                layer.weight_zero_point.data,
+                pack_reorder=self.pack_reorder,
+                is_symmetric=self.is_symmetric,
+                pack_factor=self.pack_factor,
+            )
+        output_size = layer.output_size_per_partition
+        packed_output_size = layer.packed_output_size_per_partition * self.pack_factor
+        if output_size < packed_output_size:
+            layer.weight_scale.data[:, output_size:].zero_()
+
+        _convert_awq_to_standard_format(
+            layer,
+            "weight",
+            "weight_zero_point" if not self.is_symmetric else None,
+            self.quant_type.size_bits,
+        )
+        assert self.kernel is not None
+        self.kernel.process_weights_after_loading(layer)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ):
+        assert self.kernel is not None
+        return self.kernel.apply_weights(layer, x, bias)
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py
index 0f0b2065a8bb..1118db2db137 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py
@@ -33,8 +33,7 @@
 
 
 class QuarkW4A8_MXFP4_FP8(QuarkScheme):
-    """
-    - Weights: MXFP4 with E8M0 scales per block of 32
+    """- Weights: MXFP4 with E8M0 scales per block of 32
     - Activations: FP8 E4M3 (static per-tensor quantization)
 
     Uses the AITER Triton kernel and falls back to emulation if AITER not available.
diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
index 57b08361284a..1ef7c1dcc9af 100644
--- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
@@ -9,6 +9,10 @@
 from vllm.config import get_current_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    expose_input_quant_key,
+)
 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     create_fp8_scale_parameter,
@@ -195,10 +199,12 @@ def create_weights(
             module_name=self.__class__.__name__,
         )
 
+        expose_input_quant_key(layer, self.fp8_linear)
+
     def apply_weights(
         self,
         layer: torch.nn.Module,
-        x: torch.Tensor,
+        x: torch.Tensor | QuantizedActivation,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
         return self.fp8_linear.apply_weights(layer, x, bias)
diff --git a/vllm/model_executor/layers/quantization/quark/utils.py b/vllm/model_executor/layers/quantization/quark/utils.py
index 48c978263ee7..ac455e38b83f 100644
--- a/vllm/model_executor/layers/quantization/quark/utils.py
+++ b/vllm/model_executor/layers/quantization/quark/utils.py
@@ -66,6 +66,63 @@ def should_ignore_layer(
     return all(shards_ignored)
 
 
+def parse_w4a16_int4_weight_config(
+    weight_config: Mapping[str, Any],
+) -> tuple[int, bool]:
+    """Parse required W4A16 INT4/UINT4 weight fields from Quark config."""
+    if "group_size" not in weight_config:
+        raise ValueError(
+            "Quark W4A16 INT4/UINT4 configs must specify weight.group_size"
+        )
+    if "symmetric" not in weight_config:
+        raise ValueError("Quark W4A16 INT4/UINT4 configs must specify weight.symmetric")
+
+    group_size = weight_config["group_size"]
+    is_symmetric = weight_config["symmetric"]
+    if not isinstance(group_size, int) or group_size <= 0:
+        raise ValueError(
+            f"Quark W4A16 weight.group_size must be a positive int, got {group_size!r}"
+        )
+    if not isinstance(is_symmetric, bool):
+        raise ValueError(
+            f"Quark W4A16 weight.symmetric must be a bool, got {is_symmetric!r}"
+        )
+    return group_size, is_symmetric
+
+
+def canonicalize_quark_packed_int4(
+    packed_weight: torch.Tensor,
+    *,
+    pack_reorder: bool,
+    is_symmetric: bool,
+    pack_factor: int = 8,
+) -> torch.Tensor:
+    """Convert Quark export nibble layout to AWQ checkpoint layout."""
+    from vllm.model_executor.layers.quantization.auto_awq import (
+        _REVERSE_AWQ_PACK_ORDER,
+    )
+
+    if pack_reorder:
+        source_order = torch.tensor(
+            _REVERSE_AWQ_PACK_ORDER, device=packed_weight.device, dtype=torch.int32
+        )
+    else:
+        source_order = torch.arange(
+            pack_factor, device=packed_weight.device, dtype=torch.int32
+        )
+    target_order = torch.tensor(
+        _REVERSE_AWQ_PACK_ORDER, device=packed_weight.device, dtype=torch.int32
+    )
+    source_shifts = source_order * 4
+    target_shifts = target_order * 4
+
+    values = (packed_weight.to(torch.int32)[..., None] >> source_shifts) & 0xF
+    if is_symmetric:
+        values = values ^ 0x8
+    packed = (values.to(torch.int64) << target_shifts.to(torch.int64)).sum(dim=-1)
+    return packed.to(torch.int32)
+
+
 # utility for tensor dims > 2 cases
 def quark_quantize_weight_to_mxfp4(w: torch.Tensor):
     assert w.dtype == torch.bfloat16, (
diff --git a/vllm/model_executor/layers/quantization/qutlass_utils.py b/vllm/model_executor/layers/quantization/qutlass_utils.py
index 86b0548307af..a997628a49ef 100644
--- a/vllm/model_executor/layers/quantization/qutlass_utils.py
+++ b/vllm/model_executor/layers/quantization/qutlass_utils.py
@@ -31,8 +31,7 @@ def triton_scale_swizzle(
     BLOCK_ROWS: tl.constexpr,
     BLOCK_COLS: tl.constexpr,
 ):
-    """
-    Rearranges tensor data from row-major to block-scaled swizzle format.
+    """Rearranges tensor data from row-major to block-scaled swizzle format.
 
     Args:
         scale_ptr: Pointer to the input scale tensor
@@ -43,6 +42,7 @@ def triton_scale_swizzle(
         output_block_stride: Stride between blocks in the output tensor
         BLOCK_ROWS: Number of rows in a tile (compile-time constant)
         BLOCK_COLS: Number of columns in a tile (compile-time constant)
+
     """
     pid_row = tl.program_id(0)
     pid_col = tl.program_id(1)
@@ -86,8 +86,7 @@ def triton_scale_swizzle(
 
 @torch.library.custom_op("vllm::triton_mx_block_rearrange", mutates_args=())
 def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:
-    """
-    Rearranges an E8M0 tensor scale from row-major format to
+    """Rearranges an E8M0 tensor scale from row-major format to
     block-scaled swizzle format.
 
     This format is suitable for Tmem as described in NVIDIA documentation:
@@ -98,6 +97,7 @@ def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:
 
     Returns:
         Rearranged tensor in block-scaled swizzle format
+
     """
     assert scale_tensor.element_size() == 1, (
         "Expected element size to be 1 byte (8 bits)"
@@ -154,8 +154,7 @@ def _triton_mx_block_rearrange_fake(scale_tensor: torch.Tensor) -> torch.Tensor:
 def to_blocked(
     input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton"
 ) -> torch.Tensor:
-    """
-    Rearrange a large matrix by breaking it into blocks and applying
+    """Rearrange a large matrix by breaking it into blocks and applying
     the rearrangement pattern.
 
     See:
@@ -167,6 +166,7 @@ def to_blocked(
 
     Returns:
         Rearranged flattened tensor of size (32*cdiv(H,128) * 16*cdiv(W,4))
+
     """
     if backend == "triton":
         return triton_mx_block_rearrange(input_matrix).flatten()
diff --git a/vllm/model_executor/layers/quantization/torchao.py b/vllm/model_executor/layers/quantization/torchao.py
index 8862578d6f7e..0a24c1e4a9fe 100644
--- a/vllm/model_executor/layers/quantization/torchao.py
+++ b/vllm/model_executor/layers/quantization/torchao.py
@@ -74,8 +74,7 @@ def torchao_version_at_least(torchao_version: str) -> bool:
 
 
 def should_skip(prefix: str, skip_modules: list[str]) -> bool:
-    """
-    Robust skipping logic:
+    """Robust skipping logic:
     should_skip("model.model.layers.1.q_proj",
                 ["model.model.layers.1.q_proj"])  # True
     should_skip("model.model.layers.10.o_proj", ["o_proj"])  -> True
@@ -163,14 +162,14 @@ def get_min_capability(cls) -> int:
 
     @staticmethod
     def get_config_filenames() -> list[str]:
-        """torchao doesn't require additional config files, we use
+        """Torchao doesn't require additional config files, we use
         `config.json` from huggingface: `model_config.hf_config`
         """
         return []
 
     @classmethod
     def from_config(cls, config: dict[str, Any]) -> "TorchAOConfig":
-        """Create the quant config from an hf model config"""
+        """Create the quant config from an hf model config."""
         try:
             from torchao.core.config import config_from_dict
         except ImportError as err:
@@ -287,12 +286,13 @@ def get_quant_method(
 def torchao_quantize_param_data(
     param: torch.Tensor, torchao_config: Any
 ) -> torch.nn.Parameter:
-    """Quantize a Tensor with torchao quantization specified by torchao_config
+    """Quantize a Tensor with torchao quantization specified by torchao_config.
 
     Args:
         param: weight parameter of the linear module
         torchao_config: type of quantization and their arguments we want to
             use to quantize the Tensor
+
     """
     from torchao.core.config import AOBaseConfig
     from torchao.quantization import quantize_
@@ -322,6 +322,7 @@ class TorchAOLinearMethod(LinearMethodBase):
     Args:
         quant_config: The torchao quantization config, a string that encodes
             the type of quantization and all relevant arguments.
+
     """
 
     def __init__(self, quant_config: TorchAOConfig):
diff --git a/vllm/model_executor/layers/quantization/turboquant/centroids.py b/vllm/model_executor/layers/quantization/turboquant/centroids.py
index 490265747c5b..65b1c821ccbc 100644
--- a/vllm/model_executor/layers/quantization/turboquant/centroids.py
+++ b/vllm/model_executor/layers/quantization/turboquant/centroids.py
@@ -45,6 +45,7 @@ def solve_lloyd_max(
     Returns:
         centroids: Sorted tensor of 2^bits optimal centroids.
         boundaries: Sorted tensor of 2^bits - 1 decision boundaries.
+
     """
     n_levels = 2**bits
     sigma2 = 1.0 / d
diff --git a/vllm/model_executor/layers/quantization/turboquant/config.py b/vllm/model_executor/layers/quantization/turboquant/config.py
index 84e3940d264d..5b25f32b60e2 100644
--- a/vllm/model_executor/layers/quantization/turboquant/config.py
+++ b/vllm/model_executor/layers/quantization/turboquant/config.py
@@ -83,6 +83,7 @@ class TurboQuantConfig:
         norm_correction: Re-normalize centroid vectors to unit norm before
             inverse rotation during dequant. Fixes quantization-induced norm
             distortion, improving PPL by ~0.8% at 4-bit.
+
     """
 
     head_dim: int = 128
diff --git a/vllm/model_executor/layers/quantization/utils/allspark_utils.py b/vllm/model_executor/layers/quantization/utils/allspark_utils.py
deleted file mode 100644
index 4c324682e5e6..000000000000
--- a/vllm/model_executor/layers/quantization/utils/allspark_utils.py
+++ /dev/null
@@ -1,67 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-
-from vllm.platforms import current_platform
-from vllm.scalar_type import ScalarType, scalar_types
-
-ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD = 1024
-ALLSPARK_SUPPORTED_QUANT_TYPES = [scalar_types.uint8b128]
-ALLSPARK_AMPERE_N_ALIGN = 16
-ALLSPARK_AMPERE_K_ALIGN = 16
-
-
-def check_allspark_supported_dtype_shape(
-    input_size_per_partition: int,
-    output_size_per_partition: int,
-    group_size: int,
-    weight_dtype: ScalarType,
-    act_dtype: torch.dtype,
-):
-    capability_tuple = current_platform.get_device_capability()
-    device_capability = -1 if capability_tuple is None else capability_tuple.to_int()
-
-    # For Ampere GPU
-    if device_capability >= 80 and device_capability < 90:
-        if group_size != -1:
-            return (
-                False,
-                "For Ampere GPU, AllSpark does not support group_size "
-                f"= {group_size}. Only group_size = -1 are supported.",
-            )
-
-        if weight_dtype not in ALLSPARK_SUPPORTED_QUANT_TYPES:
-            return (
-                False,
-                "For Ampere GPU, AllSpark does not support "
-                f"quant type ({weight_dtype}). Only quant type "
-                f"({ALLSPARK_SUPPORTED_QUANT_TYPES}) are supported.",
-            )
-
-        if (
-            input_size_per_partition % ALLSPARK_AMPERE_K_ALIGN != 0
-            or output_size_per_partition % ALLSPARK_AMPERE_N_ALIGN != 0
-        ):
-            return (
-                False,
-                "AllSpark needs input_size_per_partition % "
-                f"{ALLSPARK_AMPERE_K_ALIGN} = 0 and "
-                f"output_size_per_partition % {ALLSPARK_AMPERE_N_ALIGN} = 0 "
-                "for Ampere GPU optimized kernels.",
-            )
-
-        if act_dtype != torch.float16 and act_dtype != torch.bfloat16:
-            return (
-                False,
-                "AllSpark only supports act_dtype = float16 or bfloat16,"
-                f"for Ampere GPU, but got act_dtype = {act_dtype}.",
-            )
-    else:
-        return (
-            False,
-            "AllSpark currently does not support "
-            f"device_capability = {device_capability}.",
-        )
-
-    return True, None
diff --git a/vllm/model_executor/layers/quantization/utils/config_utils.py b/vllm/model_executor/layers/quantization/utils/config_utils.py
index cfee80da9627..1beb988ac2ac 100644
--- a/vllm/model_executor/layers/quantization/utils/config_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/config_utils.py
@@ -63,6 +63,7 @@ def is_shared_expert_quant_fse_compatible(
 
     Returns:
         A compatibility flag and, when incompatible, the reason.
+
     """
     if projection_names is None:
         projection_names = ["gate_up_proj", "down_proj"]
@@ -303,13 +304,11 @@ def is_equal_or_regex_match(
     check_contains: bool = False,
     use_fnmatch: bool = False,
 ) -> bool:
-    """
-    Checks whether a value is exactly equal or a regex match for target
+    """Checks whether a value is exactly equal or a regex match for target
     if target starts with 're:'. If check_contains is set to True,
     additionally checks if the target string is contained within the value.
     If use_fnmatch is set, supports shell-style patterns in target.
     """
-
     if target.startswith("re:"):
         pattern = target[3:]
         if re.match(pattern, value):
diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
index efe75fc6b611..c5c01e59ca3a 100644
--- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Utility helpers for NVFP4 + FlashInfer fused-MoE path"""
+"""Utility helpers for NVFP4 + FlashInfer fused-MoE path."""
 
 from typing import TYPE_CHECKING
 
diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py
index 1b5320615f4a..8df24ae9e6ce 100644
--- a/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py
+++ b/vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Utility helpers for MxInt4 + FlashInfer fused-MoE path"""
+"""Utility helpers for MxInt4 + FlashInfer fused-MoE path."""
 
 import functools
 
@@ -37,8 +37,7 @@ def prepare_static_weights_for_trtllm_mxint4_moe(
     gemm2_weights: torch.Tensor,
     gemm2_scales: torch.Tensor,
 ) -> dict[str, torch.Tensor]:
-    """
-    Prepare MxInt4 weights for TRT-LLM kernel.
+    """Prepare MxInt4 weights for TRT-LLM kernel.
 
     Input:
         gemm1_weights: [num_experts, 2*intermediate_size, hidden_size//8] int32
@@ -51,6 +50,7 @@ def prepare_static_weights_for_trtllm_mxint4_moe(
     Returns:
         Dict with keys 'gemm1_weights', 'gemm1_scales', 'gemm2_weights',
             'gemm2_scales' containing shuffled/packed tensors ready for kernel
+
     """
     from flashinfer import block_scale_interleave
     from flashinfer.fused_moe import (
@@ -192,8 +192,7 @@ def flashinfer_trtllm_mxint4_moe(
     routing_method_type: int | None = None,
     routing_replay_out: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """
-    Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel.
+    """Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel.
 
     Args:
         x: Input hidden states. dtype: bfloat16
@@ -209,11 +208,14 @@ def flashinfer_trtllm_mxint4_moe(
         ep_rank: Expert parallelism rank (default: 0)
         num_expert_group: Number of expert groups (default: None -> 0)
         topk_group: Top-k within groups (default: None -> 0)
+        routing_replay_out: Optional buffer receiving the routing decisions
+            so they can be replayed later
         e_score_correction_bias: Optional routing bias. dtype: bfloat16
         routing_method_type: FlashInfer RoutingMethodType enum value
 
     Returns:
         Output tensor from MoE layer. dtype: same as x (bfloat16)
+
     """
     from flashinfer import RoutingMethodType
     from flashinfer.fused_moe import trtllm_mxint4_block_scale_moe
diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
index b31309679479..1bfd7e61f90a 100644
--- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
@@ -59,7 +59,7 @@ def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:
 def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(
     gemm1_weights: torch.Tensor, gemm2_weights: torch.Tensor, is_gated_activation: bool
 ):
-    """Shuffle weights for FI TRT-LLM Format"""
+    """Shuffle weights for FI TRT-LLM Format."""
     from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a
 
     epilogue_tile_m = 128
@@ -112,8 +112,9 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Convert expert weights to FlashInfer's block layout.
 
-    This reorders W13 and W2 into the expected epilogue-tiled block layout and
-    returns the shuffled weight tensors.
+    This reorders W13 and W2 in place into the expected epilogue-tiled block
+    layout and returns views of the shuffled weight tensors. Using one expert
+    as scratch space avoids allocating another full copy of both weights.
     """
     if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16:
         raise ValueError(
@@ -147,51 +148,50 @@ def _copy_permuted_expert_to_block_layout(
             out=out,
         )
 
-    w13_rows, w13_cols = w13_weight[0].view(torch.uint8).shape
-    w2_rows, w2_cols = w2_weight[0].view(torch.uint8).shape
-    w13_weights_shuffled_tensor = torch.empty(
-        (num_experts, w13_cols // block_k, w13_rows, block_k),
-        dtype=torch.uint8,
-        device=w13_weight.device,
-    )
-    w2_weights_shuffled_tensor = torch.empty(
-        (num_experts, w2_cols // block_k, w2_rows, block_k),
-        dtype=torch.uint8,
-        device=w2_weight.device,
-    )
-
-    for i in range(num_experts):
-        w13_expert_uint8 = w13_weight[i].view(torch.uint8)
-
-        permute_indices = _maybe_get_cached_w3_w1_permute_indices(
-            cache_permute_indices,
-            w13_expert_uint8,
-            epilogue_tile_m,
-            is_gated_act_gemm=is_gated_act_gemm,
-        )
-        if is_gated_act_gemm:
-            rows = w13_expert_uint8.shape[0]
-            permute_indices = (permute_indices + rows // 2) % rows
-        _copy_permuted_expert_to_block_layout(
-            w13_weights_shuffled_tensor[i],
-            w13_expert_uint8,
-            permute_indices,
+    def _convert_weight_in_place(
+        weight: torch.Tensor,
+        is_w13: bool,
+    ) -> torch.Tensor:
+        rows, cols = weight[0].view(torch.uint8).shape
+        block_layout_shape = (num_experts, cols // block_k, rows, block_k)
+        expert_scratch = torch.empty(
+            block_layout_shape[1:],
+            dtype=torch.uint8,
+            device=weight.device,
         )
 
-        permute_indices = get_w2_permute_indices_with_cache(
-            cache_permute_indices,
-            w2_weight[i].view(torch.uint8),
-            epilogue_tile_m,
-        )
-        _copy_permuted_expert_to_block_layout(
-            w2_weights_shuffled_tensor[i],
-            w2_weight[i].view(torch.uint8),
-            permute_indices,
-        )
+        for i in range(num_experts):
+            expert_uint8 = weight[i].view(torch.uint8)
+            if is_w13:
+                permute_indices = _maybe_get_cached_w3_w1_permute_indices(
+                    cache_permute_indices,
+                    expert_uint8,
+                    epilogue_tile_m,
+                    is_gated_act_gemm=is_gated_act_gemm,
+                )
+                if is_gated_act_gemm:
+                    permute_indices = (
+                        permute_indices + expert_uint8.shape[0] // 2
+                    ) % expert_uint8.shape[0]
+            else:
+                permute_indices = get_w2_permute_indices_with_cache(
+                    cache_permute_indices,
+                    expert_uint8,
+                    epilogue_tile_m,
+                )
+
+            _copy_permuted_expert_to_block_layout(
+                expert_scratch,
+                expert_uint8,
+                permute_indices,
+            )
+            expert_uint8.view(-1).copy_(expert_scratch.view(-1))
+
+        return weight.view(torch.uint8).view(block_layout_shape).view(torch.bfloat16)
 
     return (
-        w13_weights_shuffled_tensor.view(torch.bfloat16),
-        w2_weights_shuffled_tensor.view(torch.bfloat16),
+        _convert_weight_in_place(w13_weight, is_w13=True),
+        _convert_weight_in_place(w2_weight, is_w13=False),
     )
 
 
@@ -210,7 +210,6 @@ def align_fp4_moe_weights_for_fi(
     not satisfied (e.g. with certain tensor-parallel sizes), we pad the
     gate/up and down projection weights along the intermediate dim.
     """
-
     # Current local intermediate size (per partition) is the K dimension of
     # the down projection.
     num_experts, hidden_size, intermediate = w2.shape
@@ -316,7 +315,6 @@ def align_moe_weights_for_fi(
     not satisfied (e.g. with certain tensor-parallel sizes), we pad the
     gate/up and down projection weights along the intermediate dim.
     """
-
     # Current local intermediate size (per partition) is the K dimension of
     # the down projection.
     num_experts, hidden_size, intermediate = w2.shape
@@ -453,8 +451,7 @@ def prepare_fp8_moe_layer_for_fi(
     w2_input_scale: torch.Tensor | None,
     is_trtllm: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Convert Fp8 MoE weights to flashinfer kernel format
+    """Convert Fp8 MoE weights to flashinfer kernel format.
 
     Note that for trtllm we update the model state dict
     with the scale format needed for these kernels.
@@ -462,7 +459,6 @@ def prepare_fp8_moe_layer_for_fi(
     Note that for per-tensor, we update the layer's
     intermediate size if the weights needed padding.
     """
-
     assert hasattr(layer.moe_config, "is_act_and_mul")
     block_quant = (
         hasattr(layer, "weight_block_size") and layer.weight_block_size is not None
diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py
index 199d856e3b78..c9888dc3cb6f 100644
--- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py
@@ -347,12 +347,10 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
     BLOCK_N: tl.constexpr,
 ):
     # TODO(varun) : Add expert_ids so we may early-exit no-op thread blocks.
-    """
-    Each thread block (BLOCK_N) computes [BLOCK_M, GROUP_SIZE] act-mul outputs. Then
+    """Each thread block (BLOCK_N) computes [BLOCK_M, GROUP_SIZE] act-mul outputs. Then
     the thread block quantizes the [BLOCK_M, GROUP_SIZE] block of values and fills
     the outputs tensors at the right positions.
     """
-
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
     N_2 = N // 2
@@ -421,8 +419,7 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
     alpha: float = 1.0,
     beta: float = 0.0,
 ):
-    """
-    Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate
+    """Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate
     (silu: alpha=1, beta=0; swigluoai: alpha, beta from config).
     """
     GROUP_SIZE = group_size
@@ -564,18 +561,23 @@ def per_token_group_quant_fp8(
     """Function to perform per-token-group quantization on an input tensor `x`.
     It converts the tensor values into signed float8 values and returns the
     quantized tensor along with the scaling factor used for quantization.
+
     Args:
         x: The input tensor with ndim >= 2.
         group_size: The group size used for quantization.
         eps: The minimum to avoid dividing zero.
         dtype: The dtype of output tensor. Note that only `torch.float8_e4m3fn`
-        is supported for now.
+            is supported for now.
         column_major_scales: Outputs scales in column major.
         tma_aligned_scales: Outputs scales in TMA-aligned layout.
         out_q: Optional output tensor. If not provided, function will create.
+        use_ue8m0: If True, round each scale down to a power of two (UE8M0
+            format). None defers to the platform default.
+
     Returns:
         tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the
         scaling factor.
+
     """
     if use_ue8m0 is None:
         use_ue8m0 = is_deep_gemm_e8m0_used()
@@ -695,6 +697,7 @@ def per_token_group_quant_fp8_packed_for_deepgemm(
             x_s_packed: Int32 tensor with logical shape
                         [mn, ceil(num_groups_per_row / 4)], laid out with
                         TMA-aligned stride along the packed-K dimension
+
     """
     if use_ue8m0 is None:
         use_ue8m0 = is_deep_gemm_e8m0_used()
@@ -789,7 +792,6 @@ def _w8a8_triton_block_scaled_mm(
     product) on input tensors `A` and `B` with block-wise quantization, and
     store the result in output tensor `C`.
     """
-
     pid = tl.program_id(axis=0)
     num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
@@ -842,14 +844,12 @@ def _w8a8_triton_block_scaled_mm(
 def get_w8a8_block_fp8_configs(
     N: int, K: int, block_n: int, block_k: int
 ) -> dict[int, Any] | None:
-    """
-    Return optimized configurations for the w8a8 block fp8 kernel.
+    """Return optimized configurations for the w8a8 block fp8 kernel.
     The return value will be a dictionary that maps an irregular grid of
     batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the
     kernel on a given batch size bs, the closest batch size in the grid should
     be picked and the associated configuration chosen to invoke the kernel.
     """
-
     # First look up if an optimized configuration is available in the configs
     # directory
     device_name = get_device_name_as_file_name()
@@ -889,18 +889,20 @@ def w8a8_triton_block_scaled_mm(
     quantization.
     It takes two input tensors `A` and `B` with scales `As` and `Bs`.
     The output is returned in the specified `output_dtype`.
+
     Args:
         A: The input tensor, e.g., activation.
         B: The input tensor, e.g., weight.
         As: The per-token-group quantization scale for `A`.
         Bs: The per-block quantization scale for `B`.
         block_size: The block size for per-block quantization. It should
-        be 2-dim, e.g., [128, 128].
-        output_dytpe: The dtype of the returned tensor.
+            be 2-dim, e.g., [128, 128].
+        output_dtype: The dtype of the returned tensor.
+
     Returns:
         torch.Tensor: The result of matmul.
-    """
 
+    """
     _on_gfx1250 = False
     if current_platform.is_rocm():
         from vllm.platforms.rocm import on_gfx1250
@@ -937,14 +939,13 @@ def w8a8_triton_block_scaled_mm(
     assert len(block_size) == 2
     block_n, block_k = block_size[0], block_size[1]
 
-    # Triton cannot currently bind E8M0 scale tensors directly. On ROCm,
-    # DeepSeek-V4 checkpoints store block scales in exponent-only E8M0 format,
-    # so decode them to fp32 before launching the kernel.
-    if current_platform.is_rocm() or current_platform.is_xpu():
-        if As.dtype == torch.float8_e8m0fnu:
-            As = _upcast_e8m0_to_fp32(As).contiguous()
-        if Bs.dtype == torch.float8_e8m0fnu:
-            Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
+    # Triton cannot currently bind E8M0 scale tensors directly. Checkpoints
+    # with exponent-only UE8M0 scales (e.g. DeepSeek-V4) are decoded to fp32
+    # before launching the kernel.
+    if As.dtype == torch.float8_e8m0fnu:
+        As = _upcast_e8m0_to_fp32(As).contiguous()
+    if Bs.dtype == torch.float8_e8m0fnu:
+        Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
 
     assert A.shape[-1] == B.shape[-1]
     assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
@@ -1023,6 +1024,7 @@ def requant_weight_ue8m0_inplace(
             with shape `(..., M // block_size[0], K // block_size[1])`.
         block_size: 2-element iterable `[block_m, block_k]` describing the
             block quantisation granularity.
+
     """
     if weight.numel() == 0:
         return
@@ -1089,8 +1091,6 @@ def deepgemm_post_process_weight_scale_block(
     is_sfa: bool = False,
 ) -> torch.Tensor:
     if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
-        # Scales already in E8M0 from checkpoint; upcast to fp32 and let
-        # DeepGEMM pack the layout expected by the target architecture.
         ws = _upcast_e8m0_to_fp32(ws)
     else:
         assert ws.dtype == torch.float32, (
@@ -1121,12 +1121,7 @@ def deepgemm_post_process_fp8_weight_block(
         f"to be torch.float8_e4m3fn, got {wq.dtype} instead."
     )
 
-    if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
-        # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0
-        # bits as uint8 for MXFP8) - upcast to fp32 and skip requantization
-        # (weights already have power-of-two scales).
-        ws = _upcast_e8m0_to_fp32(ws)
-    else:
+    if ws.dtype not in (torch.float8_e8m0fnu, torch.uint8):
         assert ws.dtype == torch.float32, (
             f"Expected tensor scales dtype to be torch.float32 or "
             f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
@@ -1137,7 +1132,7 @@ def deepgemm_post_process_fp8_weight_block(
     if is_bmm:
         # Reshape 2D weight/scale to 3D for grouped BMM (einsum):
         # wq: (g*r, d) -> (g, r, d)
-        # ws: (g*r/128, d/128) -> (g, r/128, d/128)
+        # Scale dimensions follow quant_block_shape.
         g = bmm_batch_size
         assert wq.ndim == 2 and ws.ndim == 2
         d = wq.size(1)
@@ -1378,7 +1373,7 @@ def process_fp8_weight_tensor_strategy(
     logical_widths: list[int],
     input_scale: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-    """Process weights for tensor-wise quantization strategy."""
+    """Requantize fused shards to one scale and return ``(K, N)`` weight."""
     from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
         normalize_e4m3fn_to_e4m3fnuz,
         requantize_with_max_scale,
@@ -1396,7 +1391,7 @@ def process_fp8_weight_tensor_strategy(
         logical_widths=logical_widths,
     )
 
-    weight = _maybe_pad_fp8_weight(weight)
+    weight = _maybe_pad_fp8_weight(weight).t()
     return weight, weight_scale, input_scale
 
 
@@ -1405,7 +1400,7 @@ def process_fp8_weight_channel_strategy(
     weight_scale: torch.Tensor,
     input_scale: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-    """Process weights for channel-wise quantization strategy."""
+    """Normalize FNUZ if needed and return ``(K, N)`` weight."""
     from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
         normalize_e4m3fn_to_e4m3fnuz,
     )
@@ -1415,14 +1410,14 @@ def process_fp8_weight_channel_strategy(
             weight=weight, weight_scale=weight_scale, input_scale=input_scale
         )
 
-    return weight, weight_scale, input_scale
+    return weight.t(), weight_scale, input_scale
 
 
 def process_fp8_weight_block_strategy(
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Process weights for block-wise quantization strategy."""
+    """Normalize FNUZ if needed and return ``(N, K)`` weight (no transpose)."""
     from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
         normalize_e4m3fn_to_e4m3fnuz,
     )
@@ -1494,7 +1489,6 @@ def process_fp8_input_tensor_strategy_moe(
     enable_eplb: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Process moe input scales for tensor-wise quantization strategy."""
-
     if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale):
         logger.info_once(
             "Found input_scales that are not equal for "
diff --git a/vllm/model_executor/layers/quantization/utils/humming/__init__.py b/vllm/model_executor/layers/quantization/utils/humming/__init__.py
new file mode 100644
index 000000000000..4303f74221f0
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/utils/humming/__init__.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Humming quantization integration."""
+
+from vllm.model_executor.layers.quantization.utils.humming.activation import (
+    get_humming_activation,
+)
+from vllm.model_executor.layers.quantization.utils.humming.linear import (
+    apply_humming_linear,
+    convert_linear_layer_to_humming_standard,
+    get_humming_linear_compute_config,
+    prepare_humming_linear_layer_config,
+)
+from vllm.model_executor.layers.quantization.utils.humming.moe import (
+    HummingMoEQuantConfig,
+    convert_to_humming_moe_kernel_format,
+    get_humming_moe_quant_config,
+    make_humming_moe_kernel,
+    make_humming_moe_quant_config,
+    select_humming_moe_experts,
+)
+from vllm.model_executor.layers.quantization.utils.humming.schema import (
+    check_and_fallback_input_schema,
+    humming_is_layer_skipped,
+    humming_update_schema_hadamard_block_size,
+    input_schema_to_quant_key,
+    quant_key_to_input_schema,
+    resolve_humming_layer_config,
+    weight_schema_to_quant_key,
+)
+
+__all__ = [
+    "check_and_fallback_input_schema",
+    "get_humming_activation",
+    "HummingMoEQuantConfig",
+    "get_humming_moe_quant_config",
+    "humming_is_layer_skipped",
+    "humming_update_schema_hadamard_block_size",
+    "make_humming_moe_quant_config",
+    "apply_humming_linear",
+    "convert_linear_layer_to_humming_standard",
+    "get_humming_linear_compute_config",
+    "prepare_humming_linear_layer_config",
+    "make_humming_moe_kernel",
+    "select_humming_moe_experts",
+    "input_schema_to_quant_key",
+    "quant_key_to_input_schema",
+    "resolve_humming_layer_config",
+    "weight_schema_to_quant_key",
+    "convert_to_humming_moe_kernel_format",
+]
diff --git a/vllm/model_executor/layers/quantization/utils/humming/activation.py b/vllm/model_executor/layers/quantization/utils/humming/activation.py
new file mode 100644
index 000000000000..3942944d2dd9
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/utils/humming/activation.py
@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MoE activation expressions for Humming input processing."""
+
+from dataclasses import replace
+
+from vllm.model_executor.layers.fused_moe.activation import (
+    ApplyMoEActivationConfig,
+    MoEActivation,
+)
+
+_SILU_IMPL = "a / (1.f + expf(-a))"
+_GELU_IMPL = "0.5f * a * (1.f + erff(a * 0.7071067811865476f))"
+_RELU2_IMPL = "fmaxf(a, 0.f) * fmaxf(a, 0.f)"
+_GELU_TANH_IMPL = """[](float a) {
+    const float cubic = 0.044715f * a * a * a;
+    const float inner = 0.7978845608028654f * (a + cubic);
+    return 0.5f * a * (1.f + tanhf(inner));
+}(a)"""
+_CLAMPED_SILU_IMPL = """[](float a, float b) {{
+    const float gate = fminf(a, {limit});
+    const float up = fmaxf(fminf(b, {limit}), -{limit});
+    const float activated_gate = gate / (1.f + expf(-{alpha} * gate));
+    return activated_gate * (up + {beta});
+}}(a, b)"""
+
+# Parameterized expressions are resolved by get_humming_activation().
+_HUMMING_ACTIVATIONS: dict[MoEActivation, dict[str, str]] = {
+    MoEActivation.SILU: {
+        "activation_type": "binary_split",
+        "activation_impl": f"({_SILU_IMPL}) * b",
+    },
+    MoEActivation.GELU: {
+        "activation_type": "binary_split",
+        "activation_impl": f"({_GELU_IMPL}) * b",
+    },
+    MoEActivation.GELU_TANH: {
+        "activation_type": "binary_split",
+        "activation_impl": f"({_GELU_TANH_IMPL}) * b",
+    },
+    MoEActivation.RELU2: {
+        "activation_type": "binary_split",
+        "activation_impl": f"({_RELU2_IMPL}) * b",
+    },
+    MoEActivation.SWIGLUOAI: {
+        "activation_type": "binary_interleaved",
+        "activation_impl": _CLAMPED_SILU_IMPL,
+    },
+    MoEActivation.SITU: {
+        "activation_type": "binary_split",
+        "activation_impl": """[](float a, float b) {{
+            const float gate = {situ_beta} * tanhf(a / {situ_beta});
+            const float activated_gate = gate / (1.f + expf(-a));
+            const float up = {situ_up};
+            return activated_gate * up;
+        }}(a, b)""",
+    },
+    MoEActivation.SWIGLUOAI_UNINTERLEAVE: {
+        "activation_type": "binary_split",
+        "activation_impl": _CLAMPED_SILU_IMPL,
+    },
+    MoEActivation.SWIGLUSTEP: {
+        "activation_type": "binary_split",
+        "activation_impl": """[](float a, float b) {
+            const float activated_gate = a / (1.f + expf(-a));
+            const float gate = fminf(activated_gate, 7.f);
+            const float up = fmaxf(fminf(b, 7.f), -7.f);
+            return gate * up;
+        }(a, b)""",
+    },
+    MoEActivation.SILU_NO_MUL: {
+        "activation_type": "unary",
+        "activation_impl": _SILU_IMPL,
+    },
+    MoEActivation.GELU_NO_MUL: {
+        "activation_type": "unary",
+        "activation_impl": _GELU_IMPL,
+    },
+    MoEActivation.GELU_TANH_NO_MUL: {
+        "activation_type": "unary",
+        "activation_impl": _GELU_TANH_IMPL,
+    },
+    MoEActivation.RELU2_NO_MUL: {
+        "activation_type": "unary",
+        "activation_impl": _RELU2_IMPL,
+    },
+}
+
+
+def get_humming_activation(
+    activation: MoEActivation,
+    config: ApplyMoEActivationConfig | None = None,
+) -> dict[str, str]:
+    config = config or ApplyMoEActivationConfig()
+    use_clamped_silu = False
+    if activation == MoEActivation.SWIGLUOAI:
+        config = ApplyMoEActivationConfig(clamp_limit=7.0, alpha=1.702, beta=1.0)
+        use_clamped_silu = True
+    elif activation == MoEActivation.SILU and config.clamp_limit is not None:
+        config = replace(config, alpha=1.0, beta=0.0)
+        use_clamped_silu = True
+    elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
+        if config.clamp_limit is None:
+            raise ValueError("SWIGLUOAI_UNINTERLEAVE requires clamp_limit")
+        use_clamped_silu = True
+
+    definition = _HUMMING_ACTIVATIONS[activation].copy()
+    if activation == MoEActivation.SITU:
+        if config.activation_situ_beta is None:
+            raise ValueError("SITU requires activation_situ_beta")
+        linear_beta = config.activation_situ_linear_beta
+        situ_up = "b"
+        if linear_beta is not None and linear_beta > 0:
+            situ_up = f"({float(linear_beta)}f) * tanhf(b / ({float(linear_beta)}f))"
+        situ_beta = f"({float(config.activation_situ_beta)}f)"
+        activation_impl = definition["activation_impl"]
+        activation_impl = activation_impl.format(situ_beta=situ_beta, situ_up=situ_up)
+        definition["activation_impl"] = activation_impl
+    elif use_clamped_silu:
+        assert config.clamp_limit is not None
+        limit = f"({float(config.clamp_limit)}f)"
+        alpha = f"({float(config.alpha)}f)"
+        beta = f"({float(config.beta)}f)"
+        activation_impl = _CLAMPED_SILU_IMPL.format(limit=limit, alpha=alpha, beta=beta)
+        definition["activation_impl"] = activation_impl
+    return definition
diff --git a/vllm/model_executor/layers/quantization/utils/humming/linear.py b/vllm/model_executor/layers/quantization/utils/humming/linear.py
new file mode 100644
index 000000000000..0c64645c4395
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/utils/humming/linear.py
@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Prepare and execute Humming linear layers."""
+
+import json
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm import envs
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.layers.quantization.utils.humming.schema import (
+    check_and_fallback_input_schema,
+)
+
+if TYPE_CHECKING:
+    from vllm.utils.humming import BaseInputSchema, LayerConfig
+
+
+def convert_linear_layer_to_humming_standard(
+    layer: LinearBase, name_map: dict[str, str]
+):
+    """Rename/reshape a linear layer's quantized params (the canonical MPLinear
+    layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names
+    and layout humming's weight schema expects (``weight`` / ``weight_scale``)."""
+    for name, checkpoint_name in name_map.items():
+        tensor = getattr(layer, checkpoint_name)
+        delattr(layer, checkpoint_name)
+
+        if name == "weight":
+            input_dim = getattr(tensor, "input_dim", 1)
+            output_dim = getattr(tensor, "output_dim", 0)
+
+            if input_dim == 0 and output_dim == 1:
+                tensor = tensor.transpose(1, 0).contiguous()
+            else:
+                assert output_dim == 0 and input_dim == 1
+
+            tensor = tensor.view(tensor.size(0), -1).view(torch.int32)
+        elif name in ["weight_scale", "zero_point"]:
+            if getattr(tensor, "output_dim", 0) == 1:
+                tensor = tensor.transpose(0, 1).contiguous()
+            if tensor.ndim == 1:
+                tensor = tensor.unsqueeze(1)
+
+            tensor = tensor.view(torch.int32) if name == "zero_point" else tensor
+
+        if isinstance(tensor, torch.nn.Parameter):
+            param = tensor
+        else:
+            param = torch.nn.Parameter(tensor, requires_grad=False)
+
+        setattr(layer, name, param)
+
+
+def prepare_humming_linear_layer_config(
+    layer: LinearBase,
+    quant_config: dict,
+    input_quant_config: dict | None = None,
+    *,
+    input_schema: "BaseInputSchema | None" = None,
+) -> "LayerConfig":
+    from vllm.utils.humming import (
+        BaseInputSchema,
+        BaseWeightSchema,
+        HummingInputSchema,
+        prepare_layer_config,
+        transform_humming_tensors,
+    )
+
+    weight_schema = BaseWeightSchema.from_config(quant_config)
+    if input_schema is None:
+        if input_quant_config is not None:
+            input_schema = BaseInputSchema.from_config(input_quant_config)
+        else:
+            input_schema = HummingInputSchema()
+
+    # ReplicatedLinear has no TP partitioning and so does not set
+    # input_size_per_partition; for it that is just input_size. Use hasattr
+    # rather than getattr's default arg, which is evaluated eagerly and would
+    # raise on layers lacking input_size (e.g. ParallelLMHead).
+    if hasattr(layer, "input_size_per_partition"):
+        input_size_per_partition = layer.input_size_per_partition
+    else:
+        input_size_per_partition = layer.input_size
+    shape_k_stacks = [input_size_per_partition]
+    shape_n_stacks = layer.output_partition_sizes
+
+    # Step 1: convert weight and input schemas to humming standard format
+    source_tensors = dict(layer.named_parameters())
+    weight_schema, tensors = weight_schema.convert_humming(
+        tensors=source_tensors,
+        shape_n_stacks=shape_n_stacks,
+        shape_k_stacks=shape_k_stacks,
+        param_dtype=layer.params_dtype,
+    )
+    input_schema, input_tensors = input_schema.convert_humming(
+        tensors=source_tensors,
+        shape_n_stacks=shape_n_stacks,
+        shape_k_stacks=shape_k_stacks,
+        param_dtype=layer.params_dtype,
+    )
+
+    if (
+        weight_schema.has_zero_point
+        and input_schema.a_dtype is not None
+        and weight_schema.b_dtype.num_bits == input_schema.a_dtype.num_bits
+    ):
+        raise ValueError(
+            "Humming does not support asymmetric (zero_point) weights "
+            "when weight and activation bit widths are equal "
+            f"(W{weight_schema.b_dtype.num_bits}"
+            f"A{input_schema.a_dtype.num_bits}). "
+            "Use symmetric weights or a different activation bit width."
+        )
+
+    tensors.update(input_tensors)
+    layer.weight_schema = weight_schema
+    input_schema = check_and_fallback_input_schema(
+        weight_schema=weight_schema,
+        input_schema=input_schema,
+        param_dtype=layer.params_dtype,
+    )
+    for name in ("input_scale", "input_scale_2"):
+        if name != input_schema.static_tensor_scale_name:
+            tensors.pop(name, None)
+            input_tensors.pop(name, None)
+
+    # Step 2: transform weight (humming standard format) for forwarding.
+    config = prepare_layer_config(
+        shape_n=sum(layer.output_partition_sizes),
+        shape_k=input_size_per_partition,
+        weight_schema=weight_schema,
+        input_schema=input_schema,
+        pad_n_to_multiple=256,
+        pad_k_to_multiple=128,
+        has_bias=layer.has_bias,
+        torch_dtype=layer.params_dtype,
+        device=tensors["weight"].device,
+    )
+    tensors = transform_humming_tensors(config, tensors)
+    tensors.update(input_tensors)
+    if "bias" in tensors:
+        zero_bias = torch.zeros_like(tensors["bias"])
+        layer.register_buffer("zero_bias", zero_bias)
+    for name, _ in list(layer.named_parameters()):
+        delattr(layer, name)
+    for name, tensor in tensors.items():
+        param = torch.nn.Parameter(tensor, requires_grad=False)
+        setattr(layer, name, param)
+
+    return config
+
+
+def get_humming_linear_compute_config() -> str:
+    return json.dumps(
+        {
+            "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
+            "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
+            "gemm_type": "dense",
+        }
+    )
+
+
+def apply_humming_linear(
+    layer: LinearBase,
+    x: torch.Tensor,
+    *,
+    skip_bias_add: bool = False,
+    layer_config: "LayerConfig",
+    compute_config: str,
+    locks: torch.Tensor,
+) -> torch.Tensor:
+    from vllm.utils.humming import humming_forward
+
+    mode = layer_config.input_quant_mode
+    input_scale = getattr(layer, "input_scale", None) if mode.should_quantize else None
+    input_scale_2 = (
+        getattr(layer, "input_scale_2", None) if mode.has_secondary_scale else None
+    )
+    flatten_inputs = x.reshape(-1, x.size(-1))
+    output = humming_forward(
+        layer_config,
+        inputs=flatten_inputs,
+        weight=layer.weight,
+        weight_scale=getattr(layer, "weight_scale", None),
+        zero_point=getattr(layer, "zero_point", None),
+        bias=getattr(layer, "zero_bias" if skip_bias_add else "bias", None),
+        weight_scale_2=getattr(layer, "weight_scale_2", None),
+        input_scale=input_scale,
+        input_scale_2=input_scale_2,
+        hadamard_block_size=layer.weight_schema.hadamard_block_size,
+        locks=locks,
+        compute_config=compute_config,
+    )
+    return output.view(*x.shape[:-1], output.size(-1))
diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming/moe.py
similarity index 50%
rename from vllm/model_executor/layers/quantization/utils/humming_utils.py
rename to vllm/model_executor/layers/quantization/utils/humming/moe.py
index 2a3c1bd4f219..37eb9c26f07d 100644
--- a/vllm/model_executor/layers/quantization/utils/humming_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/humming/moe.py
@@ -1,13 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import json
+"""Configure, prepare weights for, and assemble Humming MoE kernels."""
+
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
-import regex as re
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+import vllm.model_executor.layers.quantization.utils.humming.schema as humming_schema
 from vllm import envs
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.all2all_utils import (
@@ -23,34 +24,19 @@
     HummingGroupedExperts,
     HummingIndexedExperts,
 )
-from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.layers.quantization.utils.humming.schema import (
+    humming_is_layer_skipped,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    FP4_DTYPE,
     FP8_DTYPE,
-    INT4_DTYPE,
-    INT8_DTYPE,
-    MXFP_SCALE_DTYPE,
     GroupShape,
     QuantKey,
-    ScaleDesc,
 )
 from vllm.utils.import_utils import has_humming
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
-    from vllm.utils.humming import (
-        AWQWeightSchema,
-        BaseInputSchema,
-        BaseWeightSchema,
-        CompressedTensorsInputSchema,
-        CompressedTensorsWeightSchema,
-        Fp8WeightSchema,
-        GPTQWeightSchema,
-        HummingInputSchema,
-        HummingWeightSchema,
-        LayerConfig,
-    )
-    from vllm.utils.humming import dtypes as humming_dtypes
+    from vllm.utils.humming import LayerConfig
 
 logger = init_logger(__name__)
 
@@ -59,505 +45,12 @@
 class HummingMoEQuantConfig(FusedMoEQuantConfig):
     w1_humming_config: "LayerConfig"
     w2_humming_config: "LayerConfig"
-
-
-if has_humming():
-    from vllm.utils.humming import dtypes as humming_dtypes
-
-    _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = {
-        humming_dtypes.float4e2m1: FP4_DTYPE,
-        humming_dtypes.float8e4m3: FP8_DTYPE,
-        humming_dtypes.float8e5m2: torch.float8_e5m2,
-        humming_dtypes.int8: torch.int8,
-        humming_dtypes.uint4: INT4_DTYPE,
-        humming_dtypes.uint8: INT8_DTYPE,
-        humming_dtypes.uint2: torch.uint8,
-        humming_dtypes.uint3: torch.uint8,
-        humming_dtypes.uint5: torch.uint8,
-        humming_dtypes.uint6: torch.uint8,
-        humming_dtypes.uint7: torch.uint8,
-    }
-
-    _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = {
-        humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE,
-        humming_dtypes.float8e4m3: FP8_DTYPE,
-        humming_dtypes.float16: torch.float16,
-        humming_dtypes.bfloat16: torch.bfloat16,
-        humming_dtypes.float32: torch.float32,
-    }
-
-
-def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape:
-    """
-    Map humming group sizes to QuantKey GroupShape.
-
-    group_size:   elements per group along K (col); 0 means full dimension.
-    group_size_n: elements per group along N (row); 0 means 1 (per-row).
-
-    GroupShape convention: row = N dim, col = K dim.
-    """
-    if group_size == 0 and group_size_n == 0:
-        return GroupShape.PER_CHANNEL
-
-    row = group_size_n if group_size_n > 0 else 1
-    col = group_size if group_size > 0 else -1
-    return GroupShape(row=row, col=col)
-
-
-# ---- HummingWeightSchema (post-conversion) --------------------------------
-
-
-def _humming_weight_schema_to_quant_key(
-    schema: "HummingWeightSchema",
-) -> QuantKey:
-    from vllm.utils.humming import WeightScaleType
-
-    """Convert a HummingWeightSchema to a QuantKey."""
-    dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype]
-
-    if schema.bs_dtype is not None:
-        scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype]
-    else:
-        scale_dtype = torch.float32
-
-    group_shape = _group_shape(
-        schema.weight_scale_group_size,
-        schema.weight_scale_group_size_n,
-    )
-
-    scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)
-
-    scale2 = None
-    if schema.weight_scale_type == WeightScaleType.GROUP_TENSOR:
-        scale2 = ScaleDesc(
-            dtype=torch.float32,
-            static=True,
-            group_shape=GroupShape.PER_TENSOR,
-        )
-
-    return QuantKey(
-        dtype=dtype,
-        scale=scale,
-        scale2=scale2,
-        symmetric=not schema.has_zero_point,
-    )
-
-
-# ---- Checkpoint-format weight schemas (pre-conversion) --------------------
-
-
-def _fp8_weight_schema_to_quant_key(schema: "Fp8WeightSchema") -> QuantKey:
-    if schema.weight_block_size is not None:
-        gs_n, gs_k = schema.weight_block_size
-        group_shape = GroupShape(row=gs_n, col=gs_k)
-    else:
-        group_shape = GroupShape.PER_CHANNEL
-
-    scale = ScaleDesc(dtype=torch.float32, static=True, group_shape=group_shape)
-    return QuantKey(dtype=FP8_DTYPE, scale=scale, symmetric=True)
-
-
-def _awq_weight_schema_to_quant_key(schema: "AWQWeightSchema") -> QuantKey:
-    group_shape = _group_shape(schema.group_size)
-    scale = ScaleDesc(
-        dtype=torch.float16,
-        static=True,
-        group_shape=group_shape,
-    )
-    return QuantKey(
-        dtype=INT4_DTYPE,
-        scale=scale,
-        symmetric=not schema.zero_point,
-    )
-
-
-def _gptq_weight_schema_to_quant_key(schema: "GPTQWeightSchema") -> QuantKey:
-    group_shape = _group_shape(schema.group_size)
-    scale = ScaleDesc(
-        dtype=torch.float16,
-        static=True,
-        group_shape=group_shape,
-    )
-    return QuantKey(dtype=INT4_DTYPE, scale=scale, symmetric=schema.sym)
-
-
-def _compressed_tensors_weight_schema_to_quant_key(
-    schema: "CompressedTensorsWeightSchema",
-) -> QuantKey:
-    # Determine dtype from format/type/num_bits
-    fmt = schema.format
-    if fmt in ("int-quantized", "float-quantized", "naive-quantized"):
-        dtype = INT8_DTYPE if schema.type == "int" else FP8_DTYPE
-    elif "nvfp4" in fmt or "mxfp4" in fmt:
-        dtype = FP4_DTYPE
-    else:
-        dtype = _HUMMING_TO_QUANT_DTYPE[
-            humming_dtypes.DataType.from_str(f"uint{schema.num_bits}")
-        ]
-
-    # Determine group shape from strategy
-    if schema.strategy in ("group", "tensor_group"):
-        group_shape = _group_shape(schema.group_size or 0)
-    elif schema.strategy == "block" and schema.block_structure is not None:
-        group_shape = GroupShape(
-            row=schema.block_structure[0],
-            col=schema.block_structure[1],
-        )
-    else:
-        group_shape = GroupShape.PER_CHANNEL
-
-    # Determine scale dtype
-    if "mxfp" in fmt:
-        scale_dtype = MXFP_SCALE_DTYPE
-    elif "nvfp4" in fmt:
-        scale_dtype = FP8_DTYPE
-    else:
-        scale_dtype = torch.float32
-
-    scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)
-
-    scale2 = None
-    if "nvfp4" in fmt or schema.strategy == "tensor_group":
-        scale2 = ScaleDesc(
-            dtype=torch.float32,
-            static=True,
-            group_shape=GroupShape.PER_TENSOR,
-        )
-
-    return QuantKey(
-        dtype=dtype,
-        scale=scale,
-        scale2=scale2,
-        symmetric=schema.symmetric,
-    )
-
-
-# ---- Dispatch for any BaseWeightSchema ------------------------------------
-
-
-def weight_schema_to_quant_key(
-    schema: "BaseWeightSchema",
-) -> QuantKey:
-    from vllm.utils.humming import (
-        AWQWeightSchema,
-        BitnetWeightSchema,
-        CompressedTensorsWeightSchema,
-        Fp8WeightSchema,
-        GptOssMxfp4WeightSchema,
-        GPTQWeightSchema,
-        HummingWeightSchema,
-        ModeloptMxfp8WeightSchema,
-        ModeloptNvfp4WeightSchema,
-        Mxfp4WeightSchema,
-    )
-
-    """Convert any BaseWeightSchema to a QuantKey."""
-    if isinstance(schema, HummingWeightSchema):
-        return _humming_weight_schema_to_quant_key(schema)
-
-    # Schemas with fixed QuantKeys
-    if isinstance(schema, (Mxfp4WeightSchema, GptOssMxfp4WeightSchema)):
-        return QuantKey(
-            dtype=FP4_DTYPE,
-            scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)),
-        )
-    if isinstance(schema, ModeloptMxfp8WeightSchema):
-        return QuantKey(
-            dtype=FP8_DTYPE,
-            scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)),
-        )
-    if isinstance(schema, ModeloptNvfp4WeightSchema):
-        return QuantKey(
-            dtype=FP4_DTYPE,
-            scale=ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16)),
-            scale2=ScaleDesc(torch.float32, True, GroupShape.PER_TENSOR),
-        )
-    if isinstance(schema, BitnetWeightSchema):
-        return QuantKey(
-            dtype=torch.uint8,
-            scale=ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL),
-        )
-
-    # Schemas requiring config inspection
-    if isinstance(schema, Fp8WeightSchema):
-        return _fp8_weight_schema_to_quant_key(schema)
-    if isinstance(schema, AWQWeightSchema):
-        return _awq_weight_schema_to_quant_key(schema)
-    if isinstance(schema, GPTQWeightSchema):
-        return _gptq_weight_schema_to_quant_key(schema)
-    if isinstance(schema, CompressedTensorsWeightSchema):
-        return _compressed_tensors_weight_schema_to_quant_key(schema)
-
-    raise TypeError(f"Unsupported weight schema type: {type(schema)}")
-
-
-# ---- HummingInputSchema (post-conversion) ----------------------------------
-
-
-def _humming_input_schema_to_quant_key(
-    schema: "HummingInputSchema",
-) -> QuantKey | None:
-    """Convert a HummingInputSchema to a QuantKey. Returns None if
-    the schema represents unquantized (bf16/fp16) inputs."""
-    if schema.a_dtype is None or schema.a_dtype.num_bits >= 16:
-        return None
-
-    dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype]
-
-    gs = schema.input_scale_group_size
-    group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN
-
-    # Pick the scale dtype the Humming kernel actually consumes. An explicit
-    # input_scale_dtype always wins. Otherwise infer from the grouping: MX
-    # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8)
-    # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a
-    # float32 scale. Getting this right lets a grouped FP8 activation match
-    # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key.
-    if schema.input_scale_dtype is not None:
-        scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype]
-    elif gs == 32:
-        scale_dtype = MXFP_SCALE_DTYPE
-    else:
-        scale_dtype = torch.float32
-
-    scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape)
-
-    return QuantKey(dtype=dtype, scale=scale, symmetric=True)
-
-
-# ---- Checkpoint-format input schemas (pre-conversion) ----------------------
-
-
-def _resolve_input_quant_key(
-    origin_a_dtype: "humming_dtypes.DataType",
-    group_size: int,
-) -> QuantKey | None:
-    from vllm.utils.humming import HummingInputSchema
-
-    """Resolve the actual activation QuantKey after platform fallback."""
-    a_dtype = HummingInputSchema().get_fallback_input_dtype(origin_a_dtype)
-    if a_dtype is None or a_dtype.num_bits >= 16:
-        return None
-
-    dtype = _HUMMING_TO_QUANT_DTYPE[a_dtype]
-    gs = group_size if a_dtype == humming_dtypes.float4e2m1 else 0
-    group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN
-    scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32
-
-    scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape)
-    return QuantKey(dtype=dtype, scale=scale, symmetric=True)
-
-
-def _compressed_tensors_input_schema_to_quant_key(
-    schema: "CompressedTensorsInputSchema",
-) -> QuantKey | None:
-    type_bits_to_dtype = {
-        ("float", 8): humming_dtypes.float8e4m3,
-        ("float", 4): humming_dtypes.float4e2m1,
-        ("int", 8): humming_dtypes.int8,
-        ("int", 4): humming_dtypes.int4,
-    }
-    origin = type_bits_to_dtype.get((schema.type, schema.num_bits))
-    if origin is None:
-        return None
-    return _resolve_input_quant_key(origin, schema.group_size)
-
-
-# ---- Dispatch for any BaseInputSchema -------------------------------------
-
-
-def input_schema_to_quant_key(
-    schema: "BaseInputSchema",
-) -> QuantKey | None:
-    from vllm.utils.humming import (
-        CompressedTensorsInputSchema,
-        Fp8InputSchema,
-        HummingInputSchema,
-        ModeloptNvfp4InputSchema,
-    )
-
-    """Convert any BaseInputSchema to a QuantKey. Returns None if
-    the schema represents unquantized (bf16/fp16) inputs."""
-    if isinstance(schema, HummingInputSchema):
-        return _humming_input_schema_to_quant_key(schema)
-
-    if isinstance(schema, Fp8InputSchema):
-        return _resolve_input_quant_key(humming_dtypes.float8e4m3, 0)
-
-    if isinstance(schema, ModeloptNvfp4InputSchema):
-        return _resolve_input_quant_key(
-            humming_dtypes.float8e4m3,
-            schema.group_size,
-        )
-
-    if isinstance(schema, CompressedTensorsInputSchema):
-        return _compressed_tensors_input_schema_to_quant_key(schema)
-
-    raise TypeError(f"Unsupported input schema type: {type(schema)}")
-
-
-def humming_is_layer_skipped(config: dict[str, Any], prefix: str):
-    if not config:
-        return True
-
-    keys = ["ignored_layers", "ignore", "modules_to_not_convert"]
-    ignored_layers: list[str] = []
-    for key in keys:
-        candidate = config.get(key, []) or []
-        if candidate:
-            ignored_layers = candidate
-            break
-
-    if any(module_name in prefix for module_name in ignored_layers):
-        return True
-    if "lm_head" in prefix:
-        return True
-
-    for regex in config.get("dynamic", {}):
-        if regex[:1] != "-":
-            continue
-        if re.match(regex[2:], prefix):
-            return True
-
-    return False
-
-
-def convert_linear_layer_to_humming_standard(
-    layer: LinearBase, name_map: dict[str, str]
-):
-    """Rename/reshape a linear layer's quantized params (the canonical MPLinear
-    layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names
-    and layout humming's weight schema expects (``weight`` / ``weight_scale``)."""
-    for name, checkpoint_name in name_map.items():
-        tensor = getattr(layer, checkpoint_name)
-        delattr(layer, checkpoint_name)
-
-        if name == "weight":
-            input_dim = getattr(tensor, "input_dim", 1)
-            output_dim = getattr(tensor, "output_dim", 0)
-
-            if input_dim == 0 and output_dim == 1:
-                tensor = tensor.transpose(1, 0).contiguous()
-            else:
-                assert output_dim == 0 and input_dim == 1
-
-            tensor = tensor.view(tensor.size(0), -1).view(torch.int32)
-        elif name in ["weight_scale", "zero_point"]:
-            if getattr(tensor, "output_dim", 0) == 1:
-                tensor = tensor.transpose(0, 1).contiguous()
-            if tensor.ndim == 1:
-                tensor = tensor.unsqueeze(1)
-
-            tensor = tensor.view(torch.int32) if name == "zero_point" else tensor
-
-        if isinstance(tensor, torch.nn.Parameter):
-            param = tensor
-        else:
-            param = torch.nn.Parameter(tensor, requires_grad=False)
-
-        setattr(layer, name, param)
-
-
-def prepare_humming_linear_layer_config(
-    layer: LinearBase,
-    quant_config: dict,
-    input_quant_config: dict | None = None,
-) -> "LayerConfig":
-    from vllm.utils.humming import (
-        BaseInputSchema,
-        BaseWeightSchema,
-        HummingInputSchema,
-        prepare_layer_config,
-        transform_humming_tensors,
-    )
-
-    weight_schema = BaseWeightSchema.from_config(quant_config)
-    if input_quant_config is not None:
-        input_schema = BaseInputSchema.from_config(input_quant_config)
-    else:
-        input_schema = HummingInputSchema()
-
-    # ReplicatedLinear has no TP partitioning and so does not set
-    # input_size_per_partition; for it that is just input_size. Use hasattr
-    # rather than getattr's default arg, which is evaluated eagerly and would
-    # raise on layers lacking input_size (e.g. ParallelLMHead).
-    if hasattr(layer, "input_size_per_partition"):
-        input_size_per_partition = layer.input_size_per_partition
-    else:
-        input_size_per_partition = layer.input_size
-    shape_k_stacks = [input_size_per_partition]
-    shape_n_stacks = layer.output_partition_sizes
-
-    # Step 1: convert weight and input schemas to humming standard format
-    weight_schema, tensors = weight_schema.convert_humming(
-        tensors=dict(layer.named_parameters()),
-        shape_n_stacks=shape_n_stacks,
-        shape_k_stacks=shape_k_stacks,
-        param_dtype=layer.params_dtype,
-    )
-    input_schema, _ = input_schema.convert_humming(
-        tensors={},
-        shape_n_stacks=shape_n_stacks,
-        shape_k_stacks=shape_k_stacks,
-        param_dtype=layer.params_dtype,
-    )
-
-    layer.weight_schema = weight_schema
-
-    # Step 2: transform weight (humming standard format) for forwarding.
-    config = prepare_layer_config(
-        shape_n=sum(layer.output_partition_sizes),
-        shape_k=input_size_per_partition,
-        weight_schema=weight_schema,
-        input_schema=input_schema,
-        pad_n_to_multiple=256,
-        pad_k_to_multiple=128,
-        has_bias=layer.has_bias,
-        torch_dtype=layer.params_dtype,
-    )
-    tensors = transform_humming_tensors(config, tensors)
-    for name, _ in list(layer.named_parameters()):
-        delattr(layer, name)
-    for name, tensor in tensors.items():
-        param = torch.nn.Parameter(tensor, requires_grad=False)
-        setattr(layer, name, param)
-
-    return config
-
-
-def get_humming_linear_compute_config() -> str:
-    return json.dumps(
-        {
-            "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
-            "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
-            "gemm_type": "dense",
-        }
-    )
-
-
-def apply_humming_linear(
-    layer: LinearBase,
-    x: torch.Tensor,
-    *,
-    layer_config: "LayerConfig",
-    compute_config: str,
-    locks: torch.Tensor,
-) -> torch.Tensor:
-    from vllm.utils.humming import humming_forward
-
-    flatten_inputs = x.reshape(-1, x.size(-1))
-    output = humming_forward(
-        layer_config,
-        inputs=flatten_inputs,
-        weight=layer.weight,
-        weight_scale=getattr(layer, "weight_scale", None),
-        zero_point=getattr(layer, "zero_point", None),
-        bias=getattr(layer, "bias", None),
-        weight_scale_2=getattr(layer, "weight_scale_2", None),
-        locks=locks,
-        compute_config=compute_config,
-    )
-    return output.view(*x.shape[:-1], output.size(-1))
+    w1_hadamard_block_size: int = 1
+    w2_hadamard_block_size: int = 1
+    w1_input_scale: torch.Tensor | None = None
+    w2_input_scale: torch.Tensor | None = None
+    w1_input_scale_2: torch.Tensor | None = None
+    w2_input_scale_2: torch.Tensor | None = None
 
 
 def make_humming_moe_quant_config(
@@ -649,8 +142,12 @@ def get_humming_moe_quant_config(
         q_dtype is not None
         and q_dtype.startswith("float8")
         and input_scale_group_size == 128
+        and weight_schema.hadamard_block_size <= 1
+        and input_schema.input_quant_mode in (None, "dynamic_group")
     ):
-        q_dtype = _HUMMING_TO_QUANT_DTYPE.get(input_schema.a_dtype, FP8_DTYPE)
+        q_dtype = humming_schema._HUMMING_TO_QUANT_DTYPE.get(
+            input_schema.a_dtype, FP8_DTYPE
+        )
         activation_group_shape = GroupShape(row=1, col=input_scale_group_size)
 
     weight_scale_group_size = weight_schema.weight_scale_group_size
@@ -666,7 +163,7 @@ def get_humming_moe_quant_config(
     else:
         weight_group_shape = GroupShape(row=weight_scale_group_size, col=1)
 
-    return make_humming_moe_quant_config(
+    config = make_humming_moe_quant_config(
         quant_dtype=q_dtype,
         weight_dtype=str(weight_schema.b_dtype),
         weight_group_shape=weight_group_shape,
@@ -684,6 +181,17 @@ def get_humming_moe_quant_config(
         gemm1_clamp_limit=gemm1_clamp_limit,
         humming_configs=humming_configs,
     )
+    for prefix, sublayer in (("w1", "w13"), ("w2", "w2")):
+        setattr(
+            config,
+            f"{prefix}_hadamard_block_size",
+            layer.weight_schemas[sublayer].hadamard_block_size,
+        )
+        for name in ("input_scale", "input_scale_2"):
+            key = f"{prefix}_{name}"
+            value = getattr(layer, f"{sublayer}_{name}", None)
+            setattr(config, key, value)
+    return config
 
 
 def select_humming_moe_experts(
@@ -691,11 +199,9 @@ def select_humming_moe_experts(
     weight_key: QuantKey | None,
     activation_key: QuantKey | None,
 ) -> type[mk.FusedMoEExperts] | None:
-    """
-    Select the primary Humming MoE Experts class
+    """Select the primary Humming MoE Experts class
     Note: Shape-specific fallbacks may still occur at runtime.
     """
-
     if not has_humming():
         return None
 
@@ -809,14 +315,14 @@ def _replace_layer_parameters(
     tensors: dict[str, torch.Tensor],
     preserve_bias: bool = False,
 ) -> None:
-    """
-    Replace layer parameters for a sublayer with new tensors.
+    """Replace layer parameters for a sublayer with new tensors.
 
     Args:
         layer: The RoutedExperts layer
         sublayer_name: Name of the sublayer (e.g., "w13", "w2")
         tensors: Dict of parameter name to tensor
         preserve_bias: If True, don't delete bias parameters
+
     """
     # Delete old parameters
     for name, _ in list(layer.named_parameters()):
@@ -843,18 +349,12 @@ def _convert_sublayer_to_humming(
     num_experts: int,
     param_dtype: torch.dtype,
 ) -> tuple[Any, Any]:
-    """
-    Convert a sublayer's weights from checkpoint format to Humming format.
+    """Convert a sublayer's weights from checkpoint format to Humming format.
 
     Returns:
         Tuple of (converted_weight_schema, converted_input_schema)
-    """
-    from vllm.utils.humming import HummingWeightSchema
-
-    if isinstance(weight_schema, HummingWeightSchema):
-        # Already in Humming format
-        return weight_schema, input_schema
 
+    """
     tensors = _extract_sublayer_tensors(layer, sublayer_name)
 
     shape_k_stacks = [shape_k]
@@ -870,13 +370,14 @@ def _convert_sublayer_to_humming(
         num_experts=num_experts,
     )
 
-    converted_input_schema, _ = input_schema.convert_humming(
-        tensors=converted_tensors,
+    converted_input_schema, input_tensors = input_schema.convert_humming(
+        tensors=tensors,
         shape_n_stacks=shape_n_stacks,
         shape_k_stacks=shape_k_stacks,
         param_dtype=param_dtype,
         num_experts=num_experts,
     )
+    converted_tensors.update(input_tensors)
 
     _replace_layer_parameters(layer, sublayer_name, converted_tensors)
 
@@ -910,11 +411,16 @@ def _prepare_and_transform_sublayer(
         has_bias=has_bias,
         num_experts=num_experts,
         torch_dtype=param_dtype,
+        device=getattr(layer, f"{sublayer_name}_weight").device,
     )
-    tensors = transform_humming_tensors(
-        config,
-        _extract_sublayer_tensors(layer, sublayer_name),
-    )
+    source_tensors = _extract_sublayer_tensors(layer, sublayer_name)
+    for name in ("input_scale", "input_scale_2"):
+        if name != input_schema.static_tensor_scale_name:
+            source_tensors.pop(name, None)
+    tensors = transform_humming_tensors(config, source_tensors)
+    for name in ("input_scale", "input_scale_2"):
+        if name in source_tensors:
+            tensors[name] = source_tensors[name]
     _replace_layer_parameters(layer, sublayer_name, tensors)
     return config
 
@@ -930,9 +436,9 @@ def _process_single_sublayer(
     num_experts: int,
     param_dtype: torch.dtype,
     force_weight_schema: Any | None = None,
+    allow_input_schema_fallback: bool = True,
 ) -> tuple[Any, Any, "LayerConfig"]:
-    """
-    Process a single sublayer: convert, optionally requant, prepare, and transform.
+    """Process a single sublayer: convert, optionally requant, prepare, and transform.
 
     This combines the common logic from convert_to_humming_moe_kernel_format
     for processing a single sublayer.
@@ -948,9 +454,11 @@ def _process_single_sublayer(
         num_experts: Number of experts
         param_dtype: Parameter data type
         force_weight_schema: Optional schema to force requantization to
+        allow_input_schema_fallback: Whether incompatible input schemas may be replaced.
 
     Returns:
         Tuple of the final weight schema, input schema, and Humming layer config.
+
     """
     from vllm.utils.humming import HummingWeightSchema
 
@@ -969,19 +477,28 @@ def _process_single_sublayer(
     # Step 2: Force requant if needed
     assert isinstance(current_weight_schema, HummingWeightSchema)
     if force_weight_schema is not None and current_weight_schema != force_weight_schema:
-        tensors = _extract_sublayer_tensors(layer, sublayer_name)
+        source_tensors = _extract_sublayer_tensors(layer, sublayer_name)
 
         tensors = current_weight_schema.requant_tensors(
-            tensors=tensors,
+            tensors=source_tensors,
             target_weight_schema=force_weight_schema,
             param_dtype=param_dtype,
         )
+        name = current_input_schema.static_tensor_scale_name
+        if name is not None:
+            tensors[name] = source_tensors[name]
 
         current_weight_schema = force_weight_schema
         _replace_layer_parameters(layer, sublayer_name, tensors, preserve_bias=True)
         del tensors
 
     # Step 3: Prepare layer metadata and transform weights
+    current_input_schema = humming_schema.check_and_fallback_input_schema(
+        weight_schema=current_weight_schema,
+        input_schema=current_input_schema,
+        param_dtype=param_dtype,
+        allow_fallback=allow_input_schema_fallback,
+    )
     config = _prepare_and_transform_sublayer(
         layer=layer,
         sublayer_name=sublayer_name,
@@ -1004,9 +521,9 @@ def convert_to_humming_moe_kernel_format(
     weight_schema: Any | None = None,
     input_schema: Any | None = None,
     force_weight_schema: Any | None = None,
+    allow_input_schema_fallback: bool = True,
 ) -> dict[str, "LayerConfig"]:
-    """
-    Convert MoE weights from checkpoint format to Humming kernel format.
+    """Convert MoE weights from checkpoint format to Humming kernel format.
 
     This function processes weights for each sublayer (w13, w2) by:
     1. Converting from checkpoint format to humming format if needed
@@ -1027,13 +544,14 @@ def convert_to_humming_moe_kernel_format(
         input_schema: Optional initial input quantization schema.
                      If None, built from quant_config or env vars.
         force_weight_schema: Optional schema to force requantization to
+        allow_input_schema_fallback: Whether incompatible input schemas may be replaced.
 
     Side effects:
         - Modifies layer parameters in place
         - Sets layer.weight_schemas and layer.input_schemas
         - Sets layer.humming_configs for quant config construction
-    """
 
+    """
     # Build schemas from quant_config if not provided
     has_bias = layer.moe_config.has_bias
     num_experts = layer.moe_config.num_local_experts
@@ -1045,20 +563,23 @@ def convert_to_humming_moe_kernel_format(
                 "Must provide either weight_schema/input_schema or quant_config"
             )
 
-        from vllm.model_executor.layers.quantization.utils.humming_utils import (
-            humming_is_layer_skipped,
-        )
         from vllm.utils.humming import BaseWeightSchema, HummingInputSchema
 
         if weight_schema is None:
             weight_schema = BaseWeightSchema.from_config(quant_config)
 
         if input_schema is None:
-            input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}
+            input_quant_config = (envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}).copy()
             if humming_is_layer_skipped(input_quant_config, layer.layer_name):
                 input_schema = HummingInputSchema()
             else:
                 # TODO: read input_quant_config from quant_config
+                input_quant_config = humming_schema.resolve_humming_layer_config(
+                    input_quant_config, layer.layer_name
+                )
+                allow_input_schema_fallback = input_quant_config.pop(
+                    "allow_fallback", False
+                )
                 input_schema = HummingInputSchema.from_config(input_quant_config)
 
     # Build sublayer configs from layer properties if not provided
@@ -1093,6 +614,7 @@ def convert_to_humming_moe_kernel_format(
                 num_experts=num_experts,
                 param_dtype=param_dtype,
                 force_weight_schema=force_weight_schema,
+                allow_input_schema_fallback=allow_input_schema_fallback,
             )
         )
 
diff --git a/vllm/model_executor/layers/quantization/utils/humming/schema.py b/vllm/model_executor/layers/quantization/utils/humming/schema.py
new file mode 100644
index 000000000000..c84e8a370d22
--- /dev/null
+++ b/vllm/model_executor/layers/quantization/utils/humming/schema.py
@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Map Humming schemas and handle shared checkpoint quantization settings."""
+
+import dataclasses
+from typing import TYPE_CHECKING, Any
+
+import regex as re
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    FP4_DTYPE,
+    FP8_DTYPE,
+    INT4_DTYPE,
+    INT8_DTYPE,
+    MXFP_SCALE_DTYPE,
+    GroupShape,
+    QuantKey,
+    ScaleDesc,
+)
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+from vllm.utils.import_utils import has_humming
+
+if TYPE_CHECKING:
+    from vllm.utils.humming import (
+        BaseInputSchema,
+        BaseWeightSchema,
+        HummingInputSchema,
+        HummingWeightSchema,
+    )
+    from vllm.utils.humming import dtypes as humming_dtypes
+
+
+def humming_update_schema_hadamard_block_size(
+    weight_schema: "HummingWeightSchema",
+    input_schema: "HummingInputSchema",
+    shape_k: int,
+) -> "HummingWeightSchema":
+    assert shape_k > 0
+    block_size = 32
+    weight_group_size = weight_schema.weight_scale_group_size
+    input_group_size = input_schema.input_scale_group_size
+    for group_size in (weight_group_size, input_group_size):
+        if group_size is not None and group_size > 0:
+            block_size = min(block_size, group_size)
+    block_size = 1 << (block_size.bit_length() - 1)
+    while shape_k % block_size:
+        block_size //= 2
+    return dataclasses.replace(weight_schema, hadamard_block_size=block_size)
+
+
+if has_humming():
+    from vllm.utils.humming import dtypes as humming_dtypes
+
+    _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = {
+        humming_dtypes.float4e0m3: FP4_DTYPE,
+        humming_dtypes.float4e2m1: FP4_DTYPE,
+        humming_dtypes.float6e2m3: scalar_types.float6_e2m3f,
+        humming_dtypes.float6e3m2: scalar_types.float6_e3m2f,
+        humming_dtypes.float8e3m4: FP8_DTYPE,
+        humming_dtypes.float8e4m3: FP8_DTYPE,
+        humming_dtypes.float8e5m2: torch.float8_e5m2,
+        humming_dtypes.int8: torch.int8,
+        humming_dtypes.uint4: INT4_DTYPE,
+        humming_dtypes.uint8: INT8_DTYPE,
+        humming_dtypes.uint2: torch.uint8,
+        humming_dtypes.uint3: torch.uint8,
+        humming_dtypes.uint5: torch.uint8,
+        humming_dtypes.uint6: torch.uint8,
+        humming_dtypes.uint7: torch.uint8,
+    }
+
+    _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = {
+        humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE,
+        humming_dtypes.float8e4m3: FP8_DTYPE,
+        humming_dtypes.float16: torch.float16,
+        humming_dtypes.bfloat16: torch.bfloat16,
+        humming_dtypes.float32: torch.float32,
+    }
+
+
+logger = init_logger(__name__)
+
+
+def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape:
+    """Map humming group sizes to QuantKey GroupShape.
+
+    group_size:   elements per group along K (col); 0 means full dimension.
+    group_size_n: elements per group along N (row); 0 means 1 (per-row).
+
+    GroupShape convention: row = N dim, col = K dim.
+    """
+    if group_size == 0 and group_size_n == 0:
+        return GroupShape.PER_CHANNEL
+
+    row = group_size_n if group_size_n > 0 else 1
+    col = group_size if group_size > 0 else -1
+    return GroupShape(row=row, col=col)
+
+
+def weight_schema_to_quant_key(
+    schema: "BaseWeightSchema",
+    param_dtype: torch.dtype | None = None,
+) -> QuantKey:
+    from vllm.utils.humming import WeightScale2Type, WeightScaleType
+
+    if param_dtype is None:
+        param_dtype = torch.get_default_dtype()
+        if param_dtype not in [torch.float16, torch.bfloat16]:
+            param_dtype = torch.bfloat16
+            if not current_platform.has_device_capability(80):
+                param_dtype = torch.float16
+
+    schema = schema.to_humming_schema(param_dtype)
+    dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype]
+
+    if schema.bs_dtype is not None:
+        scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype]
+    else:
+        scale_dtype = torch.float32
+
+    group_shape = _group_shape(
+        schema.weight_scale_group_size,
+        schema.weight_scale_group_size_n,
+    )
+    if schema.weight_scale_type == WeightScaleType.TENSOR:
+        group_shape = GroupShape.PER_TENSOR
+
+    scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)
+
+    scale2 = None
+    if schema.weight_scale_2_type == WeightScale2Type.TENSOR:
+        scale2 = ScaleDesc(
+            dtype=torch.float32,
+            static=True,
+            group_shape=GroupShape.PER_TENSOR,
+        )
+
+    return QuantKey(
+        dtype=dtype,
+        scale=scale,
+        scale2=scale2,
+        symmetric=not schema.has_zero_point,
+    )
+
+
+def input_schema_to_quant_key(
+    schema: "BaseInputSchema",
+    param_dtype: torch.dtype | None = None,
+) -> QuantKey | None:
+    from vllm.utils.humming import InputQuantizationMode
+
+    if param_dtype is None:
+        param_dtype = torch.get_default_dtype()
+        if param_dtype not in [torch.float16, torch.bfloat16]:
+            param_dtype = torch.bfloat16
+            if not current_platform.has_device_capability(80):
+                param_dtype = torch.float16
+
+    schema = schema.to_humming_schema(param_dtype)
+    if schema.a_dtype is None or schema.a_dtype.num_bits >= 16:
+        return None
+
+    mode = schema.input_quant_mode
+    if mode == InputQuantizationMode.Disabled:
+        return None
+    dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype]
+
+    gs = schema.input_scale_group_size
+    group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN
+
+    # Pick the scale dtype the Humming kernel actually consumes. An explicit
+    # input_scale_dtype always wins. Otherwise infer from the grouping: MX
+    # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8)
+    # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a
+    # float32 scale. Getting this right lets a grouped FP8 activation match
+    # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key.
+    if schema.input_scale_dtype is not None:
+        scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype]
+    elif gs == 16:
+        scale_dtype = FP8_DTYPE
+    elif gs == 32:
+        scale_dtype = MXFP_SCALE_DTYPE
+    else:
+        scale_dtype = torch.float32
+
+    if mode == InputQuantizationMode.StaticTensor:
+        group_shape = GroupShape.PER_TENSOR
+    scale = ScaleDesc(
+        dtype=scale_dtype,
+        static=mode == InputQuantizationMode.StaticTensor,
+        group_shape=group_shape,
+    )
+    scale2 = None
+    has_scale2 = (
+        InputQuantizationMode.StaticTensorDynamicGroup,
+        InputQuantizationMode.DynamicGroupToken,
+    )
+    if mode in has_scale2:
+        static = mode == InputQuantizationMode.StaticTensorDynamicGroup
+        scale2 = ScaleDesc(
+            torch.float32,
+            static,
+            GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN,
+        )
+
+    return QuantKey(dtype=dtype, scale=scale, scale2=scale2, symmetric=True)
+
+
+def quant_key_to_input_schema(key: QuantKey | None) -> "HummingInputSchema":
+    from vllm.utils.humming import HummingInputSchema, InputQuantizationMode
+
+    if key is None:
+        return HummingInputSchema(input_quant_mode=InputQuantizationMode.Disabled)
+
+    if not key.symmetric:
+        raise ValueError("Humming input quantization must be symmetric")
+
+    quant_dtypes = {value: dtype for dtype, value in _HUMMING_TO_QUANT_DTYPE.items()}
+    quant_dtypes[FP4_DTYPE] = humming_dtypes.float4e2m1
+    quant_dtypes[FP8_DTYPE] = humming_dtypes.float8e4m3
+    scale_dtypes = {value: dtype for dtype, value in _HUMMING_TO_SCALE_DTYPE.items()}
+    if key.dtype not in quant_dtypes or key.scale.dtype not in scale_dtypes:
+        raise ValueError(f"Unsupported Humming input or scale dtype: {key}")
+
+    scale = key.scale
+    group_size = 0
+    if scale.static and scale.group_shape.is_per_tensor() and key.scale2 is None:
+        mode = InputQuantizationMode.StaticTensor
+    elif not scale.static and scale.group_shape.is_per_token() and key.scale2 is None:
+        mode = InputQuantizationMode.DynamicToken
+    elif not scale.static and scale.group_shape.is_per_group():
+        group_size = scale.group_shape.col
+        scale2 = key.scale2
+        if scale2 is None:
+            mode = InputQuantizationMode.DynamicGroup
+        elif scale2.dtype != torch.float32:
+            raise ValueError("Humming secondary input scales must be float32")
+        elif scale2.static and scale2.group_shape.is_per_tensor():
+            mode = InputQuantizationMode.StaticTensorDynamicGroup
+        elif not scale2.static and scale2.group_shape.is_per_token():
+            mode = InputQuantizationMode.DynamicGroupToken
+        else:
+            raise ValueError(f"Unsupported Humming secondary input scale: {scale2}")
+    else:
+        raise ValueError(f"Unsupported Humming input scale: {key}")
+
+    if group_size == 0 and scale.dtype != torch.float32:
+        raise ValueError("Humming tensor and token input scales must be float32")
+
+    return HummingInputSchema(
+        a_dtype=quant_dtypes[key.dtype],
+        input_scale_group_size=group_size,
+        input_scale_dtype=scale_dtypes[scale.dtype],
+        input_quant_mode=mode,
+    )
+
+
+def check_and_fallback_input_schema(
+    weight_schema: "BaseWeightSchema",
+    input_schema: "BaseInputSchema",
+    param_dtype: torch.dtype | None = None,
+    allow_fallback: bool = True,
+) -> "HummingInputSchema":
+    from vllm.utils.humming import HummingInputSchema, InputQuantizationMode
+
+    capability = current_platform.get_device_capability()
+    assert capability is not None
+    sm_version = capability.to_int()
+    if param_dtype is None:
+        param_dtype = torch.get_default_dtype()
+        if param_dtype not in (torch.float16, torch.bfloat16):
+            param_dtype = torch.bfloat16 if sm_version >= 80 else torch.float16
+    assert param_dtype in (torch.float16, torch.bfloat16)
+
+    dtype_fallback_order_map = {
+        humming_dtypes.float8e5m2: (
+            humming_dtypes.float8e4m3,
+            humming_dtypes.float8e3m4,
+            humming_dtypes.int8,
+        ),
+        humming_dtypes.float8e4m3: (
+            humming_dtypes.float8e3m4,
+            humming_dtypes.int8,
+        ),
+        humming_dtypes.float8e3m4: (
+            humming_dtypes.float8e4m3,
+            humming_dtypes.int8,
+        ),
+        humming_dtypes.float4e2m1: (
+            humming_dtypes.float8e4m3,
+            humming_dtypes.float8e3m4,
+            humming_dtypes.int8,
+        ),
+        humming_dtypes.float4e0m3: (
+            humming_dtypes.int4,
+            humming_dtypes.float8e4m3,
+            humming_dtypes.float8e3m4,
+            humming_dtypes.int8,
+        ),
+        humming_dtypes.int8: (
+            humming_dtypes.float8e4m3,
+            humming_dtypes.float8e3m4,
+        ),
+        humming_dtypes.int4: (
+            humming_dtypes.int8,
+            humming_dtypes.float8e4m3,
+            humming_dtypes.float8e3m4,
+        ),
+    }
+
+    weight_schema = weight_schema.to_humming_schema(param_dtype)
+    input_schema = input_schema.to_humming_schema(param_dtype)
+    fp6_dtypes = (humming_dtypes.float6e2m3, humming_dtypes.float6e3m2)
+    if input_schema.a_dtype in fp6_dtypes:
+        input_schema.a_dtype = humming_dtypes.float8e4m3
+    a_dtype = input_schema.a_dtype
+    input_bits = a_dtype.num_bits if a_dtype is not None else 16
+    weight_group_size = weight_schema.weight_scale_group_size
+    input_group_size = input_schema.input_scale_group_size
+    input_scale_dtype = input_schema.input_scale_dtype
+
+    def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool:
+        if dtype is None or dtype.num_bits == 16:
+            return False
+        is_int4_deprecated = dtype == humming_dtypes.int4 and sm_version >= 90
+        is_int8_deprecated = dtype == humming_dtypes.int8 and 103 <= sm_version < 110
+        # Humming has not yet implemented high-performance FP8 computation
+        # on SM10X and SM11X. Therefore, it is temporarily considered deprecated.
+        # TODO: removed this after Humming has high-performance FP8 computation
+        is_fp8_deprecated = dtype.num_bits < 16 and 100 <= sm_version < 120
+        return is_int4_deprecated or is_int8_deprecated or is_fp8_deprecated
+
+    if input_schema.is_compatible_with(weight_schema, param_dtype):
+        if not allow_fallback:
+            if is_deprecated(a_dtype):
+                msg = f"Humming {a_dtype} activation is slow on SM{sm_version} now."
+                logger.warning_once(msg)
+            return input_schema
+
+        if a_dtype is None or a_dtype.num_bits == 16:
+            return input_schema
+
+        if not is_deprecated(a_dtype):
+            is_mxfp8 = (
+                a_dtype.is_floating_point_type
+                and a_dtype.num_bits == 8
+                and input_group_size == 32
+                and input_scale_dtype == humming_dtypes.float8e8m0
+                and sm_version >= 120
+            )
+            is_groupwise = input_group_size > 0 or weight_group_size > 0
+            if input_bits == 8 and is_groupwise and not is_mxfp8:
+                # prefer wna16 for groupwise weight or groupwise input
+                return HummingInputSchema()
+
+            return input_schema
+
+    if allow_fallback:
+        fallback_dtypes = (
+            a_dtype,
+            *dtype_fallback_order_map.get(a_dtype, ()),
+            humming_dtypes.DataType.from_any(param_dtype),
+        )
+        for dtype in fallback_dtypes:
+            if dtype is None or is_deprecated(dtype):
+                continue
+
+            group_size = 0
+            scale_dtype = humming_dtypes.float32
+            quant_mode = InputQuantizationMode.DynamicToken
+            if dtype.num_bits == 16:
+                scale_dtype = None
+                quant_mode = InputQuantizationMode.Disabled
+            elif dtype in (humming_dtypes.float8e4m3, humming_dtypes.float8e3m4):
+                is_channelwise_weight = weight_schema.weight_scale_group_size == 0
+                weight_group_size = weight_schema.weight_scale_group_size
+                is_e8m0_scale = weight_schema.bs_dtype == humming_dtypes.float8e8m0
+                is_mx_weight = weight_group_size == 32 and is_e8m0_scale
+                if sm_version >= 120 and (is_channelwise_weight or is_mx_weight):
+                    group_size = 32
+                    scale_dtype = humming_dtypes.float8e8m0
+                    quant_mode = InputQuantizationMode.DynamicGroup
+
+            candidate = HummingInputSchema(
+                a_dtype=dtype,
+                input_scale_group_size=group_size,
+                input_scale_dtype=scale_dtype,
+                input_quant_mode=quant_mode,
+            )
+            if candidate.is_compatible_with(weight_schema, param_dtype):
+                return candidate
+
+    raise ValueError(
+        f"No compatible Humming input schema for {input_schema} with "
+        f"weight schema {weight_schema} on SM{sm_version} "
+        f"({allow_fallback=})"
+    )
+
+
+def resolve_humming_layer_config(config: dict[str, Any], prefix: str) -> dict[str, Any]:
+    layer_config = config.copy()
+    dynamic = config.get("dynamic", {})
+    if isinstance(dynamic, dict):
+        for pattern, overrides in dynamic.items():
+            if pattern.startswith("+") and re.match(pattern[2:], prefix):
+                layer_config.update(overrides)
+                break
+    return layer_config
+
+
+def humming_is_layer_skipped(config: dict[str, Any], prefix: str):
+    if not config:
+        return True
+
+    keys = ["ignored_layers", "ignore", "modules_to_not_convert"]
+    ignored_layers: list[str] = []
+    for key in keys:
+        candidate = config.get(key, []) or []
+        if candidate:
+            ignored_layers = candidate
+            break
+
+    if any(module_name in prefix for module_name in ignored_layers):
+        return True
+    if "lm_head" in prefix:
+        return True
+
+    for regex in config.get("dynamic", {}):
+        if regex[:1] != "-":
+            continue
+        if re.match(regex[2:], prefix):
+            return True
+
+    return False
diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py
index db95d3588a2a..1180e2df3619 100644
--- a/vllm/model_executor/layers/quantization/utils/int8_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py
@@ -178,6 +178,7 @@ def per_token_group_quant_int8(
     Returns:
         tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the
             scaling factor for quantization.
+
     """
     assert x.shape[-1] % group_size == 0, (
         "the last dimension of `x` cannot be divisible by `group_size`"
diff --git a/vllm/model_executor/layers/quantization/utils/machete_utils.py b/vllm/model_executor/layers/quantization/utils/machete_utils.py
index 7bf71ca99c50..b0c3573404c9 100644
--- a/vllm/model_executor/layers/quantization/utils/machete_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/machete_utils.py
@@ -17,8 +17,7 @@ def query_machete_supported_quant_types(zero_points: bool) -> list[ScalarType]:
 
 
 def query_machete_supported_group_sizes(act_type: torch.dtype) -> list[int]:
-    """
-    Queries the supported group sizes for Machete based on the activation type.
+    """Queries the supported group sizes for Machete based on the activation type.
 
     Args:
         act_type: The activation data type (torch.float16, torch.bfloat16).
@@ -27,6 +26,7 @@ def query_machete_supported_group_sizes(act_type: torch.dtype) -> list[int]:
         A list of supported group sizes. The group size must
         be divisible by `TileShapeK = 128 * 8 // num_bits(act_type)`.
         -1 indicates per-channel quantization.
+
     """
     if act_type in [torch.float16, torch.bfloat16]:
         return [-1, 64, 128]
diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py
index 22cf18f4a5a6..1ee52a17802b 100644
--- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py
@@ -8,7 +8,6 @@
 import torch
 
 import vllm.envs as envs
-from vllm import _custom_ops as ops
 from vllm.distributed.utils import verify_group_size_divides_partition
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import RoutedExperts
@@ -23,6 +22,7 @@
 from vllm.scalar_type import ScalarType, scalar_types
 from vllm.utils.math_utils import round_up
 from vllm.utils.platform_utils import num_compute_units
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from .quant_utils import pack_cols, unpack_cols
 
@@ -32,6 +32,7 @@
 GPTQ_MARLIN_MIN_THREAD_N = 64
 GPTQ_MARLIN_MIN_THREAD_K = 128
 GPTQ_MARLIN_MAX_PARALLEL = 16
+MARLIN_MAX_BLOCKS_PER_SM = 4
 
 MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
 
@@ -397,8 +398,7 @@ def check_moe_marlin_supports_layer(
 
 
 def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tensor):
-    """
-    Given Marlin packed weight matrices w1_packed, and w2_packed,
+    """Given Marlin packed weight matrices w1_packed, and w2_packed,
     return the MoE intermediate size N
     """
     marlin_tile_size = 16
@@ -408,29 +408,96 @@ def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tenso
 def marlin_make_workspace_new(
     device: torch.device,
     max_blocks_per_sm: int = 1,
-    existing: torch.Tensor | None = None,
 ) -> torch.Tensor:
     # In the new marlin kernel, we use the num of threadblocks as workspace
     # size. The num of threadblocks is sms_count * max_blocks_per_sm.
     sms = num_compute_units(device.index)
-    size = sms * max_blocks_per_sm
-    # On weight reload, reuse the existing storage so the workspace address
-    # captured by CUDA graphs stays valid.
-    if existing is not None:
-        if (
-            existing.device != device
-            or existing.dtype != torch.int
-            or existing.numel() != size
-        ):
-            raise ValueError(
-                f"Existing Marlin workspace is incompatible "
-                f"(device={existing.device}, dtype={existing.dtype}, "
-                f"numel={existing.numel()}; expected device={device}, "
-                f"dtype={torch.int}, numel={size}). Reload must reuse the "
-                f"workspace storage captured by CUDA graphs."
-            )
-        return existing.zero_()
-    return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False)
+    return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device)
+
+
+def get_marlin_workspace(device: torch.device) -> torch.Tensor:
+    """Get zero-initialized locks for the current stream, ubatch, and lane.
+
+    Marlin resets its locks after each call, so ordered calls can reuse them.
+    """
+    from vllm.utils.torch_utils import current_stream
+    from vllm.v1.worker.workspace import (
+        current_workspace_manager,
+        is_workspace_manager_initialized,
+    )
+
+    if not is_workspace_manager_initialized():
+        return marlin_make_workspace_new(device, MARLIN_MAX_BLOCKS_PER_SM)
+
+    return current_workspace_manager().get_persistent(
+        ("marlin_locks", current_stream()),
+        (num_compute_units(device.index) * MARLIN_MAX_BLOCKS_PER_SM,),
+        torch.int,
+        zero_init=True,
+    )
+
+
+def _marlin_gemm(
+    a: torch.Tensor,
+    b_q_weight: torch.Tensor,
+    b_bias: torch.Tensor | None,
+    b_scales: torch.Tensor,
+    a_scales: torch.Tensor | None,
+    global_scale: torch.Tensor | None,
+    b_zeros: torch.Tensor | None,
+    workspace: torch.Tensor | None,
+    b_q_type_id: int,
+    size_m: int,
+    size_n: int,
+    size_k: int,
+    use_atomic_add: bool = False,
+    use_fp32_reduce: bool = False,
+    is_zp_float: bool = False,
+) -> torch.Tensor:
+    if workspace is None:
+        workspace = get_marlin_workspace(a.device)
+    return torch.ops._C.marlin_gemm(
+        a,
+        None,
+        b_q_weight,
+        b_bias,
+        b_scales,
+        a_scales,
+        global_scale,
+        b_zeros,
+        workspace,
+        b_q_type_id,
+        size_m,
+        size_n,
+        size_k,
+        use_atomic_add,
+        use_fp32_reduce,
+        is_zp_float,
+    )
+
+
+def _marlin_gemm_fake(
+    a: torch.Tensor,
+    b_q_weight: torch.Tensor,
+    b_bias: torch.Tensor | None,
+    b_scales: torch.Tensor,
+    a_scales: torch.Tensor | None,
+    global_scale: torch.Tensor | None,
+    b_zeros: torch.Tensor | None,
+    workspace: torch.Tensor | None,
+    b_q_type_id: int,
+    size_m: int,
+    size_n: int,
+    size_k: int,
+    use_atomic_add: bool = False,
+    use_fp32_reduce: bool = False,
+    is_zp_float: bool = False,
+) -> torch.Tensor:
+    dtype = a.dtype if a.dtype in (torch.float16, torch.bfloat16) else b_scales.dtype
+    return torch.empty((size_m, size_n), device=a.device, dtype=dtype)
+
+
+direct_register_custom_op("marlin_gemm", _marlin_gemm, fake_impl=_marlin_gemm_fake)
 
 
 def marlin_repeat_scales_on_all_ranks(group_size: int, is_row_parallel: bool) -> bool:
@@ -687,7 +754,7 @@ def apply_gptq_marlin_linear(
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
     weight_zp: torch.Tensor,
-    workspace: torch.Tensor,
+    workspace: torch.Tensor | None,
     wtype: ScalarType,
     output_size_per_partition: int,
     input_size_per_partition: int,
@@ -724,9 +791,8 @@ def apply_gptq_marlin_linear(
 
         reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
 
-    output = ops.marlin_gemm(
+    output = torch.ops.vllm.marlin_gemm(
         reshaped_x,
-        None,
         weight,
         bias,
         weight_scale,
@@ -734,7 +800,7 @@ def apply_gptq_marlin_linear(
         None,
         weight_zp,
         workspace,
-        wtype,
+        wtype.id,
         size_m=reshaped_x.shape[0],
         size_n=padded_n,
         size_k=padded_k,
diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
index 3b6990101ed7..7fb5621f2f9b 100644
--- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
@@ -10,7 +10,6 @@
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     USE_FP32_REDUCE_DEFAULT,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
     marlin_pad_dim,
     marlin_pad_qweight,
     marlin_pad_scales,
@@ -42,7 +41,6 @@ def _nvfp4_compute_scale_factor(
     """Compute the power-of-2 scale_factor needed so that all non-zero
     values in marlin_scales * 2^7 are >= 2 after rescaling.
     Returns a Python float (power of 2, >= 1.0)."""
-
     # Since half has a smaller dynamic range compared to bfloat16,
     # no rescaling is applied here if active dtype is half.
     if a_dtype is not None and a_dtype == torch.half:
@@ -76,9 +74,12 @@ def nvfp4_marlin_process_scales(
             given value is used directly. The caller is responsible for
             dividing ``global_scale`` by the returned ``scale_factor`` to
             preserve numerical correctness.
+        a_dtype: dtype of the activations the scales will be used with. None
+            keeps the dtype of ``marlin_scales``.
 
     Returns:
         A tuple of (processed_scales, scale_factor).
+
     """
     if not (marlin_scales >= 0).all():
         logger.warning_once(
@@ -159,7 +160,7 @@ def apply_fp4_marlin_linear(
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
     weight_global_scale: torch.Tensor | None,
-    workspace: torch.Tensor,
+    workspace: torch.Tensor | None,
     size_n: int,
     size_k: int,
     bias: torch.Tensor | None = None,
@@ -194,9 +195,8 @@ def apply_fp4_marlin_linear(
 
         inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn)
 
-    output = ops.marlin_gemm(
+    output = torch.ops.vllm.marlin_gemm(
         a=inputs,
-        c=None,
         b_q_weight=weight,
         b_bias=bias,
         b_scales=weight_scale,
@@ -204,7 +204,7 @@ def apply_fp4_marlin_linear(
         global_scale=weight_global_scale,
         b_zeros=None,
         workspace=workspace,
-        b_q_type=scalar_types.float4_e2m1f,
+        b_q_type_id=scalar_types.float4_e2m1f.id,
         size_m=reshaped_x.size(0),
         size_n=padded_n,
         size_k=padded_k,
@@ -235,13 +235,6 @@ def prepare_fp4_layer_for_marlin(
 
     assert layer.weight.shape == (part_size_n, part_size_k // 2)
 
-    device = layer.weight.device
-
-    # WORKSPACE
-    layer.workspace = marlin_make_workspace_new(
-        device, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT
     # Repack weights to marlin format
     qweight = layer.weight.view(torch.int32).T.contiguous()
@@ -387,15 +380,9 @@ def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor:
             return x
         return torch.nn.functional.pad(x, (0, (padded_N - N) // packing))
 
-    device = w13.device
     param_dtype = layer.params_dtype
     is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
 
-    # WORKSPACE
-    layer.workspace = marlin_make_workspace_new(
-        device, 4, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT
     # Repack weights to marlin format
     def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
@@ -479,12 +466,7 @@ def prepare_moe_fp4_layer_for_marlin(
     k = layer.moe_config.hidden_dim
     n = layer.moe_config.intermediate_size_per_partition
 
-    # WORKSPACE
-    device = layer.w13_weight.device
     param_dtype = layer.params_dtype
-    layer.workspace = marlin_make_workspace_new(
-        device, 4, existing=getattr(layer, "workspace", None)
-    )
     is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
 
     # WEIGHT
diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
index 56841cba17c3..257c5bf44505 100644
--- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
@@ -9,7 +9,6 @@
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     USE_FP32_REDUCE_DEFAULT,
     get_marlin_input_dtype,
-    marlin_make_workspace_new,
     marlin_moe_padded_intermediate,
     marlin_pad_dim,
     marlin_pad_qweight,
@@ -50,7 +49,7 @@ def apply_fp8_marlin_linear(
     input: torch.Tensor,
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
-    workspace: torch.Tensor,
+    workspace: torch.Tensor | None,
     size_n: int,
     size_k: int,
     bias: torch.Tensor | None,
@@ -80,9 +79,8 @@ def apply_fp8_marlin_linear(
         # inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn)
         raise RuntimeError("Marlin W8A8 is not supported.")
 
-    output = ops.marlin_gemm(
+    output = torch.ops.vllm.marlin_gemm(
         a=inputs,
-        c=None,
         b_q_weight=weight,
         b_bias=bias,
         b_scales=weight_scale,
@@ -90,7 +88,7 @@ def apply_fp8_marlin_linear(
         global_scale=None,
         b_zeros=None,
         workspace=workspace,
-        b_q_type=scalar_types.float8_e4m3fn,
+        b_q_type_id=scalar_types.float8_e4m3fn.id,
         size_m=reshaped_x.size(0),
         size_n=padded_n,
         size_k=padded_k,
@@ -127,13 +125,6 @@ def prepare_fp8_layer_for_marlin(
     else:
         assert layer.weight.shape == (part_size_n, part_size_k)
 
-    device = layer.weight.device
-
-    # WORKSPACE
-    layer.workspace = marlin_make_workspace_new(
-        device, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT
     # Repack weights to marlin format
     qweight = pack_fp8_to_int32(layer.weight, size_k_first)
@@ -241,14 +232,7 @@ def prepare_fp8_moe_layer_for_marlin(
     w13_weight_scale: torch.Tensor,
     w2_weight_scale: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Shuffle weights and scales into marlin format.
-
-    Note that this function has the side effect of adding a `workspace`
-    attribute to the layer. This `workspace` does not need to be
-    registered as a Parameter as it is not used during weight reloading.
-    """
-
+    """Shuffle weights and scales into marlin format."""
     logger.warning_once(
         "Your GPU does not have native support for FP8 computation but "
         "FP8 quantization is being used. Weight-only FP8 compression will "
@@ -274,12 +258,6 @@ def prepare_fp8_moe_layer_for_marlin(
         w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n)
         w2_weight = _moe_pad_last(w2_weight, n, padded_n)
 
-    # WORKSPACE
-    device = layer.w13_weight.device
-    layer.workspace = marlin_make_workspace_new(
-        device, 4, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT
     # Repack weights to marlin format
     def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor:
@@ -372,9 +350,7 @@ def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor:
 def pack_fp8_to_int32(
     fp8_tensor: torch.Tensor, size_k_first: bool = True
 ) -> torch.Tensor:
-    """
-    Repack FP8 weights to gptq format (packed int32 elements)
-    """
+    """Repack FP8 weights to gptq format (packed int32 elements)."""
     assert fp8_tensor.dtype == torch.float8_e4m3fn
     assert fp8_tensor.ndim == 2
 
@@ -400,7 +376,7 @@ def apply_mxfp8_marlin_linear(
     input: torch.Tensor,
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
-    workspace: torch.Tensor,
+    workspace: torch.Tensor | None,
     size_n: int,
     size_k: int,
     bias: torch.Tensor | None = None,
@@ -420,9 +396,8 @@ def apply_mxfp8_marlin_linear(
         dtype=input.dtype,
     )
 
-    output = ops.marlin_gemm(
+    output = torch.ops.vllm.marlin_gemm(
         a=reshaped_x,
-        c=None,
         b_q_weight=weight,
         b_bias=bias,
         b_scales=weight_scale,
@@ -430,7 +405,7 @@ def apply_mxfp8_marlin_linear(
         global_scale=None,
         b_zeros=None,
         workspace=workspace,
-        b_q_type=scalar_types.float8_e4m3fn,
+        b_q_type_id=scalar_types.float8_e4m3fn.id,
         size_m=reshaped_x.size(0),
         size_n=padded_n,
         size_k=padded_k,
@@ -455,13 +430,6 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
     group_size = 32  # MX standard block size
     padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size)
 
-    device = layer.weight.device
-
-    # WORKSPACE
-    layer.workspace = marlin_make_workspace_new(
-        device, existing=getattr(layer, "workspace", None)
-    )
-
     # WEIGHT - repack FP8 weights to Marlin format
     qweight = pack_fp8_to_int32(layer.weight, size_k_first=False)
     qweight = qweight.T.contiguous()
@@ -516,7 +484,7 @@ def prepare_mxfp8_moe_layer_for_marlin(
     """Repack MXFP8 MoE weights and scales into Marlin kernel format.
 
     Args:
-        layer: MoE layer (used to read params_dtype and attach workspace).
+        layer: MoE layer.
         w13: [E, 2*N, K] float8_e4m3fn weights.
         w2:  [E, K, N] float8_e4m3fn weights.
         w13_scale: [E, 2*N, K//32] uint8 e8m0 scales.
@@ -524,6 +492,7 @@ def prepare_mxfp8_moe_layer_for_marlin(
 
     Returns:
         (w13, w2, w13_scale, w2_scale) in Marlin format.
+
     """
     group_size = 32
     e = w13.shape[0]
@@ -540,13 +509,8 @@ def prepare_mxfp8_moe_layer_for_marlin(
         n = padded_n
     w13_n = w13.shape[1]
 
-    device = w13.device
     param_dtype = torch.get_default_dtype()
 
-    layer.workspace = marlin_make_workspace_new(
-        device, 4, existing=getattr(layer, "workspace", None)
-    )
-
     def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
         if "w13" in name:
             size_n, size_k = w13_n, k
diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py
index d4a7c563b8f3..5f8cb9f709df 100644
--- a/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py
+++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_test.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Utility functions used for tests and benchmarks"""
+"""Utility functions used for tests and benchmarks."""
 
 import numpy as np
 import torch
diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
index f992c7e24699..ad79414e2565 100644
--- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
@@ -7,7 +7,7 @@
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
 from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
 
 logger = init_logger(__name__)
@@ -35,10 +35,32 @@ def should_use_cdna4_mx_scale_swizzle() -> bool:
     return on_gfx950() and get_tensor_model_parallel_world_size() <= 2
 
 
+def weight_mx_scale(precision_config):
+    """mxfp4 weight scale from a triton_kernels PrecisionConfig: 3.8 exposes it
+    as b_mx_scale, 3.5.1/3.6 as weight_scale.
+    """
+    if get_triton_kernels_version() == "3.8":
+        return precision_config.b_mx_scale
+    return precision_config.weight_scale
+
+
+def mx_scale_kwargs(scale):
+    """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size,
+    3.5.1/3.6 use weight_scale.
+    """
+    if get_triton_kernels_version() == "3.8":
+        return {"b_mx_scale": scale, "b_microblock_size": 32}
+    return {"weight_scale": scale}
+
+
 def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
-    """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
+    """Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel."""
     assert has_triton_kernels()
-    import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
+    if get_triton_kernels_version() == "3.8":
+        # 3.8: matmul_ogs_details -> matmul_details
+        import triton_kernels.matmul_details.opt_flags as opt_flags
+    else:
+        import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
     from triton_kernels.numerics import InFlexData
     from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
     from triton_kernels.tensor_details import layout
@@ -59,6 +81,17 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
         )
         value_layout = StridedLayout
         scale_layout = StridedLayout
+    elif current_platform.is_rocm() and get_triton_kernels_version() == "3.8":
+        # StridedLayout(-2) is triton_kernels' default CDNA mxfp4 weight layout
+        # (make_default_matmul_mxfp4_w_layout); 3.8 wants layout instances.
+        from triton_kernels.tensor_details.layout import CDNA4MXScaleLayout
+
+        value_layout = StridedLayout(-2)
+        scale_layout = (
+            CDNA4MXScaleLayout()
+            if should_use_cdna4_mx_scale_swizzle()
+            else StridedLayout()
+        )
     elif current_platform.is_rocm():
         value_layout = StridedLayout
         if should_use_cdna4_mx_scale_swizzle():
@@ -393,8 +426,7 @@ def downcast_to_mxfp(
     BLOCK_OUT_DIM: int = 128,
     BLOCK_QUANT_DIM: int = 32,
 ):
-    """
-    Convert the src weights to MXFP4. The src weight is quantized along the
+    """Convert the src weights to MXFP4. The src weight is quantized along the
     axis dimension into packed e2m1 values (torch.uint8, two values per
     byte), so the size of that dimension in the output is half of the
     logical (unpacked) size.
diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
index 0c9324a5490a..bfcc69e4b1de 100644
--- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
@@ -126,8 +126,14 @@ def _kernel(
         amax = tl.maximum(tl.max(tl.abs(x), axis=1), TINY)  # [BLOCK_M]
         sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0
         sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
-        descale = tl.exp2(sb - 127.0)
-        xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty)
+        # Scale by the reciprocal rather than dividing by ``exp2(sb - 127)``:
+        # sb == 0 (an all-zero / denormal block) makes that divisor 2**-127,
+        # which is subnormal in fp32 and flushes to zero on CDNA, turning the
+        # block's zeros into 0/0 == NaN. The reciprocal is normal for every
+        # reachable sb, and both forms are exact powers of two, so nothing
+        # else changes.
+        rescale = tl.exp2(127.0 - sb)
+        xq = (x * rescale[:, None]).to(xq_ptr.dtype.element_ty)
         tl.store(
             xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk,
             xq,
diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
index a5272dd9cb19..a66d481159fa 100644
--- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
@@ -657,8 +657,7 @@ def ref_nvfp4_quant(x, global_scale, block_size):
 def ref_nvfp4_quant_dequant(
     x: torch.Tensor, global_scale: torch.Tensor, block_size: int
 ) -> torch.Tensor:
-    """
-    NVFP4 quantize-dequantize operation.
+    """NVFP4 quantize-dequantize operation.
 
     `global_scale` is expected to have a single element.
     """
diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py
index 74e35dce566d..c446f646a18a 100644
--- a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py
@@ -11,18 +11,15 @@
 
 
 def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor:
-    """
-    Pad and block-interleave the FP4 block-scales so that they match the data
+    """Pad and block-interleave the FP4 block-scales so that they match the data
     layout expected by the CUTLASS / FlashInfer kernels.
 
-    Parameters
-    ----------
-    scale: torch.Tensor
+    Args:
+        scale: The FP4 block-scale tensor to swizzle.
+
+    Returns:
+        The swizzled tensor with the same logical shape as `scale`.
 
-    Returns
-    -------
-    torch.Tensor
-        The swizzled tensor with the same logical shape as *scale*.
     """
     assert scale.dtype == torch.float8_e4m3fn, (
         "swizzle_blockscale expects the input tensor to be in "
@@ -65,8 +62,7 @@ def pad_nvfp4_weight_for_cutlass(
     weight: torch.Tensor,
     alignment: int = 32,
 ) -> tuple[torch.Tensor, int]:
-    """
-    Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy
+    """Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy
     the alignment constraints required by CUTLASS / FlashInfer FP4 kernels.
 
     CUTLASS FP4 kernel requires both K and N matrix dimensions to be divisible
@@ -99,8 +95,7 @@ def pad_nvfp4_weight_for_cutlass(
 
 
 def nvfp4_weight_padding_bytes(layer: torch.nn.Module) -> int:
-    """
-    Compute K-dim padding (packed bytes) of an NVFP4 weight,
+    """Compute K-dim padding (packed bytes) of an NVFP4 weight,
     Aligned with `pad_nvfp4_weight_for_cutlass` above.
     """
     return layer.weight.shape[1] - layer.input_size_per_partition // 2
@@ -110,8 +105,7 @@ def pad_nvfp4_activation_for_cutlass(
     x_fp4: torch.Tensor,
     weights_padding_bytes: int,
 ) -> torch.Tensor:
-    """
-    Pad packed FP4 activations to match the K-dimension padding applied to weights.
+    """Pad packed FP4 activations to match the K-dimension padding applied to weights.
     The padding is in bytes (tensor dimension), not FP4 elements.
     """
     if weights_padding_bytes > 0:
@@ -123,9 +117,7 @@ def slice_nvfp4_output(
     out: torch.Tensor,
     output_size: int,
 ) -> torch.Tensor:
-    """
-    Slice the output tensor to remove padding in N dimension if weight was padded.
-    """
+    """Slice the output tensor to remove padding in N dimension if weight was padded."""
     if out.shape[-1] != output_size:
         return out[..., :output_size].contiguous()
     return out
diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py
index f96146e73615..4c0a8a3ff721 100644
--- a/vllm/model_executor/layers/quantization/utils/quant_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""This file is used for /tests and /benchmarks"""
+"""This file is used for /tests and /benchmarks."""
 
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
@@ -109,8 +109,7 @@ class _GroupShape(NamedTuple):
 
 
 class GroupShape(_GroupShape):
-    """
-    This class describes the quantization group shape.
+    """This class describes the quantization group shape.
     It includes static members for common shapes (per-tensor, per-token).
     """
 
@@ -139,8 +138,7 @@ def is_per_group(self) -> bool:
 
 @dataclass(frozen=True)
 class ScaleDesc:
-    """
-    Class for describing a single quantization scaling factor.
+    """Class for describing a single quantization scaling factor.
     dtype: data type of the scale
     static: static scale if True, dynamic if False
     group_shape: group shape of the scale
@@ -166,8 +164,7 @@ def __str__(self):
 
 @dataclass(frozen=True)
 class QuantKey:
-    """
-    Class for identifying the type of quantization.
+    """Class for identifying the type of quantization.
     dtype: quantized data type
     scale: scale descriptor
     scale2: second-level scale descriptor
@@ -211,6 +208,9 @@ def __str__(self):
 kNvfp4Dynamic = QuantKey(
     FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kStaticTensorScale
 )
+kNvfp4DynamicToken = QuantKey(
+    FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kDynamicTokenScale
+)
 
 kNvfp4StaticGroupScale = ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16))
 kNvfp4Static = QuantKey(
@@ -378,8 +378,7 @@ def prep_scale_for_group_broadcast(
     x: torch.Tensor,
     group_shape: GroupShape | None,
 ) -> torch.Tensor:
-    """
-    Prepare the input quantization scale for group broadcasting.
+    """Prepare the input quantization scale for group broadcasting.
 
     Args:
         scale: The scale tensor (scalar or 1D).
@@ -388,6 +387,7 @@ def prep_scale_for_group_broadcast(
 
     Returns:
         scale reshaped for correct broadcasting.
+
     """
     if scale.numel() == 1:
         # For per-tensor quant, keep the scale as a scalar (not reshaped to (1, 1)).
@@ -432,13 +432,13 @@ def scaled_quantize(
     quant_dtype: torch.dtype,
     compute_dtype: torch.dtype | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Args:
-        x: Input tensor to quantize
-        group_shape: Shape of quantization groups
-        quant_dtype: Target quantized dtype (e.g., torch.float8_e4m3fn)
-        compute_dtype: Optional dtype for intermediate computations.
-            If None, uses input dtype. Use torch.float32 for higher precision.
+    """Args:
+    x: Input tensor to quantize
+    group_shape: Shape of quantization groups
+    quant_dtype: Target quantized dtype (e.g., torch.float8_e4m3fn)
+    compute_dtype: Optional dtype for intermediate computations.
+        If None, uses input dtype. Use torch.float32 for higher precision.
+
     """
     group_shape = _normalize_quant_group_shape(x, group_shape)
     assert quant_dtype.is_floating_point, (
@@ -506,7 +506,7 @@ def get_attribute_fallback(obj, attributes: list[str]):
 def get_and_maybe_dequant_weights(
     layer: "LinearBase", out_dtype: torch.dtype = torch.float32
 ):
-    """Return layer's unquantized weights in [out, in] layout"""
+    """Return layer's unquantized weights in [out, in] layout."""
     from vllm.model_executor.layers.linear import UnquantizedLinearMethod
     from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
     from vllm.model_executor.layers.quantization.online.fp8 import (
@@ -924,8 +924,7 @@ def awq_pack(
 def convert_bf16_scales_to_fp8(
     quant_fp8: Callable, scales: torch.Tensor
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Convert a BF16 scale tensor into the pair of (fp8_scales, channel_scales)
+    """Convert a BF16 scale tensor into the pair of (fp8_scales, channel_scales)
     expected by W4A8 GEMM kernels.
     """
     assert scales.is_contiguous(), (
@@ -949,9 +948,7 @@ def convert_bf16_scales_to_fp8(
 
 
 def convert_packed_uint4b8_to_signed_int4_inplace(t: torch.Tensor) -> torch.Tensor:
-    """
-    Convert int4b8 (packed to int32) to signed int4
-    """
+    """Convert int4b8 (packed to int32) to signed int4."""
     assert t.is_cuda, "tensor must be on gpu"
     assert t.dtype == torch.int32, f"expected int32 packed weights but got {t.dtype}"
 
diff --git a/vllm/model_executor/layers/resampler.py b/vllm/model_executor/layers/resampler.py
index c9fa8054625e..320508de76de 100644
--- a/vllm/model_executor/layers/resampler.py
+++ b/vllm/model_executor/layers/resampler.py
@@ -26,8 +26,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
-Shared resampler perceiver network used in multimodal models and
+"""Shared resampler perceiver network used in multimodal models and
 related helpers for sincos positional embeddings.
 
 Example models: Qwen (Qwen-VL), MiniCPM-V 2.0
@@ -76,8 +75,7 @@ def get_abs_pos(abs_pos: torch.Tensor, tgt_size: torch.Tensor | int) -> torch.Te
 def get_1d_sincos_pos_embed_from_grid(
     embed_dim: int, pos: np.ndarray, version: tuple[int, int] = (2, 0)
 ) -> torch.Tensor:
-    """
-    embed_dim: output dimension for each position
+    """embed_dim: output dimension for each position
     pos: a list of positions to be encoded: size (M,) / (H, W)
     out: (M, D) / (H, W, D)
     """
@@ -126,8 +124,7 @@ def get_2d_sincos_pos_embed(
     cls_token: bool = False,
     version: tuple[int, int] = (2, 0),
 ) -> torch.Tensor:
-    """
-    grid_size: int of the grid height and width
+    """grid_size: int of the grid height and width
     return:
     pos_embed: [grid_size*grid_size, embed_dim] or
                 [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
@@ -154,8 +151,7 @@ def get_2d_sincos_pos_embed(
 
 
 class BaseResampler(nn.Module):
-    """
-    A 2D perceiver-resampler network with one cross attention layers by
+    """A 2D perceiver-resampler network with one cross attention layers by
         (grid_size**2) learnable queries and 2d sincos pos_emb.
     Outputs:
         A tensor with the shape of (grid_size**2, embed_dim)
diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py
index ab06448957fb..e23a5128b8a1 100644
--- a/vllm/model_executor/layers/rotary_embedding/__init__.py
+++ b/vllm/model_executor/layers/rotary_embedding/__init__.py
@@ -24,7 +24,6 @@
 from .mrope_interleaved import MRotaryEmbeddingInterleaved
 from .ntk_scaling_rope import NTKScalingRotaryEmbedding
 from .phi3_long_rope_scaled_rope import Phi3LongRoPEScaledRotaryEmbedding
-from .telechat3_scaling_rope import TeleChat3RoPEScaledRotaryEmbedding
 from .yarn_scaling_rope import YaRNScalingRotaryEmbedding
 
 _ROPE_DICT: dict[tuple[Any, ...], RotaryEmbedding] = {}
@@ -246,16 +245,15 @@ def get_rope(
             for k, v in rope_parameters.items()
             if k
             in (
-                "extrapolation_factor",
-                "attn_factor",
                 "beta_fast",
                 "beta_slow",
-                "apply_yarn_scaling",
+                "mscale",
+                "mscale_all_dim",
+                "attention_factor",
                 "truncate",
             )
         }
         if "mrope_section" in rope_parameters:
-            extra_kwargs.pop("apply_yarn_scaling", None)
             rotary_emb = MRotaryEmbedding(
                 head_size,
                 rotary_dim,
@@ -288,8 +286,6 @@ def get_rope(
             for k, v in rope_parameters.items()
             if k
             in (
-                "extrapolation_factor",
-                "attn_factor",
                 "beta_fast",
                 "beta_slow",
                 "mscale",
@@ -346,36 +342,6 @@ def get_rope(
             )
         else:
             raise ValueError("Pangu mrope lacks necessary parameters.")
-    elif scaling_type == "telechat3-yarn":
-        scaling_factor = rope_parameters["factor"]
-        if "original_max_position_embeddings" in rope_parameters:
-            original_max_position = rope_parameters["original_max_position_embeddings"]
-            scaling_factor = max_position / original_max_position
-        else:
-            original_max_position = max_position
-        extra_kwargs = {
-            k: v
-            for k, v in rope_parameters.items()
-            if k
-            in (
-                "extrapolation_factor",
-                "attn_factor",
-                "beta_fast",
-                "beta_slow",
-                "mscale",
-                "mscale_all_dim",
-            )
-        }
-        rotary_emb = TeleChat3RoPEScaledRotaryEmbedding(
-            head_size,
-            rotary_dim,
-            original_max_position,
-            base,
-            is_neox_style,
-            scaling_factor,
-            dtype,
-            **extra_kwargs,
-        )
     else:
         raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
     _ROPE_DICT[key] = rotary_emb
diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py
index 0a392812b6bc..b96ab9cb6642 100644
--- a/vllm/model_executor/layers/rotary_embedding/common.py
+++ b/vllm/model_executor/layers/rotary_embedding/common.py
@@ -70,10 +70,10 @@ def yarn_linear_ramp_mask(
     return ramp_func
 
 
-def yarn_get_mscale(scale: float = 1) -> float:
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
     if scale <= 1:
         return 1.0
-    return 0.1 * math.log(scale) + 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
 
 
 def _flashinfer_rotary_embedding(
@@ -138,14 +138,14 @@ def forward_static(
         is_neox_style: bool = True,
         enable_fp32_compute: bool = False,
     ) -> torch.Tensor:
-        """
-        Args:
-            x: [batch_size (optional), seq_len, num_heads, head_size]
-            cos: [seq_len, head_size // 2]
-            sin: [seq_len, head_size // 2]
-            is_neox_style: Whether to use the Neox-style or GPT-J-style.
-            enable_fp32_compute: Temporarily convert x, cos, sin to FP32 dtype
-                                 for higher accuracy.
+        """Args:
+        x: [batch_size (optional), seq_len, num_heads, head_size]
+        cos: [seq_len, head_size // 2]
+        sin: [seq_len, head_size // 2]
+        is_neox_style: Whether to use the Neox-style or GPT-J-style.
+        enable_fp32_compute: Temporarily convert x, cos, sin to FP32 dtype
+                             for higher accuracy.
+
         """
         origin_dtype = x.dtype
         if enable_fp32_compute:
@@ -237,6 +237,18 @@ def forward_cuda(
         output = self._post_process(output, origin_shape, origin_dtype)
         return output
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        import vllm._xpu_ops  # noqa: F401 registers torch.ops.vllm.xpu_apply_rotary_emb
+
+        x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin)
+        output = torch.ops.vllm.xpu_apply_rotary_emb(x, cos, sin, self.is_neox_style)
+        return self._post_process(output, origin_shape, origin_dtype)
+
     def forward_hip(
         self,
         x: torch.Tensor,
diff --git a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
index bfaf81ad0077..46e06371f290 100644
--- a/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+++ b/vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
@@ -39,8 +39,6 @@ def __init__(
         scaling_factor: float,
         dtype: torch.dtype,
         *,
-        extrapolation_factor: float = 1,
-        attn_factor: float = 1,
         beta_fast: int = 32,
         beta_slow: int = 1,
         mscale: float = 1,
@@ -48,15 +46,12 @@ def __init__(
         init_cache: bool = True,
     ) -> None:
         self.scaling_factor = scaling_factor
-        self.extrapolation_factor = extrapolation_factor
-        self.attn_factor = attn_factor
         self.beta_fast = beta_fast
         self.beta_slow = beta_slow
         # Get n-d magnitude scaling corrected for interpolation.
         self.mscale = float(
             yarn_get_mscale(self.scaling_factor, float(mscale))
             / yarn_get_mscale(self.scaling_factor, float(mscale_all_dim))
-            * attn_factor
         )
         self.use_flashinfer = (
             self.enabled()
@@ -96,10 +91,9 @@ def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
             self.max_position_embeddings,
         )
         # Get n-d rotational scaling corrected for extrapolation
-        inv_freq_mask = (
-            1
-            - yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float)
-        ) * self.extrapolation_factor
+        inv_freq_mask = 1 - yarn_linear_ramp_mask(
+            low, high, self.rotary_dim // 2, dtype=torch.float
+        )
         inv_freq = (
             inv_freq_interpolation * (1 - inv_freq_mask)
             + inv_freq_extrapolation * inv_freq_mask
@@ -269,7 +263,6 @@ def forward_native(
         inverse: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """PyTorch-native implementation equivalent to forward()."""
-
         head_size = query.size(-1)
         query_rot = query[..., -self.rotary_dim :]
         key_rot = key[..., -self.rotary_dim :] if key is not None else None
diff --git a/vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py b/vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py
index ec0c5a3286c7..428901596bb1 100644
--- a/vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py
+++ b/vllm/model_executor/layers/rotary_embedding/ernie45_vl_rope.py
@@ -8,7 +8,7 @@
 
 
 class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
-    """3D rotary positional embedding. 3D is t:time h:height w:width"""
+    """3D rotary positional embedding. 3D is t:time h:height w:width."""
 
     def forward_native(  # type: ignore[override]
         self,
diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py
index 4c8a7bcbfa1d..96ba482184a8 100644
--- a/vllm/model_executor/layers/rotary_embedding/fope.py
+++ b/vllm/model_executor/layers/rotary_embedding/fope.py
@@ -185,7 +185,7 @@ def forward_native(
         return query, key
 
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
-        """load fope weights"""
+        """Load fope weights."""
         world_size = get_tensor_model_parallel_world_size()
         rank = get_tensor_model_parallel_rank()
         num_key_value_heads = loaded_weight.size(0)
diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py
index 2cbffab8ddce..42013926d711 100644
--- a/vllm/model_executor/layers/rotary_embedding/mrope.py
+++ b/vllm/model_executor/layers/rotary_embedding/mrope.py
@@ -7,6 +7,7 @@
 
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from .base import RotaryEmbeddingBase
 from .yarn_scaling_rope import YaRNScalingRotaryEmbedding, yarn_get_mscale
@@ -178,8 +179,12 @@ def triton_mrope(
             (T/H/W positions with multimodal inputs)
         mrope_section: [t, h, w]
         head_size: int
+        rotary_dim: Number of leading dimensions rotary is applied to.
+        mrope_interleaved: Whether the T/H/W sections are interleaved rather
+            than concatenated.
         is_neox_style: Whether rotary pairs use split-half (NeoX) or
             adjacent (GPT-J) layout.
+
     """
     n_row, n_q_head_head_dim = q.shape
     n_q_head = n_q_head_head_dim // head_size
@@ -236,6 +241,62 @@ def apply_interleaved_rope(x: torch.Tensor, mrope_section: list[int]) -> torch.T
     return torch.where(is_width, x[2], result)
 
 
+def _mrope_apply(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    head_size: int,
+    rotary_dim: int,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    mrope_interleaved: bool,
+    is_neox_style: bool,
+) -> None:
+    cos_sin = cos_sin_cache[positions]
+    cos, sin = cos_sin.chunk(2, dim=-1)
+    query_out, key_out = triton_mrope(
+        query,
+        key,
+        cos,
+        sin,
+        [mrope_section_t, mrope_section_h, mrope_section_w],
+        head_size,
+        rotary_dim,
+        mrope_interleaved,
+        is_neox_style,
+    )
+    if query_out.data_ptr() != query.data_ptr():
+        query.copy_(query_out.view_as(query))
+    if key_out.data_ptr() != key.data_ptr():
+        key.copy_(key_out.view_as(key))
+
+
+def _mrope_apply_fake(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    head_size: int,
+    rotary_dim: int,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    mrope_interleaved: bool,
+    is_neox_style: bool,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="mrope",
+    op_func=_mrope_apply,
+    mutates_args=["query", "key"],
+    fake_impl=_mrope_apply_fake,
+)
+
+
 class MRotaryEmbedding(RotaryEmbeddingBase):
     """Rotary Embedding with Multimodal Sections."""
 
@@ -252,21 +313,28 @@ def __init__(
         # YaRN parameters.
         *,
         scaling_factor: float | None = None,
-        extrapolation_factor: float = 1,
-        attn_factor: float = 1,
         beta_fast: int = 32,
         beta_slow: int = 1,
+        mscale: float | None = None,
+        mscale_all_dim: float | None = None,
+        attention_factor: float | None = None,
         truncate: bool = True,
     ) -> None:
         self.scaling_factor = scaling_factor
-        self.extrapolation_factor = extrapolation_factor
-        self.attn_factor = attn_factor
         self.beta_fast = beta_fast
         self.beta_slow = beta_slow
         self.truncate = truncate
         if self.scaling_factor is not None:
             # Get n-d magnitude scaling corrected for interpolation
-            self.mscale = float(yarn_get_mscale(self.scaling_factor) * attn_factor)
+            if attention_factor is not None:
+                self.mscale = float(attention_factor)
+            elif mscale and mscale_all_dim:
+                self.mscale = float(
+                    yarn_get_mscale(self.scaling_factor, mscale)
+                    / yarn_get_mscale(self.scaling_factor, mscale_all_dim)
+                )
+            else:
+                self.mscale = float(yarn_get_mscale(self.scaling_factor))
         else:
             self.mscale = 1.0
 
@@ -313,6 +381,8 @@ def forward_native(
                 [3, num_tokens] (T/H/W positions with multimodal inputs)
             query: [num_tokens, num_heads * head_size]
             key: [num_tokens, num_kv_heads * head_size]
+            offsets: Optional per-token position offsets added to positions.
+
         """
         assert positions.ndim == 1 or positions.ndim == 2
         assert key is not None
@@ -413,6 +483,35 @@ def forward_cuda(
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
+    def forward_hip(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor | None = None,
+        offsets: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if positions.ndim == 2:
+            assert key is not None
+            assert self.mrope_section
+            query_shape = query.shape
+            key_shape = key.shape
+            cos_sin_cache = self._match_cos_sin_cache_dtype(query)
+            torch.ops.vllm.mrope(
+                positions,
+                query,
+                key,
+                cos_sin_cache,
+                self.head_size,
+                self.rotary_dim,
+                self.mrope_section[0],
+                self.mrope_section[1],
+                self.mrope_section[2],
+                self.mrope_interleaved,
+                self.is_neox_style,
+            )
+            return query.reshape(query_shape), key.reshape(key_shape)
+        return super().forward_hip(positions, query, key)
+
     def forward_cpu(
         self,
         positions: torch.Tensor,
diff --git a/vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py b/vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py
index a6f6c276b442..f1cf14b96d1e 100644
--- a/vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py
+++ b/vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py
@@ -85,7 +85,7 @@ def _rebuild_pos_emb(
         self,
         positions: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Interleave the rotary embedding"""
+        """Interleave the rotary embedding."""
         cos_sin = self.cos_sin_cache[positions]
         mrope_section_3d = [1] * len(self.mrope_dim)
         mrope_dim = self.mrope_dim
@@ -139,8 +139,7 @@ def forward(
     def get_mrope_interleaved_id_list(
         a: int, b: int, c: int, force_last: bool = False
     ) -> list[int]:
-        """
-        Generate an interleaved list of indices for multi-modal rotary embedding.
+        """Generate an interleaved list of indices for multi-modal rotary embedding.
 
         Args:
             a: Number of indices for first modality
@@ -150,6 +149,7 @@ def get_mrope_interleaved_id_list(
 
         Returns:
             List of interleaved indices
+
         """
         if force_last:
             a -= 1
diff --git a/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py
deleted file mode 100644
index dd2fb9c320b6..000000000000
--- a/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py
+++ /dev/null
@@ -1,57 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import math
-
-import torch
-
-from .base import RotaryEmbedding
-from .yarn_scaling_rope import YaRNScalingRotaryEmbedding
-
-
-class TeleChat3RoPEScaledRotaryEmbedding(YaRNScalingRotaryEmbedding):
-    """TeleChat3 uses a variant of YaRN method.
-
-    To achieve code reuse as much as possible, we have rewritten the
-    `get_mscale` method in the initialization function
-    """
-
-    def __init__(
-        self,
-        head_size: int,
-        rotary_dim: int,
-        max_position_embeddings: int,
-        base: int,
-        is_neox_style: bool,
-        scaling_factor: float,
-        dtype: torch.dtype,
-        *,
-        extrapolation_factor: float = 1,
-        attn_factor: float = 1,
-        beta_fast: int = 32,
-        beta_slow: int = 1,
-        truncate: bool = True,
-    ) -> None:
-        self.scaling_factor = scaling_factor
-        self.extrapolation_factor = extrapolation_factor
-        self.attn_factor = attn_factor
-        self.beta_fast = beta_fast
-        self.beta_slow = beta_slow
-        self.truncate = truncate
-
-        def get_mscale(scale, mscale=1):
-            if scale <= 1:
-                return 1.0
-            return 0.07 * mscale * math.log(scale) + 1.0
-
-        self.mscale = float(get_mscale(self.scaling_factor) * attn_factor)
-        # Initialization must be performed after mscale, otherwise mscale is useless
-        RotaryEmbedding.__init__(
-            self,
-            head_size,
-            rotary_dim,
-            max_position_embeddings,
-            base,
-            is_neox_style,
-            dtype,
-        )
diff --git a/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
index f01ca1e23121..01fbf0b2c5c2 100644
--- a/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
+++ b/vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
@@ -23,25 +23,27 @@ def __init__(
         scaling_factor: float,
         dtype: torch.dtype,
         *,
-        extrapolation_factor: float = 1,
-        attn_factor: float = 1,
         beta_fast: int = 32,
         beta_slow: int = 1,
-        apply_yarn_scaling: bool = True,
+        mscale: float | None = None,
+        mscale_all_dim: float | None = None,
+        attention_factor: float | None = None,
         truncate: bool = True,
     ) -> None:
         self.scaling_factor = scaling_factor
-        self.extrapolation_factor = extrapolation_factor
-        self.attn_factor = attn_factor
         self.beta_fast = beta_fast
         self.beta_slow = beta_slow
         self.truncate = truncate
         # Get n-d magnitude scaling corrected for interpolation
-        self.mscale = (
-            float(yarn_get_mscale(self.scaling_factor) * attn_factor)
-            if apply_yarn_scaling
-            else float(attn_factor)
-        )
+        if attention_factor is not None:
+            self.mscale = float(attention_factor)
+        elif mscale and mscale_all_dim:
+            self.mscale = float(
+                yarn_get_mscale(self.scaling_factor, mscale)
+                / yarn_get_mscale(self.scaling_factor, mscale_all_dim)
+            )
+        else:
+            self.mscale = float(yarn_get_mscale(self.scaling_factor))
         super().__init__(
             head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
         )
@@ -62,10 +64,9 @@ def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
             self.truncate,
         )
         # Get n-d rotational scaling corrected for extrapolation
-        inv_freq_mask = (
-            1
-            - yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float)
-        ) * self.extrapolation_factor
+        inv_freq_mask = 1 - yarn_linear_ramp_mask(
+            low, high, self.rotary_dim // 2, dtype=torch.float
+        )
         inv_freq = (
             inv_freq_interpolation * (1 - inv_freq_mask)
             + inv_freq_extrapolation * inv_freq_mask
diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py
index 7d656f06e85f..ae45b77dcea8 100644
--- a/vllm/model_executor/layers/sparse_attn_indexer.py
+++ b/vllm/model_executor/layers/sparse_attn_indexer.py
@@ -3,6 +3,7 @@
 """Custom Sparse Attention Indexer layers."""
 
 import torch
+import torch.distributed as dist
 
 import vllm.envs as envs
 from vllm import _custom_ops as ops
@@ -19,6 +20,10 @@
 from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
     select_candidate_blocks as _select_candidate_blocks,
 )
+from vllm.model_executor.layers.indexer_topk import (
+    RADIX_TOPK_WORKSPACE_SIZE,
+    get_indexer_topk,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     get_fp8_min_max,
 )
@@ -45,8 +50,6 @@
 
 logger = init_logger(__name__)
 
-RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
-
 # MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
 MXFP4_BLOCK_SIZE = 32
 
@@ -130,6 +133,22 @@ def _merge_dcp_topk_global(
     )
 
 
+def dcp_gather_kv_rows(
+    local_padded: torch.Tensor,
+    deinterleave_idx: torch.Tensor,
+    gathered_buf: torch.Tensor,
+    out_buf: torch.Tensor,
+) -> torch.Tensor:
+    """All-gather this rank's KV shard and return it in global token order."""
+    gathered = gathered_buf[: get_dcp_group().world_size * local_padded.shape[0]]
+    dist.all_gather_into_tensor(
+        gathered, local_padded.contiguous(), group=get_dcp_group().device_group
+    )
+    return torch.index_select(
+        gathered, 0, deinterleave_idx, out=out_buf[: deinterleave_idx.shape[0]]
+    )
+
+
 @triton.jit
 def _fused_indexer_q_rope_quant_kernel(
     positions,
@@ -325,6 +344,7 @@ def sparse_attn_indexer(
     candidate_blocks: torch.Tensor | None = None,
     candidate_block_size: int = 0,
     candidate_write: bool = False,
+    topk_backend: str = "auto",
 ) -> torch.Tensor:
     # careful! this will be None in dummy run
     forward_context = get_forward_context()
@@ -346,11 +366,19 @@ def sparse_attn_indexer(
         values_spec, scales_spec = _gather_workspace_shapes(
             total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
         )
-        current_workspace_manager().get_simultaneous(
+        profile_specs: list[tuple[tuple[int, ...], torch.dtype]] = [
             values_spec,
             scales_spec,
             ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-        )
+        ]
+        if use_pcp and dcp_world_size > 1:
+            # The PCP+DCP path takes an all-gather destination and a
+            # de-interleaved result.
+            gather_spec = _gather_workspace_shapes(
+                total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
+            )
+            profile_specs.extend(gather_spec * 2)
+        current_workspace_manager().get_simultaneous(*profile_specs)
 
         # Dummy allocation to simulate for peak logits tensor memory during inference.
         # FP8 elements so elements == bytes
@@ -401,9 +429,10 @@ def sparse_attn_indexer(
     # out-of-bounds reads in the kernel.
     # Keep PCP padding so every rank contributes the same all-gather shape.
     num_tokens = slot_mapping.shape[0]
-    if use_pcp:
-        num_tokens //= get_pcp_group().world_size
     if k is not None:
+        # MTP draft decodes are replicated and use an unexpanded slot mapping.
+        if use_pcp and num_tokens > k.shape[0]:
+            num_tokens //= get_pcp_group().world_size
         k = k[:num_tokens]
 
     if not skip_k_cache_insert:
@@ -424,7 +453,6 @@ def sparse_attn_indexer(
             quant_block_size,
             scale_fmt,
         )
-
     # The indexer and main MLA may classify the same short extend differently
     # because they use independent decode thresholds. Only the main MLA route
     # can determine whether the top-k indices will be consumed.
@@ -462,9 +490,27 @@ def sparse_attn_indexer(
         values_spec, scales_spec = _gather_workspace_shapes(
             total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
         )
-        k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
+        # PCP + DCP needs two more pairs: the rank-major all-gather destination
+        # and the de-interleaved result.
+        pcp_chunks = [
+            (c, c.pcp_deinterleave_idx)
+            for c in prefill_metadata.chunks
+            if c.pcp_deinterleave_idx is not None
+        ]
+        gather_specs: list[tuple[tuple[int, int], torch.dtype]] = []
+        if pcp_chunks:
+            gathered_rows = max(
+                dcp_world_size * c.local_total_seq_lens for c, _ in pcp_chunks
+            )
+            deinterleaved_rows = max(idx.shape[0] for _, idx in pcp_chunks)
+            for rows in (gathered_rows, deinterleaved_rows):
+                gather_specs.extend(
+                    _gather_workspace_shapes(rows, head_dim, fp8_dtype, use_fp4_cache)
+                )
+        k_quant_full, k_scale_full, *gather_bufs = workspace_manager.get_simultaneous(
             values_spec,
             scales_spec,
+            *gather_specs,
         )
         for chunk in prefill_metadata.chunks:
             cu_seqlen_ks = chunk.cu_seqlen_ks
@@ -481,6 +527,23 @@ def sparse_attn_indexer(
                     chunk.local_cu_seq_lens,
                 )
 
+            # PCP + DCP KV all-gather.
+            deinterleave_idx = chunk.pcp_deinterleave_idx
+            if deinterleave_idx is not None:
+                # local_total_seq_lens is the PCP-padded extent here, so every
+                # rank contributes an identically shaped shard.
+                gathered_values, gathered_scales, out_values, out_scales = gather_bufs
+                shard = k_quant[: chunk.local_total_seq_lens]
+                k_quant = dcp_gather_kv_rows(
+                    shard, deinterleave_idx, gathered_values, out_values
+                )
+                k_scale = dcp_gather_kv_rows(
+                    k_scale[: chunk.local_total_seq_lens],
+                    deinterleave_idx,
+                    gathered_scales,
+                    out_scales,
+                )
+
             q_slice = q_quant[chunk.token_start : chunk.token_end]
             q_scale_slice = (
                 q_scale[chunk.token_start : chunk.token_end]
@@ -491,7 +554,7 @@ def sparse_attn_indexer(
                 chunk.token_start : chunk.token_end, :topk_tokens
             ]
 
-            if chunk.local_total_seq_lens == 0:
+            if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0:
                 logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
                 topk_indices.fill_(-1)
             else:
@@ -561,21 +624,37 @@ def sparse_attn_indexer(
                     topk_tokens,
                 )
 
-            _merge_dcp_topk_global(
-                logits,
-                topk_indices,
-                topk_tokens,
-                dcp_rank,
-                dcp_world_size,
-                cp_kv_cache_interleave_size,
-                row_starts=chunk.cu_seqlen_ks,
-            )
+            if deinterleave_idx is None:
+                # Under the PCP path the top-k already ran over the whole
+                # context.
+                _merge_dcp_topk_global(
+                    logits,
+                    topk_indices,
+                    topk_tokens,
+                    dcp_rank,
+                    dcp_world_size,
+                    cp_kv_cache_interleave_size,
+                    row_starts=chunk.cu_seqlen_ks,
+                )
 
     if has_decode:
         decode_metadata = attn_metadata_narrowed.decode
         assert decode_metadata is not None
         kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
         decode_lens = decode_metadata.decode_lens
+        # requires_padding can be computed False for ragged warmup/mixed
+        # batches (seen on SM120 TP=2: 8 tokens over 6 seqs -> uniform
+        # reshape below would crash). The token-count/batch-size divisibility
+        # check is a necessary-but-not-sufficient proxy for "uniform batch"
+        # (e.g. decode_lens=[1, 2, 3] divides evenly by 3 but isn't uniform);
+        # it catches the specific case seen in practice without requiring a
+        # device sync to inspect decode_lens directly, which would break
+        # CUDA-graph-capture safety on this path. Shared by both the pack
+        # branch below and the matching unpack branch further down — they
+        # must agree on which path was taken.
+        needs_padded_path = decode_metadata.requires_padding or (
+            num_decode_tokens % decode_lens.shape[0] != 0
+        )
         if num_decode_tokens == 0:
             padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:])
             padded_q_scale = (
@@ -583,7 +662,7 @@ def sparse_attn_indexer(
                 if q_scale is not None
                 else None
             )
-        elif decode_metadata.requires_padding:
+        elif needs_padded_path:
             # pad in edge case where we have short chunked prefill length <
             # decode_threshold since we unstrictly split
             # prefill and decode by decode_threshold
@@ -686,56 +765,16 @@ def sparse_attn_indexer(
                 )
         topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
 
-        use_cooperative_topk = (
-            current_platform.is_cuda()
-            and topk_tokens in (512, 1024, 2048)
-            and num_rows <= 64
-            and logits.stride(0) % 4 == 0  # TMA 16-byte alignment
-            and current_platform.has_device_capability(90)
-            and not current_platform.is_device_capability_family(120)
-        )
-        use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
-            512,
-            1024,
-            2048,
+        # The backend comes from the layer (config is only readable at model
+        # construction); dispatchers are cached per backend.
+        get_indexer_topk(topk_backend)(
+            logits,
+            seq_lens,
+            next_n,
+            topk_indices,
+            topk_tokens,
+            attn_metadata_narrowed.max_seq_len,
         )
-        if use_cooperative_topk:
-            workspace_manager = current_workspace_manager()
-            (topk_workspace,) = workspace_manager.get_simultaneous(
-                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-            )
-            torch.ops._C.cooperative_topk(
-                logits,
-                seq_lens,
-                topk_indices,
-                topk_workspace,
-                topk_tokens,
-                attn_metadata_narrowed.max_seq_len,
-            )
-        elif use_persistent_topk:
-            workspace_manager = current_workspace_manager()
-            (topk_workspace,) = workspace_manager.get_simultaneous(
-                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-            )
-            torch.ops._C.persistent_topk(
-                logits,
-                seq_lens,
-                topk_indices,
-                topk_workspace,
-                topk_tokens,
-                logits.shape[1],
-            )
-        else:
-            ops.top_k_per_row_decode(
-                logits,
-                next_n,
-                seq_lens,
-                topk_indices,
-                num_rows,
-                logits.stride(0),
-                logits.stride(1),
-                topk_tokens,
-            )
 
         if decode_metadata.global_seq_lens is not None:
             _merge_dcp_topk_global(
@@ -747,7 +786,7 @@ def sparse_attn_indexer(
                 cp_kv_cache_interleave_size,
             )
 
-        if decode_metadata.requires_padding:
+        if needs_padded_path:
             # if padded, we need to unpack
             # the topk indices removing padded tokens
             topk_indices = unpack_seq_triton(
@@ -787,6 +826,7 @@ def sparse_attn_indexer_fake(
     candidate_blocks: torch.Tensor | None = None,
     candidate_block_size: int = 0,
     candidate_write: bool = False,
+    topk_backend: str = "auto",
 ) -> torch.Tensor:
     return topk_indices_buffer
 
@@ -853,6 +893,7 @@ def __init__(
         # than threading them through per-step metadata.
         vllm_config = get_current_vllm_config()
         parallel_config = vllm_config.parallel_config
+        self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
         self._parallel_config = parallel_config
         self.dcp_world_size = parallel_config.decode_context_parallel_size
         self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
@@ -910,6 +951,10 @@ def forward_native(
         if current_platform.is_cuda() or current_platform.is_xpu():
             return self.forward_cuda(hidden_states, q_quant, k, weights)
         elif current_platform.is_rocm():
+            if self.use_pcp and self.dcp_world_size > 1:
+                raise NotImplementedError(
+                    "The ROCm sparse-indexer path does not support PCP+DCP."
+                )
             return self.forward_hip(hidden_states, q_quant, k, weights)
         elif current_platform.is_cpu():
             return self.forward_cpu(hidden_states, q_quant, k, weights)
@@ -954,9 +999,11 @@ def forward_cuda(
             self.dcp_rank,
             self.dcp_world_size,
             self.cp_kv_cache_interleave_size,
+            False,
             candidate_blocks=self.candidate_blocks,
             candidate_block_size=self.candidate_block_size,
             candidate_write=self.candidate_write,
+            topk_backend=self.topk_backend,
         )
 
     def forward_xpu(
diff --git a/vllm/model_executor/layers/sparse_mqa_indexer.py b/vllm/model_executor/layers/sparse_mqa_indexer.py
new file mode 100644
index 000000000000..78db1b3df5a4
--- /dev/null
+++ b/vllm/model_executor/layers/sparse_mqa_indexer.py
@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sparse Attention Indexer that scores only the candidate blocks.
+
+DeepSeek V4.1 two-level selection: the candidate-source indexer publishes
+the top candidate blocks and later indexers pick their top-k inside them.
+`SparseAttnIndexer` does that by computing dense logits over the whole
+context and masking; this layer instead calls DeepGEMM's sparse MQA-logits
+kernels on the candidate blocks only, so the work is O(candidate blocks)
+instead of O(context). It requires the `DeepseekV41SparseIndexerBackend`
+metadata (see ``AttentionConfig.indexer_sparse_logits``).
+"""
+
+import torch
+from torch import nn
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+    has_deep_select,
+    sparse_mqa_logits_paged_decode,
+    sparse_mqa_logits_prefill_chunk,
+)
+from vllm.model_executor.layers.sparse_attn_indexer import (
+    _gather_workspace_shapes,
+    kv_cache_as_quant_view,
+)
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV32IndexerPrefillChunkMetadata,
+)
+from vllm.v1.attention.backends.mla.sparse_indexer import (
+    DeepseekV41SparseIndexerMetadata,
+)
+from vllm.v1.worker.workspace import current_workspace_manager
+
+
+def _prefill_k_workspaces(
+    total_seq_lens: int, head_dim: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with
+    the dense indexer layers of the same model."""
+    values_spec, scales_spec = _gather_workspace_shapes(
+        total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True
+    )
+    k_quant, k_scale = current_workspace_manager().get_simultaneous(
+        values_spec, scales_spec
+    )
+    return k_quant, k_scale
+
+
+def _gather_prefill_chunk_k(
+    kv_cache: torch.Tensor,
+    k_quant_full: torch.Tensor,
+    k_scale_full: torch.Tensor,
+    chunk: DeepseekV32IndexerPrefillChunkMetadata,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Gather one prefill chunk's paged K into the packed workspace."""
+    assert chunk.local_cu_seq_lens is not None
+    k_quant = k_quant_full[: chunk.max_local_total_seq_lens]
+    k_scale = k_scale_full[: chunk.max_local_total_seq_lens]
+    if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0:
+        ops.cp_gather_indexer_k_quant_cache(
+            kv_cache,
+            k_quant,
+            k_scale,
+            chunk.block_table,
+            chunk.local_cu_seq_lens,
+        )
+    return k_quant, k_scale
+
+
+class SparseMQAIndexer(nn.Module):
+    """Candidate-consuming indexer on DeepGEMM's sparse MQA-logits kernels.
+
+    Only valid for indexer layers that read candidate blocks with the MXFP4
+    indexer cache on SM100. The K cache is written by the model before this
+    runs; ``forward`` takes the same arguments as `SparseAttnIndexer` so the
+    attention layer can call either.
+    """
+
+    weights_dtype = torch.bfloat16
+    """Per-head weights dtype the sparse kernels take. The fused Q RoPE-quant
+    kernel writes it directly so no cast runs per step."""
+
+    def __init__(
+        self,
+        k_cache,
+        topk_tokens: int,
+        head_dim: int,
+        max_total_seq_len: int,
+        topk_indices_buffer: torch.Tensor,
+        candidate_blocks: torch.Tensor,
+        candidate_block_size: int,
+    ):
+        super().__init__()
+        if not (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(100)
+            and has_deep_gemm_sparse_mqa()
+        ):
+            raise ValueError(
+                "SparseMQAIndexer requires an SM100-class GPU and DeepGEMM >= 2.8."
+            )
+        if not has_deep_select():
+            raise ValueError(
+                "SparseMQAIndexer requires the DeepSelect top-k extension "
+                "(vllm._deepselect_C)."
+            )
+        self.k_cache = k_cache
+        self.topk_tokens = topk_tokens
+        self.head_dim = head_dim
+        self.max_total_seq_len = max_total_seq_len
+        self.topk_indices_buffer = topk_indices_buffer
+        self.candidate_blocks = candidate_blocks
+        self.candidate_block_size = candidate_block_size
+
+    def _reserve_workspaces(self, device: torch.device) -> None:
+        """Profiling run: claim the K-gather workspace and the peak sparse
+        logits allocation so the memory estimate covers them."""
+        _prefill_k_workspaces(self.max_total_seq_len, self.head_dim)
+        max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+        torch.empty(max_logits_bytes, dtype=torch.uint8, device=device)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        q_quant: tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor | None,
+        weights: torch.Tensor,
+    ) -> torch.Tensor:
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            self._reserve_workspaces(hidden_states.device)
+            return self.topk_indices_buffer
+
+        metadata = attn_metadata[self.k_cache.prefix]
+        assert isinstance(metadata, DeepseekV41SparseIndexerMetadata), (
+            "SparseMQAIndexer needs DeepseekV41SparseIndexerBackend metadata"
+        )
+        assert k is None, "the model writes the indexer K cache"
+        q_values, q_scale = q_quant
+        kv_cache = self.k_cache.kv_cache
+        topk_tokens = self.topk_tokens
+        topk_indices_buffer = self.topk_indices_buffer
+        topk_indices_buffer[: hidden_states.shape[0]] = -1
+        sparse_block_kv = metadata.sparse_block_kv
+
+        if metadata.num_prefills > 0:
+            prefill = metadata.prefill
+            sparse_chunks = metadata.sparse_prefill
+            assert prefill is not None and sparse_chunks is not None
+            k_quant_full, k_scale_full = _prefill_k_workspaces(
+                self.max_total_seq_len, self.head_dim
+            )
+            for chunk, rows in zip(prefill.chunks, sparse_chunks):
+                k_quant, k_scale = _gather_prefill_chunk_k(
+                    kv_cache, k_quant_full, k_scale_full, chunk
+                )
+                if chunk.local_total_seq_lens == 0:
+                    continue  # buffer already holds -1
+                start, end = chunk.token_start, chunk.token_end
+                rows.kernel_metadata = sparse_mqa_logits_prefill_chunk(
+                    q_values[start:end].view(torch.int8),
+                    q_scale[start:end],
+                    k_quant.view(torch.int8),
+                    k_scale.view(torch.int32).squeeze(-1),
+                    weights[start:end],
+                    chunk.cu_seqlen_ks,
+                    chunk.cu_seqlen_ke,
+                    self.candidate_blocks[start:end],
+                    self.candidate_block_size,
+                    sparse_block_kv,
+                    topk_tokens,
+                    topk_indices_buffer[start:end, :topk_tokens],
+                    sparse_indices=rows.sparse_indices,
+                    end=rows.end,
+                    col_indices=rows.col_indices,
+                    kernel_metadata=rows.kernel_metadata,
+                )
+
+        if metadata.num_decodes > 0:
+            decode_rows = metadata.sparse_decode
+            assert decode_rows is not None and decode_rows.block_table is not None
+            assert decode_rows.row_indices is not None
+            num_rows = decode_rows.row_ke.shape[0]
+            if num_rows > 0:
+                decode_rows.kernel_metadata = sparse_mqa_logits_paged_decode(
+                    q_values[:num_rows].view(torch.int8).unsqueeze(1),
+                    q_scale[:num_rows].unsqueeze(1),
+                    kv_cache_as_quant_view(kv_cache, self.head_dim, use_fp4_cache=True),
+                    weights[:num_rows],
+                    decode_rows.row_ke,
+                    decode_rows.block_table,
+                    decode_rows.row_indices,
+                    self.candidate_blocks[:num_rows],
+                    self.candidate_block_size,
+                    sparse_block_kv,
+                    topk_tokens,
+                    topk_indices_buffer[:num_rows, :topk_tokens],
+                    row_ks=decode_rows.row_ks,
+                    sparse_indices=decode_rows.sparse_indices,
+                    end=decode_rows.end,
+                    col_indices=decode_rows.col_indices,
+                    kernel_metadata=decode_rows.kernel_metadata,
+                )
+        return topk_indices_buffer
diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py
index 03239e110f7b..cff3d5b0785c 100644
--- a/vllm/model_executor/layers/utils.py
+++ b/vllm/model_executor/layers/utils.py
@@ -48,8 +48,7 @@ def apply_penalties(
     frequency_penalties: torch.Tensor,
     repetition_penalties: torch.Tensor,
 ) -> torch.Tensor:
-    """
-    Applies penalties in place to the logits tensor
+    """Applies penalties in place to the logits tensor
     logits : The input logits tensor of shape [num_seqs, vocab_size]
     prompt_tokens_tensor: A tensor containing the prompt tokens. The prompts
         are padded to the maximum prompt length within the batch using
@@ -254,6 +253,46 @@ def use_aiter_triton_gemm(n, m, k, dtype):
     )
 
 
+def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, bool]:
+    """Pick the K-shard split for wvSplitKrc and say whether the shape fits.
+
+    Mirrors wvSplitKrc() in csrc/rocm/skinny_gemms.cu, which is also where the
+    shard cap is explained. Both must pick the same chunkk or the workspace
+    check here bounds the wrong k_rnd.
+
+    Returns:
+        The CHUNKK the kernel will dispatch with, and whether the CU budget and
+        split-K workspace admit the shape at all.
+
+    """
+    # Next ^2 of n
+    N_p2 = 1 << (n - 1).bit_length()
+    # How many of 4 waves in a group can work on same 16 Ms at same time?
+    # This reduces the Ms each group works on, i.e. increasing the CUs needed.
+    GrpsShrB = min(N_p2 // 16, 4)
+    # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile), and each
+    # working on a 512-shard of K, how many CUs would we need?
+    CuNeeded = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512) * GrpsShrB
+
+    CHUNKK2_MAX_SHARDS = 11
+    shards_chunkk2 = (k + 256 - 1) // 256  # 256-wide shards
+    chunkk = (
+        2
+        if (
+            N_p2 != 16
+            and CuNeeded * 2 <= cu_count
+            and shards_chunkk2 <= CHUNKK2_MAX_SHARDS
+        )
+        else 1
+    )
+
+    # Deterministic reduction stores one fp32 partial per (M, N, k-shard); all
+    # of them must fit the split-K workspace.
+    k_rnd = (k + 512 // chunkk - 1) // (512 // chunkk)
+    fits = N_p2 * m * k_rnd <= 128 * 1024 * 12 and CuNeeded <= cu_count
+    return chunkk, fits
+
+
 def rocm_unquantized_gemm_impl(
     x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
 ) -> torch.Tensor:
@@ -265,21 +304,7 @@ def rocm_unquantized_gemm_impl(
 
     cu_count = num_compute_units()
 
-    # Next ^2 of n
-    N_p2 = 1 << (n - 1).bit_length()
-    # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
-    # and each working on a 512-shard of K, how many CUs would we need?
-    rndup_cus = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512)
-    # How many of 4 waves in a group can work on same 16 Ms at same time?
-    # This reduces the Ms each group works on, i.e. increasing the number of CUs needed.
-    GrpsShrB = min(N_p2 // 16, 4)
-    # Given the above, how many CUs would we need?
-    CuNeeded = rndup_cus * GrpsShrB
-    # Deterministic reduction stores one float workspace value per K shard.
-    fits_wvsplitkrc = (
-        N_p2 * m * ((k + 512 - 1) // 512)
-    ) <= 128 * 1024 * 12  # deterministic
-    fits_wvsplitkrc &= CuNeeded <= cu_count
+    _, fits_wvsplitkrc = wvsplitkrc_dispatch(n, k, m, cu_count)
 
     skinny_operands_compatible = weight.is_contiguous() and (
         bias is None or bias.is_contiguous()
@@ -323,7 +348,7 @@ def rocm_unquantized_gemm_impl(
         # The skinny kernels assume contiguous K elements. A shape-preserving
         # reshape can retain a transposed activation's non-contiguous strides.
         x_view = x.reshape(-1, x.size(-1)).contiguous()
-        if m > 8 and 0 < n <= 5:
+        if (m == 1 or m > 8) and 0 < n <= 5:
             cu_count = num_compute_units()
             out = ops.wvSplitK(weight, x_view, cu_count, bias)
             return out.reshape(*x.shape[:-1], weight.shape[0])
@@ -363,12 +388,18 @@ def rocm_unquantized_gemm(
 
 @functools.cache
 def warmup_rocm_skinny_gemm_workspaces(device: torch.device) -> None:
-    """Eagerly allocate wvSplitKrc's process-lifetime static workspaces.
+    """Eagerly allocate wvSplitKrc's per-device split-K workspace pool.
+
+    wvSplitKrc partitions one per-device allocation into ``kWvSlots`` slots
+    (csrc/rocm/skinny_gemms.cu) and hands each stream one on first use, so that
+    two streams never share the split-K partials and counters.
 
-    They are otherwise created lazily on the first qualifying GEMM
+    The pool is otherwise created lazily on the first qualifying GEMM
     (csrc/rocm/skinny_gemms.cu), which can be the first real request — after
-    the KV cache backing buffer exists. If one landed in that segment's
-    rounding tail, it would pin the entire segment at engine shutdown.
+    the KV cache backing buffer exists. If it landed in that segment's rounding
+    tail, it would pin the entire segment at engine shutdown; it could also land
+    inside a cudagraph capture, where it would be taken from the graph's private
+    pool and its zero-fill would become a replayed graph node.
     """
     from vllm.platforms.rocm import on_gfx950
 
@@ -483,7 +514,10 @@ def dispatch_cpu_unquantized_gemm(
             )
         )
         if remove_weight:
-            layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
+            layer.weight = torch.nn.Parameter(
+                torch.empty(0, dtype=dtype, device=layer.weight.device),
+                requires_grad=False,
+            )
         logger.debug_once(
             "CPU unquantized GEMM dispatch: using zentorch_linear_unary (prepacked=%s)",
             is_prepacked,
@@ -509,7 +543,10 @@ def dispatch_cpu_unquantized_gemm(
             bias,
         )
         if remove_weight:
-            layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
+            layer.weight = torch.nn.Parameter(
+                torch.empty(0, dtype=dtype, device=layer.weight.device),
+                requires_grad=False,
+            )
         logger.debug_once(
             "CPU unquantized GEMM dispatch: using sgl-kernel weight_packed_linear"
         )
@@ -524,7 +561,10 @@ def dispatch_cpu_unquantized_gemm(
             handler = ops.create_onednn_mm(origin_weight.t(), 32)
             layer.cpu_linear = lambda x, weight, bias: ops.onednn_mm(handler, x, bias)
             if remove_weight:
-                layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
+                layer.weight = torch.nn.Parameter(
+                    torch.empty(0, dtype=dtype, device=layer.weight.device),
+                    requires_grad=False,
+                )
             logger.debug_once("CPU unquantized GEMM dispatch: using oneDNN onednn_mm")
             return
         except RuntimeError as e:
diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py
index 456a22a7a26a..72cafceb2589 100644
--- a/vllm/model_executor/layers/vocab_parallel_embedding.py
+++ b/vllm/model_executor/layers/vocab_parallel_embedding.py
@@ -242,6 +242,7 @@ class VocabParallelEmbedding(PluggableLayer):
         disable_tp: If true, tensor parallelism will be disabled for this layer.
         quant_method: Preselected quantization method for model-specific layers.
         parallel_group: Process group used to shard and reduce the embedding.
+
     """  # noqa: E501
 
     # --8<-- [end:vocab_parallel_embedding]
@@ -581,6 +582,7 @@ class ParallelLMHead(VocabParallelEmbedding):
         org_num_embeddings: original vocabulary size (without LoRA).
         padding_size: padding size for the vocabulary.
         disable_tp: If true, tensor parallelism will be disabled for this layer.
+
     """
 
     # --8<-- [end:parallel_lm_head]
diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py
index 6305c45f5d0b..226caebe1cd4 100644
--- a/vllm/model_executor/model_loader/__init__.py
+++ b/vllm/model_executor/model_loader/__init__.py
@@ -94,6 +94,7 @@ def register_model_loader(load_format: str):
         >>> load_config = LoadConfig(load_format="my_loader")
         >>> type(get_model_loader(load_config))
         <class 'MyModelLoader'>
+
     """  # noqa: E501
 
     def _wrapper(model_loader_cls):
diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py
index 618f874fd0d6..6f356464585a 100644
--- a/vllm/model_executor/model_loader/base_loader.py
+++ b/vllm/model_executor/model_loader/base_loader.py
@@ -39,6 +39,19 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
         inplace weights loading for an already-initialized model"""
         raise NotImplementedError
 
+    def create_model(
+        self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
+    ) -> nn.Module:
+        """Create a model with the given configurations."""
+        model = initialize_model(
+            vllm_config=vllm_config,
+            model_config=model_config,
+            prefix=prefix,
+        )
+        log_online_quantization(vllm_config)
+        log_model_inspection(model)
+        return model
+
     @instrument(span_name="Load model")
     def load_model(
         self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
@@ -52,15 +65,12 @@ def load_model(
         target_device = torch.device(load_device)
         with set_default_torch_dtype(model_config.dtype):
             with target_device:
-                model = initialize_model(
+                model = self.create_model(
                     vllm_config=vllm_config,
                     model_config=model_config,
                     prefix=prefix,
                 )
 
-            log_online_quantization(vllm_config)
-            log_model_inspection(model)
-
             logger.debug("Loading weights on %s ...", load_device)
             self.load_weights(model, model_config)
 
diff --git a/vllm/model_executor/model_loader/checkpoint_weight_patch.py b/vllm/model_executor/model_loader/checkpoint_weight_patch.py
index 0425aabc7dd1..4a26c73ebc86 100644
--- a/vllm/model_executor/model_loader/checkpoint_weight_patch.py
+++ b/vllm/model_executor/model_loader/checkpoint_weight_patch.py
@@ -36,6 +36,7 @@ class CheckpointWeightPatch(NamedTuple):
             ``indices`` for a sparse patch.
         indices: Flat indices into the full checkpoint tensor described by
             ``shape``. ``None`` makes ``values`` a dense replacement.
+
     """
 
     name: str
@@ -176,8 +177,8 @@ def load_checkpoint_weight_patches(
 
     Returns:
         Union of the weight names reported by all ``model.load_weights`` calls.
-    """
 
+    """
     if max_chunk_bytes <= 0:
         raise ValueError("max_chunk_bytes must be positive")
 
diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py
index 039a577f80d1..c3e09bccd425 100644
--- a/vllm/model_executor/model_loader/default_loader.py
+++ b/vllm/model_executor/model_loader/default_loader.py
@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
 import dataclasses
 import glob
 import os
 import time
 from collections.abc import Generator, Iterable
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import torch
 from torch import nn
@@ -25,6 +27,7 @@
     fastsafetensors_weights_iterator,
     filter_duplicate_safetensors_files,
     filter_files_not_needed_for_inference,
+    filter_mm_encoder_only_safetensors_files,
     get_quant_config,
     instanttensor_weights_iterator,
     maybe_download_from_modelscope,
@@ -32,11 +35,15 @@
     multi_thread_safetensors_weights_iterator,
     np_cache_weights_iterator,
     pt_weights_iterator,
+    resolve_mm_encoder_only_lm_prefixes,
     safetensors_weights_iterator,
 )
 from vllm.tracing import instrument
 from vllm.transformers_utils.repo_utils import list_filtered_repo_files
 
+if TYPE_CHECKING:
+    from vllm.model_executor.models.utils import WeightsMapper
+
 logger = init_logger(__name__)
 
 
@@ -74,6 +81,9 @@ class Source:
     def __init__(self, load_config: LoadConfig):
         super().__init__(load_config)
         self.local_expert_ids: set[int] | None = None
+        # Set in load_weights when --mm-encoder-only; used to drop LM-only shards.
+        self._encoder_only_lm_prefixes: tuple[str, ...] | None = None
+        self._encoder_only_weights_mapper: WeightsMapper | None = None
 
         extra_config = load_config.model_loader_extra_config
         if not isinstance(extra_config, dict):
@@ -132,7 +142,7 @@ def _prepare_weights(
         revision: str | None,
         fall_back_to_pt: bool,
         allow_patterns_overrides: list[str] | None,
-    ) -> tuple[str, list[str], bool]:
+    ) -> tuple[str, list[str], bool, str]:
         """Prepare weights for the model.
 
         If the model is not local, it will be downloaded."""
@@ -239,20 +249,40 @@ def _prepare_weights(
                 f"Cannot find any model weights with `{model_name_or_path}`"
             )
 
-        return hf_folder, hf_weights_files, use_safetensors
+        return hf_folder, hf_weights_files, use_safetensors, index_file
 
     def _get_weights_iterator(
-        self, source: "Source"
+        self, source: Source
     ) -> Generator[tuple[str, torch.Tensor], None, None]:
         """Get an iterator for the model weights based on the load format."""
         extra_config = self.load_config.model_loader_extra_config
-        hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
-            source.model_or_path,
-            source.subfolder,
-            source.revision,
-            source.fall_back_to_pt,
-            source.allow_patterns_overrides,
+        hf_folder, hf_weights_files, use_safetensors, index_file = (
+            self._prepare_weights(
+                source.model_or_path,
+                source.subfolder,
+                source.revision,
+                source.fall_back_to_pt,
+                source.allow_patterns_overrides,
+            )
         )
+        if (
+            self._encoder_only_lm_prefixes is not None
+            and use_safetensors
+            and hf_weights_files
+        ):
+            hf_weights_files = filter_mm_encoder_only_safetensors_files(
+                hf_weights_files,
+                hf_folder,
+                index_file,
+                self._encoder_only_lm_prefixes,
+                weights_mapper=self._encoder_only_weights_mapper,
+            )
+            if not hf_weights_files:
+                raise RuntimeError(
+                    "mm-encoder-only shard filter removed all weight files for "
+                    f"`{source.model_or_path}`; check language_model prefixes "
+                    f"{self._encoder_only_lm_prefixes}"
+                )
         if self.load_config.load_format == "npcache":
             # Currently np_cache only support *.bin checkpoints
             assert use_safetensors is False
@@ -348,6 +378,29 @@ def download_model(self, model_config: ModelConfig) -> None:
             allow_patterns_overrides=None,
         )
 
+    def _init_mm_encoder_only_weight_filter(
+        self, model: nn.Module, model_config: ModelConfig
+    ) -> None:
+        """Skip pure language-model safetensors shards under --mm-encoder-only."""
+        mm_config = model_config.multimodal_config
+        if mm_config is None or not mm_config.mm_encoder_only:
+            self._encoder_only_lm_prefixes = None
+            self._encoder_only_weights_mapper = None
+            return
+
+        # Derive from _language_model_names; fail-closed on shared HF roots
+        # (Molmo/Phi-4-MM/Muse). Qwen nested/flat keys classified via mapper.
+        weights_mapper = cast(
+            "WeightsMapper | None", getattr(model, "hf_to_vllm_mapper", None)
+        )
+        self._encoder_only_lm_prefixes = resolve_mm_encoder_only_lm_prefixes(
+            getattr(model, "_language_model_names", None),
+            weights_mapper=weights_mapper,
+        )
+        self._encoder_only_weights_mapper = (
+            weights_mapper if self._encoder_only_lm_prefixes is not None else None
+        )
+
     def _init_ep_weight_filter(self, model_config: ModelConfig) -> None:
         """Compute local expert ids for EP weight filtering.
 
@@ -423,6 +476,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
                 self.load_config.safetensors_load_strategy = "torchao"
 
         self._init_ep_weight_filter(model_config)
+        self._init_mm_encoder_only_weight_filter(model, model_config)
 
         loaded_weights = model.load_weights(self.get_all_weights(model_config, model))
 
diff --git a/vllm/model_executor/model_loader/dummy_loader.py b/vllm/model_executor/model_loader/dummy_loader.py
index 98f82a5b7197..0f263fe11709 100644
--- a/vllm/model_executor/model_loader/dummy_loader.py
+++ b/vllm/model_executor/model_loader/dummy_loader.py
@@ -14,7 +14,6 @@
 from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo
 from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
 from vllm.model_executor.model_loader.weight_utils import (
-    initialize_dummy_weights,
     initialize_single_dummy_weight,
 )
 
@@ -41,7 +40,12 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
             else:
                 # NOTE(woosuk): For accurate performance evaluation, we assign
                 # random values to the weights.
-                initialize_dummy_weights(layer, model_config)
+                # Initialize only local tensors: state_dict() would recursively
+                # initialize descendants that model.modules() also visits.
+                # Exclude nonpersistent buffers, matching state_dict().
+                for name, tensor in get_layer_tensors(layer).items():
+                    if name not in layer._non_persistent_buffers_set:
+                        initialize_single_dummy_weight(tensor)
 
     def _process_online_quant_layer(
         self,
diff --git a/vllm/model_executor/model_loader/ep_weight_filter.py b/vllm/model_executor/model_loader/ep_weight_filter.py
index e9152f2cdb2a..0e23dfc11372 100644
--- a/vllm/model_executor/model_loader/ep_weight_filter.py
+++ b/vllm/model_executor/model_loader/ep_weight_filter.py
@@ -43,8 +43,12 @@ def compute_local_expert_ids(
     :func:`vllm.model_executor.layers.fused_moe.layer.determine_expert_map`.
 
     Args:
+        num_experts: Total number of experts in the layer.
+        ep_size: Expert parallel world size.
+        ep_rank: Rank whose local expert ids are computed.
         placement: ``"linear"`` for contiguous assignment,
             ``"round_robin"`` for interleaved assignment.
+
     """
     if ep_size <= 1:
         return None
diff --git a/vllm/model_executor/model_loader/reload/__init__.py b/vllm/model_executor/model_loader/reload/__init__.py
index 61dc1b66e6f3..a1ac282d290f 100644
--- a/vllm/model_executor/model_loader/reload/__init__.py
+++ b/vllm/model_executor/model_loader/reload/__init__.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Layerwise weight reloading utilities for vLLM.
+"""Layerwise weight reloading utilities for vLLM.
 
 This module provides functionality to reload model weights layer-by-layer,
 which is useful for weight updates without full model reconstruction.
diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py
index 5a73738e1511..062bfa60906b 100644
--- a/vllm/model_executor/model_loader/reload/layerwise.py
+++ b/vllm/model_executor/model_loader/reload/layerwise.py
@@ -54,8 +54,7 @@
 
 
 def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo:
-    """
-    Get information related to restoring and layerwise processing. If no previous
+    """Get information related to restoring and layerwise processing. If no previous
     information existed, a new entry is constructed
     """
     if layer not in LAYERWISE_INFO:
@@ -68,8 +67,7 @@ def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo:
 
 
 def record_metadata_for_reloading(model: torch.nn.Module):
-    """
-    Record layer metadata needed for later reloading.
+    """Record layer metadata needed for later reloading.
 
     Stores parameter and buffer metadata as meta tensors for restoration.
     Must be called before `initialize_layerwise_reload`.
@@ -82,8 +80,7 @@ def record_metadata_for_reloading(model: torch.nn.Module):
 
 @torch.no_grad()
 def initialize_layerwise_reload(model: torch.nn.Module):
-    """
-    Set up layerwise weight loading with deferred processing.
+    """Set up layerwise weight loading with deferred processing.
 
     Must be called after `record_metadata_for_reloading`. This function:
     1. Saves current kernel tensors for later copying
@@ -120,13 +117,13 @@ def initialize_layerwise_reload(model: torch.nn.Module):
 
 
 def initialize_online_processing(layer: torch.nn.Module):
-    """
-    Wrap a layer's weight loaders with online processing loaders.
+    """Wrap a layer's weight loaders with online processing loaders.
     Called by either `initialize_layerwise_reload` or an online quantization scheme,
     prevents double wrapping in the case of online quantization + reloading
 
     Args:
         layer: layer whose parameter weight loaders will be wrapped
+
     """
     info = get_layerwise_info(layer)
 
@@ -226,8 +223,7 @@ def online_process_loader(*args, **kwargs):
 
 
 def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelConfig):
-    """
-    Apply processing to any layers which were not layerwise processed during loading.
+    """Apply processing to any layers which were not layerwise processed during loading.
     This includes attention layers and layers which have weight elements which are not
     loaded (due to padding).
 
@@ -237,6 +233,7 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
     Args:
         model: model to finalize processing for
         model_config: config needed for applying processing to attention layers
+
     """
     if hasattr(model, "_original_do_torchao_reload"):
         model._do_torchao_reload = model._original_do_torchao_reload
@@ -330,8 +327,7 @@ def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -
 
 
 def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
-    """
-    Finalize layer loading after all weights have been buffered.
+    """Finalize layer loading after all weights have been buffered.
 
     This function:
     1. Materializes the layer onto the target device
@@ -377,7 +373,7 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
 
 
 def _get_original_loader(tensor: torch.Tensor) -> Callable:
-    """Return the weight loader with any layerwise wrappers removed"""
+    """Return the weight loader with any layerwise wrappers removed."""
     loader = _get_weight_loader(tensor)
     while loader.__name__ == "online_process_loader":
         loader = loader.__wrapped__  # type: ignore[union-attr]
diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py
index a8f8023cf26d..9874ebd67be9 100644
--- a/vllm/model_executor/model_loader/reload/meta.py
+++ b/vllm/model_executor/model_loader/reload/meta.py
@@ -9,7 +9,7 @@
 
 from .sanitize import restore_layer_refs, sanitize_layer_refs
 from .types import LayerReloadingInfo, LayerTensors
-from .utils import get_layer_params_buffers, get_layer_tensors
+from .utils import get_layer_params_buffers, get_layer_tensors, get_tensor_load_numel
 
 __all__ = [
     "to_meta_tensor",
@@ -47,8 +47,7 @@ def to_meta_tensor(tensor: torch.Tensor) -> torch.Tensor:
 
 
 def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor:
-    """
-    Materialize a meta tensor into an actual tensor on the current device.
+    """Materialize a meta tensor into an actual tensor on the current device.
     Should be called within the torch device context for the given rank.
     """
     tensor = torch.empty_strided(
@@ -119,7 +118,7 @@ def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors:
 
 
 def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo):
-    """Restore a layer to model format with tensors on the meta device"""
+    """Restore a layer to model format with tensors on the meta device."""
     if layer.__class__.__name__ in SKIP_MODULES:
         return
 
@@ -155,8 +154,7 @@ def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo):
 
 
 class CopyCounter(TorchDispatchMode):
-    """
-    Tracks total number of elements modified with `copy_`.
+    """Tracks total number of elements modified with `copy_`.
 
     Useful for keeping track of weight loading where underlying weights can be
     arbitrarily transformed (such as with `narrow`) before calling copy.
@@ -182,8 +180,7 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
 def get_numel_loaded(
     weight_loader: Callable, args: inspect.BoundArguments
 ) -> tuple[int, object]:
-    """
-    Determine how many elements would be loaded by a weight loader call.
+    """Determine how many elements would be loaded by a weight loader call.
 
     Args:
         weight_loader: used to load weights
@@ -192,6 +189,7 @@ def get_numel_loaded(
     Returns:
         number of elements loaded by the weight loader, the return value of the
         weight loader
+
     """
     with CopyCounter() as counter:
         return_value = weight_loader(*args.args, **args.kwargs)
@@ -208,5 +206,5 @@ def get_numel_loaded(
     numel = counter.copied_numel
     param = args.arguments.get("param", None)
     if isinstance(param, torch.Tensor):
-        numel = min(numel, param.numel())
+        numel = min(numel, get_tensor_load_numel(param))
     return numel, return_value
diff --git a/vllm/model_executor/model_loader/reload/sanitize.py b/vllm/model_executor/model_loader/reload/sanitize.py
index 21c47a2257f5..4e2c780e90f0 100644
--- a/vllm/model_executor/model_loader/reload/sanitize.py
+++ b/vllm/model_executor/model_loader/reload/sanitize.py
@@ -11,8 +11,7 @@
 
 
 def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor:
-    """
-    Removes references to layer held by tensor attributes. Specifically, removes the
+    """Removes references to layer held by tensor attributes. Specifically, removes the
     `__self__` attribute of weight loader methods attached to the tensor.
 
     Used by `capture_layer_to_meta` to avoid circular references to layers in
@@ -26,6 +25,7 @@ def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.T
 
     Returns:
         sanitized tensor
+
     """
     for key, value in tensor.__dict__.items():
         if isinstance(value, MethodType) and value.__self__ is layer:
@@ -35,8 +35,7 @@ def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.T
 
 
 def restore_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor:
-    """
-    Restores references to layer held by tensor attributes.
+    """Restores references to layer held by tensor attributes.
 
     Used by `restore_layer_on_meta` to add back layer references, allowing for proper
     weight loading.
@@ -47,6 +46,7 @@ def restore_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Te
 
     Returns:
         sanitized tensor
+
     """
     for key, value in tensor.__dict__.items():
         if isinstance(value, MethodType) and value.__self__ is layer_ref_sentinel:
diff --git a/vllm/model_executor/model_loader/reload/torchao_decorator.py b/vllm/model_executor/model_loader/reload/torchao_decorator.py
index 7fbc1c32944a..dc8e0acc6eaf 100644
--- a/vllm/model_executor/model_loader/reload/torchao_decorator.py
+++ b/vllm/model_executor/model_loader/reload/torchao_decorator.py
@@ -27,8 +27,7 @@ def set_torchao_reload_attrs(model: torch.nn.Module, model_config: ModelConfig):
 
 
 def support_quantized_model_reload_from_hp_weights(original_load_weights: FunctionType):
-    """
-    Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
+    """Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
     reloading high precision (bfloat16/float16/float32) weight for an already quantized
     model, this involves restoring the weights to a high precision weights and
     then online quantize the weights.
diff --git a/vllm/model_executor/model_loader/reload/utils.py b/vllm/model_executor/model_loader/reload/utils.py
index b6d129a79403..203aa82af303 100644
--- a/vllm/model_executor/model_loader/reload/utils.py
+++ b/vllm/model_executor/model_loader/reload/utils.py
@@ -11,6 +11,7 @@
     "get_layer_tensors",
     "get_layer_params_buffers",
     "get_layer_size",
+    "get_tensor_load_numel",
     "has_device_tensors",
     "get_info_size",
 ]
@@ -30,6 +31,11 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:
     )
 
 
+def get_tensor_load_numel(tensor: torch.Tensor) -> int:
+    """Count checkpoint elements, excluding padding declared by the weight creator."""
+    return getattr(tensor, "weight_loader_numel", tensor.numel())
+
+
 def get_layer_size(layer: torch.nn.Module) -> int:
     """Calculate total number of elements across loadable tensors in a layer.
 
@@ -39,21 +45,21 @@ def get_layer_size(layer: torch.nn.Module) -> int:
     from .meta import SKIP_LOAD_TENSORS
 
     return sum(
-        tensor.numel()
+        get_tensor_load_numel(tensor)
         for name, tensor in get_layer_tensors(layer).items()
         if name not in SKIP_LOAD_TENSORS
     )
 
 
 def has_device_tensors(bound_args: BoundArguments) -> bool:
-    """
-    Return True if the loaded weights exist on an accelerator device
+    """Return True if the loaded weights exist on an accelerator device.
 
     Args:
         bound_args: args to load weights
 
     Returns:
         True if weights are on accelerator device
+
     """
     return any(
         isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu")
@@ -62,14 +68,14 @@ def has_device_tensors(bound_args: BoundArguments) -> bool:
 
 
 def get_info_size(info: LayerReloadingInfo) -> int:
-    """
-    Calculate the number of bytes used by loaded weights for a given layer
+    """Calculate the number of bytes used by loaded weights for a given layer.
 
     Args:
         info: layerwise info to get size of
 
     Returns:
         number of bytes used by loaded weights
+
     """
     return sum(
         value.nbytes
diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py
index 3f30a42fe406..2f18a2e3aca8 100644
--- a/vllm/model_executor/model_loader/runai_streamer_loader.py
+++ b/vllm/model_executor/model_loader/runai_streamer_loader.py
@@ -15,12 +15,12 @@
     download_weights_from_hf,
     runai_safetensors_weights_iterator,
 )
+from vllm.transformers_utils.repo_utils import resolve_revision
 from vllm.transformers_utils.runai_utils import is_runai_obj_uri, list_safetensors
 
 
 class RunaiModelStreamerLoader(BaseModelLoader):
-    """
-    Model loader that can load safetensors
+    """Model loader that can load safetensors
     files from local FS, S3, GCS, or Azure Blob Storage.
     """
 
@@ -83,12 +83,16 @@ def _prepare_weights(
         """Prepare weights for the model.
 
         If the model is not local, it will be downloaded."""
-
         is_object_storage_path = is_runai_obj_uri(model_name_or_path)
         is_local = os.path.isdir(model_name_or_path)
         safetensors_pattern = "*.safetensors"
         index_file = SAFE_WEIGHTS_INDEX_NAME
 
+        if not is_local and not is_object_storage_path:
+            # `model_weights` can point to another repo than the one `revision` was
+            # resolved for, which does not pin this one.
+            revision = resolve_revision(model_name_or_path, revision)
+
         hf_folder = (
             model_name_or_path
             if (is_local or is_object_storage_path)
@@ -127,7 +131,7 @@ def _get_weights_iterator(
         )
 
     def download_model(self, model_config: ModelConfig) -> None:
-        """Download model if necessary"""
+        """Download model if necessary."""
         self._prepare_weights(model_config.model, model_config.revision)
 
     def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py
index 3f57fe7e0265..d17bfdc6d495 100644
--- a/vllm/model_executor/model_loader/sharded_state_loader.py
+++ b/vllm/model_executor/model_loader/sharded_state_loader.py
@@ -27,8 +27,7 @@
 
 
 class ShardedStateLoader(BaseModelLoader):
-    """
-    Model loader that directly loads each worker's model state dict, which
+    """Model loader that directly loads each worker's model state dict, which
     enables a fast load path for large tensor-parallel models where each worker
     only needs to read its own shard rather than the entire checkpoint. See
     `examples/features/sharded_state/save_sharded_state_offline.py` for creating
@@ -57,8 +56,7 @@ def __init__(self, load_config: LoadConfig):
     def _filter_subtensors(
         tensors: dict[str, torch.Tensor],
     ) -> dict[str, torch.Tensor]:
-        """
-        Filter out all tensors that share the same memory or a subset of the
+        """Filter out all tensors that share the same memory or a subset of the
         memory of another tensor.
         """
         same_storage_groups: dict[Any, list[tuple[str, torch.Tensor]]] = (
diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py
index 3ef9b6bf2a59..bb48c5dc08af 100644
--- a/vllm/model_executor/model_loader/tensorizer.py
+++ b/vllm/model_executor/model_loader/tensorizer.py
@@ -18,7 +18,7 @@
 import torch
 from torch import nn
 from torch.utils._python_dispatch import TorchDispatchMode
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 import vllm.envs as envs
 from vllm.config import ModelConfig, ParallelConfig, VllmConfig, set_current_vllm_config
@@ -173,7 +173,7 @@ class TensorizerConfig(MutableMapping):
     deserialization_kwargs: dict[str, Any] | None = None
     _extra_serialization_attrs: dict[str, Any] | None = field(init=False, default=None)
     model_class: type[torch.nn.Module] | None = field(init=False, default=None)
-    hf_config: PretrainedConfig | None = field(init=False, default=None)
+    hf_config: PreTrainedConfig | None = field(init=False, default=None)
     dtype: str | torch.dtype | None = field(init=False, default=None)
     _is_sharded: bool = field(init=False, default=False)
     _fields: ClassVar[tuple[str, ...]]
@@ -403,8 +403,7 @@ def __init__(self, tensorizer_config: TensorizerConfig):
 
     @staticmethod
     def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
-        """Tensorizer CLI arguments"""
-
+        """Tensorizer CLI arguments."""
         # Tensorizer options arg group
         group = parser.add_argument_group(
             "tensorizer options",
@@ -593,8 +592,7 @@ def tensorizer_weights_iterator(
 
 
 def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool:
-    """
-    Infer if the model is a vLLM model by checking the weights for
+    """Infer if the model is a vLLM model by checking the weights for
     a vLLM tensorized marker.
 
     Args:
@@ -603,6 +601,7 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool:
 
     Returns:
         bool: True if the model is a vLLM model, False otherwise.
+
     """
     tensorizer_args = tensorizer_config._construct_tensorizer_args()
     deserializer = TensorDeserializer(
@@ -711,7 +710,7 @@ def tensorize_vllm_model(
     tensorizer_config: TensorizerConfig,
     generate_keyfile: bool = True,
 ):
-    """Utility to load a model and then serialize it with Tensorizer
+    """Utility to load a model and then serialize it with Tensorizer.
 
     Intended to be used separately from running a vLLM server since it
     creates its own Engine instance.
@@ -773,8 +772,7 @@ def shutdown_engine_core() -> None:
 
 
 def tensorize_lora_adapter(lora_path: str, tensorizer_config: TensorizerConfig):
-    """
-    Uses tensorizer to serialize a LoRA adapter. Assumes that the files
+    """Uses tensorizer to serialize a LoRA adapter. Assumes that the files
     needed to load a LoRA adapter are a safetensors-format file called
     adapter_model.safetensors and a json config file called adapter_config.json.
 
diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py
index e12910865ec2..465703f75983 100644
--- a/vllm/model_executor/model_loader/utils.py
+++ b/vllm/model_executor/model_loader/utils.py
@@ -12,10 +12,15 @@
 from typing_extensions import assert_never
 
 import vllm.envs as envs
-from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+from vllm.config import (
+    LoadConfig,
+    ModelConfig,
+    VllmConfig,
+    replace,
+    set_current_vllm_config,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import is_deferred_attention_layer
-from vllm.model_executor.layers.hpc import HpcModule
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
@@ -24,6 +29,9 @@
     record_metadata_for_reloading,
     set_torchao_reload_attrs,
 )
+from vllm.model_executor.model_loader.weight_cache.utils import (
+    is_draft_model_cacheable,
+)
 from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings
 from vllm.model_executor.models.interfaces import SupportsQuant
 from vllm.model_executor.utils import is_weights_pre_processed
@@ -35,6 +43,33 @@
 logger = init_logger(__name__)
 
 
+def get_draft_load_config(vllm_config: VllmConfig) -> LoadConfig:
+    """Get load config for the speculative draft model."""
+    speculative_config = vllm_config.speculative_config
+    if (
+        speculative_config is not None
+        and speculative_config.draft_load_config is not None
+    ):
+        return speculative_config.draft_load_config
+    load_config = vllm_config.load_config
+    if load_config is not None and load_config.load_format != "ipc_cache":
+        return load_config
+    kwargs = (
+        # Route the draft to the daemon's draft group.
+        {
+            "model_loader_extra_config": {
+                **load_config.model_loader_extra_config,
+                "is_draft": True,
+            }
+        }
+        if is_draft_model_cacheable(speculative_config)
+        # No daemon draft group for this method; load from disk instead of
+        # hitting the target daemon with a mismatching fingerprint.
+        else {"load_format": "auto", "model_loader_extra_config": {}}
+    )
+    return replace(load_config, **kwargs)
+
+
 @instrument(span_name="Initialize model")
 def initialize_model(
     vllm_config: VllmConfig,
@@ -154,15 +189,6 @@ def process_weights_after_loading(
             with device_loading_context(module, target_device):
                 module.process_weights_after_loading(model_config.dtype)
 
-    # Process HPC modules (HpcRopeNorm, etc.) that rely on
-    # process_weights_after_loading being called from the model's
-    # load_weights(). When using DummyModelLoader (e.g. profiling or
-    # sleep/wake_up reload), the model's load_weights() is not called, so we
-    # must handle HPC modules here generically.
-    for _, module in model.named_modules():
-        if isinstance(module, HpcModule):
-            module.process_weights_after_loading(model)
-
     # Model-level post-load hook, after the per-layer quant finalize.
     if hasattr(model, "process_weights_after_loading"):
         model.process_weights_after_loading()
@@ -288,8 +314,7 @@ def get_architecture_class_name(model_config: ModelConfig) -> str:
 def configure_quant_config(
     quant_config: QuantizationConfig, model_class: type[nn.Module]
 ):
-    """
-    Pass packed_modules_mapping by reference to quant_config so that
+    """Pass packed_modules_mapping by reference to quant_config so that
     quant_config can properly match fused modules
 
     Note that model attributes are passed by reference to quant_config,
diff --git a/vllm/model_executor/model_loader/weight_cache/__init__.py b/vllm/model_executor/model_loader/weight_cache/__init__.py
index 0cea5cf72a21..6655f8913623 100644
--- a/vllm/model_executor/model_loader/weight_cache/__init__.py
+++ b/vllm/model_executor/model_loader/weight_cache/__init__.py
@@ -1,28 +1,3 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader
-from vllm.model_executor.model_loader.weight_cache.protocol import (
-    CacheConfigMismatchError,
-    TensorEntry,
-    UnsupportedPlatformForIPCError,
-    UnsupportedQuantForIPCError,
-    WeightCacheKey,
-    WeightCacheUnavailableError,
-    check_ipc_platform_support,
-    check_ipc_quant_support,
-    is_ipc_quant_supported,
-)
-
-__all__ = [
-    "CacheConfigMismatchError",
-    "IpcModelLoader",
-    "TensorEntry",
-    "UnsupportedPlatformForIPCError",
-    "UnsupportedQuantForIPCError",
-    "WeightCacheKey",
-    "WeightCacheUnavailableError",
-    "check_ipc_platform_support",
-    "check_ipc_quant_support",
-    "is_ipc_quant_supported",
-]
diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py
index 083910abb4d9..70fb31d579eb 100644
--- a/vllm/model_executor/model_loader/weight_cache/daemon.py
+++ b/vllm/model_executor/model_loader/weight_cache/daemon.py
@@ -17,8 +17,63 @@
     vllm serve /path/to/model --tensor-parallel-size 4 \\
         --load-format ipc_cache
 
-Only tensor and expert parallelism are supported; pipeline and data
-parallelism are rejected at launch.
+Tensor, expert and data parallelism are supported; pipeline parallelism is
+rejected at launch.
+
+For multi-node tensor parallelism, run one launcher per node with a shared
+rendezvous so the global TP group forms across nodes (CUDA IPC handles are
+node-local, so each node serves only its local GPUs' shards). Reuse the same
+``--nnodes``/``--node-rank``/``--master-addr`` flags you pass the engine, plus a
+``--weight-cache-master-port`` distinct from the engine's ``--master-port``:
+
+    # node 0 (8 local GPUs)
+    python -m vllm.model_executor.model_loader.weight_cache.daemon \\
+        --model /path/to/model --tensor-parallel-size 16 \\
+        --nnodes 2 --node-rank 0 --master-addr 10.0.0.1 \\
+        --weight-cache-master-port 29600
+    # node 1 (8 local GPUs)
+    python -m vllm.model_executor.model_loader.weight_cache.daemon \\
+        --model /path/to/model --tensor-parallel-size 16 \\
+        --nnodes 2 --node-rank 1 --master-addr 10.0.0.1 \\
+        --weight-cache-master-port 29600
+
+The global TP rank of local GPU ``i`` on node ``r`` is
+``r * (tp_size // nnodes) + i``, matching vLLM's contiguous per-node rank
+assignment, so each engine worker maps its shard from the daemon on its own
+node.
+
+For data parallelism (e.g. a TP1 x DP16 x EP decode fleet) run one launcher per
+node with the engine's DP placement flags. Local GPU ``i`` serves DP rank
+``start_rank + i // tp_size`` and TP rank ``i % tp_size``, and all
+``dp_size * tp_size`` daemons form one world group on
+``--data-parallel-address``/``--weight-cache-master-port`` so the expert
+shards are laid out exactly as in the engine:
+
+    # node r (4 local GPUs)
+    python -m vllm.model_executor.model_loader.weight_cache.daemon \\
+        --model /path/to/model --tensor-parallel-size 1 --enable-expert-parallel \\
+        --data-parallel-size 16 --data-parallel-size-local 4 \\
+        --data-parallel-start-rank 4r --data-parallel-address 10.0.0.1 \\
+        --weight-cache-master-port 29600
+
+Data parallelism also combines with multi-node tensor parallelism: pass both
+flag sets (``--nnodes``/``--node-rank``/``--master-addr`` and the DP flags with
+a reachable ``--data-parallel-address``). Each node then serves a contiguous
+block of the ``dp_size * tp_size`` global ranks, e.g. TP8 x DP2 on 4 nodes:
+
+    # node r (4 local GPUs)
+    python -m vllm.model_executor.model_loader.weight_cache.daemon \\
+        --model /path/to/model --tensor-parallel-size 8 --enable-expert-parallel \\
+        --nnodes 4 --node-rank r --master-addr 10.0.0.1 \\
+        --data-parallel-size 2 --data-parallel-address 10.0.0.1 \\
+        --weight-cache-master-port 29600
+
+With MTP, EAGLE or EAGLE3 speculative decoding the launcher additionally
+starts a draft daemon group that caches the draft model. It uses its own cache
+key, Unix sockets (``*_draft.sock``) and rendezvous port
+(``--weight-cache-draft-master-port``, default ``--weight-cache-master-port +
+1``), so each process serves exactly one model role. Other draft types are not
+cached and keep loading from disk in the engine.
 """
 
 import contextlib
@@ -30,15 +85,24 @@
 import socket
 import sys
 from collections.abc import Callable
+from itertools import product
 
 import torch
 
-from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
+from vllm.config import (
+    ParallelConfig,
+    VllmConfig,
+    replace,
+    set_current_vllm_config,
+)
 from vllm.distributed import (
     ensure_model_parallel_initialized,
     init_distributed_environment,
 )
+from vllm.engine.arg_utils import EngineArgs
 from vllm.logger import init_logger
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.model_loader.utils import process_weights_after_loading
 from vllm.model_executor.model_loader.weight_cache.protocol import (
     TensorEntry,
     WeightCacheKey,
@@ -46,13 +110,22 @@
     check_ipc_platform_support,
     check_ipc_quant_support,
     ensure_private_socket_dir,
-    get_physical_device_id,
+    get_current_device_uuid,
     get_socket_path,
     recv_msg,
     send_msg,
     verify_peer_is_owner,
 )
+from vllm.model_executor.model_loader.weight_cache.utils import (
+    export_model_attrs,
+    format_daemon_role,
+    is_draft_model_cacheable,
+)
 from vllm.platforms import current_platform
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm.utils.network_utils import get_distributed_init_method, get_open_port
+from vllm.utils.torch_utils import set_default_torch_dtype
+from vllm.v1.worker.workspace import init_workspace_manager
 
 logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon")
 
@@ -65,15 +138,19 @@ def export_entries(
     ``named_parameters``/``named_buffers`` are iterated with
     ``remove_duplicate=False`` so tied weights (e.g. ``lm_head.weight`` sharing
     storage with ``embed_tokens.weight``) are not silently dropped. Each unique
-    tensor is exported once; every additional name that refers to the same
+    tensor is exported once per call; every additional name that refers to the same
     tensor object is recorded in the returned alias map so the client can
     re-establish the shared identity instead of allocating uninitialized
     memory for it.
 
+    CUDA reduction arguments must be exported separately for each consumer so
+    that PyTorch registers a reference for each IPC mapping's lifetime.
+
     Returns:
         A ``(entries, aliases)`` pair where ``entries`` maps a canonical name to
         its ``TensorEntry`` and ``aliases`` maps each duplicate name to its
         canonical name.
+
     """
     entries: dict[str, TensorEntry] = {}
     aliases: dict[str, str] = {}
@@ -105,49 +182,108 @@ def __init__(
         self,
         vllm_config: VllmConfig,
         tp_rank: int,
+        local_rank: int,
         distributed_init_method: str,
         socket_dir: str | None = None,
+        is_draft: bool = False,
+        dp_rank: int = 0,
     ):
+        parallel_config = vllm_config.parallel_config
+        if parallel_config.data_parallel_size > 1:
+            # The engine sets these on each DP rank's workers; the MoE/EP
+            # layers read them to place their expert shards.
+            tp_size = parallel_config.tensor_parallel_size
+            if parallel_config.nnodes > 1:
+                world_size = parallel_config.data_parallel_size * tp_size
+                local_world = world_size // parallel_config.nnodes
+                start_dp_rank = parallel_config.node_rank * local_world // tp_size
+            else:
+                start_dp_rank = parallel_config.data_parallel_rank
+            parallel_config = replace(
+                parallel_config,
+                data_parallel_rank=dp_rank,
+                data_parallel_rank_local=dp_rank - start_dp_rank,
+            )
+            vllm_config = replace(vllm_config, parallel_config=parallel_config)
         self.vllm_config = vllm_config
+        # A draft daemon builds the draft with the target's VllmConfig, like
+        # the engine does; only the ModelConfig differs.
+        if is_draft:
+            spec = vllm_config.speculative_config
+            assert spec is not None and spec.draft_model_config is not None
+            self.model_config = spec.draft_model_config
+        else:
+            self.model_config = vllm_config.model_config
         self.tp_rank = tp_rank
+        self.dp_rank = dp_rank
+        self.tp_size = parallel_config.tensor_parallel_size
+        self.dp_size = parallel_config.data_parallel_size
+        self.global_rank = dp_rank * self.tp_size + tp_rank
+        self.world_size = self.dp_size * self.tp_size
+        self.local_rank = local_rank
         self.distributed_init_method = distributed_init_method
         self.socket_dir = socket_dir
+        self.is_draft = is_draft
+        self.role = format_daemon_role(is_draft)
         self.model: torch.nn.Module | None = None
-        self.entries: dict[str, TensorEntry] = {}
-        self.aliases: dict[str, str] = {}
         # Fingerprint before loading: process_weights_after_loading may
         # mutate hf_config.quantization_config.
         self.cache_config = WeightCacheKey.from_model_config(
-            vllm_config.model_config,
-            tp_size=vllm_config.parallel_config.tensor_parallel_size,
+            self.model_config,
+            tp_size=self.tp_size,
             tp_rank=tp_rank,
+            dp_size=self.dp_size,
+            dp_rank=dp_rank,
+            is_draft=is_draft,
         )
 
     def load_model(self) -> None:
-        from vllm.model_executor.model_loader import get_model
-
-        tp_size = self.cache_config.tp_size
-        torch.accelerator.set_device_index(self.tp_rank)
+        torch.accelerator.set_device_index(self.local_rank)
+        # Outside the config context so the daemon keeps its own rendezvous
+        # instead of the engine's; model parallel groups (TP/DP/EP) are then
+        # carved out of this world group like in the engine.
         init_distributed_environment(
-            world_size=tp_size,
-            rank=self.tp_rank,
+            world_size=self.world_size,
+            rank=self.global_rank,
             distributed_init_method=self.distributed_init_method,
-            local_rank=self.tp_rank,
+            local_rank=self.local_rank,
             backend=current_platform.dist_backend,
         )
         with set_current_vllm_config(self.vllm_config):
-            ensure_model_parallel_initialized(tp_size, 1)
-            self.model = get_model(vllm_config=self.vllm_config)
-        self._export_entries()
+            ensure_model_parallel_initialized(self.tp_size, 1)
+            self.model = self.get_model()
         logger.info(
-            "Weight cache daemon rank %d cached %d tensors",
-            self.tp_rank,
-            len(self.entries),
+            "Weight cache %s daemon rank %d loaded model",
+            self.role,
+            self.global_rank,
         )
 
-    def _export_entries(self) -> None:
-        assert self.model is not None
-        self.entries, self.aliases = export_entries(self.model)
+    def get_model(self) -> torch.nn.Module:
+        """Load the daemon's model, composed from the configured loader.
+
+        Runs the quantization check after model creation but before the slow
+        weight load, so an unsupported method fails fast. Online quantization
+        always fails the check, so load_model's finalize step for it is
+        unnecessary here.
+        """
+        vllm_config = self.vllm_config
+        model_config = self.model_config
+        load_config = vllm_config.load_config
+        loader = get_model_loader(load_config)
+        device_config = vllm_config.device_config
+        target_device = torch.device(
+            device_config.device if load_config.device is None else load_config.device
+        )
+        # Attention backends may take workspace at construction, like in the
+        # engine's worker init.
+        init_workspace_manager(target_device)
+        with set_default_torch_dtype(model_config.dtype):
+            with target_device:
+                model = loader.create_model(vllm_config, model_config)
+            check_ipc_quant_support(model)
+            loader.load_weights(model, model_config)
+            process_weights_after_loading(model, model_config, target_device)
+        return model.eval()
 
     def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> None:
         """Serve requests until terminated.
@@ -159,8 +295,9 @@ def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> Non
         Args:
             ready_callback: Invoked once the socket is bound and listening,
                 so the launcher can report overall readiness.
+
         """
-        socket_path = self._socket_path()
+        socket_path = self._socket_path
         ensure_private_socket_dir(
             os.path.dirname(socket_path), strict_perms=self.socket_dir is None
         )
@@ -174,7 +311,10 @@ def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> Non
         os.chmod(socket_path, 0o600)
         server.listen()
         logger.info(
-            "Weight cache daemon rank %d serving on %s", self.tp_rank, socket_path
+            "Weight cache %s daemon rank %d serving on %s",
+            self.role,
+            self.global_rank,
+            socket_path,
         )
         if ready_callback is not None:
             ready_callback()
@@ -219,12 +359,13 @@ def _acquire_gpu_lock(self, socket_path: str) -> int:
             ) from e
         return lock_fd
 
+    @property
     def _socket_path(self) -> str:
-        device_index = torch.accelerator.current_device_index()
-        gpu_id = get_physical_device_id(device_index)
-        if gpu_id is None:
-            gpu_id = device_index
-        return get_socket_path(gpu_id, self.socket_dir)
+        return get_socket_path(
+            get_current_device_uuid(),
+            self.socket_dir,
+            is_draft=self.is_draft,
+        )
 
     def _handle_connection(self, conn: socket.socket) -> None:
         request = recv_msg(conn)
@@ -246,73 +387,126 @@ def _handle_get_state(self, conn: socket.socket, request: dict) -> None:
             logger.warning("WeightCacheKey mismatch on fields: %s", mismatched)
             send_msg(conn, {"status": "mismatch", "fields": mismatched})
             return
-        if not self.entries:
+        if self.model is None:
             send_msg(conn, {"status": "error", "message": "Weights were released"})
             return
+        gpu_uuid = get_current_device_uuid()
+        entries, aliases = export_entries(self.model)
         send_msg(
             conn,
             {
                 "status": "ok",
-                "entries": self.entries,
-                "aliases": self.aliases,
-                "gpu_uuid": self._gpu_uuid(),
+                "entries": entries,
+                "aliases": aliases,
+                "attrs": export_model_attrs(self.model),
+                "gpu_uuid": gpu_uuid,
             },
         )
         logger.info_once(
-            "Weight cache daemon rank %d sent %d tensors (+%d aliases) to engine",
-            self.tp_rank,
-            len(self.entries),
-            len(self.aliases),
+            "Weight cache %s daemon rank %d sent %d tensors (+%d aliases) to engine",
+            self.role,
+            self.global_rank,
+            len(entries),
+            len(aliases),
         )
 
     def _handle_release(self, conn: socket.socket) -> None:
-        self.entries.clear()
-        self.aliases.clear()
         self.model = None
         torch.accelerator.empty_cache()
-        logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank)
-        send_msg(conn, {"status": "ok"})
-
-    def _gpu_uuid(self) -> str:
-        props = torch.cuda.get_device_properties(
-            torch.accelerator.current_device_index()
+        logger.info(
+            "Weight cache %s daemon rank %d released cached weights",
+            self.role,
+            self.global_rank,
         )
-        return str(props.uuid)
+        send_msg(conn, {"status": "ok"})
 
 
 def _run_daemon(
     tp_rank: int,
+    local_rank: int,
     vllm_config: VllmConfig,
     distributed_init_method: str,
     socket_dir: str | None,
-    ready_queue: "multiprocessing.Queue[int]",
+    ready_queue: "multiprocessing.Queue[tuple[str, int]]",
+    is_draft: bool = False,
+    dp_rank: int = 0,
 ) -> None:
     daemon = WeightCacheDaemon(
-        vllm_config, tp_rank, distributed_init_method, socket_dir
+        vllm_config,
+        tp_rank,
+        local_rank,
+        distributed_init_method,
+        socket_dir,
+        is_draft,
+        dp_rank,
     )
     daemon.load_model()
-    daemon.serve_forever(ready_callback=lambda: ready_queue.put(tp_rank))
+    daemon.serve_forever(
+        ready_callback=lambda: ready_queue.put((daemon.role, daemon.global_rank))
+    )
+
+
+def plan_local_ranks(parallel_config: ParallelConfig) -> list[tuple[int, int, int]]:
+    """``(local_rank, dp_rank, tp_rank)`` for every GPU this launcher serves.
+
+    Global ranks enumerate DP then TP: ``global = dp_rank * tp_size + tp_rank``.
+    With ``--nnodes`` the engine hands each node a contiguous block of global
+    ranks, so node ``r`` serves ``node_rank * local + i``; without it a
+    launcher's block starts at ``--data-parallel-start-rank * tp_size``.
+    """
+    tp_size = parallel_config.tensor_parallel_size
+    world_size = parallel_config.data_parallel_size * tp_size
+    if parallel_config.nnodes > 1:
+        local_world_size = world_size // parallel_config.nnodes
+        base = parallel_config.node_rank * local_world_size
+    else:
+        local_world_size = parallel_config.data_parallel_size_local * tp_size
+        base = parallel_config.data_parallel_rank * tp_size
+    return [
+        (i, (base + i) // tp_size, (base + i) % tp_size)
+        for i in range(local_world_size)
+    ]
+
+
+def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None:
+    """VllmConfig for the draft daemon group"""
+    if not is_draft_model_cacheable(vllm_config.speculative_config):
+        return None
+    assert vllm_config.speculative_config is not None
+    return vllm_config.speculative_config.apply_draft_overrides(vllm_config)
 
 
 def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None:
-    """Reject parallelism modes other than tensor/expert parallelism."""
-    unsupported = {
-        "pipeline parallelism": parallel_config.pipeline_parallel_size > 1,
-        "data parallelism": parallel_config.data_parallel_size > 1,
-    }
-    for name, enabled in unsupported.items():
-        if enabled:
+    """Reject pipeline parallelism and placements the daemon cannot map."""
+    if parallel_config.pipeline_parallel_size > 1:
+        raise ValueError(
+            "The weight cache daemon only supports tensor, expert and data "
+            "parallelism; pipeline parallelism is not supported"
+        )
+    dp_size = parallel_config.data_parallel_size
+    tp_size = parallel_config.tensor_parallel_size
+    if parallel_config.nnodes > 1:
+        world_size = dp_size * tp_size
+        if world_size % parallel_config.nnodes != 0:
             raise ValueError(
-                f"The weight cache daemon only supports tensor and expert "
-                f"parallelism; {name} is not supported"
+                f"--nnodes ({parallel_config.nnodes}) must evenly divide the "
+                f"daemon world size ({world_size} = data-parallel-size "
+                f"{dp_size} x tensor-parallel-size {tp_size})"
             )
+        return
+    if dp_size == 1:
+        return
+    end_rank = (
+        parallel_config.data_parallel_rank + parallel_config.data_parallel_size_local
+    )
+    if end_rank > dp_size:
+        raise ValueError(
+            "--data-parallel-start-rank + --data-parallel-size-local "
+            f"({end_rank}) exceeds --data-parallel-size ({dp_size})"
+        )
 
 
 def main() -> None:
-    from vllm.engine.arg_utils import EngineArgs
-    from vllm.utils.argparse_utils import FlexibleArgumentParser
-    from vllm.utils.network_utils import get_distributed_init_method, get_open_port
-
     parser = FlexibleArgumentParser(
         description="Launch weight cache daemons (one per TP rank)."
     )
@@ -323,6 +517,23 @@ def main() -> None:
         default=None,
         help="Directory for the daemon Unix sockets (default: tempdir).",
     )
+    parser.add_argument(
+        "--weight-cache-master-port",
+        type=int,
+        default=None,
+        help="Rendezvous port for the daemon's own TP group. Must differ from "
+        "the engine's --master-port (the daemon holds its group open while "
+        "serving) and match across nodes. Required when --nnodes > 1; defaults "
+        "to a free port for single-node.",
+    )
+    parser.add_argument(
+        "--weight-cache-draft-master-port",
+        type=int,
+        default=None,
+        help="Rendezvous port for the MTP draft daemon group. Defaults to "
+        "--weight-cache-master-port + 1 for multi-node, or a free port for "
+        "single-node.",
+    )
     args = parser.parse_args()
     engine_args = EngineArgs.from_cli_args(args)
     vllm_config = engine_args.create_engine_config()
@@ -331,30 +542,97 @@ def main() -> None:
             "The weight cache daemon itself must load from disk; use the "
             "default --load-format"
         )
-    # Checked before loading anything: an unsupported quantization method would
-    # otherwise only surface in the engine, after a full load.
-    check_ipc_platform_support(where="daemon")
-    check_ipc_quant_support(vllm_config.model_config, where="daemon")
+    # Config-only so it can fail before any model loading; the quant method
+    # check needs the created model and runs in get_daemon_model.
+    check_ipc_platform_support()
     parallel_config = vllm_config.parallel_config
     _reject_unsupported_parallelism(parallel_config)
     tp_size = parallel_config.tensor_parallel_size
+    dp_size = parallel_config.data_parallel_size
+
+    placements = plan_local_ranks(parallel_config)
+    local_world_size = len(placements)
+    if dp_size == 1:
+        nnodes = parallel_config.nnodes
+        node_rank = parallel_config.node_rank
+        master_addr = parallel_config.master_addr
+    else:
+        master_addr = parallel_config.data_parallel_master_ip
+        if parallel_config.nnodes > 1:
+            nnodes = parallel_config.nnodes
+            node_rank = parallel_config.node_rank
+        else:
+            nnodes = dp_size // parallel_config.data_parallel_size_local
+            node_rank = parallel_config.data_parallel_rank // (
+                parallel_config.data_parallel_size_local
+            )
+        if nnodes > 1 and master_addr in ("127.0.0.1", "localhost"):
+            raise ValueError(
+                "Data parallelism across nodes requires a reachable "
+                "--data-parallel-address for the daemon rendezvous"
+            )
+
+    # The daemon forms its own world group and holds it open while serving, so
+    # it needs a rendezvous port distinct from the engine's. All nodes must
+    # agree on it; single-node can auto-pick a free port. The master address is
+    # the engine's --master-addr (TP across nodes) or --data-parallel-address.
+    if nnodes > 1 and args.weight_cache_master_port is None:
+        raise ValueError(
+            "--weight-cache-master-port is required when the daemons span nodes"
+        )
+    master_port = args.weight_cache_master_port or get_open_port()
+    distributed_init_method = get_distributed_init_method(master_addr, master_port)
+
+    # (is_draft, vllm_config, rendezvous) per daemon group.
+    groups: list[tuple[bool, VllmConfig, str]] = [
+        (False, vllm_config, distributed_init_method)
+    ]
+    draft_vllm_config = get_draft_daemon_config(vllm_config)
+    if draft_vllm_config is not None:
+        draft_master_port = args.weight_cache_draft_master_port or (
+            master_port + 1 if nnodes > 1 else get_open_port()
+        )
+        if draft_master_port == master_port:
+            raise ValueError(
+                "--weight-cache-draft-master-port must differ from "
+                "--weight-cache-master-port"
+            )
+        groups.append(
+            (
+                True,
+                draft_vllm_config,
+                get_distributed_init_method(master_addr, draft_master_port),
+            )
+        )
 
-    distributed_init_method = get_distributed_init_method("127.0.0.1", get_open_port())
     ctx = multiprocessing.get_context("spawn")
-    ready_queue: multiprocessing.Queue[int] = ctx.Queue()
+    ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue()
+    # Local index == device index; global rank enumerates DP then TP.
+    expected_ready = {
+        (format_daemon_role(is_draft), dp_rank * tp_size + tp_rank)
+        for (is_draft, _, _), (_, dp_rank, tp_rank) in product(groups, placements)
+    }
     procs = [
         ctx.Process(
             target=_run_daemon,
             args=(
-                rank,
-                vllm_config,
-                distributed_init_method,
+                tp_rank,
+                local_rank,
+                config,
+                init_method,
                 args.weight_cache_socket_dir,
                 ready_queue,
+                is_draft,
+                dp_rank,
             ),
-            name=f"vllm-weight-cache-daemon-{rank}",
+            name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-"
+            f"{dp_rank * tp_size + tp_rank}",
         )
-        for rank in range(tp_size)
+        for (is_draft, config, init_method), (
+            local_rank,
+            dp_rank,
+            tp_rank,
+        ) in product(groups, placements)
     ]
     for proc in procs:
         proc.start()
@@ -366,10 +644,10 @@ def _shutdown(signum, frame):
     signal.signal(signal.SIGINT, _shutdown)
     signal.signal(signal.SIGTERM, _shutdown)
 
-    ready_ranks: set[int] = set()
-    while len(ready_ranks) < tp_size:
+    ready: set[tuple[str, int]] = set()
+    while len(ready) < len(expected_ready):
         try:
-            ready_ranks.add(ready_queue.get(timeout=1.0))
+            ready.add(ready_queue.get(timeout=1.0))
         except queue.Empty:
             dead = [p for p in procs if p.exitcode is not None]
             if dead:
@@ -383,11 +661,17 @@ def _shutdown(signum, frame):
                 for proc in procs:
                     proc.join()
                 sys.exit(max((p.exitcode or 0) for p in procs))
-    logger.info(
-        "Weight cache daemon ready: all %d ranks serving in %s",
-        tp_size,
-        args.weight_cache_socket_dir or "the default socket dir",
+    socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
+    logger.info_once(
+        "===== Weight cache daemon READY: node %d/%d serving %d local rank(s) "
+        "x %d role(s) in %s =====",
+        node_rank,
+        nnodes,
+        local_world_size,
+        len(groups),
+        socket_dir_msg,
     )
+
     for proc in procs:
         proc.join()
     sys.exit(max(proc.exitcode or 0 for proc in procs))
diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py
index 2132a55fbea5..8e37ce3ad6a8 100644
--- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py
@@ -12,6 +12,11 @@
 
 from vllm.config import ModelConfig, VllmConfig
 from vllm.config.load import LoadConfig
+from vllm.distributed import (
+    get_dp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
 from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
@@ -21,13 +26,13 @@
 )
 from vllm.model_executor.model_loader.weight_cache.protocol import (
     CacheConfigMismatchError,
-    TensorEntry,
     UnsupportedQuantForIPCError,
     WeightCacheKey,
+    WeightCacheState,
     WeightCacheUnavailableError,
     check_ipc_platform_support,
     check_ipc_quant_support,
-    get_physical_device_id,
+    get_current_device_uuid,
     get_socket_path,
     recv_msg,
     send_msg,
@@ -56,7 +61,7 @@ class IpcModelLoader(BaseModelLoader):
     Extra config keys (via --model-loader-extra-config):
 
     - socket_path: explicit daemon socket path. Defaults to a per-GPU path
-      derived from the physical GPU id.
+      derived from the physical GPU uuid and the cache role (target/draft).
     - socket_dir: directory containing the daemon sockets.
     - mode: "zero_copy" (default) or "copy".
     - fallback: fall back to disk loading when the daemon is unavailable or
@@ -74,6 +79,15 @@ def __init__(self, load_config: LoadConfig):
         extra_config = copy(load_config.model_loader_extra_config or {})
         self.socket_path: str | None = extra_config.pop("socket_path", None)
         self.socket_dir: str | None = extra_config.pop("socket_dir", None)
+        # Internal: set by the engine when routing a speculative draft to the
+        # daemon's draft group.
+        self.is_draft = bool(extra_config.pop("is_draft", False))
+        if self.is_draft and self.socket_path is not None:
+            raise ValueError(
+                "socket_path cannot be combined with the draft weight cache role; "
+                "use socket_dir so the target and draft sockets are derived "
+                "independently"
+            )
         self.mode: str = extra_config.pop("mode", "zero_copy")
         self.fallback: bool = extra_config.pop("fallback", True)
         self.connect_timeout_s: float = float(
@@ -104,7 +118,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
         loaded through this loader).
         """
         device_index = torch.accelerator.current_device_index()
-        entries, _ = self._fetch_entries(model_config)
+        entries = self._fetch_entries(model_config).entries
         params = dict(model.named_parameters())
         buffers = dict(model.named_buffers())
         for name, entry in entries.items():
@@ -119,22 +133,38 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
     def load_model(
         self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
     ) -> nn.Module:
-        # Unsupported quantization is a permanent misconfiguration rather than a
-        # transient daemon outage, so it is raised even when fallback is on.
-        self._check_supported(vllm_config, model_config)
+        # An unsupported platform is a permanent misconfiguration rather than
+        # a transient daemon outage, so it is raised even when fallback is on.
+        check_ipc_platform_support()
         state_fetched = False
         try:
-            entries, aliases = self._fetch_entries(model_config)
+            # Cross-check the routing flag against the identity of the model
+            # being loaded: a draft load that lost its flag (or a target load
+            # that got one) would hit the wrong daemon group and
+            # fingerprint-mismatch.
+            spec = vllm_config.speculative_config
+            inferred = spec is not None and model_config is spec.draft_model_config
+            if inferred != self.is_draft:
+                raise CacheConfigMismatchError(
+                    f"Weight cache role mismatch: loading "
+                    f"{'draft' if inferred else 'target'} model but the loader "
+                    f"was configured for the "
+                    f"{'draft' if self.is_draft else 'target'} group"
+                )
+            state = self._fetch_entries(model_config)
             state_fetched = True
-            return self._build_model(
-                vllm_config, model_config, prefix, entries, aliases
-            )
+            return self._build_model(vllm_config, model_config, prefix, state)
         except (WeightCacheUnavailableError, CacheConfigMismatchError) as e:
             if not self.fallback:
                 raise
             logger.warning(
                 "Weight cache unusable (%s); falling back to disk loading", e
             )
+        except UnsupportedQuantForIPCError:
+            # Unsupported quantization is a permanent misconfiguration rather
+            # than a transient daemon outage, so it is raised even when
+            # fallback is on.
+            raise
         except Exception:
             if not self.fallback:
                 raise
@@ -154,8 +184,7 @@ def _build_model(
         vllm_config: VllmConfig,
         model_config: ModelConfig,
         prefix: str,
-        entries: dict[str, TensorEntry],
-        aliases: dict[str, str],
+        state: WeightCacheState,
     ) -> nn.Module:
         device_config = vllm_config.device_config
         load_device = (
@@ -176,7 +205,12 @@ def _build_model(
                     model_config=model_config,
                     prefix=prefix,
                 )
-            self._apply_entries(model, entries, aliases, device_index)
+            check_ipc_quant_support(model)
+            self._apply_entries(model, state, device_index)
+            # Flags that load_weights would have set (e.g. EAGLE ownership of
+            # embed_tokens / lm_head); the daemon ran it, this process did not.
+            for name, value in state.attrs.items():
+                setattr(model, name, value)
             # The daemon exports tensors that already went through
             # process_weights_after_loading; re-run it in pre-processed mode
             # so quant methods only rebuild Python-side state (e.g. the MoE
@@ -192,32 +226,15 @@ def _build_model(
             self._send_release()
         logger.info(
             "Mapped %d tensors from the weight cache daemon (%s mode)",
-            len(entries),
+            len(state.entries),
             self.mode,
         )
         return model.eval()
 
-    @staticmethod
-    def _check_supported(vllm_config: VllmConfig, model_config: ModelConfig) -> None:
-        check_ipc_platform_support(where="engine")
-        check_ipc_quant_support(model_config, where="engine")
-        cache_dtype = vllm_config.cache_config.cache_dtype
-        if cache_dtype != "auto" and not str(cache_dtype).startswith("fp8"):
-            # BaseKVCacheMethod.process_weights_after_loading turns the loaded
-            # k/v scale parameters into plain float attributes. For fp8 cache
-            # dtypes those are rebuilt from the exported scale buffers when
-            # process_weights_after_loading runs in pre-processed mode; other
-            # quantized cache dtypes are not verified.
-            raise UnsupportedQuantForIPCError(
-                f"[weight_cache:engine] kv cache dtype {cache_dtype!r} is not "
-                "supported by the weight cache; use --kv-cache-dtype auto."
-            )
-
     def _apply_entries(
         self,
         model: nn.Module,
-        entries: dict[str, TensorEntry],
-        aliases: dict[str, str],
+        state: WeightCacheState,
         device_index: int,
     ) -> None:
         # remove_duplicate=False keeps tied module aliases reachable by name:
@@ -248,7 +265,7 @@ def _register(name: str, tensor: torch.Tensor, is_param: bool) -> None:
                 module.register_buffer(leaf, obj)
             registered[name] = obj
 
-        for name, entry in entries.items():
+        for name, entry in state.entries.items():
             tensor = entry.rebuild(device_index)
             if self.mode == "copy":
                 tensor = tensor.clone()
@@ -257,7 +274,7 @@ def _register(name: str, tensor: torch.Tensor, is_param: bool) -> None:
         # Re-establish tied-weight aliases by registering the *same* object the
         # canonical name resolved to, so parameter identity (and the tie) is
         # preserved instead of allocating uninitialized memory.
-        for alias_name, canonical_name in aliases.items():
+        for alias_name, canonical_name in state.aliases.items():
             obj = registered.get(canonical_name)
             if obj is None:
                 logger.warning(
@@ -268,24 +285,19 @@ def _register(name: str, tensor: torch.Tensor, is_param: bool) -> None:
                 continue
             _register(alias_name, obj, isinstance(obj, nn.Parameter))
 
-    def _fetch_entries(
-        self, model_config: ModelConfig
-    ) -> tuple[dict[str, TensorEntry], dict[str, str]]:
-        from vllm.distributed import (
-            get_tensor_model_parallel_rank,
-            get_tensor_model_parallel_world_size,
-        )
-
+    def _fetch_entries(self, model_config: ModelConfig) -> WeightCacheState:
+        dp_group = get_dp_group()
         cache_config = WeightCacheKey.from_model_config(
             model_config,
             tp_size=get_tensor_model_parallel_world_size(),
             tp_rank=get_tensor_model_parallel_rank(),
+            dp_size=dp_group.world_size,
+            dp_rank=dp_group.rank_in_group,
+            is_draft=self.is_draft,
         )
         return self._request_state(cache_config)
 
-    def _request_state(
-        self, cache_config: WeightCacheKey
-    ) -> tuple[dict[str, TensorEntry], dict[str, str]]:
+    def _request_state(self, cache_config: WeightCacheKey) -> WeightCacheState:
         with self._connect(self.state_timeout_s) as conn:
             send_msg(conn, {"cmd": "get_state", "cache_config": cache_config})
             response = recv_msg(conn)
@@ -299,7 +311,11 @@ def _request_state(
                 f"Weight cache daemon error: {response.get('message')}"
             )
         self._check_gpu_uuid(response.get("gpu_uuid"))
-        return response["entries"], response.get("aliases", {})
+        return WeightCacheState(
+            entries=response["entries"],
+            aliases=response.get("aliases", {}),
+            attrs=response.get("attrs", {}),
+        )
 
     def _connect(self, timeout: float) -> socket.socket:
         socket_path = self._resolve_socket_path()
@@ -327,22 +343,16 @@ def _connect(self, timeout: float) -> socket.socket:
     def _resolve_socket_path(self) -> str:
         if self.socket_path is not None:
             return self.socket_path
-        device_index = torch.accelerator.current_device_index()
-        gpu_id = get_physical_device_id(device_index)
-        if gpu_id is None:
-            raise WeightCacheUnavailableError(
-                "Cannot infer the physical GPU id from CUDA_VISIBLE_DEVICES; "
-                "pass socket_path via --model-loader-extra-config"
-            )
-        return get_socket_path(gpu_id, self.socket_dir)
+        return get_socket_path(
+            get_current_device_uuid(),
+            self.socket_dir,
+            is_draft=self.is_draft,
+        )
 
     def _check_gpu_uuid(self, daemon_uuid: str | None) -> None:
         if daemon_uuid is None:
             return
-        props = torch.cuda.get_device_properties(
-            torch.accelerator.current_device_index()
-        )
-        local_uuid = str(props.uuid)
+        local_uuid = get_current_device_uuid()
         if daemon_uuid != local_uuid:
             raise CacheConfigMismatchError(
                 f"Daemon GPU {daemon_uuid} != engine GPU {local_uuid}; "
@@ -361,7 +371,9 @@ def _fallback_load_config(self) -> LoadConfig:
         # DefaultModelLoader must not see load_format="ipc_cache" or the ipc
         # extra config keys.
         return dataclasses.replace(
-            self.load_config, load_format="auto", model_loader_extra_config={}
+            self.load_config,
+            load_format="auto",
+            model_loader_extra_config={},
         )
 
     def _fallback_load(
diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py
index 08e5e725bda7..6b5bbc69e5c8 100644
--- a/vllm/model_executor/model_loader/weight_cache/protocol.py
+++ b/vllm/model_executor/model_loader/weight_cache/protocol.py
@@ -20,16 +20,25 @@
 import struct
 import tempfile
 from dataclasses import dataclass, fields
-from typing import Any
+from typing import Any, NamedTuple
 
 import torch
+from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
 import vllm.version
 from vllm.config import ModelConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
+from vllm.model_executor.model_loader.weight_cache.utils import (
+    format_socket_role_suffix,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    filter_duplicate_safetensors_files,
+)
 from vllm.platforms import current_platform
 from vllm.utils.hashing import safe_hash
 
-SOCKET_NAME_TEMPLATE = "vllm_weight_cache_gpu{gpu_id}.sock"
+SOCKET_NAME_TEMPLATE = "vllm_weight_cache_{gpu_uuid}{role}.sock"
 SOCKET_DIR_TEMPLATE = "vllm_weight_cache_{uid}"
 
 _LEN_STRUCT = struct.Struct("!Q")
@@ -59,127 +68,53 @@ class UnsupportedPlatformForIPCError(Exception):
     """Raised when the current platform cannot share CUDA IPC handles."""
 
 
-def check_ipc_platform_support(*, where: str) -> None:
+def check_ipc_platform_support() -> None:
     """Hard-error unless the current platform can share CUDA IPC handles.
 
     Only CUDA/ROCm tensors get a real IPC handle from ``TensorEntry``; other
     platforms (e.g. XPU) would silently ship every tensor by value instead.
 
-    Args:
-        where: Short tag ("daemon"/"engine") used in the error message.
-
     Raises:
         UnsupportedPlatformForIPCError: If the current platform is not
             CUDA/ROCm.
+
     """
     if current_platform.is_cuda_alike():
         return
     raise UnsupportedPlatformForIPCError(
-        f"[weight_cache:{where}] platform {current_platform.device_name!r} "
-        "does not support CUDA IPC weight sharing; only CUDA and ROCm are "
-        "supported. Use the default --load-format for this platform."
+        f"platform {current_platform.device_name!r} does not support CUDA IPC "
+        "weight sharing; only CUDA and ROCm are supported. Use the default "
+        "--load-format for this platform."
     )
 
 
-# The daemon exports tensor data only, so sharing is correct just for methods
-# whose post-load effect is either fully captured by that data or rebuilt when
-# process_weights_after_loading runs under weights_already_processed. Anything
-# absent from this registry hard-errors: a method that silently repacks into
-# shapes the client cannot reproduce or stamps unrebuildable Python-side state
-# would serve wrong numerics. Extend it only after an end-to-end check against
-# a disk-loaded baseline.
-#
-# This check is config-only so the daemon fails fast before a slow full load,
-# rather than surfacing the incompatibility only after the engine has built the
-# model (where process_weights_after_loading's supports_pre_processed_weights
-# guard would eventually catch it).
-
-
-def _quant_config_field(quant_config: Any, key: str) -> Any:
-    if quant_config is None:
-        return None
-    if isinstance(quant_config, dict):
-        return quant_config.get(key)
-    return getattr(quant_config, key, None)
-
-
-def _fp8_round_trips_via_ipc(quant_config: Any) -> bool:
-    """Only block-wise FP8 is verified.
-
-    Block-wise FP8 preserves the weight shape, while per-tensor FP8 transposes
-    ``layer.weight`` during post-processing -- a shape the client's
-    meta-initialized model cannot reproduce.
-    """
-    return _quant_config_field(quant_config, "weight_block_size") is not None
-
-
-# quantization name -> predicate(quant_config) -> True when verified safe.
-IPC_QUANT_ALLOWLIST: dict[str | None, Any] = {
-    None: lambda _quant_config: True,  # unquantized
-    "fp8": _fp8_round_trips_via_ipc,
-    "modelopt_fp4": lambda _quant_config: True,
-    # Kimi-K3 routed experts (mxfp4-pack): the MegaMoE experts keep their packed
-    # weights and scales as plain parameters and defer the DeepGEMM transform to
-    # forward(), so the exported tensors are complete. Any other mxfp4 method
-    # that repacks or drops parameters still fails closed at the engine's
-    # supports_pre_processed_weights guard.
-    "mxfp4": lambda _quant_config: True,
-}
-
-
-def is_ipc_quant_supported(quantization: str | None, quant_config: Any) -> bool:
-    predicate = IPC_QUANT_ALLOWLIST.get(quantization)
-    return False if predicate is None else bool(predicate(quant_config))
-
-
-def check_ipc_quant_support(model_config: ModelConfig, *, where: str) -> None:
-    """Hard-error unless the model's quantization is verified for IPC sharing.
+# The daemon transfers post processed weights directly
+def check_ipc_quant_support(model: torch.nn.Module) -> None:
+    """Hard-error unless every quant method supports pre-processed weights.
 
     Args:
-        model_config: Model configuration to inspect.
-        where: Short tag ("daemon"/"engine") used in the error message.
+        model: The model to inspect (weights need not be loaded).
 
     Raises:
-        UnsupportedQuantForIPCError: If the quantization method is not on the
-            verified allowlist.
-    """
-    quantization = model_config.quantization
-    # Prefer the canonical, nested-aware config (multimodal models keep it under
-    # text_config); fall back to the raw hf_config for older code paths.
-    quant_config = getattr(
-        getattr(model_config, "model_arch_config", None), "quantization_config", None
-    )
-    if quant_config is None:
-        quant_config = getattr(model_config.hf_config, "quantization_config", None)
-    if is_ipc_quant_supported(quantization, quant_config):
-        return
-    verified = ", ".join(
-        "unquantized" if name is None else repr(name) for name in IPC_QUANT_ALLOWLIST
-    )
-    raise UnsupportedQuantForIPCError(
-        f"[weight_cache:{where}] quantization {quantization!r} is not verified "
-        f"for CUDA IPC weight sharing: its post-load processing may repack "
-        f"weights into shapes the client cannot reproduce or stamp Python-side "
-        f"state that tensor export cannot carry, which would silently serve "
-        f"wrong numerics. Verified: {verified} (FP8 only with weight_block_size "
-        f"set, i.e. block-wise). Use the default --load-format for this model."
-    )
+        UnsupportedQuantForIPCError: If any quant method does not declare
+            ``supports_pre_processed_weights``.
 
+    """
+    for name, module in model.named_modules():
+        quant_method = getattr(module, "quant_method", None)
+        if (
+            isinstance(quant_method, QuantizeMethodBase)
+            and not quant_method.supports_pre_processed_weights
+        ):
+            raise UnsupportedQuantForIPCError(
+                f"layer {name or '<root>'}: {type(quant_method).__name__} "
+                "does not support loading from pre-processed weights."
+            )
 
-def get_physical_device_id(device_index: int) -> int | None:
-    """Map a local CUDA device index to the physical GPU id.
 
-    Returns None if CUDA_VISIBLE_DEVICES contains non-integer entries
-    (e.g. GPU UUIDs), in which case an explicit socket path is required.
-    """
-    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
-    if not visible:
-        return device_index
-    entries = visible.split(",")
-    try:
-        return int(entries[device_index])
-    except (IndexError, ValueError):
-        return None
+def get_current_device_uuid() -> str:
+    """UUID of the physical GPU backing the current accelerator device."""
+    return current_platform.get_device_uuid(torch.accelerator.current_device_index())
 
 
 def get_socket_dir(socket_dir: str | None = None) -> str:
@@ -196,9 +131,23 @@ def get_socket_dir(socket_dir: str | None = None) -> str:
     )
 
 
-def get_socket_path(gpu_id: int, socket_dir: str | None = None) -> str:
-    directory = get_socket_dir(socket_dir)
-    return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_id=gpu_id))
+def get_socket_path(
+    gpu_uuid: str,
+    socket_dir: str | None = None,
+    *,
+    is_draft: bool = False,
+) -> str:
+    """Socket path of a daemon group; ``is_draft=False`` is the target.
+
+    The GPU uuid is hashed to keep the name well under the AF_UNIX path
+    limit (~108 bytes) even with the draft role suffix.
+    """
+    gpu_id = safe_hash(gpu_uuid.encode()).hexdigest()
+    name = SOCKET_NAME_TEMPLATE.format(
+        gpu_uuid=gpu_id,
+        role=format_socket_role_suffix(is_draft),
+    )
+    return os.path.join(get_socket_dir(socket_dir), name)
 
 
 def ensure_private_socket_dir(directory: str, strict_perms: bool = True) -> None:
@@ -309,12 +258,6 @@ def hash_checkpoint(model: str) -> str | None:
     """
     if not os.path.isdir(model):
         return None
-    from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
-
-    from vllm.model_executor.model_loader.weight_utils import (
-        filter_duplicate_safetensors_files,
-    )
-
     files = glob.glob(os.path.join(model, "*.safetensors"))
     if os.path.isfile(os.path.join(model, SAFE_WEIGHTS_INDEX_NAME)):
         files = filter_duplicate_safetensors_files(
@@ -346,10 +289,21 @@ class WeightCacheKey:
     quant_config_hash: str
     revision: str | None
     vllm_version: str
+    is_draft: bool = False
+    """Daemon group the weights come from; False is the target model."""
+    dp_size: int = 1
+    dp_rank: int = 0
 
     @classmethod
     def from_model_config(
-        cls, model_config: ModelConfig, tp_size: int, tp_rank: int
+        cls,
+        model_config: ModelConfig,
+        tp_size: int,
+        tp_rank: int,
+        *,
+        is_draft: bool = False,
+        dp_size: int = 1,
+        dp_rank: int = 0,
     ) -> "WeightCacheKey":
         """Build the fingerprint for a model configuration.
 
@@ -376,6 +330,9 @@ def from_model_config(
             quant_config_hash=_hash_quant_config(quant_config),
             revision=model_config.revision,
             vllm_version=vllm.version.__version__,
+            is_draft=is_draft,
+            dp_size=dp_size,
+            dp_rank=dp_rank,
         )
 
     def mismatched_fields(self, other: "WeightCacheKey") -> list[str]:
@@ -401,8 +358,6 @@ class TensorEntry:
 
     @classmethod
     def from_tensor(cls, tensor: torch.Tensor, kind: str) -> "TensorEntry":
-        from torch.multiprocessing.reductions import reduce_tensor
-
         tensor = tensor.detach()
         if tensor.is_cuda:
             _, ipc_args = reduce_tensor(tensor)
@@ -413,8 +368,6 @@ def rebuild(self, device_index: int) -> torch.Tensor:
         if self.ipc_args is None:
             assert self.cpu_tensor is not None
             return self.cpu_tensor
-        from torch.multiprocessing.reductions import rebuild_cuda_tensor
-
         args = list(self.ipc_args)
         # Index 6 of the args from reduce_tensor is the device index. It must
         # be retargeted to the local index since the daemon and the engine may
@@ -423,6 +376,17 @@ def rebuild(self, device_index: int) -> torch.Tensor:
         return rebuild_cuda_tensor(*args)
 
 
+class WeightCacheState(NamedTuple):
+    """Client-side decode of a daemon's get_state response payload."""
+
+    entries: dict[str, TensorEntry]
+    """Model tensors, exported as CUDA IPC handles or shipped by value."""
+    aliases: dict[str, str]
+    """Duplicate (tied) weight names aliased to their canonical entry."""
+    attrs: dict[str, bool]
+    """Python-side flags set by load_weights, e.g. EAGLE ownership flags."""
+
+
 def send_msg(sock: socket.socket, obj: Any) -> None:
     payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
     sock.sendall(_LEN_STRUCT.pack(len(payload)))
diff --git a/vllm/model_executor/model_loader/weight_cache/utils.py b/vllm/model_executor/model_loader/weight_cache/utils.py
new file mode 100644
index 000000000000..1abbf572f847
--- /dev/null
+++ b/vllm/model_executor/model_loader/weight_cache/utils.py
@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Helpers shared by the weight cache daemon, the IPC loader and the engine.
+
+These live outside ``protocol`` so callers that only need to know whether a
+draft is cached, or how a daemon group is named, do not have to import the
+wire format.
+"""
+
+from typing import Any
+
+from vllm.config import SpeculativeConfig
+from vllm.model_executor.models.interfaces import SupportsEagleBase
+
+# Speculative methods whose draft model the daemon caches in its own group.
+# Other drafts keep loading from disk in the engine.
+WEIGHT_CACHE_DRAFT_METHODS = frozenset({"mtp", "eagle", "eagle3"})
+
+# Python-side flags that weight loading sets on EAGLE-style drafts; the engine
+# never runs load_weights for cached models, so the daemon ships them.
+EXPORTED_MODEL_ATTRS = tuple(SupportsEagleBase.__annotations__)
+
+
+def export_model_attrs(model: Any) -> dict[str, bool]:
+    return {
+        name: bool(getattr(model, name))
+        for name in EXPORTED_MODEL_ATTRS
+        if hasattr(model, name)
+    }
+
+
+def is_draft_model_cacheable(speculative_config: SpeculativeConfig | None) -> bool:
+    """Whether the daemon serves the speculative draft as a separate role."""
+    return (
+        speculative_config is not None
+        and speculative_config.method in WEIGHT_CACHE_DRAFT_METHODS
+        and speculative_config.draft_model_config is not None
+    )
+
+
+def format_daemon_role(is_draft: bool) -> str:
+    """Name of a daemon group: the target model or the draft."""
+    return "draft" if is_draft else "target"
+
+
+def format_socket_role_suffix(is_draft: bool) -> str:
+    """Socket-name suffix keeping the draft group distinct from the target."""
+    return "_draft" if is_draft else ""
diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py
index 3085df85c6f8..2fa79a4e1c47 100644
--- a/vllm/model_executor/model_loader/weight_utils.py
+++ b/vllm/model_executor/model_loader/weight_utils.py
@@ -12,18 +12,17 @@
 import tempfile
 import threading
 import time
-from collections import defaultdict
 from collections.abc import Callable, Generator, Iterable
 from contextlib import contextmanager
 from pathlib import Path
-from typing import IO, Any
+from typing import IO, TYPE_CHECKING, Any
 
 import filelock
 import huggingface_hub.constants
 import numpy as np
 import regex as re
 import torch
-from safetensors.torch import load, load_file, safe_open, save_file
+from safetensors.torch import load, load_file, safe_open
 from tqdm.auto import tqdm
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
@@ -62,6 +61,9 @@
 
 from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least
 
+if TYPE_CHECKING:
+    from vllm.model_executor.models.utils import WeightsMapper
+
 logger = init_logger(__name__)
 
 # use system-level temp directory for file locks, so that multiple users
@@ -72,7 +74,7 @@
 
 
 def enable_xet_high_performance():
-    """automatically activates xet high performance mode"""
+    """Automatically activates xet high performance mode"""
     if "HF_XET_HIGH_PERFORMANCE" not in os.environ:
         huggingface_hub.constants.HF_XET_HIGH_PERFORMANCE = True
 
@@ -103,8 +105,7 @@ def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
 def atomic_writer(
     filepath: str | Path, mode: str = "w", encoding: str | None = None
 ) -> Generator[IO]:
-    """
-    Context manager that provides an atomic file writing routine.
+    """Context manager that provides an atomic file writing routine.
 
     The context manager writes to a temporary file and, if successful,
     atomically replaces the original file.
@@ -116,6 +117,7 @@ def atomic_writer(
 
     Yields:
         file object: A handle to the temporary file.
+
     """
     # Create a temporary file in the same directory as the target file
     # to ensure it's on the same filesystem for an atomic replace.
@@ -186,56 +188,6 @@ def maybe_download_from_modelscope(
     return None
 
 
-def _shared_pointers(tensors):
-    ptrs = defaultdict(list)
-    for k, v in tensors.items():
-        ptrs[v.data_ptr()].append(k)
-    failing = []
-    for _, names in ptrs.items():
-        if len(names) > 1:
-            failing.append(names)
-    return failing
-
-
-def convert_bin_to_safetensor_file(
-    pt_filename: str,
-    sf_filename: str,
-) -> None:
-    loaded = torch.load(pt_filename, map_location="cpu", weights_only=True)
-    if "state_dict" in loaded:
-        loaded = loaded["state_dict"]
-    shared = _shared_pointers(loaded)
-    for shared_weights in shared:
-        for name in shared_weights[1:]:
-            loaded.pop(name)
-
-    # For tensors to be contiguous
-    loaded = {k: v.contiguous() for k, v in loaded.items()}
-
-    dirname = os.path.dirname(sf_filename)
-    os.makedirs(dirname, exist_ok=True)
-    save_file(loaded, sf_filename, metadata={"format": "pt"})
-
-    # check file size
-    sf_size = os.stat(sf_filename).st_size
-    pt_size = os.stat(pt_filename).st_size
-    if (sf_size - pt_size) / pt_size > 0.01:
-        raise RuntimeError(
-            f"""The file size different is more than 1%:
-         - {sf_filename}: {sf_size}
-         - {pt_filename}: {pt_size}
-         """
-        )
-
-    # check if the tensors are the same
-    reloaded = load_file(sf_filename)
-    for k in loaded:
-        pt_tensor = loaded[k]
-        sf_tensor = reloaded[k]
-        if not torch.equal(pt_tensor, sf_tensor):
-            raise RuntimeError(f"The output tensors do not match for key {k}")
-
-
 # TODO(woosuk): Move this to other place.
 def get_quant_config(
     model_config: ModelConfig, load_config: LoadConfig
@@ -482,6 +434,7 @@ def download_weights_from_hf(
 
     Returns:
         str: The path to the downloaded model weights.
+
     """
     assert len(allow_patterns) > 0
     local_only = huggingface_hub.constants.HF_HUB_OFFLINE
@@ -580,6 +533,7 @@ def download_safetensors_index_file_from_hf(
         subfolder (Optional[str]): The subfolder within the model repository
             to download weights from.
         revision (Optional[str]): The revision of the model.
+
     """
     # Use file lock to prevent multiple processes from
     # downloading the same model weights at the same time.
@@ -637,8 +591,7 @@ def filter_duplicate_safetensors_files(
 
 
 def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]:
-    """
-    Exclude files that are not needed for inference.
+    """Exclude files that are not needed for inference.
 
     See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
     """
@@ -655,6 +608,171 @@ def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[s
     return hf_weights_files
 
 
+def _normalize_module_prefixes(
+    language_model_names: Iterable[str] | None,
+) -> tuple[str, ...]:
+    names = tuple(language_model_names or ())
+    return tuple(n if n.endswith(".") else f"{n}." for n in names)
+
+
+def _mapped_weight_name(weights_mapper: "WeightsMapper | None", key: str) -> str | None:
+    """Apply ``WeightsMapper.map_name`` when present; else identity."""
+    if weights_mapper is None:
+        return key
+    return weights_mapper.map_name(key)
+
+
+def _hf_prefix_maps_to_lm(
+    new_name: str | None, lm_module_prefixes: tuple[str, ...]
+) -> bool:
+    if not new_name:
+        return False
+    new_p = new_name if new_name.endswith(".") else f"{new_name}."
+    return any(new_p.startswith(p) or p.startswith(new_p) for p in lm_module_prefixes)
+
+
+def _shared_hf_root_module_prefixes(
+    lm_module_prefixes: tuple[str, ...],
+    weights_mapper: "WeightsMapper | None",
+) -> tuple[str, ...]:
+    """Module prefixes that are also an HF root for non-LM weights.
+
+    E.g. Molmo/Phi-4-MM/Muse mark LM as ``model`` while vision HF keys stay
+    under ``model.vision_*`` / ``model.embed_tokens_extend.*``. Detected via
+    ``WeightsMapper.orig_to_new_prefix`` (no hard-coded name list).
+    """
+    if weights_mapper is None:
+        return ()
+    mapping = weights_mapper.orig_to_new_prefix
+    if not mapping:
+        return ()
+
+    non_lm_origs: list[str] = []
+    for orig, new in mapping.items():
+        orig_p = orig if orig.endswith(".") else f"{orig}."
+        if not _hf_prefix_maps_to_lm(new, lm_module_prefixes):
+            non_lm_origs.append(orig_p)
+
+    if not non_lm_origs:
+        return ()
+
+    shared = [
+        p
+        for p in lm_module_prefixes
+        if any(orig.startswith(p) for orig in non_lm_origs)
+    ]
+    return tuple(shared)
+
+
+def resolve_mm_encoder_only_lm_prefixes(
+    language_model_names: Iterable[str] | None,
+    *,
+    weights_mapper: "WeightsMapper | None" = None,
+) -> tuple[str, ...] | None:
+    """Resolve vLLM LM *module* prefixes for ``--mm-encoder-only`` shard skip.
+
+    Prefixes come from ``_language_model_names``. Classification of HF index
+    keys is done later via optional ``weights_mapper`` (see
+    ``filter_mm_encoder_only_safetensors_files``), so this helper does **not**
+    hard-code Qwen/HF nest lists.
+
+    Returns ``None`` (leave the safetensors file list unchanged) when
+    ``_language_model_names`` is empty/missing, or when a module prefix is a
+    shared HF checkpoint root for both LM and non-LM weights — fail-closed so
+    Molmo / Phi-4-MM / Muse cannot under-load the encoder.
+    """
+    prefixes = _normalize_module_prefixes(language_model_names)
+    if not prefixes:
+        return None
+    shared = _shared_hf_root_module_prefixes(prefixes, weights_mapper)
+    if shared:
+        logger.warning_once(
+            "mm-encoder-only whole-shard filter not applied: LM module "
+            "prefix(es) %s share an HF checkpoint root with non-LM weights "
+            "(via WeightsMapper). Encoder-only load keeps the full "
+            "safetensors file list; finer HF prefixes or key-level "
+            "filtering can enable skip later.",
+            shared,
+        )
+        return None
+    return prefixes
+
+
+def filter_mm_encoder_only_safetensors_files(
+    hf_weights_files: list[str],
+    hf_folder: str,
+    index_file: str,
+    language_model_prefixes: Iterable[str],
+    *,
+    weights_mapper: "WeightsMapper | None" = None,
+) -> list[str]:
+    """Drop safetensors shards that only contain language-model weights.
+
+    Used with ``--mm-encoder-only`` so Encoder-only EPD instances avoid reading
+    pure LM shards from disk/DRAM.
+
+    Each HF index key is classified by mapping through ``weights_mapper`` (when
+    provided) and testing the *vLLM* name against ``language_model_prefixes``.
+    Without a mapper, HF names are compared directly (identity checkpoint
+    layout, e.g. Kimi ``language_model.*``).
+
+    A shard is kept if it contains any non-LM key. Without an index file,
+    returns ``hf_weights_files`` unchanged.
+    """
+    prefixes = tuple(language_model_prefixes)
+    if not prefixes:
+        return hf_weights_files
+
+    index_path = os.path.join(hf_folder, index_file)
+    if not os.path.isfile(index_path):
+        logger.warning_once(
+            "mm-encoder-only shard filter skipped: index file %s not found",
+            index_path,
+        )
+        return hf_weights_files
+
+    with open(index_path) as f:
+        weight_map: dict[str, str] = json.load(f)["weight_map"]
+
+    keys_by_file: dict[str, list[str]] = {}
+    for weight_name, weight_file in weight_map.items():
+        keys_by_file.setdefault(weight_file, []).append(weight_name)
+
+    def _is_lm_key(key: str) -> bool:
+        mapped = _mapped_weight_name(weights_mapper, key)
+        if mapped is None:
+            return False
+        return any(mapped.startswith(p) for p in prefixes)
+
+    def _is_lm_only(weight_file: str) -> bool:
+        keys = keys_by_file.get(weight_file)
+        if not keys:
+            # Not listed in the index (e.g. extra matched glob) — keep.
+            return False
+        return all(_is_lm_key(key) for key in keys)
+
+    kept: list[str] = []
+    skipped = 0
+    for path in hf_weights_files:
+        rel = os.path.relpath(path, hf_folder)
+        # Index weight_map values are usually basenames; tolerate either form.
+        basename = os.path.basename(path)
+        if _is_lm_only(rel) or _is_lm_only(basename):
+            skipped += 1
+            continue
+        kept.append(path)
+
+    if skipped:
+        logger.info_once(
+            "mm-encoder-only: skipped %d/%d safetensors shard(s) that only "
+            "contain language-model weights (prefixes=%s)",
+            skipped,
+            skipped + len(kept),
+            prefixes,
+        )
+    return kept
+
+
 # explicitly use pure text format, with a newline at the end
 # this makes it impossible to see the animation in the progress bar
 # but will avoid messing up with ray or multiprocessing, which wraps
@@ -727,11 +845,15 @@ def _get_available_ram_bytes() -> int:
 
     host_available = psutil.virtual_memory().available
 
-    from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit
+    from vllm.utils.cpu_resource_utils import (
+        get_cgroup_memory_limit,
+        get_cgroup_memory_usage,
+    )
 
-    cgroup_limit, cgroup_usage = get_cgroup_memory_limit()
+    cgroup_limit = get_cgroup_memory_limit()
     if cgroup_limit is None:
         return host_available
+    cgroup_usage = get_cgroup_memory_usage()
     cgroup_available = (
         cgroup_limit if cgroup_usage is None else max(0, cgroup_limit - cgroup_usage)
     )
@@ -1308,29 +1430,6 @@ def composed_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
     return composed_loader
 
 
-def initialize_dummy_weights(
-    model: torch.nn.Module,
-    model_config: ModelConfig,
-    low: float = -1e-3,
-    high: float = 1e-3,
-    seed: int = 1234,
-) -> None:
-    """Initialize model weights with random values.
-
-    The model weights must be randomly initialized for accurate performance
-    measurements. Additionally, the model weights should not cause NaNs in the
-    forward pass. We empirically found that initializing the weights with
-    values between -1e-3 and 1e-3 works well for most models.
-
-    We use per-parameter random seed, so that dummy weights are consistent,
-    even if the model is partitioned across multiple devices. When the seed
-    is fixed, the random values generated by this function only depends on
-    the parameter's number of elements and its data type.
-    """
-    for param in model.state_dict().values():
-        initialize_single_dummy_weight(param, low, high, seed)
-
-
 @torch.no_grad()
 def initialize_single_dummy_weight(
     param: torch.Tensor,
@@ -1341,6 +1440,10 @@ def initialize_single_dummy_weight(
     if param.device.type == "meta":
         return  # deferred to finalize_layerwise_processing (e.g. online quant)
 
+    if (dummy_weight_value := getattr(param, "dummy_weight_value", None)) is not None:
+        param.fill_(dummy_weight_value)
+        return
+
     if not torch.is_floating_point(param):
         if current_platform.is_rocm():
             # On ROCm, integer params (e.g. GPTQ qweight/qzeros) are left
@@ -1376,10 +1479,9 @@ def initialize_single_dummy_weight(
     generator.manual_seed(seed)
     if torch.finfo(param.data.dtype).bits < 16:
         # uniform_ doesn't support < 16-bit datatypes (FP8)
-        dtype = param.data.dtype
-        tmp_param = param.data.to(torch.float16)
-        tmp_param = tmp_param.uniform_(low, high, generator=generator).to(dtype)
-        param.data.copy_(tmp_param)
+        tmp_param = torch.empty_like(param, dtype=torch.float16)
+        tmp_param.uniform_(low, high, generator=generator)
+        param.copy_(tmp_param)
     else:
         param.uniform_(low, high, generator=generator)
 
@@ -1400,6 +1502,7 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
         str: The remapped parameter name if successful, or the original name
              if no remapping is needed.
         None: If the remapped name is not found in params_dict.
+
     """
     # Already in vLLM's expected form (e.g. weights pre-renamed by a
     # `WeightsMapper` from the quant config). Skip the regex remap, which
@@ -1499,8 +1602,7 @@ def maybe_remap_moe_expert_param_name(
     name: str,
     params_dict: dict[str, torch.nn.Parameter],
 ) -> str:
-    """
-    Remap MoE expert parameter names to account for routed_experts hierarchy.
+    """Remap MoE expert parameter names to account for routed_experts hierarchy.
 
     This handles the transition from the old FusedMoE structure where weights
     were directly in the experts module, to the new MoERunner → RoutedExperts
@@ -1522,6 +1624,7 @@ def maybe_remap_moe_expert_param_name(
     Returns:
         Remapped parameter name if routed_experts hierarchy exists,
         otherwise the original name
+
     """
     # Only remap if this looks like an expert parameter
     if ".experts." not in name:
@@ -1575,8 +1678,7 @@ def remap_moe_expert_weights(
     weights: Iterable[tuple[str, torch.Tensor]],
     params_dict: dict[str, torch.nn.Parameter],
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
-    """
-    Wrapper generator that remaps MoE expert parameter names for backward compatibility.
+    """Remap MoE expert parameter names for backward compatibility.
 
     This allows models with custom weight loading to automatically handle both old
     and new checkpoint formats without needing model-specific remapping code.
@@ -1594,6 +1696,7 @@ def remap_moe_expert_weights(
 
     Yields:
         (remapped_name, tensor) tuples
+
     """
     for name, weight in weights:
         remapped_name = maybe_remap_moe_expert_param_name(name, params_dict)
diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py
index 031808ee12b8..da38035db9bd 100644
--- a/vllm/model_executor/models/AXK1.py
+++ b/vllm/model_executor/models/AXK1.py
@@ -44,6 +44,7 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
@@ -128,11 +129,9 @@ def __init__(
                 "Only silu is supported for now."
             )
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.n_routed_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
         if config.topk_method == "noaux_tc":
@@ -324,9 +323,9 @@ def __init__(
         assert config.rope_parameters is not None
         if config.rope_parameters["rope_type"] != "default":
             config.rope_parameters["rope_type"] = (
-                "deepseek_yarn"
-                if config.rope_parameters.get("apply_yarn_scaling", True)
-                else "deepseek_llama_scaling"
+                "deepseek_llama_scaling"
+                if config.rope_parameters.get("attention_factor") == 1.0
+                else "deepseek_yarn"
             )
 
         self.rotary_emb = get_rope(
@@ -400,8 +399,7 @@ def forward(
 
 
 class AXK1MLAAttention(nn.Module):
-    """
-    Main reference: DeepseekV2 paper, and FlashInfer Implementation
+    """Main reference: DeepseekV2 paper, and FlashInfer Implementation
     (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
 
         For more info see MLACommonImpl in:
@@ -497,9 +495,9 @@ def __init__(
         assert config.rope_parameters is not None
         if config.rope_parameters["rope_type"] != "default":
             config.rope_parameters["rope_type"] = (
-                "deepseek_yarn"
-                if config.rope_parameters.get("apply_yarn_scaling", True)
-                else "deepseek_llama_scaling"
+                "deepseek_llama_scaling"
+                if config.rope_parameters.get("attention_factor") == 1.0
+                else "deepseek_yarn"
             )
 
         self.rotary_emb = get_rope(
diff --git a/vllm/model_executor/models/adapters.py b/vllm/model_executor/models/adapters.py
index ea75ff8cc0b1..cdbf628e610b 100644
--- a/vllm/model_executor/models/adapters.py
+++ b/vllm/model_executor/models/adapters.py
@@ -39,7 +39,6 @@
 
 def _load_st_projector(model_config: "ModelConfig") -> nn.Module | None:
     """Load Sentence-Transformers Dense projection layers."""
-
     dense_modules = try_get_dense_modules(
         model_config.model, revision=model_config.revision
     )
@@ -249,8 +248,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
 
 
 def as_embedding_model(cls: type[_T]) -> type[_T]:
-    """
-    Subclass an existing vLLM model to support embeddings.
+    """Subclass an existing vLLM model to support embeddings.
 
     By default, the embeddings of the whole prompt are extracted from the
     normalized hidden state corresponding to the last token.
@@ -258,6 +256,7 @@ def as_embedding_model(cls: type[_T]) -> type[_T]:
     Note:
         We assume that no extra layers are added to the original model;
         please implement your own model if this is not the case.
+
     """
     # Avoid modifying existing embedding models
     if is_pooling_model(cls):
@@ -285,7 +284,7 @@ def _init_pooler(
 def _resolve_num_labels(hf_config: Any, text_config: Any) -> int:
     """Resolve the label count for a sequence classification head.
 
-    ``PretrainedConfig.num_labels`` is derived from ``id2label``, which always
+    ``PreTrainedConfig.num_labels`` is derived from ``id2label``, which always
     carries a default of two entries. Composite configs (such as multimodal
     checkpoints) declare their label space on the top-level config, so reading
     ``num_labels`` from ``get_text_config()`` silently returns that default and
@@ -298,16 +297,15 @@ def _resolve_num_labels(hf_config: Any, text_config: Any) -> int:
     if text_config is hf_config:
         return hf_config.num_labels
 
-    from transformers import PretrainedConfig
+    from transformers import PreTrainedConfig
 
-    if hf_config.num_labels != PretrainedConfig().num_labels:
+    if hf_config.num_labels != PreTrainedConfig().num_labels:
         return hf_config.num_labels
     return text_config.num_labels
 
 
 def as_seq_cls_model(cls: type[_T]) -> type[_T]:
-    """
-    Subclass an existing vLLM model to support classify and score tasks.
+    """Subclass an existing vLLM model to support classify and score tasks.
 
     By default, the class probabilities are extracted from the softmaxed
     hidden state corresponding to the last token.
@@ -316,6 +314,7 @@ def as_seq_cls_model(cls: type[_T]) -> type[_T]:
         We assume that the classification head is a single linear layer
         stored as the attribute `score` of the top-level model;
         please implement your own model if this is not the case.
+
     """
     # Avoid modifying existing classification models
     if is_pooling_model(cls):
@@ -446,8 +445,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
 
 
 def _get_language_model_for_seq_cls(model: nn.Module) -> nn.Module:
-    """
-    Get the language model component for sequence classification conversion.
+    """Get the language model component for sequence classification conversion.
     For VLMs, returns the inner language model. For standard LLMs, returns model itself.
     """
     multimodal_model: object = model
@@ -481,8 +479,7 @@ def _get_language_model_for_seq_cls(model: nn.Module) -> nn.Module:
 
 @contextmanager
 def _disable_seq_cls_loading_on_inner_model(language_model, is_vlm: bool):
-    """
-    Context manager to temporarily disable sequence classification loading
+    """Context manager to temporarily disable sequence classification loading
     on inner VLM models to prevent recursive seq_cls_model_loader calls.
     """
     if not is_vlm:
diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py
index 808775e35fba..42643cb8c64b 100644
--- a/vllm/model_executor/models/afmoe.py
+++ b/vllm/model_executor/models/afmoe.py
@@ -19,6 +19,7 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     MoERunner,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -82,11 +83,12 @@ def __init__(
             )
 
         # Router gate
-        self.gate = nn.Linear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
-            dtype=torch.float32,
+            out_dtype=torch.float32,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.gate",
         )
         self.expert_bias = nn.Parameter(
             torch.empty(config.num_experts, dtype=torch.float32)
@@ -140,7 +142,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
 
-        router_logits = self.gate(hidden_states.to(dtype=torch.float32))
+        router_logits, _ = self.gate(hidden_states)
 
         final_hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py
index 3461b5713734..4e31b99944f1 100644
--- a/vllm/model_executor/models/aimv2.py
+++ b/vllm/model_executor/models/aimv2.py
@@ -26,7 +26,10 @@
 
 class AIMv2SwiGLUFFN(nn.Module):
     def __init__(
-        self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str
+        self,
+        config: AIMv2Config,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
     ):
         super().__init__()
         hidden_features = config.intermediate_size
@@ -91,7 +94,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class AIMv2Attention(nn.Module):
     def __init__(
-        self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str
+        self,
+        config: AIMv2Config,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
     ):
         super().__init__()
         self.config = config
@@ -144,7 +150,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class AIMv2Block(nn.Module):
     def __init__(
-        self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str
+        self,
+        config: AIMv2Config,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
     ):
         super().__init__()
         self.attn = AIMv2Attention(
@@ -166,7 +175,7 @@ class AIMv2Transformer(nn.Module):
     def __init__(
         self,
         config: AIMv2Config,
-        quant_config: QuantizationConfig,
+        quant_config: QuantizationConfig | None,
         *,
         require_post_norm: bool | None = None,
         prefix: str = "",
@@ -206,7 +215,7 @@ class AIMv2Model(torch.nn.Module):
     def __init__(
         self,
         config: AIMv2Config,
-        quant_config: QuantizationConfig,
+        quant_config: QuantizationConfig | None,
         *,
         require_post_norm: bool | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py
index 57eb62d3e26a..da91c9d9bf97 100644
--- a/vllm/model_executor/models/aria.py
+++ b/vllm/model_executor/models/aria.py
@@ -3,6 +3,7 @@
 from collections.abc import Iterable, Mapping, Sequence
 from typing import Annotated, Literal
 
+import regex as re
 import torch
 import torch.nn as nn
 from transformers import AriaConfig, AriaTextConfig, BatchFeature
@@ -10,7 +11,7 @@
 from transformers.models.aria.processing_aria import AriaProcessor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.fused_moe import FusedMoEFactory
@@ -41,13 +42,12 @@
 
 
 class AriaImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - n: Number of images
-        - c: Number of channels
-        - h: Height of each image
-        - w: Width of each image
+    """Dimensions:
+    - b: Batch size
+    - n: Number of images
+    - c: Number of channels
+    - h: Height of each image
+    - w: Width of each image
     """
 
     type: Literal["pixel_values"]
@@ -119,8 +119,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class AriaProjector(nn.Module):
-    """
-    A projection module with one cross attention layer and one FFN layer, which
+    """A projection module with one cross attention layer and one FFN layer, which
     projects ViT's outputs into MoE's inputs.
 
     Args:
@@ -129,6 +128,7 @@ class AriaProjector(nn.Module):
 
     Outputs:
         A tensor with the shape of (batch_size, query_number, output_dim)
+
     """
 
     def __init__(self, config: AriaConfig, prefix: str = "") -> None:
@@ -187,8 +187,7 @@ def forward(
 
 
 class AriaTextMoELayer(nn.Module):
-    """
-    Mixture of Experts (MoE) Layer for the AriaMoE model.
+    """Mixture of Experts (MoE) Layer for the AriaMoE model.
 
     This layer implements the MoE mechanism, which routes input tokens to
     different experts based on a routing algorithm, processes them through the
@@ -224,11 +223,11 @@ def __init__(
             intermediate_size=config.intermediate_size,
             quant_config=quant_config,
             prefix=f"{prefix}.experts",
+            is_fused_checkpoint_transposed=True,
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass of the MoE Layer.
+        """Forward pass of the MoE Layer.
 
         Args:
             hidden_states: Input tensor of shape
@@ -236,16 +235,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
         Returns:
             torch.Tensor: Output tensor after passing through the MoE layer.
-        """
 
+        """
         router_output = torch.nn.functional.linear(hidden_states, self.router_weight)
 
         return self.experts(hidden_states, router_output)
 
 
 class AriaTextDecoderLayer(LlamaDecoderLayer):
-    """
-    Custom Decoder Layer for the AriaMoE model which modifies the standard
+    """Custom Decoder Layer for the AriaMoE model which modifies the standard
     `LlamaDecoderLayer` by replacing the traditional MLP with a Mixture of
     Experts (MoE) Layer.
     """
@@ -262,8 +260,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
 
 
 class AriaTextModel(LlamaModel, SupportsQuant):
-    """
-    Custom LlamaModel for the AriaMoE model which modifies the standard
+    """Custom LlamaModel for the AriaMoE model which modifies the standard
     LlamaModel by replacing the `LlamaDecoderLayer` with `MoEDecoderLayer`.
     """
 
@@ -272,14 +269,17 @@ class AriaTextModel(LlamaModel, SupportsQuant):
         "gate_up_proj": ["gate_proj", "up_proj"],
     }
 
-    # Aria packs all experts into single (transposed) fc1/fc2 tensors, which is
-    # exactly the pre-fused checkpoint layout FusedMoE self-loads once fc1/fc2
-    # are renamed to the fused gate_up_proj/down_proj names.
+    # The fused expert loader expects names without the .weight suffix.
     hf_to_vllm_mapper = LlamaModel.hf_to_vllm_mapper | WeightsMapper(
+        orig_to_new_regex={
+            re.compile(r"experts\.fc1\.weight$"): "experts.gate_up_proj",
+            re.compile(r"experts\.fc2\.weight$"): "experts.down_proj",
+        },
+        # Quantization configs also use the expert module names.
         orig_to_new_substr={
             "experts.fc1": "experts.gate_up_proj",
             "experts.fc2": "experts.down_proj",
-        }
+        },
     )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -320,22 +320,18 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         vision_config = self.info.get_vision_config()
 
         max_image_size = vision_config.image_size
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
-        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
 
         return {
             "image": self._get_dummy_images(
                 width=max_image_size,
                 height=max_image_size,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -377,8 +373,7 @@ def _get_prompt_updates(
     dummy_inputs=AriaDummyInputsBuilder,
 )
 class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
-    """
-    Aria model for conditional generation tasks.
+    """Aria model for conditional generation tasks.
 
     This model combines a vision tower, a multi-modal projector, and a language
     model to perform tasks that involve both image and text inputs.
diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py
index 7d6e51f4d336..6308813ee29f 100644
--- a/vllm/model_executor/models/audioflamingo3.py
+++ b/vllm/model_executor/models/audioflamingo3.py
@@ -22,7 +22,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import BatchFeature, PretrainedConfig
+from transformers import BatchFeature, PreTrainedConfig
 from transformers.models.audioflamingo3 import (
     AudioFlamingo3Config,
     AudioFlamingo3Processor,
@@ -30,7 +30,7 @@
 from transformers.models.qwen2_audio import Qwen2AudioEncoder
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.models.module_mapping import MultiModelKeys
@@ -71,11 +71,10 @@
 
 
 class AudioFlamingo3FeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - num_chunks: Number of audio chunks (flattened)
-        - nmb: Number of mel bins
-        - num_audios: Number of original audio files
+    """Dimensions:
+    - num_chunks: Number of audio chunks (flattened)
+    - nmb: Number of mel bins
+    - num_audios: Number of original audio files
     """
 
     type: Literal["audio_features"]
@@ -96,12 +95,11 @@ class AudioFlamingo3FeatureInputs(TensorSchema):
 
 
 class AudioFlamingo3EmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size
-        - naf: Number of audio features
-        - hs: Hidden size (must match the hidden size of language model
-          backbone)
+    """Dimensions:
+    - bn: Batch size
+    - naf: Number of audio features
+    - hs: Hidden size (must match the hidden size of language model
+      backbone)
     """
 
     type: Literal["audio_embeds"] = "audio_embeds"
@@ -120,7 +118,7 @@ class AudioFlamingo3EmbeddingInputs(TensorSchema):
 class AudioFlamingo3Encoder(Qwen2AudioEncoder):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
     ):
         super().__init__(config)
         self.avg_pooler = nn.AvgPool1d(kernel_size=2, stride=2)
@@ -154,8 +152,7 @@ def forward(
         return hidden_states
 
     def _get_feat_extract_output_lengths(self, input_lengths: torch.Tensor):
-        """
-        Computes the output length of the convolutional layers and the output length
+        """Computes the output length of the convolutional layers and the output length
         of the audio encoder
         """
         input_lengths = (input_lengths - 1) // 2 + 1
@@ -164,7 +161,7 @@ def _get_feat_extract_output_lengths(self, input_lengths: torch.Tensor):
 
 
 class AudioFlamingo3MultiModalProjector(nn.Module):
-    def __init__(self, config: PretrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
         self.linear_1 = nn.Linear(
             config.audio_config.hidden_size,
@@ -200,6 +197,7 @@ def get_data_parser(self) -> MultiModalDataParser:
         return AudioFlamingo3MultiModalDataParser(
             target_sr=feature_extractor.sampling_rate,
             expected_hidden_size=self._get_expected_hidden_size(),
+            allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
         )
 
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
@@ -219,21 +217,18 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         hf_processor = self.info.get_hf_processor()
         feature_extractor = self.info.get_feature_extractor()
         sampling_rate = feature_extractor.sampling_rate
         audio_len = int(hf_processor.max_audio_len * sampling_rate)
-        num_audios = mm_counts.get("audio", 0)
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
 
         return {
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
@@ -243,6 +238,7 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     if chunk_counts is not None:
         return dict(
             audio_embeds=MultiModalFieldConfig.batched("audio"),
+            audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
             input_features=MultiModalFieldConfig.flat_from_sizes(
                 "audio", chunk_counts, dim=0
             ),
@@ -253,6 +249,7 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):
         )
     return dict(
         audio_embeds=MultiModalFieldConfig.batched("audio"),
+        audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
         input_features=MultiModalFieldConfig.batched("audio"),
         feature_attention_mask=MultiModalFieldConfig.batched("audio"),
         chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
@@ -360,11 +357,15 @@ def _count_audio_tokens_from_mask(
 
 
 class AudioFlamingo3MultiModalDataParser(MultiModalDataParser):
+    embedding_fields = {
+        "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
+    }
+
     def _parse_audio_data(
         self,
         data: dict[str, torch.Tensor] | ModalityData[Any],
     ) -> ModalityDataItems[Any, Any] | None:
-        if isinstance(data, dict):
+        if isinstance(data, dict) and "audio_num_tokens" not in data:
             return DictEmbeddingItems(
                 data,
                 modality="audio",
@@ -377,21 +378,9 @@ def _parse_audio_data(
 class AudioFlamingo3MultiModalProcessor(
     BaseMultiModalProcessor[AudioFlamingo3ProcessingInfo]
 ):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", None)
-        if audios is not None:
-            mm_data["audio"] = audios
-
-        return mm_data, hf_processor_mm_kwargs
-
     def _postprocess_hf_mm_data(
         self,
         mm_data: Mapping[str, object],
@@ -460,7 +449,11 @@ def _get_prompt_updates(
         chunk_counts = out_mm_data.get("chunk_counts")
 
         def get_replacement_audioflamingo3(item_idx: int):
-            if feature_attention_mask is not None:
+            if "audio_num_tokens" in out_mm_data:
+                counts = out_mm_data["audio_num_tokens"]
+                assert isinstance(counts, torch.Tensor)
+                num_features = int(counts[item_idx])
+            elif feature_attention_mask is not None:
                 num_features = _count_audio_tokens_from_mask(
                     feature_attention_mask,
                     chunk_counts,
@@ -503,9 +496,7 @@ class AudioFlamingo3ForConditionalGeneration(
     }
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model.",
             connector="multi_modal_projector.",
diff --git a/vllm/model_executor/models/bagel.py b/vllm/model_executor/models/bagel.py
index 9603dc043cc7..e9ebb0e5edc9 100644
--- a/vllm/model_executor/models/bagel.py
+++ b/vllm/model_executor/models/bagel.py
@@ -14,7 +14,7 @@
 import torch.nn as nn
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
@@ -58,12 +58,11 @@
 
 
 class BagelImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height of each image
-        - w: Width of each image
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height of each image
+    - w: Width of each image
     """
 
     type: Literal["pixel_values"]
@@ -175,12 +174,12 @@ def _get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos):
         return emb
 
     def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             position_ids: Flattened position IDs, shape (N,) where each ID
                          corresponds to a position in the flattened grid
         Returns:
             Position embeddings of shape (N, hidden_size)
+
         """
         # Ensure position_ids are on the same device as pos_embed
         position_ids = position_ids.to(self.pos_embed.device)
@@ -249,23 +248,20 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         hf_config = self.info.get_hf_config()
         vit_config = hf_config.vit_config
 
         # Use the configured image size
         image_size = vit_config.image_size
-        image_overrides = mm_options.get("image")
-        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
 
         return {
             "image": self._get_dummy_images(
                 width=image_size,
                 height=image_size,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
         }
 
@@ -322,8 +318,7 @@ def _get_mm_fields_config(
 class BagelForConditionalGeneration(
     nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP
 ):
-    """
-    BAGEL: A unified multimodal model for image understanding and generation.
+    """BAGEL: A unified multimodal model for image understanding and generation.
 
     For vLLM, we focus on the image understanding (vision-to-text) capabilities.
     The image generation part is not supported in vLLM.
@@ -507,6 +502,9 @@ def forward(
             positions: Flattened (concatenated) position ids corresponding to a batch.
             intermediate_tensors: Intermediate tensors from prior forward pass.
             inputs_embeds: Optional tensor of input embeddings.
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
+
         """
         if intermediate_tensors is not None:
             inputs_embeds = None
diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py
index 33749818b111..0a3b9c8e474b 100644
--- a/vllm/model_executor/models/bailing_moe.py
+++ b/vllm/model_executor/models/bailing_moe.py
@@ -30,7 +30,7 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -41,7 +41,10 @@
 )
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -71,7 +74,7 @@
 class BailingAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
@@ -180,7 +183,7 @@ class BailingMLP(nn.Module):
     def __init__(
         self,
         intermediate_size: int,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
         prefix: str = "",
@@ -214,7 +217,7 @@ class BailingMoE(nn.Module):
     def __init__(
         self,
         intermediate_size: int,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool | None = True,
         prefix: str = "",
@@ -243,11 +246,12 @@ def __init__(
         else:
             self.router_dtype = torch.bfloat16
 
-        self.gate = nn.Linear(
+        self.gate = GateLinear(
             self.hidden_size,
             self.num_experts,
-            bias=False,
-            dtype=self.router_dtype,
+            out_dtype=self.router_dtype,
+            params_dtype=self.router_dtype,
+            prefix=f"{prefix}.gate",
         )
 
         if getattr(config, "moe_router_enable_expert_bias", False):
@@ -312,7 +316,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = hidden_states.view(-1, hidden_size)
 
         # router_logits: (num_tokens, n_experts)
-        router_logits = self.gate(hidden_states.to(self.router_dtype))
+        router_logits, _ = self.gate(hidden_states)
         router_logits = router_logits.to(hidden_states.dtype)
 
         final_hidden_states = self.experts(
@@ -324,7 +328,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class BailingMoeBlock(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py
index 8731e91c5d6e..aea9559e6cfa 100644
--- a/vllm/model_executor/models/bailing_moe_linear.py
+++ b/vllm/model_executor/models/bailing_moe_linear.py
@@ -5,8 +5,7 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -19,6 +18,7 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -75,13 +75,11 @@ def is_linear_layer(layer_idx, layer_group_size):
 
 
 class BailingMoeV25MLAAttention(nn.Module):
-    """
-    MLA Attention for BailingMoeV2.5 full attention layers.
-    """
+    """MLA Attention for BailingMoeV2.5 full attention layers."""
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         layer_id: int = 0,
         prefix: str = "attention",
@@ -231,43 +229,12 @@ def forward(
         return self.mla_attn(positions, hidden_states)
 
 
-class BailingMoEGate(nn.Module):
-    def __init__(
-        self,
-        config: PretrainedConfig,
-        params_dtype: torch.dtype | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        if params_dtype is None:
-            params_dtype = torch.get_default_dtype()
-        self.params_dtype = params_dtype
-        self.weight = nn.Parameter(
-            torch.empty(
-                (config.num_experts, config.hidden_size),
-                dtype=self.params_dtype,
-            ),
-        )
-        if getattr(config, "moe_router_enable_expert_bias", False):
-            self.expert_bias = nn.Parameter(
-                torch.empty((config.num_experts,), dtype=torch.float32),
-            )
-        else:
-            self.expert_bias = None
-
-    def forward(self, hidden_states):
-        logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None).to(
-            hidden_states.dtype
-        )
-        return logits
-
-
 class BailingMoeV25(nn.Module):
     """Bailing MoE v2.5 - standalone implementation for linear attention model."""
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         layer_id: int = 0,
         prefix: str = "",
@@ -298,11 +265,19 @@ def __init__(
             self.router_dtype = torch.bfloat16
 
         # Gate for routing
-        self.gate = BailingMoEGate(
-            config=config,
+        self.gate = GateLinear(
+            self.hidden_size,
+            self.num_experts,
+            out_dtype=self.router_dtype,
             params_dtype=self.router_dtype,
             prefix=f"{prefix}.gate",
         )
+        if getattr(config, "moe_router_enable_expert_bias", False):
+            self.gate.expert_bias = nn.Parameter(
+                torch.empty((self.num_experts,), dtype=torch.float32),
+            )
+        else:
+            self.gate.expert_bias = None
         correction_bias = (
             self.gate.expert_bias if self.gate.expert_bias is not None else None
         )
@@ -359,7 +334,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = hidden_states.contiguous().view(-1, hidden_size)
 
         # router_logits: (num_tokens, n_experts)
-        router_logits = self.gate(hidden_states.to(self.router_dtype))
+        router_logits, _ = self.gate(hidden_states)
         router_logits = router_logits.to(hidden_states.dtype)
 
         final_hidden_states = self.experts(
@@ -374,7 +349,7 @@ class BailingMoeV25DecoderLayer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         vllm_config: VllmConfig,
         prefix: str = "layer",
         layer_id: int = 0,
@@ -594,7 +569,6 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load checkpoint weights with simplified mapping."""
-
         params_dict = dict(self.named_parameters(remove_duplicate=False))
         loaded_params: set[str] = set()
 
diff --git a/vllm/model_executor/models/bailing_moe_mtp.py b/vllm/model_executor/models/bailing_moe_mtp.py
index a403244a177e..32cc168102e7 100644
--- a/vllm/model_executor/models/bailing_moe_mtp.py
+++ b/vllm/model_executor/models/bailing_moe_mtp.py
@@ -6,7 +6,7 @@
 
 import torch
 import torch.nn as nn
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
@@ -36,7 +36,7 @@
 from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix
 
 
-def _get_draft_hf_config(vllm_config: VllmConfig) -> PretrainedConfig:
+def _get_draft_hf_config(vllm_config: VllmConfig) -> PreTrainedConfig:
     speculative_config = vllm_config.speculative_config
     if speculative_config is not None:
         draft_model_config = speculative_config.draft_model_config
@@ -48,7 +48,7 @@ def _get_draft_hf_config(vllm_config: VllmConfig) -> PretrainedConfig:
 class BailingMTPSharedHead(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         vllm_config: VllmConfig,
     ) -> None:
diff --git a/vllm/model_executor/models/bailing_moe_v3.py b/vllm/model_executor/models/bailing_moe_v3.py
index 710a37268ec2..9f5c078fac9c 100644
--- a/vllm/model_executor/models/bailing_moe_v3.py
+++ b/vllm/model_executor/models/bailing_moe_v3.py
@@ -16,9 +16,8 @@
 import regex as re
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from einops import rearrange
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
@@ -32,6 +31,7 @@
 from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -151,7 +151,7 @@ def _get_kda_state_shape_for_config(
     return shapes
 
 
-def _build_rope_parameters(config: PretrainedConfig) -> dict | None:
+def _build_rope_parameters(config: PreTrainedConfig) -> dict | None:
     rope_parameters = copy.deepcopy(getattr(config, "rope_parameters", None)) or {}
     if "rope_theta" not in rope_parameters and hasattr(config, "rope_theta"):
         rope_parameters["rope_theta"] = config.rope_theta
@@ -172,7 +172,7 @@ def _build_rope_parameters(config: PretrainedConfig) -> dict | None:
 
 
 def _build_mla_rotary_embedding(
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
     head_size: int,
 ) -> nn.Module:
     rope_parameters = _build_rope_parameters(config)
@@ -223,7 +223,7 @@ def _is_block_fp8_config(
 
 def _configure_ling_fp8_quant_config(
     quant_config: QuantizationConfig | None,
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
 ) -> None:
     if not _is_block_fp8_config(quant_config):
         return
@@ -372,7 +372,7 @@ def _pad_block_fp8_mlp_checkpoint_tensor(
 
 def _maybe_pad_block_fp8_shared_expert_checkpoint_tensor(
     quant_config: QuantizationConfig | None,
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
     name: str,
     loaded_weight: torch.Tensor,
 ) -> torch.Tensor:
@@ -410,7 +410,7 @@ class BailingMoeV3MLP(nn.Module):
     def __init__(
         self,
         intermediate_size: int,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
         prefix: str = "",
@@ -446,7 +446,7 @@ def forward(self, x):
 class BailingMoeV3MLAAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         layer_id: int = 0,
         prefix: str = "attention",
@@ -625,7 +625,7 @@ def get_state_shape(
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         layer_id: int = 0,
         prefix: str = "attention",
@@ -1118,32 +1118,10 @@ def _forward(
             core_attn_out[0, :num_actual_tokens] = out[0, :num_actual_tokens]
 
 
-class BailingMoeV3Gate(nn.Module):
-    def __init__(
-        self,
-        config: PretrainedConfig,
-        params_dtype: torch.dtype | None = None,
-    ) -> None:
-        super().__init__()
-        if params_dtype is None:
-            params_dtype = torch.float32
-        self.weight = nn.Parameter(
-            torch.empty((config.num_experts, config.hidden_size), dtype=params_dtype)
-        )
-        self.expert_bias = nn.Parameter(
-            torch.empty((config.num_experts,), dtype=torch.float32)
-        )
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return F.linear(hidden_states.to(self.weight.dtype), self.weight).to(
-            hidden_states.dtype
-        )
-
-
 class BailingMoeV3MoE(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         layer_id: int = 0,
         prefix: str = "",
@@ -1153,7 +1131,16 @@ def __init__(
         self.top_k = config.num_experts_per_tok
         self.hidden_size = config.hidden_size
         self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
-        self.gate = BailingMoeV3Gate(config)
+        self.gate = GateLinear(
+            config.hidden_size,
+            config.num_experts,
+            out_dtype=torch.float32,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.gate",
+        )
+        self.gate.expert_bias = nn.Parameter(
+            torch.empty((config.num_experts,), dtype=torch.float32)
+        )
         self.expert_swiglu_limit = _get_layer_swiglu_limit(
             getattr(config, "expert_swiglu_limit_list", None), layer_id
         )
@@ -1196,7 +1183,7 @@ def __init__(
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_size = hidden_states.shape
         hidden_states = hidden_states.contiguous().view(-1, hidden_size)
-        router_logits = self.gate(hidden_states.to(torch.float32))
+        router_logits, _ = self.gate(hidden_states)
         hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
         )
@@ -1206,7 +1193,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class BailingMoeV3DecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         layer_id: int = 0,
         prefix: str = "layer",
diff --git a/vllm/model_executor/models/bailing_moe_v3_mtp.py b/vllm/model_executor/models/bailing_moe_v3_mtp.py
index 2b4c3c8d7827..3d5b3556c237 100644
--- a/vllm/model_executor/models/bailing_moe_v3_mtp.py
+++ b/vllm/model_executor/models/bailing_moe_v3_mtp.py
@@ -6,7 +6,7 @@
 
 import torch
 import torch.nn as nn
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
@@ -47,7 +47,7 @@
 from vllm.sequence import IntermediateTensors
 
 
-def _get_draft_hf_config(vllm_config: VllmConfig) -> PretrainedConfig:
+def _get_draft_hf_config(vllm_config: VllmConfig) -> PreTrainedConfig:
     speculative_config = vllm_config.speculative_config
     if speculative_config is not None:
         draft_model_config = speculative_config.draft_model_config
@@ -59,7 +59,7 @@ def _get_draft_hf_config(vllm_config: VllmConfig) -> PretrainedConfig:
 class BailingMoeV3MTPSharedHead(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         vllm_config: VllmConfig,
     ) -> None:
diff --git a/vllm/model_executor/models/bailing_moe_v3_vl.py b/vllm/model_executor/models/bailing_moe_v3_vl.py
index c926c49ea8a1..a236d0722b4c 100644
--- a/vllm/model_executor/models/bailing_moe_v3_vl.py
+++ b/vllm/model_executor/models/bailing_moe_v3_vl.py
@@ -13,7 +13,7 @@
 from transformers import AutoProcessor, BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -103,7 +103,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         target_width, target_height = self.info.get_image_size_with_most_features()
         return {
@@ -119,7 +119,7 @@ def get_dummy_mm_data(
 class BailingMoeV3VLMultiModalProcessor(
     BaseMultiModalProcessor[BailingMoeV3VLProcessingInfo]
 ):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _apply_prompt_updates(
diff --git a/vllm/model_executor/models/bee.py b/vllm/model_executor/models/bee.py
index 4a57413b6e83..cc6c0f09abc4 100644
--- a/vllm/model_executor/models/bee.py
+++ b/vllm/model_executor/models/bee.py
@@ -8,7 +8,7 @@
 from transformers.activations import GELUActivation
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
@@ -90,21 +90,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
         }
 
diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py
index cdc7870f3ec5..85ef42675b43 100644
--- a/vllm/model_executor/models/bert.py
+++ b/vllm/model_executor/models/bert.py
@@ -458,6 +458,7 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
     Attributes:
         model: An instance of BertModel used for forward operations.
         _pooler: An instance of Pooler used for pooling operations.
+
     """
 
     is_pooling_model = True
@@ -575,8 +576,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class SPLADESparsePooler(Pooler):
-    """
-    SPLADE sparse pooling:
+    """SPLADE sparse pooling:
     logits = mlm_head(hidden_states)
             -> log1p(relu(logits))
             -> (max|sum over L)
@@ -664,8 +664,7 @@ def forward(
 
 @default_pooling_type(seq_pooling_type="CLS")
 class BertSpladeSparseEmbeddingModel(BertEmbeddingModel):
-    """
-    BertEmbeddingModel + SPLADE sparse embedding.
+    """BertEmbeddingModel + SPLADE sparse embedding.
     - Make logits by self.mlm_head
     - pooler: SPLADESparsePooler(mlm_head...)
     """
@@ -779,6 +778,7 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu
     Attributes:
         model: An instance of BertModel used for forward operations.
         _pooler: An instance of Pooler used for pooling operations.
+
     """
 
     is_pooling_model = True
diff --git a/vllm/model_executor/models/bert_with_rope.py b/vllm/model_executor/models/bert_with_rope.py
index ba6b20ec5f6a..cd35792b284b 100644
--- a/vllm/model_executor/models/bert_with_rope.py
+++ b/vllm/model_executor/models/bert_with_rope.py
@@ -4,7 +4,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -18,7 +18,11 @@
 from vllm.model_executor.layers.attention import (
     EncoderOnlyAttention,
 )
-from vllm.model_executor.layers.fused_moe import activation_without_mul, fused_topk
+from vllm.model_executor.layers.fused_moe import (
+    GateLinear,
+    activation_without_mul,
+    fused_topk,
+)
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
@@ -46,7 +50,7 @@
 
 
 class BertWithRopeEmbedding(nn.Module):
-    def __init__(self, config: PretrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
         if config.position_embedding_type not in ["rope", "rotary"]:
             raise ValueError(
@@ -254,9 +258,7 @@ def __init__(
             params_dtype = torch.get_default_dtype()
         self.params_dtype = params_dtype
 
-        self.router = ReplicatedLinear(
-            self.hidden_size, self.num_total_experts, bias=False
-        )
+        self.router = GateLinear(self.hidden_size, self.num_total_experts)
         self.w1 = nn.Parameter(
             torch.empty(
                 self.num_total_experts,
@@ -345,7 +347,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class BertWithRopeBlock(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         moe: bool = False,
diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py
index aefd17003d2f..a4e821ec6b12 100644
--- a/vllm/model_executor/models/blip.py
+++ b/vllm/model_executor/models/blip.py
@@ -82,7 +82,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
 
 
 class BlipAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
 
     def __init__(
         self,
@@ -140,8 +140,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
     ):
-        """Input shape: Batch x Time x Channel"""
-
+        """Input shape: Batch x Time x Channel."""
         qkv_states, _ = self.qkv(hidden_states)
         query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
         out = self.attn(query_states, key_states, value_states)
@@ -220,12 +219,12 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class BlipEncoder(nn.Module):
-    """
-    Transformer encoder consisting of `config.num_hidden_layers` self
+    """Transformer encoder consisting of `config.num_hidden_layers` self
     attention layers. Each layer is a [`BlipEncoderLayer`].
 
     Args:
         config: BlipConfig
+
     """
 
     def __init__(
diff --git a/vllm/model_executor/models/blip2.py b/vllm/model_executor/models/blip2.py
index 091001fc0173..cd019e41ec48 100644
--- a/vllm/model_executor/models/blip2.py
+++ b/vllm/model_executor/models/blip2.py
@@ -15,13 +15,14 @@
 )
 
 from vllm.config import CacheConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import MultiModalDataItems
@@ -49,12 +50,11 @@
 
 
 class Blip2ImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height of each image
-        - w: Width of each image
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height of each image
+    - w: Width of each image
     """
 
     type: Literal["pixel_values"]
@@ -62,11 +62,10 @@ class Blip2ImagePixelInputs(TensorSchema):
 
 
 class Blip2ImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - f: Image feature size
-        - h: Hidden size (must match the hidden size of language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - f: Image feature size
+    - h: Hidden size (must match the hidden size of language model backbone)
     """
 
     type: Literal["image_embeds"]
@@ -445,23 +444,19 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         hf_config = self.info.get_hf_config()
         vision_config = hf_config.vision_config
 
         max_image_size = vision_config.image_size
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
-        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
 
         return {
             "image": self._get_dummy_images(
                 width=max_image_size,
                 height=max_image_size,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -656,11 +651,15 @@ def forward(
         Args:
             input_ids: Flattened (concatenated) input_ids corresponding to a
                 batch.
+            positions: Flattened (concatenated) position ids corresponding to a
+                batch.
+            intermediate_tensors: Intermediate tensors from prior forward pass.
+            inputs_embeds: Optional tensor of input embeddings.
 
         Info:
             [`Blip2ImageInputs`][vllm.model_executor.models.blip2.Blip2ImageInputs]
-        """
 
+        """
         if intermediate_tensors is not None:
             inputs_embeds = None
 
@@ -687,28 +686,22 @@ def get_mm_mapping(self) -> MultiModelKeys:
             tower_model="vision_model",
         )
 
-    def get_num_mm_encoder_tokens(
+    def get_mm_lora_token_counts(
         self,
-        num_image_tokens: int,
-    ) -> int:
-        if num_image_tokens <= 0:
-            return 0
-        assert num_image_tokens % self.config.num_query_tokens == 0, (
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+        if num_mm_embeds <= 0:
+            return 0, 0
+        assert num_mm_embeds % self.config.num_query_tokens == 0, (
             "The number of image tokens must be a multiple of "
             "the number of query tokens."
         )
-        num_images = num_image_tokens / self.config.num_query_tokens
-        return num_images * self._vision_tokens_per_image
-
-    def get_num_mm_connector_tokens(
-        self,
-        num_vision_tokens: int,
-    ) -> int:
-        if num_vision_tokens <= 0:
-            return 0
-        assert num_vision_tokens % self._vision_tokens_per_image == 0, (
-            "The number of vision tokens must be a multiple of "
-            "the number of tokens per image."
+        num_images = num_mm_embeds // self.config.num_query_tokens
+        return (
+            num_images * self._vision_tokens_per_image,
+            num_images * self.config.num_query_tokens,
         )
-        num_images = num_vision_tokens / self._vision_tokens_per_image
-        return num_images * self.config.num_query_tokens
diff --git a/vllm/model_executor/models/clip.py b/vllm/model_executor/models/clip.py
index f0be4890ace4..7c5ba496ea8d 100644
--- a/vllm/model_executor/models/clip.py
+++ b/vllm/model_executor/models/clip.py
@@ -15,7 +15,7 @@
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import divide, get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict, MultiModalInput
 from vllm.model_executor.layers.activation import get_act_fn
@@ -67,12 +67,11 @@
 
 
 class CLIPImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height of each image
-        - w: Width of each image
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height of each image
+    - w: Width of each image
     """
 
     type: Literal["pixel_values"]
@@ -206,20 +205,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -417,8 +412,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
     ):
-        """Input shape: Batch x Time x Channel"""
-
+        """Input shape: Batch x Time x Channel."""
         qkv_states, _ = self.qkv_proj(hidden_states)
         query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
         out = self.attn(query_states, key_states, value_states)
@@ -506,12 +500,12 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class CLIPEncoder(nn.Module):
-    """
-    Transformer encoder consisting of `config.num_hidden_layers` self
+    """Transformer encoder consisting of `config.num_hidden_layers` self
     attention layers. Each layer is a [`CLIPEncoderLayer`].
 
     Args:
         config: CLIPConfig
+
     """
 
     def __init__(
diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py
index d657b78a5da1..7c17edc2a1db 100644
--- a/vllm/model_executor/models/cohere2_moe.py
+++ b/vllm/model_executor/models/cohere2_moe.py
@@ -16,11 +16,13 @@
 )
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -35,7 +37,7 @@
 from vllm.sequence import IntermediateTensors
 
 from .commandr import LayerNorm
-from .interfaces import SupportsPP, SupportsQuant
+from .interfaces import EagleModelMixin, SupportsEagle3, SupportsPP, SupportsQuant
 from .utils import (
     AutoWeightsLoader,
     WeightsMapper,
@@ -275,12 +277,10 @@ def __init__(
         else:
             self.custom_routing_function = None
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
             params_dtype=params_dtype,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
@@ -385,7 +385,7 @@ def forward(
 
 
 @support_torch_compile
-class Cohere2MoeModel(nn.Module):
+class Cohere2MoeModel(nn.Module, EagleModelMixin):
     """Transformer decoder for Cohere2Moe."""
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -439,7 +439,7 @@ def forward(
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
-    ) -> torch.Tensor | IntermediateTensors:
+    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
@@ -450,17 +450,28 @@ def forward(
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
+
+        aux_hidden_states = []
+        if self.start_layer in self.aux_hidden_state_layers:
+            aux_hidden_states.append(hidden_states)
+        for layer_idx, layer in enumerate(
+            islice(self.layers, self.start_layer, self.end_layer),
+            start=self.start_layer,
+        ):
             hidden_states, residual = layer(positions, hidden_states, residual)
+            if layer_idx + 1 in self.aux_hidden_state_layers:
+                aux_hidden_states.append(hidden_states)
         if not get_pp_group().is_last_rank:
             return IntermediateTensors(
                 {"hidden_states": hidden_states, "residual": residual}
             )
         hidden_states, _ = self.norm(hidden_states, residual)
+        if aux_hidden_states:
+            return hidden_states, aux_hidden_states
         return hidden_states
 
 
-class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
+class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant, SupportsEagle3):
     is_text_generation_model = True
 
     hf_to_vllm_mapper = WeightsMapper(
@@ -520,7 +531,7 @@ def forward(
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
-    ) -> torch.Tensor | IntermediateTensors:
+    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
         return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
 
     def compute_logits(
diff --git a/vllm/model_executor/models/cohere2_vision.py b/vllm/model_executor/models/cohere2_vision.py
index 42252afd67a2..bc4f78ed2cac 100644
--- a/vllm/model_executor/models/cohere2_vision.py
+++ b/vllm/model_executor/models/cohere2_vision.py
@@ -8,17 +8,17 @@
 
 import torch
 from torch import nn
-from transformers import BatchFeature, PretrainedConfig
+from transformers import BatchFeature, PreTrainedConfig
 from transformers.models.cohere2_vision import Cohere2VisionConfig
-from transformers.models.cohere2_vision.image_processing_cohere2_vision_fast import (  # noqa: E501
-    Cohere2VisionImageProcessorFast,
+from transformers.models.cohere2_vision.image_processing_cohere2_vision import (  # noqa: E501
+    Cohere2VisionImageProcessor,
 )
 from transformers.models.cohere2_vision.processing_cohere2_vision import (
     Cohere2VisionProcessor,
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import MulAndSilu
 from vllm.model_executor.layers.linear import (
@@ -61,14 +61,13 @@
 
 
 class Cohere2VisionImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - np: The total number of patches over each image over each prompt in
-              the batch
-        - c: Number of channels
-        - h: Height of each image patch
-        - w: Width of each image patch
-        - bn: Batch size * number of images
+    """Dimensions:
+    - np: The total number of patches over each image over each prompt in
+          the batch
+    - c: Number of channels
+    - h: Height of each image patch
+    - w: Width of each image patch
+    - bn: Batch size * number of images
     """
 
     type: Literal["pixel_values"]
@@ -134,6 +133,7 @@ def pixel_shuffle(self, image_features: torch.Tensor) -> torch.Tensor:
 
         Returns:
             Downsampled tensor with increased channel dimension
+
         """
         height = width = int(image_features.shape[1] ** 0.5)
         x = image_features.reshape(image_features.shape[0], width, height, -1)
@@ -175,11 +175,10 @@ def get_num_patches(
         processor: Cohere2VisionProcessor,
         mm_kwargs: Mapping[str, object],
     ) -> int:
-        """
-        Calculate the number of image patches for a given image.
+        """Calculate the number of image patches for a given image.
         Uses the HF processor to determine the actual number of patches.
         """
-        image_processor: Cohere2VisionImageProcessorFast = processor.image_processor
+        image_processor: Cohere2VisionImageProcessor = processor.image_processor
 
         return image_processor.get_number_of_image_patches(
             image_height,
@@ -203,19 +202,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         image_size = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=image_size.width,
                 height=image_size.height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -223,7 +219,7 @@ def get_dummy_mm_data(
 class Cohere2VisionMultiModalProcessor(
     BaseMultiModalProcessor[Cohere2VisionProcessingInfo]
 ):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -372,16 +368,18 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def _process_image_input(
-        self, image_input: Cohere2VisionImagePixelInputs, **kwargs
+        self, image_input: Cohere2VisionImagePixelInputs, **kwargs: object
     ) -> list[torch.Tensor]:
         """Process image pixels through vision tower and projector.
 
         Args:
             image_input: Validated image input containing pixel values and
                          patch counts
+            **kwargs: Unused; accepted for interface compatibility.
 
         Returns:
             List of flattened image embeddings, one per image
+
         """
         pixel_values = image_input["pixel_values"]
         num_patches = image_input["num_patches"]
@@ -417,7 +415,7 @@ def _parse_and_validate_image_input(
         )
 
     def _patch_quant_config(
-        self, config: PretrainedConfig, quant_config: QuantizationConfig | None
+        self, config: PreTrainedConfig, quant_config: QuantizationConfig | None
     ):
         # the awq models from OpenGVLab missing `modules_to_not_convert`
         # patch the quant_config to add `modules_to_not_convert` back
diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py
index d7f8f87fc6b1..298ccda9522b 100644
--- a/vllm/model_executor/models/cohere_asr.py
+++ b/vllm/model_executor/models/cohere_asr.py
@@ -8,11 +8,11 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from transformers import BatchFeature, PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt
@@ -50,7 +50,10 @@
     PromptUpdate,
     TimingContext,
 )
-from vllm.multimodal.processing.processor import MultiModalProcessingInfo
+from vllm.multimodal.processing.processor import (
+    HFMultiModalInputs,
+    MultiModalProcessingResult,
+)
 from vllm.renderers import TokenizeParams
 from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.transformers_utils.processors.cohere_asr import (
@@ -300,13 +303,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class FixedPositionalEncoding(nn.Module):
-    """
-    Fixed positional encoding (embedding layer) from sine and cosine functions
+    """Fixed positional encoding (embedding layer) from sine and cosine functions
     of different frequencies according to https://arxiv.org/abs/1706.03762
 
     Args:
         hidden_size: size of the embeddings in the model, also known as d_model
         max_sequence_length: maximum allowed length of the input sequence
+
     """
 
     def __init__(self, hidden_size: int, max_sequence_length: int = 512) -> None:
@@ -689,10 +692,12 @@ def forward(
 
 class PositionalEncoding(torch.nn.Module):
     """Fixed sinusoidal positional encoding.
+
     Args:
         d_model (int): embedding dim
         max_len (int): maximum input length
         xscale (bool): whether to scale the input by sqrt(d_model)
+
     """
 
     def __init__(
@@ -724,12 +729,14 @@ def forward(
         self, x: torch.Tensor, cache_len: int = 0
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Adds positional encoding.
+
         Args:
             x (torch.Tensor): Input. Its shape is (batch, time, feature_size)
             cache_len (int): the size of the cache which is used to shift positions
         Returns:
             x+pos_emb (torch.Tensor): Its shape is (batch, time, feature_size)
             pos_emb (torch.Tensor): Its shape is (1, time, feature_size)
+
         """
         input_len = x.size(1) + cache_len
         if self.xscale:
@@ -762,14 +769,15 @@ def forward(
         self, x: torch.Tensor, cache_len: int = 0
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute positional encoding.
+
         Args:
             x (torch.Tensor): Input. Its shape is (batch, time, feature_size)
             cache_len (int): the size of the cache which is used to shift positions
         Returns:
             x (torch.Tensor): Its shape is (batch, time, feature_size)
             pos_emb (torch.Tensor): Its shape is (1, time, feature_size)
-        """
 
+        """
         if self.xscale:
             x = x * self.xscale
 
@@ -783,15 +791,13 @@ def forward(
 
 
 class Swish(nn.SiLU):
-    """
-    Swish activation function introduced in 'https://arxiv.org/abs/1710.05941'
+    """Swish activation function introduced in 'https://arxiv.org/abs/1710.05941'
     Mathematically identical to SiLU. See note in nn.SiLU for references.
     """
 
 
 class ConformerFeedForward(nn.Module):
-    """
-    feed-forward module of Conformer model.
+    """feed-forward module of Conformer model.
     use_bias (bool): Apply bias to all Linear and Conv1d
         layers to improve activation flow and stabilize
         training of huge models.
@@ -819,8 +825,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class CausalConv1D(nn.Conv1d):
-    """
-    A causal version of nn.Conv1d where each step would
+    """A causal version of nn.Conv1d where each step would
     have limited access to locations on its right or left.
     All arguments are the same as nn.Conv1d except padding.
 
@@ -893,6 +898,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class ConformerConvolution(nn.Module):
     """The convolution module for the Conformer model.
+
     Args:
         d_model (int): hidden dimension
         kernel_size (int): kernel size for depthwise convolution
@@ -904,6 +910,7 @@ class ConformerConvolution(nn.Module):
         use_bias (bool): Use bias in all Linear and Conv1d
             layers to improve activation flow and stabilize
             training of huge models. Defaults to True
+
     """
 
     def __init__(
@@ -979,10 +986,12 @@ def forward(
 
 class CohereASRMultiHeadAttention(nn.Module):
     """Multi-Head Attention layer of Transformer.
+
     Args:
         n_head (int): number of heads
         n_feat (int): size of the features
         use_bias (bool): whether to remove bias in linear and conv layers
+
     """
 
     def __init__(
@@ -1010,6 +1019,7 @@ def forward_qkv(
         value: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Transforms query, key and value.
+
         Args:
             query (torch.Tensor): (batch, time1, size)
             key (torch.Tensor): (batch, time2, size)
@@ -1018,6 +1028,7 @@ def forward_qkv(
             q (torch.Tensor): (batch, head, time1, size)
             k (torch.Tensor): (batch, head, time2, size)
             v (torch.Tensor): (batch, head, time2, size)
+
         """
         n_batch = query.size(0)
         q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
@@ -1036,6 +1047,7 @@ def forward_attention(
         mask: torch.Tensor | None,
     ) -> torch.Tensor:
         """Compute attention context vector.
+
         Args:
             value (torch.Tensor): (batch, time2, size)
             scores(torch.Tensor): (batch, time1, time2)
@@ -1044,6 +1056,7 @@ def forward_attention(
             value (torch.Tensor): transformed `value`
                 (batch, time2, d_model) weighted by the
                 attention scores
+
         """
         n_batch = value.size(0)
         if mask is not None:
@@ -1071,16 +1084,19 @@ def forward(
         pos_emb: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Compute 'Scaled Dot Product Attention'.
+
         Args:
             query (torch.Tensor): (batch, time1, size)
             key (torch.Tensor): (batch, time2, size)
             value(torch.Tensor): (batch, time2, size)
             mask (torch.Tensor): (batch, time1, time2)
+            pos_emb (torch.Tensor): optional relative position embeddings
 
-        returns:
+        Returns:
             output (torch.Tensor): transformed `value`
                 (batch, time1, d_model) weighted by the
                 query dot key attention
+
         """
         q, k, v = self.forward_qkv(query, key, value)
 
@@ -1130,8 +1146,10 @@ def __init__(
 
     def rel_shift(self, x: torch.Tensor) -> torch.Tensor:
         """Compute relative positional encoding.
+
         Args:
             x (torch.Tensor): (batch, nheads, time, 2*time-1)
+
         """
         b, h, qlen, pos_len = x.size()  # (b, h, t1, t2)
         # need to add a column of zeros on the left side of
@@ -1151,17 +1169,19 @@ def forward(
         pos_emb: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
+
         Args:
             query (torch.Tensor): (batch, time1, size)
             key (torch.Tensor): (batch, time2, size)
             value(torch.Tensor): (batch, time2, size)
             mask (torch.Tensor): (batch, time1, time2)
-            pos_emb (torch.Tensor) : (batch, time1, size)
+            pos_emb (torch.Tensor): (batch, time1, size)
 
         Returns:
             output (torch.Tensor): transformed `value`
                 (batch, time1, d_model) weighted by the
                 query dot key attention
+
         """
         q, k, v = self.forward_qkv(query, key, value)
         q = q.transpose(1, 2)  # (batch, time1, head, d_k)
@@ -1186,10 +1206,18 @@ def forward(
         matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
         matrix_bd = self.rel_shift(matrix_bd)
 
-        # drops extra elements in the matrix_bd to match the matrix_ac's size
-        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
-        matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)]
-        scores = (matrix_ac + matrix_bd) / self.s_d_k  # (batch, head, time1, time2)
+        # drops extra elements in matrix_bd to match the key sequence length
+        matrix_bd = matrix_bd[:, :, :, : k.size(-2)]
+        batch, heads, query_len, key_len = matrix_bd.shape
+        scale = 1.0 / self.s_d_k
+        scores = torch.baddbmm(
+            matrix_bd.flatten(0, 1),
+            q_with_bias_u.flatten(0, 1),
+            k.transpose(-2, -1).flatten(0, 1),
+            beta=scale,
+            alpha=scale,
+        ).view(batch, heads, query_len, key_len)
+
         return self.forward_attention(v, scores, mask)
 
 
@@ -1212,6 +1240,7 @@ class ConformerLayer(torch.nn.Module):
             Conv1d layers from each ConformerLayer to
             improve activation flow and stabilize training
             of huge models. Defaults to True.
+
     """
 
     def __init__(
@@ -1279,14 +1308,14 @@ def forward(
         pos_emb: torch.Tensor | None = None,
         pad_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             x (torch.Tensor): input signals (B, T, d_model)
             att_mask (torch.Tensor): attention masks(B, T, T)
             pos_emb (torch.Tensor): (L, 1, d_model)
             pad_mask (torch.tensor): padding mask
         Returns:
             x (torch.Tensor): (B, T, d_model)
+
         """
         residual = x
         x = self.norm_feed_forward1(x)
@@ -1331,8 +1360,7 @@ def forward(
 
 
 class ConformerEncoder(nn.Module):
-    """
-    The encoder for ASR model of Conformer.
+    """The encoder for ASR model of Conformer.
     Based on this paper:
     'Conformer: Convolution-augmented Transformer for
     Speech Recognition' by Anmol Gulati et al.
@@ -1473,12 +1501,12 @@ def get_num_encoder_cross_attn_tokens(self, num_encoder_input_tokens: int) -> in
         return num_encoder_cross_attn_tokens
 
     def set_max_audio_length(self, max_audio_length: int) -> None:
-        """
-        Sets maximum input length.
+        """Sets maximum input length.
         Pre-calculates internal seq_range mask.
 
         Args:
             max_audio_length (int): New maximum sequence length.
+
         """
         device = next(self.parameters()).device
         dtype = next(self.parameters()).dtype
@@ -1833,7 +1861,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
 
 class CohereASRProcessingInfo(BaseProcessingInfo):
-    def get_hf_config(self) -> PretrainedConfig:
+    def get_hf_config(self) -> PreTrainedConfig:
         return self.ctx.get_hf_config()
 
     def get_default_tok_params(self) -> TokenizeParams:
@@ -1916,7 +1944,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
         mm_processor_kwargs=None,
     ) -> MultiModalDataDict:
         feature_extractor = self.info.get_feature_extractor()
@@ -1944,42 +1972,29 @@ def create_encoder_prompt(
     ) -> list[int]:
         return [0]
 
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
-
-        mm_data = dict(mm_data)
-        mm_data["audio"] = mm_data.pop("audios")
-
-        hf_processor_mm_kwargs = dict(
-            **hf_processor_mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+        return hf_inputs._replace(
+            hf_kwargs=dict(
+                hf_inputs.hf_kwargs,
+                sampling_rate=feature_extractor.sampling_rate,
+            )
         )
 
-        return mm_data, hf_processor_mm_kwargs
-
-    def _postprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-        processed_data: BatchFeature,
-    ) -> BatchFeature:
-        if "labels" in processed_data:
-            processed_data["input_ids"] = processed_data.pop("labels")
-
-        return processed_data
-
     def _cached_apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
+    ) -> MultiModalProcessingResult:
         # Dithering injects noise into the extracted features, so the
         # feature extractor is not a pure function of its input. Since the
         # processing cache assumes that processor outputs are invariant
diff --git a/vllm/model_executor/models/cohere_compass.py b/vllm/model_executor/models/cohere_compass.py
index 0256f747c90c..8593bedfc0d1 100644
--- a/vllm/model_executor/models/cohere_compass.py
+++ b/vllm/model_executor/models/cohere_compass.py
@@ -51,7 +51,7 @@
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_world_size,
@@ -338,8 +338,7 @@ def pos_embed_interpolate_native(
 
 
 class Cohere_VLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - np: Number of patches
         - ni: Number of images
         - cps: Number of channels * patch_size * patch_size
@@ -365,8 +364,7 @@ class Cohere_VLImagePixelInputs(TensorSchema):
 
 
 class Cohere_VLImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - nf: Number of image features
         - hs: Hidden size
         - ni: Number of images
@@ -895,6 +893,7 @@ def prepare_encoder_metadata(
                 instead of computing from cu_seqlens (needed for CUDA
                 graph capture to cover worst-case replay scenarios).
             device: Device to place tensors on. Defaults to self.device.
+
         """
         if device is None:
             device = self.device
@@ -1162,7 +1161,6 @@ def get_hf_config(self) -> CohereCompassConfig:
     def get_hf_processor(self, **kwargs: object) -> CohereCompassProcessor:
         return self.ctx.get_hf_processor(
             CohereCompassProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
@@ -1239,11 +1237,8 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        image_overrides = mm_options.get("image")
-
         target_image_width, target_image_height = (
             self.info.get_image_size_with_most_features()
         )
@@ -1252,8 +1247,8 @@ def get_dummy_mm_data(
             "image": self._get_dummy_images(
                 width=target_image_width,
                 height=target_image_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
         }
 
@@ -1744,8 +1739,7 @@ def _postprocess_image_embeds_evs(
         image_embeds_split: tuple[torch.Tensor, ...],
         image_input: Cohere_VLImageInputs,
     ) -> tuple[torch.Tensor, ...]:
-        """
-        Append mrope positions for each for images.
+        """Append mrope positions for each for images.
         This is necessary to recover correct mrope positions
 
         Args:
@@ -1757,6 +1751,7 @@ def _postprocess_image_embeds_evs(
             Tuple of image embeddings for each image item.
             Resulting embeddings will have extra 5 channels for
             computed mrope positions, consistent with video embeddings.
+
         """
         if self.is_multimodal_pruning_enabled:
             merge_size = self.visual.spatial_merge_size
@@ -1837,8 +1832,7 @@ def recompute_mrope_positions(
         mrope_positions: torch.LongTensor,
         num_computed_tokens: int,
     ) -> tuple[MultiModalEmbeddings, torch.Tensor, int]:
-        """
-        Update part of input mrope positions (starting with
+        """Update part of input mrope positions (starting with
         num_computed_tokens index). Original mrope_positions are computed
         for unpruned sequence and becomes incorrect once pruning occurs,
         so once we prune media tokens we should reflect this in the
@@ -1856,6 +1850,7 @@ def recompute_mrope_positions(
         Returns:
             Tuple of (multimodal_embeddings, mrope_positions,
                 mrope_position_delta).
+
         """
         return self._recompute_mrope_positions(
             input_ids=input_ids,
@@ -2046,8 +2041,8 @@ def forward(
                     `None` if no images are passed.
                 - image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in
                     LLM. `None` if no images are passed.
-        """
 
+        """
         if intermediate_tensors is not None:
             inputs_embeds = None
 
@@ -2083,9 +2078,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector=["visual.merger", "visual.deepstack_merger_list"],
diff --git a/vllm/model_executor/models/colbert.py b/vllm/model_executor/models/colbert.py
index 124cfb5fdbec..afdd86072100 100644
--- a/vllm/model_executor/models/colbert.py
+++ b/vllm/model_executor/models/colbert.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ColBERT late interaction model for retrieval and reranking.
+"""ColBERT late interaction model for retrieval and reranking.
 
 ColBERT uses per-token embeddings and late interaction (MaxSim) scoring
 instead of single-vector representations or cross-encoder concatenation.
@@ -75,6 +74,7 @@ def _init_colbert_components(
                 ``None``, will be inferred from weights during loading (or
                 auto-loaded from sentence-transformers Dense layers).
             head_dtype: Data type for the projection layer.
+
         """
         self.hidden_size = hidden_size
         self.colbert_dim = colbert_dim
@@ -149,6 +149,7 @@ def _load_colbert_weights(
         Returns:
             ``(remaining_weights, loaded_names)`` — the weights that were
             **not** consumed and the set of names that were loaded.
+
         """
         weights_list = list(weights)
         other_weights: list[tuple[str, torch.Tensor]] = []
diff --git a/vllm/model_executor/models/colmodernvbert.py b/vllm/model_executor/models/colmodernvbert.py
index 89e55d40e275..a14476f7747b 100644
--- a/vllm/model_executor/models/colmodernvbert.py
+++ b/vllm/model_executor/models/colmodernvbert.py
@@ -15,7 +15,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -135,18 +135,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         target_width, target_height = self.info.get_image_size_with_most_features()
-        image_overrides = mm_options.get("image")
 
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -157,15 +155,14 @@ class ColModernVBertMultiModalProcessor(
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
         if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+            return self._finalize_hf_mm_data(mm_data, hf_kwargs, passthrough_data)
 
         prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
 
@@ -193,8 +190,9 @@ def _apply_hf_processor_main(
             result.update(image_outputs)
 
         processed_data = result
-        processed_data.update(passthrough_data)
-        return processed_data
+        return self._finalize_hf_mm_data(
+            mm_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_mm_fields_config(
         self,
diff --git a/vllm/model_executor/models/colpali.py b/vllm/model_executor/models/colpali.py
index f31d63c00632..4add3b83a644 100644
--- a/vllm/model_executor/models/colpali.py
+++ b/vllm/model_executor/models/colpali.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ColPali late interaction model for multi-modal retrieval and reranking.
+"""ColPali late interaction model for multi-modal retrieval and reranking.
 
 ColPali extends PaliGemma with a ColBERT-style late interaction head,
 producing per-token embeddings for both text and image inputs. It uses
@@ -21,7 +20,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import BatchFeature, PaliGemmaProcessor
+from transformers import PaliGemmaProcessor
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
@@ -29,6 +28,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 
 from .interfaces import SupportsLateInteraction
 from .interfaces_base import default_pooling_type
@@ -60,31 +60,19 @@ def get_hf_processor(self, **kwargs: object) -> PaliGemmaProcessor:
 class ColPaliMultiModalProcessor(PaliGemmaMultiModalProcessor):
     """Multimodal processor for ColPali."""
 
-    def _apply_hf_processor_main(
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
 
         # The ColPali tokenizer_config.json ships with a small default
         # max_length (50) that truncates the 1024 image tokens inserted
         # by PaliGemmaProcessor, causing a token-count mismatch.
         # vLLM enforces its own max_model_len, so we disable HF
         # truncation to keep all image + text tokens intact.
-        processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            mm_data,
-            dict(**hf_processor_mm_kwargs, truncation=False),
-        )
-        processed_data.update(passthrough_data)
-        return processed_data
+        return hf_inputs._replace(hf_kwargs=dict(hf_inputs.hf_kwargs, truncation=False))
 
 
 @default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
diff --git a/vllm/model_executor/models/colqwen3.py b/vllm/model_executor/models/colqwen3.py
index a2012fffd2ad..0e9236300119 100644
--- a/vllm/model_executor/models/colqwen3.py
+++ b/vllm/model_executor/models/colqwen3.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ColQwen3 late interaction model for multi-modal retrieval and reranking.
+"""ColQwen3 late interaction model for multi-modal retrieval and reranking.
 
 ColQwen3 extends Qwen3-VL with a ColBERT-style late interaction head,
 producing per-token embeddings for both text and image inputs. It uses
@@ -62,7 +61,6 @@ def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor:
         # correctly for the Qwen3-VL backbone.
         return self.ctx.get_hf_processor(
             Qwen3VLProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
@@ -131,6 +129,7 @@ class ColQwen3Model(Qwen3VLForConditionalGeneration, SupportsLateInteraction):
 
     Attributes:
         custom_text_proj: Linear projection from hidden_size to embed_dim
+
     """
 
     # Mark this as a pooling model so vLLM routes to pooler path
diff --git a/vllm/model_executor/models/colqwen3_5.py b/vllm/model_executor/models/colqwen3_5.py
index dabb4885ee51..7eddbe983eaa 100644
--- a/vllm/model_executor/models/colqwen3_5.py
+++ b/vllm/model_executor/models/colqwen3_5.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ColQwen3.5 late interaction model for multi-modal retrieval and reranking.
+"""ColQwen3.5 late interaction model for multi-modal retrieval and reranking.
 
 ColQwen3.5 extends Qwen3.5 with a ColBERT-style late interaction head,
 producing per-token embeddings for both text and image inputs. It uses
@@ -59,7 +58,6 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor:
         return self.ctx.get_hf_processor(
             Qwen3VLProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
@@ -128,6 +126,7 @@ class ColQwen3_5Model(
     Attributes:
         custom_text_proj: Linear projection from hidden_size to embed_dim.
             This is passed to the pooler as pooler.head.projector.
+
     """
 
     # Mark this as a pooling model so vLLM routes to pooler path
diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py
index 3fe78f081ea8..cb17cfa32f5a 100644
--- a/vllm/model_executor/models/commandr.py
+++ b/vllm/model_executor/models/commandr.py
@@ -114,8 +114,7 @@ def forward(self, hidden_states, residuals=None):
 
 
 def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]:
-    """
-    Returns the normalization layer class and epsilon value to use.
+    """Returns the normalization layer class and epsilon value to use.
     If `config.rms_norm_eps` is present, use RMSNorm.
     Otherwise default to LayerNorm with `config.layer_norm_eps`.
     """
diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py
index 485d15354deb..3c87b061674e 100644
--- a/vllm/model_executor/models/config.py
+++ b/vllm/model_executor/models/config.py
@@ -6,7 +6,7 @@
 from vllm.utils.math_utils import round_up
 
 if TYPE_CHECKING:
-    from transformers import PretrainedConfig
+    from transformers import PreTrainedConfig
 
     from vllm.config import CacheConfig, ModelConfig, VllmConfig
     from vllm.config.cache import MambaDType
@@ -212,12 +212,15 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
         """Configure attention for heterogeneous head dimensions.
 
         Gemma4 uses different head dimensions for sliding window vs full attention
-        layers. The default FA3 on Hopper cannot handle head_dim > 256, which causes
-        mixed backend selection and numerical divergence.
-
-        When FA4 is available we force it for ALL layers, giving a uniform kernel path
-        and avoiding the mixed FA3+FA4 penalty. When FA4 is not available we fall back
-        to Triton.
+        layers. The default FA3 on Hopper cannot handle head_dim > 256.
+
+        On SM90 with FP8 KV cache, use FA3 for supported layers and let the generic
+        FlashAttention selector upgrade larger head dimensions to FA4.
+        The multimodal-prefix composite routes image masks to Triton and causal
+        requests to this per-layer FA3/FA4 selection. For other configurations,
+        force FA4 for all layers to avoid the mixed
+        FA3+FA4 penalty.
+        When FA4 is not available, fall back to Triton.
         """
         model_config = vllm_config.model_config
         arch_config = model_config.model_arch_config
@@ -230,6 +233,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
         if len(set(head_dims.values())) <= 1:
             return
 
+        from vllm.platforms import current_platform
         from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
         from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
@@ -241,12 +245,22 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
                 and vllm_config.attention_config.backend
                 in (None, AttentionBackendEnum.FLASH_ATTN)
             ):
-                vllm_config.attention_config.flash_attn_version = 4
-                logger.info(
-                    "Gemma4 model has heterogeneous head dimensions %s. Using FA4 for "
-                    "all layers to avoid mixed FA3/FA4 penalty.",
-                    head_dims,
-                )
+                use_per_layer_fa = current_platform.is_device_capability_family(
+                    90
+                ) and vllm_config.cache_config.cache_dtype.startswith("fp8")
+                if use_per_layer_fa:
+                    logger.info(
+                        "Gemma4 model has heterogeneous head dimensions %s. Using "
+                        "per-layer FA3/FA4 selection for FP8 KV cache on SM90.",
+                        head_dims,
+                    )
+                else:
+                    vllm_config.attention_config.flash_attn_version = 4
+                    logger.info(
+                        "Gemma4 model has heterogeneous head dimensions %s. Using FA4 "
+                        "for all layers to avoid mixed FA3/FA4 penalty.",
+                        head_dims,
+                    )
         elif vllm_config.attention_config.backend is None:
             vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
             logger.info(
@@ -280,11 +294,14 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
                 "FLASH_ATTN or TRITON_ATTN instead."
             )
         if attention_config.backend is None and not attention_config.use_non_causal:
+            # Only the sm_100 default backend order reads this flag (it puts
+            # FlashInfer first for causal models). Elsewhere Gemma4Config above
+            # keeps FlashInfer out: FA4 on every layer, else TRITON_ATTN.
             attention_config.use_non_causal = True
             logger.info(
-                "DiffusionGemma uses mixed causal/bidirectional attention "
-                "within a batch; setting use_non_causal=True to exclude "
-                "FlashInfer from auto-selection."
+                "DiffusionGemma mixes causal and bidirectional attention in "
+                "one batch; setting use_non_causal=True so the default backend "
+                "order prefers FlashAttention over FlashInfer."
             )
 
         # Auto-create DiffusionConfig from HF config if not provided.
@@ -330,7 +347,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         quant_config = getattr(model_config.hf_config, "quantization_config", None)
         if quant_config is not None and quant_config.get("quant_method") == "fp8":
             model_type = getattr(model_config.hf_config, "model_type", None)
-            if model_type == "deepseek_v4":
+            if model_type in ("deepseek_v4", "deepseek_v41"):
                 model_config.hf_config.quantization_config["quant_method"] = (
                     "deepseek_v4_fp8"
                 )
@@ -343,7 +360,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
             and hf_text_quant_config.get("quant_method") == "fp8"
         ):
             model_type = getattr(model_config.hf_text_config, "model_type", None)
-            if model_type == "deepseek_v4":
+            if model_type in ("deepseek_v4", "deepseek_v41"):
                 model_config.hf_text_config.quantization_config["quant_method"] = (
                     "deepseek_v4_fp8"
                 )
@@ -445,8 +462,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
 class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
     @classmethod
     def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
-        """
-        Perform early validation and setup for hybrid attention/mamba models.
+        """Perform early validation and setup for hybrid attention/mamba models.
 
         Block size alignment with mamba page sizes is handled later by
         Platform.update_block_size_for_backend(), which runs after model
@@ -454,6 +470,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
 
         Args:
             vllm_config: vLLM Config
+
         """
         # Enable FULL_AND_PIECEWISE by default
         MambaModelConfig.verify_and_update_config(vllm_config)
@@ -529,6 +546,7 @@ class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig):
     def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         config = model_config.hf_config
         config.num_labels = 1
+        config.get_text_config().num_labels = 1
         pooler_config = model_config.pooler_config
         assert pooler_config is not None
         if pooler_config.logit_mean is None:
@@ -609,12 +627,12 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
 class MambaModelConfig(VerifyAndUpdateConfig):
     @classmethod
     def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
-        """
-        Enable FULL_AND_PIECEWISE cuda graph mode by default (required
+        """Enable FULL_AND_PIECEWISE cuda graph mode by default (required
         to get good performance for mamba layers in V1).
 
         Args:
             vllm_config: vLLM Config
+
         """
         model_config = vllm_config.model_config
         cache_config = vllm_config.cache_config
@@ -663,7 +681,7 @@ class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
 
     @classmethod
     def update_mamba_ssm_cache_dtype(
-        cls, *, cache_config: "CacheConfig", hf_config: "PretrainedConfig"
+        cls, *, cache_config: "CacheConfig", hf_config: "PreTrainedConfig"
     ) -> None:
         """Update mamba_ssm_cache_dtype for NemotronH models when set to 'auto'
         (or not explicitly set), to the value specified in the HF config, or to
@@ -1003,6 +1021,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
     "ColQwen3_5": ColQwen3_5Config,
     "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,
     "DeepseekV4ForConditionalGeneration": DeepseekV4ForCausalLMConfig,
+    "DeepseekV41ForCausalLM": DeepseekV4ForCausalLMConfig,
     "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
     "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig,  # noqa: E501
     "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig,  # noqa: E501
diff --git a/vllm/model_executor/models/conformer_encoder.py b/vllm/model_executor/models/conformer_encoder.py
index 757dd1ac57c0..f4da532d5c57 100644
--- a/vllm/model_executor/models/conformer_encoder.py
+++ b/vllm/model_executor/models/conformer_encoder.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Conformer encoder components for FireRedASR2.
+"""Conformer encoder components for FireRedASR2.
 
 The audio encoder architecture is
 (Conv2dSubsampling → RelPositionalEncoding → N × RelPosEmbConformerBlock).
@@ -298,9 +297,7 @@ def forward(
 
 
 class ConformerEncoder(nn.Module):
-    """
-    Conformer encoder used by FireRedASR2.
-    """
+    """Conformer encoder used by FireRedASR2."""
 
     def __init__(
         self,
diff --git a/vllm/model_executor/models/cosmos3_edge.py b/vllm/model_executor/models/cosmos3_edge.py
index 8c5714248b69..ae65e25c901d 100644
--- a/vllm/model_executor/models/cosmos3_edge.py
+++ b/vllm/model_executor/models/cosmos3_edge.py
@@ -151,8 +151,7 @@ def patch_merging_by_param(
 
 
 class Cosmos3EdgePatchMerger(nn.Module):
-    """
-    Projector: LayerNorm -> Linear -> GELU -> Linear
+    """Projector: LayerNorm -> Linear -> GELU -> Linear.
 
     Reads config from projector_config (not vision_config).
     input_hidden_size * spatial_merge_size² -> merger_intermediate_size
@@ -490,8 +489,7 @@ class Cosmos3EdgeForConditionalGeneration(
     SupportsPP,
     SupportsMRoPE,
 ):
-    """
-    Cosmos3 Edge model with a SigLIP2 vision encoder.
+    """Cosmos3 Edge model with a SigLIP2 vision encoder.
 
     Architecture:
         - self.visual: SigLIP2 encoder + patch merger + projector
@@ -575,7 +573,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
 
         with self._mark_language_model(vllm_config):
             self.language_model = Cosmos3EdgeForCausalLM(
-                vllm_config=vllm_config,
+                vllm_config=vllm_config.with_hf_config(
+                    config.text_config, architectures=["NemotronHForCausalLM"]
+                ),
                 prefix=maybe_prefix(prefix, "language_model"),
             )
 
diff --git a/vllm/model_executor/models/dbrx.py b/vllm/model_executor/models/dbrx.py
index 0f5ffeff03e3..bfea9976981c 100644
--- a/vllm/model_executor/models/dbrx.py
+++ b/vllm/model_executor/models/dbrx.py
@@ -18,11 +18,11 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     RoutedExperts,
 )
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -62,12 +62,10 @@ def __init__(
         self.tp_size = get_tensor_model_parallel_world_size()
         self.num_total_experts = config.ffn_config.moe_num_experts
         self.d_model = config.d_model
-        self.layer = ReplicatedLinear(
+        self.layer = GateLinear(
             self.d_model,
             self.num_total_experts,
-            bias=False,
             params_dtype=params_dtype,
-            quant_config=None,
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
diff --git a/vllm/model_executor/models/deepencoder.py b/vllm/model_executor/models/deepencoder.py
index 189e8b55e510..fa972ca6157c 100644
--- a/vllm/model_executor/models/deepencoder.py
+++ b/vllm/model_executor/models/deepencoder.py
@@ -248,23 +248,23 @@ def __init__(
         global_attn_indexes: tuple[int, ...] = (),
         last_conv_output: int = 1024,
     ) -> None:
-        """
-        Args:
-            img_size (int): Input image size.
-            patch_size (int): Patch size.
-            in_chans (int): Number of input image channels.
-            embed_dim (int): Patch embedding dimension.
-            depth (int): Depth of ViT.
-            num_heads (int): Number of attention heads in each ViT block.
-            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-            qkv_bias (bool): If True, add a learnable bias to query, key, value.
-            norm_layer (nn.Module): Normalization layer.
-            act_layer (nn.Module): Activation layer.
-            use_abs_pos (bool): If True, use absolute positional embeddings.
-            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
-            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
-            window_size (int): Window size for window attention blocks.
-            global_attn_indexes (list): Indexes for blocks using global attention.
+        """Args:
+        img_size (int): Input image size.
+        patch_size (int): Patch size.
+        in_chans (int): Number of input image channels.
+        embed_dim (int): Patch embedding dimension.
+        depth (int): Depth of ViT.
+        num_heads (int): Number of attention heads in each ViT block.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value.
+        norm_layer (nn.Module): Normalization layer.
+        act_layer (nn.Module): Activation layer.
+        use_abs_pos (bool): If True, use absolute positional embeddings.
+        use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+        rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+        window_size (int): Window size for window attention blocks.
+        global_attn_indexes (list): Indexes for blocks using global attention.
+
         """  # noqa: E501
         super().__init__()
         self.img_size = img_size
@@ -378,20 +378,20 @@ def __init__(
         window_size: int = 0,
         input_size: tuple[int, int] | None = None,
     ) -> None:
-        """
-        Args:
-            dim (int): Number of input channels.
-            num_heads (int): Number of attention heads in each ViT block.
-            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
-            qkv_bias (bool): If True, add a learnable bias to query, key, value.
-            norm_layer (nn.Module): Normalization layer.
-            act_layer (nn.Module): Activation layer.
-            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
-            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
-            window_size (int): Window size for window attention blocks. If it equals 0, then
-                use global attention.
-            input_size (tuple(int, int) or None): Input resolution for calculating the relative
-                positional parameter size.
+        """Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads in each ViT block.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value.
+        norm_layer (nn.Module): Normalization layer.
+        act_layer (nn.Module): Activation layer.
+        use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+        rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+        window_size (int): Window size for window attention blocks. If it equals 0, then
+            use global attention.
+        input_size (tuple(int, int) or None): Input resolution for calculating the relative
+            positional parameter size.
+
         """  # noqa: E501
         super().__init__()
         self.norm1 = norm_layer(dim)
@@ -448,16 +448,16 @@ def __init__(
         rel_pos_zero_init: bool = True,
         input_size: tuple[int, int] | None = None,
     ) -> None:
-        """
-        Args:
-            dim (int): Number of input channels.
-            num_heads (int): Number of attention heads.
-            qkv_bias (bool):  If True, add a learnable bias to query, key, value.
-            use_triton_attention (bool): If True, fuse relative position bias in
-                a Triton attention kernel on CUDA-alike platforms.
-            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
-            input_size (tuple(int, int) or None): Input resolution for calculating the relative
-                positional parameter size.
+        """Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool):  If True, add a learnable bias to query, key, value.
+        use_triton_attention (bool): If True, fuse relative position bias in
+            a Triton attention kernel on CUDA-alike platforms.
+        rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+        input_size (tuple(int, int) or None): Input resolution for calculating the relative
+            positional parameter size.
+
         """  # noqa: E501
         super().__init__()
         self.num_heads = num_heads
@@ -540,8 +540,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 def window_partition(
     x: torch.Tensor, window_size: int
 ) -> tuple[torch.Tensor, tuple[int, int]]:
-    """
-    Partition into non-overlapping windows with padding if needed.
+    """Partition into non-overlapping windows with padding if needed.
+
     Args:
         x (tensor): input tokens with [B, H, W, C].
         window_size (int): window size.
@@ -549,6 +549,7 @@ def window_partition(
     Returns:
         windows: windows after partition with [B * num_windows, window_size, window_size, C].
         (Hp, Wp): padded height and width before partition
+
     """  # noqa: E501
     B, H, W, C = x.shape
 
@@ -571,8 +572,8 @@ def window_unpartition(
     pad_hw: tuple[int, int],
     hw: tuple[int, int],
 ) -> torch.Tensor:
-    """
-    Window unpartition into original sequences and removing padding.
+    """Window unpartition into original sequences and removing padding.
+
     Args:
         windows (tensor): input tokens with [B * num_windows, window_size, window_size, C].
         window_size (int): window size.
@@ -581,6 +582,7 @@ def window_unpartition(
 
     Returns:
         x: unpartitioned sequences with [B, H, W, C].
+
     """  # noqa: E501
     Hp, Wp = pad_hw
     H, W = hw
@@ -596,9 +598,9 @@ def window_unpartition(
 
 
 def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
-    """
-    Get relative positional embeddings according to the relative positions of
+    """Get relative positional embeddings according to the relative positions of
         query and key sizes.
+
     Args:
         q_size (int): size of query q.
         k_size (int): size of key k.
@@ -606,6 +608,7 @@ def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor
 
     Returns:
         Extracted positional embeddings according to relative positions.
+
     """
     max_rel_dist = int(2 * max(q_size, k_size) - 1)
     # Interpolate rel pos if needed.
@@ -641,8 +644,7 @@ def add_decomposed_rel_pos(
     q_size: tuple[int, int],
     k_size: tuple[int, int],
 ) -> torch.Tensor:
-    """
-    Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+    """Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
     https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py
     Args:
         q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
@@ -653,6 +655,7 @@ def add_decomposed_rel_pos(
 
     Returns:
         attn (Tensor): attention map with added relative positional embeddings.
+
     """  # noqa: E501
     q_h, q_w = q_size
     k_h, k_w = k_size
@@ -672,9 +675,7 @@ def add_decomposed_rel_pos(
 
 
 class PatchEmbed(nn.Module):
-    """
-    Image to Patch Embedding.
-    """
+    """Image to Patch Embedding."""
 
     def __init__(
         self,
@@ -684,13 +685,13 @@ def __init__(
         in_chans: int = 3,
         embed_dim: int = 768,
     ) -> None:
-        """
-        Args:
-            kernel_size (Tuple): kernel size of the projection layer.
-            stride (Tuple): stride of the projection layer.
-            padding (Tuple): padding size of the projection layer.
-            in_chans (int): Number of input image channels.
-            embed_dim (int): Patch embedding dimension.
+        """Args:
+        kernel_size (Tuple): kernel size of the projection layer.
+        stride (Tuple): stride of the projection layer.
+        padding (Tuple): padding size of the projection layer.
+        in_chans (int): Number of input image channels.
+        embed_dim (int): Patch embedding dimension.
+
         """
         super().__init__()
 
@@ -780,9 +781,7 @@ def forward(
         self, pixel_values: torch.Tensor, patch_embeds: torch.Tensor | None = None
     ) -> torch.Tensor:
         batch_size = pixel_values.shape[0]
-        if patch_embeds is not None:
-            patch_embeds = patch_embeds
-        else:
+        if patch_embeds is None:
             patch_embeds = self.patch_embedding(pixel_values)
         patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
 
diff --git a/vllm/model_executor/models/deepencoder2.py b/vllm/model_executor/models/deepencoder2.py
index 432b4eb56d1c..3b3860bc2cd4 100644
--- a/vllm/model_executor/models/deepencoder2.py
+++ b/vllm/model_executor/models/deepencoder2.py
@@ -23,8 +23,7 @@
 # --8<-- [start:qwen2_decoder]
 @PluggableLayer.register("qwen2_decoder")
 class CustomQwen2Decoder(PluggableLayer):
-    """
-    Qwen2 visual encoder
+    """Qwen2 visual encoder
     non-causal attention + causal attention
     token_type_ids ：0=non-causal, 1=causal
     """
@@ -75,7 +74,7 @@ def __init__(
         del self.model.embed_tokens
 
     def _create_custom_model(self, config):
-        """Qwen2Model"""
+        """Qwen2Model."""
 
         class CustomQwen2ModelInner(Qwen2Model):
             def forward(
@@ -194,11 +193,11 @@ def forward(
         attention_mask: torch.Tensor = None,
         **kwargs,
     ):
-        """
-        Args:
-            inputs_embeds: [batch_size, seq_len, hidden_dim]
-            token_type_ids: [batch_size, seq_len], 0=non-causal, 1=causal
-            attention_mask: [batch_size, seq_len], optional
+        """Args:
+        inputs_embeds: [batch_size, seq_len, hidden_dim]
+        token_type_ids: [batch_size, seq_len], 0=non-causal, 1=causal
+        attention_mask: [batch_size, seq_len], optional
+
         """
         return self.model(
             inputs_embeds=inputs_embeds,
@@ -209,8 +208,7 @@ def forward(
 
 
 class Qwen2Decoder2Encoder(nn.Module):
-    """
-    Decoder based on Multilingual BART
+    """Decoder based on Multilingual BART
     Set the initial weights and configuration with a pretrained multilingual BART model,
     and modify the detailed configurations as a Nougat decoder
     """
diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py
index 80c41501e6e9..9284dd59a09b 100644
--- a/vllm/model_executor/models/deepseek_eagle3.py
+++ b/vllm/model_executor/models/deepseek_eagle3.py
@@ -36,8 +36,7 @@
 
 
 class DeepseekV2Eagle3DecoderLayer(nn.Module):
-    """
-    Eagle3 decoder layer for Deepseek that:
+    """Eagle3 decoder layer for Deepseek that:
     1. Always uses MLP (not MoE)
     2. First layer accepts concatenated embeds + hidden_states
     """
@@ -305,10 +304,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.logits_processor = LogitsProcessor(
             self.config.draft_vocab_size, scale=logit_scale
         )
-        self.draft_id_to_target_id = nn.Parameter(
-            torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
-            requires_grad=False,
-        )
+        target_vocab_size = vllm_config.model_config.get_vocab_size()
+        if self.config.draft_vocab_size != target_vocab_size:
+            self.draft_id_to_target_id = nn.Parameter(
+                torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
+                requires_grad=False,
+            )
+        else:
+            self.draft_id_to_target_id = None
 
     def embed_input_ids(
         self,
diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py
index 8ac43e0bb2a1..5adc732f42cc 100644
--- a/vllm/model_executor/models/deepseek_mtp.py
+++ b/vllm/model_executor/models/deepseek_mtp.py
@@ -5,7 +5,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
@@ -39,9 +39,11 @@
     DeepseekV2MoE,
     _try_load_fp8_indexer_wk,
 )
+from .interfaces import SupportsPP
 from .utils import (
     get_pp_missing_layer_names,
     get_spec_layer_idx_from_weight_name,
+    make_empty_intermediate_tensors_factory,
     maybe_prefix,
 )
 
@@ -49,18 +51,23 @@
 class SharedHead(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         quant_config: QuantizationConfig | None = None,
+        *,
+        defer_lm_head: bool = False,
     ) -> None:
         super().__init__()
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.head = ParallelLMHead(
-            config.vocab_size,
-            config.hidden_size,
-            quant_config=quant_config,
-            prefix=maybe_prefix(prefix, "head"),
-        )
+        if defer_lm_head:
+            self.head = None
+        else:
+            self.head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "head"),
+            )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return self.norm(hidden_states)
@@ -230,7 +237,7 @@ def compute_logits(
 
 
 @support_torch_compile
-class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
+class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
@@ -240,6 +247,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         )
         # Set MoE hyperparameters
         self.set_moe_parameters()
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], self.config.hidden_size
+        )
 
     def set_moe_parameters(self):
         self.num_moe_layers = self.config.num_nextn_predict_layers
@@ -266,7 +276,7 @@ def set_moe_parameters(self):
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
 
-    def forward(
+    def forward(  # type: ignore[override]
         self,
         input_ids: torch.Tensor | None,
         positions: torch.Tensor,
@@ -328,6 +338,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 continue
             spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
             if spec_layer is None:
+                # A tied top-level embed_tokens has no spec layer to rewrite
+                # from; the draft needs its own copy under PP.
+                param = params_dict.get(name) if "embed_tokens" in name else None
+                if param is not None:
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add(name)
                 continue
             is_fusion_moe_shared_experts_layer = (
                 self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
@@ -515,8 +534,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loaded_params
 
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py
index 114498764c65..62a46fc49852 100644
--- a/vllm/model_executor/models/deepseek_ocr.py
+++ b/vllm/model_executor/models/deepseek_ocr.py
@@ -11,7 +11,7 @@
 from transformers import BatchFeature, CLIPVisionConfig
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
@@ -46,10 +46,10 @@
     PromptReplacement,
     PromptUpdate,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.sampling_params import SamplingParams
 from vllm.sequence import IntermediateTensors
 from vllm.tokenizers import cached_tokenizer_from_config
-from vllm.tokenizers.hf import HfTokenizer
 from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
 from vllm.transformers_utils.processors.deepseek_ocr import (
     BASE_SIZE,
@@ -79,13 +79,12 @@
 
 
 class DeepseekOCRImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - n: Number of images
-        - p: Number of patches
-        - base_size: Base size of the processor
-        - image_size: Image size of the processor
+    """Dimensions:
+    - b: Batch size
+    - n: Number of images
+    - p: Number of patches
+    - base_size: Base size of the processor
+    - image_size: Image size of the processor
     """
 
     type: Literal["pixel_values"]
@@ -269,17 +268,15 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         max_image_size = self.info.get_image_size_with_most_features()
 
         return {
             "image": self._get_dummy_images(
                 width=max_image_size.width,
                 height=max_image_size.height,
-                num_images=num_images,
+                num_images=mm_counts.get("image", 0),
             )
         }
 
@@ -287,34 +284,19 @@ def get_dummy_mm_data(
 class DeepseekOCRMultiModalProcessor(
     BaseMultiModalProcessor[DeepseekOCRProcessingInfo]
 ):
-    def _apply_hf_processor_main(
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
-
-        if mm_data:
-            processed_data = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(**hf_processor_mm_kwargs),
-                dict(prompt=prompt_text, **mm_data),
-                hf_processor_mm_kwargs,
-            )
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+        if "text" in hf_inputs.hf_data:
+            hf_inputs.hf_data["prompt"] = hf_inputs.hf_data.pop("text")
 
-        else:
-            tokenizer = self.info.get_tokenizer()
-            assert isinstance(tokenizer, HfTokenizer)
-            processed_data = tokenizer(
-                prompt_text, add_special_tokens=True, return_tensors="pt"
-            )
-
-        processed_data.update(passthrough_data)
-        return processed_data
+        return hf_inputs
 
     def _get_mm_fields_config(
         self,
@@ -639,9 +621,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return autoloaded_weights
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="projector",
@@ -676,8 +656,7 @@ def _get_num_input_output_tokens(
         self,
         image_spatial_crop: torch.Tensor | None = None,
     ) -> tuple[int, int, int, int]:
-        """
-        Return (num_input_tokens, num_output_tokens, global_output_token,
+        """Return (num_input_tokens, num_output_tokens, global_output_token,
         local_output_token) for a single image described by
         ``image_spatial_crop``.
         """
@@ -854,8 +833,7 @@ def _batched_encoder_forward_global_path(
         self,
         pixel_values: torch.Tensor,
     ) -> torch.Tensor:
-        """
-        Encode batched global images with newline tokens inserted.
+        """Encode batched global images with newline tokens inserted.
         Output shape: ``[B * 272, n_embed]``.
         """
         bsz = pixel_values.shape[0]
@@ -880,8 +858,7 @@ def _batched_encoder_forward_local_path(
         self,
         images_crop: torch.Tensor,
     ) -> torch.Tensor:
-        """
-        Encode local patches without newline insertion (newlines are added later
+        """Encode local patches without newline insertion (newlines are added later
         in ``postprocess_encoder_output`` via ``_assemble_patch_grid``).
         Output shape: ``[P * 100, n_embed]``.
         """
@@ -951,8 +928,7 @@ def postprocess_encoder_output(
         clone: bool = False,
         batch_mm_kwargs: dict[str, Any] | None = None,
     ) -> None:
-        """
-        Assemble per-image embeddings from global and local encoder outputs.
+        """Assemble per-image embeddings from global and local encoder outputs.
 
         ``output['global']`` contains global-image features with newlines already
         inserted (from CUDA graph replay or eager fallback):
diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py
index 0e907a344a87..938598a68bc6 100644
--- a/vllm/model_executor/models/deepseek_ocr2.py
+++ b/vllm/model_executor/models/deepseek_ocr2.py
@@ -11,7 +11,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
@@ -45,9 +45,9 @@
     PromptReplacement,
     PromptUpdate,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.sequence import IntermediateTensors
 from vllm.tokenizers import cached_tokenizer_from_config
-from vllm.tokenizers.hf import HfTokenizer
 from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
 from vllm.transformers_utils.processors.deepseek_ocr import (
     BASE_SIZE,
@@ -140,17 +140,15 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         max_image_size = self.info.get_image_size_with_most_features()
 
         return {
             "image": self._get_dummy_images(
                 width=max_image_size.width,
                 height=max_image_size.height,
-                num_images=num_images,
+                num_images=mm_counts.get("image", 0),
             )
         }
 
@@ -158,35 +156,19 @@ def get_dummy_mm_data(
 class DeepseekOCR2MultiModalProcessor(
     BaseMultiModalProcessor[DeepseekOCR2ProcessingInfo]
 ):
-    def _apply_hf_processor_main(
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
-
-        if mm_data:
-            processed_data = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(**hf_processor_mm_kwargs),
-                dict(prompt=prompt_text, **mm_data),
-                hf_processor_mm_kwargs,
-            )
-
-        else:
-            tokenizer = self.info.get_tokenizer()
-            assert isinstance(tokenizer, HfTokenizer)
-            processed_data = tokenizer(
-                prompt_text, add_special_tokens=True, return_tensors="pt"
-            )
-
-        processed_data.update(passthrough_data)
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+        if "text" in hf_inputs.hf_data:
+            hf_inputs.hf_data["prompt"] = hf_inputs.hf_data.pop("text")
 
-        return processed_data
+        return hf_inputs
 
     def _get_mm_fields_config(
         self,
@@ -456,9 +438,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return autoloaded_weights
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="projector",
diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py
index bc99509c7f45..d5728e0c3365 100644
--- a/vllm/model_executor/models/deepseek_v2.py
+++ b/vllm/model_executor/models/deepseek_v2.py
@@ -103,10 +103,14 @@
 from vllm.sequence import IntermediateTensors
 from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.backend import AttentionBackend, AttentionType
+from vllm.v1.attention.backends.mla.index_group import (
+    SparseMLAIndexGroupBuilder,
+    get_sparse_mla_index_group_max_rows,
+)
 from vllm.v1.attention.backends.mla.indexer import (
     DeepseekV32IndexerBackend,
 )
-from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
+from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, SparseCacheRole
 
 from .interfaces import (
     MixtureOfExperts,
@@ -536,9 +540,9 @@ def __init__(
         )
         if config.rope_parameters["rope_type"] != "default":
             config.rope_parameters["rope_type"] = (
-                "deepseek_yarn"
-                if config.rope_parameters.get("apply_yarn_scaling", True)
-                else "deepseek_llama_scaling"
+                "deepseek_llama_scaling"
+                if config.rope_parameters.get("attention_factor") == 1.0
+                else "deepseek_yarn"
             )
 
         self.rotary_emb = get_rope(
@@ -656,10 +660,16 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
             num_kv_heads=1,
             head_size=self.head_dim,
             dtype=self.dtype,
+            cache_role=SparseCacheRole.INDEXER,
         )  # Only has one vector instead of K + V
 
     def forward(self): ...
 
+    @property
+    def uses_shuffled_layout(self) -> bool:
+        """Whether this cache's reader expects the shuffled value layout."""
+        return False
+
     def get_attn_backend(self) -> type[AttentionBackend]:
         return DeepseekV32IndexerBackend
 
@@ -845,8 +855,7 @@ def forward(
 def _try_load_fp8_indexer_wk(
     name, tensor, buf, params_dict, loaded_params, pp_missing_layer_names
 ):
-    """
-    We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
+    """We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
     in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
     Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK
     weights and scale, and when both are available, dequantizes to BF16 and stores into
@@ -906,8 +915,7 @@ def _min_latency_fused_qkv_a_proj_impl(
     input_: torch.Tensor,
     weight: torch.Tensor,
 ) -> torch.Tensor:
-    """
-    Dynamically run min-latency gemm if num_tokens <= 16.
+    """Dynamically run min-latency gemm if num_tokens <= 16.
     This must be wrapped in a custom op because our torch.compile integration
     does not support runtime dispatching on num_tokens.
     """
@@ -988,8 +996,7 @@ def forward(
 
 
 class DeepseekV2MLAAttention(nn.Module):
-    """
-    Main reference: DeepseekV2 paper, and FlashInfer Implementation
+    """Main reference: DeepseekV2 paper, and FlashInfer Implementation
     (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
 
         For more info see MLACommonImpl in:
@@ -1012,6 +1019,7 @@ def __init__(
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
         topk_indices_buffer: torch.Tensor | None = None,
+        index_group_builder: SparseMLAIndexGroupBuilder | None = None,
         input_size: int | None = None,
         reduce_results: bool = True,
         non_causal_multi_token_decode: bool = False,
@@ -1104,9 +1112,9 @@ def __init__(
 
         if config.rope_parameters["rope_type"] != "default":
             config.rope_parameters["rope_type"] = (
-                "deepseek_yarn"
-                if config.rope_parameters.get("apply_yarn_scaling", True)
-                else "deepseek_llama_scaling"
+                "deepseek_llama_scaling"
+                if config.rope_parameters.get("attention_factor") == 1.0
+                else "deepseek_yarn"
             )
 
         self.rotary_emb = get_rope(
@@ -1199,6 +1207,7 @@ def __init__(
             indexer_rotary_emb=self.indexer_rope_emb,
             is_sparse=self.is_v32,
             topk_indices_buffer=topk_indices_buffer,
+            index_group_builder=index_group_builder,
         )
 
         self.mla_attn = MultiHeadLatentAttentionWrapper(
@@ -1242,6 +1251,7 @@ def __init__(
         prefix: str,
         config: DeepseekV2Config | None = None,
         topk_indices_buffer: torch.Tensor | None = None,
+        index_group_builder: SparseMLAIndexGroupBuilder | None = None,
     ) -> None:
         super().__init__()
 
@@ -1290,6 +1300,11 @@ def __init__(
             and parallel_config.pipeline_parallel_size == 1
             and is_moe_layer
         )
+        attn_kwargs = (
+            {"index_group_builder": index_group_builder}
+            if attn_cls is DeepseekV2MLAAttention
+            else {}
+        )
         self.self_attn = attn_cls(
             vllm_config=vllm_config,
             config=config,
@@ -1306,6 +1321,7 @@ def __init__(
             prefix=f"{prefix}.self_attn",
             topk_indices_buffer=topk_indices_buffer,
             reduce_results=not self.use_sequence_parallel_moe,
+            **attn_kwargs,
         )
 
         if is_moe_layer:
@@ -1429,6 +1445,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             )
         else:
             topk_indices_buffer = None
+        index_group_builder = (
+            SparseMLAIndexGroupBuilder(
+                topk_indices_buffer,
+                get_sparse_mla_index_group_max_rows(vllm_config),
+            )
+            if topk_indices_buffer is not None
+            else None
+        )
 
         if get_pp_group().is_first_rank:
             self.embed_tokens = VocabParallelEmbedding(
@@ -1445,6 +1469,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 vllm_config=vllm_config,
                 prefix=prefix,
                 topk_indices_buffer=topk_indices_buffer,
+                index_group_builder=index_group_builder,
             ),
             prefix=f"{prefix}.layers",
         )
diff --git a/vllm/model_executor/models/deepseek_vl2.py b/vllm/model_executor/models/deepseek_vl2.py
index e77dd3d13f59..c720f623351b 100644
--- a/vllm/model_executor/models/deepseek_vl2.py
+++ b/vllm/model_executor/models/deepseek_vl2.py
@@ -15,7 +15,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -35,7 +35,7 @@
 from vllm.multimodal.processing.processor import (
     BaseMultiModalProcessor,
     BaseProcessingInfo,
-    MultiModalProcessingInfo,
+    MultiModalProcessingResult,
     ProcessorInputs,
     PromptReplacement,
     PromptUpdate,
@@ -64,13 +64,12 @@
 
 
 class DeepseekVL2ImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * number of images * number of patches
-        - p: Number of patches
-        - c: Number of channels (3)
-        - h: Height of each image
-        - w: Width of each image
+    """Dimensions:
+    - bnp: Batch size * number of images * number of patches
+    - p: Number of patches
+    - c: Number of channels (3)
+    - h: Height of each image
+    - w: Width of each image
     """
 
     type: Literal["pixel_values"]
@@ -79,11 +78,10 @@ class DeepseekVL2ImagePixelInputs(TensorSchema):
 
 
 class DeepseekVL2VImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - f: Image feature size
-        - h: Hidden size (must match language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - f: Image feature size
+    - h: Hidden size (must match language model backbone)
     """
 
     type: Literal["image_embeds"]
@@ -214,20 +212,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         max_image_size = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=max_image_size.width,
                 height=max_image_size.height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -235,7 +229,7 @@ def get_dummy_mm_data(
 class DeepseekVL2MultiModalProcessor(
     BaseMultiModalProcessor[DeepseekVL2ProcessingInfo]
 ):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -309,7 +303,7 @@ def _cached_apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
+    ) -> MultiModalProcessingResult:
         # The processor logic is different for len(images) <= 2 vs > 2
         # Since the processing cache assumes that the processor output is
         # invariant of how many images are passed per prompt, we only
diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py
index 603d68410df6..84d52fb5ba3b 100644
--- a/vllm/model_executor/models/diffusion_gemma.py
+++ b/vllm/model_executor/models/diffusion_gemma.py
@@ -22,6 +22,7 @@
 
 import numpy as np
 import torch
+import torch._dynamo
 from torch import nn
 from torch.nn import functional as F
 from transformers import AutoModel
@@ -46,14 +47,18 @@
 )
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.transformers.utils import recursive_replace_linear
-from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    maybe_prefix,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.outputs import LogprobsTensors
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
-from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu
+from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores
@@ -107,13 +112,13 @@ class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo):
     """Processing info for DiffusionGemma.
 
     Overrides ``get_hf_config`` to accept ``DiffusionGemmaConfig``
-    (which inherits from ``PretrainedConfig``, not ``Gemma4Config``).
+    (which inherits from ``PreTrainedConfig``, not ``Gemma4Config``).
     Supports image and video modalities.
     """
 
     def get_hf_config(self):
         # DiffusionGemmaConfig doesn't inherit from Gemma4Config, so we
-        # accept any PretrainedConfig here.
+        # accept any PreTrainedConfig here.
         return self.ctx.get_hf_config()
 
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
@@ -160,14 +165,12 @@ class DiffusionGemmaForConditionalGeneration(
 
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
+            "model.decoder.self_conditioning.": "self_conditioning.",
             "model.decoder.": "model.",
             "model.encoder.language_model.": "model.",
             "model.encoder.vision_tower.": "vision_tower.",
             "model.encoder.embed_vision.": "embed_vision.",
         },
-        orig_to_new_substr={
-            ".experts.": ".moe.experts.",
-        },
     )
 
     packed_modules_mapping = {
@@ -195,6 +198,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         text_config.attention_k_eq_v = True
 
         # ---- Vision tower ----
+        # Gemma4's image path, borrowed below, reads this flag.
+        lora_config = vllm_config.lora_config
+        self._enable_mm_lora = bool(
+            lora_config is not None and lora_config.enable_tower_connector_lora
+        )
         vision_config = getattr(config, "vision_config", None)
         self.embed_vision: Gemma4MultimodalEmbedder | None
         if vision_config is not None:
@@ -336,114 +344,12 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
             logits = _softcap_logits(logits, self.final_logit_softcapping)
         return logits
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        """Load weights from checkpoint.
-
-        Checkpoint layout (HF DiffusionGemma):
-          model.encoder.vision_tower.*            → vision tower
-          model.encoder.embed_vision.*            → vision embedder
-          model.encoder.language_model.layers.*   → backbone
-          model.decoder.layers.*                  → backbone (tied)
-          model.decoder.embed_tokens.*            → embeddings
-          model.decoder.self_conditioning.*       → self-conditioning MLP
-          lm_head.*                               → LM head (tied)
-
-        We load encoder weights into our single ``Gemma4Model`` backbone,
-        skip duplicate decoder backbone weights, handle vision tower and
-        self-conditioning separately.
-        """
-
-        sc_params = dict(
-            (n, p)
-            for n, p in self.named_parameters()
-            if n.startswith("self_conditioning.")
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Some checkpoints carry a vestigial Gemma3n-style embedding table.
+        loader = AutoWeightsLoader(
+            self, ignore_unexpected_prefixes=["embed_vision.embedding."]
         )
-
-        # Collect vision tower + embedder parameters AND buffers for manual
-        # loading.  The HF vision tower registers std_bias / std_scale as
-        # buffers (not parameters) when config.standardize is True, so we
-        # must include named_buffers() to avoid "not found in model" warnings.
-        vision_params: dict[str, torch.Tensor] = {}
-        for n, p in self.named_parameters():
-            if n.startswith(("vision_tower.", "embed_vision.")):
-                vision_params[n] = p
-        for n, b in self.named_buffers():
-            if n.startswith(("vision_tower.", "embed_vision.")):
-                vision_params[n] = b
-
-        def _remap_weights():
-            # Use full weight names (including suffixes like .weight_scale,
-            # .weight_packed) for dedup instead of just the base layer name. Critical
-            # for quantized checkpoints where each weight has multiple tensors;
-            # tracking only base names skips scales as duplicates.
-            seen_weights: set[str] = set()
-            for name, weight in weights:
-                # Self-conditioning lives under model.decoder.self_conditioning.*
-                # in the checkpoint but at self_conditioning.* in our model.
-                if "self_conditioning" in name:
-                    sc_name = name.split("self_conditioning.", 1)[1]
-                    sc_name = "self_conditioning." + sc_name
-                    if sc_name in sc_params:
-                        sc_params[sc_name].data.copy_(weight)
-                    continue
-
-                # Vision tower: model.encoder.vision_tower.* → vision_tower.*
-                # In HF, the vision tower is a sibling of language_model
-                # under the encoder module.
-                if name.startswith("model.encoder.vision_tower."):
-                    vt_name = name[len("model.encoder.") :]
-                    if vt_name in vision_params:
-                        vision_params[vt_name].data.copy_(weight)
-                    else:
-                        logger.warning(
-                            "Vision tower weight %s (mapped to %s) not found in model",
-                            name,
-                            vt_name,
-                        )
-                    continue
-
-                # Vision embedder: model.encoder.embed_vision.* → embed_vision.*
-                if name.startswith("model.encoder.embed_vision."):
-                    ev_name = name[len("model.encoder.") :]
-                    if ev_name in vision_params:
-                        vision_params[ev_name].data.copy_(weight)
-                    else:
-                        logger.warning(
-                            "Embed vision weight %s (mapped to %s) not found in model",
-                            name,
-                            ev_name,
-                        )
-                    continue
-
-                # Skip vestigial embed_vision.embedding weights.
-                if "embed_vision.embedding." in name:
-                    continue
-
-                # Encoder backbone → model.*
-                if name.startswith("model.encoder.language_model."):
-                    name = name.replace("model.encoder.language_model.", "model.")
-                # Decoder backbone → model.* (skip exact duplicates)
-                elif name.startswith("model.decoder."):
-                    name = name.replace("model.decoder.", "model.")
-
-                # Skip only if we've seen the exact same weight name (including scales)
-                if name in seen_weights:
-                    continue
-                seen_weights.add(name)
-                yield name, weight
-
-        # Delegate to Gemma4ForCausalLM.load_weights for the backbone,
-        # which handles stacked params, MoE, k_eq_v, etc.
-        # Temporarily set self.config to text_config since Gemma4's
-        # load_weights expects it (e.g. tie_word_embeddings, layer_types).
-        from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
-
-        saved_config = self.config
-        self.config = self.model.config
-        try:
-            Gemma4ForCausalLM.load_weights(self, _remap_weights())
-        finally:
-            self.config = saved_config
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> str | None:
@@ -466,6 +372,34 @@ def _compute_num_rejected(
     return torch.where(is_denoise, query_lens, num_rejected)
 
 
+def _concat_logprob_stashes(
+    parts: list[LogprobsTensors], cu_num_generated_tokens: list[int]
+) -> LogprobsTensors:
+    """Join the logprobs stashed for the requests committing this step.
+
+    Each stash is as wide as the widest logprobs request in the batch at the
+    step that request converged, so stashes from different steps can differ
+    in width. Pad the narrow ones the way compute_topk_scores pads a mixed
+    batch: token id 0 at -inf, which the output processor never reports.
+    """
+    width = max(p.logprob_token_ids.shape[1] for p in parts)
+
+    def pad(t: torch.Tensor, value: float) -> torch.Tensor:
+        return F.pad(t, (0, width - t.shape[1]), value=value)
+
+    return LogprobsTensors(
+        logprob_token_ids=torch.cat([pad(p.logprob_token_ids, 0) for p in parts]),
+        logprobs=torch.cat([pad(p.logprobs, float("-inf")) for p in parts]),
+        selected_token_ranks=torch.cat([p.selected_token_ranks for p in parts]),
+        cu_num_generated_tokens=cu_num_generated_tokens,
+    )
+
+
+# Tiles specialize the graph on batch size 1, on a tile as wide as the canvas,
+# on compute_sc and on sizes that coincide with the state buffers. That set
+# is small but passes Dynamo's default of 8, after which every step would run
+# eager: about twice as slow for a self-conditioned step.
+@torch._dynamo.config.patch(recompile_limit=64)
 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
 def _compiled_sample_step(
     # Logits from the model [num_decode * CL, vocab]
@@ -486,6 +420,10 @@ def _compiled_sample_step(
     normalizer: torch.Tensor,
     history: torch.Tensor,  # [max_num_reqs, ST, CL]
     history_len_tensor: torch.Tensor,  # [max_num_reqs]
+    max_steps_tensor: torch.Tensor,  # [max_num_reqs] int32, per-slot step cap
+    pin_mask: torch.Tensor,  # [max_num_reqs, CL] bool, positions held at the seed
+    seed_canvas: torch.Tensor,  # [max_num_reqs, CL]
+    read_only: torch.Tensor,  # [max_num_reqs] bool, emit without a commit forward
     # Output tensors (modified in-place)
     sampled: torch.Tensor,  # [num_reqs, CL]
     num_sampled: torch.Tensor,  # [num_reqs]
@@ -507,6 +445,7 @@ def _compiled_sample_step(
     sc_vocab_end: int,
     tp_size: int,
     tp_group_name: str,
+    compute_sc: bool = True,
 ) -> torch.Tensor:
     """Compiled decode step: temperature → Gumbel sample → probs/confidence →
     accept/renoise → convergence, all as vectorized PyTorch ops.
@@ -573,8 +512,11 @@ def _compiled_sample_step(
         0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype
     )
 
-    # Compute denoise canvas (accept/renoise)
+    # Compute denoise canvas (accept/renoise). Pinned positions keep the seed.
     denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens)
+    denoise_canvas = torch.where(
+        pin_mask[decode_slots], seed_canvas[decode_slots], denoise_canvas
+    )
 
     # Canvas: commit → random reinit, denoise → accept/renoise result
     canvas[decode_slots] = torch.where(
@@ -599,16 +541,6 @@ def _compiled_sample_step(
     new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode))
     history_len_tensor[decode_slots] = new_hist_len
 
-    # Sampled output: commit → emit argmax_canvas, denoise → 0 (pre-zeroed)
-    sampled[decode_idx] = argmax_canvas[decode_slots].to(
-        sampled.dtype
-    ) * is_commit.unsqueeze(1).to(sampled.dtype)
-    # Commit only the real canvas length (== CL except for a canvas truncated
-    # near max_model_len); the padded tail positions are never emitted.
-    num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * valid_canvas_len.to(
-        num_sampled.dtype
-    )
-
     # ---- Phase 6: Stability + convergence ----
     ref = history[decode_slots, 0]
     mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32)
@@ -618,13 +550,17 @@ def _compiled_sample_step(
 
     step_after = step_tensor[decode_slots]
     converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | (
-        step_after >= max_denoising_steps
+        step_after >= max_steps_tensor[decode_slots]
     )
     # Commit done → denoise next (False); denoise converged → commit next (True)
     is_encoder_phase[decode_slots] = torch.where(
         is_commit, is_commit.new_zeros(num_decode), converged
     )
 
+    emit = is_commit | (is_denoise & converged & read_only[decode_slots])
+    sampled[decode_idx] = argmax_canvas[decode_slots].to(sampled.dtype) * emit[:, None]
+    num_sampled[decode_idx] = (emit * valid_canvas_len).to(num_sampled.dtype)
+
     # SC soft embedding: store ``probs @ embed_weight`` (the value the next step's
     # self-conditioning MLP consumes) only for slots that will denoise next — i.e.
     # this step denoised AND it isn't about to commit (is_encoder_phase now False).
@@ -632,24 +568,41 @@ def _compiled_sample_step(
     # sc_embeds directly. Storing the [.., hidden] soft embed instead of the full
     # [.., vocab] probs avoids a giant persistent buffer.
     sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None]
-    # Self-conditioning soft embed = probs @ embed_tokens.weight. Under tensor
-    # parallelism the embedding is vocab-sharded ([vocab/tp, hidden]) while
-    # probs spans the full vocab, so each rank multiplies its local vocab slice
-    # [sc_vocab_start, sc_vocab_end) and the partials are summed across ranks.
-    local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype)
-    soft_embeds = torch.matmul(
-        local_probs, embed_weight[: sc_vocab_end - sc_vocab_start]
-    )
-    if tp_size > 1:
-        soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name)
-    soft_embeds = soft_embeds * normalizer
-    sc_embeds[decode_slots] = soft_embeds * sc_keep
+    if compute_sc:
+        # Self-conditioning soft embed = probs @ embed_tokens.weight. Under
+        # tensor parallelism the embedding is vocab-sharded ([vocab/tp,
+        # hidden]) while probs spans the full vocab, so each rank multiplies
+        # its local vocab slice [sc_vocab_start, sc_vocab_end) and the
+        # partials are summed across ranks.
+        local_probs = probs[..., sc_vocab_start:sc_vocab_end].to(embed_weight.dtype)
+        soft_embeds = torch.matmul(
+            local_probs, embed_weight[: sc_vocab_end - sc_vocab_start]
+        )
+        if tp_size > 1:
+            soft_embeds = torch.ops.vllm.all_reduce(
+                soft_embeds, group_name=tp_group_name
+            )
+        soft_embeds = soft_embeds * normalizer
+        # A pinned position holds its seed token as input. Zero its soft embed
+        # too, or the model's own prediction there reaches the next step
+        # through self-conditioning. The buffer is fp32 while the embedding is
+        # in the model dtype: compiled code casts on the store but eager does
+        # not, and the step runs eager once torch.compile hits its recompile
+        # limit.
+        sc_pin = (~pin_mask[decode_slots]).unsqueeze(-1)
+        sc_embeds[decode_slots] = (soft_embeds * sc_keep * sc_pin).to(sc_embeds.dtype)
+    else:
+        # Every slot in this tile ends after this step, so the soft embed
+        # would never be read. The matmul is a full pass over the vocabulary
+        # matrix, so skip it.
+        sc_embeds[decode_slots] = 0
 
     # Overwrite canvas with argmax for newly converged denoise requests
     newly_converged = (converged & is_denoise).unsqueeze(1)
     canvas[decode_slots] = torch.where(
         newly_converged, argmax_canvas[decode_slots], canvas[decode_slots]
     )
+    is_encoder_phase[decode_slots] &= ~read_only[decode_slots]
 
     # ---- Phase 7: Copy canvas → draft_tokens for all slots ----
     draft_tokens[all_slots, :CL] = canvas[all_slots]
@@ -723,6 +676,32 @@ def __init__(
         # Per-slot confidence flag, set by the sampler each step.
         self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
 
+        # Per-slot step cap, lowered per request from extra_args.
+        self.max_steps = torch.full(
+            (max_num_reqs,), max_denoising_steps, dtype=torch.int32, device=device
+        )
+        # Seed canvases replace the random canvas after prefill. The host-side
+        # sets gate the seed and read-only work so plain generation skips it.
+        self.seed_canvas = torch.zeros(
+            max_num_reqs, canvas_length, dtype=torch.int64, device=device
+        )
+        self.has_seed = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
+        self.seeded_slots: set[int] = set()
+        # Pinned positions keep their seed value on every denoise step, so a
+        # multi-step read denoises the canvas the request seeded.
+        self.pin_mask = torch.zeros(
+            max_num_reqs, canvas_length, dtype=torch.bool, device=device
+        )
+        # Read-only slots emit on their converging step and skip the commit
+        # forward.
+        self.read_only = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
+        self.read_only_slots: set[int] = set()
+        # Slots capped at one denoise step never consume a soft embed.
+        self.single_step_slots: set[int] = set()
+        # Per-slot canvas width, at most canvas_length. The scheduler schedules
+        # this many draft tokens for the slot and the sampler pads the rest.
+        self.canvas_width_np = np.full(max_num_reqs, canvas_length, dtype=np.int32)
+
         # Per-slot self-conditioning soft embedding (probs @ embed_weight) from
         # the previous denoise step. Storing the [.., hidden] soft embed instead
         # of the full [.., vocab] distribution shrinks this buffer by
@@ -752,11 +731,57 @@ def add_request(self, slot_idx: int) -> None:
         self.step[slot_idx].fill_(0)
         self.accepted_canvas_history_len[slot_idx].fill_(0)
         self.self_conditioning_embeds[slot_idx] = 0
+        self.max_steps[slot_idx].fill_(self.max_denoising_steps)
+        self.has_seed[slot_idx].fill_(False)
+        self.seeded_slots.discard(slot_idx)
+        self.pin_mask[slot_idx].fill_(False)
+        self.read_only[slot_idx].fill_(False)
+        self.read_only_slots.discard(slot_idx)
+        self.single_step_slots.discard(slot_idx)
+        self.canvas_width_np[slot_idx] = self.canvas_length
 
     def remove_request(self, slot_idx: int) -> None:
+        # add_request resets the GPU flags before a slot is reused. The host
+        # sets gate work on every step, so they clear here.
         self.is_encoder_phase[slot_idx].fill_(False)
         self.accepted_canvas_history_len[slot_idx].fill_(0)
         self.self_conditioning_embeds[slot_idx] = 0
+        self.seeded_slots.discard(slot_idx)
+        self.read_only_slots.discard(slot_idx)
+        self.single_step_slots.discard(slot_idx)
+
+    def set_seed_canvas(self, slot_idx: int, ids: list[int]) -> None:
+        """``ids`` covers the slot's canvas width; positions past it are never
+        scheduled."""
+        self.seed_canvas[slot_idx, : len(ids)] = async_tensor_h2d(
+            ids, dtype=torch.int64, device=self.device
+        )
+        self.has_seed[slot_idx].fill_(True)
+        self.seeded_slots.add(slot_idx)
+
+    def set_pins(self, slot_idx: int, positions: list[int]) -> None:
+        """Hold ``positions`` of the slot's seed canvas through every denoise
+        step. Validated against the request's canvas width upstream."""
+        self.pin_mask[slot_idx].fill_(False)
+        self.pin_mask[
+            slot_idx, async_tensor_h2d(positions, dtype=torch.int64, device=self.device)
+        ] = True
+
+    def set_read_only(self, slot_idx: int) -> None:
+        self.read_only[slot_idx].fill_(True)
+        self.read_only_slots.add(slot_idx)
+
+    def apply_seed_canvases(
+        self, slots_np: np.ndarray, slots_gpu: torch.Tensor
+    ) -> None:
+        """Replace the canvas of every seeded slot among ``slots_gpu``."""
+        if self.seeded_slots.isdisjoint(slots_np.tolist()):
+            return
+        self.canvas[slots_gpu] = torch.where(
+            self.has_seed[slots_gpu, None],
+            self.seed_canvas[slots_gpu],
+            self.canvas[slots_gpu],
+        )
 
 
 class DiffusionGemmaModelState(ModelState):
@@ -1067,6 +1092,7 @@ def __init__(
         tp_group_name: str = "",
     ):
         self.sampling_states = sampler.sampling_states
+        self.logprob_token_ids_state = sampler.logprob_token_ids_state
         self.req_states = sampler.req_states
         self.logits_mode = sampler.logprobs_mode in ("raw_logits", "processed_logits")
         # Self-conditioning soft embed = probs @ embed_weight * normalizer,
@@ -1113,7 +1139,7 @@ def __init__(
         # on the subsequent commit step when num_sampled=CANVAS_LEN.
         self._pending_logprobs: dict[int, LogprobsTensors] = {}
 
-    def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> None:
+    def add_request(self, req_idx: int, sampling_params: Any) -> None:
         if use_penalty(sampling_params):
             logger.warning_once(
                 "DiffusionGemma does not support repetition/frequency/presence "
@@ -1123,9 +1149,38 @@ def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> No
         # that was aborted between its converging denoise and commit steps.
         self._pending_logprobs.pop(req_idx, None)
         self.sampling_states.add_request(req_idx, sampling_params)
+        self.logprob_token_ids_state.add_request(req_idx, sampling_params)
+        extra = getattr(sampling_params, "extra_args", None) or {}
+        states = self.diffusion_states
+        cap = extra.get("diffusion_max_steps")
+        if cap is not None:
+            cap = max(1, min(int(cap), states.max_denoising_steps))
+            states.max_steps[req_idx].fill_(cap)
+            if cap == 1:
+                states.single_step_slots.add(req_idx)
+        width = extra.get("diffusion_canvas_length")
+        if width:
+            states.canvas_width_np[req_idx] = max(
+                1, min(int(width), self.canvas_length)
+            )
+        width = int(states.canvas_width_np[req_idx])
+        seed = extra.get("diffusion_seed_canvas")
+        if seed is not None:
+            if len(seed) != width:
+                raise ValueError(
+                    f"diffusion_seed_canvas must hold exactly {width} ids, "
+                    f"got {len(seed)}"
+                )
+            states.set_seed_canvas(req_idx, seed)
+        pins = extra.get("diffusion_pinned")
+        if pins and seed is not None:
+            states.set_pins(req_idx, [int(p) for p in pins])
+        if extra.get("diffusion_read_only"):
+            states.set_read_only(req_idx)
 
     def apply_staged_writes(self) -> None:
         self.sampling_states.apply_staged_writes()
+        self.logprob_token_ids_state.apply_staged_writes()
 
     @property
     def penalties_state(self):
@@ -1159,10 +1214,11 @@ def _finish_prefills(
             return
         # Move the slot indices across once, up front: indexing a device
         # tensor with a numpy array copies them over synchronously each time.
-        ps_gpu = async_copy_to_gpu(
+        ps_gpu = async_tensor_h2d(
             ps.astype(np.int64), device=states.is_encoder_phase.device
         )
         states.init_canvas(ps_gpu)
+        states.apply_seed_canvases(ps, ps_gpu)
         self.req_states.draft_tokens[ps_gpu, : self.canvas_length] = states.canvas[
             ps_gpu
         ]
@@ -1243,7 +1299,6 @@ def __call__(
         # --- CPU/NumPy setup (outside compile): split decode vs prefill, init
         # canvas for any new prefills, and stage decode slot indices to GPU. ---
         states = self.diffusion_states
-        CL = self.canvas_length
         slots_np = input_batch.idx_mapping_np[:num_reqs]
         per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1])
 
@@ -1266,7 +1321,7 @@ def __call__(
         # was truncated near max_model_len, in which case the scheduler gave us
         # fewer than CL logits for that request.
         valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0]
-        valid_canvas_len = async_copy_to_gpu(
+        valid_canvas_len = async_tensor_h2d(
             valid_canvas_len_np.astype(np.int64), device=device
         )
 
@@ -1286,18 +1341,12 @@ def __call__(
             if top_k is not None or top_p is not None:
                 logits = apply_top_k_top_p(logits.float(), top_k, top_p)
 
-        # Pad any truncated canvas back to CL so the uniform-CL sampler math
-        # holds. Phantom (padded) positions are zeroed → uniform logits → high
-        # entropy (no premature convergence) and argmax 0 (stable); they are
-        # never committed (num_sampled == real length). masked_fill (not
-        # multiply) so -inf entries from top_k/top_p filtering above don't
-        # turn phantom rows into NaN.
-        if num_decode > 0 and valid_canvas_len_np.min() < CL:
-            ar = torch.arange(CL, device=device)
-            starts = valid_canvas_len.cumsum(0) - valid_canvas_len  # row offset per req
-            valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1)  # [num_decode, CL]
-            src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1)
-            logits = logits[src.reshape(-1)].masked_fill_(~valid.reshape(-1, 1), 0)
+        # Where each decode request's rows start in the flat logits. Tiles
+        # below gather and pad a request to the tile's width.
+        row_starts_np = np.concatenate(
+            ([0], np.cumsum(valid_canvas_len_np)[:-1])
+        ).astype(np.int64)
+        row_starts = async_tensor_h2d(row_starts_np, device=device)
 
         # Clear once: the tiled loop below only scatters its own decode slots,
         # so it must not re-clear earlier tiles' writes.
@@ -1313,109 +1362,166 @@ def __call__(
         is_committing = states.is_encoder_phase[decode_slots].clone()
 
         slots_np = input_batch.idx_mapping_np[:num_reqs]
-        is_decode_np = per_req_nlogits_np > 0
         max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)
-
-        # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in
-        # _compiled_sample_step keeps several live [group * CL, vocab] copies, so
-        # size each tile to a fraction of free memory to bound the transient at
-        # high concurrency. Tiling is bit-identical to a single pass.
-        group = max(num_decode, 1)
-        if num_decode > 0:
-            free, _ = current_platform.mem_get_info()
-            # ~10 transient fp32 copies of [group * CL, vocab] inside the step
+        # Requests may ask for specific token ids' logprobs instead of, or as
+        # well as, a top-k.
+        max_token_ids = self.logprob_token_ids_state.max_num_token_ids(slots_np)
+        want_logprobs = max_num_logprobs >= 0 or max_token_ids > 0
+        num_logprobs = max(max_num_logprobs, 0)
+
+        # Sample per tile. Decode requests are grouped by canvas width and each
+        # tile runs the compiled step at that width over [:, :W] views of the
+        # state, so a narrow read pays for its own rows rather than the served
+        # canvas. Widths ascend, so the last tile's canvas-to-draft copy (over
+        # all slots) is the widest. The fp32 pipeline keeps several live
+        # [tile * W, vocab] copies, so a tile is also bounded by free memory.
+        widths_np = states.canvas_width_np[decode_slots_np]
+        order = np.argsort(widths_np, kind="stable")
+        free = current_platform.mem_get_info()[0] if num_decode > 0 else 0
+        run_start = 0
+        while run_start < num_decode:
+            W = int(widths_np[order[run_start]])
+            run_end = run_start
+            while run_end < num_decode and widths_np[order[run_end]] == W:
+                run_end += 1
+            # ~10 transient fp32 copies of [tile * W, vocab] inside the step
             # (eager peaks at ~8; pad for allocator overhead and small tensors).
-            bytes_per_req = CL * self.vocab_size * 4 * 10
-            budget = int(free * 0.5) // max(bytes_per_req, 1)
-            group = max(1, min(num_decode, budget))
-
-        for start_req in range(0, num_decode, group):
-            end_req = min(start_req + group, num_decode)
-            tile = slice(start_req, end_req)
-            tile_slots = decode_slots[tile]
-
-            scaled = _compiled_sample_step(
-                logits[start_req * CL : end_req * CL],
-                tile_slots,
-                decode_idx[tile],
-                all_slots,
-                valid_canvas_len[tile],
-                # State
-                states.canvas,
-                states.argmax_canvas,
-                states.step,
-                states.is_encoder_phase,
-                states.confident,
-                states.self_conditioning_embeds,
-                self.embed_weight,
-                self.normalizer,
-                states.accepted_canvas_history,
-                states.accepted_canvas_history_len,
-                # Output
-                sampled,
-                num_sampled,
-                self.req_states.draft_tokens,
-                # Config
-                max_denoising_steps=float(states.max_denoising_steps),
-                t_min=self.t_min,
-                t_max=self.t_max,
-                confidence_threshold=self.confidence_threshold,
-                vocab_size=self.vocab_size,
-                CL=CL,
-                ST=states.stability_threshold,
-                entropy_bound=self.entropy_bound,
-                sc_vocab_start=self.sc_vocab_start,
-                sc_vocab_end=self.sc_vocab_end,
-                tp_size=self.tp_size,
-                tp_group_name=self.tp_group_name,
-            )
+            budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 10, 1))
+            for t0 in range(run_start, run_end, budget):
+                sel_np = order[t0 : min(t0 + budget, run_end)]
+                n = len(sel_np)
+                sel = async_tensor_h2d(sel_np.astype(np.int64), device=device)
+                tile_slots = decode_slots[sel]
+                tile_valid = valid_canvas_len[sel]
+                tile_valid_np = valid_canvas_len_np[sel_np]
+                contiguous = bool((sel_np == np.arange(sel_np[0], sel_np[0] + n)).all())
+                if contiguous and tile_valid_np.min() == W:
+                    r0 = int(row_starts_np[sel_np[0]])
+                    tile_logits = logits[r0 : r0 + n * W]
+                else:
+                    # Pad each request to W. Phantom positions are zeroed:
+                    # uniform logits, high entropy, argmax 0, never committed.
+                    # masked_fill, not multiply, so -inf from top_k/top_p
+                    # filtering above cannot turn a phantom row into NaN.
+                    ar = torch.arange(W, device=device)
+                    src = (row_starts[sel].unsqueeze(1) + ar.unsqueeze(0)).clamp_max(
+                        logits.shape[0] - 1
+                    )
+                    valid = ar.unsqueeze(0) < tile_valid.unsqueeze(1)
+                    tile_logits = logits[src.reshape(-1)].masked_fill_(
+                        ~valid.reshape(-1, 1), 0
+                    )
+                compute_sc = (
+                    not states.single_step_slots
+                    or not states.single_step_slots.issuperset(
+                        decode_slots_np[sel_np].tolist()
+                    )
+                )
 
-            # Logprobs for denoise steps that just converged (is_encoder_phase
-            # flipped False→True), stashed per tile so `scaled` is freed each tile.
-            if max_num_logprobs >= 0:
-                converged_mask = states.is_encoder_phase[tile_slots]
-                just_converged = converged_mask & ~is_committing[tile]
-                if just_converged.any():
-                    flat_logits = scaled.reshape(-1, scaled.shape[-1])
-                    argmax_tokens = scaled.argmax(dim=-1)
-                    for local_idx in just_converged.nonzero(as_tuple=True)[0]:
-                        li = local_idx.item()
-                        slot = tile_slots[local_idx]
-                        # Stash only the real canvas positions (== CL unless this
-                        # canvas was truncated near max_model_len); padded tail
-                        # positions are never emitted.
-                        k_i = int(valid_canvas_len_np[start_req + li])
-                        pos = li * CL
-                        self._pending_logprobs[slot.item()] = compute_topk_scores(
-                            flat_logits[pos : pos + k_i],
-                            max_num_logprobs,
-                            argmax_tokens[local_idx][:k_i],
-                            logits_mode=self.logits_mode,
-                        )
-
-        # Commit steps: is_committing was True at entry. Reassemble previously
-        # stashed logprobs and attach to SamplerOutput.
+                scaled = _compiled_sample_step(
+                    tile_logits,
+                    tile_slots,
+                    decode_idx[sel],
+                    all_slots,
+                    tile_valid,
+                    # State, viewed at this tile's width
+                    states.canvas[:, :W],
+                    states.argmax_canvas[:, :W],
+                    states.step,
+                    states.is_encoder_phase,
+                    states.confident,
+                    states.self_conditioning_embeds[:, :W],
+                    self.embed_weight,
+                    self.normalizer,
+                    states.accepted_canvas_history[:, :, :W],
+                    states.accepted_canvas_history_len,
+                    states.max_steps,
+                    states.pin_mask[:, :W],
+                    states.seed_canvas[:, :W],
+                    states.read_only,
+                    # Output
+                    sampled[:, :W],
+                    num_sampled,
+                    self.req_states.draft_tokens,
+                    # Config
+                    max_denoising_steps=float(states.max_denoising_steps),
+                    t_min=self.t_min,
+                    t_max=self.t_max,
+                    confidence_threshold=self.confidence_threshold,
+                    vocab_size=self.vocab_size,
+                    CL=W,
+                    ST=states.stability_threshold,
+                    entropy_bound=self.entropy_bound,
+                    sc_vocab_start=self.sc_vocab_start,
+                    sc_vocab_end=self.sc_vocab_end,
+                    tp_size=self.tp_size,
+                    tp_group_name=self.tp_group_name,
+                    compute_sc=compute_sc,
+                )
+
+                # Stash newly converged logprobs, including reads that emit now.
+                if want_logprobs:
+                    converged_mask = states.is_encoder_phase[tile_slots] | (
+                        num_sampled[decode_idx[sel]] > 0
+                    )
+                    just_converged = converged_mask & ~is_committing[sel]
+                    if just_converged.any():
+                        flat_logits = scaled.reshape(-1, scaled.shape[-1])
+                        argmax_tokens = scaled.argmax(dim=-1)
+                        raw_flat: torch.Tensor | None = None
+                        for local_idx in just_converged.nonzero(as_tuple=True)[0]:
+                            li = local_idx.item()
+                            slot = tile_slots[local_idx].item()
+                            # Stash only the real canvas positions; padded tail
+                            # positions are never emitted.
+                            k_i = int(tile_valid_np[li])
+                            pos = li * W
+                            src = flat_logits
+                            if slot in states.read_only_slots:
+                                # Read-only slots report logprobs at temperature 1.
+                                # The schedule-tempered logits share the argmax.
+                                if raw_flat is None:
+                                    raw_flat = tile_logits.float()
+                                src = raw_flat
+                            per_req_ids = max_token_ids > 0
+                            self._pending_logprobs[slot] = compute_topk_scores(
+                                src[pos : pos + k_i],
+                                num_logprobs,
+                                argmax_tokens[local_idx][:k_i],
+                                logprob_token_ids_state=(
+                                    self.logprob_token_ids_state
+                                    if per_req_ids
+                                    else None
+                                ),
+                                # every row of this stash belongs to one slot
+                                expanded_idx_mapping=(
+                                    torch.full(
+                                        (k_i,), slot, dtype=torch.int32, device=device
+                                    )
+                                    if per_req_ids
+                                    else None
+                                ),
+                                max_per_req_token_ids=max_token_ids,
+                                logits_mode=self.logits_mode,
+                            )
+            run_start = run_end
+
+        # Only emitting requests consume their stashed logprobs.
         logprobs_tensors = None
-        if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs:
-            parts_ids, parts_lp, parts_ranks = [], [], []
+        if want_logprobs and self._pending_logprobs:
+            emitting_slots = set(slots_np[num_sampled.cpu().numpy() > 0].tolist())
+            parts: list[LogprobsTensors] = []
             cu_gen: list[int] = []
             flat_offset = 0
             for i in range(num_reqs):
                 cu_gen.append(flat_offset)
                 slot = int(slots_np[i])
-                if is_decode_np[i] and slot in self._pending_logprobs:
+                if slot in emitting_slots and slot in self._pending_logprobs:
                     lp = self._pending_logprobs.pop(slot)
-                    parts_ids.append(lp.logprob_token_ids)
-                    parts_lp.append(lp.logprobs)
-                    parts_ranks.append(lp.selected_token_ranks)
+                    parts.append(lp)
                     flat_offset += lp.logprobs.shape[0]
-            if parts_ids:
-                logprobs_tensors = LogprobsTensors(
-                    logprob_token_ids=torch.cat(parts_ids),
-                    logprobs=torch.cat(parts_lp),
-                    selected_token_ranks=torch.cat(parts_ranks),
-                    cu_num_generated_tokens=cu_gen,
-                )
+            if parts:
+                logprobs_tensors = _concat_logprob_stashes(parts, cu_gen)
 
         return self._build_output(
             input_batch,
diff --git a/vllm/model_executor/models/dots_ocr.py b/vllm/model_executor/models/dots_ocr.py
index e91d6ce94c9e..adcc560ba5ea 100644
--- a/vllm/model_executor/models/dots_ocr.py
+++ b/vllm/model_executor/models/dots_ocr.py
@@ -9,7 +9,7 @@
 from transformers.models.qwen2_vl import Qwen2VLProcessor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import utils as dist_utils
 from vllm.distributed.parallel_state import (
     get_tensor_model_parallel_rank,
@@ -54,6 +54,7 @@
 )
 from vllm.model_executor.models.vision import get_vit_attn_backend
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalKwargsItem
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.dotsocr import DotsOCRConfig, DotsVisionConfig
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -65,12 +66,11 @@
 
 
 class DotsOCRImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - np: The total number of patches over each image over each prompt in
-              the batch
-        - ni: Number of images
-        - cps: Number of channels * patch_size * patch_size
+    """Dimensions:
+    - np: The total number of patches over each image over each prompt in
+          the batch
+    - ni: Number of images
+    - cps: Number of channels * patch_size * patch_size
     """
 
     type: Literal["pixel_values"]
@@ -80,11 +80,10 @@ class DotsOCRImagePixelInputs(TensorSchema):
 
 
 class DotsOCRImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - nf: Number of image features
-        - hs: Hidden size
-        - ni: Number of images
+    """Dimensions:
+    - nf: Number of image features
+    - hs: Hidden size
+    - ni: Number of images
     """
 
     type: Literal["image_embeds"]
@@ -105,20 +104,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
         }
 
@@ -718,13 +713,16 @@ def _process_image_input(
 
         return image_embeds.split(sizes)
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        merge_size = self.vision_tower.spatial_merge_size
-        return num_image_tokens * (merge_size**2)
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         merge_size = self.vision_tower.spatial_merge_size
-        return num_vision_tokens // (merge_size**2)
+        return num_mm_embeds * (merge_size**2), num_mm_embeds
 
     def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
         image_input = self._parse_and_validate_image_input(**kwargs)
@@ -764,9 +762,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="vision_tower.merger",
diff --git a/vllm/model_executor/models/eagle2_5_vl.py b/vllm/model_executor/models/eagle2_5_vl.py
index 8fdd7098f224..257907e659e3 100644
--- a/vllm/model_executor/models/eagle2_5_vl.py
+++ b/vllm/model_executor/models/eagle2_5_vl.py
@@ -8,7 +8,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -37,13 +37,12 @@
 
 
 class Eagle2_5_VLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - bnp: Batch size * number of images * (1 + num_patches)
-        - c: Number of channels (3)
-        - h: Height of each image patch
-        - w: Width of each image patch
+    """Dimensions:
+    - bn: Batch size * number of images
+    - bnp: Batch size * number of images * (1 + num_patches)
+    - c: Number of channels (3)
+    - h: Height of each image patch
+    - w: Width of each image patch
     """
 
     type: Literal["pixel_values"]
@@ -52,11 +51,10 @@ class Eagle2_5_VLImagePixelInputs(TensorSchema):
 
 
 class Eagle2_5_VLImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - n: Number of images
-        - f: Total image feature size
-        - h: Hidden size (must match the hidden size of language model backbone)
+    """Dimensions:
+    - n: Number of images
+    - f: Total image feature size
+    - h: Hidden size (must match the hidden size of language model backbone)
     """
 
     type: Literal["image_embeds"]
@@ -127,8 +125,7 @@ class Eagle2_5_VLMultiModalProcessor(
 class Eagle2_5_VLForConditionalGeneration(
     nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
 ):
-    """
-    Eagle2.5-VL model for conditional generation.
+    """Eagle2.5-VL model for conditional generation.
 
     Architecture:
         - Vision Encoder: SigLIP
@@ -196,7 +193,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
 
     def _init_vision_model(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str,
     ):
@@ -223,7 +220,7 @@ def _init_vision_model(
             prefix=prefix,
         )
 
-    def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
+    def _init_mlp1(self, config: PreTrainedConfig) -> nn.Module:
         """Initialize MLP projection layer."""
         vit_hidden_size = config.vision_config.hidden_size
         llm_hidden_size = config.text_config.hidden_size
@@ -238,8 +235,7 @@ def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
         )
 
     def pixel_shuffle(self, x: torch.Tensor, scale_factor: float = 0.5) -> torch.Tensor:
-        """
-        Pixel shuffle operation for downsampling vision features.
+        """Pixel shuffle operation for downsampling vision features.
 
         Args:
             x: Input tensor of shape (n, w, h, c)
@@ -247,6 +243,7 @@ def pixel_shuffle(self, x: torch.Tensor, scale_factor: float = 0.5) -> torch.Ten
 
         Returns:
             Downsampled tensor
+
         """
         n, w, h, c = x.size()
         # N, W, H, C --> N, W, H * scale, C // scale
@@ -264,14 +261,14 @@ def pixel_shuffle(self, x: torch.Tensor, scale_factor: float = 0.5) -> torch.Ten
         return x
 
     def extract_feature(self, pixel_values: torch.Tensor) -> torch.Tensor:
-        """
-        Extract visual features from pixel values.
+        """Extract visual features from pixel values.
 
         Args:
             pixel_values: Input pixel values of shape (batch, channels, height, width)
 
         Returns:
             Visual embeddings
+
         """
         vit_embeds = self.vision_model(pixel_values=pixel_values)
 
diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py
index 2683e94f51b9..53d1ba682744 100644
--- a/vllm/model_executor/models/ernie45_moe.py
+++ b/vllm/model_executor/models/ernie45_moe.py
@@ -29,7 +29,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
@@ -41,12 +41,15 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+    MoERunner,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -116,7 +119,7 @@ def forward(self, x):
 class Ernie4_5_MoeMoE(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
         enable_eplb: bool = False,
@@ -150,12 +153,11 @@ def __init__(
                 f"the number of experts {config.moe_num_experts}."
             )
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.moe_num_experts,
-            bias=False,
+            out_dtype=torch.float32,
             params_dtype=torch.float32,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
@@ -199,7 +201,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_dim = hidden_states.shape[-1]
         hidden_states = hidden_states.view(-1, hidden_dim)
 
-        router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
+        router_logits, _ = self.gate(hidden_states)
 
         final_hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
@@ -303,7 +305,7 @@ def forward(
 class Ernie4_5_MoeDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py
index 716aeaa3b708..0a632b66e659 100644
--- a/vllm/model_executor/models/ernie45_vl.py
+++ b/vllm/model_executor/models/ernie45_vl.py
@@ -36,7 +36,7 @@
 from transformers import BaseImageProcessor, BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions
 from vllm.distributed import parallel_state
 from vllm.distributed import utils as dist_utils
 from vllm.inputs import MultiModalDataDict
@@ -71,6 +71,7 @@
     PromptUpdate,
     cached_encode,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.sequence import IntermediateTensors
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -119,7 +120,7 @@ def all_gather_interleave(local_tensor, hidden_size: int, tp_size: int):
 
 
 class Ernie4_5_VisionAttention(nn.Module):
-    """VisionAttention using VLLM framework APIs"""
+    """VisionAttention using VLLM framework APIs."""
 
     def __init__(
         self,
@@ -497,6 +498,7 @@ def prepare_encoder_metadata(
 
         Returns:
             Dict with ``rotary_pos_emb``, ``cu_seqlens`` and ``max_seqlen``.
+
         """
         if device is None:
             device = self.device
@@ -600,12 +602,11 @@ def load_weights(self, weights) -> set[str]:
 
 
 class Ernie4_5_VLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - np: The total number of patches over each image over each prompt in
-              the batch
-        - ni: Number of images
-        - cps: Number of channels * patch_size * patch_size
+    """Dimensions:
+    - np: The total number of patches over each image over each prompt in
+          the batch
+    - ni: Number of images
+    - cps: Number of channels * patch_size * patch_size
     """
 
     type: Literal["pixel_values"]
@@ -618,13 +619,12 @@ class Ernie4_5_VLImagePixelInputs(TensorSchema):
 
 
 class Ernie4_5_VLVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - np: The total number of patches over each image over each prompt in
-              the batch
-        - ni: Number of images
-        - cps: Number of channels * temporal_patch_size * patch_size *
-              patch_size
+    """Dimensions:
+    - np: The total number of patches over each image over each prompt in
+          the batch
+    - ni: Number of images
+    - cps: Number of channels * temporal_patch_size * patch_size *
+          patch_size
     """
 
     type: Literal["pixel_values_videos"]
@@ -881,7 +881,7 @@ def get_hf_config(self):
         return self.ctx.model_config.hf_config
 
     def get_hf_processor(self, **kwargs: object):
-        return self.ctx.get_hf_processor(use_fast=True, **kwargs)
+        return self.ctx.get_hf_processor(**kwargs)
 
     def get_image_processor(self, **kwargs: object):
         return self.get_hf_processor(**kwargs).image_processor
@@ -1113,34 +1113,32 @@ def _pixel_values_norm(
         pixel_values = pixel_values.to(hf_config.dtype)
         return pixel_values
 
-    def _apply_hf_processor_main(
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if "images" not in mm_data and "videos" not in mm_data:
-            return BatchFeature(dict(passthrough_data))
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+        hf_data = hf_inputs.hf_data
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        if "images" not in hf_data and "videos" not in hf_data:
+            return hf_inputs
 
-        mm_data = dict(mm_data)
-        if "images" not in mm_data:
-            mm_data["images"] = []
-        if "videos" not in mm_data:
-            mm_data["videos"] = []
+        prompt_text = hf_data["text"]
+        hf_data["text"] = [prompt_text]
+        hf_data.setdefault("images", [])
+        hf_data.setdefault("videos", [])
 
         # Check if HF processor supports video metadata
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        hf_processor = self.info.get_hf_processor(**hf_inputs.hf_kwargs)
         supports_video_metadata = getattr(
             hf_processor, "supports_video_metadata", False
         )
 
-        videos = mm_data["videos"]
+        videos = hf_data["videos"]
         assert isinstance(videos, Sequence)
         if videos and not supports_video_metadata:
             # Old HF processor, unwrap tuple to pure frames
@@ -1148,49 +1146,43 @@ def _apply_hf_processor_main(
                 "HF processor doesn't support video metadata. "
                 "Timestamps will NOT be rendered. Please upgrade the model."
             )
-            mm_data["videos"] = [v[0] if isinstance(v, tuple) else v for v in videos]
+            hf_data["videos"] = [v[0] if isinstance(v, tuple) else v for v in videos]
 
-        processor_output = self.info.ctx.call_hf_processor(
-            hf_processor,
-            dict(
-                text=[prompt_text], images=mm_data["images"], videos=mm_data["videos"]
-            ),
-            hf_processor_mm_kwargs,
-        )
+        return hf_inputs
+
+    def _postprocess_hf_mm_data(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+        processed_data: BatchFeature,
+    ) -> BatchFeature:
+        if "images" not in hf_data and "videos" not in hf_data:
+            return processed_data
 
         # Divide the processor_output into two modalities: image and video.
-        if processor_output is not None:
-            pixel_values = processor_output["images"]
-            if pixel_values is not None:
-                processor_output["images"] = self._pixel_values_norm(
-                    pixel_values, hf_processor_mm_kwargs
-                )
-            for key in list(processor_output.keys()):
-                if processor_output[key] is None:
-                    del processor_output[key]
-                    continue
-                if key == "grid_thw":
-                    grid_thw = processor_output["grid_thw"]
-                    pixel_values_all = processor_output["images"]
-                    # Identify elements where the first
-                    # dimension is greater than 1 and
-                    # treat them as the video modality
-                    mask = grid_thw[:, 0] > 1
-                    processor_output["video_grid_thw"] = grid_thw[mask]
-                    processor_output["image_grid_thw"] = grid_thw[~mask]
-                    image_patch_num = (
-                        processor_output["image_grid_thw"].prod(dim=1).sum()
-                    )
-                    processor_output["pixel_values"] = pixel_values_all[
-                        :image_patch_num
-                    ]
-                    processor_output["pixel_values_videos"] = pixel_values_all[
-                        image_patch_num:
-                    ]
-                    del processor_output["images"]
+        pixel_values = processed_data["images"]
+        if pixel_values is not None:
+            processed_data["images"] = self._pixel_values_norm(pixel_values, hf_kwargs)
+        for key in list(processed_data.keys()):
+            if processed_data[key] is None:
+                del processed_data[key]
+                continue
+            if key == "grid_thw":
+                grid_thw = processed_data["grid_thw"]
+                pixel_values_all = processed_data["images"]
+                # Identify elements where the first
+                # dimension is greater than 1 and
+                # treat them as the video modality
+                mask = grid_thw[:, 0] > 1
+                processed_data["video_grid_thw"] = grid_thw[mask]
+                processed_data["image_grid_thw"] = grid_thw[~mask]
+                image_patch_num = processed_data["image_grid_thw"].prod(dim=1).sum()
+                processed_data["pixel_values"] = pixel_values_all[:image_patch_num]
+                processed_data["pixel_values_videos"] = pixel_values_all[
+                    image_patch_num:
+                ]
+                del processed_data["images"]
 
-        processed_data = processor_output
-        processed_data.update(passthrough_data)
         return processed_data
 
     def _get_prompt_updates(
@@ -1284,33 +1276,26 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(
             seq_len, mm_counts
         )
 
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-        assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
 
diff --git a/vllm/model_executor/models/ernie45_vl_moe.py b/vllm/model_executor/models/ernie45_vl_moe.py
index ccc4b9b868b6..4965bcdcb3b4 100644
--- a/vllm/model_executor/models/ernie45_vl_moe.py
+++ b/vllm/model_executor/models/ernie45_vl_moe.py
@@ -29,7 +29,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 # from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -192,7 +192,7 @@ def forward(
 class Ernie4_5_VLMoeMoE(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -385,7 +385,7 @@ def forward(
 class Ernie4_5_VLMoeDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/exaone.py b/vllm/model_executor/models/exaone.py
index fef05d386a9b..16f07746f392 100644
--- a/vllm/model_executor/models/exaone.py
+++ b/vllm/model_executor/models/exaone.py
@@ -30,7 +30,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -104,7 +104,7 @@ def forward(self, x):
 class ExaoneAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
@@ -191,7 +191,7 @@ def forward(
 class ExaoneBlockAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
@@ -228,7 +228,7 @@ def forward(
 class ExaoneDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/exaone4_5.py b/vllm/model_executor/models/exaone4_5.py
index b325c6a3043e..5a53e60b7d5e 100644
--- a/vllm/model_executor/models/exaone4_5.py
+++ b/vllm/model_executor/models/exaone4_5.py
@@ -311,7 +311,6 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs: object) -> Exaone4_5_Processor:
         return self.ctx.get_hf_processor(
             Exaone4_5_Processor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py
index 4c19f1ed0fcf..d5b8b4066239 100644
--- a/vllm/model_executor/models/exaone_moe.py
+++ b/vllm/model_executor/models/exaone_moe.py
@@ -20,7 +20,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
@@ -30,11 +30,13 @@
     get_tensor_model_parallel_world_size,
 )
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -63,7 +65,7 @@
 class ExaoneMoe(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         swiglu_limit: float | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -84,11 +86,9 @@ def __init__(
                 f"the number of experts {config.num_experts}."
             )
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
@@ -289,7 +289,7 @@ def forward(
 class ExaoneMoeDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         is_mtp: bool = False,
diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py
index c11c64fcba5a..653e7691fba3 100644
--- a/vllm/model_executor/models/extract_hidden_states.py
+++ b/vllm/model_executor/models/extract_hidden_states.py
@@ -45,8 +45,7 @@ def unified_kv_cache_update(
     to_cache: torch.Tensor,
     layer_name: str,
 ) -> torch.Tensor:
-    """
-    Returns a dummy that is passed to unified_attention to signal a side effect and
+    """Returns a dummy that is passed to unified_attention to signal a side effect and
     the data dependency between them to ensure torch.compile preserves ordering.
     """
     forward_context = get_forward_context()
@@ -294,6 +293,7 @@ def forward(self, to_cache: torch.Tensor) -> torch.Tensor:
 
         Returns:
             Dummy output tensor (not used)
+
         """
         # Note: we set num_heads to num_hidden_layers and
         # head_size to hidden_size for hidden states storage
@@ -373,8 +373,8 @@ def forward(self, hidden_states: torch.Tensor) -> None:
 
         Returns:
             Tuple of (dummy_output, dummy_output) - both unused
-        """
 
+        """
         # Call dummy attention layer to cache hidden states
         # Output is ignored - we only care about the KV cache side effects
         _ = self.cache_only_layers[str(self.target_num_hidden_layers)](hidden_states)
diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py
index 5b9853ece10b..f18f938774dc 100644
--- a/vllm/model_executor/models/falcon_h1.py
+++ b/vllm/model_executor/models/falcon_h1.py
@@ -140,8 +140,7 @@ def __init__(
         self._init_mup_vector()
 
     def _init_mup_vector(self):
-        """
-        Non learnable per-block scaling vector composed of element-wise
+        """Non learnable per-block scaling vector composed of element-wise
         multipliersapplied to each separate contiguous block of the output
         of the linear projection (in_proj) before further processing
         (gating, convolution, SSM):
@@ -313,8 +312,7 @@ def forward(
 
 
 class FalconH1ParallelHybrid(nn.Module):
-    """
-    A hybrid decoder layer for FalconH1 where the input is processed
+    """A hybrid decoder layer for FalconH1 where the input is processed
     in parallel through both the self-attention branch and the SSM (Mamba)
     branch. Their outputs are then summed to produce the final hidden state.
 
@@ -552,6 +550,7 @@ def get_mamba_state_shape_from_config(
             Tuple containing:
             - conv_state_shape: Shape for convolutional state cache
             - temporal_state_shape: Shape for state space model cache
+
         """
         parallel_config = vllm_config.parallel_config
         hf_config = vllm_config.model_config.hf_config
diff --git a/vllm/model_executor/models/fireredasr2.py b/vllm/model_executor/models/fireredasr2.py
index b9d99474597a..be5f5efbff55 100644
--- a/vllm/model_executor/models/fireredasr2.py
+++ b/vllm/model_executor/models/fireredasr2.py
@@ -12,7 +12,7 @@
 )
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.inputs import MultiModalDataDict, PromptType
 from vllm.logger import init_logger
@@ -38,6 +38,7 @@
     PromptUpdate,
     PromptUpdateDetails,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.transformers_utils.processor import cached_processor_from_config
 from vllm.transformers_utils.processors.fireredasr2 import (
     FireRedASR2FeatureExtractor,
@@ -63,11 +64,10 @@
 
 
 class FireRedASR2AudioInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - nmb: Number of mel bins
-        - t: Time frames (M)
+    """Dimensions:
+    - b: Batch size
+    - nmb: Number of mel bins
+    - t: Time frames (M)
     """
 
     input_features: Annotated[
@@ -207,20 +207,18 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         feature_extractor = self.info.get_feature_extractor()
 
         sampling_rate = feature_extractor.sampling_rate
         audio_len = feature_extractor.chunk_length * sampling_rate
-        num_audios = mm_counts.get("audio", 0)
-
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
 
         ret = {
             "audio": self._get_dummy_audios(
-                length=audio_len, num_audios=num_audios, overrides=audio_overrides
+                length=audio_len,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
         return ret
@@ -229,37 +227,24 @@ def get_dummy_mm_data(
 class FireRedASR2MultiModalProcessor(
     BaseMultiModalProcessor[FireRedASR2ProcessingInfo]
 ):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
-
-        mm_data = dict(mm_data)
-        mm_data["audio"] = mm_data.pop("audios")
-
-        hf_processor_mm_kwargs = dict(
-            **hf_processor_mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+        return hf_inputs._replace(
+            hf_kwargs=dict(
+                hf_inputs.hf_kwargs,
+                sampling_rate=feature_extractor.sampling_rate,
+            )
         )
 
-        return mm_data, hf_processor_mm_kwargs
-
-    def _postprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-        processed_data: BatchFeature,
-    ) -> BatchFeature:
-        if "labels" in processed_data:
-            processed_data["input_ids"] = processed_data.pop("labels")
-
-        return processed_data
-
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
diff --git a/vllm/model_executor/models/funasr.py b/vllm/model_executor/models/funasr.py
index b7981f42d364..f8db52378570 100644
--- a/vllm/model_executor/models/funasr.py
+++ b/vllm/model_executor/models/funasr.py
@@ -14,7 +14,7 @@
 )
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict, PromptType
@@ -49,7 +49,10 @@
     PromptUpdate,
     TimingContext,
 )
-from vllm.multimodal.processing.processor import MultiModalProcessingInfo
+from vllm.multimodal.processing.processor import (
+    HFMultiModalInputs,
+    MultiModalProcessingResult,
+)
 from vllm.transformers_utils.processor import cached_processor_from_config
 from vllm.transformers_utils.processors.funasr import FunASRFeatureExtractor
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -598,11 +601,10 @@ def forward(self, hidden_states: torch.Tensor, ilens: torch.Tensor):
 
 
 class FunASRAudioInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - nmb: Number of mel bins
-        - t: Time frames (M)
+    """Dimensions:
+    - b: Batch size
+    - nmb: Number of mel bins
+    - t: Time frames (M)
     """
 
     input_features: Annotated[
@@ -732,63 +734,46 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         feature_extractor = self.info.get_feature_extractor()
 
         sampling_rate = feature_extractor.sampling_rate
         audio_len = feature_extractor.chunk_length * sampling_rate
-        num_audios = mm_counts.get("audio", 0)
-
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
 
         return {
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             ),
         }
 
 
 class FunASRMultiModalProcessor(BaseMultiModalProcessor[FunASRProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
-
-        mm_data = dict(mm_data)
-        mm_data["audio"] = mm_data.pop("audios")
-
-        hf_processor_mm_kwargs = dict(
-            **hf_processor_mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+        return hf_inputs._replace(
+            hf_kwargs=dict(
+                hf_inputs.hf_kwargs,
+                sampling_rate=feature_extractor.sampling_rate,
+            )
         )
 
-        return mm_data, hf_processor_mm_kwargs
-
-    def _postprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-        processed_data: BatchFeature,
-    ) -> BatchFeature:
-        if "labels" in processed_data:
-            processed_data["input_ids"] = processed_data.pop("labels")
-
-        return processed_data
-
     def _cached_apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
+    ) -> MultiModalProcessingResult:
         # Dithering injects noise into the extracted features, so the
         # feature extractor is not a pure function of its input. Since the
         # processing cache assumes that processor outputs are invariant
diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py
index 607eb4ea5d05..8b40ddf0cd43 100644
--- a/vllm/model_executor/models/funaudiochat.py
+++ b/vllm/model_executor/models/funaudiochat.py
@@ -26,7 +26,7 @@
 from transformers.modeling_outputs import BaseModelOutput
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
 from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
@@ -580,7 +580,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         feature_extractor = self.info.get_feature_extractor()
         sampling_rate = int(feature_extractor.sampling_rate)
@@ -597,15 +597,12 @@ def get_dummy_mm_data(
             1,
             (target_num_frames * sampling_rate + token_fps - 1) // token_fps,
         )
-        num_audios = int(mm_counts.get("audio", 0))
 
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
         return {
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=int(mm_counts.get("audio", 0)),
+                overrides=mm_options.get("audio"),
             )
         }
 
@@ -616,29 +613,29 @@ class FunAudioChatMultiModalProcessor(
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
         if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+            return self._finalize_hf_mm_data(mm_data, hf_kwargs, passthrough_data)
 
         prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
 
         tokenizer = self.info.get_tokenizer()
         input_ids = torch.tensor([tokenizer.encode(prompt_text)])
 
-        audios = mm_data.get("audios", [])
+        audios = mm_data.get("audio", [])
         if not audios:
             processed_data = BatchFeature({"input_ids": input_ids})
-            processed_data.update(passthrough_data)
-            return processed_data
+            return self._finalize_hf_mm_data(
+                mm_data, hf_kwargs, passthrough_data, processed_data
+            )
         assert isinstance(audios, Sequence)
 
-        feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
         sr = int(feature_extractor.sampling_rate)
         min_samples = int(getattr(feature_extractor, "n_fft", 400) or 400)
 
@@ -690,8 +687,9 @@ def _apply_hf_processor_main(
         }
 
         processed_data = BatchFeature({"input_ids": input_ids, **mm_inputs})
-        processed_data.update(passthrough_data)
-        return processed_data
+        return self._finalize_hf_mm_data(
+            mm_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_mm_fields_config(
         self,
diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py
index b45327fbebf7..c6166bc2f423 100644
--- a/vllm/model_executor/models/gemma3_mm.py
+++ b/vllm/model_executor/models/gemma3_mm.py
@@ -11,7 +11,7 @@
 from transformers.models.gemma3.processing_gemma3 import Gemma3ProcessorKwargs
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
@@ -19,6 +19,7 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
@@ -57,14 +58,13 @@
 
 
 class Gemma3ImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - p: Number of patches total (over each image over each prompt in the
-          batch)
-        - c: Number of channels (3)
-        - h: Height of each patch
-        - w: Width of each patch
-        - bn: Batch size * number of images
+    """Dimensions:
+    - p: Number of patches total (over each image over each prompt in the
+      batch)
+    - c: Number of channels (3)
+    - h: Height of each patch
+    - w: Width of each patch
+    - bn: Batch size * number of images
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -244,26 +244,22 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
 
 class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -677,52 +673,23 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
             tower_model="vision_tower",
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        """
-        Calculate the number of tokens output by the vision encoder.
-
-        The vision encoder processes images into patch embeddings. For Gemma3,
-        the relationship between prompt placeholder tokens and actual vision
-        encoder output tokens depends on the patch grid size.
-
-        Args:
-            num_image_tokens: Number of image placeholder tokens in the prompt
-                              (typically mm_tokens_per_image per image)
-
-        Returns:
-            Number of tokens output by the vision encoder
-        """
-        # For Gemma3, the vision encoder outputs tokens_per_side x tokens_per_side
-        # tokens per image. Since num_image_tokens represents the number of
-        # connector output tokens (mm_tokens_per_image = 256), and tokens_per_side
-        # is sqrt(256) = 16, we need to account for the token expansion.
-        # Based on empirical testing, the multiplier of 16 works correctly.
-        return num_image_tokens * 16
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
-        """
-        Calculate the number of tokens output by the multimodal connector.
-
-        The connector applies projection and normalization but maintains the
-        token count for Gemma3.
-
-        Args:
-            num_vision_tokens: Number of tokens from vision encoder
-
-        Returns:
-            Number of tokens after connector processing
-        """
-        # The Gemma3 connector maintains a 1:1 token mapping
-        return num_vision_tokens
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+        tower_tokens = num_mm_embeds * 16
+        return tower_tokens, tower_tokens
 
     def get_encoder_cudagraph_config(self):
         from vllm.v1.worker.encoder_cudagraph_defs import (
diff --git a/vllm/model_executor/models/gemma3n.py b/vllm/model_executor/models/gemma3n.py
index 56d64735789a..3c9df39e674c 100644
--- a/vllm/model_executor/models/gemma3n.py
+++ b/vllm/model_executor/models/gemma3n.py
@@ -185,7 +185,7 @@ def correct(
 
 
 class Gemma3nLaurelBlock(nn.Module):
-    """Learned Augmented Residual Layer"""
+    """Learned Augmented Residual Layer."""
 
     def __init__(
         self,
@@ -593,9 +593,7 @@ def forward(
     enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill
 )
 class Gemma3nSelfDecoder(nn.Module):
-    """
-    Includes altup embedding and self decoder layers
-    """
+    """Includes altup embedding and self decoder layers."""
 
     def __init__(
         self,
@@ -776,9 +774,7 @@ def forward(
     enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill
 )
 class Gemma3nCrossDecoder(nn.Module):
-    """
-    Cross-decoder layers
-    """
+    """Cross-decoder layers."""
 
     def __init__(
         self,
diff --git a/vllm/model_executor/models/gemma3n_audio_utils.py b/vllm/model_executor/models/gemma3n_audio_utils.py
index bef9bb9a0d4e..427c98e43a17 100644
--- a/vllm/model_executor/models/gemma3n_audio_utils.py
+++ b/vllm/model_executor/models/gemma3n_audio_utils.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Lightweight utility functions for Gemma3n audio processing.
+"""Lightweight utility functions for Gemma3n audio processing.
 
 This module is separate from gemma3n_mm.py to avoid heavy CUDA dependencies,
 making it testable without a full vLLM build.
@@ -37,6 +36,7 @@ def adjust_audio_features_to_expected_length(
         Tuple of:
         - adjusted_features: Audio features adjusted to expected_tokens length
         - tokens_truncated: Number of tokens truncated (0 if padding was applied)
+
     """
     audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape
     tokens_truncated = 0
diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py
index 1c05c165dcfe..f2cb959e4714 100644
--- a/vllm/model_executor/models/gemma3n_mm.py
+++ b/vllm/model_executor/models/gemma3n_mm.py
@@ -15,10 +15,10 @@
     Gemma3nTextConfig,
     Gemma3nVisionConfig,
 )
-from transformers.models.siglip import SiglipImageProcessorFast
+from transformers.models.siglip import SiglipImageProcessor
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.inputs import MultiModalDataDict, PromptType, TextPrompt
 from vllm.logger import init_logger
@@ -74,12 +74,11 @@
 
 
 class Gemma3nImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height of each patch
-        - w: Width of each patch
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height of each patch
+    - w: Width of each patch
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -113,6 +112,7 @@ def batch_audio_features(
 
     Returns:
         The `(bn, s_max, f)` features and their `(bn, s_max)` mask.
+
     """
     if isinstance(input_features, torch.Tensor):
         assert isinstance(input_features_mask, torch.Tensor)
@@ -135,11 +135,10 @@ def batch_audio_features(
 
 
 class Gemma3nAudioInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of audios
-        - s: seq_length
-        - f: num_features
+    """Dimensions:
+    - bn: Batch size * number of audios
+    - s: seq_length
+    - f: num_features
     """
 
     type: Literal["audio"] = "audio"
@@ -187,8 +186,7 @@ def get_image_repl(
         image_height: int,
         processor: Gemma3nProcessor,
     ) -> PromptUpdateDetails:
-        """
-        Get the replacement metadata for image tokens.
+        """Get the replacement metadata for image tokens.
 
         For Gemma3n, this should return the full_image_sequence which includes
         BOI token, repeated image tokens, and EOI token.
@@ -205,8 +203,7 @@ def get_audio_repl(
         *,
         processor: Gemma3nProcessor,
     ) -> PromptUpdateDetails:
-        """
-        Get the replacement metadata for audio tokens.
+        """Get the replacement metadata for audio tokens.
 
         For Gemma3n, this should return the full_audio_sequence which includes
         BOA token, repeated audio tokens, and EOA token.
@@ -235,53 +232,36 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_audios = mm_counts.get("audio", 0)
         processor = self.info.get_hf_processor()
         audio_feature_extractor: Gemma3nAudioFeatureExtractor = (
             processor.feature_extractor
         )
         audio_len = audio_feature_extractor.fft_length
-        image_processor: SiglipImageProcessorFast = processor.image_processor
+        image_processor: SiglipImageProcessor = processor.image_processor
         img_width = image_processor.size.get("width", 224)
         img_height = image_processor.size.get("height", 224)
 
-        image_overrides = mm_options.get("image")
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
-
         return {
             "image": self._get_dummy_images(
                 width=img_width,
                 height=img_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             ),
         }
 
 
 class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        mm_data = dict(mm_data)
-        if "audios" in mm_data:
-            mm_data["audio"] = mm_data.pop("audios")
-
-        return mm_data, hf_processor_mm_kwargs
-
     def _postprocess_hf_mm_data(
         self,
         mm_data: Mapping[str, object],
@@ -316,8 +296,7 @@ def _postprocess_hf_mm_data(
         return processed_data
 
     def _get_num_audio_frames_alone(self, num_samples: int) -> int:
-        """
-        Get the number of mel frames that the feature extractor produces
+        """Get the number of mel frames that the feature extractor produces
         for a single audio clip of the given length, independent of batch
         padding.
         """
@@ -571,6 +550,7 @@ def forward(
 
         Returns:
             A torch.Tensor of embeddings with  shape `[batch_size, seq_len, self.config.text_config.hidden_size]`.
+
         """  # noqa: E501
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError(
@@ -891,9 +871,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
@@ -911,8 +889,7 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
 
     @classmethod
     def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
-        """
-        Gemma3n supports "free-form" transcription.
+        """Gemma3n supports "free-form" transcription.
         We fix its prompt here to standardize transcriptions/translations
         requests.
         """
diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py
index 961709a7e799..e3dfe65c2f58 100644
--- a/vllm/model_executor/models/gemma4.py
+++ b/vllm/model_executor/models/gemma4.py
@@ -22,12 +22,12 @@
 from dataclasses import replace
 from itertools import islice
 
-import regex as re
 import torch
 from torch import nn
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
+from vllm.config.utils import getattr_iter
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_rank,
@@ -40,7 +40,7 @@
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
     GateLinear,
-    fused_moe_make_expert_params_mapping,
+    MoERunner,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
@@ -57,10 +57,6 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
@@ -73,27 +69,63 @@
     SupportsEagle3,
     SupportsLoRA,
     SupportsPP,
+    SupportsQuant,
 )
 from .utils import (
     AutoWeightsLoader,
+    PPMissingLayer,
+    ShardIds,
     WeightsMapper,
     extract_layer_index,
-    is_pp_missing_parameter,
     make_layers,
     maybe_prefix,
 )
 
 logger = init_logger(__name__)
 
-_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
-    orig_to_new_regex={
-        re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
-    }
-)
 
+def _gemma4_layer_weights_mapper(config) -> WeightsMapper:
+    """Stack each decoder layer's projections the way its modules are built.
 
-def _remap_gemma4_expert_weight_name(name: str) -> str:
-    return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
+    gate/up pack into `mlp.gate_up_proj`. Layers with a qkv_proj pack q/k/v
+    into it; `attention_k_eq_v` full-attention layers ship k_proj only, so K
+    is loaded into both the K and V shards and no v_proj is expected. KV-shared
+    layers only have q_proj; the K/V tensors original checkpoints still ship
+    for them are dropped.
+    """
+    num_layers = config.num_hidden_layers
+    first_kv_shared = num_layers - getattr(config, "num_kv_shared_layers", 0)
+    k_eq_v_layers = (
+        {
+            i
+            for i, layer_type in enumerate(config.layer_types[:first_kv_shared])
+            if layer_type == "full_attention"
+        }
+        if getattr(config, "attention_k_eq_v", False)
+        else set()
+    )
+    qkv_shards: tuple[tuple[str, ShardIds], ...] = (("q", "q"), ("k", "k"), ("v", "v"))
+    k_eq_v_shards: tuple[tuple[str, ShardIds], ...] = (("q", "q"), ("k", ["k", "v"]))
+    stacked: dict[str, tuple[str, ShardIds]] = {
+        ".mlp.gate_proj.": (".mlp.gate_up_proj.", 0),
+        ".mlp.up_proj.": (".mlp.gate_up_proj.", 1),
+    }
+    stacked |= {
+        f"layers.{i}.self_attn.{shard}_proj.": (
+            f"layers.{i}.self_attn.qkv_proj.",
+            shard_id,
+        )
+        for i in range(first_kv_shared)
+        for shard, shard_id in (k_eq_v_shards if i in k_eq_v_layers else qkv_shards)
+    }
+    return WeightsMapper(
+        orig_to_new_substr={
+            f"layers.{i}.self_attn.{name}.": None
+            for i in range(first_kv_shared, num_layers)
+            for name in ("k_proj", "v_proj", "k_norm")
+        },
+        orig_to_new_stacked=stacked,
+    )
 
 
 @triton.jit
@@ -279,6 +311,9 @@ def __init__(
         self.norm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps, has_weight=False)
         # Per-dimension learned scale, applied after norm + root_size
         self.scale = nn.Parameter(torch.ones(self.hidden_size))
+        # Per-expert output scale folded into routing weights so that
+        # MoERunner's fused kernel computes: Σ_e (expert_e * w_e * scale_e)
+        self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts))
         # Constant 1/sqrt(hidden_size) scaling factor
         self.register_buffer(
             "root_size",
@@ -296,84 +331,36 @@ def __init__(
             prefix=f"{prefix}.proj",
         )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Returns raw router logits [T, E]."""
-        x = self.norm(x)
-        x = x * self.root_size.to(x.dtype)
-        x = x * self.scale.to(x.dtype)
-        router_logits, _ = self.proj(x)
-        return router_logits
-
-
-class Gemma4MoE(nn.Module):
-    """Mixture of Experts for Gemma4 using vLLM's MoERunner.
-
-    Wraps MoERunner with custom routing. The router projection is
-    external (Gemma4Router) — this class only handles expert dispatch.
-
-    Gemma4 routing: softmax over ALL experts → top-k → renormalize.
-    per_expert_scale is folded into routing weights for mathematical
-    correctness with MoERunner's fused kernel.
-    """
-
-    def __init__(
+    def routing_function(
         self,
-        config,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.num_experts = config.num_experts
-
-        # Per-expert output scale folded into routing weights so that
-        # MoERunner's fused kernel computes: Σ_e (expert_e * w_e * scale_e)
-        self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts))
+        hidden_states: torch.Tensor,
+        gating_output: torch.Tensor,
+        topk: int,
+        renormalize: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Softmax over ALL experts → top-k → renormalize.
 
-        # Gemma4 routing: softmax over ALL experts → top-k → renormalize.
-        # MoERunner's built-in fused_topk scopes softmax differently, so
-        # a custom routing function is needed for numerical correctness.
-        # NOTE: self.per_expert_scale is read at call time (not captured into
-        # a local) so that torch.func.functional_call parameter substitution
-        # reaches the routing function correctly.
-        def routing_function(
-            hidden_states: torch.Tensor,
-            gating_output: torch.Tensor,
-            topk: int,
-            renormalize: bool,
-        ) -> tuple[torch.Tensor, torch.Tensor]:
-            if current_platform.is_cuda_alike() or current_platform.is_xpu():
-                return gemma4_fused_routing_kernel_triton(
-                    gating_output, topk, self.per_expert_scale
-                )
+        MoERunner's built-in fused_topk scopes softmax differently, so Gemma4
+        needs its own routing for numerical correctness.
 
-            return gemma4_routing_function_torch(
+        NOTE: self.per_expert_scale is read at call time (not captured into a
+        local) so that torch.func.functional_call parameter substitution
+        reaches the routing function correctly.
+        """
+        if current_platform.is_cuda_alike() or current_platform.is_xpu():
+            return gemma4_fused_routing_kernel_triton(
                 gating_output, topk, self.per_expert_scale
             )
 
-        # MoERunner experts with custom Gemma4 routing
-        intermediate_size = getattr(
-            config,
-            "moe_intermediate_size",
-            getattr(config, "expert_intermediate_size", None),
-        )
-        if intermediate_size is None:
-            raise ValueError("Gemma4 MoE requires an expert intermediate size")
-
-        self.experts = FusedMoEFactory(
-            num_experts=config.num_experts,
-            top_k=config.top_k_experts,
-            hidden_size=config.hidden_size,
-            intermediate_size=intermediate_size,
-            renormalize=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.experts",
-            custom_routing_function=routing_function,
-            activation="gelu_tanh",
-        )
+        return gemma4_routing_function_torch(gating_output, topk, self.per_expert_scale)
 
-    def forward(self, x: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor:
-        return self.experts(x, router_logits)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Returns raw router logits [T, E]."""
+        x = self.norm(x)
+        x = x * self.root_size.to(x.dtype)
+        x = x * self.scale.to(x.dtype)
+        router_logits, _ = self.proj(x)
+        return router_logits
 
 
 class Gemma4Attention(nn.Module):
@@ -433,7 +420,7 @@ def __init__(
         else:
             # QKVParallelLinear handles GQA correctly for all layer types.
             # k_eq_v layers load K weights into both K and V slots via
-            # _weight_iterator remapping — no structural difference needed.
+            # hf_to_vllm_mapper — no structural difference needed.
             self.qkv_proj = QKVParallelLinear(
                 hidden_size,
                 self.head_dim,
@@ -645,7 +632,7 @@ def __init__(
             config, "use_second_mlp_block", False
         )
         self.router: Gemma4Router | None
-        self.moe: Gemma4MoE | None
+        self.experts: MoERunner | None
         self.post_feedforward_layernorm_1: RMSNorm | None
         self.post_feedforward_layernorm_2: RMSNorm | None
         self.pre_feedforward_layernorm_2: RMSNorm | None
@@ -655,10 +642,20 @@ def __init__(
                 quant_config=quant_config,
                 prefix=f"{prefix}.router",
             )
-            self.moe = Gemma4MoE(
-                config,
+            names = ("moe_intermediate_size", "expert_intermediate_size")
+            intermediate_size = getattr_iter(config, names)
+            if intermediate_size is None:
+                raise ValueError("Gemma4 MoE requires an expert intermediate size")
+            self.experts = FusedMoEFactory(
+                num_experts=config.num_experts,
+                top_k=config.top_k_experts,
+                hidden_size=config.hidden_size,
+                intermediate_size=intermediate_size,
+                renormalize=True,
                 quant_config=quant_config,
-                prefix=f"{prefix}.moe",
+                prefix=f"{prefix}.experts",
+                custom_routing_function=self.router.routing_function,
+                activation="gelu_tanh",
             )
             self.post_feedforward_layernorm_1 = RMSNorm(
                 config.hidden_size, eps=config.rms_norm_eps
@@ -671,7 +668,7 @@ def __init__(
             )
         else:
             self.router = None
-            self.moe = None
+            self.experts = None
             self.post_feedforward_layernorm_1 = None
             self.post_feedforward_layernorm_2 = None
             self.pre_feedforward_layernorm_2 = None
@@ -747,13 +744,13 @@ def forward(
             assert self.post_feedforward_layernorm_1 is not None
             assert self.pre_feedforward_layernorm_2 is not None
             assert self.router is not None
-            assert self.moe is not None
+            assert self.experts is not None
             assert self.post_feedforward_layernorm_2 is not None
             hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states)
 
             hidden_states_2 = self.pre_feedforward_layernorm_2(residual)
             router_logits = self.router(residual)
-            hidden_states_2 = self.moe(hidden_states_2, router_logits)
+            hidden_states_2 = self.experts(hidden_states_2, router_logits)
             hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2)
 
             # Combine MLP and MoE outputs
@@ -866,6 +863,7 @@ def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:
         Returns:
             Per-layer embeddings (num_tokens, num_layers,
             hidden_size_per_layer_input)
+
         """
         if self.embed_tokens_per_layer is None:
             return None
@@ -982,7 +980,7 @@ def forward(
 @support_torch_compile(
     enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill
 )
-class Gemma4Model(nn.Module, EagleModelMixin):
+class Gemma4Model(nn.Module, EagleModelMixin, SupportsQuant):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         config = _get_text_config(vllm_config.model_config.hf_config)
@@ -990,6 +988,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         quant_config = vllm_config.quant_config
         self.config = config
         self.quant_config = quant_config
+        self.hf_to_vllm_mapper = _gemma4_layer_weights_mapper(config)
 
         # PLE config values (default to 0 if not present — disables PLE)
         self.hidden_size_per_layer_input = getattr(
@@ -1199,6 +1198,7 @@ def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:
         Returns:
             Per-layer embeddings (num_tokens, num_layers,
             hidden_size_per_layer_input)
+
         """
         return self.self_decoder.get_per_layer_inputs(input_ids)
 
@@ -1389,152 +1389,21 @@ def forward(
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        # MoE expert weight mapping: checkpoint can have either:
-        #   1. 3D packed tensors (exploded in _weight_iterator to per-expert 2D)
-        #   2. Already per-expert 2D weights (if quantized)
-        # Map to MoERunner parameters:
-        #   moe.experts.{id}.gate_proj → MoERunner w1 (shard of w13)
-        #   moe.experts.{id}.up_proj   → MoERunner w3 (shard of w13)
-        #   moe.experts.{id}.down_proj → MoERunner w2
-        num_experts = getattr(self.config, "num_experts", None) or 0
-        # Strategy A: dot-separated suffix
-        # (standard AWQ/GPTQ e.g. .qweight, .scales, .weight)
-        dot_suffix_expert_params_mapping = fused_moe_make_expert_params_mapping(
-            self,
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=num_experts,
-        )
-        # Strategy B: underscore-separated suffix
-        # (CompressedTensors-format AWQ/W4A16 _packed, _scale)
-        underscore_suffix_expert_params_mapping = [
-            (
-                f"{param_name}weight_",
-                f"{weight_name.rstrip('.')}_",
-                expert_id,
-                shard_id,
-            )
-            for (
-                param_name,
-                weight_name,
-                expert_id,
-                shard_id,
-            ) in dot_suffix_expert_params_mapping
-        ]
-        expert_params_mapping = (
-            dot_suffix_expert_params_mapping + underscore_suffix_expert_params_mapping
-        )
-        params_dict = dict(self.named_parameters())
-        # Include buffers (e.g. layer_scalar) so they can be loaded too
-        params_dict.update(dict(self.named_buffers()))
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
-                remapped_name = maybe_remap_kv_scale_name(name, params_dict)
-                if remapped_name is not None and remapped_name in params_dict:
-                    param = params_dict[remapped_name]
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight)
-                    loaded_params.add(remapped_name)
-                    continue
-
-            for param_name, shard_name, shard_id in stacked_params_mapping:
-                if shard_name not in name:
-                    continue
-                stacked_name = name.replace(shard_name, param_name)
-                # KV-shared layers have a plain q_proj instead of a packed
-                # qkv_proj: fall through to the direct load, which also
-                # skips the redundant K/V tensors original checkpoints ship.
-                if stacked_name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(stacked_name, self):
-                    continue
-                param = params_dict[stacked_name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                loaded_params.add(stacked_name)
-                break
-            else:
-                for (
-                    param_name,
-                    weight_name,
-                    expert_id,
-                    shard_id,
-                ) in expert_params_mapping:
-                    # Match both:
-                    #  - Bare weights: "experts.0.down_proj" (from 3D explosion)
-                    #  - With suffix: "experts.0.down_proj.weight_scale" (2D quantized)
-                    # weight_name has trailing dot, so check with and without it
-                    weight_name_base = weight_name.rstrip(".")
-                    if weight_name in name:
-                        # Has suffix (e.g., .weight_scale)
-                        moe_name = name.replace(weight_name, param_name)
-                    elif name.endswith(weight_name_base):
-                        # Bare weight (no suffix)
-                        moe_name = name.replace(
-                            weight_name_base, param_name.rstrip("_") + "_weight"
-                        )
-                    else:
-                        continue
-                    if moe_name not in params_dict:
-                        continue
-                    if is_pp_missing_parameter(moe_name, self):
-                        continue
-                    param = params_dict[moe_name]
-                    # Expert weights are already in the correct
-                    # orientation for MoERunner after _weight_iterator:
-                    #   gate/up: [I, H] → w1/w3 expects [I, H]
-                    #   down:    [H, I] → w2 expects [H, I]
-                    # Scales and other quantization params may be 1D or scalar.
-                    weight_loader = param.weight_loader
-                    weight_loader(
-                        param,
-                        loaded_weight,
-                        moe_name,  # Pass mapped name (handles both weights and scales)
-                        shard_id=shard_id,
-                        expert_id=expert_id,
-                    )
-                    loaded_params.add(moe_name)
-                    break
-                else:
-                    if name.endswith(".bias") and name not in params_dict:
-                        continue
-                    remapped_name = maybe_remap_kv_scale_name(name, params_dict)
-                    if remapped_name is None:
-                        continue
-                    name = remapped_name
-                    if is_pp_missing_parameter(name, self):
-                        continue
-                    # Skip if name doesn't exist in params_dict (e.g., individual
-                    # expert weights that should have been handled above)
-                    if name not in params_dict:
-                        continue
-                    param = params_dict[name]
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 class Gemma4ForCausalLM(
     nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
 ):
-    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={
+            # Multimodal weights are handled by the multimodal wrapper.
+            "audio_tower.": None,
+            "vision_tower.": None,
+            "embed_audio.": None,
+            "embed_vision.": None,
+        },
         orig_to_new_prefix={
             # Gemma4ForConditionalGeneration already loads the text stack
             # from `model.language_model.*`. We reuse that same checkpoint
@@ -1542,13 +1411,6 @@ class Gemma4ForCausalLM(
             # so LoRA keys from the conditional wrapper map onto `model.*`.
             "model.language_model.": "model.",
         },
-        orig_to_new_substr={
-            # Gemma4ForConditionalGeneration names MoE adapter targets under
-            # `...moe.experts.*`, while the text-only model exposes them
-            # under `...moe.*`.
-            ".moe.experts.gate_up_proj": ".moe.gate_up_proj",
-            ".moe.experts.down_proj": ".moe.down_proj",
-        },
     )
     # KV-shared layers only have q_proj, so qkv_proj packing applies to
     # the other layers.
@@ -1594,26 +1456,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         )
 
         # --- MixtureOfExperts protocol ---
-        self.moe_layers: list[nn.Module] = []
-        example_moe: Gemma4MoE | None = None
-
-        for layer in self.model.layers:
-            if hasattr(layer, "moe") and isinstance(layer.moe, Gemma4MoE):
-                example_moe = layer.moe
-                self.moe_layers.append(layer.moe.experts)
-
+        self.moe_layers: list[nn.Module] = [
+            layer.experts
+            for layer in self.model.layers
+            if not isinstance(layer, PPMissingLayer) and layer.experts is not None
+        ]
         self.num_moe_layers = len(self.moe_layers)
 
-        if example_moe is not None:
-            self.num_logical_experts = example_moe.num_experts
-            self.num_physical_experts = example_moe.num_experts
-            self.num_local_physical_experts = example_moe.num_experts
-            self.num_routed_experts = example_moe.num_experts
-        else:
-            self.num_logical_experts = 0
-            self.num_physical_experts = 0
-            self.num_local_physical_experts = 0
-            self.num_routed_experts = 0
+        num_experts = config.num_experts if self.num_moe_layers else 0
+        self.num_logical_experts = num_experts
+        self.num_physical_experts = num_experts
+        self.num_local_physical_experts = num_experts
+        self.num_routed_experts = num_experts
 
         self.num_expert_groups = 1
         self.num_shared_experts = 0
@@ -1642,106 +1496,5 @@ def compute_logits(
         return self.logits_processor(self.lm_head, hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # Checkpoint weight names use "language_model." prefix (from the
-        # Gemma4ForConditionalGeneration wrapper). Strip it to map to our
-        # model tree which is just "model.*".
-        def _weight_iterator():
-            use_k_eq_v = getattr(self.config, "attention_k_eq_v", False)
-            # Build set of k_eq_v layer indices (full_attention layers
-            # when attention_k_eq_v is enabled). These layers have k_proj
-            # but no v_proj in checkpoint — we duplicate k_proj as v_proj.
-            k_eq_v_layer_indices: set[int] = set()
-            if use_k_eq_v:
-                for idx, lt in enumerate(self.config.layer_types):
-                    if lt == "full_attention":
-                        k_eq_v_layer_indices.add(idx)
-
-            for name, weight in weights:
-                # Remap "language_model." → "" to match our model tree.
-                # Checkpoint: model.language_model.layers.X.*
-                # Our model:  model.layers.X.*
-                name = name.replace("language_model.", "")
-
-                # Remap new HF checkpoint naming to internal vLLM
-                # naming: HF moved per_expert_scale to router and
-                # renamed moe → experts in the MoE block.
-                name = name.replace(
-                    ".router.per_expert_scale",
-                    ".moe.per_expert_scale",
-                )
-                if ".experts.gate_up_proj" in name:
-                    name = name.replace(
-                        ".experts.gate_up_proj",
-                        ".moe.gate_up_proj",
-                    )
-                elif ".experts.down_proj" in name:
-                    name = name.replace(
-                        ".experts.down_proj",
-                        ".moe.down_proj",
-                    )
-
-                # Remap individual 2D expert weights:
-                # .experts.{id}.{proj} → .moe.experts.{id}.{proj}
-                # (This handles per-expert 2D quantized weights)
-                name = _remap_gemma4_expert_weight_name(name)
-
-                # MoE expert weights: checkpoint stores as 3D packed
-                # tensors.  Explode into per-expert 2D weights for
-                # MoERunner weight_loader.
-                #
-                # Checkpoint format:
-                #   moe.gate_up_proj: [E, 2*I, H]  (fused gate + up)
-                #   moe.down_proj:    [E, H, I]
-                #
-                # MoERunner expects per-expert:
-                #   w1 (gate): [I, H]   — first half of gate_up
-                #   w3 (up):   [I, H]   — second half of gate_up
-                #   w2 (down): [H, I]   — as-is from checkpoint
-                #
-                # No transpose needed: checkpoint orientation already
-                # matches MoERunner's expected layout.
-                if "moe.gate_up_proj" in name and weight.dim() == 3:
-                    num_experts = weight.size(0)
-                    intermediate_size = weight.size(1) // 2
-                    for expert_id in range(num_experts):
-                        gate_weight = weight[expert_id, :intermediate_size, :]
-                        up_weight = weight[expert_id, intermediate_size:, :]
-                        base = name.replace("moe.", f"moe.experts.{expert_id}.")
-                        yield base.replace("gate_up_proj", "gate_proj"), gate_weight
-                        yield base.replace("gate_up_proj", "up_proj"), up_weight
-                    continue
-
-                if "moe.down_proj" in name and weight.dim() == 3:
-                    num_experts = weight.size(0)
-                    for expert_id in range(num_experts):
-                        expert_name = name.replace("moe.", f"moe.experts.{expert_id}.")
-                        yield expert_name, weight[expert_id]
-                    continue
-
-                # k_eq_v layers: checkpoint has k_proj but no v_proj.
-                # QKVParallelLinear expects both, so duplicate k_proj
-                # as v_proj so V gets identical weights to K.
-                # ONLY for full_attention layers — sliding layers have
-                # their own real v_proj weights.
-                if "self_attn.k_proj" in name and k_eq_v_layer_indices:
-                    m = re.search(r"layers\.(\d+)\.", name)
-                    if m and int(m.group(1)) in k_eq_v_layer_indices:
-                        yield name, weight
-                        yield name.replace("k_proj", "v_proj"), weight.clone()
-                        continue
-
-                yield name, weight
-
-        # Drop multimodal weights, which are handled by the multimodal wrapper.
-        # `_weight_iterator` already applies this model's renames by hand, so
-        # `hf_to_vllm_mapper` is deliberately not passed here.
-        mapper = WeightsMapper(
-            orig_to_new_substr={
-                "audio_tower.": None,
-                "vision_tower.": None,
-                "embed_audio.": None,
-                "embed_vision.": None,
-            }
-        )
         loader = AutoWeightsLoader(self)
-        return loader.load_weights(_weight_iterator(), mapper=mapper)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py
index 326d179cc237..796dfbb6e8fa 100644
--- a/vllm/model_executor/models/gemma4_mm.py
+++ b/vllm/model_executor/models/gemma4_mm.py
@@ -24,6 +24,7 @@
 from torch import nn
 from transformers import AutoModel, BatchFeature
 from transformers.models.gemma4 import (
+    Gemma4AudioFeatureExtractor,
     Gemma4Config,
     Gemma4Processor,
     Gemma4VisionConfig,
@@ -35,20 +36,13 @@
 
 from vllm.config import VllmConfig
 from vllm.config.model import get_served_model_name
-from vllm.config.multimodal import (
-    AudioDummyOptions,
-    BaseDummyOptions,
-    VideoDummyOptions,
-)
+from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.models.gemma3n_mm import batch_audio_features
-from vllm.model_executor.models.gemma4 import (
-    _GEMMA4_EXPERT_PARENT_MAPPER,
-    Gemma4ForCausalLM,
-)
+from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.transformers.utils import recursive_replace_linear
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -131,8 +125,7 @@ def _get_max_soft_tokens(
 
 
 class Gemma4ImagePixelInputs(TensorSchema):
-    """
-    Pre-patchified image inputs from the Gemma4 image processor.
+    """Pre-patchified image inputs from the Gemma4 image processor.
 
     Dimensions:
         - bn: Batch size * number of images
@@ -158,11 +151,10 @@ class Gemma4ImagePixelInputs(TensorSchema):
 
 
 class Gemma4AudioInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of audios
-        - s: Sequence length (MEL spectrogram frames)
-        - f: Number of features (MEL bins)
+    """Dimensions:
+    - bn: Batch size * number of audios
+    - s: Sequence length (MEL spectrogram frames)
+    - f: Number of features (MEL bins)
     """
 
     type: Literal["audio"] = "audio"
@@ -277,10 +269,7 @@ def get_mm_max_tokens_per_item(
         num_frames = _VIDEO_MAX_FRAMES
         mm_config = self.ctx.model_config.get_multimodal_config()
         video_opts = mm_config.limit_per_prompt.get("video")
-        if (
-            isinstance(video_opts, VideoDummyOptions)
-            and video_opts.num_frames is not None
-        ):
+        if video_opts is not None and video_opts.num_frames is not None:
             num_frames = min(num_frames, video_opts.num_frames)
         tokens["video"] = num_frames * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6)
         return tokens
@@ -303,9 +292,12 @@ def _compute_num_soft_tokens(
         for an image of the given dimensions, after padding is stripped.
 
         Args:
+            image_width: Width of the image in pixels.
+            image_height: Height of the image in pixels.
             max_soft_tokens: Override for the vision config's
                 ``default_output_length``.  When *None*, the value from
                 the model config is used.
+
         """
         vision_cfg = self.get_hf_config().vision_config
         patch_size = vision_cfg.patch_size
@@ -344,8 +336,12 @@ def get_image_repl(
         produce after stripping padding.
 
         Args:
+            image_width: Width of the image in pixels.
+            image_height: Height of the image in pixels.
+            processor: The Gemma4 processor, when already constructed.
             max_soft_tokens: Override for the default token budget.
                 When *None*, falls back to the model config value.
+
         """
         if processor is None:
             processor = self.get_hf_processor()
@@ -456,6 +452,28 @@ def get_video_repl(
 # ---------------------------------------------------------------------------
 
 
+def _dummy_audio_num_samples(processor: Any) -> int:
+    """Length in samples of a worst-case dummy audio item.
+
+    ``fft_length`` is set in ``Gemma4AudioFeatureExtractor.__init__`` and is not
+    defined by ``Gemma4UnifiedAudioFeatureExtractor``, which the encoder-free
+    Gemma 4 variant uses. Reading it unconditionally therefore raises
+    ``AttributeError`` during memory profiling on that variant, which is reached
+    whenever audio is the modality with the largest per-item token count.
+
+    For any other extractor, fall back to the audio budget the processor itself
+    declares -- the same pair of values used to warn about over-long audio
+    elsewhere in this file.
+    """
+    feature_extractor = processor.feature_extractor
+    if isinstance(feature_extractor, Gemma4AudioFeatureExtractor):
+        return feature_extractor.fft_length
+
+    sampling_rate = feature_extractor.sampling_rate
+    max_duration_s = processor.audio_seq_length * processor.audio_ms_per_token / 1000.0
+    return int(max_duration_s * sampling_rate)
+
+
 class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]):
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_images = mm_counts.get("image", 0)
@@ -477,9 +495,8 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+        mm_options: MultiModalDummyOptions | None = None,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         num_audios = mm_counts.get("audio", 0)
         num_videos = mm_counts.get("video", 0)
         processor = self.info.get_hf_processor()
@@ -492,27 +509,21 @@ def get_dummy_mm_data(
         img_width = size.get("width", 224)
         img_height = size.get("height", 224)
 
-        image_overrides = mm_options.get("image") if mm_options else None
-        audio_overrides = mm_options.get("audio") if mm_options else None
-        video_overrides = mm_options.get("video") if mm_options else None
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
-        assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
-
         data: dict[str, Any] = {
             "image": self._get_dummy_images(
                 width=img_width,
                 height=img_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image") if mm_options else None,
             ),
         }
 
         if num_audios > 0:
-            audio_len = processor.feature_extractor.fft_length
+            audio_len = _dummy_audio_num_samples(processor)
             data["audio"] = self._get_dummy_audios(
                 length=audio_len,
                 num_audios=num_audios,
-                overrides=audio_overrides,
+                overrides=mm_options.get("audio") if mm_options else None,
             )
 
         if num_videos > 0:
@@ -521,7 +532,7 @@ def get_dummy_mm_data(
                 height=img_height,
                 num_frames=_VIDEO_MAX_FRAMES,
                 num_videos=num_videos,
-                overrides=video_overrides,
+                overrides=mm_options.get("video") if mm_options else None,
             )
 
         return data
@@ -567,22 +578,25 @@ def _get_dummy_videos(
 
 
 class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        if not hf_data:
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        prompt_text = hf_data.pop("text")
+        assert isinstance(prompt_text, str)
 
-        merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs)
+        merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_kwargs)
         val, is_top_level_max_soft_tokens = _get_max_soft_tokens(merged_kwargs)
 
         if val is not None and val not in _SUPPORTED_SOFT_TOKENS:
@@ -591,14 +605,12 @@ def _apply_hf_processor_main(
                 f"Valid values are {_SUPPORTED_SOFT_TOKENS}."
             )
 
-        mm_data = dict(mm_data)
-
         # ---- VIDEO HANDLING ----
         # Gemma4 decomposes video into timestamped image frames.
         # Each frame is processed with max_soft_tokens=70 through the
         # same vision tower, matching transformers processing_gemma4.py.
         video_outputs: dict[str, Any] = {}
-        if videos := mm_data.pop("videos", []):
+        if videos := hf_data.pop("videos", []):
             assert isinstance(videos, list)
             processor = self.info.get_hf_processor()
 
@@ -628,7 +640,10 @@ def _apply_hf_processor_main(
                 timestamps = [idx / fps for idx in frame_indices]
 
                 # Process frames as images with max_soft_tokens=70
-                video_mm_kwargs = dict(hf_processor_mm_kwargs)
+                video_mm_kwargs = dict(hf_kwargs)
+                # Override configured image options too: these inputs are video
+                # frames, whose cache keys exclude images_kwargs.
+                video_mm_kwargs["images_kwargs"] = {}
                 video_mm_kwargs["max_soft_tokens"] = _VIDEO_MAX_SOFT_TOKENS
 
                 dummy_prompt = ("\t" + processor.image_token) * len(frames)
@@ -695,18 +710,14 @@ def _apply_hf_processor_main(
                 "video_timestamps": video_timestamps_per_video,
             }
 
-        # The processor accepts 'audio' not 'audios'.
-        if "audios" in mm_data:
-            mm_data["audio"] = mm_data.pop("audios")
-
         # Warn if any audio waveform exceeds the model's max duration.
-        if "audio" in mm_data:
+        if "audio" in hf_data:
             processor = self.info.get_hf_processor()
             sr = processor.feature_extractor.sampling_rate
             max_tokens = processor.audio_seq_length
             ms_per_tok = processor.audio_ms_per_token
             max_duration_s = max_tokens * ms_per_tok / 1000.0
-            audios = mm_data["audio"]
+            audios = hf_data["audio"]
             if not isinstance(audios, (list, tuple)):
                 audios = [audios]
             for i, waveform in enumerate(audios):
@@ -729,13 +740,13 @@ def _apply_hf_processor_main(
         # NOTE: This requires a corresponding type annotation on the
         # HF side (Gemma4ProcessorKwargs.images_kwargs) so that
         # _merge_kwargs routes max_soft_tokens into images_kwargs.
-        patched_mm_kwargs = dict(hf_processor_mm_kwargs)
+        patched_mm_kwargs = dict(hf_kwargs)
         if val is not None and is_top_level_max_soft_tokens:
             patched_mm_kwargs["max_soft_tokens"] = val
 
         processed_data = self.info.ctx.call_hf_processor(
             self.info.get_hf_processor(**patched_mm_kwargs),
-            dict(text=prompt_text, **mm_data),
+            dict(text=prompt_text, **hf_data),
             patched_mm_kwargs,
         )
 
@@ -766,9 +777,9 @@ def _apply_hf_processor_main(
             processed_data["input_features_mask"] = unpadded_masks
 
         processed_data.update(video_outputs)
-        processed_data.update(passthrough_data)
-
-        return processed_data
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_mm_fields_config(
         self,
@@ -1032,7 +1043,7 @@ class Gemma4ForConditionalGeneration(
     }
 
     # Maps checkpoint prefixes to vLLM module paths.
-    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
+    hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             # vision tower
             "model.vision_tower": "vision_tower",
diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py
index 5e25a80d6bbd..01f226ff07a5 100644
--- a/vllm/model_executor/models/gemma4_mtp.py
+++ b/vllm/model_executor/models/gemma4_mtp.py
@@ -101,6 +101,7 @@ def _select_and_score(
         Returns:
             logits: (num_tokens, num_selected) sparse logits.
             indices: (num_tokens, num_selected) corresponding vocab indices.
+
         """
         num_tokens = hidden_states.shape[0]
         _, top_k_indices = torch.topk(
@@ -246,16 +247,8 @@ def forward(
 
         q, _ = self.rotary_emb(positions, q, None)
 
-        # Attention reads K/V from the target's cache via KV sharing;
-        # these dummy tensors are never consumed but required by the API.
-        num_tokens = q.shape[0]
-        kv_dummy = torch.empty(
-            num_tokens,
-            self.num_kv_heads * self.head_dim,
-            dtype=q.dtype,
-            device=q.device,
-        )
-        attn_output = self.attn(q, kv_dummy, kv_dummy)
+        # Attention reads K/V from the target's cache via KV sharing.
+        attn_output = self.attn(q, None, None)
         output, _ = self.o_proj(attn_output)
         return output
 
diff --git a/vllm/model_executor/models/gemma4_unified.py b/vllm/model_executor/models/gemma4_unified.py
index 7c9aaa8499ec..1a2624651528 100644
--- a/vllm/model_executor/models/gemma4_unified.py
+++ b/vllm/model_executor/models/gemma4_unified.py
@@ -28,7 +28,6 @@
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import VideoDummyOptions
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.models.gemma3n_mm import batch_audio_features
 from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
@@ -177,10 +176,7 @@ def get_mm_max_tokens_per_item(
         num_frames = _VIDEO_MAX_FRAMES
         mm_config = self.ctx.model_config.get_multimodal_config()
         video_opts = mm_config.limit_per_prompt.get("video")
-        if (
-            isinstance(video_opts, VideoDummyOptions)
-            and video_opts.num_frames is not None
-        ):
+        if video_opts is not None and video_opts.num_frames is not None:
             num_frames = min(num_frames, video_opts.num_frames)
         tokens["video"] = num_frames * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6)
         return tokens
diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py
index 163115ff3d5e..95cdcc5a2063 100644
--- a/vllm/model_executor/models/glm4_1v.py
+++ b/vllm/model_executor/models/glm4_1v.py
@@ -36,9 +36,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import transformers
 from einops import rearrange
-from packaging.version import Version
 from transformers import BatchFeature, Glm4vProcessor
 from transformers.image_processing_base import ImageProcessingMixin
 from transformers.models.glm4v.configuration_glm4v import (
@@ -54,7 +52,7 @@
 from transformers.video_utils import VideoMetadata
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions
 from vllm.distributed import get_tensor_model_parallel_world_size, parallel_state
 from vllm.distributed import utils as dist_utils
 from vllm.inputs import MultiModalDataDict
@@ -83,6 +81,7 @@
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
     VideoItem,
 )
@@ -131,8 +130,6 @@
 # For profile run
 _MAX_FRAMES_PER_VIDEO = 600
 
-TRANSFORMERS_WITH_GA = Version(transformers.__version__) >= Version("5.10.0.dev0")
-
 
 def _to_video_metadata(metadata: Mapping[str, Any]) -> VideoMetadata:
     return VideoMetadata(
@@ -144,12 +141,11 @@ def _to_video_metadata(metadata: Mapping[str, Any]) -> VideoMetadata:
 
 
 class Glm4vImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - np: Number of patches
-        - cpp: Number of channels * patch_size * patch_size
-        - ni: Number of images
-        - g: Grid dimensions (3 for grid_t, grid_h, grid_w)
+    """Dimensions:
+    - np: Number of patches
+    - cpp: Number of channels * patch_size * patch_size
+    - ni: Number of images
+    - g: Grid dimensions (3 for grid_t, grid_h, grid_w)
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -159,12 +155,11 @@ class Glm4vImagePixelInputs(TensorSchema):
 
 
 class Glm4vImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - f: Number of image features (varies based on image resolution)
-        - h: Hidden size (must match language model backbone)
-        - n: Number of images
-        - g: Grid dimensions (3 for grid_t, grid_h, grid_w)
+    """Dimensions:
+    - f: Number of image features (varies based on image resolution)
+    - h: Hidden size (must match language model backbone)
+    - n: Number of images
+    - g: Grid dimensions (3 for grid_t, grid_h, grid_w)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -177,14 +172,13 @@ class Glm4vImageEmbeddingInputs(TensorSchema):
 
 
 class Glm4vVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - np: Number of patches
-        - ctpp: Number of channels * temporal_patch_size *
-            patch_size * patch_size
-        - f: Number of frames
-        - g: Grid dimensions (3 for grid_t which is usually 1 for processed
-          video, grid_h, grid_w)
+    """Dimensions:
+    - np: Number of patches
+    - ctpp: Number of channels * temporal_patch_size *
+        patch_size * patch_size
+    - f: Number of frames
+    - g: Grid dimensions (3 for grid_t which is usually 1 for processed
+      video, grid_h, grid_w)
     """
 
     type: Literal["pixel_values_videos"] = "pixel_values_videos"
@@ -194,13 +188,12 @@ class Glm4vVideoPixelInputs(TensorSchema):
 
 
 class Glm4vVideoEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - p: Number of video patches across all frames
-        - h: Hidden size (must match language model backbone)
-        - f: Number of frames
-        - g: Grid dimensions (3 for grid_t which is usually 1 for processed
-          video, grid_h, grid_w)
+    """Dimensions:
+    - p: Number of video patches across all frames
+    - h: Hidden size (must match language model backbone)
+    - f: Number of frames
+    - g: Grid dimensions (3 for grid_t which is usually 1 for processed
+      video, grid_h, grid_w)
     """
 
     type: Literal["video_embeds"] = "video_embeds"
@@ -770,7 +763,6 @@ def pos_embeds_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:
         forward graph compatible with CUDA graph replay. Coordinate generation matches
         `rot_pos_emb` exactly to guarantee spatial alignment.
         """
-
         device = self.embeddings.position_embedding.weight.device
         dtype = self.dtype
         all_embeds = []
@@ -849,6 +841,7 @@ def prepare_encoder_metadata(
                 instead of computing from cu_seqlens (needed for CUDA
                 graph capture to cover worst-case replay scenarios).
             device: Device to place tensors on. Defaults to self.device.
+
         """
         if device is None:
             device = self.device
@@ -1392,17 +1385,6 @@ def _get_video_second_idx_glmga(
         timestamps_list = full_second_idxs[::2]
         return list(timestamps_list)
 
-    def _get_video_frame_embed_token_id(self, hf_processor: object) -> int:
-        attr = (
-            "image_token_id"
-            if isinstance(hf_processor, Glm4vProcessor) or TRANSFORMERS_WITH_GA
-            else "video_token_id"
-        )
-        token_id = getattr(hf_processor, attr, None)
-        if not isinstance(token_id, int):
-            raise ValueError(f"Processor has no valid {attr}")
-        return token_id
-
     def _construct_video_placeholder(
         self,
         video_array: np.ndarray,
@@ -1440,7 +1422,7 @@ def _construct_video_placeholder(
         num_tokens_per_frame = int(H * W) // merge_length
         placeholder = []
         placeholder.append(bov_token_id)
-        frame_embed_token_id = self._get_video_frame_embed_token_id(hf_processor)
+        frame_embed_token_id = hf_processor.image_token_id
         for frame_idx in frames_idx_token:
             placeholder.append(boi_token_id)
             placeholder.extend([frame_embed_token_id] * num_tokens_per_frame)
@@ -1473,33 +1455,26 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(
             seq_len, mm_counts
         )
 
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-        assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -1584,52 +1559,54 @@ def _get_direct_path_inputs(
 
         return prepared_data, prepared_kwargs
 
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        if not hf_data:
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
-        mm_data = dict(mm_data)
+        prompt_text = hf_data.pop("text")
+        assert isinstance(prompt_text, str)
 
-        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        processor = self.info.get_hf_processor(**hf_kwargs)
 
-        use_direct_path = (
-            not isinstance(processor, Glm4vProcessor) and TRANSFORMERS_WITH_GA
-        )
-        if use_direct_path:
+        if not isinstance(processor, Glm4vProcessor):
             prepared_data, prepared_kwargs = self._get_direct_path_inputs(
-                mm_data, hf_processor_mm_kwargs
+                hf_data, hf_kwargs
             )
             processed_data = self.info.ctx.call_hf_processor(
                 self.info.get_hf_processor(**prepared_kwargs),
                 dict(text=prompt_text, **prepared_data),
                 prepared_kwargs,
             )
-            processed_data.update(passthrough_data)
-            return processed_data
+            return self._finalize_hf_mm_data(
+                hf_data, hf_kwargs, passthrough_data, processed_data
+            )
 
-        videos = mm_data.get("videos")
-        if isinstance(videos, list) and videos:
+        if (
+            "videos" in hf_data
+            and isinstance(videos := hf_data["videos"], list)
+            and len(videos) > 0
+        ):
+            # Videos are processed item-wise below; remove them from hf_data so
+            # the combined call at the end does not re-sample the raw frames.
+            del hf_data["videos"]
             video_grid_thw_lst = []
             pixel_values_videos_lst = []
-            frame_embed_token_id = self.info._get_video_frame_embed_token_id(processor)
-            swap_video_frame_tokens = frame_embed_token_id == processor.image_token_id
-            mm_data.pop("videos")
             for item in videos:
                 video_array, metadata = item
 
                 # don't update mm_kwargs inplace
-                video_mm_kwargs = dict(**hf_processor_mm_kwargs)
+                video_mm_kwargs = dict(**hf_kwargs)
                 video_mm_kwargs["do_sample_frames"] = metadata.get(
                     "do_sample_frames", True
                 )
@@ -1647,10 +1624,10 @@ def _apply_hf_processor_main(
                     video_mm_kwargs,
                 )
                 input_ids = video_outputs.pop("input_ids")
-                if swap_video_frame_tokens:
-                    input_ids[input_ids == processor.image_token_id] = (
-                        processor.video_token_id
-                    )
+
+                input_ids[input_ids == processor.image_token_id] = (
+                    processor.video_token_id
+                )
                 video_placeholder = processor.tokenizer.batch_decode(input_ids)[0]
                 prompt_text = prompt_text.replace(
                     "<|begin_of_video|><|video|><|end_of_video|>",
@@ -1666,22 +1643,20 @@ def _apply_hf_processor_main(
             )
         else:
             video_outputs = dict()
-            swap_video_frame_tokens = False
 
         processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            dict(text=prompt_text, **mm_data),
-            hf_processor_mm_kwargs,
+            self.info.get_hf_processor(**hf_kwargs),
+            dict(text=prompt_text, **hf_data),
+            hf_kwargs,
         )
-        if swap_video_frame_tokens:
-            input_ids = processed_data["input_ids"]
-            input_ids[input_ids == processor.video_token_id] = processor.image_token_id
-            processed_data["input_ids"] = input_ids
+        input_ids = processed_data["input_ids"]
+        input_ids[input_ids == processor.video_token_id] = processor.image_token_id
+        processed_data["input_ids"] = input_ids
 
         processed_data.update(video_outputs)
-        processed_data.update(passthrough_data)
-
-        return processed_data
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_mm_fields_config(
         self,
@@ -1724,7 +1699,7 @@ def get_video_replacement(item_idx: int):
             )
             return PromptUpdateDetails.select_token_id(
                 placeholder,
-                embed_token_id=self.info._get_video_frame_embed_token_id(hf_processor),
+                embed_token_id=hf_processor.image_token_id,
             )
 
         return [
@@ -2346,6 +2321,7 @@ def forward(
                 parallelism.
             inputs_embeds: Optional pre-computed input embeddings.
             **kwargs: Additional keyword arguments.
+
         """
         if intermediate_tensors is not None:
             inputs_embeds = None
@@ -2369,28 +2345,23 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model.model",
             connector="visual.merger.",
             tower_model="visual.",
         )
 
-    def get_num_mm_encoder_tokens(
-        self,
-        num_image_tokens: int,
-    ) -> int:
-        merge_size = self.config.vision_config.spatial_merge_size
-        return num_image_tokens * (merge_size**2)
-
-    def get_num_mm_connector_tokens(
+    def get_mm_lora_token_counts(
         self,
-        num_vision_tokens: int,
-    ) -> int:
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         merge_size = self.config.vision_config.spatial_merge_size
-        return num_vision_tokens // (merge_size**2)
+        return num_mm_embeds * (merge_size**2), num_mm_embeds
 
 
 @MULTIMODAL_REGISTRY.register_processor(
diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py
index 3c72a148d5ef..5e5974e209a3 100644
--- a/vllm/model_executor/models/glm4_moe.py
+++ b/vllm/model_executor/models/glm4_moe.py
@@ -43,6 +43,7 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
     is_model_fused_shared_expert_compatible,
@@ -137,14 +138,12 @@ def __init__(
                 f"Unsupported activation: {config.hidden_act}. "
                 "Only silu is supported for now."
             )
-        # NOTE In the transformers implementation, the gate isn't an nn.Linear,
-        # so we cannot use ReplicatedLinear here.
-        # See: https://github.com/huggingface/transformers/blob/v4.55.1/src/transformers/models/glm4_moe/modeling_glm4_moe.py#L260
-        self.gate = nn.Linear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.n_routed_experts,
-            bias=False,
-            dtype=torch.float32,
+            out_dtype=torch.float32,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.gate",
         )
         self.gate.e_score_correction_bias = nn.Parameter(
             torch.empty(config.n_routed_experts, dtype=torch.float32)
@@ -212,7 +211,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = hidden_states.view(-1, hidden_dim)
 
         # router_logits: (num_tokens, n_experts)
-        router_logits = self.gate(hidden_states.to(dtype=torch.float32))
+        router_logits, _ = self.gate(hidden_states)
 
         final_hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py
index d0119719af0e..019380540d08 100644
--- a/vllm/model_executor/models/glm4_moe_lite.py
+++ b/vllm/model_executor/models/glm4_moe_lite.py
@@ -67,6 +67,10 @@
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backends.mla.index_group import (
+    SparseMLAIndexGroupBuilder,
+    get_sparse_mla_index_group_max_rows,
+)
 
 from .interfaces import SupportsLoRA, SupportsPP
 from .utils import (
@@ -108,6 +112,7 @@ def __init__(
         prefix: str,
         config: "Glm4MoeLiteConfig | None" = None,
         topk_indices_buffer: torch.Tensor | None = None,
+        index_group_builder: SparseMLAIndexGroupBuilder | None = None,
     ) -> None:
         super().__init__()
 
@@ -136,6 +141,11 @@ def __init__(
             attn_cls = Glm4MoeLiteMLAAttention
         else:
             attn_cls = Glm4MoeLiteAttention
+        attn_kwargs: dict[str, typing.Any] = (
+            {"index_group_builder": index_group_builder}
+            if attn_cls is Glm4MoeLiteMLAAttention
+            else {}
+        )
 
         self.self_attn = attn_cls(
             vllm_config=vllm_config,
@@ -152,6 +162,7 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
             topk_indices_buffer=topk_indices_buffer,
+            **attn_kwargs,
         )
 
         if (
@@ -234,6 +245,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             )
         else:
             topk_indices_buffer = None
+        index_group_builder = (
+            SparseMLAIndexGroupBuilder(
+                topk_indices_buffer,
+                get_sparse_mla_index_group_max_rows(vllm_config),
+            )
+            if topk_indices_buffer is not None
+            else None
+        )
 
         if get_pp_group().is_first_rank:
             self.embed_tokens = VocabParallelEmbedding(
@@ -252,6 +271,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 config=config,
                 prefix=prefix,
                 topk_indices_buffer=topk_indices_buffer,
+                index_group_builder=index_group_builder,
             ),
             prefix=f"{prefix}.layers",
         )
diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py
index 1d30a0eaf6c6..f5a82b72314b 100644
--- a/vllm/model_executor/models/glm4_moe_lite_mtp.py
+++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py
@@ -28,7 +28,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.fused_moe import (
@@ -51,6 +51,10 @@
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backends.mla.index_group import (
+    SparseMLAIndexGroupBuilder,
+    get_sparse_mla_index_group_max_rows,
+)
 
 from .glm4_moe_lite import (
     Glm4MixtureOfExperts,
@@ -64,7 +68,7 @@
 class SharedHead(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         quant_config: QuantizationConfig | None = None,
     ) -> None:
@@ -108,6 +112,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             )
         else:
             topk_indices_buffer = None
+        index_group_builder = (
+            SparseMLAIndexGroupBuilder(
+                topk_indices_buffer,
+                get_sparse_mla_index_group_max_rows(vllm_config),
+            )
+            if topk_indices_buffer is not None
+            else None
+        )
 
         self.shared_head = SharedHead(
             config=config, prefix=prefix, quant_config=quant_config
@@ -117,6 +129,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             prefix=prefix,
             config=self.config,
             topk_indices_buffer=topk_indices_buffer,
+            index_group_builder=index_group_builder,
         )
 
     def forward(
@@ -441,8 +454,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loaded_params
 
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py
index bad4ec8cf333..330f891d1cf3 100644
--- a/vllm/model_executor/models/glm4_moe_mtp.py
+++ b/vllm/model_executor/models/glm4_moe_mtp.py
@@ -29,7 +29,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import CacheConfig, ParallelConfig, VllmConfig
 from vllm.model_executor.layers.fused_moe import (
@@ -60,7 +60,7 @@
 class SharedHead(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         quant_config: QuantizationConfig | None = None,
     ) -> None:
@@ -80,7 +80,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class Glm4MoeMultiTokenPredictorLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
@@ -412,8 +412,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loaded_params
 
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
diff --git a/vllm/model_executor/models/glm4v.py b/vllm/model_executor/models/glm4v.py
index b00e6873612a..7b42d5734f3a 100644
--- a/vllm/model_executor/models/glm4v.py
+++ b/vllm/model_executor/models/glm4v.py
@@ -16,7 +16,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import SiluAndMul, get_act_fn
@@ -65,12 +65,11 @@
 
 
 class GLMVImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - c: Number of channels (3)
-        - h: Height of image
-        - w: Width of image
+    """Dimensions:
+    - b: Batch size
+    - c: Number of channels (3)
+    - h: Height of image
+    - w: Width of image
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -90,14 +89,15 @@ def __init__(self, config):
         self.position_embedding = nn.Embedding(config.num_positions, config.hidden_size)
 
     def forward(self, images: torch.Tensor) -> torch.Tensor:
-        """
-        Parameters:
+        """Parameters
         images : torch.Tensor
             Input image tensor with shape (B, C, H, W)
 
-        Returns:
+        Returns
+        -------
         torch.Tensor
             Transformed tensor with shape (B, L, D)
+
         """
         images = images.to(device=self.proj.weight.device, dtype=self.proj.weight.dtype)
         x = self.proj(images)
@@ -246,8 +246,7 @@ def __init__(
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
-        """
-        The original implementation is the same as:
+        """The original implementation is the same as:
         ```python
         self.dense_h_to_4h = ColumnParallelLinear(
             config.hidden_size,
@@ -349,14 +348,15 @@ def __init__(
         self.scaling_factor = vision_config.scaling_factor
 
     def forward(self, images: torch.Tensor) -> torch.Tensor:
-        """
-        Parameters:
+        """Parameters
         images : torch.Tensor
             Input image tensor with shape (B, C, H, W)
 
-        Returns:
+        Returns
+        -------
         torch.Tensor
             Transformed tensor with shape (B, L, D)
+
         """
         x = self.patch_embedding(images)
         x = self.transformer(x)
@@ -445,22 +445,19 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         hf_config = self.info.get_hf_config()
         vision_config = hf_config.vision_config
 
         target_width = target_height = vision_config["image_size"]
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
 
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -518,9 +515,7 @@ class GLM4VForCausalLM(
     }
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="transformer.encoder",
             connector="transformer.vision.linear_proj",
diff --git a/vllm/model_executor/models/glm_ocr_mtp.py b/vllm/model_executor/models/glm_ocr_mtp.py
index 6c8e208c4b5c..6bcd33f3dc46 100644
--- a/vllm/model_executor/models/glm_ocr_mtp.py
+++ b/vllm/model_executor/models/glm_ocr_mtp.py
@@ -86,7 +86,8 @@ def forward(
     ) -> torch.Tensor:
         assert inputs_embeds is not None
         # masking inputs at position 0, as not needed by MTP
-        inputs_embeds[positions[0] == 0] = 0
+        token_positions = positions[0] if positions.ndim == 2 else positions
+        inputs_embeds.masked_fill_((token_positions == 0).unsqueeze(-1), 0)
 
         inputs_embeds = self.enorm(inputs_embeds)
         previous_hidden_states = self.hnorm(previous_hidden_states)
@@ -245,8 +246,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loaded_params
 
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
diff --git a/vllm/model_executor/models/glmasr.py b/vllm/model_executor/models/glmasr.py
index 219a7ce02de9..cf12dedbe5fb 100644
--- a/vllm/model_executor/models/glmasr.py
+++ b/vllm/model_executor/models/glmasr.py
@@ -11,7 +11,7 @@
 from transformers.models.whisper import WhisperFeatureExtractor
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TokensPrompt
@@ -44,6 +44,7 @@
     PromptUpdate,
     PromptUpdateDetails,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.sequence import IntermediateTensors
 from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.transformers_utils.processor import cached_processor_from_config
@@ -75,8 +76,7 @@
 
 
 class GlmAsrEncoderRotaryEmbedding(nn.Module):
-    """
-    Rotary Position Embedding for GLM-ASR encoder.
+    """Rotary Position Embedding for GLM-ASR encoder.
 
     Computes rotary position embeddings on-demand for efficiency.
     Only caches inv_freq as a buffer; cos/sin are computed during forward
@@ -115,8 +115,7 @@ def __init__(self, config) -> None:
         self.register_buffer("inv_freq", inv_freq, persistent=False)
 
     def forward(self, seq_len: int) -> torch.Tensor:
-        """
-        Compute rotary position frequencies for given sequence length.
+        """Compute rotary position frequencies for given sequence length.
 
         Args:
             seq_len: The sequence length to compute embeddings for.
@@ -124,6 +123,7 @@ def forward(self, seq_len: int) -> torch.Tensor:
         Returns:
             Frequency tensor with shape [seq_len, dim/2]. Use .cos() and
             .sin() to get the rotary embedding components.
+
         """
         # Compute on the same device as inv_freq (automatically correct after .to())
         seq = torch.arange(
@@ -134,8 +134,7 @@ def forward(self, seq_len: int) -> torch.Tensor:
 
 
 class GlmAsrEncoderAttention(nn.Module):
-    """
-    Optimized Multi-headed Grouped Query Attention for GLM-ASR encoder.
+    """Optimized Multi-headed Grouped Query Attention for GLM-ASR encoder.
 
     Uses vLLM's QKVParallelLinear for fused projections, ApplyRotaryEmb for
     rotary position embeddings, and MMEncoderAttention for hardware-optimized
@@ -208,14 +207,14 @@ def forward(
         rotary_pos_emb_cos: torch.Tensor,
         rotary_pos_emb_sin: torch.Tensor,
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             hidden_states: [batch_size, seq_len, hidden_size]
             rotary_pos_emb_cos: [seq_len, rotary_dim/2] - cosine of rotary embeddings
             rotary_pos_emb_sin: [seq_len, rotary_dim/2] - sine of rotary embeddings
 
         Returns:
             [batch_size, seq_len, hidden_size]
+
         """
         batch_size, seq_len, _ = hidden_states.shape
 
@@ -255,8 +254,7 @@ def forward(
 
 
 class GlmAsrEncoderMLP(nn.Module):
-    """
-    Optimized MLP for GLM-ASR encoder.
+    """Optimized MLP for GLM-ASR encoder.
     Uses vLLM's parallel linear layers for better performance.
     """
 
@@ -297,8 +295,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class GlmAsrEncoderLayer(nn.Module):
-    """
-    Optimized Transformer encoder layer for GLM-ASR.
+    """Optimized Transformer encoder layer for GLM-ASR.
     Combines attention and MLP with residual connections and layer norms.
     """
 
@@ -335,14 +332,14 @@ def forward(
         rotary_pos_emb_cos: torch.Tensor,
         rotary_pos_emb_sin: torch.Tensor,
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             hidden_states: [batch_size, seq_len, hidden_size]
             rotary_pos_emb_cos: [seq_len, rotary_dim/2] - cosine of rotary embeddings
             rotary_pos_emb_sin: [seq_len, rotary_dim/2] - sine of rotary embeddings
 
         Returns:
             [batch_size, seq_len, hidden_size]
+
         """
         # Self-attention with residual
         residual = hidden_states
@@ -364,8 +361,7 @@ def forward(
 
 
 class _GlmAsrEncoderOutput:
-    """
-    Simple output container compatible with transformers' BaseModelOutput.
+    """Simple output container compatible with transformers' BaseModelOutput.
 
     This lightweight container holds the encoder output and is compatible
     with the transformers library's output format while being more efficient
@@ -374,6 +370,7 @@ class _GlmAsrEncoderOutput:
     Attributes:
         last_hidden_state: Final layer hidden states from the encoder.
             Shape: [batch_size, seq_len, hidden_size]
+
     """
 
     __slots__ = ("last_hidden_state",)
@@ -383,8 +380,7 @@ def __init__(self, last_hidden_state: torch.Tensor):
 
 
 class GlmAsrEncoder(nn.Module):
-    """
-    Optimized GLM-ASR Audio Encoder with vLLM native implementation.
+    """Optimized GLM-ASR Audio Encoder with vLLM native implementation.
 
     This encoder processes audio features through convolutional layers
     followed by transformer layers with rotary position embeddings.
@@ -454,14 +450,14 @@ def __init__(
     def _get_feat_extract_output_lengths(
         self, input_lengths: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Compute the output length after convolutions.
+        """Compute the output length after convolutions.
 
         Args:
             input_lengths: Input sequence lengths [batch_size]
 
         Returns:
             Tuple of (output after conv1, output after conv2)
+
         """
         # Conv1: kernel=3, stride=1, padding=1
         output_lengths_conv1 = (input_lengths + 2 * 1 - 3) // 1 + 1
@@ -472,8 +468,7 @@ def _get_feat_extract_output_lengths(
         return output_lengths_conv1, output_lengths_conv2
 
     def forward(self, input_features: torch.Tensor) -> _GlmAsrEncoderOutput:
-        """
-        Forward pass through the encoder.
+        """Forward pass through the encoder.
 
         Args:
             input_features: [batch_size, num_mel_bins, seq_len]
@@ -482,6 +477,7 @@ def forward(self, input_features: torch.Tensor) -> _GlmAsrEncoderOutput:
             _GlmAsrEncoderOutput: Object with .last_hidden_state attribute \
                 containing [batch_size, seq_len', hidden_size] where seq_len' \
                 is the sequence length after convolutions
+
         """
         # Apply convolutional layers with GELU activation
         hidden_states = torch.nn.functional.gelu(self.conv1(input_features))
@@ -515,11 +511,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
 
 class GlmAsrFeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - num_chunks: Number of audio chunks (flattened)
-        - nmb: Number of mel bins
-        - num_audios: Number of original audio files
+    """Dimensions:
+    - num_chunks: Number of audio chunks (flattened)
+    - nmb: Number of mel bins
+    - num_audios: Number of original audio files
     """
 
     type: Literal["audio_features"]
@@ -538,12 +533,11 @@ class GlmAsrFeatureInputs(TensorSchema):
 
 
 class GlmAsrEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size
-        - naf: Number of audio features
-        - hs: Hidden size (must match the hidden size of language model
-          backbone)
+    """Dimensions:
+    - bn: Batch size
+    - naf: Number of audio features
+    - hs: Hidden size (must match the hidden size of language model
+      backbone)
     """
 
     type: Literal["audio_embeds"] = "audio_embeds"
@@ -557,8 +551,7 @@ class GlmAsrEmbeddingInputs(TensorSchema):
 
 
 class GlmAsrMultiModalProjector(nn.Module):
-    """
-    Projects audio encoder outputs to language model hidden space.
+    """Projects audio encoder outputs to language model hidden space.
 
     This projector uses a two-layer MLP to map audio features from the
     encoder's intermediate size to the language model's hidden size.
@@ -601,8 +594,7 @@ def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
 def _glmasr_field_config(
     hf_inputs: Mapping[str, torch.Tensor],
 ) -> dict[str, MultiModalFieldConfig]:
-    """
-    Configure multimodal field batching strategy for GLM-ASR.
+    """Configure multimodal field batching strategy for GLM-ASR.
 
     Determines how to batch audio inputs based on whether chunking is used.
     When chunk_counts is present, features are flattened across chunks;
@@ -614,6 +606,7 @@ def _glmasr_field_config(
     Returns:
         Dictionary mapping field names to MultiModalFieldConfig objects \
             that specify batching behavior.
+
     """
     chunk_counts = hf_inputs.get("chunk_counts")
     if chunk_counts is not None:
@@ -636,8 +629,7 @@ def _glmasr_field_config(
 
 
 class GlmAsrMultiModalDataParser(MultiModalDataParser):
-    """
-    Custom parser for GLM-ASR multimodal data.
+    """Custom parser for GLM-ASR multimodal data.
 
     Extends the base parser to handle GLM-ASR specific audio data formats,
     including both pre-computed audio embeddings and raw audio features.
@@ -658,8 +650,7 @@ def _parse_audio_data(
 
 
 class GlmAsrProcessingInfo(BaseProcessingInfo):
-    """
-    Processing information provider for GLM-ASR model.
+    """Processing information provider for GLM-ASR model.
 
     Provides access to model configuration, processor, and feature extractor
     needed for audio preprocessing and multimodal integration.
@@ -687,8 +678,7 @@ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
 
 
 class GlmAsrDummyInputsBuilder(BaseDummyInputsBuilder[GlmAsrProcessingInfo]):
-    """
-    Builder for dummy inputs used in profiling and testing.
+    """Builder for dummy inputs used in profiling and testing.
 
     Generates dummy text prompts and audio data that match the expected
     format for GLM-ASR model inputs. Used for memory profiling and
@@ -704,13 +694,10 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         feature_extractor = self.info.get_feature_extractor()
         sampling_rate = feature_extractor.sampling_rate
-        num_audios = mm_counts.get("audio", 0)
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
 
         max_audio_len = getattr(
             self.info.get_hf_processor(), "max_audio_len", DEFAULT_MAX_AUDIO_LEN_S
@@ -720,15 +707,14 @@ def get_dummy_mm_data(
         return {
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
 
 class GlmAsrMultiModalProcessor(BaseMultiModalProcessor["GlmAsrProcessingInfo"]):
-    """
-    GLM-ASR processor that inherits directly from BaseMultiModalProcessor
+    """GLM-ASR processor that inherits directly from BaseMultiModalProcessor
     for better performance and cleaner implementation.
     """
 
@@ -751,26 +737,24 @@ def _calculate_chunk_counts(
             chunk_counts.append(min(n_chunks, max_windows))
         return chunk_counts
 
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        mm_data = dict(mm_data)
-        if "audios" in mm_data:
-            mm_data["audio"] = mm_data.pop("audios")
-
-        feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
-        hf_processor_mm_kwargs = dict(
-            **hf_processor_mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+        return hf_inputs._replace(
+            hf_kwargs=dict(
+                hf_inputs.hf_kwargs,
+                sampling_rate=feature_extractor.sampling_rate,
+            )
         )
 
-        return mm_data, hf_processor_mm_kwargs
-
     def _postprocess_hf_mm_data(
         self,
         mm_data: Mapping[str, object],
diff --git a/vllm/model_executor/models/glmasr_utils.py b/vllm/model_executor/models/glmasr_utils.py
index 8dcfcfa89513..7997963ddab8 100644
--- a/vllm/model_executor/models/glmasr_utils.py
+++ b/vllm/model_executor/models/glmasr_utils.py
@@ -71,8 +71,7 @@ def _get_audio_output_lengths_for_tower(
     merge_factor: int,
     conv_params: list[tuple[int, int, int]],
 ) -> torch.Tensor:
-    """
-    Calculate the output lengths after audio processing.
+    """Calculate the output lengths after audio processing.
 
     The output length accounts for:
     1. Convolution layers (downsampling)
@@ -86,6 +85,7 @@ def _get_audio_output_lengths_for_tower(
 
     Returns:
         Output lengths after all processing [batch_size]
+
     """
     # First, calculate the output length after convolutions
     if hasattr(audio_tower, "_get_feat_extract_output_lengths"):
diff --git a/vllm/model_executor/models/gpt2.py b/vllm/model_executor/models/gpt2.py
index 07801cd098b1..c95fe5184898 100644
--- a/vllm/model_executor/models/gpt2.py
+++ b/vllm/model_executor/models/gpt2.py
@@ -327,6 +327,7 @@ class GPT2ForSequenceClassification(nn.Module, SupportsCrossEncoding):
         transformer: An instance of GPT2Model used for forward operations.
         score: A layer for calculating logits.
         _pooler: An instance of Pooler used for pooling operations.
+
     """
 
     is_pooling_model = True
diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py
index 055713566b1b..b96ecf90d3e4 100644
--- a/vllm/model_executor/models/gpt_oss.py
+++ b/vllm/model_executor/models/gpt_oss.py
@@ -22,6 +22,7 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
@@ -29,7 +30,6 @@
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -417,11 +417,10 @@ def __init__(
         self.hidden_size = config.hidden_size
         self.experts_per_token = config.num_experts_per_tok
         self.world_size = dist.get_world_size() if dist.is_initialized() else 1
-        self.router = ReplicatedLinear(
+        self.router = GateLinear(
             config.hidden_size,
             config.num_local_experts,
             bias=True,
-            quant_config=None,
             prefix=f"{prefix}.router",
             return_bias=False,
         )
@@ -850,6 +849,7 @@ def _get_moe_weight_dtype(layer_id: int = 0) -> str | None:
 
             Returns:
                 Weight dtype string (e.g., "mxfp4", "fp8") or None if not available
+
             """
             if hasattr(self.layers[layer_id].mlp.experts._quant_method, "weight_dtype"):
                 return self.layers[layer_id].mlp.experts._quant_method.weight_dtype
diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py
index 55a90f5ff04e..81ab62937e1a 100644
--- a/vllm/model_executor/models/granite.py
+++ b/vllm/model_executor/models/granite.py
@@ -36,7 +36,7 @@
 import torch
 from torch import nn
 from transformers import GraniteConfig
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -72,7 +72,7 @@
 
 
 def granite_layer_attn_params(
-    config: PretrainedConfig, layer_idx: int
+    config: PreTrainedConfig, layer_idx: int
 ) -> tuple[int | None, float, bool]:
     """Resolve one layer's sliding window, RoPE base and sink usage.
 
@@ -84,6 +84,7 @@ def granite_layer_attn_params(
     Returns:
         Sliding window size (`None` for full attention), RoPE base theta (`0`
         for NoPE), and attention sink presence/absence.
+
     """
     layer_types = getattr(config, "layer_types", None)
     sliding_window = (
diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py
index 50b0932bda8e..b62a86543414 100644
--- a/vllm/model_executor/models/granite4_vision.py
+++ b/vllm/model_executor/models/granite4_vision.py
@@ -684,6 +684,7 @@ def _get_all_layer_features(
 
         Packing on dim=-1 means the framework's token-level slicing for
         chunked prefill preserves all levels intact.
+
         """
         select_strategy = self._vision_feature_select_strategy
 
diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py
index 0fac643999ff..3ae85cfaac69 100644
--- a/vllm/model_executor/models/granite_speech.py
+++ b/vllm/model_executor/models/granite_speech.py
@@ -31,10 +31,10 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from transformers import BatchFeature, PretrainedConfig
+from transformers import BatchFeature, PreTrainedConfig
 
 from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt
 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
@@ -91,8 +91,7 @@
 
 ### Audio Input
 class GraniteSpeechAudioInputs(TensorSchema):
-    """
-    Audio input features for Granite Speech model.
+    """Audio input features for Granite Speech model.
 
     Dimensions:
         - b: Batch size
@@ -111,7 +110,7 @@ class GraniteSpeechAudioInputs(TensorSchema):
     """List of audio embedding sizes for each item in batch."""
 
 
-class GraniteSpeechMultiModalProcessingInfo(BaseProcessingInfo):
+class GraniteSpeechProcessingInfo(BaseProcessingInfo):
     def get_data_parser(self):
         feature_extractor = self.get_hf_processor().audio_processor
 
@@ -136,7 +135,7 @@ def get_max_audio_len(self):
 
 ### Input Processing  & Multimodal utils
 class GraniteSpeechMultiModalProcessor(
-    BaseMultiModalProcessor[GraniteSpeechMultiModalProcessingInfo]
+    BaseMultiModalProcessor[GraniteSpeechProcessingInfo]
 ):
     def _get_mm_fields_config(
         self,
@@ -182,23 +181,9 @@ def get_replacement(item_idx: int):
             )
         ]
 
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", [])
-
-        if audios:
-            # GraniteSpeechFeatureExtractor accepts "audio"
-            mm_data["audio"] = audios
-
-        return mm_data, hf_processor_mm_kwargs
-
     def _postprocess_hf_mm_data(
         self,
         mm_data: Mapping[str, object],
@@ -217,23 +202,19 @@ def _postprocess_hf_mm_data(
 
 
 class GraniteSpeechDummyInputsBuilder(
-    BaseDummyInputsBuilder[GraniteSpeechMultiModalProcessingInfo]
+    BaseDummyInputsBuilder[GraniteSpeechProcessingInfo]
 ):
     def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
-
         return {
             "audio": self._get_dummy_audios(
                 length=self.info.get_max_audio_len(),
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
@@ -248,7 +229,7 @@ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
 class GraniteSpeechEncoderProjector(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -306,7 +287,7 @@ class GraniteSpeechConformerFeedForward(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -342,7 +323,7 @@ class GraniteSpeechConformerAttention(nn.Module):
     for more details.
     """
 
-    def __init__(self, config: PretrainedConfig, prefix: str = ""):
+    def __init__(self, config: PreTrainedConfig, prefix: str = ""):
         super().__init__()
 
         inner_dim = config.dim_head * config.num_heads
@@ -448,7 +429,7 @@ class GraniteSpeechConformerConvModule(nn.Module):
     convolutional layers.
     """
 
-    def __init__(self, config: PretrainedConfig, prefix: str = ""):
+    def __init__(self, config: PreTrainedConfig, prefix: str = ""):
         super().__init__()
         inner_dim = config.hidden_dim * config.conv_expansion_factor
 
@@ -479,7 +460,7 @@ class GraniteSpeechConformerBlock(nn.Module):
     """Conformer block, consisting largely of linear layers,
     attention, and convolutional layers."""
 
-    def __init__(self, config: PretrainedConfig, prefix: str = ""):
+    def __init__(self, config: PreTrainedConfig, prefix: str = ""):
         super().__init__()
         self.ff1 = GraniteSpeechConformerFeedForward(config, prefix=f"{prefix}.ff1")
         self.attn = GraniteSpeechConformerAttention(config, prefix=f"{prefix}.attn")
@@ -505,7 +486,7 @@ class GraniteSpeechCTCEncoder(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         quant_config: QuantizationConfig | None = None,
     ):
@@ -565,7 +546,7 @@ def forward(self, hidden_states: torch.Tensor):
 
 @MULTIMODAL_REGISTRY.register_processor(
     GraniteSpeechMultiModalProcessor,
-    info=GraniteSpeechMultiModalProcessingInfo,
+    info=GraniteSpeechProcessingInfo,
     dummy_inputs=GraniteSpeechDummyInputsBuilder,
 )
 class GraniteSpeechForConditionalGeneration(
@@ -641,7 +622,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
     def _build_encoder(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str,
     ) -> "GraniteSpeechCTCEncoder":
@@ -732,9 +713,11 @@ def _build_input_features_mask(
         Args:
             audio_embed_sizes: torch.Tensor
                 Tensor of num features in each seq in the batch.
+
         Returns:
             torch.Tensor: Mask of shape (bsz, num_features) to be applied to
             the audio features prior to splitting the audio embeddings.
+
         """
         most_audio_features = int(torch.max(audio_embed_sizes))
         mask_indices = torch.arange(most_audio_features).view(1, -1)
@@ -763,10 +746,12 @@ def _pad_and_stack_input_features(
         Args:
             input_features: list[torch.Tensor]
                 3D Input features to be coerced into a tensor.
+
         Returns:
             torch.Tensor: Tensor of shape [bsz, num_features, 160], where
             num_features is the max number of features of any entry in the
             batch.
+
         """
         feat_lens = [feats.shape[1] for feats in input_features]
         padding = [max(feat_lens) - length for length in feat_lens]
@@ -790,8 +775,10 @@ def _process_audio_input(
             audio_input: GraniteSpeechAudioInputs
                 Audio inputs object containing Mel features, an input features
                 mask, and the (flattened) number of audio tokens per instance.
+
         Returns:
             tuple[torch.Tensor]: List of length bsz.
+
         """
         # TODO (Alex) - support embedding inputs
         encoder_embeds = self.encoder(audio_input["input_features"])
diff --git a/vllm/model_executor/models/granite_speech_plus.py b/vllm/model_executor/models/granite_speech_plus.py
index ba95cfed131e..946f3edbcc3e 100644
--- a/vllm/model_executor/models/granite_speech_plus.py
+++ b/vllm/model_executor/models/granite_speech_plus.py
@@ -3,7 +3,7 @@
 """Inference-only IBM Granite Speech Plus model."""
 
 import torch
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -12,8 +12,8 @@
     GraniteSpeechCTCEncoder,
     GraniteSpeechDummyInputsBuilder,
     GraniteSpeechForConditionalGeneration,
-    GraniteSpeechMultiModalProcessingInfo,
     GraniteSpeechMultiModalProcessor,
+    GraniteSpeechProcessingInfo,
 )
 
 ISO639_1_SUPPORTED_LANGS = {
@@ -60,7 +60,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 @MULTIMODAL_REGISTRY.register_processor(
     GraniteSpeechMultiModalProcessor,
-    info=GraniteSpeechMultiModalProcessingInfo,
+    info=GraniteSpeechProcessingInfo,
     dummy_inputs=GraniteSpeechDummyInputsBuilder,
 )
 class GraniteSpeechPlusForConditionalGeneration(GraniteSpeechForConditionalGeneration):
@@ -68,7 +68,7 @@ class GraniteSpeechPlusForConditionalGeneration(GraniteSpeechForConditionalGener
 
     def _build_encoder(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str,
     ) -> GraniteSpeechCTCEncoder:
diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py
index 3a9872a60573..5b9aa886f52b 100644
--- a/vllm/model_executor/models/granitemoe.py
+++ b/vllm/model_executor/models/granitemoe.py
@@ -29,7 +29,7 @@
 
 import torch
 from torch import nn
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -39,11 +39,13 @@
     tensor_model_parallel_all_gather,
 )
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -94,12 +96,10 @@ def __init__(
         self.is_sequence_parallel = is_sequence_parallel
 
         # Gate always runs at half / full precision for now.
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             hidden_size,
             num_experts,
-            bias=False,
             params_dtype=params_dtype,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
@@ -141,7 +141,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class GraniteMoeAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py
index 8dc47c522654..da7cd0478b55 100644
--- a/vllm/model_executor/models/granitemoehybrid.py
+++ b/vllm/model_executor/models/granitemoehybrid.py
@@ -633,6 +633,7 @@ def get_mamba_state_shape_from_config(
             Tuple containing:
             - conv_state_shape: Shape for convolutional state cache
             - temporal_state_shape: Shape for state space model cache
+
         """
         parallel_config = vllm_config.parallel_config
         hf_config = vllm_config.model_config.hf_config
diff --git a/vllm/model_executor/models/h2ovl.py b/vllm/model_executor/models/h2ovl.py
index 5c4f984c13f8..e1c59d4ea08e 100644
--- a/vllm/model_executor/models/h2ovl.py
+++ b/vllm/model_executor/models/h2ovl.py
@@ -10,18 +10,18 @@
 # --------------------------------------------------------
 
 import torch
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import BatchedTensorInputs
+from vllm.multimodal.inputs import BatchedTensorInputs, MultiModalKwargsItem
 from vllm.multimodal.parse import (
     ImageEmbeddingItems,
     ImageProcessorItems,
     MultiModalDataItems,
 )
 from vllm.multimodal.processing.processor import (
-    MultiModalProcessingInfo,
+    MultiModalProcessingResult,
     ProcessorInputs,
     PromptReplacement,
     TimingContext,
@@ -139,7 +139,7 @@ def _cached_apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
+    ) -> MultiModalProcessingResult:
         # The processor logic is different for len(images) <= 1 vs > 1
         # Since the processing cache assumes that the processor output is
         # invariant of how many images are passed per prompt, we only
@@ -158,7 +158,7 @@ def _cached_apply_hf_processor(
 class H2OVLChatModel(InternVLChatModel):
     def _init_vision_model(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
@@ -178,16 +178,17 @@ def _init_vision_model(
             prefix=prefix,
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        if num_image_tokens <= 0 or self.num_image_token <= 0:
-            return 0
-
-        num_patches = num_image_tokens // self.num_image_token
-        return num_patches * (self.patch_tokens + 1)
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
-        if num_vision_tokens <= 0 or self.num_image_token <= 0:
-            return 0
-
-        num_patches = num_vision_tokens // (self.patch_tokens + 1)
-        return num_patches * self.num_image_token
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+        if num_mm_embeds <= 0 or self.num_image_token <= 0:
+            return 0, 0
+
+        num_patches = num_mm_embeds // self.num_image_token
+        tower_tokens = num_patches * (self.patch_tokens + 1)
+        return tower_tokens, num_patches * self.num_image_token
diff --git a/vllm/model_executor/models/hrm_text.py b/vllm/model_executor/models/hrm_text.py
index db8d9c20c39d..2f11a2c45472 100644
--- a/vllm/model_executor/models/hrm_text.py
+++ b/vllm/model_executor/models/hrm_text.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-HRM-Text: Hierarchical Reasoning Model — Text variant.
+"""HRM-Text: Hierarchical Reasoning Model — Text variant.
 
 Reference Hugging Face implementation:
     src/transformers/models/hrm_text/modeling_hrm_text.py
@@ -26,7 +25,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -123,7 +122,7 @@ class HrmTextAttention(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_idx_in_stack: int,
         stack_kind: Literal["L", "H"],
         cache_config: CacheConfig | None = None,
@@ -246,7 +245,7 @@ def forward(
 class HrmTextDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_idx_in_stack: int,
         stack_kind: Literal["L", "H"],
         cache_config: CacheConfig | None = None,
@@ -311,7 +310,7 @@ class HrmTextStack(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         stack_kind: Literal["L", "H"],
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py
index 0b54c26f47a4..2e24cf65e8f3 100644
--- a/vllm/model_executor/models/hy_v3.py
+++ b/vllm/model_executor/models/hy_v3.py
@@ -31,7 +31,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
@@ -220,7 +220,7 @@ def forward(
 class HYV3Attention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
@@ -300,8 +300,11 @@ def __init__(
             prefix=f"{prefix}.attn",
         )
         if self.use_qk_norm:
-            self.q_norm = RMSNorm(self.head_dim, rms_norm_eps)
-            self.k_norm = RMSNorm(self.head_dim, rms_norm_eps)
+            # The HPC fused kernel reads norm weights directly in float32;
+            # the fallback path keeps the model default dtype.
+            norm_dtype = torch.float32 if rope_support else None
+            self.q_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype)
+            self.k_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype)
 
         # HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).
         # HunYuan V3 applies QK-Norm *before* RoPE, so NORM_THEN_ROPE.
@@ -362,7 +365,7 @@ def forward(
 class HYV3DecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/hy_v3_mtp.py b/vllm/model_executor/models/hy_v3_mtp.py
index aa9c2f35b5e2..06f6f270890c 100644
--- a/vllm/model_executor/models/hy_v3_mtp.py
+++ b/vllm/model_executor/models/hy_v3_mtp.py
@@ -29,7 +29,7 @@
 import regex as re
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
 from vllm.model_executor.layers.fused_moe import (
@@ -59,7 +59,7 @@
 )
 
 
-def _is_moe(config: PretrainedConfig) -> bool:
+def _is_moe(config: PreTrainedConfig) -> bool:
     return bool(
         getattr(config, "num_experts", None)
         and (
@@ -69,7 +69,7 @@ def _is_moe(config: PretrainedConfig) -> bool:
     )
 
 
-def _get_cla_factor(config: PretrainedConfig) -> int:
+def _get_cla_factor(config: PreTrainedConfig) -> int:
     if not getattr(config, "use_cla", False):
         return 1
     return getattr(config, "cla_share_factor", 1)
@@ -78,7 +78,7 @@ def _get_cla_factor(config: PretrainedConfig) -> int:
 class HYV3SharedHead(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
     ) -> None:
         super().__init__()
@@ -93,7 +93,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class HYV3MultiTokenPredictorLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         model_config: ModelConfig,
         cache_config: CacheConfig | None = None,
diff --git a/vllm/model_executor/models/hyperclovax_vision_v2.py b/vllm/model_executor/models/hyperclovax_vision_v2.py
index 96e718751d59..ee23d2c4df02 100644
--- a/vllm/model_executor/models/hyperclovax_vision_v2.py
+++ b/vllm/model_executor/models/hyperclovax_vision_v2.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-HyperCLOVAX V2 (32B Think Model) Implementation.
+"""HyperCLOVAX V2 (32B Think Model) Implementation.
 
 This module contains the V2 architecture that uses Qwen2.5 Vision Transformer
 instead of CLIP/SigLIP used in V1.
@@ -19,7 +18,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.forward_context import set_forward_context
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -57,8 +56,7 @@
 
 
 class HCXVisionV2ImagePixelInputs(TensorSchema):
-    """
-    V2 Image inputs using Qwen2.5-VL style grid_thw format.
+    """V2 Image inputs using Qwen2.5-VL style grid_thw format.
 
     Dimensions:
         - np: Number of patches
@@ -72,8 +70,7 @@ class HCXVisionV2ImagePixelInputs(TensorSchema):
 
 
 class HCXVisionV2ImageEmbeddingInputs(TensorSchema):
-    """
-    V2 Image embedding inputs.
+    """V2 Image embedding inputs.
 
     Dimensions:
         - nf: Number of image features
@@ -90,8 +87,7 @@ class HCXVisionV2ImageEmbeddingInputs(TensorSchema):
 
 
 class HCXVisionV2VideoPixelInputs(TensorSchema):
-    """
-    V2 Video inputs using Qwen2.5-VL style grid_thw format.
+    """V2 Video inputs using Qwen2.5-VL style grid_thw format.
 
     Dimensions:
         - np: Number of patches
@@ -105,8 +101,7 @@ class HCXVisionV2VideoPixelInputs(TensorSchema):
 
 
 class HCXVisionV2VideoEmbeddingInputs(TensorSchema):
-    """
-    V2 Video embedding inputs.
+    """V2 Video embedding inputs.
 
     Dimensions:
         - nf: Number of video features
@@ -194,67 +189,29 @@ def get_dummy_text(
         num_videos = mm_counts.get("video", 0)
         return V2_IMAGE_TOKEN * num_images + V2_VIDEO_TOKEN * num_videos
 
-    def get_dummy_processor_inputs(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions] | None = None,
-        mm_processor_kwargs: Mapping[str, object] | None = None,
-    ) -> ProcessorInputs:
-        """Build dummy processor inputs for memory profiling."""
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-        prompt_text = V2_IMAGE_TOKEN * num_images + V2_VIDEO_TOKEN * num_videos
-
-        dummy_mm_data = self.get_dummy_mm_data(
-            seq_len,
-            mm_counts,
-            mm_options,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False)
-
-        tokenizer = self.info.get_tokenizer()
-        prompt = tokenizer.encode(
-            prompt_text,
-            **self.info.default_tok_params.get_encode_kwargs(),
-        )
-
-        return ProcessorInputs(
-            prompt=prompt,
-            mm_data_items=dummy_mm_items,
-            hf_processor_mm_kwargs=mm_processor_kwargs or {},
-        )
-
     def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+        mm_options: MultiModalDummyOptions | None = None,
         mm_processor_kwargs: Mapping[str, object] | None = None,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = 16  # Default for video
 
-        image_overrides = mm_options.get("image") if mm_options else None
-        video_overrides = mm_options.get("video") if mm_options else None
-
         result: MultiModalDataDict = {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,  # type: ignore
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image") if mm_options else None,
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,  # type: ignore
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video") if mm_options else None,
             ),
         }
 
@@ -266,44 +223,30 @@ class HCXVisionV2MultiModalProcessor(
 ):
     """Multimodal processor for HyperCLOVAX V2 (32B Think model)."""
 
-    def _apply_hf_processor_main(
+    def get_dummy_inputs(
         self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
-
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
-
-        images = mm_data.get("images")
-        videos = mm_data.get("videos")
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: MultiModalDummyOptions,
+    ) -> ProcessorInputs:
+        num_images = mm_counts.get("image", 0)
+        num_videos = mm_counts.get("video", 0)
+        prompt_text = V2_IMAGE_TOKEN * num_images + V2_VIDEO_TOKEN * num_videos
 
-        # Get the HF processor
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        builder = self.dummy_inputs
+        dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options)
+        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False)
 
-        # Build data dict for HF processor (images/videos only)
-        # NOTE: We pass the prompt as-is without token normalization.
-        # Token expansion is handled by vLLM via _get_prompt_updates.
-        data: dict[str, object] = dict(
-            text=prompt_text,
-            images=images,
-            videos=videos,
+        tokenizer = self.info.get_tokenizer()
+        prompt = tokenizer.encode(
+            prompt_text,
+            **self.info.default_tok_params.get_encode_kwargs(),
         )
 
-        processed_data = self.info.ctx.call_hf_processor(
-            hf_processor=hf_processor,
-            data=data,
-            kwargs=hf_processor_mm_kwargs,
-        )
-        processed_data.update(passthrough_data)
+        return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items)
 
-        return processed_data
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _get_prompt_updates(
         self,
@@ -314,7 +257,7 @@ def _get_prompt_updates(
         hf_config = self.info.get_hf_config()
 
         # Use token IDs directly from config.
-        # This matches what get_dummy_processor_inputs uses, ensuring consistency.
+        # This matches what get_dummy_inputs uses, ensuring consistency.
         placeholder: dict[str, int] = {
             "image": hf_config.image_token_id,  # 128060 for <|IMAGE_PAD|>
             "video": hf_config.video_token_id,  # 128061 for <|VIDEO_PAD|>
@@ -419,8 +362,7 @@ def _get_mm_fields_config(
     dummy_inputs=HCXVisionV2DummyInputsBuilder,
 )
 class HCXVisionV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
-    """
-    HyperCLOVAX-SEED Vision-Language Model (V2 architecture).
+    """HyperCLOVAX-SEED Vision-Language Model (V2 architecture).
 
     Supports:
     - HyperCLOVAX-SEED-Think-32B: Vision + Text
diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py
index 3021e66f759f..09809db7022e 100644
--- a/vllm/model_executor/models/idefics2_vision_model.py
+++ b/vllm/model_executor/models/idefics2_vision_model.py
@@ -45,8 +45,7 @@
 
 
 class Idefics2VisionEmbeddings(nn.Module):
-    """
-    This is a modified version of `siglip.modelign_siglip.SiglipVisionEmbeddings
+    """This is a modified version of `siglip.modelign_siglip.SiglipVisionEmbeddings
     ` to enable images of variable
     resolution.
 
@@ -153,7 +152,7 @@ def forward(
 
 
 class Idefics2VisionAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
 
     def __init__(
         self,
@@ -308,10 +307,9 @@ def forward(
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`):
-                Input to the layer of shape `(batch, seq_len, embed_dim)`.
+        """Args:
+        hidden_states (`torch.FloatTensor`):
+            Input to the layer of shape `(batch, seq_len, embed_dim)`.
 
         """
         residual = hidden_states
@@ -326,13 +324,13 @@ def forward(
 
 
 class Idefics2Encoder(nn.Module):
-    """
-    Transformer encoder consisting of `config.num_hidden_layers` self attention
+    """Transformer encoder consisting of `config.num_hidden_layers` self attention
     layers. Each layer is a
     [`Idefics2EncoderLayer`].
 
     Args:
         config: Idefics2Config
+
     """
 
     def __init__(
@@ -368,14 +366,14 @@ def forward(
         inputs_embeds: torch.Tensor,
         attention_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        r"""
-        Args:
-            inputs_embeds (torch.Tensor):
-                Optionally, instead of passing `input_ids` you can choose to
-                directly pass an embedded representation.
-                This is useful if you want more control over how to convert
-                `input_ids` indices into associated vectorsthan the model's
-                internal embedding lookup matrix.
+        r"""Args:
+        inputs_embeds (torch.Tensor):
+            Optionally, instead of passing `input_ids` you can choose to
+            directly pass an embedded representation.
+            This is useful if you want more control over how to convert
+            `input_ids` indices into associated vectorsthan the model's
+            internal embedding lookup matrix.
+
         """
         hidden_states = inputs_embeds
         for encoder_layer in self.layers:
diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py
index baaa4106cb6b..80c71cfe7992 100644
--- a/vllm/model_executor/models/idefics3.py
+++ b/vllm/model_executor/models/idefics3.py
@@ -29,7 +29,7 @@
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -39,6 +39,7 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
@@ -51,6 +52,7 @@
     PromptUpdateDetails,
     cached_encode,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.sequence import IntermediateTensors
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -69,13 +71,12 @@
 
 
 class Idefics3ImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - bnp: Batch size * number of images * number of patches
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
+    """Dimensions:
+    - bn: Batch size * number of images
+    - bnp: Batch size * number of images * number of patches
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
     """
 
     type: Literal["pixel_values"]
@@ -85,11 +86,10 @@ class Idefics3ImagePixelInputs(TensorSchema):
 
 
 class Idefics3ImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - f: Image feature size
-        - h: Hidden size (must match the hidden size of language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - f: Image feature size
+    - h: Hidden size (must match the hidden size of language model backbone)
     """
 
     type: Literal["image_embeds"]
@@ -221,51 +221,49 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         hf_processor = self.info.get_hf_processor()
         image_processor: Idefics3ImageProcessor = hf_processor.image_processor
         longest_edge = image_processor.max_image_size["longest_edge"]
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=longest_edge,
                 height=longest_edge,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
 
 class Idefics3MultiModalProcessor(BaseMultiModalProcessor[Idefics3ProcessingInfo]):
-    def _apply_hf_processor_main(
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
 
         image_processor = self.info.get_hf_processor().image_processor
         if getattr(image_processor, "backend", "pil") == "pil":
-            hf_processor_mm_kwargs = {
-                "input_data_format": "channels_last",
-                **hf_processor_mm_kwargs,
-            }
+            hf_inputs = hf_inputs._replace(
+                hf_kwargs={
+                    "input_data_format": "channels_last",
+                    **hf_inputs.hf_kwargs,
+                }
+            )
 
-        processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            mm_data,
-            hf_processor_mm_kwargs,
-        )
+        return hf_inputs
+
+    def _postprocess_hf_mm_data(
+        self,
+        mm_data: Mapping[str, object],
+        hf_processor_mm_kwargs: Mapping[str, object],
+        processed_data: BatchFeature,
+    ) -> BatchFeature:
+        if not mm_data:
+            return processed_data
 
         images = mm_data.get("images", [])
         mm_items = self.info.parse_mm_data({"image": images}, validate=False)
@@ -290,8 +288,6 @@ def _apply_hf_processor_main(
         processed_data["pixel_values"].squeeze_(0)
         processed_data["pixel_attention_mask"].squeeze_(0)
 
-        processed_data.update(passthrough_data)
-
         return processed_data
 
     def _get_mm_fields_config(
@@ -828,29 +824,21 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="model.text_model",
             connector="model.connector",
             tower_model="model.vision_model",
         )
 
-    def get_num_mm_encoder_tokens(
+    def get_mm_lora_token_counts(
         self,
-        num_image_tokens: int,
-    ) -> int:
-        hf_config = self.config
-        scale_factor = hf_config.scale_factor
-
-        return num_image_tokens * scale_factor**2
-
-    def get_num_mm_connector_tokens(
-        self,
-        num_vision_tokens: int,
-    ) -> int:
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         hf_config = self.config
         scale_factor = hf_config.scale_factor
-
-        return num_vision_tokens // scale_factor**2
+        return num_mm_embeds * scale_factor**2, num_mm_embeds
diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py
index 254f2c635bc1..21a09cb7ba2e 100644
--- a/vllm/model_executor/models/interfaces.py
+++ b/vllm/model_executor/models/interfaces.py
@@ -31,7 +31,7 @@
 import torch.nn as nn
 from torch import Tensor
 from transformers.models.whisper.tokenization_whisper import LANGUAGES
-from typing_extensions import Self, TypeIs
+from typing_extensions import Self, TypeIs, deprecated
 
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -91,6 +91,13 @@
         tuple[int, int, int],
         tuple[int, int, int],
     ]
+    | tuple[
+        tuple[int, int],
+        tuple[int, int, int],
+        tuple[int, int, int],
+        tuple[int, int],
+        tuple[int, int, int],
+    ]
 )
 
 
@@ -112,8 +119,7 @@ def process_delta(self, text_delta: str, finished: bool) -> str:
 
 
 def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor:
-    """
-    A helper function to be used in the context of
+    """A helper function to be used in the context of
     [vllm.model_executor.models.interfaces.SupportsMultiModal.embed_input_ids][]
     to provide a better error message.
     """
@@ -225,20 +231,18 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):
 
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> str | None:
-        """
-        Get the placeholder text for the `i`th `modality` item in the prompt.
-        """
+        """Get the placeholder text for the `i`th `modality` item in the prompt."""
         ...
 
     def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
-        """
-        Returns multimodal embeddings generated from multimodal kwargs
+        """Returns multimodal embeddings generated from multimodal kwargs
         to be merged with text embeddings.
 
         Note:
             The returned multimodal embeddings must be in the same order as
             the appearances of their corresponding multimodal data item in the
             input prompt.
+
         """
         ...
 
@@ -254,14 +258,14 @@ def configure_mm_token_handling(self, vocab_size: int, mm_token_ids: list[int]):
         )
 
     def get_language_model(self) -> "VllmModel":
-        """
-        Returns the underlying language model used for text generation.
+        """Returns the underlying language model used for text generation.
 
         This is typically the `torch.nn.Module` instance responsible for
         processing the merged multimodal embeddings and producing hidden states
 
         Returns:
             torch.nn.Module: The core language model component.
+
         """
         # Cached
         assert isinstance(self, nn.Module)
@@ -302,8 +306,7 @@ def _mark_language_model(
         *,
         targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
     ):
-        """
-        Mark each child module that was assigned to this model during this context
+        """Mark each child module that was assigned to this model during this context
         as a language model component.
 
         Language model components are automatically skipped in `--mm-encoder-only`
@@ -338,8 +341,7 @@ def _mark_tower_model(
         *,
         targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
     ):
-        """
-        Mark each child module that was assigned to this model during this context
+        """Mark each child module that was assigned to this model during this context
         as a tower model component.
 
         Tower model components are automatically skipped when `--limit-mm-per-prompt`
@@ -397,8 +399,7 @@ def _mark_composite_model(
         language_targets: type[nn.Module] | tuple[type[nn.Module], ...],
         tower_targets: dict[str, type[nn.Module] | tuple[type[nn.Module], ...]],
     ):
-        """
-        Composite wrapper over `_mark_language_model` and
+        """Composite wrapper over `_mark_language_model` and
         `_mark_tower_model` by modality.
         """
         with ExitStack() as stack:
@@ -420,23 +421,16 @@ def _mark_composite_model(
 
             yield
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        """
-        Implement this function to enable LoRA support
-        for the tower module of the multi-modal model.
-        Given the number of image tokens, output the number of
-        multi-modal encoder tokens.
-        """
-        ...
+    @deprecated(
+        "get_num_mm_encoder_tokens is deprecated; use get_mm_lora_token_counts instead."
+    )
+    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: ...
 
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
-        """
-        Implement this function to enable LoRA support
-        for the connector module of the multi-modal model.
-        Given the number of vision tokens, output the number of
-        multi-modal connector tokens.
-        """
-        ...
+    @deprecated(
+        "get_num_mm_connector_tokens is deprecated; use "
+        "get_mm_lora_token_counts instead."
+    )
+    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: ...
 
     def get_mm_lora_token_counts(
         self,
@@ -445,8 +439,7 @@ def get_mm_lora_token_counts(
         mm_kwargs: "MultiModalKwargsItem | None",
         num_mm_embeds: int,
     ) -> tuple[int, int | None]:
-        """
-        Return ``(tower_tokens, connector_tokens)`` for multimodal LoRA mappings.
+        """Return ``(tower_tokens, connector_tokens)`` for multimodal LoRA mappings.
 
         MM LoRA uses these counts to build adapter mappings for the tower and
         connector forwards. Models with multiple modalities can override this
@@ -504,8 +497,7 @@ def embed_input_ids(
         *,
         is_multimodal: Tensor | None = None,
     ) -> Tensor:
-        """
-        Apply token embeddings to `input_ids`.
+        """Apply token embeddings to `input_ids`.
 
         If `multimodal_embeddings` is passed, scatter them into
         `input_ids` according to the mask `is_multimodal`.
@@ -558,8 +550,7 @@ def recompute_mrope_positions(
         mrope_positions: torch.LongTensor,
         num_computed_tokens: int,
     ) -> tuple[Sequence[torch.Tensor], Tensor, int]:
-        """
-        Update part of input mrope positions (starting with
+        """Update part of input mrope positions (starting with
         num_computed_tokens index). Original mrope_positions are computed
         for unpruned sequence and becomes incorrect once pruning occurs,
         so once we prune media tokens we should reflect this in the
@@ -578,6 +569,7 @@ def recompute_mrope_positions(
         Returns:
             Tuple of (multimodal_embeddings, mrope_positions,
                 mrope_position_delta).
+
         """
         ...
 
@@ -657,16 +649,12 @@ class SupportsScoreTemplate(Protocol):
 
     @classmethod
     def get_score_template(cls, query: str, document: str) -> str | None:
-        """
-        Generate a full prompt by populating the score template with query and document content.
-        """  # noqa: E501
+        """Generate a full prompt by populating the score template with query and document content."""  # noqa: E501
         ...
 
     @classmethod
     def post_process_tokens(cls, prompt: "TokensPrompt") -> None:
-        """
-        Perform architecture-specific manipulations on the input tokens.
-        """
+        """Perform architecture-specific manipulations on the input tokens."""
         ...
 
 
@@ -797,8 +785,7 @@ def forward(
         *,
         intermediate_tensors: "IntermediateTensors | None",
     ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]":
-        """
-        Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
+        """Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
         PP rank > 0.
 
         Return [`IntermediateTensors`][vllm.sequence.IntermediateTensors] only
@@ -960,6 +947,7 @@ def get_mamba_state_shape_from_config(
 
         Returns:
             Shapes for each state cache used by the model.
+
         """
         ...
 
@@ -973,6 +961,7 @@ def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]:
             (state, block_ids, cur_block_idx, num_accepted_tokens) and returns
             a MambaCopySpec describing the memory-copy parameters for prefix
             caching in align mode.
+
         """
         ...
 
@@ -1001,9 +990,7 @@ def is_hybrid(
 
 @runtime_checkable
 class MixtureOfExperts(Protocol):
-    """
-    Check if the model is a mixture of experts (MoE) model.
-    """
+    """Check if the model is a mixture of experts (MoE) model."""
 
     expert_weights: MutableSequence[Sequence[Tensor]]
     """
@@ -1046,8 +1033,7 @@ def set_eplb_state(
         logical_to_physical_map: Tensor,
         logical_replica_count: Tensor,
     ) -> None:
-        """
-        Register the EPLB state in the MoE model.
+        """Register the EPLB state in the MoE model.
 
         Since these are views of the actual EPLB state, any changes made by
         the EPLB algorithm are automatically reflected in the model's behavior
@@ -1061,6 +1047,7 @@ def set_eplb_state(
             expert_load_view: A view of the expert load metrics tensor.
             logical_to_physical_map: Mapping from logical to physical experts.
             logical_replica_count: Count of replicas for each logical expert.
+
         """
         self.expert_weights = []
         for layer_idx, layer in enumerate(self.moe_layers):
@@ -1093,8 +1080,8 @@ def get_mixture_of_experts_model(model: object) -> MixtureOfExperts | None:
 
     Returns:
         The MixtureOfExperts instance contained within the model, or None.
-    """
 
+    """
     if is_mixture_of_experts(model):
         return model
 
@@ -1366,8 +1353,7 @@ def get_other_languages(cls) -> Mapping[str, str]:
 
     @classmethod
     def validate_language(cls, language: str | None) -> str | None:
-        """
-        Ensure the language specified in the transcription request
+        """Ensure the language specified in the transcription request
         is a valid ISO 639-1 language code. If the request language is
         valid, but not natively supported by the model, trigger a
         warning (but not an exception).
@@ -1403,8 +1389,7 @@ def get_num_audio_tokens(
         stt_config: "SpeechToTextConfig",
         model_config: "ModelConfig",
     ) -> int | None:
-        """
-        Map from audio duration to number of audio tokens produced by the ASR
+        """Map from audio duration to number of audio tokens produced by the ASR
         model, without running a forward pass.
         This is used for estimating the amount of processing for this audio.
         """
@@ -1412,8 +1397,7 @@ def get_num_audio_tokens(
 
     @classmethod
     def post_process_output(cls, text: str) -> str:
-        """
-        Post-process the raw model output text.
+        """Post-process the raw model output text.
 
         Some ASR models output structured formats (e.g., language tags,
         special tokens) that need to be stripped before returning to the user.
@@ -1423,6 +1407,7 @@ def post_process_output(cls, text: str) -> str:
 
         Returns:
             Cleaned transcription text.
+
         """
         return text
 
@@ -1439,8 +1424,7 @@ def parse_diarized_transcript(cls, text: str) -> list[DiarizedTranscriptionSegme
     def get_streaming_post_processor_cls(
         cls,
     ) -> type[StreamingTranscriptionPostProcessor]:
-        """
-        Return a stateful post-processor class for streaming output deltas.
+        """Return a stateful post-processor class for streaming output deltas.
 
         Each instance receives the next decoded text delta and whether the
         request output is final. It returns the cleaned delta that should be
@@ -1694,12 +1678,12 @@ class SupportsEagle3(SupportsEagleBase, Protocol):
     """
 
     def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
-        """
-        Set which layers should output auxiliary hidden states for EAGLE-3.
+        """Set which layers should output auxiliary hidden states for EAGLE-3.
 
         Args:
             layers: Tuple of layer indices that should output auxiliary
                 hidden states.
+
         """
         parent_ref = self
         if hasattr(self, "get_language_model"):
@@ -1717,14 +1701,14 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         holder._set_aux_hidden_state_layers(layers)
 
     def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:
-        """
-        Get the default layer indices that should output auxiliary hidden states
+        """Get the default layer indices that should output auxiliary hidden states
         for EAGLE-3 for this model. Models can override this method to provide
         different default layers based on their architecture, but it is encouraged
         to instead include the layer specification in the model's config if possible.
 
         Returns:
             Tuple of layer indices for auxiliary hidden state outputs.
+
         """
         parent_ref = self
         if hasattr(self, "get_language_model"):
@@ -1772,8 +1756,7 @@ def get_mrope_input_positions(
         input_tokens: list[int],
         mm_features: list["MultiModalFeatureSpec"],
     ) -> tuple[torch.Tensor, int]:
-        """
-        Get M-RoPE input positions and delta value for this specific model.
+        """Get M-RoPE input positions and delta value for this specific model.
 
         This method should be implemented by each model that supports M-RoPE
         to provide model-specific logic for computing input positions.
@@ -1786,6 +1769,7 @@ def get_mrope_input_positions(
             llm_positions: Tensor of shape `[num_dims, num_tokens]`, one row
                 per M-RoPE position channel (e.g. T/H/W).
             mrope_position_delta: Delta for position calculations.
+
         """
         ...
 
@@ -1893,8 +1877,7 @@ def postprocess_encoder_output(
         clone: bool = False,
         batch_mm_kwargs: dict[str, Any] | None = None,
     ) -> None:
-        """
-        Post-process encoder output, directly call scatter_output_slices by default.
+        """Post-process encoder output, directly call scatter_output_slices by default.
 
         By default, delegates directly to scatter_output_slices.
         Override this for models that require additional processing on the raw
@@ -1920,10 +1903,17 @@ def prepare_encoder_cudagraph_capture_inputs(
         """Create dummy inputs and buffers for CUDA graph capture.
 
         Args:
+            token_budget: Token budget the capture is sized for.
+            max_batch_size: Maximum number of items in a captured batch.
+            max_frames_per_batch: Maximum number of frames in a captured batch.
+            device: Device the dummy inputs and buffers are created on.
+            dtype: Dtype of the dummy inputs and buffers.
+            path: Configured encoder path.
             axis_keys: The resolved capture-axis keys (one per axis of
                 ``EncoderCudaGraphConfig.capture_axes``) this capture is for.
                 None or empty when no capture axes are configured; models
                 without capture axes ignore it.
+
         """
         ...
 
diff --git a/vllm/model_executor/models/intern_vit.py b/vllm/model_executor/models/intern_vit.py
index 44ed3ec0fed8..873160f9687c 100644
--- a/vllm/model_executor/models/intern_vit.py
+++ b/vllm/model_executor/models/intern_vit.py
@@ -13,7 +13,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import (
     should_torch_compile_mm_encoder,
@@ -47,7 +47,7 @@
 
 
 class InternVisionEmbeddings(nn.Module):
-    def __init__(self, config: PretrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
         self.config = config
         self.embed_dim = config.hidden_size
@@ -114,40 +114,12 @@ def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
         return embeddings
 
 
-class InternVisionPatchModel(nn.Module):
-    def __init__(self, config: PretrainedConfig):
-        super().__init__()
-        self.config = config
-        self.embeddings = InternVisionEmbeddings(config)
-
-    def get_input_embeddings(self):
-        return self.embeddings
-
-    def forward(
-        self,
-        pixel_values: torch.Tensor | None = None,
-        pixel_embeds: torch.Tensor | None = None,
-    ) -> torch.FloatTensor:
-        if pixel_values is None and pixel_embeds is None:
-            raise ValueError("You have to specify pixel_values or pixel_embeds")
-
-        if pixel_embeds is not None:
-            hidden_states = pixel_embeds
-        elif pixel_values is not None:
-            if pixel_values.ndim == 4:
-                hidden_states = self.embeddings(pixel_values)
-            else:
-                raise ValueError(f"wrong pixel_values size: {pixel_values.shape}")
-
-        return hidden_states
-
-
 class InternParallelAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_dummy_heads: int = 0,
@@ -250,7 +222,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 class InternMLP(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
@@ -292,7 +264,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class InternVisionEncoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_dummy_heads: int = 0,
@@ -325,7 +297,7 @@ def __init__(
 
     def _init_attn(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         *,
         num_dummy_heads: int,
@@ -352,7 +324,7 @@ def forward(
 class InternVisionEncoder(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_hidden_layers_override: int | None = None,
@@ -397,7 +369,7 @@ class InternVisionModel(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_hidden_layers_override: int | None = None,
diff --git a/vllm/model_executor/models/internlm2.py b/vllm/model_executor/models/internlm2.py
index 29dd124b8663..e1490071c2f1 100644
--- a/vllm/model_executor/models/internlm2.py
+++ b/vllm/model_executor/models/internlm2.py
@@ -8,7 +8,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -195,7 +195,7 @@ def forward(
 class InternLMDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/interns1.py b/vllm/model_executor/models/interns1.py
index 0ed2723c32b1..b1ac9596febb 100644
--- a/vllm/model_executor/models/interns1.py
+++ b/vllm/model_executor/models/interns1.py
@@ -12,17 +12,17 @@
 import regex as re
 import torch
 import torch.nn as nn
-from transformers import BatchFeature, InternVLProcessor, PretrainedConfig
+from transformers import BatchFeature, InternVLProcessor, PreTrainedConfig
 from transformers.activations import ACT2FN
-from transformers.models.got_ocr2.image_processing_got_ocr2_fast import (
-    GotOcr2ImageProcessorFast,
+from transformers.models.got_ocr2.image_processing_got_ocr2 import (
+    GotOcr2ImageProcessor,
 )
 from transformers.models.internvl.video_processing_internvl import (
     InternVLVideoProcessor,
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.models.interns1_vit import InternS1VisionModel
@@ -89,13 +89,12 @@ def forward(self, image_features):
 
 
 class InternS1ImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * number of images * (1 + num_patches)
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
-        - bn: Batch size * number of images
+    """Dimensions:
+    - bnp: Batch size * number of images * (1 + num_patches)
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
+    - bn: Batch size * number of images
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -104,11 +103,10 @@ class InternS1ImagePixelInputs(TensorSchema):
 
 
 class InternS1ImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - ni: Number of images
-        - tifs: Total image feature size
-        - hs: Hidden size (must match language model backbone)
+    """Dimensions:
+    - ni: Number of images
+    - tifs: Total image feature size
+    - hs: Hidden size (must match language model backbone)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -119,13 +117,12 @@ class InternS1ImageEmbeddingInputs(TensorSchema):
 
 
 class InternS1VideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnv: Batch size * number of videos * number of frames
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
+    """Dimensions:
+    - bnv: Batch size * number of videos * number of frames
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
     """
 
     type: Literal["pixel_values_videos"] = "pixel_values_videos"
@@ -134,11 +131,10 @@ class InternS1VideoPixelInputs(TensorSchema):
 
 
 class InternS1VideoEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - nv: Number of videos
-        - tvfs: Total video feature size
-        - hs: Hidden size (must match language model backbone)
+    """Dimensions:
+    - nv: Number of videos
+    - tvfs: Total video feature size
+    - hs: Hidden size (must match language model backbone)
     """
 
     type: Literal["video_embeds"] = "video_embeds"
@@ -206,7 +202,7 @@ def get_num_image_tokens(
         processor: InternVLProcessor,
         mm_kwargs: Mapping[str, object],
     ) -> int:
-        image_processor: GotOcr2ImageProcessorFast = processor.image_processor
+        image_processor: GotOcr2ImageProcessor = processor.image_processor
 
         num_image_patches = image_processor.get_number_of_image_patches(
             image_height,
@@ -303,35 +299,29 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(
             seq_len, mm_counts
         )
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
 
         config = self.info.get_hf_config()
         image_size_h, image_size_w = config.vision_config.image_size
 
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-        assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=image_size_w,
                 height=image_size_h,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -339,28 +329,30 @@ def get_dummy_mm_data(
 class InternS1MultiModalProcessor(BaseMultiModalProcessor[InternS1ProcessingInfo]):
     """Basic image-only MultiModalProcessor for InternS1-style models."""
 
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        if not hf_data:
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        prompt_text = hf_data.pop("text")
+        assert isinstance(prompt_text, str)
 
-        mm_data = dict(mm_data)
-        videos = mm_data.pop("videos", [])
-        images = mm_data.pop("images", [])
+        videos = hf_data.pop("videos", [])
+        images = hf_data.pop("images", [])
         assert isinstance(videos, list)
         assert isinstance(images, list)
 
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        hf_processor = self.info.get_hf_processor(**hf_kwargs)
         tokenizer = hf_processor.tokenizer
         vocab = tokenizer.get_vocab()
 
@@ -378,9 +370,9 @@ def _apply_hf_processor_main(
             image_pixel_values = []
             for image in images:
                 processed_data = self.info.ctx.call_hf_processor(
-                    self.info.get_hf_processor(**hf_processor_mm_kwargs),
+                    self.info.get_hf_processor(**hf_kwargs),
                     dict(text=hf_processor.image_token, **{"images": image}),
-                    hf_processor_mm_kwargs,
+                    hf_kwargs,
                 )
                 image_pixel_values.append(processed_data.pop("pixel_values"))
 
@@ -402,9 +394,9 @@ def _apply_hf_processor_main(
             video_pixel_values = []
             for video in videos:
                 processed_data = self.info.ctx.call_hf_processor(
-                    self.info.get_hf_processor(**hf_processor_mm_kwargs),
+                    self.info.get_hf_processor(**hf_kwargs),
                     dict(text=hf_processor.video_token, **{"videos": video}),
-                    hf_processor_mm_kwargs,
+                    hf_kwargs,
                 )
                 video_pixel_values.append(processed_data.pop("pixel_values"))
 
@@ -434,9 +426,9 @@ def _apply_hf_processor_main(
         processed_data.update(text_outputs)
         processed_data.update(image_outputs)
         processed_data.update(video_outputs)
-        processed_data.update(passthrough_data)
-
-        return processed_data
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_mm_fields_config(
         self,
@@ -620,7 +612,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
 
     def _init_vision_model(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
@@ -633,7 +625,7 @@ def _init_vision_model(
             prefix=prefix,
         )
 
-    def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
+    def _init_mlp1(self, config: PreTrainedConfig) -> nn.Module:
         return InternS1MultiModalProjector(config)
 
     def pixel_shuffle(self, x, scale_factor=0.5):
@@ -862,9 +854,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py
index 40c5330ad320..e757241e9458 100644
--- a/vllm/model_executor/models/interns1_pro.py
+++ b/vllm/model_executor/models/interns1_pro.py
@@ -30,7 +30,7 @@
 
 import torch
 from torch import nn
-from transformers import AutoProcessor, PretrainedConfig
+from transformers import AutoProcessor, PreTrainedConfig
 
 from vllm.config import CacheConfig, VllmConfig
 from vllm.distributed import (
@@ -43,12 +43,12 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -181,10 +181,9 @@ def __init__(
             custom_routing_function=self._custom_routing_function,
         )
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
             prefix=f"{prefix}.gate",
         )
 
@@ -564,7 +563,7 @@ class InternS1ProForConditionalGeneration(
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super(Qwen3VLForConditionalGeneration, self).__init__()
-        config: PretrainedConfig = vllm_config.model_config.hf_config
+        config: PreTrainedConfig = vllm_config.model_config.hf_config
         multimodal_config = vllm_config.model_config.get_multimodal_config()
 
         self.config = config
@@ -621,7 +620,7 @@ def get_frope_params_map(self) -> dict[str, str]:
         return mapper
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        """load weights"""
+        """Load weights"""
         orig_to_new_prefix: dict[str, str | None] = {
             "model.visual.": "visual.",
             "lm_head.": "language_model.lm_head.",
diff --git a/vllm/model_executor/models/interns1_vit.py b/vllm/model_executor/models/interns1_vit.py
index a5bc781561ff..f6337339379f 100644
--- a/vllm/model_executor/models/interns1_vit.py
+++ b/vllm/model_executor/models/interns1_vit.py
@@ -11,7 +11,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.utils import torch_int
 
 from vllm.model_executor.layers.activation import get_act_fn
@@ -65,7 +65,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
 
 
 class InternS1VisionEmbeddings(nn.Module):
-    def __init__(self, config: PretrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
         self.config = config
         self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
@@ -91,15 +91,13 @@ def __init__(self, config: PretrainedConfig):
     def interpolate_pos_encoding(
         self, embeddings: torch.Tensor, height: int, width: int
     ) -> torch.Tensor:
-        """
-        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
+        """This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
         images. This method is also adapted to support torch.jit tracing.
 
         Adapted from:
         - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
         - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
         """  # noqa: E501
-
         num_patches = embeddings.shape[1] - 1
         num_positions = self.position_embeddings.shape[1] - 1
 
@@ -164,11 +162,11 @@ def forward(
 
 
 class InternSdpaAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         *,
         num_dummy_heads: int = 0,
         prefix: str = "",
@@ -225,8 +223,7 @@ def __init__(
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """x shape: (B, N, C)"""
-
+        """X shape: (B, N, C)."""
         q = self.q_proj(x)
         k = self.k_proj(x)
         v = self.v_proj(x)
@@ -245,7 +242,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 class InternS1VisionMLP(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
@@ -279,7 +276,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class InternS1VisionLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_dummy_heads: int = 0,
@@ -314,7 +311,7 @@ def __init__(
 
     def _init_attn(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         *,
         num_dummy_heads: int,
@@ -346,7 +343,7 @@ def forward(
 class InternS1VisionEncoder(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_hidden_layers_override: int | None = None,
@@ -385,7 +382,7 @@ def forward(self, inputs_embeds: torch.Tensor):
 class InternS1VisionModel(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_hidden_layers_override: int | None = None,
diff --git a/vllm/model_executor/models/interns2_mobius.py b/vllm/model_executor/models/interns2_mobius.py
index 4efaa5a08eba..c3fece062f34 100644
--- a/vllm/model_executor/models/interns2_mobius.py
+++ b/vllm/model_executor/models/interns2_mobius.py
@@ -17,7 +17,10 @@
     tensor_model_parallel_all_gather,
     tensor_model_parallel_reduce_scatter,
 )
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
@@ -93,11 +96,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
         self.n_local_physical_experts = self.n_physical_experts // self.ep_size
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
         self.experts = FusedMoEFactory(
diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py
index 1fcb6b8681d6..ec3b39b6ad84 100644
--- a/vllm/model_executor/models/internvl.py
+++ b/vllm/model_executor/models/internvl.py
@@ -14,10 +14,10 @@
 
 import torch
 import torch.nn as nn
-from transformers import BatchFeature, PretrainedConfig
+from transformers import BatchFeature, PreTrainedConfig
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
@@ -29,6 +29,7 @@
 from vllm.multimodal.inputs import (
     BatchedTensorInputs,
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import (
@@ -69,13 +70,12 @@
 
 
 class InternVLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - bnp: Batch size * number of images * (1 + num_patches)
-        - c: Number of channels (3)
-        - h: Height of each image patch
-        - w: Width of each image patch
+    """Dimensions:
+    - bn: Batch size * number of images
+    - bnp: Batch size * number of images * (1 + num_patches)
+    - c: Number of channels (3)
+    - h: Height of each image patch
+    - w: Width of each image patch
     """
 
     type: Literal["pixel_values"]
@@ -84,11 +84,10 @@ class InternVLImagePixelInputs(TensorSchema):
 
 
 class InternVLImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - n: Number of images
-        - f: Total image feature size
-        - h: Hidden size (must match the hidden size of language model backbone)
+    """Dimensions:
+    - n: Number of images
+    - f: Total image feature size
+    - h: Hidden size (must match the hidden size of language model backbone)
     """
 
     type: Literal["image_embeds"]
@@ -99,13 +98,12 @@ class InternVLImageEmbeddingInputs(TensorSchema):
 
 
 class InternVLVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bvf: Batch size * number of videos * num_frames
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height of each video frame
-        - w: Width of each video frame
+    """Dimensions:
+    - bvf: Batch size * number of videos * num_frames
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height of each video frame
+    - w: Width of each video frame
     """
 
     type: Literal["pixel_values_videos"]
@@ -114,11 +112,10 @@ class InternVLVideoPixelInputs(TensorSchema):
 
 
 class InternVLVideoEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - n: Number of videos
-        - f: Total video feature size
-        - h: Hidden size (must match the hidden size of language model backbone)
+    """Dimensions:
+    - n: Number of videos
+    - f: Total video feature size
+    - h: Hidden size (must match the hidden size of language model backbone)
     """
 
     type: Literal["video_embeds"]
@@ -206,19 +203,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         target_width, target_height = self.info.get_image_size_with_most_features()
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
 
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -326,7 +320,7 @@ def _get_prompt_updates(
 
 
 class InternVLProcessingInfo(BaseInternVLProcessingInfo):
-    """InternVL ProcessingInfo extended for video processing"""
+    """InternVL ProcessingInfo extended for video processing."""
 
     def get_image_processor(self, **kwargs):
         config = self.get_hf_config()
@@ -417,7 +411,7 @@ def get_num_frames_with_most_features(
 class InternVLDummyInputsBuilder(
     BaseInternVLDummyInputsBuilder[InternVLProcessingInfo]
 ):
-    """InternVL DummyInputsBuilder extended for video support"""
+    """InternVL DummyInputsBuilder extended for video support."""
 
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_videos = mm_counts.get("video", 0)
@@ -428,7 +422,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         dummy_image = super().get_dummy_mm_data(seq_len, mm_counts, mm_options)
         if self.info.ctx_video_token:
@@ -439,9 +433,6 @@ def get_dummy_mm_data(
             )
             num_videos = mm_counts.get("video", 0)
             video_overrides = mm_options.get("video")
-            assert video_overrides is None or isinstance(
-                video_overrides, VideoDummyOptions
-            )
             dummy_video = {
                 "video": self._get_dummy_videos(
                     width=image_size,
@@ -459,7 +450,7 @@ def get_dummy_mm_data(
 class InternVLMultiModalProcessor(
     BaseInternVLMultiModalProcessor[InternVLProcessingInfo]
 ):
-    """InternVL MultiModalProcessor extended for video support"""
+    """InternVL MultiModalProcessor extended for video support."""
 
     def _postprocess_hf_mm_data(
         self,
@@ -638,7 +629,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         )
 
     def _patch_quant_config(
-        self, config: PretrainedConfig, quant_config: QuantizationConfig | None
+        self, config: PreTrainedConfig, quant_config: QuantizationConfig | None
     ):
         # the awq models from OpenGVLab missing `modules_to_not_convert`
         # patch the quant_config to add `modules_to_not_convert` back
@@ -652,7 +643,7 @@ def _patch_quant_config(
 
     def _init_vision_model(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
@@ -672,7 +663,7 @@ def _init_vision_model(
             prefix=prefix,
         )
 
-    def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
+    def _init_mlp1(self, config: PreTrainedConfig) -> nn.Module:
         vit_hidden_size = config.vision_config.hidden_size
         llm_hidden_size = config.text_config.hidden_size
 
@@ -905,28 +896,27 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="mlp1",
             tower_model="vision_model",
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        if num_image_tokens <= 0 or self.num_image_token <= 0:
-            return 0
-
-        num_patches = num_image_tokens // self.num_image_token
-        return num_patches * (self.patch_tokens + 1)
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
-        if num_vision_tokens <= 0 or self.num_image_token <= 0:
-            return 0
-
-        num_patches = num_vision_tokens // (self.patch_tokens + 1)
-        return num_patches * self.num_image_token
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+        if num_mm_embeds <= 0 or self.num_image_token <= 0:
+            return 0, 0
+
+        num_patches = num_mm_embeds // self.num_image_token
+        tower_tokens = num_patches * (self.patch_tokens + 1)
+        return tower_tokens, num_patches * self.num_image_token
 
     # -- SupportsEncoderCudaGraph protocol methods --
 
diff --git a/vllm/model_executor/models/iquest_loopcoder.py b/vllm/model_executor/models/iquest_loopcoder.py
index 782bec0ed0df..5d39e87564bb 100644
--- a/vllm/model_executor/models/iquest_loopcoder.py
+++ b/vllm/model_executor/models/iquest_loopcoder.py
@@ -22,7 +22,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -60,7 +60,7 @@
 class LoopCoderAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
@@ -210,7 +210,7 @@ def forward(
 class LoopCoderDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -321,6 +321,7 @@ def forward(self, query: torch.Tensor) -> torch.Tensor:
 
         Returns:
             gate: [num_tokens, num_heads * head_dim] (flattened format matching q shape)
+
         """
         num_heads, num_tokens, head_dim = query.shape
 
diff --git a/vllm/model_executor/models/isaac.py b/vllm/model_executor/models/isaac.py
index a59345dc25f5..f8f6dee32c96 100644
--- a/vllm/model_executor/models/isaac.py
+++ b/vllm/model_executor/models/isaac.py
@@ -13,7 +13,7 @@
 from transformers.image_processing_utils import BatchFeature
 
 from vllm.config import ModelConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import parallel_state
 from vllm.distributed import utils as dist_utils
 from vllm.inputs import MultiModalDataDict
@@ -177,23 +177,23 @@ def create_pixel_shuffle_index_map(
     scale_factor: int = 1,
     device: torch.device | None = None,
 ) -> torch.Tensor:
-    """
-    Build a gather-index map that tells us, for every *output* token after
+    """Build a gather-index map that tells us, for every *output* token after
     pixel-shuffle, which `scale_factor**2` *input* tokens are being merged.
 
-    Args
+    Args:
     ----
     seq_sizes     : (num_images,)  - #patches in each image (row-major order)
     token_grids   : (num_images,2) - (height, width) for every image
     scale_factor  : spatial down-scale factor (≥2)
     device        : (optional) overrides `seq_sizes.device`
 
-    Returns
+    Returns:
     -------
     gather_idx : (new_total_seq_len, scale_factor**2) int64 tensor.
                  gather_idx[i, j] is the *flat* index into the *original*
                  packed sequence for the j-th sub-patch that forms the
                  i-th output token.
+
     """
     if device is None:
         device = seq_sizes.device
@@ -266,6 +266,7 @@ def pixel_shuffle_varlen(
 
     Raises:
         ValueError: If more than one batch item is provided.
+
     """
     keep_batch_dim = x.dim() == 3
     if keep_batch_dim:
@@ -385,26 +386,22 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
-        image_overrides = mm_options.get("image")
 
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
         }
 
 
 class IsaacImagePixelInputs(TensorSchema):
-    """
-    Schema for validating Isaac image inputs.
+    """Schema for validating Isaac image inputs.
 
     Dimensions:
         - np: Number of patches
@@ -687,12 +684,10 @@ def forward(
         self,
         packed_seq_patches: tuple[torch.Tensor, torch.Tensor],
     ) -> torch.Tensor:
-        r"""
-        spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
-            Tensor containing the spatial dimensions (height, width)
-            of the input images.
+        r"""spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
+        Tensor containing the spatial dimensions (height, width)
+        of the input images.
         """
-
         seq_patches, token_grids = packed_seq_patches
         seq_sizes = torch.prod(token_grids, dim=-1)
 
@@ -1015,9 +1010,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="vision_embedding.linear_fc2",  # The final linear layer
diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py
index 44ea28986981..fe9349c30b10 100644
--- a/vllm/model_executor/models/jamba.py
+++ b/vllm/model_executor/models/jamba.py
@@ -16,11 +16,11 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -74,11 +74,9 @@ def __init__(
         self.intermediate_size = config.intermediate_size
 
         if self.num_total_experts > 1:
-            self.router = ReplicatedLinear(
+            self.router = GateLinear(
                 self.hidden_size,
                 self.num_total_experts,
-                bias=False,
-                quant_config=None,
                 params_dtype=params_dtype,
                 prefix=f"{prefix}.router",
             )
diff --git a/vllm/model_executor/models/jina_vl.py b/vllm/model_executor/models/jina_vl.py
index 2895ec13b925..971284561926 100644
--- a/vllm/model_executor/models/jina_vl.py
+++ b/vllm/model_executor/models/jina_vl.py
@@ -13,7 +13,7 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import MultiModalKwargsItems
 from vllm.multimodal.processing.processor import (
-    MultiModalProcessingInfo,
+    MultiModalProcessingResult,
     ProcessorInputs,
     TimingContext,
 )
@@ -63,19 +63,19 @@ def _cached_apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
-        mm_info = super()._cached_apply_hf_processor(inputs, timing_ctx)
+    ) -> MultiModalProcessingResult:
+        mm_res = super()._cached_apply_hf_processor(inputs, timing_ctx)
 
         # Score inputs are query-first, while the prompt template is document-first.
         mm_kwargs = MultiModalKwargsItems(
             {
                 modality: list(reversed(items))
-                for modality, items in mm_info.kwargs.items()
+                for modality, items in mm_res.kwargs.items()
             }
         )
         mm_hashes = {
             modality: list(reversed(hashes))
-            for modality, hashes in mm_info.hashes.items()
+            for modality, hashes in mm_res.hashes.items()
         }
         mm_prompt_updates = {
             modality: [
@@ -85,10 +85,11 @@ def _cached_apply_hf_processor(
                 ]
                 for item_idx, updates in enumerate(reversed(item_updates))
             ]
-            for modality, item_updates in mm_info.prompt_updates.items()
+            for modality, item_updates in mm_res.prompt_updates.items()
         }
 
-        return MultiModalProcessingInfo(
+        return MultiModalProcessingResult(
+            prompt_ids=mm_res.prompt_ids,
             kwargs=mm_kwargs,
             hashes=mm_hashes,
             prompt_updates=mm_prompt_updates,
diff --git a/vllm/model_executor/models/k2_horizon.py b/vllm/model_executor/models/k2_horizon.py
index 47e87c0dcb62..db435eaf0564 100644
--- a/vllm/model_executor/models/k2_horizon.py
+++ b/vllm/model_executor/models/k2_horizon.py
@@ -49,6 +49,7 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.config import (
@@ -398,7 +399,7 @@ def __init__(
             self.physical_expert_start + self.n_local_physical_experts
         )
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
             bias=config.moe_gate_bias,
diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py
index c17228aa50fd..fd076d66e7b8 100644
--- a/vllm/model_executor/models/kanana_v.py
+++ b/vllm/model_executor/models/kanana_v.py
@@ -18,7 +18,7 @@
 from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLVisionConfig
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -48,11 +48,10 @@
 
 
 class KananaVImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - np: The total number of patches over all images in the batch
-        - cps: Number of channels * patch_size * patch_size
-        - ni: Number of images
+    """Dimensions:
+    - np: The total number of patches over all images in the batch
+    - cps: Number of channels * patch_size * patch_size
+    - ni: Number of images
     """
 
     type: Literal["pixel_values"]
@@ -443,7 +442,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         num_images = mm_counts.get("image", 0)
         return {
@@ -460,24 +459,27 @@ class KananaVMultiModalProcessor(BaseMultiModalProcessor[KananaVProcessingInfo])
     def media_token_id(self) -> int:
         return self.info.get_hf_config().text_config.eos_token_id + 1
 
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         """Run the underlying HF processor on text and image data."""
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
-        if not mm_data or not mm_data.get("images", []):
-            return BatchFeature(dict(passthrough_data))
+        if not hf_data or not hf_data.get("images", []):
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        prompt_text = hf_data.pop("text")
+        assert isinstance(prompt_text, str)
 
         # Images
-        image_inputs = mm_data.get("images", [])
+        image_inputs = hf_data.get("images", [])
         pixel_sizes = []
         if not isinstance(image_inputs[0], Image.Image):
             image_inputs = [Image.fromarray(image) for image in image_inputs]
@@ -531,8 +533,9 @@ def _apply_hf_processor_main(
             pixel_sizes=torch.tensor(pixel_sizes),
         )
         processed_data = BatchFeature(combined_outputs, tensor_type="pt")
-        processed_data.update(passthrough_data)
-        return processed_data
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_prompt_updates(
         self,
diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py
index 53cda0ed171b..7bf5d797096e 100644
--- a/vllm/model_executor/models/keye.py
+++ b/vllm/model_executor/models/keye.py
@@ -10,14 +10,14 @@
 import torch
 import torch.nn as nn
 from einops import rearrange
-from transformers import BaseImageProcessor, PretrainedConfig
+from transformers import BaseImageProcessor, PreTrainedConfig
 from transformers.activations import GELUActivation
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
 from transformers.utils import torch_int
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.logger import init_logger
@@ -40,6 +40,7 @@
     ImageItem,
     MultiModalFeatureSpec,
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
     VideoItem,
 )
@@ -124,13 +125,12 @@ def smart_resize(
 
 
 class KeyeImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * Number of patches
-        - c: Number of channels
-        - ps: Patch size
-        - ni: Number of images
-        - g: Grid dimensions (3 for t, h, w)
+    """Dimensions:
+    - bnp: Batch size * Number of patches
+    - c: Number of channels
+    - ps: Patch size
+    - ni: Number of images
+    - g: Grid dimensions (3 for t, h, w)
     """
 
     type: Literal["pixel_values"]
@@ -141,13 +141,12 @@ class KeyeImagePixelInputs(TensorSchema):
 
 
 class KeyeImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - nf: Number of image features
-        - hs: Hidden size (must match the hidden size of language model
-          backbone)
-        - ni: Number of images
-        - g: Grid dimensions (3 for t, h, w)
+    """Dimensions:
+    - nf: Number of image features
+    - hs: Hidden size (must match the hidden size of language model
+      backbone)
+    - ni: Number of images
+    - g: Grid dimensions (3 for t, h, w)
     """
 
     type: Literal["image_embeds"]
@@ -159,13 +158,12 @@ class KeyeImageEmbeddingInputs(TensorSchema):
 
 
 class KeyeVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * Number of patches
-        - c: Number of channels
-        - ps: Patch size
-        - ni: Number of images
-        - g: Grid dimensions (3 for t, h, w)
+    """Dimensions:
+    - bnp: Batch size * Number of patches
+    - c: Number of channels
+    - ps: Patch size
+    - ni: Number of images
+    - g: Grid dimensions (3 for t, h, w)
     """
 
     type: Literal["pixel_values_videos"]
@@ -176,13 +174,12 @@ class KeyeVideoPixelInputs(TensorSchema):
 
 
 class KeyeVideoEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - nf: Number of video features
-        - hs: Hidden size (must match the hidden size of language model
-          backbone)
-        - nv: Number of videos
-        - g: Grid dimensions (3 for t, h, w)
+    """Dimensions:
+    - nf: Number of video features
+    - hs: Hidden size (must match the hidden size of language model
+      backbone)
+    - nv: Number of videos
+    - g: Grid dimensions (3 for t, h, w)
     """
 
     type: Literal["video_embeds"]
@@ -194,7 +191,7 @@ class KeyeVideoEmbeddingInputs(TensorSchema):
 
 
 class KeyeVisionEmbeddings(nn.Module):
-    def __init__(self, config: PretrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
         self.config = config
         self.embed_dim = config.hidden_size
@@ -352,7 +349,7 @@ class KeyeSiglipAttention(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -490,7 +487,7 @@ def forward(self, seqlen: int) -> torch.Tensor:
 class KeyeSiglipEncoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -542,7 +539,7 @@ def forward(
 class KeyeSiglipEncoder(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -637,7 +634,7 @@ def forward(
 class KeyeSiglipVisionTransformer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -713,7 +710,7 @@ def forward(
 
 
 class KeyeSiglipVisionModel(nn.Module):
-    config_class = PretrainedConfig
+    config_class = PreTrainedConfig
     main_input_name = "pixel_values"
 
     hf_to_vllm_mapper = WeightsMapper(
@@ -727,7 +724,7 @@ class KeyeSiglipVisionModel(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -790,8 +787,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 class Projector(nn.Module):
     def __init__(
         self,
-        text_config: PretrainedConfig,
-        vision_config: PretrainedConfig,
+        text_config: PreTrainedConfig,
+        vision_config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -1127,30 +1124,24 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(seq_len)
 
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-
         mm_data = {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -1162,7 +1153,7 @@ class KeyeDummyInputsBuilder(KeyeBaseDummyInputsBuilder[KeyeProcessingInfo]):
 
 
 class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _get_prompt_updates(
@@ -1241,7 +1232,7 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        config: PretrainedConfig = vllm_config.model_config.hf_config
+        config: PreTrainedConfig = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
 
         self.config = config
@@ -1273,8 +1264,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     @abstractmethod
     def _build_projector(
         self,
-        text_config: PretrainedConfig,
-        vision_config: PretrainedConfig,
+        text_config: PreTrainedConfig,
+        vision_config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> nn.Module:
@@ -1464,6 +1455,9 @@ def forward(
                 otherwise it will be `(seq_len,)`.
             intermediate_tensors: Intermediate tensors from prior forward pass.
             inputs_embeds: Optional tensor of input embeddings.
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
+
         """
         if intermediate_tensors is not None:
             inputs_embeds = None
@@ -1495,13 +1489,16 @@ def get_mm_mapping(self) -> MultiModelKeys:
             tower_model="visual.",
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        merge_size = self.config.vision_config.spatial_merge_size
-        return num_image_tokens * merge_size**2
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         merge_size = self.config.vision_config.spatial_merge_size
-        return num_vision_tokens // merge_size**2
+        return num_mm_embeds * merge_size**2, num_mm_embeds
 
 
 @MULTIMODAL_REGISTRY.register_processor(
@@ -1514,8 +1511,8 @@ class KeyeForConditionalGeneration(
 ):
     def _build_projector(
         self,
-        text_config: PretrainedConfig,
-        vision_config: PretrainedConfig,
+        text_config: PreTrainedConfig,
+        vision_config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> nn.Module:
@@ -1584,8 +1581,7 @@ def _process_video_input(
     def _split_video_grid_thw(
         grid_thw: torch.Tensor | list[list[int]] | list[int],
     ) -> list[list[int]]:
-        """
-        Split video grid_thw along the t dimension into per-frame rows.
+        """Split video grid_thw along the t dimension into per-frame rows.
 
         This preserves Keye's current M-RoPE behavior, where a video is emitted
         as consecutive frame-level multimodal blocks rather than a single block
diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py
index 8625fd9c0188..3bc7f364205d 100644
--- a/vllm/model_executor/models/keye_vl1_5.py
+++ b/vllm/model_executor/models/keye_vl1_5.py
@@ -9,7 +9,7 @@
 import torch
 import torch.nn as nn
 from einops import rearrange
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.activations import GELUActivation
 from transformers.feature_extraction_utils import BatchFeature
 
@@ -51,8 +51,7 @@
 
 
 def split_thw(grid_thw: torch.Tensor) -> torch.Tensor:
-    """
-    Split grid_thw in t dimension.
+    """Split grid_thw in t dimension.
 
     Args:
         grid_thw: [N, 3] tensor of [t, h, w]
@@ -66,6 +65,7 @@ def split_thw(grid_thw: torch.Tensor) -> torch.Tensor:
     tensor([[1, 3, 4],
            [1, 3, 4],
            [1, 5, 6]])
+
     """
     t = grid_thw[:, 0]
     h_w = grid_thw[:, 1:]
@@ -76,8 +76,7 @@ def split_thw(grid_thw: torch.Tensor) -> torch.Tensor:
 def get_num_patches(
     grid_thw: torch.Tensor, num_frames: list[int] | torch.Tensor
 ) -> list[int]:
-    """
-    Return num_patches per video.
+    """Return num_patches per video.
 
     Args:
         grid_thw: Tensor with shape [N, 3] containing temporal, height, width
@@ -101,8 +100,8 @@ def get_num_patches(
         >>> get_num_patches(grid_thw, num_frames)
         tensor([16, 1])  # Total patches for first video: 8+8=16,
                            second video: 1.
-    """
 
+    """
     assert len(grid_thw.shape) == 2
     if isinstance(num_frames, torch.Tensor):
         num_frames = num_frames.clone().tolist()
@@ -121,13 +120,12 @@ def get_num_patches(
 
 
 class KeyeVL1_5ImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * Number of patches
-        - c: Number of channels
-        - ps: Patch size
-        - ni: Number of images
-        - g: Grid dimensions (3 for t, h, w)
+    """Dimensions:
+    - bnp: Batch size * Number of patches
+    - c: Number of channels
+    - ps: Patch size
+    - ni: Number of images
+    - g: Grid dimensions (3 for t, h, w)
     """
 
     type: Literal["pixel_values"]
@@ -140,13 +138,12 @@ class KeyeVL1_5ImagePixelInputs(TensorSchema):
 
 
 class KeyeVL1_5ImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - nf: Number of image features
-        - hs: Hidden size (must match the hidden size of language model
-          backbone)
-        - ni: Number of images
-        - g: Grid dimensions (3 for t, h, w)
+    """Dimensions:
+    - nf: Number of image features
+    - hs: Hidden size (must match the hidden size of language model
+      backbone)
+    - ni: Number of images
+    - g: Grid dimensions (3 for t, h, w)
     """
 
     type: Literal["image_embeds"]
@@ -160,13 +157,12 @@ class KeyeVL1_5ImageEmbeddingInputs(TensorSchema):
 
 
 class KeyeVL1_5VideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * Number of patches
-        - c: Number of channels
-        - ps: Patch size
-        - ni: Number of images
-        - g: Grid dimensions (3 for t, h, w)
+    """Dimensions:
+    - bnp: Batch size * Number of patches
+    - c: Number of channels
+    - ps: Patch size
+    - ni: Number of images
+    - g: Grid dimensions (3 for t, h, w)
     """
 
     type: Literal["pixel_values_videos"]
@@ -179,13 +175,12 @@ class KeyeVL1_5VideoPixelInputs(TensorSchema):
 
 
 class KeyeVL1_5VideoEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - nf: Number of video features
-        - hs: Hidden size (must match the hidden size of language model
-          backbone)
-        - nv: Number of videos
-        - g: Grid dimensions (3 for t, h, w)
+    """Dimensions:
+    - nf: Number of video features
+    - hs: Hidden size (must match the hidden size of language model
+      backbone)
+    - nv: Number of videos
+    - g: Grid dimensions (3 for t, h, w)
     """
 
     type: Literal["video_embeds"]
@@ -202,8 +197,8 @@ class KeyeVL1_5VideoEmbeddingInputs(TensorSchema):
 class KeyeVL1_5Projector(nn.Module):
     def __init__(
         self,
-        text_config: PretrainedConfig,
-        vision_config: PretrainedConfig,
+        text_config: PreTrainedConfig,
+        vision_config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -378,7 +373,7 @@ def get_supported_mm_limits(
 
 
 class KeyeVL1_5MultiModalProcessor(BaseMultiModalProcessor[KeyeVL1_5ProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _get_prompt_updates(
@@ -437,10 +432,10 @@ def _get_prompt_updates(
         cu_seqlens = torch.cumsum(torch.tensor([0] + num_frames.tolist()), dim=-1)
 
         def get_replacement_keye(item_idx: int, modality: str):
-            """
-            Args:
-                item_idx(int): The item index of modality to replace
-                modality(str): The modality
+            """Args:
+            item_idx(int): The item index of modality to replace
+            modality(str): The modality
+
             """
             if modality == "image":
                 out_item = out_mm_kwargs[modality][item_idx]
@@ -514,15 +509,15 @@ class KeyeVL1_5ForConditionalGeneration(
 ):
     def _build_projector(
         self,
-        text_config: PretrainedConfig,
-        vision_config: PretrainedConfig,
+        text_config: PreTrainedConfig,
+        vision_config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> nn.Module:
         return KeyeVL1_5Projector(text_config, vision_config, quant_config, prefix)
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        config: PretrainedConfig = vllm_config.model_config.hf_config
+        config: PreTrainedConfig = vllm_config.model_config.hf_config
         self.merge_size = config.vision_config.spatial_merge_size
         super().__init__(vllm_config=vllm_config, prefix=prefix)
 
diff --git a/vllm/model_executor/models/kimi_audio.py b/vllm/model_executor/models/kimi_audio.py
index 71cf389f98f7..a2357836f1aa 100644
--- a/vllm/model_executor/models/kimi_audio.py
+++ b/vllm/model_executor/models/kimi_audio.py
@@ -13,7 +13,7 @@
 from transformers import WhisperConfig as HFWhisperConfig
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.inputs import PromptType, TokensPrompt
 from vllm.model_executor.model_loader import DefaultModelLoader
@@ -47,6 +47,7 @@
 )
 from vllm.multimodal.processing.processor import (
     BaseMultiModalProcessor,
+    HFMultiModalInputs,
     ProcessorInputs,
 )
 from vllm.sequence import IntermediateTensors
@@ -175,29 +176,6 @@ def get_dummy_mm_data(
             ),
         }
 
-    def get_dummy_processor_inputs(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
-    ) -> ProcessorInputs:
-        dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
-        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
-
-        num_audios = mm_counts.get("audio", 0)
-        dummy_tokens = (
-            [198]
-            if num_audios == 0
-            else [
-                KimiAudioProcessor.KIMIA_MEDIA_BEGIN,
-                KimiAudioProcessor.KIMIA_TEXT_BLANK,
-                KimiAudioProcessor.KIMIA_MEDIA_END,
-            ]
-            * num_audios
-        )
-
-        return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items)
-
 
 # Field config for Kimi-Audio multimodal data
 _KIMIAUDIO_FIELD_CONFIG = {
@@ -227,27 +205,41 @@ def _parse_audio_data(
 class KimiAudioMultiModalProcessor(BaseMultiModalProcessor[KimiAudioProcessingInfo]):
     """vLLM multi-modal processor wrapper for Kimi-Audio."""
 
-    def _apply_hf_processor_main(
+    def get_dummy_inputs(
         self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        """Call the HuggingFace processor."""
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: MultiModalDummyOptions,
+    ) -> ProcessorInputs:
+        builder = self.dummy_inputs
+        dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options)
+        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
+
+        num_audios = mm_counts.get("audio", 0)
+        dummy_tokens = (
+            [198]
+            if num_audios == 0
+            else [
+                KimiAudioProcessor.KIMIA_MEDIA_BEGIN,
+                KimiAudioProcessor.KIMIA_TEXT_BLANK,
+                KimiAudioProcessor.KIMIA_MEDIA_END,
+            ]
+            * num_audios
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items)
 
-        # Convert mm_data format: {'audios': [...]} -> {'audio': ...}
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", [])
+    def _get_hf_mm_inputs(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
 
         # Convert audio format: [(array, sr), ...] -> [array, ...]
         # KimiAudioProcessor expects raw numpy arrays
-        if audios:
+        mm_data = hf_inputs.hf_data
+        if audios := mm_data.pop("audio", []):
             audio_arrays = []
             for aud in audios:
                 if isinstance(aud, (tuple, list)) and len(aud) == 2:
@@ -259,14 +251,7 @@ def _apply_hf_processor_main(
                     audio_arrays.append(aud)
             mm_data["audio"] = audio_arrays
 
-        # Use the context's call_hf_processor for proper handling
-        processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            mm_data,
-            hf_processor_mm_kwargs,
-        )
-        processed_data.update(passthrough_data)
-        return processed_data
+        return hf_inputs
 
     def _get_mm_fields_config(
         self,
diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py
index 5b0061e3cb0b..81b5ec164629 100644
--- a/vllm/model_executor/models/kimi_k25.py
+++ b/vllm/model_executor/models/kimi_k25.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Kimi-K2.5 Model Implementation for vLLM.
+"""Kimi-K2.5 Model Implementation for vLLM.
 
 Kimi-K2.5 extends Kimi-K2 with vision support.
 """
@@ -15,7 +14,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -90,8 +89,7 @@ class MaxImageTokenMeta:
 
 
 class KimiK25MediaPixelInputs(TensorSchema):
-    """
-    Media input schema for K2-VL model.
+    """Media input schema for K2-VL model.
 
     Dimensions:
         - np: Number of patches (flattened from all media items)
@@ -218,7 +216,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         # TODO: Support mm_options for vision_chunk to allow user configuration
         dummy_items = self.get_dummy_mm_items()
diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py
index eb811c9fee43..03c409a91195 100644
--- a/vllm/model_executor/models/kimi_k25_vit.py
+++ b/vllm/model_executor/models/kimi_k25_vit.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Vision tower implementation for Kimi-K2.5 model.
+"""Vision tower implementation for Kimi-K2.5 model.
 
 This module provides the vision encoder components for Kimi-K2.5,
 including 3D patch embedding, RoPE position embedding, and
@@ -284,12 +283,12 @@ def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor:
     def get_freqs_cis(
         self, grid_thws: torch.Tensor | list[list[int]], device: torch.device
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             grid_thws (torch.Tensor): grid time, height and width
 
         Returns:
             freqs_cis: tensor of shape (sum(t * height * width), dim//2)
+
         """
         if not hasattr(self, "freqs_cis"):
             self.register_buffer(
@@ -449,6 +448,10 @@ def attention_qkvpacked(
         Args:
             x (torch.Tensor): (seqlen, hidden_dim)
             cu_seqlens (torch.Tensor): cumulative sequence lengths
+            rope_freqs_cis (torch.Tensor): rotary embedding frequencies
+            max_seqlen (torch.Tensor | None): longest sequence in the batch
+            sequence_lengths (torch.Tensor | None): per-sequence lengths
+
         """
         seq_length = x.size(0)
         xqkv, _ = self.wqkv(x)
@@ -753,13 +756,13 @@ def forward(
         *,
         encoder_metadata: dict[str, torch.Tensor | None] | None = None,
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             pixel_values (torch.Tensor): The input pixel values.
             grid_thws (torch.Tensor): Temporal, height and width.
 
         Returns:
             torch.Tensor: The output tokens.
+
         """
         if encoder_metadata is not None and "pos_embeds" in encoder_metadata:
             hidden_states = self.patch_embed(
diff --git a/vllm/model_executor/models/kimi_vl.py b/vllm/model_executor/models/kimi_vl.py
index 218c257e53c5..13779042b62e 100644
--- a/vllm/model_executor/models/kimi_vl.py
+++ b/vllm/model_executor/models/kimi_vl.py
@@ -53,7 +53,7 @@
 from transformers.activations import GELUActivation
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.models.interfaces import (
@@ -135,12 +135,11 @@ def forward(self, image_features: torch.Tensor) -> torch.Tensor:
 
 
 class KimiVLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - nc: Number of channels
-        - np: Number of patches
-        - ps: Patch size
-        - ni: Number of images
+    """Dimensions:
+    - nc: Number of channels
+    - np: Number of patches
+    - ps: Patch size
+    - ni: Number of images
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -220,24 +219,20 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=MaxImageTokenMeta.width,
                 height=MaxImageTokenMeta.height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
 
 class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _get_mm_fields_config(
diff --git a/vllm/model_executor/models/laguna.py b/vllm/model_executor/models/laguna.py
index 12046379fd77..c1aaf8c9a86e 100644
--- a/vllm/model_executor/models/laguna.py
+++ b/vllm/model_executor/models/laguna.py
@@ -19,12 +19,14 @@
 )
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -54,8 +56,19 @@
 logger = init_logger(__name__)
 
 
+def get_mlp_layer_types(config) -> list[str]:
+    """Return which layers are MoE, rejecting configs predating Transformers support."""
+    mlp_layer_types = getattr(config, "mlp_layer_types", None)
+    if mlp_layer_types is None:
+        raise ValueError(
+            "Laguna requires the `LagunaConfig` from Transformers, which is "
+            "available from v5.17.0. Please update Transformers to run this model."
+        )
+    return mlp_layer_types
+
+
 class LagunaMLP(nn.Module):
-    """Dense MLP for Laguna (used in mlp_only_layers)."""
+    """Dense MLP for Laguna."""
 
     def __init__(
         self,
@@ -160,11 +173,9 @@ def __init__(
         self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
         self.n_local_physical_experts = self.n_physical_experts // self.ep_size
         # Router gate
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
@@ -206,7 +217,7 @@ def __init__(
             top_k=config.num_experts_per_tok,
             hidden_size=config.hidden_size,
             intermediate_size=config.moe_intermediate_size,
-            renormalize=config.norm_topk_prob,
+            renormalize=True,
             quant_config=quant_config,
             prefix=f"{prefix}.experts",
             scoring_func="sigmoid",
@@ -293,13 +304,12 @@ def __init__(
         else:
             self.sliding_window = None
 
-        # QKV projection (no bias for Laguna)
         self.qkv_proj = QKVParallelLinear(
             self.hidden_size,
             self.head_dim,
             self.total_num_heads,
             self.total_num_kv_heads,
-            bias=config.qkv_bias,
+            bias=config.attention_bias,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
         )
@@ -506,14 +516,9 @@ def __init__(
             ),
         )
 
-        # Check if this layer uses MoE or dense MLP (matches Qwen2/Qwen3 convention)
-        mlp_only_layers = (
-            [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers
-        )
         self.is_moe_layer = (
-            (layer_idx not in mlp_only_layers)
-            and (config.num_experts > 0)
-            and ((layer_idx + 1) % config.decoder_sparse_step == 0)
+            get_mlp_layer_types(config)[layer_idx] == "sparse"
+            and config.num_experts > 0
         )
 
         if self.is_moe_layer:
diff --git a/vllm/model_executor/models/lfm2.py b/vllm/model_executor/models/lfm2.py
index 6f03150b8461..c43a9dcbcabd 100644
--- a/vllm/model_executor/models/lfm2.py
+++ b/vllm/model_executor/models/lfm2.py
@@ -440,6 +440,7 @@ def get_mamba_state_shape_from_config(
         Returns:
             Tuple containing:
             - conv_state_shape: Shape for convolutional state cache
+
         """
         parallel_config = vllm_config.parallel_config
         hf_config = vllm_config.model_config.hf_config
diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py
index c2a21e52d18c..29bc56a69661 100644
--- a/vllm/model_executor/models/lfm2_moe.py
+++ b/vllm/model_executor/models/lfm2_moe.py
@@ -17,12 +17,12 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -125,10 +125,9 @@ def __init__(
         self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
         self.n_local_physical_experts = self.n_physical_experts // self.ep_size
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.gate",
         )
@@ -556,6 +555,7 @@ def get_mamba_state_shape_from_config(
         Returns:
             Tuple containing:
             - conv_state_shape: Shape for convolutional state cache
+
         """
         parallel_config = vllm_config.parallel_config
         hf_config = vllm_config.model_config.hf_config
diff --git a/vllm/model_executor/models/lfm2_siglip2.py b/vllm/model_executor/models/lfm2_siglip2.py
index 54fc1c5d2ff0..42c8f6264198 100644
--- a/vllm/model_executor/models/lfm2_siglip2.py
+++ b/vllm/model_executor/models/lfm2_siglip2.py
@@ -61,6 +61,7 @@ def forward(
 
         Returns:
             (1, total_tokens, embed_dim) packed embeddings.
+
         """
         assert spatial_shapes.device.type == "cpu", (
             "Expected `spatial_shapes` on CPU to avoid device-to-host sync in "
@@ -113,6 +114,7 @@ def resize_positional_embeddings_packed(
         Returns:
             (total_tokens, embed_dim) packed positional embeddings, concatenated
             in the same order as `lengths_list`.
+
         """
         assert spatial_shapes.device.type == "cpu"
 
@@ -155,7 +157,7 @@ def resize_positional_embeddings_packed(
 
 
 class Siglip2Attention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
 
     def __init__(
         self,
@@ -304,11 +306,11 @@ def forward(
         cu_seqlens: torch.Tensor,
         max_seqlen: int | torch.Tensor,
     ) -> torch.Tensor:
-        """
-        Args:
-            hidden_states: Input tensor of shape (batch, seq_len, embed_dim).
-            cu_seqlens: Cumulative sequence lengths tensor.
-            max_seqlen: Maximum sequence length.
+        """Args:
+        hidden_states: Input tensor of shape (batch, seq_len, embed_dim).
+        cu_seqlens: Cumulative sequence lengths tensor.
+        max_seqlen: Maximum sequence length.
+
         """
         residual = hidden_states
 
@@ -328,12 +330,12 @@ def forward(
 
 
 class Siglip2Encoder(nn.Module):
-    """
-    Transformer encoder consisting of `config.num_hidden_layers`
+    """Transformer encoder consisting of `config.num_hidden_layers`
     self attention layers. Each layer is a [`Siglip2EncoderLayer`].
 
     Args:
-        config: PretrainedConfig
+        config: PreTrainedConfig
+
     """
 
     def __init__(
@@ -430,8 +432,7 @@ def forward(
         max_seqlen: torch.Tensor,
         select_layers: list[int] | None = None,
     ) -> torch.Tensor:
-        r"""
-        spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
+        r"""spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
             Tensor containing the spatial dimensions (height, width)
         of the input images.
         select_layers (`list[int]` or `None`, defaults to `None`):
@@ -500,10 +501,15 @@ def forward(
         """Forward pass through the vision model.
 
         Args:
+            pixel_values_packed: Packed pixel values for all images.
+            spatial_shapes: Per-image spatial dimensions.
+            cu_seqlens: Cumulative sequence lengths of the packed images.
+            max_seqlen: Longest image sequence in the batch.
             select_layers: Layer indices to select hidden states from.
                 Supports negative indices (e.g., [-2] for second-to-last).
                 If None, returns the last layer output with post_layernorm.
                 Multiple layers can be selected and will be concatenated.
+
         """
         return self.vision_model(
             pixel_values_packed=pixel_values_packed,
diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py
index a8b57e7c28a2..dd58707b855f 100644
--- a/vllm/model_executor/models/lfm2_vl.py
+++ b/vllm/model_executor/models/lfm2_vl.py
@@ -13,15 +13,15 @@
 from transformers.activations import ACT2FN
 from transformers.models.lfm2_vl import Lfm2VlProcessor
 from transformers.models.lfm2_vl.configuration_lfm2_vl import Lfm2VlConfig
-from transformers.models.lfm2_vl.image_processing_lfm2_vl_fast import (
-    Lfm2VlImageProcessorFast,
+from transformers.models.lfm2_vl.image_processing_lfm2_vl import (
+    Lfm2VlImageProcessor,
     find_closest_aspect_ratio,
     round_by_factor,
 )
 from typing_extensions import Buffer
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.forward_context import set_forward_context
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import ReplicatedLinear
@@ -35,6 +35,7 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
@@ -80,12 +81,11 @@ def _pad_cumulative_seqlens_buffer(
 
 
 class Lfm2VLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Number of images in the prompt
-        - bn: Batch size * number of images
-        - d: Number of dimensions
-        - fd: Number of features per dimension
+    """Dimensions:
+    - b: Number of images in the prompt
+    - bn: Batch size * number of images
+    - d: Number of dimensions
+    - fd: Number of features per dimension
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -104,7 +104,7 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs):
         return self.ctx.get_hf_processor(Lfm2VlProcessor, **kwargs)
 
-    def get_image_processor(self, **kwargs: object) -> Lfm2VlImageProcessorFast:
+    def get_image_processor(self, **kwargs: object) -> Lfm2VlImageProcessor:
         return self.get_hf_processor(**kwargs).image_processor
 
     def get_default_tok_params(self) -> TokenizeParams:
@@ -213,7 +213,7 @@ def _get_image_feature_grid_size(
         processor: Lfm2VlProcessor,
         mm_kwargs: Mapping[str, object],
     ) -> tuple[int, int, int]:
-        image_processor: Lfm2VlImageProcessorFast = processor.image_processor
+        image_processor: Lfm2VlImageProcessor = processor.image_processor
 
         mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs)
         downsample_factor = mm_kwargs.get(
@@ -336,7 +336,7 @@ def get_num_image_tokens(
         processor: Lfm2VlProcessor,
         mm_kwargs: Mapping[str, object],
     ) -> tuple[int, int]:
-        image_processor: Lfm2VlImageProcessorFast = processor.image_processor
+        image_processor: Lfm2VlImageProcessor = processor.image_processor
 
         mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs)
         downsample_factor = mm_kwargs.get(
@@ -387,26 +387,22 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
         }
 
 
 class Lfm2VLMultiModalProcessor(BaseMultiModalProcessor[Lfm2VLProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -537,6 +533,7 @@ def forward(
 
         Returns:
             projected_packed: (total_projected_tokens, text_hidden_size)
+
         """
         assert spatial_shapes.device.type == "cpu", (
             "Expected `spatial_shapes` on CPU to avoid device-to-host sync in "
@@ -666,6 +663,7 @@ def get_mamba_state_shape_from_config(
         Returns:
             Tuple containing:
             - conv_state_shape: Shape for convolutional state cache
+
         """
         parallel_config = vllm_config.parallel_config
         hf_language_config = vllm_config.model_config.hf_config.text_config
@@ -1275,21 +1273,20 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
             tower_model="vision_tower",
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        downsample_factor = self.config.downsample_factor
-
-        return num_image_tokens * downsample_factor**2
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         downsample_factor = self.config.downsample_factor
-
-        return num_vision_tokens // downsample_factor**2
+        return num_mm_embeds * downsample_factor**2, num_mm_embeds
diff --git a/vllm/model_executor/models/lightonocr.py b/vllm/model_executor/models/lightonocr.py
index 0ce0d7f3d2b3..d8a124024271 100644
--- a/vllm/model_executor/models/lightonocr.py
+++ b/vllm/model_executor/models/lightonocr.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable, Mapping, Sequence
-from typing import TypeVar
 
 import torch
 import torch.nn as nn
@@ -35,8 +34,6 @@
     PromptUpdateDetails,
 )
 
-_I = TypeVar("_I", bound=Mistral3ProcessingInfo)
-
 
 class LightOnOCRProcessingInfo(Mistral3ProcessingInfo):
     def get_vision_encoder_info(
@@ -49,8 +46,8 @@ def get_vision_encoder_info(
         return PixtralHFEncoderInfo(self.get_hf_config())
 
 
-class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessingInfo]):
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -62,27 +59,6 @@ def _postprocess_hf_mm_data(
         if not mm_data:
             return processed_data
 
-        input_ids = processed_data.get("input_ids")
-        if input_ids is not None:
-            processor = self.info.get_hf_processor()
-            tokenizer = self.info.get_tokenizer()
-            vocab = tokenizer.get_vocab()
-
-            break_id = vocab.get(processor.image_break_token)
-            end_id = vocab.get(processor.image_end_token)
-
-            # create mask to remove break/end tokens
-            keep_mask = ~torch.isin(
-                input_ids,
-                torch.tensor([break_id, end_id]),
-            )
-
-            processed_data["input_ids"] = input_ids[keep_mask].unsqueeze(0)
-            if "attention_mask" in processed_data:
-                processed_data["attention_mask"] = processed_data["attention_mask"][
-                    keep_mask
-                ].unsqueeze(0)
-
         # un-pad pixel_values per-image so caches remain independent.
         pixel_values = processed_data.get("pixel_values")
         if pixel_values is not None:
diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py
index 3dc0e9a904df..0eb34c70b641 100644
--- a/vllm/model_executor/models/llama4.py
+++ b/vllm/model_executor/models/llama4.py
@@ -38,12 +38,12 @@
 )
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -98,11 +98,9 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
         self.ep_size = self.ep_group.size()
 
         intermediate_size_moe = config.intermediate_size
-        self.router = ReplicatedLinear(
+        self.router = GateLinear(
             config.hidden_size,
             config.num_local_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.router",
         )
 
@@ -406,8 +404,7 @@ def load_moe_expert_weights(
         expert_params_mapping: list[tuple[str, str, int, str]],
         fused: bool = True,
     ) -> bool:
-        """
-        Load MoE expert weights.
+        """Load MoE expert weights.
 
         Args:
             name: The name of the weight to load.
@@ -427,8 +424,8 @@ def load_moe_expert_weights(
         Returns:
             True if loaded_weight is one of MoE weights and the MoE expert
             weights are loaded successfully, False otherwise.
-        """
 
+        """
         # Whether the MoE expert weights are loaded successfully.
         expert_param_loaded = False
 
diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py
index c3d956a946d1..68e21c9dc284 100644
--- a/vllm/model_executor/models/llama_eagle3.py
+++ b/vllm/model_executor/models/llama_eagle3.py
@@ -335,10 +335,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.logits_processor = LogitsProcessor(
             self.config.draft_vocab_size, scale=logit_scale
         )
-        self.draft_id_to_target_id = nn.Parameter(
-            torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
-            requires_grad=False,
-        )
+        target_vocab_size = vllm_config.model_config.get_vocab_size()
+        if self.config.draft_vocab_size != target_vocab_size:
+            self.draft_id_to_target_id = nn.Parameter(
+                torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
+                requires_grad=False,
+            )
+        else:
+            self.draft_id_to_target_id = None
 
         self.use_parallel_drafting = speculative_config.parallel_drafting
 
diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py
index 1f1f45bfea73..ead9bbb1e72f 100644
--- a/vllm/model_executor/models/llava.py
+++ b/vllm/model_executor/models/llava.py
@@ -12,22 +12,22 @@
     CLIPVisionConfig,
     LlavaConfig,
     PixtralVisionConfig,
-    PretrainedConfig,
+    PreTrainedConfig,
     SiglipVisionConfig,
 )
 from transformers.models.llava import LlavaProcessor
 from transformers.models.pixtral import PixtralProcessor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.cache import BaseMultiModalProcessorCache
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import (
@@ -71,8 +71,7 @@
 
 
 class LlavaImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - bn: Batch size * number of images
         - c: Number of channels (3)
         - h: Height
@@ -87,8 +86,7 @@ class LlavaImagePixelInputs(TensorSchema):
 
 
 class PixtralHFImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - bn: Batch size * number of images
         - c: Number of channels
         - h: Height
@@ -106,11 +104,10 @@ class PixtralHFImagePixelInputs(TensorSchema):
 
 
 class LlavaImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - ifs: Image feature size
-        - hs: Hidden size (must match language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - ifs: Image feature size
+    - hs: Hidden size (must match language model backbone)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -159,7 +156,7 @@ def forward(self, image_features: torch.Tensor) -> torch.Tensor:
 
 
 class LlavaLikeConfig(Protocol):
-    vision_config: PretrainedConfig
+    vision_config: PreTrainedConfig
     image_token_index: int
     vision_feature_select_strategy: str
     vision_feature_layer: int | list[int]
@@ -230,20 +227,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -322,7 +315,7 @@ def get_hf_processor(self, **kwargs: object):
 
 
 class PixtralHFMultiModalProcessor(BaseMultiModalProcessor[PixtralHFProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -413,21 +406,17 @@ def _build_llava_or_pixtral_hf_info(
 def _build_llava_or_pixtral_hf_processor(
     info: _I,
     dummy_inputs: BaseDummyInputsBuilder[_I],
-    *,
-    cache: BaseMultiModalProcessorCache | None = None,
 ) -> BaseMultiModalProcessor:
     if isinstance(info, PixtralHFProcessingInfo):
         return PixtralHFMultiModalProcessor(
             info,
             dummy_inputs,  # type: ignore
-            cache=cache,
         )
 
     if isinstance(info, LlavaProcessingInfo):
         return LlavaMultiModalProcessor(
             info,
             dummy_inputs,  # type: ignore
-            cache=cache,
         )
 
     raise NotImplementedError(type(info))
@@ -439,6 +428,7 @@ def _get_num_hidden_layers(hf_config: LlavaLikeConfig) -> int:
 
     Args:
         hf_config: Model config with vision feature layer(s).
+
     """
     feature_layers = hf_config.vision_feature_layer
     num_hidden_layers = hf_config.vision_config.num_hidden_layers
@@ -706,9 +696,12 @@ def forward(
             positions: Position indices for the input tokens.
             intermediate_tensors: Intermediate tensors from prior forward pass.
             inputs_embeds: Optional tensor of input embeddings.
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
 
         Info:
             [`LlavaImageInputs`][vllm.model_executor.models.llava.LlavaImageInputs]
+
         """
         if intermediate_tensors is not None:
             inputs_embeds = None
@@ -730,27 +723,19 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
             tower_model="vision_tower",
         )
 
-    def get_num_mm_encoder_tokens(
-        self,
-        num_image_tokens: int,
-    ) -> int:
-        # LLaVA's vision encoder outputs one token per patch without
-        # spatial merging or pixel shuffle
-        return num_image_tokens
-
-    def get_num_mm_connector_tokens(
+    def get_mm_lora_token_counts(
         self,
-        num_vision_tokens: int,
-    ) -> int:
-        # LLaVA's MLP projector outputs the same number of tokens
-        # as it receives from the vision encoder (1:1 mapping)
-        return num_vision_tokens
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+        return num_mm_embeds, num_mm_embeds
diff --git a/vllm/model_executor/models/llava_next.py b/vllm/model_executor/models/llava_next.py
index fee9a4ac59a3..02400a763b22 100644
--- a/vllm/model_executor/models/llava_next.py
+++ b/vllm/model_executor/models/llava_next.py
@@ -47,8 +47,7 @@
 
 
 class LlavaNextImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - bn: Batch size * number of images
         - np: Number of patches + 1
         - c: Number of channels (3)
@@ -70,11 +69,10 @@ class LlavaNextImagePixelInputs(TensorSchema):
 
 
 class LlavaNextImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - ifs: Image feature size
-        - hs: Hidden size (must match language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - ifs: Image feature size
+    - hs: Hidden size (must match language model backbone)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -211,7 +209,7 @@ def _get_mm_fields_config(
 
 
 class LlavaNextMultiModalProcessor(BaseLlavaNextMultiModalProcessor[_I]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _get_mm_fields_config(
@@ -578,9 +576,12 @@ def forward(
             positions: Position indices for the input tokens.
             intermediate_tensors: Intermediate tensors from prior forward pass.
             inputs_embeds: Optional tensor of input embeddings.
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
 
         Info:
             [`LlavaNextImageInputs`][vllm.model_executor.models.llava_next.LlavaNextImageInputs]
+
         """
         if intermediate_tensors is not None:
             inputs_embeds = None
@@ -601,9 +602,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
diff --git a/vllm/model_executor/models/llava_next_video.py b/vllm/model_executor/models/llava_next_video.py
index 038da17b75bf..a33480ef08eb 100644
--- a/vllm/model_executor/models/llava_next_video.py
+++ b/vllm/model_executor/models/llava_next_video.py
@@ -10,13 +10,14 @@
 from transformers import BatchFeature, LlavaNextVideoConfig, LlavaNextVideoProcessor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.models.clip import CLIPVisionModel
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import (
@@ -55,8 +56,7 @@
 
 
 class LlavaNextVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - bn: Batch size * number of videos
         - f: Number of frames
         - c: Number of channels (3)
@@ -171,25 +171,20 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(
             seq_len, mm_counts
         )
 
-        video_overrides = mm_options.get("video")
-        assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
-
         return {
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             )
         }
 
@@ -380,8 +375,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
     def _parse_and_validate_video_input(
         self, **kwargs: object
     ) -> LlavaNextVideoPixelInputs | None:
-        """
-        A legal video input should have the following dimensions:
+        """A legal video input should have the following dimensions:
         {
             "pixel_values_videos" :
                 list[b, Tensor(nb_frames, nb_channels, height, width)]
@@ -456,10 +450,17 @@ def forward(
         **kwargs: object,
     ) -> torch.Tensor | IntermediateTensors:
         """Run forward pass for LlaVA-NeXT-Video.
+
         Args:
             input_ids: Flattened (concatenated) input_ids corresponding to a
                 batch.
-            pixel_values_videos: Pixels in each frames for each input videos.
+            positions: Flattened (concatenated) position ids corresponding to a
+                batch.
+            intermediate_tensors: Intermediate tensors from prior forward pass.
+            inputs_embeds: Optional tensor of input embeddings.
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
+
         """
         if intermediate_tensors is not None:
             inputs_embeds = None
@@ -485,40 +486,37 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
             tower_model="vision_tower",
         )
 
-    def get_num_mm_encoder_tokens(
+    def get_mm_lora_token_counts(
         self,
-        num_video_tokens: int,
-    ) -> int:
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+
         # Invert the spatial pooling done by `vision_resampler`: each frame
         # contributes `pooled_grid_length ** 2` tokens after the language
         # model's placeholder count, but `patch_grid_length ** 2` tokens
         # when it leaves the vision encoder.
         pooled_tokens_per_frame = self.pooled_grid_length**2
-        if num_video_tokens <= 0 or pooled_tokens_per_frame <= 0:
-            return 0
-
-        num_frames = num_video_tokens // pooled_tokens_per_frame
-        return num_frames * (self.patch_grid_length**2)
-
-    def get_num_mm_connector_tokens(
-        self,
-        num_vision_tokens: int,
-    ) -> int:
-        # The projector is length-preserving; it runs after
-        # `vision_resampler` has already pooled each frame down to
-        # `pooled_grid_length ** 2` tokens.
         patch_tokens_per_frame = self.patch_grid_length**2
-        if num_vision_tokens <= 0 or patch_tokens_per_frame <= 0:
-            return 0
-
-        num_frames = num_vision_tokens // patch_tokens_per_frame
-        return num_frames * (self.pooled_grid_length**2)
+        if (
+            num_mm_embeds <= 0
+            or pooled_tokens_per_frame <= 0
+            or patch_tokens_per_frame <= 0
+        ):
+            return 0, 0
+
+        num_frames = num_mm_embeds // pooled_tokens_per_frame
+        return (
+            num_frames * patch_tokens_per_frame,
+            num_frames * pooled_tokens_per_frame,
+        )
diff --git a/vllm/model_executor/models/llava_onevision.py b/vllm/model_executor/models/llava_onevision.py
index 631f73e64d99..6fc96c8e6774 100644
--- a/vllm/model_executor/models/llava_onevision.py
+++ b/vllm/model_executor/models/llava_onevision.py
@@ -14,7 +14,7 @@
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -53,17 +53,16 @@
 
 
 class LlavaOnevisionVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of videos
-        - f: Number of frames
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
-
-        Note that `f` may be different for each batch, and 'num_frames'
-        may be different for each video, in which case the data is passed as a
-        list instead of a batched tensor.
+    """Dimensions:
+    - bn: Batch size * number of videos
+    - f: Number of frames
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
+
+    Note that `f` may be different for each batch, and 'num_frames'
+    may be different for each video, in which case the data is passed as a
+    list instead of a batched tensor.
     """
 
     type: Literal["pixel_values_videos"] = "pixel_values_videos"
@@ -75,16 +74,15 @@ class LlavaOnevisionVideoPixelInputs(TensorSchema):
 
 
 class LlavaOnevisionImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - np: Number of patches (1 + num_patches)
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
-
-        Note that `num_patches` may be different per batch and image,
-        in which case the data is passed as a list instead of a batched tensor.
+    """Dimensions:
+    - bn: Batch size * number of images
+    - np: Number of patches (1 + num_patches)
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
+
+    Note that `num_patches` may be different per batch and image,
+    in which case the data is passed as a list instead of a batched tensor.
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -98,11 +96,10 @@ class LlavaOnevisionImagePixelInputs(TensorSchema):
 
 
 class LlavaOnevisionImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - ifs: Image feature size
-        - hs: Hidden size (must match language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - ifs: Image feature size
+    - hs: Hidden size (must match language model backbone)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -281,33 +278,26 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(
             seq_len, mm_counts
         )
 
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-        assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -327,33 +317,36 @@ def _get_mm_fields_config(
             pixel_values_videos=MultiModalFieldConfig.batched("video"),
         )
 
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        if not hf_data:
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        prompt_text = hf_data.pop("text")
+        assert isinstance(prompt_text, str)
 
-        mm_data = dict(mm_data)
-        videos = mm_data.pop("videos", [])
+        videos = hf_data.pop("videos", [])
         assert isinstance(videos, list)
 
         if not videos:
             processed_data = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(**hf_processor_mm_kwargs),
-                dict(text=prompt_text, **mm_data),
-                hf_processor_mm_kwargs,
+                self.info.get_hf_processor(**hf_kwargs),
+                dict(text=prompt_text, **hf_data),
+                hf_kwargs,
+            )
+            return self._finalize_hf_mm_data(
+                hf_data, hf_kwargs, passthrough_data, processed_data
             )
-            processed_data.update(passthrough_data)
-            return processed_data
 
         # LLaVA-OneVision processor doesn't support multiple videos
         # with different sizes when converting back to tensors
@@ -363,13 +356,13 @@ def _apply_hf_processor_main(
         image_token = processor.image_token
         video_token = processor.video_token
 
-        images = mm_data.pop("images", [])
+        images = hf_data.pop("images", [])
         assert isinstance(images, list)
         if images:
             processor_outputs = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(**hf_processor_mm_kwargs),
+                self.info.get_hf_processor(**hf_kwargs),
                 dict(text=image_token * len(images), **{"images": images}),
-                hf_processor_mm_kwargs,
+                hf_kwargs,
             )
             image_outputs = {
                 k: v
@@ -382,9 +375,9 @@ def _apply_hf_processor_main(
         pixel_values_videos = []
         for video in videos:
             item_outputs = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(**hf_processor_mm_kwargs),
+                self.info.get_hf_processor(**hf_kwargs),
                 dict(text=video_token, **{"videos": video}),
-                hf_processor_mm_kwargs,
+                hf_kwargs,
             )
 
             pixel_values_videos.append(item_outputs["pixel_values_videos"][0])
@@ -397,8 +390,9 @@ def _apply_hf_processor_main(
             **video_outputs,
         )
         processed_data = BatchFeature(combined_outputs)
-        processed_data.update(passthrough_data)
-        return processed_data
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_prompt_updates(
         self,
@@ -557,8 +551,7 @@ def _parse_and_validate_image_input(
     def _parse_and_validate_video_input(
         self, **kwargs: object
     ) -> LlavaOnevisionVideoPixelInputs | None:
-        """
-        A legal video input should have the following dimensions:
+        """A legal video input should have the following dimensions:
         {
             "pixel_values_videos" :
                 list[b, Tensor(nb_frames, nb_channels, height, width)]
@@ -891,10 +884,17 @@ def forward(
         **kwargs: object,
     ) -> torch.Tensor | IntermediateTensors:
         """Run forward pass for LlaVA-Onevision.
+
         Args:
             input_ids: Flattened (concatenated) input_ids corresponding to a
                 batch.
-            pixel_values_videos: Pixels in each frames for each input videos.
+            positions: Flattened (concatenated) position ids corresponding to a
+                batch.
+            intermediate_tensors: Intermediate tensors from prior forward pass.
+            inputs_embeds: Optional tensor of input embeddings.
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
+
         """
         if intermediate_tensors is not None:
             inputs_embeds = None
diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py
index 7eda9dc32025..a1b12b38346e 100644
--- a/vllm/model_executor/models/llava_onevision2.py
+++ b/vllm/model_executor/models/llava_onevision2.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Inference-only LLaVA-OneVision-2 (OV2) model for vLLM.
+"""Inference-only LLaVA-OneVision-2 (OV2) model for vLLM.
 
 Architecture notes:
 
@@ -55,7 +54,7 @@
     support_torch_compile,
 )
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import parallel_state
 from vllm.distributed import utils as dist_utils
 from vllm.inputs import ModalityData, MultiModalDataDict
@@ -1308,7 +1307,6 @@ def get_hf_processor(self, **kwargs: object):
         # ``_merge_mm_kwargs`` restricts caller ``mm_processor_kwargs`` to known
         # processor args and wraps values as hashable for the lru_cache.
         merged = _merge_mm_kwargs(model_config, AutoProcessor, **kwargs)
-        merged.setdefault("use_fast", True)
         return _load_ov2_processor(
             model_config.model,
             model_config.revision,
@@ -1472,7 +1470,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+        mm_options: MultiModalDummyOptions | None = None,
     ) -> MultiModalDataDict:
         n_img = mm_counts.get("image", 0)
         n_vid = mm_counts.get("video", 0)
@@ -1557,23 +1555,26 @@ def _parse_image_data(
 class LlavaOnevision2MultiModalProcessor(
     BaseMultiModalProcessor[LlavaOnevision2ProcessingInfo]
 ):
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        if not hf_data:
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        prompt_text = hf_data.pop("text")
+        assert isinstance(prompt_text, str)
 
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
-        merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs)
+        hf_processor = self.info.get_hf_processor(**hf_kwargs)
+        merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_kwargs)
         merged_kwargs.setdefault("return_tensors", "pt")
         call_kwargs = {
             k: v
@@ -1590,17 +1591,17 @@ def _apply_hf_processor_main(
                 "codec_config",
             }
         }
-        mm_data = dict(mm_data)
+
         # Explicit None + length checks: ``mm_data[...]`` may be a list, numpy
         # array, or tensor, and ``and <array>`` would raise on the ambiguous
         # truth value of a multi-element array.
-        _videos = mm_data.get("videos")
+        _videos = hf_data.get("videos")
         if _videos is not None and not isinstance(_videos, Sized):
             raise TypeError(f"object of type '{type(_videos).__name__}' has no len()")
         videos_present = _videos is not None and len(_videos) > 0
 
         codec_video_paths = (
-            _extract_codec_video_paths(mm_data["videos"]) if videos_present else None
+            _extract_codec_video_paths(hf_data["videos"]) if videos_present else None
         )
         is_codec_marker = codec_video_paths is not None
         # Fallback: caller passed video_backend=codec via mm_processor_kwargs
@@ -1614,7 +1615,7 @@ def _apply_hf_processor_main(
             and call_kwargs.get("video_backend") == "codec"
         )
         if is_codec_kwarg:
-            raw = mm_data["videos"]
+            raw = hf_data["videos"]
             if isinstance(raw, str):
                 codec_video_paths = [raw]
             elif isinstance(raw, (list, tuple)) and all(
@@ -1639,7 +1640,7 @@ def _apply_hf_processor_main(
             # Codec backend: HF processor consumes the path string directly and
             # performs decode + canvas-packing internally. The dummy ndarray
             # we attached during prepare_codec_video_input is discarded here.
-            mm_data["videos"] = (
+            hf_data["videos"] = (
                 codec_video_paths
                 if len(codec_video_paths) > 1
                 else codec_video_paths[0]
@@ -1649,11 +1650,9 @@ def _apply_hf_processor_main(
             # ``video_backend="codec"`` via mm_kwargs so the wrapped processor
             # dispatches to its codec branch.
             output = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(
-                    **{**hf_processor_mm_kwargs, "video_backend": "codec"}
-                ),
-                dict(text=prompt_text, **mm_data),
-                {**hf_processor_mm_kwargs, "video_backend": "codec"},
+                self.info.get_hf_processor(**{**hf_kwargs, "video_backend": "codec"}),
+                dict(text=prompt_text, **hf_data),
+                {**hf_kwargs, "video_backend": "codec"},
             )
             data = dict(output)
             processed_data = BatchFeature(
@@ -1661,8 +1660,9 @@ def _apply_hf_processor_main(
                     data, codec_video_paths, hf_processor
                 )
             )
-            processed_data.update(passthrough_data)
-            return processed_data
+            return self._finalize_hf_mm_data(
+                hf_data, hf_kwargs, passthrough_data, processed_data
+            )
 
         # ---- Frame backend (registered LlavaOnevision2VideoBackend) ------
         # Every non-codec video reaches here as a ``(frames_ndarray, metadata)``
@@ -1677,7 +1677,7 @@ def _apply_hf_processor_main(
         # provided by the backend's ``compute_frames_index_to_sample``; SSRF /
         # local-file gating is enforced by the connector before decoding.
         if videos_present:
-            timestamp_decimals = hf_processor_mm_kwargs.get(
+            timestamp_decimals = hf_kwargs.get(
                 "timestamp_decimals", _DEFAULT_TIMESTAMP_DECIMALS
             )
             if not isinstance(
@@ -1691,7 +1691,7 @@ def _apply_hf_processor_main(
 
             per_video_frames: list[list[Image.Image]] = []
             per_video_timestamps: list[list[float]] = []
-            videos = mm_data["videos"]
+            videos = hf_data["videos"]
             if not isinstance(videos, Iterable):
                 raise TypeError(f"Unsupported videos input: {type(videos)}")
             for item in videos:
@@ -1715,7 +1715,7 @@ def _apply_hf_processor_main(
             # order) -- otherwise mixed image+video requests bind frames to the
             # wrong placeholder. ``row_is_video`` labels each grid row so outputs
             # can be split back into per-modality keys below.
-            merged_mm_data = dict(mm_data)
+            merged_mm_data = dict(hf_data)
             existing_images = merged_mm_data.pop("images", None)
             caller_images: list[Image.Image] = []
             if existing_images:
@@ -1761,9 +1761,9 @@ def _apply_hf_processor_main(
             # **kwargs to the image processor, so passing the full merged kwarg
             # set here is a no-op beyond return_tensors/padding.
             output = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(**hf_processor_mm_kwargs),
+                self.info.get_hf_processor(**hf_kwargs),
                 dict(text=new_prompt, **merged_mm_data),
-                hf_processor_mm_kwargs,
+                hf_kwargs,
             )
             data = dict(output)
 
@@ -1821,20 +1821,22 @@ def _gather_rows(tensor, rows):
             )
 
             processed_data = BatchFeature(data)
-            processed_data.update(passthrough_data)
-            return processed_data
+            return self._finalize_hf_mm_data(
+                hf_data, hf_kwargs, passthrough_data, processed_data
+            )
 
         # ---- Image-only / text-only call --------------------------------
         # No videos present: delegate to ``ctx.call_hf_processor``, which
         # runs the wrapped processor over the (possibly empty) image set and
         # applies float-tensor dtype postprocessing automatically.
         processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            dict(text=prompt_text, **mm_data),
-            hf_processor_mm_kwargs,
+            self.info.get_hf_processor(**hf_kwargs),
+            dict(text=prompt_text, **hf_data),
+            hf_kwargs,
+        )
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
         )
-        processed_data.update(passthrough_data)
-        return processed_data
 
     def _rename_codec_outputs_to_video(
         self,
diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py
index 4178d63ff450..0c9802a8b3dc 100644
--- a/vllm/model_executor/models/longcat_flash.py
+++ b/vllm/model_executor/models/longcat_flash.py
@@ -39,7 +39,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -48,12 +48,12 @@
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -80,7 +80,7 @@
 logger = init_logger(__name__)
 
 
-class FlashConfig(PretrainedConfig):
+class FlashConfig(PreTrainedConfig):
     """Flash model configuration."""
 
     moe_intermediate_size: int
@@ -254,12 +254,12 @@ def __init__(
             else config.num_experts[0]
         )
         self.n_routed_experts = self.n_routed_experts + zero_expert_num
-        self.classifier = ReplicatedLinear(
+        self.classifier = GateLinear(
             config.hidden_size,
             self.n_routed_experts,
             bias=config.router_bias,
+            out_dtype=router_params_dtype,
             params_dtype=router_params_dtype,
-            quant_config=None,
             prefix=f"{prefix}.classifier",
         )
         self.e_score_correction_bias = nn.Parameter(
@@ -331,9 +331,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         else:
             hidden_states_padded = hidden_states
 
-        router_logits_full = self.router(
-            hidden_states_padded.to(self.router_params_dtype)
-        )
+        router_logits_full = self.router(hidden_states_padded)
 
         # MoERunner handles routing memoization and zero expert computation
         # internally. Pass full router_logits (including zero experts) so that
diff --git a/vllm/model_executor/models/longcat_flash_mtp.py b/vllm/model_executor/models/longcat_flash_mtp.py
index 360eb9c76378..770f931e81b8 100644
--- a/vllm/model_executor/models/longcat_flash_mtp.py
+++ b/vllm/model_executor/models/longcat_flash_mtp.py
@@ -7,7 +7,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -30,7 +30,7 @@
 class LongCatMultiTokenPredictorLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         vllm_config: VllmConfig,
         quant_config: QuantizationConfig | None = None,
@@ -313,8 +313,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
     def _rewrite_spec_layer_name(
         self, spec_layer: int, name: str, new_to_old_names_mapping: dict
     ) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
@@ -354,7 +353,7 @@ def _rewrite_spec_layer_name(
         return name
 
     def get_spec_layer_idx_from_weight_name(
-        self, config: PretrainedConfig, weight_name: str
+        self, config: PreTrainedConfig, weight_name: str
     ) -> int | None:
         if "model.mtp" in weight_name:
             return config.num_hidden_layers * 2
diff --git a/vllm/model_executor/models/longcat_flash_ngram.py b/vllm/model_executor/models/longcat_flash_ngram.py
index 76aa80013eb5..435a11ee17d8 100644
--- a/vllm/model_executor/models/longcat_flash_ngram.py
+++ b/vllm/model_executor/models/longcat_flash_ngram.py
@@ -156,6 +156,7 @@ def embed_batched(
             oe_ids: ``[num_tokens, num_embedders]`` global (offset) n-gram ids,
                 as produced by the ``ngram_compute_n_gram_ids`` kernel.
         Returns: ``[num_tokens, hidden]``.
+
         """
         word = self.word_embeddings(input_ids)  # [N, H]
         flat = oe_ids.permute(1, 0).contiguous()  # [num_embedders, N]
diff --git a/vllm/model_executor/models/mamba2.py b/vllm/model_executor/models/mamba2.py
index 343111ee0151..cb8b2b481631 100644
--- a/vllm/model_executor/models/mamba2.py
+++ b/vllm/model_executor/models/mamba2.py
@@ -197,6 +197,7 @@ def get_mamba_state_shape_from_config(
             Tuple containing:
             - conv_state_shape: Shape for convolutional state cache
             - temporal_state_shape: Shape for state space model cache
+
         """
         parallel_config = vllm_config.parallel_config
         hf_config = vllm_config.model_config.hf_config
diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py
index ebfbaf93fd6f..a75e3b0a8ecc 100644
--- a/vllm/model_executor/models/mellum.py
+++ b/vllm/model_executor/models/mellum.py
@@ -27,8 +27,7 @@
 
 
 class MellumAttention(Qwen3MoeAttention):
-    """
-    Differences from `Qwen3MoeAttention`:
+    """Differences from `Qwen3MoeAttention`:
     - Supports `per_layer_sliding_window` for `Attention`.
     """
 
@@ -122,8 +121,7 @@ def __init__(
 
 
 class MellumDecoderLayer(Qwen3MoeDecoderLayer):
-    """
-    Differences from `Qwen3MoeDecoderLayer`:
+    """Differences from `Qwen3MoeDecoderLayer`:
     - Supports interleaved SWA and per-layer RoPE scaling.
     """
 
@@ -185,8 +183,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
 
 @support_torch_compile
 class MellumModel(Qwen3MoeModel):
-    """
-    Differences from `Qwen3MoeModel`:
+    """Differences from `Qwen3MoeModel`:
     - Uses `MellumDecoderLayer`.
     """
 
@@ -204,8 +201,7 @@ def __init__(
 
 
 class MellumForCausalLM(Qwen3MoeForCausalLM):
-    """
-    Differences from `Qwen3MoeForCausalLM`:
+    """Differences from `Qwen3MoeForCausalLM`:
     - Uses `MellumModel`.
     """
 
diff --git a/vllm/model_executor/models/midashenglm.py b/vllm/model_executor/models/midashenglm.py
index 1129a0e4b284..6dae25ffb286 100644
--- a/vllm/model_executor/models/midashenglm.py
+++ b/vllm/model_executor/models/midashenglm.py
@@ -36,7 +36,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
@@ -61,6 +61,7 @@
     PromptUpdate,
     PromptUpdateDetails,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.midashenglm import DashengConfig
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -511,11 +512,9 @@ def forward(self, x, mask=None):
 
 # === Audio Inputs === #
 class MiDashengLMAudioInputs(TensorSchema):
-    """
-
-    Dimensions:
-        - bn: Batch size * number of audios
-        - p: Number of sampling points
+    """Dimensions:
+    - bn: Batch size * number of audios
+    - p: Number of sampling points
     """
 
     input_values: Annotated[torch.Tensor, TensorShape("n", "p")]
@@ -565,18 +564,13 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
-
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
-
         return {
             "audio": self._get_dummy_audios(
                 length=self.info.get_max_audio_len(),
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
@@ -584,35 +578,33 @@ def get_dummy_mm_data(
 class MiDashengLMMultiModalProcessor(
     BaseMultiModalProcessor[MiDashengLMProcessingInfo]
 ):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", [])
-        assert isinstance(audios, list)
-
-        min_audio_len = self.info.get_min_audio_len()
-        processed_audios = [
-            np.pad(
-                audio,
-                (0, min_audio_len - audio.shape[-1]),
-                mode="constant",
-                constant_values=0,
-            )
-            if isinstance(audio, np.ndarray) and audio.shape[-1] < min_audio_len
-            else audio
-            for audio in audios
-        ]
-
-        if processed_audios:
-            mm_data["audio"] = processed_audios
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        if audios := hf_inputs.hf_data.get("audio"):
+            assert isinstance(audios, list)
+
+            min_audio_len = self.info.get_min_audio_len()
+            hf_inputs.hf_data["audio"] = [
+                np.pad(
+                    audio,
+                    (0, min_audio_len - audio.shape[-1]),
+                    mode="constant",
+                    constant_values=0,
+                )
+                if isinstance(audio, np.ndarray) and audio.shape[-1] < min_audio_len
+                else audio
+                for audio in audios
+            ]
 
-        return mm_data, hf_processor_mm_kwargs
+        return hf_inputs
 
     def _get_mm_fields_config(
         self,
diff --git a/vllm/model_executor/models/mimo_audio.py b/vllm/model_executor/models/mimo_audio.py
index f1f901418a78..48c526bd92df 100644
--- a/vllm/model_executor/models/mimo_audio.py
+++ b/vllm/model_executor/models/mimo_audio.py
@@ -21,7 +21,7 @@
 import torch.nn.functional as F
 from einops import rearrange, repeat
 from transformers.activations import ACT2FN
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
 from transformers.models.qwen2.modeling_qwen2 import Qwen2Model
@@ -518,7 +518,7 @@ def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor:
 # ---------------------------------------------------------------------------
 
 
-class MiMoAudioTokenizerConfig(PretrainedConfig):
+class MiMoAudioTokenizerConfig(PreTrainedConfig):
     model_type = "mimo_audio_tokenizer"
 
     def __init__(
@@ -1315,6 +1315,7 @@ def process_audio(self, audio: torch.Tensor) -> torch.Tensor:
 
         Returns:
             [T//group_size, group_size, audio_channels]
+
         """
         T = audio.shape[0]
         audio = audio[:, : self.audio_channels]
@@ -1355,6 +1356,7 @@ def get_audio_feature(
             Tuple of:
             - audio_embeds: [total_tokens, out_hidden_size] concatenated embeddings
             - item_token_lens: list of int, number of tokens per input item
+
         """
         if self.audio_tokenizer is None:
             raise RuntimeError(
diff --git a/vllm/model_executor/models/mimo_mtp.py b/vllm/model_executor/models/mimo_mtp.py
index 32d3b1f662a7..65566230a562 100644
--- a/vllm/model_executor/models/mimo_mtp.py
+++ b/vllm/model_executor/models/mimo_mtp.py
@@ -24,7 +24,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -43,7 +43,7 @@
 class MiMoMultiTokenPredictorLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         model_config: ModelConfig,
         cache_config: CacheConfig | None = None,
diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py
index c7ba39bb0dd9..dfbed256551e 100644
--- a/vllm/model_executor/models/mimo_v2.py
+++ b/vllm/model_executor/models/mimo_v2.py
@@ -27,6 +27,7 @@
     FusedMoEFactory,
     fused_moe_make_expert_params_mapping,
 )
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -50,6 +51,7 @@
 )
 from vllm.model_executor.models.utils import sequence_parallel_chunk
 from vllm.sequence import IntermediateTensors
+from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backend import AttentionType
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
@@ -155,14 +157,16 @@ def __init__(
 
         dtype = getattr(config, "moe_router_dtype", "float32")
         self.gate_dtype = str_dtype_to_torch_dtype(dtype)
-        self.gate = nn.Linear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.n_routed_experts,
             bias=False,
-            dtype=self.gate_dtype,
+            params_dtype=self.gate_dtype,
+            out_dtype=torch.float32,
+            prefix=f"{prefix}.gate",
         )
         self.gate.e_score_correction_bias = nn.Parameter(
-            torch.empty(config.n_routed_experts, dtype=self.gate_dtype)
+            torch.empty(config.n_routed_experts, dtype=self.gate.out_dtype)
         )
 
         self.experts = FusedMoEFactory(
@@ -181,7 +185,7 @@ def __init__(
             num_expert_group=config.n_group,
             topk_group=config.topk_group,
             scoring_func="sigmoid",
-            router_logits_dtype=self.gate_dtype,
+            router_logits_dtype=self.gate.out_dtype,
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -197,7 +201,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             gate_input = hidden_states.to(self.gate_dtype)
         else:
             gate_input = hidden_states
-        router_logits = self.gate(gate_input)
+        router_logits, _ = self.gate(gate_input)
         final_hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
         )
@@ -460,6 +464,28 @@ def is_compressed_softmax_layer(self) -> bool:
         return self.config.hybrid_layer_pattern[self.layer_id] == 1
 
 
+def _requantize_fp8(
+    grouped: torch.Tensor, rows_rank: int, block: int, dtype: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Block-quantize a rank's ``[Q | K | V]`` rows back to fp8.
+
+    A rank's rows need not end on a block boundary (a single 1856-row slice is
+    14.5 blocks) while ``scaled_quantize`` requires both dims to be multiples of
+    the block size, so pad the tail with zeros: zero rows cannot raise a block's
+    amax, so every scale -- and the number of scale rows -- is unchanged. The
+    padding is dropped again here.
+    """
+    padded = cdiv(rows_rank, block) * block
+    if padded != rows_rank:
+        grouped = torch.cat(
+            [grouped, grouped.new_zeros(padded - rows_rank, grouped.shape[1])], dim=0
+        )
+    w_rank, s_rank = scaled_quantize(
+        grouped, GroupShape(block, block), dtype, compute_dtype=torch.float32
+    )
+    return w_rank[:rows_rank], s_rank
+
+
 def _shard_fp8_qkv_proj(
     w_full: torch.Tensor,
     s_full: torch.Tensor,
@@ -469,75 +495,136 @@ def _shard_fp8_qkv_proj(
     v_head_dim: int,
     tp_rank: int,
     tp_size: int,
+    ckpt_tp: int,
     block: int = 128,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Shard the fp8 qkv_proj weights for ``tp_rank``.
 
-    The checkpoint stores the fused QKV as ``num_kv_heads`` contiguous groups
-    (one per KV head; ``n`` below), each ordered ``[Q | K | V]``:
+    The checkpoint stores the fused QKV pre-sharded for ``ckpt_tp`` ranks (the
+    model config's ``num_key_value_heads``), each chunk holding that slice's
+    Q, K and V rows:
 
-        [Q_1 | K_1 | V_1 | Q_2 | K_2 | V_2 | ... | Q_n | K_n | V_n]
+        [Q_0 | K_0 | V_0 | Q_1 | K_1 | V_1 | ... | Q_n | K_n | V_n]   (n = ckpt_tp)
 
-    Per group, Q has ``(num_heads / num_kv_heads) * head_dim`` rows, K has
-    ``head_dim`` rows, and V has ``v_head_dim`` rows.
+    Per chunk, Q has ``(num_heads / ckpt_tp) * head_dim`` rows, K has
+    ``(num_kv_heads / ckpt_tp) * head_dim`` rows, and V has
+    ``(num_kv_heads / ckpt_tp) * v_head_dim`` rows, and the fp8 block scales
+    are tiled per chunk too (``ceil(rows_per_chunk / block)`` rows each).
 
-    Each TP rank owns ``g = num_kv_heads / tp_size`` of these groups, and the
-    forward expects them de-interleaved into a single Q, K, and V block:
+    ``ckpt_tp`` is not the layer's KV-head count: a MiMo-V2.5 SWA layer has 8
+    KV heads over 4 chunks, so each chunk carries two KV heads (3712 rows = 29
+    blocks) while a GA layer has 4 KV heads over 4 chunks (3392 rows each).
+
+    The forward expects each rank's slice de-interleaved:
 
         [Q_1 | Q_2 | ... | Q_g | K_1 | K_2 | ... | K_g | V_1 | V_2 | ... | V_g]
 
-    When ``g == 1`` the rank's slice is already ``[Q | K | V]``, so a plain
-    chunk suffices. When ``g > 1`` we cannot reach the de-interleaved layout by
-    re-permuting the fp8 block scales: each scale covers a 128-row block, and
-    since K is 192 rows (1.5 blocks) a block straddles the K/V boundary, so no
-    whole-block permutation produces it. Instead we dequantize this rank's
-    groups to float (dropping the block constraint), reorder the rows into the
-    layout above (Q, K, and V then each span a whole number of blocks), and
-    re-quantize to fp8.
+    When ``tp_size == ckpt_tp`` the checkpoint chunk *is* that layout, so a
+    plain chunk of both weight and scale suffices. Otherwise each rank's Q, K
+    and V rows are gathered from the chunks that hold them, dequantized with
+    the chunk's own scales, reordered, and re-quantized to fp8.
     """
-    assert tp_size <= num_kv_heads and num_kv_heads % tp_size == 0, (
-        "TP size must evenly split the number of KV heads."
+    assert num_heads % tp_size == 0, (
+        f"num_heads={num_heads} must be divisible by tp_size={tp_size}."
     )
-
-    kv_heads_per_rank = num_kv_heads // tp_size
-    if kv_heads_per_rank == 1:
-        # One KV head per rank. The weights and scale can be trivially sharded
-        # without re-quantization.
-        w = w_full.chunk(tp_size, dim=0)[tp_rank]
-        s = s_full.chunk(tp_size, dim=0)[tp_rank]
-        return w, s
-
-    q_rows_per_group = (num_heads // num_kv_heads) * head_dim
-    k_rows_per_group = head_dim
-    v_rows_per_group = v_head_dim
-    rows_per_group = q_rows_per_group + k_rows_per_group + v_rows_per_group
-    scale_rows_per_group = s_full.shape[0] // num_kv_heads
-    qs, ks, vs = [], [], []
-    for g_idx in range(tp_rank * kv_heads_per_rank, (tp_rank + 1) * kv_heads_per_rank):
-        row_start = g_idx * rows_per_group
-        scale_row_start = g_idx * scale_rows_per_group
-        # Dequantize this group's weights.
-        w_g = w_full[row_start : row_start + rows_per_group].to(torch.float32)
-        s_g = s_full[scale_row_start : scale_row_start + scale_rows_per_group].to(
-            torch.float32
-        )
-        s_g_expanded = s_g.repeat_interleave(block, dim=0).repeat_interleave(
-            block, dim=1
-        )[:rows_per_group]
-        w_g_dequant = w_g * s_g_expanded
-        # Track the dequantized q, k, and v weights separately.
-        qs.append(w_g_dequant[:q_rows_per_group])
-        ks.append(w_g_dequant[q_rows_per_group : q_rows_per_group + k_rows_per_group])
-        vs.append(w_g_dequant[q_rows_per_group + k_rows_per_group :])
-
-    # Combine the q, k, and v weights into the following layout:
-    # [Q_1, Q_2, .., Q_g, K_1, K_2, ..., K_g, V_1, V_2, ..., V_g]
-    grouped = torch.cat([torch.cat(qs), torch.cat(ks), torch.cat(vs)], dim=0)
-    # Quantize back to fp8.
-    return scaled_quantize(
-        grouped, GroupShape(block, block), w_full.dtype, compute_dtype=torch.float32
+    if ckpt_tp <= 0 or num_heads % ckpt_tp or num_kv_heads % ckpt_tp:
+        raise ValueError(
+            f"fused qkv_proj is pre-sharded at {ckpt_tp} chunks, which do not "
+            f"divide num_heads={num_heads} / num_kv_heads={num_kv_heads}."
+        )
+    # When there are fewer KV heads than ranks, vLLM replicates them
+    # (`num_kv_head_replicas`) and rank r owns KV head r // replicas, which
+    # keeps every Q head grouped with the KV head it attends to.
+    if tp_size <= num_kv_heads:
+        assert num_kv_heads % tp_size == 0, (
+            f"num_kv_heads={num_kv_heads} must be divisible by tp_size={tp_size}."
+        )
+        kv_head_ids = list(
+            range(
+                tp_rank * (num_kv_heads // tp_size),
+                (tp_rank + 1) * (num_kv_heads // tp_size),
+            )
+        )
+    else:
+        assert tp_size % num_kv_heads == 0, (
+            f"tp_size={tp_size} must be divisible by num_kv_heads={num_kv_heads}."
+        )
+        kv_head_ids = [tp_rank // (tp_size // num_kv_heads)]
+    q_head_ids = list(
+        range(tp_rank * (num_heads // tp_size), (tp_rank + 1) * (num_heads // tp_size))
     )
 
+    rows_per_chunk = w_full.shape[0] // ckpt_tp
+    q_per_chunk = (num_heads // ckpt_tp) * head_dim
+    k_per_chunk = (num_kv_heads // ckpt_tp) * head_dim
+    v_per_chunk = (num_kv_heads // ckpt_tp) * v_head_dim
+    if q_per_chunk + k_per_chunk + v_per_chunk != rows_per_chunk:
+        raise ValueError(
+            f"fused qkv_proj has {w_full.shape[0]} rows, not {ckpt_tp} chunks of "
+            f"{q_per_chunk} Q + {k_per_chunk} K + {v_per_chunk} V rows."
+        )
+
+    # The scales are tiled per chunk; they collapse to one continuous grid when
+    # a chunk is a whole number of blocks (the SWA chunks are: 3712 = 29 * 128).
+    chunk_scale_rows = cdiv(rows_per_chunk, block)
+    rows = torch.arange(w_full.shape[0])
+    per_chunk_scales = s_full.shape[0] == ckpt_tp * chunk_scale_rows
+    if per_chunk_scales:
+        scale_index = (rows // rows_per_chunk) * chunk_scale_rows + (
+            rows % rows_per_chunk
+        ) // block
+    elif s_full.shape[0] == cdiv(w_full.shape[0], block):
+        scale_index = rows // block
+    else:
+        raise ValueError(
+            f"fused qkv_proj scale has {s_full.shape[0]} rows, expected either "
+            f"{ckpt_tp * chunk_scale_rows} ({ckpt_tp} chunks of "
+            f"{chunk_scale_rows} rows) or {cdiv(w_full.shape[0], block)} "
+            f"(one continuous grid)"
+        )
+
+    if tp_size == ckpt_tp and per_chunk_scales:
+        # One checkpoint chunk per rank: already [Q | K | V] for that rank, and
+        # its scale rows line up with ceil(rows_per_chunk / block).
+        return (
+            w_full.chunk(ckpt_tp, dim=0)[tp_rank],
+            s_full.chunk(ckpt_tp, dim=0)[tp_rank],
+        )
+
+    # Gather this rank's Q, K and V rows from the chunks that hold them.
+    q_heads_per_chunk = num_heads // ckpt_tp
+    kv_heads_per_chunk = num_kv_heads // ckpt_tp
+    head_rows = torch.arange(head_dim)
+    v_head_rows = torch.arange(v_head_dim)
+    row_index: list[torch.Tensor] = []
+    for head in q_head_ids:
+        chunk = head // q_heads_per_chunk
+        row_index.append(
+            chunk * rows_per_chunk + (head % q_heads_per_chunk) * head_dim + head_rows
+        )
+    for head in kv_head_ids:
+        chunk = head // kv_heads_per_chunk
+        row_index.append(
+            chunk * rows_per_chunk
+            + q_per_chunk
+            + (head % kv_heads_per_chunk) * head_dim
+            + head_rows
+        )
+    for head in kv_head_ids:
+        chunk = head // kv_heads_per_chunk
+        row_index.append(
+            chunk * rows_per_chunk
+            + q_per_chunk
+            + k_per_chunk
+            + (head % kv_heads_per_chunk) * v_head_dim
+            + v_head_rows
+        )
+    index = torch.cat(row_index)
+    grouped = w_full[index].to(torch.float32) * s_full[
+        scale_index[index]
+    ].repeat_interleave(block, dim=1)
+    return _requantize_fp8(grouped, index.numel(), block, w_full.dtype)
+
 
 @support_torch_compile
 class MiMoV2Model(nn.Module, EagleModelMixin):
@@ -776,8 +863,7 @@ def _try_load_fp8_qkv_proj(
         tp_rank: int,
         tp_size: int,
     ) -> bool:
-        """
-        The fused fp8 QKV projection weights and scale are stored separately.
+        """The fused fp8 QKV projection weights and scale are stored separately.
         Special care must be taken while sharding these tensors across TP ranks.
         See _shard_fp8_qkv_proj for more details.
 
@@ -785,6 +871,7 @@ def _try_load_fp8_qkv_proj(
             True if ``tensor`` was an fp8 qkv_proj weight/scale and was consumed
             (caller should skip it); False otherwise, so the caller falls
             through to its normal loading path.
+
         """
         is_weight = (
             name.endswith("qkv_proj.weight") and tensor.dtype == torch.float8_e4m3fn
@@ -819,6 +906,8 @@ def _try_load_fp8_qkv_proj(
             v_head_dim=attn.v_head_dim,
             tp_rank=tp_rank,
             tp_size=tp_size,
+            # The fused qkv_proj is pre-sharded for this many ranks.
+            ckpt_tp=self.config.num_key_value_heads,
         )
         sharded = {"weight": w_rank, "weight_scale_inv": s_rank}
         for kind, tensor in sharded.items():
diff --git a/vllm/model_executor/models/mimo_v2_mtp.py b/vllm/model_executor/models/mimo_v2_mtp.py
index fa2cabe28860..a4e46977ec03 100644
--- a/vllm/model_executor/models/mimo_v2_mtp.py
+++ b/vllm/model_executor/models/mimo_v2_mtp.py
@@ -22,7 +22,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.distributed import (
@@ -45,7 +45,7 @@
     SupportsMultiModal,
     _require_is_multimodal,
 )
-from .mimo_v2 import MiMoV2Attention, MiMoV2MLP
+from .mimo_v2 import MiMoV2Attention, MiMoV2MLP, _shard_fp8_qkv_proj
 from .utils import _merge_multimodal_embeddings, maybe_prefix
 
 # MiMo-V2 checkpoints contain multiple MTP layers, but vLLM currently supports
@@ -62,7 +62,7 @@ class MiMoV2MTPLayer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         quant_config: QuantizationConfig | None = None,
     ) -> None:
@@ -141,11 +141,11 @@ def forward(
 
 
 class _MiMoV2MTPLayers(nn.Module):
-    """Thin wrapper so parameter paths match checkpoint: model.mtp.layers.*"""
+    """Thin wrapper so parameter paths match checkpoint: model.mtp.layers.*."""
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         num_mtp_layers: int,
         quant_config: QuantizationConfig | None,
         prefix: str,
@@ -266,6 +266,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
+        pending_qkv_proj: dict[str, dict[str, torch.Tensor]] = {}
 
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
@@ -279,14 +280,54 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ):
                 continue
 
-            # Support fused qkv_proj checkpoint (Pro format).
-            # The checkpoint is stored pre-sharded for TP=8 as
-            # [Q_rank0, K_rank0, V_rank0, Q_rank1, ...], so splitting along
-            # dim 0 with chunk(tp_size) gives each rank its Q+K+V slice for
-            # both the FP8 weight and the block weight_scale_inv. This matches
-            # how the main model loads the same layout.
+            # Fused qkv_proj: the fp8 weight and its block scales arrive as
+            # separate tensors and are pre-sharded at `config.num_key_value_heads`
+            # chunks of [Q_c | K_c | V_c] (see _shard_fp8_qkv_proj for the
+            # layout). Shard them together: a plain chunk(tp_size) only happens
+            # to be right when tp_size is that chunk count *and* the scales are
+            # tiled per chunk (they are 4 x 29 = 116 rows for MiMo-V2.5, not
+            # tp_size slices of an even split).
             if "qkv_proj" in name:
-                if name in params_dict:
+                is_fp8_weight = (
+                    name.endswith("qkv_proj.weight")
+                    and loaded_weight.dtype == torch.float8_e4m3fn
+                )
+                is_fp8_scale = name.endswith("qkv_proj.weight_scale_inv")
+                if is_fp8_weight or is_fp8_scale:
+                    prefix, kind = name.rsplit(".", 1)
+                    if (
+                        f"{prefix}.weight" not in params_dict
+                        or f"{prefix}.weight_scale_inv" not in params_dict
+                    ):
+                        continue
+                    entry = pending_qkv_proj.setdefault(prefix, {})
+                    entry[kind] = loaded_weight
+                    if "weight" not in entry or "weight_scale_inv" not in entry:
+                        # Waiting for the other half of the fused projection.
+                        continue
+                    del pending_qkv_proj[prefix]
+
+                    # Geometry comes from the module the weights belong to.
+                    attn = self.get_submodule(prefix.rsplit(".", 1)[0])
+                    w_rank, s_rank = _shard_fp8_qkv_proj(
+                        entry["weight"],
+                        entry["weight_scale_inv"],
+                        num_heads=attn.total_num_heads,
+                        num_kv_heads=attn.total_num_kv_heads,
+                        head_dim=attn.head_dim,
+                        v_head_dim=attn.v_head_dim,
+                        tp_rank=tp_rank,
+                        tp_size=tp_size,
+                        ckpt_tp=self.config.num_key_value_heads,
+                    )
+                    for k, tensor in (("weight", w_rank), ("weight_scale_inv", s_rank)):
+                        param = params_dict[f"{prefix}.{k}"]
+                        if tensor.shape[0] > param.shape[0]:
+                            tensor = tensor[: param.shape[0]]
+                        default_weight_loader(param, tensor)
+                        loaded_params.add(f"{prefix}.{k}")
+                elif name in params_dict:
+                    # Non-fp8 fused checkpoint: unchanged behaviour.
                     param = params_dict[name]
                     loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
                     default_weight_loader(param, loaded_weight)
diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py
index 93bd6a44b556..8c838c7cfb7f 100644
--- a/vllm/model_executor/models/mimo_v2_omni.py
+++ b/vllm/model_executor/models/mimo_v2_omni.py
@@ -10,16 +10,12 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import BatchFeature, PretrainedConfig
+from transformers import BatchFeature, PreTrainedConfig
 from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
 from typing_extensions import TypedDict
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import (
-    BaseDummyOptions,
-    ImageDummyOptions,
-    VideoDummyOptions,
-)
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import parallel_state
 from vllm.distributed import utils as dist_utils
 from vllm.inputs import MultiModalDataDict
@@ -46,6 +42,7 @@
     PromptUpdate,
     PromptUpdateDetails,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.transformers_utils.configs.mimo_v2_omni import Mimo_VLVisionConfig
 from vllm.transformers_utils.processors.mimo_v2_omni import (
     MiMoOmniProcessor,
@@ -55,6 +52,7 @@
 
 from .interfaces import (
     MultiModalEmbeddings,
+    SupportsEagle3,
     SupportsMultiModal,
     SupportsPP,
     SupportsQuant,
@@ -275,14 +273,14 @@ def forward(
         max_seqlen: torch.Tensor,
         full_attn: bool = True,
     ) -> torch.Tensor:
-        """
-        Args:
-            x: [seq_len, batch=1, embed_dim]  (seq-first convention)
-            cu_seqlens: cumulative sequence lengths [num_seqs+1], int32
-            rotary_pos_emb_cos: [seq_len, qk_channels // 2]
-            rotary_pos_emb_sin: [seq_len, qk_channels // 2]
-            max_seqlen: maximum sequence length
-            full_attn: if True, full attention; if False, window attention
+        """Args:
+        x: [seq_len, batch=1, embed_dim]  (seq-first convention)
+        cu_seqlens: cumulative sequence lengths [num_seqs+1], int32
+        rotary_pos_emb_cos: [seq_len, qk_channels // 2]
+        rotary_pos_emb_sin: [seq_len, qk_channels // 2]
+        max_seqlen: maximum sequence length
+        full_attn: if True, full attention; if False, window attention
+
         """
         # [seq_len, 1, embed_dim] -> QKV projection
         qkv, _ = self.qkv(x)  # [seq_len, 1, q_size + kv_size + kv_size]
@@ -420,7 +418,7 @@ class MiMoVisionTransformer(nn.Module):
 
     def __init__(
         self,
-        vision_cfg: PretrainedConfig,
+        vision_cfg: PreTrainedConfig,
         *,
         norm_eps: float = 1e-6,
         quant_config: QuantizationConfig | None = None,
@@ -531,6 +529,7 @@ def rot_pos_emb(self, grid_thw: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso
         Returns:
             cos: [total_tokens, qk_channels // 2]
             sin: [total_tokens, qk_channels // 2]
+
         """
         cos_list, sin_list = [], []
         for i in range(grid_thw.size(0)):
@@ -576,12 +575,12 @@ def rot_pos_emb(self, grid_thw: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso
         return torch.cat(cos_list, dim=0), torch.cat(sin_list, dim=0)
 
     def forward(self, x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             x: [total_tokens, C] pre-flattened patches
             grid_thw: [num_images, 3] tensor of (t, h, w) for each image/video
         Returns:
             [merged_tokens, out_hidden_size]
+
         """
         # Ensure grid_thw is a tensor
         if not isinstance(grid_thw, torch.Tensor):
@@ -929,27 +928,14 @@ def _get_mm_fields_config(
             fields["va_audio_features"] = MultiModalFieldConfig.batched("va_audio")
         return fields
 
-    def _apply_hf_processor_main(
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        """Convert numpy video arrays to (TCHW, timestamps) tuples for MiMo.
-        Also remap 'audios' → 'audio' since MiMoOmniProcessor.__call__ uses
-        the singular form.
-        """
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
-
-        # Remap audios → audio (MiMoOmniProcessor uses singular param name)
-        if "audios" in mm_data:
-            mm_data = {**mm_data, "audio": mm_data["audios"]}
-            mm_data = {k: v for k, v in mm_data.items() if k != "audios"}
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        """Convert numpy video arrays to (TCHW, timestamps) tuples for MiMo."""
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+        mm_data = hf_inputs.hf_data
 
         # Handle video_audio items: convert video part to (TCHW, timestamps) tuple
         if "video_audio" in mm_data:
@@ -964,7 +950,6 @@ def _apply_hf_processor_main(
                     assert isinstance(va_item, (list, tuple)) and len(va_item) == 2
                     vid, audio_src = va_item
                     va_item = VideoAudioInput(video=vid, audio=audio_src)
-                    vid = vid
                 # Convert video frames to (TCHW, timestamps) if needed
                 if (
                     isinstance(vid, tuple)
@@ -992,7 +977,7 @@ def _apply_hf_processor_main(
                             audio=va_item.audio,
                         )
                     )
-            mm_data = {**mm_data, "video_audio": va_converted}
+            mm_data["video_audio"] = va_converted
 
         if "videos" in mm_data:
             converted: list[tuple[torch.Tensor, torch.Tensor]] = []
@@ -1026,15 +1011,9 @@ def _apply_hf_processor_main(
                     timestamps = torch.arange(T, dtype=torch.float32) / self._INPUT_FPS
                     converted.append((frames, timestamps))
 
-            mm_data = {**mm_data, "videos": converted}
+            mm_data["videos"] = converted
 
-        processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            mm_data,
-            hf_processor_mm_kwargs,
-        )
-        processed_data.update(passthrough_data)
-        return processed_data
+        return hf_inputs
 
     def _get_prompt_updates(
         self,
@@ -1202,33 +1181,26 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(
             seq_len, mm_counts
         )
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
-        assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
 
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -1238,7 +1210,9 @@ def get_dummy_mm_data(
     info=MiMoV2OmniProcessingInfo,
     dummy_inputs=MiMoV2OmniDummyInputsBuilder,
 )
-class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant):
+class MiMoV2OmniForCausalLM(
+    nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant, SupportsEagle3
+):
     # To ensure correct weight loading and mapping.
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
@@ -1557,8 +1531,10 @@ def forward(
                 batch. **NOTE**: If mrope is enabled (default setting for
                 Qwen2.5-VL opensource models), the shape will be `(3, seq_len)`,
                 otherwise it will be `(seq_len,).
-        """
+            intermediate_tensors: Intermediate tensors from prior forward pass.
+            inputs_embeds: Optional tensor of input embeddings.
 
+        """
         if intermediate_tensors is not None:
             inputs_embeds = None
 
diff --git a/vllm/model_executor/models/minicpm.py b/vllm/model_executor/models/minicpm.py
index 0d7da44c3bf2..0c934386a7f1 100644
--- a/vllm/model_executor/models/minicpm.py
+++ b/vllm/model_executor/models/minicpm.py
@@ -31,7 +31,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -43,12 +43,15 @@
 )
 from vllm.model_executor.layers.activation import FatreluAndMul, SiluAndMul
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
+from vllm.model_executor.layers.fused_moe import (
+    GateLinear,
+    fused_experts,
+    fused_topk,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -109,12 +112,10 @@ def __init__(
             params_dtype = torch.get_default_dtype()
         self.params_dtype = params_dtype
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             self.hidden_size,
             self.num_total_experts,
-            bias=False,
             params_dtype=self.params_dtype,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
@@ -315,7 +316,7 @@ def forward(
 class MiniCPMDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -426,7 +427,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     def _init_layers(
         self,
         prefix: str,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None,
         quant_config: QuantizationConfig | None,
     ):
diff --git a/vllm/model_executor/models/minicpm3.py b/vllm/model_executor/models/minicpm3.py
index 30a9cd34ed0a..08b8ec47f33f 100644
--- a/vllm/model_executor/models/minicpm3.py
+++ b/vllm/model_executor/models/minicpm3.py
@@ -27,7 +27,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import CacheConfig, VllmConfig
 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -52,7 +52,7 @@
 class MiniCPM3Attention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         qk_nope_head_dim: int,
@@ -208,7 +208,7 @@ class MiniCPM3Model(MiniCPMModel):
     def _init_layers(
         self,
         prefix: str,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None,
         quant_config: QuantizationConfig | None,
     ):
diff --git a/vllm/model_executor/models/minicpm_eagle.py b/vllm/model_executor/models/minicpm_eagle.py
index a495fed22ce5..ebbbefd26b87 100644
--- a/vllm/model_executor/models/minicpm_eagle.py
+++ b/vllm/model_executor/models/minicpm_eagle.py
@@ -30,7 +30,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -68,7 +68,7 @@ class _EagleMiniCPMSupportsPP:
 class EagleMiniCPMDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -189,7 +189,7 @@ def __init__(
     def _init_layers(
         self,
         prefix: str,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None,
         quant_config: QuantizationConfig | None,
         start_layer: int,
diff --git a/vllm/model_executor/models/minicpmo.py b/vllm/model_executor/models/minicpmo.py
index c8b46a45f101..cc8a1066f2d9 100644
--- a/vllm/model_executor/models/minicpmo.py
+++ b/vllm/model_executor/models/minicpmo.py
@@ -40,7 +40,7 @@
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
 from vllm.multimodal.inputs import (
@@ -131,13 +131,12 @@
 
 
 class MiniCPMOAudioFeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - bns: Batch size * number of audios * number of slices
-        - bn: Batch size * number of audios
-        - c: Number of channels
-        - l: Length
-        - s: Number of slices
+    """Dimensions:
+    - bns: Batch size * number of audios * number of slices
+    - bn: Batch size * number of audios
+    - c: Number of channels
+    - l: Length
+    - s: Number of slices
     """
 
     type: Literal["audio_features"] = "audio_features"
@@ -163,8 +162,7 @@ class MiniCPMOAudioFeatureInputs(TensorSchema):
 
 
 class MiniCPMOAudioEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - bn: Batch size * number of audios
         - s: Number of slices
         - h: Hidden size (must match language model backbone)
@@ -404,22 +402,18 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
         audio_len = (
             self.info.get_max_audio_chunks_with_most_features()
             * self.info.get_default_audio_sampling_rate()
         )
 
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
-
         audio_mm_data = {
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
@@ -447,7 +441,7 @@ def process_audios(
         mm_data: Mapping[str, object],
         mm_kwargs: Mapping[str, object],
     ) -> Mapping[str, NestedTensors]:
-        if (audios := mm_data.get("audios")) is None:
+        if (audios := mm_data.get("audio")) is None:
             return {}
 
         mm_items = self.info.parse_mm_data({"audio": audios}, validate=False)
@@ -1008,8 +1002,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     dummy_inputs=MiniCPMODummyInputsBuilder,
 )
 class MiniCPMO(MiniCPMOBaseModel, MiniCPMV2_6):
-    """
-    MiniCPM-O model with audio support.
+    """MiniCPM-O model with audio support.
     Different versions use different LLM backbones:
     - Version 2.6: Uses Qwen2
     - Version 4.5: Uses Qwen3
diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py
index dac5e8367fbe..6555e2f132b6 100644
--- a/vllm/model_executor/models/minicpmv.py
+++ b/vllm/model_executor/models/minicpmv.py
@@ -36,7 +36,7 @@
 import torch.types
 from torch import nn
 from torch.nn.init import trunc_normal_
-from transformers import BatchFeature, PretrainedConfig
+from transformers import BatchFeature, PreTrainedConfig
 from transformers.dynamic_module_utils import (
     get_class_from_dynamic_module,
     resolve_trust_remote_code,
@@ -45,9 +45,8 @@
 
 from vllm.config import VllmConfig
 from vllm.config.multimodal import (
-    BaseDummyOptions,
     ImageDummyOptions,
-    VideoDummyOptions,
+    MultiModalDummyOptions,
 )
 from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -124,13 +123,12 @@
 
 
 class MiniCPMVImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bns: Batch size * number of images * number of slices
-        - bn: Batch size * number of images
-        - c: Number of channels
-        - h: Height
-        - w: Width
+    """Dimensions:
+    - bns: Batch size * number of images * number of slices
+    - bn: Batch size * number of images
+    - c: Number of channels
+    - h: Height
+    - w: Width
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -152,11 +150,10 @@ class MiniCPMVImagePixelInputs(TensorSchema):
 
 
 class MiniCPMVImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - ns: Number of slices
-        - hs: Hidden size (must match language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - ns: Number of slices
+    - hs: Hidden size (must match language model backbone)
     """
 
     type: Literal["image_embeds"]
@@ -307,8 +304,7 @@ def __init__(
     def get_1d_sincos_pos_embed_from_temporal_size(
         self, embed_dim: int, pos: np.ndarray
     ):
-        """
-        embed_dim: output dimension for each position
+        """embed_dim: output dimension for each position
         pos: a list of positions to be encoded: size (M,)
         out: (M, D)
         """
@@ -472,7 +468,7 @@ def forward(
         return x
 
 
-def get_version_by_config(config: PretrainedConfig) -> tuple[int, ...]:
+def get_version_by_config(config: PreTrainedConfig) -> tuple[int, ...]:
     version_float = getattr(config, "version", None)
 
     # The old configs do not include version number
@@ -851,9 +847,8 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         num_videos = mm_counts.get("video", 0)
 
         image_width, image_height = self.info.get_image_size_with_most_features()
@@ -862,14 +857,12 @@ def get_dummy_mm_data(
             seq_len, mm_counts
         )
 
-        image_overrides = mm_options.get("image")
         video_overrides = mm_options.get("video")
-        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
 
         # Convert video overrides to image overrides for per-frame image generation,
         # and apply num_frames override to num_video_frames.
         video_frame_overrides: ImageDummyOptions | None = None
-        if isinstance(video_overrides, VideoDummyOptions):
+        if video_overrides is not None:
             if video_overrides.num_frames:
                 num_video_frames = min(num_video_frames, video_overrides.num_frames)
             if video_overrides.width or video_overrides.height:
@@ -883,8 +876,8 @@ def get_dummy_mm_data(
             "image": self._get_dummy_images(
                 width=image_width,
                 height=image_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": [
                 self._get_dummy_images(
@@ -1062,19 +1055,18 @@ def _call_hf_processor_on_prompts(
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
         prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
 
         tokenizer = self.info.get_tokenizer()
 
         input_ids = torch.tensor([tokenizer.encode(prompt_text)])
-        mm_inputs = self.process_mm_inputs(mm_data, hf_processor_mm_kwargs)
+        mm_inputs = self.process_mm_inputs(mm_data, hf_kwargs)
 
         processed_data = BatchFeature(
             {
@@ -1082,8 +1074,9 @@ def _apply_hf_processor_main(
                 **mm_inputs,
             }
         )
-        processed_data.update(passthrough_data)
-        return processed_data
+        return self._finalize_hf_mm_data(
+            mm_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_prompt_updates(
         self,
@@ -1219,8 +1212,7 @@ def _get_mm_fields_config(
 
 
 class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):
-    """
-    The abstract class of MiniCPMV can only be inherited, but cannot be
+    """The abstract class of MiniCPMV can only be inherited, but cannot be
     instantiated.
     """
 
@@ -1243,7 +1235,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
         super().__init__()
         # All MiniCPM-V models disable `tie_word_embeddings` but
-        # `PretrainedConfig.tie_word_embeddings` defaults to True; we cannot
+        # `PreTrainedConfig.tie_word_embeddings` defaults to True; we cannot
         # check `tie_word_embeddings` until vLLM integrate MiniCPM-V model
         # and config class
         self.config = config
@@ -1419,9 +1411,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loaded
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="llm", connector="resampler", tower_model="vpm"
         )
@@ -1435,7 +1425,7 @@ def init_llm(
 
     def init_vision_module(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str = "",
     ) -> nn.Module:
@@ -1494,7 +1484,6 @@ def _mcpmv_normalize_tgt_sizes(
     slice_counts: list[int],
 ) -> torch.Tensor:
     """Normalize tgt_sizes to shape ``(total_slices, 2)``."""
-
     if isinstance(tgt_sizes, list):
         if not tgt_sizes:
             tgt_sizes = torch.zeros((0, 2), dtype=torch.long)
@@ -1968,7 +1957,7 @@ def init_llm(
 
     def init_vision_module(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str = "",
     ) -> nn.Module:
@@ -2069,7 +2058,7 @@ def init_llm(
 
     def init_vision_module(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str = "",
     ) -> nn.Module:
@@ -2165,7 +2154,7 @@ def init_llm(
 
     def init_vision_module(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> nn.Module:
@@ -2266,7 +2255,7 @@ def init_llm(
 
     def init_vision_module(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> nn.Module:
@@ -2365,7 +2354,7 @@ def init_llm(
 
     def init_vision_module(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> nn.Module:
@@ -2457,8 +2446,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
     dummy_inputs=MiniCPMVDummyInputsBuilder,
 )
 class MiniCPMV(MiniCPMVBaseModel, SupportsMultiModal, SupportsLoRA):
-    """
-    Different versions of MiniCPMV use different visual encoders and LLMs,
+    """Different versions of MiniCPMV use different visual encoders and LLMs,
     which is not conducive to the current integration logic of LoRA and
     bitsandbytes in vLLM. Therefore, it is necessary to separate them.
     """
diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py
index e43f13f79f7a..0a96c2d8479d 100644
--- a/vllm/model_executor/models/minicpmv4_6.py
+++ b/vllm/model_executor/models/minicpmv4_6.py
@@ -542,19 +542,22 @@ def _get_downsample_mode(
 
     def _compute_visual_tokens(
         self,
-        image_size,
+        image_size: ImageSize,
         max_slice_nums: int | None = None,
         downsample_mode: str | None = None,
     ) -> tuple[list[int], int, int]:
         """Compute grid, source_image_visual_tokens and patch_visual_tokens.
 
         Args:
+            image_size: Size of the source image.
+            max_slice_nums: Maximum number of slices, or None for the default.
             downsample_mode: ``"16x"`` (default, full merge) or ``"4x"``
                 (skip vit_merger, 4x more visual tokens).
 
         Returns:
             (grids, source_image_visual_tokens, patch_visual_tokens)
             grids is [0, 0] when no slicing occurs.
+
         """
         image_processor = self.get_image_processor()
         if max_slice_nums is None:
@@ -674,14 +677,6 @@ def get_num_image_tokens(
 
 
 class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module):
-    hf_to_vllm_mapper = WeightsMapper(
-        orig_to_new_stacked={
-            ".q_proj": (".qkv_proj", "q"),
-            ".k_proj": (".qkv_proj", "k"),
-            ".v_proj": (".qkv_proj", "v"),
-        }
-    )
-
     def __init__(
         self,
         config,
@@ -729,10 +724,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         out, _ = self.out_proj(attn_out)
         return out
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-
 
 class MiniCPMV4_6ViTWindowAttentionMerger(nn.Module):
     def __init__(
@@ -918,10 +909,10 @@ def forward(
         hidden_states: torch.Tensor,
         tgt_sizes: torch.Tensor,
     ) -> list[torch.Tensor]:
-        """
-        Args:
-            hidden_states: (B, max_patches, D) padded batch.
-            tgt_sizes: (B, 2) actual (H, W) per sample.
+        """Args:
+        hidden_states: (B, max_patches, D) padded batch.
+        tgt_sizes: (B, 2) actual (H, W) per sample.
+
         """
         m1, m2 = self.merge_kernel_size
         results = []
@@ -971,6 +962,11 @@ class MiniCPMV4_6ForConditionalGeneration(
     supports_encoder_tp_data = True
 
     hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+        },
         orig_to_new_prefix={
             # transformers v5.7+ uses `vision_tower` and nests `vit_merger`
             # inside it. Order matters: more specific prefix must come first.
@@ -982,7 +978,7 @@ class MiniCPMV4_6ForConditionalGeneration(
             "model.language_model.": "language_model.model.",
             "lm_head.": "language_model.lm_head.",
             "mtp.": None,
-        }
+        },
     )
 
     packed_modules_mapping = {
diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py
index 139529e55931..77f899315908 100644
--- a/vllm/model_executor/models/minimax_m2.py
+++ b/vllm/model_executor/models/minimax_m2.py
@@ -29,7 +29,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
@@ -70,7 +70,7 @@
 class MiniMaxM2MoE(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -249,7 +249,7 @@ def forward(
 class MiniMaxM2DecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         model_config: ModelConfig,
         cache_config: CacheConfig | None = None,
diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py
index 79c53b1c7e1c..a70d8e1de531 100644
--- a/vllm/model_executor/models/mistral3.py
+++ b/vllm/model_executor/models/mistral3.py
@@ -11,7 +11,7 @@
 from transformers.models.pixtral import PixtralProcessor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -58,12 +58,11 @@
 
 
 class Mistral3ImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height of each image
-        - w: Width of each image
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height of each image
+    - w: Width of each image
     """
 
     type: Literal["pixel_values_pixtral"] = "pixel_values_pixtral"
@@ -77,9 +76,7 @@ class Mistral3ImagePixelInputs(TensorSchema):
 
 
 class Mistral3PatchMerger(nn.Module):
-    """
-    Learned merging of spatial_merge_size ** 2 patches
-    """
+    """Learned merging of spatial_merge_size ** 2 patches."""
 
     def __init__(
         self, vision_hidden_size: int, spatial_merge_size: int, patch_size: int
@@ -195,18 +192,13 @@ def get_patch_grid_size(
             image_width = math.floor(image_width / ratio)
             image_height = math.floor(image_height / ratio)
 
-        patch_size = self.vision_config.patch_size
         assert isinstance(self.hf_config, Mistral3Config)
-        spatial_merge_size = self.hf_config.spatial_merge_size
-
-        # The HF processor rounds each dimension up to the vision patch size
-        # before the projector drops incomplete spatial-merge groups. This is
-        # not equivalent to rounding directly to the merged patch size.
-        num_width_patches = (image_width - 1) // patch_size + 1
-        num_height_patches = (image_height - 1) // patch_size + 1
+        merged_patch_size = (
+            self.vision_config.patch_size * self.hf_config.spatial_merge_size
+        )
 
-        ncols = num_width_patches // spatial_merge_size
-        nrows = num_height_patches // spatial_merge_size
+        ncols = (image_width - 1) // merged_patch_size + 1
+        nrows = (image_height - 1) // merged_patch_size + 1
         return ncols, nrows
 
 
@@ -251,7 +243,7 @@ def get_image_size_with_most_features(self) -> ImageSize:
         return ImageSize(width=width, height=height)
 
 
-class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder[Mistral3ProcessingInfo]):
+class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder):
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_images = mm_counts.get("image", 0)
 
@@ -264,26 +256,22 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
 
 class Mistral3MultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -365,6 +353,7 @@ def _get_num_hidden_layers(hf_config: Mistral3Config) -> int:
 
     Args:
         hf_config: Model config with vision feature layer(s).
+
     """
     feature_layers = hf_config.vision_feature_layer
     num_hidden_layers = hf_config.vision_config.num_hidden_layers
@@ -588,9 +577,12 @@ def forward(
             positions: Position indices for the input tokens.
             intermediate_tensors: Intermediate tensors from prior forward pass.
             inputs_embeds: Optional tensor of input embeddings.
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
 
         Info:
             [`Mistral3ImagePixelInputs`][vllm.model_executor.models.mistral3.Mistral3ImagePixelInputs]
+
         """
         if intermediate_tensors is not None:
             inputs_embeds = None
@@ -612,9 +604,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py
index 717c935751e4..248459ceb514 100644
--- a/vllm/model_executor/models/mixtral.py
+++ b/vllm/model_executor/models/mixtral.py
@@ -39,11 +39,13 @@
     get_tensor_model_parallel_world_size,
 )
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -106,12 +108,10 @@ def __init__(
         self.n_local_physical_experts = self.n_physical_experts // self.ep_size
         # Gate always runs at half / full precision for now.
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             hidden_size,
             num_experts,
-            bias=False,
             params_dtype=params_dtype,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py
index 5a6e65441703..fbf393b554b6 100644
--- a/vllm/model_executor/models/mllama4.py
+++ b/vllm/model_executor/models/mllama4.py
@@ -18,7 +18,6 @@
 # limitations under the License.
 import math
 from collections.abc import Hashable, Iterable, Mapping
-from itertools import tee
 from typing import Annotated, Any, Literal
 
 import torch
@@ -26,7 +25,7 @@
 from transformers import BatchFeature, Llama4Config, Llama4VisionConfig
 from transformers.image_utils import SizeDict
 from transformers.models.llama4 import Llama4Processor
-from transformers.models.llama4.image_processing_llama4_fast import (
+from transformers.models.llama4.image_processing_llama4 import (
     find_supported_resolutions,
     get_best_fit,
 )
@@ -36,7 +35,7 @@
     support_torch_compile,
 )
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.attention import MMEncoderAttention
@@ -60,6 +59,7 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
@@ -90,12 +90,11 @@
 
 
 class Llama4ImagePatchInputs(TensorSchema):
-    """
-    Dimensions:
-        - batch_size: Batch size
-        - total_num_chunks: Batch size * number of chunks
-        - num_channels: Number of channels
-        - image_size: Size of each image
+    """Dimensions:
+    - batch_size: Batch size
+    - total_num_chunks: Batch size * number of chunks
+    - num_channels: Number of channels
+    - image_size: Size of each image
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -412,16 +411,15 @@ def forward(
         self,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
-        r"""
-        Args:
-            hidden_states: Input tensor of shape
-                (batch_size, sequence_length, hidden_size).
-                Hidden states from the model embeddings, representing
-                the input tokens.
-                associated vectors than the model's internal embedding
-                lookup matrix.
-        """
+        r"""Args:
+        hidden_states: Input tensor of shape
+            (batch_size, sequence_length, hidden_size).
+            Hidden states from the model embeddings, representing
+            the input tokens.
+            associated vectors than the model's internal embedding
+            lookup matrix.
 
+        """
         for encoder_layer in self.layers:
             layer_outputs = encoder_layer(hidden_states)
             hidden_states = layer_outputs[0]
@@ -556,9 +554,7 @@ def get_hf_config(self) -> Llama4Config:
         return self.ctx.get_hf_config(Llama4Config)
 
     def get_hf_processor(self, **kwargs: object) -> Llama4Processor:
-        return self.ctx.get_hf_processor(
-            Llama4Processor, use_fast=kwargs.pop("use_fast", True), **kwargs
-        )
+        return self.ctx.get_hf_processor(Llama4Processor, **kwargs)
 
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         # Although vLLM can support more images from an infra capability
@@ -590,7 +586,7 @@ def get_image_size_with_most_features(self) -> ImageSize:
 
 
 class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -710,20 +706,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         (target_width, target_height) = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -1046,25 +1038,6 @@ def compute_logits(
     ) -> torch.Tensor | None:
         return self.language_model.compute_logits(hidden_states)
 
-    def separate_weights(
-        self,
-        weights: Iterable[tuple[str, torch.Tensor]],
-        prefix: str,
-    ) -> tuple[Iterable[tuple[str, torch.Tensor]], Iterable[tuple[str, torch.Tensor]]]:
-        weights1, weights2 = tee(weights, 2)
-
-        def get_prefix_weights() -> Iterable[tuple[str, torch.Tensor]]:
-            for name, data in weights1:
-                if name.startswith(prefix):
-                    yield (name, data)
-
-        def get_other_weights() -> Iterable[tuple[str, torch.Tensor]]:
-            for name, data in weights2:
-                if not name.startswith(prefix):
-                    yield (name, data)
-
-        return get_prefix_weights(), get_other_weights()
-
     def _consolidate_qkv_weights(
         self, weights: Iterable[tuple[str, torch.Tensor]]
     ) -> Iterable[tuple[str, torch.Tensor]]:
@@ -1269,9 +1242,7 @@ def regular_language_model_weights() -> Iterable[tuple[str, torch.Tensor]]:
         return updated_params
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector=[
@@ -1281,21 +1252,18 @@ def get_mm_mapping(self) -> MultiModelKeys:
             tower_model="vision_model.",
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         vision_config = self.config.vision_config
         patches_per_chunk = Mllama4ProcessingInfo.get_patch_per_chunk(vision_config)
-        if num_image_tokens <= 0 or patches_per_chunk <= 0:
-            return 0
+        if num_mm_embeds <= 0 or patches_per_chunk <= 0:
+            return 0, 0
         raw_patches = (vision_config.image_size // vision_config.patch_size) ** 2
-        num_chunks = num_image_tokens // patches_per_chunk
-        # Encoder processes raw_patches + 1 (CLS) per chunk
-        return num_chunks * (raw_patches + 1)
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
-        vision_config = self.config.vision_config
-        raw_patches = (vision_config.image_size // vision_config.patch_size) ** 2
-        if num_vision_tokens <= 0:
-            return 0
-        num_chunks = num_vision_tokens // (raw_patches + 1)
-        patches_per_chunk = Mllama4ProcessingInfo.get_patch_per_chunk(vision_config)
-        return num_chunks * patches_per_chunk
+        num_chunks = num_mm_embeds // patches_per_chunk
+        return num_chunks * (raw_patches + 1), num_chunks * patches_per_chunk
diff --git a/vllm/model_executor/models/mlp_speculator.py b/vllm/model_executor/models/mlp_speculator.py
index 612baba8eaa6..0e58952102d6 100644
--- a/vllm/model_executor/models/mlp_speculator.py
+++ b/vllm/model_executor/models/mlp_speculator.py
@@ -19,10 +19,10 @@
 
 
 class MLPSpeculatorLayerNorm(nn.Module):
-    """
-    A L2 normalization implementation
+    """A L2 normalization implementation
     ...
-    Args
+
+    Args:
     ----
     normalized_shape : int
         Dimensionality of input data (size of final tensor axis)
@@ -32,6 +32,7 @@ class MLPSpeculatorLayerNorm(nn.Module):
          (i.e. fp16 requires eps >= 6e-8).
     elementwise_scale_and_shift : bool
         Include a learned scaling and shift term after normalization.
+
     """
 
     def __init__(
@@ -58,8 +59,7 @@ def forward(self, x):
 
 
 class MLPSpeculator(nn.Module):
-    """
-    An implementation of the speculative models introduced in
+    """An implementation of the speculative models introduced in
     "Accelerating Production LLMs with Combined Token/Embedding
     Speculators"
     https://arxiv.org/pdf/2404.19124
diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py
index 8fe80bbe19f1..df901a6566d9 100644
--- a/vllm/model_executor/models/modernbert.py
+++ b/vllm/model_executor/models/modernbert.py
@@ -32,7 +32,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.sequence import IntermediateTensors
 
-from .interfaces import SupportsCrossEncoding
+from .interfaces import SupportsCrossEncoding, SupportsLoRA
 from .interfaces_base import attn_type, default_pooling_type
 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 
@@ -262,7 +262,7 @@ def forward(
 
 @support_torch_compile
 @default_pooling_type(seq_pooling_type="CLS")
-class ModernBertModel(nn.Module):
+class ModernBertModel(nn.Module, SupportsLoRA):
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "model.layers.": "encoder_layer.layers.",
@@ -271,6 +271,11 @@ class ModernBertModel(nn.Module):
         }
     )
 
+    packed_modules_mapping = {
+        "Wqkv": ["Wqkv"],
+        "Wi": ["Wi"],
+    }
+
     def __init__(
         self,
         vllm_config: VllmConfig,
diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py
index 502b8a5ce6ca..ad6b1c4a8b4b 100644
--- a/vllm/model_executor/models/molmo.py
+++ b/vllm/model_executor/models/molmo.py
@@ -16,12 +16,12 @@
 from transformers import (
     BaseImageProcessor,
     BatchFeature,
-    PretrainedConfig,
+    PreTrainedConfig,
 )
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_rank,
@@ -93,13 +93,12 @@
 
 
 class MolmoImageInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - bnc: Batch size * number of images * number of crops (dynamic)
-        - np: Number of patches
-        - tp: Token sequence positions
-        - pd: Patch dimension
+    """Dimensions:
+    - bn: Batch size * number of images
+    - bnc: Batch size * number of images * number of crops (dynamic)
+    - np: Number of patches
+    - tp: Token sequence positions
+    - pd: Patch dimension
     """
 
     images: Annotated[torch.Tensor, TensorShape("bnc", "np", "pd")]
@@ -382,9 +381,7 @@ def add_pos_emb(self, x: torch.Tensor, patch_num: tuple[int, int]) -> torch.Tens
     def forward(
         self, x: torch.Tensor, patch_num: tuple[int, int] | None = None
     ) -> list[torch.Tensor]:
-        """
-        : param x: (batch_size, num_patch, n_pixels)
-        """
+        """: param x: (batch_size, num_patch, n_pixels)."""
         if patch_num is None:
             patch_num = self.patch_num
         B, N, D = x.shape
@@ -408,7 +405,7 @@ class MolmoAttention(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -518,7 +515,7 @@ class LanguageModelMLP(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         input_dim: int | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -560,7 +557,7 @@ class ImageProjectorMLP(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         input_dim: int | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -601,7 +598,7 @@ def forward(
 class MolmoDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -684,7 +681,7 @@ class MolmoVisionBackbone(nn.Module, SupportsQuant):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         vision_config: VisionBackboneConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -728,9 +725,7 @@ def device(self) -> torch.device:
         return self.image_vit.patch_embedding.weight.device
 
     def encode_image(self, images: torch.Tensor) -> torch.Tensor:
-        """
-        : param images: (batch_size, num_crops, num_patch, n_pixels)
-        """
+        """: param images: (batch_size, num_crops, num_patch, n_pixels)."""
         B, T, N, D = images.shape
 
         mask = ~torch.all(images.view(B * T, N, D) == -1, dim=(1, 2), keepdim=True)
@@ -1109,24 +1104,24 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         target_width, target_height = self.info.get_image_size_with_most_features()
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
 
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
 
 class MolmoMultiModalProcessor(BaseMultiModalProcessor[MolmoProcessingInfo]):
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _postprocess_prompt(self, prompt: list[int]) -> list[int]:
         processor = self.info.get_hf_processor()
 
@@ -1145,28 +1140,28 @@ def _postprocess_prompt(self, prompt: list[int]) -> list[int]:
             dict(tokens=tokens),
         )["input_ids"].tolist()
 
-    def _apply_hf_processor_main(
+    def _call_hf_processor(
         self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        mm_counts = mm_items.get_all_counts()
-
-        valid_mm_items = mm_items.select({k for k, c in mm_counts.items() if c > 0})
-        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not processor_data:
-            return BatchFeature(dict(passthrough_data))
-
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_counts)
-
-        processed_data = self.info.ctx.call_hf_processor(
+        hf_processor = self.info.get_hf_processor(**hf_kwargs)
+        return self.info.ctx.call_hf_processor(
             hf_processor.process,
-            dict(text=prompt_text, **processor_data),
-            hf_processor_mm_kwargs,
+            hf_data,
+            hf_kwargs,
         )
 
+    def _postprocess_hf_mm_data(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+        processed_data: BatchFeature,
+    ) -> BatchFeature:
+        if not hf_data:
+            return processed_data
+
+        hf_processor = self.info.get_hf_processor(**hf_kwargs)
         tokenizer = hf_processor.tokenizer
         image_patch_id = tokenizer.vocab[IMAGE_PATCH_TOKEN]
 
@@ -1174,7 +1169,7 @@ def _apply_hf_processor_main(
 
         processed_data.pop("input_ids")
 
-        if (images := processor_data.get("images")) is not None:
+        if (images := hf_data.get("images")) is not None:
             mm_items = self.info.parse_mm_data({"image": images}, validate=False)
             parsed_images = mm_items.get_items("image", ImageProcessorItems)
             image_sizes = [
@@ -1198,8 +1193,6 @@ def _apply_hf_processor_main(
             processed_data["num_crops"] = num_crops
             processed_data["img_patch_id"] = image_patch_id
 
-        processed_data.update(passthrough_data)
-
         return processed_data
 
     def _get_mm_fields_config(
@@ -1456,9 +1449,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="model",
             connector="vision_backbone.image_projector",
diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py
index 66de39929424..4598d9da2af2 100644
--- a/vllm/model_executor/models/molmo2.py
+++ b/vllm/model_executor/models/molmo2.py
@@ -11,13 +11,12 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from PIL import ImageOps
 from PIL.Image import Image
 from transformers import (
     BaseImageProcessor,
     BaseVideoProcessor,
     BatchFeature,
-    PretrainedConfig,
+    PreTrainedConfig,
     ProcessorMixin,
 )
 from transformers.image_utils import ImageInput
@@ -26,7 +25,7 @@
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_rank,
@@ -54,6 +53,7 @@
 )
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.image import normalize_image
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalKwargsItems,
@@ -105,15 +105,14 @@
 
 
 class Molmo2ImageInputs(TensorSchema):
-    """
-    Dimensions:
-        - nc: The total number of crops (dynamic)
-        - np: The total number of patches per crop
-        - cps: Number of channels * patch_size * patch_size
-        - npp: Number of pooled patches (dynamic)
-        - pp: pooling_size * pooling_size
-        - ni: Number of images
-        - nt: Number of image tokens (dynamic)
+    """Dimensions:
+    - nc: The total number of crops (dynamic)
+    - np: The total number of patches per crop
+    - cps: Number of channels * patch_size * patch_size
+    - npp: Number of pooled patches (dynamic)
+    - pp: pooling_size * pooling_size
+    - ni: Number of images
+    - nt: Number of image tokens (dynamic)
     """
 
     pixel_values: Annotated[torch.Tensor, TensorShape("nc", "np", "cps")]
@@ -132,15 +131,14 @@ class Molmo2ImageInputs(TensorSchema):
 
 
 class Molmo2VideoInputs(TensorSchema):
-    """
-    Dimensions:
-        - nc: The total number of frames (dynamic)
-        - np: The total number of patches per frame
-        - cps: Number of channels * patch_size * patch_size
-        - npp: Number of pooled patches (dynamic)
-        - pp: pooling_size * pooling_size
-        - nv: Number of videos
-        - nt: Number of video tokens (dynamic)
+    """Dimensions:
+    - nc: The total number of frames (dynamic)
+    - np: The total number of patches per frame
+    - cps: Number of channels * patch_size * patch_size
+    - npp: Number of pooled patches (dynamic)
+    - pp: pooling_size * pooling_size
+    - nv: Number of videos
+    - nt: Number of video tokens (dynamic)
     """
 
     pixel_values_videos: Annotated[torch.Tensor, TensorShape("nc", "np", "cps")]
@@ -165,7 +163,7 @@ class Molmo2MultiModalInputs(TypedDict, total=False):
 
 @dataclass
 class VitConfig:
-    """Config for a vision transformer"""
+    """Config for a vision transformer."""
 
     hidden_size: int = 1152
     intermediate_size: int = 4304
@@ -190,7 +188,7 @@ def image_num_patch(self):
 
 @dataclass
 class AdapterConfig:
-    """Config for a vit-llm adapter"""
+    """Config for a vit-llm adapter."""
 
     vit_layers: tuple[int, int] = (-3, -9)
     pooling_attention_mask: bool = False
@@ -205,7 +203,7 @@ class AdapterConfig:
 
 @dataclass
 class TextConfig:
-    """Configuration for a text model transformer"""
+    """Configuration for a text model transformer."""
 
     hidden_size: int = 3584
     """
@@ -530,9 +528,7 @@ def forward(
         x: torch.Tensor,
         patch_num: tuple[int, int] | None = None,
     ) -> list[torch.Tensor]:
-        """
-        : param x: (batch_size, num_patch, n_pixels)
-        """
+        """: param x: (batch_size, num_patch, n_pixels)."""
         if patch_num is None:
             patch_num = self.patch_num
 
@@ -545,7 +541,7 @@ def forward(
 
 
 class ImagePoolingAttention(nn.Module):
-    """Multi-head attention used for image pooling"""
+    """Multi-head attention used for image pooling."""
 
     def __init__(
         self,
@@ -667,7 +663,7 @@ def forward(
 
 
 class ImageProjectorMLP(nn.Module):
-    """MLP used for the image projector"""
+    """MLP used for the image projector."""
 
     def __init__(
         self,
@@ -786,9 +782,7 @@ def device(self) -> torch.device:
         return self.image_vit.patch_embedding.weight.device
 
     def encode_image(self, images: torch.Tensor) -> torch.Tensor:
-        """
-        : param images: (batch_size, num_crops, num_patch, n_pixels)
-        """
+        """: param images: (batch_size, num_crops, num_patch, n_pixels)."""
         B, T, N, D = images.shape
         images = images.view(B * T, N, D)
         image_features = self.image_vit(images)
@@ -1305,13 +1299,13 @@ def exif_transpose(
             exif_transpose(img) if isinstance(img, Image) else img for img in images
         ]
     elif images is not None and isinstance(images, Image):
-        images = ImageOps.exif_transpose(images)
+        images = normalize_image(images)
     return images
 
 
 def build_flat_image_bool_length(
     image_grids: torch.LongTensor,
-    hf_config: PretrainedConfig,
+    hf_config: PreTrainedConfig,
     image_use_col_tokens: bool = True,
     use_single_crop_col_tokens: bool | None = None,
     use_single_crop_start_token: bool = True,
@@ -1399,7 +1393,7 @@ def build_flat_image_bool_length(
 
 def build_flat_video_bool_length(
     video_grids: torch.LongTensor,
-    hf_config: PretrainedConfig,
+    hf_config: PreTrainedConfig,
 ) -> tuple[torch.LongTensor, torch.LongTensor]:
     image_patch_id = hf_config.image_patch_id
     frame_start_id = hf_config.frame_start_token_id
@@ -1444,8 +1438,7 @@ def get_candidate_target_fps(
     sampling_fps: int | float,
     max_fps: int | float = _MAX_VIDEO_FPS,
 ) -> list[float]:
-    """
-    Return the subset of `video_fps` factors that remain multiples
+    """Return the subset of `video_fps` factors that remain multiples
     of `sampling_fps`.
 
     Examples:
@@ -1460,6 +1453,7 @@ def get_candidate_target_fps(
             ...
         ValueError: sampling_fps=2 must divide video_fps=5 to produce
             consistent frame steps.
+
     """
     video_fps = int(video_fps)
     sampling_fps = int(sampling_fps)
@@ -1494,9 +1488,7 @@ def get_target_fps(
     frame_sample_mode: str,
     candidate_target_fps: list[float],
 ) -> float | None:
-    """
-    Get the target fps that best spans the video and has the most frames sampled
-    """
+    """Get the target fps that best spans the video and has the most frames sampled."""
     num_frames_sampled = 0
     selected_target_fps = None
     for target_fps in candidate_target_fps:
@@ -1842,7 +1834,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         num_images = mm_counts.get("image", 0)
         num_videos = mm_counts.get("video", 0)
@@ -1853,13 +1845,11 @@ def get_dummy_mm_data(
         if num_images > 0:
             target_width, target_height = self.info.get_image_size_with_most_features()
 
-            image_overrides = mm_options.get("image")
-
             dummy_images = self._get_dummy_images(
                 width=target_width,
                 height=target_height,
                 num_images=num_images,
-                overrides=image_overrides,
+                overrides=mm_options.get("image"),
             )
 
         if num_videos > 0:
@@ -1872,7 +1862,6 @@ def get_dummy_mm_data(
             video_overrides = mm_options.get("video")
 
             if video_overrides:
-                assert isinstance(video_overrides, VideoDummyOptions)
                 num_frames_override = video_overrides.num_frames
                 if num_frames_override:
                     if num_frames_override > target_num_frames:
@@ -1940,6 +1929,9 @@ def _get_dummy_videos(
 
 
 class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _postprocess_prompt(self, prompt: list[int]) -> list[int]:
         processor = self.info.get_hf_processor()
         tokenizer = processor.tokenizer
@@ -1954,19 +1946,20 @@ def _postprocess_prompt(self, prompt: list[int]) -> list[int]:
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        mm_counts = mm_items.get_all_counts()
-
-        valid_mm_items = mm_items.select({k for k, c in mm_counts.items() if c > 0})
-        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
+        )
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_counts)
+        if not hf_data:
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
-        mm_data = dict(processor_data)
+        prompt_text = hf_data.pop("text")
+        assert isinstance(prompt_text, str)
 
         hf_config = self.info.get_hf_config()
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        hf_processor = self.info.get_hf_processor(**hf_kwargs)
 
         def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature:
             res = hf_processor(text=text, images=images, videos=videos, **kwargs)
@@ -1981,7 +1974,7 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature:
         tokenizer = hf_processor.tokenizer
         image_processor = hf_processor.image_processor
 
-        if videos := mm_data.pop("videos", []):
+        if videos := hf_data.pop("videos", []):
             assert isinstance(videos, Sequence)
             bos_token_id = tokenizer.bos_token_id or tokenizer.eos_token_id
 
@@ -2001,16 +1994,16 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature:
                 # NOTE: metadata.frames_indices indicates
                 # the sampled frames indices of pre-sampled videos, which is
                 # used to calculate the timestamps. Make sure that
-                # do_sample_frames in hf_processor_mm_kwargs is false for
+                # do_sample_frames in hf_kwargs is false for
                 # presampled videos.
 
-                # NOTE: a copy of hf_processor_mm_kwargs is created to update
+                # NOTE: a copy of hf_kwargs is created to update
                 # do_sample_frames, otherwise mm_hash for the object will be
                 # incorrect.
-                video_mm_kwargs = dict(**hf_processor_mm_kwargs)
+                video_mm_kwargs = dict(**hf_kwargs)
                 if "do_sample_frames" not in video_mm_kwargs:
                     # molmo_utils already has "do_sample_frames" in
-                    # hf_processor_mm_kwargs, don't overwrite it.
+                    # hf_kwargs, don't overwrite it.
                     video_mm_kwargs["do_sample_frames"] = metadata.get(
                         "do_sample_frames", False
                     )
@@ -2075,11 +2068,11 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature:
 
         processed_data = self.info.ctx.call_hf_processor(
             patched_call,
-            dict(text=prompt_text, **mm_data),
-            hf_processor_mm_kwargs,
+            dict(text=prompt_text, **hf_data),
+            hf_kwargs,
         )
 
-        if (images := mm_data.get("images")) is not None:
+        if (images := hf_data.get("images")) is not None:
             mm_items = self.info.parse_mm_data({"image": images}, validate=False)
             parsed_images = mm_items.get_items("image", ImageProcessorItems)
             image_sizes = [
@@ -2121,10 +2114,11 @@ def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature:
             )
 
         processed_data.update(all_video_outputs)
-        processed_data.update(passthrough_data)
         processed_data.pop("input_ids")
 
-        return processed_data
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_mm_fields_config(
         self,
@@ -2646,9 +2640,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="model",
             connector="vision_backbone.image_projector",
diff --git a/vllm/model_executor/models/moondream3.py b/vllm/model_executor/models/moondream3.py
index 818015e1ba7c..278167d0f73f 100644
--- a/vllm/model_executor/models/moondream3.py
+++ b/vllm/model_executor/models/moondream3.py
@@ -13,7 +13,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_rank,
@@ -28,12 +28,15 @@
 from vllm.model_executor.layers.attention.mm_encoder_attention import (
     MMEncoderAttention,
 )
-from vllm.model_executor.layers.fused_moe import MoEActivation, fused_experts
+from vllm.model_executor.layers.fused_moe import (
+    GateLinear,
+    MoEActivation,
+    fused_experts,
+)
 from vllm.model_executor.layers.fused_moe.config import biased_moe_quant_config
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -47,6 +50,7 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import ImageSize, MultiModalDataItems
@@ -294,6 +298,7 @@ def create_patches(self, images: torch.Tensor) -> torch.Tensor:
 
         Returns:
             patches: (batch, num_patches, patch_dim)
+
         """
         patch_size = self.config.enc_patch_size
         batch, channels, height, width = images.shape
@@ -320,6 +325,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
 
         Returns:
             features: (batch, num_patches, hidden_size)
+
         """
         # Create patches and embed
         patches = self.create_patches(pixel_values)
@@ -450,11 +456,10 @@ def __init__(
         self.num_local_experts = self.experts_per_rank
 
         # Router (gate) - use ReplicatedLinear for compatibility
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             hidden_size,
             num_experts,
             bias=True,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
@@ -488,7 +493,6 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass with expert parallelism and custom GeGLU activation."""
-
         # Get router logits and compute top-k
         router_logits, _ = self.gate(x)  # [num_tokens, num_experts]
         topk_logits, topk_ids = torch.topk(
@@ -941,15 +945,14 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+        mm_options: MultiModalDummyOptions | None = None,
         mm_processor_kwargs: Mapping[str, object] | None = None,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         return {
             "image": self._get_dummy_images(
                 width=378,
                 height=378,
-                num_images=num_images,
+                num_images=mm_counts.get("image", 0),
             )
         }
 
@@ -1104,11 +1107,15 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
     def get_language_model(self) -> nn.Module:
         return self.text
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        return num_image_tokens
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
-        return num_vision_tokens
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+        return num_mm_embeds, num_mm_embeds
 
     def _split_pixel_values(
         self,
@@ -1313,7 +1320,6 @@ def compute_logits(
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights with remapping from HuggingFace format."""
-
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
 
diff --git a/vllm/model_executor/models/moonvit.py b/vllm/model_executor/models/moonvit.py
index 9269cd29766d..4e43c9c364de 100644
--- a/vllm/model_executor/models/moonvit.py
+++ b/vllm/model_executor/models/moonvit.py
@@ -80,13 +80,14 @@ def _apply_rope_input_validation(x, freqs_cis):
 def apply_rope(
     xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Args: (The leading dimensions of all inputs should be the same)
+    """Args: (The leading dimensions of all inputs should be the same)
         xq: query, tensor of shape (..., num_heads, head_dim)
         xk: key, tensor of shape (..., num_heads, head_dim)
         freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid.
+
     Returns:
         xq_out, xk_out: tensors of shape (..., num_heads, head_dim)
+
     """
     _apply_rope_input_validation(xq, freqs_cis)
     _apply_rope_input_validation(xk, freqs_cis)
@@ -188,8 +189,7 @@ def forward(
         *,
         pos_embeds: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             x (L, Channels): input tensor
             grid_hw (N, 2): grid height and width
             pos_embeds: precomputed positional embeddings of shape
@@ -199,6 +199,7 @@ def forward(
 
         Returns:
             (L, Cout) tensor
+
         """
         x = self.proj(x).view(x.size(0), -1)
         if pos_embeds is not None:
@@ -226,6 +227,7 @@ class Rope2DPosEmb(nn.Module):
         max_width (int): the maximum width of the 2D grid
         theta_base (float): the base of the theta
         device (str): the device to store the precomputed cis
+
     """
 
     def __init__(
@@ -306,24 +308,26 @@ def get_freqs_cis_by_seqlens_list(
         return freqs_cis
 
     def get_freqs_cis_by_seqlens(self, grid_hws: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             grid_hws (torch.Tensor): containing list of (height, width) or (t, height, width) tuples.
+
         Returns:
             freqs_cis: tensor of shape (sum(t * height * width), dim//2)
+
         """
         return self.get_freqs_cis_by_seqlens_list(grid_hws.tolist())
 
     def get_freqs_cis_by_idx(
         self, pos_idx: torch.Tensor, pos_idx_mask: torch.Tensor
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             pos_idx: tensor of shape (..., 2), It contains the (h, w) position indices of each 2D token.
             pos_idx_mask: a mask of shape (...), the leading dimensions should be the same as pos_idx.
                 Rope will only be applied to the tokens with True mask. `freqs_cis` for the tokens with False mask with be ones.
+
         Return:
             freqs_cis: tensor of shape (..., dim//2)
+
         """
         assert (
             pos_idx.shape[:-1] == pos_idx_mask.shape
@@ -343,10 +347,10 @@ def get_freqs_cis_by_idx(
 
 
 class MLP2(nn.Module):
-    """
-    Args:
-        dims: [in_dim, hidden_dim, out_dim]
-        bias: whether to use bias in linear layer.
+    """Args:
+    dims: [in_dim, hidden_dim, out_dim]
+    bias: whether to use bias in linear layer.
+
     """
 
     def __init__(
@@ -441,13 +445,13 @@ def attention_qkvpacked(
         rope_freqs_cis: torch.Tensor | None = None,
         max_seqlen: torch.Tensor | None = None,
     ):
-        """
-        Args:
-            x (torch.Tensor): (seqlen, hidden_dim)
-            cu_seqlens (torch.Tensor):
-            max_seqlen: Optional precomputed scalar tensor. When omitted it
-                is derived from ``cu_seqlens``, which produces a GPU scalar
-                that breaks CUDA graph capture.
+        """Args:
+        x (torch.Tensor): (seqlen, hidden_dim)
+        cu_seqlens (torch.Tensor):
+        max_seqlen: Optional precomputed scalar tensor. When omitted it
+            is derived from ``cu_seqlens``, which produces a GPU scalar
+            that breaks CUDA graph capture.
+
         """
         seq_length = x.size(0)
         xqkv, _ = self.wqkv(x)
@@ -487,13 +491,13 @@ def forward(
         rope_freqs_cis: torch.Tensor | None = None,
         max_seqlen: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             hidden_states: non-packed (B, N, D) or packed (L, D). if non-packed, seqlens should be None, if packed, seqlens should be set
             max_seqlen: optional precomputed max-sequence-length scalar.
 
         Returns:
             output: same shape of input, non-packed (B, N, D) for non-packed input, (L, D) for packed input
+
         """
         residual = hidden_states
         hidden_states = self.norm0(hidden_states)
@@ -724,6 +728,7 @@ def prepare_encoder_metadata(
                 graph.
             device: Device for the metadata tensors. Defaults to the
                 model's parameter device.
+
         """
         if device is None:
             device = next(self.parameters()).device
@@ -785,18 +790,18 @@ def forward(
         *,
         encoder_metadata: dict[str, Any] | None = None,
     ) -> torch.Tensor | list[torch.Tensor]:
-        """
-        Args:
-            pixel_values (torch.Tensor): The input pixel values.
-            grid_hw (torch.Tensor): The grid height and width.
-            encoder_metadata: Optional precomputed metadata produced by
-                :meth:`prepare_encoder_metadata`. When provided every
-                ``.tolist()`` call in the forward path is skipped, the
-                returned tensor is the packed
-                ``(sum(new_h*new_w), kh*kw, hidden_size)`` form (suitable
-                for CUDA graph capture/replay), and ``grid_hw`` is unused.
-                When ``None`` the legacy path runs and returns a list of
-                per-image tensors.
+        """Args:
+        pixel_values (torch.Tensor): The input pixel values.
+        grid_hw (torch.Tensor): The grid height and width.
+        encoder_metadata: Optional precomputed metadata produced by
+            :meth:`prepare_encoder_metadata`. When provided every
+            ``.tolist()`` call in the forward path is skipped, the
+            returned tensor is the packed
+            ``(sum(new_h*new_w), kh*kw, hidden_size)`` form (suitable
+            for CUDA graph capture/replay), and ``grid_hw`` is unused.
+            When ``None`` the legacy path runs and returns a list of
+            per-image tensors.
+
         """
         if encoder_metadata is not None:
             hidden_states = self.patch_embed(
diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py
index d0bd8510a894..6e7f50415082 100644
--- a/vllm/model_executor/models/moss_audio.py
+++ b/vllm/model_executor/models/moss_audio.py
@@ -13,13 +13,13 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from transformers import BatchFeature, PretrainedConfig, Qwen3Config
+from transformers import BatchFeature, PreTrainedConfig, Qwen3Config
 from transformers.models.whisper import WhisperFeatureExtractor
 from typing_extensions import TypedDict
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_world_size,
@@ -147,11 +147,10 @@ def _ensure_moss_audio_placeholder_tokens(tokenizer: object) -> None:
 
 
 class MossAudioAudioInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - nmb: Number of mel bins
-        - t: Time frames
+    """Dimensions:
+    - b: Batch size
+    - nmb: Number of mel bins
+    - t: Time frames
     """
 
     audio_data: Annotated[torch.Tensor, TensorShape("b", "nmb", "t")]
@@ -264,7 +263,7 @@ def from_config(cls, config: object) -> "MossAudioEncoderConfig":
         return cls(**values)
 
 
-class MossAudioConfig(PretrainedConfig):
+class MossAudioConfig(PreTrainedConfig):
     model_type = "moss_audio"
     is_composition = True
 
@@ -1091,6 +1090,7 @@ def __call__(
             -> audio_data has shape [1, mel_dim, max_time]
             -> mel_dim is the number of mel filter-bank bins, 128 by default
             -> audio_data_seqlens stores the unpadded mel length
+
         """
         del kwargs
 
@@ -1341,47 +1341,29 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
         return {
             "audio": self._get_dummy_audios(
                 length=16000,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
 
 class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingInfo]):
-    def _apply_hf_processor_main(
+    def _call_hf_processor(
         self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
-
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", [])
-        if audios:
-            mm_data["audio"] = audios
-        hf_processor_mm_kwargs = dict(hf_processor_mm_kwargs)
-        processor_kwargs = _filter_moss_audio_processor_config(hf_processor_mm_kwargs)
-        processed_data = self.info.ctx.call_hf_processor(
+        processor_kwargs = _filter_moss_audio_processor_config(dict(hf_kwargs))
+        return self.info.ctx.call_hf_processor(
             self.info.get_hf_processor(**processor_kwargs),
-            mm_data,
+            hf_data,
             {},
         )
-        processed_data.update(passthrough_data)
-        return processed_data
 
     def _get_mm_fields_config(
         self,
@@ -1723,6 +1705,7 @@ def _process_audio_input(
             audio_data=[2, 128, 1200], audio_data_seqlens=[800, 1200]
             -> returns (audio0_embeds, audio1_embeds), split by token length
             -> DeepStack packs each item as [main, layer0, ...] on dim -1
+
         """
         audio_data = audio_input["audio_data"]
         audio_data_seqlens = audio_input["audio_data_seqlens"]
diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py
index 4dc3bdb3d743..e3fdab76ff04 100644
--- a/vllm/model_executor/models/moss_transcribe_diarize.py
+++ b/vllm/model_executor/models/moss_transcribe_diarize.py
@@ -20,7 +20,7 @@
 from transformers import BatchFeature
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt
 from vllm.model_executor.models.interfaces import (
@@ -66,7 +66,7 @@
     PromptUpdateDetails,
     cached_encode,
 )
-from vllm.multimodal.processing.processor import ProcessorInputs
+from vllm.multimodal.processing.processor import HFMultiModalInputs, ProcessorInputs
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.processor import cached_processor_from_config
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -89,12 +89,11 @@
 
 
 class MossTranscribeDiarizeAudioInputs(TensorSchema):
-    """
-    Dimensions:
-        - c: Audio chunks
-        - m: Mel bins
-        - f: Mel frames
-        - n: Number of audio items
+    """Dimensions:
+    - c: Audio chunks
+    - m: Mel bins
+    - f: Mel frames
+    - n: Number of audio items
     """
 
     type: Literal["audio_features"] = "audio_features"
@@ -114,11 +113,10 @@ class MossTranscribeDiarizeAudioInputs(TensorSchema):
 
 
 class MossTranscribeDiarizeEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - n: Number of audio items
-        - t: Number of audio tokens
-        - h: Hidden size
+    """Dimensions:
+    - n: Number of audio items
+    - t: Number of audio tokens
+    - h: Hidden size
     """
 
     type: Literal["audio_embeds"] = "audio_embeds"
@@ -425,59 +423,66 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         num_audios = mm_counts.get("audio", 0)
         if num_audios == 0:
             return {}
 
         feature_extractor = self.info.get_feature_extractor()
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
         return {
             "audio": self._get_dummy_audios(
                 length=_get_max_audio_samples(feature_extractor),
                 num_audios=num_audios,
-                overrides=audio_overrides,
+                overrides=mm_options.get("audio"),
             )
         }
 
-    def get_dummy_processor_inputs(
+
+class MossTranscribeDiarizeMultiModalProcessor(
+    BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo]
+):
+    def get_dummy_inputs(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> ProcessorInputs:
-        dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
+        builder = self.dummy_inputs
+        tokenizer = self.info.get_tokenizer()
+
+        dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options)
         dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
+
         num_audios = mm_counts.get("audio", 0)
-        tokenizer = self.info.get_tokenizer()
-        prompt = cached_encode(
-            tokenizer,
-            AUDIO_PLACEHOLDER * num_audios,
-            add_special_tokens=False,
-        ) or cached_encode(
-            tokenizer,
-            "\n",
-            add_special_tokens=False,
-        )
-        return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items)
+        if num_audios > 0:
+            prompt = cached_encode(
+                tokenizer,
+                AUDIO_PLACEHOLDER * num_audios,
+                add_special_tokens=False,
+            )
+        else:
+            prompt = cached_encode(
+                tokenizer,
+                "\n",
+                add_special_tokens=False,
+            )
 
+        return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items)
 
-class MossTranscribeDiarizeMultiModalProcessor(
-    BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo]
-):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        audios = _get_audios_from_mm_data(mm_data)
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+        if hf_inputs.hf_data:
+            hf_inputs.hf_data["audio"] = _get_audios_from_mm_data(hf_inputs.hf_data)
 
-        return dict(audio=audios), hf_processor_mm_kwargs
+        return hf_inputs
 
     def _postprocess_hf_mm_data(
         self,
diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py
index ade38cc703da..a7938bb56b84 100644
--- a/vllm/model_executor/models/muse_glimmer.py
+++ b/vllm/model_executor/models/muse_glimmer.py
@@ -45,11 +45,7 @@
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
-from vllm.config.multimodal import (
-    BaseDummyOptions,
-    ImageDummyOptions,
-    VideoDummyOptions,
-)
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import (
     divide,
     get_pp_group,
@@ -234,7 +230,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         processor = self.info.get_hf_processor()
         video_processor = processor.video_processor
@@ -245,16 +241,12 @@ def get_dummy_mm_data(
             * int(video_processor.downsample_factor)
             * video_grid
         )
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
-        assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)
         return {
             "image": self._get_dummy_images(
                 width=image_width,
                 height=image_height,
                 num_images=mm_counts.get("image", 0),
-                overrides=image_overrides,
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=video_size,
@@ -263,7 +255,7 @@ def get_dummy_mm_data(
                     seq_len, mm_counts
                 ),
                 num_videos=mm_counts.get("video", 0),
-                overrides=video_overrides,
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -274,16 +266,18 @@ class MuseGlimmerMultiModalProcessor(
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+
+        if not mm_data:
+            return self._finalize_hf_mm_data(mm_data, hf_kwargs, passthrough_data)
 
         prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
 
-        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        processor = self.info.get_hf_processor(**hf_kwargs)
         tokenizer = processor.tokenizer
         config = self.info.get_hf_config()
         images = mm_data.get("images", ())
@@ -370,8 +364,9 @@ def _apply_hf_processor_main(
                 video_feature_sizes=torch.tensor(video_sizes),
             )
         processed_data = BatchFeature(data=data, tensor_type=None)
-        processed_data.update(passthrough_data)
-        return processed_data
+        return self._finalize_hf_mm_data(
+            mm_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_mm_fields_config(
         self,
@@ -1663,9 +1658,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         return MultiModelKeys.from_string_field(
             language_model="model",
             connector=["vision_adapter.", "vision_projection."],
diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py
index ac6e254cb2df..0f4dc8b11052 100644
--- a/vllm/model_executor/models/nano_nemotron_vl.py
+++ b/vllm/model_executor/models/nano_nemotron_vl.py
@@ -12,17 +12,21 @@
 from collections.abc import Iterable, Mapping, Sequence
 from functools import cached_property
 from io import BytesIO
-from typing import Annotated, Literal, TypeAlias
+from typing import Annotated, Literal, TypeAlias, TypedDict, cast
 
 import torch
 import torch.nn as nn
-from transformers import BatchFeature, PretrainedConfig
+from transformers import BatchFeature, PreTrainedConfig
 
 from vllm import envs
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
+from vllm.config.multimodal import (
+    MultiModalDummyOptions,
+    VideoDummyOptions,
+)
 from vllm.inputs import MultiModalDataDict, MultiModalInput
 from vllm.logger import init_logger
+from vllm.lora.layers.base import BaseLayerWithLoRA
 from vllm.model_executor.layers.activation import ReLUSquaredActivation
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -30,6 +34,7 @@
     HasInnerState,
     IsHybrid,
     MultiModalEmbeddings,
+    SupportsLoRA,
     SupportsMultiModal,
     SupportsMultiModalPruning,
 )
@@ -68,6 +73,7 @@
 from vllm.multimodal.processing.processor import (
     BaseMultiModalProcessor,
     BaseProcessingInfo,
+    HFMultiModalInputs,
     PromptReplacement,
     PromptUpdate,
     cached_encode,
@@ -101,28 +107,26 @@
 
 
 class NanoNemotronVLAudioFeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - c: Number of audio clips (possibly flattened across audio items)
-        - b: Number of original audio items
-        - t: Audio feature length
-        - f: Feature size (mel bins)
+    """Dimensions:
+    - c: Number of audio clips (possibly flattened across audio items)
+    - b: Number of original audio items
+    - t: Audio feature length
+    - f: Feature size (mel bins)
     """
 
     type: Literal["audio_features"] = "audio_features"
     input_audio_features: Annotated[torch.Tensor, TensorShape("c", "t", "f")]
     feature_attention_mask: Annotated[torch.Tensor, TensorShape("c", "t")]
-    audio_num_clips: list[int]
+    audio_num_clips: list[int] | torch.Tensor
 
 
 class NanoNemotronVLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - bnp: Batch size * number of images * (1 + num_patches)
-        - c: Number of channels (3)
-        - h: Height of each image patch
-        - w: Width of each image patch
+    """Dimensions:
+    - bn: Batch size * number of images
+    - bnp: Batch size * number of images * (1 + num_patches)
+    - c: Number of channels (3)
+    - h: Height of each image patch
+    - w: Width of each image patch
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -131,8 +135,7 @@ class NanoNemotronVLImagePixelInputs(TensorSchema):
 
 
 class NanoNemotronVLImagePixelInputsDynamic(TensorSchema):
-    """
-    Dynamic-resolution image inputs.
+    """Dynamic-resolution image inputs.
 
     imgs_sizes: per-image (height, width) in pixels.
     num_tokens_per_image: per-image number of embedding tokens (post downsample).
@@ -145,11 +148,10 @@ class NanoNemotronVLImagePixelInputsDynamic(TensorSchema):
 
 
 class NanoNemotronVLImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - n: Number of images
-        - f: Total image feature size
-        - h: Hidden size (must match the hidden size of language model backbone)
+    """Dimensions:
+    - n: Number of images
+    - f: Total image feature size
+    - h: Hidden size (must match the hidden size of language model backbone)
     """
 
     type: Literal["image_embeds"]
@@ -164,14 +166,13 @@ class NanoNemotronVLImageEmbeddingInputs(TensorSchema):
 
 
 class NanoNemotronVLVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bvf: Batch size * number of videos * num_frames
-        - bn: Batch size * number of videos
-        - f: Number of frames
-        - c: Number of channels (3)
-        - h: Height of each video frame
-        - w: Width of each video frame
+    """Dimensions:
+    - bvf: Batch size * number of videos * num_frames
+    - bn: Batch size * number of videos
+    - f: Number of frames
+    - c: Number of channels (3)
+    - h: Height of each video frame
+    - w: Width of each video frame
     """
 
     type: Literal["pixel_values_videos"]
@@ -182,11 +183,10 @@ class NanoNemotronVLVideoPixelInputs(TensorSchema):
 
 
 class NanoNemotronVLVideoEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - n: Number of videos
-        - f: Total video feature size
-        - h: Hidden size (must match the hidden size of language model backbone)
+    """Dimensions:
+    - n: Number of videos
+    - f: Total video feature size
+    - h: Hidden size (must match the hidden size of language model backbone)
     """
 
     type: Literal["video_embeds"]
@@ -198,8 +198,17 @@ class NanoNemotronVLVideoEmbeddingInputs(TensorSchema):
 )
 
 
+class NanoNemotronVLMultiModalInputs(TypedDict, total=False):
+    images: NanoNemotronVLImageInputs | None
+    videos: NanoNemotronVLVideoInputs | None
+    audios: NanoNemotronVLAudioFeatureInputs | None
+
+
 class NanoNemotronVLProcessingInfo(BaseProcessingInfo):
     def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor:
+        # Not accepted by NanoNemotronVLProcessor.__init__
+        kwargs.pop("truncation", None)
+
         return self.ctx.init_processor(
             NanoNemotronVLProcessor,
             config=self.get_hf_config(),
@@ -229,7 +238,7 @@ def get_video_pruning_rate(self) -> float | None:
         return self.ctx.get_mm_config().video_pruning_rate
 
     @property
-    def sound_config(self) -> PretrainedConfig | None:
+    def sound_config(self) -> PreTrainedConfig | None:
         return getattr(self.get_hf_config(), "sound_config", None)
 
     def get_default_tok_params(self) -> TokenizeParams:
@@ -362,9 +371,23 @@ def get_mm_max_tokens_per_item(
 class NanoNemotronVLMultiModalProcessor(
     BaseMultiModalProcessor[NanoNemotronVLProcessingInfo]
 ):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
+    def _get_hf_mm_inputs(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        # NanoNemotronVLProcessor expects "audios" instead of "audio"
+        hf_data = hf_inputs.hf_data
+        if "audio" in hf_data:
+            hf_data["audios"] = hf_data.pop("audio")
+
+        return hf_inputs
+
     def _get_image_fields_config(self, hf_inputs: BatchFeature):
         if self.info.is_dynamic_tiler:
             pixel_values_flat = MultiModalFieldConfig.batched("image")
@@ -451,8 +474,11 @@ def get_image_replacement(item_idx: int):
             if isinstance(images, ImageEmbeddingItems):
                 feature_size = images.get_feature_size(item_idx)
             elif self.info.is_dynamic_tiler:
-                feature_size = out_mm_data["num_tokens_per_image"][item_idx]
+                num_tokens_per_image = out_mm_data["num_tokens_per_image"]
+                assert isinstance(num_tokens_per_image, list)
+                feature_size = num_tokens_per_image[item_idx]
             else:
+                assert isinstance(images, ImageProcessorItems)
                 image_size = images.get_image_size(item_idx)
                 max_num_tiles = hf_processor.max_num_tiles
                 feature_size = hf_processor.get_num_image_tokens(
@@ -516,10 +542,12 @@ def get_video_replacement(item_idx: int):
                 assert isinstance(num_patches, int)
 
             T = hf_processor.video_temporal_patch_size
+            num_tubelets: int | None
             if T > 1 and num_patches is not None:
                 num_tubelets = math.ceil(num_patches / T)
             else:
                 num_tubelets = num_patches
+            assert num_tubelets is not None
 
             video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate
             if video_pruning_rate is not None and video_pruning_rate > 0.0:
@@ -611,6 +639,7 @@ def _extract_audio_from_videos(
         Returns:
             A 3-tuple of (augmented mm_items, extracted audio items,
             per-video boolean mask indicating which videos have audio).
+
         """
         videos = mm_items.get_items("video", VideoProcessorItems)
         assert isinstance(videos.metadata, list)
@@ -620,6 +649,7 @@ def _extract_audio_from_videos(
         audio_items: list[AudioItem] = []
         has_audio: list[bool] = []
         for idx, metadata in enumerate(metadata_list):
+            assert metadata is not None
             video_bytes = metadata.get("original_video_bytes")
             if video_bytes is None or len(video_bytes) == 0:
                 raise ValueError(
@@ -646,21 +676,25 @@ def _extract_audio_from_videos(
 
         # Create a new VideoProcessorItems with metadata that does not contain
         # the large video bytes, to avoid modifying the input `mm_items`.
-        new_metadata_list = [
-            {k: v for k, v in meta.items() if k != "original_video_bytes"}
-            for meta in metadata_list
-        ]
+        new_metadata_list = metadata_list.copy()
+        for idx, metadata in enumerate(metadata_list):
+            assert metadata is not None
+            new_metadata_list[idx] = {
+                key: value
+                for key, value in metadata.items()
+                if key != "original_video_bytes"
+            }
         new_videos = VideoProcessorItems(data=videos.data, metadata=new_metadata_list)
 
-        audio_parsed = {}
+        audio_parsed = MultiModalDataItems()
         if audio_items:
             audio_parsed = self.data_parser.parse_mm_data({"audio": audio_items})
 
         # Create a new MultiModalDataItems with the new video and audio items.
         new_mm_items_dict = {**mm_items, **audio_parsed, "video": new_videos}
-        mm_items = MultiModalDataItems(new_mm_items_dict)
+        new_mm_items = MultiModalDataItems(new_mm_items_dict)
 
-        return mm_items, audio_items, has_audio
+        return new_mm_items, audio_items, has_audio
 
     def apply(
         self,
@@ -733,15 +767,12 @@ def apply(
         # Bypass the cached path: the HF processor must receive the
         # prompt (with injected <so_embedding>) and the audio data
         # together so it can perform audio-token replacement natively.
-        mm_info = self._apply_hf_processor(inputs, timing_ctx)
-        prompt_ids = self._postprocess_prompt(inputs.prompt)
+        mm_res = self._apply_hf_processor(inputs, timing_ctx)
 
         with timing_ctx.record("apply_prompt_updates"):
             prompt_ids, mm_placeholders = self._maybe_apply_prompt_updates(
                 mm_items=mm_items,
-                prompt_ids=prompt_ids,
-                mm_kwargs=mm_info.kwargs,
-                mm_prompt_updates=mm_info.prompt_updates,
+                mm_res=mm_res,
             )
 
         mm_placeholder_ranges = {
@@ -752,8 +783,8 @@ def apply(
         return MultiModalInput(
             type="multimodal",
             prompt_token_ids=prompt_ids,
-            mm_kwargs=mm_info.kwargs,
-            mm_hashes=mm_info.hashes,
+            mm_kwargs=mm_res.kwargs,
+            mm_hashes=mm_res.hashes,
             mm_placeholders=mm_placeholder_ranges,
         )
 
@@ -788,7 +819,7 @@ def _get_dummy_videos(
         )
         videos = [v.copy() for v in videos]
 
-        video_items = []
+        video_items: list[VideoItem] = []
         for video in videos:
             video_num_frames = video.shape[0]
             video_metadata = {
@@ -807,22 +838,19 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         (target_width, target_height), _ = (
             self.info.get_dummy_image_size_and_max_tokens(mm_counts)
         )
         processor = self.info.get_hf_processor()
 
-        image_overrides = mm_options.get("image")
-
         dummy_image = {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -897,11 +925,22 @@ def get_dummy_mm_data(
     dummy_inputs=NanoNemotronVLDummyInputsBuilder,
 )
 class NemotronH_Nano_VL_V2(
-    nn.Module, HasInnerState, IsHybrid, SupportsMultiModal, SupportsMultiModalPruning
+    nn.Module,
+    HasInnerState,
+    IsHybrid,
+    SupportsMultiModal,
+    SupportsMultiModalPruning,
+    SupportsLoRA,
 ):
     requires_sequential_video_encoding = True
     """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet"""
 
+    # LoRA covers the language model only
+    is_non_gated_moe = NemotronHForCausalLM.is_non_gated_moe
+    packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping
+    embedding_modules = NemotronHForCausalLM.embedding_modules
+    lora_skip_prefixes = NemotronHForCausalLM.lora_skip_prefixes
+
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "language_model.backbone": "language_model.model",
@@ -923,6 +962,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         model_config = vllm_config.model_config
         config = model_config.hf_config
         multimodal_config = model_config.multimodal_config
+        assert multimodal_config is not None
         image_size = config.force_image_size
         patch_size = config.patch_size
         self.patch_size = patch_size
@@ -1103,7 +1143,8 @@ def extract_feature(
     def _parse_and_validate_image_input(
         self, **kwargs: object
     ) -> NanoNemotronVLImageInputs | None:
-        if image_embeds := kwargs.pop("image_embeds", None):
+        image_embeds = kwargs.pop("image_embeds", None)
+        if image_embeds is not None:
             return NanoNemotronVLImageEmbeddingInputs(
                 type="image_embeds",
                 data=image_embeds,
@@ -1114,17 +1155,26 @@ def _parse_and_validate_image_input(
             return None
 
         if self.dynamic_resolution:
+            assert isinstance(pixel_values_flat, (list, torch.Tensor))
             pixel_values_flat = DynamicResolutionImageTiler.stack(
                 pixel_values_flat, self.patch_size
             )
+            imgs_sizes = kwargs.pop("imgs_sizes")
+            num_tokens_per_image = kwargs.pop("num_tokens_per_image")
+            assert isinstance(imgs_sizes, list)
+            assert isinstance(num_tokens_per_image, list)
             return NanoNemotronVLImagePixelInputsDynamic(
-                pixel_values_flat=pixel_values_flat, **kwargs
+                pixel_values_flat=pixel_values_flat,
+                imgs_sizes=imgs_sizes,
+                num_tokens_per_image=num_tokens_per_image,
             )
         else:
+            assert isinstance(pixel_values_flat, torch.Tensor)
+            image_num_patches = kwargs.pop("image_num_patches")
+            assert isinstance(image_num_patches, torch.Tensor)
             return NanoNemotronVLImagePixelInputs(
                 pixel_values_flat=pixel_values_flat,
-                num_patches=kwargs.pop("image_num_patches"),
-                **kwargs,
+                num_patches=image_num_patches,
             )
 
     def _process_image_input_dynamic(
@@ -1142,7 +1192,8 @@ def _process_image_input_dynamic(
         return image_embeds.split(num_tokens_per_image)
 
     def _process_image_input(
-        self, image_input: NanoNemotronVLImagePixelInputs
+        self,
+        image_input: NanoNemotronVLImagePixelInputs | NanoNemotronVLVideoPixelInputs,
     ) -> tuple[torch.Tensor, ...]:
         image_embeds = self.extract_feature(image_input["pixel_values_flat"])
         num_patches = image_input["num_patches"]
@@ -1330,7 +1381,20 @@ def _create_final_video_embeddings(
 
         # Create final video embeddings, merging text embeddings for indicator
         # tokens with video embeddings
-        text_embeddings = self.get_language_model().embed_input_ids(repl_token_ids)
+
+        # LoRA support -
+        # These replacement tokens are produced inside the encoder, so they are
+        # absent from the request-token batch that the LoRA adapter index
+        # mapping is built from, and there can be more
+        # of them than max_num_batched_tokens.
+        # Embed them with the base weights -
+        # the LoRA delta is undefined for tokens with no mapping entry.
+        language_model = cast(NemotronHForCausalLM, self.get_language_model())
+        embed_tokens = language_model.model.embed_tokens
+        if isinstance(embed_tokens, BaseLayerWithLoRA):
+            embed_tokens = embed_tokens.base_layer
+        text_embeddings = embed_tokens(repl_token_ids)
+
         final_video_embeddings = _merge_multimodal_embeddings(
             inputs_embeds=text_embeddings,
             multimodal_embeddings=video_embeddings,
@@ -1341,7 +1405,7 @@ def _create_final_video_embeddings(
 
     def _parse_and_validate_video_input(
         self, **kwargs: object
-    ) -> NanoNemotronVLVideoPixelInputs | None:
+    ) -> NanoNemotronVLVideoInputs | None:
         pixel_values_flat_video = kwargs.pop("pixel_values_flat_video", None)
         video_num_patches = kwargs.pop("video_num_patches", None)
         video_embeds = kwargs.pop("video_embeds", None)
@@ -1352,26 +1416,29 @@ def _parse_and_validate_video_input(
             return None
 
         if video_embeds is not None:
+            assert isinstance(video_embeds, (torch.Tensor, list))
             return NanoNemotronVLVideoEmbeddingInputs(
                 type="video_embeds",
                 data=video_embeds,
             )
 
         if pixel_values_flat_video is not None:
-            if torch.is_tensor(frames_indices):
+            if isinstance(frames_indices, torch.Tensor):
                 frames_indices = frames_indices.flatten()
             else:
+                assert isinstance(frames_indices, (list, tuple))
                 frames_indices = torch.cat([f.flatten() for f in frames_indices], dim=0)
 
-            if torch.is_tensor(frame_duration_ms):
+            if isinstance(frame_duration_ms, torch.Tensor):
                 frame_duration_ms = frame_duration_ms.flatten()
             else:
+                assert isinstance(frame_duration_ms, (list, tuple))
                 frame_duration_ms = torch.cat(
                     [f.flatten() for f in frame_duration_ms], dim=0
                 )
 
             if (
-                torch.is_tensor(pixel_values_flat_video)
+                isinstance(pixel_values_flat_video, torch.Tensor)
                 and pixel_values_flat_video.ndim == 5
             ):
                 # batched._reduce_data stacked same-shape videos into
@@ -1379,9 +1446,11 @@ def _parse_and_validate_video_input(
                 # same-H,W cat path below handles it uniformly.
                 pixel_values_flat_video = list(pixel_values_flat_video)
 
-            if not torch.is_tensor(pixel_values_flat_video):
+            if not isinstance(pixel_values_flat_video, torch.Tensor):
+                assert isinstance(pixel_values_flat_video, (list, tuple))
                 pixel_values_flat_video = torch.cat(pixel_values_flat_video, dim=0)
 
+            assert isinstance(video_num_patches, torch.Tensor)
             expected_h = pixel_values_flat_video.shape[-2]
             expected_w = pixel_values_flat_video.shape[-1]
             num_frames = video_num_patches[0].item()
@@ -1398,8 +1467,10 @@ def _parse_and_validate_video_input(
 
         raise AssertionError("This line should be unreachable.")
 
-    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
-        modalities = {}
+    def _parse_and_validate_multimodal_inputs(
+        self, **kwargs: object
+    ) -> NanoNemotronVLMultiModalInputs:
+        modalities: NanoNemotronVLMultiModalInputs = {}
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
@@ -1408,7 +1479,10 @@ def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
                 and "images" not in modalities
             ):
                 modalities["images"] = self._parse_and_validate_image_input(**kwargs)
-            if input_key in ("pixel_values_flat_video",) and "videos" not in modalities:
+            if (
+                input_key in ("pixel_values_flat_video", "video_embeds")
+                and "videos" not in modalities
+            ):
                 modalities["videos"] = self._parse_and_validate_video_input(**kwargs)
             if (
                 input_key
@@ -1419,8 +1493,17 @@ def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
                 )
                 and "audios" not in modalities
             ):
+                input_audio_features = kwargs.get("input_audio_features")
+                feature_attention_mask = kwargs.get("feature_attention_mask")
+                audio_num_clips = kwargs.get("audio_num_clips")
+                assert isinstance(input_audio_features, torch.Tensor)
+                assert isinstance(feature_attention_mask, torch.Tensor)
+                assert isinstance(audio_num_clips, (list, torch.Tensor))
                 modalities["audios"] = NanoNemotronVLAudioFeatureInputs(
-                    **kwargs, validate=False
+                    input_audio_features=input_audio_features,
+                    feature_attention_mask=feature_attention_mask,
+                    audio_num_clips=audio_num_clips,
+                    validate=False,
                 )
 
         return modalities
@@ -1440,20 +1523,26 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
         for modality in modalities:
             if modality == "images":
                 image_input = modalities["images"]
-                if image_input["type"] == "image_embeds":
+                assert image_input is not None
+                if isinstance(image_input, NanoNemotronVLImageEmbeddingInputs):
                     image_embeddings = image_input["data"]
-                elif self.dynamic_resolution:
-                    assert image_input["type"] == "pixel_values_dynamic"
+                elif isinstance(image_input, NanoNemotronVLImagePixelInputsDynamic):
                     image_embeddings = self._process_image_input_dynamic(image_input)
                 else:
                     image_embeddings = self._process_image_input(image_input)
                 multimodal_embeddings += tuple(image_embeddings)
             if modality == "videos":
                 video_input = modalities["videos"]
-                video_embeddings = self._process_video_input(video_input)
+                assert video_input is not None
+                if video_input["type"] == "video_embeds":
+                    video_embeddings = video_input["data"]
+                else:
+                    assert isinstance(video_input, NanoNemotronVLVideoPixelInputs)
+                    video_embeddings = self._process_video_input(video_input)
                 multimodal_embeddings += tuple(video_embeddings)
             if modality == "audios":
                 audio_input = modalities["audios"]
+                assert audio_input is not None
                 audio_embeddings = self._process_audio_input(audio_input)
                 multimodal_embeddings += tuple(audio_embeddings)
 
@@ -1481,9 +1570,7 @@ def forward(
         return hidden_states
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector=["mlp1", "sound_encoder.projection"],
@@ -1498,6 +1585,7 @@ def compute_logits(
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         mm_config = self.model_config.multimodal_config
+        assert mm_config is not None
         load_multimodal_weights = not all(
             mm_config.get_limit_per_prompt(modality) == 0
             for modality in ("image", "video", "audio")
diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py
index e65e51f3bb9d..e5451d16ba02 100644
--- a/vllm/model_executor/models/nemotron_h.py
+++ b/vllm/model_executor/models/nemotron_h.py
@@ -25,7 +25,11 @@
 from torch import nn
 
 from vllm.compilation.decorators import support_torch_compile
-from vllm.config import CacheConfig, ModelConfig, VllmConfig
+from vllm.config import (
+    CacheConfig,
+    ModelConfig,
+    VllmConfig,
+)
 from vllm.config.parallel import ParallelConfig
 from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size
 from vllm.distributed.communication_op import tensor_model_parallel_all_gather
@@ -38,12 +42,14 @@
     activation_without_mul,
     fused_moe_make_expert_params_mapping,
 )
+from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
     ReplicatedLinear,
     RowParallelLinear,
+    UnquantizedLinearMethod,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
@@ -62,6 +68,7 @@
     EagleModelMixin,
     HasInnerState,
     IsHybrid,
+    MambaStateShapes,
     MixtureOfExperts,
     SupportsEagle,
     SupportsEagle3,
@@ -118,7 +125,7 @@ def __init__(
 
     def forward(self, x: torch.Tensor):
         x, _ = self.up_proj(x)
-        x = self.act_fn(x)
+        x = maybe_fused_act_quant(self.act_fn, x, self.down_proj)
         x, _ = self.down_proj(x)
         return x
 
@@ -132,6 +139,7 @@ def __init__(
         prefix: str = "",
     ):
         super().__init__()
+        assert parallel_config is not None
         self.tp_size = get_tensor_model_parallel_world_size()
         self.routed_scaling_factor = config.routed_scaling_factor
 
@@ -182,6 +190,8 @@ def __init__(
                 prefix=f"{prefix}.shared_experts",
             )
 
+        self.fc1_latent_proj: ReplicatedLinear | None
+        self.fc2_latent_proj: ReplicatedLinear | None
         if self.use_latent_moe:
             self.fc1_latent_proj = ReplicatedLinear(
                 input_size=config.hidden_size,
@@ -199,6 +209,17 @@ def __init__(
                 disable_tp=self.is_sequence_parallel,
                 prefix=f"{prefix}.fc2_latent_proj",
             )
+            # A bias-free, unquantized linear commutes with the TP sum
+            # (sum_r W x_r == W sum_r x_r), so one reduce after the transform
+            # suffices. Test the layer, not the model-wide `quant_config`: ModelOpt
+            # excludes the latent projections, so quantized checkpoints still get
+            # an UnquantizedLinearMethod here.
+            self.fc2_latent_proj.reduce_commutative = (
+                not config.mlp_bias
+                and isinstance(
+                    self.fc2_latent_proj.quant_method, UnquantizedLinearMethod
+                )
+            )
         else:
             self.fc1_latent_proj = None
             self.fc2_latent_proj = None
@@ -705,7 +726,7 @@ class NemotronHForCausalLM(
     SupportsReplaySSM,
 ):
     # Relevant only if self.has_moe is True
-    is_non_gated_moe: bool = True
+    is_non_gated_moe = True
 
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={"backbone": "model", "mtp": None},
@@ -756,7 +777,7 @@ def get_mamba_state_dtype_from_config(
     def get_mamba_state_shape_from_config(
         cls,
         vllm_config: "VllmConfig",
-    ) -> tuple[tuple[int, ...], ...]:
+    ) -> MambaStateShapes:
         """Calculate shapes for Mamba's convolutional and state caches.
 
         Args:
@@ -767,6 +788,7 @@ def get_mamba_state_shape_from_config(
             - conv_state_shape: Shape for convolutional state cache
             - temporal_state_shape: Shape for state space model cache
             - x_cache/dt_cache/B_cache ring-buffer shapes (use_replayssm only)
+
         """
         parallel_config = vllm_config.parallel_config
         cache_config = vllm_config.cache_config
@@ -840,6 +862,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                     self.moe_layers.append(layer.mixer.experts)
 
             self.num_moe_layers = len(self.moe_layers)
+            assert example_moe is not None
             self.num_logical_experts = example_moe.n_logical_experts
             self.num_physical_experts = example_moe.n_physical_experts
             self.num_local_physical_experts = example_moe.n_local_physical_experts  # noqa: E501
diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py
index 96f1b6685ed9..6b9de6300918 100644
--- a/vllm/model_executor/models/nemotron_h_mtp.py
+++ b/vllm/model_executor/models/nemotron_h_mtp.py
@@ -92,10 +92,12 @@ def __init__(
 
     def forward(
         self,
-        inputs_embeds: torch.Tensor,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         residual: torch.Tensor | None = None,
+        *,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         # Start projections (Fusion)
         if self.has_start_projections:
@@ -179,10 +181,11 @@ def __init__(
 
     def forward(
         self,
-        inputs_embeds: torch.Tensor,
-        positions: torch.Tensor,
         hidden_states: torch.Tensor,
         residual: torch.Tensor | None = None,
+        *,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         # Start projections (Fusion)
         if self.has_start_projections:
@@ -305,13 +308,14 @@ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
 
     def forward(
         self,
-        input_ids: torch.Tensor,
+        input_ids: torch.Tensor | None,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
     ) -> torch.Tensor | IntermediateTensors:
         if inputs_embeds is None:
+            assert input_ids is not None
             inputs_embeds = self.get_input_embeddings(input_ids)
 
         residual = None
@@ -386,6 +390,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     @staticmethod
     def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
         vllm_config = kwargs.get("vllm_config")
+        assert isinstance(vllm_config, VllmConfig)
         return get_draft_quant_config(vllm_config)
 
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -393,14 +398,15 @@ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
 
     def forward(
         self,
-        input_ids: torch.Tensor,
+        input_ids: torch.Tensor | None,
         positions: torch.Tensor,
-        hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | None = None,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
         **kwargs: object,
     ) -> torch.Tensor:
         """Forward - applies attention-based MTP."""
+        assert hidden_states is not None
         hidden_states = self.model(
             input_ids,
             positions,
@@ -473,9 +479,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     name = name.replace("backbone.", "model.")
 
             if "scale" in name or "zero_point" in name:
-                name = maybe_remap_kv_scale_name(name, params_dict)
-                if name is None:
+                remapped_name = maybe_remap_kv_scale_name(name, params_dict)
+                if remapped_name is None:
                     continue
+                name = remapped_name
 
             # Handle stacked parameters (qkv_proj) for attention layers
             is_stacked = False
@@ -547,6 +554,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
             param = params_dict[name]
             weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            assert weight_loader is not None
             weight_loader(param, loaded_weight)
             loaded_params.add(name)
 
diff --git a/vllm/model_executor/models/nemotron_parse.py b/vllm/model_executor/models/nemotron_parse.py
index 12c7b580484b..9ac38c0021db 100644
--- a/vllm/model_executor/models/nemotron_parse.py
+++ b/vllm/model_executor/models/nemotron_parse.py
@@ -17,12 +17,12 @@
 from transformers import (
     BartConfig,
     BatchFeature,
-    PretrainedConfig,
+    PreTrainedConfig,
 )
 
 from vllm.config import CacheConfig, VllmConfig
 from vllm.config.lora import LoRAConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
@@ -62,8 +62,7 @@
 
 
 class BartScaledWordEmbedding(VocabParallelEmbedding):
-    """
-    This module overrides VocabParallelEmbedding's
+    """This module overrides VocabParallelEmbedding's
     forward by multiplying with embeddings scale.
     """
 
@@ -138,12 +137,13 @@ def forward(
         decoder_hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        r"""
-        Args:
+        r"""Args:
             decoder_hidden_states: torch.Tensor of *decoder* input embeddings.
             encoder_hidden_states: torch.Tensor of *encoder* input embeddings.
+
         Returns:
             Decoder layer output torch.Tensor
+
         """
         residual = decoder_hidden_states
 
@@ -218,12 +218,13 @@ def forward(
 
 
 class MBartDecoderNoPos(nn.Module):
-    """
-    Transformer decoder consisting of *config.decoder_layers* layers.
+    """Transformer decoder consisting of *config.decoder_layers* layers.
     Each layer is a [`BartDecoderLayer`]
+
     Args:
         config: BartConfig
         embed_tokens (nn.Embedding): output embedding
+
     """
 
     def __init__(
@@ -274,13 +275,13 @@ def forward(
         inputs_embeds: torch.Tensor | None = None,
         **kwargs,
     ) -> torch.Tensor:
-        r"""
-        Args:
+        r"""Args:
             decoder_input_ids: Indices of *decoder* input sequence tokens in the
                 vocabulary. Padding will be ignored by default should you provide it.
             encoder_hidden_states: Tensor of encoder output embeddings
         Returns:
             Decoder output torch.Tensor
+
         """
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(decoder_input_ids)
@@ -339,12 +340,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
 
 class NemotronParsePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
+    """Dimensions:
+    - b: Batch size
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
     """
 
     type: Literal["pixel_values"]
@@ -388,7 +388,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         num_images = mm_counts.get("image", 0)
 
@@ -440,7 +440,7 @@ class RadioWithNeck(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -480,7 +480,7 @@ def __init__(
 
     def get_vit_model_from_radio_config(
         self,
-        hf_config: PretrainedConfig,
+        hf_config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
     ) -> RadioModel:
         hf_config_vision = hf_config.encoder
@@ -630,14 +630,15 @@ def forward(
         encoder_outputs: list[torch.Tensor] | None = None,
         **kwargs,
     ) -> torch.Tensor:
-        r"""
-        Args:
+        r"""Args:
             input_ids: torch.Tensor of *decoder* input token ids.
             positions: torch.Tensor of *decoder* position indices.
             encoder_outputs: List of encoder output tensors (vision embeddings).
                 During profiling, this may be None or empty.
+
         Returns:
             Output torch.Tensor
+
         """
         inputs_embeds = None
         if encoder_outputs:
diff --git a/vllm/model_executor/models/nemotron_vl.py b/vllm/model_executor/models/nemotron_vl.py
index 818d8191971c..0ec4589e5110 100644
--- a/vllm/model_executor/models/nemotron_vl.py
+++ b/vllm/model_executor/models/nemotron_vl.py
@@ -5,7 +5,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import AutoModel, PretrainedConfig
+from transformers import AutoModel, PreTrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.linear import ReplicatedLinear
@@ -138,7 +138,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
                 prefix=maybe_prefix(prefix, "language_model"),
             )
 
-        self.img_context_token_id = None
+        self.img_context_token_id: int | None = None
 
         self.visual_token_mask = None
         self.make_empty_intermediate_tensors = (
@@ -146,7 +146,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         )
 
     def _patch_quant_config(
-        self, config: PretrainedConfig, quant_config: QuantizationConfig
+        self, config: PreTrainedConfig, quant_config: QuantizationConfig | None
     ):
         # the awq models from OpenGVLab missing `modules_to_not_convert`
         # patch the quant_config to add `modules_to_not_convert` back
@@ -160,7 +160,7 @@ def _patch_quant_config(
 
     def _init_vision_model(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
@@ -172,7 +172,7 @@ def _init_vision_model(
 
     def _init_mlp1(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         vit_hidden_size: int | None = None,
         vision_projection_hidden_size: int | None = None,
     ) -> nn.Module:
@@ -387,9 +387,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="mlp1",
@@ -467,8 +465,7 @@ def get_hf_processor(self, **kwargs: object) -> LlamaNemotronVLEmbedProcessor:
     dummy_inputs=BaseInternVLDummyInputsBuilder[LlamaNemotronVLEmbedProcessingInfo],
 )
 class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling):
-    """
-    LlamaNemotronVL model for embeddings.
+    """LlamaNemotronVL model for embeddings.
 
     Inherits from LlamaNemotronVLChatModel and specializes it for embedding tasks:
     - Uses SigLIP vision encoder instead of C-RADIO
@@ -508,7 +505,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
 
     def _init_vision_model(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config,
         *,
         prefix: str,
@@ -521,12 +518,21 @@ def _init_vision_model(
             use_head=False,
         )
 
-    def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
+    def _init_mlp1(
+        self,
+        config: PreTrainedConfig,
+        vit_hidden_size: int | None = None,
+        vision_projection_hidden_size: int | None = None,
+    ) -> nn.Module:
         """Override to use different MLP structure for embedding model."""
+        if vit_hidden_size is None:
+            vit_hidden_size = config.vision_config.hidden_size
+        if vision_projection_hidden_size is None:
+            vision_projection_hidden_size = config.get_text_config().hidden_size
         return super()._init_mlp1(
             config,
-            vit_hidden_size=config.vision_config.hidden_size,
-            vision_projection_hidden_size=config.get_text_config().hidden_size,
+            vit_hidden_size=vit_hidden_size,
+            vision_projection_hidden_size=vision_projection_hidden_size,
         )
 
     def _call_vision_model(self, pixel_values: torch.Tensor) -> torch.Tensor:
diff --git a/vllm/model_executor/models/nvlm_d.py b/vllm/model_executor/models/nvlm_d.py
index d937fade1d95..1039fea95ea0 100644
--- a/vllm/model_executor/models/nvlm_d.py
+++ b/vllm/model_executor/models/nvlm_d.py
@@ -11,9 +11,9 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -83,19 +83,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         target_width, target_height = self.info.get_image_size_with_most_features()
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
 
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -128,6 +125,7 @@ def get_replacement_nvlm(item_idx: int):
             if isinstance(images, ImageEmbeddingItems):
                 feature_size = images.get_feature_size(item_idx)
             else:
+                assert isinstance(images, ImageProcessorItems)
                 image_size = images.get_image_size(item_idx)
                 feature_size = self.info.get_num_image_tokens(
                     image_width=image_size.width,
@@ -163,7 +161,7 @@ def get_replacement_nvlm(item_idx: int):
     dummy_inputs=NVLMDummyInputsBuilder,
 )
 class NVLM_D_Model(InternVLChatModel):
-    def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
+    def _init_mlp1(self, config: PreTrainedConfig) -> nn.Module:
         vit_hidden_size = config.vision_config.hidden_size
         llm_intermediate_size = config.text_config.intermediate_size
         llm_hidden_size = config.text_config.hidden_size
@@ -181,7 +179,7 @@ def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
 
     def _init_vision_model(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
diff --git a/vllm/model_executor/models/olmo_hybrid.py b/vllm/model_executor/models/olmo_hybrid.py
index 9a24df5ec9a6..e32087ea2134 100644
--- a/vllm/model_executor/models/olmo_hybrid.py
+++ b/vllm/model_executor/models/olmo_hybrid.py
@@ -143,6 +143,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             rope_parameters["rope_theta"] is not None
         )
 
+        self.rotary_emb: nn.Module | None
         if self._use_rope:
             self.rotary_emb = get_rope(
                 self.head_dim,
@@ -183,6 +184,7 @@ def forward(
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         q, k = self._apply_qk_norm(q, k)
         if self._use_rope:
+            assert self.rotary_emb is not None
             q, k = self.rotary_emb(positions, q, k)
         attn_output = self.attn(q, k, v)
         output, _ = self.o_proj(attn_output)
diff --git a/vllm/model_executor/models/olmoe.py b/vllm/model_executor/models/olmoe.py
index 7867caa97c57..175f25259ea2 100644
--- a/vllm/model_executor/models/olmoe.py
+++ b/vllm/model_executor/models/olmoe.py
@@ -32,11 +32,13 @@
 from vllm.distributed.utils import split_tensor_along_last_dim
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -84,11 +86,9 @@ def __init__(
         self.hidden_size = hidden_size
 
         # Gate always runs at half / full precision for now.
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             hidden_size,
             num_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
diff --git a/vllm/model_executor/models/opencua.py b/vllm/model_executor/models/opencua.py
index 6a875c813cb9..6b12d1d3fbe7 100644
--- a/vllm/model_executor/models/opencua.py
+++ b/vllm/model_executor/models/opencua.py
@@ -7,6 +7,7 @@
 
 """Inference-only OpenCUA-7B model compatible with HuggingFace weights."""
 
+import typing
 from collections.abc import Mapping, Sequence
 from typing import Any
 
@@ -27,6 +28,7 @@
 )
 from vllm.multimodal.parse import MultiModalDataItems
 from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
     BaseMultiModalProcessor,
     PromptReplacement,
     PromptUpdate,
@@ -178,7 +180,13 @@ def get_replacement_opencua(item_idx: int):
         ]
 
 
-class OpenCUADummyInputsBuilder(Qwen2VLDummyInputsBuilder):
+if typing.TYPE_CHECKING:
+    _OpenCUADummyInputsBuilderBase = BaseDummyInputsBuilder[OpenCUAProcessingInfo]
+else:
+    _OpenCUADummyInputsBuilderBase = Qwen2VLDummyInputsBuilder
+
+
+class OpenCUADummyInputsBuilder(_OpenCUADummyInputsBuilderBase):
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_images = mm_counts.get("image", 0)
 
@@ -220,7 +228,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)
         config = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
-        multimodal_config = vllm_config.model_config.multimodal_config
+        multimodal_config = vllm_config.model_config.get_multimodal_config()
 
         self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
         self.config = config
diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py
index 60b5b301d56e..b6f001dbb980 100644
--- a/vllm/model_executor/models/openpangu.py
+++ b/vllm/model_executor/models/openpangu.py
@@ -26,7 +26,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ParallelConfig, VllmConfig
@@ -43,7 +43,10 @@
     Attention,
     StaticSinkAttention,
 )
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -131,7 +134,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 class OpenPanguMoE(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         parallel_config: ParallelConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -149,11 +152,9 @@ def __init__(
         self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
         check_ffn_act_fn(config.hidden_act)
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.n_routed_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
         if (
@@ -175,6 +176,7 @@ def __init__(
         self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
         self.n_local_physical_experts = self.n_physical_experts // self.ep_size
 
+        self.shared_experts: OpenPanguMLP | None
         if config.n_shared_experts is not None:
             intermediate_size = config.moe_intermediate_size * config.n_shared_experts
             self.shared_experts = OpenPanguMLP(
@@ -238,7 +240,7 @@ def forward(
 class OpenPanguMLAAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         qk_nope_head_dim: int,
@@ -283,7 +285,7 @@ def __init__(
             )
             self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
             self.q_b_proj = ColumnParallelLinear(
-                q_lora_rank,
+                self.q_lora_rank,
                 self.num_heads * self.qk_head_dim,
                 bias=False,
                 quant_config=quant_config,
@@ -387,7 +389,7 @@ def forward(
 class OpenPanguEmbeddedAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
@@ -499,7 +501,7 @@ def forward(
 
     def _init_rotary_emb(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
     ) -> None:
         is_neox_style = True
@@ -520,7 +522,7 @@ def _init_rotary_emb(
 class OpenPanguSinkAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
@@ -568,6 +570,9 @@ def __init__(
         self.qk_nope_dim = getattr(config, "qk_nope_dim", None)
         self.qk_rope_dim = getattr(config, "qk_rope_dim", None)
         self.v_channels = getattr(config, "v_channels", None)
+        assert self.qk_nope_dim is not None
+        assert self.qk_rope_dim is not None
+        assert self.v_channels is not None
         self.head_dim = self.qk_rope_dim + self.qk_nope_dim
         self.q_size = self.num_heads * self.head_dim
         self.k_size = self.num_kv_heads * self.head_dim
@@ -677,14 +682,18 @@ def __init__(
                     },
                 )
             else:
-                self.param_sink_value = torch.zeros(
-                    (
-                        self.param_sink_number,
-                        self.num_kv_heads,
-                        self.v_channels,
+                self.register_buffer(
+                    "param_sink_value",
+                    torch.zeros(
+                        (
+                            self.param_sink_number,
+                            self.num_kv_heads,
+                            self.v_channels,
+                        ),
+                        device=current_platform.current_device(),
+                        dtype=config.torch_dtype,
                     ),
-                    device=current_platform.current_device(),
-                    dtype=config.torch_dtype,
+                    persistent=False,
                 )
         # To enable dummy run with out weight
         self.post_weight_load()
@@ -734,7 +743,7 @@ def forward(
 
     def _init_rotary_emb(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         rope_parameters: dict[str, Any] | None,
         quant_config: QuantizationConfig | None,
     ) -> None:
@@ -760,7 +769,7 @@ def post_weight_load(self) -> None:
 class OpenPanguDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         vllm_config: VllmConfig,
     ) -> None:
diff --git a/vllm/model_executor/models/openpangu_mtp.py b/vllm/model_executor/models/openpangu_mtp.py
index 3a04ccdff5be..c91453342a40 100644
--- a/vllm/model_executor/models/openpangu_mtp.py
+++ b/vllm/model_executor/models/openpangu_mtp.py
@@ -52,7 +52,9 @@ class OpenPanguMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer):
     def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
         nn.Module.__init__(self)
 
-        config = vllm_config.speculative_config.draft_model_config.hf_config
+        speculative_config = vllm_config.speculative_config
+        assert speculative_config is not None
+        config = speculative_config.draft_model_config.hf_config
         self.config = config
         quant_config = vllm_config.quant_config
 
@@ -200,10 +202,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 break
             else:
                 for mapping in expert_params_mapping:
-                    param_name, weight_name, expert_id, shard_id = mapping
-                    if weight_name not in name:
+                    (
+                        expert_param_name,
+                        expert_weight_name,
+                        expert_id,
+                        expert_shard_id,
+                    ) = mapping
+                    if expert_weight_name not in name:
                         continue
-                    name = name.replace(weight_name, param_name)
+                    name = name.replace(expert_weight_name, expert_param_name)
 
                     param = params_dict[name]
                     weight_loader = param.weight_loader
@@ -211,7 +218,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                         param,
                         loaded_weight,
                         name,
-                        shard_id=shard_id,
+                        shard_id=expert_shard_id,
                         expert_id=expert_id,
                     )
                     break
@@ -235,8 +242,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loaded_params
 
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
diff --git a/vllm/model_executor/models/openpangu_vl.py b/vllm/model_executor/models/openpangu_vl.py
index 10b72f929216..b7d5c8103282 100644
--- a/vllm/model_executor/models/openpangu_vl.py
+++ b/vllm/model_executor/models/openpangu_vl.py
@@ -22,7 +22,7 @@
 
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 from functools import lru_cache, partial
-from typing import Annotated, Literal, Optional
+from typing import Annotated, Literal, Optional, TypeAlias
 
 import torch
 import torch.nn as nn
@@ -627,7 +627,7 @@ def _padding_weight(name: str, w: torch.Tensor) -> torch.Tensor:
             pad[-(dim + 1) * 2 + 1] = pad_len
             return F.pad(w, pad, mode="constant", value=0)
 
-        stacked_params_mapping = [
+        stacked_params_mapping: list[tuple[str, str, str | int]] = [
             ("attn.qkv.", "attn.q.", "q"),
             ("attn.qkv.", "attn.k.", "k"),
             ("attn.qkv.", "attn.v.", "v"),
@@ -677,20 +677,15 @@ class OpenPanguVLProcessingInfo(Qwen2_5_VLProcessingInfo):
     def get_hf_config(self):
         return self.ctx.model_config.hf_config
 
-    def get_hf_processor(
-        self,
-        *,
-        min_pixels: int | None = None,
-        max_pixels: int | None = None,
-        size: dict[str, int] | None = None,
-        fps: float | list[float] | None = None,
-        **kwargs: object,
-    ):
+    def get_hf_processor(self, **kwargs: object):
+        kwargs.pop("min_pixels", None)
+        kwargs.pop("max_pixels", None)
+        kwargs.pop("size", None)
+        fps = kwargs.pop("fps", None)
         if fps is not None:
             kwargs["fps"] = fps
 
         return self.ctx.get_hf_processor(
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
@@ -747,11 +742,19 @@ class OpenPanguVLVideoEmbeddingInputs(TensorSchema):
     ]
 
 
+OpenPanguVLImageInputs: TypeAlias = (
+    OpenPanguVLImagePixelInputs | OpenPanguVLImageEmbeddingInputs
+)
+OpenPanguVLVideoInputs: TypeAlias = (
+    OpenPanguVLVideoPixelInputs | OpenPanguVLVideoEmbeddingInputs
+)
+
+
 class OpenPanguVLMultiModalProcessor(Qwen2_5_VLMultiModalProcessor):
     def _get_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, any],
+        hf_processor_mm_kwargs: Mapping[str, object],
         out_mm_kwargs: MultiModalKwargsItems,
     ) -> Sequence[PromptUpdate]:
         hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
@@ -878,10 +881,12 @@ def _parse_preprocess_params(self, vision_config):
         self.do_rescale = image_processor.do_rescale
         self.rescale_factor = image_processor.rescale_factor
         self.do_normalize = image_processor.do_normalize
-        self.image_mean = tuple(image_processor.image_mean)
-        self.image_std = tuple(image_processor.image_std)
+        self.image_mean: tuple[float, ...] = tuple(image_processor.image_mean)
+        self.image_std: tuple[float, ...] = tuple(image_processor.image_std)
 
-    def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):
+    def _maybe_ignore_quant_config(
+        self, quant_config: QuantizationConfig | None
+    ) -> QuantizationConfig | None:
         if isinstance(quant_config, AutoGPTQConfig):
             return None
         return quant_config
@@ -904,7 +909,9 @@ def _validate_and_reshape_mm_tensor(
         else:
             return torch.concat(mm_input)
 
-    def _parse_and_validate_image_input(self, **kwargs: object):
+    def _parse_and_validate_image_input(
+        self, **kwargs: object
+    ) -> OpenPanguVLImageInputs | None:
         pixel_values = kwargs.pop("pixel_values", None)
         image_embeds = kwargs.pop("image_embeds", None)
         image_grid_thw = kwargs.pop("image_grid_thw", None)
@@ -932,26 +939,27 @@ def _parse_and_validate_image_input(self, **kwargs: object):
                 image_grid_thw=image_grid_thw,
             )
 
-        if image_embeds is not None:
-            image_embeds = self._validate_and_reshape_mm_tensor(
-                image_embeds, "image embeds"
-            )
-            image_grid_thw = self._validate_and_reshape_mm_tensor(
-                image_grid_thw, "image grid_thw"
-            )
+        assert image_embeds is not None
+        image_embeds = self._validate_and_reshape_mm_tensor(
+            image_embeds, "image embeds"
+        )
+        image_grid_thw = self._validate_and_reshape_mm_tensor(
+            image_grid_thw, "image grid_thw"
+        )
 
-            if not isinstance(image_embeds, torch.Tensor):
-                raise ValueError(
-                    "Incorrect type of image embeddings. "
-                    f"Got type: {type(image_embeds)}"
-                )
-            return OpenPanguVLImageEmbeddingInputs(
-                type="image_embeds",
-                image_embeds=image_embeds,
-                image_grid_thw=image_grid_thw,
+        if not isinstance(image_embeds, torch.Tensor):
+            raise ValueError(
+                f"Incorrect type of image embeddings. Got type: {type(image_embeds)}"
             )
+        return OpenPanguVLImageEmbeddingInputs(
+            type="image_embeds",
+            image_embeds=image_embeds,
+            image_grid_thw=image_grid_thw,
+        )
 
-    def _parse_and_validate_video_input(self, **kwargs: object):
+    def _parse_and_validate_video_input(
+        self, **kwargs: object
+    ) -> OpenPanguVLVideoInputs | None:
         pixel_values_videos = kwargs.pop("pixel_values_videos", None)
         video_embeds = kwargs.pop("video_embeds", None)
         video_grid_thw = kwargs.pop("video_grid_thw", None)
@@ -973,27 +981,30 @@ def _parse_and_validate_video_input(self, **kwargs: object):
                 video_grid_thw=video_grid_thw,
             )
 
-        if video_embeds is not None:
-            video_embeds = self._validate_and_reshape_mm_tensor(
-                video_embeds, "video embeds"
-            )
-            video_grid_thw = self._validate_and_reshape_mm_tensor(
-                video_grid_thw, "video grid_thw"
-            )
+        assert video_embeds is not None
+        video_embeds = self._validate_and_reshape_mm_tensor(
+            video_embeds, "video embeds"
+        )
+        video_grid_thw = self._validate_and_reshape_mm_tensor(
+            video_grid_thw, "video grid_thw"
+        )
 
-            if not isinstance(video_embeds, torch.Tensor):
-                raise ValueError(
-                    "Incorrect type of video embeddings. "
-                    f"Got type: {type(video_embeds)}"
-                )
-            return OpenPanguVLVideoEmbeddingInputs(
-                type="video_embeds",
-                video_embeds=video_embeds,
-                video_grid_thw=video_grid_thw,
+        if not isinstance(video_embeds, torch.Tensor):
+            raise ValueError(
+                f"Incorrect type of video embeddings. Got type: {type(video_embeds)}"
             )
+        return OpenPanguVLVideoEmbeddingInputs(
+            type="video_embeds",
+            video_embeds=video_embeds,
+            video_grid_thw=video_grid_thw,
+        )
 
-    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
-        mm_input_by_modality = {}
+    def _parse_and_validate_multimodal_inputs(
+        self, **kwargs: object
+    ) -> dict[str, OpenPanguVLImageInputs | OpenPanguVLVideoInputs | None]:
+        mm_input_by_modality: dict[
+            str, OpenPanguVLImageInputs | OpenPanguVLVideoInputs | None
+        ] = {}
         for input_key in kwargs:
             if (
                 input_key in ("pixel_values", "image_embeds")
@@ -1019,16 +1030,19 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
         multimodal_embeddings: tuple[torch.Tensor, ...] = ()
 
         for modality in mm_input_by_modality:
-            multimodal_input = mm_input_by_modality[modality]
             if modality == "image":
-                vision_embeddings = self._process_image_input(multimodal_input)
+                image_input = mm_input_by_modality["image"]
+                assert isinstance(image_input, OpenPanguVLImageInputs)
+                vision_embeddings = self._process_image_input(image_input)
                 multimodal_embeddings = (
                     multimodal_embeddings
                     if not vision_embeddings
                     else (multimodal_embeddings + vision_embeddings)
                 )
             if modality == "video":
-                video_embeddings = self._process_video_input(multimodal_input)
+                video_input = mm_input_by_modality["video"]
+                assert isinstance(video_input, OpenPanguVLVideoInputs)
+                video_embeddings = self._process_video_input(video_input)
                 multimodal_embeddings = (
                     multimodal_embeddings
                     if not video_embeddings
@@ -1041,20 +1055,19 @@ def get_input_embeddings(
         input_ids: torch.Tensor,
         multimodal_embeddings: MultiModalEmbeddings | None = None,
     ) -> torch.Tensor:
-        if multimodal_embeddings is None:
-            return self.language_model.embed_input_ids(input_ids)
-
-        mm_token_ids = input_ids.new_tensor(
-            [self.config.image_token_id, self.config.video_token_id]
-        )
-        is_multimodal = torch.isin(input_ids, mm_token_ids)
-        return self.embed_input_ids(
-            input_ids,
-            multimodal_embeddings,
-            is_multimodal=is_multimodal,
-        )
+        inputs_embeds = self.language_model.embed_input_ids(input_ids)
+        if multimodal_embeddings is not None:
+            inputs_embeds = self.embed_input_ids(  # type: ignore[call-overload]
+                input_ids,
+                inputs_embeds,
+                multimodal_embeddings,
+                [self.config.image_token_id, self.config.video_token_id],
+            )
+        return inputs_embeds
 
-    def _process_image_input(self, image_input) -> tuple[torch.Tensor, ...]:
+    def _process_image_input(
+        self, image_input: OpenPanguVLImageInputs
+    ) -> tuple[torch.Tensor, ...]:
         grid_thw = image_input["image_grid_thw"]
         if grid_thw.ndim != 2:
             raise ValueError(f"grid_thw.ndim must be 2, but it is {grid_thw.ndim}")
@@ -1085,7 +1098,9 @@ def _process_image_input(self, image_input) -> tuple[torch.Tensor, ...]:
         sizes = grid_thw.prod(-1) // merge_size // merge_size
         return image_embeds.split(sizes.tolist())
 
-    def _process_video_input(self, video_input) -> torch.Tensor:
+    def _process_video_input(
+        self, video_input: OpenPanguVLVideoInputs
+    ) -> tuple[torch.Tensor, ...]:
         grid_thw = video_input["video_grid_thw"]
         if grid_thw.ndim != 2:
             raise ValueError(f"grid_thw.ndim must be 2, but it is {grid_thw.ndim}")
@@ -1135,9 +1150,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="visual.merger.",
@@ -1160,8 +1173,13 @@ def iter_mm_grid_thw(
         for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
             offset = mm_feature.mm_position.offset
             modality = mm_feature.modality
+            feature_data = mm_feature.data
+            assert feature_data is not None
             if modality == "image":
-                t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
+                grid_thw_item = feature_data["image_grid_thw"]
+                grid_thw = grid_thw_item.data
+                assert isinstance(grid_thw, torch.Tensor)
+                t, h, w = grid_thw.tolist()
                 assert t == 1, f"Image must have 1 frame, got {t}"
                 yield (
                     modality,
@@ -1171,7 +1189,10 @@ def iter_mm_grid_thw(
                     w // spatial_merge_size,
                 )
             elif modality == "video":
-                t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
+                grid_thw_item = feature_data["video_grid_thw"]
+                grid_thw = grid_thw_item.data
+                assert isinstance(grid_thw, torch.Tensor)
+                t, h, w = grid_thw.tolist()
                 yield (
                     modality,
                     offset,
@@ -1283,14 +1304,17 @@ def normalize(image, mean, std):
 @lru_cache(maxsize=10)
 def _fuse_mean_std_and_rescale_factor(
     do_normalize: bool | None = None,
-    image_mean: float | list[float] | None = None,
-    image_std: float | list[float] | None = None,
+    image_mean: float | tuple[float, ...] | None = None,
+    image_std: float | tuple[float, ...] | None = None,
     do_rescale: bool | None = None,
     rescale_factor: float | None = None,
     device: Optional["torch.device"] = None,
 ) -> tuple:
     if do_rescale and do_normalize:
         # Fused rescale and normalize
+        assert rescale_factor is not None
+        assert image_mean is not None
+        assert image_std is not None
         image_mean = torch.tensor(image_mean, device=device) * (1.0 / rescale_factor)
         image_std = torch.tensor(image_std, device=device) * (1.0 / rescale_factor)
         do_rescale = False
@@ -1302,13 +1326,11 @@ def rescale_and_normalize(
     do_rescale: bool,
     rescale_factor: float,
     do_normalize: bool,
-    image_mean: float | list[float],
-    image_std: float | list[float],
+    image_mean: float | tuple[float, ...],
+    image_std: float | tuple[float, ...],
     dtype: torch.dtype = torch.bfloat16,
 ) -> "torch.Tensor":
-    """
-    Rescale and normalize images.
-    """
+    """Rescale and normalize images."""
     image_mean, image_std, do_rescale = _fuse_mean_std_and_rescale_factor(
         do_normalize=do_normalize,
         image_mean=image_mean,
diff --git a/vllm/model_executor/models/openvla.py b/vllm/model_executor/models/openvla.py
index a676b3447cf4..9242f7e1f51f 100644
--- a/vllm/model_executor/models/openvla.py
+++ b/vllm/model_executor/models/openvla.py
@@ -9,7 +9,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.linear import (
@@ -24,7 +24,11 @@
     SupportsPP,
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.inputs import (
+    MultiModalFieldConfig,
+    MultiModalKwargsItem,
+    MultiModalKwargsItems,
+)
 from vllm.multimodal.parse import (
     ImageEmbeddingItems,
     ImageProcessorItems,
@@ -69,12 +73,11 @@ def _get_num_image_tokens(image_size: int) -> int:
 
 
 class OpenVLAImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (6)
-        - h: Height
-        - w: Width
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (6)
+    - h: Height
+    - w: Width
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -296,18 +299,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        image_overrides = mm_options.get("image")
         image_size = self.info.get_image_size_with_most_features()
 
         return {
             "image": self._get_dummy_images(
                 width=image_size.width,
                 height=image_size.height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -346,6 +347,7 @@ def get_insertion(item_idx: int) -> PromptUpdateDetails:
             if isinstance(images, ImageEmbeddingItems):
                 num_image_tokens = images.get_feature_size(item_idx)
             else:
+                assert isinstance(images, ImageProcessorItems)
                 image_size = images.get_image_size(item_idx)
                 num_image_tokens = self.info.get_num_image_tokens(
                     image_width=image_size.width,
@@ -487,11 +489,15 @@ def get_mm_mapping(self) -> MultiModelKeys:
             tower_model="vision_backbone",
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        return num_image_tokens
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
-        return num_vision_tokens
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+        return num_mm_embeds, num_mm_embeds
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         def maybe_rename_vision_weights(
diff --git a/vllm/model_executor/models/opt.py b/vllm/model_executor/models/opt.py
index 1efe9eb8a8ca..4726c72b417b 100644
--- a/vllm/model_executor/models/opt.py
+++ b/vllm/model_executor/models/opt.py
@@ -218,6 +218,7 @@ def __init__(
         )
 
         # Project out & in will be replicated if they exist.
+        self.project_out: ReplicatedLinear | None
         if config.word_embed_proj_dim != config.hidden_size:
             self.project_out = ReplicatedLinear(
                 config.hidden_size,
@@ -229,6 +230,7 @@ def __init__(
         else:
             self.project_out = None
 
+        self.project_in: ReplicatedLinear | None
         if config.word_embed_proj_dim != config.hidden_size:
             self.project_in = ReplicatedLinear(
                 config.word_embed_proj_dim,
diff --git a/vllm/model_executor/models/orion.py b/vllm/model_executor/models/orion.py
index 2b58ee1e015c..4d50705c78f8 100644
--- a/vllm/model_executor/models/orion.py
+++ b/vllm/model_executor/models/orion.py
@@ -13,7 +13,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -163,7 +163,7 @@ def forward(
 class OrionDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/ovis.py b/vllm/model_executor/models/ovis.py
index 10b0eb60eea3..7c1668ddabcd 100644
--- a/vllm/model_executor/models/ovis.py
+++ b/vllm/model_executor/models/ovis.py
@@ -26,10 +26,10 @@
 import torch.nn as nn
 from torch import Tensor
 from torch.nn.functional import gumbel_softmax, pad, softmax
-from transformers import BatchFeature, PretrainedConfig
+from transformers import BatchFeature, PreTrainedConfig
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -59,7 +59,12 @@
 from vllm.transformers_utils.processors.ovis import OvisProcessor
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
-from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
+from .interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+    supports_pp,
+)
 
 # Cannot find the following number from hf config.
 IMAGE_TOKEN = "<image>"
@@ -88,7 +93,7 @@ def st_argmax(y_soft: torch.Tensor, dim: int):  # straight-through softmax
 class VisualTokenizer(torch.nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -115,7 +120,7 @@ def __init__(
 
     def _init_backbone(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> nn.Module:
@@ -200,7 +205,7 @@ def encode(self, pixel_values: torch.Tensor) -> torch.Tensor:
         return features
 
     def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
-        """[BatchSize, ImageShape] -> [BatchSize, Token, VocabSize]"""
+        """[BatchSize, ImageShape] -> [BatchSize, Token, VocabSize]."""
         features = self.encode(pixel_values)
         logits = self.head(features)
         tokens = self.tokenize(logits)
@@ -217,13 +222,12 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
 
 
 class OvisImagePatchInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * number of images * number of patches
-        - h: Height of each patch
-        - w: Width of each patch
-        - patch_indicators: Batch size * (number of patches + 1)
-        - bn: Batch size * number of images
+    """Dimensions:
+    - bnp: Batch size * number of images * number of patches
+    - h: Height of each patch
+    - w: Width of each patch
+    - patch_indicators: Batch size * (number of patches + 1)
+    - bn: Batch size * number of images
     """
 
     type: Literal["image_patches"]
@@ -282,7 +286,7 @@ def get_image_segment_len(self) -> int:
         # minus 1 for presented image token
         return (patch_grid_length // hidden_stride) ** 2 - 1
 
-    def get_image_pad_token(self) -> str:
+    def get_image_pad_token(self) -> str | None:
         hf_text_config = self.get_hf_config().get_text_config()
         text_model_type = hf_text_config.model_type
         return IMAGE_PAD_TOKEN_MAP.get(text_model_type)
@@ -307,32 +311,30 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         mm_data = {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
         }
         return mm_data
 
 
 class OvisMultiModalProcessor(BaseMultiModalProcessor[OvisProcessingInfo]):
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def image_indicators_to_visual_tokens(
         self,
         image_indicators: list[int],
     ) -> list[int]:
-        """
-        Filter image indicators placeholders and convert them to corresponding
+        """Filter image indicators placeholders and convert them to corresponding
         tokens in visual tokenizer.
         For example, [-301, -300, -302, -300, -303, -300, -304, -300, -305]
         should return [vocab_size-1, vocab_size-2, ..., vocab_size-5]
@@ -342,9 +344,6 @@ def image_indicators_to_visual_tokens(
         # -300 is image_atom token, filter them out
         return [vte_vocab_size + x + 300 for x in image_indicators if x < -300]
 
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
-        return self.dummy_inputs.get_dummy_text(mm_counts)
-
     def _postprocess_hf_mm_data(
         self,
         mm_data: Mapping[str, object],
@@ -420,7 +419,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         config = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
 
-        self.config: PretrainedConfig = config
+        self.config: PreTrainedConfig = config
 
         with self._mark_language_model(vllm_config):
             self.llm = init_vllm_registered_model(
@@ -441,8 +440,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         text_model_type = self.config.get_text_config().model_type
         self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type]
 
+        language_model = self.get_language_model()
+        assert supports_pp(language_model)
         self.make_empty_intermediate_tensors = (
-            self.get_language_model().make_empty_intermediate_tensors
+            language_model.make_empty_intermediate_tensors
         )
 
     def _parse_and_validate_image_input(
diff --git a/vllm/model_executor/models/ovis2_5.py b/vllm/model_executor/models/ovis2_5.py
index c55f4805377a..456ae0b7fe5e 100644
--- a/vllm/model_executor/models/ovis2_5.py
+++ b/vllm/model_executor/models/ovis2_5.py
@@ -4,14 +4,14 @@
 
 from collections.abc import Iterable, Mapping
 from functools import partial
-from typing import Annotated, Literal
+from typing import Annotated, Literal, TypedDict
 
 import torch
 import torch.nn as nn
-from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig
+from transformers import BaseImageProcessor, BatchFeature, PreTrainedConfig
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -40,7 +40,12 @@
 from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
-from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
+from .interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+    supports_pp,
+)
 
 IMAGE_TOKEN = "<image>"
 VIDEO_TOKEN = "<video>"
@@ -49,12 +54,11 @@
 
 
 class Ovis2_5ImagePatchInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * number of images * number of patches
-        - patch_size: patch_size_x * patch_size_y * num_channels
-        - patch_indicators: Batch size * (number of patches + 1)
-        - bn: Batch size * number of images
+    """Dimensions:
+    - bnp: Batch size * number of images * number of patches
+    - patch_size: patch_size_x * patch_size_y * num_channels
+    - patch_indicators: Batch size * (number of patches + 1)
+    - bn: Batch size * number of images
     """
 
     type: Literal["image_patches"]
@@ -66,12 +70,11 @@ class Ovis2_5ImagePatchInputs(TensorSchema):
 
 
 class Ovis2_5VideoPatchInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * number of videos * number of patches
-        - patch_size: patch_size_x * patch_size_y * num_channels
-        - patch_indicators: Batch size * (number of patches + 1)
-        - bn: Batch size * number of videos
+    """Dimensions:
+    - bnp: Batch size * number of videos * number of patches
+    - patch_size: patch_size_x * patch_size_y * num_channels
+    - patch_indicators: Batch size * (number of patches + 1)
+    - bn: Batch size * number of videos
     """
 
     type: Literal["video_patches"]
@@ -82,14 +85,17 @@ class Ovis2_5VideoPatchInputs(TensorSchema):
     # This is used to restore the first two dimensions of `flat_data`.
 
 
+class Ovis2_5MultiModalInputs(TypedDict, total=False):
+    images: Ovis2_5ImagePatchInputs | None
+    videos: Ovis2_5VideoPatchInputs | None
+
+
 class VisualTokenizer(torch.nn.Module):
-    """
-    VIT
-    """
+    """VIT."""
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         visual_vocab_size: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -115,7 +121,7 @@ def __init__(
 
     def _init_backbone(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -284,32 +290,26 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(
             seq_len, mm_counts
         )
 
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-
         mm_data = {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
         return mm_data
@@ -320,8 +320,7 @@ def visual_indicators_to_visual_tokens(
         self,
         visual_indicators: list[int],
     ) -> list[int]:
-        """
-        Filter image indicators placeholders and convert them to corresponding
+        """Filter image indicators placeholders and convert them to corresponding
         tokens in visual tokenizer.
         """
         hf_config = self.info.get_hf_config()
@@ -332,7 +331,7 @@ def visual_indicators_to_visual_tokens(
             if x >= INDICATOR_IDS[0]
         ]
 
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _postprocess_hf_mm_data(
@@ -439,7 +438,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         config = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
 
-        self.config: PretrainedConfig = config
+        self.config: PreTrainedConfig = config
 
         with self._mark_language_model(vllm_config):
             self.llm = init_vllm_registered_model(
@@ -458,8 +457,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         self.image_pad_token_id: int = IMAGE_PAD_TOKEN_ID
 
+        language_model = self.get_language_model()
+        assert supports_pp(language_model)
         self.make_empty_intermediate_tensors = (
-            self.get_language_model().make_empty_intermediate_tensors
+            language_model.make_empty_intermediate_tensors
         )
 
     def _parse_and_validate_image_input(
@@ -567,8 +568,10 @@ def _process_visual_input(
             vision_embeddings.append(torch.cat(vision_embeddings_per_image, dim=0))
         return tuple(vision_embeddings)
 
-    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
-        modalities = {}
+    def _parse_and_validate_multimodal_inputs(
+        self, **kwargs: object
+    ) -> Ovis2_5MultiModalInputs:
+        modalities: Ovis2_5MultiModalInputs = {}
 
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
@@ -598,10 +601,12 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
         for modality in modalities:
             if modality == "images":
                 image_input = modalities["images"]
+                assert image_input is not None
                 image_embeddings = self._process_visual_input(image_input)
                 multimodal_embeddings += tuple(image_embeddings)
             if modality == "videos":
                 video_input = modalities["videos"]
+                assert video_input is not None
                 video_embeddings = self._process_visual_input(video_input)
                 multimodal_embeddings += tuple(video_embeddings)
 
diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py
index 4a20531b76dd..1fde4cd0a11b 100644
--- a/vllm/model_executor/models/paddleocr_vl.py
+++ b/vllm/model_executor/models/paddleocr_vl.py
@@ -23,7 +23,7 @@
 import torch
 import torch.nn as nn
 from einops import rearrange
-from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig
+from transformers import BaseImageProcessor, BatchFeature, PreTrainedConfig
 from transformers.activations import GELUActivation
 from transformers.image_utils import ChannelDimension
 from transformers.modeling_outputs import (
@@ -32,7 +32,7 @@
 from transformers.utils import torch_int
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import parallel_state
 from vllm.distributed import utils as dist_utils
 from vllm.inputs import MultiModalDataDict
@@ -99,7 +99,6 @@ def smart_resize(
     3. The aspect ratio of the image is maintained as closely as possible.
 
     """
-
     if height < factor:
         width = round((width * factor) / height)
         height = factor
@@ -224,19 +223,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         max_image_size = self.info.get_image_size_with_most_features()
-        image_overrides = mm_options.get("image")
 
         return {
             "image": self._get_dummy_images(
                 width=max_image_size.width,
                 height=max_image_size.height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -244,36 +240,37 @@ def get_dummy_mm_data(
 class PaddleOCRVLMultiModalProcessor(
     BaseMultiModalProcessor[PaddleOCRVLProcessingInfo]
 ):
-    def _apply_hf_processor_main(
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
+    def _call_hf_processor(
         self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
-
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
-
-        final_mm_kwargs = dict(hf_processor_mm_kwargs or {})
+        final_mm_kwargs = dict(hf_kwargs or {})
         final_mm_kwargs.setdefault("images_kwargs", {})
         # vLLM use PIL.Image, always set channel_last
         final_mm_kwargs["input_data_format"] = ChannelDimension.LAST
-        processed_data = self.info.ctx.call_hf_processor(
+        return self.info.ctx.call_hf_processor(
             self.info.get_hf_processor(**final_mm_kwargs),
-            dict(text=prompt_text, **mm_data),
-            hf_processor_mm_kwargs,
+            hf_data,
+            hf_kwargs,
         )
+
+    def _postprocess_hf_mm_data(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+        processed_data: BatchFeature,
+    ) -> BatchFeature:
+        if not hf_data:
+            return processed_data
+
         num_patches_per_image = processed_data["image_grid_thw"].prod(-1)
         processed_data["pixel_values"] = processed_data["pixel_values"].split(
             num_patches_per_image.tolist()
         )
-        processed_data.update(passthrough_data)
-
         return processed_data
 
     def _get_mm_fields_config(
@@ -321,8 +318,8 @@ def get_replacement(item_idx: int, image_processor):
 class Projector(nn.Module):
     def __init__(
         self,
-        text_config: PretrainedConfig,
-        vision_config: PretrainedConfig,
+        text_config: PreTrainedConfig,
+        vision_config: PreTrainedConfig,
         prefix: str = "",
     ):
         super().__init__()
@@ -395,7 +392,7 @@ class PaddleOCRImagePixelInputs(TensorSchema):
 
 
 class SiglipVisionEmbeddings(nn.Module):
-    def __init__(self, config: PretrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
         self.config = config
         self.embed_dim = config.hidden_size
@@ -685,7 +682,7 @@ def forward(self, seqlen: int) -> torch.Tensor:
 class SiglipEncoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -738,7 +735,7 @@ def forward(
 class SiglipEncoder(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -840,7 +837,7 @@ def forward(
 class SiglipVisionTransformer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -1015,14 +1012,14 @@ def compute_logits(
     def iter_mm_grid_thw(
         self, mm_features: list[MultiModalFeatureSpec]
     ) -> Iterator[tuple[int, int, int, int, float]]:
-        """
-        Iterate over multimodal features and yield grid information.
+        """Iterate over multimodal features and yield grid information.
 
         Args:
             mm_features: List of multimodal feature specifications
 
         Yields:
             Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image
+
         """
         spatial_merge_size = self.config.vision_config.spatial_merge_size
         tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0)
diff --git a/vllm/model_executor/models/paligemma.py b/vllm/model_executor/models/paligemma.py
index 2fa498ee261c..d0173e4c116e 100644
--- a/vllm/model_executor/models/paligemma.py
+++ b/vllm/model_executor/models/paligemma.py
@@ -8,12 +8,13 @@
 from transformers import BatchFeature, PaliGemmaConfig
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict, MultiModalInput
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import (
@@ -56,12 +57,11 @@
 
 
 class PaliGemmaImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -69,11 +69,10 @@ class PaliGemmaImagePixelInputs(TensorSchema):
 
 
 class PaliGemmaImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - ifs: Image feature size
-        - hs: Hidden size (must match language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - ifs: Image feature size
+    - hs: Hidden size (must match language model backbone)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -131,22 +130,18 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         hf_config = self.info.get_hf_config()
         vision_config = hf_config.vision_config
         max_image_size = vision_config.image_size
 
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=max_image_size,
                 height=max_image_size,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -401,8 +396,12 @@ def get_mm_mapping(self) -> MultiModelKeys:
             tower_model="vision_tower",
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        return num_image_tokens
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
-        return num_vision_tokens
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
+        return num_mm_embeds, num_mm_embeds
diff --git a/vllm/model_executor/models/parakeet.py b/vllm/model_executor/models/parakeet.py
index ba7b96b8244a..652746e5ae6c 100644
--- a/vllm/model_executor/models/parakeet.py
+++ b/vllm/model_executor/models/parakeet.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Modules below used for the audio encoder component in: models/nano_nemotron_vl.py
-"""
+"""Modules below used for the audio encoder component in: models/nano_nemotron_vl.py."""
 
 from collections.abc import Iterable
 from functools import cache
@@ -12,7 +10,7 @@
 import torch
 import torch.nn as nn
 from transformers import ParakeetEncoder as HFParakeetEncoder
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.audio_utils import mel_filter_bank
 
 from vllm.logger import init_logger
@@ -49,7 +47,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 class ProjectedParakeet(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         *,
         dtype: torch.dtype,
         llm_hidden_size: int,
@@ -137,7 +135,7 @@ def _can_skip_missing_named_param(self, target_name: str) -> bool:
 
 
 class ParakeetExtractor:
-    def __init__(self, config: PretrainedConfig) -> None:
+    def __init__(self, config: PreTrainedConfig) -> None:
         self.config = ExtractorConfig.from_hf_config(config)
         """`config` is named *exactly* for `._get_subsampling_output_length` below"""
         self._clip_target_samples = int(
@@ -331,6 +329,6 @@ def __call__(
         }
 
     @staticmethod
-    def audio_length(raw_config: PretrainedConfig, audio_tokens: int) -> int:
+    def audio_length(raw_config: PreTrainedConfig, audio_tokens: int) -> int:
         config = ExtractorConfig.from_hf_config(raw_config)
         return int(audio_tokens * config.subsampling_factor * config.hop_length)
diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py
index 8b7beb969024..c635eab2976c 100644
--- a/vllm/model_executor/models/param2moe.py
+++ b/vllm/model_executor/models/param2moe.py
@@ -22,7 +22,6 @@
 from itertools import islice
 
 import torch
-import torch.nn.functional as F
 from torch import nn
 
 from vllm.config import CacheConfig, VllmConfig
@@ -34,6 +33,7 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     MoERunner,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -88,8 +88,7 @@ def _normalize_expert_bias(
 
 
 class Param2MoEAttention(nn.Module):
-    """
-    Grouped-Query Attention (GQA) for Param2MoE.
+    """Grouped-Query Attention (GQA) for Param2MoE.
 
     Notable differences from a vanilla GQA layer:
       * The checkpoint fuses Q, K, V into a single ``query_key_value`` weight.
@@ -252,8 +251,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class Param2MoEMoEBlock(nn.Module):
-    """
-    Mixture-of-Experts block for Param2MoE.
+    """Mixture-of-Experts block for Param2MoE.
 
     Routing:
       * Sigmoid scoring  (config.score_function = "sigmoid")
@@ -291,10 +289,12 @@ def __init__(
         self.norm_expert_prob: bool = getattr(config, "norm_topk_prob", True)
         self.score_function: str = getattr(config, "score_function", "sigmoid")
 
-        self.gate = nn.Linear(
+        self.gate = GateLinear(
             self.hidden_size,
             self.num_experts,
-            bias=False,
+            out_dtype=torch.float32,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.gate",
         )
 
         if getattr(config, "moe_router_enable_expert_bias", True):
@@ -352,15 +352,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
 
-        # Router: both input and weight must be float32 for numerical
-        # stability (mirrors the original Param2MoEGate behaviour).
-        # The gate nn.Linear weight lives in the model dtype (bfloat16),
-        # so we must cast both explicitly via F.linear instead of calling
-        # self.gate() which would hit a dtype mismatch.
-        router_logits = F.linear(
-            hidden_states.float(),
-            self.gate.weight.float(),
-        ).to(hidden_states.dtype)
+        # Router logits are accumulated in float32 for numerical stability
+        # (mirrors the original Param2MoEGate behaviour) and handed to the
+        # experts in the activation dtype.
+        router_logits, _ = self.gate(hidden_states)
+        router_logits = router_logits.to(hidden_states.dtype)
 
         expert_output = self.experts(
             hidden_states=hidden_states,
@@ -371,8 +367,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class Param2MoEDecoderLayer(nn.Module):
-    """
-    Single transformer decoder block.
+    """Single transformer decoder block.
 
     Dense for the first ``first_k_dense_replace`` layers; MoE thereafter.
     """
@@ -605,8 +600,7 @@ def update_physical_experts_metadata(
 class Param2MoEForCausalLM(
     nn.Module, SupportsPP, SupportsLoRA, Param2MoEMixtureOfExperts
 ):
-    """
-    vLLM-native Param2MoE CausalLM.
+    """vLLM-native Param2MoE CausalLM.
 
     Uses Grouped-Query Attention (GQA) with a Sigmoid-scored,
     grouped-topk Mixture-of-Experts MLP.
diff --git a/vllm/model_executor/models/phi3.py b/vllm/model_executor/models/phi3.py
index 56c8755123d3..d641f89ffa73 100644
--- a/vllm/model_executor/models/phi3.py
+++ b/vllm/model_executor/models/phi3.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 # Adapted from llama.py
-"""Inference-only Phi3 model code inherit from Llama.py"""
+"""Inference-only Phi3 model code inherit from Llama.py."""
 
 from vllm.model_executor.models.llama import LlamaForCausalLM
 
diff --git a/vllm/model_executor/models/phi3v.py b/vllm/model_executor/models/phi3v.py
index 8cb387347d2f..e300ef6a4419 100644
--- a/vllm/model_executor/models/phi3v.py
+++ b/vllm/model_executor/models/phi3v.py
@@ -24,12 +24,12 @@
 from transformers import (
     BatchFeature,
     CLIPVisionConfig,
-    PretrainedConfig,
+    PreTrainedConfig,
     ProcessorMixin,
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -95,7 +95,7 @@
 
 
 def _init_img_processor(
-    hf_config: PretrainedConfig,
+    hf_config: PreTrainedConfig,
     quant_config: QuantizationConfig | None,
     prefix: str = "",
 ) -> CLIPVisionModel:
@@ -119,13 +119,12 @@ def _init_img_processor(
 
 
 class Phi3VImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - n: Number of images
-        - p: Number of patches
-        - h: Height of each patch
-        - w: Width of each patch
+    """Dimensions:
+    - b: Batch size
+    - n: Number of images
+    - p: Number of patches
+    - h: Height of each patch
+    - w: Width of each patch
     """
 
     type: Literal["pixel_values", "image_embeds"] = "pixel_values"
@@ -143,12 +142,11 @@ class Phi3VImagePixelInputs(TensorSchema):
 
 
 class Phi3VImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - n: Number of images
-        - f: Image feature size (e.g., number of tokens per image)
-        - h: Hidden size (must match language model backbone)
+    """Dimensions:
+    - b: Batch size
+    - n: Number of images
+    - f: Image feature size (e.g., number of tokens per image)
+    - h: Hidden size (must match language model backbone)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -167,7 +165,7 @@ class Phi3HDImageEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str = "",
     ) -> None:
@@ -228,8 +226,7 @@ def get_img_features(self, img_embeds: torch.FloatTensor) -> torch.FloatTensor:
     def forward(
         self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor
     ) -> torch.FloatTensor:
-        """
-        process image and return vision embeddings.
+        """Process image and return vision embeddings.
 
         pixel_values: (num_images, num_crops, c, h, w)
         output: (num_images, num_img_tokens, hidden_size)
@@ -244,9 +241,7 @@ def forward(
         return image_features_proj
 
     def hd_feature_transform(self, image_features, image_sizes):
-        """
-        image_features: (num_images, num_crops+1, 24*24, 1024)
-        """
+        """image_features: (num_images, num_crops+1, 24*24, 1024)."""
         assert self.hd_transform_order == "sub_glb", (
             f"hd_transform_order `{self.hd_transform_order}` not implemented"
         )
@@ -304,8 +299,7 @@ def hd_feature_transform(self, image_features, image_sizes):
         return batch_image_features_proj
 
     def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
-        """
-        image_features: (num_images*num_crops, 24*24, 1024)
+        """image_features: (num_images*num_crops, 24*24, 1024)
         output: (num_images, h_crop*12, w_crop*12, 4096)
         where h_crop*w_crop == num_crops
         """
@@ -329,8 +323,7 @@ def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
         return image_features_hd
 
     def add_image_newline(self, image_features_hd):
-        """
-        image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
+        """image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
         output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)
         """
         num_images, h, w, hid_dim = image_features_hd.shape
@@ -378,47 +371,24 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
 
 class Phi3VMultiModalProcessor(BaseMultiModalProcessor[Phi3VProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _postprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-        processed_data: BatchFeature,
-    ) -> BatchFeature:
-        if not mm_data:
-            return processed_data
-
-        input_ids = processed_data["input_ids"]
-        assert isinstance(input_ids, torch.Tensor)
-
-        # Phi3v processor has inserted -1, -2 etc as placeholder in prompt_ids,
-        # which will cause OverflowError when decoding the prompt_ids.
-        # Therefore, we need to do an early replacement here
-        input_ids.masked_fill_(input_ids < 0, _IMAGE_TOKEN_ID)
-
-        return processed_data
-
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
diff --git a/vllm/model_executor/models/phi4mm.py b/vllm/model_executor/models/phi4mm.py
index 345f58e5ca5e..852445597582 100644
--- a/vllm/model_executor/models/phi4mm.py
+++ b/vllm/model_executor/models/phi4mm.py
@@ -10,14 +10,14 @@
 import torch.nn as nn
 from transformers import (
     BatchFeature,
-    PretrainedConfig,
+    PreTrainedConfig,
     ProcessorMixin,
     SequenceFeatureExtractor,
     SiglipVisionConfig,
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import get_pp_group
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -45,6 +45,7 @@
 from vllm.multimodal.processing.processor import (
     BaseMultiModalProcessor,
     BaseProcessingInfo,
+    HFMultiModalInputs,
     PromptReplacement,
     PromptUpdate,
     ResolvedPromptUpdate,
@@ -122,7 +123,7 @@ class Phi4MMImageEncoder(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str = "",
         model_dir: str = "",
@@ -246,15 +247,13 @@ def forward(
         image_sizes: torch.Tensor,
         image_attention_mask: torch.Tensor,
     ) -> list[torch.FloatTensor]:
-        """
-        process image and return vision embeddings.
+        """Process image and return vision embeddings.
 
         pixel_values: (num_images, num_crops, c, h, w)
         image_sizes: [[h1, w1], [h2, w2]]
         image_attention_mask: num_images x num_crops x 32 x 32
         output: (num_images, num_img_tokens, hidden_size)
         """
-
         # eg
         # pixel_values: torch.Size([1, 7, 3, 448, 448])
         # image_sizes: tensor([[ 896, 1344]], device='cuda:0')
@@ -463,16 +462,15 @@ def forward(
 
 
 class Phi4MMImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - p: Number of patches (1 + num_patches)
-        - c: Number of channels (3)
-        - h: Height of each image patch
-        - w: Width of each image patch
-        - nc: Number of crops
-        - H_mask: Height of attention mask
-        - W_mask: Width of attention mask
+    """Dimensions:
+    - bn: Batch size * number of images
+    - p: Number of patches (1 + num_patches)
+    - c: Number of channels (3)
+    - h: Height of each image patch
+    - w: Width of each image patch
+    - nc: Number of crops
+    - H_mask: Height of attention mask
+    - W_mask: Width of attention mask
     """
 
     type: Literal["pixel_values"]
@@ -501,10 +499,9 @@ class Phi4MMImagePixelInputs(TensorSchema):
 
 
 class Phi4MMAudioFeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of audios
-        - t: Time frames (M)
+    """Dimensions:
+    - bn: Batch size * number of audios
+    - t: Time frames (M)
     """
 
     type: Literal["audio_features"]
@@ -516,12 +513,11 @@ class Phi4MMAudioFeatureInputs(TensorSchema):
 
 
 class Phi4MMAudioEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - n: Number of audios
-        - f: Audio feature size
-        - h: Hidden size (must match language model backbone)
+    """Dimensions:
+    - b: Batch size
+    - n: Number of audios
+    - f: Audio feature size
+    - h: Hidden size (must match language model backbone)
     """
 
     type: Literal["audio_embeds"]
@@ -543,9 +539,7 @@ def stack_with_pad(
     tensors: torch.Tensor | list[torch.Tensor],
     padding_value: int | float = 0,
 ) -> torch.Tensor:
-    """
-    Stack tensors, padding dimensions that differ across items.
-    """
+    """Stack tensors, padding dimensions that differ across items."""
     if isinstance(tensors, torch.Tensor):
         return tensors
 
@@ -653,8 +647,7 @@ def _compute_num_image_tokens(
         vit_patch_size: int,
         token_compression_factor: int = 2,
     ):
-        """
-        compute the number of tokens an image is expected to take up considering
+        """Compute the number of tokens an image is expected to take up considering
         the image encoder architecture and exclude output features containing
         only padding pixels
 
@@ -776,8 +769,7 @@ def get_image_size_with_most_features(self) -> ImageSize:
         return ImageSize(height=max_side, width=vit_image_size)
 
     def get_audio_num_frames(self, audio_len: int, sr: float) -> int:
-        """
-        Compute the output size of the `extract_features` method.
+        """Compute the output size of the `extract_features` method.
 
         Args:
             audio_len (int): Length of the input waveform in samples.
@@ -787,8 +779,8 @@ def get_audio_num_frames(self, audio_len: int, sr: float) -> int:
             tuple (int, int): Output size as (T, D), where:
                 T: Number of time frames.
                 D: Number of Mel filterbank bins (80).
-        """
 
+        """
         # Resample to 16000 or 8000 if needed
         if sr > 16000:
             resample_ratio = int(sr // 16000)
@@ -812,8 +804,7 @@ def get_audio_num_frames(self, audio_len: int, sr: float) -> int:
         return num_frames
 
     def _compute_audio_embed_size(self, audio_frames: int) -> int:
-        """
-        Compute the audio embedding size based on the audio frames and
+        """Compute the audio embedding size based on the audio frames and
         compression rate.
         """
         hf_config = self.get_hf_config()
@@ -848,28 +839,21 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-        audio_overrides = mm_options.get("audio")
-        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
-
         mm_data = {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "audio": self._get_dummy_audios(
                 length=_AUDIO_MAX_SOUNDFILE_SIZE,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             ),
         }
 
@@ -877,40 +861,38 @@ def get_dummy_mm_data(
 
 
 class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]):
-    def _apply_hf_processor_main(
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
 
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        # The Phi4MM processor expects "audios" instead of "audio"
+        hf_data = hf_inputs.hf_data
+        if "audio" in hf_data:
+            hf_data["audios"] = hf_data.pop("audio")
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        sr = self.info.get_feature_extractor(**hf_inputs.hf_kwargs).sampling_rate
+        if audio_data := hf_data.get("audios", []):
+            assert isinstance(audio_data, Sequence)
+            hf_data["audios"] = [(data, sr) for data in audio_data]
 
-        sr = self.info.get_feature_extractor(**hf_processor_mm_kwargs).sampling_rate
-        raw_audio_data = mm_data.get("audios", [])
-        assert isinstance(raw_audio_data, Sequence)
-        audio_data = []
-        for data in raw_audio_data:
-            assert isinstance(data, (list, np.ndarray, torch.Tensor))
-            audio_data.append(data)
+        return hf_inputs
 
-        processor_data = dict(text=prompt_text, **mm_data)
-        if audio_data:
-            processor_data["audios"] = [(data, sr) for data in audio_data]
-
-        processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            processor_data,
-            hf_processor_mm_kwargs,
-        )
+    def _postprocess_hf_mm_data(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+        processed_data: BatchFeature,
+    ) -> BatchFeature:
+        if not hf_data:
+            return processed_data
 
-        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        hf_processor = self.info.get_hf_processor(**hf_kwargs)
         num_img_tokens = [
             self.info.get_num_image_tokens(
                 image_width=img_size[0],
@@ -921,21 +903,26 @@ def _apply_hf_processor_main(
         ]
         processed_data["num_img_tokens"] = num_img_tokens
 
-        audio_features = processed_data["input_audio_embeds"]
+        audios = hf_data.get("audios", [])
+        assert isinstance(audios, list)
+        audio_features = processed_data.get("input_audio_embeds")
+        if audio_features is None:
+            return processed_data
+
+        sr = self.info.get_feature_extractor(**hf_kwargs).sampling_rate
         feature_sizes = [
-            self.info.get_audio_num_frames(len(audio), sr) for audio in audio_data
+            self.info.get_audio_num_frames(len(audio), sr) for audio, _ in audios
         ]
         processed_data["input_audio_embeds"] = [
             audio_features[idx, :size] for idx, size in enumerate(feature_sizes)
         ]
-        processed_data.update(passthrough_data)
 
         return processed_data
 
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> Mapping[str, MultiModalFieldConfig]:
         return dict(
             input_image_embeds=MultiModalFieldConfig.batched("image"),
@@ -1046,9 +1033,7 @@ def _recompute_cached_prompt_update(
     dummy_inputs=Phi4MMDummyInputsBuilder,
 )
 class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
-    """
-    Implements the Phi-4-multimodal-instruct model in vLLM.
-    """
+    """Implements the Phi-4-multimodal-instruct model in vLLM."""
 
     packed_modules_mapping = {
         "qkv_proj": [
@@ -1135,8 +1120,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     def _parse_and_validate_audio_input(
         self, **kwargs: object
     ) -> Phi4MMAudioInputs | None:
-        """
-        Parse and validate the audio input to the model.  This handles both
+        """Parse and validate the audio input to the model.  This handles both
         audio features and audio embeddings, but only the former is used for
         now.
 
@@ -1145,6 +1129,7 @@ def _parse_and_validate_audio_input(
 
         Returns:
             Optional[Phi4MMAudioInputs]: Parsed and validated audio inputs.
+
         """
         audio_features = kwargs.pop("input_audio_embeds", None)
         audio_embeds = kwargs.pop("audio_embeds", None)
@@ -1166,16 +1151,18 @@ def _parse_and_validate_audio_input(
     def _process_audio_input(
         self, audio_input: Phi4MMAudioInputs, audio_projection_mode: str
     ) -> NestedTensors:
-        """
-        Create the audio embeddings from the audio input, where the audio input
+        """Create the audio embeddings from the audio input, where the audio input
         is pairs of audio features and audio embed lengths.  The audio input is
         created by `input_mapper_for_phi4mm_audio`.
 
         Args:
             audio_input (Phi4MMAudioInputs): Audio input.
+            audio_projection_mode (str): Projection mode to use, selecting
+                which audio projector weights are applied.
 
         Returns:
             NestedTensors: Audio embeddings
+
         """
         if audio_input["type"] == "audio_embeds":
             return audio_input["data"]
@@ -1315,9 +1302,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="model.",
             connector=["audio_projection_for_vision", "audio_projection"],
diff --git a/vllm/model_executor/models/phi4mm_audio.py b/vllm/model_executor/models/phi4mm_audio.py
index 84abee33c0b9..2a16b4ca7522 100644
--- a/vllm/model_executor/models/phi4mm_audio.py
+++ b/vllm/model_executor/models/phi4mm_audio.py
@@ -18,7 +18,7 @@
     CheckpointWrapper,
 )
 from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.model_executor.models.phi4mm_utils import (
     AbsolutePositionalEncoding,
@@ -136,6 +136,7 @@ class ConformerEncoderLayer(nn.Module):
             1 = typical Multi-Head Attention,
             1 < attn_group_sizes < attention_heads = Grouped-Query Attention
             attn_group_sizes = attention_heads = Multi-Query Attention
+
     """
 
     def __init__(
@@ -232,6 +233,7 @@ def forward(
             mask: mask for x (batch, max_time_in)
             relative_attention_bias: bias added to attention logits w.r.t.
                 relative positions (1, n_head, time1, time2)
+
         """
         x = x + 0.5 * self.feed_forward_in(x)
         norm_x = self.layer_norm_att(x)
@@ -254,7 +256,7 @@ def forward(
 
 
 class TransformerEncoderBase(abc.ABC, nn.Module):
-    """The Base class for Transformer based encoders
+    """The Base class for Transformer based encoders.
 
     Please set causal = True in streaming model
     Args:
@@ -459,7 +461,6 @@ def _chunk_size_selection(
         left_chunk: int | list[int] | None = None,
     ) -> tuple[int, int]:
         """If chunk size is a list, we will randomly select a chunk size."""
-
         if chunk_size is None:
             chunk_size = self.chunk_size
         if left_chunk is None:
@@ -595,7 +596,7 @@ def forward_embeddings(
             torch.Tensor,
         ]
     ):
-        """Forwarding the inputs through the top embedding layers
+        """Forwarding the inputs through the top embedding layers.
 
         Args:
             xs_pad: torch.Tensor
@@ -606,6 +607,7 @@ def forward_embeddings(
                             non-causal layers
             left_chunk_nc: (optional, default is None) # of left chunks for
                             non-causal layers
+
         """
         # pylint: disable=R0915
         # get new lens.
@@ -975,13 +977,14 @@ def calculate_hs_mask(
     def forward(
         self, xs_pad: torch.Tensor, masks: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Conformer Forward function
+        """Conformer Forward function.
 
         Args:
             xs_pad: torch.Tensor
                 input tensor
             masks: torch.Tensor
                 post-embedding input lengths
+
         """
         xs_pad = self.encoder_embedding(xs_pad)
         input_tensor, pos_k, pos_v, hs_mask, masks = self.forward_embeddings(
@@ -1065,7 +1068,7 @@ def forward(
 
 
 class WindowQformer(nn.Module):
-    """Window-level Qformer"""
+    """Window-level Qformer."""
 
     def __init__(
         self,
@@ -1107,7 +1110,7 @@ def forward(
         mask: torch.Tensor | None,
         embed_len: int | None = None,
     ) -> tuple[torch.Tensor, int | None]:
-        """forward decoder"""
+        """Forward decoder."""
         # audio_embed: N x T x D => N x D x T
 
         audio_embed = audio_embed.transpose(1, 2)
@@ -1149,7 +1152,7 @@ def forward(
 class AudioEmbedding(nn.Module):
     """Image embedding."""
 
-    def __init__(self, config: PretrainedConfig, **kwargs: Any) -> None:
+    def __init__(self, config: PreTrainedConfig, **kwargs: Any) -> None:
         super().__init__()
         self.config = config
         # n_embed or hidden_size for text LM
@@ -1265,9 +1268,9 @@ def get_audio_features(
         audio_attention_mask: torch.Tensor | None = None,
         audio_projection_mode: str = "speech",
     ) -> torch.Tensor:
-        """
-        arguments:
-            input_embeds: audio features (B, T, D)  B: num audios in a sequence
+        """arguments:
+        input_embeds: audio features (B, T, D)  B: num audios in a sequence
+
         """
         if self.freeze_audio_processor:
             with torch.no_grad():
@@ -1319,12 +1322,12 @@ def forward(
         audio_attention_mask: torch.Tensor | None = None,
         audio_projection_mode: str = "speech",
     ) -> torch.Tensor:
-        """
-        arguments:
+        """arguments:
             audio_features: audio features (T, D)
 
-        returns:
+        Returns:
             audio_embeds: audio embeddings (num_audio_tokens, hidden_dim)
+
         """
         audio_embeds = self.get_audio_features(
             audio_features.unsqueeze(0),
diff --git a/vllm/model_executor/models/phi4mm_utils.py b/vllm/model_executor/models/phi4mm_utils.py
index 08409af4e2f7..af54defd3fb0 100644
--- a/vllm/model_executor/models/phi4mm_utils.py
+++ b/vllm/model_executor/models/phi4mm_utils.py
@@ -13,7 +13,7 @@
 
 
 class BlockBase(nn.Module):
-    """Block abstract module"""
+    """Block abstract module."""
 
     def __init__(self, input_size: int, output_size: int) -> None:
         super().__init__()
@@ -22,13 +22,14 @@ def __init__(self, input_size: int, output_size: int) -> None:
 
 
 def get_activation(name: str = "relu") -> torch.nn.Module:
-    """Select an activation function by name
+    """Select an activation function by name.
 
     Args:
         name: str
             activation function name,
             one of ["relu", "gelu", "swish", "sigmoid"],
             default "relu".
+
     """
     name = name.lower()
     if name == "relu":
@@ -48,8 +49,7 @@ def get_activation(name: str = "relu") -> torch.nn.Module:
 def adaptive_enc_mask(
     x_len: int, chunk_start_idx: list[int], left_window: int = 0, right_window: int = 0
 ) -> torch.Tensor:
-    """
-    The function is very important for Transformer Transducer Streaming mode
+    """The function is very important for Transformer Transducer Streaming mode
     Args:
         x_len: sequence length
         chunk_start_idx: first idx of each chunk, such as [0,18,36,48].
@@ -57,7 +57,8 @@ def adaptive_enc_mask(
         left_window: how many left chunks can be seen
         right_window: how many right chunks can be seen. It is used for
         chunk overlap model.
-        Returns:
+
+    Returns:
             mask (torch.Tensor): a mask tensor for streaming model
             Torch 1.0.1
             tensor([[1., 1., 0., 0.],
@@ -67,6 +68,7 @@ def adaptive_enc_mask(
             tensor([[True., True., False., False.],
                     [False., True., True., False.],
                     [False., False., True., True.]])
+
     """
     chunk_start_idx = torch.Tensor(
         chunk_start_idx
@@ -97,7 +99,7 @@ def adaptive_enc_mask(
 
 
 class GLU(nn.Module):
-    """Implement Gated Linear Unit (GLU) module"""
+    """Implement Gated Linear Unit (GLU) module."""
 
     def __init__(self, dim: int = -1, act_name: str = "sigmoid") -> None:
         super().__init__()
@@ -184,9 +186,9 @@ def __init__(
             self.b2 = nn.Parameter(torch.zeros(1, output_dim, 1))
 
     def forward(self, x: Tensor) -> Tensor:
-        """
-        Args:
-            x: input tensor
+        """Args:
+        x: input tensor
+
         """
         # to be consistent with GLULinear, we assume the input always has the
         # #channel (#dim) in the last dimension of the tensor, so need to
@@ -273,10 +275,9 @@ def __init__(
         self.depthwise_seperable_out_channel = depthwise_seperable_out_channel
 
     def forward(self, x: Tensor) -> Tensor:
-        """
+        """Args:
+        x: input tensor
 
-        Args:
-            x: input tensor
         """
         x = self.dw_conv(x)
         if self.depthwise_seperable_out_channel != 0:
@@ -341,6 +342,7 @@ class ConvModule(nn.Module):
              or onnx export.  Typically this is set by the export program or
              the decoder program, and it isn't present in your config file.
              default False
+
     """
 
     def __init__(
@@ -406,8 +408,7 @@ def __init__(
                 self.ln2 = nn.Linear(input_dim * depthwise_multiplier, input_dim)
 
     def _add_ext_pw_layer(self) -> None:
-        """
-        This function is an extension of __init__ function
+        """This function is an extension of __init__ function
         and dedicated to the convolution module creation
         of the conformer.
         """
@@ -471,6 +472,7 @@ def forward(self, x: Tensor) -> Tensor:
 
         Args:
             x: input tensor.
+
         """
         x = self.layer_norm(x)
 
@@ -519,7 +521,7 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class GLULinear(nn.Module):
-    """Linear + GLU module
+    """Linear + GLU module.
 
     Args:
         input_dim: int
@@ -531,6 +533,7 @@ class GLULinear(nn.Module):
             default "sigmoid" (swish function).
         bias_in_glu: bool, optional
             If True, the addtive bias is added. Default False.
+
     """
 
     def __init__(
@@ -545,10 +548,11 @@ def __init__(
         self.glu_act = GLU(-1, glu_type)
 
     def forward(self, x: Tensor) -> Tensor:
-        """GLULinear forward
+        """GLULinear forward.
 
         Args:
             x: input tensor.
+
         """
         x = self.linear(x)
         return self.glu_act(x)
@@ -572,6 +576,7 @@ class FeedForward(nn.Module):
             sigmoid activation is only used with "glu_in_fnn=True",
             default "sigmoid".
         bias_in_glu: bool, optional
+
     """
 
     def __init__(
@@ -600,6 +605,7 @@ def forward(self, x: Tensor) -> Tensor:
 
         Args:
             x: input tensor.
+
         """
         out = self.net(self.layer_norm(x))
 
@@ -630,8 +636,7 @@ def _pre_hook(
 
 
 class T5RelativeAttentionLogitBias(nn.Module):
-    """
-    This module implements the relative position bias described in Section
+    """This module implements the relative position bias described in Section
     2.1 of the T5 paper: https://arxiv.org/pdf/1910.10683.pdf
 
     The Huggingface implementation is used as a reference
@@ -671,6 +676,7 @@ class T5RelativeAttentionLogitBias(nn.Module):
             Whether to use symmetric or asymmetric biases. symmetric=False uses
             2x number of bias params to distinguish L->R from R->L. This was
             found to be better for the encoder.
+
     """
 
     def __init__(
@@ -799,6 +805,7 @@ def extend_pe(self, x: torch.Tensor) -> None:
 
         Args:
             x: input tensor
+
         """
         if self.pe is not None and self.pe.size(1) >= x.size(1):
             if self.pe.dtype != x.dtype or self.pe.device != x.device:
@@ -841,6 +848,7 @@ class MeanVarianceNormLayer(nn.Module):
     Args:
         input_size: int
             layer input size.
+
     """
 
     def __init__(self, input_size: int) -> None:
@@ -850,17 +858,17 @@ def __init__(self, input_size: int) -> None:
         self.global_invstd = nn.Parameter(torch.ones(input_size))
 
     def forward(self, input_: Tensor) -> Tensor:
-        """MeanVarianceNormLayer Forward
+        """MeanVarianceNormLayer Forward.
 
         Args:
             input_: input tensor.
+
         """
         return (input_ - self.global_mean) * self.global_invstd
 
 
 class CausalConv1D(nn.Conv1d):
-    """
-    A causal version of nn.Conv1d where each step would have limited access to
+    """A causal version of nn.Conv1d where each step would have limited access to
     locations on its right or left
     All arguments are the same as nn.Conv1d except padding.
 
@@ -956,8 +964,7 @@ def forward(
 
 
 class CausalConv2D(nn.Conv2d):
-    """
-    A causal version of nn.Conv2d where each location in the 2D matrix would
+    """A causal version of nn.Conv2d where each location in the 2D matrix would
     have no access to locations on its right or down
     All arguments are the same as nn.Conv2d except padding which should be
     set as None
@@ -1041,6 +1048,7 @@ class NemoConvSubsampling(torch.nn.Module):
         activation (Module): activation function, default is nn.ReLU()
         is_causal (bool): whether to use causal Conv1/2D, where each step will
             have limited access to locations on its right or left
+
     """
 
     def __init__(
@@ -1338,8 +1346,7 @@ def get_streaming_cache_size(self) -> list[int]:
         return [0, self.subsampling_factor + 1]
 
     def forward(self, x: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor | None]:
-        """
-        Forward method for NeMo subsampling.
+        """Forward method for NeMo subsampling.
 
         Args:
             x: input tensor
@@ -1350,6 +1357,7 @@ def forward(self, x: Tensor, mask: Tensor | None) -> tuple[Tensor, Tensor | None
                 time_reduction_factor, feat_out)
             pad_mask: tensor of padded hidden state sequences (B, 1, T //
                 time_reduction_factor)
+
         """
         x = x.unsqueeze(1) if self.conv2d_subsampling else x.transpose(1, 2)
 
@@ -1427,7 +1435,7 @@ def reset_parameters(self) -> None:
                 torch.nn.init.uniform_(self.out.bias, -fc_scale, fc_scale)
 
     def conv_split_by_batch(self, x: Tensor) -> tuple[Tensor, bool]:
-        """Tries to split input by batch, run conv and concat results"""
+        """Tries to split input by batch, run conv and concat results."""
         b, _, _, _ = x.size()
         if b == 1:  # can't split if batch size is 1
             return x, False
@@ -1493,8 +1501,7 @@ def conv_split_by_channel(self, x: Tensor) -> Tensor:
     def channel_chunked_conv(
         self, conv: torch.nn.Module, chunk_size: int, x: Tensor
     ) -> Tensor:
-        """Performs channel chunked convolution"""
-
+        """Performs channel chunked convolution."""
         ind = 0
         out_chunks = []
         for chunk in torch.split(x, chunk_size, 1):
@@ -1571,14 +1578,14 @@ def calc_length_int(
 
 ####  multihead attention starts here
 class AttModule(nn.Module):
-    """Attention abstraction module"""
+    """Attention abstraction module."""
 
     def __init__(self) -> None:
         super().__init__()
         self.export_mode = False
 
     def set_export(self, mode: bool = True) -> None:
-        """set the export mode"""
+        """Set the export mode."""
         self.export_mode = mode
 
     def forward(
@@ -1588,13 +1595,14 @@ def forward(
         pos_emb: Tensor | None = None,
         att_mask: Tensor | None = None,
     ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:
-        """AttModule forward
+        """AttModule forward.
 
         Args:
             x: input tensor.
             memory: memory tensor.
             pos_emb: positional encoder embedding.
             att_mask: attention mask tensor.
+
         """
         return x, memory, pos_emb, att_mask
 
@@ -1642,6 +1650,7 @@ class MultiHeadedAttention(nn.Module):
             if group_size > 1:       GQA
             if group_size = 1:       MHA
             if group_size = n_head:  MQA
+
     """
 
     inv_sqrt_d_k: torch.jit.Final[float]
@@ -1719,6 +1728,7 @@ def forward(
             relative_attention_bias: bias added to attention logits w.r.t.
                 relative positions
                 (1, n_head, time1, time2)
+
         """
         n_batch = query.size(0)
 
@@ -1815,7 +1825,7 @@ def forward(
 
 
 class MultiSequential(torch.nn.Sequential):
-    """Multi-input multi-output torch.nn.Sequential"""
+    """Multi-input multi-output torch.nn.Sequential."""
 
     @torch.jit.ignore
     def forward(self, *args) -> tuple:
@@ -1834,6 +1844,7 @@ def get_offset(input_layer: str, time_reduction: int) -> int:
         time_reduction: time reduction factor for downsampling a feature
     Returns:
         int: offset
+
     """
     if input_layer in ("conv2d", "nemo_conv") and time_reduction == 4:
         return 3
@@ -1845,13 +1856,14 @@ def get_offset(input_layer: str, time_reduction: int) -> int:
 
 
 def unfold_tensor(xs_pad: Tensor, max_seq_len: int) -> Tensor:
-    """
-    For a given tensor with shape of (N, T, D), if sequence length T is
+    """For a given tensor with shape of (N, T, D), if sequence length T is
     longer than max_seq_len, this function unfold it to a
     (NT', max_seq_len, D) where T' is T // max_seq_len.
+
     Args:
         xs_pad: input tensor with shape (N, T, D)
         max_seq_len: maximum sequence length
+
     """
     _, _, D = xs_pad.shape
     xs_pad = xs_pad.transpose(-1, -2)  # convert to N, D, T
diff --git a/vllm/model_executor/models/phi4siglip.py b/vllm/model_executor/models/phi4siglip.py
index ff880388698c..e9ad4a795de7 100644
--- a/vllm/model_executor/models/phi4siglip.py
+++ b/vllm/model_executor/models/phi4siglip.py
@@ -11,10 +11,10 @@
 
 import torch
 import torch.nn as nn
-from transformers import BatchFeature, PretrainedConfig, Siglip2VisionConfig
+from transformers import BatchFeature, PreTrainedConfig, Siglip2VisionConfig
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -126,15 +126,14 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
         size = self.info.get_image_size_with_most_features()
         return {
             "image": self._get_dummy_images(
                 width=size.width,
                 height=size.height,
-                num_images=num_images,
+                num_images=mm_counts.get("image", 0),
                 overrides=mm_options.get("image"),
             ),
         }
@@ -143,39 +142,8 @@ def get_dummy_mm_data(
 class Phi4SiglipMultiModalProcessor(
     BaseMultiModalProcessor[Phi4SiglipProcessingInfo],
 ):
-    def _apply_hf_processor_main(
-        self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
-
-        processed = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            dict(text=prompt_text, **mm_data),
-            hf_processor_mm_kwargs,
-        )
-
-        # The HF processor's tokenizer_image_token() replaces the "<image>"
-        # string with IMAGE_TOKEN_INDEX (-200) in input_ids.  This breaks
-        # vLLM's prompt-replacement pipeline which needs to find "<image>"
-        # as normal sub-tokens.  Re-tokenize with the plain tokenizer so
-        # that "<image>" stays as sub-tokens and can be located by
-        # PromptReplacement.
-        # NOTE: tokenizer.__call__() (not .encode()) must be used so that
-        # added/special tokens like <|user|>, <|end|> are kept as single IDs.
-        tokenizer = self.info.get_tokenizer()
-        new_ids = tokenizer(prompt_text).input_ids
-        processed["input_ids"] = torch.tensor([new_ids])
-
-        processed_data = processed
-        processed_data.update(passthrough_data)
-        return processed_data
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _get_mm_fields_config(
         self,
@@ -224,11 +192,10 @@ def get_replacement(item_idx: int):
 
 
 class Phi4SiglipImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - d: Max number of patches (padded across images in the batch)
-        - fd: Features per patch (patch_size * patch_size * channels)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - d: Max number of patches (padded across images in the batch)
+    - fd: Features per patch (patch_size * patch_size * channels)
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -268,7 +235,7 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         super().__init__()
 
-        config: PretrainedConfig = vllm_config.model_config.hf_config
+        config: PreTrainedConfig = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
         self.config = config
 
diff --git a/vllm/model_executor/models/phimoe.py b/vllm/model_executor/models/phimoe.py
index b3b7b995ee3e..2cf41a701f08 100644
--- a/vllm/model_executor/models/phimoe.py
+++ b/vllm/model_executor/models/phimoe.py
@@ -29,16 +29,18 @@
 
 import torch
 from torch import nn
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -60,7 +62,7 @@
 )
 
 
-class PhiMoEConfig(PretrainedConfig):
+class PhiMoEConfig(PreTrainedConfig):
     model_type = "phimoe"
     keys_to_ignore_at_inference = ["past_key_values"]
 
@@ -262,12 +264,10 @@ def __init__(
         self.hidden_size = hidden_size
 
         # Gate always runs at half / full precision for now.
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             hidden_size,
             num_experts,
-            bias=False,
             params_dtype=params_dtype,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py
index b6848b09ded0..91ac38e11087 100644
--- a/vllm/model_executor/models/pixtral.py
+++ b/vllm/model_executor/models/pixtral.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import importlib.metadata
 import math
 from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass, fields
@@ -12,18 +13,15 @@
 from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk
 from mistral_common.protocol.instruct.messages import UserMessage
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
-from transformers import BatchFeature, PixtralVisionConfig
+from packaging.version import Version
+from transformers import PixtralVisionConfig
 from transformers.models.pixtral.image_processing_pixtral import (
     _num_image_tokens as _get_pixtral_hf_num_image_tokens,
 )
-from transformers.models.pixtral.modeling_pixtral import (
-    PixtralRotaryEmbedding,
-    apply_rotary_pos_emb,
-    position_ids_in_meshgrid,
-)
+from transformers.models.pixtral.modeling_pixtral import apply_rotary_pos_emb
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import divide, get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn
@@ -42,7 +40,7 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
-    MultiModalKwargsOptionalItems,
+    MultiModalKwargsItem,
     NestedTensors,
 )
 from vllm.multimodal.parse import (
@@ -54,7 +52,8 @@
 from vllm.multimodal.processing.processor import (
     BaseMultiModalProcessor,
     BaseProcessingInfo,
-    MultiModalPromptUpdates,
+    HFMultiModalInputs,
+    MultiModalProcessingResult,
     PlaceholderFeaturesInfo,
     ProcessorInputs,
     PromptReplacement,
@@ -68,7 +67,6 @@
     MistralCommonImageProcessor,
     MistralCommonPixtralProcessor,
 )
-from vllm.utils.collection_utils import is_list_of
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
@@ -91,6 +89,39 @@
 
 PATCH_MERGE = "patch_merge"
 
+# Transformers 5.17 renamed Pixtral's rotary embedding and switched it to axial
+# RoPE, which takes 2D (height, width) positions instead of flattened grid ids.
+TRANSFORMERS_VERSION = importlib.metadata.version("transformers")
+TRANSFORMERS_WITH_AXIAL_ROPE = Version(TRANSFORMERS_VERSION) >= Version("5.17.0.dev0")
+
+if TRANSFORMERS_WITH_AXIAL_ROPE:
+    from transformers.models.pixtral.modeling_pixtral import (
+        PixtralVisionRotaryEmbedding,
+    )
+else:
+    from transformers.models.pixtral.modeling_pixtral import (
+        PixtralRotaryEmbedding as PixtralVisionRotaryEmbedding,
+    )
+    from transformers.models.pixtral.modeling_pixtral import (
+        position_ids_in_meshgrid as flat_position_ids_in_meshgrid,
+    )
+
+
+def position_ids_in_meshgrid(
+    patch_embeds_list: list[torch.Tensor],
+    max_width: int,
+) -> torch.Tensor:
+    if not TRANSFORMERS_WITH_AXIAL_ROPE:
+        return flat_position_ids_in_meshgrid(patch_embeds_list, max_width)
+    positions = []
+    for patch in patch_embeds_list:
+        height, width = patch.shape[-2:]
+        h_ids, w_ids = torch.meshgrid(
+            torch.arange(height), torch.arange(width), indexing="ij"
+        )
+        positions.append(torch.stack([h_ids.flatten(), w_ids.flatten()], dim=-1))
+    return torch.cat(positions, dim=0)
+
 
 def _make_packed_sequence_metadata(
     sequence_lengths: list[int],
@@ -125,8 +156,7 @@ def _is_layer_none_or_staged(layer: nn.Module) -> bool:
 
 
 class PixtralImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - bn: Batch size * number of images
         - c: Number of channels (3)
         - h: Height of each image
@@ -178,35 +208,35 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
-    def get_dummy_processor_inputs(
+
+class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]):
+    def get_dummy_inputs(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
+        # For test_common.py only
         mm_data: MultiModalDataDict | None = None,
     ) -> ProcessorInputs:
+        builder = self.dummy_inputs
         tokenizer = self.info.get_tokenizer()
 
-        dummy_text = self.get_dummy_text(mm_counts)
+        dummy_text = builder.get_dummy_text(mm_counts)
         dummy_mm_data = (
-            self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
+            builder.get_dummy_mm_data(seq_len, mm_counts, mm_options)
             if mm_data is None
             else mm_data
         )
@@ -230,28 +260,24 @@ def get_dummy_processor_inputs(
 
         return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items)
 
-
-class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]):
     # The tokens are already inserted by the chat template,
     # so we just double check that they exist
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsOptionalItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_res: MultiModalProcessingResult,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
         mm_item_counts = mm_items.get_all_counts()
-        self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
-        self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
+        self._validate_mm_kwargs(mm_res.kwargs, mm_item_counts)
+        self._validate_mm_updates(mm_res.prompt_updates, mm_item_counts)
 
         mm_placeholders = self._find_mm_placeholders(
-            prompt_ids,
-            mm_prompt_updates,
+            mm_res.prompt_ids,
+            mm_res.prompt_updates,
         )
         self._validate_mm_placeholders(mm_placeholders, mm_item_counts)
 
-        return prompt_ids, mm_placeholders
+        return mm_res.prompt_ids, mm_placeholders
 
     def _get_mm_fields_config(
         self,
@@ -260,36 +286,18 @@ def _get_mm_fields_config(
     ) -> Mapping[str, MultiModalFieldConfig]:
         return dict(images=MultiModalFieldConfig.batched("image"))
 
-    def _apply_hf_processor_main(
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
-
-        outputs = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            dict(text=prompt_text, **mm_data),
-            # Avoid padding issue
-            dict(**hf_processor_mm_kwargs, return_tensors=None),
+        # Avoid padding issue
+        return hf_inputs._replace(
+            hf_kwargs=dict(hf_inputs.hf_kwargs, return_tensors=None)
         )
 
-        # Missing batch dimension
-        if is_list_of(outputs["input_ids"], int):
-            outputs["input_ids"] = [outputs["input_ids"]]
-
-        processed_data = outputs
-        processed_data.update(passthrough_data)
-        return processed_data
-
     def _get_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
@@ -623,17 +631,18 @@ def get_mm_mapping(self) -> MultiModelKeys:
             tower_model="vision_encoder",
         )
 
-    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
-        if getattr(self, "patch_merger", None) is None:
-            return num_image_tokens
-        merge_size = self.vision_args.spatial_merge_size
-        return num_image_tokens * (merge_size**2)
-
-    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         if getattr(self, "patch_merger", None) is None:
-            return num_vision_tokens
+            return num_mm_embeds, num_mm_embeds
         merge_size = self.vision_args.spatial_merge_size
-        return num_vision_tokens // (merge_size**2)
+        return num_mm_embeds * (merge_size**2), num_mm_embeds
 
 
 # Vision encoder
@@ -655,8 +664,7 @@ class VisionEncoderArgs:
 
 
 def _reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
-    """
-    freqs_cis: complex - (seq_len, head_dim / 2)
+    """freqs_cis: complex - (seq_len, head_dim / 2)
     x: complex - (bsz, seq_len, head_dim / 2)
     """
     ndim = x.ndim
@@ -675,9 +683,8 @@ def precompute_freqs_cis_2d(
     width: int,
     theta: float,
 ) -> torch.Tensor:
-    """
-    freqs_cis: 2D complex tensor of shape (height, width, dim // 2)
-        to be indexed by (height, width) position tuples
+    """freqs_cis: 2D complex tensor of shape (height, width, dim // 2)
+    to be indexed by (height, width) position tuples
     """
     # (dim / 2) frequency bases
     freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
@@ -985,13 +992,14 @@ def forward(
         self,
         images: list[torch.Tensor],
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             images: list of N_img images of variable sizes,
                 each of shape (C, H, W)
+
         Returns:
             image_features: tensor of token features for
                 all tokens of all images of shape (N_toks, D)
+
         """
         # pass images through initial convolution independently
         patch_embeds_list = [
@@ -1049,9 +1057,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class PatchMerger(nn.Module):
-    """
-    Learned merging of spatial_merge_size ** 2 patches
-    """
+    """Learned merging of spatial_merge_size ** 2 patches."""
 
     def __init__(
         self,
@@ -1091,8 +1097,7 @@ def permute(
         x: torch.Tensor,
         image_sizes: list[tuple[int, int]],
     ) -> torch.Tensor:
-        """
-        Args:
+        """Args:
             x: (N, D) where N is flattened and concatenated patch tokens
                 for all images
             image_sizes: list of tuple of (height, width) in tokens for
@@ -1101,8 +1106,8 @@ def permute(
             image_features: reorders patch tokens so each grid of
                 (spatial_merge_size, spatial_merge_size) is contiguous.
                 now (N / spatial_merge_size ** 2, D * spatial_merge_size ** 2)
-        """
 
+        """
         sub_grids = get_sub_grids(
             x=x, image_sizes=image_sizes, spatial_merge_size=self.spatial_merge_size
         )  # list of [d x sub_grid_size x sub_grid_size x n_patches]
@@ -1465,7 +1470,9 @@ def __init__(
 
         self.dtype = next(self.parameters()).dtype
         self.device = next(self.parameters()).device
-        self.patch_positional_embedding = PixtralRotaryEmbedding(config, self.device)
+        self.patch_positional_embedding = PixtralVisionRotaryEmbedding(config).to(
+            self.device
+        )
 
     def forward(
         self,
@@ -1474,8 +1481,7 @@ def forward(
         select_layers: list[int] | None = None,
         feature_select_strategy: VisionFeatureSelectStrategy | None = None,
     ) -> tuple[torch.Tensor, ...]:
-        """
-        Args:
+        """Args:
             pixel_values: Each image to be processed will be a separate tensor
                 in pixel_values. This means it will be a list of tensors
                 because multiple requests batched can have multiple images,
@@ -1487,6 +1493,7 @@ def forward(
         Returns:
             image_features: tensor of token features for
                 all tokens of all images of shape (N_toks, D)
+
         """
         # pass images through initial convolution independently
         patch_embeds_list = [
diff --git a/vllm/model_executor/models/plamo3.py b/vllm/model_executor/models/plamo3.py
index 8d870abe2f02..d30a257574e4 100644
--- a/vllm/model_executor/models/plamo3.py
+++ b/vllm/model_executor/models/plamo3.py
@@ -8,7 +8,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
@@ -54,7 +54,7 @@
 # Only used for type hinting.
 if TYPE_CHECKING:
 
-    class Plamo3Config(PretrainedConfig):  # type: ignore
+    class Plamo3Config(PreTrainedConfig):  # type: ignore
         model_type: str = "plamo3"
 
         hidden_size: int
diff --git a/vllm/model_executor/models/qianfan_ocr.py b/vllm/model_executor/models/qianfan_ocr.py
index ef2bec1e2900..8367737c9597 100644
--- a/vllm/model_executor/models/qianfan_ocr.py
+++ b/vllm/model_executor/models/qianfan_ocr.py
@@ -5,7 +5,7 @@
 # The model architecture and weights are fully compatible with InternVLChatModel,
 # only the config model_type / architectures strings differ.
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
@@ -67,7 +67,7 @@ class QianfanOCRForConditionalGeneration(InternVLChatModel):
     """
 
     def _patch_quant_config(
-        self, config: PretrainedConfig, quant_config: QuantizationConfig
+        self, config: PreTrainedConfig, quant_config: QuantizationConfig
     ) -> None:
         super()._patch_quant_config(config, quant_config)
         # ignore vit layers to preserve model performance
diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py
index 26b69bd6dd66..0ce2c37313ed 100644
--- a/vllm/model_executor/models/qwen2_5_omni_thinker.py
+++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py
@@ -29,7 +29,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniConfig,
@@ -44,7 +44,7 @@
 from transformers.models.whisper import WhisperFeatureExtractor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.forward_context import set_forward_context
 from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.logger import init_logger
@@ -82,6 +82,8 @@
 )
 from vllm.multimodal.processing.processor import (
     BaseMultiModalProcessor,
+    HFMultiModalInputs,
+    MultiModalProcessingResult,
     MultiModalPromptUpdates,
     PlaceholderFeaturesInfo,
     PromptReplacement,
@@ -139,8 +141,7 @@ def check_interleaved_audio_video(
     num_video: int,
     num_audio: int,
 ) -> bool:
-    """
-    Check if video and audio positions are interleaved in any per-video span.
+    """Check if video and audio positions are interleaved in any per-video span.
 
     For use_audio_in_video=True, each video placeholder is expanded into one
     local span containing only video/audio pad tokens, bounded by non-pad
@@ -187,8 +188,7 @@ def merge_interleaved_embeddings(
     is_audio: torch.Tensor,
     is_multimodal: torch.Tensor,
 ) -> torch.Tensor:
-    """
-    Merge embeddings for interleaved audio-in-video sequences.
+    """Merge embeddings for interleaved audio-in-video sequences.
 
     When use_audio_in_video=True, video and audio tokens are interleaved in
     the token sequence, but embeddings are provided as separate contiguous
@@ -206,6 +206,7 @@ def merge_interleaved_embeddings(
     Returns:
         The merged inputs_embeds tensor with multimodal embeddings scattered
         to their correct positions.
+
     """
     from vllm.multimodal.utils import get_mm_embedding_modalities
 
@@ -254,12 +255,11 @@ def _merge_embedding_group(
 
 
 class Qwen2_5OmniAudioFeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - na: Number of audios
-        - nmb: Number of mel bins
-        - msl: Maximum sequence length
-        - tsl: Total sequence length
+    """Dimensions:
+    - na: Number of audios
+    - nmb: Number of mel bins
+    - msl: Maximum sequence length
+    - tsl: Total sequence length
     """
 
     type: Literal["audio_features"]
@@ -282,7 +282,7 @@ def create_qwen2_5_omni_thinker_field_factory(
     def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
         audio_feature_lengths = hf_inputs.get(
             "audio_feature_lengths", torch.empty((0,))
-        )
+        ).flatten()
 
         image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
         image_pixel_grid_sizes = image_grid_thw.prod(-1)
@@ -332,6 +332,14 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
 
 
 class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
+    embedding_fields = {
+        **Qwen2VLMultiModalDataParser.embedding_fields,
+        "audio": {
+            "input_audio_features": "values",
+            "audio_feature_lengths": "metadata",
+        },
+    }
+
     def __init__(self, spatial_merge_size: int, *args, **kwargs):
         self._spatial_merge_size = spatial_merge_size
         super().__init__(self._spatial_merge_size, *args, **kwargs)
@@ -341,10 +349,15 @@ def _parse_audio_data(
         data: dict[str, torch.Tensor] | ModalityData[ImageItem],
     ) -> ModalityDataItems[Any, Any]:
         if isinstance(data, dict):
+            required, optional = self.embedding_field_sets("audio")
+            data = dict(data)
+            if "audio_feature_lengths" in data:
+                data["audio_feature_lengths"] = data["audio_feature_lengths"].flatten()
             return DictEmbeddingItems(
                 data,
                 modality="audio",
-                required_fields={"input_audio_features", "audio_feature_lengths"},
+                required_fields=required,
+                optional_fields=optional,
                 fields_factory=create_qwen2_5_omni_thinker_field_factory(
                     self._spatial_merge_size
                 ),
@@ -362,7 +375,6 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs: object) -> Qwen2_5OmniProcessor:
         return self.ctx.get_hf_processor(
             Qwen2_5OmniProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
@@ -448,12 +460,8 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         feature_extractor = self.info.get_feature_extractor()
 
         target_audio_length = (
@@ -469,28 +477,24 @@ def get_dummy_mm_data(
             seq_len, mm_counts
         )
 
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-        audio_overrides = mm_options.get("audio")
-
         mm_data = {
             "audio": self._get_dummy_audios(
                 length=target_audio_length,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             ),
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -500,6 +504,64 @@ def get_dummy_mm_data(
 class Qwen2_5OmniThinkerMultiModalProcessor(
     BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]
 ):
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
+    def _get_hf_mm_inputs(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        mm_counts = mm_items.get_all_counts()
+        if hf_kwargs.get("use_audio_in_video"):
+            num_videos = mm_counts.get("video", 0)
+            if mm_counts.get("audio", 0) < num_videos:
+                raise ValueError(
+                    "Video doesn't have audio track with `audio_in_video=True`"
+                )
+
+            if num_videos and "text" in hf_inputs.hf_data:
+                standalone_mm_counts = dict(mm_counts)
+                standalone_mm_counts["audio"] -= num_videos
+                hf_inputs.hf_data["text"] = self._get_hf_mm_text(standalone_mm_counts)
+
+        hf_data = hf_inputs.hf_data
+        normalized_kwargs = hf_inputs.hf_kwargs
+        audios = hf_data.pop("audio", [])
+
+        # NOTE: WhisperFeatureExtractor cannot handle empty list of audios
+        if audios:
+            hf_data["audio"] = audios
+        elif (
+            normalized_kwargs.get("use_audio_in_video")
+            and "audio" not in hf_data
+            and hf_data.get("videos")
+        ):
+            raise ValueError(
+                "Video doesn't have audio track with `audio_in_video=True`"
+            )
+
+        merged = self.info.ctx.get_merged_mm_kwargs(normalized_kwargs)
+        if hf_data.get("videos") and (
+            merged.keys() & {"size", "min_pixels", "max_pixels"}
+        ):
+            normalized_kwargs = dict(normalized_kwargs)
+            video_size = dict(self.info.get_hf_processor().video_processor.size)
+            size_override = merged.get("size")
+            if size_override is not None:
+                video_size = video_size | size_override
+            min_pixels = merged.get("min_pixels")
+            if min_pixels is not None:
+                video_size["shortest_edge"] = min_pixels
+            max_pixels = merged.get("max_pixels")
+            if max_pixels is not None:
+                video_size["longest_edge"] = max_pixels
+            normalized_kwargs["size"] = video_size
+
+        return hf_inputs._replace(hf_kwargs=normalized_kwargs)
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
@@ -514,8 +576,7 @@ def _derive_audio_from_video_placeholders(
         placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
         mm_prompt_updates: MultiModalPromptUpdates,
     ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
-        """
-        Helper to derive audio placeholders from video placeholders when
+        """Helper to derive audio placeholders from video placeholders when
         use_audio_in_video=True.
         """
         if "video" not in placeholders:
@@ -573,14 +634,14 @@ def _derive_audio_from_video_placeholders(
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_res: MultiModalProcessingResult,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
-        """
-        Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
-        """
+        """Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`."""
         mm_item_counts = mm_items.get_all_counts()
+        prompt_ids = mm_res.prompt_ids
+        mm_kwargs = mm_res.kwargs
+        mm_prompt_updates = mm_res.prompt_updates
+
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
         self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
 
@@ -631,7 +692,7 @@ def _maybe_apply_prompt_updates(
     @classmethod
     def omni_get_updates_use_audio_in_video(
         cls,
-        thinker_config: PretrainedConfig,
+        thinker_config: PreTrainedConfig,
         audio_len: int,
         video_grid_thw: list[int] | torch.Tensor,
         video_second_per_grid_t: float,
@@ -644,7 +705,6 @@ def omni_get_updates_use_audio_in_video(
         <|video_bos|><|VIDEO|><|video_eos|> =>
         <|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|>
         """
-
         audio_token_id = thinker_config.audio_token_index
         video_token_id = thinker_config.video_token_index
         audio_start_token_id = thinker_config.audio_start_token_id
@@ -795,103 +855,33 @@ def get_replacement_qwen2_use_audio_in_video(item_idx: int):
             ),
         ]
 
-    def _apply_hf_processor_main(
+    def _postprocess_hf_mm_data(
         self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+        processed_data: BatchFeature,
     ) -> BatchFeature:
-        """
-        Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
-        """
-        mm_counts = mm_items.get_all_counts()
-
-        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
-        if use_audio_in_video and "video" in mm_counts:
-            if mm_counts.get("audio", 0) < mm_counts["video"]:
-                raise ValueError(
-                    "Video doesn't have audio track with `audio_in_video=True`"
-                )
-            mm_counts["audio"] -= mm_counts["video"]
-
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_counts)
-
-        mm_data = dict(processor_data)
-        audios = mm_data.pop("audios", [])
-
-        # NOTE: WhisperFeatureExtractor cannot handle empty list of audios
-        if audios:
-            # NOTE: Qwen2.5-Omni processor accept "audio"
-            mm_data["audio"] = audios
-            hf_processor_mm_kwargs = dict(
-                **hf_processor_mm_kwargs,
-            )
-        elif (
-            hf_processor_mm_kwargs.get("use_audio_in_video")
-            and "audio" not in mm_data
-            and mm_data.get("videos")
-        ):
-            # A subclass (e.g. Qwen3-Omni) may have already popped "audios"
-            # and populated mm_data["audio"] itself before delegating here,
-            # in which case this isn't a genuine "no audio" case. Likewise,
-            # mm_data can be empty on a multimodal-processor-cache hit, where
-            # there's nothing to (re-)process this call and the real result
-            # comes from the cache — not a genuine "no audio" case either.
-            raise ValueError(
-                "Video doesn't have audio track with `audio_in_video=True`"
-            )
-
-        merged = self.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs)
-        if mm_data.get("videos") and (
-            merged.keys() & {"size", "min_pixels", "max_pixels"}
-        ):
-            hf_processor_mm_kwargs = dict(hf_processor_mm_kwargs)
-            video_size = dict(self.info.get_hf_processor().video_processor.size)
-            size_override = merged.get("size")
-            if size_override is not None:
-                video_size = video_size | size_override
-            min_pixels = merged.get("min_pixels")
-            if min_pixels is not None:
-                video_size["shortest_edge"] = min_pixels
-            max_pixels = merged.get("max_pixels")
-            if max_pixels is not None:
-                video_size["longest_edge"] = max_pixels
-            hf_processor_mm_kwargs["size"] = video_size
-
-        hf_inputs = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            dict(text=prompt_text, **mm_data),
-            hf_processor_mm_kwargs,
-        )
-
-        input_features = hf_inputs.pop("input_features", None)
-        feature_attention_mask = hf_inputs.get("feature_attention_mask", None)
-        if "input_audio_features" not in hf_inputs and input_features is not None:
+        input_features = processed_data.pop("input_features", None)
+        feature_attention_mask = processed_data.get("feature_attention_mask", None)
+        if "input_audio_features" not in processed_data and input_features is not None:
             if feature_attention_mask is not None:
                 input_features = input_features.permute(0, 2, 1)[
                     feature_attention_mask.bool()
                 ].permute(1, 0)
-            hf_inputs["input_audio_features"] = input_features
+            processed_data["input_audio_features"] = input_features
         if (
-            "audio_feature_lengths" not in hf_inputs
+            "audio_feature_lengths" not in processed_data
             and feature_attention_mask is not None
         ):
-            hf_inputs["audio_feature_lengths"] = feature_attention_mask.sum(-1)
+            processed_data["audio_feature_lengths"] = feature_attention_mask.sum(-1)
 
-        video_second_per_grid = hf_inputs.get("video_second_per_grid", None)
+        video_second_per_grid = processed_data.get("video_second_per_grid", None)
         if video_second_per_grid is not None:
-            hf_inputs["second_per_grid_ts"] = video_second_per_grid
+            processed_data["second_per_grid_ts"] = video_second_per_grid
 
-        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
-        hf_inputs["use_audio_in_video"] = torch.tensor(use_audio_in_video)
-
-        processed_data = hf_inputs
-        processed_data.update(passthrough_data)
-        processed_data.pop("input_ids")
+        use_audio_in_video = hf_kwargs.get("use_audio_in_video", False)
+        processed_data["use_audio_in_video"] = torch.tensor(use_audio_in_video)
+        processed_data.pop("input_ids", None)
 
         return processed_data
 
@@ -1171,8 +1161,7 @@ def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
     def _get_audio_for_video_mapping(
         self, mm_features: list[MultiModalFeatureSpec]
     ) -> tuple[dict[int, int], set[int]]:
-        """
-        Map video offset -> paired audio_feature_length for use_audio_in_video.
+        """Map video offset -> paired audio_feature_length for use_audio_in_video.
 
         When use_audio_in_video=True, audio is interleaved within video chunks.
         The pairing is based on feature order in mm_features.
@@ -1180,6 +1169,7 @@ def _get_audio_for_video_mapping(
         Returns:
             Tuple of (video_offset -> audio_feature_length mapping,
                       set of paired audio offsets to skip)
+
         """
         videos_with_audio = [
             f
@@ -1207,8 +1197,7 @@ def _compute_audio_token_count(self, audio_feature_length: int) -> int:
     def iter_mm_features(
         self, mm_features: list[MultiModalFeatureSpec]
     ) -> Iterator[tuple[int, str, dict[str, Any]]]:
-        """
-        Iterate over multimodal features sorted by position offset.
+        """Iterate over multimodal features sorted by position offset.
 
         Yields: (offset, modality, feature_data) where feature_data contains:
         - image: {"grid_t", "grid_h", "grid_w", "t_factor"}
@@ -1278,8 +1267,7 @@ def iter_mm_features(
     def _compute_interleaved_positions(
         self, start_idx: int, data: dict[str, Any]
     ) -> tuple[np.ndarray, int]:
-        """
-        Compute positions for interleaved video+audio chunks.
+        """Compute positions for interleaved video+audio chunks.
 
         Returns: (position_ids, total_token_count)
         """
@@ -1360,8 +1348,7 @@ def get_mrope_input_positions(
         input_tokens: list[int],
         mm_features: list[MultiModalFeatureSpec],
     ) -> tuple[torch.Tensor, int]:
-        """
-        Compute M-RoPE input positions using mm_features directly.
+        """Compute M-RoPE input positions using mm_features directly.
 
         Example for use_audio_in_video case:
             (V_i are vision position ids, A_i are audio position ids)
@@ -1542,9 +1529,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="merger.",
diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py
index 61d0153f86e8..28ab566a4b6a 100644
--- a/vllm/model_executor/models/qwen2_5_vl.py
+++ b/vllm/model_executor/models/qwen2_5_vl.py
@@ -41,12 +41,14 @@
     Qwen2_5_VLConfig,
     Qwen2_5_VLVisionConfig,
 )
+from transformers.video_utils import VideoMetadata
 
 from vllm.compilation.decorators import (
     should_torch_compile_mm_encoder,
     support_torch_compile,
 )
 from vllm.config import VllmConfig
+from vllm.config.multimodal import VideoDummyOptions
 from vllm.distributed import parallel_state
 from vllm.distributed import utils as dist_utils
 from vllm.logger import init_logger
@@ -70,10 +72,13 @@
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
+    VideoItem,
 )
 from vllm.multimodal.parse import MultiModalDataItems
 from vllm.multimodal.processing import PromptReplacement, PromptUpdate
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.multimodal.video_prune.evs import (
     compute_mrope_for_media,
     compute_retained_tokens_count,
@@ -101,8 +106,9 @@
     SupportsPP,
     SupportsQuant,
 )
-from .qwen2_vl import Qwen2VLDummyInputsBuilder as Qwen2_5_VLDummyInputsBuilder
 from .qwen2_vl import (
+    Qwen2VLDummyInputsBuilder,
+    Qwen2VLMultiModalDataParser,
     Qwen2VLMultiModalProcessor,
     Qwen2VLProcessingInfo,
 )
@@ -160,8 +166,7 @@ def _pad_flashinfer_cu_seqlens_buffer(
 
 
 class Qwen2_5_VLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - np: Number of patches
         - ni: Number of images
         - cps: Number of channels * patch_size * patch_size
@@ -187,8 +192,7 @@ class Qwen2_5_VLImagePixelInputs(TensorSchema):
 
 
 class Qwen2_5_VLImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - nf: Number of image features
         - hs: Hidden size
         - ni: Number of images
@@ -221,8 +225,7 @@ class Qwen2_5_VLImageEmbeddingInputs(TensorSchema):
 
 
 class Qwen2_5_VLVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - np: Number of patches
         - nv: Number of videos
         - ctps: Number of channels * temporal_patch_size * patch_size *
@@ -261,8 +264,7 @@ class Qwen2_5_VLVideoPixelInputs(TensorSchema):
 
 
 class Qwen2_5_VLVideoEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - nf: Number of video features
         - hs: Hidden size
         - nv: Number of videos
@@ -915,8 +917,8 @@ def prepare_encoder_metadata(
                 cu_window_seqlens (needed for CUDA graph capture to
                 cover worst-case replay scenarios).
             device: Device to place tensors on. Defaults to self.device.
-        """
 
+        """
         if device is None:
             device = self.device
         metadata: dict[str, torch.Tensor] = {}
@@ -1140,6 +1142,40 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
+class Qwen2_5_VLDummyInputsBuilder(Qwen2VLDummyInputsBuilder):
+    def _get_dummy_videos(
+        self,
+        *,
+        width: int,
+        height: int,
+        num_frames: int,
+        num_videos: int,
+        overrides: VideoDummyOptions | None = None,
+    ) -> list[VideoItem]:
+        videos = super()._get_dummy_videos(
+            width=width,
+            height=height,
+            num_frames=num_frames,
+            num_videos=num_videos,
+            overrides=overrides,
+        )
+        videos = [v.copy() for v in videos]
+
+        video_items = []
+        for video in videos:
+            video_num_frames = video.shape[0]
+            video_metadata = {
+                "fps": 2.0,
+                "duration": video_num_frames / 2.0,
+                "total_num_frames": video_num_frames,
+                "frames_indices": list(range(video_num_frames)),
+                "video_backend": "opencv",
+                "do_sample_frames": False,
+            }
+            video_items.append((video, video_metadata))
+        return video_items
+
+
 class Qwen2_5_VLProcessingInfo(Qwen2VLProcessingInfo):
     def get_hf_config(self):
         return self.ctx.get_hf_config(Qwen2_5_VLConfig)
@@ -1147,10 +1183,17 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs: object) -> Qwen2_5_VLProcessor:
         return self.ctx.get_hf_processor(
             Qwen2_5_VLProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
+    def get_data_parser(self):
+        return Qwen2VLMultiModalDataParser(
+            self.get_hf_config().vision_config.spatial_merge_size,
+            video_needs_metadata=True,
+            expected_hidden_size=self._get_expected_hidden_size(),
+            allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
+        )
+
 
 class Qwen2_5_VLMultiModalProcessor(Qwen2VLMultiModalProcessor):
     def _get_mm_fields_config(
@@ -1163,6 +1206,25 @@ def _get_mm_fields_config(
             second_per_grid_ts=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
         )
 
+    def _get_hf_mm_inputs(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+        hf_data = hf_inputs.hf_data
+
+        if videos := hf_data.get("videos"):
+            # HF expects video metadata as a separate argument.
+            hf_data["videos"] = [video for video, _ in videos]
+            hf_data["video_metadata"] = [
+                VideoMetadata(
+                    **{k: metadata[k] for k in metadata if k != "do_sample_frames"}
+                )
+                for _, metadata in videos
+            ]
+        return hf_inputs
+
     def _get_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
@@ -1260,14 +1322,14 @@ class Qwen2_5_VLForConditionalGeneration(
     def iter_mm_grid_thw(
         self, mm_features: list[MultiModalFeatureSpec]
     ) -> Iterator[tuple[int, int, int, int, float]]:
-        """
-        Iterate over multimodal features and yield grid information.
+        """Iterate over multimodal features and yield grid information.
 
         Args:
             mm_features: List of multimodal feature specifications
 
         Yields:
             Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image
+
         """
         spatial_merge_size = self.config.vision_config.spatial_merge_size
         tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0)
@@ -1444,7 +1506,7 @@ def _process_image_input(
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values.type(self.visual.dtype),
+                    pixel_values,
                     grid_thw_list,
                     rope_type="rope_3d",
                 )
@@ -1461,8 +1523,7 @@ def _postprocess_image_embeds_evs(
         image_embeds_split: tuple[torch.Tensor, ...],
         image_input: Qwen2_5_VLImageInputs,
     ) -> tuple[torch.Tensor, ...]:
-        """
-        Append mrope positions for each for images.
+        """Append mrope positions for each for images.
         This is necessary to recover correct mrope
         positions after video pruning
 
@@ -1475,6 +1536,7 @@ def _postprocess_image_embeds_evs(
             Tuple of image embeddings for each image item.
             Resulting embeddings will have extra 4 channels for
             computed mrope positions.
+
         """
         merge_size = self.visual.spatial_merge_size
         grid_thw = image_input["image_grid_thw"]
@@ -1504,7 +1566,7 @@ def _process_video_input(
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values_videos.type(self.visual.dtype),
+                    pixel_values_videos,
                     grid_thw_list,
                     rope_type="rope_3d",
                 )
@@ -1521,8 +1583,7 @@ def _postprocess_video_embeds_evs(
         video_embeds_split: tuple[torch.Tensor, ...],
         video_input: Qwen2_5_VLVideoInputs,
     ) -> tuple[torch.Tensor, ...]:
-        """
-        Prunes video embeddings via Efficient Video Sampling (EVS)
+        """Prunes video embeddings via Efficient Video Sampling (EVS)
         and then appends mrope positions for each retained embeddings
 
         Args:
@@ -1533,6 +1594,7 @@ def _postprocess_video_embeds_evs(
             Tuple of video embeddings for each video item.
             Resulting embeddings will have extra 4 channels for
             computed mrope positions.
+
         """
         grid_thw = video_input["video_grid_thw"]
         assert grid_thw.ndim == 2
@@ -1583,8 +1645,7 @@ def recompute_mrope_positions(
         mrope_positions: torch.LongTensor,
         num_computed_tokens: int,
     ) -> tuple[Sequence[torch.Tensor], torch.Tensor, int]:
-        """
-        Update part of input mrope positions (starting with
+        """Update part of input mrope positions (starting with
         num_computed_tokens index). Original mrope_positions are computed
         for unpruned sequence and becomes incorrect once pruning occurs,
         so once we prune media tokens we should reflect this in the
@@ -1601,6 +1662,7 @@ def recompute_mrope_positions(
         Returns:
             Tuple of (multimodal_embeddings, mrope_positions,
                 mrope_position_delta).
+
         """
         image_token_id = self.config.image_token_id
         video_token_id = self.config.video_token_id
@@ -2017,8 +2079,10 @@ def forward(
                 batch. **NOTE**: If mrope is enabled (default setting for
                 Qwen2.5-VL opensource models), the shape will be `(3, seq_len)`,
                 otherwise it will be `(seq_len,).
-        """
+            intermediate_tensors: Intermediate tensors from prior forward pass.
+            inputs_embeds: Optional tensor of input embeddings.
 
+        """
         if intermediate_tensors is not None:
             inputs_embeds = None
 
@@ -2041,30 +2105,22 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="visual.merger.",
             tower_model="visual.",
         )
 
-    def get_num_mm_encoder_tokens(
-        self,
-        num_image_tokens: int,
-    ) -> int:
-        hf_config = self.config
-        vision_config = hf_config.vision_config
-        merge_size = vision_config.spatial_merge_size
-
-        return num_image_tokens * merge_size**2
-
-    def get_num_mm_connector_tokens(
+    def get_mm_lora_token_counts(
         self,
-        num_vision_tokens: int,
-    ) -> int:
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         hf_config = self.config
         vision_config = hf_config.vision_config
         merge_size = vision_config.spatial_merge_size
-        return num_vision_tokens // merge_size**2
+        return num_mm_embeds * merge_size**2, num_mm_embeds
diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py
index 74341575e744..dca856c45f38 100644
--- a/vllm/model_executor/models/qwen2_audio.py
+++ b/vllm/model_executor/models/qwen2_audio.py
@@ -37,7 +37,7 @@
 from transformers.models.whisper import WhisperFeatureExtractor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
@@ -59,6 +59,7 @@
     PromptReplacement,
     PromptUpdate,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.sequence import IntermediateTensors
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -69,10 +70,9 @@
 
 # # === Audio Inputs === #
 class Qwen2AudioFeatureInputs(TensorSchema):
-    """
-    Dimensions:
-        - na: Number of audios
-        - nmb: Number of mel bins
+    """Dimensions:
+    - na: Number of audios
+    - nmb: Number of mel bins
     """
 
     type: Literal["audio_features"]
@@ -88,12 +88,11 @@ class Qwen2AudioFeatureInputs(TensorSchema):
 
 
 class Qwen2AudioEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size
-        - naf: Number of audio features
-        - hs: Hidden size (must match the hidden size of language model
-          backbone)
+    """Dimensions:
+    - bn: Batch size
+    - naf: Number of audio features
+    - hs: Hidden size (must match the hidden size of language model
+      backbone)
     """
 
     type: Literal["audio_embeds"] = "audio_embeds"
@@ -129,17 +128,22 @@ def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):
 def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     return dict(
         audio_embeds=MultiModalFieldConfig.batched("audio"),
+        audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
         input_features=MultiModalFieldConfig.batched("audio"),
         feature_attention_mask=MultiModalFieldConfig.batched("audio"),
     )
 
 
 class Qwen2AudioMultiModalDataParser(MultiModalDataParser):
+    embedding_fields = {
+        "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
+    }
+
     def _parse_audio_data(
         self,
         data: dict[str, torch.Tensor] | ModalityData[AudioItem],
     ) -> ModalityDataItems[Any, Any] | None:
-        if isinstance(data, dict):
+        if isinstance(data, dict) and "audio_num_tokens" not in data:
             return DictEmbeddingItems(
                 data,
                 modality="audio",
@@ -170,6 +174,7 @@ def get_data_parser(self):
             target_sr=feature_extractor.sampling_rate,
             target_channels=self.get_target_channels(),
             expected_hidden_size=self._get_expected_hidden_size(),
+            allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
         )
 
     def get_target_channels(self) -> int:
@@ -215,48 +220,42 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         feature_extractor = self.info.get_feature_extractor()
 
         sampling_rate = feature_extractor.sampling_rate
         audio_len = feature_extractor.chunk_length * sampling_rate
-        num_audios = mm_counts.get("audio", 0)
-
-        audio_overrides = mm_options.get("audio")
 
         return {
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
 
 class Qwen2AudioMultiModalProcessor(BaseMultiModalProcessor[Qwen2AudioProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
-
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", [])
-        if audios:
-            mm_data["audio"] = audios
-
-        hf_processor_mm_kwargs = dict(
-            **hf_processor_mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+        return hf_inputs._replace(
+            hf_kwargs=dict(
+                hf_inputs.hf_kwargs,
+                sampling_rate=feature_extractor.sampling_rate,
+                truncation=True,
+            )
         )
 
-        return mm_data, hf_processor_mm_kwargs
-
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
@@ -286,7 +285,9 @@ def _get_prompt_updates(
             audio_output_lengths = audio_output_lens.tolist()
 
         def get_replacement_qwen2_audio(item_idx: int):
-            if audio_output_lengths:
+            if "audio_num_tokens" in out_mm_data:
+                num_features = int(out_mm_data["audio_num_tokens"][item_idx])
+            elif audio_output_lengths:
                 num_features = audio_output_lengths[item_idx]
             else:
                 audio_embeds = out_mm_data["audio_embeds"][item_idx]
diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py
index 3d1345397f0a..69d538f329f9 100644
--- a/vllm/model_executor/models/qwen2_moe.py
+++ b/vllm/model_executor/models/qwen2_moe.py
@@ -40,7 +40,10 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -136,11 +139,9 @@ def __init__(
                 f"the number of experts {config.num_experts}."
             )
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py
index 80abdc62b576..355c61dd4caa 100644
--- a/vllm/model_executor/models/qwen2_vl.py
+++ b/vllm/model_executor/models/qwen2_vl.py
@@ -43,7 +43,7 @@
 from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import parallel_state, tensor_model_parallel_all_gather
 from vllm.distributed import utils as dist_utils
 from vllm.inputs import ModalityData, MultiModalDataDict
@@ -66,6 +66,7 @@
     ImageItem,
     MultiModalFeatureSpec,
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
     VideoItem,
 )
@@ -119,8 +120,7 @@
 
 
 class Qwen2VLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - np: The total number of patches over each image over each prompt in
               the batch
         - ni: Number of images
@@ -147,8 +147,7 @@ class Qwen2VLImagePixelInputs(TensorSchema):
 
 
 class Qwen2VLImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - nf: Number of image features
         - hs: Hidden size
         - ni: Number of images
@@ -179,8 +178,7 @@ class Qwen2VLImageEmbeddingInputs(TensorSchema):
 
 
 class Qwen2VLVideoPixelInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - np: The total number of patches over each video over each prompt in
               the batch
         - ctps: Number of channels * temporal_patch_size * patch_size *
@@ -208,8 +206,7 @@ class Qwen2VLVideoPixelInputs(TensorSchema):
 
 
 class Qwen2VLVideoEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
         - nf: Number of video features
         - hs: Hidden size
         - nv: Number of videos
@@ -851,7 +848,6 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs: object) -> Qwen2VLProcessor:
         return self.ctx.get_hf_processor(
             Qwen2VLProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
@@ -1099,32 +1095,26 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
         target_num_frames = self.info.get_num_frames_with_most_features(
             seq_len, mm_counts
         )
 
-        image_overrides = mm_options.get("image")
-        video_overrides = mm_options.get("video")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
-                overrides=video_overrides,
+                num_videos=mm_counts.get("video", 0),
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -1207,14 +1197,14 @@ class Qwen2VLForConditionalGeneration(
     def iter_mm_grid_thw(
         self, mm_features: list[MultiModalFeatureSpec]
     ) -> Iterator[tuple[int, int, int, int, float]]:
-        """
-        Iterate over multimodal features and yield grid information.
+        """Iterate over multimodal features and yield grid information.
 
         Args:
             mm_features: List of multimodal feature specifications
 
         Yields:
             Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image
+
         """
         spatial_merge_size = self.config.vision_config.spatial_merge_size
         tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0)
@@ -1382,7 +1372,7 @@ def _process_image_input(
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values.type(self.visual.dtype),
+                    pixel_values,
                     grid_thw.tolist(),
                     rope_type="rope_3d",
                 )
@@ -1407,7 +1397,7 @@ def _process_video_input(
             if self.use_data_parallel:
                 return run_dp_sharded_mrope_vision_model(
                     self.visual,
-                    pixel_values_videos.type(self.visual.dtype),
+                    pixel_values_videos,
                     grid_thw.tolist(),
                     rope_type="rope_3d",
                 )
@@ -1714,8 +1704,10 @@ def forward(
                 otherwise it will be `(seq_len,)`.
             intermediate_tensors: Intermediate tensors from prior forward pass.
             inputs_embeds: Optional tensor of input embeddings.
-        """
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
 
+        """
         if intermediate_tensors is not None:
             inputs_embeds = None
 
@@ -1738,30 +1730,22 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="visual.merger.",
             tower_model="visual.",
         )
 
-    def get_num_mm_encoder_tokens(
+    def get_mm_lora_token_counts(
         self,
-        num_image_tokens: int,
-    ) -> int:
-        hf_config = self.config
-        vision_config = hf_config.vision_config
-        merge_size = vision_config.spatial_merge_size
-
-        return num_image_tokens * merge_size**2
-
-    def get_num_mm_connector_tokens(
-        self,
-        num_vision_tokens: int,
-    ) -> int:
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         hf_config = self.config
         vision_config = hf_config.vision_config
         merge_size = vision_config.spatial_merge_size
-        return num_vision_tokens // merge_size**2
+        return num_mm_embeds * merge_size**2, num_mm_embeds
diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py
index f991f05d74d3..9204d5305dee 100644
--- a/vllm/model_executor/models/qwen3_5.py
+++ b/vllm/model_executor/models/qwen3_5.py
@@ -583,8 +583,8 @@ def forward(
                     model. `None` if no videos are passed.
                 - video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in
                     LLM. `None` if no videos are passed.
-        """
 
+        """
         if intermediate_tensors is not None:
             inputs_embeds = None
 
diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py
index 4c064bead747..7fb699b183b7 100644
--- a/vllm/model_executor/models/qwen3_5_mtp.py
+++ b/vllm/model_executor/models/qwen3_5_mtp.py
@@ -38,6 +38,7 @@
 from .interfaces import (
     MultiModalEmbeddings,
     SupportsMultiModal,
+    SupportsPP,
     _require_is_multimodal,
 )
 from .utils import (
@@ -152,7 +153,9 @@ def forward(
         inputs_embeds: torch.Tensor | None = None,
         spec_step_idx: int = 0,
     ) -> torch.Tensor:
-        if get_pp_group().is_first_rank:
+        # Branch on the inputs, not the rank: the drafter is built entirely on
+        # the last PP stage, where `is_first_rank` is False.
+        if intermediate_tensors is None:
             if inputs_embeds is None:
                 inputs_embeds = self.embed_input_ids(input_ids)
             assert hidden_states.shape[-1] == inputs_embeds.shape[-1]
@@ -162,7 +165,6 @@ def forward(
             hidden_states = self.fc(hidden_states)
             residual = None
         else:
-            assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
 
@@ -212,7 +214,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         "hidden_states": 0,
     }
 )
-class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal):
+class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP):
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",
@@ -254,6 +256,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         self.logits_processor = LogitsProcessor(config.vocab_size)
 
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
     def embed_input_ids(
         self,
         input_ids: torch.Tensor,
diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py
index e98f37ba37c2..cf8f37cebc06 100644
--- a/vllm/model_executor/models/qwen3_asr.py
+++ b/vllm/model_executor/models/qwen3_asr.py
@@ -32,7 +32,7 @@
 from transformers.models.whisper import WhisperFeatureExtractor
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TokensPrompt
 from vllm.logger import init_logger
@@ -67,6 +67,7 @@
     AudioItem,
     MultiModalFeatureSpec,
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
 )
 from vllm.multimodal.parse import (
@@ -188,7 +189,6 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs: object) -> Qwen3ASRProcessor:
         processor = self.ctx.get_hf_processor(
             Qwen3ASRProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
         if not hasattr(processor, "audio_token"):
@@ -225,10 +225,8 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
-
         feature_extractor = self.info.get_feature_extractor()
 
         target_audio_length = (
@@ -239,13 +237,11 @@ def get_dummy_mm_data(
             * feature_extractor.sampling_rate
         )
 
-        audio_overrides = mm_options.get("audio")
-
         return {
             "audio": self._get_dummy_audios(
                 length=target_audio_length,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             ),
         }
 
@@ -617,24 +613,25 @@ def get_mrope_input_positions(
         return llm_positions, mrope_position_delta
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             tower_model=["audio_tower."],
         )
 
-    def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int:
-        """Return the number of tokens processed by the audio tower encoder.
-
-        Required for LoRA support on the tower module.
-        """
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         # For Qwen3-ASR, the audio tower produces one embedding per audio
         # placeholder token inserted into the prompt (no additional
         # merge/downsample step like vision towers). Therefore, the encoder
         # token budget is identity.
-        return num_audio_tokens
+        return num_mm_embeds, None
 
     @classmethod
     def get_speech_to_text_config(
@@ -699,8 +696,7 @@ def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
 
     @classmethod
     def post_process_output(cls, text: str) -> str:
-        """
-        Post-process Qwen3-ASR raw output to extract clean transcription.
+        """Post-process Qwen3-ASR raw output to extract clean transcription.
 
         The model outputs in format: "language {lang}<asr_text>{transcription}"
         This method strips the language prefix and asr_text tags.
diff --git a/vllm/model_executor/models/qwen3_asr_realtime.py b/vllm/model_executor/models/qwen3_asr_realtime.py
index 15ecd861f14b..6eb762b6281e 100644
--- a/vllm/model_executor/models/qwen3_asr_realtime.py
+++ b/vllm/model_executor/models/qwen3_asr_realtime.py
@@ -25,29 +25,26 @@
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
 from vllm.inputs import PromptType, TokensPrompt
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import (
-    SupportsRealtime,
-)
-from vllm.model_executor.models.qwen3_asr import (
-    Qwen3ASRDummyInputsBuilder,
-    Qwen3ASRForConditionalGeneration,
-    Qwen3ASRMultiModalProcessor,
-    Qwen3ASRProcessingInfo,
-    _get_feat_extract_output_lengths,
-)
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.cache import _I, BaseMultiModalProcessorCache
-from vllm.multimodal.inputs import MultiModalKwargsOptionalItems
 from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.processing import ProcessorInputs, TimingContext
 from vllm.multimodal.processing.processor import (
-    BaseDummyInputsBuilder,
-    MultiModalPromptUpdates,
+    MultiModalProcessingResult,
     PlaceholderFeaturesInfo,
     cached_encode,
 )
 from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.transformers_utils.processor import cached_processor_from_config
 
+from .interfaces import SupportsRealtime
+from .qwen3_asr import (
+    Qwen3ASRDummyInputsBuilder,
+    Qwen3ASRForConditionalGeneration,
+    Qwen3ASRMultiModalProcessor,
+    Qwen3ASRProcessingInfo,
+    _get_feat_extract_output_lengths,
+)
+
 logger = init_logger(__name__)
 
 _PRE_ALLOCATE_BUFFER_SIZE_IN_S = 60
@@ -102,22 +99,21 @@ def flush(self) -> np.ndarray | None:
 
 
 class Qwen3ASRRealtimeMultiModalProcessor(Qwen3ASRMultiModalProcessor):
-    def __init__(
+    def _cached_apply_hf_processor(
         self,
-        info: _I,
-        dummy_inputs: BaseDummyInputsBuilder[_I],
-        *,
-        cache: BaseMultiModalProcessorCache | None = None,
-    ) -> None:
-        super().__init__(info, dummy_inputs, cache=None)
+        inputs: ProcessorInputs,
+        timing_ctx: TimingContext,
+    ) -> MultiModalProcessingResult:
+        # realtime can't make use of a cache yet
+        return self._apply_hf_processor(inputs, timing_ctx)
 
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsOptionalItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_res: MultiModalProcessingResult,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
+        mm_kwargs = mm_res.kwargs
+
         audios = mm_kwargs.get("audio", [])
         assert len(audios) == 1, (
             f"Expected only one audio input for realtime, got {len(audios)}"
@@ -147,7 +143,7 @@ def _maybe_apply_prompt_updates(
         # Find the audio_pad token position and expand it to audio_len tokens
         expanded_ids = list[int]()
         pad_start_idx = -1
-        for i, tid in enumerate(prompt_ids):
+        for i, tid in enumerate(mm_res.prompt_ids):
             if tid == audio_pad_id and pad_start_idx == -1:
                 pad_start_idx = i
                 expanded_ids.extend([audio_pad_id] * audio_len)
diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py
index ce3884a10eb8..1636e3852c0f 100644
--- a/vllm/model_executor/models/qwen3_dflash.py
+++ b/vllm/model_executor/models/qwen3_dflash.py
@@ -32,6 +32,7 @@
     VocabParallelEmbedding,
 )
 from vllm.multimodal.inputs import NestedTensors
+from vllm.platforms import current_platform
 from vllm.transformers_utils.config import set_default_rope_theta
 from vllm.transformers_utils.repo_utils import get_hf_file_bytes
 from vllm.v1.attention.backend import AttentionType
@@ -167,6 +168,7 @@ def __init__(
         rms_norm_eps: float = 1e-06,
         attention_bias: bool = False,
         add_swa_attention_sink_bias: bool = False,
+        v_scale: float | None = None,
         sliding_window: int | None = None,
         causal: bool = False,
         is_neox_style: bool = True,
@@ -192,6 +194,7 @@ def __init__(
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
+        self.v_scale = v_scale
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
@@ -263,6 +266,8 @@ def forward(
 
         q, k = self.rotary_emb(positions, q, k)
 
+        if self.v_scale is not None:
+            v = v * self.v_scale
         attn_output = self.attn(q, k, v)
         output, _ = self.o_proj(attn_output)
         return output
@@ -311,6 +316,7 @@ def __init__(
             rms_norm_eps=config.rms_norm_eps,
             attention_bias=getattr(config, "attention_bias", False),
             add_swa_attention_sink_bias=add_swa_attention_sink_bias,
+            v_scale=dflash_config.get("attention_value_scale"),
             sliding_window=sliding_window,
             causal=causal,
             is_neox_style=is_neox_style,
@@ -561,6 +567,15 @@ def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor:
         # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) ---
         # The weight is selected per layer by the outermost (layer) index.
         all_k_normed = torch.empty_like(all_k)
+        if current_platform.is_xpu():
+            for layer_idx in range(all_k.shape[0]):
+                ops.rms_norm(
+                    all_k_normed[layer_idx],
+                    all_k[layer_idx],
+                    self._k_norm_weights[layer_idx],
+                    self._rms_norm_eps,
+                )
+            return all_k_normed
         ops.rms_norm(
             all_k_normed,
             all_k,
@@ -623,6 +638,10 @@ def precompute_and_store_context_kv(
         if context_slot_mapping is None:
             return
 
+        v_scale = getattr(self.layers[0].self_attn, "v_scale", None)
+        if v_scale is not None:
+            all_v.mul_(v_scale)
+
         # --- Per-layer cache insert ---
         all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
         per_layer = isinstance(context_slot_mapping, (list, tuple))
diff --git a/vllm/model_executor/models/qwen3_dflash2.py b/vllm/model_executor/models/qwen3_dflash2.py
index cb027293f981..2043307d0250 100644
--- a/vllm/model_executor/models/qwen3_dflash2.py
+++ b/vllm/model_executor/models/qwen3_dflash2.py
@@ -11,6 +11,8 @@
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from .qwen3_dflash import (
     DFlashQwen3DecoderLayer,
@@ -20,6 +22,125 @@
 from .utils import maybe_prefix
 
 
+@triton.jit
+def _dflash2_grouped_conv_kernel(
+    x_ptr,
+    delta_ptr,
+    base_ptr,
+    output_ptr,
+    x_stride_row,
+    delta_stride_row,
+    delta_stride_tap,
+    base_stride_tap,
+    output_stride_row,
+    NUM_CHANNELS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    TAPS: tl.constexpr,
+    ELEMENT_BLOCK: tl.constexpr,
+) -> None:
+    row = tl.program_id(0) // triton.cdiv(NUM_CHANNELS, ELEMENT_BLOCK)
+    col_block = tl.program_id(0) % triton.cdiv(NUM_CHANNELS, ELEMENT_BLOCK)
+    channels = col_block * ELEMENT_BLOCK + tl.arange(0, ELEMENT_BLOCK)
+    mask = channels < NUM_CHANNELS
+    groups = channels // GROUP_SIZE
+    position = row % BLOCK_SIZE
+
+    delta_row = delta_ptr + row * delta_stride_row
+    x_row = x_ptr + row * x_stride_row
+    accumulator = (
+        tl.load(base_ptr + channels, mask=mask, other=0.0).to(tl.float32)
+        + tl.load(delta_row + groups, mask=mask, other=0.0).to(tl.float32)
+    ) * tl.load(x_row + channels, mask=mask, other=0.0).to(tl.float32)
+
+    for tap in tl.static_range(1, TAPS):
+        tap_mask = mask & (position >= tap)
+        coefficient = tl.load(
+            base_ptr + tap * base_stride_tap + channels,
+            mask=tap_mask,
+            other=0.0,
+        ).to(tl.float32) + tl.load(
+            delta_row + tap * delta_stride_tap + groups,
+            mask=tap_mask,
+            other=0.0,
+        ).to(tl.float32)
+        x = tl.load(
+            x_ptr + (row - tap) * x_stride_row + channels,
+            mask=tap_mask,
+            other=0.0,
+        ).to(tl.float32)
+        accumulator += coefficient * x
+
+    tl.store(
+        output_ptr + row * output_stride_row + channels,
+        accumulator,
+        mask=mask,
+    )
+
+
+def dflash2_grouped_conv_impl(
+    x: torch.Tensor,
+    delta: torch.Tensor,
+    base: torch.Tensor,
+    block_size: int,
+    group_size: int,
+) -> torch.Tensor:
+    """DFlash2 grouped convolution adapted from LightSeek TokenSpeed PR #1399."""
+    num_rows, num_channels = x.shape
+    output = torch.empty_like(x)
+    if num_rows == 0:
+        return output
+
+    element_block = 1024 if num_rows >= 128 and num_channels % 1024 == 0 else 512
+    grid = (num_rows * triton.cdiv(num_channels, element_block),)
+    _dflash2_grouped_conv_kernel[grid](
+        x,
+        delta,
+        base,
+        output,
+        x.stride(0),
+        delta.stride(0),
+        delta.stride(1),
+        base.stride(0),
+        output.stride(0),
+        NUM_CHANNELS=num_channels,
+        BLOCK_SIZE=block_size,
+        GROUP_SIZE=group_size,
+        TAPS=base.shape[0],
+        ELEMENT_BLOCK=element_block,
+        num_warps=4,
+    )
+    return output
+
+
+def dflash2_grouped_conv_fake(
+    x: torch.Tensor,
+    delta: torch.Tensor,
+    base: torch.Tensor,
+    block_size: int,
+    group_size: int,
+) -> torch.Tensor:
+    return torch.empty_like(x)
+
+
+direct_register_custom_op(
+    op_name="dflash2_grouped_conv",
+    op_func=dflash2_grouped_conv_impl,
+    fake_impl=dflash2_grouped_conv_fake,
+    tags=(torch.Tag.needs_fixed_stride_order,),
+)
+
+
+def dflash2_grouped_conv(
+    x: torch.Tensor,
+    delta: torch.Tensor,
+    base: torch.Tensor,
+    block_size: int,
+    group_size: int,
+) -> torch.Tensor:
+    return torch.ops.vllm.dflash2_grouped_conv(x, delta, base, block_size, group_size)
+
+
 def _grouped_conv(
     hidden_states: torch.Tensor,
     delta: torch.Tensor,
@@ -29,6 +150,9 @@ def _grouped_conv(
     group_size: int,
     taps: int,
 ) -> torch.Tensor:
+    if hidden_states.is_cuda:
+        return dflash2_grouped_conv(hidden_states, delta, base, block_size, group_size)
+
     blocks = hidden_states.unflatten(-1, (num_groups, group_size))
     coefficients = base.view(1, taps, num_groups, group_size) + delta.unsqueeze(-1)
     output = coefficients[:, 0] * blocks
diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py
index c2c252280d6a..7318d225f8d2 100644
--- a/vllm/model_executor/models/qwen3_dspark.py
+++ b/vllm/model_executor/models/qwen3_dspark.py
@@ -254,7 +254,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             self.config.draft_vocab_size, scale=logit_scale
         )
         self.target_vocab_size = vllm_config.model_config.get_vocab_size()
-        if self.config.draft_vocab_size != self.target_vocab_size:
+        # Only a pruned draft (subset of the target vocab) needs a d2t mapping to
+        # scatter ids up. A draft >= the target vocab (e.g. padded to the target's
+        # physical embedding size) maps identically, so no d2t is needed.
+        if self.config.draft_vocab_size < self.target_vocab_size:
             self.draft_id_to_target_id = nn.Parameter(
                 torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
                 requires_grad=False,
@@ -334,7 +337,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         # embed_tokens / lm_head are optional; when omitted they are shared from
         # the target by load_dspark_model, so skip the unloaded params here.
         uses_expanded_input_vocab = self.config.vocab_size > self.target_vocab_size
-        uses_reduced_vocab = self.config.draft_vocab_size != self.target_vocab_size
+        uses_reduced_vocab = self.config.draft_vocab_size < self.target_vocab_size
         if uses_expanded_input_vocab and not includes_embed_tokens:
             raise ValueError(
                 "Qwen3 DSpark checkpoints whose input vocab_size is larger than "
diff --git a/vllm/model_executor/models/qwen3_eagle3.py b/vllm/model_executor/models/qwen3_eagle3.py
index 39d0173b3ce9..745c77a97730 100644
--- a/vllm/model_executor/models/qwen3_eagle3.py
+++ b/vllm/model_executor/models/qwen3_eagle3.py
@@ -309,10 +309,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.logits_processor = LogitsProcessor(
             self.config.draft_vocab_size, scale=logit_scale
         )
-        self.draft_id_to_target_id = nn.Parameter(
-            torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
-            requires_grad=False,
-        )
+        target_vocab_size = vllm_config.model_config.get_vocab_size()
+        if self.config.draft_vocab_size != target_vocab_size:
+            self.draft_id_to_target_id = nn.Parameter(
+                torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
+                requires_grad=False,
+            )
+        else:
+            self.draft_id_to_target_id = None
 
         self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting
 
diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py
index 5df982d741ea..e2f7dcaec52d 100644
--- a/vllm/model_executor/models/qwen3_moe.py
+++ b/vllm/model_executor/models/qwen3_moe.py
@@ -42,7 +42,10 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -164,10 +167,9 @@ def __init__(
         self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
         self.n_local_physical_experts = self.n_physical_experts // self.ep_size
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.gate",
         )
diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py
index 495650eb1989..e779e608ea80 100644
--- a/vllm/model_executor/models/qwen3_next.py
+++ b/vllm/model_executor/models/qwen3_next.py
@@ -18,7 +18,10 @@
     tensor_model_parallel_reduce_scatter,
 )
 from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+)
 from vllm.model_executor.layers.fused_moe.utils import (
     is_model_fused_shared_expert_compatible,
     resolve_layer_fused_shared_expert,
@@ -180,11 +183,9 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
         self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
         self.n_local_physical_experts = self.n_physical_experts // self.ep_size
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
@@ -378,7 +379,7 @@ def __init__(
         self.use_fused_qk_norm_rope_gate = (
             self.attn_output_gate
             and getattr(self.rotary_emb, "is_neox_style", False)
-            and current_platform.is_cuda()
+            and (current_platform.is_cuda() or current_platform.is_xpu())
             and supports_dtype
             and (text_only or supports_mrope)
         )
diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py
index acc09f50382a..bd48b554cd04 100755
--- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py
+++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py
@@ -30,7 +30,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeAudioEncoderConfig,
@@ -73,6 +73,8 @@
 )
 from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems
 from vllm.multimodal.processing.processor import (
+    HFMultiModalInputs,
+    MultiModalProcessingResult,
     MultiModalPromptUpdates,
     PlaceholderFeaturesInfo,
     PromptReplacement,
@@ -98,6 +100,7 @@
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniConditionalGenerationMixin,
     Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerMultiModalDataParser,
     Qwen2_5OmniThinkerMultiModalProcessor,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
@@ -301,11 +304,11 @@ def forward(
         cu_seqlens: torch.Tensor,
         max_seqlen: torch.Tensor | None,
     ) -> torch.Tensor:
-        """
-        Args:
-            hidden_states: Input tensor of shape (seq_len, hidden_size)
-            cu_seqlens: Cumulative sequence lengths
-            max_seqlen: Maximum sequence length in the batch
+        """Args:
+        hidden_states: Input tensor of shape (seq_len, hidden_size)
+        cu_seqlens: Cumulative sequence lengths
+        max_seqlen: Maximum sequence length in the batch
+
         """
         residual = hidden_states
         hidden_states = self.self_attn_layer_norm(hidden_states)
@@ -1150,6 +1153,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         )
 
 
+class Qwen3OmniMoeThinkerMultiModalDataParser(Qwen2_5OmniThinkerMultiModalDataParser):
+    # Keep the existing raw-media fallback for vision inputs.
+    embedding_fields = {
+        "audio": Qwen2_5OmniThinkerMultiModalDataParser.embedding_fields["audio"],
+    }
+
+
 class Qwen3OmniMoeThinkerProcessingInfo(
     Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo
 ):
@@ -1159,7 +1169,6 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs: object) -> Qwen3OmniMoeProcessor:
         processor = self.ctx.get_hf_processor(
             Qwen3OmniMoeProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
         if not hasattr(processor, "audio_token"):
@@ -1176,6 +1185,15 @@ def get_feature_extractor(self, **kwargs: object):
         assert isinstance(feature_extractor, WhisperFeatureExtractor)
         return feature_extractor
 
+    def get_data_parser(self):
+        return Qwen3OmniMoeThinkerMultiModalDataParser(
+            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
+            target_sr=self.get_feature_extractor().sampling_rate,
+            target_channels=self.get_target_channels(),
+            expected_hidden_size=self._get_expected_hidden_size(),
+            allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
+        )
+
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"audio": None, "image": None, "video": None}
 
@@ -1219,13 +1237,14 @@ def get_mm_max_tokens_per_item(
 class Qwen3OmniMoeThinkerMultiModalProcessor(
     Qwen2_5OmniThinkerMultiModalProcessor,
 ):
-    def _get_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        processor_data, passthrough_data = super()._get_hf_mm_data(mm_items)
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
 
-        audios = processor_data.get("audios")
+        audios = hf_inputs.hf_data.get("audio")
         if audios:
             feature_extractor = self.info.get_feature_extractor()
             hop_length = feature_extractor.hop_length
@@ -1239,47 +1258,30 @@ def pad_to_hop_length(x: np.ndarray) -> np.ndarray:
 
             # To make sure the cache works with padding=True, we pre-padded
             # the audio to a multiple of hop_length.
-            processor_data = dict(processor_data)
-            processor_data["audios"] = [
+            hf_inputs.hf_data["audio"] = [
                 pad_to_hop_length(audio)
                 if isinstance(audio, np.ndarray)
                 else (pad_to_hop_length(audio[0]), audio[1])
                 for audio in audios
             ]
 
-        return processor_data, passthrough_data
+        return hf_inputs
 
-    def _apply_hf_processor_main(
+    def _postprocess_hf_mm_data(
         self,
-        mm_items: MultiModalDataItems,
+        mm_data: Mapping[str, object],
         hf_processor_mm_kwargs: Mapping[str, object],
+        processed_data: BatchFeature,
     ) -> BatchFeature:
-        if mm_items.get_all_counts().get("audio", 0):
-            # TODO(Isotr0py): Remove this patch after upstream fix PR
-            # released and Transformers version update:
-            # https://github.com/huggingface/transformers/pull/41473
-            hf_processor_mm_kwargs = dict(hf_processor_mm_kwargs)
-            hf_processor_mm_kwargs["audio_kwargs"] = dict(
-                hf_processor_mm_kwargs.get("audio_kwargs") or {}
-            )
-            hf_processor_mm_kwargs["text_kwargs"] = dict(
-                hf_processor_mm_kwargs.get("text_kwargs") or {}
-            )
-
-        processed_data = super()._apply_hf_processor_main(
-            mm_items,
-            hf_processor_mm_kwargs,
+        processed_data = super()._postprocess_hf_mm_data(
+            mm_data, hf_processor_mm_kwargs, processed_data
         )
 
         if (
             "audio_feature_lengths" in processed_data
             and "feature_attention_mask" in processed_data
         ):
-            valid_mm_items = mm_items.select(
-                {k for k, c in mm_items.get_all_counts().items() if c > 0}
-            )
-            processor_data, _ = self._get_hf_mm_data(valid_mm_items)
-            audios = processor_data.get("audios", [])
+            audios = mm_data.get("audio", [])
             if audios:
                 feature_extractor = self.info.get_feature_extractor(
                     **hf_processor_mm_kwargs
@@ -1312,14 +1314,14 @@ def _apply_hf_processor_main(
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_res: MultiModalProcessingResult,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
-        """
-        Qwen3-Omni reimplements this function to handle `use_audio_in_video`.
-        """
+        """Qwen3-Omni reimplements this function to handle `use_audio_in_video`."""
         mm_item_counts = mm_items.get_all_counts()
+        prompt_ids = mm_res.prompt_ids
+        mm_kwargs = mm_res.kwargs
+        mm_prompt_updates = mm_res.prompt_updates
+
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
 
         use_audio_in_video = False
@@ -1365,7 +1367,7 @@ def _maybe_apply_prompt_updates(
 
     def get_updates_use_audio_in_video(
         self,
-        thinker_config: PretrainedConfig,
+        thinker_config: PreTrainedConfig,
         audio_len: int,
         video_grid_thw: list[int] | torch.Tensor,
         video_second_per_grid_t: float,
@@ -1528,8 +1530,7 @@ def _derive_audio_from_video_placeholders(
         placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
         mm_prompt_updates: MultiModalPromptUpdates,
     ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
-        """
-        Helper to derive audio placeholders from video placeholders when
+        """Helper to derive audio placeholders from video placeholders when
         use_audio_in_video=True.
         """
         if "video" not in placeholders:
@@ -1991,8 +1992,7 @@ def _compute_audio_token_count(self, audio_feature_length: int) -> int:
     def _get_audio_for_video_mapping(
         self, mm_features: list[MultiModalFeatureSpec]
     ) -> tuple[dict[int, int], set[int]]:
-        """
-        Map video offset -> paired audio_feature_length for use_audio_in_video.
+        """Map video offset -> paired audio_feature_length for use_audio_in_video.
 
         When use_audio_in_video=True, audio is interleaved within video.
         The pairing is based on feature order in mm_features.
@@ -2000,6 +2000,7 @@ def _get_audio_for_video_mapping(
         Returns:
             Tuple of (video_offset -> audio_feature_length mapping,
                       set of paired audio offsets to skip)
+
         """
         videos_with_audio = [
             f
@@ -2022,8 +2023,7 @@ def _get_audio_for_video_mapping(
     def iter_mm_features(
         self, mm_features: list[MultiModalFeatureSpec]
     ) -> Iterator[tuple[int, str, dict[str, Any]]]:
-        """
-        Iterate over multimodal features sorted by position offset.
+        """Iterate over multimodal features sorted by position offset.
 
         Yields: (offset, modality, feature_data) where feature_data contains:
         - image: {"grid_t", "grid_h", "grid_w", "t_factor"}
@@ -2088,8 +2088,7 @@ def iter_mm_features(
     def _compute_interleaved_positions(
         self, start_idx: int, data: dict[str, Any]
     ) -> tuple[np.ndarray, int]:
-        """
-        Compute positions for interleaved video+audio using Qwen3 token-by-token
+        """Compute positions for interleaved video+audio using Qwen3 token-by-token
         interleaving logic.
 
         Returns: (position_ids [3, N], total_token_count)
@@ -2153,9 +2152,7 @@ def get_speech_to_text_config(
 
     @classmethod
     def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
-        """
-        Construct a transcription/translation prompt for Qwen3-Omni.
-        """
+        """Construct a transcription/translation prompt for Qwen3-Omni."""
         audio = stt_params.audio
         stt_config = stt_params.stt_config
         model_config = stt_params.model_config
@@ -2315,9 +2312,7 @@ def get_mrope_input_positions(
         return torch.from_numpy(llm_positions), mrope_position_delta
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="visual.merger",
diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py
index d99cfe7aa296..619f2aaad43e 100644
--- a/vllm/model_executor/models/qwen3_vl.py
+++ b/vllm/model_executor/models/qwen3_vl.py
@@ -54,8 +54,8 @@
 )
 from vllm.config import VllmConfig
 from vllm.config.multimodal import (
-    BaseDummyOptions,
     MultiModalConfig,
+    MultiModalDummyOptions,
     VideoDummyOptions,
     VideoPruningMethod,
 )
@@ -776,6 +776,7 @@ def prepare_encoder_metadata(
                 instead of computing from cu_seqlens (needed for CUDA
                 graph capture to cover worst-case replay scenarios).
             device: Device to place tensors on. Defaults to self.device.
+
         """
         if device is None:
             device = self.device
@@ -914,7 +915,6 @@ def get_hf_config(self):
     def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor:
         return self.ctx.get_hf_processor(
             Qwen3VLProcessor,
-            use_fast=kwargs.pop("use_fast", True),
             **kwargs,
         )
 
@@ -1118,11 +1118,8 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-        num_videos = mm_counts.get("video", 0)
-        image_overrides = mm_options.get("image")
         video_overrides = mm_options.get("video")
 
         target_image_width, target_image_height = (
@@ -1132,7 +1129,6 @@ def get_dummy_mm_data(
         # treat videos as special images
         target_num_frames = 2
         if video_overrides:
-            assert isinstance(video_overrides, VideoDummyOptions)
             num_frames_override = video_overrides.num_frames
             if num_frames_override:
                 if num_frames_override > target_num_frames:
@@ -1203,7 +1199,6 @@ def get_dummy_mm_data(
             target_video_size.height,
         )
         if video_overrides:
-            assert isinstance(video_overrides, VideoDummyOptions)
             width_override = video_overrides.width
             if width_override:
                 if width_override > target_video_width:
@@ -1229,14 +1224,14 @@ def get_dummy_mm_data(
             "image": self._get_dummy_images(
                 width=target_image_width,
                 height=target_image_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=target_video_width,
                 height=target_video_height,
                 num_frames=target_num_frames,
-                num_videos=num_videos,
+                num_videos=mm_counts.get("video", 0),
             ),
         }
 
@@ -1291,6 +1286,7 @@ def _replace_video_token_placeholders(
 
     Returns:
         Token IDs with every placeholder replaced.
+
     """
     result: list[int] = []
     repl_idx = 0
@@ -1324,27 +1320,25 @@ def _expands_only_video_token(hf_processor: ProcessorMixin) -> bool:
         proc_impl = getattr(type(hf_processor), "replace_video_token", None)
         return proc_impl is not None and proc_impl is not mixin_impl
 
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
-
-        mm_data = dict(mm_data)
+        if not hf_data:
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
         # Separate video processing from image processing. Because the videos
         # are processed into several image patches
         video_input_ids_lst: list[list[int]] = []
-        if videos := mm_data.pop("videos", []):
+        if videos := hf_data.pop("videos", []):
             video_grid_thw_lst = []
             pixel_values_videos_lst = []
             timestamps_per_video = []
@@ -1367,10 +1361,8 @@ def _apply_hf_processor_main(
 
                 # NOTE: a copy of is created to update do_sample_frames,
                 # otherwise mm_hash for the object will be incorrect.
-                video_mm_kwargs = dict(**hf_processor_mm_kwargs)
-                merged = self.info.ctx.get_merged_mm_kwargs(
-                    hf_processor_mm_kwargs, modality="video"
-                )
+                video_mm_kwargs = dict(**hf_kwargs)
+                merged = self.info.ctx.get_merged_mm_kwargs(hf_kwargs, modality="video")
                 if merged.keys() & {"size", "min_pixels", "max_pixels"}:
                     video_size = dict(self.info.get_video_processor().size)
                     size_override = merged.get("size")
@@ -1482,13 +1474,11 @@ def _apply_hf_processor_main(
         # fps/num_frames are video-only kwargs already consumed by the loop;
         # exclude them so the text/image processor call below never gets a list.
         non_video_mm_kwargs = {
-            k: v
-            for k, v in hf_processor_mm_kwargs.items()
-            if k not in ("fps", "num_frames")
+            k: v for k, v in hf_kwargs.items() if k not in ("fps", "num_frames")
         }
         processed_data = self.info.ctx.call_hf_processor(
             self.info.get_hf_processor(**non_video_mm_kwargs),
-            dict(text=prompt_text, **mm_data),
+            hf_data,
             non_video_mm_kwargs,
         )
 
@@ -1513,9 +1503,9 @@ def _apply_hf_processor_main(
             processed_data["input_ids"] = [expanded_ids]
 
         processed_data.update(video_outputs)
-        processed_data.update(passthrough_data)
-
-        return processed_data
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _get_mm_fields_config(
         self,
@@ -1647,9 +1637,11 @@ def get_video_repl(
             vision_start_token_id: Token ID for vision start marker
             vision_end_token_id: Token ID for vision end marker
             video_token_id: Token ID for video content
+            select_token_id: Whether to return token ids instead of text
 
         Returns:
             PromptUpdateDetails with full token sequence
+
         """
         assert len(timestamps) == len(tokens_per_frame), (
             "timestamps and tokens_per_frame must have the same length"
@@ -2350,8 +2342,7 @@ def _postprocess_image_embeds_evs(
         image_embeds_split: tuple[torch.Tensor, ...],
         image_input: Qwen2_5_VLImageInputs,
     ) -> tuple[torch.Tensor, ...]:
-        """
-        Append mrope positions for each for images.
+        """Append mrope positions for each for images.
         This is necessary to recover correct mrope
         positions after video pruning
 
@@ -2364,6 +2355,7 @@ def _postprocess_image_embeds_evs(
             Tuple of image embeddings for each image item.
             Resulting embeddings will have extra 5 channels for
             computed mrope positions, consistent with video embeddings.
+
         """
         if self.is_multimodal_pruning_enabled:
             merge_size = self.visual.spatial_merge_size
@@ -2393,8 +2385,7 @@ def _postprocess_video_embeds_evs(
         video_embeds_split: tuple[torch.Tensor, ...],
         video_input: Qwen2_5_VLVideoInputs,
     ) -> tuple[torch.Tensor, ...]:
-        """
-        Prunes video embeddings via Efficient Video Sampling (EVS)
+        """Prunes video embeddings via Efficient Video Sampling (EVS)
         and then appends mrope positions for each retained embeddings
 
         Args:
@@ -2405,6 +2396,7 @@ def _postprocess_video_embeds_evs(
             Tuple of video embeddings for each video item.
             Resulting embeddings will have extra 5 channels for computed mrope
             positions, and whether the index corresponds to a video embedding.
+
         """
         grid_thw = video_input["video_grid_thw"]
         assert grid_thw.ndim == 2
@@ -2478,7 +2470,6 @@ def _create_final_video_embeddings(
         These embeddings will replace the placeholder embeddings to create
         input_embeds for the LLM.
         """
-
         device = video_embeddings.device
 
         # Generate video replacement token IDs using get_video_repl
@@ -2664,6 +2655,7 @@ def _iter_mm_grid_hw(
             llm_grid_h: Logical grid height (may not match actual token count with EVS).
             llm_grid_w: Logical grid width (may not match actual token count with EVS).
             actual_num_tokens: Actual number of video/image tokens in the placeholder.
+
         """
         for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
             offset = mm_feature.mm_position.offset
@@ -2806,8 +2798,7 @@ def recompute_mrope_positions(
         mrope_positions: torch.LongTensor,
         num_computed_tokens: int,
     ) -> tuple[Sequence[torch.Tensor], torch.Tensor, int]:
-        """
-        Update part of input mrope positions (starting with
+        """Update part of input mrope positions (starting with
         num_computed_tokens index). Original mrope_positions are computed
         for unpruned sequence and becomes incorrect once pruning occurs,
         so once we prune media tokens we should reflect this in the
@@ -2825,6 +2816,7 @@ def recompute_mrope_positions(
         Returns:
             Tuple of (multimodal_embeddings, mrope_positions,
                 mrope_position_delta).
+
         """
         return self._recompute_mrope_positions(
             input_ids=input_ids,
@@ -3031,8 +3023,8 @@ def forward(
                     model. `None` if no videos are passed.
                 - video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in
                     LLM. `None` if no videos are passed.
-        """
 
+        """
         if intermediate_tensors is not None:
             inputs_embeds = None
 
@@ -3068,33 +3060,25 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector=["visual.merger", "visual.deepstack_merger_list"],
             tower_model="visual.",
         )
 
-    def get_num_mm_encoder_tokens(
-        self,
-        num_image_tokens: int,
-    ) -> int:
-        hf_config = self.config
-        vision_config = hf_config.vision_config
-        merge_size = vision_config.spatial_merge_size
-
-        return num_image_tokens * merge_size**2
-
-    def get_num_mm_connector_tokens(
+    def get_mm_lora_token_counts(
         self,
-        num_vision_tokens: int,
-    ) -> int:
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         hf_config = self.config
         vision_config = hf_config.vision_config
         merge_size = vision_config.spatial_merge_size
-        return num_vision_tokens // merge_size**2
+        return num_mm_embeds * merge_size**2, num_mm_embeds
 
 
 @lru_cache
diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py
index 47dc8f89d326..355bfb056a9f 100644
--- a/vllm/model_executor/models/radio.py
+++ b/vllm/model_executor/models/radio.py
@@ -18,7 +18,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -109,7 +109,7 @@ def forward(self, x: torch.Tensor):
 class ViTPatchGenerator(nn.Module):
     def __init__(
         self,
-        #  config: PretrainedConfig,
+        #  config: PreTrainedConfig,
         patch_size: int,
         embed_dim: int,
         input_dims: input_dim_t,
@@ -227,6 +227,7 @@ def forward_video(self, x: torch.Tensor) -> torch.Tensor:
 
         Returns:
             Embedded patches with temporal compression applied.
+
         """
         assert self.temporal_patch_size > 1
         T = self.temporal_patch_size
@@ -534,7 +535,7 @@ class RadioInternVisionModel(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig = None,
+        config: PreTrainedConfig = None,
         quant_config: QuantizationConfig | None = None,
         *,
         num_hidden_layers_override: int | None = None,
@@ -669,7 +670,7 @@ class RadioModel(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         *,
         num_hidden_layers_override: int | None = None,
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index a9fe7c2ec2ad..cdb13df8f5b1 100644
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Whenever you add an architecture to this page, please also update
+"""Whenever you add an architecture to this page, please also update
 `tests/models/registry.py` with example HuggingFace models for it.
 """
 
@@ -376,6 +375,10 @@
         "vllm.models.deepseek_v4",
         "DeepseekV4ForConditionalGeneration",
     ),
+    "DeepseekV41ForCausalLM": (
+        "vllm.models.deepseek_v41",
+        "DeepseekV41ForCausalLM",
+    ),
     "Dots3NoteForCausalLM": (
         "vllm.models.dots3_note",
         "Dots3NoteForCausalLM",
@@ -640,6 +643,10 @@
     "MuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
     "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
     "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
+    "DSparkV41DraftModel": (
+        "vllm.models.deepseek_v41",
+        "DSparkDeepseekV4ForCausalLM",
+    ),
     "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
     "Qwen3OmniDSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
     "K3DSparkModel": (
@@ -705,6 +712,7 @@
     "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"),
     "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"),
     "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
+    "NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"),
     "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"),
     "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"),
     "Olmo3ForCausalLM": ("transformers", "TransformersForCausalLM"),
@@ -919,9 +927,7 @@ def load_model_cls(self) -> type[nn.Module]:
 
 @dataclass(frozen=True)
 class _RegisteredModel(_BaseRegisteredModel):
-    """
-    Represents a model that has already been imported in the main process.
-    """
+    """Represents a model that has already been imported in the main process."""
 
     interfaces: _ModelInfo
     model_cls: type[nn.Module]
@@ -942,9 +948,7 @@ def load_model_cls(self) -> type[nn.Module]:
 
 @dataclass(frozen=True)
 class _LazyRegisteredModel(_BaseRegisteredModel):
-    """
-    Represents a model that has not been imported in the main process.
-    """
+    """Represents a model that has not been imported in the main process."""
 
     module_name: str
     class_name: str
@@ -1009,7 +1013,7 @@ def _load_modelinfo_from_cache(self, module_hash: str) -> _ModelInfo | None:
             return None
 
     def _save_modelinfo_to_cache(self, mi: _ModelInfo, module_hash: str) -> None:
-        """save dictionary json file to cache"""
+        """Save dictionary json file to cache."""
         from vllm.model_executor.model_loader.weight_utils import atomic_writer
 
         try:
@@ -1119,8 +1123,7 @@ def register_model(
         model_arch: str,
         model_cls: type[nn.Module] | str,
     ) -> None:
-        """
-        Register an external model to be used in vLLM.
+        """Register an external model to be used in vLLM.
 
         `model_cls` can be either:
 
diff --git a/vllm/model_executor/models/roberta.py b/vllm/model_executor/models/roberta.py
index 59054b0667dc..6a65b3a7e93f 100644
--- a/vllm/model_executor/models/roberta.py
+++ b/vllm/model_executor/models/roberta.py
@@ -275,6 +275,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
     Attributes:
         roberta: An instance of BertModel used for forward operations.
         _pooler: An instance of Pooler used for pooling operations.
+
     """
 
     is_pooling_model = True
diff --git a/vllm/model_executor/models/rvl.py b/vllm/model_executor/models/rvl.py
index bf9661bcaf8e..5e822edae1db 100644
--- a/vllm/model_executor/models/rvl.py
+++ b/vllm/model_executor/models/rvl.py
@@ -8,7 +8,7 @@
 from transformers.activations import GELUActivation
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
@@ -40,20 +40,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             ),
         }
 
diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py
index 14d68fec568f..e7db1d680130 100644
--- a/vllm/model_executor/models/sarvam.py
+++ b/vllm/model_executor/models/sarvam.py
@@ -36,7 +36,11 @@
     get_tensor_model_parallel_world_size,
 )
 from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    GateLinear,
+    MoERunner,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -306,11 +310,11 @@ def __init__(
         else:
             self.router_dtype = torch.bfloat16
 
-        self.gate = nn.Linear(
+        self.gate = GateLinear(
             self.hidden_size,
             self.num_experts,
-            bias=False,
-            dtype=self.router_dtype,
+            out_dtype=self.router_dtype,
+            prefix=f"{prefix}.gate",
         )
 
         if getattr(config, "moe_router_enable_expert_bias", True):
@@ -365,12 +369,7 @@ def maybe_get_fused_moe(self) -> MoERunner:
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
-        router_logits = self.gate(
-            hidden_states.to(self.router_dtype)
-            if self.router_dtype is not None
-            else hidden_states
-        )
-        router_logits = router_logits.to(hidden_states.dtype)
+        router_logits, _ = self.gate(hidden_states)
         final_hidden = self.experts(
             hidden_states=hidden_states,
             router_logits=router_logits,
@@ -538,6 +537,7 @@ def forward(
         Returns:
             Intermediate tensors on non-final stages. On the final stage,
             normalized hidden states, paired with auxiliary states if captured.
+
         """
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
@@ -643,9 +643,21 @@ class SarvamMLAForCausalLM(
         "gate_up_proj": ["gate_proj", "up_proj"],
     }
 
+    @staticmethod
+    def _remap_config(config) -> None:
+        """Default the routing keys the released checkpoints omit."""
+        defaults = {
+            "n_group": 1,
+            "topk_group": 1,
+        }
+        for attr, default in defaults.items():
+            if getattr(config, attr, None) is None:
+                setattr(config, attr, default)
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         super().__init__()
         config = vllm_config.model_config.hf_config
+        self._remap_config(config)
         quant_config = vllm_config.quant_config
         self.config = config
         self.quant_config = quant_config
diff --git a/vllm/model_executor/models/seed_oss.py b/vllm/model_executor/models/seed_oss.py
index df0ed7e72e5a..976a35a78ae6 100644
--- a/vllm/model_executor/models/seed_oss.py
+++ b/vllm/model_executor/models/seed_oss.py
@@ -28,7 +28,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig as SeedOssConfig
+from transformers import PreTrainedConfig as SeedOssConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py
index 6d1fd029624c..ab7146c5981d 100644
--- a/vllm/model_executor/models/siglip.py
+++ b/vllm/model_executor/models/siglip.py
@@ -16,7 +16,7 @@
 )
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import divide, get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict, MultiModalInput
 from vllm.model_executor.layers.activation import get_act_fn
@@ -73,12 +73,11 @@
 
 
 class SiglipImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height of each image
-        - w: Width of each image
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height of each image
+    - w: Width of each image
     """
 
     type: Literal["pixel_values"]
@@ -159,20 +158,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_images = mm_counts.get("image", 0)
-
         target_width, target_height = self.info.get_image_size_with_most_features()
 
-        image_overrides = mm_options.get("image")
-
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -399,7 +394,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
     ) -> tuple[torch.Tensor, None]:
-        """Input shape: Batch x Time x Channel"""
+        """Input shape: Batch x Time x Channel."""
         qkv_states, _ = self.qkv_proj(hidden_states)
         query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
         out = self.attn(query_states, key_states, value_states)
@@ -778,8 +773,9 @@ def maybe_layer_norm_and_apply_head(
         """Apply the post layer norm and head if they are enabled,
         given the last hidden states tensor.
 
-        args:
+        Args:
             encoder_outputs: The last hidden states from the visual encoder.
+
         """
         if self.post_layernorm is not None:
             encoder_outputs = self.post_layernorm(encoder_outputs)
diff --git a/vllm/model_executor/models/siglip2navit.py b/vllm/model_executor/models/siglip2navit.py
index c16619612fc3..ef0181002b1f 100644
--- a/vllm/model_executor/models/siglip2navit.py
+++ b/vllm/model_executor/models/siglip2navit.py
@@ -9,7 +9,7 @@
 from torch import nn
 from torch.nn import functional as F
 from transformers import Siglip2VisionConfig
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.distributed import divide, get_tensor_model_parallel_world_size
 from vllm.model_executor.layers.activation import get_act_fn
@@ -47,7 +47,7 @@ def forward(self, seqlen: int) -> torch.Tensor:
 
 
 class Siglip2VisionEmbeddings(nn.Module):
-    def __init__(self, config: PretrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
         self.config = config
         self.embed_dim = config.hidden_size
@@ -86,17 +86,16 @@ def forward(
         pixel_values: torch.FloatTensor,
         grid_thws: torch.LongTensor | None = None,
     ) -> torch.Tensor:
-        """
-        Args:
-            pixel_values (`torch.FloatTensor`):
-                Pixel values of shape (
-                    num_patches,
-                    num_channels * temporal_patch_size * patch_size * patch_size
-                )
-            grid_thws: (`torch.LongTensor`):
-                grid shape (num_patches, 3)
-        """
+        """Args:
+        pixel_values (`torch.FloatTensor`):
+            Pixel values of shape (
+                num_patches,
+                num_channels * temporal_patch_size * patch_size * patch_size
+            )
+        grid_thws: (`torch.LongTensor`):
+            grid shape (num_patches, 3)
 
+        """
         # Apply patch embeddings to already patchified pixel values
         target_dtype = self.patch_embedding.weight.dtype
         if isinstance(self.patch_embedding, LinearBase):
@@ -173,7 +172,7 @@ def apply_rotary_pos_emb(
 
 
 class Siglip2Attention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
 
     def __init__(
         self,
@@ -236,8 +235,7 @@ def forward(
         cu_seqlens: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Input shape: Batch x Time x Channel"""
-
+        """Input shape: Batch x Time x Channel."""
         seq_length, embed_dim = hidden_states.shape
 
         qkv_states, _ = self.qkv_proj(hidden_states)
@@ -338,11 +336,11 @@ def forward(
         cu_seqlens: torch.Tensor,
         position_embeddings: torch.Tensor,
     ) -> tuple[torch.FloatTensor]:
-        """
-        Args:
-            hidden_states: Input tensor of shape (batch, seq_len, embed_dim).
-            cu_seqlens: Cumulative sequence lengths tensor.
-            position_embeddings: Position embeddings tensor.
+        """Args:
+        hidden_states: Input tensor of shape (batch, seq_len, embed_dim).
+        cu_seqlens: Cumulative sequence lengths tensor.
+        position_embeddings: Position embeddings tensor.
+
         """
         residual = hidden_states
 
@@ -362,12 +360,12 @@ def forward(
 
 
 class Siglip2Encoder(nn.Module):
-    """
-    Transformer encoder consisting of `config.num_hidden_layers`
+    """Transformer encoder consisting of `config.num_hidden_layers`
     self attention layers. Each layer is a [`Siglip2EncoderLayer`].
 
     Args:
-        config: PretrainedConfig
+        config: PreTrainedConfig
+
     """
 
     def __init__(
@@ -486,15 +484,15 @@ def forward(
         inputs_embeds: torch.Tensor,
         grid_thws: torch.Tensor,
     ) -> torch.Tensor:
-        r"""
-        Args:
-            inputs_embeds: Input tensor of shape
-                (batch_size, sequence_length, hidden_size).
-                Embedded representation of the input tokens.
-            grid_thws: Grid tensor of shape (num_patches, 3)
-                containing grid dimensions.
-                Whether or not to return a [`~utils.ModelOutput`] instead of
-                a plain tuple.
+        r"""Args:
+        inputs_embeds: Input tensor of shape
+            (batch_size, sequence_length, hidden_size).
+            Embedded representation of the input tokens.
+        grid_thws: Grid tensor of shape (num_patches, 3)
+            containing grid dimensions.
+            Whether or not to return a [`~utils.ModelOutput`] instead of
+            a plain tuple.
+
         """
         rotary_pos_emb = self.rot_pos_emb(grid_thws)
         window_index, cu_window_seqlens = self.get_window_index(grid_thws)
@@ -575,10 +573,9 @@ def forward(
         pixel_values: torch.FloatTensor,
         grid_thws: torch.LongTensor,
     ) -> torch.Tensor:
-        r"""
-        spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
-            Tensor containing the spatial dimensions (height, width)
-            of the input images.
+        r"""spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
+        Tensor containing the spatial dimensions (height, width)
+        of the input images.
         """
         hidden_states = self.embeddings(pixel_values, grid_thws)
 
diff --git a/vllm/model_executor/models/skyworkr1v.py b/vllm/model_executor/models/skyworkr1v.py
index ef2626db1f3c..edf7fbeb4795 100644
--- a/vllm/model_executor/models/skyworkr1v.py
+++ b/vllm/model_executor/models/skyworkr1v.py
@@ -12,7 +12,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.linear import ReplicatedLinear
@@ -49,13 +49,12 @@
 
 
 class SkyworkR1VImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bnp: Batch size * number of images * (1 + num_patches)
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
-        - bn: Batch size * number of images
+    """Dimensions:
+    - bnp: Batch size * number of images * (1 + num_patches)
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
+    - bn: Batch size * number of images
     """
 
     type: Literal["pixel_values"] = "pixel_values"
@@ -72,12 +71,11 @@ class SkyworkR1VImagePixelInputs(TensorSchema):
 
 
 class SkyworkR1VImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - ni: Number of images
-        - ifs: Image feature size
-        - hs: Hidden size (must match the hidden size of language model
-          backbone)
+    """Dimensions:
+    - ni: Number of images
+    - ifs: Image feature size
+    - hs: Hidden size (must match the hidden size of language model
+      backbone)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -202,7 +200,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         )
 
     def _patch_quant_config(
-        self, config: PretrainedConfig, quant_config: QuantizationConfig | None
+        self, config: PreTrainedConfig, quant_config: QuantizationConfig | None
     ):
         # the awq models from OpenGVLab missing `modules_to_not_convert`
         # patch the quant_config to add `modules_to_not_convert` back
@@ -216,7 +214,7 @@ def _patch_quant_config(
 
     def _init_vision_model(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
@@ -238,7 +236,7 @@ def _init_vision_model(
 
     def _init_mlp1(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None,
         prefix: str = "",
     ) -> nn.Module:
diff --git a/vllm/model_executor/models/solar.py b/vllm/model_executor/models/solar.py
index cb175f7c7a3e..194d7b8ffe40 100644
--- a/vllm/model_executor/models/solar.py
+++ b/vllm/model_executor/models/solar.py
@@ -28,7 +28,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -102,7 +102,7 @@ def forward(self, x):
 class SolarAttention(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
@@ -185,7 +185,7 @@ def forward(
 class SolarDecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py
index 57c35ed8ce3a..5baeeb69eb4b 100644
--- a/vllm/model_executor/models/step3_text.py
+++ b/vllm/model_executor/models/step3_text.py
@@ -20,6 +20,7 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
@@ -77,11 +78,9 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.experts",
         )
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.moe_num_experts,
-            bias=False,
-            quant_config=None,
             prefix=f"{prefix}.gate",
         )
 
diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py
index b060ed1d15dd..ac747a7ecb39 100644
--- a/vllm/model_executor/models/step3_vl.py
+++ b/vllm/model_executor/models/step3_vl.py
@@ -11,7 +11,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
@@ -63,15 +63,14 @@
 
 
 class Step3VLImagePixelInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - c: Number of channels (3)
-        - h: Height
-        - w: Width
-        - bnp: Batch size * number of images * number of patches
-        - hp: Height of patch
-        - wp: Width of patch
+    """Dimensions:
+    - bn: Batch size * number of images
+    - c: Number of channels (3)
+    - h: Height
+    - w: Width
+    - bnp: Batch size * number of images * number of patches
+    - hp: Height of patch
+    - wp: Width of patch
     """
 
     type: Literal["pixel_values"]
@@ -81,11 +80,10 @@ class Step3VLImagePixelInputs(TensorSchema):
 
 
 class Step3VLImageEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
-        - bn: Batch size * number of images
-        - f: Image feature size
-        - h: Hidden size (must match the hidden size of language model backbone)
+    """Dimensions:
+    - bn: Batch size * number of images
+    - f: Image feature size
+    - h: Hidden size (must match the hidden size of language model backbone)
     """
 
     type: Literal["image_embeds"] = "image_embeds"
@@ -141,19 +139,16 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         target_width, target_height = self.info.get_image_size_with_most_features()
-        num_images = mm_counts.get("image", 0)
-
-        image_overrides = mm_options.get("image")
 
         return {
             "image": self._get_dummy_images(
                 width=target_width,
                 height=target_height,
-                num_images=num_images,
-                overrides=image_overrides,
+                num_images=mm_counts.get("image", 0),
+                overrides=mm_options.get("image"),
             )
         }
 
@@ -297,7 +292,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
 
 
 class Step3VisionAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
 
     def __init__(
         self,
@@ -350,7 +345,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
     ):
-        """Input shape: Batch x Time x Channel"""
+        """Input shape: Batch x Time x Channel."""
         bsz, tgt_len, _ = hidden_states.size()
 
         # get query proj
diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py
index 3a928aa59daa..b56d4d3ed8d8 100644
--- a/vllm/model_executor/models/step3p5.py
+++ b/vllm/model_executor/models/step3p5.py
@@ -8,7 +8,6 @@
 
 import torch
 from torch import nn
-from torch.nn.parameter import Parameter
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -25,6 +24,7 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEFactory,
+    GateLinear,
     MoERunner,
     fused_moe_make_expert_params_mapping,
 )
@@ -33,7 +33,6 @@
     ColumnParallelLinear,
     MergedColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -50,7 +49,7 @@
 if typing.TYPE_CHECKING:
     from vllm.transformers_utils.configs.step3p5 import Step3p5Config
 else:
-    from transformers import PretrainedConfig as Step3p5Config
+    from transformers import PreTrainedConfig as Step3p5Config
 
 from .interfaces import MixtureOfExperts, SupportsPP
 from .utils import (
@@ -68,19 +67,6 @@
 logger = init_logger(__name__)
 
 
-class FP32ReplicatedLinear(ReplicatedLinear):
-    """
-    Use FP32 for higher precision.
-    """
-
-    def forward(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
-        assert self.params_dtype == torch.float32
-        return super().forward(x.to(torch.float32))
-
-
 class Step3p5MLP(nn.Module):
     def __init__(
         self,
@@ -173,7 +159,6 @@ def __init__(
                 rope_scaling = None
 
         if sliding_window is not None and enable_sliding_window:
-            sliding_window = sliding_window
             if swa_num_attention_heads is not None:
                 num_heads = swa_num_attention_heads
                 self.total_num_heads = swa_num_attention_heads
@@ -329,12 +314,12 @@ def __init__(
                 f"the number of experts {config.moe_num_experts}."
             )
 
-        self.gate = FP32ReplicatedLinear(
+        # Router logits are accumulated in FP32 for higher precision.
+        self.gate = GateLinear(
             config.hidden_size,
             config.moe_num_experts,
-            bias=False,
-            quant_config=None,
-            params_dtype=torch.float32,  # Use FP32 for higher precision.
+            out_dtype=torch.float32,
+            params_dtype=torch.float32,
             prefix=f"{prefix}.gate",
         )
         self.use_moe_router_bias = config.use_moe_router_bias
diff --git a/vllm/model_executor/models/step3p5_mtp.py b/vllm/model_executor/models/step3p5_mtp.py
index 294b3559c011..d17b579c1a71 100644
--- a/vllm/model_executor/models/step3p5_mtp.py
+++ b/vllm/model_executor/models/step3p5_mtp.py
@@ -4,7 +4,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
@@ -30,7 +30,7 @@
 class SharedHead(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
@@ -306,8 +306,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return loaded_params
 
     def _rewrite_spec_layer_name(self, spec_layer: int | None, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         """
         spec_layer_weight_names = [
diff --git a/vllm/model_executor/models/step_vl.py b/vllm/model_executor/models/step_vl.py
index 247078e2ec09..3f31547f9824 100644
--- a/vllm/model_executor/models/step_vl.py
+++ b/vllm/model_executor/models/step_vl.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""This is basically a copy from perception_models/core/vision_encoder/pe.py"""
+"""This is basically a copy from perception_models/core/vision_encoder/pe.py."""
 
 from collections.abc import Callable
 from functools import partial
diff --git a/vllm/model_executor/models/teleflm.py b/vllm/model_executor/models/teleflm.py
index bebd7bcaa924..b8235daacb31 100644
--- a/vllm/model_executor/models/teleflm.py
+++ b/vllm/model_executor/models/teleflm.py
@@ -23,47 +23,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
-import torch.nn as nn
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.models.llama import (
-    LlamaDecoderLayer,
     LlamaForCausalLM,
-    LlamaModel,
 )
 
 
-class TeleFLMModel(LlamaModel):
-    def __init__(
-        self,
-        *,
-        vllm_config: VllmConfig,
-        prefix: str = "",
-        layer_type: type[nn.Module] = LlamaDecoderLayer,
-    ):
-        super().__init__(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)
-        """
-        This implementation is based on the µScaling paper presented at  
-        the ICLR 2025 Workshop:  
-        NanoLM: An Affordable LLM Study Benchmark \
-        via Accurate Loss Prediction across Scales
-        by Yiqun Yao et al.  
-        Available at: https://openreview.net/forum?id=IwaPYg1SCA  
-        arXiv preprint: https://arxiv.org/abs/2304.06875
-        """
-        self.use_mup = self.config.use_mup
-        if self.use_mup:
-            self.input_mult = self.config.input_mult
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        embedding = self.embed_tokens(input_ids)
-        if self.use_mup:
-            embedding = embedding * self.input_mult
-        return embedding
-
-
 class TeleFLMForCausalLM(LlamaForCausalLM):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__(vllm_config=vllm_config, prefix=prefix)
diff --git a/vllm/model_executor/models/terratorch.py b/vllm/model_executor/models/terratorch.py
index 7181680fe833..ce04acc72d37 100644
--- a/vllm/model_executor/models/terratorch.py
+++ b/vllm/model_executor/models/terratorch.py
@@ -15,7 +15,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Wrapper around `Terratorch` models"""
+"""Wrapper around `Terratorch` models."""
 
 from collections import OrderedDict
 from collections.abc import Iterable, Mapping, Sequence
@@ -33,7 +33,7 @@
 from transformers import BatchFeature
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import ModalityData, MultiModalDataDict, MultiModalInput, mm_input
 from vllm.logger import init_logger
 from vllm.model_executor.layers.pooler import IdentityPooler
@@ -153,7 +153,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         # Dummy data is generated based on the 'input' section
         # defined in the HF configuration file
@@ -202,7 +202,9 @@ def apply(
         hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs
 
         with timing_ctx.record("apply_hf_processor"):
-            _, passthrough_data = self._get_hf_mm_data(mm_items)
+            passthrough_data = self._get_hf_mm_inputs(
+                mm_items, hf_processor_mm_kwargs
+            ).passthrough_data
             mm_processed_data = BatchFeature(
                 {
                     k: torch.as_tensor(v).unsqueeze(0)
diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py
index cbfe55195233..e19de6d4e891 100644
--- a/vllm/model_executor/models/transformers/__init__.py
+++ b/vllm/model_executor/models/transformers/__init__.py
@@ -14,14 +14,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Wrapper around `transformers` models"""
+"""Wrapper around `transformers` models."""
 
 from typing import TYPE_CHECKING
 
 import torch.nn.functional as F
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 
-from vllm.model_executor.models.transformers.base import VLLM_ATTN_ATTR, Base
+from vllm.model_executor.models.transformers.base import VLLM_ATTN_ATTR
 from vllm.model_executor.models.transformers.causal import CausalMixin
 from vllm.model_executor.models.transformers.legacy import LegacyMixin
 from vllm.model_executor.models.transformers.moe import MoEMixin
@@ -126,10 +126,10 @@ def vllm_mla_attention_forward(
 
 
 # Text only models
-class TransformersForCausalLM(CausalMixin, Base): ...
+class TransformersForCausalLM(CausalMixin): ...
 
 
-class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...
+class TransformersMoEForCausalLM(MoEMixin, CausalMixin): ...
 
 
 # Multimodal models
@@ -138,7 +138,7 @@ class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...
     info=MultiModalProcessingInfo,
     dummy_inputs=MultiModalDummyInputsBuilder,
 )
-class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...
+class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin): ...
 
 
 @MULTIMODAL_REGISTRY.register_processor(
@@ -146,16 +146,14 @@ class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...
     info=MultiModalProcessingInfo,
     dummy_inputs=MultiModalDummyInputsBuilder,
 )
-class TransformersMultiModalMoEForCausalLM(
-    MoEMixin, MultiModalMixin, CausalMixin, Base
-): ...
+class TransformersMultiModalMoEForCausalLM(MoEMixin, MultiModalMixin, CausalMixin): ...
 
 
 # Embedding models
-class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ...
+class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin): ...
 
 
-class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...
+class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin): ...
 
 
 @MULTIMODAL_REGISTRY.register_processor(
@@ -163,17 +161,17 @@ class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...
     info=MultiModalProcessingInfo,
     dummy_inputs=MultiModalDummyInputsBuilder,
 )
-class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ...
+class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin): ...
 
 
 # Sequence classification models
 class TransformersForSequenceClassification(
-    SequenceClassificationMixin, LegacyMixin, Base
+    SequenceClassificationMixin, LegacyMixin
 ): ...
 
 
 class TransformersMoEForSequenceClassification(
-    SequenceClassificationMixin, MoEMixin, Base
+    SequenceClassificationMixin, MoEMixin
 ): ...
 
 
@@ -183,7 +181,7 @@ class TransformersMoEForSequenceClassification(
     dummy_inputs=MultiModalDummyInputsBuilder,
 )
 class TransformersMultiModalForSequenceClassification(
-    SequenceClassificationMixin, MultiModalMixin, Base
+    SequenceClassificationMixin, MultiModalMixin
 ): ...
 
 
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 89c6486e739e..ec0f474ee3e7 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -29,13 +29,13 @@
 import transformers
 from packaging.version import Version
 from torch import nn
-from transformers import AutoModel, PretrainedConfig
+from transformers import AutoModel, PreTrainedConfig
 from transformers.conversion_mapping import (
     WeightRenaming,
     get_model_conversion_mapping,
 )
 
-from vllm.compilation.decorators import support_torch_compile
+from vllm.compilation.decorators import DynamicArgDims, support_torch_compile
 from vllm.config.utils import getattr_iter
 from vllm.distributed import get_pp_group, get_tp_group
 from vllm.distributed.utils import get_pp_indices
@@ -76,6 +76,7 @@
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    ShardId,
     WeightsMapper,
     make_empty_intermediate_tensors_factory,
     maybe_prefix,
@@ -110,10 +111,13 @@ class Base(
     SupportsEagle,
     SupportsEagle3,
 ):
-    embedding_modules = ["embed_tokens"]  # TODO transformers will have a util to get it
+    hf_to_vllm_mapper: WeightsMapper
+
+    # TODO transformers will have a util to get these
+    embedding_modules = {"embed_tokens": "input_embeddings"}
 
     def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
-        super().__init__()
+        nn.Module.__init__(self)
         logger.info("Using Transformers modeling backend.")
 
         self.vllm_config = vllm_config
@@ -193,8 +197,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         )
 
     def _patch_config(self):
-        """
-        Patch the config to ensure that the model is created correctly:
+        """Patch the config to ensure that the model is created correctly:
 
         - Sets the attention implementation to "vllm" so the attention instances from
         `_create_attention_instances` are used
@@ -249,12 +252,11 @@ def _pre_trained_model_classes(self) -> PreTrainedModelClasses:
     def _decorate_cls_for_torch_compile(
         self,
         cls: type["PreTrainedModel"],
-        dynamic_arg_dims: dict[str, int] | None,
+        dynamic_arg_dims: DynamicArgDims | None,
         enable_if: Callable[["VllmConfig"], bool],
         is_encoder: bool,
     ):
-        """
-        Decorate `cls` to indicate to vLLM that it supports torch compile.
+        """Decorate `cls` to indicate to vLLM that it supports torch compile.
 
         Args:
             cls: The PreTrainedModel class to decorate.
@@ -265,6 +267,7 @@ def _decorate_cls_for_torch_compile(
             enable_if: A function which takes in the vLLM config and returns whether
                 torch compile should be enabled for this class.
             is_encoder: Whether the class being decorated is an encoder.
+
         """
         logger.debug(
             "Decorating `%s` as %s for torch compile with dynamic_arg_dims of %s",
@@ -285,7 +288,7 @@ def _decorate_for_torch_compile(self):
         self._decorate_cls_for_torch_compile(
             cls=self._pre_trained_model_classes.decoder,
             # Applied to a PreTrainedModel so the batch dimension will exist
-            dynamic_arg_dims=dict[str, int](
+            dynamic_arg_dims=DynamicArgDims(
                 input_ids=1,  # shape: [1, seq_len]
                 inputs_embeds=1,  # shape: [1, seq_len, hidden_size]
                 position_ids=-1,  # shape: [1, seq_len] or [3, 1, seq_len] for mrope
@@ -295,8 +298,7 @@ def _decorate_for_torch_compile(self):
         )
 
     def _create_hf_to_vllm_mapper(self):
-        """
-        Create a WeightsMapper to map checkpoint weight names to module qualnames.
+        """Create a WeightsMapper to map checkpoint weight names to module qualnames.
 
         This handles:
 
@@ -305,9 +307,8 @@ def _create_hf_to_vllm_mapper(self):
         - Checkpoints saved with no base model prefix
         - Any quantization config specific mappings
         """
-        self.hf_to_vllm_mapper = WeightsMapper()
-        orig_to_new_renaming = self.hf_to_vllm_mapper.orig_to_new_renaming
-        orig_to_new_regex = self.hf_to_vllm_mapper.orig_to_new_regex
+        orig_to_new_renaming: list[WeightRenaming] = []
+        orig_to_new_regex: dict[re.Pattern, str | None] = {}
 
         for mapping in get_model_conversion_mapping(self.model):
             # Handle weights which have been renamed in Transformers
@@ -343,22 +344,23 @@ def _create_hf_to_vllm_mapper(self):
         nested_lm_head_pattern = re.compile(r"^model\.(.+\.)*(lm_head.+)")
         orig_to_new_regex[nested_lm_head_pattern] = r"\2"
 
+        self.hf_to_vllm_mapper = WeightsMapper(
+            orig_to_new_renaming=orig_to_new_renaming,
+            orig_to_new_regex=orig_to_new_regex,
+        )
+
         # Apply mapping to quantization config if needed
         self._maybe_apply_model_mapping()
 
     def _get_tie_word_embeddings(self):
-        """
-        Check if the model has tied word embeddings.
-        """
+        """Check if the model has tied word embeddings."""
         # Models created with Transformers v4 and v5 will store this in different places
         tie_word_embeddings_v4 = getattr(self.text_config, "tie_word_embeddings", False)
         tie_word_embeddings_v5 = getattr(self.config, "tie_word_embeddings", False)
         return tie_word_embeddings_v4 or tie_word_embeddings_v5
 
     def pipeline_parallel(self):
-        """
-        Apply the model's pipeline parallelization plan.
-        """
+        """Apply the model's pipeline parallelization plan."""
         if self.pp_group.world_size <= 1:
             return
 
@@ -377,8 +379,8 @@ def pipeline_parallel(self):
                 "modeling backend will infer the split from the layers of %s in order "
                 "of declaration and keep parameter-free modules on every rank. This "
                 "may fail if the model's structure is non-standard. %s",
-                type(self.model),
-                type(module),
+                type(self.model).__name__,
+                type(module).__name__,
                 tip,
             )
 
@@ -431,9 +433,9 @@ def pipeline_parallel(self):
     def _vocab_embeddings(self) -> set[nn.Embedding]:
         """The `nn.Embedding`s in `self.model` which hold a vocab table."""
 
-        def vocab_sizes(config: PretrainedConfig):
+        def vocab_sizes(config: PreTrainedConfig):
             for key, value in vars(config).items():
-                if isinstance(value, PretrainedConfig):
+                if isinstance(value, PreTrainedConfig):
                     yield from vocab_sizes(value)
                 elif "vocab_size" in key and isinstance(value, int):
                     yield value
@@ -476,7 +478,7 @@ def recursive_replace(self):
                 "backend will shard the model the best it can during graph fusion and "
                 "replicate the rest. This may be suboptimal or fail if the model does "
                 "not fuse cleanly. %s",
-                type(self.model),
+                type(self.model).__name__,
                 tip,
             )
 
@@ -487,6 +489,8 @@ def recursive_replace(self):
 
         vocab_embeddings = self._vocab_embeddings()
 
+        orig_to_new_stacked: dict[str, tuple[str, ShardId]] = {}
+
         def register_fusion(fuser: BaseFuser, prefix: str, module: nn.Module):
             """Register a fused layer's mappings just before it is built."""
             self.fusers.setdefault(prefix, []).append(fuser)
@@ -504,8 +508,7 @@ def register_fusion(fuser: BaseFuser, prefix: str, module: nn.Module):
                     )
                 self.attention_fusers[index] = (prefix, fuser)
 
-            orig_to_new_stacked = fuser.orig_to_new_stacked(prefix)
-            self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)
+            orig_to_new_stacked.update(fuser.orig_to_new_stacked(prefix))
 
             packed_modules_mapping = fuser.packed_modules_mapping
             self.packed_modules_mapping.update(packed_modules_mapping)
@@ -570,6 +573,8 @@ def _recursive_replace(module: nn.Module, prefix: str):
 
         _recursive_replace(self.model, prefix="model")
 
+        self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked)
+
     def _create_attention_instances(self):
         """Create `Attention` instances to inform KV cache allocation."""
         text_config = self.text_config
@@ -595,8 +600,11 @@ def _create_attention_instances(self):
 
         for i in range(start, end):
             if i not in self.attention_fusers:
-                in_range = layer_types and i < len(layer_types)
-                layer = f"{i} ({layer_types[i]})" if in_range else str(i)
+                layer = (
+                    f"{i} ({layer_types[i]})"
+                    if layer_types and i < len(layer_types)
+                    else str(i)
+                )
                 raise ValueError(
                     f"Layer {layer} does not dispatch through the Transformers "
                     "attention interface and vLLM has no other way to handle it."
@@ -697,8 +705,7 @@ def _get_attn_cls(self) -> type[AttentionLayerBase]:
         return Attention
 
     def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None):
-        """
-        If a `parameter` is on the `meta` device, then its parent
+        """If a `parameter` is on the `meta` device, then its parent
         `module` is the original module created by:
 
         ```python
diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py
index 29e780557cb7..e7dba82b0f80 100644
--- a/vllm/model_executor/models/transformers/causal.py
+++ b/vllm/model_executor/models/transformers/causal.py
@@ -27,13 +27,15 @@
 from vllm.model_executor.models.interfaces_base import VllmModelForTextGeneration
 from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
 
+from .base import Base
+
 if TYPE_CHECKING:
     import torch
 
     from vllm.config import VllmConfig
 
 
-class CausalMixin(VllmModelForTextGeneration):
+class CausalMixin(VllmModelForTextGeneration, Base):
     def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         # Skip VllmModelForTextGeneration.__init__ and call the next class in MRO
         super(VllmModelForTextGeneration, self).__init__(
@@ -65,7 +67,6 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
 
     def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]:
         """A thin wrapper around `Base.load_weights` to handle the lm_head bias."""
-
         lm_head_bias = set()
 
         def auto_load_lm_head_bias(weights):
diff --git a/vllm/model_executor/models/transformers/fusers/attention.py b/vllm/model_executor/models/transformers/fusers/attention.py
index fc16527f15e3..d8d324536d27 100644
--- a/vllm/model_executor/models/transformers/fusers/attention.py
+++ b/vllm/model_executor/models/transformers/fusers/attention.py
@@ -5,6 +5,7 @@
 import ast
 import inspect
 import textwrap
+from collections.abc import Callable
 from dataclasses import dataclass
 from functools import cache
 from typing import TYPE_CHECKING, ClassVar
@@ -36,16 +37,16 @@ def _is_interface_lookup(node: ast.expr | None) -> bool:
 
 
 @cache
-def interface_call(cls: type[nn.Module]) -> ast.Call | None:
-    """The attention interface call in `cls.forward`, if it makes exactly one."""
+def interface_call(forward: Callable) -> ast.Call | None:
+    """The attention interface call in `forward`, if it makes exactly one."""
     try:
-        source = inspect.getsource(inspect.unwrap(cls.forward))
+        source = inspect.getsource(inspect.unwrap(forward))
         tree = ast.parse(textwrap.dedent(source))
     except (AttributeError, OSError, SyntaxError, TypeError):
         return None
 
     # Find the names of the local variables that read from the interface lookup
-    names = set()
+    names: set[str] = set()
     for node in ast.walk(tree):
         if not isinstance(node, (ast.Assign, ast.AnnAssign)):
             continue
@@ -97,7 +98,7 @@ def info(self, name: str) -> str:
     def match(
         cls, graph: fx.Graph | None, module: nn.Module
     ) -> "AttentionFuser | None":
-        if (call := interface_call(type(module))) is None:
+        if (call := interface_call(type(module).forward)) is None:
             return None
         scaling = [kw.value for kw in call.keywords if kw.arg == "scaling"]
         scale_expr = scaling[0] if len(scaling) == 1 else None
@@ -120,6 +121,7 @@ def fuse(
         self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
     ) -> nn.Module:
         if (sinks := self.sinks(module)) is not None:
+            assert isinstance(self.s_aux_expr, ast.Attribute)
             size = sinks.numel() // vllm_config.parallel_config.tensor_parallel_size
             device = vllm_config.device_config.device
             data = torch.empty(size, dtype=sinks.dtype, device=device)
diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py
index ad24f2542a47..937b70c5ed45 100644
--- a/vllm/model_executor/models/transformers/fusers/base.py
+++ b/vllm/model_executor/models/transformers/fusers/base.py
@@ -130,11 +130,14 @@ class StackedFuser(RewriteFuser):
     """A fuser that merges sibling projections into one stacked linear and
     rewrites the forward to call it."""
 
-    merged_name: ClassVar[str]
-    """Attribute name of the merged module created by `update_attrs`."""
     merged_cls_name: ClassVar[str]
     """Name of the vLLM class the merged projection becomes (for logging)."""
 
+    @property
+    @abstractmethod
+    def merged_name(self) -> str:
+        """Attribute name of the merged module created by `update_attrs`."""
+
     def info(self, name: str) -> str:
         sources = " + ".join(shard for shard, _ in self.shards)
         return (
diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py
index fa4b34750c5f..96a4a0b87b3d 100644
--- a/vllm/model_executor/models/transformers/fusers/mla.py
+++ b/vllm/model_executor/models/transformers/fusers/mla.py
@@ -24,6 +24,7 @@
     upstream_linear,
 )
 from vllm.model_executor.models.transformers.utils import (
+    Style,
     log_replacement,
     replace_linear_class,
 )
@@ -128,30 +129,26 @@ class MLAFuser(StackedFuser):
     merged_name: ClassVar[str] = "fused_qkv_a_proj"
     merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear"
 
-    @property
-    def has_q_lora(self) -> bool:
-        return self.q_a_proj_name is not None
-
     def info(self, name: str) -> str:
         info_str = (
             f"Fused: {name} ({self.source_cls}) -> MLAAttention (attention interface)"
         )
-        if self.has_q_lora:
+        if self.q_a_proj_name is not None:
             info_str += "; " + super().info(name).removeprefix("Fused: ")
         return info_str
 
     @property
     def shards(self) -> list[tuple[str, ShardId]]:
         """`q_a_proj` and `kv_a_proj_with_mqa` stack into one down-projection."""
-        if self.has_q_lora:
-            return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)]
-        return []
+        if self.q_a_proj_name is None:
+            return []
+        return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)]
 
     @property
     def packed_modules_mapping(self) -> dict[str, list[str]]:
-        if self.has_q_lora:
-            return super().packed_modules_mapping
-        return {}
+        if self.q_a_proj_name is None:
+            return {}
+        return super().packed_modules_mapping
 
     @classmethod
     def match(cls, graph: fx.Graph, module: nn.Module) -> "MLAFuser | None":
@@ -234,10 +231,10 @@ def update_forward(self, module: nn.Module) -> None:
         Bypass the KV expansion method so the compressed latent reaches the `vllm_mla`
         attention interface unexpanded."""
         funcdef, fn = recover_forward(type(module))
-        if self.has_q_lora:
+        if (q_a_proj_name := self.q_a_proj_name) is not None:
             # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.
             # The fused call is inserted at the top-level statement preceding both.
-            names = [self.q_a_proj_name, self.kv_a_proj_name]
+            names = [q_a_proj_name, self.kv_a_proj_name]
             calls = self._unguarded_calls(funcdef, names)
             if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]):
                 raise ValueError("down-projections read different inputs")
@@ -258,7 +255,7 @@ def update_forward(self, module: nn.Module) -> None:
     def update_attrs(self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"):
         quant_config = vllm_config.quant_config
 
-        def replace_linear_by_name(name: str, style: str):
+        def replace_linear_by_name(name: str, style: Style):
             linear = module.get_submodule(name)
             replacement = replace_linear_class(
                 linear, style, quant_config, prefix=maybe_prefix(prefix, name)
@@ -266,8 +263,10 @@ def replace_linear_by_name(name: str, style: str):
             setattr(module, name, replacement)
             log_replacement(maybe_prefix(prefix, name), linear, replacement)
 
-        if self.has_q_lora:
-            q_a = module.get_submodule(self.q_a_proj_name)
+        if (q_a_proj_name := self.q_a_proj_name) is not None:
+            # `match` sets the q-LoRA names together, or none of them
+            assert self.q_b_proj_name is not None
+            q_a = module.get_submodule(q_a_proj_name)
             kv_a = module.get_submodule(self.kv_a_proj_name)
             merged = MergedColumnParallelLinear(
                 input_size=q_a.in_features,
@@ -280,7 +279,7 @@ def replace_linear_by_name(name: str, style: str):
             )
             logger.debug(
                 "%s: %s, %s: %s -> %s: %s",
-                self.q_a_proj_name,
+                q_a_proj_name,
                 q_a,
                 self.kv_a_proj_name,
                 kv_a,
@@ -289,10 +288,11 @@ def replace_linear_by_name(name: str, style: str):
             )
             setattr(module, self.merged_name, merged)
             # The rewritten forward calls the merged projection instead.
-            delattr(module, self.q_a_proj_name)
+            delattr(module, q_a_proj_name)
             delattr(module, self.kv_a_proj_name)
             replace_linear_by_name(self.q_b_proj_name, "colwise")
         else:
+            assert self.q_proj_name is not None
             replace_linear_by_name(self.kv_a_proj_name, "replicate")
             replace_linear_by_name(self.q_proj_name, "colwise")
 
diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py
index 47d20085a2c7..f7972f03a6e7 100644
--- a/vllm/model_executor/models/transformers/fusers/rms_norm.py
+++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py
@@ -51,7 +51,7 @@ def _is_squared(node: object, x: fx.Node) -> bool:
 
 
 def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None:
-    """eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
+    """Eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
     add = peel(rsqrt.args[0])
     if not is_op(add, "add"):
         return None
@@ -216,7 +216,7 @@ def _eps_source(
             finally:
                 for name, value in candidates.items():
                     setattr(module, name, value)
-        if (name := markers.get(marked)) is not None:
+        if marked is not None and (name := markers.get(marked)) is not None:
             return name, None
         logger.debug_once(
             "%s does not hold its eps (%s) in an attribute. Every instance in this "
@@ -251,18 +251,18 @@ def fuse(
     ) -> nn.Module:
         """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""
         weight = getattr(module, "weight", None)
-        has_weight = weight is not None
-        hidden_size = weight.size(0) if has_weight else 0
+        hidden_size = weight.size(0) if weight is not None else 0
         eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps
         if not isinstance(eps, (int, float)):
             # If eps was not detected, match torch behaviour.
-            dtype = weight.dtype if has_weight else vllm_config.model_config.dtype
+            config_dtype = vllm_config.model_config.dtype
+            dtype = weight.dtype if weight is not None else config_dtype
             eps = torch.finfo(dtype).eps
         if self.zero_centered:
             return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
         return TPAwareRMSNorm(
             hidden_size=hidden_size,
             eps=eps,
-            has_weight=has_weight,
-            dtype=weight.dtype if has_weight else None,
+            has_weight=weight is not None,
+            dtype=weight.dtype if weight is not None else None,
         )
diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py
index 332bd5414fc3..cb630866a390 100644
--- a/vllm/model_executor/models/transformers/fx_utils.py
+++ b/vllm/model_executor/models/transformers/fx_utils.py
@@ -15,17 +15,19 @@
 import textwrap
 from collections.abc import Callable
 from itertools import chain
+from typing import Any
 from unittest import mock
 
 import torch
 from torch import fx, nn
 from torch.nn import functional as F
+from typing_extensions import TypeIs
 
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
-_UNKNOWN = object()
+_UNKNOWN: Any = object()
 """Sentinel meta value for proxies whose concrete value could not be inferred.
 Distinct from `None`, which is a valid concrete value (e.g. `attn_weights`)."""
 
@@ -34,7 +36,7 @@
 `call_module` nodes; meta execution must call modules for real."""
 
 
-def is_leaf_call(node: object) -> bool:
+def is_leaf_call(node: object) -> TypeIs[fx.Node]:
     """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface)."""
     return isinstance(node, fx.Node) and node.meta.get("leaf_call", False)
 
@@ -60,7 +62,7 @@ class _MetaProxy(fx.Proxy):
     executing each op on the meta values, so PyTorch's meta kernels are the
     single source of shape inference — no per-op rules."""
 
-    meta: object = _UNKNOWN
+    meta: Any = _UNKNOWN
 
     def __len__(self) -> int:
         if self.meta is not _UNKNOWN:
@@ -149,6 +151,7 @@ def meta_of(arg: object) -> object:
         if unknown:
             return _UNKNOWN
         if kind == "call_function":
+            assert callable(target)
             return target(*meta_args, **meta_kwargs)
         if kind == "call_method":
             receiver, *rest = meta_args
@@ -576,7 +579,7 @@ def _aliasing_reads(node: ast.expr) -> set[str]:
     would taint nearly every later name.
     """
     names: set[str] = set()
-    stack = [node]
+    stack: list[ast.AST] = [node]
     while stack:
         current = stack.pop()
         if isinstance(current, ast.Attribute) and current.attr in _METADATA_ATTRS:
@@ -646,7 +649,7 @@ def aliasing_names(
     """
     names = set(seed)
     assigns = [
-        node
+        (node, node.value)
         for node in ast.walk(funcdef)
         if isinstance(node, (ast.Assign, ast.AnnAssign))
         and node.value is not None
@@ -655,10 +658,10 @@ def aliasing_names(
     changed = True
     while changed:
         changed = False
-        for node in assigns:
-            if not _aliasing_reads(node.value) & names:
+        for node, value in assigns:
+            if not _aliasing_reads(value) & names:
                 continue
-            targets = getattr(node, "targets", None) or [node.target]
+            targets = node.targets if isinstance(node, ast.Assign) else [node.target]
             stores = {
                 child.id
                 for target in targets
@@ -770,7 +773,7 @@ def peel(node: object) -> object:
     return node
 
 
-def is_fn(node: object, target: Callable) -> bool:
+def is_fn(node: object, target: object) -> TypeIs[fx.Node]:
     """Is node `<target>()`."""
     return (
         isinstance(node, fx.Node)
@@ -779,17 +782,15 @@ def is_fn(node: object, target: Callable) -> bool:
     )
 
 
-def is_method(node: object, name: str) -> bool:
+def is_method(node: object, name: str) -> TypeIs[fx.Node]:
     """Is node `.<name>()`."""
     return (
         isinstance(node, fx.Node) and node.op == "call_method" and node.target == name
     )
 
 
-def is_op(node: object, name: str) -> bool:
-    """
-    Is node `torch.<name>()`, `F.<name>()`, `operator.<name>()`, or `Tensor.<name>()`.
-    """
+def is_op(node: object, name: str) -> TypeIs[fx.Node]:
+    """Is node `<mod>.<name>()` for torch, F, operator, or Tensor."""
     return any(
         is_fn(node, getattr(module, name, None)) for module in (torch, F, operator)
     ) or (hasattr(torch.Tensor, name) and is_method(node, name))
diff --git a/vllm/model_executor/models/transformers/legacy.py b/vllm/model_executor/models/transformers/legacy.py
index 9c8a5c45b6f1..7a256f29d5be 100644
--- a/vllm/model_executor/models/transformers/legacy.py
+++ b/vllm/model_executor/models/transformers/legacy.py
@@ -20,35 +20,38 @@
 
 import torch
 
+from vllm.model_executor.models.utils import WeightsMapper
 from vllm.sequence import IntermediateTensors
 
+from .base import Base
+
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
 
 
-class LegacyMixin:
+class LegacyMixin(Base):
     def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         super().__init__(vllm_config=vllm_config, prefix=prefix)
 
-        # Drop unsupported/unwanted output embeddings layers.
-        self.hf_to_vllm_mapper.orig_to_new_prefix.update(
-            {
+        self.hf_to_vllm_mapper |= WeightsMapper(
+            # Drop unsupported/unwanted output embeddings layers.
+            orig_to_new_prefix={
                 "model.lm_head.": None,
                 "model.predictions.": None,
                 "model.qa_outputs.": None,
                 "model.embeddings_project.": None,
                 "model.discriminator_predictions.": None,
-            }
+            },
+            orig_to_new_substr={
+                # Some encoder models have the position_ids buffer in the checkpoint.
+                # vLLM always passes position_ids as an argument, so drop the buffer.
+                "position_ids": None,
+                # Some encoder models have the bias of the final classifier layer in
+                # the checkpoint. vLLM does not use this bias, so drop it.
+                "score.bias": None,
+            },
         )
 
-        # Some encoder models have the position_ids buffer in the checkpoint. vLLM will
-        # always pass position_ids as an argument, so we drop the  buffer if it exists.
-        self.hf_to_vllm_mapper.orig_to_new_substr["position_ids"] = None
-
-        # Some encoder models have the bias of the final classifier layer in the
-        # checkpoint. vLLM does not use this bias, so we drop it if it exists.
-        self.hf_to_vllm_mapper.orig_to_new_substr["score.bias"] = None
-
         # roberta-like models an extra padding in positions.
         # FIXME(Isotr0py): This is quite hacky for roberta edge case,
         # we should find a better way to handle this.
@@ -61,6 +64,7 @@ def forward(
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
     ) -> torch.Tensor | IntermediateTensors:
         if self.is_roberta:
             # RoBERTa positions start at padding_idx + 1.
@@ -72,4 +76,5 @@ def forward(
             positions=positions,
             intermediate_tensors=intermediate_tensors,
             inputs_embeds=inputs_embeds,
+            **kwargs,
         )
diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py
index 480bc7f66e8d..9f63e7333670 100644
--- a/vllm/model_executor/models/transformers/moe.py
+++ b/vllm/model_executor/models/transformers/moe.py
@@ -41,6 +41,7 @@
 from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op
 
+from .base import Base
 from .utils import log_replacement, maybe_per_layer
 
 if TYPE_CHECKING:
@@ -66,14 +67,15 @@ def __init__(self, *args, moe_state: TransformersMoEState, **kwargs):
         self._moe_state = moe_state
         self._moe_state.is_sequence_parallel = self.moe_config.is_sequence_parallel
 
-    def forward(
+    def forward(  # type: ignore[override]
         self,
         hidden_states: torch.Tensor,
         topk_ids: torch.Tensor,
         topk_weights: torch.Tensor,
         **kwargs: Any,
     ) -> torch.Tensor:
-        """In Transformers `experts.forward` will have this signature.
+        """In Transformers `experts.forward` will have this signature, which is why
+        it deliberately does not match `MoERunner.forward`.
 
         We discard any extra kwargs because we cannot use them here."""
         # Note: we need to forward through a custom op so the topk_ids
@@ -126,10 +128,18 @@ def _transformers_moe_forward_fake(
 
 class TransformersRoutedExperts(RoutedExperts):
     def get_expert_mapping(
-        self, include_fused: bool = False
+        self,
+        ckpt_gate_proj_name: str | None = None,
+        ckpt_down_proj_name: str | None = None,
+        ckpt_up_proj_name: str | None = None,
+        include_fused: bool = False,
     ) -> list[tuple[str, str, int, str]]:
-        common_names = ("gate_proj", "down_proj", "up_proj")
-        common_map = super().get_expert_mapping(*common_names, include_fused)
+        common_map = super().get_expert_mapping(
+            ckpt_gate_proj_name or "gate_proj",
+            ckpt_down_proj_name or "down_proj",
+            ckpt_up_proj_name or "up_proj",
+            include_fused,
+        )
         mixtral_map = super().get_expert_mapping("w1", "w2", "w3", include_fused)
         if not include_fused:
             return common_map + mixtral_map
@@ -138,7 +148,13 @@ def get_expert_mapping(
         return common_fused + mixtral_fused + common_unfused + mixtral_unfused
 
 
-class MoEMixin(MixtureOfExperts):
+class MoEMixin(MixtureOfExperts, Base):
+    mlp_layers: list[nn.Module]
+    """MoE blocks whose experts were replaced, for the `MixtureOfExperts` methods."""
+
+    # Redeclared because `recursive_replace` assigns it below the first read
+    num_local_physical_experts: int
+
     def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         self.check_version("5.0.0", "MoE models support")
         # Skip MixtureOfExperts.__init__ and call the next class in MRO
@@ -197,7 +213,8 @@ def recursive_replace(self):
 
         # Dtype the router computes in, if it is not the activation dtype.
         config_router_dtype = getattr(text_config, "moe_router_dtype", None)
-        config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype)
+        if config_router_dtype is not None:
+            config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype)
 
         # Grouped topk routing kwargs
         num_expert_group = getattr(text_config, "n_group", None)
@@ -222,7 +239,7 @@ def recursive_replace(self):
         # MixtureOfExperts mixin settings
         ep_size = get_ep_group().world_size
 
-        self.mlp_layers = []  # Used for MixtureOfExperts methods
+        self.mlp_layers = []
         self.moe_layers = []
         self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
         self.num_logical_experts = num_experts
@@ -321,9 +338,10 @@ def _recursive_replace(module: nn.Module, prefix: str):
                         if shared_experts is not None:
                             hf_shared = shared_experts.shared_experts
                             glu_fuser = get_fuser(hf_shared, GLUFuser)
-                            down_name = glu_fuser and glu_fuser.down_name
-                            if down_name is not None:
-                                shared_down_projs.append((hf_shared, down_name))
+                            if glu_fuser is not None:
+                                down_name = glu_fuser.down_name
+                                if down_name is not None:
+                                    shared_down_projs.append((hf_shared, down_name))
                         # Prefer config, otherwise read it from fuser.
                         router_dtype = config_router_dtype or fuser.router_dtype
                         gate = fuser.gate(moe_block, prefix, router_dtype)
@@ -374,7 +392,9 @@ def custom_routing_function(
                             # Handle all gather in expert parallel
                             if topk_ids.size(0) != hidden_states.size(0):
                                 dp_metadata = get_forward_context().dp_metadata
+                                assert dp_metadata is not None
                                 sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+                                assert sizes is not None
                                 is_sp = moe_state.is_sequence_parallel
                                 group = get_ep_group() if is_sp else get_dp_group()
                                 assert sizes[group.rank_in_group] == topk_ids.shape[0]
diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py
index 1cdbadb34b2a..cecb4f564a98 100644
--- a/vllm/model_executor/models/transformers/multimodal.py
+++ b/vllm/model_executor/models/transformers/multimodal.py
@@ -17,17 +17,23 @@
 """Transformers modeling backend mixin for multi-modal models."""
 
 from collections import defaultdict
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from contextlib import ExitStack, contextmanager
 from typing import TYPE_CHECKING, Any
 
 import torch
 import transformers
 from packaging.version import Version
+from transformers.utils.generic import ModelOutput
 
 from vllm.compilation.decorators import should_torch_compile_mm_encoder
 from vllm.config.utils import getattr_iter
-from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input
+from vllm.inputs import (
+    MultiModalDataBuiltins,
+    MultiModalDataDict,
+    MultiModalInput,
+    mm_input,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
@@ -62,11 +68,13 @@
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.torch_utils import async_tensor_h2d
 
+from .base import Base
+
 if TYPE_CHECKING:
     from transformers import BatchFeature, PreTrainedModel
 
     from vllm.config import VllmConfig
-    from vllm.config.multimodal import BaseDummyOptions
+    from vllm.config.multimodal import MultiModalDummyOptions
 
 logger = init_logger(__name__)
 
@@ -148,7 +156,7 @@ def get_max_audio_tokens(self) -> int:
     def get_max_image_tokens(self) -> int:
         width, height = self.get_image_size_with_most_features()
         processor = self.get_hf_processor()
-        multimodal_config = self.ctx.model_config.multimodal_config
+        multimodal_config = self.ctx.model_config.get_multimodal_config()
         mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {}
         mm_tokens = processor._get_num_multimodal_tokens(
             image_sizes=([height, width],), **mm_processor_kwargs
@@ -186,26 +194,25 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, "BaseDummyOptions"],
+        mm_options: "MultiModalDummyOptions",
     ) -> MultiModalDataDict:
-        data: MultiModalDataDict = {}
+        data = MultiModalDataBuiltins()
         if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)):
             sampling_rate = self.info._get_audio_sampling_rate()
             sub = self.info._get_audio_processor()
             chunk_length = getattr(sub, "chunk_length", None) if sub else None
             if chunk_length is None:
                 chunk_length = 30
-            audio_len = int(chunk_length * sampling_rate)
             data["audio"] = self._get_dummy_audios(
-                length=audio_len,
+                length=int(chunk_length * sampling_rate),
                 num_audios=num_audios,
                 overrides=mm_options.get("audio"),
             )
         if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)):
-            target_width, target_height = self.info.get_image_size_with_most_features()
+            width, height = self.info.get_image_size_with_most_features()
             data["image"] = self._get_dummy_images(
-                width=target_width,
-                height=target_height,
+                width=width,
+                height=height,
                 num_images=num_images,
                 overrides=mm_options.get("image"),
             )
@@ -219,17 +226,6 @@ class _MultiModalProcessorBase(BaseMultiModalProcessor[MultiModalProcessingInfo]
     Subclasses add the strategy for locating placeholders in the prompt.
     """
 
-    def _get_hf_mm_data(
-        self,
-        mm_items: MultiModalDataItems,
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        """Rename the parser's `audios` key to the `audio` argument HF audio
-        processors take."""
-        processor_data, passthrough_data = super()._get_hf_mm_data(mm_items)
-        if self.info._is_audio_model() and "audios" in processor_data:
-            processor_data["audio"] = processor_data.pop("audios")
-        return processor_data, passthrough_data
-
     def _get_modality_field_names(self, modality: str) -> set[str]:
         """Names of the fields the sub-processor for `modality` produces."""
         # TODO: use else branch only once huggingface/transformers#44394 lands.
@@ -352,6 +348,9 @@ def _get_mm_fields_config(
 
         return mm_fields
 
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _unpad_audios(
         self,
         hf_inputs: "BatchFeature",
@@ -367,7 +366,9 @@ def _unpad_audios(
         itself is to process it by itself.
         """
         audios = mm_data.get("audio")
-        if not audios or len({len(audio) for audio in audios}) == 1:
+        if not isinstance(audios, Iterable) or not audios:
+            return
+        if len({len(audio) for audio in audios}) == 1:
             return
 
         alone = [
@@ -452,7 +453,8 @@ def _get_mm_token_ids(self, modality: str) -> list[int]:
         token_id = getattr(processor, names[0], getattr_iter(config, names))
         if token_id is None:
             token = getattr(processor, f"{modality}_token", None)
-            token_id = info.get_tokenizer().get_vocab().get(token)
+            if token is not None:
+                token_id = info.get_tokenizer().get_vocab().get(token)
         if token_id is None:
             raise ValueError(
                 f"Cannot find {modality}_token_id on processor or model config"
@@ -592,21 +594,31 @@ def apply(
             # work for Transformers. The vision path has logic tied to
             # `mm_tokens_per_modality` in _apply_vision()
             hf_processor_mm_kwargs = {
-                **hf_processor_mm_kwargs,
+                # vLLM needs the untruncated sequence to keep placeholder
+                # tokens aligned. Note that the text inputs are just dummy
+                # text, not the original prompt. The original prompt is
+                # already tokenized by the renderer.
+                "truncation": False,
                 # HF processors only accept text, and the decoded string already
                 # contains any special tokens, so don't let them be added again
                 # (the HF processor call disables `add_special_tokens`).
                 "add_special_tokens": False,
+                **hf_processor_mm_kwargs,
             }
 
-            valid_mm_items = mm_items.select(
-                {k for k, c in mm_items.get_all_counts().items() if c > 0}
+            processor_data, _, passthrough_data = self._get_hf_mm_inputs(
+                mm_items, hf_processor_mm_kwargs
             )
-            processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+
+            # The real prompt is fed to the processor below and the placeholder
+            # ranges are read back out of its output, so the dummy text meant for
+            # media-only callers is discarded here.
+            processor_data.pop("text", None)
+            has_mm_data = any(processor_data.values())
 
             # Ask which modality owns each token of the expanded prompt, which
             # is the only marking of the tokens an expansion adds around an item
-            if any(processor_data.values()):
+            if has_mm_data:
                 processor_data = {**processor_data, "return_mm_token_type_ids": True}
 
             try:
@@ -616,7 +628,7 @@ def apply(
                     hf_processor_mm_kwargs,
                 )
             except ValueError:
-                if any(processor_data.values()):
+                if has_mm_data:
                     raise
                 # Some processors reject a prompt holding placeholders with
                 # no data to go with them, so tokenize it without them
@@ -720,13 +732,15 @@ def get_target_token_ids(token: str | int | Sequence[int]) -> list[int]:
         for modality, items in out_mm_kwargs.items():
             # Popped so they are neither cached nor sent to the model; the updates
             # they produce are cached alongside the item instead
-            replacements = [
-                PromptUpdateDetails.select_token_id(
-                    (ids := item.pop(f"{modality}_replacement_ids").data).tolist(),
-                    _get_embed_token_id(ids),
+            replacements = []
+            for item in items:
+                ids = item.pop(f"{modality}_replacement_ids").data
+                assert isinstance(ids, torch.Tensor)
+                replacements.append(
+                    PromptUpdateDetails.select_token_id(
+                        ids.tolist(), _get_embed_token_id(ids)
+                    )
                 )
-                for item in items
-            ]
             token = getattr(hf_processor, f"{modality}_token")
             updates.append(
                 PromptReplacement(
@@ -775,7 +789,7 @@ def _get_num_patches_per_image(
     ) -> list[int] | None:
         """Ask the HF processor how many rows of image data each image produces."""
         images = mm_data.get("images")
-        if not images:
+        if not isinstance(images, Iterable) or not images:
             return None
         try:
             sizes = [(image.height, image.width) for image in images]
@@ -786,31 +800,37 @@ def _get_num_patches_per_image(
         except (AttributeError, KeyError, TypeError):
             return None
 
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> "BatchFeature":
-        mm_counts = mm_items.get_all_counts()
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
+        )
 
-        valid_mm_items = mm_items.select({k for k, c in mm_counts.items() if c > 0})
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        if not hf_data:
+            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
 
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_counts)
+        prompt_text = hf_data.pop("text")
+        assert isinstance(prompt_text, str)
 
         # Ask for the replacement each placeholder expands to, and record it as
         # per-item fields: its token ids, and the tokens or patches behind them
-        if has_mm_data := any(mm_data.values()):
-            mm_data = {**mm_data, "return_text_replacement_offsets": True}
+        if has_mm_data := any(hf_data.values()):
+            hf_data = {**hf_data, "return_text_replacement_offsets": True}
 
         try:
             hf_inputs = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(**hf_processor_mm_kwargs),
-                dict(text=prompt_text, **mm_data),
-                hf_processor_mm_kwargs,
+                self.info.get_hf_processor(**hf_kwargs),
+                dict(text=prompt_text, **hf_data),
+                hf_kwargs,
             )
         except ValueError:
-            if any(mm_data.values()):
+            if any(hf_data.values()):
                 raise
             # Some processors reject a prompt holding placeholders with
             # no data to go with them, so tokenize it without them
@@ -819,7 +839,7 @@ def _apply_hf_processor_main(
                 dict(input_ids=[tokenizer.encode(prompt_text)]), tensor_type="pt"
             )
         self._unpad_images(hf_inputs)
-        self._unpad_audios(hf_inputs, mm_data, hf_processor_mm_kwargs)
+        self._unpad_audios(hf_inputs, hf_data, hf_kwargs)
 
         # Drop the inputs the model would reject
         hf_inputs.pop("mm_token_type_ids", None)
@@ -839,9 +859,10 @@ def _apply_hf_processor_main(
                     "so it can be fixed, and install `transformers<5.15.0` in the "
                     "meantime to locate placeholders in the expanded prompt instead."
                 )
-            hf_inputs.update(passthrough_data)
             hf_inputs.pop("input_ids", None)
-            return hf_inputs
+            return self._finalize_hf_mm_data(
+                hf_data, hf_kwargs, passthrough_data, hf_inputs
+            )
 
         tokenizer = self.info.get_tokenizer()
         replacements = defaultdict[str, list[list[int]]](list)
@@ -859,7 +880,7 @@ def _apply_hf_processor_main(
             )
             if modality == "image":
                 hf_inputs["num_image_patches"] = self._get_num_image_patches(
-                    hf_inputs, mm_data, len(seqs)
+                    hf_inputs, hf_data, len(seqs)
                 )
             elif modality == "audio":
                 counts = []
@@ -868,10 +889,11 @@ def _apply_hf_processor_main(
                     counts.append(int(ids.eq(_get_embed_token_id(ids)).sum()))
                 hf_inputs["num_audio_tokens"] = torch.tensor(counts)
 
-        hf_inputs.update(passthrough_data)
         hf_inputs.pop("input_ids")
 
-        return hf_inputs
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, hf_inputs
+        )
 
 
 # From this version on, a processor reporting no offsets is an error rather than a
@@ -883,7 +905,7 @@ def _apply_hf_processor_main(
 )
 
 
-class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
+class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base):
     def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         # Skip SupportsMRoPE.__init__ and call the next class in MRO
         super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)
@@ -915,7 +937,7 @@ def _mark_model_components(self, vllm_config: "VllmConfig"):
         if model_config.skip_tokenizer_init:
             # Determining the supported modalities needs the HF processor, which in
             # turn needs a tokenizer
-            mm_config = model_config.multimodal_config
+            mm_config = model_config.get_multimodal_config()
             if mm_config.mm_encoder_only or any(
                 mm_config.get_limit_per_prompt(modality) == 0
                 for modality in encoder_classes
@@ -953,9 +975,8 @@ def _mark_model_components(self, vllm_config: "VllmConfig"):
             yield
 
     def _decorate_for_torch_compile(self):
-        """
-        Decorate the model's decoder and encoder classes to indicate to vLLM that they
-        support torch compile if `can_enable_torch_compile` and
+        """Decorate the model's decoder and encoder classes to indicate to vLLM
+        that they support torch compile if `can_enable_torch_compile` and
         `should_torch_compile_mm_encoder` are True respectively.
         """
         super()._decorate_for_torch_compile()
@@ -1013,7 +1034,6 @@ def get_language_model(self) -> torch.nn.Module:
         """Transformers modeling backend multimodal classes do not contain a separate
         vLLM language model class. Therefore, in order to return a language model vLLM
         class, we use a wrapper to give `self` the same interface as a text model."""
-
         # Exclude self and object
         bases = self.__class__.mro()[1:-1]
         # Keep only classes defined in `vllm.model_executor.models.transformers`
@@ -1021,7 +1041,7 @@ def get_language_model(self) -> torch.nn.Module:
         # Exclude MultiModalMixin itself
         bases = [b for b in bases if b is not MultiModalMixin]
 
-        class LanguageModel(*bases):
+        class LanguageModel(*bases):  # type: ignore[misc]
             def __init__(self, multimodal_model):
                 # Don't call super().__init__() to avoid re-initialization
                 self.__dict__.update(multimodal_model.__dict__)
@@ -1031,9 +1051,7 @@ def __init__(self, multimodal_model):
         return LanguageModel(self)
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models"""
         for name in ("language_model", "text_model"):
             if getattr(self.model, name, None) is not None:
                 return MultiModelKeys.from_string_field(language_model=f"model.{name}")
@@ -1197,7 +1215,7 @@ def _get_image_features(self, pixel_values: torch.Tensor, **kwargs) -> Any:
         # from `self.get_image_features`
         if isinstance(features, tuple):
             return features[0]
-        if isinstance(features, dict):
+        if isinstance(features, ModelOutput):
             return features.pooler_output
         return features
 
diff --git a/vllm/model_executor/models/transformers/pooling.py b/vllm/model_executor/models/transformers/pooling.py
index f4fa4b496f23..79594387b740 100644
--- a/vllm/model_executor/models/transformers/pooling.py
+++ b/vllm/model_executor/models/transformers/pooling.py
@@ -19,6 +19,7 @@
 from typing import TYPE_CHECKING
 
 import torch
+import torch.nn as nn
 from transformers import AutoModelForSequenceClassification
 
 from vllm.config.utils import getattr_iter
@@ -26,11 +27,25 @@
 from vllm.model_executor.models.interfaces import SupportsCrossEncoding
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 
+from .base import Base
+
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
 
 
-class EmbeddingMixin(VllmModelForPooling):
+class ClassifierWithReshape(nn.Module):
+    """Token extraction has already been applied in `pooler.pooling`.
+
+    Add dim to match expected input shape of `classifier.forward`.
+    """
+
+    def forward(self, *args, **kwargs):
+        if len(args) > 0:
+            args = (args[0].unsqueeze(1), *args[1:])
+        return super().forward(*args, **kwargs)
+
+
+class EmbeddingMixin(VllmModelForPooling, Base):
     default_seq_pooling_type = "CLS"
 
     def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
@@ -45,7 +60,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         self.pooler = DispatchPooler.for_embedding(pooler_config)
 
 
-class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):
+class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling, Base):
     default_seq_pooling_type = "CLS"
 
     def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
@@ -83,18 +98,13 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
             )
         self.init_parameters(self.classifier, dtype=self.model_config.head_dtype)
 
-        class ClassifierWithReshape(self.classifier.__class__):
-            """
-            Token extraction has already been applied in `pooler.pooling`.
-            Add dim to match expected input shape of `classifier.forward`.
-            """
-
-            def forward(self, *args, **kwargs):
-                if len(args) > 0:
-                    args = (args[0].unsqueeze(1), *args[1:])
-                return super().forward(*args, **kwargs)
-
-        self.classifier.__class__ = ClassifierWithReshape
+        # Order `ClassifierWithReshape` ahead of the classifier's own class so that
+        # its `super().forward(...)` reaches the original implementation.
+        self.classifier.__class__ = type(
+            "ClassifierWithReshape",
+            (ClassifierWithReshape, type(self.classifier)),
+            {},
+        )
 
         self.pooler = DispatchPooler.for_seq_cls(
             pooler_config,
diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py
index 116dd7b5166b..0443f390c465 100644
--- a/vllm/model_executor/models/transformers/utils.py
+++ b/vllm/model_executor/models/transformers/utils.py
@@ -18,11 +18,10 @@
 
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
-from functools import lru_cache
 from itertools import chain
 from operator import attrgetter
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, TypeVar
+from typing import TYPE_CHECKING, Literal, TypeAlias, TypeVar
 
 import torch
 from torch import nn
@@ -49,16 +48,15 @@
 # Copied from `accelerate`
 @contextmanager
 def init_on_device_without_buffers(device: torch.device):
-    """
-    A context manager under which models are initialized with all
+    """A context manager under which models are initialized with all
     parameters on the specified device. However buffers are not
     initialized on specified device.
 
     Args:
         device (`torch.device`):
             Device to initialize all parameters on.
-    """
 
+    """
     old_register_parameter = nn.Module.register_parameter
 
     def register_empty_parameter(module, name, param):
@@ -71,7 +69,7 @@ def register_empty_parameter(module, name, param):
                 module._parameters[name].to(device), **kwargs
             )
 
-    tensor_constructors_to_patch = {}
+    tensor_constructors_to_patch: dict[str, Callable] = {}
 
     def patch_tensor_constructor(fn):
         def wrapper(*args, **kwargs):
@@ -114,17 +112,18 @@ def replace_linear_class(
     *,
     prefix: str = "",
 ) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear:
-    """
-    Replace nn.Linear with one of vLLM's tensor parallel linear classes.
+    """Replace nn.Linear with one of vLLM's tensor parallel linear classes.
 
     Args:
         linear: `nn.Linear` to be replaced.
         style: Tensor parallel style of the new linear, e.g. "colwise".
         quant_config: Quantization config for the new linear.
+        prefix: Module prefix of the new linear, used for quantization lookup.
+
     Returns:
         The new linear.
-    """
 
+    """
     if not isinstance(style, str):
         raise ValueError(f"Unsupported parallel style type {type(style)}, expected str")
 
@@ -147,8 +146,8 @@ def replace_linear_class(
     )
 
 
-TorchConv = nn.Conv2d | nn.Conv3d
-VllmConv = Conv2dLayer | Conv3dLayer
+TorchConv: TypeAlias = nn.Conv2d | nn.Conv3d
+VllmConv: TypeAlias = Conv2dLayer | Conv3dLayer
 
 
 def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
@@ -156,9 +155,11 @@ def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
 
     Args:
         conv: `nn.Conv2d` or `nn.Conv3d` to be replaced.
+
     Returns:
         The new `Conv2dLayer` or `Conv3dLayer`. If the conv module is not supported,
         returns the original conv module.
+
     """
     # vLLM does not handle non-zero padding modes
     if conv.padding_mode != "zeros":
@@ -210,17 +211,24 @@ class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbeddin
     `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding."""
 
 
-@lru_cache
+_rebased_embedding_classes: dict[type[nn.Embedding], type[VocabParallelEmbedding]] = {}
+
+
 def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbedding]:
     """Subclass `cls` so that `VocabParallelEmbedding` supersedes its `nn.Embedding`.
 
     Args:
         cls: The `nn.Embedding` subclass to rebase. Cached, so a given `cls` always
             maps to the same class.
+
     Returns:
         The new class, to assign to `__class__` of an instance of `cls`.
+
     """
-    return type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {})
+    if cls not in _rebased_embedding_classes:
+        rebased_cls = type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {})
+        _rebased_embedding_classes[cls] = rebased_cls
+    return _rebased_embedding_classes[cls]
 
 
 def replace_embedding_class(
@@ -235,8 +243,10 @@ def replace_embedding_class(
         embedding: The `nn.Embedding` holding a vocab table.
         quant_config: Quantization config for the new embedding.
         prefix: Qualname of `embedding`, used to look up its quantization method.
+
     Returns:
         The module to install in `embedding`'s place.
+
     """
     # Tied tables are reached more than once, so don't rebase an already new embedding
     if isinstance(embedding, VocabParallelEmbedding):
@@ -272,10 +282,9 @@ def _recursive_replace(module: nn.Module, prefix: str):
             qual_name = maybe_prefix(prefix, child_name)
             # Replace modules as needed
             if isinstance(child_module, nn.Linear):
-                style = "replicate"
                 new_module = replace_linear_class(
                     child_module,
-                    style,
+                    "replicate",
                     quant_config,
                     prefix=qual_name,
                 )
@@ -320,8 +329,7 @@ def get_feature_request_tip(
 
 
 def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool:
-    """
-    Callable to be passed to `@support_torch_compile`'s `enable_if` argument.
+    """Callable to be passed to `@support_torch_compile`'s `enable_if` argument.
 
     Defaults to `True` but is disabled in the following situations:
 
diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py
index a394e5f02e05..a454ba596253 100644
--- a/vllm/model_executor/models/ultravox.py
+++ b/vllm/model_executor/models/ultravox.py
@@ -16,7 +16,7 @@
 from transformers.models.whisper import WhisperFeatureExtractor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import MulAndSilu, get_act_fn
@@ -29,6 +29,7 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
+    MultiModalKwargsItem,
     MultiModalKwargsItems,
     NestedTensors,
 )
@@ -41,6 +42,7 @@
     PromptUpdate,
     cached_encode,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.renderers import TokenizeParams
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.ultravox import UltravoxConfig
@@ -71,8 +73,7 @@
 
 
 class UltravoxAudioFeatureInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
     - b: batch size
     - n: number of chunks
     - t: Time frames (M)
@@ -95,8 +96,7 @@ class UltravoxAudioFeatureInputs(TensorSchema):
 
 
 class UltravoxAudioEmbeddingInputs(TensorSchema):
-    """
-    Dimensions:
+    """Dimensions:
     - b: batch size
     - na: number of audios
     - afs: audio feature size
@@ -114,6 +114,12 @@ class UltravoxAudioEmbeddingInputs(TensorSchema):
 )
 
 
+class UltravoxMultiModalDataParser(MultiModalDataParser):
+    embedding_fields = {
+        "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
+    }
+
+
 class UltravoxProcessingInfo(BaseProcessingInfo):
     def get_hf_processor(self, **kwargs: object) -> ProcessorMixin:
         config = self.ctx.model_config.hf_config
@@ -145,10 +151,11 @@ def get_default_tok_params(self) -> TokenizeParams:
     def get_data_parser(self):
         feature_extractor = self.get_feature_extractor()
 
-        return MultiModalDataParser(
+        return UltravoxMultiModalDataParser(
             target_sr=feature_extractor.sampling_rate,
             target_channels=self.get_target_channels(),
             expected_hidden_size=self._get_expected_hidden_size(),
+            allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
         )
 
     def get_target_channels(self) -> int:
@@ -169,7 +176,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         feature_extractor = self.info.get_feature_extractor()
 
@@ -177,86 +184,43 @@ def get_dummy_mm_data(
         audio_len = (
             feature_extractor.chunk_length * sampling_rate * _MAX_ENCODER_BATCH_SIZE
         )
-        num_audios = mm_counts.get("audio", 0)
-
-        audio_overrides = mm_options.get("audio")
 
         return {
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
 
 class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo]):
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        audios = mm_data.get("audios", [])
-        assert isinstance(audios, list)
-
-        feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
-        hf_processor_mm_kwargs = dict(
-            **hf_processor_mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
-            include_audio_num_chunks=True,
-        )
-
-        return mm_data, hf_processor_mm_kwargs
-
-    def _apply_hf_processor_main(
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if processor_data:
-            processor_data, hf_processor_mm_kwargs = self._preprocess_hf_mm_data(
-                processor_data,
-                hf_processor_mm_kwargs,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        # The Ultravox processor accepts "audios" instead of "audio"
+        hf_data = hf_inputs.hf_data
+        if "audio" in hf_data:
+            hf_data["audios"] = hf_data.pop("audio")
+
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+        normalized_kwargs = dict(hf_inputs.hf_kwargs)
+        # The remote processor already hardcodes truncation=False and forwards
+        # its kwargs to the audio processor, where another value would collide.
+        normalized_kwargs.pop("truncation", None)
+        return hf_inputs._replace(
+            hf_kwargs=dict(
+                normalized_kwargs,
+                sampling_rate=feature_extractor.sampling_rate,
+                include_audio_num_chunks=True,
             )
-
-            prompt_text = self._get_hf_processor_text(mm_items.get_all_counts())
-            if prompt_text is not None:
-                processor_data = dict(text=prompt_text, **processor_data)
-
-            hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
-
-            def patched_call(**kwargs) -> BatchFeature:
-                # The remote UltravoxProcessor does not support the
-                # `truncation` kwarg injected by `call_hf_processor`: it
-                # passes `truncation` to its audio processor by itself while
-                # also forwarding `**kwargs` there, causing a duplicate
-                # keyword argument error. Its hardcoded `truncation=False`
-                # already matches vLLM's intended behavior.
-                kwargs.pop("truncation", None)
-
-                return hf_processor(**kwargs)
-
-            processed_data = self.info.ctx.call_hf_processor(
-                patched_call,
-                processor_data,
-                hf_processor_mm_kwargs,
-            )
-            processed_data.update(passthrough_data)
-        else:
-            processed_data = BatchFeature(dict(passthrough_data))
-
-        return self._postprocess_hf_mm_data(
-            processor_data,
-            hf_processor_mm_kwargs,
-            processed_data,
         )
 
     def _postprocess_hf_mm_data(
@@ -290,6 +254,7 @@ def _get_mm_fields_config(
             # num_chunks can convert audio_chunked to audio batch dimension
             audio_num_chunks=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
             audio_embeds=MultiModalFieldConfig.batched("audio"),
+            audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
         )
 
     def _get_prompt_updates(
@@ -316,6 +281,8 @@ def _get_prompt_updates(
         )
 
         def get_replacement_ultravox(item_idx: int):
+            if "audio_num_tokens" in out_mm_data:
+                return [replacement_id] * int(out_mm_data["audio_num_tokens"][item_idx])
             start = chunks_start_idx[item_idx]
             end = chunks_start_idx[item_idx + 1]
             audio_token_len = out_mm_data["audio_token_len"][start:end].sum()
@@ -339,8 +306,7 @@ def get_replacement_ultravox(item_idx: int):
 
 
 class StackAudioFrames(nn.Module):
-    """
-    Stack the audio embedding frames to reduce the sequence length by a factor
+    """Stack the audio embedding frames to reduce the sequence length by a factor
     of `stack_factor`.
     """
 
@@ -373,9 +339,9 @@ def _build_chunk_attn_metadata(
     valid/padding boundary (equivalent to the key-padding mask the HF
     implementation uses), while every row still flows through the
     (potentially LoRA-wrapped) linears, keeping the per-chunk token counts
-    constant as required by `get_num_mm_encoder_tokens` /
-    `get_num_mm_connector_tokens`. Queries at padding positions produce
-    (garbage) outputs, which are trimmed by `audio_token_len` downstream.
+    constant as required by `get_mm_lora_token_counts`. Queries at padding
+    positions produce (garbage) outputs, which are trimmed by `audio_token_len`
+    downstream.
     """
     batch_size = feature_lens.shape[0]
     starts = np.arange(batch_size, dtype=np.int64) * seq_len
@@ -684,7 +650,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
 
         # LoRA on the tower/connector requires per-item token counts that are
         # predictable from the placeholder count alone (see
-        # `get_num_mm_encoder_tokens`), so pad every audio chunk's mel input
+        # `get_mm_lora_token_counts`), so pad every audio chunk's mel input
         # to the tower's full context instead of the batch's max length.
         self.pad_audio_to_max_context = bool(
             lora_config is not None and lora_config.enable_tower_connector_lora
@@ -751,9 +717,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         )
 
     def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
+        """Get the module prefix in multimodal models."""
         return MultiModelKeys.from_string_field(
             language_model="language_model.",
             connector="multi_modal_projector.",
@@ -769,22 +733,26 @@ def _get_max_tokens_per_chunk(self) -> int:
             self.config.audio_config.max_source_positions / self.config.stack_factor
         )
 
-    def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int:
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        del modality, mm_kwargs
         # With `pad_audio_to_max_context` (enforced when tower/connector LoRA
         # is enabled), the tower's LoRA-wrapped linears always run on
         # `max_source_positions` conv-downsampled tokens per chunk, regardless
         # of the valid frame count.
-        num_chunks = math.ceil(num_audio_tokens / self._get_max_tokens_per_chunk())
-        return num_chunks * self.config.audio_config.max_source_positions
-
-    def get_num_mm_connector_tokens(self, num_encoder_tokens: int) -> int:
+        num_chunks = math.ceil(num_mm_embeds / self._get_max_tokens_per_chunk())
+        tower_tokens = num_chunks * self.config.audio_config.max_source_positions
         # The connector runs on the frame-stacked tower output. Stacking pads
         # each chunk to a multiple of `stack_factor`, so this is
         # ceil(max_source_positions / stack_factor) tokens per chunk (188 for
-        # whisper's 1500), not `num_encoder_tokens // stack_factor` (187).
-        max_source_positions = self.config.audio_config.max_source_positions
-        num_chunks = num_encoder_tokens // max_source_positions
-        return num_chunks * self._get_max_tokens_per_chunk()
+        # whisper's 1500), not `tower_tokens // stack_factor` (187).
+        connector_tokens = num_chunks * self._get_max_tokens_per_chunk()
+        return tower_tokens, connector_tokens
 
     def _audio_features_to_embeddings(
         self,
@@ -866,7 +834,7 @@ def _process_audio_input(
             # Pad every chunk to the tower's full context so the number of
             # tokens processed by the tower/connector per chunk is constant,
             # keeping the LoRA token mappings computed by
-            # `get_num_mm_encoder_tokens`/`get_num_mm_connector_tokens` exact.
+            # `get_mm_lora_token_counts` exact.
             # Attention to the extra padding is masked via `audio_lens`.
             # Over-long inputs are not trimmed here; the tower raises on them.
             pad_len = self.audio_tower.max_context_length - audio_features.shape[-1]
@@ -934,9 +902,9 @@ def forward(
         positions: torch.Tensor,
         intermediate_tensors: torch.Tensor | None = None,
         inputs_embeds: torch.Tensor | None = None,
-        **kwargs,
+        **kwargs: object,
     ) -> torch.Tensor | IntermediateTensors:
-        """Run forward pass for Ultravox
+        """Run forward pass for Ultravox.
 
         One key thing to understand is the `input_ids` already accounts for the
         positions of the to-be-inserted audio embeddings. The to-be-inserted
@@ -951,9 +919,10 @@ def forward(
             positions: Position indices for the input tokens.
             intermediate_tensors: Intermediate tensors from prior forward pass.
             inputs_embeds: Optional tensor of input embeddings.
+            **kwargs: Multimodal inputs for this batch, forwarded to the
+                multimodal embedding path.
 
         """
-
         if intermediate_tensors is not None:
             inputs_embeds = None
 
@@ -977,8 +946,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 def pad_and_concat_to_dim3(
     features: torch.Tensor | list[torch.Tensor] | list[list[torch.Tensor]],
 ) -> torch.Tensor:
-    """
-    Pad and concatenate a list of tensors.
+    """Pad and concatenate a list of tensors.
 
     output:
         Tensor of shape [B, C, M] where M is the maximum length of the input
diff --git a/vllm/model_executor/models/unlimited_ocr.py b/vllm/model_executor/models/unlimited_ocr.py
index c96e5f293d91..d248e5c8e068 100644
--- a/vllm/model_executor/models/unlimited_ocr.py
+++ b/vllm/model_executor/models/unlimited_ocr.py
@@ -75,7 +75,7 @@
 from vllm.multimodal.processing import PromptReplacement, PromptUpdate
 from vllm.multimodal.processing.context import TimingContext
 from vllm.multimodal.processing.inputs import ProcessorInputs
-from vllm.multimodal.processing.processor import MultiModalProcessingInfo
+from vllm.multimodal.processing.processor import MultiModalProcessingResult
 from vllm.transformers_utils.processors.deepseek_ocr import (
     BASE_SIZE,
     CROP_MODE,
@@ -225,7 +225,7 @@ def _cached_apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
+    ) -> MultiModalProcessingResult:
         # The processor logic differs for single-image (crop) vs multi-image
         # (no crop) requests. The processing cache assumes per-item output is
         # invariant of how many images are passed per prompt, so we only cache
diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py
index 6cd21763a095..6375e0572ea1 100644
--- a/vllm/model_executor/models/utils.py
+++ b/vllm/model_executor/models/utils.py
@@ -32,7 +32,7 @@
 )
 
 if TYPE_CHECKING:
-    from transformers import PretrainedConfig
+    from transformers import PreTrainedConfig
     from transformers.conversion_mapping import WeightRenaming
 
     from vllm.config.model import ModelConfig
@@ -41,6 +41,9 @@
 logger = init_logger(__name__)
 
 ShardId: TypeAlias = str | int | tuple[int, ...]
+"""One shard of a stacked parameter. A tuple is a single contiguous span."""
+ShardIds: TypeAlias = ShardId | list[ShardId]
+"""One shard, or a list of shards the same weight is loaded into in turn."""
 
 
 @dataclass
@@ -52,7 +55,9 @@ class WeightsMapper:
     orig_to_new_renaming: list["WeightRenaming"] = field(default_factory=list)
     orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict)
     orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict)
-    orig_to_new_stacked: Mapping[str, tuple[str, ShardId]] = field(default_factory=dict)
+    orig_to_new_stacked: Mapping[str, tuple[str, ShardIds]] = field(
+        default_factory=dict
+    )
     orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict)
     orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict)
 
@@ -73,17 +78,23 @@ def __or__(self, other: "WeightsMapper") -> "WeightsMapper":
             orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix},
         )
 
-    def _map_name(self, key: str) -> str | None:
-        """Map a weight name (backward-compatible wrapper that discards shard_id)."""
+    def map_name(self, key: str) -> str | None:
+        """Map a weight name; returns ``None`` if the weight should be ignored."""
         result = self._map_name_with_shard(key)
         return result[0] if result is not None else None
 
-    def _map_name_with_shard(self, key: str) -> tuple[str, ShardId | None] | None:
+    def _map_name(self, key: str) -> str | None:
+        """Backward-compatible alias for :meth:`map_name`."""
+        return self.map_name(key)
+
+    def _map_name_with_shard(self, key: str) -> tuple[str, ShardIds | None] | None:
         """Map a weight name and extract any shard_id metadata.
 
         Returns:
-            (mapped_name, shard_id) if the name should be kept.
+            (mapped_name, shard_id) if the name should be kept. A list of shard
+            ids means the weight is loaded into each of those shards in turn.
             None if the name should be dropped.
+
         """
         # Deprecation warnings
         if key.endswith(".kv_scale"):
@@ -112,7 +123,7 @@ def _map_name_with_shard(self, key: str) -> tuple[str, ShardId | None] | None:
 
                 key = key.replace(substr, new_key, 1)
 
-        shard_id: ShardId | None = None
+        shard_id: ShardIds | None = None
         for substr, (new_key, new_shard_id) in self.orig_to_new_stacked.items():
             if substr in key:
                 key = key.replace(substr, new_key, 1)
@@ -142,9 +153,16 @@ def apply(
             if result is None:
                 continue
             out_name, shard_id = result
-            if shard_id is not None:
-                data.shard_id = shard_id
-            yield out_name, data
+            if shard_id is None:
+                yield out_name, data
+                continue
+            shard_ids = shard_id if isinstance(shard_id, list) else [shard_id]
+            for i, one_shard_id in enumerate(shard_ids):
+                # Each shard carries its own id, so extra shards need their own
+                # tensor object; `detach` aliases the storage instead of copying.
+                shard_data = data if i == 0 else data.detach()
+                shard_data.shard_id = one_shard_id
+                yield out_name, shard_data
 
     def apply_list(self, values: list[str]) -> list[str]:
         return [
@@ -197,8 +215,7 @@ def _get_tied_embedding_params(module: nn.Module) -> dict[str, str]:
 
 
 class AutoWeightsLoader:
-    """
-    Helper class to load weights into a [`torch.nn.Module`][]. It is able
+    """Helper class to load weights into a [`torch.nn.Module`][]. It is able
     to automatically detect child modules and parameters while iterating over
     the weights only once.
 
@@ -315,8 +332,7 @@ def _load_param(
     def _add_loadable_non_param_tensors(
         self, module: nn.Module, child_params: dict[str, torch.Tensor]
     ):
-        """
-        Add tensor names that are not in the model params that may be in the
+        """Add tensor names that are not in the model params that may be in the
         safetensors, e.g., batch normalization stats and registered buffers.
         """
         # Add persistent registered buffers.
@@ -501,6 +517,7 @@ def maybe_fuse_shared_experts(
 
     Yields:
         `(name, tensor)` pairs with shared experts routed to fused slots.
+
     """
     if enabled is None:
         from vllm._aiter_ops import rocm_aiter_ops
@@ -552,6 +569,7 @@ def get_spec_layer_idx_from_weight_name(
 
     Returns:
         The absolute layer index for an MTP-layer weight, else None.
+
     """
     if not (n := getattr(config, "num_nextn_predict_layers", 0)):
         return None
@@ -579,6 +597,7 @@ def skip_spec_layers(
 
     Yields:
         `(name, tensor)` pairs whose weight is not an MTP-layer weight.
+
     """
     return (
         (name, w)
@@ -591,11 +610,10 @@ def init_vllm_registered_model(
     vllm_config: VllmConfig,
     *,
     prefix: str = "",
-    hf_config: "PretrainedConfig | None" = None,
+    hf_config: "PreTrainedConfig | None" = None,
     architectures: list[str] | None = None,
 ) -> nn.Module:
-    """
-    Helper function to initialize an inner model registered to vLLM,
+    """Helper function to initialize an inner model registered to vLLM,
     based on the arguments passed to the outer vLLM model.
     """
     from vllm.model_executor.model_loader.utils import initialize_model
@@ -639,8 +657,7 @@ def flatten_bn(
     *,
     concat: bool = False,
 ) -> list[torch.Tensor] | torch.Tensor:
-    """
-    Flatten the `B` and `N` dimensions of batched multimodal inputs.
+    """Flatten the `B` and `N` dimensions of batched multimodal inputs.
 
     The input tensor should have shape `(B, N, ...)`.
     """
@@ -654,11 +671,9 @@ def flatten_bn(
 
 
 def _flatten_embeddings(embeddings: NestedTensors) -> torch.Tensor:
-    """
-    Recursively flattens and concatenates NestedTensors on all but the last
+    """Recursively flattens and concatenates NestedTensors on all but the last
     dimension.
     """
-
     if isinstance(embeddings, torch.Tensor):
         # Flatten all but the last dimension.
         return embeddings.flatten(0, -2)
@@ -667,11 +682,9 @@ def _flatten_embeddings(embeddings: NestedTensors) -> torch.Tensor:
 
 
 def _embedding_count_expression(embeddings: NestedTensors) -> str:
-    """
-    Constructs a debugging representation of the number of embeddings in the
+    """Constructs a debugging representation of the number of embeddings in the
     NestedTensors.
     """
-
     if isinstance(embeddings, torch.Tensor):
         return " x ".join([str(dim) for dim in embeddings.shape[:-1]])
 
@@ -691,13 +704,13 @@ def _merge_multimodal_embeddings(
     multimodal_embeddings: NestedTensors,
     is_multimodal: torch.Tensor,
 ) -> torch.Tensor:
-    """
-    Merge `multimodal_embeddings` into `inputs_embeds` by overwriting the
+    """Merge `multimodal_embeddings` into `inputs_embeds` by overwriting the
     positions in `inputs_embeds` corresponding to placeholder tokens in
     `input_ids`.
 
     Note:
         This updates `inputs_embeds` in place.
+
     """
     if len(multimodal_embeddings) == 0:
         return inputs_embeds
@@ -751,8 +764,7 @@ def collect_children(
     *,
     targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
 ):
-    """
-    Within this context, collect all direct child assignments to `module`,
+    """Within this context, collect all direct child assignments to `module`,
     returning a list of children names that is internally updated until the
     context is exited.
 
@@ -784,8 +796,7 @@ def no_init_weights(
     *,
     targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
 ):
-    """
-    Within this context, prevent weight initialization from using device memory and
+    """Within this context, prevent weight initialization from using device memory and
     replace direct child assignments to `module` with the result of `placeholder()`.
 
     If `targets` is set, instead prevent weight initialization and
@@ -824,9 +835,7 @@ def __call__(self, prefix: str) -> torch.nn.Module: ...
 
 
 class PPMissingLayer(torch.nn.Identity):
-    """
-    A placeholder layer for missing layers in a pipeline parallel model.
-    """
+    """A placeholder layer for missing layers in a pipeline parallel model."""
 
     def __init__(self, *args, **kwargs):
         super().__init__()
@@ -867,6 +876,7 @@ def make_layers(
 
     Returns:
         Tuple of (start_layer, end_layer, modules).
+
     """
     from vllm.distributed.parallel_state import get_pp_group
     from vllm.distributed.utils import get_pp_indices
@@ -945,6 +955,7 @@ def maybe_prefix(prefix: str, name: str) -> str:
 
     Returns:
         The string "prefix.name" if prefix was non-empty, otherwise just "name".
+
     """
     return name if not prefix else f"{prefix}.{name}"
 
@@ -960,6 +971,7 @@ def get_draft_quant_config(vllm_config: VllmConfig) -> "QuantizationConfig | Non
 
     Returns:
         The draft model's config if available, None otherwise.
+
     """
     draft_model_config = vllm_config.speculative_config.draft_model_config
     draft_load_config = vllm_config.load_config
@@ -972,13 +984,14 @@ def get_draft_quant_config(vllm_config: VllmConfig) -> "QuantizationConfig | Non
 
 
 def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int:
-    """
-    Extract the layer index from the module name.
+    """Extract the layer index from the module name.
+
     Examples:
     - "encoder.layers.0" -> 0
     - "encoder.layers.1.self_attn" -> 1
     - "2.self_attn" -> 2
     - "model.encoder.layers.0.sub.1" -> ValueError if num_attn_module == 1
+
     """
     subnames = layer_name.split(".")
     int_vals: list[int] = []
@@ -1013,8 +1026,7 @@ def cast_overflow_tensors(tensors: torch.Tensor, offset: float = 1000) -> torch.
 def fast_topk(
     values: torch.Tensor, topk: int, dim: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Optimized topk implementation that uses torch.max for k=1 case.
+    """Optimized topk implementation that uses torch.max for k=1 case.
 
     This function provides better performance for the common case of k=1
     by using torch.max instead of the more general torch.topk.
@@ -1027,6 +1039,7 @@ def fast_topk(
     Returns:
         Tuple of (values, indices) where values are the top-k values
         and indices are their corresponding indices in the input tensor
+
     """
     if topk == 1:
         # Use max along the specified dimension to get both value and index
@@ -1086,13 +1099,14 @@ def process_eagle_weight(
     model: nn.Module,
     name: str,
 ) -> None:
-    """
-    Update EAGLE model flags based on loaded weight name.
+    """Update EAGLE model flags based on loaded weight name.
     This should be called during weight loading to detect if a model
     has its own lm_head or embed_tokens weight.
+
     Args:
         model: The model instance (must support EAGLE)
         name: The name of the weight to process
+
     """
     if not supports_any_eagle(model):
         return
@@ -1111,6 +1125,7 @@ def get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:
     Args:
         feature_layer_index: Index of a required layer in the visual encoder.
         num_hidden_layers: The total number of hidden layers in the visual encoder.
+
     """
     if feature_layer_index < 0:
         return num_hidden_layers + feature_layer_index + 1
diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py
index 36008d999654..b65171421bc7 100644
--- a/vllm/model_executor/models/vision.py
+++ b/vllm/model_executor/models/vision.py
@@ -9,7 +9,7 @@
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import ModelConfig, MultiModalConfig, get_current_vllm_config_or_none
 from vllm.distributed import (
@@ -25,7 +25,7 @@
 
 logger = init_logger(__name__)
 
-_C = TypeVar("_C", bound=PretrainedConfig)
+_C = TypeVar("_C", bound=PreTrainedConfig)
 
 
 class _RootConfig(Protocol[_C]):
@@ -62,7 +62,7 @@ def get_patch_grid_length(self) -> int:
 
 
 class VisionLanguageConfig(Protocol):
-    vision_config: Final[PretrainedConfig]
+    vision_config: Final[PreTrainedConfig]
 
 
 def get_vision_encoder_info(hf_config: VisionLanguageConfig) -> VisionEncoderInfo:
@@ -88,9 +88,7 @@ def _get_vit_attn_backend(
     *,
     attn_backend_override: AttentionBackendEnum | None = None,
 ) -> AttentionBackendEnum:
-    """
-    Get the available attention backend for Vision Transformer.
-    """
+    """Get the available attention backend for Vision Transformer."""
     return current_platform.get_vit_attn_backend(
         head_size,
         dtype,
@@ -102,9 +100,7 @@ def get_vit_attn_backend(
     head_size: int,
     dtype: torch.dtype,
 ) -> AttentionBackendEnum:
-    """
-    Get the attention backend for Vision Transformer.
-    """
+    """Get the attention backend for Vision Transformer."""
     mm_cfg = get_multimodal_config()
     attn_backend_override = (
         mm_cfg.mm_encoder_attn_backend if mm_cfg is not None else None
@@ -142,9 +138,7 @@ def get_fp8_padded_hidden_size(num_heads: int, head_dim: int) -> int | None:
 
 
 def is_vit_use_data_parallel(num_heads: int | None = None) -> bool:
-    """
-    Get the tensor parallel type for Vision Transformer.
-    """
+    """Get the tensor parallel type for Vision Transformer."""
     mm_cfg = get_multimodal_config()
     can_split = (
         num_heads % get_tensor_model_parallel_world_size() == 0
@@ -235,6 +229,7 @@ def resolve_visual_encoder_outputs(
             concatenated with the other select_layers along the last dimension.
         feature_select_strategy: Defines how to select the hidden states
             from each layer.
+
     """
     if select_layers is None:
         if not isinstance(encoder_outputs, torch.Tensor):
@@ -302,10 +297,11 @@ def run_dp_sharded_vision_model(
     Args:
         image_input (torch.Tensor): Image input tensor.
         vision_model (torch.nn.Module): Vision model.
+
     Returns:
         torch.Tensor: Output image embeddings
-    """
 
+    """
     num_chunks = image_input.shape[0]
     mp_world_size = get_tensor_model_parallel_world_size()
     num_chunks_per_rank = (num_chunks + mp_world_size - 1) // mp_world_size
@@ -329,8 +325,7 @@ def get_load_balance_assignment(
     sizes: list[int],
     num_gpus: int = 2,
 ) -> tuple[list[int], list[int], list[int]]:
-    """
-    Generate load balancing assignment and metadata
+    """Generate load balancing assignment and metadata
     for distributing data across GPUs.
     The load is determined by the total image sizes,
     not the number of images.
@@ -354,7 +349,6 @@ def get_load_balance_assignment(
         ```
 
     """
-
     n_samples = len(sizes)
 
     # Handle edge cases
@@ -414,6 +408,7 @@ def run_dp_sharded_mrope_vision_model(
                    Different rope types have different dimension to do ViT.
                    "rope_3d" for 3D rope (e.g., Qwen2.5-VL)
                    "rope_2d" for 2D rope (e.g., Kimi-VL)
+
     Returns:
         torch.Tensor: Output image embeddings
 
@@ -486,6 +481,7 @@ def run_dp_sharded_mrope_vision_model(
     # to work
     max_len_per_rank = max(grouped_pixel_values_len) // embed_dim_reduction_factor
     local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]
+    embed_dtype = next(vision_model.parameters()).dtype
 
     # Run the vision model on the local pixel_values_local
     if rope_type == "rope_2d":
@@ -500,7 +496,7 @@ def run_dp_sharded_mrope_vision_model(
             image_embeds_local = torch.empty(
                 (0, embed_dim_reduction_factor, out_dim),
                 device=pixel_values.device,
-                dtype=pixel_values.dtype,
+                dtype=embed_dtype,
             )
     else:
         if pixel_values_local.shape[0] > 0:
@@ -510,7 +506,7 @@ def run_dp_sharded_mrope_vision_model(
             image_embeds_local = torch.empty(
                 (0, vision_model.out_hidden_size),
                 device=pixel_values.device,
-                dtype=pixel_values.dtype,
+                dtype=embed_dtype,
             )
 
     # Pad the output based on max_len_per_rank
@@ -585,8 +581,7 @@ def run_dp_sharded_mrope_vision_model(
 
 
 class FusedInputNorm(nn.Module):
-    """
-    Module that applies rescaling and normalization to input images.
+    """Module that applies rescaling and normalization to input images.
     Equivalent to: output = (input * rescale_factor - mean) / std
     """
 
diff --git a/vllm/model_executor/models/voxtral.py b/vllm/model_executor/models/voxtral.py
index b41456bcdc43..21b291344d90 100644
--- a/vllm/model_executor/models/voxtral.py
+++ b/vllm/model_executor/models/voxtral.py
@@ -16,10 +16,10 @@
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
 from mistral_common.protocol.transcription.request import TranscriptionRequest
 from mistral_common.tokens.tokenizers.audio import Audio
-from transformers import BatchFeature, WhisperConfig
+from transformers import WhisperConfig
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt
 from vllm.logger import init_logger
@@ -36,7 +36,6 @@
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalKwargsItems,
-    MultiModalKwargsOptionalItems,
     NestedTensors,
 )
 from vllm.multimodal.parse import (
@@ -48,7 +47,8 @@
 from vllm.multimodal.processing.processor import (
     BaseMultiModalProcessor,
     BaseProcessingInfo,
-    MultiModalPromptUpdates,
+    HFMultiModalInputs,
+    MultiModalProcessingResult,
     PlaceholderFeaturesInfo,
     ProcessorInputs,
     PromptReplacement,
@@ -61,7 +61,6 @@
     MistralCommonFeatureExtractor,
     MistralCommonVoxtralProcessor,
 )
-from vllm.utils.collection_utils import is_list_of
 
 from .interfaces import SupportsLoRA, SupportsMultiModal, SupportsTranscription
 from .utils import init_vllm_registered_model, maybe_prefix
@@ -138,35 +137,35 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
-
         target_length = self.info.get_max_audio_array_len()
 
-        audio_overrides = mm_options.get("audio")
-
         return {
             "audio": self._get_dummy_audios(
                 length=target_length,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
-    def get_dummy_processor_inputs(
+
+class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo]):
+    def get_dummy_inputs(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
+        # For test_common.py only
         mm_data: MultiModalDataDict | None = None,
     ) -> ProcessorInputs:
+        builder = self.dummy_inputs
         tokenizer = self.info.get_tokenizer()
         feature_extractor = self.info.get_feature_extractor()
 
-        dummy_text = self.get_dummy_text(mm_counts)
+        dummy_text = builder.get_dummy_text(mm_counts)
         dummy_mm_data = (
-            self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
+            builder.get_dummy_mm_data(seq_len, mm_counts, mm_options)
             if mm_data is None
             else mm_data
         )
@@ -202,28 +201,24 @@ def get_dummy_processor_inputs(
 
         return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items)
 
-
-class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo]):
     # The tokens are already inserted by the chat template,
     # so we just double check that they exist
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsOptionalItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_res: MultiModalProcessingResult,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
         mm_item_counts = mm_items.get_all_counts()
-        self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
-        self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
+        self._validate_mm_kwargs(mm_res.kwargs, mm_item_counts)
+        self._validate_mm_updates(mm_res.prompt_updates, mm_item_counts)
 
         mm_placeholders = self._find_mm_placeholders(
-            prompt_ids,
-            mm_prompt_updates,
+            mm_res.prompt_ids,
+            mm_res.prompt_updates,
         )
         self._validate_mm_placeholders(mm_placeholders, mm_item_counts)
 
-        return prompt_ids, mm_placeholders
+        return mm_res.prompt_ids, mm_placeholders
 
     def _get_mm_fields_config(
         self,
@@ -241,38 +236,20 @@ def _validate_mm_placeholders(
         # skip validation here
         pass
 
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", [])
-
-        if audios:
-            # MistralCommonVoxtralProcessor accepts "audio"
-            mm_data["audio"] = audios
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
 
         # Avoid padding issue
-        hf_processor_mm_kwargs = dict(**hf_processor_mm_kwargs, return_tensors=None)
-
-        return mm_data, hf_processor_mm_kwargs
-
-    def _postprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-        processed_data: BatchFeature,
-    ) -> BatchFeature:
-        # Missing batch dimension
-        input_ids = processed_data.get("input_ids")
-        if is_list_of(input_ids, int):
-            processed_data["input_ids"] = [input_ids]
-
-        return processed_data
+        return hf_inputs._replace(
+            hf_kwargs=dict(hf_inputs.hf_kwargs, return_tensors=None)
+        )
 
     def _get_prompt_updates(
         self,
@@ -502,8 +479,7 @@ def get_num_audio_tokens(
         stt_config: SpeechToTextConfig,
         model_config: ModelConfig,
     ) -> int | None:
-        """
-        Map from audio duration to number of audio tokens produced by the ASR
+        """Map from audio duration to number of audio tokens produced by the ASR
         model, without running a forward pass.
         This is used for estimating the amount of processing for this audio.
         """
@@ -586,8 +562,7 @@ def llm_weights_generator():
     def maybe_update_quant_config(
         self, quant_config: QuantizationConfig
     ) -> QuantizationConfig:
-        """
-        Update quant config to so that ignored module and target module names
+        """Update quant config to so that ignored module and target module names
         match the vLLM model names.
         Right now this is specific for compressed-tensors format and
         load_format mistral.
diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py
index 738185a7aecf..9fdc946df990 100644
--- a/vllm/model_executor/models/voxtral_realtime.py
+++ b/vllm/model_executor/models/voxtral_realtime.py
@@ -27,43 +27,37 @@
     VoxtralProcessingInfo,
 )
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.cache import _I, BaseMultiModalProcessorCache
-from vllm.multimodal.inputs import MultiModalKwargsOptionalItems
 from vllm.multimodal.parse import MultiModalDataItems
-from vllm.multimodal.processing import BaseDummyInputsBuilder
+from vllm.multimodal.processing import ProcessorInputs, TimingContext
 from vllm.multimodal.processing.processor import (
-    MultiModalPromptUpdates,
+    MultiModalProcessingResult,
     PlaceholderFeaturesInfo,
 )
 from vllm.sequence import IntermediateTensors
 from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.utils.torch_utils import is_torch_equal_or_newer
 
-from .utils import (
-    _flatten_embeddings,
-)
+from .utils import _flatten_embeddings
 
 logger = init_logger(__name__)
 
 
 class VoxtralRealtimeMultiModalProcessor(VoxtralMultiModalProcessor):
-    def __init__(
+    def _cached_apply_hf_processor(
         self,
-        info: _I,
-        dummy_inputs: BaseDummyInputsBuilder[_I],
-        *,
-        cache: BaseMultiModalProcessorCache | None = None,
-    ) -> None:
+        inputs: ProcessorInputs,
+        timing_ctx: TimingContext,
+    ) -> MultiModalProcessingResult:
         # realtime can't make use of a cache yet
-        super().__init__(info, dummy_inputs, cache=None)
+        return self._apply_hf_processor(inputs, timing_ctx)
 
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsOptionalItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_res: MultiModalProcessingResult,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
+        mm_kwargs = mm_res.kwargs
+
         # there are no placeholder audio tokens for streaming
         # so we need to build the place placeholder positions manually
 
@@ -86,11 +80,11 @@ def _maybe_apply_prompt_updates(
             * [0],  # only used for length computation, so we can take dummy inputs
             is_embed=None,
         )
-        return prompt_ids, {"audio": [features_info]}
+        return mm_res.prompt_ids, {"audio": [features_info]}
 
 
 class TimeEmbedding(torch.nn.Module):
-    """Sinusoidal Embedding for encoding time"""
+    """Sinusoidal Embedding for encoding time."""
 
     def __init__(self, dim: int, theta: float = 10000.0) -> None:
         super().__init__()
@@ -217,9 +211,15 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__(vllm_config=vllm_config, prefix=prefix)
 
-        assert (
-            not vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs()
-        ), "Voxtral realtime doesn't support full cudagraphs yet. Please use PIECEWISE."
+        # Full cudagraphs are supported for decode-only batches (the encoder's
+        # block-pooling attention builder is capture-safe for uniform
+        # single-token decode). Mixed/prefill batches are not, so block only
+        # pure FULL; FULL_DECODE_ONLY and FULL_AND_PIECEWISE are allowed.
+        cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
+        assert not cudagraph_mode.mixed_mode().has_full_cudagraphs(), (
+            "Voxtral realtime supports full cudagraphs for decode-only batches. "
+            "Use cudagraph_mode=FULL_DECODE_ONLY or FULL_AND_PIECEWISE (not FULL)."
+        )
 
         self.time_embedding: TimeEmbedding = TimeEmbedding(
             dim=self.config.text_config.hidden_size
@@ -385,7 +385,7 @@ def forward(
     def embed_multimodal(
         self, **kwargs
     ) -> list[torch.Tensor] | torch.Tensor | tuple[torch.Tensor, ...] | None:
-        """Transform audio waveforms -> initial whisper post-conv embeddings"""
+        """Transform audio waveforms -> initial whisper post-conv embeddings."""
         audio_inputs = self._parse_and_validate_audio_arrays(**kwargs)
 
         if audio_inputs is None:
diff --git a/vllm/model_executor/models/voyage.py b/vllm/model_executor/models/voyage.py
index 92d6d7633f02..9f4885a0f9db 100644
--- a/vllm/model_executor/models/voyage.py
+++ b/vllm/model_executor/models/voyage.py
@@ -8,6 +8,7 @@
 import torch.nn as nn
 
 from vllm.config import VllmConfig
+from vllm.model_executor.models.interfaces import SupportsLoRA
 from vllm.model_executor.models.qwen3 import Qwen3Model
 from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix
 from vllm.sequence import IntermediateTensors
@@ -15,9 +16,8 @@
 WeightItem = tuple[str, torch.Tensor]
 
 
-class VoyageQwen3BidirectionalEmbedModel(nn.Module):
-    """
-    Qwen3Model + Voyage embedding head + bidirectional attention.
+class VoyageQwen3BidirectionalEmbedModel(nn.Module, SupportsLoRA):
+    """Qwen3Model + Voyage embedding head + bidirectional attention.
 
     Checkpoint conventions (HF):
       - MLP: gate_proj + up_proj (unfused)
@@ -31,6 +31,11 @@ class VoyageQwen3BidirectionalEmbedModel(nn.Module):
       - No "model." prefix
     """
 
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
diff --git a/vllm/model_executor/models/whisper.py b/vllm/model_executor/models/whisper.py
index d21903bd656c..f67e3f0cb9fd 100644
--- a/vllm/model_executor/models/whisper.py
+++ b/vllm/model_executor/models/whisper.py
@@ -19,7 +19,7 @@
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.inputs import (
@@ -61,6 +61,7 @@
     PromptReplacement,
     PromptUpdate,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.renderers import TokenizeParams
 from vllm.transformers_utils.processor import cached_processor_from_config
 from vllm.utils.jsontree import json_map_leaves
@@ -94,11 +95,10 @@ class WhisperPosEmbedType(enum.Enum):
 
 
 class WhisperAudioInputs(TensorSchema):
-    """
-    Dimensions:
-        - b: Batch size
-        - nmb: Number of mel bins
-        - t: Time frames (M)
+    """Dimensions:
+    - b: Batch size
+    - nmb: Number of mel bins
+    - t: Time frames (M)
     """
 
     input_features: Annotated[
@@ -120,9 +120,8 @@ def forward(
         # Only used for FlashInfer CuDNN backend.
         sequence_lengths: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """
-        Input shape: batch_size x seq_len x hidden_size
-                     or seq_len x hidden_size
+        """Input shape: batch_size x seq_len x hidden_size
+        or seq_len x hidden_size
         """
         is_2d = query.dim() == 2
         if is_2d:
@@ -725,26 +724,26 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         feature_extractor = self.info.get_feature_extractor()
 
         sampling_rate = feature_extractor.sampling_rate
         audio_len = feature_extractor.chunk_length * sampling_rate
-        num_audios = mm_counts.get("audio", 0)
-
-        audio_overrides = mm_options.get("audio")
 
         return {
             "audio": self._get_dummy_audios(
                 length=audio_len,
-                num_audios=num_audios,
-                overrides=audio_overrides,
+                num_audios=mm_counts.get("audio", 0),
+                overrides=mm_options.get("audio"),
             )
         }
 
 
 class WhisperMultiModalProcessor(EncDecMultiModalProcessor[WhisperProcessingInfo]):
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
     def create_encoder_prompt(
         self,
         prompt: list[int],
@@ -756,37 +755,22 @@ def create_encoder_prompt(
         # for encoder profiling.
         return [0]
 
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
-        return self.dummy_inputs.get_dummy_text(mm_counts)
-
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
-
-        mm_data = dict(mm_data)
-        mm_data["audio"] = mm_data.pop("audios")
-
-        hf_processor_mm_kwargs = dict(
-            **hf_processor_mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+        return hf_inputs._replace(
+            hf_kwargs=dict(
+                hf_inputs.hf_kwargs,
+                sampling_rate=feature_extractor.sampling_rate,
+                truncation=True,
+            )
         )
 
-        return mm_data, hf_processor_mm_kwargs
-
-    def _postprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-        processed_data: BatchFeature,
-    ) -> BatchFeature:
-        if "labels" in processed_data:
-            processed_data["input_ids"] = processed_data.pop("labels")
-
-        return processed_data
-
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
@@ -931,7 +915,6 @@ def parse_language_detection_output(
         Decodes the first token ID and extracts the language code from the
         ``<|xx|>`` format. Expects a valid language token from constrained generation.
         """
-
         decoded = tokenizer.decode(
             [token_ids[0]],
             skip_special_tokens=False,
@@ -1058,8 +1041,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 def _create_fake_bias_for_k_proj(
     weights: Iterable[tuple[str, torch.Tensor]], fake_bias_key_name: str
 ) -> Iterable[tuple[str, torch.Tensor]]:
-    """
-    Create full zeros bias for k_proj weight in self-attn and x-attn layers.
+    """Create full zeros bias for k_proj weight in self-attn and x-attn layers.
     So that the bias for k_proj in qkv_proj can be initialized with zeros.
     """
     for name, weight in weights:
diff --git a/vllm/model_executor/models/whisper_causal.py b/vllm/model_executor/models/whisper_causal.py
index d76e9b0d4a69..ec84cae8cfcf 100644
--- a/vllm/model_executor/models/whisper_causal.py
+++ b/vllm/model_executor/models/whisper_causal.py
@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import copy
 import functools
 import logging
 import math
 from dataclasses import replace
 from fractions import Fraction
 from functools import partial
+from typing import ClassVar
 
 import torch
 import torch.nn.functional as F
@@ -27,6 +27,7 @@
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backend import (
     AttentionBackend,
+    AttentionCGSupport,
     AttentionMetadata,
     AttentionType,
     CommonAttentionMetadata,
@@ -123,6 +124,16 @@ def create_whisper_attention_backend_with_block_pooling(
     underlying_impl = underlying_attn_backend.get_impl_cls()
 
     class WhisperCausalAttentionWithBlockPoolingBuilder(underlying_builder):  # type: ignore
+        # Full cudagraphs only for uniform single-token decode: capture bakes in
+        # tensor addresses, so `build` writes metadata into persistent buffers
+        # (below). Prefill/mixed batches fall back to piecewise.
+        _cudagraph_support: ClassVar[AttentionCGSupport] = (
+            AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
+        )
+        # Disable the update_block_table fast path: it would bypass block pooling
+        # by splicing the raw slot_mapping/block_table into cached metadata.
+        supports_update_block_table: bool = False
+
         def __init__(
             self,
             kv_cache_spec: AttentionSpec,
@@ -150,6 +161,38 @@ def __init__(
             # num_heads_q for the encoder is the same as num_kv_heads
             # (no GQA in whisper encoder)
             self.num_heads_q = kv_cache_spec.num_kv_heads
+            # Persistent capture-safe buffers for pooled metadata; lazily sized
+            # on first `build` to match the incoming tensors' dtype/device.
+            self._pool_buffers_ready = False
+
+        def _maybe_init_pool_buffers(
+            self, common_attn_metadata: CommonAttentionMetadata
+        ) -> None:
+            if self._pool_buffers_ready:
+                return
+            device = common_attn_metadata.slot_mapping.device
+            sched = self.vllm_config.scheduler_config
+            self._qsl_buf = torch.empty(
+                sched.max_num_seqs + 1,
+                dtype=common_attn_metadata.query_start_loc.dtype,
+                device=device,
+            )
+            self._seq_buf = torch.empty(
+                sched.max_num_seqs,
+                dtype=common_attn_metadata.seq_lens.dtype,
+                device=device,
+            )
+            self._slot_buf = torch.empty(
+                sched.max_num_batched_tokens * block_pool_size,
+                dtype=common_attn_metadata.slot_mapping.dtype,
+                device=device,
+            )
+            self._pool_arange = torch.arange(
+                block_pool_size,
+                dtype=common_attn_metadata.slot_mapping.dtype,
+                device=device,
+            )
+            self._pool_buffers_ready = True
 
         def build(
             self,
@@ -157,29 +200,62 @@ def build(
             common_attn_metadata: CommonAttentionMetadata,
             fast_build: bool = False,
         ) -> AttentionMetadata:
-            new_common_attn_metadata = copy.deepcopy(common_attn_metadata)
-            new_common_attn_metadata.query_start_loc *= block_pool_size
-            new_common_attn_metadata.query_start_loc_cpu *= block_pool_size
-            new_common_attn_metadata.seq_lens *= block_pool_size
-            if new_common_attn_metadata._seq_lens_cpu is not None:
-                new_common_attn_metadata._seq_lens_cpu *= block_pool_size
-            if new_common_attn_metadata._num_computed_tokens_cpu is not None:
-                new_common_attn_metadata._num_computed_tokens_cpu *= block_pool_size
-            new_common_attn_metadata.num_actual_tokens *= block_pool_size
-            new_common_attn_metadata.max_query_len *= block_pool_size
-            new_common_attn_metadata.max_seq_len *= block_pool_size
-            original_slot_mapping = common_attn_metadata.slot_mapping
-            common_prefix_len *= block_pool_size
-            new_common_attn_metadata.slot_mapping = (
-                (
-                    original_slot_mapping.unsqueeze(1) * block_pool_size
-                    + torch.arange(block_pool_size, device=original_slot_mapping.device)
-                )
-                .flatten()
-                .clamp(min=-1)
+            cm = common_attn_metadata
+            self._maybe_init_pool_buffers(cm)
+
+            # Scale query_start_loc / seq_lens into the persistent buffers.
+            n = cm.query_start_loc.numel()
+            assert n <= self._qsl_buf.numel(), (
+                f"{n=} exceeds pooled query_start_loc buffer {self._qsl_buf.numel()}"
+            )
+            qsl_buf = self._qsl_buf[:n]
+            qsl_buf.copy_(cm.query_start_loc)
+            qsl_buf.mul_(block_pool_size)
+
+            r = cm.seq_lens.numel()
+            assert r <= self._seq_buf.numel(), (
+                f"{r=} exceeds pooled seq_lens buffer {self._seq_buf.numel()}"
+            )
+            seq_buf = self._seq_buf[:r]
+            seq_buf.copy_(cm.seq_lens)
+            seq_buf.mul_(block_pool_size)
+
+            # Expand each slot into block_pool_size contiguous slots. Padding
+            # slots (-1) stay -1 after clamp so the KV write is skipped.
+            m = cm.slot_mapping.numel()
+            assert m * block_pool_size <= self._slot_buf.numel(), (
+                f"pooled slot_mapping {m * block_pool_size} exceeds buffer "
+                f"{self._slot_buf.numel()}"
+            )
+            slot_buf = self._slot_buf[: m * block_pool_size]
+            slot_buf.view(m, block_pool_size).copy_(
+                cm.slot_mapping.unsqueeze(1) * block_pool_size + self._pool_arange
+            )
+            slot_buf.clamp_(min=-1)
+
+            # CPU tensors aren't captured; fresh scaled copies avoid mutating
+            # the runner's shared buffers.
+            qsl_cpu = cm.query_start_loc_cpu * block_pool_size
+            seq_lens_cpu_upper_bound = (
+                cm.seq_lens_cpu_upper_bound * block_pool_size
+                if cm.seq_lens_cpu_upper_bound is not None
+                else None
+            )
+
+            new_common_attn_metadata = cm.replace(
+                query_start_loc=qsl_buf,
+                query_start_loc_cpu=qsl_cpu,
+                seq_lens=seq_buf,
+                seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+                num_actual_tokens=cm.num_actual_tokens * block_pool_size,
+                max_query_len=cm.max_query_len * block_pool_size,
+                max_seq_len=cm.max_seq_len * block_pool_size,
+                slot_mapping=slot_buf,
             )
             return super().build(
-                common_prefix_len, new_common_attn_metadata, fast_build
+                common_prefix_len * block_pool_size,
+                new_common_attn_metadata,
+                fast_build,
             )
 
     # NOTE: We need a custom impl so we can use the transformed slot_mapping
diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py
index f80c8a3b0783..9398d997688a 100644
--- a/vllm/model_executor/models/zamba2.py
+++ b/vllm/model_executor/models/zamba2.py
@@ -71,6 +71,8 @@ def __init__(
             rank: LoRA rank
             output_dim: output dimension
             quant_config: Configuration for model quantization
+            prefix: Module prefix used for quantization config lookup
+
         """
         super().__init__()
 
@@ -130,6 +132,7 @@ def __init__(
             cache_config: Configuration for key-value caching
             quant_config: Configuration for model quantization
             prefix: Optional prefix for parameter names
+
         """
         super().__init__()
         tp_size = get_tensor_model_parallel_world_size()
@@ -257,6 +260,7 @@ def forward(
 
         Returns:
             Output tensor [batch_size, seq_len, hidden_size]
+
         """
         qkv, _ = self.qkv_proj(hidden_states)
         query_states, key_states, value_states = qkv.split([self.qkv_size] * 3, dim=-1)
@@ -311,6 +315,8 @@ def __init__(
             bare_block_idx: Index of the bare block in the model
             num_hybrid_layers: Total number of hybrid layers
             quant_config: Configuration for model quantization
+            prefix: Module prefix used for quantization config lookup
+
         """
         super().__init__()
         self.config = config
@@ -369,6 +375,7 @@ def forward(self, hidden_states: torch.Tensor, block_idx: int) -> torch.Tensor:
         Returns:
             Output tensor [batch_size, seq_len, hidden_size] after applying
             gated feed-forward transformation
+
         """
         # Project input to intermediate size with gating
         gate_up_states, _ = self.gate_up_proj(hidden_states)
@@ -415,6 +422,7 @@ def __init__(
             cache_config: Configuration for key-value caching
             quant_config: Configuration for model quantization
             prefix: Optional prefix for parameter names
+
         """
         super().__init__()
 
@@ -461,8 +469,8 @@ def forward(
 
         Returns:
             Transformed hidden states after attention and feed-forward
-        """
 
+        """
         # The argument original_hidden_states is concatenated with hidden_states
         # (which is the output of the previous (mamba) layer).
         # The concatenated tensor is then used as input of the pre-attention
@@ -510,7 +518,11 @@ def __init__(
 
         Args:
             config: The Zamba2 model configuration
+            model_config: The model config, when available
+            cache_config: The KV cache config, when available
             quant_config: Configuration for model quantization
+            prefix: Module prefix used for quantization config lookup
+
         """
         super().__init__()
 
@@ -555,6 +567,7 @@ def forward(
 
         Returns:
             Transformed hidden states with residual connection applied
+
         """
         # Store input for residual connection
         residual = hidden_states
@@ -603,6 +616,13 @@ def __init__(
 
         Args:
             shared_transformer: Transformer decoder layer for attention pathway
+            config: The Zamba2 model configuration
+            block_idx: Index of this hybrid block in the model
+            model_config: The model config, when available
+            cache_config: The KV cache config, when available
+            quant_config: Configuration for model quantization
+            prefix: Module prefix used for quantization config lookup
+
         """
         super().__init__()
         self.block_idx = block_idx
@@ -644,6 +664,7 @@ def forward(
 
         Returns:
             Output tensor combining transformer and Mamba representations
+
         """
         # Process through transformer pathway
         transformer_hidden_states = self.shared_transformer(
@@ -680,6 +701,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             vllm_config: Configuration object containing model, cache,
                 quantization and LoRA settings
             prefix: Optional prefix for parameter names in state dict
+
         """
         super().__init__()
 
@@ -765,6 +787,7 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
 
         Returns:
             Embedded representation of the input tokens
+
         """
         return self.embed_tokens(input_ids)
 
@@ -784,6 +807,7 @@ def forward(
         Returns:
             Either final hidden states or intermediate tensors for pipeline
             parallelism
+
         """
         # Handle pipeline parallelism for first rank
         if inputs_embeds is None:
@@ -853,8 +877,8 @@ def get_mamba_state_shape_from_config(
             Tuple containing:
             - conv_state_shape: Shape for convolutional state cache
             - temporal_state_shape: Shape for state space model cache
-        """
 
+        """
         parallel_config = vllm_config.parallel_config
         hf_config = vllm_config.model_config.hf_config
         intermediate_size = hf_config.mamba_expand * hf_config.hidden_size
@@ -884,6 +908,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         Raises:
             AssertionError: If prefix caching is enabled
                 (not supported by Mamba)
+
         """
         config = vllm_config.model_config.hf_config
 
@@ -914,10 +939,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         """Convert input token IDs to embeddings.
+
         Args:
             input_ids: Tensor of input token IDs
         Returns:
             Embedded representation of the input tokens
+
         """
         return self.model.embed_input_ids(input_ids)
 
@@ -938,6 +965,7 @@ def forward(
 
         Returns:
             Output hidden states
+
         """
         # Forward pass through model
         hidden_states = self.model(
@@ -959,6 +987,7 @@ def compute_logits(
 
         Returns:
             Logits for next token prediction
+
         """
         logits = self.logits_processor(self.lm_head, hidden_states)
         return logits
diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py
index cb9c35b79c50..3206df34e890 100644
--- a/vllm/model_executor/offloader/base.py
+++ b/vllm/model_executor/offloader/base.py
@@ -77,6 +77,7 @@ def wrap_modules(
 
         Returns:
             List of modules, potentially with offloading hooks installed.
+
         """
         pass
 
@@ -128,14 +129,14 @@ def get_offloader() -> BaseOffloader:
     return _instance
 
 
-def set_offloader(instance: BaseOffloader) -> None:
-    """Set the global offloader instance."""
+def set_offloader(instance: BaseOffloader | None) -> None:
+    """Set or reset the global offloader instance."""
     global _instance
-    _instance = instance
-    if isinstance(instance, NoopOffloader):
+    _instance = NoopOffloader() if instance is None else instance
+    if isinstance(_instance, NoopOffloader):
         logger.debug_once("Offloader set to NoopOffloader (no offloading).")
     else:
-        logger.info_once("Offloader set to %s", type(instance).__name__)
+        logger.info_once("Offloader set to %s", type(_instance).__name__)
 
 
 def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py
index e682fd0e038b..b87535ac376e 100644
--- a/vllm/model_executor/offloader/prefetch.py
+++ b/vllm/model_executor/offloader/prefetch.py
@@ -151,6 +151,7 @@ class PrefetchOffloader(BaseOffloader):
         num_in_group: Offload this many layers per group (last N of each group).
         prefetch_step: Number of layers to prefetch ahead.
         mode: Offload mode ("cpu" is currently supported).
+
     """
 
     def __init__(
diff --git a/vllm/model_executor/offloader/prefetch_ops.py b/vllm/model_executor/offloader/prefetch_ops.py
index b7f599210a19..68a78fbd0200 100644
--- a/vllm/model_executor/offloader/prefetch_ops.py
+++ b/vllm/model_executor/offloader/prefetch_ops.py
@@ -29,6 +29,7 @@ def _wait_prefetch_impl(
         input_tensor: Input to the layer (e.g., hidden_states) - declared
             as mutated to create data dependency for torch.compile.
         layer_idx: Index of the layer to wait for.
+
     """
     get_offloader()._wait_for_layer(layer_idx)
 
@@ -49,6 +50,7 @@ def _start_prefetch_impl(
             prevent torch.compile from reordering this op before the
             computation that produces output_tensor.
         layer_idx: Index of the layer to prefetch.
+
     """
     get_offloader()._start_prefetch(layer_idx)
 
diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py
index f04d4d85b1a4..5b76d2a9a1f0 100644
--- a/vllm/model_executor/offloader/uva.py
+++ b/vllm/model_executor/offloader/uva.py
@@ -34,6 +34,7 @@ class UVAOffloader(BaseOffloader):
         cpu_offload_max_bytes: Maximum bytes to offload to CPU.
         cpu_offload_params: Set of parameter name segments to selectively
             offload. If empty, all parameters are eligible up to the byte limit.
+
     """
 
     supports_tower_offload = True
diff --git a/vllm/model_executor/parameter.py b/vllm/model_executor/parameter.py
index 7f96ceda09db..de1cc7d21888 100644
--- a/vllm/model_executor/parameter.py
+++ b/vllm/model_executor/parameter.py
@@ -30,8 +30,7 @@
 
 
 class BasevLLMParameter(Parameter):
-    """
-    Base parameter for vLLM linear layers. Extends the torch.nn.parameter
+    """Base parameter for vLLM linear layers. Extends the torch.nn.parameter
     by taking in a linear weight loader. Will copy the loaded weight
     into the parameter when the provided weight loader is called.
     """
@@ -40,14 +39,13 @@ def __new__(cls, data: torch.Tensor | None, **kwargs):
         return super().__new__(cls, data=data, requires_grad=False)
 
     def __init__(self, data: torch.Tensor, weight_loader: Callable):
-        """
-        Initialize the BasevLLMParameter
+        """Initialize the BasevLLMParameter.
 
         Args:
             data: torch tensor with the parameter data
             weight_loader: weight loader callable
-        """
 
+        """
         # During weight loading, we often do something like:
         # narrowed_tensor = param.data.narrow(0, offset, len)
         # narrowed_tensor.copy_(real_weight)
@@ -127,8 +125,7 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
 
 
 class _ColumnvLLMParameter(BasevLLMParameter):
-    """
-    Private class defining weight loading functionality
+    """Private class defining weight loading functionality
     (load_merged_column_weight, load_qkv_weight)
     for parameters being loaded into linear layers with column
     parallelism. This includes QKV and MLP layers which are
@@ -202,8 +199,7 @@ def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):
 
 
 class RowvLLMParameter(BasevLLMParameter):
-    """
-    Parameter class defining weight_loading functionality
+    """Parameter class defining weight_loading functionality
     (load_row_parallel_weight) for parameters being loaded
     into linear layers with row parallel functionality.
     Requires an input_dim to be defined.
@@ -231,8 +227,7 @@ def load_row_parallel_weight(self, loaded_weight: torch.Tensor):
 
 
 class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter):
-    """
-    Parameter class for linear layer weights. Uses both column and
+    """Parameter class for linear layer weights. Uses both column and
     row parallelism.
     """
 
@@ -240,8 +235,7 @@ class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter):
 
 
 class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
-    """
-    Parameter class for weight scales loaded for weights with
+    """Parameter class for weight scales loaded for weights with
     grouped quantization. Uses both column and row parallelism.
     """
 
@@ -249,8 +243,7 @@ class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
 
 
 class ChannelQuantScaleParameter(_ColumnvLLMParameter):
-    """
-    Parameter class for weight scales loaded for weights with
+    """Parameter class for weight scales loaded for weights with
     channel-wise quantization. Equivalent to _ColumnvLLMParameter.
     """
 
@@ -258,8 +251,7 @@ class ChannelQuantScaleParameter(_ColumnvLLMParameter):
 
 
 class PerTensorScaleParameter(BasevLLMParameter):
-    """
-    Parameter class for scales where the number of scales is
+    """Parameter class for scales where the number of scales is
     equivalent to the number of logical matrices in fused linear
     layers (e.g. for QKV, there are 3 scales loaded from disk).
     This is relevant to weights with per-tensor quantization.
@@ -291,11 +283,9 @@ def load_column_parallel_weight(self, *args, **kwargs):
     def _load_into_shard_id(
         self, loaded_weight: torch.Tensor, shard_id: str | int, **kwargs
     ):
-        """
-        Slice the parameter data based on the shard id for
+        """Slice the parameter data based on the shard id for
         loading.
         """
-
         param_data = self.data
         shard_id = self._shard_id_as_int(shard_id)
 
@@ -311,8 +301,7 @@ def _load_into_shard_id(
 
 
 class PackedColumnParameter(_ColumnvLLMParameter):
-    """
-    Parameter for model parameters which are packed on disk
+    """Parameter for model parameters which are packed on disk
     and support column parallelism only. See PackedvLLMParameter
     for more details on the packed properties.
     """
@@ -351,8 +340,7 @@ def adjust_shard_indexes_for_packing(self, shard_size, shard_offset):
 
 
 class PackedvLLMParameter(ModelWeightParameter):
-    """
-    Parameter for model weights which are packed on disk.
+    """Parameter for model weights which are packed on disk.
     Example: GPTQ Marlin weights are int4 or int8, packed into int32.
     Extends the ModelWeightParameter to take in the
     packed factor, the packed dimension, and optionally, marlin
@@ -395,8 +383,7 @@ def adjust_shard_indexes_for_packing(self, shard_size, shard_offset):
 
 
 class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
-    """
-    Parameter class for weight scales loaded for weights with
+    """Parameter class for weight scales loaded for weights with
     block-wise quantization. Uses both column and row parallelism.
     """
 
@@ -404,8 +391,7 @@ class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
 
 
 class SharedWeightParameter(BasevLLMParameter):
-    """
-    Parameter for weights with many shared tensors across a model
+    """Parameter for weights with many shared tensors across a model.
 
     For example, when applying transforms to the "gate" and "up" partitions of
     `MergedColumnParallelLinear`, the transform weights must stay separate
@@ -446,8 +432,7 @@ def __init__(self, input_dim: int = 1, output_dim: int = 0, **kwargs):
             )
 
     def add_partition(self, index: int, data_key: Hashable, *args: Any, **kwargs: Any):
-        """
-        Add a partition to the weight parameter. Partitions whose `data_key`
+        """Add a partition to the weight parameter. Partitions whose `data_key`
         is the same will share tensor data
 
         Args:
@@ -455,6 +440,7 @@ def add_partition(self, index: int, data_key: Hashable, *args: Any, **kwargs: An
             data_key: hashable key used to key shared tensors
             *args: arguments for `torch.empty`
             **kwargs: keyword arguments for `torch.empty`
+
         """
         # load (shared) tensor using `data_key`
         if data_key not in self.tensors_registry:
@@ -544,8 +530,7 @@ def _fake_weight_loader(
 def permute_param_layout_(
     param: BasevLLMParameter, input_dim: int, output_dim: int, **kwargs
 ) -> BasevLLMParameter:
-    """
-    Permute a parameter's layout to the specified input and output dimensions,
+    """Permute a parameter's layout to the specified input and output dimensions,
     useful for forcing the parameter into a known layout, for example, if I need
     a packed (quantized) weight matrix to be in the layout
         {input_dim = 0, output_dim = 1, packed_dim = 0}
@@ -554,7 +539,6 @@ def permute_param_layout_(
     to ensure x is in the correct layout (permuting it to the correct layout if
     required, asserting if it cannot get it to the correct layout)
     """
-
     curr_input_dim = getattr(param, "input_dim", None)
     curr_output_dim = getattr(param, "output_dim", None)
 
diff --git a/vllm/model_executor/utils.py b/vllm/model_executor/utils.py
index 6c5c517d9ce6..a31979498890 100644
--- a/vllm/model_executor/utils.py
+++ b/vllm/model_executor/utils.py
@@ -45,6 +45,7 @@ def set_weight_attrs(
     Args:
         weight: The weight tensor.
         weight_attrs: A dictionary of attributes to set on the weight tensor.
+
     """
     if weight_attrs is None:
         return
@@ -73,9 +74,8 @@ def replace_parameter(
     new_tensor: torch.Tensor | None,
     prefer_copy: bool = False,
 ):
-    """
-    Replace a parameter of a layer while maintaining the ability to reload the weight.
-    Called within implementations of the `process_weights_after_loading` method.
+    """Replace a parameter of a layer while maintaining the ability to reload the
+    weight. Called within implementations of the `process_weights_after_loading` method.
 
     Custom attributes set on ``new_tensor`` (e.g. kernel dispatch flags such as
     ``is_shuffled``) are carried over to the replacement parameter, except
@@ -97,6 +97,7 @@ def replace_parameter(
             a new parameter. This preserves the parameter's storage address
             (``data_ptr``), which is required for captured CUDA graphs to
             remain valid across weight updates (e.g. in RL training loops).
+
     """
     # should not be used on a tied/shared param
 
diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py
index 44f2121659ea..1266f4ccb265 100644
--- a/vllm/model_executor/warmup/deep_gemm_warmup.py
+++ b/vllm/model_executor/warmup/deep_gemm_warmup.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Warmup deep_gemm kernels.
+"""Warmup deep_gemm kernels.
 DeepGEMM JIT's the kernels. The warmup aims to JIT all the kernels that would
 be used during model execution beforehand.
 """
@@ -22,9 +21,6 @@
 from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
     TritonOrDeepGemmExperts,
 )
-from vllm.model_executor.layers.linear import LinearBase
-from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
-from vllm.model_executor.layers.quantization.online.mxfp8 import Mxfp8OnlineLinearMethod
 from vllm.tracing import instrument
 from vllm.utils.deep_gemm import (
     fp8_gemm_nt,
@@ -34,21 +30,21 @@
 )
 from vllm.utils.math_utils import cdiv
 from vllm.utils.platform_utils import num_compute_units
+from vllm.v1.worker.workspace import current_workspace_manager
 
 
 def _generate_optimal_warmup_m_values(
     max_tokens: int, n: int, device: torch.device
 ) -> list[int]:
-    """
-    Generate M values that cover all possible DeepGEMM kernel configurations.
+    """Generate M values that cover all possible DeepGEMM kernel configurations.
     Reference: https://github.com/deepseek-ai/DeepGEMM/blob/79f48ee15a82dd5fad5cd9beaa393c1f755e6b55/csrc/jit_kernels/heuristics/common.hpp
 
     Args:
         max_tokens: Maximum number of tokens to warmup for
         n: The actual N dimension from the weight tensor
         device: The torch device to get properties from.
-    """
 
+    """
     # DeepGEMM's possible block sizes
     block_ms = [64, 128, 256]
     block_ns = list(range(16, min(257, n + 1), 16))
@@ -82,34 +78,10 @@ def _generate_optimal_warmup_m_values(
     return sorted(m_values)
 
 
-def _extract_data_from_linear_base_module(
-    m: torch.nn.Module,
-) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
-    """
-    Extract weights, weight scales and quantization block sizes from the given
-    LinearBase module.
-    """
-    assert isinstance(m, LinearBase)
-    assert isinstance(m.quant_method, Fp8LinearMethod)
-    assert m.quant_method.block_quant
-    assert m.quant_method.quant_config is not None
-
-    w = m.weight
-    ws = m.weight_scale_inv if hasattr(m, "weight_scale_inv") else m.weight_scale
-    quant_block_size = m.quant_method.quant_config.weight_block_size
-
-    assert isinstance(w, torch.Tensor)
-    assert isinstance(ws, torch.Tensor)
-    assert quant_block_size is not None
-    return (w, ws, quant_block_size)
-
-
 def _extract_data_from_fused_moe_module(
     m_: torch.nn.Module,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
-    """
-    Extract weights, weight scales and num_topk from MoERunner module.
-    """
+    """Extract weights, weight scales and num_topk from MoERunner module."""
     assert isinstance(m_, MoERunner)
     m = m_.routed_experts
     w13 = m.w13_weight
@@ -133,45 +105,18 @@ def _extract_data_from_fused_moe_module(
     return w13, w13_s, w2, w2_s, num_topk
 
 
-def _is_deep_gemm_backed_kernel(fp8_linear: object) -> bool:
-    """
-    Return True if the selected linear kernel dispatches to DeepGEMM, either
-    directly or as the fallback branch of a dynamic wrapper.
-    """
-    if isinstance(fp8_linear, DeepGemmFp8BlockScaledMMKernel):
-        return True
-    return isinstance(
-        getattr(fp8_linear, "fallback", None), DeepGemmFp8BlockScaledMMKernel
-    )
-
-
-def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
+def _deep_gemm_linear_data(
+    module: torch.nn.Module,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    """Return (weight, weight_scale) for a layer whose kernel stamped itself as
+    ``deep_gemm_warmup_provider`` in ``process_weights_after_loading``, else None.
     """
-    Return True if the input module/layer could be processed with DeepGEMM.
-    """
-
-    if not (
-        isinstance(module, LinearBase)
-        and isinstance(module.quant_method, Fp8LinearMethod)
-        and not isinstance(module.quant_method, Mxfp8OnlineLinearMethod)
-        and getattr(module.quant_method, "block_quant", False)
-        and not getattr(module.quant_method, "use_marlin", True)
-    ):
-        return False
-
-    fp8_linear = getattr(module.quant_method, "fp8_linear", None)
-    if not _is_deep_gemm_backed_kernel(fp8_linear):
-        return False
-
-    block_size = get_mk_alignment_for_contiguous_layout()[0]
-
-    w, _, block_sizes = _extract_data_from_linear_base_module(module)
-    return (
-        block_sizes == get_mk_alignment_for_contiguous_layout()
-        and w.ndim == 2
-        and w.shape[0] % block_size == 0
-        and w.shape[1] % block_size == 0
-    )
+    provider = getattr(module, "deep_gemm_warmup_provider", None)
+    if not isinstance(provider, DeepGemmFp8BlockScaledMMKernel):
+        return None
+    if module.weight_block_size != get_mk_alignment_for_contiguous_layout():
+        return None
+    return provider.get_deep_gemm_warmup_weights(module)
 
 
 def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
@@ -331,15 +276,15 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
     )
     if not warmup_cases:
         return
-    device = w1.device
 
     def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
         _, n, k = w.size()
-        a1q = torch.empty((MAX_M, k), device=device, dtype=torch.float8_e4m3fn)
-        a1q_scales = torch.zeros(
-            (MAX_M, k // block_m), device=device, dtype=torch.float32
+        a1q, a1q_scales, out = current_workspace_manager().get_simultaneous(
+            ((MAX_M, k), torch.float8_e4m3fn),
+            ((MAX_M, k // block_m), torch.float32),
+            ((MAX_M, n), torch.bfloat16),
         )
-        out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)
+        a1q_scales.zero_()
 
         for num_tokens, align_used, expert_ids in warmup_cases:
             with mk_alignment_scope(align_used):
@@ -361,11 +306,10 @@ def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
 def deepgemm_fp8_gemm_nt_warmup(
     model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None
 ):
-    dg_modules = [m for m in model.modules() if _fp8_linear_may_use_deep_gemm(m)]
-
-    for dgm in dg_modules:
-        w, ws, _ = _extract_data_from_linear_base_module(dgm)
-        _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens, pbar=pbar)
+    for m in model.modules():
+        if (data := _deep_gemm_linear_data(m)) is not None:
+            w, ws = data
+            _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens, pbar=pbar)
 
 
 def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
@@ -392,8 +336,8 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
 
     total = 0
     for m in model.modules():
-        if _fp8_linear_may_use_deep_gemm(m):
-            w, _, _ = _extract_data_from_linear_base_module(m)
+        if (data := _deep_gemm_linear_data(m)) is not None:
+            w, _ = data
             if w.size() not in seen_fp8_sizes:
                 total += len(_get_fp8_gemm_nt_m_values(w, max_tokens))
                 seen_fp8_sizes.add(w.size())
diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py
deleted file mode 100644
index 6ca8c94fff0a..000000000000
--- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py
+++ /dev/null
@@ -1,208 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Warm up DeepSeek V4 mHC TileLang kernels before serving requests.
-
-Ported from lucifer1004/vllm-jasl with the two env-var knobs removed
-(`VLLM_ENABLE_DEEPSEEK_V4_MHC_WARMUP`, `VLLM_DEEPSEEK_V4_MHC_WARMUP_TOKEN_SIZES`).
-Gating is intrinsic: non-DSv4 models and layers without hc_* attributes
-return early, so the warmup is a no-op except where it's needed.
-"""
-
-import time
-from collections.abc import Iterable
-
-import torch
-
-from vllm.logger import init_logger
-from vllm.tracing import instrument
-
-logger = init_logger(__name__)
-
-_AUTO_WARMUP_MAX_TOKENS = 16_384
-_DEFAULT_TOKEN_SIZE_CANDIDATES = (
-    1,
-    2,
-    4,
-    8,
-    16,
-    32,
-    64,
-    128,
-    256,
-    512,
-    1024,
-    2048,
-    4096,
-    8192,
-    16_384,
-)
-
-
-def _normalize_token_sizes(
-    token_sizes: Iterable[int],
-    *,
-    max_tokens: int,
-) -> list[int]:
-    return sorted({size for size in token_sizes if 1 <= size <= max_tokens})
-
-
-def _select_mhc_warmup_token_sizes(
-    *,
-    max_tokens: int,
-    cudagraph_capture_sizes: list[int],
-) -> list[int]:
-    if max_tokens <= 0:
-        return []
-
-    max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS)
-    candidates = list(_DEFAULT_TOKEN_SIZE_CANDIDATES)
-    candidates.extend(cudagraph_capture_sizes)
-    candidates.append(max_auto_tokens)
-    return _normalize_token_sizes(candidates, max_tokens=max_auto_tokens)
-
-
-def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None:
-    for module in model.modules():
-        if module.__class__.__name__ != "DeepseekV4DecoderLayer":
-            continue
-        if all(
-            hasattr(module, attr)
-            for attr in (
-                "hc_pre",
-                "hc_post",
-                "hc_attn_fn",
-                "hc_attn_scale",
-                "hc_attn_base",
-                "hc_ffn_fn",
-                "hc_ffn_scale",
-                "hc_ffn_base",
-            )
-        ):
-            return module
-    return None
-
-
-def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None:
-    for module in model.modules():
-        if module.__class__.__name__ != "DeepseekV4Model":
-            continue
-        if all(
-            hasattr(module, attr)
-            for attr in ("hc_head_fn", "hc_head_scale", "hc_head_base")
-        ):
-            return module
-    return None
-
-
-def _warmup_layer_mhc(
-    layer: torch.nn.Module,
-    token_sizes: list[int],
-) -> None:
-    max_tokens = max(token_sizes)
-    hidden_size = int(layer.hidden_size)
-    hc_mult = int(layer.hc_mult)
-    device = layer.hc_attn_fn.device
-    residual = torch.zeros(
-        max_tokens,
-        hc_mult,
-        hidden_size,
-        dtype=torch.bfloat16,
-        device=device,
-    )
-
-    for size in token_sizes:
-        residual_slice = residual[:size]
-        for fn, scale, base in (
-            (layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base),
-            (layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base),
-        ):
-            layer_input, post_mix, comb_mix = layer.hc_pre(
-                residual_slice,
-                fn,
-                scale,
-                base,
-            )
-            layer.hc_post(layer_input, residual_slice, post_mix, comb_mix)
-
-
-def _warmup_hc_head(
-    model: torch.nn.Module,
-    token_sizes: list[int],
-) -> None:
-    # Upstream a8887c208 ("[DSV4] aiter mhc support (ROCm)") refactored
-    # ``hc_head`` from a free function into the ``HCHeadOp`` CustomOp
-    # instance attached to the model as ``hc_head_op``. We call through
-    # that instance so the warmup exercises the same dispatched
-    # implementation as the inference path.
-    hc_head_op = getattr(model, "hc_head_op", None)
-    if hc_head_op is None:
-        return
-
-    max_tokens = max(token_sizes)
-    hidden_size = int(model.config.hidden_size)
-    hc_mult = int(model.hc_mult)
-    device = model.hc_head_fn.device
-    hidden_states = torch.zeros(
-        max_tokens,
-        hc_mult,
-        hidden_size,
-        dtype=torch.bfloat16,
-        device=device,
-    )
-
-    for size in token_sizes:
-        hc_head_op(
-            hidden_states[:size],
-            model.hc_head_fn,
-            model.hc_head_scale,
-            model.hc_head_base,
-            model.rms_norm_eps,
-            model.hc_eps,
-        )
-
-
-@instrument(span_name="DeepSeek V4 mHC warmup")
-def deepseek_v4_mhc_warmup(
-    model: torch.nn.Module,
-    *,
-    max_tokens: int,
-    cudagraph_capture_sizes: list[int] | None = None,
-) -> None:
-    # Cheap model-type gate before walking ``model.modules()``. The class
-    # walk below is O(num_layers) and shows up in startup time on very
-    # large checkpoints; bail out for any model that is not DeepSeek V4.
-    config = getattr(model, "config", None)
-    model_type = getattr(config, "model_type", None) if config is not None else None
-    if model_type is not None and model_type != "deepseek_v4":
-        return
-
-    layer = _find_first_mhc_layer(model)
-    if layer is None:
-        return
-
-    device = layer.hc_attn_fn.device
-    if device.type != "cuda":
-        return
-
-    deepseek_model = _find_deepseek_v4_model(model)
-    token_sizes = _select_mhc_warmup_token_sizes(
-        max_tokens=max_tokens,
-        cudagraph_capture_sizes=cudagraph_capture_sizes or [],
-    )
-    if not token_sizes:
-        return
-
-    started = time.perf_counter()
-    logger.info(
-        "Warming up DeepSeek V4 mHC TileLang kernels for token sizes: %s",
-        token_sizes,
-    )
-    with torch.inference_mode():
-        _warmup_layer_mhc(layer, token_sizes)
-        if deepseek_model is not None:
-            _warmup_hc_head(deepseek_model, token_sizes)
-        torch.accelerator.synchronize()
-    logger.info(
-        "DeepSeek V4 mHC TileLang warmup finished in %.2f seconds.",
-        time.perf_counter() - started,
-    )
diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py
index 37ea4812fc6c..afc3ef3ac014 100644
--- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py
+++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py
@@ -12,6 +12,7 @@
 import vllm.envs as envs
 
 if TYPE_CHECKING:
+    from vllm.distributed.parallel_state import GroupCoordinator
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 
@@ -40,6 +41,44 @@ def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
     return output_dir / "autotune_configs.json"
 
 
+def sync_flashinfer_autotune_cache(
+    runner: "GPUModelRunner",
+    group: "GroupCoordinator",
+) -> None:
+    cache: bytes | str | None = None
+    if (
+        group.rank_in_group == 0
+        and runner.vllm_config.kernel_config.enable_flashinfer_autotune
+    ):
+        try:
+            from vllm.platforms import current_platform
+            from vllm.utils.flashinfer import has_flashinfer
+
+            if has_flashinfer() and current_platform.has_device_capability(90):
+                from flashinfer.autotuner import AutoTuner
+
+                with tempfile.TemporaryDirectory() as temp_dir:
+                    path = Path(temp_dir) / "autotune_configs.json"
+                    AutoTuner.get().save_configs(str(path))
+                    cache = path.read_bytes()
+        except Exception as exc:
+            cache = f"{type(exc).__name__}: {exc}"
+
+    cache = group.broadcast_object(cache)
+    if isinstance(cache, str):
+        raise RuntimeError(f"Failed to serialize FlashInfer autotune state: {cache}")
+    if cache is None or group.rank_in_group == 0:
+        return
+
+    from flashinfer.autotuner import AutoTuner
+
+    with tempfile.NamedTemporaryFile() as f:
+        f.write(cache)
+        f.flush()
+        if not AutoTuner.get().load_configs(f.name):
+            raise RuntimeError("FlashInfer autotune cache is incompatible")
+
+
 def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None:
     cache_path.parent.mkdir(parents=True, exist_ok=True)
     fd, tmp_path = tempfile.mkstemp(
diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
index 80dfd5199614..91ae7c44c239 100644
--- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
+++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
@@ -28,20 +28,54 @@
         "FLASHMLA_SPARSE_DSV4",
         "FLASHINFER_MLA_SPARSE_DSV4",
         "ROCM_FLASHMLA_SPARSE_DSV4",
+        "FLASHMLA_SPARSE_DSV41",
+        "FLASHINFER_MLA_SPARSE_DSV41",
         "DEEPSEEK_SPARSE_SWA",
     }
 )
 _FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_SM120"})
-_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_DSV4"})
+_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset(
+    {"FLASHINFER_MLA_SPARSE_DSV4", "FLASHINFER_MLA_SPARSE_DSV41"}
+)
 
 _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = {
     "FLASHINFER_MLA_SPARSE_SM120": "DSv3.2",
     "FLASHINFER_MLA_SPARSE_DSV4": "DSv4",
+    "FLASHINFER_MLA_SPARSE_DSV41": "DSv4.1",
 }
 
 _SPARSE_MLA_MIXED_WARMUP_TOKENS = 16
 
 
+def autotune_hisparse_flashinfer_attention(runner: "GPUModelRunner") -> None:
+    """Autotune each HiSparse FlashInfer sparse-MLA configuration."""
+    from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
+        FlashInferMLASparseImpl,
+    )
+
+    tuned: set[tuple[object, ...]] = set()
+    for layer in runner.vllm_config.compilation_config.static_forward_context.values():
+        impl = getattr(layer, "impl", None)
+        if not isinstance(impl, FlashInferMLASparseImpl):
+            continue
+        if getattr(layer, "hisparse_cache", None) is None:
+            continue
+        if impl.topk_indices_buffer is None:
+            continue
+        key = (
+            impl.kv_cache_dtype,
+            impl.num_heads,
+            impl.qk_nope_head_dim,
+            impl.qk_rope_head_dim,
+            impl.kv_lora_rank,
+            impl.topk_indices_buffer.shape[1],
+        )
+        if key in tuned:
+            continue
+        impl.autotune_hisparse_decode(layer)
+        tuned.add(key)
+
+
 def _attention_backend_name(backend: object) -> str | None:
     get_name = getattr(backend, "get_name", None)
     if get_name is None:
diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py
index 0be4bddea6b3..84f713aa6489 100644
--- a/vllm/model_executor/warmup/jit_warmup.py
+++ b/vllm/model_executor/warmup/jit_warmup.py
@@ -15,14 +15,17 @@
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass, fields
+from functools import wraps
 from typing import Any, Generic, TypeVar, cast
 
 __all__ = [
     "JitWarmupRegistry",
     "VllmJitKernel",
+    "WarmupChoices",
     "WarmupIntRange",
     "get_ast_full_name",
     "get_function_source_node",
+    "kernel_launcher",
     "zip_inputs",
 ]
 
@@ -30,6 +33,26 @@
 CompileKeyT = TypeVar("CompileKeyT")
 
 
+def kernel_launcher(call_fn: Callable[..., Any]) -> Callable[..., Any]:
+    """Delegate a declarative launch specification to its backend owner."""
+    signature = inspect.signature(call_fn)
+
+    @wraps(call_fn)
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        launch_spec = call_fn(self, *args, **kwargs)
+        if not self.bind_launch_inputs:
+            return self.launch(launch_spec, {})
+
+        bound = signature.bind(self, *args, **kwargs)
+        bound.apply_defaults()
+        inputs = {
+            name: value for name, value in bound.arguments.items() if name != "self"
+        }
+        return self.launch(launch_spec, inputs)
+
+    return wrapper
+
+
 @dataclass(frozen=True)
 class WarmupIntRange:
     """Expand integers with range semantics or a custom monotonic progression."""
@@ -40,6 +63,21 @@ class WarmupIntRange:
     advance: Callable[[int], int] | None = None
 
 
+@dataclass(frozen=True)
+class WarmupChoices:
+    """Expand an explicit finite set of values in a traced warmup method."""
+
+    values: tuple[Any, ...]
+
+    def __init__(self, *values: Any) -> None:
+        object.__setattr__(self, "values", values)
+
+
+def _when(condition: bool) -> None:
+    """Filter symbolic cases in an AST-traced warmup-cases method."""
+    raise RuntimeError("_when() is only valid in an AST-traced warmup method")
+
+
 WarmupValues = Any
 CompileKeyDispatchFn = Callable[..., CompileKeyT]
 WarmupPredicateFn = Callable[..., bool]
@@ -65,6 +103,8 @@ class _WarmupInputRows:
 
 
 def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]:
+    if isinstance(values, WarmupChoices):
+        return values.values
     if isinstance(values, WarmupIntRange):
         if values.advance is None:
             return tuple(range(values.start, values.stop, values.step))
@@ -256,6 +296,27 @@ def visit_Name(self, node: ast.Name) -> Any:
     def visit_Constant(self, node: ast.Constant) -> Any:
         return node.value
 
+    def visit_Lambda(self, node: ast.Lambda) -> Callable[..., Any]:
+        arguments = node.args
+        if (
+            arguments.posonlyargs
+            or arguments.vararg is not None
+            or arguments.kwonlyargs
+            or arguments.kwarg is not None
+            or arguments.defaults
+        ):
+            raise _dispatch_expr_error(node, "Traced lambdas require positional args")
+        names = tuple(argument.arg for argument in arguments.args)
+
+        def evaluate(*values: Any) -> Any:
+            if len(values) != len(names):
+                raise TypeError(f"Expected {len(names)} lambda arguments")
+            return type(self)(
+                dict(self.values) | dict(zip(names, values)), self.globals
+            ).eval(node.body)
+
+        return evaluate
+
     def visit_IfExp(self, node: ast.IfExp) -> Any:
         return self.visit(node.body if self.visit(node.test) else node.orelse)
 
@@ -311,7 +372,12 @@ def visit_BinOp(self, node: ast.BinOp) -> Any:
         return op(self.visit(node.left), self.visit(node.right))
 
     def visit_Call(self, node: ast.Call) -> Any:
-        args = [self.visit(arg) for arg in node.args]
+        args: list[Any] = []
+        for arg in node.args:
+            if isinstance(arg, ast.Starred):
+                args.extend(self.visit(arg.value))
+            else:
+                args.append(self.visit(arg))
         fn = self.visit(node.func)
         call_kwargs: dict[str, Any] = {}
         for keyword in node.keywords:
@@ -366,6 +432,45 @@ def _eval_dispatch_expr(
     return _DispatchExprEvaluator(kwargs, globals_).eval(node)
 
 
+def _validate_dispatch_expr(node: ast.AST) -> None:
+    """Enforce the dispatch AST subset before compiling traced code."""
+    allowed_nodes = (
+        ast.Name,
+        ast.Constant,
+        ast.Lambda,
+        ast.arguments,
+        ast.arg,
+        ast.IfExp,
+        ast.Tuple,
+        ast.List,
+        ast.BoolOp,
+        ast.And,
+        ast.Or,
+        ast.Compare,
+        *tuple(_CMP_OPS),
+        ast.UnaryOp,
+        ast.Not,
+        ast.USub,
+        ast.BinOp,
+        *tuple(_BIN_OPS),
+        ast.Call,
+        ast.keyword,
+        ast.Starred,
+        ast.Attribute,
+        ast.Subscript,
+        ast.Load,
+    )
+    for child in ast.walk(node):
+        if not isinstance(child, allowed_nodes):
+            raise _dispatch_expr_error(child, "Unsupported dispatch expression")
+        if isinstance(child, ast.Call) and any(
+            keyword.arg is None for keyword in child.keywords
+        ):
+            raise _dispatch_expr_error(
+                child, "Dispatch helper calls cannot use **kwargs"
+            )
+
+
 def _collect_input_names(
     node: ast.AST,
     candidate_names: set[str],
@@ -384,6 +489,23 @@ def _collect_input_names(
     }
 
 
+def _named_assignment(
+    statement: ast.stmt,
+    subject: str,
+) -> tuple[str, ast.AST] | None:
+    match statement:
+        case ast.Assign(targets=[ast.Name(id=name)], value=value):
+            return name, value
+        case ast.AnnAssign(target=ast.Name(id=name), value=value) if value is not None:
+            return name, value
+        case ast.Assign() | ast.AnnAssign():
+            raise _dispatch_expr_error(
+                statement, f"{subject} assignments require one name"
+            )
+        case _:
+            return None
+
+
 def _collect_expression_body(
     fn: Callable[..., Any],
     function_def: ast.FunctionDef | ast.Lambda,
@@ -400,38 +522,19 @@ def _collect_expression_body(
         ):
             continue
 
-        if isinstance(statement, ast.Assign):
-            if len(statement.targets) != 1 or not isinstance(
-                statement.targets[0], ast.Name
-            ):
-                raise _dispatch_expr_error(
-                    statement,
-                    "Traced warmup helper assignments must target one local name",
-                )
-            local_exprs.append((statement.targets[0].id, statement.value))
+        if (
+            assignment := _named_assignment(statement, "Traced warmup helper")
+        ) is not None:
+            local_exprs.append(assignment)
             continue
 
-        if isinstance(statement, ast.AnnAssign):
+        if isinstance(statement, ast.Return):
             if statement.value is None:
                 raise _dispatch_expr_error(
-                    statement, "Traced warmup helper annotations must assign a value"
+                    statement, "Traced warmup helper must return an expression"
                 )
-            if not isinstance(statement.target, ast.Name):
-                raise _dispatch_expr_error(
-                    statement,
-                    "Traced warmup helper assignments must target one local name",
-                )
-            local_exprs.append((statement.target.id, statement.value))
-            continue
-
-        if isinstance(statement, ast.Return) and statement.value is not None:
             return local_exprs, statement.value
 
-        if isinstance(statement, ast.Return):
-            raise _dispatch_expr_error(
-                statement, "Traced warmup helper must return an expression"
-            )
-
         raise _dispatch_expr_error(
             statement,
             "Traced warmup helper may only contain local assignments before return",
@@ -549,12 +652,20 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
     """Kernel wrapper that owns dispatch, warmup keys, and compilation."""
 
     CompileKey: type[CompileKeyT]
+    bind_launch_inputs = True
 
     def __init__(self) -> None:
-        self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch)
+        dispatch = type(self).dispatch
+        self._dispatch_trace = (
+            None
+            if dispatch is VllmJitKernel.dispatch
+            else _trace_compile_key_dispatch(self.dispatch)
+        )
         self._compiled_cache: dict[Any, Any] = {}
 
     def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:
+        if self._dispatch_trace is None:
+            raise TypeError(f"{type(self).__name__} does not define dispatch()")
         return self._dispatch_trace.compile_key(self.CompileKey, kwargs)
 
     def _get_or_compile(
@@ -660,7 +771,184 @@ def traced(
 
         return traced
 
-    @abstractmethod
+    def _expand_warmup_cases(
+        self,
+        cases_fn: Callable[..., Any],
+        *args: Any,
+        _value_expander: Callable[[WarmupValues], tuple[Any, ...]] = (
+            _expand_warmup_values
+        ),
+        **kwargs: Any,
+    ) -> Iterator[Mapping[str, Any]]:
+        """Expand symbolic domains declared inside a warmup-cases method."""
+        function_def = get_function_source_node(cases_fn)
+        if isinstance(function_def, ast.Lambda):
+            raise _dispatch_expr_error(
+                function_def, "Warmup cases must be a function definition"
+            )
+        signature = inspect.signature(cases_fn)
+        bound = signature.bind(*args, **kwargs)
+        bound.apply_defaults()
+        static_values = dict(bound.arguments)
+        bound_self = getattr(cases_fn, "__self__", None)
+        if bound_self is not None:
+            static_values["self"] = bound_self
+        globals_ = getattr(cases_fn, "__func__", cases_fn).__globals__
+
+        domains: list[tuple[str, tuple[Any, ...]]] = []
+        local_exprs: list[tuple[str, ast.AST]] = []
+        predicates: list[ast.AST] = []
+        return_expr: ast.AST | None = None
+        for statement in function_def.body:
+            if (
+                isinstance(statement, ast.Expr)
+                and isinstance(statement.value, ast.Constant)
+                and isinstance(statement.value.value, str)
+            ):
+                continue
+            assignment = _named_assignment(statement, "Warmup case")
+            if assignment is not None:
+                name, value_expr = assignment
+                call_name = (
+                    (get_ast_full_name(value_expr.func) or "").split(".")[-1]
+                    if isinstance(value_expr, ast.Call)
+                    else ""
+                )
+                if call_name in {"WarmupIntRange", "WarmupChoices"}:
+                    value = _eval_dispatch_expr(value_expr, static_values, globals_)
+                    domains.append((name, _value_expander(value)))
+                else:
+                    local_exprs.append((name, value_expr))
+                continue
+            if (
+                isinstance(statement, ast.Expr)
+                and isinstance(statement.value, ast.Call)
+                and (get_ast_full_name(statement.value.func) or "").endswith("_when")
+            ):
+                if len(statement.value.args) != 1 or statement.value.keywords:
+                    raise _dispatch_expr_error(
+                        statement, "_when requires exactly one positional expression"
+                    )
+                predicates.append(statement.value.args[0])
+                continue
+            if isinstance(statement, ast.Return):
+                if return_expr is not None or statement.value is None:
+                    raise _dispatch_expr_error(
+                        statement, "Warmup cases require one final return expression"
+                    )
+                return_expr = statement.value
+                continue
+            raise _dispatch_expr_error(
+                statement,
+                "AST-traced warmup cases support assignments, _when, and return",
+            )
+
+        if return_expr is None or not isinstance(return_expr, ast.Call):
+            raise ValueError("AST-traced warmup cases must return one case(...) call")
+        if any(keyword.arg is None for keyword in return_expr.keywords):
+            raise _dispatch_expr_error(
+                return_expr, "Warmup case expressions do not support **kwargs"
+            )
+
+        class _InlineDomainRewriter(ast.NodeTransformer):
+            def visit_Call(self, node: ast.Call) -> ast.AST:
+                call_name = (get_ast_full_name(node.func) or "").split(".")[-1]
+                if call_name not in {"WarmupIntRange", "WarmupChoices"}:
+                    return self.generic_visit(node)
+                name = f"__warmup_domain_{len(domains)}"
+                domain = _eval_dispatch_expr(node, static_values, globals_)
+                domains.append((name, _value_expander(domain)))
+                return ast.copy_location(ast.Name(id=name, ctx=ast.Load()), node)
+
+        return_expr = cast(ast.Call, _InlineDomainRewriter().visit(return_expr))
+
+        domain_names = tuple(name for name, _ in domains)
+        dynamic_names = set(domain_names)
+        dynamic_local_exprs: list[tuple[str, ast.AST]] = []
+        for name, expr in local_exprs:
+            referenced_names = {
+                node.id for node in ast.walk(expr) if isinstance(node, ast.Name)
+            }
+            if referenced_names & dynamic_names:
+                dynamic_local_exprs.append((name, expr))
+                dynamic_names.add(name)
+            else:
+                static_values[name] = _eval_dispatch_expr(expr, static_values, globals_)
+
+        for _, expr in dynamic_local_exprs:
+            _validate_dispatch_expr(expr)
+        for predicate in predicates:
+            _validate_dispatch_expr(predicate)
+        _validate_dispatch_expr(return_expr)
+
+        case_body: list[ast.stmt] = [
+            ast.Assign(
+                targets=[ast.Name(id=name, ctx=ast.Store())],
+                value=cast(ast.expr, expr),
+            )
+            for name, expr in dynamic_local_exprs
+        ]
+        if predicates:
+            predicate = (
+                predicates[0]
+                if len(predicates) == 1
+                else ast.BoolOp(
+                    op=ast.And(),
+                    values=[cast(ast.expr, value) for value in predicates],
+                )
+            )
+            case_body.append(
+                ast.If(
+                    test=ast.UnaryOp(op=ast.Not(), operand=cast(ast.expr, predicate)),
+                    body=[ast.Return(value=ast.Constant(value=None))],
+                    orelse=[],
+                )
+            )
+        case_body.append(ast.Return(value=cast(ast.expr, return_expr)))
+        # FunctionDef fields vary across the supported Python versions, so no
+        # single constructor overload matches every mypy target.
+        case_function = ast.FunctionDef(  # type: ignore[call-overload]
+            name="__vllm_warmup_case",
+            args=ast.arguments(
+                posonlyargs=[],
+                args=[ast.arg(arg=name) for name in domain_names],
+                kwonlyargs=[],
+                kw_defaults=[],
+                defaults=[],
+            ),
+            body=case_body,
+            decorator_list=[],
+            returns=None,
+            type_comment=None,
+        )
+        case_globals = dict(globals_)
+        case_globals.update(static_values)
+        exec(
+            compile(
+                ast.fix_missing_locations(
+                    ast.Module(body=[case_function], type_ignores=[])
+                ),
+                "<jit-warmup>",
+                "exec",
+            ),
+            case_globals,
+        )
+        evaluate_case = cast(Callable[..., Any], case_globals[case_function.name])
+
+        domain_values = tuple(values for _, values in domains)
+        for values in itertools.product(*domain_values):
+            case = evaluate_case(*values)
+            if case is None:
+                continue
+            if not isinstance(case, Mapping):
+                raise TypeError("AST-traced warmup cases must return a mapping")
+            if any(
+                isinstance(value, WarmupChoices | WarmupIntRange)
+                for value in case.values()
+            ):
+                raise TypeError("Warmup domains must be direct expressions")
+            yield case
+
     def dispatch(self, **kwargs: Any) -> CompileKeyT:
         """Build one compile key from one concrete dispatch point."""
         raise NotImplementedError
@@ -675,14 +963,18 @@ def compile(self, compile_key: CompileKeyT) -> None:
         """Compile one warmup key."""
         raise NotImplementedError
 
+    def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None:
+        """Compile a batch of warmup keys, allowing backend-specific scheduling."""
+        for compile_key in compile_keys:
+            self.compile(compile_key)
+
     def register_warmup(self, *args: Any, **kwargs: Any) -> None:
         """Register this kernel with the active runner's warmup registry."""
         JitWarmupRegistry.register(self, *args, **kwargs)
 
     def warmup(self, *args: Any, **kwargs: Any) -> None:
         """Compile this kernel's warmup keys."""
-        for compile_key in self.get_warmup_keys(*args, **kwargs):
-            self.compile(compile_key)
+        self.compile_many(self.get_warmup_keys(*args, **kwargs))
 
 
 def _same_value(left: Any, right: Any) -> bool:
@@ -728,6 +1020,19 @@ def __init__(self, vllm_config: Any) -> None:
             list[tuple[tuple[Any, ...], dict[str, Any]]],
         ] = {}
 
+    @classmethod
+    def capture(cls, init_fn: Callable[..., None]) -> Callable[..., None]:
+        """Collect warmup registrations made while the decorated callable runs."""
+
+        @wraps(init_fn)
+        def wrapped(instance: Any, vllm_config: Any, *args: Any, **kwargs: Any) -> None:
+            registry = cls(vllm_config)
+            instance.jit_warmup_registry = registry
+            with registry.activate():
+                init_fn(instance, vllm_config, *args, **kwargs)
+
+        return wrapped
+
     @contextmanager
     def activate(self) -> Iterator[None]:
         """Collect registrations made in this context."""
@@ -785,9 +1090,10 @@ def warmup(self) -> None:
                 if (
                     not args
                     and not kwargs
-                    and inspect.signature(kernel.get_warmup_keys).parameters
+                    and "vllm_config"
+                    in inspect.signature(kernel.get_warmup_keys).parameters
                 ):
-                    args = (self.vllm_config,)
+                    kwargs = {"vllm_config": self.vllm_config}
                 for compile_key in kernel.get_warmup_keys(*args, **kwargs):
                     compile_keys[compile_key] = None
             if compile_keys:
@@ -809,5 +1115,4 @@ def warmup(self) -> None:
                     f"{kernel.__class__.__name__} ({len(compile_keys)} keys)",
                     refresh=False,
                 )
-                for compile_key in compile_keys:
-                    kernel.compile(compile_key)
+                kernel.compile_many(compile_keys)
diff --git a/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py b/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py
index 7311cc18a419..879c8c6998d1 100644
--- a/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py
+++ b/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py
@@ -3,19 +3,25 @@
 
 from abc import abstractmethod
 from collections.abc import Callable, Mapping
-from functools import wraps
 from typing import Any, ClassVar, Generic, TypeAlias, TypeVar
 
-from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
+from vllm.model_executor.warmup.jit_warmup import (
+    VllmJitKernel,
+)
 
 DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi"
 
 CompileKeyT = TypeVar("CompileKeyT")
-CuTeDSLLaunchSpec: TypeAlias = tuple[
-    CompileKeyT,
-    tuple[Any, ...],
-    Mapping[str, Any] | None,
-]
+CuTeDSLLaunchSpec: TypeAlias = (
+    tuple[CompileKeyT, tuple[Any, ...]]
+    | tuple[CompileKeyT, tuple[Any, ...], Any]
+    | tuple[
+        CompileKeyT,
+        tuple[Any, ...],
+        Any,
+        Callable[[], Any],
+    ]
+)
 
 
 def cutedsl_fake_stream(*, use_tvm_ffi_env_stream: bool = True) -> Any:
@@ -44,6 +50,7 @@ class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
     """CuTeDSL owner whose compiled executor is shared by warmup and runtime."""
 
     kernel: ClassVar[Any]
+    bind_launch_inputs = False
 
     @abstractmethod
     def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]:
@@ -58,23 +65,15 @@ def compile(self, compile_key: CompileKeyT) -> None:
             *self.warmup_inputs(compile_key),
         )
 
-
-def cutedsl_kernel_launcher(
-    call_fn: Callable[..., CuTeDSLLaunchSpec[CompileKeyT]],
-) -> Callable[..., Any]:
-    """Invoke a cached CuTeDSL executor from a declarative ``__call__``."""
-
-    @wraps(call_fn)
-    def wrapper(
-        self: VllmCuTeDSLJitKernel[CompileKeyT],
-        *args: Any,
-        **kwargs: Any,
+    def launch(
+        self,
+        launch_spec: CuTeDSLLaunchSpec[CompileKeyT],
+        _inputs: Mapping[str, Any],
     ) -> Any:
-        compile_key, launch_args, runtime_context = call_fn(self, *args, **kwargs)
-        executor = self._get_or_compile(
-            compile_key,
-            runtime_context=runtime_context,
-        )
-        return executor(*launch_args)
-
-    return wrapper
+        # (compile_key, args), optionally followed by output and epilogue.
+        compile_key, launch_args = launch_spec[:2]
+        executor = self._get_or_compile(compile_key)
+        result = executor(*launch_args)
+        if len(launch_spec) == 4:
+            return launch_spec[3]()
+        return launch_spec[2] if len(launch_spec) == 3 else result
diff --git a/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py
new file mode 100644
index 000000000000..188f26fdcf17
--- /dev/null
+++ b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py
@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compile-only helpers for TileLang JIT warmup."""
+
+from __future__ import annotations
+
+import logging
+from abc import abstractmethod
+from collections.abc import Callable, Iterator, Mapping, MutableMapping
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from functools import wraps
+from typing import Any, ClassVar, Generic, TypeAlias, TypeVar, cast
+
+import torch
+
+from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
+
+CompileKeyT = TypeVar("CompileKeyT")
+# (kernel_args, launch_args), optionally followed by keywords and a result.
+TileLangLaunchSpec: TypeAlias = (
+    tuple[tuple[Any, ...], tuple[Any, ...]]
+    | tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None]
+    | tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None, Any]
+)
+
+
+@contextmanager
+def _quiet_tilelang_warmup_logs() -> Iterator[None]:
+    loggers = [
+        logging.getLogger("tilelang.cache.kernel_cache"),
+        logging.getLogger("tilelang.jit.kernel"),
+    ]
+    previous_levels = [logger.level for logger in loggers]
+    try:
+        for logger in loggers:
+            logger.setLevel(logging.ERROR)
+        yield
+    finally:
+        for logger, level in zip(loggers, previous_levels):
+            logger.setLevel(level)
+
+
+@dataclass(frozen=True)
+class TileLangWarmupTensor:
+    """Minimal tensor-like object accepted by TileLang compile().
+
+    TileLang builds its cache key and TIR from tensor dtype, shape and stride.
+    This object deliberately has no storage, so compile-only warmup does not
+    allocate GPU memory and does not launch the kernel.
+    """
+
+    dtype: torch.dtype
+    shape: tuple[int, ...] = (1,)
+    strides: tuple[int, ...] | None = field(default=None)
+
+    def stride(self) -> tuple[int, ...]:
+        if self.strides is not None:
+            return self.strides
+
+        strides: list[int] = []
+        stride = 1
+        for size in reversed(self.shape):
+            strides.append(stride)
+            stride *= size
+        return tuple(reversed(strides))
+
+    def new_empty(
+        self,
+        shape: tuple[int, ...],
+        *,
+        dtype: torch.dtype | None = None,
+        device: Any = None,
+    ) -> TileLangWarmupTensor:
+        return TileLangWarmupTensor(dtype or self.dtype, shape)
+
+
+def make_tilelang_warmup_tensor(
+    dtype: torch.dtype,
+    *shape: int,
+    strides: tuple[int, ...] | None = None,
+) -> TileLangWarmupTensor:
+    return TileLangWarmupTensor(dtype=dtype, shape=tuple(shape), strides=strides)
+
+
+def _tilelang_call_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+    call_kwargs = dict(kwargs)
+    tune_params = call_kwargs.pop("__tune_params", {})
+    if isinstance(tune_params, dict):
+        call_kwargs.update(tune_params)
+    return call_kwargs
+
+
+def compile_tilelang(jit_impl: Any, *args: Any, **kwargs: Any) -> None:
+    """Compile one TileLang specialization and populate its call cache.
+
+    TileLang's ``compile()`` materializes the kernel without launching it.
+    We also store the compiled kernel in ``_kernel_cache`` using the same
+    parsed key as ``__call__`` so runtime does not report a cache miss for an
+    already materialized specialization.
+    """
+    with _quiet_tilelang_warmup_logs():
+        compiled = jit_impl.compile(*args, **kwargs)
+    func = getattr(jit_impl, "func", None)
+    parse_args = getattr(func, "parse_args", None)
+    cache = getattr(jit_impl, "_kernel_cache", None)
+    if not callable(parse_args) or not isinstance(cache, MutableMapping):
+        return
+
+    key, _ = parse_args(*args, **_tilelang_call_kwargs(kwargs))
+    cache[key] = compiled
+
+
+class VllmTileLangJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
+    """TileLang owner whose runtime launch specification is reused for warmup."""
+
+    kernel: ClassVar[Any]
+    _warming_key: CompileKeyT | None = None
+
+    @abstractmethod
+    def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
+        """Return runtime-shaped inputs that reproduce one compile key."""
+        raise NotImplementedError
+
+    def compile(self, compile_key: CompileKeyT) -> None:
+        self._warming_key = compile_key
+        try:
+            cast(Callable[..., Any], self)(**self.warmup_inputs(compile_key))
+        finally:
+            self._warming_key = None
+
+    def launch(
+        self,
+        jit_impl: Any,
+        args: tuple[Any, ...],
+        kwargs: Mapping[str, Any] | None = None,
+    ) -> Any:
+        call_kwargs = dict(kwargs or {})
+        if self._warming_key is not None:
+            return compile_tilelang(jit_impl, *args, **call_kwargs)
+        return jit_impl(*args, **call_kwargs)
+
+
+def kernel_launcher(
+    call_fn: Callable[..., TileLangLaunchSpec],
+) -> Callable[..., Any]:
+    """Launch TileLang from declarative kernel and runtime argument tuples."""
+
+    @wraps(call_fn)
+    def wrapper(
+        self: VllmTileLangJitKernel[Any],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        launch_spec = call_fn(self, *args, **kwargs)
+        kernel_args, launch_args = launch_spec[:2]
+        if self._warming_key is not None and kernel_args not in (
+            (),
+            (self._warming_key,),
+        ):
+            raise RuntimeError(
+                "TileLang warmup inputs produced a different compile key: "
+                f"expected {self._warming_key!r}, got kernel arguments "
+                f"{kernel_args!r}"
+            )
+        launch_kwargs = launch_spec[2] if len(launch_spec) > 2 else None
+        output = launch_spec[3] if len(launch_spec) > 3 else None
+        result = self.launch(self.kernel(*kernel_args), launch_args, launch_kwargs)
+        return output if len(launch_spec) > 3 else result
+
+    return wrapper
diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py
index 7fffa12ec901..198e6fdbe2f4 100644
--- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py
+++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py
@@ -3,25 +3,130 @@
 import ast
 import inspect
 from abc import abstractmethod
-from collections.abc import Callable, Mapping
-from dataclasses import dataclass
-from functools import cached_property, wraps
-from typing import Any, ClassVar, Generic, TypeVar, cast
+from collections.abc import Callable, Hashable, Iterable, Mapping
+from contextlib import nullcontext, suppress
+from dataclasses import dataclass, field
+from functools import cache, cached_property, update_wrapper
+from typing import Any, Generic, ParamSpec, Protocol, TypeVar, cast, overload
 
 from vllm.model_executor.warmup.jit_warmup import (
     VllmJitKernel,
     get_ast_full_name,
     get_function_source_node,
 )
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON, triton
 
 CompileKeyT = TypeVar("CompileKeyT")
-# ``(grid, launch_kwargs)`` or ``(grid, launch_kwargs, outputs)`` for
-# self-allocating kernels. ``grid=None`` skips the launch (e.g. empty batch)
-# but still returns the declared outputs.
+P = ParamSpec("P")
+# ``kernel_launcher`` also supports a ``None`` grid when no launch is needed.
 LaunchSpec = (
     tuple[tuple[int, ...] | None, dict[str, Any]]
     | tuple[tuple[int, ...] | None, dict[str, Any], Any]
 )
+# Dispatchers select a concrete launch configuration for a Triton kernel.
+DispatchSpec = (
+    tuple[tuple[int, ...], dict[str, Any]] | tuple[tuple[int, ...], dict[str, Any], Any]
+)
+
+
+@dataclass(frozen=True)
+class _LaunchBindingPlan:
+    input_targets: tuple[tuple[int, str], ...]
+    stride_targets: tuple[tuple[str, str, int], ...]
+
+
+@cache
+def _launch_binding_plan(
+    arg_names: tuple[str, ...],
+    input_names: tuple[str, ...],
+    launch_names: tuple[str, ...],
+) -> _LaunchBindingPlan:
+    arg_name_set = set(arg_names)
+    launch_name_set = set(launch_names)
+
+    def kernel_arg(input_name: str) -> str | None:
+        pointer_name = f"{input_name}_ptr"
+        candidates = (
+            input_name,
+            pointer_name,
+            input_name.upper(),
+            pointer_name.upper(),
+        )
+        return next(
+            (candidate for candidate in candidates if candidate in arg_name_set),
+            None,
+        )
+
+    input_targets = tuple(
+        (index, target)
+        for index, name in enumerate(input_names)
+        if (target := kernel_arg(name)) is not None and target not in launch_name_set
+    )
+    bound_names = launch_name_set | {target for _, target in input_targets}
+    stride_targets = []
+    for target in arg_names:
+        if target in bound_names:
+            continue
+        name, separator, dim = target.rpartition("_stride")
+        if not separator:
+            name, separator, dim = target.rpartition("_STRIDE")
+        dim = dim.removeprefix("_")
+        if separator and (not dim or dim.isdigit()):
+            pointer = f"{name}_PTR" if separator == "_STRIDE" else f"{name}_ptr"
+            source = name if name in bound_names else pointer
+            if source in bound_names:
+                stride_targets.append((target, source, int(dim or 0)))
+    return _LaunchBindingPlan(input_targets, tuple(stride_targets))
+
+
+def triton_warmup_inputs(
+    kernel: Any,
+    *args: Any,
+    grid: tuple[int, ...],
+    pointer_dtypes: (
+        Mapping[Any, Iterable[str]] | Iterable[tuple[Any, Iterable[str]]] | None
+    ) = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Build launcher inputs from Triton's native positional argument order."""
+    arg_names = tuple(kernel.arg_names)
+    if len(args) > len(arg_names):
+        raise ValueError(
+            f"Received {len(args)} positional inputs for {len(arg_names)} "
+            "Triton kernel arguments"
+        )
+    inputs = dict(zip(arg_names, args))
+    duplicate_names = inputs.keys() & kwargs.keys()
+    if duplicate_names:
+        raise ValueError(
+            f"Triton inputs passed twice: {', '.join(sorted(duplicate_names))}"
+        )
+    inputs.update(kwargs)
+    if pointer_dtypes is not None:
+        function_def = get_function_source_node(kernel)
+        if not isinstance(function_def, ast.FunctionDef):
+            raise ValueError("Expected Triton kernel to be defined as a function")
+        pointer_names = _pointer_arg_names(function_def, arg_names)
+        groups = (
+            pointer_dtypes.items()
+            if isinstance(pointer_dtypes, Mapping)
+            else pointer_dtypes
+        )
+        for dtype, names in groups:
+            for name in names:
+                if name not in arg_names:
+                    raise ValueError(f"Unknown Triton pointer argument: {name}")
+                if name not in pointer_names:
+                    raise ValueError(f"Triton argument is not a pointer: {name}")
+                if name in inputs:
+                    raise ValueError(f"Triton input passed twice: {name}")
+                inputs[name] = TritonWarmupTensor(dtype)
+        missing_pointers = pointer_names - inputs.keys()
+        if missing_pointers:
+            names = ", ".join(sorted(missing_pointers))
+            raise ValueError(f"Missing Triton pointer inputs: {names}")
+    return {"grid": grid, **inputs}
 
 
 def triton_scalar_specialization_rep(value: int) -> int:
@@ -72,6 +177,7 @@ class TritonWarmupTensor:
     aligned: bool = True
     shape: tuple[int, ...] = (1,)
     strides: tuple[int, ...] | None = None
+    init: Any = 0
 
     def data_ptr(self) -> int:
         return 0 if self.aligned else 1
@@ -79,6 +185,12 @@ def data_ptr(self) -> int:
     def ptr_range(self) -> int:
         return 0
 
+    @property
+    def device(self) -> Any:
+        import torch
+
+        return torch.device(current_platform.device_type)
+
     def stride(self, dim: int | None = None) -> int | tuple[int, ...]:
         if self.strides is None:
             strides: list[int] = []
@@ -95,9 +207,10 @@ def stride(self, dim: int | None = None) -> int | tuple[int, ...]:
 class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
     """Triton owner whose runtime launch specification is reused for warmup."""
 
-    kernel: ClassVar[Any]
+    kernel: Any
     _warming = False
     _warming_compile_key: CompileKeyT | None = None
+    _run_autotune = False
 
     @abstractmethod
     def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
@@ -126,23 +239,45 @@ def _kernel_arg_names(self) -> tuple[str, ...]:
             f"Cannot inspect kernel parameters for {type(self.kernel).__name__}"
         )
 
+    def _prepare_launch_kwargs(
+        self,
+        input_names: tuple[str, ...],
+        input_values: tuple[Any, ...],
+        launch_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        plan = _launch_binding_plan(
+            self._kernel_arg_names, input_names, tuple(launch_kwargs)
+        )
+        for index, target in plan.input_targets:
+            launch_kwargs[target] = input_values[index]
+        for target, source, dim in plan.stride_targets:
+            launch_kwargs[target] = launch_kwargs[source].stride(dim)
+        return launch_kwargs
+
     def launch(
         self,
-        grid: tuple[int, ...] | None,
+        launch_spec: LaunchSpec,
         inputs: Mapping[str, Any],
-        /,
-        *,
-        kernel: Any = None,
-        **kwargs: Any,
     ) -> Any:
-        kernel = kernel if kernel is not None else self.kernel
+        if len(launch_spec) == 3:
+            grid, launch_kwargs, outputs = launch_spec
+        else:
+            grid, launch_kwargs = launch_spec
+            outputs = None
+        kwargs = self._prepare_launch_kwargs(
+            tuple(inputs), tuple(inputs.values()), dict(launch_kwargs)
+        )
+        kernel = kwargs.pop("kernel", None) or self.kernel
         runtime_launcher = kwargs.pop("_runtime_launcher", None)
         runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0)
-        for name, value in inputs.items():
-            target = name if name in self._kernel_arg_names else f"{name}_ptr"
-            if target in self._kernel_arg_names and target not in kwargs:
-                kwargs[target] = value
         if self._warming:
+            if self._run_autotune:
+                assert grid is not None
+                kernel[grid](**kwargs)
+                return outputs
+            kwargs = {
+                name: _triton_metadata_arg(value) for name, value in kwargs.items()
+            }
             if (
                 self._warming_compile_key is not None
                 and "launch_pdl" in kwargs
@@ -151,51 +286,382 @@ def launch(
                 kwargs["launch_pdl"] = self._warming_compile_key.launch_pdl
             warmup = getattr(kernel, "warmup", None)
             assert warmup is not None
-            return warmup(grid=(1,), **kwargs)
+            warmup(grid=(1,), **kwargs)
+            return outputs
         if grid is None:
-            return None
+            return outputs
         if runtime_launcher is not None:
             regular_args = [
                 kwargs.pop(name)
                 for name in self._kernel_arg_names[:runtime_launcher_arg_count]
             ]
-            return runtime_launcher(kernel, grid, *regular_args, **kwargs)
-        return kernel[grid](**kwargs)
+            runtime_launcher(kernel, grid, *regular_args, **kwargs)
+        else:
+            kernel[grid](**kwargs)
+        return outputs
 
 
-def kernel_launcher(
-    call_fn: Callable[..., LaunchSpec],
-) -> Callable[..., Any]:
-    """Launch a Triton kernel from a declarative ``__call__`` specification.
+@dataclass(frozen=True)
+class TritonJitKey:
+    """Process-local identity of one Triton JIT specialization."""
+
+    jit_function_id: int
+    jit_function_key: Hashable
+    device: Hashable
+    cache_key: Hashable
 
-    ``call_fn`` returns either ``(grid, launch_kwargs)`` or, when it allocates
-    its own outputs, ``(grid, launch_kwargs, outputs)``. The declared outputs are
-    returned to the caller. A ``grid`` of ``None`` skips the launch (e.g. an
-    empty-token batch) while still returning the outputs.
-    """
-    signature = inspect.signature(call_fn)
 
-    @wraps(call_fn)
-    def wrapper(
-        self: VllmTritonJitKernel[Any],
+@dataclass(frozen=True)
+class TritonCompileKey:
+    """Triton-derived compile key with one non-comparing replay input."""
+
+    jit_keys: frozenset[TritonJitKey]
+    inputs: tuple[tuple[str, Any], ...] = field(compare=False, hash=False, repr=False)
+
+
+def _triton_key_deriver(
+    kernel: Any,
+) -> Callable[[Mapping[str, Any]], set[TritonJitKey]]:
+    """Prepare Triton's key derivation once for one kernel and device."""
+    knobs = triton.knobs
+    Autotuner = triton.runtime.autotuner.Autotuner
+    Heuristics = triton.runtime.autotuner.Heuristics
+    driver = triton.runtime.driver
+    JITFunction = triton.runtime.jit.JITFunction
+    compute_cache_key = triton.runtime.jit.compute_cache_key
+
+    if isinstance(kernel, Heuristics):
+        derive_inner = _triton_key_deriver(kernel.fn)
+
+        def derive_heuristic(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:
+            heuristic_kwargs = dict(kwargs)
+            for name, heuristic in kernel.values.items():
+                heuristic_kwargs[name] = heuristic(heuristic_kwargs)
+            return derive_inner(heuristic_kwargs)
+
+        return derive_heuristic
+
+    if isinstance(kernel, Autotuner):
+        derive_inner = _triton_key_deriver(kernel.fn)
+
+        def derive_autotuned(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:
+            previous_nargs = getattr(kernel, "nargs", None)
+            kernel.nargs = {}
+            try:
+                configs = kernel.prune_configs(dict(kwargs))
+            finally:
+                kernel.nargs = previous_nargs
+            keys: set[TritonJitKey] = set()
+            for config in configs:
+                conflicts = kwargs.keys() & config.kwargs.keys()
+                if conflicts:
+                    names = ", ".join(sorted(conflicts))
+                    raise ValueError(f"Conflicting autotune parameters: {names}")
+                keys.update(derive_inner(dict(kwargs) | config.all_kwargs()))
+            return keys
+
+        return derive_autotuned
+
+    if not isinstance(kernel, JITFunction):
+        raise TypeError(f"Unsupported Triton kernel wrapper: {type(kernel).__name__}")
+
+    device = cast(Hashable, driver.active.get_current_device())
+    _, kernel_key_cache, _, _, binder = kernel.device_caches[device]
+    jit_function_id = id(kernel)
+    jit_function_key = cast(Hashable, kernel.cache_key)
+    debug = kernel.debug or knobs.runtime.debug
+    instrumentation_mode = knobs.compilation.instrumentation_mode
+    fpsan_homomorphic_casts = getattr(
+        knobs.compilation, "fpsan_homomorphic_casts", None
+    )
+    inspection_hook = knobs.runtime.add_stages_inspection_hook
+
+    def derive_jit(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:
+        binder_kwargs = dict(kwargs)
+        binder_kwargs["debug"] = binder_kwargs.get("debug", debug) or debug
+        binder_kwargs["instrumentation_mode"] = instrumentation_mode
+        if fpsan_homomorphic_casts is not None:
+            binder_kwargs["fpsan_homomorphic_casts"] = fpsan_homomorphic_casts
+        _, specialization, options = binder(**binder_kwargs)
+        if inspection_hook is not None:
+            _, inspection_hash = inspection_hook()
+            specialization.append(f'("custom_pipeline", {inspection_hash})')
+        cache_key = cast(
+            Hashable,
+            compute_cache_key(kernel_key_cache, specialization, options),
+        )
+        return {
+            TritonJitKey(
+                jit_function_id,
+                jit_function_key,
+                device,
+                cache_key,
+            )
+        }
+
+    return derive_jit
+
+
+WarmupCases = Mapping[str, Any] | Iterable[Mapping[str, Any]]
+
+
+class _AutomaticTritonJitKernel(VllmTritonJitKernel[TritonCompileKey]):
+    def warmup_inputs(self, compile_key: TritonCompileKey) -> dict[str, Any]:
+        return dict(compile_key.inputs)
+
+    def _provider_cases(
+        self,
+        provider: Callable[..., WarmupCases],
         *args: Any,
         **kwargs: Any,
-    ) -> Any:
-        spec = call_fn(self, *args, **kwargs)
-        if len(spec) == 3:
-            grid, launch_kwargs, outputs = spec
+    ) -> Iterable[Mapping[str, Any]]:
+        cases_node = get_function_source_node(provider)
+        if any(
+            isinstance(node, ast.Call)
+            and (get_ast_full_name(node.func) or "").split(".")[-1]
+            in {"WarmupIntRange", "WarmupChoices", "_when"}
+            for node in ast.walk(cases_node)
+        ):
+            return self._expand_warmup_cases(provider, *args, **kwargs)
+        cases = provider(*args, **kwargs)
+        return (cases,) if isinstance(cases, Mapping) else cases
+
+
+def _is_autotuned(kernel: Any) -> bool:
+    if not HAS_TRITON:
+        return False
+    Autotuner = triton.runtime.autotuner.Autotuner
+
+    while kernel is not None:
+        if isinstance(kernel, Autotuner):
+            return True
+        kernel = getattr(kernel, "fn", None)
+    return False
+
+
+def _triton_metadata_arg(value: Any) -> Any:
+    from torch._subclasses.fake_tensor import FakeTensor
+
+    if not isinstance(value, FakeTensor):
+        return value
+    return TritonWarmupTensor(
+        value.dtype,
+        aligned=getattr(value, "_vllm_warmup_aligned", True),
+        shape=tuple(value.shape),
+        strides=tuple(value.stride()),
+    )
+
+
+def _materialize_warmup_case(
+    case: Mapping[str, Any],
+    *,
+    real: bool,
+    mode: Any,
+    cache: dict[TritonWarmupTensor, Any],
+) -> dict[str, Any]:
+    import torch
+
+    def materialize(value: Any) -> Any:
+        if not isinstance(value, TritonWarmupTensor):
+            return value
+        if not real:
+            try:
+                return cache[value]
+            except KeyError:
+                pass
+            except TypeError:
+                # Unhashable initializers are uncommon and safe to materialize directly.
+                pass
+
+        strides = cast(tuple[int, ...], value.stride())
+        with mode:
+            tensor = torch.empty_strided(
+                value.shape,
+                strides,
+                dtype=value.dtype,
+                device=value.device,
+            )
+        if real:
+            if callable(value.init):
+                value.init(tensor)
+            else:
+                tensor.fill_(value.init)
         else:
-            grid, launch_kwargs = spec
-            outputs = None
-        bound = signature.bind(self, *args, **kwargs)
-        bound.apply_defaults()
-        inputs = {
-            name: value for name, value in bound.arguments.items() if name != "self"
-        }
-        self.launch(grid, inputs, **launch_kwargs)
-        return outputs
+            tensor._vllm_warmup_aligned = value.aligned
+            with suppress(TypeError):
+                cache[value] = tensor
+        return tensor
+
+    return {name: materialize(value) for name, value in case.items()}
+
+
+class _DecoratedTritonJitKernel(_AutomaticTritonJitKernel):
+    def __init__(
+        self,
+        kernel: Any,
+        warmup_inputs: Callable[..., WarmupCases],
+        dispatch: Callable[..., DispatchSpec] | None,
+    ) -> None:
+        self.kernel = kernel
+        self._run_autotune = _is_autotuned(kernel)
+        self._warmup_inputs_fn = warmup_inputs
+        self._dispatch_fn = dispatch
+        self._dispatch_arg_names: tuple[str, ...] = ()
+        self._dispatch_defaults: dict[str, Any] = {}
+        if dispatch is not None:
+            parameters = tuple(inspect.signature(dispatch).parameters.values())
+            if any(
+                parameter.kind
+                in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
+                for parameter in parameters
+            ):
+                raise TypeError("Triton dispatchers require explicit parameters")
+            self._dispatch_arg_names = tuple(parameter.name for parameter in parameters)
+            self._dispatch_defaults = {
+                parameter.name: parameter.default
+                for parameter in parameters
+                if parameter.default is not inspect.Parameter.empty
+            }
+        super().__init__()
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if self._dispatch_fn is None:
+            grid = kwargs.pop("grid") if not args else args[0]
+            kernel_args = () if not args else args[1:]
+            inputs = triton_warmup_inputs(
+                self.kernel, *kernel_args, grid=grid, **kwargs
+            )
+            inputs.pop("grid")
+            return self.launch((grid, inputs), {})
+        spec = self._dispatch_fn(*args, **kwargs)
+        input_values = args
+        if len(args) != len(self._dispatch_arg_names):
+            input_values += tuple(
+                kwargs[name] if name in kwargs else self._dispatch_defaults[name]
+                for name in self._dispatch_arg_names[len(args) :]
+            )
+        return self.launch(
+            spec,
+            dict(zip(self._dispatch_arg_names, input_values)),
+        )
+
+    def __getitem__(self, grid: tuple[int, ...]) -> Callable[..., Any]:
+        return cast(Callable[..., Any], self.kernel[grid])
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.kernel, name)
+
+    def get_warmup_keys(
+        self, *, vllm_config: Any | None = None, **kwargs: Any
+    ) -> list[TritonCompileKey]:
+        keys: list[TritonCompileKey] = []
+        seen_jit_keys: set[frozenset[TritonJitKey]] = set()
+        if "vllm_config" in inspect.signature(self._warmup_inputs_fn).parameters:
+            kwargs.setdefault("vllm_config", vllm_config)
+        cases = self._provider_cases(self._warmup_inputs_fn, **kwargs)
+        derive_keys = _triton_key_deriver(self.kernel)
+        materialize = self._run_autotune
+        mode: Any = nullcontext() if self._run_autotune else None
+        materialized: dict[TritonWarmupTensor, Any] = {}
+        for case in cases:
+            inputs = (
+                _materialize_warmup_case(
+                    case,
+                    real=self._run_autotune,
+                    mode=mode,
+                    cache=materialized,
+                )
+                if materialize
+                else case
+            )
+            if self._dispatch_fn is None:
+                prepared = dict(inputs)
+                prepared.pop("grid")
+            else:
+                dispatch_inputs = {
+                    name: inputs[name]
+                    for name in self._dispatch_arg_names
+                    if name in inputs
+                }
+                try:
+                    _, launch_kwargs = self._dispatch_fn(**dispatch_inputs)[:2]
+                except (AssertionError, AttributeError, TypeError):
+                    if materialize:
+                        raise
+                    from torch._subclasses.fake_tensor import FakeTensorMode
+
+                    materialize = True
+                    mode = FakeTensorMode()
+                    inputs = _materialize_warmup_case(
+                        case,
+                        real=False,
+                        mode=mode,
+                        cache=materialized,
+                    )
+                    dispatch_inputs = {
+                        name: inputs[name]
+                        for name in self._dispatch_arg_names
+                        if name in inputs
+                    }
+                    _, launch_kwargs = self._dispatch_fn(**dispatch_inputs)[:2]
+                prepared = self._prepare_launch_kwargs(
+                    tuple(inputs),
+                    tuple(inputs.values()),
+                    dict(launch_kwargs),
+                )
+            if materialize:
+                prepared = {
+                    name: _triton_metadata_arg(value)
+                    for name, value in prepared.items()
+                }
+            jit_keys = frozenset(derive_keys(prepared))
+            if jit_keys and jit_keys not in seen_jit_keys:
+                seen_jit_keys.add(jit_keys)
+                keys.append(TritonCompileKey(jit_keys, tuple(inputs.items())))
+        return keys
+
+
+class TritonKernelDispatcher(Protocol[P]):
+    """Callable Triton launcher with compile-time warmup support."""
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Any: ...
+
+    def register_warmup(self, *args: Any, **kwargs: Any) -> None: ...
+
+
+@overload
+def triton_kernel_dispatcher_with_warmup(
+    *,
+    warmup_inputs: Callable[..., WarmupCases],
+    kernel: None = None,
+) -> Callable[[Any], Any]: ...
+
+
+@overload
+def triton_kernel_dispatcher_with_warmup(
+    *,
+    warmup_inputs: Callable[..., WarmupCases],
+    kernel: Any,
+) -> Callable[[Callable[P, DispatchSpec]], TritonKernelDispatcher[P]]: ...
+
+
+def triton_kernel_dispatcher_with_warmup(
+    *,
+    warmup_inputs: Callable[..., WarmupCases],
+    kernel: Any | None = None,
+) -> Callable[[Any], Any]:
+    """Decorate a dispatch function or a native Triton kernel for warmup."""
+
+    def decorate(dispatch: Any) -> Any:
+        owner = _DecoratedTritonJitKernel(
+            dispatch if kernel is None else kernel,
+            warmup_inputs,
+            None if kernel is None else dispatch,
+        )
+        update_wrapper(owner, dispatch, updated=())
+        return owner
 
-    return wrapper
+    return decorate
 
 
 @dataclass(frozen=True)
diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py
index c9c740f3684c..33fe7002aa2f 100644
--- a/vllm/model_executor/warmup/kernel_warmup.py
+++ b/vllm/model_executor/warmup/kernel_warmup.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Warmup kernels used during model execution.
+"""Warmup kernels used during model execution.
 This is useful specifically for JIT'ed kernels as we don't want JIT'ing to
 happen during model execution.
 """
@@ -17,14 +16,12 @@
 from vllm.model_executor.warmup.b12x_warmup import b12x_warmup
 from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
 from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
-from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
-    deepseek_v4_mhc_warmup,
-)
 from vllm.model_executor.warmup.flashinfer_autotune_cache import (
     resolve_flashinfer_autotune_file,
     write_flashinfer_autotune_cache,
 )
 from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
+    autotune_hisparse_flashinfer_attention,
     deepseek_v4_sparse_mla_attention_warmup,
     flashinfer_sparse_mla_decode_autotune_warmup,
 )
@@ -40,8 +37,8 @@
 from vllm.model_executor.warmup.replayssm_warmup import (
     replayssm_autotune_warmup,
 )
-from vllm.model_executor.warmup.spec_decode_rejection_warmup import (
-    spec_decode_rejection_warmup,
+from vllm.model_executor.warmup.watermark_sample_warmup import (
+    watermark_sample_warmup,
 )
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import is_deep_gemm_supported
@@ -139,16 +136,17 @@ def _warmup_bf16x3_router_gemm(
     logger.info_once("Warmed up BF16x3 router GEMM configs: %s.", configs)
 
 
-def _warmup_kimi_k3_gemm_rs_ar() -> None:
-    # Kimi-K3 model construction imports this module only when GEMM-RS/AR is
-    # enabled and initializes its singleton before kernel_warmup runs. Avoid
-    # importing it here so other models do not compile the RS/AR variants.
-    module = sys.modules.get("vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar")
+def _warmup_gemm_rs_ar() -> None:
+    # Model construction (Kimi-K3, DeepSeek-V4.1) imports this module only
+    # when GEMM-RS/AR is enabled and initializes its singleton before
+    # kernel_warmup runs. Avoid importing it here so other models do not
+    # compile the RS/AR variants.
+    module = sys.modules.get("vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar")
     if module is None:
         return
     compiled = module.warmup_gemm_rs_ar()
     if compiled:
-        logger.info_once("Warmed up %d Kimi-K3 GEMM-RS/AR variants.", compiled)
+        logger.info_once("Warmed up %d GEMM-RS/AR variants.", compiled)
 
 
 def _autotune_kimi_k3_kda_qkvg(model: torch.nn.Module) -> None:
@@ -174,7 +172,10 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
         logger.info("JIT kernel warmup starting.")
         jit_warmup_start = time.perf_counter()
         try:
-            worker.model_runner.jit_warmup_registry.warmup()
+            registry = (
+                worker.model_runner.jit_warmup_registry  # type: ignore[attr-defined]
+            )
+            registry.warmup()
         except Exception:
             logger.exception(
                 "JIT kernel warmup failed after %.2fs.",
@@ -193,19 +194,10 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
     compilation_config = worker.vllm_config.compilation_config
     cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or [])
 
-    # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
-    # layer per token; warm them across token sizes first so the first real
-    # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside).
-    deepseek_v4_mhc_warmup(
-        worker.get_model(),
-        max_tokens=worker.scheduler_config.max_num_batched_tokens,
-        cudagraph_capture_sizes=cudagraph_capture_sizes,
-    )
-
     # Run next so input-prep kernels JIT against pristine runner state.
     if enable_jit_warmup:
         kimi_k3_triton_warmup(worker)
-        spec_decode_rejection_warmup(worker)
+        watermark_sample_warmup(worker)
         qwen4_exp_qsa_triton_warmup(worker)
 
     if enable_jit_warmup and current_platform.is_device_capability_family(100):
@@ -217,7 +209,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
     if current_platform.has_device_capability(90):
         _warmup_ll_bf16_router_gemm(worker.get_model())
 
-    _warmup_kimi_k3_gemm_rs_ar()
+    _warmup_gemm_rs_ar()
 
     if worker.vllm_config.kernel_config.enable_cutedsl_warmup:
         # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
@@ -339,13 +331,16 @@ def _flashinfer_deferred_moe_token_counts(
     return tuple(dict.fromkeys(token_counts))
 
 
-def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]:
+def _flashinfer_autotune_token_counts(
+    runner: "GPUModelRunner", *, include_bf16: bool = True
+) -> tuple[int, ...]:
     max_tokens = runner.scheduler_config.max_num_batched_tokens
     # Tune the widest bucket set first so bounded passes reuse its configs.
     token_counts = [max_tokens]
     linear_backend = runner.vllm_config.kernel_config.linear_backend
     if (
-        linear_backend == "flashinfer_cutedsl"
+        include_bf16
+        and linear_backend == "flashinfer_cutedsl"
         and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
     ):
         token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS)
@@ -353,10 +348,13 @@ def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ..
     return tuple(dict.fromkeys(token_counts))
 
 
-def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
+def _run_flashinfer_autotune_dummy_runs(
+    runner: "GPUModelRunner", *, skip_attn: bool = False
+) -> None:
     import vllm.utils.flashinfer as fi_utils
 
-    for num_tokens in _flashinfer_autotune_token_counts(runner):
+    dummy_run_kwargs = {"skip_attn": True} if skip_attn else {}
+    for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False):
         tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
         logger.info(
             "Running FlashInfer autotune with %d tokens and token buckets %s.",
@@ -369,12 +367,42 @@ def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
                 skip_eplb=True,
                 is_profile=True,
                 randomize_inputs=True,
+                **dummy_run_kwargs,
             )
 
 
+def _run_flashinfer_bf16_autotune_dummy_run(
+    runner: "GPUModelRunner",
+    *,
+    skip_ops: set[str] | None = None,
+    skip_attn: bool = False,
+) -> None:
+    import vllm.utils.flashinfer as fi_utils
+
+    if (
+        runner.vllm_config.kernel_config.linear_backend != "flashinfer_cutedsl"
+        or runner.scheduler_config.max_num_batched_tokens
+        <= _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
+    ):
+        return
+
+    num_tokens = _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
+    tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens)
+    logger.info("Running FlashInfer BF16-only autotune with %d tokens.", num_tokens)
+    # Other ops execute from their existing cache or fallback, without the
+    # BF16 bucket cap. No full-model autotune context may enclose this pass.
+    with fi_utils.autotune_bf16_only(tuning_buckets, skip_ops=skip_ops):
+        runner._dummy_run(
+            num_tokens=num_tokens,
+            skip_eplb=True,
+            is_profile=True,
+            randomize_inputs=True,
+            **({"skip_attn": True} if skip_attn else {}),
+        )
+
+
 def flashinfer_autotune(runner: "GPUModelRunner") -> None:
-    """
-    Autotune FlashInfer operations.
+    """Autotune FlashInfer operations.
     FlashInfer have many implementations for the same operation,
     autotuning runs benchmarks for each implementation and stores
     the results. The results are cached transparently and
@@ -431,9 +459,20 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
             torch.inference_mode(),
             fi_utils.autotune(tune_mode=True, **autotune_kwargs),
         ):
-            _run_flashinfer_autotune_dummy_runs(runner)
+            hisparse_enabled = (
+                runner.vllm_config.attention_config.hisparse_config is not None
+            )
+            if hisparse_enabled:
+                # HiSparse hot-buffer attention is bounded by decode batch
+                # size, not the prefill-sized batch used for the full model.
+                autotune_hisparse_flashinfer_attention(runner)
+            _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled)
             replayssm_autotune_warmup(runner)
             _autotune_kimi_k3_kda_qkvg(runner.get_model())
+        with torch.inference_mode():
+            _run_flashinfer_bf16_autotune_dummy_run(
+                runner, skip_ops=skip_ops, skip_attn=hisparse_enabled
+            )
     finally:
         set_autotune_process_group(None)
 
diff --git a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py
index 3cbc3bf2e086..2065b3833822 100644
--- a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py
+++ b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py
@@ -19,7 +19,6 @@
 def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None:
     """Warm every reachable QSA specialization: indexer decode-query-length
     profiles plus the sparse attention split-K/merge configs."""
-
     qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.indexer_qsa")
     attn_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.qsa")
     if qsa_module is None or attn_module is None:
diff --git a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py b/vllm/model_executor/warmup/spec_decode_rejection_warmup.py
deleted file mode 100644
index 8cedbfc4db42..000000000000
--- a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py
+++ /dev/null
@@ -1,111 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Warm up spec-decode rejection-sampler Triton kernels.
-
-The rejection sampler kernels (``_compute_local_logits_stats_kernel``,
-``_rejection_kernel``, ``_resample_kernel``) are JIT-compiled by Triton on
-first use. Without warmup, the first spec-decode request pays a multi-second
-compilation cost. This pre-compiles them with dummy data matching the
-server's vocab size and speculative config.
-"""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-import torch
-
-from vllm.logger import init_logger
-
-if TYPE_CHECKING:
-    from vllm.v1.worker.gpu_worker import Worker
-
-logger = init_logger(__name__)
-
-
-@torch.inference_mode()
-def spec_decode_rejection_warmup(worker: Worker) -> None:
-    spec_config = worker.vllm_config.speculative_config
-    if spec_config is None:
-        return
-
-    from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
-        rejection_sample,
-    )
-
-    model_config = worker.vllm_config.model_config
-    vocab_size = model_config.get_vocab_size()
-    num_spec = spec_config.num_speculative_tokens
-    if num_spec <= 0 or vocab_size <= 0:
-        return
-
-    # Mirror the constexpr-relevant flags the runtime uses.
-    rejection_method = getattr(spec_config, "rejection_sample_method", None)
-    use_block_verification = rejection_method == "block"
-    use_synthetic = rejection_method == "synthetic"
-
-    device = worker.device
-    num_reqs = 1
-    tokens_per_req = num_spec + 1
-    num_logits = num_reqs * tokens_per_req
-
-    # Triton JIT-specializes on tensor dtypes. The target logits may be fp32
-    # (apply_sampling_params copies to fp32 when processing is needed) or the
-    # model dtype (pass-through otherwise), while draft logits are always the
-    # model dtype. Warm every (target, draft) combination the runtime can hit.
-    model_dtype = model_config.dtype
-    warmup_dtype_pairs = {
-        (model_dtype, model_dtype),
-        (torch.float32, torch.float32),
-        (torch.float32, model_dtype),
-        (model_dtype, torch.float32),
-    }
-
-    logger.info(
-        "Warming up spec-decode rejection sampler kernels "
-        "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s).",
-        vocab_size,
-        num_spec,
-        [(str(t), str(d)) for t, d in warmup_dtype_pairs],
-        use_block_verification,
-    )
-    for tgt_dtype, draft_dtype in warmup_dtype_pairs:
-        target_logits = torch.zeros(
-            (num_logits, vocab_size), dtype=tgt_dtype, device=device
-        )
-        draft_logits = torch.zeros(
-            (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device
-        )
-        synthetic_rates = (
-            torch.full((num_spec,), 0.5, dtype=torch.float32, device=device)
-            if use_synthetic
-            else None
-        )
-        try:
-            rejection_sample(
-                target_logits=target_logits,
-                draft_logits=draft_logits,
-                draft_sampled=torch.zeros(num_logits, dtype=torch.int64, device=device),
-                cu_num_logits=torch.tensor(
-                    [0, num_logits], dtype=torch.int32, device=device
-                ),
-                pos=torch.zeros(num_logits, dtype=torch.int64, device=device),
-                idx_mapping=torch.zeros(num_reqs, dtype=torch.int32, device=device),
-                expanded_idx_mapping=torch.zeros(
-                    num_logits, dtype=torch.int32, device=device
-                ),
-                expanded_local_pos=torch.arange(
-                    num_logits, dtype=torch.int32, device=device
-                ),
-                temperature=torch.zeros(num_reqs, dtype=torch.float32, device=device),
-                seed=torch.full((num_reqs,), 42, dtype=torch.int64, device=device),
-                num_speculative_steps=num_spec,
-                synthetic_conditional_rates=synthetic_rates,
-                use_fp64=False,
-                use_block_verification=use_block_verification,
-            )
-        except Exception:
-            logger.warning(
-                "Skipping spec-decode rejection sampler warmup.", exc_info=True
-            )
-            return
diff --git a/vllm/model_executor/warmup/watermark_sample_warmup.py b/vllm/model_executor/warmup/watermark_sample_warmup.py
new file mode 100644
index 000000000000..d315f10310b7
--- /dev/null
+++ b/vllm/model_executor/warmup/watermark_sample_warmup.py
@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Warm up the watermarked sampler's Triton kernel.
+
+``_philox_gumbel_kernel`` is JIT-compiled once per (key, logits dtype,
+skip-mask, ``USE_FP64``) specialization. The generic sampler warmup samples with
+``SamplingParams.for_sampler_warmup()``, whose feature-heavy logits processing
+forces the fp32 copy in ``apply_sampling_params``, so the first ordinary
+temperature-1.0 request is the first launch with model-dtype logits and pays
+the compilation inside inference. This pre-compiles every specialization the
+sampler path can launch for the configured watermark.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+if TYPE_CHECKING:
+    from vllm.v1.watermarking.watermarker import Watermarker
+    from vllm.v1.worker.gpu_worker import Worker
+
+logger = init_logger(__name__)
+
+_NUM_WARMUP_TOKENS = 1
+
+
+def _philox_key(watermarker: Watermarker) -> int | None:
+    """None when the watermarker takes the torch path and compiles no kernel."""
+    from vllm.v1.watermarking.gumbel import GumbelWatermarker
+    from vllm.v1.watermarking.prfs import PhiloxPRF
+
+    if not isinstance(watermarker, GumbelWatermarker):
+        return None
+    prf = watermarker.prf
+    return prf.key if type(prf) is PhiloxPRF else None
+
+
+def _philox_sampler_keys(watermarker: Watermarker) -> list[int]:
+    """Philox keys the sampler launches the kernel with."""
+    from vllm.v1.watermarking.gumbel import DualKeyGumbelWatermarker
+
+    watermarkers: list[Watermarker] = [watermarker]
+    if isinstance(watermarker, DualKeyGumbelWatermarker):
+        # ``DualKeyGumbelWatermarker.sample`` launches once per key, and
+        # collapses to a single key at the alpha bounds, where the other key's
+        # specialization is never launched.
+        if watermarker.alpha == 0:
+            watermarkers = [watermarker.draft_watermarker]
+        elif watermarker.alpha == 1:
+            watermarkers = [watermarker.target_watermarker]
+        else:
+            watermarkers = [
+                watermarker.draft_watermarker,
+                watermarker.target_watermarker,
+            ]
+    return [key for key in map(_philox_key, watermarkers) if key is not None]
+
+
+@torch.inference_mode()
+def watermark_sample_warmup(worker: Worker) -> None:
+    if getattr(worker.vllm_config, "watermark_config", None) is None:
+        return
+    # ``GumbelWatermarker`` launches the kernel only on CUDA
+    # (vllm/v1/watermarking/gumbel.py:63,:82); elsewhere the sampler takes the
+    # torch path and warming here would compile a kernel nothing launches.
+    if not current_platform.is_cuda_alike():
+        return
+
+    from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
+    from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
+
+    # Read the sampler the runner built instead of rebuilding it from config:
+    # its watermarker (the target role under a speculative config), its context
+    # deduplication setting and its fp64 flag are what the launches have to
+    # match, and cannot drift from the runner. A rank that does not sample -- a
+    # non-final pipeline stage, a pooling model -- holds no such sampler.
+    sampler = getattr(worker.model_runner, "sampler", None)
+    if not isinstance(sampler, GPUWatermarkSampler):
+        return
+
+    model_config = worker.vllm_config.model_config
+    device = worker.device
+    try:
+        keys = _philox_sampler_keys(sampler.watermarker)
+        if not keys:
+            return
+
+        # The sampler passes model-dtype logits through when no logits
+        # processing is needed and fp32 otherwise. ``ModelConfig.dtype`` is
+        # resolved to a torch dtype during validation.
+        dtypes = {cast("torch.dtype", model_config.dtype), torch.float32}
+        # ``GPUWatermarkSampler._sample_random`` builds a skip mask whenever
+        # context deduplication is on; with it off only a batch that mixes in an
+        # opted-out or greedy row does, and the mask-free specialization stays
+        # reachable.
+        with_skip_mask = (
+            (True,) if sampler.deduplicate_contexts != "none" else (True, False)
+        )
+
+        vocab_size = model_config.get_vocab_size()
+        contexts = torch.zeros(
+            (_NUM_WARMUP_TOKENS, sampler.watermarker.context_width),
+            dtype=torch.int32,
+            device=device,
+        )
+        # Dtypes mirror the runtime buffers the sampler passes and are part of
+        # the specialization key: int32 request-state token ids, int64 index and
+        # seed buffers, fp32 sampling temperatures.
+        sampling_state = {
+            "skip_mask": torch.zeros(
+                _NUM_WARMUP_TOKENS, dtype=torch.bool, device=device
+            ),
+            "expanded_idx_mapping": torch.zeros(
+                _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device
+            ),
+            "temperatures": torch.ones(
+                _NUM_WARMUP_TOKENS, dtype=torch.float32, device=device
+            ),
+            "seeds": torch.zeros(_NUM_WARMUP_TOKENS, dtype=torch.int64, device=device),
+            "positions": torch.zeros(
+                _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device
+            ),
+        }
+
+        logger.info(
+            "Warming up watermark sampler kernel (vocab=%d, keys=%d, dtypes=%s, "
+            "skip_mask=%s).",
+            vocab_size,
+            len(keys),
+            [str(dtype) for dtype in dtypes],
+            list(with_skip_mask),
+        )
+        for dtype in dtypes:
+            logits = torch.zeros(
+                (_NUM_WARMUP_TOKENS, vocab_size), dtype=dtype, device=device
+            )
+            for key in keys:
+                for use_skip_mask in with_skip_mask:
+                    philox_gumbel_sample(
+                        logits,
+                        contexts,
+                        key,
+                        # Only the mixed path forwards the sampler's setting;
+                        # ``GumbelWatermarker._sample_watermarked`` launches
+                        # the mask-free variant with the fp32 default.
+                        use_fp64=use_skip_mask and sampler.use_fp64_gumbel,
+                        **(sampling_state if use_skip_mask else {}),
+                    )
+    except Exception:
+        logger.warning("Skipping watermark sampler warmup.", exc_info=True)
diff --git a/vllm/models/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py
index 2ee6cc2854d2..e64e81821062 100644
--- a/vllm/models/common/ops/fused_qk_rmsnorm.py
+++ b/vllm/models/common/ops/fused_qk_rmsnorm.py
@@ -5,11 +5,11 @@
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
diff --git a/vllm/models/deepseek_v32/__init__.py b/vllm/models/deepseek_v32/__init__.py
index 0ac154505f2e..945623980902 100644
--- a/vllm/models/deepseek_v32/__init__.py
+++ b/vllm/models/deepseek_v32/__init__.py
@@ -7,9 +7,9 @@
 attend. The same model code serves any DSA checkpoint, including GLM-5.2
 (``glm_moe_dsa``), which reuses this architecture.
 
-The CUDA implementation selects capability-specific kernels internally and
-falls back when an optimization is unavailable. Other platforms use the
-generic implementation by default.
+The CUDA and ROCm implementations select capability-specific kernels
+internally and fall back when an optimization is unavailable. Other platforms
+use the generic implementation by default.
 """
 
 from vllm.platforms import current_platform
@@ -20,8 +20,12 @@
     from .nvidia.model import DeepseekV32ForCausalLM
     from .nvidia.model import DeepseekV32ForCausalLM as GlmMoeDsaForCausalLM
     from .nvidia.mtp import DeepseekV32MTP
+elif current_platform.is_rocm():
+    from .amd.model import DeepseekV32ForCausalLM
+    from .amd.model import DeepseekV32ForCausalLM as GlmMoeDsaForCausalLM
+    from .amd.mtp import DeepseekV32MTP
 else:
-    # ROCm, XPU, and CPU keep the generic implementation.
+    # XPU and CPU keep the generic implementation.
     from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepseekV32MTP
     from vllm.model_executor.models.deepseek_v2 import (
         DeepseekV3ForCausalLM as DeepseekV32ForCausalLM,
diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py
index 9ea86ee26b0f..46c35dec98fa 100644
--- a/vllm/models/deepseek_v32/amd/mtp.py
+++ b/vllm/models/deepseek_v32/amd/mtp.py
@@ -32,8 +32,10 @@
     _try_load_fp8_indexer_wk,
     get_spec_layer_idx_from_weight_name,
 )
+from vllm.model_executor.models.interfaces import SupportsPP
 from vllm.model_executor.models.utils import (
     get_pp_missing_layer_names,
+    make_empty_intermediate_tensors_factory,
     maybe_prefix,
 )
 from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
@@ -160,7 +162,7 @@ def compute_logits(
         return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
 
 
-class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
+class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
@@ -169,6 +171,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
         self.set_moe_parameters()
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], self.config.hidden_size
+        )
         self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
             self.model.layers.values(),
             DeepseekV2MoE,
@@ -192,7 +197,7 @@ def set_moe_parameters(self):
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
 
-    def forward(
+    def forward(  # type: ignore[override]
         self,
         input_ids: torch.Tensor | None,
         positions: torch.Tensor,
@@ -268,6 +273,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 continue
             spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
             if spec_layer is None:
+                # A tied top-level embed_tokens has no spec layer to rewrite
+                # from; the draft needs its own copy under PP.
+                param = params_dict.get(name) if "embed_tokens" in name else None
+                if param is not None:
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add(name)
                 continue
             is_fusion_moe_shared_experts_layer = (
                 self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
diff --git a/vllm/models/deepseek_v32/amd/rocm.py b/vllm/models/deepseek_v32/amd/rocm.py
index 6d8e7ea855b8..b9ac31f8be2b 100644
--- a/vllm/models/deepseek_v32/amd/rocm.py
+++ b/vllm/models/deepseek_v32/amd/rocm.py
@@ -3,12 +3,14 @@
 
 import torch
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
 from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
 from vllm.models.deepseek_v32.attention import DeepseekV32Attention, DeepseekV32Indexer
 from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q
+from vllm.platforms import current_platform
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
 from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
@@ -16,21 +18,15 @@
 )
 
 
-class DeepseekV32MLASparseBackend(ROCMAiterMLASparseBackend):
-    @staticmethod
-    def get_supported_kernel_block_sizes() -> list:
-        return [16, 32]
-
-
-class DeepseekV32ROCmIndexerBackend(DeepseekV32IndexerBackend):
-    @staticmethod
-    def get_supported_kernel_block_sizes() -> list:
-        return [16, 32]
-
-
 class DeepseekV32ROCmIndexerCache(DeepseekV32IndexerCache):
     def get_attn_backend(self):
-        return DeepseekV32ROCmIndexerBackend
+        return DeepseekV32IndexerBackend
+
+    @property
+    def uses_shuffled_layout(self) -> bool:
+        # aiter's gather/insert pair shuffles the cache above block size 1:
+        # [n_blocks, blk/16, head_dim/16, 16, 16] instead of [n_blocks, blk, head_dim].
+        return self.kv_cache.ndim == 3 and self.kv_cache.shape[1] != 1
 
 
 class DeepseekV32ROCmIndexer(DeepseekV32Indexer):
@@ -46,7 +42,7 @@ def __init__(self, vllm_config, config, prefix, topk_indices_buffer=None):
             config,
             prefix,
             topk_indices_buffer,
-            attn_backend=DeepseekV32MLASparseBackend,
+            attn_backend=ROCMAiterMLASparseBackend,
         )
 
         self.indexer_op: SparseAttnIndexer | None = None
@@ -100,14 +96,10 @@ def _compute_ql_nope(self, q_c: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso
         q_nope = q_nope.transpose(0, 1)  # (N, tokens, P)
 
         if self.is_aiter_triton_fp4_bmm_enabled:
-            from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
-
-            ql_nope = batched_gemm_a16wfp4(
+            ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4(
                 q_nope, self.W_K, self.W_K_scale, transpose_bm=True, prequant=True
             )
         elif self.is_aiter_triton_fp8_bmm_enabled:
-            from vllm._aiter_ops import rocm_aiter_ops
-
             ql_nope = rocm_aiter_ops.triton_fp8_bmm(
                 q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
             )
@@ -152,14 +144,10 @@ def _compute_uv_out(
         )
 
         if self.is_aiter_triton_fp4_bmm_enabled:
-            from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
-
-            batched_gemm_a16wfp4(
+            rocm_aiter_ops.batched_gemm_a16wfp4(
                 x, self.W_V, self.W_V_scale, out_view, transpose_bm=True, prequant=True
             )
         elif self.is_aiter_triton_fp8_bmm_enabled:
-            from vllm._aiter_ops import rocm_aiter_ops
-
             rocm_aiter_ops.triton_fp8_bmm(
                 x,
                 self.W_V,
@@ -202,6 +190,7 @@ def _fused_attention(
             indexer_k_norm_eps = self.indexer.k_norm.eps
             indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache
             indexer_k_cache = self.indexer.k_cache.kv_cache
+            indexer_cache_shuffled = self.indexer.k_cache.uses_shuffled_layout
             indexer_softmax_scale = self.indexer.softmax_scale
             indexer_n_head_scale = self.indexer.n_head**-0.5
         else:
@@ -211,6 +200,7 @@ def _fused_attention(
             indexer_k_norm_eps = 1e-6
             indexer_k_rope_cos_sin_cache = None
             indexer_k_cache = None
+            indexer_cache_shuffled = False
             indexer_softmax_scale = 0.0
             indexer_n_head_scale = 0.0
 
@@ -241,6 +231,7 @@ def _fused_attention(
             self.topk_indices_buffer,
             slot_mapping=mla_slot,
             indexer_k_cache=indexer_k_cache,
+            indexer_cache_shuffled=indexer_cache_shuffled,
             mla_kv_cache=mla_kv_cache,
             mla_kv_cache_dtype=self.kv_cache_dtype,
             mla_k_scale=mla_k_scale,
@@ -282,7 +273,7 @@ def _fused_attention(
         num_actual = attn_metadata.num_actual_tokens  # type: ignore[attr-defined]
         kv_cache = self.kv_cache
         if self._fp8_kv_needs_view:
-            kv_cache = kv_cache.view(torch.float8_e4m3fn)
+            kv_cache = kv_cache.view(current_platform.fp8_dtype())
 
         q_for_attn = self._build_q_for_attn(ql_nope, mqa_q, num_actual)
         attn_out, _ = self.impl.forward_mqa(  # type: ignore[attr-defined]
diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py
index 6a0d207ff568..e6fb00434a95 100644
--- a/vllm/models/deepseek_v32/attention.py
+++ b/vllm/models/deepseek_v32/attention.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, ClassVar, cast
 
 import torch
 import torch.nn as nn
@@ -42,6 +42,9 @@
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
+    from vllm.v1.attention.backends.mla.index_group import (
+        SparseMLAIndexGroupBuilder,
+    )
 
 
 class DeepseekV32Indexer(nn.Module):
@@ -118,6 +121,7 @@ def __init__(
 class DeepseekV32Attention(MLAAttention):
     indexer: "DeepseekV32Indexer | None"
     indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer
+    supports_pcp_dcp: ClassVar[bool] = True
 
     def __init__(
         self,
@@ -126,6 +130,7 @@ def __init__(
         prefix: str,
         topk_indices_buffer: torch.Tensor | None = None,
         attn_backend: "type | None" = None,
+        index_group_builder: "SparseMLAIndexGroupBuilder | None" = None,
     ) -> None:
         quant_config = vllm_config.quant_config
         cache_config = vllm_config.cache_config
@@ -148,9 +153,9 @@ def __init__(
         # DSA checkpoints may use plain ("default") or yarn-scaled RoPE.
         if config.rope_parameters["rope_type"] != "default":
             config.rope_parameters["rope_type"] = (
-                "deepseek_yarn"
-                if config.rope_parameters.get("apply_yarn_scaling", True)
-                else "deepseek_llama_scaling"
+                "deepseek_llama_scaling"
+                if config.rope_parameters.get("attention_factor") == 1.0
+                else "deepseek_yarn"
             )
         if config.rope_parameters["rope_type"] == "deepseek_yarn":
             mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
@@ -209,6 +214,7 @@ def __init__(
             use_sparse=True,
             indexer=indexer,
             topk_indices_buffer=topk_indices_buffer,
+            index_group_builder=index_group_builder,
             attn_backend=attn_backend,
         )
 
@@ -231,7 +237,9 @@ def __init__(
 
         fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)
         self._fp8_query = fp8_attention and self.impl.supports_quant_query_input
-        self._fp8_kv_needs_view = fp8_attention and self.kv_cache_dtype != "fp8_ds_mla"
+        self._fp8_kv_needs_view = fp8_attention and not self.kv_cache_dtype.endswith(
+            "_ds_mla"
+        )
 
         self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False)
 
@@ -304,6 +312,14 @@ def forward(  # type: ignore[override]
         slot_mapping = forward_context.slot_mapping
         assert isinstance(slot_mapping, dict)
         mla_slot = slot_mapping.get(self.layer_name)
+        indexer_slot = (
+            slot_mapping.get(self.indexer.k_cache.prefix)
+            if self.indexer is not None
+            else None
+        )
+        hisparse_cache = self.hisparse_cache
+        layer_attn_metadata, _, _, _ = get_attention_context(self.layer_name)
+        self.impl.prepare_for_batch(layer_attn_metadata)
 
         if self.indexer is not None and not self.skip_topk:
             has_indexer = True
@@ -311,6 +327,7 @@ def forward(  # type: ignore[override]
             indexer_k_norm_bias = self.indexer.k_norm.bias
             indexer_k_norm_eps = self.indexer.k_norm.eps
             indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache
+            indexer_cache_shuffled = self.indexer.k_cache.uses_shuffled_layout
             indexer_k_cache = None if self.use_pcp else self.indexer.k_cache.kv_cache
             index_k_out = torch.empty_like(index_k) if self.use_pcp else None
             indexer_softmax_scale = self.indexer.softmax_scale
@@ -322,6 +339,7 @@ def forward(  # type: ignore[override]
             indexer_k_norm_eps = 1e-6
             indexer_k_rope_cos_sin_cache = None
             indexer_k_cache = None
+            indexer_cache_shuffled = False
             index_k_out = None
             indexer_softmax_scale = 0.0
             indexer_n_head_scale = 0.0
@@ -331,12 +349,17 @@ def forward(  # type: ignore[override]
             mla_k_scale = None
             indexer_k_cache = None
             mla_slot = None
+            indexer_slot = None
         else:
-            mla_kv_cache = self.kv_cache
+            mla_kv_cache = None if hisparse_cache is not None else self.kv_cache
             mla_k_scale = self._k_scale
 
-        kv_c_out = torch.empty_like(kv_c)
-        k_pe_out = torch.empty_like(k_pe)
+        if self.use_pcp or hisparse_cache is None:
+            kv_c_out = torch.empty_like(kv_c)
+            k_pe_out = torch.empty_like(k_pe)
+        else:
+            kv_c_out = torch.empty_like(kv_c)
+            k_pe_out = torch.empty_like(k_pe)
         q_c = fused_norm_rope(
             positions,
             q_c,
@@ -354,7 +377,9 @@ def forward(  # type: ignore[override]
             indexer_k_rope_cos_sin_cache,
             self.topk_indices_buffer,
             slot_mapping=mla_slot,
+            indexer_slot_mapping=indexer_slot,
             indexer_k_cache=indexer_k_cache,
+            indexer_cache_shuffled=indexer_cache_shuffled,
             mla_kv_cache=mla_kv_cache,
             mla_kv_cache_dtype=self.kv_cache_dtype,
             mla_k_scale=mla_k_scale,
@@ -365,6 +390,18 @@ def forward(  # type: ignore[override]
             index_k_out=index_k_out,
         )
 
+        if hisparse_cache is not None and mla_slot is not None:
+            assert kv_c_out is not None and k_pe_out is not None
+            self.update_kv_cache(
+                kv_c_out,
+                k_pe_out,
+                self.kv_cache,
+                mla_slot,
+                layer_attn_metadata,
+                self.kv_cache_dtype,
+                self._k_scale,
+            )
+
         q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
         q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
         ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1)
@@ -459,6 +496,7 @@ def _sparse_indexer_and_attn(
                 ),
                 skip_topk_buffer_clear=True,
             )
+        self.impl.record_logical_topk_ready()  # type: ignore[attr-defined]
 
         attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(
             self.layer_name
@@ -467,6 +505,7 @@ def _sparse_indexer_and_attn(
             output.zero_()
             return
         attn_metadata = cast("MLACommonMetadata", attn_metadata)
+        self.impl.prepare_for_batch(attn_metadata)
 
         if self.use_pcp:
             assert kv_c is not None and k_pe is not None
@@ -508,7 +547,7 @@ def _sparse_indexer_and_attn(
             return
 
         if self._fp8_kv_needs_view:
-            kv_cache = kv_cache.view(torch.float8_e4m3fn)
+            kv_cache = kv_cache.view(current_platform.fp8_dtype())
         if self._fp8_query:
             # FlashInfer sparse: single packed fp8 query.
             mqa_q_arg: torch.Tensor | tuple[torch.Tensor, torch.Tensor] = mqa_q[
@@ -536,8 +575,9 @@ def _sparse_indexer_and_attn(
             seq_lens: torch.Tensor | None
             query_start_loc: torch.Tensor | None
             if self.use_pcp:
-                if attn_metadata.decode is not None:
-                    seq_lens = attn_metadata.decode.seq_lens
+                decode_metadata = getattr(attn_metadata, "decode", None)
+                if decode_metadata is not None:
+                    seq_lens = decode_metadata.seq_lens
                 else:
                     all_seq_lens = cast(
                         torch.Tensor,
@@ -552,12 +592,23 @@ def _sparse_indexer_and_attn(
                 # PCP-only empty-shard metadata is needed.
                 seq_lens = None
                 query_start_loc = None
-            attn_out = self.dcp_manager.combine(
-                attn_out,
-                lse,
-                seq_lens=seq_lens,  # type: ignore[arg-type]
-                query_start_loc=query_start_loc,  # type: ignore[arg-type]
-            )
+            # Under PCP+DCP the prefill rows attended over the gathered KV, so
+            # only the decode rows carry an LSE and take part in the merge.
+            num_merge_rows = lse.shape[0]
+            if num_merge_rows == attn_out.shape[0]:
+                attn_out = self.dcp_manager.combine(
+                    attn_out,
+                    lse,
+                    seq_lens=seq_lens,  # type: ignore[arg-type]
+                    query_start_loc=query_start_loc,  # type: ignore[arg-type]
+                )
+            elif num_merge_rows > 0:
+                attn_out[:num_merge_rows] = self.dcp_manager.combine(
+                    attn_out[:num_merge_rows],
+                    lse,
+                    seq_lens=seq_lens,  # type: ignore[arg-type]
+                    query_start_loc=query_start_loc,  # type: ignore[arg-type]
+                )
             if self.use_pcp:
                 attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
 
diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py
index 1c91d3015774..e0259c46faaf 100644
--- a/vllm/models/deepseek_v32/common/kernels.py
+++ b/vllm/models/deepseek_v32/common/kernels.py
@@ -12,6 +12,18 @@
 # has_indexer=False path so the indexer args don't allocate every call.
 _DUMMY_CACHE: dict[tuple, torch.Tensor] = {}
 
+# Tile shape of the indexer-K cache's shuffled layout, used when the cache reports
+# uses_shuffled_layout. The block tile is a token count; the head tile is a byte
+# count, which the kernel converts to cache elements.
+_INDEXER_CACHE_BLOCK_TILE = 16
+_INDEXER_CACHE_HEAD_TILE_BYTES = 16
+
+# FNUZ (e4m3fnuz) on gfx942, OCP (e4m3fn) elsewhere. 224.0 keeps the ue8m0
+# power-of-two scale inside the FNUZ range after its ceil.
+_FP8_DTYPE = current_platform.fp8_dtype()
+_USE_FNUZ = torch.float8_e4m3fnuz == _FP8_DTYPE
+_FP8_MAX = 224.0 if _USE_FNUZ else 448.0
+
 
 def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
     key = (shape, dtype, device)
@@ -47,16 +59,17 @@ def _get_cos_sin(
 
 
 @triton.jit
-def _fp8_ue8m0_quantize(vals):
+def _fp8_ue8m0_quantize(vals, FP8_MAX: tl.constexpr, USE_FNUZ: tl.constexpr):
     """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
 
     Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
     """
     vals = vals.to(tl.float32)
     amax = tl.max(tl.abs(vals))
-    scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
+    scale = tl.div_rn(tl.maximum(amax, 1e-4), FP8_MAX)
     scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
-    fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv)
+    fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
+    fp8_vals = tl.div_rn(vals, scale).to(fp8_dtype)
     return fp8_vals, scale
 
 
@@ -68,21 +81,35 @@ def _fp8_quant_and_cache_write(
     kv_cache_ptr,
     kv_cache_scale_ptr,
     cache_block_size,
-    cache_stride,
+    cache_block_stride,
     offsets,
     HEAD_DIM: tl.constexpr,
+    SHUFFLE: tl.constexpr,
+    BLOCK_TILE: tl.constexpr,
+    HEAD_TILE: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    USE_FNUZ: tl.constexpr,
 ):
-    k_fp8, scale = _fp8_ue8m0_quantize(vals)
+    k_fp8, scale = _fp8_ue8m0_quantize(vals, FP8_MAX, USE_FNUZ)
 
     block_idx = slot_idx // cache_block_size
     block_offset = slot_idx % cache_block_size
-    block_start = block_idx * cache_block_size * cache_stride
+    block_start = block_idx * cache_block_stride
 
-    tl.store(
-        kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
-        k_fp8,
-        mask=mask,
-    )
+    # Shuffled layout: [blk/BLOCK_TILE, HEAD_DIM/HEAD_TILE, BLOCK_TILE, HEAD_TILE],
+    # so one contiguous run holds BLOCK_TILE tokens x HEAD_TILE bytes and the
+    # reader's coalesced loads land on the bytes we wrote.
+    if SHUFFLE:
+        value_off = (
+            block_offset // BLOCK_TILE * BLOCK_TILE * HEAD_DIM
+            + block_offset % BLOCK_TILE * HEAD_TILE
+            + offsets // HEAD_TILE * BLOCK_TILE * HEAD_TILE
+            + offsets % HEAD_TILE
+        )
+    else:
+        value_off = block_offset * HEAD_DIM + offsets
+
+    tl.store(kv_cache_ptr + block_start + value_off, k_fp8, mask=mask)
     scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
     tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
 
@@ -129,13 +156,17 @@ def _fused_norm_rope_kernel(
     index_k_out_ptr,
     index_k_out_stride,
     INDEX_K_HALF_ROT_DIM: tl.constexpr,
-    # Cache params (shared by indexer K and MLA)
+    # Cache params
     slot_mapping_ptr,
+    indexer_slot_mapping_ptr,
     # Index K FP8 cache
     indexer_cache_ptr,
     indexer_cache_scale_ptr,
     indexer_cache_block_size,
     indexer_cache_stride,
+    INDEXER_CACHE_SHUFFLE: tl.constexpr,
+    INDEXER_CACHE_BLOCK_TILE: tl.constexpr,
+    INDEXER_CACHE_HEAD_TILE: tl.constexpr,
     # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
     mla_cache_ptr,
     mla_cache_block_stride,
@@ -149,6 +180,7 @@ def _fused_norm_rope_kernel(
     mla_cache_ds_scale_ptr,
     mla_cache_ds_rope_ptr,
     MLA_CACHE_DS_MLA: tl.constexpr,
+    MLA_CACHE_NVFP4: tl.constexpr,
     MLA_NUM_TILES: tl.constexpr,
     MLA_TILE_DIM: tl.constexpr,
     # Top k indices
@@ -159,12 +191,15 @@ def _fused_norm_rope_kernel(
     HAS_INDEXER: tl.constexpr,
     INDEX_ROPE_INTERLEAVE: tl.constexpr,
     USE_PDL: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    USE_FNUZ: tl.constexpr,
 ):
     tok_idx = tl.program_id(0).to(tl.int64)
     pid = tl.program_id(1)
     if USE_PDL:
         tl.extra.cuda.gdc_wait()
         tl.extra.cuda.gdc_launch_dependents()
+    fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
     if pid == 3:
         if not HAS_INDEXER:
             # Shared layer: reuse the previous indexer layer's top-k; do not
@@ -260,10 +295,60 @@ def _fused_norm_rope_kernel(
             mla_block_idx = slot_idx // mla_block_size
             mla_block_off = slot_idx % mla_block_size
 
+            if MLA_CACHE_NVFP4:
+                # nvfp4_ds_mla layout (352 B/token, KV_DIM == 512):
+                #   [0, 256)    512 e2m1 NoPE, packed 2/byte (low nibble = even)
+                #   [256, 320)   64 e4m3 RoPE, unscaled
+                #   [320, 352)   32 e4m3 NoPE scales, byte-permuted
+                # Keep in lockstep with concat_and_cache_nvfp4_ds_mla (CUDA) and
+                # nvfp4_sf_byte() in FlashMLA: byte(s) = 8*(s&3) + (s>>2).
+                byte_base = (
+                    mla_block_idx * mla_cache_block_stride
+                    + mla_block_off * mla_cache_entry_stride
+                )
+                kv_2d = tl.reshape(kv_c.to(tl.float32), (MLA_NUM_TILES, MLA_TILE_DIM))
+                amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)
+                # sf = e4m3(max(amax/6, 2^-9)), round-to-nearest like
+                # cvt_warp_fp16_to_fp4.
+                sf8 = tl.maximum(amax * (1.0 / 6.0), 0.001953125).to(tl.float8e4nv)
+                q = kv_2d * (1.0 / sf8.to(tl.float32))
+                qp = tl.reshape(q, (KV_DIM // 2, 2))
+                sel = tl.arange(0, 2)[None, :]
+                even = tl.sum(tl.where(sel == 0, qp, 0.0), axis=1)
+                odd = tl.sum(tl.where(sel == 1, qp, 0.0), axis=1)
+                # Same instruction the CUDA writer uses, so the two are bit-exact
+                # by construction: dest byte = {high: odd, low: even}.
+                packed = tl.inline_asm_elementwise(
+                    "{ .reg .b8 t; .reg .b32 w;"
+                    " cvt.rn.satfinite.e2m1x2.f32 t, $2, $1;"
+                    " mov.b32 w, {t, t, t, t}; and.b32 $0, w, 255; }",
+                    "=r,f,f",
+                    [even, odd],
+                    dtype=tl.int32,
+                    is_pure=True,
+                    pack=1,
+                ).to(tl.uint8)
+                tl.store(
+                    mla_cache_ptr + byte_base + tl.arange(0, KV_DIM // 2),
+                    packed.to(tl.float8e4nv, bitcast=True),
+                )
+                # RoPE: unscaled e4m3, interleaved like the fp8_ds_mla path
+                rope_base = mla_cache_ptr + byte_base + KV_DIM // 2
+                tl.store(rope_base + dim_off * 2, r1.to(tl.float8e4nv))
+                tl.store(rope_base + dim_off * 2 + 1, r2.to(tl.float8e4nv))
+                # scales, permuted: byte(s) = 8*(s&3) + (s>>2)
+                sf_tile = tl.arange(0, MLA_NUM_TILES)
+                sf_perm = 8 * (sf_tile % 4) + (sf_tile // 4)
+                tl.store(
+                    mla_cache_ptr + byte_base + (KV_DIM // 2 + 64) + sf_perm,
+                    tl.reshape(sf8, (MLA_NUM_TILES,)),
+                )
+                return
+
             if MLA_CACHE_DS_MLA:
                 # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element
                 # tile of the NoPE is dynamically quantized to fp8 with its own
-                # float32 scale; the RoPE tail is stored unquantized in bf16.
+                # power-of-two scale stored as float32; the RoPE tail is bf16.
                 #   bytes [0, KV_DIM)            : KV_DIM fp8 NoPE values
                 #   bytes [KV_DIM, KV_DIM + 16)  : MLA_NUM_TILES float32 scales
                 #   bytes [KV_DIM + 16, ...)     : 2 * KPE_HALF_ROT_DIM bf16 RoPE
@@ -275,10 +360,11 @@ def _fused_norm_rope_kernel(
                 )
                 kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM))
                 tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)
-                # scale = amax / 448 (fp8 e4m3 max), matching the reference
+                # scale = amax / FP8_MAX (fp8 e4m3 max), matching the reference
                 # concat_and_cache_ds_mla kernel; floored to FLT_MIN.
-                tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38)
-                kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,))
+                tile_scale = tl.maximum(tile_amax * (1.0 / FP8_MAX), 1e-4)
+                tile_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(tile_scale)))
+                kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(fp8_dtype), (KV_DIM,))
                 tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8)
                 tile_off = tl.arange(0, MLA_NUM_TILES)
                 tl.store(
@@ -298,16 +384,16 @@ def _fused_norm_rope_kernel(
             # kv_c_normed (KV_DIM elements)
             if MLA_CACHE_FP8:
                 scale = tl.load(mla_cache_scale_ptr)
-                kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)
+                kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(fp8_dtype)
                 tl.store(dst + kv_block, kv_c_fp8)
             else:
                 tl.store(dst + kv_block, kv_c)
             # k_pe_roped (from registers, interleaved layout)
             if MLA_CACHE_FP8:
-                tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))
+                tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(fp8_dtype))
                 tl.store(
                     dst + KV_DIM + dim_off * 2 + 1,
-                    (r2 / scale).to(tl.float8e4nv),
+                    (r2 / scale).to(fp8_dtype),
                 )
             else:
                 tl.store(dst + KV_DIM + dim_off * 2, r1)
@@ -400,8 +486,8 @@ def _fused_norm_rope_kernel(
             )
 
         # PCP inserts index K after gathering; other paths write it directly.
-        if indexer_cache_ptr is not None and slot_mapping_ptr is not None:
-            slot_idx = tl.load(slot_mapping_ptr + tok_idx)
+        if indexer_cache_ptr is not None and indexer_slot_mapping_ptr is not None:
+            slot_idx = tl.load(indexer_slot_mapping_ptr + tok_idx)
             _fp8_quant_and_cache_write(
                 result,
                 index_k_mask,
@@ -412,6 +498,11 @@ def _fused_norm_rope_kernel(
                 indexer_cache_stride,
                 index_k_block,
                 INDEX_K_DIM,
+                INDEXER_CACHE_SHUFFLE,
+                INDEXER_CACHE_BLOCK_TILE,
+                INDEXER_CACHE_HEAD_TILE,
+                FP8_MAX,
+                USE_FNUZ,
             )
 
 
@@ -431,9 +522,11 @@ def fused_norm_rope(
     index_k_layer_norm_eps: float,
     index_k_rope_cos_sin_cache: torch.Tensor | None,
     topk_indices_buffer: torch.Tensor,
-    # Cache params for fused writes (single slot_mapping for both caches)
+    # Cache params for fused writes
     slot_mapping: torch.Tensor | None = None,
+    indexer_slot_mapping: torch.Tensor | None = None,
     indexer_k_cache: torch.Tensor | None = None,
+    indexer_cache_shuffled: bool = False,
     mla_kv_cache: torch.Tensor | None = None,
     mla_kv_cache_dtype: str = "auto",
     mla_k_scale: torch.Tensor | None = None,
@@ -476,18 +569,36 @@ def fused_norm_rope(
     # --- Indexer K cache setup ---
     if indexer_k_cache is not None:
         assert slot_mapping is not None
+        if indexer_slot_mapping is None:
+            indexer_slot_mapping = slot_mapping
         idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
         idx_cache_block_size = indexer_k_cache.shape[1]
         idx_cache_stride = indexer_k_cache.shape[2]
+        # The caller's cache reports whether its reader expects the shuffled
+        # layout; see DeepseekV32IndexerCache.uses_shuffled_layout.
+        idx_cache_shuffle = indexer_cache_shuffled
+        if idx_cache_shuffle:
+            assert idx_cache_block_size % _INDEXER_CACHE_BLOCK_TILE == 0, (
+                f"indexer K cache block size {idx_cache_block_size} must be a "
+                f"multiple of {_INDEXER_CACHE_BLOCK_TILE} for the shuffled layout"
+            )
+        idx_cache_stride = indexer_k_cache.stride(0)
         if indexer_k_cache.dtype == torch.uint8:
-            indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
+            indexer_k_cache = indexer_k_cache.view(_FP8_DTYPE)
+        # The head tile is a byte count; the kernel indexes the cache in elements.
+        idx_cache_head_tile = (
+            _INDEXER_CACHE_HEAD_TILE_BYTES // indexer_k_cache.element_size()
+        )
     else:
         idx_cache_scale_view = None
         idx_cache_block_size = 1
         idx_cache_stride = 0
+        idx_cache_shuffle = False
+        idx_cache_head_tile = _INDEXER_CACHE_HEAD_TILE_BYTES
 
     # --- MLA KV cache setup ---
-    mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla"
+    mla_cache_nvfp4 = mla_kv_cache_dtype == "nvfp4_ds_mla"
+    mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" or mla_cache_nvfp4
     mla_cache_fp8 = is_quantized_kv_cache(mla_kv_cache_dtype) and not mla_cache_ds_mla
     mla_num_tiles = 1
     mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device)
@@ -495,22 +606,23 @@ def fused_norm_rope(
     if mla_kv_cache is not None:
         mla_block_size = mla_kv_cache.shape[1]
         if mla_cache_ds_mla:
-            # 656-byte custom layout addressed in bytes; mla_cache_ptr is the
-            # 1-byte fp8 view, so block/entry strides are byte offsets and the
-            # fp32/bf16 views share the same buffer.
-            assert kv_dim == 512, "fp8_ds_mla requires kv_lora_rank == 512"
-            mla_num_tiles = kv_dim // 128
+            # Custom byte-addressed layout: 656 B/token for fp8_ds_mla, 352 B
+            # for nvfp4_ds_mla. mla_cache_ptr is the 1-byte fp8 view, so
+            # block/entry strides are byte offsets and the fp32/bf16 views
+            # (fp8_ds_mla only) share the same buffer.
+            assert kv_dim == 512, "ds_mla layouts require kv_lora_rank == 512"
+            mla_num_tiles = kv_dim // 16 if mla_cache_nvfp4 else kv_dim // 128
             u8_cache = mla_kv_cache.view(torch.uint8)
             mla_block_stride = u8_cache.stride(0)
             mla_entry_stride = u8_cache.stride(1)
             mla_ds_scale_view = u8_cache.view(torch.float32)
             mla_ds_rope_view = u8_cache.view(torch.bfloat16)
-            mla_kv_cache = u8_cache.view(torch.float8_e4m3fn)
+            mla_kv_cache = u8_cache.view(_FP8_DTYPE)
         else:
             mla_block_stride = mla_kv_cache.stride(0)
             mla_entry_stride = mla_kv_cache.stride(1)
             if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
-                mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
+                mla_kv_cache = mla_kv_cache.view(_FP8_DTYPE)
         if mla_k_scale is None:
             mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
     else:
@@ -536,6 +648,9 @@ def fused_norm_rope(
         assert index_k_out.shape == index_k.shape
         index_k_out_stride = index_k_out.stride(0)
     use_pdl = current_platform.is_arch_support_pdl()
+    pdl_kwargs = (
+        {"USE_PDL": use_pdl, "launch_pdl": True} if use_pdl else {"USE_PDL": use_pdl}
+    )
     _fused_norm_rope_kernel[(num_tokens, 4)](
         positions,
         # Q RMS norm
@@ -578,10 +693,14 @@ def fused_norm_rope(
         index_k_rope_cos_sin_cache.shape[-1] // 2,
         # Cache params
         slot_mapping,
+        indexer_slot_mapping,
         indexer_k_cache,
         idx_cache_scale_view,
         idx_cache_block_size,
         idx_cache_stride,
+        idx_cache_shuffle,
+        _INDEXER_CACHE_BLOCK_TILE,
+        idx_cache_head_tile,
         # MLA KV cache (uses same slot_mapping)
         mla_kv_cache,
         mla_block_stride,
@@ -591,7 +710,8 @@ def fused_norm_rope(
         mla_k_scale,
         mla_ds_scale_view,
         mla_ds_rope_view,
-        mla_cache_ds_mla,
+        mla_cache_ds_mla and not mla_cache_nvfp4,
+        mla_cache_nvfp4,
         mla_num_tiles,
         kv_dim // mla_num_tiles if mla_cache_ds_mla else 1,
         # Top k indices buffer
@@ -601,8 +721,9 @@ def fused_norm_rope(
         TOPK_BLOCK_SIZE=1024,
         HAS_INDEXER=has_indexer,
         INDEX_ROPE_INTERLEAVE=index_rope_interleave,
-        USE_PDL=use_pdl,
-        launch_pdl=use_pdl,
+        FP8_MAX=_FP8_MAX,
+        USE_FNUZ=_USE_FNUZ,
+        **pdl_kwargs,
     )
     return q_c_out
 
@@ -657,6 +778,8 @@ def _fused_q_kernel(
     INDEX_ROPE_INTERLEAVE: tl.constexpr,
     QUANTIZE_MQA: tl.constexpr,
     USE_PDL: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    USE_FNUZ: tl.constexpr,
 ):
     tok_idx = tl.program_id(0).to(tl.int64)
     pid = tl.program_id(1)
@@ -664,6 +787,7 @@ def _fused_q_kernel(
     if USE_PDL:
         tl.extra.cuda.gdc_wait()
         tl.extra.cuda.gdc_launch_dependents()
+    fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
 
     if pid == 2:
         # ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16
@@ -686,7 +810,7 @@ def _fused_q_kernel(
                     + ql_nope_off,
                     mask=ql_nope_mask,
                 ).to(tl.float32)
-                ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv)
+                ql_nope_fp8 = (ql_nope / scale).to(fp8_dtype)
                 tl.store(
                     mqa_q_fp8_ptr
                     + tok_idx * mqa_q_fp8_stride0
@@ -736,7 +860,7 @@ def _fused_q_kernel(
                         + q_head_idx * mqa_q_fp8_stride1
                         + QL_NOPE_DIM
                         + rot_off * 2,
-                        (r1 / scale).to(tl.float8e4nv),
+                        (r1 / scale).to(fp8_dtype),
                     )
                     tl.store(
                         mqa_q_fp8_ptr
@@ -745,7 +869,7 @@ def _fused_q_kernel(
                         + QL_NOPE_DIM
                         + rot_off * 2
                         + 1,
-                        (r2 / scale).to(tl.float8e4nv),
+                        (r2 / scale).to(fp8_dtype),
                     )
                 else:
                     # bf16 query: write the RoPE'd q_pe unquantized.
@@ -803,7 +927,7 @@ def _fused_q_kernel(
         index_q = tl.where(in_rope, roped, index_q)
 
         # Index Q Quantize (from registers)
-        index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q)
+        index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q, FP8_MAX, USE_FNUZ)
         tl.store(
             index_q_fp8_ptr
             + tok_idx * index_q_fp8_stride0
@@ -897,7 +1021,7 @@ def fused_q(
             q_pe.shape[0],
             q_pe.shape[1],
             ql_nope.shape[2] + q_pe.shape[2],
-            dtype=torch.float8_e4m3fn,
+            dtype=_FP8_DTYPE,
             device=q_pe.device,
         )
         # Placeholder; pid 0 packs q_pe into mqa_q_fp8 instead.
@@ -909,7 +1033,7 @@ def fused_q(
         mqa_q_fp8 = q_pe_out  # unused placeholder for the fp8 pack pointer
         mqa_q = q_pe_out
 
-    index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
+    index_q_fp8 = torch.empty_like(index_q, dtype=_FP8_DTYPE)
     index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
     if cutedsl_kernel is not None:
         cutedsl_kernel(
@@ -932,6 +1056,9 @@ def fused_q(
         return index_q_fp8, index_weights_out, mqa_q
 
     use_pdl = current_platform.is_arch_support_pdl()
+    pdl_kwargs = (
+        {"USE_PDL": use_pdl, "launch_pdl": True} if use_pdl else {"USE_PDL": use_pdl}
+    )
     _fused_q_kernel[(num_tokens, 3, grid_heads)](
         positions,
         q_pe,
@@ -973,8 +1100,9 @@ def fused_q(
         HAS_INDEXER=has_indexer,
         INDEX_ROPE_INTERLEAVE=index_rope_interleave,
         QUANTIZE_MQA=quantize_mqa,
-        USE_PDL=use_pdl,
-        launch_pdl=use_pdl,
+        FP8_MAX=_FP8_MAX,
+        USE_FNUZ=_USE_FNUZ,
+        **pdl_kwargs,
         # num_warps=1 is optimal here: each program is a single 128-element
         # rope+quant, so the kernel is program-count/occupancy bound, not
         # per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).
diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py
index 76aa6ca8e161..b42ba28c8df8 100644
--- a/vllm/models/deepseek_v32/nvidia/model.py
+++ b/vllm/models/deepseek_v32/nvidia/model.py
@@ -46,6 +46,10 @@
 )
 from vllm.models.deepseek_v32.attention import DeepseekV32Attention
 from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backends.mla.index_group import (
+    SparseMLAIndexGroupBuilder,
+    get_sparse_mla_index_group_max_rows,
+)
 
 from .glm52_low_latency_gemm import enable_glm52_low_latency_gemm
 
@@ -57,6 +61,7 @@ def __init__(
         prefix: str,
         config=None,
         topk_indices_buffer: torch.Tensor | None = None,
+        index_group_builder: SparseMLAIndexGroupBuilder | None = None,
     ) -> None:
         super().__init__()
 
@@ -80,6 +85,7 @@ def __init__(
             config=config,
             prefix=f"{prefix}.self_attn",
             topk_indices_buffer=topk_indices_buffer,
+            index_group_builder=index_group_builder,
         )
 
         if (
@@ -182,12 +188,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         # DSA is always sparse (has index_topk); allocate the shared top-k
         # buffer the indexer writes and the sparse MLA backend reads.
         self.is_v32 = True
-        topk_indices_buffer = torch.empty(
+        # On the model, not a local: the MTP proposer shares it with the draft.
+        self.topk_indices_buffer = torch.empty(
             vllm_config.scheduler_config.max_num_batched_tokens,
             config.index_topk,
             dtype=torch.int32,
             device=self.device,
         )
+        index_group_builder = SparseMLAIndexGroupBuilder(
+            self.topk_indices_buffer,
+            get_sparse_mla_index_group_max_rows(vllm_config),
+        )
 
         if get_pp_group().is_first_rank:
             self.embed_tokens = make_input_embedding(
@@ -207,7 +218,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             lambda prefix: DeepseekV32DecoderLayer(
                 vllm_config=vllm_config,
                 prefix=prefix,
-                topk_indices_buffer=topk_indices_buffer,
+                topk_indices_buffer=self.topk_indices_buffer,
+                index_group_builder=index_group_builder,
             ),
             prefix=f"{prefix}.layers",
         )
diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py
index 2ea01dbdcb9c..bdc6b95c15e4 100644
--- a/vllm/models/deepseek_v32/nvidia/mtp.py
+++ b/vllm/models/deepseek_v32/nvidia/mtp.py
@@ -33,8 +33,10 @@
     _try_load_fp8_indexer_wk,
     get_spec_layer_idx_from_weight_name,
 )
+from vllm.model_executor.models.interfaces import SupportsPP
 from vllm.model_executor.models.utils import (
     get_pp_missing_layer_names,
+    make_empty_intermediate_tensors_factory,
     maybe_prefix,
 )
 from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm
@@ -46,6 +48,10 @@
 from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backends.mla.index_group import (
+    SparseMLAIndexGroupBuilder,
+    get_sparse_mla_index_group_max_rows,
+)
 
 from .glm52_low_latency_gemm import (
     build_glm52_plan,
@@ -78,6 +84,10 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
             dtype=torch.int32,
             device=current_platform.device_type,
         )
+        index_group_builder = SparseMLAIndexGroupBuilder(
+            topk_indices_buffer,
+            get_sparse_mla_index_group_max_rows(vllm_config),
+        )
         self.shared_head = SharedHead(
             config=config, prefix=prefix, quant_config=quant_config
         )
@@ -86,6 +96,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
             prefix,
             config=config,
             topk_indices_buffer=topk_indices_buffer,
+            index_group_builder=index_group_builder,
         )
 
     def forward(
@@ -260,7 +271,7 @@ def get_top_tokens(
         )
 
 
-class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
+class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
@@ -271,6 +282,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         if self.config.model_type == "glm_moe_dsa":
             enable_glm52_low_latency_gemm(self, vllm_config.model_config.dtype)
         self.set_moe_parameters()
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], self.config.hidden_size
+        )
 
     def set_moe_parameters(self):
         self.num_moe_layers = self.config.num_nextn_predict_layers
@@ -289,7 +303,7 @@ def set_moe_parameters(self):
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
 
-    def forward(
+    def forward(  # type: ignore[override]
         self,
         input_ids: torch.Tensor | None,
         positions: torch.Tensor,
@@ -368,6 +382,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 continue
             spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
             if spec_layer is None:
+                # A tied top-level embed_tokens has no spec layer to rewrite
+                # from; the draft needs its own copy under PP.
+                param = params_dict.get(name) if "embed_tokens" in name else None
+                if param is not None:
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add(name)
                 continue
             name = self._rewrite_spec_layer_name(spec_layer, name)
 
diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py
index 4170a39042ed..6bbd6e0f29b5 100644
--- a/vllm/models/deepseek_v4/__init__.py
+++ b/vllm/models/deepseek_v4/__init__.py
@@ -20,7 +20,7 @@
     )
     from .amd.model import DeepseekV4ForCausalLM
     from .amd.mtp import DeepSeekV4MTP
-    from .vl_stub import (  # type: ignore[assignment]
+    from .common.vl_model import (  # type: ignore[assignment]
         DeepseekV4ForConditionalGeneration,
     )
 elif current_platform.is_xpu():
@@ -38,14 +38,14 @@
         DeepseekV4ForConditionalGeneration,
     )
 else:
+    from .common.vl_model import (  # type: ignore[assignment]
+        DeepseekV4ForConditionalGeneration,
+    )
     from .nvidia.dspark import (  # type: ignore[assignment]
         DSparkDeepseekV4ForCausalLM,
     )
     from .nvidia.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
     from .nvidia.mtp import DeepSeekV4MTP  # type: ignore[assignment]
-    from .nvidia.vl_model import (  # type: ignore[assignment]
-        DeepseekV4ForConditionalGeneration,
-    )
 
 __all__ = [
     "DSparkDeepseekV4ForCausalLM",
diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py
index edca59032aa3..c3f249fa80d5 100644
--- a/vllm/models/deepseek_v4/amd/dspark.py
+++ b/vllm/models/deepseek_v4/amd/dspark.py
@@ -44,6 +44,7 @@
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.qwen3_dspark import (
+    DSparkConfidenceHead,
     DSparkMarkovHead,
 )
 from vllm.model_executor.models.utils import maybe_prefix
@@ -102,7 +103,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             ]
         )
 
-        # Heads: final norm + hc_head, and the Markov head
+        # Heads: final norm + hc_head, and the Markov + confidence heads
         # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         hc_dim = self.hc_mult * config.hidden_size
@@ -125,6 +126,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
             config.dspark_markov_rank,
             prefix=maybe_prefix(prefix, "markov_head"),
         )
+        self.confidence_head: DSparkConfidenceHead | None = None
+        if getattr(config, "enable_confidence_head", True):
+            self.confidence_head = DSparkConfidenceHead(
+                config.hidden_size + config.dspark_markov_rank,
+                prefix=maybe_prefix(prefix, "confidence_head"),
+            )
 
         # MHC head CustomOp dispatcher (aiter / tilelang / triton / torch),
         # replacing the direct nvidia tilelang kernel call.
@@ -358,6 +365,17 @@ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
     def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
         return self.model.markov_head.bias(markov_embed, self.logits_processor)
 
+    def compute_confidence(
+        self, head_hidden: torch.Tensor, markov_embed: torch.Tensor
+    ) -> torch.Tensor:
+        """Per-position acceptance probability for each drafted token."""
+        if self.model.confidence_head is None:
+            raise RuntimeError(
+                "compute_confidence() requires a confidence head, but the "
+                "checkpoint did not provide confidence_head weights."
+            )
+        return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))
+
     # --- Weight loading ----------------------------------------------------
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -391,6 +409,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
+        loaded_confidence_head = False
 
         tp_size = get_tensor_model_parallel_world_size()
         tp_rank = get_tensor_model_parallel_rank()
@@ -403,6 +422,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             if mapped is None:
                 continue
             name = mapped
+            if "confidence_head." in name:
+                loaded_confidence_head = True
 
             # ``.scale`` -> per-method scale suffix.
             if name.endswith(".scale"):
@@ -439,8 +460,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 continue
 
             # Stacked rules only apply to decoder-layer weights. Head-stack params
-            # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.
-            # "markov_w1" would collide with the "w1" shard rule.
+            # (main_proj/norm/hc_head/markov_head/confidence_head) load directly —
+            # otherwise e.g. "markov_w1" would collide with the "w1" shard rule.
             is_layer_param = name.startswith("model.layers.")
             for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
                 if not is_layer_param or weight_name not in name:
@@ -469,6 +490,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                 weight_loader(param, loaded_weight)
                 loaded_params.add(name)
 
+        if self.model.confidence_head is not None and not loaded_confidence_head:
+            self.model.confidence_head = None
         logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
         return loaded_params
 
@@ -482,8 +505,7 @@ def _remap_dspark_name(self, name: str) -> str | None:
             return None
         stage = int(m.group(1))
         rest = m.group(2)
-        # The confidence head is not wired into inference yet; drop its weights.
-        if rest.startswith("confidence_head."):
+        if rest.startswith("confidence_head.") and self.model.confidence_head is None:
             return None
         # Head-stack params live at model level (mtp.last), context combiner at
         # model level (mtp.0); everything else is a per-layer decoder block.
@@ -493,6 +515,7 @@ def _remap_dspark_name(self, name: str) -> str | None:
             "hc_head_base",
             "hc_head_scale",
             "markov_head.",
+            "confidence_head.",
         )
         if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
             head_prefixes
diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py
index fdedf7cf18f8..2e6922e55b98 100644
--- a/vllm/models/deepseek_v4/amd/model.py
+++ b/vllm/models/deepseek_v4/amd/model.py
@@ -12,7 +12,7 @@
 
 import vllm.envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
-from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config import ParallelConfig, VllmConfig
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_rank,
@@ -76,6 +76,8 @@
 from vllm.platforms.rocm import on_gfx950, on_gfx1250
 from vllm.sequence import IntermediateTensors
 
+from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
+
 logger = init_logger(__name__)
 
 
@@ -509,11 +511,28 @@ def forward_modular(
         return routed + shared_out
 
 
-def _fuse_shared_experts_enabled(config) -> bool:
+def _fuse_shared_experts_enabled(config, parallel_config: ParallelConfig) -> bool:
+    if (
+        getattr(config, "n_shared_experts", None)
+        and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
+        and not parallel_config.enable_expert_parallel
+        and parallel_config.data_parallel_size > 1
+    ):
+        # Fused shared experts are not supported under data parallelism:
+        # the fused path cannot load the FP8 shared expert into the MXFP4
+        # routed slot, which fails obscurely during weight loading. Fail
+        # fast with an actionable message instead.
+        raise ValueError(
+            "DeepSeek-V4 fused shared experts "
+            "(VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1) are not supported "
+            "with data parallelism (data_parallel_size > 1). Set "
+            "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 to run the shared "
+            "expert as an unfused FP8 MLP."
+        )
     return bool(
         getattr(config, "n_shared_experts", None)
         and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
-        and not get_current_vllm_config().parallel_config.enable_expert_parallel
+        and not parallel_config.enable_expert_parallel
     )
 
 
@@ -551,6 +570,10 @@ def __init__(
 
         self.gate.e_score_correction_bias = None
         self.gate.tid2eid = None
+        self.gate.bias_vl = None
+        self.image_sentinel_lo = (
+            IMAGE_SENTINEL_BASE_ID if getattr(config, "vision_n_layers", 0) > 0 else 0
+        )
         is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
         self.hash_indices_dtype = torch.int32
         if is_hash_moe:
@@ -566,12 +589,20 @@ def __init__(
                 ),
                 requires_grad=False,
             )
-        elif getattr(config, "topk_method", None) == "noaux_tc":
+        if getattr(config, "topk_method", None) == "noaux_tc" and (
+            not is_hash_moe or getattr(config, "vision_n_layers", 0) > 0
+        ):
             self.gate.e_score_correction_bias = nn.Parameter(
                 torch.empty(config.n_routed_experts, dtype=torch.float32),
                 requires_grad=False,
             )
 
+        if getattr(config, "vision_n_layers", 0) > 0:
+            self.gate.bias_vl = nn.Parameter(
+                torch.empty(config.n_routed_experts, dtype=torch.float32),
+                requires_grad=False,
+            )
+
         self.n_shared_experts = config.n_shared_experts
 
         # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1`
@@ -579,7 +610,7 @@ def __init__(
         # This should be cleaned up and use `resolve_layer_fused_shared_expert`.
         self.fuse_heterogeneous_shared_expert = fuse_heterogeneous_shared_expert
         fse_requested = (
-            _fuse_shared_experts_enabled(config)
+            _fuse_shared_experts_enabled(config, vllm_config.parallel_config)
             and not self.fuse_heterogeneous_shared_expert
         )
         fse_compatible = False
@@ -640,6 +671,8 @@ def __init__(
             routed_scaling_factor=self.routed_scaling_factor,
             e_score_correction_bias=self.gate.e_score_correction_bias,
             hash_indices_table=self.gate.tid2eid,
+            bias_vl=self.gate.bias_vl,
+            image_sentinel_lo=self.image_sentinel_lo,
             swiglu_limit=self.swiglu_limit,
             router_logits_dtype=torch.float32,
             routed_experts_cls=(
@@ -662,6 +695,8 @@ def forward(
     ) -> torch.Tensor:
         if self.gate.tid2eid is not None and input_ids is None:
             raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
+        if self.gate.bias_vl is not None and input_ids is None:
+            raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.")
 
         org_shape = hidden_states.shape
         final_hidden_states = self.experts(
@@ -941,16 +976,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.hc_dim = self.hc_mult * config.hidden_size
         self.rms_norm_eps = config.rms_norm_eps
 
-        # Three aux streams: one per non-default input GEMM in
-        # DeepseekV4Attention._run_parallel_input_projections
-        # (compressor kv_score, indexer.weights_proj, indexer.compressor
-        # kv_score). fused_wqa_wkv stays on the default stream.
-        # Disable them on ROCm because of hang issues.
-        aux_stream_list = (
-            None
-            if current_platform.is_rocm()
-            else [torch.cuda.Stream() for _ in range(3)]
-        )
+        aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
 
         self.device = current_platform.device_type
         # Reserved topk indices buffer for all Indexer layers to reuse.
@@ -1353,6 +1379,7 @@ def _make_deepseek_v4_weights_mapper(
 
 class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
     model_cls = DeepseekV4Model
+    finalizes_weights_during_load = False
 
     # Default mapper assumes the original FP4-expert checkpoint layout.
     # Overridden per-instance in __init__ when expert_dtype != "fp4".
@@ -1399,6 +1426,9 @@ def compute_logits(
         logits = self.logits_processor(self.lm_head, hidden_states)
         return logits
 
+    def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
+
     def forward(
         self,
         input_ids: torch.Tensor,
diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py
index cafb6bb0ed24..045a731321b9 100644
--- a/vllm/models/deepseek_v4/amd/mtp.py
+++ b/vllm/models/deepseek_v4/amd/mtp.py
@@ -160,7 +160,7 @@ def forward(
             inputs_embeds
         ).unsqueeze(-2)
         hidden_states, residual, post_mix, res_mix = self.mtp_block(
-            positions=positions, x=hidden_states, input_ids=None
+            positions=positions, x=hidden_states, input_ids=input_ids
         )
         if self.mtp_block.use_fused_mhc:
             hidden_states = self.mtp_block.hc_post(
@@ -496,8 +496,7 @@ def _resolve_scale_name(name: str) -> str:
         return loaded_params
 
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py
index 59b3d5faa745..fe01e3928f92 100644
--- a/vllm/models/deepseek_v4/amd/rocm.py
+++ b/vllm/models/deepseek_v4/amd/rocm.py
@@ -7,6 +7,8 @@
 
 import torch
 
+from vllm import envs
+from vllm.config import VllmConfig
 from vllm.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
@@ -23,7 +25,9 @@
 from vllm.platforms import current_platform
 from vllm.platforms.rocm import _ON_GFX950, on_gfx1250
 from vllm.triton_utils import tl, triton
+from vllm.utils.multi_stream_utils import execute_in_parallel
 from vllm.v1.attention.backend import (
+    AttentionCGSupport,
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.mla.sparse_swa import (
@@ -33,14 +37,59 @@
 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
     build_ragged_indices_from_dense,
     rocm_inv_rope_einsum,
+    rocm_inverse_rope_rows_,
     rocm_sparse_attn_decode,
     rocm_sparse_attn_prefill,
 )
+from vllm.v1.kv_cache_interface import KVCacheSpec
 from vllm.v1.worker.workspace import current_workspace_manager
 
 logger = init_logger(__name__)
 
 
+def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None:
+    """Normalize checkpoint WO_A scales to raw OCP MX E8M0 bytes.
+
+    E8M0 is an unsigned exponent-only scale format with bias 127. A finite
+    encoded byte ``b`` in ``[0, 254]`` represents ``2 ** (b - 127)``;
+    ``0xFF`` is reserved for NaN. This is the vendor-neutral OCP encoding used
+    by AMD AITER/OPUS, not an NVIDIA-specific convention. Reference: OCP
+    Microscaling Formats (MX) Specification, section 5.4.1:
+    https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
+
+    Loaders may preserve the encoded byte as ``float8_e8m0fnu``/``uint8``,
+    or decode it to a floating-point power of two. Floating-point inputs are
+    accepted only when the original E8M0 byte can be recovered losslessly.
+    This function does not quantize or round arbitrary scales.
+    """
+    if scale.dtype == torch.float8_e8m0fnu:
+        # Reinterpret the native E8M0 storage. ``to(uint8)`` would perform a
+        # numeric conversion instead of preserving the encoded exponent byte.
+        return scale.view(torch.uint8).contiguous()
+    if scale.dtype == torch.uint8:
+        # The checkpoint loader already exposed the E8M0 wire representation.
+        return scale.contiguous()
+    if not scale.dtype.is_floating_point:
+        return None
+
+    scale_f32 = scale.detach().float()
+    if not bool(torch.isfinite(scale_f32).all()) or bool((scale_f32 <= 0).any()):
+        return None
+
+    # With no mantissa, E8M0 can represent only exact powers of two. Rebuild
+    # the value before encoding so this adapter never silently quantizes a
+    # general floating-point checkpoint scale.
+    exponent = torch.round(torch.log2(scale_f32))
+    if not torch.equal(torch.exp2(exponent), scale_f32):
+        return None
+
+    encoded = exponent.to(torch.int32) + 127
+    # 0xFF is NaN in OCP E8M0, not a finite exponent.
+    if int(encoded.min()) < 0 or int(encoded.max()) > 254:
+        return None
+    return encoded.to(torch.uint8).contiguous()
+
+
 def _trust_dsv4_extra_cache_nan_free(
     kv_cache_dtype: str,
     has_kv_transfer: bool,
@@ -106,13 +155,18 @@ def _combine_topk_swa_indices_kernel(
     query_start_loc_ptr,
     seq_lens_ptr,
     gather_lens_ptr,
+    left_visible_ptr,
+    right_visible_ptr,
     M,
     N,
     TOP_K: tl.constexpr,
     COMPRESS_RATIO: tl.constexpr,
     WINDOW_SIZE: tl.constexpr,
+    SWA_WIDTH: tl.constexpr,
     TOPK_WIDTH: tl.constexpr,
     PADDED_TOP_K: tl.constexpr,
+    PADDED_SWA_WIDTH: tl.constexpr,
+    HAS_IMAGE: tl.constexpr,
 ):
     batch_idx = tl.program_id(0)
     worker_id = tl.program_id(1)
@@ -131,7 +185,20 @@ def _combine_topk_swa_indices_kernel(
         token_idx_in_query = token_idx - query_start
         pos = start_pos + token_idx_in_query
         topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
-        swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
+        if HAS_IMAGE:
+            left = tl.load(left_visible_ptr + token_idx)
+            right = tl.load(right_visible_ptr + token_idx)
+        else:
+            left = 0
+            right = 0
+        left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
+        # Prefix caching can resume inside an image span. Do not generate
+        # indices outside the SWA rows present in the gathered workspace.
+        swa_start = tl.maximum(
+            tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0), gather_start
+        )
+        swa_end = tl.minimum(pos + right + 1, seq_len)
+        swa_len = tl.maximum(swa_end - swa_start, 0)
 
         topk_offset = tl.arange(0, PADDED_TOP_K)
         topk_mask = topk_offset < topk_len
@@ -149,14 +216,14 @@ def _combine_topk_swa_indices_kernel(
             mask=topk_mask,
         )
 
-        swa_offset = tl.arange(0, WINDOW_SIZE)
+        swa_offset = tl.arange(0, PADDED_SWA_WIDTH)
         tl.store(
             combined_indices_ptr
             + token_idx * combined_indices_stride
             + topk_len
             + swa_offset,
-            M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start,
-            mask=swa_offset < swa_len,
+            M * batch_idx + N + swa_offset + swa_start - gather_start,
+            mask=(swa_offset < swa_len) & (swa_offset < SWA_WIDTH),
         )
 
         tl.store(combined_lens_ptr + token_idx, topk_len + swa_len)
@@ -172,12 +239,21 @@ def combine_topk_swa_indices(
     topk: int,
     M: int,
     N: int,
+    max_image_tokens: int = 0,
+    left_visible: torch.Tensor | None = None,
+    right_visible: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    if (left_visible is None) != (right_visible is None):
+        raise ValueError("left_visible and right_visible must be provided together")
     topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
     num_tokens = topk_indices.shape[0]
     num_reqs = seq_lens.shape[0]
+    has_image = left_visible is not None
+    # Keep the row shape fixed for a vision model even when a particular batch
+    # has no image.
+    swa_width = window_size + max_image_tokens
     combined_topk = (
-        (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
+        (topk + swa_width + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
         // _SPARSE_PREFILL_TOPK_ALIGNMENT
         * _SPARSE_PREFILL_TOPK_ALIGNMENT
     )
@@ -201,13 +277,18 @@ def combine_topk_swa_indices(
         query_start_loc,
         seq_lens,
         gather_lens,
+        left_visible if left_visible is not None else topk_indices,
+        right_visible if right_visible is not None else topk_indices,
         M,
         N,
         TOP_K=topk,
         COMPRESS_RATIO=compress_ratio,
         WINDOW_SIZE=window_size,
+        SWA_WIDTH=swa_width,
         TOPK_WIDTH=topk_indices.shape[-1],
         PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]),
+        PADDED_SWA_WIDTH=triton.next_power_of_2(swa_width),
+        HAS_IMAGE=has_image,
     )
     return combined_indices, combined_lens
 
@@ -368,6 +449,20 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
 
 
 class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
+    @classmethod
+    def get_cudagraph_support(
+        cls,
+        vllm_config: VllmConfig,
+        kv_cache_spec: KVCacheSpec,
+    ) -> AttentionCGSupport:
+        spec_config = vllm_config.speculative_config
+        if spec_config is not None and spec_config.enable_adaptive_verification:
+            # All per-token metadata is built from device query boundaries into
+            # persistent buffers, so adaptive verification can replay varlen
+            # FULL decode graphs after reallocating drafts across requests.
+            return AttentionCGSupport.ALWAYS
+        return super().get_cudagraph_support(vllm_config, kv_cache_spec)
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None
@@ -435,6 +530,20 @@ def build_for_cudagraph_capture(
 
 
 class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
+    @classmethod
+    def get_cudagraph_support(
+        cls,
+        vllm_config: VllmConfig,
+        kv_cache_spec: KVCacheSpec,
+    ) -> AttentionCGSupport:
+        spec_config = vllm_config.speculative_config
+        if spec_config is not None and spec_config.enable_adaptive_verification:
+            # SWA indices, lengths, and token-to-request mappings are built from
+            # device boundaries into persistent buffers, so adaptive verification
+            # can replay varlen FULL decode graphs safely.
+            return AttentionCGSupport.ALWAYS
+        return super().get_cudagraph_support(vllm_config, kv_cache_spec)
+
     # Keep fused multi-step decode disabled until update_draft_decode_metadata()
     # also refreshes the ROCm-specific ragged SWA indices and indptrs.
     supports_draft_decode_metadata_update = False
@@ -462,11 +571,13 @@ def build(
         common_prefix_len: int,
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
+        replay_start: torch.Tensor | None = None,
     ) -> DeepseekV4ROCMAiterSparseSWAMetadata:
         base = super().build(
             common_prefix_len=common_prefix_len,
             common_attn_metadata=common_attn_metadata,
             fast_build=fast_build,
+            replay_start=replay_start,
         )
 
         ragged_indices = None
@@ -520,10 +631,233 @@ def __init__(self, *args, **kwargs):
         # Block scale for the preshuffled weight; None = not preshuffled.
         self._wqa_wkv_scale: torch.Tensor | None = None
         self._wo_b_scale: torch.Tensor | None = None
+        self._wo_a_fp8_weight: torch.Tensor | None = None
+        self._wo_a_e8m0_scale: torch.Tensor | None = None
+        self._wo_a_cos_cache: torch.Tensor | None = None
+        self._wo_a_sin_cache: torch.Tensor | None = None
         self._fused_compressor_weight: torch.Tensor | None
         self.register_buffer("_fused_compressor_weight", None, persistent=False)
         self._fused_compressor_split_sizes: tuple[int, int] | None = None
 
+        if self.indexer is None:
+            # Dense layers have no compressor work to overlap; HCA layers
+            # (compressor, no indexer) keep the streams for the dual-stream
+            # fork below.
+            if self.compressor is None:
+                self.aux_stream_list = None
+        else:
+            # Disable indexer inner overlap.
+            self.indexer.aux_stream = None
+
+    def _enable_multi_stream_overlap(self) -> bool:
+        """ROCm multi-stream gates: streams and capture region.
+
+        Dict metadata marks piecewise cudagraph, whose eager breaks rebuild
+        the attention inputs on the owning stream. Forking side streams
+        there would rely on runtime HIP event sync, which is unreliable in
+        this overlap on ROCm (event waits can hang), so multi-stream only
+        runs where the fork/join becomes static graph edges: inside capture,
+        or with non-dict metadata (full cudagraph or the profile run), which
+        has no eager breaks. Covers both the HCA and CSA forks.
+        """
+        attn_metadata = get_forward_context().attn_metadata
+        return self.aux_stream_list is not None and (
+            torch.cuda.is_current_stream_capturing()
+            or not isinstance(attn_metadata, dict)
+        )
+
+    def _run_sequential_pipeline(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        o_padded: torch.Tensor,
+    ) -> None:
+        """Disable ROCm streams when the current execution region cannot overlap."""
+        aux_streams = self.aux_stream_list
+        self.aux_stream_list = None
+        try:
+            qr_kv, kv_score, indexer_kv_score, indexer_weights = (
+                self._run_parallel_input_projections(hidden_states)
+            )
+            qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
+            self._prepare_and_attn_fn(
+                hidden_states,
+                qr,
+                kv,
+                qr_scale,
+                kv_score,
+                indexer_kv_score,
+                indexer_weights,
+                positions,
+                o_padded,
+            )
+        finally:
+            self.aux_stream_list = aux_streams
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        llama_4_scaling: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Pre-allocate attention output with FlashMLA-padded head count.
+        # The op writes into `o_padded`; we slice to n_local_heads after.
+        num_tokens = hidden_states.shape[0]
+        o_padded = torch.empty(
+            (num_tokens, self.padded_heads, self.head_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+
+        if self._enable_multi_stream_overlap():
+            # The ROCm override consumes these sentinels inside the capture
+            # boundary, moving the stream fan-out ahead of the projections.
+            self._prepare_and_attn_fn(
+                hidden_states,
+                None,
+                None,
+                None,
+                None,
+                None,
+                None,
+                positions,
+                o_padded,
+            )
+        else:
+            self._run_sequential_pipeline(hidden_states, positions, o_padded)
+
+        o = o_padded[:, : self.n_local_heads, :]
+
+        # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
+        return self._o_proj(o, positions)
+
+    def _prepare_and_attn(
+        self,
+        hidden_states: torch.Tensor,
+        qr: torch.Tensor | None,
+        kv: torch.Tensor | None,
+        qr_scale: torch.Tensor | None,
+        kv_score: torch.Tensor | None,
+        indexer_kv_score: torch.Tensor | None,
+        indexer_weights: torch.Tensor | None,
+        positions: torch.Tensor,
+        o_padded: torch.Tensor,
+    ) -> None:
+        """Run the ROCm fork/join (HCA or CSA) inside the capture boundary."""
+        aux_streams = self.aux_stream_list
+        # The sequential pipeline disables aux_stream_list before calling
+        # back with real projection inputs; aux_streams is None ends that
+        # recursion here.
+        if aux_streams is None:
+            saved_streams = self.aux_stream_list
+            self.aux_stream_list = None
+            try:
+                super()._prepare_and_attn(
+                    hidden_states,
+                    cast(torch.Tensor, qr),
+                    cast(torch.Tensor, kv),
+                    qr_scale,
+                    cast(torch.Tensor, kv_score),
+                    cast(torch.Tensor, indexer_kv_score),
+                    cast(torch.Tensor, indexer_weights),
+                    positions,
+                    o_padded,
+                )
+            finally:
+                self.aux_stream_list = saved_streams
+            return
+
+        # Re-check: forward's gate ran inside a captured segment that
+        # _prepare_and_attn_eager (MRV1) then broke, making this region eager.
+        if not self._enable_multi_stream_overlap():
+            self._run_sequential_pipeline(hidden_states, positions, o_padded)
+            return
+
+        indexer = self.indexer
+        compressor = self.compressor
+        assert compressor is not None
+
+        def default_chain():
+            qr_kv = self._fused_wqa_wkv_gemm(hidden_states)
+            qr_out, qr_scale_out, kv_out = self._split_qkv_and_norm(qr_kv)
+            q = self._wq_b_proj(qr_out, qr_scale_out).view(
+                -1, self.n_local_heads, self.head_dim
+            )
+            attn_metadata = get_forward_context().attn_metadata
+            q = self._fused_qnorm_rope_kv_insert(q, kv_out, positions, attn_metadata)
+            return q, qr_out, qr_scale_out, kv_out
+
+        def main_compressor_chain() -> None:
+            score = torch.mm(
+                hidden_states,
+                compressor.fused_wkv_wgate.weight.T,
+                out_dtype=torch.float32,
+            )
+            compressor(score, positions, self.rotary_emb)
+
+        if indexer is None:
+            # HCA dual-stream: the main compressor runs on aux stream 0 while
+            # the default stream produces q and inserts KV into the SWA cache.
+            # Both branches only read hidden_states, so the join merely has to
+            # precede the sparse attention that consumes the compressed KV.
+            (q, _qr_out, _qr_scale_out, kv_out), _ = execute_in_parallel(
+                default_chain,
+                [main_compressor_chain],
+                self.ln_events[0],
+                [self.ln_events[1]],
+                aux_streams[:1],
+                enable=True,
+            )
+            self._sparse_indexer_and_attn(
+                hidden_states, None, None, None, q, kv_out, positions, o_padded
+            )
+            return
+
+        def indexer_compressor_chain() -> None:
+            score = torch.mm(
+                hidden_states,
+                indexer.compressor.fused_wkv_wgate.weight.T,
+                out_dtype=torch.float32,
+            )
+            indexer.compressor(score, positions, self.indexer_rotary_emb)
+
+        # CSA three-stream: the main and indexer compressors run on aux
+        # streams 0 and 1 while the default stream produces q and inserts KV
+        # into the SWA cache. Every branch only reads hidden_states plus its
+        # own state, so the join merely has to precede the indexer op and the
+        # sparse attention, which consume the compressed KV caches.
+        (q, qr_out, qr_scale_out, kv_out), _ = execute_in_parallel(
+            default_chain,
+            [main_compressor_chain, indexer_compressor_chain],
+            self.ln_events[0],
+            self.ln_events[1:3],
+            aux_streams[:2],
+            enable=True,
+        )
+
+        indexer_weights_out, _ = indexer.weights_proj(hidden_states)
+        # The indexer compressor already ran on aux stream 1; build queries only.
+        index_q, index_q_scale, weights = indexer(
+            hidden_states,
+            qr_out,
+            None,
+            indexer_weights_out,
+            positions,
+            self.indexer_rotary_emb,
+            qr_scale_out,
+            skip_compressor=True,
+        )
+        self._sparse_indexer_and_attn(
+            hidden_states,
+            index_q,
+            index_q_scale,
+            weights,
+            q,
+            kv_out,
+            positions,
+            o_padded,
+        )
+
     @classmethod
     def get_padded_num_q_heads(cls, num_heads: int) -> int:
         return num_heads
@@ -568,6 +902,93 @@ def _prep(linear) -> torch.Tensor | None:
 
         self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
         self._wo_b_scale = _prep(self.wo_b)
+        if _ON_GFX950 and envs.VLLM_ROCM_USE_AITER_FP8BMM:
+            self._prepare_fp8_wo_a()
+
+    def _prepare_fp8_wo_a(self) -> None:
+        try:
+            from aiter.ops.batched_gemm_op_a8w8 import (
+                batched_gemm_a8w8_mxscale as mxscale_op,
+            )
+            from aiter.ops.inverse_rope_group_quant import (
+                inverse_rope_group_quant as inverse_quant_op,
+            )
+        except ImportError:
+            logger.warning_once(
+                "The DeepSeek V4 FP8 WO_A path requires AITER >= 0.1.20; "
+                "falling back to BF16 WO_A."
+            )
+            return
+        del mxscale_op, inverse_quant_op
+
+        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+            get_fp8_block_weight_scale,
+            is_fp8,
+        )
+
+        weight = getattr(self.wo_a, "weight", None)
+        scale = get_fp8_block_weight_scale(self.wo_a)
+        if scale is None:
+            # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the
+            # historical ``_inv`` suffix.
+            scale = getattr(self.wo_a, "weight_scale", None)
+        if weight is None or scale is None:
+            logger.warning_once(
+                "DeepSeek V4 FP8 WO_A needs a block-scaled FP8 wo_a weight; "
+                "the layer exposes no weight/weight scale. Falling back to "
+                "BF16 WO_A."
+            )
+            return
+        if weight.dim() != 2 or scale.dim() != 2 or not is_fp8(weight.dtype):
+            logger.warning_once(
+                "DeepSeek V4 FP8 WO_A needs a 2-D FP8 wo_a weight with a 2-D "
+                "block scale, got weight %s%s and scale %s. Falling back to "
+                "BF16 WO_A.",
+                weight.dtype,
+                tuple(weight.shape),
+                tuple(scale.shape),
+            )
+            return
+
+        groups = self.n_local_groups
+        out_per_group = self.o_lora_rank
+        out_features, in_features = weight.shape
+        if (
+            out_features != groups * out_per_group
+            or out_per_group % 128 != 0
+            or in_features % 128 != 0
+            or scale.shape != (out_features // 128, in_features // 128)
+        ):
+            logger.warning_once(
+                "DeepSeek V4 FP8 WO_A needs group-128 blocks for %d groups of "
+                "%d outputs, got weight %s and scale %s. Falling back to BF16 "
+                "WO_A.",
+                groups,
+                out_per_group,
+                tuple(weight.shape),
+                tuple(scale.shape),
+            )
+            return
+
+        e8m0_scale = _wo_a_block_scale_to_e8m0(scale)
+        if e8m0_scale is None:
+            logger.warning_once(
+                "DeepSeek V4 FP8 WO_A could not losslessly encode the %s wo_a "
+                "block scale as OCP E8M0. Falling back to BF16 WO_A.",
+                scale.dtype,
+            )
+            return
+
+        self._wo_a_fp8_weight = weight.view(groups, out_per_group, in_features)
+        self._wo_a_e8m0_scale = e8m0_scale.view(
+            groups, out_per_group // 128, in_features // 128
+        )
+        cache = getattr(self.rotary_emb, "cos_sin_cache_bf16", None)
+        if cache is None:
+            cache = self.rotary_emb.cos_sin_cache.to(dtype=torch.bfloat16)
+        cos_cache, sin_cache = cache.chunk(2, dim=-1)
+        self._wo_a_cos_cache = cos_cache.contiguous()
+        self._wo_a_sin_cache = sin_cache.contiguous()
 
     def prepare_compressor_gemm_fusion(self) -> bool:
         if self._fused_compressor_weight is not None:
@@ -726,17 +1147,45 @@ def _split_qkv_and_norm(
         )
 
     def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
-        # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
-        z = rocm_inv_rope_einsum(
-            self.rotary_emb,
-            o,
-            positions,
-            self.rope_head_dim,
-            self.n_local_groups,
-            self.o_lora_rank,
-            self.wo_a,
-        )
-        zf = z.flatten(1)
+        if self._wo_a_fp8_weight is not None:
+            from aiter.ops.batched_gemm_op_a8w8 import (
+                batched_gemm_a8w8_mxscale,
+            )
+            from aiter.ops.inverse_rope_group_quant import (
+                inverse_rope_group_quant,
+            )
+
+            assert self._wo_a_cos_cache is not None
+            assert self._wo_a_sin_cache is not None
+            o_fp8, o_scale = inverse_rope_group_quant(
+                o.view(o.shape[0], self.n_local_heads, self.head_dim),
+                positions.to(torch.int64),
+                self._wo_a_cos_cache,
+                self._wo_a_sin_cache,
+                num_groups=self.n_local_groups,
+                quant_group_size=128,
+            )
+            assert self._wo_a_e8m0_scale is not None
+            zf = batched_gemm_a8w8_mxscale(
+                o_fp8,
+                self._wo_a_fp8_weight,
+                o_scale,
+                self._wo_a_e8m0_scale,
+                dtype=o.dtype,
+            ).flatten(1)
+        else:
+            # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
+            z = rocm_inv_rope_einsum(
+                self.rotary_emb,
+                o,
+                positions,
+                self.rope_head_dim,
+                self.n_local_groups,
+                self.o_lora_rank,
+                self.wo_a,
+                inverse_rope=False,
+            )
+            zf = z.flatten(1)
         if self._wo_b_scale is not None and zf.dim() == 2:
             return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
         return self.wo_b(zf)
@@ -805,9 +1254,15 @@ def forward_mqa(
                 attn_metadata=rocm_metadata,
                 swa_metadata=swa_metadata,
             )
+        # The fp8 wo_a path rotates inside inverse_rope_group_quant, so folding
+        # the rotation into the decode reduce would apply it twice. Only the
+        # BF16 einsum path hands its rotation off to the decode.
+        fuse_inv_rope = self._wo_a_fp8_weight is None
+        rotated = 0
         if num_decodes > 0:
-            self._forward_decode(
+            rotated = self._forward_decode(
                 q=q[:num_decode_tokens],
+                positions=positions[:num_decode_tokens] if fuse_inv_rope else None,
                 kv_cache=self_kv_cache,
                 swa_metadata=swa_metadata,
                 attn_metadata=rocm_metadata,
@@ -821,17 +1276,31 @@ def forward_mqa(
                     and rocm_metadata.for_cudagraph_capture
                 ),
             )
+        if fuse_inv_rope:
+            # Only the decode reduce rotates its own rows, and only the leading
+            # `rotated` of them; prefill rows and any decode path that did not
+            # fuse still owe the standalone pass. Settle that here rather than
+            # in _o_proj: the split is batch-dependent and _o_proj runs
+            # compiled, where such a value freezes at its trace-time value.
+            rocm_inverse_rope_rows_(
+                output[rotated:, : self.n_local_heads, :],
+                positions[rotated:],
+                self.rotary_emb.cos_sin_cache,
+                self.rope_head_dim,
+            )
 
     def _forward_decode(
         self,
         q: torch.Tensor,
+        positions: torch.Tensor | None,
         kv_cache: torch.Tensor | None,
         swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
         attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
         swa_only: bool,
         output: torch.Tensor,
         adaptive_splits: bool,
-    ) -> None:
+    ) -> int:
+        """Returns how many leading rows the decode epilogue inverse-RoPE'd."""
         num_decodes = swa_metadata.num_decodes
         num_decode_tokens = swa_metadata.num_decode_tokens
 
@@ -863,7 +1332,7 @@ def _forward_decode(
                 topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices
                 topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr
 
-        rocm_sparse_attn_decode(
+        return rocm_sparse_attn_decode(
             q=q,
             kv_cache=kv_cache,
             swa_k_cache=self.swa_cache_layer.kv_cache,
@@ -883,6 +1352,8 @@ def _forward_decode(
             rope_head_dim=self.rope_head_dim,
             output=output,
             adaptive_splits=adaptive_splits,
+            inv_rope_positions=positions,
+            inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
             extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
                 self.kv_cache_dtype,
                 self._has_kv_transfer,
@@ -917,6 +1388,12 @@ def _forward_prefill(
         assert query_start_loc_cpu is not None
         assert query_start_loc is not None
         prefill_token_base = query_start_loc_cpu[num_decodes]
+        left_visible = swa_metadata.prefill_left_visible
+        right_visible = swa_metadata.prefill_right_visible
+        if left_visible is not None:
+            left_visible = left_visible[num_decode_tokens:]
+            assert right_visible is not None
+            right_visible = right_visible[num_decode_tokens:]
 
         if not swa_only:
             if self.compress_ratio == 4:
@@ -995,6 +1472,17 @@ def _forward_prefill(
                 top_k,
                 M,
                 N,
+                max_image_tokens=self.max_image_tokens,
+                left_visible=(
+                    left_visible[query_start:query_end]
+                    if left_visible is not None
+                    else None
+                ),
+                right_visible=(
+                    right_visible[query_start:query_end]
+                    if right_visible is not None
+                    else None
+                ),
             )
             rocm_sparse_attn_prefill(
                 q=q[query_start:query_end],
diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py
index 12205bc0431e..8a151dba1d98 100644
--- a/vllm/models/deepseek_v4/attention.py
+++ b/vllm/models/deepseek_v4/attention.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-DeepseekV4 MLA Attention Layer
-"""
+"""DeepseekV4 MLA Attention Layer."""
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
@@ -50,6 +48,7 @@
 from vllm.model_executor.models.utils import extract_layer_index
 from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
 from vllm.models.deepseek_v4.compressor import DeepseekCompressor
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.multi_stream_utils import (
     execute_in_parallel,
@@ -317,7 +316,6 @@ def __init__(
             # graph and MRV1 produces garbage (#51430).
             self._prepare_and_attn_fn = self._prepare_and_attn_eager
 
-        # Will be None on ROCm for now.
         self.aux_stream_list = aux_stream_list
         # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
         # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
@@ -420,7 +418,6 @@ def __init__(
 
                 _COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup()
 
-            from vllm.platforms import current_platform
             from vllm.utils.import_utils import has_cutedsl
 
             _FUSED_Q_KV_RMSNORM_KERNEL.register_warmup()
@@ -447,7 +444,13 @@ def __init__(
                     )
 
                     _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL.register_warmup()
-                    if not has_cutedsl():
+                    if has_cutedsl():
+                        from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (  # noqa: E501
+                            _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
+                        )
+
+                        _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL.register_warmup()
+                    else:
                         _DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL.register_warmup()
                 elif backend_name == "FLASHINFER_MLA_SPARSE_DSV4":
                     from vllm.models.deepseek_v4.common.ops.cache_utils import (
@@ -578,7 +581,7 @@ def project_query_and_cache_kv() -> torch.Tensor:
 
         # Keep Q projection and KV insertion on the default stream. The indexer
         # and MLA compressor use aux streams 0 and 1; aux 2 is internal to the
-        # indexer. ROCm runs the same work sequentially without aux streams.
+        # indexer.
         if indexer is not None:
             assert compressor is not None
             q, (indexer_inputs, _) = execute_in_parallel(
@@ -665,7 +668,6 @@ def _run_parallel_input_projections(
         # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs
         # on aux streams 0..2 when their owning module exists. ln_events[0]
         # is the fan-out start event; ln_events[1..3] are per-aux done events.
-        # On ROCm, aux_streams is None and execute_in_parallel runs serially.
         aux_fns: list[Callable[[], Any] | None] = [None, None, None]
 
         if self.compressor is not None:
@@ -934,6 +936,21 @@ def __init__(
             "Using %s indexer cache for Lightning Indexer.",
             "MXFP4" if self.use_fp4_kv else "FP8",
         )
+        if vllm_config.kernel_config.enable_jit_warmup:
+            from vllm.utils.import_utils import has_cutedsl
+
+            if current_platform.is_cuda() and has_cutedsl():
+                from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (  # noqa: E501
+                    _INDEXER_Q_FP8_KERNEL,
+                    _INDEXER_Q_MXFP4_KERNEL,
+                )
+
+                indexer_q_kernel = (
+                    _INDEXER_Q_MXFP4_KERNEL
+                    if self.use_fp4_kv
+                    else _INDEXER_Q_FP8_KERNEL
+                )
+                indexer_q_kernel.register_warmup()
 
         # no tensor parallel, just replicated
         self.wq_b = ReplicatedLinear(
@@ -1016,7 +1033,6 @@ def __init__(
         ]
 
         if vllm_config.kernel_config.enable_jit_warmup:
-            from vllm.platforms import current_platform
             from vllm.utils.import_utils import has_cutedsl
 
             if not has_cutedsl() and not current_platform.is_xpu():
@@ -1035,11 +1051,12 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         qr: torch.Tensor,
-        compressed_kv_score: torch.Tensor,
+        compressed_kv_score: torch.Tensor | None,
         indexer_weights: torch.Tensor,
         positions: torch.Tensor,
         rotary_emb: nn.Module,
         qr_scale: torch.Tensor | None = None,
+        skip_compressor: bool = False,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
         compressor = self.compressor
 
@@ -1052,7 +1069,8 @@ def forward(
             ):
                 # candidates num smaller than topk, every candidate is selected
                 # but we still need to build k cache
-                compressor(compressed_kv_score, positions, rotary_emb)
+                if not skip_compressor:
+                    compressor(compressed_kv_score, positions, rotary_emb)
                 assert self.topk_indices_buffer is not None
                 num_tokens = (
                     indexer_metadata.num_decode_tokens
@@ -1082,15 +1100,18 @@ def wq_b_and_q_quant():
                 use_fp4=self.use_fp4_kv,
             )
 
-        # compressor returns None and writes K to the indexer KV cache; the
-        # join orders that write before indexer_op (skip_k_cache_insert=True).
-        (q_quant, weights), _ = maybe_execute_in_parallel(
-            wq_b_and_q_quant,
-            lambda: compressor(compressed_kv_score, positions, rotary_emb),
-            self.ln_events[0],
-            self.ln_events[1],
-            self.aux_stream,
-        )
+        if not skip_compressor:
+            # compressor returns None and writes K to the indexer KV cache; the
+            # join orders that write before indexer_op (skip_k_cache_insert=True).
+            (q_quant, weights), _ = maybe_execute_in_parallel(
+                wq_b_and_q_quant,
+                lambda: compressor(compressed_kv_score, positions, rotary_emb),
+                self.ln_events[0],
+                self.ln_events[1],
+                self.aux_stream,
+            )
+        else:
+            q_quant, weights = wq_b_and_q_quant()
         if isinstance(q_quant, tuple):
             q, q_scale = q_quant
         else:
diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py
index 919146c551a4..d28848bb8f8c 100644
--- a/vllm/models/deepseek_v4/common/mm_preprocess.py
+++ b/vllm/models/deepseek_v4/common/mm_preprocess.py
@@ -22,7 +22,7 @@
 
 import math
 from collections.abc import Mapping, Sequence
-from typing import Any, cast
+from typing import Any
 
 import numpy as np
 import torch
@@ -30,7 +30,7 @@
 from transformers import BatchFeature
 from typing_extensions import assert_never
 
-from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
 from vllm.multimodal.parse import ImageSize, MultiModalDataItems
@@ -375,7 +375,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         size = self.info.get_image_size_with_most_features()
         return {
@@ -383,7 +383,7 @@ def get_dummy_mm_data(
                 width=size.width,
                 height=size.height,
                 num_images=mm_counts.get("image", 0),
-                overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
+                overrides=mm_options.get("image"),
             ),
         }
 
diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py
index 1bd0706f0fe3..0cc784e34822 100644
--- a/vllm/models/deepseek_v4/common/ops/cache_utils.py
+++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index
+"""Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index
 preparation.
 
 - quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into
@@ -24,6 +23,7 @@
 )
 from vllm.model_executor.warmup.jit_warmup import (
     WarmupIntRange,
+    kernel_launcher,
     zip_inputs,
 )
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
@@ -31,7 +31,6 @@
     TritonPointerInputVariant,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
     triton_scalar_specialization_rep,
 )
 from vllm.platforms import current_platform
@@ -61,8 +60,7 @@ def quantize_and_insert_k_kernel(
     n_quant_blocks: tl.constexpr,  # 8 (7 real + 1 padding)
     use_fnuz: tl.constexpr = False,
 ):
-    """
-    Quantize K tensor and insert into paged K cache.
+    """Quantize K tensor and insert into paged K cache.
 
     K Cache block layout (block_size=64 tokens):
     - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16
@@ -174,8 +172,7 @@ def quantize_and_insert_k_cache(
     is_ue8m0: bool = True,
     use_fnuz: bool = False,
 ):
-    """
-    Quantize K tensor and insert into paged K cache.
+    """Quantize K tensor and insert into paged K cache.
 
     K Cache block layout (block_size=64 tokens):
     - First 64 * 576 = 36864 bytes: Token data
@@ -521,11 +518,17 @@ def dequantize_and_gather_k_cache(
     if has_cutedsl():
         # lazily import, otherwise some tests fail due to CUDA driver init failure.
         from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
-            dequantize_and_gather_k_cache_cutedsl,
+            _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
         )
 
-        dequantize_and_gather_k_cache_cutedsl(
-            out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
+        _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL(
+            out=out,
+            k_cache=k_cache,
+            seq_lens=seq_lens,
+            gather_lens=gather_lens,
+            block_table=block_table,
+            block_size=block_size,
+            offset=offset,
         )
         return
 
@@ -1231,7 +1234,7 @@ def build_flashinfer_mixed_sparse_indices(
 def _remap_flashinfer_index(values, block_size, block_span):
     # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
     # packed pages (#44577) need block*block_size+off -> block*block_span+off.
-    # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
+    # The launcher declares these indices with sparse_indices_are_storage_offsets=True.
     is_valid = values >= 0
     safe_values = tl.where(is_valid, values, 0)
     values = (safe_values // block_size) * block_span
diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
index e0bc2a042831..50ca86558e05 100644
--- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Fused compressor + FP8/MXFP4 UE8M0 quantization + KV cache insert kernels.
+"""Fused compressor + FP8/MXFP4 UE8M0 quantization + KV cache insert kernels.
 
 Three specialized kernels:
   - _fused_kv_compress_norm_rope_insert_sparse_attn:
@@ -26,11 +25,11 @@
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
@@ -390,7 +389,7 @@ def _compress_gather_split_sparse_attn(
     NUM_SPLITS: tl.constexpr,
     HEAD_TILE: tl.constexpr,  # HEAD_SIZE // NUM_SPLITS
 ):
-    """Stage 1: per-(token, head-split) compress gather, write to fp32 scratch
+    """Stage 1: per-(token, head-split) compress gather, write to fp32 scratch.
 
     No-overlap gather (cr>=128) on rows [0, COMPRESS_RATIO)
     """
@@ -639,7 +638,7 @@ def compress_norm_rope_store_two_stage_triton(
     num_decode_tokens: int,
     compress_scratch: torch.Tensor,
 ) -> None:
-    """Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap)
+    """Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap).
 
     Run the occupancy-fanned two-stage split for prefill [num_decodee_tokens:]
     to fill the CUs, and use the original single-pass launcher
diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
index ae014750bf90..669186edbb1b 100644
--- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
@@ -6,11 +6,11 @@
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
@@ -309,12 +309,22 @@ def __call__(
         )
 
 
+def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]:
+    """Weights dtypes the model's indexer layers ask for: fp32 for the dense
+    scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer
+    (`SparseMQAIndexer`) is enabled."""
+    if vllm_config.attention_config.indexer_sparse_logits:
+        return (torch.float32, torch.bfloat16)
+    return (torch.float32,)
+
+
 class FusedIndexerQRopeMxFp4TritonKernel(
     VllmTritonJitKernel["FusedIndexerQRopeMxFp4TritonKernel.CompileKey"]
 ):
     @dataclass(frozen=True)
     class CompileKey:
         dtype: torch.dtype
+        weights_out_dtype: torch.dtype
         num_heads: int
         index_q_half_rot_dim: int
         index_q_head_dim: int
@@ -434,12 +444,14 @@ def dispatch(  # type: ignore[override]
         self,
         *,
         dtype: torch.dtype,
+        weights_out_dtype: torch.dtype,
         num_heads: int,
         head_dim: int,
         rope_dim: int,
     ) -> CompileKey:
         return self.CompileKey(
             dtype=dtype,
+            weights_out_dtype=weights_out_dtype,
             num_heads=num_heads,
             index_q_half_rot_dim=rope_dim // 2,
             index_q_head_dim=head_dim,
@@ -456,6 +468,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
 
         return self._trace_dispatch(self.dispatch)(
             dtype=vllm_config.model_config.dtype,
+            weights_out_dtype=_indexer_weights_out_dtypes(vllm_config),
             num_heads=num_heads,
             head_dim=head_dim,
             rope_dim=rope_dim,
@@ -491,7 +504,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
                 shape=(1, compile_key.num_heads, scale_head_dim),
             ),
             index_weights_out=TritonWarmupTensor(
-                torch.float32,
+                compile_key.weights_out_dtype,
                 shape=(1, compile_key.num_heads),
             ),
         )
@@ -540,12 +553,17 @@ def fused_indexer_q_rope_quant(
     index_weights_softmax_scale: float,
     index_weights_head_scale: float,
     use_fp4: bool = False,
+    weights_out_dtype: torch.dtype = torch.float32,
 ) -> tuple[
     torch.Tensor | tuple[torch.Tensor, torch.Tensor],
     torch.Tensor,
 ]:
     """Fused RoPE + quantize Q for the sparse indexer.
 
+    ``weights_out_dtype`` is the dtype the downstream scoring kernel takes:
+    fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse
+    MQA-logits kernels (CUDA MXFP4 path only).
+
     Weight-fold semantics (important — the two paths differ):
 
     FP8 path (use_fp4=False, default):
@@ -577,7 +595,7 @@ def fused_indexer_q_rope_quant(
     num_index_q_heads = index_q.shape[1]
     index_q_head_dim = index_q.shape[2]
 
-    index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
+    index_weights_out = torch.empty_like(index_weights, dtype=weights_out_dtype)
 
     if use_fp4:
         assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
@@ -598,21 +616,22 @@ def fused_indexer_q_rope_quant(
         if has_cutedsl():
             # lazily import, otherwise some tests fail due to CUDA driver init failure.
             from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
-                fused_indexer_q_rope_quant_mxfp4_cutedsl,
+                _INDEXER_Q_MXFP4_KERNEL,
             )
 
-            fused_indexer_q_rope_quant_mxfp4_cutedsl(
-                positions,
-                index_q,
-                index_q_cos_sin_cache,
-                index_weights,
-                index_weights_softmax_scale,
-                index_weights_head_scale,
-                index_q_packed,
-                index_q_scale,
-                index_weights_out,
+            _INDEXER_Q_MXFP4_KERNEL(
+                positions=positions,
+                q=index_q,
+                cos_sin_cache=index_q_cos_sin_cache,
+                weights=index_weights,
+                weights_softmax_scale=index_weights_softmax_scale,
+                weights_head_scale=index_weights_head_scale,
+                q_packed=index_q_packed,
+                q_scale=index_q_scale,
+                weights_out=index_weights_out,
             )
         elif current_platform.is_xpu():
+            assert weights_out_dtype == torch.float32, weights_out_dtype
             torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4(
                 index_q,
                 positions,
@@ -647,6 +666,7 @@ def fused_indexer_q_rope_quant(
             index_q_scale.view(torch.int32).squeeze(-1),
         ), index_weights_out
 
+    assert weights_out_dtype == torch.float32, weights_out_dtype
     fp8_dtype = current_platform.fp8_dtype()
     use_fnuz = fp8_dtype == torch.float8_e4m3fnuz
     fp8_max = 224.0 if use_fnuz else 448.0
@@ -654,18 +674,18 @@ def fused_indexer_q_rope_quant(
     if has_cutedsl():
         # lazily import, otherwise some tests fail due to CUDA driver init failure.
         from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
-            fused_indexer_q_rope_quant_fp8_cutedsl,
+            _INDEXER_Q_FP8_KERNEL,
         )
 
-        fused_indexer_q_rope_quant_fp8_cutedsl(
-            positions,
-            index_q,
-            index_q_cos_sin_cache,
-            index_weights,
-            index_weights_softmax_scale,
-            index_weights_head_scale,
-            index_q_fp8,
-            index_weights_out,
+        _INDEXER_Q_FP8_KERNEL(
+            positions=positions,
+            q=index_q,
+            cos_sin_cache=index_q_cos_sin_cache,
+            weights=index_weights,
+            weights_softmax_scale=index_weights_softmax_scale,
+            weights_head_scale=index_weights_head_scale,
+            q_fp8=index_q_fp8.view(torch.uint8),
+            weights_out=index_weights_out,
         )
     elif current_platform.is_xpu():
         torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8(
diff --git a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
index 948d84be5c4a..dd25c73af7c2 100644
--- a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+++ b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Fused inverse RoPE + block-scaled FP8 quantization kernel for DeepseekV4 attention.
+"""Fused inverse RoPE + block-scaled FP8 quantization kernel for DeepseekV4 attention.
 
 Output scale format is pre-transformed (MN-major TMA-aligned; FP32 on SM90,
 INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layout.
@@ -12,11 +11,11 @@
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
@@ -380,6 +379,7 @@ def fused_inv_rope_fp8_quant(
     Returns:
         Rotated output in [T, G, D] and its FP8 scales. The scale tensor is
         empty when quantization is disabled.
+
     """
     from vllm.utils.deep_gemm import get_tma_aligned_size
 
diff --git a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py
index c71043f1ae09..59ac4d27c648 100644
--- a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py
+++ b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py
@@ -22,11 +22,11 @@
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import next_power_of_2
diff --git a/vllm/models/deepseek_v4/common/ops/save_partial_states.py b/vllm/models/deepseek_v4/common/ops/save_partial_states.py
index 0537feab2956..6a554b9b4e02 100644
--- a/vllm/models/deepseek_v4/common/ops/save_partial_states.py
+++ b/vllm/models/deepseek_v4/common/ops/save_partial_states.py
@@ -6,11 +6,11 @@
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import next_power_of_2
diff --git a/vllm/models/deepseek_v4/common/rope.py b/vllm/models/deepseek_v4/common/rope.py
index 46125e40f5da..5339f3548a6f 100644
--- a/vllm/models/deepseek_v4/common/rope.py
+++ b/vllm/models/deepseek_v4/common/rope.py
@@ -31,9 +31,9 @@ def build_deepseek_v4_rope(
     if compress_ratio > 1 and rope_parameters["rope_type"] != "default":
         # YaRN applies only to compressor (CSA/HCA) layers.
         rope_parameters["rope_type"] = (
-            "deepseek_yarn"
-            if rope_parameters.get("apply_yarn_scaling", True)
-            else "deepseek_llama_scaling"
+            "deepseek_llama_scaling"
+            if rope_parameters.get("attention_factor") == 1.0
+            else "deepseek_yarn"
         )
     else:
         # Sliding-window layers use plain RoPE (theta=rope_theta, no YaRN).
diff --git a/vllm/models/deepseek_v4/common/vision.py b/vllm/models/deepseek_v4/common/vision.py
index 46468897e59d..bfc02768c06c 100644
--- a/vllm/models/deepseek_v4/common/vision.py
+++ b/vllm/models/deepseek_v4/common/vision.py
@@ -40,8 +40,7 @@
 )
 
 
-@lru_cache(8)
-def get_vision_cos_sin(
+def _compute_vision_cos_sin(
     n_h: int, n_w: int, dim: int, theta: float
 ) -> tuple[torch.Tensor, torch.Tensor]:
     inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
@@ -52,6 +51,13 @@ def get_vision_cos_sin(
     return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)
 
 
+@lru_cache(8)
+def get_vision_cos_sin(
+    n_h: int, n_w: int, dim: int, theta: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return _compute_vision_cos_sin(n_h, n_w, dim, theta)
+
+
 def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
     dtype = x.dtype
     x1, x2 = x.float().chunk(2, dim=-1)
@@ -123,15 +129,23 @@ def __init__(self, config, prefix: str = ""):
         )
 
     def forward(
-        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cu_seqlens: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
     ) -> torch.Tensor:
         n = x.size(0)
         qkv, _ = self.wqkv(x)
         q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1))
         q = apply_rotary(q, cos, sin).unsqueeze(0)  # (b=1, n, h, d)
         k = apply_rotary(k, cos, sin).unsqueeze(0)
-        # One image per call: a dense batch, no varlen packing metadata.
-        o = self.attn(q, k, v.unsqueeze(0))
+        # Dense (cu_seqlens=None): one image per call. Varlen (cu_seqlens
+        # set): packed multi-image batch for encoder CUDA graph replay.
+        o = self.attn(
+            q, k, v.unsqueeze(0), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
+        )
         out, _ = self.wo(o.reshape(n, -1))
         return out
 
@@ -173,9 +187,14 @@ def __init__(self, config, prefix: str = ""):
         self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp")
 
     def forward(
-        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cu_seqlens: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.attn(self.norm1(x), cos, sin, cu_seqlens, max_seqlen)
         return x + self.mlp(self.norm2(x))
 
 
@@ -206,6 +225,115 @@ def forward(
             x = block(x, cos, sin)
         return self.norm(x)
 
+    def forward_packed(
+        self,
+        patches: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: torch.Tensor,
+    ) -> torch.Tensor:
+        """Varlen path for encoder CUDA graphs: multiple images packed along
+        rows, per-image attention via ``cu_seqlens``, RoPE tables precomputed
+        on device. Pure-tensor (no host reads, no H2D), safe to capture."""
+        x = self.patch_embed(patches)
+        for block in self.blocks:
+            x = block(x, cos, sin, cu_seqlens, max_seqlen)
+        return self.norm(x)
+
+
+def build_packed_vit_metadata(
+    grids: list[list[int]] | list[tuple[int, int]],
+    *,
+    rope_dim: int,
+    rope_theta: float,
+    device: torch.device,
+    max_seqlen_override: int | None = None,
+    cached: bool = True,
+) -> dict[str, torch.Tensor]:
+    """Precompute packed-batch ViT metadata for ``DeepseekV4ViT.forward_packed``.
+
+    Args:
+        grids: ``[n_vit_h, n_vit_w]`` per image, in packing order.
+        rope_dim: RoPE embedding dimension of the ViT.
+        rope_theta: RoPE base frequency of the ViT.
+        device: Device for the returned tensors (except ``max_seqlen``,
+            which stays on the host).
+        max_seqlen_override: Worst-case value baked in at CUDA graph capture
+            (the attention wrapper reads ``max_seqlen`` on the host, so the
+            capture-time value becomes a graph constant).
+        cached: Use the shared ``lru_cache`` for RoPE tables. Capture-time
+            dummy grids pass ``False`` to avoid evicting real entries.
+
+    """
+    cos_sin = get_vision_cos_sin if cached else _compute_vision_cos_sin
+    cos_list: list[torch.Tensor] = []
+    sin_list: list[torch.Tensor] = []
+    cu_seqlens = [0]
+    max_seqlen = 0
+    for n_h, n_w in grids:
+        cos, sin = cos_sin(n_h, n_w, rope_dim, rope_theta)
+        cos_list.append(cos)
+        sin_list.append(sin)
+        n = n_h * n_w
+        cu_seqlens.append(cu_seqlens[-1] + n)
+        max_seqlen = max(max_seqlen, n)
+    if cos_list:
+        cos = torch.cat(cos_list).to(device)
+        sin = torch.cat(sin_list).to(device)
+    else:
+        cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)
+        sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)
+    if max_seqlen_override is not None:
+        max_seqlen = max_seqlen_override
+    return {
+        "vit_cos": cos,
+        "vit_sin": sin,
+        "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32, device=device),
+        # Read on the host by the attention wrapper; keep on CPU.
+        "max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32),
+    }
+
+
+def build_packed_merge_metadata(
+    grids: list[list[int]] | list[tuple[int, int]],
+    downsample_ratio: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, torch.Tensor]:
+    """Gather indices/mask for ``DeepseekV4Aligner.forward_packed``.
+
+    For each aligner output row, the ``r x r`` source patch rows in the
+    packed ViT output; positions past the image edge (the eager path pads
+    them with zeros) get mask 0 and a clamped in-range index.
+    """
+    r = downsample_ratio
+    idx_list: list[torch.Tensor] = []
+    mask_list: list[torch.Tensor] = []
+    offset = 0
+    for n_h, n_w in grids:
+        rows = -(-n_h // r)
+        cols = -(-n_w // r)
+        bi = torch.arange(rows).repeat_interleave(cols)
+        bj = torch.arange(cols).repeat(rows)
+        di = torch.arange(r).repeat_interleave(r)
+        dj = torch.arange(r).repeat(r)
+        hi = bi[:, None] * r + di[None, :]
+        wj = bj[:, None] * r + dj[None, :]
+        valid = (hi < n_h) & (wj < n_w)
+        idx = offset + hi.clamp(max=n_h - 1) * n_w + wj.clamp(max=n_w - 1)
+        idx_list.append(idx)
+        mask_list.append(valid)
+        offset += n_h * n_w
+    if idx_list:
+        merge_idx = torch.cat(idx_list).to(device)
+        merge_mask = torch.cat(mask_list).unsqueeze(-1).to(device=device, dtype=dtype)
+    else:
+        merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device)
+        merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device)
+    return {"merge_idx": merge_idx, "merge_mask": merge_mask}
+
 
 class DeepseekV4Aligner(nn.Module):
     """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector."""
@@ -240,6 +368,27 @@ def forward(self, x: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor:
         out, _ = self.w2(F.gelu(hidden))
         return out
 
+    def forward_packed(
+        self,
+        x: torch.Tensor,
+        merge_idx: torch.Tensor,
+        merge_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Packed multi-image path for encoder CUDA graphs.
+
+        ``merge_idx``/``merge_mask`` come from ``build_packed_merge_metadata``;
+        gather+mask reproduces the eager ``F.pad`` + ``F.unfold`` merge
+        exactly (unfold rows are channel-major, so the gathered patch rows
+        are transposed before flattening).
+        """
+        r2 = self.downsample_ratio**2
+        m = merge_idx.shape[0]
+        gathered = (x[merge_idx] * merge_mask).view(m, r2, -1)
+        gathered = gathered.permute(0, 2, 1).reshape(m, -1)
+        hidden, _ = self.w1(gathered)
+        out, _ = self.w2(F.gelu(hidden))
+        return out
+
 
 def run_dp_sharded_vision_tower(
     vision_model: DeepseekV4ViT,
@@ -263,6 +412,7 @@ def run_dp_sharded_vision_tower(
 
     Returns:
         One ``(n_aligner_rows, hidden_size)`` embedding tensor per image.
+
     """
     tp_size = get_tensor_model_parallel_world_size()
     tp_rank = get_tensor_model_parallel_rank()
diff --git a/vllm/models/deepseek_v4/common/vl_model.py b/vllm/models/deepseek_v4/common/vl_model.py
new file mode 100644
index 000000000000..97701cc2d507
--- /dev/null
+++ b/vllm/models/deepseek_v4/common/vl_model.py
@@ -0,0 +1,370 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp).
+
+Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``:
+
+- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE
+  sentinel positions; four learned vectors (``image_start`` / ``image_pad`` /
+  ``image_newline`` / ``image_end``) fill the remaining sentinel positions.
+- Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor
+  in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved
+  in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>``
+  (see ``common/vision.py`` for the tower itself).
+- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before
+  its hyper-connection stream expansion. Raw ``input_ids`` still flow into
+  the model so the MoE router can apply ``bias_vl`` to image tokens
+  (``requires_raw_input_tokens``).
+"""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsEagle3,
+    SupportsLoRA,
+    SupportsMultiModal,
+    SupportsPP,
+)
+from vllm.model_executor.models.module_mapping import MultiModelKeys
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalKwargsItem
+
+from .mm_preprocess import (
+    IMAGE_PLACEHOLDER,
+    IMAGE_SENTINEL_BASE_ID,
+    DeepseekV4VLDummyInputsBuilder,
+    DeepseekV4VLMultiModalProcessor,
+    DeepseekV4VLProcessingInfo,
+    image_sentinel_mask,
+)
+from .vision import DeepseekV4Aligner, DeepseekV4ViT
+
+
+def _make_deepseek_v4_vl_weights_mapper(
+    text_mapper: WeightsMapper, image_enabled: bool
+) -> WeightsMapper:
+    """Text-checkpoint mapping rules re-rooted under ``language_model.``."""
+    orig_to_new_prefix = {
+        src: None if dst is None else f"language_model.{dst}"
+        for src, dst in text_mapper.orig_to_new_prefix.items()
+    }
+    if not image_enabled:
+        orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None})
+    return WeightsMapper(
+        orig_to_new_renaming=text_mapper.orig_to_new_renaming,
+        orig_to_new_prefix=orig_to_new_prefix,
+        orig_to_new_regex=text_mapper.orig_to_new_regex,
+        orig_to_new_stacked=text_mapper.orig_to_new_stacked,
+        orig_to_new_suffix={
+            **text_mapper.orig_to_new_suffix,
+            "head.weight": "language_model.lm_head.weight",
+        },
+        orig_to_new_substr={
+            **text_mapper.orig_to_new_substr,
+            # Draft models load the checkpoint's MTP weights separately from
+            # the target model.
+            "mtp.": None,
+        },
+    )
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    DeepseekV4VLMultiModalProcessor,
+    info=DeepseekV4VLProcessingInfo,
+    dummy_inputs=DeepseekV4VLDummyInputsBuilder,
+)
+class DeepseekV4ForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3, SupportsLoRA
+):
+    """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower.
+
+    ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
+    delegates through ``language_model`` via the protocol defaults.
+    """
+
+    packed_modules_mapping = {
+        "gate_up_proj": ["w1", "w3"],
+        "fused_wqa_wkv": ["wq_a", "wkv"],
+        "fused_wkv_wgate": ["wkv", "wgate"],
+        # for visual encoder
+        "wqkv": ["wqkv"],
+        "w1": ["w1"],
+    }
+
+    # The MoE router needs raw token ids to detect image sentinel tokens
+    # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl.
+    requires_raw_input_tokens = True
+    supports_tower_connector_lora = True
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality == "image":
+            return IMAGE_PLACEHOLDER
+        raise ValueError(f"Unsupported modality: {modality!r}")
+
+    def __init__(self, *, vllm_config, prefix: str = "") -> None:
+        super().__init__()
+        model_config = vllm_config.model_config
+        config = model_config.hf_config
+        self.config = config
+        self.multimodal_config = model_config.multimodal_config
+        assert self.multimodal_config is not None
+
+        image_enabled = (
+            config.vision_n_layers > 0
+            and self.multimodal_config.get_limit_per_prompt("image") > 0
+        )
+        with self._mark_tower_model(vllm_config, {"image"}):
+            self.vision: DeepseekV4ViT | None = None
+            self.aligner: DeepseekV4Aligner | None = None
+            self.image_start: nn.Parameter | None = None
+            self.image_end: nn.Parameter | None = None
+            self.image_newline: nn.Parameter | None = None
+            self.image_pad: nn.Parameter | None = None
+            if image_enabled:
+                self.vision = DeepseekV4ViT(config)
+                self.aligner = DeepseekV4Aligner(config)
+                for name in (
+                    "image_start",
+                    "image_end",
+                    "image_newline",
+                    "image_pad",
+                ):
+                    setattr(
+                        self,
+                        name,
+                        nn.Parameter(
+                            torch.empty(config.hidden_size, dtype=torch.float32)
+                        ),
+                    )
+                self.vision.to(dtype=model_config.dtype)
+                self.aligner.to(dtype=model_config.dtype)
+
+        with self._mark_language_model(vllm_config):
+            # The arch convertor routes any config with a vision tower to
+            # this wrapper class; mark the copy handed to the text backbone
+            # so it resolves to DeepseekV4ForCausalLM instead of recursing
+            # (with_hf_config deepcopies the config, the marker survives).
+            config._dsv4_vl_inner = True  # type: ignore[attr-defined]
+            try:
+                self.language_model = init_vllm_registered_model(
+                    vllm_config=vllm_config,
+                    hf_config=config,
+                    prefix=maybe_prefix(prefix, "language_model"),
+                    architectures=["DeepseekV4ForCausalLM"],
+                )
+            finally:
+                del config._dsv4_vl_inner  # type: ignore[attr-defined]
+        # The outer mapper (see load_weights) fully resolves HF names into
+        # this wrapper's namespace before AutoWeightsLoader strips the
+        # "language_model." prefix and delegates to the child's load_weights,
+        # so the child's own mapper must be a no-op. Its suffix rules are not
+        # idempotent (e.g. "lm_head.weight".endswith("head.weight") would
+        # re-fire "head.weight" -> "lm_head.weight").
+        text_mapper = self.language_model.hf_to_vllm_mapper
+        self.language_model.hf_to_vllm_mapper = WeightsMapper()
+        self.make_empty_intermediate_tensors = (  # type: ignore[method-assign]
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+        self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper(
+            text_mapper, image_enabled
+        )
+        self._weights_finalized = False
+
+    def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None:
+        patches = kwargs.pop("patches", None)
+        if patches is None:
+            return None
+        vit_grid = kwargs.pop("vit_grid", None)
+        llm_grid = kwargs.pop("llm_grid", None)
+        perm = kwargs.pop("perm", None)
+        assert vit_grid is not None and llm_grid is not None and perm is not None
+        return {
+            "patches": patches,
+            "vit_grid": vit_grid,
+            "llm_grid": llm_grid,
+            "perm": perm,
+        }
+
+    def _encode_image(
+        self,
+        patches: torch.Tensor,
+        n_vit_h: int,
+        n_vit_w: int,
+        perm: torch.Tensor,
+    ) -> torch.Tensor:
+        assert self.vision is not None and self.aligner is not None
+        image_embeds = self.aligner(
+            self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w
+        )
+        # Reorder into the N-layout block order used in the prompt.
+        return image_embeds[perm.to(image_embeds.device)]
+
+    def _process_image_input(
+        self,
+        patches: torch.Tensor,
+        vit_grid: torch.Tensor,
+        llm_grid: torch.Tensor,
+        perm: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        assert self.vision is not None and self.aligner is not None
+        patches = patches.to(self.aligner.w1.weight.dtype)
+
+        embeds: list[torch.Tensor] = []
+        vit_offset = 0
+        llm_offset = 0
+        for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip(
+            vit_grid.tolist(), llm_grid.tolist(), strict=True
+        ):
+            n_vit = n_vit_h * n_vit_w
+            n_llm = n_llm_h * n_llm_w
+            embeds.append(
+                self._encode_image(
+                    patches[vit_offset : vit_offset + n_vit],
+                    n_vit_h,
+                    n_vit_w,
+                    perm[llm_offset : llm_offset + n_llm],
+                )
+            )
+            vit_offset += n_vit
+            llm_offset += n_llm
+        return tuple(embeds)
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        image_input = self._parse_and_validate_image_input(**kwargs)
+        if image_input is None or self.vision is None:
+            return []
+        return self._process_image_input(
+            image_input["patches"],
+            image_input["vit_grid"],
+            image_input["llm_grid"],
+            image_input["perm"],
+        )
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.model_executor.models.utils import _merge_multimodal_embeddings
+
+        # All ids are in-vocab here: image-block sentinels are borrowed
+        # reserved tokens (their embedding rows are always overwritten below).
+        inputs_embeds = self.language_model.embed_input_ids(input_ids)
+
+        if self.image_start is not None:
+            # Branch-free sentinel overwrite: safe inside compiled/captured
+            # regions (no data-dependent control flow).
+            sentinel_mask = image_sentinel_mask(input_ids)
+            if is_multimodal is not None:
+                # IMAGE positions get vision embeddings via the merge below.
+                sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device)
+            table = torch.stack(
+                [
+                    self.image_start,
+                    self.image_pad,
+                    self.image_pad,
+                    self.image_newline,
+                    self.image_end,
+                ]
+            ).to(inputs_embeds.dtype)
+            idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4)
+            inputs_embeds = torch.where(
+                sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds
+            )
+
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        assert is_multimodal is not None
+        return _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors=None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self.language_model(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        return self.language_model.compute_logits(hidden_states)
+
+    def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.language_model.compute_logits_local(hidden_states)
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return self.language_model.get_expert_mapping()
+
+    def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
+        """Pre-hc_head residual stream buffer for the MTP/DSpark draft model."""
+        return self.language_model.get_mtp_target_hidden_states()
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        child_finalizes = getattr(
+            self.language_model, "finalizes_weights_during_load", True
+        )
+        mapped = self.hf_to_vllm_mapper.apply(weights)
+        if child_finalizes:
+            # A child which finalizes inside load_weights must see all of its
+            # weights in one contiguous delegation from AutoWeightsLoader.
+            mapped = iter(sorted(mapped, key=lambda x: x[0]))
+        loader = AutoWeightsLoader(self)
+        loaded_params = loader.load_weights(mapped)
+        self._weights_finalized = child_finalizes
+        return loaded_params
+
+    def process_weights_after_loading(self) -> None:
+        # Backbones such as the ROCm implementation require this to run only
+        # after the loader's generic per-layer quantization finalization.
+        if getattr(self, "_weights_finalized", False):
+            return
+        self.language_model.process_weights_after_loading()
+        self._weights_finalized = True
+
+    def get_mm_mapping(self) -> MultiModelKeys:
+        """Get the module prefixes in the multimodal model."""
+        return MultiModelKeys.from_string_field(
+            language_model="language_model",
+            connector="aligner",
+            tower_model="vision.",
+        )
+
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        if modality != "image":
+            raise ValueError(f"Unsupported modality: {modality!r}")
+
+        patches = mm_kwargs.get("patches") if mm_kwargs else None
+        if patches is not None and isinstance(patches.data, torch.Tensor):
+            tower_tokens = patches.data.shape[0]
+        else:
+            tower_tokens = num_mm_embeds * self.config.vision_downsample_ratio**2
+
+        return tower_tokens, num_mm_embeds
diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py
index 3bc636945a71..e6e0fd959e08 100644
--- a/vllm/models/deepseek_v4/compressor.py
+++ b/vllm/models/deepseek_v4/compressor.py
@@ -43,10 +43,12 @@ def _prefer_two_stage_compressor() -> bool:
 
 
 def _get_c128_boundary(metadata: CommonAttentionMetadata) -> bool | None:
-    starts = metadata._num_computed_tokens_cpu
-    if starts is None:
+    seq_lens_cpu = metadata.seq_lens_cpu_upper_bound
+    if seq_lens_cpu is None:
         return None
 
+    query_lens = metadata.query_start_loc_cpu[1:] - metadata.query_start_loc_cpu[:-1]
+    starts = seq_lens_cpu - query_lens
     starts_list = starts.tolist()
     query_start_loc = metadata.query_start_loc_cpu.tolist()
     return any(
@@ -64,7 +66,7 @@ def get_name() -> str:
         return "CompressorBackend"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
     @classmethod
@@ -116,6 +118,10 @@ def build(
             _, _, num_decode_tokens, _ = split_decodes_and_prefills(
                 common_attn_metadata, decode_threshold=1
             )
+        async_spec_decode = (
+            self.vllm_config.scheduler_config.async_scheduling
+            and self.vllm_config.speculative_config is not None
+        )
         return CompressorMetadata(
             block_table=common_attn_metadata.block_table_tensor.clamp_(min=0),
             slot_mapping=common_attn_metadata.slot_mapping,
@@ -124,7 +130,7 @@ def build(
             num_decode_tokens=num_decode_tokens,
             c128_boundary=(
                 _get_c128_boundary(common_attn_metadata)
-                if self.block_size == 8
+                if self.block_size == 8 and not async_spec_decode
                 else None
             ),
         )
@@ -197,8 +203,8 @@ class DeepseekCompressor(nn.Module):
     prologue (kv/score split, save_partial_states launch). The
     compress → norm → RoPE → store step is dispatched to a triton kernel
     (``compress_norm_rope_store_triton``) by default, except for the NVIDIA
-    head_dim=128 indexer path which uses the cutedsl kernel
-    (``compress_norm_rope_store_cutedsl``) for better performance.
+    head_dim=512 path which uses the CuTeDSL compressor kernels for better
+    performance.
     """
 
     def __init__(
@@ -311,7 +317,33 @@ def __init__(
                 head_dim=self.head_dim,
                 compress_ratio=self.compress_ratio,
             )
-            if self.head_dim != 512:
+            if current_platform.is_cuda() and self.head_dim == 512:
+                from vllm.models.deepseek_v4.nvidia.ops.sparse_attn_compress_cutedsl import (  # noqa: E501
+                    _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL,
+                    _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL,
+                    _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL,
+                    _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL,
+                    _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL,
+                )
+
+                store_full_kv = vllm_config.cache_config.cache_dtype != "fp8_ds_mla"
+                if self.compress_ratio == 4:
+                    (
+                        _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL
+                        if store_full_kv
+                        else _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL
+                    ).register_warmup()
+                else:
+                    _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL.register_warmup()
+                    if store_full_kv:
+                        _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL.register_warmup()
+                    else:
+                        _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL.register_warmup(
+                            vllm_config,
+                            k_cache_prefix=self.k_cache_prefix,
+                            compress_ratio=self.compress_ratio,
+                        )
+            else:
                 from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (  # noqa: E501
                     _FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL,
                 )
@@ -413,13 +445,13 @@ def forward(
         compress_norm_rope_store_fn: Any
         if current_platform.is_cuda() and self.head_dim == 512:
             from .nvidia.ops.sparse_attn_compress_cutedsl import (
-                compress_norm_rope_store_cutedsl,
+                _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL,
             )
 
             # head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla
             # layout and the plain full-cache layout. The full-cache flags
             # are consumed only here.
-            compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
+            compress_norm_rope_store_fn = _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL
             extra_kwargs: dict[str, Any] = dict(
                 store_full_kv=store_full_kv,
                 store_full_fp8=store_full_fp8,
diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py
index d51e589652f4..a2c3b00493e6 100644
--- a/vllm/models/deepseek_v4/cpu/cpu_sparse.py
+++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py
@@ -198,11 +198,12 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         qr: torch.Tensor,
-        compressed_kv_score: torch.Tensor,
+        compressed_kv_score: torch.Tensor | None,
         indexer_weights: torch.Tensor,
         positions: torch.Tensor,
         rotary_emb: nn.Module,
         qr_scale: torch.Tensor | None = None,
+        skip_compressor: bool = False,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
         """CPU override: no aux streams, so wq_b_and_q_quant and the
         compressor run straight-line instead of through
diff --git a/vllm/models/deepseek_v4/cpu/cpu_utils.py b/vllm/models/deepseek_v4/cpu/cpu_utils.py
index f03bf60b5846..9a5cff9abb89 100644
--- a/vllm/models/deepseek_v4/cpu/cpu_utils.py
+++ b/vllm/models/deepseek_v4/cpu/cpu_utils.py
@@ -23,6 +23,7 @@ def map_local_to_global_slots_cpu(
 
     Returns:
         [N, K] int64 global slot ids, -1 where ``local_indices`` was invalid.
+
     """
     valid = local_indices >= 0
     safe_local = local_indices.clamp(min=0).to(torch.int64)
diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py
index 1a3c2c780d28..987909702cf4 100644
--- a/vllm/models/deepseek_v4/nvidia/dspark.py
+++ b/vllm/models/deepseek_v4/nvidia/dspark.py
@@ -17,6 +17,7 @@
 
 import vllm.envs as envs
 from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config.kernel import MEGA_MOE_BACKENDS
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -31,7 +32,10 @@
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -48,12 +52,14 @@
     sp_padding_mask,
     sp_shard,
 )
+from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
 
 from .model import (
     DeepseekV4DecoderLayer,
     DeepseekV4Model,
     _use_sequence_parallel,
     make_deepseek_v4_expert_params_mapping,
+    prepare_mega_gate_routing_metadata,
 )
 
 logger = init_logger(__name__)
@@ -61,6 +67,24 @@
 # MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
 # fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
 _EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
+_CONTEXT_WKV_RE = re.compile(r"^mtp\.(\d+)\.attn\.wkv\.(.+)$")
+
+
+def _duplicate_context_wkv_weights(
+    weights: Iterable[tuple[str, torch.Tensor]], num_layers: int
+) -> Iterable[tuple[str, torch.Tensor]]:
+    """Load every draft layer's WKV into the cross-layer projection."""
+    for name, weight in weights:
+        yield name, weight
+        match = _CONTEXT_WKV_RE.fullmatch(name)
+        if match is None:
+            continue
+        layer_idx = int(match.group(1))
+        if layer_idx >= num_layers:
+            continue
+        stacked_weight = weight.detach()
+        stacked_weight.shard_id = layer_idx
+        yield f"context_wkv_proj.{match.group(2)}", stacked_weight
 
 
 class DSparkDeepseekV4Model(nn.Module):
@@ -75,6 +99,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.rms_norm_eps = config.rms_norm_eps
         self.num_hidden_layers = config.num_hidden_layers
         self.target_layer_ids = tuple(config.dspark_target_layer_ids)
+        self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
         self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
 
         self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3
@@ -113,6 +138,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
                 for i in range(self.num_dspark_layers)
             ]
         )
+        self.context_wkv_proj = MergedColumnParallelLinear(
+            config.hidden_size,
+            [config.head_dim] * self.num_dspark_layers,
+            bias=False,
+            return_bias=False,
+            quant_config=vllm_config.quant_config,
+            prefix=maybe_prefix(prefix, "context_wkv_proj"),
+            disable_tp=True,
+        )
 
         # Heads: final norm + hc_head, and the Markov + confidence heads
         # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
@@ -173,15 +207,16 @@ def precompute_and_store_context_kv(
         place draft layers in different groups). ``None`` (or a ``None`` entry)
         runs the projection to reserve workspace but writes nothing (profiling).
         """
-        for i, layer in enumerate(self.layers):
+        all_kv = self.context_wkv_proj(main_x).view(
+            main_x.shape[0], self.num_dspark_layers, self.config.head_dim
+        )
+        for i, (layer, kv) in enumerate(
+            zip(self.layers, all_kv.unbind(1), strict=True)
+        ):
             slot_mapping = (
                 None if context_slot_mappings is None else context_slot_mappings[i]
             )
             attn = layer.attn
-            # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
-            # (q_lora part discarded), then RoPE/quant/insert via the fused op.
-            qr_kv, _ = attn.fused_wqa_wkv(main_x)
-            kv = qr_kv[..., attn.q_lora_rank :]
             kv = attn.kv_norm(kv)
             if slot_mapping is None:
                 continue
@@ -204,6 +239,15 @@ def forward(
                 )
             inputs_embeds = sp_shard(inputs_embeds)
             input_ids = sp_shard(input_ids)
+        mega_gate_metadata = None
+        if self.use_mega_moe:
+            mega_gate_metadata = prepare_mega_gate_routing_metadata(
+                input_ids,
+                has_hash_routing=False,
+                image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
+                if getattr(self.config, "vision_n_layers", 0) > 0
+                else None,
+            )
         # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
         hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
 
@@ -216,6 +260,7 @@ def forward(
                 post_mix,
                 res_mix,
                 residual,
+                mega_gate_metadata=mega_gate_metadata,
             )
         hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
         if self.use_sequence_parallel:
@@ -429,6 +474,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         head_start = n_local_head * tp_rank
         head_end = n_local_head * (tp_rank + 1)
 
+        weights = _duplicate_context_wkv_weights(weights, len(self.model.layers))
         for name, loaded_weight in weights:
             mapped = self._remap_dspark_name(name)
             if mapped is None:
@@ -452,6 +498,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     self.quant_config, name, loaded_weight
                 )
 
+            if name.startswith("model.context_wkv_proj."):
+                param = params_dict[name]
+                param.weight_loader(param, loaded_weight, loaded_weight.shard_id)
+                loaded_params.add(name)
+                continue
+
             # E8M0 expert scales: keep raw exponent bytes.
             if ".experts." in name:
                 if (
@@ -522,6 +574,8 @@ def _remap_dspark_name(self, name: str) -> str | None:
 
         Returns None for non-mtp weights (owned by the target model).
         """
+        if name.startswith("context_wkv_proj."):
+            return f"model.{name}"
         m = re.match(r"mtp\.(\d+)\.(.*)", name)
         if m is None:
             return None
diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
index 883ef5ecafd4..bb9e42349a41 100644
--- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
@@ -111,7 +111,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [256]
 
     @staticmethod
@@ -521,6 +521,7 @@ def _forward(
                 swa_kv_cache=swa_k_cache,
                 workspace_buffer=workspace,
                 sparse_indices=sparse_indices[:num_decode_tokens],
+                sparse_indices_are_storage_offsets=True,
                 compressed_kv_cache=compressed_kv_cache,
                 sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
                 seq_lens=seq_lens[:num_decodes],
@@ -546,6 +547,7 @@ def _forward(
                 swa_kv_cache=swa_k_cache,
                 workspace_buffer=workspace,
                 sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
+                sparse_indices_are_storage_offsets=True,
                 compressed_kv_cache=compressed_kv_cache,
                 sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
                 seq_lens=seq_lens[num_decodes:num_reqs],
diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py
index 021001b9ce01..7edff8ae0202 100644
--- a/vllm/models/deepseek_v4/nvidia/model.py
+++ b/vllm/models/deepseek_v4/nvidia/model.py
@@ -98,6 +98,37 @@
 logger = init_logger(__name__)
 
 
+class MegaGateRoutingMetadata(typing.NamedTuple):
+    safe_hash_input_ids: torch.Tensor | None
+    image_token_mask: torch.Tensor | None
+    hash_token_mask: torch.Tensor | None
+
+
+def prepare_mega_gate_routing_metadata(
+    input_ids: torch.Tensor,
+    *,
+    has_hash_routing: bool,
+    image_sentinel_base_id: int | None,
+) -> MegaGateRoutingMetadata:
+    if image_sentinel_base_id is None:
+        return MegaGateRoutingMetadata(
+            input_ids if has_hash_routing else None,
+            None,
+            torch.ones_like(input_ids, dtype=torch.bool) if has_hash_routing else None,
+        )
+
+    image_token_mask = (input_ids >= image_sentinel_base_id) & (
+        input_ids < image_sentinel_base_id + 5
+    )
+    if not has_hash_routing:
+        return MegaGateRoutingMetadata(None, image_token_mask, None)
+    return MegaGateRoutingMetadata(
+        torch.where(image_token_mask, 0, input_ids),
+        image_token_mask,
+        ~image_token_mask,
+    )
+
+
 class DeepseekV4MLP(nn.Module):
     def __init__(
         self,
@@ -515,23 +546,6 @@ def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None:
 
         if self._transformed_l1_weights is None:
             self._check_runtime_supported()
-            # MegaMoE's 1x32 activation scales need 16-byte TMA rows, so
-            # pad each gate/up half and the down projection to a multiple of 512.
-            padded_size = (self.intermediate_size + 511) // 512 * 512
-            padding = padded_size - self.intermediate_size
-            if padding:
-                for param in (self.w13_weight, self.w13_weight_scale):
-                    gate_up = param.data.unflatten(1, (2, self.intermediate_size))
-                    param.data = torch.nn.functional.pad(
-                        gate_up, (0, 0, 0, padding)
-                    ).flatten(1, 2)
-                self.w2_weight.data = torch.nn.functional.pad(
-                    self.w2_weight.data, (0, padding // 2)
-                )
-                self.w2_weight_scale.data = torch.nn.functional.pad(
-                    self.w2_weight_scale.data, (0, padding // 32)
-                )
-                self.intermediate_size = padded_size
             w13_scale = deep_gemm.transform_sf_into_required_layout(
                 self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
                 2 * self.intermediate_size,
@@ -779,6 +793,7 @@ def __init__(
         use_sequence_parallel: bool = False,
         *,
         num_hash_layers: int,
+        reduce_results: bool = True,
         n_routed_experts: int | None = None,
         n_activated_experts: int | None = None,
         image_sentinel_lo: int = IMAGE_SENTINEL_BASE_ID,
@@ -790,6 +805,7 @@ def __init__(
         quant_config = vllm_config.quant_config
         self.prefix = prefix
         self.use_sequence_parallel = use_sequence_parallel
+        self.reduce_results = reduce_results
         moe_backend = vllm_config.kernel_config.moe_backend
         validate_fi_moe_ep_config(vllm_config)
         self.use_mega_moe = moe_backend in MEGA_MOE_BACKENDS
@@ -820,6 +836,10 @@ def __init__(
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
             )
+        if self.use_mega_moe and not self.renormalize:
+            raise NotImplementedError(
+                "DeepSeek V4 MegaMoE requires normalized top-k probabilities."
+            )
         if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
             raise NotImplementedError(
                 "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
@@ -975,25 +995,24 @@ def _init_fused_moe_experts(
         prefix: str,
     ) -> None:
         parallel_config = vllm_config.parallel_config
-        self.tp_rank = get_tensor_model_parallel_rank()
+        ep_group = get_ep_group()
+        ep_size = ep_group.world_size
+        ep_rank = ep_group.rank_in_group
 
         eplb_config = parallel_config.eplb_config
         self.n_redundant_experts = eplb_config.num_redundant_experts
         self.n_shared_experts = config.n_shared_experts or 0
         self.n_logical_experts = self.n_routed_experts
         self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
-        assert self.n_physical_experts % self.tp_size == 0, (
-            f"n_physical_experts={self.n_physical_experts} must be divisible by "
-            f"tp_size={self.tp_size}. Adjust num_redundant_experts."
-        )
-        self.n_local_physical_experts = self.n_physical_experts // self.tp_size
+        self.n_local_physical_experts = self.n_physical_experts // ep_size
         self.n_local_experts = self.n_local_physical_experts
-        self.experts_start_idx = self.tp_rank * self.n_local_experts
+        self.experts_start_idx = ep_rank * self.n_local_experts
         self.experts_end_idx = self.experts_start_idx + self.n_local_experts
         self.physical_expert_start = self.experts_start_idx
         self.physical_expert_end = self.experts_end_idx
 
         self.experts = FusedMoEFactory(
+            reduce_results=self.reduce_results,
             shared_experts=self.shared_experts,
             gate=self.gate,
             num_experts=self.n_routed_experts,
@@ -1017,7 +1036,10 @@ def _init_fused_moe_experts(
         )
 
     def forward(
-        self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor | None = None,
+        mega_gate_metadata: MegaGateRoutingMetadata | None = None,
     ) -> torch.Tensor:
         if self.gate.tid2eid is not None and input_ids is None:
             raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
@@ -1030,23 +1052,61 @@ def forward(
             return self._forward_fused_moe(hidden_states, input_ids)
 
         org_shape = hidden_states.shape
-        router_logits, _ = self.gate(hidden_states)
-        topk_weights, topk_ids = fused_topk_bias(
-            hidden_states=hidden_states,
-            gating_output=router_logits,
-            scoring_func=self.scoring_func,
-            e_score_correction_bias=self.gate.e_score_correction_bias.data
-            if self.gate.e_score_correction_bias is not None
-            else None,
-            topk=self.n_activated_experts,
-            renormalize=self.renormalize,
-            indices_type=self.hash_indices_dtype,
-            input_tokens=input_ids,
-            hash_indices_table=self.gate.tid2eid,
-            routed_scaling_factor=self.routed_scaling_factor,
-            bias_vl=bias_vl.data if bias_vl is not None else None,
-            image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0,
-        )
+        # Small local padded batches favor GateLinear; 128-expert gates cross earlier.
+        gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16
+        if hidden_states.shape[0] <= gate_threshold:
+            router_logits, _ = self.gate(hidden_states)
+            topk_weights, topk_ids = fused_topk_bias(
+                hidden_states=hidden_states,
+                gating_output=router_logits,
+                scoring_func=self.scoring_func,
+                e_score_correction_bias=self.gate.e_score_correction_bias.data
+                if self.gate.e_score_correction_bias is not None
+                else None,
+                topk=self.n_activated_experts,
+                renormalize=self.renormalize,
+                indices_type=self.hash_indices_dtype,
+                input_tokens=input_ids,
+                hash_indices_table=self.gate.tid2eid,
+                routed_scaling_factor=self.routed_scaling_factor,
+                bias_vl=bias_vl.data if bias_vl is not None else None,
+                image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0,
+            )
+        else:
+            from vllm.utils.deep_gemm import bf16_mega_gate
+
+            if mega_gate_metadata is None:
+                raise RuntimeError(
+                    "MegaMoE routing metadata must be prepared once by the model."
+                )
+
+            unmapped_topk_idx = None
+            fix_routing_mask = None
+            if self.gate.tid2eid is not None:
+                if mega_gate_metadata.safe_hash_input_ids is None:
+                    raise RuntimeError("Hash routing requires prepared input IDs.")
+                if mega_gate_metadata.hash_token_mask is None:
+                    raise RuntimeError("Hash routing requires a prepared routing mask.")
+                unmapped_topk_idx = self.gate.tid2eid[
+                    mega_gate_metadata.safe_hash_input_ids
+                ]
+                fix_routing_mask = mega_gate_metadata.hash_token_mask
+
+            topk_weights, topk_ids = bf16_mega_gate(
+                hidden_states,
+                self.gate.weight,
+                self.n_activated_experts,
+                scoring_func=self.scoring_func,
+                routed_scaling_factor=self.routed_scaling_factor,
+                ep_rank=self.ep_rank,
+                bias=self.gate.e_score_correction_bias.data
+                if self.gate.e_score_correction_bias is not None
+                else None,
+                image_bias=bias_vl.data if bias_vl is not None else None,
+                image_token_mask=mega_gate_metadata.image_token_mask,
+                fix_routing_mask=fix_routing_mask,
+                unmapped_topk_idx=unmapped_topk_idx,
+            )
         activation_clamp = (
             float(self.swiglu_limit) if self.swiglu_limit is not None else None
         )
@@ -1210,6 +1270,53 @@ def __init__(
             requires_grad=False,
         )
 
+        if vllm_config.kernel_config.enable_jit_warmup:
+            from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+                _HC_PRENORM_GEMM_TILELANG_KERNEL,
+                _MHC_FUSED_TILELANG_KERNEL,
+                _MHC_POST_TILELANG_KERNEL,
+                _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
+            )
+            from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+            include_pre_gemm_splits = is_deep_gemm_supported()
+            _MHC_PRE_BIG_FUSE_TILELANG_KERNEL.register_warmup(
+                vllm_config,
+                hidden_size=self.hidden_size,
+                hc_mult=self.hc_mult,
+                use_norm_weight=True,
+                include_pre_gemm_splits=include_pre_gemm_splits,
+                include_broadcast_splits=(
+                    get_pp_group().is_first_rank and extract_layer_index(prefix) == 0
+                ),
+                rms_eps=self.rms_norm_eps,
+                hc_pre_eps=self.hc_eps,
+                hc_sinkhorn_eps=self.hc_eps,
+                hc_post_mult_value=self.hc_post_alpha,
+                sinkhorn_repeat=self.hc_sinkhorn_iters,
+                norm_eps=(
+                    self.attn_norm.variance_epsilon,
+                    self.ffn_norm.variance_epsilon,
+                ),
+                broadcast_norm_eps=self.attn_norm.variance_epsilon,
+            )
+            if not include_pre_gemm_splits:
+                _HC_PRENORM_GEMM_TILELANG_KERNEL.register_warmup(
+                    vllm_config,
+                    hidden_size=self.hidden_size,
+                    hc_mult=self.hc_mult,
+                    n_out=self.hc_mult * (2 + self.hc_mult),
+                )
+            _MHC_POST_TILELANG_KERNEL.register_warmup(
+                hidden_size=self.hidden_size,
+                hc_mult=self.hc_mult,
+            )
+            _MHC_FUSED_TILELANG_KERNEL.register_warmup(
+                vllm_config,
+                hidden_size=self.hidden_size,
+                hc_mult=self.hc_mult,
+            )
+
     def forward(
         self,
         x: torch.Tensor,
@@ -1218,6 +1325,7 @@ def forward(
         post_mix: torch.Tensor | None = None,
         res_mix: torch.Tensor | None = None,
         residual: torch.Tensor | None = None,
+        mega_gate_metadata: MegaGateRoutingMetadata | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         attn_norm_weight = self.attn_norm.weight.data
         attn_norm_eps = self.attn_norm.variance_epsilon
@@ -1302,7 +1410,7 @@ def forward(
             norm_eps=ffn_norm_eps,
         )
 
-        x = self.ffn(x, input_ids)
+        x = self.ffn(x, input_ids, mega_gate_metadata)
         return x, residual, post_mix, res_mix
 
 
@@ -1364,6 +1472,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             ),
             prefix=f"{prefix}.layers",
         )
+        self.has_local_hash_moe = self.start_layer < min(
+            self.end_layer, config.num_hash_layers
+        )
 
         if get_pp_group().is_last_rank:
             self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
@@ -1402,6 +1513,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         else:
             self._mtp_hidden_buffer = None
 
+        if vllm_config.kernel_config.enable_jit_warmup and get_pp_group().is_last_rank:
+            from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+                _HC_HEAD_FUSED_TILELANG_KERNEL,
+            )
+
+            _HC_HEAD_FUSED_TILELANG_KERNEL.register_warmup(
+                hidden_size=config.hidden_size,
+                hc_mult=self.hc_mult,
+                rms_eps=self.rms_norm_eps,
+                hc_eps=self.hc_eps,
+            )
+
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
 
@@ -1454,6 +1577,16 @@ def forward(
             hidden_states = sp_shard(hidden_states)
             input_ids = sp_shard(input_ids)
 
+        mega_gate_metadata = None
+        if self.use_mega_moe:
+            mega_gate_metadata = prepare_mega_gate_routing_metadata(
+                input_ids,
+                has_hash_routing=self.has_local_hash_moe,
+                image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
+                if getattr(self.config, "vision_n_layers", 0) > 0
+                else None,
+            )
+
         residual, post_mix, res_mix = None, None, None
         remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
         aux_hidden_states: list[torch.Tensor] = []
@@ -1469,6 +1602,7 @@ def forward(
                 post_mix,
                 res_mix,
                 residual,
+                mega_gate_metadata=mega_gate_metadata,
             )
             if idx + 1 in self.aux_hidden_state_layers:
                 # Reconstruct the aux hidden state for draft models
@@ -1790,6 +1924,7 @@ class DeepseekV4ForCausalLM(
     SupportsLoRA,
     DeepseekV4MixtureOfExperts,
 ):
+    finalizes_weights_during_load = True
     model_cls = DeepseekV4Model
 
     # Default mapper assumes the original FP4-expert checkpoint layout.
diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py
index 55403559dc65..50886b0afab7 100644
--- a/vllm/models/deepseek_v4/nvidia/mtp.py
+++ b/vllm/models/deepseek_v4/nvidia/mtp.py
@@ -20,6 +20,7 @@
 
 import vllm.envs as envs
 from vllm.config import VllmConfig
+from vllm.config.kernel import MEGA_MOE_BACKENDS
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -51,6 +52,7 @@
     sp_padding_mask,
     sp_shard,
 )
+from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
 from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import (
     _FUSED_MTP_INPUT_RMSNORM_KERNEL,
     _MTP_SHARED_HEAD_RMSNORM_KERNEL,
@@ -62,6 +64,7 @@
     DeepseekV4Model,
     _use_sequence_parallel,
     make_deepseek_v4_expert_params_mapping,
+    prepare_mega_gate_routing_metadata,
 )
 
 logger = init_logger(__name__)
@@ -92,6 +95,7 @@ def __init__(
         self.config = config
         quant_config = vllm_config.quant_config
         self.rms_norm_eps = config.rms_norm_eps
+        self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
 
         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -181,8 +185,23 @@ def forward(
         hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
             inputs_embeds
         ).unsqueeze(-2)
+        mega_gate_metadata = None
+        if self.use_mega_moe:
+            routing_input_ids = input_ids
+            if self.mtp_block.use_sequence_parallel:
+                routing_input_ids = sp_shard(routing_input_ids)
+            mega_gate_metadata = prepare_mega_gate_routing_metadata(
+                routing_input_ids,
+                has_hash_routing=False,
+                image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
+                if getattr(self.config, "vision_n_layers", 0) > 0
+                else None,
+            )
         hidden_states, residual, post_mix, res_mix = self.mtp_block(
-            positions=positions, x=hidden_states, input_ids=input_ids
+            positions=positions,
+            x=hidden_states,
+            input_ids=input_ids,
+            mega_gate_metadata=mega_gate_metadata,
         )
         hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
         if self.mtp_block.use_sequence_parallel:
@@ -518,8 +537,7 @@ def process_weights_after_loading(self) -> None:
         self.finalize_mega_moe_weights()
 
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
diff --git a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
index a4514293c942..5920ed290cad 100644
--- a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
+++ b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from functools import cache
+from dataclasses import dataclass
+from typing import Any
 
 import cutlass
 import cutlass.cute as cute
@@ -12,303 +13,353 @@
 from quack.compile_utils import make_fake_tensor
 
 from vllm.cute_utils import _bf16x2_mul, cvt
-
-
-def dequantize_and_gather_k_cache_cutedsl(
-    out: torch.Tensor,
-    k_cache: torch.Tensor,
-    seq_lens: torch.Tensor,
-    gather_lens: torch.Tensor | None,
-    block_table: torch.Tensor,
-    block_size: int,
-    offset: int,
-) -> None:
-    DequantGatherKCacheKernel.compile(
-        block_size=block_size,
-        has_gather_lens=gather_lens is not None,
-    )(out, k_cache, seq_lens, gather_lens, block_table, offset)
-
-
-class DequantGatherKCacheKernel:
-    # Hard-coded for DSv4.
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
+from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
+    CuTeDSLLaunchSpec,
+    VllmCuTeDSLJitKernel,
+)
+
+
+class DequantGatherKCacheKernel(
+    VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"]
+):
+    # DSv4 geometry. The record format (how much of the row is quantized and
+    # how wide a scale group is) comes from the compile key: V4 quantizes the
+    # 448 NoPE dims in groups of 64 and keeps a bf16 RoPE tail, V4.1 quantizes
+    # all 512 dims in groups of 32 (MXFP8) and has no tail.
     head_dim = 512
-    group_size = 64  # 1 scale per 64 elems
-
-    def __init__(self, fp8_dim: int = 448, block_size: int = 64):
-        self.fp8_dim = fp8_dim
-        self.bf16_dim = self.head_dim - fp8_dim
-        self.data_dim = fp8_dim + self.bf16_dim * 2
-        self.block_size = block_size
-
-        self.num_warps = 4
-        self.tb_size = self.num_warps * 32
-        self.num_stages = 4
-
-    @cute.jit
-    def __call__(
-        self,
-        out: cute.Tensor,
-        k_cache: cute.Tensor,
-        seq_lens: cute.Tensor,
-        gather_lens: cute.Tensor | None,
-        block_table: cute.Tensor,
-        offset: Int32,
-        stream: CUstream,
-    ):
-        # Split k_cache into k_data and k_scale. Each [block_size, head_bytes]
-        # block is actually a concat of
-        # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8].
-        k_data = cute.make_tensor(
-            k_cache.iterator,
-            layout=cute.make_layout(
-                (k_cache.shape[0], self.block_size, self.data_dim),
-                stride=(k_cache.stride[0], self.data_dim, 1),
-            ),
-        )
-        k_scale = cute.make_tensor(
-            k_cache.iterator + (self.block_size * self.data_dim),
-            layout=cute.make_layout(
-                (k_cache.shape[0], self.block_size, 8),
-                stride=(k_cache.stride[0], 8, 1),
-            ),
-        )
+    num_warps = 4
+    tb_size = num_warps * 32
+    num_stages = 4
 
-        num_reqs = out.shape[0]
-        max_rows = out.shape[1] - offset
-        num_workers = cutlass.max(
-            1, cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs))
-        )
-        grid = (num_reqs, num_workers, 1)
-        self.kernel(
-            out,
-            k_data,
-            k_scale,
-            seq_lens,
-            gather_lens,
-            block_table,
-            offset,
-        ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
+    @dataclass(frozen=True)
+    class CompileKey:
+        block_size: int
+        has_gather_lens: bool
+        fp8_dim: int = 448
+        group_size: int = 64  # dims per scale
 
-    @cute.jit
-    def load_g2s(
-        self,
-        k_data_slice: cute.Tensor,
-        k_scale: cute.Tensor,
-        block_table: cute.Tensor,
-        s_kdata_slice: cute.Tensor,
-        s_kscale: cute.Tensor,
-        req_id,
-        pos,
-        lane_id,
-        stage_id,
-    ):
-        # k_data_slice: [num_blocks, block_size, (16, data_dim/16)]
-        # s_kdata_slice: [(4, data_dim/16), num_stages]
-
-        op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL)
-        cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128)
-        cp8_atom = cute.make_copy_atom(cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64)
-        page_id = block_table[req_id, pos // self.block_size]
-        block_offset = pos % self.block_size
-
-        # Load the first 512 bytes (32x16B).
-        idx = lane_id
-        src = k_data_slice[page_id, block_offset, (None, idx)]
-        cute.copy(
-            cp16_atom,
-            cute.recast_tensor(src, Uint32),
-            s_kdata_slice[(None, idx), stage_id],
-        )
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        head_dim = DequantGatherKCacheKernel.head_dim
+        fp8_dim = compile_key.fp8_dim
+        group_size = compile_key.group_size
+        num_warps = DequantGatherKCacheKernel.num_warps
+        tb_size = DequantGatherKCacheKernel.tb_size
+        num_stages = DequantGatherKCacheKernel.num_stages
+        bf16_dim = head_dim - fp8_dim
+        data_dim = fp8_dim + bf16_dim * 2
+        scale_bytes = head_dim // group_size
+        block_size = compile_key.block_size
+
+        @cute.jit
+        def load_g2s(
+            k_data_slice: cute.Tensor,
+            k_scale: cute.Tensor,
+            block_table: cute.Tensor,
+            s_kdata_slice: cute.Tensor,
+            s_kscale: cute.Tensor,
+            req_id,
+            pos,
+            lane_id,
+            stage_id,
+        ):
+            op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL)
+            cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128)
+            cp_scale_atom = cute.make_copy_atom(
+                cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=scale_bytes * 8
+            )
+            page_id = block_table[req_id, pos // block_size]
+            block_offset = pos % block_size
 
-        # Load the tail 64 bytes.
-        idx += 32
-        if idx < cutlass.const_expr(self.data_dim // 16):
+            # Load the first 512 bytes (32x16B).
+            idx = lane_id
             src = k_data_slice[page_id, block_offset, (None, idx)]
             cute.copy(
                 cp16_atom,
                 cute.recast_tensor(src, Uint32),
                 s_kdata_slice[(None, idx), stage_id],
             )
-        elif idx == cutlass.const_expr(self.data_dim // 16):
-            cute.copy(
-                cp8_atom,
-                k_scale[page_id, block_offset, None],
-                s_kscale[None, stage_id],
-            )
 
-    @cute.kernel
-    def kernel(
-        self,
-        out: cute.Tensor,
-        k_data: cute.Tensor,
-        k_scale: cute.Tensor,
-        seq_lens: cute.Tensor,
-        gather_lens: cute.Tensor | None,
-        block_table: cute.Tensor,
-        offset: Int32,
-    ):
-        req_id, worker_id, _ = cute.arch.block_idx()
-        tid, _, _ = cute.arch.thread_idx()
-        warp_id = cute.arch.make_warp_uniform(tid // 32)
-        lane_id = tid % 32
-
-        _, num_workers, _ = cute.arch.grid_dim()
-
-        # Prepare smem.
-        smem = cutlass.utils.SmemAllocator()
-        s_kdata = smem.allocate_tensor(
-            Uint32,
-            cute.make_layout((self.data_dim // 4, self.num_warps, self.num_stages)),
-            byte_alignment=16,
-        )[None, warp_id, None]
-        s_kscale = smem.allocate_tensor(
-            Uint8,
-            cute.make_layout((8, self.num_warps, self.num_stages)),
-            byte_alignment=8,
-        )[None, warp_id, None]
-
-        # Prepare for 16B cp.async, also for BF16 smem loads later.
-        k_data_slice = cute.logical_divide(k_data, (None, None, 16))
-        s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None))
-
-        # Load FP8 elems in 8B units, so once dequantized, they are 16B units.
-        s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None))
-
-        # 16B st.global.
-        out_slice = cute.logical_divide(out, (None, None, 8))
-
-        cp_op = cute.nvgpu.CopyUniversalOp()
-        cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
-        cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
-
-        seq_len = seq_lens[req_id]
-        gather_len = seq_len
-        if cutlass.const_expr(gather_lens is not None):
-            gather_len = gather_lens[req_id]  # type: ignore[index]
-        start_pos = seq_len - gather_len
-
-        # Start prefetch.
-        for i in cutlass.range_constexpr(self.num_stages - 1):
-            next_pos = (
-                start_pos
-                + worker_id * self.num_warps
-                + warp_id
-                + i * num_workers * self.num_warps
-            )
-            if next_pos < seq_len:
-                self.load_g2s(
-                    k_data_slice,
-                    k_scale,
-                    block_table,
-                    s_kdata_16B_slice,
-                    s_kscale,
-                    req_id,
-                    next_pos,
-                    lane_id,
-                    i,
+            # Load the tail 64 bytes.
+            idx += 32
+            if idx < cutlass.const_expr(data_dim // 16):
+                src = k_data_slice[page_id, block_offset, (None, idx)]
+                cute.copy(
+                    cp16_atom,
+                    cute.recast_tensor(src, Uint32),
+                    s_kdata_slice[(None, idx), stage_id],
                 )
-            cute.arch.cp_async_commit_group()
-        prefetch_stage = self.num_stages - 1
-        compute_stage = 0
-
-        # Main loop.
-        for i in range(
-            worker_id * self.num_warps + warp_id,
-            gather_len,
-            num_workers * self.num_warps,
-        ):
-            pos = start_pos + i
-
-            # Prefetch next stage.
-            next_pos = pos + num_workers * self.num_warps * (self.num_stages - 1)
-            if next_pos < seq_len:
-                self.load_g2s(
-                    k_data_slice,
-                    k_scale,
-                    block_table,
-                    s_kdata_16B_slice,
-                    s_kscale,
-                    req_id,
-                    next_pos,
-                    lane_id,
-                    prefetch_stage,
+            elif idx == cutlass.const_expr(data_dim // 16):
+                cute.copy(
+                    cp_scale_atom,
+                    k_scale[page_id, block_offset, None],
+                    s_kscale[None, stage_id],
                 )
-                prefetch_stage = (prefetch_stage + 1) % self.num_stages
-            cute.arch.cp_async_commit_group()
-
-            # Wait for gmem->smem to finish.
-            cute.arch.cp_async_wait_group(self.num_stages - 1)
-            cute.arch.sync_warp()
-
-            # There are 512 elems per token. As a warp, data0 holds the first
-            # 256 elems and data1 holds the second 256 elems, i.e. each thread
-            # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as
-            # contiguous 16B global stores. On Blackwell, this might not be
-            # necessary as we have 32B global stores, but doing it this way
-            # does not seem to be slower.
-            data0 = cute.make_rmem_tensor((2,), Uint32)
-            data1 = cute.make_rmem_tensor((2,), Uint32)
-            cute.copy(cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0)
-            cute.copy(
-                cp8_atom,
-                s_kdata_8B_slice[(None, lane_id + 32), compute_stage],
-                data1,
-            )
 
-            # Convert to bf16x2 via bit manipulation. FP8 scales are per 64
-            # elements. An 8-element chunk advances the scale index by
-            # chunk_id * 8 // group_size.
-            scale0_u32 = Uint32(s_kscale[lane_id * 8 // self.group_size, compute_stage])
-            scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7))
-            scale1_u32 = Uint32(
-                s_kscale[(lane_id + 32) * 8 // self.group_size, compute_stage]
-            )
-            scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7))
-
-            # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2
-            # (CUDA 13.2).
-            dequant0 = cute.make_rmem_tensor(4, Uint32)
-            dequant1 = cute.make_rmem_tensor(4, Uint32)
-            for j in cutlass.range_constexpr(2):
-                tmp0 = cvt.fp8x4_to_bf16x4(data0[j])
-                tmp1 = cvt.fp8x4_to_bf16x4(data1[j])
-
-                # BF16 multiply is safe because the scales are exact powers of 2.
-                dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2)
-                dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2)
-                dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2)
-                dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2)
-
-            # Last 64 elems are BF16 tail, corresponds to dequant1 of last
-            # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 +
-            # 8x 16B for BF16.
-            if lane_id + 32 >= self.fp8_dim // 8:
-                idx = self.fp8_dim // 16 + (lane_id + 32) - self.fp8_dim // 8
+        @cute.kernel
+        def device_kernel(
+            out: cute.Tensor,
+            k_data: cute.Tensor,
+            k_scale: cute.Tensor,
+            seq_lens: cute.Tensor,
+            gather_lens: cute.Tensor | None,
+            block_table: cute.Tensor,
+            offset: Int32,
+        ):
+            req_id, worker_id, _ = cute.arch.block_idx()
+            tid, _, _ = cute.arch.thread_idx()
+            warp_id = cute.arch.make_warp_uniform(tid // 32)
+            lane_id = tid % 32
+
+            _, num_workers, _ = cute.arch.grid_dim()
+
+            # Prepare smem.
+            smem = cutlass.utils.SmemAllocator()
+            s_kdata = smem.allocate_tensor(
+                Uint32,
+                cute.make_layout((data_dim // 4, num_warps, num_stages)),
+                byte_alignment=16,
+            )[None, warp_id, None]
+            s_kscale = smem.allocate_tensor(
+                Uint8,
+                cute.make_layout((scale_bytes, num_warps, num_stages)),
+                byte_alignment=scale_bytes,
+            )[None, warp_id, None]
+
+            # Prepare for 16B cp.async, also for BF16 smem loads later.
+            k_data_slice = cute.logical_divide(k_data, (None, None, 16))
+            s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None))
+
+            # Load FP8 elems in 8B units, so once dequantized, they are 16B units.
+            s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None))
+
+            # 16B st.global.
+            out_slice = cute.logical_divide(out, (None, None, 8))
+
+            cp_op = cute.nvgpu.CopyUniversalOp()
+            cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
+            cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
+
+            seq_len = seq_lens[req_id]
+            gather_len = seq_len
+            if cutlass.const_expr(gather_lens is not None):
+                gather_len = gather_lens[req_id]  # type: ignore[index]
+            start_pos = seq_len - gather_len
+
+            # Start prefetch.
+            for i in cutlass.range_constexpr(num_stages - 1):
+                next_pos = (
+                    start_pos
+                    + worker_id * num_warps
+                    + warp_id
+                    + i * num_workers * num_warps
+                )
+                if next_pos < seq_len:
+                    load_g2s(
+                        k_data_slice,
+                        k_scale,
+                        block_table,
+                        s_kdata_16B_slice,
+                        s_kscale,
+                        req_id,
+                        next_pos,
+                        lane_id,
+                        i,
+                    )
+                cute.arch.cp_async_commit_group()
+            prefetch_stage = num_stages - 1
+            compute_stage = 0
+
+            # Main loop.
+            for i in range(
+                worker_id * num_warps + warp_id,
+                gather_len,
+                num_workers * num_warps,
+            ):
+                pos = start_pos + i
+
+                # Prefetch next stage.
+                next_pos = pos + num_workers * num_warps * (num_stages - 1)
+                if next_pos < seq_len:
+                    load_g2s(
+                        k_data_slice,
+                        k_scale,
+                        block_table,
+                        s_kdata_16B_slice,
+                        s_kscale,
+                        req_id,
+                        next_pos,
+                        lane_id,
+                        prefetch_stage,
+                    )
+                    prefetch_stage = (prefetch_stage + 1) % num_stages
+                cute.arch.cp_async_commit_group()
+
+                # Wait for gmem->smem to finish.
+                cute.arch.cp_async_wait_group(num_stages - 1)
+                cute.arch.sync_warp()
+
+                # There are 512 elems per token. As a warp, data0 holds the first
+                # 256 elems and data1 holds the second 256 elems, i.e. each thread
+                # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as
+                # contiguous 16B global stores. On Blackwell, this might not be
+                # necessary as we have 32B global stores, but doing it this way
+                # does not seem to be slower.
+                data0 = cute.make_rmem_tensor((2,), Uint32)
+                data1 = cute.make_rmem_tensor((2,), Uint32)
                 cute.copy(
-                    cp16_atom,
-                    s_kdata_16B_slice[(None, idx), compute_stage],
-                    dequant1,
+                    cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0
+                )
+                cute.copy(
+                    cp8_atom,
+                    s_kdata_8B_slice[(None, lane_id + 32), compute_stage],
+                    data1,
                 )
 
-            # Store two 16B BF16 chunks per lane: first half, then second half.
-            dst = out_slice[req_id, offset + i, (None, lane_id)]
-            cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32))
+                # Convert to bf16x2 via bit manipulation. FP8 scales are per 64
+                # elements. An 8-element chunk advances the scale index by
+                # chunk_id * 8 // group_size.
+                scale0_u32 = Uint32(s_kscale[lane_id * 8 // group_size, compute_stage])
+                scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7))
+                scale1_u32 = Uint32(
+                    s_kscale[(lane_id + 32) * 8 // group_size, compute_stage]
+                )
+                scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7))
+
+                # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2
+                # (CUDA 13.2).
+                dequant0 = cute.make_rmem_tensor(4, Uint32)
+                dequant1 = cute.make_rmem_tensor(4, Uint32)
+                for j in cutlass.range_constexpr(2):
+                    tmp0 = cvt.fp8x4_to_bf16x4(data0[j])
+                    tmp1 = cvt.fp8x4_to_bf16x4(data1[j])
+
+                    # BF16 multiply is safe because the scales are exact powers of 2.
+                    dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2)
+                    dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2)
+                    dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2)
+                    dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2)
+
+                # Last 64 elems are BF16 tail, corresponds to dequant1 of last
+                # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 +
+                # 8x 16B for BF16. The V4.1 record is fully quantized and has
+                # no tail.
+                if cutlass.const_expr(bf16_dim > 0):  # noqa: SIM102
+                    if lane_id + 32 >= fp8_dim // 8:
+                        idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8
+                        cute.copy(
+                            cp16_atom,
+                            s_kdata_16B_slice[(None, idx), compute_stage],
+                            dequant1,
+                        )
+
+                # Store two 16B BF16 chunks per lane: first half, then second half.
+                dst = out_slice[req_id, offset + i, (None, lane_id)]
+                cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32))
+
+                dst = out_slice[req_id, offset + i, (None, lane_id + 32)]
+                cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32))
+
+                compute_stage = (compute_stage + 1) % num_stages
+
+        @cute.jit
+        def host_entrypoint(
+            out: cute.Tensor,
+            k_cache: cute.Tensor,
+            seq_lens: cute.Tensor,
+            gather_lens: cute.Tensor | None,
+            block_table: cute.Tensor,
+            offset: Int32,
+            stream: CUstream,
+        ):
+            # Split k_cache into k_data and k_scale. Each [block_size, head_bytes]
+            # block is actually a concat of
+            # [block_size, fp8_dim + bf16_dim * 2] and [block_size, scale_bytes].
+            k_data = cute.make_tensor(
+                k_cache.iterator,
+                layout=cute.make_layout(
+                    (k_cache.shape[0], block_size, data_dim),
+                    stride=(k_cache.stride[0], data_dim, 1),
+                ),
+            )
+            # k_data_slice: [num_blocks, block_size, (16, data_dim/16)]
+            # s_kdata_slice: [(4, data_dim/16), num_stages]
+            k_scale = cute.make_tensor(
+                k_cache.iterator + (block_size * data_dim),
+                layout=cute.make_layout(
+                    (k_cache.shape[0], block_size, scale_bytes),
+                    stride=(k_cache.stride[0], scale_bytes, 1),
+                ),
+            )
 
-            dst = out_slice[req_id, offset + i, (None, lane_id + 32)]
-            cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32))
+            num_reqs = out.shape[0]
+            max_rows = out.shape[1] - offset
+            num_workers = cutlass.max(
+                1,
+                cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs)),
+            )
+            grid = (num_reqs, num_workers, 1)
+            device_kernel(
+                out,
+                k_data,
+                k_scale,
+                seq_lens,
+                gather_lens,
+                block_table,
+                offset,
+            ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
+
+        return host_entrypoint
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        block_size: int,
+        has_gather_lens: bool,
+        fp8_dim: int = 448,
+        group_size: int = 64,
+    ) -> CompileKey:
+        return self.CompileKey(
+            block_size=block_size,
+            has_gather_lens=has_gather_lens,
+            fp8_dim=fp8_dim,
+            group_size=group_size,
+        )
 
-            compute_stage = (compute_stage + 1) % self.num_stages
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        block_size = vllm_config.cache_config.block_size
+        compress_ratios = vllm_config.model_config.hf_config.compress_ratios
+        return self._trace_dispatch(self.dispatch)(
+            zip_inputs(
+                dict(
+                    block_size=block_size,
+                    has_gather_lens=True,
+                    enabled=True,
+                ),
+                dict(
+                    block_size=max(1, block_size // 4),
+                    has_gather_lens=False,
+                    enabled=4 in compress_ratios,
+                ),
+                dict(
+                    block_size=max(1, block_size // 128),
+                    has_gather_lens=False,
+                    enabled=128 in compress_ratios,
+                ),
+            ),
+            valid_block_size=block_size > 0,
+            _when=lambda *, enabled, valid_block_size: enabled and valid_block_size,
+        )
 
-    @cache
-    @staticmethod
-    def compile(
-        fp8_dim: int = 448,
-        block_size: int = 64,
-        has_gather_lens: bool = True,
-    ):
+    def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
         num_reqs = cute.sym_int()
-        head_dim = DequantGatherKCacheKernel.head_dim
-        head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8
+        head_dim = self.head_dim
+        fp8_dim = compile_key.fp8_dim
+        block_size = compile_key.block_size
+        head_bytes = (
+            fp8_dim + (head_dim - fp8_dim) * 2 + head_dim // compile_key.group_size
+        )
 
         out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16)
         k_cache = cute.runtime.make_fake_tensor(
@@ -318,19 +369,51 @@ def compile(
             assumed_align=32,
         )
         seq_lens = make_fake_tensor(Int32, (num_reqs,))
-        gather_lens = make_fake_tensor(Int32, (num_reqs,)) if has_gather_lens else None
+        gather_lens = (
+            make_fake_tensor(Int32, (num_reqs,))
+            if compile_key.has_gather_lens
+            else None
+        )
         block_table = make_fake_tensor(Int32, (num_reqs, cute.sym_int()))
 
-        kernel = DequantGatherKCacheKernel(fp8_dim, block_size)
-        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-        return cute.compile(
-            kernel,
+        return (
             out,
             k_cache,
             seq_lens,
             gather_lens,
             block_table,
             Int32(0),
-            stream,
-            options="--enable-tvm-ffi",
         )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        *,
+        out: torch.Tensor,
+        k_cache: torch.Tensor,
+        seq_lens: torch.Tensor,
+        gather_lens: torch.Tensor | None,
+        block_table: torch.Tensor,
+        block_size: int,
+        offset: int,
+        fp8_dim: int = 448,
+        group_size: int = 64,
+    ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]:
+        compile_key = self.dispatch(
+            block_size=block_size,
+            has_gather_lens=gather_lens is not None,
+            fp8_dim=fp8_dim,
+            group_size=group_size,
+        )
+        launch_args = (
+            out,
+            k_cache,
+            seq_lens,
+            gather_lens,
+            block_table,
+            offset,
+        )
+        return compile_key, launch_args
+
+
+_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL = DequantGatherKCacheKernel()
diff --git a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
index 16e98abf23e1..f67a77e2541f 100644
--- a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
+++ b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from functools import cache
+# ruff: noqa: E501
+
+from dataclasses import dataclass
+from typing import Any
 
 import cutlass
 import cutlass.cute as cute
@@ -14,594 +17,651 @@
     _bf16x2_max,
     cvt,
     recast_val,
+    torch_to_cute_dtype,
+)
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
+from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
+    CuTeDSLLaunchSpec,
+    VllmCuTeDSLJitKernel,
 )
 from vllm.vllm_flash_attn.cute import utils as cute_utils
 
 # MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale.
 MXFP4_BLOCK_SIZE = 32
 
-_TORCH_TO_CUTE = {
-    torch.bfloat16: BFloat16,
-    torch.float32: Float32,
-}
-
-
-def fused_indexer_q_rope_quant_mxfp4_cutedsl(
-    positions: torch.Tensor,
-    index_q: torch.Tensor,
-    index_q_cos_sin_cache: torch.Tensor,
-    index_weights: torch.Tensor,
-    index_weights_softmax_scale: float,
-    index_weights_head_scale: float,
-    index_q_packed: torch.Tensor,
-    index_q_scale: torch.Tensor,
-    index_weights_out: torch.Tensor,
-) -> None:
-    num_tokens, num_heads, head_dim = index_q.shape
-    rope_dim = index_q_cos_sin_cache.shape[-1]
-    rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype]
-
-    # compile all variants at first invocation
-    for coarsen in (1, 4):
-        IndexerQMxFp4Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen)
-
-    # heuristic
-    coarsen = 1 if num_tokens < 512 else 4
-    compiled = IndexerQMxFp4Kernel.compile(
-        head_dim, rope_dim, num_heads, rope_type, coarsen
-    )
-    scale = float(index_weights_softmax_scale * index_weights_head_scale)
-    compiled(
-        positions,
-        index_q,
-        index_q_cos_sin_cache,
-        index_weights,
-        index_q_packed,
-        index_q_scale,
-        index_weights_out,
-        scale,
-    )
 
-
-def fused_indexer_q_rope_quant_fp8_cutedsl(
-    positions: torch.Tensor,
-    index_q: torch.Tensor,
-    index_q_cos_sin_cache: torch.Tensor,
-    index_weights: torch.Tensor,
-    index_weights_softmax_scale: float,
-    index_weights_head_scale: float,
-    index_q_fp8: torch.Tensor,
-    index_weights_out: torch.Tensor,
-) -> None:
-    num_tokens, num_heads, head_dim = index_q.shape
-    rope_dim = index_q_cos_sin_cache.shape[-1]
-    rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype]
-
-    for coarsen in (1, 4):
-        IndexerQFp8Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen)
-
-    coarsen = 1 if num_tokens < 512 else 4
-    compiled = IndexerQFp8Kernel.compile(
-        head_dim, rope_dim, num_heads, rope_type, coarsen
-    )
-    scale = float(index_weights_softmax_scale * index_weights_head_scale)
-    # The cute kernel treats the FP8 buffer as raw bytes (Uint8).
-    compiled(
-        positions,
-        index_q,
-        index_q_cos_sin_cache,
-        index_weights,
-        index_q_fp8.view(torch.uint8),
-        index_weights_out,
-        scale,
+@cute.jit
+def _load_q_and_rope(
+    positions: cute.Tensor,
+    q: cute.Tensor,
+    cos_sin_cache: cute.Tensor,
+    num_heads: cutlass.Constexpr,
+    coarsen: cutlass.Constexpr,
+    tb_size: cutlass.Constexpr,
+    subwarp_size: cutlass.Constexpr,
+    threads_per_token: cutlass.Constexpr,
+    nope_dim: cutlass.Constexpr,
+    rope_dim: cutlass.Constexpr,
+    cos_sin_dtype: cutlass.Constexpr,
+):
+    """Compute thread indices, load Q (BF16), and apply interleaved RoPE.
+
+    Returns a tuple
+        (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id,
+         head_start, in_bounds, num_token_heads)
+    where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed
+    bf16x2 pairs covering the 16 BF16 lanes owned by this thread for
+    each of ``coarsen`` heads. RoPE is applied in place to the
+    trailing ``rope_dim`` lanes; the leading nope lanes pass through.
+    """
+    block_id, _, _ = cute.arch.block_idx()
+    tid, _, _ = cute.arch.thread_idx()
+
+    num_tokens = q.shape[0]
+    num_token_heads = num_tokens * num_heads
+    global_tid = block_id * tb_size + tid
+
+    global_subwarp_id = global_tid // subwarp_size
+    sublane = tid % subwarp_size
+
+    token_id = global_subwarp_id // (num_heads // coarsen)
+    head_tile_id = global_subwarp_id % (num_heads // coarsen)
+    head_start = head_tile_id * coarsen
+
+    # NOTE: token_id may exceed bounds, hence we need to add load/store guards
+    # we can't do early exit because CuteDSL doesn't support it. and we also need
+    # all threads in a warp to be active since we utilize warp shuffle later.
+    # must_in_bounds is constexpr, True when a token spans a whole number of
+    # threadblocks. Then num_tokens * threads_per_token is always a multiple of
+    # tb_size, so the trailing block never straddles num_tokens and the compiler
+    # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads
+    # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half
+    # of the last block addressing token num_tokens.
+    must_in_bounds = cutlass.const_expr(threads_per_token % tb_size == 0)
+    in_bounds = must_in_bounds or (token_id < num_tokens)
+
+    cp_op = cute.nvgpu.CopyUniversalOp()
+
+    _layout = cute.make_layout((coarsen, 8), stride=(8, 1))
+    q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32)
+
+    if in_bounds:
+        # we can't do cute.copy() on the whole 2D tile directly because
+        # cute.copy() wants the 1st mode to be covered by the copy atom,
+        # and other modes as for loop. there is no fast way to
+        # "transpose" the tensor view.
+        q_tile = cute.local_tile(
+            q[token_id, None, None],
+            tiler=(coarsen, 16),
+            coord=(head_tile_id, sublane),
+        )
+        cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256)
+        for i in cutlass.range_constexpr(coarsen):
+            src = cute.recast_tensor(q_tile[i, None], Uint32)
+            cute.copy(cp_u32x8, src, q_bf16x2[i, None])
+
+    # RoPE applies only to the trailing rope_dim values. We keep the rounded
+    # BF16 result in q_bits so the later amax and quantization see BF16.
+    # cos_sin_cache layout: [max_pos, rope_dim]
+    if in_bounds and sublane * 16 >= nope_dim:
+        cos_vals = cute.make_rmem_tensor((8,), Float32)
+        sin_vals = cute.make_rmem_tensor((8,), Float32)
+
+        pos = positions[token_id]
+
+        # select 8 elems from cos and sin
+        cos_id = sublane - nope_dim // 16
+        sin_id = cos_id + rope_dim // 16
+        cos_src = cute.local_tile(cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,))
+        sin_src = cute.local_tile(cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,))
+
+        cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256)
+        cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
+
+        if const_expr(cos_sin_dtype is Float32):
+            cute.copy(cp_f32x8, cos_src, cos_vals)
+            cute.copy(cp_f32x8, sin_src, sin_vals)
+        else:
+            cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
+            sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
+            cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2)
+            cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2)
+
+            for i in cutlass.range_constexpr(4):
+                cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i])
+                sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i])
+                cos_vals[i * 2] = cos0
+                cos_vals[i * 2 + 1] = cos1
+                sin_vals[i * 2] = sin0
+                sin_vals[i * 2 + 1] = sin1
+
+        for i in cutlass.range_constexpr(coarsen):
+            for j in cutlass.range_constexpr(8):
+                q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
+                rot0 = q0 * cos_vals[j] - q1 * sin_vals[j]
+                rot1 = q0 * sin_vals[j] + q1 * cos_vals[j]
+                # convert back to BF16 to match numerics
+                q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1)
+
+    return (
+        q_bf16x2,
+        tid,
+        global_tid,
+        sublane,
+        token_id,
+        head_tile_id,
+        head_start,
+        in_bounds,
+        num_token_heads,
     )
 
 
-class IndexerQRopeQuantKernel:
-    """Shared infrastructure for indexer-Q RoPE+quant fused kernels.
+class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]):
+    tb_size = 128
 
-    Subclasses implement ``kernel`` for a particular Q quantization scheme
-    (MXFP4, FP8 e4m3, …). The base class owns the launch geometry and the
-    common preamble: thread/token addressing, the BF16 Q load, and the
-    interleaved-RoPE pass over the trailing ``rope_dim`` lanes.
-    """
-
-    def __init__(
-        self,
-        head_dim: int = 128,
-        rope_dim: int = 64,
-        num_heads: int = 64,
-        cos_sin_dtype: type[cutlass.Numeric] = Float32,
-        coarsen: int = 4,
-    ):
-        self.head_dim = head_dim
-        self.rope_dim = rope_dim
-        self.nope_dim = head_dim - rope_dim
-        self.num_heads = num_heads
-        self.cos_sin_dtype = cos_sin_dtype
+    @dataclass(frozen=True)
+    class CompileKey:
+        head_dim: int
+        rope_dim: int
+        num_heads: int
+        cos_sin_dtype: type[cutlass.Numeric]
+        weights_out_dtype: type[cutlass.Numeric]
+        coarsen: int
 
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        head_dim = compile_key.head_dim
+        rope_dim = compile_key.rope_dim
+        nope_dim = head_dim - rope_dim
+        num_heads = compile_key.num_heads
+        cos_sin_dtype = compile_key.cos_sin_dtype
         # process multiple heads at the same time to armotize RoPE load costs
+        coarsen = compile_key.coarsen
+        tb_size = IndexerQMxFp4Kernel.tb_size
+        weights_out_dtype = compile_key.weights_out_dtype
         assert num_heads % coarsen == 0
-        self.coarsen = coarsen
 
         # later we will use 32B load = 16 BF16 elems
         # thus, head_dim=128 requires 8 threads to handle.
         # let's call subwarp = 8 threads.
-        self.subwarp_size = head_dim // 16
-        self.tb_size = 128
-        self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size
-
-    @cute.jit
-    def _load_q_and_rope(
-        self,
-        positions: cute.Tensor,
-        q: cute.Tensor,
-        cos_sin_cache: cute.Tensor,
-    ):
-        """Compute thread indices, load Q (BF16), and apply interleaved RoPE.
-
-        Returns a tuple
-            (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id,
-             head_start, in_bounds, num_token_heads)
-        where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed
-        bf16x2 pairs covering the 16 BF16 lanes owned by this thread for
-        each of ``coarsen`` heads. RoPE is applied in place to the
-        trailing ``rope_dim`` lanes; the leading nope lanes pass through.
-        """
-        block_id, _, _ = cute.arch.block_idx()
-        tid, _, _ = cute.arch.thread_idx()
-
-        num_tokens = q.shape[0]
-        num_token_heads = num_tokens * self.num_heads
-        global_tid = block_id * self.tb_size + tid
-
-        global_subwarp_id = global_tid // self.subwarp_size
-        sublane = tid % self.subwarp_size
-
-        token_id = global_subwarp_id // (self.num_heads // self.coarsen)
-        head_tile_id = global_subwarp_id % (self.num_heads // self.coarsen)
-        head_start = head_tile_id * self.coarsen
-
-        # NOTE: token_id may exceed bounds, hence we need to add load/store guards
-        # we can't do early exit because CuteDSL doesn't support it. and we also need
-        # all threads in a warp to be active since we utilize warp shuffle later.
-        # must_in_bounds is constexpr, True when a token spans a whole number of
-        # threadblocks. Then num_tokens * threads_per_token is always a multiple of
-        # tb_size, so the trailing block never straddles num_tokens and the compiler
-        # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads
-        # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half
-        # of the last block addressing token num_tokens.
-        must_in_bounds = cutlass.const_expr(self.threads_per_token % self.tb_size == 0)
-        in_bounds = must_in_bounds or (token_id < num_tokens)
-
-        cp_op = cute.nvgpu.CopyUniversalOp()
-
-        _layout = cute.make_layout((self.coarsen, 8), stride=(8, 1))
-        q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32)
-
-        if in_bounds:
-            # we can't do cute.copy() on the whole 2D tile directly because
-            # cute.copy() wants the 1st mode to be covered by the copy atom,
-            # and other modes as for loop. there is no fast way to
-            # "transpose" the tensor view.
-            q_tile = cute.local_tile(
-                q[token_id, None, None],
-                tiler=(self.coarsen, 16),
-                coord=(head_tile_id, sublane),
-            )
-            cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256)
-            for i in cutlass.range_constexpr(self.coarsen):
-                src = cute.recast_tensor(q_tile[i, None], Uint32)
-                cute.copy(cp_u32x8, src, q_bf16x2[i, None])
-
-        # RoPE applies only to the trailing rope_dim values. We keep the rounded
-        # BF16 result in q_bits so the later amax and quantization see BF16.
-        # cos_sin_cache layout: [max_pos, rope_dim]
-        if in_bounds and sublane * 16 >= self.nope_dim:
-            cos_vals = cute.make_rmem_tensor((8,), Float32)
-            sin_vals = cute.make_rmem_tensor((8,), Float32)
-
-            pos = positions[token_id]
-
-            # select 8 elems from cos and sin
-            cos_id = sublane - self.nope_dim // 16
-            sin_id = cos_id + self.rope_dim // 16
-            cos_src = cute.local_tile(
-                cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,)
-            )
-            sin_src = cute.local_tile(
-                cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,)
+        subwarp_size = head_dim // 16
+        threads_per_token = (num_heads // coarsen) * subwarp_size
+
+        @cute.kernel
+        def device_kernel(
+            positions: cute.Tensor,
+            q: cute.Tensor,
+            cos_sin_cache: cute.Tensor,
+            weights: cute.Tensor,
+            q_quant: cute.Tensor,
+            q_scale: cute.Tensor,
+            weights_out: cute.Tensor,
+            scale: Float32,
+        ):
+            (
+                q_bf16x2,
+                tid,
+                global_tid,
+                sublane,
+                token_id,
+                head_tile_id,
+                head_start,
+                in_bounds,
+                num_token_heads,
+            ) = _load_q_and_rope(
+                positions,
+                q,
+                cos_sin_cache,
+                num_heads,
+                coarsen,
+                tb_size,
+                subwarp_size,
+                threads_per_token,
+                nope_dim,
+                rope_dim,
+                cos_sin_dtype,
             )
 
-            cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256)
-            cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
-
-            if const_expr(self.cos_sin_dtype is Float32):
-                cute.copy(cp_f32x8, cos_src, cos_vals)
-                cute.copy(cp_f32x8, sin_src, sin_vals)
-            else:
-                cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
-                sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
-                cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2)
-                cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2)
-
-                for i in cutlass.range_constexpr(4):
-                    cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i])
-                    sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i])
-                    cos_vals[i * 2] = cos0
-                    cos_vals[i * 2 + 1] = cos1
-                    sin_vals[i * 2] = sin0
-                    sin_vals[i * 2 + 1] = sin1
-
-            for i in cutlass.range_constexpr(self.coarsen):
-                for j in cutlass.range_constexpr(8):
-                    q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
-                    rot0 = q0 * cos_vals[j] - q1 * sin_vals[j]
-                    rot1 = q0 * sin_vals[j] + q1 * cos_vals[j]
-                    # convert back to BF16 to match numerics
-                    q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1)
-
-        return (
-            q_bf16x2,
-            tid,
-            global_tid,
-            sublane,
-            token_id,
-            head_tile_id,
-            head_start,
-            in_bounds,
-            num_token_heads,
-        )
+            cp_op = cute.nvgpu.CopyUniversalOp()
 
+            # layout: [coarsen, 8]
+            q_fp4_tile = cute.local_tile(
+                q_quant[token_id, None, None],
+                tiler=(coarsen, 8),
+                coord=(head_tile_id, sublane),
+            )
 
-class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
-    """Eight-thread subwarps process one ``(token, head)`` row."""
-
-    @cute.jit
-    def __call__(
-        self,
-        positions: cute.Tensor,
-        q: cute.Tensor,
-        cos_sin_cache: cute.Tensor,
-        weights: cute.Tensor,
-        q_quant: cute.Tensor,
-        q_scale: cute.Tensor,
-        weights_out: cute.Tensor,
-        scale: Float32,
-        stream: CUstream,
-    ):
-        total_threads = q.shape[0] * self.threads_per_token
-        grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1)
-        self.kernel(
-            positions,
-            q,
-            cos_sin_cache,
-            weights,
-            q_quant,
-            q_scale,
-            weights_out,
-            scale,
-        ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
-
-    @cute.kernel
-    def kernel(
+            for i in cutlass.range_constexpr(coarsen):
+                # compute amax in packed bf16x2 to save instructions
+                # Each thread holds 16 elems. Two adjacent threads form one 32-elem
+                # MXFP4 block, so a width-2 shuffle gives the block amax.
+                amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
+                for j in cutlass.range_constexpr(1, 8):
+                    amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
+                amax_bf16x2 = cute_utils.warp_reduce(
+                    amax_bf16x2,
+                    _bf16x2_max,
+                    width=MXFP4_BLOCK_SIZE // 16,
+                )
+                amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
+                amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
+
+                if in_bounds:
+                    # compute block scale with bit manipulation
+                    # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask
+                    # increments the exponent whenever fp4_scale is not exactly a power of 2
+                    eps = cutlass.const_expr(float.fromhex("0x6p-126"))
+                    fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0)
+                    bits = recast_val(fp4_scale, Uint32)
+                    ue8m0 = cute_utils.shr_u32(
+                        bits + Uint32(0x7FFFFF), Uint32(23)
+                    ) & Uint32(0xFF)
+
+                    # Only one of the two threads in an MXFP4 block writes the shared scale.
+                    if tid % 2 == 0:
+                        mx_block = sublane // 2
+                        q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0)
+
+                    # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent
+                    # -A + 127 = 254 - ue8m0.
+                    inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23)
+                    inv_fp4_scale = recast_val(inv_scale_bits, Float32)
+
+                    vals = cute.make_rmem_tensor(16, Float32)
+                    for j in cutlass.range_constexpr(8):
+                        q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
+                        vals[j * 2] = q0 * inv_fp4_scale
+                        vals[j * 2 + 1] = q1 * inv_fp4_scale
+
+                    # pack to FP4
+                    packed = cute.make_rmem_tensor((2,), Uint32)
+                    packed[0] = cvt.fp32x8_to_fp4x8(vals, 0)
+                    packed[1] = cvt.fp32x8_to_fp4x8(vals, 8)
+
+                    dst = q_fp4_tile[i, None]
+                    cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
+                    cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32))
+
+            # Weight scaling is independent of the Q subwarp work. The first
+            # num_tokens * num_heads logical threads cover one weight each.
+            if global_tid < num_token_heads:
+                weight_token_id = global_tid // num_heads
+                weight_head_id = global_tid % num_heads
+                weights_out[weight_token_id, weight_head_id] = (
+                    weights[weight_token_id, weight_head_id].to(Float32) * scale
+                ).to(weights_out_dtype)
+
+        @cute.jit
+        def host_entrypoint(
+            positions: cute.Tensor,
+            q: cute.Tensor,
+            cos_sin_cache: cute.Tensor,
+            weights: cute.Tensor,
+            q_quant: cute.Tensor,
+            q_scale: cute.Tensor,
+            weights_out: cute.Tensor,
+            scale: Float32,
+            stream: CUstream,
+        ):
+            total_threads = q.shape[0] * threads_per_token
+            grid = (cute.ceil_div(total_threads, tb_size), 1, 1)
+            device_kernel(
+                positions,
+                q,
+                cos_sin_cache,
+                weights,
+                q_quant,
+                q_scale,
+                weights_out,
+                scale,
+            ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
+
+        return host_entrypoint
+
+    def dispatch(  # type: ignore[override]
         self,
-        positions: cute.Tensor,
-        q: cute.Tensor,
-        cos_sin_cache: cute.Tensor,
-        weights: cute.Tensor,
-        q_quant: cute.Tensor,
-        q_scale: cute.Tensor,
-        weights_out: cute.Tensor,
-        scale: Float32,
-    ):
-        (
-            q_bf16x2,
-            tid,
-            global_tid,
-            sublane,
-            token_id,
-            head_tile_id,
-            head_start,
-            in_bounds,
-            num_token_heads,
-        ) = self._load_q_and_rope(positions, q, cos_sin_cache)
-
-        cp_op = cute.nvgpu.CopyUniversalOp()
-
-        # layout: [coarsen, 8]
-        q_fp4_tile = cute.local_tile(
-            q_quant[token_id, None, None],
-            tiler=(self.coarsen, 8),
-            coord=(head_tile_id, sublane),
+        *,
+        cos_sin_dtype: type[cutlass.Numeric],
+        weights_out_dtype: type[cutlass.Numeric],
+        **compile_key_fields: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            **compile_key_fields,
+            cos_sin_dtype=cos_sin_dtype,
+            weights_out_dtype=weights_out_dtype,
         )
 
-        for i in cutlass.range_constexpr(self.coarsen):
-            # compute amax in packed bf16x2 to save instructions
-            # Each thread holds 16 elems. Two adjacent threads form one 32-elem
-            # MXFP4 block, so a width-2 shuffle gives the block amax.
-            amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
-            for j in cutlass.range_constexpr(1, 8):
-                amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
-            amax_bf16x2 = cute_utils.warp_reduce(
-                amax_bf16x2,
-                _bf16x2_max,
-                width=MXFP4_BLOCK_SIZE // 16,
-            )
-            amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
-            amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
-
-            if in_bounds:
-                # compute block scale with bit manipulation
-                # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask
-                # increments the exponent whenever fp4_scale is not exactly a power of 2
-                eps = cutlass.const_expr(float.fromhex("0x6p-126"))
-                fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0)
-                bits = recast_val(fp4_scale, Uint32)
-                ue8m0 = cute_utils.shr_u32(
-                    bits + Uint32(0x7FFFFF), Uint32(23)
-                ) & Uint32(0xFF)
-
-                # Only one of the two threads in an MXFP4 block writes the shared scale.
-                if tid % 2 == 0:
-                    mx_block = sublane // 2
-                    q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0)
-
-                # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent
-                # -A + 127 = 254 - ue8m0.
-                inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23)
-                inv_fp4_scale = recast_val(inv_scale_bits, Float32)
-
-                vals = cute.make_rmem_tensor(16, Float32)
-                for j in cutlass.range_constexpr(8):
-                    q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
-                    vals[j * 2] = q0 * inv_fp4_scale
-                    vals[j * 2 + 1] = q1 * inv_fp4_scale
-
-                # pack to FP4
-                packed = cute.make_rmem_tensor((2,), Uint32)
-                packed[0] = cvt.fp32x8_to_fp4x8(vals, 0)
-                packed[1] = cvt.fp32x8_to_fp4x8(vals, 8)
-
-                dst = q_fp4_tile[i, None]
-                cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
-                cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32))
-
-        # Weight scaling is independent of the Q subwarp work. The first
-        # num_tokens * num_heads logical threads cover one weight each.
-        if global_tid < num_token_heads:
-            weight_token_id = global_tid // self.num_heads
-            weight_head_id = global_tid % self.num_heads
-            weights_out[weight_token_id, weight_head_id] = (
-                weights[weight_token_id, weight_head_id].to(Float32) * scale
-            )
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        from vllm.v1.attention.backends.mla.indexer import dsa_indexer_uses_fp4
+
+        hf_config = vllm_config.model_config.hf_config
+        if hf_config is None or not dsa_indexer_uses_fp4(vllm_config):
+            return []
+
+        num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0)
+        head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
+        rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
+        if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
+            return []
+
+        # The DeepSeek V4.1 sparse-logits indexer takes bf16 weights; the
+        # dense scoring kernels take fp32.
+        weights_out_dtypes: tuple[type[cutlass.Numeric], ...] = (Float32,)
+        if vllm_config.attention_config.indexer_sparse_logits:
+            weights_out_dtypes = (Float32, BFloat16)
+        return self._trace_dispatch(self.dispatch)(
+            head_dim=head_dim,
+            rope_dim=rope_dim,
+            num_heads=num_heads,
+            cos_sin_dtype=Float32,
+            weights_out_dtype=weights_out_dtypes,
+            coarsen=(1, 4),
+        )
 
-    @cache
-    @staticmethod
-    def compile(
-        head_dim: int = 128,
-        rope_dim: int = 64,
-        num_heads: int = 64,
-        cos_sin_dtype: type[cutlass.Numeric] = Float32,
-        coarsen: int = 4,
-    ):
+    def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
         num_tokens = cute.sym_int()
         max_pos = cute.sym_int()
-
         q = make_fake_tensor(
-            BFloat16, (num_tokens, num_heads, head_dim), divisibility=16
+            BFloat16,
+            (num_tokens, compile_key.num_heads, compile_key.head_dim),
+            divisibility=16,
         )
         positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1)
         cos_sin_cache = make_fake_tensor(
-            cos_sin_dtype,
-            (max_pos, rope_dim),
+            compile_key.cos_sin_dtype,
+            (max_pos, compile_key.rope_dim),
             divisibility=8,
         )
-        weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8)
-        q_fp4 = make_fake_tensor(
+        weights = make_fake_tensor(
+            BFloat16, (num_tokens, compile_key.num_heads), divisibility=8
+        )
+        q_packed = make_fake_tensor(
             Uint8,
-            (num_tokens, num_heads, head_dim // 2),
+            (num_tokens, compile_key.num_heads, compile_key.head_dim // 2),
             divisibility=16,
         )
         q_scale = make_fake_tensor(
             Uint8,
-            (num_tokens, num_heads, head_dim // MXFP4_BLOCK_SIZE),
+            (
+                num_tokens,
+                compile_key.num_heads,
+                compile_key.head_dim // MXFP4_BLOCK_SIZE,
+            ),
             divisibility=4,
         )
-        weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4)
-
-        kernel = IndexerQMxFp4Kernel(
-            head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen
+        weights_out = make_fake_tensor(
+            compile_key.weights_out_dtype,
+            (num_tokens, compile_key.num_heads),
+            divisibility=4,
         )
-        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-        return cute.compile(
-            kernel,
+        return (
             positions,
             q,
             cos_sin_cache,
             weights,
-            q_fp4,
+            q_packed,
             q_scale,
             weights_out,
             Float32(0.0),
-            stream,
-            options="--enable-tvm-ffi",
         )
 
-
-class IndexerQFp8Kernel(IndexerQRopeQuantKernel):
-    """Eight-thread subwarps process one ``(token, head)`` row and emit
-    float8 e4m3fn with a single per-(token, head) scalar scale folded
-    into the per-token weight (mirrors ``_fused_indexer_q_rope_quant_kernel``).
-    """
-
-    def __init__(
-        self,
-        head_dim: int = 128,
-        rope_dim: int = 64,
-        num_heads: int = 64,
-        cos_sin_dtype: type[cutlass.Numeric] = Float32,
-        coarsen: int = 4,
-    ):
-        super().__init__(head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen)
-        # Each subwarp owns `coarsen` heads; we use the first `coarsen`
-        # threads of the subwarp to write the per-head weights using the
-        # fp8 scale computed in the matching loop iteration.
-        assert self.coarsen <= self.subwarp_size, (
-            f"FP8 kernel requires coarsen ({self.coarsen}) <= "
-            f"subwarp_size ({self.subwarp_size}) for the weight-fold step"
-        )
-
-    @cute.jit
+    @kernel_launcher
     def __call__(
         self,
-        positions: cute.Tensor,
-        q: cute.Tensor,
-        cos_sin_cache: cute.Tensor,
-        weights: cute.Tensor,
-        q_fp8: cute.Tensor,
-        weights_out: cute.Tensor,
-        scale: Float32,
-        stream: CUstream,
-    ):
-        total_threads = q.shape[0] * self.threads_per_token
-        grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1)
-        self.kernel(
+        *,
+        positions: torch.Tensor,
+        q: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        weights: torch.Tensor,
+        weights_softmax_scale: float,
+        weights_head_scale: float,
+        q_packed: torch.Tensor,
+        q_scale: torch.Tensor,
+        weights_out: torch.Tensor,
+    ) -> CuTeDSLLaunchSpec["IndexerQMxFp4Kernel.CompileKey"]:
+        num_tokens, num_heads, head_dim = q.shape
+        compile_key = self.dispatch(
+            head_dim=head_dim,
+            rope_dim=cos_sin_cache.shape[-1],
+            num_heads=num_heads,
+            cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype),
+            weights_out_dtype=torch_to_cute_dtype(weights_out.dtype),
+            coarsen=1 if num_tokens < 512 else 4,
+        )
+        launch_args = (
             positions,
             q,
             cos_sin_cache,
             weights,
-            q_fp8,
+            q_packed,
+            q_scale,
             weights_out,
-            scale,
-        ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
+            float(weights_softmax_scale * weights_head_scale),
+        )
+        return compile_key, launch_args
 
-    @cute.kernel
-    def kernel(
-        self,
-        positions: cute.Tensor,
-        q: cute.Tensor,
-        cos_sin_cache: cute.Tensor,
-        weights: cute.Tensor,
-        q_fp8: cute.Tensor,
-        weights_out: cute.Tensor,
-        scale: Float32,
-    ):
-        (
-            q_bf16x2,
-            _tid,
-            _global_tid,
-            sublane,
-            token_id,
-            head_tile_id,
-            head_start,
-            in_bounds,
-            _num_token_heads,
-        ) = self._load_q_and_rope(positions, q, cos_sin_cache)
-
-        cp_op = cute.nvgpu.CopyUniversalOp()
-
-        # layout: [coarsen, 16] bytes (one e4m3fn per element).
-        q_fp8_tile = cute.local_tile(
-            q_fp8[token_id, None, None],
-            tiler=(self.coarsen, 16),
-            coord=(head_tile_id, sublane),
+
+class IndexerQFp8Kernel(VllmCuTeDSLJitKernel["IndexerQFp8Kernel.CompileKey"]):
+    tb_size = 128
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        head_dim: int
+        rope_dim: int
+        num_heads: int
+        cos_sin_dtype: type[cutlass.Numeric]
+        coarsen: int
+
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        head_dim = compile_key.head_dim
+        rope_dim = compile_key.rope_dim
+        nope_dim = head_dim - rope_dim
+        num_heads = compile_key.num_heads
+        cos_sin_dtype = compile_key.cos_sin_dtype
+        # process multiple heads at the same time to armotize RoPE load costs
+        coarsen = compile_key.coarsen
+        tb_size = IndexerQFp8Kernel.tb_size
+        assert num_heads % coarsen == 0
+
+        # later we will use 32B load = 16 BF16 elems
+        # thus, head_dim=128 requires 8 threads to handle.
+        # let's call subwarp = 8 threads.
+        subwarp_size = head_dim // 16
+        threads_per_token = (num_heads // coarsen) * subwarp_size
+        # Each subwarp owns `coarsen` heads; we use the first `coarsen`
+        # threads of the subwarp to write the per-head weights using the
+        # fp8 scale computed in the matching loop iteration.
+        assert coarsen <= subwarp_size, (
+            f"FP8 kernel requires coarsen ({coarsen}) <= "
+            f"subwarp_size ({subwarp_size}) for the weight-fold step"
         )
 
-        for i in cutlass.range_constexpr(self.coarsen):
-            # Reduce amax across the full head_dim: each thread already holds
-            # the max over its 16 lanes; a width=subwarp_size warp shuffle
-            # spreads the head-wide max to every lane in the subwarp.
-            amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
-            for j in cutlass.range_constexpr(1, 8):
-                amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
-            amax_bf16x2 = cute_utils.warp_reduce(
-                amax_bf16x2,
-                _bf16x2_max,
-                width=self.subwarp_size,
+        @cute.kernel
+        def device_kernel(
+            positions: cute.Tensor,
+            q: cute.Tensor,
+            cos_sin_cache: cute.Tensor,
+            weights: cute.Tensor,
+            q_fp8: cute.Tensor,
+            weights_out: cute.Tensor,
+            scale: Float32,
+        ):
+            (
+                q_bf16x2,
+                _tid,
+                _global_tid,
+                sublane,
+                token_id,
+                head_tile_id,
+                head_start,
+                in_bounds,
+                _num_token_heads,
+            ) = _load_q_and_rope(
+                positions,
+                q,
+                cos_sin_cache,
+                num_heads,
+                coarsen,
+                tb_size,
+                subwarp_size,
+                threads_per_token,
+                nope_dim,
+                rope_dim,
+                cos_sin_dtype,
+            )
+
+            cp_op = cute.nvgpu.CopyUniversalOp()
+
+            # layout: [coarsen, 16] bytes (one e4m3fn per element).
+            q_fp8_tile = cute.local_tile(
+                q_fp8[token_id, None, None],
+                tiler=(coarsen, 16),
+                coord=(head_tile_id, sublane),
             )
-            amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
-            amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
-
-            # scale = max(amax, eps) / fp8_max, then rounded UP to the next
-            # power of two. Adding the mantissa mask before shifting out the
-            # mantissa bumps the exponent whenever s isn't a pure pow2.
-            fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0)
-            bits = recast_val(fp32_scale, Uint32)
-            scale_exp = cute_utils.shr_u32(
-                bits + Uint32(0x7FFFFF), Uint32(23)
-            ) & Uint32(0xFF)
-
-            # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23
-            fp8_scale_bits = scale_exp << Uint32(23)
-            fp8_scale = recast_val(fp8_scale_bits, Float32)
-            # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23
-            inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23)
-            inv_fp8_scale = recast_val(inv_scale_bits, Float32)
-
-            # Weight fold: weights_out = weights * q_scale * scale_combined.
-            # All threads in the subwarp share the same fp8_scale after the
-            # warp_reduce above, so we let thread `sublane == i` write the
-            # weight for head `head_start + i`.
-            if in_bounds and sublane == i:
-                head_id = head_start + i
-                weights_out[token_id, head_id] = (
-                    weights[token_id, head_id].to(Float32) * scale * fp8_scale
+
+            for i in cutlass.range_constexpr(coarsen):
+                # Reduce amax across the full head_dim: each thread already holds
+                # the max over its 16 lanes; a width=subwarp_size warp shuffle
+                # spreads the head-wide max to every lane in the subwarp.
+                amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
+                for j in cutlass.range_constexpr(1, 8):
+                    amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
+                amax_bf16x2 = cute_utils.warp_reduce(
+                    amax_bf16x2,
+                    _bf16x2_max,
+                    width=subwarp_size,
                 )
+                amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
+                amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
+
+                # scale = max(amax, eps) / fp8_max, then rounded UP to the next
+                # power of two. Adding the mantissa mask before shifting out the
+                # mantissa bumps the exponent whenever s isn't a pure pow2.
+                fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0)
+                bits = recast_val(fp32_scale, Uint32)
+                scale_exp = cute_utils.shr_u32(
+                    bits + Uint32(0x7FFFFF), Uint32(23)
+                ) & Uint32(0xFF)
 
-            if in_bounds:
-                # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s
-                # (one cp.async-shaped 128-bit store per row).
-                packed = cute.make_rmem_tensor((4,), Uint32)
-                for j in cutlass.range_constexpr(4):
-                    q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2])
-                    q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1])
-                    packed[j] = cvt.fp32x4_to_fp8x4(
-                        q0 * inv_fp8_scale,
-                        q1 * inv_fp8_scale,
-                        q2 * inv_fp8_scale,
-                        q3 * inv_fp8_scale,
+                # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23
+                fp8_scale_bits = scale_exp << Uint32(23)
+                fp8_scale = recast_val(fp8_scale_bits, Float32)
+                # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23
+                inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23)
+                inv_fp8_scale = recast_val(inv_scale_bits, Float32)
+
+                # Weight fold: weights_out = weights * q_scale * scale_combined.
+                # All threads in the subwarp share the same fp8_scale after the
+                # warp_reduce above, so we let thread `sublane == i` write the
+                # weight for head `head_start + i`.
+                if in_bounds and sublane == i:
+                    head_id = head_start + i
+                    weights_out[token_id, head_id] = (
+                        weights[token_id, head_id].to(Float32) * scale * fp8_scale
                     )
 
-                dst = q_fp8_tile[i, None]
-                cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
-                cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32))
+                if in_bounds:
+                    # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s
+                    # (one cp.async-shaped 128-bit store per row).
+                    packed = cute.make_rmem_tensor((4,), Uint32)
+                    for j in cutlass.range_constexpr(4):
+                        q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2])
+                        q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1])
+                        packed[j] = cvt.fp32x4_to_fp8x4(
+                            q0 * inv_fp8_scale,
+                            q1 * inv_fp8_scale,
+                            q2 * inv_fp8_scale,
+                            q3 * inv_fp8_scale,
+                        )
+
+                    dst = q_fp8_tile[i, None]
+                    cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
+                    cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32))
+
+        @cute.jit
+        def host_entrypoint(
+            positions: cute.Tensor,
+            q: cute.Tensor,
+            cos_sin_cache: cute.Tensor,
+            weights: cute.Tensor,
+            q_fp8: cute.Tensor,
+            weights_out: cute.Tensor,
+            scale: Float32,
+            stream: CUstream,
+        ):
+            total_threads = q.shape[0] * threads_per_token
+            grid = (cute.ceil_div(total_threads, tb_size), 1, 1)
+            device_kernel(
+                positions,
+                q,
+                cos_sin_cache,
+                weights,
+                q_fp8,
+                weights_out,
+                scale,
+            ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
+
+        return host_entrypoint
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        cos_sin_dtype: type[cutlass.Numeric],
+        **compile_key_fields: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            **compile_key_fields,
+            cos_sin_dtype=cos_sin_dtype,
+        )
 
-    @cache
-    @staticmethod
-    def compile(
-        head_dim: int = 128,
-        rope_dim: int = 64,
-        num_heads: int = 64,
-        cos_sin_dtype: type[cutlass.Numeric] = Float32,
-        coarsen: int = 4,
-    ):
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        from vllm.v1.attention.backends.mla.indexer import dsa_indexer_uses_fp4
+
+        hf_config = vllm_config.model_config.hf_config
+        if hf_config is None or dsa_indexer_uses_fp4(vllm_config):
+            return []
+
+        num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0)
+        head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
+        rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
+        if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
+            return []
+
+        return self._trace_dispatch(self.dispatch)(
+            head_dim=head_dim,
+            rope_dim=rope_dim,
+            num_heads=num_heads,
+            cos_sin_dtype=Float32,
+            coarsen=(1, 4),
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
         num_tokens = cute.sym_int()
         max_pos = cute.sym_int()
-
         q = make_fake_tensor(
-            BFloat16, (num_tokens, num_heads, head_dim), divisibility=16
+            BFloat16,
+            (num_tokens, compile_key.num_heads, compile_key.head_dim),
+            divisibility=16,
         )
         positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1)
         cos_sin_cache = make_fake_tensor(
-            cos_sin_dtype,
-            (max_pos, rope_dim),
+            compile_key.cos_sin_dtype,
+            (max_pos, compile_key.rope_dim),
             divisibility=8,
         )
-        weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8)
+        weights = make_fake_tensor(
+            BFloat16, (num_tokens, compile_key.num_heads), divisibility=8
+        )
         q_fp8 = make_fake_tensor(
             Uint8,
-            (num_tokens, num_heads, head_dim),
+            (num_tokens, compile_key.num_heads, compile_key.head_dim),
             divisibility=16,
         )
-        weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4)
-
-        kernel = IndexerQFp8Kernel(
-            head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen
+        weights_out = make_fake_tensor(
+            Float32, (num_tokens, compile_key.num_heads), divisibility=4
         )
-        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-        return cute.compile(
-            kernel,
+        return (
             positions,
             q,
             cos_sin_cache,
@@ -609,6 +669,40 @@ def compile(
             q_fp8,
             weights_out,
             Float32(0.0),
-            stream,
-            options="--enable-tvm-ffi",
         )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        *,
+        positions: torch.Tensor,
+        q: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        weights: torch.Tensor,
+        weights_softmax_scale: float,
+        weights_head_scale: float,
+        q_fp8: torch.Tensor,
+        weights_out: torch.Tensor,
+    ) -> CuTeDSLLaunchSpec["IndexerQFp8Kernel.CompileKey"]:
+        num_tokens, num_heads, head_dim = q.shape
+        compile_key = self.dispatch(
+            head_dim=head_dim,
+            rope_dim=cos_sin_cache.shape[-1],
+            num_heads=num_heads,
+            cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype),
+            coarsen=1 if num_tokens < 512 else 4,
+        )
+        launch_args = (
+            positions,
+            q,
+            cos_sin_cache,
+            weights,
+            q_fp8.view(torch.uint8),
+            weights_out,
+            float(weights_softmax_scale * weights_head_scale),
+        )
+        return compile_key, launch_args
+
+
+_INDEXER_Q_MXFP4_KERNEL = IndexerQMxFp4Kernel()
+_INDEXER_Q_FP8_KERNEL = IndexerQFp8Kernel()
diff --git a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py
index 72ce43b46cbf..ec36bf763ebb 100644
--- a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py
+++ b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py
@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
+
 import torch
 import torch.nn as nn
 
@@ -31,7 +33,7 @@ def deep_gemm_fp8_o_proj(
     positions: torch.Tensor,
     cos_sin_cache: torch.Tensor,
     wo_a: nn.Module,
-    wo_b: nn.Module,
+    wo_b: Callable[[torch.Tensor], torch.Tensor],
     *,
     n_groups: int,
     heads_per_group: int,
@@ -44,7 +46,9 @@ def deep_gemm_fp8_o_proj(
     """O projection: inverse RoPE + grouped wo_a + wo_b.
 
     Shared by the FlashMLA and FlashInfer CUDA backends. The attention
-    layer selects the recipe at initialization.
+    layer selects the recipe at initialization. ``wo_b`` is any callable over
+    the flattened ``z``: the projection module itself, or a wrapper that also
+    reduce-scatters its output (DeepSeek-V4.1 GEMM-RS).
     """
     use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn
     o_proj_input, o_scale = fused_inv_rope_fp8_quant(
diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py
index d254dad9d6a1..aa9482bed643 100644
--- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py
+++ b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py
@@ -12,7 +12,7 @@
 from vllm.triton_utils import tl, triton
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["num_tokens"])
 def _prepare_megamoe_inputs_kernel(
     hidden_states,
     x_fp8,
@@ -40,27 +40,32 @@ def _prepare_megamoe_inputs_kernel(
     topk_idx_stride_k: tl.constexpr,
     topk_weights_out_stride_m: tl.constexpr,
     topk_weights_out_stride_k: tl.constexpr,
+    num_tokens,
     hidden_size: tl.constexpr,
     top_k: tl.constexpr,
+    BLOCK_M: tl.constexpr,
     BLOCK_K: tl.constexpr,
     GROUP_K: tl.constexpr,
     BLOCK_TOPK: tl.constexpr,
     SHARED_BLOCK_M: tl.constexpr,
 ) -> None:
-    token_id = tl.program_id(0)
+    token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    token_mask = token_id < num_tokens
     k_block_id = tl.program_id(1)
 
     k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
-    k_mask = k_offsets < hidden_size
+    k_mask = token_mask[:, None] & (k_offsets[None, :] < hidden_size)
     hidden = tl.load(
-        hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
+        hidden_states
+        + token_id[:, None] * hidden_stride_m
+        + k_offsets[None, :] * hidden_stride_k,
         mask=k_mask,
         other=0.0,
     ).to(tl.float32)
 
     num_groups: tl.constexpr = BLOCK_K // GROUP_K
-    hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
-    amax = tl.max(hidden_groups, axis=1)
+    hidden_groups = tl.reshape(tl.abs(hidden), [BLOCK_M, num_groups, GROUP_K])
+    amax = tl.max(hidden_groups, axis=2)
     amax = tl.maximum(amax, 1.0e-4)
 
     scale = amax / 448.0
@@ -71,21 +76,24 @@ def _prepare_megamoe_inputs_kernel(
     scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
     rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
 
-    hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
-    scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
-    scaled = tl.reshape(scaled, [BLOCK_K])
+    hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K])
+    scaled = hidden_groups * (1.0 / rounded_scale)[:, :, None]
+    scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K])
     fp8 = scaled.to(tl.float8e4nv)
     tl.store(
-        x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
+        x_fp8 + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k,
         fp8,
         mask=k_mask,
     )
 
     scale_offsets = tl.arange(0, num_groups)
-    packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
+    packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to(
+        tl.int32
+    )
     tl.store(
         x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
         packed_scale,
+        mask=token_mask,
     )
 
     # DeepGEMM's SM100 shared-expert TMA loads require the activation scales
@@ -105,41 +113,48 @@ def _prepare_megamoe_inputs_kernel(
             + shared_row * shared_x_sf_stride_m
             + k_block_id * shared_x_sf_stride_k,
             packed_scale,
+            mask=token_mask,
         )
 
     if k_block_id == 0:
         topk_offsets = tl.arange(0, BLOCK_TOPK)
-        topk_mask = topk_offsets < top_k
-        token_is_padding = False
+        topk_mask = token_mask[:, None] & (topk_offsets[None, :] < top_k)
+        token_is_padding = tl.full((BLOCK_M,), False, tl.int1)
         if is_padding is not None:
-            token_is_padding = tl.load(is_padding + token_id * is_padding_stride_m)
+            token_is_padding = tl.load(
+                is_padding + token_id * is_padding_stride_m,
+                mask=token_mask,
+                other=True,
+            )
 
         ids = tl.load(
-            topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
+            topk_ids
+            + token_id[:, None] * topk_ids_stride_m
+            + topk_offsets[None, :] * topk_ids_stride_k,
             mask=topk_mask,
             other=0,
         ).to(tl.int64)
-        ids = tl.where(token_is_padding, -1, ids)
+        ids = tl.where(token_is_padding[:, None], -1, ids)
         tl.store(
             topk_idx_out
-            + token_id * topk_idx_stride_m
-            + topk_offsets * topk_idx_stride_k,
+            + token_id[:, None] * topk_idx_stride_m
+            + topk_offsets[None, :] * topk_idx_stride_k,
             ids,
             mask=topk_mask,
         )
 
         weights = tl.load(
             topk_weights
-            + token_id * topk_weights_stride_m
-            + topk_offsets * topk_weights_stride_k,
+            + token_id[:, None] * topk_weights_stride_m
+            + topk_offsets[None, :] * topk_weights_stride_k,
             mask=topk_mask,
             other=0.0,
         )
-        weights = tl.where(token_is_padding, 0.0, weights)
+        weights = tl.where(token_is_padding[:, None], 0.0, weights)
         tl.store(
             topk_weights_out
-            + token_id * topk_weights_out_stride_m
-            + topk_offsets * topk_weights_out_stride_k,
+            + token_id[:, None] * topk_weights_out_stride_m
+            + topk_offsets[None, :] * topk_weights_out_stride_k,
             weights,
             mask=topk_mask,
         )
@@ -195,7 +210,9 @@ def prepare_megamoe_inputs(
             )
 
     block_k = 128
-    grid = (num_tokens, triton.cdiv(hidden_size, block_k))
+    # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled.
+    block_m = 8 if num_tokens >= 64 else 1
+    grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k))
     block_topk = triton.next_power_of_2(top_k)
     padding_stride_m = is_padding.stride(0) if is_padding is not None else 0
     _prepare_megamoe_inputs_kernel[grid](
@@ -225,8 +242,10 @@ def prepare_megamoe_inputs(
         topk_idx_out.stride(1),
         topk_weights_out.stride(0),
         topk_weights_out.stride(1),
+        num_tokens,
         hidden_size,
         top_k,
+        BLOCK_M=block_m,
         BLOCK_K=block_k,
         GROUP_K=32,
         BLOCK_TOPK=block_topk,
diff --git a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
index 4ff4b232d10f..74866b096089 100644
--- a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
+++ b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
@@ -7,22 +7,25 @@
 
 from __future__ import annotations
 
-from functools import cache
+from dataclasses import dataclass
 from typing import Any
 
 import cutlass
 import cutlass.cute as cute
 import torch
 from cuda.bindings.driver import CUstream
-from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr
+from cutlass import Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr
 from cutlass._mlir.dialects import llvm
 from cutlass.cutlass_dsl import T, dsl_user_op
 from quack.compile_utils import make_fake_tensor
 
-_TORCH_TO_CUTE = {
-    torch.bfloat16: BFloat16,
-    torch.float32: Float32,
-}
+from vllm.cute_utils import torch_to_cute_dtype
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
+from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
+    CuTeDSLLaunchSpec,
+    VllmCuTeDSLJitKernel,
+)
+from vllm.utils.math_utils import round_up
 
 
 @dsl_user_op
@@ -72,361 +75,406 @@ def _fp32x2_to_fp8e4m3x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint16
     return Uint16(out)
 
 
-class SparseAttnCompressNormRopeStoreC4Kernel:
+class SparseAttnCompressNormRopeStoreC4Kernel(
+    VllmCuTeDSLJitKernel["SparseAttnCompressNormRopeStoreC4Kernel.CompileKey"]
+):
     min_scale = 1.0e-4
     rcp_ln2 = 1.4426950408889634
 
-    def __init__(
-        self,
-        head_size: int,
-        state_width: int,
-        rope_head_dim: int,
-        fp8_max: float,
-        quant_block: int,
-        token_stride: int,
-        scale_dim: int,
-        compress_ratio: int,
-        overlap: bool,
-    ):
-        self.head_dim = head_size
-        self.state_width = state_width
-        self.rope_dim = rope_head_dim
-        self.nope_dim = head_size - rope_head_dim
-        self.fp8_max = fp8_max
-        self.quant_block = quant_block
-        self.token_stride = token_stride
-        self.scale_dim = scale_dim
-        self.elems_per_lane = 8
-        self.copy_elems = 4
-        self.copy_chunks = self.elems_per_lane // self.copy_elems
-        self.lanes_per_group = quant_block // self.elems_per_lane
-        self.groups_per_warp = 32 // self.lanes_per_group
-        self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1
-        self.scale_reduce_offset = self.lanes_per_group // 2
-        self.num_warps = (head_size // quant_block) // self.groups_per_warp
-        self.nope_blocks = self.nope_dim // quant_block
-        self.tb_size = self.num_warps * 32
-        self.compress_ratio = compress_ratio
-        self.overlap = overlap
-        self.window = (1 + int(overlap)) * compress_ratio
-
-    @cute.jit
-    def __call__(
-        self,
-        state_cache: cute.Tensor,
-        token_to_req_indices: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        block_table: cute.Tensor,
-        block_size: Int64,
-        rms_norm_weight: cute.Tensor,
-        rms_norm_eps: Float32,
-        cos_sin_cache: cute.Tensor,
-        k_cache: cute.Tensor,
-        kv_slot_mapping: cute.Tensor,
-        kv_cache_block_size: Int64,
-        stream: CUstream,
-    ):
-        grid = (slot_mapping.shape[0], 1, 1)
-        self.kernel(
-            state_cache,
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_size,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            k_cache,
-            kv_slot_mapping,
-            kv_cache_block_size,
-        ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
-
-    @cute.kernel
-    def kernel(
-        self,
-        state_cache: cute.Tensor,
-        token_to_req_indices: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        block_table: cute.Tensor,
-        block_size: Int64,
-        rms_norm_weight: cute.Tensor,
-        rms_norm_eps: Float32,
-        cos_sin_cache: cute.Tensor,
-        k_cache: cute.Tensor,
-        kv_slot_mapping: cute.Tensor,
-        kv_cache_block_size: Int64,
-    ):
-        token_idx, _, _ = cute.arch.block_idx()
-        tid, _, _ = cute.arch.thread_idx()
-        warp_id = cute.arch.make_warp_uniform(tid // 32)
-        lane_id = tid % 32
-        group_lane = lane_id % self.lanes_per_group
-        group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group
-        elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane
-
-        slot_id = slot_mapping[token_idx]
-        has_position = token_idx < positions.shape[0]
-        position = Int64(0)
-        if has_position:
-            position = positions[token_idx]
-        boundary = has_position and (
-            (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
-        )
-        has_req_idx = token_idx < token_to_req_indices.shape[0]
-        has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
-        kv_slot_idx = Int64(-1)
-        if has_kv_slot_idx:
-            kv_slot_idx = kv_slot_mapping[token_idx]
-        active = (
-            slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0)
-        )
+    @dataclass(frozen=True)
+    class CompileKey:
+        head_size: int
+        state_width: int
+        rope_head_dim: int
+        fp8_max: float
+        quant_block: int
+        token_stride: int
+        scale_dim: int
+        compress_ratio: int
+        overlap: bool
+        norm_weight_dtype: type[cutlass.Numeric]
 
-        if active:
-            req_idx = token_to_req_indices[token_idx]
-            start = position - Int64(self.window - 1)
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        min_scale = SparseAttnCompressNormRopeStoreC4Kernel.min_scale
+        rcp_ln2 = SparseAttnCompressNormRopeStoreC4Kernel.rcp_ln2
+        head_dim = compile_key.head_size
+        state_width = compile_key.state_width
+        rope_dim = compile_key.rope_head_dim
+        nope_dim = head_dim - rope_dim
+        fp8_max = compile_key.fp8_max
+        quant_block = compile_key.quant_block
+        token_stride = compile_key.token_stride
+        scale_dim = compile_key.scale_dim
+        elems_per_lane = 8
+        copy_elems = 4
+        copy_chunks = elems_per_lane // copy_elems
+        lanes_per_group = quant_block // elems_per_lane
+        groups_per_warp = 32 // lanes_per_group
+        scale_reduce_steps = lanes_per_group.bit_length() - 1
+        scale_reduce_offset = lanes_per_group // 2
+        num_warps = (head_dim // quant_block) // groups_per_warp
+        nope_blocks = nope_dim // quant_block
+        tb_size = num_warps * 32
+        compress_ratio = compile_key.compress_ratio
+        overlap = compile_key.overlap
+        window = (1 + int(overlap)) * compress_ratio
+
+        @cute.kernel
+        def device_kernel(
+            state_cache: cute.Tensor,
+            token_to_req_indices: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            block_table: cute.Tensor,
+            block_size: Int64,
+            rms_norm_weight: cute.Tensor,
+            rms_norm_eps: Float32,
+            cos_sin_cache: cute.Tensor,
+            k_cache: cute.Tensor,
+            kv_slot_mapping: cute.Tensor,
+            kv_cache_block_size: Int64,
+        ):
+            token_idx, _, _ = cute.arch.block_idx()
+            tid, _, _ = cute.arch.thread_idx()
+            warp_id = cute.arch.make_warp_uniform(tid // 32)
+            lane_id = tid % 32
+            group_lane = lane_id % lanes_per_group
+            group_idx = warp_id * groups_per_warp + lane_id // lanes_per_group
+            elem_base = group_idx * quant_block + group_lane * elems_per_lane
 
-            smem = cutlass.utils.SmemAllocator()
-            s_block_numbers = smem.allocate_tensor(
-                Int32, cute.make_layout((self.window,)), byte_alignment=4
-            )
-            partial_sums = smem.allocate_tensor(
-                Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
+            slot_id = slot_mapping[token_idx]
+            has_position = token_idx < positions.shape[0]
+            position = Int64(0)
+            if has_position:
+                position = positions[token_idx]
+            boundary = has_position and (
+                (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
             )
-            rrms_shared = smem.allocate_tensor(
-                Float32, cute.make_layout((1,)), byte_alignment=4
+            has_req_idx = token_idx < token_to_req_indices.shape[0]
+            has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
+            kv_slot_idx = Int64(-1)
+            if has_kv_slot_idx:
+                kv_slot_idx = kv_slot_mapping[token_idx]
+            active = (
+                slot_id >= Int64(0)
+                and has_req_idx
+                and boundary
+                and kv_slot_idx >= Int64(0)
             )
 
-            for row in cutlass.range_constexpr(self.window):
-                pos = start + Int64(row)
-                if tid == row:
-                    block_number_i32 = Int32(0)
-                    if pos >= Int64(0):
-                        block_index = pos // block_size
-                        block_number_i32 = block_table[req_idx, block_index]
-                    s_block_numbers[row] = block_number_i32
-            cute.arch.sync_threads()
+            if active:
+                req_idx = token_to_req_indices[token_idx]
+                start = position - Int64(window - 1)
 
-            local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-            local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-            local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
+                smem = cutlass.utils.SmemAllocator()
+                s_block_numbers = smem.allocate_tensor(
+                    Int32, cute.make_layout((window,)), byte_alignment=4
+                )
+                partial_sums = smem.allocate_tensor(
+                    Float32, cute.make_layout((num_warps,)), byte_alignment=4
+                )
+                rrms_shared = smem.allocate_tensor(
+                    Float32, cute.make_layout((1,)), byte_alignment=4
+                )
 
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                local_max[e] = -Float32.inf
-                local_sum[e] = Float32(0.0)
-                local_product[e] = Float32(0.0)
+                for row in cutlass.range_constexpr(window):
+                    pos = start + Int64(row)
+                    if tid == row:
+                        block_number_i32 = Int32(0)
+                        if pos >= Int64(0):
+                            block_index = pos // block_size
+                            block_number_i32 = block_table[req_idx, block_index]
+                        s_block_numbers[row] = block_number_i32
+                cute.arch.sync_threads()
+
+                local_max = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                local_product = cute.make_rmem_tensor((elems_per_lane,), Float32)
+
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    local_max[e] = -Float32.inf
+                    local_sum[e] = Float32(0.0)
+                    local_product[e] = Float32(0.0)
+
+                cp_f32x4 = cute.make_copy_atom(
+                    cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
+                )
+                copy_layout = cute.make_layout(
+                    (copy_chunks, copy_elems),
+                    stride=(copy_elems, 1),
+                )
+                kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
+                score_vals = cute.make_rmem_tensor(copy_layout, Float32)
 
-            cp_f32x4 = cute.make_copy_atom(
-                cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
-            )
-            copy_layout = cute.make_layout(
-                (self.copy_chunks, self.copy_elems),
-                stride=(self.copy_elems, 1),
-            )
-            kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
-            score_vals = cute.make_rmem_tensor(copy_layout, Float32)
-
-            for row in cutlass.range_constexpr(self.window):
-                pos = start + Int64(row)
-                if pos >= Int64(0):
-                    block_index = pos // block_size
-                    block_offset = pos - block_index * block_size
-                    block_number = s_block_numbers[row].to(Int64)
-                    head_offset = Int64((row // self.compress_ratio) * self.head_dim)
-                    row_tensor = state_cache[block_number, block_offset, None]
-                    for chunk in cutlass.range_constexpr(self.copy_chunks):
-                        copy_elem = const_expr(chunk * self.copy_elems)
-                        col_tile = (
-                            head_offset + (elem_base + Int32(copy_elem)).to(Int64)
-                        ) // Int64(self.copy_elems)
-                        kv_src = cute.local_tile(
-                            row_tensor,
-                            tiler=(self.copy_elems,),
-                            coord=(col_tile,),
-                        )
-                        score_src = cute.local_tile(
-                            row_tensor,
-                            tiler=(self.copy_elems,),
-                            coord=(
-                                col_tile + Int64(self.state_width // self.copy_elems),
-                            ),
-                        )
-                        cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
-                        cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
-
-                    for e in cutlass.range_constexpr(self.elems_per_lane):
-                        chunk = const_expr(e // self.copy_elems)
-                        copy_elem = const_expr(e % self.copy_elems)
-                        score = score_vals[chunk, copy_elem]
-                        kv = kv_vals[chunk, copy_elem]
-                        new_max = cute.arch.fmax(local_max[e], score)
-                        old_scale = cute.math.exp2(
-                            (local_max[e] - new_max) * Float32(self.rcp_ln2),
-                            fastmath=True,
-                        )
-                        new_scale = cute.math.exp2(
-                            (score - new_max) * Float32(self.rcp_ln2),
-                            fastmath=True,
-                        )
-                        local_sum[e] = local_sum[e] * old_scale + new_scale
-                        local_product[e] = local_product[e] * old_scale + kv * new_scale
-                        local_max[e] = new_max
+                for row in cutlass.range_constexpr(window):
+                    pos = start + Int64(row)
+                    if pos >= Int64(0):
+                        block_index = pos // block_size
+                        block_offset = pos - block_index * block_size
+                        block_number = s_block_numbers[row].to(Int64)
+                        head_offset = Int64((row // compress_ratio) * head_dim)
+                        row_tensor = state_cache[block_number, block_offset, None]
+                        for chunk in cutlass.range_constexpr(copy_chunks):
+                            copy_elem = const_expr(chunk * copy_elems)
+                            col_tile = (
+                                head_offset + (elem_base + Int32(copy_elem)).to(Int64)
+                            ) // Int64(copy_elems)
+                            kv_src = cute.local_tile(
+                                row_tensor,
+                                tiler=(copy_elems,),
+                                coord=(col_tile,),
+                            )
+                            score_src = cute.local_tile(
+                                row_tensor,
+                                tiler=(copy_elems,),
+                                coord=(col_tile + Int64(state_width // copy_elems),),
+                            )
+                            cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
+                            cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
+
+                        for e in cutlass.range_constexpr(elems_per_lane):
+                            chunk = const_expr(e // copy_elems)
+                            copy_elem = const_expr(e % copy_elems)
+                            score = score_vals[chunk, copy_elem]
+                            kv = kv_vals[chunk, copy_elem]
+                            new_max = cute.arch.fmax(local_max[e], score)
+                            old_scale = cute.math.exp2(
+                                (local_max[e] - new_max) * Float32(rcp_ln2),
+                                fastmath=True,
+                            )
+                            new_scale = cute.math.exp2(
+                                (score - new_max) * Float32(rcp_ln2),
+                                fastmath=True,
+                            )
+                            local_sum[e] = local_sum[e] * old_scale + new_scale
+                            local_product[e] = (
+                                local_product[e] * old_scale + kv * new_scale
+                            )
+                            local_max[e] = new_max
+
+                x = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                local_sumsq = Float32(0.0)
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    x[e] = local_product[e] / local_sum[e]
+                    local_sumsq += x[e] * x[e]
+
+                warp_sum = local_sumsq
+                for step in cutlass.range_constexpr(5):
+                    offset = const_expr(16 >> step)
+                    warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
 
-            x = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-            local_sumsq = Float32(0.0)
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                x[e] = local_product[e] / local_sum[e]
-                local_sumsq += x[e] * x[e]
+                if lane_id == 0:
+                    partial_sums[warp_id] = warp_sum
+                cute.arch.sync_threads()
+                if tid == 0:
+                    total = Float32(0.0)
+                    for i in cutlass.range_constexpr(num_warps):
+                        total += partial_sums[i]
+                    rrms_shared[0] = cute.math.rsqrt(
+                        total / Float32(head_dim) + rms_norm_eps, fastmath=True
+                    )
+                cute.arch.sync_threads()
 
-            warp_sum = local_sumsq
-            for step in cutlass.range_constexpr(5):
-                offset = const_expr(16 >> step)
-                warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
+                rrms = rrms_shared[0]
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    elem = elem_base + e
+                    x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
 
-            if lane_id == 0:
-                partial_sums[warp_id] = warp_sum
-            cute.arch.sync_threads()
-            if tid == 0:
-                total = Float32(0.0)
-                for i in cutlass.range_constexpr(self.num_warps):
-                    total += partial_sums[i]
-                rrms_shared[0] = cute.math.rsqrt(
-                    total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
+                k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
+                k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
+                page = kv_slot_idx // kv_cache_block_size
+                kv_offset = kv_slot_idx - page * kv_cache_block_size
+                value_base = page * k_cache.stride[0] + kv_offset * Int64(token_stride)
+                scale_base = (
+                    page * k_cache.stride[0]
+                    + kv_cache_block_size * Int64(token_stride)
+                    + kv_offset * Int64(scale_dim)
                 )
-            cute.arch.sync_threads()
-
-            rrms = rrms_shared[0]
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                elem = elem_base + e
-                x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
-
-            k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
-            k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
-            page = kv_slot_idx // kv_cache_block_size
-            kv_offset = kv_slot_idx - page * kv_cache_block_size
-            value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
-            scale_base = (
-                page * k_cache.stride[0]
-                + kv_cache_block_size * Int64(self.token_stride)
-                + kv_offset * Int64(self.scale_dim)
-            )
 
-            if group_idx == self.nope_blocks:
-                compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
-                    self.compress_ratio
-                )
-                for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
-                    elem = const_expr(pair * 2)
-                    pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
-                    cos_v = cos_sin_cache[compressed_pos, pair_idx]
-                    sin_v = cos_sin_cache[
-                        compressed_pos, pair_idx + Int32(self.rope_dim // 2)
-                    ]
-                    real = x[elem] * cos_v - x[elem + 1] * sin_v
-                    imag = x[elem] * sin_v + x[elem + 1] * cos_v
-                    packed = _fp32x2_to_bf16x2(real, imag)
-                    out_base = (
-                        value_base
-                        + Int64(self.nope_dim)
-                        + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64)
+                if group_idx == nope_blocks:
+                    compressed_pos = (position // Int64(compress_ratio)) * Int64(
+                        compress_ratio
                     )
-                    k_cache_u32.iterator[out_base // Int64(4)] = packed
-            else:
-                q = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-                local_absmax = Float32(0.0)
-                for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
-                    elem = const_expr(pair * 2)
-                    q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
-                    q0, q1 = _bf16x2_to_fp32(q_packed)
-                    q[elem] = q0
-                    q[elem + 1] = q1
-                    local_absmax = cute.arch.fmax(
-                        local_absmax,
-                        cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)),
+                    for pair in cutlass.range_constexpr(elems_per_lane // 2):
+                        elem = const_expr(pair * 2)
+                        pair_idx = (elem_base - nope_dim) // 2 + Int32(pair)
+                        cos_v = cos_sin_cache[compressed_pos, pair_idx]
+                        sin_v = cos_sin_cache[
+                            compressed_pos, pair_idx + Int32(rope_dim // 2)
+                        ]
+                        real = x[elem] * cos_v - x[elem + 1] * sin_v
+                        imag = x[elem] * sin_v + x[elem + 1] * cos_v
+                        packed = _fp32x2_to_bf16x2(real, imag)
+                        out_base = (
+                            value_base
+                            + Int64(nope_dim)
+                            + ((elem_base - nope_dim + Int32(elem)) * 2).to(Int64)
+                        )
+                        k_cache_u32.iterator[out_base // Int64(4)] = packed
+                else:
+                    q = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                    local_absmax = Float32(0.0)
+                    for pair in cutlass.range_constexpr(elems_per_lane // 2):
+                        elem = const_expr(pair * 2)
+                        q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
+                        q0, q1 = _bf16x2_to_fp32(q_packed)
+                        q[elem] = q0
+                        q[elem + 1] = q1
+                        local_absmax = cute.arch.fmax(
+                            local_absmax,
+                            cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)),
+                        )
+                    absmax = local_absmax
+                    group_mask_and_clamp = const_expr(
+                        (cute.arch.WARP_SIZE - lanes_per_group) << 8
+                        | (cute.arch.WARP_SIZE - 1)
                     )
-                absmax = local_absmax
-                group_mask_and_clamp = const_expr(
-                    (cute.arch.WARP_SIZE - self.lanes_per_group) << 8
-                    | (cute.arch.WARP_SIZE - 1)
-                )
-                for step in cutlass.range_constexpr(self.scale_reduce_steps):
-                    offset = const_expr(self.scale_reduce_offset >> step)
-                    absmax = cute.arch.fmax(
-                        absmax,
-                        cute.arch.shuffle_sync_bfly(
+                    for step in cutlass.range_constexpr(scale_reduce_steps):
+                        offset = const_expr(scale_reduce_offset >> step)
+                        absmax = cute.arch.fmax(
                             absmax,
-                            offset=offset,
-                            mask_and_clamp=group_mask_and_clamp,
-                        ),
-                    )
-                scale_raw = cute.arch.fmax(
-                    Float32(self.min_scale),
-                    absmax / Float32(self.fp8_max),
-                )
-                bits = _recast_val(scale_raw, Uint32)
-                ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
-                inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
-                for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
-                    elem = const_expr(pair * 2)
-                    y0 = cutlass.min(
-                        cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)),
-                        Float32(self.fp8_max),
+                            cute.arch.shuffle_sync_bfly(
+                                absmax,
+                                offset=offset,
+                                mask_and_clamp=group_mask_and_clamp,
+                            ),
+                        )
+                    scale_raw = cute.arch.fmax(
+                        Float32(min_scale),
+                        absmax / Float32(fp8_max),
                     )
-                    y1 = cutlass.min(
-                        cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)),
-                        Float32(self.fp8_max),
+                    bits = _recast_val(scale_raw, Uint32)
+                    ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
+                    inv_scale = _recast_val(
+                        (Uint32(254) - ue8m0) << Uint32(23), Float32
                     )
-                    packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
-                    out_base = value_base + (elem_base + Int32(elem)).to(Int64)
-                    k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
-                if group_lane == 0:
-                    k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8)
-                    if group_idx == 0:
-                        k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
-                            0
+                    for pair in cutlass.range_constexpr(elems_per_lane // 2):
+                        elem = const_expr(pair * 2)
+                        y0 = cutlass.min(
+                            cute.arch.fmax(q[elem] * inv_scale, Float32(-fp8_max)),
+                            Float32(fp8_max),
+                        )
+                        y1 = cutlass.min(
+                            cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-fp8_max)),
+                            Float32(fp8_max),
+                        )
+                        packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
+                        out_base = value_base + (elem_base + Int32(elem)).to(Int64)
+                        k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
+                    if group_lane == 0:
+                        k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(
+                            Uint8
                         )
+                        if group_idx == 0:
+                            k_cache.iterator[scale_base + Int64(nope_blocks)] = Uint8(0)
+
+        @cute.jit
+        def host_entrypoint(
+            state_cache: cute.Tensor,
+            token_to_req_indices: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            block_table: cute.Tensor,
+            block_size: Int64,
+            rms_norm_weight: cute.Tensor,
+            rms_norm_eps: Float32,
+            cos_sin_cache: cute.Tensor,
+            k_cache: cute.Tensor,
+            kv_slot_mapping: cute.Tensor,
+            kv_cache_block_size: Int64,
+            stream: CUstream,
+        ):
+            grid = (slot_mapping.shape[0], 1, 1)
+            device_kernel(
+                state_cache,
+                token_to_req_indices,
+                positions,
+                slot_mapping,
+                block_table,
+                block_size,
+                rms_norm_weight,
+                rms_norm_eps,
+                cos_sin_cache,
+                k_cache,
+                kv_slot_mapping,
+                kv_cache_block_size,
+            ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
+
+        return host_entrypoint
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        compress_ratio: int,
+        norm_weight_dtype: type[cutlass.Numeric],
+        head_size: int,
+        rope_head_dim: int,
+    ) -> CompileKey:
+        overlap = compress_ratio == 4
+        state_width = head_size * (2 if overlap else 1)
+        return self.CompileKey(
+            head_size=head_size,
+            state_width=state_width,
+            rope_head_dim=rope_head_dim,
+            fp8_max=448.0,
+            quant_block=64,
+            token_stride=576,
+            scale_dim=8,
+            compress_ratio=compress_ratio,
+            overlap=overlap,
+            norm_weight_dtype=norm_weight_dtype,
+        )
 
-    @cache
-    @staticmethod
-    def compile(
-        head_size: int = 512,
-        state_width: int = 1024,
-        rope_head_dim: int = 64,
-        fp8_max: float = 448.0,
-        quant_block: int = 64,
-        token_stride: int = 576,
-        scale_dim: int = 8,
-        kv_block_stride: int = 74752,
-        compress_ratio: int = 4,
-        overlap: bool = True,
-        norm_weight_dtype: type[cutlass.Numeric] = Float32,
-    ):
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        hf_config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        return self._trace_dispatch(self.dispatch)(
+            compress_ratio=4,
+            norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype),
+            head_size=hf_config.head_dim,
+            rope_head_dim=hf_config.qk_rope_head_dim,
+            enabled=(
+                cache_config.block_size > 0
+                and 4 in hf_config.compress_ratios
+                and cache_config.cache_dtype == "fp8_ds_mla"
+            ),
+            _when=lambda *, enabled: enabled,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
+        head_size = compile_key.head_size
+        state_width = compile_key.state_width
+        rope_head_dim = compile_key.rope_head_dim
+        quant_block = compile_key.quant_block
+        token_stride = compile_key.token_stride
+        scale_dim = compile_key.scale_dim
+        compress_ratio = compile_key.compress_ratio
+        overlap = compile_key.overlap
+        norm_weight_dtype = compile_key.norm_weight_dtype
         if compress_ratio != 4 or not overlap:
             raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.")
         if head_size != 512:
             raise ValueError(
-                "CuTe DSL C4 fused sparse-attn currently requires head_size=512."
+                "CuTe DSL C4 fused sparse-attn currently requires head_size=512.",
             )
         if state_width != 2 * head_size:
             raise ValueError(
-                "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size."
+                "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size.",
             )
         if quant_block != 64:
             raise ValueError(
-                "CuTe DSL C4 fused sparse-attn currently requires quant_block=64."
+                "CuTe DSL C4 fused sparse-attn currently requires quant_block=64.",
             )
         if rope_head_dim != 64:
             raise ValueError(
-                "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64."
+                "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64.",
             )
         if token_stride < head_size + rope_head_dim:
             raise ValueError("token_stride is too small for the packed FP8/BF16 row.")
         expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1
         if scale_dim < expected_scale_dim:
             raise ValueError("scale_dim is too small for the UE8M0 scale row.")
-
         num_positions = cute.sym_int()
         num_slots = cute.sym_int()
         num_req_indices = cute.sym_int()
@@ -437,7 +485,6 @@ def compile(
         block_table_width = cute.sym_int()
         max_pos = cute.sym_int()
         state_cache_width = state_width * 2
-
         state_cache = cute.runtime.make_fake_tensor(
             Float32,
             (num_state_blocks, state_cache_block_size, state_cache_width),
@@ -449,15 +496,21 @@ def compile(
             assumed_align=16,
         )
         token_to_req_indices = make_fake_tensor(
-            Int32, (num_req_indices,), divisibility=4
+            Int32,
+            (num_req_indices,),
+            divisibility=4,
         )
         positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
         slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
         block_table = make_fake_tensor(
-            Int32, (cute.sym_int(), block_table_width), divisibility=1
+            Int32,
+            (cute.sym_int(), block_table_width),
+            divisibility=1,
         )
         rms_norm_weight = make_fake_tensor(
-            norm_weight_dtype, (head_size,), divisibility=4
+            norm_weight_dtype,
+            (head_size,),
+            divisibility=4,
         )
         cos_sin_cache = cute.runtime.make_fake_tensor(
             Float32,
@@ -468,29 +521,12 @@ def compile(
         k_cache = cute.runtime.make_fake_tensor(
             Uint8,
             (num_kv_blocks, cute.sym_int(), cute.sym_int()),
-            stride=(
-                cute.sym_int64(divisibility=16),
-                cute.sym_int64(divisibility=8),
-                1,
-            ),
+            stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1),
             assumed_align=16,
         )
         kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
 
-        kernel = SparseAttnCompressNormRopeStoreC4Kernel(
-            head_size,
-            state_width,
-            rope_head_dim,
-            fp8_max,
-            quant_block,
-            token_stride,
-            scale_dim,
-            compress_ratio,
-            overlap,
-        )
-        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-        return cute.compile(
-            kernel,
+        return (
             state_cache,
             token_to_req_indices,
             positions,
@@ -503,60 +539,34 @@ def compile(
             k_cache,
             kv_slot_mapping,
             Int64(0),
-            stream,
-            options="--enable-tvm-ffi",
         )
 
-
-class SparseAttnCompressNormRopeStoreFullC4Kernel(
-    SparseAttnCompressNormRopeStoreC4Kernel
-):
-    def __init__(
+    @kernel_launcher
+    def __call__(
         self,
-        head_size: int,
-        state_width: int,
-        rope_head_dim: int,
-        fp8_max: float,
-        quant_block: int,
-        token_stride: int,
-        scale_dim: int,
+        *,
+        state_cache: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+        positions: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        rms_norm_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        kv_cache: torch.Tensor,
+        kv_slot_mapping: torch.Tensor,
         compress_ratio: int,
-        overlap: bool,
-        store_full_fp8: bool = False,
-    ):
-        super().__init__(
-            head_size,
-            state_width,
-            rope_head_dim,
-            fp8_max,
-            quant_block,
-            token_stride,
-            scale_dim,
-            compress_ratio,
-            overlap,
+        head_dim: int,
+        rope_head_dim: int,
+    ) -> CuTeDSLLaunchSpec[SparseAttnCompressNormRopeStoreC4Kernel.CompileKey]:
+        compile_key = self.dispatch(
+            compress_ratio=compress_ratio,
+            norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype),
+            head_size=head_dim,
+            rope_head_dim=rope_head_dim,
         )
-        self.store_full_fp8 = store_full_fp8
-
-    @cute.jit
-    def __call__(
-        self,
-        state_cache: cute.Tensor,
-        token_to_req_indices: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        block_table: cute.Tensor,
-        block_size: Int64,
-        rms_norm_weight: cute.Tensor,
-        rms_norm_eps: Float32,
-        cos_sin_cache: cute.Tensor,
-        k_cache: cute.Tensor,
-        kv_slot_mapping: cute.Tensor,
-        kv_cache_block_size: Int64,
-        fp8_scale: cute.Tensor,
-        stream: CUstream,
-    ):
-        grid = (slot_mapping.shape[0], 1, 1)
-        self.kernel(
+        launch_args = (
             state_cache,
             token_to_req_indices,
             positions,
@@ -566,283 +576,395 @@ def __call__(
             rms_norm_weight,
             rms_norm_eps,
             cos_sin_cache,
-            k_cache,
+            kv_cache,
             kv_slot_mapping,
-            kv_cache_block_size,
-            fp8_scale,
-        ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
-
-    @cute.kernel
-    def kernel(
-        self,
-        state_cache: cute.Tensor,
-        token_to_req_indices: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        block_table: cute.Tensor,
-        block_size: Int64,
-        rms_norm_weight: cute.Tensor,
-        rms_norm_eps: Float32,
-        cos_sin_cache: cute.Tensor,
-        k_cache: cute.Tensor,
-        kv_slot_mapping: cute.Tensor,
-        kv_cache_block_size: Int64,
-        fp8_scale: cute.Tensor,
-    ):
-        token_idx, _, _ = cute.arch.block_idx()
-        tid, _, _ = cute.arch.thread_idx()
-        warp_id = cute.arch.make_warp_uniform(tid // 32)
-        lane_id = tid % 32
-        group_lane = lane_id % self.lanes_per_group
-        group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group
-        elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane
-
-        slot_id = slot_mapping[token_idx]
-        has_position = token_idx < positions.shape[0]
-        position = Int64(0)
-        if has_position:
-            position = positions[token_idx]
-        boundary = has_position and (
-            (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
-        )
-        has_req_idx = token_idx < token_to_req_indices.shape[0]
-        has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
-        kv_slot_idx = Int64(-1)
-        if has_kv_slot_idx:
-            kv_slot_idx = kv_slot_mapping[token_idx]
-        active = (
-            slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0)
+            kv_cache.shape[1],
         )
+        return compile_key, launch_args
 
-        if active:
-            req_idx = token_to_req_indices[token_idx]
-            start = position - Int64(self.window - 1)
-
-            smem = cutlass.utils.SmemAllocator()
-            s_block_numbers = smem.allocate_tensor(
-                Int32, cute.make_layout((self.window,)), byte_alignment=4
-            )
-            partial_sums = smem.allocate_tensor(
-                Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
-            )
-            rrms_shared = smem.allocate_tensor(
-                Float32, cute.make_layout((1,)), byte_alignment=4
-            )
 
-            for row in cutlass.range_constexpr(self.window):
-                pos = start + Int64(row)
-                if tid == row:
-                    block_number_i32 = Int32(0)
-                    if pos >= Int64(0):
-                        block_index = pos // block_size
-                        block_number_i32 = block_table[req_idx, block_index]
-                    s_block_numbers[row] = block_number_i32
-            cute.arch.sync_threads()
+class SparseAttnCompressNormRopeStoreFullC4Kernel(
+    VllmCuTeDSLJitKernel["SparseAttnCompressNormRopeStoreFullC4Kernel.CompileKey"]
+):
+    rcp_ln2 = 1.4426950408889634
 
-            local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-            local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-            local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
+    @dataclass(frozen=True)
+    class CompileKey:
+        head_size: int
+        state_width: int
+        rope_head_dim: int
+        fp8_max: float
+        quant_block: int
+        compress_ratio: int
+        overlap: bool
+        store_full_fp8: bool
+        norm_weight_dtype: type[cutlass.Numeric]
 
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                local_max[e] = -Float32.inf
-                local_sum[e] = Float32(0.0)
-                local_product[e] = Float32(0.0)
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        rcp_ln2 = SparseAttnCompressNormRopeStoreFullC4Kernel.rcp_ln2
+        head_dim = compile_key.head_size
+        state_width = compile_key.state_width
+        rope_dim = compile_key.rope_head_dim
+        nope_dim = head_dim - rope_dim
+        fp8_max = compile_key.fp8_max
+        quant_block = compile_key.quant_block
+        elems_per_lane = 8
+        copy_elems = 4
+        copy_chunks = elems_per_lane // copy_elems
+        lanes_per_group = quant_block // elems_per_lane
+        groups_per_warp = 32 // lanes_per_group
+        num_warps = (head_dim // quant_block) // groups_per_warp
+        nope_blocks = nope_dim // quant_block
+        tb_size = num_warps * 32
+        compress_ratio = compile_key.compress_ratio
+        overlap = compile_key.overlap
+        window = (1 + int(overlap)) * compress_ratio
+        store_full_fp8 = compile_key.store_full_fp8
+
+        @cute.kernel
+        def device_kernel(
+            state_cache: cute.Tensor,
+            token_to_req_indices: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            block_table: cute.Tensor,
+            block_size: Int64,
+            rms_norm_weight: cute.Tensor,
+            rms_norm_eps: Float32,
+            cos_sin_cache: cute.Tensor,
+            k_cache: cute.Tensor,
+            kv_slot_mapping: cute.Tensor,
+            kv_cache_block_size: Int64,
+            fp8_scale: cute.Tensor,
+        ):
+            token_idx, _, _ = cute.arch.block_idx()
+            tid, _, _ = cute.arch.thread_idx()
+            warp_id = cute.arch.make_warp_uniform(tid // 32)
+            lane_id = tid % 32
+            group_lane = lane_id % lanes_per_group
+            group_idx = warp_id * groups_per_warp + lane_id // lanes_per_group
+            elem_base = group_idx * quant_block + group_lane * elems_per_lane
 
-            cp_f32x4 = cute.make_copy_atom(
-                cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
+            slot_id = slot_mapping[token_idx]
+            has_position = token_idx < positions.shape[0]
+            position = Int64(0)
+            if has_position:
+                position = positions[token_idx]
+            boundary = has_position and (
+                (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
             )
-            copy_layout = cute.make_layout(
-                (self.copy_chunks, self.copy_elems),
-                stride=(self.copy_elems, 1),
+            has_req_idx = token_idx < token_to_req_indices.shape[0]
+            has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
+            kv_slot_idx = Int64(-1)
+            if has_kv_slot_idx:
+                kv_slot_idx = kv_slot_mapping[token_idx]
+            active = (
+                slot_id >= Int64(0)
+                and has_req_idx
+                and boundary
+                and kv_slot_idx >= Int64(0)
             )
-            kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
-            score_vals = cute.make_rmem_tensor(copy_layout, Float32)
-
-            for row in cutlass.range_constexpr(self.window):
-                pos = start + Int64(row)
-                if pos >= Int64(0):
-                    block_index = pos // block_size
-                    block_offset = pos - block_index * block_size
-                    block_number = s_block_numbers[row].to(Int64)
-                    head_offset = Int64((row // self.compress_ratio) * self.head_dim)
-                    row_tensor = state_cache[block_number, block_offset, None]
-                    for chunk in cutlass.range_constexpr(self.copy_chunks):
-                        copy_elem = const_expr(chunk * self.copy_elems)
-                        col_tile = (
-                            head_offset + (elem_base + Int32(copy_elem)).to(Int64)
-                        ) // Int64(self.copy_elems)
-                        kv_src = cute.local_tile(
-                            row_tensor,
-                            tiler=(self.copy_elems,),
-                            coord=(col_tile,),
-                        )
-                        score_src = cute.local_tile(
-                            row_tensor,
-                            tiler=(self.copy_elems,),
-                            coord=(
-                                col_tile + Int64(self.state_width // self.copy_elems),
-                            ),
-                        )
-                        cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
-                        cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
-
-                    for e in cutlass.range_constexpr(self.elems_per_lane):
-                        chunk = const_expr(e // self.copy_elems)
-                        copy_elem = const_expr(e % self.copy_elems)
-                        score = score_vals[chunk, copy_elem]
-                        kv = kv_vals[chunk, copy_elem]
-                        new_max = cute.arch.fmax(local_max[e], score)
-                        old_scale = cute.math.exp2(
-                            (local_max[e] - new_max) * Float32(self.rcp_ln2),
-                            fastmath=True,
-                        )
-                        new_scale = cute.math.exp2(
-                            (score - new_max) * Float32(self.rcp_ln2),
-                            fastmath=True,
-                        )
-                        local_sum[e] = local_sum[e] * old_scale + new_scale
-                        local_product[e] = local_product[e] * old_scale + kv * new_scale
-                        local_max[e] = new_max
 
-            x = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-            local_sumsq = Float32(0.0)
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                x[e] = local_product[e] / local_sum[e]
-                local_sumsq += x[e] * x[e]
+            if active:
+                req_idx = token_to_req_indices[token_idx]
+                start = position - Int64(window - 1)
 
-            warp_sum = local_sumsq
-            for step in cutlass.range_constexpr(5):
-                offset = const_expr(16 >> step)
-                warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
-
-            if lane_id == 0:
-                partial_sums[warp_id] = warp_sum
-            cute.arch.sync_threads()
-            if tid == 0:
-                total = Float32(0.0)
-                for i in cutlass.range_constexpr(self.num_warps):
-                    total += partial_sums[i]
-                rrms_shared[0] = cute.math.rsqrt(
-                    total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
+                smem = cutlass.utils.SmemAllocator()
+                s_block_numbers = smem.allocate_tensor(
+                    Int32, cute.make_layout((window,)), byte_alignment=4
+                )
+                partial_sums = smem.allocate_tensor(
+                    Float32, cute.make_layout((num_warps,)), byte_alignment=4
+                )
+                rrms_shared = smem.allocate_tensor(
+                    Float32, cute.make_layout((1,)), byte_alignment=4
                 )
-            cute.arch.sync_threads()
 
-            rrms = rrms_shared[0]
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                elem = elem_base + e
-                x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
+                for row in cutlass.range_constexpr(window):
+                    pos = start + Int64(row)
+                    if tid == row:
+                        block_number_i32 = Int32(0)
+                        if pos >= Int64(0):
+                            block_index = pos // block_size
+                            block_number_i32 = block_table[req_idx, block_index]
+                        s_block_numbers[row] = block_number_i32
+                cute.arch.sync_threads()
+
+                local_max = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                local_product = cute.make_rmem_tensor((elems_per_lane,), Float32)
+
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    local_max[e] = -Float32.inf
+                    local_sum[e] = Float32(0.0)
+                    local_product[e] = Float32(0.0)
+
+                cp_f32x4 = cute.make_copy_atom(
+                    cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
+                )
+                copy_layout = cute.make_layout(
+                    (copy_chunks, copy_elems),
+                    stride=(copy_elems, 1),
+                )
+                kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
+                score_vals = cute.make_rmem_tensor(copy_layout, Float32)
 
-            page = kv_slot_idx // kv_cache_block_size
-            kv_offset = kv_slot_idx - page * kv_cache_block_size
-            value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
+                for row in cutlass.range_constexpr(window):
+                    pos = start + Int64(row)
+                    if pos >= Int64(0):
+                        block_index = pos // block_size
+                        block_offset = pos - block_index * block_size
+                        block_number = s_block_numbers[row].to(Int64)
+                        head_offset = Int64((row // compress_ratio) * head_dim)
+                        row_tensor = state_cache[block_number, block_offset, None]
+                        for chunk in cutlass.range_constexpr(copy_chunks):
+                            copy_elem = const_expr(chunk * copy_elems)
+                            col_tile = (
+                                head_offset + (elem_base + Int32(copy_elem)).to(Int64)
+                            ) // Int64(copy_elems)
+                            kv_src = cute.local_tile(
+                                row_tensor,
+                                tiler=(copy_elems,),
+                                coord=(col_tile,),
+                            )
+                            score_src = cute.local_tile(
+                                row_tensor,
+                                tiler=(copy_elems,),
+                                coord=(col_tile + Int64(state_width // copy_elems),),
+                            )
+                            cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
+                            cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
+
+                        for e in cutlass.range_constexpr(elems_per_lane):
+                            chunk = const_expr(e // copy_elems)
+                            copy_elem = const_expr(e % copy_elems)
+                            score = score_vals[chunk, copy_elem]
+                            kv = kv_vals[chunk, copy_elem]
+                            new_max = cute.arch.fmax(local_max[e], score)
+                            old_scale = cute.math.exp2(
+                                (local_max[e] - new_max) * Float32(rcp_ln2),
+                                fastmath=True,
+                            )
+                            new_scale = cute.math.exp2(
+                                (score - new_max) * Float32(rcp_ln2),
+                                fastmath=True,
+                            )
+                            local_sum[e] = local_sum[e] * old_scale + new_scale
+                            local_product[e] = (
+                                local_product[e] * old_scale + kv * new_scale
+                            )
+                            local_max[e] = new_max
+
+                x = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                local_sumsq = Float32(0.0)
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    x[e] = local_product[e] / local_sum[e]
+                    local_sumsq += x[e] * x[e]
+
+                warp_sum = local_sumsq
+                for step in cutlass.range_constexpr(5):
+                    offset = const_expr(16 >> step)
+                    warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
 
-            if const_expr(self.store_full_fp8):
-                k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
-                inv_fp8 = Float32(1.0) / fp8_scale[0]
-                if group_idx == self.nope_blocks:
-                    compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
-                        self.compress_ratio
-                    )
-                    for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
-                        elem = const_expr(pair * 2)
-                        pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
-                        cos_v = cos_sin_cache[compressed_pos, pair_idx]
-                        sin_v = cos_sin_cache[
-                            compressed_pos, pair_idx + Int32(self.rope_dim // 2)
-                        ]
-                        real = x[elem] * cos_v - x[elem + 1] * sin_v
-                        imag = x[elem] * sin_v + x[elem + 1] * cos_v
-                        packed_bf16 = _fp32x2_to_bf16x2(real, imag)
-                        b0, b1 = _bf16x2_to_fp32(packed_bf16)
-                        y0 = cutlass.min(
-                            cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
-                            Float32(self.fp8_max),
-                        )
-                        y1 = cutlass.min(
-                            cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
-                            Float32(self.fp8_max),
-                        )
-                        packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
-                        out_base = value_base + (elem_base + Int32(elem)).to(Int64)
-                        k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
-                else:
-                    for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
-                        elem = const_expr(pair * 2)
-                        packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
-                        b0, b1 = _bf16x2_to_fp32(packed_bf16)
-                        y0 = cutlass.min(
-                            cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
-                            Float32(self.fp8_max),
-                        )
-                        y1 = cutlass.min(
-                            cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
-                            Float32(self.fp8_max),
-                        )
-                        packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
-                        out_base = value_base + (elem_base + Int32(elem)).to(Int64)
-                        k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
-            else:
-                k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
-                if group_idx == self.nope_blocks:
-                    compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
-                        self.compress_ratio
+                if lane_id == 0:
+                    partial_sums[warp_id] = warp_sum
+                cute.arch.sync_threads()
+                if tid == 0:
+                    total = Float32(0.0)
+                    for i in cutlass.range_constexpr(num_warps):
+                        total += partial_sums[i]
+                    rrms_shared[0] = cute.math.rsqrt(
+                        total / Float32(head_dim) + rms_norm_eps, fastmath=True
                     )
-                    for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
-                        elem = const_expr(pair * 2)
-                        pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
-                        cos_v = cos_sin_cache[compressed_pos, pair_idx]
-                        sin_v = cos_sin_cache[
-                            compressed_pos, pair_idx + Int32(self.rope_dim // 2)
-                        ]
-                        real = x[elem] * cos_v - x[elem + 1] * sin_v
-                        imag = x[elem] * sin_v + x[elem + 1] * cos_v
-                        packed_bf16 = _fp32x2_to_bf16x2(real, imag)
-                        out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
-                            Int64
+                cute.arch.sync_threads()
+
+                rrms = rrms_shared[0]
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    elem = elem_base + e
+                    x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
+
+                page = kv_slot_idx // kv_cache_block_size
+                kv_offset = kv_slot_idx - page * kv_cache_block_size
+                value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
+
+                if const_expr(store_full_fp8):
+                    k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
+                    inv_fp8 = Float32(1.0) / fp8_scale[0]
+                    if group_idx == nope_blocks:
+                        compressed_pos = (position // Int64(compress_ratio)) * Int64(
+                            compress_ratio
                         )
-                        k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
+                        for pair in cutlass.range_constexpr(elems_per_lane // 2):
+                            elem = const_expr(pair * 2)
+                            pair_idx = (elem_base - nope_dim) // 2 + Int32(pair)
+                            cos_v = cos_sin_cache[compressed_pos, pair_idx]
+                            sin_v = cos_sin_cache[
+                                compressed_pos, pair_idx + Int32(rope_dim // 2)
+                            ]
+                            real = x[elem] * cos_v - x[elem + 1] * sin_v
+                            imag = x[elem] * sin_v + x[elem + 1] * cos_v
+                            packed_bf16 = _fp32x2_to_bf16x2(real, imag)
+                            b0, b1 = _bf16x2_to_fp32(packed_bf16)
+                            y0 = cutlass.min(
+                                cutlass.max(b0 * inv_fp8, Float32(-fp8_max)),
+                                Float32(fp8_max),
+                            )
+                            y1 = cutlass.min(
+                                cutlass.max(b1 * inv_fp8, Float32(-fp8_max)),
+                                Float32(fp8_max),
+                            )
+                            packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
+                            out_base = value_base + (elem_base + Int32(elem)).to(Int64)
+                            k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
+                    else:
+                        for pair in cutlass.range_constexpr(elems_per_lane // 2):
+                            elem = const_expr(pair * 2)
+                            packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
+                            b0, b1 = _bf16x2_to_fp32(packed_bf16)
+                            y0 = cutlass.min(
+                                cutlass.max(b0 * inv_fp8, Float32(-fp8_max)),
+                                Float32(fp8_max),
+                            )
+                            y1 = cutlass.min(
+                                cutlass.max(b1 * inv_fp8, Float32(-fp8_max)),
+                                Float32(fp8_max),
+                            )
+                            packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
+                            out_base = value_base + (elem_base + Int32(elem)).to(Int64)
+                            k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
                 else:
-                    for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
-                        elem = const_expr(pair * 2)
-                        packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
-                        out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
-                            Int64
+                    k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
+                    if group_idx == nope_blocks:
+                        compressed_pos = (position // Int64(compress_ratio)) * Int64(
+                            compress_ratio
                         )
-                        k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
+                        for pair in cutlass.range_constexpr(elems_per_lane // 2):
+                            elem = const_expr(pair * 2)
+                            pair_idx = (elem_base - nope_dim) // 2 + Int32(pair)
+                            cos_v = cos_sin_cache[compressed_pos, pair_idx]
+                            sin_v = cos_sin_cache[
+                                compressed_pos, pair_idx + Int32(rope_dim // 2)
+                            ]
+                            real = x[elem] * cos_v - x[elem + 1] * sin_v
+                            imag = x[elem] * sin_v + x[elem + 1] * cos_v
+                            packed_bf16 = _fp32x2_to_bf16x2(real, imag)
+                            out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
+                                Int64
+                            )
+                            k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
+                    else:
+                        for pair in cutlass.range_constexpr(elems_per_lane // 2):
+                            elem = const_expr(pair * 2)
+                            packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
+                            out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
+                                Int64
+                            )
+                            k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
+
+        @cute.jit
+        def host_entrypoint(
+            state_cache: cute.Tensor,
+            token_to_req_indices: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            block_table: cute.Tensor,
+            block_size: Int64,
+            rms_norm_weight: cute.Tensor,
+            rms_norm_eps: Float32,
+            cos_sin_cache: cute.Tensor,
+            k_cache: cute.Tensor,
+            kv_slot_mapping: cute.Tensor,
+            kv_cache_block_size: Int64,
+            fp8_scale: cute.Tensor,
+            stream: CUstream,
+        ):
+            grid = (slot_mapping.shape[0], 1, 1)
+            device_kernel(
+                state_cache,
+                token_to_req_indices,
+                positions,
+                slot_mapping,
+                block_table,
+                block_size,
+                rms_norm_weight,
+                rms_norm_eps,
+                cos_sin_cache,
+                k_cache,
+                kv_slot_mapping,
+                kv_cache_block_size,
+                fp8_scale,
+            ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
+
+        return host_entrypoint
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        compress_ratio: int,
+        store_full_fp8: bool,
+        norm_weight_dtype: type[cutlass.Numeric],
+        head_size: int,
+        rope_head_dim: int,
+    ) -> CompileKey:
+        overlap = compress_ratio == 4
+        return self.CompileKey(
+            head_size=head_size,
+            state_width=head_size * (2 if overlap else 1),
+            rope_head_dim=rope_head_dim,
+            fp8_max=448.0,
+            quant_block=64,
+            compress_ratio=compress_ratio,
+            overlap=overlap,
+            store_full_fp8=store_full_fp8,
+            norm_weight_dtype=norm_weight_dtype,
+        )
 
-    @cache
-    @staticmethod
-    def compile(
-        head_size: int = 512,
-        state_width: int = 1024,
-        rope_head_dim: int = 64,
-        fp8_max: float = 448.0,
-        quant_block: int = 64,
-        token_stride: int = 576,
-        scale_dim: int = 8,
-        kv_block_stride: int = 74752,
-        compress_ratio: int = 4,
-        overlap: bool = True,
-        store_full_fp8: bool = False,
-        norm_weight_dtype: type[cutlass.Numeric] = Float32,
-    ):
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        hf_config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        return self._trace_dispatch(self.dispatch)(
+            compress_ratio=4,
+            store_full_fp8=(
+                cache_config.cache_dtype.startswith("fp8")
+                and cache_config.cache_dtype != "fp8_ds_mla"
+            ),
+            norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype),
+            head_size=hf_config.head_dim,
+            rope_head_dim=hf_config.qk_rope_head_dim,
+            enabled=(
+                cache_config.block_size > 0
+                and 4 in hf_config.compress_ratios
+                and cache_config.cache_dtype != "fp8_ds_mla"
+            ),
+            _when=lambda *, enabled: enabled,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
+        head_size = compile_key.head_size
+        state_width = compile_key.state_width
+        rope_head_dim = compile_key.rope_head_dim
+        quant_block = compile_key.quant_block
+        compress_ratio = compile_key.compress_ratio
+        overlap = compile_key.overlap
+        norm_weight_dtype = compile_key.norm_weight_dtype
         if compress_ratio != 4 or not overlap:
             raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.")
         if head_size != 512:
             raise ValueError(
-                "CuTe DSL C4 fused sparse-attn currently requires head_size=512."
+                "CuTe DSL C4 fused sparse-attn currently requires head_size=512.",
             )
         if state_width != 2 * head_size:
             raise ValueError(
-                "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size."
+                "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size.",
             )
         if quant_block != 64:
             raise ValueError(
-                "CuTe DSL C4 fused sparse-attn currently requires quant_block=64."
+                "CuTe DSL C4 fused sparse-attn currently requires quant_block=64.",
             )
         if rope_head_dim != 64:
             raise ValueError(
-                "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64."
+                "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64.",
             )
         num_positions = cute.sym_int()
         num_slots = cute.sym_int()
@@ -854,7 +976,6 @@ def compile(
         block_table_width = cute.sym_int()
         max_pos = cute.sym_int()
         state_cache_width = state_width * 2
-
         state_cache = cute.runtime.make_fake_tensor(
             Float32,
             (num_state_blocks, state_cache_block_size, state_cache_width),
@@ -866,15 +987,21 @@ def compile(
             assumed_align=16,
         )
         token_to_req_indices = make_fake_tensor(
-            Int32, (num_req_indices,), divisibility=4
+            Int32,
+            (num_req_indices,),
+            divisibility=4,
         )
         positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
         slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
         block_table = make_fake_tensor(
-            Int32, (cute.sym_int(), block_table_width), divisibility=1
+            Int32,
+            (cute.sym_int(), block_table_width),
+            divisibility=1,
         )
         rms_norm_weight = make_fake_tensor(
-            norm_weight_dtype, (head_size,), divisibility=4
+            norm_weight_dtype,
+            (head_size,),
+            divisibility=4,
         )
         cos_sin_cache = cute.runtime.make_fake_tensor(
             Float32,
@@ -885,31 +1012,12 @@ def compile(
         k_cache = cute.runtime.make_fake_tensor(
             Uint8,
             (num_kv_blocks, cute.sym_int(), cute.sym_int()),
-            stride=(
-                cute.sym_int64(divisibility=16),
-                cute.sym_int64(divisibility=8),
-                1,
-            ),
+            stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1),
             assumed_align=16,
         )
         kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
         fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1)
-
-        kernel = SparseAttnCompressNormRopeStoreFullC4Kernel(
-            head_size,
-            state_width,
-            rope_head_dim,
-            fp8_max,
-            quant_block,
-            token_stride,
-            scale_dim,
-            compress_ratio,
-            overlap,
-            store_full_fp8,
-        )
-        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-        return cute.compile(
-            kernel,
+        return (
             state_cache,
             token_to_req_indices,
             positions,
@@ -923,272 +1031,348 @@ def compile(
             kv_slot_mapping,
             Int64(0),
             fp8_scale,
-            stream,
-            options="--enable-tvm-ffi",
         )
 
-
-class SparseAttnCompressC128Block8Kernel:
-    head_tile = 64
-    rows_per_warp = 16
-    elems_per_lane = 2
-    lanes_per_row = head_tile // elems_per_lane
-    num_warps = 8
-    stats_lane_stride = lanes_per_row + 1
-    final_reduce_steps = 3
-    final_reduce_initial_offset = 4
-    tb_size = num_warps * 32
-    compress_ratio = 128
-    state_block_size = 8
-    rcp_ln2 = 1.4426950408889634
-
-    def __init__(
+    @kernel_launcher
+    def __call__(
         self,
-        head_size: int,
-        state_width: int,
-    ):
-        self.head_dim = head_size
-        self.num_splits = head_size // self.head_tile
-        self.state_width = state_width
-
-    @cute.jit
-    def __call__(
-        self,
-        state_cache: cute.Tensor,
-        token_to_req_indices: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        block_table: cute.Tensor,
-        compressed_kv: cute.Tensor,
-        stream: CUstream,
-    ):
-        grid = (slot_mapping.shape[0] * self.num_splits, 1, 1)
-        self.kernel(
+        *,
+        state_cache: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+        positions: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        rms_norm_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        kv_cache: torch.Tensor,
+        kv_slot_mapping: torch.Tensor,
+        compress_ratio: int,
+        store_full_fp8: bool,
+        fp8_scale: torch.Tensor | None,
+        head_dim: int,
+        rope_head_dim: int,
+    ) -> CuTeDSLLaunchSpec[SparseAttnCompressNormRopeStoreFullC4Kernel.CompileKey]:
+        compile_key = self.dispatch(
+            compress_ratio=compress_ratio,
+            store_full_fp8=store_full_fp8,
+            norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype),
+            head_size=head_dim,
+            rope_head_dim=rope_head_dim,
+        )
+        if fp8_scale is None:
+            fp8_scale = torch.ones(1, dtype=torch.float32, device=kv_cache.device)
+        launch_args = (
             state_cache,
             token_to_req_indices,
             positions,
             slot_mapping,
             block_table,
-            compressed_kv,
-        ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
+            block_size,
+            rms_norm_weight,
+            rms_norm_eps,
+            cos_sin_cache,
+            kv_cache.view(torch.uint8),
+            kv_slot_mapping,
+            kv_cache.shape[1],
+            fp8_scale,
+        )
+        return compile_key, launch_args
 
-    @cute.kernel
-    def kernel(
-        self,
-        state_cache: cute.Tensor,
-        token_to_req_indices: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        block_table: cute.Tensor,
-        compressed_kv: cute.Tensor,
-    ):
-        block_id, _, _ = cute.arch.block_idx()
-        tid, _, _ = cute.arch.thread_idx()
-        warp_id = cute.arch.make_warp_uniform(tid // 32)
-        lane_id = tid % 32
-        col_group = lane_id % self.lanes_per_row
-
-        token_idx = block_id // self.num_splits
-        split_idx = block_id - token_idx * self.num_splits
-        col_base = split_idx * self.head_tile + col_group * self.elems_per_lane
-
-        position = Int64(0)
-        req_idx = Int32(0)
-        slot_id = Int64(-1)
-        has_position = token_idx < positions.shape[0]
-        has_req_idx = token_idx < token_to_req_indices.shape[0]
-        if lane_id == 0:
-            slot_id = slot_mapping[token_idx]
-        if lane_id == 0 and has_position:
-            position = positions[token_idx]
-        if lane_id == 0 and has_req_idx:
-            req_idx = token_to_req_indices[token_idx]
-        slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
-        position = cute.arch.shuffle_sync(position, offset=0)
-        req_idx = cute.arch.shuffle_sync(req_idx, offset=0)
-        boundary = has_position and (
-            (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
+
+class SparseAttnCompressC128Block8Kernel(
+    VllmCuTeDSLJitKernel["SparseAttnCompressC128Block8Kernel.CompileKey"]
+):
+    head_tile = 64
+    rows_per_warp = 16
+    elems_per_lane = 2
+    lanes_per_row = head_tile // elems_per_lane
+    num_warps = 8
+    stats_lane_stride = lanes_per_row + 1
+    final_reduce_steps = 3
+    final_reduce_initial_offset = 4
+    tb_size = num_warps * 32
+    compress_ratio = 128
+    state_block_size = 8
+    rcp_ln2 = 1.4426950408889634
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        head_size: int
+        state_width: int
+
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        head_tile = SparseAttnCompressC128Block8Kernel.head_tile
+        rows_per_warp = SparseAttnCompressC128Block8Kernel.rows_per_warp
+        elems_per_lane = SparseAttnCompressC128Block8Kernel.elems_per_lane
+        lanes_per_row = SparseAttnCompressC128Block8Kernel.lanes_per_row
+        num_warps = SparseAttnCompressC128Block8Kernel.num_warps
+        stats_lane_stride = SparseAttnCompressC128Block8Kernel.stats_lane_stride
+        final_reduce_steps = SparseAttnCompressC128Block8Kernel.final_reduce_steps
+        final_reduce_initial_offset = (
+            SparseAttnCompressC128Block8Kernel.final_reduce_initial_offset
         )
-        start = position - Int64(self.compress_ratio - 1)
-        active = slot_id >= Int64(0) and has_req_idx and boundary
-
-        if active:
-            smem = cutlass.utils.SmemAllocator()
-            s_max = smem.allocate_tensor(
-                Float32,
-                cute.make_layout(
-                    (
-                        self.num_warps,
-                        self.lanes_per_row,
-                        self.elems_per_lane,
-                    ),
-                    stride=(
-                        self.stats_lane_stride * self.elems_per_lane,
-                        self.elems_per_lane,
-                        1,
-                    ),
-                ),
-                byte_alignment=4,
+        tb_size = SparseAttnCompressC128Block8Kernel.tb_size
+        compress_ratio = SparseAttnCompressC128Block8Kernel.compress_ratio
+        state_block_size = SparseAttnCompressC128Block8Kernel.state_block_size
+        rcp_ln2 = SparseAttnCompressC128Block8Kernel.rcp_ln2
+        head_dim = compile_key.head_size
+        num_splits = head_dim // head_tile
+        state_width = compile_key.state_width
+
+        @cute.kernel
+        def device_kernel(
+            state_cache: cute.Tensor,
+            token_to_req_indices: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            block_table: cute.Tensor,
+            compressed_kv: cute.Tensor,
+        ):
+            block_id, _, _ = cute.arch.block_idx()
+            tid, _, _ = cute.arch.thread_idx()
+            warp_id = cute.arch.make_warp_uniform(tid // 32)
+            lane_id = tid % 32
+            col_group = lane_id % lanes_per_row
+
+            token_idx = block_id // num_splits
+            split_idx = block_id - token_idx * num_splits
+            col_base = split_idx * head_tile + col_group * elems_per_lane
+
+            position = Int64(0)
+            req_idx = Int32(0)
+            slot_id = Int64(-1)
+            has_position = token_idx < positions.shape[0]
+            has_req_idx = token_idx < token_to_req_indices.shape[0]
+            if lane_id == 0:
+                slot_id = slot_mapping[token_idx]
+            if lane_id == 0 and has_position:
+                position = positions[token_idx]
+            if lane_id == 0 and has_req_idx:
+                req_idx = token_to_req_indices[token_idx]
+            slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
+            position = cute.arch.shuffle_sync(position, offset=0)
+            req_idx = cute.arch.shuffle_sync(req_idx, offset=0)
+            boundary = has_position and (
+                (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
             )
-            s_sum = smem.allocate_tensor(
-                Float32,
-                cute.make_layout(
-                    (
-                        self.num_warps,
-                        self.lanes_per_row,
-                        self.elems_per_lane,
-                    ),
-                    stride=(
-                        self.stats_lane_stride * self.elems_per_lane,
-                        self.elems_per_lane,
-                        1,
+            start = position - Int64(compress_ratio - 1)
+            active = slot_id >= Int64(0) and has_req_idx and boundary
+
+            if active:
+                smem = cutlass.utils.SmemAllocator()
+                s_max = smem.allocate_tensor(
+                    Float32,
+                    cute.make_layout(
+                        (
+                            num_warps,
+                            lanes_per_row,
+                            elems_per_lane,
+                        ),
+                        stride=(
+                            stats_lane_stride * elems_per_lane,
+                            elems_per_lane,
+                            1,
+                        ),
                     ),
-                ),
-                byte_alignment=4,
-            )
-            s_product = smem.allocate_tensor(
-                Float32,
-                cute.make_layout(
-                    (
-                        self.num_warps,
-                        self.lanes_per_row,
-                        self.elems_per_lane,
+                    byte_alignment=4,
+                )
+                s_sum = smem.allocate_tensor(
+                    Float32,
+                    cute.make_layout(
+                        (
+                            num_warps,
+                            lanes_per_row,
+                            elems_per_lane,
+                        ),
+                        stride=(
+                            stats_lane_stride * elems_per_lane,
+                            elems_per_lane,
+                            1,
+                        ),
                     ),
-                    stride=(
-                        self.stats_lane_stride * self.elems_per_lane,
-                        self.elems_per_lane,
-                        1,
+                    byte_alignment=4,
+                )
+                s_product = smem.allocate_tensor(
+                    Float32,
+                    cute.make_layout(
+                        (
+                            num_warps,
+                            lanes_per_row,
+                            elems_per_lane,
+                        ),
+                        stride=(
+                            stats_lane_stride * elems_per_lane,
+                            elems_per_lane,
+                            1,
+                        ),
                     ),
-                ),
-                byte_alignment=4,
-            )
+                    byte_alignment=4,
+                )
 
-            row_layout = cute.make_layout(
-                (self.rows_per_warp, self.elems_per_lane),
-                stride=(self.elems_per_lane, 1),
-            )
-            kv_vals = cute.make_rmem_tensor(row_layout, Float32)
-            score_vals = cute.make_rmem_tensor(row_layout, Float32)
-            local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-            local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-            local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
-
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                local_max[e] = -Float32.inf
-                local_sum[e] = Float32(0.0)
-                local_product[e] = Float32(0.0)
-
-            first_block_index = start // Int64(self.state_block_size)
-            warp_block_index = first_block_index + (warp_id * 2).to(Int64)
-            block0_i32 = Int32(0)
-            block1_i32 = Int32(0)
-            if lane_id == 0:
-                block0_i32 = block_table[req_idx, warp_block_index]
-                block1_i32 = block_table[req_idx, warp_block_index + Int64(1)]
-            block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0)
-            block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0)
+                row_layout = cute.make_layout(
+                    (rows_per_warp, elems_per_lane),
+                    stride=(elems_per_lane, 1),
+                )
+                kv_vals = cute.make_rmem_tensor(row_layout, Float32)
+                score_vals = cute.make_rmem_tensor(row_layout, Float32)
+                local_max = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32)
+                local_product = cute.make_rmem_tensor((elems_per_lane,), Float32)
+
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    local_max[e] = -Float32.inf
+                    local_sum[e] = Float32(0.0)
+                    local_product[e] = Float32(0.0)
+
+                first_block_index = start // Int64(state_block_size)
+                warp_block_index = first_block_index + (warp_id * 2).to(Int64)
+                block0_i32 = Int32(0)
+                block1_i32 = Int32(0)
+                if lane_id == 0:
+                    block0_i32 = block_table[req_idx, warp_block_index]
+                    block1_i32 = block_table[req_idx, warp_block_index + Int64(1)]
+                block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0)
+                block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0)
 
-            cp_f32x2 = cute.make_copy_atom(
-                cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64
-            )
-            final_mask_and_clamp = const_expr(
-                (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
-            )
-            col_tile = col_base.to(Int64) // Int64(self.elems_per_lane)
-            score_col_tile = col_tile + Int64(self.state_width // self.elems_per_lane)
-
-            for i in cutlass.range_constexpr(self.rows_per_warp):
-                block_number_i32 = block0_i32
-                block_offset = Int64(i)
-                if const_expr(i >= self.state_block_size):
-                    block_number_i32 = block1_i32
-                    block_offset = Int64(i - self.state_block_size)
-                row_tensor = state_cache[block_number_i32.to(Int64), block_offset, None]
-                kv_src = cute.local_tile(
-                    row_tensor,
-                    tiler=(self.elems_per_lane,),
-                    coord=(col_tile,),
+                cp_f32x2 = cute.make_copy_atom(
+                    cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64
                 )
-                score_src = cute.local_tile(
-                    row_tensor,
-                    tiler=(self.elems_per_lane,),
-                    coord=(score_col_tile,),
+                final_mask_and_clamp = const_expr(
+                    (cute.arch.WARP_SIZE - num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
                 )
-                cute.copy(cp_f32x2, kv_src, kv_vals[i, None])
-                cute.copy(cp_f32x2, score_src, score_vals[i, None])
+                col_tile = col_base.to(Int64) // Int64(elems_per_lane)
+                score_col_tile = col_tile + Int64(state_width // elems_per_lane)
+
+                for i in cutlass.range_constexpr(rows_per_warp):
+                    block_number_i32 = block0_i32
+                    block_offset = Int64(i)
+                    if const_expr(i >= state_block_size):
+                        block_number_i32 = block1_i32
+                        block_offset = Int64(i - state_block_size)
+                    row_tensor = state_cache[
+                        block_number_i32.to(Int64), block_offset, None
+                    ]
+                    kv_src = cute.local_tile(
+                        row_tensor,
+                        tiler=(elems_per_lane,),
+                        coord=(col_tile,),
+                    )
+                    score_src = cute.local_tile(
+                        row_tensor,
+                        tiler=(elems_per_lane,),
+                        coord=(score_col_tile,),
+                    )
+                    cute.copy(cp_f32x2, kv_src, kv_vals[i, None])
+                    cute.copy(cp_f32x2, score_src, score_vals[i, None])
+
+                    for e in cutlass.range_constexpr(elems_per_lane):
+                        local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e])
 
-                for e in cutlass.range_constexpr(self.elems_per_lane):
-                    local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e])
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    for i in cutlass.range_constexpr(rows_per_warp):
+                        exp_score = cute.math.exp2(
+                            (score_vals[i, e] - local_max[e]) * Float32(rcp_ln2),
+                            fastmath=True,
+                        )
+                        local_sum[e] += exp_score
+                        local_product[e] += kv_vals[i, e] * exp_score
+
+                for e in cutlass.range_constexpr(elems_per_lane):
+                    s_max[warp_id, col_group, e] = local_max[e]
+                    s_sum[warp_id, col_group, e] = local_sum[e]
+                    s_product[warp_id, col_group, e] = local_product[e]
+                cute.arch.sync_threads()
+
+                out_group = tid // num_warps
+                final_lane = tid % num_warps
+                final_groups_per_pass = const_expr(tb_size // num_warps)
+                for pass_idx in cutlass.range_constexpr(
+                    head_tile // final_groups_per_pass
+                ):
+                    out_idx = pass_idx * final_groups_per_pass + out_group
+                    out_lane = out_idx // elems_per_lane
+                    out_elem = out_idx % elems_per_lane
+
+                    local_warp_max = s_max[final_lane, out_lane, out_elem]
+                    global_max = local_warp_max
+                    for step in cutlass.range_constexpr(final_reduce_steps):
+                        offset = const_expr(final_reduce_initial_offset >> step)
+                        global_max = cute.arch.fmax(
+                            global_max,
+                            cute.arch.shuffle_sync_bfly(
+                                global_max,
+                                offset=offset,
+                                mask_and_clamp=final_mask_and_clamp,
+                            ),
+                        )
 
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                for i in cutlass.range_constexpr(self.rows_per_warp):
-                    exp_score = cute.math.exp2(
-                        (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2),
+                    scale = cute.math.exp2(
+                        (local_warp_max - global_max) * Float32(rcp_ln2),
                         fastmath=True,
                     )
-                    local_sum[e] += exp_score
-                    local_product[e] += kv_vals[i, e] * exp_score
-
-            for e in cutlass.range_constexpr(self.elems_per_lane):
-                s_max[warp_id, col_group, e] = local_max[e]
-                s_sum[warp_id, col_group, e] = local_sum[e]
-                s_product[warp_id, col_group, e] = local_product[e]
-            cute.arch.sync_threads()
-
-            out_group = tid // self.num_warps
-            final_lane = tid % self.num_warps
-            final_groups_per_pass = const_expr(self.tb_size // self.num_warps)
-            for pass_idx in cutlass.range_constexpr(
-                self.head_tile // final_groups_per_pass
-            ):
-                out_idx = pass_idx * final_groups_per_pass + out_group
-                out_lane = out_idx // self.elems_per_lane
-                out_elem = out_idx % self.elems_per_lane
-
-                local_warp_max = s_max[final_lane, out_lane, out_elem]
-                global_max = local_warp_max
-                for step in cutlass.range_constexpr(self.final_reduce_steps):
-                    offset = const_expr(self.final_reduce_initial_offset >> step)
-                    global_max = cute.arch.fmax(
-                        global_max,
-                        cute.arch.shuffle_sync_bfly(
-                            global_max,
+                    global_sum = s_sum[final_lane, out_lane, out_elem] * scale
+                    global_product = s_product[final_lane, out_lane, out_elem] * scale
+                    for step in cutlass.range_constexpr(final_reduce_steps):
+                        offset = const_expr(final_reduce_initial_offset >> step)
+                        global_sum += cute.arch.shuffle_sync_bfly(
+                            global_sum,
                             offset=offset,
                             mask_and_clamp=final_mask_and_clamp,
-                        ),
-                    )
-
-                scale = cute.math.exp2(
-                    (local_warp_max - global_max) * Float32(self.rcp_ln2),
-                    fastmath=True,
-                )
-                global_sum = s_sum[final_lane, out_lane, out_elem] * scale
-                global_product = s_product[final_lane, out_lane, out_elem] * scale
-                for step in cutlass.range_constexpr(self.final_reduce_steps):
-                    offset = const_expr(self.final_reduce_initial_offset >> step)
-                    global_sum += cute.arch.shuffle_sync_bfly(
-                        global_sum,
-                        offset=offset,
-                        mask_and_clamp=final_mask_and_clamp,
-                    )
-                    global_product += cute.arch.shuffle_sync_bfly(
-                        global_product,
-                        offset=offset,
-                        mask_and_clamp=final_mask_and_clamp,
-                    )
+                        )
+                        global_product += cute.arch.shuffle_sync_bfly(
+                            global_product,
+                            offset=offset,
+                            mask_and_clamp=final_mask_and_clamp,
+                        )
 
-                if final_lane == 0:
-                    compressed_kv.iterator[
-                        token_idx.to(Int64) * compressed_kv.stride[0]
-                        + (split_idx * self.head_tile + out_idx).to(Int64)
-                    ] = global_product / global_sum
+                    if final_lane == 0:
+                        compressed_kv.iterator[
+                            token_idx.to(Int64) * compressed_kv.stride[0]
+                            + (split_idx * head_tile + out_idx).to(Int64)
+                        ] = global_product / global_sum
+
+        @cute.jit
+        def host_entrypoint(
+            state_cache: cute.Tensor,
+            token_to_req_indices: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            block_table: cute.Tensor,
+            compressed_kv: cute.Tensor,
+            stream: CUstream,
+        ):
+            grid = (slot_mapping.shape[0] * num_splits, 1, 1)
+            device_kernel(
+                state_cache,
+                token_to_req_indices,
+                positions,
+                slot_mapping,
+                block_table,
+                compressed_kv,
+            ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
+
+        return host_entrypoint
+
+    def dispatch(  # type: ignore[override]
+        self, *, head_size: int, state_width: int
+    ) -> CompileKey:
+        return self.CompileKey(head_size=head_size, state_width=state_width)
+
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        hf_config = vllm_config.model_config.hf_config
+        return self._trace_dispatch(self.dispatch)(
+            head_size=hf_config.head_dim,
+            state_width=hf_config.head_dim,
+            enabled=128 in hf_config.compress_ratios,
+            _when=lambda *, enabled: enabled,
+        )
 
-    @cache
-    @staticmethod
-    def compile(
-        head_size: int = 512,
-        state_width: int = 512,
-    ):
+    def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
+        head_size = compile_key.head_size
+        state_width = compile_key.state_width
         if head_size % SparseAttnCompressC128Block8Kernel.head_tile != 0:
             raise ValueError("head_size must be divisible by the 64-wide head tile.")
         num_positions = cute.sym_int()
@@ -1198,7 +1382,6 @@ def compile(
         state_cache_block_size = cute.sym_int()
         block_table_width = cute.sym_int()
         state_cache_width = state_width * 2
-
         state_cache = cute.runtime.make_fake_tensor(
             Float32,
             (num_blocks, state_cache_block_size, state_cache_width),
@@ -1210,12 +1393,16 @@ def compile(
             assumed_align=16,
         )
         token_to_req_indices = make_fake_tensor(
-            Int32, (num_req_indices,), divisibility=4
+            Int32,
+            (num_req_indices,),
+            divisibility=4,
         )
         positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
         slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
         block_table = make_fake_tensor(
-            Int32, (cute.sym_int(), block_table_width), divisibility=1
+            Int32,
+            (cute.sym_int(), block_table_width),
+            divisibility=1,
         )
         compressed_kv = cute.runtime.make_fake_tensor(
             Float32,
@@ -1223,238 +1410,332 @@ def compile(
             stride=(head_size, 1),
             assumed_align=4,
         )
-
-        kernel = SparseAttnCompressC128Block8Kernel(
-            head_size,
-            state_width,
-        )
-        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-        return cute.compile(
-            kernel,
+        return (
             state_cache,
             token_to_req_indices,
             positions,
             slot_mapping,
             block_table,
             compressed_kv,
-            stream,
-            options="--enable-tvm-ffi",
         )
 
-
-class SparseAttnNormRopeStoreKernel:
-    min_scale = 1.0e-4
-
-    def __init__(
-        self,
-        head_size: int,
-        rope_head_dim: int,
-        fp8_max: float,
-        quant_block: int,
-        token_stride: int,
-        scale_dim: int,
-        compress_ratio: int,
-        static_kv_cache_block_size: int,
-    ):
-        self.head_dim = head_size
-        self.rope_dim = rope_head_dim
-        self.nope_dim = head_size - rope_head_dim
-        self.fp8_max = fp8_max
-        self.quant_block = quant_block
-        self.token_stride = token_stride
-        self.scale_dim = scale_dim
-        self.num_warps = head_size // quant_block
-        self.nope_blocks = self.nope_dim // quant_block
-        self.tb_size = head_size // 2
-        self.compress_ratio = compress_ratio
-        self.static_kv_cache_block_size = static_kv_cache_block_size
-
-    @cute.jit
+    @kernel_launcher
     def __call__(
         self,
-        compressed_kv: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        rms_norm_weight: cute.Tensor,
-        rms_norm_eps: Float32,
-        cos_sin_cache: cute.Tensor,
-        k_cache: cute.Tensor,
-        kv_slot_mapping: cute.Tensor,
-        stream: CUstream,
-    ):
-        grid = (slot_mapping.shape[0], 1, 1)
-        self.kernel(
-            compressed_kv,
+        *,
+        state_cache: torch.Tensor,
+        num_actual: int,
+        token_to_req_indices: torch.Tensor,
+        positions: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        block_table: torch.Tensor,
+        head_dim: int,
+        compressed_kv: torch.Tensor | None = None,
+    ) -> CuTeDSLLaunchSpec[SparseAttnCompressC128Block8Kernel.CompileKey]:
+        compile_key = self.dispatch(head_size=head_dim, state_width=head_dim)
+        if compressed_kv is None:
+            compressed_kv = torch.empty(
+                (num_actual, head_dim),
+                dtype=torch.float32,
+                device=state_cache.device,
+            )
+        launch_args = (
+            state_cache,
+            token_to_req_indices,
             positions,
             slot_mapping,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            k_cache,
-            kv_slot_mapping,
-        ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
-
-    @cute.kernel
-    def kernel(
-        self,
-        compressed_kv: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        rms_norm_weight: cute.Tensor,
-        rms_norm_eps: Float32,
-        cos_sin_cache: cute.Tensor,
-        k_cache: cute.Tensor,
-        kv_slot_mapping: cute.Tensor,
-    ):
-        token_idx, _, _ = cute.arch.block_idx()
-        tid, _, _ = cute.arch.thread_idx()
-        warp_id = cute.arch.make_warp_uniform(tid // 32)
-        lane_id = tid % 32
-        elem0 = tid * 2
-
-        position = Int64(0)
-        kv_slot_idx = Int64(-1)
-        has_position = token_idx < positions.shape[0]
-        slot_id = Int64(-1)
-        if lane_id == 0:
-            slot_id = slot_mapping[token_idx]
-        if lane_id == 0 and has_position:
-            position = positions[token_idx]
-        has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
-        if lane_id == 0 and has_kv_slot_idx:
-            kv_slot_idx = kv_slot_mapping[token_idx]
-        slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
-        position = cute.arch.shuffle_sync(position, offset=0)
-        kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0)
-        boundary = has_position and (
-            (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
+            block_table,
+            compressed_kv,
         )
-        active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
-
-        if active:
-            k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
-            k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
-            static_block_size = Int64(self.static_kv_cache_block_size)
-            page = kv_slot_idx // static_block_size
-            kv_offset = kv_slot_idx - page * static_block_size
-            scale_row_offset = static_block_size * Int64(self.token_stride)
-            value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
-            scale_base = (
-                page * k_cache.stride[0]
-                + scale_row_offset
-                + kv_offset * Int64(self.scale_dim)
-            )
-            weight0 = rms_norm_weight[elem0].to(Float32)
-            weight1 = rms_norm_weight[elem0 + 1].to(Float32)
-
-            base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
-            x0 = compressed_kv.iterator[base]
-            x1 = compressed_kv.iterator[base + Int64(1)]
-
-            local_sumsq = x0 * x0 + x1 * x1
-            warp_sum = local_sumsq
-            for step in cutlass.range_constexpr(5):
-                offset = const_expr(16 >> step)
-                warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
-
-            smem = cutlass.utils.SmemAllocator()
-            partial_sums = smem.allocate_tensor(
-                Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
-            )
+        return compile_key, launch_args, compressed_kv
 
-            if lane_id == 0:
-                partial_sums[warp_id] = warp_sum
-            cute.arch.sync_threads()
 
-            total = partial_sums[lane_id % self.num_warps]
-            sum_mask_and_clamp = const_expr(
-                (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
+class SparseAttnNormRopeStoreKernel(
+    VllmCuTeDSLJitKernel["SparseAttnNormRopeStoreKernel.CompileKey"]
+):
+    min_scale = 1.0e-4
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        head_size: int
+        rope_head_dim: int
+        fp8_max: float
+        quant_block: int
+        token_stride: int
+        scale_dim: int
+        kv_block_stride: int
+        compress_ratio: int
+        norm_weight_dtype: type[cutlass.Numeric]
+        kv_cache_block_size: int
+
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        min_scale = SparseAttnNormRopeStoreKernel.min_scale
+        head_dim = compile_key.head_size
+        rope_dim = compile_key.rope_head_dim
+        nope_dim = head_dim - rope_dim
+        fp8_max = compile_key.fp8_max
+        quant_block = compile_key.quant_block
+        token_stride = compile_key.token_stride
+        scale_dim = compile_key.scale_dim
+        num_warps = head_dim // quant_block
+        nope_blocks = nope_dim // quant_block
+        tb_size = head_dim // 2
+        compress_ratio = compile_key.compress_ratio
+        static_kv_cache_block_size = compile_key.kv_cache_block_size
+
+        @cute.kernel
+        def device_kernel(
+            compressed_kv: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            rms_norm_weight: cute.Tensor,
+            rms_norm_eps: Float32,
+            cos_sin_cache: cute.Tensor,
+            k_cache: cute.Tensor,
+            kv_slot_mapping: cute.Tensor,
+        ):
+            token_idx, _, _ = cute.arch.block_idx()
+            tid, _, _ = cute.arch.thread_idx()
+            warp_id = cute.arch.make_warp_uniform(tid // 32)
+            lane_id = tid % 32
+            elem0 = tid * 2
+
+            position = Int64(0)
+            kv_slot_idx = Int64(-1)
+            has_position = token_idx < positions.shape[0]
+            slot_id = Int64(-1)
+            if lane_id == 0:
+                slot_id = slot_mapping[token_idx]
+            if lane_id == 0 and has_position:
+                position = positions[token_idx]
+            has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
+            if lane_id == 0 and has_kv_slot_idx:
+                kv_slot_idx = kv_slot_mapping[token_idx]
+            slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
+            position = cute.arch.shuffle_sync(position, offset=0)
+            kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0)
+            boundary = has_position and (
+                (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
             )
-            for step in cutlass.range_constexpr(3):
-                offset = const_expr(4 >> step)
-                total += cute.arch.shuffle_sync_bfly(
-                    total,
-                    offset,
-                    mask_and_clamp=sum_mask_and_clamp,
+            active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
+
+            if active:
+                k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
+                k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
+                static_block_size = Int64(static_kv_cache_block_size)
+                page = kv_slot_idx // static_block_size
+                kv_offset = kv_slot_idx - page * static_block_size
+                scale_row_offset = static_block_size * Int64(token_stride)
+                value_base = page * k_cache.stride[0] + kv_offset * Int64(token_stride)
+                scale_base = (
+                    page * k_cache.stride[0]
+                    + scale_row_offset
+                    + kv_offset * Int64(scale_dim)
                 )
+                weight0 = rms_norm_weight[elem0].to(Float32)
+                weight1 = rms_norm_weight[elem0 + 1].to(Float32)
 
-            rrms = cute.math.rsqrt(
-                total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
-            )
-            x0 = x0 * rrms * weight0
-            x1 = x1 * rrms * weight1
-
-            if warp_id == self.nope_blocks:
-                pair_idx = lane_id
-                compressed_pos = position - Int64(self.compress_ratio - 1)
-                cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64)
-                cos_v = cos_sin_cache.iterator[cs_base]
-                sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
-                real = x0 * cos_v - x1 * sin_v
-                imag = x0 * sin_v + x1 * cos_v
-                packed = _fp32x2_to_bf16x2(real, imag)
-                out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64)
-                k_cache_u32.iterator[out_base // Int64(4)] = packed
-            else:
-                q_packed = _fp32x2_to_bf16x2(x0, x1)
-                q0, q1 = _bf16x2_to_fp32(q_packed)
-                abs0 = cute.math.absf(q0)
-                abs1 = cute.math.absf(q1)
-                local_absmax = cute.arch.fmax(abs0, abs1)
-                absmax = local_absmax
+                base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
+                x0 = compressed_kv.iterator[base]
+                x1 = compressed_kv.iterator[base + Int64(1)]
+
+                local_sumsq = x0 * x0 + x1 * x1
+                warp_sum = local_sumsq
                 for step in cutlass.range_constexpr(5):
                     offset = const_expr(16 >> step)
-                    absmax = cute.arch.fmax(
-                        absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
-                    )
-                scale_raw = cute.arch.fmax(
-                    Float32(self.min_scale),
-                    absmax / Float32(self.fp8_max),
-                )
-                bits = _recast_val(scale_raw, Uint32)
-                ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
-                inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
-                y0 = cutlass.min(
-                    cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
-                    Float32(self.fp8_max),
+                    warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
+
+                smem = cutlass.utils.SmemAllocator()
+                partial_sums = smem.allocate_tensor(
+                    Float32, cute.make_layout((num_warps,)), byte_alignment=4
                 )
-                y1 = cutlass.min(
-                    cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
-                    Float32(self.fp8_max),
+
+                if lane_id == 0:
+                    partial_sums[warp_id] = warp_sum
+                cute.arch.sync_threads()
+
+                total = partial_sums[lane_id % num_warps]
+                sum_mask_and_clamp = const_expr(
+                    (cute.arch.WARP_SIZE - num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
                 )
-                packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
-                out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to(
-                    Int64
+                for step in cutlass.range_constexpr(3):
+                    offset = const_expr(4 >> step)
+                    total += cute.arch.shuffle_sync_bfly(
+                        total,
+                        offset,
+                        mask_and_clamp=sum_mask_and_clamp,
+                    )
+
+                rrms = cute.math.rsqrt(
+                    total / Float32(head_dim) + rms_norm_eps, fastmath=True
                 )
-                k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
-                if lane_id == 0:
-                    k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8)
-                    if warp_id == 0:
-                        k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
-                            0
+                x0 = x0 * rrms * weight0
+                x1 = x1 * rrms * weight1
+
+                if warp_id == nope_blocks:
+                    pair_idx = lane_id
+                    compressed_pos = position - Int64(compress_ratio - 1)
+                    cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(
+                        Int64
+                    )
+                    cos_v = cos_sin_cache.iterator[cs_base]
+                    sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)]
+                    real = x0 * cos_v - x1 * sin_v
+                    imag = x0 * sin_v + x1 * cos_v
+                    packed = _fp32x2_to_bf16x2(real, imag)
+                    out_base = value_base + Int64(nope_dim) + (lane_id * 4).to(Int64)
+                    k_cache_u32.iterator[out_base // Int64(4)] = packed
+                else:
+                    q_packed = _fp32x2_to_bf16x2(x0, x1)
+                    q0, q1 = _bf16x2_to_fp32(q_packed)
+                    abs0 = cute.math.absf(q0)
+                    abs1 = cute.math.absf(q1)
+                    local_absmax = cute.arch.fmax(abs0, abs1)
+                    absmax = local_absmax
+                    for step in cutlass.range_constexpr(5):
+                        offset = const_expr(16 >> step)
+                        absmax = cute.arch.fmax(
+                            absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
                         )
+                    scale_raw = cute.arch.fmax(
+                        Float32(min_scale),
+                        absmax / Float32(fp8_max),
+                    )
+                    bits = _recast_val(scale_raw, Uint32)
+                    ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
+                    inv_scale = _recast_val(
+                        (Uint32(254) - ue8m0) << Uint32(23), Float32
+                    )
+                    y0 = cutlass.min(
+                        cute.arch.fmax(q0 * inv_scale, Float32(-fp8_max)),
+                        Float32(fp8_max),
+                    )
+                    y1 = cutlass.min(
+                        cute.arch.fmax(q1 * inv_scale, Float32(-fp8_max)),
+                        Float32(fp8_max),
+                    )
+                    packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
+                    out_base = value_base + (warp_id * quant_block + lane_id * 2).to(
+                        Int64
+                    )
+                    k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
+                    if lane_id == 0:
+                        k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(
+                            Uint8
+                        )
+                        if warp_id == 0:
+                            k_cache.iterator[scale_base + Int64(nope_blocks)] = Uint8(0)
+
+        @cute.jit
+        def host_entrypoint(
+            compressed_kv: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            rms_norm_weight: cute.Tensor,
+            rms_norm_eps: Float32,
+            cos_sin_cache: cute.Tensor,
+            k_cache: cute.Tensor,
+            kv_slot_mapping: cute.Tensor,
+            stream: CUstream,
+        ):
+            grid = (slot_mapping.shape[0], 1, 1)
+            device_kernel(
+                compressed_kv,
+                positions,
+                slot_mapping,
+                rms_norm_weight,
+                rms_norm_eps,
+                cos_sin_cache,
+                k_cache,
+                kv_slot_mapping,
+            ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
+
+        return host_entrypoint
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        compress_ratio: int,
+        cache_block_size: int,
+        cache_alignment: int,
+        norm_weight_dtype: type[cutlass.Numeric],
+        head_size: int,
+        rope_head_dim: int,
+        runtime_kv_block_stride: int | None = None,
+    ) -> CompileKey:
+        raw_kv_cache_block_size = cache_block_size // compress_ratio
+        kv_cache_block_size = (
+            raw_kv_cache_block_size if raw_kv_cache_block_size >= 1 else 1
+        )
+        token_stride = 576
+        scale_dim = 8
+        return self.CompileKey(
+            head_size=head_size,
+            rope_head_dim=rope_head_dim,
+            fp8_max=448.0,
+            quant_block=64,
+            token_stride=token_stride,
+            scale_dim=scale_dim,
+            kv_block_stride=(
+                runtime_kv_block_stride
+                if runtime_kv_block_stride is not None
+                else round_up(
+                    kv_cache_block_size * (token_stride + scale_dim),
+                    cache_alignment,
+                )
+            ),
+            compress_ratio=compress_ratio,
+            norm_weight_dtype=norm_weight_dtype,
+            kv_cache_block_size=kv_cache_block_size,
+        )
 
-    @cache
-    @staticmethod
-    def compile(
-        head_size: int = 512,
-        rope_head_dim: int = 64,
-        fp8_max: float = 448.0,
-        quant_block: int = 64,
-        token_stride: int = 576,
-        scale_dim: int = 8,
-        kv_block_stride: int = 74752,
-        compress_ratio: int = 128,
-        norm_weight_dtype: type[cutlass.Numeric] = Float32,
-        static_kv_cache_block_size: int = 0,
-    ):
+    def get_warmup_keys(
+        self,
+        vllm_config: Any,
+        *,
+        k_cache_prefix: str,
+        compress_ratio: int,
+    ) -> list[CompileKey]:
+        hf_config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        k_cache_layer = vllm_config.compilation_config.static_forward_context[
+            k_cache_prefix
+        ]
+        kv_cache = k_cache_layer.kv_cache
+        return self._trace_dispatch(self.dispatch)(
+            compress_ratio=compress_ratio,
+            cache_block_size=cache_config.block_size,
+            cache_alignment=576,
+            norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype),
+            head_size=hf_config.head_dim,
+            rope_head_dim=hf_config.qk_rope_head_dim,
+            runtime_kv_block_stride=kv_cache.stride(0),
+            enabled=(
+                cache_config.block_size > 0
+                and compress_ratio == 128
+                and kv_cache.ndim == 3
+                and cache_config.cache_dtype == "fp8_ds_mla"
+            ),
+            _when=lambda *, enabled: enabled,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
+        head_size = compile_key.head_size
+        rope_head_dim = compile_key.rope_head_dim
+        quant_block = compile_key.quant_block
+        token_stride = compile_key.token_stride
+        scale_dim = compile_key.scale_dim
+        kv_block_stride = compile_key.kv_block_stride
+        norm_weight_dtype = compile_key.norm_weight_dtype
+        static_kv_cache_block_size = compile_key.kv_cache_block_size
         if quant_block != 64:
             raise ValueError(
-                "CuTe DSL sparse-attn store currently requires quant_block=64."
+                "CuTe DSL sparse-attn store currently requires quant_block=64.",
             )
         if rope_head_dim != 64:
             raise ValueError(
-                "CuTe DSL sparse-attn store currently requires rope_head_dim=64."
+                "CuTe DSL sparse-attn store currently requires rope_head_dim=64.",
             )
         if head_size % quant_block != 0:
             raise ValueError("head_size must be divisible by quant_block.")
@@ -1466,14 +1747,13 @@ def compile(
         if static_kv_cache_block_size <= 0:
             raise ValueError(
                 "CuTe DSL sparse-attn store requires a positive static "
-                "kv_cache_block_size."
+                "kv_cache_block_size.",
             )
         num_positions = cute.sym_int()
         num_slots = cute.sym_int()
         num_kv_slots = cute.sym_int()
         max_pos = cute.sym_int()
         num_blocks = cute.sym_int()
-
         compressed_kv = cute.runtime.make_fake_tensor(
             Float32,
             (num_slots, head_size),
@@ -1483,7 +1763,9 @@ def compile(
         positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
         slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
         rms_norm_weight = make_fake_tensor(
-            norm_weight_dtype, (head_size,), divisibility=4
+            norm_weight_dtype,
+            (head_size,),
+            divisibility=4,
         )
         cos_sin_cache = cute.runtime.make_fake_tensor(
             Float32,
@@ -1494,28 +1776,11 @@ def compile(
         k_cache = cute.runtime.make_fake_tensor(
             Uint8,
             (num_blocks, cute.sym_int(), cute.sym_int()),
-            stride=(
-                kv_block_stride,
-                cute.sym_int64(divisibility=8),
-                1,
-            ),
+            stride=(kv_block_stride, cute.sym_int64(divisibility=8), 1),
             assumed_align=16,
         )
         kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
-
-        kernel = SparseAttnNormRopeStoreKernel(
-            head_size,
-            rope_head_dim,
-            fp8_max,
-            quant_block,
-            token_stride,
-            scale_dim,
-            compress_ratio,
-            static_kv_cache_block_size,
-        )
-        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-        return cute.compile(
-            kernel,
+        return (
             compressed_kv,
             positions,
             slot_mapping,
@@ -1524,211 +1789,275 @@ def compile(
             cos_sin_cache,
             k_cache,
             kv_slot_mapping,
-            stream,
-            options="--enable-tvm-ffi",
         )
 
-
-class SparseAttnNormRopeStoreFullKernel:
-    def __init__(
-        self,
-        head_size: int,
-        rope_head_dim: int,
-        fp8_max: float,
-        quant_block: int,
-        token_stride: int,
-        scale_dim: int,
-        compress_ratio: int,
-        store_full_fp8: bool = False,
-    ):
-        # Standalone (not inheriting the #44230-restructured legacy kernel):
-        # set attrs directly so the full-cache C128 path is decoupled.
-        self.head_dim = head_size
-        self.rope_dim = rope_head_dim
-        self.nope_dim = head_size - rope_head_dim
-        self.fp8_max = fp8_max
-        self.quant_block = quant_block
-        self.token_stride = token_stride
-        self.scale_dim = scale_dim
-        self.num_warps = head_size // quant_block
-        self.nope_blocks = self.nope_dim // quant_block
-        self.tb_size = head_size // 2
-        self.compress_ratio = compress_ratio
-        self.store_full_fp8 = store_full_fp8
-
-    @cute.jit
+    @kernel_launcher
     def __call__(
         self,
-        compressed_kv: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        rms_norm_weight: cute.Tensor,
-        rms_norm_eps: Float32,
-        cos_sin_cache: cute.Tensor,
-        k_cache: cute.Tensor,
-        kv_slot_mapping: cute.Tensor,
-        kv_cache_block_size: Int64,
-        fp8_scale: cute.Tensor,
-        stream: CUstream,
-    ):
-        grid = (slot_mapping.shape[0], 1, 1)
-        self.kernel(
+        *,
+        compressed_kv: torch.Tensor,
+        positions: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        rms_norm_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        kv_cache: torch.Tensor,
+        kv_slot_mapping: torch.Tensor,
+        compress_ratio: int,
+        head_dim: int,
+        rope_head_dim: int,
+    ) -> CuTeDSLLaunchSpec[SparseAttnNormRopeStoreKernel.CompileKey]:
+        compile_key = self.dispatch(
+            compress_ratio=compress_ratio,
+            cache_block_size=kv_cache.shape[1] * compress_ratio,
+            cache_alignment=576,
+            norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype),
+            head_size=head_dim,
+            rope_head_dim=rope_head_dim,
+            runtime_kv_block_stride=kv_cache.stride(0),
+        )
+        launch_args = (
             compressed_kv,
             positions,
             slot_mapping,
             rms_norm_weight,
             rms_norm_eps,
             cos_sin_cache,
-            k_cache,
+            kv_cache,
             kv_slot_mapping,
-            kv_cache_block_size,
-            fp8_scale,
-        ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
-
-    @cute.kernel
-    def kernel(
-        self,
-        compressed_kv: cute.Tensor,
-        positions: cute.Tensor,
-        slot_mapping: cute.Tensor,
-        rms_norm_weight: cute.Tensor,
-        rms_norm_eps: Float32,
-        cos_sin_cache: cute.Tensor,
-        k_cache: cute.Tensor,
-        kv_slot_mapping: cute.Tensor,
-        kv_cache_block_size: Int64,
-        fp8_scale: cute.Tensor,
-    ):
-        token_idx, _, _ = cute.arch.block_idx()
-        tid, _, _ = cute.arch.thread_idx()
-        warp_id = cute.arch.make_warp_uniform(tid // 32)
-        lane_id = tid % 32
-        elem0 = tid * 2
-
-        slot_id = slot_mapping[token_idx]
-        has_position = token_idx < positions.shape[0]
-        position = Int64(0)
-        if has_position:
-            position = positions[token_idx]
-        boundary = has_position and (
-            (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
         )
-        has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
-        kv_slot_idx = Int64(-1)
-        if has_kv_slot_idx:
-            kv_slot_idx = kv_slot_mapping[token_idx]
-        active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
-
-        if active:
-            base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
-            x0 = compressed_kv.iterator[base]
-            x1 = compressed_kv.iterator[base + Int64(1)]
-
-            local_sumsq = x0 * x0 + x1 * x1
-            warp_sum = local_sumsq
-            for step in cutlass.range_constexpr(5):
-                offset = const_expr(16 >> step)
-                warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
-
-            smem = cutlass.utils.SmemAllocator()
-            partial_sums = smem.allocate_tensor(
-                Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
-            )
-            rrms_shared = smem.allocate_tensor(
-                Float32, cute.make_layout((1,)), byte_alignment=4
-            )
+        return compile_key, launch_args
 
-            if lane_id == 0:
-                partial_sums[warp_id] = warp_sum
-            cute.arch.sync_threads()
-            if tid == 0:
-                total = Float32(0.0)
-                for i in cutlass.range_constexpr(self.num_warps):
-                    total += partial_sums[i]
-                rrms_shared[0] = cute.math.rsqrt(
-                    total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
-                )
-            cute.arch.sync_threads()
 
-            rrms = rrms_shared[0]
-            x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
-            x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32)
+class SparseAttnNormRopeStoreFullKernel(
+    VllmCuTeDSLJitKernel["SparseAttnNormRopeStoreFullKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        head_size: int
+        rope_head_dim: int
+        fp8_max: float
+        quant_block: int
+        compress_ratio: int
+        store_full_fp8: bool
+        norm_weight_dtype: type[cutlass.Numeric]
 
-            page = kv_slot_idx // kv_cache_block_size
-            kv_offset = kv_slot_idx - page * kv_cache_block_size
-            value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
+    @staticmethod
+    def kernel(compile_key: CompileKey) -> Any:
+        head_dim = compile_key.head_size
+        rope_dim = compile_key.rope_head_dim
+        nope_dim = head_dim - rope_dim
+        fp8_max = compile_key.fp8_max
+        quant_block = compile_key.quant_block
+        num_warps = head_dim // quant_block
+        nope_blocks = nope_dim // quant_block
+        tb_size = head_dim // 2
+        compress_ratio = compile_key.compress_ratio
+        store_full_fp8 = compile_key.store_full_fp8
+
+        @cute.kernel
+        def device_kernel(
+            compressed_kv: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            rms_norm_weight: cute.Tensor,
+            rms_norm_eps: Float32,
+            cos_sin_cache: cute.Tensor,
+            k_cache: cute.Tensor,
+            kv_slot_mapping: cute.Tensor,
+            kv_cache_block_size: Int64,
+            fp8_scale: cute.Tensor,
+        ):
+            token_idx, _, _ = cute.arch.block_idx()
+            tid, _, _ = cute.arch.thread_idx()
+            warp_id = cute.arch.make_warp_uniform(tid // 32)
+            lane_id = tid % 32
+            elem0 = tid * 2
 
-            if const_expr(self.store_full_fp8):
-                k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
-                inv_fp8 = Float32(1.0) / fp8_scale[0]
-                fp8_v0 = x0
-                fp8_v1 = x1
-                if warp_id == self.nope_blocks:
-                    compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
-                        self.compress_ratio
-                    )
-                    pair_idx = lane_id
-                    cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(
-                        Int64
-                    )
-                    cos_v = cos_sin_cache.iterator[cs_base]
-                    sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
-                    fp8_v0 = x0 * cos_v - x1 * sin_v
-                    fp8_v1 = x0 * sin_v + x1 * cos_v
-                fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1)
-                b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16)
-                y0 = cutlass.min(
-                    cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
-                    Float32(self.fp8_max),
+            slot_id = slot_mapping[token_idx]
+            has_position = token_idx < positions.shape[0]
+            position = Int64(0)
+            if has_position:
+                position = positions[token_idx]
+            boundary = has_position and (
+                (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
+            )
+            has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
+            kv_slot_idx = Int64(-1)
+            if has_kv_slot_idx:
+                kv_slot_idx = kv_slot_mapping[token_idx]
+            active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
+
+            if active:
+                base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
+                x0 = compressed_kv.iterator[base]
+                x1 = compressed_kv.iterator[base + Int64(1)]
+
+                local_sumsq = x0 * x0 + x1 * x1
+                warp_sum = local_sumsq
+                for step in cutlass.range_constexpr(5):
+                    offset = const_expr(16 >> step)
+                    warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
+
+                smem = cutlass.utils.SmemAllocator()
+                partial_sums = smem.allocate_tensor(
+                    Float32, cute.make_layout((num_warps,)), byte_alignment=4
                 )
-                y1 = cutlass.min(
-                    cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
-                    Float32(self.fp8_max),
+                rrms_shared = smem.allocate_tensor(
+                    Float32, cute.make_layout((1,)), byte_alignment=4
                 )
-                packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
-                out_base = value_base + elem0.to(Int64)
-                k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
-            else:
-                k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
-                bf16_v0 = x0
-                bf16_v1 = x1
-                if warp_id == self.nope_blocks:
-                    compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
-                        self.compress_ratio
+
+                if lane_id == 0:
+                    partial_sums[warp_id] = warp_sum
+                cute.arch.sync_threads()
+                if tid == 0:
+                    total = Float32(0.0)
+                    for i in cutlass.range_constexpr(num_warps):
+                        total += partial_sums[i]
+                    rrms_shared[0] = cute.math.rsqrt(
+                        total / Float32(head_dim) + rms_norm_eps, fastmath=True
                     )
-                    pair_idx = lane_id
-                    cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(
-                        Int64
+                cute.arch.sync_threads()
+
+                rrms = rrms_shared[0]
+                x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
+                x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32)
+
+                page = kv_slot_idx // kv_cache_block_size
+                kv_offset = kv_slot_idx - page * kv_cache_block_size
+                value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
+
+                if const_expr(store_full_fp8):
+                    k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
+                    inv_fp8 = Float32(1.0) / fp8_scale[0]
+                    fp8_v0 = x0
+                    fp8_v1 = x1
+                    if warp_id == nope_blocks:
+                        compressed_pos = (position // Int64(compress_ratio)) * Int64(
+                            compress_ratio
+                        )
+                        pair_idx = lane_id
+                        cs_base = compressed_pos * cos_sin_cache.stride[
+                            0
+                        ] + pair_idx.to(Int64)
+                        cos_v = cos_sin_cache.iterator[cs_base]
+                        sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)]
+                        fp8_v0 = x0 * cos_v - x1 * sin_v
+                        fp8_v1 = x0 * sin_v + x1 * cos_v
+                    fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1)
+                    b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16)
+                    y0 = cutlass.min(
+                        cutlass.max(b0 * inv_fp8, Float32(-fp8_max)),
+                        Float32(fp8_max),
                     )
-                    cos_v = cos_sin_cache.iterator[cs_base]
-                    sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
-                    bf16_v0 = x0 * cos_v - x1 * sin_v
-                    bf16_v1 = x0 * sin_v + x1 * cos_v
-                bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1)
-                out_base = value_base + (elem0 * 2).to(Int64)
-                k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed
-
-    @cache
-    @staticmethod
-    def compile(
-        head_size: int = 512,
-        rope_head_dim: int = 64,
-        fp8_max: float = 448.0,
-        quant_block: int = 64,
-        token_stride: int = 576,
-        scale_dim: int = 8,
-        kv_block_stride: int = 74752,
-        compress_ratio: int = 128,
-        store_full_fp8: bool = False,
-        norm_weight_dtype: type[cutlass.Numeric] = Float32,
-    ):
+                    y1 = cutlass.min(
+                        cutlass.max(b1 * inv_fp8, Float32(-fp8_max)),
+                        Float32(fp8_max),
+                    )
+                    packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
+                    out_base = value_base + elem0.to(Int64)
+                    k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
+                else:
+                    k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
+                    bf16_v0 = x0
+                    bf16_v1 = x1
+                    if warp_id == nope_blocks:
+                        compressed_pos = (position // Int64(compress_ratio)) * Int64(
+                            compress_ratio
+                        )
+                        pair_idx = lane_id
+                        cs_base = compressed_pos * cos_sin_cache.stride[
+                            0
+                        ] + pair_idx.to(Int64)
+                        cos_v = cos_sin_cache.iterator[cs_base]
+                        sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)]
+                        bf16_v0 = x0 * cos_v - x1 * sin_v
+                        bf16_v1 = x0 * sin_v + x1 * cos_v
+                    bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1)
+                    out_base = value_base + (elem0 * 2).to(Int64)
+                    k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed
+
+        @cute.jit
+        def host_entrypoint(
+            compressed_kv: cute.Tensor,
+            positions: cute.Tensor,
+            slot_mapping: cute.Tensor,
+            rms_norm_weight: cute.Tensor,
+            rms_norm_eps: Float32,
+            cos_sin_cache: cute.Tensor,
+            k_cache: cute.Tensor,
+            kv_slot_mapping: cute.Tensor,
+            kv_cache_block_size: Int64,
+            fp8_scale: cute.Tensor,
+            stream: CUstream,
+        ):
+            grid = (slot_mapping.shape[0], 1, 1)
+            device_kernel(
+                compressed_kv,
+                positions,
+                slot_mapping,
+                rms_norm_weight,
+                rms_norm_eps,
+                cos_sin_cache,
+                k_cache,
+                kv_slot_mapping,
+                kv_cache_block_size,
+                fp8_scale,
+            ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
+
+        return host_entrypoint
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        store_full_fp8: bool,
+        norm_weight_dtype: type[cutlass.Numeric],
+        head_size: int,
+        rope_head_dim: int,
+        compress_ratio: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            head_size=head_size,
+            rope_head_dim=rope_head_dim,
+            compress_ratio=compress_ratio,
+            fp8_max=448.0,
+            quant_block=64,
+            store_full_fp8=store_full_fp8,
+            norm_weight_dtype=norm_weight_dtype,
+        )
+
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        hf_config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        return self._trace_dispatch(self.dispatch)(
+            compress_ratio=128,
+            store_full_fp8=(
+                cache_config.cache_dtype.startswith("fp8")
+                and cache_config.cache_dtype != "fp8_ds_mla"
+            ),
+            norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype),
+            head_size=hf_config.head_dim,
+            rope_head_dim=hf_config.qk_rope_head_dim,
+            enabled=(
+                cache_config.block_size > 0
+                and 128 in hf_config.compress_ratios
+                and cache_config.cache_dtype != "fp8_ds_mla"
+            ),
+            _when=lambda *, enabled: enabled,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
+        head_size = compile_key.head_size
+        rope_head_dim = compile_key.rope_head_dim
+        quant_block = compile_key.quant_block
+        norm_weight_dtype = compile_key.norm_weight_dtype
         if quant_block != 64:
             raise ValueError(
-                "CuTe DSL sparse-attn store currently requires quant_block=64."
+                "CuTe DSL sparse-attn store currently requires quant_block=64.",
             )
         if rope_head_dim != 64:
             raise ValueError(
-                "CuTe DSL sparse-attn store currently requires rope_head_dim=64."
+                "CuTe DSL sparse-attn store currently requires rope_head_dim=64.",
             )
         if head_size % quant_block != 0:
             raise ValueError("head_size must be divisible by quant_block.")
@@ -1737,7 +2066,6 @@ def compile(
         num_kv_slots = cute.sym_int()
         max_pos = cute.sym_int()
         num_blocks = cute.sym_int()
-
         compressed_kv = cute.runtime.make_fake_tensor(
             Float32,
             (num_slots, head_size),
@@ -1747,7 +2075,9 @@ def compile(
         positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
         slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
         rms_norm_weight = make_fake_tensor(
-            norm_weight_dtype, (head_size,), divisibility=4
+            norm_weight_dtype,
+            (head_size,),
+            divisibility=4,
         )
         cos_sin_cache = cute.runtime.make_fake_tensor(
             Float32,
@@ -1758,29 +2088,12 @@ def compile(
         k_cache = cute.runtime.make_fake_tensor(
             Uint8,
             (num_blocks, cute.sym_int(), cute.sym_int()),
-            stride=(
-                cute.sym_int64(divisibility=16),
-                cute.sym_int64(divisibility=8),
-                1,
-            ),
+            stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1),
             assumed_align=16,
         )
         kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
         fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1)
-
-        kernel = SparseAttnNormRopeStoreFullKernel(
-            head_size,
-            rope_head_dim,
-            fp8_max,
-            quant_block,
-            token_stride,
-            scale_dim,
-            compress_ratio,
-            store_full_fp8,
-        )
-        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-        return cute.compile(
-            kernel,
+        return (
             compressed_kv,
             positions,
             slot_mapping,
@@ -1791,83 +2104,57 @@ def compile(
             kv_slot_mapping,
             Int64(0),
             fp8_scale,
-            stream,
-            options="--enable-tvm-ffi",
         )
 
-
-def compile_split_sparse_attn_cutedsl(
-    head_size: int,
-    state_width: int,
-    block_size: int,
-    rope_head_dim: int,
-    fp8_max: float,
-    quant_block: int,
-    token_stride: int,
-    scale_dim: int,
-    kv_cache_block_size: int,
-    kv_block_stride: int,
-    compress_ratio: int,
-    overlap: bool,
-    rms_norm_weight_dtype: torch.dtype,
-    store_full_kv: bool = False,
-    store_full_fp8: bool = False,
-):
-    if not (
-        head_size == 512
-        and state_width == head_size
-        and compress_ratio == 128
-        and not overlap
-        and block_size == 8
-    ):
-        raise ValueError(
-            "CuTe DSL split sparse-attn wrapper only supports the real "
-            "DeepSeek V4 C128 layout: head_size=512, state_width=512, "
-            "compress_ratio=128, overlap=False, block_size=8."
-        )
-    compress = SparseAttnCompressC128Block8Kernel.compile(
-        head_size=head_size,
-        state_width=state_width,
-    )
-    norm_weight_dtype = _TORCH_TO_CUTE[rms_norm_weight_dtype]
-    if store_full_kv:
-        # FlashInfer contiguous bf16/fp8 cache: standalone full-cache store.
-        store = SparseAttnNormRopeStoreFullKernel.compile(
-            head_size=head_size,
-            rope_head_dim=rope_head_dim,
-            fp8_max=fp8_max,
-            quant_block=quant_block,
-            token_stride=token_stride,
-            scale_dim=scale_dim,
-            kv_block_stride=kv_block_stride,
+    @kernel_launcher
+    def __call__(
+        self,
+        *,
+        compressed_kv: torch.Tensor,
+        positions: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        rms_norm_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        kv_cache: torch.Tensor,
+        kv_slot_mapping: torch.Tensor,
+        compress_ratio: int,
+        store_full_fp8: bool,
+        fp8_scale: torch.Tensor | None,
+        head_dim: int,
+        rope_head_dim: int,
+    ) -> CuTeDSLLaunchSpec[SparseAttnNormRopeStoreFullKernel.CompileKey]:
+        compile_key = self.dispatch(
             compress_ratio=compress_ratio,
             store_full_fp8=store_full_fp8,
-            norm_weight_dtype=norm_weight_dtype,
+            norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype),
+            head_size=head_dim,
+            rope_head_dim=rope_head_dim,
         )
-    else:
-        store = SparseAttnNormRopeStoreKernel.compile(
-            head_size,
-            rope_head_dim,
-            fp8_max,
-            quant_block,
-            token_stride,
-            scale_dim,
-            kv_block_stride,
-            compress_ratio,
-            norm_weight_dtype,
-            kv_cache_block_size,
+        if fp8_scale is None:
+            fp8_scale = torch.ones(1, dtype=torch.float32, device=kv_cache.device)
+        launch_args = (
+            compressed_kv,
+            positions,
+            slot_mapping,
+            rms_norm_weight,
+            rms_norm_eps,
+            cos_sin_cache,
+            kv_cache.view(torch.uint8),
+            kv_slot_mapping,
+            kv_cache.shape[1],
+            fp8_scale,
         )
-    return compress, store
+        return compile_key, launch_args
 
 
-def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
+def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
     state_cache: torch.Tensor,
     token_to_req_indices: torch.Tensor,
     positions: torch.Tensor,
     slot_mapping: torch.Tensor,
     block_table: torch.Tensor,
     block_size: int,
-    compressed_kv: torch.Tensor,
     rms_norm_weight: torch.Tensor,
     rms_norm_eps: float,
     cos_sin_cache: torch.Tensor,
@@ -1876,102 +2163,92 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
     kv_cache_block_size: int,
     kv_block_stride: int,
     head_size: int = 512,
-    state_width: int = 512,
+    state_width: int = 1024,
     rope_head_dim: int = 64,
     fp8_max: float = 448.0,
     quant_block: int = 64,
     token_stride: int = 576,
     scale_dim: int = 8,
-    compress_ratio: int = 128,
-    overlap: bool = False,
+    compress_ratio: int = 4,
+    overlap: bool = True,
     store_full_kv: bool = False,
     store_full_fp8: bool = False,
     fp8_scale: torch.Tensor | None = None,
 ) -> None:
+    del fp8_max, quant_block, token_stride, scale_dim, kv_block_stride
+    if not (
+        head_size == 512
+        and state_width == 2 * head_size
+        and compress_ratio == 4
+        and overlap
+    ):
+        raise ValueError(
+            "CuTe DSL fused sparse-attn wrapper only supports the real "
+            "DeepSeek V4 C4 layout: head_size=512, state_width=1024, "
+            "compress_ratio=4, overlap=True."
+        )
     if k_cache.ndim != 3:
         raise ValueError(
-            "CuTe DSL sparse-attn store expects the real DeepSeek V4 "
+            "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 "
             f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
         )
-    if not store_full_kv and kv_cache_block_size != k_cache.shape[1]:
+    if kv_cache_block_size != k_cache.shape[1]:
         raise ValueError(
-            "CuTe DSL split sparse-attn wrapper expected kv_cache_block_size "
+            "CuTe DSL fused sparse-attn wrapper expected kv_cache_block_size "
             f"to match k_cache.shape[1], got {kv_cache_block_size} and "
             f"{k_cache.shape[1]}."
         )
     if positions.numel() == 0:
         return
-    if rms_norm_weight.dtype not in _TORCH_TO_CUTE:
-        raise ValueError(
-            "CuTe DSL sparse-attn store supports rms_norm_weight dtype "
-            f"bf16/fp32, got {rms_norm_weight.dtype}."
-        )
     if store_full_fp8 and not store_full_kv:
         raise ValueError("store_full_fp8 requires store_full_kv.")
-    compress, store = compile_split_sparse_attn_cutedsl(
-        head_size,
-        state_width,
-        block_size,
-        rope_head_dim,
-        fp8_max,
-        quant_block,
-        token_stride,
-        scale_dim,
-        kv_cache_block_size,
-        kv_block_stride,
-        compress_ratio,
-        overlap,
-        rms_norm_weight.dtype,
-        store_full_kv=store_full_kv,
-        store_full_fp8=store_full_fp8,
-    )
-    compress(
-        state_cache,
-        token_to_req_indices,
-        positions,
-        slot_mapping,
-        block_table,
-        compressed_kv,
-    )
 
     if store_full_kv:
-        # Byte-addressed contiguous cache; block size + per-tensor scale are
-        # passed at call time (not baked into compile).
-        if fp8_scale is None:
-            fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device)
-        store(
-            compressed_kv,
-            positions,
-            slot_mapping,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            k_cache.view(torch.uint8),
-            kv_slot_mapping,
-            kv_cache_block_size,
-            fp8_scale,
+        _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL(
+            state_cache=state_cache,
+            token_to_req_indices=token_to_req_indices,
+            positions=positions,
+            slot_mapping=slot_mapping,
+            block_table=block_table,
+            block_size=block_size,
+            rms_norm_weight=rms_norm_weight,
+            rms_norm_eps=rms_norm_eps,
+            cos_sin_cache=cos_sin_cache,
+            kv_cache=k_cache,
+            kv_slot_mapping=kv_slot_mapping,
+            compress_ratio=compress_ratio,
+            store_full_fp8=store_full_fp8,
+            fp8_scale=fp8_scale,
+            head_dim=head_size,
+            rope_head_dim=rope_head_dim,
         )
         return
-
-    store(
-        compressed_kv,
-        positions,
-        slot_mapping,
-        rms_norm_weight,
-        rms_norm_eps,
-        cos_sin_cache,
-        k_cache,
-        kv_slot_mapping,
+    _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL(
+        state_cache=state_cache,
+        token_to_req_indices=token_to_req_indices,
+        positions=positions,
+        slot_mapping=slot_mapping,
+        block_table=block_table,
+        block_size=block_size,
+        rms_norm_weight=rms_norm_weight,
+        rms_norm_eps=rms_norm_eps,
+        cos_sin_cache=cos_sin_cache,
+        kv_cache=k_cache,
+        kv_slot_mapping=kv_slot_mapping,
+        compress_ratio=compress_ratio,
+        head_dim=head_size,
+        rope_head_dim=rope_head_dim,
     )
 
 
-def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
+def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
     state_cache: torch.Tensor,
     token_to_req_indices: torch.Tensor,
     positions: torch.Tensor,
     slot_mapping: torch.Tensor,
     block_table: torch.Tensor,
     block_size: int,
+    compressed_kv: torch.Tensor,
     rms_norm_weight: torch.Tensor,
     rms_norm_eps: float,
     cos_sin_cache: torch.Tensor,
@@ -1980,185 +2257,247 @@ def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
     kv_cache_block_size: int,
     kv_block_stride: int,
     head_size: int = 512,
-    state_width: int = 1024,
+    state_width: int = 512,
     rope_head_dim: int = 64,
     fp8_max: float = 448.0,
     quant_block: int = 64,
     token_stride: int = 576,
     scale_dim: int = 8,
-    compress_ratio: int = 4,
-    overlap: bool = True,
+    compress_ratio: int = 128,
+    overlap: bool = False,
     store_full_kv: bool = False,
     store_full_fp8: bool = False,
     fp8_scale: torch.Tensor | None = None,
 ) -> None:
-    if positions.numel() == 0:
-        return
-    norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype)
-    if norm_weight_dtype is None:
+    del fp8_max, quant_block, token_stride, scale_dim, kv_block_stride
+    if not (
+        head_size == 512
+        and state_width == head_size
+        and compress_ratio == 128
+        and not overlap
+        and block_size == 8
+    ):
         raise ValueError(
-            "CuTe DSL sparse-attn fused store supports rms_norm_weight dtype "
-            f"bf16/fp32, got {rms_norm_weight.dtype}."
+            "CuTe DSL split sparse-attn wrapper only supports the real "
+            "DeepSeek V4 C128 layout: head_size=512, state_width=512, "
+            "compress_ratio=128, overlap=False, block_size=8."
         )
     if k_cache.ndim != 3:
         raise ValueError(
-            "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 "
+            "CuTe DSL sparse-attn store expects the real DeepSeek V4 "
             f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
         )
+    if not store_full_kv and kv_cache_block_size != k_cache.shape[1]:
+        raise ValueError(
+            "CuTe DSL split sparse-attn wrapper expected kv_cache_block_size "
+            f"to match k_cache.shape[1], got {kv_cache_block_size} and "
+            f"{k_cache.shape[1]}."
+        )
+    if positions.numel() == 0:
+        return
     if store_full_fp8 and not store_full_kv:
         raise ValueError("store_full_fp8 requires store_full_kv.")
+
+    _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL(
+        state_cache=state_cache,
+        num_actual=compressed_kv.shape[0],
+        token_to_req_indices=token_to_req_indices,
+        positions=positions,
+        slot_mapping=slot_mapping,
+        block_table=block_table,
+        head_dim=head_size,
+        compressed_kv=compressed_kv,
+    )
     if store_full_kv:
-        # FlashInfer contiguous bf16/fp8 cache: byte-addressed full-cache C4 store.
-        if fp8_scale is None:
-            fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device)
-        compiled = SparseAttnCompressNormRopeStoreFullC4Kernel.compile(
-            head_size=head_size,
-            state_width=state_width,
-            rope_head_dim=rope_head_dim,
-            fp8_max=fp8_max,
-            quant_block=quant_block,
-            token_stride=token_stride,
-            scale_dim=scale_dim,
-            kv_block_stride=kv_block_stride,
+        _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL(
+            compressed_kv=compressed_kv,
+            positions=positions,
+            slot_mapping=slot_mapping,
+            rms_norm_weight=rms_norm_weight,
+            rms_norm_eps=rms_norm_eps,
+            cos_sin_cache=cos_sin_cache,
+            kv_cache=k_cache,
+            kv_slot_mapping=kv_slot_mapping,
             compress_ratio=compress_ratio,
-            overlap=overlap,
             store_full_fp8=store_full_fp8,
-            norm_weight_dtype=norm_weight_dtype,
-        )
-        compiled(
-            state_cache,
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_size,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            k_cache.view(torch.uint8),
-            kv_slot_mapping,
-            kv_cache_block_size,
-            fp8_scale,
+            fp8_scale=fp8_scale,
+            head_dim=head_size,
+            rope_head_dim=rope_head_dim,
         )
         return
-
-    compiled = SparseAttnCompressNormRopeStoreC4Kernel.compile(
-        head_size=head_size,
-        state_width=state_width,
-        rope_head_dim=rope_head_dim,
-        fp8_max=fp8_max,
-        quant_block=quant_block,
-        token_stride=token_stride,
-        scale_dim=scale_dim,
-        kv_block_stride=kv_block_stride,
+    _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL(
+        compressed_kv=compressed_kv,
+        positions=positions,
+        slot_mapping=slot_mapping,
+        rms_norm_weight=rms_norm_weight,
+        rms_norm_eps=rms_norm_eps,
+        cos_sin_cache=cos_sin_cache,
+        kv_cache=k_cache,
+        kv_slot_mapping=kv_slot_mapping,
         compress_ratio=compress_ratio,
-        overlap=overlap,
-        norm_weight_dtype=norm_weight_dtype,
-    )
-    compiled(
-        state_cache,
-        token_to_req_indices,
-        positions,
-        slot_mapping,
-        block_table,
-        block_size,
-        rms_norm_weight,
-        rms_norm_eps,
-        cos_sin_cache,
-        k_cache,
-        kv_slot_mapping,
-        kv_cache_block_size,
+        head_dim=head_size,
+        rope_head_dim=rope_head_dim,
     )
 
 
-def compress_norm_rope_store_cutedsl(
-    state_cache: torch.Tensor,
-    num_actual: int,
-    token_to_req_indices: torch.Tensor,
-    positions: torch.Tensor,
-    slot_mapping: torch.Tensor,
-    block_table: torch.Tensor,
-    block_size: int,
-    state_width: int,
-    cos_sin_cache: torch.Tensor,
-    kv_cache: torch.Tensor,
-    k_cache_metadata: Any,
-    pdl_kwargs: dict,
-    head_dim: int,
-    rope_head_dim: int,
-    compress_ratio: int,
-    overlap: bool,
-    use_fp4_cache: bool,
-    rms_norm_weight: torch.Tensor,
-    rms_norm_eps: float,
-    quant_block: int,
-    token_stride: int,
-    scale_dim: int,
-    store_full_kv: bool = False,
-    store_full_fp8: bool = False,
-    fp8_scale: torch.Tensor | None = None,
-) -> None:
-    if compress_ratio == 4:
-        # For C4A, the single fused kernel is faster than the two-kernel version.
-        fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
-            state_cache,
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_size,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            kv_cache,
-            k_cache_metadata.slot_mapping,
-            kv_cache.shape[1],  # paged KV cache block size
-            kv_cache.stride(0),
-            head_size=head_dim,
-            state_width=state_width,
-            rope_head_dim=rope_head_dim,
-            fp8_max=448.0,
-            quant_block=quant_block,
-            token_stride=token_stride,
-            scale_dim=scale_dim,
-            compress_ratio=compress_ratio,
-            overlap=overlap,
-            store_full_kv=store_full_kv,
-            store_full_fp8=store_full_fp8,
-            fp8_scale=fp8_scale,
-        )
-    else:
-        # For C128, the two-kernel version is faster than the single fused kernel.
-        compressed_kv = torch.empty(
-            (num_actual, head_dim),
-            dtype=torch.float32,
-            device=state_cache.device,
+class SparseAttnCompressorCuteDSL:
+    def __init__(
+        self,
+        *,
+        c4_store: SparseAttnCompressNormRopeStoreC4Kernel,
+        c4_store_full: SparseAttnCompressNormRopeStoreFullC4Kernel,
+        c128_compress: SparseAttnCompressC128Block8Kernel,
+        c128_store: SparseAttnNormRopeStoreKernel,
+        c128_store_full: SparseAttnNormRopeStoreFullKernel,
+    ) -> None:
+        self.c4_store = c4_store
+        self.c4_store_full = c4_store_full
+        self.c128_compress = c128_compress
+        self.c128_store = c128_store
+        self.c128_store_full = c128_store_full
+
+    def __call__(
+        self,
+        state_cache: torch.Tensor,
+        num_actual: int,
+        token_to_req_indices: torch.Tensor,
+        positions: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        state_width: int,
+        cos_sin_cache: torch.Tensor,
+        kv_cache: torch.Tensor,
+        k_cache_metadata: Any,
+        pdl_kwargs: dict,
+        head_dim: int,
+        rope_head_dim: int,
+        compress_ratio: int,
+        overlap: bool,
+        use_fp4_cache: bool,
+        rms_norm_weight: torch.Tensor,
+        rms_norm_eps: float,
+        quant_block: int,
+        token_stride: int,
+        scale_dim: int,
+        store_full_kv: bool = False,
+        store_full_fp8: bool = False,
+        fp8_scale: torch.Tensor | None = None,
+    ) -> None:
+        del pdl_kwargs, state_width, overlap, use_fp4_cache
+        del quant_block, token_stride, scale_dim
+        if positions.numel() == 0:
+            return
+        if kv_cache.ndim != 3:
+            raise ValueError(
+                "CuTe DSL sparse-attn compressor expects the real DeepSeek V4 "
+                f"3D kv_cache layout, got ndim={kv_cache.ndim}.",
+            )
+        if store_full_fp8 and (not store_full_kv):
+            raise ValueError("store_full_fp8 requires store_full_kv.")
+        if compress_ratio == 4:
+            if store_full_kv:
+                self.c4_store_full(
+                    state_cache=state_cache,
+                    token_to_req_indices=token_to_req_indices,
+                    positions=positions,
+                    slot_mapping=slot_mapping,
+                    block_table=block_table,
+                    block_size=block_size,
+                    rms_norm_weight=rms_norm_weight,
+                    rms_norm_eps=rms_norm_eps,
+                    cos_sin_cache=cos_sin_cache,
+                    kv_cache=kv_cache,
+                    kv_slot_mapping=k_cache_metadata.slot_mapping,
+                    compress_ratio=compress_ratio,
+                    store_full_fp8=store_full_fp8,
+                    fp8_scale=fp8_scale,
+                    head_dim=head_dim,
+                    rope_head_dim=rope_head_dim,
+                )
+            else:
+                self.c4_store(
+                    state_cache=state_cache,
+                    token_to_req_indices=token_to_req_indices,
+                    positions=positions,
+                    slot_mapping=slot_mapping,
+                    block_table=block_table,
+                    block_size=block_size,
+                    rms_norm_weight=rms_norm_weight,
+                    rms_norm_eps=rms_norm_eps,
+                    cos_sin_cache=cos_sin_cache,
+                    kv_cache=kv_cache,
+                    kv_slot_mapping=k_cache_metadata.slot_mapping,
+                    compress_ratio=compress_ratio,
+                    head_dim=head_dim,
+                    rope_head_dim=rope_head_dim,
+                )
+            return
+        if compress_ratio != 128:
+            raise ValueError(
+                "CuTe DSL sparse-attn compressor supports compress_ratio 4 or "
+                f"128, got {compress_ratio}.",
+            )
+        if block_size != self.c128_compress.state_block_size:
+            raise ValueError(
+                "CuTe DSL C128 sparse-attn compressor requires state-cache "
+                f"block_size={self.c128_compress.state_block_size}, got "
+                f"{block_size}."
+            )
+        compressed_kv = self.c128_compress(
+            state_cache=state_cache,
+            num_actual=num_actual,
+            token_to_req_indices=token_to_req_indices,
+            positions=positions,
+            slot_mapping=slot_mapping,
+            block_table=block_table,
+            head_dim=head_dim,
         )
-        split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
-            state_cache,
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_size,
-            compressed_kv,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            kv_cache,
-            k_cache_metadata.slot_mapping,
-            kv_cache.shape[1],  # paged KV cache block size
-            kv_cache.stride(0),
-            head_size=head_dim,
-            state_width=state_width,
-            rope_head_dim=rope_head_dim,
-            fp8_max=448.0,
-            quant_block=quant_block,
-            token_stride=token_stride,
-            scale_dim=scale_dim,
+        if store_full_kv:
+            self.c128_store_full(
+                compressed_kv=compressed_kv,
+                positions=positions,
+                slot_mapping=slot_mapping,
+                rms_norm_weight=rms_norm_weight,
+                rms_norm_eps=rms_norm_eps,
+                cos_sin_cache=cos_sin_cache,
+                kv_cache=kv_cache,
+                kv_slot_mapping=k_cache_metadata.slot_mapping,
+                compress_ratio=compress_ratio,
+                store_full_fp8=store_full_fp8,
+                fp8_scale=fp8_scale,
+                head_dim=head_dim,
+                rope_head_dim=rope_head_dim,
+            )
+            return
+        self.c128_store(
+            compressed_kv=compressed_kv,
+            positions=positions,
+            slot_mapping=slot_mapping,
+            rms_norm_weight=rms_norm_weight,
+            rms_norm_eps=rms_norm_eps,
+            cos_sin_cache=cos_sin_cache,
+            kv_cache=kv_cache,
+            kv_slot_mapping=k_cache_metadata.slot_mapping,
             compress_ratio=compress_ratio,
-            overlap=overlap,
-            store_full_kv=store_full_kv,
-            store_full_fp8=store_full_fp8,
-            fp8_scale=fp8_scale,
+            head_dim=head_dim,
+            rope_head_dim=rope_head_dim,
         )
+
+
+_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL = (
+    SparseAttnCompressNormRopeStoreC4Kernel()
+)
+_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL = (
+    SparseAttnCompressNormRopeStoreFullC4Kernel()
+)
+_SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL = SparseAttnCompressC128Block8Kernel()
+_SPARSE_ATTN_NORM_ROPE_STORE_KERNEL = SparseAttnNormRopeStoreKernel()
+_SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL = SparseAttnNormRopeStoreFullKernel()
+_SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL = SparseAttnCompressorCuteDSL(
+    c4_store=_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL,
+    c4_store_full=_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL,
+    c128_compress=_SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL,
+    c128_store=_SPARSE_ATTN_NORM_ROPE_STORE_KERNEL,
+    c128_store_full=_SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL,
+)
diff --git a/vllm/models/deepseek_v4/nvidia/vl_model.py b/vllm/models/deepseek_v4/nvidia/vl_model.py
index bd40b7bce9e7..a66217f2fe65 100644
--- a/vllm/models/deepseek_v4/nvidia/vl_model.py
+++ b/vllm/models/deepseek_v4/nvidia/vl_model.py
@@ -1,333 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp).
+"""Compatibility imports for the platform-neutral DeepSeek-V4 vision model."""
 
-Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``:
-
-- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE
-  sentinel positions; four learned vectors (``image_start`` / ``image_pad`` /
-  ``image_newline`` / ``image_end``) fill the remaining sentinel positions.
-- Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor
-  in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved
-  in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>``
-  (see ``common/vision.py`` for the tower itself).
-- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before
-  its hyper-connection stream expansion. Raw ``input_ids`` still flow into
-  the model so the MoE router can apply ``bias_vl`` to image tokens
-  (``requires_raw_input_tokens``).
-"""
-
-from collections.abc import Iterable
-
-import torch
-from torch import nn
-
-from vllm.model_executor.models.interfaces import (
-    MultiModalEmbeddings,
-    SupportsEagle3,
-    SupportsMultiModal,
-    SupportsPP,
+from ..common.vl_model import (
+    DeepseekV4ForConditionalGeneration,
 )
-from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
-    WeightsMapper,
-    init_vllm_registered_model,
-    maybe_prefix,
+from ..common.vl_model import (
+    _make_deepseek_v4_vl_weights_mapper as _make_common_vl_weights_mapper,
 )
-from vllm.multimodal import MULTIMODAL_REGISTRY
-
-from ..common.mm_preprocess import (
-    IMAGE_PLACEHOLDER,
-    IMAGE_SENTINEL_BASE_ID,
-    DeepseekV4VLDummyInputsBuilder,
-    DeepseekV4VLMultiModalProcessor,
-    DeepseekV4VLProcessingInfo,
-    image_sentinel_mask,
-)
-from ..common.vision import DeepseekV4Aligner, DeepseekV4ViT
 from .model import _make_deepseek_v4_weights_mapper
 
 
-def _make_deepseek_v4_vl_weights_mapper(
-    expert_dtype: str, image_enabled: bool
-) -> WeightsMapper:
-    """Text-checkpoint mapping rules re-rooted under ``language_model.``."""
-    base = _make_deepseek_v4_weights_mapper(expert_dtype)
-    orig_to_new_prefix: dict[str, str | None] = {
-        "layers.": "language_model.model.layers.",
-        "embed.": "language_model.model.embed.",
-        "norm.": "language_model.model.norm.",
-        "hc_head": "language_model.model.hc_head",
-        "mtp.": "language_model.model.mtp.",
-    }
-    if not image_enabled:
-        orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None})
-    return WeightsMapper(
-        orig_to_new_prefix=orig_to_new_prefix,
-        orig_to_new_regex=base.orig_to_new_regex,
-        orig_to_new_suffix={
-            "head.weight": "language_model.lm_head.weight",
-            "embed.weight": "embed_tokens.weight",
-            ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
-        },
-        orig_to_new_substr={
-            ".shared_experts.w2": ".shared_experts.down_proj",
-            # The MTP/DSpark draft heads are not supported for the vision
-            # variant; drop their weights.
-            "mtp.": None,
-        },
+def _make_deepseek_v4_vl_weights_mapper(expert_dtype: str, image_enabled: bool):
+    """Retain the original NVIDIA helper signature for downstream imports."""
+    return _make_common_vl_weights_mapper(
+        _make_deepseek_v4_weights_mapper(expert_dtype), image_enabled
     )
 
 
-@MULTIMODAL_REGISTRY.register_processor(
-    DeepseekV4VLMultiModalProcessor,
-    info=DeepseekV4VLProcessingInfo,
-    dummy_inputs=DeepseekV4VLDummyInputsBuilder,
-)
-class DeepseekV4ForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3
-):
-    """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower.
-
-    ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
-    delegates through ``language_model`` via the protocol defaults.
-    """
-
-    # The MoE router needs raw token ids to detect image sentinel tokens
-    # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl.
-    requires_raw_input_tokens = True
-
-    @classmethod
-    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
-        if modality == "image":
-            return IMAGE_PLACEHOLDER
-        raise ValueError(f"Unsupported modality: {modality!r}")
-
-    def __init__(self, *, vllm_config, prefix: str = "") -> None:
-        super().__init__()
-        model_config = vllm_config.model_config
-        config = model_config.hf_config
-        self.config = config
-        self.multimodal_config = model_config.multimodal_config
-        assert self.multimodal_config is not None
-
-        image_enabled = (
-            config.vision_n_layers > 0
-            and self.multimodal_config.get_limit_per_prompt("image") > 0
-        )
-        with self._mark_tower_model(vllm_config, {"image"}):
-            self.vision: DeepseekV4ViT | None = None
-            self.aligner: DeepseekV4Aligner | None = None
-            self.image_start: nn.Parameter | None = None
-            self.image_end: nn.Parameter | None = None
-            self.image_newline: nn.Parameter | None = None
-            self.image_pad: nn.Parameter | None = None
-            if image_enabled:
-                self.vision = DeepseekV4ViT(config)
-                self.aligner = DeepseekV4Aligner(config)
-                for name in (
-                    "image_start",
-                    "image_end",
-                    "image_newline",
-                    "image_pad",
-                ):
-                    setattr(
-                        self,
-                        name,
-                        nn.Parameter(
-                            torch.empty(config.hidden_size, dtype=torch.float32)
-                        ),
-                    )
-                self.vision.to(dtype=model_config.dtype)
-                self.aligner.to(dtype=model_config.dtype)
-
-        with self._mark_language_model(vllm_config):
-            # The arch convertor routes any config with a vision tower to
-            # this wrapper class; mark the copy handed to the text backbone
-            # so it resolves to DeepseekV4ForCausalLM instead of recursing
-            # (with_hf_config deepcopies the config, the marker survives).
-            config._dsv4_vl_inner = True  # type: ignore[attr-defined]
-            try:
-                self.language_model = init_vllm_registered_model(
-                    vllm_config=vllm_config,
-                    hf_config=config,
-                    prefix=maybe_prefix(prefix, "language_model"),
-                    architectures=["DeepseekV4ForCausalLM"],
-                )
-            finally:
-                del config._dsv4_vl_inner  # type: ignore[attr-defined]
-        # The outer mapper (see load_weights) fully resolves HF names into
-        # this wrapper's namespace before AutoWeightsLoader strips the
-        # "language_model." prefix and delegates to the child's load_weights,
-        # so the child's own mapper must be a no-op. Its suffix rules are not
-        # idempotent (e.g. "lm_head.weight".endswith("head.weight") would
-        # re-fire "head.weight" -> "lm_head.weight").
-        self.language_model.hf_to_vllm_mapper = WeightsMapper()
-        self.make_empty_intermediate_tensors = (  # type: ignore[method-assign]
-            self.language_model.make_empty_intermediate_tensors
-        )
-
-        expert_dtype = getattr(config, "expert_dtype", "fp4")
-        self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper(
-            expert_dtype, image_enabled
-        )
-
-    def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None:
-        patches = kwargs.pop("patches", None)
-        if patches is None:
-            return None
-        vit_grid = kwargs.pop("vit_grid", None)
-        llm_grid = kwargs.pop("llm_grid", None)
-        perm = kwargs.pop("perm", None)
-        assert vit_grid is not None and llm_grid is not None and perm is not None
-        return {
-            "patches": patches,
-            "vit_grid": vit_grid,
-            "llm_grid": llm_grid,
-            "perm": perm,
-        }
-
-    def _encode_image(
-        self,
-        patches: torch.Tensor,
-        n_vit_h: int,
-        n_vit_w: int,
-        perm: torch.Tensor,
-    ) -> torch.Tensor:
-        assert self.vision is not None and self.aligner is not None
-        image_embeds = self.aligner(
-            self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w
-        )
-        # Reorder into the N-layout block order used in the prompt.
-        return image_embeds[perm.to(image_embeds.device)]
-
-    def _process_image_input(
-        self,
-        patches: torch.Tensor,
-        vit_grid: torch.Tensor,
-        llm_grid: torch.Tensor,
-        perm: torch.Tensor,
-    ) -> tuple[torch.Tensor, ...]:
-        assert self.vision is not None and self.aligner is not None
-        patches = patches.to(self.aligner.w1.weight.dtype)
-
-        embeds: list[torch.Tensor] = []
-        vit_offset = 0
-        llm_offset = 0
-        for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip(
-            vit_grid.tolist(), llm_grid.tolist(), strict=True
-        ):
-            n_vit = n_vit_h * n_vit_w
-            n_llm = n_llm_h * n_llm_w
-            embeds.append(
-                self._encode_image(
-                    patches[vit_offset : vit_offset + n_vit],
-                    n_vit_h,
-                    n_vit_w,
-                    perm[llm_offset : llm_offset + n_llm],
-                )
-            )
-            vit_offset += n_vit
-            llm_offset += n_llm
-        return tuple(embeds)
-
-    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
-        image_input = self._parse_and_validate_image_input(**kwargs)
-        if image_input is None or self.vision is None:
-            return []
-        return self._process_image_input(
-            image_input["patches"],
-            image_input["vit_grid"],
-            image_input["llm_grid"],
-            image_input["perm"],
-        )
-
-    def embed_input_ids(
-        self,
-        input_ids: torch.Tensor,
-        multimodal_embeddings: MultiModalEmbeddings | None = None,
-        *,
-        is_multimodal: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        from vllm.model_executor.models.utils import _merge_multimodal_embeddings
-
-        # All ids are in-vocab here: image-block sentinels are borrowed
-        # reserved tokens (their embedding rows are always overwritten below).
-        inputs_embeds = self.language_model.embed_input_ids(input_ids)
-
-        if self.image_start is not None:
-            # Branch-free sentinel overwrite: safe inside compiled/captured
-            # regions (no data-dependent control flow).
-            sentinel_mask = image_sentinel_mask(input_ids)
-            if is_multimodal is not None:
-                # IMAGE positions get vision embeddings via the merge below.
-                sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device)
-            table = torch.stack(
-                [
-                    self.image_start,
-                    self.image_pad,
-                    self.image_pad,
-                    self.image_newline,
-                    self.image_end,
-                ]
-            ).to(inputs_embeds.dtype)
-            idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4)
-            inputs_embeds = torch.where(
-                sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds
-            )
-
-        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
-            return inputs_embeds
-
-        assert is_multimodal is not None
-        return _merge_multimodal_embeddings(
-            inputs_embeds=inputs_embeds,
-            multimodal_embeddings=multimodal_embeddings,
-            is_multimodal=is_multimodal,
-        )
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        intermediate_tensors=None,
-        inputs_embeds: torch.Tensor | None = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        return self.language_model(
-            input_ids, positions, intermediate_tensors, inputs_embeds
-        )
-
-    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
-        return self.language_model.compute_logits(hidden_states)
-
-    def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return self.language_model.compute_logits_local(hidden_states)
-
-    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        return self.language_model.get_expert_mapping()
-
-    def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
-        """Pre-hc_head residual stream buffer for the MTP/DSpark draft model."""
-        return self.language_model.get_mtp_target_hidden_states()
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # Map HF names into this wrapper's namespace up front and sort, so
-        # the "language_model." group reaches the child loader as one
-        # contiguous block (AutoWeightsLoader delegates per contiguous group,
-        # and the child's load_weights finalizes fused expert weights, which
-        # must not run on a partially loaded model).
-        mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0])
-        loader = AutoWeightsLoader(self)
-        loaded_params = loader.load_weights(mapped)
-        # The child's load_weights already ran its post-load finalization.
-        self._weights_finalized = True
-        return loaded_params
-
-    def process_weights_after_loading(self) -> None:
-        # Model-level post-load hook (called by the loader after any load
-        # format). Under DummyModelLoader the child's load_weights — and
-        # hence its finalize step — is bypassed, so run it here instead.
-        if getattr(self, "_weights_finalized", False):
-            return
-        self.language_model.process_weights_after_loading()
+__all__ = ["DeepseekV4ForConditionalGeneration"]
diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py
index 81659f342a14..cdab7d23ae07 100644
--- a/vllm/models/deepseek_v4/sparse_mla.py
+++ b/vllm/models/deepseek_v4/sparse_mla.py
@@ -9,11 +9,11 @@
 
 from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
@@ -57,7 +57,7 @@ class DeepseekV4SparseMLABackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [256]
 
     @staticmethod
@@ -195,6 +195,7 @@ def build(
         if self.compress_ratio > 1:
             slot_mapping = get_compressed_slot_mapping(
                 cm.num_actual_tokens,
+                cm.slot_mapping,
                 cm.query_start_loc,
                 cm.seq_lens,
                 cm.block_table_tensor.clamp_(min=0),
diff --git a/vllm/models/deepseek_v4/vl_stub.py b/vllm/models/deepseek_v4/vl_stub.py
index 0f15a82f0190..4ade7b17d387 100644
--- a/vllm/models/deepseek_v4/vl_stub.py
+++ b/vllm/models/deepseek_v4/vl_stub.py
@@ -9,6 +9,6 @@ class DeepseekV4ForConditionalGeneration(nn.Module):
     def __init__(self, *, vllm_config, prefix: str = ""):
         super().__init__()
         raise NotImplementedError(
-            "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is only "
-            "supported on NVIDIA GPUs for now."
+            "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is not "
+            "supported on the current platform; use CUDA or ROCm."
         )
diff --git a/vllm/models/deepseek_v4/xpu/mtp.py b/vllm/models/deepseek_v4/xpu/mtp.py
index 3d1ce97687b9..a19d0d578220 100644
--- a/vllm/models/deepseek_v4/xpu/mtp.py
+++ b/vllm/models/deepseek_v4/xpu/mtp.py
@@ -493,8 +493,7 @@ def finalize_mega_moe_weights(self) -> None:
             layer.mtp_block.ffn.finalize_mega_moe_weights()
 
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
+        """Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
         and rename shared layer weights to be top level.
         """
diff --git a/vllm/models/deepseek_v4_1/__init__.py b/vllm/models/deepseek_v41/__init__.py
similarity index 100%
rename from vllm/models/deepseek_v4_1/__init__.py
rename to vllm/models/deepseek_v41/__init__.py
diff --git a/vllm/models/deepseek_v41/amd/__init__.py b/vllm/models/deepseek_v41/amd/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/deepseek_v41/amd/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/models/deepseek_v4_1/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py
similarity index 99%
rename from vllm/models/deepseek_v4_1/amd/dspark.py
rename to vllm/models/deepseek_v41/amd/dspark.py
index 3a6cef8aa74d..78f07b6b11db 100644
--- a/vllm/models/deepseek_v4_1/amd/dspark.py
+++ b/vllm/models/deepseek_v41/amd/dspark.py
@@ -259,6 +259,8 @@ def _insert_context_kv(
             attn.padded_heads,
             attn.eps,
             block_size,
+            True,  # apply_q_norm (unused: the query is a discarded dummy)
+            attn.kv_mxfp8,
         )
     elif cache_dtype == torch.bfloat16:
         swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
diff --git a/vllm/models/deepseek_v4_1/amd/model.py b/vllm/models/deepseek_v41/amd/model.py
similarity index 95%
rename from vllm/models/deepseek_v4_1/amd/model.py
rename to vllm/models/deepseek_v41/amd/model.py
index 069ff5f931a6..1383bac820cb 100644
--- a/vllm/models/deepseek_v4_1/amd/model.py
+++ b/vllm/models/deepseek_v41/amd/model.py
@@ -17,15 +17,12 @@
 )
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
-from vllm.model_executor.kernels.mhc.torch import (
-    mhc_post_torch,
-    mhc_pre_delayed_torch,
-)
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mhc import MHCPostOp, MHCPreDelayedOp
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -57,16 +54,20 @@
 from vllm.models.deepseek_v4.amd.model import (
     DeepseekV4MoE as DeepseekV4MoEBase,
 )
-from vllm.models.deepseek_v4_1.amd.rocm import DeepseekV41ROCMAiterMLAAttention
-from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
+from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention
+from vllm.models.deepseek_v41.attention import DeepseekV4Attention
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
-from ..common.engram import Engram, EngramLayout, NgramHashState
+from ..common.engram import EngramLayout, NgramHashState
 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
 
+# Engram host offload and its prefetch stream are neither ROCm- nor
+# NVIDIA-specific, so they are imported rather than duplicated.
+from ..nvidia.engram import Engram
+
 if typing.TYPE_CHECKING:
     from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
 
@@ -248,11 +249,13 @@ def __init__(
             ),
             requires_grad=False,
         )
+        self.mhc_pre_delayed = MHCPreDelayedOp()
+        self.mhc_post = MHCPostOp()
 
     @staticmethod
     def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor:
-        """Reference hyper-connection collapse used by DSpark on ROCm."""
-        return (pre_mix.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype)
+        """Hyper-connection collapse used by DSpark on ROCm."""
+        return torch.ops.vllm.hc_collapse_triton(x, pre_mix)
 
     def forward(
         self,
@@ -275,7 +278,7 @@ def forward(
                 # copies and the identity pre-mix selects copy 0.
                 assert self.hc_attn_fn_broadcast is not None
                 residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous()
-                post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
+                residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
                     residual,
                     self.hc_attn_fn_broadcast,
                     self.hc_attn_scale,
@@ -289,7 +292,7 @@ def forward(
                 )
             else:
                 residual = x
-                post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
+                residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
                     residual,
                     self.hc_attn_fn,
                     self.hc_attn_scale,
@@ -302,18 +305,7 @@ def forward(
                     pre_mix=pre_mix,
                 )
         else:
-            residual = mhc_post_torch(x, residual, post_mix, res_mix)
-            if self.engram is not None and engram_hashes is not None:
-                # Engram injection happens between the previous sublayer's
-                # post and this block's pre, on the full hc stream, so the
-                # mix coefficients see the injected stream.
-                residual = self.engram(
-                    residual,
-                    engram_hashes[:, self.engram.layer_hash_index],
-                    engram_mask,
-                )
-            post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
-                residual,
+            pre_args = (
                 self.hc_attn_fn,
                 self.hc_attn_scale,
                 self.hc_attn_base,
@@ -322,8 +314,37 @@ def forward(
                 self.hc_eps,
                 self.hc_post_alpha,
                 self.hc_sinkhorn_iters,
-                pre_mix=pre_mix,
             )
+            if self.engram is not None and engram_hashes is not None:
+                # Engram injection happens between the previous sublayer's
+                # post and this block's pre, on the full hc stream, so the
+                # mix coefficients see the injected stream. That read of the
+                # residual is what stops the post from fusing into the pre
+                # here, unlike the attention seam below.
+                residual = self.mhc_post(x, residual, post_mix, res_mix)
+                residual = self.engram(
+                    residual,
+                    engram_hashes[:, self.engram.layer_hash_index],
+                    engram_mask,
+                )
+                residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
+                    residual, *pre_args, pre_mix=pre_mix
+                )
+            else:
+                (
+                    residual,
+                    post_mix,
+                    res_mix,
+                    x,
+                    attn_pre,
+                ) = self.mhc_pre_delayed(
+                    residual,
+                    *pre_args,
+                    pre_mix=pre_mix,
+                    sublayer_out=x,
+                    post_layer_mix=post_mix,
+                    comb_res_mix=res_mix,
+                )
         x = self.attn_norm(x)
 
         if self.use_sequence_parallel:
@@ -333,8 +354,7 @@ def forward(
         if self.use_sequence_parallel:
             x = sp_reduce_scatter(x)
 
-        residual = mhc_post_torch(x, residual, post_mix, res_mix)
-        post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_torch(
+        residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed(
             residual,
             self.hc_ffn_fn,
             self.hc_ffn_scale,
@@ -345,6 +365,9 @@ def forward(
             self.hc_post_alpha,
             self.hc_sinkhorn_iters,
             pre_mix=attn_pre,
+            sublayer_out=x,
+            post_layer_mix=post_mix,
+            comb_res_mix=res_mix,
         )
         x = self.ffn_norm(x)
         x = self.ffn(x, input_ids)
@@ -453,6 +476,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         else:
             self.norm = PPMissingLayer()
 
+        self.mhc_post = MHCPostOp()
+
         spec_config = vllm_config.speculative_config
         needs_mtp_hidden_states = spec_config is not None and (
             spec_config.use_eagle() or spec_config.uses_draft_model()
@@ -603,7 +628,7 @@ def forward(
             )
             if idx + 1 in self.aux_hidden_state_layers:
                 # Reconstruct the aux hidden state for draft models
-                aux_recon = mhc_post_torch(hidden_states, residual, post_mix, res_mix)
+                aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix)
                 aux_hidden_state = aux_recon.mean(dim=1)
                 if self.use_sequence_parallel:
                     aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
@@ -614,7 +639,7 @@ def forward(
             if self.end_layer in self.aux_hidden_state_layers:
                 hidden_states = final_aux_recon
             else:
-                hidden_states = mhc_post_torch(
+                hidden_states = self.mhc_post(
                     hidden_states, residual, post_mix, res_mix
                 )
 
@@ -635,11 +660,7 @@ def forward(
         # mixes — the mix the reference applies via
         # ``last_layer.hc_pre(h, pre_mix)`` (v4.1 has no learned hc_head).
         assert pre_mix is not None
-        hidden_states = (
-            (pre_mix.unsqueeze(-1) * hidden_states.float())
-            .sum(dim=1)
-            .to(hidden_states.dtype)
-        )
+        hidden_states = torch.ops.vllm.hc_collapse_triton(hidden_states, pre_mix)
         hidden_states = self.norm(hidden_states)
         if len(aux_hidden_states) > 0:
             return hidden_states, aux_hidden_states
diff --git a/vllm/models/deepseek_v4_1/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py
similarity index 76%
rename from vllm/models/deepseek_v4_1/amd/rocm.py
rename to vllm/models/deepseek_v41/amd/rocm.py
index 9271e7ae5a14..e1588b55e716 100644
--- a/vllm/models/deepseek_v4_1/amd/rocm.py
+++ b/vllm/models/deepseek_v41/amd/rocm.py
@@ -13,9 +13,14 @@
 )
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
-from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
-from vllm.models.deepseek_v4_1.common.ops import dequantize_and_gather_k_cache
-from vllm.models.deepseek_v4_1.sparse_mla import (
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
+from vllm.models.deepseek_v41.attention import (
+    DeepseekV4Attention,
+    _replace_layer_index,
+)
+from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
+from vllm.models.deepseek_v41.sparse_mla import (
     DeepseekV4FlashMLAMetadata,
     DeepseekV4SparseMLABackend,
     DeepseekV4SparseMLAMetadataBuilder,
@@ -34,6 +39,9 @@
 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
     build_ragged_indices_from_dense,
     rocm_inv_rope_einsum,
+    rocm_inverse_rope_mxfp8_rows,
+    rocm_inverse_rope_rows_,
+    rocm_mxfp8_wo_a_bmm,
     rocm_sparse_attn_decode,
     rocm_sparse_attn_prefill,
 )
@@ -391,6 +399,10 @@ def _copy_ragged_to_graph_buffers(
     return ragged_out, indptr_out
 
 
+# (ragged_indices, ragged_indptr, lens) as consumed by rocm_sparse_attn_decode.
+_TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+
+
 @dataclass
 class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
     decode_swa_ragged_indices: torch.Tensor | None = None
@@ -425,11 +437,13 @@ def build(
         common_prefix_len: int,
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
+        replay_start: torch.Tensor | None = None,
     ) -> DeepseekV4ROCMAiterSparseSWAMetadata:
         base = super().build(
             common_prefix_len=common_prefix_len,
             common_attn_metadata=common_attn_metadata,
             fast_build=fast_build,
+            replay_start=replay_start,
         )
 
         ragged_indices = None
@@ -498,6 +512,23 @@ def __init__(self, *args, **kwargs):
         self._fused_compressor_weight: torch.Tensor | None
         self.register_buffer("_fused_compressor_weight", None, persistent=False)
         self._fused_compressor_split_sizes: tuple[int, int] | None = None
+        # Decode ragged topk metadata is a pure function of the indices its
+        # index source published, so every consumer of a source rebuilds the
+        # same thing. The source owns one cache per compress ratio, refreshed
+        # when it runs; consumers below it read through.
+        self._topk_ragged_cache: dict[int, _TopkRagged] = {}
+        self._index_source_prefix: str | None = None
+        if self.compress_ratio > 0:
+            assert self.index_source_layer_id is not None
+            self._index_source_prefix = _replace_layer_index(
+                self.prefix, self.index_source_layer_id
+            )
+            if self._index_source_prefix not in self._static_forward_context:
+                raise NotImplementedError(
+                    f"Index source {self._index_source_prefix} not found on "
+                    "this rank; PP splits inside a v4.1 index-sharing group "
+                    "are not supported."
+                )
 
     @classmethod
     def get_padded_num_q_heads(cls, num_heads: int) -> int:
@@ -636,7 +667,41 @@ def _split_qkv_and_norm(
             transpose_scale=False,
         )
 
-    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+    def _alloc_attn_out(
+        self, num_tokens: int, hidden_states: torch.Tensor
+    ) -> torch.Tensor | QuantizedActivation:
+        if not _ON_GFX950:
+            return super()._alloc_attn_out(num_tokens, hidden_states)
+        # wo_a's MXFP8 input: the decode reduce writes it directly, prefill
+        # rows are rotated and quantized after their bf16 attention.
+        width = self.n_local_heads * self.head_dim
+        data = torch.empty(
+            (num_tokens, width), dtype=torch.float8_e4m3fn, device=hidden_states.device
+        )
+        scale = torch.empty(
+            (num_tokens, width // 32), dtype=torch.uint8, device=hidden_states.device
+        )
+        return QuantizedActivation(
+            data=data,
+            scale=scale,
+            orig_dtype=hidden_states.dtype,
+            orig_shape=data.shape,
+            quant_key=kMxfp8Dynamic,
+        )
+
+    def _o_proj(
+        self, attn_out: torch.Tensor | QuantizedActivation, positions: torch.Tensor
+    ) -> torch.Tensor:
+        if isinstance(attn_out, QuantizedActivation):
+            z = rocm_mxfp8_wo_a_bmm(
+                attn_out.data,
+                attn_out.scale,
+                self.wo_a,
+                self.n_local_groups,
+                self.o_lora_rank,
+            )
+            return self._wo_b_after_wo_a(z)
+        o = attn_out[:, : self.n_local_heads, :]
         # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
         z = rocm_inv_rope_einsum(
             self.rotary_emb,
@@ -646,27 +711,31 @@ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
             self.n_local_groups,
             self.o_lora_rank,
             self.wo_a,
+            inverse_rope=False,
         )
-        zf = z.flatten(1)
+        return self._wo_b_after_wo_a(z.flatten(1))
+
+    def _wo_b_after_wo_a(self, zf: torch.Tensor) -> torch.Tensor:
         if self._wo_b_scale is not None and zf.dim() == 2:
-            result = self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
-        else:
-            result = self.wo_b(zf)
-        return result
+            return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
+        return self.wo_b(zf)
 
     def forward_mqa(
         self,
         q: torch.Tensor,
         kv: torch.Tensor,
         positions: torch.Tensor,
-        output: torch.Tensor,
+        output: torch.Tensor | QuantizedActivation,
     ) -> None:
-        assert output.shape == q.shape, (
-            f"output buffer shape {output.shape} must match q shape {q.shape}"
-        )
-        assert output.dtype == q.dtype, (
-            f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
-        )
+        mxfp8_out = output if isinstance(output, QuantizedActivation) else None
+        if mxfp8_out is None:
+            assert isinstance(output, torch.Tensor)
+            assert output.shape == q.shape, (
+                f"output buffer shape {output.shape} must match q shape {q.shape}"
+            )
+            assert output.dtype == q.dtype, (
+                f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
+            )
 
         forward_context = get_forward_context()
         attn_metadata = forward_context.attn_metadata
@@ -684,9 +753,13 @@ def forward_mqa(
             )
             M = N + self.window_size + self.max_num_batched_tokens
             current_workspace_manager().get_simultaneous(
-                ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
+                *self._prefill_workspace_shapes(M, self.max_num_batched_tokens, q)
             )
-            output.zero_()
+            if mxfp8_out is None:
+                output.zero_()
+            else:
+                mxfp8_out.data.zero_()
+                mxfp8_out.scale.zero_()
             return
 
         assert isinstance(attn_metadata, dict)
@@ -716,29 +789,112 @@ def forward_mqa(
                 positions=positions[num_decode_tokens:],
                 compressed_k_cache=self_kv_cache,
                 swa_k_cache=swa_kv_cache,
-                output=output[num_decode_tokens:],
+                output=(
+                    output[num_decode_tokens:]
+                    if isinstance(output, torch.Tensor)
+                    else None
+                ),
                 attn_metadata=rocm_metadata,
                 swa_metadata=swa_metadata,
+                mxfp8_out=(
+                    None
+                    if mxfp8_out is None
+                    else (
+                        mxfp8_out.data[num_decode_tokens:],
+                        mxfp8_out.scale[num_decode_tokens:],
+                    )
+                ),
             )
+        if mxfp8_out is not None:
+            # The reduce epilogue rotates and quantizes every decode row and
+            # the prefill path already did its own, so nothing is owed here.
+            if num_decodes > 0:
+                self._forward_decode(
+                    q=q[:num_decode_tokens],
+                    positions=positions[:num_decode_tokens],
+                    kv_cache=self_kv_cache,
+                    swa_metadata=swa_metadata,
+                    attn_metadata=rocm_metadata,
+                    swa_only=swa_only,
+                    output=None,
+                    output_mxfp8=(
+                        mxfp8_out.data[:num_decode_tokens],
+                        mxfp8_out.scale[:num_decode_tokens],
+                    ),
+                )
+            return
+        assert isinstance(output, torch.Tensor)
+        rotated = 0
         if num_decodes > 0:
-            self._forward_decode(
+            rotated = self._forward_decode(
                 q=q[:num_decode_tokens],
+                positions=positions[:num_decode_tokens],
                 kv_cache=self_kv_cache,
                 swa_metadata=swa_metadata,
                 attn_metadata=rocm_metadata,
                 swa_only=swa_only,
                 output=output[:num_decode_tokens],
             )
+        # Only the decode reduce rotates its own rows, and only the leading
+        # `rotated` of them; prefill rows and any decode path that did not
+        # fuse still owe the standalone pass. Settle that here rather than in
+        # _o_proj: the split is batch-dependent and _o_proj runs compiled,
+        # where such a value freezes at its trace-time value.
+        rocm_inverse_rope_rows_(
+            output[rotated:, : self.n_local_heads, :],
+            positions[rotated:],
+            self.rotary_emb.cos_sin_cache,
+            self.rope_head_dim,
+        )
+
+    def _decode_topk_ragged(
+        self,
+        swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        num_decodes: int,
+        num_decode_tokens: int,
+    ) -> _TopkRagged:
+        """Ragged form of the topk indices this layer's index source published.
+
+        The packing depends only on the shared ``topk_indices_buffer`` and
+        per-step metadata, apart from the layer's own compress ratio, so the
+        source memoizes one result per ratio for all the layers below it.
+        """
+        assert attn_metadata is not None
+        assert swa_metadata.is_valid_token is not None
+        assert self.topk_indices_buffer is not None
+        assert self._index_source_prefix is not None
+
+        source = self._static_forward_context[self._index_source_prefix]
+        if source is self:
+            # Fresh indices as of this layer; drop what the last step cached.
+            source._topk_ragged_cache = {}
+        cached = source._topk_ragged_cache.get(self.compress_ratio)
+        if cached is not None:
+            return cached
+
+        built = compute_global_topk_ragged_indices_and_indptr(
+            self.topk_indices_buffer[:num_decode_tokens],
+            swa_metadata.token_to_req_indices,
+            attn_metadata.block_table[:num_decodes],
+            attn_metadata.block_size // self.compress_ratio,
+            swa_metadata.is_valid_token[:num_decode_tokens],
+        )
+        source._topk_ragged_cache[self.compress_ratio] = built
+        return built
 
     def _forward_decode(
         self,
         q: torch.Tensor,
+        positions: torch.Tensor,
         kv_cache: torch.Tensor | None,
         swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
         attn_metadata: DeepseekV4FlashMLAMetadata | None,
         swa_only: bool,
-        output: torch.Tensor,
-    ) -> None:
+        output: torch.Tensor | None,
+        output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> int:
+        """Returns how many leading rows the decode epilogue inverse-RoPE'd."""
         num_decodes = swa_metadata.num_decodes
         num_decode_tokens = swa_metadata.num_decode_tokens
 
@@ -746,25 +902,18 @@ def _forward_decode(
         topk_ragged_indices = None
         topk_ragged_indptr = None
         if not swa_only:
-            # Local indices filled by the index-source layer's indexer.
-            assert attn_metadata is not None
-            assert swa_metadata.is_valid_token is not None
-            assert self.topk_indices_buffer is not None
-            block_size = attn_metadata.block_size // self.compress_ratio
-            is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
             (
                 topk_ragged_indices,
                 topk_ragged_indptr,
                 topk_lens,
-            ) = compute_global_topk_ragged_indices_and_indptr(
-                self.topk_indices_buffer[:num_decode_tokens],
-                swa_metadata.token_to_req_indices,
-                attn_metadata.block_table[:num_decodes],
-                block_size,
-                is_valid,
+            ) = self._decode_topk_ragged(
+                swa_metadata=swa_metadata,
+                attn_metadata=attn_metadata,
+                num_decodes=num_decodes,
+                num_decode_tokens=num_decode_tokens,
             )
 
-        rocm_sparse_attn_decode(
+        return rocm_sparse_attn_decode(
             q=q,
             kv_cache=kv_cache,
             swa_k_cache=self.swa_cache_layer.kv_cache,
@@ -783,6 +932,9 @@ def _forward_decode(
             nope_head_dim=self.nope_head_dim,
             rope_head_dim=self.rope_head_dim,
             output=output,
+            output_mxfp8=output_mxfp8,
+            inv_rope_positions=positions,
+            inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
             extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
                 self.kv_cache_dtype,
                 self._has_kv_transfer,
@@ -796,10 +948,16 @@ def _forward_prefill(
         positions: torch.Tensor,
         compressed_k_cache: torch.Tensor | None,
         swa_k_cache: torch.Tensor,
-        output: torch.Tensor,
+        output: torch.Tensor | None,
         attn_metadata: DeepseekV4FlashMLAMetadata | None,
         swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
+        mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> None:
+        """Sparse prefill into bf16 ``output``, or into ``mxfp8_out``.
+
+        With ``mxfp8_out`` the rows go to a bf16 workspace first and are then
+        inverse-RoPE'd and quantized into it in one pass.
+        """
         swa_only = attn_metadata is None
 
         num_prefills = swa_metadata.num_prefills
@@ -835,10 +993,13 @@ def _forward_prefill(
             self.PREFILL_CHUNK_SIZE
         )
 
-        workspace_manager = current_workspace_manager()
-        kv = workspace_manager.get_simultaneous(
-            ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
-        )[0]
+        workspace = current_workspace_manager().get_simultaneous(
+            *self._prefill_workspace_shapes(M, num_prefill_tokens, q)
+        )
+        kv = workspace[0]
+        if mxfp8_out is not None:
+            output = workspace[1]
+        assert output is not None
         for chunk_idx in range(num_chunks):
             chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
             chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
@@ -903,3 +1064,31 @@ def _forward_prefill(
                 attn_sink=self.attn_sink,
                 output=output[query_start:query_end],
             )
+        if mxfp8_out is not None:
+            rocm_inverse_rope_mxfp8_rows(
+                output,
+                positions[:num_prefill_tokens],
+                self.rotary_emb.cos_sin_cache,
+                self.rope_head_dim,
+                mxfp8_out[0][:num_prefill_tokens],
+                mxfp8_out[1][:num_prefill_tokens],
+            )
+
+    def _prefill_workspace_shapes(
+        self, M: int, num_prefill_tokens: int, q: torch.Tensor
+    ) -> list[tuple[tuple[int, ...], torch.dtype]]:
+        """The prefill gather buffer, plus bf16 rows for the MXFP8 output.
+
+        One ``get_simultaneous`` call: separate calls alias the same memory.
+        """
+        shapes: list[tuple[tuple[int, ...], torch.dtype]] = [
+            ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16)
+        ]
+        if _ON_GFX950:
+            shapes.append(
+                (
+                    (num_prefill_tokens, self.n_local_heads, self.head_dim),
+                    torch.bfloat16,
+                )
+            )
+        return shapes
diff --git a/vllm/models/deepseek_v4_1/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py
similarity index 89%
rename from vllm/models/deepseek_v4_1/amd/vl_model.py
rename to vllm/models/deepseek_v41/amd/vl_model.py
index 91a13bd72506..3c57df2ad258 100644
--- a/vllm/models/deepseek_v4_1/amd/vl_model.py
+++ b/vllm/models/deepseek_v41/amd/vl_model.py
@@ -26,6 +26,7 @@
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
     SupportsEagle3,
+    SupportsEncoderCudaGraph,
     SupportsMultiModal,
     SupportsPP,
 )
@@ -44,17 +45,12 @@
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
 from ..common.mm_preprocess import (
-    IMAGE,
-    IMAGE_END,
-    IMAGE_NEW_LINE,
-    IMAGE_PAD_ID,
     IMAGE_PLACEHOLDER,
-    IMAGE_SENTINEL_BASE_ID,
-    IMAGE_START,
     DeepseekV4VLDummyInputsBuilder,
     DeepseekV4VLMultiModalProcessor,
     DeepseekV4VLProcessingInfo,
 )
+from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin
 from .model import (
     DeepseekV41LLMForCausalLM,
     _linear_scale_param_name,
@@ -117,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper(
     info=DeepseekV4VLProcessingInfo,
     dummy_inputs=DeepseekV4VLDummyInputsBuilder,
 )
-class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3):
+class DeepseekV41ForCausalLM(
+    nn.Module,
+    DeepseekV4VLEncoderCudaGraphMixin,
+    SupportsMultiModal,
+    SupportsEncoderCudaGraph,
+    SupportsPP,
+    SupportsEagle3,
+):
     """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower.
 
     ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
     delegates through ``language_model`` via the protocol defaults.
+    ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is
+    implemented by ``DeepseekV4VLEncoderCudaGraphMixin``.
     """
 
     supports_encoder_tp_data = True
@@ -198,29 +203,6 @@ def _parse_and_validate_image_input(
             resolve_bindings={"p": self.config.vision_patch_size},
         )
 
-    def _encode_image(
-        self,
-        patches: torch.Tensor,
-        n_vit_h: int,
-        n_vit_w: int,
-    ) -> torch.Tensor:
-        # Aligner rows in reading order, one per IMAGE slot.
-        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
-
-    def _build_image_span(
-        self, image_embeds: torch.Tensor, types: torch.Tensor
-    ) -> torch.Tensor:
-        """Full image span: aligner rows at IMAGE slots, the learned
-        delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""
-        types = types.to(image_embeds.device)
-        span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])
-        dtype = image_embeds.dtype
-        span[types == IMAGE_START] = self.image_start.to(dtype)
-        span[types == IMAGE_END] = self.image_end.to(dtype)
-        span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
-        span[types == IMAGE] = image_embeds
-        return span
-
     def _process_image_input(
         self,
         image_input: DeepseekV4VLImagePixelInputs,
@@ -277,12 +259,6 @@ def embed_input_ids(
     ) -> torch.Tensor:
         from vllm.model_executor.models.utils import _merge_multimodal_embeddings
 
-        # Compressor-alignment pads borrow a reserved id; embed them as the
-        # plain image token (the checkpoint has no image_pad vector).
-        # Branch-free: safe inside compiled/captured regions.
-        input_ids = input_ids.masked_fill(
-            input_ids == IMAGE_PAD_ID, IMAGE_SENTINEL_BASE_ID
-        )
         inputs_embeds = self.language_model.embed_input_ids(input_ids)
 
         if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v41/attention.py
similarity index 73%
rename from vllm/models/deepseek_v4_1/attention.py
rename to vllm/models/deepseek_v41/attention.py
index 77e98a376e95..d0e69fb35871 100644
--- a/vllm/models/deepseek_v4_1/attention.py
+++ b/vllm/models/deepseek_v41/attention.py
@@ -1,9 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-DeepseekV4 MLA Attention Layer
-"""
+"""DeepseekV4 MLA Attention Layer"""
 
+import math
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from functools import cached_property
@@ -25,14 +24,17 @@
     RowParallelLinear,
 )
 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
+from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer
 from vllm.models.common.ops import fused_q_kv_rmsnorm
-from vllm.models.deepseek_v4_1.common.ops import (
+from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter
+from vllm.models.deepseek_v41.common.ops import (
     MXFP4_BLOCK_SIZE,
     fused_indexer_q_rope_quant,
     indexer_k_norm_rope_store,
 )
 
 if TYPE_CHECKING:
+    from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
     from vllm.v1.attention.backends.mla.sparse_swa import (
         DeepseekSparseSWAMetadata,
     )
@@ -42,6 +44,7 @@
     VllmConfig,
     get_current_vllm_config,
 )
+from vllm.config.cache import CacheDType
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
@@ -49,8 +52,9 @@
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.models.utils import extract_layer_index
-from vllm.models.deepseek_v4_1.common.rope import build_deepseek_v4_rope
-from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor
+from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope
+from vllm.models.deepseek_v41.compressor import DeepseekCompressor
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.multi_stream_utils import (
     execute_in_parallel,
@@ -58,10 +62,13 @@
 )
 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
 from vllm.v1.attention.backends.mla.indexer import (
-    DeepseekV4IndexerBackend,
+    DeepseekV41IndexerBackend,
     dsa_indexer_uses_fp4,
     get_max_prefill_buffer_size,
 )
+from vllm.v1.attention.backends.mla.sparse_indexer import (
+    DeepseekV41SparseIndexerBackend,
+)
 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
 from vllm.v1.kv_cache_interface import (
     KVCacheSpec,
@@ -109,36 +116,52 @@ def _fill_short_context_topk_indices(
     )
 
 
+# Which packed fp8_ds_mla record a V4.1 layer writes. FlashMLA decodes
+# DeepSeek's V4.1 record -- all 512 dims (RoPE included) as fp8 e4m3 with one
+# UE8M0 scale per 32 dims, 512 data bytes + 16 scale bytes per token, pages
+# rounded to the kernel's 512 B TMA stride -- only in its SM100 sparse-decode
+# kernels. Every other arch keeps the V4 record: 448 fp8 NoPE + 64 bf16 RoPE
+# plus 7 UE8M0 scales of 64 dims and a pad byte (584 B, 576 B pages).
+def _use_v41_mxfp8_kv_record() -> bool:
+    return current_platform.is_device_capability_family(100)
+
+
 def _resolve_dsv4_kv_cache_dtype(
     use_fp8_ds_mla_layout: bool,
-    kv_cache_dtype: str,
+    kv_cache_dtype: CacheDType,
     cache_config: CacheConfig | None,
-) -> tuple[str, torch.dtype]:
+    packed_kv_cache_dtype: CacheDType = "fp8_ds_mla",
+) -> tuple[CacheDType, torch.dtype]:
     """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
 
     Both layouts are paged; they differ in the per-token block format. The
-    ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the
-    canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the
-    page-size specs pick the 576B per-token slot). Plain-row backends store each
-    token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3.
+    packed formats are ``uint8``-backed: ``fp8_ds_mla`` is UE8M0 block-scaled
+    fp8 throughout, ``nvfp4_ds_mla`` keeps that sliding-window record and
+    stores the compressed cache as NVFP4. An unspecific ``--kv-cache-dtype``
+    (``auto`` / ``fp8``) resolves to ``packed_kv_cache_dtype``, the record this
+    layer's kernel prefers, and the canonical string is written back onto
+    ``cache_config`` so the page-size specs pick the right per-token slot.
+    Plain-row backends store each token's KV row in its element dtype: bf16 or
+    per-tensor FP8 E4M3.
     """
     if use_fp8_ds_mla_layout:
-        # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
-        if kv_cache_dtype == "auto":
-            kv_cache_dtype = "fp8"
-        if not kv_cache_dtype.startswith("fp8"):
+        if kv_cache_dtype in ("auto", "fp8"):
+            kv_cache_dtype = packed_kv_cache_dtype
+        elif not kv_cache_dtype.endswith("_ds_mla"):
             raise ValueError(
-                "DeepseekV4 fp8_ds_mla layout only supports fp8 "
-                f"kv-cache, got {kv_cache_dtype}. Please set "
-                "`--kv-cache-dtype fp8` or select a backend that supports "
-                "bfloat16 KV cache."
+                "DeepseekV4 packed KV layouts only support fp8 kv-cache, got "
+                f"{kv_cache_dtype}. Please set `--kv-cache-dtype fp8` or "
+                "select a backend that supports bfloat16 KV cache."
             )
-        if kv_cache_dtype != "fp8_ds_mla":
-            if cache_config is not None:
-                cache_config.cache_dtype = "fp8_ds_mla"
-            kv_cache_dtype = "fp8_ds_mla"
-            logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
+        if cache_config is not None and cache_config.cache_dtype != kv_cache_dtype:
+            cache_config.cache_dtype = kv_cache_dtype
+        logger.info_once("Using DeepSeek's %s KV cache format.", kv_cache_dtype)
         return kv_cache_dtype, torch.uint8
+    if kv_cache_dtype.endswith("_ds_mla"):
+        raise ValueError(
+            f"{kv_cache_dtype} is a packed FlashMLA DeepSeek V4.1 KV cache "
+            "format; the selected backend stores plain KV rows."
+        )
 
     # Plain bf16 / per-tensor fp8 KV row (FlashInfer).
     if kv_cache_dtype.startswith("fp8"):
@@ -156,7 +179,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
     ``DeepseekV4FlashInferSM120Attention`` /
     ``DeepseekV4FlashInferMLAAttention`` (CUDA) or
     ``DeepseekV41ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific
-    deepseek_v4_1 model module. The base is never instantiated directly.
+    deepseek_v41 model module. The base is never instantiated directly.
     """
 
     # Provided by the platform subclass.
@@ -173,6 +196,28 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
     # path to pre-reserve that workspace.
     PREFILL_CHUNK_SIZE: ClassVar[int] = 4
 
+    # ---- attention-interface contract, declared by the platform subclass ----
+
+    @property
+    def accepts_unnormed_unroped_query(self) -> bool:
+        """Whether ``forward_mqa``'s ``q`` is the raw ``wq_b`` output.
+
+        True when the attention kernel applies the Q norm and RoPE itself, and
+        reads Q in its own chunk-interleaved layout, so the layer only
+        zero-pads Q to ``padded_heads`` and inserts KV before calling it.
+        """
+        return False
+
+    @property
+    def packed_kv_cache_dtype(self) -> CacheDType:
+        """The packed KV record this layer's kernel prefers.
+
+        What an unspecific ``--kv-cache-dtype`` (``auto`` / ``fp8``) resolves
+        to. Mega attention overrides it: its kernel is the one that can read
+        an NVFP4 compressed cache.
+        """
+        return "fp8_ds_mla"
+
     @classmethod
     @abstractmethod
     def get_padded_num_q_heads(cls, num_heads: int) -> int:
@@ -197,8 +242,19 @@ def forward_mqa(
         raise NotImplementedError
 
     @abstractmethod
-    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
-        """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
+    def _o_proj(
+        self,
+        attn_out: "torch.Tensor | QuantizedActivation",
+        positions: torch.Tensor,
+    ) -> torch.Tensor:
+        """Project whatever ``_alloc_attn_out`` produced through wo_a and wo_b.
+
+        Takes the buffer whole, so each layer owns the shape it allocated: the
+        bf16 layers slice off their padding heads and apply the inverse RoPE,
+        while a layer whose attention kernel already did the inverse RoPE and
+        the FP8 cast gets a QuantizedActivation and has only wo_a and wo_b
+        left.
+        """
         raise NotImplementedError
 
     def _uses_fp8_ds_mla_layout(self) -> bool:
@@ -234,13 +290,6 @@ def __init__(
         self.n_groups = config.o_groups
         self.n_local_groups = self.n_groups // tp_size
         self.window_size = config.sliding_window
-        # Vision variant: image spans are visible bidirectionally, widening
-        # prefill SWA index rows by up to max_image_tokens columns.
-        self.max_image_tokens = (
-            getattr(config, "vision_max_n_token", 0)
-            if getattr(config, "vision_n_layers", 0) > 0
-            else 0
-        )
         # ---- v4.1 sparse-attention topology ----
         # compress_ratios has one entry per layer (MTP layers included):
         # 0 = pure sliding window, 1 = full-length compressed cache,
@@ -342,6 +391,9 @@ def __init__(
             return_bias=False,
             prefix=f"{prefix}.wo_b",
         )
+        # Set by ``bind_gemm_rs`` when the decoder layer runs sequence
+        # parallel and the fused GEMM + reduce-scatter kernel accepts wo_b.
+        self.gemm_rs: GemmRsAr | None = None
 
         # Initialize rotary embedding before the indexer/compressor consume it.
         self.rotary_emb = build_deepseek_v4_rope(
@@ -443,9 +495,56 @@ def __init__(
         # Resolve the kv-cache dtype from this backend's block format. The same
         # resolution drives the SWA cache tensor dtype below.
         self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
-            self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
+            self._uses_fp8_ds_mla_layout(),
+            cache_config.cache_dtype,
+            cache_config,
+            self.packed_kv_cache_dtype,
+        )
+        self.kv_mxfp8 = _use_v41_mxfp8_kv_record()
+        self.swa_bytes_per_token = 528 if self.kv_mxfp8 else 584
+        # nvfp4_ds_mla keeps the MXFP8 sliding-window record and stores the
+        # compressed cache as NVFP4 (256 B of e2m1 pairs + 32 e4m3 scales).
+        self.compressed_bytes_per_token = (
+            288 if self.kv_cache_dtype == "nvfp4_ds_mla" else self.swa_bytes_per_token
         )
+        # One alignment for every page in the block: the block stride is their
+        # sum, and 512 satisfies both TMA strides in play (512 for the V4.1
+        # fp8 record, 256 for NVFP4).
+        self.kv_page_alignment = 512 if self.kv_mxfp8 else 576
+        if self.kv_cache_dtype == "nvfp4_ds_mla" and not self.kv_mxfp8:
+            raise ValueError(
+                "nvfp4_ds_mla needs the V4.1 KV records, which FlashMLA "
+                "decodes only on SM100."
+            )
 
+        swa_bounded_replay = cache_config.swa_bounded_replay
+        if swa_bounded_replay and not vllm_config.use_v2_model_runner:
+            logger.warning_once(
+                "SWA bounded replay needs model runner V2 (only it skips the "
+                "paged-KV writes of replayed tokens); the sliding-window cache "
+                "takes part in prefix caching instead."
+            )
+            swa_bounded_replay = False
+        if (
+            swa_bounded_replay
+            and vllm_config.parallel_config.prefill_context_parallel_size > 1
+        ):
+            logger.warning_once(
+                "SWA bounded replay is off under prefill context parallelism "
+                "(the replayed tokens' slot padding knows the rank-local batch "
+                "only); the sliding-window cache takes part in prefix caching "
+                "instead."
+            )
+            swa_bounded_replay = False
+        if swa_bounded_replay and current_platform.is_rocm():
+            logger.warning_once(
+                "SWA bounded replay is off on ROCm (the sparse SWA metadata "
+                "builders forward replay_start, but the window clamp it relies "
+                "on lives in the FlashInfer and FlashMLA prefill kernels, so "
+                "the padded slots fault); the sliding-window cache takes part "
+                "in prefix caching instead."
+            )
+            swa_bounded_replay = False
         self.swa_cache_layer = DeepseekV4SWACache(
             head_dim=self.head_dim,
             window_size=self.window_size,
@@ -454,6 +553,9 @@ def __init__(
             cache_config=cache_config,
             backend_cls=self.swa_backend_cls,
             block_size=32,
+            packed_bytes_per_token=self.swa_bytes_per_token,
+            packed_page_alignment=self.kv_page_alignment,
+            bounded_replay=swa_bounded_replay,
         )
 
         # The attention layer itself was already registered with the
@@ -501,7 +603,6 @@ def __init__(
             _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup(
                 window_size=self.window_size,
                 block_size=self.swa_cache_layer.block_size,
-                max_image_tokens=self.max_image_tokens,
             )
 
             if self.compress_ratio > 1:
@@ -532,11 +633,15 @@ def __init__(
                     block_size=self.swa_cache_layer.block_size,
                 )
 
+            # Every backend that gathers a chunk's KV through
+            # combine_topk_swa_indices needs its Triton kernel warmed, mega
+            # attention included -- it calls it from _forward_prefill_mega.
             if self.backend_cls.get_name() in (
                 "FLASHMLA_SPARSE_DSV41",
+                "FLASHMLA_MEGA_ATTN_DSV41",
                 "ROCM_FLASHMLA_SPARSE_DSV4",
             ):
-                from vllm.models.deepseek_v4_1.common.ops.cache_utils import (
+                from vllm.models.deepseek_v41.common.ops.cache_utils import (
                     _COMBINE_TOPK_SWA_INDICES_KERNEL,
                 )
 
@@ -548,14 +653,10 @@ def forward(
         hidden_states: torch.Tensor,
         llama_4_scaling: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        # Pre-allocate attention output with FlashMLA-padded head count.
-        # The op writes into `o_padded`; we slice to n_local_heads after.
-        num_tokens = hidden_states.shape[0]
-        o_padded = torch.empty(
-            (num_tokens, self.padded_heads, self.head_dim),
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
+        # The eager attention region writes into a caller-owned buffer
+        # (breakable_cudagraph needs in-place outputs); its shape and how it is
+        # projected afterwards follow the interface contract above.
+        attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states)
 
         # Keep the attention input preparation in the captured graph. Only the
         # sparse indexer and MLA attention run in the eager break below.
@@ -572,16 +673,58 @@ def forward(
             kv_score,
             indexer_weights,
             positions,
-            o_padded,
+            attn_out,
         )
-        o = o_padded[:, : self.n_local_heads, :]
+        return self._o_proj(attn_out, positions)
 
-        # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
-        return self._o_proj(o, positions)
+    def bind_gemm_rs(self) -> None:
+        """Fuse ``wo_b`` with the sequence-parallel TP reduce-scatter.
+
+        The decoder layer calls this after the model initialized the
+        process-wide GEMM-RS workspace and before weights load: the
+        eligibility check inspects the projection's linear kernel and weight
+        shape, which online quantization may later re-layout.
+        """
+        from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
+            get_gemm_rs_ar,
+        )
+
+        # The layer owns the reduction under sequence parallel.
+        assert not self.wo_b.reduce_results
+        gemm_rs = get_gemm_rs_ar()
+        if gemm_rs.can_run(self.wo_b):
+            self.gemm_rs = gemm_rs
+        else:
+            gemm_rs.warn_incompatible_projection()
+
+    def _wo_b_proj(self, z: torch.Tensor) -> torch.Tensor:
+        """Apply ``wo_b``; with GEMM-RS bound, also reduce-scatter the result.
+
+        Every ``_o_proj`` implementation projects through this so the decoder
+        layer sees one contract: when ``gemm_rs`` is bound the output is
+        already the local sequence-parallel shard, otherwise it is the
+        unreduced TP partial.
+        """
+        if self.gemm_rs is None:
+            return self.wo_b(z)
+        if self.gemm_rs.should_run(z):
+            return self.gemm_rs.apply(z, self.wo_b)
+        # Small batches stay on the unfused path, which is faster there.
+        return sp_reduce_scatter(self.wo_b(z))
+
+    def _alloc_attn_out(
+        self, num_tokens: int, hidden_states: torch.Tensor
+    ) -> "torch.Tensor | QuantizedActivation":
+        """The buffer ``forward_mqa`` fills, per the interface contract."""
+        return torch.empty(
+            (num_tokens, self.padded_heads, self.head_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
 
     @cached_property
     def _can_fuse_query_quant(self) -> bool:
-        from vllm.models.deepseek_v4_1.common.ops.query_quant import (
+        from vllm.models.deepseek_v41.common.ops.query_quant import (
             can_fuse_query_quant,
         )
 
@@ -600,7 +743,7 @@ def _split_qkv_and_norm(
         """
         qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
         if self.q_lora_rank % 32 == 0 and self._can_fuse_query_quant:
-            from vllm.models.deepseek_v4_1.common.ops.query_quant import (
+            from vllm.models.deepseek_v41.common.ops.query_quant import (
                 fused_q_kv_rmsnorm_quant,
             )
 
@@ -631,7 +774,7 @@ def _prepare_and_attn_eager(
         kv_score: torch.Tensor,
         indexer_weights: torch.Tensor,
         positions: torch.Tensor,
-        o_padded: torch.Tensor,
+        attn_out: "torch.Tensor | QuantizedActivation",
     ) -> None:
         """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly.
 
@@ -646,7 +789,7 @@ def _prepare_and_attn_eager(
             kv_score,
             indexer_weights,
             positions,
-            o_padded,
+            attn_out,
         )
 
     def _prepare_and_attn(
@@ -658,7 +801,7 @@ def _prepare_and_attn(
         kv_score: torch.Tensor,
         indexer_weights: torch.Tensor,
         positions: torch.Tensor,
-        o_padded: torch.Tensor,
+        attn_out: "torch.Tensor | QuantizedActivation",
     ) -> None:
         """Attention input preparation followed by the sparse indexer and MLA.
 
@@ -730,7 +873,7 @@ def prepare_indexer():
             q,
             kv,
             positions,
-            o_padded,
+            attn_out,
         )
 
     def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -837,16 +980,28 @@ def _fused_qnorm_rope_kv_insert(
             dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
         ),
     ) -> torch.Tensor:
+        """Ready ``q`` for the attention kernel and publish this step's KV.
+
+        One launch does both. With ``accepts_unnormed_unroped_query`` the
+        attention kernel norms and rotates Q itself and reads it in its own
+        chunk-interleaved layout, so the Q half of the launch is a zero-pad to
+        ``padded_heads`` -- and nothing at all once the shard is that wide.
+        """
         if not isinstance(attn_metadata, dict):
             # Profile run: kernel doesn't fire; produce a padded tensor so
             # downstream FlashMLA gets the right shape.
-            if self.n_local_heads < self.padded_heads:
-                return F.pad(
-                    q,
-                    (0, 0, 0, self.padded_heads - self.n_local_heads),
-                    value=0.0,
-                )
-            return q
+            if self.n_local_heads >= self.padded_heads:
+                return q
+            if self.accepts_unnormed_unroped_query:
+                # Padding heads sit at the tail of every head-dim chunk in
+                # that layout, so no head-major pad of `q` reproduces it --
+                # and nothing reads it on a profile run.
+                return q.new_zeros((q.shape[0], self.padded_heads, q.shape[2]))
+            return F.pad(
+                q,
+                (0, 0, 0, self.padded_heads - self.n_local_heads),
+                value=0.0,
+            )
 
         swa_metadata = cast(
             "DeepseekSparseSWAMetadata | None",
@@ -865,22 +1020,38 @@ def _fused_qnorm_rope_kv_insert(
         if cache_dtype == torch.uint8:
             # fp8_ds_mla UE8M0 paged path. Horizontally fused:
             #   Q side: GPT-J RoPE, zero-filling the padding head slots; the
-            #           kernel allocates and returns the padded q tensor.
+            #           kernel allocates and returns the padded q tensor. An
+            #           interleaved Q skips the RoPE its attention kernel owns
+            #           and keeps only the pad, which q_head_padded=0 drops
+            #           too once the shard is already padded_heads wide.
             #   KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
             swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
-            return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+            pad_to = (
+                0
+                if self.accepts_unnormed_unroped_query
+                and self.n_local_heads == self.padded_heads
+                else self.padded_heads
+            )
+            q_padded = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                 q,
                 kv,
                 swa_kv_cache_2d,
                 swa_metadata.slot_mapping,
                 positions,
                 cos_sin_cache,
-                self.padded_heads,
+                pad_to,
                 self.eps,
                 swa_metadata.block_size,
-                False,
+                False,  # apply_q_norm: qr is normed before wq_b
+                self.kv_mxfp8,
+                not self.accepts_unnormed_unroped_query,  # apply_q_rope
+                self.accepts_unnormed_unroped_query,  # is_q_interleaved
             )
+            return q if pad_to == 0 else q_padded
 
+        assert not self.accepts_unnormed_unroped_query, (
+            "the chunk-interleaved Q layout only pairs with a packed KV record"
+        )
         # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
         # row in its element dtype (no Q padding). bf16 rewrites q in place;
         # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
@@ -934,10 +1105,10 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
         # DeepseekV4SWACache.
         if not self.is_kv_source:
             return None
-        # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B
-        # alignment; plain bf16 / per-tensor fp8 rows use natural element-size
-        # pages.
-        uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
+        # fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up
+        # to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows
+        # use natural element-size pages.
+        uses_fp8_ds_mla_layout = self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla")
         return MLAAttentionSpec(
             block_size=vllm_config.cache_config.block_size,
             num_kv_heads=1,
@@ -945,12 +1116,13 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
             dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
             tokens_per_state=self.compress_ratio,
             cache_dtype_str=self.kv_cache_dtype,
-            alignment=576 if uses_fp8_ds_mla_layout else 512,
+            alignment=self.kv_page_alignment if uses_fp8_ds_mla_layout else 512,
             model_version="deepseek_v4",
             kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
-            # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token;
-            # head_size stays semantic (512).
-            state_content_bytes=584 if uses_fp8_ds_mla_layout else None,
+            # Packed record width; head_size stays semantic (512).
+            state_content_bytes=(
+                self.compressed_bytes_per_token if uses_fp8_ds_mla_layout else None
+            ),
         )
 
     def _compressed_kv_cache(self) -> torch.Tensor:
@@ -979,7 +1151,9 @@ def __init__(
         self.cache_config = cache_config
         self.dtype = dtype
         self.compress_ratio = compress_ratio
-        compilation_config = get_current_vllm_config().compilation_config
+        vllm_config = get_current_vllm_config()
+        self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits
+        compilation_config = vllm_config.compilation_config
         if prefix in compilation_config.static_forward_context:
             raise ValueError(f"Duplicate layer name: {prefix}")
         compilation_config.static_forward_context[prefix] = self
@@ -991,21 +1165,42 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         # head_dim already carries the fp8 scale padding
         # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout.
-        uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
+        # nvfp4_ds_mla is packed too: its compressed record is NVFP4 but the
+        # sliding-window record stays the V4.1 MXFP8 one, so the indexer page
+        # takes the same alignment either way.
+        uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype in (
+            "fp8_ds_mla",
+            "nvfp4_ds_mla",
+        )
+        page_alignment = (
+            576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512
+        )
         return MLAAttentionSpec(
             block_size=self.cache_config.block_size,
             num_kv_heads=1,
             head_size=self.head_dim,
             dtype=self.dtype,
             tokens_per_state=self.compress_ratio,
-            # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
-            alignment=576 if uses_fp8_ds_mla_layout else 512,
+            # The indexer page shares a physical block with the main KV cache
+            # (the block stride is the sum of every page in the group), so it
+            # has to carry the same padding: the packed record's TMA stride,
+            # else 512B for FlashInfer sparse (#44577).
+            alignment=page_alignment,
+            # DeepGEMM's paged sparse MQA logits address pages by the block
+            # stride, which packs every layer's page in this layout, and need
+            # it 512B-aligned; the main KV record read from the same blocks
+            # needs its own TMA stride, so keep both.
+            block_stride_alignment=(
+                math.lcm(512, page_alignment) if self.sparse_logits else None
+            ),
         )
 
     def forward(self): ...
 
     def get_attn_backend(self) -> type[AttentionBackend]:
-        return DeepseekV4IndexerBackend
+        if self.sparse_logits:
+            return DeepseekV41SparseIndexerBackend
+        return DeepseekV41IndexerBackend
 
 
 class DeepseekV4Indexer(nn.Module):
@@ -1104,21 +1299,52 @@ def __init__(
             self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
         self.k_cache = k_cache
 
-        self.indexer_op = SparseAttnIndexer(
-            self.k_cache,
-            self.quant_block_size,
-            self.scale_fmt,
-            self.topk_tokens,
-            self.head_dim,
-            self.max_model_len,
-            self.max_total_seq_len,
-            self.topk_indices_buffer,
-            skip_k_cache_insert=True,
-            use_fp4_cache=self.use_fp4_kv,
-            compress_ratio=self.compress_ratio,
-            candidate_blocks=candidate_block_buffer,
-            candidate_block_size=candidate_block_size,
-            candidate_write=candidate_write,
+        # Candidate consumers can score only the candidate blocks (opt-in);
+        # the candidate source and pre-candidate indexers stay dense.
+        use_sparse_logits = (
+            vllm_config.attention_config.indexer_sparse_logits
+            and candidate_block_buffer is not None
+            and not candidate_write
+        )
+        self.indexer_op: SparseAttnIndexer | SparseMQAIndexer
+        if use_sparse_logits:
+            if not self.use_fp4_kv:
+                raise ValueError(
+                    "attention_config.indexer_sparse_logits requires the MXFP4 "
+                    "indexer cache (attention_config.indexer_kv_dtype='mxfp4')."
+                )
+            assert candidate_block_buffer is not None
+            assert topk_indices_buffer is not None
+            self.indexer_op = SparseMQAIndexer(
+                self.k_cache,
+                self.topk_tokens,
+                self.head_dim,
+                self.max_total_seq_len,
+                topk_indices_buffer,
+                candidate_block_buffer,
+                candidate_block_size,
+            )
+        else:
+            self.indexer_op = SparseAttnIndexer(
+                self.k_cache,
+                self.quant_block_size,
+                self.scale_fmt,
+                self.topk_tokens,
+                self.head_dim,
+                self.max_model_len,
+                self.max_total_seq_len,
+                self.topk_indices_buffer,
+                skip_k_cache_insert=True,
+                use_fp4_cache=self.use_fp4_kv,
+                compress_ratio=self.compress_ratio,
+                candidate_blocks=candidate_block_buffer,
+                candidate_block_size=candidate_block_size,
+                candidate_write=candidate_write,
+            )
+        # The fused Q kernel writes the per-head weights in the dtype the
+        # scoring kernels take, so no cast runs per step.
+        self.indexer_weights_dtype = (
+            SparseMQAIndexer.weights_dtype if use_sparse_logits else torch.float32
         )
 
     def _produce_k(
@@ -1206,6 +1432,7 @@ def forward(
             self.softmax_scale,
             self.n_head**-0.5,
             use_fp4=self.use_fp4_kv,
+            weights_out_dtype=self.indexer_weights_dtype,
         )
         if isinstance(q_quant, tuple):
             q, q_scale = q_quant
diff --git a/vllm/models/deepseek_v41/common/__init__.py b/vllm/models/deepseek_v41/common/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/deepseek_v41/common/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/models/deepseek_v4_1/common/engram.py b/vllm/models/deepseek_v41/common/engram.py
similarity index 89%
rename from vllm/models/deepseek_v4_1/common/engram.py
rename to vllm/models/deepseek_v41/common/engram.py
index f30c3a5606d1..c415bbaa6e42 100644
--- a/vllm/models/deepseek_v4_1/common/engram.py
+++ b/vllm/models/deepseek_v41/common/engram.py
@@ -51,8 +51,6 @@
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.triton_utils import tl, triton
-from vllm.utils.platform_utils import is_uva_available
-from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
 
 logger = init_logger(__name__)
 
@@ -236,6 +234,7 @@ def _write_hash_cache_kernel(
     tl.store(cache + slot, value, valid)
 
 
+# Keep request shapes out of the cache key so one warmed variant covers runtime batches.
 @triton.jit(
     do_not_specialize=[
         "num_tokens",
@@ -462,6 +461,21 @@ def ensure_cache(self) -> bool:
         self._kv_cache_ref = weakref.ref(kv_cache)
         return True
 
+    def dummy_hashes(
+        self, input_ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Participate in DP lookups without valid rows or hash-cache updates."""
+        num_tokens = input_ids.shape[0]
+        num_layers, max_ngram = self.multipliers.shape
+        num_heads = self.primes.shape[-1]
+        hashes = input_ids.new_full(
+            (num_tokens, num_layers, (max_ngram - 1) * num_heads),
+            DEAD_ID,
+            dtype=torch.int32,
+        )
+        keep = torch.zeros(num_tokens, dtype=torch.bool, device=input_ids.device)
+        return hashes, keep
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -561,7 +575,16 @@ def _engram_head_shard_weight_loader(
     param.data.copy_(shard)
 
 
-@triton.jit
+@triton.jit(
+    do_not_specialize=[
+        "vocab_start",
+        "vocab_end",
+        "num_rows",
+        "ids_stride_t",
+        "ids_stride_h",
+        "GRID",
+    ]
+)
 def _engram_lookup_kernel(
     weight,
     scales,
@@ -578,7 +601,7 @@ def _engram_lookup_kernel(
     DIM: tl.constexpr,
     QUANT_BLOCK: tl.constexpr,
     BLOCK_R: tl.constexpr,
-    GRID: tl.constexpr,
+    GRID,
 ):
     """Gather fp8 rows, apply their ue8m0 block scales, write bf16.
 
@@ -620,12 +643,9 @@ def _engram_lookup_kernel(
 
 
 class ParallelEngramEmbedding(nn.Module):
-    """The n-gram hash table, sharded by complete hash heads over TP ranks.
-    Rows stay fp8 and are dequantized with ue8m0 per-32 scales on lookup.
+    """TP-sharded hash heads with FP8 rows and per-block E8M0 scales."""
 
-    With `cpu_offload` the shard lives in pinned host memory and is read over
-    UVA instead of HBM; the TP sharding is unchanged either way.
-    """
+    _weight_loader = staticmethod(_engram_head_shard_weight_loader)
 
     def __init__(
         self,
@@ -633,85 +653,55 @@ def __init__(
         dim: int,
         head_sizes: tuple[int, ...],
         block_size: int = 32,
-        cpu_offload: bool = False,
-    ):
+    ) -> None:
         super().__init__()
-        tp_size = get_tensor_model_parallel_world_size()
-        tp_rank = get_tensor_model_parallel_rank()
         assert head_sizes and all(size > 0 for size in head_sizes)
         assert sum(head_sizes) <= num_embeddings
-        if cpu_offload and not is_uva_available():
-            raise RuntimeError("Engram CPU offload requires UVA support")
         self.num_embeddings = num_embeddings
         self.dim = dim
         self.block_size = block_size
         self.n_hash_cols = len(head_sizes)
-        self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, tp_size)
-        self.head_start = tp_rank * self.part_n_hash_cols
+        self.tp_size = get_tensor_model_parallel_world_size()
+        num_shards, head_rank = self._get_shard_info()
+        self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, num_shards)
+        # TODO: Support row-wise sharding when there are too few hash heads.
+        assert (num_shards - 1) * self.part_n_hash_cols < self.n_hash_cols, (
+            f"Engram sharding leaves ranks without hash heads: "
+            f"{self.n_hash_cols} heads over {num_shards} shards"
+        )
+        self.head_start = head_rank * self.part_n_hash_cols
         head_end = self.head_start + self.part_n_hash_cols
         self.vocab_start_idx = sum(head_sizes[: self.head_start])
         self.vocab_end_idx = sum(head_sizes[:head_end])
         self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx
-        self.tp_size = tp_size
-        self.cpu_offload = cpu_offload
-        self._views: tuple[torch.Tensor, torch.Tensor] | None = None
-        self._view_src: tuple[int, int] | None = None
         self._num_sms = torch.cuda.get_device_properties(
             torch.accelerator.current_device_index()
         ).multi_processor_count
-
-        # Explicit device: model init runs under a `torch.device("cuda")`
-        # context, which would otherwise put the shard in HBM.
-        kwargs = {"device": "cpu", "pin_memory": True} if cpu_offload else {}
-        self.weight = nn.Parameter(
-            torch.empty(
-                self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn, **kwargs
-            ),
-            requires_grad=False,
-        )
-        self.weight_scale_inv = nn.Parameter(
-            torch.empty(
-                self.part_num_embeddings,
-                dim // block_size,
-                dtype=torch.uint8,
-                **kwargs,
-            ),
-            requires_grad=False,
-        )
+        weight, scales = self._allocate_weights()
+        self.weight = nn.Parameter(weight, requires_grad=False)
+        self.weight_scale_inv = nn.Parameter(scales, requires_grad=False)
         for param in (self.weight, self.weight_scale_inv):
             set_weight_attrs(
                 param,
                 {
-                    "weight_loader": _engram_head_shard_weight_loader,
+                    "weight_loader": self._weight_loader,
                     "engram_vocab_start": self.vocab_start_idx,
                 },
             )
-        if cpu_offload:
-            logger.info(
-                "Engram table offloaded to pinned host memory: %d rows x %d, "
-                "%.2f GiB per rank",
-                self.part_num_embeddings,
-                dim,
-                self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
-            )
 
-    def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
-        """Parameters when resident, else cached UVA views of the pinned shard.
+    def _get_shard_info(self) -> tuple[int, int]:
+        return self.tp_size, get_tensor_model_parallel_rank()
 
-        Rebuilt if anything swaps `.data`, so a stale device pointer cannot
-        survive silently.
-        """
-        if not self.cpu_offload:
-            return self.weight.data, self.weight_scale_inv.data
-        src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
-        if self._view_src != src:
-            self._views = (
-                get_accelerator_view_from_cpu_tensor(self.weight.data),
-                get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
-            )
-            self._view_src = src
-        assert self._views is not None
-        return self._views
+    def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
+        return (
+            torch.empty(self.part_num_embeddings, self.dim, dtype=torch.float8_e4m3fn),
+            torch.empty(
+                self.part_num_embeddings, self.dim // self.block_size, dtype=torch.uint8
+            ),
+        )
+
+    def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.weight.data, self.weight_scale_inv.data
 
     def lookup(
         self, indices: torch.Tensor, out: torch.Tensor, background: bool = False
@@ -749,7 +739,7 @@ def lookup(
 
     def forward(self, indices: torch.Tensor) -> torch.Tensor:
         """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim]
-        bf16, gathered from all TP shards."""
+        bf16, gathered from all shards for this replica's tokens."""
         out = torch.empty(
             (indices.shape[0], self.part_n_hash_cols, self.dim),
             dtype=torch.bfloat16,
@@ -758,8 +748,7 @@ def forward(self, indices: torch.Tensor) -> torch.Tensor:
         self.lookup(indices, out)
         if self.tp_size > 1:
             out = tensor_model_parallel_all_gather(out, dim=1)
-            out = out[:, : self.n_hash_cols]
-        return out
+        return out[:, : self.n_hash_cols]
 
 
 @triton.jit(do_not_specialize=["num_kv_tokens"])
@@ -855,7 +844,7 @@ def _fused_engram_post_wkv_kernel(
 
 
 @triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"])
-def _engram_sp_rows_kernel(
+def _engram_select_rows_kernel(
     gathered,
     output,
     num_tokens,
@@ -876,6 +865,32 @@ def _engram_sp_rows_kernel(
     tl.store(output + offsets, values, offsets < num_elements)
 
 
+def _engram_select_rows(
+    gathered: torch.Tensor,
+    output: torch.Tensor,
+    source_tokens: int,
+    token_start: int,
+    local_width: int,
+) -> None:
+    """Copy one token window out of a rank-major gathered buffer.
+
+    Both gathers land rank-major ([rank][token][local width]); this walks the
+    window the rank keeps and lays its ranks out side by side as width.
+    """
+    if output.numel() == 0:
+        return
+    _engram_select_rows_kernel[(triton.cdiv(output.numel(), 1024),)](
+        gathered,
+        output,
+        source_tokens,
+        token_start,
+        output.numel(),
+        local_width,
+        output.shape[1] * output.shape[2],
+        BLOCK_SIZE=1024,
+    )
+
+
 class Engram(nn.Module):
     """Writes an n-gram lookup into the residual stream, gated by how well it
     matches that stream.
@@ -906,13 +921,7 @@ def __init__(
         # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight``
         # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight``
         # suffix rule.
-        engram_config = get_current_vllm_config().engram_config
-        self.embed_tokens = ParallelEngramEmbedding(
-            layout.num_embeddings[layer_hash_index],
-            layout.head_dim,
-            tuple(size for order in layout.primes[layer_hash_index] for size in order),
-            cpu_offload=engram_config.cpu_offload if engram_config else True,
-        )
+        self.embed_tokens = self._create_embedding(layout, layer_hash_index)
         n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
         self.wkv = ReplicatedLinear(
             n_hash_cols * layout.head_dim,
@@ -932,40 +941,54 @@ def __init__(
         )
 
         max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens
-        # Keep lookup results alive across breakable graph segments.
+        self._init_staging(max_tokens, layout.head_dim)
+
+    def _create_embedding(
+        self, layout: EngramLayout, layer_hash_index: int
+    ) -> ParallelEngramEmbedding:
+        return ParallelEngramEmbedding(
+            layout.num_embeddings[layer_hash_index],
+            layout.head_dim,
+            tuple(size for order in layout.primes[layer_hash_index] for size in order),
+        )
+
+    def _init_staging(self, max_tokens: int, head_dim: int) -> None:
+        # Persistent storage keeps lookup addresses stable across graph replays.
         self.staged_rows = torch.empty(
             max_tokens,
             self.embed_tokens.part_n_hash_cols,
-            layout.head_dim,
+            head_dim,
             dtype=torch.bfloat16,
         )
 
     def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
-        """Gather this layer's rows on the main stream before decoder layers."""
-        self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]])
+        """Look up rows before the decoder layers consume them."""
+        rows = self.staged_rows[: hash_ids.shape[0]]
+        assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small"
+        self.embed_tokens.lookup(hash_ids, rows)
+
+    def _ready_rows(self, num_tokens: int) -> torch.Tensor:
+        return self.staged_rows[:num_tokens]
 
     def embed(self, hash_ids: torch.Tensor) -> torch.Tensor:
         """Gather heads, returning only local tokens when SP is enabled."""
-        rows = self.staged_rows[: hash_ids.shape[0]]
+        rows = self._ready_rows(hash_ids.shape[0])
         if self.embed_tokens.tp_size == 1:
-            return rows
+            return rows[:, : self.embed_tokens.n_hash_cols]
         if self.use_sequence_parallel:
             tp_size = self.embed_tokens.tp_size
             num_tokens, local_heads, dim = rows.shape
             gathered = tensor_model_parallel_all_gather(rows, dim=0)
             chunk = (num_tokens + tp_size - 1) // tp_size
-            rows = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim))
-            _engram_sp_rows_kernel[(triton.cdiv(rows.numel(), 1024),)](
+            out = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim))
+            _engram_select_rows(
                 gathered,
-                rows,
+                out,
                 num_tokens,
                 get_tensor_model_parallel_rank() * chunk,
-                rows.numel(),
                 local_heads * dim,
-                self.embed_tokens.n_hash_cols * dim,
-                BLOCK_SIZE=1024,
             )
-            return rows
+            return out
         rows = tensor_model_parallel_all_gather(rows, dim=1)
         return rows[:, : self.embed_tokens.n_hash_cols]
 
diff --git a/vllm/models/deepseek_v4_1/common/mm_preprocess.py b/vllm/models/deepseek_v41/common/mm_preprocess.py
similarity index 71%
rename from vllm/models/deepseek_v4_1/common/mm_preprocess.py
rename to vllm/models/deepseek_v41/common/mm_preprocess.py
index 8bd8d70caea3..7710edb43f85 100644
--- a/vllm/models/deepseek_v4_1/common/mm_preprocess.py
+++ b/vllm/models/deepseek_v41/common/mm_preprocess.py
@@ -12,27 +12,21 @@
 rows in reading order; the delimiters take the learned ``image_start`` /
 ``image_newline`` / ``image_end`` vectors.
 
-One vLLM-side deviation from the reference token stream: a leading
-compressor-alignment pad (``COMPRESS_PAD_TO - 1 - start % COMPRESS_PAD_TO``
-positions, so the span always starts at the same compressor phase) is
-prepended when the block is spliced into the final prompt. Pad positions
-borrow the reserved in-vocab token ``<|place_holder_mm_span_0436|>`` so
-they stay distinguishable from real span positions; they embed as the
-plain image token (v4.1 has no ``image_pad`` vector) and are routed and
-engram-deadened like image tokens.
+The token stream matches the reference exactly: no compressor-alignment
+pad is inserted (the reference pools image tokens across compressor-group
+boundaries freely).
 """
 
 import math
 from collections.abc import Mapping, Sequence
-from typing import Any, cast
+from typing import Any
 
 import numpy as np
 import torch
 from PIL import Image, ImageOps
 from transformers import BatchFeature
-from typing_extensions import assert_never
 
-from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
 from vllm.multimodal.parse import ImageSize, MultiModalDataItems
@@ -44,42 +38,26 @@
     PromptUpdate,
     PromptUpdateDetails,
 )
-from vllm.multimodal.processing.processor import (
-    MultiModalPromptUpdates,
-    MultiModalPromptUpdatesApplyResult,
-    PlaceholderFeaturesInfo,
-    UpdateMode,
-    _plan_prompt_updates,
-)
 from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
 
 IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4)
-# Image spans are padded so they always start at the same compressor group
-# phase; v4.1 uses ratio-2 compressors (v4.0 used 4).
-COMPRESS_PAD_TO = 2
 
 IMAGE_PLACEHOLDER = "<｜deepseek_image｜>"
 
 # The checkpoint tokenizer's id for IMAGE_PLACEHOLDER (the config's
 # ``image_token_id``). Every image-span position carries this id, exactly
 # like the reference; the MoE router detects image tokens as the range
-# [IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID + 5), which also covers
-# IMAGE_PAD_ID below.
+# [IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID + 5).
 IMAGE_SENTINEL_BASE_ID = 129264
-# Leading compressor-alignment pads borrow this reserved in-vocab token
-# (``<|place_holder_mm_span_0436|>``) so they can be told apart from real
-# span positions when the per-position role is decided from token ids.
-IMAGE_PAD_ID = 129265
-IMAGE_PAD_TOKEN_NAME = "<|place_holder_mm_span_0436|>"
 
 
 def image_sentinel_mask(token_ids: torch.Tensor) -> torch.Tensor:
-    """Boolean mask for image-span positions (span tokens and align pads)."""
-    return (token_ids == IMAGE_SENTINEL_BASE_ID) | (token_ids == IMAGE_PAD_ID)
+    """Boolean mask for image-span positions."""
+    return token_ids == IMAGE_SENTINEL_BASE_ID
 
 
 def validate_image_sentinel_ids(tokenizer) -> None:
-    """Check the image/pad token ids against the tokenizer."""
+    """Check the image token id against the tokenizer."""
     image_id = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER)
     if image_id != IMAGE_SENTINEL_BASE_ID:
         raise ValueError(
@@ -88,13 +66,6 @@ def validate_image_sentinel_ids(tokenizer) -> None:
             "image_token_id); the DeepSeek-V4.1 vision path keys image "
             "routing and engram masking off this id."
         )
-    pad_id = tokenizer.convert_tokens_to_ids(IMAGE_PAD_TOKEN_NAME)
-    if pad_id != IMAGE_PAD_ID:
-        raise ValueError(
-            f"Image pad token {IMAGE_PAD_TOKEN_NAME!r} has id {pad_id}, "
-            f"expected {IMAGE_PAD_ID}; the DeepSeek-V4.1 vision path "
-            "borrows this reserved id for compressor-alignment pads."
-        )
 
 
 def llm_grid(best_height, best_width, patch_size, downsample_ratio):
@@ -134,8 +105,7 @@ def safe_resize(
     height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token
 ):
     """Shrink the pixel size until the image costs at most max_n_token LLM
-    tokens (minus the reservation for the compressor-alignment pad)."""
-    max_n_token -= COMPRESS_PAD_TO - 1
+    tokens."""
     n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio)
     if num_image_tokens(n_llm_h, n_llm_w) > max_n_token:
         best_height, best_width = solve_resize_ratio(
@@ -293,13 +263,7 @@ def get_mm_max_tokens_per_item(
         seq_len: int,
         mm_counts: Mapping[str, int],
     ) -> Mapping[str, int]:
-        # ``safe_resize`` reserves COMPRESS_PAD_TO - 1 tokens of
-        # vision_max_n_token for the compressor-alignment pad, so the full
-        # image span is bounded by vision_max_n_token; the margin is kept
-        # in case that reservation changes.
-        return {
-            "image": self.get_hf_config().vision_max_n_token + COMPRESS_PAD_TO - 1,
-        }
+        return {"image": self.get_hf_config().vision_max_n_token}
 
     def get_image_placeholder_token_id(self) -> int:
         token_id = self.get_tokenizer().convert_tokens_to_ids(IMAGE_PLACEHOLDER)
@@ -314,7 +278,7 @@ def get_image_size_with_most_features(self) -> ImageSize:
         # A square maximizes the ViT patch count (area) within the token
         # budget; solve the budget-derived size directly to keep the dummy
         # image small.
-        budget = hf_config.vision_max_n_token - (COMPRESS_PAD_TO - 1)
+        budget = hf_config.vision_max_n_token
         side = budget * patch_size * downsample_ratio
         best_h, best_w = solve_resize_ratio(
             side, side, patch_size, downsample_ratio, budget
@@ -332,7 +296,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         size = self.info.get_image_size_with_most_features()
         return {
@@ -340,7 +304,7 @@ def get_dummy_mm_data(
                 width=size.width,
                 height=size.height,
                 num_images=mm_counts.get("image", 0),
-                overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
+                overrides=mm_options.get("image"),
             ),
         }
 
@@ -397,67 +361,3 @@ def get_image_replacement(item_idx: int) -> PromptUpdateDetails:
                 replacement=get_image_replacement,
             ),
         ]
-
-    def _apply_token_matches_with_placeholders(
-        self,
-        token_ids: list[int],
-        mm_prompt_updates: MultiModalPromptUpdates,
-    ) -> tuple[
-        list[int],
-        MultiModalPromptUpdatesApplyResult,
-        Mapping[str, list[PlaceholderFeaturesInfo]],
-    ]:
-        # Same as the base implementation, except that each image block gets
-        # its compressor-alignment pad (``COMPRESS_PAD_TO - 1 - start %
-        # COMPRESS_PAD_TO`` IMAGE_PAD_ID tokens) prepended while splicing:
-        # the pad depends on the block's final position in the prompt, which
-        # is unknown when the (cacheable) prompt updates are built. Pad
-        # positions are not embed positions (their id differs from
-        # image_token_id), so ``is_embed`` stays aligned with the pad-free
-        # replacement content.
-        matched_updates, result = _plan_prompt_updates(token_ids, mm_prompt_updates)
-        placeholders: dict[str, list[PlaceholderFeaturesInfo]] = {
-            modality: [] for modality in mm_prompt_updates
-        }
-
-        new_token_ids = list[int]()
-        prev_end_idx = 0
-        for matched_update in matched_updates:
-            update = matched_update.update
-            match = matched_update.match
-
-            if update.mode == UpdateMode.INSERT:
-                end_idx_to_insert = match.end_idx
-            elif update.mode == UpdateMode.REPLACE:
-                end_idx_to_insert = match.start_idx
-            else:
-                assert_never(update.mode)
-
-            new_token_ids.extend(token_ids[prev_end_idx:end_idx_to_insert])
-            start_idx = len(new_token_ids)
-
-            tokens = list(update.content.full)
-            if tokens and update.modality == "image":
-                compress_pad = COMPRESS_PAD_TO - 1 - start_idx % COMPRESS_PAD_TO
-                tokens = [IMAGE_PAD_ID] * compress_pad + tokens
-
-            if tokens:
-                content_is_embed = update.content.is_embed
-                is_embed = (
-                    content_is_embed(tokens) if content_is_embed is not None else None
-                )
-                placeholders[update.modality].append(
-                    PlaceholderFeaturesInfo(
-                        modality=update.modality,
-                        item_idx=update.item_idx,
-                        start_idx=start_idx,
-                        tokens=tokens,
-                        is_embed=is_embed,
-                    )
-                )
-                new_token_ids.extend(tokens)
-
-            prev_end_idx = match.end_idx
-
-        new_token_ids.extend(token_ids[prev_end_idx:])
-        return new_token_ids, result, placeholders
diff --git a/vllm/models/deepseek_v4_1/common/ops/__init__.py b/vllm/models/deepseek_v41/common/ops/__init__.py
similarity index 100%
rename from vllm/models/deepseek_v4_1/common/ops/__init__.py
rename to vllm/models/deepseek_v41/common/ops/__init__.py
diff --git a/vllm/models/deepseek_v4_1/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py
similarity index 75%
rename from vllm/models/deepseek_v4_1/common/ops/cache_utils.py
rename to vllm/models/deepseek_v41/common/ops/cache_utils.py
index e2da6f5e2bed..ad03c3294524 100644
--- a/vllm/models/deepseek_v4_1/common/ops/cache_utils.py
+++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index
+"""Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index
 preparation.
 
 - quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into
@@ -36,6 +35,24 @@
 from vllm.utils.import_utils import has_cutedsl
 from vllm.utils.math_utils import next_power_of_2
 
+# Per-token byte width of the two paged fp8 records. Both put a page's whole
+# data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the
+# width the KV-cache spec publishes -- identifies the record:
+#   V4   (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64
+#                 dims each plus a pad byte.
+#   V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0
+#                 scales of 32 dims each. FlashMLA's ``ModelType::V41``.
+#   V4.1 NVFP4 (288 B): 512 e2m1 values packed two per byte (even element in
+#                 the low nibble), then 32 e4m3 scales of 16 dims each.
+#                 FlashMLA's ``ModelType::V41_FP4``, compressed cache only.
+V4_BYTES_PER_TOKEN = 584
+V41_BYTES_PER_TOKEN = 528
+V41_QUANT_BLOCK = 32
+V41_NUM_SCALES = 512 // V41_QUANT_BLOCK  # 16
+V41_NVFP4_BYTES_PER_TOKEN = 288
+V41_NVFP4_QUANT_BLOCK = 16
+V41_NVFP4_NUM_SCALES = 512 // V41_NVFP4_QUANT_BLOCK  # 32
+
 
 @triton.jit
 def quantize_and_insert_k_kernel(
@@ -58,8 +75,7 @@ def quantize_and_insert_k_kernel(
     n_quant_blocks: tl.constexpr,  # 8 (7 real + 1 padding)
     use_fnuz: tl.constexpr = False,
 ):
-    """
-    Quantize K tensor and insert into paged K cache.
+    """Quantize K tensor and insert into paged K cache.
 
     K Cache block layout (block_size=64 tokens):
     - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16
@@ -163,6 +179,67 @@ def quantize_and_insert_k_kernel(
         tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)
 
 
+@triton.jit
+def _quantize_and_insert_k_mxfp8_kernel(
+    k_ptr,  # [num_tokens, 512] bf16
+    slot_mapping_ptr,  # [num_tokens] int64
+    k_cache_ptr,  # [num_blocks, block_bytes] as uint8 (flattened view)
+    num_tokens,
+    head_dim: tl.constexpr,  # 512, RoPE dims included
+    scale_dim: tl.constexpr,  # 16
+    quant_block: tl.constexpr,  # 32 dims per UE8M0 scale
+    cache_block_size: tl.constexpr,  # paged cache block size, in tokens
+    block_stride: tl.constexpr,  # total bytes per block (padded)
+    fp8_max: tl.constexpr,
+    use_fnuz: tl.constexpr = False,
+):
+    """Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache.
+
+    One program per token. Unlike the V4 record there is no bf16 tail, so every
+    dim is scaled.
+    """
+    pid = tl.program_id(0)
+    if pid >= num_tokens:
+        return
+
+    slot_idx = tl.load(slot_mapping_ptr + pid)
+    if slot_idx == -1:
+        return
+
+    block_idx = slot_idx // cache_block_size
+    pos_in_block = slot_idx % cache_block_size
+
+    # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks.
+    cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride
+    token_data_ptr = cache_block_ptr + pos_in_block * head_dim
+    token_scale_ptr = (
+        cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim
+    )
+
+    d = tl.arange(0, head_dim)
+    x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32)
+
+    tiles = tl.reshape(x, (scale_dim, quant_block))
+    # Floor the absmax like the CUDA encoder so an all-zero tile still lands on
+    # a finite scale.
+    amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4)
+    exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max)))
+    scaled = tl.clamp(
+        tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max
+    )
+    if use_fnuz:  # noqa: SIM108
+        fp8 = scaled.to(tl.float8e4b8)
+    else:
+        fp8 = scaled.to(tl.float8e4nv)
+    tl.store(
+        token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,))
+    )
+
+    # UE8M0 encoding: stored_value = exponent + 127 (bias).
+    encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0)
+    tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8))
+
+
 def quantize_and_insert_k_cache(
     k: torch.Tensor,  # [num_tokens, 512] bf16
     k_cache: torch.Tensor,  # [num_blocks, block_bytes] uint8
@@ -170,16 +247,12 @@ def quantize_and_insert_k_cache(
     block_size: int = 64,
     is_ue8m0: bool = True,
     use_fnuz: bool = False,
+    bytes_per_token: int = V4_BYTES_PER_TOKEN,
 ):
-    """
-    Quantize K tensor and insert into paged K cache.
+    """Quantize K tensor and insert into paged K cache.
 
-    K Cache block layout (block_size=64 tokens):
-    - First 64 * 576 = 36864 bytes: Token data
-      - Each token: 448 bytes (fp8) + 128 bytes (bf16)
-    - Next 64 * 8 = 512 bytes: Scales
-      - Each token: 8 bytes (uint8 scales, 7 real + 1 padding)
-    - Padded to multiple of 576
+    ``bytes_per_token`` picks the record (see the module header): the V4 one,
+    or V4.1's all-dims MXFP8 one.
 
     ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on
     platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3,
@@ -196,6 +269,26 @@ def quantize_and_insert_k_cache(
     num_tokens = slot_mapping.shape[0]
     block_stride = k_cache.stride(0)  # bytes per block
 
+    if bytes_per_token == V41_BYTES_PER_TOKEN:
+        if use_fnuz and not current_platform.is_fp8_fnuz():
+            raise ValueError("use_fnuz=True requires a platform using FNUZ FP8")
+        _quantize_and_insert_k_mxfp8_kernel[(num_tokens,)](
+            k,
+            slot_mapping,
+            k_cache,
+            num_tokens,
+            head_dim=512,
+            scale_dim=V41_NUM_SCALES,
+            quant_block=V41_QUANT_BLOCK,
+            cache_block_size=block_size,
+            block_stride=block_stride,
+            fp8_max=get_fp8_min_max()[1]
+            if use_fnuz
+            else torch.finfo(torch.float8_e4m3fn).max,
+            use_fnuz=use_fnuz,
+        )
+        return
+
     TOKEN_FP8_DIM = 448
     TOKEN_BF16_DIM = 64
     TOKEN_SCALE_DIM = 8
@@ -343,6 +436,136 @@ def _dequantize_and_gather_k_kernel(
             tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals)
 
 
+@triton.jit
+def _dequantize_and_gather_k_mxfp8_kernel(
+    out_ptr,
+    out_stride0,
+    out_stride1,
+    k_cache_ptr,
+    seq_lens_ptr,
+    block_table_ptr,
+    offset,
+    gather_lens_ptr,
+    max_blocks_per_seq: tl.constexpr,
+    head_dim: tl.constexpr,  # 512
+    scale_dim: tl.constexpr,  # 16
+    quant_block: tl.constexpr,  # 32
+    cache_block_size: tl.constexpr,
+    block_stride: tl.constexpr,
+    use_fnuz: tl.constexpr = False,
+):
+    """Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place."""
+    batch_idx = tl.program_id(0)
+    worker_id = tl.program_id(1)
+    num_workers = tl.num_programs(1)
+
+    seq_len = tl.load(seq_lens_ptr + batch_idx)
+    if gather_lens_ptr is not None:  # noqa: SIM108
+        gather_len = tl.load(gather_lens_ptr + batch_idx)
+    else:
+        gather_len = seq_len
+    start_pos = seq_len - gather_len
+
+    d = tl.arange(0, head_dim)
+    s = tl.arange(0, scale_dim)
+    for i in range(worker_id, gather_len, num_workers):
+        pos = start_pos + i
+        block_in_seq = pos // cache_block_size
+        pos_in_block = pos % cache_block_size
+
+        block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
+        physical_block_idx = tl.load(block_table_row_ptr + block_in_seq)  # int32
+
+        # int64: physical_block_idx * block_stride can exceed 2^31.
+        cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
+        token_data_ptr = cache_block_ptr + pos_in_block * head_dim
+        token_scale_ptr = (
+            cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim
+        )
+
+        x_uint8 = tl.load(token_data_ptr + d)
+        if use_fnuz:
+            x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
+        else:
+            x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
+        tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block))
+
+        # UE8M0: scale = 2^(stored_value - 127).
+        encoded = tl.load(token_scale_ptr + s)
+        scale = tl.exp2(encoded.to(tl.float32) - 127.0)
+        dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,))
+
+        output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
+        tl.store(output_row_ptr + d, dequant.to(tl.bfloat16))
+
+
+@triton.jit
+def _dequantize_and_gather_k_nvfp4_kernel(
+    out_ptr,
+    out_stride0,
+    out_stride1,
+    k_cache_ptr,
+    seq_lens_ptr,
+    block_table_ptr,
+    offset,
+    gather_lens_ptr,
+    max_blocks_per_seq: tl.constexpr,
+    head_dim: tl.constexpr,  # 512
+    scale_dim: tl.constexpr,  # 32
+    quant_block: tl.constexpr,  # 16
+    cache_block_size: tl.constexpr,
+    block_stride: tl.constexpr,
+):
+    """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace."""
+    batch_idx = tl.program_id(0)
+    worker_id = tl.program_id(1)
+    num_workers = tl.num_programs(1)
+
+    seq_len = tl.load(seq_lens_ptr + batch_idx)
+    if gather_lens_ptr is not None:  # noqa: SIM108
+        gather_len = tl.load(gather_lens_ptr + batch_idx)
+    else:
+        gather_len = seq_len
+    start_pos = seq_len - gather_len
+
+    packed_bytes: tl.constexpr = head_dim // 2
+    d = tl.arange(0, head_dim)
+    for i in range(worker_id, gather_len, num_workers):
+        pos = start_pos + i
+        block_in_seq = pos // cache_block_size
+        pos_in_block = pos % cache_block_size
+
+        physical_block_idx = tl.load(
+            block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq
+        )
+        page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
+
+        packed = tl.load(
+            page + pos_in_block * packed_bytes + tl.arange(0, packed_bytes)
+        )
+        # Even element in the low nibble, odd in the high nibble.
+        codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32))
+        # e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6.
+        mag_code = codes & 7
+        e = (mag_code >> 1).to(tl.float32)
+        m = (mag_code & 1).to(tl.float32)
+        mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0))
+        vals = tl.where(codes >= 8, -mag, mag)
+
+        sf = tl.load(
+            page
+            + cache_block_size * packed_bytes
+            + pos_in_block * scale_dim
+            + tl.arange(0, scale_dim)
+        )
+        scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32)
+        tiles = tl.reshape(vals, (scale_dim, quant_block))
+        dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,))
+
+        output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
+        tl.store(output_row_ptr + d, dequant.to(tl.bfloat16))
+
+
 def dequantize_and_gather_k_cache_triton(
     # [num_reqs, max_num_tokens, head_size]
     out: torch.Tensor,
@@ -358,6 +581,51 @@ def dequantize_and_gather_k_cache_triton(
     offset: int,
     use_fnuz: bool = False,
 ) -> None:
+    num_reqs = seq_lens.shape[0]
+    NUM_WORKERS = 128
+    if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN:
+        # Keep short gathers unchanged; cap total workers for multi-request chunks.
+        num_workers = min(
+            max(128, 2048 // max(1, num_reqs)), max(128, out.shape[1] - offset)
+        )
+        _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, num_workers)](
+            out,
+            out.stride(0),
+            out.stride(1),
+            k_cache,
+            seq_lens,
+            block_table,
+            offset,
+            gather_lens,
+            max_blocks_per_seq=block_table.shape[-1],
+            head_dim=512,
+            scale_dim=V41_NVFP4_NUM_SCALES,
+            quant_block=V41_NVFP4_QUANT_BLOCK,
+            cache_block_size=block_size,
+            block_stride=k_cache.stride(0),
+        )
+        return
+
+    if k_cache.shape[-1] == V41_BYTES_PER_TOKEN:
+        _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)](
+            out,
+            out.stride(0),
+            out.stride(1),
+            k_cache,
+            seq_lens,
+            block_table,
+            offset,
+            gather_lens,
+            max_blocks_per_seq=block_table.shape[-1],
+            head_dim=512,
+            scale_dim=V41_NUM_SCALES,
+            quant_block=V41_QUANT_BLOCK,
+            cache_block_size=block_size,
+            block_stride=k_cache.stride(0),
+            use_fnuz=use_fnuz,
+        )
+        return
+
     TOKEN_FP8_DIM = 448
     TOKEN_BF16_DIM = 64
     TOKEN_SCALE_DIM = 8
@@ -365,8 +633,6 @@ def dequantize_and_gather_k_cache_triton(
     FP8_MAX = 448.0
     TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
 
-    num_reqs = seq_lens.shape[0]
-    NUM_WORKERS = 128
     _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)](
         out,
         out.stride(0),
@@ -408,19 +674,32 @@ def dequantize_and_gather_k_cache(
 ) -> None:
     """Dequantize and gather a paged DSv4 K cache.
 
+    The record is read off ``k_cache.shape[-1]``; see the module header. Only
+    the fp8 records have a CuteDSL gather, so NVFP4 always takes the Triton
+    path.
+
     ``use_fnuz`` MUST match the encoder of the specific cache being read:
     ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere),
     ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder
     writes FNUZ on gfx942 and OCP on gfx950).
     """
-    if has_cutedsl():
+    if has_cutedsl() and k_cache.shape[-1] != V41_NVFP4_BYTES_PER_TOKEN:
         # lazily import, otherwise some tests fail due to CUDA driver init failure.
         from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
-            dequantize_and_gather_k_cache_cutedsl,
+            _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
         )
 
-        dequantize_and_gather_k_cache_cutedsl(
-            out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
+        mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN
+        _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL(
+            out=out,
+            k_cache=k_cache,
+            seq_lens=seq_lens,
+            gather_lens=gather_lens,
+            block_table=block_table,
+            block_size=block_size,
+            offset=offset,
+            fp8_dim=512 if mxfp8 else 448,
+            group_size=V41_QUANT_BLOCK if mxfp8 else 64,
         )
         return
 
@@ -540,16 +819,12 @@ def combine_topk_swa_indices(
     M: int,
     N: int,
     out: tuple[torch.Tensor, torch.Tensor] | None = None,
-    left_visible: torch.Tensor | None = None,
-    right_visible: torch.Tensor | None = None,
-    max_image_tokens: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     num_tokens = topk_indices.shape[0]
-    # max_image_tokens widens the SWA column region for in-image
-    # bidirectional visibility; the width is fixed per model so the
-    # caller-provided workspace stays valid on image-free batches.
+    # The width is fixed per model so the caller-provided workspace stays
+    # valid on every batch.
     combined_topk = (
-        (topk + window_size + max_image_tokens + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
+        (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
         // _SPARSE_PREFILL_TOPK_ALIGNMENT
         * _SPARSE_PREFILL_TOPK_ALIGNMENT
     )
@@ -578,9 +853,6 @@ def combine_topk_swa_indices(
         TOP_K=topk,
         COMPRESS_RATIO=compress_ratio,
         WINDOW_SIZE=window_size,
-        left_visible=left_visible,
-        right_visible=right_visible,
-        max_image_tokens=max_image_tokens,
     )
     return combined_indices, combined_lens
 
@@ -649,7 +921,6 @@ class CompileKey:
         TOP_K: int
         COMPRESS_RATIO: int
         WINDOW_SIZE: int
-        IMAGE_WIDTH: int
         PADDED_TOP_K: int
         input_variant: TritonPointerInputVariant
 
@@ -671,14 +942,11 @@ def kernel(
         query_start_loc_ptr,
         seq_lens_ptr,
         gather_lens_ptr,
-        left_visible_ptr,
-        right_visible_ptr,
         M,
         N,
         TOP_K: tl.constexpr,
         COMPRESS_RATIO: tl.constexpr,
         WINDOW_SIZE: tl.constexpr,
-        IMAGE_WIDTH: tl.constexpr,
         PADDED_TOP_K: tl.constexpr,
     ):
         batch_idx = tl.program_id(0)
@@ -713,19 +981,12 @@ def kernel(
                 # (integer div by constexpr 0 is UB and yields garbage, which
                 # turns the store index below negative).
                 topk_len = 0
-            if IMAGE_WIDTH > 0:
-                # In-image bidirectional visibility: the window starts up to
-                # max(left - (window - 1), 0) positions earlier and extends
-                # `right` positions past the query token (both 0 outside
-                # image spans, reducing to the plain causal window).
-                left = tl.load(left_visible_ptr + token_idx)
-                right = tl.load(right_visible_ptr + token_idx)
-            else:
-                left = 0
-                right = 0
-            left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
-            swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
-            swa_len = pos + right - swa_start + 1
+            swa_start = tl.maximum(pos - (WINDOW_SIZE - 1), 0)
+            # gather_len already excludes context below the request's replay
+            # start (SWA bounded replay), so the window cannot start before
+            # the gathered buffer does.
+            swa_start = tl.maximum(swa_start, gather_start)
+            swa_len = pos - swa_start + 1
 
             offset = tl.arange(0, PADDED_TOP_K)
             mask = offset < topk_len
@@ -739,18 +1000,17 @@ def kernel(
                 mask=mask,
             )
             # Index into gathered buffer: N + (position - gather_start)
-            # For positions [swa_start, pos + right], the buffer indices are:
-            # [N + swa_start - gather_start, N + pos + right - gather_start]
-            for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE):
-                swa_offset = i + tl.arange(0, WINDOW_SIZE)
-                tl.store(
-                    combined_indices_ptr
-                    + token_idx * combined_indices_stride
-                    + topk_len
-                    + swa_offset,
-                    M * batch_idx + N + swa_start + swa_offset - gather_start,
-                    mask=swa_offset < swa_len,
-                )
+            # For positions [swa_start, pos], the buffer indices are:
+            # [N + swa_start - gather_start, N + pos - gather_start]
+            swa_offset = tl.arange(0, WINDOW_SIZE)
+            tl.store(
+                combined_indices_ptr
+                + token_idx * combined_indices_stride
+                + topk_len
+                + swa_offset,
+                M * batch_idx + N + swa_start + swa_offset - gather_start,
+                mask=swa_offset < swa_len,
+            )
 
             combined_len = topk_len + swa_len
             tl.store(combined_lens_ptr + token_idx, combined_len)
@@ -766,7 +1026,6 @@ def dispatch(  # type: ignore[override]
         topk: int,
         compress_ratio: int,
         WINDOW_SIZE: int,
-        image_width: int = 0,
     ) -> CompileKey:
         padded_topk = next_power_of_2(topk_width)
         input_variant = TritonPointerInputVariant.from_alignment(
@@ -779,7 +1038,6 @@ def dispatch(  # type: ignore[override]
             TOP_K=topk,
             COMPRESS_RATIO=compress_ratio,
             WINDOW_SIZE=WINDOW_SIZE,
-            IMAGE_WIDTH=image_width,
             PADDED_TOP_K=padded_topk,
             input_variant=input_variant,
         )
@@ -789,19 +1047,10 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
             return []
 
         window_size = _hf_config_int(vllm_config, "sliding_window", 128)
-        image_width = (
-            _hf_config_int(vllm_config, "vision_max_n_token", 0)
-            if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0
-            else 0
-        )
-        # Warm both the plain-window variant (batches without image spans)
-        # and the in-image bidirectional variant.
-        image_widths = [0, image_width] if image_width > 0 else [0]
         return self._trace_dispatch(self.dispatch)(
             _combine_topk_swa_warmup_inputs(vllm_config),
             _COMBINE_TOPK_SWA_POINTER_INPUTS,
             WINDOW_SIZE=window_size,
-            image_width=image_widths,
         )
 
     def compile(self, compile_key: CompileKey) -> None:
@@ -818,14 +1067,11 @@ def compile(self, compile_key: CompileKey) -> None:
             input_variant.pointer("query_start_loc", torch.int32),
             input_variant.pointer("seq_lens", torch.int32),
             input_variant.pointer("gather_lens", torch.int32),
-            int32_ptr,
-            int32_ptr,
             1,  # do not specialize M
             1,  # do not specialize N
             TOP_K=compile_key.TOP_K,
             COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
             WINDOW_SIZE=compile_key.WINDOW_SIZE,
-            IMAGE_WIDTH=compile_key.IMAGE_WIDTH,
             PADDED_TOP_K=compile_key.PADDED_TOP_K,
             grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS),
         )
@@ -844,13 +1090,8 @@ def __call__(
         TOP_K: int,
         COMPRESS_RATIO: int,
         WINDOW_SIZE: int,
-        left_visible: torch.Tensor | None = None,
-        right_visible: torch.Tensor | None = None,
-        max_image_tokens: int = 0,
     ) -> None:
         num_reqs = seq_lens.shape[0]
-        has_image = left_visible is not None
-        image_width = max_image_tokens if has_image else 0
         self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)](
             combined_indices,
             combined_indices.stride(0),
@@ -860,14 +1101,11 @@ def __call__(
             query_start_loc,
             seq_lens,
             gather_lens,
-            left_visible if has_image else topk_indices,
-            right_visible if has_image else topk_indices,
             M,
             N,
             TOP_K=TOP_K,
             COMPRESS_RATIO=COMPRESS_RATIO,
             WINDOW_SIZE=WINDOW_SIZE,
-            IMAGE_WIDTH=image_width,
             PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]),
         )
 
@@ -894,9 +1132,8 @@ def build_flashinfer_mixed_sparse_indices(
     decode_is_valid_token: torch.Tensor | None = None,
     swa_block_span: int | None = None,
     compressed_block_span: int | None = None,
-    prefill_left_visible: torch.Tensor | None = None,
-    prefill_right_visible: torch.Tensor | None = None,
-    max_image_tokens: int = 0,
+    *,
+    replay_start: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
 
@@ -907,20 +1144,13 @@ def build_flashinfer_mixed_sparse_indices(
     tokens derive their SWA window from the position and translate local
     compressed indices to global slots via the block tables.
 
-    When ``prefill_left_visible``/``prefill_right_visible`` are given (vision
-    variant), the SWA column region widens by ``max_image_tokens`` and prefill
-    tokens inside an image span get a bidirectionally widened window; decode
-    rows are padded with -1 across the extra columns.
+    ``replay_start`` ([num_reqs], SWA bounded replay) lower-bounds every
+    prefill token's window: positions below it hold no window KV.
     """
     assert decode_swa_indices.dtype == torch.int32
     assert decode_swa_indices.dim() == 2
-    swa_index_width = decode_swa_indices.shape[-1]
-    assert swa_index_width >= window_size
-    has_image = prefill_left_visible is not None
-    image_width = max_image_tokens if has_image else 0
-    swa_total_width = swa_index_width + image_width
-    if has_image:
-        assert prefill_right_visible is not None
+    swa_total_width = decode_swa_indices.shape[-1]
+    assert swa_total_width >= window_size
     if decode_compressed_topk_lens is not None:
         assert decode_compressed_topk_lens.dtype == torch.int32
     assert prefill_topk_indices.dtype == torch.int32
@@ -979,7 +1209,7 @@ def build_flashinfer_mixed_sparse_indices(
     if num_tokens == 0:
         return sparse_indices, sparse_topk_lens
 
-    window_block_size = triton.next_power_of_2(max(swa_index_width, 1))
+    window_block_size = triton.next_power_of_2(max(swa_total_width, 1))
     topk_block_size = triton.next_power_of_2(max(padded_topk, 1))
     max_block_size = max(window_block_size, topk_block_size)
     num_warps = 4 if max_block_size >= 256 else 1
@@ -1003,11 +1233,10 @@ def build_flashinfer_mixed_sparse_indices(
         decode_is_valid_token,
         prefill_topk_indices,
         prefill_topk_indices.stride(0),
-        prefill_left_visible if has_image else token_to_req_indices,
-        prefill_right_visible if has_image else token_to_req_indices,
         query_start_loc,
         seq_lens,
         token_to_req_indices,
+        replay_start,
         swa_block_table,
         swa_block_table.stride(0),
         swa_block_size,
@@ -1018,8 +1247,7 @@ def build_flashinfer_mixed_sparse_indices(
         compressed_span,
         NUM_DECODE_TOKENS=num_decode_tokens,
         WINDOW_SIZE=window_size,
-        SWA_INDEX_WIDTH=swa_index_width,
-        IMAGE_WIDTH=image_width,
+        SWA_INDEX_WIDTH=swa_total_width,
         COMPRESS_RATIO=compress_ratio,
         TOP_K=topk,
         PADDED_TOP_K=padded_topk,
@@ -1038,7 +1266,7 @@ def build_flashinfer_mixed_sparse_indices(
 def _remap_flashinfer_index(values, block_size, block_span):
     # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
     # packed pages (#44577) need block*block_size+off -> block*block_span+off.
-    # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
+    # The launcher declares these indices with sparse_indices_are_storage_offsets=True.
     is_valid = values >= 0
     safe_values = tl.where(is_valid, values, 0)
     values = (safe_values // block_size) * block_span
@@ -1074,11 +1302,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
     decode_is_valid_token_ptr,
     prefill_topk_indices_ptr,
     prefill_topk_stride,
-    left_visible_ptr,
-    right_visible_ptr,
     query_start_loc_ptr,
     seq_lens_ptr,
     token_to_req_indices_ptr,
+    replay_start_ptr,
     swa_block_table_ptr,
     swa_block_table_stride,
     swa_block_size,
@@ -1090,7 +1317,6 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
     NUM_DECODE_TOKENS,
     WINDOW_SIZE: tl.constexpr,
     SWA_INDEX_WIDTH: tl.constexpr,
-    IMAGE_WIDTH: tl.constexpr,
     COMPRESS_RATIO: tl.constexpr,
     TOP_K: tl.constexpr,
     PADDED_TOP_K: tl.constexpr,
@@ -1102,18 +1328,14 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
     TOPK_BLOCK_SIZE: tl.constexpr,
 ):
     token_idx = tl.program_id(0)
-    # Total SWA column count; > SWA_INDEX_WIDTH only for in-image
-    # bidirectional visibility (vision variant), where the extra columns are
-    # -1-padded for decode rows.
-    SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH
 
     if token_idx < NUM_DECODE_TOKENS:
-        for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
+        for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):
             offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
-            mask = offset < SWA_TOTAL_WIDTH
+            mask = offset < SWA_INDEX_WIDTH
             values = tl.load(
                 decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
-                mask=offset < SWA_INDEX_WIDTH,
+                mask=mask,
                 other=-1,
             )
             values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
@@ -1156,7 +1378,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
             tl.store(
                 sparse_indices_ptr
                 + token_idx * sparse_indices_stride
-                + SWA_TOTAL_WIDTH
+                + SWA_INDEX_WIDTH
                 + offset,
                 values,
                 mask=mask,
@@ -1170,7 +1392,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
             else:
                 compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
 
-        tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len)
+        tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len)
         return
 
     prefill_idx = token_idx - NUM_DECODE_TOKENS
@@ -1182,26 +1404,19 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
     start_pos = seq_len - query_len
     token_idx_in_query = token_idx - query_start
     pos = start_pos + token_idx_in_query
-    if IMAGE_WIDTH > 0:
-        # In-image bidirectional visibility: window starts up to
-        # max(left - (window - 1), 0) earlier and extends `right` past pos.
-        left = tl.load(left_visible_ptr + token_idx)
-        right = tl.load(right_visible_ptr + token_idx)
-    else:
-        left = 0
-        right = 0
-    left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
-    swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
-    swa_len = pos + right - swa_start_pos + 1
+    swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1), 0)
+    # SWA bounded replay: no window KV exists below the request's replay start.
+    swa_start_pos = tl.maximum(swa_start_pos, tl.load(replay_start_ptr + req_idx))
+    swa_len = pos - swa_start_pos + 1
     if COMPRESS_RATIO > 0:
         topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
     else:
         # SWA-only layers (ratio 0): no indexer topk; div by constexpr 0 is UB.
         topk_len = 0
 
-    for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
+    for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):
         offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
-        mask = offset < SWA_TOTAL_WIDTH
+        mask = offset < SWA_INDEX_WIDTH
         pos_offset = swa_start_pos + offset
         block_indices = pos_offset // swa_block_size
         block_numbers = tl.load(
@@ -1245,10 +1460,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
         tl.store(
             sparse_indices_ptr
             + token_idx * sparse_indices_stride
-            + SWA_TOTAL_WIDTH
+            + SWA_INDEX_WIDTH
             + offset,
             slot_ids,
             mask=mask,
         )
 
-    tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len)
+    tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len)
diff --git a/vllm/models/deepseek_v4_1/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py
similarity index 74%
rename from vllm/models/deepseek_v4_1/common/ops/fused_compress_quant_cache.py
rename to vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py
index 5950c4b6ba5c..2ba9dbea3de7 100644
--- a/vllm/models/deepseek_v4_1/common/ops/fused_compress_quant_cache.py
+++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py
@@ -4,6 +4,7 @@
 
 import torch
 
+from vllm.models.deepseek_v4.common.ops.fused_indexer_q import _fp32x2_to_fp4x2
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
@@ -57,6 +58,7 @@ def fused_save_compress_norm(
         rms_norm_eps: RMSNorm epsilon.
         compress_ratio: Group size, either 1 or 2.
         latent_out: BF16 [tokens, 512], written only at valid group boundaries.
+
     """
     assert compress_ratio in (1, 2)
     assert kv_score.dtype == torch.float32
@@ -232,8 +234,12 @@ def rope_quant_insert(
 
     The BF16 latent supplies both NoPE quantization and RoPE input. It is read
     only for valid slots at group boundaries. The cache dtype selects the
-    layout: ``uint8`` is the fp8_ds_mla paged layout (576 value bytes and eight
-    segregated UE8M0 scale bytes per token, including one zero padding scale);
+    layout: ``uint8`` is a paged FlashMLA layout, whose record the per-token
+    byte width names -- 584 B for V4 (576 value bytes and eight segregated
+    UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1
+    (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of
+    32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32
+    e4m3 scales of 16 dims each), which only the compressed cache uses.
     ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows
     read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``.
     """
@@ -246,8 +252,15 @@ def rope_quant_insert(
         return
     launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
     if kv_cache.dtype == torch.uint8:
-        assert kv_cache.shape[-1] == 584
-        _rope_quant_insert_kernel[(num_tokens,)](
+        kernel = {
+            584: _rope_quant_insert_kernel,
+            528: _rope_quant_insert_mxfp8_kernel,
+            288: _rope_quant_insert_nvfp4_kernel,
+        }.get(kv_cache.shape[-1])
+        assert kernel is not None, (
+            f"unsupported paged KV record width {kv_cache.shape[-1]}"
+        )
+        kernel[(num_tokens,)](
             latent,
             positions,
             cos_sin_cache,
@@ -338,6 +351,116 @@ def _rope_quant_insert_kernel(
     tl.store(rope_dst + d - 448, rotated.to(tl.bfloat16), d >= 448)
 
 
+@triton.jit
+def _rope_quant_insert_mxfp8_kernel(
+    latent,
+    positions,
+    cos_sin,
+    cache,
+    cache_slots,
+    COS_STRIDE: tl.constexpr,
+    CACHE_STRIDE: tl.constexpr,
+    CACHE_BLOCK: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    SANITIZE_CACHE_NANS: tl.constexpr,
+):
+    """V4.1 record: RoPE first, then MXFP8-quantize all 512 dims.
+
+    The RoPE dims are quantized here too, so unlike the V4 kernel the rotation
+    has to happen before the scales are picked.
+    """
+    t = tl.program_id(0)
+    slot = tl.load(cache_slots + t)
+    if slot < 0:
+        return
+    position = tl.load(positions + t)
+    if (position + 1) % COMPRESS_RATIO != 0:
+        return
+    d = tl.arange(0, 512)
+    normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
+
+    # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
+    even, odd = tl.split(tl.reshape(normed, (256, 2)))
+    pair = tl.arange(0, 256) - 224
+    cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
+    c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
+    s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
+    rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
+    if SANITIZE_CACHE_NANS:
+        rotated = tl.where(rotated == rotated, rotated, 0.0)
+
+    page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
+    values = page + (slot % CACHE_BLOCK) * 512
+    scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16
+
+    quant = tl.reshape(rotated, (16, 32))
+    amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4)
+    exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0)))
+    scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1))
+    fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv)
+    tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,)))
+
+    max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
+    encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded)
+    tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8))
+
+
+@triton.jit
+def _rope_quant_insert_nvfp4_kernel(
+    latent,
+    positions,
+    cos_sin,
+    cache,
+    cache_slots,
+    COS_STRIDE: tl.constexpr,
+    CACHE_STRIDE: tl.constexpr,
+    CACHE_BLOCK: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    SANITIZE_CACHE_NANS: tl.constexpr,
+):
+    """V4.1 NVFP4 record: RoPE, then e2m1 with one e4m3 scale per 16 dims.
+
+    The scale is ``amax / 6`` (6 is e2m1's largest magnitude) clamped to the
+    e4m3 range, with no per-tensor scale on top.
+    """
+    t = tl.program_id(0)
+    slot = tl.load(cache_slots + t)
+    if slot < 0:
+        return
+    position = tl.load(positions + t)
+    if (position + 1) % COMPRESS_RATIO != 0:
+        return
+    d = tl.arange(0, 512)
+    normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
+
+    # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
+    even, odd = tl.split(tl.reshape(normed, (256, 2)))
+    pair = tl.arange(0, 256) - 224
+    cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
+    c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
+    s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
+    rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
+    if SANITIZE_CACHE_NANS:
+        rotated = tl.where(rotated == rotated, rotated, 0.0)
+
+    tiles = tl.reshape(rotated, (32, 16))
+    amax = tl.max(tl.abs(tiles), 1)
+    # 2**-9 is the smallest normal e4m3 magnitude; 448 the largest.
+    scale = tl.clamp(amax * (1.0 / 6.0), 0.001953125, 448.0).to(tl.float8e4nv)
+    # Round-to-nearest division: Triton's default div.full misplaces values
+    # that land exactly on an e2m1 tie.
+    scaled = tl.math.div_rn(tiles, tl.reshape(scale.to(tl.float32), (32, 1)))
+    lo, hi = tl.split(tl.reshape(tl.reshape(scaled, (512,)), (256, 2)))
+
+    page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
+    tl.store(
+        page + (slot % CACHE_BLOCK) * 256 + tl.arange(0, 256),
+        _fp32x2_to_fp4x2(lo, hi),
+    )
+    scales = page + CACHE_BLOCK * 256 + (slot % CACHE_BLOCK) * 32
+    tl.store(scales + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True))
+
+
 @triton.jit
 def _rope_plain_insert_kernel(
     latent,
diff --git a/vllm/models/deepseek_v41/common/ops/fused_layout.py b/vllm/models/deepseek_v41/common/ops/fused_layout.py
new file mode 100644
index 000000000000..071afde6bca4
--- /dev/null
+++ b/vllm/models/deepseek_v41/common/ops/fused_layout.py
@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Weight permutations for FlashMLA's mega-attention kernel.
+
+The kernel reads Q with 16-element head-dim chunks interleaved across heads and
+writes O with 32-element chunks interleaved across the 8 heads of a ``wo_a``
+group. Every permutation here satisfies ``fused = standard[perm]``, and is
+applied once to ``wq_b`` rows and ``wo_a`` columns at load time so the
+surrounding GEMMs produce and consume the kernel's layouts directly -- no
+per-step shuffle.
+"""
+
+import torch
+
+HEAD_DIM = 512
+Q_CHUNK = 16
+O_CHUNK = 32
+WV_GROUP_SIZE = 8
+
+
+def q_fused_permutation(num_heads: int, head_dim: int = HEAD_DIM) -> torch.Tensor:
+    """``fused[(d // 16) * (H * 16) + h * 16 + d % 16] = standard[h * D + d]``."""
+    h = torch.arange(num_heads).view(num_heads, 1)
+    d = torch.arange(head_dim).view(1, head_dim)
+    fused_index = (d // Q_CHUNK) * (num_heads * Q_CHUNK) + h * Q_CHUNK + d % Q_CHUNK
+    perm = torch.empty(num_heads * head_dim, dtype=torch.long)
+    perm[fused_index.reshape(-1)] = torch.arange(num_heads * head_dim)
+    return perm
+
+
+def o_fused_permutation(
+    heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM
+) -> torch.Tensor:
+    """``fused[(c * G + h) * 32 + j] = standard[h * D + c * 32 + j]`` per group."""
+    h = torch.arange(heads_per_group).view(-1, 1, 1)
+    c = torch.arange(head_dim // O_CHUNK).view(1, -1, 1)
+    j = torch.arange(O_CHUNK).view(1, 1, -1)
+    fused_index = (c * heads_per_group + h) * O_CHUNK + j
+    perm = torch.empty(heads_per_group * head_dim, dtype=torch.long)
+    perm[fused_index.reshape(-1)] = torch.arange(heads_per_group * head_dim)
+    return perm
+
+
+def o_fused_chunk_permutation(
+    heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM
+) -> torch.Tensor:
+    """Per-32-element-chunk form of :func:`o_fused_permutation`, for scales."""
+    return o_fused_permutation(heads_per_group, head_dim)[::O_CHUNK] // O_CHUNK
+
+
+def permute_q_to_fused(q: torch.Tensor) -> torch.Tensor:
+    """``[N, H, D]`` standard layout -> the same shape in the fused layout."""
+    n, h, d = q.shape
+    perm = q_fused_permutation(h, d).to(q.device)
+    return q.reshape(n, h * d)[:, perm].view(n, h, d)
+
+
+def _bytes_view(t: torch.Tensor) -> torch.Tensor:
+    """Byte view of a 1-byte-element tensor, so fp8/ue8m0 can be gathered."""
+    return t.view(torch.uint8) if t.element_size() == 1 else t
+
+
+def permute_wq_b_(
+    weight: torch.Tensor, weight_scale: torch.Tensor, num_local_heads: int
+) -> None:
+    """Permute an MXFP8 ``wq_b`` shard's rows and per-row scales, in place."""
+    head_dim = weight.shape[0] // num_local_heads
+    perm = q_fused_permutation(num_local_heads, head_dim).to(weight.device)
+    for t in (weight, weight_scale):
+        b = _bytes_view(t)
+        b.copy_(b[perm])
+
+
+def permute_wo_a_(
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    heads_per_group: int = WV_GROUP_SIZE,
+) -> None:
+    """Permute an MXFP8 ``wo_a`` shard's input columns and scales, in place."""
+    head_dim = weight.shape[1] // heads_per_group
+    perm = o_fused_permutation(heads_per_group, head_dim).to(weight.device)
+    w = _bytes_view(weight)
+    w.copy_(w[:, perm])
+    chunk_perm = o_fused_chunk_permutation(heads_per_group, head_dim)
+    s = _bytes_view(weight_scale)
+    s.copy_(s[:, chunk_perm.to(weight.device)])
diff --git a/vllm/models/deepseek_v4_1/common/ops/indexer_k_store.py b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py
similarity index 99%
rename from vllm/models/deepseek_v4_1/common/ops/indexer_k_store.py
rename to vllm/models/deepseek_v41/common/ops/indexer_k_store.py
index c4bd98b9e99e..ba8ddbd95a9a 100644
--- a/vllm/models/deepseek_v4_1/common/ops/indexer_k_store.py
+++ b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py
@@ -46,11 +46,13 @@ def indexer_k_norm_rope_store(
         cos_sin_cache: [max_pos, rope_head_dim] GPT-J layout (cos half, then
             sin half), from the layer's compress-RoPE instance.
         rms_norm_weight: [128] k_norm weight.
+        rms_norm_eps: Epsilon of the k_norm RMSNorm.
         k_cache: uint8 paged indexer cache [num_blocks, block_size, row_bytes].
         kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip).
         compress_ratio: group size; keys are emitted at group boundaries.
         use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else
             per-token FP8 with a single fp32 scale.
+
     """
     num_tokens = kv_slot_mapping.numel()
     assert k_pre.ndim == 2 and k_pre.shape[1] == 128
diff --git a/vllm/models/deepseek_v4_1/common/ops/query_quant.py b/vllm/models/deepseek_v41/common/ops/query_quant.py
similarity index 100%
rename from vllm/models/deepseek_v4_1/common/ops/query_quant.py
rename to vllm/models/deepseek_v41/common/ops/query_quant.py
diff --git a/vllm/models/deepseek_v4_1/common/rope.py b/vllm/models/deepseek_v41/common/rope.py
similarity index 93%
rename from vllm/models/deepseek_v4_1/common/rope.py
rename to vllm/models/deepseek_v41/common/rope.py
index 51ff87926992..75c2ae1a063e 100644
--- a/vllm/models/deepseek_v4_1/common/rope.py
+++ b/vllm/models/deepseek_v41/common/rope.py
@@ -32,9 +32,9 @@ def build_deepseek_v4_rope(
         # YaRN applies to all compressed-KV layers (compress_ratio > 0); only
         # pure sliding-window layers (compress_ratio == 0) use plain RoPE.
         rope_parameters["rope_type"] = (
-            "deepseek_yarn"
-            if rope_parameters.get("apply_yarn_scaling", True)
-            else "deepseek_llama_scaling"
+            "deepseek_llama_scaling"
+            if rope_parameters.get("attention_factor") == 1.0
+            else "deepseek_yarn"
         )
     else:
         # Sliding-window layers use plain RoPE (theta=rope_theta, no YaRN).
diff --git a/vllm/models/deepseek_v41/common/vl_cudagraph.py b/vllm/models/deepseek_v41/common/vl_cudagraph.py
new file mode 100644
index 000000000000..e2a718b7453c
--- /dev/null
+++ b/vllm/models/deepseek_v41/common/vl_cudagraph.py
@@ -0,0 +1,390 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Encoder CUDA graph support for the DeepSeek-V4.1 vision tower.
+
+Mixin implementing the ``SupportsEncoderCudaGraph`` protocol
+(``vllm/model_executor/models/interfaces.py``) on top of the shared
+``DeepseekV4ViT``/``DeepseekV4Aligner``. The captured graph packs a whole
+batch of images into one varlen run: ViT blocks attend per image via
+``cu_seqlens``, and the aligner's spatial merge becomes a gather+mask
+against precomputed indices. Span assembly (IMAGE_START/IMAGE_NEW_LINE/
+IMAGE_END delimiters) stays eager in ``postprocess_encoder_output`` but is
+batched over all items of the packed batch.
+
+Enabled via ``--compilation-config '{"cudagraph_mm_encoder": true}'``; the
+manager (``vllm/v1/worker/encoder_cudagraph.py``) handles budget packing, DP
+sharding and eager fallback.
+"""
+
+from collections.abc import Hashable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.models.deepseek_v4.common.vision import (
+    build_packed_merge_metadata,
+    build_packed_vit_metadata,
+)
+from vllm.v1.worker.encoder_cudagraph_defs import (
+    EncoderCudaGraphCaptureInputs,
+    EncoderCudaGraphConfig,
+    EncoderCudaGraphReplayBuffers,
+    EncoderItemSpec,
+)
+
+from .mm_preprocess import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.config.multimodal import MultiModalConfig
+    from vllm.models.deepseek_v4.common.vision import (
+        DeepseekV4Aligner,
+        DeepseekV4ViT,
+    )
+    from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
+
+
+class DeepseekV4VLEncoderCudaGraphMixin:
+    """``SupportsEncoderCudaGraph`` for ``DeepseekV41ForCausalLM``.
+
+    Expects ``self.vision``/``self.aligner``/``self.config``/
+    ``self.multimodal_config`` and the ``image_start``/``image_newline``/
+    ``image_end`` parameters from the host model.
+    """
+
+    # Provided by the host model; declared here so mixin methods type-check.
+    vision: "DeepseekV4ViT"
+    aligner: "DeepseekV4Aligner"
+    config: "DeepseekV41Config"
+    multimodal_config: "MultiModalConfig | None"
+    image_start: torch.nn.Parameter
+    image_end: torch.nn.Parameter
+    image_newline: torch.nn.Parameter
+
+    # -- helpers shared with the eager multimodal path --
+
+    def _encode_image(
+        self,
+        patches: torch.Tensor,
+        n_vit_h: int,
+        n_vit_w: int,
+    ) -> torch.Tensor:
+        # Aligner rows in reading order, one per IMAGE slot.
+        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
+
+    def _build_image_span(
+        self, image_embeds: torch.Tensor, types: torch.Tensor
+    ) -> torch.Tensor:
+        """Full image span: aligner rows at IMAGE slots, the learned
+        delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""
+        types = types.to(image_embeds.device)
+        span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])
+        dtype = image_embeds.dtype
+        span[types == IMAGE_START] = self.image_start.to(dtype)
+        span[types == IMAGE_END] = self.image_end.to(dtype)
+        span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
+        span[types == IMAGE] = image_embeds
+        return span
+
+    @staticmethod
+    def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]:
+        grid = mm_kwargs[key]
+        if isinstance(grid, torch.Tensor):
+            grid = grid.tolist()
+        return [[int(v) for v in row] for row in grid]
+
+    _encoder_cg_pad_totals: dict[int, int] = {}
+
+    # -- SupportsEncoderCudaGraph protocol --
+
+    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
+        mm_cfg = self.multimodal_config
+        if (
+            mm_cfg is not None
+            and mm_cfg.mm_encoder_attn_dtype == "fp8"
+            and mm_cfg.mm_encoder_fp8_scale_path is None
+        ):
+            raise ValueError(
+                "cudagraph_mm_encoder is incompatible with dynamic FP8 "
+                "scaling for the vision encoder (the amax history is "
+                "host-side state that cannot be captured). Provide a static "
+                "scale file via --mm-encoder-fp8-scale-path or disable "
+                "cudagraph_mm_encoder."
+            )
+
+        from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+        backend = self.vision.blocks[0].attn.attn.attn_backend
+        # FlashInfer reads max_seqlen on the host and TORCH_SDPA calls
+        # .tolist() on CUDA cu_seqlens; neither can be captured.
+        if backend in (
+            AttentionBackendEnum.FLASHINFER,
+            AttentionBackendEnum.TORCH_SDPA,
+        ):
+            raise ValueError(
+                f"cudagraph_mm_encoder is not supported with the {backend.name} "
+                "ViT attention backend for DeepSeek-V4.1. Set "
+                "--mm-encoder-attn-backend FLASH_ATTN or disable "
+                "cudagraph_mm_encoder."
+            )
+
+        pad_totals = self._encoder_cg_pad_totals
+
+        def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None:
+            # Varlen attention requires cu_seqlens[-1] to equal the number of
+            # rows actually passed in. The captured buffers are sized for the
+            # full token budget, so a smaller real batch is completed with one
+            # trailing padding sequence; declaring fewer rows than the buffer
+            # holds is undefined behaviour and returns NaN on FlashAttn.
+            total = pad_totals.get(dst.data_ptr())
+            if total is None:
+                raise RuntimeError(
+                    "cu_seqlens replay buffer was not registered at capture "
+                    "time; the manager must replay into the capture-time "
+                    "buffers."
+                )
+            n = min(src.shape[0], dst.shape[0])
+            dst[:n].copy_(src[:n])
+            dst[n:] = total
+
+        return EncoderCudaGraphConfig(
+            modalities=["image"],
+            buffer_keys=[
+                "patches",
+                "vit_cos",
+                "vit_sin",
+                "cu_seqlens",
+                "max_seqlen",
+                "merge_idx",
+                "merge_mask",
+            ],
+            out_hidden_size=self.config.hidden_size,
+            padding_logics={"cu_seqlens": pad_cu_seqlens},
+        )
+
+    def get_encoder_cudagraph_budget_range(
+        self,
+        vllm_config: "VllmConfig",
+    ) -> tuple[int, int]:
+        # Min: comfortably above the smallest possible image span (4 tokens
+        # for a 1x1 grid); also sets max_batch_size = 64 via the manager's
+        # auto-inference (max_batch_size <= min_budget invariant).
+        min_budget = 64
+        max_budget = min(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            vllm_config.model_config.max_model_len,
+        )
+        return (min_budget, max_budget)
+
+    def get_encoder_cudagraph_item_specs(
+        self,
+        mm_kwargs: dict[str, Any],
+    ) -> list[EncoderItemSpec]:
+        vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")
+        llm_grid = self._get_grid_list(mm_kwargs, "llm_grid")
+        return [
+            EncoderItemSpec(
+                input_size=h * w,
+                output_tokens=lh * (lw + 1) + 2,
+            )
+            for (h, w), (lh, lw) in zip(vit_grid, llm_grid, strict=True)
+        ]
+
+    def select_encoder_cudagraph_items(
+        self,
+        mm_kwargs: dict[str, Any],
+        indices: list[int],
+    ) -> dict[str, Any]:
+        vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")
+        llm_grid = self._get_grid_list(mm_kwargs, "llm_grid")
+        patches = mm_kwargs["patches"]
+        types = mm_kwargs["types"]
+
+        if len(indices) == 0:
+            return {
+                "patches": patches[:0],
+                "vit_grid": torch.zeros((0, 2), dtype=torch.int64),
+                "llm_grid": torch.zeros((0, 2), dtype=torch.int64),
+                "types": types[:0],
+            }
+
+        cum_patches = [0]
+        for h, w in vit_grid:
+            cum_patches.append(cum_patches[-1] + h * w)
+        cum_spans = [0]
+        for lh, lw in llm_grid:
+            cum_spans.append(cum_spans[-1] + lh * (lw + 1) + 2)
+
+        return {
+            "patches": torch.cat(
+                [patches[cum_patches[i] : cum_patches[i + 1]] for i in indices]
+            ),
+            "vit_grid": torch.tensor([vit_grid[i] for i in indices], dtype=torch.int64),
+            "llm_grid": torch.tensor([llm_grid[i] for i in indices], dtype=torch.int64),
+            "types": torch.cat(
+                [types[cum_spans[i] : cum_spans[i + 1]] for i in indices]
+            ),
+        }
+
+    def prepare_encoder_cudagraph_capture_inputs(
+        self,
+        token_budget: int,
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        path: str = "default",
+        axis_keys: tuple[Hashable, ...] | None = None,
+    ) -> EncoderCudaGraphCaptureInputs:
+        if (
+            self.multimodal_config is not None
+            and self.multimodal_config.get_limit_per_prompt("image") == 0
+        ):
+            raise RuntimeError(
+                "cudagraph_mm_encoder requires the vision tower, but it was "
+                "stubbed out because the image limit is 0. Disable "
+                "cudagraph_mm_encoder."
+            )
+
+        config = self.config
+        r = config.vision_downsample_ratio
+        # Span rows per dummy item (ceiling so the buffers tile the budget).
+        # Each dummy image is a single aligner row of per_item_out blocks:
+        # r x (r * per_item_out) patches. One real image using the whole
+        # budget needs at most r^2 * token_budget patches, which this fits.
+        per_item_out = -(-token_budget // max_batch_size)
+        grids = [[r, r * per_item_out]] * max_batch_size
+        total_patches = r * r * per_item_out * max_batch_size
+
+        p = config.vision_patch_size
+        patches = torch.zeros(total_patches, 3, p, p, device=device, dtype=dtype)
+
+        metadata = build_packed_vit_metadata(
+            grids,
+            rope_dim=self.vision.rope_dim,
+            rope_theta=self.vision.rope_theta,
+            device=device,
+            max_seqlen_override=total_patches,
+            cached=False,
+        )
+        # Spare cu_seqlens slots let replay append a padding sequence that
+        # covers the rows a smaller real batch does not fill.
+        real_cu = metadata.pop("cu_seqlens")
+        cu_seqlens = torch.full(
+            (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device
+        )
+        cu_seqlens[: real_cu.numel()] = real_cu
+
+        merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype)
+
+        self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches
+
+        return EncoderCudaGraphCaptureInputs(
+            values={
+                "patches": patches,
+                **metadata,
+                "cu_seqlens": cu_seqlens,
+                **merge,
+            }
+        )
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self,
+        mm_kwargs: dict[str, Any],
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        path: str = "default",
+    ) -> EncoderCudaGraphReplayBuffers:
+        vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")
+        patches = mm_kwargs["patches"]
+        dtype = self.aligner.w1.weight.dtype
+        if patches.dtype != dtype:
+            patches = patches.to(dtype)
+
+        # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and
+        # pad_cu_seqlens appends the padding sequence covering the tail rows.
+        metadata = build_packed_vit_metadata(
+            vit_grid,
+            rope_dim=self.vision.rope_dim,
+            rope_theta=self.vision.rope_theta,
+            device=patches.device,
+        )
+        merge = build_packed_merge_metadata(
+            vit_grid,
+            self.config.vision_downsample_ratio,
+            device=patches.device,
+            dtype=dtype,
+        )
+        return EncoderCudaGraphReplayBuffers(
+            values={"patches": patches, **metadata, **merge}
+        )
+
+    def encoder_cudagraph_forward(
+        self,
+        inputs: dict[str, torch.Tensor],
+        path: str = "default",
+    ) -> torch.Tensor:
+        vit_out = self.vision.forward_packed(
+            inputs["patches"],
+            inputs["vit_cos"],
+            inputs["vit_sin"],
+            inputs["cu_seqlens"],
+            inputs["max_seqlen"],
+        )
+        return self.aligner.forward_packed(
+            vit_out, inputs["merge_idx"], inputs["merge_mask"]
+        )
+
+    def encoder_eager_forward(
+        self,
+        mm_kwargs: dict[str, Any],
+        path: str = "default",
+    ) -> torch.Tensor:
+        patches = mm_kwargs["patches"].to(self.aligner.w1.weight.dtype)
+        vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")
+        outs: list[torch.Tensor] = []
+        offset = 0
+        for n_vit_h, n_vit_w in vit_grid:
+            n_vit = n_vit_h * n_vit_w
+            outs.append(
+                self._encode_image(patches[offset : offset + n_vit], n_vit_h, n_vit_w)
+            )
+            offset += n_vit
+        if not outs:
+            return patches.new_zeros((0, self.config.hidden_size))
+        return torch.cat(outs, dim=0)
+
+    def postprocess_encoder_output(
+        self,
+        outputs: dict[str, torch.Tensor],
+        indices: list[int],
+        per_item_out_tokens: list[int],
+        dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
+        clone: bool = False,
+        batch_mm_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        if not indices:
+            return
+        assert batch_mm_kwargs is not None
+        aligner_out = outputs["default"]
+        r = self.config.vision_downsample_ratio
+        vit_grid = self._get_grid_list(batch_mm_kwargs, "vit_grid")
+        llm_grid = self._get_grid_list(batch_mm_kwargs, "llm_grid")
+        types = batch_mm_kwargs["types"].to(aligner_out.device)
+
+        n_rows = sum(-(-h // r) * (-(-w // r)) for h, w in vit_grid)
+        span_lens = [lh * (lw + 1) + 2 for lh, lw in llm_grid]
+
+        # Batched span assembly: one masked fill per role across all items.
+        dtype = aligner_out.dtype
+        span = aligner_out.new_empty(sum(span_lens), self.config.hidden_size)
+        span[types == IMAGE] = aligner_out[:n_rows]
+        span[types == IMAGE_START] = self.image_start.to(dtype)
+        span[types == IMAGE_END] = self.image_end.to(dtype)
+        span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
+
+        # Freshly allocated, so later replays cannot clobber the results.
+        offset = 0
+        for idx, span_len in zip(indices, span_lens):
+            dest[idx] = span[offset : offset + span_len]
+            offset += span_len
diff --git a/vllm/models/deepseek_v4_1/compressor.py b/vllm/models/deepseek_v41/compressor.py
similarity index 98%
rename from vllm/models/deepseek_v4_1/compressor.py
rename to vllm/models/deepseek_v41/compressor.py
index d09d97d04c27..42335ae84b42 100644
--- a/vllm/models/deepseek_v4_1/compressor.py
+++ b/vllm/models/deepseek_v41/compressor.py
@@ -12,7 +12,7 @@
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import MergedColumnParallelLinear
-from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import (
+from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
     fused_save_compress_norm,
     rope_quant_insert,
 )
@@ -37,7 +37,7 @@ def get_name() -> str:
         return "CompressorBackend"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
     @classmethod
diff --git a/vllm/models/deepseek_v41/nvidia/__init__.py b/vllm/models/deepseek_v41/nvidia/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/deepseek_v41/nvidia/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/models/deepseek_v4_1/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py
similarity index 89%
rename from vllm/models/deepseek_v4_1/nvidia/dspark.py
rename to vllm/models/deepseek_v41/nvidia/dspark.py
index 46fe2558ab7f..a477d1b5ecd8 100644
--- a/vllm/models/deepseek_v4_1/nvidia/dspark.py
+++ b/vllm/models/deepseek_v41/nvidia/dspark.py
@@ -18,6 +18,7 @@
 
 import vllm.envs as envs
 from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config.kernel import MEGA_MOE_BACKENDS
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -50,12 +51,15 @@
     sp_shard,
 )
 
+from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
 from .model import (
     DeepseekV4DecoderLayer,
     DeepseekV4Model,
     _linear_scale_param_name,
     _use_sequence_parallel,
     make_deepseek_v4_expert_params_mapping,
+    maybe_init_gemm_rs,
+    prepare_mega_gate_routing_metadata,
 )
 
 logger = init_logger(__name__)
@@ -77,6 +81,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.rms_norm_eps = config.rms_norm_eps
         self.num_hidden_layers = config.num_hidden_layers
         self.target_layer_ids = tuple(config.dspark_target_layer_ids)
+        self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
         self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
 
         self.num_dspark_layers = (
@@ -109,12 +114,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         )
 
         current_vllm_config = get_current_vllm_config()
+        # The target model already holds the GEMM-RS workspace (same hidden
+        # size, same TP group); this only re-checks and binds the draft layers.
+        run_gemm_rs = maybe_init_gemm_rs(
+            current_vllm_config, self.use_sequence_parallel
+        )
         self.layers = nn.ModuleList(
             [
                 DeepseekV4DecoderLayer(
                     current_vllm_config,
                     prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
                     topk_indices_buffer=self.topk_indices_buffer,
+                    run_gemm_rs=run_gemm_rs,
                 )
                 for i in range(self.num_dspark_layers)
             ]
@@ -202,12 +213,21 @@ def forward(
                 )
             inputs_embeds = sp_shard(inputs_embeds)
             input_ids = sp_shard(input_ids)
+        mega_gate_metadata = None
+        if self.use_mega_moe:
+            mega_gate_metadata = prepare_mega_gate_routing_metadata(
+                input_ids,
+                has_hash_routing=False,
+                image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
+                if getattr(self.config, "vision_n_layers", 0) > 0
+                else None,
+            )
         # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
         hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
 
         residual = post_mix = res_mix = pre_mix = None
         for layer in self.layers:
-            hidden_states, residual, post_mix, res_mix, pre_mix = layer(
+            hidden_states, residual, post_mix, res_mix, pre_mix, _ = layer(
                 hidden_states,
                 positions,
                 input_ids,
@@ -215,17 +235,17 @@ def forward(
                 post_mix,
                 res_mix,
                 residual,
+                mega_gate_metadata=mega_gate_metadata,
             )
         hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
-        if self.use_sequence_parallel:
-            hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
-            pre_mix = sp_all_gather(pre_mix)[:full_num_tokens]
         # Collapse the hc copies with the pre-mix from the last layer's FFN
         # mixes — the mix the reference forward_head applies (hc_pre with the
         # last block's ffn pre-mix). Return the PRE-norm head hidden;
         # compute_logits applies self.norm.
         assert pre_mix is not None
         hidden_states = hc_collapse_triton(hidden_states, pre_mix)
+        if self.use_sequence_parallel:
+            hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
         return hidden_states
 
 
@@ -235,65 +255,24 @@ def _insert_context_kv(
     positions: torch.Tensor,
     slot_mapping: torch.Tensor,
 ) -> None:
-    """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
-
-    Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
-    query and discard it, since context tokens have no query). Mirrors
-    ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
-    """
-    swa_cache = attn.swa_cache_layer.kv_cache
-    block_size = attn.swa_cache_layer.block_size
-    cos_sin_cache = attn.rotary_emb.cos_sin_cache
-    cache_dtype = swa_cache.dtype
-    n_ctx = kv.shape[0]
-    dummy_q = torch.zeros(
-        (n_ctx, attn.n_local_heads, attn.head_dim),
-        dtype=kv.dtype,
-        device=kv.device,
+    """Insert normalized context KV without constructing an unused query."""
+    # The bound cache already is [num_blocks, block_size, bytes_per_token], so
+    # the record width needs no per-layer attribute -- one whose name the
+    # attention layer is free to change (kv_bytes_per_token has to become
+    # swa_bytes_per_token once the sliding-window and compressed records can
+    # differ). block_size stays the layer's, so the op's size(1) check still
+    # cross-checks the two rather than restating the tensor.
+    cache = attn.swa_cache_layer.kv_cache
+    torch.ops._C.fused_deepseek_v4_kv_rope_insert(
+        kv,
+        cache,
+        slot_mapping,
+        positions,
+        attn.rotary_emb.cos_sin_cache,
+        attn.swa_cache_layer.block_size,
+        attn._flashinfer_fp8_kv_scale if cache.dtype == torch.float8_e4m3fn else None,
+        attn.kv_mxfp8,
     )
-    if cache_dtype == torch.uint8:
-        # fp8_ds_mla UE8M0 paged layout
-        swa_2d = swa_cache.view(swa_cache.shape[0], -1)
-        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-            dummy_q,
-            kv,
-            swa_2d,
-            slot_mapping,
-            positions,
-            cos_sin_cache,
-            attn.padded_heads,
-            attn.eps,
-            block_size,
-        )
-    elif cache_dtype == torch.bfloat16:
-        swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
-        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
-            dummy_q,
-            kv,
-            swa_3d,
-            slot_mapping,
-            positions,
-            cos_sin_cache,
-            attn.eps,
-            block_size,
-        )
-    else:  # per-tensor fp8 (torch.float8_e4m3fn)
-        # TODO(ben): double-check if this is being dispatched correctly for FI backend
-        swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
-        dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
-        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
-            dummy_q,
-            kv,
-            dummy_q_fp8,
-            swa_3d,
-            slot_mapping,
-            positions,
-            cos_sin_cache,
-            attn._flashinfer_fp8_kv_scale,
-            attn._flashinfer_fp8_q_scale_inv,
-            attn.eps,
-            block_size,
-        )
 
 
 class DSparkDeepseekV4ForCausalLM(nn.Module):
@@ -519,8 +498,22 @@ def _finalize_moe(self) -> None:
         for layer in self.model.layers:
             layer.ffn.finalize_mega_moe_weights()
 
+    def _finalize_attn(self) -> None:
+        """Run the attention backend's post-load weight step on the draft layers.
+
+        They are ordinary DeepseekV4DecoderLayers, so they take the engine-wide
+        attention backend and owe it whatever the target model owes it -- mega
+        attention permutes wq_b / wo_a here and refuses to run without it.
+        Idempotent, like the target's.
+        """
+        for layer in self.model.layers:
+            finalize = getattr(layer.attn, "finalize_loaded_weights", None)
+            if finalize is not None:
+                finalize()
+
     def process_weights_after_loading(self) -> None:
         self._finalize_moe()
+        self._finalize_attn()
 
     def _remap_dspark_name(self, name: str) -> str | None:
         """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
diff --git a/vllm/models/deepseek_v41/nvidia/engram.py b/vllm/models/deepseek_v41/nvidia/engram.py
new file mode 100644
index 000000000000..ddb57d3bd24e
--- /dev/null
+++ b/vllm/models/deepseek_v41/nvidia/engram.py
@@ -0,0 +1,417 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
+
+import mmap
+import os
+import shutil
+import tempfile
+import weakref
+from contextlib import ExitStack
+
+import numpy as np
+import torch
+
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.config import get_current_vllm_config
+from vllm.distributed import (
+    get_dp_group,
+    get_engram_dp_group,
+    get_engram_dp_size,
+    get_tensor_model_parallel_rank,
+    tensor_model_parallel_all_gather,
+)
+from vllm.distributed.device_communicators.shm_broadcast import (
+    SHM_PATH,
+    check_shm_free_space,
+)
+from vllm.distributed.parallel_state import GroupCoordinator
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.models.deepseek_v41.common.engram import (
+    DEAD_ID,
+    EngramLayout,
+    _engram_head_shard_weight_loader,
+    _engram_select_rows,
+)
+from vllm.models.deepseek_v41.common.engram import (
+    Engram as BaseEngram,
+)
+from vllm.models.deepseek_v41.common.engram import (
+    ParallelEngramEmbedding as BaseParallelEngramEmbedding,
+)
+from vllm.utils.platform_utils import is_uva_available
+from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
+
+logger = init_logger(__name__)
+
+
+def engram_head_shard_rank() -> int:
+    """This rank's slot among the hash-head shards of one engram table.
+
+    TP-major, so the shards a DP gather brings in are contiguous heads and
+    the following TP gather completes the head order.
+    """
+    dp_group = get_engram_dp_group()
+    dp_size = dp_group.world_size if dp_group is not None else 1
+    dp_rank = dp_group.rank_in_group if dp_group is not None else 0
+    return get_tensor_model_parallel_rank() * dp_size + dp_rank
+
+
+def engram_gathered_num_tokens() -> int:
+    """Per-replica token slot for the node-local Engram DP group."""
+    dp_metadata = get_forward_context().dp_metadata
+    if dp_metadata is None:
+        raise RuntimeError("a DP-shared engram table needs DP token metadata")
+    group = get_engram_dp_group()
+    assert group is not None
+    # Engram groups are contiguous slices of the full DP group.
+    start = get_dp_group().rank_in_group - group.rank_in_group
+    return int(
+        dp_metadata.num_tokens_across_dp_cpu[start : start + group.world_size].max()
+    )
+
+
+def gather_engram_hashes(
+    hash_ids: torch.Tensor, *, dp_shared_memory: bool = False
+) -> torch.Tensor:
+    """Collect the n-gram ids of every DP replica sharing one table.
+
+    Replicas are padded to a common token slot, so the gathered shape is
+    static under CUDA graph capture (where DP already pads alike).
+    """
+    dp_group = get_engram_dp_group()
+    if dp_group is None or dp_shared_memory:
+        return hash_ids
+    slot = engram_gathered_num_tokens()
+    if hash_ids.shape[0] > slot:
+        raise ValueError("Engram token count exceeds the DP token slot")
+    if hash_ids.shape[0] < slot:
+        pad = hash_ids.new_full(
+            (slot - hash_ids.shape[0], *hash_ids.shape[1:]), DEAD_ID
+        )
+        hash_ids = torch.cat((hash_ids, pad))
+    return dp_group.all_gather(hash_ids, dim=0)
+
+
+def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool:
+    """Whether co-located DP replicas exist and /dev/shm can hold the full tables."""
+    if not os.path.isdir(SHM_PATH):
+        logger.warning_once(
+            "%s is not mounted; storing the offloaded Engram tables per rank "
+            "instead of sharing them across DP replicas.",
+            SHM_PATH,
+        )
+        return False
+    if get_engram_dp_size() == 1:
+        logger.warning_once(
+            "Engram DP replicas are not co-located on one node; "
+            "storing the offloaded tables per rank instead of sharing them."
+        )
+        return False
+    num_bytes = sum(layout.num_embeddings) * (
+        layout.head_dim + layout.head_dim // block_size
+    )
+    if shutil.disk_usage(SHM_PATH).total < num_bytes:
+        logger.warning_once(
+            "Sharing Engram tables across DP replicas needs %.1f GiB of %s "
+            "(--shm-size or --ipc=host); sharding them across replicas instead.",
+            num_bytes / 1024**3,
+            SHM_PATH,
+        )
+        return False
+    return True
+
+
+class DPSharedEngramStorage:
+    """Registered host weights shared by a node-local DP group with one writer."""
+
+    def __init__(
+        self, num_rows: int, dim: int, block_size: int, group: GroupCoordinator
+    ) -> None:
+        self.group = group
+        weight_bytes = num_rows * dim
+        storage = self._allocate(weight_bytes + weight_bytes // block_size)
+        self.weight = storage[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim)
+        self.weight_scale_inv = storage[weight_bytes:].view(-1, dim // block_size)
+        self._views: tuple[torch.Tensor, torch.Tensor] | None = None
+
+    def _allocate(self, num_bytes: int) -> torch.Tensor:
+        """Map and register one physical allocation across a node-local DP group."""
+        group = self.group
+        with ExitStack() as stack:
+            path = error = None
+            if group.rank_in_group == 0:
+                try:
+                    check_shm_free_space(num_bytes)
+                    backing_file = stack.enter_context(
+                        tempfile.NamedTemporaryFile(prefix="vllm_engram_", dir=SHM_PATH)
+                    )
+                    backing_file.truncate(num_bytes)
+                    path = backing_file.name
+                except Exception as exc:
+                    error = f"{type(exc).__name__}: {exc}"
+            path, error = group.broadcast_object((path, error))
+            if error is not None:
+                raise RuntimeError(
+                    "Engram shared-memory creation failed on EDP rank 0: " + error
+                )
+
+            mapping = owner = tensor = None
+            try:
+                with open(path, "r+b") as file:
+                    mapping = mmap.mmap(file.fileno(), num_bytes, flags=mmap.MAP_SHARED)
+                owner = np.frombuffer(mapping, dtype=np.uint8)
+                pointer = owner.ctypes.data
+                tensor = torch.from_numpy(owner)
+                result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
+                if result.value != 0:
+                    raise RuntimeError(f"cudaHostRegister failed: {result}")
+                finalizer = weakref.finalize(owner, self._unregister, mapping, pointer)
+                finalizer.atexit = False  # type: ignore[misc]
+                # The UVA helper otherwise allocates a private pinned copy.
+                if not tensor.is_pinned():
+                    raise RuntimeError(
+                        "cudaHostRegister did not pin the shared Engram mapping"
+                    )
+            except Exception as exc:
+                error = f"{type(exc).__name__}: {exc}"
+
+            # Also fences peer mappings before the leader unlinks the file.
+            errors: list[str | None] = [None] * group.world_size
+            torch.distributed.all_gather_object(errors, error, group=group.cpu_group)
+            failures = "; ".join(
+                f"EDP rank {rank}: {error}"
+                for rank, error in enumerate(errors)
+                if error is not None
+            )
+            if failures:
+                # Dropping the owner runs the finalizer; the mapping then unmaps.
+                del tensor, owner, mapping
+                raise RuntimeError(
+                    "Engram shared-memory initialization failed: " + failures
+                )
+            assert tensor is not None
+            return tensor
+
+    @staticmethod
+    def _unregister(mapping: mmap.mmap, pointer: int) -> None:
+        # Torch storage retains the numpy owner, including through cached UVA views.
+        # Keep its mmap alive until CUDA has released the registration.
+        result = torch.cuda.cudart().cudaHostUnregister(pointer)
+        if result.value != 0:
+            logger.warning("Engram cudaHostUnregister failed: %s", result)
+
+    def load_weight(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        if self.group.rank_in_group == 0:
+            _engram_head_shard_weight_loader(param, loaded_weight)
+        # Read order may differ across ranks. Equal load counts ensure all shared
+        # weights are ready after the last weight-loader call returns.
+        torch.distributed.barrier(group=self.group.cpu_group)
+
+    def get_views(
+        self, weight: torch.Tensor, scales: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if (weight.data_ptr(), scales.data_ptr()) != (
+            self.weight.data_ptr(),
+            self.weight_scale_inv.data_ptr(),
+        ):
+            raise RuntimeError("Shared Engram parameter storage must not be replaced")
+        if self._views is None:
+            self._views = (
+                get_accelerator_view_from_cpu_tensor(self.weight),
+                get_accelerator_view_from_cpu_tensor(self.weight_scale_inv),
+            )
+        return self._views
+
+
+class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
+    """Extend TP lookup with DP head sharding or shared, CPU-offloaded TP slices."""
+
+    _shared_memory: DPSharedEngramStorage | None = None
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        dim: int,
+        head_sizes: tuple[int, ...],
+        block_size: int = 32,
+        cpu_offload: bool = False,
+        dp_shared_memory: bool = False,
+    ) -> None:
+        self.cpu_offload = cpu_offload
+        self.dp_shared_memory = dp_shared_memory
+        self.dp_size = get_engram_dp_size()
+        if dp_shared_memory:
+            if not cpu_offload:
+                raise ValueError("dp_shared_memory requires cpu_offload=True")
+            if self.dp_size <= 1:
+                raise ValueError(
+                    "dp_shared_memory requires a node-local Engram DP "
+                    f"group with size > 1; effective Engram DP size is {self.dp_size}. "
+                    "Check that the node layout and rank placement allow complete "
+                    "DP replicas to be co-located."
+                )
+            self.dp_size = 1
+        if cpu_offload and not is_uva_available():
+            raise RuntimeError("Engram CPU offload requires UVA support")
+        self._views: tuple[torch.Tensor, torch.Tensor] | None = None
+        self._view_src: tuple[int, int] | None = None
+        super().__init__(num_embeddings, dim, head_sizes, block_size)
+        if cpu_offload:
+            # Constant dummy values avoid randomizing huge CPU lookup tables.
+            set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
+            # The ue8m0 encoding of scale 1.0 is exponent byte 127.
+            set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127})
+            logger.info(
+                "Engram table offloaded to pinned host memory: %d rows x %d, "
+                "%.2f GiB %s",
+                self.part_num_embeddings,
+                dim,
+                self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
+                "shared across DP replicas" if dp_shared_memory else "per rank",
+            )
+
+    def _get_shard_info(self) -> tuple[int, int]:
+        if self.dp_size == 1:
+            return super()._get_shard_info()
+        return self.tp_size * self.dp_size, engram_head_shard_rank()
+
+    def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.dp_shared_memory:
+            group = get_engram_dp_group()
+            assert group is not None
+            storage = DPSharedEngramStorage(
+                self.part_num_embeddings, self.dim, self.block_size, group
+            )
+            self._shared_memory = storage
+            self._weight_loader = storage.load_weight
+            return storage.weight, storage.weight_scale_inv
+        if not self.cpu_offload:
+            return super()._allocate_weights()
+        # Model initialization may be inside a CUDA device context.
+        return (
+            torch.empty(
+                self.part_num_embeddings,
+                self.dim,
+                dtype=torch.float8_e4m3fn,
+                device="cpu",
+                pin_memory=True,
+            ),
+            torch.empty(
+                self.part_num_embeddings,
+                self.dim // self.block_size,
+                dtype=torch.uint8,
+                device="cpu",
+                pin_memory=True,
+            ),
+        )
+
+    def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
+        if self._shared_memory is not None:
+            return self._shared_memory.get_views(self.weight, self.weight_scale_inv)
+        if not self.cpu_offload:
+            return super()._storage()
+        src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
+        if self._view_src != src:
+            self._views = (
+                get_accelerator_view_from_cpu_tensor(self.weight.data),
+                get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
+            )
+            self._view_src = src
+        assert self._views is not None
+        return self._views
+
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        if self.dp_size == 1:
+            return super().forward(indices)
+        num_tokens = indices.shape[0]
+        indices = gather_engram_hashes(indices)
+        out = torch.empty(
+            (indices.shape[0], self.part_n_hash_cols, self.dim),
+            dtype=torch.bfloat16,
+            device=indices.device,
+        )
+        self.lookup(indices, out)
+        out = _gather_engram_rows(out, num_tokens)
+        if self.tp_size > 1:
+            out = tensor_model_parallel_all_gather(out, dim=1)
+        return out[:, : self.n_hash_cols]
+
+
+def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor:
+    """Exchange DP tokens for heads, retaining only this replica's tokens."""
+    dp_group = get_engram_dp_group()
+    assert dp_group is not None
+    slot, remainder = divmod(staged.shape[0], dp_group.world_size)
+    assert remainder == 0 and 0 <= num_tokens <= slot
+    gathered = dp_group.all_gather(staged, dim=0)
+    local_heads, dim = staged.shape[1:]
+    rows = staged.new_empty((num_tokens, dp_group.world_size * local_heads, dim))
+    _engram_select_rows(
+        gathered,
+        rows,
+        staged.shape[0],
+        dp_group.rank_in_group * slot,
+        local_heads * dim,
+    )
+    return rows
+
+
+class Engram(BaseEngram):
+    """NVIDIA Engram with asynchronous offload and node-local DP lookup."""
+
+    _prefetch_stream: torch.cuda.Stream | None = None
+
+    def _create_embedding(
+        self, layout: EngramLayout, layer_hash_index: int
+    ) -> ParallelEngramEmbedding:
+        engram_config = get_current_vllm_config().engram_config
+        assert engram_config is not None
+        return ParallelEngramEmbedding(
+            layout.num_embeddings[layer_hash_index],
+            layout.head_dim,
+            tuple(size for order in layout.primes[layer_hash_index] for size in order),
+            cpu_offload=engram_config.cpu_offload,
+            dp_shared_memory=bool(engram_config.dp_shared_memory),
+        )
+
+    def _init_staging(self, max_tokens: int, head_dim: int) -> None:
+        super()._init_staging(max_tokens * self.embed_tokens.dp_size, head_dim)
+        if self.embed_tokens.cpu_offload:
+            self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device)
+
+    def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
+        """Prefetch local shared rows or the DP group's gathered hash IDs."""
+        if self._prefetch_stream is None:
+            return super().prepare_embeddings(hash_ids)
+        rows = self.staged_rows[: hash_ids.shape[0]]
+        assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small"
+        self._start_prefetch(hash_ids, rows, self._prefetch_stream)
+
+    @eager_break_during_capture
+    def _start_prefetch(
+        self, hash_ids: torch.Tensor, rows: torch.Tensor, stream: torch.cuda.Stream
+    ) -> None:
+        # Eager boundaries let the lookup span piecewise graph segments.
+        stream.wait_stream(torch.cuda.current_stream())
+        # Keep temporary hash storage alive until lookup finishes reading it.
+        hash_ids.record_stream(stream)
+        with torch.cuda.stream(stream):
+            self.embed_tokens.lookup(hash_ids, rows, background=True)
+
+    @eager_break_during_capture
+    def _finish_prefetch(self, stream: torch.cuda.Stream) -> None:
+        torch.cuda.current_stream().wait_stream(stream)
+
+    def _ready_rows(self, num_tokens: int) -> torch.Tensor:
+        if self._prefetch_stream is not None:
+            self._finish_prefetch(self._prefetch_stream)
+        if self.embed_tokens.dp_size > 1:
+            slot = engram_gathered_num_tokens()
+            staged = self.staged_rows[: slot * self.embed_tokens.dp_size]
+            return _gather_engram_rows(staged, num_tokens)
+        return super()._ready_rows(num_tokens)
diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
new file mode 100644
index 000000000000..8efc987739b1
--- /dev/null
+++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
@@ -0,0 +1,482 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V4.1 attention on FlashMLA's mega-attention kernel (SM100).
+
+One kernel does Q RoPE, sparse attention, the inverse RoPE of the output and
+its FP8 cast, and writes straight into the buffer ``wo_a`` consumes. So the
+layer declares ``accepts_unnormed_unroped_query`` -- the kernel RoPEs Q itself,
+leaving the fused KV insert only the zero-pad to the kernel's head count -- and
+its ``_alloc_attn_out`` / ``_o_proj`` pair speaks QuantizedActivation instead
+of bf16, since the output needs no rotation or quantization here.
+
+``wq_b`` rows and ``wo_a`` columns are permuted once at load so the surrounding
+GEMMs speak the kernel's chunk-interleaved layouts directly. A step's prefill
+and decode segments write disjoint token ranges of one output buffer pair, so
+a single ``wo_a`` einsum covers the whole step.
+"""
+
+from dataclasses import replace
+from typing import TYPE_CHECKING, cast
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.cache import CacheDType
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
+from vllm.models.deepseek_v41.common.ops import (
+    combine_topk_swa_indices,
+    compute_global_topk_indices_and_lens,
+    dequantize_and_gather_k_cache,
+)
+from vllm.models.deepseek_v41.common.ops.fused_layout import (
+    WV_GROUP_SIZE,
+    permute_wo_a_,
+    permute_wq_b_,
+)
+from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention
+from vllm.models.deepseek_v41.sparse_mla import (
+    DeepseekV4FlashMLAMetadata,
+    FlashMLAMegaAttnBackend,
+)
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size
+from vllm.utils.math_utils import round_up
+from vllm.v1.attention.ops.flashmla import is_flashmla_sparse_supported
+from vllm.v1.worker.workspace import current_workspace_manager
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
+
+_HEAD_DIM_V = 512
+_ROPE_DIM = 64
+# The kernel always emits per-32 scales, whatever num_per_channels says.
+_QUANT_GROUP = 32
+# Fixed kernel conventions: GPT-J (non-neox) RoPE, TMA-aligned col-major
+# scales, round-to-nearest, packed ue8m0 -- i.e. what DeepGEMM's fp8_einsum
+# consumes. Q norm is off: the V4.1 layer norms qr before wq_b.
+_ROPE_ARGS = (False, _ROPE_DIM)
+_SF_ARGS = (_QUANT_GROUP, True, True, True)
+
+# Decode always goes through the mega kernel: there is no batch size at which
+# falling back to the split-KV decode pays, so there is no minimum-token
+# threshold and no fallback path to maintain. What decides whether this layer
+# is worth selecting at all is TP, not the batch size.
+#
+# Measured on one GB300 with benchmarks/kernels/benchmark_dsv41_mega_attn.py,
+# CUDA graphs on, counting every per-step op each path runs (64-head model,
+# topk_swa=128, topk_extra=512, NVFP4 compressed cache), as mega / split-KV
+# microseconds:
+#
+#     s_q                     1      32     128     256     512
+#     TP1 (64 live heads)  29/29   31/31   31/40   53/72   96/140
+#     TP2 (32 live)        29/29   31/30   35/35   61/63  110/119
+#     TP4 (16 live)        29/29   31/29   34/34   59/58  109/109
+#     TP8 (8 live)         29/29   31/29   33/33   59/57  107/104
+#
+# The kernel absorbs the inverse RoPE and the FP8 cast, whose cost scales with
+# *live* heads, while its own cost scales with *padded* heads -- always 64,
+# since FlashMLA's FP8 decode takes only h_q in {64, 128}. So it wins at TP1
+# (up to 1.45x, and TP1 is also where the fused insert skips the Q pad
+# entirely) and from TP2 down the two cancel to within a few percent, with a
+# fixed ~2 us penalty at small s_q where the padded-head cost dominates.
+# Measure before assuming a win at high TP -- and measure under CUDA graphs:
+# in eager mode launch overhead swamps both paths and overstates the fused one
+# several-fold.
+
+
+def is_flashmla_mega_attn_supported() -> tuple[bool, str | None]:
+    """Whether FlashMLA's fused mega-attention kernels are usable here."""
+    is_available, maybe_reason = is_flashmla_sparse_supported()
+    if not is_available:
+        return False, maybe_reason
+    if not current_platform.is_device_capability_family(100):
+        return False, "FlashMLA mega attention requires sm_10x GPUs."
+    if not hasattr(torch.ops._flashmla_C, "fused_norm_rope_attn_rope_cast_decode"):
+        return False, "vllm._flashmla_C was built without the mega attention ops."
+    if not hasattr(torch.ops._flashmla_C, "permute_q_b_proj"):
+        return False, "vllm._flashmla_C predates the mega attention weight permutes."
+    return True, None
+
+
+def alloc_mega_attn_output(
+    num_tokens: int,
+    n_wv_group: int,
+    device: torch.device,
+) -> QuantizedActivation:
+    """Allocate the output buffer pair the mega-attention kernels write into.
+
+    One pair per forward step: the prefill and decode segments fill disjoint
+    token ranges, so a single ``fp8_einsum`` over all ``num_tokens`` can
+    consume the result instead of one call per segment.
+
+    The scale buffer is MN-major -- stride 1 along tokens, head-dim stride
+    ``ceil4(num_tokens)`` -- which is both what the kernel requires of a
+    caller-provided buffer and what DeepGEMM expects for this N.
+    """
+    d = WV_GROUP_SIZE * _HEAD_DIM_V
+    data = torch.empty(
+        (num_tokens, n_wv_group, d), dtype=torch.float8_e4m3fn, device=device
+    )
+    aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize)
+    scale = torch.empty(
+        (n_wv_group, d // (_QUANT_GROUP * 4), aligned),
+        dtype=torch.int32,
+        device=device,
+    ).permute(2, 0, 1)[:num_tokens]
+    return QuantizedActivation(
+        data=data,
+        scale=scale,
+        orig_dtype=torch.bfloat16,
+        orig_shape=data.shape,
+        quant_key=kMxfp8Dynamic,
+    )
+
+
+def _token_slice(out: QuantizedActivation, start: int, end: int) -> QuantizedActivation:
+    """The ``[start, end)`` token slice of a mega-attention output buffer."""
+    data = out.data[start:end]
+    return replace(out, data=data, scale=out.scale[start:end], orig_shape=data.shape)
+
+
+class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention):
+    """FlashMLA mega-attention layer for DeepSeek V4.1 (SM100)."""
+
+    backend_cls = FlashMLAMegaAttnBackend
+
+    @classmethod
+    def is_available_for(cls, vllm_config: VllmConfig) -> bool:
+        """Whether this layer can serve the configured model on this device.
+
+        The kernel is SM100-only and wants ``WV_GROUP_SIZE`` heads per ``wo_a``
+        group, which stops holding once TP divides the head count far enough
+        (TP16 on a 64-head model leaves 4). ``__init__`` raises on both, so the
+        default-backend selector asks here first rather than turning an
+        unsupported topology into a startup crash.
+        """
+        if not is_flashmla_mega_attn_supported()[0]:
+            return False
+        config = vllm_config.model_config.hf_config
+        config = getattr(config, "text_config", config)
+        n_heads = getattr(config, "num_attention_heads", 0) or 0
+        n_groups = getattr(config, "o_groups", 0) or 0
+        tp_size = vllm_config.parallel_config.tensor_parallel_size
+        if n_heads % tp_size or n_groups % tp_size:
+            return False
+        n_local_heads, n_local_groups = n_heads // tp_size, n_groups // tp_size
+        if not n_local_groups or n_local_heads % WV_GROUP_SIZE:
+            return False
+        return n_local_heads // n_local_groups == WV_GROUP_SIZE
+
+    def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
+        super().__init__(vllm_config, *args, **kwargs)
+        if self.n_local_heads % WV_GROUP_SIZE:
+            raise ValueError(
+                f"{self.prefix}: mega attention needs the local head count "
+                f"({self.n_local_heads}) to be a multiple of {WV_GROUP_SIZE}."
+            )
+        if self.n_local_heads // self.n_local_groups != WV_GROUP_SIZE:
+            raise ValueError(
+                f"{self.prefix}: mega attention needs {WV_GROUP_SIZE} heads per "
+                "wo_a group."
+            )
+        self.n_wv_group = self.padded_heads // WV_GROUP_SIZE
+        self._fused_layouts_ready = False
+
+    # ---- interface contract ------------------------------------------------
+
+    @property
+    def accepts_unnormed_unroped_query(self) -> bool:
+        return True
+
+    @property
+    def packed_kv_cache_dtype(self) -> CacheDType:
+        # This kernel is the only one that reads an NVFP4 compressed cache, so
+        # it is what an unspecific --kv-cache-dtype resolves to here.
+        return "nvfp4_ds_mla"
+
+    def _alloc_attn_out(
+        self, num_tokens: int, hidden_states: torch.Tensor
+    ) -> QuantizedActivation:
+        return alloc_mega_attn_output(num_tokens, self.n_wv_group, hidden_states.device)
+
+    def _o_proj(
+        self, attn_out: QuantizedActivation, positions: torch.Tensor
+    ) -> torch.Tensor:
+        """Grouped ``wo_a`` then ``wo_b`` over the kernel's quantized output.
+
+        The inverse RoPE and the FP8 cast happened inside the attention
+        kernel, and ``wo_a`` is permuted for its output layout, so there is
+        nothing to rotate or quantize here and ``positions`` is unused. The
+        scale already is DeepGEMM's packed-ue8m0 MN-major layout, so the
+        einsum consumes the kernel's output with no repacking.
+        """
+        groups = self.n_local_groups
+        z = torch.empty(
+            (attn_out.data.shape[0], groups, self.o_lora_rank),
+            dtype=torch.bfloat16,
+            device=attn_out.data.device,
+        )
+        fp8_einsum(
+            "bhr,hdr->bhd",
+            (attn_out.data[:, :groups], attn_out.scale[:, :groups]),
+            (self.wo_a.weight, self.wo_a.weight_scale),
+            z,
+            recipe=self._einsum_recipe,
+        )
+        return self._wo_b_proj(z.flatten(1))
+
+    # ---- weights -----------------------------------------------------------
+
+    def finalize_loaded_weights(self) -> None:
+        """Permute wq_b rows / wo_a columns into the kernel's layouts.
+
+        Idempotent: a second post-load pass must not permute twice.
+        """
+        if self._fused_layouts_ready:
+            return
+        permute_wq_b_(
+            self.wq_b.weight.data, self.wq_b.weight_scale.data, self.n_local_heads
+        )
+        permute_wo_a_(
+            self.wo_a.weight.data,
+            self.wo_a.weight_scale.data,
+            self.n_local_heads // self.n_local_groups,
+        )
+        self._fused_layouts_ready = True
+
+    # ---- forward -----------------------------------------------------------
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        positions: torch.Tensor,
+        output: QuantizedActivation,
+    ) -> None:
+        if not self._fused_layouts_ready:
+            raise RuntimeError(
+                f"{self.prefix}: wq_b / wo_a were never permuted for the mega "
+                "attention kernel; refusing to run with mismatched layouts."
+            )
+        attn_metadata = get_forward_context().attn_metadata
+        if attn_metadata is None:
+            # Warmup dummy run: reserve the prefill workspace, produce zeros.
+            self._reserve_prefill_workspace(q)
+            output.data.zero_()
+            output.scale.zero_()
+            return
+
+        assert isinstance(attn_metadata, dict)
+        # q was zero-padded to the kernel's head count by the fused KV insert.
+        assert q.shape[1] == self.padded_heads
+        flashmla_metadata = cast(
+            DeepseekV4FlashMLAMetadata | None,
+            attn_metadata.get(self.compressed_cache_prefix)
+            if self.compressed_cache_prefix is not None
+            else None,
+        )
+        swa_metadata = cast(
+            "DeepseekSparseSWAMetadata", attn_metadata[self.swa_cache_layer.prefix]
+        )
+        # The kernel takes int32 RoPE positions.
+        positions_int32 = positions.to(torch.int32)
+        num_decode_tokens = swa_metadata.num_decode_tokens
+
+        if swa_metadata.num_prefills > 0:
+            self._forward_prefill_mega(
+                q[num_decode_tokens:],
+                positions_int32[num_decode_tokens:],
+                flashmla_metadata,
+                swa_metadata,
+                output,
+                num_decode_tokens,
+            )
+        if swa_metadata.num_decodes > 0:
+            self._forward_decode_mega(
+                q[:num_decode_tokens],
+                positions_int32[:num_decode_tokens],
+                flashmla_metadata,
+                swa_metadata,
+                _token_slice(output, 0, num_decode_tokens),
+            )
+
+    def _reserve_prefill_workspace(self, q: torch.Tensor) -> None:
+        swa_only = self.compress_ratio == 0
+        n = 0 if swa_only else -(-self.max_model_len // self.compress_ratio)
+        m = n + self.window_size + self.max_num_batched_tokens
+        if swa_only:
+            top_k = 0
+        else:
+            assert self.topk_indices_buffer is not None
+            top_k = self.topk_indices_buffer.shape[-1]
+        combined_topk = round_up(top_k + self.window_size, 128)
+        current_workspace_manager().get_simultaneous(
+            ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16),
+            ((self.max_num_batched_tokens, combined_topk), torch.int32),
+            ((self.max_num_batched_tokens,), torch.int32),
+        )
+
+    def _decode_compressed_kv_and_topk(
+        self,
+        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
+        """The compressed cache and its ``[n, topk]`` slot indices / lengths."""
+        if self.compress_ratio == 0:
+            return None, None, None
+        assert flashmla_metadata is not None
+        assert swa_metadata.is_valid_token is not None
+        assert self.topk_indices_buffer is not None
+        num_decode_tokens = swa_metadata.num_decode_tokens
+        indices, lens = compute_global_topk_indices_and_lens(
+            self.topk_indices_buffer[:num_decode_tokens],
+            swa_metadata.token_to_req_indices,
+            flashmla_metadata.block_table[: swa_metadata.num_decodes],
+            flashmla_metadata.block_size // self.compress_ratio,
+            swa_metadata.is_valid_token[:num_decode_tokens],
+        )
+        return self._compressed_kv_cache().unsqueeze(-2), indices, lens
+
+    def _forward_decode_mega(
+        self,
+        q: torch.Tensor,
+        positions_int32: torch.Tensor,
+        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+        out: QuantizedActivation,
+    ) -> None:
+        """Mega attention over the paged quantized caches, in place into ``out``.
+
+        The caches are ``[num_blocks, page, 1, bytes]``, the record taken from
+        ``bytes`` (528 V4.1 fp8, 288 V4.1 NVFP4; NVFP4 only as the compressed
+        cache beside an fp8 SWA one), and the indices ``[s_q, topk]`` int32
+        slot ids (``block * page + offset``, ``-1`` invalid).
+        """
+        extra_cache, extra_idx, extra_len = self._decode_compressed_kv_and_topk(
+            flashmla_metadata, swa_metadata
+        )
+        assert swa_metadata.decode_swa_indices is not None
+        torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode(
+            q,
+            self.swa_cache_layer.kv_cache.unsqueeze(-2),
+            swa_metadata.decode_swa_indices.view(q.shape[0], -1),
+            self.scale,
+            _HEAD_DIM_V,
+            self.attn_sink,
+            swa_metadata.decode_swa_lens,
+            extra_cache,
+            extra_idx,
+            extra_len,
+            False,  # enable_q_norm
+            0.0,  # rms_norm_eps
+            positions_int32,
+            *_ROPE_ARGS,
+            self.rotary_emb.cos_sin_cache,
+            self.n_wv_group,
+            *_SF_ARGS,
+            out.data,
+            out.scale,
+        )
+
+    def _forward_prefill_mega(
+        self,
+        q: torch.Tensor,
+        positions_int32: torch.Tensor,
+        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+        out: QuantizedActivation,
+        token_base: int,
+    ) -> None:
+        swa_only = self.compress_ratio == 0
+        num_decodes = swa_metadata.num_decodes
+        num_decode_tokens = swa_metadata.num_decode_tokens
+        seq_lens = swa_metadata.prefill_seq_lens
+        gather_lens = swa_metadata.prefill_gather_lens
+        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
+        query_start_loc = swa_metadata.query_start_loc
+        assert seq_lens is not None and gather_lens is not None
+        assert query_start_loc_cpu is not None and query_start_loc is not None
+        prefill_token_base = query_start_loc_cpu[num_decodes]
+        assert self.topk_indices_buffer is not None
+        topk_indices = self.topk_indices_buffer[num_decode_tokens:][
+            : swa_metadata.num_prefill_tokens
+        ]
+        top_k = 0 if swa_only else topk_indices.shape[-1]
+        chunk_plan = swa_metadata.get_prefill_chunk_plan(
+            compress_ratio=self.compress_ratio,
+            prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
+            has_compressed=not swa_only,
+        )
+        assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
+        workspace_manager = current_workspace_manager()
+        combined_topk = round_up(top_k + self.window_size, 128)
+        for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan:
+            chunk_size = chunk_end - chunk_start
+            kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous(
+                ((chunk_size, chunk_m, q.shape[-1]), torch.bfloat16),
+                ((self.max_num_batched_tokens, combined_topk), torch.int32),
+                ((self.max_num_batched_tokens,), torch.int32),
+            )
+            if not swa_only:
+                assert flashmla_metadata is not None
+                dequantize_and_gather_k_cache(
+                    kv_ws[:chunk_size, :chunk_n],
+                    self._compressed_kv_cache(),
+                    seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
+                    gather_lens=None,
+                    block_table=flashmla_metadata.block_table[num_decodes:][
+                        chunk_start:chunk_end
+                    ],
+                    block_size=flashmla_metadata.block_size // self.compress_ratio,
+                    offset=0,
+                )
+            dequantize_and_gather_k_cache(
+                kv_ws[:chunk_size],
+                self.swa_cache_layer.kv_cache,
+                seq_lens=seq_lens[chunk_start:chunk_end],
+                gather_lens=gather_lens[chunk_start:chunk_end],
+                block_table=swa_metadata.block_table[num_decodes:][
+                    chunk_start:chunk_end
+                ],
+                block_size=swa_metadata.block_size,
+                offset=chunk_n,
+            )
+            first, last = num_decodes + chunk_start, num_decodes + chunk_end
+            qs = int(query_start_loc_cpu[first] - prefill_token_base)
+            qe = int(query_start_loc_cpu[last] - prefill_token_base)
+            combined_indices, combined_lens = combine_topk_swa_indices(
+                topk_indices[qs:qe],
+                query_start_loc[first : last + 1],
+                seq_lens[chunk_start:chunk_end],
+                gather_lens[chunk_start:chunk_end],
+                self.window_size,
+                self.compress_ratio,
+                top_k,
+                chunk_m,
+                chunk_n,
+                out=(idx_ws[: qe - qs], lens_ws[: qe - qs]),
+            )
+            chunk_out = _token_slice(out, token_base + qs, token_base + qe)
+            # Mega attention over the gathered non-paged bf16 KV (RoPE already
+            # applied), in place into this chunk's token range. `indices`
+            # entries outside [0, s_kv) skip.
+            torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_fwd(
+                q[qs:qe],
+                kv_ws.view(-1, 1, q.shape[-1]),
+                combined_indices.unsqueeze(1),
+                self.scale,
+                _HEAD_DIM_V,
+                self.attn_sink,
+                combined_lens,
+                False,  # enable_q_norm
+                0.0,  # rms_norm_eps
+                positions_int32[qs:qe],
+                *_ROPE_ARGS,
+                self.rotary_emb.cos_sin_cache,
+                self.n_wv_group,
+                *_SF_ARGS,
+                chunk_out.data,
+                chunk_out.scale,
+            )
diff --git a/vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
similarity index 90%
rename from vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py
rename to vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
index 802ddede9c64..423849dcf1eb 100644
--- a/vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py
+++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
@@ -13,12 +13,12 @@
     compute_fp8_einsum_recipe,
     deep_gemm_fp8_o_proj,
 )
-from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
-from vllm.models.deepseek_v4_1.common.ops import (
+from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+from vllm.models.deepseek_v41.common.ops import (
     build_flashinfer_mixed_sparse_indices,
     compute_global_topk_indices_and_lens,
 )
-from vllm.models.deepseek_v4_1.sparse_mla import (
+from vllm.models.deepseek_v41.sparse_mla import (
     DeepseekV4FlashMLAMetadata,
     DeepseekV4SparseMLABackend,
     DeepseekV4SparseMLAMetadataBuilder,
@@ -26,7 +26,11 @@
 )
 from vllm.platforms.interface import DeviceCapability
 from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
-from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf
+from vllm.v1.attention.backend import (
+    AttentionCGSupport,
+    CommonAttentionMetadata,
+    MultipleOf,
+)
 from vllm.v1.attention.backends.mla.compressor_utils import (
     get_dspark_swa_index_width,
 )
@@ -109,7 +113,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [128]
 
     @staticmethod
@@ -183,6 +187,42 @@ class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBui
 
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Graphs retain these addresses while each build refreshes their contents.
+        self._decode_topk_lens = torch.empty(
+            self._max_tokens, dtype=torch.int32, device=self.device
+        )
+        self._decode_seq_lens = torch.empty_like(self._decode_topk_lens)
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+        replay_start: torch.Tensor | None = None,
+    ) -> "DeepseekSparseSWAMetadata":
+        metadata = super().build(
+            common_prefix_len, common_attn_metadata, fast_build, replay_start
+        )
+        num_tokens = metadata.num_decode_tokens
+        if not common_attn_metadata.causal and num_tokens > 0:
+            assert metadata.decode_swa_lens is not None
+            assert metadata.seq_lens is not None
+            assert metadata.token_to_req_indices is not None
+            topk_lens = self._decode_topk_lens[:num_tokens]
+            seq_lens = self._decode_seq_lens[:num_tokens]
+            torch.clamp(metadata.decode_swa_lens, min=self.window_size, out=topk_lens)
+            torch.index_select(
+                metadata.seq_lens,
+                0,
+                metadata.token_to_req_indices[:num_tokens],
+                out=seq_lens,
+            )
+            metadata.flashinfer_decode_topk_lens = topk_lens
+            metadata.flashinfer_decode_seq_lens = seq_lens
+        return metadata
+
 
 class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
     @staticmethod
@@ -201,13 +241,14 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
     def get_padded_num_q_heads(cls, num_heads: int) -> int:
         return _pad_to_supported_q_heads(num_heads)
 
-    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+    def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        o = attn_out[:, : self.n_local_heads, :]
         return deep_gemm_fp8_o_proj(
             o,
             positions,
             self.rotary_emb.cos_sin_cache,
             self.wo_a,
-            self.wo_b,
+            self._wo_b_proj,
             n_groups=self.n_local_groups,
             heads_per_group=self.n_local_heads // self.n_local_groups,
             nope_dim=self.nope_head_dim,
@@ -337,6 +378,7 @@ def _build_sparse_index_metadata(
         assert swa_metadata.token_to_req_indices is not None
         assert swa_metadata.decode_swa_indices is not None
         assert swa_metadata.block_table is not None
+        assert swa_metadata.replay_start is not None
 
         decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
             num_decode_tokens, swa_metadata.decode_swa_width
@@ -416,10 +458,7 @@ def _build_sparse_index_metadata(
                 decode_is_valid_token=decode_is_valid_token,
                 swa_block_span=swa_block_span,
                 compressed_block_span=compressed_block_span,
-                prefill_left_visible=swa_metadata.prefill_left_visible,
-                prefill_right_visible=swa_metadata.prefill_right_visible,
-                # getattr for tests that bypass __init__ via object.__new__.
-                max_image_tokens=getattr(self, "max_image_tokens", 0),
+                replay_start=swa_metadata.replay_start[:num_reqs],
             )
             if swa_only:
                 swa_metadata.flashinfer_sparse_index_cache["swa_only"] = (
@@ -496,21 +535,41 @@ def _forward(
         # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
         # uniform-q batches, and this avoids flattening mixed batches into one call.
         if num_decode_tokens > 0:
+            decode_query = query[:num_decode_tokens]
+            decode_output = output[:num_decode_tokens]
             decode_cu = query_start_loc[: num_decodes + 1]
+            decode_seq_lens = seq_lens[:num_decodes]
+            decode_topk_lens = sparse_topk_lens[:num_decode_tokens]
+            max_decode_query_len = swa_metadata.max_decode_query_len
+            if swa_metadata.decode_swa_width > self.window_size:
+                # DSpark's non-causal window extends past the fixed 128 SWA
+                # columns into the aliased compressed pool. Exclude padding,
+                # and expose the full block to each query instead of letting
+                # TRTLLM derive a causal SWA length from its query position.
+                assert swa_only
+                assert swa_metadata.flashinfer_decode_topk_lens is not None
+                assert swa_metadata.flashinfer_decode_seq_lens is not None
+                decode_topk_lens = swa_metadata.flashinfer_decode_topk_lens
+                decode_seq_lens = swa_metadata.flashinfer_decode_seq_lens
+                decode_query = decode_query.unsqueeze(1)
+                decode_output = decode_output.unsqueeze(1)
+                decode_cu = None
+                max_decode_query_len = 1
             flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
-                query=query[:num_decode_tokens],
+                query=decode_query,
                 swa_kv_cache=swa_k_cache,
                 workspace_buffer=workspace,
                 sparse_indices=sparse_indices[:num_decode_tokens],
+                sparse_indices_are_storage_offsets=True,
                 compressed_kv_cache=compressed_kv_cache,
-                sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
-                seq_lens=seq_lens[:num_decodes],
-                out=output[:num_decode_tokens],
+                sparse_topk_lens=decode_topk_lens,
+                seq_lens=decode_seq_lens,
+                out=decode_output,
                 bmm1_scale=bmm1_scale,
                 bmm2_scale=bmm2_scale,
                 sinks=self.attn_sink,
                 cum_seq_lens_q=decode_cu,
-                max_q_len=swa_metadata.max_decode_query_len,
+                max_q_len=max_decode_query_len,
             )
 
         if num_prefill_tokens > 0:
@@ -527,6 +586,7 @@ def _forward(
                 swa_kv_cache=swa_k_cache,
                 workspace_buffer=workspace,
                 sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
+                sparse_indices_are_storage_offsets=True,
                 compressed_kv_cache=compressed_kv_cache,
                 sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
                 seq_lens=seq_lens[num_decodes:num_reqs],
@@ -562,13 +622,14 @@ def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor:
     def get_padded_num_q_heads(cls, num_heads: int) -> int:
         return _pad_to_supported_q_heads(num_heads)
 
-    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+    def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        o = attn_out[:, : self.n_local_heads, :]
         return deep_gemm_fp8_o_proj(
             o,
             positions,
             self.rotary_emb.cos_sin_cache,
             self.wo_a,
-            self.wo_b,
+            self._wo_b_proj,
             n_groups=self.n_local_groups,
             heads_per_group=self.n_local_heads // self.n_local_groups,
             nope_dim=self.nope_head_dim,
diff --git a/vllm/models/deepseek_v4_1/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py
similarity index 92%
rename from vllm/models/deepseek_v4_1/nvidia/flashmla.py
rename to vllm/models/deepseek_v41/nvidia/flashmla.py
index 86d8fb2031ca..96dde36c61d1 100644
--- a/vllm/models/deepseek_v4_1/nvidia/flashmla.py
+++ b/vllm/models/deepseek_v41/nvidia/flashmla.py
@@ -10,13 +10,13 @@
     compute_fp8_einsum_recipe,
     deep_gemm_fp8_o_proj,
 )
-from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
-from vllm.models.deepseek_v4_1.common.ops import (
+from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+from vllm.models.deepseek_v41.common.ops import (
     combine_topk_swa_indices,
     compute_global_topk_indices_and_lens,
     dequantize_and_gather_k_cache,
 )
-from vllm.models.deepseek_v4_1.sparse_mla import (
+from vllm.models.deepseek_v41.sparse_mla import (
     DeepseekV4FlashMLABackend,
     DeepseekV4FlashMLAMetadata,
     DeepseekV41SparseSWAMetadataBuilder,
@@ -58,13 +58,14 @@ def __init__(self, *args, **kwargs) -> None:
             self._o_proj_block_size
         )
 
-    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+    def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        o = attn_out[:, : self.n_local_heads, :]
         return deep_gemm_fp8_o_proj(
             o,
             positions,
             self.rotary_emb.cos_sin_cache,
             self.wo_a,
-            self.wo_b,
+            self._wo_b_proj,
             n_groups=self.n_local_groups,
             heads_per_group=self.n_local_heads // self.n_local_groups,
             nope_dim=self.nope_head_dim,
@@ -119,9 +120,7 @@ def forward_mqa(
             else:
                 assert self.topk_indices_buffer is not None
                 top_k = self.topk_indices_buffer.shape[-1]
-            combined_topk = round_up(
-                top_k + self.window_size + self.max_image_tokens, 128
-            )
+            combined_topk = round_up(top_k + self.window_size, 128)
             current_workspace_manager().get_simultaneous(
                 ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
                 ((self.max_num_batched_tokens, combined_topk), torch.int32),
@@ -302,7 +301,7 @@ def _forward_prefill(
         )
         assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
         workspace_manager = current_workspace_manager()
-        combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128)
+        combined_topk = round_up(top_k + self.window_size, 128)
         for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
             chunk_size = chunk_end - chunk_start
             workspace = workspace_manager.get_simultaneous(
@@ -360,21 +359,6 @@ def _forward_prefill(
                 chunk_M,
                 chunk_N,
                 out=(combined_indices_out, combined_lens_out),
-                left_visible=(
-                    swa_metadata.prefill_left_visible[
-                        num_decode_tokens + query_start : num_decode_tokens + query_end
-                    ]
-                    if swa_metadata.prefill_left_visible is not None
-                    else None
-                ),
-                right_visible=(
-                    swa_metadata.prefill_right_visible[
-                        num_decode_tokens + query_start : num_decode_tokens + query_end
-                    ]
-                    if swa_metadata.prefill_right_visible is not None
-                    else None
-                ),
-                max_image_tokens=self.max_image_tokens,
             )
             flash_mla_sparse_fwd(
                 q=q[query_start:query_end],
diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py
similarity index 75%
rename from vllm/models/deepseek_v4_1/nvidia/model.py
rename to vllm/models/deepseek_v41/nvidia/model.py
index 6d803736ec5c..22674bf7054c 100644
--- a/vllm/models/deepseek_v4_1/nvidia/model.py
+++ b/vllm/models/deepseek_v41/nvidia/model.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import typing
 from collections.abc import Callable, Iterable
+from functools import partial
 from itertools import islice
 
 import regex as re
@@ -12,11 +13,17 @@
 from vllm.config import VllmConfig
 from vllm.config.kernel import MEGA_MOE_BACKENDS
 from vllm.distributed import (
+    get_engram_dp_size,
     get_pp_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.forward_context import (
+    get_forward_context,
+    in_piecewise_cudagraph,
+    is_forward_context_available,
 )
-from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.mhc.tilelang import (
     mhc_post_tilelang,
@@ -60,22 +67,35 @@
     DeepseekV4MoE as DeepseekV4MoEBase,
 )
 from vllm.models.deepseek_v4.nvidia.model import (
+    MegaGateRoutingMetadata,
     make_deepseek_v4_expert_params_mapping,
+    prepare_mega_gate_routing_metadata,
+)
+from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
+    DeepseekV4MegaAttnAttention,
 )
-from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
-from vllm.models.deepseek_v4_1.nvidia.flashinfer_sparse import (
+from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import (
     DeepseekV4FlashInferMLAAttention,
     DeepseekV4FlashInferSM120Attention,
 )
-from vllm.models.deepseek_v4_1.nvidia.flashmla import DeepseekV4FlashMLAAttention
+from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
-from ..common.engram import Engram, EngramLayout, NgramHashState
+from ..common.engram import EngramLayout, NgramHashState
 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
+from .engram import Engram, can_share_engram_tables, gather_engram_hashes
+from .ops.mhc import (
+    MHC_OVERLAP_MAX_TOKENS,
+    mhc_pre_delayed_overlap,
+    mhc_shifted_post_pre,
+    supports_mhc_all_reduce,
+    supports_mhc_overlap,
+)
 
 if typing.TYPE_CHECKING:
     from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -89,6 +109,7 @@ def __init__(
         vllm_config: VllmConfig,
         prefix: str = "",
         use_sequence_parallel: bool = False,
+        reduce_results: bool = True,
     ):
         config = vllm_config.model_config.hf_config
         n_routed_experts = config.n_routed_experts
@@ -107,6 +128,7 @@ def __init__(
             n_routed_experts=n_routed_experts,
             n_activated_experts=n_activated_experts,
             num_hash_layers=0,
+            reduce_results=reduce_results,
             image_sentinel_lo=IMAGE_SENTINEL_BASE_ID,
         )
 
@@ -116,8 +138,9 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
 
     The generic CUDA backend selector does not instantiate DSv4 layers directly,
     so map generic sparse-MLA choices to the DSv4-specialized attention class.
-    Without an explicit backend, SM12 defaults to FlashInfer while the other
-    CUDA arches keep the FlashMLA path.
+    Without an explicit backend: SM12 takes FlashInfer, SM100 takes mega
+    attention where the topology allows it, and everything else keeps the
+    FlashMLA path.
     """
     backend = vllm_config.attention_config.backend
     device_capability = current_platform.get_device_capability()
@@ -137,6 +160,8 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
         if device_capability is not None and device_capability.major == 12:
             return DeepseekV4FlashInferSM120Attention
         return DeepseekV4FlashInferMLAAttention
+    if backend is AttentionBackendEnum.FLASHMLA_MEGA_ATTN_DSV41:
+        return DeepseekV4MegaAttnAttention
     if backend in (
         AttentionBackendEnum.FLASHMLA_SPARSE,
         AttentionBackendEnum.FLASHMLA_SPARSE_DSV4,
@@ -146,6 +171,14 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
 
     if device_capability is not None and device_capability.major == 12:
         return DeepseekV4FlashInferSM120Attention
+    # Mega attention is the SM100 default: it fuses Q RoPE, sparse attention,
+    # the output's inverse RoPE and its FP8 cast into one launch, and brings
+    # the 288 B NVFP4 compressed record -- the format the reference
+    # implementation itself stores. It declines topologies it cannot serve
+    # (non-SM100, TP that leaves fewer than WV_GROUP_SIZE heads per wo_a
+    # group, a build without the kernel), which then fall through to FlashMLA.
+    if DeepseekV4MegaAttnAttention.is_available_for(vllm_config):
+        return DeepseekV4MegaAttnAttention
     return DeepseekV4FlashMLAAttention
 
 
@@ -160,6 +193,30 @@ def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
     )
 
 
+def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool:
+    """Set up the fused ``wo_b`` GEMM + reduce-scatter when opted in.
+
+    Sequence parallel is the only topology where ``wo_b`` ends in a
+    reduce-scatter, so the kernel is bound to RS mode. The workspace is a
+    process-wide singleton shared with any other model code in this worker
+    (the DSpark drafter reuses it), and its NVLink multicast rendezvous is
+    collective, so every TP rank must take the same decision here.
+    """
+    if not (use_sequence_parallel and envs.VLLM_ENABLE_GEMM_RS):
+        return False
+
+    # The kernel module pulls in cute_dsl, so import it only once opted in.
+    from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
+        maybe_init_gemm_rs_ar,
+    )
+
+    hidden_size = vllm_config.model_config.hf_config.hidden_size
+    if not maybe_init_gemm_rs_ar(vllm_config, N=hidden_size, all_reduce=False):
+        return False
+    logger.info_once("To disable DeepSeek-V4.1 GEMM-RS, set VLLM_ENABLE_GEMM_RS=0.")
+    return True
+
+
 class DeepseekV4DecoderLayer(nn.Module):
     def __init__(
         self,
@@ -169,11 +226,16 @@ def __init__(
         aux_stream_list: list[torch.cuda.Stream] | None = None,
         candidate_block_buffer: torch.Tensor | None = None,
         engram_layout: EngramLayout | None = None,
+        run_gemm_rs: bool = False,
+        mhc_stream: torch.cuda.Stream | None = None,
+        fuse_mhc_all_reduce: bool = False,
     ):
         super().__init__()
 
         config = vllm_config.model_config.hf_config
         self.hidden_size = config.hidden_size
+        self.mhc_stream = mhc_stream
+        self.fuse_mhc_all_reduce = fuse_mhc_all_reduce
         self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
 
         self.engram: Engram | None = None
@@ -197,12 +259,17 @@ def __init__(
             aux_stream_list=aux_stream_list,
             candidate_block_buffer=candidate_block_buffer,
         )
-        if self.use_sequence_parallel:
+        if self.use_sequence_parallel or fuse_mhc_all_reduce:
             self.attn.wo_b.reduce_results = False
+            if run_gemm_rs:
+                # Binds only when wo_b's kernel and shape qualify; otherwise
+                # forward keeps the separate reduce-scatter below.
+                self.attn.bind_gemm_rs()
         self.ffn = DeepseekV4MoE(
             vllm_config,
             prefix=f"{prefix}.ffn",
             use_sequence_parallel=self.use_sequence_parallel,
+            reduce_results=not fuse_mhc_all_reduce,
         )
 
         self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
@@ -212,30 +279,68 @@ def __init__(
         self.hc_eps = config.hc_eps
         self.hc_post_alpha = 2.0
         if vllm_config.kernel_config.enable_jit_warmup and current_platform.is_cuda():
+            from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+                mhc_fused_post_pre_splits,
+            )
             from vllm.model_executor.kernels.mhc.warmup import MHC_PRE_NORM_KERNEL
 
             max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
             if self.use_sequence_parallel:
                 tp_size = vllm_config.parallel_config.tensor_parallel_size
                 max_tokens = (max_tokens + tp_size - 1) // tp_size
-            for input_size, use_pre_mix in (
-                (self.hidden_size, False),
-                (self.hc_mult * self.hidden_size, False),
-                (self.hc_mult * self.hidden_size, True),
-            ):
-                MHC_PRE_NORM_KERNEL.register_warmup(
-                    max_tokens=max_tokens,
-                    hidden_size=self.hidden_size,
-                    rms_eps=self.rms_norm_eps,
-                    hc_pre_eps=self.hc_eps,
-                    hc_sinkhorn_eps=self.hc_eps,
-                    hc_post_mult_value=self.hc_post_alpha,
-                    sinkhorn_repeat=self.hc_sinkhorn_iters,
-                    norm_eps=self.rms_norm_eps,
-                    hc_mult=self.hc_mult,
-                    use_pre_mix_in=use_pre_mix,
-                    rms_numel=input_size,
+            # The epilogue compiles per projection width and per pre-mix mode.
+            # The first layer projects the broadcast embedding, so it reads one
+            # hidden_size-wide row and selects stream zero; every later sublayer
+            # projects the full hc stream, and collapses it with the pre-mix the
+            # previous sublayer carried in. Those later shapes also arrive
+            # through the fused post + pre-norm GEMM, which picks split-k
+            # factors the token sweep below never produces.
+            broadcast_embedding = {
+                "rms_numel": self.hidden_size,
+                "use_pre_mix_in": False,
+                "extra_splits": (),
+            }
+            hc_stream = {
+                "rms_numel": self.hc_mult * self.hidden_size,
+                "extra_splits": mhc_fused_post_pre_splits(
+                    self.hidden_size, self.hc_mult
+                ),
+            }
+            variants = [
+                broadcast_embedding,
+                {**hc_stream, "use_pre_mix_in": False},
+                {**hc_stream, "use_pre_mix_in": True},
+            ]
+            if vllm_config.speculative_config is not None:
+                # Draft setups read aux hidden states out of the same collapse,
+                # which is its own specialization of the kernel.
+                variants.append(
+                    {**hc_stream, "use_pre_mix_in": True, "write_aux": True}
                 )
+            # Shifted mHC splits the epilogue in two: the input collapse on the
+            # caller stream and coefficient generation on a side stream.
+            split_modes = (
+                ("fused", "stats", "input") if mhc_stream is not None else ("fused",)
+            )
+            for variant in variants:
+                for split_mode in split_modes:
+                    MHC_PRE_NORM_KERNEL.register_warmup(
+                        split_mode=split_mode,
+                        max_tokens=(
+                            max_tokens
+                            if split_mode == "fused"
+                            else min(max_tokens, MHC_OVERLAP_MAX_TOKENS)
+                        ),
+                        hidden_size=self.hidden_size,
+                        rms_eps=self.rms_norm_eps,
+                        hc_pre_eps=self.hc_eps,
+                        hc_sinkhorn_eps=self.hc_eps,
+                        hc_post_mult_value=self.hc_post_alpha,
+                        sinkhorn_repeat=self.hc_sinkhorn_iters,
+                        norm_eps=self.rms_norm_eps,
+                        hc_mult=self.hc_mult,
+                        **variant,
+                    )
         mix_hc = (2 + self.hc_mult) * self.hc_mult
         hc_dim = self.hc_mult * self.hidden_size
         self.hc_attn_fn = nn.Parameter(
@@ -293,7 +398,31 @@ def forward(
         residual: torch.Tensor | None = None,
         engram_hashes: torch.Tensor | None = None,
         engram_mask: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        *,
+        capture_previous_aux: bool = False,
+        mega_gate_metadata: MegaGateRoutingMetadata | None = None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | None,
+    ]:
+        previous_aux: torch.Tensor | None = None
+        mhc_stream = self.mhc_stream
+        if mhc_stream is not None and (
+            in_piecewise_cudagraph()
+            or not 0 < positions.shape[0] <= MHC_OVERLAP_MAX_TOKENS
+            or not torch.cuda.is_current_stream_capturing()
+        ):
+            # Use overlap only in FULL graphs; eager warmup initializes Mega mHC.
+            mhc_stream = None
+        mhc_pre = (
+            partial(mhc_pre_delayed_overlap, stream=mhc_stream)
+            if mhc_stream is not None
+            else mhc_pre_delayed_tilelang
+        )
         # The reference collapses each sublayer's input with the *previous*
         # sublayer's pre-mix: attention uses the pre-mix carried in (identity
         # for the first layer), the FFN uses this layer's attention pre-mix.
@@ -303,7 +432,7 @@ def forward(
                 # copies and the identity pre-mix selects copy 0.
                 assert self.hc_attn_fn_broadcast is not None
                 residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous()
-                post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(
+                post_mix, res_mix, x, attn_pre = mhc_pre(
                     residual,
                     self.hc_attn_fn_broadcast,
                     self.hc_attn_scale,
@@ -319,7 +448,7 @@ def forward(
                 )
             else:
                 residual = x
-                post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(
+                post_mix, res_mix, x, attn_pre = mhc_pre(
                     residual,
                     self.hc_attn_fn,
                     self.hc_attn_scale,
@@ -333,19 +462,43 @@ def forward(
                     norm_weight=self.attn_norm.weight,
                     norm_eps=self.attn_norm.variance_epsilon,
                 )
+        elif self.engram is not None and engram_hashes is not None:
+            # Engram injection happens between the previous sublayer's post
+            # and this block's pre, on the full hc stream, so the mix
+            # coefficients see the injected stream. The injection also keeps
+            # the post out of the pre-norm GEMM's fused prologue.
+            if self.fuse_mhc_all_reduce:
+                x = tensor_model_parallel_all_reduce(x)
+            previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix)
+            if capture_previous_aux:
+                previous_aux = previous_post.mean(dim=1)
+            residual = self.engram(
+                previous_post,
+                engram_hashes[:, self.engram.layer_hash_index],
+                engram_mask,
+            )
+            post_mix, res_mix, x, attn_pre = mhc_pre(
+                residual,
+                self.hc_attn_fn,
+                self.hc_attn_scale,
+                self.hc_attn_base,
+                self.rms_norm_eps,
+                self.hc_eps,
+                self.hc_eps,
+                self.hc_post_alpha,
+                self.hc_sinkhorn_iters,
+                pre_mix=pre_mix,
+                norm_weight=self.attn_norm.weight,
+                norm_eps=self.attn_norm.variance_epsilon,
+            )
         else:
-            residual = mhc_post_tilelang(x, residual, post_mix, res_mix)
-            if self.engram is not None and engram_hashes is not None:
-                # Engram injection happens between the previous sublayer's
-                # post and this block's pre, on the full hc stream, so the
-                # mix coefficients see the injected stream.
-                residual = self.engram(
-                    residual,
-                    engram_hashes[:, self.engram.layer_hash_index],
-                    engram_mask,
-                )
-            post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(
+            # The collapse already reads the post-mapped streams, so the mean
+            # aux consumers want comes out of the same kernel.
+            residual, post_mix, res_mix, x, attn_pre, aux = mhc_shifted_post_pre(
+                x,
                 residual,
+                post_mix,
+                res_mix,
                 self.hc_attn_fn,
                 self.hc_attn_scale,
                 self.hc_attn_base,
@@ -357,18 +510,29 @@ def forward(
                 pre_mix=pre_mix,
                 norm_weight=self.attn_norm.weight,
                 norm_eps=self.attn_norm.variance_epsilon,
+                capture_aux=capture_previous_aux,
+                stream=mhc_stream,
+                reduce_results=self.fuse_mhc_all_reduce,
             )
+            if capture_previous_aux:
+                previous_aux = aux
 
         if self.use_sequence_parallel:
             x = sp_all_gather(x)[: positions.shape[0]]
 
         x = self.attn(positions, x, None)
-        if self.use_sequence_parallel:
+        # With GEMM-RS bound, the attention output is already this rank's
+        # sequence-parallel shard (see DeepseekV4Attention._wo_b_proj).
+        if self.use_sequence_parallel and self.attn.gemm_rs is None:
             x = sp_reduce_scatter(x)
 
-        residual = mhc_post_tilelang(x, residual, post_mix, res_mix)
-        post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_tilelang(
+        if mhc_stream is not None:
+            torch.cuda.current_stream().wait_stream(mhc_stream)
+        residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre(
+            x,
             residual,
+            post_mix,
+            res_mix,
             self.hc_ffn_fn,
             self.hc_ffn_scale,
             self.hc_ffn_base,
@@ -380,9 +544,13 @@ def forward(
             pre_mix=attn_pre,
             norm_weight=self.ffn_norm.weight,
             norm_eps=self.ffn_norm.variance_epsilon,
+            stream=mhc_stream,
+            reduce_results=self.fuse_mhc_all_reduce,
         )
-        x = self.ffn(x, input_ids)
-        return x, residual, post_mix, res_mix, ffn_pre
+        x = self.ffn(x, input_ids, mega_gate_metadata)
+        if mhc_stream is not None:
+            torch.cuda.current_stream().wait_stream(mhc_stream)
+        return x, residual, post_mix, res_mix, ffn_pre, previous_aux
 
 
 class DeepseekV4Model(nn.Module, EagleModelMixin):
@@ -413,6 +581,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         # (compressor kv_score, indexer.weights_proj). fused_wqa_wkv stays on
         # the default stream.
         aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
+        # Keep mHC independent of the streams used inside attention.
+        mhc_stream = torch.cuda.Stream() if supports_mhc_overlap(vllm_config) else None
+        self.fuse_mhc_all_reduce = mhc_stream is not None and supports_mhc_all_reduce(
+            vllm_config
+        )
 
         # Reserved topk indices buffer for all Indexer layers to reuse.
         self.topk_indices_buffer = torch.empty(
@@ -447,6 +620,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             self.embed_tokens = PPMissingLayer()
 
         self.engram_layout = EngramLayout.from_config(config)
+        engram_config = vllm_config.engram_config
+        if (
+            self.engram_layout is not None
+            and engram_config is not None
+            and engram_config.dp_shared_memory
+        ):
+            engram_config.dp_shared_memory = can_share_engram_tables(self.engram_layout)
+
+        # GEMM-RS uses NCCL symmetric-memory multicast, which requires all TP
+        # ranks to belong to one NVLink domain. Collective: run before layers.
+        self.run_gemm_rs = maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel)
 
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
@@ -457,6 +641,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 aux_stream_list=aux_stream_list,
                 candidate_block_buffer=self.candidate_block_buffer,
                 engram_layout=self.engram_layout,
+                run_gemm_rs=self.run_gemm_rs,
+                mhc_stream=mhc_stream,
+                fuse_mhc_all_reduce=self.fuse_mhc_all_reduce,
             ),
             prefix=f"{prefix}.layers",
         )
@@ -466,6 +653,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         # layer's sliding-window KV cache. Only PP ranks owning an engram
         # layer need it.
         self.engram_hash: NgramHashState | None = None
+        self.engram_dp_shared_memory = bool(
+            vllm_config.engram_config and vllm_config.engram_config.dp_shared_memory
+        )
         self.engram_swa_prefix: str | None = None
         if self.engram_layout is not None:
             local_engram = any(
@@ -594,12 +784,21 @@ def forward(
                     swa_metadata.slot_mapping,
                     swa_metadata.block_table,
                 )
+            elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1:
+                # DP-sharded lookups are collective, so a replica skipping the
+                # hash still has to reach them.
+                engram_hashes, engram_mask = self.engram_hash.dummy_hashes(input_ids)
+            if engram_hashes is not None:
                 # Gather all Engram rows before entering the decoder layers.
+                # One gather feeds every layer sharing the DP-split table.
+                gathered_hashes = gather_engram_hashes(
+                    engram_hashes, dp_shared_memory=self.engram_dp_shared_memory
+                )
                 for layer in islice(self.layers, self.start_layer, self.end_layer):
                     engram = getattr(layer, "engram", None)
                     if engram is not None:
                         engram.prepare_embeddings(
-                            engram_hashes[:, engram.layer_hash_index]
+                            gathered_hashes[:, engram.layer_hash_index]
                         )
 
         full_num_tokens = positions.shape[0]
@@ -612,18 +811,29 @@ def forward(
             hidden_states = sp_shard(hidden_states)
             input_ids = sp_shard(input_ids)
 
+        mega_gate_metadata = None
+        if self.use_mega_moe:
+            mega_gate_metadata = prepare_mega_gate_routing_metadata(
+                input_ids,
+                has_hash_routing=False,
+                image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
+                if getattr(self.config, "vision_n_layers", 0) > 0
+                else None,
+            )
+
         residual, post_mix, res_mix = None, None, None
         pre_mix: torch.Tensor | None = None
         if not get_pp_group().is_first_rank:
             assert intermediate_tensors is not None
             pre_mix = intermediate_tensors["pre_mix"]
-        aux_hidden_states: list[torch.Tensor] = []
-        final_aux_recon: torch.Tensor | None = None  # avoid duplicate mhc_post call
+        # Every layer's post runs inside the next layer's fused pre, so aux
+        # hidden states are read back from there instead of recomputed.
+        aux_hidden_by_layer: dict[int, torch.Tensor] = {}
         for idx, layer in enumerate(
             islice(self.layers, self.start_layer, self.end_layer),
             start=self.start_layer,
         ):
-            hidden_states, residual, post_mix, res_mix, pre_mix = layer(
+            hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = layer(
                 hidden_states,
                 positions,
                 input_ids,
@@ -633,25 +843,32 @@ def forward(
                 residual,
                 engram_hashes,
                 engram_mask,
+                capture_previous_aux=idx in self.aux_hidden_state_layers,
+                mega_gate_metadata=mega_gate_metadata,
             )
-            if idx + 1 in self.aux_hidden_state_layers:
-                # Reconstruct the aux hidden state for draft models
-                aux_recon = mhc_post_tilelang(
-                    hidden_states, residual, post_mix, res_mix
-                )
-                aux_hidden_state = aux_recon.mean(dim=1)
+            if previous_aux is not None:
+                # idx is the one-based id of the layer whose post this is.
                 if self.use_sequence_parallel:
-                    aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
-                aux_hidden_states.append(aux_hidden_state)
-                final_aux_recon = aux_recon
+                    previous_aux = sp_all_gather(previous_aux)[:full_num_tokens]
+                aux_hidden_by_layer[idx] = previous_aux
         if layer is not None:
-            # Reuse if the last layer was captured as an aux hidden state
+            # The last layer has no successor to fold its post into.
+            if self.fuse_mhc_all_reduce:
+                hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+            hidden_states = mhc_post_tilelang(
+                hidden_states, residual, post_mix, res_mix
+            )
             if self.end_layer in self.aux_hidden_state_layers:
-                hidden_states = final_aux_recon
-            else:
-                hidden_states = mhc_post_tilelang(
-                    hidden_states, residual, post_mix, res_mix
-                )
+                final_aux = hidden_states.mean(dim=1)
+                if self.use_sequence_parallel:
+                    final_aux = sp_all_gather(final_aux)[:full_num_tokens]
+                aux_hidden_by_layer[self.end_layer] = final_aux
+
+        aux_hidden_states = [
+            aux_hidden_by_layer[layer_id]
+            for layer_id in self.aux_hidden_state_layers
+            if layer_id in aux_hidden_by_layer
+        ]
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors(
@@ -852,6 +1069,17 @@ def finalize_mega_moe_weights(self) -> None:
         for layer in islice(self.layers, self.start_layer, self.end_layer):
             layer.ffn.finalize_mega_moe_weights()
 
+    def finalize_mega_attn_weights(self) -> None:
+        """Permute wq_b / wo_a into FlashMLA's mega-attention layouts.
+
+        A no-op for every other attention layer, and idempotent, so a second
+        post-load pass cannot permute twice.
+        """
+        for layer in islice(self.layers, self.start_layer, self.end_layer):
+            finalize = getattr(layer.attn, "finalize_loaded_weights", None)
+            if finalize is not None:
+                finalize()
+
     def finalize_mhc_broadcast_weights(self) -> None:
         if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer:
             return
@@ -1106,6 +1334,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
     def process_weights_after_loading(self) -> None:
         self.model.finalize_mega_moe_weights()
         self.model.finalize_mhc_broadcast_weights()
+        self.model.finalize_mega_attn_weights()
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         return self.model.get_expert_mapping()
diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py
new file mode 100644
index 000000000000..af20976dedbf
--- /dev/null
+++ b/vllm/models/deepseek_v41/nvidia/model_state.py
@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
+from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
+from vllm.v1.worker.gpu.states import RequestState
+from vllm.v1.worker.utils import AttentionGroup
+
+
+@triton.jit
+def _gather_lookback_kernel(
+    lookback_ptr,
+    idx_mapping_ptr,
+    num_computed_tokens_ptr,
+    all_token_ids_ptr,
+    all_token_ids_stride,
+    num_reqs,
+    DEPTH: tl.constexpr,
+    BLOCK_DEPTH: tl.constexpr,
+):
+    # One program per lookback row; rows past the batch are filled with -1.
+    batch_idx = tl.program_id(0)
+    in_batch = batch_idx < num_reqs
+    req_state_idx = tl.load(idx_mapping_ptr + batch_idx, mask=in_batch, other=0)
+    num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
+
+    offs = tl.arange(0, BLOCK_DEPTH)
+    pos = num_computed - 1 - offs
+    valid = in_batch & (offs < DEPTH) & (pos >= 0)
+    ids = tl.load(
+        all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos,
+        mask=valid,
+        other=-1,
+    )
+    tl.store(lookback_ptr + batch_idx * DEPTH + offs, ids, mask=offs < DEPTH)
+
+
+@triton.jit
+def _pad_replayed_slots_kernel(
+    slot_mappings_ptr,  # [num_groups, num_tokens_padded]
+    group_stride,
+    cacheable_groups_ptr,  # [NUM_GROUPS] indices of the prefix-cacheable groups
+    query_start_loc_ptr,  # [num_reqs + 1]
+    positions_ptr,  # [num_tokens]
+    replay_start_ptr,  # [num_reqs]
+    window,
+    pad_slot_id,
+    NUM_GROUPS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    req = tl.program_id(0)
+    start = tl.load(replay_start_ptr + req)
+    if start == 0:
+        return
+    begin = tl.load(query_start_loc_ptr + req)
+    end = tl.load(query_start_loc_ptr + req + 1)
+    for tok in range(begin, end, BLOCK):
+        offs = tok + tl.arange(0, BLOCK)
+        pos = tl.load(positions_ptr + offs, mask=offs < end, other=0)
+        replayed = (offs < end) & (pos >= start) & (pos < start + window)
+        for i in tl.static_range(NUM_GROUPS):
+            group = tl.load(cacheable_groups_ptr + i)
+            tl.store(
+                slot_mappings_ptr + group * group_stride + offs,
+                pad_slot_id,
+                mask=replayed,
+            )
+
+
+class ReplayAttnMetadata(ModelSpecificAttnMetadata):
+    """Hands the batch's replay starts to the sliding-window builders."""
+
+    def __init__(self, replay_start: torch.Tensor) -> None:
+        self.replay_start = replay_start
+
+    def get_extra_attn_kwargs(
+        self, attn_metadata_builder: Any, num_reqs: int
+    ) -> dict[str, Any]:
+        if isinstance(attn_metadata_builder, DeepseekSparseSWAMetadataBuilder):
+            return {"replay_start": self.replay_start[:num_reqs]}
+        return {}
+
+
+class DeepseekV41ModelState(DefaultModelState):
+    """DefaultModelState plus the engram lookback window and SWA bounded replay.
+
+    The engram n-gram hash needs the ids of the ``depth`` tokens preceding
+    each request's chunk start (see ``common/engram.py``). The runner keeps
+    the full token history on device, so the window is gathered there every
+    step: exact for prompt and generated tokens alike, whatever instance
+    produced their KV.
+
+    SWA bounded replay (``CacheConfig.swa_bounded_replay``) keeps the
+    sliding-window KV out of prefix caching and rebuilds it after a prefix hit
+    by recomputing the hit's last window; the scheduler tells each request the
+    position from which it holds window KV (``NewRequestData.replay_start``).
+    ``prepare_attn`` gathers those per batch, pads the replayed tokens' slots
+    in the prefix-cacheable groups so the cached KV stays as is, and hands the
+    starts to the sliding-window metadata builders, whose kernels read no
+    window KV below them.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        model: nn.Module,
+        encoder_cache: EncoderCache | None,
+        device: torch.device,
+    ):
+        super().__init__(vllm_config, model, encoder_cache, device)
+        depth = model.token_lookback_depth
+        self.lookback_token_ids: torch.Tensor | None = None
+        if depth > 0:
+            # Persistent so a captured graph can read it on replay.
+            self.lookback_token_ids = torch.full(
+                (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device
+            )
+
+        # Per request state index; batches gather from it (see prepare_attn).
+        self._replay_start_np = np.zeros(self.max_num_reqs, dtype=np.int32)
+        self._replay_start = torch.zeros(
+            self.max_num_reqs, dtype=torch.int32, device=device
+        )
+        self._replay_start_staging = UvaBufferPool(self.max_num_reqs, torch.int32)
+        # The replay window and the indices of the prefix-cacheable groups, from
+        # the KV cache config on first use; None until then.
+        self._replay: tuple[int, torch.Tensor] | None = None
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        super().add_request(req_index, new_req_data)
+        self._replay_start_np[req_index] = new_req_data.replay_start
+
+    def prepare_inputs(
+        self, input_batch: InputBatch, req_states: RequestState
+    ) -> dict[str, torch.Tensor | None]:
+        model_inputs = super().prepare_inputs(input_batch, req_states)
+        window = self.lookback_token_ids
+        if window is None:
+            return model_inputs
+        all_token_ids = req_states.all_token_ids.gpu
+        depth = window.shape[1]
+        _gather_lookback_kernel[(window.shape[0],)](
+            window,
+            input_batch.idx_mapping,
+            req_states.num_computed_tokens.gpu,
+            all_token_ids,
+            all_token_ids.stride(0),
+            input_batch.idx_mapping.shape[0],
+            DEPTH=depth,
+            BLOCK_DEPTH=triton.next_power_of_2(depth),
+        )
+        model_inputs["lookback_token_ids"] = window
+        return model_inputs
+
+    def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
+        model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens)
+        if self.lookback_token_ids is not None:
+            # The captured graph reads this buffer; replays refill it in place.
+            self.lookback_token_ids.fill_(-1)
+            model_inputs["lookback_token_ids"] = self.lookback_token_ids
+        return model_inputs
+
+    def prepare_attn(
+        self,
+        input_batch: InputBatch,
+        cudagraph_mode: CUDAGraphMode,
+        block_tables: tuple[torch.Tensor, ...],
+        slot_mappings: torch.Tensor,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+        for_capture: bool = False,
+        ubatch_idx: int = 0,
+        model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
+    ) -> dict[str, Any]:
+        if self._replay is None:
+            specs = [group.kv_cache_spec for group in kv_cache_config.kv_cache_groups]
+            self._replay = (
+                max(spec.prefix_replay_tokens for spec in specs),
+                torch.tensor(
+                    [i for i, spec in enumerate(specs) if spec.prefix_cacheable],
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+            )
+        window, cacheable_groups = self._replay
+        if window:
+            num_reqs = input_batch.num_reqs
+            # Decode rows sit above the hit, so only prefills carry a replay
+            # start; dummy batches (captures, profiling) carry none.
+            replay_start_np = np.where(
+                input_batch.is_prefilling_np[:num_reqs],
+                self._replay_start_np[input_batch.idx_mapping_np[:num_reqs]],
+                0,
+            ).astype(np.int32)
+            replay_start = self._replay_start_staging.copy_to_gpu(
+                replay_start_np, out=self._replay_start[:num_reqs]
+            )
+            if replay_start_np.any():
+                # The replayed tokens rebuild window KV only: their slots in the
+                # prefix-cacheable groups are padded so the cached KV stays as is.
+                _pad_replayed_slots_kernel[(num_reqs,)](
+                    slot_mappings,
+                    slot_mappings.stride(0),
+                    cacheable_groups,
+                    input_batch.query_start_loc,
+                    input_batch.positions,
+                    replay_start,
+                    window,
+                    PAD_SLOT_ID,
+                    NUM_GROUPS=cacheable_groups.numel(),
+                    BLOCK=1024,
+                )
+            assert model_specific_attn_metadata is None
+            model_specific_attn_metadata = ReplayAttnMetadata(replay_start)
+        return super().prepare_attn(
+            input_batch,
+            cudagraph_mode,
+            block_tables,
+            slot_mappings,
+            attn_groups,
+            kv_cache_config,
+            for_capture=for_capture,
+            ubatch_idx=ubatch_idx,
+            model_specific_attn_metadata=model_specific_attn_metadata,
+        )
diff --git a/vllm/models/deepseek_v41/nvidia/ops/__init__.py b/vllm/models/deepseek_v41/nvidia/ops/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/deepseek_v41/nvidia/ops/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py
new file mode 100644
index 000000000000..bdb5b95bbeb9
--- /dev/null
+++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py
@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import functools
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import (
+    _import_deep_gemm,
+    is_deep_gemm_supported,
+    mega_mhc,
+)
+
+
+@functools.cache
+def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool:
+    """Return whether DeepGEMM can run the DSV4.1 shifted mHC kernel."""
+    if not (
+        is_deep_gemm_supported()
+        and current_platform.is_device_capability_family(100)
+        and hidden_size > 0
+        and hidden_size % 1024 == 0
+        and hc_mult == 4
+    ):
+        return False
+    deep_gemm = _import_deep_gemm()
+    return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None))
+
+
+def can_use_mega_mhc(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    pre_mix: torch.Tensor | None,
+    norm_weight: torch.Tensor | None,
+    capture_aux: bool,
+) -> bool:
+    return (
+        pre_mix is not None
+        and norm_weight is not None
+        and not capture_aux
+        and x.shape[0] <= 1 << 20
+        and is_mega_mhc_supported(x.shape[1], residual.shape[1])
+    )
+
+
+def mhc_shifted_post_pre_deep_gemm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    shifted_prev_mix: torch.Tensor,
+    post_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    mix_scales: torch.Tensor,
+    mix_bases: torch.Tensor,
+    hc_norm_eps: float,
+    hc_pre_eps: float,
+    hc_post_scale: float,
+    sinkhorn_eps: float,
+    num_sinkhorn_iters: int,
+    rmsnorm_weight: torch.Tensor,
+    rmsnorm_eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Run DSV4.1 shifted post, next pre, and BF16 RMSNorm with Mega mHC."""
+    num_tokens, hidden_size = x.shape
+    hc_mult = residual.shape[1]
+    new_residual = torch.empty_like(residual)
+    new_prev_mix = x.new_empty(num_tokens, hc_mult, 1, dtype=torch.float32)
+    new_post_mix = torch.empty_like(post_mix)
+    new_comb_res_mix = torch.empty_like(comb_res_mix)
+    y_bf16 = x.new_empty(num_tokens, hidden_size, dtype=torch.bfloat16)
+    mega_mhc(
+        x=x,
+        residual=residual,
+        shifted_prev_mix=shifted_prev_mix.unsqueeze(-1),
+        post_mix=post_mix,
+        comb_res_mix=comb_res_mix,
+        fn=fn,
+        mix_scales=mix_scales,
+        mix_bases=mix_bases,
+        hc_mult=hc_mult,
+        hc_norm_eps=hc_norm_eps,
+        hc_pre_eps=hc_pre_eps,
+        hc_post_scale=hc_post_scale,
+        sinkhorn_eps=sinkhorn_eps,
+        num_sinkhorn_iters=num_sinkhorn_iters,
+        rmsnorm_weight=rmsnorm_weight,
+        rmsnorm_eps=rmsnorm_eps,
+        rmsnorm_scale=1.0,
+        new_residual=new_residual,
+        new_prev_mix=new_prev_mix,
+        new_post_mix=new_post_mix,
+        new_comb_res_mix=new_comb_res_mix,
+        y_bf16=y_bf16,
+    )
+    return (
+        new_residual,
+        new_post_mix,
+        new_comb_res_mix,
+        y_bf16,
+        new_prev_mix.squeeze(-1),
+    )
diff --git a/vllm/models/deepseek_v41/nvidia/ops/mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mhc.py
new file mode 100644
index 000000000000..88288ead4d9d
--- /dev/null
+++ b/vllm/models/deepseek_v41/nvidia/ops/mhc.py
@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dispatch DSV4.1 mHC operations and overlap coefficient generation."""
+
+from functools import partial
+from typing import TYPE_CHECKING, cast
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed import get_tp_group
+from vllm.model_executor.kernels.mhc.tilelang import (
+    mhc_fused_post_pre_delayed_tilelang,
+    mhc_post_tilelang,
+)
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+from .mega_mhc import (
+    can_use_mega_mhc,
+    mhc_shifted_post_pre_deep_gemm,
+)
+
+if TYPE_CHECKING:
+    from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
+
+# GB200 TP4/FlashInfer improves through 16 tokens; larger screens tie or regress.
+MHC_OVERLAP_MAX_TOKENS = 16
+
+
+def supports_mhc_overlap(vllm_config: VllmConfig) -> bool:
+    """Check kernel requirements and safety of sharing the coefficient stream."""
+    config = vllm_config.model_config.hf_config
+    # DeepGEMM's prenorm kernel requires K % 64 == 0, N % 8 == 0, and N <= 32.
+    mix_size = config.hc_mult * (config.hc_mult + 2)
+    return (
+        current_platform.is_device_capability_family(100)
+        and is_deep_gemm_supported()
+        and config.hidden_size % 64 == 0
+        and 0 < mix_size <= 32
+        and mix_size % 8 == 0
+        and not vllm_config.parallel_config.use_ubatching
+    )
+
+
+def supports_mhc_all_reduce(vllm_config: VllmConfig) -> bool:
+    parallel = vllm_config.parallel_config
+    config = vllm_config.model_config.hf_config
+    if (
+        parallel.tensor_parallel_size != 4
+        or parallel.enable_expert_parallel
+        or config.hidden_size != 5120
+        or config.hc_mult != 4
+    ):
+        return False
+    comm = cast("CudaCommunicator", get_tp_group().device_communicator).ca_comm
+    return comm is not None and bool(comm.mnnvl_lamport_ag_multicast_ptr)
+
+
+def mhc_pre_delayed_overlap(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    pre_mix: torch.Tensor | None = None,
+    x: torch.Tensor | None = None,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+    *,
+    stream: torch.cuda.Stream,
+    layer_input: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Prepare the input on the caller stream and coefficients on another stream.
+
+    Returns post mix, residual mix, normalized input, and next pre-mix. Only the
+    input is ready on the caller stream; join stream before using coefficients.
+    The caller must retain the weights and coefficient outputs until the join.
+    """
+    from vllm.model_executor.kernels.mhc.warmup import (
+        MHC_PRE_NORM_KERNEL,
+        compute_mhc_pre_num_splits,
+    )
+    from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
+
+    n, hc, hidden = residual.shape
+    assert residual.is_contiguous() and residual.dtype == torch.bfloat16
+    assert norm_weight is not None
+    if x is None:
+        x = residual.view(n, hc * hidden)
+    assert x.is_contiguous()
+    post = torch.empty((n, hc), device=residual.device, dtype=torch.float32)
+    comb = torch.empty((n, hc * hc), device=residual.device, dtype=torch.float32)
+    next_pre = torch.empty_like(post)
+    compute_input = layer_input is None
+    if layer_input is None:
+        layer_input = torch.empty(
+            (n, hidden), device=residual.device, dtype=torch.bfloat16
+        )
+    outputs = post.unsqueeze(-1), comb.view(n, hc, hc), layer_input, next_pre
+    if n == 0:
+        return outputs
+    splits = compute_mhc_pre_num_splits(x.shape[1], n)
+    mix = torch.empty(
+        (splits, n, hc * (hc + 2)), device=residual.device, dtype=torch.float32
+    )
+    sqr = torch.empty((splits, n), device=residual.device, dtype=torch.float32)
+    epilogue = partial(
+        MHC_PRE_NORM_KERNEL,
+        mix,
+        sqr,
+        hc_scale,
+        hc_base,
+        residual,
+        post,
+        comb,
+        layer_input,
+        norm_weight,
+        pre_mix if pre_mix is not None else post,
+        next_pre,
+        layer_input,  # Unused aux output in split modes.
+        hidden_size=hidden,
+        rms_eps=rms_eps,
+        hc_pre_eps=hc_pre_eps,
+        hc_sinkhorn_eps=hc_sinkhorn_eps,
+        hc_post_mult_value=hc_post_mult_value,
+        sinkhorn_repeat=sinkhorn_repeat,
+        norm_eps=norm_eps,
+        hc_mult=hc,
+        use_pre_mix_in=pre_mix is not None,
+        save_pre_mix=True,
+        rms_numel=x.shape[1],
+    )
+    main = torch.cuda.current_stream()
+    # Prioritize input readiness for small batches before releasing statistics.
+    if compute_input and n <= 8:
+        epilogue(split_mode="input")
+    stream.wait_stream(main)
+    with torch.cuda.stream(stream):
+        tf32_hc_prenorm_gemm(x, fn, mix, sqr, splits)
+        epilogue(split_mode="stats")
+    for tensor in (x, mix, sqr):
+        tensor.record_stream(stream)
+    if compute_input and n > 8:
+        epilogue(split_mode="input")
+    return outputs
+
+
+def mhc_shifted_post_pre(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    pre_mix: torch.Tensor | None = None,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+    capture_aux: bool = False,
+    *,
+    stream: torch.cuda.Stream | None = None,
+    reduce_results: bool = False,
+) -> tuple[
+    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+]:
+    """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang.
+
+    When stream is supplied, join it before consuming the returned coefficients.
+    """
+    layer_input = None
+    if reduce_results:
+        tp = get_tp_group()
+        if stream is not None and 0 < x.shape[0] <= MHC_OVERLAP_MAX_TOKENS:
+            comm = cast("CudaCommunicator", tp.device_communicator).ca_comm
+            assert comm is not None and comm.mnnvl_lamport_epochs is not None
+            output = torch.empty_like(residual)
+            layer_input = torch.empty_like(x)
+            torch.ops._C_custom_ar.all_reduce_mhc(
+                x,
+                residual,
+                post_layer_mix,
+                comb_res_mix,
+                pre_mix,
+                norm_weight,
+                output,
+                layer_input,
+                comm.mnnvl_lamport_ag_local_ptr,
+                comm.mnnvl_lamport_ag_multicast_ptr,
+                comm.mnnvl_lamport_epochs[0],
+                comm.rank,
+                comm.mnnvl_buffer_size,
+                norm_eps,
+            )
+            residual = output
+        else:
+            x = tp.all_reduce(x)
+    if stream is not None:
+        if layer_input is None:
+            residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix)
+        aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1])
+        pre_outputs = mhc_pre_delayed_overlap(
+            residual,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            pre_mix=pre_mix,
+            norm_weight=norm_weight,
+            norm_eps=norm_eps,
+            stream=stream,
+            layer_input=layer_input,
+        )
+        return residual, *pre_outputs, aux
+
+    if can_use_mega_mhc(x, residual, pre_mix, norm_weight, capture_aux):
+        assert pre_mix is not None and norm_weight is not None
+        outputs = mhc_shifted_post_pre_deep_gemm(
+            x,
+            residual,
+            pre_mix,
+            post_layer_mix,
+            comb_res_mix,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_pre_eps,
+            hc_post_mult_value,
+            hc_sinkhorn_eps,
+            sinkhorn_repeat,
+            norm_weight,
+            norm_eps,
+        )
+        return *outputs, x.new_empty(0, x.shape[1])
+
+    return mhc_fused_post_pre_delayed_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_mult_value,
+        sinkhorn_repeat,
+        pre_mix=pre_mix,
+        norm_weight=norm_weight,
+        norm_eps=norm_eps,
+        capture_aux=capture_aux,
+    )
diff --git a/vllm/models/deepseek_v4_1/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py
similarity index 89%
rename from vllm/models/deepseek_v4_1/nvidia/vl_model.py
rename to vllm/models/deepseek_v41/nvidia/vl_model.py
index 111936b43115..901985bdd7fb 100644
--- a/vllm/models/deepseek_v4_1/nvidia/vl_model.py
+++ b/vllm/models/deepseek_v41/nvidia/vl_model.py
@@ -26,6 +26,7 @@
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
     SupportsEagle3,
+    SupportsEncoderCudaGraph,
     SupportsMultiModal,
     SupportsPP,
 )
@@ -44,17 +45,12 @@
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
 from ..common.mm_preprocess import (
-    IMAGE,
-    IMAGE_END,
-    IMAGE_NEW_LINE,
-    IMAGE_PAD_ID,
     IMAGE_PLACEHOLDER,
-    IMAGE_SENTINEL_BASE_ID,
-    IMAGE_START,
     DeepseekV4VLDummyInputsBuilder,
     DeepseekV4VLMultiModalProcessor,
     DeepseekV4VLProcessingInfo,
 )
+from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin
 from .model import (
     DeepseekV41LLMForCausalLM,
     _linear_scale_param_name,
@@ -117,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper(
     info=DeepseekV4VLProcessingInfo,
     dummy_inputs=DeepseekV4VLDummyInputsBuilder,
 )
-class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3):
+class DeepseekV41ForCausalLM(
+    nn.Module,
+    DeepseekV4VLEncoderCudaGraphMixin,
+    SupportsMultiModal,
+    SupportsEncoderCudaGraph,
+    SupportsPP,
+    SupportsEagle3,
+):
     """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower.
 
     ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
     delegates through ``language_model`` via the protocol defaults.
+    ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is
+    implemented by ``DeepseekV4VLEncoderCudaGraphMixin``.
     """
 
     supports_encoder_tp_data = True
@@ -198,29 +203,6 @@ def _parse_and_validate_image_input(
             resolve_bindings={"p": self.config.vision_patch_size},
         )
 
-    def _encode_image(
-        self,
-        patches: torch.Tensor,
-        n_vit_h: int,
-        n_vit_w: int,
-    ) -> torch.Tensor:
-        # Aligner rows in reading order, one per IMAGE slot.
-        return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
-
-    def _build_image_span(
-        self, image_embeds: torch.Tensor, types: torch.Tensor
-    ) -> torch.Tensor:
-        """Full image span: aligner rows at IMAGE slots, the learned
-        delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""
-        types = types.to(image_embeds.device)
-        span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])
-        dtype = image_embeds.dtype
-        span[types == IMAGE_START] = self.image_start.to(dtype)
-        span[types == IMAGE_END] = self.image_end.to(dtype)
-        span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
-        span[types == IMAGE] = image_embeds
-        return span
-
     def _process_image_input(
         self,
         image_input: DeepseekV4VLImagePixelInputs,
@@ -277,12 +259,6 @@ def embed_input_ids(
     ) -> torch.Tensor:
         from vllm.model_executor.models.utils import _merge_multimodal_embeddings
 
-        # Compressor-alignment pads borrow a reserved id; embed them as the
-        # plain image token (the checkpoint has no image_pad vector).
-        # Branch-free: safe inside compiled/captured regions.
-        input_ids = input_ids.masked_fill(
-            input_ids == IMAGE_PAD_ID, IMAGE_SENTINEL_BASE_ID
-        )
         inputs_embeds = self.language_model.embed_input_ids(input_ids)
 
         if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
diff --git a/vllm/models/deepseek_v4_1/quant_config.py b/vllm/models/deepseek_v41/quant_config.py
similarity index 100%
rename from vllm/models/deepseek_v4_1/quant_config.py
rename to vllm/models/deepseek_v41/quant_config.py
diff --git a/vllm/models/deepseek_v4_1/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py
similarity index 81%
rename from vllm/models/deepseek_v4_1/sparse_mla.py
rename to vllm/models/deepseek_v41/sparse_mla.py
index 6a0f09f952e4..1130eec0903b 100644
--- a/vllm/models/deepseek_v4_1/sparse_mla.py
+++ b/vllm/models/deepseek_v41/sparse_mla.py
@@ -86,7 +86,7 @@ class DeepseekV4SparseMLABackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64 if current_platform.is_device_capability_family(90) else 128]
 
     @staticmethod
@@ -195,6 +195,7 @@ def build(
         if self.compress_ratio > 1:
             slot_mapping = get_compressed_slot_mapping(
                 cm.num_actual_tokens,
+                cm.slot_mapping,
                 cm.query_start_loc,
                 cm.seq_lens,
                 cm.block_table_tensor.clamp_(min=0),
@@ -229,3 +230,54 @@ def get_name() -> str:
     @staticmethod
     def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]:
         return DeepseekV4FlashMLAMetadataBuilder
+
+
+class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend):
+    """FlashMLA's mega-attention kernel: Q RoPE + sparse attention + inverse
+    RoPE + FP8 cast of the output, in one launch.
+
+    Same metadata and KV-cache geometry as ``FLASHMLA_SPARSE_DSV41`` -- what
+    differs is the attention layer's interface contract (it takes an unnormed,
+    unroped Q and returns an already-inverse-RoPE'd, quantized output) and the
+    extra ``nvfp4_ds_mla`` compressed-cache record only this kernel can read.
+    SM100 only; the kernel has no SM90 instantiation.
+    """
+
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "fp8_ds_mla",
+        "fp8",  # alias for fp8_ds_mla
+        "nvfp4_ds_mla",  # V4.1 fp8 SWA cache + NVFP4 compressed cache
+    ]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHMLA_MEGA_ATTN_DSV41"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
+        return [128]
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major == 10
+
+    @classmethod
+    def supports_combination(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        device_capability: DeviceCapability,
+    ) -> str | None:
+        # Imported here: the layer module imports this backend class.
+        from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
+            is_flashmla_mega_attn_supported,
+        )
+
+        return is_flashmla_mega_attn_supported()[1]
diff --git a/vllm/models/deepseek_v4_1/nvidia/model_state.py b/vllm/models/deepseek_v4_1/nvidia/model_state.py
deleted file mode 100644
index 1ba0f4160333..000000000000
--- a/vllm/models/deepseek_v4_1/nvidia/model_state.py
+++ /dev/null
@@ -1,98 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-from vllm.config import VllmConfig
-from vllm.triton_utils import tl, triton
-from vllm.v1.worker.gpu.input_batch import InputBatch
-from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
-from vllm.v1.worker.gpu.model_states.default import DefaultModelState
-from vllm.v1.worker.gpu.states import RequestState
-
-
-@triton.jit
-def _gather_lookback_kernel(
-    lookback_ptr,
-    idx_mapping_ptr,
-    num_computed_tokens_ptr,
-    all_token_ids_ptr,
-    all_token_ids_stride,
-    num_reqs,
-    DEPTH: tl.constexpr,
-    BLOCK_DEPTH: tl.constexpr,
-):
-    # One program per lookback row; rows past the batch are filled with -1.
-    batch_idx = tl.program_id(0)
-    in_batch = batch_idx < num_reqs
-    req_state_idx = tl.load(idx_mapping_ptr + batch_idx, mask=in_batch, other=0)
-    num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
-
-    offs = tl.arange(0, BLOCK_DEPTH)
-    pos = num_computed - 1 - offs
-    valid = in_batch & (offs < DEPTH) & (pos >= 0)
-    ids = tl.load(
-        all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos,
-        mask=valid,
-        other=-1,
-    )
-    tl.store(lookback_ptr + batch_idx * DEPTH + offs, ids, mask=offs < DEPTH)
-
-
-class DeepseekV41ModelState(DefaultModelState):
-    """DefaultModelState plus the engram lookback window.
-
-    The engram n-gram hash needs the ids of the ``depth`` tokens preceding
-    each request's chunk start (see ``common/engram.py``). The runner keeps
-    the full token history on device, so the window is gathered there every
-    step: exact for prompt and generated tokens alike, whatever instance
-    produced their KV.
-    """
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        model: nn.Module,
-        encoder_cache: EncoderCache | None,
-        device: torch.device,
-    ):
-        super().__init__(vllm_config, model, encoder_cache, device)
-        depth = model.token_lookback_depth
-        self.lookback_token_ids: torch.Tensor | None = None
-        if depth > 0:
-            # Persistent so a captured graph can read it on replay.
-            self.lookback_token_ids = torch.full(
-                (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device
-            )
-
-    def prepare_inputs(
-        self, input_batch: InputBatch, req_states: RequestState
-    ) -> dict[str, torch.Tensor | None]:
-        model_inputs = super().prepare_inputs(input_batch, req_states)
-        window = self.lookback_token_ids
-        if window is None:
-            return model_inputs
-        all_token_ids = req_states.all_token_ids.gpu
-        depth = window.shape[1]
-        _gather_lookback_kernel[(window.shape[0],)](
-            window,
-            input_batch.idx_mapping,
-            req_states.num_computed_tokens.gpu,
-            all_token_ids,
-            all_token_ids.stride(0),
-            input_batch.idx_mapping.shape[0],
-            DEPTH=depth,
-            BLOCK_DEPTH=triton.next_power_of_2(depth),
-        )
-        model_inputs["lookback_token_ids"] = window
-        return model_inputs
-
-    def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
-        model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens)
-        if self.lookback_token_ids is not None:
-            # The captured graph reads this buffer; replays refill it in place.
-            self.lookback_token_ids.fill_(-1)
-            model_inputs["lookback_token_ids"] = self.lookback_token_ids
-        return model_inputs
diff --git a/vllm/models/dots3_note/common/processor.py b/vllm/models/dots3_note/common/processor.py
index e8d365aa380d..4589874ec0a1 100644
--- a/vllm/models/dots3_note/common/processor.py
+++ b/vllm/models/dots3_note/common/processor.py
@@ -13,10 +13,7 @@
 from transformers import BatchFeature
 
 from vllm.config.multimodal import (
-    AudioDummyOptions,
-    BaseDummyOptions,
-    ImageDummyOptions,
-    VideoDummyOptions,
+    MultiModalDummyOptions,
 )
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
@@ -34,6 +31,7 @@
     PromptUpdate,
     PromptUpdateDetails,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
 
 from .video import preprocess_dots3_note_video
@@ -583,7 +581,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         data: dict[str, Any] = {}
         num_images = mm_counts.get("image", 0)
@@ -593,7 +591,7 @@ def get_dummy_mm_data(
                 width=width,
                 height=height,
                 num_images=num_images,
-                overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
+                overrides=mm_options.get("image"),
             )
         num_audios = mm_counts.get("audio", 0)
         if num_audios:
@@ -610,7 +608,7 @@ def get_dummy_mm_data(
             data["audio"] = self._get_dummy_audios(
                 length=min(chunk_samples, max(1, seq_len) * stride),
                 num_audios=num_audios,
-                overrides=cast(AudioDummyOptions | None, mm_options.get("audio")),
+                overrides=mm_options.get("audio"),
             )
         num_videos = mm_counts.get("video", 0)
         if num_videos:
@@ -624,32 +622,39 @@ def get_dummy_mm_data(
                 height=height,
                 num_frames=num_frames,
                 num_videos=num_videos,
-                overrides=cast(VideoDummyOptions | None, mm_options.get("video")),
+                overrides=mm_options.get("video"),
             )
         return data
 
 
 class Dots3NoteMultiModalProcessor(BaseMultiModalProcessor[Dots3NoteProcessingInfo]):
-    def _get_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        processor_data, passthrough_data = super()._get_hf_mm_data(mm_items)
-        if "video" not in mm_items:
-            return processor_data, passthrough_data
-
-        videos = mm_items.get_items("video", VideoProcessorItems)
-        raw_videos: list[object] = []
-        for index, item in enumerate(videos.data):
-            if isinstance(item, MediaWithBytes):
-                raw_videos.append(item.original_bytes)
-            else:
-                raw_videos.append(videos.get(index))
-        processor_data = dict(processor_data)
-        processor_data["videos"] = raw_videos
-        return processor_data, passthrough_data
-
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+
+        # The Dots3Note processor accepts "audios" instead of "audio"
+        hf_data = hf_inputs.hf_data
+        if "audio" in hf_data:
+            hf_data["audios"] = hf_data.pop("audio")
+
+        if "video" in mm_items:
+            videos = mm_items.get_items("video", VideoProcessorItems)
+
+            raw_videos: list[object] = []
+            for index, item in enumerate(videos.data):
+                if isinstance(item, MediaWithBytes):
+                    raw_videos.append(item.original_bytes)
+                else:
+                    raw_videos.append(videos.get(index))
+
+            hf_data["videos"] = raw_videos
+
+        return hf_inputs
+
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
     def _get_mm_fields_config(
diff --git a/vllm/models/dots3_note/nvidia/attention.py b/vllm/models/dots3_note/nvidia/attention.py
index cbd9ad7f4d78..561a3bfd5c1f 100644
--- a/vllm/models/dots3_note/nvidia/attention.py
+++ b/vllm/models/dots3_note/nvidia/attention.py
@@ -19,7 +19,11 @@
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.utils.torch_utils import (
+    PIN_MEMORY,
+    async_tensor_h2d,
+    is_quantized_kv_cache,
+)
 from vllm.v1.attention.backend import (
     AttentionCGSupport,
     AttentionLayer,
@@ -231,8 +235,12 @@ def _build_sliding_window_metadata(
         query_lens = query_lens_cpu[req_start:req_end]
         kv_lens = kv_lens_cpu[req_start:req_end]
         num_reqs = req_end - req_start
-        cu_seq_lens_q_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32)
-        cu_seq_lens_k_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32)
+        cu_seq_lens_q_cpu = torch.zeros(
+            num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY
+        )
+        cu_seq_lens_k_cpu = torch.zeros(
+            num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY
+        )
         torch.cumsum(query_lens, 0, out=cu_seq_lens_q_cpu[1:])
         torch.cumsum(kv_lens, 0, out=cu_seq_lens_k_cpu[1:])
         token_to_seq_cpu = torch.repeat_interleave(
@@ -248,8 +256,8 @@ def _build_sliding_window_metadata(
                 query_end=query_end,
                 cu_seq_lens_q=cu_seq_lens_q_cpu.to(device, non_blocking=True),
                 cu_seq_lens_k=cu_seq_lens_k_cpu.to(device, non_blocking=True),
-                starts=starts_cpu[req_start:req_end].to(device, non_blocking=True),
-                token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
+                starts=async_tensor_h2d(starts_cpu[req_start:req_end], device),
+                token_to_seq=async_tensor_h2d(token_to_seq_cpu, device),
                 num_kv_tokens=num_kv_tokens,
                 max_seq_len_q=int(query_lens.max().item()),
                 max_seq_len_k=int(kv_lens.max().item()),
@@ -345,9 +353,10 @@ def _build_decode(
         query_start_loc_cpu: torch.Tensor,
         query_start_loc_device: torch.Tensor,
         num_decode_tokens: int,
+        max_query_len: int,
         dcp_tot_seq_lens_device: torch.Tensor | None,
     ) -> Dots3NoteDecodeMetadata:
-        del max_seq_len, query_start_loc_device, num_decode_tokens
+        del max_seq_len, query_start_loc_device, num_decode_tokens, max_query_len
         query_len = int(query_start_loc_cpu[1] - query_start_loc_cpu[0])
         return Dots3NoteDecodeMetadata(
             block_table=block_table_tensor,
@@ -396,8 +405,8 @@ def build(
                 seq_lens = common_attn_metadata.seq_lens[
                     reqs_start + req_start : reqs_start + req_end
                 ]
-                query_lens = query_lens_cpu[req_start:req_end].to(
-                    self.device, non_blocking=True
+                query_lens = async_tensor_h2d(
+                    query_lens_cpu[req_start:req_end], self.device
                 )
                 kv_lens = torch.minimum(seq_lens, query_lens + self.sliding_window - 1)
                 chunk.cu_seq_lens_k.zero_()
diff --git a/vllm/models/dots3_note/nvidia/audio_encoder.py b/vllm/models/dots3_note/nvidia/audio_encoder.py
index dbf7eaded26e..c690c87b7f69 100644
--- a/vllm/models/dots3_note/nvidia/audio_encoder.py
+++ b/vllm/models/dots3_note/nvidia/audio_encoder.py
@@ -194,7 +194,7 @@ def forward(self, input_ids, past_key_values_length=0):
 
 
 class WhisperAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
 
     def __init__(
         self,
@@ -346,8 +346,7 @@ def forward(
         rotary_cos: torch.Tensor | None = None,
         rotary_sin: torch.Tensor | None = None,
     ) -> tuple[Any, ...]:
-        """
-        Args:
+        """Args:
             hidden_states: Input to the layer of shape `(batch, seq_len,
                 embed_dim)`, or `(total_tokens, embed_dim)` when `cu_seqlens_q`
                 is given.
@@ -365,6 +364,7 @@ def forward(
         Returns:
             A tuple of the output hidden states, followed by the attention
             weights if `output_attentions` is `True`.
+
         """
         residual = hidden_states
         hidden_states = self.self_attn_layer_norm(hidden_states)
@@ -435,12 +435,12 @@ def _init_weights(self, module):
 
 
 class DotsSpeechEncoder(DotsSpeechPreTrainedModel):
-    """
-    Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+    """Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
     [`WhisperEncoderLayer`].
 
     Args:
         config: WhisperConfig
+
     """
 
     def __init__(self, config: WhisperConfig):
diff --git a/vllm/models/dots3_note/nvidia/vision.py b/vllm/models/dots3_note/nvidia/vision.py
index 4aba14d3c671..774d88830125 100644
--- a/vllm/models/dots3_note/nvidia/vision.py
+++ b/vllm/models/dots3_note/nvidia/vision.py
@@ -9,7 +9,7 @@
 import torch.nn.functional as F
 from torch import nn
 from torch.nn import LayerNorm
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.modeling_utils import PreTrainedModel
 
 from vllm.third_party.deep_gemm import per_block_cast_to_fp8
@@ -27,7 +27,7 @@
 from .vision_moe import note_vision_fused_moe_fp8
 
 
-class DotsMoEVitConfig(PretrainedConfig):
+class DotsMoEVitConfig(PreTrainedConfig):
     model_type: str = "dots_moe_vit"
 
     def __init__(
diff --git a/vllm/models/glm5next/__init__.py b/vllm/models/glm5next/__init__.py
index a1146e2602eb..b38806260bd4 100644
--- a/vllm/models/glm5next/__init__.py
+++ b/vllm/models/glm5next/__init__.py
@@ -6,8 +6,8 @@
 if current_platform.is_xpu():
     raise NotImplementedError("GLM-5.3-Flash does not currently support XPU.")
 
-from .nvidia.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
-from .nvidia.mtp import Glm5NextMTP
+from .common.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
+from .common.mtp import Glm5NextMTP
 
 __all__ = [
     "Glm5NextForCausalLM",
diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py
index 64e6fd46d73f..97d8314c0929 100644
--- a/vllm/models/glm5next/amd/ops/kpool_compress.py
+++ b/vllm/models/glm5next/amd/ops/kpool_compress.py
@@ -65,50 +65,6 @@ def _hadamard128(x):
     return x * 0.08838834764831845  # 1/sqrt(128)
 
 
-# Map pool ids to physical cache slots.
-
-
-def compute_pooled_write_locs(
-    page_table_64: torch.Tensor,
-    pool_ids: torch.Tensor,
-    pool_size: int,
-) -> torch.Tensor:
-    """Map logical pooled-K ids to physical flat cache slots.
-
-    ``pool_size`` consecutive tokens share one pool slot that lives at the
-    *first* token page of each page-group. ``page_table_64`` maps token pages
-    to physical block ids; we gather the block id of each pool's page-group
-    and add the in-block pool offset.
-    """
-    assert page_table_64.ndim == 1
-    pool_ids = pool_ids.to(torch.int64)
-    block_size = 64
-    pool_page_group = torch.div(pool_ids, block_size, rounding_mode="floor")
-    token_page_row = pool_page_group * pool_size
-    packed_page = page_table_64.index_select(0, token_page_row.to(torch.int64))
-    return packed_page.to(torch.int64) * block_size + torch.remainder(
-        pool_ids, block_size
-    )
-
-
-def build_pooled_page_table(
-    page_table: torch.Tensor,
-    pool_size: int,
-) -> torch.Tensor:
-    """Build a pool-granular page table by taking every ``pool_size``-th
-    token-page column (one pool maps to ``pool_size`` token pages).
-
-    Uses gather (not strided slicing) so the result is always a fresh
-    row-major tensor — some downstream kernels require stride(-1) == 1.
-    """
-    block_size = page_table.shape[-1]
-    assert block_size % pool_size == 0, (
-        f"pool_size ({pool_size}) must divide page columns ({block_size})"
-    )
-    idx = torch.arange(0, block_size, pool_size, device=page_table.device)
-    return page_table[..., idx].contiguous()
-
-
 # Fused pool compression and cache write.
 
 
@@ -256,6 +212,13 @@ def kpool_compress_and_write_cache(
         slot_score: ``[n_pools, pool_size, head_dim]`` — per-token gate score.
         ape: ``[pool_size, head_dim]`` fp32 — per-slot position bias.
         loc: ``[n_pools]`` int64 — flat physical slot per pool.
+        pool_size: Number of tokens compressed into one cache entry.
+        head_dim: Indexer head dimension.
+        write_mask: ``[n_pools]`` bool — pools to write, or None for all.
+        round_scale: Round each fp8 scale down to a power of two.
+        return_compressed: Also return the compressed K and scales.
+        write_cache: Write the compressed result into ``kv_cache``.
+
     """
     assert slot_k.ndim == 3
     assert slot_score.shape == slot_k.shape
@@ -636,6 +599,10 @@ def kpool_decode_update_and_maybe_write_cache_batched(
         ape: ``[pool_size, head_dim]`` fp32.
         slot_mapping: ``[num_requests, next_n]`` int32.
         positions: ``[num_requests, next_n]`` int32.
+        pool_size: Number of tokens compressed into one cache entry.
+        head_dim: Indexer head dimension.
+        round_scale: Round each fp8 scale down to a power of two.
+
     """
     num_requests, next_n = key.shape[0], key.shape[1]
     if num_requests == 0 or next_n == 0:
diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py b/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py
index d7de85984f85..d8b25211e0c1 100644
--- a/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py
+++ b/vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py
@@ -564,8 +564,7 @@ def fused_recurrent_gated_delta_rule(
     num_accepted_tokens: torch.Tensor | None = None,
     use_qk_l2norm_in_kernel: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    r"""
-    Args:
+    r"""Args:
         q (torch.Tensor):
             queries of shape `[B, T, H, K]`.
         k (torch.Tensor):
@@ -627,6 +626,7 @@ def fused_recurrent_gated_delta_rule(
             initial_state=h0,
             cu_seqlens=cu_seqlens
         )
+
     """
     if cu_seqlens is not None and q.shape[0] != 1:
         raise ValueError(
diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py
index be417d6b0526..ea86750e7d0a 100644
--- a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py
+++ b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py
@@ -406,19 +406,24 @@ def chunk_kda_scaled_dot_kkt_fwd(
     chunk_size: int = FLA_CHUNK_SIZE,
     output_dtype: torch.dtype = torch.float32,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    r"""
-    Compute beta * K * K^T.
+    r"""Compute beta * K * K^T.
 
     Args:
+        q (torch.Tensor):
+            The query tensor of shape `[B, T, H, K]`.
         k (torch.Tensor):
             The key tensor of shape `[B, T, H, K]`.
         beta (torch.Tensor):
             The beta tensor of shape `[B, T, H]`.
         gk (torch.Tensor):
             The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.
+        scale (float):
+            Scale applied to the query-key products. Default: `None`.
         cu_seqlens (torch.Tensor):
             The cumulative sequence lengths of the input tensor.
             Default: None
+        chunk_indices (torch.Tensor):
+            Precomputed chunk indices for `cu_seqlens`. Default: `None`.
         chunk_size (int):
             The chunk size. Default: 64.
         output_dtype (torch.dtype):
@@ -426,6 +431,7 @@ def chunk_kda_scaled_dot_kkt_fwd(
 
     Returns:
         beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
+
     """
     B, T, H, K = k.shape
     assert K <= 256
@@ -1318,16 +1324,15 @@ def fused_kda_gate(
     safe_gate: bool = False,
     lower_bound: float | None = -5.0,
 ) -> torch.Tensor:
-    """
-    Forward pass for KDA gate:
-      input g: [..., H*D]
-      param A: [H] or [1, 1, H, 1]
-      beta: softplus beta parameter (softplus branch only)
-      threshold: softplus threshold parameter (softplus branch only)
-      safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias);
-        when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias))
-      lower_bound: floor for the safe_gate branch (default -5.0)
-      return  : [..., H, D]
+    """Forward pass for KDA gate:
+    input g: [..., H*D]
+    param A: [H] or [1, 1, H, 1]
+    beta: softplus beta parameter (softplus branch only)
+    threshold: softplus threshold parameter (softplus branch only)
+    safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias);
+      when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias))
+    lower_bound: floor for the safe_gate branch (default -5.0)
+    return  : [..., H, D]
     """
     orig_shape = g.shape[:-1]
 
diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py
new file mode 100644
index 000000000000..761f79d95de7
--- /dev/null
+++ b/vllm/models/glm5next/amd/sparse_indexer.py
@@ -0,0 +1,752 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Custom Sparse Attention Indexer layers."""
+
+import torch
+
+import vllm.envs as envs
+from vllm import _custom_ops  # noqa: F401  # registers the torch.ops._C kernels
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.config import get_current_vllm_config_or_none
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.custom_op import CustomOp
+from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops
+from vllm.models.glm5next.common.sparse_indexer import (
+    RADIX_TOPK_WORKSPACE_SIZE,
+    _build_decode_scatter_indices,
+    _decode_topk_seq_lens,
+    _fill_causal_indices,
+    _fill_short_decode_causal_indices,
+    _gather_workspace_shapes,
+    _scatter_decode_tokens_by_request,
+    kv_cache_as_quant_view,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import (
+    LayerNameType,
+    _encode_layer_name,
+    _resolve_layer_name,
+)
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV32IndexerMetadata,
+)
+from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
+from vllm.v1.worker.workspace import current_workspace_manager
+
+logger = init_logger(__name__)
+
+# kpool write helper: form pools from the current token batch and compress them
+# into the index K cache via the fused Triton kernel.
+
+
+def _kpool_compress_insert(
+    k: torch.Tensor,
+    gate_score: torch.Tensor,
+    ape: torch.Tensor,
+    kv_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kpool: int,
+    head_dim: int,
+    round_scale: bool,
+) -> None:
+    """Pool ``kpool`` consecutive tokens into one fp8 K and write at pool slots.
+
+    ``slot_mapping`` is pool-granular (compress_ratio == kpool on the spec):
+    only the *last* token of each complete pool carries a valid (>=0) slot;
+    intra-pool tokens are -1. Every position is treated as a pool-completion
+    candidate and non-completions are masked off inside the kernel. Compacting
+    the valid rows first costs two device syncs on the eager prefill path and
+    buys nothing numerically. Assumes pool-aligned chunk starts.
+    """
+    n = slot_mapping.shape[0]
+    # No pool can complete in a batch smaller than one pool; also keeps the
+    # clamped gather indices below in bounds.
+    if n < kpool:
+        return
+    pos = torch.arange(n, device=k.device)
+    valid = slot_mapping >= 0
+    # Drop pools whose start falls before the batch (leading padding); their
+    # gate/k data is undefined anyway.
+    write_mask = valid & (pos >= kpool - 1)
+    offs = torch.arange(kpool, device=k.device)
+    idx = (pos - (kpool - 1)).clamp_min(0)[:, None] + offs[None, :]
+    kpool_ops.kpool_compress_and_write_cache(
+        kv_cache,
+        k[idx],  # [n, kpool, head_dim]
+        gate_score[idx],
+        ape,
+        slot_mapping.to(torch.int64),
+        pool_size=kpool,
+        head_dim=head_dim,
+        write_mask=write_mask,
+        round_scale=round_scale,
+        write_cache=True,
+        return_compressed=False,
+    )
+
+
+@eager_break_during_capture
+def sparse_attn_indexer_kpool(
+    hidden_states: torch.Tensor,
+    k_cache_prefix: LayerNameType,
+    kv_cache: torch.Tensor,
+    q_quant: torch.Tensor,
+    q_scale: torch.Tensor | None,
+    k: torch.Tensor,
+    weights: torch.Tensor,
+    quant_block_size: int,
+    scale_fmt: str | None,
+    topk_tokens: int,
+    head_dim: int,
+    max_pool_len: int,
+    total_seq_lens: int,
+    topk_indices_buffer: torch.Tensor,
+    skip_k_cache_insert: bool,
+    use_fp4_cache: bool = False,
+    # kpool params (Plan-A: gate is consumed at write time and read back at
+    # topk time to softmax-weight the pool).
+    gate_score: torch.Tensor | None = None,
+    compress_ape: torch.Tensor | None = None,
+    index_kpool: int = 1,
+    positions: torch.Tensor | None = None,
+    # Paged tail cache (in-progress pool's raw K + gate score), replacing the
+    # transient _DECODE_TAIL ring. tail_prefix resolves attn_metadata[tail_prefix]
+    # for the tail group's token-granular slot_mapping. None on the dummy/profiling
+    # path and when the tail cache is disabled.
+    tail_kv_cache: torch.Tensor | None = None,
+    tail_prefix: str | None = None,
+) -> torch.Tensor:
+    # careful! this will be None in dummy run
+    attn_metadata = get_forward_context().attn_metadata
+    fp8_dtype = current_platform.fp8_dtype()
+    k_cache_prefix = _resolve_layer_name(k_cache_prefix)
+
+    # assert isinstance(attn_metadata, dict)
+    if not isinstance(attn_metadata, dict):
+        # Reserve workspace for indexer during profiling run
+        values_spec, scales_spec = _gather_workspace_shapes(
+            total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
+        )
+        current_workspace_manager().get_simultaneous(
+            values_spec,
+            scales_spec,
+            ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+        )
+
+        # Reserve profiler-visible memory for the worst-case decode logits,
+        # whose shape is [B * next_n, max_pool_len]. This profiling branch
+        # returns before invoking the logits kernel itself.
+        cfg = get_current_vllm_config_or_none()
+        worst_decode_tokens = 0
+        if cfg is not None:
+            sched = cfg.scheduler_config
+            num_spec = (
+                cfg.speculative_config.num_speculative_tokens
+                if cfg.speculative_config is not None
+                else 0
+            )
+            worst_decode_tokens = min(
+                sched.max_num_seqs * (num_spec + 1),
+                sched.max_num_batched_tokens,
+            )
+        # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
+        decode_logits_elems = worst_decode_tokens * max_pool_len * 4
+        prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+        max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
+        _ = torch.empty(
+            max_logits_elems, dtype=torch.uint8, device=hidden_states.device
+        )
+
+        return topk_indices_buffer
+    attn_metadata_narrowed = attn_metadata[k_cache_prefix]
+    assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
+    slot_mapping = attn_metadata_narrowed.slot_mapping
+    has_decode = attn_metadata_narrowed.num_decodes > 0
+    has_prefill = attn_metadata_narrowed.num_prefills > 0
+    num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
+
+    # q_scale is required iff the FP4 cache path is enabled; the FP8 path
+    # folds the Q scale into `weights` inside fused_indexer_q_rope_quant.
+    if use_fp4_cache:
+        assert q_scale is not None, "use_fp4_cache=True requires q_scale"
+    else:
+        assert q_scale is None, "q_scale must be None when use_fp4_cache=False"
+
+    # During speculative decoding, k may be padded to the CUDA graph batch
+    # size while slot_mapping only covers actual tokens. Truncate k to avoid
+    # out-of-bounds reads in the kernel.
+    num_tokens = slot_mapping.shape[0]
+    if k is not None:
+        k = k[:num_tokens]
+
+    if not skip_k_cache_insert:
+        assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
+        if index_kpool > 1 and gate_score is not None and compress_ape is not None:
+            # kpool prefill write: pool kpool consecutive prefill tokens via
+            # softmax(gate+ape)-weighted sum -> Hadamard -> fp8 -> pool slots.
+            # Decode tokens (the first num_decode_tokens in the batch) cannot be
+            # pooled here — their pool's earlier tokens are not in this batch —
+            # so they are deferred to the tail-buffer kernel in has_decode.
+            # compress_ratio == index_kpool makes slot_mapping pool-granular.
+            n_prefill = num_tokens - num_decode_tokens
+            if n_prefill > 0:
+                # decode tokens are batched first; prefill tokens follow.
+                prefill_slice = slice(num_decode_tokens, num_tokens)
+                _kpool_compress_insert(
+                    k[prefill_slice],
+                    gate_score[prefill_slice],
+                    compress_ape,
+                    kv_cache,
+                    slot_mapping[prefill_slice],
+                    index_kpool,
+                    head_dim,
+                    round_scale=(scale_fmt is not None),
+                )
+                # Persist each request's incomplete prefill pool so decode can
+                # finish it, including after PD transfer. Tail slots use
+                # ``pos % kpool`` within the request's tail block. Processing
+                # only the batch's trailing tokens would miss all but the last
+                # request in a multi-request prefill.
+                if tail_kv_cache is not None and tail_prefix is not None:
+                    tail_meta = attn_metadata.get(_resolve_layer_name(tail_prefix))
+                    if tail_meta is not None:
+                        assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
+                        kpool_ops.kpool_seed_tail_cache(
+                            tail_kv_cache,
+                            k[prefill_slice],
+                            gate_score[prefill_slice],
+                            tail_meta.slot_mapping[prefill_slice],
+                            index_kpool,
+                            head_dim,
+                        )
+        else:
+            # standard: per-token fp8 quant + scatter (all tokens).
+            assert scale_fmt is not None
+            from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+                indexer_k_quant_and_cache_triton,
+            )
+
+            indexer_k_quant_and_cache_triton(
+                k,
+                kv_cache,
+                slot_mapping,
+                quant_block_size,
+                scale_fmt,
+            )
+
+    topk_indices_buffer[: hidden_states.shape[0]] = -1
+    if has_prefill:
+        prefill_metadata = attn_metadata_narrowed.prefill
+        assert prefill_metadata is not None
+
+        # Short sequences select every pool, so skip sparse scoring and fill
+        # the top-k buffer with all causal token indices. The index-K cache was
+        # already written above.
+        n_prefill_sf = num_tokens - num_decode_tokens
+        # Host-side short-prefill predicate: max_prefill_seq_len is computed
+        # in the metadata builder (exact for prefill rows) and equals
+        # positions[prefill_slice].max() + 1, so this replaces a
+        # positions.max().item() device sync per layer. -1 (unknown metadata)
+        # falls back to the device-side check.
+        if prefill_metadata.max_prefill_seq_len >= 0:
+            short_prefill = (
+                n_prefill_sf > 0
+                and positions is not None
+                and prefill_metadata.max_prefill_seq_len <= topk_tokens
+            )
+        else:
+            short_prefill = (
+                n_prefill_sf > 0
+                and positions is not None
+                and int(positions[num_decode_tokens:num_tokens].max().item()) + 1
+                <= topk_tokens
+            )
+        if short_prefill:
+            # short_prefill is only True when positions is not None (above),
+            # but narrow explicitly for the indexer below.
+            assert positions is not None
+            _pos = positions[num_decode_tokens:num_tokens].to(torch.int32)
+            _buf = topk_indices_buffer[num_decode_tokens:num_tokens]
+            _fill_causal_indices(_buf, _pos)
+
+        # Get the full shared workspace buffers once (will allocate on first use).
+        # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and
+        # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0
+        # scales) based on use_fp4_cache.
+        workspace_manager = current_workspace_manager()
+        values_spec, scales_spec = _gather_workspace_shapes(
+            total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
+        )
+        k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
+            values_spec,
+            scales_spec,
+        )
+        for chunk in prefill_metadata.chunks if not short_prefill else ():
+            k_quant = k_quant_full[: chunk.total_seq_lens]
+            k_scale = k_scale_full[: chunk.total_seq_lens]
+
+            if not chunk.skip_kv_gather:
+                from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+                    cp_gather_indexer_k_quant_cache_triton,
+                )
+
+                cp_gather_indexer_k_quant_cache_triton(
+                    kv_cache,
+                    k_quant,
+                    k_scale,
+                    chunk.block_table,
+                    chunk.cu_seq_lens,
+                    token_to_seq=chunk.token_to_seq,
+                )
+
+            q_slice = q_quant[chunk.token_start : chunk.token_end]
+            q_scale_slice = (
+                q_scale[chunk.token_start : chunk.token_end]
+                if q_scale is not None
+                else None
+            )
+            # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
+            # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
+            if use_fp4_cache:
+                q_slice_cast = q_slice.view(torch.int8)
+                k_quant_cast = k_quant.view(torch.int8)
+                k_scale_cast = k_scale.view(torch.int32).squeeze(-1)
+            else:
+                q_slice_cast = q_slice
+                k_quant_cast = k_quant
+                k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
+            from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+                rocm_fp8_mqa_logits,
+            )
+
+            assert q_scale_slice is None
+            logits = rocm_fp8_mqa_logits(
+                q_slice_cast,
+                (k_quant_cast, k_scale_cast),
+                weights[chunk.token_start : chunk.token_end],
+                chunk.cu_seqlen_ks,
+                chunk.cu_seqlen_ke,
+            )
+            num_rows = logits.shape[0]
+
+            # kpool: logits are pool-granular (compress_ratio == index_kpool),
+            # so topk selects pools. We pick topk_tokens // kpool pools then
+            # expand each pool back to its kpool constituent tokens.
+            select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens
+            if index_kpool > 1:
+                pool_topk = torch.empty(
+                    (num_rows, select_k), dtype=torch.int32, device=logits.device
+                )
+                topk_dst = pool_topk
+            else:
+                topk_dst = topk_indices_buffer[
+                    chunk.token_start : chunk.token_end, :topk_tokens
+                ]
+
+            torch.ops._C.top_k_per_row_prefill(
+                logits,
+                chunk.cu_seqlen_ks,
+                chunk.cu_seqlen_ke,
+                topk_dst,
+                num_rows,
+                logits.stride(0),
+                logits.stride(1),
+                select_k,
+            )
+
+            if index_kpool > 1:
+                pool_ids = pool_topk.to(torch.int64)
+                if positions is not None:
+                    # Fused expand-pools + append-tail into one Triton kernel
+                    # (replaces ~25 elementwise ops). seq_len is token-granular
+                    # (pos+1); the kernel derives pool_len internally.
+                    q_seq = (
+                        positions[chunk.token_start : chunk.token_end].to(torch.int32)
+                        + 1
+                    )
+                    expanded = kpool_ops.expand_pools_and_append_tail(
+                        pool_ids, q_seq, index_kpool
+                    )
+                else:
+                    valid = pool_ids >= 0
+                    expanded = kpool_ops.expand_pools_to_tokens(
+                        pool_ids, valid, topk_tokens, index_kpool
+                    )
+                topk_indices_buffer[
+                    chunk.token_start : chunk.token_end, : expanded.shape[-1]
+                ] = expanded
+
+    if has_decode:
+        decode_metadata = attn_metadata_narrowed.decode
+        assert decode_metadata is not None
+        kv_cache_raw = kv_cache  # raw [num_blocks, block_size, head_dim+4] for writes
+        kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
+
+        # Update the tail before reading logits; completed pools are compressed
+        # into the slot supplied by slot_mapping.
+        # Spec verification groups tokens by request and preserves position
+        # order so each token is stashed before the next completes its pool.
+        # Positions must remain token-granular because the kernel derives the
+        # pool phase and tail index from ``pos % kpool``.
+        if (
+            index_kpool > 1
+            and gate_score is not None
+            and compress_ape is not None
+            and positions is not None
+            and not skip_k_cache_insert
+        ):
+            num_requests = attn_metadata_narrowed.num_decodes
+            # Kpool writes must recover the original request grouping after the
+            # indexer's flattened decode path. Host metadata avoids a CUDA graph
+            # sync when choosing the uniform or padded layout.
+            per_req_lens = decode_metadata.per_req_decode_lens
+            if per_req_lens is not None:
+                use_uniform = (
+                    decode_metadata.decode_is_uniform
+                    and num_decode_tokens
+                    == num_requests * decode_metadata.write_max_decode_len
+                )
+                group_lens = per_req_lens
+                lmax = decode_metadata.write_max_decode_len
+            else:
+                # Legacy metadata without per-request lens: fall back to the
+                # host-side requires_padding flag. Unreached now (per-request
+                # lens is always populated for decode), kept defensive.
+                use_uniform = not decode_metadata.requires_padding
+                group_lens = decode_metadata.decode_lens
+                lmax = int(decode_metadata.decode_lens.max().item())
+            if not use_uniform:
+                # Non-uniform decode_lens (mixed plain-decode + spec-verify, or
+                # a variable MTP-verify batch): scatter actual tokens into a
+                # padded [B, lmax] layout. int32 tensors can't go through
+                # pack_seq_triton (float/uint8 only). The scatter indices are
+                # shared by all five scatters below (and the tail slot one).
+                scatter_idx = _build_decode_scatter_indices(
+                    group_lens, num_requests, num_decode_tokens
+                )
+                dec_k = _scatter_decode_tokens_by_request(
+                    k[:num_decode_tokens], 0, num_requests, lmax, scatter_idx
+                )
+                dec_gate = _scatter_decode_tokens_by_request(
+                    gate_score[:num_decode_tokens],
+                    0,
+                    num_requests,
+                    lmax,
+                    scatter_idx,
+                )
+                dec_slot = _scatter_decode_tokens_by_request(
+                    slot_mapping[:num_decode_tokens],
+                    -1,
+                    num_requests,
+                    lmax,
+                    scatter_idx,
+                )
+                dec_pos = _scatter_decode_tokens_by_request(
+                    positions[:num_decode_tokens].to(torch.int32),
+                    -1,
+                    num_requests,
+                    lmax,
+                    scatter_idx,
+                )
+            else:
+                next_n = num_decode_tokens // num_requests
+                shape2 = (num_requests, next_n)
+                dec_k = k[:num_decode_tokens].view(*shape2, head_dim)
+                dec_gate = gate_score[:num_decode_tokens].view(*shape2, head_dim)
+                dec_slot = slot_mapping[:num_decode_tokens].view(shape2)
+                dec_pos = positions[:num_decode_tokens].to(torch.int32).view(shape2)
+            tail_meta = (
+                attn_metadata.get(_resolve_layer_name(tail_prefix))
+                if tail_prefix is not None
+                else None
+            )
+            # Paged tail cache replaces the transient _DECODE_TAIL ring. Group
+            # the tail group's token-granular slot_mapping per-request, mirroring
+            # dec_slot / dec_pos, so the kernel gets each request's current-token
+            # tail slot (block * kpool + pos % kpool).
+            if tail_meta is not None:
+                assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
+            if tail_meta is None or tail_kv_cache is None:
+                dec_tail_slot = None
+            elif not use_uniform:
+                dec_tail_slot = _scatter_decode_tokens_by_request(
+                    tail_meta.slot_mapping[:num_decode_tokens],
+                    -1,
+                    num_requests,
+                    lmax,
+                    scatter_idx,
+                )
+            else:
+                dec_tail_slot = tail_meta.slot_mapping[:num_decode_tokens].view(shape2)
+            # The compress kernel writes the raw fp8 cache (not the quant view);
+            # pass the underlying kv_cache, not kv_cache_quant_view.
+            if dec_tail_slot is not None:
+                # Single batched launch over [num_requests, next_n] replaces the
+                # per-token sequential loop. The kernel iterates each request's
+                # tokens in position order internally, preserving the
+                # pool-completion read-after-stash dependency that the loop
+                # provided. Inputs are already grouped per request (uniform:
+                # view; non-uniform: _scatter_decode_tokens_by_request padded to
+                # [B, lmax]) — no per-token .contiguous() copies needed.
+                kpool_ops.kpool_decode_update_and_maybe_write_cache_batched(
+                    kv_cache_raw,
+                    tail_kv_cache,
+                    dec_tail_slot,
+                    dec_k,
+                    dec_gate,
+                    compress_ape,
+                    dec_slot,
+                    dec_pos,
+                    index_kpool,
+                    head_dim,
+                    round_scale=(scale_fmt is not None),
+                )
+        if current_platform.is_cuda_alike() and _fill_short_decode_causal_indices(
+            topk_indices_buffer,
+            positions,
+            num_decode_tokens,
+            attn_metadata_narrowed.max_seq_len,
+            topk_tokens,
+        ):
+            return topk_indices_buffer
+        decode_lens = decode_metadata.decode_lens
+        if decode_metadata.requires_padding:
+            # Padding also covers short chunked prefills classified as decode.
+            # MXFP4 uses zero-byte padding so padded slots dequantize to zero.
+            if q_scale is not None:
+                padded_q_quant_decode_tokens = pack_seq_triton(
+                    q_quant[:num_decode_tokens], decode_lens, pad_value=0
+                )
+                padded_q_scale = pack_seq_triton(
+                    q_scale[:num_decode_tokens], decode_lens, pad_value=0
+                )
+            else:
+                padded_q_quant_decode_tokens = pack_seq_triton(
+                    q_quant[:num_decode_tokens], decode_lens
+                )
+                padded_q_scale = None
+            padded_weights = pack_seq_triton(
+                weights[:num_decode_tokens], decode_lens, pad_value=0
+            ).reshape(-1, *weights.shape[1:])
+        else:
+            padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape(
+                decode_lens.shape[0], -1, *q_quant.shape[1:]
+            )
+            if q_scale is not None:
+                padded_q_scale = q_scale[:num_decode_tokens].reshape(
+                    decode_lens.shape[0], -1, *q_scale.shape[1:]
+                )
+            else:
+                padded_q_scale = None
+            padded_weights = weights[:num_decode_tokens]
+        # TODO: move and optimize below logic with triton kernels
+        batch_size = padded_q_quant_decode_tokens.shape[0]
+        next_n = padded_q_quant_decode_tokens.shape[1]
+        num_padded_tokens = batch_size * next_n
+        seq_lens = decode_metadata.seq_lens[:batch_size]
+        # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1)
+        # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens;
+        # the downstream topk kernels accept both 1D and 2D.
+        padded_q_quant_cast = (
+            padded_q_quant_decode_tokens.view(torch.int8)
+            if use_fp4_cache
+            else padded_q_quant_decode_tokens
+        )
+        from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+            rocm_fp8_paged_mqa_logits,
+        )
+
+        assert padded_q_scale is None
+        logits = rocm_fp8_paged_mqa_logits(
+            padded_q_quant_cast,
+            kv_cache,
+            padded_weights[:num_padded_tokens],
+            seq_lens,
+            decode_metadata.block_table,
+            decode_metadata.schedule_metadata,
+            max_model_len=max_pool_len,
+        )
+        num_rows = logits.shape[0]
+        # kpool: logits are pool-granular -> select topk_tokens//kpool pools,
+        # then expand each pool back to its kpool tokens.
+        select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens
+        if index_kpool > 1:
+            pool_topk = torch.empty(
+                (num_rows, select_k), dtype=torch.int32, device=logits.device
+            )
+            topk_dst = pool_topk
+        else:
+            topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
+
+        torch.ops._C.top_k_per_row_decode(
+            logits,
+            next_n,
+            seq_lens,
+            topk_dst,
+            num_rows,
+            logits.stride(0),
+            logits.stride(1),
+            select_k,
+        )
+
+        # Resolve to token-level indices in the output buffer.
+        if index_kpool > 1:
+            pool_ids = pool_topk.to(torch.int64)
+            n = pool_topk.shape[0]
+            # Decode seq_lens are pool-granular; recover token lengths from
+            # positions using the padded [B, next_n] row layout when needed.
+            if positions is not None:
+                dec_seq = _decode_topk_seq_lens(
+                    positions,
+                    decode_lens,
+                    num_decode_tokens,
+                    batch_size,
+                    next_n,
+                    decode_metadata.requires_padding,
+                )
+            else:
+                dec_seq = decode_metadata.seq_lens[:n]
+                if dec_seq.ndim == 2:
+                    dec_seq = dec_seq[:, -1]
+                dec_seq = dec_seq.to(torch.int32)
+            out = kpool_ops.expand_pools_and_append_tail(pool_ids, dec_seq, index_kpool)
+        else:
+            out = topk_dst
+
+        if decode_metadata.requires_padding:
+            # Drop padded query rows introduced by the next_n padding above.
+            out = unpack_seq_triton(
+                out.reshape(batch_size, -1, out.shape[-1]), decode_lens
+            )
+        topk_indices_buffer[: out.shape[0], : out.shape[-1]] = out
+
+    return topk_indices_buffer
+
+
+@CustomOp.register("sparse_attn_indexer_kpool")
+class SparseAttnIndexerKpool(CustomOp):
+    """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a
+    separate custom op since it involves heavy custom kernels like `mqa_logits`,
+    `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires
+    specific memory layout or implementation for different hardware backends to
+    achieve optimal performance.
+
+    For now, the default native path will use CUDA backend path. Other platform
+    may requires add the corresponding Custom Op name `sparse_attn_indexer` to
+    `custom_ops` in `CompilationConfig` to enable the platform specific path.
+    """
+
+    def __init__(
+        self,
+        k_cache,
+        quant_block_size: int,
+        scale_fmt: str,
+        topk_tokens: int,
+        head_dim: int,
+        max_pool_len: int,
+        max_total_seq_len: int,
+        topk_indices_buffer: torch.Tensor,
+        skip_k_cache_insert: bool = False,
+        use_fp4_cache: bool = False,
+        tail_cache=None,
+    ):
+        super().__init__()
+        self.k_cache = k_cache
+        self.tail_cache = tail_cache
+        self.quant_block_size = quant_block_size
+        self.scale_fmt = scale_fmt
+        self.topk_tokens = topk_tokens
+        self.head_dim = head_dim
+        self.max_pool_len = max_pool_len
+        self.max_total_seq_len = max_total_seq_len
+        self.topk_indices_buffer = topk_indices_buffer
+        self.skip_k_cache_insert = skip_k_cache_insert
+        self.use_fp4_cache = use_fp4_cache
+
+    def forward_hip(
+        self,
+        hidden_states: torch.Tensor,
+        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        weights: torch.Tensor,
+        *,
+        gate_score: torch.Tensor | None = None,
+        compress_ape: torch.Tensor | None = None,
+        index_kpool: int = 1,
+        positions: torch.Tensor | None = None,
+    ):
+        assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
+        assert isinstance(q_quant, torch.Tensor), (
+            "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
+        )
+        if not rocm_aiter_ops.is_enabled():
+            raise RuntimeError(
+                "Sparse attention indexer ROCm path is only supported on AITER. "
+                "Please enable aiter with VLLM_ROCM_USE_AITER=1"
+            )
+        if index_kpool <= 1:
+            return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
+                hidden_states,
+                _encode_layer_name(self.k_cache.prefix),
+                self.k_cache.kv_cache,
+                q_quant,
+                k,
+                weights,
+                self.quant_block_size,
+                self.scale_fmt,
+                self.topk_tokens,
+                self.head_dim,
+                self.max_pool_len,
+                self.max_total_seq_len,
+                self.topk_indices_buffer,
+                skip_k_cache_insert=self.skip_k_cache_insert,
+            )
+        return sparse_attn_indexer_kpool(
+            hidden_states,
+            self.k_cache.prefix,
+            self.k_cache.kv_cache,
+            q_quant,
+            None,
+            k,
+            weights,
+            self.quant_block_size,
+            self.scale_fmt,
+            self.topk_tokens,
+            self.head_dim,
+            self.max_pool_len,
+            self.max_total_seq_len,
+            self.topk_indices_buffer,
+            self.skip_k_cache_insert,
+            self.use_fp4_cache,
+            gate_score,
+            compress_ape,
+            index_kpool,
+            positions,
+            self.tail_cache.kv_cache if self.tail_cache is not None else None,
+            self.tail_cache.prefix if self.tail_cache is not None else None,
+        )
+
+    def forward_native(
+        self,
+        hidden_states: torch.Tensor,
+        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        k: torch.Tensor,
+        weights: torch.Tensor,
+        *,
+        gate_score: torch.Tensor | None = None,
+        compress_ape: torch.Tensor | None = None,
+        index_kpool: int = 1,
+        positions: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.forward_hip(
+            hidden_states,
+            q_quant,
+            k,
+            weights,
+            gate_score=gate_score,
+            compress_ape=compress_ape,
+            index_kpool=index_kpool,
+            positions=positions,
+        )
diff --git a/vllm/models/glm5next/common/__init__.py b/vllm/models/glm5next/common/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/glm5next/common/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/common/attention.py
similarity index 98%
rename from vllm/models/glm5next/nvidia/attention.py
rename to vllm/models/glm5next/common/attention.py
index 0be749348494..4f776e8a8e76 100644
--- a/vllm/models/glm5next/nvidia/attention.py
+++ b/vllm/models/glm5next/common/attention.py
@@ -23,7 +23,6 @@
 from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope
-from vllm.model_executor.layers.sparse_attn_indexer_kpool import SparseAttnIndexerKpool
 from vllm.model_executor.models.deepseek_v2 import (
     DeepSeekV2FusedQkvAProjLinear,
     DeepseekV32IndexerCache,
@@ -31,6 +30,7 @@
 )
 from vllm.model_executor.utils import maybe_disable_graph_partition
 from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8
+from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool
 from vllm.platforms import current_platform
 from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
 from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES
@@ -296,7 +296,7 @@ def __init__(
             cache_config=cache_config,
             index_kpool=self.index_kpool,
         )
-        self.max_model_len = vllm_config.model_config.max_model_len
+        self.max_pool_len = vllm_config.model_config.max_model_len // self.index_kpool
         self.prefix = prefix
         from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
 
@@ -307,7 +307,7 @@ def __init__(
             self.scale_fmt,
             self.topk_tokens,
             self.head_dim,
-            self.max_model_len,
+            self.max_pool_len,
             self.max_total_seq_len,
             self.topk_indices_buffer,
             tail_cache=self.tail_cache,
@@ -496,9 +496,9 @@ def __init__(
             assert config.rope_parameters is not None
             if config.rope_parameters["rope_type"] != "default":
                 config.rope_parameters["rope_type"] = (
-                    "deepseek_yarn"
-                    if config.rope_parameters.get("apply_yarn_scaling", True)
-                    else "deepseek_llama_scaling"
+                    "deepseek_llama_scaling"
+                    if config.rope_parameters.get("attention_factor") == 1.0
+                    else "deepseek_yarn"
                 )
 
             self.rotary_emb: RotaryEmbedding | None = get_rope(
diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/common/kda.py
similarity index 78%
rename from vllm/models/glm5next/nvidia/kda.py
rename to vllm/models/glm5next/common/kda.py
index e6ba6ac5c659..436b9685cb12 100644
--- a/vllm/models/glm5next/nvidia/kda.py
+++ b/vllm/models/glm5next/common/kda.py
@@ -37,6 +37,7 @@
 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
 from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
+from vllm.v1.worker.workspace import current_workspace_manager
 
 if current_platform.is_rocm():
     from vllm.models.glm5next.amd.ops.third_party.kda import (
@@ -123,6 +124,32 @@ def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor:
     return x.float().sigmoid()
 
 
+def _resolve_kda_prefill_backend(
+    backend: str, head_dim: int, dtype: torch.dtype, lower_bound: float | None
+) -> str:
+    """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on
+    SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton
+    ``chunk_kda_with_fused_gate`` path. ``backend`` comes from
+    ``additional_config.kda_prefill_backend`` (auto / triton / flashkda)."""
+    if backend not in ("auto", "triton", "flashkda"):
+        raise ValueError(f"Unsupported KDA prefill backend: {backend}")
+    capability = current_platform.get_device_capability()
+    supported = (
+        current_platform.is_cuda()
+        and capability is not None
+        and capability.major in (9, 10, 12)
+        and head_dim == 128
+        and dtype == torch.bfloat16
+        and lower_bound is not None
+    )
+    if backend == "flashkda" and not supported:
+        raise RuntimeError(
+            "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, head_dim=128 "
+            "and a bounded KDA gate."
+        )
+    return "flashkda" if supported and backend != "triton" else "triton"
+
+
 class Glm5NextLinearAttention(GatedDeltaNetAttention):
     head_dim: int
     num_heads: int
@@ -175,6 +202,15 @@ def __init__(
         projection_size = self.head_dim * self.num_heads
         self.local_projection_size = divide(projection_size, self.tp_size)
 
+        _, recurrent_state_shape = self.get_state_shape()
+        _, recurrent_state_dtype = self.get_state_dtype()
+        scatter_states.register_warmup(
+            state_shape=recurrent_state_shape,
+            state_dtype=recurrent_state_dtype,
+            indices_dtype=torch.int32,
+            max_num_tokens=vllm_config.scheduler_config.max_num_seqs,
+        )
+
         # Merge q, k, v, b, f_a, g_a projections into one GEMM (6→1 launches).
         # Order matches checkpoint's fused_qkvbfg_a_proj convention.
         # Shards 4 (f_a) and 5 (g_a) are replicated across TP ranks.
@@ -282,6 +318,85 @@ def _a_log_weight_loader(param, loaded_weight):
         # every _forward call (it reads an env-derived flag each time).
         self._conv_state_dim_first = is_conv_state_dim_first()
 
+        additional_config = vllm_config.additional_config
+        self.kda_prefill_backend = _resolve_kda_prefill_backend(
+            additional_config.get("kda_prefill_backend", "auto")
+            if isinstance(additional_config, dict)
+            else "auto",
+            self.head_dim,
+            vllm_config.model_config.dtype,
+            self.kda_lower_bound,
+        )
+        self._flashkda_buffer_specs: (
+            tuple[tuple[tuple[int, ...], torch.dtype], ...] | None
+        ) = None
+        if self.kda_prefill_backend == "flashkda":
+            import vllm._flashkda_C  # noqa: F401
+
+            max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+            max_seqs = vllm_config.scheduler_config.max_num_seqs
+            workspace_size = torch.ops._flashkda_C.get_workspace_size(
+                max_tokens, self.local_num_heads, max_seqs
+            )
+            self._flashkda_buffer_specs = (
+                (
+                    (max_seqs, self.local_num_heads, self.head_dim, self.head_dim),
+                    self.get_state_dtype()[1],
+                ),
+                ((workspace_size,), torch.uint8),
+                # Output buffer for steps that also carry spec-decode tokens:
+                # the non-spec tokens are then scattered by non_spec_token_indx.
+                (
+                    (1, max_tokens, self.local_num_heads, self.head_dim),
+                    vllm_config.model_config.dtype,
+                ),
+            )
+
+    def _flashkda_prefill(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        initial_state: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        out: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Fused KDA chunked prefill (FlashKDA). Takes the raw gate logits ``g``
+        and raw ``beta`` logits, l2-normalizes q/k in-kernel and applies the
+        bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``,
+        matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the
+        attention output into ``out`` (a workspace buffer when ``None``) and
+        returns ``(out, final_state)``."""
+        assert self._flashkda_buffer_specs is not None
+        final_state, workspace, workspace_out = (
+            current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs)
+        )
+        final_state = final_state[: initial_state.shape[0]]
+        if out is None:
+            out = workspace_out[:, : q.shape[1]]
+        # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided.
+        torch.ops._flashkda_C.fwd(
+            q.contiguous(),
+            k.contiguous(),
+            v.contiguous(),
+            g.contiguous(),
+            beta,
+            self.head_dim**-0.5,
+            out,
+            workspace,
+            self.A_log.view(-1),
+            self.dt_bias.view(-1, self.head_dim),
+            self.kda_lower_bound,
+            initial_state.contiguous(),
+            final_state,
+            cu_seqlens.contiguous(),
+            None,
+            None,
+        )
+        return out, final_state
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -533,26 +648,42 @@ def _rearr(x):
             initial_state = gather_initial_states(
                 recurrent_state, non_spec_state_indices_tensor, has_initial_state
             )
-            (
-                core_attn_out_non_spec,
-                last_recurrent_state,
-            ) = chunk_kda_with_fused_gate(
-                q=_rearr(q_ns),
-                k=_rearr(k_ns),
-                v=_rearr(v_ns),
-                raw_g=g1_ns,
-                # Chunk path wants the pre-sigmoided fp32 beta (its kernels
-                # don't sigmoid); beta_ns is raw bf16 from forward.
-                beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0),
-                A_log=self.A_log,
-                g_bias=self.dt_bias,
-                initial_state=initial_state,
-                output_final_state=True,
-                use_qk_l2norm_in_kernel=True,
-                cu_seqlens=non_spec_query_start_loc,
-                safe_gate=safe_gate,
-                lower_bound=lower_bound,
-            )
+            if self.kda_prefill_backend == "flashkda":
+                # Non-spec step: write straight into the layer output buffer
+                # (dense token order, no merge copy). Step with spec-decode
+                # tokens: write to the workspace buffer and scatter below.
+                ns_out = None if use_spec else core_attn_out[:, :num_actual_tokens]
+                core_attn_out_non_spec, last_recurrent_state = self._flashkda_prefill(
+                    q=_rearr(q_ns),
+                    k=_rearr(k_ns),
+                    v=_rearr(v_ns),
+                    g=g1_ns,
+                    beta=beta_ns,
+                    initial_state=initial_state,
+                    cu_seqlens=non_spec_query_start_loc,
+                    out=ns_out,
+                )
+            else:
+                (
+                    core_attn_out_non_spec,
+                    last_recurrent_state,
+                ) = chunk_kda_with_fused_gate(
+                    q=_rearr(q_ns),
+                    k=_rearr(k_ns),
+                    v=_rearr(v_ns),
+                    raw_g=g1_ns,
+                    # Chunk path wants the pre-sigmoided fp32 beta (its
+                    # kernels don't sigmoid); beta_ns is raw bf16 from forward.
+                    beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0),
+                    A_log=self.A_log,
+                    g_bias=self.dt_bias,
+                    initial_state=initial_state,
+                    output_final_state=True,
+                    use_qk_l2norm_in_kernel=True,
+                    cu_seqlens=non_spec_query_start_loc,
+                    safe_gate=safe_gate,
+                    lower_bound=lower_bound,
+                )
             # Init cache
             scatter_states(
                 recurrent_state,
@@ -591,14 +722,8 @@ def _rearr(x):
         # --- merge spec / non-spec outputs back into token order ---
         if use_spec and core_attn_out_non_spec is not None:
             assert core_attn_out_spec is not None
-            merged = torch.empty(
-                (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
-                dtype=core_attn_out_non_spec.dtype,
-                device=core_attn_out_non_spec.device,
-            )
-            merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
-            merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
-            core_attn_out[0, :num_actual_tokens] = merged.squeeze(0)
+            core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
+            core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
         elif use_spec:
             assert core_attn_out_spec is not None
             if spec_out is None:
diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/common/model.py
similarity index 94%
rename from vllm/models/glm5next/nvidia/model.py
rename to vllm/models/glm5next/common/model.py
index 8444c8141611..1c826f58be3b 100644
--- a/vllm/models/glm5next/nvidia/model.py
+++ b/vllm/models/glm5next/common/model.py
@@ -355,6 +355,7 @@ def __init__(
                 intermediate_size=config.intermediate_size,
                 hidden_act=config.hidden_act,
                 quant_config=quant_config,
+                is_sequence_parallel=self.is_sequence_parallel,
                 prefix=f"{prefix}.mlp",
                 swiglu_limit=config.swiglu_limit,
             )
@@ -403,6 +404,50 @@ def __init__(
             self.mhc_post_op = MHCPostOp()
             self.mhc_fused_post_pre_op = MHCFusedPostPreOp()
 
+            if vllm_config.kernel_config.enable_jit_warmup:
+                from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+                    _HC_PRENORM_GEMM_TILELANG_KERNEL,
+                    _MHC_FUSED_TILELANG_KERNEL,
+                    _MHC_POST_TILELANG_KERNEL,
+                    _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
+                )
+                from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+                include_pre_gemm_splits = is_deep_gemm_supported()
+                _MHC_PRE_BIG_FUSE_TILELANG_KERNEL.register_warmup(
+                    vllm_config,
+                    hidden_size=self.hidden_size,
+                    hc_mult=self.n,
+                    use_norm_weight=True,
+                    include_pre_gemm_splits=include_pre_gemm_splits,
+                    include_broadcast_splits=False,
+                    rms_eps=self.rms_norm_eps,
+                    hc_pre_eps=self.hc_eps,
+                    hc_sinkhorn_eps=self.hc_eps,
+                    hc_post_mult_value=self.mhc_post_mult_value,
+                    sinkhorn_repeat=self.mhc_sinkhorn_iterations,
+                    norm_eps=(
+                        self.input_layernorm.variance_epsilon,
+                        self.post_attention_layernorm.variance_epsilon,
+                    ),
+                )
+                if not include_pre_gemm_splits:
+                    _HC_PRENORM_GEMM_TILELANG_KERNEL.register_warmup(
+                        vllm_config,
+                        hidden_size=self.hidden_size,
+                        hc_mult=self.n,
+                        n_out=self.n * (2 + self.n),
+                    )
+                _MHC_POST_TILELANG_KERNEL.register_warmup(
+                    hidden_size=self.hidden_size,
+                    hc_mult=self.n,
+                )
+                _MHC_FUSED_TILELANG_KERNEL.register_warmup(
+                    vllm_config,
+                    hidden_size=self.hidden_size,
+                    hc_mult=self.n,
+                )
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -992,6 +1037,17 @@ class Glm5NextForConditionalGeneration(
     has_inner_state: ClassVar[Literal[True]] = True
     is_hybrid: ClassVar[Literal[True]] = True
 
+    # GLM-5.3-Flash stores the dense-MLP gate/up as separate tensors (like
+    # ``Glm4vMoeForConditionalGeneration``, ``glm4_moe`` and ``deepseek_v2``),
+    # so the fused ``gate_up_proj`` must expand to its real shard names for
+    # per-layer quant-scheme resolution. The identity ``gate_up_proj`` entry
+    # inherited from ``Glm4vForConditionalGeneration`` (pre-fused gate_up_proj)
+    # would otherwise route the module to ``global_quant_config`` and mismatch
+    # at load for mixed-precision Quark checkpoints.
+    packed_modules_mapping = {
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
     # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration
     # (``model.visual.`` -> ``visual.``, ``model.language_model.`` ->
     # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``),
@@ -1215,7 +1271,9 @@ def _try_load_fp8_attn_proj(
         return False
     suffix, (key, target_base, shard_id, is_kva) = matched
     is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
-    is_scale = "weight_scale_inv" in name
+    # Need to accept both the DeepSeek-native ``weight_scale_inv`` and the Quark
+    # ``weight_scale`` names before feeding the shared block dequant below.
+    is_scale = "weight_scale_inv" in name or name.endswith(".weight_scale")
     if not is_weight and not is_scale:
         return False
 
diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/common/mtp.py
similarity index 99%
rename from vllm/models/glm5next/nvidia/mtp.py
rename to vllm/models/glm5next/common/mtp.py
index 35b452f44f79..8361a053b8f6 100644
--- a/vllm/models/glm5next/nvidia/mtp.py
+++ b/vllm/models/glm5next/common/mtp.py
@@ -22,6 +22,7 @@
 from vllm.model_executor.models.deepseek_mtp import SharedHead
 from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts
 from vllm.model_executor.models.utils import maybe_prefix
+from vllm.models.glm5next.nvidia.ops.fused_eh_norm import fused_eh_norm
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 
@@ -33,7 +34,6 @@
     _try_load_fp8_indexer_wk,
     get_spec_layer_idx_from_weight_name,
 )
-from .ops.fused_eh_norm import fused_eh_norm
 
 
 class Glm5NextMultiTokenPredictorLayer(nn.Module):
@@ -42,7 +42,6 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
         assert vllm_config.speculative_config is not None
         config = vllm_config.speculative_config.draft_model_config.hf_config
         self.config = config
-        quant_config = vllm_config.quant_config
 
         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -66,7 +65,9 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
             device=current_platform.device_type,
         )
         self.shared_head = SharedHead(
-            config=config, prefix=prefix, quant_config=quant_config
+            config=config,
+            prefix=prefix,
+            defer_lm_head=True,
         )
         # MTP layers sit past the base model's hidden layers; parse the index
         # from the prefix (e.g. "...layers.32") so the decoder builds an MLA
diff --git a/vllm/models/glm5next/nvidia/multimodal.py b/vllm/models/glm5next/common/multimodal.py
similarity index 100%
rename from vllm/models/glm5next/nvidia/multimodal.py
rename to vllm/models/glm5next/common/multimodal.py
diff --git a/vllm/models/glm5next/common/sparse_indexer.py b/vllm/models/glm5next/common/sparse_indexer.py
new file mode 100644
index 000000000000..ad3845457063
--- /dev/null
+++ b/vllm/models/glm5next/common/sparse_indexer.py
@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared helpers for the glm5next sparse attention indexer (kpool) layers."""
+
+import torch
+
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
+
+# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
+MXFP4_BLOCK_SIZE = 32
+
+
+def _build_decode_scatter_indices(
+    decode_lens: torch.Tensor,
+    num_requests: int,
+    n: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-token (request id, intra-request index) for a non-uniform decode
+    batch, with ``n == decode_lens.sum()`` as a host int (avoids a
+    device sync and keeps both repeat_interleaves sync-free).
+
+    Shared by every ``_scatter_decode_tokens_by_request`` call in a step:
+    building it per call would repeat the same repeat_interleave/cumsum chain
+    up to 5x per layer on the eager decode break.
+    """
+    device = decode_lens.device
+    dl = decode_lens.to(torch.int64)
+    req_id = torch.repeat_interleave(
+        torch.arange(num_requests, device=device, dtype=torch.int64),
+        dl,
+        output_size=n,
+    )
+    req_starts = torch.cumsum(
+        torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]),
+        dim=0,
+    )
+    # Broadcast the per-request start offsets to per-token (length n ==
+    # dl.sum()) so each token's intra-request index subtracts its own
+    # request's start.
+    starts = torch.repeat_interleave(req_starts, dl, output_size=n)
+    intra = torch.arange(n, device=device, dtype=torch.int64) - starts
+    return req_id, intra
+
+
+def _scatter_decode_tokens_by_request(
+    tokens: torch.Tensor,
+    pad_value,
+    num_requests: int,
+    lmax: int,
+    scatter_indices: tuple[torch.Tensor, torch.Tensor],
+) -> torch.Tensor:
+    """Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]``
+    layout: request ``r``'s tokens at row ``r`` in order; short requests padded.
+
+    Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32
+    slot/pos tensors) — it scatters with the shared per-step indices from
+    ``_build_decode_scatter_indices``. Used only for the non-uniform
+    (``requires_padding``) decode batch; uniform batches use a zero-copy
+    reshape.
+    """
+    req_id, intra = scatter_indices
+    out = torch.full(
+        (num_requests, lmax, *tokens.shape[1:]),
+        pad_value,
+        dtype=tokens.dtype,
+        device=tokens.device,
+    )
+    out[req_id, intra] = tokens
+    return out
+
+
+def _decode_topk_seq_lens(
+    positions: torch.Tensor,
+    decode_lens: torch.Tensor,
+    num_decode_tokens: int,
+    batch_size: int,
+    next_n: int,
+    requires_padding: bool,
+) -> torch.Tensor:
+    """Token-granular seq_len (pos + 1) per pool-topk row, layout-aware.
+
+    ``pool_topk`` (and the logits it comes from) follow the padded
+    ``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row
+    ``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT
+    ``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]``
+    (the uniform-layout shortcut) misaligns every row after the first
+    non-uniform request and, past the decode region, reads prefill tokens'
+    positions; ``expand_pools_and_append_tail`` then anchors the tail at
+    another request's length, dropping the row's real tail tokens or emitting
+    indices past its sequence (out-of-bounds block-table reads). Padded rows
+    get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway.
+    """
+    n = batch_size * next_n
+    if not requires_padding:
+        return positions[:n].to(torch.int32) + 1
+    scatter_idx = _build_decode_scatter_indices(
+        decode_lens, batch_size, num_decode_tokens
+    )
+    padded = _scatter_decode_tokens_by_request(
+        positions[:num_decode_tokens].to(torch.int32),
+        -1,
+        batch_size,
+        next_n,
+        scatter_idx,
+    )
+    return padded.reshape(n) + 1  # pad rows: -1 + 1 = 0 -> empty tail
+
+
+def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None:
+    causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32)
+    positions = positions.to(torch.int32)
+    rows[:] = causal_range[None, :]
+    rows[causal_range[None, :] > positions[:, None]] = -1
+
+
+def _fill_short_decode_causal_indices(
+    topk_indices_buffer: torch.Tensor,
+    positions: torch.Tensor | None,
+    num_decode_tokens: int,
+    max_seq_len: int,
+    topk_tokens: int,
+) -> bool:
+    """Fill exact causal rows when sparse decode would select every token."""
+    if positions is None or positions.numel() == 0 or max_seq_len > topk_tokens:
+        return False
+    _fill_causal_indices(
+        topk_indices_buffer[:num_decode_tokens], positions[:num_decode_tokens]
+    )
+    return True
+
+
+def _gather_workspace_shapes(
+    total_seq_lens: int,
+    head_dim: int,
+    fp8_dtype: torch.dtype,
+    use_fp4_cache: bool,
+) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]:
+    """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for
+    the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32
+    scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 +
+    (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales."""
+    if use_fp4_cache:
+        return (
+            ((total_seq_lens, head_dim // 2), torch.uint8),
+            ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8),
+        )
+    return (
+        ((total_seq_lens, head_dim), fp8_dtype),
+        ((total_seq_lens, 4), torch.uint8),
+    )
+
+
+def kv_cache_as_quant_view(
+    kv_cache: torch.Tensor,
+    head_dim: int,
+    use_fp4_cache: bool,
+) -> torch.Tensor:
+    """4D ``[num_blocks, block_size, 1, head_width]`` view expected by
+    DeepGEMM, from the 3D indexer kv-cache allocation."""
+    if use_fp4_cache:
+        assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8
+        num_blocks, block_size, _ = kv_cache.shape
+        page_bytes = int(kv_cache.stride(0))
+        fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
+        return torch.as_strided(
+            kv_cache,
+            size=(num_blocks, block_size, 1, fp4_bytes),
+            stride=(page_bytes, fp4_bytes, fp4_bytes, 1),
+        )
+    return kv_cache.unsqueeze(-2)
diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py
index 70b41021aab0..416d3b5c48ea 100644
--- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py
+++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py
@@ -116,6 +116,7 @@ def fwht128_quant_fp8(q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
 
     Returns:
         (q_fp8 ``[rows, 128]`` float8_e4m3fn, scale ``[rows, 1]`` float32).
+
     """
     assert q.ndim == 2 and q.shape[1] == 128, q.shape
     assert q.dtype == torch.bfloat16
@@ -276,6 +277,13 @@ def kpool_compress_and_write_cache(
         slot_score: ``[n_pools, pool_size, head_dim]`` — per-token gate score.
         ape: ``[pool_size, head_dim]`` fp32 — per-slot position bias.
         loc: ``[n_pools]`` int64 — flat physical slot per pool.
+        pool_size: Number of tokens compressed into one cache entry.
+        head_dim: Indexer head dimension.
+        write_mask: ``[n_pools]`` bool — pools to write, or None for all.
+        round_scale: Round each fp8 scale down to a power of two.
+        return_compressed: Also return the compressed K and scales.
+        write_cache: Write the compressed result into ``kv_cache``.
+
     """
     assert slot_k.ndim == 3
     assert slot_score.shape == slot_k.shape
@@ -375,6 +383,8 @@ def _kpool_tail_seed_kernel(
     tslot_ptr,
     tail_ptr,
     n_tokens,
+    TAIL_BLOCK_ELEMS: tl.constexpr,
+    KPOOL_HEAD: tl.constexpr,
     HEAD_DIM: tl.constexpr,
     KPOOL: tl.constexpr,
     BLOCK_D: tl.constexpr,
@@ -385,6 +395,11 @@ def _kpool_tail_seed_kernel(
     ahead belongs to a different tail block (or is past the batch / padding,
     slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is
     ``tail[block, {0:K, 1:score}, pos % KPOOL, :]``.
+
+    The tail cache aliases the indexer cache with the indexer's (padded) block
+    stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` /
+    ``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense
+    ``[num_blocks, 2, KPOOL, HEAD_DIM]`` array.
     """
     i = tl.program_id(0)
     t = tl.load(tslot_ptr + i).to(tl.int64)
@@ -401,11 +416,11 @@ def _kpool_tail_seed_kernel(
         return
     offs = tl.arange(0, BLOCK_D)
     m = offs < HEAD_DIM
-    base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM
+    base = blk * TAIL_BLOCK_ELEMS + (t % KPOOL) * HEAD_DIM
     k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m)
     s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m)
     tl.store(tail_ptr + base + offs, k, mask=m)
-    tl.store(tail_ptr + base + KPOOL * HEAD_DIM + offs, s, mask=m)
+    tl.store(tail_ptr + base + KPOOL_HEAD + offs, s, mask=m)
 
 
 def kpool_seed_tail_cache(
@@ -418,6 +433,8 @@ def kpool_seed_tail_cache(
 ) -> None:
     """Seed the paged tail cache from a prefill batch (see the kernel)."""
     assert tail_kv_cache.dtype == torch.bfloat16
+    assert tail_kv_cache.ndim == 4 and tail_kv_cache.shape[1] == 2
+    assert tail_kv_cache.stride(3) == 1 and tail_kv_cache.stride(2) == head_dim
     assert key.dtype == torch.bfloat16
     n = tslot.shape[0]
     if n == 0:
@@ -428,6 +445,8 @@ def kpool_seed_tail_cache(
         tslot,
         tail_kv_cache,
         n,
+        TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),
+        KPOOL_HEAD=tail_kv_cache.stride(1),
         HEAD_DIM=head_dim,
         KPOOL=kpool,
         BLOCK_D=triton.next_power_of_2(head_dim),
@@ -642,6 +661,10 @@ def kpool_decode_update_and_maybe_write_cache_batched(
         ape: ``[pool_size, head_dim]`` fp32.
         slot_mapping: ``[num_requests, next_n]`` int32.
         positions: ``[num_requests, next_n]`` int32.
+        pool_size: Number of tokens compressed into one cache entry.
+        head_dim: Indexer head dimension.
+        round_scale: Round each fp8 scale down to a power of two.
+
     """
     num_requests, next_n = key.shape[0], key.shape[1]
     if num_requests == 0 or next_n == 0:
diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py
index 152412b8ca8f..a9d0fba476ef 100644
--- a/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py
+++ b/vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py
@@ -590,8 +590,7 @@ def fused_recurrent_gated_delta_rule(
     num_accepted_tokens: torch.Tensor | None = None,
     use_qk_l2norm_in_kernel: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    r"""
-    Args:
+    r"""Args:
         q (torch.Tensor):
             queries of shape `[B, T, H, K]`.
         k (torch.Tensor):
@@ -653,6 +652,7 @@ def fused_recurrent_gated_delta_rule(
             initial_state=h0,
             cu_seqlens=cu_seqlens
         )
+
     """
     if cu_seqlens is not None and q.shape[0] != 1:
         raise ValueError(
diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py
index b792f606d4e0..9805b3747130 100644
--- a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py
+++ b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py
@@ -416,19 +416,24 @@ def chunk_kda_scaled_dot_kkt_fwd(
     chunk_size: int = FLA_CHUNK_SIZE,
     output_dtype: torch.dtype = torch.float32,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    r"""
-    Compute beta * K * K^T.
+    r"""Compute beta * K * K^T.
 
     Args:
+        q (torch.Tensor):
+            The query tensor of shape `[B, T, H, K]`.
         k (torch.Tensor):
             The key tensor of shape `[B, T, H, K]`.
         beta (torch.Tensor):
             The beta tensor of shape `[B, T, H]`.
         gk (torch.Tensor):
             The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.
+        scale (float):
+            Scale applied to the query-key products. Default: `None`.
         cu_seqlens (torch.Tensor):
             The cumulative sequence lengths of the input tensor.
             Default: None
+        chunk_indices (torch.Tensor):
+            Precomputed chunk indices for `cu_seqlens`. Default: `None`.
         chunk_size (int):
             The chunk size. Default: 64.
         output_dtype (torch.dtype):
@@ -436,6 +441,7 @@ def chunk_kda_scaled_dot_kkt_fwd(
 
     Returns:
         beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
+
     """
     B, T, H, K = k.shape
     assert K <= 256
@@ -1328,16 +1334,15 @@ def fused_kda_gate(
     safe_gate: bool = False,
     lower_bound: float | None = -5.0,
 ) -> torch.Tensor:
-    """
-    Forward pass for KDA gate:
-      input g: [..., H*D]
-      param A: [H] or [1, 1, H, 1]
-      beta: softplus beta parameter (softplus branch only)
-      threshold: softplus threshold parameter (softplus branch only)
-      safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias);
-        when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias))
-      lower_bound: floor for the safe_gate branch (default -5.0)
-      return  : [..., H, D]
+    """Forward pass for KDA gate:
+    input g: [..., H*D]
+    param A: [H] or [1, 1, H, 1]
+    beta: softplus beta parameter (softplus branch only)
+    threshold: softplus threshold parameter (softplus branch only)
+    safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias);
+      when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias))
+    lower_bound: floor for the safe_gate branch (default -5.0)
+    return  : [..., H, D]
     """
     orig_shape = g.shape[:-1]
 
diff --git a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py b/vllm/models/glm5next/nvidia/sparse_indexer.py
similarity index 65%
rename from vllm/model_executor/layers/sparse_attn_indexer_kpool.py
rename to vllm/models/glm5next/nvidia/sparse_indexer.py
index ba7a5adff160..c80e44201571 100644
--- a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py
+++ b/vllm/models/glm5next/nvidia/sparse_indexer.py
@@ -2,30 +2,30 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Custom Sparse Attention Indexer layers."""
 
-from typing import TYPE_CHECKING
-
 import torch
 
 import vllm.envs as envs
-from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import get_current_vllm_config_or_none
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.indexer_topk import get_indexer_topk
+from vllm.models.glm5next.common.sparse_indexer import (
+    RADIX_TOPK_WORKSPACE_SIZE,
+    _build_decode_scatter_indices,
+    _decode_topk_seq_lens,
+    _fill_causal_indices,
+    _fill_short_decode_causal_indices,
+    _gather_workspace_shapes,
+    _scatter_decode_tokens_by_request,
+    kv_cache_as_quant_view,
+)
+from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
 from vllm.platforms import current_platform
-
-if TYPE_CHECKING:
-    from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
-elif current_platform.is_rocm():
-    from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops
-else:
-    from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
-
 from vllm.utils.deep_gemm import has_deep_gemm
 from vllm.utils.torch_utils import (
     LayerNameType,
-    _encode_layer_name,
     _resolve_layer_name,
 )
 from vllm.v1.attention.backends.mla.indexer import (
@@ -36,15 +36,9 @@
 
 if current_platform.is_cuda_alike():
     from vllm import _custom_ops as ops
-elif current_platform.is_xpu():
-    from vllm._xpu_ops import xpu_ops
 
 logger = init_logger(__name__)
 
-RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
-
-# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
-MXFP4_BLOCK_SIZE = 32
 
 # kpool write helper: form pools from the current token batch and compress them
 # into the index K cache via the fused Triton kernel.
@@ -96,166 +90,6 @@ def _kpool_compress_insert(
     )
 
 
-def _build_decode_scatter_indices(
-    decode_lens: torch.Tensor,
-    num_requests: int,
-    n: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Per-token (request id, intra-request index) for a non-uniform decode
-    batch, with ``n == decode_lens.sum()`` as a host int (avoids a
-    device sync and keeps both repeat_interleaves sync-free).
-
-    Shared by every ``_scatter_decode_tokens_by_request`` call in a step:
-    building it per call would repeat the same repeat_interleave/cumsum chain
-    up to 5x per layer on the eager decode break.
-    """
-    device = decode_lens.device
-    dl = decode_lens.to(torch.int64)
-    req_id = torch.repeat_interleave(
-        torch.arange(num_requests, device=device, dtype=torch.int64),
-        dl,
-        output_size=n,
-    )
-    req_starts = torch.cumsum(
-        torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]),
-        dim=0,
-    )
-    # Broadcast the per-request start offsets to per-token (length n ==
-    # dl.sum()) so each token's intra-request index subtracts its own
-    # request's start.
-    starts = torch.repeat_interleave(req_starts, dl, output_size=n)
-    intra = torch.arange(n, device=device, dtype=torch.int64) - starts
-    return req_id, intra
-
-
-def _scatter_decode_tokens_by_request(
-    tokens: torch.Tensor,
-    pad_value,
-    num_requests: int,
-    lmax: int,
-    scatter_indices: tuple[torch.Tensor, torch.Tensor],
-) -> torch.Tensor:
-    """Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]``
-    layout: request ``r``'s tokens at row ``r`` in order; short requests padded.
-
-    Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32
-    slot/pos tensors) — it scatters with the shared per-step indices from
-    ``_build_decode_scatter_indices``. Used only for the non-uniform
-    (``requires_padding``) decode batch; uniform batches use a zero-copy
-    reshape.
-    """
-    req_id, intra = scatter_indices
-    out = torch.full(
-        (num_requests, lmax, *tokens.shape[1:]),
-        pad_value,
-        dtype=tokens.dtype,
-        device=tokens.device,
-    )
-    out[req_id, intra] = tokens
-    return out
-
-
-def _decode_topk_seq_lens(
-    positions: torch.Tensor,
-    decode_lens: torch.Tensor,
-    num_decode_tokens: int,
-    batch_size: int,
-    next_n: int,
-    requires_padding: bool,
-) -> torch.Tensor:
-    """Token-granular seq_len (pos + 1) per pool-topk row, layout-aware.
-
-    ``pool_topk`` (and the logits it comes from) follow the padded
-    ``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row
-    ``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT
-    ``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]``
-    (the uniform-layout shortcut) misaligns every row after the first
-    non-uniform request and, past the decode region, reads prefill tokens'
-    positions; ``expand_pools_and_append_tail`` then anchors the tail at
-    another request's length, dropping the row's real tail tokens or emitting
-    indices past its sequence (out-of-bounds block-table reads). Padded rows
-    get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway.
-    """
-    n = batch_size * next_n
-    if not requires_padding:
-        return positions[:n].to(torch.int32) + 1
-    scatter_idx = _build_decode_scatter_indices(
-        decode_lens, batch_size, num_decode_tokens
-    )
-    padded = _scatter_decode_tokens_by_request(
-        positions[:num_decode_tokens].to(torch.int32),
-        -1,
-        batch_size,
-        next_n,
-        scatter_idx,
-    )
-    return padded.reshape(n) + 1  # pad rows: -1 + 1 = 0 -> empty tail
-
-
-def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None:
-    causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32)
-    positions = positions.to(torch.int32)
-    rows[:] = causal_range[None, :]
-    rows[causal_range[None, :] > positions[:, None]] = -1
-
-
-def _fill_short_decode_causal_indices(
-    topk_indices_buffer: torch.Tensor,
-    positions: torch.Tensor | None,
-    num_decode_tokens: int,
-    max_seq_len: int,
-    topk_tokens: int,
-) -> bool:
-    """Fill exact causal rows when sparse decode would select every token."""
-    if positions is None or positions.numel() == 0 or max_seq_len > topk_tokens:
-        return False
-    _fill_causal_indices(
-        topk_indices_buffer[:num_decode_tokens], positions[:num_decode_tokens]
-    )
-    return True
-
-
-def _gather_workspace_shapes(
-    total_seq_lens: int,
-    head_dim: int,
-    fp8_dtype: torch.dtype,
-    use_fp4_cache: bool,
-) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]:
-    """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for
-    the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32
-    scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 +
-    (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales."""
-    if use_fp4_cache:
-        return (
-            ((total_seq_lens, head_dim // 2), torch.uint8),
-            ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8),
-        )
-    return (
-        ((total_seq_lens, head_dim), fp8_dtype),
-        ((total_seq_lens, 4), torch.uint8),
-    )
-
-
-def kv_cache_as_quant_view(
-    kv_cache: torch.Tensor,
-    head_dim: int,
-    use_fp4_cache: bool,
-) -> torch.Tensor:
-    """4D ``[num_blocks, block_size, 1, head_width]`` view expected by
-    DeepGEMM, from the 3D indexer kv-cache allocation."""
-    if use_fp4_cache:
-        assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8
-        num_blocks, block_size, _ = kv_cache.shape
-        page_bytes = int(kv_cache.stride(0))
-        fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
-        return torch.as_strided(
-            kv_cache,
-            size=(num_blocks, block_size, 1, fp4_bytes),
-            stride=(page_bytes, fp4_bytes, fp4_bytes, 1),
-        )
-    return kv_cache.unsqueeze(-2)
-
-
 @eager_break_during_capture
 def sparse_attn_indexer_kpool(
     hidden_states: torch.Tensor,
@@ -269,7 +103,7 @@ def sparse_attn_indexer_kpool(
     scale_fmt: str | None,
     topk_tokens: int,
     head_dim: int,
-    max_model_len: int,
+    max_pool_len: int,
     total_seq_lens: int,
     topk_indices_buffer: torch.Tensor,
     skip_k_cache_insert: bool,
@@ -286,6 +120,7 @@ def sparse_attn_indexer_kpool(
     # path and when the tail cache is disabled.
     tail_kv_cache: torch.Tensor | None = None,
     tail_prefix: str | None = None,
+    topk_backend: str = "auto",
 ) -> torch.Tensor:
     # careful! this will be None in dummy run
     attn_metadata = get_forward_context().attn_metadata
@@ -305,7 +140,7 @@ def sparse_attn_indexer_kpool(
         )
 
         # Reserve profiler-visible memory for the worst-case decode logits,
-        # whose shape is [B * next_n, max_model_len]. This profiling branch
+        # whose shape is [B * next_n, max_pool_len]. This profiling branch
         # returns before invoking the logits kernel itself.
         cfg = get_current_vllm_config_or_none()
         worst_decode_tokens = 0
@@ -321,7 +156,7 @@ def sparse_attn_indexer_kpool(
                 sched.max_num_batched_tokens,
             )
         # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
-        decode_logits_elems = worst_decode_tokens * max_model_len * 4
+        decode_logits_elems = worst_decode_tokens * max_pool_len * 4
         prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
         max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
         _ = torch.empty(
@@ -393,26 +228,13 @@ def sparse_attn_indexer_kpool(
         else:
             # standard: per-token fp8 quant + scatter (all tokens).
             assert scale_fmt is not None
-            if current_platform.is_rocm():
-                from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
-                    indexer_k_quant_and_cache_triton,
-                )
-
-                indexer_k_quant_and_cache_triton(
-                    k,
-                    kv_cache,
-                    slot_mapping,
-                    quant_block_size,
-                    scale_fmt,
-                )
-            else:
-                ops.indexer_k_quant_and_cache(
-                    k,
-                    kv_cache,
-                    slot_mapping,
-                    quant_block_size,
-                    scale_fmt,
-                )
+            ops.indexer_k_quant_and_cache(
+                k,
+                kv_cache,
+                slot_mapping,
+                quant_block_size,
+                scale_fmt,
+            )
 
     topk_indices_buffer[: hidden_states.shape[0]] = -1
     if has_prefill:
@@ -466,27 +288,13 @@ def sparse_attn_indexer_kpool(
             k_scale = k_scale_full[: chunk.total_seq_lens]
 
             if not chunk.skip_kv_gather:
-                if current_platform.is_rocm():
-                    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
-                        cp_gather_indexer_k_quant_cache_triton,
-                    )
-
-                    cp_gather_indexer_k_quant_cache_triton(
-                        kv_cache,
-                        k_quant,
-                        k_scale,
-                        chunk.block_table,
-                        chunk.cu_seq_lens,
-                        token_to_seq=chunk.token_to_seq,
-                    )
-                else:
-                    ops.cp_gather_indexer_k_quant_cache(
-                        kv_cache,
-                        k_quant,
-                        k_scale,
-                        chunk.block_table,
-                        chunk.cu_seq_lens,
-                    )
+                ops.cp_gather_indexer_k_quant_cache(
+                    kv_cache,
+                    k_quant,
+                    k_scale,
+                    chunk.block_table,
+                    chunk.cu_seq_lens,
+                )
 
             q_slice = q_quant[chunk.token_start : chunk.token_end]
             q_scale_slice = (
@@ -504,30 +312,16 @@ def sparse_attn_indexer_kpool(
                 q_slice_cast = q_slice
                 k_quant_cast = k_quant
                 k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
-            if current_platform.is_rocm():
-                from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
-                    rocm_fp8_mqa_logits,
-                )
-
-                assert q_scale_slice is None
-                logits = rocm_fp8_mqa_logits(
-                    q_slice_cast,
-                    (k_quant_cast, k_scale_cast),
-                    weights[chunk.token_start : chunk.token_end],
-                    chunk.cu_seqlen_ks,
-                    chunk.cu_seqlen_ke,
-                )
-            else:
-                from vllm.utils.deep_gemm import fp8_fp4_mqa_logits
-
-                logits = fp8_fp4_mqa_logits(
-                    (q_slice_cast, q_scale_slice),
-                    (k_quant_cast, k_scale_cast),
-                    weights[chunk.token_start : chunk.token_end],
-                    chunk.cu_seqlen_ks,
-                    chunk.cu_seqlen_ke,
-                    clean_logits=False,
-                )
+            from vllm.utils.deep_gemm import fp8_fp4_mqa_logits
+
+            logits = fp8_fp4_mqa_logits(
+                (q_slice_cast, q_scale_slice),
+                (k_quant_cast, k_scale_cast),
+                weights[chunk.token_start : chunk.token_end],
+                chunk.cu_seqlen_ks,
+                chunk.cu_seqlen_ke,
+                clean_logits=False,
+            )
             num_rows = logits.shape[0]
 
             # kpool: logits are pool-granular (compress_ratio == index_kpool),
@@ -544,28 +338,16 @@ def sparse_attn_indexer_kpool(
                     chunk.token_start : chunk.token_end, :topk_tokens
                 ]
 
-            if current_platform.is_xpu():
-                xpu_ops.top_k_per_row_prefill(  # type: ignore[attr-defined]
-                    logits,
-                    chunk.cu_seqlen_ks,
-                    chunk.cu_seqlen_ke,
-                    topk_dst,
-                    num_rows,
-                    logits.stride(0),
-                    logits.stride(1),
-                    select_k,
-                )
-            else:
-                torch.ops._C.top_k_per_row_prefill(
-                    logits,
-                    chunk.cu_seqlen_ks,
-                    chunk.cu_seqlen_ke,
-                    topk_dst,
-                    num_rows,
-                    logits.stride(0),
-                    logits.stride(1),
-                    select_k,
-                )
+            torch.ops._C.top_k_per_row_prefill(
+                logits,
+                chunk.cu_seqlen_ks,
+                chunk.cu_seqlen_ke,
+                topk_dst,
+                num_rows,
+                logits.stride(0),
+                logits.stride(1),
+                select_k,
+            )
 
             if index_kpool > 1:
                 pool_ids = pool_topk.to(torch.int64)
@@ -765,34 +547,18 @@ def sparse_attn_indexer_kpool(
             if use_fp4_cache
             else padded_q_quant_decode_tokens
         )
-        if current_platform.is_rocm():
-            from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
-                rocm_fp8_paged_mqa_logits,
-            )
-
-            assert padded_q_scale is None
-            logits = rocm_fp8_paged_mqa_logits(
-                padded_q_quant_cast,
-                kv_cache,
-                padded_weights[:num_padded_tokens],
-                seq_lens,
-                decode_metadata.block_table,
-                decode_metadata.schedule_metadata,
-                max_model_len=max_model_len,
-            )
-        else:
-            from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits
+        from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits
 
-            logits = fp8_fp4_paged_mqa_logits(
-                (padded_q_quant_cast, padded_q_scale),
-                kv_cache,
-                padded_weights[:num_padded_tokens],
-                seq_lens,
-                decode_metadata.block_table,
-                decode_metadata.schedule_metadata,
-                max_model_len=max_model_len,
-                clean_logits=False,
-            )
+        logits = fp8_fp4_paged_mqa_logits(
+            (padded_q_quant_cast, padded_q_scale),
+            kv_cache,
+            padded_weights[:num_padded_tokens],
+            seq_lens,
+            decode_metadata.block_table,
+            decode_metadata.schedule_metadata,
+            max_model_len=max_pool_len,
+            clean_logits=False,
+        )
         num_rows = logits.shape[0]
         # kpool: logits are pool-granular -> select topk_tokens//kpool pools,
         # then expand each pool back to its kpool tokens.
@@ -805,42 +571,18 @@ def sparse_attn_indexer_kpool(
         else:
             topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
 
-        if current_platform.is_cuda() and select_k in (512, 1024, 2048):
-            workspace_manager = current_workspace_manager()
-            (topk_workspace,) = workspace_manager.get_simultaneous(
-                ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-            )
-            torch.ops._C.persistent_topk(
-                logits,
-                seq_lens,
-                topk_dst,
-                topk_workspace,
-                select_k,
-                attn_metadata_narrowed.max_seq_len,
-            )
-        else:
-            if current_platform.is_xpu():
-                xpu_ops.top_k_per_row_decode(  # type: ignore[attr-defined]
-                    logits,
-                    next_n,
-                    seq_lens,
-                    topk_dst,
-                    num_rows,
-                    logits.stride(0),
-                    logits.stride(1),
-                    select_k,
-                )
-            else:
-                torch.ops._C.top_k_per_row_decode(
-                    logits,
-                    next_n,
-                    seq_lens,
-                    topk_dst,
-                    num_rows,
-                    logits.stride(0),
-                    logits.stride(1),
-                    select_k,
-                )
+        # Shared dispatcher with the DSA sparse indexer, so the kpool select
+        # honors kernel_config.sparse_indexer_topk_backend and picks between
+        # cooperative/persistent by batch size instead of always taking the
+        # same kernel.
+        get_indexer_topk(topk_backend)(
+            logits,
+            seq_lens,
+            next_n,
+            topk_dst,
+            select_k,
+            attn_metadata_narrowed.max_seq_len,
+        )
 
         # Resolve to token-level indices in the output buffer.
         if index_kpool > 1:
@@ -896,7 +638,7 @@ def __init__(
         scale_fmt: str,
         topk_tokens: int,
         head_dim: int,
-        max_model_len: int,
+        max_pool_len: int,
         max_total_seq_len: int,
         topk_indices_buffer: torch.Tensor,
         skip_k_cache_insert: bool = False,
@@ -910,7 +652,7 @@ def __init__(
         self.scale_fmt = scale_fmt
         self.topk_tokens = topk_tokens
         self.head_dim = head_dim
-        self.max_model_len = max_model_len
+        self.max_pool_len = max_pool_len
         self.max_total_seq_len = max_total_seq_len
         self.topk_indices_buffer = topk_indices_buffer
         self.skip_k_cache_insert = skip_k_cache_insert
@@ -919,6 +661,21 @@ def __init__(
             raise RuntimeError(
                 "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed."
             )
+        _cfg = get_current_vllm_config_or_none()
+        self.topk_backend = (
+            _cfg.kernel_config.sparse_indexer_topk_backend
+            if _cfg is not None
+            else "auto"
+        )
+        _parallel = _cfg.parallel_config if _cfg is not None else None
+        if (
+            _parallel is not None
+            and _parallel.prefill_context_parallel_size > 1
+            and _parallel.decode_context_parallel_size > 1
+        ):
+            raise NotImplementedError(
+                "SparseAttnIndexerKpool does not support PCP+DCP."
+            )
 
     def forward_native(
         self,
@@ -932,33 +689,16 @@ def forward_native(
         index_kpool: int = 1,
         positions: torch.Tensor | None = None,
     ):
-        if current_platform.is_cuda() or current_platform.is_xpu():
-            return self.forward_cuda(
-                hidden_states,
-                q_quant,
-                k,
-                weights,
-                gate_score=gate_score,
-                compress_ape=compress_ape,
-                index_kpool=index_kpool,
-                positions=positions,
-            )
-        elif current_platform.is_rocm():
-            return self.forward_hip(
-                hidden_states,
-                q_quant,
-                k,
-                weights,
-                gate_score=gate_score,
-                compress_ape=compress_ape,
-                index_kpool=index_kpool,
-                positions=positions,
-            )
-        else:
-            raise NotImplementedError(
-                "SparseAttnIndexer native forward is only implemented for "
-                "CUDA, ROCm and XPU platforms."
-            )
+        return self.forward_cuda(
+            hidden_states,
+            q_quant,
+            k,
+            weights,
+            gate_score=gate_score,
+            compress_ape=compress_ape,
+            index_kpool=index_kpool,
+            positions=positions,
+        )
 
     def forward_cuda(
         self,
@@ -990,7 +730,7 @@ def forward_cuda(
             self.scale_fmt,
             self.topk_tokens,
             self.head_dim,
-            self.max_model_len,
+            self.max_pool_len,
             self.max_total_seq_len,
             self.topk_indices_buffer,
             self.skip_k_cache_insert,
@@ -1001,53 +741,5 @@ def forward_cuda(
             positions,
             self.tail_cache.kv_cache if self.tail_cache is not None else None,
             self.tail_cache.prefix if self.tail_cache is not None else None,
-        )
-
-    def forward_hip(
-        self,
-        hidden_states: torch.Tensor,
-        q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
-        k: torch.Tensor,
-        weights: torch.Tensor,
-        *,
-        gate_score: torch.Tensor | None = None,
-        compress_ape: torch.Tensor | None = None,
-        index_kpool: int = 1,
-        positions: torch.Tensor | None = None,
-    ):
-        assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
-        assert isinstance(q_quant, torch.Tensor), (
-            "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
-        )
-        if rocm_aiter_ops.is_enabled():
-            if index_kpool <= 1:
-                return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
-                    hidden_states,
-                    _encode_layer_name(self.k_cache.prefix),
-                    self.k_cache.kv_cache,
-                    q_quant,
-                    k,
-                    weights,
-                    self.quant_block_size,
-                    self.scale_fmt,
-                    self.topk_tokens,
-                    self.head_dim,
-                    self.max_model_len,
-                    self.max_total_seq_len,
-                    self.topk_indices_buffer,
-                    skip_k_cache_insert=self.skip_k_cache_insert,
-                )
-            return self.forward_cuda(
-                hidden_states,
-                q_quant,
-                k,
-                weights,
-                gate_score=gate_score,
-                compress_ape=compress_ape,
-                index_kpool=index_kpool,
-                positions=positions,
-            )
-        raise RuntimeError(
-            "Sparse attention indexer ROCm path is only supported on AITER. "
-            "Please enable aiter with VLLM_ROCM_USE_AITER=1"
+            self.topk_backend,
         )
diff --git a/vllm/models/glm5next/sparse_indexer.py b/vllm/models/glm5next/sparse_indexer.py
new file mode 100644
index 000000000000..5aa0e47367f5
--- /dev/null
+++ b/vllm/models/glm5next/sparse_indexer.py
@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sparse attention indexer (kpool) — hardware-isolated entry point.
+
+The implementation lives under ``nvidia/`` and ``amd/`` (with shared helpers
+in ``common/``); this module picks the right one for the current platform.
+"""
+
+from vllm.platforms import current_platform
+
+if current_platform.is_rocm():
+    from .amd.sparse_indexer import (
+        SparseAttnIndexerKpool,
+    )
+else:
+    from .nvidia.sparse_indexer import (
+        SparseAttnIndexerKpool,
+    )
+
+__all__ = [
+    "SparseAttnIndexerKpool",
+]
diff --git a/vllm/models/hy_v4/__init__.py b/vllm/models/hy_v4/__init__.py
index f09f25bfb21f..f6ffafb1c516 100644
--- a/vllm/models/hy_v4/__init__.py
+++ b/vllm/models/hy_v4/__init__.py
@@ -14,20 +14,27 @@
 only public entry point and dispatches on the current platform, so registry
 entries never reach into a platform subpackage.
 
-NVIDIA and ROCm are supported. The port also drops the reference
-implementation's HPC/TPCP fusion paths, which depend on infrastructure that
-does not exist in this tree.
+NVIDIA and ROCm are supported, the latter through `amd/`. The port also drops
+the reference implementation's HPC/TPCP fusion paths, which depend on
+infrastructure that does not exist in this tree.
 """
 
+from typing import TYPE_CHECKING
+
 from vllm.platforms import current_platform
 
 if current_platform.is_xpu():
     raise NotImplementedError("hy_v4 does not yet support XPU.")
-else:
-    # The implementation is shared by CUDA and ROCm; platform-specific
-    # attention behavior is selected by the normal attention backend router.
+
+# The NVIDIA branch is the static default that type-checkers see; the ROCm
+# branch overrides it at runtime (kept type-compatible via type: ignore).
+if TYPE_CHECKING or not current_platform.is_rocm():
     from .nvidia.model import HYV4ForCausalLM
-    from .nvidia.mtp import HYV4MTP
+else:
+    from .amd.model import HYV4ForCausalLM  # type: ignore[assignment]
+
+# The draft head has no ROCm-specific implementation yet.
+from .nvidia.mtp import HYV4MTP  # noqa: E402
 
 __all__ = [
     "HYV4ForCausalLM",
diff --git a/vllm/models/hy_v4/amd/__init__.py b/vllm/models/hy_v4/amd/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/models/hy_v4/amd/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/models/hy_v4/amd/model.py b/vllm/models/hy_v4/amd/model.py
new file mode 100644
index 000000000000..772e4e5b74ff
--- /dev/null
+++ b/vllm/models/hy_v4/amd/model.py
@@ -0,0 +1,718 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Adapted from
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only HY V4 model compatible with HuggingFace weights (ROCm)."""
+
+import typing
+from collections.abc import Callable, Iterable, MutableSequence, Sequence
+from itertools import islice
+
+import regex as re
+import torch
+from torch import nn
+from transformers import PreTrainedConfig
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.deepseek_v2 import _try_load_fp8_indexer_wk
+from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    get_pp_missing_layer_names,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+
+from ..nvidia.attention import (
+    HYV4MLAAttention,
+    compute_skip_topk_layers,
+    is_skip_topk_indexer_weight,
+)
+from ..nvidia.hc import HYV4HCHeadLayer, HYV4HCLayer
+from ..nvidia.moe import HYV4FeedForward, HYV4MoEFused
+
+logger = init_logger(__name__)
+
+
+def _normalize_hyv4_config(config: PreTrainedConfig) -> PreTrainedConfig:
+    """Populate the aliases consumed by the shared MoE implementation."""
+    config.router_scaling_factor = config.routed_scaling_factor
+    config.num_experts = config.n_routed_experts
+    config.expert_hidden_dim = config.moe_intermediate_size
+    config.num_shared_experts = config.n_shared_experts
+    config.route_norm = config.norm_topk_prob
+    return config
+
+
+class HYV4DecoderLayer(nn.Module):
+    """One HY V4 decoder layer: MLA attention plus a dense or MoE MLP.
+
+    When``config.enable_ihc`` is set the layer runs on ``hc_mult`` residual
+    channels and each sub-block is wrapped by an `HYV4HCLayer` boundary;
+    otherwise it uses the standard single-stream residual.
+    """
+
+    def __init__(
+        self,
+        config: PreTrainedConfig,
+        vllm_config: VllmConfig,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        topk_indices_buffer: torch.Tensor | None = None,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.config = config
+        self.enable_ihc = getattr(config, "enable_ihc", False)
+        layer_idx = int(prefix.split(".")[-1])
+        self.layer_idx = layer_idx
+
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        self.self_attn = HYV4MLAAttention(
+            vllm_config=vllm_config,
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            qk_nope_head_dim=config.qk_nope_head_dim,
+            qk_rope_head_dim=config.qk_rope_head_dim,
+            v_head_dim=config.v_head_dim,
+            q_lora_rank=config.q_lora_rank,
+            kv_lora_rank=config.kv_lora_rank,
+            max_position_embeddings=max_position_embeddings,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+            layer_idx=layer_idx,
+            topk_indices_buffer=topk_indices_buffer,
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        if config.mlp_layer_types[layer_idx] == "dense":
+            self.mlp = HYV4FeedForward(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+            self.block_type = "feedforward"
+        else:
+            self.mlp = HYV4MoEFused(
+                config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
+            )
+            self.block_type = "moe"
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.hc_attn_layer = HYV4HCLayer(
+            config, layer_idx, prefix=f"{prefix}.hc_attn_layer"
+        )
+        self.hc_mlp_layer = HYV4HCLayer(
+            config, layer_idx, prefix=f"{prefix}.hc_mlp_layer"
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if self.enable_ihc:
+            return self._forward_ihc(positions, hidden_states)
+        return self._forward_normal(positions, hidden_states, residual)
+
+    def _forward_normal(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Standard single-residual-stream forward (iHC disabled)."""
+        if residual is not None:
+            hidden_states = hidden_states + residual
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        hidden_states = hidden_states + residual
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+
+        hidden_states = self.mlp(hidden_states)
+
+        return hidden_states, residual
+
+    def _forward_ihc(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, None]:
+        """IHC forward: each sub-block reduces and re-scatters the channels."""
+        hidden_states = self.hc_attn_layer.prepare_input(hidden_states)
+        hidden_states, post_gates, residual = self.hc_attn_layer.pre(hidden_states)
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+        hidden_states = self.hc_attn_layer.post(hidden_states, residual, post_gates)
+
+        hidden_states = self.hc_mlp_layer.prepare_input(hidden_states)
+        hidden_states, post_gates, residual = self.hc_mlp_layer.pre(hidden_states)
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.hc_mlp_layer.post(hidden_states, residual, post_gates)
+
+        # Under iHC the residual is carried inside hidden_states.
+        return hidden_states, None
+
+
+@support_torch_compile
+class HYV4Model(nn.Module):
+    """HY V4 backbone."""
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = _normalize_hyv4_config(vllm_config.model_config.hf_config)
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        parallel_config = vllm_config.parallel_config
+        eplb_config = parallel_config.eplb_config
+        self.num_redundant_experts = eplb_config.num_redundant_experts
+        self.device = current_platform.device_type
+        self.vocab_size = config.vocab_size
+        self.is_sparse = hasattr(config, "index_topk")
+        if self.is_sparse:
+            self.topk_indices_buffer = torch.empty(
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                config.index_topk,
+                dtype=torch.int32,
+                device=self.device,
+            )
+        else:
+            self.topk_indices_buffer = None
+        self.config = config
+        self.quant_config = quant_config
+        self.enable_ihc = getattr(config, "enable_ihc", False)
+
+        if get_pp_group().is_first_rank or (
+            self.config.tie_word_embeddings and get_pp_group().is_last_rank
+        ):
+            self.embed_tokens = VocabParallelEmbedding(
+                config.vocab_size,
+                config.hidden_size,
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: HYV4DecoderLayer(
+                config=config,
+                vllm_config=vllm_config,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=prefix,
+                topk_indices_buffer=self.topk_indices_buffer,
+            ),
+            prefix=f"{prefix}.layers",
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+        if self.enable_ihc:
+            # iHC head layer: merge the residual channels into one stream.
+            if get_pp_group().is_last_rank:
+                self.hc_head = HYV4HCHeadLayer(
+                    config,
+                    hidden_size=config.hidden_size,
+                    hc_mult=config.hc_mult,
+                    hc_eps=config.hc_eps,
+                    prefix=f"{prefix}.hc_head",
+                )
+            else:
+                self.hc_head = PPMissingLayer()
+            self.make_empty_intermediate_tensors = (
+                make_empty_intermediate_tensors_factory(
+                    ["hidden_states"], config.hc_mult * config.hidden_size
+                )
+            )
+        else:
+            self.make_empty_intermediate_tensors = (
+                make_empty_intermediate_tensors_factory(
+                    ["hidden_states", "residual"], config.hidden_size
+                )
+            )
+
+        # MoE hyperparameters (consumed by EPLB).
+        self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = []
+        self.num_expert_groups = 1
+        self.moe_layers: list[nn.Module] = []
+        example_layer: HYV4MoEFused | None = None
+        for layer in self.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+
+            assert isinstance(layer, HYV4DecoderLayer)
+            if layer.block_type == "moe":
+                assert isinstance(layer.mlp, HYV4MoEFused)
+                example_layer = layer.mlp
+                self.moe_layers.append(layer.mlp.experts)
+
+        if example_layer is None:
+            self.num_moe_layers = 0
+            self.num_logical_experts = 0
+            self.num_physical_experts = 0
+            self.num_local_physical_experts = 0
+            self.num_routed_experts = 0
+            self.num_redundant_experts = 0
+            return
+
+        self.num_moe_layers = len(self.moe_layers)
+        self.num_logical_experts = example_layer.n_logical_experts
+        self.num_physical_experts = example_layer.n_physical_experts
+        self.num_local_physical_experts = example_layer.n_local_physical_experts
+        self.num_routed_experts = example_layer.n_routed_experts
+        self.num_redundant_experts = example_layer.n_redundant_experts
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ) -> None:
+        assert self.num_local_physical_experts == num_local_physical_experts
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+        for layer in self.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+            if isinstance(layer.mlp, HYV4MoEFused):
+                moe = layer.mlp
+                moe.n_local_physical_experts = num_local_physical_experts
+                moe.n_physical_experts = num_physical_experts
+                moe.n_redundant_experts = self.num_redundant_experts
+                moe.experts.update_expert_map()
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        # (param_name, weight_name, expert_id, shard_id) for weights, fp8
+        # weight scales and fp8 activation scales.
+        if not hasattr(self, "_cached_expert_params_mapping"):
+            self._cached_expert_params_mapping = fused_moe_make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=self.config.num_experts,
+                num_redundant_experts=self.num_redundant_experts,
+            )
+        return self._cached_expert_params_mapping
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            # In iHC mode the flattened [num_tokens, hc*h] tensor from the
+            # previous PP stage is reshaped back to 3D by the first layer's
+            # prepare_input, and residual is unused.
+            residual = None if self.enable_ihc else intermediate_tensors["residual"]
+
+        for layer in islice(self.layers, self.start_layer, self.end_layer):
+            hidden_states, residual = layer(positions, hidden_states, residual)
+
+        if not get_pp_group().is_last_rank:
+            if self.enable_ihc:
+                # hidden_states is [num_tokens, hc, h]; flatten the channel dim
+                # for PP transfer (matches the 2D receive buffer).
+                return IntermediateTensors({"hidden_states": hidden_states.flatten(1)})
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+        if self.enable_ihc:
+            hidden_states = self.hc_head(hidden_states)
+        else:
+            hidden_states = hidden_states + residual
+
+        return self.norm(hidden_states)
+
+    def load_fused_expert_weights(
+        self,
+        name: str,
+        params_dict: dict,
+        loaded_weight: torch.Tensor,
+        shard_id: str,
+        num_experts: int,
+    ) -> bool:
+        param = params_dict[name]
+        weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
+        loaded_local_expert = False
+        for expert_id in range(num_experts):
+            curr_expert_weight = loaded_weight[expert_id]
+            success = weight_loader(
+                param,
+                curr_expert_weight,
+                name,
+                shard_id,
+                expert_id,
+                return_success=True,
+            )
+            if success:
+                loaded_local_expert = True
+
+        return loaded_local_expert
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # FP8 indexer wk dequant buffer (weight and scale arrive separately).
+        pending_wk_fp8: dict[str, dict[str, torch.Tensor]] = {}
+        pp_missing_layer_names = get_pp_missing_layer_names(self)
+        skip_topk_layers = compute_skip_topk_layers(self.config)
+
+        # Must not be cached on `self`: `process_weights_after_loading` swaps in
+        # kernel-specific expert weights via `replace_parameter`, and a cache
+        # outliving this call would pin the pre-shuffle storage (OOM on large
+        # MoE) and make a later reload target orphaned tensors.
+        params_dict = dict(self.named_parameters())
+        # Split per-expert mapping (V3 style): experts.0.gate_proj.weight
+        split_expert_params_mapping = self.get_expert_mapping()
+        loaded_params: set[str] = set()
+
+        # Sink weights are sharded like the q/k/v linears.
+        sink_tp_size = get_tensor_model_parallel_world_size()
+        sink_tp_rank = get_tensor_model_parallel_rank()
+        n_local_head = self.config.num_attention_heads // sink_tp_size
+        head_rank_start = n_local_head * sink_tp_rank
+        head_rank_end = n_local_head * (sink_tp_rank + 1)
+        base_layer = (
+            "base_layer." if any(".base_layer." in name for name in params_dict) else ""
+        )
+        # Fused expert mapping: experts.gate_up_proj (all experts in one tensor).
+        # The packed weights are owned by the RoutedExperts submodule, so the
+        # targets are experts.routed_experts.[base_layer.]w{13,2}_weight.
+        fused_expert_prefix = f".experts.routed_experts.{base_layer}"
+        fused_expert_params_mapping = [
+            (f"{fused_expert_prefix}w13_weight", ".experts.gate_up_proj", 0, "w1"),
+            (f"{fused_expert_prefix}w2_weight", ".experts.down_proj", 0, "w2"),
+        ]
+        num_experts = getattr(self.config, "num_experts", 0)
+
+        def _should_skip_missing_param(param_name: str) -> bool:
+            # Sparse checkpoints may contain indexer weights for layers that
+            # fall back to dense MLA on unsupported GPUs.
+            return ".indexer." in param_name
+
+        def _is_split_expert_weight(weight_name: str) -> bool:
+            """Whether this weight is in split (per-expert) format."""
+            # Split format: mlp.experts.<id>.gate_proj.weight
+            return bool(re.search(r"\.experts\.\d+\.", weight_name))
+
+        def _is_fused_expert_weight(weight_name: str) -> bool:
+            """Whether this weight is in fused (all-experts-packed) format."""
+            return ".experts.gate_up_proj" in weight_name or (
+                ".experts.down_proj" in weight_name
+                and not _is_split_expert_weight(weight_name)
+            )
+
+        for name, loaded_weight in weights:
+            if is_skip_topk_indexer_weight(name, skip_topk_layers):
+                continue
+            if _try_load_fp8_indexer_wk(
+                name,
+                loaded_weight,
+                pending_wk_fp8,
+                params_dict,
+                loaded_params,
+                pp_missing_layer_names,
+            ):
+                continue
+            if self.config.tie_word_embeddings and "lm_head.weight" in name:
+                continue
+            # Names and stacked shard ids are normalized upstream by
+            # AutoWeightsLoader, through quant_config.get_cache_scale_mapper()
+            # and hf_to_vllm_mapper.
+            stacked_shard_id = getattr(loaded_weight, "shard_id", None)
+            if ".indexer.wk." in name:
+                name = name.replace(".indexer.wk.", ".indexer.wk_weights_proj.")
+                stacked_shard_id = 0
+            # Skip loading extra bias for GPTQ models.
+            if name.endswith(".bias") and name not in params_dict:
+                continue
+
+            # Determine per-weight whether this is fused or split format.
+            is_fused_expert = _is_fused_expert_weight(name)
+            expert_params_mapping = (
+                fused_expert_params_mapping
+                if is_fused_expert
+                else split_expert_params_mapping
+            )
+
+            is_expert_weight = False
+            loaded_expert_param_names: set[str] = set()
+            for mapping in expert_params_mapping:
+                param_name, weight_name, expert_id, shard_id = mapping
+                if weight_name not in name:
+                    continue
+
+                # This is an expert weight and must not be attempted as any
+                # other kind of weight later on.
+                is_expert_weight = True
+
+                # Do not modify `name`: the loop may continue past this point.
+                name_mapped = name.replace(weight_name, param_name)
+                if is_pp_missing_parameter(name_mapped, self):
+                    continue
+                if is_fused_expert:
+                    if "experts.gate_up_proj" in name:
+                        chunks = loaded_weight.chunk(2, dim=-2)
+                        success_w1 = self.load_fused_expert_weights(
+                            name_mapped, params_dict, chunks[0], "w1", num_experts
+                        )
+                        success_w3 = self.load_fused_expert_weights(
+                            name_mapped, params_dict, chunks[1], "w3", num_experts
+                        )
+                        success = success_w1 and success_w3
+                    else:
+                        success = self.load_fused_expert_weights(
+                            name_mapped,
+                            params_dict,
+                            loaded_weight,
+                            shard_id,
+                            num_experts,
+                        )
+                    if success:
+                        name = name_mapped
+                        break
+                else:
+                    # Split per-expert format (V3 style).
+                    if name_mapped not in params_dict:
+                        if _should_skip_missing_param(name_mapped):
+                            continue
+                        logger.warning_once(
+                            "Skipping unknown checkpoint weight: %s",
+                            name_mapped,
+                        )
+                        continue
+                    param = params_dict[name_mapped]
+                    # Ask the weight loader whether it succeeded, otherwise we
+                    # may skip experts that have other available replicas.
+                    weight_loader: Callable[..., typing.Any] = param.weight_loader
+                    success = weight_loader(
+                        param,
+                        loaded_weight,
+                        name_mapped,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                        return_success=True,
+                    )
+                if success:
+                    if not is_fused_expert:
+                        loaded_expert_param_names.add(name_mapped)
+                        continue
+                    name = name_mapped
+                    break
+            else:
+                if loaded_expert_param_names:
+                    loaded_params.update(loaded_expert_param_names)
+                    continue
+                if "learnable_sink_param" in name:
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    narrow_weight = loaded_weight[head_rank_start:head_rank_end]
+                    n = narrow_weight.shape[0]
+                    with torch.no_grad():
+                        params_dict[name][:n].copy_(narrow_weight)
+                else:
+                    if is_expert_weight:
+                        # An expert weight that is not mapped to this rank.
+                        continue
+                    if "gate.e_score_correction_bias" in name:
+                        name = name.replace(
+                            "gate.e_score_correction_bias", "expert_bias"
+                        )
+                    if "router.gate." in name:
+                        name = name.replace("router.", "")
+                    if "hc_fn" in name:
+                        name = name.replace("hc_fn", "hc_fn.weight")
+                    if "hc_head_fn" in name:
+                        name = name.replace("hc_head_fn", "hc_head_fn.weight")
+
+                    if is_pp_missing_parameter(name, self):
+                        continue
+
+                    if name not in params_dict:
+                        if _should_skip_missing_param(name):
+                            continue
+                        logger.warning_once(
+                            "Skipping unknown checkpoint weight: %s",
+                            name,
+                        )
+                        continue
+
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    if stacked_shard_id is None:
+                        weight_loader(param, loaded_weight)
+                    else:
+                        weight_loader(param, loaded_weight, stacked_shard_id)
+            loaded_params.add(name)
+
+        return loaded_params
+
+
+class HYV4ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_a_proj": (".fused_qkv_a_proj", 0),
+            ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1),
+            ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
+            ".mlp.up_proj": (".mlp.gate_up_proj", 1),
+            ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
+            ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
+            ".indexer.weights_proj.": (".indexer.wk_weights_proj.", 1),
+        }
+    )
+    packed_modules_mapping = {
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        # MLA runs both latent down-projections as one GEMM.
+        "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
+        # The indexer fuses wk and weights_proj into one GEMM.
+        "wk_weights_proj": ["wk", "weights_proj"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+        if quant_config is not None:
+            quant_config.packed_modules_mapping = self.packed_modules_mapping
+
+        parallel_config = vllm_config.parallel_config
+        eplb_config = parallel_config.eplb_config
+        self.num_redundant_experts = eplb_config.num_redundant_experts
+
+        self.model = HYV4Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        if get_pp_group().is_last_rank:
+            # The head stays in the model dtype; ``enable_lm_head_fp32`` is
+            # surfaced as ``head_dtype`` on the config (see HYV4Config), which
+            # makes LogitsProcessor accumulate the projection straight into
+            # fp32 instead of materializing an fp32 copy of the weight.
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+            # With tie_word_embeddings, embed_tokens is kept on the last rank
+            # (see HYV4Model.__init__) so the weight can be shared here.
+            if self.config.tie_word_embeddings:
+                self.lm_head.weight = self.model.embed_tokens.weight
+        else:
+            self.lm_head = PPMissingLayer()
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        logits = self.logits_processor(self.lm_head, hidden_states)
+
+        if getattr(self.config, "soft_logits_capping", False):
+            soft_cap = self.config.soft_logits_capping_logits
+            logits = soft_cap * torch.nn.functional.tanh(logits / soft_cap)
+
+        return logits
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # MTP layers are loaded by the MTP head, not by the target model. They
+        # appear as model.mtp_layers.<i>.* and as the layer ids past the backbone.
+        mtp_start = self.config.num_hidden_layers
+        drop_prefixes: dict[str, str | None] = {"model.mtp_layers.": None}
+        for i in range(getattr(self.config, "num_nextn_predict_layers", 0)):
+            drop_prefixes[f"model.layers.{mtp_start + i}."] = None
+        if self.config.tie_word_embeddings:
+            drop_prefixes["lm_head."] = None
+        drop_mapper = WeightsMapper(orig_to_new_prefix=drop_prefixes)
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper | drop_mapper)
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return self.model.get_expert_mapping()
diff --git a/vllm/models/hy_v4/nvidia/attention.py b/vllm/models/hy_v4/nvidia/attention.py
index 40a05f2c6407..0e7fb61173d1 100644
--- a/vllm/models/hy_v4/nvidia/attention.py
+++ b/vllm/models/hy_v4/nvidia/attention.py
@@ -18,7 +18,7 @@
 import regex as re
 import torch
 from torch import nn
-from transformers import DeepseekV2Config, DeepseekV3Config, PretrainedConfig
+from transformers import DeepseekV2Config, DeepseekV3Config, PreTrainedConfig
 
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
@@ -51,7 +51,7 @@
 _WEIGHT_LAYER_INDEX_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)")
 
 
-def compute_skip_topk_layers(config: PretrainedConfig) -> set[int]:
+def compute_skip_topk_layers(config: PreTrainedConfig) -> set[int]:
     """Return the backbone layers that reuse a previous layer's top-k indices.
 
     A "shared" indexer layer performs sparse attention with the indices computed
@@ -67,6 +67,7 @@ def compute_skip_topk_layers(config: PretrainedConfig) -> set[int]:
     Raises:
         ValueError: If ``indexer_types`` has the wrong length or an unknown
             entry, or if ``index_topk_freq`` is not a positive integer.
+
     """
     if not hasattr(config, "index_topk"):
         return set()
@@ -115,6 +116,7 @@ def is_skip_topk_indexer_weight(weight_name: str, skip_topk_layers: set[int]) ->
     Returns:
         True when the weight is an indexer weight of a layer that has no
         indexer module and therefore must be dropped.
+
     """
     if ".indexer." not in weight_name or not skip_topk_layers:
         return False
@@ -278,7 +280,7 @@ class HYV4MLAAttention(nn.Module):
     def __init__(
         self,
         vllm_config: VllmConfig,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         num_heads: int,
         qk_nope_head_dim: int,
@@ -557,6 +559,7 @@ def _resolve_sink_backend(
             The backend class to bind, or None when no sink-capable backend is
             available; the caller then loads the sink weight but disables the
             bias.
+
         """
         head_size = self.kv_lora_rank + self.qk_rope_head_dim
         dtype = torch.get_default_dtype()
@@ -755,6 +758,8 @@ def _indexer_and_attn(
         """
         if self.indexer is not None and self.is_sparse and not self.skip_topk:
             self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb)
+        if self.is_sparse:
+            self.mla_attn.impl.record_logical_topk_ready()  # type: ignore[attr-defined]
         out.copy_(
             self.mla_attn(
                 q,
diff --git a/vllm/models/hy_v4/nvidia/flashmla_sparse.py b/vllm/models/hy_v4/nvidia/flashmla_sparse.py
index 952dd8de5c08..1d62f2f556de 100644
--- a/vllm/models/hy_v4/nvidia/flashmla_sparse.py
+++ b/vllm/models/hy_v4/nvidia/flashmla_sparse.py
@@ -34,6 +34,7 @@
     flash_mla_sparse_fwd,
     flash_mla_with_kvcache,
 )
+from vllm.v1.worker.workspace import current_workspace_manager
 
 if TYPE_CHECKING:
     from vllm.model_executor.models.deepseek_v2 import Indexer
@@ -99,6 +100,7 @@ def _validate_sinks(sinks: torch.Tensor | None, num_heads: int) -> None:
         Raises:
             ValueError: If the dtype is not float32 or the shape is not
                 ``(num_heads,)``.
+
         """
         if sinks is None:
             return
@@ -134,6 +136,7 @@ def _sinks_for_query(
         Raises:
             ValueError: If the sink and query head layouts disagree, or if they
                 live on different devices.
+
         """
         sinks = self.sinks
         if sinks is None:
@@ -218,6 +221,7 @@ def _bf16_flash_mla_kernel(
         topk_indices: torch.Tensor,
         topk_length: torch.Tensor | None = None,
         actual_num_heads: int | None = None,
+        out: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         num_tokens = q.shape[0]
         kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
@@ -248,7 +252,14 @@ def _bf16_flash_mla_kernel(
             # Zero (not new_empty) the padded lanes: topk_indices is shared by
             # all heads, so the kernel reduces across the head group and NaNs
             # from uninitialized memory would leak into the real heads.
-            q_padded = q.new_zeros((q.shape[0], padded_num_heads, q.shape[2]))
+            if out is None:
+                q_padded = q.new_zeros((q.shape[0], padded_num_heads, q.shape[2]))
+            else:
+                *_, q_padded, _, _ = current_workspace_manager().get_simultaneous(
+                    *self.workspace_specs
+                )
+                q_padded = q_padded[:num_tokens]
+                q_padded[:, actual_num_heads:].zero_()
             q_padded[:, :actual_num_heads, :] = q
             q = q_padded
 
@@ -260,6 +271,7 @@ def _bf16_flash_mla_kernel(
             self.softmax_scale,
             attn_sink=attn_sink,
             topk_length=topk_length,
+            out=out,
         )
 
         output = output[:, :actual_num_heads, :]
diff --git a/vllm/models/hy_v4/nvidia/hc.py b/vllm/models/hy_v4/nvidia/hc.py
index 130d3842ecfc..617b361e754f 100644
--- a/vllm/models/hy_v4/nvidia/hc.py
+++ b/vllm/models/hy_v4/nvidia/hc.py
@@ -17,7 +17,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.model_executor.layers.hpc import HpcIHCHead, HpcIHCPost, HpcIHCPre
 from vllm.model_executor.layers.linear import ReplicatedLinear
@@ -40,7 +40,7 @@ class HYV4HCPreLayer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_dim: int,
         hc_mult: int = 4,
         magnitude: float = 2.0,
@@ -107,6 +107,7 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         Returns:
             A tuple of the pre-gated reduction ``[num_tokens, d]`` and the post
             gates ``[num_tokens, hc]`` consumed by `HYV4HCPostLayer`.
+
         """
         if self.hpc_op is not None:
             return self.hpc_op(x)
@@ -162,7 +163,7 @@ class HYV4HCPostLayer(nn.Module):
         y[n, i, d] = post[n, i] * x[n, d] + residual[n, i, d]
     """
 
-    def __init__(self, config: PretrainedConfig):
+    def __init__(self, config: PreTrainedConfig):
         super().__init__()
         self.config = config
 
@@ -186,6 +187,7 @@ def forward(
 
         Returns:
             The updated residual channels ``[num_tokens, hc, d]``.
+
         """
         if self.hpc_op is not None:
             return self.hpc_op(x, residual, post)
@@ -211,7 +213,7 @@ class HYV4HCHeadLayer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         hidden_size: int,
         hc_mult: int = 4,
         hc_eps: float = 1e-6,
@@ -272,6 +274,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         Returns:
             The merged hidden state ``[num_tokens, d]``.
+
         """
         if self.hpc_op is not None:
             return self.hpc_op(x)
@@ -298,7 +301,7 @@ class HYV4HCLayer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_idx: int,
         init_std: float = 6e-3,
         base_noise_std: float = 0.0,
@@ -367,6 +370,7 @@ def pre(
             A tuple of the reduced hidden states ``[num_tokens, d]``, the post
             gates ``[num_tokens, hc]`` (``None`` when iHC is disabled) and the
             residual (the untouched input).
+
         """
         if not self.enable_ihc:
             return hidden_states, None, hidden_states
diff --git a/vllm/models/hy_v4/nvidia/model.py b/vllm/models/hy_v4/nvidia/model.py
index 626d704f0e74..e9313989fa57 100644
--- a/vllm/models/hy_v4/nvidia/model.py
+++ b/vllm/models/hy_v4/nvidia/model.py
@@ -30,7 +30,7 @@
 import regex as re
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import CacheConfig, VllmConfig
 from vllm.distributed import (
@@ -74,7 +74,7 @@
 logger = init_logger(__name__)
 
 
-def _normalize_hyv4_config(config: PretrainedConfig) -> PretrainedConfig:
+def _normalize_hyv4_config(config: PreTrainedConfig) -> PreTrainedConfig:
     """Populate the aliases consumed by the shared MoE implementation."""
     config.router_scaling_factor = config.routed_scaling_factor
     config.num_experts = config.n_routed_experts
@@ -94,7 +94,7 @@ class HYV4DecoderLayer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         vllm_config: VllmConfig,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
@@ -189,7 +189,7 @@ def _forward_ihc(
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
     ) -> tuple[torch.Tensor, None]:
-        """iHC forward: each sub-block reduces and re-scatters the channels."""
+        """IHC forward: each sub-block reduces and re-scatters the channels."""
         hidden_states = self.hc_attn_layer.prepare_input(hidden_states)
         hidden_states, post_gates, residual = self.hc_attn_layer.pre(hidden_states)
         hidden_states = self.input_layernorm(hidden_states)
diff --git a/vllm/models/hy_v4/nvidia/moe.py b/vllm/models/hy_v4/nvidia/moe.py
index be8e49f21eaf..4c5907f76501 100644
--- a/vllm/models/hy_v4/nvidia/moe.py
+++ b/vllm/models/hy_v4/nvidia/moe.py
@@ -4,7 +4,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import get_current_vllm_config
 from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size
@@ -87,7 +87,7 @@ class HYV4MoEFused(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
         enable_eplb: bool = False,
diff --git a/vllm/models/hy_v4/nvidia/mtp.py b/vllm/models/hy_v4/nvidia/mtp.py
index 49f1b3d7a508..8683777cf41f 100644
--- a/vllm/models/hy_v4/nvidia/mtp.py
+++ b/vllm/models/hy_v4/nvidia/mtp.py
@@ -8,7 +8,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
 from vllm.distributed import (
@@ -59,7 +59,7 @@
 
 
 def _get_spec_layer_idx_from_weight_name(
-    config: PretrainedConfig, weight_name: str
+    config: PreTrainedConfig, weight_name: str
 ) -> int | None:
     """Return the MTP layer index a checkpoint weight belongs to, or None.
 
@@ -106,6 +106,7 @@ def _resolve_fused_expert_param(
 
     Returns:
         The matching draft parameter name, or None when there is none.
+
     """
     if not ckpt_suffix:
         return param_base if param_base in params_dict else None
@@ -146,7 +147,7 @@ def _prepare_mtp_fp8_expert_scale(
 
 
 def _create_mtp_quant_config(
-    hf_config: PretrainedConfig,
+    hf_config: PreTrainedConfig,
     backbone_quant_config: QuantizationConfig | None = None,
 ) -> QuantizationConfig | None:
     """Create the quantization config for the MTP layers.
@@ -164,6 +165,7 @@ def _create_mtp_quant_config(
     Returns:
         The quantization config to use for the MTP layers, or None when the MTP
         layers are unquantized.
+
     """
     mtp_quant_algo = getattr(hf_config, "mtp_quant_algo", None)
 
@@ -241,6 +243,7 @@ def _remap_mtp_quant_exclusions(
 
     Returns:
         A shallow copy with the translated exclusions, or the input unchanged.
+
     """
     if quant_config is None:
         return None
@@ -300,7 +303,7 @@ class HYV4SharedHead(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
     ) -> None:
         super().__init__()
@@ -326,7 +329,7 @@ class HYV4MultiTokenPredictorLayer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         vllm_config: VllmConfig,
         model_config: ModelConfig,
@@ -671,6 +674,7 @@ def _load_expert_weight(
         Returns:
             True when the weight was consumed (even if this rank holds none of
             the addressed experts).
+
         """
         base = name.split(".experts.")[0]
         for ckpt_proj, tag in (
diff --git a/vllm/models/inkling/amd/logits_processor.py b/vllm/models/inkling/amd/logits_processor.py
index 9cf1c663c891..4302c6f674f8 100644
--- a/vllm/models/inkling/amd/logits_processor.py
+++ b/vllm/models/inkling/amd/logits_processor.py
@@ -37,6 +37,7 @@ class InklingLogitsProcessor(LogitsProcessor):
         soft_cap: Optional logit soft cap (``None`` for the served checkpoint).
         logits_mup_width_multiplier: muP width divisor for the final logits;
             ``None`` or ``0`` disables it.
+
     """
 
     def __init__(
diff --git a/vllm/models/inkling/amd/ops/lamport.py b/vllm/models/inkling/amd/ops/lamport.py
deleted file mode 100644
index 9705a9a3df37..000000000000
--- a/vllm/models/inkling/amd/ops/lamport.py
+++ /dev/null
@@ -1,766 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm.
-
-The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``.
-
-Liveness
---------
-Large grids are deliberately split at the two communication dependencies:
-
-  1. ``_publish_input_kernel`` only publishes rank partials.
-  2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard.
-  3. ``_sconv_publish_kernel`` only computes and publishes the local result.
-  4. ``_gather_norm_kernel`` waits, gathers, and normalizes.
-
-Without PDL, CUDA stream order completes a producer before its consumer.  With
-PDL, every producer CTA posts all peer stores before triggering its dependent,
-and the consumer executes ``gdc_wait`` before polling.  A consumer therefore
-waits only for stores from a producer that is already running (or complete) on
-another GPU.  It never waits for another block in its own grid.  Consequently
-spinning consumers cannot occupy resources needed by any producer, and the
-wait-for graph has no cycle.  The proof is independent of grid size, block
-dispatch order, and occupancy.
-
-For one token, the first three phases use eight independent channel slices to
-expose enough CTA parallelism for decode latency.  Each slice has exclusive
-ownership of its cache and Lamport columns; the gather/RMSNorm phase retains
-one CTA per token so no cross-CTA reduction or completion counter is needed.
-
-Immediate buffer reuse (including replay of a captured CUDA graph) is also
-safe.  Rank R cannot republish input for call n+1 until its gather for n has
-finished; that gather waited for owner O's output, which O publishes only
-after consuming R's input for n.  Likewise, R cannot republish output for n+1
-until its reduction for n+1 has observed destination D's input; D publishes
-that input only after its gather consumed R's output for n.  Thus every prior
-read happens-before a same-slot rewrite.  Three generations reduce incidental
-coupling for ordinary launches, but correctness does not depend on rotation.
-
-The payload itself is the Lamport flag.  A 32-bit store publishes two bf16s
-atomically; 0x80008000 (two negative zeroes) denotes an empty pair.  Real
-negative zeroes are changed to positive zero before publication.  Consumers
-use volatile 32-bit loads and restore the sentinel after consuming a slot.
-"""
-
-from __future__ import annotations
-
-import os
-
-import torch
-
-from vllm.distributed import get_tp_group
-from vllm.distributed.parallel_state import in_the_same_node_as
-from vllm.logger import init_logger
-from vllm.triton_utils import tl, triton
-
-logger = init_logger(__name__)
-
-_MAX_TOKENS = 16384
-_EMPTY_PAIR = tl.constexpr(0x80008000)
-
-
-@triton.jit
-def _pack_bf16_pairs(values):
-    """Pack bf16 values into atomic u32 pairs and reserve negative zero."""
-    lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2]))
-    lo = lo.to(tl.uint16, bitcast=True)
-    hi = hi.to(tl.uint16, bitcast=True)
-    lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32)
-    hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32)
-    return lo | (hi << 16)
-
-
-@triton.jit
-def _unpack_bf16_pairs(values):
-    lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
-    hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
-    return tl.interleave(lo, hi)
-
-
-@triton.jit
-def _wait_pairs(ptr, offsets, mask):
-    values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
-    while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0:
-        values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True)
-    return values
-
-
-@triton.jit
-def _publish_input_kernel(
-    stage_ptr,
-    peer_ptrs,
-    peer_offset_u32,
-    stride_stage_t,
-    C: tl.constexpr,
-    CS: tl.constexpr,
-    CS_P2: tl.constexpr,
-    SPLITS: tl.constexpr,
-    RANK: tl.constexpr,
-    WORLD: tl.constexpr,
-    USE_PDL: tl.constexpr,
-    launch_pdl: tl.constexpr,
-):
-    """Publish this rank's full partial row into every shard owner's slots."""
-    token = tl.program_id(0).to(tl.int64)
-    split = tl.program_id(1).to(tl.int64)
-    CSS: tl.constexpr = CS // SPLITS
-    pair = tl.arange(0, CS_P2 // 2)
-    pair_mask = pair < CSS // 2
-    elem = tl.arange(0, CS_P2)
-    elem_mask = elem < CSS
-    ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
-    # Address generation is independent of the preceding stream kernel.  The
-    # acquire remains before stage/peer loads, which may consume its writes.
-    if USE_PDL:
-        tl.extra.cuda.gdc_wait()
-
-    for owner in tl.static_range(WORLD):
-        values = tl.load(
-            stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem,
-            mask=elem_mask,
-            other=0.0,
-        )
-        packed = _pack_bf16_pairs(values)
-        base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32))
-        dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair
-        tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask)
-    if USE_PDL:
-        tl.extra.cuda.gdc_launch_dependents()
-
-
-@triton.jit
-def _reduce_insert_kernel(
-    input_peer_ptrs,
-    input_peer_offset_u32,
-    cache_ptr,
-    slot_ptr,
-    stride_cache_block,
-    stride_cache_head,
-    stride_cache_token,
-    stride_cache_dim,
-    block_size,
-    C: tl.constexpr,
-    CS: tl.constexpr,
-    CS_P2: tl.constexpr,
-    SPLITS: tl.constexpr,
-    HEAD_SIZE: tl.constexpr,
-    CACHE_OFFSET: tl.constexpr,
-    RANK: tl.constexpr,
-    WORLD: tl.constexpr,
-    USE_PDL: tl.constexpr,
-    launch_pdl: tl.constexpr,
-):
-    """Consume all partials for this rank and insert the reduced cache row."""
-    token = tl.program_id(0).to(tl.int64)
-    split = tl.program_id(1).to(tl.int64)
-    CSS: tl.constexpr = CS // SPLITS
-    source = tl.arange(0, WORLD)
-    pair = tl.arange(0, CS_P2 // 2)
-    pair_mask = pair < CSS // 2
-    offsets = (
-        (token * WORLD + source)[:, None] * (CS // 2)
-        + split * (CSS // 2)
-        + pair[None, :]
-    )
-    mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :]
-    input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
-    input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32))
-    input_u32 += input_peer_offset_u32
-    # Slot metadata and the cache destination do not depend on input publish.
-    slot = tl.load(slot_ptr + token)
-    valid = slot >= 0
-    channel = tl.arange(0, CS_P2)
-    channel_mask = channel < CSS
-    global_channel = split * CSS + channel
-    head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
-    dim = CACHE_OFFSET + global_channel % HEAD_SIZE
-    safe_slot = tl.maximum(slot, 0).to(tl.int64)
-    dst = (
-        cache_ptr
-        + (safe_slot // block_size) * stride_cache_block
-        + head * stride_cache_head
-        + (safe_slot % block_size) * stride_cache_token
-        + dim * stride_cache_dim
-    )
-    if USE_PDL:
-        tl.extra.cuda.gdc_wait()
-    packed = _wait_pairs(input_u32, offsets, mask)
-
-    lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True)
-    hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True)
-    reduced = tl.interleave(
-        tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16),
-        tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16),
-    )
-    tl.store(
-        input_u32 + offsets,
-        tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32),
-        mask=mask,
-    )
-
-    tl.store(dst, reduced, mask=valid & channel_mask)
-    if USE_PDL:
-        tl.extra.cuda.gdc_launch_dependents()
-
-
-@triton.jit
-def _sconv_publish_kernel(
-    peer_ptrs,
-    peer_offset_u32,
-    residual_ptr,
-    weight_ptr,
-    cache_ptr,
-    position_ptr,
-    sequence_ptr,
-    slot_ptr,
-    block_table_ptr,
-    stride_residual_t,
-    stride_cache_block,
-    stride_cache_head,
-    stride_cache_token,
-    stride_cache_dim,
-    stride_block_table_r,
-    max_blocks,
-    block_size,
-    C: tl.constexpr,
-    CS: tl.constexpr,
-    CS_P2: tl.constexpr,
-    SPLITS: tl.constexpr,
-    HEAD_SIZE: tl.constexpr,
-    CACHE_OFFSET: tl.constexpr,
-    RANK: tl.constexpr,
-    WORLD: tl.constexpr,
-    WINDOW: tl.constexpr,
-    USE_PDL: tl.constexpr,
-    launch_pdl: tl.constexpr,
-):
-    """Compute this rank's shard, then publish it to every rank."""
-    tl.static_assert(WINDOW == 4)
-    token = tl.program_id(0).to(tl.int64)
-    split = tl.program_id(1).to(tl.int64)
-    CSS: tl.constexpr = CS // SPLITS
-    channel = tl.arange(0, CS_P2)
-    channel_mask = channel < CSS
-    global_channel = split * CSS + channel
-    head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1)
-    dim = CACHE_OFFSET + global_channel % HEAD_SIZE
-    slot = tl.load(slot_ptr + token)
-    valid = slot >= 0
-    position = tl.load(position_ptr + token)
-    sequence = tl.load(sequence_ptr + token)
-    safe_slot = tl.maximum(slot, 0).to(tl.int64)
-    own_ptr = (
-        cache_ptr
-        + (safe_slot // block_size) * stride_cache_block
-        + head * stride_cache_head
-        + (safe_slot % block_size) * stride_cache_token
-        + dim * stride_cache_dim
-    )
-    # These loads were made visible before reduce/insert could signal us and
-    # are independent of its cache store.  Hoisting them gives PDL useful work
-    # to overlap while retaining the acquire before every cache read.
-    residual = tl.load(
-        residual_ptr + token * stride_residual_t + RANK * CS + global_channel,
-        mask=channel_mask,
-        other=0.0,
-    )
-    weight0 = tl.load(
-        weight_ptr + global_channel * WINDOW,
-        mask=channel_mask,
-        other=0.0,
-    ).to(tl.float32)
-    weight1 = tl.load(
-        weight_ptr + global_channel * WINDOW + 1,
-        mask=channel_mask,
-        other=0.0,
-    ).to(tl.float32)
-    weight2 = tl.load(
-        weight_ptr + global_channel * WINDOW + 2,
-        mask=channel_mask,
-        other=0.0,
-    ).to(tl.float32)
-    weight3 = tl.load(
-        weight_ptr + global_channel * WINDOW + 3,
-        mask=channel_mask,
-        other=0.0,
-    ).to(tl.float32)
-    ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64))
-    if USE_PDL:
-        tl.extra.cuda.gdc_wait()
-    current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0)
-
-    conv = tl.zeros([CS_P2], tl.float32)
-    for tap_idx in tl.static_range(WINDOW):
-        source_position = position - (WINDOW - 1) + tap_idx
-        take = valid & (source_position >= 0)
-        if tap_idx == WINDOW - 1:
-            value = tl.where(take, current.to(tl.float32), 0.0)
-        else:
-            safe_position = tl.maximum(source_position, 0)
-            logical_block = tl.minimum(safe_position // block_size, max_blocks - 1)
-            physical_block = tl.load(
-                block_table_ptr + sequence * stride_block_table_r + logical_block,
-                mask=take,
-                other=0,
-            ).to(tl.int64)
-            source_ptr = (
-                cache_ptr
-                + physical_block * stride_cache_block
-                + head * stride_cache_head
-                + (safe_position % block_size) * stride_cache_token
-                + dim * stride_cache_dim
-            )
-            cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0)
-            value = tl.where(take, cached.to(tl.float32), 0.0)
-        if tap_idx == 0:
-            weight = weight0
-        elif tap_idx == 1:
-            weight = weight1
-        elif tap_idx == 2:
-            weight = weight2
-        else:
-            weight = weight3
-        conv += value * weight
-
-    # Preserve both bf16 rounding points of the original sublayer.
-    short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16)
-    output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to(
-        tl.bfloat16
-    )
-
-    pair = tl.arange(0, CS_P2 // 2)
-    pair_mask = pair < CSS // 2
-    packed = _pack_bf16_pairs(output)
-    row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair
-
-    for destination in tl.static_range(WORLD):
-        base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32))
-        tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask)
-    if USE_PDL:
-        tl.extra.cuda.gdc_launch_dependents()
-
-
-@triton.jit
-def _gather_norm_kernel(
-    output_peer_ptrs,
-    output_peer_offset_u32,
-    norm_weight_ptr,
-    normed_ptr,
-    residual_out_ptr,
-    eps,
-    stride_output_t,
-    C: tl.constexpr,
-    C_P2: tl.constexpr,
-    RANK: tl.constexpr,
-    HAS_NORM: tl.constexpr,
-    USE_PDL: tl.constexpr,
-    launch_pdl: tl.constexpr,
-):
-    """Consume a complete row; one CTA owns all outputs for one token."""
-    token = tl.program_id(0).to(tl.int64)
-    pair = tl.arange(0, C_P2 // 2)
-    pair_mask = pair < C // 2
-    output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
-    output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32))
-    output_u32 += output_peer_offset_u32
-    offsets = token * (C // 2) + pair
-    channel = tl.arange(0, C_P2)
-    channel_mask = channel < C
-    # Gamma is independent of the preceding sconv publication.  Keep the
-    # acquire immediately before polling the Lamport output slots.
-    if HAS_NORM:
-        weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0)
-    if USE_PDL:
-        tl.extra.cuda.gdc_wait()
-    packed = _wait_pairs(output_u32, offsets, pair_mask)
-    row = _unpack_bf16_pairs(packed)
-    tl.store(
-        residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask
-    )
-    if HAS_NORM:
-        row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0)
-        inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps)
-        tl.store(
-            normed_ptr + token * stride_output_t + channel,
-            (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16),
-            mask=channel_mask,
-        )
-    tl.store(
-        output_u32 + offsets,
-        tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32),
-        mask=pair_mask,
-    )
-    if USE_PDL:
-        tl.extra.cuda.gdc_launch_dependents()
-
-
-@triton.jit
-def _validate_lamport_init_kernel(
-    input_peer_ptrs,
-    output_peer_ptrs,
-    bad_ptr,
-    num_pairs,
-    RANK: tl.constexpr,
-    BLOCK: tl.constexpr,
-):
-    """Validate both complete local allocations through their fabric pointers."""
-    offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
-    mask = offsets < num_pairs
-    ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64))
-    ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64))
-    local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32))
-    local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32))
-    value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR)
-    value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR)
-    bad = tl.max(
-        tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0)
-    )
-    if bad != 0:
-        tl.atomic_max(bad_ptr, 1)
-
-
-class LamportRSConv:
-    """Persistent symmetric buffers for one TP group."""
-
-    def _initialize_lamport_buffers(self) -> None:
-        """Arm every slot and collectively verify the exact sentinel bits.
-
-        Do not replace this with a zero-fill: Lamport readiness distinguishes
-        the bf16 bit pattern for -0.0 (0x8000) from every published payload.
-        The validation is deliberately collective so one rank cannot enter the
-        first polling kernel while another rank still has an unarmed buffer.
-        """
-        if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous():
-            raise RuntimeError("Lamport symmetric buffers must be contiguous")
-        if self.buf_in.numel() % 2 or self.buf_out.numel() % 2:
-            raise RuntimeError("Lamport buffers must contain whole uint32 pairs")
-
-        # Fill through int16 so each bf16 lane receives the exact -0.0 bits.
-        # This covers all three generations and all max-token slots, including
-        # slots that a smaller first invocation does not touch.
-        self.buf_in.view(torch.int16).fill_(-0x8000)
-        self.buf_out.view(torch.int16).fill_(-0x8000)
-        torch.accelerator.synchronize(self.device)
-        self.tp.barrier()
-
-        # Scan the complete local allocations.  Since every rank performs the
-        # scan and participates in MAX, success proves every symmetric backing
-        # allocation was armed before any rank is allowed to use generation 0.
-        bad = torch.logical_or(
-            self.buf_in.view(torch.int16).ne(-0x8000).any(),
-            self.buf_out.view(torch.int16).ne(-0x8000).any(),
-        ).to(dtype=torch.float32)
-        bad = self.tp.all_reduce(bad)
-        if int(bad.item()) != 0:
-            raise RuntimeError("Lamport sentinel initialization failed on a TP rank")
-        self.tp.barrier()
-
-    def _initialize_mnnvl_buffers(self) -> None:
-        """Initialize and validate FlashInfer fabric-mapped Lamport storage."""
-        self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16)
-        self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16)
-        torch.accelerator.synchronize(self.device)
-        self.tp.barrier()
-
-        bad = torch.zeros((), dtype=torch.int32, device=self.device)
-        num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2
-        _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)](
-            self.input_peer_ptrs,
-            self.output_peer_ptrs,
-            bad,
-            num_pairs,
-            RANK=self.rank,
-            BLOCK=256,
-            num_warps=4,
-        )
-        bad = self.tp.all_reduce(bad.to(torch.float32))
-        if int(bad.item()) != 0:
-            raise RuntimeError("MNNVL Lamport sentinel initialization failed")
-        self.tp.barrier()
-
-    def __init__(
-        self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS
-    ) -> None:
-        import torch.distributed._symmetric_memory as symm_mem
-
-        tp = get_tp_group()
-        self.tp = tp
-        self.group = tp.device_group
-        self.world_size = tp.world_size
-        self.rank = tp.rank_in_group
-        self.device = torch.device(tp.device)
-        if self.world_size not in (2, 4, 8):
-            raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}")
-        if hidden_size % (2 * self.world_size) != 0:
-            raise ValueError("hidden size must produce an even shard on every rank")
-        if window_size != 4:
-            raise ValueError(f"short-conv window size must be 4, got {window_size}")
-        if max_tokens < 1 or max_tokens > _MAX_TOKENS:
-            raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]")
-
-        is_cross_node = not all(in_the_same_node_as(tp.cpu_group))
-        self.hidden_size = hidden_size
-        self.window_size = window_size
-        self.max_tokens = max_tokens
-        self.shard_size = hidden_size // self.world_size
-        # Three generations follow FlashInfer's Lamport layout.  A generation
-        # is reused only after two intervening collective calls.
-        self.num_buffers = 3
-        self.input_generation_bytes = max_tokens * hidden_size * 2
-        self.output_generation_bytes = max_tokens * hidden_size * 2
-        self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9
-        if is_cross_node:
-            try:
-                from flashinfer.comm.mnnvl import (
-                    McastGPUBuffer,
-                    TorchDistBackend,
-                    is_mnnvl_fabric_supported,
-                )
-            except ImportError as error:
-                raise RuntimeError(
-                    "cross-node TP requires FlashInfer MNNVL support"
-                ) from error
-
-            local_supported = int(
-                is_mnnvl_fabric_supported(torch.accelerator.current_device_index())
-            )
-            unsupported = torch.tensor(
-                1 - local_supported, dtype=torch.float32, device=self.device
-            )
-            unsupported = tp.all_reduce(unsupported)
-            if int(unsupported.item()) != 0:
-                raise RuntimeError("cross-node TP is supported only on MNNVL fabric")
-
-            comm_backend = TorchDistBackend(self.group)
-            allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2
-            self._mnnvl_input_handle = McastGPUBuffer(
-                allocation_bytes,
-                self.world_size,
-                self.rank,
-                self.device,
-                comm_backend,
-            )
-            self._mnnvl_output_handle = McastGPUBuffer(
-                allocation_bytes,
-                self.world_size,
-                self.rank,
-                self.device,
-                comm_backend,
-            )
-            self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev()
-            self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev()
-            self._initialize_mnnvl_buffers()
-            logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers")
-        else:
-            self.buf_in = symm_mem.empty(
-                self.num_buffers,
-                max_tokens,
-                self.world_size,
-                self.shard_size,
-                dtype=torch.bfloat16,
-                device=self.device,
-            )
-            self.buf_out = symm_mem.empty(
-                self.num_buffers,
-                max_tokens,
-                hidden_size,
-                dtype=torch.bfloat16,
-                device=self.device,
-            )
-            group_name = self.group.group_name
-            input_handle = symm_mem.rendezvous(self.buf_in, group_name)
-            output_handle = symm_mem.rendezvous(self.buf_out, group_name)
-            self.input_peer_ptrs = input_handle.buffer_ptrs_dev
-            self.output_peer_ptrs = output_handle.buffer_ptrs_dev
-            self._input_handle = input_handle
-            self._output_handle = output_handle
-            self._initialize_lamport_buffers()
-        self.generation = 0
-
-    def usable(self, num_tokens: int) -> bool:
-        return 0 < num_tokens <= self.max_tokens
-
-    def rs_sconv_ag_add_norm(
-        self,
-        input_tensor: torch.Tensor,
-        residual: torch.Tensor,
-        conv_weight: torch.Tensor,
-        norm_weight: torch.Tensor | None,
-        eps: float,
-        cache: torch.Tensor,
-        positions: torch.Tensor,
-        block_table: torch.Tensor,
-        seq_idx: torch.Tensor,
-        slot_mapping: torch.Tensor,
-        off_s: int,
-        ws: int,
-        block_size: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor]:
-        """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``."""
-        tokens, hidden_size = residual.shape
-        if not self.usable(tokens):
-            raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]")
-        if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16:
-            raise ValueError("residual must be bf16 [T, 6144]")
-        if (
-            input_tensor.shape != residual.shape
-            or input_tensor.dtype != torch.bfloat16
-            or input_tensor.stride(1) != 1
-        ):
-            raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]")
-        shard_size = hidden_size // self.world_size
-        if conv_weight.shape != (shard_size, self.window_size):
-            raise ValueError(
-                f"conv_weight must have shape [{shard_size}, {self.window_size}]"
-            )
-        if (
-            conv_weight.dtype != torch.bfloat16
-            or conv_weight.stride(0) != self.window_size
-        ):
-            raise ValueError("conv_weight must be contiguous bf16")
-        if norm_weight is not None and (
-            norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16
-        ):
-            raise ValueError("norm_weight must be bf16 [6144] or None")
-        if cache.dtype != torch.bfloat16 or cache.ndim != 4:
-            raise ValueError("cache must be a 4-D bf16 tensor")
-        if shard_size % ws != 0 or cache.shape[1] != shard_size // ws:
-            raise ValueError("cache head layout is inconsistent with ws")
-        if off_s < 0 or off_s + ws > cache.shape[3]:
-            raise ValueError("cache channel offset is out of bounds")
-
-        index = self.generation
-        input_offset = index * self.input_generation_bytes // 4
-        output_offset = index * self.output_generation_bytes // 4
-        normed = torch.empty_like(residual) if norm_weight is not None else None
-        residual_out = torch.empty_like(residual)
-        phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1
-        phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits)
-        phase_grid = (tokens, phase_splits)
-        # Wide CTAs win before the grid saturates; smaller CTAs reduce register
-        # pressure once high-throughput batches provide enough parallelism.
-        if 128 <= tokens <= 2048:
-            phase_warps = 16
-        elif tokens > 2048:
-            phase_warps = 8
-        else:
-            phase_warps = 4
-        gather_warps = 4 if tokens >= 256 else 8
-        _publish_input_kernel[phase_grid](
-            input_tensor,
-            self.input_peer_ptrs,
-            input_offset,
-            input_tensor.stride(0),
-            C=hidden_size,
-            CS=shard_size,
-            CS_P2=phase_tile_p2,
-            SPLITS=phase_splits,
-            RANK=self.rank,
-            WORLD=self.world_size,
-            USE_PDL=self.use_pdl,
-            launch_pdl=self.use_pdl,
-            num_warps=phase_warps,
-        )
-        _reduce_insert_kernel[phase_grid](
-            self.input_peer_ptrs,
-            input_offset,
-            cache,
-            slot_mapping,
-            cache.stride(0),
-            cache.stride(1),
-            cache.stride(2),
-            cache.stride(3),
-            block_size,
-            C=hidden_size,
-            CS=shard_size,
-            CS_P2=phase_tile_p2,
-            SPLITS=phase_splits,
-            HEAD_SIZE=ws,
-            CACHE_OFFSET=off_s,
-            RANK=self.rank,
-            WORLD=self.world_size,
-            USE_PDL=self.use_pdl,
-            launch_pdl=self.use_pdl,
-            num_warps=phase_warps,
-        )
-        _sconv_publish_kernel[phase_grid](
-            self.output_peer_ptrs,
-            output_offset,
-            residual,
-            conv_weight,
-            cache,
-            positions,
-            seq_idx,
-            slot_mapping,
-            block_table,
-            residual.stride(0),
-            cache.stride(0),
-            cache.stride(1),
-            cache.stride(2),
-            cache.stride(3),
-            block_table.stride(0),
-            block_table.shape[1],
-            block_size,
-            C=hidden_size,
-            CS=shard_size,
-            CS_P2=phase_tile_p2,
-            SPLITS=phase_splits,
-            HEAD_SIZE=ws,
-            CACHE_OFFSET=off_s,
-            RANK=self.rank,
-            WORLD=self.world_size,
-            WINDOW=self.window_size,
-            USE_PDL=self.use_pdl,
-            launch_pdl=self.use_pdl,
-            num_warps=phase_warps,
-        )
-        _gather_norm_kernel[(tokens,)](
-            self.output_peer_ptrs,
-            output_offset,
-            norm_weight if norm_weight is not None else residual,
-            normed if normed is not None else residual_out,
-            residual_out,
-            eps,
-            residual.stride(0),
-            C=hidden_size,
-            C_P2=triton.next_power_of_2(hidden_size),
-            RANK=self.rank,
-            HAS_NORM=norm_weight is not None,
-            USE_PDL=self.use_pdl,
-            launch_pdl=self.use_pdl,
-            num_warps=gather_warps,
-        )
-        self.generation = (index + 1) % self.num_buffers
-        return normed, residual_out
-
-
-_STATE: LamportRSConv | None = None
-_STATE_FAILED = False
-
-
-def initialize_lamport_rs_conv(
-    hidden_size: int, window_size: int, max_num_batched_tokens: int
-) -> None:
-    """Collectively initialize the TP-group state during model construction."""
-    global _STATE, _STATE_FAILED
-    if _STATE is not None:
-        if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size:
-            raise RuntimeError("all Lamport users must share hidden and window sizes")
-        return
-    if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0":
-        return
-    try:
-        max_tokens = min(_MAX_TOKENS, max_num_batched_tokens)
-        _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens)
-    except Exception:
-        _STATE_FAILED = True
-        logger.exception("fused collective unavailable; use the NCCL fallback")
-
-
-def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None:
-    """Return the state initialized with the model, or ``None`` for fallback."""
-    if _STATE is not None and (
-        _STATE.hidden_size != hidden_size or _STATE.window_size != window_size
-    ):
-        raise RuntimeError("all Lamport users must share hidden and window sizes")
-    return _STATE
diff --git a/vllm/models/inkling/amd/ops/mm_towers.py b/vllm/models/inkling/amd/ops/mm_towers.py
deleted file mode 100644
index 021cf5274bd2..000000000000
--- a/vllm/models/inkling/amd/ops/mm_towers.py
+++ /dev/null
@@ -1,190 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Fused CUDA kernels for the Inkling vision/audio towers.
-
-Both kernels keep the reference paths' fp32 accumulation and per-op bf16
-rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently
-bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order
-(real-checkpoint-weight cosine vs reference > 0.9999998).
-"""
-
-from __future__ import annotations
-
-import torch
-
-from vllm.triton_utils import tl, tldevice, triton
-
-from .norm import _get_num_warps_from_block_size
-
-
-@triton.jit
-def _dmel_embed_sum_norm_kernel(
-    idx_ptr,  # [T, NB] int32 dMel bin indices (values in [0, VOCAB))
-    w_ptr,  # [NB * VOCAB, D] bf16 embedding table
-    norm_w_ptr,  # [D] (unused if HAS_NORM=False)
-    out_ptr,  # [T, D] bf16
-    eps,
-    D,
-    stride_idx_t,
-    NB: tl.constexpr,
-    VOCAB: tl.constexpr,
-    D_P2: tl.constexpr,
-    HAS_NORM: tl.constexpr,
-):
-    t = tl.program_id(0).to(tl.int64)
-    offs = tl.arange(0, D_P2)
-    mask = offs < D
-    # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)),
-    # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()).
-    acc = tl.zeros([D_P2], dtype=tl.float32)
-    for b in tl.static_range(NB):
-        v = tl.load(idx_ptr + t * stride_idx_t + b)
-        row = (b * VOCAB + v).to(tl.int64)
-        acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32)
-    h = acc.to(tl.bfloat16)
-    if HAS_NORM:
-        # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded
-        # bf16 multiply with the bf16 weight.
-        x32 = h.to(tl.float32)
-        var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D
-        xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16)
-        w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0)
-        h = xn * w
-    tl.store(out_ptr + t * D + offs, h, mask=mask)
-
-
-def dmel_embed_sum_norm(
-    dmel_idx: torch.Tensor,  # [T, NB] int32
-    weight: torch.Tensor,  # [NB * VOCAB, D] bf16
-    norm_weight: torch.Tensor | None,
-    eps: float,
-) -> torch.Tensor:
-    """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no
-    [T, NB, D] intermediate)."""
-    T, nb = dmel_idx.shape
-    D = weight.shape[1]
-    assert weight.shape[0] % nb == 0
-    vocab = weight.shape[0] // nb
-    out = torch.empty((T, D), dtype=weight.dtype, device=weight.device)
-    if T == 0:
-        return out
-    d_p2 = triton.next_power_of_2(D)
-    _dmel_embed_sum_norm_kernel[(T,)](
-        dmel_idx,
-        weight,
-        norm_weight if norm_weight is not None else weight,
-        out,
-        eps,
-        D,
-        dmel_idx.stride(0),
-        NB=nb,
-        VOCAB=vocab,
-        D_P2=d_p2,
-        HAS_NORM=norm_weight is not None,
-        # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4%
-        # at large T (the 80-row gather chain is latency- not lane-bound).
-        num_warps=8,
-    )
-    return out
-
-
-@triton.jit
-def _rmsnorm_gelu_kernel(
-    x_ptr,  # [R, D] bf16
-    w_ptr,  # [D]
-    out_ptr,  # [R, D] bf16 (or the folded layout when FOLD)
-    eps,
-    R,
-    D,
-    D_P2: tl.constexpr,
-    BLOCK_M: tl.constexpr,  # rows per block (>1 for small D)
-    HAS_GELU: tl.constexpr,
-    FOLD: tl.constexpr,
-    # fold geometry: input rows index [N, T, H, W]; the store scatters each
-    # row to (out_row, slot) of fold_timespace_to_depth's output layout.
-    FT: tl.constexpr,
-    FH: tl.constexpr,
-    FW: tl.constexpr,
-    TF: tl.constexpr,
-    HF: tl.constexpr,
-):
-    pid = tl.program_id(0).to(tl.int64)
-    rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
-    rmask = rows < R
-    offs = tl.arange(0, D_P2)
-    mask = rmask[:, None] & (offs < D)[None, :]
-    x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to(
-        tl.float32
-    )
-    var = tl.sum(x32 * x32, axis=1) / D
-    xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16)
-    w = tl.load(w_ptr + offs, mask=offs < D, other=0.0)
-    h = xn * w[None, :]  # bf16 multiply, matching ir.ops.rms_norm
-    if HAS_GELU:
-        # Exact (erf) GELU on the bf16-rounded norm output, fp32 math,
-        # matching F.gelu's opmath on a bf16 tensor.
-        g32 = h.to(tl.float32)
-        h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16)
-    if FOLD:
-        # Store directly into the next layer's folded layout (a pure
-        # permutation — replaces the separate fold copy pass).
-        t = (rows // (FH * FW)) % FT
-        hh = (rows // FW) % FH
-        ww = rows % FW
-        n = rows // (FT * FH * FW)
-        slot = ((t % TF) * HF + hh % HF) * HF + ww % HF
-        out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * (
-            FW // HF
-        ) + ww // HF
-        base = (out_row * (TF * HF * HF) + slot) * D
-        tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask)
-    else:
-        tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask)
-
-
-def rmsnorm_gelu(
-    x: torch.Tensor,  # [..., D] bf16 contiguous
-    weight: torch.Tensor,
-    eps: float,
-    gelu: bool = True,
-    fold: tuple[int, int] | None = None,  # (t_fold, hw_fold) of the NEXT fold
-) -> torch.Tensor:
-    """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block
-    when D is small. With ``fold``, x must be [N, T, H, W, D] and the output
-    comes back as ``fold_timespace_to_depth(result, *fold)``."""
-    D = x.shape[-1]
-    flat = x.reshape(-1, D)
-    assert flat.stride(1) == 1 and flat.stride(0) == D
-    R = flat.shape[0]
-    if fold is None:
-        out = torch.empty_like(flat)
-        ft = fh = fw = tf = hf = 1
-        out_shape = x.shape
-    else:
-        tf, hf = fold
-        N, ft, fh, fw, _ = x.shape
-        out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D)
-        out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
-    if R == 0:
-        return out.reshape(out_shape)
-    d_p2 = triton.next_power_of_2(D)
-    block_m = max(1, 4096 // d_p2)
-    _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)](
-        flat,
-        weight,
-        out,
-        eps,
-        R,
-        D,
-        D_P2=d_p2,
-        BLOCK_M=block_m,
-        HAS_GELU=gelu,
-        FOLD=fold is not None,
-        FT=ft,
-        FH=fh,
-        FW=fw,
-        TF=tf,
-        HF=hf,
-        num_warps=_get_num_warps_from_block_size(d_p2 * block_m),
-    )
-    return out.reshape(out_shape)
diff --git a/vllm/models/inkling/common/mm_preprocess.py b/vllm/models/inkling/common/mm_preprocess.py
index 249d7edeb43f..a949e4ec7403 100644
--- a/vllm/models/inkling/common/mm_preprocess.py
+++ b/vllm/models/inkling/common/mm_preprocess.py
@@ -13,9 +13,7 @@
 from transformers.feature_extraction_utils import BatchFeature
 
 from vllm.config.multimodal import (
-    AudioDummyOptions,
-    BaseDummyOptions,
-    ImageDummyOptions,
+    MultiModalDummyOptions,
 )
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal.inputs import (
@@ -145,7 +143,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         config = self.info.get_hf_config()
         num_images = mm_counts.get("image", 0)
@@ -156,12 +154,11 @@ def get_dummy_mm_data(
             patch_size = getattr(config.vision_config, "patch_size", 40)
             # A square image ~4 patches wide so the dummy emits several patches.
             side = patch_size * 4
-            image_overrides = mm_options.get("image")
             mm_data["image"] = self._get_dummy_images(
                 width=side,
                 height=side,
                 num_images=num_images,
-                overrides=cast(ImageDummyOptions | None, image_overrides),
+                overrides=mm_options.get("image"),
             )
         if num_audios:
             # Size the dummy at the maximum allowed audio so memory/encoder
@@ -169,11 +166,10 @@ def get_dummy_mm_data(
             params = self.info.get_hf_processor().audio_feature_extractor.params
             hop = int(round(params.audio_token_duration_s * params.sample_rate))
             audio_len = MAX_AUDIO_TOKENS * hop
-            audio_overrides = mm_options.get("audio")
             mm_data["audio"] = self._get_dummy_audios(
                 length=audio_len,
                 num_audios=num_audios,
-                overrides=cast(AudioDummyOptions | None, audio_overrides),
+                overrides=mm_options.get("audio"),
             )
         return mm_data
 
@@ -182,20 +178,19 @@ class InklingMultiModalProcessor(BaseMultiModalProcessor[InklingProcessingInfo])
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
         )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
         prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
 
-        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        processor = self.info.get_hf_processor(**hf_kwargs)
         tokenizer = self.info.get_tokenizer()
 
         images = mm_data.get("images") or []
-        audios = mm_data.get("audios") or []
+        audios = mm_data.get("audio") or []
         if not isinstance(images, list):
             images = list(cast(Iterable[Any], images))
         if not isinstance(audios, list):
@@ -235,8 +230,9 @@ def _apply_hf_processor_main(
             data["num_audio_tokens"] = torch.tensor(num_audio_tokens, dtype=torch.int64)
 
         processed_data = BatchFeature(data=data, tensor_type=None)
-        processed_data.update(passthrough_data)
-        return processed_data
+        return self._finalize_hf_mm_data(
+            mm_data, hf_kwargs, passthrough_data, processed_data
+        )
 
     def _tokenize_with_placeholders(
         self,
diff --git a/vllm/models/inkling/configs.py b/vllm/models/inkling/configs.py
index 2f8a247aec6b..58ed5ea00a82 100644
--- a/vllm/models/inkling/configs.py
+++ b/vllm/models/inkling/configs.py
@@ -8,10 +8,10 @@
 from typing import Any, ClassVar, Literal, cast
 
 import torch
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class InklingModelConfig(PretrainedConfig):
+class InklingModelConfig(PreTrainedConfig):
     model_type = "inkling_model"
     keys_to_ignore_at_inference = ["past_key_values"]
 
@@ -215,7 +215,7 @@ def tp_local_kv_conv_dim(num_kv_heads: int, head_dim: int) -> int:
         return TMLConvCacheParams(shape=shape, layers=self.conv_layer_ids, dtype=dtype)
 
 
-class InklingAudioConfig(PretrainedConfig):
+class InklingAudioConfig(PreTrainedConfig):
     model_type = "inkling_audio_model"
 
     def __init__(
@@ -255,7 +255,7 @@ def __init__(
         super().__init__(**kwargs)
 
 
-class InklingVisionConfig(PretrainedConfig):
+class InklingVisionConfig(PreTrainedConfig):
     model_type = "inkling_vision_model"
 
     def __init__(
@@ -295,10 +295,10 @@ def __init__(
         super().__init__(**kwargs)
 
 
-class InklingMMConfig(PretrainedConfig):
+class InklingMMConfig(PreTrainedConfig):
     model_type = "inkling_mm_model"
     keys_to_ignore_at_inference = ["past_key_values"]
-    sub_configs: ClassVar[dict[str, type[PretrainedConfig]]] = {
+    sub_configs: ClassVar[dict[str, type[PreTrainedConfig]]] = {
         "text_config": InklingModelConfig,
         "audio_config": InklingAudioConfig,
         "vision_config": InklingVisionConfig,
diff --git a/vllm/models/inkling/nvidia/logits_processor.py b/vllm/models/inkling/nvidia/logits_processor.py
index 9cf1c663c891..4302c6f674f8 100644
--- a/vllm/models/inkling/nvidia/logits_processor.py
+++ b/vllm/models/inkling/nvidia/logits_processor.py
@@ -37,6 +37,7 @@ class InklingLogitsProcessor(LogitsProcessor):
         soft_cap: Optional logit soft cap (``None`` for the served checkpoint).
         logits_mup_width_multiplier: muP width divisor for the final logits;
             ``None`` or ``0`` disables it.
+
     """
 
     def __init__(
diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py
index a02bf90d38cd..ed631aea7e85 100644
--- a/vllm/models/kimi_k3/amd/kda.py
+++ b/vllm/models/kimi_k3/amd/kda.py
@@ -271,8 +271,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         positions: torch.Tensor,
-        output: torch.Tensor,
-    ) -> None:
+    ) -> torch.Tensor:
         num_tokens = hidden_states.size(0)
         projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
 
@@ -306,7 +305,7 @@ def forward(
             core_attn_out=core_attn_out,
         )
         core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
-        output[:] = self.o_proj(core_attn_out)[0]
+        return self.o_proj(core_attn_out)[0]
 
     @eager_break_during_capture
     def _forward(
diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py
index 83c953e1a1d4..96e71700cd7d 100644
--- a/vllm/models/kimi_k3/amd/latent_moe_runner.py
+++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py
@@ -61,9 +61,7 @@ def _shard_up_proj_tail(
         shared_output: torch.Tensor,
         trunc_size: int | None,
     ) -> torch.Tensor:
-        """
-        Tier 2: column-parallel up-projection folded into the final reduce.
-        """
+        """Tier 2: column-parallel up-projection folded into the final reduce."""
         if not self._logged_sharded_tail:
             self._logged_sharded_tail = True
             logger.info_once(
diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py
index a74859efb884..65dad661ab12 100644
--- a/vllm/models/kimi_k3/amd/linear.py
+++ b/vllm/models/kimi_k3/amd/linear.py
@@ -317,9 +317,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 class KimiMLAAttention(nn.Module):
-    """
-    Main reference: DeepseekV2 vllm Implementation
-    """
+    """Main reference: DeepseekV2 vllm Implementation."""
 
     def __init__(
         self,
@@ -458,9 +456,8 @@ def forward(
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
-        output: torch.Tensor,
-    ) -> None:
-        output[:] = self.mla_attn(positions, hidden_states)
+    ) -> torch.Tensor:
+        return self.mla_attn(positions, hidden_states)
 
 
 class KimiDecoderLayer(nn.Module):
@@ -584,25 +581,29 @@ def _run_self_attn(
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
-        attn_output = torch.empty_like(hidden_states)
-        self.self_attn(
+        return self.self_attn(
             hidden_states=hidden_states,
             positions=positions,
-            output=attn_output,
         )
-        return attn_output
 
     def forward(
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         residual: torch.Tensor | None,
+        prefix_delta: torch.Tensor | None = None,
         **kwargs,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    ):
         if self.use_attn_residuals:
             assert residual is not None
-            return self.forward_attn_residual(positions, hidden_states, residual)
+            return self.forward_attn_residual(
+                positions, hidden_states, residual, prefix_delta
+            )
 
+        assert prefix_delta is None
         # Self Attention
         if residual is None:
             residual = hidden_states
@@ -622,7 +623,8 @@ def forward_attn_residual(
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         block_residual: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prefix_delta: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         prefix_sum = hidden_states
         hidden_states = _apply_attn_res(
             prefix_sum,
@@ -630,6 +632,7 @@ def forward_attn_residual(
             self.self_attention_res_proj,
             self.self_attention_res_norm,
             self.prev_valid_blocks,
+            delta=prefix_delta,
             output_norm=self.input_layernorm,
             block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1),
         )
@@ -659,8 +662,7 @@ def forward_attn_residual(
         )
 
         hidden_states = self.mlp(hidden_states)
-        prefix_sum = prefix_sum + hidden_states
-        return prefix_sum, block_residual
+        return prefix_sum, block_residual, hidden_states
 
 
 class KimiLinearModel(nn.Module, EagleModelMixin):
@@ -815,24 +817,28 @@ def forward(
         if residual is not None:
             block_residual[:, : residual.size(1), :].copy_(residual)
         residual = block_residual
+        prefix_delta = None
 
         for layer_idx, layer in enumerate(
             self.layers[self.start_layer : self.end_layer],
             start=self.start_layer,
         ):
-            hidden_states, residual = layer(
+            hidden_states, residual, prefix_delta = layer(
                 positions=positions,
                 hidden_states=hidden_states,
                 residual=residual,
+                prefix_delta=prefix_delta,
             )
             if (layer_idx + 1) in self.aux_hidden_state_layers:
-                # AMD attn-res layer already returns prefix_sum + MLP delta as
-                # hidden_states; the override drops the block bank in residual.
                 self._maybe_add_hidden_state(
-                    aux_hidden_states, layer_idx + 1, hidden_states, residual
+                    aux_hidden_states,
+                    layer_idx + 1,
+                    hidden_states + prefix_delta,
+                    residual,
                 )
 
         if not get_pp_group().is_last_rank:
+            hidden_states = hidden_states + prefix_delta
             return IntermediateTensors(
                 {"hidden_states": hidden_states, "residual": residual}
             )
@@ -843,6 +849,7 @@ def forward(
             self.output_attn_res_proj,
             self.output_attn_res_norm,
             attn_res_block_num,
+            delta=prefix_delta,
         )
         # NOTE: the final norm is applied in compute_logits instead of here, so
         # the MTP draft model receives the pre-norm hidden states.
diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py
index efcdb1bc75d1..556373c084aa 100644
--- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py
+++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py
@@ -213,9 +213,17 @@ def fused_kda_chunk(
 
     Args:
         qg: ``q * exp2(gk_cumsum)``, ``[1, T, H, 128]``.
+        w: chunk-local WY representation of the gated keys.
+        u: chunk-local WY representation of the values.
         kg_t: chunk-major transposed gated keys, ``[chunks, H, 128, 64]``.
+        aqk: intra-chunk query-key attention, ``[chunks, H, 64, 64]``.
         decay: ``exp2`` of each chunk's last gate row, ``[chunks, H, 128]``.
         out: output buffer, ``[1, T, H, 128]``; may alias ``u``'s source.
+        scale: scale applied to the query-key products.
+        cu_seqlens: int32 cumulative sequence lengths.
+        initial_state: fp32 per-sequence initial recurrent state, or ``None``.
+        output_final_state: whether to return the final recurrent state.
+        chunk_offsets: int32 per-sequence first chunk index.
         checkpoint_state: destination for the mid-prefill state snapshots,
             fp32 ``[rows, H, 128, 128]``. Without
             ``checkpoint_state_indices`` it is a staging buffer indexed by
@@ -239,6 +247,7 @@ def fused_kda_chunk(
         has_initial_state: bool ``[N]``; a false entry starts that sequence
             from a zero state and its cache row is read only, never before the
             walk writes it.
+
     """
     if state_cache is not None:
         if initial_state is not None or output_final_state:
@@ -286,9 +295,9 @@ def fused_kda_chunk(
         None if checkpoint_offsets is None else checkpoint_offsets.to(torch.int32),
         None
         if checkpoint_state_indices is None
-        else checkpoint_state_indices.to(torch.int32),
+        else checkpoint_state_indices.to(torch.int32).contiguous(),
         state_cache,
-        None if state_indices is None else state_indices.to(torch.int32),
+        None if state_indices is None else state_indices.to(torch.int32).contiguous(),
         has_initial_state,
     )
     return out, final_state
diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py
index 4c85ce8624f6..56cb3e148291 100644
--- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py
+++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py
@@ -52,6 +52,22 @@ def chunk_kda_prefill(
     """Run chunk KDA from raw gate and beta projections.
 
     Args:
+        q: query tensor, ``[1, T, H, K]``.
+        k: key tensor, ``[1, T, H, K]``.
+        v: value tensor, ``[1, T, H, V]``.
+        raw_g: raw gate projection, before the activation.
+        raw_beta: raw beta projection, before the activation.
+        A_log: log of the per-head gate decay.
+        g_bias: optional per-head gate bias.
+        scale: scale applied to the query-key products. Defaults to
+            ``k.shape[-1] ** -0.5``.
+        initial_state: fp32 per-sequence initial recurrent state, or ``None``.
+        output_final_state: whether to return the final recurrent state.
+        lower_bound: optional floor applied to the gate.
+        use_qk_l2norm_in_kernel: fold the q/k L2 norm into the kernel.
+        cu_seqlens: int32 cumulative sequence lengths.
+        chunk_indices: precomputed chunk indices for ``cu_seqlens``.
+        chunk_offsets: int32 per-sequence first chunk index.
         use_fused_chunk: request the two-kernel ROCm path. It is used only when
             every one of its preconditions holds; otherwise the Triton path
             runs unchanged.
@@ -71,6 +87,7 @@ def chunk_kda_prefill(
 
     Returns:
         The output and, when requested, the final recurrent state.
+
     """
     if scale is None:
         scale = k.shape[-1] ** -0.5
diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py
index 7be50dcd04ec..99eb8c45bdd8 100644
--- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py
+++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py
@@ -911,13 +911,12 @@ def fused_kda_gate(
     threshold: float = 20.0,
     lower_bound: float | None = None,
 ) -> torch.Tensor:
-    """
-    Forward pass for KDA gate:
-      input g: [..., H*D]
-      param A: [H] or [1, 1, H, 1]
-      beta: softplus beta parameter
-      threshold: softplus threshold parameter
-      return  : [..., H, D]
+    """Forward pass for KDA gate:
+    input g: [..., H*D]
+    param A: [H] or [1, 1, H, 1]
+    beta: softplus beta parameter
+    threshold: softplus threshold parameter
+    return  : [..., H, D]
     """
     orig_shape = g.shape[:-1]
 
diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py
index ca3bc8f75e33..8b9aa86f9ed9 100644
--- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py
+++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py
@@ -59,8 +59,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
     USE_SAFE_GATE: tl.constexpr,
     SOLVE_TRIL_DOT_PRECISION: tl.constexpr,
 ):
-    """
-    Fused kernel: compute inter-subchunk Akk + solve_tril in one pass.
+    """Fused kernel: compute inter-subchunk Akk + solve_tril in one pass.
     Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd.
 
     This kernel:
diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py
index cc448d95c0eb..86654f264431 100644
--- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py
+++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py
@@ -151,8 +151,7 @@ def chunk_kda_fwd_intra_token_parallel(
     chunk_size: int = 64,
     sub_chunk_size: int = 16,
 ) -> None:
-    """
-    Token-parallel implementation: each token gets its own thread block.
+    """Token-parallel implementation: each token gets its own thread block.
     Supports both fixed-length and variable-length sequences.
     Reduces wasted computation on padding.
 
@@ -164,10 +163,12 @@ def chunk_kda_fwd_intra_token_parallel(
         gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA)
         beta: [B, T, HV]
         Aqk: [B, T, HV, BT] output tensor to write to
+        cu_seqlens: cumulative sequence lengths for variable-length input
         Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32)
         scale: attention scale
         chunk_size: BT (default 64)
         sub_chunk_size: BC (default 16)
+
     """
     B, T, H, K, HV = *q.shape, gk.shape[2]
     N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
diff --git a/vllm/models/kimi_k3/common/mm_preprocess.py b/vllm/models/kimi_k3/common/mm_preprocess.py
index cfe603b47b41..dfd90fe9a969 100644
--- a/vllm/models/kimi_k3/common/mm_preprocess.py
+++ b/vllm/models/kimi_k3/common/mm_preprocess.py
@@ -9,7 +9,7 @@
 import torch
 from transformers import BatchFeature
 
-from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.config.multimodal import ImageDummyOptions, MultiModalDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.multimodal.inputs import (
@@ -204,7 +204,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+        mm_options: MultiModalDummyOptions | None = None,
     ) -> MultiModalDataDict:
         media_proc_cfg = self.info.image_processor.media_proc_cfg
         max_size = self.info.get_max_image_size(
@@ -214,7 +214,6 @@ def get_dummy_mm_data(
             media_proc_cfg["patch_limit_on_one_side"],
             media_proc_cfg["fixed_output_tokens"],
         )
-        num_images = mm_counts.get("image", 0)
         image_overrides = cast(
             ImageDummyOptions | None,
             mm_options.get("image") if mm_options else None,
@@ -223,7 +222,7 @@ def get_dummy_mm_data(
             "image": self._get_dummy_images(
                 width=max_size.width,
                 height=max_size.height,
-                num_images=num_images,
+                num_images=mm_counts.get("image", 0),
                 overrides=image_overrides,
             )
         }
diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py
index 97febda0c86e..0baff884ccaa 100644
--- a/vllm/models/kimi_k3/nvidia/dspark_mla.py
+++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py
@@ -15,7 +15,10 @@
     ReplicatedLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
+from vllm.model_executor.models.qwen3_dspark import (
+    DSparkConfidenceHead,
+    DSparkMarkovHead,
+)
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     WeightsMapper,
@@ -28,6 +31,10 @@
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.worker.workspace import current_workspace_manager
 
+_GROUPED_KV_CACHE_DTYPES = frozenset(
+    {"auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_e5m2"}
+)
+
 
 def _duplicate_context_kv_weights(
     weights: Iterable[tuple[str, torch.Tensor]], num_layers: int
@@ -184,6 +191,18 @@ def __init__(
             self.config.markov_rank,
             prefix=maybe_prefix(prefix, "markov_head"),
         )
+        self.confidence_head: DSparkConfidenceHead | None = None
+        if getattr(self.config, "enable_confidence_head", False):
+            with_markov = getattr(self.config, "confidence_head_with_markov", False)
+            input_dim = self.config.hidden_size + (
+                self.config.markov_rank if with_markov else 0
+            )
+            self.confidence_head = DSparkConfidenceHead(
+                input_dim,
+                prefix=maybe_prefix(prefix, "confidence_head"),
+                bias=True,
+                with_markov=with_markov,
+            )
         self._max_num_context_tokens = (
             vllm_config.scheduler_config.max_num_batched_tokens
         )
@@ -233,6 +252,13 @@ def _build_fused_context_kv_metadata(self) -> None:
         self._context_kv_lora_rank = attn0.kv_lora_rank
         self._context_rope_dim = attn0.qk_rope_head_dim
         self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon
+        self._context_kv_scales: torch.Tensor | None = None
+        if attn0.kv_cache_dtype in _GROUPED_KV_CACHE_DTYPES and is_quantized_kv_cache(
+            attn0.kv_cache_dtype
+        ):
+            self._context_kv_scales = torch.stack(
+                [attn._k_scale.reshape(()) for attn in attentions]
+            )
 
     def _precompute_fused_context_kv(
         self,
@@ -286,16 +312,18 @@ def _precompute_fused_context_kv(
             return
 
         cache_layers = [layer.self_attn for layer in self.layers]
+        cache_dtype = cache_layers[0].kv_cache_dtype
         if (
-            not is_quantized_kv_cache(cache_layers[0].kv_cache_dtype)
+            cache_dtype in _GROUPED_KV_CACHE_DTYPES
+            and all(cl.kv_cache_dtype == cache_dtype for cl in cache_layers)
             and self._has_uniform_block_layout(cache_layers)
             and (
                 isinstance(context_slot_mapping, torch.Tensor)
                 or all(s is not None for s in context_slot_mapping)
             )
         ):
-            # Grouped context KV insert only supports unquantized (bf16) KV cache
-            # and assumes that all layers share the same block layout.
+            # Grouped context KV insert assumes that all layers share the same
+            # cache dtype and block layout.
 
             if isinstance(context_slot_mapping, (list, tuple)):
                 per_layer_slot_mappings = [
@@ -321,6 +349,8 @@ def _precompute_fused_context_kv(
                 ref_cache.size(1),
                 ref_cache.stride(0),
                 ref_cache.stride(1),
+                self._context_kv_scales,
+                cache_dtype,
             )
             return
 
@@ -421,6 +451,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         assert vllm_config.speculative_config is not None
         self.draft_model_config = vllm_config.speculative_config.draft_model_config
         self.config = self.draft_model_config.hf_config
+        if getattr(self.config, "enable_confidence_head", False):
+            # The class mapper drops the head as training-only; keep it for drafts
+            # that actually built one, since adaptive verification reads it.
+            self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
+                orig_to_new_substr={"confidence_head": "confidence_head"}
+            )
         target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()
         self.model = K3DSparkModel(
             vllm_config=vllm_config,
@@ -479,10 +515,18 @@ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
     def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
         return self.model.markov_head.bias(markov_embed, self.logits_processor)
 
+    def compute_confidence(
+        self, head_hidden: torch.Tensor, markov_embed: torch.Tensor
+    ) -> torch.Tensor:
+        """Per-position acceptance probability for each drafted token."""
+        assert self.model.confidence_head is not None
+        return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
         # read: 1. all weights. 2. context kv weights
         weights = _duplicate_context_kv_weights(weights, len(self.model.layers))
-        loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
+    def process_weights_after_loading(self) -> None:
         self.model._build_fused_context_kv_metadata()
-        return loaded_weights
diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py
index 92326944f592..97b402908454 100644
--- a/vllm/models/kimi_k3/nvidia/kda.py
+++ b/vllm/models/kimi_k3/nvidia/kda.py
@@ -21,6 +21,10 @@
     RowParallelLinear,
 )
 from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
+from vllm.model_executor.layers.mamba.kda_checkpoint import (
+    FlashKDAPrefillCheckpointExporter,
+    kda_prefill_checkpoint_alignment,
+)
 from vllm.model_executor.layers.mamba.mamba_utils import (
     MambaStateDtypeCalculator,
     MambaStateShapeCalculator,
@@ -49,7 +53,6 @@
 from vllm.platforms import current_platform
 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
-from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import (
     flashinfer_fused_kda_decode,
     flashinfer_recurrent_kda,
@@ -57,7 +60,6 @@
     has_flashinfer_recurrent_kda,
 )
 from vllm.v1.attention.backend import AttentionBackend
-from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.kv_cache_interface import MambaSpec
 from vllm.v1.worker.workspace import current_workspace_manager
 
@@ -389,69 +391,6 @@ def _flashinfer_kda_prefill(
     return output, initial_state
 
 
-@triton.jit
-def _store_cache_checkpoints_kernel(
-    x_ptr,
-    conv_state_ptr,
-    recurrent_checkpoint_ptr,
-    recurrent_state_ptr,
-    query_start_loc_ptr,
-    checkpoint_offsets_ptr,
-    checkpoint_state_indices_ptr,
-    x_stride_0: tl.constexpr,
-    x_stride_1: tl.constexpr,
-    state_stride_0: tl.constexpr,
-    state_stride_1: tl.constexpr,
-    state_stride_2: tl.constexpr,
-    checkpoint_stride_0: tl.constexpr,
-    recurrent_state_stride_0: tl.constexpr,
-    checkpoint_offset_stride: tl.constexpr,
-    STATE_LEN: tl.constexpr,
-    WIDTH: tl.constexpr,
-    RECURRENT_ROW_SIZE: tl.constexpr,
-    NULL_STATE_IDX: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    # store checkpoints to cache
-    seq_idx = tl.program_id(0)
-    cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64)
-    checkpoint_offset = tl.load(
-        checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
-    )
-    valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)
-    valid_conv = (
-        (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)
-    )
-    width_idx = cols // STATE_LEN
-    history_idx = cols % STATE_LEN
-    checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset
-    token_idx = checkpoint_end - STATE_LEN + history_idx
-    values = tl.load(
-        x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
-        mask=valid_conv,
-    )
-    tl.store(
-        conv_state_ptr
-        + state_idx * state_stride_0
-        + width_idx * state_stride_1
-        + history_idx * state_stride_2,
-        values,
-        mask=valid_conv,
-    )
-
-    valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint
-    recurrent = tl.load(
-        recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols,
-        mask=valid_recurrent,
-    )
-    tl.store(
-        recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,
-        recurrent,
-        mask=valid_recurrent,
-    )
-
-
 def resolve_kda_prefill_backend(
     backend: str,
     head_dim: int,
@@ -735,6 +674,7 @@ def __init__(
         self._flashinfer_kda_output_spec: tuple[tuple[int, ...], torch.dtype] | None = (
             None
         )
+        self._checkpoint_exporter: FlashKDAPrefillCheckpointExporter | None = None
         if self.kda_prefill_backend == "flashkda":
             T = vllm_config.scheduler_config.max_num_batched_tokens
             N = vllm_config.scheduler_config.max_num_seqs
@@ -748,6 +688,7 @@ def __init__(
                 ((N, H, D, D), self.get_state_dtype()[1]),
                 ((workspace_size,), torch.uint8),
             )
+            self._checkpoint_exporter = FlashKDAPrefillCheckpointExporter()
         elif self.kda_prefill_backend == "flashinfer":
             T = vllm_config.scheduler_config.max_num_batched_tokens
             H, D = self.local_num_heads, self.head_dim
@@ -783,7 +724,7 @@ def __init__(
         )
         self.gemm_rs_ar = None
         if run_gemm_rs_ar:
-            from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
+            from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
                 get_gemm_rs_ar,
             )
 
@@ -800,12 +741,11 @@ def __init__(
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
         spec = super().get_kv_cache_spec(vllm_config)
         assert isinstance(spec, MambaSpec)
+        alignment = kda_prefill_checkpoint_alignment(self.kda_prefill_backend)
         return replace(
             spec,
-            num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"),
-            prefill_checkpoint_alignment=(
-                16 if self.kda_prefill_backend == "flashkda" else None
-            ),
+            num_prefill_checkpoint_blocks=int(alignment is not None),
+            prefill_checkpoint_alignment=alignment,
         )
 
     def forward(
@@ -869,7 +809,7 @@ def forward(
         )
         core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
         if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(core_attn_out):
-            return self.gemm_rs_ar(core_attn_out, self.o_proj.weight)
+            return self.gemm_rs_ar.apply(core_attn_out, self.o_proj)
         return self.o_proj(core_attn_out)[0]
 
     @eager_break_during_capture
@@ -1173,39 +1113,14 @@ def _prefill_conv(
                         )
                         core_attn_out_non_spec = flashkda_out
                         last_recurrent_state = final_state
-                        state_len = conv_state.shape[-1]
-                        width = mixed_qkv_ns.shape[-1]
-                        recurrent_row_size = checkpoint_state[0].numel()
-                        block_size = 256
-                        _store_cache_checkpoints_kernel[
-                            (
-                                checkpoint_offsets.numel(),
-                                triton.cdiv(
-                                    max(width * state_len, recurrent_row_size),
-                                    block_size,
-                                ),
-                            )
-                        ](
-                            mixed_qkv_ns,
-                            conv_state,
-                            checkpoint_state,
-                            recurrent_state,
-                            non_spec_query_start_loc,
-                            checkpoint_offsets,
-                            checkpoint.state_indices,
-                            mixed_qkv_ns.stride(0),
-                            mixed_qkv_ns.stride(1),
-                            conv_state.stride(0),
-                            conv_state.stride(1),
-                            conv_state.stride(2),
-                            checkpoint_state.stride(0),
-                            recurrent_state.stride(0),
-                            checkpoint_offsets.stride(0),
-                            state_len,
-                            width,
-                            recurrent_row_size,
-                            NULL_BLOCK_ID,
-                            block_size,
+                        assert self._checkpoint_exporter is not None
+                        self._checkpoint_exporter.export(
+                            checkpoint,
+                            raw_qkv=mixed_qkv_ns,
+                            conv_state=conv_state,
+                            recurrent_checkpoint=checkpoint_state,
+                            recurrent_state=recurrent_state,
+                            cu_seqlens=non_spec_query_start_loc,
                         )
                     else:
                         (
diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py
index 0b395b0c6c8f..2562171109d9 100644
--- a/vllm/models/kimi_k3/nvidia/kda_metadata.py
+++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py
@@ -18,11 +18,14 @@
 import torch
 
 from vllm.config import VllmConfig
+from vllm.model_executor.layers.mamba.checkpoint import (
+    MambaPrefillCheckpointBuilder,
+    MambaPrefillCheckpointMetadata,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import async_tensor_h2d
-from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
 from vllm.v1.attention.backends.gdn_attn import (
     GDNAttentionBackend,
     GDNAttentionMetadata,
@@ -37,11 +40,7 @@
     compute_causal_conv1d_metadata,
     split_decodes_and_prefills,
 )
-from vllm.v1.kv_cache_interface import (
-    MambaSpec,
-    get_mamba_prefill_checkpoint_position,
-    is_mamba_prefill_checkpoint_valid,
-)
+from vllm.v1.kv_cache_interface import MambaSpec
 
 if TYPE_CHECKING:
     from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
@@ -262,12 +261,6 @@ class KDARecoverSSMCommitMetadata:
     align: KDARecoverSSMAlignMetadata | None
 
 
-@dataclass
-class KDACheckpointMetadata:
-    checkpoint_offsets: torch.Tensor
-    state_indices: torch.Tensor
-
-
 @dataclass
 class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
     spec_token_start: int | None = None
@@ -278,7 +271,7 @@ class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
     recoverssm_context: "KDARecoverSSMCommitContext | None" = field(
         default=None, repr=False, compare=False
     )
-    checkpoint: KDACheckpointMetadata | None = None
+    checkpoint: MambaPrefillCheckpointMetadata | None = None
 
     def commit_recoverssm_state(
         self, num_accepted_tokens: torch.Tensor
@@ -312,6 +305,10 @@ def commit_recoverssm_state(
 
 
 class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
+    # Overrides GDN's UNIFORM_BATCH: adaptive verification requires ALWAYS from every
+    # builder, and KDA reads per-request offsets off device within a fixed k+1 window,
+    # so one k+1 graph replays any 1..k+1 mix.
+    _cudagraph_support = AttentionCGSupport.ALWAYS
     mamba_aligned_state_indices: torch.Tensor | None = None
 
     def __init__(
@@ -322,6 +319,9 @@ def __init__(
         device: torch.device,
     ) -> None:
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self.checkpoint_builder = MambaPrefillCheckpointBuilder(
+            vllm_config, kv_cache_spec
+        )
         additional_config = vllm_config.additional_config
         self.use_flashinfer_prefill = (
             isinstance(additional_config, dict)
@@ -400,8 +400,10 @@ def build(  # type: ignore[override]
         if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
             spec_sequence_masks_cpu = None
             num_spec_decodes = 0
+            active_non_spec_mask_cpu = None
         else:
             spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
+            active_non_spec_mask_cpu = None
             if self.use_recoverssm:
                 assert m.is_prefilling is not None
                 assert m.is_prefilling.device.type == "cpu"
@@ -409,6 +411,12 @@ def build(  # type: ignore[override]
                     query_start_loc_cpu.diff() > 0
                 )
                 spec_sequence_masks_cpu |= active_decode_mask_cpu
+                # Adaptive's cost-table profiler probes eager dummy rows longer than
+                # RecoverSSM's num_spec+1 activation capacity; route them out instead
+                # of tripping the capacity check below.
+                spec_sequence_masks_cpu &= (
+                    query_start_loc_cpu.diff() <= self.num_spec + 1
+                )
             # Native KDA can use its regular decode path when no draft token
             # was scheduled. RecoverSSM must preserve its extended conv window.
             if (
@@ -418,6 +426,7 @@ def build(  # type: ignore[override]
             ):
                 spec_sequence_masks_cpu = None
                 num_spec_decodes = 0
+                active_non_spec_mask_cpu = None
             else:
                 num_spec_decodes = spec_sequence_masks_cpu.sum().item()
                 if num_spec_decodes == 0:
@@ -427,10 +436,33 @@ def build(  # type: ignore[override]
         spec_token_start = None
         non_spec_token_start = None
         if num_spec_decodes == 0:
-            # The runner orders ordinary decodes before prefills.
+            # V2 already excludes prefills from full decode graphs via has_prefill.
+            # Classify first chunks as prefills to mask recycled state;
+            # resumed one-token chunks can still use the decode kernels.
+            assert m.seq_lens_cpu_upper_bound is not None
+            query_lens_cpu = query_start_loc_cpu.diff()
+            no_prior_state = (query_lens_cpu > 0) & (
+                m.seq_lens_cpu_upper_bound <= query_lens_cpu
+            )
+            # Capture batches also have seq_len == query_len, but are not prefills.
+            if m.is_prefilling is not None:
+                no_prior_state &= m.is_prefilling
+            else:
+                no_prior_state = torch.zeros_like(no_prior_state)
             num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
-                split_decodes_and_prefills(m, decode_threshold=1)
+                split_decodes_and_prefills(
+                    m.replace(is_prefilling=no_prior_state),
+                    decode_threshold=1,
+                    treat_short_extends_as_decodes=False,
+                )
             )
+            # Exclude trailing padding from both prefill counts.
+            if num_prefills:
+                num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum())
+                num_prefill_tokens = (
+                    int(query_start_loc_cpu[num_decodes + num_prefills])
+                    - num_decode_tokens
+                )
             num_spec_decode_tokens = 0
             spec_token_indx = None
             non_spec_token_indx = None
@@ -601,80 +633,6 @@ def build(  # type: ignore[override]
         else:
             has_initial_state = None
 
-        checkpoint = None
-        checkpoint_enabled = (
-            self.vllm_config.cache_config.mamba_cache_mode == "align"
-            and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0
-        )
-        if num_prefills > 0 and checkpoint_enabled:
-            # prepare checkpoint metadata
-            assert m.seq_lens_cpu_upper_bound is not None
-            request_rows = list(range(m.num_reqs))
-            if spec_sequence_masks_cpu is not None:
-                # get non spec request rows
-                request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist()
-            all_query_lens = query_start_loc_cpu.diff().tolist()
-            query_lens = [all_query_lens[row] for row in request_rows]
-            seq_lens = m.seq_lens_cpu_upper_bound.tolist()
-            block_size = self.kv_cache_spec.block_size
-            hash_block_size = (
-                self.vllm_config.cache_config.prefix_match_unit or block_size
-            )
-            speculative_config = self.vllm_config.speculative_config
-            drop_eagle_block = (
-                speculative_config is not None
-                and speculative_config.use_eagle_block_drop()
-            )
-            checkpoint_splits = []
-            checkpoint_cols = []
-            for row, query_len in zip(request_rows, query_lens):
-                seq_len = seq_lens[row]
-                query_start = seq_len - query_len
-                checkpoint_position = get_mamba_prefill_checkpoint_position(
-                    seq_len,
-                    hash_block_size,
-                    drop_eagle_block=drop_eagle_block,
-                )
-                offset = checkpoint_position - query_start
-                valid = is_mamba_prefill_checkpoint_valid(
-                    query_start=query_start,
-                    query_end=seq_len,
-                    checkpoint_position=checkpoint_position,
-                    hash_block_size=hash_block_size,
-                    mamba_block_size=block_size,
-                    checkpoint_alignment=(
-                        self.kv_cache_spec.prefill_checkpoint_alignment
-                    ),
-                )
-                offset = offset if valid else 0
-                first_len = offset or query_len
-                checkpoint_splits.append((first_len, query_len - first_len))
-                checkpoint_cols.append(cdiv(seq_len, block_size) - 2 if valid else -1)
-            if any(tail for _, tail in checkpoint_splits):
-                checkpoint_offsets_tensor = async_tensor_h2d(
-                    [first if tail else 0 for first, tail in checkpoint_splits],
-                    dtype=torch.int32,
-                    device=query_start_loc.device,
-                )
-                request_rows_tensor = async_tensor_h2d(
-                    request_rows, dtype=torch.int64, device=query_start_loc.device
-                )
-                checkpoint_cols_tensor = async_tensor_h2d(
-                    checkpoint_cols, dtype=torch.int64, device=query_start_loc.device
-                )
-                checkpoint_state_indices = m.block_table_tensor[
-                    request_rows_tensor, checkpoint_cols_tensor
-                ]
-                checkpoint_state_indices = torch.where(
-                    checkpoint_cols_tensor >= 0,
-                    checkpoint_state_indices,
-                    NULL_BLOCK_ID,
-                )
-                checkpoint = KDACheckpointMetadata(
-                    checkpoint_offsets_tensor,
-                    checkpoint_state_indices,
-                )
-
         # Prepare per-request tensors for cudagraph replay. num_actual_tokens
         # may be token-padded, while state/query/acceptance metadata is indexed
         # by request.
@@ -753,6 +711,13 @@ def build(  # type: ignore[override]
                     flashinfer_prefill_query_start_loc.diff(), descending=True
                 ).to(torch.int32)
 
+        checkpoint = None
+        if num_prefills > 0:
+            request_rows = list(range(m.num_reqs))
+            if active_non_spec_mask_cpu is not None:
+                request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist()
+            checkpoint = self.checkpoint_builder.build(m, request_rows)
+
         return KimiK3KDAMetadata(
             num_prefills=num_prefills,
             num_prefill_tokens=num_prefill_tokens,
@@ -795,3 +760,9 @@ def get_name() -> str:
     @staticmethod
     def get_builder_cls() -> type[KimiK3KDAMetadataBuilder]:
         return KimiK3KDAMetadataBuilder
+
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        # Pure spec-decode reads its plan from DEVICE offsets, so CPU/device
+        # mismatch is fine; the mixed branch still assumes CPU totals match.
+        return True
diff --git a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py
index 450e46a8e4ec..8eedb5ba89ed 100644
--- a/vllm/models/kimi_k3/nvidia/latent_moe_runner.py
+++ b/vllm/models/kimi_k3/nvidia/latent_moe_runner.py
@@ -275,9 +275,7 @@ def _shard_up_proj_tail(
         shared_output: torch.Tensor,
         trunc_size: int | None,
     ) -> torch.Tensor:
-        """
-        Tier 2: column-parallel up-projection folded into the final reduce.
-        """
+        """Tier 2: column-parallel up-projection folded into the final reduce."""
         transform = self.routed_output_transform
         assert transform is not None
 
diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py
index 02be89aa311a..cd55362759aa 100644
--- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py
@@ -65,6 +65,20 @@
 _KDA_TP_SIZE = 8
 
 
+def _kda_qkvg_flashinfer_backend() -> Literal["cute-dsl"] | None:
+    """Return the tested FlashInfer backend for the low-M QKVG branch.
+
+    FlashInfer's ``cute-dsl`` BF16 GEMM supports SM100 and SM103, but rejects
+    SM90 during backend validation. Other capabilities use ``torch.mm`` for
+    QKVG while retaining the concurrent F_A/beta and F_B branch.
+    """
+    if current_platform.is_device_capability(
+        (10, 0)
+    ) or current_platform.is_device_capability((10, 3)):
+        return "cute-dsl"
+    return None
+
+
 @dataclass(frozen=True, slots=True)
 class ProjectionSpec:
     n: int
@@ -737,6 +751,7 @@ def run_kda_projection_overlap(
 
     Returns:
         The Q/K/V/G projection, F_B output, and beta projection.
+
     """
     num_tokens = hidden_states.shape[0]
     qkvg_weight = packed_weight[:_KDA_QKVG_SIZE]
@@ -749,14 +764,16 @@ def run_qkvg() -> torch.Tensor:
                 config,
                 None,
             )
-        if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS:
+        if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS and (
+            backend := _kda_qkvg_flashinfer_backend()
+        ):
             from flashinfer.gemm import mm_bf16
 
             return mm_bf16(
                 hidden_states,
                 qkvg_weight.t(),
                 pdl=True,
-                backend="cute-dsl",
+                backend=backend,
             )
         return torch.mm(hidden_states, qkvg_weight.t())
 
@@ -807,11 +824,10 @@ def run_fab_fb() -> tuple[torch.Tensor, torch.Tensor]:
 
 
 def autotune_kda_qkvg(model: nn.Module) -> None:
-    """Autotune the FlashInfer QKVG GEMM before CUDA graph capture."""
-    from flashinfer.gemm import mm_bf16
-
+    """Autotune the supported QKVG GEMM before CUDA graph capture."""
     from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
 
+    backend = _kda_qkvg_flashinfer_backend()
     children: list[KimiK3DeltaAttention] = []
     weights_by_shape: dict[
         tuple[int, int, int, torch.dtype, torch.device], torch.Tensor
@@ -822,29 +838,33 @@ def autotune_kda_qkvg(model: nn.Module) -> None:
         if child._projection_overlap_max_tokens <= 0:
             continue
         children.append(child)
-        qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
-        shape = (
-            KDA_PROJECTION_OVERLAP_MAX_TOKENS,
-            qkvg_weight.shape[0],
-            qkvg_weight.shape[1],
-            qkvg_weight.dtype,
-            qkvg_weight.device,
-        )
-        weights_by_shape.setdefault(shape, qkvg_weight)
-
-    for shape, qkvg_weight in weights_by_shape.items():
-        num_tokens = shape[0]
-        hidden_states = torch.empty(
-            (num_tokens, qkvg_weight.shape[1]),
-            dtype=qkvg_weight.dtype,
-            device=qkvg_weight.device,
-        )
-        mm_bf16(
-            hidden_states,
-            qkvg_weight.t(),
-            pdl=True,
-            backend="cute-dsl",
-        )
+        if backend is not None:
+            qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
+            shape = (
+                KDA_PROJECTION_OVERLAP_MAX_TOKENS,
+                qkvg_weight.shape[0],
+                qkvg_weight.shape[1],
+                qkvg_weight.dtype,
+                qkvg_weight.device,
+            )
+            weights_by_shape.setdefault(shape, qkvg_weight)
+
+    if backend is not None:
+        from flashinfer.gemm import mm_bf16
+
+        for shape, qkvg_weight in weights_by_shape.items():
+            num_tokens = shape[0]
+            hidden_states = torch.empty(
+                (num_tokens, qkvg_weight.shape[1]),
+                dtype=qkvg_weight.dtype,
+                device=qkvg_weight.device,
+            )
+            mm_bf16(
+                hidden_states,
+                qkvg_weight.t(),
+                pdl=True,
+                backend=backend,
+            )
     for child in children:
         child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS
 
diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py
index 87bf59749c1d..188f72e66016 100644
--- a/vllm/models/kimi_k3/nvidia/mla.py
+++ b/vllm/models/kimi_k3/nvidia/mla.py
@@ -58,6 +58,7 @@
     accumulate_mla_context_chunk,
     init_mla_context_partial,
     neutralize_empty_context_partials,
+    split_kv_b_proj,
 )
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -73,9 +74,6 @@
     QuantizationConfig,
     resolve_quant_method,
 )
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    get_and_maybe_dequant_weights,
-)
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope
 from vllm.model_executor.utils import replace_parameter
 from vllm.models.common.ops import fused_q_kv_rmsnorm
@@ -165,9 +163,9 @@ def __init__(
             rope_parameters = dict(config.rope_parameters)
             if rope_parameters["rope_type"] != "default":
                 rope_parameters["rope_type"] = (
-                    "deepseek_yarn"
-                    if rope_parameters.get("apply_yarn_scaling", True)
-                    else "deepseek_llama_scaling"
+                    "deepseek_llama_scaling"
+                    if rope_parameters.get("attention_factor") == 1.0
+                    else "deepseek_yarn"
                 )
             self.rotary_emb = get_rope(
                 qk_rope_head_dim,
@@ -301,7 +299,7 @@ def __init__(
         )
         self.gemm_rs_ar = None
         if run_gemm_rs_ar:
-            from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
+            from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
                 get_gemm_rs_ar,
             )
 
@@ -441,20 +439,13 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
         projected into latent space by ``W_UK_T`` and the attention output is
         projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V.
         """
-        kv_b_proj_weight = get_and_maybe_dequant_weights(
-            self.kv_b_proj, out_dtype=act_dtype
-        ).T
-        assert kv_b_proj_weight.shape == (
-            self.kv_lora_rank,
-            self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim),
-        ), f"{kv_b_proj_weight.shape=}"
-        kv_b_proj_weight = kv_b_proj_weight.view(
+        W_UK, W_UV = split_kv_b_proj(
+            self.kv_b_proj,
+            act_dtype,
             self.kv_lora_rank,
             self.num_local_heads,
-            self.qk_nope_head_dim + self.v_head_dim,
-        )
-        W_UK, W_UV = kv_b_proj_weight.split(
-            [self.qk_nope_head_dim, self.v_head_dim], dim=-1
+            self.qk_nope_head_dim,
+            self.v_head_dim,
         )
         # (L, N, V) -> (N, L, V)
         replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
@@ -631,7 +622,7 @@ def forward(
             attn_out = _gate_sigmoid_mul(attn_out, gate)
 
         if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(attn_out):
-            return self.gemm_rs_ar(attn_out, self.o_proj.weight)
+            return self.gemm_rs_ar.apply(attn_out, self.o_proj)
 
         return self.o_proj(attn_out)[0]
 
diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py
index c4f0092f8d9b..c1818a444b53 100644
--- a/vllm/models/kimi_k3/nvidia/model.py
+++ b/vllm/models/kimi_k3/nvidia/model.py
@@ -160,54 +160,23 @@ def maybe_init_gemm_rs_ar(vllm_config: VllmConfig, use_sequence_parallel: bool)
     # its singleton to exactly one mode.
     all_reduce = not use_sequence_parallel
     mode = "GEMM-AR" if all_reduce else "GEMM-RS"
-    enabled = envs.VLLM_KIMI_K3_GEMM_AR if all_reduce else envs.VLLM_KIMI_K3_GEMM_RS
+    enabled = envs.VLLM_KIMI_K3_GEMM_AR if all_reduce else envs.VLLM_ENABLE_GEMM_RS
     if not enabled:
         return False
 
-    parallel_config = vllm_config.parallel_config
-    tp_size = parallel_config.tensor_parallel_size
-    if parallel_config.use_ubatching:
-        reason = "ubatching is enabled"
-    elif vllm_config.model_config.dtype != torch.bfloat16:
-        reason = "the model dtype is not BF16"
-    elif not current_platform.is_cuda():
-        reason = "the device is not CUDA"
-    elif not current_platform.is_device_capability_family(100):
-        reason = "the device is not SM100-family"
-    elif not 1 < tp_size <= 16:
-        reason = "TP size is not in the supported range 2-16"
-    elif 128 % tp_size != 0:
-        reason = "TP size does not divide 128"
-    else:
-        reason = None
-
-    if reason is not None:
-        logger.warning_once("%s is disabled because %s.", mode, reason)
-        return False
-
-    from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import init_gemm_rs_ar
+    # The kernel module pulls in cute_dsl, so import it only once the flag
+    # asks for it.
+    from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
+        maybe_init_gemm_rs_ar as maybe_init_shared_gemm_rs_ar,
+    )
 
     config = vllm_config.model_config.hf_text_config
-    try:
-        init_gemm_rs_ar(
-            max_M=vllm_config.scheduler_config.max_num_batched_tokens,
-            N=config.hidden_size,
-            all_reduce=all_reduce,
-        )
-    except RuntimeError as e:
-        logger.warning_once(
-            "%s is disabled because initialization failed: %s. This may mean "
-            "the TP ranks do not share one NVLink domain.",
-            mode,
-            e,
-        )
+    if not maybe_init_shared_gemm_rs_ar(
+        vllm_config, N=config.hidden_size, all_reduce=all_reduce
+    ):
         return False
-    if all_reduce:
-        logger.info_once(
-            "GEMM-AR is enabled. To disable it, set VLLM_KIMI_K3_GEMM_AR=0."
-        )
-    else:
-        logger.info_once("GEMM-RS is enabled.")
+    flag = "VLLM_KIMI_K3_GEMM_AR" if all_reduce else "VLLM_ENABLE_GEMM_RS"
+    logger.info_once("To disable %s, set %s=0.", mode, flag)
     return True
 
 
@@ -276,7 +245,7 @@ def __init__(
             not use_sequence_parallel and reduce_results
         )
         if use_gemm_rs_ar and run_gemm_rs_ar:
-            from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import (
+            from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
                 get_gemm_rs_ar,
             )
 
@@ -309,7 +278,7 @@ def forward(self, x):
         x = self.act_fn(gate_up)
 
         if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(x):
-            return self.gemm_rs_ar(x, self.down_proj.weight)
+            return self.gemm_rs_ar.apply(x, self.down_proj)
 
         x, _ = self.down_proj(x)
         if self.shard_sequence_parallel:
@@ -339,6 +308,7 @@ def forward(
             residual: Optional tensor of the up-projection's output shape to
                 accumulate into. It is consumed in the GEMM's beta-add
                 epilogue, so adding it costs no extra kernel.
+
         """
         if self.norm is not None:
             hidden_states = self.norm(hidden_states)
@@ -532,8 +502,10 @@ def forward(
             symm_buffer,
             activation_clamp=activation_clamp,
             activation=self.activation,
-            situ_beta=self.activation_beta,
-            situ_linear_beta=self.activation_linear_beta,
+            # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the
+            # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched.
+            activation_alpha=self.activation_beta or 1.0,
+            activation_beta=self.activation_linear_beta or 0.0,
             fast_math=fast_math,
         )
         return y
@@ -674,7 +646,7 @@ def __init__(
                 hidden_size,
                 self.moe_hidden_size,
                 bias=False,
-                quant_config=None,
+                quant_config=quant_config,
                 prefix=f"{prefix}.routed_expert_down_proj",
             )
             self.routed_expert_norm = (
@@ -691,7 +663,7 @@ def __init__(
                 self.moe_hidden_size,
                 hidden_size,
                 bias=False,
-                quant_config=None,
+                quant_config=quant_config,
                 prefix=f"{prefix}.routed_expert_up_proj",
             )
 
@@ -758,6 +730,7 @@ def __init__(
                 routed_input_transform=None,
                 routed_output_transform=self.routed_output_transform,
                 is_sequence_parallel=use_sequence_parallel,
+                skip_padding=True,
                 runner_cls=LatentMoERunner if self.use_latent_moe else None,
             )
         if self.padded_moe_intermediate_size != moe_intermediate_size:
@@ -794,6 +767,7 @@ def _maybe_overlap_router_and_down_proj(
             ``router_output`` holds the grouped top-k weights and ``topk_ids``
             the selected experts; otherwise ``router_output`` holds the raw gate
             logits and ``topk_ids`` is ``None``.
+
         """
 
         def _router(
@@ -917,14 +891,12 @@ def __init__(
                     aux_stream=aux_stream,
                     run_gemm_rs_ar=run_gemm_rs_ar,
                 )
-                self._self_attn_writes_output = False
             else:
                 self.self_attn = KimiLinearGatedDeltaNetAttention(
                     config,
                     vllm_config,
                     prefix=f"{prefix}.self_attn",
                 )
-                self._self_attn_writes_output = True
         else:
             qk_nope_head_dim = config.qk_nope_head_dim
             qk_rope_head_dim = config.qk_rope_head_dim
@@ -954,7 +926,6 @@ def __init__(
                 aux_stream=aux_stream,
                 run_gemm_rs_ar=run_gemm_rs_ar,
             )
-            self._self_attn_writes_output = False
 
         if self.use_sequence_parallel:
             self.self_attn.o_proj.reduce_results = False
@@ -1020,14 +991,6 @@ def _run_self_attn(
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
-        if self._self_attn_writes_output:
-            output = torch.empty_like(hidden_states)
-            self.self_attn(
-                hidden_states=hidden_states,
-                positions=positions,
-                output=output,
-            )
-            return output
         return self.self_attn(
             hidden_states=hidden_states,
             positions=positions,
diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py
index da69e2631bd9..9e8f93330018 100644
--- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py
+++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py
@@ -18,6 +18,8 @@
 from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
 from cutlass.cutlass_dsl import T, dsl_user_op
 
+from vllm.platforms import current_platform
+
 _FAB_N = 144
 _FAB_K = 7168
 _FAB_VECTOR_WIDTH = 16
@@ -91,14 +93,61 @@ def _fma_f32_bf16(
     return Float32(result)
 
 
+@dsl_user_op
+def _fma_f32_bf16_portable(
+    a: BFloat16,
+    b: BFloat16,
+    acc: Float32,
+    *,
+    loc=None,
+    ip=None,
+) -> Float32:
+    """BF16 multiply with FP32 accumulation for pre-SM100 GPUs.
+
+    PTX ``fma.f32.bf16`` requires SM100 or newer. Converting the operands to
+    FP32 first keeps this kernel available on Hopper, where an unconditional
+    mixed-precision FMA fails libNVVM compilation for ``sm_90a``.
+    """
+    a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
+    b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
+    result = llvm.inline_asm(
+        T.f32(),
+        [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)],
+        "{\n\t"
+        ".reg .f32 a_f32, b_f32;\n\t"
+        "cvt.f32.bf16 a_f32, $1;\n\t"
+        "cvt.f32.bf16 b_f32, $2;\n\t"
+        "fma.rn.f32 $0, a_f32, b_f32, $3;\n\t"
+        "}",
+        "=f,h,h,f",
+        has_side_effects=False,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+    return Float32(result)
+
+
+def _has_mixed_precision_bf16_fma() -> bool:
+    """Whether PTX ``fma.f32.bf16`` is supported by the current GPU."""
+    return current_platform.has_device_capability(100)
+
+
 class _KdaSkinnyNGemm:
-    def __init__(self, num_rows: int, split_n: int) -> None:
+    def __init__(
+        self,
+        num_rows: int,
+        split_n: int,
+        use_mixed_precision_fma: bool,
+    ) -> None:
         if _FAB_N % split_n:
             raise ValueError("split_n must divide the FAB output size")
         if num_rows * split_n > _FAB_BLOCK_SIZE:
             raise ValueError("the FAB output tile must fit in one block")
         self.num_rows = num_rows
         self.split_n = split_n
+        self.use_mixed_precision_fma = use_mixed_precision_fma
 
     @cute.jit
     def __call__(
@@ -191,11 +240,18 @@ def kernel(
         for mi in cutlass.range_constexpr(num_rows):
             for ni in cutlass.range_constexpr(split_n):
                 for vi in cutlass.range_constexpr(_FAB_VECTOR_WIDTH):
-                    acc[mi, ni] = _fma_f32_bf16(
-                        a_regs[mi, vi],
-                        weight_regs[ni, vi],
-                        acc[mi, ni],
-                    )
+                    if cutlass.const_expr(self.use_mixed_precision_fma):
+                        acc[mi, ni] = _fma_f32_bf16(
+                            a_regs[mi, vi],
+                            weight_regs[ni, vi],
+                            acc[mi, ni],
+                        )
+                    else:
+                        acc[mi, ni] = _fma_f32_bf16_portable(
+                            a_regs[mi, vi],
+                            weight_regs[ni, vi],
+                            acc[mi, ni],
+                        )
                 acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
 
         partials_layout = cute.make_layout(
@@ -232,13 +288,19 @@ def kernel(
 
 
 class _KdaSkinnyKGemm:
-    def __init__(self, num_rows: int, split_n: int) -> None:
+    def __init__(
+        self,
+        num_rows: int,
+        split_n: int,
+        use_mixed_precision_fma: bool,
+    ) -> None:
         if _FB_N % split_n:
             raise ValueError("split_n must divide the FB output size")
         if split_n * _FB_K_LANES > 1024:
             raise ValueError("the FB output tile exceeds the block-size limit")
         self.num_rows = num_rows
         self.split_n = split_n
+        self.use_mixed_precision_fma = use_mixed_precision_fma
 
     @cute.jit
     def __call__(
@@ -314,11 +376,18 @@ def kernel(
         for mi in cutlass.range_constexpr(num_rows):
             acc = Float32(0.0)
             for vi in cutlass.range_constexpr(_FB_VECTOR_WIDTH):
-                acc = _fma_f32_bf16(
-                    input_regs[mi, vi],
-                    weight_regs[vi],
-                    acc,
-                )
+                if cutlass.const_expr(self.use_mixed_precision_fma):
+                    acc = _fma_f32_bf16(
+                        input_regs[mi, vi],
+                        weight_regs[vi],
+                        acc,
+                    )
+                else:
+                    acc = _fma_f32_bf16_portable(
+                        input_regs[mi, vi],
+                        weight_regs[vi],
+                        acc,
+                    )
             for offset in (8, 4, 2, 1):
                 acc += cute.arch.shuffle_sync_bfly(acc, offset)
             output[mi, n_idx] = acc.to(BFloat16)
@@ -341,7 +410,11 @@ def _compile_n(self, num_rows: int) -> None:
             return
         split_n = SKINNY_N_SPLIT_N[num_rows]
         self._compiled_n[key] = cute.compile(
-            _KdaSkinnyNGemm(num_rows, split_n),
+            _KdaSkinnyNGemm(
+                num_rows,
+                split_n,
+                use_mixed_precision_fma=_has_mixed_precision_bf16_fma(),
+            ),
             make_fake_tensor(
                 BFloat16,
                 (num_rows, _FAB_K),
@@ -371,7 +444,11 @@ def _compile_k(self, num_rows: int) -> None:
             return
         split_n = SKINNY_K_SPLIT_N[num_rows]
         self._compiled_k[key] = cute.compile(
-            _KdaSkinnyKGemm(num_rows, split_n),
+            _KdaSkinnyKGemm(
+                num_rows,
+                split_n,
+                use_mixed_precision_fma=_has_mixed_precision_bf16_fma(),
+            ),
             make_fake_tensor(
                 BFloat16,
                 (num_rows, _FB_K),
diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py
index b8eed0d7f4a0..67d5cfdf081d 100644
--- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py
+++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py
@@ -59,7 +59,6 @@ def _mapping(
 
 def validate_shape(*, tp_size: int, latent_dim: int, hidden_dim: int) -> None:
     """Validate constraints imposed by the fused collective mapping."""
-
     if tp_size <= 0 or latent_dim <= 0 or hidden_dim <= 0:
         raise ValueError("collective dimensions must be positive")
     if hidden_dim % tp_size:
@@ -86,7 +85,6 @@ def _select_routed_schedule(
     max_m: int,
 ) -> tuple[int, int]:
     """Match upstream MNNVL's one-cluster-per-token occupancy policy."""
-
     _, cluster_ctas, threads, _ = _mapping(tp_size, latent_dim, hidden_dim)
     sm_count = torch.cuda.get_device_properties(
         torch.accelerator.current_device_index()
@@ -886,7 +884,6 @@ def compile_kernel(
     top_k: int = 0,
 ) -> None:
     """Compile the rank/M specialization without retaining caller tensors."""
-
     key = _compile_key(
         rank,
         tp_size,
diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py
index 6a1ba9d77b44..ea9d86d6fa85 100644
--- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py
+++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py
@@ -79,7 +79,6 @@ def validate_configuration(
     b_prime_stages: int,
 ) -> None:
     """Validate constraints imposed by this BF16 SM100 GEMM."""
-
     if latent_dim <= 0 or shard_dim <= 0:
         raise ValueError("GEMM K and N must be positive")
     if latent_dim % 8 or shard_dim % 8:
diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py
index 82782fcb96d2..3840caf24a27 100644
--- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py
+++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py
@@ -52,7 +52,6 @@ def to_cute_dynamic_m(
     mode is symbolic, so changing M within an Op's capacity reuses the same
     compiled kernel.
     """
-
     return to_cute(tensor, assumed_align).mark_compact_shape_dynamic(
         mode=mode,
         stride_order=tensor.dim_order(),
@@ -98,7 +97,6 @@ def load_global_u32x4(
     side-effecting prevents loop-invariant motion and common-subexpression
     elimination across polling iterations.
     """
-
     address = pointer.toint(loc=loc, ip=ip)
     opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32"
     out = llvm.inline_asm(
@@ -194,7 +192,6 @@ def finalize_top16_bf16(
     ip=None,
 ):
     """Finalize one Kimi K3 top-16 vector with wide metadata loads."""
-
     addresses = [
         pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
         for pointer in (gemm2_vector, route_indices, route_weights)
@@ -229,7 +226,6 @@ def store_global_u32x4(
     ip=None,
 ) -> None:
     """Store four packed words to an ordinary or NVLS multicast global VA."""
-
     words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)]
     opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32"
     llvm.inline_asm(
@@ -294,7 +290,6 @@ def store_global_u32(
 @dsl_user_op
 def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None:
     """Publish eight BF16 values through an NVLS multicast mapping."""
-
     words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)]
     llvm.inline_asm(
         None,
@@ -312,7 +307,6 @@ def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None:
 @dsl_user_op
 def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None:
     """Reset one Lamport fragment to four FP32 negative-zero bit patterns."""
-
     address = pointer.toint(loc=loc, ip=ip)
     value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip)
     llvm.inline_asm(
@@ -333,7 +327,6 @@ def red_async_release_gpu_add_u32(
     pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None
 ) -> None:
     """The exact SM100 arrival primitive used by upstream LamportFlags."""
-
     address = pointer.toint(loc=loc, ip=ip)
     llvm.inline_asm(
         None,
@@ -378,7 +371,6 @@ def map_shared_to_peer(
     ip=None,
 ) -> Int32:
     """Map a local shared-memory slot to the same slot in a peer CTA."""
-
     smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value()
     return Int32(
         llvm.inline_asm(
@@ -422,7 +414,6 @@ def store_shared_cluster_f32(
 @dsl_user_op
 def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None):
     """Load two aligned FP32 DSM partials from local shared memory."""
-
     address = Int32(pointer.toint(loc=loc, ip=ip))
     out = llvm.inline_asm(
         llvm.StructType.get_literal([T.f32()] * 2),
@@ -443,7 +434,6 @@ def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None):
 @dsl_user_op
 def load_shared_f32x4(pointer: cute.Pointer, *, loc=None, ip=None):
     """Load four aligned FP32 DSM partials from local shared memory."""
-
     address = Int32(pointer.toint(loc=loc, ip=ip))
     out = llvm.inline_asm(
         llvm.StructType.get_literal([T.f32()] * 4),
@@ -530,7 +520,6 @@ def sanitize_negative_zero_u32x2(packed):
 @cute.jit
 def sanitize_negative_zero(packed):
     """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel."""
-
     result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
     for i in cutlass.range_constexpr(4):
         word = packed[i]
@@ -547,7 +536,6 @@ def sanitize_negative_zero(packed):
 @cute.jit
 def fragment_is_dirty(packed):
     """Bit-exact upstream sentinel check: one comparison per 32-bit word."""
-
     dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS)
     for i in cutlass.range_constexpr(1, 4):
         dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS))
@@ -564,7 +552,6 @@ def warp_sum_specialized(
     last_warp_mask: cutlass.Constexpr[int],
 ) -> Float32:
     """Warp sum supporting a compile-time partial final warp."""
-
     if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32):
         for offset in cutlass.range_constexpr(16, 0, -1):
             # range_constexpr does not provide powers-of-two stepping.
@@ -587,34 +574,3 @@ def warp_sum_specialized(
                     mask_and_clamp=31,
                 )
     return value
-
-
-@cute.jit
-def block_sum_specialized(
-    value: Float32,
-    warp_sums: cute.Tensor,
-    tidx: Int32,
-    warps: cutlass.Constexpr[int],
-    last_warp_lanes: cutlass.Constexpr[int],
-    last_warp_mask: cutlass.Constexpr[int],
-) -> Float32:
-    """Upstream-equivalent FP32 block reduction."""
-
-    lane = cute.arch.lane_idx()
-    warp_idx = cute.arch.warp_idx()
-    value = warp_sum_specialized(
-        value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask
-    )
-    if lane == 0:
-        warp_sums[warp_idx] = value
-    cute.arch.barrier()
-
-    block_sum = Float32(0.0)
-    if warp_idx == 0:
-        if lane < Int32(warps):
-            block_sum = warp_sums[lane]
-        block_sum = cute.arch.warp_reduction_sum(block_sum)
-        if lane == 0:
-            warp_sums[0] = block_sum
-    cute.arch.barrier()
-    return warp_sums[0]
diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py
index 24ae1a1a820c..79de3af0bc22 100644
--- a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py
+++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py
@@ -76,9 +76,9 @@ def fused_mla_key_concat_ds_mla_insert(
     """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout.
 
     The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of
-    128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no
-    scale argument. Returns the bf16 full key; optionally rotates ``q`` and
-    writes ``kv_cache`` in place.
+    128 with per-tile power-of-two scales stored as float32, RoPE as bf16) --
+    self-scaling, so no scale argument. Returns the bf16 full key; optionally
+    rotates ``q`` and writes ``kv_cache`` in place.
     """
     k_pe = k_pe.reshape(k_pe.shape[0], -1)
     tp, num_heads, qk_nope_head_dim = k_nope.shape
diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py
index 416c316a70e7..3911d20d0358 100644
--- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py
+++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py
@@ -539,9 +539,7 @@ def fused_kda_gate_chunk_cumsum(
         )
     B, T, H, D = raw_g.shape
     if raw_beta.shape != (B, T, H):
-        raise ValueError(
-            f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}"
-        )
+        raise ValueError(f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}")
     if chunk_indices is None and cu_seqlens is not None:
         chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
     NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices)
@@ -901,13 +899,12 @@ def fused_kda_gate(
     threshold: float = 20.0,
     lower_bound: float | None = None,
 ) -> torch.Tensor:
-    """
-    Forward pass for KDA gate:
-      input g: [..., H*D]
-      param A: [H] or [1, 1, H, 1]
-      beta: softplus beta parameter
-      threshold: softplus threshold parameter
-      return  : [..., H, D]
+    """Forward pass for KDA gate:
+    input g: [..., H*D]
+    param A: [H] or [1, 1, H, 1]
+    beta: softplus beta parameter
+    threshold: softplus threshold parameter
+    return  : [..., H, D]
     """
     orig_shape = g.shape[:-1]
 
diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py
index 087197d10a56..8b9aa86f9ed9 100644
--- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py
+++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py
@@ -23,18 +23,20 @@
 ################################################################################
 
 
-@triton.heuristics({
-    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
-})
+@triton.heuristics(
+    {
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
 @triton.autotune(
     configs=[
-        triton.Config({'BK': BK}, num_warps=num_warps)
+        triton.Config({"BK": BK}, num_warps=num_warps)
         for BK in [32, 64]
         for num_warps in [1, 2, 4]
     ],
     key=["H", "HV", "K", "BC"],
 )
-@triton.jit(do_not_specialize=['T'])
+@triton.jit(do_not_specialize=["T"])
 def chunk_kda_fwd_kernel_inter_solve_fused(
     q,
     k,
@@ -57,8 +59,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
     USE_SAFE_GATE: tl.constexpr,
     SOLVE_TRIL_DOT_PRECISION: tl.constexpr,
 ):
-    """
-    Fused kernel: compute inter-subchunk Akk + solve_tril in one pass.
+    """Fused kernel: compute inter-subchunk Akk + solve_tril in one pass.
     Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd.
 
     This kernel:
@@ -74,8 +75,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
     i_h = i_hv // (HV // H)
 
     if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
     else:
         bos, eos = i_b * T, i_b * T + T
@@ -122,21 +129,31 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
         o_k = i_k * BK + tl.arange(0, BK)
         m_k = o_k < K
 
-        p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0))
-        p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0))
+        p_k0 = tl.make_block_ptr(
+            k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
+        )
+        p_g0 = tl.make_block_ptr(
+            g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
+        )
         b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32)
         b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32)
 
         if i_tc1 < T:
-            p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
-            p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
-            p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
+            p_q1 = tl.make_block_ptr(
+                q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
+            )
+            p_k1 = tl.make_block_ptr(
+                k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
+            )
+            p_g1 = tl.make_block_ptr(
+                g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
+            )
             # [BC, BK]
             b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32)
             b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32)
             b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32)
             # [BK]
-            b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
+            b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32)
             # [BC, BK]
             b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0)
             # [BK, BC]
@@ -146,15 +163,23 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
             b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt)
 
             if i_tc2 < T:
-                p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
-                p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
-                p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
+                p_q2 = tl.make_block_ptr(
+                    q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
+                )
+                p_k2 = tl.make_block_ptr(
+                    k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
+                )
+                p_g2 = tl.make_block_ptr(
+                    g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
+                )
                 # [BC, BK]
                 b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32)
                 b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32)
                 b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32)
                 # [BK]
-                b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
+                b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to(
+                    tl.float32
+                )
                 # [BC, BK]
                 b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0)
                 b_qg2 = b_q2 * b_gqn2
@@ -170,15 +195,23 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
                 b_Akk21 += tl.dot(b_kg2, b_kgt)
 
                 if i_tc3 < T:
-                    p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
-                    p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
-                    p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
+                    p_q3 = tl.make_block_ptr(
+                        q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
+                    )
+                    p_k3 = tl.make_block_ptr(
+                        k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
+                    )
+                    p_g3 = tl.make_block_ptr(
+                        g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
+                    )
                     # [BC, BK]
                     b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32)
                     b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32)
                     b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32)
                     # [BK]
-                    b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
+                    b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to(
+                        tl.float32
+                    )
                     # [BC, BK]
                     b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0)
                     b_qg3 = b_q3 * b_gqn3
@@ -203,40 +236,78 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
     # save off-diagonal Aqk blocks and prepare Akk
     ################################################################################
     if i_tc1 < T:
-        p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
-        tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
+        p_Aqk10 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)
+        )
+        tl.store(
+            p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
 
-        p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,))
+        p_b1 = tl.make_block_ptr(
+            beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)
+        )
         b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32)
         b_Akk10 = b_Akk10 * b_b1[:, None]
     if i_tc2 < T:
-        p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0))
-        p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0))
-        tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
-        tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
+        p_Aqk20 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)
+        )
+        p_Aqk21 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)
+        )
+        tl.store(
+            p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+        tl.store(
+            p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
 
-        p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,))
+        p_b2 = tl.make_block_ptr(
+            beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)
+        )
         b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32)
         b_Akk20 = b_Akk20 * b_b2[:, None]
         b_Akk21 = b_Akk21 * b_b2[:, None]
     if i_tc3 < T:
-        p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0))
-        p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0))
-        p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0))
-        tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
-        tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
-        tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
-
-        p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,))
+        p_Aqk30 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)
+        )
+        p_Aqk31 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)
+        )
+        p_Aqk32 = tl.make_block_ptr(
+            Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
+        )
+        tl.store(
+            p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+        tl.store(
+            p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+        tl.store(
+            p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
+        )
+
+        p_b3 = tl.make_block_ptr(
+            beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)
+        )
         b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32)
         b_Akk30 = b_Akk30 * b_b3[:, None]
         b_Akk31 = b_Akk31 * b_b3[:, None]
         b_Akk32 = b_Akk32 * b_b3[:, None]
 
-    p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0))
-    p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0))
-    p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0))
-    p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0))
+    p_Akk00 = tl.make_block_ptr(
+        Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)
+    )
+    p_Akk11 = tl.make_block_ptr(
+        Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)
+    )
+    p_Akk22 = tl.make_block_ptr(
+        Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)
+    )
+    p_Akk33 = tl.make_block_ptr(
+        Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)
+    )
     b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32)
     b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32)
     b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32)
@@ -256,25 +327,25 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
         b_Ai33 = -tl.where(m_A, b_Ai33, 0)
 
         for i in range(2, min(BC, T - i_tc0)):
-            b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
-            b_a00 = tl.where(o_i < i, b_a00, 0.)
+            b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
+            b_a00 = tl.where(o_i < i, b_a00, 0.0)
             b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0)
             b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
-        for i in range(BC + 2, min(2*BC, T - i_tc0)):
-            b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
-            b_a11 = tl.where(o_i < i - BC, b_a11, 0.)
+        for i in range(BC + 2, min(2 * BC, T - i_tc0)):
+            b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
+            b_a11 = tl.where(o_i < i - BC, b_a11, 0.0)
             b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0)
             b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11)
-        for i in range(2*BC + 2, min(3*BC, T - i_tc0)):
-            b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
-            b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.)
+        for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)):
+            b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
+            b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.0)
             b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0)
-            b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22)
-        for i in range(3*BC + 2, min(4*BC, T - i_tc0)):
-            b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
-            b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.)
+            b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22)
+        for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)):
+            b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
+            b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.0)
             b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0)
-            b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33)
+            b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33)
 
         b_Ai00 += m_I
         b_Ai11 += m_I
@@ -289,53 +360,73 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
     b_Ai10 = -tl.dot(
         tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION),
         b_Ai00,
-        input_precision=SOLVE_TRIL_DOT_PRECISION
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
     )
     b_Ai21 = -tl.dot(
         tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION),
         b_Ai11,
-        input_precision=SOLVE_TRIL_DOT_PRECISION
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
     )
     b_Ai32 = -tl.dot(
         tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION),
         b_Ai22,
-        input_precision=SOLVE_TRIL_DOT_PRECISION
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
     )
 
     b_Ai20 = -tl.dot(
         b_Ai22,
-        tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
-        tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
-        input_precision=SOLVE_TRIL_DOT_PRECISION
+        tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION)
+        + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
     )
     b_Ai31 = -tl.dot(
         b_Ai33,
-        tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) +
-        tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
-        input_precision=SOLVE_TRIL_DOT_PRECISION
+        tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION)
+        + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
     )
     b_Ai30 = -tl.dot(
         b_Ai33,
-        tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
-        tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) +
-        tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
-        input_precision=SOLVE_TRIL_DOT_PRECISION
+        tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION)
+        + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION)
+        + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
     )
 
     ################################################################################
     # store full Akk_inv to Akk
     ################################################################################
 
-    p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0))
-    p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
-    p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0))
-    p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0))
-    p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0))
-    p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0))
-    p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0))
-    p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0))
-    p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0))
-    p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0))
+    p_Akk00 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)
+    )
+    p_Akk10 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)
+    )
+    p_Akk11 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)
+    )
+    p_Akk20 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)
+    )
+    p_Akk21 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)
+    )
+    p_Akk22 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)
+    )
+    p_Akk30 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)
+    )
+    p_Akk31 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)
+    )
+    p_Akk32 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
+    )
+    p_Akk33 = tl.make_block_ptr(
+        Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)
+    )
 
     tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1))
     tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1))
@@ -349,18 +440,20 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
     tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1))
 
 
-@triton.heuristics({
-    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
-})
+@triton.heuristics(
+    {
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
 @triton.autotune(
     configs=[
         triton.Config({}, num_warps=num_warps, num_stages=num_stages)
         for num_warps in [1, 2, 4, 8]
         for num_stages in [2, 3, 4]
     ],
-    key=['BK', 'NC', 'BT', 'HV'],
+    key=["BK", "NC", "BT", "HV"],
 )
-@triton.jit(do_not_specialize=['B', 'T'])
+@triton.jit(do_not_specialize=["B", "T"])
 def chunk_kda_fwd_kernel_intra_sub_chunk(
     q,
     k,
@@ -386,8 +479,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
     i_h = i_hv // (HV // H)
 
     if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
     else:
         bos, eos = i_b * T, i_b * T + T
@@ -406,9 +505,9 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
     Aqk = Aqk + (bos * HV + i_hv) * BT
     Akk = Akk + (bos * HV + i_hv) * BC
 
-    p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0))
-    p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0))
-    p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, 0), (BC, BK), (1, 0))
+    p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
+    p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
+    p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0))
 
     p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,))
 
@@ -418,10 +517,12 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
     b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
 
     if USE_GATHER:
-        b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0)
+        b_gn = gather(
+            b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0
+        )
     else:
         # caculate offset
-        p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + tl.arange(0, BK)
+        p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + tl.arange(0, BK)
         b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0)
         b_gn = b_gn[None, :]
 
@@ -429,8 +530,8 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
     # less than 85 to avoid overflow in exp2
     b_gm = (b_g - b_gn).to(tl.float32)
 
-    b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.)
-    b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.)
+    b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.0)
+    b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.0)
 
     b_kgt = tl.trans(b_k * b_gk)
 
@@ -445,8 +546,10 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
     b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0)
     b_Akk = tl.where(m_Akk, b_Akk, 0.0)
 
-    p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0))
-    p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV*BC, 1), (i_ti, 0), (BC, BC), (1, 0))
+    p_Aqk = tl.make_block_ptr(
+        Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)
+    )
+    p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0))
     tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))
     tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1))
 
@@ -458,8 +561,8 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
 
     b_Ai = -b_Akk
     for i in range(2, min(BC, T - i_ti)):
-        b_a = -tl.load(Akk + (i_ti + i) * HV*BC + o_i)
-        b_a = tl.where(o_i < i, b_a, 0.)
+        b_a = -tl.load(Akk + (i_ti + i) * HV * BC + o_i)
+        b_a = tl.where(o_i < i, b_a, 0.0)
         b_a += tl.sum(b_a[:, None] * b_Ai, 0)
         b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai)
     b_Ai += m_I
@@ -531,8 +634,7 @@ def chunk_kda_fwd_intra(
     # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen)
     solve_tril_dot_precision = (
         "tf32"
-        if current_platform.is_cuda()
-        and current_platform.has_device_capability(80)
+        if current_platform.is_cuda() and current_platform.has_device_capability(80)
         else "ieee"
     )
     grid = (NT, B * HV)
diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py
index ecd00a51f9dd..86654f264431 100644
--- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py
+++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py
@@ -16,18 +16,20 @@
 from vllm.triton_utils import tl, triton
 
 
-@triton.heuristics({
-    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
-})
+@triton.heuristics(
+    {
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
 @triton.autotune(
     configs=[
-        triton.Config({'BH': BH}, num_warps=num_warps)
+        triton.Config({"BH": BH}, num_warps=num_warps)
         for BH in [1, 2, 4, 8]
         for num_warps in [1, 2, 4, 8]
     ],
     key=["K", "H", "HV"],
 )
-@triton.jit(do_not_specialize=['T', 'N'])
+@triton.jit(do_not_specialize=["T", "N"])
 def chunk_kda_fwd_kernel_intra_token_parallel(
     q,
     k,
@@ -65,7 +67,10 @@ def chunk_kda_fwd_kernel_intra_token_parallel(
                     left = mid + 1
         i_n = left
 
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
         i_t = i_tg - bos
     else:
@@ -82,11 +87,11 @@ def chunk_kda_fwd_kernel_intra_token_parallel(
 
     G: tl.constexpr = HV // H
 
-    q += bos * H*K
-    k += bos * H*K
-    g += bos * HV*K
-    Aqk += bos * HV*BT
-    Akk += bos * HV*BC
+    q += bos * H * K
+    k += bos * H * K
+    g += bos * HV * K
+    Aqk += bos * HV * BT
+    Akk += bos * HV * BC
     beta += bos * HV
 
     BK: tl.constexpr = triton.next_power_of_2(K)
@@ -103,7 +108,9 @@ def chunk_kda_fwd_kernel_intra_token_parallel(
     b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
 
     # g: [B, T, HV, K], beta: [B, T, HV]
-    p_g = tl.make_block_ptr(g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0))
+    p_g = tl.make_block_ptr(
+        g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)
+    )
     p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,))
     b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
     b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
@@ -111,15 +118,25 @@ def chunk_kda_fwd_kernel_intra_token_parallel(
 
     for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))):
         b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
-        p_gj = tl.make_block_ptr(g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0))
+        p_gj = tl.make_block_ptr(
+            g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)
+        )
         b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32)
 
         b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0)
         b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale
         b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0)
 
-        tl.store(Aqk + i_t * HV * BT + o_hv * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_hv)
-        tl.store(Akk + i_t * HV * BC + o_hv * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_hv)
+        tl.store(
+            Aqk + i_t * HV * BT + o_hv * BT + j % BT,
+            b_Aqk.to(Aqk.dtype.element_ty),
+            mask=m_hv,
+        )
+        tl.store(
+            Akk + i_t * HV * BC + o_hv * BC + j - i_ts,
+            b_Akk.to(Akk.dtype.element_ty),
+            mask=m_hv,
+        )
 
 
 def chunk_kda_fwd_intra_token_parallel(
@@ -134,8 +151,7 @@ def chunk_kda_fwd_intra_token_parallel(
     chunk_size: int = 64,
     sub_chunk_size: int = 16,
 ) -> None:
-    """
-    Token-parallel implementation: each token gets its own thread block.
+    """Token-parallel implementation: each token gets its own thread block.
     Supports both fixed-length and variable-length sequences.
     Reduces wasted computation on padding.
 
@@ -147,17 +163,21 @@ def chunk_kda_fwd_intra_token_parallel(
         gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA)
         beta: [B, T, HV]
         Aqk: [B, T, HV, BT] output tensor to write to
+        cu_seqlens: cumulative sequence lengths for variable-length input
         Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32)
         scale: attention scale
         chunk_size: BT (default 64)
         sub_chunk_size: BC (default 16)
+
     """
     B, T, H, K, HV = *q.shape, gk.shape[2]
     N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
     BT = chunk_size
     BC = sub_chunk_size
 
-    def grid(meta): return (B * T, triton.cdiv(HV, meta['BH']))
+    def grid(meta):
+        return (B * T, triton.cdiv(HV, meta["BH"]))
+
     chunk_kda_fwd_kernel_intra_token_parallel[grid](
         q=q,
         k=k,
diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py
index 7f5f98400f2d..e65c95670d62 100644
--- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py
+++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py
@@ -483,9 +483,7 @@ def fused_recurrent_kda(
     )
 
 
-@triton.jit(
-    do_not_specialize=["stride_beta_token", "stride_state_indices"]
-)
+@triton.jit(do_not_specialize=["stride_beta_token", "stride_state_indices"])
 def fused_recurrent_kda_packed_decode_kernel(
     mixed_qkv,
     raw_g,
diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py
index 50cdb5518381..4871e66a9ab3 100644
--- a/vllm/models/minimax_m3/amd/model.py
+++ b/vllm/models/minimax_m3/amd/model.py
@@ -21,7 +21,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm import _custom_ops as ops
 from vllm import envs
@@ -134,7 +134,7 @@
 logger = init_logger(__name__)
 
 
-def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:
+def _sparse_attention_layer_ids(config: PreTrainedConfig) -> set[int]:
     """Layer ids whose attention runs the extra sparse "index" branch."""
     cfg = getattr(config, "sparse_attention_config", None)
     if not cfg:
@@ -145,7 +145,7 @@ def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:
     return {i for i, f in enumerate(freq) if f != 0}
 
 
-def _sparse_attention_layer_ordinals(config: PretrainedConfig) -> dict[int, int]:
+def _sparse_attention_layer_ordinals(config: PreTrainedConfig) -> dict[int, int]:
     """Map each sparse-attention layer id to its ordinal among sparse layers."""
     return {
         lid: ordinal
@@ -153,7 +153,7 @@ def _sparse_attention_layer_ordinals(config: PretrainedConfig) -> dict[int, int]
     }
 
 
-def _should_skip_index_topk(config: PretrainedConfig, layer_id: int) -> bool:
+def _should_skip_index_topk(config: PreTrainedConfig, layer_id: int) -> bool:
     """ATOM ``index_topk_freq`` (cross-layer index sharing).
 
     Only 1 of every ``index_topk_freq`` sparse-attention layers recomputes the
@@ -176,7 +176,7 @@ def _should_skip_index_topk(config: PretrainedConfig, layer_id: int) -> bool:
     return max(ordinal - offset, 0) % freq != 0
 
 
-def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool:
+def _is_moe_layer(config: PreTrainedConfig, layer_id: int) -> bool:
     """Whether this layer's MLP is a sparse MoE block (vs a dense MLP)."""
     moe_layer_freq = getattr(config, "moe_layer_freq", None)
     if moe_layer_freq is None:
@@ -184,7 +184,7 @@ def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool:
     return moe_layer_freq[layer_id] != 0
 
 
-def _build_rotary_emb(config: PretrainedConfig, head_dim: int):
+def _build_rotary_emb(config: PreTrainedConfig, head_dim: int):
     """Build the (partial NeoX) RoPE, honoring an optional ``rope_scaling`` config.
 
     Without scaling the cos/sin cache is sized to ``max_position_embeddings``
@@ -262,7 +262,7 @@ class MiniMaxM3MLP(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         intermediate_size: int,
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
@@ -338,7 +338,7 @@ class MiniMaxM3MoE(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_id: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -488,7 +488,7 @@ class MiniMaxM3Attention(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_id: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -579,6 +579,32 @@ def forward(
         return output
 
 
+# Widest row pitch aiter's reshape_and_cache can carry in its `int` stride args.
+_MAX_KV_INSERT_ROW_STRIDE = 2**31 - 1
+
+
+def _kv_insert_operand(t: torch.Tensor) -> torch.Tensor:
+    """K or V operand for ``aiter.reshape_and_cache``, without a needless copy.
+
+    The AITER sparse-PA insert takes K and V as column slices of the fused
+    ``[q | k | v | index_q | index_k]`` projection, so they are row-strided
+    views and never contiguous. ``reshape_and_cache`` does not need them to be:
+    it passes ``key.stride(0)`` to the kernel as the row pitch and reads
+    ``key[token_idx * key_stride + i]`` for ``i`` over ``num_heads *
+    head_size``, so the only requirement is that the trailing
+    ``(num_heads, head_size)`` block be contiguous within a row, which such a
+    slice always satisfies. Any other layout falls back to a copy.
+    """
+    if (
+        t.dim() == 3
+        and t.stride(2) == 1
+        and t.stride(1) == t.shape[2]
+        and t.stride(0) <= _MAX_KV_INSERT_ROW_STRIDE
+    ):
+        return t
+    return t.contiguous()
+
+
 class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
     """Block-sparse attention layer with the lightning-indexer branch.
 
@@ -597,7 +623,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_id: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -864,6 +890,30 @@ def get_aiter_sparse_pa_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]:
         self._ensure_aiter_sparse_pa_kv_cache()
         return self.kv_cache_k, self.kv_cache_v
 
+    def _get_aiter_sparse_pa_slot_mapping(
+        self,
+        slot_mapping: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+    ) -> torch.Tensor:
+        if value_cache.shape[0] == key_cache.shape[0]:
+            return slot_mapping
+
+        attn_metadata = get_forward_context().attn_metadata
+        page16_slot_mapping = None
+        if isinstance(attn_metadata, dict):
+            main_md = attn_metadata[self.layer_name]
+            assert isinstance(main_md, MiniMaxM3SparseMetadata)
+            page16_slot_mapping = main_md.page16_slot_mapping
+        if (
+            page16_slot_mapping is None
+            or page16_slot_mapping.shape != slot_mapping.shape
+        ):
+            page16_slot_mapping = minimax_m3_rebase_slots_to_page16(
+                slot_mapping, self.kv_cache.shape[2]
+            )
+        return page16_slot_mapping
+
     def _insert_aiter_sparse_pa_kv(
         self,
         k: torch.Tensor,
@@ -871,6 +921,8 @@ def _insert_aiter_sparse_pa_kv(
         index_k: torch.Tensor | None,
         slot_mapping: torch.Tensor,
         index_slot_mapping: torch.Tensor | None,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
     ) -> None:
         if self.kv_cache.numel() == 0:
             return
@@ -880,30 +932,14 @@ def _insert_aiter_sparse_pa_kv(
             minimax_m3_insert_index_cache,
         )
 
-        key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache()
-        if value_cache.shape[0] != key_cache.shape[0]:
-            attn_metadata = get_forward_context().attn_metadata
-            page16_slot_mapping = None
-            if isinstance(attn_metadata, dict):
-                main_md = attn_metadata[self.layer_name]
-                assert isinstance(main_md, MiniMaxM3SparseMetadata)
-                page16_slot_mapping = main_md.page16_slot_mapping
-            if (
-                page16_slot_mapping is None
-                or page16_slot_mapping.shape != slot_mapping.shape
-            ):
-                page16_slot_mapping = minimax_m3_rebase_slots_to_page16(
-                    slot_mapping, self.kv_cache.shape[2]
-                )
-            slot_mapping = page16_slot_mapping
         kv_cache_dtype = (
             self.kv_cache_dtype
             if is_quantized_kv_cache(self.kv_cache_dtype)
             else "auto"
         )
         reshape_and_cache(
-            k.contiguous(),
-            v.contiguous(),
+            _kv_insert_operand(k),
+            _kv_insert_operand(v),
             key_cache,
             value_cache,
             slot_mapping,
@@ -950,38 +986,78 @@ def forward(
 
         main_slot_mapping = fwd_slot_mapping[self.layer_name]
         q = qkv.new_empty((num_tokens, self.q_size))
+        key_cache = value_cache = page16_slot_mapping = None
+        k_scale = v_scale = None
+        use_fused_qknorm = False
+        if self.use_aiter_sparse_pa:
+            key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache()
+            page16_slot_mapping = self._get_aiter_sparse_pa_slot_mapping(
+                main_slot_mapping, key_cache, value_cache
+            )
+            k_scale = getattr(self, "_k_scale", None)
+            v_scale = getattr(self, "_v_scale", None)
+            use_fused_qknorm = rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled(
+                self.kv_cache_dtype, k_scale, v_scale
+            )
         if self.skip_index_topk:
             index_q = None
             if self.use_aiter_sparse_pa:
-                ops.fused_minimax_m3_qknorm_rope_kv_insert(
-                    qkv,
-                    self.q_norm.weight,
-                    self.k_norm.weight,
-                    cos_sin_cache,
-                    positions,
-                    self.num_heads,
-                    self.num_kv_heads,
-                    rotary_dim,
-                    eps,
-                    num_index_heads=self.num_idx_heads,
-                    q_out=q,
-                    skip_index_branch=True,
-                )
-                k_start = self.q_size
-                v_start = k_start + self.kv_size
-                k = qkv[:, k_start:v_start].view(
-                    num_tokens, self.num_kv_heads, self.head_dim
-                )
-                v = qkv[:, v_start : v_start + self.kv_size].view(
-                    num_tokens, self.num_kv_heads, self.head_dim
-                )
-                self._insert_aiter_sparse_pa_kv(
-                    k,
-                    v,
-                    None,
-                    main_slot_mapping,
-                    None,
-                )
+                assert key_cache is not None
+                assert value_cache is not None
+                assert page16_slot_mapping is not None
+                if use_fused_qknorm:
+                    rocm_aiter_ops.fused_qknorm_idxrqknorm(
+                        qkv,
+                        self.q_norm.weight,
+                        self.k_norm.weight,
+                        cos_sin_cache,
+                        positions,
+                        self.num_heads,
+                        self.num_kv_heads,
+                        rotary_dim,
+                        eps,
+                        page16_slot_mapping,
+                        key_cache,
+                        value_cache,
+                        q,
+                        self.kv_cache_dtype,
+                        k_scale,
+                        v_scale,
+                        num_index_heads=self.num_idx_heads,
+                        skip_index_branch=True,
+                    )
+                else:
+                    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+                        qkv,
+                        self.q_norm.weight,
+                        self.k_norm.weight,
+                        cos_sin_cache,
+                        positions,
+                        self.num_heads,
+                        self.num_kv_heads,
+                        rotary_dim,
+                        eps,
+                        num_index_heads=self.num_idx_heads,
+                        q_out=q,
+                        skip_index_branch=True,
+                    )
+                    k_start = self.q_size
+                    v_start = k_start + self.kv_size
+                    k = qkv[:, k_start:v_start].view(
+                        num_tokens, self.num_kv_heads, self.head_dim
+                    )
+                    v = qkv[:, v_start : v_start + self.kv_size].view(
+                        num_tokens, self.num_kv_heads, self.head_dim
+                    )
+                    self._insert_aiter_sparse_pa_kv(
+                        k,
+                        v,
+                        None,
+                        page16_slot_mapping,
+                        None,
+                        key_cache,
+                        value_cache,
+                    )
             else:
                 ops.fused_minimax_m3_qknorm_rope_kv_insert(
                     qkv,
@@ -1011,42 +1087,74 @@ def forward(
                 dtype=self.indexer.index_cache.dtype,
             )
             if self.use_aiter_sparse_pa:
-                ops.fused_minimax_m3_qknorm_rope_kv_insert(
-                    qkv,
-                    self.q_norm.weight,
-                    self.k_norm.weight,
-                    cos_sin_cache,
-                    positions,
-                    self.num_heads,
-                    self.num_kv_heads,
-                    rotary_dim,
-                    eps,
-                    self.index_q_norm.weight,
-                    self.index_k_norm.weight,
-                    self.num_idx_heads,
-                    q_out=q,
-                    index_q_out=index_q,
-                    kv_cache_dtype=self.kv_cache_dtype,
-                )
-                k_start = self.q_size
-                v_start = k_start + self.kv_size
-                index_k_start = v_start + self.kv_size + self.index_q_size
-                k = qkv[:, k_start:v_start].view(
-                    num_tokens, self.num_kv_heads, self.head_dim
-                )
-                v = qkv[:, v_start : v_start + self.kv_size].view(
-                    num_tokens, self.num_kv_heads, self.head_dim
-                )
-                index_k = qkv[
-                    :, index_k_start : index_k_start + self.idx_head_dim
-                ].view(num_tokens, self.idx_head_dim)
-                self._insert_aiter_sparse_pa_kv(
-                    k,
-                    v,
-                    index_k,
-                    main_slot_mapping,
-                    index_slot_mapping,
-                )
+                assert key_cache is not None
+                assert value_cache is not None
+                assert page16_slot_mapping is not None
+                index_cache = self.indexer.index_cache.kv_cache
+                if use_fused_qknorm:
+                    rocm_aiter_ops.fused_qknorm_idxrqknorm(
+                        qkv,
+                        self.q_norm.weight,
+                        self.k_norm.weight,
+                        cos_sin_cache,
+                        positions,
+                        self.num_heads,
+                        self.num_kv_heads,
+                        rotary_dim,
+                        eps,
+                        page16_slot_mapping,
+                        key_cache,
+                        value_cache,
+                        q,
+                        self.kv_cache_dtype,
+                        k_scale,
+                        v_scale,
+                        self.index_q_norm.weight,
+                        self.index_k_norm.weight,
+                        self.num_idx_heads,
+                        index_cache,
+                        index_q,
+                        index_slot_mapping,
+                    )
+                else:
+                    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+                        qkv,
+                        self.q_norm.weight,
+                        self.k_norm.weight,
+                        cos_sin_cache,
+                        positions,
+                        self.num_heads,
+                        self.num_kv_heads,
+                        rotary_dim,
+                        eps,
+                        self.index_q_norm.weight,
+                        self.index_k_norm.weight,
+                        self.num_idx_heads,
+                        q_out=q,
+                        index_q_out=index_q,
+                        kv_cache_dtype=self.kv_cache_dtype,
+                    )
+                    k_start = self.q_size
+                    v_start = k_start + self.kv_size
+                    index_k_start = v_start + self.kv_size + self.index_q_size
+                    k = qkv[:, k_start:v_start].view(
+                        num_tokens, self.num_kv_heads, self.head_dim
+                    )
+                    v = qkv[:, v_start : v_start + self.kv_size].view(
+                        num_tokens, self.num_kv_heads, self.head_dim
+                    )
+                    index_k = qkv[
+                        :, index_k_start : index_k_start + self.idx_head_dim
+                    ].view(num_tokens, self.idx_head_dim)
+                    self._insert_aiter_sparse_pa_kv(
+                        k,
+                        v,
+                        index_k,
+                        page16_slot_mapping,
+                        index_slot_mapping,
+                        key_cache,
+                        value_cache,
+                    )
             else:
                 ops.fused_minimax_m3_qknorm_rope_kv_insert(
                     qkv,
@@ -1156,7 +1264,7 @@ def _run_attention(
 class MiniMaxM3DecoderLayer(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         prefix: str,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
@@ -1630,7 +1738,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         with self._mark_tower_model(vllm_config, {"image", "video"}):
             vision_config = config.vision_config
             self.vision_tower = MiniMaxVLVisionModel(
-                config=PretrainedConfig.from_dict(vision_config),
+                config=PreTrainedConfig.from_dict(vision_config),
                 text_hidden_size=text_hidden_size,
                 projector_hidden_size=projector_hidden_size,
                 quant_config=self.quant_config,
diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py
index 54a166f066db..d2b4602c56b8 100644
--- a/vllm/models/minimax_m3/amd/ops/index_topk.py
+++ b/vllm/models/minimax_m3/amd/ops/index_topk.py
@@ -37,6 +37,7 @@
 MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 64
 MAX_DECODE_SCORE_KV_CHUNKS = 256
 DECODE_TOPK_BLOCKS_PER_CHUNK = 512
+DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK = 128
 MAX_DECODE_TOPK_FAST_CHUNKS = 16
 DECODE_TOPK_TARGET_GRID = 64
 
@@ -103,14 +104,23 @@ def _decode_topk_launch_policy(
     topk: int,
     *,
     is_gfx950: bool,
-) -> tuple[int, bool, bool]:
-    """Choose the graph grid and bounded-context selector specialization."""
+) -> tuple[int, int, int, int, bool, bool]:
+    """Choose the selector grid and compile-time launch configuration."""
     if (
         is_gfx950
         and topk == 16
         and 0 < max_block <= MAX_DECODE_TOPK_FAST_CHUNKS * DECODE_TOPK_BLOCKS_PER_CHUNK
     ):
-        return MAX_DECODE_TOPK_FAST_CHUNKS, True, True
+        if max_block <= DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK:
+            return 1, DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK, 2, 1, True, True
+        return (
+            MAX_DECODE_TOPK_FAST_CHUNKS,
+            DECODE_TOPK_BLOCKS_PER_CHUNK,
+            4,
+            2,
+            True,
+            True,
+        )
 
     target = max(
         1,
@@ -119,7 +129,14 @@ def _decode_topk_launch_policy(
             DECODE_TOPK_TARGET_GRID // max(1, total_q * num_idx_heads),
         ),
     )
-    return 1 << (target.bit_length() - 1), False, False
+    return (
+        1 << (target.bit_length() - 1),
+        DECODE_TOPK_BLOCKS_PER_CHUNK,
+        8,
+        2,
+        False,
+        False,
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -1526,6 +1543,9 @@ def minimax_m3_index_decode(
     # only the number of chunks needed for its live context.
     (
         num_topk_chunks,
+        topk_block_size,
+        selector_num_warps,
+        selector_num_stages,
         single_tile_guaranteed,
         adaptive_final_merge,
     ) = _decode_topk_launch_policy(
@@ -1609,12 +1629,12 @@ def minimax_m3_index_decode(
         pages_per_sparse_block=PAGES_PER_SPARSE_BLOCK,
         block_page_stride=selector_block_page_stride,
         NUM_TOPK_CHUNKS=num_topk_chunks,
-        BLOCK_SIZE_K=512,
+        BLOCK_SIZE_K=topk_block_size,
         BLOCK_SIZE_T=block_size_t,
         EMIT_SPARSE_TABLE=emit_sparse_table,
         SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
         ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
-        num_warps=4 if adaptive_final_merge else 8,
-        num_stages=2,
+        num_warps=selector_num_warps,
+        num_stages=selector_num_stages,
     )
     return topk_idx
diff --git a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py
index 836649b725b4..4dd2b91475fd 100644
--- a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py
+++ b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py
@@ -83,6 +83,8 @@ def _swiglu_oai_quant_kernel(
     limit,
     HAS_LIMIT: tl.constexpr,
     BLOCK_M: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    TINY: tl.constexpr,
 ):
     """SwiGLU-OAI (split layout) fused with per-32-block MXFP8 (E4M3 + E8M0)
     quant. Each program handles ``[BLOCK_M, 32]`` of the ``[M, I]`` output (one
@@ -109,10 +111,12 @@ def _swiglu_oai_quant_kernel(
         gate = tl.minimum(gate, limit)
         up = tl.minimum(tl.maximum(up, -limit), limit)
     act = gate * tl.sigmoid(alpha * gate) * (up + beta)  # [BLOCK_M, 32] fp32
-    amax = tl.maximum(tl.max(tl.abs(act), axis=1), 1e-30)  # [BLOCK_M]
-    sb = tl.minimum(tl.maximum(tl.floor(tl.log2(amax)) + 127.0, 0.0), 254.0)
-    descale = tl.exp2(sb - 127.0)
-    aq = (act / descale[:, None]).to(aq_ptr.dtype.element_ty)
+    amax = tl.maximum(tl.max(tl.abs(act), axis=1), TINY)  # [BLOCK_M]
+    sb = tl.minimum(tl.maximum(tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0, 0.0), 254.0)
+    # Computing the E8M0 byte-0 divisor (2^-127) can flush to zero on CDNA.
+    # Use its finite reciprocal so all-zero blocks do not produce 0/0.
+    rescale = tl.exp2(127.0 - sb)
+    aq = (act * rescale[:, None]).to(aq_ptr.dtype.element_ty)
     tl.store(
         aq_ptr + offs_m[:, None] * stride_qm + offs_c[None, :] * stride_qn,
         aq,
@@ -173,6 +177,8 @@ def swiglu_oai_quantize_mxfp8(
         0.0 if limit is None else float(limit),
         HAS_LIMIT=limit is not None,
         BLOCK_M=block_m,
+        FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
+        TINY=float(torch.finfo(torch.float32).tiny),
         num_warps=4,
     )
     return aq, asc
diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py
index f5a8db0c2443..2459d4ef856a 100644
--- a/vllm/models/minimax_m3/common/indexer.py
+++ b/vllm/models/minimax_m3/common/indexer.py
@@ -55,6 +55,7 @@
     AttentionSpec,
     KVCacheSpec,
     MLAAttentionSpec,
+    SparseCacheRole,
 )
 
 logger = init_logger(__name__)
@@ -90,7 +91,7 @@ def get_supported_head_sizes(cls) -> list[int]:
         return [128]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [128]
 
     @classmethod
@@ -148,6 +149,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
             num_kv_heads=1,
             head_size=self.head_dim,
             dtype=self.dtype,
+            cache_role=SparseCacheRole.INDEXER,
         )
 
     def forward(self) -> None: ...
diff --git a/vllm/models/minimax_m3/common/mm_preprocess.py b/vllm/models/minimax_m3/common/mm_preprocess.py
index 9f1d08011a03..958e4f189b78 100644
--- a/vllm/models/minimax_m3/common/mm_preprocess.py
+++ b/vllm/models/minimax_m3/common/mm_preprocess.py
@@ -11,9 +11,7 @@
 from transformers.video_utils import VideoMetadata
 
 from vllm.config.multimodal import (
-    BaseDummyOptions,
-    ImageDummyOptions,
-    VideoDummyOptions,
+    MultiModalDummyOptions,
 )
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal.inputs import (
@@ -32,6 +30,7 @@
     PromptUpdate,
     PromptUpdateDetails,
 )
+from vllm.multimodal.processing.processor import HFMultiModalInputs
 from vllm.multimodal.video import (
     VIDEO_LOADER_REGISTRY,
     VideoBackend,
@@ -282,7 +281,7 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
         size = self.info.get_image_size_with_most_features()
         video_size = self.info.get_video_size_with_most_features()
@@ -292,14 +291,14 @@ def get_dummy_mm_data(
                 width=size.width,
                 height=size.height,
                 num_images=mm_counts.get("image", 0),
-                overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
+                overrides=mm_options.get("image"),
             ),
             "video": self._get_dummy_videos(
                 width=video_size.width,
                 height=video_size.height,
                 num_frames=num_frames,
                 num_videos=mm_counts.get("video", 0),
-                overrides=cast(VideoDummyOptions | None, mm_options.get("video")),
+                overrides=mm_options.get("video"),
             ),
         }
 
@@ -307,26 +306,21 @@ def get_dummy_mm_data(
 class MiniMaxM3VLMultiModalProcessor(
     BaseMultiModalProcessor[MiniMaxM3VLProcessingInfo]
 ):
-    def _apply_hf_processor_main(
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
+        return self.dummy_inputs.get_dummy_text(mm_counts)
+
+    def _get_hf_mm_inputs(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
-
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
-
-        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+        hf_data = hf_inputs.hf_data
 
-        mm_data = dict(mm_data)
         # With ``video_needs_metadata=True`` each video arrives as a
         # ``(frames, metadata)`` tuple. Split the frames back out and forward the
         # metadata as ``VideoMetadata`` so the processor emits timestamps.
-        videos = cast(list | None, mm_data.get("videos"))
+        videos = cast(list | None, hf_data.get("videos"))
         video_metadata: list[VideoMetadata] | None = None
         if videos:
             frames_only = []
@@ -346,21 +340,25 @@ def _apply_hf_processor_main(
                 # stays consistent with _get_prompt_updates.
                 meta.setdefault("total_num_frames", len(frames))
                 video_metadata.append(VideoMetadata(**meta))
-            mm_data["videos"] = frames_only
+            hf_data["videos"] = frames_only
 
-        # Override the video processor's default do_resize=False (set for a
-        # pre-resized pipeline) to True for vLLM's raw-frame inputs.
-        merged = dict(do_resize=True, **hf_processor_mm_kwargs)
-        data = dict(text=prompt_text, **mm_data)
         if video_metadata is not None:
-            data["video_metadata"] = video_metadata
-        processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            data,
+            hf_data["video_metadata"] = video_metadata
+
+        return hf_inputs
+
+    def _call_hf_processor(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        # Override the video processor's default for vLLM's raw-frame inputs.
+        merged = dict(do_resize=True, **hf_kwargs)
+        return self.info.ctx.call_hf_processor(
+            self.info.get_hf_processor(**hf_kwargs),
+            hf_data,
             merged,
         )
-        processed_data.update(passthrough_data)
-        return processed_data
 
     def _get_mm_fields_config(
         self,
diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py
index 50c03dd9bf0a..75aa90100c8f 100644
--- a/vllm/models/minimax_m3/common/sparse_attention.py
+++ b/vllm/models/minimax_m3/common/sparse_attention.py
@@ -159,7 +159,7 @@ def get_supported_head_sizes(cls) -> list[int]:
         return [128]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # Page size == sparse block size (one sparse block per KV page).
         return [128]
 
@@ -314,7 +314,10 @@ def build(
             prefill_cu_seqlens_k = torch.empty(
                 num_prefills + 1, dtype=torch.int32, device=seq_lens.device
             )
-            prefill_cu_seqlens_k[0] = 0
+            if current_platform.is_rocm():
+                prefill_cu_seqlens_k[:1].zero_()
+            else:
+                prefill_cu_seqlens_k[0] = 0
             torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:])
             prefill_cu_seqlens_q = (
                 query_start_loc[num_decodes:] - num_decode_tokens
diff --git a/vllm/models/minimax_m3/common/vision_tower.py b/vllm/models/minimax_m3/common/vision_tower.py
index 23b8b3ed3197..468fee42371d 100644
--- a/vllm/models/minimax_m3/common/vision_tower.py
+++ b/vllm/models/minimax_m3/common/vision_tower.py
@@ -6,7 +6,7 @@
 import torch
 import torch.nn as nn
 from einops import rearrange
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.distributed import parallel_state
 from vllm.distributed import utils as dist_utils
@@ -14,6 +14,7 @@
 from vllm.model_executor.layers.attention.mm_encoder_attention import (
     MMEncoderAttention,
 )
+from vllm.model_executor.layers.conv import Conv3dLayer
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -43,7 +44,7 @@ class MiniMaxVLPatchEmbed(nn.Module):
     and projects each to a hidden-size embedding.
     """
 
-    def __init__(self, config: PretrainedConfig) -> None:
+    def __init__(self, config: PreTrainedConfig) -> None:
         super().__init__()
         compression = config.img_token_compression_config
         temporal_patch_size = compression.get("temporal_patch_size", 2)
@@ -55,7 +56,7 @@ def __init__(self, config: PretrainedConfig) -> None:
         self.num_channels = num_channels
         self.hidden_size = config.hidden_size
 
-        self.patch_embedding = nn.Conv3d(
+        self.patch_embedding = Conv3dLayer(
             in_channels=num_channels,
             out_channels=config.hidden_size,
             kernel_size=(temporal_patch_size, patch_size, patch_size),
@@ -231,7 +232,7 @@ class MiniMaxVLEncoderLayer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
@@ -294,7 +295,7 @@ def forward(
 class MiniMaxVLEncoder(nn.Module):
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         num_hidden_layers_override: int | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -348,7 +349,7 @@ class MiniMaxVLVisionTransformer(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         num_hidden_layers_override: int | None = None,
         require_post_norm: bool | None = None,
         quant_config: QuantizationConfig | None = None,
@@ -678,7 +679,7 @@ class MiniMaxVLVisionModel(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         text_hidden_size: int,
         projector_hidden_size: int | None = None,
         quant_config: QuantizationConfig | None = None,
diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py
index 970b06d407ea..4fb260330f14 100644
--- a/vllm/models/minimax_m3/nvidia/model.py
+++ b/vllm/models/minimax_m3/nvidia/model.py
@@ -16,7 +16,7 @@
 
 import torch
 from torch import nn
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm import _custom_ops as ops
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
@@ -92,7 +92,7 @@
 )
 
 
-def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:
+def _sparse_attention_layer_ids(config: PreTrainedConfig) -> set[int]:
     """Layer ids whose attention runs the extra sparse "index" branch."""
     cfg = getattr(config, "sparse_attention_config", None)
     if not cfg:
@@ -103,7 +103,7 @@ def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:
     return {i for i, f in enumerate(freq) if f != 0}
 
 
-def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool:
+def _is_moe_layer(config: PreTrainedConfig, layer_id: int) -> bool:
     """Whether this layer's MLP is a sparse MoE block (vs a dense MLP)."""
     moe_layer_freq = getattr(config, "moe_layer_freq", None)
     if moe_layer_freq is None:
@@ -147,7 +147,7 @@ class MiniMaxM3MLP(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         intermediate_size: int,
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
@@ -194,7 +194,7 @@ class MiniMaxM3MoE(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_id: int,
         quant_config: QuantizationConfig | None = None,
         reduce_results: bool = True,
@@ -290,7 +290,7 @@ class MiniMaxM3Attention(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_id: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -403,7 +403,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         layer_id: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -1130,7 +1130,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         with self._mark_tower_model(vllm_config, {"image", "video"}):
             vision_config = config.vision_config
             self.vision_tower = MiniMaxVLVisionModel(
-                config=PretrainedConfig.from_dict(vision_config),
+                config=PreTrainedConfig.from_dict(vision_config),
                 text_hidden_size=text_hidden_size,
                 projector_hidden_size=projector_hidden_size,
                 quant_config=self.quant_config,
diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/indexer_qsa.py
index 6b7d8f9f9380..2b8c7b144420 100644
--- a/vllm/models/qwen4_exp/amd/indexer_qsa.py
+++ b/vllm/models/qwen4_exp/amd/indexer_qsa.py
@@ -33,7 +33,6 @@ def apply_qsa_rope(
     tensor: torch.Tensor,
 ) -> torch.Tensor:
     """Apply the main attention's exact 1D/MRoPE composition to QSA heads."""
-
     num_tokens, _, head_dim = tensor.shape
     rotary_dim = rotary_emb.rotary_dim
     cache = rotary_emb._match_cos_sin_cache_dtype(tensor)  # noqa: SLF001
@@ -67,7 +66,6 @@ def apply_qsa_rmsnorm(
     tensor: torch.Tensor,
 ) -> torch.Tensor:
     """Use vLLM's portable RMSNorm implementation on ROCm."""
-
     return cast(torch.Tensor, norm(tensor))
 
 
@@ -153,7 +151,6 @@ def project_qk(
         positions: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Project replicated Q/K, normalize+rotate Q, and preserve raw K."""
-
         qk, _ = self.index_qk_proj(hidden_states)
         q_raw, token_k = qk.split(
             (
@@ -176,7 +173,6 @@ def normalize_compressed_keys(
         first_rope_positions: torch.Tensor,
     ) -> torch.Tensor:
         """Normalize pooled K and apply the first token's exact group position."""
-
         keys = compressed_keys.reshape(-1, self.index_head_dim)
         keys = apply_qsa_rmsnorm(self.k_layernorm, keys).reshape(
             -1, 1, self.index_head_dim
@@ -284,7 +280,6 @@ def forward(
         out: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Return fixed-width request-relative token indices padded with ``-1``."""
-
         metadata = self._metadata()
         if metadata is None:
             # Preserve step-0 indices when later MTP steps reuse the buffer.
diff --git a/vllm/models/qwen4_exp/amd/low_latency_gemm.py b/vllm/models/qwen4_exp/amd/low_latency_gemm.py
index e8c5c2f917ec..c6e54885e2f6 100644
--- a/vllm/models/qwen4_exp/amd/low_latency_gemm.py
+++ b/vllm/models/qwen4_exp/amd/low_latency_gemm.py
@@ -11,5 +11,4 @@ def enable_qwen4_exp_low_latency_gemm(
     dtype: torch.dtype,
 ) -> None:
     """Keep the standard vLLM linear methods on AMD ROCm."""
-
     del module, dtype
diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py
index aeb3a14dadb8..b3209678922d 100644
--- a/vllm/models/qwen4_exp/amd/model.py
+++ b/vllm/models/qwen4_exp/amd/model.py
@@ -76,7 +76,7 @@
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.kv_cache_interface import MambaSpec
 
-from ..config import Qwen4ExpConfig
+from ..config import ATTENTION_LAYER_TYPES, QSA_LAYER_TYPE, Qwen4ExpConfig
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
 from .ple_layer import Qwen4ExpPLELayer
@@ -87,7 +87,6 @@ def without_modelopt_fp4(
     quant_config: QuantizationConfig | None,
 ) -> QuantizationConfig | None:
     """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4."""
-
     if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
         return None
     return quant_config
@@ -103,7 +102,6 @@ def _remap_qsa_cache_scale_name(
     that cache directly, so only QSA layers need the final path component
     moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers.
     """
-
     scale_suffixes = {
         "k_proj.k_scale": "_k_scale",
         "k_proj.output_scale": "_k_scale",
@@ -214,8 +212,11 @@ def __init__(
                 prefix=f"{prefix}.linear_attn",
                 gqa_interleaved_layout=False,
             )
-        elif layer_type == "full_attention":
-            use_qsa = getattr(config, "indexer_n_heads", None) is not None
+        elif layer_type in ATTENTION_LAYER_TYPES:
+            use_qsa = (
+                layer_type == QSA_LAYER_TYPE
+                or getattr(config, "indexer_n_heads", None) is not None
+            )
             if not use_qsa:
                 self.self_attn = Qwen3NextAttention(
                     config,
@@ -311,7 +312,7 @@ def forward(
 
         if self.layer_type == "linear_attention":
             attn_out = self.linear_attn(hidden_states=block_input)
-        elif self.layer_type == "full_attention":
+        elif self.layer_type in ATTENTION_LAYER_TYPES:
             attn_out = self.self_attn(
                 hidden_states=block_input,
                 positions=positions,
@@ -402,8 +403,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self._qsa_layer_ids = frozenset(
             layer_idx
             for layer_idx, layer_type in enumerate(config.layer_types)
-            if layer_type == "full_attention"
-            and getattr(config, "indexer_n_heads", None) is not None
+            if layer_type == QSA_LAYER_TYPE
+            or (
+                layer_type == "full_attention"
+                and getattr(config, "indexer_n_heads", None) is not None
+            )
         )
         self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size)
 
diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py
index 40ee6a6d4797..a2b420603f06 100644
--- a/vllm/models/qwen4_exp/amd/mtp.py
+++ b/vllm/models/qwen4_exp/amd/mtp.py
@@ -91,7 +91,6 @@ def _remap_quantized_layers(
 
 def _remap_mtp_weight_name(name: str) -> str | None:
     """Map Qwen4Exp checkpoint paths into the standalone draft model."""
-
     for checkpoint_prefix in (
         "model.language_model.",
         "language_model.",
@@ -260,7 +259,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
 
     def _iter_qsa_attentions(self):
         """Yield MTP attention modules that own a QSA indexer."""
-
         for layer in self.layers:
             attention = getattr(layer, "self_attn", None)
             if (
@@ -271,13 +269,11 @@ def _iter_qsa_attentions(self):
 
     def set_skip_topk(self, skip: bool) -> None:
         """Select on MTP step 0 and reuse its QSA indices on later steps."""
-
         for attention in self._iter_qsa_attentions():
             attention.indexer.skip_topk = skip
 
     def compact_topk_indices(self, row_indices: torch.Tensor) -> None:
         """Keep each request's target-aligned step-0 sparse-index row."""
-
         num_rows = row_indices.numel()
         for attention in self._iter_qsa_attentions():
             buffer = attention.topk_indices_buffer
diff --git a/vllm/models/qwen4_exp/amd/ops/qsa.py b/vllm/models/qwen4_exp/amd/ops/qsa.py
index 2e45db711d85..6bdb1fe592df 100644
--- a/vllm/models/qwen4_exp/amd/ops/qsa.py
+++ b/vllm/models/qwen4_exp/amd/ops/qsa.py
@@ -600,7 +600,6 @@ def qsa_mqa_paged(
     score_scale: float | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Compute QSA scores directly from a paged compressed-key cache."""
-
     _validate_mqa(q)
     if not q.is_cuda or not HAS_TRITON:
         raise RuntimeError("paged QSA scoring requires a GPU and Triton")
@@ -678,7 +677,6 @@ def expand_qsa_block_indices_cuda(
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Expand compressed blocks and compact the causal tail of the open group."""
-
     if not block_indices.is_cuda or not HAS_TRITON:
         raise RuntimeError("QSA index expansion requires a GPU and Triton")
     if token_topk % compress_ratio:
@@ -738,7 +736,6 @@ def qsa_select_paged_tokens(
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Score, select, and expand QSA indices without host synchronization."""
-
     rows = q.shape[0]
     output_width = token_topk + compress_ratio - 1
     if out is None:
@@ -829,7 +826,6 @@ def qsa_sparse_paged_attention(
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Run sparse GQA directly over paged BF16 K/V caches."""
-
     if not q.is_cuda or not HAS_TRITON:
         raise RuntimeError("paged QSA sparse attention requires a GPU and Triton")
     if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape:
@@ -971,7 +967,6 @@ def qsa_store_cache_rows(
     rows: torch.Tensor,
 ) -> None:
     """Store fixed-width rows in a QSA cache without boolean indexing."""
-
     if not cache.is_cuda or not HAS_TRITON:
         raise RuntimeError("QSA cache stores require a GPU and Triton")
     if cache.ndim != 4 or cache.shape[2] != 1:
@@ -1017,7 +1012,6 @@ def qsa_compress_groups_with_ratio(
     rope_cache: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Pool completed groups from the compressor-state ring and raw token rows."""
-
     if not raw_keys.is_cuda or not HAS_TRITON:
         raise RuntimeError("QSA compression requires a GPU and Triton")
     rows = token_to_req.numel()
diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py
index fc494be055cf..2616a3e9e02d 100644
--- a/vllm/models/qwen4_exp/amd/ple_layer.py
+++ b/vllm/models/qwen4_exp/amd/ple_layer.py
@@ -357,7 +357,6 @@ def forward(
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load hash buffers and checkpoint-split embedding rows."""
-
         persistent_buffers = {
             "layer_multipliers": self.layer_multipliers,
             "ngram_heads_offsets": self.ngram_heads_offsets,
diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py
index c0c766aad141..3c8a61bd547b 100644
--- a/vllm/models/qwen4_exp/amd/qsa.py
+++ b/vllm/models/qwen4_exp/amd/qsa.py
@@ -75,7 +75,7 @@ def get_name() -> str:
         return "QWEN4_EXP_QSA_TRITON"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # QSA consumes manager pages directly and does not use FA4 paged attention.
         return [MultipleOf(16)]
 
@@ -430,7 +430,6 @@ def qwen4_exp_qsa_with_output(
     layer_name: LayerNameType,
 ) -> None:
     """Run the complete QSA state/update/attend transaction."""
-
     layer_name = _resolve_layer_name(layer_name)
     layer = get_forward_context().no_compile_layers[layer_name]
     if not isinstance(layer, Qwen4ExpQSAAttention):
diff --git a/vllm/models/qwen4_exp/common/ple.py b/vllm/models/qwen4_exp/common/ple.py
index bd132685b528..72fc494b6037 100644
--- a/vllm/models/qwen4_exp/common/ple.py
+++ b/vllm/models/qwen4_exp/common/ple.py
@@ -28,7 +28,6 @@ def compute_ple_shard_overlap(
     tp_end: int,
 ) -> PLEShardOverlap | None:
     """Compute the overlap of a checkpoint shard and one TP vocabulary range."""
-
     if checkpoint_start < 0 or checkpoint_rows < 0:
         raise ValueError("checkpoint shard bounds must be non-negative")
     if tp_start < 0 or tp_end < tp_start:
@@ -54,7 +53,6 @@ def copy_ple_embedding_shard_(
     tp_end: int,
 ) -> int:
     """Copy the overlapping rows of a PLE checkpoint shard into a TP table."""
-
     if destination.ndim == 0 or loaded_weight.ndim != destination.ndim:
         raise ValueError("destination and loaded weight must have matching ranks")
     if destination.shape[1:] != loaded_weight.shape[1:]:
diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py
index e806dc0375de..eb0758025fb3 100644
--- a/vllm/models/qwen4_exp/common/qsa_cache.py
+++ b/vllm/models/qwen4_exp/common/qsa_cache.py
@@ -47,7 +47,6 @@
 
 def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor:
     """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows."""
-
     if positions.ndim == 1:
         positions = positions.unsqueeze(0).expand(3, -1)
     elif positions.ndim != 2 or positions.shape[0] not in (1, 3):
@@ -115,7 +114,6 @@ def circular_qsa_slot_mapping(
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Map each request to its fixed physical block as a circular token ring."""
-
     if compressor_state_size <= 0:
         raise ValueError("QSA circular buffer size must be positive")
     if block_table.ndim != 2:
@@ -167,7 +165,6 @@ def compressed_qsa_slot_mapping(
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Build boundary-only slots for an ``MLAAttentionSpec`` QSA cache."""
-
     if storage_block_size <= 0 or compress_ratio <= 0:
         raise ValueError("QSA block size and compression ratio must be positive")
     compressed_positions = torch.div(
diff --git a/vllm/models/qwen4_exp/config.py b/vllm/models/qwen4_exp/config.py
index a8eedab83f6b..dfac816d5e2c 100644
--- a/vllm/models/qwen4_exp/config.py
+++ b/vllm/models/qwen4_exp/config.py
@@ -4,7 +4,7 @@
 
 from typing import Any, ClassVar, cast
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.models.qwen3_vl.configuration_qwen3_vl import (
     Qwen3VLVisionConfig,
 )
@@ -19,6 +19,12 @@
     "indexer_compress_ratio",
 )
 
+# Transformers labels Qwen4Exp's sparse-attention layers `qwen_sparse_attention`.
+# Older checkpoints label them `full_attention` and mark QSA by setting the
+# indexer fields for the whole model.
+QSA_LAYER_TYPE = "qwen_sparse_attention"
+ATTENTION_LAYER_TYPES = ("full_attention", QSA_LAYER_TYPE)
+
 
 class Qwen4ExpVisionConfig(Qwen3VLVisionConfig):
     model_type = "qwen4_exp"
@@ -134,6 +140,11 @@ def _validate_ple_layer_ids(self) -> None:
     def _validate_qsa_config(self) -> None:
         configured = {name: getattr(self, name, None) for name in _QSA_CONFIG_FIELDS}
         if all(value is None for value in configured.values()):
+            if QSA_LAYER_TYPE in (self.layer_types or ()):
+                raise ValueError(
+                    f"{QSA_LAYER_TYPE} layers require the QSA config fields: "
+                    f"{list(_QSA_CONFIG_FIELDS)}"
+                )
             return
 
         missing = [name for name, value in configured.items() if value is None]
@@ -165,7 +176,7 @@ def _validate_qsa_config(self) -> None:
     @property
     def layers_block_type(self) -> list[str]:
         return [
-            "attention" if layer_type == "full_attention" else layer_type
+            "attention" if layer_type in ATTENTION_LAYER_TYPES else layer_type
             for layer_type in self.layer_types
         ]
 
@@ -190,9 +201,9 @@ def ngram_context_len(self) -> int:
         return max(int(self.ngram_size) - 1, 0)
 
 
-class Qwen4ExpConfig(PretrainedConfig):
+class Qwen4ExpConfig(PreTrainedConfig):
     model_type = "qwen4_exp"
-    sub_configs: ClassVar[dict[str, type[PretrainedConfig]]] = {
+    sub_configs: ClassVar[dict[str, type[PreTrainedConfig]]] = {
         "vision_config": Qwen4ExpVisionConfig,
         "text_config": Qwen4ExpTextConfig,
     }
diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py
index bbdc5e4feba9..f4b968f15201 100644
--- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py
@@ -32,7 +32,6 @@ def apply_qsa_rope(
     tensor: torch.Tensor,
 ) -> torch.Tensor:
     """Apply the main attention's exact 1D/MRoPE composition to QSA heads."""
-
     num_tokens, _, head_dim = tensor.shape
     rotary_dim = rotary_emb.rotary_dim
     cache = rotary_emb._match_cos_sin_cache_dtype(tensor)  # noqa: SLF001
@@ -245,7 +244,6 @@ def forward(
         valid-entry count (the attention kernel's loop bound, never a token
         index).
         """
-
         metadata = self._metadata()
         if metadata is None:
             # Preserve step-0 indices when later MTP steps reuse the buffer.
diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py
index 28c3ea69d11b..2e3833965574 100644
--- a/vllm/models/qwen4_exp/nvidia/model.py
+++ b/vllm/models/qwen4_exp/nvidia/model.py
@@ -75,7 +75,7 @@
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.kv_cache_interface import MambaSpec
 
-from ..config import Qwen4ExpConfig
+from ..config import ATTENTION_LAYER_TYPES, QSA_LAYER_TYPE, Qwen4ExpConfig
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
 from .ple_layer import Qwen4ExpPLELayer
@@ -86,7 +86,6 @@ def without_modelopt_fp4(
     quant_config: QuantizationConfig | None,
 ) -> QuantizationConfig | None:
     """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4."""
-
     if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
         return None
     return quant_config
@@ -102,7 +101,6 @@ def _remap_qsa_cache_scale_name(
     that cache directly, so only QSA layers need the final path component
     moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers.
     """
-
     scale_suffixes = {
         "k_proj.k_scale": "_k_scale",
         "k_proj.output_scale": "_k_scale",
@@ -215,8 +213,11 @@ def __init__(
                 prefix=f"{prefix}.linear_attn",
                 gqa_interleaved_layout=False,
             )
-        elif layer_type == "full_attention":
-            use_qsa = getattr(config, "indexer_n_heads", None) is not None
+        elif layer_type in ATTENTION_LAYER_TYPES:
+            use_qsa = (
+                layer_type == QSA_LAYER_TYPE
+                or getattr(config, "indexer_n_heads", None) is not None
+            )
             if not use_qsa:
                 self.self_attn = Qwen3NextAttention(
                     config,
@@ -297,7 +298,7 @@ def forward(
 
             if input_ids is None or query_start_loc is None or ngram_context is None:
                 raise RuntimeError("PLE inputs were not prepared")
-            hidden_states = hidden_states + self.ple(
+            hidden_states = self.ple(
                 hidden_states,
                 input_ids,
                 query_start_loc,
@@ -314,7 +315,7 @@ def forward(
 
         if self.layer_type == "linear_attention":
             attn_out = self.linear_attn(hidden_states=block_input)
-        elif self.layer_type == "full_attention":
+        elif self.layer_type in ATTENTION_LAYER_TYPES:
             attn_out = self.self_attn(
                 hidden_states=block_input,
                 positions=positions,
@@ -394,8 +395,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self._qsa_layer_ids = frozenset(
             layer_idx
             for layer_idx, layer_type in enumerate(config.layer_types)
-            if layer_type == "full_attention"
-            and getattr(config, "indexer_n_heads", None) is not None
+            if layer_type == QSA_LAYER_TYPE
+            or (
+                layer_type == "full_attention"
+                and getattr(config, "indexer_n_heads", None) is not None
+            )
         )
         self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size)
 
diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py
index d12b851d688e..af3bcd4659f8 100644
--- a/vllm/models/qwen4_exp/nvidia/mtp.py
+++ b/vllm/models/qwen4_exp/nvidia/mtp.py
@@ -90,7 +90,6 @@ def _remap_quantized_layers(
 
 def _remap_mtp_weight_name(name: str) -> str | None:
     """Map Qwen4Exp checkpoint paths into the standalone draft model."""
-
     for checkpoint_prefix in (
         "model.language_model.",
         "language_model.",
@@ -250,7 +249,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
 
     def _iter_qsa_attentions(self):
         """Yield MTP attention modules that own a QSA indexer."""
-
         for layer in self.layers:
             attention = getattr(layer, "self_attn", None)
             if (
@@ -261,13 +259,11 @@ def _iter_qsa_attentions(self):
 
     def set_skip_topk(self, skip: bool) -> None:
         """Select on MTP step 0 and reuse its QSA indices on later steps."""
-
         for attention in self._iter_qsa_attentions():
             attention.indexer.skip_topk = skip
 
     def compact_topk_indices(self, row_indices: torch.Tensor) -> None:
         """Keep each request's target-aligned step-0 sparse-index row."""
-
         num_rows = row_indices.numel()
         for attention in self._iter_qsa_attentions():
             buffer = attention.topk_indices_buffer
diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py
index 84c4e7c5e1bb..c26fd87cff4f 100644
--- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py
+++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py
@@ -869,7 +869,6 @@ def start_prefetch(
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load hash buffers and checkpoint-split embedding rows."""
-
         persistent_buffers = {
             "layer_multipliers": self.layer_multipliers,
             "ngram_heads_offsets": self.ngram_heads_offsets,
diff --git a/vllm/models/qwen4_exp/nvidia/ops/ple.py b/vllm/models/qwen4_exp/nvidia/ops/ple.py
index d76e2705a7bf..bc4c99957d90 100644
--- a/vllm/models/qwen4_exp/nvidia/ops/ple.py
+++ b/vllm/models/qwen4_exp/nvidia/ops/ple.py
@@ -322,6 +322,7 @@ def _ple_conv_kernel(
     state_ptr,
     w_ptr,
     residual_ptr,
+    outer_residual_ptr,
     state_idx_ptr,
     qsl_ptr,
     num_acc_ptr,
@@ -442,11 +443,21 @@ def _ple_conv_kernel(
         mask=c_mask,
         other=0.0,
     )
+    # Preserve the original eager operation boundaries: short convolution is
+    # first accumulated into the BF16/FP16 PLE output, then the outer residual
+    # is added to that rounded value.
+    ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty)
+    outer_residual = tl.load(
+        outer_residual_ptr + output_t * C + c_offs,
+        mask=c_mask,
+        other=0.0,
+    )
+    ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32)
     if launch_pdl:
         tl.extra.cuda.gdc_launch_dependents()
     tl.store(
         residual_ptr + output_t * C + c_offs,
-        residual + conv_output,
+        ple_output,
         mask=c_mask,
     )
 
@@ -564,6 +575,7 @@ def ple_conv(
     conv_state: torch.Tensor,
     conv_weights: torch.Tensor,
     state_indices: torch.Tensor,
+    outer_residual: torch.Tensor,
     *,
     mode: Literal["decode", "spec", "prefill"],
     dilation: int,
@@ -573,7 +585,7 @@ def ple_conv(
     spec_query_len: int = 1,
     token_indices: torch.Tensor | None = None,
 ) -> None:
-    """Add short-convolution output to ``residual`` and update its state."""
+    """Add short convolution and the outer residual; update state."""
     BLOCK_C = 512
     kernel_spec_query_len = spec_query_len if mode == "spec" else 1
     T, C = inputs.shape
@@ -624,6 +636,7 @@ def ple_conv(
         conv_state,
         conv_weights,
         residual,
+        outer_residual,
         state_indices,
         query_start_loc,
         num_accepted_tokens,
diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
index 709acbde00a4..401048e02c55 100644
--- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py
@@ -4,15 +4,30 @@
 
 from __future__ import annotations
 
+from functools import lru_cache
+
 import torch
 
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     TritonWarmupTensor,
     triton_scalar_specialization_rep,
 )
+from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON, tl, triton
 
 
+@lru_cache(maxsize=1)
+def _is_sm120() -> bool:
+    """True on sm_120 (RTX PRO 6000 Blackwell): selects the sm_120 tuning table."""
+    return current_platform.get_device_capability() == (12, 0)
+
+
+@lru_cache(maxsize=1)
+def _is_sm90() -> bool:
+    """True on sm_90 (H100/H200/H20): selects the sm_90 tuning table."""
+    return current_platform.get_device_capability() == (9, 0)
+
+
 @triton.jit(do_not_specialize=["num_rows", "num_requests"])
 def _qsa_sparse_paged_gqa_splitk_kernel(
     q_ptr,
@@ -24,6 +39,9 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
     partial_output_ptr,
     partial_lse_ptr,
     output_ptr,
+    softmax_scale,
+    output_scale,
+    output_gate_ptr,
     stride_q_row,
     stride_q_head,
     stride_k_block,
@@ -36,6 +54,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
     stride_table_req,
     stride_output_row,
     stride_output_head,
+    stride_output_gate_row,
+    stride_output_gate_head,
     num_rows,
     num_cache_blocks,
     num_requests,
@@ -49,6 +69,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
     NUM_TILES: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    IS_FP8: tl.constexpr,
 ) -> None:
     row = tl.program_id(0)
     kv_head = tl.program_id(1)
@@ -78,7 +99,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
     max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32)
     normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32)
     accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)
-    softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634
+    # softmax_scale is the host-side attention scale (1/sqrt(head_dim), with the
+    # fp8 K dequant scale already pre-multiplied in); convert to log2 units once
+    # here for the exp2-based online softmax.
+    score_scale = softmax_scale * 1.4426950408889634
 
     tile_end = tl.minimum(NUM_TILES, tl.cdiv(tl.minimum(valid_count, TOPK), BLOCK_N))
 
@@ -126,15 +150,26 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
             mask=valid[:, None],
             other=0.0,
         )
+        if IS_FP8:
+            # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK
+            # measured slower here and less accurate).
+            keys = keys.to(query.dtype)
         scores = tl.dot(query, keys)
-        # Scaling scores avoids re-quantizing a scaled query to BF16.
-        scores *= softmax_scale_log2
+        # Scaling scores avoids re-quantizing a scaled query to BF16; for fp8
+        # caches the K dequant scale is already folded into softmax_scale on the
+        # host.
+        scores *= score_scale
         scores = tl.where(valid[None, :], scores, -1.0e20)
         next_max = tl.maximum(max_value, tl.max(scores, axis=1))
         alpha = tl.math.exp2(max_value - next_max)
         probabilities = tl.where(
             valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0
         )
+        if IS_FP8:
+            # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online
+            # softmax) so fp16 has the range, its wider mantissa is more
+            # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster.
+            values = values.to(tl.float16)
         accumulator = tl.dot(
             probabilities.to(values.dtype),
             values,
@@ -144,13 +179,30 @@ def _qsa_sparse_paged_gqa_splitk_kernel(
         max_value = next_max
 
     has_values = normalizer > 0
+    # Fold the fp8 V dequant scale (output_scale, 1.0 for bf16) into a per-row
+    # reciprocal normalizer, so the output is a per-row multiply rather than a
+    # HEAD_DIM-wide scale. The split-K LSE below keeps the unscaled normalizer,
+    # so the merge stays correct.
+    inv_normalizer = output_scale / tl.maximum(normalizer, 1.0e-20)
     normalized_output = tl.where(
         has_values[:, None],
-        accumulator / tl.maximum(normalizer[:, None], 1.0e-20),
+        accumulator * inv_normalizer[:, None],
         0.0,
     )
     output_mask = head_offsets[:, None] < GROUP_SIZE
     if NUM_SPLITS == 1:
+        # Preserve the unfused path's BF16 attention-output rounding before
+        # applying the gate in FP32.
+        normalized_output = normalized_output.to(output_ptr.dtype.element_ty)
+        output_gate = tl.load(
+            output_gate_ptr
+            + row * stride_output_gate_row
+            + (first_head + head_offsets[:, None]) * stride_output_gate_head
+            + dim_offsets[None, :],
+            mask=output_mask,
+            other=0.0,
+        ).to(tl.float32)
+        normalized_output = normalized_output.to(tl.float32) * tl.sigmoid(output_gate)
         tl.store(
             output_ptr
             + row * stride_output_row
@@ -186,8 +238,11 @@ def _qsa_merge_splitk_kernel(
     partial_output_ptr,
     partial_lse_ptr,
     output_ptr,
+    output_gate_ptr,
     stride_output_row,
     stride_output_head,
+    stride_output_gate_row,
+    stride_output_gate_head,
     num_rows,
     HEAD_DIM: tl.constexpr,
     NUM_QUERY_HEADS: tl.constexpr,
@@ -219,6 +274,16 @@ def _qsa_merge_splitk_kernel(
     )
     merged = tl.sum(partial_output * weights[:, None], axis=0)
     merged = tl.where(denominator > 0, merged / denominator, 0.0)
+    # Preserve the unfused path's BF16 attention-output rounding before
+    # applying the gate in FP32.
+    merged = merged.to(output_ptr.dtype.element_ty)
+    output_gate = tl.load(
+        output_gate_ptr
+        + row * stride_output_gate_row
+        + head * stride_output_gate_head
+        + dim_offsets
+    ).to(tl.float32)
+    merged = merged.to(tl.float32) * tl.sigmoid(output_gate)
     tl.store(
         output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets,
         merged,
@@ -411,18 +476,98 @@ def _compress_qsa_groups_kernel(
     )
 
 
+def _select_sm120_config(
+    base_programs: int, use_prefill_config: bool, is_fp8: bool
+) -> tuple[int, int, int]:
+    """(block_n, target_splits, num_warps) retuned on sm_120 (RTX PRO 6000
+    Blackwell), split by cache dtype: bf16 and fp8 favour different configs on
+    sm_120, most visibly on the large-prefill region. Each entry is the fastest
+    config that stays correct on its own path; the main win is more warps.
+    """
+    if is_fp8:
+        if base_programs > 2048:
+            return (32, 2, 1) if use_prefill_config else (64, 1, 2)
+        if base_programs <= 24:
+            return 64, 64, 8
+        if base_programs <= 32:
+            return 128, 8, 4
+        if base_programs <= 64:
+            return 64, 8, 4
+        if base_programs <= 128:
+            return 32, 4, 4
+        if base_programs <= 256:
+            return 128, 4, 4
+        if base_programs <= 512:
+            return 64, 4, 4
+        return 32, 2, 1
+    # bf16 K/V cache.
+    if base_programs > 2048:
+        return 64, 2, 2
+    if base_programs <= 24:
+        return 64, 64, 8
+    if base_programs <= 64:
+        return 64, 8, 8
+    if base_programs <= 128:
+        return 32, 8, 8
+    if base_programs <= 256:
+        return 32, 8, 1
+    if base_programs <= 512:
+        return 64, 4, 2
+    return 32, 4, 1
+
+
+def _select_sm90_config(
+    base_programs: int, use_prefill_config: bool, is_fp8: bool
+) -> tuple[int, int, int]:
+    """(block_n, target_splits, num_warps) tuned on sm_90 (H20)."""
+    if base_programs > 2048:
+        return (32, 1, 1) if use_prefill_config else (64, 1, 2)
+    if is_fp8:
+        if base_programs <= 24:
+            return 64, 64, 2
+        if base_programs <= 32:
+            return 32, 16, 1
+        if base_programs <= 64:
+            return 32, 8, 1
+        if base_programs <= 128:
+            return 32, 4, 1
+        if base_programs <= 256:
+            return 32, 8, 1
+        return 32, 4, 1
+    if base_programs <= 32:
+        return 64, 64, 2
+    if base_programs <= 64:
+        return 32, 16, 1
+    if base_programs <= 256:
+        return 32, 8, 1
+    return 32, 4, 1
+
+
 def _select_config(
-    num_rows: int, num_kv_heads: int, use_prefill_config: bool, num_columns: int
+    num_rows: int,
+    num_kv_heads: int,
+    use_prefill_config: bool,
+    num_columns: int,
+    is_fp8: bool = False,
 ) -> tuple[int, int, int, int]:
     """Select (block_n, num_warps, num_tiles, num_splits) for the kernel.
 
-    Tuned on GB300 for the Qwen3.8-Flash-Next TP1/TP2/TP4 shapes, keyed on
-    base_programs = num_rows * num_kv_heads. The bp > 2048 region splits on
-    use_prefill_config (capture-stable: at FULL-graph capture max_query_len is the
-    uniform decode/verify length).
+    Keyed on base_programs = num_rows * num_kv_heads. The bp > 2048 region splits
+    on use_prefill_config (capture-stable: at FULL-graph capture max_query_len is
+    the uniform decode/verify length). This default table was tuned on GB300;
+    sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead,
+    and sm_90 (Hopper) to _select_sm90_config.
     """
     base_programs = num_rows * num_kv_heads
-    if base_programs > 2048:
+    if _is_sm120():
+        BLOCK_N, target_splits, num_warps = _select_sm120_config(
+            base_programs, use_prefill_config, is_fp8
+        )
+    elif _is_sm90():
+        BLOCK_N, target_splits, num_warps = _select_sm90_config(
+            base_programs, use_prefill_config, is_fp8
+        )
+    elif base_programs > 2048:
         BLOCK_N, target_splits, num_warps = (
             (32, 1, 1) if use_prefill_config else (64, 1, 2)
         )
@@ -455,8 +600,17 @@ def qsa_sparse_paged_attention(
     token_to_req: torch.Tensor,
     use_prefill_config: bool,
     out: torch.Tensor | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+    *,
+    output_gate: torch.Tensor,
 ) -> torch.Tensor:
-    """Run sparse GQA directly over paged BF16 K/V caches.
+    """Run sparse GQA directly over paged BF16 or FP8-e4m3 K/V caches.
+
+    With fp8 caches, k_scale/v_scale are the layer's per-tensor dequant scales
+    as host floats (e.g. layer._k_scale_float/_v_scale_float): k_scale is
+    pre-multiplied into the softmax scale and v_scale becomes the kernel's
+    output scale, so the kernel needs no device scale buffers.
 
     logical_indices is the PACKED selection buffer: [rows, selection_width + 1]
     with the trailing column holding each row's valid-entry count (written by
@@ -480,7 +634,19 @@ def qsa_sparse_paged_attention(
         raise ValueError("QSA sparse attention requires valid grouped-query heads")
     head_dim = q.shape[2]
     assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0
-    assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16
+    assert q.dtype == torch.bfloat16
+    assert k_cache.dtype == v_cache.dtype
+    is_fp8 = k_cache.dtype == torch.float8_e4m3fn
+    if is_fp8:
+        assert k_scale is not None and v_scale is not None
+        # Host pre-multiply: fold the K dequant scale into the attention scale
+        # and pass V's dequant scale as the kernel's output scale.
+        softmax_scale = (head_dim**-0.5) * float(k_scale)
+        output_scale = float(v_scale)
+    else:
+        assert k_cache.dtype == torch.bfloat16
+        softmax_scale = head_dim**-0.5
+        output_scale = 1.0
     assert logical_indices.dtype == block_table.dtype == torch.int32
     assert token_to_req.dtype == torch.int32
     assert q.device == k_cache.device == v_cache.device
@@ -496,6 +662,8 @@ def qsa_sparse_paged_attention(
         raise ValueError("QSA sparse output must match its query")
     assert out.dtype == q.dtype and out.device == q.device
     assert out.stride(2) == 1
+    assert output_gate.is_contiguous()
+    output_gate_view = output_gate.view_as(q)
     if not q.shape[0]:
         return out
 
@@ -503,7 +671,7 @@ def qsa_sparse_paged_attention(
     block_m = triton.next_power_of_2(group_size)
     selection_width = logical_indices.shape[1] - 1  # trailing column is the count
     block_n, partial_warps, num_tiles, num_splits = _select_config(
-        q.shape[0], k_cache.shape[2], use_prefill_config, selection_width
+        q.shape[0], k_cache.shape[2], use_prefill_config, selection_width, is_fp8
     )
 
     # Split=1 writes output directly and compiles out all workspace accesses.
@@ -533,6 +701,9 @@ def qsa_sparse_paged_attention(
         partial_output,
         partial_lse,
         out,
+        softmax_scale,
+        output_scale,
+        output_gate_view,
         q.stride(0),
         q.stride(1),
         k_cache.stride(0),
@@ -545,6 +716,8 @@ def qsa_sparse_paged_attention(
         block_table.stride(0),
         out.stride(0),
         out.stride(1),
+        output_gate_view.stride(0),
+        output_gate_view.stride(1),
         q.shape[0],
         k_cache.shape[0],
         block_table.shape[0],
@@ -558,6 +731,7 @@ def qsa_sparse_paged_attention(
         NUM_TILES=num_tiles,
         BLOCK_M=block_m,
         BLOCK_N=block_n,
+        IS_FP8=is_fp8,
         num_warps=partial_warps,
         num_stages=2,
     )
@@ -568,8 +742,11 @@ def qsa_sparse_paged_attention(
         partial_output,
         partial_lse,
         out,
+        output_gate_view,
         out.stride(0),
         out.stride(1),
+        output_gate_view.stride(0),
+        output_gate_view.stride(1),
         q.shape[0],
         HEAD_DIM=q.shape[2],
         NUM_QUERY_HEADS=q.shape[1],
@@ -589,16 +766,20 @@ def warmup_qsa_sparse_paged_attention(
     selection_width: int,
 ) -> tuple[tuple[int, int, int], ...]:
     """Compile every production-reachable split-K/merge specialization."""
-
     head_dim = kv_cache.shape[-1] // 2
     key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1)
+    # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time.
+    is_fp8 = kv_cache.dtype == torch.uint8
+    cache_dtype = torch.float8_e4m3fn if is_fp8 else key_cache.dtype
     num_kv_heads = key_cache.shape[2]
     group_size = num_query_heads // num_kv_heads
     block_m = triton.next_power_of_2(group_size)
 
     # Every config the dispatch can pick for this group size.
     profiles = {
-        _select_config(num_rows, num_kv_heads, use_prefill_config, selection_width)
+        _select_config(
+            num_rows, num_kv_heads, use_prefill_config, selection_width, is_fp8
+        )
         for num_rows in range(1, 8193)
         for use_prefill_config in (False, True)
     }
@@ -612,12 +793,12 @@ def warmup_qsa_sparse_paged_attention(
         torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)
     )
     k_cache_ptr = TritonWarmupTensor(
-        key_cache.dtype,
+        cache_dtype,
         shape=tuple(key_cache.shape),
         strides=tuple(key_cache.stride()),
     )
     v_cache_ptr = TritonWarmupTensor(
-        value_cache.dtype,
+        cache_dtype,
         shape=tuple(value_cache.shape),
         strides=tuple(value_cache.stride()),
     )
@@ -632,6 +813,10 @@ def warmup_qsa_sparse_paged_attention(
     output_ptr = TritonWarmupTensor(
         torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)
     )
+    # The output gate is mandatory at runtime; warm the gated specialization.
+    output_gate_ptr = TritonWarmupTensor(
+        torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)
+    )
     head_stride = head_dim
     row_stride = num_query_heads * head_dim
     num_cache_blocks = triton_scalar_specialization_rep(kv_cache.shape[0])
@@ -659,6 +844,9 @@ def warmup_qsa_sparse_paged_attention(
             partial_output_ptr,
             partial_lse_ptr,
             output_ptr,
+            1.0,
+            1.0,
+            output_gate_ptr,
             row_stride,
             head_stride,
             key_cache.stride(0),
@@ -671,6 +859,8 @@ def warmup_qsa_sparse_paged_attention(
             block_table.stride(0),
             row_stride,
             head_stride,
+            row_stride,
+            head_stride,
             num_rows,
             num_cache_blocks,
             num_requests,
@@ -684,6 +874,7 @@ def warmup_qsa_sparse_paged_attention(
             NUM_TILES=num_tiles,
             BLOCK_M=block_m,
             BLOCK_N=block_n,
+            IS_FP8=is_fp8,
             num_warps=warps,
             num_stages=2,
             grid=(num_rows, num_kv_heads, num_splits),
@@ -693,6 +884,9 @@ def warmup_qsa_sparse_paged_attention(
                 partial_output_ptr,
                 partial_lse_ptr,
                 output_ptr,
+                output_gate_ptr,
+                row_stride,
+                head_stride,
                 row_stride,
                 head_stride,
                 num_rows,
@@ -714,7 +908,6 @@ def qsa_store_cache_rows(
     rows: torch.Tensor,
 ) -> None:
     """Store fixed-width rows in a QSA cache without boolean indexing."""
-
     if not cache.is_cuda or not HAS_TRITON:
         raise RuntimeError("QSA CUDA cache stores require Triton")
     if cache.ndim != 4 or cache.shape[2] != 1:
@@ -760,7 +953,6 @@ def qsa_compress_groups_with_ratio(
     rope_cache: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Pool completed groups from the compressor-state ring and raw token rows."""
-
     if not raw_keys.is_cuda or not HAS_TRITON:
         raise RuntimeError("QSA CUDA compression requires Triton")
     rows = token_to_req.numel()
diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py
index 9ddd65d15935..077dea91f028 100644
--- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py
+++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py
@@ -317,7 +317,6 @@ def warmup_qsa_mqa_paged_decode(
     max_num_batched_tokens: int,
 ) -> tuple[tuple[int, int], ...]:
     """Compile every reachable decode specialization without launching it."""
-
     page_size = k_cache.shape[1]
     page_table_width = page_table.shape[1]
     columns = page_table_width * page_size
@@ -444,7 +443,6 @@ def expand_qsa_block_indices(
     out: torch.Tensor,
 ) -> None:
     """Expand compressed blocks and compact the causal tail of the open group."""
-
     assert token_topk % compress_ratio == 0
     block_topk = token_topk // compress_ratio
     output_width = token_topk + compress_ratio - 1
@@ -524,8 +522,8 @@ def qsa_select_paged_decode(
         compress_ratio: Number of logical tokens represented by a cache row.
         decode_query_len: Number of query tokens per request.
         block_indices: Compressed-index output buffer.
-    """
 
+    """
     assert token_topk % compress_ratio == 0
     assert block_indices.shape == (q.shape[0], token_topk // compress_ratio)
     assert decode_query_len > 0 and q.shape[0] % decode_query_len == 0
@@ -599,8 +597,8 @@ def qsa_select_paged_prefill(
         max_query_len: Maximum number of query tokens in one request.
         block_indices: Compressed-index output buffer.
         max_seq_len: Longest context length in the batch this step.
-    """
 
+    """
     assert token_topk % compress_ratio == 0
     assert block_indices.shape == (q.shape[0], token_topk // compress_ratio)
     assert q.dtype == k_cache.dtype, "Q and the compressed K cache must match"
diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py
index 26622357ca50..b04cdaad2128 100644
--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py
+++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py
@@ -146,7 +146,6 @@ def _dequantize_embeddings(
         output_dtype: torch.dtype,
     ) -> torch.Tensor:
         """Dequantize PLE lookup output."""
-
         return self.ple_embedding.ngram_embedding.dequantize(
             embeddings,
             output_dtype,
@@ -195,6 +194,7 @@ def _short_conv_dilated_dispatch(
         self,
         inputs: torch.Tensor,
         residual: torch.Tensor,
+        outer_residual: torch.Tensor,
         metadata: PleShortConvAttentionMetadata,
         conv_state: torch.Tensor,
         conv_weights: torch.Tensor,
@@ -209,6 +209,7 @@ def _short_conv_dilated_dispatch(
         has_non_spec = has_prefill or has_decode
         inputs = inputs[: metadata.num_actual_tokens]
         residual = residual[: metadata.num_actual_tokens]
+        outer_residual = outer_residual[: metadata.num_actual_tokens]
 
         spec_token_indices = None
         non_spec_token_indices = None
@@ -235,6 +236,7 @@ def _short_conv_dilated_dispatch(
                 conv_state=conv_state,
                 conv_weights=conv_weights,
                 state_indices=spec_state_indices,
+                outer_residual=outer_residual,
                 mode="spec",
                 dilation=self.short_conv_dilation,
                 query_start_loc=query_start_loc,
@@ -259,11 +261,17 @@ def _short_conv_dilated_dispatch(
                 residual_d, residual_p = torch.split(
                     residual, [num_decode_tokens, num_prefill_tokens], dim=0
                 )
+                outer_residual_d, outer_residual_p = torch.split(
+                    outer_residual,
+                    [num_decode_tokens, num_prefill_tokens],
+                    dim=0,
+                )
                 token_indices_d = None
                 token_indices_p = None
             else:
                 inputs_d = inputs_p = inputs
                 residual_d = residual_p = residual
+                outer_residual_d = outer_residual_p = outer_residual
                 token_indices_d, token_indices_p = torch.split(
                     non_spec_token_indices,
                     [num_decode_tokens, num_prefill_tokens],
@@ -277,6 +285,7 @@ def _short_conv_dilated_dispatch(
                     conv_state=conv_state,
                     conv_weights=conv_weights,
                     state_indices=state_indices_d,
+                    outer_residual=outer_residual_d,
                     mode="decode",
                     dilation=self.short_conv_dilation,
                     has_initial_states=metadata.has_initial_states_d,
@@ -298,6 +307,7 @@ def _short_conv_dilated_dispatch(
                 conv_state=conv_state,
                 conv_weights=conv_weights,
                 state_indices=state_indices_p,
+                outer_residual=outer_residual_p,
                 mode="prefill",
                 dilation=self.short_conv_dilation,
                 query_start_loc=query_start_loc,
@@ -316,6 +326,7 @@ def _short_conv_dilated_dispatch(
                 conv_state=conv_state,
                 conv_weights=conv_weights,
                 state_indices=state_indices[:num_decode_rows],
+                outer_residual=outer_residual,
                 mode="decode",
                 dilation=self.short_conv_dilation,
                 has_initial_states=metadata.has_initial_states_d,
@@ -324,12 +335,18 @@ def _short_conv_dilated_dispatch(
 
     # State routing consumes the current request metadata on every replay.
     @eager_break_during_capture
-    def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None:
+    def _short_conv(
+        self,
+        inputs: torch.Tensor,
+        residual: torch.Tensor,
+        outer_residual: torch.Tensor,
+    ) -> None:
         forward_context = get_forward_context()
         attn_metadata = forward_context.attn_metadata
-        # Profiling omits all metadata or this Mamba entry. The residual
-        # already contains the gated output, so short convolution is a no-op.
+        # Profiling omits all metadata or this Mamba entry. Short convolution
+        # is a no-op there, but preserve the outer residual addition.
         if attn_metadata is None:
+            residual.add_(outer_residual)
             return
 
         if not isinstance(attn_metadata, dict):
@@ -340,6 +357,7 @@ def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None:
 
         layer_attn_metadata = attn_metadata.get(self.prefix)
         if layer_attn_metadata is None:
+            residual.add_(outer_residual)
             return
         if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):
             raise TypeError(
@@ -367,6 +385,7 @@ def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None:
         self._short_conv_dilated_dispatch(
             inputs=inputs,
             residual=residual,
+            outer_residual=outer_residual,
             metadata=layer_attn_metadata,
             conv_state=conv_state,
             conv_weights=conv_weights.to(dtype=inputs.dtype),
@@ -404,7 +423,9 @@ def forward(
             self.norm_conv.weight,
             self.norm_key.eps,
         )
-        self._short_conv(conv_input, gated_output)
+        # State routing depends on runtime request metadata and remains outside
+        # the graph; short convolution accumulates into gated_output.
+        self._short_conv(conv_input, gated_output, hidden_states)
         return gated_output
 
 
diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py
index 90449d5706fb..b6da89a863bd 100644
--- a/vllm/models/qwen4_exp/nvidia/qsa.py
+++ b/vllm/models/qwen4_exp/nvidia/qsa.py
@@ -24,6 +24,7 @@
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope
 from vllm.model_executor.models.qwen3_next import Qwen3NextAttention
 from vllm.platforms import current_platform
+from vllm.platforms.interface import DeviceCapability
 from vllm.transformers_utils.configs.qwen4_exp import (
     Qwen4ExpTextConfig,
 )
@@ -64,14 +65,45 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
     """FullAttentionSpec backend used by the merged QSA owner."""
 
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
-    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
+    # fp8/fp8_e4m3: e4m3 bytes in a uint8 cache, written by reshape_and_cache
+    # with the layer's per-tensor scales and dequantized on load inside the QSA
+    # Triton kernel. flash-attn never runs over this cache, so its fp8 probe
+    # does not apply (see supports_kv_cache_dtype and the impl constructor).
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
+    ]
+
+    @classmethod
+    def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
+        return kv_cache_dtype is None or kv_cache_dtype in cls.supported_kv_cache_dtypes
+
+    @classmethod
+    def supports_combination(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        device_capability: DeviceCapability,
+    ) -> str | None:
+        # QSA dequantizes the fp8 KV in its own Triton kernel and never runs
+        # flash-attn over the quantized cache, so the parent's fp8-KV rejection
+        # does not apply and every combination it is handed is accepted here.
+        return None
 
     @staticmethod
     def get_name() -> str:
         return "QWEN4_EXP_QSA_TRITON"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # QSA consumes manager pages directly and does not use FA4 paged attention.
         return [MultipleOf(16)]
 
@@ -98,16 +130,53 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
     supports_dcp: bool = False
     supports_pcp: bool = False
 
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None = None,
+        attn_type: AttentionType = AttentionType.DECODER,
+        kv_sharing_target_layer_name: str | None = None,
+        sinks: torch.Tensor | None = None,
+    ) -> None:
+        # The parent constructor probes flash-attn for quantized-KV support and
+        # raises where it is unavailable (sm120), but QSA dequantizes fp8 inside
+        # its own Triton kernel and never runs flash-attn over the cache. Hand
+        # the parent "auto" for that probe and restore the real dtype afterwards:
+        # the parent only uses it there, and do_kv_cache_update reads the
+        # attribute at call time.
+        real_kv_cache_dtype = kv_cache_dtype
+        if kv_cache_dtype in ("fp8", "fp8_e4m3"):
+            kv_cache_dtype = "auto"
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            sinks,
+        )
+        self.kv_cache_dtype = real_kv_cache_dtype
         if not is_flash_attn_varlen_func_available():
             raise NotImplementedError("Qwen4Exp QSA requires FlashAttention")
         if self.dcp_world_size != 1:
             raise NotImplementedError(
                 "Qwen4Exp QSA does not support decode context parallelism"
             )
-        if self.kv_cache_dtype not in ("auto", "bfloat16"):
-            raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
+        if self.kv_cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"):
+            raise NotImplementedError(
+                "Qwen4Exp QSA requires a BF16 or FP8-e4m3 main KV cache"
+            )
         self.supports_quant_query_input = False
 
     def forward_qsa(
@@ -121,6 +190,7 @@ def forward_qsa(
         output: torch.Tensor,
         token_to_req: torch.Tensor,
         use_prefill_config: bool,
+        output_gate: torch.Tensor,
         output_scale: torch.Tensor | None = None,
         output_block_scale: torch.Tensor | None = None,
     ) -> torch.Tensor:
@@ -143,8 +213,23 @@ def forward_qsa(
         logical_indices = topk_buffer[:num_tokens]
         token_to_req = token_to_req[:num_tokens]
         key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
-        if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16:
-            raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V")
+        k_scale = v_scale = None
+        if self.kv_cache_dtype in ("fp8", "fp8_e4m3"):
+            # The cache is allocated as uint8; reinterpret the e4m3 bytes
+            # (same itemsize, so shape and strides are preserved).
+            key_cache = key_cache.view(torch.float8_e4m3fn)
+            value_cache = value_cache.view(torch.float8_e4m3fn)
+            # Host-side per-tensor dequant scales (Python floats), as used by
+            # other host-scale backends; folded into the kernel's scales.
+            k_scale = layer._k_scale_float
+            v_scale = layer._v_scale_float
+        if query.dtype != torch.bfloat16 or key_cache.dtype not in (
+            torch.bfloat16,
+            torch.float8_e4m3fn,
+        ):
+            raise NotImplementedError(
+                "Qwen4Exp QSA requires BF16 Q and BF16 or FP8-e4m3 K/V"
+            )
 
         from .ops.qsa import qsa_sparse_paged_attention
 
@@ -157,6 +242,9 @@ def forward_qsa(
             token_to_req,
             use_prefill_config,
             output[:num_tokens],
+            k_scale=k_scale,
+            v_scale=v_scale,
+            output_gate=output_gate[:num_tokens],
         )
         return output
 
@@ -183,8 +271,15 @@ def __init__(
             raise ValueError("Qwen4Exp QSA requires a paged KV cache")
         if model_config.dtype != torch.bfloat16:
             raise NotImplementedError("Qwen4Exp QSA currently requires BF16")
-        if cache_config.cache_dtype not in ("auto", "bfloat16"):
-            raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
+        if cache_config.cache_dtype not in (
+            "auto",
+            "bfloat16",
+            "fp8",
+            "fp8_e4m3",
+        ):
+            raise NotImplementedError(
+                "Qwen4Exp QSA requires a BF16 or FP8-e4m3 main KV cache"
+            )
         if getattr(quant_config, "kv_cache_scheme", None) is not None:
             raise NotImplementedError("Qwen4Exp QSA does not support KV quantization")
         parallel_config = vllm_config.parallel_config
@@ -283,8 +378,10 @@ def __init__(
         self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
             self.kv_cache_dtype, model_config
         )
-        if self.kv_cache_torch_dtype != torch.bfloat16:
-            raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage")
+        if self.kv_cache_torch_dtype not in (torch.bfloat16, torch.uint8):
+            raise NotImplementedError(
+                "Qwen4Exp QSA requires BF16 or FP8-e4m3 (uint8) cache storage"
+            )
         self.kv_sharing_target_layer_name = None
         self.kv_cache = torch.tensor([])
         set_default_quant_scales(self, register_buffer=True)
@@ -353,6 +450,7 @@ def _run_qsa(
         key: torch.Tensor,
         value: torch.Tensor,
         output: torch.Tensor,
+        output_gate: torch.Tensor,
     ) -> None:
         metadata = get_forward_context().attn_metadata
         if isinstance(metadata, list):
@@ -396,6 +494,7 @@ def _run_qsa(
             output,
             token_to_req=side_metadata.token_to_req,
             use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len,
+            output_gate=output_gate,
         )
 
     def forward(
@@ -405,6 +504,7 @@ def forward(
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v, gate = self._project_qkv_gate(qkv, positions)
+        assert gate is not None
         num_tokens = hidden_states.shape[0]
         query = q.view(num_tokens, self.num_heads, self.head_dim)
         key = k.view(num_tokens, self.num_kv_heads, self.head_dim)
@@ -419,10 +519,9 @@ def forward(
             key,
             value,
             attn_output,
+            gate,
         )
         flat_output = attn_output.view(num_tokens, -1)
-        if gate is not None:
-            flat_output = flat_output * torch.sigmoid(gate)
         output, _ = self.o_proj(flat_output)
         return output
 
diff --git a/vllm/multimodal/audio.py b/vllm/multimodal/audio.py
index 59f020230372..be026810bf77 100644
--- a/vllm/multimodal/audio.py
+++ b/vllm/multimodal/audio.py
@@ -45,6 +45,7 @@ def get_audio_duration(*, y: npt.NDArray[np.floating], sr: float = 22050) -> flo
 
     Returns:
         Duration of the audio in seconds.
+
     """
     n_samples = y.shape[-1]
     return float(n_samples) / sr
@@ -71,6 +72,7 @@ class AudioSpec:
             (no normalization). 1 = mono, 2 = stereo, etc.
         channel_reduction: Method to reduce channels when input has more
             channels than target. Only used when reducing channels.
+
     """
 
     target_channels: int | None = 1
@@ -119,6 +121,7 @@ def normalize_audio(
     Raises:
         ValueError: If audio has unsupported dimensions or channel expansion
             is requested (e.g., mono to stereo).
+
     """
     if not spec.needs_normalization:
         return audio
@@ -195,6 +198,7 @@ def resample_audio_pyav(
 
     Returns:
         Resampled audio with the same shape as the input (1D → 1D, 2D → 2D).
+
     """
     orig_sr_int = int(round(orig_sr))
     target_sr_int = int(round(target_sr))
@@ -312,6 +316,7 @@ def resample_audio_torchaudio(
 
     Returns:
         Resampled audio with the same shape as the input (1D → 1D, 2D → 2D).
+
     """
     orig_sr_int = int(round(orig_sr))
     target_sr_int = int(round(target_sr))
@@ -431,6 +436,7 @@ def split_audio(
         ... )
         >>> len(chunks)
         3
+
     """
     if audio_data.ndim > 1:
         raise ValueError(
@@ -500,6 +506,7 @@ def find_split_point(
         ... )
         >>> 16000 <= split_idx <= 17600
         True
+
     """
     segment = wav[start_idx:end_idx]
 
diff --git a/vllm/multimodal/cache/__init__.py b/vllm/multimodal/cache/__init__.py
new file mode 100644
index 000000000000..29411c185922
--- /dev/null
+++ b/vllm/multimodal/cache/__init__.py
@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .base import (
+    BaseMultiModalProcessorCache,
+    BaseMultiModalReceiverCache,
+    MultiModalCacheMissError,
+    MultiModalProcessorOnlyCache,
+)
+from .factories import (
+    engine_receiver_cache_from_config,
+    processor_cache_from_config,
+    processor_only_cache_from_config,
+    worker_receiver_cache_from_config,
+)
+from .lru import (
+    LruKeyReplicatedReceiverCache,
+    LruKeyReplicatedSenderCache,
+)
+from .shm import (
+    ShmObjectStoreReceiverCache,
+    ShmObjectStoreSenderCache,
+)
+
+__all__ = [
+    "MultiModalCacheMissError",
+    "BaseMultiModalProcessorCache",
+    "BaseMultiModalReceiverCache",
+    "MultiModalProcessorOnlyCache",
+    "LruKeyReplicatedSenderCache",
+    "LruKeyReplicatedReceiverCache",
+    "ShmObjectStoreReceiverCache",
+    "ShmObjectStoreSenderCache",
+    "processor_cache_from_config",
+    "processor_only_cache_from_config",
+    "engine_receiver_cache_from_config",
+    "worker_receiver_cache_from_config",
+]
diff --git a/vllm/multimodal/cache.py b/vllm/multimodal/cache/base.py
similarity index 52%
rename from vllm/multimodal/cache.py
rename to vllm/multimodal/cache/base.py
index cc48b58298bb..1b9b757917a3 100644
--- a/vllm/multimodal/cache.py
+++ b/vllm/multimodal/cache/base.py
@@ -4,26 +4,19 @@
 import sys
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
-from multiprocessing.synchronize import Lock as LockType
-from typing import TYPE_CHECKING, Generic, TypeAlias, TypeVar, cast
+from typing import TYPE_CHECKING, Generic, TypeAlias, TypeVar
 
 import torch
 from typing_extensions import override
 
-import vllm.envs as envs
-from vllm.distributed.device_communicators.shm_object_storage import (
-    MsgpackSerde,
-    SingleWriterShmObjectStorage,
-    SingleWriterShmRingBuffer,
-)
+from vllm.config import ModelConfig
 from vllm.logger import init_logger
 from vllm.utils.cache import CacheInfo, LRUCache
 from vllm.utils.jsontree import json_count_leaves, json_map_leaves, json_reduce_leaves
-from vllm.utils.mem_constants import GiB_bytes, MiB_bytes
+from vllm.utils.mem_constants import GiB_bytes
 from vllm.utils.mem_utils import format_gib
 
-from .inputs import (
-    MultiModalBatchedField,
+from ..inputs import (
     MultiModalFeatureSpec,
     MultiModalFieldElem,
     MultiModalKwargsItem,
@@ -32,9 +25,7 @@
 )
 
 if TYPE_CHECKING:
-    from vllm.config import ModelConfig, VllmConfig
-
-    from .processing.processor import ResolvedPromptUpdate
+    from ..processing.processor import ResolvedPromptUpdate
 
 logger = init_logger(__name__)
 
@@ -63,12 +54,12 @@ def __init__(self, mm_hashes: list[str]) -> None:
 
 
 class MultiModalProcessorCacheItem:
-    """
-    The data to store inside `MultiModalProcessorOnlyCache`.
+    """The data to store inside `MultiModalProcessorOnlyCache`.
 
     Args:
         item: The processed tensor data corresponding to a multi-modal item.
         prompt_updates: The prompt updates corresponding to `item`.
+
     """
 
     def __init__(
@@ -83,8 +74,7 @@ def __init__(
 
 
 class MultiModalProcessorCacheItemMetadata:
-    """
-    The metadata to store inside `MultiModalProcessorSenderCache`.
+    """The metadata to store inside `MultiModalProcessorSenderCache`.
 
     Args:
         item: The processed tensor data corresponding to a multi-modal item.
@@ -94,6 +84,7 @@ class MultiModalProcessorCacheItemMetadata:
         prompt_updates: The prompt updates corresponding to `item`.
             This needs to stay on P0 because for some models, they are
             dependent on the processed tensor data (cached on P1).
+
     """
 
     def __init__(
@@ -163,8 +154,7 @@ def get_item_size(
 
     @classmethod
     def get_item_complexity(cls, value: MultiModalCacheValue) -> int:
-        """
-        Get the number of leaf elements in a multi-modal cache value.
+        """Get the number of leaf elements in a multi-modal cache value.
 
         This provides a measure of structural complexity that can be useful
         for debugging cache performance and understanding data patterns.
@@ -174,6 +164,7 @@ def get_item_complexity(cls, value: MultiModalCacheValue) -> int:
 
         Returns:
             The number of leaf elements in the nested structure.
+
         """
         return json_count_leaves(value)
 
@@ -196,8 +187,7 @@ def get_lru_cache(
 
 
 class BaseMultiModalCache(ABC, Generic[_I, _O]):
-    """
-    Abstract base class to read/write multi-modal items from cache.
+    """Abstract base class to read/write multi-modal items from cache.
 
     The idea of multi-modal caching is based on having a client and server
     where the client executes in the frontend process (=P0) and
@@ -226,8 +216,7 @@ def get_and_update_item(
         mm_item: _I,
         mm_hash: str,
     ) -> _O:
-        """
-        Possibly update a multi-modal item based on whether it is
+        """Possibly update a multi-modal item based on whether it is
         in the underlying cache.
 
         This update is done out-of-place and updates the cache eviction order.
@@ -238,6 +227,7 @@ def get_and_update_item(
 
         Returns:
             The update multi-modal item.
+
         """
         raise NotImplementedError
 
@@ -246,8 +236,7 @@ def get_and_update(
         mm_items: Sequence[_I],
         mm_hashes: list[str],
     ) -> list[_O]:
-        """
-        Possibly update a sequence of multi-modal items based on whether they
+        """Possibly update a sequence of multi-modal items based on whether they
         are in the underlying cache.
 
         This update is done out-of-place and updates the cache eviction order.
@@ -258,6 +247,7 @@ def get_and_update(
 
         Returns:
             A new list of updated multi-modal items.
+
         """
         assert len(mm_items) == len(mm_hashes)
 
@@ -289,6 +279,7 @@ def cache_if_fits(
 
         Returns:
             `True` if the item was cached, otherwise `False`.
+
         """
         if cache.put_if_fits(key, value):
             return True
@@ -325,8 +316,7 @@ class BaseMultiModalProcessorCache(
 
     @abstractmethod
     def is_cached_item(self, mm_hash: str) -> bool:
-        """
-        Check whether a multi-modal item is
+        """Check whether a multi-modal item is
         in the underlying cache.
 
         This **DOES NOT** update the cache eviction order.
@@ -336,12 +326,12 @@ def is_cached_item(self, mm_hash: str) -> bool:
 
         Returns:
             `True` if the item is cached, otherwise `False`.
+
         """
         raise NotImplementedError
 
     def is_cached(self, mm_hashes: list[str]) -> list[bool]:
-        """
-        Check whether a sequence of multi-modal items are
+        """Check whether a sequence of multi-modal items are
         in the underlying cache.
 
         This **DOES NOT** update the cache eviction order.
@@ -351,6 +341,7 @@ def is_cached(self, mm_hashes: list[str]) -> list[bool]:
 
         Returns:
             For each item, `True` if the item is cached, otherwise `False`.
+
         """
         return [self.is_cached_item(mm_hash) for mm_hash in mm_hashes]
 
@@ -367,31 +358,30 @@ def close(self) -> None:
 
     @abstractmethod
     def touch_sender_cache_item(self, mm_hash: str) -> None:
-        """
-        Update the cache eviction order for a multi-modal item.
+        """Update the cache eviction order for a multi-modal item.
 
         This is used to touch the item in the cache without changing
         its value.
 
         Args:
             mm_hash: The hash of the multi-modal item.
+
         """
         raise NotImplementedError
 
     @abstractmethod
     def make_stats(self, *, delta: bool = False) -> CacheInfo:
-        """
-        Get (and reset) the multi-modal cache stats.
+        """Get (and reset) the multi-modal cache stats.
 
         Returns:
             The current multi-modal caching stats.
+
         """
         raise NotImplementedError
 
 
 class MultiModalProcessorOnlyCache(BaseMultiModalProcessorCache):
-    """
-    The cache which is used on P0 when IPC caching is disabled.
+    """The cache which is used on P0 when IPC caching is disabled.
 
     How to update each item:
 
@@ -400,7 +390,7 @@ class MultiModalProcessorOnlyCache(BaseMultiModalProcessorCache):
       tensor data and metadata) into the cache, and return the input.
     """
 
-    def __init__(self, model_config: "ModelConfig") -> None:
+    def __init__(self, model_config: ModelConfig) -> None:
         super().__init__()
 
         mm_config = model_config.get_multimodal_config()
@@ -441,218 +431,6 @@ def make_stats(self, *, delta: bool = False) -> CacheInfo:
         return self._cache.stat(delta=delta)
 
 
-class MultiModalProcessorSenderCache(BaseMultiModalProcessorCache):
-    """
-    The cache which is used on P0 when IPC caching is enabled.
-
-    How to update each item:
-
-    - If the item is already in the cache, clear the input to avoid
-      unnecessary IPC.
-
-    - If the item is not in the cache, store the metadata of that item so
-      that the eviction policy remains the same as the cache on P1,
-      and return the input.
-      By only storing the metadata, we avoid keeping the data itself in
-      memory inside P0.
-    """
-
-    def __init__(self, model_config: "ModelConfig") -> None:
-        super().__init__()
-
-        mm_config = model_config.get_multimodal_config()
-
-        self._cache = MultiModalCache.get_lru_cache(
-            mm_config.mm_processor_cache_gb,
-            MultiModalProcessorCacheItemMetadata,
-        )
-
-    @override
-    def is_cached_item(self, mm_hash: str) -> bool:
-        return mm_hash in self._cache
-
-    @override
-    def get_and_update_item(
-        self,
-        mm_item: MultiModalProcessorCacheInItem,
-        mm_hash: str,
-    ) -> MultiModalProcessorCacheOutItem:
-        if (cached_item := self._cache.get(mm_hash)) is not None:
-            return None, cached_item.prompt_updates
-
-        assert mm_item is not None, f"Expected a cached item for {mm_hash=}"
-
-        self.cache_if_fits(
-            self._cache, mm_hash, MultiModalProcessorCacheItemMetadata(*mm_item)
-        )
-        return mm_item
-
-    @override
-    def touch_sender_cache_item(self, mm_hash: str) -> None:
-        self._cache.touch(mm_hash)
-
-    @override
-    def clear_cache(self) -> None:
-        self._cache.clear()
-
-    @override
-    def make_stats(self, *, delta: bool = False) -> CacheInfo:
-        return self._cache.stat(delta=delta)
-
-    @override
-    def invalidate(self, mm_hash: str) -> None:
-        # Drop our stale shadow entry so the next request for this hash re-sends the
-        # data and repopulates P1 (see MultiModalCacheMissError).
-        self._cache.pop(mm_hash, None)
-
-
-class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
-    """
-    The cache which is used on P0 when IPC caching is enabled.
-
-    How to update each item:
-
-    - If the item is already in the cache, clear the input to avoid
-      unnecessary IPC.
-
-    - If the item is not in the cache, store the data in shared memory.
-    """
-
-    def __init__(self, vllm_config: "VllmConfig") -> None:
-        super().__init__()
-
-        self.world_size = vllm_config.parallel_config.world_size
-        mm_config = vllm_config.model_config.get_multimodal_config()
-
-        ring_buffer = SingleWriterShmRingBuffer(
-            data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes),
-            name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME,
-            create=True,  # sender is the writer
-        )
-        self._shm_cache = SingleWriterShmObjectStorage(
-            max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes,
-            n_readers=self.world_size,
-            ring_buffer=ring_buffer,
-            serde_class=MsgpackSerde,
-        )
-        # cache prompt_updates for P0 only
-        self._p0_cache: dict[str, Sequence[ResolvedPromptUpdate]] = {}
-
-        self._hits = 0
-        self._total = 0
-        self._last_info = CacheInfo(hits=0, total=0)
-
-    def _stat(self, *, delta: bool = False) -> CacheInfo:
-        info = CacheInfo(hits=self._hits, total=self._total)
-
-        if delta:
-            info_delta = info - self._last_info
-            self._last_info = info
-            info = info_delta
-
-        return info
-
-    @override
-    def is_cached_item(self, mm_hash: str) -> bool:
-        return self._shm_cache.is_cached(mm_hash)
-
-    @override
-    def get_and_update_item(
-        self,
-        mm_item: MultiModalProcessorCacheInItem,
-        mm_hash: str,
-    ) -> MultiModalProcessorCacheOutItem:
-        if self._shm_cache.is_cached(mm_hash):
-            self._hits += 1
-            self._total += 1
-
-            address, monotonic_id = self._shm_cache.get_cached(mm_hash)
-            prompt_updates = self._p0_cache[mm_hash]
-            return self.address_as_item(address, monotonic_id), prompt_updates
-
-        assert mm_item is not None, f"Expected a cached item for {mm_hash=}"
-        item, prompt_updates = mm_item
-
-        self._total += 1
-
-        try:
-            address, monotonic_id = self._shm_cache.put(mm_hash, item)
-            # Try to remove dangling items if p0 cache is too large.
-            if len(self._p0_cache) >= 2 * len(self._shm_cache.key_index):
-                self.remove_dangling_items()
-
-            self._p0_cache[mm_hash] = prompt_updates
-            return self.address_as_item(address, monotonic_id), prompt_updates
-        except ValueError as e:
-            # `put` raises ValueError either for an oversize item or for a
-            # duplicate key (concurrent insert); the latter is benign so we
-            # only warn on the oversize case. Subsequent UUID-only requests
-            # for an oversize item will fail with a cache miss.
-            if "already exists" not in str(e):
-                logger.warning_once(
-                    "mm_input %s too large to cache; "
-                    "raise --mm-shm-cache-max-object-size-mb. (%s)",
-                    mm_hash,
-                    str(e),
-                )
-            return mm_item
-        except MemoryError as e:
-            # Cache full and protected items prevent eviction.
-            logger.debug(
-                "mm_input %s not cached; shm cache full, "
-                "consider raising --mm-processor-cache-gb. (%s)",
-                mm_hash,
-                str(e),
-            )
-            return mm_item
-
-    @override
-    def touch_sender_cache_item(self, mm_hash: str) -> None:
-        """Touch the item in shared memory cache to prevent eviction.
-        Increments writer_flag on sender side."""
-        self._shm_cache.touch(mm_hash)
-
-    @override
-    def clear_cache(self) -> None:
-        self._shm_cache.clear()
-        self._p0_cache.clear()
-
-        self._hits = 0
-        self._total = 0
-        self._last_info = CacheInfo(hits=0, total=0)
-
-    @override
-    def make_stats(self, *, delta: bool = False) -> CacheInfo:
-        return self._stat(delta=delta)
-
-    @override
-    def close(self) -> None:
-        self._shm_cache.close()
-
-    def remove_dangling_items(self) -> None:
-        """Remove items that are no longer in the shared memory cache."""
-        cached_hashes = self._shm_cache.key_index.keys()
-        dangling_hashes = set(self._p0_cache.keys()) - cached_hashes
-        for mm_hash in dangling_hashes:
-            del self._p0_cache[mm_hash]
-
-    def address_as_item(
-        self,
-        address: int,
-        monotonic_id: int,
-    ) -> MultiModalKwargsItem:
-        addr_elem = MultiModalFieldElem(
-            data=address,
-            field=MultiModalBatchedField(),
-        )
-        id_elem = MultiModalFieldElem(
-            data=monotonic_id,
-            field=MultiModalBatchedField(),
-        )
-
-        return MultiModalKwargsItem({"address": addr_elem, "monotonic_id": id_elem})
-
-
 class BaseMultiModalReceiverCache(
     BaseMultiModalCache[MultiModalKwargsItem | None, MultiModalKwargsItem]
 ):
@@ -662,8 +440,7 @@ def get_and_update_features(
         self,
         mm_features: list["MultiModalFeatureSpec"],
     ) -> list["MultiModalFeatureSpec"]:
-        """
-        Update multimodal features with cached encoder outputs.
+        """Update multimodal features with cached encoder outputs.
         Touch all identifier at first before update to avoid
         item in updated list evict during update.
 
@@ -695,8 +472,7 @@ def touch_receiver_cache_item(
         mm_hash: str,
         mm_item: MultiModalKwargsItem | None = None,
     ) -> None:
-        """
-        Update the cache eviction order for a multi-modal item.
+        """Update the cache eviction order for a multi-modal item.
 
         This is used to touch the item in the cache without changing
         its value.
@@ -705,135 +481,6 @@ def touch_receiver_cache_item(
             mm_hash: The hash of the multi-modal item.
             mm_item: The multi-modal item itself. This is optional and
                 may not be needed by some cache implementations.
+
         """
         raise NotImplementedError
-
-
-class MultiModalReceiverCache(BaseMultiModalReceiverCache):
-    """
-    The cache which is used on P1 when IPC caching is enabled.
-
-    How to update each item:
-
-    - If the item is in the cache, replace the input with the cached item.
-    - If the item is not in the cache, store that item (which includes tensor
-      data) into the cache, and return the input.
-    """
-
-    def __init__(self, model_config: "ModelConfig") -> None:
-        super().__init__()
-
-        mm_config = model_config.get_multimodal_config()
-
-        self._cache = MultiModalCache.get_lru_cache(
-            mm_config.mm_processor_cache_gb,
-            MultiModalKwargsItem,
-        )
-
-    @override
-    def get_and_update_item(
-        self,
-        mm_item: MultiModalKwargsItem | None,
-        mm_hash: str,
-    ) -> MultiModalKwargsItem:
-        if (cached_item := self._cache.get(mm_hash)) is not None:
-            return cached_item
-
-        # No data and not cached here: P0 sent data=None trusting its shadow, but
-        # the P0/P1 caches have drifted. Raise a retryable error (not assert) so the
-        # engine can have P0 drop the stale entry and the client resend the data.
-        if mm_item is None:
-            raise MultiModalCacheMissError([mm_hash])
-
-        self.cache_if_fits(self._cache, mm_hash, mm_item)
-        return mm_item
-
-    @override
-    def touch_receiver_cache_item(
-        self,
-        mm_hash: str,
-        mm_item: MultiModalKwargsItem | None = None,
-    ) -> None:
-        self._cache.touch(mm_hash)
-
-    @override
-    def clear_cache(self) -> None:
-        self._cache.clear()
-
-
-class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):
-    """
-    The cache which is used on P1 Worker Process when IPC caching is enabled.
-
-    How to update each item:
-
-    - If the item has an address, replace the input with the cached item.
-    - If not, return the input.
-    """
-
-    def __init__(
-        self,
-        vllm_config: "VllmConfig",
-        shared_worker_lock: LockType,
-    ) -> None:
-        super().__init__()
-
-        self.world_size = vllm_config.parallel_config.world_size
-        mm_config = vllm_config.model_config.get_multimodal_config()
-
-        ring_buffer = SingleWriterShmRingBuffer(
-            data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes),
-            name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME,
-            create=False,  # Server is a reader
-        )
-        self._shm_cache = SingleWriterShmObjectStorage(
-            max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes,
-            n_readers=self.world_size,
-            ring_buffer=ring_buffer,
-            serde_class=MsgpackSerde,
-            reader_lock=shared_worker_lock,
-        )
-
-    @override
-    def get_and_update_features(
-        self, mm_features: list["MultiModalFeatureSpec"]
-    ) -> list["MultiModalFeatureSpec"]:
-        # strip_covered_mm_data preserves address items, so None here represents
-        # a stripped uncached payload and has no SHM reference to acknowledge.
-        features_with_data = [
-            feature for feature in mm_features if feature.data is not None
-        ]
-        super().get_and_update_features(features_with_data)
-        return mm_features
-
-    @override
-    def get_and_update_item(
-        self,
-        mm_item: MultiModalKwargsItem | None,
-        mm_hash: str,
-    ) -> MultiModalKwargsItem:
-        assert mm_item is not None, f"Expected an address item for {mm_hash=}"
-        if "address" in mm_item:
-            address = cast(int, mm_item["address"].data)
-            monotonic_id = cast(int, mm_item["monotonic_id"].data)
-            return self._shm_cache.get(address, monotonic_id)
-
-        return mm_item
-
-    @override
-    def touch_receiver_cache_item(
-        self,
-        mm_hash: str,
-        mm_item: MultiModalKwargsItem | None = None,
-    ) -> None:
-        """Touch the item in shared memory cache to prevent eviction.
-        Increments reader_count on receiver side."""
-        assert mm_item is not None
-        if "address" in mm_item:
-            address = cast(int, mm_item["address"].data)
-            monotonic_id = cast(int, mm_item["monotonic_id"].data)
-            self._shm_cache.touch(mm_hash, address=address, monotonic_id=monotonic_id)
-
-    @override
-    def clear_cache(self) -> None:
-        self._shm_cache.clear()
diff --git a/vllm/multimodal/cache/factories.py b/vllm/multimodal/cache/factories.py
new file mode 100644
index 000000000000..aa4f5a4d2cb5
--- /dev/null
+++ b/vllm/multimodal/cache/factories.py
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from multiprocessing.synchronize import Lock as LockType
+from typing import Literal
+
+from vllm.config import VllmConfig
+
+from .base import (
+    BaseMultiModalProcessorCache,
+    BaseMultiModalReceiverCache,
+    MultiModalProcessorOnlyCache,
+)
+from .lru import (
+    LruKeyReplicatedReceiverCache,
+    LruKeyReplicatedSenderCache,
+)
+from .shm import (
+    ShmObjectStoreReceiverCache,
+    ShmObjectStoreSenderCache,
+)
+
+
+def _get_cache_type(
+    vllm_config: VllmConfig,
+) -> Literal[None, "processor_only", "lru", "shm"]:
+    model_config = vllm_config.model_config
+    if not model_config.supports_multimodal_inputs:
+        return None
+
+    # Check if the cache is disabled.
+    mm_config = model_config.get_multimodal_config()
+    if mm_config.mm_processor_cache_gb <= 0:
+        return None
+
+    # Check if IPC caching is supported.
+    parallel_config = vllm_config.parallel_config
+    is_ipc_supported = parallel_config._api_process_count == 1 and (
+        parallel_config.data_parallel_size == 1
+        or parallel_config.data_parallel_external_lb
+    )
+
+    if not is_ipc_supported:
+        return "processor_only"
+
+    mm_config = model_config.get_multimodal_config()
+    return mm_config.mm_processor_cache_type
+
+
+def processor_cache_from_config(
+    vllm_config: VllmConfig,
+) -> BaseMultiModalProcessorCache | None:
+    """Return a `BaseMultiModalProcessorCache`, if enabled."""
+    cache_type = _get_cache_type(vllm_config)
+    if cache_type is None:
+        return None
+    elif cache_type == "processor_only":
+        return MultiModalProcessorOnlyCache(vllm_config.model_config)
+    elif cache_type == "lru":
+        return LruKeyReplicatedSenderCache(vllm_config.model_config)
+    elif cache_type == "shm":
+        return ShmObjectStoreSenderCache(vllm_config)
+    else:
+        raise ValueError(f"Unknown cache type: {cache_type!r}")
+
+
+def processor_only_cache_from_config(
+    vllm_config: VllmConfig,
+) -> MultiModalProcessorOnlyCache | None:
+    """Return a `MultiModalProcessorOnlyCache`, if enabled."""
+    cache_type = _get_cache_type(vllm_config)
+    if cache_type is None:
+        return None
+
+    return MultiModalProcessorOnlyCache(vllm_config.model_config)
+
+
+def engine_receiver_cache_from_config(
+    vllm_config: VllmConfig,
+) -> BaseMultiModalReceiverCache | None:
+    """Return a `BaseMultiModalReceiverCache` for the engine process."""
+    cache_type = _get_cache_type(vllm_config)
+    if cache_type in (None, "processor_only", "shm"):
+        return None
+    elif cache_type == "lru":
+        return LruKeyReplicatedReceiverCache(vllm_config.model_config)
+    else:
+        raise ValueError(f"Unknown cache type: {cache_type!r}")
+
+
+def worker_receiver_cache_from_config(
+    vllm_config: VllmConfig,
+    shared_worker_lock: LockType | None,
+) -> BaseMultiModalReceiverCache | None:
+    """Return a `BaseMultiModalReceiverCache` for the worker process."""
+    cache_type = _get_cache_type(vllm_config)
+    if cache_type in (None, "processor_only", "lru"):
+        return None
+    elif cache_type == "shm":
+        if shared_worker_lock is None:
+            raise ValueError(
+                "Missing `shared_worker_lock` argument from executor. "
+                "This argument is needed for mm_processor_cache_type='shm'."
+            )
+
+        return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock)
+    else:
+        raise ValueError(f"Unknown cache type: {cache_type!r}")
diff --git a/vllm/multimodal/cache/lru.py b/vllm/multimodal/cache/lru.py
new file mode 100644
index 000000000000..6b5ec3516a57
--- /dev/null
+++ b/vllm/multimodal/cache/lru.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Implementation of Key-Replicated Cache (see docs/configuration/optimization.md)."""
+
+from typing_extensions import override
+
+from vllm.config import ModelConfig
+from vllm.utils.cache import CacheInfo
+
+from ..inputs import MultiModalKwargsItem
+from .base import (
+    BaseMultiModalProcessorCache,
+    BaseMultiModalReceiverCache,
+    MultiModalCache,
+    MultiModalCacheMissError,
+    MultiModalProcessorCacheInItem,
+    MultiModalProcessorCacheItemMetadata,
+    MultiModalProcessorCacheOutItem,
+)
+
+
+class LruKeyReplicatedSenderCache(BaseMultiModalProcessorCache):
+    """The cache which is used on P0 when LRU caching is enabled.
+
+    How to update each item:
+
+    - If the item is already in the cache, clear the input to avoid
+      unnecessary IPC.
+
+    - If the item is not in the cache, store the metadata of that item so
+      that the eviction policy remains the same as the cache on P1,
+      and return the input.
+      By only storing the metadata, we avoid keeping the data itself in
+      memory inside P0.
+    """
+
+    def __init__(self, model_config: ModelConfig) -> None:
+        super().__init__()
+
+        mm_config = model_config.get_multimodal_config()
+
+        self._cache = MultiModalCache.get_lru_cache(
+            mm_config.mm_processor_cache_gb,
+            MultiModalProcessorCacheItemMetadata,
+        )
+
+    @override
+    def is_cached_item(self, mm_hash: str) -> bool:
+        return mm_hash in self._cache
+
+    @override
+    def get_and_update_item(
+        self,
+        mm_item: MultiModalProcessorCacheInItem,
+        mm_hash: str,
+    ) -> MultiModalProcessorCacheOutItem:
+        if (cached_item := self._cache.get(mm_hash)) is not None:
+            return None, cached_item.prompt_updates
+
+        assert mm_item is not None, f"Expected a cached item for {mm_hash=}"
+
+        self.cache_if_fits(
+            self._cache, mm_hash, MultiModalProcessorCacheItemMetadata(*mm_item)
+        )
+        return mm_item
+
+    @override
+    def touch_sender_cache_item(self, mm_hash: str) -> None:
+        self._cache.touch(mm_hash)
+
+    @override
+    def clear_cache(self) -> None:
+        self._cache.clear()
+
+    @override
+    def make_stats(self, *, delta: bool = False) -> CacheInfo:
+        return self._cache.stat(delta=delta)
+
+    @override
+    def invalidate(self, mm_hash: str) -> None:
+        # Drop our stale shadow entry so the next request for this hash re-sends the
+        # data and repopulates P1 (see MultiModalCacheMissError).
+        self._cache.pop(mm_hash, None)
+
+
+class LruKeyReplicatedReceiverCache(BaseMultiModalReceiverCache):
+    """The cache which is used on P1 when LRU caching is enabled.
+
+    How to update each item:
+
+    - If the caller sent tensor data, store it (replacing any cached item
+      under the same key) and return that data. P0 can miss after independent
+      LRU eviction and resend a different item for the same identity.
+    - If the caller sent no data and the item is cached, return the cached item.
+    - If the caller sent no data and the item is not cached, raise
+      `MultiModalCacheMissError`.
+    """
+
+    def __init__(self, model_config: ModelConfig) -> None:
+        super().__init__()
+
+        mm_config = model_config.get_multimodal_config()
+
+        self._cache = MultiModalCache.get_lru_cache(
+            mm_config.mm_processor_cache_gb,
+            MultiModalKwargsItem,
+        )
+
+    @override
+    def get_and_update_item(
+        self,
+        mm_item: MultiModalKwargsItem | None,
+        mm_hash: str,
+    ) -> MultiModalKwargsItem:
+        if mm_item is not None:
+            # P0 sent a payload. Never keep a stale cached tensor under this
+            # identity: a later P0 hit would pair new placeholders with the old
+            # item and crash the engine. Drop first so a too-large replacement
+            # cannot leave the previous entry in place.
+            self._cache.pop(mm_hash, None)
+            self.cache_if_fits(self._cache, mm_hash, mm_item)
+            return mm_item
+
+        if (cached_item := self._cache.get(mm_hash)) is not None:
+            return cached_item
+
+        # No data and not cached here: P0 sent data=None trusting its shadow, but
+        # the P0/P1 caches have drifted. Raise a retryable error (not assert) so the
+        # engine can have P0 drop the stale entry and the client resend the data.
+        raise MultiModalCacheMissError([mm_hash])
+
+    @override
+    def touch_receiver_cache_item(
+        self,
+        mm_hash: str,
+        mm_item: MultiModalKwargsItem | None = None,
+    ) -> None:
+        self._cache.touch(mm_hash)
+
+    @override
+    def clear_cache(self) -> None:
+        self._cache.clear()
diff --git a/vllm/multimodal/cache/shm.py b/vllm/multimodal/cache/shm.py
new file mode 100644
index 000000000000..dddf3494cf5d
--- /dev/null
+++ b/vllm/multimodal/cache/shm.py
@@ -0,0 +1,261 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Implementation of Shared Memory Cache (see docs/configuration/optimization.md)."""
+
+from collections.abc import Sequence
+from multiprocessing.synchronize import Lock as LockType
+from typing import TYPE_CHECKING, cast
+
+from typing_extensions import override
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.distributed.device_communicators.shm_object_storage import (
+    MsgpackSerde,
+    SingleWriterShmObjectStorage,
+    SingleWriterShmRingBuffer,
+)
+from vllm.logger import init_logger
+from vllm.utils.cache import CacheInfo
+from vllm.utils.mem_constants import GiB_bytes, MiB_bytes
+
+from ..inputs import (
+    MultiModalBatchedField,
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+)
+from .base import (
+    BaseMultiModalProcessorCache,
+    BaseMultiModalReceiverCache,
+    MultiModalProcessorCacheInItem,
+    MultiModalProcessorCacheOutItem,
+)
+
+if TYPE_CHECKING:
+    from ..processing.processor import ResolvedPromptUpdate
+
+logger = init_logger(__name__)
+
+
+class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
+    """The cache which is used on P0 when SHM caching is enabled.
+
+    How to update each item:
+
+    - If the item is already in the cache, clear the input to avoid
+      unnecessary IPC.
+
+    - If the item is not in the cache, store the data in shared memory.
+    """
+
+    def __init__(self, vllm_config: VllmConfig) -> None:
+        super().__init__()
+
+        self.world_size = vllm_config.parallel_config.world_size
+        mm_config = vllm_config.model_config.get_multimodal_config()
+
+        ring_buffer = SingleWriterShmRingBuffer(
+            data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes),
+            name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME,
+            create=True,  # sender is the writer
+        )
+        self._shm_cache = SingleWriterShmObjectStorage(
+            max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes,
+            n_readers=self.world_size,
+            ring_buffer=ring_buffer,
+            serde_class=MsgpackSerde,
+        )
+        # cache prompt_updates for P0 only
+        self._p0_cache: dict[str, Sequence[ResolvedPromptUpdate]] = {}
+
+        self._hits = 0
+        self._total = 0
+        self._last_info = CacheInfo(hits=0, total=0)
+
+    def _stat(self, *, delta: bool = False) -> CacheInfo:
+        info = CacheInfo(hits=self._hits, total=self._total)
+
+        if delta:
+            info_delta = info - self._last_info
+            self._last_info = info
+            info = info_delta
+
+        return info
+
+    @override
+    def is_cached_item(self, mm_hash: str) -> bool:
+        return self._shm_cache.is_cached(mm_hash)
+
+    @override
+    def get_and_update_item(
+        self,
+        mm_item: MultiModalProcessorCacheInItem,
+        mm_hash: str,
+    ) -> MultiModalProcessorCacheOutItem:
+        if self._shm_cache.is_cached(mm_hash):
+            self._hits += 1
+            self._total += 1
+
+            address, monotonic_id = self._shm_cache.get_cached(mm_hash)
+            prompt_updates = self._p0_cache[mm_hash]
+            return self.address_as_item(address, monotonic_id), prompt_updates
+
+        assert mm_item is not None, f"Expected a cached item for {mm_hash=}"
+        item, prompt_updates = mm_item
+
+        self._total += 1
+
+        try:
+            address, monotonic_id = self._shm_cache.put(mm_hash, item)
+            # Try to remove dangling items if p0 cache is too large.
+            if len(self._p0_cache) >= 2 * len(self._shm_cache.key_index):
+                self.remove_dangling_items()
+
+            self._p0_cache[mm_hash] = prompt_updates
+            return self.address_as_item(address, monotonic_id), prompt_updates
+        except ValueError as e:
+            # `put` raises ValueError either for an oversize item or for a
+            # duplicate key (concurrent insert); the latter is benign so we
+            # only warn on the oversize case. Subsequent UUID-only requests
+            # for an oversize item will fail with a cache miss.
+            if "already exists" not in str(e):
+                logger.warning_once(
+                    "mm_input %s too large to cache; "
+                    "raise --mm-shm-cache-max-object-size-mb. (%s)",
+                    mm_hash,
+                    str(e),
+                )
+            return mm_item
+        except MemoryError as e:
+            # Cache full and protected items prevent eviction.
+            logger.debug(
+                "mm_input %s not cached; shm cache full, "
+                "consider raising --mm-processor-cache-gb. (%s)",
+                mm_hash,
+                str(e),
+            )
+            return mm_item
+
+    @override
+    def touch_sender_cache_item(self, mm_hash: str) -> None:
+        """Touch the item in shared memory cache to prevent eviction.
+        Increments writer_flag on sender side."""
+        self._shm_cache.touch(mm_hash)
+
+    @override
+    def clear_cache(self) -> None:
+        self._shm_cache.clear()
+        self._p0_cache.clear()
+
+        self._hits = 0
+        self._total = 0
+        self._last_info = CacheInfo(hits=0, total=0)
+
+    @override
+    def make_stats(self, *, delta: bool = False) -> CacheInfo:
+        return self._stat(delta=delta)
+
+    @override
+    def close(self) -> None:
+        self._shm_cache.close()
+
+    def remove_dangling_items(self) -> None:
+        """Remove items that are no longer in the shared memory cache."""
+        cached_hashes = self._shm_cache.key_index.keys()
+        dangling_hashes = set(self._p0_cache.keys()) - cached_hashes
+        for mm_hash in dangling_hashes:
+            del self._p0_cache[mm_hash]
+
+    def address_as_item(
+        self,
+        address: int,
+        monotonic_id: int,
+    ) -> MultiModalKwargsItem:
+        addr_elem = MultiModalFieldElem(
+            data=address,
+            field=MultiModalBatchedField(),
+        )
+        id_elem = MultiModalFieldElem(
+            data=monotonic_id,
+            field=MultiModalBatchedField(),
+        )
+
+        return MultiModalKwargsItem({"address": addr_elem, "monotonic_id": id_elem})
+
+
+class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):
+    """The cache which is used on P1 Worker Process when SHM caching is enabled.
+
+    How to update each item:
+
+    - If the item has an address, replace the input with the cached item.
+    - If not, return the input.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        shared_worker_lock: LockType,
+    ) -> None:
+        super().__init__()
+
+        self.world_size = vllm_config.parallel_config.world_size
+        mm_config = vllm_config.model_config.get_multimodal_config()
+
+        ring_buffer = SingleWriterShmRingBuffer(
+            data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes),
+            name=envs.VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME,
+            create=False,  # Server is a reader
+        )
+        self._shm_cache = SingleWriterShmObjectStorage(
+            max_object_size=mm_config.mm_shm_cache_max_object_size_mb * MiB_bytes,
+            n_readers=self.world_size,
+            ring_buffer=ring_buffer,
+            serde_class=MsgpackSerde,
+            reader_lock=shared_worker_lock,
+        )
+
+    @override
+    def get_and_update_features(
+        self, mm_features: list["MultiModalFeatureSpec"]
+    ) -> list["MultiModalFeatureSpec"]:
+        # strip_covered_mm_data preserves address items, so None here represents
+        # a stripped uncached payload and has no SHM reference to acknowledge.
+        features_with_data = [
+            feature for feature in mm_features if feature.data is not None
+        ]
+        super().get_and_update_features(features_with_data)
+        return mm_features
+
+    @override
+    def get_and_update_item(
+        self,
+        mm_item: MultiModalKwargsItem | None,
+        mm_hash: str,
+    ) -> MultiModalKwargsItem:
+        assert mm_item is not None, f"Expected an address item for {mm_hash=}"
+        if "address" in mm_item:
+            address = cast(int, mm_item["address"].data)
+            monotonic_id = cast(int, mm_item["monotonic_id"].data)
+            return self._shm_cache.get(address, monotonic_id)
+
+        return mm_item
+
+    @override
+    def touch_receiver_cache_item(
+        self,
+        mm_hash: str,
+        mm_item: MultiModalKwargsItem | None = None,
+    ) -> None:
+        """Touch the item in shared memory cache to prevent eviction.
+        Increments reader_count on receiver side."""
+        assert mm_item is not None
+        if "address" in mm_item:
+            address = cast(int, mm_item["address"].data)
+            monotonic_id = cast(int, mm_item["monotonic_id"].data)
+            self._shm_cache.touch(mm_hash, address=address, monotonic_id=monotonic_id)
+
+    @override
+    def clear_cache(self) -> None:
+        self._shm_cache.clear()
diff --git a/vllm/multimodal/encoder_budget.py b/vllm/multimodal/encoder_budget.py
index 3042091611f9..32db904bfa14 100644
--- a/vllm/multimodal/encoder_budget.py
+++ b/vllm/multimodal/encoder_budget.py
@@ -2,8 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Mapping
 
-from vllm.config import ModelConfig, VllmConfig
+from vllm.config import VllmConfig
 from vllm.logger import init_logger
+from vllm.multimodal.cache import (
+    BaseMultiModalProcessorCache,
+    processor_only_cache_from_config,
+)
 from vllm.multimodal.inputs import MultiModalKwargsItem
 from vllm.multimodal.processing import BaseMultiModalProcessor
 from vllm.multimodal.registry import MultiModalRegistry
@@ -14,27 +18,21 @@
 
 
 def get_mm_max_toks_per_item(
-    model_config: ModelConfig,
-    mm_registry: MultiModalRegistry,
     processor: BaseMultiModalProcessor,
     mm_counts: Mapping[str, int],
+    cache: BaseMultiModalProcessorCache | None,
 ) -> Mapping[str, int]:
-    """
-    Get the maximum number of tokens per data item from each modality based
+    """Get the maximum number of tokens per data item from each modality based
     on underlying model configuration.
     """
     max_tokens_per_item = processor.info.get_mm_max_tokens_per_item(
-        seq_len=model_config.max_model_len,
+        seq_len=processor.info.ctx.model_config.max_model_len,
         mm_counts=mm_counts,
     )
     if max_tokens_per_item is not None:
         return max_tokens_per_item
 
-    mm_inputs = mm_registry.get_dummy_mm_inputs(
-        model_config,
-        mm_counts=mm_counts,
-        processor=processor,
-    )
+    mm_inputs = processor.get_dummy_mm_inputs(mm_counts, cache=cache)
 
     return {
         modality: sum(item.get_num_embeds() for item in placeholders)
@@ -62,11 +60,9 @@ def __init__(
 
         with set_default_torch_num_threads():  # Avoid hang during startup
             cache = (
-                mm_registry.processor_only_cache_from_config(vllm_config)
-                if enable_cache
-                else None
+                processor_only_cache_from_config(vllm_config) if enable_cache else None
             )
-            processor = mm_registry.create_processor(model_config, cache=cache)
+            processor = mm_registry.create_processor(model_config)
 
             self.cache = cache
             self.processor = processor
@@ -92,9 +88,8 @@ def __init__(
             active_modalities = tower_modalities | embed_only_modalities
 
             all_mm_max_toks_per_item = get_mm_max_toks_per_item(
-                model_config,
-                mm_registry,
                 processor,
+                cache=cache,
                 mm_counts=dict.fromkeys(active_modalities, 1),
             )
 
@@ -195,26 +190,32 @@ def get_modality_with_max_tokens(self) -> str:
     def get_encoder_budget(self) -> int:
         return min(self.encoder_compute_budget, self.encoder_cache_size)
 
-    def reset_cache(self) -> None:
-        if self.cache is not None:
-            self.cache.clear_cache()
-
+    def get_dummy_encoder_profile_inputs(
+        self,
+        modality: str | None = None,
+        max_items_per_batch: int | None = None,
+    ) -> list[tuple[str, MultiModalKwargsItem]]:
+        if self.get_encoder_budget() <= 0 or not self.mm_max_toks_per_item:
+            return []
+
+        if modality is None:
+            modality = self.get_modality_with_max_tokens()
+        if max_items_per_batch is None:
+            max_items_per_batch = self.mm_max_items_per_batch[modality]
+
+        # Don't use `max_items_per_batch` here to avoid redundant computation
+        dummy_mm_inputs = self.processor.get_dummy_mm_inputs(
+            mm_counts={modality: 1},
+            cache=self.cache,
+        )
+        dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0]
 
-def get_dummy_encoder_profile_inputs(
-    mm_registry: MultiModalRegistry,
-    budget: MultiModalBudget,
-) -> list[tuple[str, MultiModalKwargsItem]]:
-    if budget.get_encoder_budget() <= 0 or not budget.mm_max_toks_per_item:
-        return []
+        # We use the cache so that the item is saved to the cache,
+        # but not read from the cache
+        assert dummy_mm_item is not None, "Item should not already be cached"
 
-    modality = budget.get_modality_with_max_tokens()
-    max_items_per_batch = budget.mm_max_items_per_batch[modality]
-    dummy_mm_inputs = mm_registry.get_dummy_mm_inputs(
-        budget.model_config,
-        mm_counts={modality: 1},
-        processor=budget.processor,
-    )
-    dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0]
-    assert dummy_mm_item is not None, "Dummy item should be generated"
+        return [(modality, dummy_mm_item)] * max_items_per_batch
 
-    return [(modality, dummy_mm_item)] * max_items_per_batch
+    def reset_cache(self) -> None:
+        if self.cache is not None:
+            self.cache.clear_cache()
diff --git a/vllm/multimodal/gpu_ipc_memory.py b/vllm/multimodal/gpu_ipc_memory.py
index 58f807303700..7ccc1bdaa9d0 100644
--- a/vllm/multimodal/gpu_ipc_memory.py
+++ b/vllm/multimodal/gpu_ipc_memory.py
@@ -183,6 +183,7 @@ def reserve_mm_ipc_gpu_memory(
 
     Raises:
         ValueError: If the reservation leaves no memory for the KV cache.
+
     """
     if mm_config is None:
         return available_kv_cache_memory_bytes
diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py
index 3d86b33d5138..ae52ceefc409 100644
--- a/vllm/multimodal/hasher.py
+++ b/vllm/multimodal/hasher.py
@@ -18,13 +18,34 @@
 
 logger = init_logger(__name__)
 
+# Framing for the digest input. The hash is built by feeding a stream of byte
+# chunks to `hasher.update`, so the stream must be uniquely decodable: without
+# an explicit length in front of every chunk and a tag in front of every
+# container, distinct inputs serialize to identical bytes and share a cache
+# entry. See `test_hash_collision_*` in tests/multimodal/test_hasher.py.
+_LENGTH_BYTES = 8
+_TAG_NONE = b"\x00"
+_TAG_SEQUENCE = b"\x01"
+_TAG_MAPPING = b"\x02"
+_TAG_LEAF = b"\x03"
+
+
+def _encode_length(value: int) -> bytes:
+    return value.to_bytes(_LENGTH_BYTES, "little")
+
+
+def _framed(chunk: bytes | memoryview) -> Iterable[bytes | memoryview]:
+    """Yield *chunk* preceded by its size in bytes."""
+    size = chunk.nbytes if isinstance(chunk, memoryview) else len(chunk)
+    yield _encode_length(size)
+    yield chunk
+
 
 @functools.lru_cache(maxsize=3)
 def _get_hasher_factory(
     algorithm: MMHasherAlgorithm,
 ) -> Callable[[], "hashlib._Hash"]:
-    """
-    Get the hasher factory based on the configured algorithm.
+    """Get the hasher factory based on the configured algorithm.
 
     Args:
         algorithm: Hash algorithm name (blake3, sha256, or sha512)
@@ -33,8 +54,8 @@ def _get_hasher_factory(
     Supports blake3 (default), sha256, and sha512 for FIPS compliance.
 
     See: https://github.com/vllm-project/vllm/issues/18334
-    """
 
+    """
     if algorithm == "blake3":
         from blake3 import blake3
 
@@ -48,23 +69,41 @@ def _get_hasher_factory(
         raise ValueError(f"Unsupported hash algorithm: {algorithm}")
 
 
+def _get_image_id_bytes(image: Image.Image) -> bytes | None:
+    try:
+        exif = image.getexif()
+        image_id = exif.get(Image.ExifTags.Base.ImageID)
+        if isinstance(image_id, uuid.UUID):
+            return image_id.bytes
+    except Exception:
+        # Tolerate malformed EXIF metadata (e.g. invalid TIFF header)
+        # and fall back to serializing raw image data or bytes.
+        pass
+    return None
+
+
 class MultiModalHasher:
+    """Derives multi-modal cache keys.
+
+    Every method here yields *framed* chunks: each chunk is preceded by its
+    length and each container by its kind, so that the concatenation fed to the
+    digest is uniquely decodable and distinct inputs cannot share a key.
+    """
+
     @classmethod
     def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]:
         # Simple cases
         if isinstance(obj, (bytes, memoryview)):
-            return (obj,)
+            return _framed(obj)
         if isinstance(obj, str):
-            return (obj.encode("utf-8"),)
+            return _framed(obj.encode("utf-8"))
         if isinstance(obj, (int, float)):
-            return (np.array(obj).tobytes(),)
+            return _framed(np.array(obj).tobytes())
 
         if isinstance(obj, Image.Image):
-            exif = obj.getexif()
-            if Image.ExifTags.Base.ImageID in exif and isinstance(
-                exif[Image.ExifTags.Base.ImageID], uuid.UUID
-            ):
-                return (exif[Image.ExifTags.Base.ImageID].bytes,)
+            image_id = _get_image_id_bytes(obj)
+            if image_id is not None:
+                return _framed(image_id)
 
             data = {"mode": obj.mode, "data": np.asarray(obj)}
             palette = obj.palette
@@ -76,11 +115,9 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]:
             return cls.iter_item_to_bytes("image", data)
 
         if isinstance(obj, MediaWithBytes) and isinstance(obj.media, Image.Image):
-            exif = obj.media.getexif()
-            if Image.ExifTags.Base.ImageID in exif and isinstance(
-                exif[Image.ExifTags.Base.ImageID], uuid.UUID
-            ):
-                return (exif[Image.ExifTags.Base.ImageID].bytes,)
+            image_id = _get_image_id_bytes(obj.media)
+            if image_id is not None:
+                return _framed(image_id)
 
             if obj.io_config:
                 return cls.iter_item_to_bytes(
@@ -89,8 +126,11 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]:
                 )
             return cls.iter_item_to_bytes("image", obj.original_bytes)
 
-        if isinstance(obj, MediaWithBytes) and isinstance(obj.media, np.ndarray):
+        if isinstance(obj, MediaWithBytes) and isinstance(
+            obj.media, (np.ndarray, torch.Tensor)
+        ):
             frames = obj.media
+            # Both np.ndarray and torch.Tensor expose .nbytes.
             if frames.nbytes < len(obj.original_bytes):
                 return cls.iter_item_to_bytes("video", frames)
             return cls.iter_item_to_bytes("video", obj.original_bytes)
@@ -140,7 +180,7 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]:
             "No serialization method found for %s. Falling back to pickle.", type(obj)
         )
 
-        return (pickle.dumps(obj),)
+        return _framed(pickle.dumps(obj))
 
     @classmethod
     def iter_item_to_bytes(
@@ -148,18 +188,36 @@ def iter_item_to_bytes(
         key: str,
         obj: object,
     ) -> Iterable[bytes | memoryview]:
+        """Yield the digest input for a single ``key``/``obj`` pair."""
+        yield from _framed(key.encode("utf-8"))
+        yield from cls.iter_value_to_bytes(obj)
+
+    @classmethod
+    def iter_value_to_bytes(
+        cls,
+        obj: object,
+    ) -> Iterable[bytes | memoryview]:
+        """Yield the digest input for a value, tagged by its container kind.
+
+        Containers carry their kind and length so that a nested structure can
+        never serialize to the same bytes as a differently shaped one (a list
+        and a mapping keyed by stringified indices, for example).
+        """
         if obj is None:
-            yield key.encode("utf-8")
-            return
-        # Recursive cases
-        if isinstance(obj, (list, tuple)):
-            for i, elem in enumerate(obj):
-                yield from cls.iter_item_to_bytes(f"{key}.{i}", elem)
+            yield _TAG_NONE
+        elif isinstance(obj, (list, tuple)):
+            yield _TAG_SEQUENCE
+            yield _encode_length(len(obj))
+            for elem in obj:
+                yield from cls.iter_value_to_bytes(elem)
         elif isinstance(obj, dict):
+            yield _TAG_MAPPING
+            yield _encode_length(len(obj))
             for k, v in obj.items():
-                yield from cls.iter_item_to_bytes(f"{key}.{k}", v)
+                yield from _framed(str(k).encode("utf-8"))
+                yield from cls.iter_value_to_bytes(v)
         else:
-            yield key.encode("utf-8")
+            yield _TAG_LEAF
             yield from cls.serialize_item(obj)
 
     @classmethod
diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py
index bd1fc169d144..91fa34e86422 100644
--- a/vllm/multimodal/inputs.py
+++ b/vllm/multimodal/inputs.py
@@ -120,11 +120,9 @@ class VisionChunkVideo(TypedDict):
 
 @dataclass(frozen=True)
 class PlaceholderRange:
-    """
-    Placeholder location information for multi-modal data.
+    """Placeholder location information for multi-modal data.
 
     Example:
-
     Prompt: `AAAA BBBB What is in these images?`
 
     Images A and B will have:
@@ -133,6 +131,7 @@ class PlaceholderRange:
     A: PlaceholderRange(offset=0, length=4)
     B: PlaceholderRange(offset=5, length=4)
     ```
+
     """
 
     offset: int
@@ -161,8 +160,7 @@ def get_num_embeds(self) -> int:
     def get_embeds_indices_in_range(
         self, start_idx: int, end_idx: int
     ) -> tuple[int, int]:
-        """
-        Returns the starting and ending indices of the embeddings of encoder outputs
+        """Returns the starting and ending indices of the embeddings of encoder outputs
         in the range of [start_idx, end_idx) in the placeholders.
 
         For example, given:
@@ -190,6 +188,7 @@ def extract_embeds_range(self) -> list[tuple[int, int]]:
             A tuple `(start, end)` representing the start and end
             indices (inclusive) of the embedded region.
             Returns full placeholder range if `is_embed` is `None`.
+
         """
         if self.is_embed is None:
             return [(self.offset, self.offset + self.length - 1)]
@@ -234,8 +233,7 @@ def nested_tensors_equal(
     b: NestedTensors,
     check_dtype: bool = True,
 ) -> bool:
-    """
-    Equality check between
+    """Equality check between
     [`NestedTensors`][vllm.multimodal.inputs.NestedTensors] objects.
 
     If `check_dtype` is `True`, the tensors must have the same dtype.
@@ -336,8 +334,7 @@ def _h2d(x: torch.Tensor) -> torch.Tensor:
 
 
 def batched_tensors_equal(a: BatchedTensorInputs, b: BatchedTensorInputs) -> bool:
-    """
-    Equality check between
+    """Equality check between
     [`BatchedTensorInputs`][vllm.multimodal.inputs.BatchedTensorInputs] objects.
     """
     return all(k in b and nested_tensors_equal(a[k], b[k]) for k in a)
@@ -345,8 +342,7 @@ def batched_tensors_equal(a: BatchedTensorInputs, b: BatchedTensorInputs) -> boo
 
 @dataclass
 class MultiModalFeatureSpec:
-    """
-    Represents a single multimodal input with its processed data and metadata.
+    """Represents a single multimodal input with its processed data and metadata.
 
     Used to track multimodal data through processing and caching.
     A request containing multiple multimodal items will have one
@@ -392,8 +388,7 @@ def gather_kwargs(features: list["MultiModalFeatureSpec"], keys: set[str]):
 
 @dataclass
 class MultiModalFieldElem:
-    """
-    Represents a processed keyword argument to pass to a model for a
+    """Represents a processed keyword argument to pass to a model for a
     [`MultiModalKwargsItem`][vllm.multimodal.inputs.MultiModalKwargsItem].
     """
 
@@ -429,8 +424,7 @@ def __eq__(self, other: object) -> bool:
 
 @dataclass(frozen=True, kw_only=True)
 class BaseMultiModalField(ABC):
-    """
-    Defines how to interpret tensor data belonging to a keyword argument for
+    """Defines how to interpret tensor data belonging to a keyword argument for
     [`MultiModalKwargsItems`][vllm.multimodal.inputs.MultiModalKwargsItems],
     and vice versa.
     """
@@ -458,8 +452,7 @@ def build_elems(
         key: str,
         data: NestedTensors,
     ) -> Sequence[MultiModalFieldElem]:
-        """
-        Construct
+        """Construct
         [`MultiModalFieldElem`][vllm.multimodal.inputs.MultiModalFieldElem]
         instances to represent the provided data.
 
@@ -481,8 +474,7 @@ def reduce_data(
         device: torch.types.Device = None,
         pin_memory: bool = False,
     ) -> NestedTensors:
-        """
-        Merge the data from multiple instances of
+        """Merge the data from multiple instances of
         [`MultiModalFieldElem`][vllm.multimodal.inputs.MultiModalFieldElem].
 
         This is the inverse of
@@ -511,9 +503,8 @@ def reduce_data(
 
 @dataclass(frozen=True, kw_only=True)
 class MultiModalBatchedField(BaseMultiModalField):
-    """
-    Info:
-        [`MultiModalFieldConfig.batched`][vllm.multimodal.inputs.MultiModalFieldConfig.batched]
+    """Info:
+    [`MultiModalFieldConfig.batched`][vllm.multimodal.inputs.MultiModalFieldConfig.batched]
     """
 
     def build_elems(
@@ -547,10 +538,9 @@ def _reduce_data(
 
 @dataclass(frozen=True, kw_only=True)
 class MultiModalFlatField(BaseMultiModalField):
-    """
-    Info:
-        [`MultiModalFieldConfig.flat`][vllm.multimodal.inputs.MultiModalFieldConfig.flat]
-        [`MultiModalFieldConfig.flat_from_sizes`][vllm.multimodal.inputs.MultiModalFieldConfig.flat_from_sizes]
+    """Info:
+    [`MultiModalFieldConfig.flat`][vllm.multimodal.inputs.MultiModalFieldConfig.flat]
+    [`MultiModalFieldConfig.flat_from_sizes`][vllm.multimodal.inputs.MultiModalFieldConfig.flat_from_sizes]
     """
 
     slices: Sequence[slice] | Sequence[Sequence[slice]]
@@ -637,9 +627,8 @@ def _shape_before_after(tensor: torch.Tensor):
 
 @dataclass(frozen=True, kw_only=True)
 class MultiModalSharedField(BaseMultiModalField):
-    """
-    Info:
-        [`MultiModalFieldConfig.shared`][vllm.multimodal.inputs.MultiModalFieldConfig.shared]
+    """Info:
+    [`MultiModalFieldConfig.shared`][vllm.multimodal.inputs.MultiModalFieldConfig.shared]
     """
 
     batch_size: int
@@ -666,8 +655,7 @@ def _reduce_data(
 class MultiModalFieldConfig:
     @staticmethod
     def batched(modality: str, *, keep_on_cpu: bool = False):
-        """
-        Defines a field where an element in the batch is obtained by
+        """Defines a field where an element in the batch is obtained by
         indexing into the first dimension of the underlying data.
 
         Args:
@@ -676,7 +664,6 @@ def batched(modality: str, *, keep_on_cpu: bool = False):
             keep_on_cpu: Whether to keep this field on the CPU for the model inputs.
 
         Example:
-
         ```
         Input:
             Data: [[AAAA]
@@ -688,6 +675,7 @@ def batched(modality: str, *, keep_on_cpu: bool = False):
             Element 2: [BBBB]
             Element 3: [CCCC]
         ```
+
         """
         return MultiModalFieldConfig(
             field=MultiModalBatchedField(keep_on_cpu=keep_on_cpu),
@@ -702,8 +690,7 @@ def flat(
         *,
         keep_on_cpu: bool = False,
     ):
-        """
-        Defines a field where an element in the batch is obtained by
+        """Defines a field where an element in the batch is obtained by
         slicing along the first dimension of the underlying data.
 
         Args:
@@ -716,7 +703,6 @@ def flat(
             keep_on_cpu: Whether to keep this field on the CPU for the model inputs.
 
         Example:
-
         ```
         Given:
             slices: [slice(0, 3), slice(3, 7), slice(7, 9)]
@@ -746,6 +732,7 @@ def flat(
             Element 2: [[B],[B],[B],[B]]
             Element 3: [[C],[C]]
         ```
+
         """
         return MultiModalFieldConfig(
             field=MultiModalFlatField(
@@ -764,8 +751,7 @@ def flat_from_sizes(
         *,
         keep_on_cpu: bool = False,
     ):
-        """
-        Defines a field where an element in the batch is obtained by
+        """Defines a field where an element in the batch is obtained by
         slicing along the first dimension of the underlying data.
 
         Args:
@@ -777,7 +763,6 @@ def flat_from_sizes(
             keep_on_cpu: Whether to keep this field on the CPU for the model inputs.
 
         Example:
-
         ```
         Given:
             size_per_item: [3, 4, 2]
@@ -807,8 +792,8 @@ def flat_from_sizes(
 
         Info:
             [`MultiModalFieldConfig.flat`][vllm.multimodal.inputs.MultiModalFieldConfig.flat]
-        """
 
+        """
         if size_per_item.ndim != 1:
             raise ValueError(
                 "size_per_item should be a 1-D tensor, "
@@ -836,8 +821,7 @@ def shared(
         *,
         keep_on_cpu: bool = False,
     ):
-        """
-        Defines a field where an element in the batch is obtained by
+        """Defines a field where an element in the batch is obtained by
         taking the entirety of the underlying data.
 
         This means that the data is the same for each element in the batch.
@@ -849,7 +833,6 @@ def shared(
             keep_on_cpu: Whether to keep this field on the CPU for the model inputs.
 
         Example:
-
         ```
         Given:
             batch_size: 4
@@ -863,6 +846,7 @@ def shared(
             Element 3: [XYZ]
             Element 4: [XYZ]
         ```
+
         """
         return MultiModalFieldConfig(
             field=MultiModalSharedField(
@@ -884,8 +868,7 @@ def build_elems(
 
 
 class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]):
-    """
-    A dictionary of processed keyword arguments to pass to the model,
+    """A dictionary of processed keyword arguments to pass to the model,
     corresponding to a single item in
     [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems].
     """
@@ -912,8 +895,7 @@ def get_data(self) -> dict[str, NestedTensors]:
 
 
 class MultiModalKwargsItems(UserDict[str, Sequence[_I]]):
-    """
-    A dictionary of processed multi-modal inputs by modality.
+    """A dictionary of processed multi-modal inputs by modality.
 
     For example, given a processor that processes
     images into `pixel_values` and `image_grid_thw`,
diff --git a/vllm/multimodal/media/audio.py b/vllm/multimodal/media/audio.py
index e1e75dd8d744..1f1cda8dccec 100644
--- a/vllm/multimodal/media/audio.py
+++ b/vllm/multimodal/media/audio.py
@@ -92,6 +92,10 @@ def load_audio_pyav(
     Args:
         path: A :class:`~io.BytesIO` buffer, a filesystem
             :class:`~pathlib.Path`, or a string path.
+        sr: Target sample rate, or None to keep the native rate.
+        mono: Whether to average the channels down to mono.
+        max_decode_bytes: If set, abort decoding once this many bytes have
+            been read from the source.
         max_duration_s: If set, abort decoding once the accumulated
             sample count exceeds this many seconds of audio.  Prevents
             decompression-bomb attacks where a small compressed file
@@ -100,6 +104,7 @@ def load_audio_pyav(
     Returns:
         ``(waveform, sample_rate)`` where *waveform* is a 1-D float32
         NumPy array and *sample_rate* is the native sample rate in Hz.
+
     """
     try:
         container = av.open(path)
@@ -209,7 +214,7 @@ def load_audio_soundfile(
     max_duration_s: float | None = None,
     max_decode_bytes: int | None = None,
 ) -> tuple[np.ndarray, int]:
-    """Load audio via soundfile"""
+    """Load audio via soundfile."""
     with soundfile.SoundFile(path) as f:
         native_sr = f.samplerate
         if max_duration_s is not None:
@@ -274,6 +279,7 @@ def load_audio_torchcodec(
         ``(waveform, sample_rate)`` where *waveform* is a float32 NumPy
         array (1-D when ``mono=True``) and *sample_rate* is the output
         sample rate in Hz.
+
     """
     if AudioDecoder is None:
         # Unify "torchcodec not installed" and "system ffmpeg missing" into
@@ -412,10 +418,17 @@ def load_audio(
     """Load audio using the selected decoding backend.
 
     Args:
+        path: Audio file path or in-memory buffer.
+        sr: Target sample rate, or None to keep the native rate.
+        mono: Whether to downmix to a single channel.
+        max_duration_s: Reject audio longer than this many seconds.
+        max_decode_bytes: Reject audio that would decode to more than this
+            many bytes.
         backend: One of ``AUDIO_BACKENDS``. ``None`` (default) selects
             ``"auto"``, which tries soundfile, then torchcodec, then PyAV;
             the other values select a single
             backend with no fallback.
+
     """
     backend = backend or "auto"
     if backend not in AUDIO_BACKENDS:
diff --git a/vllm/multimodal/media/base.py b/vllm/multimodal/media/base.py
index 5a58b9ff72d8..d3e622156fc8 100644
--- a/vllm/multimodal/media/base.py
+++ b/vllm/multimodal/media/base.py
@@ -14,8 +14,7 @@
 
 @dataclass
 class MediaWithBytes(Generic[_T]):
-    """
-    Wrapper that couples a media object with its original encoded bytes.
+    """Wrapper that couples a media object with its original encoded bytes.
 
     This ensures the raw bytes and media object remain synchronized,
     preventing cache corruption from in-place modifications.
@@ -89,8 +88,7 @@ def load_bytes(self, data: bytes) -> _T:
 
     @abstractmethod
     def load_base64(self, media_type: str, data: str) -> _T:
-        """
-        List of media types:
+        """List of media types:
         https://www.iana.org/assignments/media-types/media-types.xhtml
         """
         raise NotImplementedError
diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py
index 8980418b3275..297dc91af232 100644
--- a/vllm/multimodal/media/connector.py
+++ b/vllm/multimodal/media/connector.py
@@ -73,6 +73,7 @@ def _wrap_media_fetch_error(
             408/429, invalid URL)
         Original exception for transient errors (5xx, 408, 429, network blips)
             or other exceptions
+
     """
     if isinstance(exc, VLLMValidationError):
         return exc
@@ -162,16 +163,16 @@ def __init__(
         allowed_local_media_path: str = "",
         allowed_media_domains: list[str] | None = None,
     ) -> None:
-        """
-        Args:
-            media_io_kwargs: Additional args passed to process media
-                             inputs, keyed by modalities. For example,
-                             to set num_frames for video, set
-                             `--media-io-kwargs '{"video":{"num_frames":40}}'`
-            connection: HTTP connection client to download media contents.
-            allowed_local_media_path: A local directory to load media files from.
-            allowed_media_domains: If set, only media URLs that belong to this
-                                   domain can be used for multi-modal inputs.
+        """Args:
+        media_io_kwargs: Additional args passed to process media
+                         inputs, keyed by modalities. For example,
+                         to set num_frames for video, set
+                         `--media-io-kwargs '{"video":{"num_frames":40}}'`
+        connection: HTTP connection client to download media contents.
+        allowed_local_media_path: A local directory to load media files from.
+        allowed_media_domains: If set, only media URLs that belong to this
+                               domain can be used for multi-modal inputs.
+
         """
         super().__init__()
 
@@ -468,9 +469,7 @@ def fetch_audio(
         self,
         audio_url: str,
     ) -> tuple[np.ndarray, int | float]:
-        """
-        Load audio from a URL.
-        """
+        """Load audio from a URL."""
         audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
 
         return self.load_from_url(
@@ -483,9 +482,7 @@ async def fetch_audio_async(
         self,
         audio_url: str,
     ) -> tuple[np.ndarray, int | float]:
-        """
-        Asynchronously fetch audio from a URL.
-        """
+        """Asynchronously fetch audio from a URL."""
         audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
 
         return await self.load_from_url_async(
@@ -500,8 +497,7 @@ def fetch_image(
         *,
         image_mode: str | None = "RGB",
     ) -> Image.Image:
-        """
-        Load a PIL image from an HTTP or base64 data URL.
+        """Load a PIL image from an HTTP or base64 data URL.
 
         By default, the image is converted into RGB format. Set
         `media_io_kwargs={"image": {"image_mode": None}}` to keep the
@@ -527,8 +523,7 @@ async def fetch_image_async(
         *,
         image_mode: str | None = "RGB",
     ) -> Image.Image:
-        """
-        Asynchronously load a PIL image from an HTTP or base64 data URL.
+        """Asynchronously load a PIL image from an HTTP or base64 data URL.
 
         By default, the image is converted into RGB format. Set
         `media_io_kwargs={"image": {"image_mode": None}}` to keep the
@@ -555,9 +550,7 @@ def fetch_video(
         image_mode: str | None = "RGB",
         video_processor: str | None = None,
     ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
-        """
-        Load video from an HTTP or base64 data URL.
-        """
+        """Load video from an HTTP or base64 data URL."""
         image_io = ImageMediaIO(
             **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
@@ -581,8 +574,7 @@ async def fetch_video_async(
         image_mode: str | None = "RGB",
         video_processor: str | None = None,
     ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
-        """
-        Asynchronously load video from an HTTP or base64 data URL.
+        """Asynchronously load video from an HTTP or base64 data URL.
 
         By default, the image is converted into RGB format. Set
         `media_io_kwargs={"image": {"image_mode": None}}` to keep the
@@ -608,9 +600,7 @@ def fetch_image_embedding(
         self,
         data: str,
     ) -> torch.Tensor:
-        """
-        Load image embedding from a URL.
-        """
+        """Load image embedding from a URL."""
         image_embedding_io = ImageEmbeddingMediaIO()
 
         return image_embedding_io.load_base64("", data)
@@ -619,9 +609,7 @@ async def fetch_image_embedding_async(
         self,
         data: str,
     ) -> torch.Tensor:
-        """
-        Asynchronously load image embedding from a URL.
-        """
+        """Asynchronously load image embedding from a URL."""
         image_embedding_io = ImageEmbeddingMediaIO()
         loop = asyncio.get_running_loop()
 
@@ -633,9 +621,7 @@ def fetch_audio_embedding(
         self,
         data: str,
     ) -> torch.Tensor:
-        """
-        Load audio embedding from a URL.
-        """
+        """Load audio embedding from a URL."""
         audio_embedding_io = AudioEmbeddingMediaIO()
 
         return audio_embedding_io.load_base64("", data)
@@ -644,9 +630,7 @@ async def fetch_audio_embedding_async(
         self,
         data: str,
     ) -> torch.Tensor:
-        """
-        Asynchronously load audio embedding from a URL.
-        """
+        """Asynchronously load audio embedding from a URL."""
         audio_embedding_io = AudioEmbeddingMediaIO()
         loop = asyncio.get_running_loop()
 
@@ -658,9 +642,7 @@ def fetch_video_embedding(
         self,
         data: str,
     ) -> torch.Tensor:
-        """
-        Load video embedding from a URL.
-        """
+        """Load video embedding from a URL."""
         video_embedding_io = VideoEmbeddingMediaIO()
 
         return video_embedding_io.load_base64("", data)
@@ -669,9 +651,7 @@ async def fetch_video_embedding_async(
         self,
         data: str,
     ) -> torch.Tensor:
-        """
-        Asynchronously load video embedding from a URL.
-        """
+        """Asynchronously load video embedding from a URL."""
         video_embedding_io = VideoEmbeddingMediaIO()
         loop = asyncio.get_running_loop()
 
diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py
index e5c073f5f382..7528331cdfd0 100644
--- a/vllm/multimodal/media/video.py
+++ b/vllm/multimodal/media/video.py
@@ -19,17 +19,23 @@
     safe_to_dense,
 )
 
-from ..video import VIDEO_LOADER_REGISTRY
+from ..video import VIDEO_LOADER_REGISTRY, DecodedFrames
+from ..video_decoders import incompatible_backend_options
 from .base import MediaIO, MediaWithBytes
 from .image import MAGIC_NUMPY_PREFIX, ImageMediaIO
 
 logger = init_logger(__name__)
 
 
-class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]):
+class VideoMediaIO(MediaIO[MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]]):
     """Configuration values can be user-provided either by --media-io-kwargs or
     by the runtime API field "media_io_kwargs". Ensure proper validation and
     error handling.
+
+    Use --media-io-kwargs '{"video": {"backend": "torchcodec",
+    "device": "cuda"}}' to decode with NVDEC and keep frames on the GPU,
+    which pairs with a device-side multi-modal processor
+    (--mm-processor-device, EPD encode-only instances).
     """
 
     @classmethod
@@ -43,6 +49,9 @@ def merge_kwargs(
             runtime_kwargs = dict(runtime_kwargs)
             runtime_kwargs.pop("hw_decoders", None)
             runtime_kwargs.pop("pool_size", None)
+            # The decode device determines which device the decoded frames
+            # (and thus the processor pipeline) live on; startup-only.
+            runtime_kwargs.pop("device", None)
 
             # Block request-level selection of GPU video backends that
             # were not configured (and VRAM-reserved) at startup.
@@ -70,6 +79,14 @@ def merge_kwargs(
                 merged.pop("fps", None)
             elif "fps" in runtime_kwargs and "num_frames" not in runtime_kwargs:
                 merged.pop("num_frames", None)
+
+            requested_backend = runtime_kwargs.get("backend")
+            static_backend = (default_kwargs or {}).get("backend")
+            if requested_backend and requested_backend != static_backend:
+                for name in incompatible_backend_options(
+                    static_backend, requested_backend
+                ):
+                    merged.pop(name, None)
         return merged
 
     def __init__(
@@ -101,7 +118,7 @@ def __init__(
 
     def load_bytes(
         self, data: bytes
-    ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
+    ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]:
         video = self.video_loader.load_bytes(
             data, num_frames=self.num_frames, **self.kwargs
         )
@@ -109,7 +126,7 @@ def load_bytes(
 
     def load_base64(
         self, media_type: str, data: str
-    ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
+    ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]:
         if media_type.lower() == "video/jpeg":
             load_frame = partial(
                 self.image_io.load_base64,
@@ -173,11 +190,11 @@ def load_base64(
             }
             return MediaWithBytes((frames, metadata), data.encode())
 
-        return self.load_bytes(pybase64.b64decode(data))
+        return self.load_bytes(pybase64.b64decode(data, validate=True))
 
     def load_file(
         self, filepath: Path
-    ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
+    ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]:
         with filepath.open("rb") as f:
             data = f.read()
 
diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py
index d96f19f39b95..98f8a8b2f3e9 100644
--- a/vllm/multimodal/parse.py
+++ b/vllm/multimodal/parse.py
@@ -36,6 +36,7 @@
     VideoItem,
 )
 from .media import MediaWithBytes
+from .video import DecodedFrames
 
 _T = TypeVar("_T")
 _I = TypeVar("_I")
@@ -57,8 +58,7 @@
 
 
 class ModalityDataItems(ABC, Generic[_T, _I]):
-    """
-    Represents data items for a modality in
+    """Represents data items for a modality in
     [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems].
     """
 
@@ -150,6 +150,7 @@ def validate_embedding_ndim(
         tensor: The tensor to validate.
         modality: The modality name for error messages (e.g., "image", "audio").
         index: Optional index for list items, included in error messages.
+
     """
     if tensor.ndim < 2 or tensor.ndim > 3:
         idx_str = f" [{index}]" if index is not None else ""
@@ -163,8 +164,7 @@ def validate_embedding_ndim(
 class EmbeddingItems(
     ModalityDataItems[torch.Tensor | list[torch.Tensor], torch.Tensor]
 ):
-    """
-    Base class for data items that are expressed as a batched embedding tensor,
+    """Base class for data items that are expressed as a batched embedding tensor,
     or a list of embedding tensors (one per item).
     """
 
@@ -250,8 +250,7 @@ def get_feature_size(self, item_idx: int) -> int:
 class DictEmbeddingItems(
     ModalityDataItems[Mapping[str, torch.Tensor], Mapping[str, torch.Tensor]]
 ):
-    """
-    Base class for data items that are expressed as a dictionary of tensors.
+    """Base class for data items that are expressed as a dictionary of tensors.
 
     Usually, the dictionary keys correspond to the outputs of HF processor.
     """
@@ -267,18 +266,18 @@ def __init__(
         ],
         optional_fields: Set[str] = frozenset(),
     ) -> None:
-        """
-        Args:
-            data: The dictionary of tensors for this modality.
-            modality: The modality these items belong to.
-            required_fields: Fields `data` must contain.
-            fields_factory: Builds the field config from the data.
-            optional_fields: Fields `data` may omit. Which fields these are is
-                the caller's decision -- see
-                `MultiModalDataParser.embedding_field_sets`, where a deployment
-                that receives embeddings through an EC connector makes the
-                embeddings optional. They still need a field config, since they
-                are used whenever they *are* supplied.
+        """Args:
+        data: The dictionary of tensors for this modality.
+        modality: The modality these items belong to.
+        required_fields: Fields `data` must contain.
+        fields_factory: Builds the field config from the data.
+        optional_fields: Fields `data` may omit. Which fields these are is
+            the caller's decision -- see
+            `MultiModalDataParser.embedding_field_sets`, where a deployment
+            that receives embeddings through an EC connector makes the
+            embeddings optional. They still need a field config, since they
+            are used whenever they *are* supplied.
+
         """
         from transformers.feature_extraction_utils import BatchFeature
 
@@ -470,14 +469,12 @@ def __init__(self, data: Sequence[Any]) -> None:
 
 
 class MultiModalDataItems(UserDict[str, ModalityDataItems[Any, Any]]):
-    """
-    A normalized [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict]
+    """A normalized [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict]
     such that each entry corresponds to a list.
     """
 
     def select(self, modalities: Set[str]):
-        """
-        Construct a new `MultiModalDataItems` instance containing only the
+        """Construct a new `MultiModalDataItems` instance containing only the
         selected modalities.
         """
         return MultiModalDataItems(
@@ -485,8 +482,7 @@ def select(self, modalities: Set[str]):
         )
 
     def get_count(self, modality: str, *, strict: bool = True) -> int:
-        """
-        Get the number of data items belonging to a modality.
+        """Get the number of data items belonging to a modality.
 
         If `strict=False`, return `0` instead of raising [`KeyError`][]
         even if the modality is not found.
@@ -512,8 +508,7 @@ def get_items(
         modality: str,
         typ: type[_D] | tuple[type[_D], ...],
     ) -> _D:
-        """
-        Get the data items belonging to a modality,
+        """Get the data items belonging to a modality,
         requiring that they belong to a certain type.
         """
         if modality not in self:
@@ -540,8 +535,7 @@ def get_items(
 
 
 class MultiModalDataParser:
-    """
-    Parses [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict]
+    """Parses [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict]
     into [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems].
 
     Args:
@@ -560,6 +554,7 @@ class MultiModalDataParser:
         allow_missing_mm_embeddings (bool): Whether pre-computed embedding
             tensors may be absent from the request on a disaggregated consumer.
             Derived by `BaseProcessingInfo.allow_missing_mm_embeddings`.
+
     """
 
     embedding_fields: Mapping[str, Mapping[str, EmbeddingFieldRole]] = {}
@@ -653,7 +648,7 @@ def _get_audio_with_sr(
     def _get_video_with_metadata(
         self,
         video: VideoItem,
-    ) -> tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any] | None]:
+    ) -> tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any] | None]:
         if isinstance(video, MediaWithBytes):
             new_video, metadata = self._get_video_with_metadata(video.media)
             return MediaWithBytes(new_video, video.original_bytes), metadata
@@ -661,16 +656,54 @@ def _get_video_with_metadata(
             return video
         if isinstance(video, list):
             return np.array(video), None
-        if isinstance(video, np.ndarray):
+        if isinstance(video, (np.ndarray, torch.Tensor)):
+            # Tensors pass through untouched: HF video processors accept them
+            # directly, and device tensors (e.g. NVDEC-decoded frames) must
+            # stay on-device to avoid a D2H round-trip.
             return video, None
-        if isinstance(video, torch.Tensor):
-            return video.numpy(), None
 
         assert_never(video)
 
+    def _parse_audio_embedding_data(
+        self, data: dict[str, torch.Tensor]
+    ) -> DictEmbeddingItems:
+        counts = data.get("audio_num_tokens")
+        if (
+            not isinstance(counts, torch.Tensor)
+            or counts.dtype not in (torch.int32, torch.int64)
+            or counts.ndim not in (1, 2)
+            or (counts.ndim == 2 and counts.shape[1] != 1)
+            or not torch.all(counts > 0)
+        ):
+            raise ValueError(
+                "audio_num_tokens must contain one positive integer per audio"
+            )
+
+        counts = counts.flatten()
+        if "audio_embeds" in data:
+            embeds = data["audio_embeds"]
+            if len(embeds) != len(counts) or any(
+                len(embedding) != count for embedding, count in zip(embeds, counts)
+            ):
+                raise ValueError("audio_num_tokens does not match audio_embeds")
+
+        required, optional = self.embedding_field_sets("audio")
+        return DictEmbeddingItems(
+            {**data, "audio_num_tokens": counts},
+            modality="audio",
+            required_fields=required,
+            optional_fields=optional,
+            fields_factory=lambda _: {
+                "audio_embeds": MultiModalFieldConfig.batched("audio"),
+                "audio_num_tokens": MultiModalFieldConfig.batched(
+                    "audio", keep_on_cpu=True
+                ),
+            },
+        )
+
     def _parse_audio_data(
         self,
-        data: ModalityData[AudioItem],
+        data: dict[str, torch.Tensor] | ModalityData[AudioItem],
     ) -> ModalityDataItems[Any, Any] | None:
         if data is None:
             return None
@@ -678,7 +711,12 @@ def _parse_audio_data(
         if self.is_embeddings(data):
             return AudioEmbeddingItems(data, self.expected_hidden_size)
 
-        data_items: list[AudioItem]
+        if isinstance(data, dict) and "audio_num_tokens" in self.embedding_fields.get(
+            "audio", {}
+        ):
+            return self._parse_audio_embedding_data(data)
+
+        data_items: list[AudioItem | None]
         if (
             (is_list_of(data, float) and len(data) > 0)
             or (isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 1)
@@ -690,8 +728,13 @@ def _parse_audio_data(
         else:
             data_items = data  # type: ignore[assignment]
 
-        new_audios = list[np.ndarray]()
+        new_audios = list[np.ndarray | None]()
         for data_item in data_items:
+            # Requests can omit audio samples when reusing a cached UUID.
+            if data_item is None:
+                new_audios.append(None)
+                continue
+
             audio, orig_sr = self._get_audio_with_sr(data_item)
             if orig_sr is None:
                 new_audio = audio
@@ -746,8 +789,14 @@ def _parse_video_data(
             return VideoEmbeddingItems(data, self.expected_hidden_size)
 
         data_items: list[VideoItem]
-        if (is_list_of(data, PILImage.Image) and len(data) > 0) or (
-            isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4
+        if (
+            (is_list_of(data, PILImage.Image) and len(data) > 0)
+            or (
+                is_list_of(data, (np.ndarray, torch.Tensor), check="all")
+                and len(data) > 0
+                and all(item.ndim == 3 for item in data)
+            )
+            or (isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4)
         ):
             data_items = [data]
         elif isinstance(data, (np.ndarray, torch.Tensor)):
@@ -758,9 +807,9 @@ def _parse_video_data(
             data_items = data  # type: ignore[assignment]
 
         new_videos = list[
-            np.ndarray
-            | MediaWithBytes[np.ndarray]
-            | tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any]]
+            DecodedFrames
+            | MediaWithBytes[DecodedFrames]
+            | tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any]]
             | None
         ]()
         metadata_lst: list[dict[str, Any] | None] = []
diff --git a/vllm/multimodal/processing/context.py b/vllm/multimodal/processing/context.py
index 07149dbc4a26..db1cffefd6ea 100644
--- a/vllm/multimodal/processing/context.py
+++ b/vllm/multimodal/processing/context.py
@@ -26,14 +26,14 @@
 from vllm.utils.mistral import is_mistral_tokenizer
 
 if TYPE_CHECKING:
-    from transformers.configuration_utils import PretrainedConfig
+    from transformers.configuration_utils import PreTrainedConfig
     from transformers.feature_extraction_utils import BatchFeature
     from transformers.processing_utils import ProcessorMixin
 
     from vllm.config import ModelConfig
     from vllm.renderers import TokenizeParams
 else:
-    PretrainedConfig = object
+    PreTrainedConfig = object
     BatchFeature = object
     ProcessorMixin = object
 
@@ -70,6 +70,7 @@ def overlay_modality_mm_kwargs(
 
     Returns:
         A new dict with the modality-scoped keys overlaid when present.
+
     """
     merged = dict(kwargs)
     if modality is None:
@@ -122,14 +123,13 @@ def get_stats_dict(self):
 
 
 _T = TypeVar("_T")
-_C = TypeVar("_C", bound=PretrainedConfig, default=PretrainedConfig)
+_C = TypeVar("_C", bound=PreTrainedConfig, default=PreTrainedConfig)
 _P = TypeVar("_P", bound=ProcessorMixin, default=ProcessorMixin)
 
 
 @dataclass(frozen=True)
 class InputProcessingContext:
-    """
-    Contains information about the model which may be used to
+    """Contains information about the model which may be used to
     modify the inputs.
     """
 
@@ -148,7 +148,7 @@ def get_tokenizer(self) -> TokenizerLike:
         return self.tokenizer
 
     @overload
-    def get_hf_config(self, /) -> PretrainedConfig: ...
+    def get_hf_config(self, /) -> PreTrainedConfig: ...
 
     @overload
     def get_hf_config(
@@ -162,18 +162,18 @@ def get_hf_config(
         typ: type[Any] | tuple[type[Any], ...] | None = None,
         /,
     ) -> Any:
-        """
-        Get the HuggingFace configuration
-        (`transformers.PretrainedConfig`) of the model,
+        """Get the HuggingFace configuration
+        (`transformers.PreTrainedConfig`) of the model,
         additionally checking its type.
 
         Raises:
             TypeError: If the configuration is not of the specified type.
+
         """
         if typ is None:
-            from transformers.configuration_utils import PretrainedConfig
+            from transformers.configuration_utils import PreTrainedConfig
 
-            typ = PretrainedConfig
+            typ = PreTrainedConfig
 
         hf_config = self.model_config.hf_config
         if not isinstance(hf_config, typ):
@@ -186,17 +186,15 @@ def get_hf_config(
         return hf_config
 
     def get_hf_image_processor_config(self) -> dict[str, Any]:
-        """
-        Get the HuggingFace image processor configuration of the model.
-        """
+        """Get the HuggingFace image processor configuration of the model."""
         return self.model_config.hf_image_processor_config
 
     def get_mm_config(self):
-        """
-        Get the multimodal config of the model.
+        """Get the multimodal config of the model.
 
         Raises:
             RuntimeError: If the model is not a multimodal model.
+
         """
         mm_config = self.model_config.multimodal_config
         if mm_config is None:
@@ -221,13 +219,13 @@ def get_hf_processor(
         /,
         **kwargs: object,
     ) -> Any:
-        """
-        Get the HuggingFace processor
+        """Get the HuggingFace processor
         (`transformers.ProcessorMixin`) of the model,
         additionally checking its type.
 
         Raises:
             TypeError: If the processor is not of the specified type.
+
         """
         if typ is None:
             from transformers.processing_utils import ProcessorMixin
@@ -254,8 +252,7 @@ def init_processor(
         /,
         **kwargs: object,
     ) -> _T:
-        """
-        Initialize a HuggingFace-like processor class, merging the
+        """Initialize a HuggingFace-like processor class, merging the
         keyword arguments with those in the model's configuration.
         """
         merged_kwargs = self.get_merged_mm_kwargs(kwargs)
@@ -317,20 +314,13 @@ def call_hf_processor(
         data: Mapping[str, object],
         kwargs: Mapping[str, object] = {},
     ) -> BatchFeature:
-        """
-        Call `hf_processor` on the prompt `data`
+        """Call `hf_processor` on the prompt `data`
         (text, image, audio...) with configurable options `kwargs`.
         """
         assert callable(hf_processor)
 
         merged_kwargs = self.get_merged_mm_kwargs(kwargs)
 
-        # vLLM needs the full untruncated sequence to keep multi-modal
-        # placeholder tokens aligned; note that the text inputs in
-        # call_hf_processor are just dummy text, not the original prompt.
-        # The original prompt is already tokenized by the renderer.
-        merged_kwargs.setdefault("truncation", False)
-
         allowed_kwargs = get_allowed_kwarg_only_overrides(
             hf_processor,
             merged_kwargs,
@@ -381,12 +371,11 @@ def model_id(self) -> str:
     def get_tokenizer(self) -> TokenizerLike:
         return self.ctx.get_tokenizer()
 
-    def get_hf_config(self) -> PretrainedConfig:
+    def get_hf_config(self) -> PreTrainedConfig:
         return self.ctx.get_hf_config()
 
     def get_hf_processor(self, **kwargs: object) -> ProcessorMixin:
-        """
-        Subclasses can override this method to handle
+        """Subclasses can override this method to handle
         specific kwargs from model config or user inputs.
         """
         return self.ctx.get_hf_processor(**kwargs)
@@ -409,8 +398,7 @@ def default_tok_params(self) -> TokenizeParams:
         return self.get_default_tok_params()
 
     def _get_expected_hidden_size(self) -> int | None:
-        """
-        Get expected hidden size for embedding validation if `mm_embeds` are enabled.
+        """Get expected hidden size for embedding validation if `mm_embeds` are enabled.
 
         This validates hidden dimensions to prevent a vulnerability where embeddings
         with correct `ndim` but wrong `shape` could cause crashes at inference time.
@@ -430,10 +418,9 @@ def allow_missing_mm_embeddings(self) -> bool:
         return mm_config is not None and mm_config.allow_missing_mm_embeddings
 
     def get_data_parser(self) -> MultiModalDataParser:
-        """
-        Constructs a parser to preprocess multi-modal data items
+        """Constructs a parser to preprocess multi-modal data items
         before passing them to
-        [`_get_hf_mm_data`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_data].
+        [`_get_hf_mm_inputs`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs].
 
         You can support additional modalities by creating a subclass
         of [`MultiModalDataParser`][vllm.multimodal.parse.MultiModalDataParser]
@@ -453,8 +440,7 @@ def skip_prompt_length_check(self) -> bool:
         return False
 
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
-        """
-        Return the maximum supported number of items for each modality.
+        """Return the maximum supported number of items for each modality.
 
         A value of `None` means unlimited number of items.
 
@@ -486,8 +472,7 @@ def allowed_mm_limits(self) -> Mapping[str, int]:
         return allowed_limits
 
     def validate_num_items(self, modality: str, num_items: int) -> None:
-        """
-        Raise `ValueError` if the number of input items for the given modality
+        """Raise `ValueError` if the number of input items for the given modality
         is invalid.
         """
         supported_limit = self.supported_mm_limits.get(modality, 0)
@@ -512,11 +497,10 @@ def parse_mm_data(
         *,
         validate: bool = True,
     ) -> MultiModalDataItems:
-        """
-        Normalize [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict]
+        """Normalize [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict]
         to [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems]
         before passing them to
-        [`_get_hf_mm_data`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_data].
+        [`_get_hf_mm_inputs`][vllm.multimodal.processing.BaseMultiModalProcessor._get_hf_mm_inputs].
         """
         mm_items = self.data_parser.parse_mm_data(mm_data)
 
@@ -546,8 +530,7 @@ def get_mm_max_tokens_per_item(
         seq_len: int,
         mm_counts: Mapping[str, int],
     ) -> Mapping[str, int] | None:
-        """
-        Return the maximum number of tokens per item of for each modality.
+        """Return the maximum number of tokens per item of for each modality.
 
         When `None` (the default) is returned, vLLM will generate dummy inputs
         (images/videos) at maximum possible sizes and process them to determine
@@ -564,5 +547,6 @@ def get_mm_max_tokens_per_item(
             length and the maximum number of items of each modality allowed,
             and agree with dummy inputs (images/videos) at maximum possible
             sizes.
+
         """
         return None
diff --git a/vllm/multimodal/processing/dummy_inputs.py b/vllm/multimodal/processing/dummy_inputs.py
index 29f373ccc897..dd49b6c4fabc 100644
--- a/vllm/multimodal/processing/dummy_inputs.py
+++ b/vllm/multimodal/processing/dummy_inputs.py
@@ -10,15 +10,14 @@
 
 from vllm.config.multimodal import (
     AudioDummyOptions,
-    BaseDummyOptions,
     ImageDummyOptions,
+    MultiModalDummyOptions,
     VideoDummyOptions,
 )
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 
 from .context import BaseProcessingInfo
-from .inputs import ProcessorInputs
 
 _I = TypeVar("_I", bound=BaseProcessingInfo)
 
@@ -26,8 +25,7 @@
 
 
 class BaseDummyInputsBuilder(ABC, Generic[_I]):
-    """
-    Abstract base class that constructs the dummy data to profile
+    """Abstract base class that constructs the dummy data to profile
     multi-modal models.
     """
 
@@ -38,9 +36,7 @@ def __init__(self, info: _I) -> None:
 
     @abstractmethod
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
-        """
-        Build the text input corresponding to `mm_counts`.
-        """
+        """Build the text input corresponding to `mm_counts`."""
         raise NotImplementedError
 
     @abstractmethod
@@ -48,10 +44,9 @@ def get_dummy_mm_data(
         self,
         seq_len: int,
         mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
+        mm_options: MultiModalDummyOptions,
     ) -> MultiModalDataDict:
-        """
-        Build the multimodal input which, after processing, results in
+        """Build the multimodal input which, after processing, results in
         the maximum possible number of placeholder tokens.
 
         Args:
@@ -61,44 +56,9 @@ def get_dummy_mm_data(
                        If None, use model defaults for backward compatibility.
                        If provided, models can use these to customize dummy
                        data generation.
-        """
-        raise NotImplementedError
-
-    def get_dummy_processor_inputs(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions],
-    ) -> ProcessorInputs:
-        """
-        Build the input which, after processing, results in
-        the maximum possible number of placeholder tokens.
 
-        Args:
-            seq_len: Sequence length
-            mm_counts: Count of items per modality
-            mm_options: Configurable options per modality (optional)
         """
-        dummy_text = self.get_dummy_text(mm_counts)
-        dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
-        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False)
-
-        tokenizer = self.info.ctx.tokenizer
-        dummy_prompt: list[int]
-        if tokenizer is None:
-            # Tokenizer-less models (e.g. `skip_tokenizer_init=True`) only
-            # accept embeddings and have an empty dummy text, so there are no
-            # prompt tokens.
-            dummy_prompt = []
-        else:
-            from .processor import cached_encode
-
-            dummy_prompt = cached_encode(tokenizer, dummy_text, truncation=False)
-
-        return ProcessorInputs(
-            prompt=dummy_prompt,
-            mm_data_items=dummy_mm_items,
-        )
+        raise NotImplementedError
 
     def _get_dummy_audios(
         self,
@@ -127,9 +87,8 @@ def _get_dummy_images(
         width: int,
         height: int,
         num_images: int,
-        overrides: BaseDummyOptions | None = None,
+        overrides: ImageDummyOptions | None = None,
     ) -> list[Image.Image]:
-        assert overrides is None or isinstance(overrides, ImageDummyOptions)
         if num_images == 0:
             return []
         if overrides:
@@ -193,5 +152,5 @@ def _get_dummy_videos(
                         height,
                     )
                 height = min(height, overrides.height)
-        video = np.full((num_frames, width, height, 3), 255, dtype=np.uint8)
+        video = np.full((num_frames, height, width, 3), 255, dtype=np.uint8)
         return [video] * num_videos
diff --git a/vllm/multimodal/processing/inputs.py b/vllm/multimodal/processing/inputs.py
index 56750c3087cb..6e547c000434 100644
--- a/vllm/multimodal/processing/inputs.py
+++ b/vllm/multimodal/processing/inputs.py
@@ -6,6 +6,7 @@
 from vllm.config.multimodal import MMHasherAlgorithm
 from vllm.inputs import MultiModalHashes
 
+from ..cache import BaseMultiModalProcessorCache
 from ..hasher import MultiModalHasher
 from ..parse import MultiModalDataItems, MultiModalUUIDItems
 
@@ -18,8 +19,7 @@
 
 @dataclass
 class ProcessorInputs:
-    """
-    Represents the keyword arguments to
+    """Represents the keyword arguments to
     [`vllm.multimodal.processing.BaseMultiModalProcessor.apply`][].
     """
 
@@ -29,6 +29,8 @@ class ProcessorInputs:
     hf_processor_mm_kwargs: Mapping[str, object] = field(default_factory=dict)
     media_io_kwargs: Mapping[str, Mapping[str, object]] = field(default_factory=dict)
 
+    cache: BaseMultiModalProcessorCache | None = None
+
     def get_mm_hashes(
         self,
         model_id: str,
diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py
index 3db820458880..eb53023f0f35 100644
--- a/vllm/multimodal/processing/processor.py
+++ b/vllm/multimodal/processing/processor.py
@@ -18,6 +18,7 @@
 import torch
 from typing_extensions import TypeVar, assert_never
 
+from vllm.config import SchedulerConfig
 from vllm.inputs import (
     MultiModalEncDecInput,
     MultiModalHashes,
@@ -38,7 +39,7 @@
 )
 from ..parse import MultiModalDataItems, MultiModalUUIDItems
 from .context import BaseProcessingInfo, TimingContext
-from .dummy_inputs import BaseDummyInputsBuilder
+from .dummy_inputs import BaseDummyInputsBuilder, MultiModalDummyOptions
 from .inputs import ProcessorInputs
 
 if TYPE_CHECKING:
@@ -104,8 +105,7 @@ class PromptIndex:
 class PromptIndexTargets:
     @staticmethod
     def start() -> PromptIndex:
-        """
-        Resolves to the start of the prompt (before the first token).
+        """Resolves to the start of the prompt (before the first token).
 
         This results in a match even if the prompt is empty.
         """
@@ -116,9 +116,7 @@ def start() -> PromptIndex:
 
     @staticmethod
     def prefix(seq: list[int]) -> PromptIndex:
-        """
-        Resolves to a location in the prompt after the given prefix.
-        """
+        """Resolves to a location in the prompt after the given prefix."""
 
         def get_match_index(
             prompt: list[int],
@@ -146,8 +144,7 @@ def get_text_match_index(
 
     @staticmethod
     def end() -> PromptIndex:
-        """
-        Resolves to the end of the prompt (after the last token).
+        """Resolves to the end of the prompt (after the last token).
 
         This results in a match even if the prompt is empty.
         """
@@ -247,8 +244,7 @@ class UpdateMode(str, Enum):
 
 @dataclass
 class PromptUpdate(ABC):
-    """
-    Defines how to update a prompt with placeholder tokens.
+    """Defines how to update a prompt with placeholder tokens.
 
     Note:
         The target and content are token sequences. When converting text
@@ -256,6 +252,7 @@ class PromptUpdate(ABC):
         `add_special_tokens=False` to match how the prompt itself is
         tokenized. Otherwise, the updated prompt may contain duplicated
         special tokens or fail to match the target.
+
     """
 
     modality: str
@@ -294,8 +291,7 @@ def _resolve_content(self, item_idx: int) -> PromptUpdateDetails:
         return content
 
     def resolve(self, item_idx: int) -> "ResolvedPromptUpdate":
-        """
-        Given the index of the processed item within
+        """Given the index of the processed item within
         [`modality`][vllm.multimodal.processing.PromptUpdate.modality],
         output a copy of this object with its lazy attributes resolved.
         """
@@ -310,11 +306,9 @@ def resolve(self, item_idx: int) -> "ResolvedPromptUpdate":
 
 @dataclass
 class PromptInsertion(PromptUpdate):
-    """
-    Defines how to insert placeholder tokens into a prompt.
+    """Defines how to insert placeholder tokens into a prompt.
 
     Example:
-
     For each image, insert a number of `<image>` feature placeholders
     equal to the feature size of the vision encoder after the `<s>` token:
 
@@ -355,6 +349,7 @@ class PromptInsertion(PromptUpdate):
         insertion=[image_token_id] * image_feature_size,
     )
     ```
+
     """
 
     insertion: PromptUpdateContent = field(repr=False)
@@ -379,11 +374,9 @@ def mode(self) -> UpdateMode:
 
 @dataclass
 class PromptReplacement(PromptUpdate):
-    """
-    Defines how to replace portions of an input prompt with placeholder tokens.
+    """Defines how to replace portions of an input prompt with placeholder tokens.
 
     Example:
-
     For each image, replace one `<image>` input placeholder in the prompt
     with a number of `<image>` feature placeholders
     equal to the feature size of the vision encoder:
@@ -412,6 +405,7 @@ class PromptReplacement(PromptUpdate):
         ),
     )
     ```
+
     """
 
     replacement: PromptUpdateContent = field(repr=False)
@@ -447,8 +441,7 @@ def modality(self) -> str: ...
 
 
 def full_groupby_modality(values: Iterable[_M]) -> ItemsView[str, list[_M]]:
-    """
-    Convenience function to apply
+    """Convenience function to apply
     [`full_groupby`][vllm.utils.collection_utils.full_groupby]
     based on modality.
     """
@@ -462,8 +455,7 @@ class PromptTargetMatch(NamedTuple):
 
 @dataclass(frozen=True)
 class ResolvedPromptUpdate:
-    """
-    A [`PromptUpdate`][vllm.multimodal.processing.PromptUpdate] with its
+    """A [`PromptUpdate`][vllm.multimodal.processing.PromptUpdate] with its
     lazy attributes resolved.
     """
 
@@ -522,8 +514,7 @@ def iter_token_matches(
     *,
     start_idx: int = 0,
 ) -> Generator[_TokenMatch]:
-    """
-    Yield each occurrence of `match_ids` in `token_ids`.
+    """Yield each occurrence of `match_ids` in `token_ids`.
 
     Note that empty matches are ignored.
     """
@@ -562,8 +553,7 @@ def replace_token_matches(
     match_ids: list[int],
     new_ids: list[int],
 ) -> list[int]:
-    """
-    Replace each occurrence of `match_ids` in `token_ids`
+    """Replace each occurrence of `match_ids` in `token_ids`
     with `new_ids`.
 
     Note that empty matches are ignored.
@@ -846,8 +836,7 @@ def apply_token_matches(
     prompt: list[int],
     mm_prompt_updates: "MultiModalPromptUpdates",
 ) -> tuple[list[int], "MultiModalPromptUpdatesApplyResult"]:
-    """
-    Apply the updates in `mm_prompt_updates` to `prompt`.
+    """Apply the updates in `mm_prompt_updates` to `prompt`.
 
     Matches are exclusive even when multiple modalities share
     the same placeholder tokens. In that case, the modality that
@@ -914,8 +903,7 @@ def _iter_placeholders(
     prompt: list[int],
     mm_prompt_updates: "MultiModalPromptUpdates",
 ) -> Iterable[PlaceholderFeaturesInfo]:
-    """
-    Yield each set of placeholder tokens found in `prompt`.
+    """Yield each set of placeholder tokens found in `prompt`.
 
     Matches are exclusive even when multiple modalities share
     the same placeholder tokens. In that case, the modality that
@@ -1021,15 +1009,27 @@ def find_mm_placeholders(
 _I = TypeVar("_I", bound=BaseProcessingInfo)
 
 
-class MultiModalProcessingInfo(NamedTuple):
+class HFMultiModalInputs(NamedTuple):
+    hf_data: dict[str, object]
+    """The data inputs to the HF processor."""
+
+    hf_kwargs: Mapping[str, object]
+    """The non-data inputs to the HF processor."""
+
+    passthrough_data: dict[str, object]
+    """The inputs to assign to the output of the HF processor,
+    instead of passing them to the HF processor."""
+
+
+class MultiModalProcessingResult(NamedTuple):
+    prompt_ids: list[int]
     kwargs: MultiModalKwargsOptionalItems
     hashes: MultiModalHashes
     prompt_updates: MultiModalPromptUpdates
 
 
 class BaseMultiModalProcessor(ABC, Generic[_I]):
-    """
-    Abstract base class to process multi-modal inputs to be used in vLLM.
+    """Abstract base class to process multi-modal inputs to be used in vLLM.
 
     Not to be confused with `transformers.ProcessorMixin`.
     """
@@ -1040,8 +1040,6 @@ def __init__(
         self,
         info: _I,
         dummy_inputs: "BaseDummyInputsBuilder[_I]",
-        *,
-        cache: BaseMultiModalProcessorCache | None = None,
     ) -> None:
         super().__init__()
 
@@ -1054,16 +1052,85 @@ def __init__(
 
         self.info = info
         self.dummy_inputs = dummy_inputs
-        self.cache = cache
 
         self.data_parser = self.info.get_data_parser()
 
+    def get_dummy_inputs(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: MultiModalDummyOptions,
+    ) -> ProcessorInputs:
+        """Build the input which, after processing, results in
+        the maximum possible number of placeholder tokens.
+
+        Args:
+            seq_len: Sequence length
+            mm_counts: Count of items per modality
+            mm_options: Configurable options per modality (optional)
+
+        """
+        builder = self.dummy_inputs
+        dummy_text = builder.get_dummy_text(mm_counts)
+        dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options)
+        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False)
+
+        tokenizer = self.info.ctx.tokenizer
+        dummy_prompt: list[int]
+        if tokenizer is None:
+            # Tokenizer-less models (e.g. `skip_tokenizer_init=True`) only
+            # accept embeddings and have an empty dummy text, so there are no
+            # prompt tokens.
+            dummy_prompt = []
+        else:
+            dummy_prompt = cached_encode(tokenizer, dummy_text, truncation=False)
+
+        return ProcessorInputs(
+            prompt=dummy_prompt,
+            mm_data_items=dummy_mm_items,
+        )
+
+    def get_dummy_mm_inputs(
+        self,
+        mm_counts: Mapping[str, int],
+        *,
+        cache: BaseMultiModalProcessorCache | None = None,
+        scheduler_config: "SchedulerConfig | None" = None,
+    ) -> MultiModalInput:
+        """Create dummy data for profiling the memory usage of a model."""
+        model_config = self.info.ctx.model_config
+        seq_len = model_config.max_model_len
+        if scheduler_config is not None and scheduler_config.enable_chunked_prefill:
+            seq_len = min(seq_len, scheduler_config.max_num_batched_tokens)
+
+        mm_config = model_config.get_multimodal_config()
+
+        processor_inputs = self.get_dummy_inputs(
+            seq_len=seq_len,
+            mm_counts=mm_counts,
+            mm_options=mm_config.limit_per_prompt,
+        )
+        processor_inputs.cache = cache
+
+        mm_inputs = self.apply(
+            processor_inputs,
+            timing_ctx=TimingContext(enabled=False),
+        )
+
+        prompt_token_ids = mm_inputs["prompt_token_ids"]
+        total_len = len(prompt_token_ids)
+        if total_len < seq_len:
+            prompt_token_ids.extend([0] * (seq_len - total_len))
+
+        return mm_inputs
+
     def __call__(
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
         mm_uuid_items: MultiModalUUIDItems | None = None,
         hf_processor_mm_kwargs: Mapping[str, object] | None = None,
+        cache: BaseMultiModalProcessorCache | None = None,
     ) -> MultiModalInput:
         if isinstance(prompt, str):
             tokenizer = self.info.get_tokenizer()
@@ -1077,6 +1144,7 @@ def __call__(
             mm_items,
             mm_uuid_items,
             hf_processor_mm_kwargs=hf_processor_mm_kwargs or {},
+            cache=cache,
         )
 
         return self.apply(processor_inputs, TimingContext(enabled=False))
@@ -1097,8 +1165,7 @@ def _get_prompt_updates(
         hf_processor_mm_kwargs: Mapping[str, object],
         out_mm_kwargs: MultiModalKwargsItems,
     ) -> Sequence[PromptUpdate]:
-        """
-        Given the original multi-modal items for this modality
+        """Given the original multi-modal items for this modality
         and HF-processed data, output the updates to perform.
 
         The information returned by this method is used to update token inputs
@@ -1152,24 +1219,8 @@ def _find_mm_placeholders(
     ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
         return find_mm_placeholders(new_token_ids, mm_prompt_updates)
 
-    def _get_hf_mm_data(
-        self,
-        mm_items: MultiModalDataItems,
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        """Extract processor and passthrough data from multi-modal items."""
-        processor_data = dict[str, object]()
-        passthrough_data = dict[str, object]()
-
-        for items in mm_items.values():
-            processor_data.update(items.get_processor_data())
-            passthrough_data.update(items.get_passthrough_data())
-
-        return processor_data, passthrough_data
-
-    def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str | None:
-        """
-        Get the text to pass to the HF processor alongside the multi-modal
-        data.
+    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str | None:
+        """Get the text to pass to the HF processor alongside the multi-modal data.
 
         By default, no text is passed. If the HF processor requires that
         text and multi-modal items correspond to each other, you should
@@ -1178,31 +1229,53 @@ def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str | None:
         """
         return None
 
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        """
-        Pre-process the multi-modal data and HF processor keyword arguments
-        before they are passed to the HF processor.
-
-        By default, both are returned as-is. If the HF processor expects the
-        multi-modal data under different keys than those provided by the
-        multi-modal items (e.g. `audio` instead of `audios`), or requires
-        additional keyword arguments (e.g. `sampling_rate`), you should
-        override this method.
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> HFMultiModalInputs:
+        """Extract the multi-modal data and corresponding keyword arguments
+        for HF processing.
+
+        If the HF processor expects the multi-modal data under different keys
+        than those provided by the multi-modal items, or requires additional
+        keyword arguments, you should override this method.
         """
-        return mm_data, hf_processor_mm_kwargs
+        hf_data = dict[str, object]()
+        passthrough_data = dict[str, object]()
+
+        for items in mm_items.values():
+            if not items:
+                continue
+
+            hf_data.update(items.get_processor_data())
+            passthrough_data.update(items.get_passthrough_data())
+
+        if hf_data:
+            prompt_text = self._get_hf_mm_text(mm_items.get_all_counts())
+            if prompt_text is not None:
+                hf_data = dict(text=prompt_text, **hf_data)
+
+                # vLLM needs the untruncated sequence to keep placeholder
+                # tokens aligned. Note that the text inputs are just dummy
+                # text, not the original prompt. The original prompt is
+                # already tokenized by the renderer.
+                hf_kwargs = {"truncation": False, **hf_kwargs}
+
+        # HF processors accept (images, text, videos, audio)
+        # so we need to remap vLLM's "audios" -> HF "audio"
+        if "audios" in hf_data:
+            hf_data["audio"] = hf_data.pop("audios")
+
+        return HFMultiModalInputs(hf_data, hf_kwargs, passthrough_data)
 
     def _postprocess_hf_mm_data(
         self,
-        mm_data: Mapping[str, object],
+        mm_inputs: Mapping[str, object],
         hf_processor_mm_kwargs: Mapping[str, object],
         processed_data: BatchFeature,
     ) -> BatchFeature:
-        """
-        Post-process the output of the HF processor.
+        """Post-process the combined processor and passthrough data.
 
         By default, the output is returned as-is. If you need to modify the
         output of the HF processor before it is converted into multi-modal
@@ -1210,48 +1283,55 @@ def _postprocess_hf_mm_data(
         """
         return processed_data
 
-    def _apply_hf_processor_main(
+    def _call_hf_processor(
         self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        """
-        Apply the HF processor on the multi-modal data.
-        """
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        """Call the HF processor on the normalized multi-modal inputs."""
+        return self.info.ctx.call_hf_processor(
+            self.info.get_hf_processor(**hf_kwargs),
+            hf_data,
+            hf_kwargs,
         )
-        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
-        if processor_data:
-            processor_data, hf_processor_mm_kwargs = self._preprocess_hf_mm_data(
-                processor_data, hf_processor_mm_kwargs
-            )
+    def _finalize_hf_mm_data(
+        self,
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
+        passthrough_data: Mapping[str, object],
+        processed_data: BatchFeature | None = None,
+    ) -> BatchFeature:
+        """Merge passthrough fields and apply model-specific postprocessing."""
+        if processed_data is None:
+            from transformers.feature_extraction_utils import BatchFeature
 
-            prompt_text = self._get_hf_processor_text(mm_items.get_all_counts())
-            if prompt_text is not None:
-                processor_data = dict(text=prompt_text, **processor_data)
+            processed_data = BatchFeature()
 
-            processed_data = self.info.ctx.call_hf_processor(
-                self.info.get_hf_processor(**hf_processor_mm_kwargs),
-                processor_data,
-                hf_processor_mm_kwargs,
-            )
-            processed_data.update(passthrough_data)
-        else:
-            from transformers.feature_extraction_utils import BatchFeature
+        processed_data.update(passthrough_data)
+        return self._postprocess_hf_mm_data(hf_data, hf_kwargs, processed_data)
 
-            processed_data = BatchFeature(dict(passthrough_data))
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        """Apply the HF processor on the multi-modal data."""
+        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
+            mm_items, hf_kwargs
+        )
+
+        if hf_data:
+            processed_data = self._call_hf_processor(hf_data, hf_kwargs)
+        else:
+            processed_data = None
 
-        return self._postprocess_hf_mm_data(
-            processor_data,
-            hf_processor_mm_kwargs,
-            processed_data,
+        return self._finalize_hf_mm_data(
+            hf_data, hf_kwargs, passthrough_data, processed_data
         )
 
     def _postprocess_prompt(self, prompt: list[int]) -> list[int]:
-        """
-        Post-process the prompt token IDs before locating or applying
+        """Post-process the prompt token IDs before locating or applying
         multi-modal prompt updates.
 
         By default, the prompt is returned as-is. If the HF processor (or
@@ -1282,6 +1362,9 @@ def _get_cache_missing_items(
 
         mm_missing_data = {}
         for modality, idxs in mm_missing_idxs.items():
+            if not idxs:
+                continue
+
             missing_modality_data = []
             for idx in idxs:
                 data = mm_data_items[modality][idx]
@@ -1292,6 +1375,7 @@ def _get_cache_missing_items(
                     )
                 else:
                     missing_modality_data.append(data)
+
             mm_missing_data[modality] = missing_modality_data
 
         mm_missing_items = self.info.parse_mm_data(mm_missing_data, validate=False)
@@ -1303,8 +1387,7 @@ def _recompute_cached_prompt_update(
         cached_update: ResolvedPromptUpdate,
         new_item_idx: int,
     ) -> ResolvedPromptUpdate:
-        """
-        Override this if other attributes of `ResolvedPromptUpdate`
+        """Override this if other attributes of `ResolvedPromptUpdate`
         also need to be recomputed after retrieving from the cache.
         """
         return replace(cached_update, item_idx=new_item_idx)
@@ -1364,11 +1447,11 @@ def _apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
+    ) -> MultiModalProcessingResult:
         with timing_ctx.record("apply_hf_processor"):
             mm_processed_data = self._apply_hf_processor_main(
                 mm_items=inputs.mm_data_items,
-                hf_processor_mm_kwargs=inputs.hf_processor_mm_kwargs,
+                hf_kwargs=inputs.hf_processor_mm_kwargs,
             )
 
         mm_kwargs = MultiModalKwargsItems.from_hf_inputs(
@@ -1391,27 +1474,28 @@ def _apply_hf_processor(
             mm_kwargs,
         )
 
-        mm_info = MultiModalProcessingInfo(
+        return MultiModalProcessingResult(
+            prompt_ids=self._postprocess_prompt(inputs.prompt),
             kwargs=mm_kwargs,
             hashes=mm_hashes,
             prompt_updates=mm_prompt_updates,
         )
 
-        return mm_info
-
     def _cached_apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
-        """
-        Apply the HF processor on the full prompt text,
+    ) -> MultiModalProcessingResult:
+        """Apply the HF processor on the full prompt text,
         caching the results and reusing cached results.
         """
-        cache = self.cache
+        cache = inputs.cache
+        has_passthrough_data = any(
+            len(items.get_passthrough_data()) > 0
+            for items in inputs.mm_data_items.values()
+        )
 
-        _, passthrough_data = self._get_hf_mm_data(inputs.mm_data_items)
-        if cache is None or passthrough_data:
+        if cache is None or has_passthrough_data:
             return self._apply_hf_processor(inputs, timing_ctx)
 
         with timing_ctx.record("get_mm_hashes"):
@@ -1433,7 +1517,7 @@ def _cached_apply_hf_processor(
         with timing_ctx.record("apply_hf_processor"):
             mm_missing_processed_data = self._apply_hf_processor_main(
                 mm_items=mm_missing_data_items,
-                hf_processor_mm_kwargs=inputs.hf_processor_mm_kwargs,
+                hf_kwargs=inputs.hf_processor_mm_kwargs,
             )
 
         mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs(
@@ -1458,14 +1542,13 @@ def _cached_apply_hf_processor(
                 mm_missing_prompt_updates=mm_missing_prompt_updates,
             )
 
-        mm_info = MultiModalProcessingInfo(
+        return MultiModalProcessingResult(
+            prompt_ids=self._postprocess_prompt(inputs.prompt),
             kwargs=mm_kwargs,
             hashes=mm_hashes,
             prompt_updates=mm_prompt_updates,
         )
 
-        return mm_info
-
     def _apply_token_matches(
         self,
         prompt: list[int],
@@ -1522,8 +1605,7 @@ def _apply_prompt_updates_via_text(
         *,
         encode_segments_separately: bool = False,
     ) -> tuple[list[int], MultiModalPromptUpdatesApplyResult]:
-        """
-        Apply multi-modal prompt updates in text space, then encode the
+        """Apply multi-modal prompt updates in text space, then encode the
         result.
 
         A non-special-token target may be tokenized differently inside the
@@ -1694,17 +1776,15 @@ def _validate_mm_placeholders(
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsOptionalItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_res: MultiModalProcessingResult,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
         mm_item_counts = mm_items.get_all_counts()
-        self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
-        self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
+        self._validate_mm_kwargs(mm_res.kwargs, mm_item_counts)
+        self._validate_mm_updates(mm_res.prompt_updates, mm_item_counts)
 
         prompt_ids, mm_placeholders = self._apply_prompt_updates(
-            prompt_ids,
-            mm_prompt_updates,
+            mm_res.prompt_ids,
+            mm_res.prompt_updates,
         )
 
         self._validate_mm_placeholders(mm_placeholders, mm_item_counts)
@@ -1716,8 +1796,7 @@ def apply(
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
     ) -> MultiModalInput:
-        """
-        Process multi-modal inputs to be used in vLLM.
+        """Process multi-modal inputs to be used in vLLM.
 
         The main steps are:
 
@@ -1730,15 +1809,12 @@ def apply(
         4. Extract information about the placeholder tokens from the
            processed token IDs.
         """
-        prompt_ids = self._postprocess_prompt(inputs.prompt)
-        mm_info = self._cached_apply_hf_processor(inputs, timing_ctx)
+        mm_res = self._cached_apply_hf_processor(inputs, timing_ctx)
 
         with timing_ctx.record("apply_prompt_updates"):
             prompt_ids, mm_placeholders = self._maybe_apply_prompt_updates(
                 mm_items=inputs.mm_data_items,
-                prompt_ids=prompt_ids,
-                mm_kwargs=mm_info.kwargs,
-                mm_prompt_updates=mm_info.prompt_updates,
+                mm_res=mm_res,
             )
 
         mm_placeholder_ranges = {
@@ -1748,8 +1824,8 @@ def apply(
 
         return mm_input(
             prompt_token_ids=prompt_ids,
-            mm_kwargs=mm_info.kwargs,
-            mm_hashes=mm_info.hashes,
+            mm_kwargs=mm_res.kwargs,
+            mm_hashes=mm_res.hashes,
             mm_placeholders=mm_placeholder_ranges,
         )
 
@@ -1763,8 +1839,7 @@ def create_encoder_prompt(
         prompt: list[int],
         mm_items: MultiModalDataItems,
     ) -> list[int]:
-        """
-        Create input prompt for the encoder. HF processor will be applied on
+        """Create input prompt for the encoder. HF processor will be applied on
         this prompt during profiling and generation.
         """
         raise NotImplementedError
@@ -1805,8 +1880,7 @@ def apply(
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
     ) -> MultiModalEncDecInput:
-        """
-        Process multi-modal inputs to be used in vLLM.
+        """Process multi-modal inputs to be used in vLLM.
         The main processing steps are modified to fit encoder-decoder model:
         1. Create encoder prompt from input prompt text.
         2. Apply the HF processor on encoder prompt.
@@ -1822,6 +1896,7 @@ def apply(
             inputs.mm_uuid_items,
             hf_processor_mm_kwargs=inputs.hf_processor_mm_kwargs,
             media_io_kwargs=inputs.media_io_kwargs,
+            cache=inputs.cache,
         )
 
         encoder_inputs = super().apply(encoder_processor_inputs, timing_ctx)
diff --git a/vllm/multimodal/registry.py b/vllm/multimodal/registry.py
index 1d541b975136..28383979f917 100644
--- a/vllm/multimodal/registry.py
+++ b/vllm/multimodal/registry.py
@@ -2,24 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import threading
 from collections import defaultdict
-from collections.abc import Mapping
 from dataclasses import dataclass
-from multiprocessing.synchronize import Lock as LockType
-from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeVar, cast
+from typing import TYPE_CHECKING, Generic, Protocol, TypeVar, cast
 
-from vllm.inputs import MultiModalInput
 from vllm.logger import init_logger
 from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
 
-from .cache import (
-    BaseMultiModalProcessorCache,
-    BaseMultiModalReceiverCache,
-    MultiModalProcessorOnlyCache,
-    MultiModalProcessorSenderCache,
-    MultiModalReceiverCache,
-    ShmObjectStoreReceiverCache,
-    ShmObjectStoreSenderCache,
-)
 from .processing import (
     BaseDummyInputsBuilder,
     BaseMultiModalProcessor,
@@ -29,12 +17,7 @@
 )
 
 if TYPE_CHECKING:
-    from vllm.config import (
-        ModelConfig,
-        ObservabilityConfig,
-        SchedulerConfig,
-        VllmConfig,
-    )
+    from vllm.config import ModelConfig, ObservabilityConfig
     from vllm.model_executor.models.interfaces import SupportsMultiModal
 
 logger = init_logger(__name__)
@@ -45,8 +28,7 @@
 
 
 class ProcessingInfoFactory(Protocol[_I_co]):
-    """
-    Constructs a
+    """Constructs a
     [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor]
     instance from the context.
     """
@@ -58,8 +40,7 @@ def __call__(
 
 
 class DummyInputsBuilderFactory(Protocol[_I]):  # type: ignore[misc]
-    """
-    Constructs a
+    """Constructs a
     [`BaseDummyInputsBuilder`][vllm.multimodal.processing.BaseDummyInputsBuilder]
     instance from the context.
     """
@@ -68,8 +49,7 @@ def __call__(self, info: _I) -> BaseDummyInputsBuilder[_I]: ...
 
 
 class MultiModalProcessorFactory(Protocol[_I]):  # type: ignore[misc]
-    """
-    Constructs a
+    """Constructs a
     [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor]
     instance from the context.
     """
@@ -78,8 +58,6 @@ def __call__(
         self,
         info: _I,
         dummy_inputs: BaseDummyInputsBuilder[_I],
-        *,
-        cache: BaseMultiModalProcessorCache | None = None,
     ) -> BaseMultiModalProcessor[_I]: ...
 
 
@@ -92,64 +70,14 @@ class _ProcessorFactories(Generic[_I]):
     def build_processor(
         self,
         ctx: InputProcessingContext,
-        *,
-        cache: BaseMultiModalProcessorCache | None = None,
     ):
         info = self.info(ctx)
         dummy_inputs_builder = self.dummy_inputs(info)
-        return self.processor(info, dummy_inputs_builder, cache=cache)
+        return self.processor(info, dummy_inputs_builder)
 
 
 class MultiModalRegistry:
-    """
-    A registry that dispatches data processing according to the model.
-    """
-
-    def supports_multimodal_inputs(self, model_config: "ModelConfig") -> bool:
-        """
-        Checks if the model supports multimodal inputs.
-        Returns True if the model is multimodal with any non-zero supported
-        modalities, otherwise returns False, effectively running in
-        text-only mode.
-        """
-        if not model_config.is_multimodal_model:
-            return False
-
-        mm_config = model_config.get_multimodal_config()
-        try:
-            info = self._create_processing_info(model_config, tokenizer=None)
-        except ValueError:
-            # Speculative drafters for multimodal targets (e.g. Qwen3_5MTP,
-            # Exaone4_5_MTP, MiMoV2OmniMTP) declare `SupportsMultiModal` so
-            # that they can consume the embeddings merged by the target model,
-            # but they never run a multi-modal processor of their own. Running
-            # in text-only mode is the expected outcome for them, not a
-            # misconfiguration worth warning about.
-            if model_config.runner_type != "draft":
-                logger.warning_once(
-                    "Model %s is treated as multimodal but has no registered "
-                    "multimodal processor; running in text-only mode.",
-                    model_config.model,
-                )
-            return False
-
-        # Check if all supported modalities have limit == 0
-        if all(
-            mm_config.get_limit_per_prompt(modality) == 0
-            for modality in info.supported_mm_limits
-        ):
-            # If enable_mm_embeds is True, we still need MM infrastructure
-            # to process pre-computed embeddings even though encoder won't run
-            if mm_config.enable_mm_embeds:
-                return True
-
-            logger.info_once(
-                "All limits of multimodal modalities supported by the model "
-                "are set to 0, running in text-only mode."
-            )
-            return False
-
-        return True
+    """A registry that dispatches data processing according to the model."""
 
     def register_processor(
         self,
@@ -158,8 +86,7 @@ def register_processor(
         info: ProcessingInfoFactory[_I],
         dummy_inputs: DummyInputsBuilderFactory[_I],
     ):
-        """
-        Register a multi-modal processor to a model class. The processor
+        """Register a multi-modal processor to a model class. The processor
         is constructed lazily, hence a factory method should be passed.
 
         When the model receives multi-modal data, the provided function is
@@ -225,11 +152,8 @@ def create_processor(
         model_config: "ModelConfig",
         *,
         tokenizer: TokenizerLike | None = None,
-        cache: BaseMultiModalProcessorCache | None = None,
     ) -> BaseMultiModalProcessor[BaseProcessingInfo]:
-        """
-        Create a multi-modal processor for a specific model and tokenizer.
-        """
+        """Create a multi-modal processor for a specific model and tokenizer."""
         if not model_config.is_multimodal_model:
             model_name = model_config.served_model_name or model_config.model
             raise ValueError(f"{model_name} is not a multimodal model")
@@ -239,127 +163,7 @@ def create_processor(
 
         ctx = self._create_processing_ctx(model_config, tokenizer)
 
-        return factories.build_processor(ctx, cache=cache)
-
-    def get_dummy_mm_inputs(
-        self,
-        model_config: "ModelConfig",
-        mm_counts: Mapping[str, int],
-        *,
-        cache: BaseMultiModalProcessorCache | None = None,
-        processor: BaseMultiModalProcessor | None = None,
-        scheduler_config: "SchedulerConfig | None" = None,
-    ) -> MultiModalInput:
-        """
-        Create dummy data for profiling the memory usage of a model.
-
-        The model is identified by `model_config`.
-        """
-        seq_len = model_config.max_model_len
-        if scheduler_config is not None and scheduler_config.enable_chunked_prefill:
-            seq_len = min(seq_len, scheduler_config.max_num_batched_tokens)
-
-        if processor is None:
-            processor = self.create_processor(model_config, cache=cache)
-
-        mm_config = model_config.get_multimodal_config()
-        processor_inputs = processor.dummy_inputs.get_dummy_processor_inputs(
-            seq_len=seq_len,
-            mm_counts=mm_counts,
-            mm_options=mm_config.limit_per_prompt,
-        )
-        mm_inputs = processor.apply(
-            processor_inputs,
-            timing_ctx=TimingContext(enabled=False),
-        )
-
-        prompt_token_ids = mm_inputs["prompt_token_ids"]
-        total_len = len(prompt_token_ids)
-        if total_len < seq_len:
-            prompt_token_ids.extend([0] * (seq_len - total_len))
-
-        return mm_inputs
-
-    def _get_cache_type(
-        self,
-        vllm_config: "VllmConfig",
-    ) -> Literal[None, "processor_only", "lru", "shm"]:
-        model_config = vllm_config.model_config
-        if not self.supports_multimodal_inputs(model_config):
-            return None
-
-        # Check if the cache is disabled.
-        mm_config = model_config.get_multimodal_config()
-        if mm_config.mm_processor_cache_gb <= 0:
-            return None
-
-        # Check if IPC caching is supported.
-        parallel_config = vllm_config.parallel_config
-        is_ipc_supported = parallel_config._api_process_count == 1 and (
-            parallel_config.data_parallel_size == 1
-            or parallel_config.data_parallel_external_lb
-        )
-
-        if not is_ipc_supported:
-            return "processor_only"
-
-        mm_config = model_config.get_multimodal_config()
-        return mm_config.mm_processor_cache_type
-
-    def processor_cache_from_config(
-        self,
-        vllm_config: "VllmConfig",
-    ) -> BaseMultiModalProcessorCache | None:
-        """Return a `BaseMultiModalProcessorCache`, if enabled."""
-        cache_type = self._get_cache_type(vllm_config)
-        if cache_type is None:
-            return None
-        elif cache_type == "processor_only":
-            return MultiModalProcessorOnlyCache(vllm_config.model_config)
-        elif cache_type == "lru":
-            return MultiModalProcessorSenderCache(vllm_config.model_config)
-        elif cache_type == "shm":
-            return ShmObjectStoreSenderCache(vllm_config)
-        else:
-            raise ValueError(f"Unknown cache type: {cache_type!r}")
-
-    def processor_only_cache_from_config(
-        self,
-        vllm_config: "VllmConfig",
-    ) -> MultiModalProcessorOnlyCache | None:
-        """Return a `MultiModalProcessorOnlyCache`, if enabled."""
-        cache_type = self._get_cache_type(vllm_config)
-        if cache_type is None:
-            return None
-
-        return MultiModalProcessorOnlyCache(vllm_config.model_config)
-
-    def engine_receiver_cache_from_config(
-        self,
-        vllm_config: "VllmConfig",
-    ) -> BaseMultiModalReceiverCache | None:
-        """Return a `BaseMultiModalReceiverCache` for the engine process."""
-        cache_type = self._get_cache_type(vllm_config)
-        if cache_type in (None, "processor_only", "shm"):
-            return None
-        elif cache_type == "lru":
-            return MultiModalReceiverCache(vllm_config.model_config)
-        else:
-            raise ValueError(f"Unknown cache type: {cache_type!r}")
-
-    def worker_receiver_cache_from_config(
-        self,
-        vllm_config: "VllmConfig",
-        shared_worker_lock: LockType,
-    ) -> BaseMultiModalReceiverCache | None:
-        """Return a `BaseMultiModalReceiverCache` for the worker process."""
-        cache_type = self._get_cache_type(vllm_config)
-        if cache_type in (None, "processor_only", "lru"):
-            return None
-        elif cache_type == "shm":
-            return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock)
-        else:
-            raise ValueError(f"Unknown cache type: {cache_type!r}")
+        return factories.build_processor(ctx)
 
 
 class MultiModalTimingRegistry:
diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py
index 4de183e0dcec..3627dd9355ed 100644
--- a/vllm/multimodal/utils.py
+++ b/vllm/multimodal/utils.py
@@ -67,8 +67,7 @@ def encode_image_base64(
     image_mode: str | None = "RGB",
     format: str = "PNG",
 ) -> str:
-    """
-    Encode a pillow image to base64 format.
+    """Encode a pillow image to base64 format.
 
     By default, the image is converted into RGB format before being encoded.
     Pass `image_mode=None` to keep the original image mode.
@@ -83,8 +82,7 @@ def encode_image_url(
     image_mode: str | None = "RGB",
     format: str = "PNG",
 ) -> str:
-    """
-    Encode a pillow image as a data URL.
+    """Encode a pillow image as a data URL.
 
     By default, the image is converted into RGB format before being encoded.
     Pass `image_mode=None` to keep the original image mode.
@@ -147,14 +145,14 @@ def get_mm_features_in_window(
 def argsort_mm_positions(
     mm_positions: MultiModalPlaceholders,
 ) -> list[tuple[str, int]]:
-    """
-    Given a `MultiModalPlaceholders`, output a sequence of keys to
+    """Given a `MultiModalPlaceholders`, output a sequence of keys to
     sort the dictionary by `offset` (starting index in the input sequence)
     in ascending order.
 
     Returns:
         A list of `(modality, idx)`, which can be used to access an item
         by `mm_positions[modality][idx]`.
+
     """
     flat_items = (
         (modality, idx, item)
@@ -253,8 +251,7 @@ def group_and_batch_mm_items(
     device: torch.types.Device = None,
     pin_memory: bool = False,
 ) -> Generator[tuple[int, BatchedTensorInputs]]:
-    """
-    Group consecutive items (possibly from different requests) into batches.
+    """Group consecutive items (possibly from different requests) into batches.
 
     Items must be split across groups if any of the following occurs,
     as the batch would otherwise be invalid:
@@ -270,6 +267,7 @@ def group_and_batch_mm_items(
         A tuple `(num_items, grouped_kwargs)`, where:
         - `kwargs` is a dictionary of keyword arguments to pass to the model;
         - `num_items` is the corresponding number of items.
+
     """
     start_idx = 0
     for end_idx in range(1, len(items) + 1):
@@ -296,8 +294,7 @@ def group_and_batch_mm_kwargs(
     device: torch.types.Device = None,
     pin_memory: bool = False,
 ) -> Generator[tuple[str, int, BatchedTensorInputs], None, None]:
-    """
-    Group consecutive items (possibly from different requests) into batches.
+    """Group consecutive items (possibly from different requests) into batches.
 
     Items must be split across groups if any of the following occurs,
     as the batch would otherwise be invalid:
@@ -317,6 +314,7 @@ def group_and_batch_mm_kwargs(
         - `modality` is the modality of the batch;
         - `kwargs` is a dictionary of keyword arguments to pass to the model;
         - `num_items` is the corresponding number of items.
+
     """
     for modality, group in groupby(mm_kwargs, key=lambda x: x[0]):
         items_lst = [item for _, item in group]
@@ -333,14 +331,14 @@ def fetch_audio(
     audio_url: str,
     audio_io_kwargs: dict[str, Any] | None = None,
 ) -> tuple[np.ndarray, int | float]:
-    """
-    Args:
+    """Args:
         audio_url: URL of the audio file to fetch.
         audio_io_kwargs: Additional kwargs passed to handle audio IO.
 
     Warning:
         This method has direct access to local files and is only intended
         to be called by user code. Never call this from the online server!
+
     """
     media_io_kwargs = None if not audio_io_kwargs else {"audio": audio_io_kwargs}
     media_connector = MediaConnector(
@@ -354,14 +352,14 @@ def fetch_image(
     image_url: str,
     image_io_kwargs: dict[str, Any] | None = None,
 ) -> Image.Image:
-    """
-    Args:
+    """Args:
         image_url: URL of the image file to fetch.
         image_io_kwargs: Additional kwargs passed to handle image IO.
 
     Warning:
         This method has direct access to local files and is only intended
         to be called by user code. Never call this from the online server!
+
     """
     media_io_kwargs = None if not image_io_kwargs else {"image": image_io_kwargs}
     media_connector = MediaConnector(
@@ -375,14 +373,14 @@ def fetch_video(
     video_url: str,
     video_io_kwargs: dict[str, Any] | None = None,
 ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
-    """
-    Args:
+    """Args:
         video_url: URL of the video file to fetch.
         video_io_kwargs: Additional kwargs passed to handle video IO.
 
     Warning:
         This method has direct access to local files and is only intended
         to be called by user code. Never call this from the online server!
+
     """
     media_io_kwargs = None if not video_io_kwargs else {"video": video_io_kwargs}
     media_connector = MediaConnector(
diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py
index 854b5572b3ef..ff74183a55a7 100644
--- a/vllm/multimodal/video.py
+++ b/vllm/multimodal/video.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import math
 from abc import abstractmethod
-from typing import Any, ClassVar, cast
+from typing import Any, ClassVar, TypeAlias, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -28,6 +28,10 @@
 
 logger = init_logger(__name__)
 
+DecodedFrames: TypeAlias = npt.NDArray | torch.Tensor
+"""Decoded video frames: a host ``np.ndarray``, or a device ``torch.Tensor``
+when a GPU decoding codec is used (e.g. torchcodec with ``device="cuda"``)."""
+
 
 class VideoLoaderRegistry(ExtensionManager):
     def __init__(self) -> None:
@@ -164,8 +168,13 @@ def load_bytes(
         cls,
         data: bytes,
         **kwargs,
-    ) -> tuple[npt.NDArray, dict[str, Any]]:
-        """Load video frames from bytes and return (frames_array, metadata_dict)."""
+    ) -> tuple[DecodedFrames, dict[str, Any]]:
+        """Load video frames from bytes and return (frames, metadata_dict).
+
+        ``frames`` is a CPU ``np.ndarray`` unless a GPU decoding codec is
+        used (e.g. torchcodec with ``device="cuda"``), in which case it is
+        a ``torch.Tensor`` living on that device.
+        """
         raise NotImplementedError
 
     @classmethod
@@ -239,7 +248,7 @@ def load_bytes(
         *,
         backend: VideoDecoderBackend = "opencv",
         **kwargs: Any,
-    ) -> tuple[npt.NDArray, dict[str, Any]]:
+    ) -> tuple[DecodedFrames, dict[str, Any]]:
         """Load sampled frames from raw video bytes.
 
         Args:
@@ -265,6 +274,11 @@ def load_bytes(
                   creation at the cost of relying on the file's metadata. See
                   https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html
                   for details.
+                - ``device`` (TorchCodec): ``"cpu"`` (default) decodes on the
+                  host and returns a ``np.ndarray``; ``"cuda"`` decodes with
+                  NVDEC and returns the frames as a CUDA ``torch.Tensor``,
+                  which a device-side HF video processor can consume without
+                  a host round-trip.
                 - ``hw_decoders`` (PyNvVideoCodec): maximum number of
                   concurrent decoder slots. Defaults to 2 and must be a
                   positive integer.
@@ -272,7 +286,9 @@ def load_bytes(
                   size and pool acquisition timeout in seconds.
 
         Returns:
-            Tuple of ``(frames_array, metadata_dict)``.
+            Tuple of ``(frames, metadata_dict)``, where ``frames`` is a
+            CPU ``np.ndarray`` unless TorchCodec decodes on ``device="cuda"``.
+
         """
         target = VideoTargetMetadata(
             num_frames=num_frames, fps=fps, max_duration=max_duration
@@ -342,6 +358,9 @@ def load_bytes(
     video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"),
 )
 class Qwen3VLVideoBackend(VideoBackend):
+    _MAX_FRAMES: ClassVar[int] = 768
+    _MAX_FPS: ClassVar[int] = 30
+
     @classmethod
     def compute_frames_index_to_sample(
         cls,
@@ -351,10 +370,10 @@ def compute_frames_index_to_sample(
     ) -> list[int]:
         total_frames_num = source.total_frames_num
         original_fps = source.original_fps
-        fps = target.fps
+        fps = min(target.fps, cls._MAX_FPS)
         max_frame_idx = source.total_frames_num - 1
         min_frames = kwargs.get("min_frames", 4)
-        max_frames = kwargs.get("max_frames", 768)
+        max_frames = min(kwargs.get("max_frames", cls._MAX_FRAMES), cls._MAX_FRAMES)
 
         # Refer to:
         # https://github.com/huggingface/transformers/blob/v5.9.0/src/transformers/models/qwen3_vl/video_processing_qwen3_vl.py#L119-L125
@@ -405,6 +424,9 @@ class Qwen2VLVideoBackend(VideoBackend):
     clip); it is clamped to the last valid frame.
     """
 
+    _MAX_FRAMES: ClassVar[int] = 768
+    _MAX_FPS: ClassVar[int] = 30
+
     @classmethod
     def compute_frames_index_to_sample(
         cls,
@@ -418,7 +440,7 @@ def compute_frames_index_to_sample(
         original_fps = source.original_fps
         temporal_patch_size = kwargs.get("temporal_patch_size", 2)
         min_frames = kwargs.get("min_frames", 4)
-        max_frames = kwargs.get("max_frames", 768)
+        max_frames = min(kwargs.get("max_frames", cls._MAX_FRAMES), cls._MAX_FRAMES)
 
         # vLLM reports original_fps == 0 for clips with unknown/variable fps
         # (VFR, malformed, streaming); fail loudly instead of dividing by zero.
@@ -432,7 +454,7 @@ def compute_frames_index_to_sample(
             math.floor(min(max_frames, total_frames_num) / temporal_patch_size)
             * temporal_patch_size
         )
-        n = total_frames_num / original_fps * target.fps
+        n = total_frames_num / original_fps * min(target.fps, cls._MAX_FPS)
         n = min(max(n, min_frames), max_frames, total_frames_num)
         n = math.floor(n / temporal_patch_size) * temporal_patch_size
 
@@ -894,8 +916,7 @@ def get_candidate_target_fps(
         sampling_fps: float,
         max_fps: float = 8.0,
     ) -> list[float]:
-        """
-        Return the subset of `video_fps` factors that remain multiples
+        """Return the subset of `video_fps` factors that remain multiples
         of `sampling_fps`.
 
         Examples:
@@ -910,6 +931,7 @@ def get_candidate_target_fps(
                 ...
             ValueError: sampling_fps=2 must divide video_fps=5 to produce
                 consistent frame steps.
+
         """
         if sampling_fps is None:
             raise ValueError("sampling_fps must be provided")
@@ -946,8 +968,8 @@ def get_target_fps(
         frame_sample_mode: str,
         candidate_target_fps: list[float],
     ) -> float | None:
-        """
-        Get the target fps that best spans the videoand has the most frames sampled
+        """Get the target fps that best spans the video and samples the most
+        frames.
         """
         num_frames_sampled = 0
         selected_target_fps = None
@@ -1251,8 +1273,7 @@ def load_bytes(
         frame_recovery: bool = False,
         **kwargs,
     ) -> tuple[npt.NDArray, dict[str, Any]]:
-        """
-        Load video frames with dynamic sampling based on duration.
+        """Load video frames with dynamic sampling based on duration.
 
         Args:
             data: Raw video bytes
@@ -1263,6 +1284,7 @@ def load_bytes(
 
         Returns:
             Tuple of (frames_array, metadata_dict)
+
         """
         # recompute source metadata with adjusted duration to ensure correct
         # sampling indices computation
diff --git a/vllm/multimodal/video_decoders/__init__.py b/vllm/multimodal/video_decoders/__init__.py
index 1e6106e19937..9931f60f871d 100644
--- a/vllm/multimodal/video_decoders/__init__.py
+++ b/vllm/multimodal/video_decoders/__init__.py
@@ -19,6 +19,7 @@
     "torchcodec": {
         "num_ffmpeg_threads": 0,
         "seek_mode": "exact",
+        "device": "cpu",
     },
     PYNVVIDEOCODEC_VIDEO_BACKEND: {
         "hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
@@ -40,6 +41,17 @@ def _get_backend_option_defaults(backend: str) -> dict[str, Any]:
         ) from None
 
 
+# TODO(Isotr0py): Make kwargs structured.
+def incompatible_backend_options(
+    old_backend: str | None,
+    new_backend: str,
+) -> set[str]:
+    """Option names accepted by ``old_backend`` but not by ``new_backend``."""
+    old = _BACKEND_OPTION_DEFAULTS.get(old_backend or "", {})
+    new = _BACKEND_OPTION_DEFAULTS.get(new_backend, {})
+    return old.keys() - new.keys()
+
+
 def resolve_video_backend_kwargs(
     backend: str,
     kwargs: dict[str, Any],
@@ -103,5 +115,6 @@ def decode_video(
     "VideoTargetMetadata",
     "check_frame_pixel_limit",
     "decode_video",
+    "incompatible_backend_options",
     "resolve_video_backend_kwargs",
 ]
diff --git a/vllm/multimodal/video_decoders/opencv.py b/vllm/multimodal/video_decoders/opencv.py
index 2b86d17a28d9..3ca141c738c4 100644
--- a/vllm/multimodal/video_decoders/opencv.py
+++ b/vllm/multimodal/video_decoders/opencv.py
@@ -128,8 +128,7 @@ def _read_frames_with_recovery(
         frame_indices: list[int],
         total_frames: int,
     ) -> tuple[npt.NDArray, list[int], dict[int, int]]:
-        """
-        Read frames with dynamic window forward-scan recovery.
+        """Read frames with dynamic window forward-scan recovery.
 
         When a target frame fails to load, the next successfully grabbed
         frame (before the next target frame) will be used to recover it.
@@ -144,6 +143,7 @@ def _read_frames_with_recovery(
             - frames_array: Array of loaded frames
             - valid_frame_indices: List of frame indices that were loaded
             - recovered_map: Dict mapping recovered_idx -> source_idx
+
         """
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
diff --git a/vllm/multimodal/video_decoders/torchcodec.py b/vllm/multimodal/video_decoders/torchcodec.py
index f1f3290629a2..346f7b645b5b 100644
--- a/vllm/multimodal/video_decoders/torchcodec.py
+++ b/vllm/multimodal/video_decoders/torchcodec.py
@@ -5,7 +5,10 @@
 
 import numpy as np
 import numpy.typing as npt
+import torch
 
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available
 
 from .base import (
@@ -21,6 +24,8 @@
         "decoders.VideoDecoder"
     )
 
+logger = init_logger(__name__)
+
 
 def decode_torchcodec(
     loader_cls,
@@ -30,12 +35,14 @@ def decode_torchcodec(
     *,
     num_ffmpeg_threads: int = 0,
     seek_mode: Literal["exact", "approximate"] = "exact",
-) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
+    device: str = "cpu",
+) -> tuple[npt.NDArray | torch.Tensor, VideoSourceMetadata, list[int], list[int]]:
     check_torchcodec_available()
     decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder(
         data,
         num_ffmpeg_threads=num_ffmpeg_threads,
         seek_mode=seek_mode,
+        device=device,
     )
     check_frame_pixel_limit(
         decoder.metadata.width or 0,
@@ -48,7 +55,7 @@ def decode_torchcodec(
         source=source, target=target, **sampling_kwargs
     )
     frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames(
-        decoder, frame_idx
+        decoder, frame_idx, device=device
     )
     return frames, source, frame_idx, valid
 
@@ -67,7 +74,19 @@ def make_torchcodec_decoder(
         *,
         num_ffmpeg_threads: int = 0,
         seek_mode: Literal["exact", "approximate"] = "exact",
+        device: str = "cpu",
     ) -> "VideoDecoder":
+        torch_device = torch.device(device)
+        if torch_device.type == "cuda" and not current_platform.is_cuda():
+            raise ValueError(
+                f"torchcodec video decoding on device {device!r} requires "
+                "a CUDA-capable platform."
+            )
+        elif torch_device.type not in ("cpu", "cuda"):
+            raise ValueError(
+                f"torchcodec video decoding only supports 'cpu' and 'cuda' "
+                f"devices, got {device!r}."
+            )
         # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest
         # of the pipeline expects, avoiding a transpose.
         return VideoDecoder(
@@ -75,6 +94,7 @@ def make_torchcodec_decoder(
             dimension_order="NHWC",
             num_ffmpeg_threads=num_ffmpeg_threads,
             seek_mode=seek_mode,
+            device=device,
         )
 
     @staticmethod
@@ -91,10 +111,23 @@ def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata:
     def decode_torchcodec_frames(
         decoder: "VideoDecoder",
         frame_indices: list[int],
-    ) -> tuple[npt.NDArray, list[int]]:
+        *,
+        device: str = "cpu",
+    ) -> tuple[npt.NDArray | torch.Tensor, list[int]]:
         """Decode the requested indices in one batched, index-exact call."""
         if not frame_indices:
             return np.empty((0,), dtype=np.uint8), []
         # Note: torchcodec releases the GIL for the entire call
         batch = decoder.get_frames_at(frame_indices)
-        return batch.data.numpy(), list(frame_indices)
+        frames = batch.data
+        if (
+            torch.device(device).type != "cpu"
+            and (status := getattr(decoder, "cpu_fallback", None)) is not None
+            and bool(status)
+        ):
+            logger.warning_once(
+                "torchcodec could not use NVDEC for this video and "
+                "decoded on CPU instead; check codec support and "
+                "libnvcuvid."
+            )
+        return frames, list(frame_indices)
diff --git a/vllm/multimodal/video_prune/evs.py b/vllm/multimodal/video_prune/evs.py
index 8584c08299fc..56049b6e9d20 100644
--- a/vllm/multimodal/video_prune/evs.py
+++ b/vllm/multimodal/video_prune/evs.py
@@ -16,8 +16,7 @@
 def compute_retained_tokens_count(
     tokens_per_frame: int, num_frames: int, q: float
 ) -> int:
-    """
-    Compute the number of retained tokens for a given video.
+    """Compute the number of retained tokens for a given video.
     Method ensures that we retain all the tokens from the first frame
     regardless of the pruning rate.
 
@@ -28,6 +27,7 @@ def compute_retained_tokens_count(
 
     Returns:
         The number of retained tokens.
+
     """
     total_tokens = tokens_per_frame * num_frames
     evs_num_tokens = int(total_tokens * (1 - q))
@@ -41,8 +41,7 @@ def compute_retention_mask(
     spatial_merge_size: int,
     q: float,
 ) -> torch.Tensor:
-    """
-    Computes the retention mask for input video embeddings.
+    """Computes the retention mask for input video embeddings.
 
     Args:
         video_embeds (`torch.Tensor`): The input video embeddings
@@ -55,6 +54,7 @@ def compute_retention_mask(
     Returns:
         `torch.Tensor`: The retention mask for the video embeddings of
             `(T * H * W // spatial_merge_size ^ 2)` shape.
+
     """
     T, H, W = map(int, video_size_thw)
 
@@ -98,8 +98,7 @@ def compute_mrope_for_media(
     tokens_per_second: float = 1.0,
     video_second_per_grid: float = 1.0,
 ) -> torch.Tensor:
-    """
-    Computes the mrope for video embeddings based on the grid dimensions.
+    """Computes the mrope for video embeddings based on the grid dimensions.
     Computed mrope positions match original qwen 2.5 implementation,
     but positions are built for media being the first element in sequence.
 
@@ -113,6 +112,7 @@ def compute_mrope_for_media(
         Tensor of shape `(T * H * W, 4)` where last dimension
         represents mrope positions [0:3), while the last channel
         contains value of llm_grid_w repeated for all positions.
+
     """
     llm_grid_t = video_size_thw[0]
     llm_grid_h = video_size_thw[1] // spatial_merge_size
@@ -160,8 +160,7 @@ def recompute_mrope_positions(
     image_token_id: int,
     video_token_id: int,
 ) -> tuple[torch.LongTensor, int]:
-    """
-    Update part of input mrope positions.
+    """Update part of input mrope positions.
     Original mrope_positions are computed incorrectly, so once we prune media
     tokens we should reflect this in the mrope positions for the LLM.
 
@@ -204,8 +203,8 @@ def recompute_mrope_positions(
 
     Returns:
         Tuple of (mrope_positions, mrope_position_delta).
-    """
 
+    """
     # Tensors
     positions: torch.LongTensor = typing.cast(
         torch.LongTensor, mrope_positions.clone()
diff --git a/vllm/multimodal/video_prune/vidcom2.py b/vllm/multimodal/video_prune/vidcom2.py
index 9379f26ca74b..4f583473c9b2 100644
--- a/vllm/multimodal/video_prune/vidcom2.py
+++ b/vllm/multimodal/video_prune/vidcom2.py
@@ -48,6 +48,7 @@ def compute_retention_mask(
         Flat bool tensor of shape `(T*H*W/merge^2,)`, True for retained
         tokens. The True count equals `compute_retained_tokens_count` so
         placeholders sized at prompt-processing time match exactly.
+
     """
     T, H, W = map(int, video_size_thw)
     rows = H // spatial_merge_size
diff --git a/vllm/outputs.py b/vllm/outputs.py
index 84b9fccdc88d..e9a08f7c4628 100644
--- a/vllm/outputs.py
+++ b/vllm/outputs.py
@@ -53,6 +53,7 @@ class CompletionOutput:
             ``--per-request-spec-decode-metrics`` is enabled; None otherwise.
             Surfaced in the response as ``metrics.speculative_decoding`` for
             single-sequence (``n == 1``) requests.
+
     """
 
     index: int
@@ -90,6 +91,7 @@ class PoolingOutput:
 
     Args:
         data: The extracted hidden states.
+
     """
 
     data: torch.Tensor
@@ -128,6 +130,7 @@ class RequestOutput:
             prefix-cache writes for this request.
         kv_transfer_params: The params for remote K/V transfer.
         ec_transfer_params: The params for remote encoder-cache transfer.
+
     """
 
     def __init__(
@@ -172,7 +175,6 @@ def __init__(
 
     def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
         """Merge subsequent RequestOutput into this one"""
-
         self.finished |= next_output.finished
         self.kv_transfer_params = next_output.kv_transfer_params
         self.ec_transfer_params = next_output.ec_transfer_params
@@ -192,6 +194,10 @@ def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
                         completion.cumulative_logprob = (
                             next_completion.cumulative_logprob
                         )
+                        # R3 is returned on the terminal output and must survive
+                        # aggregation with earlier chunks that have no R3.
+                        if next_completion.routed_experts is not None:
+                            completion.routed_experts = next_completion.routed_experts
                         completion.finish_reason = next_completion.finish_reason
                         completion.stop_reason = next_completion.stop_reason
                     else:
@@ -232,8 +238,7 @@ def __repr__(self) -> str:
 
 
 class PoolingRequestOutput(Generic[_O]):
-    """
-    The output data of a pooling request to the LLM.
+    """The output data of a pooling request to the LLM.
 
     Args:
         request_id (str): A unique identifier for the pooling request.
@@ -241,6 +246,7 @@ class PoolingRequestOutput(Generic[_O]):
         prompt_token_ids (list[int]): A list of token IDs used in the prompt.
         num_cached_tokens: The number of tokens with prefix cache hit.
         finished (bool): A flag indicating whether the pooling is completed.
+
     """
 
     def __init__(
@@ -274,6 +280,7 @@ class EmbeddingOutput:
     Args:
         embedding: The embedding vector, which is a list of floats.
             Its length depends on the hidden dimension of the model.
+
     """
 
     embedding: list[float]
@@ -315,6 +322,7 @@ class ClassificationOutput:
     Args:
         probs: The probability vector, which is a list of floats.
             Its length depends on the number of classes.
+
     """
 
     probs: list[float]
@@ -356,6 +364,7 @@ class ScoringOutput:
 
     Args:
         score: The similarity score, which is a scalar value.
+
     """
 
     score: float
diff --git a/vllm/parser/AGENTS.md b/vllm/parser/AGENTS.md
new file mode 100644
index 000000000000..d9e6b79707a4
--- /dev/null
+++ b/vllm/parser/AGENTS.md
@@ -0,0 +1,41 @@
+# vllm/parser
+
+Model-output parsing for the serving layer. `Parser` (`abstract_parser.py`)
+unifies reasoning extraction and tool-call extraction behind one object; the
+serving layer calls `parse` / `parse_delta` and `is_reasoning_end` on it.
+
+## Layout
+
+- `abstract_parser.py` — `Parser`, and `DelegatingParser` for wrapping a legacy
+  `ReasoningParser` + `ToolParser` pair.
+- `parser_manager.py` — resolves parser names to classes.
+- `engine/` — the declarative parser engine. A `ParserEngineConfig`
+  (`parser_engine_config.py`) is a set of terminals, transitions and content
+  events; `parser_engine.py` and `streaming_parser_engine.py` drive it for the
+  non-streaming and streaming paths; `token_id_scanner.py` pre-lexes
+  special-token ids and `incremental_lexer.py` lexes text; `adapters.py`
+  exposes an engine through the legacy reasoning/tool parser interfaces.
+- `<model>.py` — one file per model family. Most subclass `ParserEngine`
+  (`deepseek_v4.py`, `qwen3.py`, ...); a few are `DelegatingParser` wrappers
+  over legacy parsers (`harmony.py`, `kimi_k3.py`).
+
+## Guidelines
+
+- A new parser either subclasses `ParserEngine` with a declarative config, or
+  subclasses `Parser` and implements `parse` / `parse_delta` directly. Use the
+  engine when the format is expressible as terminals and transitions, and
+  subclass an existing config when it is a variant of one (`seed_oss.py`,
+  `ling3.py`). `DelegatingParser` exists for backward compatibility with the
+  separate `ReasoningParser` / `ToolParser` classes; do not use it for new
+  formats.
+- Keep a model-specific fix inside that model's `<model>.py` where possible.
+  Change `engine/` or `abstract_parser.py` only when the behavior is wrong for
+  every parser, and say so in the PR.
+- Text emitted in a state with no `content_events` entry is discarded. Make
+  that deliberate when adding a state.
+- Special tokens reach the engine as token ids, not text. A text terminal that
+  spans a special token never matches in streaming; list it in
+  `token_id_terminals`.
+- Comment only where the code alone would leave a reader with a question —
+  a non-obvious transition, a deliberate swallow, a model quirk being worked
+  around. Keep it minimal and concise.
diff --git a/CLAUDE.md b/vllm/parser/CLAUDE.md
similarity index 100%
rename from CLAUDE.md
rename to vllm/parser/CLAUDE.md
diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py
index cf83bb996e63..0e84e710ed79 100644
--- a/vllm/parser/abstract_parser.py
+++ b/vllm/parser/abstract_parser.py
@@ -6,7 +6,6 @@
 from abc import abstractmethod
 from collections.abc import Sequence
 from dataclasses import dataclass, field
-from functools import cached_property
 
 from openai.types.responses import ToolChoiceFunction
 from pydantic import TypeAdapter, ValidationError
@@ -37,6 +36,7 @@
     extract_named_tool_call_streaming,
     extract_required_tool_call_streaming,
 )
+from vllm.tool_parsers.tool_strict_level import ToolStrictLevel
 
 logger = init_logger(__name__)
 
@@ -84,30 +84,20 @@ def commit(
 
 
 class Parser:
-    """
-    Abstract Parser class that unifies ReasoningParser and ToolParser into
-    a single interface for parsing model output.
-
-    This class provides a unified way to handle both reasoning extraction
-    (e.g., chain-of-thought content in <think> tags) and tool call extraction
-    (e.g., function calls in XML/JSON format) from model outputs.
-
-    Subclasses can either:
-    1. Override the abstract methods directly for custom parsing logic
-    2. Set `reasoning_parser` and `tool_parser` properties to delegate to
-       existing parser implementations
-
-    Class Attributes:
-        reasoning_parser_cls: The ReasoningParser class to use (for compatibility
-            with code that needs the class, not instance).
-        tool_parser_cls: The ToolParser class to use (for compatibility with
-            code that needs the class, not instance).
+    """Parse model output into reasoning, content and tool calls.
+
+    The serving layer holds one ``Parser`` per request and calls only the
+    members defined here. ``ParserEngine`` implements them over a single
+    declarative engine; ``DelegatingParser`` composes a legacy
+    ``ReasoningParser`` / ``ToolParser`` pair.
     """
 
     # Class-level parser classes for compatibility with existing patterns
     # Subclasses should override these if they use specific parser classes
     reasoning_parser_cls: type[ReasoningParser] | None = None
     tool_parser_cls: type[ToolParser] | None = None
+    # Server-side floor for tool-call structural tags (--tool-strict-level).
+    tool_strict_level: ToolStrictLevel = ToolStrictLevel.AUTO
 
     def __init__(
         self,
@@ -151,29 +141,16 @@ def __init__(
             engine_based=self._engine_based,
         )
 
-    @cached_property
-    def vocab(self) -> dict[str, int]:
-        """Get the vocabulary mapping from tokens to IDs."""
-        return self.model_tokenizer.get_vocab()
-
     @property
     def reasoning_parser(self) -> ReasoningParser | None:
         """The underlying reasoning parser, if any."""
         return self._reasoning_parser
 
-    @reasoning_parser.setter
-    def reasoning_parser(self, parser: ReasoningParser | None) -> None:
-        self._reasoning_parser = parser
-
     @property
     def tool_parser(self) -> ToolParser | None:
         """The underlying tool parser, if any."""
         return self._tool_parser
 
-    @tool_parser.setter
-    def tool_parser(self, parser: ToolParser | None) -> None:
-        self._tool_parser = parser
-
     def _initialize_history_tool_call_cnt(
         self,
         request: ChatCompletionRequest | ResponsesRequest,
@@ -187,120 +164,10 @@ def _initialize_history_tool_call_cnt(
         state.history_tool_call_cnt = count_history_tool_calls(request)
         state.history_tool_call_cnt_initialized = True
 
-    # ========== Reasoning Parser Methods ==========
-
-    @abstractmethod
-    def is_reasoning_end(self, input_ids: list[int]) -> bool:
-        """
-        Check if the reasoning content ends in the input_ids.
-
-        Used by structured engines like `xgrammar` to check if the
-        reasoning content ends in the model output.
-
-        Args:
-            input_ids: The token IDs of the model output.
-
-        Returns:
-            True if the reasoning content ends in the input_ids.
-        """
-
-    def is_reasoning_end_streaming(
-        self, input_ids: list[int], delta_ids: list[int]
-    ) -> bool:
-        """
-        Check if the reasoning content ends during a decode step.
-
-        Args:
-            input_ids: The entire model output token IDs.
-            delta_ids: The last few computed tokens at the current decode step.
-
-        Returns:
-            True if the reasoning content ends in the delta_ids.
-        """
-        return self.is_reasoning_end(input_ids)
-
-    def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None:
-        """
-        Locate the token that ends reasoning within one window of tokens.
-
-        Args:
-            token_ids: The tokens to examine.
-
-        Returns:
-            The offset within ``token_ids`` of the last reasoning token, or
-            ``None`` when reasoning does not end in this window. Parsers that
-            cannot answer from a window alone always return ``None``; callers
-            fall back to :meth:`is_reasoning_end_streaming`.
-        """
-        return None
-
-    @abstractmethod
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        """
-        Extract content token IDs from the input_ids.
-
-        This extracts the non-reasoning content (e.g., everything after
-        the </think> tag).
-
-        Args:
-            input_ids: The token IDs of the model output.
-
-        Returns:
-            The extracted content token IDs.
-        """
-
-    @abstractmethod
-    def extract_reasoning(
-        self,
-        model_output: str,
-        request: ChatCompletionRequest | ResponsesRequest,
-    ) -> tuple[str | None, str | None]:
-        """
-        Extract reasoning content from a complete model-generated string.
-
-        Used for non-streaming responses where we have the entire model
-        response available before sending to the client.
-
-        Args:
-            model_output: The complete model-generated string.
-            request: The request object used to generate the output.
-
-        Returns:
-            A tuple of (reasoning, response_content).
-        """
-
-    @abstractmethod
-    def extract_reasoning_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-    ) -> DeltaMessage | None:
-        """
-        Extract reasoning content from a streaming delta message.
-
-        Args:
-            previous_text: Text from all previous tokens.
-            current_text: Text including the current delta.
-            delta_text: The new text in this delta.
-            previous_token_ids: Token IDs from previous generation.
-            current_token_ids: All token IDs including current.
-            delta_token_ids: The new token IDs in this delta.
-
-        Returns:
-            A DeltaMessage with reasoning and/or content fields, or None.
-        """
-
-    # ========== Tool Parser Methods ==========
-
     def adjust_request(
         self, request: ChatCompletionRequest | ResponsesRequest
     ) -> ChatCompletionRequest | ResponsesRequest:
-        """
-        Adjust the request parameters for tool calling.
+        """Adjust the request parameters for tool calling.
 
         Can be overridden by subclasses to modify request parameters
         (e.g., setting structured output schemas for tool calling).
@@ -310,53 +177,23 @@ def adjust_request(
 
         Returns:
             The adjusted request.
+
         """
         return request
 
     @abstractmethod
-    def extract_tool_calls(
-        self,
-        model_output: str,
-        request: ChatCompletionRequest | ResponsesRequest,
-    ) -> ExtractedToolCallInformation:
-        """
-        Extract tool calls from a complete model-generated string.
+    def is_reasoning_end(self, input_ids: list[int]) -> bool:
+        """Check if the reasoning content ends in the input_ids.
 
-        Used for non-streaming responses.
+        Called with the rendered prompt to decide whether generation starts
+        after reasoning. Must be a pure function of the input_ids.
 
         Args:
-            model_output: The complete model-generated string.
-            request: The request object used to generate the output.
+            input_ids: The token IDs of the model output.
 
         Returns:
-            ExtractedToolCallInformation containing the tool calls.
-        """
-
-    @abstractmethod
-    def extract_tool_calls_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-        request: ChatCompletionRequest | ResponsesRequest,
-    ) -> DeltaMessage | None:
-        """
-        Extract tool calls from a streaming delta message.
-
-        Args:
-            previous_text: Text from all previous tokens.
-            current_text: Text including the current delta.
-            delta_text: The new text in this delta.
-            previous_token_ids: Token IDs from previous generation.
-            current_token_ids: All token IDs including current.
-            delta_token_ids: The new token IDs in this delta.
-            request: The request object.
+            True if the reasoning content ends in the input_ids.
 
-        Returns:
-            A DeltaMessage with tool_calls field, or None.
         """
 
     @abstractmethod
@@ -377,6 +214,7 @@ def parse(
 
         Returns:
             A tuple of (reasoning, content, tool_calls).
+
         """
 
     @abstractmethod
@@ -399,8 +237,7 @@ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
 
 
 class DelegatingParser(Parser):
-    """
-    A Parser implementation that delegates to separate ReasoningParser and
+    """A Parser implementation that delegates to separate ReasoningParser and
     ToolParser instances.
 
     This is the recommended base class for creating model-specific parsers
@@ -581,6 +418,7 @@ def _apply_structural_tag(
         structure_tag = self._tool_parser.get_structural_tag(
             request,
             reasoning=False,
+            strict_level=self.tool_strict_level,
         )
         if structure_tag is None:
             return request
@@ -761,14 +599,14 @@ def is_reasoning_end(self, input_ids: list[int]) -> bool:
             return False
         return self._reasoning_parser.is_reasoning_end(input_ids)
 
-    def is_reasoning_end_streaming(
+    def _is_reasoning_end_streaming(
         self, input_ids: list[int], delta_ids: list[int]
     ) -> bool:
         if self._reasoning_parser is None:
             return False
         return self._reasoning_parser.is_reasoning_end_streaming(input_ids, delta_ids)
 
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
+    def _extract_content_ids(self, input_ids: list[int]) -> list[int]:
         if self._reasoning_parser is None:
             return input_ids
         return self._reasoning_parser.extract_content_ids(input_ids)
@@ -882,13 +720,13 @@ def parse_delta(
                     reasoning_parser.has_engine_confirmed_reasoning_end()
                 )
             else:
-                should_transition = self.is_reasoning_end_streaming(
+                should_transition = self._is_reasoning_end_streaming(
                     current_token_ids, delta_token_ids
                 )
             if should_transition:
                 state.reasoning_ended = True
                 reasoning_transitioned = True
-                current_token_ids = self.extract_content_ids(delta_token_ids)
+                current_token_ids = self._extract_content_ids(delta_token_ids)
                 # Flush whenever the reasoning parser is engine-based (not only
                 # when _engine_based is True): it buffers the post-marker text
                 # (e.g. the "<" of "<tool_call>"), surfaced via finish_streaming().
diff --git a/vllm/parser/cohere_command.py b/vllm/parser/cohere_command.py
new file mode 100644
index 000000000000..e5ed7aed33ae
--- /dev/null
+++ b/vllm/parser/cohere_command.py
@@ -0,0 +1,683 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cohere Command parser.
+
+Uses ``cohere_melody`` as a stateful parser filter.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, TypeGuard
+
+import regex as re
+import xgrammar as xgr
+
+try:
+    from cohere_melody import (
+        AccumulatedToolCall,
+        FilterAggregatedResult,
+        PyFilter,
+        PyFilterOptions,
+    )
+except ImportError as e:
+    raise ImportError(
+        "The Cohere Command parser requires the `cohere_melody` package. "
+        "Install it with:\n    pip install 'cohere-melody>=0.14.0'"
+    ) from e
+
+from vllm.entrypoints.cohere.cohere_chat_message import (
+    Citation,
+    CitationSource,
+    CohereDeltaMessage,
+)
+from vllm.entrypoints.generate.base.protocol import (
+    AnyResponseFormat,
+    DeltaFunctionCall,
+    DeltaMessage,
+    DeltaToolCall,
+    FunctionCall,
+)
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.parser.abstract_parser import DelegatingParser
+from vllm.reasoning.cohere_command_reasoning_parser import (
+    BaseCohereCommandReasoningParser,
+)
+from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
+from vllm.sampling_params import StructuredOutputsParams
+from vllm.tool_parsers.cohere_command_tool_parser import BaseCohereCommandToolParser
+
+if TYPE_CHECKING:
+    # Declared in melody's stub only; not exported at runtime.
+    from cohere_melody import FilterCitation, Source
+
+    from vllm.config import ModelConfig
+    from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+    from vllm.tokenizers import TokenizerLike
+    from vllm.tool_parsers.utils import Tool
+
+PositionToSource = Mapping[tuple[int, int], CitationSource]
+
+
+class CohereTagRegistry(NamedTuple):
+    """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``)."""
+
+    trigger: str
+    end: str
+
+
+class CohereTagStyle(NamedTuple):
+    """The structural tags style for a given model architecture.
+
+    ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses
+    both response and text delimiters). ``tools`` is the tool-call wrapper.
+    """
+
+    json_tags: tuple[CohereTagRegistry, ...]
+    tools: CohereTagRegistry
+
+
+class CohereNormalizedTool(TypedDict):
+    """A tool definition normalized to the shape ``collect_tool_schema`` expects.
+
+    ``parameters`` is a JSON Schema object (possibly empty) describing the tool's
+    call signature.
+    """
+
+    name: str
+    parameters: dict[str, Any]
+
+
+COMMAND_A_TOOLS_TAG = CohereTagRegistry(
+    trigger="<|START_ACTION|>",
+    end="<|END_ACTION|>",
+)
+COMMAND_A_JSON_TAG = CohereTagRegistry(
+    trigger="<|START_RESPONSE|>",
+    end="<|END_RESPONSE|>",
+)
+COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry(
+    trigger="<|START_TEXT|>",
+    end="<|END_TEXT|>",
+)
+
+MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = {
+    "Cohere2ForCausalLM": CohereTagStyle(
+        json_tags=(COMMAND_A_JSON_TAG,),
+        tools=COMMAND_A_TOOLS_TAG,
+    ),
+    "Cohere2VisionForConditionalGeneration": CohereTagStyle(
+        json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
+        tools=COMMAND_A_TOOLS_TAG,
+    ),
+    "Cohere2MoeForCausalLM": CohereTagStyle(
+        json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
+        tools=COMMAND_A_TOOLS_TAG,
+    ),
+}
+
+
+def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str:
+    """Build an xgrammar EBNF grammar that matches a JSON array of tool calls.
+
+    The grammar shape is architecture-independent; callers are responsible for
+    wrapping it in the correct structural tag (see ``CohereTagStyle.tools``).
+    """
+    tool_dictionary: dict[str, str] = {}
+    for tool in tool_schema:
+        tool_name = tool["name"]
+        tool_parameters = json.dumps(tool["parameters"])
+        json_schema = f"""{{
+                        "type": "object",
+                        "properties": {{
+                            "tool_call_id": {{
+                                "type": "string",
+                                "pattern": "^[0-9]+$"
+                            }},
+                            "tool_name": {{
+                                "type": "string",
+                                "const": "{tool_name}"
+                            }},
+                            "parameters": {tool_parameters}
+                            }}
+                            }}"""
+        tool_grammar = str(xgr.Grammar.from_json_schema(json_schema))
+        for match in re.findall(r"\b(\w+)\s*::=", tool_grammar):
+            tool_grammar = re.sub(
+                rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar
+            )
+        tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}"
+    # Emitted grammar shape:
+    #   root  ::= tools
+    #   tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws
+    #   ws    ::= (" " | "\t" | "\n")*
+    #   tool  ::= <tool_a> | <tool_b> | ...         (one alternative per input)
+    #   <tool_x>     ::= <tool_x>root               (per-tool xgrammar rules)
+    #   <tool_x>root ::= ...                        (from xgr.Grammar.from_json_schema)
+    tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys())
+    tool_rules = "\n    ".join(tool_dictionary.values())
+    grammar = f"""root ::= tools
+    tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws
+    ws    ::= (" " | "\\t" | "\\n")*
+    {tool_alternatives}
+    {tool_rules}
+    """
+    return grammar
+
+
+def _tool_definitions_to_schema_list(
+    tools: str | list[Any],
+) -> list[CohereNormalizedTool]:
+    """Build the list of ``CohereNormalizedTool`` dicts expected by
+    ``collect_tool_schema``.
+
+    Accepts:
+    - JSON string
+    - list of dicts with top-level ``name`` / ``parameters``
+    - list of Chat Completions-style ``{"type": "function", "function": {...}}``
+    - list of Pydantic models with ``model_dump()``
+    """
+    if isinstance(tools, str):
+        try:
+            parsed = json.loads(tools)
+        except json.JSONDecodeError:
+            return []
+        if not isinstance(parsed, list):
+            return []
+    else:
+        parsed = list(tools)
+
+    out: list[CohereNormalizedTool] = []
+    for raw in parsed:
+        t = raw.model_dump() if hasattr(raw, "model_dump") else raw
+        if not isinstance(t, dict):
+            continue
+        # Unwrap Chat Completions' ``{"type": "function", "function": {...}}``
+        # shape; otherwise take the dict as-is.
+        if t.get("type") == "function" and isinstance(t.get("function"), dict):
+            t = t["function"]
+        name = t.get("name")
+        if not isinstance(name, str):
+            continue
+        params = t.get("parameters")
+        out.append(
+            CohereNormalizedTool(
+                name=name,
+                parameters=params if isinstance(params, dict) else {},
+            )
+        )
+    return out
+
+
+def _has_effective_tools(
+    tools: str | list[Any] | None,
+) -> TypeGuard[str | list[Any]]:
+    """True when ``tools`` contains at least one tool definition to convert.
+
+    ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest``
+    uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are
+    treated as effective only when non-blank.
+    """
+    if tools is None:
+        return False
+    if isinstance(tools, str):
+        return bool(tools.strip())
+    return len(tools) > 0
+
+
+# Builder: produces vLLM response_format in xgrammar's canonical format.
+# See xgrammar docs: type "structural_tag" with "format" = triggered_tags
+# and tag content type = json_schema | grammar.
+def convert_schema_to_structural_tags(
+    schema: dict | None = None,
+    tools: str | list[Any] | None = None,
+    model_architecture: str | None = None,
+) -> str | None:
+    """Returns a response_format string accepted by xgrammar's structural tag format.
+    Uses the canonical shape: {"type": "structural_tag", "format": {...}} with
+    format.type "triggered_tags" and tag content type "json_schema" or "grammar".
+
+    Callers that are not on an engine path (e.g. the reasoning parser) must pass
+    ``model_architecture`` explicitly.
+    """
+    if model_architecture is None or model_architecture not in MODEL_TO_TAG_STYLE:
+        return None
+    style = MODEL_TO_TAG_STYLE[model_architecture]
+
+    tags: list[dict] = []
+    triggers: list[str] = []
+
+    def _add_tag(tag: CohereTagRegistry, content: dict) -> None:
+        tags.append({"begin": tag.trigger, "content": content, "end": tag.end})
+        triggers.append(tag.trigger)
+
+    if schema is not None:
+        # One structural tag per JSON wrapper (e.g. MOE: response + text).
+        # Same for schema-only and "tools plus JSON mode" (North: schema when
+        # the model does not call tools).
+        for jt in style.json_tags:
+            _add_tag(jt, {"type": "json_schema", "json_schema": schema})
+
+    if _has_effective_tools(tools):
+        # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS)
+        # or a list (Chat Completions ``request.tools`` as Pydantic models or dicts).
+        tool_schema_list = _tool_definitions_to_schema_list(tools)
+        if not tool_schema_list:
+            raise ValueError(
+                "No valid tool definitions could be parsed from the request for "
+                "structural tag conversion."
+            )
+        tool_grammar = collect_tool_schema(tool_schema_list)
+        _add_tag(style.tools, {"type": "grammar", "grammar": tool_grammar})
+
+    if not tags:
+        return None
+    return json.dumps(
+        {
+            "type": "structural_tag",
+            "format": {
+                "type": "triggered_tags",
+                "triggers": triggers,
+                "tags": tags,
+            },
+        }
+    )
+
+
+def _response_format_type(
+    response_format: AnyResponseFormat | dict | None,
+) -> str | None:
+    if response_format is None:
+        return None
+    if isinstance(response_format, dict):
+        t = response_format.get("type")
+        return t if isinstance(t, str) else None
+    return response_format.type
+
+
+def _maybe_parse_json_dict(value: Any) -> dict | None:
+    """If value is a JSON string, parse to dict; otherwise require dict."""
+    if isinstance(value, dict):
+        return value
+    if isinstance(value, str):
+        try:
+            parsed = json.loads(value)
+        except (TypeError, json.JSONDecodeError):
+            return None
+        return parsed if isinstance(parsed, dict) else None
+    return None
+
+
+def _unwrap_nested_schema(candidate: Any) -> dict | None:
+    """Return ``candidate`` as a dict, unwrapping a nested ``schema`` if present.
+
+    Returns ``None`` if ``candidate`` is not (and cannot be parsed into) a dict.
+    """
+    cand = _maybe_parse_json_dict(candidate)
+    if not isinstance(cand, dict):
+        return None
+    nested = cand.get("schema")
+    return nested if isinstance(nested, dict) else cand
+
+
+def _schema_from_json_schema_field(js_wr: Any) -> dict | None:
+    """Extract the JSON Schema object from Chat Completions ``json_schema`` payload.
+
+    Accepts:
+    - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field
+    - dict in OpenAI shape ``{"name": ..., "schema": {...}}``
+    - dict with ``json_schema`` key holding either the schema or a nested wrapper
+    - dict that is already a JSON Schema document (some clients omit the wrapper)
+    - JSON strings for any of the above
+    """
+    if js_wr is None:
+        return None
+
+    parsed_wr = _maybe_parse_json_dict(js_wr)
+    if parsed_wr is not None:
+        js_wr = parsed_wr
+
+    if hasattr(js_wr, "model_dump"):
+        for by_alias in (True, False):
+            try:
+                data = js_wr.model_dump(by_alias=by_alias, exclude_none=False)
+            except TypeError:
+                data = js_wr.model_dump(by_alias=by_alias)
+            out = _unwrap_nested_schema(data.get("schema") or data.get("json_schema"))
+            if out is not None:
+                return out
+        inner_attr = getattr(js_wr, "json_schema", None)
+        return inner_attr if isinstance(inner_attr, dict) else None
+
+    if isinstance(js_wr, dict):
+        for key in ("schema", "json_schema"):
+            out = _unwrap_nested_schema(js_wr.get(key))
+            if out is not None:
+                return out
+        return js_wr
+
+    return None
+
+
+def _schema_dict_from_chat_response_format(
+    rf: AnyResponseFormat | dict | None,
+) -> dict | None:
+    """JSON schema dict from Chat Completions ``request.response_format`` only."""
+    if rf is None:
+        return None
+    rf_type = _response_format_type(rf)
+    if rf_type == "json_object":
+        return {"type": "object"}
+    if rf_type != "json_schema":
+        return None
+    js_wr = (
+        rf.get("json_schema")
+        if isinstance(rf, dict)
+        else getattr(rf, "json_schema", None)
+    )
+    return _schema_from_json_schema_field(js_wr)
+
+
+def _schema_dict_from_structured_outputs(
+    so: StructuredOutputsParams | None,
+) -> dict | None:
+    """Schema dict from ``structured_outputs`` (``json`` / ``json_object``).
+
+    Same unwrapping as ``json_schema``. ``json`` is expected to be ``str`` or
+    ``dict`` (enforced by ``StructuredOutputsParams`` / request models); other
+    types raise ``ValueError`` only if a caller bypasses that validation.
+    """
+    if so is None:
+        return None
+    if so.json_object:
+        return {"type": "object"}
+    raw: Any = so.json
+    if raw is None:
+        return None
+
+    if hasattr(raw, "model_dump"):
+        out = _schema_from_json_schema_field(raw)
+        if out is None:
+            raise ValueError(
+                "structured_outputs.json model has no extractable JSON Schema."
+            )
+        return out
+
+    if isinstance(raw, str):
+        if not raw.strip():
+            raise ValueError("structured_outputs.json cannot be empty.")
+        try:
+            raw = json.loads(raw)
+        except json.JSONDecodeError as e:
+            raise ValueError("structured_outputs.json must be valid JSON.") from e
+        if not isinstance(raw, dict):
+            raise ValueError("structured_outputs.json must decode to a JSON object.")
+
+    if isinstance(raw, Mapping):
+        body = raw if isinstance(raw, dict) else dict(raw)
+        return _schema_from_json_schema_field(body) or body
+
+    raise ValueError(
+        f"structured_outputs.json has unsupported type {type(raw).__name__}."
+    )
+
+
+def _melody_sources_to_vllm(
+    sources: list[Source], position_to_source: PositionToSource | None
+) -> list[CitationSource]:
+    """Resolve melody ``Source`` addresses into :class:`CitationSource` objects.
+
+    Melody reports each source as ``(tool_call_index, tool_result_indices)``.
+    ``position_to_source`` is built by
+    :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source`
+    and forwarded through ``chat_template_kwargs``; each index fans out to one
+    resolved source, and unresolved positions are skipped. Without a map (a
+    request that did not come through the Cohere v2 handler) every source is
+    dropped and the serving layer discards the citation at wire-coercion time.
+    """
+    if not position_to_source:
+        return []
+    out: list[CitationSource] = []
+    for source in sources:
+        for idx in source.tool_result_indices:
+            info = position_to_source.get((source.tool_call_index, idx))
+            if info is not None:
+                out.append(info)
+    return out
+
+
+def _melody_citations_to_vllm(
+    citations: list[FilterCitation], position_to_source: PositionToSource | None
+) -> list[Citation] | None:
+    """Convert melody ``FilterCitation`` objects into :class:`Citation`.
+
+    Citations whose sources all fail to resolve are still returned with
+    ``sources=[]``; the serving layer drops them so the fail-closed policy
+    lives in one place.
+    """
+    if not citations:
+        return None
+    return [
+        Citation(
+            start=c.start_index,
+            end=c.end_index,
+            text=c.text,
+            sources=_melody_sources_to_vllm(c.sources, position_to_source),
+            type="THINKING_CONTENT" if c.is_thinking else "TEXT_CONTENT",
+        )
+        for c in citations
+    ]
+
+
+class CohereCommandParser(DelegatingParser):
+    def __init__(
+        self,
+        tokenizer: TokenizerLike,
+        tools: list[Tool] | None = None,
+        *args,
+        model_config: ModelConfig | None = None,
+        **kwargs,
+    ):
+        super().__init__(tokenizer, tools, *args, model_config=model_config, **kwargs)
+        self._model_config = model_config
+        self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
+        self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
+        shim = self._reasoning_parser or self._tool_parser
+        assert isinstance(
+            shim, (BaseCohereCommandReasoningParser, BaseCohereCommandToolParser)
+        )
+        options = getattr(PyFilterOptions(), shim.melody_preset)()
+        self._streaming = PyFilter(options)
+        self._unary = PyFilter(options)
+        # Citations from the most recent ``parse`` call, surfaced on
+        # ``CohereChatMessage`` by ``CohereServingChatV2``.
+        self.last_unary_citations: list[Citation] | None = None
+        # Melody may emit a tool-call id before the function name; hold it until
+        # the first name delta so clients receive both together.
+        self._pending_tool_ids: dict[int, str] = {}
+        ctk = kwargs.get("chat_template_kwargs") or {}
+        raw_map = ctk.get(POSITION_TO_SOURCE_KEY)
+        self._position_to_source: PositionToSource | None = (
+            raw_map if isinstance(raw_map, Mapping) else None
+        )
+
+    def adjust_request(
+        self, request: ChatCompletionRequest | ResponsesRequest
+    ) -> ChatCompletionRequest | ResponsesRequest:
+        # keep special tokens to ensure melody can correctly identify boundaries.
+        request.skip_special_tokens = False
+        request = self._apply_structural_tags(request)
+        return super().adjust_request(request)
+
+    def _apply_structural_tags(
+        self, request: ChatCompletionRequest | ResponsesRequest
+    ) -> ChatCompletionRequest | ResponsesRequest:
+        """Fold ``tools`` and any JSON schema into a Cohere ``structural_tag``."""
+        so = request.structured_outputs
+        if so is not None and so.structural_tag:
+            return request
+        # Schema: prefer ``response_format`` (OpenAI Chat Completions), then
+        # ``structured_outputs.json`` / ``json_object`` (vLLM direct). Tools stay
+        # on ``request.tools``.
+        rf = (
+            request.response_format
+            if isinstance(request, ChatCompletionRequest)
+            else None
+        )
+        if rf is not None and _response_format_type(rf) == "structural_tag":
+            return request
+        model_architecture = (
+            self._model_config.architecture if self._model_config is not None else None
+        )
+        tools = request.tools
+        # ``response_format`` wins if both it and ``structured_outputs`` supply JSON.
+        schema = _schema_dict_from_chat_response_format(rf)
+        if schema is None:
+            schema = _schema_dict_from_structured_outputs(so)
+        if schema is None and not _has_effective_tools(tools):
+            return request
+        if model_architecture is None:
+            return request
+        result = convert_schema_to_structural_tags(
+            schema=schema,
+            tools=tools,
+            model_architecture=model_architecture,
+        )
+        if result is None:
+            # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``.
+            raise ValueError(
+                "Failed to build structural_tag guided decoding constraints from "
+                "this request's JSON schema and/or tools. The configured model "
+                f"architecture ({model_architecture!r}) does not support Cohere "
+                "command structural tags, or the schema cannot be expressed in "
+                "that format."
+            )
+        request.structured_outputs = StructuredOutputsParams(structural_tag=result)
+        # Folded JSON constraints into ``structural_tag``; drop ``response_format``
+        # when it was the source so ``to_sampling_params`` does not also set ``json`` /
+        # ``json_object`` (mutually exclusive in ``StructuredOutputsParams``).
+        if isinstance(request, ChatCompletionRequest) and rf is not None:
+            rf_type = _response_format_type(rf)
+            if rf_type in ("json_schema", "json_object"):
+                request.response_format = None
+        return request
+
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        count = 0
+        depth = 0
+        for token_id in token_ids:
+            if token_id == self.start_token_id:
+                depth += 1
+                continue
+            if token_id == self.end_token_id:
+                if depth > 0:
+                    depth -= 1
+                continue
+            if depth > 0:
+                count += 1
+        return count
+
+    def parse(
+        self,
+        model_output: str,
+        request: ChatCompletionRequest | ResponsesRequest,
+        enable_auto_tools: bool = False,
+        model_output_token_ids: Sequence[int] = (),
+    ) -> tuple[str | None, str | None, list[FunctionCall] | None]:
+        result = self._unary.process_full_text(model_output)
+        self.last_unary_citations = _melody_citations_to_vllm(
+            result.citations, self._position_to_source
+        )
+        tool_calls: list[FunctionCall] | None = None
+        if self._tools_allowed(request):
+            tool_calls = [
+                FunctionCall(id=tc.id, name=tc.name, arguments=tc.arguments)
+                for tc in result.tool_calls
+            ] or None
+        return result.reasoning, result.content, tool_calls
+
+    def parse_delta(
+        self,
+        delta_text: str,
+        delta_token_ids: list[int],
+        request: ChatCompletionRequest | ResponsesRequest,
+        prompt_token_ids: list[int] | None = None,
+        *,
+        finished: bool,
+    ) -> DeltaMessage | None:
+        results = [self._streaming.write_decoded(delta_text)]
+        if finished:
+            results.append(self._streaming.flush_partials())
+        return self._to_delta(results, request)
+
+    def _tools_allowed(self, request: ChatCompletionRequest | ResponsesRequest) -> bool:
+        return self._tool_parser is not None and request.tool_choice != "none"
+
+    def _to_delta(
+        self,
+        results: list[FilterAggregatedResult],
+        request: ChatCompletionRequest | ResponsesRequest,
+    ) -> CohereDeltaMessage | None:
+        content = self._join(r.content for r in results)
+        reasoning = (
+            self._join(r.reasoning for r in results)
+            if request.include_reasoning
+            else None
+        )
+        tool_calls: list[DeltaToolCall] = []
+        citations: list[Citation] = []
+        for result in results:
+            if result.tool_calls and self._tools_allowed(request):
+                tool_calls.extend(self._tool_call_deltas(result.tool_calls))
+            citations.extend(
+                _melody_citations_to_vllm(result.citations, self._position_to_source)
+                or []
+            )
+        if not (content or reasoning or tool_calls or citations):
+            return None
+        # ``CohereDeltaMessage`` carries citations to the wire via
+        # ``SerializeAsAny[DeltaMessage]``; it drops the field when unset so the
+        # shape matches a plain ``DeltaMessage``.
+        message = CohereDeltaMessage()
+        if content is not None:
+            message.content = content
+        if reasoning is not None:
+            message.reasoning = reasoning
+        if tool_calls:
+            message.tool_calls = tool_calls
+        if citations:
+            message.citations = citations
+        return message
+
+    @staticmethod
+    def _join(parts: Any) -> str | None:
+        present = [p for p in parts if p is not None]
+        return "".join(present) if present else None
+
+    def _tool_call_deltas(
+        self, melody_tool_calls: list[AccumulatedToolCall]
+    ) -> list[DeltaToolCall]:
+        tool_calls: list[DeltaToolCall] = []
+        for tc in melody_tool_calls:
+            name = tc.name or None
+            arguments = tc.arguments or None
+            if tc.id:
+                self._pending_tool_ids[tc.index] = tc.id
+            # Empty strings are melody streaming placeholders; skip fragments
+            # that carry neither a name nor argument text.
+            if name is None and arguments is None:
+                continue
+            function = DeltaFunctionCall()
+            if name is not None:
+                function.name = name
+            if arguments is not None:
+                function.arguments = arguments
+            delta = DeltaToolCall(index=tc.index, function=function)
+            if name is not None:
+                resolved_id = tc.id or self._pending_tool_ids.pop(tc.index, None)
+                if resolved_id is not None:
+                    delta.id = resolved_id
+                delta.type = "function"
+            tool_calls.append(delta)
+        return tool_calls
diff --git a/vllm/parser/deepseek_v32.py b/vllm/parser/deepseek_v32.py
index f76c954a93d6..de6634cf7368 100644
--- a/vllm/parser/deepseek_v32.py
+++ b/vllm/parser/deepseek_v32.py
@@ -53,6 +53,7 @@ def deepseek_v32_config() -> ParserEngineConfig:
     return ParserEngineConfig(
         name="deepseek_v32",
         initial_state=ParserState.CONTENT,
+        wait_for_reasoning=False,
         terminals={
             "TOOL_START": DSML_FUNC_START,
             "TOOL_END": DSML_FUNC_END,
diff --git a/vllm/parser/deepseek_v4.py b/vllm/parser/deepseek_v4.py
index 13eb83e6b804..49ef63736531 100644
--- a/vllm/parser/deepseek_v4.py
+++ b/vllm/parser/deepseek_v4.py
@@ -56,18 +56,24 @@
     f"<{_DSML}tool>",
 )
 
-_ESCAPED_DSML = re.escape(_DSML)
-_PARAM_RE = re.compile(
-    rf'<{_ESCAPED_DSML}parameter\s+name="([^"]+)"\s+string="(true|false)">'
-    rf"(.*?)"
-    rf"(?:</{_ESCAPED_DSML}parameter>|(?=<{_ESCAPED_DSML}parameter\s+name=))",
-    re.DOTALL,
-)
-_PARTIAL_PARAM_RE = re.compile(
-    rf'<{_ESCAPED_DSML}parameter\s+name="([^"]+)"\s+string="(true|false)">'
-    rf"(.*)$",
-    re.DOTALL,
-)
+
+def _param_patterns(
+    param_start: str, param_close: str
+) -> tuple[re.Pattern, re.Pattern]:
+    """Complete and trailing-partial parameter regexes for one DSML dialect.
+
+    The ``string`` attribute is optional: the model sometimes omits it, and
+    dropping such a parameter hands the client a tool call with no arguments.
+    """
+    start, close = re.escape(param_start), re.escape(param_close)
+    head = rf'{start}\s+name="([^"]+)"(?:\s+string="([^"]*)")?\s*>'
+    return (
+        re.compile(rf"{head}(.*?)(?:{close}|(?={start}\s+name=))", re.DOTALL),
+        re.compile(rf"{head}(.*)$", re.DOTALL),
+    )
+
+
+_PARAM_RE, _PARTIAL_PARAM_RE = _param_patterns(DSML_PARAM_START, DSML_PARAM_CLOSE)
 
 
 def _dsml_arg_converter(
@@ -80,6 +86,8 @@ def _dsml_arg_converter(
     params: dict[str, object] = {}
 
     last_end = 0
+    # A missing ``string`` attribute is treated like ``string="false"``: try
+    # JSON, fall back to the literal text.
     for m in param_re.finditer(raw_args):
         name, is_str, value = m.group(1), m.group(2), m.group(3)
         if is_str == "true":
@@ -140,6 +148,7 @@ def deepseek_v4_config(thinking: bool = False) -> ParserEngineConfig:
     return ParserEngineConfig(
         name="deepseek_v4",
         initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,
+        wait_for_reasoning=thinking,
         terminals={
             "THINK_START": DSML_THINK_START,
             "THINK_END": DSML_THINK_END,
diff --git a/vllm/parser/deepseek_v41.py b/vllm/parser/deepseek_v41.py
index 668bdd4a6bed..84216005102a 100644
--- a/vllm/parser/deepseek_v41.py
+++ b/vllm/parser/deepseek_v41.py
@@ -5,11 +5,10 @@
 import functools
 from dataclasses import replace
 
-import regex as re
-
 from vllm.parser.deepseek_v4 import (
     DeepSeekV4Parser,
     _dsml_arg_converter,
+    _param_patterns,
     deepseek_v4_config,
 )
 from vllm.parser.engine.parser_engine_config import ParserEngineConfig
@@ -21,17 +20,7 @@
 DSML_PARAM_START = "<｜DSML｜ parameter"
 DSML_PARAM_CLOSE = "</｜DSML｜ parameter>"
 
-_PARAM_RE = re.compile(
-    r'<｜DSML｜ parameter\s+name="([^"]+)"\s+string="(true|false)">'
-    r"(.*?)"
-    r"(?:</｜DSML｜ parameter>|(?=<｜DSML｜ parameter\s+name=))",
-    re.DOTALL,
-)
-_PARTIAL_PARAM_RE = re.compile(
-    r'<｜DSML｜ parameter\s+name="([^"]+)"\s+string="(true|false)">'
-    r"(.*)$",
-    re.DOTALL,
-)
+_PARAM_RE, _PARTIAL_PARAM_RE = _param_patterns(DSML_PARAM_START, DSML_PARAM_CLOSE)
 
 
 @functools.cache
diff --git a/vllm/parser/engine/__init__.py b/vllm/parser/engine/__init__.py
index 0bd26020bddd..a19f51f8ced4 100644
--- a/vllm/parser/engine/__init__.py
+++ b/vllm/parser/engine/__init__.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Streaming parser engine framework for tool call and reasoning extraction.
+"""Streaming parser engine framework for tool call and reasoning extraction.
 
 Instead of hand-rolling a parser for every model's tool-call / reasoning
 format, each format is declared as a ParserEngineConfig (terminals,
diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py
index 4817dedb4663..527c310c1485 100644
--- a/vllm/parser/engine/parser_engine.py
+++ b/vllm/parser/engine/parser_engine.py
@@ -645,34 +645,57 @@ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None:
         for offset, token_id in enumerate(token_ids):
             if token_id in end_ids:
                 return offset
-        return None
+        return len(token_ids)
 
     def is_reasoning_end(self, input_ids: list[int]) -> bool:
-        end_id = self._reasoning_end_token_id
+        config = self.parser_engine_config
+        wait_for_reasoning = config.wait_for_reasoning
+        if wait_for_reasoning is None:
+            wait_for_reasoning = config.initial_state is ParserState.REASONING
+        start_transition = config.transitions.get((config.initial_state, "THINK_START"))
+        start_opens_reasoning = (
+            start_transition is not None
+            and start_transition.next_state is ParserState.REASONING
+        )
+        end_ids = self._reasoning_end_token_ids
         start_id = self._reasoning_start_token_id
-        if end_id is not None:
-            if not input_ids:
-                return self.parser_engine_config.initial_state != ParserState.REASONING
-            boundary_ids = self._turn_boundary_token_ids
-            for i in range(len(input_ids) - 1, -1, -1):
-                token_id = input_ids[i]
-                if token_id == end_id:
-                    return True
-                if start_id is not None and token_id == start_id:
+        boundary_ids = self._turn_boundary_token_ids
+        for token_id in reversed(input_ids):
+            if token_id in end_ids:
+                return True
+            if token_id == start_id:
+                if start_opens_reasoning:
                     return False
-                if token_id in boundary_ids:
-                    return (
-                        self.parser_engine_config.initial_state != ParserState.REASONING
-                    )
-            return False
-        return self._reasoning_ended
+                break
+            if token_id in boundary_ids:
+                break
+        return not wait_for_reasoning
 
     def extract_content_ids(self, input_ids: list[int]) -> list[int]:
+        config = self.parser_engine_config
+        wait_for_reasoning = config.wait_for_reasoning
+        if wait_for_reasoning is None:
+            wait_for_reasoning = config.initial_state is ParserState.REASONING
+        if not wait_for_reasoning:
+            return input_ids
+
         end_id = self._reasoning_end_token_id
         if end_id is not None:
             for i in range(len(input_ids) - 1, -1, -1):
                 if input_ids[i] == end_id:
                     return input_ids[i + 1 :]
+
+        end_ids = self._reasoning_end_token_ids
+        if end_ids:
+            turn_start = 0
+            boundary_ids = self._turn_boundary_token_ids
+            for i in range(len(input_ids) - 1, -1, -1):
+                if input_ids[i] in boundary_ids:
+                    turn_start = i + 1
+                    break
+            for i in range(turn_start, len(input_ids)):
+                if input_ids[i] in end_ids:
+                    return input_ids[i:]
         return input_ids
 
     def get_streaming_fallback_content(
diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py
index 5c5cd2b467fd..215fa0d9bdf5 100644
--- a/vllm/parser/engine/parser_engine_config.py
+++ b/vllm/parser/engine/parser_engine_config.py
@@ -74,12 +74,20 @@ class ParserEngineConfig:
 
     initial_state: ParserState = ParserState.CONTENT
 
+    wait_for_reasoning: bool | None = None
+
     arg_converter: Callable[[str, bool], str] | None = None
 
     stream_arg_deltas: bool = True
 
     tool_args_json: bool = True
 
+    # The tool region body is a single JSON array of ``{"name", "arguments"}``
+    # objects (one marker → N calls, no closing terminal). The engine splits
+    # the array into one tool call per element; the array's ``]`` ends the
+    # region and any trailing text becomes content. Used by Granite.
+    tool_call_body_array: bool = False
+
     arg_structural_chars: frozenset[str] | None = None
 
     # Special tokens exempt from auto-drop but not state-machine terminals.
diff --git a/vllm/parser/engine/registered_adapters.py b/vllm/parser/engine/registered_adapters.py
index 8e10a1b5efbe..f7f749948d23 100644
--- a/vllm/parser/engine/registered_adapters.py
+++ b/vllm/parser/engine/registered_adapters.py
@@ -13,6 +13,7 @@
 from vllm.parser.engine.adapters import make_adapters
 from vllm.parser.gemma4 import Gemma4Parser
 from vllm.parser.glm47_moe import Glm47MoeParser
+from vllm.parser.granite import GraniteParser
 from vllm.parser.inkling import InklingParser
 from vllm.parser.kimi_k2 import KimiK2Parser
 from vllm.parser.minimax_m2 import MinimaxM2Parser
@@ -46,6 +47,11 @@
     Gemma4ParserToolAdapter,
 ) = make_adapters(Gemma4Parser)
 
+(
+    GraniteParserReasoningAdapter,
+    GraniteParserToolAdapter,
+) = make_adapters(GraniteParser)
+
 (
     NemotronV3ParserReasoningAdapter,
     NemotronV3ParserToolAdapter,
diff --git a/vllm/parser/engine/streaming_parser_engine.py b/vllm/parser/engine/streaming_parser_engine.py
index 21d053ef28da..3fd1ecdcb656 100644
--- a/vllm/parser/engine/streaming_parser_engine.py
+++ b/vllm/parser/engine/streaming_parser_engine.py
@@ -190,6 +190,12 @@ def _reset_args_state(self) -> None:
         self._args_in_string: bool = False
         self._args_escape_next: bool = False
 
+    def _reset_array_state(self) -> None:
+        # State for ``tool_call_body_array``: whether the opening ``[`` has been
+        # seen, and whether a ``{...}`` element is currently being streamed.
+        self._array_started: bool = False
+        self._array_in_element: bool = False
+
     def reset(self, initial_state: ParserState | None = None) -> None:
         """Reset mutable state for reuse across requests.
 
@@ -213,6 +219,7 @@ def reset(self, initial_state: ParserState | None = None) -> None:
         self._message_header_token_count = 0
         self._in_skipped_tool_span = False
         self._reset_args_state()
+        self._reset_array_state()
 
     def feed(
         self,
@@ -306,7 +313,12 @@ def finish(self) -> list[SemanticEvent]:
             ParserState.TOOL_NAME,
             ParserState.TOOL_BETWEEN,
         ):
-            if self.tool_index >= 0:
+            # In array mode a closed element has already emitted its
+            # TOOL_CALL_END; only a truncated (still-open) element needs one.
+            end_open_call = (
+                self._array_in_element if self.config.tool_call_body_array else True
+            )
+            if end_open_call and self.tool_index >= 0:
                 events.append(
                     SemanticEvent(
                         EventType.TOOL_CALL_END,
@@ -474,6 +486,8 @@ def _emit_for_state(self, text: str, token_count: int = 0) -> list[SemanticEvent
             self._message_header_token_count += token_count
             return []
         if self.state == ParserState.TOOL_ARGS:
+            if self.config.tool_call_body_array:
+                return self._feed_array_text(text)
             if self.config.tool_args_json:
                 return self._feed_args_text(text)
             return [
@@ -564,6 +578,8 @@ def _apply_transition(
             self._args_in_string = False
             self._args_escape_next = False
             self._args_safe_end = 0
+            if self.config.tool_call_body_array:
+                self._reset_array_state()
 
         return events
 
@@ -634,3 +650,88 @@ def _flush_safe_args(self) -> list[SemanticEvent]:
                 tool_index=self.tool_index,
             )
         ]
+
+    # ── JSON-array tool body (tool_call_body_array) ─────────────────────
+
+    def _feed_array_text(self, text: str) -> list[SemanticEvent]:
+        """Split a JSON-array tool body into one tool call per element.
+
+        Granite emits ``[{"name":..,"arguments":{..}}, ...]`` after the tool
+        marker. Element interiors are streamed through the string/brace-aware
+        :meth:`_feed_args_char`, so each element behaves exactly like a
+        single-object tool body; the array's ``]`` ends the region and any
+        trailing text becomes content.
+        """
+        events: list[SemanticEvent] = []
+        for i, ch in enumerate(text):
+            if self.state != ParserState.TOOL_ARGS:
+                # The array closed mid-chunk; the rest of the text is content.
+                events.extend(self._emit_for_state(text[i:]))
+                break
+            events.extend(self._feed_array_char(ch))
+        return events
+
+    def _feed_array_char(self, ch: str) -> list[SemanticEvent]:
+        if self._array_in_element:
+            events = self._feed_args_char(ch)
+            # Element object closed: top-level '}' with all braces balanced.
+            if (
+                ch == "}"
+                and not self._args_in_string
+                and not self._args_escape_next
+                and self._args_brace_depth == 0
+            ):
+                events.extend(self._end_array_element())
+            return events
+
+        # Between elements only structural array characters are expected.
+        if ch in " \t\r\n":
+            return []
+        if not self._array_started:
+            if ch == "[":
+                self._array_started = True
+                return []
+            return self._bail_array_to_content(ch)
+        if ch == ",":
+            return []
+        if ch == "{":
+            return self._start_array_element(ch)
+        if ch == "]":
+            self.state = ParserState.CONTENT
+            self._reset_array_state()
+            return []
+        return self._bail_array_to_content(ch)
+
+    def _start_array_element(self, ch: str) -> list[SemanticEvent]:
+        self.tool_index += 1
+        self._reset_args_state()
+        self._array_in_element = True
+        events = [SemanticEvent(EventType.TOOL_CALL_START, tool_index=self.tool_index)]
+        events.extend(self._feed_args_char(ch))
+        return events
+
+    def _end_array_element(self) -> list[SemanticEvent]:
+        self._array_in_element = False
+        events: list[SemanticEvent] = []
+        # Flush the held closing '}' so the slot body is the full element.
+        if self._args_buffer:
+            events.append(
+                SemanticEvent(
+                    EventType.ARG_VALUE_CHUNK,
+                    value=self._args_buffer,
+                    tool_index=self.tool_index,
+                )
+            )
+            self._args_buffer = ""
+            self._args_safe_end = 0
+        events.append(
+            SemanticEvent(EventType.TOOL_CALL_END, tool_index=self.tool_index)
+        )
+        return events
+
+    def _bail_array_to_content(self, ch: str) -> list[SemanticEvent]:
+        # Marker not followed by a JSON array (malformed) — revert to content
+        # so the remaining tokens are emitted as plain text, no tool call.
+        self.state = ParserState.CONTENT
+        self._reset_array_state()
+        return self._emit_for_state(ch)
diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py
index da1718332b21..185610126f5b 100644
--- a/vllm/parser/gemma4.py
+++ b/vllm/parser/gemma4.py
@@ -83,6 +83,7 @@ def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict:
             (e.g. partial boolean parsed as bare string).
 
     Returns a dict ready for ``json.dumps()``.
+
     """
     if not args_str or not args_str.strip():
         return {}
@@ -293,16 +294,21 @@ def _gemma4_arg_converter(raw_args: str, partial: bool) -> str:
 
 
 @functools.cache
-def gemma4_config() -> ParserEngineConfig:
+def gemma4_config(thinking: bool = False) -> ParserEngineConfig:
     return ParserEngineConfig(
         name="gemma4",
         initial_state=ParserState.CONTENT,
+        # Decides a silent tail: the E-series HF templates write no marker
+        # with thinking off, and no template does after a tool response.
+        wait_for_reasoning=thinking,
+        turn_boundary_tokens=frozenset({"<|turn>", "<|tool_response>"}),
         terminals={
             "THINK_START": CHANNEL_START,
             "THINK_END": CHANNEL_END,
             "TOOL_START": TOOL_CALL_START,
             "TOOL_END": TOOL_CALL_END,
             "CALL_PREFIX": "call:",
+            "COLON": ":",
             "OPEN_BRACE": "{",
             "OPEN_PAREN": "(",
         },
@@ -348,6 +354,13 @@ def gemma4_config() -> ParserEngineConfig:
                 ParserState.TOOL_NAME,
                 (),
             ),
+            # Bare opener: some checkpoints emit "<|tool_call>:name{...}"
+            # without the "call" prefix.  Match the colon directly so the
+            # tool call is not silently dropped.
+            (ParserState.TOOL_PREAMBLE, "COLON"): Transition(
+                ParserState.TOOL_NAME,
+                (),
+            ),
             (ParserState.TOOL_NAME, "OPEN_BRACE"): Transition(
                 ParserState.TOOL_ARGS,
                 (),
@@ -409,11 +422,11 @@ def __init__(
         **kwargs,
     ) -> None:
         chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
-        self._thinking_enabled = chat_kwargs.get("enable_thinking", False)
+        thinking = chat_kwargs.get("enable_thinking", False)
         super().__init__(
             tokenizer,
             tools,
-            parser_engine_config=gemma4_config(),
+            parser_engine_config=gemma4_config(thinking=thinking),
             **kwargs,
         )
         vocab = self.vocab
@@ -464,30 +477,6 @@ def _preprocess_feed(
 
         return delta_text, delta_token_ids
 
-    def is_reasoning_end(self, input_ids: list[int]) -> bool:
-        end_id = self._reasoning_end_token_id
-        start_id = self._reasoning_start_token_id
-        tool_call_id = self._tool_call_token_id
-        new_turn_id = self._new_turn_token_id
-        tool_response_id = self._tool_response_token_id
-
-        if end_id is not None and not input_ids:
-            return self.parser_engine_config.initial_state != ParserState.REASONING
-
-        for i in range(len(input_ids) - 1, -1, -1):
-            tid = input_ids[i]
-            if start_id is not None and tid == start_id:
-                return False
-            if tool_call_id is not None and tid == tool_call_id:
-                return True
-            if new_turn_id is not None and tid == new_turn_id:
-                return not self._thinking_enabled
-            if tool_response_id is not None and tid == tool_response_id:
-                return not self._thinking_enabled
-            if end_id is not None and tid == end_id:
-                return True
-        return True
-
     def _prompt_ends_in_open_reasoning(self, prompt_token_ids: Sequence[int]) -> bool:
         """Whether the prompt tail is inside an open ``<|channel>`` block.
 
diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py
index 8aa4feef2599..00e517f2ea68 100644
--- a/vllm/parser/glm47_moe.py
+++ b/vllm/parser/glm47_moe.py
@@ -41,6 +41,13 @@
 ARG_VALUE_START = "<arg_value>"
 ARG_VALUE_END = "</arg_value>"
 
+# Special tokens that delimit conversation turns in a rendered GLM prompt.
+# Reasoning markers belonging to earlier turns must not be mistaken for the
+# state of the turn currently being generated.
+GLM_TURN_BOUNDARIES = frozenset(
+    ("<|system|>", "<|user|>", "<|assistant|>", "<|observation|>")
+)
+
 _ARG_RE = re.compile(
     r"<arg_key>(?P<key>.*?)</arg_key>\s*"
     r"<arg_value>(?P<value>.*?)</arg_value>",
@@ -165,6 +172,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig:
             ),
             **arg_tag_transitions,
         },
+        turn_boundary_tokens=GLM_TURN_BOUNDARIES,
         arg_converter=_glm47_arg_converter,
         stream_arg_deltas=True,
         tool_args_json=False,
@@ -206,16 +214,6 @@ def _handle_tool_end(self, event, deltas) -> None:
             self._tool_slots[idx].name = self._tool_slots[idx].name.strip()
         super()._handle_tool_end(event, deltas)
 
-    def is_reasoning_end(self, input_ids: list[int]) -> bool:
-        if not self.thinking_enabled:
-            return True
-        return super().is_reasoning_end(input_ids)
-
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        if not self.thinking_enabled:
-            return input_ids
-        return super().extract_content_ids(input_ids)
-
     def extract_reasoning(
         self,
         model_output: str,
diff --git a/vllm/parser/granite.py b/vllm/parser/granite.py
new file mode 100644
index 000000000000..afa202122506
--- /dev/null
+++ b/vllm/parser/granite.py
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Granite parser for JSON-array tool calls.
+
+Granite 3.0/3.1 emit tool calls as a JSON array following a marker::
+
+    <|tool_call|> [{"name": "get_weather", "arguments": {"city": "SF"}}]   # 3.0
+    <tool_call> [{"name": "get_weather", "arguments": {"city": "SF"}}]      # 3.1
+
+``<|tool_call|>`` is a single special token; ``<tool_call>`` is plain text.
+The array may hold several calls, has no closing marker (the ``]`` ends it),
+and any surrounding prose is content. The engine's ``tool_call_body_array``
+mode splits the array into one call per element and extracts each name via the
+name-from-args path; ``_granite_arg_converter`` carves the ``arguments`` object
+out of each ``{"name":..,"arguments":{..}}`` element. Granite has no reasoning.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+from vllm.parser.engine.parser_engine import ParserEngine
+from vllm.parser.engine.parser_engine_config import (
+    ParserEngineConfig,
+    ParserState,
+    Transition,
+)
+from vllm.parser.utils import extract_json_member_object
+
+if TYPE_CHECKING:
+    from vllm.tokenizers import TokenizerLike
+    from vllm.tool_parsers.abstract_tool_parser import Tool
+
+TOOL_TOKEN = "<|tool_call|>"  # granite 3.0 (single special token)
+TOOL_STRING = "<tool_call>"  # granite 3.1 (plain text)
+
+
+def _granite_arg_converter(raw_args: str, partial: bool) -> str:
+    """Carve the ``arguments`` object out of a ``{"name":..,"arguments":{..}}``
+    array element (see :func:`extract_json_member_object`)."""
+    span = extract_json_member_object(raw_args, "arguments")
+    if span is None:
+        return "" if partial else "{}"
+    return span
+
+
+@functools.cache
+def granite_config() -> ParserEngineConfig:
+    return ParserEngineConfig(
+        name="granite",
+        initial_state=ParserState.CONTENT,
+        terminals={
+            "TOOL_TOKEN": TOOL_TOKEN,
+            "TOOL_STRING": TOOL_STRING,
+        },
+        token_id_terminals={
+            "TOOL_TOKEN": TOOL_TOKEN,
+        },
+        transitions={
+            # The marker only switches into the tool region; the array splitter
+            # emits TOOL_CALL_START per element (empty events here).
+            (ParserState.CONTENT, "TOOL_TOKEN"): Transition(
+                ParserState.TOOL_ARGS,
+                (),
+            ),
+            (ParserState.CONTENT, "TOOL_STRING"): Transition(
+                ParserState.TOOL_ARGS,
+                (),
+            ),
+        },
+        arg_converter=_granite_arg_converter,
+        tool_call_body_array=True,
+        stream_arg_deltas=True,
+    )
+
+
+class GraniteParser(ParserEngine):
+    CONFIG_NAME = "granite"
+
+    def __init__(
+        self,
+        tokenizer: TokenizerLike,
+        tools: list[Tool] | None = None,
+        **kwargs,
+    ) -> None:
+        kwargs.setdefault("parser_engine_config", granite_config())
+        super().__init__(tokenizer, tools, **kwargs)
diff --git a/vllm/parser/inkling.py b/vllm/parser/inkling.py
index 16883b58aa7f..cc786c66dd3e 100644
--- a/vllm/parser/inkling.py
+++ b/vllm/parser/inkling.py
@@ -37,6 +37,7 @@
     ParserState,
     Transition,
 )
+from vllm.parser.utils import extract_json_member_object
 
 if TYPE_CHECKING:
     from vllm.tokenizers import TokenizerLike
@@ -62,84 +63,6 @@
     END_MESSAGE,
 )
 
-_WS = " \t\r\n"
-
-
-def _scan_json_value(raw: str, start: int) -> int | None:
-    """Return the end index (exclusive) of the JSON object starting at
-    ``raw[start]``, or ``None`` when the object is still unterminated."""
-    depth = 0
-    in_string = False
-    escape = False
-    for i in range(start, len(raw)):
-        ch = raw[i]
-        if escape:
-            escape = False
-            continue
-        if in_string:
-            if ch == "\\":
-                escape = True
-            elif ch == '"':
-                in_string = False
-            continue
-        if ch == '"':
-            in_string = True
-        elif ch == "{":
-            depth += 1
-        elif ch == "}":
-            depth -= 1
-            if depth == 0:
-                return i + 1
-    return None
-
-
-def _args_value_span(raw: str) -> str | None:
-    """Extract the raw text span of the top-level ``"args"`` value from a
-    (possibly incomplete) ``{"name":...,"args":{...}}`` wrapper.
-
-    Returns the verbatim substring (prefix-stable across growing input,
-    which the engine's argument-delta diffing relies on), possibly an
-    unterminated object prefix; ``None`` when the value has not started.
-    Raises ``ValueError`` when the value is not a JSON object.
-    """
-    depth = 0
-    in_string = False
-    escape = False
-    string_start = -1
-    last_string: str | None = None
-    for i, ch in enumerate(raw):
-        if escape:
-            escape = False
-            continue
-        if in_string:
-            if ch == "\\":
-                escape = True
-            elif ch == '"':
-                in_string = False
-                if depth == 1:
-                    last_string = raw[string_start + 1 : i]
-            continue
-        if ch == '"':
-            in_string = True
-            string_start = i
-        elif ch == ":" and depth == 1 and last_string == "args":
-            value_start = i + 1
-            while value_start < len(raw) and raw[value_start] in _WS:
-                value_start += 1
-            if value_start >= len(raw):
-                return None
-            if raw[value_start] != "{":
-                raise ValueError("Inkling tool call args must be a JSON object")
-            value_end = _scan_json_value(raw, value_start)
-            if value_end is None:
-                return raw[value_start:]
-            return raw[value_start:value_end]
-        elif ch in "{[":
-            depth += 1
-        elif ch in "}]":
-            depth -= 1
-    return None
-
 
 def _inkling_arg_converter(raw_args: str, partial: bool) -> str:
     """Carve the ``args`` object out of the tool-call JSON wrapper.
@@ -150,19 +73,10 @@ def _inkling_arg_converter(raw_args: str, partial: bool) -> str:
     converter, ``_compute_arg_delta`` streams the wrapper verbatim into
     the OpenAI ``arguments`` field (``converter is None -> raw delta``,
     unconditionally; ``stream_arg_deltas=False`` does not stop it).
-
-    Why a hand-rolled scanner instead of (partial) ``json.loads`` +
-    ``json.dumps``: the engine diffs successive converter outputs and
-    requires each to extend the previous one (``startswith``); a
-    violation silently drops argument deltas. Re-serialization changes
-    whitespace and closes unterminated structures differently across
-    ticks, so the only prefix-stable output is a verbatim substring of
-    the input. The scanner is also string/escape-aware so an ``"args"``
-    literal inside the name or a string value cannot mislead it, and it
-    still recovers a partial span when EOS truncates the wrapper (where
-    ``json.loads`` would fail).
+    :func:`extract_json_member_object` returns a prefix-stable verbatim
+    substring, which the engine's argument-delta diffing requires.
     """
-    span = _args_value_span(raw_args)
+    span = extract_json_member_object(raw_args, "args")
     if span is None:
         # No args value yet (streaming) or none at all (treat as empty).
         return "" if partial else "{}"
@@ -282,6 +196,7 @@ def inkling_config() -> ParserEngineConfig:
         # `<|message_model|>`. Non-streaming parsing receives only the generated
         # suffix, so begin in the corresponding message-header state as well.
         initial_state=ParserState.MESSAGE_HEADER,
+        wait_for_reasoning=True,
         terminals=terminals,
         # Inkling content-kind markers are the grammar. When the engine is
         # used through DelegatingParser, the reasoning pass can hand the tool
diff --git a/vllm/parser/kimi_k2.py b/vllm/parser/kimi_k2.py
index 1e1067a8c2ee..7c5f30aafef4 100644
--- a/vllm/parser/kimi_k2.py
+++ b/vllm/parser/kimi_k2.py
@@ -81,6 +81,7 @@ def kimi_k2_config(thinking: bool = True) -> ParserEngineConfig:
     return ParserEngineConfig(
         name="kimi_k2",
         initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,
+        wait_for_reasoning=thinking,
         terminals={
             **reasoning_terminals,
             "TOOL_SECTION_START": TOOL_SECTION_START,
@@ -236,24 +237,6 @@ def _handle_arg_chunk(self, event, deltas) -> None:
     def _extract_args_json(self, raw_args: str, func_name: str) -> str:
         return raw_args.strip() or "{}"
 
-    def is_reasoning_end(self, input_ids: list[int]) -> bool:
-        if not self.thinking_enabled:
-            return True
-
-        start_id = self._start_token_id
-        end_id = self._end_token_id
-        tool_section_id = self._tool_section_start_token_id
-
-        for i in range(len(input_ids) - 1, -1, -1):
-            token_id = input_ids[i]
-            if start_id is not None and token_id == start_id:
-                return False
-            if end_id is not None and token_id == end_id:
-                return True
-            if tool_section_id is not None and token_id == tool_section_id:
-                return True
-        return False
-
     def extract_content_ids(self, input_ids: list[int]) -> list[int]:
         if not self.thinking_enabled:
             return input_ids
diff --git a/vllm/parser/minimax_m2.py b/vllm/parser/minimax_m2.py
index ab1b7eb11427..c96f852b9df6 100644
--- a/vllm/parser/minimax_m2.py
+++ b/vllm/parser/minimax_m2.py
@@ -95,6 +95,7 @@ def minimax_m2_config() -> ParserEngineConfig:
     return ParserEngineConfig(
         name="minimax_m2",
         initial_state=ParserState.REASONING,
+        wait_for_reasoning=True,
         terminals={
             "THINK_START": THINK_START,
             "THINK_END": THINK_END,
diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py
index fc9b8fbefe37..247ef2803fbe 100644
--- a/vllm/parser/mistral.py
+++ b/vllm/parser/mistral.py
@@ -136,6 +136,7 @@ def mistral_config(
 
     Returns:
         A frozen :class:`ParserEngineConfig` with ``initial_state=CONTENT``.
+
     """
     if reasoning_encoding == "special_token":
         think_start = _THINK_START_SPECIAL
@@ -193,6 +194,7 @@ def mistral_config(
     return ParserEngineConfig(
         name=name,
         initial_state=ParserState.CONTENT,
+        wait_for_reasoning=reasoning_encoding != "none",
         terminals={
             **reasoning_terminals,
             "TOOL_CALLS": _TOOL_CALLS,
@@ -466,6 +468,7 @@ def _build_guided_schema_pre_v11(
 
         Returns:
             A JSON Schema dict, or ``None`` if the named tool is not found.
+
         """
         tool_choice = request.tool_choice
         tools = request.tools or []
@@ -891,24 +894,6 @@ def _split_delta(
 
         return (delta_to_be_parsed, "")
 
-    def is_reasoning_end(self, input_ids: list[int]) -> bool:
-        if self._reasoning_encoding == "none":
-            return True
-        if super().is_reasoning_end(input_ids):
-            return True
-        # [TOOL_CALLS] acts as an implicit reasoning-end marker
-        if self.bot_token_id is not None:
-            reasoning_start_id = self._reasoning_start_token_id
-            for i in range(len(input_ids) - 1, -1, -1):
-                if (
-                    reasoning_start_id is not None
-                    and input_ids[i] == reasoning_start_id
-                ):
-                    return False
-                if input_ids[i] == self.bot_token_id:
-                    return True
-        return False
-
     def extract_reasoning(
         self,
         model_output: str,
diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py
index 420873a58da8..41b27967a6cb 100644
--- a/vllm/parser/parser_manager.py
+++ b/vllm/parser/parser_manager.py
@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING
 
 from vllm.logger import init_logger
+from vllm.tool_parsers.tool_strict_level import ToolStrictLevel
 
 if TYPE_CHECKING:
     from vllm.parser.abstract_parser import Parser
@@ -16,9 +17,7 @@
 
 
 class ParserManager:
-    """
-    Provides a unified Parser by composing reasoning and tool parser adapters.
-    """
+    """Provides a unified Parser by composing reasoning and tool parser adapters."""
 
     @classmethod
     def get_tool_parser(
@@ -79,9 +78,9 @@ def get_parser(
         enable_auto_tools: bool = False,
         model_name: str | None = None,
         is_harmony: bool = False,
+        tool_strict_level: str = "auto",
     ) -> type[Parser] | None:
-        """
-        Get a Parser that handles both reasoning and tool parsing.
+        """Get a Parser that handles both reasoning and tool parsing.
 
         Composes the individual parsers into a ``DelegatingParser`` subclass.
 
@@ -92,9 +91,12 @@ def get_parser(
             model_name: The model name for parser-specific warnings.
             is_harmony: Whether the selected model uses the Harmony format.
                         If True, HarmonyParser is always returned.
+            tool_strict_level: Server-side floor for tool-call structural
+                tags (``--tool-strict-level``).
 
         Returns:
             A Parser class, or None if neither parser is specified.
+
         """
         if not tool_parser_name and not reasoning_parser_name:
             return None
@@ -107,11 +109,14 @@ def get_parser(
         if reasoning_parser_cls is None and tool_parser_cls is None:
             return None
 
+        strict_level = ToolStrictLevel.from_name(tool_strict_level)
+
         if is_harmony:
             from vllm.parser.harmony import HarmonyParser
 
             HarmonyParser.reasoning_parser_cls = reasoning_parser_cls
             HarmonyParser.tool_parser_cls = tool_parser_cls
+            HarmonyParser.tool_strict_level = strict_level
             return HarmonyParser
 
         if reasoning_parser_name == "kimi_k3" or tool_parser_name == "kimi_k3":
@@ -123,9 +128,26 @@ def get_parser(
             class _KimiK3Parser(KimiK3Parser):
                 reasoning_parser_cls = r_cls
                 tool_parser_cls = t_cls
+                tool_strict_level = strict_level
 
             return _KimiK3Parser
 
+        if {reasoning_parser_name, tool_parser_name} & {
+            "cohere_command3",
+            "cohere_command4",
+        }:
+            from vllm.parser.cohere_command import CohereCommandParser
+
+            r_cls = reasoning_parser_cls
+            t_cls = tool_parser_cls
+
+            class _CohereCommandParser(CohereCommandParser):
+                reasoning_parser_cls = r_cls
+                tool_parser_cls = t_cls
+                tool_strict_level = strict_level
+
+            return _CohereCommandParser
+
         from vllm.parser.abstract_parser import DelegatingParser
 
         r_cls = reasoning_parser_cls
@@ -134,5 +156,6 @@ class _KimiK3Parser(KimiK3Parser):
         class _Parser(DelegatingParser):
             reasoning_parser_cls = r_cls
             tool_parser_cls = t_cls
+            tool_strict_level = strict_level
 
         return _Parser
diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py
index 8edf42e26363..a11551a209a0 100644
--- a/vllm/parser/qwen3.py
+++ b/vllm/parser/qwen3.py
@@ -100,6 +100,7 @@ def qwen3_config(
     return ParserEngineConfig(
         name=name,
         initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,
+        wait_for_reasoning=thinking,
         turn_boundary_tokens=turn_boundary_tokens,
         terminals={
             # Reasoning terminals
@@ -205,8 +206,8 @@ class Qwen3Parser(ParserEngine):
     """Qwen3 parser: ``<think>``/``</think>`` reasoning +
     ``<tool_call>`` XML tool calls in a single engine.
 
-    - ``<tool_call>`` as implicit reasoning end
-    - Unpaired ``<tool_call>`` token ID detection for ``is_reasoning_end``
+    - ``<tool_call>`` as implicit reasoning end (a grammar transition, so it
+      also feeds ``is_reasoning_end`` and the structured-output gate)
 
     Subclasses that share the grammar but differ only in the four wrapper
     token strings (reasoning + tool-call) override the class attributes
@@ -245,9 +246,6 @@ def __init__(
             tools,
             **kwargs,
         )
-        vocab = self.vocab
-        self._tool_call_token_id: int | None = vocab.get(self.TOOL_START)
-        self._tool_call_end_token_id: int | None = vocab.get(self.TOOL_END)
 
     def extract_reasoning(
         self,
@@ -257,28 +255,3 @@ def extract_reasoning(
         if not self.thinking_enabled:
             return None, model_output
         return super().extract_reasoning(model_output, request)
-
-    def is_reasoning_end(self, input_ids: list[int]) -> bool:
-        if super().is_reasoning_end(input_ids):
-            return True
-        tool_call_id = self._tool_call_token_id
-        tool_call_end_id = self._tool_call_end_token_id
-        reasoning_start_id = self._reasoning_start_token_id
-        boundary_ids = self._turn_boundary_token_ids
-        if tool_call_id is not None:
-            for i in range(len(input_ids) - 1, -1, -1):
-                if (
-                    reasoning_start_id is not None
-                    and input_ids[i] == reasoning_start_id
-                ):
-                    return False
-                if input_ids[i] in boundary_ids:
-                    return False
-                if input_ids[i] == tool_call_id:
-                    if tool_call_end_id is not None and any(
-                        input_ids[j] == tool_call_end_id
-                        for j in range(i + 1, len(input_ids))
-                    ):
-                        continue
-                    return True
-        return False
diff --git a/vllm/parser/utils.py b/vllm/parser/utils.py
index 51382cd29094..9194b7ef1821 100644
--- a/vllm/parser/utils.py
+++ b/vllm/parser/utils.py
@@ -63,3 +63,88 @@ def count_history_tool_calls(
         return 0
 
     return count_response_history_tool_calls(request_input)
+
+
+_JSON_WS = " \t\r\n"
+
+
+def scan_json_object_end(raw: str, start: int) -> int | None:
+    """Return the end index (exclusive) of the JSON object starting at
+    ``raw[start]`` (which must be ``{``), or ``None`` when it is still
+    unterminated. String- and escape-aware so braces inside strings do not
+    mislead the depth count."""
+    depth = 0
+    in_string = False
+    escape = False
+    for i in range(start, len(raw)):
+        ch = raw[i]
+        if escape:
+            escape = False
+            continue
+        if in_string:
+            if ch == "\\":
+                escape = True
+            elif ch == '"':
+                in_string = False
+            continue
+        if ch == '"':
+            in_string = True
+        elif ch == "{":
+            depth += 1
+        elif ch == "}":
+            depth -= 1
+            if depth == 0:
+                return i + 1
+    return None
+
+
+def extract_json_member_object(raw: str, key: str) -> str | None:
+    """Extract the raw text span of the top-level object-valued member ``key``
+    from a (possibly incomplete) ``{...}`` JSON wrapper.
+
+    Returns the verbatim substring (prefix-stable as the input grows, which
+    the engine's argument-delta diffing relies on), possibly an unterminated
+    object prefix; ``None`` when the value has not started. Raises
+    ``ValueError`` when the member is present but its value is not a JSON
+    object. A verbatim substring (not re-serialized JSON) is required because
+    the engine diffs successive converter outputs and expects each to extend
+    the previous one; the scanner is string/escape-aware so a matching literal
+    inside another value cannot mislead it, and it recovers a partial span when
+    the wrapper is truncated (where ``json.loads`` would fail)."""
+    depth = 0
+    in_string = False
+    escape = False
+    string_start = -1
+    last_string: str | None = None
+    for i, ch in enumerate(raw):
+        if escape:
+            escape = False
+            continue
+        if in_string:
+            if ch == "\\":
+                escape = True
+            elif ch == '"':
+                in_string = False
+                if depth == 1:
+                    last_string = raw[string_start + 1 : i]
+            continue
+        if ch == '"':
+            in_string = True
+            string_start = i
+        elif ch == ":" and depth == 1 and last_string == key:
+            value_start = i + 1
+            while value_start < len(raw) and raw[value_start] in _JSON_WS:
+                value_start += 1
+            if value_start >= len(raw):
+                return None
+            if raw[value_start] != "{":
+                raise ValueError(f"member {key!r} must be a JSON object")
+            value_end = scan_json_object_end(raw, value_start)
+            if value_end is None:
+                return raw[value_start:]
+            return raw[value_start:value_end]
+        elif ch in "{[":
+            depth += 1
+        elif ch in "}]":
+            depth -= 1
+    return None
diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py
index b34b0b9a7dd9..add01346550b 100644
--- a/vllm/platforms/__init__.py
+++ b/vllm/platforms/__init__.py
@@ -16,9 +16,7 @@
 
 
 def vllm_version_matches_substr(substr: str) -> bool:
-    """
-    Check to see if the vLLM version matches a substring.
-    """
+    """Check to see if the vLLM version matches a substring."""
     from importlib.metadata import PackageNotFoundError, version
 
     try:
@@ -217,6 +215,15 @@ def cpu_platform_plugin() -> str | None:
                 "AMD Zen CPU detected but zentorch not installed, "
                 "falling back to CpuPlatform."
             )
+        except OSError:
+            # An ABI-mismatched build fails here with an undefined-symbol
+            # error; other failures are not known to be safe to recover from.
+            logger.warning(
+                "AMD Zen CPU detected but zentorch failed to import, falling "
+                "back to CpuPlatform. This usually means the zentorch build "
+                "does not match the installed torch version.",
+                exc_info=True,
+            )
 
     return "vllm.platforms.cpu.CpuPlatform"
 
diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py
index ecd2dc5fe623..ae8501f821ed 100644
--- a/vllm/platforms/cpu.py
+++ b/vllm/platforms/cpu.py
@@ -181,9 +181,7 @@ def get_device_total_memory(cls, device_id: int = 0) -> int:
 
     @classmethod
     def set_device(cls, device: torch.device) -> None:
-        """
-        Set the device for the current platform.
-        """
+        """Set the device for the current platform."""
         torch.cpu.set_device(device)
 
     @classmethod
@@ -201,6 +199,15 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         if model_config is not None:
             model_config.disable_cascade_attn = True
 
+        # Import lazily: vllm.triton_utils imports vllm.platforms.current_platform,
+        # which is still being resolved while this platform class is loading.
+        from vllm.triton_utils import HAS_TRITON
+
+        if cls.get_cpu_architecture() == CpuArchEnum.X86 and not HAS_TRITON:
+            logger.warning_once(
+                "Triton is not installed. triton-cpu is expected on x86 CPUs."
+            )
+
         cache_config = vllm_config.cache_config
 
         is_deepseek_v4 = (
@@ -397,11 +404,23 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         # Avoid inductor generates num_thread() and breaks the thread binding
         os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"
 
-        # For efficient conv state memory access. The C++ causal_conv1d
-        # kernels (VDPBF16PS, no AMX tiles) consume the SD layout on any
-        # AVX-512BF16 CPU, so apply it beyond AMX (e.g. AMD Zen5/Turin).
-        if torch.cpu._is_avx512_bf16_supported():
-            os.environ["VLLM_SSM_CONV_STATE_LAYOUT"] = "SD"
+        # NIXL's Mamba descriptors require DS conv state storage. Select it
+        # before cache shapes are created, while preserving an explicit layout.
+        conv_state_layout_env = "VLLM_SSM_CONV_STATE_LAYOUT"
+        if conv_state_layout_env not in os.environ:
+            kv_transfer_config = vllm_config.kv_transfer_config
+            uses_nixl = kv_transfer_config is not None and any(
+                kv_transfer_config.has_connector(name)
+                for name in (
+                    "NixlConnector",
+                    "NixlPullConnector",
+                    "NixlPushConnector",
+                )
+            )
+            if uses_nixl:
+                os.environ[conv_state_layout_env] = "DS"
+            elif torch.cpu._is_avx512_bf16_supported():
+                os.environ[conv_state_layout_env] = "SD"
 
         ld_preload_str = os.getenv("LD_PRELOAD", "")
         cpu_architecture = Platform.get_cpu_architecture()
@@ -500,16 +519,15 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
             return
 
         # reconcile attention and mamba page sizes
-        backend_cls = cls._find_non_ssm_backend(vllm_config)
-        if backend_cls is None:
+        backend_classes = cls._find_non_ssm_backends(vllm_config)
+        if not backend_classes:
             return
 
-        cls._align_hybrid_block_size(vllm_config, backend_cls)
+        cls._align_hybrid_block_size(vllm_config, backend_classes[0])
 
     @classmethod
     def discover_numa_topology(cls) -> list[list[int]]:
-        """
-        Discover NUMA topology and keep the last physical core of each numa
+        """Discover NUMA topology and keep the last physical core of each numa
         into one core group list for nixl start_kv_load()
         """
         SYS_NODE = "/sys/devices/system/node"
@@ -518,6 +536,7 @@ def discover_numa_topology(cls) -> list[list[int]]:
         if not (os.path.exists(SYS_NODE) and os.path.exists(SYS_CPU)):
             return []
 
+        use_highest_sibling = cls.get_cpu_architecture() == CpuArchEnum.X86
         core_rsv_for_kv = []
         for node in os.listdir(SYS_NODE):
             if not node.startswith("node") or not node[4:].isdigit():
@@ -550,7 +569,7 @@ def discover_numa_topology(cls) -> list[list[int]]:
                 else:
                     siblings = [cpu_id]
 
-                phys = min(siblings)
+                phys = max(siblings) if use_highest_sibling else min(siblings)
 
                 if phys not in seen_phys:
                     seen_phys.add(phys)
@@ -570,9 +589,7 @@ def get_punica_wrapper(cls) -> str:
 
     @classmethod
     def get_device_communicator_cls(cls) -> str:
-        """
-        Get device specific communicator class for distributed communication.
-        """
+        """Get device specific communicator class for distributed communication."""
         return "vllm.distributed.device_communicators.cpu_communicator.CpuCommunicator"  # noqa
 
     @classmethod
@@ -605,26 +622,30 @@ def import_kernels(cls) -> None:
                     try:
                         import vllm._C  # noqa: F401
                     except ImportError as e:
-                        logger.warning_once("Failed to import from vllm._C: %r", e)
+                        logger.warning_once(
+                            "Failed to import from vllm._C: %s", repr(e)
+                        )
                 else:
                     try:
                         import vllm._C_AVX512  # noqa: F401
                     except ImportError as e:
                         if ignored_msg not in e.msg:
                             logger.warning_once(
-                                "Failed to import from vllm._C_AVX512: %r", e
+                                "Failed to import from vllm._C_AVX512: %s", repr(e)
                             )
             else:
                 try:
                     import vllm._C_AVX2  # noqa: F401
                 except ImportError as e:
                     if ignored_msg not in e.msg:
-                        logger.warning_once("Failed to import from vllm._C_AVX2: %r", e)
+                        logger.warning_once(
+                            "Failed to import from vllm._C_AVX2: %s", repr(e)
+                        )
         else:
             try:
                 import vllm._C  # noqa: F401
             except ImportError as e:
-                logger.warning_once("Failed to import from vllm._C: %r", e)
+                logger.warning_once("Failed to import from vllm._C: %s", repr(e))
 
     @classmethod
     def pack_kv_cache(
@@ -632,9 +653,7 @@ def pack_kv_cache(
         kv_cache: torch.Tensor,
         indices: torch.Tensor,
     ) -> None:
-        """
-        Rewrite the kv cache shape for the current platform.
-        """
+        """Rewrite the kv cache shape for the current platform."""
         # Import lazily: cpu_attn pulls in _custom_ops, which needs a fully
         # initialized vllm.platforms (avoid circular import while CpuPlatform loads).
         from vllm._custom_ops import cpu_attn_reshape_and_cache
diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py
index 1e5a5a3ad02b..fb1540b60066 100644
--- a/vllm/platforms/cuda.py
+++ b/vllm/platforms/cuda.py
@@ -87,6 +87,7 @@ def _get_backend_priorities(
     kv_cache_dtype: CacheDType | None = None,
     use_non_causal: bool = False,
     head_size: int | None = None,
+    use_mm_prefix: bool = False,
 ) -> list[AttentionBackendEnum]:
     """Get backend priorities with lazy import to avoid circular dependency."""
     from vllm.utils.torch_utils import is_quantized_kv_cache
@@ -156,6 +157,7 @@ def _get_backend_priorities(
         # So prefer FlashAttention when non-causal on SM100f.
         if device_capability.major == 10 and not use_non_causal:
             return [
+                *([AttentionBackendEnum.TRITON_FLASHINFER] if use_mm_prefix else []),
                 AttentionBackendEnum.FLASHINFER,
                 AttentionBackendEnum.FLASH_ATTN,
                 AttentionBackendEnum.TRITON_ATTN,
@@ -164,6 +166,11 @@ def _get_backend_priorities(
             ]
         else:
             return [
+                *(
+                    [AttentionBackendEnum.TRITON_FLASH_ATTN]
+                    if device_capability.major == 9 and use_mm_prefix
+                    else []
+                ),
                 AttentionBackendEnum.FLASH_ATTN,
                 AttentionBackendEnum.FLASHINFER,
                 AttentionBackendEnum.TRITON_ATTN,
@@ -232,7 +239,9 @@ def import_kernels(cls) -> None:
         try:
             import vllm._C_stable_libtorch  # noqa: F401
         except ImportError as e:
-            logger.warning_once("Failed to import from vllm._C_stable_libtorch: %r", e)
+            logger.warning_once(
+                "Failed to import from vllm._C_stable_libtorch: %s", repr(e)
+            )
         with contextlib.suppress(ImportError):
             import vllm._moe_C_stable_libtorch  # noqa: F401
         with contextlib.suppress(ImportError):
@@ -256,9 +265,7 @@ def supported_dtypes(self) -> list[torch.dtype]:
 
     @classmethod
     def set_device(cls, device: torch.device) -> None:
-        """
-        Set the device for the current platform.
-        """
+        """Set the device for the current platform."""
         torch.cuda.set_device(device)
         # With this trick we can force the device to be set eagerly
         # see https://github.com/pytorch/pytorch/issues/155668
@@ -322,12 +329,6 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         parallel_config = vllm_config.parallel_config
         model_config = vllm_config.model_config
 
-        if (
-            parallel_config.prefill_context_parallel_size > 1
-            and parallel_config.data_parallel_size > 1
-        ):
-            raise ValueError("PCP does not support data parallelism on CUDA yet.")
-
         if parallel_config.worker_cls == "auto":
             parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
 
@@ -394,6 +395,7 @@ def get_valid_backends(
             kv_cache_dtype=attn_selector_config.kv_cache_dtype,
             use_non_causal=attn_selector_config.use_non_causal,
             head_size=attn_selector_config.head_size,
+            use_mm_prefix=attn_selector_config.use_mm_prefix,
         )
         for priority, backend in enumerate(backend_priorities):
             try:
@@ -747,7 +749,10 @@ def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConf
             rms_norm = ["oink"] + default
 
         return IrOpPriorityConfig.with_default(
-            default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
+            default,
+            rms_norm=rms_norm,
+            fused_add_rms_norm=rms_norm,
+            gelu_and_mul_sparse=["triton", "native"],
         )
 
     @classmethod
@@ -825,9 +830,7 @@ def get_device_total_memory(cls, device_id: int = 0) -> int:
     @classmethod
     @with_nvml_context
     def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
-        """
-        query if the set of gpus are fully connected by nvlink (1 hop)
-        """
+        """Query if the set of gpus are fully connected by nvlink (1 hop)."""
         handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
         for i, handle in enumerate(handles):
             for j, peer_handle in enumerate(handles):
diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py
index 7ad0b42da774..76d96846e155 100644
--- a/vllm/platforms/interface.py
+++ b/vllm/platforms/interface.py
@@ -3,6 +3,8 @@
 import contextlib
 import enum
 import functools
+import itertools
+import math
 import os
 import platform
 import sys
@@ -122,8 +124,7 @@ def as_version_str(self) -> str:
         return f"{self.major}.{self.minor}"
 
     def to_int(self) -> int:
-        """
-        Express device capability as an integer `<major><minor>`.
+        """Express device capability as an integer `<major><minor>`.
 
         It is assumed that the minor version is always a single digit.
         """
@@ -175,7 +176,7 @@ class Platform:
 
     @property
     def pass_key(self) -> str:
-        """Inductor config key for the PassManager custom pass"""
+        """Inductor config key for the PassManager custom pass."""
         return "post_grad_custom_post_pass"
 
     @property
@@ -238,23 +239,19 @@ def is_cumem_allocator_available(self) -> bool:
 
     @classmethod
     def get_pass_manager_cls(cls) -> str:
-        """
-        Get the pass manager class for this platform.
+        """Get the pass manager class for this platform.
         It will be registered as a custom pass under the current_platform.pass_key.
         """
         return "vllm.compilation.passes.pass_manager.PostGradPassManager"
 
     @classmethod
     def get_compile_backend(cls) -> str:
-        """
-        Get the custom compile backend for current platform.
-        """
+        """Get the custom compile backend for current platform."""
         return cls.simple_compile_backend
 
     @classmethod
     def import_ir_kernels(cls) -> None:
-        """
-        The default implementation imports ``vllm.kernels``, which registers
+        """The default implementation imports ``vllm.kernels``, which registers
         the built-in IR op implementations. Out-of-tree (OOT) platforms should
         override this method to import their own kernel modules.
         """
@@ -365,7 +362,7 @@ def import_kernels(cls) -> None:
         try:
             import vllm._C  # noqa: F401
         except ImportError as e:
-            logger.warning_once("Failed to import from vllm._C: %r", e)
+            logger.warning_once("Failed to import from vllm._C: %s", repr(e))
         with contextlib.suppress(ImportError):
             import vllm._moe_C_stable_libtorch  # noqa: F401
 
@@ -392,8 +389,7 @@ def get_vit_attn_backend(
         dtype: torch.dtype,
         backend: "AttentionBackendEnum | None" = None,
     ) -> "AttentionBackendEnum":
-        """
-        Get the vision attention backend class of a device.
+        """Get the vision attention backend class of a device.
 
         NOTE: ViT Attention should be checked and override in the platform-specific
         implementation. we should not override this in any other places, like
@@ -426,6 +422,7 @@ def get_device_capability(
         Args:
             device_id: Device index in the visible device namespace, matching
                 the argument accepted by torch.cuda.
+
         """
         return None
 
@@ -435,8 +432,7 @@ def has_device_capability(
         capability: tuple[int, int] | int,
         device_id: int = 0,
     ) -> bool:
-        """
-        Test whether this platform is compatible with a device capability.
+        """Test whether this platform is compatible with a device capability.
 
         The `capability` argument can either be:
 
@@ -459,8 +455,7 @@ def is_device_capability(
         capability: tuple[int, int] | int,
         device_id: int = 0,
     ) -> bool:
-        """
-        Test whether this platform has exactly the specified device capability.
+        """Test whether this platform has exactly the specified device capability.
 
         The `capability` argument can either be:
 
@@ -483,8 +478,7 @@ def is_device_capability_family(
         capability: int,
         device_id: int = 0,
     ) -> bool:
-        """
-        Returns True if the device capability is any <major>.x.
+        """Returns True if the device capability is any <major>.x.
         Mirrors CUDA 13 'family' architecture semantics (e.g. 10.x, 11.x, 12.x).
         """
         current_capability = cls.get_device_capability(device_id=device_id)
@@ -532,9 +526,7 @@ def inference_mode(cls):
 
     @classmethod
     def set_device(cls, device: torch.device) -> None:
-        """
-        Set the device for the current platform.
-        """
+        """Set the device for the current platform."""
         raise NotImplementedError
 
     @classmethod
@@ -546,8 +538,7 @@ def manual_seed_all(cls, seed: int) -> None:
     def pre_register_and_update(
         cls, parser: FlexibleArgumentParser | None = None
     ) -> None:
-        """
-        Do some pre-registration or update action for the current platform.
+        """Do some pre-registration or update action for the current platform.
 
         This function is called before global VllmConfig is initialized or cli
         arguments are parsed. It's used for out-of-tree platforms to register or
@@ -560,8 +551,7 @@ def pre_register_and_update(
 
     @classmethod
     def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
-        """
-        Apply the platform-specific default values to the config.
+        """Apply the platform-specific default values to the config.
 
         This function is called during the initialization of global VllmConfig, after
         parsing cli arguments.
@@ -574,8 +564,7 @@ def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
-        """
-        Check and update the configuration for the current platform.
+        """Check and update the configuration for the current platform.
 
         It can raise an exception if the configuration is not compatible with
         the current platform, or it can update the configuration to make it
@@ -586,10 +575,10 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
         pass
 
     @classmethod
-    def _find_non_ssm_backend(
+    def _find_non_ssm_backends(
         cls, vllm_config: "VllmConfig"
-    ) -> "type[AttentionBackend] | None":
-        """Find the first non-SSM attention backend from model layers."""
+    ) -> "list[type[AttentionBackend]]":
+        """Distinct non-SSM attention backends, in layer order."""
         from vllm.config.vllm import get_layers_from_vllm_config
         from vllm.model_executor.layers.attention_layer_base import (
             AttentionLayerBase,
@@ -599,20 +588,69 @@ def _find_non_ssm_backend(
             vllm_config,
             AttentionLayerBase,  # type: ignore[type-abstract]
         )
+        backends: list[type[AttentionBackend]] = []
         for layer in attn_layers.values():
             b = layer.get_attn_backend()
-            if not b.is_ssm():
-                return b
-        return None
+            if not b.is_ssm() and b not in backends:
+                backends.append(b)
+        return backends
 
     @classmethod
-    def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
+    def _preferred_block_size_for_backends(
+        cls,
+        backend_classes: "list[type[AttentionBackend]]",
+        default_block_size: int,
+        vllm_config: "VllmConfig",
+    ) -> int:
+        """Smallest block size every backend accepts.
+
+        ``supports_block_size`` may be overridden to accept exact sizes only
+        (CPU_MLA takes 16 and no multiple of it), so candidates are the LCMs of
+        one supported size per backend, each checked against every backend.
+
+        Raises:
+            ValueError: If the backends share no supported block size.
+
         """
-        Ensure block_size is compatible with the attention backend.
+        from vllm.config.vllm import set_current_vllm_config
+        from vllm.v1.attention.backend import MultipleOf
+
+        # Backends may read the current config to decide a preference.
+        with set_current_vllm_config(vllm_config):
+            if len(backend_classes) == 1:
+                return backend_classes[0].get_preferred_block_size(default_block_size)
+            if all(b.supports_block_size(default_block_size) for b in backend_classes):
+                return default_block_size
+            # A backend declaring no sizes accepts any, so it contributes 1.
+            per_backend_sizes = [
+                [
+                    s.base if isinstance(s, MultipleOf) else s
+                    for s in b.get_supported_kernel_block_sizes()
+                ]
+                or [1]
+                for b in backend_classes
+            ]
+            candidates = sorted(
+                {math.lcm(*sizes) for sizes in itertools.product(*per_backend_sizes)}
+            )
+            for candidate in candidates:
+                if all(b.supports_block_size(candidate) for b in backend_classes):
+                    return candidate
+        raise ValueError(
+            "The attention backends share no supported KV cache block size ("
+            + "; ".join(
+                f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}"
+                for b in backend_classes
+            )
+            + ")."
+        )
+
+    @classmethod
+    def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
+        """Ensure block_size is compatible with the attention backend.
         For hybrid models, also aligns block_size with mamba page sizes.
         """
         from vllm.config.cache import CacheConfig
-        from vllm.config.vllm import set_current_vllm_config
 
         cache_config = vllm_config.cache_config
         model_config = vllm_config.model_config
@@ -621,34 +659,36 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
         if not model_config:
             return
 
-        backend_cls = cls._find_non_ssm_backend(vllm_config)
-        if backend_cls is None:
+        backend_classes = cls._find_non_ssm_backends(vllm_config)
+        if not backend_classes:
             return
 
-        # Phase 1: Pick block size from backend (skip if user set --block-size)
+        # Phase 1: Pick a block size every attention backend supports (skip if
+        # user set --block-size). Models can mix backends with disjoint
+        # preferences, and a size from the first alone later fails
+        # select_common_block_size().
         if not cache_config.user_specified_block_size:
-            with set_current_vllm_config(vllm_config):
-                preferred = backend_cls.get_preferred_block_size(
-                    CacheConfig.DEFAULT_BLOCK_SIZE
-                )
+            preferred = cls._preferred_block_size_for_backends(
+                backend_classes, CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config
+            )
             if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:
                 logger.info(
-                    "Setting kv cache block size to %d for %s backend.",
+                    "Setting kv cache block size to %d for %s backend(s).",
                     preferred,
-                    backend_cls.get_name(),
+                    "/".join(b.get_name() for b in backend_classes),
                 )
             cache_config.block_size = preferred
 
         # Phase 2: Align block/mamba sizes for hybrid models
         # (may override user settings).
         if model_config.is_hybrid:
-            cls._align_hybrid_block_size(vllm_config, backend_cls)
+            cls._align_hybrid_block_size(vllm_config, backend_classes[0])
 
         # Phase 3: Align block/page sizes when multiple KV dtypes share the
         # block pool (e.g. nvfp4 primary + unquantized skip layers).
         # May override the user's --block-size.
         if cache_config.kv_cache_dtype_skip_layers:
-            cls._align_heterogeneous_kv_block_size(vllm_config, backend_cls)
+            cls._align_heterogeneous_kv_block_size(vllm_config, backend_classes[0])
 
     @classmethod
     def _align_heterogeneous_kv_block_size(
@@ -780,8 +820,7 @@ def _align_hybrid_block_size(
         vllm_config: "VllmConfig",
         backend_cls: "type[AttentionBackend]",
     ) -> None:
-        """
-        For hybrid attention/mamba models, ensure that the attention page
+        """For hybrid attention/mamba models, ensure that the attention page
         size is >= the mamba page size, and pad the mamba page size to match.
         """
         from math import lcm
@@ -963,15 +1002,12 @@ def _align_hybrid_block_size(
 
     @classmethod
     def register_custom_kv_cache_specs(cls, vllm_config: "VllmConfig") -> None:
-        """
-        Register custom KVCacheSpec class on current platform.
-        """
+        """Register custom KVCacheSpec class on current platform."""
         pass
 
     @classmethod
     def verify_model_arch(cls, model_arch: str) -> None:
-        """
-        Verify whether the current platform supports the specified model
+        """Verify whether the current platform supports the specified model
         architecture.
 
         - This will raise an Error or Warning based on the model support on
@@ -982,9 +1018,7 @@ def verify_model_arch(cls, model_arch: str) -> None:
 
     @classmethod
     def verify_quantization(cls, quant: str) -> None:
-        """
-        Verify whether the quantization is supported by the current platform.
-        """
+        """Verify whether the quantization is supported by the current platform."""
         if cls.supported_quantization and quant not in cls.supported_quantization:
             raise ValueError(
                 f"{quant} quantization is currently not supported in {cls.device_name}."
@@ -992,8 +1026,7 @@ def verify_quantization(cls, quant: str) -> None:
 
     @classmethod
     def get_cpu_architecture(cls) -> CpuArchEnum:
-        """
-        Determine the CPU architecture of the current system.
+        """Determine the CPU architecture of the current system.
         Returns CpuArchEnum indicating the architecture type.
         """
         machine = platform.machine().lower()
@@ -1031,50 +1064,37 @@ def is_pin_memory_available(cls) -> bool:
     def get_current_memory_usage(
         cls, device: torch.types.Device | None = None
     ) -> float:
-        """
-        Return the memory usage in bytes.
-        """
+        """Return the memory usage in bytes."""
         raise NotImplementedError
 
     @classmethod
     def get_punica_wrapper(cls) -> str:
-        """
-        Return the punica wrapper for current platform.
-        """
+        """Return the punica wrapper for current platform."""
         raise NotImplementedError
 
     @classmethod
     def get_infinity_values(cls, dtype: torch.dtype) -> tuple[float, float]:
-        """
-        Return the platform specific values for (-inf, inf)
-        """
+        """Return the platform specific values for (-inf, inf)."""
         return float("-inf"), float("inf")
 
     @classmethod
     def can_update_inplace(cls) -> bool:
-        """
-        Checks if the platform allows inplace memory updates
-        """
+        """Checks if the platform allows inplace memory updates."""
         return True
 
     @classmethod
     def get_lora_vocab_padding_size(cls) -> int:
-        """
-        Returns how much padding the LoRA logits need for kernels
-        """
+        """Returns how much padding the LoRA logits need for kernels."""
         return 256
 
     @classmethod
     def get_device_communicator_cls(cls) -> str:
-        """
-        Get device specific communicator class for distributed communication.
-        """
+        """Get device specific communicator class for distributed communication."""
         return "vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase"  # noqa
 
     @classmethod
     def is_integrated_gpu(cls, device_id: int = 0) -> bool:
-        """
-        Returns whether the GPU is an integrated (UMA) device that shares
+        """Returns whether the GPU is an integrated (UMA) device that shares
         system memory with the CPU.
 
         On UMA systems (e.g. NVIDIA GH200, DGX Spark, Jetson Orin),
@@ -1085,22 +1105,17 @@ def is_integrated_gpu(cls, device_id: int = 0) -> bool:
 
     @classmethod
     def supports_mx(cls) -> bool:
-        """
-        Returns whether the current platform supports MX types.
-        """
+        """Returns whether the current platform supports MX types."""
         return False
 
     @classmethod
     def supports_fp8(cls) -> bool:
-        """
-        Returns whether the current platform supports FP8 types.
-        """
+        """Returns whether the current platform supports FP8 types."""
         return False
 
     @classmethod
     def is_fp8_fnuz(cls) -> bool:
-        """
-        Returns whether the preferred FP8 type is FNUZ on the current platform.
+        """Returns whether the preferred FP8 type is FNUZ on the current platform.
 
         There are two representations of FP8, OCP FP8 and FNUZ FP8.
         The OCP specification can be found at https://tinyurl.com/b7jvwpft.
@@ -1113,8 +1128,7 @@ def is_fp8_fnuz(cls) -> bool:
 
     @classmethod
     def fp8_dtype(cls) -> torch.dtype:
-        """
-        Returns the preferred FP8 type on the current platform.
+        """Returns the preferred FP8 type on the current platform.
 
         See the documentation for is_fp8_fnuz for details.
         """
@@ -1122,22 +1136,17 @@ def fp8_dtype(cls) -> torch.dtype:
 
     @classmethod
     def use_all_gather(cls) -> bool:
-        """
-        Whether to use allgather in LogitsProcessor to gather the logits.
-        """
+        """Whether to use allgather in LogitsProcessor to gather the logits."""
         return True
 
     @classmethod
     def use_custom_allreduce(cls) -> bool:
-        """
-        Returns if custom allreduce is supported on the current platform
-        """
+        """Returns if custom allreduce is supported on the current platform."""
         return False
 
     @classmethod
     def opaque_attention_op(cls) -> bool:
-        """
-        Returns True if we register attention as one giant opaque custom op
+        """Returns True if we register attention as one giant opaque custom op
         on the current platform
         """
         return False
@@ -1148,7 +1157,7 @@ def validate_request(
         processed_inputs: "EngineInput",
         params: "SamplingParams | PoolingParams",
     ) -> None:
-        """Raises if this request is unsupported on this platform"""
+        """Raises if this request is unsupported on this platform."""
 
     def __getattr__(self, key: str):
         # Pickle checks dunder methods like __getstate__. If we return None
@@ -1172,9 +1181,7 @@ def __getattr__(self, key: str):
         return None
 
     def get_global_graph_pool(self) -> Any:
-        """
-        Return the global graph pool for this platform.
-        """
+        """Return the global graph pool for this platform."""
         cls = self.__class__
         if cls._global_graph_pool is None:
             cls._global_graph_pool = self.graph_pool_handle()
@@ -1182,9 +1189,7 @@ def get_global_graph_pool(self) -> Any:
 
     @classmethod
     def get_static_graph_wrapper_cls(cls) -> str:
-        """
-        Get static graph wrapper class for static graph.
-        """
+        """Get static graph wrapper class for static graph."""
         return "vllm.compilation.base_static_graph.AbstractStaticGraphWrapper"
 
     @classmethod
@@ -1196,36 +1201,27 @@ def stateless_init_device_torch_dist_pg(
         group_size: int,
         timeout: timedelta,
     ) -> "ProcessGroup":
-        """
-        Init platform-specific torch distributed process group.
-        """
+        """Init platform-specific torch distributed process group."""
         raise NotImplementedError
 
     @classmethod
     def check_if_supports_dtype(cls, dtype: torch.dtype):
-        """
-        Check if the dtype is supported by the current platform.
-        """
+        """Check if the dtype is supported by the current platform."""
         raise NotImplementedError
 
     @classmethod
     def support_hybrid_kv_cache(cls) -> bool:
-        """
-        Returns if the hybrid kv cache is supported by the current platform.
-        """
+        """Returns if the hybrid kv cache is supported by the current platform."""
         return False
 
     @classmethod
     def support_static_graph_mode(cls) -> bool:
-        """
-        Returns if the graph mode is supported by the current platform.
-        """
+        """Returns if the graph mode is supported by the current platform."""
         return False
 
     @classmethod
     def check_runner_kv_caches_multi_layer(cls) -> None:
-        """
-        Check whether the platform's ModelRunner can handle multiple attention
+        """Check whether the platform's ModelRunner can handle multiple attention
         layers that share the same layer index (e.g. cross attention and self
         attention in the same decoder block of an encoder-decoder model such as
         BART).
@@ -1241,15 +1237,12 @@ def check_runner_kv_caches_multi_layer(cls) -> None:
 
     @classmethod
     def support_deep_gemm(cls) -> bool:
-        """
-        Returns if DeepGEMM is supported by the current platform.
-        """
+        """Returns if DeepGEMM is supported by the current platform."""
         return False
 
     @classmethod
     def use_custom_op_collectives(cls) -> bool:
-        """
-        Whether this platform should use torch.ops.vllm.* custom ops for collectives.
+        """Whether this platform should use torch.ops.vllm.* custom ops for collectives.
 
         Returns False by default - platforms must explicitly opt-in.
         """
@@ -1257,16 +1250,12 @@ def use_custom_op_collectives(cls) -> bool:
 
     @classmethod
     def use_sync_weight_loader(cls) -> bool:
-        """
-        Returns if the current platform needs to sync weight loader.
-        """
+        """Returns if the current platform needs to sync weight loader."""
         return False
 
     @classmethod
     def make_synced_weight_loader(cls, original_weight_loader):
-        """
-        Wrap the original weight loader to make it synced.
-        """
+        """Wrap the original weight loader to make it synced."""
         if not cls.use_sync_weight_loader():
             return original_weight_loader
 
@@ -1280,37 +1269,29 @@ def _synced_weight_loader(param, *args, **kwargs):
 
     @classmethod
     def get_nixl_supported_devices(cls) -> dict[str, tuple[str, ...]]:
-        """
-        Returns a mapping from device_type to a tuple of supported
+        """Returns a mapping from device_type to a tuple of supported
         kv_buffer_device for nixl.
         """
         return {}
 
     @classmethod
     def get_nixl_memory_type(cls) -> str | None:
-        """
-        Returns the nixl memory type for the current platform.
-        """
+        """Returns the nixl memory type for the current platform."""
         return None
 
     @classmethod
     def check_max_model_len(cls, max_model_len: int) -> int:
-        """
-        Check max_model_len for the current platform.
-        """
+        """Check max_model_len for the current platform."""
         return max_model_len
 
     @classmethod
     def set_additional_forward_context(cls, *args, **kwargs) -> dict[str, Any]:
-        """
-        Set some additional forward context for the current platform if needs.
-        """
+        """Set some additional forward context for the current platform if needs."""
         return {}
 
     @classmethod
     def num_compute_units(cls, device_id: int = 0) -> int:
-        """
-        Get the number of compute units for the current platform.
+        """Get the number of compute units for the current platform.
         (NVIDIA SM / AMD CU / Intel EU)
         """
         raise NotImplementedError(
@@ -1329,11 +1310,25 @@ def get_default_ir_op_priority(
 
     @classmethod
     def is_arch_support_pdl(cls) -> bool:
-        """
-        Does the current platform support PDL (Programmatic Dependent Launch)?
-        """
+        """Does the current platform support PDL (Programmatic Dependent Launch)?"""
         return False
 
+    @classmethod
+    def validate_environ(cls, hard_fail: bool) -> None:
+        """Validate environment variables for the current platform."""
+        from vllm import envs
+
+        for env in os.environ:
+            if env.startswith("VLLM_") and env not in envs.environment_variables:
+                if hard_fail:
+                    raise ValueError(
+                        f"Unknown vLLM environment variable detected: {env}"
+                    )
+                else:
+                    logger.warning(
+                        "Unknown vLLM environment variable detected: %s", env
+                    )
+
 
 class UnspecifiedPlatform(Platform):
     _enum = PlatformEnum.UNSPECIFIED
diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
index bdaa86c5decc..726673b29b71 100644
--- a/vllm/platforms/rocm.py
+++ b/vllm/platforms/rocm.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import importlib.metadata
 import os
 import platform
 from datetime import timedelta
@@ -21,10 +22,15 @@
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
     from vllm.config.kernel import IrOpPriorityConfig
+    from vllm.utils.argparse_utils import FlexibleArgumentParser
+    from vllm.v1.attention.backend import AttentionBackend
     from vllm.v1.attention.selector import AttentionSelectorConfig
 
 logger = init_logger(__name__)
 
+_KV_CACHE_DTYPE_REASON = "kv_cache_dtype not supported"
+_TURBOQUANT_LAYOUT_REASON = "no KV cache layout in common with TURBOQUANT"
+
 try:
     from amdsmi import (
         AmdSmiException,
@@ -130,14 +136,6 @@ def _sync_hip_cuda_env_vars():
     hip_val = os.environ.get("HIP_VISIBLE_DEVICES") or None
     cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES") or None
 
-    if cuda_val is not None:
-        logger.warning_once(
-            "Using CUDA_VISIBLE_DEVICES on ROCm is deprecated and support "
-            "will be removed in vLLM v0.26.0. Please use HIP_VISIBLE_DEVICES "
-            "instead.",
-            scope="process",
-        )
-
     if hip_val is not None and cuda_val is not None:
         if hip_val != cuda_val:
             raise ValueError(
@@ -148,8 +146,6 @@ def _sync_hip_cuda_env_vars():
             )
     elif hip_val is not None:
         os.environ["CUDA_VISIBLE_DEVICES"] = hip_val
-    elif cuda_val is not None:
-        os.environ["HIP_VISIBLE_DEVICES"] = cuda_val
 
 
 # Sync at import time - catches misconfigurations from process start.
@@ -197,8 +193,7 @@ def _query_total_memory_from_amdsmi(physical_device_id: int) -> int:
 
 
 def _get_gcn_arch() -> str:
-    """
-    Get GCN arch via amdsmi (no CUDA init), fallback to torch.cuda.
+    """Get GCN arch via amdsmi (no CUDA init), fallback to torch.cuda.
     Called once at module level; result stored in _GCN_ARCH.
     """
     try:
@@ -233,8 +228,7 @@ def _get_gcn_arch() -> str:
 
 
 def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None:
-    """
-    Parse (major, minor) from a GCN arch string, mirroring how
+    """Parse (major, minor) from a GCN arch string, mirroring how
     HIP derives hipDeviceProp_t.major / .minor.
 
     Format: gfx<MAJOR><MINOR><STEPPING>
@@ -248,6 +242,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None:
     Returns None only when the string is not gfx-prefixed at all
     (i.e. not a ROCm arch string). Raises on any string that looks
     like a GCN arch but does not match a known layout.
+
     """
     m = re.match(r"gfx(\d+)", gcn_arch)
     if not m:
@@ -371,6 +366,25 @@ def get_cdna_version() -> int:
     return 0
 
 
+@cache
+def get_rocm_version() -> tuple[int, ...] | None:
+    """Return the installed ROCm release as (major, minor, patch), or None."""
+    # ROCm 10+ ships as the `rocm` pip SDK; older releases install to /opt/rocm.
+    try:
+        version = importlib.metadata.version("rocm")
+    except importlib.metadata.PackageNotFoundError:
+        rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
+        try:
+            with open(os.path.join(rocm_path, ".info", "version")) as f:
+                version = f.read()
+        except OSError:
+            return None
+    match = re.match(r"\s*(\d+)\.(\d+)(?:\.(\d+))?", version)
+    if match is None:
+        return None
+    return tuple(int(part) for part in match.groups() if part is not None)
+
+
 # Enable HIP online tuning early, before hipBLASLt initializes.
 # Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM.
 if (
@@ -495,6 +509,53 @@ def _get_backend_priorities(
     return backends
 
 
+def _uses_turboquant(vllm_config: "VllmConfig | None") -> bool:
+    """Whether the run's KV cache dtype is one of the turboquant_* presets."""
+    cache_config = getattr(vllm_config, "cache_config", None)
+    return cache_config is not None and str(cache_config.cache_dtype).startswith(
+        "turboquant_"
+    )
+
+
+def _shares_layout_with_turboquant(backend_class: type["AttentionBackend"]) -> bool:
+    """Whether this backend reads a KV cache layout TURBOQUANT reads too.
+
+    A turboquant_* run keeps its boundary layers at the native dtype (see
+    TurboQuantConfig.get_boundary_skip_layers), so those layers pick a backend of
+    their own while every other layer picks TURBOQUANT. One layout has to serve
+    the whole worker, and get_supported_kv_cache_layouts() turns an empty
+    intersection into a hard error at engine startup.
+    """
+    layouts = backend_class.supported_kv_cache_layouts()
+    turboquant_layouts = (
+        AttentionBackendEnum.TURBOQUANT.get_class().supported_kv_cache_layouts()
+    )
+    if layouts is None or turboquant_layouts is None:
+        return True
+    return not set(layouts).isdisjoint(turboquant_layouts)
+
+
+def _get_invalid_reasons(
+    backend_class: type["AttentionBackend"],
+    device_capability: DeviceCapability,
+    attn_selector_config: "AttentionSelectorConfig",
+    *,
+    is_turboquant_run: bool,
+) -> list[str]:
+    """Why this backend cannot serve the layer, empty when it can."""
+    invalid_reasons = backend_class.validate_configuration(
+        device_capability=device_capability,
+        **attn_selector_config._asdict(),
+    )
+    if (
+        not invalid_reasons
+        and is_turboquant_run
+        and not _shares_layout_with_turboquant(backend_class)
+    ):
+        invalid_reasons = [_TURBOQUANT_LAYOUT_REASON]
+    return invalid_reasons
+
+
 class RocmPlatform(Platform):
     _enum = PlatformEnum.ROCM
     device_name: str = "rocm"
@@ -504,6 +565,9 @@ class RocmPlatform(Platform):
     dist_backend: str = "nccl"
     # rocm shares the same device control env var as CUDA
     device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
+    # Set in pre_register_and_update, so it exists only on the driver; Ray
+    # workers are separate processes and copy env vars by allowlist.
+    additional_env_vars: list[str] = ["GPU_PINNED_MIN_XFER_SIZE"]
     ray_noset_device_env_vars: list[str] = [
         "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES",
         "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
@@ -598,12 +662,14 @@ def get_valid_backends(
         # TODO: Make this explicit in the selector in a future PR.
         if is_encoder_decoder and AttentionBackendEnum.ROCM_ATTN in backend_priorities:
             backend_priorities.remove(AttentionBackendEnum.ROCM_ATTN)
+        is_turboquant_run = _uses_turboquant(vllm_config)
         for priority, backend in enumerate(backend_priorities):
             try:
-                backend_class = backend.get_class()
-                invalid_reasons_i = backend_class.validate_configuration(
-                    device_capability=device_capability,
-                    **attn_selector_config._asdict(),
+                invalid_reasons_i = _get_invalid_reasons(
+                    backend.get_class(),
+                    device_capability,
+                    attn_selector_config,
+                    is_turboquant_run=is_turboquant_run,
                 )
             except ImportError:
                 invalid_reasons_i = ["ImportError"]
@@ -626,11 +692,16 @@ def get_attn_backend_cls(
 
         # First try checking just the selected backend, if there is one.
         if selected_backend is not None:
+            # Keep lazy: vllm.config imports current_platform during initialization.
+            from vllm.config import get_current_vllm_config_or_none
+
+            is_turboquant_run = _uses_turboquant(get_current_vllm_config_or_none())
             try:
-                backend_class = selected_backend.get_class()
-                sel_invalid_reasons = backend_class.validate_configuration(
-                    device_capability=device_capability,
-                    **attn_selector_config._asdict(),
+                sel_invalid_reasons = _get_invalid_reasons(
+                    selected_backend.get_class(),
+                    device_capability,
+                    attn_selector_config,
+                    is_turboquant_run=is_turboquant_run,
                 )
             except ImportError:
                 sel_invalid_reasons = ["ImportError"]
@@ -640,21 +711,29 @@ def get_attn_backend_cls(
                     selected_backend.name,
                 )
                 return selected_backend.get_path()
-            # Only tolerate the mismatch for turboquant_* KV-cache layers:
-            # boundary layers keep the native dtype (served by the selected
-            # backend) while turboquant_* layers need TURBOQUANT, so no single
-            # --attention-backend can serve every layer. For any other dtype
-            # the explicit selection is genuinely invalid -> fail loud.
+            # Only tolerate the mismatch when turboquant is in play: boundary
+            # layers keep the native dtype while every other layer needs
+            # TURBOQUANT, so no single --attention-backend can serve every layer.
+            # For any other dtype the selection is genuinely invalid -> fail loud.
             kv_dtype = attn_selector_config.kv_cache_dtype
-            if not (kv_dtype is not None and str(kv_dtype).startswith("turboquant")):
+            layer_is_turboquant = kv_dtype is not None and str(kv_dtype).startswith(
+                "turboquant"
+            )
+            is_turboquant_fallback = (
+                is_turboquant_run or layer_is_turboquant
+            ) and sel_invalid_reasons in (
+                [_KV_CACHE_DTYPE_REASON],
+                [_TURBOQUANT_LAYOUT_REASON],
+            )
+            if not is_turboquant_fallback:
                 raise ValueError(
                     f"Selected backend {selected_backend} is not valid for "
                     f"this configuration. Reason: {sel_invalid_reasons}"
                 )
             # NOTE: pass a str (not the list) -- info_once hashes its args.
             logger.info_once(
-                "Selected backend %s is incompatible with this turboquant "
-                "layer (%s); using the auto-selected per-layer backend. "
+                "Selected backend %s is incompatible with this layer (%s) of "
+                "the turboquant run; using the auto-selected per-layer backend. "
                 "Reason: %s",
                 selected_backend.name,
                 attn_selector_config.attn_type,
@@ -788,9 +867,7 @@ def get_vit_attn_backend(
 
     @classmethod
     def set_device(cls, device: torch.device) -> None:
-        """
-        Set the device for the current platform.
-        """
+        """Set the device for the current platform."""
         torch.cuda.set_device(device)
 
     @classmethod
@@ -815,9 +892,7 @@ def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
     @classmethod
     @with_amdsmi_context
     def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
-        """
-        Query if the set of gpus are fully connected by xgmi (1 hop)
-        """
+        """Query if the set of gpus are fully connected by xgmi (1 hop)."""
         handles = [amdsmi_get_processor_handles()[i] for i in physical_device_ids]
         for i, handle in enumerate(handles):
             for j, peer_handle in enumerate(handles):
@@ -877,6 +952,16 @@ def get_device_total_memory(cls, device_id: int = 0) -> int:
             )
         return torch.cuda.get_device_properties(device_id).total_memory
 
+    @classmethod
+    def pre_register_and_update(
+        cls, parser: "FlexibleArgumentParser | None" = None
+    ) -> None:
+        # Keep mmap'd weight pages on the HIP staging path: above this
+        # threshold the runtime registers the pageable source instead, and each
+        # registration's MMU notifier makes KFD suspend our queues. In KB, so
+        # 4 GiB.
+        os.environ.setdefault("GPU_PINNED_MIN_XFER_SIZE", str(4 * 1024 * 1024))
+
     @classmethod
     def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
         from vllm._aiter_ops import rocm_aiter_ops
@@ -915,12 +1000,6 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
         compilation_config = vllm_config.compilation_config
         parallel_config = vllm_config.parallel_config
 
-        if (
-            parallel_config.prefill_context_parallel_size > 1
-            and parallel_config.data_parallel_size > 1
-        ):
-            raise ValueError("PCP does not support data parallelism on ROCm yet.")
-
         if (
             compilation_config.cudagraph_mode.has_full_cudagraphs()
             and parallel_config.prefill_context_parallel_size > 1
@@ -936,6 +1015,21 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
         if parallel_config.worker_cls == "auto":
             parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
 
+        model_config = vllm_config.model_config
+        scheduler_config = vllm_config.scheduler_config
+        # Note: model_config may be None during testing
+        if (
+            model_config is not None
+            and model_config.is_mm_prefix_lm
+            and scheduler_config.is_multimodal_model
+            and not scheduler_config.disable_chunked_mm_input
+        ):
+            logger.warning_once(
+                "Forcing --disable_chunked_mm_input for models "
+                "with multimodal-bidirectional attention."
+            )
+            scheduler_config.disable_chunked_mm_input = True
+
     @classmethod
     def verify_model_arch(cls, model_arch: str) -> None:
         if model_arch in _ROCM_UNSUPPORTED_MODELS:
@@ -1139,7 +1233,10 @@ def get_default_ir_op_priority(
             rms_norm = default
 
         return IrOpPriorityConfig.with_default(
-            default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
+            default,
+            rms_norm=rms_norm,
+            fused_add_rms_norm=rms_norm,
+            gelu_and_mul_sparse=["native"],
         )
 
     @classmethod
diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py
index 8cb29fbdbd32..d9277fb2c544 100644
--- a/vllm/platforms/xpu.py
+++ b/vllm/platforms/xpu.py
@@ -31,8 +31,7 @@
 def get_mem_info_wrapper(
     device: int | str | torch.device | None = None,
 ) -> tuple[int, int]:
-    """
-    Get memory info for a device, compatible with torch.accelerator.get_memory_info API.
+    """Get memory info for a device, matching `torch.accelerator.get_memory_info`.
 
     Args:
         device: Device specification. Can be:
@@ -43,6 +42,7 @@ def get_mem_info_wrapper(
 
     Returns:
         Tuple[int, int]: (free_memory, total_memory) in bytes
+
     """
     # Handle None - use current device
     if device is None:
@@ -162,6 +162,25 @@ def get_attn_backend_cls(
         if selected_backend == AttentionBackendEnum.TRITON_ATTN:
             logger.info_once("Using Triton backend.")
             return AttentionBackendEnum.TRITON_ATTN.get_path()
+        elif attn_selector_config.use_batch_invariant:
+            # Flash Attention on XPU has not been validated for batch
+            # invariance. Honor an explicit Flash Attention request;
+            # otherwise fall back to Triton Attention, which implements
+            # batch-invariant kernels.
+            if selected_backend == AttentionBackendEnum.FLASH_ATTN:
+                logger.warning_once(
+                    "Using Flash Attention on XPU with batch invariance "
+                    "enabled because it was explicitly requested. This "
+                    "backend has not been validated for batch invariance "
+                    "on XPU and may produce non-deterministic results "
+                    "across batch sizes."
+                )
+                return AttentionBackendEnum.FLASH_ATTN.get_path()
+            logger.info_once(
+                "VLLM_BATCH_INVARIANT is enabled. Using Triton Attention "
+                "backend on XPU, which implements batch-invariant kernels."
+            )
+            return AttentionBackendEnum.TRITON_ATTN.get_path()
         elif attn_selector_config.use_mm_prefix:
             # Flash Attention on XPU has no FA4 kernel, so it cannot apply the
             # multimodal prefix-LM bidirectional mask. Honor an explicit Flash
@@ -237,9 +256,7 @@ def get_vit_attn_backend(
 
     @classmethod
     def set_device(cls, device: torch.device) -> None:
-        """
-        Set the device for the current platform.
-        """
+        """Set the device for the current platform."""
         torch.xpu.set_device(device)
 
     @classmethod
@@ -285,6 +302,27 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         # lazy import to avoid circular import
         from vllm.config import CUDAGraphMode
 
+        if envs.VLLM_BATCH_INVARIANT:
+            model_config = vllm_config.model_config
+            if model_config is not None and (
+                model_config.quantization is not None
+                or vllm_config.quant_config is not None
+            ):
+                raise ValueError(
+                    "XPU batch invariance currently supports only unquantized "
+                    f"models; got quantization={model_config.quantization!r}. "
+                    "Use an unquantized model or disable VLLM_BATCH_INVARIANT."
+                )
+
+            cache_dtype = vllm_config.cache_config.cache_dtype
+            if cache_dtype not in ("auto", "float16", "bfloat16"):
+                raise ValueError(
+                    "XPU batch invariance currently does not support quantized "
+                    f"KV caches; got kv_cache_dtype={cache_dtype!r}. "
+                    "Use an unquantized KV cache dtype or disable "
+                    "VLLM_BATCH_INVARIANT."
+                )
+
         compilation_config = vllm_config.compilation_config
         if compilation_config.compile_sizes is None:
             compilation_config.compile_sizes = []
@@ -298,16 +336,15 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                 "XPU Graph is not supported in the current PyTorch version, "
                 "disabling cudagraph_mode."
             )
-        elif not envs.VLLM_XPU_ENABLE_XPU_GRAPH:
+
+        if (
+            vllm_config.model_config is not None
+            and vllm_config.model_config.enable_sleep_mode
+            and compilation_config.cudagraph_mode != CUDAGraphMode.NONE
+        ):
             compilation_config.cudagraph_mode = CUDAGraphMode.NONE
             logger.warning_once(
-                "XPU Graph is disabled by environment variable, "
-                "please set VLLM_XPU_ENABLE_XPU_GRAPH=1 to enable it."
-            )
-        else:
-            logger.warning_once(
-                "XPU Graph support is experimental and currently only supports "
-                "single-GPU execution."
+                "XPU Graph is not compatible with sleep mode, disabling cudagraph_mode."
             )
 
         # Disable fusion passes not yet supported on XPU.
@@ -491,7 +528,7 @@ def get_default_ir_op_priority(
         using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
         default = ["native"] if using_inductor else ["vllm_c", "native"]
 
-        return IrOpPriorityConfig.with_default(default)
+        return IrOpPriorityConfig.with_default(default, gelu_and_mul_sparse=["native"])
 
     @classmethod
     def device_count(cls) -> int:
diff --git a/vllm/plugins/__init__.py b/vllm/plugins/__init__.py
index 95e895c279be..7bebfdf03e4d 100644
--- a/vllm/plugins/__init__.py
+++ b/vllm/plugins/__init__.py
@@ -115,6 +115,7 @@ def load_endpoint_plugins(
 
     Returns:
         Instantiated plugins that passed gating in discovery order.
+
     """
     from importlib.metadata import entry_points
 
diff --git a/vllm/plugins/lora_resolvers/filesystem_resolver.py b/vllm/plugins/lora_resolvers/filesystem_resolver.py
index 31db0ac4838e..822f410f7dad 100644
--- a/vllm/plugins/lora_resolvers/filesystem_resolver.py
+++ b/vllm/plugins/lora_resolvers/filesystem_resolver.py
@@ -47,8 +47,7 @@ async def _get_lora_req_from_path(
 
 
 def register_filesystem_resolver():
-    """Register the filesystem LoRA Resolver with vLLM"""
-
+    """Register the filesystem LoRA Resolver with vLLM."""
     lora_cache_dir = envs.VLLM_LORA_RESOLVER_CACHE_DIR
     if lora_cache_dir:
         if not os.path.exists(lora_cache_dir) or not os.path.isdir(lora_cache_dir):
diff --git a/vllm/plugins/lora_resolvers/hf_hub_resolver.py b/vllm/plugins/lora_resolvers/hf_hub_resolver.py
index 57ccdc2e48d8..9ab249b1987e 100644
--- a/vllm/plugins/lora_resolvers/hf_hub_resolver.py
+++ b/vllm/plugins/lora_resolvers/hf_hub_resolver.py
@@ -67,6 +67,7 @@ async def _resolve_repo(self, lora_name: str) -> str | None:
             lora_name: Path to LoRA in HF Hub, e.g., <org>/<repo>/<subpath>,
                 match on <org>/<repo> (if it contains an adapter directly) or
                 <org>/<repo>/ if it may have one in subdirs.
+
         """
         for potential_repo in self.repo_list:
             if lora_name.startswith(potential_repo) and (
@@ -86,6 +87,7 @@ async def _resolve_repo_subpath(
         Args:
             lora_name: Path to LoRA in HF Hub, e.g., <org>/<repo>/<subpath>
             maybe_repo: Path to the repo to match against if one exists.
+
         """
         if maybe_repo is None:
             return None
@@ -108,6 +110,7 @@ async def _get_adapter_dirs(self, repo_name: str) -> set[str]:
 
         Args:
             repo_name: Name of the HF hub repo to inspect.
+
         """
         repo_files = await asyncio.to_thread(
             hf_api().list_repo_files,
@@ -124,8 +127,7 @@ async def _get_adapter_dirs(self, repo_name: str) -> set[str]:
 
 
 def register_hf_hub_resolver():
-    """Register the Hf hub LoRA Resolver with vLLM"""
-
+    """Register the Hf hub LoRA Resolver with vLLM."""
     hf_repo_list = envs.VLLM_LORA_RESOLVER_HF_REPO_LIST
     is_enabled = (
         envs.VLLM_PLUGINS is not None and "lora_hf_hub_resolver" in envs.VLLM_PLUGINS
diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py
index 272552d61b8d..319b77d5986a 100644
--- a/vllm/pooling_params.py
+++ b/vllm/pooling_params.py
@@ -28,6 +28,7 @@ class LateInteractionParams(
             - "score_doc": score a document against a cached query.
         query_key: stable key used for both DP routing and worker cache lookup.
         query_uses: expected number of document requests
+
     """
 
     mode: str
@@ -47,6 +48,7 @@ class PoolingParams(
             `None` uses the pooler's default, which is `True` in most cases.
         dimensions: Reduce the dimensions of embeddings
             if model support matryoshka representation.
+
     """
 
     # --8<-- [start:common-pooling-params]
diff --git a/vllm/profiler/layerwise_profile.py b/vllm/profiler/layerwise_profile.py
index a36e4611f3ce..1393e123fbcc 100644
--- a/vllm/profiler/layerwise_profile.py
+++ b/vllm/profiler/layerwise_profile.py
@@ -224,7 +224,7 @@ def _get_kineto_gpu_event(self, node: _ModuleTreeNode):
         return next(iterator, None)
 
     def _cumulative_cuda_time(self, node: _ModuleTreeNode):
-        "Return cuda time in microseconds"
+        """Return cuda time in microseconds."""
 
         def _cumulative_cuda_time_recursive(node: _ModuleTreeNode):
             if node.is_leaf and (gpu_kineto_event := self._get_kineto_gpu_event(node)):
@@ -372,13 +372,13 @@ def df_traversal(
 
 class layerwise_profile(profile):
     def __init__(self, num_running_seqs: int | None = None):
-        """
-        layerwise profile constructor.
+        """Layerwise profile constructor.
 
         Args:
             num_running_seqs (Optional[int], optional): When given,
                 num_running_seqs will be passed to LayerProfileResults
                 for metadata update. Defaults to None.
+
         """
         super().__init__(
             activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py
index 877a390d6d50..bd9b6c231db7 100644
--- a/vllm/profiler/wrapper.py
+++ b/vllm/profiler/wrapper.py
@@ -6,9 +6,9 @@
 import json
 import os
 from abc import ABC, abstractmethod
-from collections.abc import Callable
-from contextlib import nullcontext
-from typing import Literal
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress
+from glob import glob
 from uuid import uuid4
 
 import torch
@@ -17,8 +17,9 @@
 
 import vllm.version
 from vllm.config import ProfilerConfig
-from vllm.config.profiler import _is_uri_path
+from vllm.config.profiler import ProfilerKind, TorchProfilerActivity, _is_uri_path
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
 
@@ -55,6 +56,11 @@ def is_running(self) -> bool:
         """Whether the underlying profiler is currently collecting data."""
         return self._running
 
+    @property
+    def should_annotate(self) -> bool:
+        """Whether worker iterations should receive profiler annotations."""
+        return self._running
+
     @abstractmethod
     def _start(self) -> None:
         """Start the profiler."""
@@ -136,6 +142,7 @@ def _profiler_step(self) -> bool:
         Returns:
             True if the step was an active profiling step (data recorded),
             False if the step was a warmup step (data discarded).
+
         """
         return True
 
@@ -159,12 +166,20 @@ def shutdown(self) -> None:
         if self._running:
             self.stop()
 
+    @property
+    def has_cuda_graph_session(self) -> bool:
+        """Whether a capture-time session is retained to attribute replays."""
+        return False
+
+    def capture_cuda_graphs(self) -> AbstractContextManager[None]:
+        """Observe graph creation for backends that attribute replay activity."""
+        return nullcontext()
+
     def annotate_context_manager(self, name: str):
         """Return a context manager to annotate profiler traces."""
         return nullcontext()
 
 
-TorchProfilerActivity = Literal["CPU", "CUDA", "PrivateUse1", "XPU"]
 TorchProfilerActivityMap = {
     "CPU": torch.profiler.ProfilerActivity.CPU,
     "CUDA": torch.profiler.ProfilerActivity.CUDA,
@@ -179,7 +194,7 @@ def __init__(
         profiler_config: ProfilerConfig,
         worker_name: str,
         local_rank: int,
-        activities: list[TorchProfilerActivity],
+        activities: Sequence[TorchProfilerActivity],
         on_trace_ready: Callable[[torch.profiler.profile], None] | None = None,
     ) -> None:
         super().__init__(profiler_config)
@@ -212,7 +227,12 @@ def __init__(
                 use_gzip=profiler_config.torch_profiler_use_gzip,
             )
 
-        self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1
+        self._records_cpu_activity = "CPU" in activities
+        self.dump_device_time_total = (
+            any(activity != "CPU" for activity in activities)
+            and profiler_config.torch_profiler_dump_cuda_time_total
+        )
+        self.dump_cpu_time_total = self._records_cpu_activity and len(activities) == 1
 
         # Create profiler schedule if warmup or wait iterations are configured
         profiler_schedule = None
@@ -232,7 +252,7 @@ def __init__(
                     profiler_config.active_iterations,
                 )
 
-        self.profiler = torch.profiler.profile(
+        self._profiler_kwargs = dict(
             activities=[TorchProfilerActivityMap[activity] for activity in activities],
             schedule=profiler_schedule,
             record_shapes=profiler_config.torch_profiler_record_shapes,
@@ -241,6 +261,7 @@ def __init__(
             with_flops=profiler_config.torch_profiler_with_flops,
             on_trace_ready=trace_handler,
         )
+        self.profiler: torch.profiler.profile
 
         # Track if we're using a schedule (need to call step())
         self._uses_schedule = profiler_schedule is not None
@@ -248,10 +269,11 @@ def __init__(
         # Subtract 1 because profiler.start() already consumes step 0
         # (WAIT or WARMUP), so only wait + warmup - 1 non-active steps
         # remain to be advanced through via profiler.step() calls.
-        self._warmup_steps_remaining = max(
+        self._initial_warmup_steps_remaining = max(
             profiler_config.wait_iterations + profiler_config.warmup_iterations - 1,
             0,
         )
+        self._warmup_steps_remaining = self._initial_warmup_steps_remaining
         self._version_metadata_added = False
 
     def _build_profiler_table(
@@ -259,9 +281,14 @@ def _build_profiler_table(
         sort_key: str,
         row_limit: int | None = None,
     ) -> str:
+        group_by_input_shape = (
+            current_platform.is_cpu()
+            and self.profiler_config.torch_profiler_record_shapes
+        )
+        averages = self.profiler.key_averages(group_by_input_shape=group_by_input_shape)
         if row_limit is None:  # use profiler default row limit of 100
-            return self.profiler.key_averages().table(sort_by=sort_key)
-        return self.profiler.key_averages().table(
+            return averages.table(sort_by=sort_key)
+        return averages.table(
             sort_by=sort_key,
             row_limit=row_limit,
         )
@@ -302,6 +329,9 @@ def _maybe_add_version_metadata(self) -> None:
 
     @override
     def _start(self) -> None:
+        self.profiler = torch.profiler.profile(**self._profiler_kwargs)
+        self._warmup_steps_remaining = self._initial_warmup_steps_remaining
+        self._version_metadata_added = False
         self.profiler.start()
         # No-schedule case: Kineto is live immediately. With a schedule this
         # no-ops and _profiler_step stamps it once WAIT ends.
@@ -311,10 +341,9 @@ def _start(self) -> None:
     def _stop(self) -> None:
         self.profiler.stop()
 
-        profiler_config = self.profiler_config
         rank = self.local_rank
-        if profiler_config.torch_profiler_dump_cuda_time_total:
-            table = self._build_profiler_table(sort_key="self_cuda_time_total")
+        if self.dump_device_time_total:
+            table = self._build_profiler_table(sort_key="self_device_time_total")
             self._write_profiler_table(rank, table)
 
             # only print profiler results on rank 0
@@ -338,6 +367,7 @@ def _profiler_step(self) -> bool:
         Returns:
             True if the step was an active profiling step (data recorded),
             False if the step was a warmup step (data discarded).
+
         """
         if self._uses_schedule:
             self.profiler.step()
@@ -349,8 +379,15 @@ def _profiler_step(self) -> bool:
                 return False
         return True
 
+    @property
+    @override
+    def should_annotate(self) -> bool:
+        return self._running and self._records_cpu_activity
+
     @override
     def annotate_context_manager(self, name: str):
+        if not self.should_annotate:
+            return nullcontext()
         return torch.profiler.record_function(name)
 
 
@@ -381,6 +418,7 @@ def __init__(
         self._mode = profiler_config.proton_mode
         self._hook = profiler_config.proton_hook
         self._output_format = profiler_config.proton_output_format
+        self._graph_attribution = profiler_config.proton_graph_attribution
         self._triton_version_string = getattr(triton, "__version__", "unknown")
         try:
             self._triton_version = Version(self._triton_version_string)
@@ -392,6 +430,13 @@ def __init__(
         # worker cannot overwrite profiles left by an earlier server process.
         self._instance_id = f"pid{os.getpid()}_{uuid4().hex}"
         self._run_id = 0
+        self._graph_session = False
+        self._phase = 0
+        self._active_output_path: str | None = None
+        self._session_storage_path = os.path.join(
+            self._output_dir,
+            f".proton_cuda_graph_session_{worker_name}_{self._instance_id}",
+        )
 
         logger.info_once(
             "Proton profiling enabled. Output will be saved under: %s",
@@ -406,6 +451,10 @@ def _require_triton_version(self, feature: str, minimum: Version) -> None:
             )
 
     def _validate_capabilities(self) -> None:
+        if self._graph_attribution:
+            self._require_triton_version(
+                "CUDA graph attribution", _TRITON_PROTON_3_7_VERSION
+            )
         if self._output_format is not None:
             parameters = inspect.signature(self._proton.finalize).parameters
             supports_output_format = "output_format" in parameters or any(
@@ -422,7 +471,7 @@ def _validate_capabilities(self) -> None:
             self._require_triton_version(
                 "hatchet_msgpack output", _TRITON_PROTON_3_7_VERSION
             )
-        if self._mode and self._mode.split(":", 1)[0] == "periodic_flushing":
+        if self._mode and self._mode.split(":", 1)[0].lower() == "periodic_flushing":
             self._require_triton_version(
                 "periodic flushing", _TRITON_PROTON_3_7_VERSION
             )
@@ -441,28 +490,112 @@ def _create_session(self, output_path: str) -> int:
             raise RuntimeError("Proton did not create a profiling session")
         return session_id
 
+    @property
+    def has_cuda_graph_session(self) -> bool:
+        return self._graph_session
+
+    def set_output_name(self, worker_name: str) -> None:
+        """Set the next run's output name after startup graph capture."""
+        if self._active:
+            return
+        self._output_path = os.path.join(self._output_dir, f"proton_{worker_name}")
+
+    @override
+    @contextmanager
+    def capture_cuda_graphs(self) -> Iterator[None]:
+        """Keep a Proton session active while vLLM captures CUDA graphs."""
+        if not self._graph_attribution:
+            yield
+            return
+        if self._session_id is None:
+            self._session_id = self._create_session(self._session_storage_path)
+        else:
+            self._proton.activate(session=self._session_id)
+        self._graph_session = True
+
+        try:
+            yield
+        finally:
+            captured_phase = self._phase
+            try:
+                self._phase = self._proton.data.advance_phase(self._session_id)
+            finally:
+                self._proton.deactivate(session=self._session_id, flushing=True)
+            self._proton.data.clear(self._session_id, captured_phase)
+
     @override
     def _start(self) -> None:
-        output_path = f"{self._output_path}_{self._instance_id}_run{self._run_id}"
-        self._session_id = self._create_session(output_path)
+        self._active_output_path = (
+            f"{self._output_path}_{self._instance_id}_run{self._run_id}"
+        )
         self._run_id += 1
+        if self._graph_session:
+            assert self._session_id is not None
+            self._proton.activate(session=self._session_id)
+        else:
+            self._session_id = self._create_session(self._active_output_path)
+
+    def _write_graph_phase(self, phase: int) -> None:
+        assert self._active_output_path is not None
+        output_format = self._output_format or "hatchet"
+        output_path = f"{self._active_output_path}.{output_format}"
+        if output_format == "hatchet_msgpack":
+            with open(output_path, "wb") as output_file:
+                output_file.write(
+                    self._proton.data.get_msgpack(self._session_id, phase)
+                )
+        else:
+            with open(output_path, "w", encoding="utf-8") as output_file:
+                json.dump(self._proton.data.get(self._session_id, phase), output_file)
+
+    def _finalize_session(self, session_id: int) -> None:
+        if self._output_format is None:
+            self._proton.finalize(session=session_id)
+        else:
+            self._proton.finalize(session=session_id, output_format=self._output_format)
 
     @override
     def _stop(self) -> None:
         assert self._session_id is not None
         session_id = self._session_id
+        if self._graph_session:
+            completed_phase = self._phase
+            try:
+                self._phase = self._proton.data.advance_phase(session_id)
+            finally:
+                self._proton.deactivate(session=session_id, flushing=True)
+            try:
+                self._write_graph_phase(completed_phase)
+            finally:
+                self._proton.data.clear(session_id, completed_phase)
+                self._active_output_path = None
+            return
+
         try:
             self._proton.deactivate(session=session_id)
         finally:
             try:
-                if self._output_format is None:
-                    self._proton.finalize(session=session_id)
-                else:
-                    self._proton.finalize(
-                        session=session_id, output_format=self._output_format
-                    )
+                self._finalize_session(session_id)
             finally:
                 self._session_id = None
+                self._active_output_path = None
+
+    @override
+    def shutdown(self) -> None:
+        super().shutdown()
+        if self._graph_session and self._session_id is not None:
+            session_id = self._session_id
+            self._session_id = None
+            try:
+                self._finalize_session(session_id)
+            except Exception:
+                logger.exception(
+                    "Failed to finalize Proton CUDA graph session during shutdown."
+                )
+            finally:
+                for output_path in glob(f"{self._session_storage_path}.*"):
+                    with suppress(FileNotFoundError):
+                        os.remove(output_path)
 
     @override
     def annotate_context_manager(self, name: str):
@@ -490,3 +623,100 @@ def _stop(self) -> None:
     @override
     def annotate_context_manager(self, name: str):
         return torch.cuda.nvtx.range(name)
+
+
+_DEFAULT_TORCH_PROFILER_ACTIVITIES: dict[str, tuple[TorchProfilerActivity, ...]] = {
+    "cpu": ("CPU",),
+    "cuda": ("CPU", "CUDA"),
+    "xpu": ("CPU", "XPU"),
+}
+_SUPPORTED_TORCH_PROFILER_ACTIVITIES = {
+    device: frozenset(activities)
+    for device, activities in _DEFAULT_TORCH_PROFILER_ACTIVITIES.items()
+}
+_SUPPORTED_PROFILER_KINDS: dict[str, frozenset[ProfilerKind]] = {
+    "cpu": frozenset(("torch",)),
+    "cuda": frozenset(("torch", "cuda", "proton")),
+    "xpu": frozenset(("torch",)),
+}
+
+
+def validate_worker_profiler_config(profiler_config: ProfilerConfig) -> None:
+    """Validate profiler selections against the worker's capabilities."""
+    profiler_type = profiler_config.profiler
+    if profiler_type is None:
+        return
+    device_type = current_platform.device_type
+    if device_type not in _SUPPORTED_PROFILER_KINDS:
+        raise ValueError(f"Unsupported profiler device type: {device_type}")
+    supported_kinds = _SUPPORTED_PROFILER_KINDS[device_type]
+    if profiler_type not in supported_kinds:
+        supported_names = ", ".join(sorted(supported_kinds))
+        raise ValueError(
+            f"Unsupported profiler type for {device_type}: "
+            f"{profiler_type}. Supported profiler types: {supported_names}."
+        )
+    if profiler_type == "torch":
+        default_activities = _DEFAULT_TORCH_PROFILER_ACTIVITIES[device_type]
+        supported_activities = _SUPPORTED_TORCH_PROFILER_ACTIVITIES[device_type]
+        configured = profiler_config.torch_profiler_activities
+        activities = default_activities if configured is None else configured
+        unsupported = set(activities) - supported_activities
+        if unsupported:
+            unsupported_names = ", ".join(sorted(unsupported))
+            supported_names = ", ".join(sorted(supported_activities))
+            raise ValueError(
+                f"Unsupported torch profiler activities for "
+                f"{device_type}: {unsupported_names}. "
+                f"Supported activities: {supported_names}."
+            )
+
+
+def create_worker_profiler(
+    profiler_config: ProfilerConfig,
+    *,
+    worker_name: str,
+    local_rank: int,
+) -> WorkerProfiler:
+    """Create a profiler using a validated config and platform defaults."""
+    profiler_type = profiler_config.profiler
+    if profiler_type == "torch":
+        default_activities = _DEFAULT_TORCH_PROFILER_ACTIVITIES[
+            current_platform.device_type
+        ]
+        configured = profiler_config.torch_profiler_activities
+        logger.debug("Starting torch profiler with trace name: %s", worker_name)
+        return TorchProfilerWrapper(
+            profiler_config,
+            worker_name=worker_name,
+            local_rank=local_rank,
+            activities=default_activities if configured is None else tuple(configured),
+        )
+    if profiler_type == "cuda":
+        logger.debug("Starting CUDA profiler")
+        return CudaProfilerWrapper(profiler_config)
+
+    assert profiler_type == "proton", f"Unknown profiler type: {profiler_type}"
+    logger.debug("Starting Proton profiler with trace name: %s", worker_name)
+    return ProtonProfilerWrapper(profiler_config, worker_name=worker_name)
+
+
+def create_graph_capture_profiler(
+    profiler_config: ProfilerConfig, global_rank: int
+) -> WorkerProfiler | None:
+    """Create a profiler to observe CUDA graph capture, if configured.
+
+    Applies only to backends that attribute graph replay activity which
+    need a session around capture (Proton with ``proton_graph_attribution``).
+    """
+    if (
+        profiler_config.profiler == "proton"
+        and profiler_config.proton_graph_attribution
+    ):
+        from vllm.distributed.utils import get_worker_rank_suffix
+
+        return ProtonProfilerWrapper(
+            profiler_config,
+            worker_name=get_worker_rank_suffix(global_rank=global_rank),
+        )
+    return None
diff --git a/vllm/ray/lazy_utils.py b/vllm/ray/lazy_utils.py
index 06c91cc3943a..674e8429e811 100644
--- a/vllm/ray/lazy_utils.py
+++ b/vllm/ray/lazy_utils.py
@@ -16,7 +16,6 @@ def is_ray_initialized():
 
 def is_in_ray_actor():
     """Check if we are in a Ray actor."""
-
     try:
         import ray
 
diff --git a/vllm/ray/ray_env.py b/vllm/ray/ray_env.py
index 4103b93172ce..2e2ed48b85c7 100644
--- a/vllm/ray/ray_env.py
+++ b/vllm/ray/ray_env.py
@@ -76,6 +76,7 @@ def get_env_vars_to_copy(
         additional_vars: Extra individual env var names to copy.  Useful
             for caller-specific vars (e.g. platform env vars).
         destination: Label used in log messages only.
+
     """
     exclude = (exclude_vars or set()) | RAY_NON_CARRY_OVER_ENV_VARS
 
diff --git a/vllm/reasoning/AGENTS.md b/vllm/reasoning/AGENTS.md
new file mode 100644
index 000000000000..1f7c8b3557b1
--- /dev/null
+++ b/vllm/reasoning/AGENTS.md
@@ -0,0 +1,54 @@
+# vllm/reasoning
+
+Every `--reasoning-parser` name resolves here. The parser has two consumers:
+`ParserManager.get_parser` composes it with the tool parser into the unified
+`Parser` the serving layer uses, and the structured-output manager
+(`vllm/v1/structured_output/__init__.py`) instantiates it directly and calls
+`is_reasoning_end` / `is_reasoning_end_streaming` every decode step to decide
+when the grammar bitmask starts applying. The `ReasoningParser` interface
+itself is legacy: it is kept so existing in-tree parsers and out-of-tree
+plugins keep working, not as the base for new ones.
+
+## Layout
+
+- `abs_reasoning_parsers.py` — `ReasoningParser` (`extract_reasoning`,
+  `extract_reasoning_streaming`, `is_reasoning_end`,
+  `is_reasoning_end_streaming`, `extract_content_ids`, `adjust_request`) and
+  `ReasoningParserManager`.
+- `__init__.py` — `_REASONING_PARSERS_TO_REGISTER`, the name → (file, class)
+  table that lazily registers every built-in parser.
+- `basic_parsers.py` — `BaseThinkingReasoningParser`, the base for parsers
+  that only need a start/end token pair.
+- `<model>_reasoning_parser.py` — one file per family. Either a subclass of a
+  `ParserReasoningAdapter` generated from a `vllm/parser/<model>.py` engine
+  grammar (`mistral_reasoning_parser.py`, `kimi_k2_reasoning_parser.py`; the
+  `_engine_` suffix is not applied consistently), or a self-contained
+  `ReasoningParser` subclass.
+
+## Guidelines
+
+- Do not change the `ReasoningParser` interface in `abs_reasoning_parsers.py`.
+  Out-of-tree plugins loaded through
+  `ReasoningParserManager.import_reasoning_parser` subclass it. New
+  capabilities go on `Parser` in `vllm/parser/`; if a legacy parser needs the
+  same capability, `DelegatingParser` supplies the default.
+- For a new model, implement a unified `Parser` in `vllm/parser/` first, then
+  expose it here as a thin adapter subclass plus an entry in
+  `_REASONING_PARSERS_TO_REGISTER`.
+- The `_REASONING_PARSERS_TO_REGISTER` key is what users pass on the command
+  line; treat renames as breaking.
+- Keep adapters thin: registration, `adjust_request` tweaks, and class
+  attributes. Parsing logic belongs in `vllm/parser/`.
+- Keep a model-specific fix inside that model's `<model>_reasoning_parser.py`,
+  or its grammar in `vllm/parser/<model>.py`, where possible. Touch
+  `abs_reasoning_parsers.py` or `basic_parsers.py` only for behavior every
+  parser needs, and say so in the PR.
+- `is_reasoning_end_streaming` runs on the engine core for every decode step
+  of every structured-output request. Work on `delta_ids`; re-scanning the
+  full sequence makes the step cost grow with sequence length.
+- The structured-output gate and the frontend's extraction must agree on
+  where reasoning ends, whether the turn continues with content or a tool
+  call. Test the end-detection methods with the same inputs as extraction.
+- Comment only where the code alone would leave a reader with a question —
+  a request tweak in `adjust_request`, a marker handled unusually, a model
+  quirk being worked around. Keep it minimal and concise.
diff --git a/vllm/reasoning/CLAUDE.md b/vllm/reasoning/CLAUDE.md
new file mode 100644
index 000000000000..43c994c2d361
--- /dev/null
+++ b/vllm/reasoning/CLAUDE.md
@@ -0,0 +1 @@
+@AGENTS.md
diff --git a/vllm/reasoning/abs_reasoning_parsers.py b/vllm/reasoning/abs_reasoning_parsers.py
index 6fb694545a4a..accc8c9c2727 100644
--- a/vllm/reasoning/abs_reasoning_parsers.py
+++ b/vllm/reasoning/abs_reasoning_parsers.py
@@ -23,8 +23,7 @@
 
 
 class ReasoningParser:
-    """
-    Abstract reasoning parser class that should not be used directly.
+    """Abstract reasoning parser class that should not be used directly.
     Provided and methods should be used in derived classes.
 
     It is used to extract reasoning content from the model output.
@@ -71,26 +70,27 @@ def has_engine_confirmed_reasoning_end(self) -> bool:
 
     @abstractmethod
     def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
-        """
-        Check if the reasoning content ends in the input_ids.
+        """Check if the reasoning content ends in the input_ids.
 
         It is used in structured engines like `xgrammar` to check if the
         reasoning content ends in the model output.
 
-        Parameters:
+        Parameters
+        ----------
         input_ids: list[int]
             The input_ids of the model output.
 
-        Returns:
+        Returns
+        -------
         bool
             True if the reasoning content ends in the input_ids.
+
         """
 
     def is_reasoning_end_streaming(
         self, input_ids: Sequence[int], delta_ids: Iterable[int]
     ) -> bool:
-        """
-        Check if the reasoning content ends in the input_ids on a
+        """Check if the reasoning content ends in the input_ids on a
         decode step.
 
         It is used in structured engines like `xgrammar` to check if the
@@ -98,29 +98,36 @@ def is_reasoning_end_streaming(
         `input_ids` the entire model output and `delta_ids` are the last few
         computed tokens of the model output (like during a decode step).
 
-        Parameters:
+        Parameters
+        ----------
         input_ids: list[int]
             The entire model output.
         delta_ids: list[int]
             The last few computed tokens of the model output at the current decode step.
 
-        Returns:
+        Returns
+        -------
         bool
             True if the reasoning content ends in the `delta_ids` on a
             decode step.
+
         """
         return self.is_reasoning_end(input_ids)
 
     @abstractmethod
     def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        """
-        Extract content token ids from the input_ids.
-        Parameters:
+        """Extract content token ids from the input_ids.
+
+        Parameters
+        ----------
         input_ids: list[int]
             The input_ids of the model output.
-        Returns:
+
+        Returns
+        -------
         list[int]
             The extracted content from the input_ids.
+
         """
 
     def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
@@ -137,8 +144,8 @@ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
 
         Returns:
             int: Number of tokens that belong to reasoning content.
-        """
 
+        """
         # By default, assume the parser cannot detect reasoning spans.
         return 0
 
@@ -148,18 +155,18 @@ def extract_reasoning(
         model_output: str,
         request: "ChatCompletionRequest | ResponsesRequest",
     ) -> tuple[str | None, str | None]:
-        """
-        Extract reasoning content from a complete model-generated string.
+        """Extract reasoning content from a complete model-generated string.
 
         Used for non-streaming responses where we have the entire model response
         available before sending to the client.
 
-        Parameters:
+        Args:
             model_output: The model-generated string to extract reasoning content from.
             request: The request object that was used to generate the model_output.
 
         Returns:
             A tuple containing the reasoning content and the content.
+
         """
 
     @abstractmethod
@@ -172,8 +179,7 @@ def extract_reasoning_streaming(
         current_token_ids: Sequence[int],
         delta_token_ids: Sequence[int],
     ) -> "DeltaMessage | None":
-        """
-        Instance method that should be implemented for extracting reasoning
+        """Instance method that should be implemented for extracting reasoning
         from an incomplete response; for use when handling reasoning calls and
         streaming. Has to be an instance method because  it requires state -
         the current tokens/diffs, but also the information about what has
@@ -203,15 +209,12 @@ def prepare_structured_tag(
         original_tag: str | None,
         tool_server: ToolServer | None,
     ) -> str | None:
-        """
-        Instance method that is implemented for preparing the structured tag
-        """
+        """Instance method that is implemented for preparing the structured tag"""
         return original_tag
 
 
 class ReasoningParserManager:
-    """
-    Central registry for ReasoningParser implementations.
+    """Central registry for ReasoningParser implementations.
 
     Supports two registration modes:
       - Eager registration via `register_module`
@@ -225,14 +228,14 @@ class ReasoningParserManager:
 
     @classmethod
     def get_reasoning_parser(cls, name: str) -> type[ReasoningParser]:
-        """
-        Retrieve a registered or lazily registered ReasoningParser class.
+        """Retrieve a registered or lazily registered ReasoningParser class.
 
         If the parser is lazily registered, it will be imported and cached
         on first access.
 
         Raises:
             KeyError: if no parser is found under the given name.
+
         """
         if name in cls.reasoning_parsers:
             return cls.reasoning_parsers[name]
@@ -303,8 +306,7 @@ def _register_module(
 
     @classmethod
     def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None:
-        """
-        Register a lazy module mapping for delayed import.
+        """Register a lazy module mapping for delayed import.
 
         Example:
             ReasoningParserManager.register_lazy_module(
@@ -312,6 +314,7 @@ def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> N
                 module_path="vllm.reasoning.qwen3_engine_reasoning_parser",
                 class_name="Qwen3ParserReasoningAdapter",
             )
+
         """
         cls.lazy_parsers[name] = (module_path, class_name)
 
@@ -324,8 +327,7 @@ def register_module(
     ) -> (
         type[ReasoningParser] | Callable[[type[ReasoningParser]], type[ReasoningParser]]
     ):
-        """
-        Register module with the given name or name list. it can be used as a
+        """Register module with the given name or name list. it can be used as a
         decoder(with module as None) or normal function(with module as not
         None).
         """
diff --git a/vllm/reasoning/basic_parsers.py b/vllm/reasoning/basic_parsers.py
index f85ccd010a8d..e1d309250648 100644
--- a/vllm/reasoning/basic_parsers.py
+++ b/vllm/reasoning/basic_parsers.py
@@ -16,8 +16,7 @@
 
 
 class BaseThinkingReasoningParser(ReasoningParser):
-    """
-    Base class for reasoning parsers that use thinking tokens.
+    """Base class for reasoning parsers that use thinking tokens.
 
     This class provides common functionality for parsers that use start and end
     tokens to delimit reasoning content (
@@ -87,9 +86,7 @@ def is_reasoning_end_streaming(
         return end_token_id in delta_ids
 
     def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        """
-        Extract the content after the end tokens
-        """
+        """Extract the content after the end tokens."""
         if self.end_token_id not in islice(input_ids, 0, max(0, len(input_ids) - 1)):
             return []
         else:
@@ -104,8 +101,7 @@ def extract_reasoning_streaming(
         current_token_ids: Sequence[int],
         delta_token_ids: Sequence[int],
     ) -> DeltaMessage | None:
-        """
-        Extract reasoning content from a delta message.
+        """Extract reasoning content from a delta message.
         Handles streaming output where previous + delta = current.
         Uses token IDs for faster processing.
         """
@@ -157,8 +153,7 @@ def extract_reasoning_streaming(
     def extract_reasoning(
         self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
     ) -> tuple[str | None, str | None]:
-        """
-        Extract reasoning content from the model output.
+        """Extract reasoning content from the model output.
 
         This is the base implementation that works for most models.
         Subclasses can override this method for specific behavior.
diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py
index 4dada76ee7c2..d02091e6003e 100644
--- a/vllm/reasoning/cohere_command_reasoning_parser.py
+++ b/vllm/reasoning/cohere_command_reasoning_parser.py
@@ -3,531 +3,28 @@
 
 from __future__ import annotations
 
-import json
-from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict, TypeGuard
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, ClassVar
 
-import regex as re
-import xgrammar as xgr
-
-try:
-    from cohere_melody import PyFilter, PyFilterOptions
-except ImportError as e:
-    raise ImportError(
-        "The Cohere reasoning parser requires the `cohere_melody` "
-        "package, which is not installed. Install it with:\n"
-        "    pip install 'cohere-melody>=0.11.1'"
-    ) from e
-
-
-from vllm.entrypoints.cohere.cohere_chat_message import (
-    Citation,
-    CitationSource,
-    CohereDeltaMessage,
-)
-from vllm.entrypoints.generate.base.protocol import (
-    AnyResponseFormat,
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-)
-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.reasoning import ReasoningParser
-from vllm.sampling_params import StructuredOutputsParams
-from vllm.tokenizers import TokenizerLike
 
 if TYPE_CHECKING:
-    from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-
-REPLACEMENT_CHAR = "\ufffd"
-
-
-class CohereTagRegistry(NamedTuple):
-    """A single ``structural_tag`` trigger / end pair (``begin`` uses ``trigger``)."""
-
-    trigger: str
-    end: str
-
-
-class CohereTagStyle(NamedTuple):
-    """The structural tags style for a given model architecture.
-
-    ``json_tags`` lists every JSON-schema wrapper the model may emit (MOE uses
-    both response and text delimiters). ``tools`` is the tool-call wrapper.
-    """
-
-    json_tags: tuple[CohereTagRegistry, ...]
-    tools: CohereTagRegistry
-
-
-class CohereNormalizedTool(TypedDict):
-    """A tool definition normalized to the shape ``collect_tool_schema`` expects.
-
-    ``parameters`` is a JSON Schema object (possibly empty) describing the tool's
-    call signature.
-    """
-
-    name: str
-    parameters: dict[str, Any]
-
-
-COMMAND_A_TOOLS_TAG = CohereTagRegistry(
-    trigger="<|START_ACTION|>",
-    end="<|END_ACTION|>",
-)
-COMMAND_A_JSON_TAG = CohereTagRegistry(
-    trigger="<|START_RESPONSE|>",
-    end="<|END_RESPONSE|>",
-)
-COMMAND_A_PLUS_JSON_TAG = CohereTagRegistry(
-    trigger="<|START_TEXT|>",
-    end="<|END_TEXT|>",
-)
-
-MODEL_TO_TAG_STYLE: dict[str, CohereTagStyle] = {
-    "Cohere2ForCausalLM": CohereTagStyle(
-        json_tags=(COMMAND_A_JSON_TAG,),
-        tools=COMMAND_A_TOOLS_TAG,
-    ),
-    "Cohere2VisionForConditionalGeneration": CohereTagStyle(
-        json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
-        tools=COMMAND_A_TOOLS_TAG,
-    ),
-    "Cohere2MoeForCausalLM": CohereTagStyle(
-        json_tags=(COMMAND_A_JSON_TAG, COMMAND_A_PLUS_JSON_TAG),
-        tools=COMMAND_A_TOOLS_TAG,
-    ),
-}
-
-
-def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str:
-    """Build an xgrammar EBNF grammar that matches a JSON array of tool calls.
-
-    The grammar shape is architecture-independent; callers are responsible for
-    wrapping it in the correct structural tag (see ``CohereTagStyle.tools``).
-    """
-    tool_dictionary: dict[str, str] = {}
-    for tool in tool_schema:
-        tool_name = tool["name"]
-        tool_parameters = json.dumps(tool["parameters"])
-        json_schema = f"""{{
-                        "type": "object",
-                        "properties": {{
-                            "tool_call_id": {{
-                                "type": "string",
-                                "pattern": "^[0-9]+$"
-                            }},
-                            "tool_name": {{
-                                "type": "string",
-                                "const": "{tool_name}"
-                            }},
-                            "parameters": {tool_parameters}
-                            }}
-                            }}"""
-        tool_grammar = str(xgr.Grammar.from_json_schema(json_schema))
-        for match in re.findall(r"\b(\w+)\s*::=", tool_grammar):
-            tool_grammar = re.sub(
-                rf"\b{re.escape(match)}\b", tool_name + match, tool_grammar
-            )
-        tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}"
-    # Emitted grammar shape:
-    #   root  ::= tools
-    #   tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws
-    #   ws    ::= (" " | "\t" | "\n")*
-    #   tool  ::= <tool_a> | <tool_b> | ...         (one alternative per input)
-    #   <tool_x>     ::= <tool_x>root               (per-tool xgrammar rules)
-    #   <tool_x>root ::= ...                        (from xgr.Grammar.from_json_schema)
-    tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys())
-    tool_rules = "\n    ".join(tool_dictionary.values())
-    grammar = f"""root ::= tools
-    tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws
-    ws    ::= (" " | "\\t" | "\\n")*
-    {tool_alternatives}
-    {tool_rules}
-    """
-    return grammar
-
-
-def _tool_definitions_to_schema_list(
-    tools: str | list[Any],
-) -> list[CohereNormalizedTool]:
-    """
-    Build the list of ``CohereNormalizedTool`` dicts expected by
-    ``collect_tool_schema``.
-
-    Accepts:
-    - JSON string
-    - list of dicts with top-level ``name`` / ``parameters``
-    - list of Chat Completions-style ``{"type": "function", "function": {...}}``
-    - list of Pydantic models with ``model_dump()``
-    """
-    if isinstance(tools, str):
-        try:
-            parsed = json.loads(tools)
-        except json.JSONDecodeError:
-            return []
-        if not isinstance(parsed, list):
-            return []
-    else:
-        parsed = list(tools)
-
-    out: list[CohereNormalizedTool] = []
-    for raw in parsed:
-        t = raw.model_dump() if hasattr(raw, "model_dump") else raw
-        if not isinstance(t, dict):
-            continue
-        # Unwrap Chat Completions' ``{"type": "function", "function": {...}}``
-        # shape; otherwise take the dict as-is.
-        if t.get("type") == "function" and isinstance(t.get("function"), dict):
-            t = t["function"]
-        name = t.get("name")
-        if not isinstance(name, str):
-            continue
-        params = t.get("parameters")
-        out.append(
-            CohereNormalizedTool(
-                name=name,
-                parameters=params if isinstance(params, dict) else {},
-            )
-        )
-    return out
-
-
-def _has_effective_tools(
-    tools: str | list[Any] | None,
-) -> TypeGuard[str | list[Any]]:
-    """
-    True when ``tools`` contains at least one tool definition to convert.
-
-    ``ResponsesRequest`` defaults ``tools`` to ``[]``; ``ChatCompletionRequest``
-    uses ``None``. Both mean "no tools" here. Strings (e.g. a JSON blob) are
-    treated as effective only when non-blank.
-    """
-    if tools is None:
-        return False
-    if isinstance(tools, str):
-        return bool(tools.strip())
-    return len(tools) > 0
-
-
-# Builder: produces vLLM response_format in xgrammar's canonical format.
-# See xgrammar docs: type "structural_tag" with "format" = triggered_tags
-# and tag content type = json_schema | grammar.
-def convert_schema_to_structural_tags(
-    schema: dict | None = None,
-    tools: str | list[Any] | None = None,
-    model_architecture: str | None = None,
-) -> str | None:
-    """
-    Returns a response_format string accepted by xgrammar's structural tag format.
-    Uses the canonical shape: {"type": "structural_tag", "format": {...}} with
-    format.type "triggered_tags" and tag content type "json_schema" or "grammar".
-
-    Callers that are not on an engine path (e.g. the reasoning parser) must pass
-    ``model_architecture`` explicitly.
-    """
-    if model_architecture is None or model_architecture not in MODEL_TO_TAG_STYLE:
-        return None
-    style = MODEL_TO_TAG_STYLE[model_architecture]
-
-    tags: list[dict] = []
-    triggers: list[str] = []
-
-    def _add_tag(tag: CohereTagRegistry, content: dict) -> None:
-        tags.append({"begin": tag.trigger, "content": content, "end": tag.end})
-        triggers.append(tag.trigger)
-
-    if schema is not None:
-        # One structural tag per JSON wrapper (e.g. MOE: response + text).
-        # Same for schema-only and "tools plus JSON mode" (North: schema when
-        # the model does not call tools).
-        for jt in style.json_tags:
-            _add_tag(jt, {"type": "json_schema", "json_schema": schema})
-
-    if _has_effective_tools(tools):
-        # ``tools`` may be a JSON string (poseidon / RESPONSE_FORMAT_TOOL_DEFINITIONS)
-        # or a list (Chat Completions ``request.tools`` as Pydantic models or dicts).
-        tool_schema_list = _tool_definitions_to_schema_list(tools)
-        if not tool_schema_list:
-            raise ValueError(
-                "No valid tool definitions could be parsed from the request for "
-                "structural tag conversion."
-            )
-        tool_grammar = collect_tool_schema(tool_schema_list)
-        _add_tag(style.tools, {"type": "grammar", "grammar": tool_grammar})
+    from vllm.tokenizers import TokenizerLike
 
-    if not tags:
-        return None
-    return json.dumps(
-        {
-            "type": "structural_tag",
-            "format": {
-                "type": "triggered_tags",
-                "triggers": triggers,
-                "tags": tags,
-            },
-        }
-    )
 
-
-def _response_format_type(
-    response_format: AnyResponseFormat | dict | None,
-) -> str | None:
-    if response_format is None:
-        return None
-    if isinstance(response_format, dict):
-        t = response_format.get("type")
-        return t if isinstance(t, str) else None
-    return response_format.type
-
-
-def _maybe_parse_json_dict(value: Any) -> dict | None:
-    """If value is a JSON string, parse to dict; otherwise require dict."""
-    if isinstance(value, dict):
-        return value
-    if isinstance(value, str):
-        try:
-            parsed = json.loads(value)
-        except (TypeError, json.JSONDecodeError):
-            return None
-        return parsed if isinstance(parsed, dict) else None
-    return None
-
-
-def _unwrap_nested_schema(candidate: Any) -> dict | None:
-    """Return ``candidate`` as a dict, unwrapping a nested ``schema`` if present.
-
-    Returns ``None`` if ``candidate`` is not (and cannot be parsed into) a dict.
-    """
-    cand = _maybe_parse_json_dict(candidate)
-    if not isinstance(cand, dict):
-        return None
-    nested = cand.get("schema")
-    return nested if isinstance(nested, dict) else cand
-
-
-def _schema_from_json_schema_field(js_wr: Any) -> dict | None:
-    """
-    Extract the JSON Schema object from Chat Completions ``json_schema`` payload.
-
-    Accepts:
-    - ``JsonSchemaResponseFormat`` (Pydantic) with ``schema`` / ``json_schema`` field
-    - dict in OpenAI shape ``{"name": ..., "schema": {...}}``
-    - dict with ``json_schema`` key holding either the schema or a nested wrapper
-    - dict that is already a JSON Schema document (some clients omit the wrapper)
-    - JSON strings for any of the above
-    """
-    if js_wr is None:
-        return None
-
-    parsed_wr = _maybe_parse_json_dict(js_wr)
-    if parsed_wr is not None:
-        js_wr = parsed_wr
-
-    if hasattr(js_wr, "model_dump"):
-        for by_alias in (True, False):
-            try:
-                data = js_wr.model_dump(by_alias=by_alias, exclude_none=False)
-            except TypeError:
-                data = js_wr.model_dump(by_alias=by_alias)
-            out = _unwrap_nested_schema(data.get("schema") or data.get("json_schema"))
-            if out is not None:
-                return out
-        inner_attr = getattr(js_wr, "json_schema", None)
-        return inner_attr if isinstance(inner_attr, dict) else None
-
-    if isinstance(js_wr, dict):
-        for key in ("schema", "json_schema"):
-            out = _unwrap_nested_schema(js_wr.get(key))
-            if out is not None:
-                return out
-        return js_wr
-
-    return None
-
-
-def _schema_dict_from_chat_response_format(
-    rf: AnyResponseFormat | dict | None,
-) -> dict | None:
-    """JSON schema dict from Chat Completions ``request.response_format`` only."""
-    if rf is None:
-        return None
-    rf_type = _response_format_type(rf)
-    if rf_type == "json_object":
-        return {"type": "object"}
-    if rf_type != "json_schema":
-        return None
-    js_wr = (
-        rf.get("json_schema")
-        if isinstance(rf, dict)
-        else getattr(rf, "json_schema", None)
-    )
-    return _schema_from_json_schema_field(js_wr)
-
-
-def _schema_dict_from_structured_outputs(
-    so: StructuredOutputsParams | None,
-) -> dict | None:
-    """Schema dict from ``structured_outputs`` (``json`` / ``json_object``).
-
-    Same unwrapping as ``json_schema``. ``json`` is expected to be ``str`` or
-    ``dict`` (enforced by ``StructuredOutputsParams`` / request models); other
-    types raise ``ValueError`` only if a caller bypasses that validation.
-    """
-    if so is None:
-        return None
-    if so.json_object:
-        return {"type": "object"}
-    raw: Any = so.json
-    if raw is None:
-        return None
-
-    if hasattr(raw, "model_dump"):
-        out = _schema_from_json_schema_field(raw)
-        if out is None:
-            raise ValueError(
-                "structured_outputs.json model has no extractable JSON Schema."
-            )
-        return out
-
-    if isinstance(raw, str):
-        if not raw.strip():
-            raise ValueError("structured_outputs.json cannot be empty.")
-        try:
-            raw = json.loads(raw)
-        except json.JSONDecodeError as e:
-            raise ValueError("structured_outputs.json must be valid JSON.") from e
-        if not isinstance(raw, dict):
-            raise ValueError("structured_outputs.json must decode to a JSON object.")
-
-    if isinstance(raw, Mapping):
-        body = raw if isinstance(raw, dict) else dict(raw)
-        return _schema_from_json_schema_field(body) or body
-
-    raise ValueError(
-        f"structured_outputs.json has unsupported type {type(raw).__name__}."
-    )
-
-
-def _melody_sources_to_vllm(
-    raw_sources: Any,
-    position_to_source: Mapping[tuple[int, int], CitationSource] | None,
-) -> list[CitationSource]:
-    """Convert melody's ``Source`` objects into resolved :class:`CitationSource`.
-
-    Melody's ``Source`` shape is
-    ``{tool_call_index, tool_result_indices, document_ids}``. Without
-    :meth:`PyFilterOptions.with_message_history` the ``document_ids``
-    list is empty, so we resolve the numeric address ourselves using
-    ``position_to_source``, which is
-    :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source`
-    applied to the inbound request and forwarded through
-    ``chat_template_kwargs`` at parser-construction time. Each melody
-    source with ``tool_result_indices=[i, j, ...]`` fans out to one
-    :class:`CitationSource` per resolved position, populating the
-    wire-shape ``type`` / ``id`` / ``document`` / ``tool_output``
-    fields; unresolved positions are silently skipped (the calling
-    citation is dropped at wire-coercion time if every source in it
-    ends up unresolved -- see ``_to_wire_citation`` in
-    :mod:`vllm.entrypoints.cohere.serving`).
-
-    When ``position_to_source`` is ``None`` (parser wired outside of
-    :class:`CohereServingChatV2`, e.g. a plain OpenAI chat completion
-    request that happens to hit a cohere reasoning parser) sources
-    can't be resolved and get dropped; citations then never reach any
-    citation-aware handler downstream.
-    """
-    if not position_to_source:
-        return []
-    out: list[CitationSource] = []
-    for s in raw_sources or []:
-        bucket = getattr(s, "tool_call_index", None)
-        indices = list(getattr(s, "tool_result_indices", None) or [])
-        if bucket is None or not indices:
-            continue
-        for idx in indices:
-            info = position_to_source.get((bucket, idx))
-            if info is not None:
-                out.append(info)
-    return out
-
-
-def _melody_citations_to_vllm(
-    raw_citations: Any,
-    position_to_source: Mapping[tuple[int, int], CitationSource] | None,
-) -> list[Citation] | None:
-    """Convert melody's ``FilterCitation`` objects into :class:`Citation`.
-
-    Resolves each source in-place via ``position_to_source`` (see
-    :func:`_melody_sources_to_vllm`). Citations whose sources all fail
-    to resolve are still emitted here with ``sources=[]``; the serving
-    layer drops them at wire-coercion time so the fail-closed policy
-    lives in exactly one place.
+class BaseCohereCommandReasoningParser(ReasoningParser):
+    """Expose is_reasoning_end for the structured-output engine.
+    Everything else runs in
+    :class:`vllm.parser.cohere_command.CohereCommandParser`.
     """
-    if not raw_citations:
-        return None
-    out: list[Citation] = []
-    for c in raw_citations:
-        out.append(
-            Citation(
-                start=getattr(c, "start_index", None),
-                end=getattr(c, "end_index", None),
-                text=getattr(c, "text", None),
-                sources=_melody_sources_to_vllm(
-                    getattr(c, "sources", None), position_to_source
-                ),
-                type=(
-                    "THINKING_CONTENT"
-                    if getattr(c, "is_thinking", False)
-                    else "TEXT_CONTENT"
-                ),
-            )
-        )
-    return out
 
+    melody_preset: ClassVar[str]
 
-class BaseCohereCommandReasoningParser(ReasoningParser):
-    def __init__(
-        self,
-        tokenizer: TokenizerLike,
-        *args,
-        streaming_opts: PyFilterOptions,
-        unary_opts: PyFilterOptions,
-        **kwargs,
-    ):
+    def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
         super().__init__(tokenizer, *args, **kwargs)
         self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>")
         self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>")
         self.chatbot_token_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>")
-        self.unary_opts = unary_opts
-        self.melody_unary = PyFilter(unary_opts)
-        self.melody_streaming = PyFilter(streaming_opts)
-        # Citations extracted by the most recent ``extract_reasoning`` call.
-        # Citation-aware handlers (e.g. :class:`CohereServingChatV2` via its
-        # ``_finalize_response_message`` hook) read this back from the
-        # parser instance (which is constructed per-request) and attach
-        # the result to :class:`CohereChatMessage.citations`. ``None`` when
-        # the last parse produced no citations.
-        self.last_unary_citations: list[Citation] | None = None
-        # Request-scoped ``(tool_call_index, tool_result_idx) -> resolved
-        # CitationSource`` map, forwarded from
-        # :meth:`CohereServingChatV2._apply_cohere_template_kwargs` via
-        # ``chat_template_kwargs``. See :func:`_melody_sources_to_vllm`
-        # for how it's consumed. Absent for parser instances that
-        # weren't created by ``CohereServingChatV2`` (e.g. an OpenAI
-        # chat completion request routed through a cohere reasoning
-        # parser); citations from those requests are dropped since we
-        # have no way to attribute their sources.
-        ctk = kwargs.get("chat_template_kwargs") or {}
-        # Lazy import: ``renderers.cohere`` transitively imports
-        # ``cohere_melody`` and other heavy deps; hoisting the constant
-        # to module scope would tie parser importability to that graph.
-        from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
-
-        raw_map = ctk.get(POSITION_TO_SOURCE_KEY)
-        self._position_to_source: Mapping[tuple[int, int], CitationSource] | None = (
-            raw_map if isinstance(raw_map, Mapping) else None
-        )
 
     @property
     def reasoning_start_str(self) -> str | None:
@@ -537,80 +34,6 @@ def reasoning_start_str(self) -> str | None:
     def reasoning_end_str(self) -> str | None:
         return "<|END_THINKING|>"
 
-    def extract_reasoning_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-    ) -> DeltaMessage | None:
-        r = self.melody_streaming.write_decoded(delta_text)
-        citations = _melody_citations_to_vllm(
-            getattr(r, "citations", None), self._position_to_source
-        )
-        if (
-            r.content is None
-            and r.reasoning is None
-            and not r.tool_calls
-            and not citations
-        ):
-            return None
-        # Always emit CohereDeltaMessage so citations reach the wire via
-        # ``SerializeAsAny[DeltaMessage]`` on the streaming choice envelope.
-        # When ``citations`` is unset, ``CohereDeltaMessage._serialize``
-        # drops the field so the emitted shape matches plain
-        # :class:`DeltaMessage`
-        msg = CohereDeltaMessage()
-        if r.content is not None:
-            msg.content = r.content
-        if r.reasoning is not None:
-            msg.reasoning = r.reasoning
-        if r.tool_calls:
-            msg.tool_calls = [
-                DeltaToolCall(
-                    id=tc.id,
-                    index=tc.index,
-                    type="function",
-                    function=DeltaFunctionCall(name=tc.name, arguments=tc.arguments),
-                )
-                for tc in r.tool_calls
-            ]
-        if citations:
-            msg.citations = citations
-        return msg
-
-    def extract_reasoning(
-        self, model_output: str, request: ChatCompletionRequest | ResponsesRequest
-    ) -> tuple[str | None, str | None]:
-        result = self.melody_unary.process_full_text(model_output)
-        # Cache citations so citation-aware handlers can surface them
-        # on :class:`CohereChatMessage.citations`. The base
-        # :meth:`ReasoningParser.extract_reasoning` contract only
-        # returns ``(reasoning, content)``, so citations are passed
-        # back via parser-instance state -- safe because the parser is
-        # constructed per-request.
-        self.last_unary_citations = _melody_citations_to_vllm(
-            getattr(result, "citations", None), self._position_to_source
-        )
-        return result.reasoning, result.content
-
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        token_buf: list[int] = []
-        content_ids: list[int] = []
-        content_filter = PyFilter(self.unary_opts)
-        for t in input_ids:
-            token_buf.append(t)
-            s = self.model_tokenizer.decode(token_buf, skip_special_tokens=False)
-            if s.endswith(REPLACEMENT_CHAR):
-                continue
-            r = content_filter.write_decoded(s)
-            if r.content is not None:
-                content_ids.extend(token_buf)
-            token_buf = []
-        return content_ids
-
     def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
         chatbot = self.chatbot_token_id
         start = self.start_token_id
@@ -628,91 +51,10 @@ def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
 
         return has_end_token
 
-    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
-        count = 0
-        depth = 0
-        for token_id in token_ids:
-            if token_id == self.start_token_id:
-                depth += 1
-                continue
-            if token_id == self.end_token_id:
-                if depth > 0:
-                    depth -= 1
-                continue
-            if depth > 0:
-                count += 1
-        return count
-
-    def adjust_request(
-        self, request: ChatCompletionRequest | ResponsesRequest
-    ) -> ChatCompletionRequest | ResponsesRequest:
-        so = request.structured_outputs
-        if so is not None and so.structural_tag:
-            return request
-        # Schema: prefer ``response_format`` (OpenAI Chat Completions), then
-        # ``structured_outputs.json`` / ``json_object`` (vLLM direct). Tools stay
-        # on ``request.tools``.
-        rf = (
-            request.response_format
-            if isinstance(request, ChatCompletionRequest)
-            else None
-        )
-        if rf is not None and _response_format_type(rf) == "structural_tag":
-            return request
-        model_architecture = (
-            self._model_config.architecture if self._model_config is not None else None
-        )
-        tools = request.tools
-        # ``response_format`` wins if both it and ``structured_outputs`` supply JSON.
-        schema = _schema_dict_from_chat_response_format(rf)
-        if schema is None:
-            schema = _schema_dict_from_structured_outputs(so)
-        if schema is None and not _has_effective_tools(tools):
-            return request
-        if model_architecture is None:
-            return request
-        result = convert_schema_to_structural_tags(
-            schema=schema,
-            tools=tools,
-            model_architecture=model_architecture,
-        )
-        if result is None:
-            # Unsupported architectures are not in ``MODEL_TO_TAG_STYLE``.
-            raise ValueError(
-                "Failed to build structural_tag guided decoding constraints from "
-                "this request's JSON schema and/or tools. The configured model "
-                f"architecture ({model_architecture!r}) does not support Cohere "
-                "command structural tags, or the schema cannot be expressed in "
-                "that format."
-            )
-        request.structured_outputs = StructuredOutputsParams(structural_tag=result)
-        # Folded JSON constraints into ``structural_tag``; drop ``response_format``
-        # when it was the source so ``to_sampling_params`` does not also set ``json`` /
-        # ``json_object`` (mutually exclusive in ``StructuredOutputsParams``).
-        if isinstance(request, ChatCompletionRequest) and rf is not None:
-            rf_type = _response_format_type(rf)
-            if rf_type in ("json_schema", "json_object"):
-                request.response_format = None
-        return request
-
 
 class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser):
-    def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
-        super().__init__(
-            tokenizer,
-            *args,
-            streaming_opts=PyFilterOptions().cmd3().no_tools(),
-            unary_opts=PyFilterOptions().cmd3().no_tools(),
-            **kwargs,
-        )
+    melody_preset = "cmd3"
 
 
 class CohereCommand4ReasoningParser(BaseCohereCommandReasoningParser):
-    def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
-        super().__init__(
-            tokenizer,
-            *args,
-            streaming_opts=PyFilterOptions().cmd4().no_tools(),
-            unary_opts=PyFilterOptions().cmd4().no_tools(),
-            **kwargs,
-        )
+    melody_preset = "cmd4"
diff --git a/vllm/reasoning/deepseek_r1_reasoning_parser.py b/vllm/reasoning/deepseek_r1_reasoning_parser.py
index f16f87f6b0a8..991cf016599d 100644
--- a/vllm/reasoning/deepseek_r1_reasoning_parser.py
+++ b/vllm/reasoning/deepseek_r1_reasoning_parser.py
@@ -8,8 +8,7 @@
 
 
 class DeepSeekR1ReasoningParser(BaseThinkingReasoningParser):
-    """
-    Reasoning parser for DeepSeek R1 model.
+    """Reasoning parser for DeepSeek R1 model.
 
     The DeepSeek R1 model uses <think>...</think> tokens to denote reasoning
     text. This parser extracts the reasoning content from the model output.
diff --git a/vllm/reasoning/deepseek_v3_reasoning_parser.py b/vllm/reasoning/deepseek_v3_reasoning_parser.py
index 302024a8e77d..072048af91a4 100644
--- a/vllm/reasoning/deepseek_v3_reasoning_parser.py
+++ b/vllm/reasoning/deepseek_v3_reasoning_parser.py
@@ -18,8 +18,7 @@
 
 
 class DeepSeekV3ReasoningParser(ReasoningParser):
-    """
-    V3 parser that delegates to either DeepSeekR1ReasoningParser or
+    """V3 parser that delegates to either DeepSeekR1ReasoningParser or
     IdentityReasoningParser based on `thinking` and `separate_reasoning`.
     """
 
@@ -81,9 +80,7 @@ def extract_reasoning_streaming(
 
 
 class DeepSeekV3ReasoningWithThinkingParser(DeepSeekV3ReasoningParser):
-    """
-    DeepSeekV3ReasoningParser that defaults to thinking mode.
-    """
+    """DeepSeekV3ReasoningParser that defaults to thinking mode."""
 
     def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
         chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
diff --git a/vllm/reasoning/ernie45_reasoning_parser.py b/vllm/reasoning/ernie45_reasoning_parser.py
index 0aa926281988..444467b68e66 100644
--- a/vllm/reasoning/ernie45_reasoning_parser.py
+++ b/vllm/reasoning/ernie45_reasoning_parser.py
@@ -15,8 +15,7 @@
 
 
 class Ernie45ReasoningParser(BaseThinkingReasoningParser):
-    """
-    Reasoning parser for Ernie45 thinking model.
+    """Reasoning parser for Ernie45 thinking model.
     The Ernie45 thinking model output format is
         abc\n</think>\n\n<response>\ndef\n</response>\n
     or  abc\n</think>\ndef
@@ -60,8 +59,7 @@ def extract_reasoning_streaming(
         current_token_ids: Sequence[int],
         delta_token_ids: Sequence[int],
     ) -> DeltaMessage | None:
-        """
-        Extract reasoning content from a delta message.
+        """Extract reasoning content from a delta message.
         Handles streaming output where previous + delta = current.
         Uses token IDs for faster processing.
         The Ernie45 thinking model output format is
@@ -138,8 +136,7 @@ def extract_reasoning_streaming(
     def extract_reasoning(
         self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
     ) -> tuple[str | None, str | None]:
-        """
-        Extract reasoning content from the model output.
+        """Extract reasoning content from the model output.
         The Ernie45 thinking model output format is
             abc\n</think>\n\n\n<response>\ndef\n</response>\n
         or  abc\n</think>\ndef
diff --git a/vllm/reasoning/gemma4_utils.py b/vllm/reasoning/gemma4_utils.py
index 9cdac72039e8..9566b8af56b8 100644
--- a/vllm/reasoning/gemma4_utils.py
+++ b/vllm/reasoning/gemma4_utils.py
@@ -75,6 +75,7 @@ def parse_thinking_output(text: str) -> dict[str, str | None]:
         >>> result = parse_thinking_output(output_text)
         >>> print(result["thinking"])  # chain-of-thought reasoning or None
         >>> print(result["answer"])    # final answer
+
     """
     if _THINKING_END_TAG in text:
         parts = text.split(_THINKING_END_TAG, 1)
diff --git a/vllm/reasoning/gptoss_reasoning_parser.py b/vllm/reasoning/gptoss_reasoning_parser.py
index a6f108015812..5505f1de446b 100644
--- a/vllm/reasoning/gptoss_reasoning_parser.py
+++ b/vllm/reasoning/gptoss_reasoning_parser.py
@@ -14,8 +14,7 @@
 
 
 class GptOssReasoningParser(ReasoningParser):
-    """
-    Reasoning parser for GptOss model.
+    """Reasoning parser for GptOss model.
 
     The GptOss model uses harmony to extract reasoning content and this parser
     is only used for detecting the end of the reasoning content.
diff --git a/vllm/reasoning/granite_reasoning_parser.py b/vllm/reasoning/granite_reasoning_parser.py
index 463659f8d03d..dc6118224f7e 100644
--- a/vllm/reasoning/granite_reasoning_parser.py
+++ b/vllm/reasoning/granite_reasoning_parser.py
@@ -16,8 +16,7 @@
 
 
 class GraniteReasoningParser(ReasoningParser):
-    """
-    Reasoning parser for IBM Granite.
+    """Reasoning parser for IBM Granite.
 
     IBM granite models currently use "Here is my thought process:"
     and "Here is my response:" to separate its thinking / response outputs.
@@ -65,6 +64,7 @@ def extract_reasoning(
         Returns:
             tuple[Optional[str], Optional[str]]: Tuple pair containing the
             reasoning content and non-reasoning content.
+
         """
         re_match = self.reasoning_regex.findall(model_output)
         if not re_match:
@@ -108,6 +108,7 @@ def extract_reasoning_streaming(
         Returns:
             Union[DeltaMessage, None]
                 DeltaMessage with either reasoning content or content, or None.
+
         """
         reasoning, resp_seq_len, content = self._get_content_sections(current_text)
         # Either we haven't finished the start of the reasoning sequence,
@@ -142,6 +143,7 @@ def _is_reasoning_start_substr(self, text: str) -> bool:
 
         Returns:
             bool: True if any of the possible reasoning start seqs match.
+
         """
         return any(
             think_start.startswith(text) for think_start in self.valid_think_starts
@@ -155,6 +157,7 @@ def _is_response_start_substr(self, text: str) -> bool:
 
         Returns:
             bool: True if any of the possible response start seqs match.
+
         """
         return any(
             response_start.startswith(text)
@@ -175,6 +178,7 @@ def _get_delta_message_with_no_reasoning_bounds(
 
         Returns:
             DeltaMessage: Message containing the parsed content.
+
         """
         prev_longest_length = len(current_text) - len(delta_text)
         is_substr = self._is_reasoning_start_substr(current_text)
@@ -213,6 +217,7 @@ def _get_delta_message_with_no_response_bounds(
 
         Returns:
             DeltaMessage: Message containing the parsed content.
+
         """
         # If we have no reasoning content or explicitly end with the start of
         # response sequence, we are in transition to the response; need to be
@@ -293,6 +298,7 @@ def _get_delta_message_with_both_bounds(
 
         Returns:
             DeltaMessage: Message containing the parsed content.
+
         """
         # Always have content; take length to the end
         delta_content = delta_text[-len(response_content) :]
@@ -329,6 +335,7 @@ def _get_content_sections(
             tuple[Optional[str], Optional[int], Optional[str]]: Tuple of len 3
             containing the reasoning content, the length of the response seq
             (if there is one) and the non-reasoning content.
+
         """
         current_chunk_start = 0
         start_reasoning = None
diff --git a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py
index b15b5608fbe9..0a7e021a464c 100644
--- a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py
+++ b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py
@@ -16,8 +16,7 @@
 
 
 class HunyuanA13BReasoningParser(ReasoningParser):
-    """
-    Reasoning parser for Hunyuan A13B Model
+    """Reasoning parser for Hunyuan A13B Model.
 
     HunyuanReasoningParser
 
@@ -102,8 +101,8 @@ def extract_reasoning(
         Returns:
             tuple[Optional[str], Optional[str]]: Tuple pair containing the
             reasoning content and non-reasoning content.
-        """
 
+        """
         re_match = self.full_match_reasoning_regex.findall(model_output)
         if re_match:
             reasoning, response_content = re_match[0]
@@ -139,7 +138,7 @@ def extract_reasoning_streaming(
         current_token_ids: Sequence[int],
         delta_token_ids: Sequence[int],
     ) -> DeltaMessage | None:
-        """Extract content using token ID sequence state machine"""
+        """Extract content using token ID sequence state machine."""
         # Define sequences
         think_start_sequence = self.think_start_ids
         response_start_sequence = self.response_start_ids
diff --git a/vllm/reasoning/hy_v3_reasoning_parser.py b/vllm/reasoning/hy_v3_reasoning_parser.py
index 0765be74af73..77a1a271d4a8 100644
--- a/vllm/reasoning/hy_v3_reasoning_parser.py
+++ b/vllm/reasoning/hy_v3_reasoning_parser.py
@@ -18,8 +18,7 @@
 
 
 class HYV3ReasoningParser(BaseThinkingReasoningParser):
-    """
-    HYV3 parser that delegates to either HYV3ReasoningParser or
+    """HYV3 parser that delegates to either HYV3ReasoningParser or
     IdentityReasoningParser based on `reasoning_effort`.
 
     The HYV3 model uses <think>...</think> tokens to denote reasoning text.
diff --git a/vllm/reasoning/hy_v4_reasoning_parser.py b/vllm/reasoning/hy_v4_reasoning_parser.py
index b65db92b92f5..65035ed53ad4 100644
--- a/vllm/reasoning/hy_v4_reasoning_parser.py
+++ b/vllm/reasoning/hy_v4_reasoning_parser.py
@@ -46,8 +46,8 @@ def detect_token_suffix(tokenizer: TokenizerLike) -> str:
         RuntimeError: The tokenizer declares the structural tokens through
             ``model_specific_special_tokens``, which transformers 5 no longer
             round-trips.
-    """
 
+    """
     import transformers
 
     if int(transformers.__version__.split(".")[0]) >= 5:
diff --git a/vllm/reasoning/identity_reasoning_parser.py b/vllm/reasoning/identity_reasoning_parser.py
index 52fab9914105..f2c9a1e21e33 100644
--- a/vllm/reasoning/identity_reasoning_parser.py
+++ b/vllm/reasoning/identity_reasoning_parser.py
@@ -15,8 +15,7 @@
 
 
 class IdentityReasoningParser(ReasoningParser):
-    """
-    Identity reasoning parser.
+    """Identity reasoning parser.
 
     This parser does not attempt to parse or strip out reasoning tokens.
     It treats the entire model output as content and ignores reasoning.
diff --git a/vllm/reasoning/kimi_k3_reasoning_parser.py b/vllm/reasoning/kimi_k3_reasoning_parser.py
index d3ae2b3abe38..e5618d47be4a 100644
--- a/vllm/reasoning/kimi_k3_reasoning_parser.py
+++ b/vllm/reasoning/kimi_k3_reasoning_parser.py
@@ -256,11 +256,13 @@ def extract_reasoning(
     ) -> tuple[str | None, str | None]:
         """Split full text into ``(reasoning, rest)`` for the non-streaming path.
 
-        Handles three shapes:
-          * no think channel at all   -> ``(None, model_output)`` (all content)
+        Handles four shapes:
+          * response opener without think markers -> no think channel; all content
           * open marker present       -> reasoning starts after ``<|open|>think<|sep|>``
           * open marker absent but a   close marker exists (gen-prefix consumed
             the open) -> reasoning starts at offset 0
+          * neither marker present     -> truncated reasoning after a consumed
+            generation prefix
         ``rest`` is whatever follows the close marker, fed on to the tool parser.
         """
         if not self._thinking_enabled:
@@ -270,11 +272,14 @@ def extract_reasoning(
         # reasoning content begins right after think-open (or at start if the
         # open marker was already consumed as a generation prefix)
         content_start = m_open.end() if m_open is not None else 0
-        # if there is no think channel at all, everything is content
-        if m_open is None and self._think_close_re.search(model_output) is None:
-            return None, self._content_after_reasoning(model_output, request)
 
         m_close = self._think_close_re.search(model_output, content_start)
+        if (
+            m_open is None
+            and m_close is None
+            and self._response_open_re.search(model_output) is not None
+        ):
+            return None, self._content_after_reasoning(model_output, request)
         if m_close is not None:
             reasoning = model_output[content_start : m_close.start()]
             rest = model_output[m_close.end() :]
diff --git a/vllm/reasoning/minimax_m2_reasoning_parser.py b/vllm/reasoning/minimax_m2_reasoning_parser.py
index b26ea8889ae0..8fa5b45d5853 100644
--- a/vllm/reasoning/minimax_m2_reasoning_parser.py
+++ b/vllm/reasoning/minimax_m2_reasoning_parser.py
@@ -15,8 +15,7 @@
 
 
 class MiniMaxM2ReasoningParser(MinimaxM2ParserReasoningAdapter):  # type: ignore[valid-type, misc]
-    """
-    Reasoning parser for MiniMax M2 model.
+    """Reasoning parser for MiniMax M2 model.
 
     MiniMax M2 models don't generate <think> start token, only </think> end
     token. All content before </think> is reasoning, content after is the
@@ -25,9 +24,7 @@ class MiniMaxM2ReasoningParser(MinimaxM2ParserReasoningAdapter):  # type: ignore
 
 
 class MiniMaxM2AppendThinkReasoningParser(ReasoningParser):
-    """
-    Reasoning parser for MiniMax M2 model.
-    """
+    """Reasoning parser for MiniMax M2 model."""
 
     def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
         super().__init__(tokenizer, *args, **kwargs)
diff --git a/vllm/reasoning/olmo3_reasoning_parser.py b/vllm/reasoning/olmo3_reasoning_parser.py
index b658b92bcbe1..8ac75f91ff76 100644
--- a/vllm/reasoning/olmo3_reasoning_parser.py
+++ b/vllm/reasoning/olmo3_reasoning_parser.py
@@ -32,8 +32,7 @@ def __len__(self):
 
 
 def string_overlap(a: str, b: str) -> tuple[Indices | None, Indices | None]:
-    """
-    Find the longest overlap where the end of string a matches the start
+    """Find the longest overlap where the end of string a matches the start
     of string b.
 
     Args:
@@ -43,8 +42,8 @@ def string_overlap(a: str, b: str) -> tuple[Indices | None, Indices | None]:
     Returns:
         Tuple of IndicesTuples representing the overlapping portions in each
         string, or a tuple of None if no overlap exists
-    """
 
+    """
     # swap so a is always the shorter string
     a, b, swap = (a, b, False) if len(a) < len(b) else (b, a, True)
 
@@ -188,8 +187,7 @@ def add_text(self, delta_text: str) -> DeltaMessage | None:
 
 
 class Olmo3ReasoningParser(ReasoningParser):
-    """
-    Reasoning parser for Olmo 3 model
+    """Reasoning parser for Olmo 3 model.
 
     Olmo3ReasoningParser
 
@@ -285,8 +283,8 @@ def extract_reasoning(
         Returns:
             tuple[Optional[str], Optional[str]]: Tuple pair containing the
             reasoning content and non-reasoning content.
-        """
 
+        """
         re_match = self.reasoning_regex.match(model_output)
         if re_match:
             reasoning = re_match.group("reasoning") or None
@@ -305,8 +303,7 @@ def extract_reasoning_streaming(
         current_token_ids: Sequence[int],
         delta_token_ids: Sequence[int],
     ) -> DeltaMessage | None:
-        """Extract content using token ID sequence state machine"""
-
+        """Extract content using token ID sequence state machine."""
         delta_message = self.buffer.add_text(delta_text)
         if delta_message is None and self.buffer.think_end in self.buffer.buffer:
             # this is a bit hacky, but, because of how the buffer is
diff --git a/vllm/reasoning/poolside_v1_reasoning_parser.py b/vllm/reasoning/poolside_v1_reasoning_parser.py
index 30031d8513a9..e854afe4bc46 100644
--- a/vllm/reasoning/poolside_v1_reasoning_parser.py
+++ b/vllm/reasoning/poolside_v1_reasoning_parser.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Laguna reasoning parser.
+"""Laguna reasoning parser.
 
 ``DeepSeekV3ReasoningParser.is_reasoning_end`` walks the entire
 token sequence backwards and returns ``True`` on the first ``</think>`` it
diff --git a/vllm/reasoning/step3_reasoning_parser.py b/vllm/reasoning/step3_reasoning_parser.py
index d2a575c74ee7..16ab1cf9f400 100644
--- a/vllm/reasoning/step3_reasoning_parser.py
+++ b/vllm/reasoning/step3_reasoning_parser.py
@@ -17,8 +17,7 @@
 
 
 class Step3ReasoningParser(ReasoningParser):
-    """
-    Reasoning parser for Step3 model.
+    """Reasoning parser for Step3 model.
 
     The Step3 model uses </think> token to denote the end of reasoning
     text. This parser extracts all content before </think> as reasoning content.
@@ -62,8 +61,7 @@ def extract_reasoning_streaming(
         current_token_ids: Sequence[int],
         delta_token_ids: Sequence[int],
     ) -> DeltaMessage | None:
-        """
-        Extract reasoning content from a delta message.
+        """Extract reasoning content from a delta message.
         Handles streaming output where previous + delta = current.
         Uses token IDs for faster processing.
         For text "abc</think>xyz":
diff --git a/vllm/reasoning/step3p5_reasoning_parser.py b/vllm/reasoning/step3p5_reasoning_parser.py
index cc7c227e6d7e..5969ca5228cc 100644
--- a/vllm/reasoning/step3p5_reasoning_parser.py
+++ b/vllm/reasoning/step3p5_reasoning_parser.py
@@ -14,8 +14,7 @@
 
 
 class Step3p5ReasoningParser(BaseThinkingReasoningParser):
-    """
-    Reasoning parser for Step3p5 model.
+    """Reasoning parser for Step3p5 model.
 
     Step3p5 uses the <think>...</think> format, but it tends to emit an extra
     newline immediately before and/or after the </think> token. This parser
diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py
index 0fd1b7610bae..87d11b7b92d4 100644
--- a/vllm/renderers/base.py
+++ b/vllm/renderers/base.py
@@ -7,11 +7,13 @@
 from collections.abc import Mapping, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import ExitStack
+from dataclasses import replace
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Generic, overload
 
 from typing_extensions import TypeVar
 
+from vllm.exceptions import VLLMValidationError
 from vllm.inputs import (
     EmbedsInput,
     EmbedsPrompt,
@@ -30,7 +32,11 @@
 )
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY as mm_registry
-from vllm.multimodal.cache import BaseMultiModalProcessorCache
+from vllm.multimodal.cache import (
+    BaseMultiModalProcessorCache,
+    processor_cache_from_config,
+    processor_only_cache_from_config,
+)
 from vllm.multimodal.gpu_ipc_memory import maybe_init_mm_gpu_ipc_pool
 from vllm.multimodal.parse import (
     MultiModalDataItems,
@@ -127,8 +133,7 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None:
         self._async_tokenizer_decode = make_async(self._decode, executor=self._executor)
 
         self.mm_processor: BaseMultiModalProcessor | None = None
-        self._readonly_mm_processor: BaseMultiModalProcessor | None = None
-        self._mm_cache_stats: MultiModalCacheStats | None = None
+
         self._clear_mm_cache_async = make_async(
             self.clear_mm_cache, executor=self._mm_executor
         )
@@ -138,7 +143,17 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None:
         self._safe_load_prompt_embeds_async = make_async(
             safe_load_prompt_embeds, executor=self._executor
         )
-        if mm_registry.supports_multimodal_inputs(config.model_config):
+
+        self._mm_processor_cache: BaseMultiModalProcessorCache | None = None
+
+        # A second processor with its own processor-only cache.
+        # Used by the tokenize endpoint so that tokenize-only
+        # requests don't pollute the sender cache.
+        self._mm_processor_only_cache: BaseMultiModalProcessorCache | None = None
+
+        self._mm_cache_stats: MultiModalCacheStats | None = None
+
+        if config.model_config.supports_multimodal_inputs:
             # Install the process-global GPU memory pool used to gate
             # frontend GPU-side multimodal decoding (no-op when the budget
             # is 0). Lives in the API-server process only.
@@ -149,30 +164,14 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None:
                     config.parallel_config._api_process_count,
                 )
 
-            mm_processor_cache = mm_registry.processor_cache_from_config(config)
-
             with set_default_torch_num_threads():
                 self.mm_processor = mm_registry.create_processor(
                     config.model_config,
                     tokenizer=self.tokenizer,
-                    cache=mm_processor_cache,
                 )
 
-            if mm_processor_cache:
-                self._mm_cache_stats = MultiModalCacheStats()
-                self._resources.callback(mm_processor_cache.close)
-
-            # A second processor with its own processor-only cache.
-            # Used by the tokenize endpoint so that tokenize-only
-            # requests don't pollute the sender cache.
-            ro_cache = mm_registry.processor_only_cache_from_config(config)
-            if ro_cache is not None:
-                with set_default_torch_num_threads():
-                    self._readonly_mm_processor = mm_registry.create_processor(
-                        config.model_config,
-                        tokenizer=self.tokenizer,
-                        cache=ro_cache,
-                    )
+            self._mm_processor_cache = processor_cache_from_config(config)
+            self._mm_processor_only_cache = processor_only_cache_from_config(config)
 
             # This is used to generate internal request ID for MM processing
             # It has no relation to the request ID for engine core
@@ -181,6 +180,12 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None:
                 config.observability_config
             )
 
+        if self._mm_processor_cache:
+            self._mm_cache_stats = MultiModalCacheStats()
+            self._resources.callback(self._mm_processor_cache.close)
+        if self._mm_processor_only_cache:
+            self._resources.callback(self._mm_processor_only_cache.close)
+
     def get_tokenizer(self) -> _T:
         tokenizer = self.tokenizer
         if tokenizer is None:
@@ -188,6 +193,10 @@ def get_tokenizer(self) -> _T:
 
         return tokenizer
 
+    def render_completion_suffix(self, prompt: str, suffix: str) -> str | None:
+        """Render OpenAI completion suffix input when the renderer supports FIM."""
+        return None
+
     def _decode(self, *args, **kwargs):
         return self.get_tokenizer().decode(*args, **kwargs)
 
@@ -198,11 +207,8 @@ def get_mm_processor(self) -> "BaseMultiModalProcessor":
         return self.mm_processor
 
     @property
-    def mm_processor_cache(self) -> "BaseMultiModalProcessorCache | None":
-        if self.mm_processor is None:
-            return None
-
-        return self.mm_processor.cache
+    def mm_processor_cache(self) -> BaseMultiModalProcessorCache | None:
+        return self._mm_processor_cache
 
     def stat_mm_cache(self) -> MultiModalCacheStats | None:
         mm_cache_stats = self._mm_cache_stats
@@ -214,7 +220,7 @@ def stat_mm_cache(self) -> MultiModalCacheStats | None:
         return mm_cache_stats
 
     def update_mm_cache_stats(self) -> None:
-        mm_processor_cache = self.mm_processor_cache
+        mm_processor_cache = self._mm_processor_cache
         mm_cache_stats = self._mm_cache_stats
 
         if mm_processor_cache and mm_cache_stats:
@@ -222,38 +228,26 @@ def update_mm_cache_stats(self) -> None:
             mm_cache_stats.record(delta.total, delta.hits)
 
     def clear_mm_cache(self) -> None:
-        mm_processor_cache = self.mm_processor_cache
+        mm_processor_cache = self._mm_processor_cache
         if mm_processor_cache is not None:
             mm_processor_cache.clear_cache()
 
         if self._mm_cache_stats is not None:
             self._mm_cache_stats.reset = True
 
-    @staticmethod
-    def _clear_processor_cache(
-        processor: "BaseMultiModalProcessor | None",
-    ) -> None:
-        if processor is None:
-            return
-
-        processor_cache = processor.cache
-        if processor_cache is not None:
-            processor_cache.clear_cache()
-
     def _warmup_mm_processor(
         self,
-        processor: "BaseMultiModalProcessor",
+        processor: BaseMultiModalProcessor,
+        cache: BaseMultiModalProcessorCache | None,
         *,
         log_prefix: str,
     ) -> None:
-        model_config = self.model_config
         mm_limits = {k: v for k, v in processor.info.allowed_mm_limits.items() if v > 0}
 
         start_time = time.perf_counter()
-        _ = mm_registry.get_dummy_mm_inputs(
-            model_config,
-            mm_counts=dict.fromkeys(mm_limits, 1),
-            processor=processor,
+        _ = processor.get_dummy_mm_inputs(
+            dict.fromkeys(mm_limits, 1),
+            cache=cache,
             scheduler_config=self.config.scheduler_config,
         )
 
@@ -274,28 +268,20 @@ def warmup_mm(self) -> None:
             # prevent MM processor hangs
             with set_default_torch_num_threads(1):
                 if self.mm_processor:
+                    mm_processor_only_cache = self._mm_processor_only_cache
+
                     try:
                         logger.debug("Warming up multi-modal processing...")
                         self._warmup_mm_processor(
                             self.mm_processor,
+                            mm_processor_only_cache,
                             log_prefix="Multi-modal",
                         )
                     except Exception:
                         logger.warning("Multi-modal warmup failed")
                     finally:
-                        self.clear_mm_cache()
-
-                if self._readonly_mm_processor is not None:
-                    try:
-                        logger.debug("Warming up readonly multi-modal processing...")
-                        self._warmup_mm_processor(
-                            self._readonly_mm_processor,
-                            log_prefix="Readonly multi-modal",
-                        )
-                    except Exception:
-                        logger.warning("Readonly multi-modal warmup failed")
-                    finally:
-                        self._clear_processor_cache(self._readonly_mm_processor)
+                        if mm_processor_only_cache:
+                            mm_processor_only_cache.clear_cache()
         finally:
             self._mm_warmup_done = True
 
@@ -323,7 +309,7 @@ def start_mm_warmup_in_background(self) -> None:
         # which is not designed for concurrent access (its get_and_update must
         # stay ordered with the engine-core cache) — is never touched
         # concurrently by the warmup and the serving path.
-        if self.mm_processor is None and self._readonly_mm_processor is None:
+        if self.mm_processor is None:
             return  # text-only model: nothing to warm up
         if self._mm_warmup_future is not None:
             return  # already launched
@@ -340,8 +326,7 @@ def _join_mm_warmup(self) -> None:
             self._mm_warmup_future = None
 
     def warmup(self, chat_params: ChatParams) -> None:
-        """
-        Warm up this renderer to avoid first-request latency.
+        """Warm up this renderer to avoid first-request latency.
 
         For chat requests:
         - Jinja2 template compilation
@@ -408,8 +393,7 @@ def get_eos_token_id(self) -> int | None:
         return self.tokenizer.eos_token_id
 
     def get_dec_start_token_id(self) -> int:
-        """
-        Obtain the decoder start token id employed by an encoder/decoder model,
+        """Obtain the decoder start token id employed by an encoder/decoder model,
         raising an error if it is not available.
         """
         dec_start_token_id = getattr(
@@ -824,20 +808,22 @@ def _process_mm_uuids(
     ) -> MultiModalUUIDItems:
         model_config = self.model_config
 
-        # NOTE: When users explicitly turn off BOTH prefix caching and input
-        # processing caching, no multimodal features or embeddings will be
-        # reused across requests, therefore identifying multimodal data items
-        # by their content is no longer necessary, and we create uuids with
-        # `<mm_req_id>-<modality>-<index>`, overriding even user-provided ones.
+        # Skip content hashing when both caches are disabled, but preserve
+        # explicit UUIDs used to identify embeddings across EC instances.
         if (
             model_config.multimodal_config
             and model_config.multimodal_config.mm_processor_cache_gb == 0
             and not self.config.cache_config.enable_prefix_caching
         ):
-            mm_uuid_items = {
-                modality: [f"{mm_req_id}-{modality}-{i}" for i in range(data_count)]
-                for modality, data_count in mm_data_items.get_all_counts().items()
-            }
+            mm_uuid_items = dict(mm_uuid_items)
+            for modality, data_count in mm_data_items.get_all_counts().items():
+                uuids = mm_uuid_items.get(modality)
+                if uuids is None:
+                    uuids = [None] * data_count
+                mm_uuid_items[modality] = [
+                    uuid if uuid is not None else f"{mm_req_id}-{modality}-{i}"
+                    for i, uuid in enumerate(uuids)
+                ]
 
         self._validate_mm_uuids(mm_data, mm_data_items, mm_uuid_items)
 
@@ -853,10 +839,7 @@ def _process_multimodal(
         *,
         skip_mm_cache: bool = False,
     ) -> "MultiModalInput":
-        if skip_mm_cache and self._readonly_mm_processor is not None:
-            mm_processor = self._readonly_mm_processor
-        else:
-            mm_processor = self.get_mm_processor()
+        mm_processor = self.get_mm_processor()
 
         mm_req_id = f"renderer{self.api_process_rank}-mm-{self._mm_req_counter.inc(1)}"
 
@@ -873,6 +856,11 @@ def _process_multimodal(
             mm_uuid_items,
             hf_processor_mm_kwargs=mm_processor_kwargs or {},
             media_io_kwargs=media_io_kwargs or {},
+            cache=(
+                self._mm_processor_only_cache
+                if skip_mm_cache
+                else self._mm_processor_cache
+            ),
         )
         mm_timing_ctx = self._mm_timing_registry.get(mm_req_id)
 
@@ -982,27 +970,103 @@ async def _process_tokens_async(
 
         return engine_input
 
+    def _truncate_expanded_prompt(
+        self,
+        engine_input: TokensInput | MultiModalInput,
+        params: TokenizeParams | None,
+    ) -> None:
+        """Re-apply `truncate_prompt_tokens` once placeholders are expanded.
+
+        `TokenizeParams.apply_post_tokenization` truncates the prompt before
+        multimodal expansion, so a single placeholder token can expand into
+        many and push the result back over the requested bound.
+
+        Multimodal spans are never cut: a request whose spans would not all
+        survive intact is rejected, matching the post-expansion rejection in
+        `vllm.v1.engine.input_processor._validate_model_input`.
+
+        Only multimodal inputs can grow after `apply_post_tokenization`, so
+        everything else returns immediately.
+        """
+        if engine_input["type"] != "multimodal":
+            return
+
+        if params is None or params.truncate_prompt_tokens is None:
+            return
+
+        prompt_token_ids = engine_input.get("prompt_token_ids")
+        if prompt_token_ids is None:
+            return
+
+        num_tokens = len(prompt_token_ids)
+        truncation = params._truncation_slice(self.tokenizer, num_tokens)
+        if truncation is None:
+            return
+
+        start, stop, _ = truncation.indices(num_tokens)
+
+        mm_placeholders = engine_input["mm_placeholders"]
+        for modality, positions in mm_placeholders.items():
+            for position in positions:
+                end = position.offset + position.length
+                if position.offset >= start and end <= stop:
+                    continue
+                raise VLLMValidationError(
+                    f"truncate_prompt_tokens={params.truncate_prompt_tokens} "
+                    f"would split or drop a(n) {modality} item, which spans "
+                    f"tokens {position.offset} to {end} of the "
+                    f"{num_tokens}-token prompt. Multimodal placeholders "
+                    f"expand after truncation is applied, so the limit must "
+                    f"leave every multimodal item intact. Please raise "
+                    f"truncate_prompt_tokens or shorten the multimodal input.",
+                    parameter="truncate_prompt_tokens",
+                    value=params.truncate_prompt_tokens,
+                )
+
+        if start:
+            engine_input["mm_placeholders"] = {
+                modality: [
+                    replace(position, offset=position.offset - start)
+                    for position in positions
+                ]
+                for modality, positions in mm_placeholders.items()
+            }
+
+        engine_input["prompt_token_ids"] = prompt_token_ids[truncation]
+
     def _process_singleton(
         self,
         prompt: SingletonTokPrompt,
         *,
         skip_mm_cache: bool = False,
+        tok_params: TokenizeParams | None = None,
     ) -> SingletonInput:
         if "prompt_embeds" in prompt:
             return self._process_embeds(prompt)  # type: ignore[arg-type]
 
-        return self._process_tokens(prompt, skip_mm_cache=skip_mm_cache)  # type: ignore[arg-type]
+        engine_input = self._process_tokens(  # type: ignore[arg-type]
+            prompt,  # type: ignore[arg-type]
+            skip_mm_cache=skip_mm_cache,
+        )
+        self._truncate_expanded_prompt(engine_input, tok_params)
+        return engine_input
 
     async def _process_singleton_async(
         self,
         prompt: SingletonTokPrompt,
         *,
         skip_mm_cache: bool = False,
+        tok_params: TokenizeParams | None = None,
     ) -> SingletonInput:
         if "prompt_embeds" in prompt:
             return self._process_embeds(prompt)  # type: ignore[arg-type]
 
-        return await self._process_tokens_async(prompt, skip_mm_cache=skip_mm_cache)  # type: ignore[arg-type]
+        engine_input = await self._process_tokens_async(  # type: ignore[arg-type]
+            prompt,  # type: ignore[arg-type]
+            skip_mm_cache=skip_mm_cache,
+        )
+        self._truncate_expanded_prompt(engine_input, tok_params)
+        return engine_input
 
     def _get_skip_decoder_start_token(self) -> bool:
         """Whether the multimodal processor supplies a complete decoder prefix."""
@@ -1068,12 +1132,15 @@ def process_for_engine(
         arrival_time: float,
         *,
         skip_mm_cache: bool = False,
+        tok_params: TokenizeParams | None = None,
     ) -> EngineInput:
         engine_input: EngineInput
         if "encoder_prompt" in prompt:
             engine_input = self._process_enc_dec(prompt, skip_mm_cache=skip_mm_cache)  # type: ignore[arg-type]
         else:
-            engine_input = self._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
+            engine_input = self._process_singleton(
+                prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params
+            )
 
         engine_input["arrival_time"] = arrival_time
 
@@ -1085,6 +1152,7 @@ async def process_for_engine_async(
         arrival_time: float,
         *,
         skip_mm_cache: bool = False,
+        tok_params: TokenizeParams | None = None,
     ) -> EngineInput:
         engine_input: EngineInput
         if "encoder_prompt" in prompt:
@@ -1094,7 +1162,7 @@ async def process_for_engine_async(
             )
         else:
             engine_input = await self._process_singleton_async(
-                prompt, skip_mm_cache=skip_mm_cache
+                prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params
             )
 
         engine_input["arrival_time"] = arrival_time
@@ -1121,7 +1189,9 @@ def render_cmpl(
         self._apply_prompt_extras(tok_prompts, prompt_extras)
 
         return [
-            self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache)
+            self.process_for_engine(
+                prompt, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params
+            )
             for prompt in tok_prompts
         ]
 
@@ -1146,7 +1216,7 @@ async def render_cmpl_async(
         return await asyncio.gather(
             *(
                 self.process_for_engine_async(
-                    p, arrival_time, skip_mm_cache=skip_mm_cache
+                    p, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params
                 )
                 for p in tok_prompts
             )
@@ -1184,7 +1254,9 @@ def render_chat(
         self._apply_prompt_extras(tok_prompts, prompt_extras)
 
         eng_prompts = [
-            self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache)
+            self.process_for_engine(
+                prompt, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params
+            )
             for prompt in tok_prompts
         ]
 
@@ -1224,7 +1296,7 @@ async def render_chat_async(
         eng_prompts = await asyncio.gather(
             *(
                 self.process_for_engine_async(
-                    p, arrival_time, skip_mm_cache=skip_mm_cache
+                    p, arrival_time, skip_mm_cache=skip_mm_cache, tok_params=tok_params
                 )
                 for p in tok_prompts
             )
diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py
index 8e6c4a5ac875..e22ebf1ebf1b 100644
--- a/vllm/renderers/cohere.py
+++ b/vllm/renderers/cohere.py
@@ -244,7 +244,7 @@ def _try_import_melody():
         raise ImportError(
             "The `cohere` tokenizer/renderer mode requires the "
             "`cohere_melody` package. Install it via "
-            "`pip install 'cohere-melody>=0.11.1'` or build from "
+            "`pip install 'cohere-melody>=0.14.0'` or build from "
             "https://github.com/cohere-ai/melody."
         ) from e
 
diff --git a/vllm/renderers/deepseek_v4.py b/vllm/renderers/deepseek_v4.py
index e93069209a20..470e21f1eb11 100644
--- a/vllm/renderers/deepseek_v4.py
+++ b/vllm/renderers/deepseek_v4.py
@@ -16,6 +16,10 @@
 from .inputs.preprocess import parse_dec_only_prompt
 from .params import ChatParams
 
+_FIM_BEGIN = "<｜fim▁begin｜>"
+_FIM_HOLE = "<｜fim▁hole｜>"
+_FIM_END = "<｜fim▁end｜>"
+
 
 class DeepseekV4Renderer(BaseRenderer[DeepseekV4Tokenizer]):
     def __init__(
@@ -32,6 +36,9 @@ def __init__(
     def _apply_chat_template(self, *args, **kwargs):
         return self.get_tokenizer().apply_chat_template(*args, **kwargs)
 
+    def render_completion_suffix(self, prompt: str, suffix: str) -> str | None:
+        return f"{_FIM_BEGIN}{prompt}{_FIM_HOLE}{suffix}{_FIM_END}"
+
     def render_messages(
         self,
         messages: list[ChatCompletionMessageParam],
@@ -40,7 +47,7 @@ def render_messages(
         conversation, mm_data, mm_uuids = parse_chat_messages(
             messages,
             self.model_config,
-            content_format="string",
+            content_format="openai",
             media_io_kwargs=params.media_io_kwargs,
             mm_processor_kwargs=params.mm_processor_kwargs,
         )
@@ -67,7 +74,7 @@ async def render_messages_async(
         conversation, mm_data, mm_uuids = await parse_chat_messages_async(
             messages,
             self.model_config,
-            content_format="string",
+            content_format="openai",
             media_io_kwargs=params.media_io_kwargs,
             mm_processor_kwargs=params.mm_processor_kwargs,
         )
diff --git a/vllm/renderers/embed_utils.py b/vllm/renderers/embed_utils.py
index c4079005a54a..981aeee3fb09 100644
--- a/vllm/renderers/embed_utils.py
+++ b/vllm/renderers/embed_utils.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from io import BytesIO
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Final
 
 import pybase64
 import torch
@@ -17,6 +17,27 @@
     from vllm.config import ModelConfig
 
 
+# How much of torch's own reason is repeated back to the caller. The reason is
+# built from bytes the caller sent, so it is bounded rather than echoed whole.
+_MAX_EMBED_ERROR_REASON_CHARS: Final = 200
+
+
+def _truncated_reason(exc: Exception) -> str:
+    """`exc`'s message, capped, saying how much was left out when it is capped.
+
+    Without the count a capped reason reads like the whole one, so a caller
+    debugging a large payload cannot tell that torch said more than this.
+    """
+    reason: Final = str(exc).strip() or type(exc).__name__
+    omitted: Final = len(reason) - _MAX_EMBED_ERROR_REASON_CHARS
+    if omitted <= 0:
+        return reason
+    return (
+        f"{reason[:_MAX_EMBED_ERROR_REASON_CHARS]}"
+        f"... ({omitted} more characters truncated)"
+    )
+
+
 def safe_load_prompt_embeds(
     model_config: "ModelConfig",
     embed: bytes,
@@ -27,12 +48,38 @@ def safe_load_prompt_embeds(
             parameter="prompt_embeds",
         )
 
+    # Decoded outside the try below so that a body which is not base64 keeps
+    # raising binascii.Error, which is a ValueError and already answered 400.
+    payload: Final = pybase64.b64decode(embed, validate=True)
+
     with check_sparse_tensor_invariants_threadsafe():
-        tensor = torch.load(
-            BytesIO(pybase64.b64decode(embed, validate=True)),
-            weights_only=True,
-            map_location=torch.device("cpu"),
-        )
+        try:
+            tensor = torch.load(
+                BytesIO(payload),
+                weights_only=True,
+                map_location=torch.device("cpu"),
+            )
+        except (torch.OutOfMemoryError, MemoryError):
+            # Running out of memory is the server's condition, not the caller's.
+            raise
+        except Exception as exc:
+            # Everything else here is about the bytes the caller sent, and torch
+            # has no single error type for "this is not a tensor file": the zip
+            # reader raises `RuntimeError`, the legacy pickle path raises
+            # `UnpicklingError`, `KeyError` or `IndexError`, an empty payload
+            # raises `EOFError`. None of them is a `ValueError`, so the
+            # entrypoints' fallback mapped them all to 500. Catching by
+            # behaviour rather than by type also survives torch changing which
+            # error it raises for a given malformed payload.
+            #
+            # torch's reason is worth keeping -- for a malformed sparse tensor
+            # it names the offending index -- but it is built from the caller's
+            # own bytes, so it is truncated rather than echoed whole.
+            raise VLLMValidationError(
+                "`prompt_embeds` could not be deserialized as a torch tensor: "
+                f"{_truncated_reason(exc)}",
+                parameter="prompt_embeds",
+            ) from exc
         tensor = safe_to_dense(tensor, parameter="prompt_embeds")
 
     if tensor.dim() > 2:
diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py
index 268a9aa46367..4ef14aee58cf 100644
--- a/vllm/renderers/hf.py
+++ b/vllm/renderers/hf.py
@@ -7,7 +7,7 @@
 import itertools
 import weakref
 from collections import defaultdict, deque
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
 
@@ -758,7 +758,6 @@ def safe_apply_chat_template(
     tools: list[dict[str, Any]] | None = ...,
     chat_template: str | None = ...,
     tokenize: Literal[True] = ...,
-    return_assistant_tokens_mask: Literal[False] = ...,
     **kwargs,
 ) -> list[int]: ...
 @overload
@@ -770,20 +769,8 @@ def safe_apply_chat_template(
     tools: list[dict[str, Any]] | None = ...,
     chat_template: str | None = ...,
     tokenize: Literal[False] = ...,
-    return_assistant_tokens_mask: Literal[False] = ...,
     **kwargs,
 ) -> str: ...
-@overload
-def safe_apply_chat_template(
-    model_config: ModelConfig,
-    tokenizer: HfTokenizer,
-    conversation: list[ConversationMessage],
-    *,
-    tools: list[dict[str, Any]] | None = ...,
-    chat_template: str | None = ...,
-    return_assistant_tokens_mask: Literal[True],
-    **kwargs,
-) -> tuple[list[int], list[int] | None]: ...
 def safe_apply_chat_template(
     model_config: ModelConfig,
     tokenizer: HfTokenizer,
@@ -792,9 +779,8 @@ def safe_apply_chat_template(
     tools: list[dict[str, Any]] | None = None,
     chat_template: str | None = None,
     tokenize: bool = True,
-    return_assistant_tokens_mask: bool = False,
     **kwargs,
-) -> str | list[int] | tuple[list[int], list[int] | None]:
+) -> str | list[int]:
     chat_template = resolve_chat_template(
         tokenizer,
         chat_template=chat_template,
@@ -822,38 +808,6 @@ def safe_apply_chat_template(
         chat_template_kwargs=kwargs,
     )
 
-    # assistant_tokens_mask requires tokenized output — force tokenize=True.
-    if return_assistant_tokens_mask:
-        tokenize = True
-
-    # When return_assistant_tokens_mask is requested and the template supports it,
-    # request assistant_tokens_mask via return_dict.
-    # Check for the actual Jinja tag, not just the word "generation"
-    # (which also appears in add_generation_prompt).
-    if return_assistant_tokens_mask and "{% generation %}" in chat_template:
-        resolved_kwargs["return_assistant_tokens_mask"] = True
-        resolved_kwargs["return_dict"] = True
-        resolved_kwargs.pop("tokenize", None)
-        try:
-            result = tokenizer.apply_chat_template(
-                conversation=conversation,  # type: ignore[arg-type]
-                tools=tools,  # type: ignore[arg-type]
-                chat_template=chat_template,
-                tokenize=True,
-                **resolved_kwargs,
-            )
-        except (TypeError, ValueError) as exc:
-            logger.warning(
-                "apply_chat_template failed for assistant_tokens_mask: %s", exc
-            )
-        else:
-            if isinstance(result, Mapping):
-                token_ids = list(result.get("input_ids", []))
-                mask_raw = result.get("assistant_masks")
-                mask = list(mask_raw) if mask_raw is not None else None
-                return token_ids, mask
-            return list(result), None
-
     # transformers v5 changed the default of `return_dict` to True, which
     # makes `apply_chat_template(tokenize=True)` return a `BatchEncoding`
     # instead of `list[int]`. Force `return_dict=False` so downstream code
@@ -879,9 +833,6 @@ def safe_apply_chat_template(
         logger.warning("Chat template rejected the request: %s", e)
         raise VLLMValidationError(_template_error_reason(e)) from e
 
-    if return_assistant_tokens_mask:
-        assert isinstance(plain, list), f"Expected list[int], got {type(plain)}"
-        return plain, None
     return plain
 
 
@@ -900,6 +851,7 @@ def rebuild_mm_uuids_from_mm_data(
 
     Returns:
         Updated UUIDs dictionary with chunk UUIDs
+
     """
     vision_chunks = mm_data.get("vision_chunk")
     if vision_chunks is None:
@@ -932,6 +884,7 @@ def build_video_prompts_from_mm_data(
 
     Returns:
         List of video prompts, one per video.
+
     """
     vision_chunks = mm_data.get("vision_chunk")
     if vision_chunks is None:
@@ -1013,8 +966,7 @@ def __init__(
             )
 
     def _can_produce_offsets(self) -> bool:
-        # HF tokenizers may be slow (use_fast=False); only fast tokenizers
-        # expose offset_mapping.
+        # Only fast tokenizers expose offset_mapping.
         return self.tokenizer is not None and self.tokenizer.is_fast
 
     def render_messages(
@@ -1064,22 +1016,12 @@ def render_messages(
                 logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
             chat_template_kwargs["tokenize"] = True
 
-        assistant_tokens_mask: list[int] | None = None
-        if params.return_assistant_tokens_mask:
-            prompt_raw, assistant_tokens_mask = safe_apply_chat_template(
-                model_config,
-                tokenizer,
-                conversation,
-                return_assistant_tokens_mask=True,
-                **chat_template_kwargs,
-            )
-        else:
-            prompt_raw = safe_apply_chat_template(
-                model_config,
-                tokenizer,
-                conversation,
-                **chat_template_kwargs,
-            )
+        prompt_raw = safe_apply_chat_template(
+            model_config,
+            tokenizer,
+            conversation,
+            **chat_template_kwargs,
+        )
 
         # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
         # model which uses unified vision chunks for both images and videos.
@@ -1105,9 +1047,6 @@ def render_messages(
 
         prompt = parse_dec_only_prompt(prompt_raw)
 
-        if assistant_tokens_mask is not None:
-            cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask
-
         # When `prompt_embeds` is mixed with other modality data,
         # `_process_tokens` runs `_process_multimodal` first (expanding
         # `<|AUDIO|>` / `<|IMAGE|>` placeholders) and then
@@ -1181,30 +1120,12 @@ async def render_messages_async(
                 logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
             chat_template_kwargs["tokenize"] = True
 
-        assistant_tokens_mask: list[int] | None = None
-        if params.return_assistant_tokens_mask:
-            result_with_mask = cast(
-                tuple[list[int], list[int] | None],
-                await make_async(
-                    safe_apply_chat_template,
-                    executor=self._executor,
-                )(
-                    model_config,
-                    tokenizer,
-                    conversation,
-                    return_assistant_tokens_mask=True,  # type: ignore[arg-type]
-                    **chat_template_kwargs,
-                ),
-            )
-            prompt_raw: str | list[int] = result_with_mask[0]
-            assistant_tokens_mask = result_with_mask[1]
-        else:
-            prompt_raw = await self._apply_chat_template_async(
-                model_config,
-                tokenizer,
-                conversation,
-                **chat_template_kwargs,
-            )
+        prompt_raw = await self._apply_chat_template_async(
+            model_config,
+            tokenizer,
+            conversation,
+            **chat_template_kwargs,
+        )
 
         # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
         # model which uses unified vision chunks for both images and videos.
@@ -1230,9 +1151,6 @@ async def render_messages_async(
 
         prompt = parse_dec_only_prompt(prompt_raw)
 
-        if assistant_tokens_mask is not None:
-            cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask
-
         # See `render_messages` for the rationale.
         if prompt_embeds_tensors and mm_data:
             assert prompt_embeds_placeholder_token_id is not None
@@ -1274,24 +1192,22 @@ def _process_tokens(
         processor records all placeholder offsets in the final (post-expansion)
         coordinate space, no offset shifting needed afterwards.
         """
-        assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None)
         prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
-        if prompt_embeds_info is not None:
-            tensors, placeholder_token_id = prompt_embeds_info
-            mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
-            cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
-                list(prompt["prompt_token_ids"]), mm_updates
-            )
+        if prompt_embeds_info is None:
+            return super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache)
+
+        tensors, placeholder_token_id = prompt_embeds_info
+        mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
+        cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
+            list(prompt["prompt_token_ids"]), mm_updates
+        )
         engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache)
-        if prompt_embeds_info is not None:
-            tensors, _ = prompt_embeds_info
-            self._apply_prompt_embeds_to_engine_input(
-                cast(MultiModalInput, engine_input),
-                tensors,
-                mm_updates,
-            )
-        if assistant_tokens_mask is not None:
-            engine_input["assistant_tokens_mask"] = assistant_tokens_mask
+        self._apply_prompt_embeds_to_engine_input(
+            cast(MultiModalInput, engine_input),
+            tensors,
+            mm_updates,
+        )
+
         return engine_input
 
     @override
@@ -1302,26 +1218,26 @@ async def _process_tokens_async(
         skip_mm_cache: bool = False,
     ) -> TokensInput | MultiModalInput:
         """Async equivalent of `_process_tokens`."""
-        assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None)
         prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
-        if prompt_embeds_info is not None:
-            tensors, placeholder_token_id = prompt_embeds_info
-            mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
-            cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
-                list(prompt["prompt_token_ids"]), mm_updates
+        if prompt_embeds_info is None:
+            return await super()._process_tokens_async(
+                prompt, skip_mm_cache=skip_mm_cache
             )
+
+        tensors, placeholder_token_id = prompt_embeds_info
+        mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
+        cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
+            list(prompt["prompt_token_ids"]), mm_updates
+        )
         engine_input = await super()._process_tokens_async(
             prompt, skip_mm_cache=skip_mm_cache
         )
-        if prompt_embeds_info is not None:
-            tensors, _ = prompt_embeds_info
-            self._apply_prompt_embeds_to_engine_input(
-                cast(MultiModalInput, engine_input),
-                tensors,
-                mm_updates,
-            )
-        if assistant_tokens_mask is not None:
-            engine_input["assistant_tokens_mask"] = assistant_tokens_mask
+        self._apply_prompt_embeds_to_engine_input(
+            cast(MultiModalInput, engine_input),
+            tensors,
+            mm_updates,
+        )
+
         return engine_input
 
     @staticmethod
diff --git a/vllm/renderers/inputs/preprocess.py b/vllm/renderers/inputs/preprocess.py
index dcf80ac73899..82a49f5638cb 100644
--- a/vllm/renderers/inputs/preprocess.py
+++ b/vllm/renderers/inputs/preprocess.py
@@ -1,6 +1,4 @@
-"""
-Schemas and utilities for preprocessing inputs.
-"""
+"""Schemas and utilities for preprocessing inputs."""
 
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -90,8 +88,7 @@ def conversation_to_seq(
 
 
 class EncoderDecoderDictPrompt(TypedDict):
-    """
-    A [`EncoderDecoderPrompt`][vllm.inputs.llm.EncoderDecoderPrompt]
+    """A [`EncoderDecoderPrompt`][vllm.inputs.llm.EncoderDecoderPrompt]
     that has been standardized into a dictionary.
     """
 
@@ -130,9 +127,7 @@ def _validate_prompt_dict(prompt: Mapping[str, object]) -> None:
 
 
 def parse_dec_only_prompt(prompt: PromptType | object) -> DecoderOnlyDictPrompt:
-    """
-    Parse a prompt for a decoder-only model and normalize it to a dictionary.
-    """
+    """Parse a prompt for a decoder-only model and normalize it to a dictionary."""
     if isinstance(prompt, str):
         return TextPrompt(prompt=prompt)
 
@@ -216,9 +211,7 @@ def _parse_dec_prompt(prompt: PromptType | object) -> DecoderDictPrompt:
 
 
 def parse_enc_dec_prompt(prompt: PromptType | object) -> EncoderDecoderDictPrompt:
-    """
-    Parse a prompt for an encoder-decoder model and normalize it to a dictionary.
-    """
+    """Parse a prompt for an encoder-decoder model and normalize it to a dictionary."""
     if isinstance(prompt, dict) and "encoder_prompt" in prompt:
         enc_prompt = prompt["encoder_prompt"]  # type: ignore[typeddict-item]
         dec_prompt = prompt["decoder_prompt"]  # type: ignore[typeddict-item]
diff --git a/vllm/renderers/inputs/tokenize.py b/vllm/renderers/inputs/tokenize.py
index 4168e201203e..bf85051c379d 100644
--- a/vllm/renderers/inputs/tokenize.py
+++ b/vllm/renderers/inputs/tokenize.py
@@ -1,6 +1,4 @@
-"""
-Schemas and utilities for tokenization inputs.
-"""
+"""Schemas and utilities for tokenization inputs."""
 
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -30,8 +28,7 @@
 
 
 class EncoderDecoderTokPrompt(TypedDict):
-    """
-    A
+    """A
     [`EncoderDecoderDictPrompt`][vllm.renderers.inputs.preprocess.EncoderDecoderDictPrompt]
     that has been tokenized.
     """
diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py
index ca70bb6ace44..6f7cd1075819 100644
--- a/vllm/renderers/online_derenderer.py
+++ b/vllm/renderers/online_derenderer.py
@@ -4,7 +4,10 @@
 
 from vllm.config import ModelConfig
 from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
-from vllm.entrypoints.generate.base.protocol import DeltaMessage, ToolCall
+from vllm.entrypoints.generate.base.protocol import (
+    DeltaMessage,
+    ToolCall,
+)
 from vllm.entrypoints.generate.base.serving import resolve_token_id_placeholder
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionLogProbs,
@@ -29,6 +32,9 @@
 )
 from vllm.entrypoints.serve.engine.protocol import UsageInfo
 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
+from vllm.entrypoints.serve.utils.tool_calls_utils import (
+    maybe_filter_parallel_tool_calls,
+)
 from vllm.logger import init_logger
 from vllm.parser import Parser, ParserManager
 from vllm.renderers import BaseRenderer
@@ -54,6 +60,7 @@ def __init__(
         exclude_tools_when_tool_choice_none: bool = False,
         tool_parser: str | None = None,
         reasoning_parser: str | None = None,
+        tool_strict_level: str = "auto",
         default_chat_template_kwargs: dict[str, Any] | None = None,
         log_error_stack: bool = False,
     ) -> None:
@@ -68,6 +75,7 @@ def __init__(
             tool_parser_name=tool_parser,
             reasoning_parser_name=reasoning_parser,
             enable_auto_tools=enable_auto_tools,
+            tool_strict_level=tool_strict_level,
             model_name=model_config.model,
             is_harmony=self.use_harmony,
         )
@@ -93,6 +101,13 @@ def __init__(
         self._derender_completion_async = make_async(
             self._derender_completion, executor=renderer._executor
         )
+        self._detokenize_delta_async = make_async(
+            self._detokenize_delta, executor=renderer._executor
+        )
+        # Replay is O(n) per chunk, so it must not run on the event loop.
+        self._derender_chat_stream_parsed_async = make_async(
+            self._derender_chat_stream_parsed, executor=renderer._executor
+        )
 
     async def derender_chat(
         self,
@@ -142,6 +157,7 @@ def _derender_chat(
                     tokenizer,
                     chat_request.tools,
                     chat_template_kwargs=chat_template_kwargs,
+                    model_config=self.model_config,
                 )
                 reasoning, content, tool_calls = parser.parse(
                     decoded_text,
@@ -233,6 +249,7 @@ def _detokenize_delta(
         Returns:
             (new_text, updated_state) — the delta text for this chunk and the
             state to pass to the next call.
+
         """
         prev_tokens = list(state.prev_tokens)
         prefix_offset = state.prefix_offset
@@ -275,53 +292,69 @@ async def derender_chat_stream(
         state: DerenderStreamState | None = None,
         chat_request: ChatCompletionRequest | None = None,
         prompt_tokens: int | None = None,
+        prompt_token_ids: list[int] | None = None,
     ) -> tuple[ChatCompletionStreamResponse, DerenderStreamState]:
         """Process one GenerateStreamResponse chunk for streaming chat derender.
 
-        TODO: parse path for reasoning and tool calls is implemented in future PR.
-
-        Unlike OpenAI's API, which always emits ``role: "assistant"`` on the
+        Unlike OpenAI's API, which always emits `role: "assistant"` on the
         very first chunk, this emits it on the first chunk with a non empty
-        ``choices`` list. A leading usage only chunk therefore defers the
+        `choices` list. A leading usage only chunk therefore defers the
         role to the following content chunk instead of sending an empty
         role only delta.
 
         Args:
             model: Model name for the response object.
-            generate_chunk: One SSE chunk from ``/inference/v1/generate``.
-            state: Client carried detok state (``None`` for first call).
-            chat_request: Original ChatCompletionRequest from ``/render``.
+            generate_chunk: One SSE chunk from `/inference/v1/generate`.
+            state: Client carried detok state (`None` for first call).
+            chat_request: Original ChatCompletionRequest from `/render`.
+                Required when a reasoning or tool parser is configured
+                (validated by the caller — see `ServingDerender`) because
+                plain detokenization would leak raw parser markup into `content`.
             prompt_tokens: Prompt token count for the usage chunk.
+            prompt_token_ids: Prompt token IDs. Parser path only. Lets
+                `parse_delta` settle its initial reasoning state. Required
+                when a reasoning or tool parser is configured (validated by
+                the caller — see ``ServingDerender``). Without it reasoning
+                left open by the prompt would be misclassified as content.
 
         Returns:
             (chunk, updated_state) — the derendered SSE chunk and the state
             the client must pass to the next call.
+
         """
         if state is None:
             state = DerenderStreamState()
 
-        if self.parser is not None:
-            # TODO: Follow on PR will implement the parse path.  Check on the
-            # parser alone (fail closed). A parser configured model must never
-            # fall through to plain detok on the streaming path, even when
-            # ``chat_request`` is omitted or reasoning/tool markup would leak
-            # into ``delta.content``.
-            raise NotImplementedError(
-                "Streaming chat derender is not yet supported for models with "
-                "a reasoning or tool parser configured. Use the non-streaming "
-                "derender endpoint (stream=false) for parsed output."
-            )
-
         # A single DerenderStreamState is threaded through every choice in
         # this chunk. Correct only when there is at most one choice per SSE
         # event (n=1, one call per index), as the streaming derender
         # protocol assumes. Multiple choices sharing one chunk would corrupt
-        # each other's detok window.
+        # each other's detok/parser state.
         if len(generate_chunk.choices) > 1:
             raise ValueError(
                 "derender_chat_stream expects at most one choice per chunk"
             )
 
+        parser_cls = self.parser
+        if parser_cls is not None:
+            # Fail (mirrors ServingDerender's pre-check) because a parser
+            # configured model must never fall through to plain detok or
+            # reasoning/tool markup would leak into `delta.content`.
+            if chat_request is None:
+                raise ValueError(
+                    "chat_request is required for streaming chat derender "
+                    "when a tool or reasoning parser is configured"
+                )
+            return await self._derender_chat_stream_parsed_async(
+                parser_cls,
+                model,
+                generate_chunk,
+                state,
+                chat_request,
+                prompt_tokens,
+                prompt_token_ids,
+            )
+
         tokenizer = self.renderer.get_tokenizer()
         skip_special = (
             chat_request.skip_special_tokens if chat_request is not None else True
@@ -331,7 +364,7 @@ async def derender_chat_stream(
 
         for choice in generate_chunk.choices:
             delta_tids = choice.token_ids or []
-            new_text, updated_state = self._detokenize_delta(
+            new_text, updated_state = await self._detokenize_delta_async(
                 tokenizer, delta_tids, updated_state, skip_special_tokens=skip_special
             )
 
@@ -367,6 +400,205 @@ async def derender_chat_stream(
             model=model,
             choices=stream_choices,
             usage=usage,
+            metrics=generate_chunk.metrics,
+        )
+        return chunk, updated_state
+
+    def _derender_chat_stream_parsed(
+        self,
+        parser_cls: type[Parser],
+        model: str,
+        generate_chunk: GenerateStreamResponse,
+        state: DerenderStreamState,
+        chat_request: ChatCompletionRequest,
+        prompt_tokens: int | None,
+        prompt_token_ids: list[int] | None,
+    ) -> tuple[ChatCompletionStreamResponse, DerenderStreamState]:
+        """Parser path for streaming chat derender: replay + `parse_delta`.
+
+        Parser internal state (buffered markup, reasoning/tool phase, etc.)
+        cannot be serialized into `DerenderStreamState`, so each call
+        builds a fresh parser and replays every prior output token through
+        `parse_delta` (discarding the result) before processing this
+        chunk's tokens for real.
+
+        Every chunk, replayed or live, goes through one `parse_delta` call
+        with the tokens it arrived with (more than one under e.g.
+        speculative decoding). Replay recovers those boundaries from
+        `state.output_chunk_lens`, so the rebuilt parser state matches
+        standard serving, which calls `parse_delta` once per engine step.
+        Text comes from a fresh incremental detokenizer with special tokens
+        preserved (``skip_special_tokens=False``).
+        """
+        tokenizer = self.renderer.get_tokenizer()
+
+        chat_template_kwargs: dict[str, Any] = {}
+        if not self.use_harmony:
+            chat_template_kwargs = (
+                chat_request.build_chat_params(
+                    self.chat_template,
+                    self.chat_template_content_format,
+                )
+                .with_defaults(self.default_chat_template_kwargs)
+                .chat_template_kwargs
+            )
+
+        parser = parser_cls(
+            tokenizer,
+            chat_request.tools,
+            chat_template_kwargs=chat_template_kwargs,
+            model_config=self.model_config,
+        )
+
+        # Ephemeral incremental detok window, local to this call. Threaded
+        # across both the replay and current chunk phases (via `nonlocal`)
+        # so multi-byte characters split across that boundary still decode
+        # correctly. Discarded once the call returns.
+        detok_state = DerenderStreamState()
+
+        def _replay(token_ids: list[int], chunk_lens: list[int]) -> None:
+            """Replay prior chunks through `parse_delta` to rebuild parser
+            state, discarding the result. Never `finished` since that only
+            applies to the current chunk.
+            """
+            nonlocal detok_state
+            start = 0
+            for chunk_len in chunk_lens:
+                chunk = token_ids[start : start + chunk_len]
+                start += chunk_len
+                text, detok_state = self._detokenize_delta(
+                    tokenizer, chunk, detok_state, skip_special_tokens=False
+                )
+                parser.parse_delta(
+                    text,
+                    chunk,
+                    chat_request,
+                    prompt_token_ids=prompt_token_ids,
+                    finished=False,
+                )
+
+        # Replay history to reconstruct parser state. The result is thrown
+        # away and only the current chunk's emission goes to the client.
+        _replay(state.output_token_ids, state.output_chunk_lens)
+
+        stream_choices: list[ChatCompletionResponseStreamChoice] = []
+        output_token_ids = list(state.output_token_ids)
+        output_chunk_lens = list(state.output_chunk_lens)
+        role_sent = state.role_sent
+        tools_streamed = state.tools_streamed
+        last_tool_call_ids = list(state.last_tool_call_ids)
+
+        # At most one choice: the caller (derender_chat_stream) already
+        # rejects >1 before dispatching here. role_sent/tools_streamed/
+        # output_token_ids below are updated for that single choice, not
+        # accumulated across choices. Looping over generate_chunk.choices
+        # could silently corrupt output if chunks were ever allowed to
+        # contain multiple choices.
+        if generate_chunk.choices:
+            choice = generate_chunk.choices[0]
+            delta_tids = choice.token_ids or []
+            is_finished = choice.finish_reason is not None
+
+            if delta_tids:
+                # One parse_delta call for the whole chunk (producer
+                # granularity), not one per token. See the granularity note
+                # in this method's docstring.
+                text, detok_state = self._detokenize_delta(
+                    tokenizer, delta_tids, detok_state, skip_special_tokens=False
+                )
+                delta_message = parser.parse_delta(
+                    text,
+                    delta_tids,
+                    chat_request,
+                    prompt_token_ids=prompt_token_ids,
+                    finished=is_finished,
+                )
+            elif is_finished:
+                # Finish only chunk (no new tokens). Still flush any
+                # buffered tool call arguments.
+                delta_message = parser.parse_delta(
+                    "",
+                    [],
+                    chat_request,
+                    prompt_token_ids=prompt_token_ids,
+                    finished=True,
+                )
+            else:
+                delta_message = None
+
+            if delta_tids:
+                output_token_ids.extend(delta_tids)
+                output_chunk_lens.append(len(delta_tids))
+
+            if delta_message is None:
+                delta_message = DeltaMessage()
+
+            if delta_message.tool_calls:
+                tools_streamed = True
+                for tc in delta_message.tool_calls:
+                    if tc.id is None:
+                        continue
+                    if tc.index < len(last_tool_call_ids):
+                        # Pin: reuse the ID already recorded for this index
+                        # rather than one a from scratch replay regenerated.
+                        # Real trigger not just defensive with
+                        # tool_choice="required",
+                        # extract_required_tool_call_streaming resets
+                        # function_name_returned to False whenever the
+                        # partial JSON transiently fails to parse which
+                        # re-emits id+name for the same index on replay.
+                        tc.id = last_tool_call_ids[tc.index]
+                    else:
+                        last_tool_call_ids.append(tc.id)
+
+            if not role_sent:
+                delta_message.role = "assistant"
+                role_sent = True
+
+            finish_reason = choice.finish_reason
+            if finish_reason is not None:
+                is_named_tool_choice = (
+                    type(chat_request.tool_choice) is ChatCompletionNamedToolChoiceParam
+                )
+                if tools_streamed and not is_named_tool_choice:
+                    finish_reason = "tool_calls"
+
+            stream_choice = ChatCompletionResponseStreamChoice(
+                index=choice.index,
+                delta=delta_message,
+                finish_reason=finish_reason,
+            )
+            stream_choices.append(
+                maybe_filter_parallel_tool_calls(stream_choice, chat_request)
+            )
+
+        updated_state = state.model_copy(
+            update={
+                "output_token_ids": output_token_ids,
+                "output_chunk_lens": output_chunk_lens,
+                "role_sent": role_sent,
+                "tools_streamed": tools_streamed,
+                "last_tool_call_ids": last_tool_call_ids,
+            }
+        )
+
+        usage: UsageInfo | None = None
+        if generate_chunk.usage is not None:
+            u = generate_chunk.usage
+            pt = prompt_tokens if prompt_tokens is not None else (u.prompt_tokens or 0)
+            ct = u.completion_tokens or 0
+            usage = UsageInfo(
+                prompt_tokens=pt,
+                completion_tokens=ct,
+                total_tokens=pt + ct,
+            )
+
+        chunk = ChatCompletionStreamResponse(
+            id=generate_chunk.request_id,
+            model=model,
+            choices=stream_choices,
+            usage=usage,
+            metrics=generate_chunk.metrics,
         )
         return chunk, updated_state
 
@@ -461,6 +693,7 @@ async def derender_completion_stream(
 
         Returns:
             (chunk, updated_state) — the derendered chunk and updated state.
+
         """
         if state is None:
             state = DerenderStreamState()
@@ -485,7 +718,7 @@ async def derender_completion_stream(
 
         for choice in generate_chunk.choices:
             delta_tids = choice.token_ids or []
-            new_text, updated_state = self._detokenize_delta(
+            new_text, updated_state = await self._detokenize_delta_async(
                 tokenizer, delta_tids, updated_state, skip_special_tokens=skip_special
             )
             stream_choices.append(
@@ -512,6 +745,7 @@ async def derender_completion_stream(
             model=model,
             choices=stream_choices,
             usage=usage,
+            metrics=generate_chunk.metrics,
         )
         return chunk, updated_state
 
diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py
index 2ab5725b73fc..e5169d70ba91 100644
--- a/vllm/renderers/online_renderer.py
+++ b/vllm/renderers/online_renderer.py
@@ -4,7 +4,7 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
 from http import HTTPStatus
-from typing import TYPE_CHECKING, Any, TypeAlias
+from typing import TYPE_CHECKING, Any, TypeAlias, cast
 
 from openai.types.responses import ResponseFunctionToolCall, ResponseOutputItem
 from openai.types.responses.tool import Mcp, Tool
@@ -127,6 +127,7 @@ def __init__(
         exclude_tools_when_tool_choice_none: bool = False,
         tool_parser: str | None = None,
         reasoning_parser: str | None = None,
+        tool_strict_level: str = "auto",
         default_chat_template_kwargs: dict[str, Any] | None = None,
         log_error_stack: bool = False,
     ) -> None:
@@ -141,6 +142,7 @@ def __init__(
             tool_parser_name=tool_parser,
             reasoning_parser_name=reasoning_parser,
             enable_auto_tools=enable_auto_tools,
+            tool_strict_level=tool_strict_level,
             model_name=model_config.model,
             is_harmony=self.use_harmony,
         )
@@ -546,9 +548,58 @@ async def render_completion(
         Called directly by render_completion_request and delegated to by
         OpenAIServingCompletion.render_completion_request after its engine-aware checks.
         """
-        # Return error for unsupported features.
+        prompt_input = request.prompt
         if request.suffix is not None:
-            return self.create_error_response("suffix is not currently supported")
+            if request.echo:
+                return self.create_error_response(
+                    "Echo is unsupported with suffix.",
+                    param="suffix",
+                )
+
+            if request.prompt_embeds is not None:
+                return self.create_error_response(
+                    "suffix is not supported with prompt_embeds",
+                    param="suffix",
+                )
+
+            if request.truncate_prompt_tokens is not None:
+                return self.create_error_response(
+                    "suffix is not supported with truncate_prompt_tokens",
+                    param="suffix",
+                )
+
+            if isinstance(request.prompt, str):
+                rendered_prompt = self.renderer.render_completion_suffix(
+                    request.prompt, request.suffix
+                )
+                if rendered_prompt is None:
+                    return self.create_error_response(
+                        "suffix is only supported for models with FIM completion "
+                        "rendering",
+                        param="suffix",
+                    )
+                prompt_input = rendered_prompt
+            elif isinstance(request.prompt, list) and all(
+                isinstance(prompt, str) for prompt in request.prompt
+            ):
+                rendered_prompts = []
+                for prompt in cast(list[str], request.prompt):
+                    rendered_prompt = self.renderer.render_completion_suffix(
+                        prompt, request.suffix
+                    )
+                    if rendered_prompt is None:
+                        return self.create_error_response(
+                            "suffix is only supported for models with FIM completion "
+                            "rendering",
+                            param="suffix",
+                        )
+                    rendered_prompts.append(rendered_prompt)
+                prompt_input = rendered_prompts
+            else:
+                return self.create_error_response(
+                    "suffix requires text prompt input for FIM completion rendering",
+                    param="suffix",
+                )
 
         if request.echo and request.prompt_embeds is not None:
             return self.create_error_response("Echo is unsupported with prompt embeds.")
@@ -560,7 +611,7 @@ async def render_completion(
 
         engine_inputs = await self.preprocess_completion(
             request,
-            prompt_input=request.prompt,
+            prompt_input=prompt_input,
             prompt_embeds=request.prompt_embeds,
             skip_mm_cache=skip_mm_cache,
         )
diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py
index b651a060955f..9f76a2ed49a4 100644
--- a/vllm/renderers/params.py
+++ b/vllm/renderers/params.py
@@ -27,13 +27,9 @@
 # Prompt keys whose entry i describes token i of the prompt, so they must be
 # reduced with the same slice as the prompt tokens themselves. Anything added
 # here is truncated by `TokenizeParams.apply_post_tokenization`.
-# `_assistant_tokens_mask` is renderer-internal: `HfRenderer.render_messages`
-# stashes it on the prompt so `_process_tokens` can move it onto the engine
-# input.
 _PARALLEL_TO_PROMPT_TOKENS = (
     "prompt_token_offsets",
     "prompt_is_token_ids",
-    "_assistant_tokens_mask",
 )
 
 
@@ -99,9 +95,6 @@ class ChatParams:
     mm_processor_kwargs: dict[str, Any] | None = None
     """The kwargs to pass to the multi-modal processor."""
 
-    return_assistant_tokens_mask: bool = False
-    """Request a per-token assistant mask from apply_chat_template."""
-
     tool_choice: Any | None = None
     """Request-level tool choice for renderers that need API metadata."""
 
@@ -136,7 +129,6 @@ def with_defaults(
                 default_mm_processor_kwargs,
                 self.mm_processor_kwargs,
             ),
-            return_assistant_tokens_mask=self.return_assistant_tokens_mask,
             tool_choice=self.tool_choice,
             response_format=self.response_format,
         )
@@ -425,8 +417,7 @@ def apply_pre_tokenization(
         tokenizer: TokenizerLike | None,
         prompt: TextPrompt,
     ) -> TextPrompt:
-        """
-        Ensure that the prompt meets the requirements set out by this config.
+        """Ensure that the prompt meets the requirements set out by this config.
         If that is not possible, raise a `VLLMValidationError`.
 
         This method is run before tokenization occurs.
@@ -538,8 +529,7 @@ def apply_post_tokenization(
         tokenizer: TokenizerLike | None,
         prompt: TokensPrompt | EmbedsPrompt,
     ) -> TokensPrompt | EmbedsPrompt:
-        """
-        Ensure that the prompt meets the requirements set out by this config.
+        """Ensure that the prompt meets the requirements set out by this config.
         If that is not possible, raise a `VLLMValidationError`.
 
         This method is run after tokenization occurs.
diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py
index 52a65aefbce1..77dd91b6b62c 100644
--- a/vllm/sampling_params.py
+++ b/vllm/sampling_params.py
@@ -127,9 +127,7 @@ def __post_init__(self):
             )
 
     def all_constraints_none(self) -> bool:
-        """
-        Returns True if all structured-output constraint fields are None.
-        """
+        """Returns True if all structured-output constraint fields are None."""
         return all(
             getattr(self, field) is None
             for field in (
@@ -143,9 +141,7 @@ def all_constraints_none(self) -> bool:
         )
 
     def all_non_structural_tag_constraints_none(self) -> bool:
-        """
-        Returns True if all structured-output constraint fields are None.
-        """
+        """Returns True if all structured-output constraint fields are None."""
         return all(
             getattr(self, field) is None
             for field in (
@@ -530,6 +526,16 @@ def __post_init__(self) -> None:
 
         self._verify_args()
 
+        if (
+            envs.VLLM_BATCH_INVARIANT
+            and self.temperature >= _SAMPLING_EPS
+            and self.seed is None
+        ):
+            logger.warning_once(
+                "Random sampling without an explicit seed may not be batch "
+                "invariant. Set seed in SamplingParams or use temperature=0."
+            )
+
         if self.temperature < _SAMPLING_EPS:
             # Zero temperature means greedy sampling.
             self.top_p = 1.0
@@ -692,7 +698,7 @@ def update_from_generation_config(
         generation_config: dict[str, Any],
         eos_token_id: int | None = None,
     ) -> None:
-        """Update if there are non-default values from generation_config"""
+        """Update if there are non-default values from generation_config."""
         if not self.ignore_eos:
             self._eos_token_id = eos_token_id
 
@@ -827,7 +833,6 @@ def verify(
         self._validate_logit_bias(model_config)
         self._validate_trace_replay(model_config, speculative_config)
         self._validate_stop_token_ids(model_config)
-        self._validate_logits_processors(model_config)
         self._validate_allowed_token_ids(model_config)
         self._validate_spec_decode(speculative_config)
         self._validate_diffusion(model_config)
@@ -997,13 +1002,6 @@ def _validate_trace_replay(
                 value=invalid_token_ids,
             )
 
-    def _validate_logits_processors(self, model_config: ModelConfig) -> None:
-        from vllm.v1.sample.logits_processor import (
-            validate_logits_processors_parameters,
-        )
-
-        validate_logits_processors_parameters(model_config.logits_processors, self)
-
     def _validate_allowed_token_ids(self, model_config: ModelConfig) -> None:
         allowed_token_ids = self.allowed_token_ids
         if allowed_token_ids is None:
@@ -1156,6 +1154,18 @@ def _validate_structured_outputs(
             raise VLLMValidationError(
                 "structured_outputs.regex must not contain a NUL character ('\\x00')"
             )
+        # Note(arpera):
+        # We do NOT check here structured output regex on emptiness because
+        # empty regex is indeed compiles to a valid grammar as well as
+        # whitespace-only regexps, for instance, regex="\n" or regex=" "
+        # are valid patterns and we MUST process them.
+        if (
+            isinstance(self.structured_outputs.structural_tag, str)
+            and self.structured_outputs.structural_tag.strip() == ""
+        ):
+            raise VLLMValidationError(
+                "structured_outputs.structural_tag cannot be an empty string"
+            )
 
         from vllm.v1.structured_output.backend_guidance import (
             has_guidance_unsupported_json_features,
@@ -1313,6 +1323,7 @@ class BeamSearchParams(
     length_penalty: float = 1.0
     include_stop_str_in_output: bool = False
     structured_outputs: StructuredOutputsParams | None = None
+    skip_special_tokens: bool = True
 
     def __post_init__(self) -> None:
         _verify_num_sequences(self.beam_width, "beam_width")
diff --git a/vllm/scalar_type.py b/vllm/scalar_type.py
index db52e93465ca..311b328c40cf 100644
--- a/vllm/scalar_type.py
+++ b/vllm/scalar_type.py
@@ -21,8 +21,7 @@ class NanRepr(Enum):
 # in sync until the inductor fully supports custom C++ classes.
 @dataclass(frozen=True)
 class ScalarType:
-    """
-    ScalarType can represent a wide range of floating point and integer
+    """ScalarType can represent a wide range of floating point and integer
     types, in particular it can be used to represent sub-byte data types
     (something that torch.dtype currently does not support). It is also
     capable of  representing types with a bias, i.e.:
@@ -135,8 +134,7 @@ def _raw_min(self) -> int | float:
 
     @functools.cached_property
     def id(self) -> int:
-        """
-        Convert the ScalarType to an int which can be passed to pytorch custom
+        """Convert the ScalarType to an int which can be passed to pytorch custom
         ops. This layout of the int must be kept in sync with the C++
         ScalarType's from_id method.
         """
@@ -168,56 +166,51 @@ def size_bits(self) -> int:
         return self.exponent + self.mantissa + int(self.signed)
 
     def min(self) -> int | float:
-        """
-        Min representable value for this scalar type.
+        """Min representable value for this scalar type.
         (accounting for bias if there is one)
         """
         return self._raw_min() - self.bias
 
     def max(self) -> int | float:
-        """
-        Max representable value for this scalar type.
+        """Max representable value for this scalar type.
         (accounting for bias if there is one)
         """
         return self._raw_max() - self.bias
 
     def is_signed(self) -> bool:
-        """
-        If the type is signed (i.e. has a sign bit), same as `signed`
+        """If the type is signed (i.e. has a sign bit), same as `signed`
         added for consistency with:
         https://pytorch.org/docs/stable/generated/torch.Tensor.is_signed.html
         """
         return self.signed
 
     def is_floating_point(self) -> bool:
-        "If the type is a floating point type"
+        """If the type is a floating point type."""
         return self.exponent != 0
 
     def is_integer(self) -> bool:
-        "If the type is an integer type"
+        """If the type is an integer type."""
         return self.exponent == 0
 
     def has_bias(self) -> bool:
-        "If the type has a non-zero bias"
+        """If the type has a non-zero bias."""
         return self.bias != 0
 
     def has_infs(self) -> bool:
-        "If the type is floating point and supports infinity"
+        """If the type is floating point and supports infinity."""
         return not self._finite_values_only
 
     def has_nans(self) -> bool:
         return self.nan_repr != NanRepr.NONE.value
 
     def is_ieee_754(self) -> bool:
-        """
-        If the type is a floating point type that follows IEEE 754
+        """If the type is a floating point type that follows IEEE 754
         conventions
         """
         return self.nan_repr == NanRepr.IEEE_754.value and not self._finite_values_only
 
     def __str__(self) -> str:
-        """
-        naming generally follows: https://github.com/jax-ml/ml_dtypes
+        """Naming generally follows: https://github.com/jax-ml/ml_dtypes
         for floating point types (leading f) the scheme is:
         `float<size_bits>_e<exponent_bits>m<mantissa_bits>[flags]`
         flags:
@@ -265,7 +258,7 @@ def __len__(self) -> int:
 
     @classmethod
     def int_(cls, size_bits: int, bias: int | None) -> "ScalarType":
-        "Create a signed integer scalar type (size_bits includes sign-bit)."
+        """Create a signed integer scalar type (size_bits includes sign-bit)."""
         ret = cls(0, size_bits - 1, True, bias if bias else 0)
         ret.id  # noqa B018: make sure the id is cached
         return ret
@@ -279,8 +272,7 @@ def uint(cls, size_bits: int, bias: int | None) -> "ScalarType":
 
     @classmethod
     def float_IEEE754(cls, exponent: int, mantissa: int) -> "ScalarType":
-        """
-        Create a standard floating point type
+        """Create a standard floating point type
         (i.e. follows IEEE 754 conventions).
         """
         assert mantissa > 0 and exponent > 0
@@ -292,8 +284,7 @@ def float_IEEE754(cls, exponent: int, mantissa: int) -> "ScalarType":
     def float_(
         cls, exponent: int, mantissa: int, finite_values_only: bool, nan_repr: NanRepr
     ) -> "ScalarType":
-        """
-        Create a non-standard floating point type
+        """Create a non-standard floating point type
         (i.e. does not follow IEEE 754 conventions).
         """
         assert mantissa > 0 and exponent > 0
@@ -325,8 +316,13 @@ def from_id(cls, scalar_type_id: int):
 
 
 class scalar_types:
+    uint2 = ScalarType.uint(2, None)
+    uint3 = ScalarType.uint(3, None)
     int4 = ScalarType.int_(4, None)
     uint4 = ScalarType.uint(4, None)
+    uint5 = ScalarType.uint(5, None)
+    uint6 = ScalarType.uint(6, None)
+    uint7 = ScalarType.uint(7, None)
     int8 = ScalarType.int_(8, None)
     uint8 = ScalarType.uint(8, None)
     float8_e4m3fn = ScalarType.float_(4, 3, True, NanRepr.EXTD_RANGE_MAX_MIN)
diff --git a/vllm/transformers_utils/configs/mellum.py b/vllm/snapshot/__init__.py
similarity index 51%
rename from vllm/transformers_utils/configs/mellum.py
rename to vllm/snapshot/__init__.py
index 2bed53394b25..8c4fde128286 100644
--- a/vllm/transformers_utils/configs/mellum.py
+++ b/vllm/snapshot/__init__.py
@@ -1,7 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from transformers import Qwen3MoeConfig
 
-
-class MellumConfig(Qwen3MoeConfig):
-    model_type = "mellum"
+"""Create and restore explicitly compatible vLLM engine snapshots."""
diff --git a/vllm/snapshot/controller.py b/vllm/snapshot/controller.py
new file mode 100644
index 000000000000..d5564fb8632d
--- /dev/null
+++ b/vllm/snapshot/controller.py
@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Create and restore transaction orchestration for local snapshots."""
+
+import argparse
+import shutil
+import sys
+from pathlib import Path
+
+from vllm.snapshot.manifest import (
+    read_manifest,
+    validate_artifact_root,
+    validate_identity,
+)
+from vllm.snapshot.runtime import (
+    LocalSnapshotTools,
+    SnapshotCreateError,
+    SnapshotRestoreError,
+    _error_detail,
+)
+from vllm.snapshot.types import Oracle, oracles_match
+
+
+def _child_engine_argv(
+    args: argparse.Namespace, engine_argv: tuple[str, ...] | None
+) -> tuple[str, ...]:
+    """Build the child argv from the first argv source that is present.
+
+    In order: an explicit ``engine_argv`` from the caller, the process argv
+    tail when invoked as ``vllm snapshot create ...``, then the bare model from
+    the parsed args. Whichever source wins, snapshot options are stripped and
+    sleep mode is appended.
+    """
+    if not engine_argv:
+        process_argv = tuple(sys.argv[1:])
+        if process_argv[:2] == ("snapshot", "create"):
+            engine_argv = process_argv[2:]
+        else:
+            model = getattr(args, "model_tag", None) or getattr(args, "model", None)
+            if not model:
+                raise SnapshotCreateError("snapshot create requires a model argument")
+            engine_argv = (str(model),)
+    child_argv: list[str] = []
+    iterator = iter(engine_argv)
+    for item in iterator:
+        option, separator, _ = item.partition("=")
+        option = option.replace("_", "-")
+        if len(option) > 2 and "--snapshot-dir".startswith(option):
+            if not separator:
+                next(iterator, None)
+        else:
+            child_argv.append(item)
+    if "--enable-sleep-mode" not in child_argv:
+        child_argv.append("--enable-sleep-mode")
+    return tuple(child_argv)
+
+
+def create_snapshot(
+    args: argparse.Namespace,
+    *,
+    engine_argv: tuple[str, ...] | None = None,
+    tools: LocalSnapshotTools | None = None,
+) -> None:
+    target = Path(args.snapshot_dir).absolute()
+    toolset = tools or LocalSnapshotTools()
+    toolset.preflight("create", target)
+    if target.exists() or target.is_symlink():
+        raise SnapshotCreateError(f"snapshot target already exists: {target}")
+    validate_artifact_root(target, creating=True)
+    target.parent.mkdir(mode=0o700, parents=False, exist_ok=True)
+    validate_artifact_root(target.parent, creating=False)
+    target.mkdir(mode=0o700)
+    validate_artifact_root(target, creating=False)
+    published = False
+    root_pid: int | None = None
+    try:
+        child_argv = _child_engine_argv(args, engine_argv)
+        root_pid = toolset.launch_child(target, child_argv)
+        oracle = toolset.wait_ready(target, root_pid)
+        inventory = toolset.inventory(root_pid, target)
+        toolset.dump(target, inventory)
+        toolset.verify_dead(inventory)
+        manifest = toolset.make_manifest(args, child_argv, inventory, oracle, target)
+        toolset.publish(target, manifest)
+        published = True
+    except BaseException:
+        if root_pid is not None:
+            toolset.abort_create(root_pid)
+        raise
+    finally:
+        # Runs even when abort_create raised, so no partial artifact is left.
+        if not published and target.exists():
+            shutil.rmtree(target)
+
+
+def restore_snapshot(
+    args: argparse.Namespace, *, tools: LocalSnapshotTools | None = None
+) -> None:
+    artifact = Path(args.snapshot_dir).absolute()
+    toolset = tools or LocalSnapshotTools()
+    toolset.preflight("restore", artifact)
+    manifest = read_manifest(artifact)
+    validate_identity(manifest, toolset.current_identity(manifest.gpu_uuid))
+
+    root_pid: int | None = None
+    try:
+        root_pid = toolset.restore(artifact, manifest)
+        toolset.release(artifact, args.host, args.port)
+        toolset.wait_listener(root_pid, args.host, args.port)
+        actual = toolset.request_oracle(args.host, args.port, manifest)
+        expected = Oracle(
+            token_ids=manifest.oracle_token_ids,
+            text=manifest.oracle_text,
+            sampled_token_logprob=manifest.oracle_sampled_token_logprob,
+        )
+        if not oracles_match(expected, actual):
+            raise SnapshotRestoreError(
+                f"snapshot oracle mismatch: expected={expected!r}, actual={actual!r}"
+            )
+        toolset.complete_restore(root_pid)
+    except BaseException as error:
+        if root_pid is not None:
+            try:
+                toolset.cleanup(root_pid)
+            except BaseException as cleanup_error:
+                raise SnapshotRestoreError(
+                    f"snapshot restore failed: {_error_detail(error)}; "
+                    f"cleanup failed: {_error_detail(cleanup_error)}"
+                ) from error
+        if isinstance(error, SnapshotRestoreError):
+            raise
+        raise SnapshotRestoreError(str(error)) from error
diff --git a/vllm/snapshot/manifest.py b/vllm/snapshot/manifest.py
new file mode 100644
index 000000000000..a029c544c52e
--- /dev/null
+++ b/vllm/snapshot/manifest.py
@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import math
+import os
+import stat
+from pathlib import Path
+from typing import Annotated, Any, Literal
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    StrictFloat,
+    StrictInt,
+    StrictStr,
+    ValidationError,
+    field_validator,
+    model_validator,
+)
+from typing_extensions import Self
+
+
+class SnapshotCompatibilityError(RuntimeError):
+    """The snapshot identity does not match the requested runtime."""
+
+
+class SnapshotSecurityError(RuntimeError):
+    """The snapshot artifact does not meet the private-path contract."""
+
+
+_MAX_PID = 2**31 - 1
+_PositivePid = Annotated[int, Field(gt=0, le=_MAX_PID)]
+_OracleTokenId = Annotated[int, Field(ge=0)]
+
+
+class SnapshotRuntimeIdentity(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+
+    vllm_version: str
+    python_version: str
+    torch_version: str
+    cuda_runtime: str
+    driver_version: str
+    criu_version: str
+    cuda_checkpoint_sha256: str
+    kernel_release: str
+    host_id: str
+    gpu_name: str
+    gpu_uuid: str
+    environment: tuple[tuple[str, str], ...]
+
+
+class SnapshotManifest(SnapshotRuntimeIdentity):
+    schema_version: Literal[1]
+    boundary: Literal["post-engine-init-reloadable-state-released"]
+    created_at: str
+    artifact_bytes: Annotated[int, Field(ge=0)]
+    model: str
+    served_model_name: str
+    model_revision: str
+    tokenizer_revision: str
+    engine_argv: tuple[str, ...]
+    process_tree: Annotated[tuple[_PositivePid, ...], Field(min_length=1)]
+    cuda_holders: Annotated[tuple[_PositivePid, ...], Field(min_length=1)]
+    external_cache_files: tuple[tuple[str, str], ...] = ()
+    oracle_token_ids: Annotated[
+        tuple[_OracleTokenId, ...], Field(min_length=1, max_length=1)
+    ]
+    oracle_text: str
+    oracle_sampled_token_logprob: float
+
+    @field_validator("schema_version", mode="before")
+    @classmethod
+    def _require_schema_version_one(cls, value: Any) -> Any:
+        if type(value) is not int or value != 1:
+            raise ValueError("schema version must be integer 1")
+        return value
+
+    @field_validator("oracle_sampled_token_logprob", mode="before")
+    @classmethod
+    def _require_finite_float(cls, value: Any) -> Any:
+        if type(value) is not float or not math.isfinite(value):
+            raise ValueError("sampled token log probability must be a finite float")
+        return value
+
+    @model_validator(mode="after")
+    def _validate_process_relationships(self) -> Self:
+        if len(set(self.process_tree)) != len(self.process_tree):
+            raise ValueError("process_tree contains duplicate PIDs")
+        if len(set(self.cuda_holders)) != len(self.cuda_holders):
+            raise ValueError("cuda_holders contains duplicate PIDs")
+        if not set(self.cuda_holders).issubset(self.process_tree):
+            raise ValueError("cuda_holders must be a subset of process_tree")
+        return self
+
+
+class ReleaseMarker(BaseModel):
+    """Release marker the controller writes; unknown keys are ignored."""
+
+    release: Literal[True]
+    host: StrictStr | None = None
+    port: Annotated[StrictInt, Field(ge=1, le=65535)]
+
+    @field_validator("release", mode="before")
+    @classmethod
+    def _require_json_true(cls, value: Any) -> Any:
+        # A literal accepts 1 as True; the barrier requires JSON true itself.
+        if value is not True:
+            raise ValueError("release must be JSON true")
+        return value
+
+
+class ReadyMarker(BaseModel):
+    """Ready marker the child writes; unknown keys are ignored."""
+
+    token_ids: tuple[StrictInt, ...]
+    text: StrictStr
+    sampled_token_logprob: StrictFloat | StrictInt
+
+
+def validate_identity(
+    expected: SnapshotRuntimeIdentity, actual: SnapshotRuntimeIdentity
+) -> None:
+    """Require an exact match for every field that can affect compatibility."""
+    # ``expected`` may be a manifest, so compare the identity fields only.
+    fields = set(SnapshotRuntimeIdentity.model_fields)
+    expected_values = expected.model_dump(include=fields)
+    actual_values = actual.model_dump(include=fields)
+    differing = [
+        name
+        for name in SnapshotRuntimeIdentity.model_fields
+        if expected_values[name] != actual_values[name]
+    ]
+    if differing:
+        raise SnapshotCompatibilityError(f"snapshot mismatch: {', '.join(differing)}")
+
+
+def _existing_path_chain(path: Path):
+    current = path
+    while True:
+        if current.exists() or current.is_symlink():
+            yield current
+        if current.parent == current:
+            break
+        current = current.parent
+
+
+def _is_trusted_sticky_directory(path_stat: os.stat_result) -> bool:
+    return (
+        path_stat.st_uid == 0
+        and stat.S_ISDIR(path_stat.st_mode)
+        and bool(path_stat.st_mode & stat.S_ISVTX)
+    )
+
+
+def validate_artifact_root(path: Path, *, creating: bool) -> None:
+    """Validate that an artifact path cannot be replaced by another user."""
+    path = Path(os.path.abspath(path))
+    if path.is_symlink():
+        raise SnapshotSecurityError(f"snapshot path is a symlink: {path}")
+    if not creating and not path.exists():
+        raise SnapshotSecurityError(f"snapshot path does not exist: {path}")
+
+    effective_uid = os.geteuid()
+    for component in _existing_path_chain(path):
+        component_stat = component.lstat()
+        if stat.S_ISLNK(component_stat.st_mode):
+            raise SnapshotSecurityError(
+                f"snapshot path contains a symlink: {component}"
+            )
+        if component_stat.st_uid not in (0, effective_uid):
+            raise SnapshotSecurityError(
+                f"snapshot path has an untrusted owner: {component}"
+            )
+        if component_stat.st_mode & (
+            stat.S_IWGRP | stat.S_IWOTH
+        ) and not _is_trusted_sticky_directory(component_stat):
+            raise SnapshotSecurityError(
+                f"snapshot path has a group- or world-writable ancestor: {component}"
+            )
+
+    if path.exists():
+        path_stat = path.lstat()
+        if not stat.S_ISDIR(path_stat.st_mode):
+            raise SnapshotSecurityError(f"snapshot path is not a directory: {path}")
+        if path_stat.st_uid != effective_uid:
+            raise SnapshotSecurityError(
+                f"snapshot directory is not owned by the current user: {path}"
+            )
+        if stat.S_IMODE(path_stat.st_mode) & 0o077:
+            raise SnapshotSecurityError(
+                f"snapshot directory must have mode 0700: {path}"
+            )
+
+
+def _validation_path(error: ValidationError) -> str:
+    detail = error.errors(include_url=False)[0]
+    location = detail["loc"]
+    if not location:
+        return "JSON" if detail["type"] == "json_invalid" else "root"
+    return ".".join(map(str, location))
+
+
+def _fsync_directory(path: Path) -> None:
+    directory_fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
+    try:
+        os.fsync(directory_fd)
+    finally:
+        os.close(directory_fd)
+
+
+def _write_json_atomic(path: Path, payload: object, *, overwrite: bool = False) -> None:
+    temporary = path.with_name(f"{path.name}.tmp")
+    if overwrite:
+        temporary.unlink(missing_ok=True)
+    elif path.exists() or path.is_symlink():
+        raise FileExistsError(path)
+    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
+    descriptor = os.open(temporary, flags, 0o600)
+    try:
+        with os.fdopen(descriptor, "wb") as output:
+            output.write(
+                json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
+            )
+            output.flush()
+            os.fsync(output.fileno())
+        os.replace(temporary, path)
+        _fsync_directory(path.parent)
+    except BaseException:
+        temporary.unlink(missing_ok=True)
+        raise
+
+
+def write_manifest_atomic(path: Path, manifest: SnapshotManifest) -> None:
+    """Write a new private manifest and make its directory entry durable."""
+    path = Path(path)
+    validate_artifact_root(path, creating=False)
+    destination = path / "manifest.json"
+    try:
+        _write_json_atomic(destination, manifest.model_dump(mode="json"))
+    except FileExistsError as error:
+        raise SnapshotSecurityError(
+            f"manifest already exists: {destination}"
+        ) from error
+
+
+def read_manifest(path: Path) -> SnapshotManifest:
+    """Read a snapshot published by atomically installing its manifest."""
+    path = Path(path)
+    validate_artifact_root(path, creating=False)
+    manifest_path = path / "manifest.json"
+    if manifest_path.is_symlink():
+        raise SnapshotSecurityError(f"manifest is a symlink: {manifest_path}")
+    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
+    try:
+        file_descriptor = os.open(manifest_path, flags)
+    except FileNotFoundError as error:
+        raise SnapshotSecurityError(
+            f"manifest does not exist: {manifest_path}"
+        ) from error
+    try:
+        with os.fdopen(file_descriptor, encoding="utf-8") as manifest_file:
+            payload = manifest_file.read()
+    except UnicodeDecodeError as error:
+        raise SnapshotCompatibilityError("invalid snapshot manifest: JSON") from error
+    try:
+        return SnapshotManifest.model_validate_json(payload, strict=True)
+    except ValidationError as error:
+        raise SnapshotCompatibilityError(
+            f"invalid snapshot manifest: {_validation_path(error)}"
+        ) from error
+
+
+def inspect_snapshot(path: Path) -> dict[str, object]:
+    return read_manifest(path).model_dump(mode="json")
diff --git a/vllm/snapshot/runtime.py b/vllm/snapshot/runtime.py
new file mode 100644
index 000000000000..582bc9a7d9cb
--- /dev/null
+++ b/vllm/snapshot/runtime.py
@@ -0,0 +1,1246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Host runtime operations for local snapshots."""
+
+import argparse
+import ast
+import hashlib
+import importlib.metadata
+import importlib.util
+import json
+import os
+import platform
+import select
+import shutil
+import signal
+import socket
+import stat
+import struct
+import subprocess
+import sys
+import tempfile
+import time
+import urllib.request
+from collections.abc import Iterable
+from contextlib import suppress
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import NoReturn
+
+import regex as re
+from pydantic import ValidationError
+
+from vllm.snapshot.manifest import (
+    ReadyMarker,
+    SnapshotManifest,
+    SnapshotRuntimeIdentity,
+    _fsync_directory,
+    _validation_path,
+    _write_json_atomic,
+    validate_artifact_root,
+    write_manifest_atomic,
+)
+from vllm.snapshot.types import Oracle
+
+
+class SnapshotCreateError(RuntimeError):
+    """Snapshot creation did not produce a complete artifact."""
+
+
+class SnapshotRestoreError(RuntimeError):
+    """Snapshot restore failed."""
+
+
+_COMMON_CRIU_OPTIONS = (
+    "--shell-job",
+    "--ext-unix-sk",
+    "--tcp-established",
+    "--link-remap",
+    "--file-locks",
+)
+_TCP_TABLES = (
+    ("AF_INET", Path("/proc/net/tcp")),
+    ("AF_INET6", Path("/proc/net/tcp6")),
+)
+
+_CACHE_REMEDY = "recreate the snapshot or restore with the same cache state"
+# Recorded instead of a digest for files whose content is volatile by design.
+_CACHE_LOCKED = "locked"
+# Recorded instead of a digest for a file capture found but could not read.
+_CACHE_UNREADABLE = "unreadable"
+
+
+def _generated_cache_prefixes() -> tuple[str, ...]:
+    """Path prefixes of the generated caches an engine process keeps open.
+
+    vLLM redirects Inductor and Triton under ``VLLM_CACHE_ROOT`` when it
+    compiles, but a process can still hold their upstream defaults open.
+    A cache root is resolved and separator-terminated so that it still matches
+    the kernel-canonical descriptor targets under a symlinked root, and so that
+    a sibling directory sharing its name does not. Inductor's default location
+    is a bare prefix instead, because its directory carries a per-user suffix.
+    """
+    import vllm.envs as envs
+
+    inductor_root = os.environ.get("TORCHINDUCTOR_CACHE_DIR")
+    roots = tuple(
+        f"{os.path.realpath(os.path.expanduser(root)).rstrip(os.sep)}{os.sep}"
+        for root in (
+            envs.VLLM_CACHE_ROOT,
+            os.environ.get("TRITON_CACHE_DIR") or "~/.triton/cache",
+            "~/.cache/flashinfer",
+            *((inductor_root,) if inductor_root else ()),
+        )
+    )
+    if inductor_root:
+        return roots
+    return roots + (
+        os.path.join(os.path.realpath(tempfile.gettempdir()), "torchinductor_"),
+    )
+
+
+def _error_detail(error: BaseException) -> str:
+    message = str(error)
+    return f"{type(error).__name__}: {message}" if message else type(error).__name__
+
+
+def _close_pidfds(handles: Iterable[tuple[int, int]]) -> None:
+    for _pid, pidfd in handles:
+        with suppress(OSError):
+            os.close(pidfd)
+
+
+@dataclass(frozen=True)
+class ProcessInventory:
+    root_pid: int
+    process_tree: tuple[int, ...]
+    cuda_holders: tuple[int, ...]
+    gpu_uuid: str
+    external_cache_files: tuple[tuple[str, str], ...] = ()
+
+
+@dataclass(frozen=True)
+class _TcpSocketRecord:
+    family: str
+    local_raw: str
+    remote_raw: str
+    inode: int
+
+
+def _process_comm(pid: int) -> str:
+    try:
+        return (Path("/proc") / str(pid) / "comm").read_text().strip()
+    except OSError:
+        return "unknown"
+
+
+def _decode_endpoint(family: str, raw: str) -> str:
+    """Render a ``/proc/net/tcp`` ``HEXADDR:HEXPORT`` field as ``ip:port``.
+
+    The kernel prints the address as little-endian 32-bit words: one for IPv4,
+    four in order for IPv6.
+    """
+    hex_address, _, hex_port = raw.partition(":")
+    port = int(hex_port, 16)
+    if family == "AF_INET":
+        packed = struct.pack("<I", int(hex_address, 16))
+        return f"{socket.inet_ntop(socket.AF_INET, packed)}:{port}"
+    packed = struct.pack(
+        "<4I", *(int(hex_address[index : index + 8], 16) for index in (0, 8, 16, 24))
+    )
+    return f"[{socket.inet_ntop(socket.AF_INET6, packed)}]:{port}"
+
+
+def _validate_tcp_connections(
+    records: tuple[_TcpSocketRecord, ...], socket_owners: dict[int, int]
+) -> None:
+    owned = tuple(record for record in records if record.inode in socket_owners)
+    endpoints = {
+        (record.family, record.local_raw, record.remote_raw) for record in owned
+    }
+    offenders = tuple(
+        record
+        for record in owned
+        if (record.family, record.remote_raw, record.local_raw) not in endpoints
+    )
+    if not offenders:
+        return
+    named = "; ".join(
+        f"pid {socket_owners[record.inode]} "
+        f"({_process_comm(socket_owners[record.inode])}) "
+        f"{_decode_endpoint(record.family, record.local_raw)} -> "
+        f"{_decode_endpoint(record.family, record.remote_raw)}"
+        for record in offenders[:5]
+    )
+    remaining = f" (+{len(offenders) - 5} more)" if len(offenders) > 5 else ""
+    raise SnapshotCreateError(
+        f"snapshot tree has an external established TCP connection: {named}{remaining}"
+    )
+
+
+class LocalSnapshotTools:
+    """Host-owned CRIU and cuda-checkpoint implementation."""
+
+    def __init__(self) -> None:
+        self.criu = shutil.which("criu") or "criu"
+        self.cuda_checkpoint = shutil.which("cuda-checkpoint") or "cuda-checkpoint"
+        self.nvidia_smi = shutil.which("nvidia-smi") or "nvidia-smi"
+        plugin_dir = os.environ.get("CRIU_CUDA_PLUGIN_DIR", "")
+        self.plugin_dir = Path(plugin_dir) if plugin_dir else None
+        self.shm_dir = Path("/dev/shm")
+        self.timeout_s = float(os.environ.get("VLLM_SNAPSHOT_TIMEOUT_S", "900"))
+        if not 0 < self.timeout_s < float("inf"):
+            raise ValueError("VLLM_SNAPSHOT_TIMEOUT_S must be a positive finite number")
+        self._children: dict[int, subprocess.Popen[bytes]] = {}
+        self._restored_processes: dict[int, tuple[tuple[int, int], ...]] = {}
+
+    def _privileged(self) -> list[str]:
+        return [] if os.geteuid() == 0 else ["sudo", "-n"]
+
+    def _run(
+        self,
+        command: list[str],
+        *,
+        timeout: float | None = None,
+    ) -> subprocess.CompletedProcess[str]:
+        return subprocess.run(
+            command,
+            check=True,
+            capture_output=True,
+            text=True,
+            timeout=timeout or self.timeout_s,
+        )
+
+    def preflight(self, action: str, artifact: Path) -> None:
+        if platform.system() != "Linux" or platform.machine() != "x86_64":
+            raise RuntimeError("snapshot requires Linux x86_64")
+        if action == "restore" and (
+            not callable(getattr(os, "pidfd_open", None))
+            or not callable(getattr(signal, "pidfd_send_signal", None))
+        ):
+            raise RuntimeError("snapshot requires Linux pidfd support")
+        for command in (self.criu, self.cuda_checkpoint, self.nvidia_smi):
+            if shutil.which(command) is None:
+                raise RuntimeError(f"snapshot dependency not found: {command}")
+        if self.plugin_dir is None:
+            raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required")
+        if not (self.plugin_dir / "cuda_plugin.so").is_file():
+            raise RuntimeError(f"CRIU CUDA plugin not found in {self.plugin_dir}")
+        if os.geteuid() != 0:
+            self._run(["sudo", "-n", "true"], timeout=5)
+        if action == "restore":
+            validate_artifact_root(artifact, creating=False)
+        else:
+            validate_artifact_root(artifact, creating=True)
+
+    def launch_child(
+        self,
+        workdir: Path,
+        engine_argv: tuple[str, ...],
+    ) -> int:
+        command = [
+            sys.executable,
+            "-m",
+            "vllm.snapshot.server",
+            "--ready-file",
+            str(workdir / "ready.json"),
+            "--release-file",
+            str(workdir / "release.json"),
+            "--release-timeout-s",
+            str(self.timeout_s),
+            "--",
+            *engine_argv,
+        ]
+        child_environment = os.environ.copy()
+        child_environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+        with (workdir / "child.log").open("wb") as log_file:
+            process = subprocess.Popen(
+                command,
+                env=child_environment,
+                stdin=subprocess.DEVNULL,
+                stdout=log_file,
+                stderr=subprocess.STDOUT,
+                start_new_session=True,
+            )
+        self._children[process.pid] = process
+        return process.pid
+
+    def wait_ready(self, workdir: Path, root_pid: int) -> Oracle:
+        ready_file = workdir / "ready.json"
+        deadline = time.monotonic() + self.timeout_s
+        waitid = os.waitid  # type: ignore[attr-defined]
+        while time.monotonic() < deadline:
+            if ready_file.is_file() and ready_file.stat().st_size:
+                try:
+                    marker = ReadyMarker.model_validate_json(ready_file.read_text())
+                except ValidationError as error:
+                    raise SnapshotCreateError(
+                        f"snapshot ready marker invalid: {_validation_path(error)}"
+                    ) from error
+                return Oracle(
+                    token_ids=marker.token_ids,
+                    text=marker.text,
+                    sampled_token_logprob=marker.sampled_token_logprob,
+                )
+            process = self._children.get(root_pid)
+            if process is not None and waitid(
+                os.P_PID,
+                root_pid,
+                os.WEXITED | os.WNOHANG | os.WNOWAIT,
+            ):
+                log = (workdir / "child.log").read_text(errors="replace")
+                raise SnapshotCreateError(f"snapshot child exited before ready:\n{log}")
+            try:
+                os.kill(root_pid, 0)
+            except ProcessLookupError as error:
+                log = (workdir / "child.log").read_text(errors="replace")
+                raise SnapshotCreateError(
+                    f"snapshot child exited before ready:\n{log}"
+                ) from error
+            time.sleep(0.05)
+        raise SnapshotCreateError("snapshot child did not become ready before timeout")
+
+    def _tree_pids(self, root_pid: int) -> tuple[int, ...]:
+        rows = self._run(["ps", "-eo", "pid=,ppid="]).stdout.splitlines()
+        children: dict[int, list[int]] = {}
+        for row in rows:
+            pid, parent_pid = map(int, row.split())
+            children.setdefault(parent_pid, []).append(pid)
+        stack = [root_pid]
+        found: list[int] = []
+        while stack:
+            pid = stack.pop()
+            if pid in found:
+                continue
+            found.append(pid)
+            stack.extend(children.get(pid, []))
+        return tuple(found)
+
+    def _cuda_process_rows(self) -> tuple[str, ...]:
+        output = self._run(
+            [
+                self.nvidia_smi,
+                "--query-compute-apps=pid,gpu_uuid",
+                "--format=csv,noheader,nounits",
+            ]
+        ).stdout
+        return tuple(line for line in output.splitlines() if line.strip())
+
+    def _cuda_pids(self, rows: tuple[str, ...]) -> tuple[int, ...]:
+        return tuple(sorted({int(line.split(",", 1)[0].strip()) for line in rows}))
+
+    def _descriptor_targets(self, pid: int) -> tuple[str, ...]:
+        descriptor_dir = Path("/proc") / str(pid) / "fd"
+        try:
+            descriptors = tuple(descriptor_dir.iterdir())
+        except FileNotFoundError as error:
+            raise SnapshotCreateError(
+                f"snapshot process exited during descriptor inventory: {pid}"
+            ) from error
+        targets: list[str] = []
+        for descriptor in descriptors:
+            try:
+                targets.append(os.readlink(descriptor))
+            except FileNotFoundError:
+                continue
+        return tuple(targets)
+
+    def _descriptor_inventory(
+        self, process_tree: tuple[int, ...]
+    ) -> tuple[tuple[int, ...], dict[int, int]]:
+        """Collect io_uring holders and each socket inode's owning pid."""
+        io_uring_pids: list[int] = []
+        socket_owners: dict[int, int] = {}
+        for pid in process_tree:
+            targets = self._descriptor_targets(pid)
+            if "anon_inode:[io_uring]" in targets:
+                io_uring_pids.append(pid)
+            for target in targets:
+                if target.startswith("socket:[") and target.endswith("]"):
+                    # Shared descriptors name one holder; any of them locates
+                    # the connection for the operator.
+                    socket_owners.setdefault(int(target[len("socket:[") : -1]), pid)
+        return tuple(io_uring_pids), socket_owners
+
+    def _tcp_records(self) -> tuple[_TcpSocketRecord, ...]:
+        records: list[_TcpSocketRecord] = []
+        for family, table in _TCP_TABLES:
+            if not table.exists():
+                continue
+            for row in table.read_text().splitlines()[1:]:
+                fields = row.split()
+                if fields[3] != "01":
+                    continue
+                records.append(
+                    _TcpSocketRecord(
+                        family=family,
+                        local_raw=fields[1],
+                        remote_raw=fields[2],
+                        inode=int(fields[9]),
+                    )
+                )
+        return tuple(records)
+
+    def _cache_fingerprint(self, path: Path) -> str | None:
+        """Digest a cache file, or mark the ones whose content is volatile."""
+        try:
+            if not stat.S_ISREG(path.stat().st_mode):
+                return None
+        except OSError:
+            return None
+        if path.name.endswith(".lock"):
+            return _CACHE_LOCKED
+        try:
+            return self._sha256(path)
+        except OSError:
+            # Restore can only check that this one still exists, so record the
+            # downgrade rather than passing it off as a volatile file.
+            return _CACHE_UNREADABLE
+
+    def _external_cache_files(
+        self, process_tree: tuple[int, ...], artifact: Path
+    ) -> tuple[tuple[str, str], ...]:
+        """Fingerprint the generated-cache files the captured tree holds open.
+
+        CRIU reopens these by path on restore, so a JIT cache that rotates
+        between capture and restore fails the restore permanently. Recording
+        them lets restore fail early and name the file that changed. Files
+        inside the artifact are excluded: they belong to the snapshot itself.
+        Deleted targets are excluded too, because CRIU remaps those links.
+        """
+        prefixes = _generated_cache_prefixes()
+        artifact_prefix = f"{os.path.abspath(artifact)}{os.sep}"
+        candidates: set[str] = set()
+        for pid in process_tree:
+            for target in self._descriptor_targets(pid):
+                if target.endswith(" (deleted)") or target.startswith(artifact_prefix):
+                    continue
+                if target.endswith(".lock") or target.startswith(prefixes):
+                    candidates.add(target)
+        recorded = (
+            (path, self._cache_fingerprint(Path(path))) for path in sorted(candidates)
+        )
+        return tuple(
+            (path, fingerprint)
+            for path, fingerprint in recorded
+            if fingerprint is not None
+        )
+
+    def inventory(self, root_pid: int, artifact: Path) -> ProcessInventory:
+        process_tree = self._tree_pids(root_pid)
+        cuda_rows = self._cuda_process_rows()
+        cuda_holders = tuple(
+            pid for pid in self._cuda_pids(cuda_rows) if pid in process_tree
+        )
+        if not cuda_holders:
+            raise SnapshotCreateError("snapshot tree has no CUDA-holding process")
+        io_uring_pids, socket_owners = self._descriptor_inventory(process_tree)
+        if io_uring_pids:
+            raise SnapshotCreateError(
+                "snapshot tree owns io_uring state that CRIU cannot dump; set "
+                "kernel.io_uring_disabled=1 before starting an unprivileged "
+                "vLLM process, or use =2 to disable it host-wide "
+                f"(pids: {', '.join(map(str, io_uring_pids))})"
+            )
+        _validate_tcp_connections(self._tcp_records(), socket_owners)
+        gpu_uuid = self._gpu_uuid_for_pids(cuda_holders, cuda_rows)
+        return ProcessInventory(
+            root_pid=root_pid,
+            process_tree=process_tree,
+            cuda_holders=cuda_holders,
+            gpu_uuid=gpu_uuid,
+            external_cache_files=self._external_cache_files(process_tree, artifact),
+        )
+
+    def _gpu_uuid_for_pids(self, pids: tuple[int, ...], rows: tuple[str, ...]) -> str:
+        wanted = set(pids)
+        matches: set[str] = set()
+        try:
+            for line in rows:
+                pid, gpu_uuid = (item.strip() for item in line.split(",", 1))
+                if int(pid) in wanted:
+                    matches.add(gpu_uuid)
+        except (TypeError, ValueError) as error:
+            raise SnapshotCreateError(
+                "could not identify the GPU used by the snapshot process tree"
+            ) from error
+        if len(matches) != 1:
+            raise SnapshotCreateError(
+                "snapshot process tree must use exactly one identifiable GPU"
+            )
+        return matches.pop()
+
+    def _criu(
+        self,
+        action: str,
+        artifact: Path,
+        arguments: list[str],
+    ) -> None:
+        if self.plugin_dir is None:
+            raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required")
+        env = os.environ.copy()
+        env["PATH"] = f"{Path(self.cuda_checkpoint).parent}:{env['PATH']}"
+        env["VLLM_SNAPSHOT_IGNORED_FD_MAP"] = str(artifact / "ignored-nvidia-fds.tsv")
+        env["RIG_CRIU_IGNORED_FD_MAP"] = env["VLLM_SNAPSHOT_IGNORED_FD_MAP"]
+        command = [
+            *self._privileged(),
+            "env",
+            f"PATH={env['PATH']}",
+            f"VLLM_SNAPSHOT_IGNORED_FD_MAP={env['VLLM_SNAPSHOT_IGNORED_FD_MAP']}",
+            f"RIG_CRIU_IGNORED_FD_MAP={env['RIG_CRIU_IGNORED_FD_MAP']}",
+            self.criu,
+            action,
+            "-L",
+            str(self.plugin_dir),
+            *arguments,
+        ]
+        self._run(command)
+
+    def _record_child_log_size(self, artifact: Path) -> None:
+        child_log = artifact / "child.log"
+        size_path = artifact / "child.log.snapshot-size"
+        payload = f"{child_log.stat().st_size}\n".encode()
+        descriptor = os.open(
+            size_path,
+            os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0),
+            0o600,
+        )
+        with os.fdopen(descriptor, "wb") as output:
+            output.write(payload)
+            output.flush()
+            os.fsync(output.fileno())
+
+    def _link_remap_names(self) -> set[str]:
+        return {
+            path.name
+            for path in self.shm_dir.glob("link_remap.*")
+            if path.name.removeprefix("link_remap.").isdigit()
+        }
+
+    def _capture_link_remaps(self, artifact: Path, names: set[str]) -> None:
+        if not names:
+            return
+        destination_dir = artifact / "link-remaps"
+        destination_dir.mkdir(mode=0o700)
+        for name in sorted(names):
+            source = self.shm_dir / name
+            source_fd = os.open(
+                source,
+                os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
+            )
+            try:
+                metadata = os.fstat(source_fd)
+                if not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != os.getuid():
+                    raise SnapshotCreateError(f"invalid CRIU link remap: {source}")
+                destination_fd = os.open(
+                    destination_dir / name,
+                    os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0),
+                    0o600,
+                )
+                with os.fdopen(destination_fd, "wb") as destination:
+                    with os.fdopen(os.dup(source_fd), "rb") as contents:
+                        shutil.copyfileobj(contents, destination)
+                    destination.flush()
+                    os.fsync(destination.fileno())
+            finally:
+                os.close(source_fd)
+            source.unlink()
+
+    def _rollback_link_remaps(self, created: dict[str, tuple[int, int]]) -> None:
+        for name, identity in created.items():
+            target = self.shm_dir / name
+            try:
+                metadata = target.lstat()
+            except FileNotFoundError:
+                continue
+            if (metadata.st_dev, metadata.st_ino) == identity:
+                target.unlink()
+
+    def _restore_link_remaps(self, artifact: Path) -> dict[str, tuple[int, int]]:
+        source_dir = artifact / "link-remaps"
+        if not source_dir.exists():
+            return {}
+        created: dict[str, tuple[int, int]] = {}
+        try:
+            for source in sorted(source_dir.iterdir()):
+                if (
+                    not source.name.removeprefix("link_remap.").isdigit()
+                    or source.is_symlink()
+                    or not source.is_file()
+                ):
+                    raise SnapshotRestoreError(f"invalid saved link remap: {source}")
+                payload = source.read_bytes()
+                target = self.shm_dir / source.name
+                try:
+                    descriptor = os.open(
+                        target,
+                        os.O_WRONLY
+                        | os.O_CREAT
+                        | os.O_EXCL
+                        | getattr(os, "O_NOFOLLOW", 0),
+                        0o600,
+                    )
+                except FileExistsError as error:
+                    try:
+                        descriptor = os.open(
+                            target,
+                            os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
+                        )
+                    except OSError:
+                        raise SnapshotRestoreError(
+                            f"conflicting CRIU link remap exists: {target}"
+                        ) from error
+                    with os.fdopen(descriptor, "rb") as existing:
+                        metadata = os.fstat(existing.fileno())
+                        if (
+                            not stat.S_ISREG(metadata.st_mode)
+                            or metadata.st_uid != os.getuid()
+                            or existing.read() != payload
+                        ):
+                            raise SnapshotRestoreError(
+                                f"conflicting CRIU link remap exists: {target}"
+                            ) from error
+                    continue
+                metadata = os.fstat(descriptor)
+                created[source.name] = (metadata.st_dev, metadata.st_ino)
+                with os.fdopen(descriptor, "wb") as output:
+                    output.write(payload)
+                    output.flush()
+                    os.fsync(output.fileno())
+        except BaseException:
+            self._rollback_link_remaps(created)
+            raise
+        return created
+
+    def _reset_child_log(self, artifact: Path) -> None:
+        size_path = artifact / "child.log.snapshot-size"
+        descriptor = os.open(
+            size_path,
+            os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
+        )
+        with os.fdopen(descriptor, "r") as source:
+            size = int(source.read().strip())
+        if size < 0:
+            raise SnapshotRestoreError("captured child log size is invalid")
+
+        child_log = artifact / "child.log"
+        descriptor = os.open(
+            child_log,
+            os.O_WRONLY | getattr(os, "O_NOFOLLOW", 0),
+        )
+        try:
+            metadata = os.fstat(descriptor)
+            if not stat.S_ISREG(metadata.st_mode) or metadata.st_size < size:
+                raise SnapshotRestoreError("captured child log is invalid")
+            os.ftruncate(descriptor, size)
+            os.fsync(descriptor)
+        finally:
+            os.close(descriptor)
+
+    def dump(self, workdir: Path, inventory: ProcessInventory) -> None:
+        images = workdir / "images"
+        images.mkdir(mode=0o700)
+        fd_map = workdir / "ignored-nvidia-fds.tsv"
+        fd_map.touch(mode=0o600)
+        remaps_before = self._link_remap_names()
+        try:
+            self._criu(
+                "dump",
+                workdir,
+                [
+                    "--tree",
+                    str(inventory.root_pid),
+                    "--images-dir",
+                    str(images),
+                    "--log-file",
+                    "dump.log",
+                    *_COMMON_CRIU_OPTIONS,
+                ],
+            )
+            self._capture_link_remaps(
+                workdir,
+                self._link_remap_names() - remaps_before,
+            )
+        except BaseException:
+            for name in self._link_remap_names() - remaps_before:
+                (self.shm_dir / name).unlink(missing_ok=True)
+            raise
+        self._record_child_log_size(workdir)
+
+    def verify_dead(self, inventory: ProcessInventory) -> None:
+        process = self._children.get(inventory.root_pid)
+        if process is not None:
+            process.wait(timeout=10)
+        for pid in inventory.process_tree:
+            try:
+                os.kill(pid, 0)
+            except ProcessLookupError:
+                continue
+            raise SnapshotCreateError(f"process survived CRIU dump: {pid}")
+        self._children.pop(inventory.root_pid, None)
+
+    def _version(self, command: list[str]) -> str:
+        try:
+            result = self._run(command, timeout=10)
+        except subprocess.CalledProcessError as error:
+            return (error.stdout or error.stderr).splitlines()[0]
+        return result.stdout.splitlines()[0] if result.stdout else "unknown"
+
+    def _sha256(self, path: Path) -> str:
+        digest = hashlib.sha256()
+        with path.open("rb") as source:
+            for chunk in iter(lambda: source.read(1024 * 1024), b""):
+                digest.update(chunk)
+        return digest.hexdigest()
+
+    def _plugin_identity(self) -> tuple[tuple[str, str], ...]:
+        if self.plugin_dir is None:
+            return ()
+        return tuple(
+            (f"snapshot_plugin_sha256:{path.name}", self._sha256(path))
+            for path in sorted(self.plugin_dir.glob("*.so"))
+        )
+
+    def _torch_identity(self) -> tuple[str, str]:
+        torch_version = importlib.metadata.version("torch")
+        torch_spec = importlib.util.find_spec("torch")
+        locations = torch_spec and torch_spec.submodule_search_locations
+        if not locations:
+            raise RuntimeError("installed torch package could not be located")
+        version_path = Path(next(iter(locations))) / "version.py"
+        # Parse the version file rather than importing torch: snapshot create's
+        # CLI path stays import-light (vllm/entrypoints/cli/snapshot.py).
+        syntax = ast.parse(version_path.read_text(), filename=str(version_path))
+        cuda_runtime = ""
+        for statement in syntax.body:
+            target = None
+            value = None
+            if isinstance(statement, ast.Assign) and len(statement.targets) == 1:
+                target = statement.targets[0]
+                value = statement.value
+            elif isinstance(statement, ast.AnnAssign):
+                target = statement.target
+                value = statement.value
+            if isinstance(target, ast.Name) and target.id == "cuda" and value:
+                parsed = ast.literal_eval(value)
+                cuda_runtime = str(parsed or "")
+                break
+        return torch_version, cuda_runtime
+
+    def _gpu_identity(self, gpu_uuid: str) -> tuple[str, str, str]:
+        output = self._run(
+            [
+                self.nvidia_smi,
+                "--query-gpu=name,uuid,driver_version",
+                "--format=csv,noheader,nounits",
+                f"--id={gpu_uuid}",
+            ]
+        ).stdout.strip()
+        name, uuid, driver = (item.strip() for item in output.split(",", 2))
+        return name, uuid, driver
+
+    def _artifact_bytes(self, workdir: Path) -> int:
+        return sum(path.stat().st_size for path in workdir.rglob("*") if path.is_file())
+
+    def _environment_identity(self) -> tuple[tuple[str, str], ...]:
+        prefixes = ("VLLM_", "CUDA_", "NCCL_", "TORCH_", "TRITON_")
+        environment = os.environ.copy()
+        environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+        selected = tuple(
+            sorted(
+                (key, hashlib.sha256(key.encode() + b"\0" + value.encode()).hexdigest())
+                for key, value in environment.items()
+                if key.startswith(prefixes)
+                and key not in {"VLLM_API_KEY", "VLLM_SNAPSHOT_TIMEOUT_S"}
+            )
+        )
+        return selected + self._plugin_identity()
+
+    def make_manifest(
+        self,
+        args: argparse.Namespace,
+        engine_argv: tuple[str, ...],
+        inventory: ProcessInventory,
+        oracle: Oracle,
+        workdir: Path,
+    ) -> SnapshotManifest:
+        identity = self.current_identity(inventory.gpu_uuid)
+        redacted_argv = list(engine_argv)
+        redact_api_keys = False
+        for index, argument in enumerate(redacted_argv):
+            option, separator, _ = argument.partition("=")
+            normalized_option = option.replace("_", "-")
+            # Parsed argv can contain unambiguous long-option abbreviations.
+            if len(normalized_option) > 2 and "--api-key".startswith(normalized_option):
+                redact_api_keys = not separator
+                if separator:
+                    redacted_argv[index] = f"{option}=***"
+            elif len(normalized_option) > 2 and "--hf-token".startswith(
+                normalized_option
+            ):
+                if separator:
+                    redacted_argv[index] = f"{option}=***"
+                elif index + 1 < len(redacted_argv):
+                    redacted_argv[index + 1] = "***"
+                redact_api_keys = False
+            elif argument.startswith("-"):
+                redact_api_keys = False
+            elif redact_api_keys:
+                redacted_argv[index] = "***"
+        engine_argv = tuple(redacted_argv)
+        revision = str(getattr(args, "revision", None) or "")
+        tokenizer_revision = str(getattr(args, "tokenizer_revision", None) or revision)
+        source_model = str(getattr(args, "model_tag", None) or args.model)
+        from vllm.config.model import get_served_model_name
+
+        served_model_name = get_served_model_name(
+            source_model, getattr(args, "served_model_name", None)
+        )
+        return SnapshotManifest(
+            schema_version=1,
+            boundary="post-engine-init-reloadable-state-released",
+            created_at=datetime.now(timezone.utc).isoformat(),
+            artifact_bytes=self._artifact_bytes(workdir),
+            **identity.model_dump(),
+            model=source_model,
+            served_model_name=str(served_model_name),
+            model_revision=revision,
+            tokenizer_revision=tokenizer_revision,
+            engine_argv=engine_argv,
+            process_tree=inventory.process_tree,
+            cuda_holders=inventory.cuda_holders,
+            external_cache_files=inventory.external_cache_files,
+            oracle_token_ids=oracle.token_ids,
+            oracle_text=oracle.text,
+            oracle_sampled_token_logprob=oracle.sampled_token_logprob,
+        )
+
+    def publish(self, workdir: Path, manifest: SnapshotManifest) -> None:
+        write_manifest_atomic(workdir, manifest)
+        _fsync_directory(workdir.parent)
+
+    def current_identity(self, gpu_uuid: str) -> SnapshotRuntimeIdentity:
+        gpu_name, gpu_uuid, driver_version = self._gpu_identity(gpu_uuid)
+        torch_version, cuda_runtime = self._torch_identity()
+        host_id = Path("/etc/machine-id").read_text().strip()
+        return SnapshotRuntimeIdentity(
+            vllm_version=importlib.metadata.version("vllm"),
+            python_version=platform.python_version(),
+            torch_version=torch_version,
+            cuda_runtime=cuda_runtime,
+            driver_version=driver_version,
+            criu_version=self._version([self.criu, "--version"]),
+            cuda_checkpoint_sha256=self._sha256(Path(self.cuda_checkpoint)),
+            kernel_release=platform.release(),
+            host_id=host_id,
+            gpu_name=gpu_name,
+            gpu_uuid=gpu_uuid,
+            environment=self._environment_identity(),
+        )
+
+    def _read_restored_pid(self, pidfile: Path) -> int:
+        try:
+            payload = self._run([*self._privileged(), "cat", str(pidfile)]).stdout
+        except Exception as error:
+            raise SnapshotRestoreError("restored PID file is missing") from error
+        match = re.fullmatch(r"([1-9][0-9]{0,9})\n?", payload)
+        if match is None:
+            raise SnapshotRestoreError("restored PID file is invalid")
+        restored_pid = int(match.group(1))
+        if restored_pid > 2**31 - 1:
+            raise SnapshotRestoreError("restored PID file is invalid")
+        return restored_pid
+
+    def _process_state(self, pid: int) -> tuple[int, int, int, int]:
+        payload = (Path("/proc") / str(pid) / "stat").read_text()
+        command_end = payload.rfind(")")
+        fields = payload[command_end + 2 :].split()
+        if command_end < 0 or len(fields) < 20:
+            raise ValueError(f"invalid process state for PID {pid}")
+        return (int(fields[1]), int(fields[2]), int(fields[3]), int(fields[19]))
+
+    def _process_states(self) -> dict[int, tuple[int, int, int, int]]:
+        states: dict[int, tuple[int, int, int, int]] = {}
+        for entry in Path("/proc").iterdir():
+            if not entry.name.isdigit():
+                continue
+            pid = int(entry.name)
+            try:
+                states[pid] = self._process_state(pid)
+            except (OSError, ValueError):
+                continue
+        return states
+
+    def _process_command(self, pid: int) -> tuple[str, ...]:
+        payload = (Path("/proc") / str(pid) / "cmdline").read_bytes()
+        return tuple(
+            os.fsdecode(argument) for argument in payload.split(b"\0") if argument
+        )
+
+    def _pidfd_open(self, pid: int) -> int:
+        pidfd_open = getattr(os, "pidfd_open", None)
+        if not callable(pidfd_open):
+            raise SnapshotRestoreError("snapshot requires Linux pidfd support")
+        return pidfd_open(pid)
+
+    def _pidfd_exited(self, pidfd: int) -> bool:
+        poller = select.poll()
+        poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR)
+        return any(self._pidfd_event_is_exit(event) for _fd, event in poller.poll(0))
+
+    def _pidfd_event_is_exit(self, event: int) -> bool:
+        if event & (select.POLLERR | select.POLLNVAL):
+            raise SnapshotRestoreError("restored process pidfd poll failed")
+        return bool(event & (select.POLLIN | select.POLLHUP))
+
+    def _pin_restored_tree(
+        self,
+        artifact: Path,
+        process_tree: tuple[int, ...],
+        cuda_holders: tuple[int, ...],
+    ) -> None:
+        root_pid = process_tree[0]
+        expected_pids = set(process_tree)
+        expected_cuda_holders = set(cuda_holders)
+        if not expected_cuda_holders or not expected_cuda_holders.issubset(
+            expected_pids
+        ):
+            raise SnapshotRestoreError(
+                "captured CUDA holders are outside the captured process tree"
+            )
+        if root_pid in self._restored_processes:
+            raise SnapshotRestoreError("restored PID already has a pinned transaction")
+        handles: list[tuple[int, int]] = []
+        try:
+            for pid in process_tree:
+                handles.append((pid, self._pidfd_open(pid)))
+            handles_by_pid = dict(handles)
+            if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles):
+                raise SnapshotRestoreError(
+                    "restored process exited during verification"
+                )
+
+            root_state = self._process_state(root_pid)
+            if root_state[1:3] != (root_pid, root_pid):
+                raise SnapshotRestoreError(
+                    "restored root is not the session and process-group leader"
+                )
+            command = self._process_command(root_pid)
+            release_file = str(artifact / "release.json")
+            if not any(
+                command[index : index + 2] == ("-m", "vllm.snapshot.server")
+                for index in range(len(command) - 1)
+            ) or not any(
+                command[index : index + 2] == ("--release-file", release_file)
+                for index in range(len(command) - 1)
+            ):
+                raise SnapshotRestoreError(
+                    "restored root command does not match the snapshot barrier"
+                )
+            if self._process_state(root_pid) != root_state:
+                raise SnapshotRestoreError("restored PID changed during verification")
+            if self._pidfd_exited(handles[0][1]):
+                raise SnapshotRestoreError("restored root exited during verification")
+
+            first_states = self._process_states()
+            live_session = {
+                pid for pid, state in first_states.items() if state[2] == root_pid
+            }
+            if not expected_cuda_holders.issubset(live_session):
+                raise SnapshotRestoreError(
+                    "restored session is missing captured CUDA holders"
+                )
+            if live_session != expected_pids:
+                raise SnapshotRestoreError(
+                    "restored session does not match the captured process tree"
+                )
+            if first_states[root_pid] != root_state:
+                raise SnapshotRestoreError("restored PID changed during verification")
+            for pid in process_tree[1:]:
+                state = first_states[pid]
+                pidfd = handles_by_pid[pid]
+                if self._process_state(pid) != state or self._pidfd_exited(pidfd):
+                    raise SnapshotRestoreError(
+                        "restored process tree changed during verification"
+                    )
+
+            second_states = self._process_states()
+            second_session = {
+                pid for pid, state in second_states.items() if state[2] == root_pid
+            }
+            if second_session != expected_pids or any(
+                second_states[pid] != first_states[pid] for pid in process_tree
+            ):
+                raise SnapshotRestoreError(
+                    "restored process tree changed during verification"
+                )
+            if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles):
+                raise SnapshotRestoreError(
+                    "restored process exited during verification"
+                )
+        except BaseException as error:
+            _close_pidfds(handles)
+            if isinstance(error, SnapshotRestoreError):
+                raise
+            if not isinstance(error, Exception):
+                raise
+            raise SnapshotRestoreError(
+                "restored process tree could not be pinned"
+            ) from error
+        self._restored_processes[root_pid] = tuple(handles)
+
+    def _abort_failed_restore(
+        self,
+        artifact: Path,
+        root_pid: int,
+        staged_remaps: dict[str, tuple[int, int]],
+        primary: BaseException,
+    ) -> NoReturn:
+        secondary: list[str] = []
+        try:
+            _write_json_atomic(
+                artifact / "release.json", {"release": False}, overwrite=True
+            )
+        except BaseException as error:
+            secondary.append(f"abort marker failed: {_error_detail(error)}")
+        # The marker above is only a request. Terminate and wait for the tree we
+        # own, before releasing the shared state it may still be using. A tree
+        # that could not be pinned is not ours to signal.
+        if root_pid in self._restored_processes:
+            try:
+                self.cleanup(root_pid)
+            except BaseException as error:
+                secondary.append(f"restore cleanup failed: {_error_detail(error)}")
+        try:
+            self._rollback_link_remaps(staged_remaps)
+        except BaseException as error:
+            secondary.append(f"remap rollback failed: {_error_detail(error)}")
+        if secondary:
+            raise SnapshotRestoreError(
+                f"snapshot restore failed: {_error_detail(primary)}; "
+                + "; ".join(secondary)
+            ) from primary
+        raise primary
+
+    def _verify_external_cache_files(self, manifest: SnapshotManifest) -> None:
+        """Reject a rotated generated cache before CRIU reopens it by path.
+
+        Like the PID probe above this reports a stale host early rather than
+        guarding it: CRIU stays authoritative for a cache replaced afterwards.
+        """
+        for path, fingerprint in manifest.external_cache_files:
+            target = Path(path)
+            if not target.is_file():
+                raise SnapshotRestoreError(
+                    f"external cache file vanished since capture: {path}; "
+                    f"{_CACHE_REMEDY}"
+                )
+            if fingerprint in (_CACHE_LOCKED, _CACHE_UNREADABLE):
+                continue
+            try:
+                current = self._sha256(target)
+            except OSError as error:
+                raise SnapshotRestoreError(
+                    f"external cache file is unreadable since capture: {path}; "
+                    f"{_CACHE_REMEDY}"
+                ) from error
+            if current != fingerprint:
+                raise SnapshotRestoreError(
+                    f"external cache file changed since capture: {path}; "
+                    f"{_CACHE_REMEDY}"
+                )
+
+    def restore(self, artifact: Path, manifest: SnapshotManifest) -> int:
+        # This is necessarily a best-effort TOCTOU check. CRIU restore and the
+        # rollback below remain authoritative if a PID is claimed after it.
+        for pid in manifest.process_tree:
+            try:
+                os.kill(pid, 0)
+            except ProcessLookupError:
+                continue
+            except PermissionError:
+                pass
+            except Exception as error:
+                raise SnapshotRestoreError(
+                    f"captured PID availability probe failed: {pid}"
+                ) from error
+            raise SnapshotRestoreError(f"captured PID is already occupied: {pid}")
+        self._verify_external_cache_files(manifest)
+        expected_root_pid = manifest.process_tree[0]
+        (artifact / "release.json").unlink(missing_ok=True)
+        pidfile = artifact / "restored.pid"
+        pidfile.unlink(missing_ok=True)
+        self._reset_child_log(artifact)
+        staged_remaps = self._restore_link_remaps(artifact)
+        try:
+            self._criu(
+                "restore",
+                artifact,
+                [
+                    "--images-dir",
+                    str(artifact / "images"),
+                    "--log-file",
+                    "restore.log",
+                    *_COMMON_CRIU_OPTIONS,
+                    "--restore-detached",
+                    "--pidfile",
+                    str(pidfile),
+                ],
+            )
+            # Pin before reading the PID file so that a failure between restore
+            # and the read cannot orphan the tree.
+            self._pin_restored_tree(
+                artifact, manifest.process_tree, manifest.cuda_holders
+            )
+            restored_pid = self._read_restored_pid(pidfile)
+            if restored_pid != expected_root_pid:
+                raise SnapshotRestoreError(
+                    "restored PID does not match the captured process tree"
+                )
+        except BaseException as error:
+            self._abort_failed_restore(
+                artifact, expected_root_pid, staged_remaps, error
+            )
+        return restored_pid
+
+    def release(self, artifact: Path, host: str | None, port: int) -> None:
+        if not 1 <= port <= 65535:
+            raise SnapshotRestoreError(f"snapshot restore port is invalid: {port}")
+        family = socket.AF_INET6 if host and ":" in host else socket.AF_INET
+        with socket.socket(family=family, type=socket.SOCK_STREAM) as listener_probe:
+            listener_probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            try:
+                listener_probe.bind((host or "", port))
+            except OSError as error:
+                address = f"{host or '0.0.0.0'}:{port}"
+                raise SnapshotRestoreError(
+                    f"snapshot restore address is already in use: {address}"
+                ) from error
+
+        _write_json_atomic(
+            artifact / "release.json",
+            {"release": True, "host": host, "port": port},
+            overwrite=True,
+        )
+
+    def _connect_host(self, host: str | None) -> str:
+        if not host or host == "0.0.0.0":
+            return "127.0.0.1"
+        if host == "::":
+            return "::1"
+        return host
+
+    def wait_listener(self, root_pid: int, host: str | None, port: int) -> None:
+        deadline = time.monotonic() + self.timeout_s
+        connect_host = self._connect_host(host)
+        while time.monotonic() < deadline:
+            try:
+                with socket.create_connection((connect_host, port), timeout=1):
+                    return
+            except OSError as error:
+                handles = self._restored_processes.get(root_pid)
+                if handles is None or self._pidfd_exited(handles[0][1]):
+                    raise SnapshotRestoreError(
+                        "restored process exited before binding HTTP"
+                    ) from error
+                time.sleep(0.1)
+        raise SnapshotRestoreError("restored HTTP listener did not become ready")
+
+    def request_oracle(
+        self, host: str | None, port: int, manifest: SnapshotManifest
+    ) -> Oracle:
+        connect_host = self._connect_host(host)
+        url_host = f"[{connect_host}]" if ":" in connect_host else connect_host
+        payload = json.dumps(
+            {
+                "model": manifest.served_model_name,
+                "prompt": "The capital of France is",
+                "min_tokens": 1,
+                "max_tokens": 1,
+                "temperature": 0,
+                "seed": 0,
+                "logprobs": 0,
+                "return_token_ids": True,
+            }
+        ).encode()
+        request = urllib.request.Request(
+            f"http://{url_host}:{port}/v1/completions",
+            data=payload,
+            headers={"Content-Type": "application/json"},
+        )
+        with urllib.request.urlopen(request, timeout=120) as response:
+            body = json.load(response)
+        choice = body["choices"][0]
+        try:
+            return Oracle(
+                token_ids=tuple(choice["token_ids"]),
+                text=choice["text"],
+                sampled_token_logprob=choice["logprobs"]["token_logprobs"][0],
+            )
+        except ValueError as error:
+            raise SnapshotRestoreError(
+                "snapshot HTTP canary must return exactly one finite token logprob"
+            ) from error
+        except (IndexError, KeyError, TypeError) as error:
+            raise SnapshotRestoreError(
+                "snapshot HTTP canary response is missing sampled token logprob"
+            ) from error
+
+    def cleanup(self, root_pid: int) -> None:
+        handles = self._restored_processes.get(root_pid)
+        if handles is None:
+            raise SnapshotRestoreError("no pinned restore transaction for cleanup")
+        pidfd_send_signal = getattr(signal, "pidfd_send_signal", None)
+        if not callable(pidfd_send_signal):
+            raise SnapshotRestoreError("restored process cleanup is incomplete")
+        try:
+            for _pid, pidfd in handles:
+                with suppress(ProcessLookupError):
+                    pidfd_send_signal(pidfd, signal.SIGTERM, None, 0)
+            time.sleep(min(0.5, self.timeout_s))
+            for _pid, pidfd in handles:
+                with suppress(ProcessLookupError):
+                    pidfd_send_signal(pidfd, signal.SIGKILL, None, 0)
+        except BaseException as error:
+            raise SnapshotRestoreError(
+                f"restored process cleanup is incomplete: {_error_detail(error)}"
+            ) from error
+
+        poller = select.poll()
+        for _pid, pidfd in handles:
+            poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR)
+        pending = {pidfd for _pid, pidfd in handles}
+        deadline = time.monotonic() + self.timeout_s
+        while pending:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise SnapshotRestoreError("restored process cleanup is incomplete")
+            for pidfd, event in poller.poll(max(1, int(remaining * 1000))):
+                if not self._pidfd_event_is_exit(event):
+                    continue
+                pending.discard(pidfd)
+                with suppress(KeyError, OSError):
+                    poller.unregister(pidfd)
+        self._restored_processes.pop(root_pid)
+        _close_pidfds(handles)
+
+    def complete_restore(self, root_pid: int) -> None:
+        handles = self._restored_processes.pop(root_pid, None)
+        if handles is None:
+            raise SnapshotRestoreError("no pinned restore transaction to complete")
+        _close_pidfds(handles)
+
+    def abort_create(self, root_pid: int) -> None:
+        process = self._children.get(root_pid)
+        if process is None:
+            return
+        with suppress(ProcessLookupError):
+            os.killpg(root_pid, signal.SIGKILL)
+        try:
+            process.wait(timeout=10)
+        except subprocess.TimeoutExpired as error:
+            raise SnapshotCreateError(
+                f"snapshot child process group {root_pid} survived SIGKILL"
+            ) from error
+        self._children.pop(root_pid, None)
diff --git a/vllm/snapshot/server.py b/vllm/snapshot/server.py
new file mode 100644
index 000000000000..2dc0f3384aeb
--- /dev/null
+++ b/vllm/snapshot/server.py
@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import asyncio
+import json
+import os
+import stat
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from pydantic import ValidationError
+
+from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic
+from vllm.snapshot.types import Oracle, oracles_match
+
+_CANARY_PROMPT = "The capital of France is"
+
+
+class SnapshotCanaryError(RuntimeError):
+    """The initialized engine did not produce a valid snapshot oracle."""
+
+
+class SnapshotBarrierError(RuntimeError):
+    """The controller supplied an invalid snapshot release marker."""
+
+
+@dataclass(frozen=True)
+class ListenerConfig:
+    host: str | None
+    port: int
+
+
+@dataclass(frozen=True)
+class ControlArgs:
+    ready_file: Path
+    release_file: Path
+    release_timeout_s: float
+
+
+def write_ready_atomic(path: Path, oracle: Oracle) -> None:
+    _write_json_atomic(
+        path,
+        {
+            "sampled_token_logprob": oracle.sampled_token_logprob,
+            "token_ids": oracle.token_ids,
+            "text": oracle.text,
+        },
+    )
+
+
+def read_release_marker(path: Path) -> ListenerConfig:
+    path = Path(path)
+    path_stat = path.lstat()
+    if stat.S_ISLNK(path_stat.st_mode):
+        raise SnapshotBarrierError("release marker must not be a symlink")
+    try:
+        payload = json.loads(path.read_text())
+    except (OSError, json.JSONDecodeError) as error:
+        raise SnapshotBarrierError("release marker is not valid JSON") from error
+    try:
+        marker = ReleaseMarker.model_validate(payload)
+    except ValidationError as error:
+        raise SnapshotBarrierError(
+            f"release marker is invalid: {_validation_path(error)}"
+        ) from error
+    return ListenerConfig(host=marker.host, port=marker.port)
+
+
+def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]:
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--ready-file", type=Path, required=True)
+    parser.add_argument("--release-file", type=Path, required=True)
+    parser.add_argument("--release-timeout-s", type=float, default=900.0)
+    control_args, remaining = parser.parse_known_args(argv)
+    if remaining and remaining[0] == "--":
+        remaining = remaining[1:]
+    if control_args.release_timeout_s <= 0:
+        raise ValueError("release timeout must be positive")
+    return (
+        ControlArgs(
+            ready_file=control_args.ready_file,
+            release_file=control_args.release_file,
+            release_timeout_s=control_args.release_timeout_s,
+        ),
+        remaining,
+    )
+
+
+async def _release_reloadable_state(engine: Any) -> None:
+    """Discard model and KV state before the process image is captured."""
+    await engine.sleep(level=2)
+
+
+async def _restore_reloadable_state(engine: Any) -> None:
+    """Rebuild state discarded by ``_release_reloadable_state``."""
+    await engine.wake_up(tags=["weights"])
+    await engine.collective_rpc("reload_weights")
+    await engine.wake_up(tags=["kv_cache"])
+
+
+def oracle_from_request_output(request_output: Any) -> Oracle:
+    try:
+        candidate = request_output.outputs[0]
+        token_ids = tuple(candidate.token_ids)
+        (sampled_token_id,) = token_ids
+        return Oracle(
+            token_ids=token_ids,
+            text=candidate.text,
+            sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob,
+        )
+    except ValueError as error:
+        raise SnapshotCanaryError(
+            "snapshot canary must produce exactly one finite token logprob"
+        ) from error
+    except (AttributeError, IndexError, KeyError, TypeError) as error:
+        raise SnapshotCanaryError(
+            "snapshot canary did not return sampled token logprob"
+        ) from error
+
+
+async def run_engine_canary(engine: Any) -> Oracle:
+    from vllm import SamplingParams
+
+    final_output = None
+    sampling_params = SamplingParams(
+        temperature=0,
+        min_tokens=1,
+        max_tokens=1,
+        seed=0,
+        logprobs=0,
+    )
+    async for output in engine.generate(
+        _CANARY_PROMPT,
+        sampling_params,
+        request_id="vllm-snapshot-canary",
+    ):
+        final_output = output
+    if final_output is None:
+        raise SnapshotCanaryError("canary generation returned no output")
+    return oracle_from_request_output(final_output)
+
+
+def detach_snapshot_streams() -> None:
+    """Remove launch-log descriptors from the reusable process image."""
+    sys.stdout.flush()
+    sys.stderr.flush()
+    sink = os.open(os.devnull, os.O_WRONLY)
+    try:
+        os.dup2(sink, sys.stdout.fileno())
+        os.dup2(sink, sys.stderr.fileno())
+    finally:
+        os.close(sink)
+
+
+async def wait_for_release_marker(
+    path: Path,
+    *,
+    timeout_s: float,
+    poll_interval_s: float = 0.05,
+) -> ListenerConfig:
+    if timeout_s <= 0:
+        raise ValueError("release timeout must be positive")
+    if poll_interval_s <= 0:
+        raise ValueError("release poll interval must be positive")
+    remaining_s = timeout_s
+    while not path.exists():
+        if remaining_s <= 0:
+            raise TimeoutError(f"release marker not found before timeout: {path}")
+        delay_s = min(poll_interval_s, remaining_s)
+        await asyncio.sleep(delay_s)
+        remaining_s -= delay_s
+    return read_release_marker(path)
+
+
+def parse_vllm_args(argv: list[str]) -> Any:
+    from vllm.entrypoints.launchers.cli_args import (
+        make_arg_parser,
+        validate_parsed_serve_args,
+    )
+    from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+    parser = FlexibleArgumentParser(prog="python -m vllm.snapshot.server")
+    args = make_arg_parser(parser).parse_args(argv)
+    if args.model_tag is not None:
+        args.model = args.model_tag
+    if args.grpc or args.headless:
+        raise ValueError("snapshot server requires the HTTP frontend")
+    if args.api_server_count not in (None, 1):
+        raise ValueError("snapshot server supports one API server")
+    args.api_server_count = None
+    validate_parsed_serve_args(args)
+    return args
+
+
+async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None:
+    from vllm.entrypoints.launchers.api_server.entry import (
+        build_and_serve,
+        build_async_engine_client,
+    )
+    from vllm.entrypoints.launchers.launcher import setup_server
+
+    args.enable_sleep_mode = True
+
+    async with build_async_engine_client(args) as engine:
+        oracle = await run_engine_canary(engine)
+        try:
+            await _release_reloadable_state(engine)
+            await _restore_reloadable_state(engine)
+            rehearsal_oracle = await run_engine_canary(engine)
+            if not oracles_match(oracle, rehearsal_oracle):
+                raise SnapshotCanaryError("snapshot rehearsal changed canary output")
+            await _release_reloadable_state(engine)
+        except SnapshotCanaryError:
+            raise
+        except Exception as error:
+            raise SnapshotCanaryError("snapshot rehearsal failed") from error
+
+        detach_snapshot_streams()
+        write_ready_atomic(control.ready_file, oracle)
+        listener = await wait_for_release_marker(
+            control.release_file,
+            timeout_s=control.release_timeout_s,
+        )
+        args.host = listener.host
+        args.port = listener.port
+        await _restore_reloadable_state(engine)
+        listen_address, sock = setup_server(args, reuse_port=False)
+        try:
+            shutdown_task = await build_and_serve(
+                engine,
+                listen_address,
+                sock,
+                args,
+            )
+            await shutdown_task
+        finally:
+            sock.close()
+
+
+def main(argv: list[str] | None = None) -> None:
+    control, remaining = parse_control_args(sys.argv[1:] if argv is None else argv)
+    args = parse_vllm_args(remaining)
+    import uvloop
+
+    uvloop.run(run_vllm_snapshot_child(control, args))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/vllm/snapshot/types.py b/vllm/snapshot/types.py
new file mode 100644
index 000000000000..236ba124365a
--- /dev/null
+++ b/vllm/snapshot/types.py
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Oracle:
+    token_ids: tuple[int, ...]
+    text: str
+    sampled_token_logprob: float
+
+    def __post_init__(self) -> None:
+        if len(self.token_ids) != 1:
+            raise ValueError("snapshot oracle must contain exactly one token")
+        value = self.sampled_token_logprob
+        if (
+            isinstance(value, bool)
+            or not isinstance(value, (int, float))
+            or not math.isfinite(value)
+        ):
+            raise ValueError("sampled token logprob must be finite")
+        object.__setattr__(self, "sampled_token_logprob", float(value))
+
+
+def oracles_match(expected: Oracle, actual: Oracle) -> bool:
+    return (
+        expected.token_ids == actual.token_ids
+        and expected.text == actual.text
+        and math.isclose(
+            expected.sampled_token_logprob,
+            actual.sampled_token_logprob,
+            rel_tol=0.0,
+            abs_tol=1e-3,
+        )
+    )
diff --git a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py
index 60fc53350702..8f58ddcf8c7f 100644
--- a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py
+++ b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py
@@ -15,6 +15,7 @@
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv, next_power_of_2
 
+
 @triton.heuristics(
     {
         "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,
diff --git a/vllm/third_party/flash_linear_attention/ops/index.py b/vllm/third_party/flash_linear_attention/ops/index.py
index cc6f54d8af0b..880be7747dcc 100644
--- a/vllm/third_party/flash_linear_attention/ops/index.py
+++ b/vllm/third_party/flash_linear_attention/ops/index.py
@@ -10,6 +10,7 @@
 import torch
 
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.triton_utils import triton
 
 from .utils import tensor_cache
@@ -27,8 +28,8 @@ def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Te
         chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()
     indices = torch.cat([torch.arange(n) for n in chunk_counts])
     chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1)
-    return chunk_indices.to(
-        device=cu_seqlens.device, dtype=cu_seqlens.dtype, non_blocking=True
+    return async_tensor_h2d(
+        chunk_indices, device=cu_seqlens.device, dtype=cu_seqlens.dtype
     )
 
 
diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py
index 1701cd47aca8..b5c3a92b65e7 100644
--- a/vllm/third_party/flash_linear_attention/ops/kda.py
+++ b/vllm/third_party/flash_linear_attention/ops/kda.py
@@ -529,6 +529,22 @@ def forward_cuda(
             residual_in_fp32=residual_in_fp32,
         )
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        g: torch.Tensor,
+        residual: torch.Tensor | None = None,
+        prenorm: bool = False,
+        residual_in_fp32: bool = False,
+    ) -> torch.Tensor:
+        return self.forward_cuda(
+            x,
+            g,
+            residual=residual,
+            prenorm=prenorm,
+            residual_in_fp32=residual_in_fp32,
+        )
+
 
 @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 @triton.autotune(
@@ -1333,6 +1349,7 @@ def grid(meta):
     )
     return y
 
+
 def _chunk_kda_fwd_with_cumulative_g(
     q: torch.Tensor,
     k: torch.Tensor,
diff --git a/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py b/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py
index 79f5abf031fb..997df5a173b1 100644
--- a/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py
+++ b/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py
@@ -335,9 +335,7 @@ def compile(self, compile_key: CompileKey) -> None:
         assert warmup is not None
         variant = compile_key.input_variant
 
-        def pointer(
-            name: str, dtype: torch.dtype | None
-        ) -> TritonWarmupTensor | None:
+        def pointer(name: str, dtype: torch.dtype | None) -> TritonWarmupTensor | None:
             return None if dtype is None else variant.pointer(name, dtype)
 
         warmup(
diff --git a/vllm/third_party/pynvml.py b/vllm/third_party/pynvml.py
index 6aabbc217dd0..3bb53e8d9358 100644
--- a/vllm/third_party/pynvml.py
+++ b/vllm/third_party/pynvml.py
@@ -45,52 +45,54 @@
 ## C Type mappings ##
 ## Enums
 _nvmlEnableState_t = c_uint
-NVML_FEATURE_DISABLED    = 0
-NVML_FEATURE_ENABLED     = 1
+NVML_FEATURE_DISABLED = 0
+NVML_FEATURE_ENABLED = 1
 
 _nvmlBrandType_t = c_uint
-NVML_BRAND_UNKNOWN             = 0
-NVML_BRAND_QUADRO              = 1
-NVML_BRAND_TESLA               = 2
-NVML_BRAND_NVS                 = 3
-NVML_BRAND_GRID                = 4   # Deprecated from API reporting. Keeping definition for backward compatibility.
-NVML_BRAND_GEFORCE             = 5
-NVML_BRAND_TITAN               = 6
-NVML_BRAND_NVIDIA_VAPPS        = 7   # NVIDIA Virtual Applications
-NVML_BRAND_NVIDIA_VPC          = 8   # NVIDIA Virtual PC
-NVML_BRAND_NVIDIA_VCS          = 9   # NVIDIA Virtual Compute Server
-NVML_BRAND_NVIDIA_VWS          = 10  # NVIDIA RTX Virtual Workstation
+NVML_BRAND_UNKNOWN = 0
+NVML_BRAND_QUADRO = 1
+NVML_BRAND_TESLA = 2
+NVML_BRAND_NVS = 3
+NVML_BRAND_GRID = (
+    4  # Deprecated from API reporting. Keeping definition for backward compatibility.
+)
+NVML_BRAND_GEFORCE = 5
+NVML_BRAND_TITAN = 6
+NVML_BRAND_NVIDIA_VAPPS = 7  # NVIDIA Virtual Applications
+NVML_BRAND_NVIDIA_VPC = 8  # NVIDIA Virtual PC
+NVML_BRAND_NVIDIA_VCS = 9  # NVIDIA Virtual Compute Server
+NVML_BRAND_NVIDIA_VWS = 10  # NVIDIA RTX Virtual Workstation
 NVML_BRAND_NVIDIA_CLOUD_GAMING = 11  # NVIDIA Cloud Gaming
-NVML_BRAND_NVIDIA_VGAMING      = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility.
-NVML_BRAND_QUADRO_RTX          = 12
-NVML_BRAND_NVIDIA_RTX          = 13
-NVML_BRAND_NVIDIA              = 14
-NVML_BRAND_GEFORCE_RTX         = 15  # Unused
-NVML_BRAND_TITAN_RTX           = 16  # Unused
-NVML_BRAND_COUNT               = 17
+NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING  # Deprecated from API reporting. Keeping definition for backward compatibility.
+NVML_BRAND_QUADRO_RTX = 12
+NVML_BRAND_NVIDIA_RTX = 13
+NVML_BRAND_NVIDIA = 14
+NVML_BRAND_GEFORCE_RTX = 15  # Unused
+NVML_BRAND_TITAN_RTX = 16  # Unused
+NVML_BRAND_COUNT = 17
 
 _nvmlTemperatureThresholds_t = c_uint
-NVML_TEMPERATURE_THRESHOLD_SHUTDOWN      = 0
-NVML_TEMPERATURE_THRESHOLD_SLOWDOWN      = 1
-NVML_TEMPERATURE_THRESHOLD_MEM_MAX       = 2
-NVML_TEMPERATURE_THRESHOLD_GPU_MAX       = 3
-NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN  = 4
+NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0
+NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1
+NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2
+NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3
+NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4
 NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5
-NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX  = 6
-NVML_TEMPERATURE_THRESHOLD_GPS_CURR      = 7
-NVML_TEMPERATURE_THRESHOLD_COUNT         = 8
+NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6
+NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7
+NVML_TEMPERATURE_THRESHOLD_COUNT = 8
 
 _nvmlTemperatureSensors_t = c_uint
-NVML_TEMPERATURE_GPU     = 0
-NVML_TEMPERATURE_COUNT   = 1
+NVML_TEMPERATURE_GPU = 0
+NVML_TEMPERATURE_COUNT = 1
 
 
 _nvmlComputeMode_t = c_uint
-NVML_COMPUTEMODE_DEFAULT           = 0
-NVML_COMPUTEMODE_EXCLUSIVE_THREAD  = 1  ## Support Removed
-NVML_COMPUTEMODE_PROHIBITED        = 2
+NVML_COMPUTEMODE_DEFAULT = 0
+NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1  ## Support Removed
+NVML_COMPUTEMODE_PROHIBITED = 2
 NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3
-NVML_COMPUTEMODE_COUNT             = 4
+NVML_COMPUTEMODE_COUNT = 4
 
 _nvmlMemoryLocation_t = c_uint
 NVML_MEMORY_LOCATION_L1_CACHE = 0
@@ -127,158 +129,158 @@
 _nvmlNvLinkCapability_t = c_uint
 NVML_NVLINK_CAP_P2P_SUPPORTED = 0
 NVML_NVLINK_CAP_SYSMEM_ACCESS = 1
-NVML_NVLINK_CAP_P2P_ATOMICS   = 2
-NVML_NVLINK_CAP_SYSMEM_ATOMICS= 3
-NVML_NVLINK_CAP_SLI_BRIDGE    = 4
-NVML_NVLINK_CAP_VALID         = 5
-NVML_NVLINK_CAP_COUNT         = 6
+NVML_NVLINK_CAP_P2P_ATOMICS = 2
+NVML_NVLINK_CAP_SYSMEM_ATOMICS = 3
+NVML_NVLINK_CAP_SLI_BRIDGE = 4
+NVML_NVLINK_CAP_VALID = 5
+NVML_NVLINK_CAP_COUNT = 6
 
 _nvmlNvLinkUtilizationCountPktTypes_t = c_uint
-NVML_NVLINK_COUNTER_PKTFILTER_NOP        = 0x1
-NVML_NVLINK_COUNTER_PKTFILTER_READ       = 0x2
-NVML_NVLINK_COUNTER_PKTFILTER_WRITE      = 0x4
-NVML_NVLINK_COUNTER_PKTFILTER_RATOM      = 0x8
-NVML_NVLINK_COUNTER_PKTFILTER_NRATOM     = 0x10
-NVML_NVLINK_COUNTER_PKTFILTER_FLUSH      = 0x20
-NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA   = 0x40
+NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1
+NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2
+NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4
+NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8
+NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10
+NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20
+NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40
 NVML_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80
-NVML_NVLINK_COUNTER_PKTFILTER_ALL        = 0xFF
+NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF
 
 _nvmlNvLinkUtilizationCountUnits_t = c_uint
-NVML_NVLINK_COUNTER_UNIT_CYCLES   = 0
-NVML_NVLINK_COUNTER_UNIT_PACKETS  = 1
-NVML_NVLINK_COUNTER_UNIT_BYTES    = 2
+NVML_NVLINK_COUNTER_UNIT_CYCLES = 0
+NVML_NVLINK_COUNTER_UNIT_PACKETS = 1
+NVML_NVLINK_COUNTER_UNIT_BYTES = 2
 NVML_NVLINK_COUNTER_UNIT_RESERVED = 3
-NVML_NVLINK_COUNTER_UNIT_COUNT    = 4
+NVML_NVLINK_COUNTER_UNIT_COUNT = 4
 
 _nvmlNvLinkDeviceType_t = c_uint
-NVML_NVLINK_DEVICE_TYPE_GPU     = 0x00
-NVML_NVLINK_DEVICE_TYPE_IBMNPU  = 0x01
-NVML_NVLINK_DEVICE_TYPE_SWITCH  = 0x02
+NVML_NVLINK_DEVICE_TYPE_GPU = 0x00
+NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01
+NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02
 NVML_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF
 
 # These are deprecated, instead use _nvmlMemoryErrorType_t
 _nvmlEccBitType_t = c_uint
-NVML_SINGLE_BIT_ECC    = 0
-NVML_DOUBLE_BIT_ECC    = 1
+NVML_SINGLE_BIT_ECC = 0
+NVML_DOUBLE_BIT_ECC = 1
 NVML_ECC_ERROR_TYPE_COUNT = 2
 
 _nvmlEccCounterType_t = c_uint
-NVML_VOLATILE_ECC      = 0
-NVML_AGGREGATE_ECC     = 1
+NVML_VOLATILE_ECC = 0
+NVML_AGGREGATE_ECC = 1
 NVML_ECC_COUNTER_TYPE_COUNT = 2
 
 _nvmlMemoryErrorType_t = c_uint
-NVML_MEMORY_ERROR_TYPE_CORRECTED   = 0
+NVML_MEMORY_ERROR_TYPE_CORRECTED = 0
 NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1
-NVML_MEMORY_ERROR_TYPE_COUNT       = 2
+NVML_MEMORY_ERROR_TYPE_COUNT = 2
 
 _nvmlClockType_t = c_uint
-NVML_CLOCK_GRAPHICS  = 0
-NVML_CLOCK_SM        = 1
-NVML_CLOCK_MEM       = 2
-NVML_CLOCK_VIDEO     = 3
-NVML_CLOCK_COUNT     = 4
+NVML_CLOCK_GRAPHICS = 0
+NVML_CLOCK_SM = 1
+NVML_CLOCK_MEM = 2
+NVML_CLOCK_VIDEO = 3
+NVML_CLOCK_COUNT = 4
 
 _nvmlClockId_t = c_uint
-NVML_CLOCK_ID_CURRENT            = 0
-NVML_CLOCK_ID_APP_CLOCK_TARGET   = 1
-NVML_CLOCK_ID_APP_CLOCK_DEFAULT  = 2
+NVML_CLOCK_ID_CURRENT = 0
+NVML_CLOCK_ID_APP_CLOCK_TARGET = 1
+NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2
 NVML_CLOCK_ID_CUSTOMER_BOOST_MAX = 3
-NVML_CLOCK_ID_COUNT              = 4
+NVML_CLOCK_ID_COUNT = 4
 
 _nvmlDriverModel_t = c_uint
-NVML_DRIVER_WDDM       = 0
-NVML_DRIVER_WDM        = 1
-NVML_DRIVER_MCDM       = 2
+NVML_DRIVER_WDDM = 0
+NVML_DRIVER_WDM = 1
+NVML_DRIVER_MCDM = 2
 
 NVML_MAX_GPU_PERF_PSTATES = 16
 
 _nvmlPstates_t = c_uint
-NVML_PSTATE_0               = 0
-NVML_PSTATE_1               = 1
-NVML_PSTATE_2               = 2
-NVML_PSTATE_3               = 3
-NVML_PSTATE_4               = 4
-NVML_PSTATE_5               = 5
-NVML_PSTATE_6               = 6
-NVML_PSTATE_7               = 7
-NVML_PSTATE_8               = 8
-NVML_PSTATE_9               = 9
-NVML_PSTATE_10              = 10
-NVML_PSTATE_11              = 11
-NVML_PSTATE_12              = 12
-NVML_PSTATE_13              = 13
-NVML_PSTATE_14              = 14
-NVML_PSTATE_15              = 15
-NVML_PSTATE_UNKNOWN         = 32
+NVML_PSTATE_0 = 0
+NVML_PSTATE_1 = 1
+NVML_PSTATE_2 = 2
+NVML_PSTATE_3 = 3
+NVML_PSTATE_4 = 4
+NVML_PSTATE_5 = 5
+NVML_PSTATE_6 = 6
+NVML_PSTATE_7 = 7
+NVML_PSTATE_8 = 8
+NVML_PSTATE_9 = 9
+NVML_PSTATE_10 = 10
+NVML_PSTATE_11 = 11
+NVML_PSTATE_12 = 12
+NVML_PSTATE_13 = 13
+NVML_PSTATE_14 = 14
+NVML_PSTATE_15 = 15
+NVML_PSTATE_UNKNOWN = 32
 
 _nvmlInforomObject_t = c_uint
-NVML_INFOROM_OEM            = 0
-NVML_INFOROM_ECC            = 1
-NVML_INFOROM_POWER          = 2
-NVML_INFOROM_DEN            = 3
-NVML_INFOROM_COUNT          = 4
+NVML_INFOROM_OEM = 0
+NVML_INFOROM_ECC = 1
+NVML_INFOROM_POWER = 2
+NVML_INFOROM_DEN = 3
+NVML_INFOROM_COUNT = 4
 
 _nvmlReturn_t = c_uint
-NVML_SUCCESS                         = 0
-NVML_ERROR_UNINITIALIZED             = 1
-NVML_ERROR_INVALID_ARGUMENT          = 2
-NVML_ERROR_NOT_SUPPORTED             = 3
-NVML_ERROR_NO_PERMISSION             = 4
-NVML_ERROR_ALREADY_INITIALIZED       = 5
-NVML_ERROR_NOT_FOUND                 = 6
-NVML_ERROR_INSUFFICIENT_SIZE         = 7
-NVML_ERROR_INSUFFICIENT_POWER        = 8
-NVML_ERROR_DRIVER_NOT_LOADED         = 9
-NVML_ERROR_TIMEOUT                   = 10
-NVML_ERROR_IRQ_ISSUE                 = 11
-NVML_ERROR_LIBRARY_NOT_FOUND         = 12
-NVML_ERROR_FUNCTION_NOT_FOUND        = 13
-NVML_ERROR_CORRUPTED_INFOROM         = 14
-NVML_ERROR_GPU_IS_LOST               = 15
-NVML_ERROR_RESET_REQUIRED            = 16
-NVML_ERROR_OPERATING_SYSTEM          = 17
-NVML_ERROR_LIB_RM_VERSION_MISMATCH   = 18
-NVML_ERROR_IN_USE                    = 19
-NVML_ERROR_MEMORY                    = 20
-NVML_ERROR_NO_DATA                   = 21
-NVML_ERROR_VGPU_ECC_NOT_SUPPORTED    = 22
-NVML_ERROR_INSUFFICIENT_RESOURCES    = 23
-NVML_ERROR_FREQ_NOT_SUPPORTED        = 24
+NVML_SUCCESS = 0
+NVML_ERROR_UNINITIALIZED = 1
+NVML_ERROR_INVALID_ARGUMENT = 2
+NVML_ERROR_NOT_SUPPORTED = 3
+NVML_ERROR_NO_PERMISSION = 4
+NVML_ERROR_ALREADY_INITIALIZED = 5
+NVML_ERROR_NOT_FOUND = 6
+NVML_ERROR_INSUFFICIENT_SIZE = 7
+NVML_ERROR_INSUFFICIENT_POWER = 8
+NVML_ERROR_DRIVER_NOT_LOADED = 9
+NVML_ERROR_TIMEOUT = 10
+NVML_ERROR_IRQ_ISSUE = 11
+NVML_ERROR_LIBRARY_NOT_FOUND = 12
+NVML_ERROR_FUNCTION_NOT_FOUND = 13
+NVML_ERROR_CORRUPTED_INFOROM = 14
+NVML_ERROR_GPU_IS_LOST = 15
+NVML_ERROR_RESET_REQUIRED = 16
+NVML_ERROR_OPERATING_SYSTEM = 17
+NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18
+NVML_ERROR_IN_USE = 19
+NVML_ERROR_MEMORY = 20
+NVML_ERROR_NO_DATA = 21
+NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22
+NVML_ERROR_INSUFFICIENT_RESOURCES = 23
+NVML_ERROR_FREQ_NOT_SUPPORTED = 24
 NVML_ERROR_ARGUMENT_VERSION_MISMATCH = 25
-NVML_ERROR_DEPRECATED                = 26
-NVML_ERROR_NOT_READY                 = 27
-NVML_ERROR_GPU_NOT_FOUND             = 28
-NVML_ERROR_INVALID_STATE             = 29
-NVML_ERROR_UNKNOWN                   = 999
+NVML_ERROR_DEPRECATED = 26
+NVML_ERROR_NOT_READY = 27
+NVML_ERROR_GPU_NOT_FOUND = 28
+NVML_ERROR_INVALID_STATE = 29
+NVML_ERROR_UNKNOWN = 999
 
 _nvmlFanState_t = c_uint
-NVML_FAN_NORMAL             = 0
-NVML_FAN_FAILED             = 1
+NVML_FAN_NORMAL = 0
+NVML_FAN_FAILED = 1
 
 _nvmlFanControlPolicy_t = c_uint
 NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0
-NVML_FAN_POLICY_MANUAL                   = 1
+NVML_FAN_POLICY_MANUAL = 1
 
 _nvmlLedColor_t = c_uint
-NVML_LED_COLOR_GREEN        = 0
-NVML_LED_COLOR_AMBER        = 1
+NVML_LED_COLOR_GREEN = 0
+NVML_LED_COLOR_AMBER = 1
 
 _nvmlGpuOperationMode_t = c_uint
-NVML_GOM_ALL_ON                 = 0
-NVML_GOM_COMPUTE                = 1
-NVML_GOM_LOW_DP                 = 2
+NVML_GOM_ALL_ON = 0
+NVML_GOM_COMPUTE = 1
+NVML_GOM_LOW_DP = 2
 
 _nvmlPageRetirementCause_t = c_uint
 NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0
-NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR           = 1
-NVML_PAGE_RETIREMENT_CAUSE_COUNT                          = 2
+NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1
+NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2
 
 _nvmlRestrictedAPI_t = c_uint
-NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS                = 0
-NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS               = 1
-NVML_RESTRICTED_API_COUNT                                 = 2
+NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0
+NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1
+NVML_RESTRICTED_API_COUNT = 2
 
 _nvmlBridgeChipType_t = c_uint
 NVML_BRIDGE_CHIP_PLX = 0
@@ -365,9 +367,9 @@
 NVML_TOPOLOGY_SYSTEM = 50
 
 _nvmlGpuP2PCapsIndex_t = c_uint
-NVML_P2P_CAPS_INDEX_READ = 0,
+NVML_P2P_CAPS_INDEX_READ = (0,)
 NVML_P2P_CAPS_INDEX_WRITE = 1
-NVML_P2P_CAPS_INDEX_NVLINK =2
+NVML_P2P_CAPS_INDEX_NVLINK = 2
 NVML_P2P_CAPS_INDEX_ATOMICS = 3
 #
 # NVML_P2P_CAPS_INDEX_PROP is deprecated.
@@ -378,39 +380,39 @@
 NVML_P2P_CAPS_INDEX_UNKNOWN = 5
 
 _nvmlGpuP2PStatus_t = c_uint
-NVML_P2P_STATUS_OK     = 0
+NVML_P2P_STATUS_OK = 0
 NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1
 NVML_P2P_STATUS_CHIPSET_NOT_SUPPORTED = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED
 NVML_P2P_STATUS_GPU_NOT_SUPPORTED = 2
-NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED =3
-NVML_P2P_STATUS_DISABLED_BY_REGKEY =4
-NVML_P2P_STATUS_NOT_SUPPORTED =5
-NVML_P2P_STATUS_UNKNOWN =6
+NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED = 3
+NVML_P2P_STATUS_DISABLED_BY_REGKEY = 4
+NVML_P2P_STATUS_NOT_SUPPORTED = 5
+NVML_P2P_STATUS_UNKNOWN = 6
 
 _nvmlDeviceArchitecture_t = c_uint
-NVML_DEVICE_ARCH_KEPLER   = 2
-NVML_DEVICE_ARCH_MAXWELL  = 3
-NVML_DEVICE_ARCH_PASCAL   = 4
-NVML_DEVICE_ARCH_VOLTA    = 5
-NVML_DEVICE_ARCH_TURING   = 6
-NVML_DEVICE_ARCH_AMPERE   = 7
-NVML_DEVICE_ARCH_ADA      = 8
-NVML_DEVICE_ARCH_HOPPER   = 9
-NVML_DEVICE_ARCH_BLACKWELL   = 10
-NVML_DEVICE_ARCH_T23X     = 11
-NVML_DEVICE_ARCH_UNKNOWN  = 0xffffffff
+NVML_DEVICE_ARCH_KEPLER = 2
+NVML_DEVICE_ARCH_MAXWELL = 3
+NVML_DEVICE_ARCH_PASCAL = 4
+NVML_DEVICE_ARCH_VOLTA = 5
+NVML_DEVICE_ARCH_TURING = 6
+NVML_DEVICE_ARCH_AMPERE = 7
+NVML_DEVICE_ARCH_ADA = 8
+NVML_DEVICE_ARCH_HOPPER = 9
+NVML_DEVICE_ARCH_BLACKWELL = 10
+NVML_DEVICE_ARCH_T23X = 11
+NVML_DEVICE_ARCH_UNKNOWN = 0xFFFFFFFF
 
 # PCI bus Types
 _nvmlBusType_t = c_uint
 NVML_BUS_TYPE_UNKNOWN = 0
-NVML_BUS_TYPE_PCI     = 1
-NVML_BUS_TYPE_PCIE    = 2
-NVML_BUS_TYPE_FPCI    = 3
-NVML_BUS_TYPE_AGP     = 4
+NVML_BUS_TYPE_PCI = 1
+NVML_BUS_TYPE_PCIE = 2
+NVML_BUS_TYPE_FPCI = 3
+NVML_BUS_TYPE_AGP = 4
 
 _nvmlPowerSource_t = c_uint
-NVML_POWER_SOURCE_AC         = 0x00000000
-NVML_POWER_SOURCE_BATTERY    = 0x00000001
+NVML_POWER_SOURCE_AC = 0x00000000
+NVML_POWER_SOURCE_BATTERY = 0x00000001
 NVML_POWER_SOURCE_UNDERSIZED = 0x00000002
 
 _nvmlAdaptiveClockInfoStatus_t = c_uint
@@ -418,196 +420,290 @@
 NVML_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001
 
 _nvmlClockLimitId_t = c_uint
-NVML_CLOCK_LIMIT_ID_RANGE_START = 0xffffff00
-NVML_CLOCK_LIMIT_ID_TDP         = 0xffffff01
-NVML_CLOCK_LIMIT_ID_UNLIMITED   = 0xffffff02
+NVML_CLOCK_LIMIT_ID_RANGE_START = 0xFFFFFF00
+NVML_CLOCK_LIMIT_ID_TDP = 0xFFFFFF01
+NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xFFFFFF02
 
 _nvmlPcieLinkMaxSpeed_t = c_uint
-NVML_PCIE_LINK_MAX_SPEED_INVALID   = 0x00000000
-NVML_PCIE_LINK_MAX_SPEED_2500MBPS  = 0x00000001
-NVML_PCIE_LINK_MAX_SPEED_5000MBPS  = 0x00000002
-NVML_PCIE_LINK_MAX_SPEED_8000MBPS  = 0x00000003
+NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000
+NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001
+NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002
+NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003
 NVML_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004
 NVML_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005
 NVML_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006
 
 _nvmlPcieAtomicsCapability_t = c_uint
-NVML_PCIE_ATOMICS_CAP_FETCHADD32  = 0x01
-NVML_PCIE_ATOMICS_CAP_FETCHADD64  = 0x02
-NVML_PCIE_ATOMICS_CAP_SWAP32      = 0x04
-NVML_PCIE_ATOMICS_CAP_SWAP64      = 0x08
-NVML_PCIE_ATOMICS_CAP_CAS32       = 0x10
-NVML_PCIE_ATOMICS_CAP_CAS64       = 0x20
-NVML_PCIE_ATOMICS_CAP_CAS128      = 0x40
-NVML_PCIE_ATOMICS_OPS_MAX         = 7
+NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01
+NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02
+NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04
+NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08
+NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10
+NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20
+NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40
+NVML_PCIE_ATOMICS_OPS_MAX = 7
 
 _nvmlAffinityScope_t = c_uint
-NVML_AFFINITY_SCOPE_NODE   = 0
+NVML_AFFINITY_SCOPE_NODE = 0
 NVML_AFFINITY_SCOPE_SOCKET = 1
 
 _nvmlDeviceGpuRecoveryAction_t = c_uint
-NVML_GPU_RECOVERY_ACTION_NONE        = 0
-NVML_GPU_RECOVERY_ACTION_GPU_RESET   = 1
+NVML_GPU_RECOVERY_ACTION_NONE = 0
+NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1
 NVML_GPU_RECOVERY_ACTION_NODE_REBOOT = 2
-NVML_GPU_RECOVERY_ACTION_DRAIN_P2P   = 3
+NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3
 NVML_GPU_RECOVERY_ACTION_DRAIN_AND_RESET = 4
 
 # C preprocessor defined values
-nvmlFlagDefault             = 0
-nvmlFlagForce               = 1
-NVML_INIT_FLAG_NO_GPUS      = 1
-NVML_INIT_FLAG_NO_ATTACH    = 2
+nvmlFlagDefault = 0
+nvmlFlagForce = 1
+NVML_INIT_FLAG_NO_GPUS = 1
+NVML_INIT_FLAG_NO_ATTACH = 2
 
-NVML_MAX_GPC_COUNT          = 32
+NVML_MAX_GPC_COUNT = 32
 
 # buffer size
-NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE      = 16
-NVML_DEVICE_UUID_BUFFER_SIZE                 = 80
-NVML_DEVICE_UUID_V2_BUFFER_SIZE              = 96
-NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE       = 80
-NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE         = 80
-NVML_DEVICE_NAME_BUFFER_SIZE                 = 64
-NVML_DEVICE_NAME_V2_BUFFER_SIZE              = 96
-NVML_DEVICE_SERIAL_BUFFER_SIZE               = 30
-NVML_DEVICE_PART_NUMBER_BUFFER_SIZE          = 80
-NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE      = 80
-NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE        = 32
-NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE           = 32
-NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE        = 16
-NVML_GRID_LICENSE_BUFFER_SIZE                = 128
-NVML_VGPU_NAME_BUFFER_SIZE                   = 64
-NVML_GRID_LICENSE_FEATURE_MAX_COUNT          = 3
-NVML_VGPU_METADATA_OPAQUE_DATA_SIZE          = sizeof(c_uint) + 256
-NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE     = 256
-NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE  = 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH
-NVML_PERF_MODES_BUFFER_SIZE                  = 2048
+NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16
+NVML_DEVICE_UUID_BUFFER_SIZE = 80
+NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96
+NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80
+NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80
+NVML_DEVICE_NAME_BUFFER_SIZE = 64
+NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96
+NVML_DEVICE_SERIAL_BUFFER_SIZE = 30
+NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80
+NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80
+NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32
+NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32
+NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16
+NVML_GRID_LICENSE_BUFFER_SIZE = 128
+NVML_VGPU_NAME_BUFFER_SIZE = 64
+NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3
+NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256
+NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256
+NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = (
+    0x14  # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH
+)
+NVML_PERF_MODES_BUFFER_SIZE = 2048
 
 # Format strings
-NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT   = "%04X:%02X:%02X.0"
-NVML_DEVICE_PCI_BUS_ID_FMT          = "%08X:%02X:%02X.0"
+NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0"
+NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0"
 
 NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1)
 NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1)
 
-'''
+"""
  Field Identifiers.
 
  All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change.
-'''
-NVML_FI_DEV_ECC_CURRENT          = 1   # Current ECC mode. 1=Active. 0=Inactive
-NVML_FI_DEV_ECC_PENDING          = 2   # Pending ECC mode. 1=Active. 0=Inactive
-
-#ECC Count Totals
-NVML_FI_DEV_ECC_SBE_VOL_TOTAL    = 3   # Total single bit volatile ECC errors
-NVML_FI_DEV_ECC_DBE_VOL_TOTAL    = 4   # Total double bit volatile ECC errors
-NVML_FI_DEV_ECC_SBE_AGG_TOTAL    = 5   # Total single bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_DBE_AGG_TOTAL    = 6   # Total double bit aggregate (persistent) ECC errors
-#Individual ECC locations
-NVML_FI_DEV_ECC_SBE_VOL_L1       = 7   # L1 cache single bit volatile ECC errors
-NVML_FI_DEV_ECC_DBE_VOL_L1       = 8   # L1 cache double bit volatile ECC errors
-NVML_FI_DEV_ECC_SBE_VOL_L2       = 9   # L2 cache single bit volatile ECC errors
-NVML_FI_DEV_ECC_DBE_VOL_L2       = 10  # L2 cache double bit volatile ECC errors
-NVML_FI_DEV_ECC_SBE_VOL_DEV      = 11  # Device memory single bit volatile ECC errors
-NVML_FI_DEV_ECC_DBE_VOL_DEV      = 12  # Device memory double bit volatile ECC errors
-NVML_FI_DEV_ECC_SBE_VOL_REG      = 13  # Register file single bit volatile ECC errors
-NVML_FI_DEV_ECC_DBE_VOL_REG      = 14  # Register file double bit volatile ECC errors
-NVML_FI_DEV_ECC_SBE_VOL_TEX      = 15  # Texture memory single bit volatile ECC errors
-NVML_FI_DEV_ECC_DBE_VOL_TEX      = 16  # Texture memory double bit volatile ECC errors
-NVML_FI_DEV_ECC_DBE_VOL_CBU      = 17  # CBU double bit volatile ECC errors
-NVML_FI_DEV_ECC_SBE_AGG_L1       = 18  # L1 cache single bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_DBE_AGG_L1       = 19  # L1 cache double bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_SBE_AGG_L2       = 20  # L2 cache single bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_DBE_AGG_L2       = 21  # L2 cache double bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_SBE_AGG_DEV      = 22  # Device memory single bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_DBE_AGG_DEV      = 23  # Device memory double bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_SBE_AGG_REG      = 24  # Register File single bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_DBE_AGG_REG      = 25  # Register File double bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_SBE_AGG_TEX      = 26  # Texture memory single bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_DBE_AGG_TEX      = 27  # Texture memory double bit aggregate (persistent) ECC errors
-NVML_FI_DEV_ECC_DBE_AGG_CBU      = 28  # CBU double bit aggregate ECC errors
+"""
+NVML_FI_DEV_ECC_CURRENT = 1  # Current ECC mode. 1=Active. 0=Inactive
+NVML_FI_DEV_ECC_PENDING = 2  # Pending ECC mode. 1=Active. 0=Inactive
+
+# ECC Count Totals
+NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3  # Total single bit volatile ECC errors
+NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4  # Total double bit volatile ECC errors
+NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5  # Total single bit aggregate (persistent) ECC errors
+NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6  # Total double bit aggregate (persistent) ECC errors
+# Individual ECC locations
+NVML_FI_DEV_ECC_SBE_VOL_L1 = 7  # L1 cache single bit volatile ECC errors
+NVML_FI_DEV_ECC_DBE_VOL_L1 = 8  # L1 cache double bit volatile ECC errors
+NVML_FI_DEV_ECC_SBE_VOL_L2 = 9  # L2 cache single bit volatile ECC errors
+NVML_FI_DEV_ECC_DBE_VOL_L2 = 10  # L2 cache double bit volatile ECC errors
+NVML_FI_DEV_ECC_SBE_VOL_DEV = 11  # Device memory single bit volatile ECC errors
+NVML_FI_DEV_ECC_DBE_VOL_DEV = 12  # Device memory double bit volatile ECC errors
+NVML_FI_DEV_ECC_SBE_VOL_REG = 13  # Register file single bit volatile ECC errors
+NVML_FI_DEV_ECC_DBE_VOL_REG = 14  # Register file double bit volatile ECC errors
+NVML_FI_DEV_ECC_SBE_VOL_TEX = 15  # Texture memory single bit volatile ECC errors
+NVML_FI_DEV_ECC_DBE_VOL_TEX = 16  # Texture memory double bit volatile ECC errors
+NVML_FI_DEV_ECC_DBE_VOL_CBU = 17  # CBU double bit volatile ECC errors
+NVML_FI_DEV_ECC_SBE_AGG_L1 = 18  # L1 cache single bit aggregate (persistent) ECC errors
+NVML_FI_DEV_ECC_DBE_AGG_L1 = 19  # L1 cache double bit aggregate (persistent) ECC errors
+NVML_FI_DEV_ECC_SBE_AGG_L2 = 20  # L2 cache single bit aggregate (persistent) ECC errors
+NVML_FI_DEV_ECC_DBE_AGG_L2 = 21  # L2 cache double bit aggregate (persistent) ECC errors
+NVML_FI_DEV_ECC_SBE_AGG_DEV = (
+    22  # Device memory single bit aggregate (persistent) ECC errors
+)
+NVML_FI_DEV_ECC_DBE_AGG_DEV = (
+    23  # Device memory double bit aggregate (persistent) ECC errors
+)
+NVML_FI_DEV_ECC_SBE_AGG_REG = (
+    24  # Register File single bit aggregate (persistent) ECC errors
+)
+NVML_FI_DEV_ECC_DBE_AGG_REG = (
+    25  # Register File double bit aggregate (persistent) ECC errors
+)
+NVML_FI_DEV_ECC_SBE_AGG_TEX = (
+    26  # Texture memory single bit aggregate (persistent) ECC errors
+)
+NVML_FI_DEV_ECC_DBE_AGG_TEX = (
+    27  # Texture memory double bit aggregate (persistent) ECC errors
+)
+NVML_FI_DEV_ECC_DBE_AGG_CBU = 28  # CBU double bit aggregate ECC errors
 
 # Page Retirement
-NVML_FI_DEV_RETIRED_SBE          = 29  # Number of retired pages because of single bit errors
-NVML_FI_DEV_RETIRED_DBE          = 30  # Number of retired pages because of double bit errors
-NVML_FI_DEV_RETIRED_PENDING      = 31  # If any pages are pending retirement. 1=yes. 0=no.
+NVML_FI_DEV_RETIRED_SBE = 29  # Number of retired pages because of single bit errors
+NVML_FI_DEV_RETIRED_DBE = 30  # Number of retired pages because of double bit errors
+NVML_FI_DEV_RETIRED_PENDING = 31  # If any pages are pending retirement. 1=yes. 0=no.
 
 # NvLink Flit Error Counters
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0   = 32 # NVLink flow control CRC  Error Counter for Lane 0
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1   = 33 # NVLink flow control CRC  Error Counter for Lane 1
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2   = 34 # NVLink flow control CRC  Error Counter for Lane 2
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3   = 35 # NVLink flow control CRC  Error Counter for Lane 3
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4   = 36 # NVLink flow control CRC  Error Counter for Lane 4
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5   = 37 # NVLink flow control CRC  Error Counter for Lane 5
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = 38 # NVLink flow control CRC  Error Counter total for all Lanes
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = (
+    32  # NVLink flow control CRC  Error Counter for Lane 0
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = (
+    33  # NVLink flow control CRC  Error Counter for Lane 1
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = (
+    34  # NVLink flow control CRC  Error Counter for Lane 2
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = (
+    35  # NVLink flow control CRC  Error Counter for Lane 3
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = (
+    36  # NVLink flow control CRC  Error Counter for Lane 4
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = (
+    37  # NVLink flow control CRC  Error Counter for Lane 5
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = (
+    38  # NVLink flow control CRC  Error Counter total for all Lanes
+)
 
 # NvLink CRC Data Error Counters
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0   = 39 # NVLink data CRC Error Counter for Lane 0
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1   = 40 # NVLink data CRC Error Counter for Lane 1
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2   = 41 # NVLink data CRC Error Counter for Lane 2
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3   = 42 # NVLink data CRC Error Counter for Lane 3
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4   = 43 # NVLink data CRC Error Counter for Lane 4
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5   = 44 # NVLink data CRC Error Counter for Lane 5
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = 45 # NvLink data CRC Error Counter total for all Lanes
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = (
+    39  # NVLink data CRC Error Counter for Lane 0
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = (
+    40  # NVLink data CRC Error Counter for Lane 1
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = (
+    41  # NVLink data CRC Error Counter for Lane 2
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = (
+    42  # NVLink data CRC Error Counter for Lane 3
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = (
+    43  # NVLink data CRC Error Counter for Lane 4
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = (
+    44  # NVLink data CRC Error Counter for Lane 5
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = (
+    45  # NvLink data CRC Error Counter total for all Lanes
+)
 
 # NvLink Replay Error Counters
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0     = 46 # NVLink Replay Error Counter for Lane 0
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1     = 47 # NVLink Replay Error Counter for Lane 1
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2     = 48 # NVLink Replay Error Counter for Lane 2
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3     = 49 # NVLink Replay Error Counter for Lane 3
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4     = 50 # NVLink Replay Error Counter for Lane 4
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5     = 51 # NVLink Replay Error Counter for Lane 5
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL  = 52 # NVLink Replay Error Counter total for all Lanes
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46  # NVLink Replay Error Counter for Lane 0
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47  # NVLink Replay Error Counter for Lane 1
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48  # NVLink Replay Error Counter for Lane 2
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49  # NVLink Replay Error Counter for Lane 3
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50  # NVLink Replay Error Counter for Lane 4
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51  # NVLink Replay Error Counter for Lane 5
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = (
+    52  # NVLink Replay Error Counter total for all Lanes
+)
 
 # NvLink Recovery Error Counters
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0   = 53 # NVLink Recovery Error Counter for Lane 0
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1   = 54 # NVLink Recovery Error Counter for Lane 1
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2   = 55 # NVLink Recovery Error Counter for Lane 2
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3   = 56 # NVLink Recovery Error Counter for Lane 3
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4   = 57 # NVLink Recovery Error Counter for Lane 4
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5   = 58 # NVLink Recovery Error Counter for Lane 5
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = 59 # NVLink Recovery Error Counter total for all Lanes
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = (
+    53  # NVLink Recovery Error Counter for Lane 0
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = (
+    54  # NVLink Recovery Error Counter for Lane 1
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = (
+    55  # NVLink Recovery Error Counter for Lane 2
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = (
+    56  # NVLink Recovery Error Counter for Lane 3
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = (
+    57  # NVLink Recovery Error Counter for Lane 4
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = (
+    58  # NVLink Recovery Error Counter for Lane 5
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = (
+    59  # NVLink Recovery Error Counter total for all Lanes
+)
 
 # NvLink Bandwidth Counters
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0    = 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1    = 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2    = 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3    = 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4    = 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5    = 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = (
+    60  # NVLink Bandwidth Counter for Counter Set 0, Lane 0
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = (
+    61  # NVLink Bandwidth Counter for Counter Set 0, Lane 1
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = (
+    62  # NVLink Bandwidth Counter for Counter Set 0, Lane 2
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = (
+    63  # NVLink Bandwidth Counter for Counter Set 0, Lane 3
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = (
+    64  # NVLink Bandwidth Counter for Counter Set 0, Lane 4
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = (
+    65  # NVLink Bandwidth Counter for Counter Set 0, Lane 5
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = (
+    66  # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes
+)
 
 # NvLink Bandwidth Counters
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0    = 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1    = 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2    = 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3    = 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4    = 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5    = 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = (
+    67  # NVLink Bandwidth Counter for Counter Set 1, Lane 0
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = (
+    68  # NVLink Bandwidth Counter for Counter Set 1, Lane 1
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = (
+    69  # NVLink Bandwidth Counter for Counter Set 1, Lane 2
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = (
+    70  # NVLink Bandwidth Counter for Counter Set 1, Lane 3
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = (
+    71  # NVLink Bandwidth Counter for Counter Set 1, Lane 4
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = (
+    72  # NVLink Bandwidth Counter for Counter Set 1, Lane 5
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = (
+    73  # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes
+)
 
 # Perf Policy Counters
-NVML_FI_DEV_PERF_POLICY_POWER             = 74   # Perf Policy Counter for Power Policy
-NVML_FI_DEV_PERF_POLICY_THERMAL           = 75   # Perf Policy Counter for Thermal Policy
-NVML_FI_DEV_PERF_POLICY_SYNC_BOOST        = 76   # Perf Policy Counter for Sync boost Policy
-NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT       = 77   # Perf Policy Counter for Board Limit
-NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION   = 78   # Perf Policy Counter for Low GPU Utilization Policy
-NVML_FI_DEV_PERF_POLICY_RELIABILITY       = 79   # Perf Policy Counter for Reliability Policy
-NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS  = 80   # Perf Policy Counter for Total App Clock Policy
-NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = 81   # Perf Policy Counter for Total Base Clocks Policy
+NVML_FI_DEV_PERF_POLICY_POWER = 74  # Perf Policy Counter for Power Policy
+NVML_FI_DEV_PERF_POLICY_THERMAL = 75  # Perf Policy Counter for Thermal Policy
+NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76  # Perf Policy Counter for Sync boost Policy
+NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77  # Perf Policy Counter for Board Limit
+NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = (
+    78  # Perf Policy Counter for Low GPU Utilization Policy
+)
+NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79  # Perf Policy Counter for Reliability Policy
+NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = (
+    80  # Perf Policy Counter for Total App Clock Policy
+)
+NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = (
+    81  # Perf Policy Counter for Total Base Clocks Policy
+)
 
 # Memory temperatures
-NVML_FI_DEV_MEMORY_TEMP  = 82 # Memory temperature for the device
+NVML_FI_DEV_MEMORY_TEMP = 82  # Memory temperature for the device
 
 # Energy Counter
-NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded
+NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = (
+    83  # Total energy consumption for the GPU in mJ since the driver was last reloaded
+)
 
 # NVLink Speed
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L0     = 84
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L1     = 85
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L2     = 86
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L3     = 87
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L4     = 88
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L5     = 89
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89
 NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90
 
 # NVLink Link Count
@@ -622,307 +718,451 @@
 NVML_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95
 
 # NvLink Flit Error Counters
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6   = 96 # NVLink flow control CRC  Error Counter for Lane 6
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7   = 97 # NVLink flow control CRC  Error Counter for Lane 7
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8   = 98 # NVLink flow control CRC  Error Counter for Lane 8
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9   = 99 # NVLink flow control CRC  Error Counter for Lane 9
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10  = 100 # NVLink flow control CRC  Error Counter for Lane 10
-NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11  = 101 # NVLink flow control CRC  Error Counter for Lane 11
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = (
+    96  # NVLink flow control CRC  Error Counter for Lane 6
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = (
+    97  # NVLink flow control CRC  Error Counter for Lane 7
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = (
+    98  # NVLink flow control CRC  Error Counter for Lane 8
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = (
+    99  # NVLink flow control CRC  Error Counter for Lane 9
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = (
+    100  # NVLink flow control CRC  Error Counter for Lane 10
+)
+NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = (
+    101  # NVLink flow control CRC  Error Counter for Lane 11
+)
 
 # NvLink CRC Data Error Counters
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6   = 102 # NVLink data CRC Error Counter for Lane 6
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7   = 103 # NVLink data CRC Error Counter for Lane 7
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8   = 104 # NVLink data CRC Error Counter for Lane 8
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9   = 105 # NVLink data CRC Error Counter for Lane 9
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10  = 106 # NVLink data CRC Error Counter for Lane 10
-NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11  = 107 # NVLink data CRC Error Counter for Lane 11
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = (
+    102  # NVLink data CRC Error Counter for Lane 6
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = (
+    103  # NVLink data CRC Error Counter for Lane 7
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = (
+    104  # NVLink data CRC Error Counter for Lane 8
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = (
+    105  # NVLink data CRC Error Counter for Lane 9
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = (
+    106  # NVLink data CRC Error Counter for Lane 10
+)
+NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = (
+    107  # NVLink data CRC Error Counter for Lane 11
+)
 
 # NvLink Replay Error Counters
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6     = 108 # NVLink Replay Error Counter for Lane 6
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7     = 109 # NVLink Replay Error Counter for Lane 7
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8     = 110 # NVLink Replay Error Counter for Lane 8
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9     = 111 # NVLink Replay Error Counter for Lane 9
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10    = 112 # NVLink Replay Error Counter for Lane 10
-NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11    = 113 # NVLink Replay Error Counter for Lane 11
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108  # NVLink Replay Error Counter for Lane 6
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109  # NVLink Replay Error Counter for Lane 7
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110  # NVLink Replay Error Counter for Lane 8
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111  # NVLink Replay Error Counter for Lane 9
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = (
+    112  # NVLink Replay Error Counter for Lane 10
+)
+NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = (
+    113  # NVLink Replay Error Counter for Lane 11
+)
 
 # NvLink Recovery Error Counters
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6   = 114 # NVLink Recovery Error Counter for Lane 6
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7   = 115 # NVLink Recovery Error Counter for Lane 7
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8   = 116 # NVLink Recovery Error Counter for Lane 8
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9   = 117 # NVLink Recovery Error Counter for Lane 9
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10  = 118 # NVLink Recovery Error Counter for Lane 10
-NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11  = 119 # NVLink Recovery Error Counter for Lane 11
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = (
+    114  # NVLink Recovery Error Counter for Lane 6
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = (
+    115  # NVLink Recovery Error Counter for Lane 7
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = (
+    116  # NVLink Recovery Error Counter for Lane 8
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = (
+    117  # NVLink Recovery Error Counter for Lane 9
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = (
+    118  # NVLink Recovery Error Counter for Lane 10
+)
+NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = (
+    119  # NVLink Recovery Error Counter for Lane 11
+)
 
 # NvLink Bandwidth Counters
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6    = 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7    = 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8    = 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9    = 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10   = 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10
-NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11   = 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = (
+    120  # NVLink Bandwidth Counter for Counter Set 0, Lane 6
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = (
+    121  # NVLink Bandwidth Counter for Counter Set 0, Lane 7
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = (
+    122  # NVLink Bandwidth Counter for Counter Set 0, Lane 8
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = (
+    123  # NVLink Bandwidth Counter for Counter Set 0, Lane 9
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = (
+    124  # NVLink Bandwidth Counter for Counter Set 0, Lane 10
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = (
+    125  # NVLink Bandwidth Counter for Counter Set 0, Lane 11
+)
 
 # NvLink Bandwidth Counters
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6    = 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7    = 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8    = 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9    = 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10   = 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10
-NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11   = 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = (
+    126  # NVLink Bandwidth Counter for Counter Set 1, Lane 6
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = (
+    127  # NVLink Bandwidth Counter for Counter Set 1, Lane 7
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = (
+    128  # NVLink Bandwidth Counter for Counter Set 1, Lane 8
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = (
+    129  # NVLink Bandwidth Counter for Counter Set 1, Lane 9
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = (
+    130  # NVLink Bandwidth Counter for Counter Set 1, Lane 10
+)
+NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = (
+    131  # NVLink Bandwidth Counter for Counter Set 1, Lane 11
+)
 
 # NVLink Speed
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L6     = 132
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L7     = 133
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L8     = 134
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L9     = 135
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L10    = 136
-NVML_FI_DEV_NVLINK_SPEED_MBPS_L11    = 137
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136
+NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137
 
 # NVLink Throughput Counters
-NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB
-NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB
-NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX  = 140 # NVLink TX Data + protocol overhead in KiB
-NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX  = 141 # NVLink RX Data + protocol overhead in KiB
+NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138  # NVLink TX Data throughput in KiB
+NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139  # NVLink RX Data throughput in KiB
+NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140  # NVLink TX Data + protocol overhead in KiB
+NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141  # NVLink RX Data + protocol overhead in KiB
 
 # Row Remapper
-NVML_FI_DEV_REMAPPED_COR        = 142
-NVML_FI_DEV_REMAPPED_UNC        = 143
-NVML_FI_DEV_REMAPPED_PENDING    = 144
-NVML_FI_DEV_REMAPPED_FAILURE    = 145
+NVML_FI_DEV_REMAPPED_COR = 142
+NVML_FI_DEV_REMAPPED_UNC = 143
+NVML_FI_DEV_REMAPPED_PENDING = 144
+NVML_FI_DEV_REMAPPED_FAILURE = 145
 
-#Remote device NVLink ID
+# Remote device NVLink ID
 NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146
 
 # Number of NVLinks connected to NVSwitch
 NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147
 
 # NvLink ECC Data Error Counters
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0    = 148 #< NVLink data ECC Error Counter for Link 0
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1    = 149 #< NVLink data ECC Error Counter for Link 1
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2    = 150 #< NVLink data ECC Error Counter for Link 2
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3    = 151 #< NVLink data ECC Error Counter for Link 3
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4    = 152 #< NVLink data ECC Error Counter for Link 4
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5    = 153 #< NVLink data ECC Error Counter for Link 5
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6    = 154 #< NVLink data ECC Error Counter for Link 6
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7    = 155 #< NVLink data ECC Error Counter for Link 7
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8    = 156 #< NVLink data ECC Error Counter for Link 8
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9    = 157 #< NVLink data ECC Error Counter for Link 9
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10   = 158 #< NVLink data ECC Error Counter for Link 10
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11   = 159 #< NVLink data ECC Error Counter for Link 11
-NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = 160 #< NvLink data ECC Error Counter total for all Links
-
-NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY            = 161
-NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY          = 162
-NVML_FI_DEV_NVLINK_ERROR_DL_CRC               = 163
-NVML_FI_DEV_NVLINK_GET_SPEED                  = 164
-NVML_FI_DEV_NVLINK_GET_STATE                  = 165
-NVML_FI_DEV_NVLINK_GET_VERSION                = 166
-
-NVML_FI_DEV_NVLINK_GET_POWER_STATE            = 167
-NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD        = 168
-
-NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER       = 169
-
-NVML_FI_DEV_C2C_LINK_COUNT                    = 170
-NVML_FI_DEV_C2C_LINK_GET_STATUS               = 171
-NVML_FI_DEV_C2C_LINK_GET_MAX_BW               = 172
-
-NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS     = 173
-NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED          = 174
-NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR         = 175
-NVML_FI_DEV_PCIE_COUNT_BAD_TLP                = 176
-NVML_FI_DEV_PCIE_COUNT_NAKS_SENT              = 177
-NVML_FI_DEV_PCIE_COUNT_BAD_DLLP               = 178
-NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR        = 179
-NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR            = 180
-NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ        = 181
-NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR             = 182
-NVML_FI_DEV_PCIE_COUNT_LANE_ERROR             = 183
-
-NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED        = 184
-
-NVML_FI_DEV_POWER_AVERAGE                     = 185
-NVML_FI_DEV_POWER_INSTANT                     = 186
-NVML_FI_DEV_POWER_MIN_LIMIT                   = 187
-NVML_FI_DEV_POWER_MAX_LIMIT                   = 188
-NVML_FI_DEV_POWER_DEFAULT_LIMIT               = 189
-NVML_FI_DEV_POWER_CURRENT_LIMIT               = 190
-NVML_FI_DEV_ENERGY                            = 191
-NVML_FI_DEV_POWER_REQUESTED_LIMIT             = 192
-
-NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT       = 193
-NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT       = 194
-NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT        = 195
-NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT        = 196
-
-NVML_FI_DEV_PCIE_COUNT_TX_BYTES               = 197
-NVML_FI_DEV_PCIE_COUNT_RX_BYTES               = 198
-
-NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE   = 199
-
-NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX              = 200
-
-NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS                    = 201
-NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES                      = 202
-NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS                     = 203
-NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES                       = 204
-NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED                    = 205 # Deprecated, do not use
-NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS         = 206
-NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS           = 207
-NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS                      = 208
-NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS               = 209
-NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS              = 210
-NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS     = 211
-NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS                   = 212
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = (
+    148  # < NVLink data ECC Error Counter for Link 0
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = (
+    149  # < NVLink data ECC Error Counter for Link 1
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = (
+    150  # < NVLink data ECC Error Counter for Link 2
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = (
+    151  # < NVLink data ECC Error Counter for Link 3
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = (
+    152  # < NVLink data ECC Error Counter for Link 4
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = (
+    153  # < NVLink data ECC Error Counter for Link 5
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = (
+    154  # < NVLink data ECC Error Counter for Link 6
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = (
+    155  # < NVLink data ECC Error Counter for Link 7
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = (
+    156  # < NVLink data ECC Error Counter for Link 8
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = (
+    157  # < NVLink data ECC Error Counter for Link 9
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = (
+    158  # < NVLink data ECC Error Counter for Link 10
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = (
+    159  # < NVLink data ECC Error Counter for Link 11
+)
+NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = (
+    160  # < NvLink data ECC Error Counter total for all Links
+)
+
+NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161
+NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162
+NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163
+NVML_FI_DEV_NVLINK_GET_SPEED = 164
+NVML_FI_DEV_NVLINK_GET_STATE = 165
+NVML_FI_DEV_NVLINK_GET_VERSION = 166
+
+NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167
+NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168
+
+NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169
+
+NVML_FI_DEV_C2C_LINK_COUNT = 170
+NVML_FI_DEV_C2C_LINK_GET_STATUS = 171
+NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172
+
+NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173
+NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174
+NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175
+NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176
+NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177
+NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178
+NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179
+NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180
+NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181
+NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182
+NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183
+
+NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184
+
+NVML_FI_DEV_POWER_AVERAGE = 185
+NVML_FI_DEV_POWER_INSTANT = 186
+NVML_FI_DEV_POWER_MIN_LIMIT = 187
+NVML_FI_DEV_POWER_MAX_LIMIT = 188
+NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189
+NVML_FI_DEV_POWER_CURRENT_LIMIT = 190
+NVML_FI_DEV_ENERGY = 191
+NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192
+
+NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193
+NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194
+NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195
+NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196
+
+NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197
+NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198
+
+NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199
+
+NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200
+
+NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201
+NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202
+NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203
+NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204
+NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205  # Deprecated, do not use
+NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206
+NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207
+NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208
+NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209
+NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210
+NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211
+NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212
 
 NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_SUCCESSFUL_EVENTS = 213
-NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS     = 214
-NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS            = 215
-
-NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0                   = 216  # Deprecated, do not use
-NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1                   = 217  # Deprecated, do not use
-NVML_FI_DEV_NVLINK_COUNT_RAW_BER                         = 218  # Deprecated, do not use
-NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS                = 219
-NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER                   = 220
-NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS                   = 221
-NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER                      = 222
-
-NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN               = 223
-NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS             = 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_*
-NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED         = 225
-
-NVML_FI_DEV_RESET_STATUS                                 = 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead 
-NVML_FI_DEV_DRAIN_AND_RESET_STATUS                       = 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead
-NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK                   = 228
-NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK                    = 229
-NVML_FI_DEV_GET_GPU_RECOVERY_ACTION                      = 230
-
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0                   = 235
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1                   = 236
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2                   = 237
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3                   = 238
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4                   = 239
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5                   = 240
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6                   = 241
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7                   = 242
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8                   = 243
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9                   = 244
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10                  = 245
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11                  = 246
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12                  = 247
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13                  = 248
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14                  = 249
-NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15                  = 250
-NVML_FI_PWR_SMOOTHING_ENABLED                                   = 251 # Enablement (0/DISABLED or 1/ENABLED)
-NVML_FI_PWR_SMOOTHING_PRIV_LVL                                  = 252 # Current privilege level
-NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED                     = 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED)
-NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL                          = 254 # Applied TMP ceiling value
-NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR                         = 255 # Applied TMP floor value
-NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING             = 256 # Max % TMP Floor value
-NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING             = 257 # Min % TMP Floor value
-NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING   = 258 # HW Circuitry % lifetime remaining
-NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES                   = 259 # Max number of preset profiles
-NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR                 = 260 # % TMP floor for a given profile
-NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE                      = 261 # Ramp up rate in mW/s for a given profile
-NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE                    = 262 # Ramp down rate in mW/s for a given profile
-NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL                = 263 # Ramp down hysteresis value in ms for a given profile
-NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE                     = 264 # Active preset profile number
-NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR          = 265 # % TMP floor for a given profile
-NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE               = 266 # Ramp up rate in mW/s for a given profile
-NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE             = 267 # Ramp down rate in mW/s for a given profile
-NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL         = 268 # Ramp down hysteresis value in ms for a given profile
-
-NVML_FI_MAX = 269 # One greater than the largest field ID defined above
+NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214
+NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215
+
+NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216  # Deprecated, do not use
+NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217  # Deprecated, do not use
+NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218  # Deprecated, do not use
+NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219
+NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220
+NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221
+NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222
+
+NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223
+NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = (
+    224  # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_*
+)
+NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225
+
+NVML_FI_DEV_RESET_STATUS = (
+    226  # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead
+)
+NVML_FI_DEV_DRAIN_AND_RESET_STATUS = (
+    227  # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead
+)
+NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228
+NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229
+NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230
+
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249
+NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250
+NVML_FI_PWR_SMOOTHING_ENABLED = 251  # Enablement (0/DISABLED or 1/ENABLED)
+NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252  # Current privilege level
+NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = (
+    253  # Immediate ramp down enablement (0/DISABLED or 1/ENABLED)
+)
+NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254  # Applied TMP ceiling value
+NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255  # Applied TMP floor value
+NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 256  # Max % TMP Floor value
+NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 257  # Min % TMP Floor value
+NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = (
+    258  # HW Circuitry % lifetime remaining
+)
+NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 259  # Max number of preset profiles
+NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 260  # % TMP floor for a given profile
+NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = (
+    261  # Ramp up rate in mW/s for a given profile
+)
+NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = (
+    262  # Ramp down rate in mW/s for a given profile
+)
+NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = (
+    263  # Ramp down hysteresis value in ms for a given profile
+)
+NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 264  # Active preset profile number
+NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = (
+    265  # % TMP floor for a given profile
+)
+NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = (
+    266  # Ramp up rate in mW/s for a given profile
+)
+NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = (
+    267  # Ramp down rate in mW/s for a given profile
+)
+NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = (
+    268  # Ramp down hysteresis value in ms for a given profile
+)
+
+NVML_FI_MAX = 269  # One greater than the largest field ID defined above
 
 # NVML_FI_DEV_NVLINK_GET_STATE state enums
 NVML_NVLINK_STATE_INACTIVE = 0x0
-NVML_NVLINK_STATE_ACTIVE   = 0x1
-NVML_NVLINK_STATE_SLEEP    = 0x2
+NVML_NVLINK_STATE_ACTIVE = 0x1
+NVML_NVLINK_STATE_SLEEP = 0x2
 
-NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS
-NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US  = 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS
+NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = (
+    0  # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS
+)
+NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = (
+    1  # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS
+)
 
 ## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode
-NVML_GPU_VIRTUALIZATION_MODE_NONE        = 0  # Represents Bare Metal GPU
-NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = 1  # Device is associated with GPU-Passthorugh
-NVML_GPU_VIRTUALIZATION_MODE_VGPU        = 2  # Device is associated with vGPU inside virtual machine.
-NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU   = 3  # Device is associated with VGX hypervisor in vGPU mode
-NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA   = 4  # Device is associated with VGX hypervisor in vSGA mode
+NVML_GPU_VIRTUALIZATION_MODE_NONE = 0  # Represents Bare Metal GPU
+NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = (
+    1  # Device is associated with GPU-Passthorugh
+)
+NVML_GPU_VIRTUALIZATION_MODE_VGPU = (
+    2  # Device is associated with vGPU inside virtual machine.
+)
+NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = (
+    3  # Device is associated with VGX hypervisor in vGPU mode
+)
+NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = (
+    4  # Device is associated with VGX hypervisor in vSGA mode
+)
 
 ## Lib loading ##
 nvmlLib = None
 libLoadLock = threading.Lock()
-_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown
+_nvmlLib_refcount = 0  # Incremented on each nvmlInit and decremented on nvmlShutdown
 
 ## vGPU Management
-_nvmlVgpuTypeId_t   = c_uint
+_nvmlVgpuTypeId_t = c_uint
 _nvmlVgpuInstance_t = c_uint
 
 _nvmlVgpuVmIdType_t = c_uint
-NVML_VGPU_VM_ID_DOMAIN_ID    = 0
-NVML_VGPU_VM_ID_UUID         = 1
+NVML_VGPU_VM_ID_DOMAIN_ID = 0
+NVML_VGPU_VM_ID_UUID = 1
 
 _nvmlGridLicenseFeatureCode_t = c_uint
-NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN      = 0
-NVML_GRID_LICENSE_FEATURE_CODE_VGPU         = 1
-NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX   = 2
-NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX.
-NVML_GRID_LICENSE_FEATURE_CODE_GAMING       = 3
-NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE      = 4
+NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0
+NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1
+NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2
+NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = (
+    2  # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX.
+)
+NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3
+NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4
 
 _nvmlGridLicenseExpiryStatus_t = c_uint8
-NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE    = 0,   # Expiry information not available
-NVML_GRID_LICENSE_EXPIRY_INVALID          = 1,   # Invalid expiry or error fetching expiry
-NVML_GRID_LICENSE_EXPIRY_VALID            = 2,   # Valid expiry
-NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE   = 3,   # Expiry not applicable
-NVML_GRID_LICENSE_EXPIRY_PERMANENT        = 4,   # Permanent expiry
+NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = (0,)  # Expiry information not available
+NVML_GRID_LICENSE_EXPIRY_INVALID = (1,)  # Invalid expiry or error fetching expiry
+NVML_GRID_LICENSE_EXPIRY_VALID = (2,)  # Valid expiry
+NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = (3,)  # Expiry not applicable
+NVML_GRID_LICENSE_EXPIRY_PERMANENT = (4,)  # Permanent expiry
 
 _nvmlVgpuCapability_t = c_uint
-NVML_VGPU_CAP_NVLINK_P2P                    = 0  # vGPU P2P over NVLink is supported
-NVML_VGPU_CAP_GPUDIRECT                     = 1  # GPUDirect capability is supported
-NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE          = 2  # vGPU profile cannot be mixed with other vGPU profiles in same VM
-NVML_VGPU_CAP_EXCLUSIVE_TYPE                = 3  # vGPU profile cannot run on a GPU alongside other profiles of different type
-NVML_VGPU_CAP_EXCLUSIVE_SIZE                = 4  # vGPU profile cannot run on a GPU alongside other profiles of different size
-NVML_VGPU_CAP_COUNT                         = 5
+NVML_VGPU_CAP_NVLINK_P2P = 0  # vGPU P2P over NVLink is supported
+NVML_VGPU_CAP_GPUDIRECT = 1  # GPUDirect capability is supported
+NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = (
+    2  # vGPU profile cannot be mixed with other vGPU profiles in same VM
+)
+NVML_VGPU_CAP_EXCLUSIVE_TYPE = (
+    3  # vGPU profile cannot run on a GPU alongside other profiles of different type
+)
+NVML_VGPU_CAP_EXCLUSIVE_SIZE = (
+    4  # vGPU profile cannot run on a GPU alongside other profiles of different size
+)
+NVML_VGPU_CAP_COUNT = 5
 
 _nvmlVgpuDriverCapability_t = c_uint
-NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU   = 0  # Supports mixing of different vGPU profiles within one guest VM
-NVML_VGPU_DRIVER_CAP_WARM_UPDATE                = 1  # Supports FSR and warm update of vGPU host driver without terminating the running guest VM
-NVML_VGPU_DRIVER_CAP_COUNT                      = 2
+NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = (
+    0  # Supports mixing of different vGPU profiles within one guest VM
+)
+NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1  # Supports FSR and warm update of vGPU host driver without terminating the running guest VM
+NVML_VGPU_DRIVER_CAP_COUNT = 2
 
 _nvmlDeviceVgpuCapability_t = c_uint
-NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU             = 0  # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations
-NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES  = 1  # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types
-NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES     = 2  # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes
-NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW             = 3  # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second
-NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW            = 4  # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second
-NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING                  = 5  # Query whether the vGPU profiles on the GPU supports migration data streaming
-NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU                  = 6  # Set/Get support of mini-quarter vGPU profiles
-NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU          = 7  # Set/Get support for compute media engine vGPU profiles
-NVML_DEVICE_VGPU_CAP_WARM_UPDATE                       = 8  # Query whether the GPU supports FSR and warm update
-NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS            = 9  # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes
-NVML_DEVICE_VGPU_CAP_COUNT                             = 10
+NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0  # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations
+NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1  # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types
+NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2  # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes
+NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3  # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second
+NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4  # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second
+NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = (
+    5  # Query whether the vGPU profiles on the GPU supports migration data streaming
+)
+NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = (
+    6  # Set/Get support of mini-quarter vGPU profiles
+)
+NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = (
+    7  # Set/Get support for compute media engine vGPU profiles
+)
+NVML_DEVICE_VGPU_CAP_WARM_UPDATE = (
+    8  # Query whether the GPU supports FSR and warm update
+)
+NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9  # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes
+NVML_DEVICE_VGPU_CAP_COUNT = 10
 
 _nvmlVgpuGuestInfoState_t = c_uint
 NVML_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0
-NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED   = 1
+NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1
 
 _nvmlVgpuVmCompatibility_t = c_uint
-NVML_VGPU_VM_COMPATIBILITY_NONE         = 0x0
-NVML_VGPU_VM_COMPATIBILITY_COLD         = 0x1
-NVML_VGPU_VM_COMPATIBILITY_HIBERNATE    = 0x2
-NVML_VGPU_VM_COMPATIBILITY_SLEEP        = 0x4
-NVML_VGPU_VM_COMPATIBILITY_LIVE         = 0x8
+NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0
+NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1
+NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2
+NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4
+NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8
 
 _nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint
-NVML_VGPU_COMPATIBILITY_LIMIT_NONE          = 0x0
-NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER   = 0x1
-NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER  = 0x2
-NVML_VGPU_COMPATIBILITY_LIMIT_GPU           = 0x4
-NVML_VGPU_COMPATIBILITY_LIMIT_OTHER         = 0x80000000
+NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0
+NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1
+NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2
+NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4
+NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000
 
 _nvmlHostVgpuMode_t = c_uint
-NVML_HOST_VGPU_MODE_NON_SRIOV   = 0
-NVML_HOST_VGPU_MODE_SRIOV       = 1
+NVML_HOST_VGPU_MODE_NON_SRIOV = 0
+NVML_HOST_VGPU_MODE_SRIOV = 1
 
 _nvmlConfComputeGpusReadyState_t = c_uint
 NVML_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0
@@ -945,11 +1185,11 @@
 
 NVML_CC_SYSTEM_MULTIGPU_NONE = 0
 NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1
- 
+
 NVML_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0
 NVML_CC_SYSTEM_ENVIRONMENT_SIM = 1
 NVML_CC_SYSTEM_ENVIRONMENT_PROD = 2
- 
+
 _nvmlConfComputeCcFeature_t = c_uint
 NVML_CC_SYSTEM_FEATURE_DISABLED = 0
 NVML_CC_SYSTEM_FEATURE_ENABLED = 1
@@ -961,62 +1201,71 @@
 # GSP firmware
 NVML_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40
 
+
 class NVMLLibraryMismatchError(Exception):
     pass
 
+
 ## Error Checking ##
 class NVMLError(Exception):
     _valClassMapping = dict()
     # List of currently known error codes
     _errcode_to_string = {
-        NVML_ERROR_UNINITIALIZED:       "Uninitialized",
-        NVML_ERROR_INVALID_ARGUMENT:    "Invalid Argument",
-        NVML_ERROR_NOT_SUPPORTED:       "Not Supported",
-        NVML_ERROR_NO_PERMISSION:       "Insufficient Permissions",
+        NVML_ERROR_UNINITIALIZED: "Uninitialized",
+        NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument",
+        NVML_ERROR_NOT_SUPPORTED: "Not Supported",
+        NVML_ERROR_NO_PERMISSION: "Insufficient Permissions",
         NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized",
-        NVML_ERROR_NOT_FOUND:           "Not Found",
-        NVML_ERROR_INSUFFICIENT_SIZE:   "Insufficient Size",
-        NVML_ERROR_INSUFFICIENT_POWER:  "Insufficient External Power",
-        NVML_ERROR_DRIVER_NOT_LOADED:   "Driver Not Loaded",
-        NVML_ERROR_TIMEOUT:             "Timeout",
-        NVML_ERROR_IRQ_ISSUE:           "Interrupt Request Issue",
-        NVML_ERROR_LIBRARY_NOT_FOUND:   "NVML Shared Library Not Found",
-        NVML_ERROR_FUNCTION_NOT_FOUND:  "Function Not Found",
-        NVML_ERROR_CORRUPTED_INFOROM:   "Corrupted infoROM",
-        NVML_ERROR_GPU_IS_LOST:         "GPU is lost",
-        NVML_ERROR_RESET_REQUIRED:      "GPU requires restart",
-        NVML_ERROR_OPERATING_SYSTEM:    "The operating system has blocked the request.",
+        NVML_ERROR_NOT_FOUND: "Not Found",
+        NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size",
+        NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power",
+        NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded",
+        NVML_ERROR_TIMEOUT: "Timeout",
+        NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue",
+        NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found",
+        NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found",
+        NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM",
+        NVML_ERROR_GPU_IS_LOST: "GPU is lost",
+        NVML_ERROR_RESET_REQUIRED: "GPU requires restart",
+        NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.",
         NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.",
-        NVML_ERROR_MEMORY:              "Insufficient Memory",
-        NVML_ERROR_UNKNOWN:             "Unknown Error",
-        }
+        NVML_ERROR_MEMORY: "Insufficient Memory",
+        NVML_ERROR_UNKNOWN: "Unknown Error",
+    }
+
     def __new__(typ, value):
-        '''
+        """
         Maps value to a proper subclass of NVMLError.
         See _extractNVMLErrorsAsClasses function for more details
-        '''
+        """
         if typ == NVMLError:
             typ = NVMLError._valClassMapping.get(value, typ)
         obj = Exception.__new__(typ)
         obj.value = value
         return obj
+
     def __str__(self):
         try:
             if self.value not in NVMLError._errcode_to_string:
-                NVMLError._errcode_to_string[self.value] = str(nvmlErrorString(self.value))
+                NVMLError._errcode_to_string[self.value] = str(
+                    nvmlErrorString(self.value)
+                )
             return NVMLError._errcode_to_string[self.value]
         except NVMLError:
             return "NVML Error with code %d" % self.value
+
     def __eq__(self, other):
         return self.value == other.value
 
+
 def nvmlExceptionClass(nvmlErrorCode):
     if nvmlErrorCode not in NVMLError._valClassMapping:
-        raise ValueError('nvmlErrorCode %s is not valid' % nvmlErrorCode)
+        raise ValueError("nvmlErrorCode %s is not valid" % nvmlErrorCode)
     return NVMLError._valClassMapping[nvmlErrorCode]
 
+
 def _extractNVMLErrorsAsClasses():
-    '''
+    """
     Generates a hierarchy of classes on top of NVMLError class.
 
     Each NVML Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate
@@ -1024,31 +1273,44 @@ def _extractNVMLErrorsAsClasses():
 
     NVMLError is a parent class. Each NVML_ERROR_* gets its own subclass.
     e.g. NVML_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized
-    '''
+    """
     this_module = sys.modules[__name__]
     nvmlErrorsNames = [x for x in dir(this_module) if x.startswith("NVML_ERROR_")]
     for err_name in nvmlErrorsNames:
         # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized
-        class_name = "NVMLError_" + string.capwords(err_name.replace("NVML_ERROR_", ""), "_").replace("_", "")
+        class_name = "NVMLError_" + string.capwords(
+            err_name.replace("NVML_ERROR_", ""), "_"
+        ).replace("_", "")
         err_val = getattr(this_module, err_name)
+
         def gen_new(val):
             def new(typ):
                 obj = NVMLError.__new__(typ, val)
                 return obj
+
             return new
-        new_error_class = type(class_name, (NVMLError,), {'__new__': gen_new(err_val)})
+
+        new_error_class = type(class_name, (NVMLError,), {"__new__": gen_new(err_val)})
         new_error_class.__module__ = __name__
         setattr(this_module, class_name, new_error_class)
         NVMLError._valClassMapping[err_val] = new_error_class
+
+
 _extractNVMLErrorsAsClasses()
 
+
 def _nvmlCheckReturn(ret):
-    if (ret != NVML_SUCCESS):
+    if ret != NVML_SUCCESS:
         raise NVMLError(ret)
     return ret
 
+
 ## Function access ##
-_nvmlGetFunctionPointer_cache = dict() # function pointers are cached to prevent unnecessary libLoadLock locking
+_nvmlGetFunctionPointer_cache = (
+    dict()
+)  # function pointers are cached to prevent unnecessary libLoadLock locking
+
+
 def _nvmlGetFunctionPointer(name):
     global nvmlLib
 
@@ -1058,7 +1320,7 @@ def _nvmlGetFunctionPointer(name):
     libLoadLock.acquire()
     try:
         # ensure library was loaded
-        if (nvmlLib == None):
+        if nvmlLib == None:
             raise NVMLError(NVML_ERROR_UNINITIALIZED)
         try:
             _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name)
@@ -1069,6 +1331,7 @@ def _nvmlGetFunctionPointer(name):
         # lock is always freed
         libLoadLock.release()
 
+
 ## Alternative object
 # Allows the object to be printed
 # Allows mismatched types to be assigned
@@ -1077,9 +1340,11 @@ class nvmlFriendlyObject(object):
     def __init__(self, dictionary):
         for x in dictionary:
             setattr(self, x, dictionary[x])
+
     def __str__(self):
         return self.__dict__.__str__()
 
+
 def nvmlStructToFriendlyObject(struct):
     d = {}
     for x in struct._fields_:
@@ -1090,6 +1355,7 @@ def nvmlStructToFriendlyObject(struct):
     obj = nvmlFriendlyObject(d)
     return obj
 
+
 # pack the object so it can be passed to the NVML library
 def nvmlFriendlyObjectToStruct(obj, model):
     for x in model._fields_:
@@ -1102,11 +1368,15 @@ def nvmlFriendlyObjectToStruct(obj, model):
             setattr(model, key, value)
     return model
 
+
 ## Unit structures
 class struct_c_nvmlUnit_t(Structure):
-    pass # opaque handle
+    pass  # opaque handle
+
+
 c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t)
 
+
 class _PrintableStructure(Structure):
     """
     Abstract class that produces nicer __str__ output than ctypes.Structure.
@@ -1126,7 +1396,9 @@ class _PrintableStructure(Structure):
 
     Exact format of returned str from this class is subject to change in the future.
     """
+
     _fmt_ = {}
+
     def __str__(self):
         result = []
         for x in self._fields_:
@@ -1138,7 +1410,7 @@ def __str__(self):
             elif "<default>" in self._fmt_:
                 fmt = self._fmt_["<default>"]
             result.append(("%s: " + fmt) % (key, value))
-        return self.__class__.__name__ + "(" +  ", ".join(result) + ")"
+        return self.__class__.__name__ + "(" + ", ".join(result) + ")"
 
     def __getattribute__(self, name):
         res = super(_PrintableStructure, self).__getattribute__(name)
@@ -1159,185 +1431,200 @@ def __setattr__(self, name, value):
             value = value.encode()
         super(_PrintableStructure, self).__setattr__(name, value)
 
+
 class c_nvmlUnitInfo_t(_PrintableStructure):
     _fields_ = [
-        ('name', c_char * 96),
-        ('id', c_char * 96),
-        ('serial', c_char * 96),
-        ('firmwareVersion', c_char * 96),
+        ("name", c_char * 96),
+        ("id", c_char * 96),
+        ("serial", c_char * 96),
+        ("firmwareVersion", c_char * 96),
     ]
 
+
 class c_nvmlC2cModeInfo_v1_t(_PrintableStructure):
-    _fields_ = [
-        ('isC2cEnabled', c_uint)
-    ]
+    _fields_ = [("isC2cEnabled", c_uint)]
+
+
+nvmlC2cModeInfo_v1 = 0x1000008
 
-nvmlC2cModeInfo_v1 = 0x1000008;
 
 class c_nvmlLedState_t(_PrintableStructure):
     _fields_ = [
-        ('cause', c_char * 256),
-        ('color', _nvmlLedColor_t),
+        ("cause", c_char * 256),
+        ("color", _nvmlLedColor_t),
     ]
 
+
 class c_nvmlPSUInfo_t(_PrintableStructure):
     _fields_ = [
-        ('state', c_char * 256),
-        ('current', c_uint),
-        ('voltage', c_uint),
-        ('power', c_uint),
+        ("state", c_char * 256),
+        ("current", c_uint),
+        ("voltage", c_uint),
+        ("power", c_uint),
     ]
 
+
 class c_nvmlUnitFanInfo_t(_PrintableStructure):
     _fields_ = [
-        ('speed', c_uint),
-        ('state', _nvmlFanState_t),
+        ("speed", c_uint),
+        ("state", _nvmlFanState_t),
     ]
 
+
 class c_nvmlUnitFanSpeeds_t(_PrintableStructure):
-    _fields_ = [
-        ('fans', c_nvmlUnitFanInfo_t * 24),
-        ('count', c_uint)
-    ]
+    _fields_ = [("fans", c_nvmlUnitFanInfo_t * 24), ("count", c_uint)]
+
 
 ## Device structures
 class struct_c_nvmlDevice_t(Structure):
-    pass # opaque handle
+    pass  # opaque handle
+
+
 c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t)
 
+
 class nvmlPciInfoExt_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('domain', c_uint),
-        ('bus', c_uint),
-        ('device', c_uint),
-        ('pciDeviceId', c_uint),
-        ('pciSubSystemId', c_uint),
-        ('baseClass', c_uint),
-        ('subClass', c_uint),
-        ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE),
+        ("version", c_uint),
+        ("domain", c_uint),
+        ("bus", c_uint),
+        ("device", c_uint),
+        ("pciDeviceId", c_uint),
+        ("pciSubSystemId", c_uint),
+        ("baseClass", c_uint),
+        ("subClass", c_uint),
+        ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE),
     ]
     _fmt_ = {
-            'version'        : "0x%04X",
-            'domain'         : "0x%04X",
-            'bus'            : "0x%02X",
-            'device'         : "0x%02X",
-            'pciDeviceId'    : "0x%08X",
-            'pciSubSystemId' : "0x%08X",
-            'baseClass'      : "0x%01X",
-            'subClass'       : "0x%01X",
-            }
+        "version": "0x%04X",
+        "domain": "0x%04X",
+        "bus": "0x%02X",
+        "device": "0x%02X",
+        "pciDeviceId": "0x%08X",
+        "pciSubSystemId": "0x%08X",
+        "baseClass": "0x%01X",
+        "subClass": "0x%01X",
+    }
+
 
 nvmlPciInfoExt_v1 = 0x1000040
 
+
 # Legacy pciInfo used for _v1 and _v2
 class nvmlPciInfo_v2_t(_PrintableStructure):
     _fields_ = [
-        ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE),
-        ('domain', c_uint),
-        ('bus', c_uint),
-        ('device', c_uint),
-        ('pciDeviceId', c_uint),
-
+        ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE),
+        ("domain", c_uint),
+        ("bus", c_uint),
+        ("device", c_uint),
+        ("pciDeviceId", c_uint),
         # Added in 2.285
-        ('pciSubSystemId', c_uint),
-        ('reserved0', c_uint),
-        ('reserved1', c_uint),
-        ('reserved2', c_uint),
-        ('reserved3', c_uint),
+        ("pciSubSystemId", c_uint),
+        ("reserved0", c_uint),
+        ("reserved1", c_uint),
+        ("reserved2", c_uint),
+        ("reserved3", c_uint),
     ]
     _fmt_ = {
-            'domain'         : "0x%04X",
-            'bus'            : "0x%02X",
-            'device'         : "0x%02X",
-            'pciDeviceId'    : "0x%08X",
-            'pciSubSystemId' : "0x%08X",
-            }
+        "domain": "0x%04X",
+        "bus": "0x%02X",
+        "device": "0x%02X",
+        "pciDeviceId": "0x%08X",
+        "pciSubSystemId": "0x%08X",
+    }
+
 
 class nvmlPciInfo_t(_PrintableStructure):
     _fields_ = [
         # Moved to the new busId location below
-        ('busIdLegacy', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE),
-        ('domain', c_uint),
-        ('bus', c_uint),
-        ('device', c_uint),
-        ('pciDeviceId', c_uint),
-
+        ("busIdLegacy", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE),
+        ("domain", c_uint),
+        ("bus", c_uint),
+        ("device", c_uint),
+        ("pciDeviceId", c_uint),
         # Added in 2.285
-        ('pciSubSystemId', c_uint),
+        ("pciSubSystemId", c_uint),
         # New busId replaced the long deprecated and reserved fields with a
         # field of the same size in 9.0
-        ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE),
+        ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE),
     ]
     _fmt_ = {
-            'domain'         : "0x%08X",
-            'bus'            : "0x%02X",
-            'device'         : "0x%02X",
-            'pciDeviceId'    : "0x%08X",
-            'pciSubSystemId' : "0x%08X",
-            }
+        "domain": "0x%08X",
+        "bus": "0x%02X",
+        "device": "0x%02X",
+        "pciDeviceId": "0x%08X",
+        "pciSubSystemId": "0x%08X",
+    }
+
 
 class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
+        ("version", c_uint),
         ("branch", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
     ]
 
+
 SystemDriverBranchInfo_v1 = 0x1000054
 
+
 class c_nvmlExcludedDeviceInfo_t(_PrintableStructure):
-    _fields_ = [
-        ('pci', nvmlPciInfo_t),
-        ('uuid', c_char * NVML_DEVICE_UUID_BUFFER_SIZE)
-    ]
+    _fields_ = [("pci", nvmlPciInfo_t), ("uuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE)]
+
 
 class nvmlNvLinkUtilizationControl_t(_PrintableStructure):
     _fields_ = [
-        ('units', _nvmlNvLinkUtilizationCountUnits_t),
-        ('pktfilter', _nvmlNvLinkUtilizationCountPktTypes_t),
+        ("units", _nvmlNvLinkUtilizationCountUnits_t),
+        ("pktfilter", _nvmlNvLinkUtilizationCountPktTypes_t),
     ]
 
+
 class c_nvmlMemory_t(_PrintableStructure):
     _fields_ = [
-        ('total', c_ulonglong),
-        ('free', c_ulonglong),
-        ('used', c_ulonglong),
+        ("total", c_ulonglong),
+        ("free", c_ulonglong),
+        ("used", c_ulonglong),
     ]
-    _fmt_ = {'<default>': "%d B"}
+    _fmt_ = {"<default>": "%d B"}
+
 
 class c_nvmlMemory_v2_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('total', c_ulonglong),
-        ('reserved', c_ulonglong),
-        ('free', c_ulonglong),
-        ('used', c_ulonglong),
+        ("version", c_uint),
+        ("total", c_ulonglong),
+        ("reserved", c_ulonglong),
+        ("free", c_ulonglong),
+        ("used", c_ulonglong),
     ]
-    _fmt_ = {'<default>': "%d B"}
+    _fmt_ = {"<default>": "%d B"}
+
 
 nvmlMemory_v2 = 0x02000028
 
+
 class c_nvmlBAR1Memory_t(_PrintableStructure):
     _fields_ = [
-        ('bar1Total', c_ulonglong),
-        ('bar1Free', c_ulonglong),
-        ('bar1Used', c_ulonglong),
+        ("bar1Total", c_ulonglong),
+        ("bar1Free", c_ulonglong),
+        ("bar1Used", c_ulonglong),
     ]
-    _fmt_ = {'<default>': "%d B"}
+    _fmt_ = {"<default>": "%d B"}
+
 
 class nvmlClkMonFaultInfo_t(Structure):
-    _fields_ = [("clkApiDomain", c_uint),
-                ("clkDomainFaultMask", c_uint)
-    ]
+    _fields_ = [("clkApiDomain", c_uint), ("clkDomainFaultMask", c_uint)]
+
 
 MAX_CLK_DOMAINS = 32
 
+
 class nvmlClkMonStatus_t(Structure):
-    _fields_ = [("bGlobalStatus", c_uint),
-                ("clkMonListSize", c_uint),
-                ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS)
+    _fields_ = [
+        ("bGlobalStatus", c_uint),
+        ("clkMonListSize", c_uint),
+        ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS),
     ]
 
+
 # On Windows with the WDDM driver, usedGpuMemory is reported as None
 # Code that processes this structure should check for None, I.E.
 #
@@ -1351,887 +1638,1031 @@ class nvmlClkMonStatus_t(Structure):
 # See NVML documentation for more information
 class c_nvmlProcessInfo_v2_t(_PrintableStructure):
     _fields_ = [
-        ('pid', c_uint),
-        ('usedGpuMemory', c_ulonglong),
-        ('gpuInstanceId', c_uint),
-        ('computeInstanceId', c_uint),
+        ("pid", c_uint),
+        ("usedGpuMemory", c_ulonglong),
+        ("gpuInstanceId", c_uint),
+        ("computeInstanceId", c_uint),
     ]
-    _fmt_ = {'usedGpuMemory': "%d B"}
+    _fmt_ = {"usedGpuMemory": "%d B"}
+
 
 c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t
 
 c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t
 
 _nvmlProcessMode_t = c_uint
-NVML_PROCESS_MODE_COMPUTE  = 0
+NVML_PROCESS_MODE_COMPUTE = 0
 NVML_PROCESS_MODE_GRAPHICS = 1
-NVML_PROCESS_MODE_MPS      = 2
+NVML_PROCESS_MODE_MPS = 2
+
 
 class c_nvmlProcessDetail_v1_t(Structure):
     _fields_ = [
-        ('pid', c_uint),
-        ('usedGpuMemory', c_ulonglong),
-        ('gpuInstanceId', c_uint),
-        ('computeInstanceId', c_uint),
-        ('usedGpuCcProtectedMemory', c_ulonglong),
+        ("pid", c_uint),
+        ("usedGpuMemory", c_ulonglong),
+        ("gpuInstanceId", c_uint),
+        ("computeInstanceId", c_uint),
+        ("usedGpuCcProtectedMemory", c_ulonglong),
     ]
 
+
 class c_nvmlProcessDetailList_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('mode', _nvmlProcessMode_t),
-        ('numProcArrayEntries', c_uint),
-        ('procArray', POINTER(c_nvmlProcessDetail_v1_t)),
+        ("version", c_uint),
+        ("mode", _nvmlProcessMode_t),
+        ("numProcArrayEntries", c_uint),
+        ("procArray", POINTER(c_nvmlProcessDetail_v1_t)),
     ]
-    _fmt_ = {'numProcArrayEntries': "%d B"}
+    _fmt_ = {"numProcArrayEntries": "%d B"}
+
 
 c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t
 
 nvmlProcessDetailList_v1 = 0x1000018
 
+
 class c_nvmlBridgeChipInfo_t(_PrintableStructure):
     _fields_ = [
-        ('type', _nvmlBridgeChipType_t),
-        ('fwVersion', c_uint),
+        ("type", _nvmlBridgeChipType_t),
+        ("fwVersion", c_uint),
     ]
 
+
 class c_nvmlBridgeChipHierarchy_t(_PrintableStructure):
     _fields_ = [
-        ('bridgeCount', c_uint),
-        ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128),
+        ("bridgeCount", c_uint),
+        ("bridgeChipInfo", c_nvmlBridgeChipInfo_t * 128),
     ]
 
+
 class c_nvmlEccErrorCounts_t(_PrintableStructure):
     _fields_ = [
-        ('l1Cache', c_ulonglong),
-        ('l2Cache', c_ulonglong),
-        ('deviceMemory', c_ulonglong),
-        ('registerFile', c_ulonglong),
+        ("l1Cache", c_ulonglong),
+        ("l2Cache", c_ulonglong),
+        ("deviceMemory", c_ulonglong),
+        ("registerFile", c_ulonglong),
     ]
 
+
 class c_nvmlUtilization_t(_PrintableStructure):
     _fields_ = [
-        ('gpu', c_uint),
-        ('memory', c_uint),
+        ("gpu", c_uint),
+        ("memory", c_uint),
     ]
-    _fmt_ = {'<default>': "%d %%"}
+    _fmt_ = {"<default>": "%d %%"}
+
 
 # Added in 2.285
 class c_nvmlHwbcEntry_t(_PrintableStructure):
     _fields_ = [
-        ('hwbcId', c_uint),
-        ('firmwareVersion', c_char * 32),
+        ("hwbcId", c_uint),
+        ("firmwareVersion", c_char * 32),
     ]
 
+
 class c_nvmlValue_t(Union):
     _fields_ = [
-        ('dVal', c_double),
-        ('uiVal', c_uint),
-        ('ulVal', c_ulong),
-        ('ullVal', c_ulonglong),
-        ('sllVal', c_longlong),
-        ('siVal', c_int),
-        ('usVal', c_ushort),
+        ("dVal", c_double),
+        ("uiVal", c_uint),
+        ("ulVal", c_ulong),
+        ("ullVal", c_ulonglong),
+        ("sllVal", c_longlong),
+        ("siVal", c_int),
+        ("usVal", c_ushort),
     ]
 
+
 class c_nvmlSample_t(_PrintableStructure):
     _fields_ = [
-        ('timeStamp', c_ulonglong),
-        ('sampleValue', c_nvmlValue_t),
+        ("timeStamp", c_ulonglong),
+        ("sampleValue", c_nvmlValue_t),
     ]
 
+
 class c_nvmlViolationTime_t(_PrintableStructure):
     _fields_ = [
-        ('referenceTime', c_ulonglong),
-        ('violationTime', c_ulonglong),
+        ("referenceTime", c_ulonglong),
+        ("violationTime", c_ulonglong),
     ]
 
+
 class c_nvmlFieldValue_t(_PrintableStructure):
     _fields_ = [
-        ('fieldId', c_uint32),
-        ('scopeId', c_uint32),
-        ('timestamp', c_int64),
-        ('latencyUsec', c_int64),
-        ('valueType', _nvmlValueType_t),
-        ('nvmlReturn', _nvmlReturn_t),
-        ('value', c_nvmlValue_t)
+        ("fieldId", c_uint32),
+        ("scopeId", c_uint32),
+        ("timestamp", c_int64),
+        ("latencyUsec", c_int64),
+        ("valueType", _nvmlValueType_t),
+        ("nvmlReturn", _nvmlReturn_t),
+        ("value", c_nvmlValue_t),
     ]
 
+
 NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES = 23
 
-nvmlNvlinkSupportedBwModes_v1 = 0x100001c
+nvmlNvlinkSupportedBwModes_v1 = 0x100001C
+
+
 class c_nvmlNvlinkSupportedBwModes_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('bwModes', c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES),
-        ('totalBwModes', c_uint8)
+        ("version", c_uint),
+        ("bwModes", c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES),
+        ("totalBwModes", c_uint8),
     ]
 
     def __init__(self):
-        super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__(version=nvmlNvlinkSupportedBwModes_v1)
+        super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__(
+            version=nvmlNvlinkSupportedBwModes_v1
+        )
+
+
+nvmlNvlinkGetBwMode_v1 = 0x100000C
+
 
-nvmlNvlinkGetBwMode_v1 = 0x100000c
 class c_nvmlNvlinkGetBwMode_v1_t(_PrintableStructure):
-    _fields_ = [
-        ('version', c_uint),
-        ('bIsBest', c_uint),
-        ('bwMode', c_uint8)
-    ]
+    _fields_ = [("version", c_uint), ("bIsBest", c_uint), ("bwMode", c_uint8)]
 
     def __init__(self):
         super(c_nvmlNvlinkGetBwMode_v1_t, self).__init__(version=nvmlNvlinkGetBwMode_v1)
 
-nvmlNvlinkSetBwMode_v1 = 0x100000c
+
+nvmlNvlinkSetBwMode_v1 = 0x100000C
+
+
 class c_nvmlNvlinkSetBwMode_v1_t(_PrintableStructure):
-    _fields_ = [
-        ('version', c_uint),
-        ('bSetBest', c_uint),
-        ('bwMode', c_uint8)
-    ]
+    _fields_ = [("version", c_uint), ("bSetBest", c_uint), ("bwMode", c_uint8)]
 
     def __init__(self):
         super(c_nvmlNvlinkSetBwMode_v1_t, self).__init__(version=nvmlNvlinkSetBwMode_v1)
 
+
 class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('mode', c_uint),
+        ("version", c_uint),
+        ("mode", c_uint),
     ]
 
+
 VgpuHeterogeneousMode_v1 = 0x1000008
 
+
 class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('placementId', c_uint),
+        ("version", c_uint),
+        ("placementId", c_uint),
     ]
 
+
 VgpuPlacementId_v1 = 0x1000008
 
+
 class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('count', c_uint),
-        ('placementSize', c_uint),
-        ('placementIds', POINTER(c_uint)),
+        ("version", c_uint),
+        ("count", c_uint),
+        ("placementSize", c_uint),
+        ("placementIds", POINTER(c_uint)),
     ]
 
+
 VgpuPlacementList_v1 = 0x1000018
 
-NVML_VGPU_PGPU_HETEROGENEOUS_MODE   = 0
-NVML_VGPU_PGPU_HOMOGENEOUS_MODE     = 1
+NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0
+NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1
+
 
 class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('placementSize', c_uint),
-        ('count', c_uint),
-        ('placementIds', POINTER(c_uint)),
-        ('mode', c_uint),
+        ("version", c_uint),
+        ("placementSize", c_uint),
+        ("count", c_uint),
+        ("placementIds", POINTER(c_uint)),
+        ("mode", c_uint),
     ]
 
+
 VgpuPlacementList_v2 = 0x2000020
 
+
 class c_nvmlVgpuTypeBar1Info_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('bar1Size', c_ulonglong),
+        ("version", c_uint),
+        ("bar1Size", c_ulonglong),
     ]
 
+
 VgpuTypeBar1Info_v1 = 0x1000010
 
+
 class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure):
     _fields_ = [
-        ('vgpuInstance', _nvmlVgpuInstance_t),
-        ('timeStamp', c_ulonglong),
-        ('smUtil', c_nvmlValue_t),
-        ('memUtil', c_nvmlValue_t),
-        ('encUtil', c_nvmlValue_t),
-        ('decUtil', c_nvmlValue_t),
+        ("vgpuInstance", _nvmlVgpuInstance_t),
+        ("timeStamp", c_ulonglong),
+        ("smUtil", c_nvmlValue_t),
+        ("memUtil", c_nvmlValue_t),
+        ("encUtil", c_nvmlValue_t),
+        ("decUtil", c_nvmlValue_t),
     ]
 
+
 class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('timeStamp', c_ulonglong),
-        ('vgpuInstance', _nvmlVgpuInstance_t),
-        ('smUtil', c_nvmlValue_t),
-        ('memUtil', c_nvmlValue_t),
-        ('encUtil', c_nvmlValue_t),
-        ('decUtil', c_nvmlValue_t),
-        ('jpgUtil', c_nvmlValue_t),
-        ('ofaUtil', c_nvmlValue_t),
+        ("timeStamp", c_ulonglong),
+        ("vgpuInstance", _nvmlVgpuInstance_t),
+        ("smUtil", c_nvmlValue_t),
+        ("memUtil", c_nvmlValue_t),
+        ("encUtil", c_nvmlValue_t),
+        ("decUtil", c_nvmlValue_t),
+        ("jpgUtil", c_nvmlValue_t),
+        ("ofaUtil", c_nvmlValue_t),
     ]
 
+
 class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('sampleValType', _nvmlValueType_t),
-        ('vgpuInstanceCount', c_uint),
-        ('lastSeenTimeStamp', c_ulonglong),
-        ('vgpuUtilArray', POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)),
+        ("version", c_uint),
+        ("sampleValType", _nvmlValueType_t),
+        ("vgpuInstanceCount", c_uint),
+        ("lastSeenTimeStamp", c_ulonglong),
+        ("vgpuUtilArray", POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)),
     ]
 
+
 VgpuInstancesUtilizationInfo_v1 = 0x01000020
 
+
 class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure):
     _fields_ = [
-        ('vgpuInstance', _nvmlVgpuInstance_t),
-        ('pid', c_uint),
-        ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE),
-        ('timeStamp', c_ulonglong),
-        ('smUtil', c_uint),
-        ('memUtil', c_uint),
-        ('encUtil', c_uint),
-        ('decUtil', c_uint),
+        ("vgpuInstance", _nvmlVgpuInstance_t),
+        ("pid", c_uint),
+        ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE),
+        ("timeStamp", c_ulonglong),
+        ("smUtil", c_uint),
+        ("memUtil", c_uint),
+        ("encUtil", c_uint),
+        ("decUtil", c_uint),
     ]
 
+
 class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE),
-        ('timeStamp', c_ulonglong),
-        ('vgpuInstance', _nvmlVgpuInstance_t),
-        ('pid', c_uint),
-        ('smUtil', c_uint),
-        ('memUtil', c_uint),
-        ('encUtil', c_uint),
-        ('decUtil', c_uint),
-        ('jpgUtil', c_uint),
-        ('ofaUtil', c_uint),
+        ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE),
+        ("timeStamp", c_ulonglong),
+        ("vgpuInstance", _nvmlVgpuInstance_t),
+        ("pid", c_uint),
+        ("smUtil", c_uint),
+        ("memUtil", c_uint),
+        ("encUtil", c_uint),
+        ("decUtil", c_uint),
+        ("jpgUtil", c_uint),
+        ("ofaUtil", c_uint),
     ]
 
+
 class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('vgpuProcessCount', c_uint),
-        ('lastSeenTimeStamp', c_ulonglong),
-        ('vgpuProcUtilArray', POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)),
+        ("version", c_uint),
+        ("vgpuProcessCount", c_uint),
+        ("lastSeenTimeStamp", c_ulonglong),
+        ("vgpuProcUtilArray", POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)),
     ]
 
+
 VgpuProcessesUtilizationInfo_v1 = 0x01000018
 
+
 class nvmlVgpuRuntimeState_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('size', c_ulonglong),
+        ("version", c_uint),
+        ("size", c_ulonglong),
     ]
 
+
 VgpuRuntimeState_v1 = 0x1000010
 
+
 class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure):
     _fields_ = [
-        ('year',    c_uint32),
-        ('month',   c_uint16),
-        ('day',     c_uint16),
-        ('hour',    c_uint16),
-        ('min',     c_uint16),
-        ('sec',     c_uint16),
-        ('status',  c_uint8),
+        ("year", c_uint32),
+        ("month", c_uint16),
+        ("day", c_uint16),
+        ("hour", c_uint16),
+        ("min", c_uint16),
+        ("sec", c_uint16),
+        ("status", c_uint8),
     ]
 
-NVML_GRID_LICENSE_STATE_UNKNOWN                 = 0
-NVML_GRID_LICENSE_STATE_UNINITIALIZED           = 1
+
+NVML_GRID_LICENSE_STATE_UNKNOWN = 0
+NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1
 NVML_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2
-NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED   = 3
-NVML_GRID_LICENSE_STATE_UNLICENSED              = 4
-NVML_GRID_LICENSE_STATE_LICENSED                = 5
+NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3
+NVML_GRID_LICENSE_STATE_UNLICENSED = 4
+NVML_GRID_LICENSE_STATE_LICENSED = 5
+
 
 class c_nvmlVgpuLicenseInfo_t(_PrintableStructure):
     _fields_ = [
-        ('isLicensed',      c_uint8),
-        ('licenseExpiry',   c_nvmlVgpuLicenseExpiry_t),
-        ('currentState',    c_uint),
+        ("isLicensed", c_uint8),
+        ("licenseExpiry", c_nvmlVgpuLicenseExpiry_t),
+        ("currentState", c_uint),
     ]
 
+
 class c_nvmlEncoderSession_t(_PrintableStructure):
     _fields_ = [
-        ('sessionId', c_uint),
-        ('pid', c_uint),
-        ('vgpuInstance', _nvmlVgpuInstance_t),
-        ('codecType', c_uint),
-        ('hResolution', c_uint),
-        ('vResolution', c_uint),
-        ('averageFps', c_uint),
-        ('encodeLatency', c_uint),
+        ("sessionId", c_uint),
+        ("pid", c_uint),
+        ("vgpuInstance", _nvmlVgpuInstance_t),
+        ("codecType", c_uint),
+        ("hResolution", c_uint),
+        ("vResolution", c_uint),
+        ("averageFps", c_uint),
+        ("encodeLatency", c_uint),
     ]
 
+
 class c_nvmlProcessUtilizationSample_t(_PrintableStructure):
     _fields_ = [
-        ('pid', c_uint),
-        ('timeStamp', c_ulonglong),
-        ('smUtil', c_uint),
-        ('memUtil', c_uint),
-        ('encUtil', c_uint),
-        ('decUtil', c_uint),
+        ("pid", c_uint),
+        ("timeStamp", c_ulonglong),
+        ("smUtil", c_uint),
+        ("memUtil", c_uint),
+        ("encUtil", c_uint),
+        ("decUtil", c_uint),
     ]
 
+
 class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('timeStamp', c_ulonglong),
-        ('pid', c_uint),
-        ('smUtil', c_uint),
-        ('memUtil', c_uint),
-        ('encUtil', c_uint),
-        ('decUtil', c_uint),
-        ('jpgUtil', c_uint),
-        ('ofaUtil', c_uint),
+        ("timeStamp", c_ulonglong),
+        ("pid", c_uint),
+        ("smUtil", c_uint),
+        ("memUtil", c_uint),
+        ("encUtil", c_uint),
+        ("decUtil", c_uint),
+        ("jpgUtil", c_uint),
+        ("ofaUtil", c_uint),
     ]
 
+
 class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('processSamplesCount', c_uint),
-        ('lastSeenTimeStamp', c_ulonglong),
-        ('procUtilArray', POINTER(c_nvmlProcessUtilizationInfo_v1_t)),
+        ("version", c_uint),
+        ("processSamplesCount", c_uint),
+        ("lastSeenTimeStamp", c_ulonglong),
+        ("procUtilArray", POINTER(c_nvmlProcessUtilizationInfo_v1_t)),
     ]
 
+
 ProcessesUtilizationInfo_v1 = 0x01000018
 
+
 class c_nvmlGridLicenseExpiry_t(_PrintableStructure):
     _fields_ = [
-        ('year',    c_uint32),
-        ('month',   c_uint16),
-        ('day',     c_uint16),
-        ('hour',    c_uint16),
-        ('min',     c_uint16),
-        ('sec',     c_uint16),
-        ('status',  c_uint8),
+        ("year", c_uint32),
+        ("month", c_uint16),
+        ("day", c_uint16),
+        ("hour", c_uint16),
+        ("min", c_uint16),
+        ("sec", c_uint16),
+        ("status", c_uint8),
     ]
 
+
 class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure):
     _fields_ = [
-        ('featureCode',    _nvmlGridLicenseFeatureCode_t),
-        ('featureState',   c_uint),
-        ('licenseInfo',    c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
-        ('productName',    c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
-        ('featureEnabled', c_uint),
-        ('licenseExpiry',  c_nvmlGridLicenseExpiry_t),
+        ("featureCode", _nvmlGridLicenseFeatureCode_t),
+        ("featureState", c_uint),
+        ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
+        ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
+        ("featureEnabled", c_uint),
+        ("licenseExpiry", c_nvmlGridLicenseExpiry_t),
     ]
 
+
 class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure):
     _fields_ = [
-        ('isGridLicenseSupported',  c_int),
-        ('licensableFeaturesCount', c_uint),
-        ('gridLicensableFeatures',  c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT),
+        ("isGridLicenseSupported", c_int),
+        ("licensableFeaturesCount", c_uint),
+        (
+            "gridLicensableFeatures",
+            c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT,
+        ),
     ]
 
+
 class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure):
     _fields_ = [
-        ('featureCode', _nvmlGridLicenseFeatureCode_t),
-        ('featureState', c_uint),
-        ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
-        ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
-        ('featureEnabled', c_uint),
+        ("featureCode", _nvmlGridLicenseFeatureCode_t),
+        ("featureState", c_uint),
+        ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
+        ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
+        ("featureEnabled", c_uint),
     ]
 
+
 class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure):
     _fields_ = [
-        ('isGridLicenseSupported', c_int),
-        ('licensableFeaturesCount', c_uint),
-        ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT),
+        ("isGridLicenseSupported", c_int),
+        ("licensableFeaturesCount", c_uint),
+        (
+            "gridLicensableFeatures",
+            c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT,
+        ),
     ]
 
+
 class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure):
     _fields_ = [
-        ('featureCode', _nvmlGridLicenseFeatureCode_t),
-        ('featureState', c_uint),
-        ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
-        ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
+        ("featureCode", _nvmlGridLicenseFeatureCode_t),
+        ("featureState", c_uint),
+        ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
+        ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
     ]
 
+
 class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure):
     _fields_ = [
-        ('isGridLicenseSupported', c_int),
-        ('licensableFeaturesCount', c_uint),
-        ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT),
+        ("isGridLicenseSupported", c_int),
+        ("licensableFeaturesCount", c_uint),
+        (
+            "gridLicensableFeatures",
+            c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT,
+        ),
     ]
 
+
 class c_nvmlGridLicensableFeature_t(_PrintableStructure):
     _fields_ = [
-        ('featureCode', _nvmlGridLicenseFeatureCode_t),
-        ('featureState', c_uint),
-        ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
+        ("featureCode", _nvmlGridLicenseFeatureCode_t),
+        ("featureState", c_uint),
+        ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE),
     ]
 
+
 class c_nvmlGridLicensableFeatures_t(_PrintableStructure):
     _fields_ = [
-        ('isGridLicenseSupported', c_int),
-        ('licensableFeaturesCount', c_uint),
-        ('gridLicensableFeatures', c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT),
+        ("isGridLicenseSupported", c_int),
+        ("licensableFeaturesCount", c_uint),
+        (
+            "gridLicensableFeatures",
+            c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT,
+        ),
     ]
 
+
 class c_nvmlMarginTemperature_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('marginTemperature', c_int),
+        ("version", c_uint),
+        ("marginTemperature", c_int),
     ]
 
+
 nvmlMarginTemperature_v1 = 0x1000008
 
+
 ## Event structures
 class struct_c_nvmlEventSet_t(Structure):
-    pass # opaque handle
+    pass  # opaque handle
+
+
 c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t)
 
-nvmlEventTypeSingleBitEccError      = 0x0000000000000001
-nvmlEventTypeDoubleBitEccError      = 0x0000000000000002
-nvmlEventTypePState                 = 0x0000000000000004
-nvmlEventTypeXidCriticalError       = 0x0000000000000008
-nvmlEventTypeClock                  = 0x0000000000000010
-nvmlEventTypePowerSourceChange      = 0x0000000000000080
-nvmlEventMigConfigChange            = 0x0000000000000100
+nvmlEventTypeSingleBitEccError = 0x0000000000000001
+nvmlEventTypeDoubleBitEccError = 0x0000000000000002
+nvmlEventTypePState = 0x0000000000000004
+nvmlEventTypeXidCriticalError = 0x0000000000000008
+nvmlEventTypeClock = 0x0000000000000010
+nvmlEventTypePowerSourceChange = 0x0000000000000080
+nvmlEventMigConfigChange = 0x0000000000000100
 nvmlEventTypeSingleBitEccErrorStorm = 0x0000000000000200
-nvmlEventTypeDramRetirementEvent    = 0x0000000000000400
-nvmlEventTypeDramRetirementFailure  = 0x0000000000000800
-nvmlEventTypeNonFatalPoisonError    = 0x0000000000001000
-nvmlEventTypeFatalPoisonError       = 0x0000000000002000
-nvmlEventTypeGpuUnavailableError    = 0x0000000000004000
-nvmlEventTypeGpuRecoveryAction      = 0x0000000000008000
-nvmlEventTypeNone                   = 0x0000000000000000
-nvmlEventTypeAll                    = (
-                                        nvmlEventTypeNone
-                                        | nvmlEventTypeSingleBitEccError
-                                        | nvmlEventTypeDoubleBitEccError
-                                        | nvmlEventTypePState
-                                        | nvmlEventTypeClock
-                                        | nvmlEventTypePowerSourceChange
-                                        | nvmlEventTypeXidCriticalError
-                                        | nvmlEventMigConfigChange
-                                        | nvmlEventTypeSingleBitEccErrorStorm
-                                        | nvmlEventTypeDramRetirementEvent
-                                        | nvmlEventTypeDramRetirementFailure
-                                        | nvmlEventTypeNonFatalPoisonError
-                                        | nvmlEventTypeFatalPoisonError
-                                        | nvmlEventTypeGpuUnavailableError
-                                        | nvmlEventTypeGpuRecoveryAction
-                                        )
+nvmlEventTypeDramRetirementEvent = 0x0000000000000400
+nvmlEventTypeDramRetirementFailure = 0x0000000000000800
+nvmlEventTypeNonFatalPoisonError = 0x0000000000001000
+nvmlEventTypeFatalPoisonError = 0x0000000000002000
+nvmlEventTypeGpuUnavailableError = 0x0000000000004000
+nvmlEventTypeGpuRecoveryAction = 0x0000000000008000
+nvmlEventTypeNone = 0x0000000000000000
+nvmlEventTypeAll = (
+    nvmlEventTypeNone
+    | nvmlEventTypeSingleBitEccError
+    | nvmlEventTypeDoubleBitEccError
+    | nvmlEventTypePState
+    | nvmlEventTypeClock
+    | nvmlEventTypePowerSourceChange
+    | nvmlEventTypeXidCriticalError
+    | nvmlEventMigConfigChange
+    | nvmlEventTypeSingleBitEccErrorStorm
+    | nvmlEventTypeDramRetirementEvent
+    | nvmlEventTypeDramRetirementFailure
+    | nvmlEventTypeNonFatalPoisonError
+    | nvmlEventTypeFatalPoisonError
+    | nvmlEventTypeGpuUnavailableError
+    | nvmlEventTypeGpuRecoveryAction
+)
 
 ## Clock Event Reasons defines
-nvmlClocksEventReasonGpuIdle              = 0x0000000000000001
+nvmlClocksEventReasonGpuIdle = 0x0000000000000001
 nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002
-nvmlClocksEventReasonUserDefinedClocks         = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting
-nvmlClocksEventReasonSwPowerCap           = 0x0000000000000004
-nvmlClocksEventReasonHwSlowdown           = 0x0000000000000008
-nvmlClocksEventReasonSyncBoost            = 0x0000000000000010
-nvmlClocksEventReasonSwThermalSlowdown    = 0x0000000000000020
-nvmlClocksEventReasonHwThermalSlowdown    = 0x0000000000000040
+nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting  # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting
+nvmlClocksEventReasonSwPowerCap = 0x0000000000000004
+nvmlClocksEventReasonHwSlowdown = 0x0000000000000008
+nvmlClocksEventReasonSyncBoost = 0x0000000000000010
+nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020
+nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040
 nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080
-nvmlClocksEventReasonDisplayClockSetting  = 0x0000000000000100
-nvmlClocksEventReasonNone                 = 0x0000000000000000
-nvmlClocksEventReasonAll                  = (
-                                                  nvmlClocksEventReasonNone |
-                                                  nvmlClocksEventReasonGpuIdle |
-                                                  nvmlClocksEventReasonApplicationsClocksSetting |
-                                                  nvmlClocksEventReasonSwPowerCap |
-                                                  nvmlClocksEventReasonHwSlowdown |
-                                                  nvmlClocksEventReasonSyncBoost |
-                                                  nvmlClocksEventReasonSwThermalSlowdown |
-                                                  nvmlClocksEventReasonHwThermalSlowdown |
-                                                  nvmlClocksEventReasonHwPowerBrakeSlowdown |
-                                                  nvmlClocksEventReasonDisplayClockSetting
-                                               )
+nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100
+nvmlClocksEventReasonNone = 0x0000000000000000
+nvmlClocksEventReasonAll = (
+    nvmlClocksEventReasonNone
+    | nvmlClocksEventReasonGpuIdle
+    | nvmlClocksEventReasonApplicationsClocksSetting
+    | nvmlClocksEventReasonSwPowerCap
+    | nvmlClocksEventReasonHwSlowdown
+    | nvmlClocksEventReasonSyncBoost
+    | nvmlClocksEventReasonSwThermalSlowdown
+    | nvmlClocksEventReasonHwThermalSlowdown
+    | nvmlClocksEventReasonHwPowerBrakeSlowdown
+    | nvmlClocksEventReasonDisplayClockSetting
+)
 
 ## Following have been deprecated
-nvmlClocksThrottleReasonGpuIdle              = 0x0000000000000001
+nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001
 nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002
-nvmlClocksThrottleReasonUserDefinedClocks         = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting
-nvmlClocksThrottleReasonSwPowerCap           = 0x0000000000000004
-nvmlClocksThrottleReasonHwSlowdown           = 0x0000000000000008
-nvmlClocksThrottleReasonSyncBoost            = 0x0000000000000010
-nvmlClocksThrottleReasonSwThermalSlowdown    = 0x0000000000000020
-nvmlClocksThrottleReasonHwThermalSlowdown    = 0x0000000000000040
+nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting  # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting
+nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004
+nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008
+nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010
+nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020
+nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040
 nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080
-nvmlClocksThrottleReasonDisplayClockSetting  = 0x0000000000000100
-nvmlClocksThrottleReasonNone                 = 0x0000000000000000
-nvmlClocksThrottleReasonAll                  = (
-                                                  nvmlClocksThrottleReasonNone |
-                                                  nvmlClocksThrottleReasonGpuIdle |
-                                                  nvmlClocksThrottleReasonApplicationsClocksSetting |
-                                                  nvmlClocksThrottleReasonSwPowerCap |
-                                                  nvmlClocksThrottleReasonHwSlowdown |
-                                                  nvmlClocksThrottleReasonSyncBoost |
-                                                  nvmlClocksThrottleReasonSwThermalSlowdown |
-                                                  nvmlClocksThrottleReasonHwThermalSlowdown |
-                                                  nvmlClocksThrottleReasonHwPowerBrakeSlowdown |
-                                                  nvmlClocksThrottleReasonDisplayClockSetting
-                                               )
+nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100
+nvmlClocksThrottleReasonNone = 0x0000000000000000
+nvmlClocksThrottleReasonAll = (
+    nvmlClocksThrottleReasonNone
+    | nvmlClocksThrottleReasonGpuIdle
+    | nvmlClocksThrottleReasonApplicationsClocksSetting
+    | nvmlClocksThrottleReasonSwPowerCap
+    | nvmlClocksThrottleReasonHwSlowdown
+    | nvmlClocksThrottleReasonSyncBoost
+    | nvmlClocksThrottleReasonSwThermalSlowdown
+    | nvmlClocksThrottleReasonHwThermalSlowdown
+    | nvmlClocksThrottleReasonHwPowerBrakeSlowdown
+    | nvmlClocksThrottleReasonDisplayClockSetting
+)
+
 
 class c_nvmlEventData_t(_PrintableStructure):
     _fields_ = [
-        ('device', c_nvmlDevice_t),
-        ('eventType', c_ulonglong),
-        ('eventData', c_ulonglong),
-        ('gpuInstanceId', c_uint),
-        ('computeInstanceId', c_uint)
+        ("device", c_nvmlDevice_t),
+        ("eventType", c_ulonglong),
+        ("eventData", c_ulonglong),
+        ("gpuInstanceId", c_uint),
+        ("computeInstanceId", c_uint),
     ]
-    _fmt_ = {'eventType': "0x%08X"}
+    _fmt_ = {"eventType": "0x%08X"}
+
 
 class c_nvmlAccountingStats_t(_PrintableStructure):
     _fields_ = [
-        ('gpuUtilization', c_uint),
-        ('memoryUtilization', c_uint),
-        ('maxMemoryUsage', c_ulonglong),
-        ('time', c_ulonglong),
-        ('startTime', c_ulonglong),
-        ('isRunning', c_uint),
-        ('reserved', c_uint * 5)
+        ("gpuUtilization", c_uint),
+        ("memoryUtilization", c_uint),
+        ("maxMemoryUsage", c_ulonglong),
+        ("time", c_ulonglong),
+        ("startTime", c_ulonglong),
+        ("isRunning", c_uint),
+        ("reserved", c_uint * 5),
     ]
 
+
 class c_nvmlVgpuVersion_t(Structure):
-    _fields_ = [("minVersion", c_uint),
-                ("maxVersion", c_uint)
-               ]
+    _fields_ = [("minVersion", c_uint), ("maxVersion", c_uint)]
+
 
 class c_nvmlVgpuMetadata_t(_PrintableStructure):
-    _fields_ = [("version", c_uint),
-                ("revision", c_uint),
-                ("guestInfoState", _nvmlVgpuGuestInfoState_t),
-                ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
-                ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
-                ("reserved", c_uint * 6),
-                ("vgpuVirtualizationCaps", c_uint),
-                ("guestVgpuVersion", c_uint),
-                ("opaqueDataSize", c_uint),
-                ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE)
-               ]
+    _fields_ = [
+        ("version", c_uint),
+        ("revision", c_uint),
+        ("guestInfoState", _nvmlVgpuGuestInfoState_t),
+        ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
+        ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
+        ("reserved", c_uint * 6),
+        ("vgpuVirtualizationCaps", c_uint),
+        ("guestVgpuVersion", c_uint),
+        ("opaqueDataSize", c_uint),
+        ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE),
+    ]
+
 
 class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure):
-    _fields_ = [("version", c_uint),
-                ("revision", c_uint),
-                ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
-                ("pgpuVirtualizationCaps", c_uint),
-                ("reserved", c_uint * 5),
-                ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t),
-                ("opaqueDataSize", c_uint),
-                ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE)
-               ]
+    _fields_ = [
+        ("version", c_uint),
+        ("revision", c_uint),
+        ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
+        ("pgpuVirtualizationCaps", c_uint),
+        ("reserved", c_uint * 5),
+        ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t),
+        ("opaqueDataSize", c_uint),
+        ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE),
+    ]
+
 
 class c_nvmlVgpuPgpuCompatibility_t(Structure):
-    _fields_ = [("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t),
-                ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t)
-               ]
+    _fields_ = [
+        ("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t),
+        ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t),
+    ]
+
 
 ## vGPU scheduler policy defines
-NVML_VGPU_SCHEDULER_POLICY_UNKNOWN      = 0
-NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT  = 1
-NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE  = 2
-NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE  = 3
+NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0
+NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1
+NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2
+NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3
 
 ## Supported vGPU scheduler policy count
-NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT  = 3
+NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3
 
-NVML_SCHEDULER_SW_MAX_LOG_ENTRIES           = 200
+NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200
+
+NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0
+NVML_VGPU_SCHEDULER_ARR_DISABLE = 1
+NVML_VGPU_SCHEDULER_ARR_ENABLE = 2
 
-NVML_VGPU_SCHEDULER_ARR_DEFAULT   = 0
-NVML_VGPU_SCHEDULER_ARR_DISABLE   = 1
-NVML_VGPU_SCHEDULER_ARR_ENABLE    = 2
 
 class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure):
     _fields_ = [
-        ('avgFactor',   c_uint),
-        ('timeslice',   c_uint),
+        ("avgFactor", c_uint),
+        ("timeslice", c_uint),
     ]
 
+
 class c_nvmlVgpuSchedData_t(_PrintableStructure):
     _fields_ = [
-        ('timeslice',   c_uint),
+        ("timeslice", c_uint),
     ]
 
+
 class c_nvmlVgpuSchedulerParams_t(Union):
     _fields_ = [
-        ('vgpuSchedDataWithARR', c_nvmlVgpuSchedDataWithARR_t),
-        ('vgpuSchedData',        c_nvmlVgpuSchedData_t),
+        ("vgpuSchedDataWithARR", c_nvmlVgpuSchedDataWithARR_t),
+        ("vgpuSchedData", c_nvmlVgpuSchedData_t),
     ]
 
+
 class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure):
     _fields_ = [
-        ('timestamp',                   c_ulonglong),
-        ('timeRunTotal',                c_ulonglong),
-        ('timeRun',                     c_ulonglong),
-        ('swRunlistId',                 c_uint),
-        ('targetTimeSlice',             c_ulonglong),
-        ('cumulativePreemptionTime',    c_ulonglong),
+        ("timestamp", c_ulonglong),
+        ("timeRunTotal", c_ulonglong),
+        ("timeRun", c_ulonglong),
+        ("swRunlistId", c_uint),
+        ("targetTimeSlice", c_ulonglong),
+        ("cumulativePreemptionTime", c_ulonglong),
     ]
 
+
 class c_nvmlVgpuSchedulerLog_t(_PrintableStructure):
     _fields_ = [
-        ('engineId',        c_uint),
-        ('schedulerPolicy', c_uint),
-        ('arrMode',         c_uint),
-        ('schedulerParams', c_nvmlVgpuSchedulerParams_t),
-        ('entriesCount',    c_uint),
-        ('logEntries',      c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES),
+        ("engineId", c_uint),
+        ("schedulerPolicy", c_uint),
+        ("arrMode", c_uint),
+        ("schedulerParams", c_nvmlVgpuSchedulerParams_t),
+        ("entriesCount", c_uint),
+        (
+            "logEntries",
+            c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES,
+        ),
     ]
 
+
 class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure):
     _fields_ = [
-        ('schedulerPolicy', c_uint),
-        ('arrMode',         c_uint),
-        ('schedulerParams', c_nvmlVgpuSchedulerParams_t),
+        ("schedulerPolicy", c_uint),
+        ("arrMode", c_uint),
+        ("schedulerParams", c_nvmlVgpuSchedulerParams_t),
     ]
 
+
 class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure):
     _fields_ = [
-        ('avgFactor',   c_uint),
-        ('frequency',   c_uint),
+        ("avgFactor", c_uint),
+        ("frequency", c_uint),
     ]
 
+
 class c_nvmlVgpuSchedSetData_t(_PrintableStructure):
     _fields_ = [
-        ('timeslice',   c_uint),
+        ("timeslice", c_uint),
     ]
 
+
 class c_nvmlVgpuSchedulerSetParams_t(Union):
     _fields_ = [
-        ('vgpuSchedDataWithARR', c_nvmlVgpuSchedSetDataWithARR_t),
-        ('vgpuSchedData',        c_nvmlVgpuSchedSetData_t),
+        ("vgpuSchedDataWithARR", c_nvmlVgpuSchedSetDataWithARR_t),
+        ("vgpuSchedData", c_nvmlVgpuSchedSetData_t),
     ]
 
+
 class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure):
     _fields_ = [
-        ('schedulerPolicy', c_uint),
-        ('enableARRMode',   c_uint),
-        ('schedulerParams', c_nvmlVgpuSchedulerSetParams_t),
+        ("schedulerPolicy", c_uint),
+        ("enableARRMode", c_uint),
+        ("schedulerParams", c_nvmlVgpuSchedulerSetParams_t),
     ]
 
+
 class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure):
     _fields_ = [
-        ('supportedSchedulers', c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT),
-        ('maxTimeslice',        c_uint),
-        ('minTimeslice',        c_uint),
-        ('isArrModeSupported',  c_uint),
-        ('maxFrequencyForARR',  c_uint),
-        ('minFrequencyForARR',  c_uint),
-        ('maxAvgFactorForARR',  c_uint),
-        ('minAvgFactorForARR',  c_uint),
+        ("supportedSchedulers", c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT),
+        ("maxTimeslice", c_uint),
+        ("minTimeslice", c_uint),
+        ("isArrModeSupported", c_uint),
+        ("maxFrequencyForARR", c_uint),
+        ("minFrequencyForARR", c_uint),
+        ("maxAvgFactorForARR", c_uint),
+        ("minAvgFactorForARR", c_uint),
     ]
 
+
 class c_nvmlFBCStats_t(Structure):
-    _fields_ = [("sessionsCount", c_uint),
-                ("averageFPS", c_uint),
-                ("averageLatency", c_uint)
-               ]
+    _fields_ = [
+        ("sessionsCount", c_uint),
+        ("averageFPS", c_uint),
+        ("averageLatency", c_uint),
+    ]
+
 
 class c_nvmlFBCSession_t(_PrintableStructure):
     _fields_ = [
-        ('sessionId', c_uint),
-        ('pid', c_uint),
-        ('vgpuInstance', _nvmlVgpuInstance_t),
-        ('displayOrdinal', c_uint),
-        ('sessionType', c_uint),
-        ('sessionFlags', c_uint),
-        ('hMaxResolution', c_uint),
-        ('vMaxResolution', c_uint),
-        ('hResolution', c_uint),
-        ('vResolution', c_uint),
-        ('averageFPS', c_uint),
-        ('averageLatency', c_uint),
+        ("sessionId", c_uint),
+        ("pid", c_uint),
+        ("vgpuInstance", _nvmlVgpuInstance_t),
+        ("displayOrdinal", c_uint),
+        ("sessionType", c_uint),
+        ("sessionFlags", c_uint),
+        ("hMaxResolution", c_uint),
+        ("vMaxResolution", c_uint),
+        ("hResolution", c_uint),
+        ("vResolution", c_uint),
+        ("averageFPS", c_uint),
+        ("averageLatency", c_uint),
     ]
 
+
 NVML_DEVICE_MIG_DISABLE = 0x0
-NVML_DEVICE_MIG_ENABLE  = 0x1
-
-NVML_GPU_INSTANCE_PROFILE_1_SLICE      = 0x0
-NVML_GPU_INSTANCE_PROFILE_2_SLICE      = 0x1
-NVML_GPU_INSTANCE_PROFILE_3_SLICE      = 0x2
-NVML_GPU_INSTANCE_PROFILE_4_SLICE      = 0x3
-NVML_GPU_INSTANCE_PROFILE_7_SLICE      = 0x4
-NVML_GPU_INSTANCE_PROFILE_8_SLICE      = 0x5
-NVML_GPU_INSTANCE_PROFILE_6_SLICE      = 0x6
+NVML_DEVICE_MIG_ENABLE = 0x1
+
+NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0
+NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1
+NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2
+NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3
+NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4
+NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5
+NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6
 NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7
 NVML_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8
 NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9
-NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX  = 0xA
-NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX  = 0xB
-NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX  = 0xC
-NVML_GPU_INSTANCE_PROFILE_COUNT        = 0xD
+NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA
+NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB
+NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC
+NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD
+
 
 class c_nvmlGpuInstancePlacement_t(Structure):
-    _fields_ = [("start", c_uint),
-                ("size", c_uint)
-               ]
+    _fields_ = [("start", c_uint), ("size", c_uint)]
+
 
 class c_nvmlGpuInstanceProfileInfo_t(Structure):
-    _fields_ = [("id", c_uint),
-                ("isP2pSupported", c_uint),
-                ("sliceCount", c_uint),
-                ("instanceCount", c_uint),
-                ("multiprocessorCount", c_uint),
-                ("copyEngineCount", c_uint),
-                ("decoderCount", c_uint),
-                ("encoderCount", c_uint),
-                ("jpegCount", c_uint),
-                ("ofaCount", c_uint),
-                ("memorySizeMB", c_ulonglong),
-               ]
+    _fields_ = [
+        ("id", c_uint),
+        ("isP2pSupported", c_uint),
+        ("sliceCount", c_uint),
+        ("instanceCount", c_uint),
+        ("multiprocessorCount", c_uint),
+        ("copyEngineCount", c_uint),
+        ("decoderCount", c_uint),
+        ("encoderCount", c_uint),
+        ("jpegCount", c_uint),
+        ("ofaCount", c_uint),
+        ("memorySizeMB", c_ulonglong),
+    ]
+
 
 nvmlGpuInstanceProfileInfo_v2 = 0x02000098
 
+
 class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure):
-    _fields_ = [("version", c_uint),
-                ("id", c_uint),
-                ("isP2pSupported", c_uint),
-                ("sliceCount", c_uint),
-                ("instanceCount", c_uint),
-                ("multiprocessorCount", c_uint),
-                ("copyEngineCount", c_uint),
-                ("decoderCount", c_uint),
-                ("encoderCount", c_uint),
-                ("jpegCount", c_uint),
-                ("ofaCount", c_uint),
-                ("memorySizeMB", c_ulonglong),
-                ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE)
-               ]
-    
+    _fields_ = [
+        ("version", c_uint),
+        ("id", c_uint),
+        ("isP2pSupported", c_uint),
+        ("sliceCount", c_uint),
+        ("instanceCount", c_uint),
+        ("multiprocessorCount", c_uint),
+        ("copyEngineCount", c_uint),
+        ("decoderCount", c_uint),
+        ("encoderCount", c_uint),
+        ("jpegCount", c_uint),
+        ("ofaCount", c_uint),
+        ("memorySizeMB", c_ulonglong),
+        ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE),
+    ]
+
     def __init__(self):
-        super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__(version=nvmlGpuInstanceProfileInfo_v2)
+        super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__(
+            version=nvmlGpuInstanceProfileInfo_v2
+        )
+
 
 class c_nvmlGpuInstanceInfo_t(Structure):
-    _fields_ = [("device", c_nvmlDevice_t),
-                ("id", c_uint),
-                ("profileId", c_uint),
-                ("placement", c_nvmlGpuInstancePlacement_t)
-               ]
+    _fields_ = [
+        ("device", c_nvmlDevice_t),
+        ("id", c_uint),
+        ("profileId", c_uint),
+        ("placement", c_nvmlGpuInstancePlacement_t),
+    ]
+
 
 class struct_c_nvmlGpuInstance_t(Structure):
-    pass # opaque handle
+    pass  # opaque handle
+
+
 c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t)
 
-NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE      = 0x0
-NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE      = 0x1
-NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE      = 0x2
-NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE      = 0x3
-NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE      = 0x4
-NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE      = 0x5
-NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE      = 0x6
+NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0
+NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1
+NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2
+NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3
+NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4
+NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5
+NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6
 NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7
-NVML_COMPUTE_INSTANCE_PROFILE_COUNT        = 0x8
+NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8
 
 NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0
 NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1
 
+
 class c_nvmlComputeInstancePlacement_t(Structure):
-    _fields_ = [("start", c_uint),
-                ("size", c_uint)
-               ]
+    _fields_ = [("start", c_uint), ("size", c_uint)]
+
 
 class c_nvmlComputeInstanceProfileInfo_t(Structure):
-    _fields_ = [("id", c_uint),
-                ("sliceCount", c_uint),
-                ("instanceCount", c_uint),
-                ("multiprocessorCount", c_uint),
-                ("sharedCopyEngineCount", c_uint),
-                ("sharedDecoderCount", c_uint),
-                ("sharedEncoderCount", c_uint),
-                ("sharedJpegCount", c_uint),
-                ("sharedOfaCount", c_uint)
-               ]
+    _fields_ = [
+        ("id", c_uint),
+        ("sliceCount", c_uint),
+        ("instanceCount", c_uint),
+        ("multiprocessorCount", c_uint),
+        ("sharedCopyEngineCount", c_uint),
+        ("sharedDecoderCount", c_uint),
+        ("sharedEncoderCount", c_uint),
+        ("sharedJpegCount", c_uint),
+        ("sharedOfaCount", c_uint),
+    ]
+
 
 nvmlComputeInstanceProfileInfo_v2 = 0x02000088
 
+
 class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure):
-    _fields_ = [("version", c_uint),
-                ("id", c_uint),
-                ("sliceCount", c_uint),
-                ("instanceCount", c_uint),
-                ("multiprocessorCount", c_uint),
-                ("sharedCopyEngineCount", c_uint),
-                ("sharedDecoderCount", c_uint),
-                ("sharedEncoderCount", c_uint),
-                ("sharedJpegCount", c_uint),
-                ("sharedOfaCount", c_uint),
-                ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE)
-               ]
+    _fields_ = [
+        ("version", c_uint),
+        ("id", c_uint),
+        ("sliceCount", c_uint),
+        ("instanceCount", c_uint),
+        ("multiprocessorCount", c_uint),
+        ("sharedCopyEngineCount", c_uint),
+        ("sharedDecoderCount", c_uint),
+        ("sharedEncoderCount", c_uint),
+        ("sharedJpegCount", c_uint),
+        ("sharedOfaCount", c_uint),
+        ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE),
+    ]
 
     def __init__(self):
-        super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__(version=nvmlComputeInstanceProfileInfo_v2)
+        super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__(
+            version=nvmlComputeInstanceProfileInfo_v2
+        )
+
 
 class c_nvmlComputeInstanceInfo_t(Structure):
-    _fields_ = [("device", c_nvmlDevice_t),
-                ("gpuInstance", c_nvmlGpuInstance_t),
-                ("id", c_uint),
-                ("profileId", c_uint),
-                ("placement", c_nvmlComputeInstancePlacement_t)
-               ]
+    _fields_ = [
+        ("device", c_nvmlDevice_t),
+        ("gpuInstance", c_nvmlGpuInstance_t),
+        ("id", c_uint),
+        ("profileId", c_uint),
+        ("placement", c_nvmlComputeInstancePlacement_t),
+    ]
+
 
 NVML_MAX_GPU_UTILIZATIONS = 8
-NVML_GPU_UTILIZATION_DOMAIN_GPU    = 0
-NVML_GPU_UTILIZATION_DOMAIN_FB     = 1
-NVML_GPU_UTILIZATION_DOMAIN_VID    = 2
-NVML_GPU_UTILIZATION_DOMAIN_BUS    = 3
+NVML_GPU_UTILIZATION_DOMAIN_GPU = 0
+NVML_GPU_UTILIZATION_DOMAIN_FB = 1
+NVML_GPU_UTILIZATION_DOMAIN_VID = 2
+NVML_GPU_UTILIZATION_DOMAIN_BUS = 3
+
+
 class c_nvmlGpuDynamicPstatesUtilization_t(Structure):
-    _fields_ = [("bIsPresent", c_uint, 1),
-                ("percentage", c_uint),
-                ("incThreshold", c_uint),
-                ("decThreshold", c_uint)]
+    _fields_ = [
+        ("bIsPresent", c_uint, 1),
+        ("percentage", c_uint),
+        ("incThreshold", c_uint),
+        ("decThreshold", c_uint),
+    ]
+
+
 class c_nvmlGpuDynamicPstatesInfo_t(Structure):
-    _fields_ = [("flags", c_uint),
-                ("utilization", c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS)]
+    _fields_ = [
+        ("flags", c_uint),
+        (
+            "utilization",
+            c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS,
+        ),
+    ]
+
 
 NVML_MAX_THERMAL_SENSORS_PER_GPU = 3
 
-NVML_THERMAL_TARGET_NONE          = 0
-NVML_THERMAL_TARGET_GPU           = 1
-NVML_THERMAL_TARGET_MEMORY        = 2
-NVML_THERMAL_TARGET_POWER_SUPPLY  = 4
-NVML_THERMAL_TARGET_BOARD         = 8
-NVML_THERMAL_TARGET_VCD_BOARD     = 9
-NVML_THERMAL_TARGET_VCD_INLET     = 10
-NVML_THERMAL_TARGET_VCD_OUTLET    = 11
-NVML_THERMAL_TARGET_ALL           = 15
-NVML_THERMAL_TARGET_UNKNOWN       = -1
-
-NVML_THERMAL_CONTROLLER_NONE            = 0
-NVML_THERMAL_CONTROLLER_GPU_INTERNAL    = 1
-NVML_THERMAL_CONTROLLER_ADM1032         = 2
-NVML_THERMAL_CONTROLLER_ADT7461         = 3
-NVML_THERMAL_CONTROLLER_MAX6649         = 4
-NVML_THERMAL_CONTROLLER_MAX1617         = 5
-NVML_THERMAL_CONTROLLER_LM99            = 6
-NVML_THERMAL_CONTROLLER_LM89            = 7
-NVML_THERMAL_CONTROLLER_LM64            = 8
-NVML_THERMAL_CONTROLLER_G781            = 9
-NVML_THERMAL_CONTROLLER_ADT7473         = 10
-NVML_THERMAL_CONTROLLER_SBMAX6649       = 11
-NVML_THERMAL_CONTROLLER_VBIOSEVT        = 12
-NVML_THERMAL_CONTROLLER_OS              = 13
+NVML_THERMAL_TARGET_NONE = 0
+NVML_THERMAL_TARGET_GPU = 1
+NVML_THERMAL_TARGET_MEMORY = 2
+NVML_THERMAL_TARGET_POWER_SUPPLY = 4
+NVML_THERMAL_TARGET_BOARD = 8
+NVML_THERMAL_TARGET_VCD_BOARD = 9
+NVML_THERMAL_TARGET_VCD_INLET = 10
+NVML_THERMAL_TARGET_VCD_OUTLET = 11
+NVML_THERMAL_TARGET_ALL = 15
+NVML_THERMAL_TARGET_UNKNOWN = -1
+
+NVML_THERMAL_CONTROLLER_NONE = 0
+NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1
+NVML_THERMAL_CONTROLLER_ADM1032 = 2
+NVML_THERMAL_CONTROLLER_ADT7461 = 3
+NVML_THERMAL_CONTROLLER_MAX6649 = 4
+NVML_THERMAL_CONTROLLER_MAX1617 = 5
+NVML_THERMAL_CONTROLLER_LM99 = 6
+NVML_THERMAL_CONTROLLER_LM89 = 7
+NVML_THERMAL_CONTROLLER_LM64 = 8
+NVML_THERMAL_CONTROLLER_G781 = 9
+NVML_THERMAL_CONTROLLER_ADT7473 = 10
+NVML_THERMAL_CONTROLLER_SBMAX6649 = 11
+NVML_THERMAL_CONTROLLER_VBIOSEVT = 12
+NVML_THERMAL_CONTROLLER_OS = 13
 NVML_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14
-NVML_THERMAL_CONTROLLER_NVSYSCON_E551   = 15
-NVML_THERMAL_CONTROLLER_MAX6649R        = 16
-NVML_THERMAL_CONTROLLER_ADT7473S        = 17
-NVML_THERMAL_CONTROLLER_UNKNOWN         = -1
+NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15
+NVML_THERMAL_CONTROLLER_MAX6649R = 16
+NVML_THERMAL_CONTROLLER_ADT7473S = 17
+NVML_THERMAL_CONTROLLER_UNKNOWN = -1
+
 
 class c_nvmlGpuThermalSensor_t(Structure):
-    _fields_ = [("controller", c_int),
-                ("defaultMinTemp", c_int),
-                ("defaultMaxTemp", c_int),
-                ("currentTemp", c_int),
-                ("target", c_int)]
+    _fields_ = [
+        ("controller", c_int),
+        ("defaultMinTemp", c_int),
+        ("defaultMaxTemp", c_int),
+        ("currentTemp", c_int),
+        ("target", c_int),
+    ]
+
+
 class c_nvmlGpuThermalSettings_t(Structure):
-    _fields_ = [("count", c_uint),
-                ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU)]
+    _fields_ = [
+        ("count", c_uint),
+        ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU),
+    ]
+
 
 _nvmlCoolerControl_t = c_uint
-NVML_THERMAL_COOLER_SIGNAL_NONE        = 0
-NVML_THERMAL_COOLER_SIGNAL_TOGGLE      = 1
-NVML_THERMAL_COOLER_SIGNAL_VARIABLE    = 2
-NVML_THERMAL_COOLER_SIGNAL_COUNT       = 3
+NVML_THERMAL_COOLER_SIGNAL_NONE = 0
+NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1
+NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2
+NVML_THERMAL_COOLER_SIGNAL_COUNT = 3
 
 _nvmlCoolerTarget_t = c_uint
-NVML_THERMAL_COOLER_TARGET_NONE          = (1 << 0)
-NVML_THERMAL_COOLER_TARGET_GPU           = (1 << 1)
-NVML_THERMAL_COOLER_TARGET_MEMORY        = (1 << 2)
-NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY  = (1 << 3)
-NVML_THERMAL_COOLER_TARGET_GPU_RELATED   = (NVML_THERMAL_COOLER_TARGET_GPU | NVML_THERMAL_COOLER_TARGET_MEMORY | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY)
+NVML_THERMAL_COOLER_TARGET_NONE = 1 << 0
+NVML_THERMAL_COOLER_TARGET_GPU = 1 << 1
+NVML_THERMAL_COOLER_TARGET_MEMORY = 1 << 2
+NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = 1 << 3
+NVML_THERMAL_COOLER_TARGET_GPU_RELATED = (
+    NVML_THERMAL_COOLER_TARGET_GPU
+    | NVML_THERMAL_COOLER_TARGET_MEMORY
+    | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY
+)
+
 
 class c_nvmlCoolerInfo_t(_PrintableStructure):
-    _fields_ = [("version", c_uint),
-                ("index", c_uint),
-                ("coolerControlType", _nvmlCoolerControl_t),
-                ("coolerTarget", _nvmlCoolerTarget_t)
-               ]
+    _fields_ = [
+        ("version", c_uint),
+        ("index", c_uint),
+        ("coolerControlType", _nvmlCoolerControl_t),
+        ("coolerTarget", _nvmlCoolerTarget_t),
+    ]
+
 
 nvmlCoolerInfo_v1 = 0x1000010
 
+
 def nvmlDeviceGetCoolerInfo(handle):
     c_coolerInfo = c_nvmlCoolerInfo_t()
     c_coolerInfo.version = nvmlCoolerInfo_v1
@@ -2241,105 +2672,138 @@ def nvmlDeviceGetCoolerInfo(handle):
     _nvmlCheckReturn(ret)
     return [c_coolerInfo.coolerControlType, c_coolerInfo.coolerTarget]
 
+
 class struct_c_nvmlComputeInstance_t(Structure):
-    pass # opaque handle
+    pass  # opaque handle
+
+
 c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t)
 
+
 class c_nvmlDeviceAttributes(Structure):
-    _fields_ = [("multiprocessorCount", c_uint),
-                ("sharedCopyEngineCount", c_uint),
-                ("sharedDecoderCount", c_uint),
-                ("sharedEncoderCount", c_uint),
-                ("sharedJpegCount", c_uint),
-                ("sharedOfaCount", c_uint),
-                ("gpuInstanceSliceCount", c_uint),
-                ("computeInstanceSliceCount", c_uint),
-                ("memorySizeMB", c_ulonglong),
-               ]
+    _fields_ = [
+        ("multiprocessorCount", c_uint),
+        ("sharedCopyEngineCount", c_uint),
+        ("sharedDecoderCount", c_uint),
+        ("sharedEncoderCount", c_uint),
+        ("sharedJpegCount", c_uint),
+        ("sharedOfaCount", c_uint),
+        ("gpuInstanceSliceCount", c_uint),
+        ("computeInstanceSliceCount", c_uint),
+        ("memorySizeMB", c_ulonglong),
+    ]
+
 
 class c_nvmlRowRemapperHistogramValues(Structure):
-    _fields_ = [("max", c_uint),
-                ("high", c_uint),
-                ("partial", c_uint),
-                ("low", c_uint),
-                ("none", c_uint)
-               ]
-
-NVML_GPU_CERT_CHAIN_SIZE                = 0x1000
-NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE    = 0x1400
-NVML_CC_GPU_CEC_NONCE_SIZE              = 0x20
-NVML_CC_GPU_ATTESTATION_REPORT_SIZE     = 0x2000
+    _fields_ = [
+        ("max", c_uint),
+        ("high", c_uint),
+        ("partial", c_uint),
+        ("low", c_uint),
+        ("none", c_uint),
+    ]
+
+
+NVML_GPU_CERT_CHAIN_SIZE = 0x1000
+NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400
+NVML_CC_GPU_CEC_NONCE_SIZE = 0x20
+NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000
 NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000
 NVML_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0
-NVML_CC_CEC_ATTESTATION_REPORT_PRESENT     = 1
+NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1
+
 
 class c_nvmlConfComputeSystemState_t(Structure):
-    _fields_ = [('environment', c_uint),
-                ('ccFeature', c_uint),
-                ('devToolsMode', c_uint),
-               ]
+    _fields_ = [
+        ("environment", c_uint),
+        ("ccFeature", c_uint),
+        ("devToolsMode", c_uint),
+    ]
+
 
 nvmlSystemConfComputeSettings_v1 = 0x1000014
 
+
 class c_nvmlSystemConfComputeSettings_v1_t(Structure):
-    _fields_ = [('version', c_uint),
-                ('environment', c_uint),
-                ('ccFeature', c_uint),
-                ('devToolsMode', c_uint),
-                ('multiGpuMode', c_uint),
-               ]
+    _fields_ = [
+        ("version", c_uint),
+        ("environment", c_uint),
+        ("ccFeature", c_uint),
+        ("devToolsMode", c_uint),
+        ("multiGpuMode", c_uint),
+    ]
+
     def __init__(self):
-        super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__(version=nvmlSystemConfComputeSettings_v1)
+        super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__(
+            version=nvmlSystemConfComputeSettings_v1
+        )
+
 
 class c_nvmlConfComputeSystemCaps_t(Structure):
-    _fields_ = [('cpuCaps', c_uint),
-                ('gpusCaps', c_uint),
-               ]
+    _fields_ = [
+        ("cpuCaps", c_uint),
+        ("gpusCaps", c_uint),
+    ]
+
 
 class c_nvmlConfComputeMemSizeInfo_t(Structure):
-    _fields_ = [('protectedMemSizeKib', c_ulonglong),
-                ('unprotectedMemSizeKib', c_ulonglong),
-               ]
+    _fields_ = [
+        ("protectedMemSizeKib", c_ulonglong),
+        ("unprotectedMemSizeKib", c_ulonglong),
+    ]
+
 
 class c_nvmlConfComputeGpuCertificate_t(Structure):
-    _fields_ = [('certChainSize', c_uint),
-                ('attestationCertChainSize', c_uint),
-                ('certChain', c_uint8 * NVML_GPU_CERT_CHAIN_SIZE),
-                ('attestationCertChain', c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE),
-               ]
+    _fields_ = [
+        ("certChainSize", c_uint),
+        ("attestationCertChainSize", c_uint),
+        ("certChain", c_uint8 * NVML_GPU_CERT_CHAIN_SIZE),
+        ("attestationCertChain", c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE),
+    ]
+
 
 class c_nvmlConfComputeGpuAttestationReport_t(Structure):
-    _fields_ = [('isCecAttestationReportPresent', c_uint),
-                ('attestationReportSize', c_uint),
-                ('cecAttestationReportSize', c_uint),
-                ('nonce', c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE),
-                ('attestationReport', c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE),
-                ('cecAttestationReport', c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE),
-               ]
+    _fields_ = [
+        ("isCecAttestationReportPresent", c_uint),
+        ("attestationReportSize", c_uint),
+        ("cecAttestationReportSize", c_uint),
+        ("nonce", c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE),
+        ("attestationReport", c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE),
+        ("cecAttestationReport", c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE),
+    ]
+
 
 class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure):
-    _fields_ = [('version', c_uint),
-                ('maxAttackerAdvantage', c_ulong),
-               ]
+    _fields_ = [
+        ("version", c_uint),
+        ("maxAttackerAdvantage", c_ulong),
+    ]
+
+
 ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010
 
+
 class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure):
-    _fields_ = [('version', c_uint),
-                ('attackerAdvantage', c_ulong),
-               ]
+    _fields_ = [
+        ("version", c_uint),
+        ("attackerAdvantage", c_ulong),
+    ]
+
+
 ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010
 
 
 ## string/bytes conversion for ease of use
 def convertStrBytes(func):
-    '''
+    """
     In python 3, strings are unicode instead of bytes, and need to be converted for ctypes
     Args from caller: (1, 'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF>)
     Args passed to function: (1, b'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF)>
     ----
     Returned from function: b'returned string'
     Returned to caller: 'returned string'
-    '''
+    """
+
     @wraps(func)
     def wrapper(*args, **kwargs):
         # encoding a str returns bytes in python 2 and 3
@@ -2358,18 +2822,23 @@ def wrapper(*args, **kwargs):
         return wrapper
     return func
 
+
 def throwOnVersionMismatch(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
         except NVMLError_FunctionNotFound:
-            raise NVMLLibraryMismatchError("Unversioned function called and the "
-                                           "pyNVML version does not match the NVML lib version. "
-                                           "Either use matching pyNVML and NVML lib versions or "
-                                           "use a versioned function such as " + func.__name__ + "_v2")
+            raise NVMLLibraryMismatchError(
+                "Unversioned function called and the "
+                "pyNVML version does not match the NVML lib version. "
+                "Either use matching pyNVML and NVML lib versions or "
+                "use a versioned function such as " + func.__name__ + "_v2"
+            )
+
     return wrapper
 
+
 ## C function wrappers ##
 def nvmlInitWithFlags(flags):
     _LoadNvmlLibrary()
@@ -2388,44 +2857,57 @@ def nvmlInitWithFlags(flags):
     libLoadLock.release()
     return None
 
+
 def nvmlInit():
     nvmlInitWithFlags(0)
     return None
 
+
 def _LoadNvmlLibrary():
-    '''
+    """
     Load the library if it isn't loaded already
-    '''
+    """
     global nvmlLib
 
-    if (nvmlLib == None):
+    if nvmlLib == None:
         # lock to ensure only one caller loads the library
         libLoadLock.acquire()
 
         try:
             # ensure the library still isn't loaded
-            if (nvmlLib == None):
+            if nvmlLib == None:
                 try:
-                    if (sys.platform[:3] == "win"):
+                    if sys.platform[:3] == "win":
                         # cdecl calling convention
                         try:
                             # Check for nvml.dll in System32 first for DCH drivers
-                            nvmlLib = CDLL(os.path.join(os.getenv("WINDIR", "C:/Windows"), "System32/nvml.dll"))
+                            nvmlLib = CDLL(
+                                os.path.join(
+                                    os.getenv("WINDIR", "C:/Windows"),
+                                    "System32/nvml.dll",
+                                )
+                            )
                         except OSError as ose:
                             # If nvml.dll is not found in System32, it should be in ProgramFiles
                             # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll
-                            nvmlLib = CDLL(os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll"))
+                            nvmlLib = CDLL(
+                                os.path.join(
+                                    os.getenv("ProgramFiles", "C:/Program Files"),
+                                    "NVIDIA Corporation/NVSMI/nvml.dll",
+                                )
+                            )
                     else:
                         # assume linux
                         nvmlLib = CDLL("libnvidia-ml.so.1")
                 except OSError as ose:
                     _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
-                if (nvmlLib == None):
+                if nvmlLib == None:
                     _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
         finally:
             # lock is always freed
             libLoadLock.release()
 
+
 def nvmlShutdown():
     #
     # Leave the library loaded, but shutdown the interface
@@ -2437,19 +2919,21 @@ def nvmlShutdown():
     # Atomically update refcount
     global _nvmlLib_refcount
     libLoadLock.acquire()
-    if (0 < _nvmlLib_refcount):
+    if 0 < _nvmlLib_refcount:
         _nvmlLib_refcount -= 1
     libLoadLock.release()
     return None
 
+
 # Added in 2.285
 @convertStrBytes
 def nvmlErrorString(result):
     fn = _nvmlGetFunctionPointer("nvmlErrorString")
-    fn.restype = c_char_p # otherwise return is an int
+    fn.restype = c_char_p  # otherwise return is an int
     ret = fn(result)
     return ret
 
+
 # Added in 2.285
 @convertStrBytes
 def nvmlSystemGetNVMLVersion():
@@ -2459,6 +2943,7 @@ def nvmlSystemGetNVMLVersion():
     _nvmlCheckReturn(ret)
     return c_version.value
 
+
 def nvmlSystemGetCudaDriverVersion():
     c_cuda_version = c_int()
     fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion")
@@ -2466,6 +2951,7 @@ def nvmlSystemGetCudaDriverVersion():
     _nvmlCheckReturn(ret)
     return c_cuda_version.value
 
+
 def nvmlSystemGetCudaDriverVersion_v2():
     c_cuda_version = c_int()
     fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion_v2")
@@ -2473,6 +2959,7 @@ def nvmlSystemGetCudaDriverVersion_v2():
     _nvmlCheckReturn(ret)
     return c_cuda_version.value
 
+
 # Added in 2.285
 @convertStrBytes
 def nvmlSystemGetProcessName(pid):
@@ -2482,6 +2969,7 @@ def nvmlSystemGetProcessName(pid):
     _nvmlCheckReturn(ret)
     return c_name.value
 
+
 @convertStrBytes
 def nvmlSystemGetDriverVersion():
     c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)
@@ -2490,6 +2978,7 @@ def nvmlSystemGetDriverVersion():
     _nvmlCheckReturn(ret)
     return c_version.value
 
+
 # Added in 2.285
 def nvmlSystemGetHicVersion():
     c_count = c_uint(0)
@@ -2500,12 +2989,11 @@ def nvmlSystemGetHicVersion():
     ret = fn(byref(c_count), None)
 
     # this should only fail with insufficient size
-    if ((ret != NVML_SUCCESS) and
-        (ret != NVML_ERROR_INSUFFICIENT_SIZE)):
+    if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE):
         raise NVMLError(ret)
 
     # If there are no hics
-    if (c_count.value == 0):
+    if c_count.value == 0:
         return []
 
     hic_array = c_nvmlHwbcEntry_t * c_count.value
@@ -2514,14 +3002,16 @@ def nvmlSystemGetHicVersion():
     _nvmlCheckReturn(ret)
     return hics
 
+
 def nvmlSystemGetDriverBranch():
     c_branchInfo = c_nvmlSystemDriverBranchInfo_v1_t(0)
     c_branchInfo.version = SystemDriverBranchInfo_v1
-    fn  = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch")
+    fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch")
     ret = fn(byref(c_branchInfo), c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE))
     _nvmlCheckReturn(ret)
     return c_branchInfo
 
+
 ## Unit get functions
 def nvmlUnitGetCount():
     c_count = c_uint()
@@ -2530,6 +3020,7 @@ def nvmlUnitGetCount():
     _nvmlCheckReturn(ret)
     return c_count.value
 
+
 def nvmlUnitGetHandleByIndex(index):
     c_index = c_uint(index)
     unit = c_nvmlUnit_t()
@@ -2538,6 +3029,7 @@ def nvmlUnitGetHandleByIndex(index):
     _nvmlCheckReturn(ret)
     return unit
 
+
 def nvmlUnitGetUnitInfo(unit):
     c_info = c_nvmlUnitInfo_t()
     fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo")
@@ -2545,13 +3037,15 @@ def nvmlUnitGetUnitInfo(unit):
     _nvmlCheckReturn(ret)
     return c_info
 
+
 def nvmlUnitGetLedState(unit):
-    c_state =  c_nvmlLedState_t()
+    c_state = c_nvmlLedState_t()
     fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState")
     ret = fn(unit, byref(c_state))
     _nvmlCheckReturn(ret)
     return c_state
 
+
 def nvmlUnitGetPsuInfo(unit):
     c_info = c_nvmlPSUInfo_t()
     fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo")
@@ -2559,6 +3053,7 @@ def nvmlUnitGetPsuInfo(unit):
     _nvmlCheckReturn(ret)
     return c_info
 
+
 def nvmlUnitGetTemperature(unit, type):
     c_temp = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature")
@@ -2566,6 +3061,7 @@ def nvmlUnitGetTemperature(unit, type):
     _nvmlCheckReturn(ret)
     return c_temp.value
 
+
 def nvmlUnitGetFanSpeedInfo(unit):
     c_speeds = c_nvmlUnitFanSpeeds_t()
     fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo")
@@ -2573,17 +3069,19 @@ def nvmlUnitGetFanSpeedInfo(unit):
     _nvmlCheckReturn(ret)
     return c_speeds
 
+
 # added to API
 def nvmlUnitGetDeviceCount(unit):
     c_count = c_uint(0)
     # query the unit to determine device count
     fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices")
     ret = fn(unit, byref(c_count), None)
-    if (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    if ret == NVML_ERROR_INSUFFICIENT_SIZE:
         ret = NVML_SUCCESS
     _nvmlCheckReturn(ret)
     return c_count.value
 
+
 def nvmlUnitGetDevices(unit):
     c_count = c_uint(nvmlUnitGetDeviceCount(unit))
     device_array = c_nvmlDevice_t * c_count.value
@@ -2593,6 +3091,7 @@ def nvmlUnitGetDevices(unit):
     _nvmlCheckReturn(ret)
     return c_devices
 
+
 ## Device get functions
 def nvmlDeviceGetCount():
     c_count = c_uint()
@@ -2601,6 +3100,7 @@ def nvmlDeviceGetCount():
     _nvmlCheckReturn(ret)
     return c_count.value
 
+
 def nvmlDeviceGetHandleByIndex(index):
     c_index = c_uint(index)
     device = c_nvmlDevice_t()
@@ -2609,6 +3109,7 @@ def nvmlDeviceGetHandleByIndex(index):
     _nvmlCheckReturn(ret)
     return device
 
+
 @convertStrBytes
 def nvmlDeviceGetHandleBySerial(serial):
     c_serial = c_char_p(serial)
@@ -2618,6 +3119,7 @@ def nvmlDeviceGetHandleBySerial(serial):
     _nvmlCheckReturn(ret)
     return device
 
+
 @convertStrBytes
 def nvmlDeviceGetHandleByUUID(uuid):
     c_uuid = c_char_p(uuid)
@@ -2627,6 +3129,7 @@ def nvmlDeviceGetHandleByUUID(uuid):
     _nvmlCheckReturn(ret)
     return device
 
+
 @convertStrBytes
 def nvmlDeviceGetHandleByPciBusId(pciBusId):
     c_busId = c_char_p(pciBusId)
@@ -2636,6 +3139,7 @@ def nvmlDeviceGetHandleByPciBusId(pciBusId):
     _nvmlCheckReturn(ret)
     return device
 
+
 @convertStrBytes
 def nvmlDeviceGetName(handle):
     c_name = create_string_buffer(NVML_DEVICE_NAME_V2_BUFFER_SIZE)
@@ -2644,14 +3148,17 @@ def nvmlDeviceGetName(handle):
     _nvmlCheckReturn(ret)
     return c_name.value
 
+
 class c_nvmlDevicePerfModes_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE),
+        ("version", c_uint),
+        ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE),
     ]
 
+
 nvmlDevicePerfModes_v1 = 0x1000804
 
+
 @convertStrBytes
 def nvmlDeviceGetPerformanceModes(handle):
     perfModes = c_nvmlDevicePerfModes_v1_t()
@@ -2661,14 +3168,17 @@ def nvmlDeviceGetPerformanceModes(handle):
     _nvmlCheckReturn(ret)
     return perfModes.str
 
+
 class c_nvmlDeviceCurrentClockFreqs_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE),
+        ("version", c_uint),
+        ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE),
     ]
 
+
 nvmlDeviceCurrentClockFreqs_v1 = 0x1000804
 
+
 @convertStrBytes
 def nvmlDeviceGetCurrentClockFreqs(handle):
     currentClockFreqs = c_nvmlDeviceCurrentClockFreqs_v1_t()
@@ -2678,20 +3188,23 @@ def nvmlDeviceGetCurrentClockFreqs(handle):
     _nvmlCheckReturn(ret)
     return currentClockFreqs.str
 
+
 def nvmlDeviceGetBoardId(handle):
-    c_id = c_uint();
+    c_id = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId")
     ret = fn(handle, byref(c_id))
     _nvmlCheckReturn(ret)
     return c_id.value
 
+
 def nvmlDeviceGetMultiGpuBoard(handle):
-    c_multiGpu = c_uint();
+    c_multiGpu = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard")
     ret = fn(handle, byref(c_multiGpu))
     _nvmlCheckReturn(ret)
     return c_multiGpu.value
 
+
 def nvmlDeviceGetBrand(handle):
     c_type = _nvmlBrandType_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand")
@@ -2699,6 +3212,7 @@ def nvmlDeviceGetBrand(handle):
     _nvmlCheckReturn(ret)
     return c_type.value
 
+
 def nvmlDeviceGetC2cModeInfoV1(handle):
     c_info = c_nvmlC2cModeInfo_v1_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetC2cModeInfoV")
@@ -2706,9 +3220,11 @@ def nvmlDeviceGetC2cModeInfoV1(handle):
     _nvmlCheckReturn(ret)
     return c_info
 
+
 def nvmlDeviceGetC2cModeInfoV(handle):
     return nvmlDeviceGetC2cModeInfoV1(handle)
 
+
 @convertStrBytes
 def nvmlDeviceGetBoardPartNumber(handle):
     c_part_number = create_string_buffer(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE)
@@ -2717,6 +3233,7 @@ def nvmlDeviceGetBoardPartNumber(handle):
     _nvmlCheckReturn(ret)
     return c_part_number.value
 
+
 @convertStrBytes
 def nvmlDeviceGetSerial(handle):
     c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE)
@@ -2725,6 +3242,7 @@ def nvmlDeviceGetSerial(handle):
     _nvmlCheckReturn(ret)
     return c_serial.value
 
+
 def nvmlDeviceGetModuleId(handle, moduleId=c_uint()):
     isReference = type(moduleId) is not c_uint
     moduleIdRef = moduleId if isReference else byref(moduleId)
@@ -2736,6 +3254,7 @@ def nvmlDeviceGetModuleId(handle, moduleId=c_uint()):
         _nvmlCheckReturn(ret)
         return moduleId.value
 
+
 def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope):
     affinity_array = c_ulonglong * nodeSetSize
     c_affinity = affinity_array()
@@ -2744,6 +3263,7 @@ def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope):
     _nvmlCheckReturn(ret)
     return c_affinity
 
+
 def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope):
     affinity_array = c_ulonglong * cpuSetSize
     c_affinity = affinity_array()
@@ -2752,6 +3272,7 @@ def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope):
     _nvmlCheckReturn(ret)
     return c_affinity
 
+
 def nvmlDeviceGetCpuAffinity(handle, cpuSetSize):
     affinity_array = c_ulonglong * cpuSetSize
     c_affinity = affinity_array()
@@ -2760,18 +3281,21 @@ def nvmlDeviceGetCpuAffinity(handle, cpuSetSize):
     _nvmlCheckReturn(ret)
     return c_affinity
 
+
 def nvmlDeviceSetCpuAffinity(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity")
     ret = fn(handle)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceClearCpuAffinity(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity")
     ret = fn(handle)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetNumaNodeId(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumaNodeId")
     node = c_int()
@@ -2779,6 +3303,7 @@ def nvmlDeviceGetNumaNodeId(handle):
     _nvmlCheckReturn(ret)
     return node.value
 
+
 def nvmlDeviceGetMinorNumber(handle):
     c_minor_number = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber")
@@ -2786,6 +3311,7 @@ def nvmlDeviceGetMinorNumber(handle):
     _nvmlCheckReturn(ret)
     return c_minor_number.value
 
+
 @convertStrBytes
 def nvmlDeviceGetUUID(handle):
     c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE)
@@ -2794,15 +3320,21 @@ def nvmlDeviceGetUUID(handle):
     _nvmlCheckReturn(ret)
     return c_uuid.value
 
+
 @convertStrBytes
 def nvmlDeviceGetInforomVersion(handle, infoRomObject):
     c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion")
-    ret = fn(handle, _nvmlInforomObject_t(infoRomObject),
-                 c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE))
+    ret = fn(
+        handle,
+        _nvmlInforomObject_t(infoRomObject),
+        c_version,
+        c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE),
+    )
     _nvmlCheckReturn(ret)
     return c_version.value
 
+
 # Added in 4.304
 @convertStrBytes
 def nvmlDeviceGetInforomImageVersion(handle):
@@ -2812,6 +3344,7 @@ def nvmlDeviceGetInforomImageVersion(handle):
     _nvmlCheckReturn(ret)
     return c_version.value
 
+
 # Added in 4.304
 def nvmlDeviceGetInforomConfigurationChecksum(handle):
     c_checksum = c_uint()
@@ -2820,6 +3353,7 @@ def nvmlDeviceGetInforomConfigurationChecksum(handle):
     _nvmlCheckReturn(ret)
     return c_checksum.value
 
+
 # Added in 4.304
 def nvmlDeviceValidateInforom(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom")
@@ -2827,6 +3361,7 @@ def nvmlDeviceValidateInforom(handle):
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetLastBBXFlushTime(handle):
     c_timestamp = c_ulonglong()
     c_durationUs = c_ulong()
@@ -2835,6 +3370,7 @@ def nvmlDeviceGetLastBBXFlushTime(handle):
     _nvmlCheckReturn(ret)
     return [c_timestamp.value, c_durationUs.value]
 
+
 def nvmlDeviceGetDisplayMode(handle):
     c_mode = _nvmlEnableState_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode")
@@ -2842,6 +3378,7 @@ def nvmlDeviceGetDisplayMode(handle):
     _nvmlCheckReturn(ret)
     return c_mode.value
 
+
 def nvmlDeviceGetDisplayActive(handle):
     c_mode = _nvmlEnableState_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive")
@@ -2857,12 +3394,14 @@ def nvmlDeviceGetPersistenceMode(handle):
     _nvmlCheckReturn(ret)
     return c_state.value
 
+
 def nvmlDeviceGetPciInfoExt(handle, c_info):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfoExt")
     ret = fn(handle, c_info)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetPciInfo_v3(handle):
     c_info = nvmlPciInfo_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v3")
@@ -2870,9 +3409,11 @@ def nvmlDeviceGetPciInfo_v3(handle):
     _nvmlCheckReturn(ret)
     return c_info
 
+
 def nvmlDeviceGetPciInfo(handle):
     return nvmlDeviceGetPciInfo_v3(handle)
 
+
 def nvmlDeviceGetClockInfo(handle, type):
     c_clock = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo")
@@ -2880,6 +3421,7 @@ def nvmlDeviceGetClockInfo(handle, type):
     _nvmlCheckReturn(ret)
     return c_clock.value
 
+
 # Added in 2.285
 def nvmlDeviceGetMaxClockInfo(handle, type):
     c_clock = c_uint()
@@ -2888,6 +3430,7 @@ def nvmlDeviceGetMaxClockInfo(handle, type):
     _nvmlCheckReturn(ret)
     return c_clock.value
 
+
 # Added in 4.304
 def nvmlDeviceGetApplicationsClock(handle, type):
     c_clock = c_uint()
@@ -2896,6 +3439,7 @@ def nvmlDeviceGetApplicationsClock(handle, type):
     _nvmlCheckReturn(ret)
     return c_clock.value
 
+
 def nvmlDeviceGetMaxCustomerBoostClock(handle, type):
     c_clock = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxCustomerBoostClock")
@@ -2903,6 +3447,7 @@ def nvmlDeviceGetMaxCustomerBoostClock(handle, type):
     _nvmlCheckReturn(ret)
     return c_clock.value
 
+
 def nvmlDeviceGetClock(handle, type, id):
     c_clock = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetClock")
@@ -2910,6 +3455,7 @@ def nvmlDeviceGetClock(handle, type, id):
     _nvmlCheckReturn(ret)
     return c_clock.value
 
+
 # Added in 5.319
 def nvmlDeviceGetDefaultApplicationsClock(handle, type):
     c_clock = c_uint()
@@ -2918,6 +3464,7 @@ def nvmlDeviceGetDefaultApplicationsClock(handle, type):
     _nvmlCheckReturn(ret)
     return c_clock.value
 
+
 # Added in 4.304
 def nvmlDeviceGetSupportedMemoryClocks(handle):
     # first call to get the size
@@ -2925,10 +3472,10 @@ def nvmlDeviceGetSupportedMemoryClocks(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks")
     ret = fn(handle, byref(c_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no clocks
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         clocks_array = c_uint * c_count.value
         c_clocks = clocks_array()
@@ -2946,6 +3493,7 @@ def nvmlDeviceGetSupportedMemoryClocks(handle):
         # error case
         raise NVMLError(ret)
 
+
 # Added in 4.304
 def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz):
     # first call to get the size
@@ -2953,10 +3501,10 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks")
     ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no clocks
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         clocks_array = c_uint * c_count.value
         c_clocks = clocks_array()
@@ -2974,6 +3522,7 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetFanSpeed(handle):
     c_speed = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed")
@@ -2981,6 +3530,7 @@ def nvmlDeviceGetFanSpeed(handle):
     _nvmlCheckReturn(ret)
     return c_speed.value
 
+
 def nvmlDeviceGetFanSpeed_v2(handle, fan):
     c_speed = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed_v2")
@@ -2988,15 +3538,18 @@ def nvmlDeviceGetFanSpeed_v2(handle, fan):
     _nvmlCheckReturn(ret)
     return c_speed.value
 
+
 class c_nvmlFanSpeedInfo_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('fan', c_uint),
-        ('speed', c_uint),
+        ("version", c_uint),
+        ("fan", c_uint),
+        ("speed", c_uint),
     ]
 
+
 nvmlFanSpeedInfo_v1 = 0x100000C
 
+
 def nvmlDeviceGetFanSpeedRPM(handle):
     c_fanSpeed = c_nvmlFanSpeedInfo_t()
     c_fanSpeed.fan = 0
@@ -3006,6 +3559,7 @@ def nvmlDeviceGetFanSpeedRPM(handle):
     _nvmlCheckReturn(ret)
     return c_fanSpeed.speed
 
+
 def nvmlDeviceGetTargetFanSpeed(handle, fan):
     c_speed = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetTargetFanSpeed")
@@ -3013,6 +3567,7 @@ def nvmlDeviceGetTargetFanSpeed(handle, fan):
     _nvmlCheckReturn(ret)
     return c_speed.value
 
+
 def nvmlDeviceGetNumFans(device):
     c_numFans = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumFans")
@@ -3020,12 +3575,14 @@ def nvmlDeviceGetNumFans(device):
     _nvmlCheckReturn(ret)
     return c_numFans.value
 
+
 def nvmlDeviceSetDefaultFanSpeed_v2(handle, index):
-    fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2");
+    fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2")
     ret = fn(handle, index)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()):
     isReference = (type(minSpeed) is not c_uint) or (type(maxSpeed) is not c_uint)
     minSpeedRef = minSpeed if isReference else byref(minSpeed)
@@ -3035,6 +3592,7 @@ def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else [minSpeed.value, maxSpeed.value]
 
+
 def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()):
     isReference = type(fanControlPolicy) is not c_uint
     fanControlPolicyRef = fanControlPolicy if isReference else byref(fanControlPolicy)
@@ -3043,35 +3601,42 @@ def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else fanControlPolicy.value
 
+
 def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanControlPolicy")
     ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy))
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 class c_nvmlTemperature_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('sensorType', _nvmlTemperatureSensors_t),
-        ('temperature', c_int),
+        ("version", c_uint),
+        ("sensorType", _nvmlTemperatureSensors_t),
+        ("temperature", c_int),
     ]
+
+
 nvmlTemperature_v1 = 0x100000C
 
+
 def nvmlDeviceGetTemperatureV1(handle, sensor):
     c_temp = c_nvmlTemperature_v1_t()
     c_temp.version = nvmlTemperature_v1
-    c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) 
+    c_temp.sensorType = _nvmlTemperatureSensors_t(sensor)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureV")
     ret = fn(handle, byref(c_temp))
     _nvmlCheckReturn(ret)
     return c_temp.temperature
 
+
 def nvmlDeviceGetTemperatureV(handle, sensor, version=nvmlTemperature_v1):
     if version == nvmlTemperature_v1:
         return nvmlDeviceGetTemperatureV1(handle, sensor)
     else:
         raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH)
 
+
 # DEPRECATED use nvmlDeviceGetTemperatureV instead
 def nvmlDeviceGetTemperature(handle, sensor):
     c_temp = c_uint()
@@ -3080,6 +3645,7 @@ def nvmlDeviceGetTemperature(handle, sensor):
     _nvmlCheckReturn(ret)
     return c_temp.value
 
+
 def nvmlDeviceGetTemperatureThreshold(handle, threshold):
     c_temp = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold")
@@ -3087,6 +3653,7 @@ def nvmlDeviceGetTemperatureThreshold(handle, threshold):
     _nvmlCheckReturn(ret)
     return c_temp.value
 
+
 def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp):
     c_temp = c_uint()
     c_temp.value = temp
@@ -3095,6 +3662,7 @@ def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp):
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetMarginTemperature(handle):
     c_marginTempInfo = c_nvmlMarginTemperature_v1_t()
     c_marginTempInfo.version = nvmlMarginTemperature_v1
@@ -3103,6 +3671,7 @@ def nvmlDeviceGetMarginTemperature(handle):
     _nvmlCheckReturn(ret)
     return c_marginTempInfo.marginTemperature
 
+
 # DEPRECATED use nvmlDeviceGetPerformanceState
 def nvmlDeviceGetPowerState(handle):
     c_pstate = _nvmlPstates_t()
@@ -3111,6 +3680,7 @@ def nvmlDeviceGetPowerState(handle):
     _nvmlCheckReturn(ret)
     return c_pstate.value
 
+
 def nvmlDeviceGetPerformanceState(handle):
     c_pstate = _nvmlPstates_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState")
@@ -3118,6 +3688,7 @@ def nvmlDeviceGetPerformanceState(handle):
     _nvmlCheckReturn(ret)
     return c_pstate.value
 
+
 def nvmlDeviceGetPowerManagementMode(handle):
     c_pcapMode = _nvmlEnableState_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode")
@@ -3125,6 +3696,7 @@ def nvmlDeviceGetPowerManagementMode(handle):
     _nvmlCheckReturn(ret)
     return c_pcapMode.value
 
+
 def nvmlDeviceGetPowerManagementLimit(handle):
     c_limit = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit")
@@ -3132,6 +3704,7 @@ def nvmlDeviceGetPowerManagementLimit(handle):
     _nvmlCheckReturn(ret)
     return c_limit.value
 
+
 # Added in 4.304
 def nvmlDeviceGetPowerManagementLimitConstraints(handle):
     c_minLimit = c_uint()
@@ -3141,6 +3714,7 @@ def nvmlDeviceGetPowerManagementLimitConstraints(handle):
     _nvmlCheckReturn(ret)
     return [c_minLimit.value, c_maxLimit.value]
 
+
 # Added in 4.304
 def nvmlDeviceGetPowerManagementDefaultLimit(handle):
     c_limit = c_uint()
@@ -3158,6 +3732,7 @@ def nvmlDeviceGetEnforcedPowerLimit(handle):
     _nvmlCheckReturn(ret)
     return c_limit.value
 
+
 def nvmlDeviceGetPowerUsage(handle):
     c_watts = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage")
@@ -3165,6 +3740,7 @@ def nvmlDeviceGetPowerUsage(handle):
     _nvmlCheckReturn(ret)
     return c_watts.value
 
+
 def nvmlDeviceGetTotalEnergyConsumption(handle):
     c_millijoules = c_uint64()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEnergyConsumption")
@@ -3172,6 +3748,7 @@ def nvmlDeviceGetTotalEnergyConsumption(handle):
     _nvmlCheckReturn(ret)
     return c_millijoules.value
 
+
 # Added in 4.304
 def nvmlDeviceGetGpuOperationMode(handle):
     c_currState = _nvmlGpuOperationMode_t()
@@ -3181,14 +3758,17 @@ def nvmlDeviceGetGpuOperationMode(handle):
     _nvmlCheckReturn(ret)
     return [c_currState.value, c_pendingState.value]
 
+
 # Added in 4.304
 def nvmlDeviceGetCurrentGpuOperationMode(handle):
     return nvmlDeviceGetGpuOperationMode(handle)[0]
 
+
 # Added in 4.304
 def nvmlDeviceGetPendingGpuOperationMode(handle):
     return nvmlDeviceGetGpuOperationMode(handle)[1]
 
+
 def nvmlDeviceGetMemoryInfo(handle, version=None):
     if not version:
         c_memory = c_nvmlMemory_t()
@@ -3201,6 +3781,7 @@ def nvmlDeviceGetMemoryInfo(handle, version=None):
     _nvmlCheckReturn(ret)
     return c_memory
 
+
 def nvmlDeviceGetBAR1MemoryInfo(handle):
     c_bar1_memory = c_nvmlBAR1Memory_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo")
@@ -3208,6 +3789,7 @@ def nvmlDeviceGetBAR1MemoryInfo(handle):
     _nvmlCheckReturn(ret)
     return c_bar1_memory
 
+
 def nvmlDeviceGetComputeMode(handle):
     c_mode = _nvmlComputeMode_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode")
@@ -3215,6 +3797,7 @@ def nvmlDeviceGetComputeMode(handle):
     _nvmlCheckReturn(ret)
     return c_mode.value
 
+
 def nvmlDeviceGetCudaComputeCapability(handle):
     c_major = c_int()
     c_minor = c_int()
@@ -3223,6 +3806,7 @@ def nvmlDeviceGetCudaComputeCapability(handle):
     _nvmlCheckReturn(ret)
     return (c_major.value, c_minor.value)
 
+
 def nvmlDeviceGetEccMode(handle):
     c_currState = _nvmlEnableState_t()
     c_pendingState = _nvmlEnableState_t()
@@ -3231,14 +3815,17 @@ def nvmlDeviceGetEccMode(handle):
     _nvmlCheckReturn(ret)
     return [c_currState.value, c_pendingState.value]
 
+
 # added to API
 def nvmlDeviceGetCurrentEccMode(handle):
     return nvmlDeviceGetEccMode(handle)[0]
 
+
 # added to API
 def nvmlDeviceGetPendingEccMode(handle):
     return nvmlDeviceGetEccMode(handle)[1]
 
+
 def nvmlDeviceGetDefaultEccMode(handle):
     c_defaultState = _nvmlEnableState_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultEccMode")
@@ -3246,35 +3833,49 @@ def nvmlDeviceGetDefaultEccMode(handle):
     _nvmlCheckReturn(ret)
     return [c_defaultState.value]
 
+
 def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType):
     c_count = c_ulonglong()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors")
-    ret = fn(handle, _nvmlMemoryErrorType_t(errorType),
-                 _nvmlEccCounterType_t(counterType), byref(c_count))
+    ret = fn(
+        handle,
+        _nvmlMemoryErrorType_t(errorType),
+        _nvmlEccCounterType_t(counterType),
+        byref(c_count),
+    )
     _nvmlCheckReturn(ret)
     return c_count.value
 
+
 # This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter
 def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType):
     c_counts = c_nvmlEccErrorCounts_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors")
-    ret = fn(handle, _nvmlMemoryErrorType_t(errorType),
-                 _nvmlEccCounterType_t(counterType), byref(c_counts))
+    ret = fn(
+        handle,
+        _nvmlMemoryErrorType_t(errorType),
+        _nvmlEccCounterType_t(counterType),
+        byref(c_counts),
+    )
     _nvmlCheckReturn(ret)
     return c_counts
 
+
 # Added in 4.304
 def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType):
     c_count = c_ulonglong()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter")
-    ret = fn(handle,
-             _nvmlMemoryErrorType_t(errorType),
-             _nvmlEccCounterType_t(counterType),
-             _nvmlMemoryLocation_t(locationType),
-             byref(c_count))
+    ret = fn(
+        handle,
+        _nvmlMemoryErrorType_t(errorType),
+        _nvmlEccCounterType_t(counterType),
+        _nvmlMemoryLocation_t(locationType),
+        byref(c_count),
+    )
     _nvmlCheckReturn(ret)
     return c_count.value
 
+
 def nvmlDeviceGetUtilizationRates(handle):
     c_util = c_nvmlUtilization_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates")
@@ -3282,6 +3883,7 @@ def nvmlDeviceGetUtilizationRates(handle):
     _nvmlCheckReturn(ret)
     return c_util
 
+
 def nvmlDeviceGetEncoderUtilization(handle):
     c_util = c_uint()
     c_samplingPeriod = c_uint()
@@ -3290,6 +3892,7 @@ def nvmlDeviceGetEncoderUtilization(handle):
     _nvmlCheckReturn(ret)
     return [c_util.value, c_samplingPeriod.value]
 
+
 def nvmlDeviceGetDecoderUtilization(handle):
     c_util = c_uint()
     c_samplingPeriod = c_uint()
@@ -3298,6 +3901,7 @@ def nvmlDeviceGetDecoderUtilization(handle):
     _nvmlCheckReturn(ret)
     return [c_util.value, c_samplingPeriod.value]
 
+
 def nvmlDeviceGetJpgUtilization(handle):
     c_util = c_uint()
     c_samplingPeriod = c_uint()
@@ -3306,6 +3910,7 @@ def nvmlDeviceGetJpgUtilization(handle):
     _nvmlCheckReturn(ret)
     return [c_util.value, c_samplingPeriod.value]
 
+
 def nvmlDeviceGetOfaUtilization(handle):
     c_util = c_uint()
     c_samplingPeriod = c_uint()
@@ -3314,6 +3919,7 @@ def nvmlDeviceGetOfaUtilization(handle):
     _nvmlCheckReturn(ret)
     return [c_util.value, c_samplingPeriod.value]
 
+
 def nvmlDeviceGetPcieReplayCounter(handle):
     c_replay = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter")
@@ -3321,6 +3927,7 @@ def nvmlDeviceGetPcieReplayCounter(handle):
     _nvmlCheckReturn(ret)
     return c_replay.value
 
+
 def nvmlDeviceGetDriverModel(handle):
     c_currModel = _nvmlDriverModel_t()
     c_pendingModel = _nvmlDriverModel_t()
@@ -3329,14 +3936,17 @@ def nvmlDeviceGetDriverModel(handle):
     _nvmlCheckReturn(ret)
     return [c_currModel.value, c_pendingModel.value]
 
+
 # added to API
 def nvmlDeviceGetCurrentDriverModel(handle):
     return nvmlDeviceGetDriverModel(handle)[0]
 
+
 # added to API
 def nvmlDeviceGetPendingDriverModel(handle):
     return nvmlDeviceGetDriverModel(handle)[1]
 
+
 # Added in 2.285
 @convertStrBytes
 def nvmlDeviceGetVbiosVersion(handle):
@@ -3346,16 +3956,17 @@ def nvmlDeviceGetVbiosVersion(handle):
     _nvmlCheckReturn(ret)
     return c_version.value
 
+
 # Added in 2.285
 def nvmlDeviceGetComputeRunningProcesses_v2(handle):
     # first call to get the size
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v2")
     ret = fn(handle, byref(c_count), None)
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no running processes
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         # oversize the array incase more processes are created
         c_count.value = c_count.value * 2 + 5
@@ -3368,7 +3979,7 @@ def nvmlDeviceGetComputeRunningProcesses_v2(handle):
         for i in range(c_count.value):
             # use an alternative struct for this object
             obj = nvmlStructToFriendlyObject(c_procs[i])
-            if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+            if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
                 # special case for WDDM on Windows, see comment above
                 obj.usedGpuMemory = None
             procs.append(obj)
@@ -3377,6 +3988,7 @@ def nvmlDeviceGetComputeRunningProcesses_v2(handle):
         # error case
         raise NVMLError(ret)
 
+
 # Added in 2.285
 def nvmlDeviceGetComputeRunningProcesses_v3(handle):
     # first call to get the size
@@ -3384,10 +3996,10 @@ def nvmlDeviceGetComputeRunningProcesses_v3(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v3")
     ret = fn(handle, byref(c_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no running processes
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         # oversize the array incase more processes are created
         c_count.value = c_count.value * 2 + 5
@@ -3402,7 +4014,7 @@ def nvmlDeviceGetComputeRunningProcesses_v3(handle):
         for i in range(c_count.value):
             # use an alternative struct for this object
             obj = nvmlStructToFriendlyObject(c_procs[i])
-            if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+            if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
                 # special case for WDDM on Windows, see comment above
                 obj.usedGpuMemory = None
             procs.append(obj)
@@ -3412,19 +4024,21 @@ def nvmlDeviceGetComputeRunningProcesses_v3(handle):
         # error case
         raise NVMLError(ret)
 
+
 @throwOnVersionMismatch
 def nvmlDeviceGetComputeRunningProcesses(handle):
     return nvmlDeviceGetComputeRunningProcesses_v3(handle)
 
+
 def nvmlDeviceGetGraphicsRunningProcesses_v2(handle):
     # first call to get the size
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v2")
     ret = fn(handle, byref(c_count), None)
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no running processes
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         # oversize the array incase more processes are created
         c_count.value = c_count.value * 2 + 5
@@ -3437,7 +4051,7 @@ def nvmlDeviceGetGraphicsRunningProcesses_v2(handle):
         for i in range(c_count.value):
             # use an alternative struct for this object
             obj = nvmlStructToFriendlyObject(c_procs[i])
-            if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+            if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
                 # special case for WDDM on Windows, see comment above
                 obj.usedGpuMemory = None
             procs.append(obj)
@@ -3446,16 +4060,17 @@ def nvmlDeviceGetGraphicsRunningProcesses_v2(handle):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetGraphicsRunningProcesses_v3(handle):
     # first call to get the size
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v3")
     ret = fn(handle, byref(c_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no running processes
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         # oversize the array incase more processes are created
         c_count.value = c_count.value * 2 + 5
@@ -3470,7 +4085,7 @@ def nvmlDeviceGetGraphicsRunningProcesses_v3(handle):
         for i in range(c_count.value):
             # use an alternative struct for this object
             obj = nvmlStructToFriendlyObject(c_procs[i])
-            if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+            if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
                 # special case for WDDM on Windows, see comment above
                 obj.usedGpuMemory = None
             procs.append(obj)
@@ -3480,24 +4095,27 @@ def nvmlDeviceGetGraphicsRunningProcesses_v3(handle):
         # error case
         raise NVMLError(ret)
 
+
 @throwOnVersionMismatch
 def nvmlDeviceGetGraphicsRunningProcesses(handle):
     return nvmlDeviceGetGraphicsRunningProcesses_v3(handle)
 
+
 @throwOnVersionMismatch
 def nvmlDeviceGetMPSComputeRunningProcesses(handle):
     return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle)
 
+
 def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle):
     # first call to get the size
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v2")
     ret = fn(handle, byref(c_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no running processes
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         # oversize the array incase more processes are created
         c_count.value = c_count.value * 2 + 5
@@ -3512,7 +4130,7 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle):
         for i in range(c_count.value):
             # use an alternative struct for this object
             obj = nvmlStructToFriendlyObject(c_procs[i])
-            if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+            if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
                 # special case for WDDM on Windows, see comment above
                 obj.usedGpuMemory = None
             procs.append(obj)
@@ -3522,16 +4140,17 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle):
     # first call to get the size
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v3")
     ret = fn(handle, byref(c_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no running processes
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         # oversize the array in case more processes are created
         c_count.value = c_count.value * 2 + 5
@@ -3546,7 +4165,7 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle):
         for i in range(c_count.value):
             # use an alternative struct for this object
             obj = nvmlStructToFriendlyObject(c_procs[i])
-            if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+            if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
                 # special case for WDDM on Windows, see comment above
                 obj.usedGpuMemory = None
             procs.append(obj)
@@ -3556,6 +4175,7 @@ def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetRunningProcessDetailList(handle, version, mode):
     c_processDetailList = c_nvmlProcessDetailList_t()
     c_processDetailList.version = version
@@ -3565,12 +4185,14 @@ def nvmlDeviceGetRunningProcessDetailList(handle, version, mode):
 
     # first call to get the size
     ret = fn(handle, byref(c_processDetailList))
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no running processes
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         c_procs = c_nvmlProcessDetail_v1_t * c_processDetailList.numProcArrayEntries
-        c_processDetailList.procArray = cast((c_procs)(), POINTER(c_nvmlProcessDetail_v1_t))
+        c_processDetailList.procArray = cast(
+            (c_procs)(), POINTER(c_nvmlProcessDetail_v1_t)
+        )
 
         # make the call again
         ret = fn(handle, byref(c_processDetailList))
@@ -3580,9 +4202,9 @@ def nvmlDeviceGetRunningProcessDetailList(handle, version, mode):
         for i in range(c_processDetailList.numProcArrayEntries):
             # use an alternative struct for this object
             obj = c_processDetailList.procArray[i]
-            if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+            if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
                 obj.usedGpuMemory = None
-            if (obj.usedGpuCcProtectedMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+            if obj.usedGpuCcProtectedMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
                 obj.usedGpuCcProtectedMemory = None
             procs.append(obj)
 
@@ -3591,6 +4213,7 @@ def nvmlDeviceGetRunningProcessDetailList(handle, version, mode):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetAutoBoostedClocksEnabled(handle):
     c_isEnabled = _nvmlEnableState_t()
     c_defaultIsEnabled = _nvmlEnableState_t()
@@ -3598,7 +4221,8 @@ def nvmlDeviceGetAutoBoostedClocksEnabled(handle):
     ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled))
     _nvmlCheckReturn(ret)
     return [c_isEnabled.value, c_defaultIsEnabled.value]
-    #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+    # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+
 
 ## Set functions
 def nvmlUnitSetLedState(unit, color):
@@ -3607,49 +4231,57 @@ def nvmlUnitSetLedState(unit, color):
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceSetPersistenceMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode")
     ret = fn(handle, _nvmlEnableState_t(mode))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceSetComputeMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode")
     ret = fn(handle, _nvmlComputeMode_t(mode))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceSetEccMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode")
     ret = fn(handle, _nvmlEnableState_t(mode))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceClearEccErrorCounts(handle, counterType):
     fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts")
     ret = fn(handle, _nvmlEccCounterType_t(counterType))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceSetDriverModel(handle, model):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel")
     ret = fn(handle, _nvmlDriverModel_t(model))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled")
     ret = fn(handle, _nvmlEnableState_t(enabled))
     _nvmlCheckReturn(ret)
     return None
-    #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+    # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+
 
 def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled")
     ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags))
     _nvmlCheckReturn(ret)
     return None
-    #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+    # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+
 
 def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuLockedClocks")
@@ -3657,24 +4289,28 @@ def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz):
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceResetGpuLockedClocks(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceResetGpuLockedClocks")
     ret = fn(handle)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemoryLockedClocks")
     ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceResetMemoryLockedClocks(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceResetMemoryLockedClocks")
     ret = fn(handle)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()):
     isReference = type(c_clkMonInfo) is not nvmlClkMonStatus_t
     c_clkMonInfoRef = c_clkMonInfo if isReference else byref(c_clkMonInfo)
@@ -3683,6 +4319,7 @@ def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else c_clkMonInfo
 
+
 # Added in 4.304
 def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks")
@@ -3690,6 +4327,7 @@ def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz)
     _nvmlCheckReturn(ret)
     return None
 
+
 # Added in 4.304
 def nvmlDeviceResetApplicationsClocks(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks")
@@ -3697,6 +4335,7 @@ def nvmlDeviceResetApplicationsClocks(handle):
     _nvmlCheckReturn(ret)
     return None
 
+
 # Added in 4.304
 def nvmlDeviceSetPowerManagementLimit(handle, limit):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit")
@@ -3704,6 +4343,7 @@ def nvmlDeviceSetPowerManagementLimit(handle, limit):
     _nvmlCheckReturn(ret)
     return None
 
+
 # Added in 4.304
 def nvmlDeviceSetGpuOperationMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode")
@@ -3711,6 +4351,7 @@ def nvmlDeviceSetGpuOperationMode(handle, mode):
     _nvmlCheckReturn(ret)
     return None
 
+
 # Added in 2.285
 def nvmlEventSetCreate():
     fn = _nvmlGetFunctionPointer("nvmlEventSetCreate")
@@ -3719,6 +4360,7 @@ def nvmlEventSetCreate():
     _nvmlCheckReturn(ret)
     return eventSet
 
+
 # Added in 2.285
 def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet):
     fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents")
@@ -3726,6 +4368,7 @@ def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet):
     _nvmlCheckReturn(ret)
     return None
 
+
 # Added in 2.285
 def nvmlDeviceGetSupportedEventTypes(handle):
     c_eventTypes = c_ulonglong()
@@ -3734,6 +4377,7 @@ def nvmlDeviceGetSupportedEventTypes(handle):
     _nvmlCheckReturn(ret)
     return c_eventTypes.value
 
+
 # raises NVML_ERROR_TIMEOUT exception on timeout
 def nvmlEventSetWait_v2(eventSet, timeoutms):
     fn = _nvmlGetFunctionPointer("nvmlEventSetWait_v2")
@@ -3742,9 +4386,11 @@ def nvmlEventSetWait_v2(eventSet, timeoutms):
     _nvmlCheckReturn(ret)
     return data
 
+
 def nvmlEventSetWait(eventSet, timeoutms):
     return nvmlEventSetWait_v2(eventSet, timeoutms)
 
+
 # Added in 2.285
 def nvmlEventSetFree(eventSet):
     fn = _nvmlGetFunctionPointer("nvmlEventSetFree")
@@ -3752,13 +4398,15 @@ def nvmlEventSetFree(eventSet):
     _nvmlCheckReturn(ret)
     return None
 
+
 # Added in 3.295
 def nvmlDeviceOnSameBoard(handle1, handle2):
     fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard")
     onSameBoard = c_int()
     ret = fn(handle1, handle2, byref(onSameBoard))
     _nvmlCheckReturn(ret)
-    return (onSameBoard.value != 0)
+    return onSameBoard.value != 0
+
 
 # Added in 3.295
 def nvmlDeviceGetCurrPcieLinkGeneration(handle):
@@ -3768,6 +4416,7 @@ def nvmlDeviceGetCurrPcieLinkGeneration(handle):
     _nvmlCheckReturn(ret)
     return gen.value
 
+
 # Added in 3.295
 def nvmlDeviceGetMaxPcieLinkGeneration(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration")
@@ -3776,6 +4425,7 @@ def nvmlDeviceGetMaxPcieLinkGeneration(handle):
     _nvmlCheckReturn(ret)
     return gen.value
 
+
 # Added in 3.295
 def nvmlDeviceGetCurrPcieLinkWidth(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth")
@@ -3784,6 +4434,7 @@ def nvmlDeviceGetCurrPcieLinkWidth(handle):
     _nvmlCheckReturn(ret)
     return width.value
 
+
 # Added in 3.295
 def nvmlDeviceGetMaxPcieLinkWidth(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth")
@@ -3792,6 +4443,7 @@ def nvmlDeviceGetMaxPcieLinkWidth(handle):
     _nvmlCheckReturn(ret)
     return width.value
 
+
 def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuMaxPcieLinkGeneration")
     gen = c_uint()
@@ -3799,36 +4451,41 @@ def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle):
     _nvmlCheckReturn(ret)
     return gen.value
 
+
 # Added in 4.304
 def nvmlDeviceGetSupportedClocksThrottleReasons(handle):
-    c_reasons= c_ulonglong()
+    c_reasons = c_ulonglong()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons")
     ret = fn(handle, byref(c_reasons))
     _nvmlCheckReturn(ret)
     return c_reasons.value
 
+
 def nvmlDeviceGetSupportedClocksEventReasons(handle):
-    c_reasons= c_ulonglong()
+    c_reasons = c_ulonglong()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksEventReasons")
     ret = fn(handle, byref(c_reasons))
     _nvmlCheckReturn(ret)
     return c_reasons.value
 
+
 # Added in 4.304
 def nvmlDeviceGetCurrentClocksThrottleReasons(handle):
-    c_reasons= c_ulonglong()
+    c_reasons = c_ulonglong()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons")
     ret = fn(handle, byref(c_reasons))
     _nvmlCheckReturn(ret)
     return c_reasons.value
 
+
 def nvmlDeviceGetCurrentClocksEventReasons(handle):
-    c_reasons= c_ulonglong()
+    c_reasons = c_ulonglong()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksEventReasons")
     ret = fn(handle, byref(c_reasons))
     _nvmlCheckReturn(ret)
     return c_reasons.value
 
+
 # Added in 5.319
 def nvmlDeviceGetIndex(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex")
@@ -3837,6 +4494,7 @@ def nvmlDeviceGetIndex(handle):
     _nvmlCheckReturn(ret)
     return c_index.value
 
+
 # Added in 5.319
 def nvmlDeviceGetAccountingMode(handle):
     c_mode = _nvmlEnableState_t()
@@ -3845,35 +4503,40 @@ def nvmlDeviceGetAccountingMode(handle):
     _nvmlCheckReturn(ret)
     return c_mode.value
 
+
 def nvmlDeviceSetAccountingMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode")
     ret = fn(handle, _nvmlEnableState_t(mode))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceClearAccountingPids(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids")
     ret = fn(handle)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetAccountingStats(handle, pid):
     stats = c_nvmlAccountingStats_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats")
     ret = fn(handle, c_uint(pid), byref(stats))
     _nvmlCheckReturn(ret)
-    if (stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value):
+    if stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
         # special case for WDDM on Windows, see comment above
         stats.maxMemoryUsage = None
     return stats
 
+
 def nvmlDeviceGetAccountingPids(handle):
     count = c_uint(nvmlDeviceGetAccountingBufferSize(handle))
     pids = (c_uint * count.value)()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids")
     ret = fn(handle, byref(count), pids)
     _nvmlCheckReturn(ret)
-    return list(map(int, pids[0:count.value]))
+    return list(map(int, pids[0 : count.value]))
+
 
 def nvmlDeviceGetAccountingBufferSize(handle):
     bufferSize = c_uint()
@@ -3882,6 +4545,7 @@ def nvmlDeviceGetAccountingBufferSize(handle):
     _nvmlCheckReturn(ret)
     return int(bufferSize.value)
 
+
 def nvmlDeviceGetRetiredPages(device, sourceFilter):
     c_source = _nvmlPageRetirementCause_t(sourceFilter)
     c_count = c_uint(0)
@@ -3891,8 +4555,7 @@ def nvmlDeviceGetRetiredPages(device, sourceFilter):
     ret = fn(device, c_source, byref(c_count), None)
 
     # this should only fail with insufficient size
-    if ((ret != NVML_SUCCESS) and
-        (ret != NVML_ERROR_INSUFFICIENT_SIZE)):
+    if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE):
         raise NVMLError(ret)
 
     # call again with a buffer
@@ -3903,7 +4566,8 @@ def nvmlDeviceGetRetiredPages(device, sourceFilter):
     c_pages = page_array()
     ret = fn(device, c_source, byref(c_count), c_pages)
     _nvmlCheckReturn(ret)
-    return list(map(int, c_pages[0:c_count.value]))
+    return list(map(int, c_pages[0 : c_count.value]))
+
 
 def nvmlDeviceGetRetiredPages_v2(device, sourceFilter):
     c_source = _nvmlPageRetirementCause_t(sourceFilter)
@@ -3914,8 +4578,7 @@ def nvmlDeviceGetRetiredPages_v2(device, sourceFilter):
     ret = fn(device, c_source, byref(c_count), None)
 
     # this should only fail with insufficient size
-    if ((ret != NVML_SUCCESS) and
-        (ret != NVML_ERROR_INSUFFICIENT_SIZE)):
+    if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE):
         raise NVMLError(ret)
 
     # call again with a buffer
@@ -3928,7 +4591,11 @@ def nvmlDeviceGetRetiredPages_v2(device, sourceFilter):
     c_times = times_array()
     ret = fn(device, c_source, byref(c_count), c_pages, c_times)
     _nvmlCheckReturn(ret)
-    return [ { 'address': int(c_pages[i]), 'timestamp': int(c_times[i]) } for i in range(c_count.value) ];
+    return [
+        {"address": int(c_pages[i]), "timestamp": int(c_times[i])}
+        for i in range(c_count.value)
+    ]
+
 
 def nvmlDeviceGetRetiredPagesPendingStatus(device):
     c_pending = _nvmlEnableState_t()
@@ -3937,6 +4604,7 @@ def nvmlDeviceGetRetiredPagesPendingStatus(device):
     _nvmlCheckReturn(ret)
     return int(c_pending.value)
 
+
 def nvmlDeviceGetAPIRestriction(device, apiType):
     c_permission = _nvmlEnableState_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction")
@@ -3944,12 +4612,14 @@ def nvmlDeviceGetAPIRestriction(device, apiType):
     _nvmlCheckReturn(ret)
     return int(c_permission.value)
 
+
 def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction")
     ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetBridgeChipInfo(handle):
     bridgeHierarchy = c_nvmlBridgeChipHierarchy_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo")
@@ -3957,6 +4627,7 @@ def nvmlDeviceGetBridgeChipInfo(handle):
     _nvmlCheckReturn(ret)
     return bridgeHierarchy
 
+
 def nvmlDeviceGetSamples(device, sampling_type, timeStamp):
     c_sampling_type = _nvmlSamplingType_t(sampling_type)
     c_time_stamp = c_ulonglong(timeStamp)
@@ -3965,17 +4636,32 @@ def nvmlDeviceGetSamples(device, sampling_type, timeStamp):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSamples")
 
     ## First Call gets the size
-    ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), None)
+    ret = fn(
+        device,
+        c_sampling_type,
+        c_time_stamp,
+        byref(c_sample_value_type),
+        byref(c_sample_count),
+        None,
+    )
 
     # Stop if this fails
-    if (ret != NVML_SUCCESS):
+    if ret != NVML_SUCCESS:
         raise NVMLError(ret)
 
     sampleArray = c_sample_count.value * c_nvmlSample_t
     c_samples = sampleArray()
-    ret = fn(device, c_sampling_type, c_time_stamp,  byref(c_sample_value_type), byref(c_sample_count), c_samples)
+    ret = fn(
+        device,
+        c_sampling_type,
+        c_time_stamp,
+        byref(c_sample_value_type),
+        byref(c_sample_count),
+        c_samples,
+    )
     _nvmlCheckReturn(ret)
-    return (c_sample_value_type.value, c_samples[0:c_sample_count.value])
+    return (c_sample_value_type.value, c_samples[0 : c_sample_count.value])
+
 
 def nvmlDeviceGetViolationStatus(device, perfPolicyType):
     c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType)
@@ -3987,6 +4673,7 @@ def nvmlDeviceGetViolationStatus(device, perfPolicyType):
     _nvmlCheckReturn(ret)
     return c_violTime
 
+
 def nvmlDeviceGetPcieThroughput(device, counter):
     c_util = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput")
@@ -3994,6 +4681,7 @@ def nvmlDeviceGetPcieThroughput(device, counter):
     _nvmlCheckReturn(ret)
     return c_util.value
 
+
 def nvmlSystemGetTopologyGpuSet(cpuNumber):
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet")
@@ -4008,7 +4696,8 @@ def nvmlSystemGetTopologyGpuSet(cpuNumber):
     c_devices = device_array()
     ret = fn(cpuNumber, byref(c_count), c_devices)
     _nvmlCheckReturn(ret)
-    return list(c_devices[0:c_count.value])
+    return list(c_devices[0 : c_count.value])
+
 
 def nvmlDeviceGetTopologyNearestGpus(device, level):
     c_count = c_uint(0)
@@ -4025,7 +4714,8 @@ def nvmlDeviceGetTopologyNearestGpus(device, level):
     c_devices = device_array()
     ret = fn(device, level, byref(c_count), c_devices)
     _nvmlCheckReturn(ret)
-    return list(c_devices[0:c_count.value])
+    return list(c_devices[0 : c_count.value])
+
 
 def nvmlDeviceGetTopologyCommonAncestor(device1, device2):
     c_level = _nvmlGpuTopologyLevel_t()
@@ -4034,6 +4724,7 @@ def nvmlDeviceGetTopologyCommonAncestor(device1, device2):
     _nvmlCheckReturn(ret)
     return c_level.value
 
+
 def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter):
     c_rxcounter = c_ulonglong()
     c_txcounter = c_ulonglong()
@@ -4042,24 +4733,28 @@ def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter):
     _nvmlCheckReturn(ret)
     return (c_rxcounter.value, c_txcounter.value)
 
+
 def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze):
     fn = _nvmlGetFunctionPointer("nvmlDeviceFreezeNvLinkUtilizationCounter")
     ret = fn(device, link, counter, freeze)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter):
     fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkUtilizationCounter")
     ret = fn(device, link, counter)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceSetNvLinkUtilizationControl(device, link, counter, control, reset):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkUtilizationControl")
     ret = fn(device, link, counter, byref(control), reset)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter):
     c_control = nvmlNvLinkUtilizationControl_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationControl")
@@ -4067,6 +4762,7 @@ def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter):
     _nvmlCheckReturn(ret)
     return c_control
 
+
 def nvmlDeviceGetNvLinkCapability(device, link, capability):
     c_capResult = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkCapability")
@@ -4074,6 +4770,7 @@ def nvmlDeviceGetNvLinkCapability(device, link, capability):
     _nvmlCheckReturn(ret)
     return c_capResult.value
 
+
 def nvmlDeviceGetNvLinkErrorCounter(device, link, counter):
     c_result = c_ulonglong()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkErrorCounter")
@@ -4081,12 +4778,14 @@ def nvmlDeviceGetNvLinkErrorCounter(device, link, counter):
     _nvmlCheckReturn(ret)
     return c_result.value
 
+
 def nvmlDeviceResetNvLinkErrorCounters(device, link):
     fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkErrorCounters")
     ret = fn(device, link)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetNvLinkRemotePciInfo(device, link):
     c_pci = nvmlPciInfo_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemotePciInfo_v2")
@@ -4094,6 +4793,7 @@ def nvmlDeviceGetNvLinkRemotePciInfo(device, link):
     _nvmlCheckReturn(ret)
     return c_pci
 
+
 def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link):
     c_type = _nvmlNvLinkDeviceType_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemoteDeviceType")
@@ -4101,6 +4801,7 @@ def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link):
     _nvmlCheckReturn(ret)
     return c_type.value
 
+
 def nvmlDeviceGetNvLinkState(device, link):
     c_isActive = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkState")
@@ -4108,6 +4809,7 @@ def nvmlDeviceGetNvLinkState(device, link):
     _nvmlCheckReturn(ret)
     return c_isActive.value
 
+
 def nvmlDeviceGetNvLinkVersion(device, link):
     c_version = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkVersion")
@@ -4115,12 +4817,14 @@ def nvmlDeviceGetNvLinkVersion(device, link):
     _nvmlCheckReturn(ret)
     return c_version.value
 
+
 def nvmlDeviceModifyDrainState(pciInfo, newState):
     fn = _nvmlGetFunctionPointer("nvmlDeviceModifyDrainState")
     ret = fn(pointer(pciInfo), newState)
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceQueryDrainState(pciInfo):
     c_newState = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceQueryDrainState")
@@ -4128,18 +4832,21 @@ def nvmlDeviceQueryDrainState(pciInfo):
     _nvmlCheckReturn(ret)
     return c_newState.value
 
+
 def nvmlDeviceRemoveGpu(pciInfo):
     fn = _nvmlGetFunctionPointer("nvmlDeviceRemoveGpu")
     ret = fn(pointer(pciInfo))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceDiscoverGpus(pciInfo):
     fn = _nvmlGetFunctionPointer("nvmlDeviceDiscoverGpus")
     ret = fn(pointer(pciInfo))
     _nvmlCheckReturn(ret)
     return None
 
+
 def nvmlDeviceGetFieldValues(handle, fieldIds):
     values_arr = c_nvmlFieldValue_t * len(fieldIds)
     values = values_arr()
@@ -4155,6 +4862,7 @@ def nvmlDeviceGetFieldValues(handle, fieldIds):
     _nvmlCheckReturn(ret)
     return values
 
+
 def nvmlDeviceClearFieldValues(handle, fieldIds):
     values_arr = c_nvmlFieldValue_t * len(fieldIds)
     values = values_arr()
@@ -4170,6 +4878,7 @@ def nvmlDeviceClearFieldValues(handle, fieldIds):
     _nvmlCheckReturn(ret)
     return values
 
+
 def nvmlDeviceGetVirtualizationMode(handle):
     c_virtualization_mode = c_ulonglong()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetVirtualizationMode")
@@ -4177,10 +4886,12 @@ def nvmlDeviceGetVirtualizationMode(handle):
     _nvmlCheckReturn(ret)
     return c_virtualization_mode.value
 
+
 def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetVirtualizationMode")
     return fn(handle, virtualization_mode)
 
+
 def nvmlDeviceGetVgpuHeterogeneousMode(handle):
     c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0)
     c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1
@@ -4189,6 +4900,7 @@ def nvmlDeviceGetVgpuHeterogeneousMode(handle):
     _nvmlCheckReturn(ret)
     return c_vgpuHeterogeneousMode.mode
 
+
 def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode):
     c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0)
     c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1
@@ -4198,6 +4910,7 @@ def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlVgpuInstanceGetPlacementId(vgpuInstance):
     c_placement = c_nvmlVgpuPlacementId_v1_t(0)
     c_placement.version = VgpuPlacementId_v1
@@ -4206,9 +4919,10 @@ def nvmlVgpuInstanceGetPlacementId(vgpuInstance):
     _nvmlCheckReturn(ret)
     return c_placement.placementId
 
+
 def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version=1):
     c_max_instances = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances")
     ret = fn(handle, vgpuTypeId, byref(c_max_instances))
     _nvmlCheckReturn(ret)
 
@@ -4225,14 +4939,15 @@ def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version
 
     c_placements = c_uint * c_max_instances.value
     c_vgpu_placements.placementIds = c_placements()
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements")
     ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements))
     _nvmlCheckReturn(ret)
     return c_vgpu_placements
 
+
 def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1):
     c_max_instances = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances")
     ret = fn(handle, vgpuTypeId, byref(c_max_instances))
     _nvmlCheckReturn(ret)
 
@@ -4246,11 +4961,12 @@ def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1):
 
     c_placements = c_uint * c_max_instances.value
     c_vgpu_placements.placementIds = c_placements()
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements")
     ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements))
     _nvmlCheckReturn(ret)
     return c_vgpu_placements
 
+
 def nvmlGetVgpuDriverCapabilities(capability):
     c_capResult = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlGetVgpuDriverCapabilities")
@@ -4258,6 +4974,7 @@ def nvmlGetVgpuDriverCapabilities(capability):
     _nvmlCheckReturn(ret)
     return c_capResult.value
 
+
 def nvmlDeviceGetVgpuCapabilities(handle, capability):
     c_capResult = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuCapabilities")
@@ -4265,23 +4982,25 @@ def nvmlDeviceGetVgpuCapabilities(handle, capability):
     _nvmlCheckReturn(ret)
     return c_capResult.value
 
+
 def nvmlDeviceSetVgpuCapabilities(handle, capability, state):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuCapabilities")
     ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetSupportedVgpus(handle):
     # first call to get the size
     c_vgpu_count = c_uint(0)
 
-    fn =  _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus")
     ret = fn(handle, byref(c_vgpu_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no supported vGPUs
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value
         c_vgpu_type_ids = vgpu_type_ids_array()
@@ -4297,17 +5016,18 @@ def nvmlDeviceGetSupportedVgpus(handle):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetCreatableVgpus(handle):
     # first call to get the size
     c_vgpu_count = c_uint(0)
 
-    fn =  _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus")
     ret = fn(handle, byref(c_vgpu_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no supported vGPUs
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value
         c_vgpu_type_ids = vgpu_type_ids_array()
@@ -4323,132 +5043,148 @@ def nvmlDeviceGetCreatableVgpus(handle):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId):
     c_profile_id = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId")
     ret = fn(vgpuTypeId, byref(c_profile_id))
     _nvmlCheckReturn(ret)
-    return (c_profile_id.value)
+    return c_profile_id.value
+
 
 @convertStrBytes
 def nvmlVgpuTypeGetClass(vgpuTypeId):
     c_class = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE)
     c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass")
     ret = fn(vgpuTypeId, c_class, byref(c_buffer_size))
     _nvmlCheckReturn(ret)
     return c_class.value
 
+
 @convertStrBytes
 def nvmlVgpuTypeGetName(vgpuTypeId):
     c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE)
     c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName")
     ret = fn(vgpuTypeId, c_name, byref(c_buffer_size))
     _nvmlCheckReturn(ret)
     return c_name.value
 
+
 def nvmlVgpuTypeGetDeviceID(vgpuTypeId):
-    c_device_id    = c_ulonglong(0)
+    c_device_id = c_ulonglong(0)
     c_subsystem_id = c_ulonglong(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID")
     ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id))
     _nvmlCheckReturn(ret)
     return (c_device_id.value, c_subsystem_id.value)
 
+
 def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId):
     c_fb_size = c_ulonglong(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize")
     ret = fn(vgpuTypeId, byref(c_fb_size))
     _nvmlCheckReturn(ret)
     return c_fb_size.value
 
+
 def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId):
     c_num_heads = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads")
     ret = fn(vgpuTypeId, byref(c_num_heads))
     _nvmlCheckReturn(ret)
     return c_num_heads.value
 
+
 def nvmlVgpuTypeGetResolution(vgpuTypeId):
     c_xdim = c_uint(0)
     c_ydim = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution")
     ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim))
     _nvmlCheckReturn(ret)
     return (c_xdim.value, c_ydim.value)
 
+
 @convertStrBytes
 def nvmlVgpuTypeGetLicense(vgpuTypeId):
     c_license = create_string_buffer(NVML_GRID_LICENSE_BUFFER_SIZE)
     c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense")
     ret = fn(vgpuTypeId, c_license, c_buffer_size)
     _nvmlCheckReturn(ret)
     return c_license.value
 
+
 def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId):
     c_frl_config = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit")
     ret = fn(vgpuTypeId, byref(c_frl_config))
     _nvmlCheckReturn(ret)
     return c_frl_config.value
 
+
 def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId):
     c_gsp_heap = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize")
     ret = fn(vgpuTypeId, byref(c_gsp_heap))
     _nvmlCheckReturn(ret)
     return c_gsp_heap.value
 
+
 def nvmlVgpuTypeGetFbReservation(vgpuTypeId):
     c_fb_reservation = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation")
     ret = fn(vgpuTypeId, byref(c_fb_reservation))
     _nvmlCheckReturn(ret)
     return c_fb_reservation.value
 
+
 def nvmlVgpuInstanceGetRuntimeStateSize(vgpuInstance):
     c_runtime_state = nvmlVgpuRuntimeState_v1_t()
     c_runtime_state.version = VgpuRuntimeState_v1
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize")
     ret = fn(vgpuInstance, byref(c_runtime_state))
     _nvmlCheckReturn(ret)
     return c_runtime_state
 
+
 def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId):
     c_max_instances = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances")
     ret = fn(handle, vgpuTypeId, byref(c_max_instances))
     _nvmlCheckReturn(ret)
     return c_max_instances.value
 
+
 def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId):
     c_max_instances_per_vm = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm")
     ret = fn(vgpuTypeId, byref(c_max_instances_per_vm))
     _nvmlCheckReturn(ret)
     return c_max_instances_per_vm.value
 
+
 def nvmlVgpuTypeGetBAR1Info(vgpuTypeId):
     c_bar1Info = c_nvmlVgpuTypeBar1Info_v1_t(0)
     c_bar1Info.version = VgpuTypeBar1Info_v1
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info")
     ret = fn(vgpuTypeId, byref(c_bar1Info))
     _nvmlCheckReturn(ret)
     return c_bar1Info
 
+
 def nvmlDeviceGetActiveVgpus(handle):
     # first call to get the size
     c_vgpu_count = c_uint(0)
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus")
     ret = fn(handle, byref(c_vgpu_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no active vGPUs
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_count.value
         c_vgpu_instances = vgpu_instance_array()
@@ -4464,67 +5200,76 @@ def nvmlDeviceGetActiveVgpus(handle):
         # error case
         raise NVMLError(ret)
 
+
 @convertStrBytes
 def nvmlVgpuInstanceGetVmID(vgpuInstance):
     c_vm_id = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE)
     c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE)
-    c_vm_id_type  = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID")
+    c_vm_id_type = c_uint(0)
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID")
     ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type))
     _nvmlCheckReturn(ret)
     return (c_vm_id.value, c_vm_id_type.value)
 
+
 @convertStrBytes
 def nvmlVgpuInstanceGetUUID(vgpuInstance):
     c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE)
     c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID")
     ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size)
     _nvmlCheckReturn(ret)
     return c_uuid.value
 
+
 @convertStrBytes
 def nvmlVgpuInstanceGetMdevUUID(vgpuInstance):
     c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE)
     c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID")
     ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size)
     _nvmlCheckReturn(ret)
     return c_uuid.value
 
+
 @convertStrBytes
 def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance):
     c_driver_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)
     c_buffer_size = c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion")
     ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size)
     _nvmlCheckReturn(ret)
     return c_driver_version.value
 
+
 def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance):
     c_license_status = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus")
     ret = fn(vgpuInstance, byref(c_license_status))
     _nvmlCheckReturn(ret)
     return c_license_status.value
 
+
 def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance):
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2")
     c_license_info = c_nvmlVgpuLicenseInfo_t()
     ret = fn(vgpuInstance, byref(c_license_info))
     _nvmlCheckReturn(ret)
     return c_license_info
 
+
 def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance):
     return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance)
 
+
 def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance):
     c_frl = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit")
     ret = fn(vgpuInstance, byref(c_frl))
     _nvmlCheckReturn(ret)
     return c_frl.value
 
+
 def nvmlVgpuInstanceGetEccMode(vgpuInstance):
     c_mode = _nvmlEnableState_t()
     fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEccMode")
@@ -4532,79 +5277,98 @@ def nvmlVgpuInstanceGetEccMode(vgpuInstance):
     _nvmlCheckReturn(ret)
     return c_mode.value
 
+
 def nvmlVgpuInstanceGetType(vgpuInstance):
     c_vgpu_type = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType")
     ret = fn(vgpuInstance, byref(c_vgpu_type))
     _nvmlCheckReturn(ret)
     return c_vgpu_type.value
 
+
 def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance):
     c_encoder_capacity = c_ulonglong(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity")
     ret = fn(vgpuInstance, byref(c_encoder_capacity))
     _nvmlCheckReturn(ret)
     return c_encoder_capacity.value
 
+
 def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity):
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity")
     return fn(vgpuInstance, encoder_capacity)
 
+
 def nvmlVgpuInstanceGetFbUsage(vgpuInstance):
     c_fb_usage = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage")
     ret = fn(vgpuInstance, byref(c_fb_usage))
     _nvmlCheckReturn(ret)
     return c_fb_usage.value
 
+
 def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability):
     c_cap_result = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities")
     ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result))
     _nvmlCheckReturn(ret)
-    return (c_cap_result.value)
+    return c_cap_result.value
+
 
 def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance):
     c_id = c_uint(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId")
     ret = fn(vgpuInstance, byref(c_id))
     _nvmlCheckReturn(ret)
-    return (c_id.value)
+    return c_id.value
+
 
 @convertStrBytes
 def nvmlVgpuInstanceGetGpuPciId(vgpuInstance):
     c_vgpuPciId = create_string_buffer(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE)
     fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuPciId")
-    ret = fn(vgpuInstance, c_vgpuPciId, byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE)))
+    ret = fn(
+        vgpuInstance, c_vgpuPciId, byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE))
+    )
     _nvmlCheckReturn(ret)
     return c_vgpuPciId.value
 
+
 def nvmlDeviceGetVgpuUtilization(handle, timeStamp):
     # first call to get the size
     c_vgpu_count = c_uint(0)
     c_time_stamp = c_ulonglong(timeStamp)
     c_sample_value_type = _nvmlValueType_t()
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization")
-    ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None)
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization")
+    ret = fn(
+        handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None
+    )
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no active vGPUs
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         sampleArray = c_vgpu_count.value * c_nvmlVgpuInstanceUtilizationSample_t
         c_samples = sampleArray()
 
         # make the call again
-        ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), c_samples)
+        ret = fn(
+            handle,
+            c_time_stamp,
+            byref(c_sample_value_type),
+            byref(c_vgpu_count),
+            c_samples,
+        )
         _nvmlCheckReturn(ret)
 
-        return c_samples[0:c_vgpu_count.value]
+        return c_samples[0 : c_vgpu_count.value]
     else:
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp):
     # first call to get the size
     c_time_stamp = c_ulonglong(timeStamp)
@@ -4614,15 +5378,17 @@ def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp):
     c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0)
     c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo")
     ret = fn(handle, byref(c_vgpuUtilInfo))
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no active vGPUs
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
-        sampleArray = c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t
+        sampleArray = (
+            c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t
+        )
         c_samples = sampleArray()
         c_vgpuUtilInfo.vgpuUtilArray = c_samples
 
@@ -4630,29 +5396,33 @@ def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp):
         ret = fn(handle, byref(c_vgpuUtilInfo))
         _nvmlCheckReturn(ret)
 
-        return c_samples[0:c_vgpuUtilInfo.vgpuInstanceCount]
+        return c_samples[0 : c_vgpuUtilInfo.vgpuInstanceCount]
     else:
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex):
     c_p2pstatus = _nvmlGpuP2PStatus_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetP2PStatus")
-    ret = fn(device1, device2,p2pIndex, byref(c_p2pstatus))
+    ret = fn(device1, device2, p2pIndex, byref(c_p2pstatus))
     _nvmlCheckReturn(ret)
     return c_p2pstatus.value
 
+
 def nvmlDeviceGetGridLicensableFeatures_v4(handle):
     c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGridLicensableFeatures_v4")
     ret = fn(handle, byref(c_get_grid_licensable_features))
     _nvmlCheckReturn(ret)
 
-    return (c_get_grid_licensable_features)
+    return c_get_grid_licensable_features
+
 
 def nvmlDeviceGetGridLicensableFeatures(handle):
     return nvmlDeviceGetGridLicensableFeatures_v4(handle)
 
+
 def nvmlDeviceGetGspFirmwareVersion(handle, version=None):
     isUserDefined = version is not None
     if not isUserDefined:
@@ -4662,6 +5432,7 @@ def nvmlDeviceGetGspFirmwareVersion(handle, version=None):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isUserDefined else version.value
 
+
 def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint()):
     isReference = type(isEnabled) is not c_uint
     isEnabledRef = isEnabled if isReference else byref(isEnabled)
@@ -4671,6 +5442,7 @@ def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint(
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else [isEnabled.value, defaultMode.value]
 
+
 def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType):
     c_encoder_capacity = c_ulonglong(0)
     c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType)
@@ -4680,18 +5452,19 @@ def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType):
     _nvmlCheckReturn(ret)
     return c_encoder_capacity.value
 
+
 def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp):
     # first call to get the size
     c_vgpu_count = c_uint(0)
     c_time_stamp = c_ulonglong(timeStamp)
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization")
     ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None)
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no active vGPUs
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         sampleArray = c_vgpu_count.value * c_nvmlVgpuProcessUtilizationSample_t
         c_samples = sampleArray()
@@ -4700,11 +5473,12 @@ def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp):
         ret = fn(handle, c_time_stamp, byref(c_vgpu_count), c_samples)
         _nvmlCheckReturn(ret)
 
-        return c_samples[0:c_vgpu_count.value]
+        return c_samples[0 : c_vgpu_count.value]
     else:
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp):
     # first call to get the size
     c_time_stamp = c_ulonglong(timeStamp)
@@ -4713,15 +5487,17 @@ def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp):
     c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0)
     c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo")
     ret = fn(handle, byref(c_vgpuProcUtilInfo))
 
-    if (ret == NVML_SUCCESS):
+    if ret == NVML_SUCCESS:
         # special case, no active vGPUs
         return []
-    elif (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
-        sampleArray = c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t
+        sampleArray = (
+            c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t
+        )
         c_samples = sampleArray()
         c_vgpuProcUtilInfo.vgpuProcUtilArray = c_samples
 
@@ -4729,11 +5505,12 @@ def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp):
         ret = fn(handle, byref(c_vgpuProcUtilInfo))
         _nvmlCheckReturn(ret)
 
-        return c_samples[0:c_vgpuProcUtilInfo.vgpuProcessCount]
+        return c_samples[0 : c_vgpuProcUtilInfo.vgpuProcessCount]
     else:
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetEncoderStats(handle):
     c_encoderCount = c_ulonglong(0)
     c_encodeFps = c_ulonglong(0)
@@ -4743,15 +5520,16 @@ def nvmlDeviceGetEncoderStats(handle):
     _nvmlCheckReturn(ret)
     return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value)
 
+
 def nvmlDeviceGetEncoderSessions(handle):
     # first call to get the size
     c_session_count = c_uint(0)
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions")
     ret = fn(handle, byref(c_session_count), None)
 
-    if (ret == NVML_SUCCESS):
-        if (c_session_count.value != 0):
+    if ret == NVML_SUCCESS:
+        if c_session_count.value != 0:
             # typical case
             session_array = c_nvmlEncoderSession_t * c_session_count.value
             c_sessions = session_array()
@@ -4769,6 +5547,7 @@ def nvmlDeviceGetEncoderSessions(handle):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetFBCStats(handle):
     c_fbcStats = c_nvmlFBCStats_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCStats")
@@ -4776,15 +5555,16 @@ def nvmlDeviceGetFBCStats(handle):
     _nvmlCheckReturn(ret)
     return c_fbcStats
 
+
 def nvmlDeviceGetFBCSessions(handle):
     # first call to get the size
     c_session_count = c_uint(0)
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions")
     ret = fn(handle, byref(c_session_count), None)
 
-    if (ret == NVML_SUCCESS):
-        if (c_session_count.value != 0):
+    if ret == NVML_SUCCESS:
+        if c_session_count.value != 0:
             # typical case
             session_array = c_nvmlFBCSession_t * c_session_count.value
             c_sessions = session_array()
@@ -4802,24 +5582,28 @@ def nvmlDeviceGetFBCSessions(handle):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlVgpuInstanceGetEncoderStats(vgpuInstance):
-    c_encoderCount    = c_ulonglong(0)
-    c_encodeFps       = c_ulonglong(0)
-    c_encoderLatency  = c_ulonglong(0)
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats")
-    ret = fn(vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency))
+    c_encoderCount = c_ulonglong(0)
+    c_encodeFps = c_ulonglong(0)
+    c_encoderLatency = c_ulonglong(0)
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats")
+    ret = fn(
+        vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)
+    )
     _nvmlCheckReturn(ret)
     return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value)
 
+
 def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance):
     # first call to get the size
     c_session_count = c_uint(0)
 
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions")
     ret = fn(vgpuInstance, byref(c_session_count), None)
 
-    if (ret == NVML_SUCCESS):
-        if (c_session_count.value != 0):
+    if ret == NVML_SUCCESS:
+        if c_session_count.value != 0:
             # typical case
             session_array = c_nvmlEncoderSession_t * c_session_count.value
             c_sessions = session_array()
@@ -4837,6 +5621,7 @@ def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlVgpuInstanceGetFBCStats(vgpuInstance):
     c_fbcStats = c_nvmlFBCStats_t()
     fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCStats")
@@ -4844,15 +5629,16 @@ def nvmlVgpuInstanceGetFBCStats(vgpuInstance):
     _nvmlCheckReturn(ret)
     return c_fbcStats
 
+
 def nvmlVgpuInstanceGetFBCSessions(vgpuInstance):
     # first call to get the size
     c_session_count = c_uint(0)
 
-    fn  = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions")
+    fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions")
     ret = fn(vgpuInstance, byref(c_session_count), None)
 
-    if (ret == NVML_SUCCESS):
-        if (c_session_count.value != 0):
+    if ret == NVML_SUCCESS:
+        if c_session_count.value != 0:
             # typical case
             session_array = c_nvmlFBCSession_t * c_session_count.value
             c_sessions = session_array()
@@ -4870,15 +5656,16 @@ def nvmlVgpuInstanceGetFBCSessions(vgpuInstance):
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetProcessUtilization(handle, timeStamp):
     # first call to get the size
     c_count = c_uint(0)
     c_time_stamp = c_ulonglong(timeStamp)
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization")
     ret = fn(handle, None, byref(c_count), c_time_stamp)
 
-    if (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    if ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
         sampleArray = c_count.value * c_nvmlProcessUtilizationSample_t
         c_samples = sampleArray()
@@ -4887,11 +5674,12 @@ def nvmlDeviceGetProcessUtilization(handle, timeStamp):
         ret = fn(handle, c_samples, byref(c_count), c_time_stamp)
         _nvmlCheckReturn(ret)
 
-        return c_samples[0:c_count.value]
+        return c_samples[0 : c_count.value]
     else:
         # error case
         raise NVMLError(ret)
 
+
 def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp):
     # first call to get the size
     c_time_stamp = c_ulonglong(timeStamp)
@@ -4900,12 +5688,14 @@ def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp):
     c_processesUtilInfo.processSamplesCount = c_uint(0)
     c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp
 
-    fn  = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo")
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo")
     ret = fn(handle, byref(c_processesUtilInfo))
 
-    if (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    if ret == NVML_ERROR_INSUFFICIENT_SIZE:
         # typical case
-        sampleArray = c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t
+        sampleArray = (
+            c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t
+        )
         c_samples = sampleArray()
         c_processesUtilInfo.procUtilArray = c_samples
 
@@ -4913,11 +5703,12 @@ def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp):
         ret = fn(handle, byref(c_processesUtilInfo))
         _nvmlCheckReturn(ret)
 
-        return c_samples[0:c_processesUtilInfo.processSamplesCount]
+        return c_samples[0 : c_processesUtilInfo.processSamplesCount]
     else:
         # error case
         raise NVMLError(ret)
 
+
 def nvmlVgpuInstanceGetMetadata(vgpuInstance):
     fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMetadata")
     c_vgpuMetadata = c_nvmlVgpuMetadata_t()
@@ -4925,13 +5716,14 @@ def nvmlVgpuInstanceGetMetadata(vgpuInstance):
     # Make the first NVML API call to get the c_bufferSize value.
     # We have already allocated required buffer above.
     ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize))
-    if (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    if ret == NVML_ERROR_INSUFFICIENT_SIZE:
         ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize))
         _nvmlCheckReturn(ret)
     else:
         raise NVMLError(ret)
     return c_vgpuMetadata
 
+
 def nvmlDeviceGetVgpuMetadata(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuMetadata")
     c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t()
@@ -4939,13 +5731,14 @@ def nvmlDeviceGetVgpuMetadata(handle):
     # Make the first NVML API call to get the c_bufferSize value.
     # We have already allocated required buffer above.
     ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize))
-    if (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    if ret == NVML_ERROR_INSUFFICIENT_SIZE:
         ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize))
         _nvmlCheckReturn(ret)
     else:
         raise NVMLError(ret)
     return c_vgpuPgpuMetadata
 
+
 def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata):
     fn = _nvmlGetFunctionPointer("nvmlGetVgpuCompatibility")
     c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t()
@@ -4953,6 +5746,7 @@ def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata):
     _nvmlCheckReturn(ret)
     return c_vgpuPgpuCompatibility
 
+
 @convertStrBytes
 def nvmlDeviceGetPgpuMetadataString(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPgpuMetadataString")
@@ -4961,13 +5755,14 @@ def nvmlDeviceGetPgpuMetadataString(handle):
     # Make the first NVML API call to get the c_bufferSize value.
     # We have already allocated required buffer above.
     ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize))
-    if (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    if ret == NVML_ERROR_INSUFFICIENT_SIZE:
         ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize))
         _nvmlCheckReturn(ret)
     else:
         raise NVMLError(ret)
     return (c_pgpuMetadata.value, c_bufferSize.value)
 
+
 def nvmlDeviceGetVgpuSchedulerLog(handle):
     c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerLog")
@@ -4975,6 +5770,7 @@ def nvmlDeviceGetVgpuSchedulerLog(handle):
     _nvmlCheckReturn(ret)
     return c_vgpu_sched_log
 
+
 def nvmlDeviceGetVgpuSchedulerState(handle):
     c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerState")
@@ -4982,6 +5778,7 @@ def nvmlDeviceGetVgpuSchedulerState(handle):
     _nvmlCheckReturn(ret)
     return c_vgpu_sched_state
 
+
 def nvmlDeviceGetVgpuSchedulerCapabilities(handle):
     c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerCapabilities")
@@ -4989,18 +5786,21 @@ def nvmlDeviceGetVgpuSchedulerCapabilities(handle):
     _nvmlCheckReturn(ret)
     return c_vgpu_sched_caps
 
+
 def nvmlDeviceSetVgpuSchedulerState(handle, sched_state):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuSchedulerState")
     ret = fn(handle, byref(sched_state))
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlSetVgpuVersion(vgpuVersion):
     fn = _nvmlGetFunctionPointer("nvmlSetVgpuVersion")
     ret = fn(byref(vgpuVersion))
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlGetVgpuVersion(supported=None, current=None):
     isUserDefined = (supported is not None) or (current is not None)
     if not isUserDefined:
@@ -5009,10 +5809,15 @@ def nvmlGetVgpuVersion(supported=None, current=None):
     fn = _nvmlGetFunctionPointer("nvmlGetVgpuVersion")
     ret = fn(byref(supported), byref(current))
     _nvmlCheckReturn(ret)
-    return NVML_SUCCESS if isUserDefined else [(supported.minVersion,
-                                                supported.maxVersion),
-                                               (current.minVersion,
-                                                current.maxVersion)]
+    return (
+        NVML_SUCCESS
+        if isUserDefined
+        else [
+            (supported.minVersion, supported.maxVersion),
+            (current.minVersion, current.maxVersion),
+        ]
+    )
+
 
 def nvmlVgpuInstanceGetAccountingMode(vgpuInstance):
     c_mode = _nvmlEnableState_t()
@@ -5021,11 +5826,12 @@ def nvmlVgpuInstanceGetAccountingMode(vgpuInstance):
     _nvmlCheckReturn(ret)
     return c_mode.value
 
+
 def nvmlVgpuInstanceGetAccountingPids(vgpuInstance):
     c_pidCount = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingPids")
     ret = fn(vgpuInstance, byref(c_pidCount), None)
-    if (ret == NVML_ERROR_INSUFFICIENT_SIZE):
+    if ret == NVML_ERROR_INSUFFICIENT_SIZE:
         sampleArray = c_pidCount.value * c_uint
         c_pidArray = sampleArray()
         ret = fn(vgpuInstance, byref(c_pidCount), byref(c_pidArray))
@@ -5034,6 +5840,7 @@ def nvmlVgpuInstanceGetAccountingPids(vgpuInstance):
         raise NVMLError(ret)
     return (c_pidCount, c_pidArray)
 
+
 def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid):
     c_accountingStats = c_nvmlAccountingStats_t()
     fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingStats")
@@ -5041,12 +5848,14 @@ def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid):
     _nvmlCheckReturn(ret)
     return c_accountingStats
 
+
 def nvmlVgpuInstanceClearAccountingPids(vgpuInstance):
     fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceClearAccountingPids")
     ret = fn(vgpuInstance)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlGetExcludedDeviceCount():
     c_count = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceCount")
@@ -5054,6 +5863,7 @@ def nvmlGetExcludedDeviceCount():
     _nvmlCheckReturn(ret)
     return c_count.value
 
+
 def nvmlGetExcludedDeviceInfoByIndex(index):
     c_index = c_uint(index)
     info = c_nvmlExcludedDeviceInfo_t()
@@ -5062,6 +5872,7 @@ def nvmlGetExcludedDeviceInfoByIndex(index):
     _nvmlCheckReturn(ret)
     return info
 
+
 def nvmlDeviceGetHostVgpuMode(handle):
     c_host_vgpu_mode = _nvmlHostVgpuMode_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetHostVgpuMode")
@@ -5069,6 +5880,7 @@ def nvmlDeviceGetHostVgpuMode(handle):
     _nvmlCheckReturn(ret)
     return c_host_vgpu_mode.value
 
+
 def nvmlDeviceSetMigMode(device, mode):
     c_activationStatus = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetMigMode")
@@ -5076,6 +5888,7 @@ def nvmlDeviceSetMigMode(device, mode):
     _nvmlCheckReturn(ret)
     return c_activationStatus.value
 
+
 def nvmlDeviceGetMigMode(device):
     c_currentMode = c_uint()
     c_pendingMode = c_uint()
@@ -5084,6 +5897,7 @@ def nvmlDeviceGetMigMode(device):
     _nvmlCheckReturn(ret)
     return [c_currentMode.value, c_pendingMode.value]
 
+
 def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2):
     if version == 2:
         c_info = c_nvmlGpuInstanceProfileInfo_v2_t()
@@ -5097,9 +5911,11 @@ def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2):
     _nvmlCheckReturn(ret)
     return c_info
 
+
 # Define function alias for the API exposed by NVML
 nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo
 
+
 def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId):
     c_count = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceRemainingCapacity")
@@ -5107,12 +5923,16 @@ def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId):
     _nvmlCheckReturn(ret)
     return c_count.value
 
-def nvmlDeviceGetGpuInstancePossiblePlacements(device, profileId, placementsRef, countRef):
+
+def nvmlDeviceGetGpuInstancePossiblePlacements(
+    device, profileId, placementsRef, countRef
+):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstancePossiblePlacements_v2")
     ret = fn(device, profileId, placementsRef, countRef)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceCreateGpuInstance(device, profileId):
     c_instance = c_nvmlGpuInstance_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstance")
@@ -5120,6 +5940,7 @@ def nvmlDeviceCreateGpuInstance(device, profileId):
     _nvmlCheckReturn(ret)
     return c_instance
 
+
 def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement):
     c_instance = c_nvmlGpuInstance_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstanceWithPlacement")
@@ -5127,18 +5948,21 @@ def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement):
     _nvmlCheckReturn(ret)
     return c_instance
 
+
 def nvmlGpuInstanceDestroy(gpuInstance):
     fn = _nvmlGetFunctionPointer("nvmlGpuInstanceDestroy")
     ret = fn(gpuInstance)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstances")
     ret = fn(device, profileId, gpuInstancesRef, countRef)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId):
     c_instance = c_nvmlGpuInstance_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceById")
@@ -5146,6 +5970,7 @@ def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId):
     _nvmlCheckReturn(ret)
     return c_instance
 
+
 def nvmlGpuInstanceGetInfo(gpuInstance):
     c_info = c_nvmlGpuInstanceInfo_t()
     fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetInfo")
@@ -5153,7 +5978,10 @@ def nvmlGpuInstanceGetInfo(gpuInstance):
     _nvmlCheckReturn(ret)
     return c_info
 
-def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, version=2):
+
+def nvmlGpuInstanceGetComputeInstanceProfileInfo(
+    device, profile, engProfile, version=2
+):
     if version == 2:
         c_info = c_nvmlComputeInstanceProfileInfo_v2_t()
         fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfoV")
@@ -5161,13 +5989,17 @@ def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, ve
         c_info = c_nvmlComputeInstanceProfileInfo_t()
         fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfo")
     else:
-        raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) 
+        raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND)
     ret = fn(device, profile, engProfile, byref(c_info))
     _nvmlCheckReturn(ret)
     return c_info
 
+
 # Define function alias for the API exposed by NVML
-nvmlGpuInstanceGetComputeInstanceProfileInfoV = nvmlGpuInstanceGetComputeInstanceProfileInfo
+nvmlGpuInstanceGetComputeInstanceProfileInfoV = (
+    nvmlGpuInstanceGetComputeInstanceProfileInfo
+)
+
 
 def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId):
     c_count = c_uint()
@@ -5176,12 +6008,16 @@ def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId):
     _nvmlCheckReturn(ret)
     return c_count.value
 
-def nvmlGpuInstanceGetComputeInstancePossiblePlacements(gpuInstance, profileId, placementsRef, countRef):
+
+def nvmlGpuInstanceGetComputeInstancePossiblePlacements(
+    gpuInstance, profileId, placementsRef, countRef
+):
     fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstancePossiblePlacements")
     ret = fn(gpuInstance, profileId, placementsRef, countRef)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId):
     c_instance = c_nvmlComputeInstance_t()
     fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstance")
@@ -5189,25 +6025,33 @@ def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId):
     _nvmlCheckReturn(ret)
     return c_instance
 
-def nvmlGpuInstanceCreateComputeInstanceWithPlacement(gpuInstance, profileId, placement):
+
+def nvmlGpuInstanceCreateComputeInstanceWithPlacement(
+    gpuInstance, profileId, placement
+):
     c_instance = c_nvmlComputeInstance_t()
     fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstanceWithPlacement")
     ret = fn(gpuInstance, profileId, placement, byref(c_instance))
     _nvmlCheckReturn(ret)
     return c_instance
 
+
 def nvmlComputeInstanceDestroy(computeInstance):
     fn = _nvmlGetFunctionPointer("nvmlComputeInstanceDestroy")
     ret = fn(computeInstance)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
-def nvmlGpuInstanceGetComputeInstances(gpuInstance, profileId, computeInstancesRef, countRef):
+
+def nvmlGpuInstanceGetComputeInstances(
+    gpuInstance, profileId, computeInstancesRef, countRef
+):
     fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstances")
     ret = fn(gpuInstance, profileId, computeInstancesRef, countRef)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId):
     c_instance = c_nvmlComputeInstance_t()
     fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceById")
@@ -5215,6 +6059,7 @@ def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId):
     _nvmlCheckReturn(ret)
     return c_instance
 
+
 def nvmlComputeInstanceGetInfo_v2(computeInstance):
     c_info = c_nvmlComputeInstanceInfo_t()
     fn = _nvmlGetFunctionPointer("nvmlComputeInstanceGetInfo_v2")
@@ -5222,9 +6067,11 @@ def nvmlComputeInstanceGetInfo_v2(computeInstance):
     _nvmlCheckReturn(ret)
     return c_info
 
+
 def nvmlComputeInstanceGetInfo(computeInstance):
     return nvmlComputeInstanceGetInfo_v2(computeInstance)
 
+
 def nvmlDeviceIsMigDeviceHandle(device):
     c_isMigDevice = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceIsMigDeviceHandle")
@@ -5232,6 +6079,7 @@ def nvmlDeviceIsMigDeviceHandle(device):
     _nvmlCheckReturn(ret)
     return c_isMigDevice
 
+
 def nvmlDeviceGetGpuInstanceId(device):
     c_gpuInstanceId = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceId")
@@ -5239,6 +6087,7 @@ def nvmlDeviceGetGpuInstanceId(device):
     _nvmlCheckReturn(ret)
     return c_gpuInstanceId.value
 
+
 def nvmlDeviceGetComputeInstanceId(device):
     c_computeInstanceId = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeInstanceId")
@@ -5246,6 +6095,7 @@ def nvmlDeviceGetComputeInstanceId(device):
     _nvmlCheckReturn(ret)
     return c_computeInstanceId.value
 
+
 def nvmlDeviceGetMaxMigDeviceCount(device):
     c_count = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxMigDeviceCount")
@@ -5253,6 +6103,7 @@ def nvmlDeviceGetMaxMigDeviceCount(device):
     _nvmlCheckReturn(ret)
     return c_count.value
 
+
 def nvmlDeviceGetMigDeviceHandleByIndex(device, index):
     c_index = c_uint(index)
     migDevice = c_nvmlDevice_t()
@@ -5261,6 +6112,7 @@ def nvmlDeviceGetMigDeviceHandleByIndex(device, index):
     _nvmlCheckReturn(ret)
     return migDevice
 
+
 def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice):
     device = c_nvmlDevice_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetDeviceHandleFromMigDeviceHandle")
@@ -5268,6 +6120,7 @@ def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice):
     _nvmlCheckReturn(ret)
     return device
 
+
 def nvmlDeviceGetAttributes_v2(device):
     c_attrs = c_nvmlDeviceAttributes()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetAttributes_v2")
@@ -5275,9 +6128,11 @@ def nvmlDeviceGetAttributes_v2(device):
     _nvmlCheckReturn(ret)
     return c_attrs
 
+
 def nvmlDeviceGetAttributes(device):
     return nvmlDeviceGetAttributes_v2(device)
 
+
 def nvmlDeviceGetRemappedRows(device):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetRemappedRows")
     c_corr = c_uint()
@@ -5288,6 +6143,7 @@ def nvmlDeviceGetRemappedRows(device):
     _nvmlCheckReturn(ret)
     return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value)
 
+
 def nvmlDeviceGetRowRemapperHistogram(device):
     c_vals = c_nvmlRowRemapperHistogramValues()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetRowRemapperHistogram")
@@ -5295,6 +6151,7 @@ def nvmlDeviceGetRowRemapperHistogram(device):
     _nvmlCheckReturn(ret)
     return c_vals
 
+
 def nvmlDeviceGetArchitecture(device):
     arch = _nvmlDeviceArchitecture_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetArchitecture")
@@ -5302,6 +6159,7 @@ def nvmlDeviceGetArchitecture(device):
     _nvmlCheckReturn(ret)
     return arch.value
 
+
 def nvmlDeviceGetBusType(device):
     c_busType = _nvmlBusType_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetBusType")
@@ -5309,6 +6167,7 @@ def nvmlDeviceGetBusType(device):
     _nvmlCheckReturn(ret)
     return c_busType.value
 
+
 def nvmlDeviceGetIrqNum(device):
     c_irqNum = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetIrqNum")
@@ -5316,6 +6175,7 @@ def nvmlDeviceGetIrqNum(device):
     _nvmlCheckReturn(ret)
     return c_irqNum.value
 
+
 def nvmlDeviceGetNumGpuCores(device):
     c_numCores = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumGpuCores")
@@ -5323,6 +6183,7 @@ def nvmlDeviceGetNumGpuCores(device):
     _nvmlCheckReturn(ret)
     return c_numCores.value
 
+
 def nvmlDeviceGetPowerSource(device):
     c_powerSource = _nvmlPowerSource_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerSource")
@@ -5330,6 +6191,7 @@ def nvmlDeviceGetPowerSource(device):
     _nvmlCheckReturn(ret)
     return c_powerSource.value
 
+
 def nvmlDeviceGetMemoryBusWidth(device):
     c_memBusWidth = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryBusWidth")
@@ -5337,6 +6199,7 @@ def nvmlDeviceGetMemoryBusWidth(device):
     _nvmlCheckReturn(ret)
     return c_memBusWidth.value
 
+
 def nvmlDeviceGetPcieLinkMaxSpeed(device):
     c_speed = _nvmlPcieLinkMaxSpeed_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieLinkMaxSpeed")
@@ -5344,6 +6207,7 @@ def nvmlDeviceGetPcieLinkMaxSpeed(device):
     _nvmlCheckReturn(ret)
     return c_speed.value
 
+
 def nvmlDeviceGetAdaptiveClockInfoStatus(device):
     c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetAdaptiveClockInfoStatus")
@@ -5351,6 +6215,7 @@ def nvmlDeviceGetAdaptiveClockInfoStatus(device):
     _nvmlCheckReturn(ret)
     return c_adaptiveClockInfoStatus.value
 
+
 def nvmlDeviceGetPcieSpeed(device):
     c_speed = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieSpeed")
@@ -5358,64 +6223,87 @@ def nvmlDeviceGetPcieSpeed(device):
     _nvmlCheckReturn(ret)
     return c_speed.value
 
-def nvmlDeviceGetDynamicPstatesInfo(device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t()):
+
+def nvmlDeviceGetDynamicPstatesInfo(
+    device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t()
+):
     isReference = type(c_dynamicpstatesinfo) is not c_nvmlGpuDynamicPstatesInfo_t
-    dynamicpstatesinfoRef = c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo)
+    dynamicpstatesinfoRef = (
+        c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo)
+    )
 
-    fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo");
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo")
     ret = fn(device, dynamicpstatesinfoRef)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else c_dynamicpstatesinfo
 
+
 def nvmlDeviceSetFanSpeed_v2(handle, index, speed):
-    fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2");
+    fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2")
     ret = fn(handle, index, speed)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
-def nvmlDeviceGetThermalSettings(device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t()):
+
+def nvmlDeviceGetThermalSettings(
+    device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t()
+):
     isReference = type(c_thermalsettings) is not c_nvmlGpuThermalSettings_t
     thermalsettingsRef = c_thermalsettings if isReference else byref(c_thermalsettings)
-    fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings");
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings")
     ret = fn(device, sensorindex, thermalsettingsRef)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else c_thermalsettings.sensor[:]
 
-def nvmlDeviceGetMinMaxClockOfPState(device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint()):
+
+def nvmlDeviceGetMinMaxClockOfPState(
+    device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint()
+):
     isReference = (type(minClockMHz) is not c_uint) or (type(maxClockMHz) is not c_uint)
     minClockMHzRef = minClockMHz if isReference else byref(minClockMHz)
     maxClockMHzRef = maxClockMHz if isReference else byref(maxClockMHz)
-    fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState");
-    ret = fn(device, _nvmlClockType_t(clockType), _nvmlClockType_t(pstate), minClockMHzRef, maxClockMHzRef)
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState")
+    ret = fn(
+        device,
+        _nvmlClockType_t(clockType),
+        _nvmlClockType_t(pstate),
+        minClockMHzRef,
+        maxClockMHzRef,
+    )
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else (minClockMHz.value, maxClockMHz.value)
 
+
 class c_nvmlClockOffset_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('type', _nvmlClockType_t),
-        ('pstate', _nvmlPstates_t),
-        ('clockOffsetMHz', c_int),
-        ('minClockOffsetMHz', c_int),
-        ('maxClockOffsetMHz', c_int),
+        ("version", c_uint),
+        ("type", _nvmlClockType_t),
+        ("pstate", _nvmlPstates_t),
+        ("clockOffsetMHz", c_int),
+        ("minClockOffsetMHz", c_int),
+        ("maxClockOffsetMHz", c_int),
     ]
 
+
 nvmlClockOffset_v1 = 0x1000018
 
+
 def nvmlDeviceGetClockOffsets(device, info):
-    fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets");
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets")
     ret = fn(device, info)
     return NVML_SUCCESS
 
+
 def nvmlDeviceSetClockOffsets(device, info):
-    fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets");
+    fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets")
     ret = fn(device, info)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetSupportedPerformanceStates(device):
     pstates = []
     c_count = c_uint(NVML_MAX_GPU_PERF_PSTATES)
-    c_size = sizeof(c_uint)*c_count.value
+    c_size = sizeof(c_uint) * c_count.value
 
     # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration.
     pstates_array = _nvmlPstates_t * c_count.value
@@ -5431,6 +6319,7 @@ def nvmlDeviceGetSupportedPerformanceStates(device):
 
     return pstates
 
+
 def nvmlDeviceGetGpcClkVfOffset(device):
     offset = c_int32()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkVfOffset")
@@ -5438,6 +6327,7 @@ def nvmlDeviceGetGpcClkVfOffset(device):
     _nvmlCheckReturn(ret)
     return offset.value
 
+
 def nvmlDeviceSetGpcClkVfOffset(device, offset):
     c_offset = c_int32(offset)
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpcClkVfOffset")
@@ -5445,6 +6335,7 @@ def nvmlDeviceSetGpcClkVfOffset(device, offset):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()):
     isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int)
     minOffsetRef = minOffset if isReference else byref(minOffset)
@@ -5454,6 +6345,7 @@ def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value)
 
+
 def nvmlDeviceGetMemClkVfOffset(device):
     offset = c_int32()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkVfOffset")
@@ -5461,6 +6353,7 @@ def nvmlDeviceGetMemClkVfOffset(device):
     _nvmlCheckReturn(ret)
     return offset.value
 
+
 def nvmlDeviceSetMemClkVfOffset(device, offset):
     c_offset = c_int32(offset)
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemClkVfOffset")
@@ -5468,6 +6361,7 @@ def nvmlDeviceSetMemClkVfOffset(device, offset):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()):
     isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int)
     minOffsetRef = minOffset if isReference else byref(minOffset)
@@ -5478,6 +6372,7 @@ def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value)
 
+
 def nvmlSystemSetConfComputeGpusReadyState(state):
     c_state = c_uint(state)
     fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeGpusReadyState")
@@ -5485,6 +6380,7 @@ def nvmlSystemSetConfComputeGpusReadyState(state):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlSystemGetConfComputeGpusReadyState():
     c_state = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeGpusReadyState")
@@ -5492,6 +6388,7 @@ def nvmlSystemGetConfComputeGpusReadyState():
     _nvmlCheckReturn(ret)
     return c_state.value
 
+
 def nvmlSystemGetConfComputeCapabilities():
     c_ccSysCaps = c_nvmlConfComputeSystemCaps_t()
     fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeCapabilities")
@@ -5499,6 +6396,7 @@ def nvmlSystemGetConfComputeCapabilities():
     _nvmlCheckReturn(ret)
     return c_ccSysCaps
 
+
 def nvmlSystemGetConfComputeState():
     c_state = c_nvmlConfComputeSystemState_t()
     fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeState")
@@ -5506,16 +6404,19 @@ def nvmlSystemGetConfComputeState():
     _nvmlCheckReturn(ret)
     return c_state
 
+
 def nvmlSystemGetConfComputeSettings(settings):
     fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeSettings")
     return fn(settings)
 
+
 def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetConfComputeUnprotectedMemSize")
     ret = fn(device, c_ccMemSize)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetConfComputeMemSizeInfo(device):
     c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeMemSizeInfo")
@@ -5523,6 +6424,7 @@ def nvmlDeviceGetConfComputeMemSizeInfo(device):
     _nvmlCheckReturn(ret)
     return c_ccMemSize
 
+
 def nvmlDeviceGetConfComputeProtectedMemoryUsage(device):
     c_memory = c_nvmlMemory_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeProtectedMemoryUsage")
@@ -5530,6 +6432,7 @@ def nvmlDeviceGetConfComputeProtectedMemoryUsage(device):
     _nvmlCheckReturn(ret)
     return c_memory
 
+
 def nvmlDeviceGetConfComputeGpuCertificate(device):
     c_cert = c_nvmlConfComputeGpuCertificate_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuCertificate")
@@ -5537,15 +6440,17 @@ def nvmlDeviceGetConfComputeGpuCertificate(device):
     _nvmlCheckReturn(ret)
     return c_cert
 
+
 def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce):
     c_attestReport = c_nvmlConfComputeGpuAttestationReport_t()
     c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce))
-    setattr(c_attestReport, 'nonce', c_nonce_arr)
+    setattr(c_attestReport, "nonce", c_nonce_arr)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuAttestationReport")
     ret = fn(device, byref(c_attestReport))
     _nvmlCheckReturn(ret)
     return c_attestReport
 
+
 def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv):
     c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0)
     c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1
@@ -5555,6 +6460,7 @@ def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlSystemGetConfComputeKeyRotationThresholdInfo():
     c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0)
     c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1
@@ -5563,98 +6469,182 @@ def nvmlSystemGetConfComputeKeyRotationThresholdInfo():
     _nvmlCheckReturn(ret)
     return c_keyRotationThrInfo
 
+
 ## GPM ##
 #########
 
 ## Enums/defines
 
 #### GPM Metric Identifiers
-NVML_GPM_METRIC_GRAPHICS_UTIL               = 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0
-NVML_GPM_METRIC_SM_UTIL                     = 2 # Percentage of SMs that were busy. 0.0 - 100.0
-NVML_GPM_METRIC_SM_OCCUPANCY                = 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0
-NVML_GPM_METRIC_INTEGER_UTIL                = 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0
-NVML_GPM_METRIC_ANY_TENSOR_UTIL             = 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0
-NVML_GPM_METRIC_DFMA_TENSOR_UTIL            = 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0
-NVML_GPM_METRIC_HMMA_TENSOR_UTIL            = 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0
-NVML_GPM_METRIC_IMMA_TENSOR_UTIL            = 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0
-NVML_GPM_METRIC_DRAM_BW_UTIL                = 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0
-NVML_GPM_METRIC_FP64_UTIL                   = 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0
-NVML_GPM_METRIC_FP32_UTIL                   = 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0
-NVML_GPM_METRIC_FP16_UTIL                   = 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0
-NVML_GPM_METRIC_PCIE_TX_PER_SEC             = 20 # PCIe traffic from this GPU in MiB/sec
-NVML_GPM_METRIC_PCIE_RX_PER_SEC             = 21 # PCIe traffic to this GPU in MiB/sec
-NVML_GPM_METRIC_NVDEC_0_UTIL                = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0
-NVML_GPM_METRIC_NVDEC_1_UTIL                = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0
-NVML_GPM_METRIC_NVDEC_2_UTIL                = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0
-NVML_GPM_METRIC_NVDEC_3_UTIL                = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0
-NVML_GPM_METRIC_NVDEC_4_UTIL                = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0
-NVML_GPM_METRIC_NVDEC_5_UTIL                = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0
-NVML_GPM_METRIC_NVDEC_6_UTIL                = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0
-NVML_GPM_METRIC_NVDEC_7_UTIL                = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0
-NVML_GPM_METRIC_NVJPG_0_UTIL                = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0
-NVML_GPM_METRIC_NVJPG_1_UTIL                = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0
-NVML_GPM_METRIC_NVJPG_2_UTIL                = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0
-NVML_GPM_METRIC_NVJPG_3_UTIL                = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0
-NVML_GPM_METRIC_NVJPG_4_UTIL                = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0
-NVML_GPM_METRIC_NVJPG_5_UTIL                = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0
-NVML_GPM_METRIC_NVJPG_6_UTIL                = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0
-NVML_GPM_METRIC_NVJPG_7_UTIL                = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0
-NVML_GPM_METRIC_NVOFA_0_UTIL                = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0
-NVML_GPM_METRIC_NVOFA_1_UTIL                = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0
-NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC     = 60 # NvLink read bandwidth for all links in MiB/sec
-NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC     = 61 # NvLink write bandwidth for all links in MiB/sec
-NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC        = 62 # NvLink read bandwidth for link 0 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC        = 63 # NvLink write bandwidth for link 0 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC        = 64 # NvLink read bandwidth for link 1 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC        = 65 # NvLink write bandwidth for link 1 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC        = 66 # NvLink read bandwidth for link 2 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC        = 67 # NvLink write bandwidth for link 2 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC        = 68 # NvLink read bandwidth for link 3 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC        = 69 # NvLink write bandwidth for link 3 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC        = 70 # NvLink read bandwidth for link 4 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC        = 71 # NvLink write bandwidth for link 4 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC        = 72 # NvLink read bandwidth for link 5 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC        = 73 # NvLink write bandwidth for link 5 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC        = 74 # NvLink read bandwidth for link 6 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC        = 75 # NvLink write bandwidth for link 6 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC        = 76 # NvLink read bandwidth for link 7 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC        = 77 # NvLink write bandwidth for link 7 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC        = 78 # NvLink read bandwidth for link 8 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC        = 79 # NvLink write bandwidth for link 8 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC        = 80 # NvLink read bandwidth for link 9 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC        = 81 # NvLink write bandwidth for link 9 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC       = 82 # NvLink read bandwidth for link 10 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC       = 83 # NvLink write bandwidth for link 10 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC       = 84 # NvLink read bandwidth for link 11 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC       = 85 # NvLink write bandwidth for link 11 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC       = 86 # NvLink read bandwidth for link 12 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC       = 87 # NvLink write bandwidth for link 12 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC       = 88 # NvLink read bandwidth for link 13 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC       = 89 # NvLink write bandwidth for link 13 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC       = 90 # NvLink read bandwidth for link 14 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC       = 91 # NvLink write bandwidth for link 14 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC       = 92 # NvLink read bandwidth for link 15 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC       = 93 # NvLink write bandwidth for link 15 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC       = 94 # NvLink read bandwidth for link 16 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC       = 95 # NvLink write bandwidth for link 16 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC       = 96 # NvLink read bandwidth for link 17 in MiB/sec
-NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC       = 97 # NvLink write bandwidth for link 17 in MiB/sec
-NVML_GPM_METRIC_MAX                         = 98
+NVML_GPM_METRIC_GRAPHICS_UTIL = (
+    1  # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0
+)
+NVML_GPM_METRIC_SM_UTIL = 2  # Percentage of SMs that were busy. 0.0 - 100.0
+NVML_GPM_METRIC_SM_OCCUPANCY = (
+    3  # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0
+)
+NVML_GPM_METRIC_INTEGER_UTIL = (
+    4  # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0
+)
+NVML_GPM_METRIC_ANY_TENSOR_UTIL = (
+    5  # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0
+)
+NVML_GPM_METRIC_DFMA_TENSOR_UTIL = (
+    6  # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0
+)
+NVML_GPM_METRIC_HMMA_TENSOR_UTIL = (
+    7  # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0
+)
+NVML_GPM_METRIC_IMMA_TENSOR_UTIL = (
+    9  # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0
+)
+NVML_GPM_METRIC_DRAM_BW_UTIL = (
+    10  # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0
+)
+NVML_GPM_METRIC_FP64_UTIL = (
+    11  # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0
+)
+NVML_GPM_METRIC_FP32_UTIL = (
+    12  # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0
+)
+NVML_GPM_METRIC_FP16_UTIL = (
+    13  # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0
+)
+NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20  # PCIe traffic from this GPU in MiB/sec
+NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21  # PCIe traffic to this GPU in MiB/sec
+NVML_GPM_METRIC_NVDEC_0_UTIL = 30  # Percent utilization of NVDEC 0. 0.0 - 100.0
+NVML_GPM_METRIC_NVDEC_1_UTIL = 31  # Percent utilization of NVDEC 1. 0.0 - 100.0
+NVML_GPM_METRIC_NVDEC_2_UTIL = 32  # Percent utilization of NVDEC 2. 0.0 - 100.0
+NVML_GPM_METRIC_NVDEC_3_UTIL = 33  # Percent utilization of NVDEC 3. 0.0 - 100.0
+NVML_GPM_METRIC_NVDEC_4_UTIL = 34  # Percent utilization of NVDEC 4. 0.0 - 100.0
+NVML_GPM_METRIC_NVDEC_5_UTIL = 35  # Percent utilization of NVDEC 5. 0.0 - 100.0
+NVML_GPM_METRIC_NVDEC_6_UTIL = 36  # Percent utilization of NVDEC 6. 0.0 - 100.0
+NVML_GPM_METRIC_NVDEC_7_UTIL = 37  # Percent utilization of NVDEC 7. 0.0 - 100.0
+NVML_GPM_METRIC_NVJPG_0_UTIL = 40  # Percent utilization of NVJPG 0. 0.0 - 100.0
+NVML_GPM_METRIC_NVJPG_1_UTIL = 41  # Percent utilization of NVJPG 1. 0.0 - 100.0
+NVML_GPM_METRIC_NVJPG_2_UTIL = 42  # Percent utilization of NVJPG 2. 0.0 - 100.0
+NVML_GPM_METRIC_NVJPG_3_UTIL = 43  # Percent utilization of NVJPG 3. 0.0 - 100.0
+NVML_GPM_METRIC_NVJPG_4_UTIL = 44  # Percent utilization of NVJPG 4. 0.0 - 100.0
+NVML_GPM_METRIC_NVJPG_5_UTIL = 45  # Percent utilization of NVJPG 5. 0.0 - 100.0
+NVML_GPM_METRIC_NVJPG_6_UTIL = 46  # Percent utilization of NVJPG 6. 0.0 - 100.0
+NVML_GPM_METRIC_NVJPG_7_UTIL = 47  # Percent utilization of NVJPG 7. 0.0 - 100.0
+NVML_GPM_METRIC_NVOFA_0_UTIL = 50  # Percent utilization of NVOFA 0. 0.0 - 100.0
+NVML_GPM_METRIC_NVOFA_1_UTIL = 51  # Percent utilization of NVOFA 1. 0.0 - 100.0
+NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = (
+    60  # NvLink read bandwidth for all links in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = (
+    61  # NvLink write bandwidth for all links in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62  # NvLink read bandwidth for link 0 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = (
+    63  # NvLink write bandwidth for link 0 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64  # NvLink read bandwidth for link 1 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = (
+    65  # NvLink write bandwidth for link 1 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66  # NvLink read bandwidth for link 2 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = (
+    67  # NvLink write bandwidth for link 2 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68  # NvLink read bandwidth for link 3 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = (
+    69  # NvLink write bandwidth for link 3 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70  # NvLink read bandwidth for link 4 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = (
+    71  # NvLink write bandwidth for link 4 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72  # NvLink read bandwidth for link 5 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = (
+    73  # NvLink write bandwidth for link 5 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74  # NvLink read bandwidth for link 6 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = (
+    75  # NvLink write bandwidth for link 6 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76  # NvLink read bandwidth for link 7 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = (
+    77  # NvLink write bandwidth for link 7 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78  # NvLink read bandwidth for link 8 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = (
+    79  # NvLink write bandwidth for link 8 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80  # NvLink read bandwidth for link 9 in MiB/sec
+NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = (
+    81  # NvLink write bandwidth for link 9 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = (
+    82  # NvLink read bandwidth for link 10 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = (
+    83  # NvLink write bandwidth for link 10 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = (
+    84  # NvLink read bandwidth for link 11 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = (
+    85  # NvLink write bandwidth for link 11 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = (
+    86  # NvLink read bandwidth for link 12 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = (
+    87  # NvLink write bandwidth for link 12 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = (
+    88  # NvLink read bandwidth for link 13 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = (
+    89  # NvLink write bandwidth for link 13 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = (
+    90  # NvLink read bandwidth for link 14 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = (
+    91  # NvLink write bandwidth for link 14 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = (
+    92  # NvLink read bandwidth for link 15 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = (
+    93  # NvLink write bandwidth for link 15 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = (
+    94  # NvLink read bandwidth for link 16 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = (
+    95  # NvLink write bandwidth for link 16 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = (
+    96  # NvLink read bandwidth for link 17 in MiB/sec
+)
+NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = (
+    97  # NvLink write bandwidth for link 17 in MiB/sec
+)
+NVML_GPM_METRIC_MAX = 98
 
 ## Structs
 
+
 class c_nvmlUnitInfo_t(_PrintableStructure):
     _fields_ = [
-        ('name', c_char * 96),
-        ('id', c_char * 96),
-        ('serial', c_char * 96),
-        ('firmwareVersion', c_char * 96),
+        ("name", c_char * 96),
+        ("id", c_char * 96),
+        ("serial", c_char * 96),
+        ("firmwareVersion", c_char * 96),
     ]
 
+
 class struct_c_nvmlGpmSample_t(Structure):
-    pass # opaque handle
+    pass  # opaque handle
+
+
 c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t)
 
+
 class c_metricInfo_t(Structure):
     _fields_ = [
         ("shortName", c_char_p),
@@ -5662,47 +6652,55 @@ class c_metricInfo_t(Structure):
         ("unit", c_char_p),
     ]
 
+
 class c_nvmlGpmMetric_t(_PrintableStructure):
     _fields_ = [
-        ('metricId', c_uint),
-        ('nvmlReturn', _nvmlReturn_t),
-        ('value', c_double),
-        ('metricInfo', c_metricInfo_t)
+        ("metricId", c_uint),
+        ("nvmlReturn", _nvmlReturn_t),
+        ("value", c_double),
+        ("metricInfo", c_metricInfo_t),
     ]
 
+
 class c_nvmlGpmMetricsGet_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('numMetrics', c_uint),
-        ('sample1', c_nvmlGpmSample_t),
-        ('sample2', c_nvmlGpmSample_t),
-        ('metrics', c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX)
+        ("version", c_uint),
+        ("numMetrics", c_uint),
+        ("sample1", c_nvmlGpmSample_t),
+        ("sample2", c_nvmlGpmSample_t),
+        ("metrics", c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX),
     ]
 
+
 NVML_GPM_METRICS_GET_VERSION = 1
 
+
 class c_nvmlGpmSupport_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('isSupportedDevice', c_uint),
+        ("version", c_uint),
+        ("isSupportedDevice", c_uint),
     ]
 
+
 NVML_GPM_SUPPORT_VERSION = 1
 
 ## Functions
 
+
 def nvmlGpmMetricsGet(metricsGet):
     fn = _nvmlGetFunctionPointer("nvmlGpmMetricsGet")
     ret = fn(byref(metricsGet))
     _nvmlCheckReturn(ret)
     return metricsGet
 
+
 def nvmlGpmSampleFree(gpmSample):
     fn = _nvmlGetFunctionPointer("nvmlGpmSampleFree")
     ret = fn(gpmSample)
     _nvmlCheckReturn(ret)
     return
 
+
 def nvmlGpmSampleAlloc():
     gpmSample = c_nvmlGpmSample_t()
     fn = _nvmlGetFunctionPointer("nvmlGpmSampleAlloc")
@@ -5710,18 +6708,21 @@ def nvmlGpmSampleAlloc():
     _nvmlCheckReturn(ret)
     return gpmSample
 
+
 def nvmlGpmSampleGet(device, gpmSample):
     fn = _nvmlGetFunctionPointer("nvmlGpmSampleGet")
     ret = fn(device, gpmSample)
     _nvmlCheckReturn(ret)
     return gpmSample
 
+
 def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample):
     fn = _nvmlGetFunctionPointer("nvmlGpmMigSampleGet")
     ret = fn(device, gpuInstanceId, gpmSample)
     _nvmlCheckReturn(ret)
     return gpmSample
 
+
 def nvmlGpmQueryDeviceSupport(device):
     gpmSupport = c_nvmlGpmSupport_t()
     gpmSupport.version = NVML_GPM_SUPPORT_VERSION
@@ -5730,6 +6731,7 @@ def nvmlGpmQueryDeviceSupport(device):
     _nvmlCheckReturn(ret)
     return gpmSupport
 
+
 def nvmlGpmSetStreamingEnabled(device, state):
     c_state = c_uint(state)
     fn = _nvmlGetFunctionPointer("nvmlGpmSetStreamingEnabled")
@@ -5737,6 +6739,7 @@ def nvmlGpmSetStreamingEnabled(device, state):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlGpmQueryIfStreamingEnabled(device):
     c_state = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlGpmQueryIfStreamingEnabled")
@@ -5744,21 +6747,24 @@ def nvmlGpmQueryIfStreamingEnabled(device):
     _nvmlCheckReturn(ret)
     return c_state.value
 
+
 # Low Power Structure and Function
 
-NVML_NVLINK_POWER_STATE_HIGH_SPEED    = 0x0
-NVML_NVLINK_POWER_STATE_LOW           = 0x1
+NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0
+NVML_NVLINK_POWER_STATE_LOW = 0x1
 
-NVML_NVLINK_LOW_POWER_THRESHOLD_MIN     = 0x1
-NVML_NVLINK_LOW_POWER_THRESHOLD_MAX     = 0x1FFF
-NVML_NVLINK_LOW_POWER_THRESHOLD_RESET   = 0xFFFFFFFF
+NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1
+NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF
+NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF
 NVML_NVLINK_LOW_POWER_THRESHOLD_DEFAULT = NVML_NVLINK_LOW_POWER_THRESHOLD_RESET
 
+
 class c_nvmlNvLinkPowerThres_t(Structure):
     _fields_ = [
         ("lowPwrThreshold", c_uint),
     ]
 
+
 def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold):
     c_info = c_nvmlNvLinkPowerThres_t()
     c_info.lowPwrThreshold = l1threshold
@@ -5767,48 +6773,52 @@ def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 NVML_GPU_FABRIC_UUID_LEN = 16
 
 _nvmlGpuFabricState_t = c_uint
 NVML_GPU_FABRIC_STATE_NOT_SUPPORTED = 0
-NVML_GPU_FABRIC_STATE_NOT_STARTED   = 1
-NVML_GPU_FABRIC_STATE_IN_PROGRESS   = 2
-NVML_GPU_FABRIC_STATE_COMPLETED     = 3
+NVML_GPU_FABRIC_STATE_NOT_STARTED = 1
+NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2
+NVML_GPU_FABRIC_STATE_COMPLETED = 3
+
 
 class c_nvmlGpuFabricInfo_t(_PrintableStructure):
     _fields_ = [
         ("clusterUuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE),
         ("status", _nvmlReturn_t),
         ("cliqueId", c_uint32),
-        ("state", _nvmlGpuFabricState_t)
+        ("state", _nvmlGpuFabricState_t),
     ]
 
+
 NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_NOT_SUPPORTED = 0
-NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE          = 1
-NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE         = 2
-NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW         = 0
-NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW         = 0x11
-
-NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED   = 0
-NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE            = 1
-NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE           = 2
-NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY           = 2
-NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY           = 0x11
-
-NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED  = 0
-NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE           = 1
-NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE          = 2
-NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY          = 4
-NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY          = 0x11
+NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1
+NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2
+NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0
+NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11
+
+NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0
+NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1
+NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2
+NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2
+NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11
+
+NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0
+NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1
+NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2
+NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4
+NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11
 
 NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_NOT_SUPPORTED = 0
-NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE          = 1
-NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE         = 2
-NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY         = 6
-NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY         = 0x11
+NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1
+NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2
+NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6
+NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11
 
 nvmlGpuFabricInfo_v2 = 0x02000024
 
+
 class c_nvmlGpuFabricInfoV_t(_PrintableStructure):
     _fields_ = [
         ("version", c_uint),
@@ -5816,33 +6826,37 @@ class c_nvmlGpuFabricInfoV_t(_PrintableStructure):
         ("status", _nvmlReturn_t),
         ("cliqueId", c_uint32),
         ("state", _nvmlGpuFabricState_t),
-        ("healthMask", c_uint32)
+        ("healthMask", c_uint32),
     ]
 
     def __init__(self):
         super(c_nvmlGpuFabricInfoV_t, self).__init__(version=nvmlGpuFabricInfo_v2)
 
+
 def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo):
-    fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo");
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo")
     ret = fn(device, gpuFabricInfo)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo):
-    fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV");
+    fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV")
     ret = fn(device, gpuFabricInfo)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 ######################
 ## Enums/defines
 #### NVML GPU NVLINK BW MODE
-NVML_GPU_NVLINK_BW_MODE_FULL      = 0x0
-NVML_GPU_NVLINK_BW_MODE_OFF       = 0x1
-NVML_GPU_NVLINK_BW_MODE_MIN       = 0x2
-NVML_GPU_NVLINK_BW_MODE_HALF      = 0x3
-NVML_GPU_NVLINK_BW_MODE_3QUARTER  = 0x4
-NVML_GPU_NVLINK_BW_MODE_COUNT     = 0x5
+NVML_GPU_NVLINK_BW_MODE_FULL = 0x0
+NVML_GPU_NVLINK_BW_MODE_OFF = 0x1
+NVML_GPU_NVLINK_BW_MODE_MIN = 0x2
+NVML_GPU_NVLINK_BW_MODE_HALF = 0x3
+NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4
+NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5
+
 
 def nvmlSystemSetNvlinkBwMode(mode):
     fn = _nvmlGetFunctionPointer("nvmlSystemSetNvlinkBwMode")
@@ -5850,6 +6864,7 @@ def nvmlSystemSetNvlinkBwMode(mode):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlSystemGetNvlinkBwMode():
     mode = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlSystemGetNvlinkBwMode")
@@ -5857,22 +6872,28 @@ def nvmlSystemGetNvlinkBwMode():
     _nvmlCheckReturn(ret)
     return mode.value
 
+
 _nvmlPowerScopeType_t = c_uint
-NVML_POWER_SCOPE_GPU     = 0
-NVML_POWER_SCOPE_MODULE  = 1
-NVML_POWER_SCOPE_MEMORY  = 2
+NVML_POWER_SCOPE_GPU = 0
+NVML_POWER_SCOPE_MODULE = 1
+NVML_POWER_SCOPE_MEMORY = 2
+
 
 class c_nvmlPowerValue_v2_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('powerScope', _nvmlPowerScopeType_t),
-        ('powerValueMw', c_uint),
+        ("version", c_uint),
+        ("powerScope", _nvmlPowerScopeType_t),
+        ("powerValueMw", c_uint),
     ]
-    _fmt_ = {'<default>': "%d B"}
+    _fmt_ = {"<default>": "%d B"}
+
 
 nvmlPowerValue_v2 = 0x0200000C
 
-def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version=nvmlPowerValue_v2):
+
+def nvmlDeviceSetPowerManagementLimit_v2(
+    device, powerScope, powerLimit, version=nvmlPowerValue_v2
+):
     c_powerScope = _nvmlPowerScopeType_t(powerScope)
     c_powerValue = c_nvmlPowerValue_v2_t()
     c_powerValue.version = c_uint(version)
@@ -5882,139 +6903,178 @@ def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version
     ret = fn(device, byref(c_powerValue))
     return NVML_SUCCESS
 
+
 class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('aggregateUncParity', c_ulonglong),
-        ('aggregateUncSecDed', c_ulonglong),
-        ('aggregateCor', c_ulonglong),
-        ('volatileUncParity', c_ulonglong),
-        ('volatileUncSecDed', c_ulonglong),
-        ('volatileCor', c_ulonglong),
-        ('aggregateUncBucketL2', c_ulonglong),
-        ('aggregateUncBucketSm', c_ulonglong),
-        ('aggregateUncBucketPcie', c_ulonglong),
-        ('aggregateUncBucketMcu', c_ulonglong),
-        ('aggregateUncBucketOther', c_ulonglong),
-        ('bThresholdExceeded', c_uint)
+        ("version", c_uint),
+        ("aggregateUncParity", c_ulonglong),
+        ("aggregateUncSecDed", c_ulonglong),
+        ("aggregateCor", c_ulonglong),
+        ("volatileUncParity", c_ulonglong),
+        ("volatileUncSecDed", c_ulonglong),
+        ("volatileCor", c_ulonglong),
+        ("aggregateUncBucketL2", c_ulonglong),
+        ("aggregateUncBucketSm", c_ulonglong),
+        ("aggregateUncBucketPcie", c_ulonglong),
+        ("aggregateUncBucketMcu", c_ulonglong),
+        ("aggregateUncBucketOther", c_ulonglong),
+        ("bThresholdExceeded", c_uint),
     ]
 
     def __init__(self):
-        super(c_nvmlEccSramErrorStatus_v1_t, self).__init__(version=nvmlEccSramErrorStatus_v1)
+        super(c_nvmlEccSramErrorStatus_v1_t, self).__init__(
+            version=nvmlEccSramErrorStatus_v1
+        )
+
 
 nvmlEccSramErrorStatus_v1 = 0x1000068
+
+
 def nvmlDeviceGetSramEccErrorStatus(device, status):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramEccErrorStatus")
     ret = fn(device, status)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
-NVML_DEV_CAP_EGM = (1 << 0)
+
+NVML_DEV_CAP_EGM = 1 << 0
 nvmlDeviceCapabilities_v1 = 0x1000008
 
+
 class c_nvmlDeviceCapabilities_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('capMask', c_uint),
+        ("version", c_uint),
+        ("capMask", c_uint),
     ]
 
     def __init__(self):
-        super(c_nvmlDeviceCapabilities_v1_t, self).__init__(version=nvmlDeviceCapabilities_v1)
+        super(c_nvmlDeviceCapabilities_v1_t, self).__init__(
+            version=nvmlDeviceCapabilities_v1
+        )
 
 
 def nvmlDeviceGetCapabilities(device, caps):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetCapabilities")
     return fn(device, caps)
 
+
 class c_nvmlPlatformInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('ibGuid', c_char * 16),
-        ('rackGuid', c_char * 16),
-        ('chassisPhysicalSlotNumber', c_char),
-        ('computeSlotIndex', c_char),
-        ('nodeIndex', c_char),
-        ('peerType', c_char),
-        ('moduleId', c_char)
+        ("version", c_uint),
+        ("ibGuid", c_char * 16),
+        ("rackGuid", c_char * 16),
+        ("chassisPhysicalSlotNumber", c_char),
+        ("computeSlotIndex", c_char),
+        ("nodeIndex", c_char),
+        ("peerType", c_char),
+        ("moduleId", c_char),
     ]
 
     def __init__(self):
         super(c_nvmlPlatformInfo_v1_t, self).__init__(version=nvmlPlatformInfo_v1)
 
-nvmlPlatformInfo_v1 = 0x100002c
+
+nvmlPlatformInfo_v1 = 0x100002C
+
+
 def nvmlDeviceGetPlatformInfo(device, platformInfo):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPlatformInfo")
     ret = fn(device, platformInfo)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 class c_nvmlMask255_t(_PrintableStructure):
     _fields_ = [
-        ('mask', c_uint * 8),
+        ("mask", c_uint * 8),
     ]
 
-NVML_WORKLOAD_POWER_MAX_PROFILES    = 255
-NVML_POWER_PROFILE_MAX_P            = 0
-NVML_POWER_PROFILE_MAX_Q            = 1
-NVML_POWER_PROFILE_COMPUTE          = 2
-NVML_POWER_PROFILE_MEMORY_BOUND     = 3
-NVML_POWER_PROFILE_NETWORK          = 4
-NVML_POWER_PROFILE_BALANCED         = 5
-NVML_POWER_PROFILE_LLM_INFERENCE    = 6
-NVML_POWER_PROFILE_LLM_TRAINING     = 7
-NVML_POWER_PROFILE_RBM              = 8
-NVML_POWER_PROFILE_DCPCIE           = 9
-NVML_POWER_PROFILE_HMMA_SPARSE      = 10
-NVML_POWER_PROFILE_HMMA_DENSE       = 11
-NVML_POWER_PROFILE_SYNC_BALANCED    = 12
-NVML_POWER_PROFILE_HPC              = 13
-NVML_POWER_PROFILE_MIG              = 14
-NVML_POWER_PROFILE_MAX              = 15
-
-nvmlWorkloadPowerProfileInfo_v1 = 0x100002c
+
+NVML_WORKLOAD_POWER_MAX_PROFILES = 255
+NVML_POWER_PROFILE_MAX_P = 0
+NVML_POWER_PROFILE_MAX_Q = 1
+NVML_POWER_PROFILE_COMPUTE = 2
+NVML_POWER_PROFILE_MEMORY_BOUND = 3
+NVML_POWER_PROFILE_NETWORK = 4
+NVML_POWER_PROFILE_BALANCED = 5
+NVML_POWER_PROFILE_LLM_INFERENCE = 6
+NVML_POWER_PROFILE_LLM_TRAINING = 7
+NVML_POWER_PROFILE_RBM = 8
+NVML_POWER_PROFILE_DCPCIE = 9
+NVML_POWER_PROFILE_HMMA_SPARSE = 10
+NVML_POWER_PROFILE_HMMA_DENSE = 11
+NVML_POWER_PROFILE_SYNC_BALANCED = 12
+NVML_POWER_PROFILE_HPC = 13
+NVML_POWER_PROFILE_MIG = 14
+NVML_POWER_PROFILE_MAX = 15
+
+nvmlWorkloadPowerProfileInfo_v1 = 0x100002C
+
+
 class c_nvmlWorkloadPowerProfileInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('profileId', c_uint),
-        ('priority', c_uint),
-        ('conflictingmask', c_nvmlMask255_t)
+        ("version", c_uint),
+        ("profileId", c_uint),
+        ("priority", c_uint),
+        ("conflictingmask", c_nvmlMask255_t),
     ]
 
     def __init__(self):
-        super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileInfo_v1)
+        super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__(
+            version=nvmlWorkloadPowerProfileInfo_v1
+        )
+
+
+nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002BF8
+
 
-nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002bf8
 class c_nvmlWorkloadPowerProfileProfilesInfo_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('perfProfilesMask', c_nvmlMask255_t),
-        ('perfProfile', c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES)
+        ("version", c_uint),
+        ("perfProfilesMask", c_nvmlMask255_t),
+        (
+            "perfProfile",
+            c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES,
+        ),
     ]
 
     def __init__(self):
-        super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileProfilesInfo_v1)
+        super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__(
+            version=nvmlWorkloadPowerProfileProfilesInfo_v1
+        )
+
 
 nvmlWorkloadPowerProfileCurrentProfiles_v1 = 0x1000064
+
+
 class c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('perfProfilesMask', c_nvmlMask255_t),
-        ('requestedProfilesMask', c_nvmlMask255_t),
-        ('enforcedProfilesMask', c_nvmlMask255_t)
+        ("version", c_uint),
+        ("perfProfilesMask", c_nvmlMask255_t),
+        ("requestedProfilesMask", c_nvmlMask255_t),
+        ("enforcedProfilesMask", c_nvmlMask255_t),
     ]
 
     def __init__(self):
-        super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileCurrentProfiles_v1)
+        super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__(
+            version=nvmlWorkloadPowerProfileCurrentProfiles_v1
+        )
+
 
 nvmlWorkloadPowerProfileRequestedProfiles_v1 = 0x1000024
+
+
 class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('requestedProfilesMask', c_nvmlMask255_t),
+        ("version", c_uint),
+        ("requestedProfilesMask", c_nvmlMask255_t),
     ]
 
     def __init__(self):
-        super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileRequestedProfiles_v1)
+        super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__(
+            version=nvmlWorkloadPowerProfileRequestedProfiles_v1
+        )
+
 
 def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo):
     fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetProfilesInfo")
@@ -6022,52 +7082,63 @@ def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo):
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceWorkloadPowerProfileGetCurrentProfiles(device, currentProfiles):
     fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetCurrentProfiles")
     ret = fn(device, currentProfiles)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceWorkloadPowerProfileSetRequestedProfiles(device, requestedProfiles):
     fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileSetRequestedProfiles")
     ret = fn(device, requestedProfiles)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceWorkloadPowerProfileClearRequestedProfiles(device, requestedProfiles):
     fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileClearRequestedProfiles")
     ret = fn(device, requestedProfiles)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetNvlinkSupportedBwModes(device, supportedBwModes):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkSupportedBwModes")
     ret = fn(device, supportedBwModes)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceGetNvlinkBwMode(device, getBwMode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkBwMode")
     ret = fn(device, getBwMode)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 def nvmlDeviceSetNvlinkBwMode(device, setBwMode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvlinkBwMode")
     ret = fn(device, setBwMode)
     _nvmlCheckReturn(ret)
     return NVML_SUCCESS
 
+
 nvmlDramEncryptionInfo_v1 = 0x01000008
 
+
 class c_nvmlDramEncryptionInfo_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('encryptionState',  _nvmlEnableState_t),
+        ("version", c_uint),
+        ("encryptionState", _nvmlEnableState_t),
     ]
 
     def __init__(self):
-        super(c_nvmlDramEncryptionInfo_t, self).__init__(version=nvmlDramEncryptionInfo_v1)
+        super(c_nvmlDramEncryptionInfo_t, self).__init__(
+            version=nvmlDramEncryptionInfo_v1
+        )
+
 
 def nvmlDeviceGetDramEncryptionMode(handle):
     c_currState = c_nvmlDramEncryptionInfo_t()
@@ -6077,64 +7148,79 @@ def nvmlDeviceGetDramEncryptionMode(handle):
     _nvmlCheckReturn(ret)
     return [c_currState.encryptionState, c_pendingState.encryptionState]
 
+
 # added to API
 def nvmlDeviceGetCurrentDramEncryptionMode(handle):
     return nvmlDeviceGetDramEncryptionMode(handle)[0]
 
+
 # added to API
 def nvmlDeviceGetPendingDramEncryptionMode(handle):
     return nvmlDeviceGetDramEncryptionMode(handle)[1]
 
+
 def nvmlDeviceSetDramEncryptionMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetDramEncryptionMode")
     c_dramEncryptionMode = c_nvmlDramEncryptionInfo_t()
-    c_dramEncryptionMode.encryptionState = mode;
+    c_dramEncryptionMode.encryptionState = mode
     ret = fn(handle, byref(c_dramEncryptionMode))
     _nvmlCheckReturn(ret)
     return None
 
+
 # Power Smoothing defines
-NVML_POWER_SMOOTHING_MAX_NUM_PROFILES                   = 5
-NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET             = 0xFFFFFFFF
-NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR    = 0
-NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE         = 1
-NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE       = 2
+NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5
+NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF
+NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0
+NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1
+NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2
 NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_HYSTERESIS = 3
 
-nvmlPowerSmoothingState_v1=0x1000008
+nvmlPowerSmoothingState_v1 = 0x1000008
+
+
 class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('state', c_uint),
+        ("version", c_uint),
+        ("state", c_uint),
     ]
 
     def __init__(self):
-        super(c_nvmlPowerSmoothingState_v1_t, self).__init__(version=nvmlPowerSmoothingState_v1)
+        super(c_nvmlPowerSmoothingState_v1_t, self).__init__(
+            version=nvmlPowerSmoothingState_v1
+        )
+
+
+nvmlPowerSmoothingProfile_v1 = 0x1000018
+
 
-nvmlPowerSmoothingProfile_v1=0x1000018
 class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure):
     _fields_ = [
-        ('version', c_uint),
-        ('profileId', c_uint),
-        ('paramId', c_uint),
-        ('value', c_double),
+        ("version", c_uint),
+        ("profileId", c_uint),
+        ("paramId", c_uint),
+        ("value", c_double),
     ]
 
     def __init__(self):
-        super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__(version=nvmlPowerSmoothingProfile_v1)
+        super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__(
+            version=nvmlPowerSmoothingProfile_v1
+        )
+
 
 def nvmlDevicePowerSmoothingActivatePresetProfile(device, profile):
     fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingActivatePresetProfile")
     ret = fn(device, profile)
     _nvmlCheckReturn(ret)
 
+
 def nvmlDevicePowerSmoothingUpdatePresetProfileParam(device, profile):
     fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingUpdatePresetProfileParam")
     ret = fn(device, profile)
     _nvmlCheckReturn(ret)
 
+
 def nvmlDevicePowerSmoothingSetState(device, state):
     fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingSetState")
     ret = fn(device, state)
     _nvmlCheckReturn(ret)
-
diff --git a/vllm/tilelang_utils/__init__.py b/vllm/tilelang_utils/__init__.py
index 803409b1ec65..5ac03c858ed0 100644
--- a/vllm/tilelang_utils/__init__.py
+++ b/vllm/tilelang_utils/__init__.py
@@ -31,6 +31,7 @@ def _ensure_tilelang_imported() -> None:
 
     Raises:
         ImportError: If TileLang is not installed.
+
     """
     global T, tilelang
 
@@ -60,8 +61,41 @@ def _get_pass_configs() -> dict[Any, Any]:
     return pass_configs
 
 
+class _DeferredTileLangJitKernel:
+    """Stand-in for a `tilelang.jit` kernel that decorates on first use.
+
+    Both attribute access and calling apply the decoration. Required for
+    compile-only JIT warmup on platforms that defer import of tilelang.
+    """
+
+    _kernel_function: Callable[..., Any] | None = None
+    _jit_kernel: Any = None
+
+    def __init__(self, kernel_function: Callable[..., Any]) -> None:
+        self._kernel_function = kernel_function
+        functools.update_wrapper(self, kernel_function)
+
+    def _ensure_jit_kernel(self) -> Any:
+        if self._jit_kernel is None:
+            _ensure_tilelang_imported()
+            kernel_function = self._kernel_function
+            assert kernel_function is not None
+            kernel_function.__globals__["tilelang"] = tilelang
+            kernel_function.__globals__["T"] = T
+            self._jit_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
+                kernel_function
+            )
+        return self._jit_kernel
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._ensure_jit_kernel()(*args, **kwargs)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._ensure_jit_kernel(), name)
+
+
 def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
-    """Apply `tilelang.jit`, deferring until first call on ROCm.
+    """Apply `tilelang.jit`, deferring until first use on ROCm.
 
     ROCm defers JIT decoration so importing the caller's module does not
     require TileLang immediately. CUDA keeps the eager decoration behavior.
@@ -74,18 +108,4 @@ def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
         _ensure_tilelang_imported()
         return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function)
 
-    compiled_kernel: Callable[..., Any] | None = None
-
-    @functools.wraps(kernel_function)
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
-        nonlocal compiled_kernel
-        if compiled_kernel is None:
-            _ensure_tilelang_imported()
-            kernel_function.__globals__["tilelang"] = tilelang
-            kernel_function.__globals__["T"] = T
-            compiled_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
-                kernel_function
-            )
-        return compiled_kernel(*args, **kwargs)
-
-    return wrapper
+    return _DeferredTileLangJitKernel(kernel_function)
diff --git a/vllm/tokenizers/deepseek_v32.py b/vllm/tokenizers/deepseek_v32.py
index 0e9aaa2cae5a..3f0fff938ad9 100644
--- a/vllm/tokenizers/deepseek_v32.py
+++ b/vllm/tokenizers/deepseek_v32.py
@@ -13,9 +13,7 @@
 
 
 def get_deepseek_v32_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
-    """
-    Wraps a tokenizer to use the custom DeepSeek V3.2 chat template encoding.
-    """
+    """Wraps a tokenizer to use the custom DeepSeek V3.2 chat template encoding."""
     dsv32_tokenizer = copy.copy(tokenizer)
 
     added_vocab = tokenizer.get_added_vocab()
diff --git a/vllm/tokenizers/deepseek_v4.py b/vllm/tokenizers/deepseek_v4.py
index 2a0db215b80f..532ca0df7c9b 100644
--- a/vllm/tokenizers/deepseek_v4.py
+++ b/vllm/tokenizers/deepseek_v4.py
@@ -13,9 +13,7 @@
 
 
 def get_deepseek_v4_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
-    """
-    Wraps a tokenizer to use the custom DeepSeek V4 chat template encoding.
-    """
+    """Wraps a tokenizer to use the custom DeepSeek V4 chat template encoding."""
     dsv4_tokenizer = copy.copy(tokenizer)
 
     added_vocab = tokenizer.get_added_vocab()
@@ -37,8 +35,18 @@ def apply_chat_template(
             conversation = kwargs.get("conversation", messages)
             messages = conversation.copy()
             if tools is not None and len(tools) > 0:
-                messages.insert(0, {"role": "system"})
-                messages[0]["tools"] = tools  # type: ignore[typeddict-unknown-key]
+                # Match the Rust renderer: request tools attach to the first
+                # system message; synthesize one only when none exists.
+                system_idx = next(
+                    (i for i, m in enumerate(messages) if m.get("role") == "system"),
+                    None,
+                )
+                if system_idx is None:
+                    messages.insert(0, {"role": "system"})
+                    system_idx = 0
+                else:
+                    messages[system_idx] = copy.copy(messages[system_idx])
+                messages[system_idx]["tools"] = tools  # type: ignore[typeddict-unknown-key]
 
             reasoning_effort = kwargs.get("reasoning_effort")
             if not isinstance(reasoning_effort, str):
diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py
index 3a577172708e..a1449e59e203 100644
--- a/vllm/tokenizers/deepseek_v41.py
+++ b/vllm/tokenizers/deepseek_v41.py
@@ -32,7 +32,7 @@ def _normalize_messages(
             parts = []
             for block in content:
                 part_type = block.get("type")
-                if part_type == "text":
+                if part_type in ("text", "input_text", "output_text"):
                     parts.append(block.get("text", ""))
                 elif part_type in ("image_url", "input_image", "image_pil"):
                     parts.append(IMAGE_PLACEHOLDER)
diff --git a/vllm/tokenizers/deepseek_v4_encoding.py b/vllm/tokenizers/deepseek_v4_encoding.py
index fb5bc44e854b..37dc55936688 100644
--- a/vllm/tokenizers/deepseek_v4_encoding.py
+++ b/vllm/tokenizers/deepseek_v4_encoding.py
@@ -224,7 +224,7 @@ def flatten_content_blocks(content: Any) -> Any:
             raise ValueError(
                 f"Unsupported content block type: {block.get('type')!r}"
             )
-    return "".join(parts)
+    return "\n\n".join(parts)
 
 
 def render_message(
@@ -267,7 +267,7 @@ def render_message(
     )
 
     role = msg.get("role")
-    content = msg.get("content")
+    content = flatten_content_blocks(msg.get("content"))
     tools = msg.get("tools")
     response_format = msg.get("response_format")
     tool_calls = msg.get("tool_calls")
@@ -337,7 +337,7 @@ def render_message(
                     parts.append(f"[Unsupported {block_type}]")
             prompt += "\n\n".join(parts)
         else:
-            prompt += flatten_content_blocks(content) or ""
+            prompt += content or ""
 
     elif role == "latest_reminder":
         prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
diff --git a/vllm/tokenizers/detokenizer_utils.py b/vllm/tokenizers/detokenizer_utils.py
index 516f1b70e732..35e73bc4100d 100644
--- a/vllm/tokenizers/detokenizer_utils.py
+++ b/vllm/tokenizers/detokenizer_utils.py
@@ -205,6 +205,7 @@ def detokenize_incrementally(
         skip_special_tokens: Whether to skip special tokens.
         spaces_between_special_tokens: Whether to add spaces between special
             tokens.
+
     """
     new_token_id = all_input_ids[-1]
     # This is the first iteration for this sequence
diff --git a/vllm/tokenizers/hf.py b/vllm/tokenizers/hf.py
index bdc767acd66c..a091d8e10547 100644
--- a/vllm/tokenizers/hf.py
+++ b/vllm/tokenizers/hf.py
@@ -23,8 +23,7 @@ class ThreadSafeHFTokenizerMixin:
 
 
 def maybe_make_thread_pool(tokenizer: _T, copies: int = 1):
-    """
-    If `tokenizer` is a `TokenizersBackend`, modify the tokenizer
+    """If `tokenizer` is a `TokenizersBackend`, modify the tokenizer
     in-place to make the public interface thread-safe by routing calls
     through a deep-copied tokenizer pool.
 
@@ -105,8 +104,7 @@ def __reduce__(self):
 
 
 def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
-    """
-    By default, transformers will recompute multiple tokenizer properties
+    """By default, transformers will recompute multiple tokenizer properties
     each time they are called, leading to a significant slowdown.
     This proxy caches these properties for faster access.
     """
@@ -116,17 +114,10 @@ def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
     tokenizer_all_special_tokens = tokenizer.all_special_tokens
     tokenizer_vocab = tokenizer.get_vocab()
     tokenizer_len = len(tokenizer)
-    # The underlying tokenizer class could be MistralCommonBackend,
-    # which does not implement is_fast in Transformers
+    # The underlying tokenizer class could be a specific backend,
+    # which does not always implement is_fast in Transformers
     tokenizer_is_fast = getattr(tokenizer, "is_fast", True)
 
-    # MistralCommonBackend is tekken-backed and needs byte-fallback-aware tokenization.
-    mistral_tekkenizer = None
-    if getattr(getattr(tokenizer, "tokenizer", None), "instruct_tokenizer", None):
-        from vllm.tokenizers.mistral import mistral_common_tekkenizer
-
-        mistral_tekkenizer = mistral_common_tekkenizer(tokenizer)
-
     max_token_id = max(tokenizer_vocab.values())
     max_chars_per_token = max(len(tok) for tok in tokenizer_vocab)
 
@@ -159,27 +150,6 @@ def max_chars_per_token(self) -> int:
         def is_fast(self) -> bool:
             return tokenizer_is_fast
 
-        def convert_ids_to_tokens(self, ids, skip_special_tokens: bool = False):
-            if mistral_tekkenizer is not None:
-                from vllm.tokenizers.mistral import tekken_convert_ids_to_tokens
-
-                return tekken_convert_ids_to_tokens(mistral_tekkenizer, ids)
-            return super().convert_ids_to_tokens(
-                ids, skip_special_tokens=skip_special_tokens
-            )
-
-        def convert_tokens_to_string(self, tokens: list[str]) -> str:
-            if mistral_tekkenizer is not None:
-                from vllm.tokenizers.mistral import tekken_convert_tokens_to_string
-
-                return tekken_convert_tokens_to_string(mistral_tekkenizer, tokens)
-            try:
-                return super().convert_tokens_to_string(tokens)
-            except NotImplementedError:
-                # The underlying tokenizer class could be MistralCommonBackend,
-                # which does not implement convert_tokens_to_string in Transformers
-                return "".join(tokens)
-
         def get_vocab(self) -> dict[str, int]:
             return tokenizer_vocab
 
diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py
index 66b69089f0e6..4ffefe0db6d9 100644
--- a/vllm/tokenizers/mistral.py
+++ b/vllm/tokenizers/mistral.py
@@ -354,22 +354,13 @@ def __call__(
                 "`text_pair` is not supported by `MistralTokenizer.__call__`."
             )
 
-        encoded = self.transformers_tokenizer(
+        return self.transformers_tokenizer(
             text=text,
             text_pair=text_pair,
             add_special_tokens=add_special_tokens,
             truncation=truncation,
             max_length=max_length,
         )
-        # TODO(juliendenize): once https://github.com/huggingface/transformers/pull/41962
-        # is in, revert to only call self.transformers_tokenizer(...).
-        # Hack to fix wrongly added eos token, when fix will be supported the condition
-        # below will be False even before the revert is done.
-        if encoded["input_ids"] and encoded["input_ids"][-1] == self.eos_token_id:
-            encoded["input_ids"].pop(-1)
-            if attention_mask := encoded.get("attention_mask"):
-                attention_mask.pop(-1)
-        return encoded
 
     @property
     def vocab(self) -> list[str]:
@@ -389,8 +380,10 @@ def encode(
         max_length: int | None = None,
         add_special_tokens: bool = True,
     ) -> list[int]:
-        # TODO(juliendenize): once https://github.com/huggingface/transformers/pull/41962
-        # is in, directly call self.transformers_tokenizer.encode(...).
+        # NOTE: transformers guards truncation with `if max_length and ...`, so
+        # `max_length=0` returns the full sequence instead of an empty one. vLLM
+        # treats `truncate_prompt_tokens=0` as an empty prompt, so keep slicing
+        # here until that is fixed upstream.
         encoded = self.tokenizer.encode(text, bos=add_special_tokens, eos=False)
 
         if truncation is not False and max_length is not None:
@@ -437,11 +430,6 @@ def apply_chat_template(
     def decode(
         self, ids: Sequence[int] | int, skip_special_tokens: bool = False
     ) -> str:
-        # TODO(juliendenize): once https://github.com/huggingface/transformers/pull/41962
-        # is in, directly call self.transformers_tokenizer.decode(...).
-        if isinstance(ids, int):
-            ids = [ids]
-
         return self.transformers_tokenizer.decode(
             ids, skip_special_tokens=skip_special_tokens
         )
diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py
index b9c9226675e1..09a715c117de 100644
--- a/vllm/tokenizers/registry.py
+++ b/vllm/tokenizers/registry.py
@@ -7,9 +7,15 @@
 from typing import TYPE_CHECKING
 
 import huggingface_hub
+import transformers
+from packaging.version import Version
+from transformers.models.auto.tokenization_auto import (
+    TOKENIZER_MAPPING_NAMES as _HF_TOKENIZER_MAPPING_NAMES,
+)
+from transformers.utils.import_utils import is_mistral_common_available
 from typing_extensions import TypeVar, assert_never
 
-import vllm.envs as envs
+from vllm import envs
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import _maybe_register_hf_config, get_config
 from vllm.transformers_utils.repo_utils import (
@@ -26,6 +32,8 @@
 
 logger = init_logger(__name__)
 
+_MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL = Version("5.15.0")
+
 
 # Model types whose hub tokenizer_class is incorrect and should be overridden with
 # TokenizersBackend (the generic fast tokenizer). Adding a model type here is always a
@@ -74,8 +82,6 @@ def register(self, tokenizer_mode: str, module: str, class_name: str) -> None:
 
         self.tokenizers[tokenizer_mode] = (module, class_name)
 
-        return None
-
     def load_tokenizer_cls(self, tokenizer_mode: str) -> type[TokenizerLike]:
         if tokenizer_mode not in self.tokenizers:
             raise ValueError(f"No tokenizer registered for {tokenizer_mode=!r}.")
@@ -139,23 +145,19 @@ def resolve_tokenizer_args(
         else:
             assert_never(runner_type)
 
-    if tokenizer_mode == "slow":
-        if kwargs.get("use_fast", False):
-            raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
-
-        tokenizer_mode = "hf"
-        kwargs["use_fast"] = False
-
     # Try to use official Mistral tokenizer if possible
     if (
         tokenizer_mode == "auto"
         and is_mistral_model_repo(
-            model_name_or_path=str(tokenizer_name), revision=revision
+            model_name_or_path=str(tokenizer_name),
+            revision=revision,
+            token=kwargs.get("token"),
         )
         and any_pattern_in_repo_files(
             model_name_or_path=str(tokenizer_name),
             allow_patterns=["tekken.json", "tokenizer.model.v*"],
             revision=revision,
+            token=kwargs.get("token"),
         )
     ):
         tokenizer_mode = "mistral"
@@ -164,6 +166,13 @@ def resolve_tokenizer_args(
     if tokenizer_mode == "auto":
         tokenizer_mode = "hf"
 
+    if tokenizer_mode == "hf":
+        if kwargs.pop("mistral_format", False):
+            raise ValueError(
+                "mistral_format=True is not supported with tokenizer_mode='hf'"
+            )
+        kwargs["mistral_format"] = False
+
     return tokenizer_mode, tokenizer_name, args, kwargs
 
 
@@ -231,6 +240,27 @@ def get_tokenizer(
             trust_remote_code=trust_remote_code,
             revision=revision,
             config_format=config_format,
+            token=kwargs.get("token"),
+        )
+
+    # TODO: delete when Transformers version dependency is bumped >= 5.15.0
+    if (
+        tokenizer_mode == "hf"
+        and Version(transformers.__version__) < _MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL
+        and is_mistral_common_available()
+        and _HF_TOKENIZER_MAPPING_NAMES.get(getattr(config, "model_type", None))
+        == "MistralCommonBackend"
+        and any_pattern_in_repo_files(
+            model_name_or_path=str(tokenizer_name),
+            allow_patterns=["tekken.json", "tokenizer.model.v*"],
+            revision=revision,
+            token=kwargs.get("token"),
+        )
+    ):
+        raise ValueError(
+            "Loading Mistral models with tokenizer_mode='hf' requires "
+            f"transformers>={_MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL}. "
+            "Please upgrade transformers or delete Mistral tokenizer files."
         )
 
     # Some models have an incorrect tokenizer_class on the hub.
diff --git a/vllm/tool_parsers/AGENTS.md b/vllm/tool_parsers/AGENTS.md
new file mode 100644
index 000000000000..2bd487997c57
--- /dev/null
+++ b/vllm/tool_parsers/AGENTS.md
@@ -0,0 +1,45 @@
+# vllm/tool_parsers
+
+Every `--tool-call-parser` name resolves here. `ParserManager.get_parser`
+composes the tool parser chosen by name with the reasoning parser into the
+unified `Parser` the serving layer uses. The `ToolParser` interface itself is
+legacy: it is kept so existing in-tree parsers and out-of-tree plugins keep
+working, not as the base for new ones.
+
+## Layout
+
+- `abstract_tool_parser.py` — `ToolParser` (`extract_tool_calls`,
+  `extract_tool_calls_streaming`, `adjust_request`) and `ToolParserManager`.
+- `__init__.py` — `_TOOL_PARSERS_TO_REGISTER`, the name → (file, class) table
+  that lazily registers every built-in parser.
+- `structural_tag_registry.py` — structural-tag builders used to constrain
+  tool-call output for `tool_choice="required"` / named tools.
+- `<model>_tool_parser.py` — one file per family. Either a subclass of a
+  `ParserToolAdapter` generated from a `vllm/parser/<model>.py` engine grammar
+  (`mistral_tool_parser.py`, `kimi_k2_tool_parser.py`; the `_engine_` suffix
+  is not applied consistently), or a self-contained `ToolParser` subclass.
+
+## Guidelines
+
+- Do not change the `ToolParser` interface in `abstract_tool_parser.py`.
+  Out-of-tree plugins loaded through `ToolParserManager.import_tool_parser`
+  subclass it. New capabilities go on `Parser` in `vllm/parser/`; if a legacy
+  parser needs the same capability, `DelegatingParser` supplies the default.
+- For a new model, implement a unified `Parser` in `vllm/parser/` first, then
+  expose it here as a thin adapter subclass plus an entry in
+  `_TOOL_PARSERS_TO_REGISTER`.
+- The `_TOOL_PARSERS_TO_REGISTER` key is what users pass on the command line;
+  treat renames as breaking.
+- Keep adapters thin: registration, `adjust_request` tweaks, and class
+  attributes. Parsing logic belongs in `vllm/parser/`.
+- Keep a model-specific fix inside that model's `<model>_tool_parser.py`, or
+  its grammar in `vllm/parser/<model>.py`, where possible. Touch
+  `abstract_tool_parser.py` or `structural_tag_registry.py` only for behavior
+  every parser needs, and say so in the PR.
+- Set `structural_tag_model` only to a key in `SUPPORTED_STRUCTURAL_TAG_MODELS`
+  (an xgrammar builtin, or one registered with `register_vllm_structural_tag`).
+  It only takes effect under `VLLM_ENFORCE_STRICT_TOOL_CALLING`, where it also
+  disables the generic required/named handling for that parser.
+- Comment only where the code alone would leave a reader with a question —
+  a request tweak in `adjust_request`, a marker handled unusually, a model
+  quirk being worked around. Keep it minimal and concise.
diff --git a/vllm/tool_parsers/CLAUDE.md b/vllm/tool_parsers/CLAUDE.md
new file mode 100644
index 000000000000..43c994c2d361
--- /dev/null
+++ b/vllm/tool_parsers/CLAUDE.md
@@ -0,0 +1 @@
+@AGENTS.md
diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py
index c43a34772fc9..2515009d0e81 100644
--- a/vllm/tool_parsers/__init__.py
+++ b/vllm/tool_parsers/__init__.py
@@ -75,8 +75,8 @@
         "Granite20bFCToolParser",
     ),
     "granite": (
-        "granite_tool_parser",
-        "GraniteToolParser",
+        "granite_engine_tool_parser",
+        "GraniteEngineToolParser",
     ),
     "granite4": (
         "granite4_tool_parser",
diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py
index 3ee35365e4a8..953782daa96f 100644
--- a/vllm/tool_parsers/abstract_tool_parser.py
+++ b/vllm/tool_parsers/abstract_tool_parser.py
@@ -30,6 +30,7 @@
     StructuredOutputsParams,
 )
 from vllm.tokenizers import TokenizerLike
+from vllm.tool_parsers.tool_strict_level import ToolStrictLevel
 from vllm.tool_parsers.utils import Tool, get_json_schema_from_tools
 from vllm.utils.collection_utils import is_list_of
 from vllm.utils.import_utils import import_plugin
@@ -40,8 +41,7 @@
 
 
 class ToolParser:
-    """
-    Abstract ToolParser class that should not be used directly. Provided
+    """Abstract ToolParser class that should not be used directly. Provided
     properties and methods should be used in
     derived classes.
     """
@@ -169,6 +169,7 @@ def get_structural_tag(
         request: ChatCompletionRequest | ResponsesRequest,
         *,
         reasoning: bool = False,
+        strict_level: ToolStrictLevel = ToolStrictLevel.AUTO,
     ):
         if self.structural_tag_model is None:
             return None
@@ -181,13 +182,13 @@ def get_structural_tag(
             tools=request.tools,
             tool_choice=request.tool_choice,
             reasoning=reasoning,
+            strict_level=strict_level,
         )
 
     def extract_tool_calls(
         self, model_output: str, request: ChatCompletionRequest
     ) -> ExtractedToolCallInformation:
-        """
-        Static method that should be implemented for extracting tool calls from
+        """Static method that should be implemented for extracting tool calls from
         a complete model-generated string.
         Used for non-streaming responses where we have the entire model response
         available before sending to the client.
@@ -207,8 +208,7 @@ def extract_tool_calls_streaming(
         delta_token_ids: Sequence[int],
         request: ChatCompletionRequest,
     ) -> DeltaMessage | None:
-        """
-        Instance method that should be implemented for extracting tool calls
+        """Instance method that should be implemented for extracting tool calls
         from an incomplete response; for use when handling tool calls and
         streaming. Has to be an instance method because  it requires state -
         the current tokens/diffs, but also the information about what has
@@ -220,8 +220,7 @@ def extract_tool_calls_streaming(
 
 
 class ToolParserManager:
-    """
-    Central registry for ToolParser implementations.
+    """Central registry for ToolParser implementations.
 
     Supports two modes:
       - Eager (immediate) registration via `register_module`
@@ -233,8 +232,7 @@ class ToolParserManager:
 
     @classmethod
     def get_tool_parser(cls, name: str) -> type[ToolParser]:
-        """
-        Retrieve a registered or lazily registered ToolParser class.
+        """Retrieve a registered or lazily registered ToolParser class.
 
         If the parser is lazily registered,
         it will be imported and cached on first access.
@@ -301,8 +299,7 @@ def _register_module(
 
     @classmethod
     def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None:
-        """
-        Register a lazy module mapping.
+        """Register a lazy module mapping.
 
         Example:
             ToolParserManager.register_lazy_module(
@@ -310,6 +307,7 @@ def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> N
                 module_path="vllm.tool_parsers.kimi_k2_parser",
                 class_name="KimiK2ToolParser",
             )
+
         """
         cls.lazy_parsers[name] = (module_path, class_name)
 
@@ -320,8 +318,7 @@ def register_module(
         force: bool = True,
         module: type[ToolParser] | None = None,
     ) -> type[ToolParser] | Callable[[type[ToolParser]], type[ToolParser]]:
-        """
-        Register module immediately or lazily (as a decorator).
+        """Register module immediately or lazily (as a decorator).
 
         Usage:
             @ToolParserManager.register_module("kimi_k2")
diff --git a/vllm/tool_parsers/apertus_tool_parser.py b/vllm/tool_parsers/apertus_tool_parser.py
index f23b8601aa61..6825b3ef39e1 100644
--- a/vllm/tool_parsers/apertus_tool_parser.py
+++ b/vllm/tool_parsers/apertus_tool_parser.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Tool call parser for Apertus models.
+"""Tool call parser for Apertus models.
 
 Extracts tool calls from the format:
 <|tools_prefix|>[{"function_name": {"arg1": "value1", ...}}, ...]<|tools_suffix|>
@@ -45,8 +44,7 @@
 
 
 class ApertusToolParser(ToolParser):
-    """
-    Tool call parser for Apertus models.
+    """Tool call parser for Apertus models.
 
     Handles the extraction of tool calls from text in both non-streaming
     (complete string) and streaming (chunked token) environments.
@@ -68,11 +66,11 @@ class ApertusToolParser(ToolParser):
         "get_weather"
         >>> info.tool_calls[0].function.arguments
         '{"city": "Paris"}'
+
     """
 
     def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
-        """
-        Initializes the ApertusToolParser.
+        """Initializes the ApertusToolParser.
 
         Args:
             tokenizer: The model's tokenizer.
@@ -82,6 +80,7 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         Raises:
             ValueError: If the `model_tokenizer`
                 is not successfully passed to the base class.
+
         """
         super().__init__(tokenizer, tools)
 
@@ -101,8 +100,7 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         self._reset_streaming_state()
 
     def _reset_streaming_state(self) -> None:
-        """
-        Resets all streaming state variables for a new completion request.
+        """Resets all streaming state variables for a new completion request.
 
         This clears the delta text buffer and resets the pointers used to
         track the currently streaming tool index and arguments. Called implicitly
@@ -116,8 +114,7 @@ def _reset_streaming_state(self) -> None:
     def adjust_request(
         self, request: ChatCompletionRequest | ResponsesRequest
     ) -> ChatCompletionRequest | ResponsesRequest:
-        """
-        Adjusts the generation request to ensure special tool tokens are not skipped.
+        """Adjusts the generation request to ensure special tool tokens are not skipped.
 
         Forces `skip_special_tokens=False` if tools are actively being evaluated,
         ensuring the tools special tokens are surfaced to the engine for parsing.
@@ -127,6 +124,7 @@ def adjust_request(
 
         Returns:
             The potentially modified chat completion request.
+
         """
         request = super().adjust_request(request)
         if request.tools and request.tool_choice != "none":
@@ -134,8 +132,7 @@ def adjust_request(
         return request
 
     def _buffer_delta_text(self, delta_text: str) -> str:
-        """
-        Buffers incoming delta chunks to prevent
+        """Buffers incoming delta chunks to prevent
         fragmentation of multi-token special tags.
 
         If a chunk ends with a partial match of
@@ -155,6 +152,7 @@ def _buffer_delta_text(self, delta_text: str) -> str:
             "Let me check "  # "<|tool" is buffered internally
             >>> parser._buffer_delta_text("s_prefix|>" \
             "<|tools_prefix|>"  # Buffer released on completion
+
         """
         self.buffered_delta_text += delta_text
         text = self.buffered_delta_text
@@ -178,8 +176,7 @@ def extract_tool_calls(
         model_output: str,
         request: ChatCompletionRequest,
     ) -> ExtractedToolCallInformation:
-        """
-        Extracts tool calls from a completely generated model response (Non-Streaming).
+        """Extract tool calls from a complete model response (non-streaming).
 
         Args:
             model_output: The full completion string generated by the model.
@@ -200,6 +197,7 @@ def extract_tool_calls(
             'Let me see.'
             >>> info.tool_calls[0].function.name
             'get_weather'
+
         """
         match = self.tool_call_regex.search(model_output)
         if not match:
@@ -280,8 +278,7 @@ def extract_tool_calls_streaming(
         delta_token_ids: Sequence[int],
         request: ChatCompletionRequest,
     ) -> DeltaMessage | None:
-        """
-        Handles streaming chunks
+        """Handles streaming chunks.
 
         Args:
             previous_text: The complete model text generated prior to this chunk.
@@ -305,6 +302,7 @@ def extract_tool_calls_streaming(
             ... )
             >>> msg.tool_calls[0].function.arguments
             'ation": "Paris"}'
+
         """
         delta_text = self._buffer_delta_text(delta_text)
         if not delta_text:
@@ -323,8 +321,7 @@ def extract_tool_calls_streaming(
     def _extract_streaming(
         self, current_text: str, delta_text: str
     ) -> DeltaMessage | None:
-        """
-        Core streaming logic.
+        """Core streaming logic.
         Separates visible chat text from JSON blocks and computes diffs.
 
         Args:
@@ -333,6 +330,7 @@ def _extract_streaming(
 
         Returns:
             A `DeltaMessage` containing the `content` delta and/or `tool_calls` delta.
+
         """
         prefix_idx = current_text.rfind(TOOL_CALLS_PREFIX)
         suffix_idx = current_text.rfind(TOOL_CALLS_SUFFIX)
@@ -404,8 +402,7 @@ def _extract_streaming(
     def _parse_and_diff_json(
         self, json_str: str, is_final: bool
     ) -> list[DeltaToolCall]:
-        """
-        Parses an isolated, potentially incomplete streaming JSON array and returns
+        """Parses an isolated, potentially incomplete streaming JSON array and returns
         newly accumulated tool call diffs.
 
         Args:
@@ -417,6 +414,7 @@ def _parse_and_diff_json(
             A list of `DeltaToolCall`
             items representing string diffs in function arguments
             to stream back to the client.
+
         """
         try:
             parsed, _ = partial_json_loads(json_str, Allow.ALL)
@@ -460,8 +458,7 @@ def _parse_and_diff_json(
     def _emit_tool_name(
         self, parsed: list, index: int, tool_calls: list[DeltaToolCall]
     ) -> None:
-        """
-        Extracts and emits the function name mapped to a new tool call ID.
+        """Extracts and emits the function name mapped to a new tool call ID.
 
         Args:
             parsed: The partially parsed JSON list containing tool dictionaries.
@@ -472,6 +469,7 @@ def _emit_tool_name(
             Appends `DeltaToolCall(index=0,
                 function=DeltaFunctionCall(name="get_weather", ...))`
             to the `tool_calls` list and marks the name as sent.
+
         """
         obj = parsed[index]
         if isinstance(obj, dict) and obj:
@@ -491,8 +489,7 @@ def _emit_tool_name(
     def _get_tool_diff(
         self, parsed: list, index: int, is_final: bool
     ) -> DeltaToolCall | None:
-        """
-        Calculates the exact string difference to safely append new tool parameters.
+        """Calculates the exact string difference to safely append new tool parameters.
 
         This ensures characters like `{`, `}`, and `"` don't jump around unevenly
         in the UI frontend while streaming incomplete JSON arguments.
@@ -513,6 +510,7 @@ def _get_tool_diff(
             >>> # Returns diff (closing bracket suppressed until final):
             >>> parser._get_tool_diff(parsed, index=0, is_final=False)
             DeltaToolCall(index=0, function=DeltaFunctionCall(arguments='s'))
+
         """
         obj = parsed[index]
         if not isinstance(obj, dict) or not obj:
diff --git a/vllm/tool_parsers/cohere_command_tool_parser.py b/vllm/tool_parsers/cohere_command_tool_parser.py
index 3739b222e5ff..eca5c5496729 100644
--- a/vllm/tool_parsers/cohere_command_tool_parser.py
+++ b/vllm/tool_parsers/cohere_command_tool_parser.py
@@ -1,148 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from collections.abc import Sequence
+from typing import ClassVar
 
-try:
-    from cohere_melody import PyFilter, PyFilterOptions
-except ImportError as e:
-    raise ImportError(
-        "The Cohere tool parser requires the `cohere_melody` "
-        "package, which is not installed. Install it with:\n"
-        "    pip install 'cohere-melody>=0.11.1'"
-    ) from e
-
-from vllm.entrypoints.generate.base.protocol import (
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-    ExtractedToolCallInformation,
-    FunctionCall,
-    ToolCall,
-)
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionRequest,
-)
-from vllm.entrypoints.openai.responses.protocol import (
-    ResponsesRequest,
-)
-from vllm.tokenizers import TokenizerLike
 from vllm.tool_parsers import ToolParser
-from vllm.tool_parsers.utils import Tool
 
 
 class BaseCohereCommandToolParser(ToolParser):
-    def __init__(
-        self,
-        tokenizer: TokenizerLike,
-        streaming_opts: PyFilterOptions,
-        unary_opts: PyFilterOptions,
-    ):
-        super().__init__(tokenizer)
-        self.melody_streaming = PyFilter(streaming_opts)
-        self.melody_unary = PyFilter(unary_opts)
-        # Melody can emit the tool-call id before the function name. Keep it
-        # until the first real name delta so clients receive both together.
-        self._pending_streaming_tool_call_ids: dict[int, str] = {}
-
-    def adjust_request(
-        self, request: ChatCompletionRequest | ResponsesRequest
-    ) -> ChatCompletionRequest | ResponsesRequest:
-        request = super().adjust_request(request)
-        request.skip_special_tokens = False
-        return request
-
-    def extract_tool_calls_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-        request: ChatCompletionRequest,
-    ) -> DeltaMessage | None:
-        r = self.melody_streaming.write_decoded(delta_text)
-        if r.content is not None:
-            return DeltaMessage(content=r.content)
-        if r.reasoning is not None:
-            return DeltaMessage(reasoning=r.reasoning)
-        if r.tool_calls:
-            tool_calls: list[DeltaToolCall] = []
-            for tc in r.tool_calls:
-                if tc.id:
-                    self._pending_streaming_tool_call_ids[tc.index] = tc.id
-                name = tc.name or None
-                arguments = tc.arguments or None
-                # Empty strings are placeholders in Melody's streaming output;
-                # omit them from OpenAI-compatible deltas instead of sending
-                # invalid tool names or empty argument fragments.
-                if name is None and arguments is None:
-                    continue
-
-                function_kwargs = {}
-                if name is not None:
-                    function_kwargs["name"] = name
-                if arguments is not None:
-                    function_kwargs["arguments"] = arguments
-                tool_call_kwargs = {
-                    "index": tc.index,
-                    "function": DeltaFunctionCall(**function_kwargs),
-                }
-                if name is not None:
-                    tool_call_id = tc.id or self._pending_streaming_tool_call_ids.pop(
-                        tc.index, None
-                    )
-                    if tool_call_id is not None:
-                        tool_call_kwargs["id"] = tool_call_id
-                    tool_call_kwargs["type"] = "function"
-                tool_calls.append(DeltaToolCall(**tool_call_kwargs))
-            if tool_calls:
-                return DeltaMessage(tool_calls=tool_calls)
-        return None
-
-    def extract_tool_calls(
-        self,
-        model_output: str,
-        request: ChatCompletionRequest,
-    ) -> ExtractedToolCallInformation:
-        result = self.melody_unary.process_full_text(model_output)
-        tool_calls = [
-            ToolCall(
-                id=tc.id,
-                type="function",
-                function=FunctionCall(name=tc.name, arguments=tc.arguments),
-            )
-            for tc in result.tool_calls
-        ]
-        return ExtractedToolCallInformation(
-            tools_called=len(tool_calls) > 0,
-            tool_calls=tool_calls,
-            content=result.content,
-        )
+    melody_preset: ClassVar[str]
 
 
 class CohereCommand3ToolParser(BaseCohereCommandToolParser):
-    def __init__(
-        self,
-        tokenizer: TokenizerLike,
-        tools: list[Tool] | None = None,
-    ):
-        super().__init__(
-            tokenizer,
-            streaming_opts=PyFilterOptions().cmd3().start_in_answer(),
-            unary_opts=PyFilterOptions().cmd3().start_in_answer(),
-        )
+    melody_preset = "cmd3"
 
 
 class CohereCommand4ToolParser(BaseCohereCommandToolParser):
-    def __init__(
-        self,
-        tokenizer: TokenizerLike,
-        tools: list[Tool] | None = None,
-    ):
-        super().__init__(
-            tokenizer,
-            streaming_opts=PyFilterOptions().cmd4().start_in_answer(),
-            unary_opts=PyFilterOptions().cmd4().start_in_answer(),
-        )
+    melody_preset = "cmd4"
diff --git a/vllm/tool_parsers/deepseekv41_engine_tool_parser.py b/vllm/tool_parsers/deepseekv41_engine_tool_parser.py
index a4c699aeabc0..a94883d06a52 100644
--- a/vllm/tool_parsers/deepseekv41_engine_tool_parser.py
+++ b/vllm/tool_parsers/deepseekv41_engine_tool_parser.py
@@ -5,4 +5,4 @@
 
 
 class DeepSeekV41EngineToolParser(DeepSeekV41ParserToolAdapter):  # type: ignore[valid-type, misc]
-    structural_tag_model = "deepseek_v41"
+    structural_tag_model = "deepseek_v4_1"
diff --git a/vllm/tool_parsers/ernie45_tool_parser.py b/vllm/tool_parsers/ernie45_tool_parser.py
index 74624b14048d..3c654b44897e 100644
--- a/vllm/tool_parsers/ernie45_tool_parser.py
+++ b/vllm/tool_parsers/ernie45_tool_parser.py
@@ -29,8 +29,7 @@
 
 class Ernie45ToolParser(ToolParser):
     def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
-        """
-        Ernie thinking model format:
+        """Ernie thinking model format:
         abc\n</think>\n\n\n<tool_call>\ndef\n</tool_call>\n
         """
         super().__init__(tokenizer, tools)
diff --git a/vllm/tool_parsers/functiongemma_tool_parser.py b/vllm/tool_parsers/functiongemma_tool_parser.py
index a975f68368a5..708ecb1dd997 100644
--- a/vllm/tool_parsers/functiongemma_tool_parser.py
+++ b/vllm/tool_parsers/functiongemma_tool_parser.py
@@ -27,8 +27,7 @@
 
 
 class FunctionGemmaToolParser(ToolParser):
-    """
-    Tool parser for Google's FunctionGemma model (google/functiongemma-270m-it).
+    """Tool parser for Google's FunctionGemma model (google/functiongemma-270m-it).
 
     Handles the FunctionGemma function call format:
     <start_function_call>call:func_name{param:<escape>value<escape>}<end_function_call>
diff --git a/vllm/tool_parsers/gemma4_utils.py b/vllm/tool_parsers/gemma4_utils.py
index d9aad254c728..33ad20b19396 100644
--- a/vllm/tool_parsers/gemma4_utils.py
+++ b/vllm/tool_parsers/gemma4_utils.py
@@ -52,6 +52,7 @@ def _parse_tool_arguments(args_str: str) -> dict[str, str]:
 
     Returns:
         Dictionary of argument name → string value.
+
     """
     if not args_str or not args_str.strip():
         return {}
@@ -95,6 +96,7 @@ def parse_tool_calls(text: str, *, strict: bool = False) -> list[dict]:
         >>> tool_calls = parse_tool_calls(output)
         >>> for tc in tool_calls:
         ...     print(f"Call: {tc['name']}({tc['arguments']})")
+
     """
     results = []
 
@@ -150,6 +152,7 @@ def has_tool_response_tag(text: str) -> bool:
         >>> if not has_tool_response_tag(model_output):
         ...     # Model used <eos> instead — inject <|tool_response> manually
         ...     next_prompt = "<|tool_response>" + tool_result
+
     """
     stripped = text.rstrip()
     return stripped.endswith(_TOOL_RESPONSE_START_TAG)
diff --git a/vllm/tool_parsers/gptoss_tool_parser.py b/vllm/tool_parsers/gptoss_tool_parser.py
index e147ff1c7418..80881f69a5f5 100644
--- a/vllm/tool_parsers/gptoss_tool_parser.py
+++ b/vllm/tool_parsers/gptoss_tool_parser.py
@@ -15,8 +15,7 @@
 
 
 class GptOssToolParser(ToolParser):
-    """
-    Stub tool parser for gpt-oss/harmony models.
+    """Stub tool parser for gpt-oss/harmony models.
 
     All output parsing is handled by HarmonyParser. This stub exists as a
     capability declaration via HarmonyParser.tool_parser_cls.
diff --git a/vllm/tool_parsers/granite_20b_fc_tool_parser.py b/vllm/tool_parsers/granite_20b_fc_tool_parser.py
index 432f94b61e5d..4c6c0a1b0efe 100644
--- a/vllm/tool_parsers/granite_20b_fc_tool_parser.py
+++ b/vllm/tool_parsers/granite_20b_fc_tool_parser.py
@@ -38,8 +38,7 @@
 
 
 class Granite20bFCToolParser(ToolParser):
-    """
-    Tool call parser for the granite-20b-functioncalling model intended
+    """Tool call parser for the granite-20b-functioncalling model intended
     for use with the examples/tool_chat_template_granite20b_fc.jinja
     template.
 
diff --git a/vllm/tool_parsers/granite_engine_tool_parser.py b/vllm/tool_parsers/granite_engine_tool_parser.py
new file mode 100644
index 000000000000..1fb701d5a422
--- /dev/null
+++ b/vllm/tool_parsers/granite_engine_tool_parser.py
@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.parser.engine.registered_adapters import GraniteParserToolAdapter
+
+
+class GraniteEngineToolParser(GraniteParserToolAdapter):  # type: ignore[valid-type, misc]
+    structural_tag_model = None
diff --git a/vllm/tool_parsers/granite_tool_parser.py b/vllm/tool_parsers/granite_tool_parser.py
deleted file mode 100644
index c1d852dfc352..000000000000
--- a/vllm/tool_parsers/granite_tool_parser.py
+++ /dev/null
@@ -1,258 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import json
-from collections.abc import Sequence
-
-import partial_json_parser
-from partial_json_parser.core.options import Allow
-
-from vllm.entrypoints.chat_utils import make_tool_call_id
-from vllm.entrypoints.generate.base.protocol import (
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-    ExtractedToolCallInformation,
-    FunctionCall,
-    ToolCall,
-)
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionRequest,
-)
-from vllm.logger import init_logger
-from vllm.tokenizers import TokenizerLike
-from vllm.tool_parsers.abstract_tool_parser import (
-    Tool,
-    ToolParser,
-)
-from vllm.tool_parsers.utils import (
-    consume_space,
-    find_common_prefix,
-    is_complete_json,
-    partial_json_loads,
-)
-
-logger = init_logger(__name__)
-
-
-class GraniteToolParser(ToolParser):
-    """
-    Tool call parser for the granite 3.0 models. Intended
-    for use with the examples/tool_chat_template_granite.jinja
-    template.
-
-    Used when --enable-auto-tool-choice --tool-call-parser granite
-    are all set
-    """
-
-    def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
-        super().__init__(tokenizer, tools)
-        # for granite 3.0, the token `<|tool_call|>`
-        self.bot_token = "<|tool_call|>"
-        # for granite 3.1, the string `<tool_call>`
-        self.bot_string = "<tool_call>"
-
-    def extract_tool_calls(
-        self, model_output: str, request: ChatCompletionRequest
-    ) -> ExtractedToolCallInformation:
-        stripped = (
-            model_output.strip()
-            .removeprefix(self.bot_token)
-            .removeprefix(self.bot_string)
-            .lstrip()
-        )
-        if not stripped or stripped[0] != "[":
-            return ExtractedToolCallInformation(
-                tools_called=False, tool_calls=[], content=model_output
-            )
-        try:
-            raw_function_calls = json.loads(stripped)
-            if not isinstance(raw_function_calls, list):
-                raise Exception(
-                    f"Expected dict or list, got {type(raw_function_calls)}"
-                )
-
-            logger.debug("Extracted %d tool calls", len(raw_function_calls))
-            tool_calls = [
-                ToolCall(
-                    type="function",
-                    function=FunctionCall(
-                        name=function_call["name"],
-                        # function call args are JSON but as a string
-                        arguments=json.dumps(
-                            function_call["arguments"], ensure_ascii=False
-                        ),
-                    ),
-                )
-                for function_call in raw_function_calls
-            ]
-
-            return ExtractedToolCallInformation(
-                tools_called=True,
-                tool_calls=tool_calls,
-                content=None,
-            )
-
-        except Exception as e:
-            logger.error("Error in extracting tool call from response %s", e)
-            return ExtractedToolCallInformation(
-                tools_called=False, tool_calls=[], content=model_output
-            )
-
-    def extract_tool_calls_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-        request: ChatCompletionRequest,
-    ) -> DeltaMessage | None:
-        start_idx = consume_space(0, current_text)
-        if current_text[start_idx:].startswith(self.bot_token):
-            start_idx = consume_space(start_idx + len(self.bot_token), current_text)
-        if current_text[start_idx:].startswith(self.bot_string):
-            start_idx = consume_space(start_idx + len(self.bot_string), current_text)
-        if (
-            not current_text
-            or start_idx >= len(current_text)
-            or current_text[start_idx] != "["
-        ):
-            return DeltaMessage(content=delta_text)
-
-        # bit mask flags for partial JSON parsing. If the name hasn't been
-        # sent yet, don't allow sending
-        # an incomplete string since OpenAI only ever (as far as I have
-        # seen) allows sending the entire tool/ function name at once.
-        flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR
-        try:
-            tool_call_arr = None
-            is_complete = None
-            try:
-                tool_calls, end_idx = partial_json_loads(
-                    current_text[start_idx:], flags
-                )
-                if type(tool_calls) is list:
-                    tool_call_arr = tool_calls
-                else:
-                    return DeltaMessage(content=delta_text)
-
-                is_complete = [True] * len(tool_calls)
-                if not is_complete_json(current_text[start_idx : start_idx + end_idx]):
-                    is_complete[-1] = False
-            except partial_json_parser.core.exceptions.MalformedJSON:
-                logger.debug("not enough tokens to parse into JSON yet")
-                return None
-
-            # case -- if no tokens have been streamed for the tool, e.g.
-            #   only the array brackets, stream nothing
-            if not tool_call_arr:
-                return None
-
-            # select as the current tool call the one we're on the state at
-            current_tool_call: dict = tool_call_arr[self.current_tool_id]
-
-            delta = None
-            # Only advance once the current tool name is streamed; granite
-            # emits arguments before name, so advancing early would drop it.
-            if len(tool_call_arr) > self.current_tool_id + 1 and (
-                self.current_tool_id < 0 or self.current_tool_name_sent
-            ):
-                # if we're moving on to a new call, first make sure we
-                # haven't missed anything in the previous one that was
-                # auto-generated due to JSON completions, but wasn't
-                # streamed to the client yet.
-                if self.current_tool_id >= 0:
-                    cur_arguments = current_tool_call.get("arguments")
-                    if cur_arguments:
-                        cur_args_json = json.dumps(cur_arguments, ensure_ascii=False)
-                        sent = len(self.streamed_args_for_tool[self.current_tool_id])
-                        argument_diff = cur_args_json[sent:]
-
-                        logger.debug("got arguments diff: %s", argument_diff)
-                        delta = DeltaMessage(
-                            tool_calls=[
-                                DeltaToolCall(
-                                    index=self.current_tool_id,
-                                    function=DeltaFunctionCall(
-                                        arguments=argument_diff
-                                    ).model_dump(exclude_none=True),
-                                )
-                            ]
-                        )
-                        self.streamed_args_for_tool[self.current_tool_id] += (
-                            argument_diff
-                        )
-
-                # re-set stuff pertaining to progress in the current tool
-                self.current_tool_id += 1
-                self.current_tool_name_sent = False
-                self.streamed_args_for_tool.append("")
-                logger.debug("starting on new tool %d", self.current_tool_id)
-                return delta
-
-            # if the current tool name hasn't been sent, send if available
-            # - otherwise send nothing
-            elif not self.current_tool_name_sent:
-                function_name = current_tool_call.get("name")
-                if function_name:
-                    delta = DeltaMessage(
-                        tool_calls=[
-                            DeltaToolCall(
-                                index=self.current_tool_id,
-                                type="function",
-                                id=make_tool_call_id(),
-                                function=DeltaFunctionCall(
-                                    name=function_name
-                                ).model_dump(exclude_none=True),
-                            )
-                        ]
-                    )
-                    self.current_tool_name_sent = True
-
-            # now we know we're on the same tool call and we're streaming
-            # arguments
-            else:
-                cur_arguments = current_tool_call.get("arguments")
-
-                if cur_arguments:
-                    sent = len(self.streamed_args_for_tool[self.current_tool_id])
-                    cur_args_json = json.dumps(cur_arguments, ensure_ascii=False)
-                    prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(
-                        "arguments"
-                    )
-
-                    argument_diff = None
-                    if is_complete[self.current_tool_id]:
-                        argument_diff = cur_args_json[sent:]
-                    elif prev_arguments:
-                        prev_args_json = json.dumps(prev_arguments, ensure_ascii=False)
-                        if cur_args_json != prev_args_json:
-                            prefix = find_common_prefix(prev_args_json, cur_args_json)
-                            argument_diff = prefix[sent:]
-
-                    if argument_diff is not None:
-                        delta = DeltaMessage(
-                            tool_calls=[
-                                DeltaToolCall(
-                                    index=self.current_tool_id,
-                                    function=DeltaFunctionCall(
-                                        arguments=argument_diff
-                                    ).model_dump(exclude_none=True),
-                                )
-                            ]
-                        )
-                        self.streamed_args_for_tool[self.current_tool_id] += (
-                            argument_diff
-                        )
-
-            self.prev_tool_call_arr = tool_call_arr
-            return delta
-
-        except Exception as e:
-            logger.error("Error trying to handle streaming tool call: %s", e)
-            logger.debug(
-                "Skipping chunk as a result of tool streaming extraction error"
-            )
-            return None
diff --git a/vllm/tool_parsers/hunyuan_a13b_tool_parser.py b/vllm/tool_parsers/hunyuan_a13b_tool_parser.py
index 88b2a20d3a0d..90e291084fd3 100644
--- a/vllm/tool_parsers/hunyuan_a13b_tool_parser.py
+++ b/vllm/tool_parsers/hunyuan_a13b_tool_parser.py
@@ -100,9 +100,7 @@ def preprocess_model_output(
     def extract_tool_calls(
         self, model_output: str, request: ChatCompletionRequest
     ) -> ExtractedToolCallInformation:
-        """
-        Extract tool calls from a complete model output.
-        """
+        """Extract tool calls from a complete model output."""
         try:
             # Preprocess the model output
             content, potential_tool_calls = self.preprocess_model_output(model_output)
@@ -177,10 +175,7 @@ def extract_tool_calls_streaming(
         delta_token_ids: Sequence[int],
         request: ChatCompletionRequest,
     ) -> DeltaMessage | None:
-        """
-        Extract tool calls for streaming mode.
-        """
-
+        """Extract tool calls for streaming mode."""
         start_idx = consume_space(0, current_text)
         if current_text[start_idx:].startswith(self.bot_string):
             start_idx = consume_space(start_idx + len(self.bot_string), current_text)
diff --git a/vllm/tool_parsers/hy_v3_tool_parser.py b/vllm/tool_parsers/hy_v3_tool_parser.py
index d73473673fdc..0effb377d0a6 100644
--- a/vllm/tool_parsers/hy_v3_tool_parser.py
+++ b/vllm/tool_parsers/hy_v3_tool_parser.py
@@ -181,6 +181,7 @@ def _try_parse_wildcard_number(value: str) -> int | float | None:
             "-3"   → int(-3)
 
         Return None on failure.
+
         """
         try:
             if "." in value or "e" in value or "E" in value:
diff --git a/vllm/tool_parsers/hy_v4_tool_parser.py b/vllm/tool_parsers/hy_v4_tool_parser.py
index 729463162693..b9c857dbd9e6 100644
--- a/vllm/tool_parsers/hy_v4_tool_parser.py
+++ b/vllm/tool_parsers/hy_v4_tool_parser.py
@@ -25,6 +25,7 @@
 from vllm.logger import init_logger
 from vllm.tokenizers import TokenizerLike
 from vllm.tool_parsers.abstract_tool_parser import ToolParser
+from vllm.tool_parsers.tool_strict_level import ToolStrictLevel
 from vllm.tool_parsers.utils import partial_tag_overlap
 
 if TYPE_CHECKING:
@@ -101,8 +102,8 @@ def detect_token_suffix(tokenizer: TokenizerLike) -> str:
         RuntimeError: The tokenizer declares the structural tokens through
             ``model_specific_special_tokens``, which transformers 5 no longer
             round-trips.
-    """
 
+    """
     import transformers
 
     if int(transformers.__version__.split(".")[0]) >= 5:
@@ -611,6 +612,7 @@ def extract_tool_calls_streaming(
         Returns:
             A streaming delta carrying content and/or the tool calls drained
             from the buffer, or None when nothing can be emitted yet.
+
         """
         content_delta: str | None = None
         tool_calls: list[StreamToolCall] = []
@@ -993,6 +995,7 @@ def get_structural_tag(
         request: ChatCompletionRequest | ResponsesRequest,
         *,
         reasoning: bool = False,
+        strict_level: ToolStrictLevel = ToolStrictLevel.AUTO,
     ) -> StructuralTag | None:
         """Build a structural tag matching HYV4's tool tokens.
 
@@ -1008,9 +1011,11 @@ def get_structural_tag(
         Args:
             request: The request being adjusted.
             reasoning: Whether the grammar also covers the reasoning phase.
+            strict_level: Server-side floor from ``--tool-strict-level``.
 
         Returns:
             The structural tag, or None when structural tagging does not apply.
+
         """
         if not envs.VLLM_ENFORCE_STRICT_TOOL_CALLING:
             return None
@@ -1039,6 +1044,7 @@ def get_structural_tag(
                 tool_choice=request.tool_choice,
                 reasoning=reasoning,
                 token_suffix=self._extractor.token_suffix,
+                strict_level=strict_level,
             )
         except Exception:
             logger.warning(
@@ -1070,6 +1076,7 @@ def _is_guided(request: ChatCompletionRequest) -> bool:
 
         Returns:
             True when the streaming parser must use the string-marker path.
+
         """
         structured_outputs = getattr(request, "structured_outputs", None)
         return (
diff --git a/vllm/tool_parsers/lfm2_tool_parser.py b/vllm/tool_parsers/lfm2_tool_parser.py
index 6f4db492be29..2a0bab07e2a0 100644
--- a/vllm/tool_parsers/lfm2_tool_parser.py
+++ b/vllm/tool_parsers/lfm2_tool_parser.py
@@ -42,8 +42,7 @@
 
 
 class Lfm2ToolParser(ToolParser):
-    """
-    Tool call parser for LiquidAI LFM2/LFM2.5 models that produce pythonic
+    """Tool call parser for LiquidAI LFM2/LFM2.5 models that produce pythonic
     tool calls wrapped in <|tool_call_start|> and <|tool_call_end|> tokens.
 
     Example model output:
diff --git a/vllm/tool_parsers/llama4_pythonic_tool_parser.py b/vllm/tool_parsers/llama4_pythonic_tool_parser.py
index 18c697f04829..fdffb48ff7fe 100644
--- a/vllm/tool_parsers/llama4_pythonic_tool_parser.py
+++ b/vllm/tool_parsers/llama4_pythonic_tool_parser.py
@@ -31,8 +31,7 @@
 
 
 class Llama4PythonicToolParser(ToolParser):
-    """
-    Toolcall parser for Llama4 that produce tool calls in a pythonic style
+    """Toolcall parser for Llama4 that produce tool calls in a pythonic style
     Use --enable-auto-tool-choice --tool-call-parser llama4_pythonic
     """
 
@@ -67,10 +66,7 @@ def current_tool_index(self, value: int) -> None:
     def extract_tool_calls(
         self, model_output: str, request: ChatCompletionRequest
     ) -> ExtractedToolCallInformation:
-        """
-        Extract the tool calls from a complete model response.
-        """
-
+        """Extract the tool calls from a complete model response."""
         # remove <|python_start|> and <|python_end|>
         # as Llama 4 model sometime will output those tokens
         if model_output.startswith("<|python_start|>"):
diff --git a/vllm/tool_parsers/llama_tool_parser.py b/vllm/tool_parsers/llama_tool_parser.py
index 5ad7b1cbdd29..93b0980c3775 100644
--- a/vllm/tool_parsers/llama_tool_parser.py
+++ b/vllm/tool_parsers/llama_tool_parser.py
@@ -37,8 +37,7 @@
 
 
 class Llama3JsonToolParser(ToolParser):
-    """
-    Tool call parser for Llama 3.x and 4 models intended for use with the
+    """Tool call parser for Llama 3.x and 4 models intended for use with the
     examples/tool_chat_template_llama.jinja template.
 
     Used when --enable-auto-tool-choice --tool-call-parser llama3_json or
@@ -77,8 +76,7 @@ def __init__(
     def extract_tool_calls(
         self, model_output: str, request: ChatCompletionRequest
     ) -> ExtractedToolCallInformation:
-        """
-        Extract the tool calls from a complete model response.
+        """Extract the tool calls from a complete model response.
         Only extracts JSON content and ignores any surrounding plain text.
         Supports both single JSON and multiple JSONs separated by semicolons.
         """
diff --git a/vllm/tool_parsers/olmo3_tool_parser.py b/vllm/tool_parsers/olmo3_tool_parser.py
index 49a871653399..cb7bc632baa8 100644
--- a/vllm/tool_parsers/olmo3_tool_parser.py
+++ b/vllm/tool_parsers/olmo3_tool_parser.py
@@ -31,8 +31,7 @@
 
 
 class Olmo3PythonicToolParser(ToolParser):
-    """
-    Tool call parser for Olmo 3 models that produce tool calls as
+    """Tool call parser for Olmo 3 models that produce tool calls as
     newline-separated pythonic strings.
     Used when --enable-auto-tool-choice --tool-call-parser pythonic are all set
     Code copied from pythonic_tool_parser.py and updated to handle
@@ -71,9 +70,7 @@ def current_tool_index(self, value: int) -> None:
     def extract_tool_calls(
         self, model_output: str, request: ChatCompletionRequest
     ) -> ExtractedToolCallInformation:
-        """
-        Extract the tool calls from a complete model response.
-        """
+        """Extract the tool calls from a complete model response."""
         original_model_output = model_output
         # Remove xml tags.
         match = re.search(
diff --git a/vllm/tool_parsers/phi4mini_tool_parser.py b/vllm/tool_parsers/phi4mini_tool_parser.py
index f5b794f0f67b..820e797f4d9f 100644
--- a/vllm/tool_parsers/phi4mini_tool_parser.py
+++ b/vllm/tool_parsers/phi4mini_tool_parser.py
@@ -28,8 +28,7 @@
 
 
 class Phi4MiniJsonToolParser(ToolParser):
-    """
-    Tool call parser for phi-4-mini models intended for use with the
+    """Tool call parser for phi-4-mini models intended for use with the
     examples/tool_chat_template_llama.jinja template.
 
     Used when --enable-auto-tool-choice --tool-call-parser phi4_mini_json
@@ -55,9 +54,7 @@ def __init__(
     def extract_tool_calls(
         self, model_output: str, request: ChatCompletionRequest
     ) -> ExtractedToolCallInformation:
-        """
-        Extract the tool calls from a complete model response.
-        """
+        """Extract the tool calls from a complete model response."""
         logger.debug("Model output: %s", model_output)
 
         pattern = r"functools\[(.*?)\]"
diff --git a/vllm/tool_parsers/poolside_v1_tool_parser.py b/vllm/tool_parsers/poolside_v1_tool_parser.py
index 693580e85a2e..178036760634 100644
--- a/vllm/tool_parsers/poolside_v1_tool_parser.py
+++ b/vllm/tool_parsers/poolside_v1_tool_parser.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-GLM-4 Tool Call Parser with incremental string streaming support.
+"""GLM-4 Tool Call Parser with incremental string streaming support.
 
 This parser fixes the streaming issue reported in Issue #32829 where long string
 parameters (e.g., file content with 4000+ characters of code) are buffered until
diff --git a/vllm/tool_parsers/pythonic_tool_parser.py b/vllm/tool_parsers/pythonic_tool_parser.py
index cc584d3d7c95..76b263cfc5bb 100644
--- a/vllm/tool_parsers/pythonic_tool_parser.py
+++ b/vllm/tool_parsers/pythonic_tool_parser.py
@@ -31,8 +31,7 @@
 
 
 class PythonicToolParser(ToolParser):
-    """
-    Tool call parser for models that produce tool calls in a pythonic style,
+    """Tool call parser for models that produce tool calls in a pythonic style,
     such as Llama 3.2 and Llama 4 models.
 
     Used when --enable-auto-tool-choice --tool-call-parser pythonic are all set
@@ -69,9 +68,7 @@ def current_tool_index(self, value: int) -> None:
     def extract_tool_calls(
         self, model_output: str, request: ChatCompletionRequest
     ) -> ExtractedToolCallInformation:
-        """
-        Extract the tool calls from a complete model response.
-        """
+        """Extract the tool calls from a complete model response."""
         is_tool_call_pattern = False
         try:
             is_tool_call_pattern = (
diff --git a/vllm/tool_parsers/step3_tool_parser.py b/vllm/tool_parsers/step3_tool_parser.py
index 929e0c19c600..514945cd0d70 100644
--- a/vllm/tool_parsers/step3_tool_parser.py
+++ b/vllm/tool_parsers/step3_tool_parser.py
@@ -32,8 +32,7 @@
 
 
 class Step3ToolParser(ToolParser):
-    """
-    Tool parser for a model that uses a specific XML-like format for tool calls.
+    """Tool parser for a model that uses a specific XML-like format for tool calls.
     This version uses a robust, stateful, cursor-based streaming parser and
     consolidates tool arguments into a single message.
     """
diff --git a/vllm/tool_parsers/step3p5_tool_parser.py b/vllm/tool_parsers/step3p5_tool_parser.py
index 74fcb56ec46f..2639d2f2043a 100644
--- a/vllm/tool_parsers/step3p5_tool_parser.py
+++ b/vllm/tool_parsers/step3p5_tool_parser.py
@@ -28,8 +28,7 @@
 
 
 class StreamingXMLToolCallParser:
-    """
-    Simplified streaming XML tool call parser
+    """Simplified streaming XML tool call parser
     Supports streaming input, parsing, and output
     """
 
@@ -46,8 +45,7 @@ def __init__(self):
         self.parameter_end_token: str = "</parameter>"
 
     def reset_streaming_state(self):
-        """Reset streaming parsing state"""
-
+        """Reset streaming parsing state."""
         self.deltas = []
         # state for streaming
         self.tool_call_index = 0
@@ -80,8 +78,7 @@ def reset_streaming_state(self):
         self.setup_parser()
 
     def parse_single_streaming_chunks(self, xml_chunk: str) -> DeltaMessage:
-        """
-        Parse single streaming XML chunk and return Delta response
+        """Parse single streaming XML chunk and return Delta response
         This is the actual streaming interface that receives chunks
         one by one and maintains internal state
 
@@ -90,6 +87,7 @@ def parse_single_streaming_chunks(self, xml_chunk: str) -> DeltaMessage:
         Returns:
             DeltaMessage: Contains delta information generated by this chunk,
             returns empty response if no complete elements
+
         """
         # Record delta count before processing
         initial_delta_count = len(self.deltas)
@@ -218,8 +216,7 @@ def parse_single_streaming_chunks(self, xml_chunk: str) -> DeltaMessage:
             return DeltaMessage(content=None)
 
     def _escape_xml_special_chars(self, text: str) -> str:
-        """
-        Escape XML special characters
+        """Escape XML special characters
         Args:
             text: Original text
         Returns:
@@ -239,11 +236,11 @@ def _escape_xml_special_chars(self, text: str) -> str:
         return text
 
     def _process_complete_xml_elements(self) -> bool:
-        """
-        Process complete XML elements in buffer
+        """Process complete XML elements in buffer.
 
         Returns:
             bool: Whether complete elements were found and processed
+
         """
         found_any = False
 
@@ -320,8 +317,7 @@ def _process_complete_xml_elements(self) -> bool:
         return found_any
 
     def _fix_incomplete_tag_in_chunk(self, chunk: str) -> str:
-        """
-        Fallback: fix incomplete <parameter=xxx or <function=xxx tags
+        """Fallback: fix incomplete <parameter=xxx or <function=xxx tags
         (missing >)
         Examples: <parameter=-C: -> <parameter=-C>, <parameter=parameter=-n:
         -> <parameter=-n>
@@ -386,8 +382,7 @@ def _fix_incomplete_tag_in_chunk(self, chunk: str) -> str:
         return chunk
 
     def _fix_missing_equals_in_function_tag(self, chunk: str) -> str:
-        """
-        Fix missing = in function tags: <function xxx> or <functionxxx>
+        """Fix missing = in function tags: <function xxx> or <functionxxx>
         Examples:
           <function execute_bash> -> <function=execute_bash>
           <functionexecute_bash> -> <function=execute_bash>
@@ -425,7 +420,7 @@ def _fix_missing_equals_in_function_tag(self, chunk: str) -> str:
         return chunk
 
     def _validate_function_name(self, func_name: str) -> bool:
-        """Check if function name exists in tool definitions"""
+        """Check if function name exists in tool definitions."""
         if not self.tools:
             return False
 
@@ -442,7 +437,7 @@ def _validate_function_name(self, func_name: str) -> bool:
         return False
 
     def _validate_parameter_name(self, param_name: str) -> bool:
-        """Check if parameter exists in current function's tool definition"""
+        """Check if parameter exists in current function's tool definition."""
         if not self.tools or not self.current_function_name:
             return True
 
@@ -465,16 +460,15 @@ def _validate_parameter_name(self, param_name: str) -> bool:
         return True
 
     def _should_skip_element(self, element: str) -> bool:
-        """
-        Determine whether an element should be skipped
+        """Determine whether an element should be skipped.
 
         Args:
             element: Element to evaluate
 
         Returns:
             bool: True means should skip, False means should process
-        """
 
+        """
         # If it's a tool_call XML tag, don't skip
         if (
             element.startswith(self.tool_call_start_token)
@@ -501,8 +495,7 @@ def _should_skip_element(self, element: str) -> bool:
         return not element
 
     def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]:
-        """
-        Find next complete XML element from specified position
+        """Find next complete XML element from specified position.
 
         Args:
             start_pos: Position to start searching
@@ -510,6 +503,7 @@ def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]:
         Returns:
             (Complete element string, element end position),
             returns (None, start_pos) if no complete element found
+
         """
         buffer = self.streaming_buffer[start_pos:]
 
@@ -588,8 +582,7 @@ def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]:
                 return remaining, start_pos + len(remaining)
 
     def _merge_new_deltas_to_single_response(self, initial_count: int) -> DeltaMessage:
-        """
-        Merge newly generated deltas from this processing
+        """Merge newly generated deltas from this processing
         into a single DeltaMessage
 
         Args:
@@ -597,6 +590,7 @@ def _merge_new_deltas_to_single_response(self, initial_count: int) -> DeltaMessa
 
         Returns:
             Merged DeltaMessage containing all newly generated delta information
+
         """
         if len(self.deltas) <= initial_count:
             return DeltaMessage(content=None)
@@ -652,8 +646,7 @@ def _merge_new_deltas_to_single_response(self, initial_count: int) -> DeltaMessa
         )
 
     def _preprocess_xml_chunk(self, chunk: str) -> str:
-        """
-        Preprocess XML chunk, handle non-standard formats,
+        """Preprocess XML chunk, handle non-standard formats,
         and escape special characters
 
         Args:
@@ -661,8 +654,8 @@ def _preprocess_xml_chunk(self, chunk: str) -> str:
 
         Returns:
             Processed XML chunk
-        """
 
+        """
         # Check if this is a tool_call related element
         is_tool_call = False
         if chunk.startswith(self.tool_call_start_token) or chunk.startswith(
@@ -785,7 +778,7 @@ def _preprocess_xml_chunk(self, chunk: str) -> str:
         return processed
 
     def _emit_delta(self, delta: DeltaMessage):
-        """Emit Delta response (streaming output)"""
+        """Emit Delta response (streaming output)."""
         self.deltas.append(delta)
 
     def _auto_close_open_parameter_if_needed(self, incoming_tag: str | None = None):
@@ -814,8 +807,7 @@ def _auto_close_open_parameter_if_needed(self, incoming_tag: str | None = None):
             self._end_element("tool_call")
 
     def _start_element(self, name: str, attrs: dict[str, str]):
-        """Handle XML start element events"""
-
+        """Handle XML start element events."""
         if name == "root":
             return
 
@@ -904,7 +896,7 @@ def _start_element(self, name: str, attrs: dict[str, str]):
                     self.current_param_is_first = False
 
     def _char_data(self, data: str):
-        """Handle XML character data events"""
+        """Handle XML character data events."""
         if data and self.current_param_name:
             # If preprocessing stage determines deferred parsing is needed,
             # only cache character data, no streaming output
@@ -979,8 +971,7 @@ def _char_data(self, data: str):
             self._emit_delta(delta)
 
     def _end_element(self, name: str):
-        """Handle XML end element events"""
-
+        """Handle XML end element events."""
         if name == "root":
             return
 
@@ -1154,18 +1145,18 @@ def _end_element(self, name: str):
             self._reset_xml_parser_after_tool_call()
 
     def setup_parser(self):
-        """Set up XML parser event handlers"""
+        """Set up XML parser event handlers."""
         self.parser.buffer_text = True
         self.parser.StartElementHandler = self._start_element
         self.parser.EndElementHandler = self._end_element
         self.parser.CharacterDataHandler = self._char_data
 
     def set_tools(self, tools: list[Tool] | None):
-        """Set tool configuration information"""
+        """Set tool configuration information."""
         self.tools = tools
 
     def _extract_function_name(self, name: str, attrs: dict[str, str]) -> str | None:
-        """Extract function name from various formats"""
+        """Extract function name from various formats."""
         if attrs and "name" in attrs:
             return attrs["name"]
 
@@ -1177,7 +1168,7 @@ def _extract_function_name(self, name: str, attrs: dict[str, str]) -> str | None
         return None
 
     def _extract_parameter_name(self, name: str, attrs: dict[str, str]) -> str | None:
-        """Extract parameter name from various formats"""
+        """Extract parameter name from various formats."""
         if attrs and "name" in attrs:
             return attrs["name"]
 
@@ -1195,6 +1186,7 @@ def _get_param_type(self, param_name: str) -> str:
 
         Returns:
             Parameter type
+
         """
         if not self.tools or not self.current_function_name:
             return "string"
@@ -1235,6 +1227,7 @@ def repair_param_type(self, param_type: str) -> str:
 
         Returns:
             Repaired parameter type
+
         """
         if (
             param_type in ["string", "str", "text", "varchar", "char", "enum"]
@@ -1264,6 +1257,7 @@ def _convert_param_value(self, param_value: str, param_type: str) -> Any:
 
         Returns:
             Converted value
+
         """
         if param_value.lower() == "null":
             return None
@@ -1315,6 +1309,7 @@ def _convert_for_json_streaming(self, converted_value: Any, param_type: str) ->
 
         Returns:
             Converted string for streaming output
+
         """
         # Check if value is empty, but exclude numeric 0
         if converted_value is None or converted_value == "":
@@ -1331,11 +1326,9 @@ def _convert_for_json_streaming(self, converted_value: Any, param_type: str) ->
                 return converted_value
 
     def _reset_xml_parser_after_tool_call(self):
-        """
-        Each tool_call is treated as a separate XML document,
+        """Each tool_call is treated as a separate XML document,
         so we need to reset the parser after each tool_call.
         """
-
         # recreate XML parser
         self.parser = ParserCreate()
         self.setup_parser()
@@ -1513,7 +1506,5 @@ def extract_tool_calls_streaming(
         return result
 
     def parser_should_check_for_unstreamed_tool_arg_tokens(self) -> bool:
-        """
-        Skip the remaining_call calculation in serving_chat
-        """
+        """Skip the remaining_call calculation in serving_chat."""
         return False
diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py
index 34e57840bed7..ba86395b7a38 100644
--- a/vllm/tool_parsers/structural_tag_registry.py
+++ b/vllm/tool_parsers/structural_tag_registry.py
@@ -34,6 +34,7 @@
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionToolsParam,
 )
+from vllm.tool_parsers.tool_strict_level import ToolStrictLevel
 
 ToolChoice: TypeAlias = (
     Literal["none", "auto", "required"]
@@ -69,11 +70,10 @@
         "deepseek_v3_2",
         "glm_4_7",
         "deepseek_v4",
+        "deepseek_v4_1",
     }
 )
-VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset(
-    {"deepseek_v41", "hermes", "hy_v4", "kimi_k3"}
-)
+VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes", "hy_v4", "kimi_k3"})
 SUPPORTED_STRUCTURAL_TAG_MODELS = (
     XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS
 )
@@ -93,15 +93,57 @@ def decorator(func: StructuralTagBuilder) -> StructuralTagBuilder:
     return decorator
 
 
+def _tool_is_strict(tool: ChatCompletionToolsParam | ResponsesTool) -> bool:
+    if isinstance(tool, FunctionTool):
+        return tool.strict is True
+    if isinstance(tool, ChatCompletionToolsParam):
+        return tool.function.strict is True
+    return False
+
+
 def _any_tool_strict(
     tools: Sequence[ChatCompletionToolsParam | ResponsesTool],
 ) -> bool:
-    for tool in tools:
-        if isinstance(tool, FunctionTool) and tool.strict is True:
-            return True
-        if isinstance(tool, ChatCompletionToolsParam) and tool.function.strict is True:
-            return True
-    return False
+    return any(_tool_is_strict(tool) for tool in tools)
+
+
+def _with_tool_strict(
+    tool: ChatCompletionToolsParam | ResponsesTool,
+    strict: bool,
+) -> ChatCompletionToolsParam | ResponsesTool:
+    """Return a copy of ``tool`` with ``strict`` set, leaving the request's alone."""
+    if isinstance(tool, FunctionTool):
+        return tool.model_copy(update={"strict": strict})
+    if isinstance(tool, ChatCompletionToolsParam):
+        return tool.model_copy(
+            update={"function": tool.function.model_copy(update={"strict": strict})}
+        )
+    return tool
+
+
+def _resolve_tool_strictness(
+    tools: Sequence[ChatCompletionToolsParam | ResponsesTool],
+    tool_choice: ToolChoice,
+    strict_level: ToolStrictLevel,
+) -> Sequence[ChatCompletionToolsParam | ResponsesTool] | None:
+    """Decide whether a structural tag applies and pin each tool's ``strict``.
+
+    A tool without an explicit ``strict`` is treated as non-strict: its call
+    envelope is still constrained, but its arguments stay free unless the
+    server level is PARAMETER. ``None`` means no structural tag.
+    """
+    if (
+        tool_choice == "auto"
+        and strict_level == ToolStrictLevel.AUTO
+        and not _any_tool_strict(tools)
+    ):
+        return None
+    return [
+        _with_tool_strict(
+            tool, strict_level >= ToolStrictLevel.PARAMETER or _tool_is_strict(tool)
+        )
+        for tool in tools
+    ]
 
 
 def get_model_structural_tag(
@@ -110,13 +152,14 @@ def get_model_structural_tag(
     tool_choice: ToolChoice,
     reasoning: bool,
     token_suffix: str = "",
+    strict_level: ToolStrictLevel = ToolStrictLevel.AUTO,
 ) -> StructuralTag | None:
     """Build a structural tag with xgrammar's builtin model templates."""
-
     if not tools or tool_choice == "none":
         return None
 
-    if tool_choice == "auto" and not _any_tool_strict(tools):
+    tools = _resolve_tool_strictness(tools, tool_choice, strict_level)
+    if tools is None:
         return None
 
     dumped_tools = [_dump_tool_for_xgrammar(tool) for tool in tools]
@@ -157,7 +200,6 @@ def _dump_tool_for_xgrammar(
     tool: ChatCompletionToolsParam | ResponsesTool,
 ) -> dict[str, Any]:
     """Convert tool objects to xgrammar's Chat Completions tool protocol."""
-
     if isinstance(tool, FunctionTool):
         function: dict[str, Any] = {"name": tool.name}
         if tool.description is not None:
@@ -177,7 +219,6 @@ def _dump_tool_choice_for_xgrammar(
     tool_choice: ToolChoice,
 ) -> dict[str, Any] | str | None:
     """Convert tool_choice objects to xgrammar's expected protocol."""
-
     if tool_choice is None:
         return None
 
@@ -227,99 +268,6 @@ def get_function_parameters(function) -> dict[str, Any] | bool:
     return function.parameters if function.parameters is not None else True
 
 
-_V41_CALLS_START = "<｜DSML｜ calls>"
-_V41_CALLS_END = "</｜DSML｜ calls>"
-_V41_INVOKE_END = "</｜DSML｜ invoke>"
-_V41_PARAMETER_END = "</｜DSML｜ parameter>"
-
-
-@register_vllm_structural_tag("deepseek_v41")
-def get_deepseek_v41_structural_tag(
-    tools: list[FunctionToolParam],
-    builtin_tools: list[BuiltinToolParam],
-    tool_choice: SimplifiedToolChoice,
-    reasoning: bool,
-    token_suffix: str = "",
-) -> StructuralTag:
-    # Serving enables this visible-text grammar after the reasoning boundary.
-    del builtin_tools, reasoning, token_suffix
-
-    # TODO: lower parameter schemas into DSML constraints. This builder constrains
-    # tool names and DSML/value syntax only. Parameter names, presence, uniqueness
-    # and value schemas remain unconstrained, including for strict=true. The request
-    # layer still uses strict to decide whether auto tool choice activates a grammar.
-    parameter = TagFormat(
-        begin='<｜DSML｜ parameter name="',
-        content=SequenceFormat(
-            elements=[
-                RegexFormat(pattern=r'[^"]+'),
-                ConstStringFormat(value='" string="'),
-                OrFormat(
-                    elements=[
-                        SequenceFormat(
-                            elements=[
-                                ConstStringFormat(value='true">'),
-                                AnyTextFormat(
-                                    excludes=[
-                                        _V41_PARAMETER_END,
-                                        _V41_INVOKE_END,
-                                        _V41_CALLS_END,
-                                    ]
-                                ),
-                            ]
-                        ),
-                        SequenceFormat(
-                            elements=[
-                                ConstStringFormat(value='false">'),
-                                JSONSchemaFormat(json_schema=True),
-                            ]
-                        ),
-                    ]
-                ),
-            ]
-        ),
-        end=f"{_V41_PARAMETER_END}\n",
-    )
-    calls = TagsWithSeparatorFormat(
-        tags=[
-            TagFormat(
-                begin=f'<｜DSML｜ invoke name="{tool.function.name}">\n',
-                content=StarFormat(content=parameter),
-                end=f"{_V41_INVOKE_END}\n",
-            )
-            for tool in tools
-        ],
-        separator="",
-        at_least_one=True,
-        stop_after_first=tool_choice == "forced",
-    )
-    if tool_choice == "auto":
-        return StructuralTag(
-            format=TriggeredTagsFormat(
-                triggers=[_V41_CALLS_START],
-                tags=[
-                    TagFormat(
-                        begin=f"{_V41_CALLS_START}\n",
-                        content=calls,
-                        end=_V41_CALLS_END,
-                    )
-                ],
-                excludes=["<think>", "</think>"],
-            )
-            if tools
-            else AnyTextFormat(excludes=["<think>", "</think>"])
-        )
-    return StructuralTag(
-        format=SequenceFormat(
-            elements=[
-                ConstStringFormat(value=f"\n\n{_V41_CALLS_START}\n"),
-                calls,
-                ConstStringFormat(value=_V41_CALLS_END),
-            ]
-        )
-    )
-
-
 def _hermes_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]:
     arguments_field_prefix = '", "arguments": '
     formats = [
@@ -771,6 +719,7 @@ def get_hy_v4_structural_tag(
             leading colon (e.g. ``":6124c78e"``), or ``""`` when the checkpoint
             uses unsuffixed tokens. The HYV4 tool parser reads it off the
             tokenizer vocab and passes it to ``get_model_structural_tag``.
+
     """
     del builtin_tools
 
diff --git a/vllm/tool_parsers/tool_strict_level.py b/vllm/tool_parsers/tool_strict_level.py
new file mode 100644
index 000000000000..03a7d8b170a1
--- /dev/null
+++ b/vllm/tool_parsers/tool_strict_level.py
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import enum
+from typing import Literal
+
+ToolStrictLevelName = Literal["auto", "function", "parameter"]
+
+
+class ToolStrictLevel(enum.IntEnum):
+    """Server-side floor for tool-call structural tags (``--tool-strict-level``).
+
+    AUTO:      derive call constraints from ``tool_choice`` and argument-schema
+               enforcement from each tool's explicit ``strict`` field.
+    FUNCTION:  constrain the tool-call envelope for every request with tools.
+    PARAMETER: additionally pin argument schemas, as if every tool were
+               ``strict: true``.
+    """
+
+    AUTO = 0
+    FUNCTION = 1
+    PARAMETER = 2
+
+    @classmethod
+    def from_name(cls, name: str) -> "ToolStrictLevel":
+        try:
+            return cls[name.strip().upper()]
+        except KeyError:
+            expected = ", ".join(level.name.lower() for level in cls)
+            raise ValueError(
+                f"Unknown tool strict level {name!r}; expected one of {expected}."
+            ) from None
diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py
index d39c83969f41..d2c8c6f41220 100644
--- a/vllm/tool_parsers/utils.py
+++ b/vllm/tool_parsers/utils.py
@@ -56,8 +56,7 @@ def partial_tag_overlap(text: str, tag: str) -> int:
 
 
 def find_common_prefix(s1: str, s2: str) -> str:
-    """
-    Finds a common prefix that is shared between two strings, if there is one.
+    """Finds a common prefix that is shared between two strings, if there is one.
     Order of arguments is NOT important.
 
     This function is provided as a UTILITY for extracting information from JSON
@@ -79,8 +78,7 @@ def find_common_prefix(s1: str, s2: str) -> str:
 
 
 def find_common_suffix(s1: str, s2: str) -> str:
-    """
-    Finds a common suffix shared between two strings, if there is one. Order of
+    """Finds a common suffix shared between two strings, if there is one. Order of
     arguments is NOT important.
     Stops when the suffix ends OR it hits an alphanumeric character
 
@@ -97,8 +95,7 @@ def find_common_suffix(s1: str, s2: str) -> str:
 
 
 def extract_intermediate_diff(curr: str, old: str) -> str:
-    """
-    Given two strings, extract the difference in the middle between two strings
+    """Given two strings, extract the difference in the middle between two strings
     that are known to have a common prefix and/or suffix.
 
     This function is provided as a UTILITY for extracting information from JSON
@@ -450,6 +447,7 @@ def get_parameter_value(val: ast.expr) -> Any:
 
     Raises:
         UnexpectedAstError: If the AST node is not a supported literal type.
+
     """
     if isinstance(val, ast.Constant):
         if val.value is None or isinstance(val.value, (str, int, float)):
@@ -528,6 +526,7 @@ def _ast_callable_dotted_name(node: ast.expr) -> str:
     Raises:
         UnexpectedAstError: If the chain does not bottom out in an
             ``ast.Name`` (e.g. subscript or call expression as receiver).
+
     """
     parts: list[str] = []
     current: ast.expr = node
@@ -550,6 +549,7 @@ def handle_single_tool(call: ast.Call) -> ToolCall:
     Raises:
         UnexpectedAstError: If the call target is neither a simple name
             nor a chain of attribute accesses bottoming out in a name.
+
     """
     if not isinstance(call.func, (ast.Name, ast.Attribute)):
         logger.warning(
@@ -921,6 +921,7 @@ def make_valid_python(text: str) -> tuple[str, str] | None:
     Raises:
         UnexpectedAstError: If mismatched brackets or parentheses
             are detected.
+
     """
     bracket_stack: list[str] = []
     for index, char in enumerate(text):
@@ -1088,6 +1089,7 @@ def coerce_to_schema_type(value: str, schema_type: str | list[str]) -> Any:
         value: The raw string value from the model output.
         schema_type: One or more JSON Schema type strings
             (e.g. ``"string"`` or ``["string", "null"]``).
+
     """
     if isinstance(schema_type, str):
         schema_type = [schema_type]
@@ -1176,6 +1178,7 @@ def compute_tool_delta(
     Returns:
         A DeltaToolCall with only the new argument characters, or None
         if there is no difference from what was previously sent.
+
     """
     new_call_args = new_call.function.arguments
     if withheld_suffix:
diff --git a/vllm/tracing/__init__.py b/vllm/tracing/__init__.py
index 4f025970ecbd..9725b57ba979 100644
--- a/vllm/tracing/__init__.py
+++ b/vllm/tracing/__init__.py
@@ -94,9 +94,7 @@ def instrument(
     attributes: dict[str, str] | None = None,
     record_exception: bool = True,
 ):
-    """
-    Generic decorator to instrument functions.
-    """
+    """Generic decorator to instrument functions."""
     if obj is None:
         return functools.partial(
             instrument,
@@ -135,6 +133,7 @@ def instrument_manual(
         attributes: Optional dict of span attributes.
         context: Optional trace context (e.g., from extract_trace_context).
         kind: Optional SpanKind (e.g., SpanKind.SERVER).
+
     """
     is_available, _, _, _, manual_instrument_fn = _REGISTERED_TRACING_BACKENDS["otel"]
     if is_available():
@@ -146,8 +145,7 @@ def instrument_manual(
 
 
 def is_tracing_available() -> bool:
-    """
-    Returns True if any tracing backend (OTel, Profiler, etc.) is available.
+    """Returns True if any tracing backend (OTel, Profiler, etc.) is available.
     Use this to guard expensive tracing logic in the main code.
     """
     check_available = [
diff --git a/vllm/tracing/otel.py b/vllm/tracing/otel.py
index ac06ae97255f..36c189b88324 100644
--- a/vllm/tracing/otel.py
+++ b/vllm/tracing/otel.py
@@ -107,9 +107,7 @@ def init_otel_worker_tracer(
     process_kind: str,
     process_name: str,
 ) -> Tracer:
-    """
-    Backend-specific initialization for OpenTelemetry in a worker process.
-    """
+    """Backend-specific initialization for OpenTelemetry in a worker process."""
     # Initialize the tracer if an OTLP endpoint is configured.
     # The endpoint is propagated via environment variable from the main process.
     otlp_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT")
@@ -133,7 +131,6 @@ def extract_trace_context(headers: Mapping[str, str] | None) -> Context | None:
 
 def instrument_otel(func, span_name, attributes, record_exception):
     """Internal wrapper logic for sync and async functions."""
-
     # Pre-calculate static code attributes once (these don't change)
     code_attrs = {
         LoadingSpanAttributes.CODE_FUNCTION: func.__qualname__,
@@ -214,8 +211,7 @@ def manual_instrument_otel(
 
 
 def _get_smart_context() -> Context | None:
-    """
-    Determines the parent context.
+    """Determines the parent context.
     1. If a Span is already active in this process, use it.
     2. If not, extract from os.environ, handling the case-sensitivity mismatch.
     """
@@ -239,8 +235,7 @@ def _get_smart_context() -> Context | None:
 
 @contextmanager
 def propagate_trace_to_env():
-    """
-    Temporarily injects the current OTel context into os.environ.
+    """Temporarily injects the current OTel context into os.environ.
     This ensures that any subprocesses (like vLLM workers) spawned
     within this context inherit the correct traceparent.
     """
diff --git a/vllm/tracing/utils.py b/vllm/tracing/utils.py
index 0e11f850b442..138b801f5125 100644
--- a/vllm/tracing/utils.py
+++ b/vllm/tracing/utils.py
@@ -13,8 +13,7 @@
 
 
 class SpanAttributes:
-    """
-    Standard attributes for spans.
+    """Standard attributes for spans.
 
     These are largely based on OpenTelemetry Semantic Conventions but are defined
     here as constants so they can be used by any backend or logger.
@@ -60,8 +59,7 @@ def contains_trace_headers(headers: Mapping[str, str]) -> bool:
 
 
 def extract_trace_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
-    """
-    Extract only trace-related headers from a larger header dictionary.
+    """Extract only trace-related headers from a larger header dictionary.
     Useful for logging or passing context to a non-OTel client.
     """
     return {h: headers[h] for h in TRACE_HEADERS if h in headers}
diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py
index 93e3ad4de88d..b7f4f54d938b 100644
--- a/vllm/transformers_utils/config.py
+++ b/vllm/transformers_utils/config.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import math
 import os
 from collections.abc import Callable, Iterator, Sequence
 from contextlib import contextmanager
@@ -12,10 +13,11 @@
 
 import huggingface_hub
 import torch
+import transformers.configuration_utils as hf_configuration_utils
 from huggingface_hub import constants
 from packaging.version import Version
 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
-from transformers import GenerationConfig, PretrainedConfig
+from transformers import GenerationConfig, PreTrainedConfig
 from transformers.configuration_utils import ALLOWED_LAYER_TYPES
 from transformers.models.auto.image_processing_auto import get_image_processor_config
 from transformers.models.auto.modeling_auto import (
@@ -84,7 +86,7 @@ def __getitem__(self, key):
         return getattr(configs, value)
 
 
-_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
+_CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict(
     afmoe="AfmoeConfig",
     axk1="AXK1Config",
     bagel="BagelConfig",
@@ -125,7 +127,6 @@ def __getitem__(self, key):
     RefinedWebModel="RWConfig",  # For tiiuae/falcon-7b(-instruct)
     mlp_speculator="MLPSpeculatorConfig",
     medusa="MedusaConfig",
-    mellum="MellumConfig",
     midashenglm="MiDashengLMConfig",
     minimax_m3_vl="MiniMaxM3Config",
     minimax_m3_mtp="MiniMaxM3MTPConfig",
@@ -150,7 +151,6 @@ def __getitem__(self, key):
     qwen3_5_text="Qwen3_5TextConfig",
     qwen3_5_moe="Qwen3_5MoeConfig",
     qwen3_5_moe_text="Qwen3_5MoeTextConfig",
-    laguna="LagunaConfig",
     lfm2_moe="Lfm2MoeConfig",
     **{"unlimited-ocr": "UnlimitedOCRConfig"},
     **{"deepseek_v41": "DeepseekV41Config"},
@@ -162,6 +162,13 @@ def __getitem__(self, key):
 
 _PATCH_HF_VALIDATE_ROPE: set[str] = {"sarvam_mla"}
 
+# Model types whose checkpoints carry shared RoPE parameters alongside the
+# per-layer-type dicts. Since transformers 5.17, `validate_rope` treats every
+# top-level value of such a dict as a layer's parameters and raises on the
+# shared ones. `laguna` gets them injected by `convert_rope_params_to_dict`;
+# `gemma4_text` ships them in the checkpoint itself.
+_PATCH_HF_NESTED_ROPE_VALIDATION: set[str] = {"laguna", "gemma4_text"}
+
 # Model types whose checkpoints declare `layer_types` entries that upstream
 # transformers has not added to `ALLOWED_LAYER_TYPES` yet, so its strict config
 # validation rejects them (e.g.  GLM-5.2 `glm_moe_dsa` use
@@ -182,13 +189,13 @@ def __getitem__(self, key):
 
 
 def _register_config_class(
-    model_type: str, config_class: type[PretrainedConfig]
+    model_type: str, config_class: type[PreTrainedConfig]
 ) -> None:
     config_class.model_type = model_type
     AutoConfig.register(model_type, config_class, exist_ok=True)
 
 
-def _maybe_register_hf_config(config: PretrainedConfig | None) -> None:
+def _maybe_register_hf_config(config: PreTrainedConfig | None) -> None:
     if config is None:
         return
 
@@ -218,6 +225,35 @@ def _mistral_patch_hf_hub_constants() -> Iterator[None]:
         constants.SAFETENSORS_INDEX_FILE = hf_safetensors_index_file
 
 
+def _install_hf_config_validator(
+    name: str, validator: Callable, supersedes: Callable
+) -> None:
+    """Replace a ``PreTrainedConfig`` validator on every ``@strict`` snapshot.
+
+    ``@strict`` snapshots each ``validate_*`` method into ``__class_validators__``
+    at class creation and automatic post-``__init__`` validation dispatches off
+    that frozen list, so assigning the class attribute alone only affects
+    explicit ``config.validate_*()`` calls. Every ``@strict``-decorated config
+    class owns a snapshot, so rewrite them all, matching ``supersedes`` by
+    identity to leave a genuine per-model override in place.
+    """
+    setattr(PreTrainedConfig, name, validator)
+
+    seen: set[int] = set()
+    stack = [PreTrainedConfig]
+    while stack:
+        cls = stack.pop()
+        if id(cls) in seen:
+            continue
+        seen.add(id(cls))
+        validators = cls.__dict__.get("__class_validators__")
+        if isinstance(validators, list) and any(v is supersedes for v in validators):
+            cls.__class_validators__ = [
+                validator if v is supersedes else v for v in validators
+            ]
+        stack.extend(cls.__subclasses__())
+
+
 def _patch_hf_transformers_validate_rope():
     """Transformers v5 moved the ignore_keys option from the method signature of
     validate_rope and replaced it with the ignore_keys_at_rope_validation parameter
@@ -225,22 +261,64 @@ def _patch_hf_transformers_validate_rope():
     validate_rope() with the ignore_keys parameter work with newer versions of
     hf transformers (from v5 onwards)
     """
-
-    if hasattr(PretrainedConfig.validate_rope, "__vllm_patched__"):
+    if hasattr(PreTrainedConfig.validate_rope, "__vllm_patched__"):
         return
 
-    _original_validate_rope = PretrainedConfig.validate_rope
+    _original_validate_rope = PreTrainedConfig.validate_rope
 
     @wraps(_original_validate_rope)
     def patched_validate_rope(self, *args, **kwargs):
         ignore_keys_param = kwargs.pop("ignore_keys", None)
         original_ignore_keys = self.ignore_keys_at_rope_validation
+        if ignore_keys_param is not None:
+            logger.warning(
+                "validate_rope() was called with the legacy 'ignore_keys' "
+                "argument; use 'ignore_keys_at_rope_validation' instead"
+            )
         self.ignore_keys_at_rope_validation = original_ignore_keys or ignore_keys_param
         result = _original_validate_rope(self, *args, **kwargs)
         return result
 
     patched_validate_rope.__vllm_patched__ = True  # type: ignore[attr-defined]
-    PretrainedConfig.validate_rope = patched_validate_rope
+    _install_hf_config_validator(
+        "validate_rope", patched_validate_rope, _original_validate_rope
+    )
+
+
+def _patch_hf_transformers_nested_rope_validation() -> None:
+    """Drop shared entries sitting alongside a nested ``rope_parameters`` dict.
+
+    Transformers validates a dict with any layer-type key by iterating all of
+    its values, so a shared entry next to the per-layer dicts raises an
+    ``AttributeError``. The per-layer dicts already carry their own defaults by
+    the time validation runs, so the shared entries can be dropped.
+    """
+    if hasattr(PreTrainedConfig.validate_rope, "__vllm_nested_rope_patched__"):
+        return
+
+    _original_validate_rope = PreTrainedConfig.validate_rope
+
+    @wraps(_original_validate_rope)
+    def patched_validate_rope(self, *args, **kwargs):
+        rope_parameters = getattr(self, "rope_parameters", None)
+        if isinstance(rope_parameters, dict):
+            layer_types = set(rope_parameters) & set(
+                hf_configuration_utils.ALLOWED_LAYER_TYPES
+            )
+            if layer_types and (shared_keys := set(rope_parameters) - layer_types):
+                for key in shared_keys:
+                    del rope_parameters[key]
+                logger.warning(
+                    "Dropping %s from rope_parameters; entries shared across "
+                    "layer types are not supported",
+                    sorted(shared_keys),
+                )
+        return _original_validate_rope(self, *args, **kwargs)
+
+    patched_validate_rope.__vllm_nested_rope_patched__ = True  # type: ignore[attr-defined]
+    _install_hf_config_validator(
+        "validate_rope", patched_validate_rope, _original_validate_rope
+    )
 
 
 def _patch_hf_transformers_allowed_layer_types(
@@ -269,11 +347,11 @@ def parse(
         revision: str | None = None,
         code_revision: str | None = None,
         **kwargs,
-    ) -> tuple[dict, PretrainedConfig]:
+    ) -> tuple[dict, PreTrainedConfig]:
         kwargs["local_files_only"] = huggingface_hub.constants.HF_HUB_OFFLINE
         trust_remote_code |= kwargs.get("trust_remote_code", False)
         kwargs = without_trust_remote_code(kwargs)
-        config_dict, _ = PretrainedConfig.get_config_dict(
+        config_dict, _ = PreTrainedConfig.get_config_dict(
             model,
             revision=revision,
             code_revision=code_revision,
@@ -296,13 +374,16 @@ def parse(
                 # through and remain unchanged by this elif block
                 dummy_model_type = f"dummy_{model_type}"
                 dummy_kwargs = dict(architectures=[""], model_type=dummy_model_type)
-                dummy_config = PretrainedConfig(**dummy_kwargs)
+                dummy_config = PreTrainedConfig(**dummy_kwargs)
                 dummy_model_type = hf_overrides(dummy_config).model_type
                 model_type = dummy_model_type.removeprefix("dummy_")
 
         if model_type in _PATCH_HF_VALIDATE_ROPE:
             _patch_hf_transformers_validate_rope()
 
+        if model_type in _PATCH_HF_NESTED_ROPE_VALIDATION:
+            _patch_hf_transformers_nested_rope_validation()
+
         if extra_layer_types := _PATCH_HF_ALLOWED_LAYER_TYPES.get(model_type):
             _patch_hf_transformers_allowed_layer_types(extra_layer_types)
 
@@ -369,7 +450,7 @@ def parse(
         revision: str | None = None,
         code_revision: str | None = None,
         **kwargs,
-    ) -> tuple[dict, PretrainedConfig]:
+    ) -> tuple[dict, PreTrainedConfig]:
         # This function loads a params.json config which
         # should be used when loading models in mistral format
         config_dict = _download_mistral_config_file(model, revision)
@@ -385,7 +466,7 @@ def parse(
 
         # Get missing fields from HF config if available
         try:
-            hf_config_dict, _ = PretrainedConfig.get_config_dict(
+            hf_config_dict, _ = PreTrainedConfig.get_config_dict(
                 model,
                 revision=revision,
                 code_revision=code_revision,
@@ -441,10 +522,11 @@ def register_config_parser(config_format: str):
     """Register a customized vllm config parser.
      When a config format is not supported by vllm, you can register a customized
     config parser to support it.
-     Args:
+
+    Args:
          config_format (str): The config parser format name.
-     Examples:
 
+    Examples:
          >>> from vllm.transformers_utils.config import (get_config_parser,
                                                          register_config_parser)
          >>> from vllm.transformers_utils.config_parser_base import ConfigParserBase
@@ -458,11 +540,12 @@ def register_config_parser(config_format: str):
          ...         revision: str | None = None,
          ...         code_revision: str | None = None,
          ...         **kwargs,
-         ...     ) -> tuple[dict, PretrainedConfig]:
+         ...     ) -> tuple[dict, PreTrainedConfig]:
          ...         raise NotImplementedError
          >>>
          >>> type(get_config_parser("custom_config_parser"))
          <class 'CustomConfigParser'>
+
     """  # noqa: E501
 
     def _wrapper(config_parser_cls):
@@ -488,7 +571,7 @@ def _wrapper(config_parser_cls):
     return _wrapper
 
 
-def set_default_rope_theta(config: PretrainedConfig, default_theta: float) -> None:
+def set_default_rope_theta(config: PreTrainedConfig, default_theta: float) -> None:
     """Some models may have no rope_theta in their config but still use RoPE.
     This function sets a default rope_theta if it's missing."""
     if getattr(config, "rope_parameters", None) is None:
@@ -500,7 +583,6 @@ def set_default_rope_theta(config: PretrainedConfig, default_theta: float) -> No
 def patch_legacy_rope_type(rope_parameters: dict[str, Any] | None) -> None:
     """Patch legacy RoPE type fields for backwards compatibility with
     older custom models which would otherwise fail to load."""
-
     # No RoPE parameters to patch
     if rope_parameters is None:
         return
@@ -539,6 +621,14 @@ def _patch_legacy_rope_type(rope_parameters: dict[str, Any]) -> None:
                 )
             rope_parameters["rope_type"] = "default"
             logger.warning("Replacing legacy rope_type 'mrope' with 'default'")
+        elif rope_parameters["rope_type"] == "telechat3-yarn":
+            # TeleChat3 is YaRN with 0.07 in place of YaRN's 0.1 attention
+            # scaling coefficient. Precompute it so the config is plain YaRN.
+            factor = rope_parameters["factor"]
+            rope_parameters["rope_type"] = "yarn"
+            rope_parameters["attention_factor"] = 0.07 * math.log(factor) + 1.0
+            rope_parameters.pop("type", None)
+            logger.warning("Replacing rope_type 'telechat3-yarn' with 'yarn'")
 
     # Handle nested rope_parameters in interleaved sliding attention models
     if is_rope_parameters_nested(rope_parameters):
@@ -549,7 +639,7 @@ def _patch_legacy_rope_type(rope_parameters: dict[str, Any]) -> None:
         _patch_legacy_rope_type(rope_parameters)
 
 
-def patch_rope_parameters(config: PretrainedConfig) -> None:
+def patch_rope_parameters(config: PreTrainedConfig) -> None:
     """Provide backwards compatibility for RoPE."""
     from vllm.config.utils import getattr_iter
 
@@ -571,7 +661,7 @@ def patch_rope_parameters(config: PretrainedConfig) -> None:
         config.validate_rope()
 
 
-def _iter_rope_parameters(config: PretrainedConfig) -> Iterator[dict[str, Any]]:
+def _iter_rope_parameters(config: PreTrainedConfig) -> Iterator[dict[str, Any]]:
     """Yield a config's rope parameters, one dict per layer type if nested."""
     rope_parameters = getattr(config, "rope_parameters", None)
     if not isinstance(rope_parameters, dict):
@@ -583,7 +673,7 @@ def _iter_rope_parameters(config: PretrainedConfig) -> Iterator[dict[str, Any]]:
         yield rope_parameters
 
 
-def _mrope_section(config: PretrainedConfig) -> Sequence[int] | None:
+def _mrope_section(config: PreTrainedConfig) -> Sequence[int] | None:
     """Return the M-RoPE section this config declares, if any.
 
     `xdrope_section` is the legacy name HunYuan-VL checkpoints use for the same
@@ -610,11 +700,11 @@ def _mrope_section(config: PretrainedConfig) -> Sequence[int] | None:
     return section if isinstance(section, (list, tuple)) else None
 
 
-def _uses_mrope(config: PretrainedConfig) -> bool:
+def _uses_mrope(config: PreTrainedConfig) -> bool:
     return _mrope_section(config) is not None
 
 
-def uses_mrope(config: PretrainedConfig) -> bool:
+def uses_mrope(config: PreTrainedConfig) -> bool:
     """Detect if the model with this config uses M-ROPE."""
     return (
         _uses_mrope(config)
@@ -623,7 +713,7 @@ def uses_mrope(config: PretrainedConfig) -> bool:
     )
 
 
-def thinker_uses_mrope(config: PretrainedConfig) -> bool:
+def thinker_uses_mrope(config: PreTrainedConfig) -> bool:
     """Detect if the model contains a thinker config and it uses M-ROPE."""
     thinker_config = getattr(config, "thinker_config", None)
     if thinker_config is None:
@@ -636,7 +726,7 @@ def thinker_uses_mrope(config: PretrainedConfig) -> bool:
     return uses_mrope(thinker_text_config)
 
 
-def mrope_num_dims(config: PretrainedConfig) -> int:
+def mrope_num_dims(config: PreTrainedConfig) -> int:
     """Number of M-RoPE position channels the model consumes.
 
     Each section entry sizes one channel, so the section length is the channel
@@ -655,25 +745,23 @@ def mrope_num_dims(config: PretrainedConfig) -> int:
     return 3
 
 
-def is_encoder_decoder(config: PretrainedConfig) -> bool:
+def is_encoder_decoder(config: PreTrainedConfig) -> bool:
     """Detect if the model with this config is used as an encoder/decoder."""
 
-    def _is_encoder_decoder(config: PretrainedConfig) -> bool:
+    def _is_encoder_decoder(config: PreTrainedConfig) -> bool:
         return getattr(config, "is_encoder_decoder", False)
 
     return _is_encoder_decoder(config) or _is_encoder_decoder(config.get_text_config())
 
 
 def _maybe_update_auto_config_kwargs(kwargs: dict[str, Any], model_type: str):
-    """
-    Update kwargs for AutoConfig initialization based on model_type
-    """
+    """Update kwargs for AutoConfig initialization based on model_type."""
     if model_type in _AUTO_CONFIG_KWARGS_OVERRIDES:
         kwargs.update(_AUTO_CONFIG_KWARGS_OVERRIDES[model_type])
     return kwargs
 
 
-def _maybe_remap_hf_config_attrs(config: PretrainedConfig) -> PretrainedConfig:
+def _maybe_remap_hf_config_attrs(config: PreTrainedConfig) -> PreTrainedConfig:
     """Remap config attributes to match the expected names."""
     for old_attr, new_attr in _CONFIG_ATTRS_MAPPING.items():
         if hasattr(config, old_attr):
@@ -692,8 +780,7 @@ def maybe_override_with_speculators(
     hf_token: bool | str | None = None,
     **kwargs,
 ) -> tuple[str, str | None, dict[str, Any] | None]:
-    """
-    Resolve model configuration when speculators are detected.
+    """Resolve model configuration when speculators are detected.
 
     Checks if the provided model is a speculators model and if so, extracts
     the target model configuration and builds the speculative config.
@@ -708,9 +795,10 @@ def maybe_override_with_speculators(
 
     Returns:
         Tuple of (resolved_model, resolved_tokenizer, speculative_config)
+
     """
     kwargs["local_files_only"] = huggingface_hub.constants.HF_HUB_OFFLINE
-    config_dict, _ = PretrainedConfig.get_config_dict(
+    config_dict, _ = PreTrainedConfig.get_config_dict(
         model,
         revision=revision,
         token=hf_token,
@@ -752,9 +840,9 @@ def get_config(
     code_revision: str | None = None,
     config_format: str | ConfigFormat = "auto",
     hf_overrides_kw: dict[str, Any] | None = None,
-    hf_overrides_fn: Callable[[PretrainedConfig], PretrainedConfig] | None = None,
+    hf_overrides_fn: Callable[[PreTrainedConfig], PreTrainedConfig] | None = None,
     **kwargs,
-) -> PretrainedConfig:
+) -> PreTrainedConfig:
     if config_format == "auto":
         try:
             # First check for Mistral to avoid defaulting to
@@ -865,7 +953,7 @@ def get_config(
     # Exhaustively patch RoPE parameters everywhere they might be
     patch_rope_parameters(config)
     patch_rope_parameters(config.get_text_config())
-    SubConfigs: TypeAlias = dict[str, PretrainedConfig]
+    SubConfigs: TypeAlias = dict[str, PreTrainedConfig]
     sub_configs: SubConfigs | None = getattr(config, "sub_configs", None)
     if sub_configs:
         for sub_config in sub_configs:
@@ -882,8 +970,7 @@ def get_pooling_config(
     model: str,
     revision: str | None = "main",
 ) -> dict[str, Any] | None:
-    """
-    This function gets the pooling and normalize
+    """This function gets the pooling and normalize
     config from the model - only applies to
     sentence-transformers models.
 
@@ -895,6 +982,7 @@ def get_pooling_config(
     Returns:
         A dictionary containing the pooling type and whether
             normalization is used, or None if no pooling configuration is found.
+
     """
     modules_file_name = "modules.json"
 
@@ -979,19 +1067,17 @@ def parse_pooling_type(pooling_name: str):
 def get_sentence_transformer_tokenizer_config(
     model: str | Path, revision: str | None = "main"
 ) -> dict[str, Any] | None:
-    """
-    Returns the tokenization configuration dictionary for a
+    """Returns the tokenization configuration dictionary for a
     given Sentence Transformer BERT model.
 
-    Parameters:
-    - model (str|Path): The name of the Sentence Transformer
-    BERT model.
-    - revision (str, optional): The revision of the m
-    odel to use. Defaults to 'main'.
+    Args:
+        model (str|Path): The name of the Sentence Transformer BERT model.
+        revision (str, optional): The revision of the model to use. Defaults to 'main'.
 
     Returns:
     - dict: A dictionary containing the configuration parameters
     for the Sentence Transformer BERT model.
+
     """
     sentence_transformer_config_files = [
         "sentence_bert_config.json",
@@ -1126,7 +1212,7 @@ def get_hf_image_processor_config(
     )
 
 
-def get_hf_text_config(config: PretrainedConfig):
+def get_hf_text_config(config: PreTrainedConfig):
     """Get the "sub" config relevant to llm for multi modal models.
     No op for pure text models.
     """
@@ -1252,9 +1338,7 @@ def get_safetensors_params_metadata(
     *,
     revision: str | None = None,
 ) -> dict[str, Any]:
-    """
-    Get the safetensors parameters metadata for remote/local model repository.
-    """
+    """Get the safetensors parameters metadata for remote/local model repository."""
     if (model_path := Path(model)).exists():
         return _read_safetensors_metadata_in_dir(model_path)
 
@@ -1298,6 +1382,7 @@ def checkpoint_has_lm_head(model: str, *, revision: str | None = None) -> bool |
     Returns:
         `None` if the checkpoint contents could not be determined, for example
         because it is not stored as safetensors.
+
     """
     metadata = get_safetensors_params_metadata(model, revision=revision)
     if not metadata:
diff --git a/vllm/transformers_utils/config_parser_base.py b/vllm/transformers_utils/config_parser_base.py
index 79d47ff56042..46cbab5be9ee 100644
--- a/vllm/transformers_utils/config_parser_base.py
+++ b/vllm/transformers_utils/config_parser_base.py
@@ -4,7 +4,7 @@
 from abc import ABC, abstractmethod
 from pathlib import Path
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
 class ConfigParserBase(ABC):
@@ -16,5 +16,5 @@ def parse(
         revision: str | None = None,
         code_revision: str | None = None,
         **kwargs,
-    ) -> tuple[dict, PretrainedConfig]:
+    ) -> tuple[dict, PreTrainedConfig]:
         raise NotImplementedError
diff --git a/vllm/transformers_utils/configs/AXK1.py b/vllm/transformers_utils/configs/AXK1.py
index 09a3a0f46ae6..50bc891573c8 100644
--- a/vllm/transformers_utils/configs/AXK1.py
+++ b/vllm/transformers_utils/configs/AXK1.py
@@ -2,18 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
-class AXK1Config(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`AXK1Model`].
+class AXK1Config(PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a [`AXK1Model`].
     It is used to instantiate an A.X model according to the specified arguments,
     defining the model architecture. Instantiating a configuration with the defaults
     will yield a similar configuration to that of the A.X K1.
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control
-    the model outputs. Read the documentation from [`PretrainedConfig`] for more
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control
+    the model outputs. Read the documentation from [`PreTrainedConfig`] for more
     information.
+
     Args:
         vocab_size (`int`, *optional*, defaults to 163840):
             Vocabulary size of the A.X K1 model. Defines the number of different
@@ -112,6 +112,7 @@ class AXK1Config(PretrainedConfig):
             layers during self-attention.
         attention_dropout (`float`, *optional*, defaults to 0.0):
             The dropout ratio for the attention probabilities.
+
     """
 
     model_type = "axk1"
diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py
index 971b6ffcb2b5..a302d5d72928 100644
--- a/vllm/transformers_utils/configs/__init__.py
+++ b/vllm/transformers_utils/configs/__init__.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Model configs may be defined in this directory for the following reasons:
+"""Model configs may be defined in this directory for the following reasons:
 
 - There is no configuration file defined by HF Hub or Transformers library.
 - There is a need to override the existing config to support vLLM.
@@ -55,10 +54,8 @@
     # tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the
     # `FalconConfig` class from the official HuggingFace transformers library.
     "RWConfig": "vllm.transformers_utils.configs.falcon",
-    "LagunaConfig": "vllm.transformers_utils.configs.laguna",
     "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe",
     "MedusaConfig": "vllm.transformers_utils.configs.medusa",
-    "MellumConfig": "vllm.transformers_utils.configs.mellum",
     "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
     "MiniMaxM3Config": "vllm.transformers_utils.configs.minimax_m3",
     "MiniMaxM3MTPConfig": "vllm.transformers_utils.configs.minimax_m3",
@@ -151,10 +148,8 @@
     "HyperCLOVAXConfig",
     "IsaacConfig",
     "RWConfig",
-    "LagunaConfig",
     "Lfm2MoeConfig",
     "MedusaConfig",
-    "MellumConfig",
     "MiDashengLMConfig",
     "MiniMaxM3Config",
     "MiniMaxM3MTPConfig",
diff --git a/vllm/transformers_utils/configs/afmoe.py b/vllm/transformers_utils/configs/afmoe.py
index 47fee9882f9f..b3d10dc5f274 100644
--- a/vllm/transformers_utils/configs/afmoe.py
+++ b/vllm/transformers_utils/configs/afmoe.py
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class AfmoeConfig(PretrainedConfig):
+class AfmoeConfig(PreTrainedConfig):
     model_type = "afmoe"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/bagel.py b/vllm/transformers_utils/configs/bagel.py
index 53347ef45213..a972000db4a5 100644
--- a/vllm/transformers_utils/configs/bagel.py
+++ b/vllm/transformers_utils/configs/bagel.py
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from transformers import PretrainedConfig, SiglipVisionConfig
+from transformers import PreTrainedConfig, SiglipVisionConfig
 from transformers.models.qwen2 import Qwen2Config
 
 
-class BagelConfig(PretrainedConfig):
+class BagelConfig(PreTrainedConfig):
     """Configuration class for BAGEL model."""
 
     model_type = "bagel"
diff --git a/vllm/transformers_utils/configs/bailing_moe_v3_vl.py b/vllm/transformers_utils/configs/bailing_moe_v3_vl.py
index 484e27991ef9..e009a4db14d8 100644
--- a/vllm/transformers_utils/configs/bailing_moe_v3_vl.py
+++ b/vllm/transformers_utils/configs/bailing_moe_v3_vl.py
@@ -9,11 +9,11 @@
 
 from typing import Any
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
 def _normalize_mrope_parameters(
-    config: PretrainedConfig,
+    config: PreTrainedConfig,
     mrope_section: list[int],
 ) -> None:
     """Add the top-level Bailing M-RoPE split to a text config."""
@@ -50,7 +50,7 @@ def _build_layer_types(
     ]
 
 
-class BailingMoeV3TextConfig(PretrainedConfig):
+class BailingMoeV3TextConfig(PreTrainedConfig):
     model_type = "bailing_hybrid"
     base_config_key = "text_config"
     keys_to_ignore_at_inference = ["past_key_values"]
@@ -230,7 +230,7 @@ def __init__(
         self.validate_layer_type()
 
 
-class BailingMoeV3VisionConfig(PretrainedConfig):
+class BailingMoeV3VisionConfig(PreTrainedConfig):
     model_type = "qwen3_moe_vit"
     base_config_key = "vision_config"
 
@@ -273,7 +273,7 @@ def __init__(
         self.initializer_range = initializer_range
 
 
-class BailingMoeV3VLConfig(PretrainedConfig):
+class BailingMoeV3VLConfig(PreTrainedConfig):
     model_type = "bailing_moe_v3_vl"
     keys_to_ignore_at_inference = ["past_key_values"]
     sub_configs = {
diff --git a/vllm/transformers_utils/configs/chatglm.py b/vllm/transformers_utils/configs/chatglm.py
index 1d795b55c8bc..3908df740a47 100644
--- a/vllm/transformers_utils/configs/chatglm.py
+++ b/vllm/transformers_utils/configs/chatglm.py
@@ -3,10 +3,10 @@
 
 # Adapted from
 # https://github.com/zai-org/ChatGLM2-6B
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
-class ChatGLMConfig(PretrainedConfig):
+class ChatGLMConfig(PreTrainedConfig):
     model_type = "chatglm"
     attribute_map = {
         "num_hidden_layers": "num_layers",
diff --git a/vllm/transformers_utils/configs/colmodernvbert.py b/vllm/transformers_utils/configs/colmodernvbert.py
index 17558aafcec7..a8305417571c 100644
--- a/vllm/transformers_utils/configs/colmodernvbert.py
+++ b/vllm/transformers_utils/configs/colmodernvbert.py
@@ -8,10 +8,10 @@
 Reference: https://huggingface.co/ModernVBERT/colmodernvbert-merged
 """
 
-from transformers import ModernBertConfig, PretrainedConfig, SiglipVisionConfig
+from transformers import ModernBertConfig, PreTrainedConfig, SiglipVisionConfig
 
 
-class ColModernVBertConfig(PretrainedConfig):
+class ColModernVBertConfig(PreTrainedConfig):
     model_type = "colmodernvbert"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/colpali.py b/vllm/transformers_utils/configs/colpali.py
index c40c58b25ce1..a4e0a83e2929 100644
--- a/vllm/transformers_utils/configs/colpali.py
+++ b/vllm/transformers_utils/configs/colpali.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ColPali configuration that extends PaliGemmaConfig with embedding projection
+"""ColPali configuration that extends PaliGemmaConfig with embedding projection
 fields. This allows ColPali models to be loaded without trust_remote_code
 by mapping their custom model_type (colpali) to a standard config class
 that vLLM understands.
diff --git a/vllm/transformers_utils/configs/colqwen3.py b/vllm/transformers_utils/configs/colqwen3.py
index 1c09a0a91845..c98995491d8c 100644
--- a/vllm/transformers_utils/configs/colqwen3.py
+++ b/vllm/transformers_utils/configs/colqwen3.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ColQwen3 configuration that extends Qwen3VLConfig with embedding projection
+"""ColQwen3 configuration that extends Qwen3VLConfig with embedding projection
 fields. This allows ColQwen3 models to be loaded without trust_remote_code
 by mapping their custom model_type (colqwen3, ops_colqwen3, etc.) to a
 standard config class that vLLM understands.
diff --git a/vllm/transformers_utils/configs/cosmos3_edge.py b/vllm/transformers_utils/configs/cosmos3_edge.py
index 743e7d5debe7..447b64cc99ed 100644
--- a/vllm/transformers_utils/configs/cosmos3_edge.py
+++ b/vllm/transformers_utils/configs/cosmos3_edge.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig
 
 from .nemotron_h import NemotronHConfig
@@ -50,7 +50,7 @@ def to_dict(self) -> dict:
         return config_dict
 
 
-class Cosmos3EdgeProjectorConfig(PretrainedConfig):
+class Cosmos3EdgeProjectorConfig(PreTrainedConfig):
     model_type = "cosmos3_edge_projector"
 
     def __init__(
@@ -74,7 +74,7 @@ class Cosmos3EdgeVisionConfig(Siglip2VisionConfig):
     model_type = "cosmos3_edge_vision"
 
 
-class Cosmos3EdgeConfig(PretrainedConfig):
+class Cosmos3EdgeConfig(PreTrainedConfig):
     model_type = "cosmos3_edge"
     keys_to_ignore_at_inference = ["past_key_values"]
     has_no_defaults_at_init = True
diff --git a/vllm/transformers_utils/configs/deepseek_v4.py b/vllm/transformers_utils/configs/deepseek_v4.py
index 15346e35e589..278f47786466 100755
--- a/vllm/transformers_utils/configs/deepseek_v4.py
+++ b/vllm/transformers_utils/configs/deepseek_v4.py
@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
-class DeepseekV4Config(PretrainedConfig):
+class DeepseekV4Config(PreTrainedConfig):
     model_type = "deepseek_v4"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/deepseek_v41.py b/vllm/transformers_utils/configs/deepseek_v41.py
index 3a406d24bdc5..f1ed6f9aa895 100644
--- a/vllm/transformers_utils/configs/deepseek_v41.py
+++ b/vllm/transformers_utils/configs/deepseek_v41.py
@@ -2,17 +2,17 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
-class DeepseekV41Config(PretrainedConfig):
+class DeepseekV41Config(PreTrainedConfig):
     """DeepSeek V4.1 config.
 
     The HF config nests the text model under ``text_config`` and the vision
     tower under ``vision_config``. The vLLM model code (ported from
     ``deepseek_v4``) reads flat attributes, so both sub-configs are flattened
     onto the top level here: text fields are exposed as-is, vision fields
-    with the ``vision_*`` naming used by ``deepseek_v4_1.common.vision``.
+    with the ``vision_*`` naming used by ``deepseek_v41.common.vision``.
     """
 
     model_type = "deepseek_v41"
@@ -34,7 +34,7 @@ def __init__(
         for key, value in text_config.items():
             if key == "model_type":
                 continue
-            # Don't clobber PretrainedConfig properties (e.g. is_encoder_decoder).
+            # Don't clobber PreTrainedConfig properties (e.g. is_encoder_decoder).
             if isinstance(getattr(type(self), key, None), property):
                 continue
             setattr(self, key, value)
@@ -60,17 +60,3 @@ def __init__(
         self.vision_max_n_token = vision_config.get("max_image_tokens", 1024)
         self.vision_min_pixels = vision_config.get("min_pixels", 295936)
         self.vision_max_wh_ratio = vision_config.get("max_wh_ratio")
-        # The vision variant needs the mm-prefix plumbing (atomic image-span
-        # prefill + sparse-SWA window widening); the base arch-config
-        # convertor reads this config-level attribute.
-        self.is_mm_prefix_lm = vision_n_layers > 0
-        # The sparse-SWA index kernels widen the window within image spans
-        # in-kernel, so mm-prefix ranges longer than sliding_window must be
-        # kept (they are only consumed by those kernels).
-        self.mm_prefix_clamp_sliding_window = vision_n_layers > 0
-        # The visibility span covers the image block [IMAGE_START,
-        # IMAGE_END]; the mm placeholder additionally carries a leading
-        # compressor-alignment pad of ``COMPRESS_PAD_TO - 1 - offset %
-        # COMPRESS_PAD_TO`` tokens (see common/mm_preprocess.py), with
-        # COMPRESS_PAD_TO = 2 for v4.1's ratio-2 compressors.
-        self.mm_prefix_span_leading_pad_modulus = 2 if vision_n_layers > 0 else 0
diff --git a/vllm/transformers_utils/configs/deepseek_vl2.py b/vllm/transformers_utils/configs/deepseek_vl2.py
index 3b980a824d01..3543db833a4d 100644
--- a/vllm/transformers_utils/configs/deepseek_vl2.py
+++ b/vllm/transformers_utils/configs/deepseek_vl2.py
@@ -4,10 +4,10 @@
 # adapted from https://github.com/deepseek-ai/DeepSeek-VL2/blob/faf18023f24b962b32d9f0a2d89e402a8d383a78/deepseek_vl2/models/modeling_deepseek_vl_v2.py#L115-L268
 
 from huggingface_hub.dataclasses import strict
-from transformers import DeepseekV2Config, PretrainedConfig
+from transformers import DeepseekV2Config, PreTrainedConfig
 
 
-class VisionEncoderConfig(PretrainedConfig):
+class VisionEncoderConfig(PreTrainedConfig):
     model_type: str = "vision"
 
     model_name: str = "vit_so400m_patch14_siglip_384.webli"
@@ -58,7 +58,7 @@ def __init__(
         super().__init__(**kwargs)
 
 
-class MlpProjectorConfig(PretrainedConfig):
+class MlpProjectorConfig(PreTrainedConfig):
     model_type = "mlp_projector"
     projector_type: str = "downsample_mlp_gelu"
     input_dim: int = 1152
@@ -98,7 +98,7 @@ class DeepseekVLV2TextConfig(DeepseekV2Config):
     kv_lora_rank: int | None = DeepseekV2Config.kv_lora_rank
 
 
-class DeepseekVLV2Config(PretrainedConfig):
+class DeepseekVLV2Config(PreTrainedConfig):
     model_type = "deepseek_vl_v2"
 
     tile_tag: str = "2D"
diff --git a/vllm/transformers_utils/configs/diffusion_gemma.py b/vllm/transformers_utils/configs/diffusion_gemma.py
index 246a25b32c6d..884d176cd0b4 100644
--- a/vllm/transformers_utils/configs/diffusion_gemma.py
+++ b/vllm/transformers_utils/configs/diffusion_gemma.py
@@ -2,12 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig
 
 
-def _init_text_config(self: PretrainedConfig, **kwargs: Any) -> None:
-    PretrainedConfig.__init__(self, **kwargs)
+def _init_text_config(self: PreTrainedConfig, **kwargs: Any) -> None:
+    PreTrainedConfig.__init__(self, **kwargs)
     # DiffusionGemma always uses MoE and K=V sharing for full_attention
     # layers. The HF reference removed these config fields entirely.
     if getattr(self, "num_experts", None):
@@ -15,14 +15,14 @@ def _init_text_config(self: PretrainedConfig, **kwargs: Any) -> None:
     self.attention_k_eq_v = True
 
 
-class DiffusionGemmaTextConfig(PretrainedConfig):
+class DiffusionGemmaTextConfig(PreTrainedConfig):
     model_type = "diffusion_gemma_text"
 
     def __init__(self, **kwargs: Any):
         _init_text_config(self, **kwargs)
 
 
-class DiffusionGemmaConfig(PretrainedConfig):
+class DiffusionGemmaConfig(PreTrainedConfig):
     model_type = "diffusion_gemma"
 
     def __init__(
@@ -41,4 +41,4 @@ def __init__(
         else:
             self.vision_config = vision_config
         self.audio_config = None
-        PretrainedConfig.__init__(self, **kwargs)
+        PreTrainedConfig.__init__(self, **kwargs)
diff --git a/vllm/transformers_utils/configs/dotsocr.py b/vllm/transformers_utils/configs/dotsocr.py
index 1e42cb2fd859..918179dcb8c6 100644
--- a/vllm/transformers_utils/configs/dotsocr.py
+++ b/vllm/transformers_utils/configs/dotsocr.py
@@ -2,11 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.models.qwen2 import Qwen2Config
 
 
-class DotsVisionConfig(PretrainedConfig):
+class DotsVisionConfig(PreTrainedConfig):
     model_type: str = "dots_vit"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/eagle.py b/vllm/transformers_utils/configs/eagle.py
index 682828c0a594..6b218cd0f3f8 100644
--- a/vllm/transformers_utils/configs/eagle.py
+++ b/vllm/transformers_utils/configs/eagle.py
@@ -3,22 +3,22 @@
 
 import os
 
-from transformers import AutoConfig, DeepseekV2Config, PretrainedConfig
+from transformers import AutoConfig, DeepseekV2Config, PreTrainedConfig
 
 from vllm.transformers_utils.utils import without_trust_remote_code
 
 
-class EAGLEConfig(PretrainedConfig):
+class EAGLEConfig(PreTrainedConfig):
     model_type = "eagle"
 
     def __init__(
         self,
-        model: PretrainedConfig | dict | None = None,
+        model: PreTrainedConfig | dict | None = None,
         truncated_vocab_size: int | None = None,
         method: str | None = "eagle",
         **kwargs,
     ):
-        model_config: PretrainedConfig | DeepseekV2Config | None
+        model_config: PreTrainedConfig | DeepseekV2Config | None
         if isinstance(model, dict):
             model_config = AutoConfig.for_model(**model)
         else:
diff --git a/vllm/transformers_utils/configs/extract_hidden_states.py b/vllm/transformers_utils/configs/extract_hidden_states.py
index 2beec0e3081b..0b984c772208 100644
--- a/vllm/transformers_utils/configs/extract_hidden_states.py
+++ b/vllm/transformers_utils/configs/extract_hidden_states.py
@@ -5,17 +5,17 @@
 
 import os
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.transformers_utils.utils import without_trust_remote_code
 
 
-class ExtractHiddenStatesConfig(PretrainedConfig):
+class ExtractHiddenStatesConfig(PreTrainedConfig):
     model_type = "extract_hidden_states"
 
     def __init__(
         self,
-        model: PretrainedConfig | dict | None = None,
+        model: PreTrainedConfig | dict | None = None,
         method: str | None = "extract_hidden_states",
         **kwargs,
     ):
@@ -24,7 +24,7 @@ def __init__(
         if isinstance(model, dict):
             model_dict = model
             source_text_config = None
-        elif isinstance(model, PretrainedConfig):
+        elif isinstance(model, PreTrainedConfig):
             model_dict = model.to_dict()
             text_config = model.get_text_config()
             source_text_config = text_config if text_config is not model else None
@@ -40,7 +40,7 @@ def __init__(
         combined["architectures"] = ["ExtractHiddenStatesModel"]
 
         # to_dict() and kwargs both flatten text_config to a plain dict;
-        # downstream get_hf_text_config() needs it as a PretrainedConfig
+        # downstream get_hf_text_config() needs it as a PreTrainedConfig
         # for attribute access. Re-insert the original object.
         if source_text_config is not None:
             combined["text_config"] = source_text_config
diff --git a/vllm/transformers_utils/configs/falcon.py b/vllm/transformers_utils/configs/falcon.py
index c646d241d4eb..7dfe21d0b8bc 100644
--- a/vllm/transformers_utils/configs/falcon.py
+++ b/vllm/transformers_utils/configs/falcon.py
@@ -18,12 +18,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Falcon configuration"""
+"""Falcon configuration."""
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class RWConfig(PretrainedConfig):
+class RWConfig(PreTrainedConfig):
     model_type = "falcon"
     keys_to_ignore_at_inference = ["past_key_values"]
     attribute_map = {
diff --git a/vllm/transformers_utils/configs/funaudiochat.py b/vllm/transformers_utils/configs/funaudiochat.py
index 36a446860c56..dd2573aac426 100644
--- a/vllm/transformers_utils/configs/funaudiochat.py
+++ b/vllm/transformers_utils/configs/funaudiochat.py
@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from transformers import CONFIG_MAPPING, PretrainedConfig
+from transformers import CONFIG_MAPPING, PreTrainedConfig
 
 # NOTE: Temporary shim for FunAudioChat checkpoints.
 # These checkpoints use `model_type="funaudiochat"`, which is not currently
@@ -13,7 +13,7 @@
 # provides an `auto_map`) and vLLM can rely on `AutoConfig.from_pretrained()`.
 
 
-class FunAudioChatAudioEncoderConfig(PretrainedConfig):
+class FunAudioChatAudioEncoderConfig(PreTrainedConfig):
     model_type = "funaudiochat_audio_encoder"
 
     def __init__(
@@ -74,7 +74,7 @@ def __init__(
         self.pad_token_id = pad_token_id
 
 
-class FunAudioChatConfig(PretrainedConfig):
+class FunAudioChatConfig(PreTrainedConfig):
     model_type = "funaudiochat"
     attribute_map = {
         "audio_token_id": "audio_token_index",
@@ -82,8 +82,8 @@ class FunAudioChatConfig(PretrainedConfig):
 
     def __init__(
         self,
-        audio_config: PretrainedConfig | dict | None = None,
-        text_config: PretrainedConfig | dict | None = None,
+        audio_config: PreTrainedConfig | dict | None = None,
+        text_config: PreTrainedConfig | dict | None = None,
         audio_token_index: int = 151646,
         ignore_index: int = -100,
         hidden_size: int | None = None,
diff --git a/vllm/transformers_utils/configs/gemma4.py b/vllm/transformers_utils/configs/gemma4.py
index c57ba5b7ecf2..69cbc06d0482 100644
--- a/vllm/transformers_utils/configs/gemma4.py
+++ b/vllm/transformers_utils/configs/gemma4.py
@@ -4,12 +4,12 @@
 
 from copy import copy
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
 def gemma4_layer_config(
-    text_config: PretrainedConfig, layer_idx: int
-) -> PretrainedConfig:
+    text_config: PreTrainedConfig, layer_idx: int
+) -> PreTrainedConfig:
     """The Gemma4 text config as it applies to one layer.
 
     Gemma4 uses a larger head dimension on its full attention layers than on its
diff --git a/vllm/transformers_utils/configs/glm5_next.py b/vllm/transformers_utils/configs/glm5_next.py
index 95a0081bdd45..4a565884a26d 100644
--- a/vllm/transformers_utils/configs/glm5_next.py
+++ b/vllm/transformers_utils/configs/glm5_next.py
@@ -1,14 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
 
-class Glm5NextTextConfig(PretrainedConfig):
+class Glm5NextTextConfig(PreTrainedConfig):
     model_type = "glm5_next_text"
     base_config_key = "text_config"
     keys_to_ignore_at_inference = ["past_key_values"]
@@ -271,7 +271,7 @@ def layers_block_type(self):
         ]
 
 
-class Glm5NextVisionConfig(PretrainedConfig):
+class Glm5NextVisionConfig(PreTrainedConfig):
     model_type = "glm5_next_vision"
     base_config_key = "vision_config"
 
@@ -324,7 +324,7 @@ def __init__(
         self.swiglu_limit = swiglu_limit
 
 
-class Glm5NextConfig(PretrainedConfig):
+class Glm5NextConfig(PreTrainedConfig):
     model_type = "glm5_next"
     sub_configs = {
         "vision_config": Glm5NextVisionConfig,
@@ -345,7 +345,7 @@ def __init__(
         **kwargs,
     ):
         # Init super() first so base-class defaults don't clobber text-config
-        # values set below (PretrainedConfig has many text-related defaults
+        # values set below (PreTrainedConfig has many text-related defaults
         # that differ from Glm5NextTextConfig).
         super().__init__(**kwargs)
 
@@ -379,7 +379,7 @@ def __init__(
     # and must NOT be mirrored onto text_config: ``architectures`` /
     # ``torch_dtype`` differ between the top-level config and the text
     # sub-config, and mirroring them makes the top-level ``architectures``
-    # silently read back as None (PretrainedConfig initializes both to None),
+    # silently read back as None (PreTrainedConfig initializes both to None),
     # which then fails model-class resolution ("No model architectures are
     # specified").
     _UNMIRRORED_KEYS = [
diff --git a/vllm/transformers_utils/configs/granite4_vision.py b/vllm/transformers_utils/configs/granite4_vision.py
index a67e6ffeb8ee..8522bdd47a3e 100644
--- a/vllm/transformers_utils/configs/granite4_vision.py
+++ b/vllm/transformers_utils/configs/granite4_vision.py
@@ -6,7 +6,7 @@
 import transformers
 
 
-class Granite4VisionConfig(transformers.PretrainedConfig):
+class Granite4VisionConfig(transformers.PreTrainedConfig):
     """Configuration for Granite 4 Vision model.
 
     This config is needed because the granite4_vision model type is not yet
@@ -81,7 +81,7 @@ def __init__(
                 **vision_config
             )
         else:
-            self.vision_config = transformers.PretrainedConfig(**vision_config)
+            self.vision_config = transformers.PreTrainedConfig(**vision_config)
 
         text_model_type = text_config.get("model_type", "granite")
         if text_model_type in transformers.CONFIG_MAPPING:
@@ -89,6 +89,6 @@ def __init__(
                 **text_config
             )
         else:
-            self.text_config = transformers.PretrainedConfig(**text_config)
+            self.text_config = transformers.PreTrainedConfig(**text_config)
 
         super().__init__(**kwargs)
diff --git a/vllm/transformers_utils/configs/hy_v3.py b/vllm/transformers_utils/configs/hy_v3.py
index 9425caf4e03f..939dd5d78ceb 100644
--- a/vllm/transformers_utils/configs/hy_v3.py
+++ b/vllm/transformers_utils/configs/hy_v3.py
@@ -2,17 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class HYV3Config(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`HYV3Model`].
+class HYV3Config(PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a [`HYV3Model`].
     It is used to instantiate a HYV3 model (HY V3 MoE language model) according to
     the specified arguments.
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to
-    control the model outputs. Read the documentation from [`PretrainedConfig`]
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to
+    control the model outputs. Read the documentation from [`PreTrainedConfig`]
     for more information.
 
     Args:
@@ -93,6 +92,7 @@ class HYV3Config(PretrainedConfig):
         >>> config = HYV3Config()
         >>> model = HYV3Model(config)
         ```
+
     """
 
     model_type = "hy_v3"
diff --git a/vllm/transformers_utils/configs/hy_v4.py b/vllm/transformers_utils/configs/hy_v4.py
index beca23ee7437..687a2206c878 100644
--- a/vllm/transformers_utils/configs/hy_v4.py
+++ b/vllm/transformers_utils/configs/hy_v4.py
@@ -52,6 +52,7 @@ class HYV4Config(PreTrainedConfig):
             below zero disable the clamp.
         num_nextn_predict_layers: Number of MTP layers in the checkpoint.
         mtp_loss_factor: Training-time MTP loss weight (unused at inference).
+
     """
 
     model_type = "hy_v4"
diff --git a/vllm/transformers_utils/configs/hyperclovax.py b/vllm/transformers_utils/configs/hyperclovax.py
index 9fa823743d66..5470408b5803 100644
--- a/vllm/transformers_utils/configs/hyperclovax.py
+++ b/vllm/transformers_utils/configs/hyperclovax.py
@@ -17,17 +17,16 @@
 # limitations under the License.
 """HyperCLOVA X model configuration."""
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class HyperCLOVAXConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a
+class HyperCLOVAXConfig(PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a
     [`HyperCLOVAXModel`]. It is used to instantiate a HyperCLOVAX model
     according to the specified arguments, defining the model architecture.
-    Configuration objects inherit from [`PretrainedConfig`] and can be used
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used
     to control the model outputs. Read the documentation from
-    [`PretrainedConfig`] for more information.
+    [`PreTrainedConfig`] for more information.
 
     Args:
         vocab_size (`int`, *optional*, defaults to 32000):
@@ -166,6 +165,7 @@ class HyperCLOVAXConfig(PretrainedConfig):
             `get_rope`. When provided, takes precedence over `rope_theta`
             and `rope_scaling`. If `None`, it is derived from `rope_theta`
             and `rope_scaling` automatically.
+
     """
 
     model_type = "hyperclovax"
diff --git a/vllm/transformers_utils/configs/kimi_k25.py b/vllm/transformers_utils/configs/kimi_k25.py
index 710e9b56367f..ed3d05d7cc21 100644
--- a/vllm/transformers_utils/configs/kimi_k25.py
+++ b/vllm/transformers_utils/configs/kimi_k25.py
@@ -1,17 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Kimi-K2.5 Model Configuration.
+"""Kimi-K2.5 Model Configuration.
 
 This configuration supports video-chunk as an internal modality type.
 A video-chunk is the smallest independently processable unit of video.
 """
 
 from transformers import DeepseekV3Config
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class KimiK25VisionConfig(PretrainedConfig):
+class KimiK25VisionConfig(PreTrainedConfig):
     model_type = "kimi_k25_vision"
 
     def __init__(
@@ -60,7 +59,7 @@ def __init__(
         self.projector_ln_eps = projector_ln_eps
 
 
-class KimiK25Config(PretrainedConfig):
+class KimiK25Config(PreTrainedConfig):
     """Kimi-K2.5 model configuration.
 
     Kimi-K2.5 extends Kimi-K2 with vision support using video-chunks.
@@ -73,6 +72,7 @@ class KimiK25Config(PretrainedConfig):
         ignore_index: The ignore index for the loss function.
         media_placeholder_token_id: The token ID for media placeholders.
         pad_token_id: The token ID for padding.
+
     """
 
     model_type = "kimi_k25"
diff --git a/vllm/transformers_utils/configs/kimi_k3.py b/vllm/transformers_utils/configs/kimi_k3.py
index 62de6dc03b40..6154a438bbf5 100644
--- a/vllm/transformers_utils/configs/kimi_k3.py
+++ b/vllm/transformers_utils/configs/kimi_k3.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Kimi-K3 multimodal configuration."""
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.logger import init_logger
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
@@ -10,7 +10,7 @@
 logger = init_logger(__name__)
 
 
-class KimiK3VisionConfig(PretrainedConfig):
+class KimiK3VisionConfig(PreTrainedConfig):
     model_type = "kimi_k3_vision"
 
     def __init__(
@@ -83,7 +83,7 @@ def __init__(
         self.intermediate_size = vt_intermediate_size
 
 
-class KimiK3Config(PretrainedConfig):
+class KimiK3Config(PreTrainedConfig):
     model_type = "kimi_k3"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/kimi_linear.py b/vllm/transformers_utils/configs/kimi_linear.py
index 6c2fd335ead8..62b3963e1c77 100644
--- a/vllm/transformers_utils/configs/kimi_linear.py
+++ b/vllm/transformers_utils/configs/kimi_linear.py
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class KimiLinearConfig(PretrainedConfig):
+class KimiLinearConfig(PreTrainedConfig):
     model_type = "kimi_linear"
     keys_to_ignore_at_inference = ["past_key_values"]
 
diff --git a/vllm/transformers_utils/configs/kimi_vl.py b/vllm/transformers_utils/configs/kimi_vl.py
index 6d992464cbe8..5733afd6687d 100644
--- a/vllm/transformers_utils/configs/kimi_vl.py
+++ b/vllm/transformers_utils/configs/kimi_vl.py
@@ -3,12 +3,12 @@
 # Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py
 
 from transformers import DeepseekV2Config
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 from vllm.transformers_utils.configs.moonvit import MoonViTConfig
 
 
-class KimiVLConfig(PretrainedConfig):
+class KimiVLConfig(PreTrainedConfig):
     model_type = "kimi_vl"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/laguna.py b/vllm/transformers_utils/configs/laguna.py
deleted file mode 100644
index 2702d3af5aa1..000000000000
--- a/vllm/transformers_utils/configs/laguna.py
+++ /dev/null
@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from transformers.configuration_utils import PretrainedConfig
-
-
-class LagunaConfig(PretrainedConfig):
-    model_type = "laguna"
-    keys_to_ignore_at_inference = ["past_key_values"]
-    base_model_tp_plan = {
-        "layers.*.self_attn.q_proj": "colwise",
-        "layers.*.self_attn.k_proj": "colwise",
-        "layers.*.self_attn.v_proj": "colwise",
-        "layers.*.self_attn.g_proj": "colwise",
-        "layers.*.self_attn.o_proj": "rowwise",
-        "layers.*.mlp.gate_proj": "colwise",
-        "layers.*.mlp.up_proj": "colwise",
-        "layers.*.mlp.down_proj": "rowwise",
-    }
-    base_model_pp_plan = {
-        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
-        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
-        "norm": (["hidden_states"], ["hidden_states"]),
-    }
-
-    def __init__(
-        self,
-        vocab_size: int = 100352,
-        hidden_size: int = 2048,
-        intermediate_size: int = 8192,
-        num_hidden_layers: int = 40,
-        num_attention_heads: int = 48,
-        num_key_value_heads: int = 8,
-        head_dim: int = 128,
-        qkv_bias: bool = False,
-        attention_bias: bool = False,
-        gating: bool | str = True,
-        hidden_act: str = "silu",
-        max_position_embeddings: int = 131072,
-        initializer_range: float = 0.02,
-        rms_norm_eps: float = 1e-6,
-        use_cache: bool = True,
-        tie_word_embeddings: bool = False,
-        rope_theta: float = 500000.0,
-        rope_scaling: dict | None = None,
-        rope_parameters: dict | None = None,
-        partial_rotary_factor: float = 1.0,
-        attention_dropout: float = 0.0,
-        sliding_window: int | None = None,
-        layer_types: list[str] | None = None,
-        swa_attention_sink_enabled: bool = False,
-        swa_rope_parameters: dict | None = None,
-        num_attention_heads_per_layer: list[int] | None = None,
-        num_experts: int = 256,
-        num_experts_per_tok: int = 8,
-        moe_intermediate_size: int = 512,
-        shared_expert_intermediate_size: int = 512,
-        norm_topk_prob: bool = True,
-        decoder_sparse_step: int = 1,
-        mlp_only_layers: list[int] | None = None,
-        router_aux_loss_coef: float = 0.001,
-        output_router_logits: bool = False,
-        moe_routed_scaling_factor: float = 1.0,
-        moe_apply_router_weight_on_input: bool = False,
-        **kwargs,
-    ):
-        if mlp_only_layers is None:
-            mlp_only_layers = [0]
-
-        # Accept either v4-style (rope_theta + rope_scaling) or v5-style
-        # (rope_parameters). Translate v5 → v4 so downstream code has one path.
-        if rope_parameters is not None:
-            rp = dict(rope_parameters)
-            rope_theta = float(rp.pop("rope_theta", rope_theta))
-            rt = rp.pop("rope_type", None)
-            if rt is not None and rt != "default":
-                rope_scaling = {"rope_type": rt, **rp}
-            elif rp and rope_scaling is None:
-                rope_scaling = {"rope_type": "default", **rp}
-
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim
-        self.qkv_bias = qkv_bias
-        self.attention_bias = attention_bias
-        self.gating = gating
-        self.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.partial_rotary_factor = partial_rotary_factor
-        self.attention_dropout = attention_dropout
-        self.sliding_window = sliding_window
-        self.layer_types = layer_types
-        self.swa_attention_sink_enabled = swa_attention_sink_enabled
-        self.swa_rope_parameters = swa_rope_parameters
-        self.num_attention_heads_per_layer = num_attention_heads_per_layer
-        self.num_experts = num_experts
-        self.num_experts_per_tok = num_experts_per_tok
-        self.moe_intermediate_size = moe_intermediate_size
-        self.shared_expert_intermediate_size = shared_expert_intermediate_size
-        self.norm_topk_prob = norm_topk_prob
-        self.decoder_sparse_step = decoder_sparse_step
-        self.mlp_only_layers = mlp_only_layers
-        self.router_aux_loss_coef = router_aux_loss_coef
-        self.output_router_logits = output_router_logits
-        self.moe_routed_scaling_factor = moe_routed_scaling_factor
-        self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-
-__all__ = ["LagunaConfig"]
diff --git a/vllm/transformers_utils/configs/lfm2_moe.py b/vllm/transformers_utils/configs/lfm2_moe.py
index b399a03c030f..903ae9d54700 100644
--- a/vllm/transformers_utils/configs/lfm2_moe.py
+++ b/vllm/transformers_utils/configs/lfm2_moe.py
@@ -2,18 +2,17 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class Lfm2MoeConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Lfm2MoeModel`]. It is used to instantiate a LFM2 Moe
+class Lfm2MoeConfig(PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a [`Lfm2MoeModel`]. It is used to instantiate a LFM2 Moe
     model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
     defaults will yield a similar configuration to that of the LFM2-8B-A1B model.
     e.g. [LiquidAI/LFM2-8B-A1B](https://huggingface.co/LiquidAI/LFM2-8B-A1B)
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
 
 
     Args:
@@ -85,7 +84,9 @@ class Lfm2MoeConfig(PretrainedConfig):
 
     >>> # Accessing the model configuration
     >>> configuration = model.config
-    ```"""  # noqa: E501
+    ```
+
+    """  # noqa: E501
 
     model_type = "lfm2_moe"
     keys_to_ignore_at_inference = ["past_key_values"]
diff --git a/vllm/transformers_utils/configs/medusa.py b/vllm/transformers_utils/configs/medusa.py
index f146c4c5f5d9..e8cef7a35fca 100644
--- a/vllm/transformers_utils/configs/medusa.py
+++ b/vllm/transformers_utils/configs/medusa.py
@@ -3,12 +3,12 @@
 
 import os
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.transformers_utils.utils import without_trust_remote_code
 
 
-class MedusaConfig(PretrainedConfig):
+class MedusaConfig(PreTrainedConfig):
     model_type = "medusa"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/midashenglm.py b/vllm/transformers_utils/configs/midashenglm.py
index f1bbd057103e..76d8393e010e 100644
--- a/vllm/transformers_utils/configs/midashenglm.py
+++ b/vllm/transformers_utils/configs/midashenglm.py
@@ -22,13 +22,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniTextConfig,
 )
 
 
-class DashengConfig(PretrainedConfig):
+class DashengConfig(PreTrainedConfig):
     model_type = "midashenglm_dasheng_encoder"
 
     def __init__(
@@ -80,7 +80,7 @@ def __init__(
         super().__init__(**kwargs)
 
 
-class MiDashengLMConfig(PretrainedConfig):
+class MiDashengLMConfig(PreTrainedConfig):
     model_type = "midashenglm"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/mimo_v2_omni.py b/vllm/transformers_utils/configs/mimo_v2_omni.py
index b87ca22a9a83..c0e2f4286c28 100644
--- a/vllm/transformers_utils/configs/mimo_v2_omni.py
+++ b/vllm/transformers_utils/configs/mimo_v2_omni.py
@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
-class Mimo_VLVisionConfig(PretrainedConfig):
+class Mimo_VLVisionConfig(PreTrainedConfig):
     model_type = "mimovl"
     base_config_key = "vision_config"
 
diff --git a/vllm/transformers_utils/configs/minimax_m3.py b/vllm/transformers_utils/configs/minimax_m3.py
index c340dda85a6c..374b854a662e 100644
--- a/vllm/transformers_utils/configs/minimax_m3.py
+++ b/vllm/transformers_utils/configs/minimax_m3.py
@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
-class MiniMaxM3TextConfig(PretrainedConfig):
+class MiniMaxM3TextConfig(PreTrainedConfig):
     """Config for the MiniMax M3 text backbone (MiniMaxM3SparseForCausalLM).
 
     Defaults mirror the ``text_config`` of the MiniMax-M3-preview checkpoint.
@@ -120,7 +120,7 @@ def __init__(self, num_hidden_layers: int = 1, **kwargs):
         super().__init__(num_hidden_layers=num_hidden_layers, **kwargs)
 
 
-class MiniMaxM3Config(PretrainedConfig):
+class MiniMaxM3Config(PreTrainedConfig):
     """Top-level MiniMax M3 (VL) config.
 
     Holds the text backbone as ``text_config`` so that
diff --git a/vllm/transformers_utils/configs/mistral.py b/vllm/transformers_utils/configs/mistral.py
index 2b0796691472..aaeae9cbcfaa 100644
--- a/vllm/transformers_utils/configs/mistral.py
+++ b/vllm/transformers_utils/configs/mistral.py
@@ -2,9 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from packaging.version import Version
-from transformers import PretrainedConfig, WhisperConfig
-from transformers import __version__ as TRANSFORMERS_VERSION
+from transformers import PreTrainedConfig, WhisperConfig
 
 from vllm.logger import init_logger
 
@@ -14,7 +12,7 @@
 def adapt_config_dict(
     config_dict: dict[str, Any],
     defaults: dict[str, Any],
-) -> PretrainedConfig:
+) -> PreTrainedConfig:
     config_dict = _remap_general_mistral_args(config_dict)
     config_dict = _remap_mistral_sliding_window(config_dict)
 
@@ -88,7 +86,7 @@ def adapt_config_dict(
     for k, v in defaults.items():
         config_dict.setdefault(k, v)
 
-    config = PretrainedConfig.from_dict(config_dict)
+    config = PreTrainedConfig.from_dict(config_dict)
 
     logger.debug("Initialized config %s", config)
 
@@ -105,8 +103,8 @@ def _remap_mistral_vision_args(config: dict) -> dict:
     config = {
         "model_type": "pixtral",
         "architectures": ["PixtralForConditionalGeneration"],
-        "text_config": PretrainedConfig.from_dict(config),
-        "vision_config": PretrainedConfig.from_dict(vision_config),
+        "text_config": PreTrainedConfig.from_dict(config),
+        "vision_config": PreTrainedConfig.from_dict(vision_config),
     }
     if quant_config:
         config["quantization_config"] = quant_config
@@ -119,7 +117,6 @@ def _remap_mistral_yarn_args(config: dict) -> dict:
         "original_max_position_embeddings": ("original_max_position_embeddings", int),
         "beta": ("beta_fast", float),
         "alpha": ("beta_slow", float),
-        "apply_scale": ("apply_yarn_scaling", bool),
     }
 
     yarn_config = config.get("yarn") or {}
@@ -136,9 +133,10 @@ def _remap_mistral_yarn_args(config: dict) -> dict:
             # Cast to remove Transformers > v5 type warnings
             config["rope_parameters"][new_name] = cast(yarn_config.pop(old_name))
 
-    # Ignore apply_yarn_scaling in Transformers > v5 RoPE validation to remove warnings
-    if Version(TRANSFORMERS_VERSION) >= Version("5.3.0.dev0"):
-        config["ignore_keys_at_rope_validation"] = {"apply_yarn_scaling"}
+    # `apply_scale: false` means no magnitude correction, which Transformers
+    # spells as an explicit attention_factor of 1.
+    if not yarn_config.pop("apply_scale", True):
+        config["rope_parameters"]["attention_factor"] = 1.0
 
     assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}"
 
@@ -250,7 +248,7 @@ def _remap_mistral_audio_args(config: dict) -> dict:
     config = {
         "model_type": "voxtral",
         "architectures": [architecture],
-        "text_config": PretrainedConfig.from_dict(config),
+        "text_config": PreTrainedConfig.from_dict(config),
         "audio_config": WhisperConfig(
             num_mel_bins=encoder_args["audio_encoding_args"]["num_mel_bins"],
             window_size=encoder_args["audio_encoding_args"]["window_size"],
diff --git a/vllm/transformers_utils/configs/mlp_speculator.py b/vllm/transformers_utils/configs/mlp_speculator.py
index 75745f227f48..f668e5ef5202 100644
--- a/vllm/transformers_utils/configs/mlp_speculator.py
+++ b/vllm/transformers_utils/configs/mlp_speculator.py
@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
-class MLPSpeculatorConfig(PretrainedConfig):
+class MLPSpeculatorConfig(PreTrainedConfig):
     model_type = "mlp_speculator"
 
     attribute_map = {
@@ -24,8 +24,7 @@ def __init__(
         scale_input: bool = False,
         **kwargs,
     ):
-        """
-        Initialize an MLPSpeculatorConfig
+        """Initialize an MLPSpeculatorConfig.
 
         Args:
             vocab_size: int
@@ -52,6 +51,7 @@ def __init__(
             scale_input: bool
                 if True, will scale the initial hidden states from
                 the base model.
+
         """
         if top_k_tokens_per_head is None:
             top_k_tokens_per_head = [5, 4, 3]
diff --git a/vllm/transformers_utils/configs/moondream3.py b/vllm/transformers_utils/configs/moondream3.py
index 307bb2977206..e8a80082c1ca 100644
--- a/vllm/transformers_utils/configs/moondream3.py
+++ b/vllm/transformers_utils/configs/moondream3.py
@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Configuration for Moondream3 model."""
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 
-class Moondream3VisionConfig(PretrainedConfig):
+class Moondream3VisionConfig(PreTrainedConfig):
     """Vision encoder configuration for Moondream3."""
 
     model_type = "moondream3_vision"
@@ -43,7 +43,7 @@ def __init__(
         self.image_size = crop_size
 
 
-class Moondream3TextConfig(PretrainedConfig):
+class Moondream3TextConfig(PreTrainedConfig):
     """Text decoder configuration for Moondream3."""
 
     model_type = "moondream3_text"
@@ -101,7 +101,7 @@ def __init__(
         self.num_experts_per_tok = self.moe_experts_per_token
 
 
-class Moondream3Config(PretrainedConfig):
+class Moondream3Config(PreTrainedConfig):
     """Combined configuration for Moondream3 multimodal model."""
 
     model_type = "moondream3"
@@ -148,5 +148,6 @@ def get_text_config(self, decoder: bool = False) -> "Moondream3TextConfig":
 
         Args:
             decoder: Ignored. Only used for encoder-decoder models.
+
         """
         return self.text_config
diff --git a/vllm/transformers_utils/configs/moonvit.py b/vllm/transformers_utils/configs/moonvit.py
index 6e9b2897f4cc..c7cd34d41df1 100644
--- a/vllm/transformers_utils/configs/moonvit.py
+++ b/vllm/transformers_utils/configs/moonvit.py
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class MoonViTConfig(PretrainedConfig):
+class MoonViTConfig(PreTrainedConfig):
     model_type = "moonvit"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/moss_transcribe_diarize.py b/vllm/transformers_utils/configs/moss_transcribe_diarize.py
index 20abbd8feec9..6d57642f39e4 100644
--- a/vllm/transformers_utils/configs/moss_transcribe_diarize.py
+++ b/vllm/transformers_utils/configs/moss_transcribe_diarize.py
@@ -3,11 +3,11 @@
 
 from typing import Any
 
-from transformers import PretrainedConfig, Qwen3Config
+from transformers import PreTrainedConfig, Qwen3Config
 from transformers.models.whisper.configuration_whisper import WhisperConfig
 
 
-class MossTranscribeDiarizeConfig(PretrainedConfig):
+class MossTranscribeDiarizeConfig(PreTrainedConfig):
     """Configuration for MOSS-Transcribe-Diarize."""
 
     model_type = "moss_transcribe_diarize"
diff --git a/vllm/transformers_utils/configs/muse_glimmer.py b/vllm/transformers_utils/configs/muse_glimmer.py
index 9022c84cf0cc..80c4c4e171f2 100644
--- a/vllm/transformers_utils/configs/muse_glimmer.py
+++ b/vllm/transformers_utils/configs/muse_glimmer.py
@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 from transformers import Qwen3Config
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
 def _default_no_rope_layers(num_hidden_layers: int) -> list[int]:
@@ -27,7 +27,7 @@ def _default_no_rope_layers(num_hidden_layers: int) -> list[int]:
     ]
 
 
-class MuseGlimmerTextConfig(PretrainedConfig):
+class MuseGlimmerTextConfig(PreTrainedConfig):
     model_type = "muse_glimmer_text"
     keys_to_ignore_at_inference = ["past_key_values"]
 
@@ -127,7 +127,7 @@ def __init__(
         )
 
 
-class MuseGlimmerVisionConfig(PretrainedConfig):
+class MuseGlimmerVisionConfig(PreTrainedConfig):
     model_type = "muse_glimmer_vision"
 
     def __init__(
@@ -177,7 +177,7 @@ def __init__(
         self.layer_types = layer_types
 
 
-class MuseGlimmerConfig(PretrainedConfig):
+class MuseGlimmerConfig(PreTrainedConfig):
     model_type = "muse_glimmer"
     sub_configs = {
         "text_config": MuseGlimmerTextConfig,
diff --git a/vllm/transformers_utils/configs/nemotron.py b/vllm/transformers_utils/configs/nemotron.py
index 62f52703029b..135a0e1596dc 100644
--- a/vllm/transformers_utils/configs/nemotron.py
+++ b/vllm/transformers_utils/configs/nemotron.py
@@ -15,25 +15,24 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Nemotron model configuration"""
+"""Nemotron model configuration."""
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
 
 
-class NemotronConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a
+class NemotronConfig(PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a
     [`NemotronModel`]. It is used to instantiate a Nemotron model
     according to the specified arguments, defining the model architecture.
     Instantiating a configuration with the defaults will yield a similar
     configuration to that of the Nemotron-8B.
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be
+    Configuration objects inherit from [`PreTrainedConfig`] and can be
     used to control the model outputs. Read the documentation from
-    [`PretrainedConfig`] for more information.
+    [`PreTrainedConfig`] for more information.
 
 
     Args:
@@ -114,7 +113,9 @@ class NemotronConfig(PretrainedConfig):
     >>> model = NemotronModel(configuration)
     >>> # Accessing the model configuration
     >>> configuration = model.config
-    ```"""
+    ```
+
+    """
 
     model_type = "nemotron"
     keys_to_ignore_at_inference = ["past_key_values"]
@@ -193,9 +194,7 @@ def __init__(
         )
 
     def _rope_parameters_validation(self):
-        """
-        Validate the `rope_parameters` configuration.
-        """
+        """Validate the `rope_parameters` configuration."""
         if self.rope_parameters is None:
             return
 
diff --git a/vllm/transformers_utils/configs/nemotron_h.py b/vllm/transformers_utils/configs/nemotron_h.py
index ed62b5d294b3..c25a5d54e3d7 100644
--- a/vllm/transformers_utils/configs/nemotron_h.py
+++ b/vllm/transformers_utils/configs/nemotron_h.py
@@ -15,22 +15,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""NemotronH model configuration"""
+"""NemotronH model configuration."""
 
 import regex as re
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
 
 
-class NemotronHConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a
+class NemotronHConfig(PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a
     [`NemotronHModel`]. It is used to instantiate a NemotronH model according
     to the specified arguments, defining the model architecture. Instantiating
     a configuration with the defaults will yield a similar configuration to
     that of the NemotronH-v0.1 model.
+
     Args:
         vocab_size (`int`, *optional*, defaults to 131072):
             Vocabulary size of the NemotronH model. Defines the number of
@@ -139,6 +139,7 @@ class NemotronHConfig(PretrainedConfig):
             Size of chunks for Mamba processing.
         rescale_prenorm_residual (`bool`, *optional*, defaults to `True`):
             Whether to rescale the pre-normalization residual connections.
+
     """
 
     model_type = "nemotron_h"
diff --git a/vllm/transformers_utils/configs/olmo_hybrid.py b/vllm/transformers_utils/configs/olmo_hybrid.py
index cdca81757e70..62d3c4e0d5d1 100644
--- a/vllm/transformers_utils/configs/olmo_hybrid.py
+++ b/vllm/transformers_utils/configs/olmo_hybrid.py
@@ -2,12 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class OlmoHybridConfig(PretrainedConfig):
-    r"""
-        Configuration class for [`OlmoHybridModel`]. It is used to
+class OlmoHybridConfig(PreTrainedConfig):
+    r"""Configuration class for [`OlmoHybridModel`]. It is used to
         instantiate an OLMo Hybrid model according to the specified
         arguments, defining the model architecture. Instantiating a
         configuration with the defaults will yield a similar
@@ -19,7 +18,7 @@ class OlmoHybridConfig(PretrainedConfig):
         can be used to control the model outputs. Read the
         documentation from [`PreTrainedConfig`] for more information.
 
-        Args:
+    Args:
             vocab_size (`int`, *optional*, defaults to 100352):
                 Vocabulary size of the OlmoHybrid model. Defines
                 the number of different tokens that can be
@@ -159,6 +158,7 @@ class OlmoHybridConfig(PretrainedConfig):
         >>> model = OlmoHybridModel(configuration)
         >>> configuration = model.config
     ```
+
     """
 
     model_type = "olmo_hybrid"
diff --git a/vllm/transformers_utils/configs/openvla.py b/vllm/transformers_utils/configs/openvla.py
index 897af8f705dd..dba2f7f5a26d 100644
--- a/vllm/transformers_utils/configs/openvla.py
+++ b/vllm/transformers_utils/configs/openvla.py
@@ -9,10 +9,10 @@
 
 from typing import Any
 
-from transformers import LlamaConfig, PretrainedConfig
+from transformers import LlamaConfig, PreTrainedConfig
 
 
-class OpenVLAConfig(PretrainedConfig):
+class OpenVLAConfig(PreTrainedConfig):
     """Configuration class for OpenVLA models."""
 
     model_type = "openvla"
diff --git a/vllm/transformers_utils/configs/ovis.py b/vllm/transformers_utils/configs/ovis.py
index 294b4c9037aa..bc84bfe812fc 100644
--- a/vllm/transformers_utils/configs/ovis.py
+++ b/vllm/transformers_utils/configs/ovis.py
@@ -7,13 +7,14 @@
 # Ovis Config with AimV2 config registration removed for Transformers compatibility
 from typing import Any
 
-from transformers import AutoConfig, PretrainedConfig
+from transformers import AutoConfig, PreTrainedConfig
 
 
-class AIMv2Config(PretrainedConfig):
+class AIMv2Config(PreTrainedConfig):
     """This is the configuration class to store the configuration of an [`AIMv2Model`].
     Instantiating a configuration with the defaults will yield a similar configuration
     to that of the [apple/aimv2-large-patch14-224](https://huggingface.co/apple/aimv2-large-patch14-224).
+
     Args:
         hidden_size: Dimension of the hidden representations.
         intermediate_size: Dimension of the SwiGLU representations.
@@ -28,7 +29,8 @@ class AIMv2Config(PretrainedConfig):
         projection_dropout: Dropout ratio for the projection layer after the attention.
         qkv_bias: Whether to add a bias to the queries, keys and values.
         use_bias: Whether to add a bias in the feed-forward and projection layers.
-        kwargs: Keyword arguments for the [`PretrainedConfig`].
+        kwargs: Keyword arguments for the [`PreTrainedConfig`].
+
     """
 
     model_type: str = "aimv2"
@@ -68,7 +70,7 @@ def __init__(
 # ----------------------------------------------------------------------
 #                     Visual Tokenizer Configuration
 # ----------------------------------------------------------------------
-class BaseVisualTokenizerConfig(PretrainedConfig):
+class BaseVisualTokenizerConfig(PreTrainedConfig):
     def __init__(
         self,
         vocab_size=16384,
@@ -76,7 +78,7 @@ def __init__(
         tau=1.0,
         depths=None,
         drop_cls_token=False,
-        backbone_config: PretrainedConfig | dict | None = None,
+        backbone_config: PreTrainedConfig | dict | None = None,
         hidden_stride: int = 1,
         **kwargs,
     ):
@@ -90,10 +92,10 @@ def __init__(
         self.backbone_kwargs = dict[str, Any]()
         self.drop_cls_token = drop_cls_token
         if backbone_config is not None:
-            assert isinstance(backbone_config, (PretrainedConfig, dict)), (
-                f"expect `backbone_config` to be instance of PretrainedConfig or dict, but got {type(backbone_config)} type"
+            assert isinstance(backbone_config, (PreTrainedConfig, dict)), (
+                f"expect `backbone_config` to be instance of PreTrainedConfig or dict, but got {type(backbone_config)} type"
             )
-            if not isinstance(backbone_config, PretrainedConfig):
+            if not isinstance(backbone_config, PreTrainedConfig):
                 model_type = backbone_config["model_type"]
                 if model_type != "aimv2":
                     backbone_config.pop("model_type")
@@ -137,13 +139,13 @@ def __init__(self, **kwargs):
 # ----------------------------------------------------------------------
 #                           Ovis Configuration
 # ----------------------------------------------------------------------
-class OvisConfig(PretrainedConfig):
+class OvisConfig(PreTrainedConfig):
     model_type = "ovis"
 
     def __init__(
         self,
-        llm_config: PretrainedConfig | dict | None = None,
-        visual_tokenizer_config: PretrainedConfig | dict | None = None,
+        llm_config: PreTrainedConfig | dict | None = None,
+        visual_tokenizer_config: PreTrainedConfig | dict | None = None,
         multimodal_max_length=8192,
         hidden_size=None,
         conversation_formatter_class=None,
@@ -153,10 +155,10 @@ def __init__(
     ):
         super().__init__(**kwargs)
         if llm_config is not None:
-            assert isinstance(llm_config, (PretrainedConfig, dict)), (
-                f"expect `llm_config` to be instance of PretrainedConfig or dict, but got {type(llm_config)} type"
+            assert isinstance(llm_config, (PreTrainedConfig, dict)), (
+                f"expect `llm_config` to be instance of PreTrainedConfig or dict, but got {type(llm_config)} type"
             )
-            if not isinstance(llm_config, PretrainedConfig):
+            if not isinstance(llm_config, PreTrainedConfig):
                 model_type = llm_config["model_type"]
                 llm_config.pop("model_type")
                 llm_config = AutoConfig.for_model(model_type, **llm_config)
@@ -164,10 +166,10 @@ def __init__(
         # map llm_config to text_config
         self.text_config = llm_config
         if visual_tokenizer_config is not None:
-            assert isinstance(visual_tokenizer_config, (PretrainedConfig, dict)), (
-                f"expect `visual_tokenizer_config` to be instance of PretrainedConfig or dict, but got {type(visual_tokenizer_config)} type"
+            assert isinstance(visual_tokenizer_config, (PreTrainedConfig, dict)), (
+                f"expect `visual_tokenizer_config` to be instance of PreTrainedConfig or dict, but got {type(visual_tokenizer_config)} type"
             )
-            if not isinstance(visual_tokenizer_config, PretrainedConfig):
+            if not isinstance(visual_tokenizer_config, PreTrainedConfig):
                 model_type = visual_tokenizer_config["model_type"]
                 visual_tokenizer_config.pop("model_type")
                 visual_tokenizer_config = AutoConfig.for_model(
diff --git a/vllm/transformers_utils/configs/parakeet.py b/vllm/transformers_utils/configs/parakeet.py
index febd1da446a2..c574317f81de 100644
--- a/vllm/transformers_utils/configs/parakeet.py
+++ b/vllm/transformers_utils/configs/parakeet.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
 
-from transformers import ParakeetEncoderConfig, PretrainedConfig
+from transformers import ParakeetEncoderConfig, PreTrainedConfig
 
 
 class ParakeetConfig(ParakeetEncoderConfig):
@@ -24,9 +24,9 @@ def __init__(
 
     @staticmethod
     def from_hf_config(
-        config: PretrainedConfig, *, llm_hidden_size: int, max_model_len: int
+        config: PreTrainedConfig, *, llm_hidden_size: int, max_model_len: int
     ) -> "ParakeetConfig":
-        assert isinstance(config, PretrainedConfig)
+        assert isinstance(config, PreTrainedConfig)
         return ParakeetConfig(
             **config.to_dict(),
             scale_input=False,
@@ -55,8 +55,8 @@ class ExtractorConfig:
     padding_value: float = 0.0
 
     @classmethod
-    def from_hf_config(cls, config: PretrainedConfig) -> "ExtractorConfig":
-        assert isinstance(config, PretrainedConfig)
+    def from_hf_config(cls, config: PreTrainedConfig) -> "ExtractorConfig":
+        assert isinstance(config, PreTrainedConfig)
         defaults = ("hop_length", "win_length", "preemphasis", "n_fft", "padding_value")
         optional_kwargs = {
             name: getattr(config, name) for name in defaults if hasattr(config, name)
diff --git a/vllm/transformers_utils/configs/qianfan_ocr.py b/vllm/transformers_utils/configs/qianfan_ocr.py
index da004bb90f4f..73df82c8fe0c 100644
--- a/vllm/transformers_utils/configs/qianfan_ocr.py
+++ b/vllm/transformers_utils/configs/qianfan_ocr.py
@@ -3,11 +3,11 @@
 
 from typing import Any
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 from transformers.models.auto import CONFIG_MAPPING
 
 
-class QianfanOCRVisionConfig(PretrainedConfig):
+class QianfanOCRVisionConfig(PreTrainedConfig):
     model_type = "qianfan_ocr_vision"
 
     def __init__(
@@ -53,7 +53,7 @@ def __init__(
         self.use_mean_pooling = use_mean_pooling
 
 
-class QianfanOCRConfig(PretrainedConfig):
+class QianfanOCRConfig(PreTrainedConfig):
     model_type = "qianfan_ocr"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/qwen3_5.py b/vllm/transformers_utils/configs/qwen3_5.py
index d5820a5783c6..57f64f0f7814 100644
--- a/vllm/transformers_utils/configs/qwen3_5.py
+++ b/vllm/transformers_utils/configs/qwen3_5.py
@@ -14,12 +14,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Qwen3.5 model configuration"""
+"""Qwen3.5 model configuration."""
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class Qwen3_5TextConfig(PretrainedConfig):
+class Qwen3_5TextConfig(PreTrainedConfig):
     model_type = "qwen3_5_text"
     keys_to_ignore_at_inference = ["past_key_values"]
 
@@ -108,7 +108,7 @@ def __init__(
         self.linear_num_value_heads = linear_num_value_heads
         super().__init__(**kwargs)
         # Set these AFTER super().__init__() because transformers v4's
-        # PretrainedConfig.__init__ has these as explicit params with different
+        # PreTrainedConfig.__init__ has these as explicit params with different
         # defaults (e.g. tie_word_embeddings=True) that would overwrite our values.
         self.pad_token_id = pad_token_id
         self.bos_token_id = bos_token_id
@@ -116,7 +116,7 @@ def __init__(
         self.tie_word_embeddings = tie_word_embeddings
 
 
-class Qwen3_5VisionConfig(PretrainedConfig):
+class Qwen3_5VisionConfig(PreTrainedConfig):
     model_type = "qwen3_5"
     base_config_key = "vision_config"
 
@@ -152,7 +152,7 @@ def __init__(
         self.initializer_range = initializer_range
 
 
-class Qwen3_5Config(PretrainedConfig):
+class Qwen3_5Config(PreTrainedConfig):
     model_type = "qwen3_5"
     sub_configs = {
         "vision_config": Qwen3_5VisionConfig,
@@ -186,7 +186,7 @@ def __init__(
         self.vision_start_token_id = vision_start_token_id
         self.vision_end_token_id = vision_end_token_id
         super().__init__(**kwargs)
-        # Set after super().__init__() to avoid v4 PretrainedConfig overwrite
+        # Set after super().__init__() to avoid v4 PreTrainedConfig overwrite
         self.tie_word_embeddings = tie_word_embeddings
 
 
diff --git a/vllm/transformers_utils/configs/qwen3_5_moe.py b/vllm/transformers_utils/configs/qwen3_5_moe.py
index ec229ce81426..34e66c825ac5 100644
--- a/vllm/transformers_utils/configs/qwen3_5_moe.py
+++ b/vllm/transformers_utils/configs/qwen3_5_moe.py
@@ -14,12 +14,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Qwen3.5-MoE model configuration"""
+"""Qwen3.5-MoE model configuration."""
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class Qwen3_5MoeTextConfig(PretrainedConfig):
+class Qwen3_5MoeTextConfig(PreTrainedConfig):
     model_type = "qwen3_5_moe_text"
     keys_to_ignore_at_inference = ["past_key_values"]
 
@@ -120,7 +120,7 @@ def __init__(
         self.router_aux_loss_coef = router_aux_loss_coef
         super().__init__(**kwargs)
         # Set these AFTER super().__init__() because transformers v4's
-        # PretrainedConfig.__init__ has these as explicit params with different
+        # PreTrainedConfig.__init__ has these as explicit params with different
         # defaults (e.g. tie_word_embeddings=True) that would overwrite our values.
         self.pad_token_id = pad_token_id
         self.bos_token_id = bos_token_id
@@ -128,7 +128,7 @@ def __init__(
         self.tie_word_embeddings = tie_word_embeddings
 
 
-class Qwen3_5MoeVisionConfig(PretrainedConfig):
+class Qwen3_5MoeVisionConfig(PreTrainedConfig):
     model_type = "qwen3_5_moe"
     base_config_key = "vision_config"
 
@@ -164,7 +164,7 @@ def __init__(
         self.initializer_range = initializer_range
 
 
-class Qwen3_5MoeConfig(PretrainedConfig):
+class Qwen3_5MoeConfig(PreTrainedConfig):
     model_type = "qwen3_5_moe"
     sub_configs = {
         "vision_config": Qwen3_5MoeVisionConfig,
@@ -198,7 +198,7 @@ def __init__(
         self.vision_start_token_id = vision_start_token_id
         self.vision_end_token_id = vision_end_token_id
         super().__init__(**kwargs)
-        # Set after super().__init__() to avoid v4 PretrainedConfig overwrite
+        # Set after super().__init__() to avoid v4 PreTrainedConfig overwrite
         self.tie_word_embeddings = tie_word_embeddings
 
 
diff --git a/vllm/transformers_utils/configs/qwen3_asr.py b/vllm/transformers_utils/configs/qwen3_asr.py
index 652c02525fe6..77125e4b22df 100644
--- a/vllm/transformers_utils/configs/qwen3_asr.py
+++ b/vllm/transformers_utils/configs/qwen3_asr.py
@@ -17,14 +17,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.modeling_rope_utils import rope_config_validation
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
 
 
-class Qwen3ASRAudioEncoderConfig(PretrainedConfig):
+class Qwen3ASRAudioEncoderConfig(PreTrainedConfig):
     r"""
     This is the configuration class to store the configuration of a [`Qwen3ASRAudioEncoder`]. It is used to instantiate a
     Qwen3-ASR audio encoder according to the specified arguments, defining the model architecture. Instantiating a
@@ -33,8 +33,8 @@ class Qwen3ASRAudioEncoderConfig(PretrainedConfig):
 
     e.g. [Qwen/Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B)
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
 
     Args:
         num_mel_bins (`int`, *optional*, defaults to 128):
@@ -130,15 +130,15 @@ def __init__(
         self.downsample_hidden_size = downsample_hidden_size
 
 
-class Qwen3ASRTextConfig(PretrainedConfig):
+class Qwen3ASRTextConfig(PreTrainedConfig):
     r"""
     This is the configuration class to store the configuration of a [`Qwen3ASRTextModel`]. It is used to instantiate a
     Qwen3-ASR model according to the specified arguments, defining the model architecture. Instantiating a configuration
     with the defaults will yield a similar configuration to that of
     Qwen3-ASR-1.7B [Qwen/Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B)
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
 
     Args:
         vocab_size (`int`, *optional*, defaults to 151936):
@@ -285,7 +285,7 @@ def __init__(
         super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
 
 
-class Qwen3ASRThinkerConfig(PretrainedConfig):
+class Qwen3ASRThinkerConfig(PreTrainedConfig):
     r"""
     This is the configuration class to store the configuration of a [`Qwen3ASRThinker`]. It is used to instantiate a
     Qwen3-ASR-Thinker model according to the specified arguments, defining the model architecture. Instantiating a
@@ -294,8 +294,8 @@ class Qwen3ASRThinkerConfig(PretrainedConfig):
 
     e.g. [Qwen/Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B)
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
 
     Args:
         audio_config (`dict`, *optional*):
@@ -365,7 +365,7 @@ def __init__(
         self.audio_token_id = audio_token_id
 
 
-class Qwen3ASRConfig(PretrainedConfig):
+class Qwen3ASRConfig(PreTrainedConfig):
     """
     This is the configuration class to store the configuration of a [`Qwen3ASRForConditionalGeneration`]. It is used to instantiate a Qwen3ASR
     model according to the specified sub-models configurations, defining the model architecture.
@@ -373,8 +373,8 @@ class Qwen3ASRConfig(PretrainedConfig):
     Instantiating a configuration with the defaults will yield a similar configuration to that of the
     [Qwen/Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B) architecture.
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
 
     Args:
         thinker_config (`dict`, *optional*): Configuration of the underlying thinker sub-model.
@@ -437,7 +437,7 @@ def __init__(
         self.timestamp_segment_time = timestamp_segment_time
         super().__init__(**kwargs)
 
-    def get_text_config(self, decoder=False) -> "PretrainedConfig":
+    def get_text_config(self, decoder=False) -> "PreTrainedConfig":
         """
         Returns the config that is meant to be used with text IO. On most models, it is the original config instance
         itself. On specific composite models, it is under a set of valid names.
diff --git a/vllm/transformers_utils/configs/qwen3_next.py b/vllm/transformers_utils/configs/qwen3_next.py
index de579ed2cf3f..8b083ff9633d 100644
--- a/vllm/transformers_utils/configs/qwen3_next.py
+++ b/vllm/transformers_utils/configs/qwen3_next.py
@@ -14,23 +14,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Qwen3-Next model configuration"""
+"""Qwen3-Next model configuration."""
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
 
 
-class Qwen3NextConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a
+class Qwen3NextConfig(PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a
     Qwen3-Next model according to the specified arguments, defining the model architecture.
     Instantiating a configuration with the defaults will yield a similar configuration to that of
     Qwen3-Next-80B-A3B-Instruct [Qwen/Qwen3-Next-80B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct).
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
 
 
     Args:
@@ -157,6 +156,7 @@ class Qwen3NextConfig(PretrainedConfig):
     >>> # Accessing the model configuration
     >>> configuration = model.config
     ```
+
     """  # noqa: E501
 
     model_type = "qwen3_next"
diff --git a/vllm/transformers_utils/configs/radio.py b/vllm/transformers_utils/configs/radio.py
index e668c5c5e7f2..89a56a85ef9e 100644
--- a/vllm/transformers_utils/configs/radio.py
+++ b/vllm/transformers_utils/configs/radio.py
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Radio vision model configuration"""
+"""Radio vision model configuration."""
 
 from typing import Any
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
@@ -20,9 +20,8 @@
 OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
 
 
-class RadioConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a Radio
+class RadioConfig(PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a Radio
     vision model. It is used to instantiate a Radio model according to the
     specified arguments, defining the model architecture.
 
@@ -55,6 +54,7 @@ class RadioConfig(PretrainedConfig):
             dedicated video patch embedder (3*T*P*P -> hidden) separate from the
             image embedder (3*P*P -> hidden). When False, a single embedder with
             input size 3*T*P*P is used for both (images are duplicated T times).
+
     """
 
     model_type = "radio"
diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py
index 9f45183328a9..6407fc4868fa 100644
--- a/vllm/transformers_utils/configs/speculators/algos.py
+++ b/vllm/transformers_utils/configs/speculators/algos.py
@@ -14,8 +14,7 @@ def decorator(fn):
 
 @register_speculator("eagle3")
 def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
-    """
-    Apply Eagle-3 specific configuration transformations to the `dict` used to
+    """Apply Eagle-3 specific configuration transformations to the `dict` used to
     construct the Transformers PreTrainedConfig.
 
     Eagle-3 specific fields:
@@ -28,7 +27,6 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
         provide intermediate hidden states that help the drafter make better
         predictions. This is the standard field used in Eagle3 checkpoints.
     """
-
     pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
     if config_dict.get("target_hidden_size") is not None:
         pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
@@ -54,8 +52,7 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
 
 @register_speculator("peagle")
 def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
-    """
-    Apply PEagle (Parallel Eagle) specific configuration transformations to
+    """Apply PEagle (Parallel Eagle) specific configuration transformations to
     the `dict` used to construct the Transformers PreTrainedConfig.
 
     PEagle specific fields:
@@ -92,8 +89,7 @@ def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
 
 @register_speculator("dflash")
 def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
-    """
-    Apply DFlash specific configuration transformations to the `dict` used to
+    """Apply DFlash specific configuration transformations to the `dict` used to
     construct the Transformers PreTrainedConfig.
 
     DFlash specific fields:
@@ -132,8 +128,7 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
 
 @register_speculator("dflash2")
 def update_dflash2(config_dict: dict, pre_trained_config: dict) -> None:
-    """
-    Apply DFlash2 specific configuration transformations to the `dict` used
+    """Apply DFlash2 specific configuration transformations to the `dict` used
     to construct the Transformers PreTrainedConfig.
 
     DFlash2 extends DFlash with a grouped convolution layer and a
@@ -156,8 +151,7 @@ def update_dflash2(config_dict: dict, pre_trained_config: dict) -> None:
 
 @register_speculator("dspark")
 def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
-    """
-    Apply DSpark specific configuration transformations to the `dict` used to
+    """Apply DSpark specific configuration transformations to the `dict` used to
     construct the Transformers PreTrainedConfig.
 
     DSpark extends DFlash with a Markov logit-bias head, reusing the same
diff --git a/vllm/transformers_utils/configs/speculators/base.py b/vllm/transformers_utils/configs/speculators/base.py
index 1cbaf6e801e1..cb188a1ef0a2 100644
--- a/vllm/transformers_utils/configs/speculators/base.py
+++ b/vllm/transformers_utils/configs/speculators/base.py
@@ -4,7 +4,7 @@
 from dataclasses import fields
 from typing import Any
 
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm.transformers_utils.configs.speculators.algos import (
     SUPPORTED_SPECULATORS_TYPES,
@@ -12,13 +12,13 @@
 from vllm.transformers_utils.utils import without_trust_remote_code
 
 
-class SpeculatorsConfig(PretrainedConfig):
+class SpeculatorsConfig(PreTrainedConfig):
     model_type = "speculators"
 
     def __init__(self, **kwargs):
         # super().__init__ performs some validation before setting all kwargs as
         # attributes, so we set them first to be safe
-        pre_trained_config_fields = {f.name for f in fields(PretrainedConfig)}
+        pre_trained_config_fields = {f.name for f in fields(PreTrainedConfig)}
         super_kwargs = dict()
         for key, value in kwargs.items():
             if key == "model_type":
@@ -47,9 +47,7 @@ def from_pretrained(
     def extract_transformers_pre_trained_config(
         cls, config_dict: dict[str, Any]
     ) -> dict[str, Any]:
-        """
-        Extract standard Transformers PreTrainedConfig config from speculators config.
-        """
+        """Extract the Transformers `PreTrainedConfig` from a speculators config."""
         speculators_model_type = config_dict.get("speculators_model_type")
         if speculators_model_type not in SUPPORTED_SPECULATORS_TYPES:
             raise ValueError(
@@ -99,8 +97,7 @@ def validate_speculators_config(cls, config_dict: dict[str, Any]) -> None:
     def build_vllm_speculative_config(
         cls, config_dict: dict[str, Any]
     ) -> dict[str, Any]:
-        """
-        Build vLLM-compatible speculative configuration from speculators format.
+        """Build vLLM-compatible speculative configuration from speculators format.
 
         This method extracts and transforms speculative configuration from the
         speculators format into the structure expected by vLLM.
@@ -110,6 +107,7 @@ def build_vllm_speculative_config(
 
         Returns:
             Dictionary with vLLM-compatible speculative configuration
+
         """
         # Extract speculators configuration
         spec_config = config_dict["speculators_config"]
diff --git a/vllm/transformers_utils/configs/step3_vl.py b/vllm/transformers_utils/configs/step3_vl.py
index 0ee650a70451..c19030e55448 100644
--- a/vllm/transformers_utils/configs/step3_vl.py
+++ b/vllm/transformers_utils/configs/step3_vl.py
@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class Step3VisionEncoderConfig(PretrainedConfig):
+class Step3VisionEncoderConfig(PreTrainedConfig):
     model_type = "step3_vision_encoder"
 
     def __init__(
@@ -35,7 +35,7 @@ def __init__(
         super().__init__(**kwargs)
 
 
-class Step3TextConfig(PretrainedConfig):
+class Step3TextConfig(PreTrainedConfig):
     model_type = "step3_text"
     architectures = ["Step3TextForCausalLM"]
 
@@ -146,7 +146,7 @@ def __init__(
         super().__init__(**kwargs)
 
 
-class Step3VLConfig(PretrainedConfig):
+class Step3VLConfig(PreTrainedConfig):
     model_type = "step3_vl"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/step3p5.py b/vllm/transformers_utils/configs/step3p5.py
index d67b9dfff5f6..1b8ac82a7955 100644
--- a/vllm/transformers_utils/configs/step3p5.py
+++ b/vllm/transformers_utils/configs/step3p5.py
@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import Any
 
-from transformers.configuration_utils import PretrainedConfig
+from transformers.configuration_utils import PreTrainedConfig
 
 
-class Step3p5Config(PretrainedConfig):
+class Step3p5Config(PreTrainedConfig):
     model_type = "step3p5"
 
     def __init__(
diff --git a/vllm/transformers_utils/configs/ultravox.py b/vllm/transformers_utils/configs/ultravox.py
index 31b49b9d993f..669fe7be485a 100644
--- a/vllm/transformers_utils/configs/ultravox.py
+++ b/vllm/transformers_utils/configs/ultravox.py
@@ -7,15 +7,14 @@
 import transformers
 
 
-class UltravoxConfig(transformers.PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a
+class UltravoxConfig(transformers.PreTrainedConfig):
+    r"""This is the configuration class to store the configuration of a
     [`UltravoxForConditionalGeneration`]. It is used to instantiate an
     Ultravox model according to the specified arguments, defining the model
     architecture.
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to
-    control the model outputs. Read the documentation from [`PretrainedConfig`]
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to
+    control the model outputs. Read the documentation from [`PreTrainedConfig`]
     for more information.
 
     Args:
@@ -41,6 +40,7 @@ class UltravoxConfig(transformers.PretrainedConfig):
             Whether to apply layer normalization at the middle of the
             projector or at the end. Versions v0.4.1 and below
             use `False`, but v0.5 and above use `True`.
+
     """
 
     model_type = "ultravox"
@@ -74,7 +74,7 @@ def __init__(
         self.num_projector_layers = num_projector_layers
 
         # N.B. May set the wrapped_model_config below.
-        self.wrapped_model_config: transformers.PretrainedConfig
+        self.wrapped_model_config: transformers.PreTrainedConfig
         self.text_model_id = text_model_id
         if text_model_id is None:
             text_config = text_config or {}
@@ -113,7 +113,7 @@ def __setattr__(self, key, value):
         return super().__setattr__(key, value)
 
     @property
-    def text_config(self) -> transformers.PretrainedConfig:
+    def text_config(self) -> transformers.PreTrainedConfig:
         # When Ultravox wraps a multi-modal model (e.g. Gemma), we instantiate
         # the full model, but the text config is the text config of the inner
         # model.
diff --git a/vllm/transformers_utils/dynamic_module.py b/vllm/transformers_utils/dynamic_module.py
index f0702b15bb7e..f42d5752bf57 100644
--- a/vllm/transformers_utils/dynamic_module.py
+++ b/vllm/transformers_utils/dynamic_module.py
@@ -29,8 +29,7 @@ def try_get_class_from_dynamic_module(
     warn_on_fail: bool = True,
     **kwargs,
 ) -> type | None:
-    """
-    As `transformers.dynamic_module_utils.get_class_from_dynamic_module`,
+    """As `transformers.dynamic_module_utils.get_class_from_dynamic_module`,
     but ignoring any errors.
     """
     try:
diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py
index 72f2920480a1..9386e9a328fa 100644
--- a/vllm/transformers_utils/model_arch_config_convertor.py
+++ b/vllm/transformers_utils/model_arch_config_convertor.py
@@ -5,7 +5,7 @@
 
 import torch
 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
-from transformers import PretrainedConfig
+from transformers import PreTrainedConfig
 
 from vllm import envs
 from vllm.config.model_arch import (
@@ -26,8 +26,8 @@
 class ModelArchConfigConvertorBase:
     def __init__(
         self,
-        hf_config: PretrainedConfig,
-        hf_text_config: PretrainedConfig,
+        hf_config: PreTrainedConfig,
+        hf_text_config: PreTrainedConfig,
         revision: str | None = None,
     ):
         self.hf_config = hf_config
@@ -36,7 +36,7 @@ def __init__(
 
     def get_per_layer_hf_configs(
         self,
-    ) -> list[tuple[PretrainedConfig, PretrainedConfig]] | None:
+    ) -> list[tuple[PreTrainedConfig, PreTrainedConfig]] | None:
         """`(hf_config, hf_text_config)` per layer, or `None` if homogeneous.
 
         This is the only place that decides whether a checkpoint is heterogeneous
@@ -219,7 +219,7 @@ def get_num_experts_per_token(self) -> int:
     @classmethod
     def get_torch_dtype(
         cls,
-        hf_config: PretrainedConfig,
+        hf_config: PreTrainedConfig,
         model_id: str,
         revision: str | None,
         config_format: str | ConfigFormat,
@@ -257,7 +257,7 @@ def get_torch_dtype(
 
         return config_dtype
 
-    def _normalize_quantization_config(self, config: PretrainedConfig):
+    def _normalize_quantization_config(self, config: PreTrainedConfig):
         quant_cfg = getattr(config, "quantization_config", None)
         if quant_cfg is None:
             # compressed-tensors uses a "compression_config" key
@@ -585,8 +585,8 @@ def get_num_hidden_layers(self) -> int:
 class DeepseekV4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
     def __init__(
         self,
-        hf_config: PretrainedConfig,
-        hf_text_config: PretrainedConfig,
+        hf_config: PreTrainedConfig,
+        hf_text_config: PreTrainedConfig,
         revision: str | None = None,
     ):
         # DeepSeek-V4-Flash-Vision-Exp ships the same architectures/model_type
@@ -622,7 +622,7 @@ def get_num_hidden_layers(self) -> int:
 
 
 def _strip_mimo_v2_attention_chunk_size(
-    hf_config: PretrainedConfig, hf_text_config: PretrainedConfig
+    hf_config: PreTrainedConfig, hf_text_config: PreTrainedConfig
 ) -> None:
     # MiMo-V2-Flash's config.json sets `attention_chunk_size=128` but the
     # architecture does not actually use chunked local attention. Leaving it
@@ -635,8 +635,8 @@ def _strip_mimo_v2_attention_chunk_size(
 class MimoV2ModelArchConfigConvertor(ModelArchConfigConvertorBase):
     def __init__(
         self,
-        hf_config: PretrainedConfig,
-        hf_text_config: PretrainedConfig,
+        hf_config: PreTrainedConfig,
+        hf_text_config: PreTrainedConfig,
         revision: str | None = None,
     ):
         if getattr(hf_config, "vision_config", None):
@@ -648,8 +648,8 @@ def __init__(
 class MimoV2MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
     def __init__(
         self,
-        hf_config: PretrainedConfig,
-        hf_text_config: PretrainedConfig,
+        hf_config: PreTrainedConfig,
+        hf_text_config: PreTrainedConfig,
         revision: str | None = None,
     ):
         super().__init__(hf_config, hf_text_config, revision)
@@ -742,7 +742,7 @@ def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
 
     def get_per_layer_hf_configs(
         self,
-    ) -> list[tuple[PretrainedConfig, PretrainedConfig]] | None:
+    ) -> list[tuple[PreTrainedConfig, PreTrainedConfig]] | None:
         # Gemma4 uses a larger head dimension, and sometimes more KV heads, on its
         # full attention layers than on its sliding ones. Transformers >= 5.15.0
         # says so in the config; before that the values are flat attributes picked
@@ -763,7 +763,7 @@ def get_per_layer_hf_configs(
 
 
 class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase):
-    def _language_config(self) -> PretrainedConfig:
+    def _language_config(self) -> PreTrainedConfig:
         return self.hf_config.language_config
 
     def get_num_hidden_layers(self) -> int:
@@ -813,6 +813,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]:
     "dbrx": DbrxModelArchConfigConvertor,
     "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,
     "deepseek_v4": DeepseekV4ModelArchConfigConvertor,
+    "diffusion_gemma": Gemma4ModelArchConfigConvertor,
     "diffusion_gemma_text": Gemma4ModelArchConfigConvertor,
     "ernie_mtp": ErnieMTPModelArchConfigConvertor,
     "falcon": FalconModelArchConfigConvertor,
diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py
index 8fd924fcf381..64af034d9ac4 100644
--- a/vllm/transformers_utils/processor.py
+++ b/vllm/transformers_utils/processor.py
@@ -80,9 +80,7 @@ def __init__(self, *args, **kwargs):
 
 
 class HashableDict(dict):
-    """
-    A dictionary that can be hashed by lru_cache.
-    """
+    """A dictionary that can be hashed by lru_cache."""
 
     # NOTE: pythonic dict is not hashable,
     # we override on it directly for simplicity
@@ -91,9 +89,7 @@ def __hash__(self) -> int:  # type: ignore[override]
 
 
 class HashableList(list):
-    """
-    A list that can be hashed by lru_cache.
-    """
+    """A list that can be hashed by lru_cache."""
 
     def __hash__(self) -> int:  # type: ignore[override]
         return hash(tuple(self))
diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py
index 249a4e58e662..b46183f4e12a 100644
--- a/vllm/transformers_utils/processors/__init__.py
+++ b/vllm/transformers_utils/processors/__init__.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Multi-modal processors may be defined in this directory for the following
+"""Multi-modal processors may be defined in this directory for the following
 reasons:
 
 - There is no processing file defined by HF Hub or Transformers library.
diff --git a/vllm/transformers_utils/processors/bagel.py b/vllm/transformers_utils/processors/bagel.py
index 3226d7b0c83d..b6fefe032759 100644
--- a/vllm/transformers_utils/processors/bagel.py
+++ b/vllm/transformers_utils/processors/bagel.py
@@ -18,8 +18,7 @@ class BagelProcessorKwargs(ProcessingKwargs, total=False):  # type: ignore[call-
 
 
 class BagelProcessor(ProcessorMixin):
-    """
-    Constructs a BAGEL processor which wraps a
+    """Constructs a BAGEL processor which wraps a
     SigLIP image processor and a Qwen2 tokenizer.
     """
 
@@ -36,9 +35,7 @@ def __call__(
         images: ImageInput = None,
         **kwargs: Unpack[BagelProcessorKwargs],
     ):
-        """
-        Main method to prepare for the model one or several sequences(s) and image(s).
-        """
+        """Prepare one or more sequences and images for the model."""
         output_kwargs = self._merge_kwargs(
             BagelProcessorKwargs,
             tokenizer_init_kwargs=self.tokenizer.init_kwargs,
@@ -62,15 +59,11 @@ def __call__(
         return BatchFeature(data={**pixel_values, **text_inputs})
 
     def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's batch_decode.
-        """
+        """Forward all arguments to `Qwen2TokenizerFast.batch_decode`."""
         return self.tokenizer.batch_decode(*args, **kwargs)
 
     def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's decode.
-        """
+        """This method forwards all its arguments to Qwen2TokenizerFast's decode."""
         return self.tokenizer.decode(*args, **kwargs)
 
     @property
diff --git a/vllm/transformers_utils/processors/cohere_asr.py b/vllm/transformers_utils/processors/cohere_asr.py
index e1257de4e735..5665d3cbaeb0 100644
--- a/vllm/transformers_utils/processors/cohere_asr.py
+++ b/vllm/transformers_utils/processors/cohere_asr.py
@@ -228,7 +228,7 @@ def filter_banks(self):
         return self.fb
 
     def splice_frames(self, x, frame_splicing):
-        """Stacks frames together across feature dim
+        """Stacks frames together across feature dim.
 
         input is batch_size, feature_dim, num_frames
         output is batch_size, feature_dim*frame_splicing, num_frames
diff --git a/vllm/transformers_utils/processors/cosmos3_edge.py b/vllm/transformers_utils/processors/cosmos3_edge.py
index 95ca2cccd531..a8751cdf6d52 100644
--- a/vllm/transformers_utils/processors/cosmos3_edge.py
+++ b/vllm/transformers_utils/processors/cosmos3_edge.py
@@ -70,8 +70,7 @@ def _resize_image(
         resample: PILImageResampling | InterpolationMode | int | None = None,
         **kwargs,
     ) -> torch.Tensor:
-        """
-        Rescales the image so that the following conditions are met:
+        """Rescales the image so that the following conditions are met:
 
         1. Both dimensions (height and width) are divisible by 'factor'.
         2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
@@ -322,6 +321,7 @@ class Cosmos3EdgeProcessor(Qwen3VLProcessor):
             normalize, and patch video inputs.
         chat_template: An optional Jinja template used to render chat messages
             into a tokenizable string.
+
     """
 
     attributes = ["image_processor", "tokenizer", "video_processor"]
diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py
index 8a9e36448d83..55af43a70957 100644
--- a/vllm/transformers_utils/processors/deepseek_ocr.py
+++ b/vllm/transformers_utils/processors/deepseek_ocr.py
@@ -234,9 +234,7 @@ def process_one(
         images: list[Image.Image],
         crop_mode: bool = CROP_MODE,
     ):
-        """
-
-        Args:
+        """Args:
             prompt (str): the formatted prompt;
             images (List[ImageType]): the list of images;
             crop_mode (bool): if True, then crop the image;
@@ -248,8 +246,8 @@ def process_one(
                 - pixel_values (torch.FloatTensor): [n_patches, 3, H, W]
                 - image_id (int): the id of the image token
                 - num_image_tokens (List[int]): the number of image tokens
-        """
 
+        """
         if prompt is None or images is None:
             raise ValueError("prompt and images must be used at the same time.")
 
@@ -309,7 +307,6 @@ def tokenize_with_images(
         cropping: bool = True,
     ):
         """Tokenize text with <image> tags."""
-
         num_image_tags = conversation.count(self.image_token)
         if num_image_tags != len(images):
             raise ValueError(
diff --git a/vllm/transformers_utils/processors/deepseek_vl2.py b/vllm/transformers_utils/processors/deepseek_vl2.py
index be6d9800a918..6c9966abf8d1 100644
--- a/vllm/transformers_utils/processors/deepseek_vl2.py
+++ b/vllm/transformers_utils/processors/deepseek_vl2.py
@@ -186,9 +186,7 @@ def process_one(
         inference_mode: bool = True,
         **kwargs: Any,
     ):
-        """
-
-        Args:
+        """Args:
             prompt (str): the formatted prompt;
             images (list[ImageType]): the list of images;
             inference_mode (bool): if True, then remove the last eos token;
@@ -201,8 +199,8 @@ def process_one(
                 - pixel_values (torch.FloatTensor): [n_patches, 3, H, W]
                 - image_id (int): the id of the image token
                 - num_image_tokens (list[int]): the number of image tokens
-        """
 
+        """
         if prompt is None or images is None:
             raise ValueError("prompt and images must be used at the same time.")
 
@@ -276,9 +274,7 @@ def __call__(
         inference_mode: bool = True,
         **kwargs: Any,
     ):
-        """
-
-        Args:
+        """Args:
             text (str): the formatted prompt;
             images (list[ImageType]): the list of images;
             inference_mode (bool): if True, then remove the last eos token;
@@ -290,8 +286,8 @@ def __call__(
                 - images (torch.FloatTensor): [n_images, 3, H, W]
                 - image_id (int): the id of the image token
                 - num_image_tokens (list[int]): the number of image tokens
-        """
 
+        """
         prepare = self.process_one(
             prompt=text,
             images=images,
diff --git a/vllm/transformers_utils/processors/fireredasr2.py b/vllm/transformers_utils/processors/fireredasr2.py
index bba7e7ee0495..6225e88e6a70 100644
--- a/vllm/transformers_utils/processors/fireredasr2.py
+++ b/vllm/transformers_utils/processors/fireredasr2.py
@@ -68,8 +68,7 @@ def __call__(self, sample_rate, wav_np, is_train=False):
 
 
 class FireRedASR2FeatureExtractor(SequenceFeatureExtractor):
-    r"""
-    Constructs a FireRedASR2 feature extractor.
+    r"""Constructs a FireRedASR2 feature extractor.
 
     This feature extractor inherits from [`~feature_extraction_sequence_
         utils.SequenceFeatureExtractor`] which contains most of the main
@@ -99,6 +98,7 @@ class FireRedASR2FeatureExtractor(SequenceFeatureExtractor):
             Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces
             the high log_mel_fbank values for signals with hard-zero sections,
             when VAD cutoff is present in the signal.
+
     """
 
     model_input_names = ["input_features"]
@@ -230,9 +230,8 @@ def padding_position_is_0(padded_input, input_lengths):
 
 
 class FireRedASR2Processor(ProcessorMixin):
-    r"""
-    Constructs a FireRedASR2 processor which wraps a FireRedASR2 feature extractor and
-    a FireRedASR2 tokenizer into a single processor.
+    r"""Constructs a FireRedASR2 processor which wraps a FireRedASR2 feature extractor
+    and a FireRedASR2 tokenizer into a single processor.
 
     [`FireRedASR2Processor`] offers all the functionalities of
     [`FireRedASR2FeatureExtractor`] and [`Qwen2Tokenizer`]. See the
@@ -246,6 +245,7 @@ class FireRedASR2Processor(ProcessorMixin):
         tokenizer (`Qwen2Tokenizer`):
             An instance of [`Qwen2Tokenizer`]. The tokenizer is a required
             input.
+
     """
 
     feature_extractor_class = "FireRedASR2FeatureExtractor"
@@ -271,8 +271,7 @@ def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
         )
 
     def __call__(self, *args, **kwargs):
-        """
-        Forwards the `audio` argument to FireRedASR2FeatureExtractor's
+        """Forwards the `audio` argument to FireRedASR2FeatureExtractor's
         [`~FireRedASR2FeatureExtractor.__call__`] and the `text` argument to
         [`~Qwen2Tokenizer.__call__`]. Please refer to the docstring of the
         above two methods for more information.
diff --git a/vllm/transformers_utils/processors/funasr.py b/vllm/transformers_utils/processors/funasr.py
index d7a3c4060ceb..0a2f9ca6e14e 100644
--- a/vllm/transformers_utils/processors/funasr.py
+++ b/vllm/transformers_utils/processors/funasr.py
@@ -21,10 +21,7 @@
 
 
 def apply_cmvn(inputs, cmvn):  # noqa
-    """
-    Apply CMVN with mvn data
-    """
-
+    """Apply CMVN with mvn data."""
     device = inputs.device
     # dtype = inputs.dtype
     frame, dim = inputs.shape
@@ -167,60 +164,9 @@ def forward(
             feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
         return feats_pad, feats_lens
 
-    def forward_fbank(
-        self, input: torch.Tensor, input_lengths: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        batch_size = input.size(0)
-        feats = []
-        feats_lens = []
-        for i in range(batch_size):
-            waveform_length = input_lengths[i]
-            waveform = input[i][:waveform_length]
-            waveform = waveform * (1 << 15)
-            waveform = waveform.unsqueeze(0)
-            mat = kaldi.fbank(
-                waveform,
-                num_mel_bins=self.n_mels,
-                frame_length=self.frame_length,
-                frame_shift=self.frame_shift,
-                dither=self.dither,
-                energy_floor=0.0,
-                window_type=self.window,
-                sample_frequency=self.fs,
-            )
-
-            feat_length = mat.size(0)
-            feats.append(mat)
-            feats_lens.append(feat_length)
-
-        feats_lens = torch.as_tensor(feats_lens)
-        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
-        return feats_pad, feats_lens
-
-    def forward_lfr_cmvn(
-        self, input: torch.Tensor, input_lengths: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        batch_size = input.size(0)
-        feats = []
-        feats_lens = []
-        for i in range(batch_size):
-            mat = input[i, : input_lengths[i], :]
-            if self.lfr_m != 1 or self.lfr_n != 1:
-                mat = apply_lfr(mat, self.lfr_m, self.lfr_n)
-            if self.cmvn is not None:
-                mat = apply_cmvn(mat, self.cmvn)
-            feat_length = mat.size(0)
-            feats.append(mat)
-            feats_lens.append(feat_length)
-
-        feats_lens = torch.as_tensor(feats_lens)
-        feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0)
-        return feats_pad, feats_lens
-
 
 class FunASRFeatureExtractor(SequenceFeatureExtractor):
-    r"""
-    Constructs a FunASR feature extractor.
+    r"""Constructs a FunASR feature extractor.
 
     This feature extractor inherits from [`~feature_extraction_sequence_
         utils.SequenceFeatureExtractor`] which contains most of the main
@@ -255,6 +201,7 @@ class FunASRFeatureExtractor(SequenceFeatureExtractor):
             Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces
             the high log_mel_fbank values for signals with hard-zero sections,
             when VAD cutoff is present in the signal.
+
     """
 
     model_input_names = ["input_features"]
@@ -369,8 +316,7 @@ def __call__(
 
 
 class FunASRProcessor(ProcessorMixin):
-    r"""
-    Constructs a FunASR processor which wraps a FunASR feature extractor and
+    r"""Constructs a FunASR processor which wraps a FunASR feature extractor and
     a FunASR tokenizer into a single processor.
 
     [`FunASRProcessor`] offers all the functionalities of
@@ -385,6 +331,7 @@ class FunASRProcessor(ProcessorMixin):
         tokenizer (`Qwen2Tokenizer`):
             An instance of [`Qwen2Tokenizer`]. The tokenizer is a required
             input.
+
     """
 
     feature_extractor_class = "FunASRFeatureExtractor"
@@ -410,8 +357,7 @@ def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
         )
 
     def __call__(self, *args, **kwargs):
-        """
-        Forwards the `audio` argument to FunASRFeatureExtractor's
+        """Forwards the `audio` argument to FunASRFeatureExtractor's
         [`~FunASRFeatureExtractor.__call__`] and the `text` argument to
         [`~Qwen2Tokenizer.__call__`]. Please refer to the docstring of the
         above two methods for more information.
diff --git a/vllm/transformers_utils/processors/glm4v.py b/vllm/transformers_utils/processors/glm4v.py
index a8da23955259..605cd7e34a1c 100644
--- a/vllm/transformers_utils/processors/glm4v.py
+++ b/vllm/transformers_utils/processors/glm4v.py
@@ -4,14 +4,13 @@
 # Adapted from
 # https://github.com/zai-org/CogAgent
 from transformers import PythonBackend
-from transformers.image_processing_utils_fast import BaseImageProcessorFast
+from transformers.image_processing_backends import TorchvisionBackend
 from transformers.image_utils import PILImageResampling
 from transformers.processing_utils import ProcessorMixin
 
 
-class GLM4VImageProcessorFast(BaseImageProcessorFast):
-    """
-    Port of https://huggingface.co/zai-org/glm-4v-9b/blob/main/tokenization_chatglm.py#L177
+class GLM4VImageProcessorFast(TorchvisionBackend):
+    """Port of https://huggingface.co/zai-org/glm-4v-9b/blob/main/tokenization_chatglm.py#L177
     to HF Transformers.
     """
 
diff --git a/vllm/transformers_utils/processors/glm5next.py b/vllm/transformers_utils/processors/glm5next.py
index b6988d3b8fed..a6cba9d77f17 100644
--- a/vllm/transformers_utils/processors/glm5next.py
+++ b/vllm/transformers_utils/processors/glm5next.py
@@ -7,12 +7,12 @@
 import numpy as np
 import torch
 from torchvision.transforms.v2 import functional as tvF
-from transformers.image_processing_utils import BatchFeature
-from transformers.image_processing_utils_fast import (
-    BaseImageProcessorFast,
+from transformers.image_processing_backends import (
+    TorchvisionBackend,
     group_images_by_shape,
     reorder_images,
 )
+from transformers.image_processing_utils import BatchFeature
 from transformers.image_utils import (
     OPENAI_CLIP_MEAN,
     OPENAI_CLIP_STD,
@@ -327,7 +327,7 @@ class Glm5NextImageProcessorKwargs(ImagesKwargs, total=False):  # type: ignore[c
     max_image_tokens: int | None
 
 
-class Glm5NextImageProcessor(BaseImageProcessorFast):
+class Glm5NextImageProcessor(TorchvisionBackend):
     """Fast torchvision image processor for GLM-5.3-Flash.
 
     ``patch_expand_factor`` multiplies into the ``smart_resize`` spatial
@@ -759,7 +759,7 @@ class Glm5NextProcessor(ProcessorMixin):
     attributes = ["image_processor", "tokenizer", "video_processor"]
     image_processor_class = "AutoImageProcessor"
     video_processor_class = "AutoVideoProcessor"
-    tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
+    tokenizer_class = ("PythonBackend", "TokenizersBackend")
 
     def __init__(
         self,
diff --git a/vllm/transformers_utils/processors/internvl.py b/vllm/transformers_utils/processors/internvl.py
index 92c1834c74ee..84976cb178ce 100644
--- a/vllm/transformers_utils/processors/internvl.py
+++ b/vllm/transformers_utils/processors/internvl.py
@@ -358,8 +358,7 @@ def __call__(
 
 
 class InternVLProcessor(ProcessorMixin):
-    """
-    This model doesn't define its own HF processor,
+    """This model doesn't define its own HF processor,
     so we implement our own one here.
 
     The code to insert image tokens is based on:
diff --git a/vllm/transformers_utils/processors/isaac.py b/vllm/transformers_utils/processors/isaac.py
index 2d791df01673..7d10e9be86e9 100644
--- a/vllm/transformers_utils/processors/isaac.py
+++ b/vllm/transformers_utils/processors/isaac.py
@@ -91,6 +91,7 @@ def get_image_size_for_max_num_patches(
         `tuple[int, int]`: Height and width (in pixels) that are multiples of
         `patch_size * pixel_shuffle_scale` and respect both the maximum and
         optional minimum patch-count constraints.
+
     """
 
     def get_scaled_image_size(scale, original_size, patch_size, pixel_shuffle_scale):
@@ -177,9 +178,11 @@ def prepare_image_tensor(
             The tensor is converted to floating point if needed.
         scale (`float`, *optional*, defaults to `VISION_SCALE`):
             Scalar multiplier applied before normalization.
+
     Returns:
         `torch.Tensor`: Normalized tensor with the same shape as the input and
         dtype `torch.float32`.
+
     """
     if not torch.is_floating_point(image):
         image = image.float()
@@ -209,6 +212,7 @@ def patchify_vision(image: torch.Tensor, patch_size: int) -> torch.Tensor:
 
     Raises:
         ValueError: If `height` or `width` is not divisible by `patch_size`.
+
     """
     num_images, height, width, channels = image.shape
     if height % patch_size or width % patch_size:
@@ -265,6 +269,7 @@ def process_vision_for_patches(
         / patch_size, channels * patch_size**2)` and `dims_virtual` encodes
         effective `(images, height, width)` dimensions after optional pixel
         shuffling.
+
     """
     # Add batch dim if single image
     if images.dim() == 3:
diff --git a/vllm/transformers_utils/processors/kimi_k25.py b/vllm/transformers_utils/processors/kimi_k25.py
index edee9734ce42..13aed85de4ec 100644
--- a/vllm/transformers_utils/processors/kimi_k25.py
+++ b/vllm/transformers_utils/processors/kimi_k25.py
@@ -28,8 +28,7 @@ def __call__(
         return_tensors: str | TensorType | None = None,
         **kwargs,
     ) -> BatchFeature:
-        """
-        Args:
+        """Args:
             text: The text to be field to the model.
             vision_chunks: List of `VisionChunk` items to be processed.
                 For image: `VisionChunkImage` with
@@ -44,6 +43,7 @@ def __call__(
               Returned when `vision_chunks` is not `None`.
             - **grid_thws** -- list of image 3D grid in LLM.
               Returned when `vision_chunks` is not `None`.
+
         """
         if vision_chunks is not None:
             mm_inputs = self.image_processor.preprocess(
diff --git a/vllm/transformers_utils/processors/mimo_v2_omni.py b/vllm/transformers_utils/processors/mimo_v2_omni.py
index 22349f751364..0555c941b2b4 100644
--- a/vllm/transformers_utils/processors/mimo_v2_omni.py
+++ b/vllm/transformers_utils/processors/mimo_v2_omni.py
@@ -1071,6 +1071,8 @@ def __call__(
                     tuples (pre-decoded).
             audio: list of ``str`` (path/url/base64), ``bytes``,
                    ``(waveform_1D, sample_rate)`` tuples, or ``np.ndarray``.
+            video_audio: audio track extracted from ``videos``, in the same
+                   formats accepted by ``audio``.
             return_tensors: Passed to :class:`BatchFeature`.
 
         Returns:
@@ -1079,6 +1081,7 @@ def __call__(
             - ``pixel_values`` + ``image_grid_thw``
             - ``pixel_values_videos`` + ``video_grid_thw`` + ``second_per_grid_ts``
             - ``audio_features``
+
         """
         if isinstance(text, list):
             text = text[0] if len(text) == 1 else "\n".join(text)
diff --git a/vllm/transformers_utils/processors/minicpmo.py b/vllm/transformers_utils/processors/minicpmo.py
index d9ac841cd86b..17a915b2d6b4 100644
--- a/vllm/transformers_utils/processors/minicpmo.py
+++ b/vllm/transformers_utils/processors/minicpmo.py
@@ -14,17 +14,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
-Processor class for MiniCPMO.
-"""
+"""Processor class for MiniCPMO."""
 
 import math
-from typing import Literal, TypeAlias
+from typing import TypeAlias
 
 import numpy as np
 import regex
 import torch
-import torchaudio
 from transformers.image_processing_utils import BatchFeature
 from transformers.image_utils import ImageInput
 from transformers.processing_utils import ProcessorMixin
@@ -35,8 +32,7 @@
 
 
 class MiniCPMOProcessor(ProcessorMixin):
-    r"""
-    Constructs a MiniCPMV processor which wraps a MiniCPMV image
+    r"""Constructs a MiniCPMV processor which wraps a MiniCPMV image
     processor and a MiniCPMV tokenizer into a single processor.
 
     [`MiniCPMVProcessor`] offers all the functionalities of
@@ -49,6 +45,7 @@ class MiniCPMOProcessor(ProcessorMixin):
             The image processor is a required input.
         tokenizer ([`LlamaTokenizerWrapper`], *optional*):
             The tokenizer is a required input.
+
     """
 
     attributes = ["image_processor", "feature_extractor", "tokenizer"]
@@ -253,8 +250,7 @@ def audio_feature_extract(
     # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode
     # with CLIP->Llama
     def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to LlamaTokenizerFast's
+        """This method forwards all its arguments to LlamaTokenizerFast's
         [`~PythonBackend.batch_decode`]. Please refer to the
         docstring of this method for more information.
         """
@@ -274,8 +270,7 @@ def batch_decode(self, *args, **kwargs):
     # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode
     # with CLIP->Llama
     def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to LlamaTokenizerFast's
+        """This method forwards all its arguments to LlamaTokenizerFast's
         [`~PythonBackend.decode`]. Please refer to the docstring
         of this method for more information.
         """
@@ -526,70 +521,3 @@ def pad(
             padding_length.append(tensor.shape[-1] - len(item))
 
         return tensor, padding_length
-
-
-class MelSpectrogramFeatures(torch.nn.Module):
-    def __init__(
-        self,
-        sample_rate=24000,
-        n_fft=1024,
-        hop_length=256,
-        n_mels=100,
-        padding: Literal["center", "same"] = "center",
-    ):
-        super().__init__()
-        if padding not in ["center", "same"]:
-            raise ValueError("Padding must be 'center' or 'same'.")
-        self.padding = padding
-        self.mel_spec = torchaudio.transforms.MelSpectrogram(
-            sample_rate=sample_rate,
-            n_fft=n_fft,
-            hop_length=hop_length,
-            n_mels=n_mels,
-            center=padding == "center",
-            power=1,
-        )
-
-    def __call__(self, audio: torch.Tensor) -> torch.Tensor:
-        """
-        audio: Tensor([num_channels, num_samples])
-        """
-        return super().__call__(audio)
-
-    def forward(self, audio: torch.Tensor) -> torch.Tensor:
-        """
-        audio: Tensor([num_channels, num_samples])
-        """
-        mel: torch.Tensor = self.mel_spec(audio)
-        features = torch.log(torch.clip(mel, min=1e-5))
-        return features
-
-
-class ChatTTSProcessor:
-    def __init__(self, text_tokenizer):
-        self.audio_processor = MelSpectrogramFeatures()
-        self.text_tokenizer = text_tokenizer
-
-    def __call__(self, text_list, audio_list):
-        assert len(text_list) == len(audio_list)
-        input_ids_varlen = []
-        for text in text_list:
-            input_ids_ = self.text_tokenizer.encode(
-                text, return_tensors="pt", add_special_tokens=False
-            )  # [1, seq_len]
-            input_ids_ = input_ids_.squeeze(0)  # [seq_len]
-            input_ids_varlen.append(input_ids_)
-
-        audio_features_varlen = []
-        for audio in audio_list:
-            assert audio.shape.__len__() == 1  # [seq_len]
-            try:
-                mel = self.audio_processor(audio)  # [100(num_mel_bins), seq_len_mel]
-            except Exception as e:
-                raise e
-            audio_features_varlen.append(mel)
-
-        return {
-            "tts_input_ids_varlen": input_ids_varlen,  # return List[Tensor]
-            "tts_input_features_varlen": audio_features_varlen,  # return List[Tensor]
-        }
diff --git a/vllm/transformers_utils/processors/minicpmv.py b/vllm/transformers_utils/processors/minicpmv.py
index 712cbbb47a01..aaedaa376780 100644
--- a/vllm/transformers_utils/processors/minicpmv.py
+++ b/vllm/transformers_utils/processors/minicpmv.py
@@ -13,9 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
-Processor class for MiniCPMV.
-"""
+"""Processor class for MiniCPMV."""
 
 from typing import TypeAlias
 
@@ -36,8 +34,7 @@
 
 
 class MiniCPMVProcessor(ProcessorMixin):
-    r"""
-    Constructs a MiniCPMV processor which wraps a MiniCPMV image
+    r"""Constructs a MiniCPMV processor which wraps a MiniCPMV image
     processor and a MiniCPMV tokenizer into a single processor.
 
     [`MiniCPMVProcessor`] offers all the functionalities of
@@ -50,6 +47,7 @@ class MiniCPMVProcessor(ProcessorMixin):
             The image processor is a required input.
         tokenizer ([`LlamaTokenizerWrapper`], *optional*):
             The tokenizer is a required input.
+
     """
 
     attributes = ["image_processor", "tokenizer"]
@@ -93,8 +91,7 @@ def __call__(
     # Copied from transformers.models.clip.processing_clip.CLIPProcessor
     # .batch_decode with CLIP->Llama
     def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to LlamaTokenizerFast's
+        """This method forwards all its arguments to LlamaTokenizerFast's
         [`~PythonBackend.batch_decode`]. Please refer to the
         docstring of this method for more information.
         """
@@ -126,8 +123,7 @@ def batch_decode(self, *args, **kwargs):
     # Copied from transformers.models.clip.processing_clip.CLIPProcessor
     # .decode with CLIP->Llama
     def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to LlamaTokenizerFast's
+        """This method forwards all its arguments to LlamaTokenizerFast's
         [`~PythonBackend.decode`]. Please refer to the docstring
         of this method for more information.
         """
diff --git a/vllm/transformers_utils/processors/minimax_m3.py b/vllm/transformers_utils/processors/minimax_m3.py
index 60fa7ed536b7..237b1b5f1e1b 100644
--- a/vllm/transformers_utils/processors/minimax_m3.py
+++ b/vllm/transformers_utils/processors/minimax_m3.py
@@ -18,8 +18,8 @@
 import torch
 from torchvision.transforms import InterpolationMode
 from transformers import AutoTokenizer, BatchFeature
-from transformers.image_processing_utils_fast import (
-    BaseImageProcessorFast,
+from transformers.image_processing_backends import (
+    TorchvisionBackend,
     group_images_by_shape,
     reorder_images,
 )
@@ -156,7 +156,7 @@ class MiniMaxM3VLImageProcessorKwargs(ImagesKwargs, total=False):  # type: ignor
     max_long_side_pixel: int
 
 
-class MiniMaxM3VLImageProcessor(BaseImageProcessorFast):
+class MiniMaxM3VLImageProcessor(TorchvisionBackend):
     do_resize = True
     resample = PILImageResampling.BICUBIC
     # required by base-class validation, not used as the resize bound
@@ -522,9 +522,14 @@ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
         # register() API now stores classes as {"pil": cls} dicts in
         # _extra_content, but get_possibly_dynamic_module() still calls
         # .__name__ on the raw value, crashing with AttributeError on dicts.
-        tokenizer = AutoTokenizer.from_pretrained(
-            pretrained_model_name_or_path, **kwargs
-        )
+        #
+        # Reuse the tokenizer passed by vLLM instead of loading another one,
+        # and keep it out of the kwargs forwarded to the sub-processor loaders.
+        tokenizer = kwargs.pop("tokenizer", None)
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(
+                pretrained_model_name_or_path, **kwargs
+            )
         image_processor = MiniMaxM3VLImageProcessor.from_pretrained(
             pretrained_model_name_or_path, **kwargs
         )
diff --git a/vllm/transformers_utils/processors/moondream3.py b/vllm/transformers_utils/processors/moondream3.py
index ae6833b20079..321fc9517172 100644
--- a/vllm/transformers_utils/processors/moondream3.py
+++ b/vllm/transformers_utils/processors/moondream3.py
@@ -67,8 +67,7 @@ def select_tiling(
 
 
 class Moondream3Processor(ProcessorMixin):
-    """
-    Constructs a Moondream3 processor which handles image preprocessing
+    """Constructs a Moondream3 processor which handles image preprocessing
     and tokenization for the Moondream3 multimodal model.
 
     Args:
@@ -78,6 +77,7 @@ class Moondream3Processor(ProcessorMixin):
         max_crops: Maximum number of crops per image.
         overlap_margin: Margin for overlapping crops in patches.
         patch_size: Size of each patch.
+
     """
 
     attributes = ["tokenizer"]
@@ -189,8 +189,7 @@ def from_pretrained(
         pretrained_model_name_or_path,
         **kwargs,
     ):
-        """
-        Load the processor, using a separate tokenizer repo.
+        """Load the processor, using a separate tokenizer repo.
 
         The moondream3 model uses a custom tokenizer from 'moondream/starmie-v1'
         instead of having tokenizer files in the model repo.
@@ -210,7 +209,6 @@ def from_pretrained(
             "revision",
             "subfolder",
             "token",
-            "use_fast",
         ):
             if key in kwargs:
                 tokenizer_kwargs[key] = kwargs[key]
@@ -294,8 +292,7 @@ def __call__(
         | list[PreTokenizedInput] = None,
         **kwargs: Unpack[Moondream3ProcessorKwargs],
     ) -> BatchFeature:
-        """
-        Process images and text for Moondream3 model.
+        """Process images and text for Moondream3 model.
 
         Args:
             images: Input images (PIL Image, numpy array, or list thereof).
@@ -304,6 +301,7 @@ def __call__(
 
         Returns:
             BatchFeature with processed inputs.
+
         """
         output_kwargs = self._merge_kwargs(
             Moondream3ProcessorKwargs,
@@ -428,8 +426,7 @@ def preprocess_image(
         convert_to_rgb: bool = True,
         return_tensors: str = "pt",
     ) -> tuple[torch.Tensor, tuple[int, int]]:
-        """
-        Preprocess an image using overlap-and-resize cropping strategy.
+        """Preprocess an image using overlap-and-resize cropping strategy.
 
         Args:
             image: Input PIL image, numpy array, or torch tensor.
@@ -442,6 +439,7 @@ def preprocess_image(
 
         Returns:
             Tuple of (pixel_values tensor, tiling tuple).
+
         """
         image = self._to_pil_image(image)
         if convert_to_rgb:
diff --git a/vllm/transformers_utils/processors/muse_glimmer.py b/vllm/transformers_utils/processors/muse_glimmer.py
index c1a5326ea1f1..baacbcf0c184 100644
--- a/vllm/transformers_utils/processors/muse_glimmer.py
+++ b/vllm/transformers_utils/processors/muse_glimmer.py
@@ -437,7 +437,7 @@ class MuseGlimmerProcessor(ProcessorMixin):
     attributes = ["image_processor", "video_processor", "tokenizer"]
     image_processor_class = "AutoImageProcessor"
     video_processor_class = "AutoVideoProcessor"
-    tokenizer_class = "PreTrainedTokenizerFast"
+    tokenizer_class = "TokenizersBackend"
 
     def __init__(
         self,
diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py
index 3cddbec6c866..9572b7c4a76d 100644
--- a/vllm/transformers_utils/processors/nano_nemotron_vl.py
+++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py
@@ -20,7 +20,7 @@
 import regex as re
 import torch
 from PIL import Image
-from transformers import BatchFeature, PretrainedConfig, TensorType
+from transformers import BatchFeature, PreTrainedConfig, TensorType
 
 from vllm.model_executor.models.parakeet import ParakeetExtractor
 from vllm.multimodal.inputs import AudioItem
@@ -292,8 +292,7 @@ def width_and_height_for_max_num_tokens_available(
         self,
         target_num_tokens_post_shuffle: int,
     ) -> tuple[int, int]:
-        """
-        TODO: optimize this so it squeezes closer to target number of tokens.
+        """TODO: optimize this so it squeezes closer to target number of tokens.
         Calculate image dimensions that produce approximately `target` tokens after
         pixel_shuffle.
 
@@ -318,6 +317,7 @@ def width_and_height_for_max_num_tokens_available(
         ...     height // PATCH_SIZE
         ... ) // 2**2 == 8100  # tokens post-shuffle
         >>> assert tiler._get_num_embeddings(width=width, height=height) == 8100
+
         """
         side_pixels = (
             math.isqrt(target_num_tokens_post_shuffle)
@@ -386,6 +386,7 @@ def process_media(
             num_tokens_available: Number of tokens available for this media
         Returns:
             DynamicResolutionParams for the media
+
         """
         current_num_tokens_available = num_tokens_available
         assert isinstance(media, Image.Image), (
@@ -473,6 +474,7 @@ def compute_params(
             num_tokens_available: Total number of tokens available across all media
         Returns:
             List of ImageTilingParams for each media item
+
         """
         num_tokens_available = (
             num_tokens_available
@@ -549,7 +551,9 @@ def compute_params(
         )
 
     @staticmethod
-    def stack(images: list[torch.Tensor], patch_size: int) -> torch.Tensor:
+    def stack(
+        images: list[torch.Tensor] | torch.Tensor, patch_size: int
+    ) -> torch.Tensor:
         assert len(images) > 0, "No images to stack"
 
         def rearrange_img(x):
@@ -571,8 +575,7 @@ def rearrange_img(x):
 
 
 class BaseNanoNemotronVLProcessor(ABC):
-    """
-    This model doesn't define its own HF processor,
+    """This model doesn't define its own HF processor,
     so we implement our own one here.
 
     The code to insert image tokens is based on:
@@ -581,7 +584,7 @@ class BaseNanoNemotronVLProcessor(ABC):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         tokenizer: HfTokenizer,
         *args,
         max_model_len: int,
@@ -620,7 +623,7 @@ def __init__(
         self.dtype: torch.dtype = getattr(config, "dtype", torch.float32)
 
     @staticmethod
-    def use_dynamic_resolution(config: PretrainedConfig) -> bool:
+    def use_dynamic_resolution(config: PreTrainedConfig) -> bool:
         return "min_num_patches" in config.vision_config.args
 
     @property
@@ -761,15 +764,14 @@ def __call__(
 
 
 class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
-    """
-    HF Processor with extended video processing logic.
+    """HF Processor with extended video processing logic.
     Code for video processing is adapted from video example:
     https://huggingface.co/OpenGVLab/InternVL3-1B#inference-with-transformers
     """
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        config: PreTrainedConfig,
         tokenizer: HfTokenizer,
         *,
         max_model_len: int,
@@ -1125,8 +1127,7 @@ def get_video_repl(
         img_context_token_ids: list[int],
         video_temporal_patch_size: int = 1,
     ) -> PromptUpdateDetails:
-        """
-        Build prompt replacement for a video.
+        """Build prompt replacement for a video.
         The replacement returned is not actually used to replace the placeholder
         tokens - it's just used to make sure we allocate the correct number
         of tokens.
diff --git a/vllm/transformers_utils/processors/nemotron_vl.py b/vllm/transformers_utils/processors/nemotron_vl.py
index 9c5436aa2729..9504f5ee2ba5 100644
--- a/vllm/transformers_utils/processors/nemotron_vl.py
+++ b/vllm/transformers_utils/processors/nemotron_vl.py
@@ -191,8 +191,7 @@ def _images_to_pixel_values_lst(
 
 
 class LlamaNemotronNanoVLProcessor(InternVLProcessor):
-    """
-    This model doesn't define its own HF processor,
+    """This model doesn't define its own HF processor,
     so we implement our own one here.
 
     The image processor is given by:
@@ -287,8 +286,7 @@ def _images_to_pixel_values_lst(
 
 
 class LlamaNemotronVLEmbedProcessor(InternVLProcessor):
-    """
-    Processor for LlamaNemotronVL embedding model.
+    """Processor for LlamaNemotronVL embedding model.
 
     Inherits from NemotronVLProcessor and specializes it for embedding tasks:
     - Uses SigLIP transform with normalization instead of base transform
diff --git a/vllm/transformers_utils/processors/ovis.py b/vllm/transformers_utils/processors/ovis.py
index 907c68035fcf..9ad320be3935 100644
--- a/vllm/transformers_utils/processors/ovis.py
+++ b/vllm/transformers_utils/processors/ovis.py
@@ -50,10 +50,10 @@ class OvisProcessorKwargs(ProcessingKwargs, total=False):  # type: ignore[call-a
 
 
 class OvisProcessor(ProcessorMixin):
-    r"""
-    Constructs an Ovis processor which wraps an Ovis image processor and a Qwen2 tokenizer into a single processor.
+    r"""Constructs an Ovis processor which wraps an Ovis image processor and a Qwen2 tokenizer into a single processor.
     [`OvisProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the
     [`~OvisProcessor.__call__`] and [`~OvisProcessor.decode`] for more information.
+
     Args:
         image_processor ([`Qwen2VLImageProcessor`], *optional*):
             The image processor is a required input.
@@ -61,6 +61,7 @@ class OvisProcessor(ProcessorMixin):
             The tokenizer is a required input.
         chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
             in a chat into a tokenizable string.
+
     """
 
     attributes = ["image_processor", "tokenizer"]
@@ -107,12 +108,12 @@ def __call__(
         | list[PreTokenizedInput] = None,
         **kwargs: Unpack[OvisProcessorKwargs],
     ) -> BatchFeature:
-        """
-        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+        """Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
         and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
-        the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to
+        the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwargs` arguments to
         Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`.
-            Args:
+
+        Args:
                 images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
                     The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                     tensor. Both channels-first and channels-last formats are supported.
@@ -129,7 +130,8 @@ def __call__(
                     - `'pt'`: Return PyTorch `torch.Tensor` objects.
                     - `'np'`: Return NumPy `np.ndarray` objects.
                     - `'jax'`: Return JAX `jnp.ndarray` objects.
-            Returns:
+
+        Returns:
                 [`BatchFeature`]: A [`BatchFeature`] with the following fields:
                 - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
                 - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
@@ -140,8 +142,8 @@ def __call__(
                 - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
                 - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`.
                 - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`.
-        """
 
+        """
         max_partition = kwargs.pop("max_partition", 9)
         covering_threshold = kwargs.pop("covering_threshold", 0.9)
 
@@ -416,28 +418,28 @@ def _get_best_grid(img, side):
         return torch.tensor(pixel_values), image_placeholders, torch.tensor(grid)
 
     def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.batch_decode`]. Please
+        """This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.batch_decode`]. Please
         refer to the docstring of this method for more information.
         """
         return self.tokenizer.batch_decode(*args, **kwargs)
 
     def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.decode`]. Please refer to
+        """This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.decode`]. Please refer to
         the docstring of this method for more information.
         """
         return self.tokenizer.decode(*args, **kwargs)
 
     def post_process_image_text_to_text(self, generated_outputs):
-        """
-        Post-process the output of the model to decode the text.
+        """Post-process the output of the model to decode the text.
+
         Args:
             generated_outputs (`torch.Tensor` or `np.ndarray`):
                 The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
                 or `(sequence_length,)`.
+
         Returns:
             `list[str]`: The decoded text.
+
         """
         return self.tokenizer.batch_decode(
             generated_outputs,
diff --git a/vllm/transformers_utils/processors/ovis2_5.py b/vllm/transformers_utils/processors/ovis2_5.py
index b9c2d112409c..2d5764bfd03a 100644
--- a/vllm/transformers_utils/processors/ovis2_5.py
+++ b/vllm/transformers_utils/processors/ovis2_5.py
@@ -33,13 +33,13 @@ class Ovis2_5ProcessorKwargs(ProcessingKwargs, total=False):  # type: ignore[cal
 
 
 class Ovis2_5Processor(ProcessorMixin):
-    r"""
-    Constructs an Ovis processor which wraps an Ovis image processor
+    r"""Constructs an Ovis processor which wraps an Ovis image processor
     and a Qwen2 tokenizer into a single processor.
     [`OvisProcessor`] offers all the functionalities of
     [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`].
     See the [`~OvisProcessor.__call__`] and [`~OvisProcessor.decode`]
     for more information.
+
     Args:
         image_processor ([`Qwen2VLImageProcessor`], *optional*):
             The image processor is a required input.
@@ -48,6 +48,7 @@ class Ovis2_5Processor(ProcessorMixin):
         chat_template (`str`, *optional*): A Jinja template which will
             be used to convert lists of messages in a chat into
             a tokenizable string.
+
     """
 
     attributes = ["image_processor", "tokenizer"]
@@ -110,43 +111,39 @@ def __call__(
         | list[PreTokenizedInput] = None,
         **kwargs: Unpack[Ovis2_5ProcessorKwargs],
     ) -> BatchFeature:
-        """
-        Main method to prepare for the model one or several sequences(s)
+        """Main method to prepare for the model one or several sequences(s)
         and image(s). This method forwards the `text`and `kwargs` arguments
         to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text`
         is not `None` to encode the text. To prepare the vision inputs,
-        this method forwards the `vision_infos` and `kwrags` arguments to
+        this method forwards the `vision_infos` and `kwargs` arguments to
         Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`]
         if `vision_infos` is not `None`.
-            Args:
-                images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`,
-                    `list[PIL.Image.Image]`, `list[np.ndarray]`,
-                    `list[torch.Tensor]`):
-                    The image or batch of images to be prepared.
-                    Each image can be a PIL image, NumPy array or PyTorch
-                    tensor. Both channels-first and channels-last formats
-                    are supported.
-                text (`str`, `list[str]`, `list[list[str]]`):
-                    The sequence or batch of sequences to be encoded.
-                    Each sequence can be a string or a list of strings
-                    (pretokenized string). If the sequences are provided as
-                    list of strings (pretokenized), you must set
-                    `is_split_into_words=True` (to lift the ambiguity with
-                    a batch of sequences).
-                videos (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`,
-                    `list[torch.Tensor]`):
-                    The image or batch of videos to be prepared. Each video
-                    can be a 4D NumPy array or PyTorch tensor, or a nested
-                    list of 3D frames. Both channels-first and channels-last
-                    formats are supported.
-                return_tensors (`str` or [`~utils.TensorType`], *optional*):
-                    If set, will return tensors of a particular framework.
-                    Acceptable values are:
-                    - `'tf'`: Return TensorFlow `tf.constant` objects.
-                    - `'pt'`: Return PyTorch `torch.Tensor` objects.
-                    - `'np'`: Return NumPy `np.ndarray` objects.
-                    - `'jax'`: Return JAX `jnp.ndarray` objects.
-            Returns:
+
+        Args:
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, or a list):
+                The image or batch of images to be prepared. Each image can
+                be a PIL image, NumPy array or PyTorch tensor. Both
+                channels-first and channels-last formats are supported.
+            videos (`np.ndarray`, `torch.Tensor`, or a list of either):
+                The video or batch of videos to be prepared. Each video can
+                be a 4D NumPy array or PyTorch tensor, or a nested list of
+                3D frames. Both channels-first and channels-last formats are
+                supported.
+            text (`str`, `list[str]`, `list[list[str]]`):
+                The sequence or batch of sequences to be encoded. Each
+                sequence can be a string or a list of strings (pretokenized
+                string). If the sequences are provided as list of strings
+                (pretokenized), you must set `is_split_into_words=True` (to
+                lift the ambiguity with a batch of sequences).
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework.
+                Acceptable values are:
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
                 [`BatchFeature`]: A [`BatchFeature`] with the following fields:
                 - **input_ids** -- list of token ids to be fed to a model.
                   Returned when `text` is not `None`.
@@ -164,6 +161,7 @@ def __call__(
                   when `videos` is not `None`.
                 - **second_per_grid_ts** -- list of video seconds per time grid.
                   Returned when `videos` is not `None`.
+
         """
         output_kwargs = self._merge_kwargs(
             Ovis2_5ProcessorKwargs,
diff --git a/vllm/transformers_utils/processors/pixtral.py b/vllm/transformers_utils/processors/pixtral.py
index 588ad6b7fad8..85604eb5b71d 100644
--- a/vllm/transformers_utils/processors/pixtral.py
+++ b/vllm/transformers_utils/processors/pixtral.py
@@ -11,8 +11,7 @@
 
 
 class MistralCommonImageProcessor(ImageProcessingMixin):
-    """
-    Provide a HF-compatible interface for
+    """Provide a HF-compatible interface for
     `mistral_common.tokens.tokenizers.multimodal.ImageEncoder`.
     """
 
diff --git a/vllm/transformers_utils/processors/voxtral.py b/vllm/transformers_utils/processors/voxtral.py
index 5403a7e47f99..4cc23c47754f 100644
--- a/vllm/transformers_utils/processors/voxtral.py
+++ b/vllm/transformers_utils/processors/voxtral.py
@@ -18,8 +18,7 @@
 
 
 class MistralCommonFeatureExtractor(SequenceFeatureExtractor):
-    """
-    Provide a HF-compatible interface for
+    """Provide a HF-compatible interface for
     `mistral_common.tokens.tokenizers.multimodal.AudioEncoder`.
     """
 
diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py
index ca20390aad51..b6a23b007aa2 100644
--- a/vllm/transformers_utils/repo_utils.py
+++ b/vllm/transformers_utils/repo_utils.py
@@ -63,9 +63,15 @@ def resolve_revision(
     requested revision that also carries the commit hash, so downstream error
     messages stay readable and offline loads reuse the cached `refs/` entry.
 
+    A commit hash only means something for the repo it was resolved against, so
+    a repo that is not `repo_id` needs its own call. Passing a revision resolved
+    for another repo back in is safe: the revision initially requested is
+    resolved again, against `repo_id` this time.
+
     Returns:
         The resolved revision, or `revision` unchanged if it cannot be resolved
         (local path, ModelScope, or any Hub error).
+
     """
     if Path(repo_id).exists() or envs.VLLM_USE_MODELSCOPE:
         return revision
@@ -106,6 +112,7 @@ def with_retry(
         retry_delay: Seconds to wait after the first failure, doubled each time.
         fatal_errors: Exceptions that are a definitive answer rather than a
             transient failure. These are raised immediately, without logging.
+
     """
     for attempt in range(max_retries):
         try:
@@ -353,8 +360,7 @@ def get_hf_file_bytes(
 def try_get_local_file(
     model: str | Path, file_name: str, revision: str | None = "main"
 ) -> Path | Any | None:
-    """
-    Try to get a local file from the HuggingFace repository.
+    """Try to get a local file from the HuggingFace repository.
 
     The possible return values are:
 
@@ -385,20 +391,19 @@ def try_get_local_file(
 def get_hf_file_to_dict(
     file_name: str, model: str | Path, revision: str | None = "main"
 ):
-    """
-    Downloads a file from the Hugging Face Hub and returns
+    """Downloads a file from the Hugging Face Hub and returns
     its contents as a dictionary.
 
-    Parameters:
-    - file_name (str): The name of the file to download.
-    - model (str): The name of the model on the Hugging Face Hub.
-    - revision (str): The specific version of the model.
+    Args:
+        file_name (str): The name of the file to download.
+        model (str): The name of the model on the Hugging Face Hub.
+        revision (str): The specific version of the model.
 
     Returns:
     - config_dict (dict): A dictionary containing
     the contents of the downloaded file.
-    """
 
+    """
     file_path = try_get_local_file(model=model, file_name=file_name, revision=revision)
 
     if file_path is None:
diff --git a/vllm/transformers_utils/runai_utils.py b/vllm/transformers_utils/runai_utils.py
index 248ede6a6f1d..d02992b10b1c 100644
--- a/vllm/transformers_utils/runai_utils.py
+++ b/vllm/transformers_utils/runai_utils.py
@@ -25,14 +25,14 @@
 
 
 def list_safetensors(path: str = "") -> list[str]:
-    """
-    List full file names from object path and filter by allow pattern.
+    """List full file names from object path and filter by allow pattern.
 
     Args:
         path: The object storage path to list from.
 
     Returns:
         list[str]: List of full object storage paths allowed by the pattern
+
     """
     return runai_list_safetensors(path)
 
@@ -42,8 +42,7 @@ def is_runai_obj_uri(model_or_path: str) -> bool:
 
 
 class ObjectStorageModel:
-    """
-    A class representing an ObjectStorage model mirrored into a
+    """A class representing an ObjectStorage model mirrored into a
     temporary directory.
 
     Attributes:
@@ -51,6 +50,7 @@ class ObjectStorageModel:
 
     Methods:
         pull_files(): Pull model from object storage to the temporary directory.
+
     """
 
     def __init__(self, url: str) -> None:
@@ -86,8 +86,7 @@ def pull_files(
         allow_pattern: list[str] | None = None,
         ignore_pattern: list[str] | None = None,
     ) -> None:
-        """
-        Pull files from object storage into the temporary directory.
+        """Pull files from object storage into the temporary directory.
 
         Args:
             model_path: The object storage path of the model.
diff --git a/vllm/transformers_utils/s3_utils.py b/vllm/transformers_utils/s3_utils.py
index 7dfee9fabe36..d0f403c6908d 100644
--- a/vllm/transformers_utils/s3_utils.py
+++ b/vllm/transformers_utils/s3_utils.py
@@ -36,8 +36,7 @@ def glob(
     path: str = "",
     allow_pattern: list[str] | None = None,
 ) -> list[str]:
-    """
-    List full file names from S3 path and filter by allow pattern.
+    """List full file names from S3 path and filter by allow pattern.
 
     Args:
         s3: S3 client to use.
@@ -46,6 +45,7 @@ def glob(
 
     Returns:
         list[str]: List of full S3 paths allowed by the pattern
+
     """
     if s3 is None:
         s3 = boto3.client("s3")
@@ -61,8 +61,7 @@ def list_files(
     allow_pattern: list[str] | None = None,
     ignore_pattern: list[str] | None = None,
 ) -> tuple[str, str, list[str]]:
-    """
-    List files from S3 path and filter by pattern.
+    """List files from S3 path and filter by pattern.
 
     Args:
         s3: S3 client to use.
@@ -77,6 +76,7 @@ def list_files(
               and the prefix as a dir like string
             - The third element is a list of files allowed or
               disallowed by pattern
+
     """
     parts = path.removeprefix("s3://").split("/")
     prefix = "/".join(parts[1:])
diff --git a/vllm/transformers_utils/utils.py b/vllm/transformers_utils/utils.py
index 674f6a05bf98..ad07d901b4fe 100644
--- a/vllm/transformers_utils/utils.py
+++ b/vllm/transformers_utils/utils.py
@@ -80,16 +80,15 @@ def _maybe_space_split_dict(path: str | PathLike) -> dict[str, str]:
 
 @cache
 def maybe_model_redirect(model: str) -> str:
-    """
-    Use model_redirect to redirect the model name to a local folder.
+    """Use model_redirect to redirect the model name to a local folder.
 
     Args:
         model: hf model name
 
     Returns:
         maybe redirect to a local folder
-    """
 
+    """
     model_redirect_path = envs.VLLM_MODEL_REDIRECT_PATH
 
     if not model_redirect_path:
diff --git a/vllm/triton_utils/importing.py b/vllm/triton_utils/importing.py
index 966ab55f326b..e312564c0bd4 100644
--- a/vllm/triton_utils/importing.py
+++ b/vllm/triton_utils/importing.py
@@ -124,7 +124,7 @@ def decorator(*args, **kwargs):
 class TritonLanguagePlaceholder(types.ModuleType):
     def __init__(self):
         super().__init__("triton.language")
-        self.constexpr = None
+        self.constexpr = lambda value: value
         self.dtype = None
         self.int64 = None
         self.int32 = None
diff --git a/vllm/utils/argparse_utils.py b/vllm/utils/argparse_utils.py
index fc906cf22727..b91b1f26fd13 100644
--- a/vllm/utils/argparse_utils.py
+++ b/vllm/utils/argparse_utils.py
@@ -34,6 +34,7 @@ def human_readable_int(value: str) -> int:
     - '1k' -> 1,000
     - '1K' -> 1,024
     - '25.6k' -> 25,600
+
     """
     value = value.strip()
 
@@ -82,6 +83,7 @@ def human_readable_int_or_auto(value: str) -> int:
     - '1K' -> 1,024
     - '25.6k' -> 25,600
     - '-1' or 'auto' -> -1 (special value for auto-detection)
+
     """
     value = value.strip()
 
@@ -91,12 +93,16 @@ def human_readable_int_or_auto(value: str) -> int:
     return human_readable_int(value)
 
 
+def summary_line(description: str | None) -> str:
+    """Return the first paragraph of a config docstring as a single line."""
+    return " ".join((description or "").split("\n\n", 1)[0].split())
+
+
 class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter):
     """SortedHelpFormatter that sorts arguments by their option strings."""
 
     def _split_lines(self, text, width):
-        """
-        1. Sentences split across lines have their single newlines removed.
+        """1. Sentences split across lines have their single newlines removed.
         2. Paragraphs and lists are split into separate lines.
         3. Each line is wrapped to the specified width (width of terminal).
         """
@@ -202,7 +208,7 @@ def format_help(self):
             for group in self._action_groups:
                 if group.title and group.title.lower() == search_keyword:
                     formatter.start_section(group.title)
-                    formatter.add_text(group.description)
+                    formatter.add_text(summary_line(group.description))
                     formatter.add_arguments(group._group_actions)
                     formatter.end_section()
                     formatter.add_text(self._json_tip)
@@ -241,12 +247,18 @@ def format_help(self):
         # positionals, optionals and user-defined groups
         formatter.start_section("Config Groups")
         config_groups = ""
-        for group in self._action_groups:
-            if not group._group_actions:
-                continue
-            title = group.title
-            description = group.description or ""
-            config_groups += f"{title: <24}{description}\n"
+        groups = [g for g in self._action_groups if g._group_actions]
+        titles = [g.title or "" for g in groups]
+        title_width = max((len(title) for title in titles), default=0) + 1
+        width = max(formatter._width - formatter._current_indent, title_width + 20)
+        for title, group in zip(titles, groups):
+            lines = textwrap.wrap(
+                summary_line(group.description),
+                width,
+                initial_indent=f"{title: <{title_width}}",
+                subsequent_indent=" " * title_width,
+            )
+            config_groups += ("\n".join(lines) or title) + "\n"
         formatter.add_text(config_groups)
         formatter.end_section()
 
@@ -466,7 +478,7 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]:
         The arguments in config file will be inserted between
         the argument list.
 
-        example:
+        Example:
         ```yaml
             port: 12323
             tensor-parallel-size: 4
@@ -492,6 +504,7 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]:
         Please note how the config args are inserted after the sub command.
         this way the order of priorities is maintained when these are args
         parsed by super().
+
         """
         assert args.count("--config") <= 1, "More than one config file specified!"
 
@@ -506,7 +519,7 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]:
         config_args = self.load_config_file(file_path)
 
         # 0th index might be the sub command {serve,chat,complete,...}
-        # optionally followed by model_tag (only for serve)
+        # optionally followed by model_tag (serve or snapshot create)
         # followed by config args
         # followed by rest of cli args.
         # maintaining this order will enforce the precedence
@@ -514,8 +527,11 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]:
         if args[0].startswith("-"):
             # No sub command (e.g., api_server entry point)
             args = config_args + args[0:index] + args[index + 2 :]
-        elif args[0] == "serve":
-            model_in_cli = len(args) > 1 and not args[1].startswith("-")
+        elif args[0] == "serve" or args[:2] == ["snapshot", "create"]:
+            model_index = 1 if args[0] == "serve" else 2
+            model_in_cli = len(args) > model_index and not args[model_index].startswith(
+                "-"
+            )
             model_in_config = any(arg == "--model" for arg in config_args)
 
             if not model_in_cli and not model_in_config:
@@ -527,15 +543,19 @@ def _pull_args_from_config(self, args: list[str]) -> list[str]:
             if model_in_cli:
                 # Model specified as positional arg, keep CLI version
                 args = (
-                    [args[0]]
-                    + [args[1]]
+                    args[: model_index + 1]
                     + config_args
-                    + args[2:index]
+                    + args[model_index + 1 : index]
                     + args[index + 2 :]
                 )
             else:
                 # No model in CLI, use config if available
-                args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
+                args = (
+                    args[:model_index]
+                    + config_args
+                    + args[model_index:index]
+                    + args[index + 2 :]
+                )
         else:
             args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
 
diff --git a/vllm/utils/async_utils.py b/vllm/utils/async_utils.py
index 3faf728384e6..b2c0d57dc921 100644
--- a/vllm/utils/async_utils.py
+++ b/vllm/utils/async_utils.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Contains helpers related to asynchronous code.
+"""Contains helpers related to asynchronous code.
 
 This is similar in concept to the `asyncio` module.
 """
@@ -29,8 +28,7 @@ def make_async(
     func: Callable[P, T],
     executor: Executor | None = None,
 ) -> Callable[P, Awaitable[T]]:
-    """
-    Take a blocking function, and run it on in an executor thread.
+    """Take a blocking function, and run it on in an executor thread.
 
     This function prevents the blocking function from blocking the
     asyncio event loop.
@@ -49,8 +47,7 @@ def make_async_with_semaphore(
     func: Callable[P, T],
     executor: ThreadPoolExecutor,
 ) -> Callable[P, Awaitable[T]]:
-    """
-    Take a blocking function, and run it on in an executor thread.
+    """Take a blocking function, and run it on in an executor thread.
 
     This function prevents the blocking function from blocking the
     asyncio event loop.
@@ -59,7 +56,6 @@ def make_async_with_semaphore(
     The function is wrapped in a semaphore to limit the number of
     concurrent executions making it easier to cancel tasks before they start.
     """
-
     semaphore = asyncio.Semaphore(executor._max_workers)
 
     async def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
diff --git a/vllm/utils/cache.py b/vllm/utils/cache.py
index 791b8d848d4f..9fca5ec46541 100644
--- a/vllm/utils/cache.py
+++ b/vllm/utils/cache.py
@@ -102,8 +102,7 @@ def usage(self) -> float:
         return self.currsize / self.maxsize
 
     def stat(self, *, delta: bool = False) -> CacheInfo:
-        """
-        Gets the cumulative number of hits and queries against this cache.
+        """Gets the cumulative number of hits and queries against this cache.
 
         If `delta=True`, instead gets these statistics
         since the last call that also passed `delta=True`.
@@ -165,6 +164,7 @@ def put_if_fits(self, key: _K, value: _V) -> bool:
 
         Returns:
             `True` if the item was cached, otherwise `False`.
+
         """
         try:
             self[key] = value
@@ -175,8 +175,7 @@ def put_if_fits(self, key: _K, value: _V) -> bool:
         return True
 
     def pin(self, key: _K) -> None:
-        """
-        Pins a key in the cache preventing it from being
+        """Pins a key in the cache preventing it from being
         evicted in the LRU order.
         """
         if key not in self:
@@ -184,8 +183,7 @@ def pin(self, key: _K) -> None:
         self.pinned_items.add(key)
 
     def _unpin(self, key: _K) -> None:
-        """
-        Unpins a key in the cache allowing it to be
+        """Unpins a key in the cache allowing it to be
         evicted in the LRU order.
         """
         self.pinned_items.remove(key)
diff --git a/vllm/utils/collection_utils.py b/vllm/utils/collection_utils.py
index 5852aff549c4..de32ed0c15c2 100644
--- a/vllm/utils/collection_utils.py
+++ b/vllm/utils/collection_utils.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Contains helpers that are applied to collections.
+"""Contains helpers that are applied to collections.
 
 This is similar in concept to the `collections` module.
 """
@@ -20,8 +19,7 @@
 
 
 class LazyDict(Mapping[str, _V], Generic[_V]):
-    """
-    Evaluates dictionary items only when they are accessed.
+    """Evaluates dictionary items only when they are accessed.
 
     Adapted from: https://stackoverflow.com/a/47212782/5082708
     """
@@ -116,8 +114,7 @@ def flatten_2d_lists(lists: Iterable[Iterable[T]]) -> list[T]:
 
 
 def full_groupby(values: Iterable[_V], *, key: Callable[[_V], _K]):
-    """
-    Unlike [`itertools.groupby`][], groups are not broken by
+    """Unlike [`itertools.groupby`][], groups are not broken by
     non-contiguous data.
     """
     groups = defaultdict[_K, list[_V]](list)
diff --git a/vllm/utils/counter.py b/vllm/utils/counter.py
index c2dce32e97e1..6251cdcfb1b3 100644
--- a/vllm/utils/counter.py
+++ b/vllm/utils/counter.py
@@ -19,10 +19,10 @@ def reset(self) -> None:
 
 
 class AtomicCounter:
-    """An atomic, thread-safe counter"""
+    """An atomic, thread-safe counter."""
 
     def __init__(self, initial: int = 0) -> None:
-        """Initialize a new atomic counter to given initial value"""
+        """Initialize a new atomic counter to given initial value."""
         super().__init__()
 
         self._value = initial
@@ -33,13 +33,13 @@ def value(self) -> int:
         return self._value
 
     def inc(self, num: int = 1) -> int:
-        """Atomically increment the counter by num and return the new value"""
+        """Atomically increment the counter by num and return the new value."""
         with self._lock:
             self._value += num
             return self._value
 
     def dec(self, num: int = 1) -> int:
-        """Atomically decrement the counter by num and return the new value"""
+        """Atomically decrement the counter by num and return the new value."""
         with self._lock:
             self._value -= num
             return self._value
diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py
index 53e74610659b..296fe466460f 100644
--- a/vllm/utils/cpu_resource_utils.py
+++ b/vllm/utils/cpu_resource_utils.py
@@ -12,8 +12,12 @@
 import psutil
 import regex as re
 
+from vllm.logger import init_logger
+
 DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES"
 
+logger = init_logger(__name__)
+
 
 @dataclass
 class LogicalCPUInfo:
@@ -63,21 +67,20 @@ def _read_int_file(path: str) -> int | None:
 
 
 @cache
-def get_cgroup_memory_limit() -> tuple[int | None, int | None]:
-    """Return (limit, usage) in bytes from cgroup, or (None, None).
+def get_cgroup_memory_limit() -> int | None:
+    """Return the cgroup memory limit in bytes, or None.
 
-    Supports both cgroup v2 (unified) and v1. Returns (None, None) when
+    Supports both cgroup v2 (unified) and v1. Returns None when
     not running under a constrained cgroup (e.g. bare metal, or limit
     reported as `max`/an unrealistically large value).
     """
     if sys.platform != "linux":
-        return None, None
+        return None
 
     # cgroup v2 unified hierarchy
     v2_limit = _read_int_file("/sys/fs/cgroup/memory.max")
     if v2_limit is not None:
-        v2_usage = _read_int_file("/sys/fs/cgroup/memory.current")
-        return v2_limit, v2_usage
+        return v2_limit
 
     # cgroup v1
     v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes")
@@ -85,11 +88,75 @@ def get_cgroup_memory_limit() -> tuple[int | None, int | None]:
         # cgroup v1 reports a huge sentinel (close to PAGE_COUNTER_MAX)
         # when unlimited. Treat absurdly large values as "no limit".
         if v1_limit >= (1 << 62):
-            return None, None
-        v1_usage = _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes")
-        return v1_limit, v1_usage
+            return None
+        return v1_limit
+
+    return None
+
+
+def get_cgroup_memory_usage() -> int | None:
+    """Return the current cgroup memory usage in bytes, or None.
+
+    The usage value is intentionally read on every call because cgroup
+    memory usage changes while the process is running.
+    """
+    if sys.platform != "linux":
+        return None
+
+    # cgroup v2 unified hierarchy
+    if _read_int_file("/sys/fs/cgroup/memory.max") is not None:
+        return _read_int_file("/sys/fs/cgroup/memory.current")
+
+    # cgroup v1
+    v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes")
+    if v1_limit is not None and v1_limit < (1 << 62):
+        return _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes")
+
+    return None
+
 
-    return None, None
+def check_cgroup_memory_available(
+    required_bytes: int,
+    allocation_name: str,
+) -> None:
+    """Log cgroup memory headroom for an upcoming allocation.
+
+    Args:
+        required_bytes: Bytes required by the allocation.
+        allocation_name: Human-readable name used in log messages.
+
+    Low headroom logs a warning, but does not reject the allocation because
+    cgroup usage can include reclaimable memory. If the cgroup limit or usage
+    cannot be read, the check is skipped.
+
+    """
+    cgroup_limit = get_cgroup_memory_limit()
+    cgroup_usage = get_cgroup_memory_usage()
+    if cgroup_limit is None or cgroup_usage is None:
+        return
+
+    cgroup_available = max(0, cgroup_limit - cgroup_usage)
+    mib = 1 << 20
+    remaining_bytes = cgroup_available - required_bytes
+    log_fn = logger.debug if remaining_bytes >= 0 else logger.warning
+    status = (
+        "current headroom meets the requested allocation"
+        if remaining_bytes >= 0
+        else "current headroom is below the requested allocation; allocation "
+        "will still be attempted because cgroup usage may be reclaimable"
+    )
+    log_fn(
+        "Cgroup memory preflight for %s: %.0f MiB required, %.0f MiB current "
+        "usage, %.0f MiB available under %.0f MiB limit, %.0f MiB remaining "
+        "after allocation based on current usage; %s.",
+        allocation_name,
+        required_bytes / mib,
+        cgroup_usage / mib,
+        cgroup_available / mib,
+        cgroup_limit / mib,
+        remaining_bytes / mib,
+        status,
+    )
 
 
 def get_memory_affinity(pid: int = 0) -> list[int]:
@@ -105,7 +172,7 @@ def get_memory_affinity(pid: int = 0) -> list[int]:
 
 
 def parse_id_list(raw_str: str) -> list[int]:
-    """Parses strings like '0-2,4,7-8' into [0, 1, 2, 4, 7, 8]"""
+    """Parses strings like '0-2,4,7-8' into [0, 1, 2, 4, 7, 8]."""
     result: list[int] = []
     if not raw_str:
         return result
@@ -161,7 +228,8 @@ def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:
     # would be applied to host RAM instead of the pod's limit. cgroup
     # does not expose per-NUMA-node limits, so we just clamp the totals
     # against the pod-wide limit here.
-    cgroup_limit, cgroup_usage = get_cgroup_memory_limit()
+    cgroup_limit = get_cgroup_memory_limit()
+    cgroup_usage = get_cgroup_memory_usage()
     if cgroup_limit is not None and cgroup_limit < total_memory:
         total_memory = cgroup_limit
         cgroup_available = cgroup_limit - (cgroup_usage or 0)
diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py
index 9ce33b835707..d1696317aab8 100644
--- a/vllm/utils/cpu_triton_utils.py
+++ b/vllm/utils/cpu_triton_utils.py
@@ -1,9 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Contains replacement functions to fallback Triton usages in CPU backend
-"""
+"""Contains replacement functions to fallback Triton usages in CPU backend."""
 
 import ctypes
 from collections.abc import Callable
diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py
index 3fc70c0f5da3..0eb979ce05a5 100644
--- a/vllm/utils/deep_gemm.py
+++ b/vllm/utils/deep_gemm.py
@@ -64,7 +64,7 @@ class DeepGemmQuantScaleFMT(Enum):
 
     @classmethod
     def init_oracle_cache(cls) -> None:
-        """Initialize the oracle decision and store it in the class cache"""
+        """Initialize the oracle decision and store it in the class cache."""
         cached = getattr(cls, "_oracle_cache", None)
         if cached is not None:
             return
@@ -112,7 +112,7 @@ def is_deep_gemm_supported() -> bool:
     Currently, only Hopper and Blackwell GPUs are supported.
     """
     is_supported_arch = current_platform.support_deep_gemm()
-    return envs.VLLM_USE_DEEP_GEMM and has_deep_gemm() and is_supported_arch
+    return is_supported_arch and envs.VLLM_USE_DEEP_GEMM and has_deep_gemm()
 
 
 @functools.cache
@@ -158,7 +158,13 @@ def _missing(*_: Any, **__: Any) -> NoReturn:
 _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None
 _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None
 _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
+_fp8_fp4_sparse_mqa_logits_impl: Callable[..., Any] | None = None
+_fp8_fp4_paged_sparse_mqa_logits_impl: Callable[..., Any] | None = None
+_get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None
+_get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None
 _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
+_mega_mhc_impl: Callable[..., Any] | None = None
+_bf16_mega_gate_impl: Callable[..., Any] | None = None
 _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
 _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
 _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = (
@@ -231,7 +237,11 @@ def _lazy_init() -> None:
     global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl
     global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl
     global _get_paged_mqa_logits_metadata_impl
-    global _tf32_hc_prenorm_gemm_impl
+    global _fp8_fp4_sparse_mqa_logits_impl, _fp8_fp4_paged_sparse_mqa_logits_impl
+    global _get_sparse_mqa_logits_metadata_impl
+    global _get_paged_sparse_mqa_logits_metadata_impl
+    global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl
+    global _bf16_mega_gate_impl
     global _get_mn_major_tma_aligned_tensor_impl
     global _get_mk_alignment_for_contiguous_layout_impl
     global _get_theoretical_mk_alignment_for_contiguous_layout_impl
@@ -251,6 +261,8 @@ def _lazy_init() -> None:
         or _fp8_fp4_paged_mqa_logits_impl is not None
         or _get_paged_mqa_logits_metadata_impl is not None
         or _tf32_hc_prenorm_gemm_impl is not None
+        or _mega_mhc_impl is not None
+        or _bf16_mega_gate_impl is not None
         or _get_mk_alignment_for_contiguous_layout_impl is not None
         or _transform_sf_into_required_layout_impl is not None
         or _pack_ue8m0_to_int_impl is not None
@@ -280,7 +292,11 @@ def _lazy_init() -> None:
     _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
     _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None)
     _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None)
-    _grouped_masked_impl = getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None)
+    # DeepGEMM 2.8.0 dropped the legacy `fp8_*` aliases; keep both spellings so
+    # an externally pinned older DeepGEMM still resolves.
+    _grouped_masked_impl = getattr(
+        _dg, "m_grouped_fp8_gemm_nt_masked", None
+    ) or getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None)
     _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None)
     # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that
     # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`.
@@ -289,7 +305,20 @@ def _lazy_init() -> None:
     _get_paged_mqa_logits_metadata_impl = getattr(
         _dg, "get_paged_mqa_logits_metadata", None
     )
+    # Sparse-indexer kernels (DeepGEMM >= 2.8, SM100 only).
+    _fp8_fp4_sparse_mqa_logits_impl = getattr(_dg, "fp8_fp4_sparse_mqa_logits", None)
+    _fp8_fp4_paged_sparse_mqa_logits_impl = getattr(
+        _dg, "fp8_fp4_paged_sparse_mqa_logits", None
+    )
+    _get_sparse_mqa_logits_metadata_impl = getattr(
+        _dg, "get_sparse_mqa_logits_metadata", None
+    )
+    _get_paged_sparse_mqa_logits_metadata_impl = getattr(
+        _dg, "get_paged_sparse_mqa_logits_metadata", None
+    )
     _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None)
+    _mega_mhc_impl = getattr(_dg, "mega_mhc", None)
+    _bf16_mega_gate_impl = getattr(_dg, "bf16_mega_gate", None)
     _get_mn_major_tma_aligned_tensor_impl = getattr(
         _dg, "get_mn_major_tma_aligned_tensor", None
     )
@@ -401,7 +430,7 @@ def mk_alignment_scope(value: int):
 
 
 def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
-    """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor"""
+    """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor."""
     _lazy_init()
     if _get_mn_major_tma_aligned_tensor_impl is None:
         return _missing()
@@ -457,6 +486,55 @@ def cublaslt_gemm_nt(*args, **kwargs):
     return _cublaslt_gemm_nt_impl(*args, **kwargs)
 
 
+def bf16_mega_gate(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    num_topk: int,
+    *,
+    scoring_func: str,
+    routed_scaling_factor: float,
+    ep_rank: int,
+    bias: torch.Tensor | None = None,
+    image_bias: torch.Tensor | None = None,
+    image_token_mask: torch.Tensor | None = None,
+    fix_routing_mask: torch.Tensor | None = None,
+    unmapped_topk_idx: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run DeepGEMM Mega Gate.
+
+    The output order follows vLLM's router contract: weights, then indices.
+    Passing PyTorch-owned outputs to DeepGEMM makes their storage part of CUDA
+    graph capture.
+    """
+    _lazy_init()
+    if _bf16_mega_gate_impl is None:
+        raise RuntimeError(
+            "DeepGEMM MegaMoE requires a DeepGEMM build with bf16_mega_gate."
+        )
+
+    topk_idx = torch.empty((x.shape[0], num_topk), dtype=torch.int64, device=x.device)
+    topk_weights = torch.empty(
+        (x.shape[0], num_topk), dtype=torch.float32, device=x.device
+    )
+    _bf16_mega_gate_impl(
+        x,
+        weight,
+        num_topk,
+        use_shared_as_routed=False,
+        num_shared_experts=0,
+        routed_scaling_factor=routed_scaling_factor,
+        ep_rank=ep_rank,
+        scoring_func=scoring_func,
+        bias=bias,
+        image_bias=image_bias,
+        image_token_mask=image_token_mask,
+        fix_routing_mask=fix_routing_mask,
+        unmapped_topk_idx=unmapped_topk_idx,
+        out=(topk_idx, topk_weights),
+    )
+    return topk_weights, topk_idx
+
+
 def fp8_gemm_nt(*args, **kwargs):
     _lazy_init()
     if _fp8_gemm_nt_impl is None:
@@ -543,6 +621,7 @@ def fp8_fp4_mqa_logits(
 
     Returns:
         Logits tensor of shape [M, N], dtype `torch.float32`.
+
     """
     _lazy_init()
     if _fp8_fp4_mqa_logits_impl is None:
@@ -598,6 +677,7 @@ def get_paged_mqa_logits_metadata(
     Returns:
         Tensor of shape [slots + 1, 2] consumed by `fp8_fp4_paged_mqa_logits`
         to schedule work across SMs.
+
     """
     _lazy_init()
     if _get_paged_mqa_logits_metadata_impl is None:
@@ -649,6 +729,7 @@ def fp8_fp4_paged_mqa_logits(
     Returns:
         Logits tensor of shape [B * next_n, max_model_len], dtype
         `torch.float32`.
+
     """
     _lazy_init()
     if _fp8_fp4_paged_mqa_logits_impl is None:
@@ -674,6 +755,179 @@ def fp8_fp4_paged_mqa_logits(
     )
 
 
+def has_deep_gemm_sparse_mqa() -> bool:
+    """Whether the installed DeepGEMM provides the sparse-indexer kernels
+    (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only)."""
+    _lazy_init()
+    return (
+        _fp8_fp4_sparse_mqa_logits_impl is not None
+        and _fp8_fp4_paged_sparse_mqa_logits_impl is not None
+        and _get_sparse_mqa_logits_metadata_impl is not None
+        and _get_paged_sparse_mqa_logits_metadata_impl is not None
+    )
+
+
+def get_sparse_mqa_logits_metadata(
+    cu_seqlen_ks: torch.Tensor,
+    cu_seqlen_ke: torch.Tensor,
+    num_kv_tokens: int,
+    sparse_kv_block_indices: torch.Tensor,
+    qk_dtype: torch.dtype,
+    sparse_block_kv: int,
+) -> torch.Tensor:
+    """Build scheduling metadata for `fp8_fp4_sparse_mqa_logits`.
+
+    Always uses the unaligned-ks variant, which degenerates to the aligned
+    variant when ``cu_seqlen_ks % sparse_block_kv == 0``, so callers never
+    need a host-side sync to pick one.
+
+    Args:
+        cu_seqlen_ks: Per-row K range start bounds in the packed KV
+            workspace, shape [num_q_tokens], dtype int32.
+        cu_seqlen_ke: Per-row K range end bounds in the packed KV workspace,
+            shape [num_q_tokens], dtype int32.
+        num_kv_tokens: Total KV tokens in the packed workspace.
+        sparse_kv_block_indices: Per-row candidate block ids, shape
+            [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's
+            valid prefix must be sorted ascending (repeats are tolerated but
+            wasteful); pad by repeating the last valid block. Block ``i``
+            covers tokens
+            ``[i * sparse_block_kv + ks % sparse_block_kv, ...)``.
+        qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4,
+            ``torch.float8_e4m3fn`` for FP8).
+        sparse_block_kv: Tokens per sparse block, 8 or 16.
+
+    """
+    _lazy_init()
+    if _get_sparse_mqa_logits_metadata_impl is None:
+        return _missing()
+    return _get_sparse_mqa_logits_metadata_impl(
+        cu_seqlen_ks,
+        cu_seqlen_ke,
+        num_kv_tokens,
+        sparse_kv_block_indices,
+        qk_dtype,
+        sparse_block_kv,
+        use_unaligned_ks=True,
+    )
+
+
+def get_paged_sparse_mqa_logits_metadata(
+    context_lens: torch.Tensor,
+    block_table: torch.Tensor,
+    indices: torch.Tensor,
+    page_kv: int,
+    sparse_kv_block_indices: torch.Tensor,
+    qk_dtype: torch.dtype,
+    sparse_block_kv: int,
+) -> torch.Tensor:
+    """Build scheduling metadata for `fp8_fp4_paged_sparse_mqa_logits`.
+
+    Args:
+        context_lens: Per-row context length, shape [num_q_tokens], int32.
+        block_table: Per-row page table, shape [num_q_tokens, max_pages],
+            int32; rows of queries paired with the same request must be
+            identical (expand it per query row before calling).
+        indices: Request index for each query row, shape [num_q_tokens].
+        page_kv: Tokens per KV page; must be a multiple of sparse_block_kv.
+        sparse_kv_block_indices: Per-row logical (context-relative) candidate
+            block ids, [num_q_tokens, num_max_sparse_blocks], int32.
+        qk_dtype: dtype of the packed Q values.
+        sparse_block_kv: Tokens per sparse block, 8 or 16.
+
+    """
+    _lazy_init()
+    if _get_paged_sparse_mqa_logits_metadata_impl is None:
+        return _missing()
+    return _get_paged_sparse_mqa_logits_metadata_impl(
+        context_lens,
+        block_table,
+        indices,
+        page_kv,
+        sparse_kv_block_indices,
+        qk_dtype,
+        sparse_block_kv,
+    )
+
+
+def fp8_fp4_sparse_mqa_logits(
+    q: tuple[torch.Tensor, torch.Tensor],
+    kv: tuple[torch.Tensor, torch.Tensor],
+    weights: torch.Tensor,
+    metadata: torch.Tensor,
+    num_max_sparse_blocks: int,
+    sparse_block_kv: int,
+) -> torch.Tensor:
+    """Compute MQA logits only at the candidate blocks (prefill path).
+
+    Args:
+        q: ``(q_values, q_scale)``. q_values is [M, H, D] (packed FP4 viewed
+            as int8, or FP8); q_scale is the packed UE8M0 scale tensor [M, H]
+            int32 (mandatory, unlike the dense kernels).
+        kv: ``(kv_values, kv_scale)`` — the packed KV workspace with UE8M0
+            int32 scales.
+        weights: [M, H] ``torch.bfloat16`` (the Q scale is NOT folded in).
+        metadata: From `get_sparse_mqa_logits_metadata`.
+        num_max_sparse_blocks: Width of each sparse-index row.
+        sparse_block_kv: Tokens per sparse block, 8 or 16.
+
+    Returns:
+        bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv];
+        column ``j * sparse_block_kv + o`` scores the token at
+        ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``.
+
+    """
+    _lazy_init()
+    if _fp8_fp4_sparse_mqa_logits_impl is None:
+        return _missing()
+    return _fp8_fp4_sparse_mqa_logits_impl(
+        q,
+        kv,
+        weights,
+        metadata,
+        num_max_sparse_blocks,
+        sparse_block_kv,
+        use_unaligned_ks=True,
+    )
+
+
+def fp8_fp4_paged_sparse_mqa_logits(
+    q: tuple[torch.Tensor, torch.Tensor],
+    kv_cache: torch.Tensor,
+    weights: torch.Tensor,
+    metadata: torch.Tensor,
+    num_max_sparse_blocks: int,
+    sparse_block_kv: int,
+) -> torch.Tensor:
+    """Compute MQA logits only at the candidate blocks (paged decode path).
+
+    Args:
+        q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D].
+        kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8,
+            page stride 512B-aligned.
+        weights: [num_q_tokens, H] ``torch.bfloat16``.
+        metadata: From `get_paged_sparse_mqa_logits_metadata`.
+        num_max_sparse_blocks: Candidate blocks per row.
+        sparse_block_kv: Tokens per sparse block, 8 or 16.
+
+    Returns:
+        bf16 logits of shape
+        [num_q_tokens, num_max_sparse_blocks * sparse_block_kv].
+
+    """
+    _lazy_init()
+    if _fp8_fp4_paged_sparse_mqa_logits_impl is None:
+        return _missing()
+    return _fp8_fp4_paged_sparse_mqa_logits_impl(
+        q,
+        kv_cache,
+        weights,
+        metadata,
+        num_max_sparse_blocks,
+        sparse_block_kv,
+    )
+
+
 def tf32_hc_prenorm_gemm(
     x: torch.Tensor,
     fn: torch.Tensor,
@@ -681,8 +935,7 @@ def tf32_hc_prenorm_gemm(
     sqrsum: torch.Tensor,
     num_split: int,
 ) -> torch.Tensor:
-    """
-    Perform the following computation:
+    """Perform the following computation:
         out = x.float() @ fn.T
         sqrsum = x.float().square().sum(-1)
 
@@ -700,6 +953,15 @@ def tf32_hc_prenorm_gemm(
     )
 
 
+def mega_mhc(*args: Any, **kwargs: Any) -> None:
+    """Run DeepGEMM Mega mHC with caller-owned output tensors."""
+    _lazy_init()
+    if _mega_mhc_impl is None:
+        _missing()
+        return
+    _mega_mhc_impl(*args, **kwargs)
+
+
 def _ceil_to_ue8m0(x: torch.Tensor):
     return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
 
@@ -749,7 +1011,6 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor):
     and report `1 - sim`.  Once kernel accuracy improves this helper can be
     removed.
     """
-
     x, y = x.double(), y.double()
     denominator = (x * x + y * y).sum()
     sim = 2 * (x * y).sum() / denominator
@@ -780,6 +1041,7 @@ def should_use_deepgemm_for_fp8_linear(
 
 __all__ = [
     "calc_diff",
+    "bf16_mega_gate",
     "DeepGemmQuantScaleFMT",
     "fp8_gemm_nt",
     "fp8_einsum",
@@ -793,6 +1055,7 @@ def should_use_deepgemm_for_fp8_linear(
     "per_block_cast_to_fp8",
     "is_deep_gemm_e8m0_used",
     "is_deep_gemm_supported",
+    "mega_mhc",
     "get_num_sms",
     "set_num_sms",
     "should_use_deepgemm_for_fp8_linear",
diff --git a/vllm/utils/diffusion.py b/vllm/utils/diffusion.py
new file mode 100644
index 000000000000..278c17d83823
--- /dev/null
+++ b/vllm/utils/diffusion.py
@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Request validation for discrete diffusion models."""
+
+from typing import TYPE_CHECKING
+
+from vllm.exceptions import VLLMValidationError
+
+if TYPE_CHECKING:
+    from vllm.sampling_params import SamplingParams
+
+
+def validate_diffusion_sampling_params(
+    params: "SamplingParams",
+    *,
+    canvas_length: int | None,
+    vocab_size: int,
+    async_scheduling: bool | None,
+) -> None:
+    """Validate request options and normalize one-canvas reads at admission."""
+    extra = params.extra_args
+    if not extra:
+        return
+
+    def check(valid: bool, message: str) -> None:
+        if not valid:
+            raise VLLMValidationError(message, parameter="extra_args")
+
+    width = extra.get("diffusion_canvas_length")
+    if width is not None:
+        check(
+            isinstance(width, int)
+            and not isinstance(width, bool)
+            and width > 0
+            and (canvas_length is None or width <= canvas_length),
+            "diffusion_canvas_length must be a positive integer no larger "
+            "than the served canvas.",
+        )
+    # The request's canvas: its own width, else the served one.
+    expected_len = width or canvas_length
+
+    seed = extra.get("diffusion_seed_canvas")
+    if seed is not None:
+        check(
+            isinstance(seed, (list, tuple))
+            and all(isinstance(t, int) and not isinstance(t, bool) for t in seed),
+            "diffusion_seed_canvas must be a list of token ids.",
+        )
+        check(
+            all(0 <= t < vocab_size for t in seed),
+            f"diffusion_seed_canvas ids must be in [0, {vocab_size}).",
+        )
+        check(
+            expected_len is None or len(seed) == expected_len,
+            f"diffusion_seed_canvas must hold exactly {expected_len} ids, "
+            f"got {len(seed)}.",
+        )
+
+    pins = extra.get("diffusion_pinned")
+    if pins is not None:
+        check(
+            isinstance(pins, (list, tuple))
+            and all(isinstance(p, int) and not isinstance(p, bool) for p in pins),
+            "diffusion_pinned must be a list of canvas positions.",
+        )
+        check(
+            seed is not None,
+            "diffusion_pinned needs a diffusion_seed_canvas to hold.",
+        )
+        check(
+            all(p >= 0 and (expected_len is None or p < expected_len) for p in pins),
+            "diffusion_pinned positions must lie inside the canvas.",
+        )
+
+    cap = extra.get("diffusion_max_steps")
+    if cap is not None:
+        check(
+            isinstance(cap, int) and not isinstance(cap, bool) and cap > 0,
+            "diffusion_max_steps must be a positive integer.",
+        )
+
+    # The OpenAI server's vllm_xargs narrows JSON booleans to 0/1.
+    read_only = extra.get("diffusion_read_only")
+    if read_only is not None:
+        check(
+            isinstance(read_only, (bool, int)) and read_only in (0, 1),
+            "diffusion_read_only must be a boolean (or 0/1).",
+        )
+    if read_only:
+        # One canvas is the whole output. Cap max_tokens so the scheduler
+        # ends the request there, and ignore EOS so an end-of-turn token
+        # drawn into a noise slot does not end it early.
+        if expected_len is not None:
+            params.max_tokens = min(params.max_tokens or expected_len, expected_len)
+        params.ignore_eos = True
+
+    check(
+        bool(async_scheduling)
+        or width is None
+        or canvas_length is None
+        or width == canvas_length,
+        "A diffusion_canvas_length smaller than the served canvas "
+        "requires --async-scheduling.",
+    )
diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py
index 7c76dd7ba2ff..abc4d75500c1 100644
--- a/vllm/utils/flashinfer.py
+++ b/vllm/utils/flashinfer.py
@@ -11,7 +11,8 @@
 import importlib.util
 import os
 import shutil
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
+from contextvars import ContextVar
 from typing import Any, NoReturn
 
 import requests
@@ -26,6 +27,24 @@
 logger = init_logger(__name__)
 
 
+_bf16_autotune_buckets: ContextVar[tuple[int, ...] | None] = ContextVar(
+    "flashinfer_bf16_autotune_buckets", default=None
+)
+
+
+@contextlib.contextmanager
+def autotune_bf16_only(
+    tuning_buckets: tuple[int, ...], *, skip_ops: set[str] | None = None
+) -> Iterator[None]:
+    """Tune BF16 calls with bounded buckets, outside full-model autotuning."""
+    token = _bf16_autotune_buckets.set(tuning_buckets)
+    try:
+        with autotune(tune_mode=False, skip_ops=skip_ops):
+            yield
+    finally:
+        _bf16_autotune_buckets.reset(token)
+
+
 def flashinfer_bf16_mm(
     a: torch.Tensor,
     b: torch.Tensor,
@@ -35,14 +54,21 @@ def flashinfer_bf16_mm(
 ) -> torch.Tensor:
     from flashinfer import mm_bf16
 
-    return mm_bf16(
-        a,
-        b,
-        bias=bias,
-        pdl=pdl,
-        out_dtype=torch.bfloat16,
-        backend=backend,
+    tuning_buckets = _bf16_autotune_buckets.get()
+    tuning = (
+        autotune(tune_mode=True, tuning_buckets=tuning_buckets)
+        if tuning_buckets is not None
+        else contextlib.nullcontext()
     )
+    with tuning:
+        return mm_bf16(
+            a,
+            b,
+            bias=bias,
+            pdl=pdl,
+            out_dtype=torch.bfloat16,
+            backend=backend,
+        )
 
 
 # This is the storage path for the cubins, it can be replaced
@@ -589,8 +615,7 @@ def supports_trtllm_attention(is_prefill: bool = False) -> bool:
 
 
 def force_use_trtllm_attention() -> bool | None:
-    """
-    This function should only be called during initialization stage when vllm config
+    """This function should only be called during initialization stage when vllm config
     is set.
     Return `None` if --attention-config.use_trtllm_attention is not set,
     return `True` if TRTLLM attention is forced to be used,
@@ -628,7 +653,6 @@ def use_trtllm_attention(
     has_spec: bool = False,
 ) -> bool:
     """Return `True` if TRTLLM attention is used."""
-
     # CLI argument is set to 0 - respect it
     if force_use_trtllm is not None and not force_use_trtllm:
         return False
@@ -736,6 +760,7 @@ def _flashinfer_concat_mla_k(
             k_nope: The nope part of k, shape [num_tokens, num_heads, nope_dim].
             k_pe: The rope part of k (shared), shape [num_tokens, 1, rope_dim].
                   This is broadcast to all heads.
+
         """
         from flashinfer.concat_ops import concat_mla_k
 
@@ -1238,6 +1263,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
 __all__ = [
     "has_flashinfer",
     "flashinfer_bf16_mm",
+    "autotune_bf16_only",
     "has_flashinfer_bf16_gemm",
     "is_flashinfer_bf16_gemm_supported",
     "is_flashinfer_cutedsl_bf16_gemm_supported",
diff --git a/vllm/utils/func_utils.py b/vllm/utils/func_utils.py
index 5ce23e6a0074..69600decda49 100644
--- a/vllm/utils/func_utils.py
+++ b/vllm/utils/func_utils.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Contains helpers that are applied to functions.
+"""Contains helpers that are applied to functions.
 
 This is similar in concept to the `functools` module.
 """
@@ -127,8 +126,7 @@ def get_allowed_kwarg_only_overrides(
     requires_kw_only: bool = True,
     allow_var_kwargs: bool = False,
 ) -> dict[str, Any]:
-    """
-    Given a callable which has one or more keyword only params and a dict
+    """Given a callable which has one or more keyword only params and a dict
     mapping param names to values, drop values that can be not be kwarg
     expanded to overwrite one or more keyword-only args. This is used in a
     few places to handle custom processor overrides for multimodal models,
@@ -139,12 +137,14 @@ def get_allowed_kwarg_only_overrides(
         callable: Callable which takes 0 or more keyword only arguments.
                   If None is provided, all overrides names are allowed.
         overrides: Potential overrides to be used when invoking the callable.
+        requires_kw_only: If True, only keyword-only params may be overridden.
         allow_var_kwargs: Allows overrides that are expandable for var kwargs.
 
     Returns:
         Dictionary containing the kwargs to be leveraged which may be used
         to overwrite one or more keyword only arguments when invoking the
         callable.
+
     """
     if not overrides:
         return {}
diff --git a/vllm/utils/gc_utils.py b/vllm/utils/gc_utils.py
index 32bdfb764861..be6ab4a23fa1 100644
--- a/vllm/utils/gc_utils.py
+++ b/vllm/utils/gc_utils.py
@@ -14,8 +14,7 @@
 
 
 class GCDebugConfig:
-    """
-    Config for GC Debugger.
+    """Config for GC Debugger.
     - 0: disable GC debugger
     - 1: enable GC debugger with gc.collect elapsed times
     - '{"top_objects":5}': enable GC debugger with top 5 collected objects
@@ -44,8 +43,7 @@ def __repr__(self) -> str:
 
 
 class GCDebugger:
-    """
-    Debugger for GC which logs helpful information for GC understanding.
+    """Debugger for GC which logs helpful information for GC understanding.
     To enable, you should call maybe_attach_gc_debug_callback in the process.
     """
 
@@ -59,9 +57,7 @@ def __init__(self, config: GCDebugConfig) -> None:
         self.gc_top_collected_objects: str = ""
 
     def handle(self, phase: str, info: dict[str, int]) -> None:
-        """
-        Handles a GC event (e.g. GC start or GC finish)
-        """
+        """Handles a GC event (e.g. GC start or GC finish)."""
         generation = info.get("generation")
         if generation is None:
             return
@@ -122,8 +118,7 @@ def freeze_gc_for_cudagraph_capture():
 
 
 def freeze_gc_heap() -> None:
-    """
-    Freeze all objects tracked by the garbage collector. It should be invoked
+    """Freeze all objects tracked by the garbage collector. It should be invoked
     after server init / warmup, to reduce GC overhead from static objects
     during serving time.
     """
@@ -137,9 +132,7 @@ def freeze_gc_heap() -> None:
 
 
 def maybe_attach_gc_debug_callback() -> None:
-    """
-    Attached a callback for GC debug when VLLM_GC_DEBUG is enabled.
-    """
+    """Attached a callback for GC debug when VLLM_GC_DEBUG is enabled."""
     config = GCDebugConfig(envs.VLLM_GC_DEBUG)
     if config.enabled:
         debugger: GCDebugger = GCDebugger(config)
@@ -151,8 +144,7 @@ def gc_callback(phase: str, info: dict[str, int]) -> None:
 
 
 def _compute_detailed_type(o: Any) -> str:
-    """
-    Detailed object type.
+    """Detailed object type.
 
     TODO(Jialin): Further enhance the detailed type with element types for
     easier debugging. We tried but occasionally it would run into signals
@@ -167,9 +159,7 @@ def _compute_detailed_type(o: Any) -> str:
 
 
 def _compute_top_gc_collected_objects(objects: list[Any], top: int) -> str:
-    """
-    Group collected objects by types.
-    """
+    """Group collected objects by types."""
     if top <= 0:
         return ""
     object_types = [_compute_detailed_type(o) for o in objects]
diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py
index 53fcdfe71004..8d0e7ae1871f 100644
--- a/vllm/utils/gpu_sync_debug.py
+++ b/vllm/utils/gpu_sync_debug.py
@@ -225,6 +225,16 @@ def _install_compile_time_sync_suppressors() -> None:
 _torch_cuda = torch.Tensor.cuda
 _torch_copy_ = torch.Tensor.copy_
 
+
+# Once ``Tensor.to`` is replaced below, Dynamo sees the saved C++ descriptor
+# through this module global instead of as a recognized tensor method. Put the
+# descriptor behind a weak-referenceable Python callable so ``allow_in_graph``
+# can admit it and the downstream compiler can lower the original operation.
+@torch.compiler.allow_in_graph
+def _torch_to_in_graph(self, *args, **kwargs):
+    return _torch_to(self, *args, **kwargs)
+
+
 # torch's own parser for `Tensor.to`'s overloaded signature, returning
 # (device, dtype, non_blocking, memory_format) (also used by torch._dynamo).
 _parse_to = getattr(torch._C._nn, "_parse_to", None)
@@ -232,7 +242,7 @@ def _install_compile_time_sync_suppressors() -> None:
 
 def _active_check_mode() -> str | None:
     """The mode the calling thread is being checked in, if any."""
-    if not _sync_check_enabled or _allow_depth.get() or torch.compiler.is_compiling():
+    if not _sync_check_enabled or torch.compiler.is_compiling() or _allow_depth.get():
         return None
     return _checking.get()
 
@@ -296,7 +306,7 @@ def _checked_to(self, *args, **kwargs):
         else:
             if non_blocking and (reason := _cpu_copy_stall_reason(self, device)):
                 _report_implicit_sync("to", "H2D", reason, mode)
-    return _torch_to(self, *args, **kwargs)
+    return _torch_to_in_graph(self, *args, **kwargs)
 
 
 def _checked_cuda(self, *args, **kwargs):
diff --git a/vllm/utils/hashing.py b/vllm/utils/hashing.py
index f01c6b074ffe..9caee37493ca 100644
--- a/vllm/utils/hashing.py
+++ b/vllm/utils/hashing.py
@@ -35,6 +35,7 @@ def sha256(input: Any) -> bytes:
 
     Returns:
         Bytes representing the SHA-256 hash of the serialized input.
+
     """
     input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
     return hashlib.sha256(input_bytes).digest()
@@ -53,6 +54,7 @@ def sha256_cbor(input: Any) -> bytes:
 
     Returns:
         Bytes representing the SHA-256 hash of the CBOR serialized input.
+
     """
     input_bytes = cbor2.dumps(input, canonical=True)
     return hashlib.sha256(input_bytes).digest()
@@ -87,6 +89,7 @@ def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]:
 
     Returns:
         A hash function.
+
     """
     if hash_fn_name == "sha256":
         return sha256
@@ -110,6 +113,7 @@ def safe_hash(data: bytes, usedforsecurity: bool = True) -> HASH:
 
     Returns:
         Hash object
+
     """
     try:
         return hashlib.md5(data, usedforsecurity=usedforsecurity)
diff --git a/vllm/utils/humming.py b/vllm/utils/humming.py
index 014603ff39c3..574a00f16b8a 100644
--- a/vllm/utils/humming.py
+++ b/vllm/utils/humming.py
@@ -4,12 +4,56 @@
 
 vLLM code should import humming symbols from here so that ``import humming``
 (which has import-time side effects) is deferred until first use. Add new
-symbols by appending one entry to ``_EXPORTS`` as ``"module.path:attr"``,
-or ``"module.path"`` for a whole-module re-export.
+symbols to both the ``TYPE_CHECKING`` imports and ``_EXPORTS``, using
+``"module.path:attr"`` or ``"module.path"`` for a whole-module re-export.
 """
 
 import importlib
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from humming import dtypes as dtypes
+    from humming.config import GemmType as GemmType
+    from humming.config import InputQuantizationMode as InputQuantizationMode
+    from humming.config import LayerConfig as LayerConfig
+    from humming.config import MmaType as MmaType
+    from humming.config import WeightScale2Type as WeightScale2Type
+    from humming.config import WeightScaleType as WeightScaleType
+    from humming.dtypes import DataType as DataType
+    from humming.forward import humming_forward as humming_forward
+    from humming.forward import may_process_input as may_process_input
+    from humming.forward import may_quant_input as may_quant_input
+    from humming.ops import process_input as process_input
+    from humming.schema import AWQWeightSchema as AWQWeightSchema
+    from humming.schema import BaseInputSchema as BaseInputSchema
+    from humming.schema import BaseWeightSchema as BaseWeightSchema
+    from humming.schema import BitnetWeightSchema as BitnetWeightSchema
+    from humming.schema import (
+        CompressedTensorsInputSchema as CompressedTensorsInputSchema,
+    )
+    from humming.schema import (
+        CompressedTensorsWeightSchema as CompressedTensorsWeightSchema,
+    )
+    from humming.schema import Fp8InputSchema as Fp8InputSchema
+    from humming.schema import GptOssMxfp4WeightSchema as GptOssMxfp4WeightSchema
+    from humming.schema import GPTQWeightSchema as GPTQWeightSchema
+    from humming.schema import HummingInputSchema as HummingInputSchema
+    from humming.schema import HummingWeightSchema as HummingWeightSchema
+    from humming.schema import Mxfp4WeightSchema as Mxfp4WeightSchema
+    from humming.schema.fp8 import Fp8WeightSchema as Fp8WeightSchema
+    from humming.schema.modelopt import (
+        ModeloptMxfp8WeightSchema as ModeloptMxfp8WeightSchema,
+    )
+    from humming.schema.modelopt import (
+        ModeloptNvfp4InputSchema as ModeloptNvfp4InputSchema,
+    )
+    from humming.schema.modelopt import (
+        ModeloptNvfp4WeightSchema as ModeloptNvfp4WeightSchema,
+    )
+    from humming.transform import prepare_layer_config as prepare_layer_config
+    from humming.transform import transform_humming_tensors as transform_humming_tensors
+    from humming.tune import get_heuristics_config as get_heuristics_config
+    from humming.utils.weight import quantize_weight as quantize_weight
 
 _EXPORTS: dict[str, str] = {
     "dtypes": "humming.dtypes",
@@ -17,7 +61,12 @@
     "GemmType": "humming.config:GemmType",
     "LayerConfig": "humming.config:LayerConfig",
     "WeightScaleType": "humming.config:WeightScaleType",
+    "WeightScale2Type": "humming.config:WeightScale2Type",
+    "InputQuantizationMode": "humming.config:InputQuantizationMode",
+    "MmaType": "humming.config:MmaType",
     "humming_forward": "humming.forward:humming_forward",
+    "may_process_input": "humming.forward:may_process_input",
+    "process_input": "humming.ops:process_input",
     "may_quant_input": "humming.forward:may_quant_input",
     "prepare_layer_config": "humming.transform:prepare_layer_config",
     "transform_humming_tensors": "humming.transform:transform_humming_tensors",
diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py
index 20c510edc025..aaacccc7cb5a 100644
--- a/vllm/utils/import_utils.py
+++ b/vllm/utils/import_utils.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Contains helpers related to importing modules.
+"""Contains helpers related to importing modules.
 
 This is similar in concept to the `importlib` module.
 """
@@ -24,8 +23,7 @@
 
 
 def import_pynvml():
-    """
-    Historical comments:
+    """Historical comments:
 
     libnvml.so is the library behind nvidia-smi, and
     pynvml is a Python wrapper around it. We use it to get GPU
@@ -57,8 +55,7 @@ def import_pynvml():
 
 @cache
 def import_triton_kernels():
-    """
-    For convenience, prioritize triton_kernels that is available in
+    """For convenience, prioritize triton_kernels that is available in
     `site-packages`. Use `vllm.third_party.triton_kernels` as a fall-back.
     """
     if _has_module("triton_kernels"):
@@ -83,8 +80,7 @@ def import_triton_kernels():
 
 
 def import_plugin(plugin_path: str) -> ModuleType | None:
-    """
-    Import a user-defined plugin.
+    """Import a user-defined plugin.
 
     Plugin can be either:
     * a module in site-packages
@@ -104,8 +100,7 @@ def import_plugin(plugin_path: str) -> ModuleType | None:
 
 
 def import_from_path(module_name: str, file_path: str | os.PathLike) -> ModuleType:
-    """
-    Import a Python file according to its file path.
+    """Import a Python file according to its file path.
 
     Based on the official recipe:
     https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
@@ -123,9 +118,7 @@ def import_from_path(module_name: str, file_path: str | os.PathLike) -> ModuleTy
 
 
 def resolve_obj_by_qualname(qualname: str) -> Any:
-    """
-    Resolve an object by its fully-qualified class name.
-    """
+    """Resolve an object by its fully-qualified class name."""
     module_name, obj_name = qualname.rsplit(".", 1)
     module = importlib.import_module(module_name)
     return getattr(module, obj_name)
@@ -148,8 +141,7 @@ def get_vllm_optional_dependencies():
 
 
 class _PlaceholderBase:
-    """
-    Disallows downstream usage of placeholder modules.
+    """Disallows downstream usage of placeholder modules.
 
     We need to explicitly override each dunder method because
     [`__getattr__`][vllm.utils.import_utils._PlaceholderBase.__getattr__]
@@ -160,8 +152,7 @@ class _PlaceholderBase:
     """
 
     def __getattr__(self, key: str) -> Never:
-        """
-        The main class should implement this to throw an error
+        """The main class should implement this to throw an error
         for attribute accesses representing downstream usage.
         """
         raise NotImplementedError
@@ -304,8 +295,7 @@ def __exit__(self, *args: object, **kwargs: object):
 
 
 class PlaceholderModule(_PlaceholderBase):
-    """
-    A placeholder object to use when a module does not exist.
+    """A placeholder object to use when a module does not exist.
 
     This enables more informative errors when trying to access attributes
     of a module that does not exist.
@@ -360,8 +350,7 @@ def __getattr__(self, key: str) -> Never:
 
 
 class LazyLoader(ModuleType):
-    """
-    `LazyLoader` module borrowed from [Tensorflow]
+    """`LazyLoader` module borrowed from [Tensorflow]
     (https://github.com/tensorflow/tensorflow/blob/main/tensorflow/python/util/lazy_loader.py)
     with an addition of "module caching".
 
@@ -523,6 +512,34 @@ def has_nixl_ep() -> bool:
     return _has_module("nixl_ep")
 
 
+def has_moonep() -> bool:
+    """Whether the optional `moonep` package is available."""
+    return _has_module("moonep")
+
+
+def check_moonep_system_support() -> None:
+    """Raise if the current device cannot run MoonEP.
+
+    MoonEP's symmetric-memory buffers require CUDA VMM plus NVSwitch multicast
+    (SHARP) on every EP rank; NVLink-only topologies without NVSwitch (e.g.
+    4x H100 NV6) fail deep inside ``moonep.Buffer`` otherwise.
+    """
+    if not has_moonep():
+        raise RuntimeError(
+            "MoonEP not available. Install it from "
+            "https://github.com/MoonshotAI/MoonEP."
+        )
+    from moonep._C import (  # type: ignore[import-not-found]
+        nvl_multicast_supported,
+    )
+
+    if not nvl_multicast_supported():
+        raise RuntimeError(
+            "MoonEP requires NVSwitch multicast (SHARP) support on the "
+            "current device; it is not supported on this GPU/topology."
+        )
+
+
 def is_numba_available() -> bool:
     """Whether the optional `numba` package is available."""
     return _has_module("numba")
@@ -538,6 +555,26 @@ def has_triton_kernels() -> bool:
     return is_available
 
 
+@cache
+def get_triton_kernels_version() -> str | None:
+    """The triton_kernels MoE-API generation ("3.5.1"/"3.6"/"3.8"), or None.
+
+    Inferred by capability since the package exposes no usable version: 3.8
+    replaced ``matmul_ogs`` with ``matmul``, and 3.5.1 predates ``SparseMatrix``.
+    """
+    if not has_triton_kernels():
+        return None
+    try:
+        import triton_kernels.matmul_ogs  # noqa: F401
+    except ImportError:
+        return "3.8"
+    try:
+        from triton_kernels.tensor import SparseMatrix  # noqa: F401
+    except ImportError:
+        return "3.5.1"
+    return "3.6"
+
+
 @cache
 def has_tilelang() -> bool:
     """Whether the optional `tilelang` package is available.
@@ -562,7 +599,6 @@ def has_tilelang() -> bool:
 
 def has_arctic_inference() -> bool:
     """Whether the optional `arctic_inference` package is available."""
-
     return _has_module("arctic_inference")
 
 
diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py
index bb0bb8df9a1e..ffb63bc493d4 100644
--- a/vllm/utils/jit_monitor.py
+++ b/vllm/utils/jit_monitor.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Monitor unexpected kernel JIT compilation during inference.
+"""Monitor unexpected kernel JIT compilation during inference.
 
 After server warmup completes, any kernel JIT compilation or autotuning event
 indicates a cache miss or unexpected input shape that causes a latency spike.
diff --git a/vllm/utils/jsontree.py b/vllm/utils/jsontree.py
index fe757c2f3374..0f430935acdf 100644
--- a/vllm/utils/jsontree.py
+++ b/vllm/utils/jsontree.py
@@ -142,8 +142,7 @@ def json_reduce_leaves(
     initial: _U = ...,  # type: ignore[assignment]
     /,
 ) -> _T | _U:
-    """
-    Apply a function of two arguments cumulatively to each leaf in a
+    """Apply a function of two arguments cumulatively to each leaf in a
     nested JSON structure, from left to right, so as to reduce the
     sequence to a single value.
     """
diff --git a/vllm/utils/math_utils.py b/vllm/utils/math_utils.py
index c681e708b5c1..8fb0009d60bf 100644
--- a/vllm/utils/math_utils.py
+++ b/vllm/utils/math_utils.py
@@ -13,7 +13,7 @@ def cdiv(a: int, b: int) -> int:
 
 
 def next_power_of_2(n: int) -> int:
-    """The next power of 2 (inclusive)"""
+    """The next power of 2 (inclusive)."""
     return 1 if n < 1 else 1 << (n - 1).bit_length()
 
 
diff --git a/vllm/utils/mem_utils.py b/vllm/utils/mem_utils.py
index c092dff41045..1b449058b27d 100644
--- a/vllm/utils/mem_utils.py
+++ b/vllm/utils/mem_utils.py
@@ -54,6 +54,7 @@ def release_device_memory_under_pressure(device: torch.device) -> bool:
 
     Returns:
         True if memory was released.
+
     """
     if device.type != "cuda" or not current_platform.is_integrated_gpu(device.index):
         return False
@@ -230,8 +231,7 @@ def memory_profiling(
     baseline_snapshot: MemorySnapshot,
     weights_memory: int = 0,
 ) -> Generator[MemoryProfilingResult, None, None]:
-    """
-    Memory profiling context manager.
+    """Memory profiling context manager.
 
     baseline_snapshot: the memory snapshot before the current vLLM instance.
     weights_memory: memory used by PyTorch when loading the model weights.
diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py
index a8286c84804e..6190b7699178 100644
--- a/vllm/utils/multi_stream_utils.py
+++ b/vllm/utils/multi_stream_utils.py
@@ -41,6 +41,7 @@ def maybe_execute_in_parallel(
 
     Returns:
         Tuple of (fn0_result, fn1_result).
+
     """
     if aux_stream is not None:
         from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
@@ -100,6 +101,7 @@ def execute_in_parallel(
     Returns:
         Tuple of (default_result, aux_results) where aux_results[i] is the
         result of aux_fns[i] (or None when skipped).
+
     """
     aux_results: list[Any]
     if aux_streams is None or not enable:
diff --git a/vllm/utils/network_utils.py b/vllm/utils/network_utils.py
index 356d913879a1..fa24d4d5f9ab 100644
--- a/vllm/utils/network_utils.py
+++ b/vllm/utils/network_utils.py
@@ -172,8 +172,7 @@ def _get_reserved_port_range() -> range:
 
 
 def get_open_port() -> int:
-    """
-    Get an open port for the vLLM process to listen on.
+    """Get an open port for the vLLM process to listen on.
     An edge case to handle, is when we run data parallel,
     we need to avoid ports that are potentially used by
     the data parallel master process.
@@ -298,6 +297,7 @@ def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str:
 
     Returns:
         A properly formatted ZMQ path string.
+
     """
     if port is None:
         return f"{scheme}://{host}"
@@ -317,7 +317,6 @@ def make_zmq_socket(
     router_handover: bool = False,
 ) -> zmq.Socket | zmq.asyncio.Socket:  # type: ignore[name-defined]
     """Make a ZMQ socket with the proper bind/connect semantics."""
-
     mem = psutil.virtual_memory()
     socket = ctx.socket(socket_type)
 
@@ -377,8 +376,7 @@ def zmq_socket_ctx(
     identity: bytes | None = None,
     router_handover: bool = False,
 ) -> Iterator[zmq.Socket]:
-    """Context manager for a ZMQ socket"""
-
+    """Context manager for a ZMQ socket."""
     ctx = zmq.Context()  # type: ignore[attr-defined]
     try:
         yield make_zmq_socket(
diff --git a/vllm/utils/nvtx_pytorch_hooks.py b/vllm/utils/nvtx_pytorch_hooks.py
index 25f76edbfaa2..5d979b04cf84 100644
--- a/vllm/utils/nvtx_pytorch_hooks.py
+++ b/vllm/utils/nvtx_pytorch_hooks.py
@@ -25,7 +25,7 @@ def print_tensor(tensor_obj, prefix, tensor_list=None):
 
 
 def process_layer_params(module_obj):
-    """Extract the static parameters from LLM and VLM relevant layer types"""
+    """Extract the static parameters from LLM and VLM relevant layer types."""
     param_info = {}
     # Extract parameters for layers commonly used in LLMs and VLMs
     if isinstance(module_obj, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)):
@@ -186,6 +186,7 @@ def layerwise_nvtx_marker_context(module_name, module_obj, in_tensor=None, kwarg
                                  kwargs=kwargs) as ctx:
             ctx.result = module(*args, **kwargs)
         return ctx.result
+
     """
     holder = ResultHolder()
 
@@ -221,9 +222,9 @@ class PytHooks:
     a PytHook object. Then call the register_hooks method.
 
     Example:
-
         my_hook = PytHook()
         my_hook.register_hooks(my_network_model)
+
     """
 
     def __init__(self):
diff --git a/vllm/utils/ompmultiprocessing.py b/vllm/utils/ompmultiprocessing.py
index 0a704e6d2676..34395489fb44 100644
--- a/vllm/utils/ompmultiprocessing.py
+++ b/vllm/utils/ompmultiprocessing.py
@@ -46,12 +46,12 @@ def __init__(self, config: "VllmConfig"):
         self.reserve_cpu_num = (
             self.local_world_size
             if current_platform.get_cpu_architecture()
-            in (CpuArchEnum.ARM, CpuArchEnum.RISCV)
+            in (CpuArchEnum.X86, CpuArchEnum.ARM, CpuArchEnum.RISCV)
             else 1
         )
         # reserve at one more core for nixl_connector under p/d case
         if config.kv_transfer_config:
-            self.reserve_cpu_num += 1
+            self.reserve_cpu_num += self.local_world_size
 
         if envs.VLLM_CPU_NUM_OF_RESERVED_CPU is not None:
             if self.reserve_cpu_num > envs.VLLM_CPU_NUM_OF_RESERVED_CPU:
@@ -187,18 +187,19 @@ def _parse_omp_threads_bind_env(self):
     def _get_autobind_cpu_ids(
         self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]]
     ) -> tuple[list[list[LogicalCPUInfo]], list[LogicalCPUInfo]]:
-        """
-        Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for
+        """Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for
         other processes.
         Currently for rank N, only CPU ids on the N-th node in available NUMA
         node list will be selected.
+
         Args:
             cpu_selector: a callable object to select CPUs from a CPU list
-            of a physical core. The input is a LogicalCPUInfo list contains
-            logical CPUs of a physical CPU, sorted by the LogicalCPUInfo.id.
-            A selected LogicalCPUInfo list should be returned.
-        """
+                of a physical core. The input is a LogicalCPUInfo list contains
+                logical CPUs of a physical CPU, sorted by the
+                LogicalCPUInfo.id. A selected LogicalCPUInfo list should be
+                returned.
 
+        """
         # this memory node list has been sliced for DP offset
         allowed_numa_nodes = cr_utils.get_visible_memory_node()
         logical_cpu_list = cr_utils.get_allowed_cpu_list()
diff --git a/vllm/utils/registry.py b/vllm/utils/registry.py
index a136d450e7b1..6df0bead0db1 100644
--- a/vllm/utils/registry.py
+++ b/vllm/utils/registry.py
@@ -6,8 +6,7 @@
 
 
 class ExtensionManager:
-    """
-    A registry for managing pluggable extension classes.
+    """A registry for managing pluggable extension classes.
 
     This class provides a simple mechanism to register and instantiate
     extension classes by name. It is commonly used to implement plugin
@@ -26,15 +25,11 @@ class ExtensionManager:
     """
 
     def __init__(self) -> None:
-        """
-        Initialize an empty extension registry.
-        """
+        """Initialize an empty extension registry."""
         self.name2class: dict[str, type] = {}
 
     def register(self, name: str):
-        """
-        Decorator to register a class with the given name.
-        """
+        """Decorator to register a class with the given name."""
 
         def wrap(cls_to_register: _T) -> _T:
             self.name2class[name] = cls_to_register
@@ -43,9 +38,7 @@ def wrap(cls_to_register: _T) -> _T:
         return wrap
 
     def load(self, cls_name: str, *args, **kwargs) -> Any:
-        """
-        Instantiate and return a registered extension class by name.
-        """
+        """Instantiate and return a registered extension class by name."""
         cls = self.name2class.get(cls_name)
         assert cls is not None, f"Extension class {cls_name} not found"
         return cls(*args, **kwargs)
diff --git a/vllm/utils/sparse_utils.py b/vllm/utils/sparse_utils.py
index 79330271dd51..582d91093eae 100644
--- a/vllm/utils/sparse_utils.py
+++ b/vllm/utils/sparse_utils.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Safe deserialization helpers for client-supplied embedding tensors.
+"""Safe deserialization helpers for client-supplied embedding tensors.
 
 `check_sparse_tensor_invariants_threadsafe()` is a thread-safe wrapper for
 sparse tensor invariant validation. `safe_to_dense()` bounds the memory a
diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py
index 6d56df378f13..cc83a62a1690 100644
--- a/vllm/utils/system_utils.py
+++ b/vllm/utils/system_utils.py
@@ -60,8 +60,7 @@ def set_env_var(key: str, value: str) -> Iterator[None]:
 
 @contextlib.contextmanager
 def suppress_stdout():
-    """
-    Suppress stdout from C libraries at the file descriptor level.
+    """Suppress stdout from C libraries at the file descriptor level.
 
     Only suppresses stdout, not stderr, to preserve error messages.
     Suppression is disabled when VLLM_LOGGING_LEVEL is set to DEBUG.
@@ -70,6 +69,7 @@ def suppress_stdout():
         with suppress_stdout():
             # C library calls that would normally print to stdout
             torch.distributed.new_group(ranks, backend="gloo")
+
     """
     # Don't suppress if logging level is DEBUG
     if envs.VLLM_LOGGING_LEVEL == "DEBUG":
@@ -114,7 +114,6 @@ def unique_filepath(fn: Callable[[int], Path]) -> Path:
 
 def _sync_visible_devices_env_vars():
     """Sync HIP/CUDA visibility env vars before spawning (ROCm only)."""
-
     if not current_platform.is_rocm():
         return
 
@@ -254,11 +253,11 @@ def decorate_logs(
 
 
 def kill_process_tree(pid: int):
-    """
-    Kills all descendant processes of the given pid by sending SIGKILL.
+    """Kills all descendant processes of the given pid by sending SIGKILL.
 
     Args:
         pid (int): Process ID of the parent process
+
     """
     try:
         parent = psutil.Process(pid)
diff --git a/vllm/utils/tensor_schema.py b/vllm/utils/tensor_schema.py
index 526dfd38bac4..02ddda03f1fc 100644
--- a/vllm/utils/tensor_schema.py
+++ b/vllm/utils/tensor_schema.py
@@ -161,7 +161,6 @@ def _validate_tensor_shape_expected(
         dynamic_dims: set[str],
     ) -> None:
         """Validate that the actual tensor shape matches the expected shape."""
-
         if len(actual_shape) != len(expected_shape):
             raise ValueError(
                 f"{field_name} has rank {len(actual_shape)} "
diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py
index b453487a01a8..849d0a38f55c 100644
--- a/vllm/utils/torch_utils.py
+++ b/vllm/utils/torch_utils.py
@@ -96,8 +96,7 @@ def is_meta_module(module: torch.nn.Module) -> bool:
 
 
 def is_strictly_contiguous(t: torch.Tensor) -> bool:
-    """
-    Check if tensor is contiguous AND has no degenerate strides.
+    """Check if tensor is contiguous AND has no degenerate strides.
 
     A degenerate stride occurs when a dimension has size 1 but the stride
     doesn't match the canonical contiguous layout. This can cause issues
@@ -294,8 +293,7 @@ def set_torch_threads_for_runtime() -> None:
 
 @contextlib.contextmanager
 def set_default_torch_num_threads(num_threads: int | None = None):
-    """
-    Sets the default number of threads for PyTorch to the given value.
+    """Sets the default number of threads for PyTorch to the given value.
 
     `None` means using the value of the environment variable `OMP_NUM_THREADS`
     (or `1` if that is not available).
@@ -363,8 +361,7 @@ def _get_precision_level(dtype: torch.dtype) -> int:
 
 
 def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype):
-    """
-    Test whether it is lossless to cast a tensor from
+    """Test whether it is lossless to cast a tensor from
     `src_dtype` to `tgt_dtype`.
     """
     if src_dtype == tgt_dtype:
@@ -395,8 +392,7 @@ def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype):
 
 
 def common_broadcastable_dtype(dtypes: Collection[torch.dtype]):
-    """
-    Get the common `dtype` where all of the other `dtypes` can be
+    """Get the common `dtype` where all of the other `dtypes` can be
     cast to it without losing any information.
     """
     return max(
@@ -567,6 +563,7 @@ def nvfp4_split_data_scale(
     Returns:
         ``(data, scale)`` where *data* is uint8 and *scale* is
         float8_e4m3fn, both views of the same storage.
+
     """
     num_pages = kv_side.shape[0]
     dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2]
@@ -705,17 +702,30 @@ def create_kv_caches_with_random(
 
 def async_tensor_h2d(
     data: list | np.ndarray | torch.Tensor,
-    device: str | torch.device,
+    device: str | torch.device | None = None,
     dtype: torch.dtype | None = None,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Copy list/numpy array/tensor async from host to device."""
+    if dtype is None and out is not None:
+        dtype = out.dtype
     if isinstance(data, np.ndarray):
         data = torch.from_numpy(data)
     if isinstance(data, torch.Tensor):
-        t = data.pin_memory() if PIN_MEMORY else data
+        t = data
+        if PIN_MEMORY and not t.is_pinned():
+            # Stage in pinned, contiguous buffer to ensure fully async copy.
+            t = torch.empty(
+                t.shape, dtype=dtype or t.dtype, device="cpu", pin_memory=True
+            ).copy_(t)
     else:
         t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu")
     assert t.is_cpu
+
+    if out is not None:
+        assert out.dtype == dtype
+        return out.copy_(t, non_blocking=True)
+    assert device is not None, "must provide destination tensor or device"
     return t.to(device=device, dtype=dtype, non_blocking=True)
 
 
@@ -731,8 +741,7 @@ def make_ndarray_with_pad(
     *,
     max_len: int | None = None,
 ) -> npt.NDArray:
-    """
-    Make a padded array from 2D inputs.
+    """Make a padded array from 2D inputs.
 
     The padding is applied to the end of each inner list until it reaches
     `max_len`.
@@ -758,8 +767,7 @@ def make_tensor_with_pad(
     device: str | torch.device | None = None,
     pin_memory: bool = False,
 ) -> torch.Tensor:
-    """
-    Make a padded tensor from 2D inputs.
+    """Make a padded tensor from 2D inputs.
 
     The padding is applied to the end of each inner list until it reaches
     `max_len`.
@@ -793,8 +801,7 @@ def __init__(self):
 
 
 def current_stream() -> torch.cuda.Stream:
-    """
-    replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
+    """Replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
     it turns out that `torch.cuda.current_stream()` is quite expensive,
     as it will construct a new stream object at each call.
     here we patch `torch.cuda.set_stream` to keep track of the current stream
@@ -816,8 +823,11 @@ def current_stream() -> torch.cuda.Stream:
         # https://github.com/pytorch/pytorch/blob/42ad9edfb754743fdae3276ade43de000beb4f60/aten/src/ATen/cuda/CUDAGraph.cpp#L77
         # for more details. Therefore, we create a dedicated stream per process.
         if current_platform.is_rocm() or current_platform.is_cuda():
+            # Ensure new stream is ordered w.r.t. replaced stream's work.
+            new_stream = torch.cuda.Stream()
+            new_stream.wait_stream(torch.cuda.current_stream())
             # torch.cuda.set_stream here is the alias of _pathed_set_stream
-            torch.cuda.set_stream(torch.cuda.Stream())
+            torch.cuda.set_stream(new_stream)
         elif current_platform.is_cpu():
             _current_stream_tls.value = _StreamPlaceholder()
         else:
@@ -842,9 +852,7 @@ def current_stream() -> torch.cuda.Stream:
 
 
 def aux_stream() -> torch.cuda.Stream | None:
-    """
-    Ensures aux_stream is initialized only once
-    """
+    """Ensures aux_stream is initialized only once."""
     global _aux_stream
 
     from vllm.platforms import current_platform
@@ -856,8 +864,7 @@ def aux_stream() -> torch.cuda.Stream | None:
 
 
 def weak_ref_tensor(tensor: Any) -> Any:
-    """
-    Create a weak reference to a tensor.
+    """Create a weak reference to a tensor.
     The new tensor will share the same data as the original tensor,
     but will not keep the original tensor alive.
     This ignores 0-size tensors as those don't allocate any memory.
@@ -874,8 +881,7 @@ def weak_ref_tensors(
     | tuple[torch.Tensor]
     | IntermediateTensors,
 ) -> torch.Tensor | list[Any] | tuple[Any] | Any:
-    """
-    Convenience function to create weak references to tensors,
+    """Convenience function to create weak references to tensors,
     for single tensor, list of tensors or tuple of tensors.
     """
     if isinstance(tensors, torch.Tensor):
@@ -897,9 +903,7 @@ def weak_ref_tensors(
 
 
 def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
-    """
-    Get an accelerator view of a CPU tensor using Unified Virtual Addressing (UVA).
-    """
+    """Get an accelerator view of a CPU tensor via Unified Virtual Addressing."""
     from vllm.platforms import current_platform
 
     if current_platform.is_xpu():
@@ -935,6 +939,7 @@ def is_torch_equal_or_newer(target: str) -> bool:
 
     Returns:
         Whether the condition meets.
+
     """
     try:
         return _is_torch_equal_or_newer(str(torch.__version__), target)
@@ -963,6 +968,7 @@ def is_torch_equal(target: str) -> bool:
 
     Returns:
         Whether the condition meets.
+
     """
     try:
         return _is_torch_equal(target)
@@ -1048,8 +1054,7 @@ def direct_register_custom_op(
     dispatch_key: str | None = None,
     tags: tuple[torch.Tag, ...] = (),
 ):
-    """
-    `torch.library.custom_op` can have significant overhead because it
+    """`torch.library.custom_op` can have significant overhead because it
     needs to consider complicated dispatching logic. This function
     directly registers a custom op and dispatches it to the CUDA backend.
     See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py
index fcc937fe8a67..8b51e3916986 100644
--- a/vllm/v1/attention/backend.py
+++ b/vllm/v1/attention/backend.py
@@ -8,7 +8,6 @@
 
 import numpy as np
 import torch
-from typing_extensions import deprecated
 
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8Dynamic64Sym,
@@ -23,6 +22,7 @@
     from vllm.model_executor.layers.linear import ColumnParallelLinear
     from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
     from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.hisparse.runtime import HiSparseCacheHandle
     from vllm.v1.kv_cache_interface import (
         AttentionSpec,
         KVCacheLayout,
@@ -34,8 +34,7 @@
 
 
 class AttentionType(str, Enum):
-    """
-    Attention type.
+    """Attention type.
     Use string to be compatible with `torch.compile`.
     """
 
@@ -70,7 +69,9 @@ class AttentionBackend(ABC):
     forward_includes_kv_cache_update: bool = True
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(
+        kv_cache_spec: "KVCacheSpec | None" = None,
+    ) -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
     @staticmethod
@@ -185,6 +186,10 @@ def supports_per_head_quant_scales(cls) -> bool:
     def supports_sliding_window(cls) -> bool:
         return False
 
+    @classmethod
+    def supports_rswa(cls) -> bool:
+        return False
+
     @classmethod
     def supports_non_causal(cls) -> bool:
         """Check if backend supports non-causal (bidirectional) attention
@@ -287,6 +292,7 @@ def validate_configuration(
         use_pcp: bool = False,
         use_adaptive_verification: bool = False,
         use_dcp: bool = False,
+        use_rswa: bool = False,
     ) -> list[str]:
         invalid_reasons = []
         if not cls.supports_head_size(head_size):
@@ -321,6 +327,8 @@ def validate_configuration(
             invalid_reasons.append(f"attention type {attn_type} not supported")
         if has_sliding_window and not cls.supports_sliding_window():
             invalid_reasons.append("sliding window not supported")
+        if use_rswa and not cls.supports_rswa():
+            invalid_reasons.append("R-SWA not supported")
         if use_non_causal and not cls.supports_non_causal():
             invalid_reasons.append("non-causal attention not supported")
         if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp():
@@ -376,8 +384,7 @@ class AttentionMetadata:
 
 @dataclass
 class CommonAttentionMetadata:
-    """
-    Per-batch attention metadata, shared across layers and backends.
+    """Per-batch attention metadata, shared across layers and backends.
     AttentionMetadataBuilder instances use it to construct per-layer metadata.
 
     For many of the tensors we keep both GPU and CPU versions.
@@ -415,9 +422,14 @@ class CommonAttentionMetadata:
     encoder_seq_lens_cpu: np.ndarray | None = None
 
     dcp_local_seq_lens: torch.Tensor | None = None
-    dcp_local_seq_lens_cpu: torch.Tensor | None = None
     """Sequence lengths of the local rank in decode context parallelism world"""
 
+    dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None
+    """(batch_size,) CPU upper bound on dcp_local_seq_lens. Under PCP+DCP it
+    holds, on every row, the largest shard any DCP rank has of the row's
+    whole request, identical on every PCP rank, so the sparse backends pad
+    their KV gather to it."""
+
     positions: torch.Tensor | None = None
     """(num_actual_tokens,) token positions.  Optional; set when the caller
     has positions available so that builders can pre-compute position-dependent
@@ -434,6 +446,11 @@ class CommonAttentionMetadata:
     decode rows (assumes every draft was accepted). Not safe for kernels
     that need exact per-row context lengths on decode rows."""
 
+    req_idx: np.ndarray | None = None
+    """(batch_size,) index of each row's request in the runner's request
+    table. Rows of one request are adjacent, so equal neighbours are PCP
+    chunks sharing one KV context."""
+
     mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
     """PrefixLM bidirectional ranges for multimodal tokens. Maps
     request index to list of (start, end) token position ranges
@@ -452,10 +469,6 @@ class CommonAttentionMetadata:
     at the current decode run's last full-state write. write_pos counts from
     here, so a preemption-resumed request re-anchors past the prompt boundary."""
 
-    # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
-    _seq_lens_cpu: torch.Tensor | None = None
-    _num_computed_tokens_cpu: torch.Tensor | None = None
-
     _num_computed_tokens_cache: torch.Tensor | None = None
     _token_to_req_indices_cache: torch.Tensor | None = None
 
@@ -469,36 +482,6 @@ def naive_query_lens(self) -> torch.Tensor:
     def replace(self, **kwargs) -> "CommonAttentionMetadata":
         return replace(self, **kwargs)
 
-    @property
-    @deprecated(
-        """
-    Prefer using device seq_lens directly to avoid implicit H<>D sync.
-    If a CPU copy is needed, use `seq_lens.cpu()` instead.
-    Will be removed in a future release, please migrate as soon as possible.
-    """
-    )
-    def seq_lens_cpu(self) -> torch.Tensor:
-        if self._seq_lens_cpu is None:
-            self._seq_lens_cpu = self.seq_lens.to("cpu")
-        return self._seq_lens_cpu
-
-    @property
-    @deprecated(
-        """
-    Prefer using device seq_lens directly to avoid implicit H<>D sync which breaks full
-    async scheduling. If a CPU copy is needed, it can be derived from 
-    query_start_loc_cpu and seq_lens.
-    Will be removed in a future release, please migrate as soon as possible.
-    """
-    )
-    def num_computed_tokens_cpu(self) -> torch.Tensor:
-        if self._num_computed_tokens_cpu is None:
-            query_seq_lens = (
-                self.query_start_loc_cpu[1:] - self.query_start_loc_cpu[:-1]
-            )
-            self._num_computed_tokens_cpu = self.seq_lens_cpu - query_seq_lens
-        return self._num_computed_tokens_cpu
-
     def compute_num_computed_tokens(self) -> torch.Tensor:
         """Compute num_computed_tokens on device (seq_lens - query_lens)."""
         if self._num_computed_tokens_cache is None:
@@ -518,19 +501,21 @@ def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor:
         # Built from the device query_start_loc: adaptive verification decides the
         # per-request draft split on device, so the CPU copy carries the right total
         # but not the right per-request boundaries. Padding requests have a query
-        # length of zero and drop out of the repeat.
+        # length of zero and are skipped by the device boundary search.
         num_mapped_tokens = int(self.query_start_loc_cpu[-1])
-        query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
-        assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens)
-        token_to_req_indices = torch.repeat_interleave(
-            torch.arange(query_lens.shape[0], dtype=torch.int32, device=buffer.device),
-            query_lens,
-            output_size=num_mapped_tokens,
+        from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel
+
+        num_output_tokens = max(num_mapped_tokens, num_tokens)
+        assert buffer.shape[0] >= num_output_tokens
+        _token_request_mapping_kernel[((num_output_tokens + 255) // 256,)](
+            self.query_start_loc,
+            buffer,
+            self.query_start_loc.shape[0] - 1,
+            num_mapped_tokens,
+            num_output_tokens,
+            num_warps=4,
         )
-        buffer[:num_mapped_tokens].copy_(token_to_req_indices)
-        if num_mapped_tokens < num_tokens:
-            buffer[num_mapped_tokens:num_tokens].zero_()
-        self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)]
+        self._token_to_req_indices_cache = buffer[:num_output_tokens]
         return self._token_to_req_indices_cache[:num_tokens]
 
     # TODO(lucas): remove once we have FULL-CG spec-decode support
@@ -542,12 +527,6 @@ def unpadded(
             query_start_loc=self.query_start_loc[: num_actual_reqs + 1],
             query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1],
             seq_lens=self.seq_lens[:num_actual_reqs],
-            _seq_lens_cpu=self._seq_lens_cpu[:num_actual_reqs]
-            if self._seq_lens_cpu is not None
-            else None,
-            _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs]
-            if self._num_computed_tokens_cpu is not None
-            else None,
             num_reqs=num_actual_reqs,
             num_actual_tokens=num_actual_tokens,
             max_query_len=self.max_query_len,
@@ -563,8 +542,12 @@ def unpadded(
             encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),
             encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),
             dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
-            dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),
+            dcp_local_seq_lens_cpu_upper_bound=maybe_slice_reqs(
+                self.dcp_local_seq_lens_cpu_upper_bound
+            ),
+            seq_lens_cpu_upper_bound=maybe_slice_reqs(self.seq_lens_cpu_upper_bound),
             is_prefilling=maybe_slice_reqs(self.is_prefilling),
+            req_idx=maybe_slice_reqs(self.req_idx),
             rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens),
             replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu),
         )
@@ -672,8 +655,7 @@ def build(
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
     ) -> M:
-        """
-        Central method that builds attention metadata.
+        """Central method that builds attention metadata.
         Some builders (MLA) require reorder_batch to be called prior to build.
 
         Args:
@@ -682,6 +664,7 @@ def build(
             fast_build: The meta-data will prioritize speed of building over
                 then speed at execution. Can be used for spec-decode where the
                 result of a build call may only be used for few layers/iters.
+
         """
         raise NotImplementedError
 
@@ -691,8 +674,7 @@ def update_block_table(
         blk_table: torch.Tensor,
         slot_mapping: torch.Tensor,
     ) -> M:
-        """
-        Update the block table for the attention metadata.
+        """Update the block table for the attention metadata.
         Faster when theres multiple kv-cache groups that create virtually the
         same metadata but just with different block tables.
 
@@ -703,8 +685,7 @@ def update_block_table(
     def build_for_cudagraph_capture(
         self, common_attn_metadata: CommonAttentionMetadata
     ) -> M:
-        """
-        Build attention metadata for CUDA graph capture. Uses build by default.
+        """Build attention metadata for CUDA graph capture. Uses build by default.
         Subclasses that override this method should call self.build or
         super().build_for_cudagraph_capture.
         """
@@ -717,8 +698,7 @@ def build_for_drafting(
         common_attn_metadata: CommonAttentionMetadata,
         draft_index: int,
     ) -> M:
-        """
-        Build attention metadata for draft model. Uses build by default.
+        """Build attention metadata for draft model. Uses build by default.
 
         Args:
             common_attn_metadata: The common attention metadata.
@@ -727,6 +707,7 @@ def build_for_drafting(
                 draft attempt for the i-th token.
                 For tree-based attention, this index instead refers to the
                 draft attempt for the i-th level in the tree of tokens.
+
         """
         return self.build(
             common_prefix_len=0,
@@ -878,6 +859,12 @@ def __new__(cls, *args, **kwargs):
     def process_weights_after_loading(self, act_dtype: torch.dtype):
         pass
 
+    def get_impl_variants(self) -> tuple["AttentionImplBase", ...]:
+        return (self,)
+
+    def prepare_for_batch(self, attn_metadata: T | None) -> None:
+        """Prepare implementation-specific state for the current batch."""
+
 
 class AttentionImpl(AttentionImplBase[T], Generic[T]):
     """Standard attention implementation with forward method."""
@@ -921,8 +908,7 @@ def forward(
         raise NotImplementedError
 
     def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool:
-        """
-        Does this attention implementation support fused output quantization.
+        """Does this attention implementation support fused output quantization.
         This is used by the AttnFusionPass to only fuse output quantization
         onto implementations that support it.
 
@@ -931,20 +917,23 @@ def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool:
 
         Returns:
             is fusion supported for this type of quantization
+
         """
         return False
 
     def fused_qk_norm_rope_kvcache_supported(self):
-        """
-        Does this attention implementation support fused QKNorm+RoPE+KVCache fusion.
+        """Does this attention implementation support fused QKNorm+RoPE+KVCache fusion.
         This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops
         with the RoPE ops and the KV cache update for implementations that support it.
         """
         return False
 
+    def fused_qk_norm_mrope_kvcache_supported(self):
+        """Whether this implementation supports fused QKNorm+MRoPE+KVCache."""
+        return False
+
     def fused_rope_kvcache_supported(self):
-        """
-        Does this attention implementation support RoPE+KVCache fusion.
+        """Does this attention implementation support RoPE+KVCache fusion.
         This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops
         with the KV cache update for implementations that support it.
         """
@@ -965,14 +954,33 @@ def do_qk_norm_rope_kvcache_update(
         kv_cache: torch.Tensor,
         layer_slot_mapping: torch.Tensor,
     ):
-        """
-        If `fused_qk_norm_rope_kvcache_supported` returns True, this method
+        """If `fused_qk_norm_rope_kvcache_supported` returns True, this method
         will be called by the fused custom op. Applies QK-norm + RoPE and
         writes K/V to the KV cache. Results are written to the pre-allocated
         q_out and k_out tensors; V is split from QKV at the graph level.
         """
         raise NotImplementedError
 
+    def do_qk_norm_mrope_kvcache_update(
+        self,
+        layer: AttentionLayer,
+        qkv: torch.Tensor,
+        q_out: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        mrope_section: tuple[int, int, int],
+        is_interleaved: bool,
+        rotary_dim: int,
+        kv_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+    ):
+        """Apply QK-norm and MRoPE, then write K/V to the cache."""
+        raise NotImplementedError
+
     def do_rope_and_kv_cache_update(
         self,
         layer: AttentionLayer,
@@ -985,8 +993,7 @@ def do_rope_and_kv_cache_update(
         kv_cache: torch.Tensor,
         layer_slot_mapping: torch.Tensor,
     ):
-        """
-        If `fused_rope_kvcache_supported` returns True, this method will be called
+        """If `fused_rope_kvcache_supported` returns True, this method will be called
         by torch.ops.vllm.fused_rope_and_unified_kv_cache_update
         to perform the inplace RoPE and KV cache update.
         """
@@ -996,6 +1003,7 @@ def do_rope_and_kv_cache_update(
 class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
     """MLA attention implementation with forward_mqa and forward_mha methods."""
 
+    hisparse_cache: "HiSparseCacheHandle | None" = None
     supports_pcp: bool = True
 
     @abstractmethod
@@ -1050,8 +1058,7 @@ def forward_mqa(
         raise NotImplementedError
 
     def fused_output_quant_supported(self, quant_key: "QuantKey"):
-        """
-        Does this attention implementation support fused output quantization.
+        """Does this attention implementation support fused output quantization.
         Since MLA quantization is done manually in forward_impl (common code),
         all MLA backends support it by default.
         """
@@ -1090,9 +1097,7 @@ def subclass_attention_backend(
     attention_backend_cls: type[AttentionBackend],
     builder_cls: type[AttentionMetadataBuilder[M]],
 ) -> type[AttentionBackend]:
-    """
-    Return a new subclass where `get_builder_cls` returns `builder_cls`.
-    """
+    """Return a new subclass where `get_builder_cls` returns `builder_cls`."""
     name: str = name_prefix + attention_backend_cls.__name__  # type: ignore
 
     return type(
diff --git a/vllm/v1/attention/backends/b12x.py b/vllm/v1/attention/backends/b12x.py
index d2b28eed0ac6..833372ab446b 100644
--- a/vllm/v1/attention/backends/b12x.py
+++ b/vllm/v1/attention/backends/b12x.py
@@ -183,7 +183,7 @@ def get_builder_cls() -> type[B12xPagedMetadataBuilder]:
         return B12xPagedMetadataBuilder
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return list(_B12X_SUPPORTED_PAGE_SIZES)
 
     @classmethod
diff --git a/vllm/v1/attention/backends/composite.py b/vllm/v1/attention/backends/composite.py
new file mode 100644
index 000000000000..0ba40e99f20a
--- /dev/null
+++ b/vllm/v1/attention/backends/composite.py
@@ -0,0 +1,515 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reusable attention backend composition and routing policies."""
+
+import functools
+import math
+from typing import ClassVar, Protocol, cast
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionImpl,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+)
+from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
+from vllm.v1.kv_cache_interface import KVCacheSpec
+
+
+def requires_mm_prefix(
+    metadata: CommonAttentionMetadata, *, unclamped_window: bool = False
+) -> bool:
+    ranges = metadata.mm_req_doc_ranges
+    # With one query per request there are no future keys to unmask. This
+    # also keeps single-token prefills on the captured FlashInfer path.
+    if not ranges or (metadata.max_query_len <= 1 and not unclamped_window):
+        return False
+    seq_lens = metadata.seq_lens_cpu_upper_bound
+    assert seq_lens is not None
+    starts = metadata.query_start_loc_cpu
+    for req_idx, spans in ranges.items():
+        query_len = int(starts[req_idx + 1] - starts[req_idx])
+        if query_len == 0 or (query_len == 1 and not unclamped_window):
+            continue
+        end = int(seq_lens[req_idx])
+        begin = end - query_len
+        # An unclamped image can also expose keys behind the sliding window.
+        if unclamped_window and any(
+            start < stop and start < end and stop >= begin for start, stop in spans
+        ):
+            return True
+        # Prefill lengths are exact. Historical image ranges must not divert
+        # later text/decode queries to the image backend.
+        if any(
+            start < stop and start < end - 1 and stop > begin for start, stop in spans
+        ):
+            return True
+    return False
+
+
+def _has_unclamped_window(layers) -> bool:
+    return any(
+        getattr(layer, "sliding_window", None) is not None
+        and not getattr(layer, "mm_prefix_clamp_sliding_window", False)
+        for layer in layers
+    )
+
+
+class CompositeAttentionRouting(Protocol):
+    """Select a child backend and declare routing constraints for composition."""
+
+    capture_variant: ClassVar[int]
+    """Child used during graph capture; replay must select the same child."""
+
+    def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None: ...
+
+    def select(self, metadata: CommonAttentionMetadata) -> int:
+        """Return 0 for the first child or 1 for the second child."""
+        ...
+
+    @staticmethod
+    def get_cudagraph_support(
+        vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
+    ) -> AttentionCGSupport:
+        """Limit graph support to batches whose routing stays invariant."""
+        ...
+
+    @staticmethod
+    def variant_uses_mm_prefix(variant: int) -> bool:
+        """Whether this child must support the model's image-mask requirement."""
+        ...
+
+
+class MMPrefixAttentionRouting(CompositeAttentionRouting):
+    """Route image masks to variant zero and causal queries to variant one."""
+
+    def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None:
+        layers = vllm_config.compilation_config.static_forward_context
+        self.unclamped_window = _has_unclamped_window(
+            layers[name] for name in layer_names
+        )
+
+    def select(self, metadata: CommonAttentionMetadata) -> int:
+        return int(
+            not requires_mm_prefix(metadata, unclamped_window=self.unclamped_window)
+        )
+
+    capture_variant: ClassVar[int] = 1
+
+    @staticmethod
+    def get_cudagraph_support(
+        vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
+    ) -> AttentionCGSupport:
+        if _has_unclamped_window(
+            vllm_config.compilation_config.static_forward_context.values()
+        ):
+            return AttentionCGSupport.NEVER
+        return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
+
+    @staticmethod
+    def variant_uses_mm_prefix(variant: int) -> bool:
+        return variant == 0
+
+
+def _intersect_kernel_block_sizes(
+    first: type[AttentionBackend],
+    second: type[AttentionBackend],
+    kv_cache_spec: KVCacheSpec | None = None,
+):
+    from vllm.v1.attention.backend import MultipleOf
+
+    result: list[int | MultipleOf] = []
+    seen: set[tuple[type, int]] = set()
+    for lhs in first.get_supported_kernel_block_sizes(kv_cache_spec):
+        for rhs in second.get_supported_kernel_block_sizes(kv_cache_spec):
+            candidate: int | MultipleOf | None = None
+            if isinstance(lhs, MultipleOf) and isinstance(rhs, MultipleOf):
+                candidate = MultipleOf(math.lcm(lhs.base, rhs.base))
+            elif isinstance(lhs, MultipleOf) and isinstance(rhs, int):
+                candidate = rhs if rhs % lhs.base == 0 else None
+            elif isinstance(lhs, int) and isinstance(rhs, MultipleOf):
+                candidate = lhs if lhs % rhs.base == 0 else None
+            elif lhs == rhs:
+                candidate = lhs
+            if candidate is None:
+                continue
+            value = candidate.base if isinstance(candidate, MultipleOf) else candidate
+            key = (type(candidate), value)
+            if key not in seen:
+                result.append(candidate)
+                seen.add(key)
+    return result
+
+
+@functools.cache
+def create_composite_attention_backend(
+    general_backend: type[AttentionBackend],
+    causal_backend: type[AttentionBackend],
+    *,
+    name: str,
+    backend_name: str,
+    module: str,
+    routing_policy: type[CompositeAttentionRouting],
+    head_sizes: tuple[int, ...] = (),
+    kernel_block_sizes: tuple[int, ...] = (),
+    device_major: int | None = None,
+) -> type[AttentionBackend]:
+    """Compose two backends with shared storage and an explicit routing policy."""
+    if general_backend.full_cls_name() == causal_backend.full_cls_name():
+        return general_backend
+    general_impl_cls = general_backend.get_impl_cls()
+    causal_impl_cls = causal_backend.get_impl_cls()
+    general_builder_cls = general_backend.get_builder_cls()
+    causal_builder_cls = causal_backend.get_builder_cls()
+    if (
+        general_backend.forward_includes_kv_cache_update
+        or causal_backend.forward_includes_kv_cache_update
+    ):
+        raise ValueError("Composite backends require separate KV cache updates")
+
+    class CompositeAttentionImpl(AttentionImpl):
+        def __init__(self, *args, **kwargs) -> None:
+            self.general_impl = cast(AttentionImpl, general_impl_cls(*args, **kwargs))
+            try:
+                self.causal_impl = cast(AttentionImpl, causal_impl_cls(*args, **kwargs))
+            except TypeError as e:
+                raise ValueError(
+                    f"Causal attention backend {causal_backend.get_name()} "
+                    f"does not accept the attention arguments required by "
+                    f"the general backend {general_backend.get_name()}: {e}"
+                ) from e
+            for name in (
+                "num_heads",
+                "num_kv_heads",
+                "head_size",
+                "scale",
+                "kv_cache_dtype",
+            ):
+                if hasattr(self.general_impl, name):
+                    setattr(self, name, getattr(self.general_impl, name))
+            self.supports_quant_query_input = (
+                self.general_impl.supports_quant_query_input
+                and self.causal_impl.supports_quant_query_input
+            )
+            self.supports_pcp = self.general_impl.supports_pcp
+            self.supports_dcp = self.causal_impl.supports_dcp
+            self.can_return_lse_for_decode = self.causal_impl.can_return_lse_for_decode
+            self.lse_base_on_e = self.causal_impl.lse_base_on_e
+            self.need_to_return_lse_for_decode = (
+                self.causal_impl.need_to_return_lse_for_decode
+            )
+            self.supports_mtp_with_cp_non_trivial_interleave_size = (
+                self.causal_impl.supports_mtp_with_cp_non_trivial_interleave_size
+            )
+
+        def get_impl_for_metadata(self, attn_metadata):
+            return (
+                self.causal_impl
+                if getattr(attn_metadata, "_attention_backend_variant", 0)
+                else self.general_impl
+            )
+
+        def get_impl_variants(self):
+            return self.general_impl, self.causal_impl
+
+        def forward(
+            self,
+            layer,
+            query,
+            key,
+            value,
+            kv_cache,
+            attn_metadata,
+            output,
+            output_scale=None,
+            output_block_scale=None,
+        ):
+            impl = self.get_impl_for_metadata(attn_metadata)
+            return impl.forward(
+                layer,
+                query,
+                key,
+                value,
+                kv_cache,
+                attn_metadata,
+                output,
+                output_scale,
+                output_block_scale,
+            )
+
+        def process_weights_after_loading(self, act_dtype: torch.dtype):
+            for impl in self.get_impl_variants():
+                impl.process_weights_after_loading(act_dtype)
+
+        def fused_output_quant_supported(self, quant_key):
+            return all(
+                impl.fused_output_quant_supported(quant_key)
+                for impl in self.get_impl_variants()
+            )
+
+        def fused_qk_norm_rope_kvcache_supported(self):
+            return self.general_impl.fused_qk_norm_rope_kvcache_supported()
+
+        def do_qk_norm_rope_kvcache_update(self, *args, **kwargs):
+            method = self.general_impl.do_qk_norm_rope_kvcache_update
+            return method(*args, **kwargs)
+
+        def fused_rope_kvcache_supported(self):
+            return self.general_impl.fused_rope_kvcache_supported()
+
+        def do_rope_and_kv_cache_update(self, *args, **kwargs):
+            method = self.general_impl.do_rope_and_kv_cache_update
+            return method(*args, **kwargs)
+
+        def do_kv_cache_update(self, *args, **kwargs):
+            method = self.general_impl.do_kv_cache_update  # type: ignore[attr-defined]
+            return method(*args, **kwargs)
+
+    class CompositeAttentionMetadataBuilder(AttentionMetadataBuilder):
+        requires_block_table_width = (
+            general_builder_cls.requires_block_table_width
+            or causal_builder_cls.requires_block_table_width
+        )
+
+        def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs):
+            super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+            args = (kv_cache_spec, layer_names, vllm_config, device)
+            # Only the variant that asked for the extra ctor kwargs receives them.
+            self.general_builder = general_builder_cls(
+                *args,
+                **(kwargs if general_builder_cls.requires_block_table_width else {}),
+            )
+            self.causal_builder = causal_builder_cls(
+                *args,
+                **(kwargs if causal_builder_cls.requires_block_table_width else {}),
+            )
+            self.reorder_batch_threshold = min(
+                (
+                    builder.reorder_batch_threshold
+                    for builder in self._builders
+                    if builder.reorder_batch_threshold is not None
+                ),
+                default=None,
+            )
+            self.routing = routing_policy(layer_names, vllm_config)
+
+        def set_kernel_block_size(self, kernel_block_size):
+            super().set_kernel_block_size(kernel_block_size)
+            for builder in self._builders:
+                builder.set_kernel_block_size(kernel_block_size)
+
+        @classmethod
+        def get_cudagraph_support(cls, vllm_config, kv_cache_spec):
+            return min(
+                general_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec),
+                causal_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec),
+                routing_policy.get_cudagraph_support(vllm_config, kv_cache_spec),
+                key=lambda support: support.value,
+            )
+
+        def _builder(self, common_attn_metadata):
+            backend_variant = self.routing.select(common_attn_metadata)
+            builder = self.causal_builder if backend_variant else self.general_builder
+            return builder, backend_variant
+
+        def _build(self, method, common_attn_metadata, **kwargs):
+            builder, backend_variant = self._builder(common_attn_metadata)
+            metadata = getattr(builder, method)(
+                common_attn_metadata=common_attn_metadata, **kwargs
+            )
+            metadata._attention_backend_variant = backend_variant
+            return metadata
+
+        def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
+            return self._build(
+                "build",
+                common_attn_metadata,
+                common_prefix_len=common_prefix_len,
+                fast_build=fast_build,
+            )
+
+        def build_for_cudagraph_capture(self, common_attn_metadata):
+            variant = routing_policy.capture_variant
+            metadata = self._builders[variant].build_for_cudagraph_capture(
+                common_attn_metadata
+            )
+            metadata._attention_backend_variant = variant
+            return metadata
+
+        @property
+        def _builders(self):
+            return self.general_builder, self.causal_builder
+
+        # Forward the optional workspace protocol so that a wrapped builder
+        # which allocates one (FlashInfer) still joins the runner's cross-group
+        # sharing instead of allocating a second buffer behind the composite.
+        def _get_workspace_buffer(self):
+            for builder in self._builders:
+                if hasattr(builder, "_get_workspace_buffer"):
+                    workspace = builder._get_workspace_buffer()
+                    # The runner only calls the getter for the group that
+                    # provides the buffer, so hand it to the sibling here.
+                    self.set_workspace_buffer(workspace)
+                    return workspace
+            return None
+
+        def set_workspace_buffer(self, workspace_buffer):
+            for builder in self._builders:
+                if hasattr(builder, "set_workspace_buffer"):
+                    builder.set_workspace_buffer(workspace_buffer)
+
+    class CompositeAttentionBackend(AttentionBackend):
+        forward_includes_kv_cache_update = False
+        general_backend_cls = general_backend
+        causal_backend_cls = causal_backend
+
+        @staticmethod
+        def get_name():
+            return backend_name
+
+        @staticmethod
+        def get_impl_cls():
+            return CompositeAttentionImpl
+
+        @staticmethod
+        def get_builder_cls():
+            return CompositeAttentionMetadataBuilder
+
+        @staticmethod
+        def get_supported_kernel_block_sizes(
+            kv_cache_spec: KVCacheSpec | None = None,
+        ):
+            if kernel_block_sizes:
+                return [
+                    size
+                    for size in kernel_block_sizes
+                    if all(
+                        backend.supports_block_size(size)
+                        for backend in (general_backend, causal_backend)
+                    )
+                ]
+            return _intersect_kernel_block_sizes(
+                general_backend, causal_backend, kv_cache_spec
+            )
+
+        @classmethod
+        def supports_block_size(cls, block_size):
+            return bool(
+                cls.get_supported_kernel_block_sizes()
+            ) and super().supports_block_size(block_size)
+
+        @classmethod
+        def get_supported_head_sizes(cls):
+            return list(head_sizes)
+
+        @classmethod
+        def supports_head_size(cls, head_size):
+            return (not head_sizes or head_size in head_sizes) and all(
+                backend.supports_head_size(head_size)
+                for backend in (general_backend, causal_backend)
+            )
+
+        @classmethod
+        def supports_compute_capability(cls, capability):
+            return (device_major is None or capability.major == device_major) and all(
+                backend.supports_compute_capability(capability)
+                for backend in (general_backend, causal_backend)
+            )
+
+        @classmethod
+        def supports_dtype(cls, dtype):
+            return all(
+                backend.supports_dtype(dtype)
+                for backend in (general_backend, causal_backend)
+            )
+
+        @classmethod
+        def supports_kv_cache_dtype(cls, dtype):
+            return all(
+                backend.supports_kv_cache_dtype(dtype)
+                for backend in (general_backend, causal_backend)
+            )
+
+        @classmethod
+        def supports_mm_prefix(cls):
+            return all(
+                not routing_policy.variant_uses_mm_prefix(variant)
+                or backend.supports_mm_prefix()
+                for variant, backend in enumerate((general_backend, causal_backend))
+            )
+
+        @classmethod
+        def supports_sliding_window(cls):
+            return all(
+                backend.supports_sliding_window()
+                for backend in (general_backend, causal_backend)
+            )
+
+        @classmethod
+        def supports_rswa(cls):
+            return all(
+                backend.supports_rswa() for backend in (general_backend, causal_backend)
+            )
+
+        @classmethod
+        def supports_pcp(cls):
+            return False
+
+        @classmethod
+        def supports_dcp(cls):
+            return False
+
+        @classmethod
+        def supports_device_cpu_query_lens_mismatch(cls):
+            return False
+
+        @classmethod
+        def supports_combination(
+            cls,
+            head_size,
+            dtype,
+            kv_cache_dtype,
+            block_size,
+            use_mla,
+            has_sink,
+            use_sparse,
+            use_mm_prefix,
+            device_capability,
+        ):
+            for variant, backend in enumerate((general_backend, causal_backend)):
+                reason = backend.supports_combination(
+                    head_size,
+                    dtype,
+                    kv_cache_dtype,
+                    block_size,
+                    use_mla,
+                    has_sink,
+                    use_sparse,
+                    use_mm_prefix and routing_policy.variant_uses_mm_prefix(variant),
+                    device_capability,
+                )
+                if reason:
+                    return f"{backend.get_name()}: {reason}"
+            return None
+
+        @classmethod
+        def customize_spec(cls, spec):
+            general_spec = general_backend.customize_spec(spec)
+            if general_spec != causal_backend.customize_spec(spec):
+                raise ValueError("Composite backends require the same KV cache spec")
+            return general_spec
+
+        @classmethod
+        def supported_kv_cache_layouts(cls):
+            return tuple(
+                get_supported_kv_cache_layouts((general_backend, causal_backend))
+            )
+
+    CompositeAttentionBackend.__name__ = name
+    CompositeAttentionBackend.__qualname__ = name
+    CompositeAttentionBackend.__module__ = module
+    return CompositeAttentionBackend
diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py
index dd3f34c9167e..46b01cc5519a 100644
--- a/vllm/v1/attention/backends/cpu_attn.py
+++ b/vllm/v1/attention/backends/cpu_attn.py
@@ -13,6 +13,7 @@
 from vllm import envs
 from vllm.config import (
     VllmConfig,
+    get_current_vllm_config,
     get_layers_from_vllm_config,
 )
 from vllm.logger import init_logger
@@ -31,6 +32,10 @@
 from vllm.v1.attention.backends.utils import (
     get_num_attention_heads_from_layers,
 )
+from vllm.v1.attention.backends.zentorch_sdpa import (
+    should_use_zentorch_sdpa,
+    zentorch_sdpa_attn,
+)
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     CrossAttentionSpec,
@@ -57,7 +62,7 @@ class CPUAttentionBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(32)]
 
     @classmethod
@@ -334,6 +339,12 @@ def __init__(
                 "heads in the layer"
             )
 
+        vllm_config = get_current_vllm_config()
+        self.use_zentorch_sdpa = should_use_zentorch_sdpa(
+            attn_type,
+            vllm_config.model_config.dtype,
+        )
+
     def forward(
         self,
         layer: AttentionLayer,
@@ -349,14 +360,22 @@ def forward(
         """Forward pass for CPU attention backend.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
                 [num_blocks, num_kv_heads, block_size, 2 * head_size]
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization; not supported
+                by this backend.
+            output_block_scale: Block scale for fused output quantization;
+                not supported by this backend.
+
         Returns:
             shape = [num_tokens, num_heads * head_size]
+
         """
         if output_scale is not None or output_block_scale is not None:
             raise NotImplementedError(
@@ -375,6 +394,21 @@ def forward(
             AttentionType.ENCODER,
         )
         if is_encoder_attention:
+            if self.use_zentorch_sdpa:
+                # Encoder attention never reads the KV cache back, so the
+                # zentorch path attends the packed QKV directly instead of
+                # staging it through the scratch encoder cache.
+                zentorch_sdpa_attn(
+                    query[:num_actual_tokens],
+                    key[:num_actual_tokens],
+                    value[:num_actual_tokens],
+                    output[:num_actual_tokens],
+                    attn_metadata,
+                    self.scale,
+                    self.sliding_window,
+                    self.alibi_slopes,
+                )
+                return output
             # For encoder attention,
             kv_cache = attn_metadata.encoder_cache
 
@@ -398,6 +432,15 @@ def forward(
                 kv_cache_dtype=self.kv_cache_dtype,
             )
 
+        # The CPU kernel executes attention sinks natively in bf16. If the
+        # sinks tensor is anything other than bf16, cast it to fp32 so it is
+        # executed in full float precision (done lazily here, after weights
+        # are loaded, rather than at __init__ time).
+        if self.sinks is not None and self.sinks.dtype not in [
+            torch.bfloat16,
+            torch.float32,
+        ]:
+            self.sinks = self.sinks.to(torch.float32)
         ops.cpu_attention_with_kv_cache(
             query=query[:num_actual_tokens],
             key_cache=key_cache,
@@ -504,11 +547,10 @@ def _get_attn_isa(
         )
     if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
         return "amx"
+    elif supports_arm:
+        return "neon"
     elif block_size % 32 == 0:
-        if supports_arm:
-            # support ARM NEON FMLA and BFMMLA (bf16) for block size 32
-            return "neon"
-        elif supports_riscv and _riscv_supports_rvv():
+        if supports_riscv and _riscv_supports_rvv():
             return "rvv"
         elif supports_vxe:
             return "vxe"
diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py
index 0de6b85217b1..83348bb94a42 100644
--- a/vllm/v1/attention/backends/fa_utils.py
+++ b/vllm/v1/attention/backends/fa_utils.py
@@ -303,7 +303,11 @@ def flash_attn_supports_kv_cache_dtype(
         kv_cache_block_size=kv_cache_block_size,
         supports_fa4_hd256=supports_fa4_hd256,
     )
-    return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (
+    is_sm90 = current_platform.is_device_capability_family(90)
+    sm90_fp8_kv_supported = fa_version == 3 or (
+        fa_version == 4 and head_size in (None, 512)
+    )
+    return (sm90_fp8_kv_supported and is_sm90) or (
         fa_version == 4 and current_platform.is_device_capability_family(100)
     )
 
@@ -357,6 +361,7 @@ def is_flash_attn_varlen_func_available() -> bool:
 
     Returns:
         bool: True if a working flash_attn_varlen_func implementation is available.
+
     """
     if current_platform.is_cuda() or current_platform.is_xpu():
         # CUDA and XPU always have flash_attn_varlen_func available
diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py
index 9ff543024fde..129b3e39c31a 100755
--- a/vllm/v1/attention/backends/flash_attn.py
+++ b/vllm/v1/attention/backends/flash_attn.py
@@ -5,12 +5,13 @@
 import copy
 import functools
 from dataclasses import dataclass
-from typing import ClassVar
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
 
 from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, WarmupChoices
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import (
     PIN_MEMORY,
@@ -46,6 +47,7 @@
 
 if is_flash_attn_varlen_func_available():
     from vllm.v1.attention.backends.fa_utils import (
+        compile_flash_attn_varlen_func_from_specs,
         flash_attn_supports_sinks,
         flash_attn_varlen_func,
         get_scheduler_metadata,
@@ -67,7 +69,13 @@
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
 )
-from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    FullAttentionSpec,
+    KVCacheSpec,
+    KVQuantMode,
+    SlidingWindowSpec,
+)
 from vllm.v1.worker.cp_utils import (
     run_split_fa2_dcp_context_attention,
     should_skip_dcp_context_attention,
@@ -77,6 +85,200 @@
 
 logger = init_logger(__name__)
 
+FA4_DENSE_FLOAT_DTYPES = (torch.bfloat16, torch.float16)
+FA4_DENSE_Q_TILE = 128
+FA4_DENSE_NUM_BLOCKS = 256
+FA4_DENSE_MAX_SEQLEN_K = 8192
+
+
+class FA4DenseAttentionKernel(VllmJitKernel["FA4DenseAttentionKernel.CompileKey"]):
+    """Warm paged, causal FA4 on Blackwell and Hopper."""
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        q_stage: int
+        dtype: torch.dtype
+        qhead_per_kvhead: int
+        head_dim: int
+        page_size: int
+        has_window_left: bool
+        has_window_right: bool
+        softcap: float
+        is_single_batch: bool
+        num_splits: int
+
+    @staticmethod
+    def kernel(*args: Any, **kwargs: Any) -> Any:
+        assert flash_attn_varlen_func is not None
+        return flash_attn_varlen_func(*args, **kwargs)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        q_stage: int,
+        dtype: torch.dtype,
+        num_qo_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        page_size: int,
+        window_size: tuple[int, int],
+        softcap: float,
+        is_single_batch: bool = True,
+        num_splits: int = 1,
+    ) -> CompileKey:
+        return self.CompileKey(
+            q_stage=q_stage,
+            dtype=dtype,
+            qhead_per_kvhead=num_qo_heads // num_kv_heads,
+            head_dim=head_dim,
+            page_size=page_size,
+            has_window_left=window_size[0] >= 0,
+            has_window_right=window_size[1] >= 0,
+            softcap=softcap,
+            is_single_batch=is_single_batch,
+            num_splits=num_splits,
+        )
+
+    def get_warmup_keys(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        kv_cache_spec: AttentionSpec,
+        num_qo_heads: int,
+    ) -> list[CompileKey]:
+        capability = current_platform.get_device_capability()
+        major = capability.major if capability is not None else None
+        if (
+            vllm_config.parallel_config.decode_context_parallel_size != 1
+            or vllm_config.model_config.rswa_window is not None
+            or kv_cache_spec.kv_quant_mode != KVQuantMode.NONE
+            or kv_cache_spec.dtype not in FA4_DENSE_FLOAT_DTYPES
+            or not isinstance(kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec))
+        ):
+            return []
+        head_dim = kv_cache_spec.head_size
+        page_size = kv_cache_spec.block_size
+        sliding_window = kv_cache_spec.sliding_window
+        window_size = (
+            (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
+        )
+        softcap = float(
+            getattr(
+                vllm_config.model_config.hf_text_config, "attn_logit_softcapping", 0
+            )
+            or 0
+        )
+        if (
+            major in (10, 11)
+            and uses_fa4_hd256_kernel(head_dim)
+            and page_size == FA4_HD256_PAGE_SIZE
+        ):
+            q_stages = WarmupChoices(1, 2)
+            single_batches = WarmupChoices(True, False)
+            split_counts = WarmupChoices(1)
+        elif (
+            major == 9 and head_dim == 512 and window_size == (-1, -1) and softcap == 0
+        ):
+            # SM90 forward has split/non-split variants; its transitive
+            # combine specializes at 32/64/128/256 splits.
+            q_stages = WarmupChoices(1)
+            single_batches = WarmupChoices(False)
+            split_counts = WarmupChoices(1, 32, 64, 128, 256)
+        else:
+            return []
+
+        return self._trace_dispatch(self.dispatch)(
+            q_stage=q_stages,
+            is_single_batch=single_batches,
+            num_splits=split_counts,
+            dtype=vllm_config.model_config.dtype,
+            num_qo_heads=num_qo_heads,
+            num_kv_heads=kv_cache_spec.num_kv_heads,
+            head_dim=head_dim,
+            page_size=page_size,
+            window_size=WarmupChoices(window_size),
+            softcap=softcap,
+        )
+
+    def compile(self, compile_key: CompileKey) -> None:
+        assert compile_flash_attn_varlen_func_from_specs is not None
+        max_seqlen_q = FA4_DENSE_Q_TILE + 1 if compile_key.q_stage == 2 else 1
+        batch_size = 1 if compile_key.is_single_batch else 2
+        kv_shape = (
+            FA4_DENSE_NUM_BLOCKS,
+            compile_key.page_size,
+            1,
+            compile_key.head_dim,
+        )
+        kv_stride = (
+            2 * compile_key.page_size * compile_key.head_dim,
+            2 * compile_key.head_dim,
+            2 * compile_key.head_dim,
+            1,
+        )
+        compile_flash_attn_varlen_func_from_specs(
+            q_shape=(
+                batch_size * max_seqlen_q,
+                compile_key.qhead_per_kvhead,
+                compile_key.head_dim,
+            ),
+            k_shape=kv_shape,
+            v_shape=kv_shape,
+            k_stride=kv_stride,
+            v_stride=kv_stride,
+            q_dtype=compile_key.dtype,
+            cu_seqlens_q_shape=(batch_size + 1,),
+            seqused_k_shape=(batch_size,),
+            page_table_shape=(
+                batch_size,
+                FA4_DENSE_MAX_SEQLEN_K // compile_key.page_size,
+            ),
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=FA4_DENSE_MAX_SEQLEN_K,
+            causal=True,
+            window_size=[
+                1 if compile_key.has_window_left else -1,
+                1 if compile_key.has_window_right else -1,
+            ],
+            softcap=compile_key.softcap,
+            num_splits=2 if compile_key.num_splits > 1 else 1,
+            fa_version=4,
+            return_softmax_lse=False,
+        )
+        if compile_key.head_dim == 512 and compile_key.num_splits > 1:
+            from torch._subclasses.fake_tensor import FakeTensorMode
+
+            from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd_combine
+
+            # FA4's forward compile-only entry returns before its transitive
+            # SplitKV combine. Call the native combine entry in FakeTensorMode:
+            # it populates the native cache without launching the kernel.
+            with FakeTensorMode():
+                _flash_attn_fwd_combine(
+                    torch.empty(
+                        (compile_key.num_splits, 2, 1, 512),
+                        dtype=torch.float32,
+                        device="cuda",
+                    ),
+                    torch.empty(
+                        (compile_key.num_splits, 2, 1),
+                        dtype=torch.float32,
+                        device="cuda",
+                    ),
+                    torch.empty((2, 1, 512), dtype=compile_key.dtype, device="cuda"),
+                    cu_seqlens=torch.empty((3,), dtype=torch.int32, device="cuda"),
+                )
+
+    def __call__(
+        self,
+        *,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        **kwargs: Any,
+    ) -> Any:
+        return self.kernel(q=q, k=k, v=v, **kwargs)
+
 
 class FlashAttentionBackend(AttentionBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
@@ -89,18 +291,56 @@ class FlashAttentionBackend(AttentionBackend):
     ]
     head_size_v: int | None = None
 
+    @staticmethod
+    def _get_sm90_fa4_fp8_kv_block_size(
+        kv_cache_spec: KVCacheSpec | None = None,
+    ) -> int | None:
+        if kv_cache_spec is None:
+            vllm_config = get_current_vllm_config_or_none()
+            if vllm_config is None or vllm_config.model_config is None:
+                return None
+            head_size = vllm_config.model_config.get_head_size()
+            uses_fp8_kv = vllm_config.cache_config.cache_dtype in (
+                "fp8",
+                "fp8_e4m3",
+            )
+        else:
+            if not isinstance(kv_cache_spec, AttentionSpec):
+                return None
+            head_size = kv_cache_spec.head_size
+            uses_fp8_kv = kv_cache_spec.kv_quant_mode == KVQuantMode.FP8_PER_TENSOR
+
+        if (
+            current_platform.is_device_capability_family(90)
+            and uses_fp8_kv
+            and head_size == 512
+            and get_flash_attn_version(head_size=head_size) == 4
+        ):
+            # The SM90 FP8-KV-dequant kernel uses a 64-token TMA tile/page.
+            return 64
+        return None
+
     @classmethod
-    def _get_fa4_hd256_block_size(cls) -> int | None:
-        vllm_config = get_current_vllm_config_or_none()
-        if vllm_config is None or vllm_config.model_config is None:
-            return None
+    def _get_fa4_hd256_block_size(
+        cls, kv_cache_spec: KVCacheSpec | None = None
+    ) -> int | None:
+        if kv_cache_spec is None:
+            vllm_config = get_current_vllm_config_or_none()
+            if vllm_config is None or vllm_config.model_config is None:
+                return None
+            head_size = vllm_config.model_config.get_head_size()
+            head_size_v = cls.head_size_v
+        else:
+            if not isinstance(kv_cache_spec, AttentionSpec):
+                return None
+            head_size = kv_cache_spec.head_size
+            head_size_v = kv_cache_spec.head_size_v
 
-        head_size = vllm_config.model_config.get_head_size()
         if (
-            uses_fa4_hd256_kernel(head_size, cls.head_size_v)
+            uses_fa4_hd256_kernel(head_size, head_size_v)
             and get_flash_attn_version(
                 head_size=head_size,
-                head_size_v=cls.head_size_v,
+                head_size_v=head_size_v,
                 supports_fa4_hd256=True,
             )
             == 4
@@ -109,8 +349,12 @@ def _get_fa4_hd256_block_size(cls) -> int | None:
         return None
 
     @classmethod
-    def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]:
-        if block_size := cls._get_fa4_hd256_block_size():
+    def get_supported_kernel_block_sizes(
+        cls, kv_cache_spec: KVCacheSpec | None = None
+    ) -> list[int | MultipleOf]:
+        if block_size := cls._get_sm90_fa4_fp8_kv_block_size(kv_cache_spec):
+            return [block_size]
+        if block_size := cls._get_fa4_hd256_block_size(kv_cache_spec):
             # Sliding-window specs select the smallest advertised size.
             return [block_size]
         return [MultipleOf(16)]
@@ -119,6 +363,8 @@ def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]:
 
     @classmethod
     def get_preferred_block_size(cls, default_block_size: int) -> int:
+        if block_size := cls._get_sm90_fa4_fp8_kv_block_size():
+            return max(default_block_size, block_size)
         if block_size := cls._get_fa4_hd256_block_size():
             return max(default_block_size, block_size)
         if current_platform.is_xpu():
@@ -129,6 +375,10 @@ def get_preferred_block_size(cls, default_block_size: int) -> int:
     def get_name() -> str:
         return "FLASH_ATTN"
 
+    @classmethod
+    def supports_rswa(cls) -> bool:
+        return True
+
     @classmethod
     def supports_sliding_window(cls) -> bool:
         return True
@@ -213,6 +463,13 @@ def supports_combination(
     ) -> str | None:
         if has_sink and device_capability < DeviceCapability(9, 0):
             return "sink not supported on compute capability < 9.0"
+        if (
+            use_mm_prefix
+            and kv_cache_dtype is not None
+            and is_quantized_kv_cache(kv_cache_dtype)
+            and device_capability == DeviceCapability(9, 0)
+        ):
+            return "SM90 FP8 KV with mm_prefix requires Triton"
         if (
             kv_cache_dtype is not None
             and is_quantized_kv_cache(kv_cache_dtype)
@@ -225,7 +482,10 @@ def supports_combination(
                 supports_fa4_hd256=True,
             )
         ):
-            return "FP8 KV cache requires FA3 on SM90 or FA4 on SM100"
+            return (
+                "FP8 KV cache requires FA3 on SM90, FA4 with head_size=512 "
+                "on SM90, or FA4 on SM100"
+            )
         if (
             use_mm_prefix
             and get_flash_attn_version(
@@ -308,15 +568,16 @@ def _get_sliding_window_configs(
 ) -> set[tuple[int, int] | None]:
     """Get the set of all sliding window configs used in the model.
 
-    Only inspects FlashAttentionImpl layers. Other backends (e.g.
-    TurboQuant, MLA) use their own metadata builders and are skipped.
+    Only inspects FlashAttentionImpl layers, including implementations wrapped
+    by a composite backend. Other backends (e.g. TurboQuant, MLA) use their own
+    metadata builders and are skipped.
     """
     sliding_window_configs: set[tuple[int, int] | None] = set()
     layers = get_layers_from_vllm_config(vllm_config, Attention)
     for layer in layers.values():
-        if not isinstance(layer.impl, FlashAttentionImpl):
-            continue
-        sliding_window_configs.add(layer.impl.sliding_window)
+        for impl in layer.impl.get_impl_variants():
+            if isinstance(impl, FlashAttentionImpl):
+                sliding_window_configs.add(impl.sliding_window)
     return sliding_window_configs
 
 
@@ -443,15 +704,27 @@ def __init__(
 
         self.max_num_splits = 0  # No upper bound on the number of splits.
         self.aot_schedule = get_flash_attn_version() == 3
+        head_size_v = getattr(kv_cache_spec, "head_size_v", None)
+        fa_version = get_flash_attn_version(
+            head_size=self.headdim,
+            head_size_v=head_size_v,
+            kv_cache_block_size=self.block_size,
+            supports_fa4_hd256=True,
+        )
+        self.fa4_hd256 = fa_version == 4 and uses_fa4_hd256_kernel(
+            self.headdim, head_size_v
+        )
 
-        self.fa4_hd256 = uses_fa4_hd256_kernel(self.headdim) and (
-            get_flash_attn_version(
-                head_size=self.headdim,
-                kv_cache_block_size=self.block_size,
-                supports_fa4_hd256=True,
+        if (
+            vllm_config.kernel_config.enable_jit_warmup
+            and self.model_config.hf_config.model_type in ("gemma4", "gemma4_unified")
+            and fa_version == 4
+        ):
+            _FA4_DENSE_ATTENTION_KERNEL.register_warmup(
+                vllm_config=vllm_config,
+                kv_cache_spec=kv_cache_spec,
+                num_qo_heads=self.num_heads_q,
             )
-            == 4
-        )
 
         try:
             from vllm.distributed.parallel_state import get_dcp_group
@@ -548,8 +821,7 @@ def build(
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
     ) -> FlashAttentionMetadata:
-        """
-        fast_build disables AOT scheduling, used when there will be few
+        """fast_build disables AOT scheduling, used when there will be few
         iterations i.e. spec-decode
         """
         num_reqs = common_attn_metadata.num_reqs
@@ -928,7 +1200,17 @@ def __init__(
                 "heads in the layer"
             )
 
-        self.supports_quant_query_input = flash_attn_supports_quant_query_input()
+        # FA4's SM90 FP8-KV path consumes native FP16/BF16 Q and dequantizes
+        # FP8 K/V in-kernel. Other FA4 paths (notably SM100) still require Q,
+        # K, and V to have the same FP8 dtype.
+        uses_sm90_fa4_fp8_kv_dequant = (
+            self.vllm_flash_attn_version == 4
+            and current_platform.is_device_capability_family(90)
+            and self.kv_cache_dtype in ("fp8", "fp8_e4m3")
+        )
+        self.supports_quant_query_input = flash_attn_supports_quant_query_input() and (
+            not uses_sm90_fa4_fp8_kv_dequant
+        )
 
         dcp_a2a = (
             vllm_config is not None
@@ -960,17 +1242,25 @@ def forward(
         """Forward pass with FlashAttention.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
                 [num_blocks, num_kv_heads, block_size, 2 * head_size]
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization; not supported
+                by this backend.
+            output_block_scale: Block scale for fused output quantization;
+                not supported by this backend.
+
         Returns:
             shape = [num_tokens, num_heads * head_size]
         NOTE: FP8 quantization, flash-attn expect the size of
               {q,k,v}_descale to be (num_sequences, num_kv_heads).
               We use torch's .expand() to avoid duplicating values
+
         """
         assert self.vllm_flash_attn_version is not None, (
             "FlashAttention version not detected."
@@ -1160,7 +1450,7 @@ def forward(
                     block_table = block_table[:, :num_pages]
                     num_splits = 1
 
-                flash_attn_varlen_func(
+                _FA4_DENSE_ATTENTION_KERNEL(
                     q=query[:num_actual_tokens],
                     k=key_cache,
                     v=value_cache,
@@ -1447,6 +1737,7 @@ def _forward_encoder_attention(
             output: shape = [num_encoder_tokens, num_heads, head_size]
             attn_metadata: Encoder attention metadata
             layer: The attention layer
+
         """
         assert self.vllm_flash_attn_version is not None, (
             "FlashAttention version not detected."
@@ -1833,3 +2124,6 @@ def cascade_attention(
 
     # Merge prefix and suffix outputs, and store the result in output.
     merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
+
+
+_FA4_DENSE_ATTENTION_KERNEL = FA4DenseAttentionKernel()
diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py
index 40846758d833..4bb58ef4ea43 100644
--- a/vllm/v1/attention/backends/flash_attn_diffkv.py
+++ b/vllm/v1/attention/backends/flash_attn_diffkv.py
@@ -131,17 +131,25 @@ def forward(
         """Forward pass with FlashAttention.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size_v]
             kv_cache: shape =
                 [num_blocks, block_size, num_kv_heads, head_size + head_size_v]
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization; not supported
+                by this backend.
+            output_block_scale: Block scale for fused output quantization;
+                not supported by this backend.
+
         Returns:
             shape = [num_tokens, num_heads * head_size_v]
         NOTE: FP8 quantization, flash-attn expect the size of
               {q,k,v}_descale to be (num_sequences, num_kv_heads).
               We use torch's .expand() to avoid duplicating values
+
         """
         assert self.vllm_flash_attn_version is not None, (
             "FlashAttention version not detected."
diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py
index 4b79135caba6..08920fa3ca35 100755
--- a/vllm/v1/attention/backends/flashinfer.py
+++ b/vllm/v1/attention/backends/flashinfer.py
@@ -423,7 +423,7 @@ def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # Page sizes >= 128 only run on the trtllm-gen dynamic kernel (GQA/MQA
         # on Blackwell); advertise them only when usable so selection never
         # picks a large kernel block we cannot serve.
@@ -519,7 +519,7 @@ def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
 
     @classmethod
     def supports_sink(cls) -> bool:
-        """FlashInfer supports sinks on SM12x XQA and SM100 trtllm-gen."""
+        """Whether FlashInfer can serve attention sinks on this platform."""
         from vllm.utils.flashinfer import (
             force_use_trtllm_attention,
         )
@@ -527,7 +527,9 @@ def supports_sink(cls) -> bool:
         if force_use_trtllm_attention() is False:
             return False
 
-        if current_platform.is_device_capability_family(120):
+        if current_platform.is_device_capability(
+            90
+        ) or current_platform.is_device_capability_family(120):
             return supports_trtllm_attention(is_prefill=False)
 
         if not current_platform.is_device_capability_family(100):
@@ -570,6 +572,10 @@ class FlashInferDecodeKernel(Enum):
     TRTLLM_GEN = "trtllm-gen"
 
 
+def _is_xqa_head_dim_supported(head_dim: int) -> bool:
+    return 16 <= head_dim <= 256 and head_dim % 16 == 0
+
+
 @dataclass
 class TRTLLMPrefill:
     """Metadata for the TRTLLM prefill pathway."""
@@ -827,7 +833,7 @@ def __init__(
         # the wider groups fall back to native FlashInfer decode.
         if (
             self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
-            and not (16 <= self.head_dim <= 256 and self.head_dim % 16 == 0)
+            and not _is_xqa_head_dim_supported(self.head_dim)
         ):
             logger.warning_once(
                 "FlashInfer XQA decode does not support head_dim=%d; "
@@ -847,9 +853,8 @@ def __init__(
             )
             self.use_trtllm_decode_attention = False
             self.flashinfer_trtllm_api_decode_kernel = None
-        self.use_dedicated_xqa = (
-            current_platform.is_device_capability_family(120)
-            and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
+        self.use_xqa = (
+            self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
         )
         # Adaptive verification trims drafts on device, so decode query lengths
         # must come from the device qo_indptr; only trtllm-gen supports that
@@ -861,14 +866,11 @@ def __init__(
             == FlashInferDecodeKernel.TRTLLM_GEN
             and not self.use_dcp
         )
-        supports_spec_as_decode = (
-            self.flashinfer_trtllm_api_decode_kernel
-            == FlashInferDecodeKernel.TRTLLM_GEN
-            or self.use_dedicated_xqa
-        )
         self._init_reorder_batch_threshold(
             1,
-            supports_spec_as_decode=supports_spec_as_decode,
+            supports_spec_as_decode=(
+                self.flashinfer_trtllm_api_decode_kernel is not None
+            ),
             # trtllm-gen decode receives no cp_rank/global-seq-len information,
             # so its end-aligned causal mask is wrong for q_len > 1 over the
             # DCP-interleaved local KV shard (spec token i misses up to
@@ -886,16 +888,6 @@ def __init__(
         per_layer_parameters = get_per_layer_parameters(
             vllm_config, layer_names, FlashInferImpl
         )
-        if current_platform.is_device_capability(90) and any(
-            params.window_left != -1 for params in per_layer_parameters.values()
-        ):
-            # FlashInfer SM90 sliding-window prefill is not reliable with FP8-Q:
-            # https://github.com/flashinfer-ai/flashinfer/issues/3578
-            raise NotImplementedError(
-                "FlashInfer backend on SM90 currently crashes with "
-                "sliding-window attention layers. Use the default attention "
-                "backend."
-            )
         self.global_hyperparameters = infer_global_hyperparameters(per_layer_parameters)
         self.sm_scale = self.global_hyperparameters.sm_scale
         self.window_left = self.global_hyperparameters.window_left
@@ -950,7 +942,9 @@ def get_q_data_type(self, is_prefill: bool) -> torch.dtype:
         # if cache_config requests a quantized dtype globally.
         cache_dtype = self.cache_dtype
 
-        # XQA decode requires BF16/FP16-Q even with FP8 KV cache.
+        # On SM90/SM12x, XQA decode requires BF16/FP16-Q even with FP8 KV cache.
+        # FI native prefill on SM90 still uses FP8-Q in that case; SM12x prefill
+        # is fa2-only and keeps the model dtype (handled below).
         if (
             (
                 current_platform.is_device_capability(90)
@@ -964,7 +958,9 @@ def get_q_data_type(self, is_prefill: bool) -> torch.dtype:
 
         # Otherwise, match Q dtype to the KV cache dtype.
         if cache_dtype.startswith("fp8"):
-            # FP8-Q requires an fp8 tensor-core attention path.
+            # FP8-Q requires an fp8 tensor-core attention path. SM90 uses
+            # native FA3 for prefill; XQA decode on SM90/SM12x is handled
+            # above with model-dtype Q, while SM100 uses trtllm-gen.
             # Architectures with only fa2 (e.g. SM89, SM120) cannot
             # consume FP8 queries, so keep the model dtype for Q there.
             if current_platform.is_device_capability(
@@ -983,24 +979,19 @@ def get_cudagraph_support(
         vllm_config: VllmConfig,
         kv_cache_spec: KVCacheSpec,
     ) -> AttentionCGSupport:
-        """Get the cudagraph support level for FlashInfer attention.
-
-        SM90 XQA supports only single-token decode. SM12x uses the dedicated
-        XQA API, which supports speculative and non-causal decode.
-        """
-        if current_platform.is_device_capability(90):
-            return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
-
-        is_sm12x = current_platform.is_device_capability_family(120)
-        # XQA does not return LSE and therefore does not support DCP.
-        if is_sm12x and vllm_config.parallel_config.decode_context_parallel_size > 1:
+        """Get the cudagraph support level for FlashInfer attention."""
+        # XQA lacks LSE for DCP; DCP also cannot graph variable-length trtllm-gen.
+        if vllm_config.parallel_config.decode_context_parallel_size > 1:
             return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
 
         kv_specs = iter_layer_specs(kv_cache_spec)
         num_qo_heads = vllm_config.model_config.get_num_attention_heads(
             vllm_config.parallel_config
         )
-        has_trtllm_support: bool = len(kv_specs) > 0
+        is_xqa_arch = current_platform.is_device_capability(
+            90
+        ) or current_platform.is_device_capability_family(120)
+        has_uniform_batch_support: bool = len(kv_specs) > 0
         for spec in kv_specs:
             if not isinstance(spec, AttentionSpec):
                 # FlashInfer only applies to attention, so we don't consider other types
@@ -1010,13 +1001,12 @@ def get_cudagraph_support(
                 num_qo_heads=num_qo_heads,
                 num_kv_heads=spec.num_kv_heads,
                 is_prefill=False,
-            ):
-                has_trtllm_support = False
+            ) or (is_xqa_arch and not _is_xqa_head_dim_supported(spec.head_size)):
+                has_uniform_batch_support = False
                 break
 
-        if has_trtllm_support and (
-            is_sm12x or not vllm_config.attention_config.use_non_causal
-        ):
+        use_non_causal = vllm_config.attention_config.use_non_causal
+        if has_uniform_batch_support and (not use_non_causal or is_xqa_arch):
             return AttentionCGSupport.UNIFORM_BATCH
         else:
             return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
@@ -1066,10 +1056,27 @@ def _compute_decode_query_lens(
         num_decode_tokens: int,
     ) -> tuple[int, torch.Tensor | None, list[int] | None]:
         """Return the query width, ragged offsets, and effective query lengths."""
-        assert self.use_dedicated_xqa
         if num_decodes == 0 or num_decode_tokens == 0:
             return 1, None, None
 
+        if not self.use_xqa:
+            decode_q_lens = (
+                qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]
+            )
+            q_len_per_req = num_decode_tokens // num_decodes
+            if num_decode_tokens % num_decodes == 0 and bool(
+                (decode_q_lens == q_len_per_req).all().item()
+            ):
+                return q_len_per_req, None, None
+
+            # CUDA-graph padding can leave zero-length requests after a
+            # uniform speculative decode batch. Keep the packed query and
+            # describe it through the trtllm-gen varlen API.
+            max_q_len = int(decode_q_lens.max().item())
+            if max_q_len > 1:
+                return max_q_len, qo_indptr[: num_decodes + 1], None
+            return 1, None, None
+
         decode_q_lens = qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]
         nonzero = decode_q_lens[decode_q_lens > 0]
         q_len_per_req = int(nonzero.max().item()) if nonzero.numel() > 0 else 1
@@ -1128,7 +1135,10 @@ def _get_prefill_wrapper(
                     "NVFP4 KV cache."
                 )
             if self._noncausal_prefill_wrapper is None:
-                if self.has_sinks and current_platform.is_device_capability_family(120):
+                if self.has_sinks and (
+                    current_platform.is_device_capability(90)
+                    or current_platform.is_device_capability_family(120)
+                ):
                     self._noncausal_prefill_wrapper = (
                         BatchAttentionWithAttentionSinkWrapper(
                             self._get_workspace_buffer(),
@@ -1159,7 +1169,10 @@ def _get_prefill_wrapper(
                     dcp_a2a=self.dcp_a2a,
                 )
             else:
-                if self.has_sinks and current_platform.is_device_capability_family(120):
+                if self.has_sinks and (
+                    current_platform.is_device_capability(90)
+                    or current_platform.is_device_capability_family(120)
+                ):
                     assert not self.is_kvcache_nvfp4
                     self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper(
                         self._get_workspace_buffer(),
@@ -1240,9 +1253,7 @@ def _compute_flashinfer_kv_metadata(
         num_reqs: int,
         page_size: int,
     ) -> torch.Tensor:
-        """
-        Compute paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len for FlashInfer
-        attention.
+        """Compute paged_kv_indptr, paged_kv_indices and paged_kv_last_page_len.
 
         Results are stored in self.paged_kv_indptr,
         self.paged_kv_indices, self.paged_kv_last_page_len buffers.
@@ -1251,15 +1262,9 @@ def _compute_flashinfer_kv_metadata(
         """
         # write self.paged_kv_indptr_cpu inplace (0-index is always 0)
         np.cumsum(
-            num_blocks_np,
-            dtype=np.int32,
-            out=self.paged_kv_indptr.np[1 : num_reqs + 1],
+            num_blocks_np, dtype=np.int32, out=self.paged_kv_indptr.np[1 : num_reqs + 1]
         )
-        paged_kv_indptr = self.paged_kv_indptr.gpu[: num_reqs + 1]
-        paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_reqs + 1]
-        if PIN_MEMORY:
-            paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory()
-        paged_kv_indptr.copy_(paged_kv_indptr_cpu, non_blocking=True)
+        paged_kv_indptr = self.paged_kv_indptr.copy_to_gpu(num_reqs + 1)
 
         # write self.paged_kv_indices inplace
         num_actual_pages = self.paged_kv_indptr.np[num_reqs]
@@ -1279,12 +1284,7 @@ def _compute_flashinfer_kv_metadata(
             page_size,
             paged_kv_last_page_len_np,
         )
-        paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs]
-        if PIN_MEMORY:
-            paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory()
-        self.paged_kv_last_page_len.gpu[:num_reqs].copy_(
-            paged_kv_last_page_len_cpu, non_blocking=True
-        )
+        self.paged_kv_last_page_len.copy_to_gpu(num_reqs)
         return paged_kv_indices
 
     def build(
@@ -1296,13 +1296,13 @@ def build(
         num_reqs = common_attn_metadata.num_reqs
         num_actual_tokens = common_attn_metadata.num_actual_tokens
         causal = common_attn_metadata.causal
-        route_decode = causal or self.use_dedicated_xqa
+        route_decode = causal or self.use_xqa
         if route_decode:
             num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
                 split_decodes_and_prefills(
                     common_attn_metadata,
                     decode_threshold=self.reorder_batch_threshold,
-                    require_uniform=not self.use_dedicated_xqa,
+                    require_uniform=not self.use_xqa,
                 )
             )
         else:
@@ -1342,7 +1342,7 @@ def build(
             has_spec=uses_spec_reorder,
         )
         decode_with_flashinfer_trtllm_api = self.use_trtllm_decode_attention and (
-            causal or self.use_dedicated_xqa
+            causal or self.use_xqa
         )
 
         if not causal and self.use_dcp:
@@ -1360,13 +1360,12 @@ def build(
         )
 
         if not all_uses_trtllm:
-            if self.has_sinks and (
-                not self.use_dedicated_xqa or self.use_dcp or use_cascade
-            ):
+            sinks_would_be_dropped = self.use_dcp or use_cascade
+            if self.has_sinks and sinks_would_be_dropped:
                 raise NotImplementedError(
-                    "FlashInfer backend currently does not support attention "
-                    "sinks, please use trtllm on blackwell or flash attention "
-                    "on earlier GPUs."
+                    "FlashInfer BatchDCPPrefillWrapper and "
+                    "MultiLevelCascadeAttentionWrapper do not support attention sinks. "
+                    "Please use TRTLLM on Blackwell or FlashAttention on earlier GPUs."
                 )
 
             if not self.global_hyperparameters.has_same_window_lefts:
@@ -1409,7 +1408,7 @@ def build(
         needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm
         if needs_seq_lens_cpu:
             with gpu_sync_allowed():
-                seq_lens_cpu = common_attn_metadata.seq_lens_cpu
+                seq_lens_cpu = common_attn_metadata.seq_lens.cpu()
         else:
             seq_lens_cpu = None
 
@@ -1564,8 +1563,30 @@ def build(
                     qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]
                 )
                 max_q_len_prefill = int(query_lens_prefill_cpu.max().item())
+                prefill_block_tables = block_table_tensor[prefill_start:]
+                if (
+                    self.q_data_type_prefill != FP8_DTYPE
+                    and self.cache_dtype.startswith("fp8")
+                ):
+                    seq_lens_cpu_upper_bound = (
+                        common_attn_metadata.seq_lens_cpu_upper_bound
+                    )
+                    max_prefill_seq_len = max_seq_len
+                    if seq_lens_cpu_upper_bound is not None:
+                        max_prefill_seq_len = int(
+                            seq_lens_cpu_upper_bound[prefill_start:num_reqs]
+                            .max()
+                            .item()
+                        )
+                    # Dequantization allocates one page per table entry and
+                    # indexes rows by their logical width.
+                    prefill_block_tables = canonicalize_singleton_dim_strides(
+                        prefill_block_tables[
+                            :, : cdiv(max_prefill_seq_len, page_size)
+                        ].contiguous()
+                    )
                 attn_metadata.prefill = TRTLLMPrefill(
-                    block_tables=block_table_tensor[prefill_start:],
+                    block_tables=prefill_block_tables,
                     seq_lens=prefill_seq_lens,
                     cum_seq_lens_q=qo_indptr_prefill_gpu,
                     cum_seq_lens_kv=paged_kv_indptr_prefill_gpu,
@@ -1663,29 +1684,19 @@ def build(
         if num_decodes > 0:
             if decode_with_flashinfer_trtllm_api:
                 assert self.flashinfer_trtllm_api_decode_kernel is not None
-                if not self.use_dedicated_xqa and not self.use_trtllm_gen_varlen_decode:
-                    assert num_decode_tokens % num_decodes == 0, (
-                        "XQA/trtllm-gen decode requires uniform query lengths "
-                        f"per request. Got {num_decode_tokens=} and {num_decodes=}."
-                    )
                 seq_lens_decode = seq_lens[:num_decodes]
                 if self.use_dcp:
                     assert common_attn_metadata.dcp_local_seq_lens is not None
                     seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[
                         :num_decodes
                     ]
-                q_len_per_req = 1
-                q_cu_seq_lens = None
-                decode_mask = None
-                if self.use_dedicated_xqa:
-                    q_len_per_req, q_cu_seq_lens, ragged_q_lens = (
-                        self._compute_decode_query_lens(
-                            qo_indptr,
-                            qo_indptr_cpu,
-                            num_decodes,
-                            num_decode_tokens,
-                        )
+                q_len_per_req, q_cu_seq_lens, ragged_q_lens = (
+                    self._compute_decode_query_lens(
+                        qo_indptr, qo_indptr_cpu, num_decodes, num_decode_tokens
                     )
+                )
+                decode_mask = None
+                if self.use_xqa:
                     decode_mask = self._get_decode_mask(
                         q_len_per_req,
                         ragged_q_lens,
@@ -1962,8 +1973,8 @@ def forward(
         self,
         layer: torch.nn.Module,
         query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
+        key: torch.Tensor | None,
+        value: torch.Tensor | None,
         kv_cache: torch.Tensor,
         attn_metadata: FlashInferMetadata,
         output: torch.Tensor,
@@ -1973,13 +1984,23 @@ def forward(
         """Forward pass with FlashInfer.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
-            key: shape = [num_tokens, num_kv_heads, head_size]
-            value: shape = [num_tokens, num_kv_heads, head_size]
+            key: shape = [num_tokens, num_kv_heads, head_size], or None for a
+                KV-sharing decoder layer.
+            value: shape = [num_tokens, num_kv_heads, head_size], or None for a
+                KV-sharing decoder layer.
             kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size]
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization. Enables the
+                attention+quantization fusion path when provided.
+            output_block_scale: Block scale for fused output quantization,
+                required for nvfp4 output and rejected for fp8 output.
+
         Returns:
             shape = [num_tokens, num_heads * head_size]
+
         """
         if attn_metadata is None:
             # Profiling run.
@@ -2064,10 +2085,11 @@ def forward(
             if fp8_view_dtype is not None:
                 kv_cache = kv_cache.view(fp8_view_dtype)
 
-        # Inputs and outputs may be padded for CUDA graphs
+        # Inputs and outputs may be padded for CUDA graphs. Keep the original
+        # query for single-token XQA, whose non-ragged API requires one row for
+        # every (including padded) decode request.
+        query_padded = query
         query = query[:num_actual_tokens]
-        key = key[:num_actual_tokens]
-        value = value[:num_actual_tokens]
         output_padded = output
         output = output[:num_actual_tokens]
 
@@ -2131,10 +2153,7 @@ def forward(
             kv_cache_tuple = kv_cache_permute.split(hs, dim=-1)
 
         use_dcp = self.dcp_world_size > 1
-        decode_with_dedicated_xqa = (
-            decode_with_xqa and current_platform.is_device_capability_family(120)
-        )
-        if decode_with_dedicated_xqa:
+        if decode_with_xqa:
             assert not use_dcp
             assert not self.is_kvcache_nvfp4
             assert self.o_sf_scale is None
@@ -2158,6 +2177,10 @@ def forward(
                 prefill_wrapper = attn_metadata.prefill.wrapper
                 assert prefill_wrapper is not None
                 if use_dcp:
+                    if key is None or value is None:
+                        raise NotImplementedError(
+                            "FlashInfer DCP prefill does not support KV-sharing layers"
+                        )
                     assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)
                     assert prefill_wrapper._context._window_left == self.window_left
                     assert prefill_wrapper._context._logits_soft_cap == (
@@ -2176,8 +2199,8 @@ def forward(
                         layer,
                         prefill_query,
                         kv_cache_tuple,
-                        key[num_decode_tokens:],
-                        value[num_decode_tokens:],
+                        key[num_decode_tokens:num_actual_tokens],
+                        value[num_decode_tokens:num_actual_tokens],
                         out=output[num_decode_tokens:],
                     )
                 else:
@@ -2347,8 +2370,13 @@ def forward(
                     ].copy_(out[:num_prefill_tokens])
 
         if num_decode_tokens > 0:
-            decode_query = query[:num_decode_tokens]
-            assert decode_query.shape[0] == num_decode_tokens
+            decode_query_tokens = num_decode_tokens
+            if decode_with_xqa:
+                assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode)
+                if attn_metadata.decode.q_len_per_req == 1:
+                    decode_query_tokens = attn_metadata.num_decodes
+            decode_query = query_padded[:decode_query_tokens]
+            assert decode_query.shape[0] == decode_query_tokens
 
             # Convert query to the expected dtype for decode if needed.
             decode_query = self.maybe_quant_query(
@@ -2463,7 +2491,7 @@ def forward(
                     )
                     decode_query = canonicalize_singleton_dim_strides(decode_query)
 
-                if decode_with_dedicated_xqa:
+                if decode_with_xqa:
                     bmm1_scale = self.get_xqa_bmm1_scale(
                         layer, attn_metadata.q_data_type_decode
                     )
@@ -2479,7 +2507,7 @@ def forward(
                         bmm1_scale=bmm1_scale,
                         bmm2_scale=self.bmm2_scale,
                         window_left=self.window_left,
-                        out=output[:num_decode_tokens],
+                        out=output_padded[:decode_query_tokens],
                         sinks=self.sinks,
                         kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),
                         q_len_per_req=q_len_per_req,
@@ -2488,6 +2516,7 @@ def forward(
                     )
                     return output_padded
 
+                assert decode_with_trtllm_gen
                 if output.dtype == FP4_DTYPE:
                     assert self.o_sf_scale is not None
                     out = FP4Tensor(
@@ -2519,19 +2548,6 @@ def forward(
                 else:
                     q_len_per_req = num_decode_tokens // attn_metadata.num_decodes
 
-                if decode_with_xqa and q_len_per_req is not None and q_len_per_req > 1:
-                    raise NotImplementedError(
-                        "FlashInfer XQA speculative decode is not wired in vLLM yet."
-                    )
-
-                # XQA decode can use model-dtype Q with FP8 KV, so only include
-                # q_scale when the decode query is actually FP8.
-                bmm1_scale = (
-                    self.get_xqa_bmm1_scale(layer, attn_metadata.q_data_type_decode)
-                    if decode_with_xqa
-                    else self.bmm1_scale
-                )
-
                 lse = None
                 if use_dcp:
                     out = torch.empty(
@@ -2554,7 +2570,7 @@ def forward(
                     block_tables=block_tables_decode,
                     seq_lens=seq_lens_decode,
                     max_seq_len=attn_metadata.decode.max_seq_len,
-                    bmm1_scale=bmm1_scale,
+                    bmm1_scale=self.bmm1_scale,
                     bmm2_scale=self.bmm2_scale,
                     window_left=self.window_left,
                     sinks=self.sinks,
@@ -2650,8 +2666,7 @@ def fast_plan_decode(
     fixed_split_size: int = -1,
     disable_split_kv: bool = False,
 ) -> None:
-    """
-    A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for
+    """A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for
     cudagraph capture/replay, while the no cudagraph version turns back
     to the original plan.
     using original plan after passing host-side buffers:
diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py
index 7529d8d20f8e..1b9a21b65578 100644
--- a/vllm/v1/attention/backends/flex_attention.py
+++ b/vllm/v1/attention/backends/flex_attention.py
@@ -105,6 +105,10 @@ class FlexAttentionBackend(AttentionBackend):
     def get_name() -> str:
         return "FLEX_ATTENTION"
 
+    @classmethod
+    def supports_rswa(cls) -> bool:
+        return True
+
     @classmethod
     def supports_sliding_window(cls) -> bool:
         return True
@@ -150,7 +154,7 @@ def get_supported_head_sizes(cls) -> list[int]:
         return []
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
 
@@ -161,8 +165,7 @@ def physical_to_logical_mapping(
     block_size: int,
     total_blocks: int,
 ) -> torch.Tensor:
-    """
-    Creates an inverse mapping from physical block locations to logical indices.
+    """Creates an inverse mapping from physical block locations to logical indices.
 
     The original block_table maps from logical blocks to physical locations:
 
@@ -235,6 +238,7 @@ def physical_to_logical_mapping(
         A tensor of shape [max_reqs, total_blocks] where each entry
         physical_to_logical[req_id, physical_block] contains the logical
         block index for that physical block, or -1 if unused.
+
     """
     max_reqs, max_num_blocks = block_table.shape
     device = block_table.device
@@ -271,8 +275,7 @@ def unique_static_unsorted(
     ignored_val: int = 0,  # value to ignore
     pad_val: int = -1,  # sentinel for unused slots
 ) -> torch.Tensor:
-    """
-    - Keeps the first occurrence of each non-zero value while preserving order,
+    """- Keeps the first occurrence of each non-zero value while preserving order,
       then left-packs those uniques and fills the rest with `pad_val`.
     - Returns (packed, keep_mask) with the *same shape* as `x`.
     - Requires that all values be in the range [0, M]
@@ -282,6 +285,7 @@ def unique_static_unsorted(
 
     Example:
     x =[3, 1, 0, 1, 2], M=3, ignored_val=0 => [3, 1, 2, -1, -1]
+
     """
     if not (-1 <= pad_val <= M):
         raise ValueError("`pad_val` must lie in [-1, M]")
@@ -345,6 +349,7 @@ class BlockSparsityHint(NamedTuple):
         hint_fn: (q_block_idx [num_tokens, 1], kv_block_idx [1, num_kv_blocks],
             block_size int) -> bool Tensor [num_tokens, num_kv_blocks].
             Returns True for block pairs that may contain non-masked elements.
+
     """
 
     hint_fn: _block_sparsity_hint_signature
@@ -431,6 +436,7 @@ def _convert_physical_to_logical(
 
         Returns:
             tuple of (is_valid, logical_q_idx, logical_kv_idx)
+
         """
         # Map query indices to corresponding request indices
         q_req = request_lookup[q_idx]
@@ -508,7 +514,6 @@ def get_sliding_window_mask_mod(self) -> _mask_mod_signature:
         Note that the sliding window mask here is bidirectional, we need
         to mask it with the bidirectional/causal mask for encoder/decoder.
         """
-
         if self.sliding_window is None:
             raise ValueError("sliding_window must be set for sliding window attention")
 
@@ -536,7 +541,6 @@ def final_mask_mod(
 
     def get_prefix_lm_mask_mod(self) -> _mask_mod_signature:
         """Creates the prefix LM mask_mod function for FlexAttention."""
-
         assert self.doc_ids is not None
         request_lookup = self.doc_ids
 
@@ -583,7 +587,6 @@ def get_rswa_mask_mod(self) -> _mask_mod_signature:
         far-away generated tokens that fall outside the window and outside the
         prefix.
         """
-
         assert self.doc_ids is not None
         assert self.rswa_prefix_lens is not None
         assert self.rswa_window is not None
@@ -1287,14 +1290,22 @@ def forward(
         """Forward pass with FLexAttention.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
                 [num_blocks, num_kv_heads, block_size, 2 * head_size]
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization; not supported
+                by this backend.
+            output_block_scale: Block scale for fused output quantization;
+                not supported by this backend.
+
         Returns:
             shape = [num_tokens, num_heads * head_size]
+
         """
         if output_scale is not None or output_block_scale is not None:
             raise NotImplementedError(
diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py
index 23daff54300b..c2cfd003d90d 100644
--- a/vllm/v1/attention/backends/gdn_attn.py
+++ b/vllm/v1/attention/backends/gdn_attn.py
@@ -247,9 +247,34 @@ def build(  # type: ignore[override]
                 )
 
         if spec_sequence_masks is None:
+            # V2 already excludes prefills from full decode graphs via
+            # has_prefill. Classify first chunks as prefills to mask recycled
+            # state; resumed one-token chunks can still use the decode kernels.
+            assert m.seq_lens_cpu_upper_bound is not None
+            query_lens_cpu = query_start_loc_cpu.diff()
+            no_prior_state = (query_lens_cpu > 0) & (
+                m.seq_lens_cpu_upper_bound <= query_lens_cpu
+            )
+            # Capture batches also have seq_len == query_len, but are not
+            # prefills.
+            if m.is_prefilling is not None:
+                no_prior_state &= m.is_prefilling
+            else:
+                no_prior_state = torch.zeros_like(no_prior_state)
             num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
-                split_decodes_and_prefills(m, decode_threshold=1)
+                split_decodes_and_prefills(
+                    m.replace(is_prefilling=no_prior_state),
+                    decode_threshold=1,
+                    treat_short_extends_as_decodes=False,
+                )
             )
+            # Exclude trailing padding from both prefill counts.
+            if num_prefills:
+                num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum())
+                num_prefill_tokens = (
+                    int(query_start_loc_cpu[num_decodes + num_prefills])
+                    - num_decode_tokens
+                )
             num_spec_decode_tokens = 0
             spec_token_indx = None
             non_spec_token_indx = None
@@ -524,8 +549,7 @@ def build(  # type: ignore[override]
     def build_for_cudagraph_capture(
         self, common_attn_metadata: CommonAttentionMetadata
     ):
-        """
-        This method builds the metadata for full cudagraph capture.
+        """This method builds the metadata for full cudagraph capture.
         Currently, only decode is supported for full cudagraphs with Mamba.
         """
         m = common_attn_metadata
diff --git a/vllm/v1/attention/backends/hpc_attn.py b/vllm/v1/attention/backends/hpc_attn.py
index 14545307fef0..cbe97da45c9d 100644
--- a/vllm/v1/attention/backends/hpc_attn.py
+++ b/vllm/v1/attention/backends/hpc_attn.py
@@ -288,7 +288,7 @@ class HpcAttentionBackend(AttentionBackend):
     forward_includes_kv_cache_update: bool = True
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64]
 
     @staticmethod
diff --git a/vllm/v1/attention/backends/mamba2_attn.py b/vllm/v1/attention/backends/mamba2_attn.py
index 189ef54b236d..0a4584632dca 100644
--- a/vllm/v1/attention/backends/mamba2_attn.py
+++ b/vllm/v1/attention/backends/mamba2_attn.py
@@ -23,8 +23,7 @@ def compute_varlen_chunk_metadata(
     query_start_loc: torch.Tensor,
     chunk_size: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Build chunk-aligned, variable-length metadata used by Mamba2 SSD kernels.
+    """Build chunk-aligned, variable-length metadata used by Mamba2 SSD kernels.
 
     Given per-sequence cumulative token starts `query_start_loc` of shape [B+1]
     and a physical `chunk_size`, returns three tensors on the same device:
diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py
index 1b9ee263db5b..b5f0b9b182e1 100644
--- a/vllm/v1/attention/backends/mamba_attn.py
+++ b/vllm/v1/attention/backends/mamba_attn.py
@@ -227,8 +227,7 @@ def __init__(
     def build_for_cudagraph_capture(
         self, common_attn_metadata: CommonAttentionMetadata
     ) -> M:
-        """
-        This method builds the metadata for full cudagraph capture.
+        """This method builds the metadata for full cudagraph capture.
         Currently, only decode is supported for full cudagraphs with Mamba.
         """
         m = common_attn_metadata
@@ -276,8 +275,7 @@ def build(
         num_decode_draft_tokens_cpu: torch.Tensor | None = None,
         **kwargs: Any,
     ) -> M:
-        """
-        Default build implementation for Mamba-like attention backends.
+        """Default build implementation for Mamba-like attention backends.
         Subclasses (e.g., Mamba2) can override to add additional metadata.
         """
         return self._compute_common_metadata(
@@ -294,8 +292,7 @@ def _compute_chunk_metadata(
         num_computed_tokens_p_cpu: torch.Tensor,
         query_start_loc_p_cpu: torch.Tensor,
     ) -> tuple[list[int], list[int], list[int]]:
-        """
-        Compute chunk-specific metadata for Mamba models.
+        """Compute chunk-specific metadata for Mamba models.
 
         The code below carefully constructs the chunks such that:
         1. Chunks contain tokens from a *single* sequence only.
@@ -382,8 +379,7 @@ def _build_chunk_metadata_tensors(
         common: M,
         common_attn_metadata: CommonAttentionMetadata,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Compute chunk metadata and return as device tensors.
+        """Compute chunk metadata and return as device tensors.
         Returns (cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p).
         """
         num_prefills = common.num_prefills
@@ -463,9 +459,7 @@ def _compute_common_metadata(
         prev_last_scheduled_idx: torch.Tensor | None = None,
         num_decode_draft_tokens_cpu: torch.Tensor | None = None,
     ) -> M:
-        """
-        Compute metadata common to both Mamba1 and Mamba2.
-        """
+        """Compute metadata common to both Mamba1 and Mamba2."""
         num_reqs = common_attn_metadata.num_reqs
 
         # Treat multi-token queries as decode requests when
@@ -633,13 +627,21 @@ def _compute_common_metadata(
 
         if self.use_replayssm and not self.use_flashinfer_replayssm and num_decodes > 0:
             decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu
-            num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu
-            if decode_base_cpu is None or num_computed_tokens_cpu is None:
+            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+            async_spec_decode = (
+                self.vllm_config.scheduler_config.async_scheduling
+                and self.vllm_config.speculative_config is not None
+            )
+            if decode_base_cpu is None or seq_lens_cpu is None or async_spec_decode:
                 raise ValueError(
-                    "--use-replayssm requires CPU decode-base and "
-                    "computed-token counts to derive decode write positions"
+                    "--use-replayssm requires exact CPU sequence lengths and "
+                    "decode-base counts to derive decode write positions"
                 )
-            num_computed_d = num_computed_tokens_cpu[:num_decodes]
+            query_lens_cpu = (
+                common_attn_metadata.query_start_loc_cpu[1 : num_decodes + 1]
+                - common_attn_metadata.query_start_loc_cpu[:num_decodes]
+            )
+            num_computed_d = seq_lens_cpu[:num_decodes] - query_lens_cpu
             decode_base_d = decode_base_cpu[:num_decodes]
             align_mode = self.vllm_config.cache_config.mamba_cache_mode == "align"
             block_size = self.kv_cache_spec.block_size
@@ -656,10 +658,6 @@ def _compute_common_metadata(
             # write_pos counts decode steps since the ring's last full-state
             # write (the anchor), so a resumed request re-anchors correctly.
             decode_steps_cpu = num_computed_d - effective_base
-            query_lens_cpu = (
-                common_attn_metadata.query_start_loc_cpu[1 : num_decodes + 1]
-                - common_attn_metadata.query_start_loc_cpu[:num_decodes]
-            )
             valid_decode_rows = query_lens_cpu > 0
             # A single-token prefill row replayed as decode (query_len==1 with
             # prior state) has decode_steps < 0; force it to a one-token flush
@@ -748,8 +746,7 @@ def _update_metadata_for_cudagraph_capture(
         self,
         metadata: M,
     ) -> M:
-        """
-        Update the metadata for cudagraph capture.
+        """Update the metadata for cudagraph capture.
         Currently, only decode is supported for full cudagraphs with Mamba.
         """
         state_indices_tensor_d = metadata.state_indices_tensor_d
diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py
index 5f9ac7f17adf..16689887af0a 100644
--- a/vllm/v1/attention/backends/mla/amx_mla.py
+++ b/vllm/v1/attention/backends/mla/amx_mla.py
@@ -61,6 +61,44 @@
 _SPLIT_OCCUPANCY_MULTIPLIER = 2
 
 
+def _block_quant_to_tensor_quant(
+    weight: torch.Tensor,  # FP8 [N, K] block-quantized weight
+    weight_scale: torch.Tensor,  # [ceil(N/bn), ceil(K/bk)] block scales
+    block_size: list[int],  # [block_n, block_k]
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Re-quantize FP8 block-quant weight to FP8 per-tensor quant.
+
+    Returns (fp8_weight, dequant_scale) where fp8_weight * dequant_scale
+    approximates the original float weight.  dequant_scale is a 0-dim
+    scalar tensor as expected by bmm_cpu's ``scale`` argument.
+
+    NOTE: block->tensor re-quantization may slightly affect accuracy
+    (equivalent to SGLang's block_quant_to_tensor_quant for CPU).
+    """
+    block_n, block_k = block_size[0], block_size[1]
+    n, k = weight.shape
+
+    # Expand block scales [ceil(N/bn), ceil(K/bk)] -> [N, K]
+    scale = weight_scale.to(torch.float32)
+    if scale.shape[0] < n:
+        scale = scale.repeat_interleave(block_n, dim=0)[:n, :]
+    if scale.shape[1] < k:
+        scale = scale.repeat_interleave(block_k, dim=1)[:, :k]
+
+    # Dequantize to float32
+    x_f32 = weight.to(torch.float32) * scale
+
+    # Re-quantize per-tensor back to FP8
+    fp8_max = torch.finfo(weight.dtype).max  # 448.0 for e4m3fn
+    amax = x_f32.abs().amax().clamp(min=1e-12)
+    quant_scale = fp8_max / amax
+    x_fp8 = (x_f32 * quant_scale).clamp(-fp8_max, fp8_max).to(weight.dtype)
+
+    # dequant_scale = 1 / quant_scale (0-dim tensor for bmm_cpu)
+    dequant_scale = (amax / fp8_max).to(torch.float32)
+    return x_fp8, dequant_scale
+
+
 def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int:
     """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread
     count in place of SM count."""
@@ -92,7 +130,7 @@ def get_supported_head_sizes(cls) -> list[int]:
         return [576]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(32)]
 
     @classmethod
@@ -208,6 +246,9 @@ def __init__(
             )
         self._w_uk_packed: torch.Tensor | None = None
         self._w_uv_packed: torch.Tensor | None = None
+        self._w_scale: torch.Tensor | None = (
+            None  # per-tensor dequant scale (FP8 W8A16)
+        )
 
     def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
         # forward_mha (unlike forward_mqa) receives no `layer` argument, so it
@@ -216,12 +257,40 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
         # impl derives its own copies from self.kv_b_proj (already available
         # via the constructor) and VNNI-packs them once, ahead of time, for
         # use with bmm_cpu in forward_mha.
-        kv_b_proj_weight = get_and_maybe_dequant_weights(
-            self.kv_b_proj, out_dtype=act_dtype
-        ).T
-        kv_b_proj_weight = kv_b_proj_weight.view(
-            self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
-        )
+        #
+        # kv_b_proj carries `_cpu_skip_gemm_dispatch`, so the CPU FP8 kernels
+        # leave its weight in the checkpoint [N, K] layout. For block-quant
+        # FP8 checkpoints, re-quantize per-tensor and keep the BMM in FP8
+        # (W8A16); everything else goes through the BF16 path.
+        weight = getattr(self.kv_b_proj, "weight", None)
+        block_size = getattr(self.kv_b_proj, "weight_block_size", None)
+        if (
+            weight is not None
+            and weight.dtype == torch.float8_e4m3fn
+            and block_size is not None
+        ):
+            block_scale = getattr(self.kv_b_proj, "weight_scale_inv", None)
+            if block_scale is None:
+                block_scale = self.kv_b_proj.weight_scale
+            fp8_weight, self._w_scale = _block_quant_to_tensor_quant(
+                weight, block_scale, block_size
+            )
+            # [N, K] -> [K, N] = [kv_lora_rank, heads * (nope + v)]
+            kv_b_proj_weight = fp8_weight.T.contiguous().view(
+                self.kv_lora_rank,
+                self.num_heads,
+                self.qk_nope_head_dim + self.v_head_dim,
+            )
+        else:
+            kv_b_proj_weight = get_and_maybe_dequant_weights(
+                self.kv_b_proj, out_dtype=act_dtype
+            ).T
+            kv_b_proj_weight = kv_b_proj_weight.view(
+                self.kv_lora_rank,
+                self.num_heads,
+                self.qk_nope_head_dim + self.v_head_dim,
+            )
+            self._w_scale = None
         w_uk, w_uv = kv_b_proj_weight.split(
             [self.qk_nope_head_dim, self.v_head_dim], dim=-1
         )
@@ -348,7 +417,9 @@ def forward_mha(
         ql_nope = torch.empty(
             num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device
         )
-        ops.bmm_cpu(ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, None)
+        ops.bmm_cpu(
+            ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, self._w_scale
+        )
         mqa_q = torch.cat([ql_nope, q_pe], dim=-1)
 
         kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size)
@@ -402,4 +473,4 @@ def forward_mha(
         output_view = output.view(num_tokens, num_heads, self.v_head_dim).transpose(
             0, 1
         )
-        ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, None)
+        ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, self._w_scale)
diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py
index a7d801b0406c..25e24f347c0e 100644
--- a/vllm/v1/attention/backends/mla/compressor_utils.py
+++ b/vllm/v1/attention/backends/mla/compressor_utils.py
@@ -5,12 +5,11 @@
 
 import torch
 
-from vllm.model_executor.warmup.jit_warmup import zip_inputs
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
     triton_scalar_specialization_rep,
 )
 from vllm.triton_utils import tl, triton
@@ -43,6 +42,8 @@ class CompileKey:
     @triton.jit(do_not_specialize=["block_table_stride"])
     def kernel(
         # [num_tokens]
+        compressed_slot_mapping_ptr,
+        # [num_tokens] the tokens' own slots
         slot_mapping_ptr,
         # [num_reqs + 1]
         query_start_loc_ptr,
@@ -70,7 +71,12 @@ def kernel(
             mask = offset < query_len
 
             pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
-            is_valid = (pos + 1) % COMPRESS_RATIO == 0
+            # A replayed token has a PAD slot (SWA bounded replay: its KV is
+            # cached already); its compressed KV is cached too, so PAD here.
+            slot = tl.load(
+                slot_mapping_ptr + query_start + offset, mask=mask, other=PAD_ID
+            )
+            is_valid = ((pos + 1) % COMPRESS_RATIO == 0) & (slot != PAD_ID)
             pos_after_compress = pos // COMPRESS_RATIO
 
             block_ids = pos_after_compress // block_size
@@ -82,7 +88,9 @@ def kernel(
 
             # NOTE
             slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
-            tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)
+            tl.store(
+                compressed_slot_mapping_ptr + query_start + offset, slot_ids, mask=mask
+            )
 
     def dispatch(  # type: ignore[override]
         self,
@@ -97,10 +105,10 @@ def dispatch(  # type: ignore[override]
         )
 
     def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
-        hf_config = vllm_config.model_config.hf_config
+        hf_text_config = vllm_config.model_config.hf_text_config
         configured_ratios = (
-            *(getattr(hf_config, "compress_ratios", None) or ()),
-            getattr(hf_config, "index_kpool", 1) or 1,
+            *(getattr(hf_text_config, "compress_ratios", None) or ()),
+            getattr(hf_text_config, "index_kpool", 1) or 1,
         )
         compress_ratios = tuple(
             dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1)
@@ -122,6 +130,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
     def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
         int32_ptr = TritonWarmupTensor(torch.int32)
         return dict(
+            compressed_slot_mapping=TritonWarmupTensor(torch.int64),
             slot_mapping=TritonWarmupTensor(torch.int64),
             query_start_loc=int32_ptr,
             seq_lens=int32_ptr,
@@ -133,6 +142,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
     @kernel_launcher
     def __call__(
         self,
+        compressed_slot_mapping: torch.Tensor,
         slot_mapping: torch.Tensor,
         query_start_loc: torch.Tensor,
         seq_lens: torch.Tensor,
@@ -150,6 +160,7 @@ def __call__(
 
 def get_compressed_slot_mapping(
     num_tokens: int,
+    slot_mapping: torch.Tensor,
     query_start_loc: torch.Tensor,
     seq_lens: torch.Tensor,
     block_table: torch.Tensor,
@@ -157,19 +168,28 @@ def get_compressed_slot_mapping(
     compress_ratio: int,
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
+    """Slot mapping for writing the compressed states of ``num_tokens`` tokens.
+
+    Every ``compress_ratio`` tokens share one compressed state, written by the
+    last of them: that token maps to the state's slot, the others to PAD. A
+    token whose own ``slot_mapping`` entry is PAD maps to PAD too: SWA bounded
+    replay recomputes tokens whose KV is cached already, and their compressed
+    states must not be rewritten either.
+    """
     if out is not None:
         # Guard: for padded / invalid sequences.
         # Negative positions produce bogus block indices that lead to illegal memory
         # accesses inside the block_table load.
         # NOTE: Fill -1 to the whole tensor, not just the first `num_tokens`.
         out.fill_(-1)
-        slot_mapping = out[:num_tokens]
+        compressed_slot_mapping = out[:num_tokens]
     else:
-        slot_mapping = torch.full(
+        compressed_slot_mapping = torch.full(
             (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device
         )
 
     _COMPRESSED_SLOT_MAPPING_KERNEL(
+        compressed_slot_mapping,
         slot_mapping,
         query_start_loc,
         seq_lens,
@@ -177,7 +197,7 @@ def get_compressed_slot_mapping(
         block_size,
         compress_ratio,
     )
-    return slot_mapping
+    return compressed_slot_mapping
 
 
 _COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel()
diff --git a/vllm/v1/attention/backends/mla/cpu_mla.py b/vllm/v1/attention/backends/mla/cpu_mla.py
index 91d3cb333986..431f6c318533 100644
--- a/vllm/v1/attention/backends/mla/cpu_mla.py
+++ b/vllm/v1/attention/backends/mla/cpu_mla.py
@@ -64,7 +64,7 @@ def get_supported_head_sizes(cls) -> list[int]:
         return [576]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [16]
 
     @classmethod
diff --git a/vllm/v1/attention/backends/mla/cutlass_mla.py b/vllm/v1/attention/backends/mla/cutlass_mla.py
index 575b9d32a9b2..b88a4a17a822 100644
--- a/vllm/v1/attention/backends/mla/cutlass_mla.py
+++ b/vllm/v1/attention/backends/mla/cutlass_mla.py
@@ -46,7 +46,7 @@ class CutlassMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [128]
 
     @staticmethod
diff --git a/vllm/v1/attention/backends/mla/flashattn_mla.py b/vllm/v1/attention/backends/mla/flashattn_mla.py
index 084a1f4ce6c9..b685be102ed6 100644
--- a/vllm/v1/attention/backends/mla/flashattn_mla.py
+++ b/vllm/v1/attention/backends/mla/flashattn_mla.py
@@ -49,7 +49,7 @@ class FlashAttnMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
     @staticmethod
@@ -194,8 +194,11 @@ def _build_decode(
         query_start_loc_cpu: torch.Tensor,
         query_start_loc_device: torch.Tensor,
         num_decode_tokens: int,
+        max_query_len: int,
         dcp_tot_seq_lens_device: torch.Tensor | None,
     ) -> FlashAttnMLADecodeMetadata:
+        # Measured, not the promised bound: _schedule_decode asserts metadata size
+        # against a persistent buffer.
         query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
         max_query_len = query_lens_cpu.max().item()
 
diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
index 4a78a700cfaa..6048f6b5fd1b 100644
--- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
@@ -8,9 +8,9 @@
 
 from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
-from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
 from vllm.model_executor.layers.attention.sparse_mla_attention import (
     SparseMLACommonImpl,
+    SparseMLACommonMetadata,
     SparseMLACommonMetadataBuilder,
 )
 from vllm.platforms.interface import DeviceCapability
@@ -18,11 +18,11 @@
     AttentionBackend,
     AttentionCGSupport,
     AttentionLayer,
-    AttentionMetadata,
     MLAAttentionImpl,
     MultipleOf,
 )
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.sparse_utils import (
     flat_kv_row_view,
     triton_convert_req_index_to_global_index,
@@ -40,7 +40,7 @@ class FlashAttnMLASparseBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64]
 
     @staticmethod
@@ -99,33 +99,15 @@ def supports_combination(
             if vllm_config.parallel_config.decode_context_parallel_size > 1:
                 return "FlashAttention MLA Sparse does not support DCP for now"
 
-            hf_config = vllm_config.model_config.hf_config
-            if not hasattr(hf_config, "index_topk"):
+            hf_text_config = vllm_config.model_config.hf_text_config
+            if not hasattr(hf_text_config, "index_topk"):
                 return "FlashAttention MLA Sparse requires model with index_topk"
         return None
 
 
 @dataclass
-class FlashAttnMLASparseMetadata(AttentionMetadata):
-    num_reqs: int
-    max_query_len: int
-    max_seq_len: int
-
-    num_actual_tokens: int
-    query_start_loc: torch.Tensor
-    slot_mapping: torch.Tensor
-
-    block_table: torch.Tensor
-    req_id_per_token: torch.Tensor
-    seq_lens: torch.Tensor
-    block_size: int = 64
-    topk_tokens: int = 2048
-    num_decodes: int = 0
-    num_prefills: int = 0
-    num_decode_tokens: int = 0
-    prefill_max_seq_len: int = 0
-    prefill: MLACommonPrefillMetadata | None = None
-    cp_kv_cache_interleave_size: int = 1
+class FlashAttnMLASparseMetadata(SparseMLACommonMetadata):
+    pass
 
 
 class FlashAttnMLASparseMetadataBuilder(
@@ -214,6 +196,76 @@ def forward_mqa(
 
         assert self.topk_indices_buffer is not None
         topk_indices = self.topk_indices_buffer[:num_actual_toks]
+        index_group = self.index_group
+        if isinstance(index_group, HiSparseMLAIndexGroup):
+            num_decode_tokens = attn_metadata.num_decode_tokens
+            outputs = []
+            if num_decode_tokens:
+                physical_topk, valid_counts = (
+                    index_group.convert_decode_logical_to_physical_topk(
+                        self.index_group_index,
+                        topk_indices[:num_decode_tokens],
+                        attn_metadata,
+                        return_valid_counts=True,
+                    )
+                )
+                outputs.append(
+                    self._run_mqa_kernel(
+                        q_nope[:num_decode_tokens],
+                        q_rope[:num_decode_tokens],
+                        index_group.physical_kv_cache(self.index_group_index).view(
+                            kv_c_and_k_pe_cache.dtype
+                        ),
+                        physical_topk,
+                        valid_counts,
+                        attn_metadata.block_size,
+                    )
+                )
+            if num_decode_tokens < num_actual_toks:
+                cache = index_group.cache(self.index_group_index)
+                if num_decode_tokens == 0 and cache.all_context_pages_resident:
+                    physical_topk, valid_counts = (
+                        index_group.convert_logical_to_physical_topk(
+                            self.index_group_index,
+                            topk_indices,
+                            attn_metadata,
+                            block_stride_rows=None,
+                            return_valid_counts=True,
+                        )
+                    )
+                    prefill_cache = index_group.physical_kv_cache(
+                        self.index_group_index
+                    ).view(kv_c_and_k_pe_cache.dtype)
+                else:
+                    prefill_cache, block_table, req_ids = (
+                        index_group.stage_prefill_rows(
+                            self.index_group_index,
+                            kv_c_and_k_pe_cache,
+                            attn_metadata,
+                        )
+                    )
+                    physical_topk, valid_counts = (
+                        triton_convert_req_index_to_global_index(
+                            req_ids,
+                            block_table,
+                            topk_indices[num_decode_tokens:],
+                            BLOCK_SIZE=attn_metadata.block_size,
+                            NUM_TOPK_TOKENS=topk_indices.shape[1],
+                            return_valid_counts=True,
+                        )
+                    )
+                outputs.append(
+                    self._run_mqa_kernel(
+                        q_nope[num_decode_tokens:],
+                        q_rope[num_decode_tokens:],
+                        prefill_cache,
+                        physical_topk,
+                        valid_counts,
+                        attn_metadata.block_size,
+                    )
+                )
+            return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None
+
         kv_rows, block_stride_rows = flat_kv_row_view(
             kv_c_and_k_pe_cache, attn_metadata.block_size
         )
@@ -226,12 +278,47 @@ def forward_mqa(
             NUM_TOPK_TOKENS=topk_indices.shape[1],
             return_valid_counts=True,
         )
+        return (
+            self._run_mqa_kernel(
+                q_nope,
+                q_rope,
+                kv_rows,
+                topk_indices,
+                valid_counts,
+                attn_metadata.block_size,
+                cache_is_flat=True,
+            ),
+            None,
+        )
+
+    def _run_mqa_kernel(
+        self,
+        q_nope: torch.Tensor,
+        q_rope: torch.Tensor,
+        kv_cache: torch.Tensor,
+        topk_indices: torch.Tensor,
+        valid_counts: torch.Tensor,
+        block_size: int,
+        *,
+        cache_is_flat: bool = False,
+    ) -> torch.Tensor:
+        kv_rows = (
+            kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0]
+        )
 
         cu_seqlens_q = torch.arange(
-            0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device
+            0, q_rope.shape[0] + 1, dtype=torch.int32, device=q_rope.device
         )
-        k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
         v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1)
+        if self.qk_rope_head_dim == 0:
+            # FA3's QV path requires the 64-wide Q/K specialization. For NoPE
+            # MLA, zero Q preserves QV-only attention scores while providing a
+            # valid TMA shape.
+            _FA3_QV_HEAD_DIM = 64
+            q_rope = q_nope.new_zeros(*q_nope.shape[:-1], _FA3_QV_HEAD_DIM)
+            k_cache = v_cache[..., :_FA3_QV_HEAD_DIM]
+        else:
+            k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
 
         out = flash_attn_varlen_func(
             q=q_rope,
@@ -247,4 +334,4 @@ def forward_mqa(
             causal=True,
             fa_version=3,
         )
-        return out, None
+        return out
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py
index 61414bba49df..ca78961eb591 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, ClassVar
 
 import torch
@@ -15,6 +16,7 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention.mla_attention import (
     MLACommonBackend,
+    MLACommonDecodeMetadata,
     MLACommonImpl,
     MLACommonMetadata,
     MLACommonMetadataBuilder,
@@ -115,9 +117,28 @@ def _get_multi_ctas_kv_counter_buffer(
     return _fi_multi_ctas_kv_counter
 
 
-class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
-    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
-    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
+@dataclass
+class FlashInferMLADecodeMetadata(MLACommonDecodeMetadata):
+    flattened_block_table: torch.Tensor | None = None
+    flattened_seq_lens: torch.Tensor | None = None
+    flattened_row_req: torch.Tensor | None = None
+    # Row count the flattened tensors were built for (0 = not built), not a query len.
+    query_len: int = 0
+    query_start_loc: torch.Tensor | None = None
+    max_query_len: int = 1
+
+
+@dataclass
+class FlashInferMLAMetadata(MLACommonMetadata[FlashInferMLADecodeMetadata]):
+    pass
+
+
+class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[FlashInferMLAMetadata]):
+    # Adaptive verification requires ALWAYS from every builder, matching upstream's
+    # DeepseekV4FlashMLAMetadataBuilder. The kernels tile ragged queries from the device
+    # query offsets (flashinfer #3238), so one k+1 graph replays any 1..k+1 mix.
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
+    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.VARLEN
     # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.
     supports_non_causal_multi_token_decode: ClassVar[bool] = True
     supports_non_causal_multi_token_dcp: ClassVar[bool] = True
@@ -143,10 +164,31 @@ def __init__(
             layer_names,
             vllm_config,
             device,
-            MLACommonMetadata,
+            FlashInferMLAMetadata,
             supports_dcp_with_varlen=True,
         )
 
+    def _build_decode(
+        self,
+        block_table_tensor: torch.Tensor,
+        seq_lens_device: torch.Tensor,
+        max_seq_len: int,
+        query_start_loc_cpu: torch.Tensor,
+        query_start_loc_device: torch.Tensor,
+        num_decode_tokens: int,
+        max_query_len: int,
+        dcp_tot_seq_lens_device: torch.Tensor | None,
+    ) -> FlashInferMLADecodeMetadata:
+        # Promised bound, not measured: a capture dummy puts one token/req, so
+        # measuring would bake an undersized graph. Ragged only when max_query_len>1.
+        return FlashInferMLADecodeMetadata(
+            block_table=block_table_tensor,
+            seq_lens=seq_lens_device,
+            dcp_tot_seq_lens=dcp_tot_seq_lens_device,
+            query_start_loc=query_start_loc_device,
+            max_query_len=max_query_len,
+        )
+
 
 class FlashInferMLABackend(MLACommonBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
@@ -159,7 +201,7 @@ class FlashInferMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [32, 64]
 
     @staticmethod
@@ -208,7 +250,7 @@ def supports_combination(
         return None
 
 
-class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
+class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):
     can_return_lse_for_decode: bool = True
     supports_dcp: bool = True
     # DCP is the only path that consumes LSE. It uses monolithic CuTeDSL,
@@ -276,7 +318,7 @@ def forward_mqa(
         self,
         q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
         kv_c_and_k_pe_cache: torch.Tensor,
-        attn_metadata: MLACommonMetadata,
+        attn_metadata: FlashInferMLAMetadata,
         layer: AttentionLayer,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         assert kv_c_and_k_pe_cache.numel() > 0
@@ -289,14 +331,32 @@ def forward_mqa(
         block_table = attn_metadata.decode.block_table
         seq_lens = attn_metadata.decode.seq_lens
 
+        # Led by the promised bound (baked into the graph), not
+        # num_decode_tokens // num_decodes -- that average is a per-request length
+        # only for uniform batches.
+        multi_token_decode = (
+            attn_metadata.decode.max_query_len > 1
+            or attn_metadata.num_decode_tokens > attn_metadata.num_decodes
+        )
+        cum_seq_lens_q: torch.Tensor | None = None
+        max_q_len: int | None = None
+        row_req: torch.Tensor | None = None
+
         if not attn_metadata.causal:
             # FlashInfer decode has no causal flag. Flatten each non-causal
             # query block into independent single-token rows.
-            query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
             q = q.unsqueeze(1)
-            if query_len > 1:
-                block_table = block_table.repeat_interleave(query_len, dim=0)
-                seq_lens = seq_lens.repeat_interleave(query_len)
+            if multi_token_decode:
+                block_table, seq_lens, row_req = self._flattened_decode_metadata(
+                    attn_metadata, q.shape[0]
+                )
+        elif attn_metadata.decode.max_query_len > 1 and self.dcp_world_size == 1:
+            # Causal spec decode: keep q compact and let the kernel tile each
+            # request's length (uniform 1+k and adaptive ragged). Mirrors vllm #52157.
+            # DCP keeps the uniform reshape below: it needs LSE, which the kernels
+            # do not return on the ragged path (flashinfer #3238).
+            cum_seq_lens_q = attn_metadata.decode.query_start_loc
+            max_q_len = attn_metadata.decode.max_query_len
         # trtllm API requires extra dimension q_len_per_request for MTP
         elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
             logger.warning_once(
@@ -327,10 +387,8 @@ def forward_mqa(
         if self.dcp_world_size > 1:
             causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens
             assert causal_seqlens_kv_global is not None
-            if not attn_metadata.causal and query_len > 1:
-                causal_seqlens_kv_global = causal_seqlens_kv_global.repeat_interleave(
-                    query_len
-                )
+            if row_req is not None:
+                causal_seqlens_kv_global = causal_seqlens_kv_global[row_req]
             extra_kwargs.update(
                 enable_dcp=True,
                 cp_world=self.dcp_world_size,
@@ -342,6 +400,16 @@ def forward_mqa(
             # trtllm-gen rejects MLA head counts it can't tile (e.g. 96);
             # fall back to cute-dsl for those.
             decode_backend = _select_mla_decode_backend(runtime_num_heads)
+        if cum_seq_lens_q is not None:
+            # Neither decode backend returns LSE on the ragged path
+            # (flashinfer #3238); DCP, the only LSE consumer, took the uniform
+            # branch above, so this only guards a future caller wiring the two.
+            assert not return_lse, (
+                "FlashInferMLA ragged decode cannot return LSE; DCP and adaptive "
+                "variable-length decode are mutually exclusive."
+            )
+            extra_kwargs["cum_seq_lens_q"] = cum_seq_lens_q
+            extra_kwargs["max_q_len"] = max_q_len
         if decode_backend:
             extra_kwargs["backend"] = decode_backend
         elif kv_c_and_k_pe_cache.shape[-2] in (32, 64):
@@ -386,3 +454,37 @@ def forward_mqa(
         o = o.view(-1, o.shape[-2], o.shape[-1])
 
         return o, lse
+
+    def _flattened_decode_metadata(
+        self,
+        attn_metadata: FlashInferMLAMetadata,
+        num_rows: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Expand per-request decode tensors to one row per query token.
+
+        Cached across the layers of a group, which all see the same batch.
+        """
+        decode = attn_metadata.decode
+        assert decode is not None
+        if decode.query_len != num_rows:
+            cu = decode.query_start_loc
+            assert cu is not None
+            # searchsorted on the device offsets, not repeat_interleave(uniform_len):
+            # the latter only lines up for uniform batches and hands ragged rows
+            # another request's KV (silent garbage).
+            rows = torch.arange(num_rows, device=cu.device, dtype=cu.dtype)
+            row_req = torch.searchsorted(cu[1:], rows, right=True).clamp_(
+                max=decode.block_table.shape[0] - 1
+            )
+            decode.flattened_row_req = row_req
+            decode.flattened_block_table = decode.block_table[row_req]
+            decode.flattened_seq_lens = decode.seq_lens[row_req]
+            decode.query_len = num_rows
+        assert decode.flattened_block_table is not None
+        assert decode.flattened_seq_lens is not None
+        assert decode.flattened_row_req is not None
+        return (
+            decode.flattened_block_table,
+            decode.flattened_seq_lens,
+            decode.flattened_row_req,
+        )
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
index b34c98375261..32f67e7e936c 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
@@ -11,24 +11,26 @@
 from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
-from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
 from vllm.model_executor.layers.attention.sparse_mla_attention import (
     SparseMLACommonImpl,
+    SparseMLACommonMetadata,
     SparseMLACommonMetadataBuilder,
 )
+from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
     AttentionLayer,
-    AttentionMetadata,
     AttentionType,
     MLAAttentionImpl,
     MultipleOf,
 )
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.sparse_utils import (
     flat_kv_row_view,
+    prepare_sparse_mla_safe_lengths,
     triton_convert_req_index_to_global_index,
     triton_filter_and_convert_dcp_index,
 )
@@ -81,7 +83,7 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [32, 64]
 
     @staticmethod
@@ -112,6 +114,16 @@ def supports_combination(
         from vllm.config import get_current_vllm_config
 
         vllm_config = get_current_vllm_config()
+        parallel_config = vllm_config.parallel_config
+        if (
+            parallel_config.prefill_context_parallel_size > 1
+            and parallel_config.decode_context_parallel_size > 1
+        ):
+            return (
+                "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; "
+                "use FLASHMLA_SPARSE, which gathers each DCP KV shard before "
+                "running the rank-local PCP prefill queries"
+            )
         if kv_cache_dtype == "fp8_ds_mla":
             return (
                 "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype"
@@ -160,7 +172,7 @@ def get_name() -> str:
         return "FLASHINFER_MLA_SPARSE_SM120"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64, 256]
 
     @staticmethod
@@ -224,33 +236,9 @@ def supports_combination(
 
 
 @dataclass
-class FlashInferMLASparseMetadata(AttentionMetadata):
+class FlashInferMLASparseMetadata(SparseMLACommonMetadata):
     """Attention metadata for FlashInfer MLA Sparse backend."""
 
-    num_reqs: int
-    max_query_len: int
-    max_seq_len: int
-    num_actual_tokens: int
-
-    # Query start locations
-    query_start_loc: torch.Tensor
-    slot_mapping: torch.Tensor
-    block_table: torch.Tensor
-    req_id_per_token: torch.Tensor
-
-    # Sequence lengths for all requests (context + query)
-    seq_lens: torch.Tensor
-    num_decodes: int
-    num_prefills: int
-    num_decode_tokens: int
-    prefill_max_seq_len: int = 0
-    prefill: MLACommonPrefillMetadata | None = None
-
-    # Sparse-specific
-    block_size: int = 64
-    topk_tokens: int = 2048
-    cp_kv_cache_interleave_size: int = 1
-
 
 class FlashInferMLASparseMetadataBuilder(
     SparseMLACommonMetadataBuilder[FlashInferMLASparseMetadata]
@@ -286,6 +274,32 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde
     """Metadata builder for the SM100 TRT-LLM sparse MLA kernel."""
 
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
+    hisparse_supports_multi_token_decode: ClassVar[bool] = True
+
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+
+        # Under DCP the workspace must hold the trtllm-gen softmax-stats slab.
+        # The buffer address is baked into CUDA graphs, so it has to reach its
+        # final size here, before the first forward/capture (no lazy regrow).
+        if self.dcp_world_size > 1:
+            num_q_heads = vllm_config.model_config.get_num_attention_heads(
+                vllm_config.parallel_config
+            )
+            _get_workspace_buffer(
+                device,
+                _required_workspace_bytes(
+                    self.dcp_world_size,
+                    num_q_heads,
+                    vllm_config.scheduler_config.max_num_batched_tokens,
+                ),
+            )
 
     def _build_req_id_per_token(
         self,
@@ -297,14 +311,98 @@ def _build_req_id_per_token(
 # Global workspace buffer (lazily initialized)
 _fi_sparse_workspace: torch.Tensor | None = None
 
+# trtllm-gen carves a softmax-stats slab from the workspace whenever LSE is
+# requested (FlashInfer csrc/trtllm_fmha_kernel_launcher.cu, unchanged from
+# v0.6.14 through current main):
+#   sizeof(float2) * num_qo_heads * batch_size * round_up(max_q_len, 256)
+#   + 1 MiB guard
+# forward_mqa always passes q_len == 1, so each (head, token) pair costs
+# round_up(1, 256) == 256 slots.
+_TRTLLM_GEN_SOFTMAX_STAT_BYTES = 8  # sizeof(float2)
+_TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN = 256  # round_up(q_len=1, 256)
+_TRTLLM_GEN_SOFTMAX_GUARD_BYTES = 1024 * 1024
+
+# Keep in sync with the VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE default in
+# vllm/envs.py.
+_DEFAULT_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024
+
+
+def compute_trtllm_sparse_mla_workspace_bytes(
+    base_workspace_bytes: int,
+    dcp_world_size: int,
+    num_heads_per_rank: int,
+    max_num_batched_tokens: int,
+) -> int:
+    """Workspace bytes needed by the trtllm-gen sparse MLA decode kernel.
+
+    Under DCP the query is all-gathered across the DCP group in the head dim
+    and the kernel is asked for LSE, so trtllm-gen carves the softmax-stats
+    slab described above from the workspace before its (batch-independent)
+    counter and scratch regions. ``base_workspace_bytes`` must stay available
+    for those regions, so the slab is added on top of it (see #50781).
+
+    Without DCP no LSE is requested, no slab is carved, and the base size is
+    returned unchanged.
+    """
+    if dcp_world_size <= 1:
+        return base_workspace_bytes
+    softmax_bytes = (
+        _TRTLLM_GEN_SOFTMAX_STAT_BYTES
+        * (num_heads_per_rank * dcp_world_size)
+        * max_num_batched_tokens
+        * _TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN
+        + _TRTLLM_GEN_SOFTMAX_GUARD_BYTES
+    )
+    return base_workspace_bytes + softmax_bytes
+
+
+def _required_workspace_bytes(
+    dcp_world_size: int,
+    num_heads_per_rank: int,
+    max_num_batched_tokens: int,
+) -> int:
+    """Resolve the workspace size, honoring an explicit env override."""
+    computed = compute_trtllm_sparse_mla_workspace_bytes(
+        _DEFAULT_WORKSPACE_BUFFER_SIZE,
+        dcp_world_size,
+        num_heads_per_rank,
+        max_num_batched_tokens,
+    )
+    if not envs.is_set("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE"):
+        return computed
+    env_bytes = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
+    if env_bytes < computed:
+        logger.warning_once(
+            "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d is below the %d bytes "
+            "computed for the sparse MLA workspace with dcp_world_size=%d, "
+            "%d heads/rank and max_num_batched_tokens=%d. Respecting the "
+            "override, but the trtllm-gen kernel may crash with a workspace "
+            "overflow; set VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d or unset "
+            "it to use the computed size.",
+            env_bytes,
+            computed,
+            dcp_world_size,
+            num_heads_per_rank,
+            max_num_batched_tokens,
+            computed,
+        )
+    return env_bytes
+
 
-def _get_workspace_buffer(device: torch.device) -> torch.Tensor:
+def _get_workspace_buffer(
+    device: torch.device, min_bytes: int | None = None
+) -> torch.Tensor:
     global _fi_sparse_workspace
-    if _fi_sparse_workspace is None:
+    required = (
+        min_bytes
+        if min_bytes is not None
+        else envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
+    )
+    if _fi_sparse_workspace is None or _fi_sparse_workspace.numel() < required:
         # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace;
         # the trtllm-gen path views it as uint8, so int8 is safe for all backends.
         _fi_sparse_workspace = torch.zeros(
-            envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
+            required,
             dtype=torch.int8,
             device=device,
         )
@@ -402,6 +500,78 @@ def forward_mqa(
         assert self.topk_indices_buffer is not None
         topk_indices = self.topk_indices_buffer[:num_actual_toks]
 
+        self._prepare_mqa_kernel(layer, q.device)
+
+        index_group = self.index_group
+        if isinstance(index_group, HiSparseMLAIndexGroup):
+            num_decode_tokens = attn_metadata.num_decode_tokens
+            decode_out: torch.Tensor | None = None
+            decode_lse: torch.Tensor | None = None
+            if num_decode_tokens > 0:
+                physical_topk, valid_counts = (
+                    index_group.convert_decode_logical_to_physical_topk(
+                        self.index_group_index,
+                        topk_indices[:num_decode_tokens],
+                        attn_metadata,
+                        return_valid_counts=True,
+                    )
+                )
+                decode_out, decode_lse = self._run_mqa_kernel(
+                    q[:num_decode_tokens],
+                    index_group.physical_kv_cache(self.index_group_index).view(
+                        kv_c_and_k_pe_cache.dtype
+                    ),
+                    physical_topk,
+                    valid_counts,
+                )
+                if num_decode_tokens == num_actual_toks:
+                    return decode_out, decode_lse
+
+            cache = index_group.cache(self.index_group_index)
+            if num_decode_tokens == 0 and cache.all_context_pages_resident:
+                physical_topk, valid_counts = (
+                    index_group.convert_logical_to_physical_topk(
+                        self.index_group_index,
+                        topk_indices,
+                        attn_metadata,
+                        block_stride_rows=None,
+                        return_valid_counts=True,
+                    )
+                )
+                return self._run_mqa_kernel(
+                    q,
+                    index_group.physical_kv_cache(self.index_group_index).view(
+                        kv_c_and_k_pe_cache.dtype
+                    ),
+                    physical_topk,
+                    valid_counts,
+                )
+
+            prefill_cache, block_table, req_ids = index_group.stage_prefill_rows(
+                self.index_group_index, kv_c_and_k_pe_cache, attn_metadata
+            )
+            prefill_indices, prefill_lens = triton_convert_req_index_to_global_index(
+                req_ids,
+                block_table,
+                topk_indices[num_decode_tokens:],
+                BLOCK_SIZE=attn_metadata.block_size,
+                NUM_TOPK_TOKENS=topk_indices.shape[1],
+                return_valid_counts=True,
+            )
+            prefill_out, prefill_lse = self._run_mqa_kernel(
+                q[num_decode_tokens:],
+                prefill_cache,
+                prefill_indices,
+                prefill_lens,
+            )
+            if decode_out is None:
+                return prefill_out, prefill_lse
+            output = torch.cat((decode_out, prefill_out))
+            if decode_lse is None:
+                return output, None
+            assert prefill_lse is not None
+            return output, torch.cat((decode_lse, prefill_lse))
+
         _, block_stride_rows = flat_kv_row_view(
             kv_c_and_k_pe_cache, attn_metadata.block_size
         )
@@ -420,18 +590,27 @@ def forward_mqa(
                 return_valid_counts=True,
             )
         else:
-            topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(
-                attn_metadata.req_id_per_token[:num_actual_toks],
-                attn_metadata.block_table,
+            topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk(
                 topk_indices,
-                BLOCK_SIZE=attn_metadata.block_size,
-                BLOCK_STRIDE_ROWS=block_stride_rows,
-                NUM_TOPK_TOKENS=topk_indices.shape[1],
+                attn_metadata,
+                block_stride_rows=block_stride_rows,
                 return_valid_counts=True,
             )
 
+        return self._run_mqa_kernel(
+            q,
+            kv_c_and_k_pe_cache,
+            topk_indices_physical,
+            seq_lens,
+        )
+
+    def _prepare_mqa_kernel(
+        self,
+        layer: AttentionLayer,
+        device: torch.device,
+    ) -> None:
         if self._workspace_buffer is None:
-            self._workspace_buffer = _get_workspace_buffer(q.device)
+            self._workspace_buffer = _get_workspace_buffer(device)
 
         if self.bmm1_scale is None:
             self.bmm1_scale = self.scale
@@ -442,15 +621,71 @@ def forward_mqa(
             if is_quantized_kv_cache(self.kv_cache_dtype):
                 self.bmm2_scale *= layer._k_scale_float
 
+    def autotune_hisparse_decode(self, layer: AttentionLayer) -> None:
+        """Autotune the largest legal HiSparse decode batch."""
+        assert isinstance(self.index_group, HiSparseMLAIndexGroup)
+        cache = self.index_group.cache(self.index_group_index)
+        assert self.topk_indices_buffer is not None
+
+        runtime = cache.runtime
+        kv_cache = runtime.hot.attention_cache
+        num_tokens = runtime.max_num_reqs
+        topk_tokens = self.topk_indices_buffer.shape[1]
+        self._prepare_mqa_kernel(layer, kv_cache.device)
+
+        q_dtype = (
+            current_platform.fp8_dtype()
+            if is_quantized_kv_cache(self.kv_cache_dtype)
+            else kv_cache.dtype
+        )
+        q = torch.zeros(
+            (
+                num_tokens,
+                self.num_heads,
+                self.kv_lora_rank + self.qk_rope_head_dim,
+            ),
+            dtype=q_dtype,
+            device=kv_cache.device,
+        )
+        topk_indices = (
+            torch.arange(
+                topk_tokens,
+                dtype=torch.int32,
+                device=kv_cache.device,
+            )
+            .expand(num_tokens, -1)
+            .contiguous()
+        )
+        seq_lens = torch.full(
+            (num_tokens,),
+            topk_tokens,
+            dtype=torch.int32,
+            device=kv_cache.device,
+        )
+        self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens)
+
+    def _run_mqa_kernel(
+        self,
+        q: torch.Tensor,
+        kv_cache: torch.Tensor,
+        topk_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        assert self._workspace_buffer is not None
+        assert self.bmm1_scale is not None
+        assert self.bmm2_scale is not None
+
         from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
 
+        kv_cache = kv_cache.view(q.dtype)
+
         # Single-token sparse decode. trtllm-gen requires the q_len_per_request
         # dim, but the sparse attention mask is fully per-token (each query token
         # carries its own top-k index row), so unsqueeze is sufficient and
         # correct. The MTP/multi-token q_len grouping is a perf-only layout and is
         # deferred until MTP is validated end-to-end for this backend.
         query = q.unsqueeze(1)
-        block_tables = topk_indices_physical.unsqueeze(1)
+        block_tables = topk_indices.unsqueeze(1)
         seq_lens_arg = seq_lens
 
         # page_table width = topk buffer width, which kpool widens past
@@ -460,27 +695,23 @@ def forward_mqa(
         # count), so the -1 padding slots past seq_lens are never attended to.
         # Use the actual buffer width instead of the fixed topk_tokens, which
         # mismatches the page_table when index_kpool > 1.
-        sparse_topk_capacity = topk_indices_physical.shape[1]
+        sparse_topk_capacity = topk_indices.shape[1]
 
         extra_kwargs: dict[str, torch.Tensor] = {}
-        empty_rows: torch.Tensor | None = None
+        needs_empty_query_guard = self.need_to_return_lse_for_decode or isinstance(
+            self.index_group, HiSparseMLAIndexGroup
+        )
         if self.is_nope_mla:
-            # The native no-rope kernel takes the active top-k length per query
-            # token (``seq_lens`` here is already the compacted per-token valid
-            # count, int32) and rejects zero-length rows. Point empty rows at a
-            # single valid dummy slot with length 1 and zero their output after
-            # the launch. ``triton_convert_req_index_to_global_index`` packs the
-            # valid indices into a contiguous prefix, which is what the kernel
-            # requires of the page table.
-            empty_rows = seq_lens == 0
-            topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill(
-                empty_rows, 0
-            )
-            extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1)
+            # Resident TP queries have nonempty selections. Preserve empty-query
+            # handling for DCP's local selections and host-backed HiSparse.
+            topk_lens = seq_lens
+            if needs_empty_query_guard:
+                topk_lens = prepare_sparse_mla_safe_lengths(topk_indices, seq_lens)
+            extra_kwargs["sparse_mla_top_k_lens"] = topk_lens
 
         kernel_out = trtllm_batch_decode_with_kv_cache_mla(
             query=query,
-            kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
+            kv_cache=kv_cache.unsqueeze(1),
             workspace_buffer=self._workspace_buffer,
             qk_nope_head_dim=self.qk_nope_head_dim,
             kv_lora_rank=self.kv_lora_rank,
@@ -505,12 +736,17 @@ def forward_mqa(
         out = o.view(-1, o.shape[-2], o.shape[-1])
         if lse is not None:
             lse = self._normalize_lse(lse, out.shape[0], out.shape[1])
-        if empty_rows is None and lse is not None:
-            empty_rows = (topk_indices_physical == -1).all(dim=-1)
-        if empty_rows is not None:
-            out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
+        if self.is_nope_mla and needs_empty_query_guard:
+            empty_queries = seq_lens == 0
             if lse is not None:
-                lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
+                # DCP combine already suppresses outputs with zero LSE weight.
+                lse.masked_fill_(empty_queries[:, None], float("-inf"))
+            else:
+                out.masked_fill_(empty_queries[:, None, None], 0.0)
+        elif lse is not None:
+            empty_rows = (topk_indices == -1).all(dim=-1)
+            out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
+            lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
         return out, lse
 
     @staticmethod
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
index 4de2c2f53f3b..e59c9ce46db9 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
@@ -6,15 +6,15 @@
 
 import torch
 
-from vllm.v1.attention.backend import (
-    AttentionLayer,
-    AttentionType,
-    MLAAttentionImpl,
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SparseMLACommonImpl,
 )
+from vllm.v1.attention.backend import AttentionLayer, AttentionType
 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
     FlashInferMLASparseMetadata,
     _get_workspace_buffer,
 )
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.sparse_utils import (
     triton_convert_req_index_to_global_index,
 )
@@ -29,7 +29,7 @@ def _kv_scale_format_for_model(model_type: str | None) -> str:
     return "pow2_fp32"
 
 
-class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata]):
+class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
     """SM120 FlashInfer sparse-MLA implementation."""
 
     is_sparse = True
@@ -60,20 +60,28 @@ def __init__(
                 "FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention"
             )
 
-        self.num_heads = num_heads
-        self.head_size = head_size
-        self.scale = float(scale)
-        self.num_kv_heads = num_kv_heads
-        self.kv_cache_dtype = kv_cache_dtype
-        if self.kv_cache_dtype != "fp8_ds_mla":
+        if kv_cache_dtype != "fp8_ds_mla":
             raise NotImplementedError(
                 "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla "
                 f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}."
             )
 
-        self.kv_lora_rank: int = mla_args["kv_lora_rank"]
-        self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]
-        self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
+        topk_indices_buffer = mla_args.pop("topk_indices_buffer", None)
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            indexer=indexer,
+            topk_indices_buffer=topk_indices_buffer,
+            **mla_args,
+        )
         from vllm.config import get_current_vllm_config
 
         vllm_config = get_current_vllm_config()
@@ -84,13 +92,6 @@ def __init__(
             )
         self.kv_scale_format = _kv_scale_format_for_model(model_type)
 
-        # Skip-topk layers are built with indexer=None and get the shared
-        # buffer via mla_args instead (cf. FLASHMLA_SPARSE).
-        self.topk_indices_buffer: torch.Tensor | None = (
-            indexer.topk_indices_buffer
-            if indexer is not None
-            else mla_args.get("topk_indices_buffer")
-        )
         from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
 
         if not has_flashinfer_sparse_mla_sm120():
@@ -117,6 +118,70 @@ def forward_mqa(
 
         assert self.topk_indices_buffer is not None
         topk_indices = self.topk_indices_buffer[:num_actual_toks]
+        index_group = self.index_group
+        if isinstance(index_group, HiSparseMLAIndexGroup):
+            num_decode_tokens = attn_metadata.num_decode_tokens
+            outputs = []
+            if num_decode_tokens:
+                topk_indices_physical = cast(
+                    torch.Tensor,
+                    index_group.convert_decode_logical_to_physical_topk(
+                        self.index_group_index,
+                        topk_indices[:num_decode_tokens],
+                        attn_metadata,
+                        return_valid_counts=False,
+                    ),
+                )
+                outputs.append(
+                    self._run_mqa_kernel(
+                        q[:num_decode_tokens],
+                        index_group.physical_kv_cache(self.index_group_index),
+                        topk_indices_physical,
+                    )
+                )
+            if num_decode_tokens < num_actual_toks:
+                cache = index_group.cache(self.index_group_index)
+                if num_decode_tokens == 0 and cache.all_context_pages_resident:
+                    topk_indices_physical = cast(
+                        torch.Tensor,
+                        index_group.convert_logical_to_physical_topk(
+                            self.index_group_index,
+                            topk_indices,
+                            attn_metadata,
+                            block_stride_rows=None,
+                            return_valid_counts=False,
+                        ),
+                    )
+                    prefill_cache = index_group.physical_kv_cache(
+                        self.index_group_index
+                    )
+                else:
+                    prefill_cache, block_table, req_ids = (
+                        index_group.stage_prefill_rows(
+                            self.index_group_index,
+                            kv_c_and_k_pe_cache,
+                            attn_metadata,
+                        )
+                    )
+                    topk_indices_physical = cast(
+                        torch.Tensor,
+                        triton_convert_req_index_to_global_index(
+                            req_ids,
+                            block_table,
+                            topk_indices[num_decode_tokens:],
+                            BLOCK_SIZE=attn_metadata.block_size,
+                            NUM_TOPK_TOKENS=topk_indices.shape[1],
+                        ),
+                    )
+                outputs.append(
+                    self._run_mqa_kernel(
+                        q[num_decode_tokens:],
+                        prefill_cache,
+                        topk_indices_physical,
+                    )
+                )
+            output = torch.cat(outputs) if len(outputs) > 1 else outputs[0]
+            return output, None
 
         topk_indices_physical = cast(
             torch.Tensor,
@@ -128,6 +193,22 @@ def forward_mqa(
                 NUM_TOPK_TOKENS=topk_indices.shape[1],
             ),
         )
+        return (
+            self._run_mqa_kernel(
+                q,
+                kv_c_and_k_pe_cache,
+                topk_indices_physical,
+            ),
+            None,
+        )
+
+    def _run_mqa_kernel(
+        self,
+        q: torch.Tensor,
+        kv_cache: torch.Tensor,
+        topk_indices_physical: torch.Tensor,
+    ) -> torch.Tensor:
+        num_actual_toks = q.shape[0]
 
         output = q.new_empty(
             (num_actual_toks, self.num_heads, self.kv_lora_rank),
@@ -141,20 +222,21 @@ def forward_mqa(
             flashinfer_trtllm_batch_decode_with_kv_cache_mla,
         )
 
+        sparse_capacity = topk_indices_physical.shape[1]
         out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
             query=q.unsqueeze(1),
-            kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1),
+            kv_cache=kv_cache.view(torch.uint8).unsqueeze(1),
             workspace_buffer=self._workspace_buffer,
             qk_nope_head_dim=self.qk_nope_head_dim,
             kv_lora_rank=self.kv_lora_rank,
             qk_rope_head_dim=self.qk_rope_head_dim,
             block_tables=topk_indices_physical.unsqueeze(1),
             seq_lens=None,
-            max_seq_len=attn_metadata.topk_tokens,
+            max_seq_len=sparse_capacity,
             out=output.unsqueeze(1),
             bmm1_scale=self.scale,
             bmm2_scale=1.0,
-            sparse_mla_top_k=attn_metadata.topk_tokens,
+            sparse_mla_top_k=sparse_capacity,
             kv_scale_format=self.kv_scale_format,
         )
-        return out.squeeze(1), None
+        return out.squeeze(1)
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
index 90dc4cd94738..27fd633601b6 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
@@ -43,6 +43,7 @@
     SparseMLACommonImpl,
 )
 from vllm.platforms.interface import DeviceCapability
+from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import has_flashinfer_sm90_nope_mla
 from vllm.v1.attention.backend import (
     AttentionBackend,
@@ -64,6 +65,23 @@
 _WORKSPACE_BYTES = 128 * 1024 * 1024
 
 
+@triton.jit
+def _pack_topk_indices(
+    slots,
+    indptr,
+    indices,
+    WIDTH: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    start = tl.load(indptr + row)
+    end = tl.load(indptr + row + 1)
+    mask = (cols < end - start) & (cols < WIDTH)
+    values = tl.load(slots + row * WIDTH + cols, mask=mask, other=0)
+    tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask)
+
+
 class FlashInferMLASparseSM90Backend(AttentionBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -74,7 +92,7 @@ class FlashInferMLASparseSM90Backend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(64)]
 
     @staticmethod
@@ -192,10 +210,11 @@ def __init__(
         self.kv_len_arr = torch.full(
             (max_tokens,), topk_width, dtype=torch.int32, device=device
         )
+        self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32, device=device)
         self.wrapper = BatchMLAPagedAttentionWrapper(
             self.workspace,
             qo_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
-            kv_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
+            kv_indptr=self.kv_indptr,
             kv_indices=self.kv_indices,
             kv_len_arr=self.kv_len_arr,
             use_cuda_graph=True,
@@ -234,12 +253,13 @@ def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
         # fixed (max_tokens+1)-sized buffers with exact-size copy_, so the
         # indptr must always be full-size. Rows past num_tokens are padded
         # empty (qo_indptr flat at num_tokens) — zero-query rows read no q
-        # and schedule no work. Padded rows keep the full width lens; the
-        # value is never dereferenced.
+        # and schedule no work. FlashInfer 0.7 requires packed CSR offsets
+        # matching the exact lengths, including zero-length padded rows.
         torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)
-        torch.mul(self._qo_cpu, self.topk_width, out=self._kv_cpu)
-        self._lens_cpu.fill_(self.topk_width)
+        self._lens_cpu.zero_()
         self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32)
+        self._kv_cpu[0] = 0
+        torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:])
         self.wrapper.plan(
             self._qo_cpu,
             self._kv_cpu,
@@ -255,6 +275,12 @@ def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
             kv_data_type=self.kv_dtype,
         )
 
+    def pack_indices(self, topk_slots: torch.Tensor) -> None:
+        num_tokens, width = topk_slots.shape
+        _pack_topk_indices[(num_tokens, triton.cdiv(width, 256))](
+            topk_slots, self.kv_indptr, self.kv_indices, width, 256
+        )
+
 
 @dataclass
 class FlashInferMLASparseSM90Metadata(FlashInferMLASparseMetadata):
@@ -277,6 +303,7 @@ def __init__(
         attention_layer = vllm_config.compilation_config.static_forward_context[
             layer_names[0]
         ]
+        self._attention_layer = attention_layer
         impl = attention_layer.impl
         if not isinstance(impl, FlashInferMLASparseSM90Impl):
             raise TypeError(
@@ -376,6 +403,13 @@ def build(
         # Replan every step outside any CUDA graph capture with this step's
         # exact per-row lengths; captured runs read the refreshed buffers.
         num_rows, kv_lens = self._kv_lens_host(common_attn_metadata)
+        # MHA prefills are removed from q before forward_mqa. Plan only the
+        # decode prefix in that case, using the same routing decision as MLA.
+        if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha(
+            metadata
+        ):
+            num_rows = metadata.num_decode_tokens
+            kv_lens = kv_lens[:num_rows]
         self.state.plan(num_rows, kv_lens)
         metadata.state = self.state
         return metadata
@@ -454,12 +488,8 @@ def forward_mqa(
         )
         state = attn_metadata.state
         assert state is not None
-        # Refresh top-k rows in graph and clamp masked tails to a valid slot;
-        # per-row lengths are already baked into the host-side plan.
-        width = topk_slots.shape[1]
-        state.kv_indices[: num_tokens * width].copy_(
-            topk_slots.reshape(-1).clamp_(min=0).to(torch.int32)
-        )
+        # Pack valid prefixes at the offsets refreshed by plan() before replay.
+        state.pack_indices(topk_slots)
 
         flat = (
             kv_c_and_k_pe_cache.view(torch.float8_e4m3fn)
diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py
index 8ebfbbba796c..c8cd31d564d1 100644
--- a/vllm/v1/attention/backends/mla/flashmla.py
+++ b/vllm/v1/attention/backends/mla/flashmla.py
@@ -55,7 +55,7 @@ class FlashMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64]
 
     @staticmethod
@@ -158,6 +158,7 @@ def _build_decode(
         query_start_loc_cpu: torch.Tensor,
         query_start_loc_device: torch.Tensor,
         num_decode_tokens: int,
+        max_query_len: int,
         dcp_tot_seq_lens_device: torch.Tensor | None,
     ) -> FlashMLADecodeMetadata:
         query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py
index b837cc366baf..014358f530e0 100644
--- a/vllm/v1/attention/backends/mla/flashmla_sparse.py
+++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py
@@ -3,35 +3,45 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar
 
+import numpy as np
 import torch
+import torch.distributed as dist
 
 from vllm import _custom_ops as ops
 from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.config.cache import CacheDType
+from vllm.distributed.parallel_state import get_dcp_group
 from vllm.logger import init_logger
-from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
 from vllm.model_executor.layers.attention.sparse_mla_attention import (
     SparseMLACommonImpl,
+    SparseMLACommonMetadata,
     SparseMLACommonMetadataBuilder,
 )
 from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
+from vllm.utils.math_utils import round_up
 from vllm.utils.platform_utils import num_compute_units
-from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.utils.torch_utils import (
+    async_tensor_h2d,
+    current_stream,
+    is_quantized_kv_cache,
+)
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
     AttentionLayer,
-    AttentionMetadata,
     CommonAttentionMetadata,
     MultipleOf,
 )
+from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.sparse_utils import (
     flat_kv_row_view,
+    request_row_bounds,
     triton_convert_req_index_to_global_index,
     triton_filter_and_convert_dcp_index,
 )
 from vllm.v1.attention.backends.utils import (
+    get_dcp_local_seq_lens,
     reshape_attn_output_for_spec_decode,
     reshape_query_for_spec_decode,
     split_prefill_chunks,
@@ -85,6 +95,17 @@
     The first `ue8m0` is the scale for the first 64 `float8_e4m3` values,
     the second for the next 64, and so on.
 
+For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as:
+-   **First 512 bytes:** all 512 dims as `float8_e4m3`. Unlike V4 the 64 RoPE
+    dims are quantized too, so there is no `bfloat16` part.
+-   **Last 16 bytes:** Scale factors, containing 16 `ue8m0` values, one per 32
+    consecutive `float8_e4m3` values (i.e. MXFP8).
+
+The V4 and V4.1 records are not laid out token-by-token within a page: a page
+holds all its data rows first and all its scale rows after, so a page is
+`block_size * bytes_per_token` bytes rounded up to the decode kernel's TMA
+stride (576 B for V4, 512 B for V4.1).
+
 In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each
 token's KV cache is 352 Bytes, structured as:
 -   **First 256 bytes:** 512 `e2m1` NoPE values packed 2/byte (low nibble =
@@ -118,7 +139,7 @@ class FlashMLASparseBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64]
 
     @staticmethod
@@ -135,8 +156,8 @@ def get_impl_cls() -> type["FlashMLASparseImpl"]:
 
     @classmethod
     def get_supported_head_sizes(cls) -> list[int]:
-        # DeepSeek V3.2 layout: 512 NoPE + 64 RoPE = 576.
-        return [576]
+        # DeepSeek V3.2: 512 NoPE + 64 RoPE = 576; GLM5Next NoPE: 512.
+        return [576, 512]
 
     @classmethod
     def is_mla(cls) -> bool:
@@ -163,6 +184,37 @@ def supports_combination(
         use_mm_prefix: bool,
         device_capability: DeviceCapability,
     ) -> str | None:
+        if head_size == 512:
+            # GLM5Next NoPE (qk_rope_head_dim == 0, kv_lora_rank == 512) has
+            # head_size 512 and is served here by the direct 512-wide bf16
+            # cache on SM90. Quantized DS-MLA caches need the zero-padded
+            # 576/656B envelope, which is wired up separately; plain fp8,
+            # SM100 bf16, and rope-carrying 512 models must fall through to
+            # FlashInfer/TRITON.
+            if (
+                kv_cache_dtype in (None, "auto", "bfloat16", "float16")
+                and device_capability.major == 9
+            ):
+                # Direct bf16 NoPE-512 is only correct for rope-free models.
+                # Precedent for reading hf_text_config in supports_combination:
+                # flashinfer_mla_sparse.py.
+                from vllm.config import get_current_vllm_config
+
+                vllm_config = get_current_vllm_config()
+                if vllm_config.model_config is not None:
+                    hf_text_config = vllm_config.model_config.hf_text_config
+                    if getattr(hf_text_config, "qk_rope_head_dim", 64) != 0:
+                        return (
+                            "FLASHMLA_SPARSE supports head_size 512 only for "
+                            "rope-free (NoPE) models"
+                        )
+            else:
+                return (
+                    "FLASHMLA_SPARSE supports head_size 512 only with bf16 "
+                    "kv-cache on SM90 (NoPE), got "
+                    f"kv_cache_dtype={kv_cache_dtype}, "
+                    f"capability={device_capability}"
+                )
         if kv_cache_dtype == "nvfp4_ds_mla" and device_capability.major != 10:
             return (
                 f"FLASHMLA_SPARSE only supports the {kv_cache_dtype} kv-cache "
@@ -171,29 +223,30 @@ def supports_combination(
         return None
 
 
-@dataclass
-class FlashMLASparseMetadata(AttentionMetadata):
-    num_reqs: int
-    max_query_len: int
-    max_seq_len: int
-
-    num_actual_tokens: int  # Number of tokens excluding padding.
-    query_start_loc: torch.Tensor
-    slot_mapping: torch.Tensor
-
-    block_table: torch.Tensor
-    req_id_per_token: torch.Tensor
-    block_size: int = 64
-    topk_tokens: int = 2048
-
-    num_decodes: int = 0
-    num_prefills: int = 0
-    num_decode_tokens: int = 0
-    seq_lens: torch.Tensor | None = None
-    prefill_max_seq_len: int = 0
-    prefill: MLACommonPrefillMetadata | None = None
-    cp_kv_cache_interleave_size: int = 1
+def gathered_prefill_shards(
+    row_req_idx: np.ndarray, row_shard_rows: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+    """Under PCP+DCP, group this rank's prefill rows by request and take each
+    request's slice of this rank's KV shard.
+
+    PCP gives a rank two adjacent chunk rows of a split prefill; they share
+    one context, so the workspace holds it once. ``row_shard_rows`` is
+    ``CommonAttentionMetadata.dcp_local_seq_lens_cpu_upper_bound``: the
+    largest DCP shard of the whole request on every row, identical on every
+    PCP rank, so the all-gather is shaped the same everywhere.
+
+    Returns the request row bounds and, per request, its shard rows.
+    """
+    row_bounds = request_row_bounds(row_req_idx)
+    shard_rows = row_shard_rows[row_bounds[:-1]].astype(np.int64)
+    assert np.all(shard_rows > 0), (
+        f"PCP+DCP prefill got an empty context: {shard_rows.tolist()}"
+    )
+    return row_bounds, shard_rows
 
+
+@dataclass
+class FlashMLASparseMetadata(SparseMLACommonMetadata):
     @dataclass
     class FP8KernelMetadata:
         scheduler_metadata: FlashMLASchedMeta
@@ -233,6 +286,7 @@ class Chunk:
                 req_start_idx: int
                 workspace_starts: torch.Tensor
                 chunk_tot_seqlen: int
+                seq_lens: torch.Tensor | None = None
 
             chunks: list[Chunk]
 
@@ -263,6 +317,7 @@ class FlashMLASparseMetadataBuilder(
 ):
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
     require_uniform_decodes: ClassVar[bool] = True
+    hisparse_supports_multi_token_decode: ClassVar[bool] = True
     metadata_cls = FlashMLASparseMetadata
 
     def __init__(
@@ -283,6 +338,9 @@ def __init__(
             )
         else:
             threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256)
+        self.use_hisparse = vllm_config.attention_config.hisparse_config is not None
+        if self.use_hisparse:
+            threshold = 1
         # Varlen decodes are safe under DCP: causality comes from the
         # indexer's top-k indices, not from the kernel metadata.
         self._init_reorder_batch_threshold(
@@ -347,7 +405,14 @@ def __init__(
             device=device,
         )
 
-        self.fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL
+        # PCP+DCP attends prefill rows over the DCP-gathered KV, which only the
+        # separate prefill/decode path can do.
+        self.pcp_dcp_kv_gather = self.use_pcp and self.dcp_world_size > 1
+        self.fp8_use_mixed_batch = (
+            self.num_heads < MIN_HEADS_FOR_BF16_PREFILL
+            and not self.use_hisparse
+            and not self.pcp_dcp_kv_gather
+        )
 
         if parallel_config.decode_context_parallel_size > 1:
             if parallel_config.dcp_comm_backend != "ag_rs":
@@ -356,7 +421,12 @@ def __init__(
                     "default 'ag_rs' DCP comm backend; got "
                     f"'{parallel_config.dcp_comm_backend}'"
                 )
-            if not self.fp8_use_mixed_batch:
+            if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla":
+                raise NotImplementedError(
+                    "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla "
+                    f"upconvert; got a {cache_config.cache_dtype} cache"
+                )
+            if not self.fp8_use_mixed_batch and not self.pcp_dcp_kv_gather:
                 raise NotImplementedError(
                     "DCP for FlashMLA sparse is only supported on the "
                     "mixed-batch fp8 path (num_heads < "
@@ -367,9 +437,16 @@ def __init__(
             # Head padding (and the tile-scheduler metadata sized from it) is
             # computed from the local head count, but the kernel runs on the
             # DCP-gathered heads.
-            gathered_num_heads = (
-                self.num_heads * parallel_config.decode_context_parallel_size
-            )
+            if self.use_pcp:
+                gathered_num_heads = (
+                    self.num_heads * parallel_config.tensor_parallel_size
+                    if self.dcp_world_size > self.pcp_world_size
+                    else self.num_heads
+                )
+            else:
+                gathered_num_heads = (
+                    self.num_heads * parallel_config.decode_context_parallel_size
+                )
             gathered_padded_heads = FlashMLASparseImpl._compute_fp8_decode_padded_heads(
                 gathered_num_heads
             )
@@ -385,37 +462,20 @@ def __init__(
 
     def _build_fp8_mixed_decode_prefill(
         self,
-        common_attn_metadata: CommonAttentionMetadata,
-    ) -> "FlashMLASparseMetadata.FP8KernelMetadata":
+    ) -> FlashMLASparseMetadata.FP8KernelMetadata:
         """Build FP8 metadata treating MQA tokens as one batch.
 
         The scheduler initializes lazily from the runtime query shape, which may
         be the full batch or only decodes when prefills use dense MHA. This avoids
         the BF16 prefill kernel's head-padding overhead at high TP.
         """
-        num_tokens = common_attn_metadata.num_actual_tokens
-
-        # Use padded head count since that's what the kernel will see
-        padded_heads = self.fp8_decode_padded_heads
-
-        # Build metadata for all tokens as a single batch
-        scheduler_metadata, _ = get_mla_metadata(
-            cache_seqlens=self.topk_tokens_tensor[:1],  # Single batch
-            num_q_tokens_per_head_k=num_tokens * padded_heads,
-            topk=self.topk_tokens,
-            num_heads_q=padded_heads,
-            num_heads_k=1,
-            is_fp8_kvcache=True,
-        )
-
-        fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
+        scheduler_metadata, _ = get_mla_metadata()
+        return FlashMLASparseMetadata.FP8KernelMetadata(
             scheduler_metadata=scheduler_metadata,
             cache_lens=self.max_model_len_tensor[:1],
             dummy_block_table=self.dummy_block_table[:1],
         )
 
-        return fp8_metadata
-
     def _build_fp8_separate_prefill_decode(
         self,
         common_attn_metadata: CommonAttentionMetadata,
@@ -463,6 +523,40 @@ def _build_fp8_separate_prefill_decode(
             query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
 
             prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:]
+            # Chunk prefill requests to fit within workspace size
+            max_prefill_buffer_size = get_prefill_workspace_size(
+                self.vllm_config.model_config.max_model_len
+            )
+            # Workspace rows per entry, and the prefill rows each entry covers.
+            workspace_rows = prefill_seq_lens_cpu
+            row_bounds = np.arange(num_prefills + 1)
+            prefill_block_table = common_attn_metadata.block_table_tensor[num_decodes:]
+            prefill_seq_lens = common_attn_metadata.seq_lens[num_decodes:]
+            if self.pcp_dcp_kv_gather:
+                # One entry per request, holding this rank's KV shard: the
+                # all-gather of every rank's shard is the whole context.
+                # A dummy batch bypasses the PCP manager: one row per request,
+                # so the row's own extent is the request's.
+                row_req_idx = common_attn_metadata.req_idx
+                if row_req_idx is None:
+                    row_req_idx = np.arange(common_attn_metadata.num_reqs)
+                shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound
+                if shard_rows_cpu is None:
+                    shard_rows_cpu = get_dcp_local_seq_lens(
+                        seq_lens_cpu,
+                        self.dcp_world_size,
+                        0,
+                        self.cp_kv_cache_interleave_size,
+                    )
+                row_bounds, rows_per_rank = gathered_prefill_shards(
+                    row_req_idx[num_decodes:], shard_rows_cpu[num_decodes:].numpy()
+                )
+                workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32))
+                max_prefill_buffer_size //= self.dcp_world_size
+                entry_rows = async_tensor_h2d(row_bounds[:-1], device=self.device)
+                prefill_block_table = prefill_block_table.index_select(0, entry_rows)
+                prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows)
+            num_entries = len(workspace_rows)
 
             # Build prefill_request_id: -1 for decode, request index for
             # prefill. This enables a single
@@ -471,33 +565,28 @@ def _build_fp8_separate_prefill_decode(
                 (num_tokens,), -1, dtype=torch.int32, device=self.device
             )
             # Map prefill tokens to their request IDs (0, 1, 2, ...)
-            for req_idx in range(num_prefills):
+            for req_idx in range(num_entries):
                 # Get query token range for this prefill request
-                global_req_idx = num_decodes + req_idx
-                req_query_start = query_start_loc_cpu[global_req_idx]
-                req_query_end = query_start_loc_cpu[global_req_idx + 1]
+                req_query_start = query_start_loc_cpu[
+                    num_decodes + int(row_bounds[req_idx])
+                ]
+                req_query_end = query_start_loc_cpu[
+                    num_decodes + int(row_bounds[req_idx + 1])
+                ]
                 prefill_request_id[req_query_start:req_query_end] = req_idx
 
             # will be adjusted by chunk loop
             prefill_workspace_starts_cpu = torch.zeros(
-                num_prefills, dtype=torch.int32, pin_memory=True
-            )
-            prefill_workspace_starts_cpu[1:] = torch.cumsum(
-                prefill_seq_lens_cpu[:-1], dim=0
+                num_entries, dtype=torch.int32, pin_memory=True
             )
+            prefill_workspace_starts_cpu[1:] = torch.cumsum(workspace_rows[:-1], dim=0)
             # populated by non-blocking copy after prefill_workspace_starts_cpu is
             # updated by each chunk
             prefill_workspace_starts = torch.empty(
-                num_prefills, dtype=torch.int32, device=self.device
+                num_entries, dtype=torch.int32, device=self.device
             )
 
-            # Chunk prefill requests to fit within workspace size
-            max_prefill_buffer_size = get_prefill_workspace_size(
-                self.vllm_config.model_config.max_model_len
-            )
-            chunk_bounds = split_prefill_chunks(
-                prefill_seq_lens_cpu, max_prefill_buffer_size
-            )
+            chunk_bounds = split_prefill_chunks(workspace_rows, max_prefill_buffer_size)
 
             prefill_chunks = []
             for chunk_start, chunk_end in chunk_bounds:
@@ -510,16 +599,19 @@ def _build_fp8_separate_prefill_decode(
                 offset = prefill_workspace_starts_cpu[chunk_start].item()
                 prefill_workspace_starts_cpu[chunk_start:chunk_end] -= offset
 
-                chunk_tot_seqlen = prefill_seq_lens_cpu[chunk_start:chunk_end].sum()
-                token_start = query_start_loc_cpu[num_decodes + chunk_start].item()
-                token_end = query_start_loc_cpu[num_decodes + chunk_end].item()
+                chunk_tot_seqlen = workspace_rows[chunk_start:chunk_end].sum()
+                token_start = query_start_loc_cpu[
+                    num_decodes + int(row_bounds[chunk_start])
+                ].item()
+                token_end = query_start_loc_cpu[
+                    num_decodes + int(row_bounds[chunk_end])
+                ].item()
                 tokens_slice = slice(token_start, token_end)
 
                 # Create chunk view of gpu tensor
                 chunk_workspace_starts = prefill_workspace_starts[chunk_start:chunk_end]
-                chunk_block_table = common_attn_metadata.block_table_tensor[
-                    num_decodes + chunk_start : num_decodes + chunk_end
-                ]
+                chunk_block_table = prefill_block_table[chunk_start:chunk_end]
+                chunk_seq_lens = prefill_seq_lens[chunk_start:chunk_end]
 
                 prefill_chunks.append(
                     FP8Meta.Prefill.Chunk(
@@ -528,6 +620,7 @@ def _build_fp8_separate_prefill_decode(
                         req_start_idx=chunk_start,
                         workspace_starts=chunk_workspace_starts,
                         chunk_tot_seqlen=chunk_tot_seqlen,
+                        seq_lens=chunk_seq_lens,
                     )
                 )
 
@@ -542,14 +635,19 @@ def _build_fp8_separate_prefill_decode(
             )
 
         if num_decodes > 0:
-            # Use padded head count since that's what the kernel will see
-            scheduler_metadata, _ = get_mla_metadata()
+            if self.pcp_dcp_kv_gather:
+                # Decode rows take the mixed-batch kernel call, which returns
+                # the LSE the DCP merge needs.
+                kernel_meta = self._build_fp8_mixed_decode_prefill()
+            else:
+                # Use padded head count since that's what the kernel will see
+                scheduler_metadata, _ = get_mla_metadata()
 
-            kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata(
-                scheduler_metadata=scheduler_metadata,
-                dummy_block_table=self.dummy_block_table[:active_num_decodes],
-                cache_lens=self.max_model_len_tensor[:active_num_decodes],
-            )
+                kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata(
+                    scheduler_metadata=scheduler_metadata,
+                    dummy_block_table=self.dummy_block_table[:active_num_decodes],
+                    cache_lens=self.max_model_len_tensor[:active_num_decodes],
+                )
             fp8_metadata.decode = FP8Meta.Decode(
                 seq_lens=common_attn_metadata.seq_lens[:active_num_decodes],
                 kernel_metadata=kernel_meta,
@@ -569,9 +667,7 @@ def build(
         metadata.fp8_use_mixed_batch = self.fp8_use_mixed_batch
         if self.use_fp8_kv_cache:
             if self.fp8_use_mixed_batch:
-                metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(
-                    common_attn_metadata
-                )
+                metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill()
             else:
                 metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode(
                     common_attn_metadata, metadata
@@ -583,6 +679,7 @@ def build(
 class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
     can_return_lse_for_decode: bool = True
     supports_dcp: bool = True
+    supports_mtp_with_cp_non_trivial_interleave_size: bool = True
 
     @staticmethod
     def _compute_fp8_decode_padded_heads(num_heads: int) -> int:
@@ -654,23 +751,49 @@ def __init__(
                 "the bf16 sparse path is not supported under DCP."
             )
 
+        self.pcp_dcp_kv_gather = False
+        self.workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [
+            (q_concat_shape, torch.bfloat16)
+        ]
         if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS:
             # Reserve workspace during initialization
             assert vllm_config is not None and vllm_config.model_config is not None
             prefill_workspace_size = get_prefill_workspace_size(
                 vllm_config.model_config.max_model_len
             )
-            self.prefill_workspace_shape = (prefill_workspace_size, head_size)
-            self.q_concat_buffer, self.prefill_bf16_workspace = (
-                current_workspace_manager().get_simultaneous(
-                    (q_concat_shape, torch.bfloat16),
-                    (self.prefill_workspace_shape, torch.bfloat16),
+            parallel_config = vllm_config.parallel_config
+            self.pcp_dcp_kv_gather = (
+                parallel_config.prefill_context_parallel_size > 1
+                and parallel_config.decode_context_parallel_size > 1
+            )
+            shard_rows = prefill_workspace_size
+            if self.pcp_dcp_kv_gather:
+                # PCP+DCP upconverts this rank's KV shard, then all-gathers the
+                # shards into a workspace of the full prefill size.
+                shard_rows //= parallel_config.decode_context_parallel_size
+            self.prefill_workspace_shape = (shard_rows, head_size)
+            self.workspace_specs.append((self.prefill_workspace_shape, torch.bfloat16))
+            if self.pcp_dcp_kv_gather:
+                self.workspace_specs.append(
+                    ((prefill_workspace_size, head_size), torch.bfloat16)
                 )
+            prefill_query_heads = num_heads
+            if self.pcp_dcp_kv_gather and self.dcp_world_size > self.pcp_world_size:
+                prefill_query_heads *= parallel_config.tensor_parallel_size
+            padded_prefill_query_heads = round_up(
+                prefill_query_heads, self.prefill_padding
             )
-        else:
-            (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous(
-                (q_concat_shape, torch.bfloat16),
+            self.workspace_specs.extend(
+                (shape, torch.bfloat16)
+                for shape in (
+                    (max_tokens, padded_prefill_query_heads, head_size),
+                    (max_tokens, padded_prefill_query_heads, self.kv_lora_rank),
+                    (max_tokens, prefill_query_heads, self.kv_lora_rank),
+                )
             )
+        # Reserve capacity without retaining views that prevent old storage
+        # from being released when another layer grows the shared workspace.
+        current_workspace_manager().get_simultaneous(*self.workspace_specs)
 
     def _forward_bf16_kv(
         self,
@@ -679,30 +802,120 @@ def _forward_bf16_kv(
         topk_indices: torch.Tensor,
         attn_metadata: FlashMLASparseMetadata,
         actual_num_heads: int,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        # Convert per-request indices to global slots (decode) or workspace
-        # offsets (prefill). req_id_per_token covers the whole batch; slice it
-        # to the MQA tokens (q may exclude prefill tokens routed to dense MHA).
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        index_group = self.index_group
+        if isinstance(index_group, HiSparseMLAIndexGroup):
+            cache = index_group.cache(self.index_group_index)
+        else:
+            cache = None
+        block_table = attn_metadata.block_table
+        # req_id_per_token covers the whole batch; slice it to the MQA tokens
+        # (q may exclude prefill tokens routed to dense MHA).
+        req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]]
+        decode_out: torch.Tensor | None = None
+        if cache is not None:
+            assert isinstance(index_group, HiSparseMLAIndexGroup)
+            num_decode_tokens = attn_metadata.num_decode_tokens
+            if num_decode_tokens > 0:
+                decode_topk, decode_lengths = (
+                    index_group.convert_decode_logical_to_physical_topk(
+                        self.index_group_index,
+                        topk_indices[:num_decode_tokens],
+                        attn_metadata,
+                        return_valid_counts=True,
+                    )
+                )
+                decode_out, _ = self._bf16_flash_mla_kernel(
+                    q[:num_decode_tokens],
+                    index_group.physical_kv_cache(self.index_group_index),
+                    decode_topk,
+                    decode_lengths,
+                    actual_num_heads,
+                )
+                if num_decode_tokens == q.shape[0]:
+                    return decode_out, None
+                q = q[num_decode_tokens:]
+                topk_indices = topk_indices[num_decode_tokens:]
+            kv_c_and_k_pe_cache, block_table, req_id_per_token = (
+                index_group.stage_prefill_rows(
+                    self.index_group_index, kv_c_and_k_pe_cache, attn_metadata
+                )
+            )
+        # Convert per-request indices to global slots (decode) or workspace offsets.
         kv_rows, block_stride_rows = flat_kv_row_view(
             kv_c_and_k_pe_cache, attn_metadata.block_size
         )
-        topk_indices, topk_length = triton_convert_req_index_to_global_index(
-            attn_metadata.req_id_per_token[: topk_indices.shape[0]],
-            attn_metadata.block_table,
-            topk_indices,
-            BLOCK_SIZE=attn_metadata.block_size,
-            BLOCK_STRIDE_ROWS=block_stride_rows,
-            NUM_TOPK_TOKENS=topk_indices.shape[1],
-            return_valid_counts=True,
+        decode_only = (
+            attn_metadata.num_decode_tokens
+            == attn_metadata.num_actual_tokens
+            == topk_indices.shape[0]
         )
+        uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup)
+        if not uses_host_cache and decode_only:
+            topk_indices, topk_length = self._convert_logical_to_physical_topk(
+                topk_indices,
+                attn_metadata,
+                block_stride_rows=block_stride_rows,
+                return_valid_counts=True,
+            )
+        else:
+            topk_indices, topk_length = triton_convert_req_index_to_global_index(
+                req_id_per_token,
+                block_table,
+                topk_indices,
+                BLOCK_SIZE=attn_metadata.block_size,
+                BLOCK_STRIDE_ROWS=block_stride_rows,
+                NUM_TOPK_TOKENS=topk_indices.shape[1],
+                return_valid_counts=True,
+            )
 
-        return self._bf16_flash_mla_kernel(
+        attn_out, lse = self._bf16_flash_mla_kernel(
             q,
             kv_rows,
             topk_indices,
             topk_length,
             actual_num_heads,
         )
+        if decode_out is None:
+            return attn_out, lse
+        return torch.cat([decode_out, attn_out], dim=0), None
+
+    def _gather_prefill_chunk(
+        self,
+        chunk: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill.Chunk",
+        shard: torch.Tensor,
+        topk_indices: torch.Tensor,
+        attn_metadata: FlashMLASparseMetadata,
+        prefill_meta: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill",
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """All-gather this rank's upconverted KV shard so the chunk's rows attend
+        the whole context, and map their top-k onto the rank-major result."""
+        shard_rows = int(chunk.chunk_tot_seqlen)
+        _, _, gathered_kv_workspace, *_ = current_workspace_manager().get_simultaneous(
+            *self.workspace_specs
+        )
+        gathered_kv = gathered_kv_workspace[: self.dcp_world_size * shard_rows]
+        dist.all_gather_into_tensor(
+            gathered_kv, shard, group=get_dcp_group().device_group
+        )
+
+        tokens = chunk.tokens_slice
+        topk_indices, topk_length = triton_convert_req_index_to_global_index(
+            attn_metadata.req_id_per_token[tokens],
+            attn_metadata.block_table,
+            topk_indices,
+            BLOCK_SIZE=attn_metadata.block_size,
+            NUM_TOPK_TOKENS=topk_indices.shape[1],
+            HAS_PREFILL_WORKSPACE=True,
+            prefill_workspace_request_ids=prefill_meta.request_ids[tokens],
+            prefill_workspace_starts=prefill_meta.workspace_starts,
+            prefill_workspace_rank_stride=shard_rows,
+            dcp_size=self.dcp_world_size,
+            dcp_rank=self.dcp_rank,
+            cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size,
+            return_valid_counts=True,
+        )
+        return gathered_kv, topk_indices, topk_length
 
     def _forward_fp8_kv_separate_prefill_decode(
         self,
@@ -710,7 +923,8 @@ def _forward_fp8_kv_separate_prefill_decode(
         kv_c_and_k_pe_cache: torch.Tensor,
         topk_indices: torch.Tensor,
         attn_metadata: FlashMLASparseMetadata,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """The lse covers the decode rows and is only returned under PCP+DCP."""
         fp8_metadata = attn_metadata.fp8_extra_metadata
         assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode)
         num_decodes = fp8_metadata.num_decodes
@@ -721,6 +935,31 @@ def _forward_fp8_kv_separate_prefill_decode(
             "FP8 sparse MLA expects either the decode subset or the full batch"
         )
 
+        decode_topk: torch.Tensor | None = None
+        index_group = self.index_group
+        uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup)
+        if uses_host_cache and num_decode_tokens > 0:
+            decode_topk = topk_indices[:num_decode_tokens]
+
+        prefill_ready = None
+        if num_prefill_tokens > 0 and uses_host_cache:
+            assert fp8_metadata.prefill is not None
+            first_chunk = fp8_metadata.prefill.chunks[0]
+            assert isinstance(index_group, HiSparseMLAIndexGroup)
+            _, prefill_bf16_workspace, *_ = (
+                current_workspace_manager().get_simultaneous(*self.workspace_specs)
+            )
+            prefill_ready = index_group.gather_fp8_prefill(
+                self.index_group_index,
+                kv_c_and_k_pe_cache,
+                prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen],
+                first_chunk.block_table,
+                first_chunk.workspace_starts,
+                len(first_chunk.block_table),
+                attn_metadata,
+                first_chunk.req_start_idx,
+            )
+
         prefill_request_ids = None
         prefill_workspace_starts = None
         has_prefill_workspace = False
@@ -736,17 +975,31 @@ def _forward_fp8_kv_separate_prefill_decode(
         # For BF16 cache: always use global cache slots (no workspace)
         # prefill_workspace_starts has been adjusted in-place per chunk so
         # prefill indices automatically come out chunk-local
-        topk_indices, topk_length = triton_convert_req_index_to_global_index(
-            attn_metadata.req_id_per_token[: topk_indices.shape[0]],
-            attn_metadata.block_table,
-            topk_indices,
-            BLOCK_SIZE=attn_metadata.block_size,
-            NUM_TOPK_TOKENS=topk_indices.shape[1],
-            HAS_PREFILL_WORKSPACE=has_prefill_workspace,
-            prefill_workspace_request_ids=prefill_request_ids,
-            prefill_workspace_starts=prefill_workspace_starts,
-            return_valid_counts=True,
-        )
+        # Under PCP+DCP the decode rows are converted by the mixed-batch call
+        # and the prefill rows per chunk, against the gathered workspace.
+        lse: torch.Tensor | None = None
+        topk_length: torch.Tensor | None = None
+        if self.pcp_dcp_kv_gather:
+            pass
+        elif num_prefill_tokens == 0 and not uses_host_cache:
+            topk_indices, topk_length = self._convert_logical_to_physical_topk(
+                topk_indices,
+                attn_metadata,
+                block_stride_rows=None,
+                return_valid_counts=True,
+            )
+        elif num_prefill_tokens > 0:
+            topk_indices, topk_length = triton_convert_req_index_to_global_index(
+                attn_metadata.req_id_per_token[: topk_indices.shape[0]],
+                attn_metadata.block_table,
+                topk_indices,
+                BLOCK_SIZE=attn_metadata.block_size,
+                NUM_TOPK_TOKENS=topk_indices.shape[1],
+                HAS_PREFILL_WORKSPACE=has_prefill_workspace,
+                prefill_workspace_request_ids=prefill_request_ids,
+                prefill_workspace_starts=prefill_workspace_starts,
+                return_valid_counts=True,
+            )
 
         fp8_metadata = attn_metadata.fp8_extra_metadata
         assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode)
@@ -755,6 +1008,27 @@ def _fp8_decode(
             q: torch.Tensor,
             topk_indices: torch.Tensor,
         ) -> torch.Tensor:
+            assert fp8_metadata.decode is not None
+            if self.pcp_dcp_kv_gather:
+                # Attend this rank's KV shard; the LSE feeds the DCP merge.
+                nonlocal lse
+                out, lse = self._forward_fp8_kv_mixed_batch(
+                    q,
+                    kv_c_and_k_pe_cache,
+                    topk_indices,
+                    attn_metadata,
+                    kernel_metadata=fp8_metadata.decode.kernel_metadata,
+                )
+                return out
+            if uses_host_cache:
+                return self._host_backed_fp8_decode(
+                    q,
+                    topk_indices,
+                    attn_metadata,
+                    fp8_metadata.decode.kernel_metadata,
+                    num_decodes,
+                    fp8_metadata.decode.decode_query_len,
+                )
             # Reshape q: (num_decode_tokens, num_heads, head_dim)
             #         -> (num_decodes, seq_len, num_heads, head_dim)
             q = reshape_query_for_spec_decode(q, num_decodes)
@@ -762,7 +1036,6 @@ def _fp8_decode(
             # Reshape topk_indices: (num_decode_tokens, topk)
             #                    -> (num_decodes, seq_len, topk)
             topk_indices = topk_indices.view(num_decodes, seq_len, -1)
-            assert fp8_metadata.decode is not None
             attn_out, _ = self._fp8_flash_mla_kernel(
                 q=q,
                 kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,
@@ -776,25 +1049,42 @@ def _fp8_decode(
         # Pure decode: direct call without allocation
         if num_decode_tokens > 0 and num_prefill_tokens == 0:
             assert fp8_metadata.decode is not None
-            attn_out = _fp8_decode(q, topk_indices)
+            attn_out = _fp8_decode(
+                q, decode_topk if decode_topk is not None else topk_indices
+            )
         else:
-            # Mixed or pure prefill: allocate output tensor
-            attn_out = q.new_empty(
-                (num_mqa_tokens, self.num_heads, self.kv_lora_rank),
-                dtype=q.dtype,
-                device=q.device,
+            _, prefill_bf16_workspace, *_, padded_out, compact_out = (
+                current_workspace_manager().get_simultaneous(*self.workspace_specs)
             )
+            attn_out = compact_out[:num_mqa_tokens]
 
             if num_decode_tokens > 0:
                 attn_out[:num_decode_tokens] = _fp8_decode(
                     q[:num_decode_tokens],
-                    topk_indices[:num_decode_tokens],
+                    decode_topk
+                    if decode_topk is not None
+                    else topk_indices[:num_decode_tokens],
                 )
 
             assert fp8_metadata.prefill is not None
-            for chunk in fp8_metadata.prefill.chunks:
-                chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen]
-                if self.kv_cache_dtype == "fp8_ds_mla":
+            for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks):
+                chunk_workspace = prefill_bf16_workspace[: chunk.chunk_tot_seqlen]
+                if uses_host_cache and chunk_index > 0:
+                    assert isinstance(index_group, HiSparseMLAIndexGroup)
+                    prefill_ready = index_group.gather_fp8_prefill(
+                        self.index_group_index,
+                        kv_c_and_k_pe_cache,
+                        chunk_workspace,
+                        chunk.block_table,
+                        chunk.workspace_starts,
+                        len(chunk.block_table),
+                        attn_metadata,
+                        chunk.req_start_idx,
+                    )
+                if uses_host_cache:
+                    assert prefill_ready is not None
+                    current_stream().wait_event(prefill_ready)
+                elif self.kv_cache_dtype == "fp8_ds_mla":
                     ops.cp_gather_and_upconvert_fp8_kv_cache(
                         kv_c_and_k_pe_cache,
                         chunk_workspace,
@@ -812,17 +1102,35 @@ def _fp8_decode(
                     )
 
                 chunk_q = q[chunk.tokens_slice]
-                chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]
-                chunk_topk_length = topk_length[chunk.tokens_slice]
+                if self.pcp_dcp_kv_gather:
+                    chunk_workspace, chunk_topk_indices_workspace, chunk_topk_length = (
+                        self._gather_prefill_chunk(
+                            chunk,
+                            chunk_workspace,
+                            topk_indices[chunk.tokens_slice],
+                            attn_metadata,
+                            fp8_metadata.prefill,
+                        )
+                    )
+                else:
+                    chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]
+                    assert topk_length is not None
+                    chunk_topk_length = topk_length[chunk.tokens_slice]
 
-                attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel(
+                chunk_out, chunk_lse = self._bf16_flash_mla_kernel(
                     chunk_q,
                     chunk_workspace,
                     chunk_topk_indices_workspace,
                     chunk_topk_length,
+                    out=padded_out[chunk.tokens_slice],
                 )
+                attn_out[chunk.tokens_slice].copy_(chunk_out)
+                del chunk_out, chunk_lse
 
-        return attn_out
+        if self.pcp_dcp_kv_gather and lse is None:
+            # No decode rows: the DCP merge still expects an LSE, of no rows.
+            lse = q.new_empty((0, q.shape[1]), dtype=torch.float32)
+        return attn_out, lse
 
     def _forward_fp8_kv_mixed_batch(
         self,
@@ -830,6 +1138,7 @@ def _forward_fp8_kv_mixed_batch(
         kv_c_and_k_pe_cache: torch.Tensor,
         topk_indices: torch.Tensor,
         attn_metadata: FlashMLASparseMetadata,
+        kernel_metadata: "FlashMLASparseMetadata.FP8KernelMetadata | None" = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Mixed batch FP8 forward path that treats all tokens as one batch.
 
@@ -839,6 +1148,13 @@ def _forward_fp8_kv_mixed_batch(
 
         The lse is only returned when DCP needs it, otherwise None.
         """
+        fp8_metadata = kernel_metadata or attn_metadata.fp8_extra_metadata
+        assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8KernelMetadata)
+
+        block_table = attn_metadata.block_table
+        # req_id_per_token covers the whole batch; slice it to the MQA tokens
+        # (q may exclude prefill tokens routed to dense MHA).
+        req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]]
         if self.dcp_world_size > 1:
             # The indexer emits global token ids; keep this rank's shard and
             # convert to local slots. compact_valid_to_front=False keeps the
@@ -846,8 +1162,8 @@ def _forward_fp8_kv_mixed_batch(
             # empty-row neutralization below relies on. req_id is sliced to
             # topk_indices rows (the converter grids from req_id).
             topk_indices = triton_filter_and_convert_dcp_index(
-                attn_metadata.req_id_per_token[: topk_indices.shape[0]],
-                attn_metadata.block_table,
+                req_id_per_token,
+                block_table,
                 topk_indices,
                 dcp_size=self.dcp_world_size,
                 dcp_rank=self.dcp_rank,
@@ -859,19 +1175,26 @@ def _forward_fp8_kv_mixed_batch(
         else:
             # Convert per-request indices to global slots (decode) or workspace
             # offsets (prefill).
-            topk_indices = triton_convert_req_index_to_global_index(
-                attn_metadata.req_id_per_token[: topk_indices.shape[0]],
-                attn_metadata.block_table,
-                topk_indices,
-                BLOCK_SIZE=attn_metadata.block_size,
-                NUM_TOPK_TOKENS=topk_indices.shape[1],
+            decode_only = (
+                attn_metadata.num_decode_tokens
+                == attn_metadata.num_actual_tokens
+                == topk_indices.shape[0]
             )
-
-        assert attn_metadata.fp8_extra_metadata is not None
-        assert isinstance(
-            attn_metadata.fp8_extra_metadata, FlashMLASparseMetadata.FP8KernelMetadata
-        )
-        fp8_metadata = attn_metadata.fp8_extra_metadata
+            if decode_only:
+                topk_indices = self._convert_logical_to_physical_topk(
+                    topk_indices,
+                    attn_metadata,
+                    block_stride_rows=None,
+                    return_valid_counts=False,
+                )
+            else:
+                topk_indices = triton_convert_req_index_to_global_index(
+                    req_id_per_token,
+                    block_table,
+                    topk_indices,
+                    BLOCK_SIZE=attn_metadata.block_size,
+                    NUM_TOPK_TOKENS=topk_indices.shape[1],
+                )
 
         _attn_out, _lse = self._fp8_flash_mla_kernel(
             q=q.unsqueeze(0),  # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D)
@@ -937,6 +1260,37 @@ def _fp8_flash_mla_kernel(
 
         return out, lse
 
+    def _host_backed_fp8_decode(
+        self,
+        q: torch.Tensor,
+        topk_indices: torch.Tensor,
+        attn_metadata: FlashMLASparseMetadata,
+        kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata,
+        num_decodes: int,
+        decode_query_len: int,
+    ) -> torch.Tensor:
+        assert isinstance(self.index_group, HiSparseMLAIndexGroup)
+        physical_topk = self.index_group.convert_decode_logical_to_physical_topk(
+            self.index_group_index,
+            topk_indices,
+            attn_metadata,
+            return_valid_counts=False,
+            num_decodes=num_decodes,
+            decode_query_len=decode_query_len,
+        )
+        assert isinstance(physical_topk, torch.Tensor)
+        q = reshape_query_for_spec_decode(q, num_decodes)
+        physical_topk = physical_topk.view(num_decodes, q.shape[1], -1)
+        output, _ = self._fp8_flash_mla_kernel(
+            q=q,
+            kv_c_and_k_pe_cache=self.index_group.physical_kv_cache(
+                self.index_group_index
+            ),
+            topk_indices=physical_topk,
+            kernel_metadata=kernel_metadata,
+        )
+        return reshape_attn_output_for_spec_decode(output)
+
     def _bf16_flash_mla_kernel(
         self,
         q: torch.Tensor,
@@ -944,6 +1298,7 @@ def _bf16_flash_mla_kernel(
         topk_indices: torch.Tensor,
         topk_length: torch.Tensor | None = None,
         actual_num_heads: int | None = None,
+        out: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         num_tokens = q.shape[0]
         kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
@@ -965,7 +1320,14 @@ def _bf16_flash_mla_kernel(
                 f"Padding num_heads from {actual_num_heads} to "
                 f"{padded_num_heads} for BF16 sparse prefill kernel"
             )
-            q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2]))
+            if out is None:
+                q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2]))
+            else:
+                *_, q_padded, _, _ = current_workspace_manager().get_simultaneous(
+                    *self.workspace_specs
+                )
+                q_padded = q_padded[:num_tokens]
+                q_padded[:, actual_num_heads:].zero_()
             q_padded[:, :actual_num_heads, :] = q
             q = q_padded
 
@@ -976,6 +1338,7 @@ def _bf16_flash_mla_kernel(
             topk_indices,
             self.softmax_scale,
             topk_length=topk_length,
+            out=out,
         )
 
         output = output[:, :actual_num_heads, :]
@@ -996,8 +1359,16 @@ def forward_mqa(
         actual_num_heads = self.num_heads
         if isinstance(q, tuple):
             ql_nope, q_pe = q
-            q = self.q_concat_buffer[: ql_nope.shape[0]]
-            ops.concat_mla_q(ql_nope, q_pe, q)
+            q_concat_buffer, *_ = current_workspace_manager().get_simultaneous(
+                *self.workspace_specs
+            )
+            q = q_concat_buffer[: ql_nope.shape[0]]
+            if q_pe.size(-1) == 0:
+                # NoPE (GLM5Next): concat_mla_q requires rope_dim == 64,
+                # copy directly into the head-padded buffer instead.
+                q[:, : ql_nope.shape[1]].copy_(ql_nope)
+            else:
+                ops.concat_mla_q(ql_nope, q_pe, q)
         else:
             actual_num_heads = q.shape[1]
 
@@ -1026,7 +1397,7 @@ def forward_mqa(
                 q, kv_c_and_k_pe_cache, topk_indices, attn_metadata
             )
         else:
-            attn_out = self._forward_fp8_kv_separate_prefill_decode(
+            attn_out, lse = self._forward_fp8_kv_separate_prefill_decode(
                 q, kv_c_and_k_pe_cache, topk_indices, attn_metadata
             )
 
diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py
new file mode 100644
index 000000000000..1b8455a2328a
--- /dev/null
+++ b/vllm/v1/attention/backends/mla/index_group.py
@@ -0,0 +1,512 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.config import VllmConfig
+from vllm.forward_context import in_piecewise_cudagraph
+from vllm.utils.torch_utils import current_stream
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_convert_req_index_to_global_index,
+)
+from vllm.v1.hisparse.runtime import (
+    FP8_DS_MLA_ROW_BYTES,
+    HiSparseCacheHandle,
+    HiSparsePrefillStagingPlan,
+    create_hisparse_cache_handle,
+)
+
+
+def _create_side_stream(device: torch.device) -> torch.Stream:
+    return torch.Stream(device=device)
+
+
+def _create_event() -> torch.Event:
+    return torch.Event()
+
+
+@dataclass
+class SparseMLAIndexGroup:
+    """Layers that consume one sparse-indexer result."""
+
+    logical_topk_indices: torch.Tensor
+    physical_topk_indices: torch.Tensor
+    valid_topk_counts: torch.Tensor
+    request_ids: torch.Tensor
+    side_stream: torch.Stream
+    logical_topk_ready: torch.Event
+    physical_topk_ready: torch.Event
+    has_indexer: bool
+    num_layers: int = 0
+
+    def register_layer(
+        self,
+        vllm_config: VllmConfig | None = None,
+        *,
+        head_size: int | None = None,
+        kv_cache_dtype: str | None = None,
+    ) -> int:
+        layer_index = self.num_layers
+        self.num_layers += 1
+        return layer_index
+
+    def set_logical_topk_ready(self, layer_index: int) -> None:
+        if layer_index != 0 or not self.has_indexer:
+            return
+        # In piecewise cudagraph mode the consumer converts on the side stream
+        # from an eager break; an event recorded inside a captured segment is
+        # graph-local and any eager wait on it raises cudaErrorInvalidValue.
+        # The convert side falls back to wait_stream there instead (and never
+        # touches this event), so skip the record as well.
+        if in_piecewise_cudagraph():
+            return
+        self.logical_topk_ready.record(current_stream())
+
+    def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None:
+        pass
+
+    def gather_fp8_prefill(
+        self,
+        layer_index: int,
+        source_cache: torch.Tensor,
+        dst: torch.Tensor,
+        block_table: torch.Tensor,
+        workspace_starts: torch.Tensor,
+        batch_size: int,
+        attn_metadata: Any,
+        request_start: int,
+    ) -> torch.Event | None:
+        ops.cp_gather_and_upconvert_fp8_kv_cache(
+            source_cache,
+            dst,
+            block_table,
+            workspace_starts,
+            batch_size,
+        )
+        return None
+
+    def convert_logical_to_physical_topk(
+        self,
+        layer_index: int,
+        logical_topk_indices: torch.Tensor,
+        attn_metadata: Any,
+        *,
+        block_stride_rows: int | None,
+        return_valid_counts: bool,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """Convert one index result once for all layers in this group."""
+        return self._convert_once(
+            layer_index,
+            logical_topk_indices,
+            attn_metadata.req_id_per_token[: logical_topk_indices.shape[0]],
+            attn_metadata.block_table,
+            attn_metadata.block_size,
+            block_stride_rows=block_stride_rows,
+            return_valid_counts=return_valid_counts,
+        )
+
+    def _convert_once(
+        self,
+        layer_index: int,
+        logical_topk_indices: torch.Tensor,
+        req_id_per_token: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        *,
+        block_stride_rows: int | None,
+        return_valid_counts: bool,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        num_tokens = logical_topk_indices.shape[0]
+        if num_tokens > self.physical_topk_indices.shape[0]:
+            return triton_convert_req_index_to_global_index(
+                req_id_per_token,
+                block_table,
+                logical_topk_indices,
+                BLOCK_SIZE=block_size,
+                BLOCK_STRIDE_ROWS=block_stride_rows,
+                NUM_TOPK_TOKENS=logical_topk_indices.shape[1],
+                return_valid_counts=return_valid_counts,
+            )
+
+        physical_topk_indices = self.physical_topk_indices[:num_tokens]
+        valid_topk_counts = self.valid_topk_counts[:num_tokens]
+        compute_stream = current_stream()
+        if layer_index == 0:
+            if self.has_indexer and not in_piecewise_cudagraph():
+                self.side_stream.wait_event(self.logical_topk_ready)
+            else:
+                self.side_stream.wait_stream(compute_stream)
+            with self.side_stream:
+                triton_convert_req_index_to_global_index(
+                    req_id_per_token,
+                    block_table,
+                    logical_topk_indices,
+                    BLOCK_SIZE=block_size,
+                    BLOCK_STRIDE_ROWS=block_stride_rows,
+                    NUM_TOPK_TOKENS=logical_topk_indices.shape[1],
+                    return_valid_counts=return_valid_counts,
+                    out=physical_topk_indices,
+                    valid_counts_out=(
+                        valid_topk_counts if return_valid_counts else None
+                    ),
+                )
+                self.physical_topk_ready.record(self.side_stream)
+        compute_stream.wait_event(self.physical_topk_ready)
+
+        if return_valid_counts:
+            return physical_topk_indices, valid_topk_counts
+        return physical_topk_indices
+
+
+@dataclass
+class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
+    """Sparse index group with per-layer host-backed KV caches."""
+
+    caches: list[HiSparseCacheHandle] = field(default_factory=list)
+    hisparse_group: Any | None = None
+    prefill_stream: torch.Stream | None = None
+    prefill_ready_events: list[torch.Event] = field(default_factory=list)
+
+    def register_layer(
+        self,
+        vllm_config: VllmConfig | None = None,
+        *,
+        head_size: int | None = None,
+        kv_cache_dtype: str | None = None,
+    ) -> int:
+        assert vllm_config is not None
+        assert head_size is not None
+        assert kv_cache_dtype is not None
+        layer_index = super().register_layer()
+        if self.prefill_stream is None:
+            self.prefill_stream = _create_side_stream(self.logical_topk_indices.device)
+        if kv_cache_dtype == "fp8_ds_mla":
+            row_width = FP8_DS_MLA_ROW_BYTES
+            kv_dtype = torch.uint8
+        else:
+            from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
+
+            row_width = head_size
+            kv_dtype = kv_cache_dtype_str_to_dtype(
+                kv_cache_dtype, vllm_config.model_config
+            )
+        cache = create_hisparse_cache_handle(
+            vllm_config,
+            self.logical_topk_indices.shape[1],
+            is_index_group_leader=layer_index == 0,
+            row_width=row_width,
+            kv_dtype=kv_dtype,
+            index_group=self,
+        )
+        assert cache is not None
+        self.caches.append(cache)
+        cache.index_group_caches = self.caches
+        self.prefill_ready_events.append(_create_event())
+        return layer_index
+
+    def cache(self, layer_index: int) -> HiSparseCacheHandle:
+        return self.caches[layer_index]
+
+    def physical_kv_cache(self, layer_index: int) -> torch.Tensor:
+        cache = self.cache(layer_index)
+        return cache.runtime.hot.attention_cache
+
+    def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None:
+        if layer_index == 0:
+            self.cache(0).prepare_group_for_batch(attn_metadata)
+
+    def convert_logical_to_physical_topk(
+        self,
+        layer_index: int,
+        logical_topk_indices: torch.Tensor,
+        attn_metadata: Any,
+        *,
+        block_stride_rows: int | None,
+        return_valid_counts: bool,
+        req_id_per_token: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        cache = self.cache(layer_index)
+        num_tokens = logical_topk_indices.shape[0]
+        if req_id_per_token is None:
+            req_id_per_token = attn_metadata.req_id_per_token[:num_tokens]
+        if num_tokens > self.physical_topk_indices.shape[0]:
+            # Prefill-sized batches do not fit the decode residency workspace.
+            # Non-resident prefills are staged before reaching this path.
+            assert cache.all_context_pages_resident
+            leader = self.cache(0)
+            assert leader.view is not None and leader.block_table is not None
+            return self._convert_once(
+                layer_index,
+                logical_topk_indices,
+                req_id_per_token,
+                leader.block_table,
+                leader.view.block_size,
+                block_stride_rows=leader.view.attention_block_stride,
+                return_valid_counts=return_valid_counts,
+            )
+        source_block_table = cache.source_block_table
+        assert source_block_table is not None
+        return cache.swap_in(
+            req_id_per_token,
+            block_table=source_block_table,
+            logical_topk_indices=logical_topk_indices,
+            block_size=attn_metadata.block_size,
+            return_valid_counts=return_valid_counts,
+        )
+
+    def convert_decode_logical_to_physical_topk(
+        self,
+        layer_index: int,
+        logical_topk_indices: torch.Tensor,
+        attn_metadata: Any,
+        *,
+        return_valid_counts: bool,
+        num_decodes: int | None = None,
+        decode_query_len: int | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        num_tokens = logical_topk_indices.shape[0]
+        if num_decodes is None:
+            num_decodes = attn_metadata.num_decodes
+        if decode_query_len is None:
+            decode_query_len = attn_metadata.decode_max_query_len
+        if decode_query_len == 1:
+            return self.convert_logical_to_physical_topk(
+                layer_index,
+                logical_topk_indices,
+                attn_metadata,
+                block_stride_rows=None,
+                return_valid_counts=return_valid_counts,
+                req_id_per_token=self.request_ids[:num_decodes],
+            )
+
+        if num_tokens != num_decodes * decode_query_len:
+            query_start_loc = attn_metadata.query_start_loc[: num_decodes + 1]
+            physical_topk_indices = self.physical_topk_indices[: num_tokens + 1]
+            valid_topk_counts = self.valid_topk_counts[: num_tokens + 1]
+            request_ids = self.request_ids[:num_decodes]
+            cache = self.cache(layer_index)
+            source_block_table = cache.source_block_table
+            assert source_block_table is not None
+            for step in range(decode_query_len):
+                token_indices = query_start_loc[:-1] + step
+                active = (token_indices < query_start_loc[1:]) & (
+                    token_indices < num_tokens
+                )
+                output_indices = torch.where(
+                    active, token_indices, torch.full_like(token_indices, num_tokens)
+                ).long()
+                token_indices = token_indices.clamp(0, num_tokens - 1).long()
+                step_topk = logical_topk_indices.index_select(
+                    0, token_indices
+                ).masked_fill(~active.unsqueeze(1), -1)
+                step_result = cache.swap_in(
+                    request_ids,
+                    block_table=source_block_table,
+                    logical_topk_indices=step_topk,
+                    block_size=attn_metadata.block_size,
+                    return_valid_counts=return_valid_counts,
+                )
+                if layer_index == 0:
+                    if return_valid_counts:
+                        step_indices, step_counts = step_result
+                        valid_topk_counts.index_copy_(0, output_indices, step_counts)
+                    else:
+                        step_indices = step_result
+                    physical_topk_indices.index_copy_(0, output_indices, step_indices)
+            physical_topk_indices = physical_topk_indices[:num_tokens]
+            if return_valid_counts:
+                return physical_topk_indices, valid_topk_counts[:num_tokens]
+            return physical_topk_indices
+
+        assert num_tokens == num_decodes * decode_query_len
+        logical_topk_by_request = logical_topk_indices.view(
+            num_decodes, decode_query_len, -1
+        )
+        physical_topk_by_request = self.physical_topk_indices[:num_tokens].view(
+            num_decodes, decode_query_len, -1
+        )
+        valid_topk_by_request = self.valid_topk_counts[:num_tokens].view(
+            num_decodes, decode_query_len
+        )
+        request_ids = self.request_ids[:num_decodes]
+        for step in range(decode_query_len):
+            cache = self.cache(layer_index)
+            source_block_table = cache.source_block_table
+            assert source_block_table is not None
+            cache.swap_in(
+                request_ids,
+                block_table=source_block_table,
+                logical_topk_indices=logical_topk_by_request[:, step],
+                block_size=attn_metadata.block_size,
+                return_valid_counts=return_valid_counts,
+                attention_indices_out=(
+                    physical_topk_by_request[:, step] if layer_index == 0 else None
+                ),
+                valid_counts_out=(
+                    valid_topk_by_request[:, step]
+                    if layer_index == 0 and return_valid_counts
+                    else None
+                ),
+            )
+        physical_topk_indices = physical_topk_by_request.view(num_tokens, -1)
+        if return_valid_counts:
+            return physical_topk_indices, valid_topk_by_request.view(num_tokens)
+        return physical_topk_indices
+
+    def stage_prefill_rows(
+        self, layer_index: int, kv_cache: torch.Tensor, attn_metadata: Any
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        cache = self.cache(layer_index)
+        prefill = attn_metadata.prefill
+        staging_plan = prefill.host_staging_plan if prefill is not None else None
+        assert staging_plan is not None
+        resident_cache = None
+        if cache.view is not None and cache.block_table is not None:
+            staging_plan.ensure_gpu_sources(
+                cache.block_table[attn_metadata.num_decodes :],
+                cache.view.block_size,
+            )
+            resident_cache = cache.view.cache
+        staged_cache = cache.runtime.gather_prefill_cache(
+            kv_cache,
+            staging_plan,
+            resident_cache=resident_cache,
+        )
+        req_ids = attn_metadata.req_id_per_token[attn_metadata.num_decode_tokens :]
+        if attn_metadata.num_decodes > 0:
+            req_ids = req_ids - attn_metadata.num_decodes
+        return staged_cache, staging_plan.block_table, req_ids
+
+    def _prefill_gather_plan(
+        self, layer_index: int, attn_metadata: Any
+    ) -> HiSparsePrefillStagingPlan:
+        cache = self.cache(layer_index)
+        prefill = attn_metadata.prefill
+        plan = prefill.host_staging_plan if prefill is not None else None
+        assert plan is not None
+        assert cache.view is not None and cache.block_table is not None
+        plan.ensure_gpu_sources(
+            cache.block_table[attn_metadata.num_decodes :],
+            cache.view.block_size,
+        )
+        assert plan.gpu_row_ids is not None
+        return plan
+
+    def gather_fp8_prefill(
+        self,
+        layer_index: int,
+        source_cache: torch.Tensor,
+        dst: torch.Tensor,
+        block_table: torch.Tensor,
+        workspace_starts: torch.Tensor,
+        batch_size: int,
+        attn_metadata: Any,
+        request_start: int,
+    ) -> torch.Event | None:
+        cache = self.cache(layer_index)
+        plan = self._prefill_gather_plan(layer_index, attn_metadata)
+        assert cache.view is not None and plan.gpu_row_ids is not None
+        row_width = source_cache.shape[-1]
+        prefill_stream = self.prefill_stream
+        assert prefill_stream is not None
+        block_table = plan.block_table[request_start : request_start + batch_size]
+        prefill_stream.wait_stream(current_stream())
+        with prefill_stream:
+            ops.cp_gather_and_upconvert_fp8_kv_cache(
+                cache.view.cache,
+                dst,
+                block_table,
+                workspace_starts,
+                batch_size,
+                host_cache=source_cache.view(-1, row_width),
+                host_row_ids=plan.row_ids,
+                device_row_ids=plan.gpu_row_ids,
+            )
+            ready = self.prefill_ready_events[layer_index]
+            ready.record(prefill_stream)
+        return ready
+
+
+class SparseMLAIndexGroupBuilder:
+    """Assign sparse MLA layers to their index-producing layer."""
+
+    def __init__(
+        self, logical_topk_indices: torch.Tensor, max_decode_rows: int | None = None
+    ) -> None:
+        self.logical_topk_indices = logical_topk_indices
+        self.max_decode_rows = (
+            logical_topk_indices.shape[0]
+            if max_decode_rows is None
+            else max_decode_rows
+        )
+        self.current_group: SparseMLAIndexGroup | None = None
+
+    def register_layer(
+        self,
+        is_index_producing_layer: bool,
+        vllm_config: VllmConfig | None = None,
+        *,
+        head_size: int | None = None,
+        kv_cache_dtype: str | None = None,
+    ) -> tuple[SparseMLAIndexGroup, int]:
+        if is_index_producing_layer or self.current_group is None:
+            group_cls = (
+                HiSparseMLAIndexGroup
+                if vllm_config is not None
+                and vllm_config.attention_config.hisparse_config is not None
+                else SparseMLAIndexGroup
+            )
+            workspace_rows = self.max_decode_rows
+            physical_topk_indices = torch.empty(
+                (workspace_rows + 1, self.logical_topk_indices.shape[1]),
+                dtype=self.logical_topk_indices.dtype,
+                device=self.logical_topk_indices.device,
+            )
+            self.current_group = group_cls(
+                logical_topk_indices=self.logical_topk_indices,
+                physical_topk_indices=physical_topk_indices,
+                valid_topk_counts=torch.empty(
+                    workspace_rows + 1,
+                    dtype=torch.int32,
+                    device=self.logical_topk_indices.device,
+                ),
+                request_ids=torch.arange(
+                    workspace_rows,
+                    dtype=torch.int32,
+                    device=self.logical_topk_indices.device,
+                ),
+                side_stream=_create_side_stream(self.logical_topk_indices.device),
+                logical_topk_ready=_create_event(),
+                physical_topk_ready=_create_event(),
+                has_indexer=is_index_producing_layer,
+            )
+        group = self.current_group
+        group_index = group.register_layer(
+            vllm_config,
+            head_size=head_size,
+            kv_cache_dtype=kv_cache_dtype,
+        )
+        return group, group_index
+
+
+def get_sparse_mla_index_group_max_rows(vllm_config: VllmConfig) -> int:
+    max_query_len = 1
+    speculative_config = vllm_config.speculative_config
+    if (
+        speculative_config is not None
+        and speculative_config.num_speculative_tokens is not None
+    ):
+        max_query_len += speculative_config.num_speculative_tokens * (
+            2 if speculative_config.parallel_drafting else 1
+        )
+    scheduler_config = vllm_config.scheduler_config
+    return min(
+        scheduler_config.max_num_batched_tokens,
+        scheduler_config.max_num_seqs * max_query_len,
+    )
diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py
index 5ec07ad211bb..e7eb67b61f48 100644
--- a/vllm/v1/attention/backends/mla/indexer.py
+++ b/vllm/v1/attention/backends/mla/indexer.py
@@ -3,18 +3,19 @@
 from dataclasses import dataclass
 from typing import Any
 
+import numpy as np
 import torch
 
 import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonPointerInputVariant,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
     triton_scalar_specialization_rep,
 )
 from vllm.platforms import current_platform
@@ -24,7 +25,9 @@
     has_deep_gemm,
     native_next_n_supported,
 )
+from vllm.utils.math_utils import round_down
 from vllm.utils.platform_utils import num_compute_units
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -33,6 +36,7 @@
     MultipleOf,
 )
 from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
+from vllm.v1.attention.backends.mla.sparse_utils import request_row_bounds
 from vllm.v1.attention.backends.utils import (
     get_dcp_local_seq_lens,
     split_decodes_and_prefills,
@@ -196,7 +200,7 @@ def get_name() -> str:
         return "DEEPSEEK_V32_INDEXER"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [1, MultipleOf(16)] if current_platform.is_rocm() else [64]
 
     @classmethod
@@ -224,7 +228,7 @@ def get_supported_head_sizes(cls) -> list[int]:
         return []
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
     @staticmethod
@@ -233,6 +237,15 @@ def get_builder_cls() -> type["KpoolTailMetadataBuilder"]:  # type: ignore[overr
 
 
 class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        # ROCm runs adaptive verification through the per-token flattened
+        # indexer path, which derives row ownership from device decode lengths.
+        return (
+            current_platform.is_rocm()
+            or super().supports_device_cpu_query_lens_mismatch()
+        )
+
     @staticmethod
     def get_name() -> str:
         return "DEEPSEEK_V4_INDEXER"
@@ -244,10 +257,112 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
         return (KVCacheLayout.BLHNC, KVCacheLayout.BLNHC)
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
+        # Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows.
         return [256]
 
 
+class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        # The ROCm flattened-query support above is validated for the
+        # DeepSeek-V4 adaptive DSpark path only.
+        return DeepseekV32IndexerBackend.supports_device_cpu_query_lens_mismatch()
+
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_V41_INDEXER"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
+        return [64 if current_platform.is_device_capability_family(90) else 128]
+
+
+@dataclass(frozen=True)
+class PCPGlobalChunkPlan:
+    """PCP packing for one indexer prefill chunk under PCP + DCP."""
+
+    # [num_reqs+1] cumsum of the padded per-request context (shard rows x W):
+    # the global layout.
+    row_start_cu: torch.Tensor
+    # [num_reqs+1] like row_start_cu, but only a region's FIRST row carries its
+    # extent.
+    global_cu: torch.Tensor
+    # [num_reqs+1] cumsum of shard rows: this rank's padded layout.
+    padded_local_cu: torch.Tensor
+    padded_local_total: int
+    total: int
+    # [total] for each global row, its position in the rank-major gathered buffer.
+    deinterleave_idx: torch.Tensor
+
+
+def build_pcp_global_chunk_plan(
+    row_req_idx: np.ndarray,
+    row_shard_rows: np.ndarray,
+    dcp_world_size: int,
+    device: torch.device,
+    interleave: int = 1,
+) -> PCPGlobalChunkPlan:
+    """Plan the PCP packing for one chunk from its requests' KV shard rows.
+
+    The global layout is padded to ``shard rows x W`` per request. Positions
+    past a request's true extent lie beyond every token's causal bound, so
+    the kernel never reads them.
+    """
+    shard_rows = np.ascontiguousarray(row_shard_rows, dtype=np.int64)
+    assert row_req_idx.shape == shard_rows.shape
+    num_rows = len(shard_rows)
+
+    row_bounds = request_row_bounds(row_req_idx)
+    region_first_row = row_bounds[:-1]
+    region_of_row = np.repeat(np.arange(len(region_first_row)), np.diff(row_bounds))
+
+    region_padded = shard_rows[region_first_row]
+    assert np.all(region_padded > 0), (
+        f"PCP+DCP prefill got an empty context: {region_padded.tolist()}"
+    )
+
+    # Pinned host buffers, filled in place so the copies below are async.
+    cu = torch.zeros((3, num_rows + 1), dtype=torch.int32, pin_memory=PIN_MEMORY)
+    row_start_cu, global_cu, padded_cu = cu
+    # Only a region's first row carries its extent.
+    first_row = torch.from_numpy(region_first_row)
+    row_padded = torch.zeros(num_rows, dtype=torch.int32)
+    row_padded[first_row] = torch.from_numpy(region_padded).int()
+    torch.cumsum(row_padded * dcp_world_size, 0, out=global_cu[1:])
+    torch.cumsum(row_padded, 0, out=padded_cu[1:])
+    row_start_cu[:num_rows] = global_cu[first_row][torch.from_numpy(region_of_row)]
+    row_start_cu[num_rows] = global_cu[num_rows]
+    total = int(global_cu[num_rows])
+    padded_total = int(padded_cu[num_rows])
+
+    idx = torch.empty(total, dtype=torch.int64, pin_memory=PIN_MEMORY)
+    idx_np = idx.numpy()
+    region_start = global_cu[first_row].numpy()
+    region_padded_start = padded_cu[first_row].numpy()
+    for i in range(len(region_first_row)):
+        g = int(region_padded[i]) * dcp_world_size
+        t = np.arange(g, dtype=np.int64)
+        local = round_down(t // dcp_world_size, interleave) + t % interleave
+        # Padded positions are outside causal bounds but must remain in-bounds.
+        local = np.minimum(local, region_padded[i] - 1)
+        idx_np[region_start[i] : region_start[i] + g] = (
+            ((t // interleave) % dcp_world_size) * padded_total
+            + region_padded_start[i]
+            + local
+        )
+
+    cu = async_tensor_h2d(cu, device=device)
+    return PCPGlobalChunkPlan(
+        row_start_cu=cu[0],
+        global_cu=cu[1],
+        padded_local_cu=cu[2],
+        padded_local_total=padded_total,
+        total=total,
+        deinterleave_idx=async_tensor_h2d(idx, device=device),
+    )
+
+
 @dataclass
 class DeepseekV32IndexerPrefillChunkMetadata:
     block_table: torch.Tensor
@@ -266,6 +381,8 @@ class DeepseekV32IndexerPrefillChunkMetadata:
     local_total_seq_lens: int = 0
     max_local_total_seq_lens: int = 0
 
+    pcp_deinterleave_idx: torch.Tensor | None = None
+
 
 class BuildPrefillChunkMetadataKernel(
     VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
@@ -385,35 +502,59 @@ def dispatch(  # type: ignore[override]
 
     def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
         max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8))
-        hf_config = vllm_config.model_config.hf_config
+        hf_text_config = vllm_config.model_config.hf_text_config
         parallel_config = vllm_config.parallel_config
         dcp_world = parallel_config.decode_context_parallel_size
         dcp_interleave = parallel_config.cp_kv_cache_interleave_size
         dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0
+        dcp_cases = [dict(DCP_RANK=dcp_rank, DCP_WORLD=dcp_world)]
+        if parallel_config.prefill_context_parallel_size > 1 and dcp_world > 1:
+            # PCP+DCP prefill constructs a global PCP chunk plan and therefore
+            # dispatches this kernel with already-localized row starts. The
+            # runtime key is normalized to rank 0/world 1 in that path; warm it
+            # explicitly so one PCP rank cannot enter a collective while
+            # another rank is still compiling the first real inference batch.
+            dcp_cases.append(dict(DCP_RANK=0, DCP_WORLD=1))
         compress_ratios = tuple(
             dict.fromkeys(
                 max(1, int(ratio))
                 for ratio in (
-                    *(getattr(hf_config, "compress_ratios", None) or (1,)),
-                    getattr(hf_config, "index_kpool", 1) or 1,
+                    *(getattr(hf_text_config, "compress_ratios", None) or (1,)),
+                    getattr(hf_text_config, "index_kpool", 1) or 1,
                 )
             )
         )
-        index_kpool = getattr(hf_config, "index_kpool", None)
+        index_kpool = getattr(hf_text_config, "index_kpool", None)
         if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios:
             compress_ratios = compress_ratios + (index_kpool,)
         return self._trace_dispatch(self.dispatch)(
+            zip_inputs(*dcp_cases),
             # Cover Triton's divisible, exact-one, and generic i32 classes.
             query_slice_start=(0, 1, 2),
             query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
-            DCP_RANK=dcp_rank,
-            DCP_WORLD=dcp_world,
             DCP_INTERLEAVE=dcp_interleave,
             BLOCK_SIZE=self.BLOCK_SIZE,
             COMPRESS_RATIO=list(compress_ratios),
+            # PCP's global cumulative lengths are the second row of one packed
+            # allocation, so their pointer is not always 16-byte aligned.
+            # Prefill chunking can independently unalign uncompressed lengths.
             input_variant=(
-                TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
-                TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
+                TritonPointerInputVariant.from_alignment(
+                    uncompressed_seq_lens=True,
+                    cu_compressed_seq_lens=True,
+                ),
+                TritonPointerInputVariant.from_alignment(
+                    uncompressed_seq_lens=True,
+                    cu_compressed_seq_lens=False,
+                ),
+                TritonPointerInputVariant.from_alignment(
+                    uncompressed_seq_lens=False,
+                    cu_compressed_seq_lens=True,
+                ),
+                TritonPointerInputVariant.from_alignment(
+                    uncompressed_seq_lens=False,
+                    cu_compressed_seq_lens=False,
+                ),
             ),
         )
 
@@ -424,7 +565,9 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
             uncompressed_seq_lens=compile_key.input_variant.pointer(
                 "uncompressed_seq_lens", torch.int32
             ),
-            cu_compressed_seq_lens=int32_ptr,
+            cu_compressed_seq_lens=compile_key.input_variant.pointer(
+                "cu_compressed_seq_lens", torch.int32
+            ),
             row_start_cu_compressed_seq_lens=int32_ptr,
             token_to_seq=int32_ptr,
             cu_compressed_seq_len_ks=int32_ptr,
@@ -506,6 +649,38 @@ class DeepseekV32IndexerMetadata:
     prefill: DeepseekV32IndexerPrefillMetadata | None = None
 
 
+@triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"])
+def _kpool_tail_slot_mapping_kernel(
+    slot_mapping_ptr,
+    block_table_ptr,
+    block_table_stride,
+    query_start_loc_ptr,
+    positions_ptr,
+    out_ptr,
+    num_reqs,
+    num_actual_tokens,
+    num_tokens,
+    kpool,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    if pid < num_reqs:
+        start = tl.load(query_start_loc_ptr + pid)
+        # Tokens past the last request's boundary (if any) also map to it.
+        end = tl.load(query_start_loc_ptr + pid + 1)
+        end = tl.where(pid == num_reqs - 1, num_actual_tokens, end)
+        own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64)
+        for i in range(start, end, BLOCK):
+            offs = i + tl.arange(0, BLOCK)
+            mask = offs < end
+            pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64)
+            tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask)
+    else:
+        offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK)
+        mask = offs < num_tokens
+        tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask)
+
+
 def compute_kpool_tail_slot_mapping(
     slot_mapping: torch.Tensor,
     block_table: torch.Tensor,
@@ -518,10 +693,33 @@ def compute_kpool_tail_slot_mapping(
 ) -> torch.Tensor:
     """Map every token to its request's one circular tail block."""
     if out is None:
-        out = slot_mapping.clone()
+        out = torch.empty_like(slot_mapping)
     else:
         assert out.shape == slot_mapping.shape
-        out.copy_(slot_mapping)
+    if slot_mapping.is_cuda and slot_mapping.dim() == 1 and num_reqs > 0:
+        block = 256
+        num_tokens = slot_mapping.shape[0]
+        num_actual_tokens = min(num_actual_tokens, num_tokens)
+        grid = (num_reqs + triton.cdiv(num_tokens - num_actual_tokens, block),)
+        _kpool_tail_slot_mapping_kernel[grid](
+            slot_mapping,
+            block_table,
+            block_table.stride(0),
+            query_start_loc,
+            positions,
+            out,
+            num_reqs,
+            num_actual_tokens,
+            num_tokens,
+            kpool,
+            BLOCK=block,
+            num_warps=4,
+        )
+        return out
+    # Torch fallback: CPU tensors (the CPU unit tests), non-1D or empty inputs.
+    # Production always takes the Triton path above — spec-decode tokens arrive
+    # flattened token-major, so slot_mapping is 1D there too.
+    out.copy_(slot_mapping)
     if num_actual_tokens == 0:
         return out
     tokens = torch.arange(num_actual_tokens, device=slot_mapping.device)
@@ -619,6 +817,15 @@ def _supports_flattened_device_query_lens() -> bool:
     )
 
 
+def _rocm_supports_flattened_device_query_lens(vllm_config: VllmConfig) -> bool:
+    model_config = vllm_config.model_config
+    return (
+        current_platform.is_rocm()
+        and model_config is not None
+        and "DeepseekV4ForCausalLM" in model_config.architectures
+    )
+
+
 def _supports_native_decode(next_n: int) -> bool:
     """Whether decode can pass `next_n` Q rows per request to the kernel
     instead of flattening to one single-token row per query, which re-reads
@@ -639,7 +846,10 @@ def _use_flattening(vllm_config: VllmConfig) -> bool:
     return not _supports_native_decode(next_n) or (
         speculative_config is not None
         and speculative_config.enable_adaptive_verification
-        and _supports_flattened_device_query_lens()
+        and (
+            _supports_flattened_device_query_lens()
+            or _rocm_supports_flattened_device_query_lens(vllm_config)
+        )
     )
 
 
@@ -667,16 +877,8 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None:
         self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
         self.pcp_world_size = parallel_config.prefill_context_parallel_size
         self.use_pcp = self.pcp_world_size > 1
+        self.pcp_rank = get_pcp_group().rank_in_group if self.use_pcp else 0
         self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
-        # The DCP sparse-indexer code is parameterized by interleave size, but
-        # interleave > 1 is not yet validated end-to-end (gsm8k parity fails),
-        # so fail closed here rather than silently produce wrong output.
-        if self.dcp_world_size > 1 and self.cp_kv_cache_interleave_size > 1:
-            raise NotImplementedError(
-                "DCP sparse indexer currently supports only "
-                f"cp_kv_cache_interleave_size=1 (got "
-                f"{self.cp_kv_cache_interleave_size})."
-            )
         # NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
         self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config)
         self.num_speculative_tokens = (
@@ -684,7 +886,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None:
             if self.vllm_config.speculative_config
             else 0
         )
-        self.use_fp4_indexer_cache = dsa_indexer_uses_fp4(self.vllm_config)
+        self.indexer_uses_fp4 = dsa_indexer_uses_fp4(self.vllm_config)
 
         next_n = self.num_speculative_tokens + 1
         self.decode_threshold = next_n
@@ -697,7 +899,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None:
             self.use_flattening,
             self.supports_varlen,
             next_n,
-            self.use_fp4_indexer_cache,
+            self.indexer_uses_fp4,
         )
 
         sm_count = num_compute_units(self.device.index)
@@ -742,6 +944,12 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None:
             dtype=torch.int32,
             device=self.device,
         )
+        # Materialize the rank on device during builder initialization. Creating
+        # this scalar in build() would introduce a GPU<->CPU sync in the decode
+        # hot path.
+        self.dcp_rank_tensor = torch.tensor(
+            self.dcp_rank, dtype=torch.int32, device=self.device
+        )
         self.expanded_block_table_buffer = torch.zeros(
             (scheduler_config.max_num_batched_tokens, block_table_width),
             dtype=torch.int32,
@@ -795,7 +1003,7 @@ def _dcp_localize_decode_seq_lens(
         local_seq_lens = get_dcp_local_seq_lens(
             seq_lens,
             self.dcp_world_size,
-            self.dcp_rank,
+            self.dcp_rank_tensor,
             self.cp_kv_cache_interleave_size,
         )
         if seq_lens_is_buffer_view:
@@ -949,27 +1157,52 @@ def _prepare_global_decode_seq_lens(
         self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0
         return self.global_decode_seq_lens_buffer[:num_decode_tokens]
 
-    def _build_varlen_decode_indices(
+    def _split_pcp_dcp_prefill_chunks(
         self,
-        decode_lens: torch.Tensor,
-        decode_lens_cpu: torch.Tensor,
-        num_decode_tokens: int,
-    ) -> torch.Tensor:
-        """Build request ids for flattened SM100 varlen rows."""
-        indices = self.decode_indices_buffer[:num_decode_tokens]
-        actual_expanded = int(decode_lens_cpu.sum().item())
-        num_decodes = decode_lens.shape[0]
-        indices[:actual_expanded] = torch.repeat_interleave(
-            self.arange_buffer[:num_decodes],
-            decode_lens,
-            output_size=actual_expanded,
+        row_req_idx: np.ndarray,
+        row_shard_rows: np.ndarray,
+        row_query_lens_cpu: torch.Tensor,
+        max_logits_bytes: int,
+        request_offset: int,
+    ) -> list[tuple[slice, slice]]:
+        """Chunk by request rather than by row, so a split prefill's two rows
+        charge their shared context once, then widen each chunk back to rows.
+
+        ``row_shard_rows`` holds the whole request's largest DCP shard on
+        every row, so the plan is identical on every PCP rank.
+        """
+        row_bounds = request_row_bounds(row_req_idx)
+        first_rows = row_bounds[:-1]
+        # Each request's context, padded to whole DCP shards.
+        seq_lens = row_shard_rows[first_rows] * self.dcp_world_size
+        # Every rank holds a full-size chunk of a split request (shorter ones
+        # are replicated), so rows x longest row is the same on every rank,
+        # whichever row holds the short tail.
+        row_query_lens = row_query_lens_cpu.numpy()
+        query_lens = np.diff(row_bounds) * np.maximum.reduceat(
+            row_query_lens, first_rows
+        )
+        chunk_specs = self._split_indexer_prefill_chunks(
+            torch.from_numpy(seq_lens.astype(np.int32)),
+            torch.from_numpy(query_lens.astype(np.int32)),
+            self.max_prefill_buffer_size,
+            max_logits_bytes,
         )
-        if actual_expanded < num_decode_tokens:
-            pad = num_decode_tokens - actual_expanded
-            indices[actual_expanded:num_decode_tokens] = (
-                num_decodes + self.arange_buffer[:pad]
+        return [
+            (
+                slice(
+                    request_offset + int(row_bounds[request_slice.start]),
+                    request_offset + int(row_bounds[request_slice.stop]),
+                ),
+                query_slice,
             )
-        return indices
+            for request_slice, query_slice in chunk_specs
+        ]
+
+    def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor:
+        """Per-request KV lengths the prefill chunker budgets logits with;
+        subclasses whose logits rows are wider than the context override."""
+        return seq_lens_cpu
 
     @staticmethod
     def _split_indexer_prefill_chunks(
@@ -1055,7 +1288,6 @@ def build(
         assert num_decode_tokens + num_prefill_tokens == num_tokens
 
         compressed_slot_mapping = slot_mapping
-        compressed_seq_lens = seq_lens
         indexer_block_table = block_table
         if self.compress_ratio > 1:
             kernel_block_size = self.kernel_block_size
@@ -1067,10 +1299,18 @@ def build(
                 factor = self.kv_cache_spec.block_size // kernel_block_size
                 indexer_block_table = (block_table[:, ::factor] // factor).contiguous()
             padded_num_tokens = num_tokens
-            if self.pcp_world_size > 1:
+            local_slot_mapping = slot_mapping
+            if self.use_pcp:
+                # The gathered layout holds each rank's local tokens, padded, in
+                # rank order, so this rank's segment lines up with query_start_loc.
                 padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size
+                local_slot_mapping = slot_mapping[
+                    self.pcp_rank * padded_num_tokens : (self.pcp_rank + 1)
+                    * padded_num_tokens
+                ]
             compressed_slot_mapping = get_compressed_slot_mapping(
                 num_tokens,
+                local_slot_mapping,
                 query_start_loc,
                 seq_lens,
                 indexer_block_table,
@@ -1083,10 +1323,12 @@ def build(
                     self.compressed_slot_mapping_buffer[:padded_num_tokens],
                     dim=0,
                 )
-            compressed_seq_lens = seq_lens // self.compress_ratio
 
         prefill_metadata = None
         if num_prefills > 0:
+            compressed_seq_lens = (
+                seq_lens // self.compress_ratio if self.compress_ratio > 1 else seq_lens
+            )
             # This CPU value is an upper bound for async-spec extend rows.  It
             # is safe for chunking/allocation because CUDA metadata below is
             # built from exact device seq_lens and gather ignores the tail.
@@ -1105,16 +1347,53 @@ def build(
             # slice below).
             assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
             seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
-            chunk_specs = self._split_indexer_prefill_chunks(
-                compressed_seq_lens_cpu[num_decodes:],
-                prefill_query_lens_cpu,
-                self.max_prefill_buffer_size,
-                max_logits_bytes,
-                request_offset=num_decodes,
-            )
+            req_idx = None
+            shard_rows = None
+            if self.use_pcp and self.dcp_world_size > 1:
+                # The gathered KV must be packed identically on every PCP rank:
+                # chunk by request from its DCP shard rows, which every rank
+                # holds. A dummy batch bypasses the PCP manager and has one
+                # row per request, so its own extent is the request's.
+                req_idx = common_attn_metadata.req_idx
+                if req_idx is None:
+                    req_idx = np.arange(num_reqs)
+                shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound
+                if shard_rows_cpu is None:
+                    shard_rows_cpu = get_dcp_local_seq_lens(
+                        seq_lens_cpu,
+                        self.dcp_world_size,
+                        0,
+                        self.cp_kv_cache_interleave_size,
+                    )
+                shard_rows = shard_rows_cpu.numpy()
+                chunk_specs = self._split_pcp_dcp_prefill_chunks(
+                    req_idx[num_decodes:],
+                    shard_rows[num_decodes:],
+                    prefill_query_lens_cpu,
+                    max_logits_bytes,
+                    request_offset=num_decodes,
+                )
+            else:
+                chunk_specs = self._split_indexer_prefill_chunks(
+                    self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]),
+                    prefill_query_lens_cpu,
+                    self.max_prefill_buffer_size,
+                    max_logits_bytes,
+                    request_offset=num_decodes,
+                )
 
             chunks = []
             for req_slice, query_slice in chunk_specs:
+                pcp_plan = None
+                if req_idx is not None:
+                    assert shard_rows is not None
+                    pcp_plan = build_pcp_global_chunk_plan(
+                        req_idx[req_slice],
+                        shard_rows[req_slice],
+                        self.dcp_world_size,
+                        self.device,
+                        self.cp_kv_cache_interleave_size,
+                    )
                 metadata = build_prefill_chunk_metadata(
                     req_slice.start,
                     req_slice.stop,
@@ -1130,6 +1409,7 @@ def build(
                     dcp_rank=self.dcp_rank,
                     dcp_world_size=self.dcp_world_size,
                     cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+                    pcp_plan=pcp_plan,
                 )
                 # Skip when total_seq_lens is 0 (i.e., no compressed token).
                 if metadata is not None:
@@ -1145,12 +1425,15 @@ def build(
 
         decode_metadata = None
         if num_decodes > 0:
-            torch.diff(
-                common_attn_metadata.query_start_loc[: num_decodes + 1],
-                out=self.decode_lens_buffer[:num_decodes],
-            )
+            if not self.supports_varlen:
+                torch.diff(
+                    common_attn_metadata.query_start_loc[: num_decodes + 1],
+                    out=self.decode_lens_buffer[:num_decodes],
+                )
+                self.per_req_decode_lens_buffer[:num_decodes].copy_(
+                    self.decode_lens_buffer[:num_decodes]
+                )
             decode_lens = self.decode_lens_buffer[:num_decodes]
-            self.per_req_decode_lens_buffer[:num_decodes].copy_(decode_lens)
             decode_lens_cpu = torch.diff(
                 common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
             )
@@ -1185,38 +1468,74 @@ def build(
                 and step_next_n_ok
             )
 
-            global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
-                global_seq_lens=global_seq_lens_for_decode,
-                decode_lens=decode_lens,
-                decode_lens_cpu=decode_lens_cpu,
-                query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
-                num_decode_tokens=num_decode_tokens,
-                use_native=use_native,
-                max_decode_len=max_decode_len,
-            )
-
-            decode_indices = None
-            if self.supports_varlen:
-                decode_indices = self._build_varlen_decode_indices(
-                    decode_lens=decode_lens,
-                    decode_lens_cpu=decode_lens_cpu,
-                    num_decode_tokens=num_decode_tokens,
-                )
-
-            seq_lens, block_table, decode_lens, batch_size, requires_padding = (
-                self._prepare_decode_tensors(
-                    seq_lens=seq_lens,
-                    block_table=block_table,
+            if not self.supports_varlen:
+                global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
+                    global_seq_lens=global_seq_lens_for_decode,
                     decode_lens=decode_lens,
                     decode_lens_cpu=decode_lens_cpu,
                     query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
-                    num_decodes=num_decodes,
                     num_decode_tokens=num_decode_tokens,
                     use_native=use_native,
-                    next_n=next_n,
                     max_decode_len=max_decode_len,
                 )
-            )
+
+            decode_indices = None
+            if self.supports_varlen:
+                from vllm.v1.attention.ops.metadata import (
+                    _indexer_decode_metadata_kernel,
+                )
+
+                capacity = self.decode_seq_lens_buffer.numel()
+                grid = max(
+                    num_decodes,
+                    num_decode_tokens + triton.cdiv(capacity - num_decode_tokens, 256),
+                )
+                _indexer_decode_metadata_kernel[(grid,)](
+                    query_start_loc,
+                    seq_lens,
+                    block_table,
+                    self.decode_seq_lens_buffer,
+                    self.expanded_block_table_buffer,
+                    self.decode_lens_buffer,
+                    self.decode_indices_buffer,
+                    self.per_req_decode_lens_buffer,
+                    num_decodes,
+                    num_decode_tokens,
+                    capacity,
+                    block_table.stride(0),
+                    self.expanded_block_table_buffer.stride(0),
+                    BLOCK_COLS=block_table.shape[1],
+                    num_warps=4,
+                )
+                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
+                block_table = self.expanded_block_table_buffer[:num_decode_tokens]
+                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
+                decode_indices = self.decode_indices_buffer[:num_decode_tokens]
+                requires_padding = False
+                if global_seq_lens_for_decode is not None and max_decode_len > 1:
+                    self.global_decode_seq_lens_buffer[:num_decode_tokens].copy_(
+                        seq_lens
+                    )
+                    global_seq_lens_for_decode = self.global_decode_seq_lens_buffer[
+                        :num_decode_tokens
+                    ]
+            else:
+                seq_lens, block_table, decode_lens, batch_size, requires_padding = (
+                    self._prepare_decode_tensors(
+                        seq_lens=seq_lens,
+                        block_table=block_table,
+                        decode_lens=decode_lens,
+                        decode_lens_cpu=decode_lens_cpu,
+                        query_start_loc=common_attn_metadata.query_start_loc[
+                            :num_decodes
+                        ],
+                        num_decodes=num_decodes,
+                        num_decode_tokens=num_decode_tokens,
+                        use_native=use_native,
+                        next_n=next_n,
+                        max_decode_len=max_decode_len,
+                    )
+                )
 
             if self.compress_ratio > 1:
                 kernel_block_size = self.kernel_block_size
@@ -1239,9 +1558,9 @@ def build(
                     )
                     block_table = self.indexer_decode_block_table_buffer[:rows, :cols]
 
-            seq_lens_is_buffer_view = (use_native and next_n > 1) or (
-                not use_native and max_decode_len > 1
-            )
+            # Flattening always returns a buffer view, including single-token
+            # batches. Keep its address stable across varlen graph replays.
+            seq_lens_is_buffer_view = not use_native or next_n > 1
 
             # DCP: localize the now-expanded per-token global bounds to this
             # rank's owned KV. Done here (after expansion) so each token's global
@@ -1325,8 +1644,12 @@ def build_prefill_chunk_metadata(
     dcp_rank: int = 0,
     dcp_world_size: int = 1,
     cp_kv_cache_interleave_size: int = 1,
+    pcp_plan: PCPGlobalChunkPlan | None = None,
 ) -> DeepseekV32IndexerPrefillChunkMetadata | None:
-    total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
+    if pcp_plan is not None:
+        total_seq_lens = pcp_plan.total
+    else:
+        total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
     if total_seq_lens == 0:
         return None
 
@@ -1334,15 +1657,22 @@ def build_prefill_chunk_metadata(
     device = block_table.device
     token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device)
 
-    cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
-    # Assigning to slice avoids cpu sync.
-    cu_seq_lens[:1] = 0
-    torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
+    if pcp_plan is not None:
+        cu_seq_lens = pcp_plan.global_cu
+    else:
+        cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
+        # Assigning to slice avoids cpu sync.
+        cu_seq_lens[:1] = 0
+        torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
 
     local_cu_seq_lens = cu_seq_lens
     local_total_seq_lens = total_seq_lens
     max_local_total_seq_lens = total_seq_lens
-    if dcp_world_size > 1:
+    if pcp_plan is not None:
+        local_cu_seq_lens = pcp_plan.padded_local_cu
+        local_total_seq_lens = pcp_plan.padded_local_total
+        max_local_total_seq_lens = pcp_plan.padded_local_total
+    elif dcp_world_size > 1:
         # Per-rank local KV length under interleave-aware DCP sharding, shape
         # [num_reqs, dcp_world_size]. Reuse the canonical CP helper so the
         # sharding matches the rest of the DCP pipeline (decode/prefill).
@@ -1366,35 +1696,47 @@ def build_prefill_chunk_metadata(
         (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
     )
     if query_slice is not None:
-        qs_start = query_slice.start
-        qs_stop = query_slice.stop
+        qs_start = min(query_slice.start, total_query_len)
+        qs_stop = min(query_slice.stop, total_query_len)
     else:
         qs_start = 0
         qs_stop = total_query_len
     output_query_len = qs_stop - qs_start
+    assert 0 <= output_query_len <= total_query_len - qs_start, (
+        f"query slice [{qs_start}, {qs_stop}) does not fit this rank's "
+        f"{total_query_len} tokens; the metadata tail would stay uninitialized"
+    )
 
     cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
     cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
 
-    # Under DCP the kernel writes this rank's local row bounds into
-    # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
+    if pcp_plan is not None:
+        row_start_cu = pcp_plan.row_start_cu
+        kernel_dcp_rank, kernel_dcp_world = 0, 1
+    else:
+        # Under DCP the kernel writes this rank's local row bounds into
+        # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
+        row_start_cu = local_cu_seq_lens
+        kernel_dcp_rank, kernel_dcp_world = dcp_rank, dcp_world_size
     _BUILD_PREFILL_CHUNK_METADATA_KERNEL(
         query_start_loc,
         uncompressed_seq_lens[start_idx:end_idx],
         cu_seq_lens,
-        local_cu_seq_lens,
+        row_start_cu,
         token_to_seq,
         cu_seq_len_ks,
         cu_seq_len_ke,
         qs_start,
         qs_stop,
-        dcp_rank,
-        dcp_world_size,
+        kernel_dcp_rank,
+        kernel_dcp_world,
         cp_kv_cache_interleave_size,
         num_reqs=num_reqs,
         COMPRESS_RATIO=compress_ratio,
     )
 
+    pcp_deinterleave_idx = pcp_plan.deinterleave_idx if pcp_plan is not None else None
+
     token_start = query_start_loc_cpu[start_idx].item()
     if query_slice is not None:
         token_end = token_start + qs_stop
@@ -1404,6 +1746,7 @@ def build_prefill_chunk_metadata(
         token_end = query_start_loc_cpu[end_idx].item()
 
     return DeepseekV32IndexerPrefillChunkMetadata(
+        pcp_deinterleave_idx=pcp_deinterleave_idx,
         cu_seqlen_ks=cu_seq_len_ks,
         cu_seqlen_ke=cu_seq_len_ke,
         cu_seq_lens=cu_seq_lens,
diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py
index 3ba5a5a5fad3..1d03b9e77392 100644
--- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py
+++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py
@@ -24,7 +24,7 @@
 
 
 class AiterFlashAttnPrefillBackend(MLAPrefillBackend):
-    """AITER FlashAttention backend for MLA prefill"""
+    """AITER FlashAttention backend for MLA prefill."""
 
     @staticmethod
     def get_name() -> str:
diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py
index 3e0d544195e1..3af623b1f13b 100644
--- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py
+++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py
@@ -327,6 +327,13 @@ def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:
                 qk_rope_head_dim=64,
                 v_head_dim=128,
             ),
+            # GLM5Next NoPE layout: qk_head_dim 256 + 0 and v_head_dim 256 run
+            # the same kernels as the (192, 64, 256) DeepSeek-V3.2 layout.
+            MLADimensions(
+                qk_nope_head_dim=256,
+                qk_rope_head_dim=0,
+                v_head_dim=256,
+            ),
         ]
 
     def __init__(
diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py
index e7e00a0d0143..1cf1761a7708 100644
--- a/vllm/v1/attention/backends/mla/prefill/registry.py
+++ b/vllm/v1/attention/backends/mla/prefill/registry.py
@@ -65,6 +65,7 @@ def get_path(self) -> str:
 
         Raises:
             ValueError: If Backend.CUSTOM is used without being registered
+
         """
         path = _MLA_PREFILL_OVERRIDES.get(self, self.value)
         if not path:
@@ -85,6 +86,7 @@ def get_class(self) -> "type[MLAPrefillBackend]":
         Raises:
             ImportError: If the backend class cannot be imported
             ValueError: If CUSTOM is used without being registered
+
         """
         return resolve_obj_by_qualname(self.get_path())
 
@@ -130,6 +132,7 @@ class MyCustomPrefillBackend(MLAPrefillBackend):
             MLAPrefillBackendEnum.CUSTOM,
             "my.module.MyCustomPrefillBackend"
         )
+
     """
 
     def decorator(cls: type) -> type:
diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py
index c7e935fa75bb..a2b74e4327eb 100644
--- a/vllm/v1/attention/backends/mla/prefill/selector.py
+++ b/vllm/v1/attention/backends/mla/prefill/selector.py
@@ -57,6 +57,7 @@ def _get_mla_prefill_backend_priorities(
 
     Returns:
         List of backends in priority order (highest priority first).
+
     """
     from vllm.platforms import current_platform
 
@@ -104,6 +105,7 @@ def get_mla_prefill_backend(
 
     Returns:
         The selected prefill backend class.
+
     """
     from vllm.platforms import current_platform
 
@@ -173,6 +175,7 @@ def _auto_select_mla_prefill_backend(
 
     Returns:
         The selected prefill backend class.
+
     """
     priorities = _get_mla_prefill_backend_priorities(
         device_capability,
diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
index c341bfc3486e..1ab4df0664d5 100644
--- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
@@ -4,7 +4,7 @@
 import functools
 from dataclasses import dataclass
 from pathlib import Path
-from typing import ClassVar, Final
+from typing import TYPE_CHECKING, ClassVar, Final
 
 import torch
 
@@ -22,6 +22,7 @@
     QueryLenSupport,
 )
 from vllm.triton_utils import tl, triton
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.math_utils import cdiv, largest_power_of_2_divisor
 from vllm.v1.attention.backend import (
     AttentionCGSupport,
@@ -37,6 +38,9 @@
 )
 from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
 
+if TYPE_CHECKING:
+    from vllm.platforms.interface import DeviceCapability
+
 logger = init_logger(__name__)
 
 
@@ -143,6 +147,20 @@ def _aiter_mla_native_h24_supported() -> bool:
     )
 
 
+def _aiter_mla_non_causal_asm_kernels() -> bool:
+    """Whether this arch ships non-causal MLA decode ASM kernels.
+
+    The Python `causal=` probe cannot see the per-arch manifest. gfx950 has
+    the kernels; gfx942 and everything else must fall through to another
+    backend.
+    """
+    try:
+        from vllm.platforms.rocm import on_gfx950
+    except Exception:  # noqa: BLE001
+        return False
+    return bool(on_gfx950())
+
+
 @functools.lru_cache(maxsize=1)
 def _gluon_mla_decode_supported() -> bool:
     """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build.
@@ -160,6 +178,30 @@ def _gluon_mla_decode_supported() -> bool:
     return on_gfx950()
 
 
+# Past 2 GiB, mla_gluon swaps its masked buffer_load for an unmasked
+# global_load that reads outside the cache and aborts the process. Mirrors the
+# kernel's own within_2gb bound, so the fallback engages exactly where it does.
+_GLUON_MAX_KV_CACHE_BYTES = 1 << 31
+
+
+def _gluon_kv_cache_in_bounds(kv_cache_bytes: int | None) -> bool:
+    """Whether this layer's KV cache keeps Gluon on its bounds-checked path.
+
+    The cache is one flat tensor per layer, so the bound disqualifies the layer
+    outright; no batch shape brings an oversized cache back in range. ``None``
+    is the not-yet-sized case during profiling, before any kernel runs.
+    """
+    if kv_cache_bytes is None or kv_cache_bytes <= _GLUON_MAX_KV_CACHE_BYTES:
+        return True
+    logger.warning_once(
+        "KV cache is %.1f GiB per layer, past the 2 GiB where the Gluon MLA "
+        "kernel drops its KV bounds mask; using the padded ASM decode instead. "
+        "Lower --gpu-memory-utilization or --max-model-len to get Gluon back.",
+        kv_cache_bytes / (1 << 30),
+    )
+    return False
+
+
 @functools.lru_cache(maxsize=1)
 def _get_segmented_mla_decode():
     """Load AITER's segmented MLA decode with unreduced partial output."""
@@ -240,7 +282,7 @@ def get_supported_head_sizes(cls) -> list[int]:
         return []
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # The aiter MLA decode kernel always operates with page_size=1
         # internally (the wrapper flattens kv_buffer via .view(-1, 1, 1, H)).
         # We support any kernel_block_size by expanding block-level indices
@@ -251,6 +293,72 @@ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
     def get_name() -> str:
         return "ROCM_AITER_MLA"
 
+    @classmethod
+    def supports_non_causal(cls) -> bool:
+        # causal= on mla_decode_fwd is necessary but not sufficient: pinned
+        # AITER v0.1.21.post2 still exports that argument on gfx942, whose ASM
+        # manifest has no non-causal decode kernels. Selecting this backend
+        # there reaches `cannot get heuristic kernel!` instead of TRITON_MLA.
+        return bool(
+            rocm_aiter_ops.mla_decode_supports_non_causal()
+            and _aiter_mla_non_causal_asm_kernels()
+        )
+
+    @classmethod
+    def validate_configuration(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        use_per_head_quant_scales: bool,
+        device_capability: "DeviceCapability",
+        attn_type: str,
+        has_sliding_window: bool = False,
+        use_non_causal: bool = False,
+        use_batch_invariant: bool = False,
+        use_kv_connector: bool = False,
+        use_pcp: bool = False,
+        use_adaptive_verification: bool = False,
+        use_dcp: bool = False,
+        use_rswa: bool = False,
+    ) -> list[str]:
+        invalid_reasons = super().validate_configuration(
+            head_size=head_size,
+            dtype=dtype,
+            kv_cache_dtype=kv_cache_dtype,
+            block_size=block_size,
+            use_mla=use_mla,
+            has_sink=has_sink,
+            use_sparse=use_sparse,
+            use_mm_prefix=use_mm_prefix,
+            use_per_head_quant_scales=use_per_head_quant_scales,
+            device_capability=device_capability,
+            attn_type=attn_type,
+            has_sliding_window=has_sliding_window,
+            use_non_causal=use_non_causal,
+            use_batch_invariant=use_batch_invariant,
+            use_kv_connector=use_kv_connector,
+            use_pcp=use_pcp,
+            use_adaptive_verification=use_adaptive_verification,
+            use_dcp=use_dcp,
+            use_rswa=use_rswa,
+        )
+        # Arch/signature cannot express this: the backend still advertises
+        # fp16, so a DSpark draft with kv_cache_dtype="auto" would select
+        # AITER and abort on fp16 query. Causal fp16 is fine.
+        if (
+            use_non_causal
+            and dtype == torch.float16
+            and kv_cache_dtype in (None, "auto", "float16")
+        ):
+            invalid_reasons.append("non-causal fp16 MLA decode not supported")
+        return invalid_reasons
+
     @staticmethod
     def get_impl_cls() -> type["AiterMLAImpl"]:
         return AiterMLAImpl
@@ -346,6 +454,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
     #  https://github.com/vllm-project/vllm/issues/22945
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
     query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
+    # Served by passing the mask to the kernel; _build_decode turns away the
+    # shapes AITER has no non-causal kernel for.
+    supports_non_causal_multi_token_decode: ClassVar[bool] = True
+    # Set from the common metadata every build; a batch is causal unless the
+    # drafter says otherwise.
+    _decode_causal: bool = True
 
     @staticmethod
     def _uniform_padded_mtp_qo_len(
@@ -474,6 +588,13 @@ def __init__(
         # _build_decode needs the cache dtype to pick the decode kernel; keep
         # the normalized string instead of dropping it at the end of __init__.
         self._kv_cache_dtype_str = kv_cache_dtype_str
+        # Sized per layer, matching the flat KV tensor the kernel is handed.
+        num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
+        self._kv_cache_bytes = (
+            None
+            if num_gpu_blocks is None
+            else num_gpu_blocks * kv_cache_spec.page_size_bytes
+        )
         # MLAAttention quantizes decode Q to FP8 before calling this backend
         # whenever the KV cache is FP8 and supports_quant_query_input is true.
         q_dtype = (
@@ -620,6 +741,7 @@ def _init_fp8_prefill_ps_buffers(
             attn_out_dtype: Dtype of the attention output buffer, used to size
                 the padded-head output scratch (small head counts only).
             device: Target device for the buffers.
+
         """
         from aiter import get_ps_metadata_info_v1
 
@@ -681,9 +803,18 @@ def _init_fp8_prefill_ps_buffers(
             device=device,
         )
 
+        from vllm.platforms import current_platform
         from vllm.v1.worker.workspace import current_workspace_manager
 
-        max_num_partial_tiles = reduce_partial_map_size
+        # AITER metadata sizing assumes all requests can carry max_prefill_qlen tokens,
+        # which is a loose worst case in the number of QO tiles. Rather, the sum of
+        # qlens is bounded by max_num_batched_tokens. Manually compute the number of
+        # qo tiles to avoid OOM at startup.
+        # TODO: AITER should give us this budget constrained value
+        qo_tile_cnt = (
+            cdiv(max_num_batched_tokens, _FP8_PREFILL_TILE_Q) + max_num_reqs - 1
+        )
+        max_num_partial_tiles = qo_tile_cnt + current_platform.num_compute_units()
         reservations: list[tuple[tuple[int, ...], torch.dtype]] = [
             (
                 (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),
@@ -796,7 +927,8 @@ def _build_fp8_prefill_ps_metadata(
         # build) keeps it off the per-layer forward path where a sync would
         # break CUDA Graph capture.  Using the device-side reduce_indptr is
         # acceptable since build is allowed to incur an occasional sync.
-        num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item())
+        with gpu_sync_allowed():
+            num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item())
 
         # Attach PS metadata to the metadata object so forward_mha can read it.
         metadata.fp8_prefill_qo_indptr = qo_indptr
@@ -915,8 +1047,10 @@ def _build_decode(
         query_start_loc_cpu: torch.Tensor,
         query_start_loc_device: torch.Tensor,
         num_decode_tokens: int,
+        max_query_len: int,
         dcp_tot_seq_lens_device: torch.Tensor | None,
     ) -> AiterMLADecodeMetadata:
+        causal = self._decode_causal
         device = self.device
         num_reqs = seq_lens_device.size(0)
         qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
@@ -962,12 +1096,15 @@ def _build_decode(
             self._decode_num_heads,
             int(max_qo_len),
             self._kv_cache_dtype_str,
+            self._kv_cache_bytes,
         )
         use_gluon_verify = AiterMLAHelper.use_gluon_verify(
             self._decode_num_heads,
             int(max_qo_len),
             self._kv_cache_dtype_str,
             self.dcp_world_size,
+            causal,
+            self._kv_cache_bytes,
         )
         use_segmented_dcp_verify = (
             self._supports_segmented_dcp_verify and max_qo_len > 1
@@ -1045,14 +1182,36 @@ def _build_decode(
             # A padded rank has no bf16 persistent kernel past qlen 4 where the
             # gfx950 fold is absent; the non-persistent entry covers it. fp8
             # keeps the schedule -- its fold rejects non-persistent outright.
+            # A non-causal block keeps it too: what the fold drops past qlen 4
+            # is the block's causal staircase, which a non-causal block does
+            # not have, and the schedule is the only thing carrying its mask.
             and (
-                self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
+                not causal
+                or self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
                 or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
                 or is_quantized_kv_cache(self._kv_cache_dtype_str)
             )
             and max_qo_len >= 1
             and max_qo_len <= self._mtp_decode_qlen
         )
+        if (
+            not causal
+            and max_qo_len == 2
+            and is_quantized_kv_cache(self._kv_cache_dtype_str)
+            and not AiterMLAHelper.has_fp8_non_causal_qlen2_kernel(
+                self._decode_num_heads
+            )
+        ):
+            # AITER's fp8 dispatch folds (gqa 16, qlen 3 or 4) onto the
+            # qseqlen-4 kernel but never lists 2. Only 32/64/96/128 heads
+            # at qlen 2 fold onto that 16-head / 4-token non-causal kernel;
+            # 16-head qlen 2, and padded 48/80/112, keep Q2 and abort the
+            # process rather than raising. The bf16 fold has no such hole.
+            raise ValueError(
+                "AITER has no non-causal fp8 MLA kernel for this 2-token "
+                "query block. Pin the draft to TRITON_MLA for this "
+                "speculative config."
+            )
         if use_persistent_metadata:
             from aiter import get_mla_metadata_v1
 
@@ -1065,7 +1224,7 @@ def _build_decode(
                 paged_kv_last_page_len,
                 self._num_attention_heads,
                 1,
-                True,
+                causal,
                 self._mla_work_meta_data,
                 self._mla_work_info_set,
                 self._mla_work_indptr,
@@ -1135,6 +1294,11 @@ def build(
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
     ) -> AiterMLAMetadata:
+        # The common MLA builder owns the causal flag but its _build_decode
+        # contract is shared with CUDA backends. Keep this ROCm-specific state
+        # on the per-ubatch Aiter builder instead of changing that shared API;
+        # super().build invokes _build_decode synchronously.
+        self._decode_causal = common_attn_metadata.causal
         attn_metadata = super().build(
             common_prefix_len, common_attn_metadata, fast_build
         )
@@ -1208,8 +1372,7 @@ def _expand_page_indices_kernel(
 
 
 class AiterMLAHelper:
-    """
-    AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned
+    """AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned
     head counts through 128 are padded to the next multiple of 16 by tiling the
     query heads and slicing to the padded size. Native H24 AITER builds bypass
     that padding. Small divisors of 16 retain the existing repeat_interleave and
@@ -1224,6 +1387,10 @@ class AiterMLAHelper:
     # fold that reaches a persistent one is gfx950-only.
     _ASM_PADDED_MAX_PS_QLEN: Final = 4
     _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()
+    # Pinned AITER v0.1.21.post2 folds these fp8 qlen-2 counts onto the
+    # non-causal 16-head / 4-token kernel. Other multiples of 16 (48/80/112)
+    # fold to H16 while keeping Q2, which has no non-causal fp8 entry.
+    _AITER_FP8_NON_CAUSAL_QLEN2_HEADS: ClassVar[tuple[int, ...]] = (32, 64, 96, 128)
 
     @staticmethod
     def qo_indptr_for_uniform_qlen(
@@ -1261,6 +1428,18 @@ def is_valid_num_heads(num_heads: int) -> bool:
             )
         )
 
+    @staticmethod
+    def has_fp8_non_causal_qlen2_kernel(num_heads: int) -> bool:
+        """Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel.
+
+        Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the
+        non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2
+        block has no matching entry, and padded 48/80/112 keep Q2 after the
+        H16 fold.
+        """
+        kernel_heads = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
+        return kernel_heads in AiterMLAHelper._AITER_FP8_NON_CAUSAL_QLEN2_HEADS
+
     @staticmethod
     def get_actual_mla_num_heads(num_heads: int) -> int:
         if num_heads == 24 and _aiter_mla_native_h24_supported():
@@ -1335,7 +1514,12 @@ def get_mla_unpadded_lse(num_heads: int, lse: torch.Tensor) -> torch.Tensor:
         return AiterMLAHelper._get_mla_unpadded_heads(num_heads, lse)
 
     @staticmethod
-    def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
+    def use_gluon_decode(
+        num_heads: int,
+        max_qo_len: int,
+        kv_cache_dtype: str,
+        kv_cache_bytes: int | None = None,
+    ) -> bool:
         # Small-head (<16) single-token decode takes either the Gluon kernel or
         # the padded asm persistent decode, selected by
         # VLLM_ROCM_AITER_MLA_ASM_PADDING and the arch (Gluon is gfx950 only).
@@ -1351,10 +1535,12 @@ def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bo
         mode = _aiter_mla_small_head_mode()
         if mode == "asm":
             return False
-        gluon_supported = _gluon_mla_decode_supported()
-        if mode == "gluon":
-            return gluon_supported
-        return m % num_heads == 0 and gluon_supported
+        if not _gluon_mla_decode_supported():
+            return False
+        if mode != "gluon" and m % num_heads != 0:
+            return False
+        # Last, so the size warning only fires where Gluon would have run.
+        return _gluon_kv_cache_in_bounds(kv_cache_bytes)
 
     @staticmethod
     def use_gluon_verify(
@@ -1362,6 +1548,8 @@ def use_gluon_verify(
         max_qo_len: int,
         kv_cache_dtype: str,
         dcp_world_size: int = 1,
+        causal: bool = True,
+        kv_cache_bytes: int | None = None,
     ) -> bool:
         """Whether a small-head multi-token verify uses native Gluon MTP.
 
@@ -1374,7 +1562,13 @@ def use_gluon_verify(
 
         DCP verify is excluded: its per-row causal windows are served by the
         segmented path, which Gluon's MTP entry cannot express.
+
+        Gluon masks the block causally with no way to turn it off, so a
+        non-causal block takes the padded asm decode whatever the head count
+        -- padding is what gives it a kernel there.
         """
+        if not causal:
+            return False
         if max_qo_len <= 1 or dcp_world_size > 1:
             return False
         if is_quantized_kv_cache(kv_cache_dtype):
@@ -1383,8 +1577,10 @@ def use_gluon_verify(
             return False
         if not _gluon_mla_decode_supported():
             return False
-        # Same arch and mode gating as use_gluon_decode.
-        return _aiter_mla_small_head_mode() != "asm"
+        # Same arch, mode and cache-size gating as use_gluon_decode.
+        if _aiter_mla_small_head_mode() == "asm":
+            return False
+        return _gluon_kv_cache_in_bounds(kv_cache_bytes)
 
     @staticmethod
     def dcp_local_verify_row_lens(
@@ -1972,6 +2168,7 @@ def forward_mqa(
                 decode.paged_kv_indptr,
                 decode.paged_kv_indices,
                 decode.paged_kv_last_page_len,
+                causal=attn_metadata.causal,
                 **mla_kwargs,
             )
 
diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
index 3a9a4c92b350..f6258759d082 100644
--- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
@@ -15,8 +15,12 @@
 from vllm.model_executor.layers.attention.mla_attention import (
     get_mla_dims,
 )
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SharedTopkIndicesBuffer,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.torch_utils import np_to_pinned_tensor
 from vllm.v1.attention.backend import (
     AttentionBackend,
@@ -53,13 +57,16 @@ def _use_rocm_sparse_triton(
     num_decode_tokens: int,
     max_query_len: int,
 ) -> bool:
-    """Select the rope-free BF16 path not supported by AITER sparse MLA."""
-    plain_decode = num_decode_tokens == num_decodes
+    """Select the rope-free BF16 path not supported by AITER sparse MLA.
+
+    The ragged Triton kernel indexes metadata per query token, so multi-token
+    speculative verification rows have the same capability requirements as
+    plain decode rows.
+    """
     return (
         not kv_cache_dtype.startswith("fp8")
         and head_size == kv_lora_rank
-        and plain_decode
-        and (num_prefills > 0 or (num_decodes > 0 and max_query_len == 1))
+        and (num_prefills > 0 or num_decodes > 0)
     )
 
 
@@ -160,8 +167,7 @@ def triton_convert_req_index_to_global_index(
     NUM_TOPK_TOKENS: int = 2048,
     BLOCK_N: int = 128,  # tile width along columns
 ):
-    """
-    out[token_id, indice_id] =
+    """out[token_id, indice_id] =
         block_table[req_id[token_id],
             token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE
         + token_indices[token_id, indice_id] % BLOCK_SIZE
@@ -270,48 +276,6 @@ def generate_sparse_seqlen_triton(
     return out
 
 
-@triton.jit
-def fetch_id_to_ragged_kernel(
-    in_tensor_ptr,  # [num_seq, topk]
-    cumsum_ptr,  # [num_seq + 1]
-    out_tensor_ptr,  # [max_num_seq * topk]
-    in_tensor_ptr_stride,
-    TOPK: tl.constexpr,
-    TOKEN_NUM: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    seq_id = tl.program_id(0)
-    block_id = tl.program_id(1)
-    offset = tl.arange(0, BLOCK_SIZE)
-    token_start = tl.load(cumsum_ptr + seq_id)
-    token_end = tl.load(cumsum_ptr + seq_id + 1)
-    token_num = token_end - token_start
-    row_offset = block_id * BLOCK_SIZE
-    if row_offset >= token_num:
-        return
-    in_tensor_offset = seq_id * in_tensor_ptr_stride + row_offset + offset
-    in_tensor_mask = (row_offset + offset) < TOPK
-    in_tensor_val = tl.load(in_tensor_ptr + in_tensor_offset, mask=in_tensor_mask)
-    out_tensor_offset = token_start + row_offset + offset
-    out_tensor_mask = (out_tensor_offset < token_end) & in_tensor_mask
-    tl.store(out_tensor_ptr + out_tensor_offset, in_tensor_val, mask=out_tensor_mask)
-
-
-def fetch_id_to_ragged_triton(
-    in_tensor: torch.Tensor, cumsum: torch.Tensor, out_tensor: torch.Tensor, topk
-):
-    num_tokens = in_tensor.size(0)
-    block_size = 64
-    num_block_per_row = triton.cdiv(topk, block_size)
-    grid = (
-        num_tokens,
-        num_block_per_row,
-    )
-    fetch_id_to_ragged_kernel[grid](
-        in_tensor, cumsum, out_tensor, in_tensor.stride(0), topk, num_tokens, block_size
-    )
-
-
 class ROCMAiterMLASparseBackend(AttentionBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -323,7 +287,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [1, MultipleOf(16)]
 
     @staticmethod
@@ -428,7 +392,7 @@ def __init__(
 
         self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
         self.mla_dims = get_mla_dims(self.model_config)
-        self.topk_tokens = vllm_config.model_config.hf_config.index_topk
+        self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk
         attention_context = vllm_config.compilation_config.static_forward_context
         # Sink decode must use AITER's nonpersistent path. In particular,
         # gfx942 has no persistent+LSE kernel, and its metadata heuristic
@@ -635,12 +599,14 @@ def build(
         reduce_partial_map = None
         if self._use_persistent_metadata and not use_triton_sparse:
             num_reqs = common_attn_metadata.num_reqs
+            with gpu_sync_allowed():
+                seq_lens_cpu = common_attn_metadata.seq_lens[:num_reqs].cpu().numpy()
             clamped_seq_lens = np.minimum(
-                common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(),
+                seq_lens_cpu,
                 self.topk_tokens,
             )
             clamped_context_lens = np.minimum(
-                common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() - seg_lengths,
+                seq_lens_cpu - seg_lengths,
                 self.topk_tokens,
             )
             metadata_key = (
@@ -744,7 +710,9 @@ def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor:
     return (result, lse)
 
 
-class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
+class ROCMAiterMLASparseImpl(
+    MLAAttentionImpl[ROCMAiterMLASparseMetadata], SharedTopkIndicesBuffer
+):
     is_sparse = True
     supports_dense_mha_prefill = False
     supports_dcp = False
@@ -789,12 +757,7 @@ def __init__(
         self.kv_cache_dtype = kv_cache_dtype
         self.kv_lora_rank: int = mla_args["kv_lora_rank"]
         self.softmax_scale = scale
-        # The indexer carries the shared buffer for normal layers and tests;
-        # the explicitly-passed buffer covers backbone skip layers, whose
-        # indexer is not constructed (see deepseek_v2.py).
-        self.topk_indices_buffer: torch.Tensor | None = (
-            indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
-        )
+        self.init_topk_indices_buffer(indexer, topk_indices_buffer)
 
         vllm_config = get_current_vllm_config()
         max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -803,6 +766,11 @@ def __init__(
             (q_concat_shape, vllm_config.model_config.dtype),
         )
 
+    def record_logical_topk_ready(self) -> None:
+        # This impl shares the top-k indices buffer via SharedTopkIndicesBuffer
+        # but does not participate in sparse-MLA index groups.
+        pass
+
     def _forward_mla(
         self,
         layer: AttentionLayer,
diff --git a/vllm/v1/attention/backends/mla/sparse_indexer.py b/vllm/v1/attention/backends/mla/sparse_indexer.py
new file mode 100644
index 000000000000..c0b445534d22
--- /dev/null
+++ b/vllm/v1/attention/backends/mla/sparse_indexer.py
@@ -0,0 +1,235 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V4.1 indexer backend for the DeepGEMM sparse MQA-logits path.
+
+Candidate-consuming indexer layers score only the candidate blocks published
+by the candidate-source layer (``fp8_fp4_(paged_)sparse_mqa_logits``) instead
+of dense logits over the whole context. This backend extends the dense
+indexer metadata with the per-row scratch that path needs, so the layer op
+allocates nothing per step, and lets all consumer layers of a step share one
+candidate expansion and DeepGEMM schedule (they all read the same candidate
+blocks). Opt in with ``AttentionConfig.indexer_sparse_logits``.
+"""
+
+from dataclasses import dataclass, fields
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
+    check_deep_select_layout,
+    has_deep_select,
+    pick_sparse_block_kv,
+)
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV32IndexerMetadata,
+    DeepseekV32IndexerMetadataBuilder,
+    DeepseekV41IndexerBackend,
+    dsa_indexer_uses_fp4,
+)
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class SparseMQARowsMetadata:
+    """Per-row state for one sparse-logits call: the decode rows, or one
+    prefill chunk. Buffers are builder-owned views sized to the rows.
+
+    ``kernel_metadata`` (the DeepGEMM schedule) is filled by the first
+    consumer indexer layer of the step together with ``sparse_indices`` /
+    ``end``; later consumer layers reuse all three.
+    """
+
+    row_ks: torch.Tensor
+    """[rows] int32 K-range start per row; zeros for paged decode."""
+    row_ke: torch.Tensor
+    """[rows] int32 K-range end (prefill) or compressed context length."""
+    sparse_indices: torch.Tensor
+    """[rows, num_sparse_blocks] int32 DeepGEMM sparse block ids."""
+    end: torch.Tensor
+    """[rows] int32 valid sparse-column count per row."""
+    col_indices: torch.Tensor
+    """[rows, topk] int32 scratch for the sparse-column top-k."""
+    block_table: torch.Tensor | None = None
+    """Decode only: [rows, pages] per-row page table."""
+    row_indices: torch.Tensor | None = None
+    """Decode only: [rows] row -> request index."""
+    kernel_metadata: torch.Tensor | None = None
+
+
+@dataclass
+class DeepseekV41SparseIndexerMetadata(DeepseekV32IndexerMetadata):
+    sparse_block_kv: int = 0
+    sparse_decode: SparseMQARowsMetadata | None = None
+    sparse_prefill: list[SparseMQARowsMetadata] | None = None
+    """Parallel to ``prefill.chunks``."""
+
+
+class DeepseekV41SparseIndexerBackend(DeepseekV41IndexerBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "DEEPSEEK_V41_SPARSE_INDEXER"
+
+    @staticmethod
+    def get_builder_cls() -> type["DeepseekV41SparseIndexerMetadataBuilder"]:
+        return DeepseekV41SparseIndexerMetadataBuilder
+
+
+class DeepseekV41SparseIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder):
+    """Dense indexer metadata plus the sparse-logits row scratch.
+
+    Every requirement of the sparse kernels is checked here, at engine start,
+    rather than falling back to the dense path per step.
+    """
+
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+        **kwargs,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs)
+        hf_config = vllm_config.model_config.hf_text_config
+        candidate_block_size = getattr(hf_config, "candidate_block_size", 0)
+        candidate_topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0)
+        topk_tokens = getattr(hf_config, "index_topk", 0)
+        prefix = "attention_config.indexer_sparse_logits"
+        if candidate_block_size <= 0 or candidate_topk_blocks <= 0:
+            raise ValueError(
+                f"{prefix} requires a two-level (candidate block) indexer model "
+                "such as DeepSeek-V4.1."
+            )
+        if not dsa_indexer_uses_fp4(vllm_config):
+            raise ValueError(
+                f"{prefix} requires the MXFP4 indexer cache "
+                "(attention_config.indexer_kv_dtype='mxfp4'): the sparse kernels "
+                "take UE8M0-packed granularity-32 scales."
+            )
+        if not (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(100)
+        ):
+            raise ValueError(f"{prefix} requires an SM100-class GPU.")
+        if not has_deep_gemm_sparse_mqa():
+            raise ValueError(
+                f"{prefix} requires DeepGEMM >= 2.8 "
+                "(fp8_fp4_sparse_mqa_logits not found)."
+            )
+        if self.dcp_world_size > 1 or self.use_pcp:
+            raise NotImplementedError(
+                f"{prefix} is not supported with decode/prefill context parallel."
+            )
+        if not self.supports_varlen:
+            raise ValueError(
+                f"{prefix} requires the varlen paged MQA logits decode layout."
+            )
+        # The top-k runs on the kernels' bf16 logits with DeepSelect; there is
+        # no fp32 fallback, so the other top-k backends cannot be honored.
+        topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend
+        if not has_deep_select():
+            raise ValueError(
+                f"{prefix} requires the DeepSelect top-k extension "
+                "(vllm._deepselect_C) for this GPU."
+            )
+        if topk_backend not in ("auto", "deep_select"):
+            raise ValueError(
+                f"{prefix} uses DeepSelect for the top-k; "
+                f"kernel_config.sparse_indexer_topk_backend='{topk_backend}' "
+                "is not supported with it."
+            )
+        self.sparse_block_kv = pick_sparse_block_kv(candidate_block_size)
+        if kv_cache_spec.block_size % self.sparse_block_kv != 0:
+            raise ValueError(
+                f"{prefix}: indexer cache pages ({kv_cache_spec.block_size} "
+                f"tokens) must be a multiple of {self.sparse_block_kv} tokens."
+            )
+        self.num_sparse_blocks = candidate_topk_blocks * (
+            candidate_block_size // self.sparse_block_kv
+        )
+        num_sparse_cols = self.num_sparse_blocks * self.sparse_block_kv
+        check_deep_select_layout(num_sparse_cols, topk_tokens)
+        # Sparse logits are bf16 [rows, num_sparse_cols]; make the prefill
+        # chunker (which budgets fp32 [rows, seq_len]) size chunks as if every
+        # row were this wide.
+        self.min_split_seq_len = (num_sparse_cols * 2 + 3) // 4
+
+        # Rows are indexed by batch token (decode rows first, then prefill
+        # chunks by token range), like ``arange_buffer``.
+        max_rows = self.arange_buffer.shape[0]
+        int32 = dict(dtype=torch.int32, device=device)
+        self.sparse_indices_buffer = torch.zeros(
+            (max_rows, self.num_sparse_blocks), **int32
+        )
+        self.end_buffer = torch.zeros(max_rows, **int32)
+        self.col_indices_buffer = torch.zeros((max_rows, topk_tokens), **int32)
+        self.zero_row_ks = torch.zeros(max_rows, **int32)
+        logger.info_once(
+            "DeepSeek V4.1 indexer: candidate consumers score only the candidate "
+            "blocks with DeepGEMM sparse MQA logits (%d sparse blocks of %d "
+            "tokens per row).",
+            self.num_sparse_blocks,
+            self.sparse_block_kv,
+        )
+
+    def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor:
+        return seq_lens_cpu.clamp(min=self.min_split_seq_len)
+
+    def _rows(self, start: int, end: int, row_ks, row_ke, **kwargs):
+        return SparseMQARowsMetadata(
+            row_ks=row_ks,
+            row_ke=row_ke,
+            sparse_indices=self.sparse_indices_buffer[start:end],
+            end=self.end_buffer[start:end],
+            col_indices=self.col_indices_buffer[start:end],
+            **kwargs,
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> DeepseekV41SparseIndexerMetadata:
+        base = super().build(common_prefix_len, common_attn_metadata, fast_build)
+        metadata = DeepseekV41SparseIndexerMetadata(
+            **{f.name: getattr(base, f.name) for f in fields(base)},
+            sparse_block_kv=self.sparse_block_kv,
+        )
+        if base.decode is not None:
+            decode = base.decode
+            rows = decode.seq_lens.shape[0]
+            # The varlen layout flattens every query to one row, so the base
+            # decode tensors are already per row.
+            assert decode.seq_lens.shape[1] == 1 and not decode.requires_padding
+            assert decode.block_table.shape[0] == rows
+            metadata.sparse_decode = self._rows(
+                0,
+                rows,
+                row_ks=self.zero_row_ks[:rows],
+                row_ke=decode.seq_lens.view(-1),
+                block_table=decode.block_table,
+                row_indices=(
+                    decode.indices
+                    if decode.indices is not None
+                    else self.arange_buffer[:rows]
+                ),
+            )
+        if base.prefill is not None:
+            metadata.sparse_prefill = [
+                self._rows(
+                    chunk.token_start,
+                    chunk.token_end,
+                    row_ks=chunk.cu_seqlen_ks,
+                    row_ke=chunk.cu_seqlen_ke,
+                )
+                for chunk in base.prefill.chunks
+            ]
+        return metadata
diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py
index 02e6336a3835..78727c2eb561 100644
--- a/vllm/v1/attention/backends/mla/sparse_swa.py
+++ b/vllm/v1/attention/backends/mla/sparse_swa.py
@@ -9,13 +9,13 @@
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.warmup.jit_warmup import (
     WarmupIntRange,
+    kernel_launcher,
     zip_inputs,
 )
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
@@ -49,7 +49,7 @@
 # v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but
 # their compressed page block is block_size // ratio, so each gets its own plan.
 # v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see
-# deepseek_v4_1/sparse_mla.py.
+# deepseek_v41/sparse_mla.py.
 _LAYER_TYPE_C1A = "c1a"
 _LAYER_TYPE_C2A = "c2a"
 
@@ -79,15 +79,22 @@ def __init__(
         cache_config: CacheConfig,
         backend_cls: "type[AttentionBackend] | None" = None,
         block_size: int = 64,
+        packed_bytes_per_token: int = 584,
+        packed_page_alignment: int = 576,
+        bounded_replay: bool = False,
     ):
         super().__init__()
         self.backend_cls = backend_cls or DeepseekSparseSWABackend
+        # DeepseekV4.1 SWA bounded replay.
+        self.bounded_replay = bounded_replay
         self.kv_cache = torch.tensor([])
         self.head_dim = head_dim
         self.window_size = window_size
         self.prefix = prefix
         self.cache_config = cache_config
         self.dtype = dtype
+        self.packed_bytes_per_token = packed_bytes_per_token
+        self.packed_page_alignment = packed_page_alignment
         compilation_config = get_current_vllm_config().compilation_config
         if prefix in compilation_config.static_forward_context:
             raise ValueError(f"Duplicate layer name: {prefix}")
@@ -105,20 +112,26 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
         self.kv_cache = kv_cache.squeeze(1)
 
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
-        # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous
-        # bf16/fp8 cache uses the natural element-size page.
-        uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla"
+        # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode
+        # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural
+        # element-size page.
+        uses_fp8_ds_mla_layout = self.cache_config.cache_dtype in (
+            "fp8_ds_mla",
+            "nvfp4_ds_mla",
+        )
         return SlidingWindowMLASpec(
+            bounded_replay=self.bounded_replay,
             block_size=self.block_size,
             num_kv_heads=1,
             head_size=self.head_dim,
             dtype=self.dtype,
             sliding_window=self.window_size,
             cache_dtype_str=self.cache_config.cache_dtype,
-            # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales)
-            state_content_bytes=584 if uses_fp8_ds_mla_layout else None,
-            # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
-            alignment=576 if uses_fp8_ds_mla_layout else 512,
+            state_content_bytes=(
+                self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None
+            ),
+            # FlashMLA packing stride; 512B for FlashInfer sparse (#44577).
+            alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512,
             model_version="deepseek_v4",
             kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
         )
@@ -135,7 +148,7 @@ def get_name() -> str:
         return "DEEPSEEK_SPARSE_SWA"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(32)]
 
     @classmethod
@@ -189,6 +202,10 @@ class DeepseekSparseSWAMetadata:
     # None when the model is text-only or the batch has no image spans.
     prefill_left_visible: torch.Tensor | None = None
     prefill_right_visible: torch.Tensor | None = None
+    # SWA bounded replay: [num_reqs] lower bound of the positions a request's
+    # window attention may read (SWA bounded replay);
+    # zeros when nothing replays.
+    replay_start: torch.Tensor | None = None
 
     # Number of decode/prefill requests/tokens (batch is reordered: decodes first)
     num_decodes: int = 0
@@ -196,6 +213,8 @@ class DeepseekSparseSWAMetadata:
     num_decode_tokens: int = 0
     num_prefill_tokens: int = 0
     max_decode_query_len: int = 1
+    flashinfer_decode_topk_lens: torch.Tensor | None = None
+    flashinfer_decode_seq_lens: torch.Tensor | None = None
 
     # Pre-computed prefill metadata shared across all DeepseekV4 attention layers.
     prefill_seq_lens: torch.Tensor | None = None
@@ -245,10 +264,12 @@ def get_prefill_chunk_plan(
             has_compressed = compress_ratio > 1
 
         # query_len <= max_num_batched_tokens and
-        # gather_len = query_len + min(prefix_len, window_size - 1), so the
-        # worst-case gathered width is bounded by
-        # max_num_batched_tokens + window_size - 1. The compressed prefix pool
-        # is bounded by ceil(max_model_len / compress_ratio).
+        # gather_len <= query_len + min(prefix_len, window_size - 1) (a replay
+        # boundary only shrinks it), so the worst-case gathered width is
+        # bounded by max_num_batched_tokens + window_size - 1. The compressed
+        # prefix pool is bounded by ceil(max_model_len / compress_ratio). The
+        # CPU gather lengths below are this upper bound, for chunk planning
+        # only; the kernels use the exact per-request values.
         max_workspace_area = prefill_chunk_size * (
             (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0)
             + self.prefill_window_size
@@ -320,6 +341,7 @@ def kernel(
         # Inputs
         seq_lens_ptr,
         query_start_loc_ptr,
+        replay_start_ptr,
         num_prefills,
         num_decodes,
         window_size,
@@ -341,7 +363,10 @@ def kernel(
 
         query_len = qsl_end - qsl_start
         prefix_len = seq_len - query_len
-        gather_len = query_len + tl.minimum(prefix_len, window_size - 1)
+        # Context below replay_start has no window KV (SWA bounded replay).
+        replay_start = tl.load(replay_start_ptr + num_decodes + safe_offset, mask=mask)
+        visible_prefix = tl.maximum(prefix_len - replay_start, 0)
+        gather_len = query_len + tl.minimum(visible_prefix, window_size - 1)
 
         tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
 
@@ -373,6 +398,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
             prefill_gather_lens=int32_ptr,
             seq_lens=int32_ptr,
             query_start_loc=int32_ptr,
+            replay_start=int32_ptr,
             num_prefills=compile_key.block_size,
             num_decodes=0,
             window_size=1,
@@ -384,6 +410,7 @@ def __call__(
         prefill_gather_lens: torch.Tensor,
         seq_lens: torch.Tensor,
         query_start_loc: torch.Tensor,
+        replay_start: torch.Tensor,
         num_prefills: int,
         num_decodes: int,
         window_size: int,
@@ -443,13 +470,15 @@ def __init__(self, *args, **kwargs):
         assert hasattr(hf_config, "sliding_window")
         self.window_size = hf_config.sliding_window
 
-        # Vision variant: image spans (up to vision_max_n_token tokens) are
+        # V4 vision variant: image spans (up to vision_max_n_token tokens) are
         # visible bidirectionally, so prefill index rows widen from
-        # window_size to window_size + max_image_tokens. Text-only models keep
-        # max_image_tokens == 0 and take the original code paths everywhere.
+        # window_size to window_size + max_image_tokens. The V4 config sets
+        # mm_prefix_clamp_sliding_window exactly for these in-kernel-widened
+        # ranges; V4.1 image tokens use the plain causal window, and text-only
+        # models keep max_image_tokens == 0 everywhere.
         self.max_image_tokens = (
             getattr(hf_config, "vision_max_n_token", 0)
-            if getattr(hf_config, "vision_n_layers", 0) > 0
+            if getattr(hf_config, "mm_prefix_clamp_sliding_window", False)
             else 0
         )
         self.prefill_index_width = self.window_size + self.max_image_tokens
@@ -537,12 +566,19 @@ def __init__(self, *args, **kwargs):
         )
         self.decode_swa_indices_noncausal: torch.Tensor | None = None
         self._max_tokens = max_tokens
+        # replay_start when no group replays: no lower bound.
+        self.no_replay_start = torch.zeros(
+            self.vllm_config.scheduler_config.max_num_seqs + 1,
+            dtype=torch.int32,
+            device=self.device,
+        )
 
     def build(
         self,
         common_prefix_len: int,
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
+        replay_start: torch.Tensor | None = None,
     ) -> DeepseekSparseSWAMetadata:
         """Build SWA metadata for mixed decode/prefill batches.
 
@@ -551,6 +587,10 @@ def build(
         separate window_topk_idxs for each portion.
 
         For prefill, we use chunked prefill to align with the indexer's chunking.
+
+        ``replay_start`` ([num_reqs], SWA bounded replay): the position from
+        which each request holds window KV; the model state passes it for every
+        batch of a replaying group, graph captures pass nothing.
         """
         seq_lens = common_attn_metadata.seq_lens
         seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
@@ -573,7 +613,10 @@ def build(
         )
 
         is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
-        is_valid_token.copy_(slot_mapping >= 0)
+        torch.ge(slot_mapping, 0, out=is_valid_token)
+        if replay_start is None:
+            # Graph captures build without it; their dummy batches replay nothing.
+            replay_start = self.no_replay_start
 
         non_causal = not common_attn_metadata.causal
         decode_swa_width = (
@@ -596,6 +639,8 @@ def build(
                         device=self.device,
                     )
                 decode_swa_indices = self.decode_swa_indices_noncausal
+                # No replay_start: decode rows sit above the hit, so a
+                # request's replayed window never reaches them.
                 _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL(
                     decode_swa_indices,
                     self.decode_swa_lens,
@@ -624,6 +669,7 @@ def build(
                     is_valid_token,
                     block_table,
                     self.block_size,
+                    replay_start,
                     num_tokens=num_decode_tokens,
                     token_offset=0,
                 )
@@ -670,6 +716,7 @@ def build(
                 is_valid_token,
                 block_table,
                 self.block_size,
+                replay_start,
                 num_tokens=num_prefill_tokens,
                 token_offset=num_decode_tokens,
                 has_image=has_image,
@@ -683,6 +730,7 @@ def build(
             seq_lens_cpu,
             query_start_loc,
             query_start_loc_cpu,
+            replay_start,
         )
 
         # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta
@@ -714,6 +762,7 @@ def build(
             ),
             prefill_left_visible=prefill_left_visible,
             prefill_right_visible=prefill_right_visible,
+            replay_start=replay_start,
             block_size=self.block_size,
             num_decodes=num_decodes,
             num_prefills=num_prefills,
@@ -790,6 +839,7 @@ def update_draft_decode_metadata(
         assert metadata.is_valid_token is not None
         assert metadata.decode_swa_indices is not None
         assert metadata.decode_swa_lens is not None
+        assert metadata.replay_start is not None
 
         _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
             metadata.decode_swa_indices,
@@ -804,6 +854,7 @@ def update_draft_decode_metadata(
             metadata.is_valid_token,
             metadata.block_table,
             self.block_size,
+            metadata.replay_start,
             num_tokens=metadata.num_decode_tokens,
             token_offset=0,
         )
@@ -859,6 +910,7 @@ def _build_deepseek_v4_metadata(
         seq_lens_cpu: torch.Tensor | None,
         query_start_loc: torch.Tensor,
         query_start_loc_cpu: torch.Tensor,
+        replay_start: torch.Tensor,
     ) -> dict[str, torch.Tensor | int | None]:
         """Pre-compute DeepseekV4 prefill metadata during the metadata build phase.
 
@@ -880,6 +932,7 @@ def _build_deepseek_v4_metadata(
                 pfx_gather_lens,
                 seq_lens,
                 query_start_loc,
+                replay_start,
                 num_prefills,
                 num_decodes,
                 self.window_size,
@@ -969,6 +1022,7 @@ def _compute_swa_indices_and_lens_kernel(
     block_table_ptr,
     block_table_stride,
     block_size,
+    replay_start_ptr,
     token_offset,
     HAS_IMAGE: tl.constexpr,
     TRITON_BLOCK_SIZE: tl.constexpr,
@@ -1009,6 +1063,8 @@ def _compute_swa_indices_and_lens_kernel(
         right = 0
     left_add = tl.maximum(left - (window_size - 1), 0)
     start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0)
+    # SWA bounded replay: no window KV exists below the request's replay start.
+    start_pos = tl.maximum(start_pos, tl.load(replay_start_ptr + req_idx))
     end_pos = pos + right + 1
 
     swa_len = end_pos - start_pos
@@ -1098,6 +1154,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
             is_valid_token=TritonWarmupTensor(torch.bool),
             block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
             block_size=compile_key.block_size,
+            replay_start=int32_ptr,
             num_tokens=1,
             token_offset=0,
             has_image=bool(compile_key.has_image),
@@ -1118,6 +1175,7 @@ def __call__(
         is_valid_token: torch.Tensor,
         block_table: torch.Tensor,
         block_size: int,
+        replay_start: torch.Tensor,
         *,
         num_tokens: int,
         token_offset: int,
diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py
index 83e0e780e4c0..2ec10218d99f 100644
--- a/vllm/v1/attention/backends/mla/sparse_utils.py
+++ b/vllm/v1/attention/backends/mla/sparse_utils.py
@@ -5,18 +5,19 @@
 from dataclasses import dataclass
 from typing import Any
 
+import numpy as np
 import torch
 
 from vllm.config import VllmConfig
 from vllm.distributed import get_dcp_group
 from vllm.model_executor.warmup.jit_warmup import (
+    kernel_launcher,
     zip_inputs,
 )
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
     triton_scalar_specialization_rep,
 )
 from vllm.triton_utils import tl, triton
@@ -24,6 +25,21 @@
 from vllm.v1.worker.block_table import get_block_table_width
 
 
+def request_row_bounds(req_idx: np.ndarray) -> np.ndarray:
+    """Bounds of the runs of adjacent rows that belong to one request: run
+    ``r`` is rows ``[bounds[r], bounds[r + 1])``.
+
+    Under PCP a rank holds two adjacent chunk rows of a split prefill; the
+    sparse backends give such a run one KV region.
+    """
+    assert req_idx.size > 0
+    bounds = np.flatnonzero(np.r_[True, req_idx[1:] != req_idx[:-1], True])
+    assert bounds.size - 1 == np.unique(req_idx).size, (
+        "rows of one request must be adjacent"
+    )
+    return bounds
+
+
 def flat_kv_row_view(
     kv_cache: torch.Tensor,  # [num_blocks, block_size, head_dim]
     block_size: int,
@@ -69,8 +85,8 @@ class CompileKey:
         block_table_stride: int
 
     @staticmethod
-    @triton.jit(do_not_specialize=["max_num_blocks_per_req"])
-    def kernel(
+    @triton.jit(do_not_specialize=["max_num_blocks_per_req", "workspace_rank_stride"])
+    def sparse_mla_index_remap_kernel(
         req_id_ptr,  # int32 [num_tokens]
         block_table_ptr,  # int32 [num_requests, max_num_blocks_per_req]
         token_indices_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
@@ -78,20 +94,22 @@ def kernel(
         valid_count_ptr,  # int32 [num_tokens] - output valid count per row
         prefill_request_id_ptr,  # int32 [num_tokens], -1 for decode, >=0 for prefill
         workspace_starts_ptr,  # int32 [num_prefill_reqs+1] or nullptr
+        workspace_rank_stride,  # rows per DCP rank of a gathered prefill workspace
         # shapes (compile-time where possible)
         max_num_blocks_per_req,
         BLOCK_SIZE: tl.constexpr,
         BLOCK_STRIDE_ROWS: tl.constexpr,
         BLOCK_N: tl.constexpr,  # tile width along columns
+        NUM_TOPK_TOKENS: tl.constexpr,
         HAS_PREFILL: tl.constexpr,
         COUNT_VALID: tl.constexpr,  # whether to count valid indices
-        # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count
+        # BLOCK_N >= NUM_TOPK_TOKENS: one program owns the query, so its count
         # is an in-register reduction and needs no atomic.
         SINGLE_TILE: tl.constexpr,
         # When set, scatter valid slots to a contiguous prefix [0, valid_count) using
         # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
         # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).
-        # Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the
+        # Requires COUNT_VALID; multi-tile output must be pre-filled with -1. The
         # prefix is unspecified (only the selected set matters).
         COMPACT_TO_FRONT: tl.constexpr,
         # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op
@@ -119,7 +137,7 @@ def kernel(
 
         # Load token indices for this tile
         ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
-        tok = tl.load(ti_ptr)  # int32
+        tok = tl.load(ti_ptr, mask=indice_id < NUM_TOPK_TOKENS, other=-1)
 
         # Only token == -1 should propagate as -1
         is_invalid_tok = tok < 0
@@ -146,7 +164,7 @@ def kernel(
         # Guard block_table access
         valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
         bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
-        is_invalid_tok |= ~valid_block | is_remote
+        is_invalid_tok |= ~valid_block | (is_remote & ~is_prefill)
         base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
         out_val = base * BLOCK_STRIDE_ROWS + inblock_off
 
@@ -155,7 +173,11 @@ def kernel(
             workspace_start = tl.load(
                 workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0
             )
-            prefill_out = workspace_start + tok
+            # Under DCP the prefill workspace is the all-gather of every rank's
+            # shard, rank-major. With DCP_SIZE == 1 this is workspace_start + tok.
+            prefill_out = (
+                owning_rank * workspace_rank_stride + workspace_start + local_idx
+            )
             out_val = tl.where(is_prefill, prefill_out, out_val)
         out_val = tl.where(is_invalid_tok, -1, out_val)
 
@@ -176,10 +198,19 @@ def kernel(
             dest = base + local_offset
             out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
             tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
+            if SINGLE_TILE:
+                # This program owns both the compacted prefix and its padding.
+                # The stores are disjoint, so no separate initialization is needed.
+                tl.store(
+                    out_ptr + token_id * out_stride0 + indice_id * out_stride1,
+                    -1,
+                    mask=(indice_id >= tile_valid_count)
+                    & (indice_id < NUM_TOPK_TOKENS),
+                )
         else:
             # Store results in place (input column == output column).
             out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
-            tl.store(out_ptr_ij, out_val)
+            tl.store(out_ptr_ij, out_val, mask=indice_id < NUM_TOPK_TOKENS)
 
             # Accumulate the tile's valid count into the row total; a single tile's
             # reduction *is* the total.
@@ -190,6 +221,9 @@ def kernel(
                 else:
                     tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
 
+    # Keep the warmup wrapper interface while exposing a descriptive CUDA symbol.
+    kernel = sparse_mla_index_remap_kernel
+
     def dispatch(  # type: ignore[override]
         self,
         *,
@@ -232,12 +266,28 @@ def get_warmup_keys(
         dcp_size = vllm_config.parallel_config.decode_context_parallel_size
         dcp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
         dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0
-        num_topk_tokens = vllm_config.model_config.hf_config.index_topk
+        num_topk_tokens = vllm_config.model_config.hf_text_config.index_topk
         max_num_blocks = cdiv(
             vllm_config.model_config.max_model_len,
             block_size * dcp_size,
         )
         max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size)
+        pcp_dcp_combos = []
+        if (
+            vllm_config.parallel_config.prefill_context_parallel_size > 1
+            and dcp_size > 1
+        ):
+            # PCP + DCP attends prefill rows over a DCP-gathered workspace.
+            pcp_dcp_combos.append(
+                dict(
+                    HAS_PREFILL_WORKSPACE=True,
+                    COUNT_VALID=True,
+                    COMPACT_TO_FRONT=True,
+                    DCP_SIZE=dcp_size,
+                    DCP_RANK=dcp_rank,
+                    DCP_INTERLEAVE=dcp_interleave,
+                )
+            )
         return self._trace_dispatch(self.dispatch)(
             zip_inputs(
                 dict(
@@ -280,6 +330,7 @@ def get_warmup_keys(
                     DCP_RANK=dcp_rank,
                     DCP_INTERLEAVE=dcp_interleave,
                 ),
+                *pcp_dcp_combos,
             ),
             BLOCK_SIZE=block_size,
             BLOCK_STRIDE_ROWS=block_stride_rows,
@@ -316,6 +367,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
                 int32_ptr if compile_key.has_prefill_workspace else None
             ),
             max_num_blocks_per_req=1,
+            workspace_rank_stride=1,
             BLOCK_SIZE=compile_key.block_size,
             BLOCK_STRIDE_ROWS=compile_key.block_stride_rows,
             BLOCK_N=compile_key.block_n,
@@ -339,6 +391,7 @@ def __call__(
         prefill_workspace_request_ids: torch.Tensor | None,
         prefill_workspace_starts: torch.Tensor | None,
         *,
+        workspace_rank_stride: int,
         max_num_blocks_per_req: int,
         BLOCK_SIZE: int,
         BLOCK_STRIDE_ROWS: int,
@@ -381,17 +434,19 @@ def _remap_tiling(
     Counting the valid slots per row is the only reason the column tiles have to
     talk to each other, so when counting give one program the whole row: the
     count becomes an in-register reduction plus a plain store, needing neither
-    atomics nor a zero-initialized counter. The row is one ``tl.arange``, so this
-    needs a power-of-two width; other top-k sizes stay tiled and atomic.
+    atomics nor a zero-initialized counter. Pad modest non-power-of-two widths
+    (including GLM's 2176 entries) to one tile; larger widths stay tiled and atomic.
 
     Returns:
         (single_tile, block_n, tiles_per_row, num_warps)
+
     """
-    single_tile = (
-        count_valid and triton.next_power_of_2(NUM_TOPK_TOKENS) == NUM_TOPK_TOKENS
+    padded_width = triton.next_power_of_2(NUM_TOPK_TOKENS)
+    single_tile = count_valid and (
+        padded_width == NUM_TOPK_TOKENS or padded_width <= 4096
     )
     if single_tile:
-        return True, NUM_TOPK_TOKENS, 1, 8
+        return True, padded_width, 1, 8
     return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4
 
 
@@ -406,10 +461,15 @@ def triton_convert_req_index_to_global_index(
     HAS_PREFILL_WORKSPACE: bool = False,
     prefill_workspace_request_ids: torch.Tensor | None = None,
     prefill_workspace_starts: torch.Tensor | None = None,
+    prefill_workspace_rank_stride: int | None = None,
+    dcp_size: int = 1,
+    dcp_rank: int = 0,
+    cp_kv_cache_interleave_size: int = 1,
     return_valid_counts: bool = False,
+    out: torch.Tensor | None = None,
+    valid_counts_out: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    out[token_id, indice_id] =
+    """out[token_id, indice_id] =
         block_table[req_id[token_id],
             token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE
         + token_indices[token_id, indice_id] % BLOCK_SIZE
@@ -458,20 +518,37 @@ def triton_convert_req_index_to_global_index(
     req_id_c = req_id.contiguous()
     block_table_c = block_table.contiguous()
     token_indices_c = token_indices.contiguous()
-    # When return_valid_counts, the kernel scatters valid entries to a
-    # contiguous prefix [0, valid_count) and leaves the tail unwritten, so
-    # pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to
-    # [:topk_length] == exactly the valid set (no dropped tokens).
-    if return_valid_counts:
-        out = torch.full_like(token_indices_c, -1)
+    # Only multi-tile compaction needs separate padding initialization.
+    if out is None:
+        out = (
+            torch.full_like(token_indices_c, -1)
+            if return_valid_counts and not single_tile
+            else torch.empty_like(token_indices_c)
+        )
     else:
-        out = torch.empty_like(token_indices_c)
+        assert out.dtype == token_indices_c.dtype
+        assert out.device == token_indices_c.device
+        assert out.shape == token_indices_c.shape
+        assert out.is_contiguous()
+        if return_valid_counts and not single_tile:
+            out.fill_(-1)
 
     valid_counts: torch.Tensor | None = None
     if return_valid_counts:
-        # Zero-init only matters for the atomic accumulation path.
-        alloc = torch.empty if single_tile else torch.zeros
-        valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
+        if valid_counts_out is None:
+            # Zero-init only matters for the atomic accumulation path.
+            alloc = torch.empty if single_tile else torch.zeros
+            valid_counts = alloc(
+                num_tokens, dtype=torch.int32, device=token_indices.device
+            )
+        else:
+            assert valid_counts_out.dtype == torch.int32
+            assert valid_counts_out.device == token_indices.device
+            assert valid_counts_out.shape == (num_tokens,)
+            assert valid_counts_out.is_contiguous()
+            valid_counts = valid_counts_out
+            if not single_tile:
+                valid_counts.zero_()
 
     # Prepare prefill pointers
     if HAS_PREFILL_WORKSPACE:
@@ -489,6 +566,7 @@ def triton_convert_req_index_to_global_index(
         prefill_workspace_request_ids,
         prefill_workspace_starts,
         # shapes / constexprs
+        workspace_rank_stride=prefill_workspace_rank_stride or 0,
         max_num_blocks_per_req=max_num_blocks_per_req,
         BLOCK_SIZE=BLOCK_SIZE,
         BLOCK_STRIDE_ROWS=(
@@ -499,10 +577,9 @@ def triton_convert_req_index_to_global_index(
         HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
         COUNT_VALID=return_valid_counts,
         COMPACT_TO_FRONT=return_valid_counts,
-        # DCP disabled (no-op de-interleave)
-        DCP_SIZE=1,
-        DCP_RANK=0,
-        DCP_INTERLEAVE=1,
+        DCP_SIZE=dcp_size,
+        DCP_RANK=dcp_rank,
+        DCP_INTERLEAVE=cp_kv_cache_interleave_size,
     )
 
     if return_valid_counts:
@@ -570,13 +647,13 @@ def triton_filter_and_convert_dcp_index(
     token_indices_c = token_indices.contiguous()
 
     # The compaction uses the valid-count buffer as a slot allocator, so it
-    # requires counting. Pre-fill out with -1 so the unwritten tail stays -1.
+    # requires counting. Only the multi-tile path needs pre-filled padding.
     count_valid = return_valid_counts or compact_valid_to_front
 
     # The compaction builds on the counting, so it shares the tiling.
     single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid)
 
-    if compact_valid_to_front:
+    if compact_valid_to_front and not single_tile:
         out = torch.full_like(token_indices_c, -1)
     else:
         out = torch.empty_like(token_indices_c)
@@ -596,6 +673,7 @@ def triton_filter_and_convert_dcp_index(
         # No prefill workspace on the DCP decode path.
         None,
         None,
+        workspace_rank_stride=0,
         max_num_blocks_per_req=max_num_blocks_per_req,
         BLOCK_SIZE=BLOCK_SIZE,
         BLOCK_STRIDE_ROWS=(
@@ -618,3 +696,46 @@ def triton_filter_and_convert_dcp_index(
 
 
 _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL = ConvertReqIndexToGlobalIndexKernel()
+
+
+@triton.jit(do_not_specialize=["num_tokens"])
+def sparse_mla_prepare_safe_lengths_kernel(
+    indices_ptr,
+    counts_ptr,
+    safe_lengths_ptr,
+    num_tokens,
+    index_stride0: tl.constexpr,
+    count_stride: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    token = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    count = tl.load(counts_ptr + token * count_stride, token < num_tokens, other=1)
+    tl.store(safe_lengths_ptr + token, tl.maximum(count, 1), token < num_tokens)
+    # TRTLLM's NoPE sparse kernel requires at least one valid KV-cache entry.
+    tl.store(
+        indices_ptr + token * index_stride0, 0, (token < num_tokens) & (count == 0)
+    )
+
+
+def prepare_sparse_mla_safe_lengths(
+    physical_indices: torch.Tensor, valid_counts: torch.Tensor
+) -> torch.Tensor:
+    """Install dummy slots for empty queries and return nonzero kernel lengths.
+
+    Preserve the raw counts so empty outputs and LSE can be neutralized later.
+    """
+    num_tokens = valid_counts.numel()
+    safe_lengths = torch.empty(
+        (num_tokens,), dtype=valid_counts.dtype, device=valid_counts.device
+    )
+    if num_tokens:
+        sparse_mla_prepare_safe_lengths_kernel[(triton.cdiv(num_tokens, 256),)](
+            physical_indices,
+            valid_counts,
+            safe_lengths,
+            num_tokens,
+            physical_indices.stride(0),
+            valid_counts.stride(0),
+            BLOCK=256,
+        )
+    return safe_lengths
diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py
index 20cddfa4ca9c..b3d982eac3c5 100644
--- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py
+++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py
@@ -87,7 +87,7 @@ class TokenspeedMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [32, 64]
 
     @staticmethod
diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py
index 78787d5c264d..950d798f2210 100644
--- a/vllm/v1/attention/backends/mla/triton_mla.py
+++ b/vllm/v1/attention/backends/mla/triton_mla.py
@@ -6,7 +6,6 @@
 import torch
 
 import vllm.envs as envs
-from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention.mla_attention import (
@@ -14,6 +13,7 @@
     MLACommonImpl,
     MLACommonMetadata,
     MLACommonMetadataBuilder,
+    QueryLenSupport,
 )
 from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
@@ -26,7 +26,6 @@
     MultipleOf,
 )
 from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
-from vllm.v1.kv_cache_interface import KVCacheSpec
 from vllm.v1.worker.workspace import (
     current_workspace_manager,
     is_workspace_manager_initialized,
@@ -50,50 +49,16 @@ def _compute_num_kv_splits(max_seq_len: int, sm_count: int) -> int:
 
 
 class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
-    _cudagraph_support: ClassVar[AttentionCGSupport] = (
-        AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
-    )
-    # Non-causal DSpark block is flattened to one decode row per query token in
-    # forward_mqa, so no intra-block causal masking is required.
+    # forward_mqa flattens a uniform multi-token block to one decode row per
+    # query token, so causal and non-causal blocks both take the decode path.
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
     supports_non_causal_multi_token_decode: ClassVar[bool] = True
 
-    @classmethod
-    def get_cudagraph_support(
-        cls,
-        vllm_config: VllmConfig,
-        kv_cache_spec: KVCacheSpec,
-    ) -> AttentionCGSupport:
-        """Report UNIFORM_BATCH where a non-causal multi-token block is served.
-
-        ``_cudagraph_support`` is a class constant, so serving the DSpark
-        draft's (1 + num_spec) block through the decode path reports
-        UNIFORM_SINGLE_TOKEN_DECODE and, because the engine takes the minimum
-        over all attention groups, downgrades the *whole* engine off full
-        cudagraphs. ``forward_mqa`` flattens that block with
-        ``repeat_interleave`` on a Python int and performs no device->host
-        sync, so it does satisfy the UNIFORM_BATCH contract.
-
-        ``non_causal_multi_token_decode`` is a KV-cache-group property, not a
-        per-layer one: ``MLAAttentionSpec.merge`` ORs it over every layer in
-        the group, so a group holding both a draft and its target reports it
-        for both. That is the same predicate ``__init__`` below already uses to
-        raise ``reorder_batch_threshold``, so the two stay consistent, but it
-        does mean this lifts a causal target sharing the draft's KV cache group
-        as well.
-        """
-        if getattr(kv_cache_spec, "non_causal_multi_token_decode", False):
-            return AttentionCGSupport.UNIFORM_BATCH
-        return cls._cudagraph_support
-
     def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
         # DCP local sequence lengths are not advanced between draft steps.
         self.supports_draft_decode_metadata_update = self.dcp_world_size == 1
-        # Only the non-causal DSpark draft group serves multi-token blocks via
-        # the decode path; raise its reorder threshold to the spec block length
-        # so full-cudagraph capture admits it. Causal usage stays single-token.
-        if getattr(self, "non_causal_multi_token_decode", False):
-            self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
         self._reserve_attn_logits_workspace()
 
     def update_draft_decode_metadata(self, _metadata: MLACommonMetadata) -> None:
@@ -109,12 +74,12 @@ def _reserve_attn_logits_workspace(self) -> None:
         """
         if not is_workspace_manager_initialized():
             return
-        # Decode reorder threshold is 1, so decode tokens <= max_num_seqs.
-        B = self.vllm_config.scheduler_config.max_num_seqs
-        # Non-causal DSpark draft flattens each request's block to query_len
-        # decode rows; cover max_num_seqs * block_len rows.
-        if getattr(self, "non_causal_multi_token_decode", False):
-            B *= self.reorder_batch_threshold
+        # forward_mqa flattens each request's block to query_len decode rows,
+        # and query_len is bounded by the reorder threshold.
+        B = (
+            self.vllm_config.scheduler_config.max_num_seqs
+            * self.reorder_batch_threshold
+        )
         # DCP all-gathers the query heads before forward_mqa.
         q_num_heads = self.num_heads * self.dcp_world_size
         max_splits = _compute_num_kv_splits(
@@ -142,7 +107,7 @@ def get_supported_head_sizes(cls) -> list[int]:
         return []
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
     @classmethod
@@ -313,14 +278,32 @@ def forward_mqa(
 
         block_table = attn_metadata.decode.block_table
         seq_lens = attn_metadata.decode.seq_lens
-        if not attn_metadata.causal:
-            # Non-causal DSpark block: flatten to one decode row per query token.
-            # Each row attends to the same committed KV prefix (per-row seq_lens)
-            # and never to sibling block tokens = non-causal block semantics.
-            # Mirrors FlashInferMLA's non-causal path.
-            query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
-            if query_len > 1:
-                block_table = block_table.repeat_interleave(query_len, dim=0)
+        # decode_attention_fwd has no causal flag: it launches one program per
+        # row of q and reads that row's KV extent from seq_lens, so intra-block
+        # causality is expressed as per-row extents. Deriving query_len from the
+        # tensors the kernel indexes keeps the three row counts in step.
+        num_decodes = seq_lens.shape[0]
+        query_len, remainder = divmod(B, num_decodes) if num_decodes else (1, 0)
+        assert remainder == 0, (
+            f"non-uniform decode block: {B} query rows over {num_decodes} requests"
+        )
+        if query_len > 1:
+            block_table = block_table.repeat_interleave(query_len, dim=0)
+            if attn_metadata.causal:
+                # Per-row extents are offsets off the global sequence length;
+                # under DCP seq_lens is this rank's local slice instead.
+                assert self.dcp_world_size == 1, (
+                    "causal multi-token decode is not supported with DCP"
+                )
+                # Row t attends the prefix plus block tokens 0..t. Clamp holds
+                # padding rows at extent 0; the kernel skips them either way
+                # (split_kv_end == split_kv_start).
+                offsets = torch.arange(
+                    1 - query_len, 1, device=seq_lens.device, dtype=seq_lens.dtype
+                )
+                seq_lens = (seq_lens.unsqueeze(1) + offsets).flatten().clamp(min=0)
+            else:
+                # Non-causal draft block: every row sees the same prefix.
                 seq_lens = seq_lens.repeat_interleave(query_len)
 
         # Run MQA — always pass layer scales. When KV cache is
diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
index 8aaa4749d506..036dfe01730b 100644
--- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -13,6 +13,9 @@
 from vllm.model_executor.layers.attention.mla_attention import (
     get_mla_dims,
 )
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SharedTopkIndicesBuffer,
+)
 from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
 from vllm.v1.attention.backend import (
     AttentionBackend,
@@ -125,7 +128,7 @@ def __init__(
 
         self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
         self.mla_dims = get_mla_dims(self.model_config)
-        self.topk_tokens = vllm_config.model_config.hf_config.index_topk
+        self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk
         self.topk_tokens_tensor = torch.tensor(
             [self.topk_tokens], device=device, dtype=torch.int32
         )
@@ -183,7 +186,7 @@ def build(
         return metadata
 
 
-class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]):
+class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata], SharedTopkIndicesBuffer):
     is_sparse = True
 
     def __init__(
@@ -210,12 +213,7 @@ def __init__(
         self.kv_cache_dtype = kv_cache_dtype
         self.kv_lora_rank: int = mla_args["kv_lora_rank"]
         self.softmax_scale = scale
-        # The indexer carries the shared buffer for normal layers and tests;
-        # the explicitly-passed buffer covers backbone skip layers, whose
-        # indexer is not constructed (see deepseek_v2.py).
-        self.topk_indices_buffer: torch.Tensor | None = (
-            indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
-        )
+        self.init_topk_indices_buffer(indexer, topk_indices_buffer)
 
     def _forward_bf16_kv(
         self,
diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py
index c1c63305767c..32e2fabeb567 100644
--- a/vllm/v1/attention/backends/registry.py
+++ b/vllm/v1/attention/backends/registry.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Attention backend registry"""
+"""Attention backend registry."""
 
 from collections.abc import Callable
 from enum import Enum, EnumMeta
@@ -49,6 +49,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     TRITON_ATTN_DIFFKV = (
         "vllm.v1.attention.backends.triton_attn_diffkv.TritonAttentionDiffKVBackend"
     )
+    TRITON_FLASH_ATTN = (
+        "vllm.v1.attention.backends.triton_flash_attn.TritonFlashAttentionBackend"
+    )
     ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
     ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"
     ROCM_AITER_TRITON_MLA = (
@@ -63,6 +66,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     XPU_MLA_SPARSE = "vllm.v1.attention.backends.mla.xpu_mla_sparse.XPUMLASparseBackend"
     TORCH_SDPA = ""  # this tag is only used for ViT
     FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
+    TRITON_FLASHINFER = (
+        "vllm.v1.attention.backends.triton_flashinfer.TritonFlashInferBackend"
+    )
     FLASHINFER_MLA = (
         "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
     )
@@ -102,12 +108,15 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     # layer). Separate names from DSV4 so a V4.1 model never resolves the
     # V4.0 backend classes through this enum.
     FLASHMLA_SPARSE_DSV41 = (
-        "vllm.models.deepseek_v4_1.sparse_mla.DeepseekV4FlashMLABackend"
+        "vllm.models.deepseek_v41.sparse_mla.DeepseekV4FlashMLABackend"
     )
     FLASHINFER_MLA_SPARSE_DSV41 = (
-        "vllm.models.deepseek_v4_1.nvidia.flashinfer_sparse."
+        "vllm.models.deepseek_v41.nvidia.flashinfer_sparse."
         "DeepseekV4FlashInferMLASparseBackend"
     )
+    FLASHMLA_MEGA_ATTN_DSV41 = (
+        "vllm.models.deepseek_v41.sparse_mla.FlashMLAMegaAttnBackend"
+    )
     B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend"
     FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
     FLASH_ATTN_MLA_SPARSE = (
@@ -152,6 +161,7 @@ def get_path(self, include_classname: bool = True) -> str:
 
         Raises:
             ValueError: If Backend.CUSTOM is used without being registered
+
         """
         path = _ATTN_OVERRIDES.get(self, self.value)
         if not path:
@@ -172,6 +182,7 @@ def get_class(self) -> "type[AttentionBackend]":
         Raises:
             ImportError: If the backend class cannot be imported
             ValueError: If Backend.CUSTOM is used without being registered
+
         """
         return resolve_obj_by_qualname(self.get_path())
 
@@ -180,6 +191,7 @@ def is_overridden(self) -> bool:
 
         Returns:
             True if the backend has a registered override
+
         """
         return self in _ATTN_OVERRIDES
 
@@ -215,6 +227,7 @@ def get_path(self, include_classname: bool = True) -> str:
 
         Raises:
             ValueError: If Backend.CUSTOM is used without being registered
+
         """
         path = _MAMBA_ATTN_OVERRIDES.get(self, self.value)
         if not path:
@@ -235,6 +248,7 @@ def get_class(self) -> "type[AttentionBackend]":
         Raises:
             ImportError: If the backend class cannot be imported
             ValueError: If Backend.CUSTOM is used without being registered
+
         """
         return resolve_obj_by_qualname(self.get_path())
 
@@ -243,6 +257,7 @@ def is_overridden(self) -> bool:
 
         Returns:
             True if the backend has a registered override
+
         """
         return self in _MAMBA_ATTN_OVERRIDES
 
@@ -266,6 +281,7 @@ def register_backend(
         backend: The AttentionBackendEnum member to register
         class_path: Optional class path. If not provided and used as
             decorator, will be auto-generated from the class.
+        is_mamba: Whether the backend is a Mamba attention backend.
 
     Returns:
         Decorator function if class_path is None, otherwise a no-op
@@ -291,6 +307,7 @@ class MyCustomBackend:
             AttentionBackendEnum.CUSTOM,
             "my.module.MyCustomBackend"
         )
+
     """
 
     def decorator(cls: type) -> type:
diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py
index 062aa82285a1..3df3b26bbaaa 100644
--- a/vllm/v1/attention/backends/rocm_aiter_fa.py
+++ b/vllm/v1/attention/backends/rocm_aiter_fa.py
@@ -770,8 +770,7 @@ def build_for_drafting(
         common_attn_metadata: CommonAttentionMetadata,
         draft_index: int,
     ) -> AiterFlashAttentionMetadata:
-        """
-        Build attention metadata for draft model without CPU-GPU sync.
+        """Build attention metadata for draft model without CPU-GPU sync.
 
         During EAGLE drafting all requests are uniform decodes, so we can
         skip split_decodes_prefills_and_extends() and avoid all .cpu() /
@@ -853,7 +852,7 @@ def supports_attn_type(cls, attn_type: str) -> bool:
         return attn_type in (AttentionType.DECODER,)
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():
             return [16, 32]
         # Only gluon serves 128-token pages; the pa_fwd_asm/ll4mi decode
@@ -990,6 +989,20 @@ def __init__(
                 "cross-attention."
             )
 
+    def _get_kv_cache_descales(
+        self,
+        layer: AttentionLayer,
+        num_decodes: int,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        # AITER treats non-null descales as an instruction to dequantize.
+        if not is_quantized_kv_cache(self.kv_cache_dtype):
+            return None, None
+        descale_shape = (num_decodes, self.num_kv_heads)
+        return (
+            layer._k_scale.expand(descale_shape),
+            layer._v_scale.expand(descale_shape),
+        )
+
     def extend_for_sliding_window(
         self,
         attn_metadata: AiterFlashAttentionMetadata,
@@ -1189,17 +1202,25 @@ def forward(
         """Forward pass with AiterFlashAttention.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
                 [num_blocks, 2, block_size, num_kv_heads, head_size]
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization; not supported
+                by this backend.
+            output_block_scale: Block scale for fused output quantization;
+                not supported by this backend.
+
         Returns:
             shape = [num_tokens, num_heads * head_size]
         NOTE: FP8 quantization, flash-attn expect the size of
               {q,k,v}_descale to be (num_sequences, num_kv_heads).
               We use torch's .expand() to avoid duplicating values
+
         """
         if output_scale is not None or output_block_scale is not None:
             raise NotImplementedError(
@@ -1360,6 +1381,9 @@ def forward(
                     or decode_max_query_len > 1
                     or self.sinks is not None
                 ):
+                    k_descale, v_descale = self._get_kv_cache_descales(
+                        layer, num_decodes
+                    )
                     assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                         "Shuffle KV cache layout is not supported with sliding "
                         "window, sinks, or speculative decoding (multi-token decode)."
@@ -1369,7 +1393,6 @@ def forward(
                             flash_attn_with_kvcache,
                         )
 
-                        descale_shape = (num_decodes, key_cache.shape[2])
                         decode_query = query[:num_decode_tokens].reshape(
                             num_decodes,
                             decode_max_query_len,
@@ -1386,8 +1409,8 @@ def forward(
                             window_size=self.sliding_window,
                             softcap=self.logits_soft_cap,
                             q_descale=None,
-                            k_descale=layer._k_scale.expand(descale_shape),
-                            v_descale=layer._v_scale.expand(descale_shape),
+                            k_descale=k_descale,
+                            v_descale=v_descale,
                             page_table=attn_metadata.block_table[:num_decodes],
                         )
                         output[:num_decode_tokens].copy_(
@@ -1405,10 +1428,6 @@ def forward(
                             unified_attention,
                         )
 
-                        descale_shape = (
-                            num_decodes,
-                            key_cache.shape[2],
-                        )
                         unified_attention(
                             q=query[:num_decode_tokens],
                             k=key_cache,
@@ -1427,8 +1446,8 @@ def forward(
                             block_table=attn_metadata.block_table[:num_decodes],
                             softcap=self.logits_soft_cap,
                             q_descale=None,
-                            k_descale=layer._k_scale.expand(descale_shape),
-                            v_descale=layer._v_scale.expand(descale_shape),
+                            k_descale=k_descale,
+                            v_descale=v_descale,
                             sinks=self.sinks,
                         )
                     return
@@ -1442,6 +1461,9 @@ def forward(
                 use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI
 
                 if use_unified_attention:
+                    k_descale, v_descale = self._get_kv_cache_descales(
+                        layer, num_decodes
+                    )
                     assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                         "unified_attention fallback with shuffle layout "
                         "is not supported yet."
@@ -1453,10 +1475,6 @@ def forward(
                     decode_cu_seqlens_q = attn_metadata.query_start_loc[
                         : num_decodes + 1
                     ]
-                    descale_shape = (
-                        num_decodes,
-                        key_cache.shape[2],
-                    )
                     unified_attention(
                         q=query[:num_decode_tokens],
                         k=key_cache,
@@ -1473,8 +1491,8 @@ def forward(
                         block_table=attn_metadata.block_table[:num_decodes],
                         softcap=self.logits_soft_cap,
                         q_descale=None,
-                        k_descale=layer._k_scale.expand(descale_shape),
-                        v_descale=layer._v_scale.expand(descale_shape),
+                        k_descale=k_descale,
+                        v_descale=v_descale,
                     )
                 elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
                     _, num_heads, head_size = query.shape
diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
index 3f6c1916f136..d18ea8c29be6 100644
--- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py
@@ -7,7 +7,8 @@
 import torch
 
 from vllm import _custom_ops as ops
-from vllm._aiter_ops import rocm_aiter_ops
+from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops
+from vllm.config import get_current_vllm_config_or_none
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -33,6 +34,14 @@
 logger = init_logger(__name__)
 
 
+def _kv_connector_enabled() -> bool:
+    vllm_config = get_current_vllm_config_or_none()
+    if vllm_config is None:
+        return False
+    kv_transfer_config = vllm_config.kv_transfer_config
+    return kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance
+
+
 class RocmAiterUnifiedAttentionMetadataBuilder(RocmAttentionMetadataBuilder):
     def build_for_cudagraph_capture(
         self, common_attn_metadata: CommonAttentionMetadata
@@ -56,7 +65,7 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
     @classmethod
@@ -85,6 +94,10 @@ def supports_sink(cls) -> bool:
     def supports_non_causal(cls) -> bool:
         return False
 
+    @classmethod
+    def supports_kv_connector(cls) -> bool:
+        return True
+
     forward_includes_kv_cache_update: bool = False
 
     @staticmethod
@@ -111,6 +124,9 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
         # K and V come out of the content dim as transposed views rather than
         # copies, so the head dim may sit on either side of the block dim, but
         # the layer must stay outermost.
+        if _kv_connector_enabled():
+            # Connectors like MoRI assume contiguous blocks
+            return (KVCacheLayout.LBHNC,)
         return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)
 
     @staticmethod
@@ -214,14 +230,21 @@ def forward(
         """Forward pass with FlashAttention.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
                 [num_blocks, 2, block_size, num_kv_heads, head_size]
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization.
+            output_block_scale: Block scale for fused output quantization;
+                not supported by this backend.
+
         Returns:
             shape = [num_tokens, num_heads * head_size]
+
         """
         if output_block_scale is not None:
             raise NotImplementedError(
@@ -384,6 +407,9 @@ def fused_rope_kvcache_supported(self):
     def fused_qk_norm_rope_kvcache_supported(self):
         return rocm_aiter_ops.is_enabled()
 
+    def fused_qk_norm_mrope_kvcache_supported(self) -> bool:
+        return IS_AITER_FOUND
+
     def do_qk_norm_rope_kvcache_update(
         self,
         layer: AttentionLayer,
@@ -422,6 +448,47 @@ def do_qk_norm_rope_kvcache_update(
             use_shuffle_layout=self._use_shuffle_kv,
         )
 
+    def do_qk_norm_mrope_kvcache_update(
+        self,
+        layer: AttentionLayer,
+        qkv: torch.Tensor,
+        q_out: torch.Tensor,
+        positions: torch.Tensor,
+        q_weight: torch.Tensor,
+        k_weight: torch.Tensor,
+        rms_norm_eps: float,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        mrope_section: tuple[int, int, int],
+        is_interleaved: bool,
+        rotary_dim: int,
+        kv_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+    ):
+        key_cache, value_cache = self._split_kv_cache(kv_cache)
+        rocm_aiter_ops.do_qk_norm_mrope_kvcache_update(
+            qkv=qkv,
+            q_weight=q_weight,
+            k_weight=k_weight,
+            cos_sin_cache=cos_sin_cache,
+            positions=positions,
+            num_heads_q=self.num_heads,
+            num_heads_k=self.num_kv_heads,
+            head_dim=self.head_size,
+            is_neox=is_neox,
+            mrope_section=list(mrope_section),
+            is_interleaved=is_interleaved,
+            rms_norm_eps=rms_norm_eps,
+            q_out=q_out,
+            key_cache=key_cache,
+            value_cache=value_cache,
+            slot_mapping=layer_slot_mapping,
+            k_scale=layer._k_scale_cpu,
+            v_scale=layer._v_scale_cpu,
+            kv_cache_dtype=self.kv_cache_dtype,
+            rotary_dim=rotary_dim,
+        )
+
     def do_rope_and_kv_cache_update(
         self,
         layer: AttentionLayer,
diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py
index f3bb3b351451..76382aa75e76 100644
--- a/vllm/v1/attention/backends/rocm_attn.py
+++ b/vllm/v1/attention/backends/rocm_attn.py
@@ -176,7 +176,7 @@ class RocmAttentionBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # ROCM paged attention native C++ kernel only supports block sizes 16 and 32
         # due to shared memory (LDS) constraints on AMD GPUs.
         # See csrc/rocm/attention.cu CALL_CUSTOM_LAUNCHER_BLK macro.
@@ -344,6 +344,7 @@ def _forward_encoder_attention(
             output: shape = [num_encoder_tokens, num_heads, head_size]
             attn_metadata: Encoder attention metadata
             layer: The attention layer
+
         """
         # For encoder attention, process FP8 quantization if needed
         if is_quantized_kv_cache(self.kv_cache_dtype):
@@ -390,14 +391,21 @@ def forward(
         """Forward pass with FlashAttention.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: logical [num_blocks, 2, block_size, num_kv_heads *
                 head_size] under LHBNC (physically K/V-group-first)
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization.
+            output_block_scale: Block scale for fused output quantization;
+                not supported by this backend.
+
         Returns:
             shape = [num_tokens, num_heads * head_size]
+
         """
         if output_block_scale is not None:
             raise NotImplementedError(
diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py
index 48dd30d66635..8a2fab4cc7dd 100644
--- a/vllm/v1/attention/backends/triton_attn.py
+++ b/vllm/v1/attention/backends/triton_attn.py
@@ -304,7 +304,7 @@ def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
     @classmethod
@@ -323,6 +323,10 @@ def supports_non_causal(cls) -> bool:
     def get_name() -> str:
         return "TRITON_ATTN"
 
+    @classmethod
+    def supports_rswa(cls) -> bool:
+        return True
+
     @classmethod
     def supports_sliding_window(cls) -> bool:
         return True
@@ -553,14 +557,21 @@ def forward(
         """Forward pass with Paged Attention impl. in Triton.
 
         Args:
+            layer: The attention layer, providing the q/k/v quantization scales.
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
                 [num_blocks, num_kv_heads, block_size, 2 * head_size]
             attn_metadata: Metadata for attention.
+            output: Tensor that the attention result is written into.
+            output_scale: Scale for fused output quantization.
+            output_block_scale: Block scale for fused output quantization;
+                not supported by this backend.
+
         Returns:
             shape = [num_tokens, num_heads * head_size]
+
         """
         if output_block_scale is not None:
             raise NotImplementedError(
@@ -720,6 +731,7 @@ def _forward_encoder_attention(
             output: shape = [num_encoder_tokens, num_heads, head_size]
             attn_metadata: Encoder attention metadata
             layer: The attention layer
+
         """
         # Quantized KV cache is not supported for encoder attention.
         if is_quantized_kv_cache(self.kv_cache_dtype):
diff --git a/vllm/v1/attention/backends/triton_flash_attn.py b/vllm/v1/attention/backends/triton_flash_attn.py
new file mode 100644
index 000000000000..152a92048e07
--- /dev/null
+++ b/vllm/v1/attention/backends/triton_flash_attn.py
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Triton image-mask and FlashAttention causal attention composite."""
+
+import torch
+
+from vllm.config.cache import CacheDType
+from vllm.platforms.interface import DeviceCapability
+from vllm.v1.attention.backends.composite import (
+    MMPrefixAttentionRouting,
+    create_composite_attention_backend,
+)
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
+
+
+class _FABackend(FlashAttentionBackend):
+    @classmethod
+    def supports_combination(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        device_capability: DeviceCapability,
+    ) -> str | None:
+        reason = super().supports_combination(
+            head_size,
+            dtype,
+            kv_cache_dtype,
+            block_size,
+            use_mla,
+            has_sink,
+            use_sparse,
+            use_mm_prefix,
+            device_capability,
+        )
+        if reason is not None:
+            return reason
+        fa_version = get_flash_attn_version(
+            head_size=head_size,
+            has_sinks=has_sink,
+            kv_cache_block_size=block_size,
+            supports_fa4_hd256=True,
+        )
+        if fa_version not in (3, 4):
+            return "causal route requires FlashAttention v3 or v4"
+        return None
+
+
+TritonFlashAttentionBackend = create_composite_attention_backend(
+    TritonAttentionBackend,
+    _FABackend,
+    name="TritonFlashAttentionBackend",
+    backend_name="TRITON_FLASH_ATTN",
+    module=__name__,
+    routing_policy=MMPrefixAttentionRouting,
+    head_sizes=(256, 512),
+    device_major=9,
+)
diff --git a/vllm/v1/attention/backends/triton_flashinfer.py b/vllm/v1/attention/backends/triton_flashinfer.py
new file mode 100644
index 000000000000..07912275913e
--- /dev/null
+++ b/vllm/v1/attention/backends/triton_flashinfer.py
@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Triton image-mask and FlashInfer causal attention composite."""
+
+from vllm.v1.attention.backends.composite import (
+    MMPrefixAttentionRouting,
+    create_composite_attention_backend,
+)
+from vllm.v1.attention.backends.flashinfer import FlashInferBackend
+from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
+
+TritonFlashInferBackend = create_composite_attention_backend(
+    TritonAttentionBackend,
+    FlashInferBackend,
+    name="TritonFlashInferBackend",
+    backend_name="TRITON_FLASHINFER",
+    module=__name__,
+    routing_policy=MMPrefixAttentionRouting,
+    head_sizes=(256, 512),
+    kernel_block_sizes=(64,),
+    device_major=10,
+)
diff --git a/vllm/v1/attention/backends/turboquant_attn.py b/vllm/v1/attention/backends/turboquant_attn.py
index 8542e6164f03..056e42cf62d5 100644
--- a/vllm/v1/attention/backends/turboquant_attn.py
+++ b/vllm/v1/attention/backends/turboquant_attn.py
@@ -161,7 +161,7 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
         return (KVCacheLayout.LBNHC,)
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [16, 32, 64, 128]
 
     @classmethod
diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py
index 716f3689592f..e918774088f4 100644
--- a/vllm/v1/attention/backends/utils.py
+++ b/vllm/v1/attention/backends/utils.py
@@ -8,6 +8,7 @@
     TYPE_CHECKING,
     Any,
     Protocol,
+    TypeVar,
     cast,
 )
 
@@ -45,6 +46,19 @@
 PAD_SLOT_ID = -1
 NULL_BLOCK_ID = 0
 
+
+ImplT = TypeVar("ImplT", bound=AttentionImpl)
+
+
+def find_attention_impl_variant(
+    impl: AttentionImpl, impl_cls: type[ImplT]
+) -> ImplT | None:
+    for variant in impl.get_impl_variants():
+        if isinstance(variant, impl_cls):
+            return cast(ImplT, variant)
+    return None
+
+
 _LN_2 = math.log(2.0)
 
 
@@ -310,8 +324,7 @@ def resolve_kv_cache_layout(
 
 @dataclass
 class PerLayerParameters:
-    """
-    Currently, FlashInfer backend only support models in which all layers share
+    """Currently, FlashInfer backend only support models in which all layers share
     the same values for the following hyperparameters. Should not be used for
     trtllm-gen backend since it supports different values for the following
     hyperparameters.
@@ -329,11 +342,9 @@ class PerLayerParameters:
 def get_per_layer_parameters(
     vllm_config: VllmConfig, layer_names: list[str], cls_: type["AttentionImpl"]
 ) -> dict[str, PerLayerParameters]:
-    """
-    Scan layers in `layer_names` and determine some hyperparameters
+    """Scan layers in `layer_names` and determine some hyperparameters
     to use during `plan`.
     """
-
     layers = get_layers_from_vllm_config(
         vllm_config,
         AttentionLayerBase,  # type: ignore[type-abstract]
@@ -342,8 +353,8 @@ def get_per_layer_parameters(
     per_layer_params: dict[str, PerLayerParameters] = {}
 
     for key, layer in layers.items():
-        impl = layer.impl
-        assert isinstance(impl, cls_)
+        impl = find_attention_impl_variant(layer.impl, cls_)
+        assert impl is not None
 
         # Infer hyperparameters from the attention layer
         window_size = getattr(impl, "sliding_window", None)
@@ -391,8 +402,7 @@ def get_num_attention_heads_from_layers(
 def infer_global_hyperparameters(
     per_layer_params: dict[str, PerLayerParameters],
 ) -> PerLayerParameters:
-    """
-    Currently, FlashInfer backend other than trtllm-gen
+    """Currently, FlashInfer backend other than trtllm-gen
     only support models in which all layers share
     the same values for the following hyperparameters:
     - `window_left`
@@ -402,7 +412,6 @@ def infer_global_hyperparameters(
     So this function asserts that all layers share the same values for these
     hyperparameters and returns the global values.
     """
-
     assert len(per_layer_params) > 0, "No attention layers found in the model."
 
     param_sets = list(per_layer_params.values())
@@ -478,7 +487,7 @@ def make_local_attention_virtual_batches(
     query_start_loc_np = common_attn_metadata.query_start_loc_cpu.numpy()
     with gpu_sync_allowed():
         # TODO see https://github.com/vllm-project/vllm/pull/31852
-        seq_lens_np = common_attn_metadata.seq_lens_cpu.numpy()
+        seq_lens_np = common_attn_metadata.seq_lens.cpu().numpy()
     block_table = common_attn_metadata.block_table_tensor
     device = common_attn_metadata.query_start_loc.device
 
@@ -541,8 +550,6 @@ def make_local_attention_virtual_batches(
     #   seqlens_k_local = [4, 2, 4, 4, 4, 1, 4, 1]
     seqlens_k_local = np.full(cu_num_blocks[-1], attn_chunk_size, dtype=np.int32)
     seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block
-    num_computed_tokens_local = seqlens_k_local - seqlens_q_local
-
     k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - (
         rarange * attn_chunk_size + np.repeat(tokens_in_last_block, local_blocks)
     )
@@ -612,9 +619,7 @@ def make_local_attention_virtual_batches(
         block_table_tensor=block_table_local,
         slot_mapping=common_attn_metadata.slot_mapping,
         causal=True,
-        seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound,
-        _seq_lens_cpu=seq_lens_cpu,
-        _num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local),
+        seq_lens_cpu_upper_bound=seq_lens_cpu,
     ), make_block_table
 
 
@@ -626,8 +631,13 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
         # Skip computing fast prefill path
         return common_attn_metadata
 
-    assert common_attn_metadata.logits_indices_padded is not None
-    assert common_attn_metadata.num_logits_indices is not None
+    if (
+        common_attn_metadata.logits_indices_padded is None
+        or common_attn_metadata.num_logits_indices is None
+    ):
+        # Fast prefill not armed for this step (e.g. cudagraph capture, or a
+        # pure-decode step): run the KV-sharing layers on the full batch.
+        return common_attn_metadata
 
     logits_indices_padded = common_attn_metadata.logits_indices_padded
     num_logits_indices = common_attn_metadata.num_logits_indices
@@ -687,8 +697,6 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
         slot_mapping=common_attn_metadata.slot_mapping,
         causal=True,
         seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound,
-        _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,
-        _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
     )
     return common_attn_metadata
 
@@ -697,8 +705,7 @@ def split_decodes_prefills_and_extends(
     common_attn_metadata: CommonAttentionMetadata,
     decode_threshold: int = 1,
 ) -> tuple[int, int, int, int, int, int]:
-    """
-    Assuming a reordered batch, finds the boundary between prefill and decode
+    """Assuming a reordered batch, finds the boundary between prefill and decode
     requests.
 
     Args:
@@ -713,6 +720,7 @@ def split_decodes_prefills_and_extends(
         num_decode_tokens: The number of tokens in the decode requests.
         num_extend_tokens: The number of tokens in the extend requests.
         num_prefill_tokens: The number of tokens in the prefill requests.
+
     """
     max_query_len = common_attn_metadata.max_query_len
     num_reqs = common_attn_metadata.num_reqs
@@ -771,8 +779,7 @@ def split_decodes_and_prefills(
     require_uniform: bool = False,
     treat_short_extends_as_decodes: bool = True,
 ) -> tuple[int, int, int, int]:
-    """
-    Assuming a reordered batch, finds the boundary between prefill and decode
+    """Assuming a reordered batch, finds the boundary between prefill and decode
     requests.
 
     The batch is expected to be ordered as:
@@ -794,6 +801,7 @@ def split_decodes_and_prefills(
         num_prefills: The number of prefill requests.
         num_decode_tokens: The number of tokens in the decode requests.
         num_prefill_tokens: The number of tokens in the prefill requests.
+
     """
     max_query_len = common_attn_metadata.max_query_len
     num_reqs = common_attn_metadata.num_reqs
@@ -820,7 +828,7 @@ def split_decodes_and_prefills(
             (query_lens == query_lens[0]) | (query_lens == 0)
         ):
             return num_reqs, 0, num_tokens, 0  # all decodes
-        is_prefill = query_lens != query_lens[0]
+        is_prefill = (query_lens != query_lens[0]) & (query_lens != 0)
     else:
         is_prefill = query_lens > decode_threshold
 
@@ -842,16 +850,17 @@ def split_decodes_and_prefills(
 def split_prefill_chunks(
     seq_lens_cpu: torch.Tensor, workspace_size: int, request_offset: int = 0
 ) -> list[tuple[int, int]]:
-    """
-    Split the prefill requests into chunks such that the total sequence length
+    """Split the prefill requests into chunks such that the total sequence length
     of each chunk is less than or equal to the workspace size.
 
     Args:
         seq_lens_cpu: The sequence lengths of the prefill requests on CPU.
         workspace_size: The maximum workspace size (in tokens) per chunk.
         request_offset: The offset to add to the request indices.
+
     Returns:
         A list of tuples of (reqs_start, reqs_end) representing chunk boundaries.
+
     """
     chunk_bounds = []
     i, n = 0, len(seq_lens_cpu)
@@ -871,8 +880,7 @@ def reorder_batch_to_split_decodes_and_prefills(
     scheduler_output: "SchedulerOutput",
     decode_threshold: int = 1,
 ) -> bool:
-    """
-    Reorders the batch to split into prefill and decode requests; places all
+    """Reorders the batch to split into prefill and decode requests; places all
     requests with <= decode_threshold tokens at the front of the batch.
 
     The batch is reordered into 4 regions:
@@ -883,6 +891,7 @@ def reorder_batch_to_split_decodes_and_prefills(
 
     Returns:
         True if the batch was modified, False otherwise.
+
     """
     num_reqs = len(input_batch.req_ids)
     num_scheduled_tokens = [
@@ -947,8 +956,7 @@ def reorder_batch_to_split_decodes_and_prefills(
 
 
 def reshape_query_for_spec_decode(query: torch.Tensor, batch_size: int) -> torch.Tensor:
-    """
-    Reshapes the query tensor for the specified batch size, so that
+    """Reshapes the query tensor for the specified batch size, so that
     it has shape (batch_size, seq_len, num_heads, head_dim).
     """
     assert query.dim() == 3, f"query must be 3D, got {query.dim()}D"
@@ -963,8 +971,7 @@ def reshape_query_for_spec_decode(query: torch.Tensor, batch_size: int) -> torch
 
 
 def reshape_attn_output_for_spec_decode(attn_output: torch.Tensor) -> torch.Tensor:
-    """
-    Reshapes the attention output tensor, so that
+    """Reshapes the attention output tensor, so that
     the batch_size and seq_len dimensions are combined.
     """
     if attn_output.dim() == 3:
@@ -1014,6 +1021,9 @@ def __init__(self, metadata, common_attn_metadata):
                         common_attn_metadata.logits_indices_padded
                     )
                     self.num_logits_indices = common_attn_metadata.num_logits_indices
+                    self._attention_backend_variant = getattr(
+                        metadata, "_attention_backend_variant", 0
+                    )
 
             return KVSharingFastPrefillAttentionMetadata(metadata, common_attn_metadata)
 
@@ -1078,7 +1088,7 @@ def compute_causal_conv1d_metadata(
 def get_dcp_local_seq_lens(
     seq_lens: torch.Tensor,
     dcp_size: int = 1,
-    dcp_rank: int | None = None,
+    dcp_rank: int | torch.Tensor | None = None,
     cp_kv_cache_interleave_size: int = 1,
 ) -> torch.Tensor:
     """While using dcp, kv_cache size stored on each rank may be different,
@@ -1086,6 +1096,7 @@ def get_dcp_local_seq_lens(
     Only consider dcp now, we can extend the case of cp based on this.
     """
     seq_lens_i32 = seq_lens.to(torch.int32)
+    rank_offsets: int | torch.Tensor
     if dcp_rank is None:
         rank_offsets = torch.arange(
             dcp_size,
@@ -1096,8 +1107,15 @@ def get_dcp_local_seq_lens(
             dcp_size,
         )
         seq_lens_tiled = seq_lens_i32.unsqueeze(-1)
+    elif isinstance(dcp_rank, torch.Tensor):
+        assert dcp_rank.dtype == torch.int32
+        assert dcp_rank.device == seq_lens.device
+        assert dcp_rank.numel() == 1
+        rank_offsets = dcp_rank
+        seq_lens_tiled = seq_lens_i32
     else:
-        rank_offsets = torch.tensor(dcp_rank, dtype=torch.int32, device=seq_lens.device)
+        # Use the Python scalar directly to avoid a synchronizing H2D copy.
+        rank_offsets = dcp_rank
         seq_lens_tiled = seq_lens_i32
     base = (
         seq_lens_tiled
@@ -1121,8 +1139,7 @@ def mamba_get_block_table_tensor(
     kv_cache_spec: KVCacheSpec,
     mamba_cache_mode: str,
 ) -> torch.Tensor:
-    """
-    Get the block table tensor for mamba kernels from the input
+    """Get the block table tensor for mamba kernels from the input
     common_attn_metadata.block_table_tensor given different mamba cache modes.
 
     - "all":   input  (#requests, cdiv(max_model_len, block_size)
diff --git a/vllm/v1/attention/backends/zentorch_sdpa.py b/vllm/v1/attention/backends/zentorch_sdpa.py
new file mode 100644
index 000000000000..cafb524b0e43
--- /dev/null
+++ b/vllm/v1/attention/backends/zentorch_sdpa.py
@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Encoder attention through the zentorch SDPA kernel on Zen CPUs."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
+from vllm.v1.attention.backend import AttentionType
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.cpu_attn import CPUAttentionMetadata
+
+__all__ = ["should_use_zentorch_sdpa", "zentorch_sdpa_attn"]
+
+
+def should_use_zentorch_sdpa(
+    attn_type: str,
+    dtype: torch.dtype,
+) -> bool:
+    """True when encoder attention should dispatch to zentorch_sdpa.
+
+    Full, sliding-window, and ALiBi encoder layers all dispatch. A window is
+    a single-head [1, 1, S, S] additive mask; ALiBi is a per-query-head
+    [1, H, S, S] bias. The zentorch_sdpa_attn op builds both.
+
+    dtype must match the op's ISA gate (bf16: AVX512-BF16, fp32: AVX512);
+    fp16 needs AVX512-FP16, which torch exposes no query for, so it stays
+    native.
+    """
+    if attn_type not in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
+        return False
+    if not has_zentorch_op(["zentorch_sdpa_attn"]):
+        return False
+    if dtype == torch.bfloat16:
+        return torch.cpu._is_avx512_bf16_supported()
+    if dtype == torch.float32:
+        return torch.cpu._is_avx512_supported()
+    return False
+
+
+def zentorch_sdpa_attn(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    attn_metadata: CPUAttentionMetadata,
+    scale: float,
+    sliding_window: int = -1,
+    alibi_slopes: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run encoder / encoder-only attention with zentorch_sdpa.
+
+    Encoder attention reads no KV cache, so the packed
+    [num_tokens, num_heads, head_size] query/key/value are attended in place
+    and written straight into `output`.
+
+    Args:
+        query: Query of shape [num_tokens, num_heads, head_size].
+        key: Key of shape [num_tokens, num_kv_heads, head_size].
+        value: Value of shape [num_tokens, num_kv_heads, head_size].
+        output: Pre-allocated output, same shape as `query`.
+        attn_metadata: Metadata carrying the per-sequence token offsets and
+            whether attention is causal.
+        scale: Softmax scale.
+        sliding_window: Encoder window size, or -1 for full bidirectional
+            attention. Handed to the op, which masks token i to
+            [i-(W-1), i+(W-1)].
+        alibi_slopes: Per-query-head ALiBi slopes, or None. Handed to the op,
+            which applies them as an additive per-head bias.
+
+    Returns:
+        `output`, filled in place.
+
+    """
+    # zentorch_sdpa_attn builds the ALiBi mask in query.dtype. A bf16 per-head
+    # mask is wrong when the query head count is not a power of two, so run
+    # that path in fp32 and write bf16 back.
+    if alibi_slopes is not None and query.dtype == torch.bfloat16:
+        out32 = torch.empty_like(query, dtype=torch.float32)
+        zentorch_sdpa_attn(
+            query.float(),
+            key.float(),
+            value.float(),
+            out32,
+            attn_metadata,
+            scale,
+            sliding_window,
+            alibi_slopes,
+        )
+        output.copy_(out32.to(dtype=output.dtype))
+        return output
+
+    # The op counts window tokens on each side of the diagonal; vLLM's window
+    # includes the diagonal itself.
+    window = -1 if sliding_window in (None, -1) else sliding_window - 1
+
+    # zentorch_sdpa takes no enable_gqa argument but derives the KV head count
+    # from the key tensor and maps each query head onto its KV head, so GQA/MQA
+    # needs no expansion here.
+    query = query.movedim(0, query.dim() - 2)
+    key = key.movedim(0, key.dim() - 2)
+    value = value.movedim(0, value.dim() - 2)
+
+    start_loc = attn_metadata.query_start_loc.numpy()
+    seq_lens = start_loc[1:] - start_loc[:-1]
+
+    # Encoder sequences are packed along the token dimension. When they all
+    # have the same length, recover a dense BHSD batch and invoke zentorch_sdpa
+    # once for the whole scheduler batch instead of once per sequence.
+    if len(seq_lens) > 0 and (seq_lens == seq_lens[0]).all():
+        batch_size = len(seq_lens)
+        seq_len = int(seq_lens[0])
+
+        def _packed_hsd_to_bhsd(tensor: torch.Tensor) -> torch.Tensor:
+            return tensor.unflatten(1, (batch_size, seq_len)).permute(1, 0, 2, 3)
+
+        torch.ops.zentorch.zentorch_sdpa_attn(
+            _packed_hsd_to_bhsd(query),
+            _packed_hsd_to_bhsd(key),
+            _packed_hsd_to_bhsd(value),
+            # The op writes [B, H, S, D]; hand it the matching view of the
+            # packed buffer so it stores there instead of into a temporary.
+            output.unflatten(0, (batch_size, seq_len)).permute(0, 2, 1, 3),
+            scale=scale,
+            is_causal=attn_metadata.causal,
+            alibi_slopes=alibi_slopes,
+            left_window_size=window,
+            right_window_size=window,
+        )
+        return output
+
+    for start_q, end_q in zip(start_loc[:-1], start_loc[1:], strict=True):
+        torch.ops.zentorch.zentorch_sdpa_attn(
+            query[None, :, start_q:end_q, :],
+            key[None, :, start_q:end_q, :],
+            value[None, :, start_q:end_q, :],
+            output[start_q:end_q, :, :].movedim(0, 1).unsqueeze(0),
+            scale=scale,
+            is_causal=attn_metadata.causal,
+            alibi_slopes=alibi_slopes,
+            left_window_size=window,
+            right_window_size=window,
+        )
+    return output
diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py
index 88ce37122626..5c1677cac479 100644
--- a/vllm/v1/attention/ops/common.py
+++ b/vllm/v1/attention/ops/common.py
@@ -5,11 +5,11 @@
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
 )
 from vllm.triton_utils import tl, triton
 
@@ -178,6 +178,7 @@ def pack_seq_triton(
 
     Returns:
         packed: [B, Lmax, ...] — packed tensor.
+
     """
     is_uint8 = x.dtype == torch.uint8
     if is_uint8:
@@ -334,8 +335,7 @@ def unpack_seq_triton(
     block_t: int = 64,
     block_d: int = 64,
 ) -> torch.Tensor:
-    """
-    Unpack a packed decode query tensor back to the original format.
+    """Unpack a packed decode query tensor back to the original format.
     Efficient Triton implementation.
 
     Args:
@@ -346,8 +346,8 @@ def unpack_seq_triton(
 
     Returns:
         unpacked_tensor: [N, ...] where N = sum(lengths)
-    """
 
+    """
     # Handle multi-dimensional input by reshaping to (B, Lmax, -1)
     original_shape = packed_tensor.shape
     if len(original_shape) > 3:
diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py
index 96b25870b8d0..8f844b647641 100644
--- a/vllm/v1/attention/ops/dcp.py
+++ b/vllm/v1/attention/ops/dcp.py
@@ -18,12 +18,12 @@
 from vllm.logger import init_logger
 from vllm.model_executor.warmup.jit_warmup import (
     WarmupIntRange,
+    kernel_launcher,
 )
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
     triton_scalar_specialization_rep,
 )
 from vllm.triton_utils import tl, triton
@@ -45,13 +45,18 @@
 # LSE/output combine
 
 
-def mask_dcp_empty_shards_(
-    lse: torch.Tensor,
+def _validate_dcp_empty_shard_args(
     seq_lens: torch.Tensor | None,
     query_start_loc: torch.Tensor | None,
-) -> None:
+) -> bool:
+    """Validate the empty-shard mask inputs.
+
+    Returns True when masking is requested, False when it is disabled. Raises on
+    an inconsistent pair. Performs no device work so it is safe to call from a
+    CUDA-graph-captured region.
+    """
     if seq_lens is None and query_start_loc is None:
-        return
+        return False
     if seq_lens is None or query_start_loc is None:
         raise ValueError("seq_lens and query_start_loc must be provided together")
     if (
@@ -60,6 +65,24 @@ def mask_dcp_empty_shards_(
         or query_start_loc.shape[0] != seq_lens.shape[0] + 1
     ):
         raise ValueError("query_start_loc must contain one boundary per sequence")
+    return True
+
+
+def mask_dcp_empty_shards_(
+    lse: torch.Tensor,
+    seq_lens: torch.Tensor | None,
+    query_start_loc: torch.Tensor | None,
+) -> None:
+    if not _validate_dcp_empty_shard_args(seq_lens, query_start_loc):
+        return
+    assert seq_lens is not None and query_start_loc is not None
+
+    # A DCP rank can receive no local sequences during CUDA graph warmup even
+    # though the padded LSE buffer still has rows. In that case every row is an
+    # empty shard; avoid indexing the empty seq_lens tensor below.
+    if seq_lens.shape[0] == 0:
+        lse.fill_(float("-inf"))
+        return
 
     row_indices = torch.arange(
         lse.shape[0], device=lse.device, dtype=query_start_loc.dtype
@@ -110,8 +133,7 @@ def kernel(
         N_ROUNDED: tl.constexpr,
         IS_BASE_E: tl.constexpr,
     ):
-        """
-        Apply the all-gathered lses to correct each local rank's attention
+        """Apply the all-gathered lses to correct each local rank's attention
         output. we still need perform a cross-rank reduction to obtain the
         final attention output.
 
@@ -124,6 +146,17 @@ def kernel(
                 Pointer to output tensor of shape [ B, H, D ]
             vlse_ptr (triton.PointerType):
                 Pointer to output tensor of shape [ B, H ]
+            outputs_stride_B (int): Batch stride of ``outputs_ptr``
+            outputs_stride_H (int): Head stride of ``outputs_ptr``
+            outputs_stride_D (int): Head-dim stride of ``outputs_ptr``
+            lses_stride_N (int): Rank stride of ``lses_ptr``
+            lses_stride_B (int): Batch stride of ``lses_ptr``
+            lses_stride_H (int): Head stride of ``lses_ptr``
+            lse_idx (int): Index of this rank's lse within the all-gathered tensor
+            HEAD_DIM: Head dimension, as a constexpr
+            N_ROUNDED: Rank count rounded to a power of two, as a constexpr
+            IS_BASE_E: Whether the lses are natural-log based, as a constexpr
+
         """
         batch_idx = tl.program_id(axis=0).to(tl.int64)
         head_idx = tl.program_id(axis=1).to(tl.int64)
@@ -353,9 +386,11 @@ def correct_attn_out(
         lses: Tensor of shape [ N, B, H ]
         cp_rank: Current rank in the context-parallel group
         ctx: Triton context to avoid recompilation
+        is_lse_base_on_e: Whether the lses use base e rather than base 2
 
     Returns:
         Tuple of (out, lse) with corrected attention and final log-sum-exp.
+
     """
     if ctx is None:
         ctx = CPTritonContext()
@@ -408,8 +443,7 @@ def _cp_lse_common(
     seq_lens: torch.Tensor | None = None,
     query_start_loc: torch.Tensor | None = None,
 ):
-    """
-    cp_attn_out: [ B, H, D ]
+    """cp_attn_out: [ B, H, D ]
     cp_attn_lse: [ B, H ]
     """
     if cp_group.world_size == 1:
@@ -443,8 +477,7 @@ def cp_lse_ag_out_rs(
     seq_lens: torch.Tensor | None = None,
     query_start_loc: torch.Tensor | None = None,
 ):
-    """
-    cp_attn_out: [ B, H, D ]
+    """cp_attn_out: [ B, H, D ]
     cp_attn_lse: [ B, H ]
     """
     out, lse = _cp_lse_common(
@@ -476,8 +509,7 @@ def cp_lse_ag_out_ar(
     seq_lens: torch.Tensor | None = None,
     query_start_loc: torch.Tensor | None = None,
 ):
-    """
-    cp_attn_out: [ B, H, D ]
+    """cp_attn_out: [ B, H, D ]
     cp_attn_lse: [ B, H ]
     """
     out, lse = _cp_lse_common(
@@ -505,8 +537,7 @@ def _lse_weighted_combine(
     return_lse: bool = False,
     is_lse_base_on_e: bool = True,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    CPU reference implementation for LSE-weighted combination.
+    """CPU reference implementation for LSE-weighted combination.
 
     This is a pure PyTorch implementation used for testing and validation.
 
@@ -522,6 +553,7 @@ def _lse_weighted_combine(
 
     Returns:
         Combined output [B, H, D], and optionally global LSE [B, H]
+
     """
     N, B, H, D = outputs.shape
 
@@ -603,6 +635,9 @@ def _dcp_a2a_pack_send_kernel(
     out_ptr,
     lse_ptr,
     send_ptr,
+    seq_lens_ptr,
+    query_start_loc_ptr,
+    num_seqs,
     out_stride_B,
     out_stride_H,
     out_stride_D,
@@ -616,11 +651,33 @@ def _dcp_a2a_pack_send_kernel(
     HEAD_DIM: tl.constexpr,
     H_PER_RANK: tl.constexpr,
     LSE_PACK_DIM: tl.constexpr,
+    HAS_MASK: tl.constexpr,
+    NUM_SEQS_POW2: tl.constexpr,
 ):
     batch_idx = tl.program_id(0).to(tl.int64)
     local_head_idx = tl.program_id(1).to(tl.int64)
     d_offsets = tl.arange(0, HEAD_DIM)
 
+    # Empty-shard masking, fused in to avoid a separate per-layer pass over the
+    # LSE. Rows that belong to a request with no local KV shard, and rows past
+    # the end of the batch (CUDA-graph padding), must not contribute to the
+    # combine, which the reduction achieves by giving them an LSE of -inf.
+    row_is_empty = False
+    if HAS_MASK:
+        # searchsorted(query_start_loc[1:], batch_idx, right=True), i.e. the
+        # number of request boundaries at or below this row.
+        seq_offsets = tl.arange(0, NUM_SEQS_POW2)
+        boundaries = tl.load(
+            query_start_loc_ptr + 1 + seq_offsets,
+            mask=seq_offsets < num_seqs,
+            other=2**31 - 1,
+        )
+        seq_idx = tl.sum((boundaries <= batch_idx).to(tl.int32), axis=0)
+        seq_idx = tl.minimum(seq_idx, num_seqs - 1)
+        num_rows = tl.load(query_start_loc_ptr + num_seqs)
+        seq_len = tl.load(seq_lens_ptr + seq_idx)
+        row_is_empty = (batch_idx >= num_rows) | (seq_len == 0)
+
     for rank_idx in tl.static_range(N):
         src_head_idx = rank_idx * H_PER_RANK + local_head_idx
         send_base = (
@@ -642,6 +699,8 @@ def _dcp_a2a_pack_send_kernel(
         lse_val = tl.load(
             lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H
         ).to(tl.float32)
+        if HAS_MASK:
+            lse_val = tl.where(row_is_empty, -float("inf"), lse_val)
         if LSE_PACK_DIM == 1:
             tl.store(
                 send_ptr + send_base + HEAD_DIM * send_stride_D,
@@ -797,12 +856,21 @@ def _dcp_a2a_pack_send(
     seq_lens: torch.Tensor | None = None,
     query_start_loc: torch.Tensor | None = None,
 ) -> None:
-    mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
+    # The empty-shard mask is fused into the pack kernel below rather than
+    # applied as a separate pass: the mask is a pure function of seq_lens and
+    # query_start_loc, which are identical for every MLA layer in a step, while
+    # the eager form costs eight launch-bound kernels per layer.
+    has_mask = _validate_dcp_empty_shard_args(seq_lens, query_start_loc)
+    num_seqs = seq_lens.shape[0] if has_mask and seq_lens is not None else 0
+
     grid = (cp_attn_out.shape[0], h_per_rank, 1)
     _dcp_a2a_pack_send_kernel[grid](
         cp_attn_out,
         cp_attn_lse,
         send_buffer,
+        seq_lens,
+        query_start_loc,
+        num_seqs,
         cp_attn_out.stride(0),
         cp_attn_out.stride(1),
         cp_attn_out.stride(2),
@@ -816,6 +884,11 @@ def _dcp_a2a_pack_send(
         HEAD_DIM=head_dim,
         H_PER_RANK=h_per_rank,
         LSE_PACK_DIM=lse_pack_dim,
+        HAS_MASK=has_mask,
+        # Only the arange bound is a constexpr: num_seqs itself is a runtime
+        # argument so the kernel specializes on at most log2(max_num_seqs)
+        # variants instead of one per distinct batch size.
+        NUM_SEQS_POW2=triton.next_power_of_2(max(num_seqs, 1)),
     )
 
 
@@ -872,8 +945,7 @@ def dcp_a2a_lse_reduce(
     seq_lens: torch.Tensor | None = None,
     query_start_loc: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-    """
-    Combine partial attention outputs across DCP ranks using All-to-All.
+    """Combine partial attention outputs across DCP ranks using All-to-All.
 
     The output and LSE are packed into a single output-dtype buffer, sent
     with one All-to-All, then unpacked and combined with exact LSE weighting.
@@ -891,6 +963,7 @@ def dcp_a2a_lse_reduce(
     Returns:
         Combined output [B, H/N, D] (head-scattered)
         If return_lse=True, also returns global_lse [B, H/N]
+
     """
     world_size = cp_group.world_size
 
diff --git a/vllm/v1/attention/ops/flashmla.py b/vllm/v1/attention/ops/flashmla.py
index c84a495859ba..eb5383e5d13a 100644
--- a/vllm/v1/attention/ops/flashmla.py
+++ b/vllm/v1/attention/ops/flashmla.py
@@ -49,9 +49,7 @@ def _is_flashmla_available() -> tuple[bool, str | None]:
 
 
 def is_flashmla_dense_supported() -> tuple[bool, str | None]:
-    """
-    Return: is_supported_flag, unsupported_reason (optional).
-    """
+    """Return: is_supported_flag, unsupported_reason (optional)."""
     is_available, maybe_reason = _is_flashmla_available()
     if not is_available:
         return False, maybe_reason
@@ -61,9 +59,7 @@ def is_flashmla_dense_supported() -> tuple[bool, str | None]:
 
 
 def is_flashmla_sparse_supported() -> tuple[bool, str | None]:
-    """
-    Return: is_supported_flag, unsupported_reason (optional).
-    """
+    """Return: is_supported_flag, unsupported_reason (optional)."""
     is_available, maybe_reason = _is_flashmla_available()
     if not is_available:
         return False, maybe_reason
diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py
index 6791e1f5fd1e..d28c78bdb9fa 100644
--- a/vllm/v1/attention/ops/merge_attn_states.py
+++ b/vllm/v1/attention/ops/merge_attn_states.py
@@ -106,8 +106,8 @@ def merge_attn_states(
             are treated as having context.
         output_scale: Optional scalar tensor for FP8 static quantization.
             When provided, output must be FP8 dtype.
-    """
 
+    """
     # Both the CUDA and Triton kernels derive the suffix head stride from
     # prefix_output, so suffix_output must share the same head stride.
     assert prefix_output.stride(1) == suffix_output.stride(1), (
diff --git a/vllm/v1/attention/ops/metadata.py b/vllm/v1/attention/ops/metadata.py
new file mode 100644
index 000000000000..f5baa51a9297
--- /dev/null
+++ b/vllm/v1/attention/ops/metadata.py
@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Device-side request mapping and sparse indexer metadata."""
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _token_request(query_start_loc, token, num_reqs):
+    lo = tl.full(token.shape, 0, tl.int32)
+    hi = tl.full(token.shape, num_reqs, tl.int32)
+    while tl.sum((lo < hi).to(tl.int32), 0) > 0:
+        mid = (lo + hi) // 2
+        end = tl.load(query_start_loc + mid + 1, mid < num_reqs, other=0x7FFFFFFF)
+        right = (lo < hi) & (end <= token)
+        lo = tl.where(right, mid + 1, lo)
+        hi = tl.where((lo < hi) & ~right, mid, hi)
+    return tl.minimum(lo, num_reqs - 1)
+
+
+@triton.jit(do_not_specialize=["num_reqs", "num_mapped", "num_tokens"])
+def _token_request_mapping_kernel(
+    query_start_loc, output, num_reqs, num_mapped, num_tokens
+):
+    token = tl.program_id(0) * 256 + tl.arange(0, 256)
+    req = _token_request(query_start_loc, token, num_reqs)
+    tl.store(output + token, tl.where(token < num_mapped, req, 0), token < num_tokens)
+
+
+@triton.jit(
+    do_not_specialize=[
+        "num_reqs",
+        "num_tokens",
+        "capacity",
+        "input_stride",
+        "output_stride",
+    ]
+)
+def _indexer_decode_metadata_kernel(
+    query_start_loc,
+    seq_lens,
+    block_table,
+    output_seq_lens,
+    output_block_table,
+    decode_lens,
+    indices,
+    per_req_lens,
+    num_reqs,
+    num_tokens,
+    capacity,
+    input_stride,
+    output_stride,
+    BLOCK_COLS: tl.constexpr,
+):
+    token = tl.program_id(0)
+    cols = tl.arange(0, triton.next_power_of_2(BLOCK_COLS))
+    actual_tokens = tl.load(query_start_loc + num_reqs)
+    if token < num_tokens:
+        req = _token_request(query_start_loc, tl.full((1,), token, tl.int32), num_reqs)
+        req = tl.sum(req, 0)
+        valid = token < actual_tokens
+        start = tl.load(query_start_loc + req)
+        end = tl.load(query_start_loc + req + 1)
+        seq = tl.load(seq_lens + req)
+        length = tl.where(valid, seq - (end - start) + token - start + 1, 0)
+        tl.store(output_seq_lens + token, length)
+        tl.store(decode_lens + token, 1)
+        tl.store(
+            indices + token, tl.where(valid, req, num_reqs + token - actual_tokens)
+        )
+        blocks = tl.load(
+            block_table + req.to(tl.int64) * input_stride + cols,
+            valid & (cols < BLOCK_COLS),
+            other=0,
+        )
+        tl.store(
+            output_block_table + token.to(tl.int64) * output_stride + cols,
+            blocks,
+            cols < BLOCK_COLS,
+        )
+    else:
+        tail = num_tokens + (token - num_tokens) * 256 + tl.arange(0, 256)
+        tl.store(output_seq_lens + tail, 0, tail < capacity)
+    if token < num_reqs:
+        start = tl.load(query_start_loc + token)
+        end = tl.load(query_start_loc + token + 1)
+        tl.store(per_req_lens + token, end - start)
diff --git a/vllm/v1/attention/ops/pcp.py b/vllm/v1/attention/ops/pcp.py
index 75ab1c9e8e13..7e7dbe224575 100644
--- a/vllm/v1/attention/ops/pcp.py
+++ b/vllm/v1/attention/ops/pcp.py
@@ -18,8 +18,15 @@ def _gather_prefill_cache_inputs(
     assert all(tensor.shape[0] == local_num_tokens for tensor in tensors)
     assert 0 <= num_decode_tokens <= local_num_tokens
 
-    if num_decode_tokens == local_num_tokens:
-        return tensors, slot_mapping[:num_decode_tokens]
+    # Replicated draft decodes use unexpanded slot mappings, even with DP padding.
+    if (
+        num_decode_tokens == local_num_tokens
+        or slot_mapping.shape[0] <= local_num_tokens
+    ):
+        return (
+            tuple(tensor[:num_decode_tokens] for tensor in tensors),
+            slot_mapping[:num_decode_tokens],
+        )
 
     pcp_group = get_pcp_group()
     gathered_prefills = tuple(
diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
index 0259e5080fc7..a7e69a8a0c65 100644
--- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
@@ -11,12 +11,12 @@
 
 import vllm.envs as envs
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
-from vllm.config import CUDAGraphMode
+from vllm.config import CUDAGraphMode, get_current_vllm_config
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.torch_utils import LayerNameType
+from vllm.utils.torch_utils import LayerNameType, direct_register_custom_op
 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
 from vllm.v1.worker.workspace import current_workspace_manager
@@ -29,6 +29,8 @@
 
 logger = init_logger(__name__)
 
+FP8_DTYPE = current_platform.fp8_dtype()
+
 
 @functools.cache
 def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | None:
@@ -58,17 +60,25 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None:
 
 _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024
 _GFX950_C4A_NATIVE_MAX_ROWS = 256
+_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024
 # Conservative perf gate, not a correctness bound: OPUS is correct for any query
 # count, but Triton stays faster below this measured crossover.
 _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024
 
 
+def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool:
+    """V4.0 C4A is block-flat (NORMAL). Ratio 1 and 2 are 16×16 SHUFFLE."""
+    return compress_ratio == 4
+
+
 def _get_aiter_top_k_kernel(
     *,
     is_prefill: bool,
     compress_ratio: int,
     num_rows: int,
     max_valid_seq_len: int | None = None,
+    num_columns: int | None = None,
+    topk_tokens: int = 1024,
     on_gfx950: bool = _ON_GFX950,
 ) -> Callable[..., None] | None:
     if compress_ratio <= 1 or not on_gfx950:
@@ -76,6 +86,13 @@ def _get_aiter_top_k_kernel(
 
     if not is_prefill:
         assert max_valid_seq_len is not None
+        if (
+            topk_tokens == 512
+            and 0 < num_rows <= 384
+            and num_columns is not None
+            and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS
+        ):
+            return None
         # AITER v0.1.19 decode is one-block only. This measured gfx950
         # FP32/k=1024 compressed-row boundary is independent of the native
         # split-count boundary in sampler.cu.
@@ -250,8 +267,7 @@ def indexer_k_quant_and_cache_triton(
     # In real layout, we store the first portion as kv cache value
     # and second portion as kv cache scale
     kv_cache = kv_cache.view(num_blocks, -1)
-    fp8_dtype = current_platform.fp8_dtype()
-    kv_cache_value = kv_cache[:, : block_size * head_dim].view(fp8_dtype)
+    kv_cache_value = kv_cache[:, : block_size * head_dim].view(FP8_DTYPE)
     kv_cache_scale = kv_cache[:, block_size * head_dim :].view(torch.float32)
     head_tile_size = head_tile_size // kv_cache.element_size()
     layout = "NORMAL" if block_size == 1 else "SHUFFLE"
@@ -269,7 +285,7 @@ def indexer_k_quant_and_cache_triton(
         layout,
         block_tile_size,
         head_tile_size,
-        IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,
+        IS_FNUZ=torch.float8_e4m3fnuz == FP8_DTYPE,
         USE_UE8M0=scale_fmt == "ue8m0",
     )
 
@@ -292,16 +308,16 @@ def _cp_gather_indexer_quant_cache_kernel(
     HEAD_DIM: tl.constexpr,
     BLOCK_TILE_SIZE: tl.constexpr,
     HEAD_TILE_SIZE: tl.constexpr,
-    NUM_TOKENS: tl.constexpr,
-    NUM_BATCHES: tl.constexpr,
-    BLOCK_TABLE_WIDTH: tl.constexpr,
-    NUM_BLOCKS: tl.constexpr,
+    num_tokens,
+    num_batches,
+    block_table_width,
+    num_blocks,
 ):
     tid = tl.program_id(0)
     offset = tl.arange(0, HEAD_DIM)
-    valid_tid = tid < NUM_TOKENS
+    valid_tid = tid < num_tokens
     batch_id = tl.load(token_to_seq_ptr + tid, mask=valid_tid, other=-1)
-    valid_batch = (batch_id >= 0) & (batch_id < NUM_BATCHES)
+    valid_batch = (batch_id >= 0) & (batch_id < num_batches)
     safe_batch_id = tl.where(valid_batch, batch_id, 0)
     batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0)
     batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0)
@@ -314,7 +330,7 @@ def _cp_gather_indexer_quant_cache_kernel(
     valid_block_table = (
         valid_token
         & (block_table_id >= 0)
-        & (block_table_id < BLOCK_TABLE_WIDTH)
+        & (block_table_id < block_table_width)
         & (block_offset >= 0)
         & (block_offset < block_size)
     )
@@ -323,7 +339,7 @@ def _cp_gather_indexer_quant_cache_kernel(
     block_id = tl.load(
         block_table_ptr + block_table_offset, mask=valid_block_table, other=-1
     )
-    valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS)
+    valid_block = valid_block_table & (block_id >= 0) & (block_id < num_blocks)
     # The packed KV layout makes per-block strides large
     # enough that block_id * stride can exceed 32-bit range.
     safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64)
@@ -447,6 +463,7 @@ def cp_gather_indexer_k_quant_cache_triton(
     token_to_seq: torch.Tensor,
     block_tile_size: int = 16,
     head_tile_size: int = 16,
+    cache_layout: str | None = None,
 ):
     num_tokens = k_fp8.size(0)
     block_size = k_cache.size(1)
@@ -455,12 +472,14 @@ def cp_gather_indexer_k_quant_cache_triton(
     num_blocks = k_cache.shape[0]
     # we assume the kv cache already been split to 2 portion
     k_cache = k_cache.view(num_blocks, -1)
-    fp8_dtype = current_platform.fp8_dtype()
-    k_cache_value = k_cache[:, : block_size * head_dim].view(fp8_dtype)
+    k_cache_value = k_cache[:, : block_size * head_dim].view(FP8_DTYPE)
     k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32)
     grid = (num_tokens,)
     k_fp8_scale = k_fp8_scale.view(torch.float32)
-    layout = "NORMAL" if block_size == 1 else "SHUFFLE"
+    if cache_layout is None:
+        layout = "NORMAL" if block_size == 1 else "SHUFFLE"
+    else:
+        layout = cache_layout
     kernel_args = (
         k_cache_value,
         k_cache_scale,
@@ -506,7 +525,6 @@ def fp8_paged_mqa_logits_torch(
 ):
     from vllm.utils.math_utils import cdiv
 
-    fp8_dtype = current_platform.fp8_dtype()
     batch_size, next_n, _, dim = q.size()
     if next_n == 1:
         block_size = kv_cache.shape[1]
@@ -530,7 +548,7 @@ def fp8_paged_mqa_logits_torch(
             cache = kv_cache_flat[pages]
             scale_offset = block_size * dim
             cache_value = (
-                cache[..., :scale_offset].view(dtype=fp8_dtype).to(torch.float32)
+                cache[..., :scale_offset].view(dtype=FP8_DTYPE).to(torch.float32)
             )
             cache_scale = (
                 cache[..., scale_offset:].view(dtype=torch.float32).contiguous()
@@ -548,7 +566,7 @@ def fp8_paged_mqa_logits_torch(
     kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
     scale = scale.contiguous().view(torch.float)
     q = q.float()
-    kv_cache = kv_cache.view(fp8_dtype).float() * scale
+    kv_cache = kv_cache.view(FP8_DTYPE).float() * scale
     num_block, block_size, _, dim = kv_cache.size()
     logits = torch.full(
         [batch_size * next_n, max_model_len],
@@ -598,6 +616,158 @@ def fp8_paged_mqa_logits_torch(
     return logits
 
 
+@triton.jit
+def _fp8_paged_mqa_logits_decode_kernel(
+    q_ptr,  # fp8 [B, NEXT_N, H, D]
+    kv_val_ptr,  # fp8, block-flat: [num_blocks, block_size*(D+4)]
+    kv_scale_ptr,  # fp32, block-flat: [num_blocks, block_size*(D+4)//4]
+    weights_ptr,  # fp32 [B*NEXT_N, H]
+    ctx_lens_ptr,  # int32 [B*NEXT_N] if CTX_PER_ROW else [B]
+    block_tables_ptr,  # int32 [B, max_blocks]
+    logits_ptr,  # fp32 [B*NEXT_N, max_model_len]
+    stride_q_b,
+    stride_q_n,
+    stride_q_h,
+    stride_w_row,
+    stride_kvblk_fp8,
+    stride_kvblk_f32,
+    scale_region_off,  # block_size*D // 4 (fp32 offset of scale region within a block)
+    stride_bt_b,
+    stride_logits_row,
+    max_blocks,  # block_tables width; guards the block-table gather
+    max_model_len,
+    NUM_HEADS: tl.constexpr,
+    HEAD_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,  # paged-cache page size
+    BLOCK_KV: tl.constexpr,  # positions per tile; must divide BLOCK_SIZE
+    N_SPLITS: tl.constexpr,  # KV-tile parallelism factor (grid dim 1)
+    NEXT_N: tl.constexpr,  # query positions per batch (1 = decode; >1 = MTP verify)
+    CTX_PER_ROW: tl.constexpr,  # ctx_lens is already per-(b, n)
+):
+    # scale_region_off = block_size*D//4 needs D % 4 == 0 (D=128).
+    tl.static_assert(HEAD_SIZE % 4 == 0)
+    # Grid (B*NEXT_N, N_SPLITS): disjoint KV tiles per program, no host sync.
+    row = tl.program_id(0)
+    split = tl.program_id(1)
+    b = row // NEXT_N
+    n = row % NEXT_N
+    if CTX_PER_ROW:
+        seq_len = tl.load(ctx_lens_ptr + row)
+    else:
+        seq_len = tl.load(ctx_lens_ptr + b) - NEXT_N + n + 1
+    seq_len = tl.minimum(tl.maximum(seq_len, 0), max_model_len)
+
+    h = tl.arange(0, NUM_HEADS)
+    d = tl.arange(0, HEAD_SIZE)
+    # keep q/kv fp8 -> fp8 MFMA (f32 accum), not the slow f32 path
+    q = tl.load(
+        q_ptr + b * stride_q_b + n * stride_q_n + h[:, None] * stride_q_h + d[None, :]
+    )
+    w = tl.load(weights_ptr + row * stride_w_row + h).to(tl.float32)  # [H]
+
+    kv_col = tl.arange(0, BLOCK_KV)
+    for kv_start in tl.range(split * BLOCK_KV, seq_len, N_SPLITS * BLOCK_KV):
+        pos = kv_start + kv_col
+        # Tiles are page-aligned (BLOCK_KV | BLOCK_SIZE), so a tile is one page.
+        logical_blk = kv_start // BLOCK_SIZE
+        blk_ok = logical_blk < max_blocks  # guard against a wild page -> fault
+        mask_pos = (pos < seq_len) & blk_ok
+        page = tl.load(
+            block_tables_ptr + b * stride_bt_b + logical_blk, mask=blk_ok, other=0
+        ).to(tl.int64)
+        pos_in_blk = pos - logical_blk * BLOCK_SIZE  # [BLOCK_KV] == pos % BLOCK_SIZE
+
+        # block-flat layout: values region, then scales region.
+        val_off = page * stride_kvblk_fp8 + pos_in_blk[:, None] * HEAD_SIZE + d[None, :]
+        kv = tl.load(
+            kv_val_ptr + val_off, mask=mask_pos[:, None], other=0.0
+        )  # [BLOCK_KV, D] fp8
+        sc = tl.load(
+            kv_scale_ptr + page * stride_kvblk_f32 + scale_region_off + pos_in_blk,
+            mask=mask_pos,
+            other=0.0,
+        )  # [BLOCK_KV]
+
+        # fp8 upcasts exactly; bit-identical on gfx942, ~1e-5 rel on gfx950
+        # (accum order), top-k unchanged.
+        s = tl.dot(kv, tl.trans(q))  # [BLOCK_KV, H]
+        s = tl.maximum(s, 0.0)
+        s = s * w[None, :]
+        s = tl.sum(s, axis=1)  # [BLOCK_KV]
+        s = s * sc
+        tl.store(logits_ptr + row * stride_logits_row + pos, s, mask=mask_pos)
+
+
+def rocm_fp8_paged_mqa_logits_triton(
+    q_fp8: torch.Tensor,
+    kv_cache_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_tables: torch.Tensor,
+    max_model_len: int,
+) -> torch.Tensor:
+    """Triton paged MQA-logits for decode and MTP; matches the torch ref but
+    has no host sync, so it is safe to capture under a full CUDA graph."""
+    batch_size, next_n, num_heads, head_size = q_fp8.shape
+    block_size = kv_cache_fp8.shape[1]
+    BLOCK_KV = 64
+    assert block_size % BLOCK_KV == 0
+
+    fp8_dtype = current_platform.fp8_dtype()
+    num_blocks = kv_cache_fp8.shape[0]
+    kv_flat = kv_cache_fp8.reshape(
+        num_blocks, -1
+    )  # uint8 [num_blocks, block_size*(D+4)]
+    kv_val = kv_flat.view(fp8_dtype)  # [num_blocks, block_size*(D+4)] fp8
+    kv_scale = kv_flat.view(torch.float32)  # [num_blocks, block_size*(D+4)//4] fp32
+
+    cl = context_lens.reshape(-1)
+    ctx_per_row = not (next_n > 1 and cl.numel() == batch_size)
+
+    max_blocks = block_tables.shape[1]
+    (out_logits,) = current_workspace_manager().get_simultaneous(
+        ((batch_size * next_n, max_model_len), torch.float32),
+    )
+
+    # Memory-bound over the KV range: split each row's keys across programs so
+    # few-row / long-context launches still fill the GPU. Cap splits at the
+    # device CU count (304 on gfx942, 256 on gfx950) rather than a gfx950-sized
+    # constant. All terms are static at launch, so the grid stays CUDA-graph-safe.
+    rows = batch_size * next_n
+    tiles_cap = (max_model_len + BLOCK_KV - 1) // BLOCK_KV
+    N_SPLITS = max(1, min(max(1, _decode_cu_count()), tiles_cap, 1024 // rows))
+    _fp8_paged_mqa_logits_decode_kernel[(rows, N_SPLITS)](
+        q_fp8,
+        kv_val,
+        kv_scale,
+        weights,
+        cl,
+        block_tables,
+        out_logits,
+        q_fp8.stride(0),
+        q_fp8.stride(1),
+        q_fp8.stride(2),
+        weights.stride(0),
+        kv_val.stride(0),
+        kv_scale.stride(0),
+        (block_size * head_size) // 4,
+        block_tables.stride(0),
+        out_logits.stride(0),
+        max_blocks,
+        max_model_len,
+        NUM_HEADS=num_heads,
+        HEAD_SIZE=head_size,
+        BLOCK_SIZE=block_size,
+        BLOCK_KV=BLOCK_KV,
+        N_SPLITS=N_SPLITS,
+        NEXT_N=next_n,
+        CTX_PER_ROW=ctx_per_row,
+        num_warps=4,
+        num_stages=2,
+    )
+    return out_logits
+
+
 @functools.lru_cache
 def paged_mqa_logits_module():
     paged_mqa_logits_module_path = None
@@ -623,6 +793,8 @@ def rocm_fp8_paged_mqa_logits(
     block_tables: torch.Tensor,
     schedule_metadata: torch.Tensor,
     max_model_len: int,
+    *,
+    compress_ratio: int = 1,
 ) -> torch.Tensor:
     """Compute FP8 MQA logits using paged KV-cache.
 
@@ -630,8 +802,7 @@ def rocm_fp8_paged_mqa_logits(
         q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to
             `torch.float8_e4m3fn` by caller.
         kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape
-            [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last
-            4 bytes per (block,pos) store the `float` dequant scale.
+            [num_blocks, block_size, 1, D+4], dtype `torch.uint8`.
         weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
         context_lens: Tensor of shape [B], dtype int32; effective context length
             for each batch element.
@@ -640,18 +811,35 @@ def rocm_fp8_paged_mqa_logits(
         schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
             used to distribute work across SMs.
         max_model_len: Maximum sequence length used to size the logits output.
+        compress_ratio: C4A (4) takes block-flat Triton; 1 and 2 stay on AITER.
 
     Returns:
         Logits tensor of shape [B * next_n, max_model_len], dtype
         `torch.float32`.
+
     """
     from vllm._aiter_ops import rocm_aiter_ops
 
-    aiter_paged_mqa_logits_module = None
-    # if rocm_aiter_ops.is_enabled():
     batch_size, next_n = q_fp8.shape[:2]
     block_size = kv_cache_fp8.shape[1]
 
+    # C4A only: Flash/DSv3.2 also skip insert but still write SHUFFLE.
+    if (
+        (_ON_GFX950 or _ON_GFX942)
+        and _indexer_k_is_c4a_block_flat(compress_ratio)
+        and block_size > 1
+    ):
+        if block_size % 64 == 0:
+            return rocm_fp8_paged_mqa_logits_triton(
+                q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
+            )
+        # Non 64-aligned page size (not used in prod): eager torch ref.
+        return fp8_paged_mqa_logits_torch(
+            q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
+        )
+
+    aiter_paged_mqa_logits_module = None
+
     if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():
         aiter_paged_mqa_logits_module = paged_mqa_logits_module()
 
@@ -727,6 +915,7 @@ def fp8_mqa_logits_torch(
 
     Returns:
         Logits tensor of shape [M, N], dtype `torch.float32`.
+
     """
     k_fp8, scale = kv
     seq_len_kv = k_fp8.shape[0]
@@ -795,8 +984,8 @@ def rocm_fp8_mqa_logits(
 
     Returns:
         Logits tensor of shape [M, N], dtype `torch.float32`.
-    """
 
+    """
     from vllm._aiter_ops import rocm_aiter_ops
 
     k_fp8, scale = kv
@@ -819,6 +1008,149 @@ def rocm_fp8_mqa_logits(
         return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
 
 
+# Programs along the column axis of the decode mask grid. The shared kernel in
+# model_executor/kernels/attention/dsa/candidate_blocks.py launches one program
+# per 1024-column tile, which is 1024 of them per row at max_model_len 1048576
+# with only the first few doing work. This is a fixed count that strides
+# instead, measured on gfx950; it is not a portable choice, which is why this
+# variant lives here rather than replacing the shared one.
+_MASK_GRID_COLS = 128
+# Columns each program handles per iteration. Work stops at the tile boundary
+# containing a row's end rather than at the end itself; the overshoot is
+# masked the same way the shared kernel masks it, so it costs a tile and
+# changes nothing.
+_MASK_TILE = 1024
+
+
+@triton.jit(do_not_specialize=["width", "nblocks"])
+def _mask_candidates_strided_kernel(
+    logits,
+    starts,
+    ends,
+    flags,
+    stride_row,
+    stride_col,
+    stride_start,
+    stride_end,
+    width,
+    nblocks,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_STARTS: tl.constexpr,
+    ROW_REPEAT: tl.constexpr,
+    TILE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
+    end = tl.load(ends + row // ROW_REPEAT * stride_end)
+    edge = tl.load(flags + row * (nblocks + 1) + nblocks)
+    step = tl.num_programs(1) * TILE
+    tile_start = tl.program_id(1) * TILE
+    # The shared kernel also sanitizes columns at or past `end`. Nothing reads
+    # them: top_k_per_row_decode bounds its scan by the same row ends passed
+    # here, and persistent_topk/cooperative_topk clamp by the same lengths.
+    # Stopping at `end` is what makes the cost track the live context rather
+    # than max_model_len.
+    while tile_start < end:
+        cols = tile_start + tl.arange(0, TILE)
+        valid = (cols >= start) & (cols < end) & (cols < width)
+        block = (cols - start) // BLOCK_SIZE
+        keep = tl.load(flags + row * (nblocks + 1) + block, valid, other=0)
+        keep = (keep != 0) | ((cols == width - 1) & (edge != 0))
+        tl.store(
+            logits + row * stride_row + cols * stride_col,
+            -float("inf"),
+            (cols < width) & ~(valid & keep),
+        )
+        tile_start += step
+
+
+def _apply_candidate_mask_strided(
+    logits: torch.Tensor,
+    row_ks: torch.Tensor | None,
+    row_ke: torch.Tensor,
+    candidate_blocks: torch.Tensor,
+    block_size: int,
+    row_repeat: int = 1,
+) -> None:
+    """ROCm decode variant of ``apply_candidate_mask``.
+
+    Same masking semantics over ``[0, end)``, but the grid is sized by a fixed
+    program count rather than by the logits width. Only worth using where the
+    width is the ``max_model_len`` workspace and the live context is far
+    shorter, i.e. the paged decode path below; the prefill chunks pass
+    chunk-sized logits and stay on the shared kernel.
+    """
+    from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+        _candidate_flags_kernel,
+    )
+
+    rows, width = logits.shape
+    if not rows or not width:
+        return
+    nblocks = triton.cdiv(width, block_size)
+    flags = torch.empty((rows, nblocks + 1), device=logits.device, dtype=torch.uint8)
+    start_stride = row_ks.stride(0) if row_ks is not None else 0
+    _candidate_flags_kernel[(rows,)](
+        candidate_blocks,
+        row_ks,
+        flags,
+        *candidate_blocks.stride(),
+        start_stride,
+        width,
+        nblocks,
+        block_size,
+        candidate_blocks.shape[1],
+        row_ks is not None,
+        row_repeat,
+    )
+    # Derived from width, which is a tensor shape, so the grid stays static and
+    # a FULL cudagraph capture remains valid across replays; only the loop trip
+    # count inside the kernel is data-dependent. The min keeps narrow widths
+    # from launching programs that would only fall through.
+    grid_cols = min(_MASK_GRID_COLS, triton.cdiv(width, _MASK_TILE))
+    _mask_candidates_strided_kernel[(rows, grid_cols)](
+        logits,
+        row_ks,
+        row_ke,
+        flags,
+        *logits.stride(),
+        start_stride,
+        row_ke.stride(0),
+        width,
+        nblocks,
+        block_size,
+        row_ks is not None,
+        row_repeat,
+        _MASK_TILE,
+    )
+
+
+def _max_decode_logits_rows(num_batched_tokens: int) -> int:
+    """Upper bound on decode rows the paged-MQA logits buffer can ever hold.
+
+    ``rocm_fp8_paged_mqa_logits`` sizes its workspace as
+    ``(batch_size * next_n, max_model_len)``. ``batch_size`` is bounded by
+    ``max_num_seqs`` and ``next_n`` by ``1 + num_speculative_tokens``, which is
+    far tighter than ``max_num_batched_tokens`` -- 192 vs 16384 for a typical
+    32-seq DSpark-5 deployment. The loose bound is harmless at short contexts
+    but scales with ``max_model_len``, so at the model's full context it asks
+    for tens of TiB and the engine cannot start. Take whichever valid bound is
+    smaller; the workspace is locked after profiling, so it must not be under-
+    estimated.
+    """
+    try:
+        vllm_config = get_current_vllm_config()
+    except Exception:
+        return num_batched_tokens
+    scheduler_config = getattr(vllm_config, "scheduler_config", None)
+    max_num_seqs = getattr(scheduler_config, "max_num_seqs", None)
+    if not max_num_seqs:
+        return num_batched_tokens
+    speculative_config = getattr(vllm_config, "speculative_config", None)
+    num_spec = getattr(speculative_config, "num_speculative_tokens", 0) or 0
+    return min(num_batched_tokens, max_num_seqs * (1 + num_spec))
+
+
 def rocm_aiter_sparse_attn_indexer_fake(
     hidden_states: torch.Tensor,
     k_cache_prefix: LayerNameType,
@@ -872,7 +1204,6 @@ def rocm_aiter_sparse_attn_indexer(
     # careful! this will be None in dummy run
     forward_context = get_forward_context()
     attn_metadata = forward_context.attn_metadata
-    fp8_dtype = current_platform.fp8_dtype()
     from vllm.utils.torch_utils import _resolve_layer_name
 
     k_cache_prefix = _resolve_layer_name(k_cache_prefix)
@@ -888,20 +1219,20 @@ def rocm_aiter_sparse_attn_indexer(
         # Prefill k_fp8 and k_scale buffers, used by
         # rocm_aiter_sparse_attn_indexer's prefill path
         workspace_manager.get_simultaneous(
-            ((total_seq_lens, head_dim), fp8_dtype),
+            ((total_seq_lens, head_dim), FP8_DTYPE),
             ((total_seq_lens, 4), torch.uint8),
         )
 
         # Decode logits buffer, used by rocm_fp8_paged_mqa_logits.
-        # batch_size * next_n <= hidden_states.shape[0] == max_num_batched_tokens
+        decode_rows = _max_decode_logits_rows(hidden_states.shape[0])
         if _ON_GFX942 or _ON_GFX950:
             workspace_manager.get_simultaneous(
-                ((hidden_states.shape[0], max_model_len), torch.float32),
+                ((decode_rows, max_model_len), torch.float32),
             )
         else:
             workspace_manager.get_simultaneous(
                 (
-                    (q_fp8.shape[1], hidden_states.shape[0], max_model_len),
+                    (q_fp8.shape[1], decode_rows, max_model_len),
                     torch.float32,
                 ),
             )
@@ -931,6 +1262,9 @@ def rocm_aiter_sparse_attn_indexer(
             topk_indices_buffer,
             skip_k_cache_insert,
             compress_ratio,
+            candidate_blocks,
+            candidate_block_size,
+            candidate_write,
         )
     layer_attn_metadata = attn_metadata[k_cache_prefix]
     assert isinstance(layer_attn_metadata, DeepseekV32IndexerMetadata)
@@ -965,7 +1299,7 @@ def rocm_aiter_sparse_attn_indexer(
 
         workspace_manager = current_workspace_manager()
         k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
-            ((total_seq_lens, head_dim), fp8_dtype),
+            ((total_seq_lens, head_dim), FP8_DTYPE),
             ((total_seq_lens, 4), torch.uint8),
         )
         for chunk in prefill_metadata.chunks:
@@ -978,6 +1312,9 @@ def rocm_aiter_sparse_attn_indexer(
                 chunk.block_table,
                 chunk.cu_seq_lens,
                 token_to_seq=chunk.token_to_seq,
+                cache_layout=(
+                    "NORMAL" if _indexer_k_is_c4a_block_flat(compress_ratio) else None
+                ),
             )
             logits = rocm_fp8_mqa_logits(
                 q_fp8[chunk.token_start : chunk.token_end],
@@ -986,6 +1323,30 @@ def rocm_aiter_sparse_attn_indexer(
                 chunk.cu_seqlen_ks,
                 chunk.cu_seqlen_ke,
             )
+            if candidate_blocks is not None:
+                from vllm.model_executor.layers.sparse_attn_indexer import (
+                    _apply_candidate_mask,
+                    _select_candidate_blocks,
+                )
+
+                chunk_candidates = candidate_blocks[chunk.token_start : chunk.token_end]
+                if candidate_write:
+                    _select_candidate_blocks(
+                        logits,
+                        chunk.cu_seqlen_ks,
+                        chunk.cu_seqlen_ke,
+                        chunk_candidates.shape[1],
+                        candidate_block_size,
+                        chunk_candidates,
+                    )
+                else:
+                    _apply_candidate_mask(
+                        logits,
+                        chunk.cu_seqlen_ks,
+                        chunk.cu_seqlen_ke,
+                        chunk_candidates,
+                        candidate_block_size,
+                    )
             topk_indices = topk_indices_buffer[
                 chunk.token_start : chunk.token_end, :topk_tokens
             ]
@@ -1051,8 +1412,39 @@ def rocm_aiter_sparse_attn_indexer(
             decode_metadata.block_table,
             decode_metadata.schedule_metadata,
             max_model_len=max_model_len,
+            compress_ratio=compress_ratio,
         )
 
+        if candidate_blocks is not None:
+            from vllm.model_executor.layers.sparse_attn_indexer import (
+                _select_candidate_blocks,
+            )
+
+            num_rows = logits.shape[0]
+            visible = decode_metadata.seq_lens.reshape(-1)
+            if visible.numel() != num_rows:
+                visible = visible.repeat_interleave(next_n)
+            visible = visible[:num_rows].to(torch.int64)
+            row_starts = torch.zeros_like(visible)
+            decode_candidates = candidate_blocks[:num_rows]
+            if candidate_write:
+                _select_candidate_blocks(
+                    logits,
+                    row_starts,
+                    visible,
+                    decode_candidates.shape[1],
+                    candidate_block_size,
+                    decode_candidates,
+                )
+            else:
+                _apply_candidate_mask_strided(
+                    logits,
+                    row_starts,
+                    visible,
+                    decode_candidates,
+                    candidate_block_size,
+                )
+
         topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
         num_rows = logits.shape[0]
 
@@ -1067,6 +1459,8 @@ def rocm_aiter_sparse_attn_indexer(
             compress_ratio=compress_ratio,
             num_rows=num_rows,
             max_valid_seq_len=max_compressed_seq_len,
+            num_columns=logits.shape[1],
+            topk_tokens=topk_tokens,
         )
         if aiter_topk_kernel is not None:
             _launch_aiter_top_k_per_row_decode(
@@ -1109,6 +1503,10 @@ def _decode_e8m0_scales(scale: torch.Tensor) -> torch.Tensor:
         )
 
         return _upcast_e8m0_to_fp32(scale).contiguous()
+    if scale.dtype == torch.uint8:
+        # MXFP8 parameters preserve E8M0 scales as their raw exponent bytes.
+        # They are biased exponents, not numeric uint8 scale values.
+        return torch.exp2(scale.to(torch.int16).to(torch.float32) - 127.0)
     return scale.to(torch.float32)
 
 
@@ -1180,8 +1578,13 @@ def _fused_inverse_rope_gptj(
     positions: torch.Tensor,
     cos_sin_cache: torch.Tensor,
     rope_head_dim: int,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """bf16 inverse GPT-J RoPE via a single fused Triton kernel."""
+    """bf16 inverse GPT-J RoPE via a single fused Triton kernel.
+
+    ``out`` may alias ``o``: the rotation is a per-row bijection whose kernel
+    reads both lanes of a pair before storing either.
+    """
     assert o.dim() == 3 and o.stride(-1) == 1, (
         "_fused_inverse_rope_gptj expects a [T, H, D] input with a contiguous last dim"
     )
@@ -1193,9 +1596,14 @@ def _fused_inverse_rope_gptj(
         f"[P, {rope_head_dim}] = cos | sin, got {tuple(cos_sin_cache.shape)}"
     )
     num_tokens, num_heads, head_dim = o.shape
-    out = torch.empty(
-        (num_tokens, num_heads, head_dim), dtype=torch.bfloat16, device=o.device
-    )
+    if out is None:
+        out = torch.empty(
+            (num_tokens, num_heads, head_dim), dtype=torch.bfloat16, device=o.device
+        )
+    else:
+        assert out.dtype == torch.bfloat16, (
+            f"inverse RoPE writes bf16, got an output buffer of {out.dtype}"
+        )
     if num_tokens == 0:
         return out
     _inverse_rope_gptj_kernel[(num_tokens, num_heads)](
@@ -1216,6 +1624,109 @@ def _fused_inverse_rope_gptj(
     return out
 
 
+def rocm_inverse_rope_rows_(
+    o: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    rope_head_dim: int,
+) -> None:
+    """Inverse-RoPE attention output rows in place.
+
+    For rows no attention kernel rotated in its epilogue. Call it from the
+    eager attention segment: which rows still owe a rotation depends on the
+    prefill/decode split, and the o_proj that used to do this runs inside the
+    compiled region, where a batch-dependent Python value would be frozen at
+    trace time.
+    """
+    if o.shape[0] == 0:
+        return
+    _fused_inverse_rope_gptj(o, positions, cos_sin_cache, rope_head_dim, out=o)
+
+
+@triton.jit
+def _inverse_rope_mxfp8_quant_kernel(
+    o_ptr,  # [T, H, D] bf16
+    q_ptr,  # [T, H * D] e4m3
+    s_ptr,  # [T, H * D // 32] uint8 (E8M0)
+    pos_ptr,
+    cos_sin_ptr,
+    s_t,
+    s_h,
+    qs_t,
+    ss_t,
+    cs_stride,
+    HEAD_DIM: tl.constexpr,
+    NOPE: tl.constexpr,
+    HALF: tl.constexpr,
+):
+    t = tl.program_id(0)
+    h = tl.program_id(1)
+    lanes = tl.arange(0, HEAD_DIM)
+    x = tl.load(o_ptr + t * s_t + h * s_h + lanes).to(tl.float32)
+    # Same lane-pair formulation as the reduce epilogue: NoPE pairs take
+    # cos=1 / sin=0, so one expression rotates the whole row.
+    pos = tl.load(pos_ptr + t).to(tl.int64)
+    pair_idx = tl.arange(0, HEAD_DIM // 2) - (NOPE // 2)
+    is_rope = pair_idx >= 0
+    k = tl.where(is_rope, pair_idx, 0)
+    cos = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + k), 1.0)
+    sin = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + HALF + k), 0.0)
+    even, odd = tl.split(tl.reshape(x, (HEAD_DIM // 2, 2)))
+    x = tl.reshape(
+        tl.join(even * cos + odd * sin, odd * cos - even * sin), (1, HEAD_DIM)
+    )
+    xq, bits = _mxfp8_quantize_rows(x, 1, HEAD_DIM)
+    tl.store(
+        q_ptr + t * qs_t + h * HEAD_DIM + lanes,
+        tl.reshape(xq, (HEAD_DIM,)).to(q_ptr.dtype.element_ty),
+    )
+    scale_offsets = tl.arange(0, HEAD_DIM // 32)
+    tl.store(
+        s_ptr + t * ss_t + h * (HEAD_DIM // 32) + scale_offsets,
+        tl.reshape(bits, (HEAD_DIM // 32,)).to(tl.uint8),
+    )
+
+
+def rocm_inverse_rope_mxfp8_rows(
+    o: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    rope_head_dim: int,
+    out_data: torch.Tensor,
+    out_scale: torch.Tensor,
+) -> None:
+    """Inverse-RoPE bf16 attention rows and MXFP8-quantize them for wo_a.
+
+    The counterpart of ``rocm_inverse_rope_rows_`` for layers whose attention
+    output is MXFP8: rows the decode reduce did not emit (prefill) go through
+    here. ``o`` is [T, H, D]; ``out_data`` [T, H * D] e4m3 and ``out_scale``
+    [T, H * D // 32] E8M0, the layout the reduce epilogue writes.
+    """
+    num_tokens, num_heads, head_dim = o.shape
+    if num_tokens == 0:
+        return
+    assert o.stride(-1) == 1 and out_data.stride(-1) == 1 and out_scale.stride(-1) == 1
+    assert out_data.shape == (num_tokens, num_heads * head_dim)
+    assert out_scale.shape == (num_tokens, num_heads * head_dim // 32)
+    assert cos_sin_cache.shape[-1] == rope_head_dim
+    _inverse_rope_mxfp8_quant_kernel[(num_tokens, num_heads)](
+        o,
+        out_data,
+        out_scale,
+        positions,
+        cos_sin_cache,
+        o.stride(0),
+        o.stride(1),
+        out_data.stride(0),
+        out_scale.stride(0),
+        cos_sin_cache.stride(0),
+        HEAD_DIM=head_dim,
+        NOPE=head_dim - rope_head_dim,
+        HALF=rope_head_dim // 2,
+        num_warps=4,
+    )
+
+
 def _get_cached_wo_a_bf16(
     wo_a: torch.nn.Module,
     n_local_groups: int,
@@ -1238,7 +1749,15 @@ def _get_cached_wo_a_bf16(
     )
 
     wo_a_scale_param = get_fp8_block_weight_scale(wo_a)
-    if wo_a_scale_param is not None:
+    if wo_a_scale_param is None:
+        # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the
+        # historical ``_inv`` suffix.
+        wo_a_scale_param = getattr(wo_a, "weight_scale", None)
+    # Emulated MXFP8 kernels can replace the original one-byte weight with an
+    # already-dequantized BF16 tensor while retaining the scale attribute for
+    # metadata. Applying that retained scale again would double-dequantize the
+    # weight. Block scaling is only valid while the one-byte FP8 storage remains.
+    if wo_a_scale_param is not None and wo_a.weight.element_size() == 1:
         wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
             torch.float32
         )
@@ -1264,16 +1783,27 @@ def rocm_inv_rope_einsum(
     n_local_groups: int,
     o_lora_rank: int,
     wo_a: torch.nn.Module,
+    inverse_rope: bool = True,
 ) -> torch.Tensor:
     """Inverse-RoPE + WO_A bmm path used on ROCm.
 
     Fuses the inverse GPT-J RoPE into one Triton kernel and caches the bf16
-    wo_a weight so the per-step dequant disappears.
+    wo_a weight so the per-step dequant disappears. Callers whose attention
+    already rotated every row pass ``inverse_rope=False``; that is a property
+    of the attention backend, not of the batch, so it stays constant across
+    steps and is safe to read from compiled code.
     """
-    o_ref = _fused_inverse_rope_gptj(
-        o, positions, rotary_emb.cos_sin_cache, rope_head_dim
-    )
-    o_ref = o_ref.view(o.shape[0], n_local_groups, -1)
+    if inverse_rope:
+        o_ref = _fused_inverse_rope_gptj(
+            o, positions, rotary_emb.cos_sin_cache, rope_head_dim
+        )
+    else:
+        assert o.dtype == torch.bfloat16, (
+            "a pre-rotated attention output feeds the wo_a bmm directly, so it "
+            f"must already be bf16, got {o.dtype}"
+        )
+        o_ref = o
+    o_ref = o_ref.reshape(o.shape[0], n_local_groups, -1)
 
     wo_a_weight = _get_cached_wo_a_bf16(
         wo_a, n_local_groups, o_lora_rank, o_ref.shape[-1]
@@ -1282,6 +1812,175 @@ def rocm_inv_rope_einsum(
     return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
 
 
+@triton.jit
+def _mxfp8_wo_a_bmm_kernel(
+    a_ptr,  # [T, G * K] e4m3
+    as_ptr,  # [T, G * K // 32] E8M0
+    w_ptr,  # [G * R, K] e4m3
+    ws_ptr,  # [G * R, K // 32] E8M0
+    out_ptr,  # [T, G * R]
+    num_tokens,
+    stride_at,
+    stride_ast,
+    stride_wn,
+    stride_wsn,
+    stride_out,
+    R: tl.constexpr,
+    K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    # Output columns are group-major, so a BLOCK_N tile never straddles two
+    # groups (R % BLOCK_N == 0) and picks its group's K slice of A.
+    pid_n = tl.program_id(0)
+    pid_m = tl.program_id(1)
+    group = (pid_n * BLOCK_N) // R
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    m_mask = offs_m < num_tokens
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_sk = tl.arange(0, BLOCK_K // 32)
+
+    a_ptrs = a_ptr + offs_m[:, None] * stride_at + (group * K + offs_k)[None, :]
+    as_ptrs = (
+        as_ptr + offs_m[:, None] * stride_ast + (group * (K // 32) + offs_sk)[None, :]
+    )
+    w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :]
+    ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :]
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for _ in range(K // BLOCK_K):
+        a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
+        a_s = tl.load(as_ptrs, mask=m_mask[:, None], other=127)
+        w = tl.load(w_ptrs)
+        w_s = tl.load(ws_ptrs)
+        acc += tl.dot_scaled(a, a_s, "e4m3", w.T, w_s, "e4m3")
+        a_ptrs += BLOCK_K
+        as_ptrs += BLOCK_K // 32
+        w_ptrs += BLOCK_K
+        ws_ptrs += BLOCK_K // 32
+
+    tl.store(
+        out_ptr + offs_m[:, None] * stride_out + offs_n[None, :],
+        acc.to(out_ptr.dtype.element_ty),
+        mask=m_mask[:, None],
+    )
+
+
+def _mxfp8_wo_a_bmm_config(num_tokens: int, n_groups: int) -> tuple[int, ...]:
+    """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) for gfx950.
+
+    Tuned under HIP graphs with a cold weight at G = 4 and 2, over every
+    decode shape of conc 1-128 x 0-5 spec tokens plus prefill chunks up to
+    8K tokens. The best tile tracks the total work T * G, so the tiers are
+    keyed on it.
+
+    This will be replaced after new GEMM kernel from AITER with proper 32x32 scale
+    shape GEMM fp8 enabled.
+    """
+    work = num_tokens * n_groups
+    if work <= 64:
+        return 16, 16, 1024, 2, 3
+    if work <= 128:
+        return 32, 16, 1024, 2, 3
+    if work <= 256:
+        return 32, 32, 512, 2, 3
+    if work <= 512:
+        return 64, 32, 512, 2, 3
+    if work <= 1024:
+        return 64, 64, 512, 4, 2
+    if work <= 2048:
+        return 64, 64, 256, 4, 2
+    if work <= 3072:
+        return 64, 64, 256, 2, 1
+    if work <= 4096:
+        return 128, 128, 256, 8, 2
+    return 128, 128, 128, 4, 2
+
+
+def _rocm_mxfp8_wo_a_bmm_impl(
+    a: torch.Tensor,
+    a_scale: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    n_groups: int,
+    o_lora_rank: int,
+) -> torch.Tensor:
+    num_tokens = a.shape[0]
+    group_dim = a.shape[1] // n_groups
+    out = torch.empty(
+        (num_tokens, n_groups * o_lora_rank), dtype=torch.bfloat16, device=a.device
+    )
+    if num_tokens == 0:
+        return out
+    block_m, block_n, block_k, num_warps, num_stages = _mxfp8_wo_a_bmm_config(
+        num_tokens, n_groups
+    )
+    grid = (n_groups * o_lora_rank // block_n, triton.cdiv(num_tokens, block_m))
+    _mxfp8_wo_a_bmm_kernel[grid](
+        a,
+        a_scale,
+        weight,
+        weight_scale,
+        out,
+        num_tokens,
+        a.stride(0),
+        a_scale.stride(0),
+        weight.stride(0),
+        weight_scale.stride(0),
+        out.stride(0),
+        R=o_lora_rank,
+        K=group_dim,
+        BLOCK_M=block_m,
+        BLOCK_N=block_n,
+        BLOCK_K=block_k,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return out
+
+
+def _rocm_mxfp8_wo_a_bmm_fake(
+    a: torch.Tensor,
+    a_scale: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    n_groups: int,
+    o_lora_rank: int,
+) -> torch.Tensor:
+    return a.new_empty((a.shape[0], n_groups * o_lora_rank), dtype=torch.bfloat16)
+
+
+# An opaque op: the tile choice branches on the token count, which the
+# compiled o_proj must not freeze at its trace-time value.
+direct_register_custom_op(
+    op_name="rocm_dsv41_mxfp8_wo_a_bmm",
+    op_func=_rocm_mxfp8_wo_a_bmm_impl,
+    fake_impl=_rocm_mxfp8_wo_a_bmm_fake,
+)
+
+
+def rocm_mxfp8_wo_a_bmm(
+    a: torch.Tensor,
+    a_scale: torch.Tensor,
+    wo_a: torch.nn.Module,
+    n_groups: int,
+    o_lora_rank: int,
+) -> torch.Tensor:
+    """Grouped MXFP8 wo_a: ``out[t, g, :] = a[t, g, :] @ W[g].T``, bf16 out.
+
+    ``a`` is the [T, G * K] e4m3 attention output and ``a_scale`` its
+    [T, G * K // 32] E8M0 scales, as the sparse decode reduce writes them.
+    The weight is the checkpoint's MXFP8 ``wo_a`` as loaded, [G * R, K] with
+    [G * R, K // 32] scales, so there is no dequantized copy to keep.
+    Returns [T, G * R].
+    """
+    return torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm(
+        a, a_scale, wo_a.weight, wo_a.weight_scale, n_groups, o_lora_rank
+    )
+
+
 _DSV4_SPARSE_NOPE_DIM = 448
 _DSV4_SPARSE_ROPE_DIM = 64
 
@@ -2482,6 +3181,33 @@ def _sparse_attn_decode_gfx950_partial_kernel(
     )
 
 
+@triton.jit
+def _mxfp8_scale_bits(amax):
+    """Biased E8M0 exponent that puts ``amax`` at the top of the e4m3 range.
+
+    Same rounding as ``mxfp8_e4m3_quantize``, so the output is bit-identical to
+    quantizing the tensor there.
+    """
+    amax = tl.maximum(amax, 1.1754943508222875e-38)
+    bits = tl.ceil(tl.log2(amax / 448.0)) + 127.0
+    return tl.minimum(tl.maximum(bits, 0.0), 254.0)
+
+
+@triton.jit
+def _mxfp8_quantize_rows(x, ROWS: tl.constexpr, COLS: tl.constexpr):
+    """MXFP8-quantize ``x`` [ROWS, COLS] in registers, one scale per 32 lanes.
+
+    Returns the rescaled fp32 values (to be cast to e4m3 on store) and the
+    [ROWS, COLS // 32] biased E8M0 exponents.
+    """
+    blocks = tl.reshape(x, (ROWS, COLS // 32, 32))
+    bits = _mxfp8_scale_bits(tl.max(tl.abs(blocks), axis=2))
+    # Multiply by the reciprocal: a divisor of 2**-127 would be subnormal and
+    # flush to zero, turning an all-zero block into NaN.
+    q = blocks * tl.exp2(127.0 - bits)[:, :, None]
+    return tl.reshape(q, (ROWS, COLS)), bits
+
+
 @triton.jit
 def _sparse_attn_decode_reduce_kernel(
     part_m_ptr,
@@ -2489,13 +3215,19 @@ def _sparse_attn_decode_reduce_kernel(
     part_acc_ptr,
     attn_sink_ptr,
     out_ptr,
+    pos_ptr,
+    cos_sin_ptr,
+    out_scale_ptr,
     out_stride0,
     out_stride1,
+    os_stride0,
+    os_stride1,
     pm_stride0,
     pm_stride_s,
     pa_stride0,
     pa_stride_s,
     pa_stride_h,
+    cs_stride,
     num_heads,
     HAS_ATTN_SINK: tl.constexpr,
     ADAPTIVE_SPLITS: tl.constexpr,
@@ -2503,6 +3235,10 @@ def _sparse_attn_decode_reduce_kernel(
     BLOCK_H: tl.constexpr,
     NUM_SPLITS: tl.constexpr,
     SPLITS_PAD: tl.constexpr,
+    FUSE_INV_ROPE: tl.constexpr,
+    NOPE: tl.constexpr,
+    HALF: tl.constexpr,
+    QUANT_OUT: tl.constexpr,
 ):
     query_idx = tl.program_id(0)
     pid_h = tl.program_id(1)
@@ -2592,12 +3328,45 @@ def _sparse_attn_decode_reduce_kernel(
 
     out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0)
 
+    if FUSE_INV_ROPE:
+        # Inverse GPT-J RoPE on the trailing rope lanes, straight out of the
+        # combine registers: out_even = a*cos + b*sin, out_odd = b*cos - a*sin.
+        # NoPE lanes take cos=1/sin=0 so one expression covers the whole row
+        # and the o_proj rotation pass disappears for these tokens.
+        pos = tl.load(pos_ptr + query_idx).to(tl.int64)
+        pair_idx = tl.arange(0, COMB_DIM // 2) - (NOPE // 2)
+        is_rope = pair_idx >= 0
+        k = tl.where(is_rope, pair_idx, 0)
+        cos = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + k), 1.0)
+        sin = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + HALF + k), 0.0)
+        even, odd = tl.split(tl.reshape(out, (BLOCK_H, COMB_DIM // 2, 2)))
+        out = tl.reshape(
+            tl.join(
+                even * cos[None, :] + odd * sin[None, :],
+                odd * cos[None, :] - even * sin[None, :],
+            ),
+            (BLOCK_H, COMB_DIM),
+        )
+
     out_row_ptr = (
         out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1
     )
+    if QUANT_OUT:
+        # Emit wo_a's MXFP8 input directly instead of a bf16 row that a
+        # separate pass would read back to quantize.
+        out, scale_bits = _mxfp8_quantize_rows(out, BLOCK_H, COMB_DIM)
+        scale_offsets = tl.arange(0, COMB_DIM // 32)
+        tl.store(
+            out_scale_ptr
+            + query_idx * os_stride0
+            + head_offsets[:, None] * os_stride1
+            + scale_offsets[None, :],
+            scale_bits.to(tl.uint8),
+            mask=head_mask[:, None],
+        )
     tl.store(
         out_row_ptr + comb_offsets[None, :],
-        out,
+        out.to(out_ptr.dtype.element_ty),
         mask=head_mask[:, None],
     )
 
@@ -2912,7 +3681,16 @@ def _rocm_sparse_attn_decode_ragged_triton(
     out: torch.Tensor | None = None,
     extra_cache_nan_free: bool = False,
     adaptive_splits: bool = False,
+    inv_rope_positions: torch.Tensor | None = None,
+    inv_rope_cos_sin_cache: torch.Tensor | None = None,
+    out_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
 ) -> torch.Tensor:
+    """Split-K sparse decode; returns the attention output.
+
+    With ``out_mxfp8 = (data, scale)`` the reduce writes MXFP8 instead of
+    bf16: ``data`` is [b, h * d] e4m3 and ``scale`` [b, h * d // 32] E8M0, and
+    ``data`` viewed as [b, h, d] is returned.
+    """
     assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}"
     assert main_cache.ndim == 3, (
         f"expected main_cache=[blocks,block,bytes], got {main_cache.shape}"
@@ -2977,7 +3755,28 @@ def _rocm_sparse_attn_decode_ragged_triton(
         extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32)
 
     block_h = 16
-    if out is None:
+    out_scale = None
+    if out_mxfp8 is not None:
+        assert out is None, "out and out_mxfp8 are mutually exclusive"
+        assert _ON_GFX950, "the MXFP8 reduce epilogue is gfx950-only"
+        assert inv_rope_positions is not None, (
+            "the MXFP8 output feeds wo_a, so it must be inverse-RoPE'd first"
+        )
+        out_data, out_scale = out_mxfp8
+        assert out_data.dtype == torch.float8_e4m3fn and out_scale.dtype == (
+            torch.uint8
+        ), f"expected e4m3/uint8 MXFP8 buffers, got {out_data.dtype}/{out_scale.dtype}"
+        assert out_data.shape == (num_queries, num_heads * head_dim), (
+            f"expected MXFP8 data [{num_queries}, {num_heads * head_dim}], "
+            f"got {tuple(out_data.shape)}"
+        )
+        assert out_scale.shape == (num_queries, num_heads * head_dim // 32), (
+            f"expected MXFP8 scale [{num_queries}, {num_heads * head_dim // 32}], "
+            f"got {tuple(out_scale.shape)}"
+        )
+        assert out_data.stride(-1) == 1 and out_scale.stride(-1) == 1
+        out = out_data.view(num_queries, num_heads, head_dim)
+    elif out is None:
         out = torch.empty_like(q, dtype=torch.bfloat16)
     else:
         assert out.shape == q.shape, f"expected out shape {q.shape}, got {out.shape}"
@@ -3154,19 +3953,38 @@ def _rocm_sparse_attn_decode_ragged_triton(
             num_warps=4,
         )
 
+    fuse_inv_rope = inv_rope_positions is not None
+    if fuse_inv_rope:
+        assert inv_rope_cos_sin_cache is not None
+        assert inv_rope_cos_sin_cache.shape[-1] == rope_head_dim, (
+            "fused inverse RoPE expects cos_sin_cache laid out as "
+            f"[P, {rope_head_dim}] = cos | sin, got "
+            f"{tuple(inv_rope_cos_sin_cache.shape)}"
+        )
+        assert nope_head_dim % 2 == 0 and rope_head_dim % 2 == 0, (
+            "fused inverse RoPE pairs adjacent lanes, so both head dims must "
+            f"be even, got nope={nope_head_dim} rope={rope_head_dim}"
+        )
+
     _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)](
         part_m,
         part_l,
         part_acc,
         attn_sink,
         out,
+        inv_rope_positions,
+        inv_rope_cos_sin_cache,
+        out_scale,
         out.stride(0),
         out.stride(1),
+        out_scale.stride(0) if out_scale is not None else 0,
+        head_dim // 32,
         part_m.stride(0),
         part_m.stride(1),
         part_acc.stride(0),
         part_acc.stride(1),
         part_acc.stride(2),
+        inv_rope_cos_sin_cache.stride(0) if inv_rope_cos_sin_cache is not None else 0,
         num_heads,
         HAS_ATTN_SINK=has_attn_sink,
         ADAPTIVE_SPLITS=adaptive_splits,
@@ -3174,6 +3992,10 @@ def _rocm_sparse_attn_decode_ragged_triton(
         BLOCK_H=1,
         NUM_SPLITS=num_splits,
         SPLITS_PAD=triton.next_power_of_2(num_splits),
+        FUSE_INV_ROPE=fuse_inv_rope,
+        NOPE=nope_head_dim,
+        HALF=rope_head_dim // 2,
+        QUANT_OUT=out_scale is not None,
         num_warps=4,
     )
     return out
@@ -3198,6 +4020,9 @@ def _rocm_sparse_attn_decode_triton(
     out: torch.Tensor | None = None,
     extra_cache_nan_free: bool = False,
     adaptive_splits: bool = False,
+    inv_rope_positions: torch.Tensor | None = None,
+    inv_rope_cos_sin_cache: torch.Tensor | None = None,
+    out_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
 ) -> torch.Tensor:
     if main_ragged_indices is None or main_ragged_indptr is None:
         main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense(
@@ -3236,6 +4061,9 @@ def _rocm_sparse_attn_decode_triton(
         out=out,
         extra_cache_nan_free=extra_cache_nan_free,
         adaptive_splits=adaptive_splits,
+        inv_rope_positions=inv_rope_positions,
+        inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,
+        out_mxfp8=out_mxfp8,
     )
 
 
@@ -3334,10 +4162,25 @@ def rocm_sparse_attn_decode(
     head_dim: int,
     nope_head_dim: int,
     rope_head_dim: int,
-    output: torch.Tensor,
+    output: torch.Tensor | None,
     extra_cache_nan_free: bool = False,
     adaptive_splits: bool = False,
-) -> None:
+    inv_rope_positions: torch.Tensor | None = None,
+    inv_rope_cos_sin_cache: torch.Tensor | None = None,
+    output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
+) -> int:
+    """Run sparse MLA decode into ``output``.
+
+    Passing ``inv_rope_positions`` folds the inverse RoPE into the reduce
+    epilogue. Returns how many leading rows of ``output`` came back rotated,
+    so a caller mixing in a decode path that does not fuse still knows what it
+    owes the standalone pass. Read it from the eager attention segment only.
+
+    ``output_mxfp8 = (data, scale)`` replaces ``output``: the reduce also
+    MXFP8-quantizes the rotated rows for the FP8 wo_a (see
+    ``_rocm_sparse_attn_decode_ragged_triton``). It needs gfx950 and the fused
+    inverse RoPE, and always covers every row.
+    """
     assert swa_k_cache.dtype == torch.uint8, (
         "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, "
         f"got {swa_k_cache.dtype}"
@@ -3366,7 +4209,12 @@ def rocm_sparse_attn_decode(
         if topk_indices is not None:
             extra_indices = topk_indices.reshape(topk_indices.shape[0], -1)
 
-    direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None
+    if output_mxfp8 is not None:
+        assert output is None, "output and output_mxfp8 are mutually exclusive"
+        direct_out = None
+    else:
+        assert output is not None
+        direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None
     attn_out = _rocm_sparse_attn_decode_triton(
         q=q,
         main_cache=swa_k_cache,
@@ -3386,6 +4234,13 @@ def rocm_sparse_attn_decode(
         out=direct_out,
         extra_cache_nan_free=extra_cache_nan_free,
         adaptive_splits=adaptive_splits,
+        inv_rope_positions=inv_rope_positions,
+        inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,
+        out_mxfp8=output_mxfp8,
     )
+    if output_mxfp8 is not None:
+        return q.shape[0]
+    assert output is not None
     if direct_out is None:
         output.copy_(attn_out.to(output.dtype))
+    return output.shape[0] if inv_rope_positions is not None else 0
diff --git a/vllm/v1/attention/ops/triton_attention_helpers.py b/vllm/v1/attention/ops/triton_attention_helpers.py
index dff0b21c6423..7804b013171e 100644
--- a/vllm/v1/attention/ops/triton_attention_helpers.py
+++ b/vllm/v1/attention/ops/triton_attention_helpers.py
@@ -28,12 +28,16 @@ def cdiv_fn(x, y):
 def apply_softcap(S, x):
     """Softcap (aka tanh-style clamp) used to bound attention scores.
 
-    ``x * tanh(S / x)`` rewritten to avoid a direct ``tanh`` call.
+    ``x * tanh(S / x)`` computed without a direct ``tanh`` call.
+
+    Use ``tanh(y) = 1 - 2 / (exp(2y) + 1)`` rather than
+    ``(exp(y) - exp(-y)) / (exp(y) + exp(-y))``. The symmetric-exponential
+    form overflows to ``inf`` for ``|y| > ~88`` and then evaluates to
+    ``inf / inf = NaN``, poisoning the attention row. The single-exponential
+    form saturates cleanly at +/-1 for either sign.
     """
-    Sdiv = S / x
-    p1 = tl.exp(Sdiv)
-    p2 = tl.exp(-Sdiv)
-    return x * (p1 - p2) / (p1 + p2)
+    y = S / x
+    return x * (1.0 - 2.0 / (tl.exp(2.0 * y) + 1.0))
 
 
 # ===========================================================================
diff --git a/vllm/v1/attention/ops/triton_decode_attention.py b/vllm/v1/attention/ops/triton_decode_attention.py
index 6aec1db2e59c..130760d1ded8 100644
--- a/vllm/v1/attention/ops/triton_decode_attention.py
+++ b/vllm/v1/attention/ops/triton_decode_attention.py
@@ -24,8 +24,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""
-Memory-efficient attention for decoding.
+"""Memory-efficient attention for decoding.
 It supports page size >= 1.
 """
 
diff --git a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py b/vllm/v1/attention/ops/triton_fp8_mqa_logits.py
deleted file mode 100644
index 619d0ec50a9d..000000000000
--- a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py
+++ /dev/null
@@ -1,262 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Temporary gfx942 fallback for AITER's fp8_mqa_logits kernel.
-
-This module vendors AITER's Triton fp8_mqa_logits kernel with the gfx942
-tile-size workaround from ROCm/aiter#3257. It is used only while vLLM's
-pinned AITER version lacks that fix.
-
-TODO: Remove this vendored copy once vLLM pins an AITER version that includes
-ROCm/aiter#3257 bugfix for gfx942.
-"""
-
-import torch
-
-from vllm.triton_utils import tl, triton
-
-# gfx942 (MI300X) has 64 KiB of LDS per CU. We accept the default
-# (BLOCK_KV=128, num_stages=2) tile only when *both* of these hold:
-#
-# 1. Occupancy gate. With waves_per_eu=2 and num_warps=4 we target two
-#    workgroups co-resident on a CU -> per-WG LDS budget = 32 KiB. Triton
-#    keeps Q in registers (loop-invariant) and the fp32 scores accumulator
-#    in VGPRs (heavy VALU), so only the double-buffered KV tile is
-#    expected to live in LDS. A 0.9 safety factor leaves headroom for any
-#    LDS overhead the compiler may add.
-#
-# 2. Hardware ceiling. Defensive upper bound that also counts Q and
-#    scores against the 64 KiB CU limit, in case a Triton version (older
-#    or future) decides to spill them to LDS. False positives here only
-#    shrink the tile; false negatives are JIT-aborts, so we lean
-#    conservative.
-_GFX942_CU_LDS_BYTES = 64 * 1024
-_GFX942_PER_WG_LDS_BUDGET_BYTES = _GFX942_CU_LDS_BYTES * 9 // 20  # ~28.8 KiB
-
-
-def _gfx942_default_tile_fits_lds(num_heads: int, head_size: int) -> bool:
-    """Return True iff (BLOCK_KV=128, num_stages=2) fits in MI300X LDS."""
-    BLOCK_KV = 128
-    NUM_STAGES = 2
-    kv_bytes = head_size * BLOCK_KV * NUM_STAGES
-    scores_bytes = num_heads * BLOCK_KV * 4
-    q_bytes = num_heads * head_size
-    fits_occupancy = kv_bytes < _GFX942_PER_WG_LDS_BUDGET_BYTES
-    fits_hardware = q_bytes + kv_bytes + scores_bytes <= _GFX942_CU_LDS_BYTES
-    return fits_occupancy and fits_hardware
-
-
-@triton.jit
-def _fp8_mqa_logits_kernel(
-    Q_ptr,  # fp8e4m3 [seq_len, H, D]
-    KV_ptr,  # fp8e4m3 [seq_len_kv, D]
-    kv_scales_ptr,  # fp32 [seq_len_kv]
-    weights_ptr,  # fp32 [seq_len, H]
-    cu_start_ptr,  # int32 [seq_len]
-    cu_end_ptr,  # int32 [seq_len]
-    logits_ptr,  # fp32 [seq_len, seq_len_kv]
-    seq_len,
-    seq_len_kv,
-    NUM_HEADS: tl.constexpr,
-    HEAD_SIZE: tl.constexpr,
-    # strides
-    stride_q_s: tl.int64,
-    stride_q_h: tl.constexpr,
-    stride_q_d: tl.constexpr,
-    stride_kv_s: tl.int64,
-    stride_kv_d: tl.constexpr,
-    stride_w_s: tl.int64,
-    stride_w_h: tl.constexpr,
-    stride_logits_s: tl.int64,
-    stride_logits_k: tl.int64,
-    # block sizes
-    BLOCK_KV: tl.constexpr,
-):
-    row_id = tl.program_id(0)
-    # go from larger to smaller in terms of work
-    # to reduce the tail effect
-    row_id = tl.num_programs(0) - row_id - 1
-    tl.assume(row_id >= 0)
-    tl.assume(stride_q_s > 0)
-    tl.assume(stride_q_h > 0)
-    tl.assume(stride_q_d > 0)
-    tl.assume(stride_kv_s > 0)
-    tl.assume(stride_kv_d > 0)
-    tl.assume(stride_w_s > 0)
-    tl.assume(stride_w_h > 0)
-
-    logits_row_ptrs = logits_ptr + row_id * stride_logits_s
-
-    h_inds = tl.arange(0, NUM_HEADS)[:, None]
-    d_inds = tl.arange(0, HEAD_SIZE)
-
-    # load Q[BLOCK_Q, NUM_HEADS, HEAD_SIZE]
-    q_ptrs = (
-        Q_ptr + row_id * stride_q_s + h_inds * stride_q_h + d_inds[None, :] * stride_q_d
-    )
-
-    q_block = tl.load(q_ptrs, cache_modifier=".cg")
-    w_ptrs = weights_ptr + row_id * stride_w_s + h_inds * stride_w_h
-    w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32)
-
-    # Load start/end for each row in this block
-    start_ind = tl.load(cu_start_ptr + row_id)
-    end_ind = tl.load(cu_end_ptr + row_id)
-
-    start_ind = tl.maximum(start_ind, 0)
-    end_ind = tl.minimum(end_ind, seq_len_kv)
-    shifted_end = end_ind - start_ind
-    shifted_unmasked_end = shifted_end // BLOCK_KV * BLOCK_KV
-
-    kv_col_offsets = tl.arange(0, BLOCK_KV) + start_ind
-    kv_ptrs = (
-        KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d
-    )
-
-    kv_scales_ptrs = kv_scales_ptr + kv_col_offsets
-
-    logits_ptrs = logits_row_ptrs + kv_col_offsets * stride_logits_k
-
-    # Loop over KV tiles
-    for _ in tl.range(0, shifted_unmasked_end, BLOCK_KV):
-        kv_block = tl.load(kv_ptrs)
-        kv_scales = tl.load(kv_scales_ptrs)
-
-        # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
-        scores = tl.dot(q_block, kv_block, input_precision="ieee")
-        # Multiply by kv_scales (broadcast along rows)
-        scores = scores * kv_scales[None, :]
-        # ReLU
-        scores = tl.maximum(scores, 0.0)
-        scores = scores * w_block
-        # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]
-        scores = tl.sum(scores, axis=0)
-        tl.store(logits_ptrs, scores)
-
-        kv_ptrs += BLOCK_KV * stride_kv_s
-        kv_scales_ptrs += BLOCK_KV
-        logits_ptrs += BLOCK_KV * stride_logits_k
-        kv_col_offsets += BLOCK_KV
-
-    # masked load
-    kv_col_mask = kv_col_offsets < end_ind
-    kv_block = tl.load(kv_ptrs, mask=kv_col_mask[None, :], other=0.0)
-    kv_scales = tl.load(kv_scales_ptrs, mask=kv_col_mask, other=0.0)
-
-    # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV]
-    scores = tl.dot(q_block, kv_block, input_precision="ieee")
-    # Multiply by kv_scales (broadcast along rows)
-    scores = scores * kv_scales[None, :]
-    # ReLU
-    scores = tl.maximum(scores, 0.0)
-    scores = scores * w_block
-    # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ]
-    scores = tl.sum(scores, axis=0)
-    # masked store
-    in_window = (kv_col_offsets >= start_ind) & (kv_col_offsets < end_ind)
-    tl.store(logits_ptrs, scores, mask=in_window)
-
-
-def fp8_mqa_logits_gfx942(
-    q: torch.Tensor,
-    k_fp8: torch.Tensor,
-    kv_scales: torch.Tensor,
-    weights: torch.Tensor,
-    cu_starts: torch.Tensor,
-    cu_ends: torch.Tensor,
-) -> torch.Tensor:
-    """Compute FP8 MQA logits on MI300X (gfx942) using the vendored kernel.
-
-    Drop-in replacement for ``aiter.ops.triton.attention.fp8_mqa_logits.
-    fp8_mqa_logits`` on MI300X. Selects ``(BLOCK_KV, num_stages)`` based on
-    whether the default tile fits within the 64 KiB LDS budget of a gfx942
-    CU (see module docstring).
-
-    Args:
-        q: Query tensor of shape ``[M, H, D]``, FP8 dtype.
-        k_fp8: Key tensor of shape ``[N, D]``, FP8 dtype.
-        kv_scales: K scales of shape ``[N]`` (or ``[N, 1]`` -- viewed as
-            ``[N]``), float32.
-        weights: Per-head weights of shape ``[M, H]``, float32.
-        cu_starts: Start indices (inclusive) of shape ``[M]``, int32.
-        cu_ends: End indices (exclusive) of shape ``[M]``, int32.
-
-    Returns:
-        Logits of shape ``[M, N]``, float32 -- positions outside
-        ``[cu_starts[i], cu_ends[i])`` for row ``i`` are pre-filled with
-        ``-inf`` so the caller can run a top-k without masking.
-    """
-    seq_len, num_heads, head_size = q.shape
-    seq_len_kv = k_fp8.shape[0]
-    assert num_heads & (num_heads - 1) == 0, (
-        f"num_heads must be a power of two (got {num_heads})"
-    )
-    assert head_size & (head_size - 1) == 0, (
-        f"head_size must be a power of two (got {head_size})"
-    )
-
-    # The kernel walks ``kv_scales`` as a 1-D contiguous array of size N
-    # (it indexes by ``kv_scales_ptr + kv_col_offsets``). The vLLM caller
-    # passes a ``[N, 4]`` uint8 view-cast-to-float32 which lands as
-    # ``[N, 1]`` contiguous -- byte-identical to ``[N]`` -- but flatten
-    # explicitly to keep the kernel's pointer arithmetic intent clear.
-    kv_scales_1d = kv_scales.reshape(-1)
-
-    # Initialise with -inf so positions outside [cu_starts, cu_ends) read
-    # as ``-inf`` after the masked store path -- this matches AITER's
-    # ``fp8_mqa_logits`` semantics and is what the top-k consumer expects.
-    logits = torch.full(
-        (seq_len, seq_len_kv),
-        fill_value=-float("inf"),
-        dtype=torch.float32,
-        device=q.device,
-    )
-
-    if _gfx942_default_tile_fits_lds(num_heads, head_size):
-        block_kv = 128
-        num_stages = 2
-    else:
-        # DSv4 sparse indexer (NUM_HEADS=64, HEAD_SIZE=128) lands here:
-        # default tile spills past gfx942's 64 KiB LDS budget. (64, 1)
-        # needs ~33 KiB and clears the per-WG budget with margin.
-        block_kv = 64
-        num_stages = 1
-
-    # heuristic for MFMA instruction shape, identical to AITER's choice
-    matrix_instr_nonkdim = 32
-    if seq_len <= 1024:
-        matrix_instr_nonkdim = 16
-
-    stride_q_s, stride_q_h, stride_q_d = q.stride()
-    stride_kv_s, stride_kv_d = k_fp8.stride()
-    stride_w_s, stride_w_h = weights.stride()
-    stride_logits_s, stride_logits_k = logits.stride()
-
-    _fp8_mqa_logits_kernel[(seq_len,)](
-        Q_ptr=q,
-        KV_ptr=k_fp8,
-        kv_scales_ptr=kv_scales_1d,
-        weights_ptr=weights,
-        cu_start_ptr=cu_starts,
-        cu_end_ptr=cu_ends,
-        logits_ptr=logits,
-        seq_len=seq_len,
-        seq_len_kv=seq_len_kv,
-        NUM_HEADS=num_heads,
-        HEAD_SIZE=head_size,
-        stride_q_s=stride_q_s,
-        stride_q_h=stride_q_h,
-        stride_q_d=stride_q_d,
-        stride_kv_s=stride_kv_s,
-        stride_kv_d=stride_kv_d,
-        stride_w_s=stride_w_s,
-        stride_w_h=stride_w_h,
-        stride_logits_s=stride_logits_s,
-        stride_logits_k=stride_logits_k,
-        BLOCK_KV=block_kv,
-        num_warps=4,
-        num_stages=num_stages,
-        waves_per_eu=2,
-        matrix_instr_nonkdim=matrix_instr_nonkdim,
-    )
-
-    return logits
diff --git a/vllm/v1/attention/ops/triton_prefill_attention.py b/vllm/v1/attention/ops/triton_prefill_attention.py
index 0fb31ce7e1dc..1ed55c3e28b5 100644
--- a/vllm/v1/attention/ops/triton_prefill_attention.py
+++ b/vllm/v1/attention/ops/triton_prefill_attention.py
@@ -19,8 +19,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""
-Memory-efficient attention for prefill.
+"""Memory-efficient attention for prefill.
 It supports page size = 1.
 """
 
@@ -33,6 +32,16 @@
 from vllm.utils.math_utils import RCP_LN2
 
 
+def _prefer_narrow_kv_tile() -> bool:
+    """RDNA3/RDNA4 prefer a narrower KV tile than the shared default."""
+    if not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import on_gfx1x
+
+    # RDNA3/RDNA4 only: on_gfx1x() excludes gfx10xx and gfx1250.
+    return on_gfx1x()
+
+
 @triton.jit
 def _fwd_kernel(
     Q,
@@ -229,8 +238,7 @@ def context_attention_fwd(
     sinks: torch.Tensor | None = None,
     sinks_bias_key0: bool = False,
 ):
-    """
-    q, k, v: [b * s, head, head_dim]
+    """q, k, v: [b * s, head, head_dim]
     b_start_loc: [b]
     b_seq_len: [b]
     out: [b * s, head, head_dim]
@@ -251,6 +259,10 @@ def context_attention_fwd(
     grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
     num_warps = 4 if Lk <= 64 else 8
 
+    # BLOCK_M, num_warps and num_stages stay at the shared defaults; min()
+    # leaves dtypes whose default tile is already 32, such as float32, alone.
+    BLOCK_N = min(BLOCK, 32) if _prefer_narrow_kv_tile() else BLOCK
+
     sliding_window_q = sliding_window_q if sliding_window_q is not None else 0
     sliding_window_k = sliding_window_k if sliding_window_k is not None else 0
 
@@ -274,7 +286,7 @@ def context_attention_fwd(
         kv_group_num=kv_group_num,
         BLOCK_M=BLOCK,
         BLOCK_DMODEL=triton.next_power_of_2(Lk),
-        BLOCK_N=BLOCK,
+        BLOCK_N=BLOCK_N,
         IS_CAUSAL=is_causal,
         SLIDING_WINDOW_Q=sliding_window_q,
         SLIDING_WINDOW_K=sliding_window_k,
diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py
index 321f42bcb199..68ce0e0de2e2 100644
--- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py
+++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py
@@ -21,15 +21,11 @@
 """
 
 import math
-from typing import Any
 
 import torch
 
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.v1.attention.ops.triton_decode_attention import _fwd_kernel_stage2
-
-from .triton_turboquant_decode import _use_fp8_e4b15
 
 # ROCm prefers num_stages=1 in attention-like kernels to reduce shared-memory
 # pressure (mirrors the pattern used in triton_decode_attention.py and
@@ -66,387 +62,6 @@ def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
     return lut.reshape(N * N, 2).contiguous()
 
 
-# ---------------------------------------------------------------------------
-# Stage 1 v2: Grouped-Q + pair LUT + tl.dot + exp2
-# ---------------------------------------------------------------------------
-
-
-@triton.jit
-def _tq_decode_stage1_v2(
-    Q_rot_ptr,
-    KV_cache_ptr,
-    KV_cache_u16_ptr,  # uint16 view — Opt#3 SoA loads
-    Block_table_ptr,
-    Seq_lens_ptr,
-    Centroids_ptr,
-    Pair_lut_ptr,
-    Mid_o_ptr,
-    stride_qb,
-    stride_qh,
-    stride_cache_block,  # bytes per block (bs*H*slot_aligned)
-    stride_bt_b,
-    stride_mid_b,
-    stride_mid_h,
-    stride_mid_s,
-    NUM_Q_HEADS: tl.constexpr,
-    NUM_KV_HEADS: tl.constexpr,
-    HEAD_DIM: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-    PADDED_SLOT: tl.constexpr,
-    MAX_NUM_BLOCKS: tl.constexpr,
-    NUM_KV_SPLITS: tl.constexpr,
-    KV_GROUP_SIZE: tl.constexpr,
-    MSE_BITS: tl.constexpr,
-    MSE_BYTES: tl.constexpr,
-    VQB: tl.constexpr,
-    VAL_DATA_BYTES: tl.constexpr,
-    # Opt#3 SoA layout
-    KEY_DATA_BYTES: tl.constexpr,
-    META_REGION_OFFSET: tl.constexpr,
-    NUM_SOA_FIELDS: tl.constexpr,
-    SOA_K_NORM: tl.constexpr,
-    SOA_V_SCALE: tl.constexpr,
-    SOA_V_ZERO: tl.constexpr,
-    N_CENTROIDS: tl.constexpr,
-    ATTN_SCALE: tl.constexpr,
-    BLOCK_D: tl.constexpr,
-    TILE_SIZE: tl.constexpr,
-    BLOCK_M: tl.constexpr,
-    KEY_FP8: tl.constexpr,
-    NORM_CORRECTION: tl.constexpr = 0,
-    FP8_E4B15: tl.constexpr = 0,
-    USE_PAIR_LUT: tl.constexpr = 0,
-    USE_BF16_DOT: tl.constexpr = 0,
-):
-    # Grid: (B, Hk, NUM_KV_SPLITS) — one program per KV head group
-    bid = tl.program_id(0)
-    kv_hid = tl.program_id(1)
-    sid = tl.program_id(2)
-
-    seq_len = tl.load(Seq_lens_ptr + bid)
-    split_len = tl.cdiv(seq_len, NUM_KV_SPLITS)
-    split_start = split_len * sid
-    split_end = tl.minimum(split_start + split_len, seq_len)
-
-    if split_start >= split_end:
-        return
-
-    # Dimension and tile offsets
-    d_offs = tl.arange(0, BLOCK_D)
-    d_mask = d_offs < HEAD_DIM
-    m_offs = tl.arange(0, BLOCK_M)
-    t_offs = tl.arange(0, TILE_SIZE)
-
-    # Q head indices for this KV group: [kv_hid * GS, kv_hid * GS + GS)
-    q_head_offs = kv_hid * KV_GROUP_SIZE + m_offs
-    q_mask = m_offs < KV_GROUP_SIZE
-
-    # Load Q tile: [BLOCK_M, BLOCK_D] float32
-    q_addrs = bid * stride_qb + q_head_offs[:, None] * stride_qh + d_offs[None, :]
-    Q = tl.load(
-        Q_rot_ptr + q_addrs,
-        mask=q_mask[:, None] & d_mask[None, :],
-        other=0.0,
-    )
-
-    # Pre-warm centroids in L1
-    if not KEY_FP8:
-        _c_warm = tl.load(Centroids_ptr + tl.arange(0, N_CENTROIDS))  # noqa: F841
-
-    # Loop-invariant bit-extraction constants for MSE keys
-    if not KEY_FP8:
-        if MSE_BITS == 4:
-            pass  # 4-bit byte extraction is computed inline in the tile loop
-        else:
-            # Generic bit extraction
-            mse_bit_off = d_offs * MSE_BITS
-            mse_byte_idx = mse_bit_off // 8
-            mse_bit_shift = mse_bit_off % 8
-            mse_mask_val = (1 << MSE_BITS) - 1
-
-    # Loop-invariant bit-extraction for 3-bit values
-    if VQB == 3:
-        val_bit_off = d_offs * 3
-        val_byte_idx = val_bit_off // 8
-        val_bit_shift = val_bit_off % 8
-
-    # exp2 pre-scaling
-    RCP_LN2: tl.constexpr = 1.4426950408889634
-    LN2: tl.constexpr = 0.6931471805599453
-    QK_SCALE = ATTN_SCALE * RCP_LN2
-
-    # Online softmax accumulators: [BLOCK_M] for max, [BLOCK_M] for L
-    M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
-    L = tl.zeros([BLOCK_M], dtype=tl.float32)
-
-    acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32)
-
-    bt_base = bid * stride_bt_b
-
-    # ================================================================
-    # TILED LOOP: TILE_SIZE KV tokens per iteration
-    # ================================================================
-    for start_n in range(split_start, split_end, TILE_SIZE):
-        kv_offs = start_n + t_offs
-        kv_mask_1d = kv_offs < split_end
-
-        page_idx = kv_offs // BLOCK_SIZE
-        page_off = kv_offs % BLOCK_SIZE
-        block_nums = tl.load(
-            Block_table_ptr + bt_base + page_idx,
-            mask=kv_mask_1d,
-            other=0,
-        ).to(tl.int64)
-
-        # Opt#3 SoA addressing: data region then per-block SoA metadata.
-        slot_within_block = page_off.to(tl.int64)
-        block_base = block_nums * stride_cache_block
-        DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES
-        data_bases = (
-            block_base
-            + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)
-            + tl.cast(kv_hid, tl.int64) * DATA_BYTES_PER_SLOT
-        )
-        head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(
-            kv_hid, tl.int64
-        ) * (NUM_SOA_FIELDS * BLOCK_SIZE)
-        knorm_u16_addrs = (
-            head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block
-        )
-        vscale_u16_addrs = (
-            head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block
-        )
-        vzero_u16_addrs = (
-            head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block
-        )
-
-        # ============================================================
-        # KEY DEQUANT → K_T: [BLOCK_D, TILE_SIZE] for tl.dot(Q, K_T)
-        # ============================================================
-        if KEY_FP8:
-            k_addrs = data_bases[:, None] + d_offs[None, :]
-            k_raw = tl.load(
-                KV_cache_ptr + k_addrs,
-                mask=kv_mask_1d[:, None] & d_mask[None, :],
-                other=0,
-            )
-            if FP8_E4B15:
-                K_f = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
-            else:
-                K_f = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
-            # K_f: [TILE_SIZE, BLOCK_D] → transpose for tl.dot
-            K_T = tl.trans(K_f)
-            # S: [BLOCK_M, TILE_SIZE]  (bf16 dot on ROCm for better dynamic range)
-            if USE_BF16_DOT:
-                S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16))
-            else:
-                S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16))
-        else:
-            # --- MSE dequantization ---
-            if MSE_BITS == 4 and USE_PAIR_LUT:
-                # OPT#1 (load-halving): FLUTE §3.2 pair LUT, but with each
-                # packed byte loaded EXACTLY ONCE instead of twice (the
-                # previous code built byte_addrs from d_offs // 2 which has
-                # the pattern [0,0,1,1,2,2,...], duplicating every byte
-                # load). We load the packed bytes as a [TILE_SIZE, HALF_D]
-                # tile, decode both nibbles, then issue a single 3-D gather
-                # against pair_lut which returns (T[lo], T[hi]) in slots
-                # 0/1. Reshape [TILE_SIZE, HALF_D, 2] -> [TILE_SIZE, BLOCK_D]
-                # interleaves the pairs into the expected dim order
-                # (d=2k -> T[lo_k], d=2k+1 -> T[hi_k]).
-                HALF_D: tl.constexpr = BLOCK_D // 2
-                half_offs = tl.arange(0, HALF_D)
-                # Each byte k covers dims (2k, 2k+1); it is in-range iff
-                # 2k < HEAD_DIM. HEAD_DIM is always even in our layouts,
-                # so this is equivalent to k < HEAD_DIM // 2.
-                byte_mask = (half_offs * 2) < HEAD_DIM
-
-                byte_addrs = data_bases[:, None] + half_offs[None, :]
-                byte_raw = tl.load(
-                    KV_cache_ptr + byte_addrs,
-                    mask=kv_mask_1d[:, None] & byte_mask[None, :],
-                    other=0,
-                ).to(tl.int32)
-
-                lo_idx = byte_raw & 0xF  # [TILE_SIZE, HALF_D]
-                hi_idx = (byte_raw >> 4) & 0xF  # [TILE_SIZE, HALF_D]
-                pair_key = lo_idx * N_CENTROIDS + hi_idx  # [TILE_SIZE, HALF_D]
-
-                # Single 3-D gather: pair_lut[pair_key*2 + {0,1}]
-                # -> [TILE_SIZE, HALF_D, 2] = (T[lo_k], T[hi_k]) per element.
-                pair_slot = tl.arange(0, 2)
-                c_pair = tl.load(
-                    Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :],
-                    mask=(kv_mask_1d[:, None, None] & byte_mask[None, :, None]),
-                    other=0.0,
-                )
-                # Row-major reshape interleaves pairs:
-                #   out[t, 2k]   = c_pair[t, k, 0] = T[lo_k]
-                #   out[t, 2k+1] = c_pair[t, k, 1] = T[hi_k]
-                c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D])
-            elif MSE_BITS == 4:
-                # 4-bit without pair LUT: exploit byte alignment
-                half_idx = d_offs // 2
-                nibble_shift = (d_offs % 2) * 4
-                mse_addrs = data_bases[:, None] + half_idx[None, :]
-                mse_raw = tl.load(
-                    KV_cache_ptr + mse_addrs,
-                    mask=kv_mask_1d[:, None] & d_mask[None, :],
-                    other=0,
-                ).to(tl.int32)
-                mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF
-                c_vals = tl.load(
-                    Centroids_ptr + mse_idx,
-                    mask=kv_mask_1d[:, None] & d_mask[None, :],
-                    other=0.0,
-                )
-            else:
-                # Generic 3-bit path: two byte loads → 16-bit → shift/mask
-                mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :]
-                mse_raw0 = tl.load(
-                    KV_cache_ptr + mse_addrs0,
-                    mask=kv_mask_1d[:, None] & d_mask[None, :],
-                    other=0,
-                ).to(tl.int32)
-                mse_raw1 = tl.load(
-                    KV_cache_ptr + mse_addrs0 + 1,
-                    mask=kv_mask_1d[:, None] & d_mask[None, :],
-                    other=0,
-                ).to(tl.int32)
-                raw16 = mse_raw0 | (mse_raw1 << 8)
-                mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val
-                c_vals = tl.load(
-                    Centroids_ptr + mse_idx,
-                    mask=kv_mask_1d[:, None] & d_mask[None, :],
-                    other=0.0,
-                )
-
-            # Opt#1: norm-correction is pre-folded into the stored per-token
-            # scalar at store time, so no per-tile sum+sqrt+divide here. The
-            # stored vec_norm carries ||k||/||c_vec||; multiply by c_vals
-            # below gives the original c_vec/||c_vec|| * ||k||.
-
-            # Opt#3: K-norms from per-block SoA region. Single u16 load per
-            # token; contiguous across tokens within a block.
-            norm_u16 = tl.load(
-                KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask_1d, other=0
-            )
-            vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)
-
-            # Reconstruct K: K_recon[t, d] = norm[t] * centroid[t, d]
-            K_recon = c_vals * vec_norms[:, None]
-            # K_T: [BLOCK_D, TILE_SIZE]
-            K_T = tl.trans(K_recon)
-
-            # S: [BLOCK_M, TILE_SIZE] via tensor core MMA
-            if USE_BF16_DOT:
-                S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16))
-            else:
-                S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16))
-
-        # Mask out-of-range positions
-        S = tl.where(kv_mask_1d[None, :], S, float("-inf"))
-
-        # ============================================================
-        # ONLINE SOFTMAX with exp2
-        # ============================================================
-        m_j = tl.maximum(M, tl.max(S, axis=1))
-        m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
-        P = tl.math.exp2(S - m_j[:, None])
-        l_j = tl.sum(P, axis=1)
-        alpha = tl.math.exp2(M - m_j)
-
-        acc = acc * alpha[:, None]
-        L = L * alpha + l_j
-        M = m_j
-
-        # ============================================================
-        # VALUE DEQUANT → V: [TILE_SIZE, BLOCK_D] for tl.dot(P, V)
-        # Opt#3: V data lives at data_base + KEY_DATA_BYTES; scale/zero
-        # move to the SoA metadata region (one u16 load each per token).
-        # ============================================================
-        val_bases = data_bases + KEY_DATA_BYTES
-
-        if VQB == 3:
-            val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]
-            val_raw0 = tl.load(
-                KV_cache_ptr + val_addrs0,
-                mask=kv_mask_1d[:, None] & d_mask[None, :],
-                other=0,
-            ).to(tl.int32)
-            val_raw1 = tl.load(
-                KV_cache_ptr + val_addrs0 + 1,
-                mask=kv_mask_1d[:, None] & d_mask[None, :],
-                other=0,
-            ).to(tl.int32)
-            raw16 = val_raw0 | (val_raw1 << 8)
-            v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)
-        else:  # VQB == 4
-            # OPT#3 (load-halving, value path): same pattern as OPT#1 for
-            # keys -- load each byte exactly once as [TILE_SIZE, HALF_D],
-            # decode both nibbles, then tl.interleave to [TILE_SIZE, BLOCK_D]
-            # with pattern [v_lo0, v_hi0, v_lo1, v_hi1, ...].
-            V_HALF_D: tl.constexpr = BLOCK_D // 2
-            v_half_offs = tl.arange(0, V_HALF_D)
-            v_byte_mask = (v_half_offs * 2) < HEAD_DIM
-
-            val_addrs = val_bases[:, None] + v_half_offs[None, :]
-            val_byte = tl.load(
-                KV_cache_ptr + val_addrs,
-                mask=kv_mask_1d[:, None] & v_byte_mask[None, :],
-                other=0,
-            ).to(tl.int32)
-            v_lo = (val_byte & 0xF).to(tl.float32)  # [T, HALF_D]
-            v_hi = ((val_byte >> 4) & 0xF).to(tl.float32)  # [T, HALF_D]
-            v_idx = tl.interleave(v_lo, v_hi)  # [T, BLOCK_D]
-
-        scale_u16 = tl.load(
-            KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask_1d, other=0
-        )
-        zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask_1d, other=0)
-        v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32)
-        v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32)
-        V = v_idx * v_scales[:, None] + v_zeros[:, None]
-
-        # P·V accumulation via tensor core MMA
-        # P: [BLOCK_M, TILE_SIZE], V: [TILE_SIZE, BLOCK_D]
-        if USE_BF16_DOT:
-            acc += tl.dot(P.to(tl.bfloat16), V.to(tl.bfloat16))
-        else:
-            acc += tl.dot(P.to(tl.float16), V.to(tl.float16))
-
-    # ================================================================
-    # EPILOGUE: Store per-Q-head partial results for stage2
-    # ================================================================
-    # acc: [BLOCK_M, BLOCK_D], L: [BLOCK_M], M: [BLOCK_M]
-    # Stage2 expects per-Q-head: mid_o[B, Hq, split, D+1]
-    safe_L = tl.where(L > 0.0, L, 1.0)
-    one_over_L = 1.0 / safe_L
-    acc = acc * one_over_L[:, None]
-
-    # Convert M from log2 space back to natural log for stage2 compatibility
-    lse = M * LN2 + tl.log(safe_L)
-
-    # Store each Q head's result using 2D scatter
-    # q_head_offs: [BLOCK_M] = kv_hid * KV_GROUP_SIZE + m_offs
-    out_addrs = (
-        bid * stride_mid_b
-        + q_head_offs[:, None] * stride_mid_h
-        + sid * stride_mid_s
-        + d_offs[None, :]
-    )
-    tl.store(
-        Mid_o_ptr + out_addrs,
-        acc,
-        mask=q_mask[:, None] & d_mask[None, :],
-    )
-    # Store LSE per Q head
-    lse_addrs = (
-        bid * stride_mid_b + q_head_offs * stride_mid_h + sid * stride_mid_s + HEAD_DIM
-    )
-    tl.store(Mid_o_ptr + lse_addrs, lse, mask=q_mask)
-
-
 # ---------------------------------------------------------------------------
 # Launcher v2
 # ---------------------------------------------------------------------------
@@ -480,190 +95,3 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
     shows up on a profile, cache by a hash-of-values fingerprint instead.
     """
     return build_pair_lut(centroids)
-
-
-def triton_turboquant_decode_attention_v2(
-    query: torch.Tensor,
-    kv_cache: torch.Tensor,
-    block_table: torch.Tensor,
-    seq_lens: torch.Tensor,
-    Pi: torch.Tensor,
-    centroids: torch.Tensor,
-    scale: float,
-    mse_bits: int,
-    key_packed_size: int,
-    value_quant_bits: int,
-    value_packed_size: int,
-    max_seq_len: int = 0,  # unused; kept for backward compatibility
-    key_fp8: bool = False,
-    norm_correction: bool = False,
-    PiT: torch.Tensor | None = None,
-    # Pre-allocated buffers (optional, required for CUDA-graph stability).
-    mid_o_buf: torch.Tensor | None = None,
-    output_buf: torch.Tensor | None = None,
-    lse_buf: torch.Tensor | None = None,
-    buf_holder: Any = None,
-    # Fixed split count — MUST be a compile-time constant across iterations
-    # for CUDA-graph capture/replay to work. Mirrors the v1 launcher contract.
-    max_num_kv_splits: int = 32,
-) -> torch.Tensor:
-    """Launch optimized TQ decode attention (v2 stage1 + shared stage2).
-
-    Follows the same buffer-reuse + fixed-grid contract as the v1 launcher
-    (triton_turboquant_decode_attention) so the backend can capture a CUDA
-    graph across both versions.
-    """
-    B, Hq, D = query.shape
-    Hk = kv_cache.shape[2]
-    block_size = kv_cache.shape[1]
-    padded_slot = kv_cache.shape[3]
-    max_num_blocks = block_table.shape[1]
-    n_centroids = centroids.shape[0]
-    kv_group_size = Hq // Hk
-    device = query.device
-    del max_seq_len  # no longer used: splits is fixed via max_num_kv_splits
-
-    cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size)
-
-    # Opt#3 SoA layout constants (match store-side computation).
-    key_data_bytes = D if key_fp8 else cfg["mse_bytes"]
-    data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"]
-    meta_region_offset = block_size * Hk * data_bytes_per_slot
-    num_soa_fields = 2 if key_fp8 else 3
-    soa_k_norm = 0
-    soa_v_scale = 0 if key_fp8 else 1
-    soa_v_zero = 1 if key_fp8 else 2
-    kv_cache_u16 = kv_cache.view(torch.uint16)
-
-    # Compute q_rot = q @ Pi.T
-    if key_fp8:
-        q_rot = query.float().contiguous()
-    else:
-        q_float = query.float()
-        if PiT is None:
-            PiT = Pi.T.contiguous()
-        q_rot = (q_float @ PiT).contiguous()
-
-    # BLOCK_M: pad KV_GROUP_SIZE to power of 2, minimum 16 for tensor cores
-    BLOCK_M = max(16, triton.next_power_of_2(kv_group_size))
-
-    # TILE_SIZE (BLOCK_KV): tokens per inner-loop iteration
-    TILE_SIZE = 16
-
-    # Fixed split count — must be constant across calls for cudagraph replay.
-    NUM_KV_SPLITS = max_num_kv_splits
-
-    # --- mid_o buffer reuse (same pattern as v1 launcher) ---
-    if (
-        mid_o_buf is not None
-        and mid_o_buf.shape[0] >= B
-        and mid_o_buf.shape[2] >= NUM_KV_SPLITS
-    ):
-        mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :]
-    else:
-        mid_o = torch.empty(
-            B,
-            Hq,
-            NUM_KV_SPLITS,
-            D + 1,
-            dtype=torch.float32,
-            device=device,
-        )
-        if buf_holder is not None:
-            buf_holder._tq_mid_o_buf = mid_o
-
-    fp8_e4b15 = _use_fp8_e4b15(device.index or 0)
-
-    # Build pair LUT for 4-bit MSE (FLUTE §3.2)
-    use_pair_lut = mse_bits == 4 and not key_fp8
-    pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids
-
-    # Platform-dependent pipelining depth: ROCm prefers num_stages=1.
-    stage1_num_stages = 1 if _is_hip else 2
-
-    # --- Stage 1: v2 kernel ---
-    # Grid over KV heads (not Q heads) — each program handles all Q heads
-    grid = (B, Hk, NUM_KV_SPLITS)
-    _tq_decode_stage1_v2[grid](
-        q_rot,
-        kv_cache,
-        kv_cache_u16,
-        block_table,
-        seq_lens,
-        centroids,
-        pair_lut,
-        mid_o,
-        q_rot.stride(0),
-        q_rot.stride(1),
-        kv_cache.stride(0),
-        block_table.stride(0),
-        mid_o.stride(0),
-        mid_o.stride(1),
-        mid_o.stride(2),
-        NUM_Q_HEADS=Hq,
-        NUM_KV_HEADS=Hk,
-        HEAD_DIM=D,
-        BLOCK_SIZE=block_size,
-        PADDED_SLOT=padded_slot,
-        MAX_NUM_BLOCKS=max_num_blocks,
-        NUM_KV_SPLITS=NUM_KV_SPLITS,
-        KV_GROUP_SIZE=kv_group_size,
-        MSE_BITS=mse_bits,
-        MSE_BYTES=cfg["mse_bytes"],
-        VQB=value_quant_bits,
-        VAL_DATA_BYTES=cfg["val_data_bytes"],
-        KEY_DATA_BYTES=key_data_bytes,
-        META_REGION_OFFSET=meta_region_offset,
-        NUM_SOA_FIELDS=num_soa_fields,
-        SOA_K_NORM=soa_k_norm,
-        SOA_V_SCALE=soa_v_scale,
-        SOA_V_ZERO=soa_v_zero,
-        N_CENTROIDS=n_centroids,
-        ATTN_SCALE=scale,
-        BLOCK_D=cfg["BLOCK_D"],
-        TILE_SIZE=TILE_SIZE,
-        BLOCK_M=BLOCK_M,
-        KEY_FP8=1 if key_fp8 else 0,
-        NORM_CORRECTION=1 if norm_correction else 0,
-        FP8_E4B15=fp8_e4b15,
-        USE_PAIR_LUT=1 if use_pair_lut else 0,
-        USE_BF16_DOT=1 if _is_hip else 0,
-        num_warps=4,
-        num_stages=stage1_num_stages,
-    )
-
-    # --- output / lse buffer reuse (same pattern as v1 launcher) ---
-    if output_buf is not None and output_buf.shape[0] >= B:
-        output = output_buf[:B, :Hq, :D]
-    else:
-        output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)
-        if buf_holder is not None:
-            buf_holder._tq_output_buf = output
-    if lse_buf is not None and lse_buf.shape[0] >= B:
-        lse = lse_buf[:B, :Hq]
-    else:
-        lse = torch.empty(B, Hq, dtype=torch.float32, device=device)
-        if buf_holder is not None:
-            buf_holder._tq_lse_buf = lse
-
-    # --- Stage 2: reduce across KV splits (unchanged) ---
-    grid2 = (B, Hq)
-    _fwd_kernel_stage2[grid2](
-        mid_o,
-        output,
-        lse,
-        seq_lens,
-        mid_o.stride(0),
-        mid_o.stride(1),
-        mid_o.stride(2),
-        output.stride(0),
-        output.stride(1),
-        lse.stride(0),
-        NUM_KV_SPLITS=NUM_KV_SPLITS,
-        BLOCK_DV=cfg["BLOCK_D"],
-        Lv=D,
-        num_warps=4,
-        num_stages=2,
-    )
-
-    return output.to(query.dtype)
diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py
index 03ef36e5b68c..d77592010a4b 100644
--- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py
+++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py
@@ -876,21 +876,6 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
     return build_pair_lut(centroids)
 
 
-# PiT in Q's dtype cached per (PiT-id, dtype) to avoid a host cast on every
-# launcher call. Keyed by the fp32 PiT's data_ptr so ownership/lifetime is
-# tied to the caller's tensor.
-_pit_qdtype_cache: dict = {}
-
-
-def _get_pit_in_query_dtype(PiT: torch.Tensor, qdtype: torch.dtype) -> torch.Tensor:
-    key = (PiT.data_ptr(), qdtype)
-    cached = _pit_qdtype_cache.get(key)
-    if cached is None:
-        cached = PiT.to(qdtype).contiguous()
-        _pit_qdtype_cache[key] = cached
-    return cached
-
-
 def triton_turboquant_unified_attention(
     query: torch.Tensor,  # [num_tokens, Hq, D] - fp16/bf16
     kv_cache: torch.Tensor,  # [num_blocks, block_size, Hk, padded_slot] uint8
diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py
index b7aeb8c6a769..dccc670ce54c 100644
--- a/vllm/v1/attention/ops/vit_attn_wrappers.py
+++ b/vllm/v1/attention/ops/vit_attn_wrappers.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file contains ops for ViT attention to be compatible with torch.compile
+"""This file contains ops for ViT attention to be compatible with torch.compile
 as there are operations here not supported by torch.compile (for instance,
 `.item()` in flash attention)
 
@@ -238,14 +237,21 @@ def apply_sdpa(
     scale: float | None = None,
     enable_gqa: bool = False,
 ) -> torch.Tensor:
-    """
-    Input shape:
+    """Input shape:
     (batch_size x seq_len x num_heads x head_size)
     """
     q, k, v = (einops.rearrange(x, "b s h d -> b h s d") for x in [q, k, v])
-    output = F.scaled_dot_product_attention(
-        q, k, v, dropout_p=0.0, scale=scale, enable_gqa=enable_gqa
-    )
+    if current_platform.is_zen_cpu() and hasattr(torch.ops.zentorch, "zentorch_sdpa"):
+        # Schema: (query, key, value, dropout_p=0., is_causal=False, *,
+        # attn_mask=None, scale=None) -> (attention, logsumexp).
+        # No enable_gqa: GQA is inferred from key.size(1) vs query.size(1).
+        output, _ = torch.ops.zentorch.zentorch_sdpa(
+            q, k, v, dropout_p=0.0, is_causal=False, scale=scale
+        )
+    else:
+        output = F.scaled_dot_product_attention(
+            q, k, v, dropout_p=0.0, scale=scale, enable_gqa=enable_gqa
+        )
     output = einops.rearrange(output, "b h s d -> b s h d ")
     return output
 
diff --git a/vllm/v1/attention/ops/xpu_mla_sparse.py b/vllm/v1/attention/ops/xpu_mla_sparse.py
index 9c1a92e10f57..1e7a9ec3fead 100644
--- a/vllm/v1/attention/ops/xpu_mla_sparse.py
+++ b/vllm/v1/attention/ops/xpu_mla_sparse.py
@@ -185,15 +185,20 @@ def triton_bf16_mla_sparse_interface(
     d_v: int = 512,
     block_dpe: int = 64,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    out : [num_tokens, num_heads_q, d_v]
+    """Out : [num_tokens, num_heads_q, d_v]
     max_logits : [num_tokens, num_heads_q]
     lse : logsumexp, [num_tokens, num_heads_q]
 
     Args:
+        q: Query tensor, [num_tokens, num_heads_q, dim_qk].
+        kv: Latent KV cache tensor.
+        indices: Sparse KV indices selected for each query token.
+        sm_scale: Softmax scale applied to the QK product.
+        d_v: Size of the value dimension.
         block_dpe: Size of positional embedding portion of dim_qk.
             Set to 0 when q/kv contain only the nope latent (e.g. DSv4
             prefill where RoPE is not split out).
+
     """
     num_tokens, num_heads_q, dim_qk = q.shape
     _, num_heads_kv, _ = kv.shape
diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py
index c9631fd45715..a7e8cd0fd22a 100644
--- a/vllm/v1/attention/selector.py
+++ b/vllm/v1/attention/selector.py
@@ -36,6 +36,7 @@ class AttentionSelectorConfig(NamedTuple):
     use_pcp: bool = False
     use_adaptive_verification: bool = False
     use_dcp: bool = False
+    use_rswa: bool = False
 
     def __repr__(self):
         return (
@@ -55,7 +56,8 @@ def __repr__(self):
             f"use_kv_connector={self.use_kv_connector}, "
             f"use_adaptive_verification={self.use_adaptive_verification}, "
             f"use_pcp={self.use_pcp}, "
-            f"use_dcp={self.use_dcp})"
+            f"use_dcp={self.use_dcp}, "
+            f"use_rswa={self.use_rswa})"
         )
 
 
@@ -83,6 +85,7 @@ def get_attn_spec_kind(
 
     Returns:
         The ``KVCacheSpecKind`` the layer maps to.
+
     """
     from vllm.v1.kv_cache_interface import KVCacheSpecKind
 
@@ -113,7 +116,6 @@ def get_attn_backend(
     has_sliding_window: bool = False,
 ) -> type[AttentionBackend]:
     """Selects which attention backend to use and lazily imports it."""
-
     if kv_cache_dtype is not None:
         valid_cache_dtypes = get_args(CacheDType)
         assert kv_cache_dtype in valid_cache_dtypes, (
@@ -168,6 +170,10 @@ def get_attn_backend(
         use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1,
         use_adaptive_verification=use_adaptive_verification,
         use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1,
+        use_rswa=(
+            vllm_config.model_config is not None
+            and vllm_config.model_config.rswa_window is not None
+        ),
     )
 
     # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over
@@ -188,6 +194,9 @@ def get_attn_backend(
         backend=backend,
         attn_selector_config=attn_selector_config,
         num_heads=num_heads,
+        _run_kv_cache_dtype=(
+            cache_config.cache_dtype if cache_config is not None else None
+        ),
     )
 
 
@@ -196,7 +205,11 @@ def _cached_get_attn_backend(
     backend,
     attn_selector_config: AttentionSelectorConfig,
     num_heads: int | None = None,
+    *,
+    _run_kv_cache_dtype: CacheDType | None = None,
 ) -> type[AttentionBackend]:
+    # Some platform selectors inspect the run-wide cache dtype. Keep it in the
+    # cache key even though they read the value from the current config.
     from vllm.platforms import current_platform
 
     attention_cls = current_platform.get_attn_backend_cls(
diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py
index 1efa0e0cea6a..b4c9dcaa2991 100644
--- a/vllm/v1/core/block_pool.py
+++ b/vllm/v1/core/block_pool.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from typing import Any
 
 from vllm.distributed.kv_events import (
@@ -31,8 +31,7 @@
 
 
 class BlockHashToBlockMap:
-    """
-    Cache of blocks that are used for prefix caching. It caches blocks
+    """Cache of blocks that are used for prefix caching. It caches blocks
     from hash directly to a block or multiple blocks
     (i.e. {block_hash: KVCacheBlocks})
     - Mostly block_hash maps to a single KVCacheBlock, and KVCacheBlocks
@@ -59,9 +58,7 @@ def __init__(self):
         ] = {}
 
     def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None:
-        """
-        Gets any block with the given block hash key.
-        """
+        """Gets any block with the given block hash key."""
         blocks = self._cache.get(key)
         if blocks is not None:
             if isinstance(blocks, KVCacheBlock):
@@ -72,9 +69,7 @@ def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None:
         return None
 
     def contain(self, key: BlockHashWithGroupId, block_id: int) -> bool:
-        """
-        Checks whether the key maps to the given block ID.
-        """
+        """Checks whether the key maps to the given block ID."""
         blocks = self._cache.get(key)
         if blocks is None:
             return False
@@ -86,9 +81,7 @@ def contain(self, key: BlockHashWithGroupId, block_id: int) -> bool:
         return False
 
     def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None:
-        """
-        Inserts the KVCacheBlock to the cache
-        """
+        """Inserts the KVCacheBlock to the cache"""
         blocks = self._cache.get(key)
         if blocks is None:
             # When key is not found, attach a single block to the key
@@ -104,9 +97,7 @@ def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None:
             self._unexpected_blocks_type(blocks)
 
     def pop(self, key: BlockHashWithGroupId, block_id: int) -> KVCacheBlock | None:
-        """
-        Checks if block_hash exists and pop block_id from the cache
-        """
+        """Checks if block_hash exists and pop block_id from the cache"""
         blocks = self._cache.pop(key, None)
         if blocks is None:
             # block_hash not found in the cache
@@ -157,6 +148,8 @@ class BlockPool:
             actual block size can be a multiple of hash_block_size.
         enable_kv_cache_events: Whether to enable kv cache events.
         metrics_collector: Optional metrics collector for tracking block residency.
+        medium: Storage medium reported in KV cache events.
+
     """
 
     def __init__(
@@ -166,14 +159,16 @@ def __init__(
         hash_block_size: int,
         enable_kv_cache_events: bool = False,
         metrics_collector: KVCacheMetricsCollector | None = None,
+        medium: str = MEDIUM_GPU,
     ):
         assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
         self.num_gpu_blocks = num_gpu_blocks
+        self.medium = medium
         self.enable_caching = enable_caching
         self.hash_block_size = hash_block_size
         # All kv-cache blocks.
         self.blocks: list[KVCacheBlock] = [
-            KVCacheBlock(idx) for idx in range(num_gpu_blocks)
+            KVCacheBlock(idx, pool=self) for idx in range(num_gpu_blocks)
         ]
         # Free block queue that constructs and manipulates a doubly linked
         # list of free blocks (including eviction candidates when caching is
@@ -194,6 +189,9 @@ def __init__(
         self.kv_event_queue: list[KVCacheEvent] = []
 
         self.metrics_collector = metrics_collector
+        # Callbacks for blocks released with ``unpin_blocks`` whose contents
+        # are still being read until the pool reuses them.
+        self._reuse_watchers: dict[int, Callable[[KVCacheBlock], None]] = {}
 
     def get_cached_block(
         self, block_hash: BlockHash, kv_cache_group_ids: list[int]
@@ -208,6 +206,7 @@ def get_cached_block(
 
         Returns:
             The cached blocks if exists, or None.
+
         """
         cached_blocks = []
         for group_id in kv_cache_group_ids:
@@ -255,6 +254,7 @@ def cache_full_blocks(
                 consults a subset of blocks (e.g. SWA tail-window), so blocks
                 that can never serve a hit stay out of the prefix-cache hash
                 map.
+
         """
         if num_cached_blocks >= num_full_blocks:
             return
@@ -364,7 +364,7 @@ def _build_block_stored_event(
             token_ids=request.all_token_ids[start_token_idx:end_token_idx],
             block_size=block_size,
             lora_id=request.lora_request.adapter_id if request.lora_request else None,
-            medium=MEDIUM_GPU,
+            medium=self.medium,
             lora_name=request.lora_request.name if request.lora_request else None,
             extra_keys=extra_keys_list if extra_keys_list else None,
             group_idx=kv_cache_group_id,
@@ -389,6 +389,7 @@ def emit_cached_block_events(
             num_cached_blocks: Number of blocks that were cache hits.
             block_size: Number of tokens per block.
             kv_cache_group_id: The KV cache group ID.
+
         """
         if not self.enable_kv_cache_events or num_cached_blocks == 0:
             return
@@ -485,6 +486,7 @@ def cache_partial_block(
         Returns:
             The hash key with group ID if a partial entry can be registered;
             otherwise ``None`` for null blocks.
+
         """
         if block.is_null:
             return None
@@ -543,7 +545,7 @@ def cache_partial_block(
                     lora_id=request.lora_request.adapter_id
                     if request.lora_request
                     else None,
-                    medium=MEDIUM_GPU,
+                    medium=self.medium,
                     lora_name=request.lora_request.name
                     if request.lora_request
                     else None,
@@ -610,7 +612,7 @@ def _emit_block_removed_events(
             self.kv_event_queue.append(
                 BlockRemoved(
                     block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))],
-                    medium=MEDIUM_GPU,
+                    medium=self.medium,
                     group_idx=get_group_id(block_hash),
                 )
             )
@@ -665,12 +667,16 @@ def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
 
         Returns:
             A list of new block.
+
         """
         if num_blocks > self.get_num_free_blocks():
             raise ValueError(f"Cannot get {num_blocks} free blocks from the pool")
 
         ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)
 
+        if self._reuse_watchers:
+            self._notify_reuse(ret)
+
         # In order to only iterate the list once, we duplicated code a bit
         if self.enable_caching:
             for block in ret:
@@ -687,9 +693,35 @@ def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
                     self.metrics_collector.on_block_allocated(block)
         return ret
 
-    def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
+    def _notify_reuse(self, blocks: Iterable[KVCacheBlock]) -> None:
+        for block in blocks:
+            watcher = self._reuse_watchers.pop(block.block_id, None)
+            if watcher is not None:
+                watcher(block)
+
+    def unpin_blocks(
+        self,
+        blocks: Iterable[KVCacheBlock],
+        on_reuse: Callable[[KVCacheBlock], None],
+    ) -> None:
+        """Release references to blocks that stay readable until reused.
+
+        The blocks become last-resort eviction candidates regardless of prefix
+        caching: they join the tail of the free queue and count as free.
+        ``on_reuse`` fires when ``get_new_blocks`` hands a block out (or the
+        cache is reset), so the caller can stop reading it.
         """
-        If a block is cached in `cached_block_hash_to_block`, we reset its hash
+        released: list[KVCacheBlock] = []
+        for block in blocks:
+            assert block.ref_cnt > 0 and not block.is_null
+            block.ref_cnt -= 1
+            self._reuse_watchers[block.block_id] = on_reuse
+            if block.ref_cnt == 0:
+                released.append(block)
+        self.free_block_queue.append_n(released)
+
+    def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
+        """If a block is cached in `cached_block_hash_to_block`, we reset its hash
         metadata and evict it from the cache.
 
         Args:
@@ -697,6 +729,7 @@ def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
 
         Returns:
             True if the block is evicted, False otherwise.
+
         """
         # Clean up metrics tracking first to prevent leaks
         if self.metrics_collector:
@@ -717,6 +750,7 @@ def touch(self, blocks: Sequence[KVCacheBlock]) -> None:
 
         Args:
             blocks: A list of blocks to touch.
+
         """
         for block in blocks:
             # ref_cnt=0 means this block is in the free list (i.e. eviction
@@ -738,11 +772,16 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
         Args:
             ordered_blocks: A list of blocks to free ordered by their eviction
                 priority.
+
         """
         # Identify blocks with hash (LRU cache) and without it (never match APC)
         blocks_to_evict_last = []
         blocks_to_evict_first = []
+        other_pools: dict[BlockPool, list[KVCacheBlock]] = {}
         for block in ordered_blocks:
+            if block.pool is not None and block.pool is not self:
+                other_pools.setdefault(block.pool, []).append(block)
+                continue
             block.ref_cnt -= 1
             if block.ref_cnt == 0 and not block.is_null:
                 if block.block_hash is None or not self.enable_caching:
@@ -756,9 +795,11 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
         self.free_block_queue.prepend_n(blocks_to_evict_first)
         # Blocks to reuse last are appended to the end of the free queue.
         self.free_block_queue.append_n(blocks_to_evict_last)
+        for pool, blocks in other_pools.items():
+            pool.free_blocks(blocks)
 
     def evict_blocks(self, block_ids: set[int]) -> None:
-        """evict blocks from the prefix cache by their block IDs.
+        """Evict blocks from the prefix cache by their block IDs.
 
         only evicts blocks that are currently cached (have a hash). blocks
         with ref_cnt > 0 are not freed from the block pool, only evicted
@@ -766,6 +807,7 @@ def evict_blocks(self, block_ids: set[int]) -> None:
 
         Args:
             block_ids: Set of block IDs to evict from cache.
+
         """
         for block_id in block_ids:
             assert block_id < len(self.blocks), (
@@ -784,6 +826,7 @@ def reset_prefix_cache(self) -> bool:
         Returns:
             bool: True if the prefix cache is successfully reset,
             False otherwise.
+
         """
         num_used_blocks = self.num_gpu_blocks - self.get_num_free_blocks()
         if num_used_blocks != 1:  # The null block is always marked as used
@@ -797,6 +840,10 @@ def reset_prefix_cache(self) -> bool:
         # Remove all hashes so that no new blocks will hit.
         self.cached_block_hash_to_block = BlockHashToBlockMap()
         self.cached_block_hashes_by_block.clear()
+        if self._reuse_watchers:
+            self._notify_reuse(
+                [self.blocks[block_id] for block_id in list(self._reuse_watchers)]
+            )
 
         # Remove all hashes from all blocks.
         for block in self.blocks:
@@ -817,6 +864,7 @@ def get_num_free_blocks(self) -> int:
 
         Returns:
             The number of free blocks.
+
         """
         return self.free_block_queue.num_free_blocks
 
@@ -825,8 +873,8 @@ def get_usage(self) -> float:
 
         Returns:
             The KV cache usage (between 0.0 and 1.0).
-        """
 
+        """
         # Subtract 1 to account for null block.
         total_gpu_blocks = self.num_gpu_blocks - 1
         if not total_gpu_blocks:
@@ -838,6 +886,7 @@ def take_events(self) -> list[KVCacheEvent]:
 
         Returns:
             A list of KV cache events.
+
         """
         if not self.enable_kv_cache_events:
             return []
diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py
index a3ed0c7aa15a..24cf609d22b9 100644
--- a/vllm/v1/core/encoder_cache_manager.py
+++ b/vllm/v1/core/encoder_cache_manager.py
@@ -59,11 +59,14 @@ class EncoderCacheManager:
             reference the cached entry. If the set is empty, the entry exists
             but is not referenced by any request and is eligible for
             reclamation.
+        num_embeds_by_hash: Mapping from mm_hash to the number of encoder
+            embeddings stored for that cache entry.
         freeable: List of tuples (mm_hash, num_encoder_embeds) representing entries
             whose no current running request is needed and that can be freed to
             make space when needed.
         freed: List of mm_hash strings that were actually evicted since the
             last call to get_freed_mm_hashes(). This list is cleared on return.
+
     """
 
     @classmethod
@@ -79,6 +82,8 @@ def __init__(self, cache_size: int):
 
         # mm_hash of mm_data => ids of requests that reference the mm_data
         self.cached: dict[str, set[str]] = {}
+        # mm_hash of mm_data => number of cached encoder embeddings
+        self.num_embeds_by_hash: dict[str, int] = {}
         # request_id => set of input_ids cached for that request
         self.request_cached_ids: dict[str, set[int]] = {}
 
@@ -93,6 +98,7 @@ def reset(self) -> None:
         Called when model weights are updated to invalidate stale embeddings.
         """
         self.cached.clear()
+        self.num_embeds_by_hash.clear()
         self.request_cached_ids.clear()
         self.freeable.clear()
         self.freed.clear()
@@ -113,6 +119,7 @@ def check_and_update_cache(self, request: Request, input_id: int) -> bool:
 
         Returns:
             True if the encoder output for this input is already cached
+
         """
         mm_hash = request.mm_features[input_id].identifier
         # Not cached at all
@@ -128,6 +135,13 @@ def check_and_update_cache(self, request: Request, input_id: int) -> bool:
         self.request_cached_ids.setdefault(request.request_id, set()).add(input_id)
         return True
 
+    def get_cached_num_encoder_embeds(
+        self, request: Request, input_id: int
+    ) -> int | None:
+        """Return the encoder-embedding count stored for a cached input."""
+        mm_hash = request.mm_features[input_id].identifier
+        return self.num_embeds_by_hash.get(mm_hash)
+
     def can_allocate(
         self,
         request: Request,
@@ -162,6 +176,7 @@ def can_allocate(
 
         Note: This method does not allocate physical memory for the encoder
         output but only the state of EncoderCacheManager.
+
         """
         num_embeds = request.get_num_encoder_embeds(input_id)
 
@@ -185,6 +200,7 @@ def can_allocate(
         while num_embeds > self.num_free_slots:
             mm_hash, num_free_embeds = self.freeable.popitem(last=False)
             del self.cached[mm_hash]
+            del self.num_embeds_by_hash[mm_hash]
             self.freed.append(mm_hash)
             self.num_free_slots += num_free_embeds
         return True
@@ -198,14 +214,20 @@ def allocate(self, request: Request, input_id: int) -> None:
 
         Note:
             This method assumes can_allocate() returned True for the same input.
-        """
 
+        """
         mm_hash = request.mm_features[input_id].identifier
         request_id = request.request_id
         if mm_hash not in self.cached:
             self.cached[mm_hash] = set()
 
         num_encoder_embeds = request.get_num_encoder_embeds(input_id)
+        cached_num_encoder_embeds = self.num_embeds_by_hash.setdefault(
+            mm_hash, num_encoder_embeds
+        )
+        # _try_schedule_encoder_inputs rejects local-hit and same-step
+        # identifier/count mismatches before allocation.
+        assert cached_num_encoder_embeds == num_encoder_embeds
 
         # NOTE: Encoder cache should always have enough space for encoder inputs
         # that are scheduled since eviction takes place at can_allocate().
@@ -222,7 +244,7 @@ def get_cached_input_ids(self, request: Request) -> set[int]:
         return self.request_cached_ids.get(request.request_id, set())
 
     def free_encoder_input(self, request: Request, input_id: int) -> None:
-        """Free the request's reference to the encoder input (`mm_data`)
+        """Free the request's reference to the encoder input (`mm_data`).
 
         When the reference set for the corresponding `mm_hash` becomes empty,
         the entry is appended to `freeable` and `num_freeable_slots` is
@@ -255,7 +277,7 @@ def free_encoder_input(self, request: Request, input_id: int) -> None:
             return
         self.cached[mm_hash].discard(req_id)
         if not self.cached[mm_hash]:
-            num_encoder_embeds = request.get_num_encoder_embeds(input_id)
+            num_encoder_embeds = self.num_embeds_by_hash[mm_hash]
             self.freeable[mm_hash] = num_encoder_embeds
             self.num_freeable_slots += num_encoder_embeds
 
@@ -279,6 +301,7 @@ def get_freed_mm_hashes(self) -> list[str]:
             call to be used by the scheduler to notify workers about which
             encoder outputs can be removed from their caches. The internal
             list is cleared after this call.
+
         """
         # An entry evicted early in the scheduling pass can be allocated again
         # later in the same pass. Keep its worker-side tensor in that case.
@@ -307,8 +330,8 @@ def compute_mm_encoder_budget(
             from the input sequence.
         - Space budget for encoder cache size, measured in number of tokens
             from the input sequence.
-    """
 
+    """
     if not mm_max_toks_per_item:
         logger.warning(
             "All non-text modalities supported by the model have been "
diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py
index a495f997bf3b..8061999409c4 100644
--- a/vllm/v1/core/kv_cache_coordinator.py
+++ b/vllm/v1/core/kv_cache_coordinator.py
@@ -64,9 +64,7 @@ def _validate_prefix_cache_retention_interval(
 
 
 class KVCacheCoordinator(ABC):
-    """
-    Coordinate the KV cache of different KV cache groups.
-    """
+    """Coordinate the KV cache of different KV cache groups."""
 
     enable_partial_hash_hits = False
 
@@ -139,6 +137,7 @@ def __init__(
                 max_in_flight_tokens=max_in_flight_tokens,
                 max_model_len=max_model_len,
                 block_pool=self.block_pool,
+                role=kv_cache_group.role,
                 enable_caching=enable_caching,
                 kv_cache_group_id=i,
                 dcp_world_size=dcp_world_size_for_kv_cache_spec(
@@ -155,6 +154,9 @@ def __init__(
             for manager in self.single_type_managers:
                 if isinstance(manager, MambaManager):
                     manager.drop_eagle_checkpoint_block = True
+        self.group_block_sizes = tuple(
+            manager.block_size for manager in self.single_type_managers
+        )
 
         # A positive retention interval must be a multiple of the base hit granularity
         # (``scheduler_block_size``) to land on real cache-hit boundaries.
@@ -175,8 +177,7 @@ def get_num_blocks_to_allocate(
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
     ) -> int:
-        """
-        Get the number of blocks needed to be allocated for the request.
+        """Get the number of device blocks needed to be allocated for the request.
 
         Args:
             request_id: The request ID.
@@ -199,6 +200,7 @@ def get_num_blocks_to_allocate(
 
         Returns:
             The number of blocks to allocate.
+
         """
         num_blocks_to_allocate = 0
         for i, manager in enumerate(self.single_type_managers):
@@ -233,8 +235,7 @@ def allocate_new_computed_blocks(
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
     ) -> None:
-        """
-        Add the new computed blocks to the request. Optionally allocate new
+        """Add the new computed blocks to the request. Optionally allocate new
             blocks for external computed tokens (if any).
 
         Args:
@@ -243,6 +244,7 @@ def allocate_new_computed_blocks(
                 prefix cache.
             num_local_computed_tokens: The number of local computed tokens.
             num_external_computed_tokens: The number of external computed tokens.
+
         """
         # A running request is already tracked in num_cached_block and won't
         # have new prefix-cache hits, so this is a no-op for it.
@@ -279,8 +281,7 @@ def allocate_new_blocks(
         num_tokens_main_model: int,
         num_encoder_tokens: int = 0,
     ) -> tuple[list[KVCacheBlock], ...]:
-        """
-        Allocate new blocks for the request to give it at least `num_tokens`
+        """Allocate new blocks for the request to give it at least `num_tokens`
         token slots.
 
         Args:
@@ -295,6 +296,7 @@ def allocate_new_blocks(
 
         Returns:
             The new allocated blocks.
+
         """
         return tuple(
             manager.allocate_new_blocks(
@@ -329,17 +331,19 @@ def get_replay_boundaries(self, request: Request) -> tuple[int, ...]:
         return tuple(sorted({max(resend - block, 0), max(extension - block, 0)}))
 
     def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
-        """
-        Cache the blocks for the request.
+        """Cache the blocks for the request.
 
         Args:
             request: The request.
             num_computed_tokens: The total number of tokens
                 that need to be cached
                 (including tokens that are already cached).
+
         """
         boundaries = self.get_replay_boundaries(request)
         for manager in self.single_type_managers:
+            if not manager.enable_caching:
+                continue
             # Only cache tokens with finalized KV. The last num_reprefillable_tokens
             # tokens can be re-prefilled during multi-module MTP.
             num_tokens_to_cache = max(
@@ -352,19 +356,39 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
                 replay_boundaries=boundaries,
             )
 
+    def emit_cached_block_events(
+        self, request: Request, computed_blocks: tuple[list[KVCacheBlock], ...]
+    ) -> None:
+        for group_idx, group_blocks in enumerate(computed_blocks):
+            if group_blocks:
+                manager = self.single_type_managers[group_idx]
+                group = self.kv_cache_config.kv_cache_groups[group_idx]
+                manager.block_pool.emit_cached_block_events(
+                    request,
+                    len(group_blocks),
+                    group.kv_cache_spec.block_size,
+                    group_idx,
+                )
+
+    def reset_prefix_cache(self) -> bool:
+        """Reset each manager's pool once."""
+        pools = dict.fromkeys(
+            manager.block_pool for manager in self.single_type_managers
+        )
+        return all([pool.reset_prefix_cache() for pool in pools])
+
     def free(self, request_id: str) -> None:
-        """
-        Free the blocks for the request.
+        """Free the blocks for the request.
 
         Args:
             request_id: The request ID.
+
         """
         for manager in self.single_type_managers:
             manager.free(request_id)
 
     def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
-        """
-        Pop the request's bookkeeping from all single-type managers and
+        """Pop the request's bookkeeping from all single-type managers and
         return its blocks without returning them to the block pool. The
         caller must eventually pass the returned blocks to
         `block_pool.free_blocks`, freeing them in reverse order (so that
@@ -375,6 +399,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
 
         Returns:
             The request's blocks in allocation order.
+
         """
         blocks: list[KVCacheBlock] = []
         for manager in self.single_type_managers:
@@ -382,8 +407,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
         return blocks
 
     def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
-        """
-        Get the number of common prefix blocks for all requests with allocated
+        """Get the number of common prefix blocks for all requests with allocated
         KV cache for each kv cache group.
 
         Args:
@@ -392,6 +416,7 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
 
         Returns:
             list[int]: The number of common prefix blocks for each kv cache group.
+
         """
         return [
             manager.get_num_common_prefix_blocks(running_request_id)
@@ -404,8 +429,7 @@ def remove_skipped_blocks(
         processed_computed_tokens: int,
         num_prompt_tokens: int | None = None,
     ) -> None:
-        """
-        Remove the blocks that are no longer needed from `blocks` and replace
+        """Remove the blocks that are no longer needed from `blocks` and replace
         the removed blocks with null_block.
 
         Args:
@@ -415,6 +439,7 @@ def remove_skipped_blocks(
             num_prompt_tokens: Optional prompt length. R-SWA managers use this to
                 free gap blocks between the prefill tail and decode window; other
                 manager types ignore it.
+
         """
         for manager in self.single_type_managers:
             manager.remove_skipped_blocks(
@@ -422,9 +447,7 @@ def remove_skipped_blocks(
             )
 
     def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]:
-        """
-        Get the blocks for the request.
-        """
+        """Get the blocks for the request."""
         return tuple(
             manager.req_to_blocks.get(request_id) or []
             for manager in self.single_type_managers
@@ -448,8 +471,7 @@ def new_step_starts(self) -> None:
 
 
 class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
-    """
-    KV cache coordinator to use if prefix caching is disabled or unsupported.
+    """KV cache coordinator to use if prefix caching is disabled or unsupported.
     In contrast to UnitaryKVCacheCoordinator and HybridKVCacheCoordinator,
     supports arbitrary numbers of KV cache groups (including 0 groups).
     Does not implement any features related to prefix caching.
@@ -500,8 +522,7 @@ def find_longest_cache_hit(
 
 
 class UnitaryKVCacheCoordinator(KVCacheCoordinator):
-    """
-    KV cache coordinator for models with only one KV cache group. This is the
+    """KV cache coordinator for models with only one KV cache group. This is the
     case for models with only one KV cache type, e.g., all attention layers use
     full attention or all attention layers use sliding window attention.
     """
@@ -586,8 +607,7 @@ class SpecGroup(NamedTuple):
 
 
 class HybridKVCacheCoordinator(KVCacheCoordinator):
-    """
-    KV cache coordinator for hybrid models with multiple KV cache types, and
+    """KV cache coordinator for hybrid models with multiple KV cache types, and
     thus multiple kv cache groups.
     """
 
@@ -605,6 +625,7 @@ def __init__(
         hash_block_size: int,
         metrics_collector: KVCacheMetricsCollector | None = None,
         num_prefill_lookahead: int = 0,
+        allow_partial_hash_hits: bool = True,
     ):
         super().__init__(
             kv_cache_config,
@@ -629,7 +650,7 @@ def __init__(
         # Only groups that participate in prefix caching must satisfy the
         # divisibility constraint; groups that opt out (e.g. GLM-5.3-Flash kpool
         # tail, block_size=kpool) are scratch buffers and excluded.
-        group_block_sizes = [
+        cacheable_block_sizes = [
             manager.block_size
             for manager, group in zip(
                 self.single_type_managers, kv_cache_config.kv_cache_groups
@@ -637,10 +658,10 @@ def __init__(
             if group.kv_cache_spec.prefix_cacheable
         ]
         assert all(
-            block_size % hash_block_size == 0 for block_size in group_block_sizes
+            block_size % hash_block_size == 0 for block_size in cacheable_block_sizes
         ), (
             "Each KV cache group's real block_size must be divisible by "
-            f"hash_block_size. block_sizes={group_block_sizes}, "
+            f"hash_block_size. block_sizes={cacheable_block_sizes}, "
             f"hash_block_size={hash_block_size}"
         )
         assert pcp_world_size == 1, "PCP not support hybrid attn now."
@@ -669,7 +690,9 @@ def __init__(
             )
             for g in kv_cache_config.kv_cache_groups
         )
-        self.enable_partial_hash_hits = has_partial_mamba_group
+        self.enable_partial_hash_hits = (
+            allow_partial_hash_hits and has_partial_mamba_group
+        )
         if self.enable_partial_hash_hits:
             unsupported_partial_hit_managers = {
                 type(manager).__name__
@@ -703,8 +726,7 @@ def _cache_hit_alignment_tokens(self) -> int:
         )
 
     def verify_and_split_kv_cache_groups(self) -> None:
-        """
-        Groups KV cache groups by their spec type for efficient batch processing
+        """Groups KV cache groups by their spec type for efficient batch processing
         during cache hit lookup.
         """
         self.attention_groups: list[SpecGroup] = []
@@ -722,10 +744,12 @@ def verify_and_split_kv_cache_groups(self) -> None:
 
             # Try to find an existing group with the same spec
             for idx, group in enumerate(self.attention_groups):
-                if group.spec == spec:
-                    assert manager_cls is group.manager_cls, (
-                        "Expected same manager class for identical KV cache specs."
-                    )
+                if (
+                    group.spec == spec
+                    and group.manager_cls is manager_cls
+                    and self.single_type_managers[group.group_ids[0]].block_pool
+                    is self.single_type_managers[i].block_pool
+                ):
                     group.group_ids.append(i)
                     if use_eagle and not group.use_eagle:
                         self.attention_groups[idx] = group._replace(use_eagle=True)
@@ -776,6 +800,8 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
         cached_num_computed_tokens = self._align_cacheable(num_computed_tokens)
         boundaries = self.get_replay_boundaries(request)
         for manager in self.single_type_managers:
+            if not manager.enable_caching:
+                continue
             num_tokens_to_cache = cached_num_computed_tokens
             # EAGLE groups match one block past each aligned boundary and drop
             # it, so make that lookahead block eligible to be cached.
@@ -809,8 +835,7 @@ def find_longest_cache_hit(
         block_hashes: list[BlockHash],
         max_cache_hit_length: int,
     ) -> tuple[tuple[list[KVCacheBlock], ...], int, int]:
-        """
-        Find the longest cache hit using an iterative fixed-point algorithm.
+        """Find the longest cache hit using an iterative fixed-point algorithm.
 
         Each attention type either accepts the current candidate length or
         reduces it. If any type reduces the length, restart checks over all
@@ -827,8 +852,8 @@ def find_longest_cache_hit(
                 - The number of tokens of the reconciled (combined) cache hit.
                 - ``num_uncached_common_prefix_tokens``: a shared prefix that a
                   sparse-retention group has not cached yet (0 unless hybrid).
-        """
 
+        """
         num_groups = len(self.kv_cache_config.kv_cache_groups)
         hit_length = max_cache_hit_length
         longest_hit_length = 0
@@ -889,7 +914,7 @@ def find_longest_cache_hit(
                     block_hashes=block_hashes,
                     max_length=_max_length,
                     kv_cache_group_ids=group_ids,
-                    block_pool=self.block_pool,
+                    block_pool=self.single_type_managers[first_group_id].block_pool,
                     kv_cache_spec=spec,
                     drop_eagle_block=drop_eagle_block,
                     alignment_tokens=self._cache_hit_alignment_tokens,
@@ -923,7 +948,10 @@ def find_longest_cache_hit(
         for group in self.attention_groups:
             if not isinstance(group.spec, FullAttentionSpec):
                 continue
-            group_block_size = self.single_type_managers[group.group_ids[0]].block_size
+            manager = self.single_type_managers[group.group_ids[0]]
+            if manager.retains_longer_hit:
+                continue
+            group_block_size = manager.block_size
             num_blocks = cdiv(hit_length, group_block_size)
             for group_id in group.group_ids:
                 if (blks := hit_blocks_by_group[group_id]) is not None:
@@ -948,8 +976,8 @@ def find_longest_cache_hit_per_group(
 
         Returns:
             (blocks_per_group, hit_lengths_per_group)
-        """
 
+        """
         num_groups = len(self.kv_cache_config.kv_cache_groups)
         hit_blocks: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)]
         hit_lengths: list[int] = [0] * num_groups
@@ -960,7 +988,7 @@ def find_longest_cache_hit_per_group(
                 block_hashes=block_hashes,
                 max_length=max_cache_hit_length,
                 kv_cache_group_ids=group_ids,
-                block_pool=self.block_pool,
+                block_pool=manager.block_pool,
                 kv_cache_spec=spec,
                 drop_eagle_block=use_eagle,
                 alignment_tokens=self._cache_hit_alignment_tokens,
@@ -987,6 +1015,7 @@ def get_kv_cache_coordinator(
     hash_block_size: int,
     metrics_collector: KVCacheMetricsCollector | None = None,
     num_prefill_lookahead: int = 0,
+    allow_partial_hash_hits: bool = True,
 ) -> KVCacheCoordinator:
     if not enable_caching:
         return KVCacheCoordinatorNoPrefixCache(
@@ -1030,4 +1059,5 @@ def get_kv_cache_coordinator(
         hash_block_size=hash_block_size,
         metrics_collector=metrics_collector,
         num_prefill_lookahead=num_prefill_lookahead,
+        allow_partial_hash_hits=allow_partial_hash_hits,
     )
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index 084ba24311c0..30bf584f8c8e 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -33,8 +33,7 @@
 
 @dataclass
 class KVCacheBlocks:
-    """
-    The allocation result of KVCacheManager, work as the interface between
+    """The allocation result of KVCacheManager, work as the interface between
     Scheduler and KVCacheManager, to hide KVCacheManager's internal data
     structure from the Scheduler.
     """
@@ -67,30 +66,45 @@ def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks":
     def get_block_ids(
         self,
         allow_none: Literal[False] = False,
+        *,
+        group_ids: Sequence[int] | None = None,
     ) -> tuple[list[int], ...]: ...
 
     @overload
     def get_block_ids(
         self,
         allow_none: Literal[True] = True,
+        *,
+        group_ids: Sequence[int] | None = None,
     ) -> tuple[list[int], ...] | None: ...
 
     def get_block_ids(
         self,
         allow_none: bool = False,
+        *,
+        group_ids: Sequence[int] | None = None,
     ) -> tuple[list[int], ...] | None:
-        """
-        Converts the KVCacheBlocks instance to block_ids.
+        """Converts the KVCacheBlocks instance to block_ids.
+
+        Args:
+            allow_none: Return None when every selected group is empty.
+            group_ids: KV cache groups to include. Includes all groups by default.
 
         Returns:
             tuple[list[int], ...]: A tuple of lists where:
                 - the outer tuple corresponds to KV cache groups
                 - each inner list contains the block_ids of the blocks in that
                   group
+
         """
-        if allow_none and all(len(group) == 0 for group in self.blocks):
+        groups = (
+            self.blocks
+            if group_ids is None
+            else tuple(self.blocks[group_id] for group_id in group_ids)
+        )
+        if allow_none and all(len(group) == 0 for group in groups):
             return None
-        return tuple([blk.block_id for blk in group] for group in self.blocks)
+        return tuple([blk.block_id for blk in group] for group in groups)
 
     def get_unhashed_block_ids(self) -> list[int]:
         """Get block_ids of unhashed blocks from KVCacheBlocks instance."""
@@ -110,9 +124,7 @@ def get_unhashed_block_ids_all_groups(self) -> list[list[int]]:
         ]
 
     def new_empty(self) -> "KVCacheBlocks":
-        """
-        Creates a new KVCacheBlocks instance with no blocks.
-        """
+        """Creates a new KVCacheBlocks instance with no blocks."""
         return KVCacheBlocks(tuple(() for _ in range(len(self.blocks))))
 
 
@@ -133,7 +145,7 @@ def __init__(
         pcp_world_size: int = 1,
         metrics_collector: KVCacheMetricsCollector | None = None,
         watermark: float = 0.0,
-        enable_mamba_fine_grained_prefix_cache: bool = False,
+        enable_mamba_shared_prefix_checkpoint: bool = False,
     ) -> None:
         self.max_model_len = max_model_len
         # When unset, fall back to `max_model_len` so the recycling-aware cap
@@ -171,18 +183,23 @@ def __init__(
         # junction costs a forward pass and displaces the block-boundary stop.
         # Multi-module MTP is excluded because ``cache_blocks`` then hands the
         # manager ``num_computed - num_reprefillable`` rather than the chunk end.
-        self.mamba_fine_grained_prefix_cache = (
-            enable_mamba_fine_grained_prefix_cache
+        self.mamba_shared_prefix_checkpoint = (
+            enable_mamba_shared_prefix_checkpoint
             and bool(self.coordinator.eagle_group_ids)
             and self.coordinator.enable_partial_hash_hits
             and self.coordinator.num_reprefillable_tokens == 0
         )
-        if self.mamba_fine_grained_prefix_cache:
+        if self.mamba_shared_prefix_checkpoint:
             for manager in self.coordinator.single_type_managers:
                 if isinstance(manager, MambaManager):
-                    manager.fine_grained_prefix_cache = True
+                    manager.shared_prefix_checkpoint = True
         self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
         self.block_pool = self.coordinator.block_pool
+        self.retained_hit_group_ids = tuple(
+            manager.kv_cache_group_id
+            for manager in self.coordinator.single_type_managers
+            if manager.retains_longer_hit
+        )
         self.kv_cache_config = kv_cache_config
 
         # Watermark: minimum number of KV cache blocks to keep free when
@@ -212,6 +229,7 @@ def usage(self) -> float:
 
         Returns:
             The KV cache usage (between 0.0 and 1.0).
+
         """
         return self.block_pool.get_usage()
 
@@ -220,6 +238,7 @@ def make_prefix_cache_stats(self) -> PrefixCacheStats | None:
 
         Returns:
             The current prefix caching stats, or None if logging is disabled.
+
         """
         if not self.log_stats:
             return None
@@ -258,6 +277,7 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int
                   window) has not cached yet (Marconi-style APC), or 0 if none.
                   Pinned so sparse prefix-cache retention does not drop
                   the junction and defeat cross-request reuse.
+
         """
         # We skip finding the prefix cache hit when prefix caching is
         # disabled or the request is marked as skipping kv cache read
@@ -287,17 +307,7 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int
             and self.enable_kv_cache_events
             and getattr(request, "kv_cache_report_mode", "incremental") == "full"
         ):
-            for group_idx, group_blocks in enumerate(computed_blocks):
-                num_blocks = len(group_blocks)
-                if num_blocks > 0:
-                    group = self.kv_cache_config.kv_cache_groups[group_idx]
-                    block_size = group.kv_cache_spec.block_size
-                    self.block_pool.emit_cached_block_events(
-                        request,
-                        num_blocks,
-                        block_size,
-                        group_idx,
-                    )
+            self.coordinator.emit_cached_block_events(request, computed_blocks)
 
         # The junction to pin is where the lagging sparse-retention group stops
         # (``num_new_computed_tokens``) plus the uncached shared prefix -- i.e.
@@ -330,6 +340,7 @@ def get_computed_blocks_for_connector(
         Returns:
             The ``get_computed_blocks`` triple (blocks, number of local computed
             tokens, shared-prefix boundary) plus ``hit_diverged``.
+
         """
         coordinator = self.coordinator
         if not (
@@ -426,7 +437,7 @@ def allocate_slots(
         ----------------------------------------------------------------------
         ```
 
-        Abbrivations:
+        Abbreviations:
 
         ```
         comp      = request.num_computed_tokens
@@ -452,17 +463,33 @@ def allocate_slots(
 
         Returns:
             A list of new allocated blocks.
+
         """
-        # When loading KV data asynchronously, we may have zero new tokens to
-        # compute while still allocating slots for externally computed tokens.
-        if num_new_tokens == 0 and num_external_computed_tokens == 0:
+        # A step may need no slots of its own while still adopting computed
+        # tokens: an async KV load, or a chunk that ends inside the replayed
+        # range of a hit (SWA bounded replay).
+        if (
+            num_new_tokens == 0
+            and num_external_computed_tokens == 0
+            and num_new_computed_tokens == 0
+        ):
             raise ValueError(
                 "num_new_tokens must be greater than 0 when there are no "
-                "external computed tokens"
+                "computed tokens to adopt"
             )
 
         if new_computed_blocks is not None:
             new_computed_block_list = new_computed_blocks.blocks
+            if self.retained_hit_group_ids:
+                # Only retain the prefix covered by the local and external hits.
+                reused = num_new_computed_tokens + num_external_computed_tokens
+                groups = list(new_computed_block_list)
+                for group_id in self.retained_hit_group_ids:
+                    manager = self.coordinator.single_type_managers[group_id]
+                    groups[group_id] = groups[group_id][
+                        : cdiv(reused, manager.block_size)
+                    ]
+                new_computed_block_list = tuple(groups)
         else:
             new_computed_block_list = self.empty_kv_cache_blocks.blocks
 
@@ -564,7 +591,7 @@ def allocate_slots(
 
         # P/D: delay caching blocks if we have to recv from
         # remote. Update state for locally cached blocks.
-        if not self.enable_caching or delay_cache_blocks:
+        if delay_cache_blocks:
             return self.create_kv_cache_blocks(new_blocks)
 
         # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens
@@ -587,6 +614,7 @@ def free(self, request: Request) -> None:
 
         Args:
             request: The request to free the blocks.
+
         """
         self.coordinator.free(request.request_id)
 
@@ -604,6 +632,7 @@ def remove_skipped_blocks(
             processed_computed_tokens: Computed-token prefix length covering
                 fully processed and committed tokens only (safe to free).
             num_prompt_tokens: Optional prompt length for R-SWA gap eviction.
+
         """
         self.coordinator.remove_skipped_blocks(
             request_id, processed_computed_tokens, num_prompt_tokens
@@ -619,14 +648,16 @@ def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]:
 
         Returns:
             The request's blocks in allocation order.
+
         """
         return self.coordinator.pop_blocks_for_free(request.request_id)
 
     def evict_blocks(self, block_ids: set[int]) -> None:
-        """evict blocks from the prefix cache by their block IDs.
+        """Evict blocks from the prefix cache by their block IDs.
 
         Args:
             block_ids: Set of block IDs to evict from cache.
+
         """
         self.block_pool.evict_blocks(block_ids)
 
@@ -638,8 +669,9 @@ def reset_prefix_cache(self) -> bool:
         Returns:
             bool: True if the prefix cache is successfully reset,
             False otherwise.
+
         """
-        if not self.block_pool.reset_prefix_cache():
+        if not self.coordinator.reset_prefix_cache():
             return False
         if self.log_stats:
             assert self.prefix_cache_stats is not None
@@ -677,6 +709,7 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
         Returns:
             list[int]: The number of common prefix blocks for each kv cache
             group.
+
         """
         return self.coordinator.get_num_common_prefix_blocks(running_request_id)
 
@@ -685,6 +718,7 @@ def take_events(self) -> list[KVCacheEvent]:
 
         Returns:
             A list of KV cache events.
+
         """
         events = self.block_pool.take_events()
         for event in events:
@@ -741,6 +775,8 @@ def estimate_cached_tokens(self, request: Request) -> int:
             self.kv_cache_config.kv_cache_groups,
             self.get_blocks(request.request_id).blocks,
         ):
+            if not group.kv_cache_spec.prefix_cacheable:
+                continue
             if isinstance(
                 group.kv_cache_spec,
                 (CrossAttentionSpec, EncoderOnlyAttentionSpec),
@@ -770,6 +806,7 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
             request: The request to cache the blocks.
             num_computed_tokens: The number of computed tokens, including tokens
                 that are already cached and tokens to be cached.
+
         """
         if self.enable_caching:
             self.coordinator.cache_blocks(request, num_computed_tokens)
diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py
index 38bae039095d..af1112c54aaf 100644
--- a/vllm/v1/core/kv_cache_utils.py
+++ b/vllm/v1/core/kv_cache_utils.py
@@ -8,22 +8,30 @@
 import os
 from collections import defaultdict
 from collections.abc import Callable, Iterable, Iterator, Sequence
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 from functools import partial
-from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload
+from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload
 
 from vllm import envs
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.utils.hashing import xxhash, xxhash_cbor
-from vllm.utils.math_utils import cdiv
+from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.mem_utils import format_gib
 from vllm.utils.torch_utils import get_dtype_size
+from vllm.v1.hisparse.layout import (
+    get_hisparse_gpu_memory_usage,
+    get_hisparse_host_pool_bytes,
+    get_hisparse_kv_cache_config,
+    get_hisparse_kv_cache_groups,
+)
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     ChunkedLocalAttentionSpec,
+    CircularBufferSpec,
     FullAttentionSpec,
     HiddenStateCacheSpec,
+    HiSparseHotSpec,
     KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -43,6 +51,10 @@
 from vllm.v1.request import Request
 from vllm.v1.utils import tensor_data
 
+if TYPE_CHECKING:
+    from vllm.v1.core.block_pool import BlockPool
+
+
 # BlockHash represents the hash of a single KV-cache block used for
 # prefix caching.  Treating it as a distinct type from `bytes` helps
 # catch accidental misuse when passing around raw byte strings.
@@ -168,6 +180,10 @@ class KVCacheBlock:
     block_id: int
     # Reference count.
     ref_cnt: int = 0
+    # Deferred frees may contain blocks from multiple pools.
+    pool: "BlockPool | None" = field(
+        default=None, repr=False, compare=False, kw_only=True
+    )
     # The hash key (block hash + group id) of the block, only available
     # when the block is full and cached.
     _block_hash: BlockHashWithGroupId | None = None
@@ -247,6 +263,7 @@ class FreeKVCacheBlockQueue:
 
     Args:
         blocks: A list of KVCacheBlock objects.
+
     """
 
     def __init__(self, blocks: list[KVCacheBlock]) -> None:
@@ -284,6 +301,7 @@ def popleft(self) -> KVCacheBlock:
 
         Returns:
             The first free block.
+
         """
         if (
             self.fake_free_list_head.next_free_block is self.fake_free_list_tail
@@ -324,6 +342,7 @@ def popleft_n(self, n: int) -> list[KVCacheBlock]:
 
         Returns:
             A list of n free blocks.
+
         """
         if n == 0:
             return []
@@ -354,6 +373,7 @@ def remove(self, block: KVCacheBlock) -> None:
 
         Args:
             block: The block to remove.
+
         """
         if block.prev_free_block is None or block.next_free_block is None:
             # This should not happen if the block is from the free list.
@@ -375,6 +395,7 @@ def append(self, block: KVCacheBlock) -> None:
 
         Args:
             block: The block to append.
+
         """
         if self.fake_free_list_tail.prev_free_block is None:
             raise RuntimeError(
@@ -418,6 +439,7 @@ def append_n(self, blocks: list[KVCacheBlock]) -> None:
 
         Args:
             blocks: The blocks to append.
+
         """
         if len(blocks) == 0:
             return
@@ -443,6 +465,7 @@ def get_all_free_blocks(self) -> list[KVCacheBlock]:
 
         Returns:
             A list of free blocks.
+
         """
         ret = []
         if self.fake_free_list_head.next_free_block is None:
@@ -489,6 +512,7 @@ def _gen_mm_extra_hash_keys(
 
     Returns:
         A tuple of extra keys and the next multi-modal index.
+
     """
     extra_keys: list[Any] = []
 
@@ -549,6 +573,7 @@ def _gen_lora_extra_hash_keys(request: Request) -> list[str]:
     Returns:
         Return LoRA name of the request if it is a LoRA request. Return empty
         list otherwise.
+
     """
     if not request.lora_request:
         return []
@@ -568,6 +593,7 @@ def _gen_prompt_embeds_extra_hash_keys(
     Returns:
         Return a stable hash of the block prompt embeddings if prompt embeds
         are present. Return empty list otherwise.
+
     """
     if request.prompt_embeds is None:
         return []
@@ -596,6 +622,7 @@ def generate_block_hash_extra_keys(
 
     Returns:
         A tuple of extra keys and the next multi-modal index.
+
     """
     mm_extra_keys: list[Any]
     mm_extra_keys, new_start_mm_idx = _gen_mm_extra_hash_keys(
@@ -629,6 +656,7 @@ def hash_block_tokens(
     the contents of the preceding block(s). The hash value is used for
     prefix caching. We use LRU cache for this function to avoid recomputing
     hash values for the same block contents.
+
     Args:
         hash_function: The hash function used to compute block hash.
         parent_block_hash: The hash of the parent block. None
@@ -636,9 +664,11 @@ def hash_block_tokens(
         curr_block_token_ids: A list of token ids in the current
             block. The current block is assumed to be full.
         extra_keys: Extra keys for the block.
+
     Returns:
         The hash value of the block and the token ids in the block.
         The entire tuple is used as the hash key of the block.
+
     """
     if not parent_block_hash:
         parent_block_hash = NONE_HASH
@@ -797,8 +827,7 @@ def get_request_block_hasher(
     hash_block_size: int,
     caching_hash_fn: Callable[[Any], bytes],
 ) -> Callable[[Request], list[BlockHash]]:
-    """
-    Returns a function which computes the list of un-computed block hashes
+    """Returns a function which computes the list of un-computed block hashes
     of a request.
 
     Hashes are computed at ``hash_block_size`` granularity and chained over the
@@ -896,9 +925,7 @@ def _check_enough_kv_cache_memory(
 def max_memory_usage_bytes(
     vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec]
 ) -> int:
-    """
-    Get the maximum memory usage in bytes for the given KV cache specs.
-    """
+    """Get the maximum memory usage in bytes for the given KV cache specs."""
     return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs)
 
 
@@ -907,8 +934,7 @@ def estimate_max_model_len(
     kv_cache_spec: dict[str, KVCacheSpec],
     available_memory: int,
 ) -> int:
-    """
-    Estimates the maximum model length that can fit in the available memory
+    """Estimates the maximum model length that can fit in the available memory
     using binary search.
 
     This function temporarily modifies max_model_len during estimation but
@@ -921,6 +947,7 @@ def estimate_max_model_len(
 
     Returns:
         The estimated maximum model length that can fit in the available memory.
+
     """
     # Save the original max_model_len to restore after estimation
     original_max_model_len = vllm_config.model_config.max_model_len
@@ -961,8 +988,7 @@ def check_enough_kv_cache_memory(
     kv_cache_spec: dict[str, KVCacheSpec],
     available_memory: int,
 ):
-    """
-    Checks whether `available_memory` is enough for the KV cache to hold at
+    """Checks whether `available_memory` is enough for the KV cache to hold at
     least one request with the model's max_model_len.
 
     Args:
@@ -972,8 +998,8 @@ def check_enough_kv_cache_memory(
 
     Raises:
         ValueError: If there is not enough memory available for the KV cache.
-    """
 
+    """
     # No need to check for available memory if the kv_cache_spec is empty
     if kv_cache_spec:
         # Reserve the null block BlockPool permanently holds back, so the check
@@ -996,8 +1022,7 @@ def check_enough_kv_cache_memory(
 def create_kv_cache_group_specs(
     kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]]
 ) -> list[KVCacheGroupSpec]:
-    """
-    Create KVCacheGroupSpec object for each kv cache group layer.
+    """Create KVCacheGroupSpec object for each kv cache group layer.
     The layers in the same group should share the same
     KVCacheSpec.
 
@@ -1008,8 +1033,10 @@ def create_kv_cache_group_specs(
             A list of kv cache groups, where each element is a list of layer
             names that belong to the same group and should share the same
             KVCacheSpec.
+
     Returns:
         A list of KVCacheGroupSpec objects, one for each group.
+
     """
     kv_cache_groups = []
     for layer_names_one_group in grouped_layer_names:
@@ -1024,8 +1051,7 @@ def create_kv_cache_group_specs(
 
 
 def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
-    """
-    Whether all layers in the given KVCacheSpec have the same KV cache spec.
+    """Whether all layers in the given KVCacheSpec have the same KV cache spec.
     Note that we regard FullAttentionSpec with and without sliding window as
     the same type.
 
@@ -1034,8 +1060,8 @@ def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
 
     Returns:
         True if all layers have the same type, False otherwise.
-    """
 
+    """
     if not kv_cache_spec:
         # Encoder-only models do not have KV cache, kv_cache_type can be
         # regarded as uniform.
@@ -1051,31 +1077,40 @@ def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
 def get_max_concurrency_for_kv_cache_config(
     vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
 ) -> float:
-    """
-    Get the maximum concurrency for the given KV cache configuration.
+    """Get the maximum concurrency for the given KV cache configuration.
 
     A request at max_model_len consumes whole blocks from each group's block
     table — cdiv(per-request bytes, page bytes) of the group's spec — and all
-    groups draw those block ids from one shared pool, so the per-request
+    device groups draw those block ids from one shared pool, so the per-request
     total is the sum over groups. The memory/page ratio is identical whether
     a group carries an aggregated UniformTypeKVCacheSpecs (worker config) or
     a representative per-layer spec (scheduler config), so both capacity
     call sites agree.
+
+    Host groups use a separate pool; the smaller concurrency limit applies.
     """
-    num_blocks_per_request = sum(
-        cdiv(
+    num_blocks_per_request = 0
+    host_blocks_per_request = 0
+    for group in kv_cache_config.kv_cache_groups:
+        required = cdiv(
             group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
             group.kv_cache_spec.page_size_bytes,
         )
-        for group in kv_cache_config.kv_cache_groups
-    )
-    max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request
-    return max_concurrency
+        if group.host_resident:
+            host_blocks_per_request += required
+        else:
+            num_blocks_per_request += required
+    limits = [kv_cache_config.num_blocks / num_blocks_per_request]
+    if host_blocks_per_request:
+        assert kv_cache_config.hisparse_host_num_blocks is not None
+        limits.append(
+            kv_cache_config.hisparse_host_num_blocks / host_blocks_per_request
+        )
+    return min(limits)
 
 
 def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
-    """
-    Override the number of kv cache blocks if `num_gpu_blocks_override` is set.
+    """Override the number of kv cache blocks if `num_gpu_blocks_override` is set.
     The override is logged once, at the call site in `get_kv_cache_configs`.
     """
     if vllm_config.cache_config.num_gpu_blocks_override is not None:
@@ -1084,8 +1119,7 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
 
 
 def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
-    """
-    Bytes consumed by one block in the worker's shared KV cache pool, mirroring
+    """Bytes consumed by one block in the worker's shared KV cache pool, mirroring
     the divisor used by `get_kv_cache_config_from_groups` to convert
     `available_memory` into `num_blocks`. Used to compute the effective KV cache
     capacity once `num_gpu_blocks_override` is applied.
@@ -1094,9 +1128,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
 
 
 def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:
-    """
-    Get the page size of the KV cache.
-    """
+    """Get the page size of the KV cache."""
     page_sizes = {layer.page_size_bytes for layer in kv_cache_specs}
     assert len(page_sizes) == 1
     return page_sizes.pop()
@@ -1105,8 +1137,7 @@ def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:
 def _get_kv_cache_groups_uniform_spec(
     kv_cache_specs: dict[str, KVCacheSpec],
 ) -> list[KVCacheGroupSpec]:
-    """
-    Generates the KV cache configuration for a model with the same KV cache
+    """Generates the KV cache configuration for a model with the same KV cache
     spec for all layers.
 
     Args:
@@ -1114,16 +1145,15 @@ def _get_kv_cache_groups_uniform_spec(
 
     Returns:
         The generated KVCacheGroupSpecs
-    """
 
+    """
     return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())])
 
 
 def _get_kv_cache_groups_uniform_type(
     spec: UniformTypeKVCacheSpecs,
 ) -> list[KVCacheGroupSpec]:
-    """
-    Generates the KV cache configuration for a model with one type of KV cache
+    """Generates the KV cache configuration for a model with one type of KV cache
     but different hidden sizes. All layers are merged into one group.
 
     Args:
@@ -1131,8 +1161,8 @@ def _get_kv_cache_groups_uniform_type(
 
     Returns:
         The generated KVCacheGroupSpecs
-    """
 
+    """
     return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)]
 
 
@@ -1338,8 +1368,7 @@ def _glm5_next_tensor_layout(
 def unify_kv_cache_spec_page_size(
     kv_cache_spec: dict[str, KVCacheSpec],
 ) -> dict[str, KVCacheSpec]:
-    """
-    Unify the page size of the given KVCacheSpec. If the page size of all layers
+    """Unify the page size of the given KVCacheSpec. If the page size of all layers
     are the same, return the original KVCacheSpec. If not same, unify the page
     size by increasing the block size of layers with smaller page size. Two
     cases cannot be unified by block size alone and pad their physical page to
@@ -1358,6 +1387,7 @@ def unify_kv_cache_spec_page_size(
 
     Returns:
         The updated KVCacheSpec with the same page_size_bytes.
+
     """
     page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()}
     if len(page_sizes) <= 1:
@@ -1408,8 +1438,7 @@ def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bo
 def _get_kv_cache_groups_uniform_page_size(
     kv_cache_spec: dict[str, KVCacheSpec],
 ) -> list[KVCacheGroupSpec]:
-    """
-    Generates the KV cache groups for hybrid models with multiple
+    """Generates the KV cache groups for hybrid models with multiple
     attention types but still with a uniform page size (physical memory per
     block per layer) for all layers.
 
@@ -1469,6 +1498,7 @@ def _get_kv_cache_groups_uniform_page_size(
         kv_cache_spec: The KVCacheSpec of each attention layer in the model
     Returns:
         The generated KVCacheGroupSpecs
+
     """
     # Group all layers by kv_cache_spec.
     # E.g., 2 full attention layers and 3 sliding window attention layers,
@@ -1574,6 +1604,22 @@ def _get_kv_cache_bytes_per_block(
         for group in kv_cache_groups
     )
     assert bytes_per_block > 0
+    hot_page_sizes = [
+        group.kv_cache_spec.page_size_bytes
+        for group in kv_cache_groups
+        if isinstance(group.kv_cache_spec, HiSparseHotSpec)
+    ]
+    if hot_page_sizes:
+        bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes))
+    stride_alignments = [
+        spec.block_stride_alignment
+        for group in kv_cache_groups
+        for layer_name in group.layer_names
+        if isinstance(spec := _get_per_layer_spec(group, layer_name), MLAAttentionSpec)
+        and spec.block_stride_alignment
+    ]
+    if stride_alignments:
+        bytes_per_block = round_up(bytes_per_block, math.lcm(*stride_alignments))
     return bytes_per_block
 
 
@@ -1616,8 +1662,7 @@ def get_kv_cache_config_from_groups(
     kv_cache_groups: list[KVCacheGroupSpec],
     available_memory: int,
 ) -> KVCacheConfig:
-    """
-    Generate the KV cache configuration from the KV cache groups and spec
+    """Generate the KV cache configuration from the KV cache groups and spec
     of each layer.
 
     Args:
@@ -1626,6 +1671,7 @@ def get_kv_cache_config_from_groups(
         available_memory: Memory available for KV cache in bytes
     Returns:
         The generated KVCacheConfig
+
     """
     if len(kv_cache_groups) == 0:
         # Attention free models do not have KV cache.
@@ -1639,6 +1685,12 @@ def get_kv_cache_config_from_groups(
             ),
         )
 
+    if vllm_config.attention_config.hisparse_config is not None:
+        host_budget = get_hisparse_host_pool_bytes(vllm_config)
+        return get_hisparse_kv_cache_config(
+            vllm_config, kv_cache_groups, available_memory, host_budget
+        )
+
     if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
         (
             attn_group,
@@ -1877,15 +1929,14 @@ def _try_get_full_allocation_fallback_groups(
 
 
 def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
-    """
-    This function tries to convert the KV cache specs to one type if the model
+    """This function tries to convert the KV cache specs to one type if the model
     is a hybrid model with multiple type of KV cache. It will convert all
     SlidingWindowSpec to FullAttentionSpec if both types are present.
 
     Args:
         kv_cache_spec: The kv cache spec of each attention layer in the model
-    """
 
+    """
     if is_kv_cache_spec_uniform(
         kv_cache_spec
     ) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec):
@@ -2013,15 +2064,22 @@ def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int:
             cdiv(len(names), n) * page for page, names in page_size_layers.items()
         )
 
-    # Bytes a block must hold however the mamba buckets end up split: a mamba
-    # bucket can go down to one state per group, every other bucket's split is
-    # already fixed by the repeat pattern.
+    # Bytes a block must hold however the state buckets end up split: mamba,
+    # circular-buffer and unsplit sliding-window buckets can go down to one
+    # state per group, every other bucket's split is fixed by the repeat pattern.
+    def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool:
+        if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)):
+            return True
+        return repeats_per_group is None and isinstance(
+            spec.first_spec, SlidingWindowSpec
+        )
+
     anchor_bytes = max(
         (
             widest_group_bytes(
                 page_size_layers,
                 len(spec.kv_cache_specs)
-                if isinstance(spec.first_spec, MambaSpec)
+                if is_state_bucket(spec)
                 else num_groups_for(spec, balanced),
             )
             for spec, page_size_layers, balanced in bucketed
@@ -2032,10 +2090,9 @@ def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int:
     groups = []
     for spec, page_size_layers, balanced in bucketed:
         num_groups = num_groups_for(spec, balanced)
-        # `_align_hybrid_block_size` pads a mamba state up to one attention
-        # page, so cap a mamba group at the states a block already fits rather
-        # than let it widen the block.
-        if anchor_bytes and isinstance(spec.first_spec, MambaSpec):
+        # Cap a state group at the states a block already fits rather than let
+        # it widen the block.
+        if anchor_bytes and is_state_bucket(spec):
             states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1)
             num_groups = max(
                 num_groups, cdiv(len(spec.kv_cache_specs), states_per_block)
@@ -2060,27 +2117,49 @@ def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int:
         vllm_config,
         kv_cache_spec,
         groups,
-        use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config),
+        use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config),
     )
     _warn_if_unannotated_eagle_mamba(vllm_config, groups)
     return groups
 
 
-def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool:
+def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool:
+    """Whether the drafter's KV layers can be located positionally.
+
+    MTP drafters register their KV layers after the target layers, without a
+    distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too.
+    The annotator separately checks that the groups partition the layers exactly.
+    """
     spec_config = vllm_config.speculative_config
     if spec_config is None or not spec_config.use_eagle():
         return False
+    if spec_config.method == "mtp":
+        return True
     model_config = vllm_config.model_config
-    return (
-        model_config is not None and model_config.hf_config.model_type == "deepseek_v4"
+    return model_config is not None and model_config.hf_config.model_type in (
+        "deepseek_v4",
+        "deepseek_v41",
     )
 
 
+def _groups_partition_layers_exactly(
+    kv_cache_spec: dict[str, KVCacheSpec],
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> bool:
+    """Whether the groups cover every layer of ``kv_cache_spec`` exactly once.
+
+    The trailing-layer draft fallback is only meaningful when the last
+    registered layer is guaranteed to appear in exactly one group.
+    """
+    names = [name for group in kv_cache_groups for name in group.layer_names]
+    return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec)
+
+
 def _annotate_eagle_groups(
     vllm_config: VllmConfig,
     kv_cache_spec: dict[str, KVCacheSpec],
     kv_cache_groups: list[KVCacheGroupSpec],
-    use_deepseek_v4_fallback: bool = False,
+    use_trailing_layer_fallback: bool = False,
 ) -> None:
     """Flag the KV cache groups that hold drafter attention layers.
 
@@ -2093,14 +2172,14 @@ def _annotate_eagle_groups(
        spec merging, wherever grouping happens to land. It is sufficient but
        not necessary: a drafter whose spec is indistinguishable from the
        target's cannot be found this way.
-    2. Model-scoped positional fallback for DeepseekV4, whose MTP block reuses
-       the target's own decoder layer and so carries no spec marker. Its draft
-       attention layer is always the last registered layer, so flag whichever
-       group holds it. This rule is only valid where the groups partition
-       exactly the layers of ``kv_cache_spec``, which is true on the packed
-       grouping path and not in general; other callers must leave
-       ``use_deepseek_v4_fallback`` False. The caller gates this fallback on
-       the configured model type.
+    2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose
+       MTP block reuses the target's own decoder layer and so carries no spec
+       marker. Their draft attention layers always register after every
+       target layer, so flag whichever group holds the last registered layer.
+       This rule is only valid where the groups partition exactly the layers
+       of ``kv_cache_spec``, which is re-checked here before applying it.
+       When a drafter's trailing caches span several groups, this rule flags
+       only the group holding the very last layer and must be generalized.
        FIXME(yifan): avoid/generalize this hacky check.
 
     Args:
@@ -2108,7 +2187,9 @@ def _annotate_eagle_groups(
         kv_cache_spec: The kv cache spec of each attention layer, in layer
             registration order. Only read by rule 2.
         kv_cache_groups: Groups to annotate in place.
-        use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group.
+        use_trailing_layer_fallback: Enable rule 2. Callers gate this on
+            ``_uses_trailing_mtp_layers``.
+
     """
     spec_config = vllm_config.speculative_config
     if spec_config is None or not spec_config.use_eagle_block_drop():
@@ -2121,7 +2202,9 @@ def _annotate_eagle_groups(
         ):
             group.is_eagle_group = True
 
-    if not use_deepseek_v4_fallback:
+    if not use_trailing_layer_fallback:
+        return
+    if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups):
         return
     last_layer = next(reversed(kv_cache_spec))
     for group in kv_cache_groups:
@@ -2134,16 +2217,12 @@ def _warn_if_unannotated_eagle_mamba(
     vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
 ) -> None:
-    """Warn when the flag-all eagle fallback will silently disable reuse.
-
-    With no group annotated, consumers flag every group as a draft group. That
-    widens a Mamba group's required lookup window to two consecutive chunks,
-    which align-mode checkpointing never produces, so reuse drops to zero with
-    no error and no metric to show it.
+    """Warn when no KV cache group could be identified as the draft model's.
 
     Args:
         vllm_config: Config supplying the speculative method, if any.
         kv_cache_groups: Groups as they will be handed to consumers.
+
     """
     spec_config = vllm_config.speculative_config
     if spec_config is None or not spec_config.use_eagle():
@@ -2162,13 +2241,8 @@ def _warn_if_unannotated_eagle_mamba(
         return
     logger.warning(
         "Speculative decoding (method=%s) is enabled but no KV cache group "
-        "could be identified as the draft model's, so every group -- "
-        "including Mamba groups %s -- will be treated as a draft group. A "
-        "Mamba group cannot satisfy the widened lookup window that implies, "
-        "so prefix-cache reuse across requests will be disabled and any "
-        "external KV offload tier will store without ever serving a hit.",
+        "could be identified as the draft model's.",
         spec_config.method,
-        mamba_groups,
     )
 
 
@@ -2179,12 +2253,43 @@ def _largest_divisor_at_most(value: int, limit: int) -> int:
     return 1
 
 
+def _ensure_min_page_size(
+    groups: list[KVCacheGroupSpec],
+    common_page: int,
+    hidden_specs: dict[str, HiddenStateCacheSpec],
+) -> tuple[list[KVCacheGroupSpec], int]:
+    """Scale up group block sizes so the common page is at least as large as
+    the biggest hidden-state per-token cost.
+
+    This protects the hidden state extraction feature, where we store hidden states in
+    the KV cache, and return them to clients using a custom hidden state connector.
+
+    Returns the (possibly scaled) groups and updated common page size.
+    """
+    max_per_token = max(
+        spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)
+        for spec in hidden_specs.values()
+    )
+    if max_per_token <= common_page:
+        return groups, common_page
+
+    scale = cdiv(max_per_token, common_page)
+    common_page *= scale
+    scaled: list[KVCacheGroupSpec] = []
+    for g in groups:
+        s = g.kv_cache_spec
+        kw: dict[str, int] = {"block_size": s.block_size * scale}
+        if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None:
+            kw["page_size_padded"] = s.page_size_padded * scale
+        scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw)))
+    return scaled, common_page
+
+
 def get_kv_cache_groups(
     vllm_config: VllmConfig,
     kv_cache_spec: dict[str, KVCacheSpec],
 ) -> list[KVCacheGroupSpec]:
-    """
-    Split the layers in the model into groups with the same KV cache spec.
+    """Split the layers in the model into groups with the same KV cache spec.
 
     Args:
         vllm_config: The global VllmConfig
@@ -2192,6 +2297,7 @@ def get_kv_cache_groups(
 
     Returns:
         The generated KVCacheGroups
+
     """
     if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
         unify_hybrid_kv_cache_specs(kv_cache_spec)
@@ -2201,6 +2307,9 @@ def get_kv_cache_groups(
         # attention free models.
         return []
 
+    if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec):
+        return hisparse_groups
+
     if is_kv_cache_spec_uniform(kv_cache_spec):
         # KV cache of all layers are the same, which is true for
         # most models. Allocate the same amount of memory for
@@ -2247,6 +2356,8 @@ def get_kv_cache_groups(
     # Add hidden-state layers back with page aligned to the common page.
     if hidden_specs:
         common_page = get_uniform_page_size([g.kv_cache_spec for g in groups])
+        # TP may shrink the common page below hidden-state per-token cost.
+        groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs)
         group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups))
         for name, spec in hidden_specs.items():
             per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)
@@ -2264,7 +2375,12 @@ def get_kv_cache_groups(
             aligned = replace(spec, block_size=new_bs, page_size_padded=common_page)
             groups.append(KVCacheGroupSpec([name], aligned))
 
-    _annotate_eagle_groups(vllm_config, kv_cache_spec, groups)
+    _annotate_eagle_groups(
+        vllm_config,
+        kv_cache_spec,
+        groups,
+        use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config),
+    )
     _warn_if_unannotated_eagle_mamba(vllm_config, groups)
     return groups
 
@@ -2272,12 +2388,14 @@ def get_kv_cache_groups(
 def generate_scheduler_kv_cache_config(
     kv_cache_configs: list[KVCacheConfig],
 ) -> KVCacheConfig:
-    """
-    Generate the KV cache configuration for the scheduler.
-    """
+    """Generate the KV cache configuration for the scheduler."""
     assert all(
         [cfg.num_blocks == kv_cache_configs[0].num_blocks for cfg in kv_cache_configs]
     )
+    assert all(
+        cfg.hisparse_host_num_blocks == kv_cache_configs[0].hisparse_host_num_blocks
+        for cfg in kv_cache_configs
+    )
     # All workers have the same kv_cache_config except layer names, so use
     # an arbitrary one to initialize the scheduler.
     cfg = copy.deepcopy(kv_cache_configs[0])
@@ -2294,9 +2412,7 @@ def generate_scheduler_kv_cache_config(
 def get_kv_cache_capacity(
     vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
 ) -> tuple[int, float]:
-    """
-    Get the group-aware KV cache token capacity and max concurrency.
-    """
+    """Get the group-aware KV cache token capacity and max concurrency."""
     max_model_len = vllm_config.model_config.max_model_len
     max_concurrency = get_max_concurrency_for_kv_cache_config(
         vllm_config, kv_cache_config
@@ -2312,9 +2428,11 @@ def update_kv_cache_capacity(
     vllm_config.cache_config.kv_cache_size_tokens = num_tokens
     vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency
     max_model_len = vllm_config.model_config.max_model_len
+    device_type = vllm_config.device_config.device_type
     logger.info_once(
-        "GPU KV cache size: %s tokens, "
+        "%s KV cache size: %s tokens, "
         "Maximum concurrency for %s tokens per request: %.2fx",
+        "GPU" if device_type == "cuda" else device_type.upper(),
         f"{num_tokens:,}",
         f"{max_model_len:,}",
         max_concurrency,
@@ -2325,8 +2443,7 @@ def _max_memory_usage_bytes_from_groups(
     vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
 ) -> int:
-    """
-    Calculate maximum memory usage in bytes from KV cache groups.
+    """Calculate maximum memory usage in bytes from KV cache groups.
 
     This correctly accounts for padding in hybrid models. For example, if a
     model has 8 full attention layers and 9 sliding window layers, they will
@@ -2338,6 +2455,9 @@ def _max_memory_usage_bytes_from_groups(
     if not kv_cache_groups:
         return 0
 
+    if vllm_config.attention_config.hisparse_config is not None:
+        return get_hisparse_gpu_memory_usage(vllm_config, kv_cache_groups)
+
     if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
         (
             attn_group,
@@ -2382,14 +2502,25 @@ def _estimate_max_model_len_from_groups(
     kv_cache_groups: list[KVCacheGroupSpec],
     available_memory: int,
 ) -> int:
-    """
-    Binary search for the maximum model length that fits in available memory.
+    """Binary search for the maximum model length that fits in available memory.
     Returns 0 if even 1 token doesn't fit.
     """
     original_max = vllm_config.model_config.max_model_len
+    hisparse_enabled = (
+        vllm_config.attention_config.hisparse_config is not None
+        and bool(kv_cache_groups)
+    )
 
     def fits(model_len: int) -> bool:
         vllm_config.model_config.max_model_len = model_len
+        if hisparse_enabled:
+            try:
+                config = get_kv_cache_config_from_groups(
+                    vllm_config, kv_cache_groups, available_memory
+                )
+            except ValueError:
+                return False
+            return get_max_concurrency_for_kv_cache_config(vllm_config, config) >= 1
         return (
             _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups)
             <= available_memory
@@ -2417,8 +2548,7 @@ def _auto_fit_max_model_len(
     projected_groups_per_worker: list[list[KVCacheGroupSpec]],
     available_memory: list[int],
 ) -> None:
-    """
-    When max_model_len is set to -1, this function estimates the largest
+    """When max_model_len is set to -1, this function estimates the largest
     context length that can be supported with the available GPU memory.
     It uses binary search to find the maximum length that fits across all
     workers.
@@ -2428,6 +2558,7 @@ def _auto_fit_max_model_len(
         projected_groups_per_worker: KV cache groups projected to each worker.
         available_memory: Memory available for KV cache in bytes for each
             worker.
+
     """
     original_max = vllm_config.model_config.max_model_len
 
@@ -2480,8 +2611,7 @@ def _project_kv_cache_groups_to_worker(
     global_kv_cache_groups: list[KVCacheGroupSpec],
     worker_spec: dict[str, KVCacheSpec],
 ) -> list[KVCacheGroupSpec]:
-    """
-    Projects global KV cache groups onto a single worker's assigned layers.
+    """Projects global KV cache groups onto a single worker's assigned layers.
 
     In pipeline parallelism, each worker only owns a subset of layers. This
     function filters the global groups to include only layers present on the
@@ -2493,6 +2623,7 @@ def _project_kv_cache_groups_to_worker(
 
     Returns:
         The projected KV cache groups containing only this worker's layers.
+
     """
     projected_groups: list[KVCacheGroupSpec] = []
     for group in global_kv_cache_groups:
@@ -2523,8 +2654,7 @@ def get_kv_cache_configs(
     kv_cache_specs: list[dict[str, KVCacheSpec]],
     available_memory: list[int],
 ) -> list[KVCacheConfig]:
-    """
-    Generates the KV cache configurations for a model.
+    """Generates the KV cache configurations for a model.
     Since we use a shared centralized controller for all workers, we need the
     `kv_cache_config` to be consistent across all workers to make sure
     the KV cache allocation can be applied to all workers. However, different
@@ -2551,8 +2681,8 @@ def get_kv_cache_configs(
 
     Returns:
         The generated KVCacheConfigs for each worker.
-    """
 
+    """
     # Merge the KV cache specs of all workers. Different PP stages may have
     # different layer names, and different TP ranks of the same PP stage should
     # have the same KV cache spec.
@@ -2573,12 +2703,7 @@ def get_kv_cache_configs(
 
     # When speculating with more than 1 speculative module (e.g. multi-layered MTP)
     # tag every SlidingWindowSpec with how many extra tokens to retain in the window.
-    extra_retained_tokens = (
-        vllm_config.speculative_config.num_speculative_tokens - 1
-        if vllm_config.speculative_config is not None
-        and vllm_config.speculative_config.use_multi_module_mtp()
-        else 0
-    )
+    extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1)
     for layer_name, layer_spec in merged_kv_cache_specs.items():
         if isinstance(layer_spec, SlidingWindowSpec):
             merged_kv_cache_specs[layer_name] = replace(
@@ -2620,6 +2745,9 @@ def get_kv_cache_configs(
             adjusted_memory.append(override * bytes_per_block)
         available_memory = adjusted_memory
 
+    if vllm_config.attention_config.hisparse_config is not None:
+        available_memory = [min(available_memory)] * len(available_memory)
+
     # Reserve the null block BlockPool permanently holds back, so auto-fit and
     # the capacity check both plan against usable blocks. Allocation below
     # still uses the full memory.
@@ -2675,8 +2803,7 @@ def get_kv_cache_configs(
 
 
 class BlockHashListWithBlockSize:
-    """
-    Convert block-hash granularity from `hash_block_size` to `target_block_size`.
+    """Convert block-hash granularity from `hash_block_size` to `target_block_size`.
     Used when KV cache groups have different block sizes: `hash_block_size`
     is the size used to compute the original `block_hashes`; `target_block_size`
     is the group's actual block size.
@@ -2706,6 +2833,7 @@ class BlockHashListWithBlockSize:
         block_hashes: Block hashes to convert, computed at `hash_block_size`.
         hash_block_size: Block size at which `block_hashes` were computed.
         target_block_size: Desired block size; must be a multiple of `hash_block_size`.
+
     """
 
     def __init__(
diff --git a/vllm/v1/core/sched/diffusion_scheduler.py b/vllm/v1/core/sched/diffusion_scheduler.py
new file mode 100644
index 000000000000..96462aa6abe1
--- /dev/null
+++ b/vllm/v1/core/sched/diffusion_scheduler.py
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Async scheduling for diffusion requests.
+
+These rules live here so Scheduler.schedule() and AsyncScheduler stay
+unchanged. VllmConfig selects this class for a diffusion model under async
+scheduling. A sync scheduler creates no output placeholders, which both rules
+read.
+"""
+
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.request import Request
+
+
+def diffusion_canvas_width(request: Request, canvas_length: int) -> int:
+    """The canvas width a diffusion request asked for, else the served one."""
+    params = request.sampling_params
+    extra = params.extra_args if params is not None else None
+    width = extra.get("diffusion_canvas_length") if extra else None
+    return int(width) if width else canvas_length
+
+
+def _read_in_flight(request: Request, width: int) -> bool:
+    """True when a read-only request has all its denoise steps in flight.
+
+    The request emits its canvas on the last step and ends, so a further
+    step is discarded.
+    """
+    params = request.sampling_params
+    extra = params.extra_args if params is not None else None
+    if not extra or not extra.get("diffusion_read_only"):
+        return False
+    steps = extra.get("diffusion_max_steps")
+    if not steps:
+        return False
+    # Each in-flight denoise step holds one canvas of placeholders.
+    return request.num_output_placeholders >= steps * width
+
+
+class DiffusionAsyncScheduler(AsyncScheduler):
+    def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
+        for request in self.running:
+            width = diffusion_canvas_width(request, self.num_spec_tokens)
+            # The placeholders are as wide as the served canvas.
+            if len(request.spec_token_ids) > width:
+                request.spec_token_ids = request.spec_token_ids[:width]
+            if _read_in_flight(request, width):
+                # Scheduler.schedule()'s max_tokens guard cannot be reached
+                # from a subclass. schedule() advances current_step before it
+                # checks decode eligibility, so current_step + 2 skips this
+                # call alone. max keeps a longer pipeline-parallel wait.
+                request.next_decode_eligible_step = max(
+                    request.next_decode_eligible_step, self.current_step + 2
+                )
+        return super().schedule(throttle_prefills)
diff --git a/vllm/v1/core/sched/interface.py b/vllm/v1/core/sched/interface.py
index c15296dd0516..c868d249c8a5 100644
--- a/vllm/v1/core/sched/interface.py
+++ b/vllm/v1/core/sched/interface.py
@@ -79,6 +79,7 @@ def schedule(self, throttle_prefills: bool = False) -> "SchedulerOutput":
         Returns:
             A SchedulerOutput object containing information about the scheduled
             requests.
+
         """
         raise NotImplementedError
 
@@ -105,6 +106,7 @@ def update_from_output(
         Returns:
             A dict of client index to EngineCoreOutputs object containing the
             outputs for each request originating from that client.
+
         """
         raise NotImplementedError
 
@@ -115,6 +117,7 @@ def update_draft_token_ids(self, draft_token_ids: "DraftTokenIds") -> None:
 
         Args:
             draft_token_ids: The input draft token ids for each request.
+
         """
         raise NotImplementedError
 
@@ -129,6 +132,7 @@ def update_draft_token_ids_in_output(
             draft_token_ids: The input draft token ids for each request.
             scheduler_output: Update the given scheduler_output
                 with the corresponding draft token ids.
+
         """
         raise NotImplementedError
 
@@ -138,6 +142,7 @@ def add_request(self, request: "Request") -> None:
 
         Args:
             request: The new request being added.
+
         """
         raise NotImplementedError
 
@@ -162,6 +167,7 @@ def finish_requests(
         Returns:
             List of requests that were aborted. Will not include any that were
             already finished.
+
         """
         raise NotImplementedError
 
@@ -218,6 +224,8 @@ def reset_prefix_cache(
                 preempted and moved to the waiting queue. Otherwise, this method
                 will only reset the KV prefix cache when there is no running request
                 taking KV cache.
+            reset_connector: If True, also reset any KV connector state.
+
         """
         raise NotImplementedError
 
diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py
index 93fc2d83c151..0da64cf27e0f 100644
--- a/vllm/v1/core/sched/output.py
+++ b/vllm/v1/core/sched/output.py
@@ -14,6 +14,9 @@
     import numpy.typing as npt
     import torch
 
+    from vllm.distributed.aux_output_connector.connector import (
+        AuxOutputConnectorMetadata,
+    )
     from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata
     from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
     from vllm.lora.request import LoRARequest
@@ -23,6 +26,7 @@
     from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
     from vllm.v1.request import Request
 else:
+    AuxOutputConnectorMetadata = object
     ECConnectorMetadata = object
     KVConnectorMetadata = object
     KVCacheBlockCopy = object
@@ -48,6 +52,8 @@ class NewRequestData:
 
     # Only used for v2 model runner.
     prefill_token_ids: list[int] | None = None
+    # DeepSeek-V4.1 only: SWA bounded replay; see Request.replay_start.
+    replay_start: int = 0
 
     @classmethod
     def from_request(
@@ -73,6 +79,7 @@ def from_request(
             prompt_embeds=request.prompt_embeds,
             prompt_is_token_ids=request.prompt_is_token_ids,
             prefill_token_ids=prefill_token_ids,
+            replay_start=request.replay_start,
         )
 
     @property
@@ -282,6 +289,9 @@ class SchedulerOutput:
     # synchronously during this step (load_async=False).
     has_sync_kv_loads: bool = False
 
+    # Execution-auxiliary output control metadata consumed by the worker connector.
+    aux_output_connector_metadata: AuxOutputConnectorMetadata | None = None
+
     # EC Cache Connector metadata
     ec_connector_metadata: ECConnectorMetadata | None = None
     # EC Cache Manager metadata
@@ -294,6 +304,9 @@ class SchedulerOutput:
     # CoW copies to apply after zeroing new blocks and before forward.
     kv_cache_block_copies: list[KVCacheBlockCopy] | None = None
 
+    # Complete block-table rows that replace incrementally appended block IDs.
+    block_table_updates: dict[str, tuple[list[int], ...]] | None = None
+
     # Scheduler-local; always None by the time this reaches a worker.
     kv_connector_block_state: KVConnectorBlockState | None = None
 
diff --git a/vllm/v1/core/sched/request_queue.py b/vllm/v1/core/sched/request_queue.py
index 38c7db94a45e..ce58e6e1c7c3 100644
--- a/vllm/v1/core/sched/request_queue.py
+++ b/vllm/v1/core/sched/request_queue.py
@@ -129,8 +129,7 @@ def __iter__(self) -> Iterator[Request]:
 
 
 class PriorityRequestQueue(RequestQueue):
-    """
-    A priority queue that supports heap operations.
+    """A priority queue that supports heap operations.
 
     Respects the ordering defined in the Request class, where
     requests with a smaller value of `priority` are processed first.
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 01423212999f..ce832a18f655 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -9,12 +9,14 @@
 
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.config import KVEventsConfig, VllmConfig
+from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
 from vllm.distributed.ec_transfer.ec_connector.base import (
     ECConnectorBase,
     ECConnectorMetadata,
     ECConnectorRole,
 )
 from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory
+from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats
 from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch
 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
 from vllm.distributed.kv_transfer.kv_connector.v1 import (
@@ -25,12 +27,10 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    RoutedExpertsManager,
-)
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
 from vllm.multimodal.encoder_budget import MultiModalBudget
 from vllm.multimodal.utils import get_mm_features_in_window
+from vllm.utils.math_utils import cdiv
 from vllm.v1.core.encoder_cache_manager import (
     EncoderCacheManager,
     EncoderDecoderCacheManager,
@@ -70,8 +70,7 @@
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
-from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager
-from vllm.v1.structured_output.utils import strip_speculative_padding
+from vllm.v1.structured_output import StructuredOutputManager
 from vllm.v1.utils import record_function_or_nullcontext
 
 logger = init_logger(__name__)
@@ -123,11 +122,21 @@ def __init__(
 
         # Scheduling constraints.
         self.max_num_running_reqs = self.scheduler_config.max_num_seqs
+        # Admission into RUNNING. Separate from max_num_running_reqs, which
+        # is the model-runner slot count.
+        self.max_num_active_reqs = (
+            self.scheduler_config.max_num_active_seqs
+            if self.scheduler_config.max_num_active_seqs is not None
+            else self.max_num_running_reqs
+        )
         self.max_num_scheduled_tokens = (
             self.scheduler_config.max_num_scheduled_tokens
             if self.scheduler_config.max_num_scheduled_tokens is not None
             else self.scheduler_config.max_num_batched_tokens
         )
+        self.adaptive_long_prefill_threshold = (
+            self.scheduler_config.long_prefill_token_threshold_adaptive
+        )
         self.max_model_len = vllm_config.model_config.max_model_len
         self.enable_kv_cache_events = (
             self.kv_events_config is not None
@@ -225,12 +234,11 @@ def __init__(
         # Grammar compilation failures to finish as per-request errors in
         # update_from_output.
         self.grammar_compile_error_reqs: set[str] = set()
+        self.encoder_cache_mismatch_reqs: set[str] = set()
 
         # Encoder-related.
         # Calculate encoder cache size if applicable
-        supports_mm_inputs = mm_registry.supports_multimodal_inputs(
-            vllm_config.model_config
-        )
+        supports_mm_inputs = vllm_config.model_config.supports_multimodal_inputs
         mm_budget = (
             MultiModalBudget(vllm_config, mm_registry) if supports_mm_inputs else None
         )
@@ -263,13 +271,19 @@ def __init__(
         self.use_eagle_block_drop = False
         self.num_spec_tokens = vllm_config.num_speculative_tokens
         self.num_lookahead_tokens = vllm_config.num_lookahead_tokens
-        # Positions past the computed tokens that the drafter reads mid-prefill.
-        # Eagle-family drafters read 1 ahead, but multi-module MTP reads
-        # num_spec_tokens ahead at chunked-prefill boundaries. Determines the
-        # encoder scheduling shift, the deferred encoder free, the KV cache
-        # manager's re-prefillable window (this minus 1), and how many tokens to
-        # reserve between a chunk boundary and the prefill end.
-        self.num_prefill_lookahead = 0
+        # DSV41 SWA bounded replay: groups that declare a replay window are rebuilt
+        # after a prefix hit by recomputing its trailing tokens. One window
+        # for all such groups, so the rewind matches every group's allocation.
+        replay_windows = {
+            group.kv_cache_spec.prefix_replay_tokens
+            for group in kv_cache_config.kv_cache_groups
+            if group.kv_cache_spec.prefix_replay_tokens > 0
+        }
+        assert len(replay_windows) <= 1, (
+            f"Prefix replay windows should agree: {sorted(replay_windows)}"
+        )
+        self.prefix_replay_tokens = replay_windows.pop() if replay_windows else 0
+        self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens
         self.dynamic_sd_lookup: list[int] | None = None
         if speculative_config is not None:
             if speculative_config.num_speculative_tokens_per_batch_size:
@@ -279,12 +293,6 @@ def __init__(
                     vllm_num_speculative_tokens=self.num_spec_tokens,
                 )
             self.use_eagle = speculative_config.use_eagle()
-            if self.use_eagle:
-                self.num_prefill_lookahead = (
-                    self.num_spec_tokens
-                    if speculative_config.use_multi_module_mtp()
-                    else 1
-                )
             self.use_eagle_block_drop = speculative_config.use_eagle_block_drop()
             if self.use_eagle and not self.use_eagle_block_drop:
                 logger.warning(
@@ -312,14 +320,13 @@ def __init__(
             hash_block_size=hash_block_size,
             metrics_collector=self.kv_metrics_collector,
             watermark=self.scheduler_config.watermark,
-            enable_mamba_fine_grained_prefix_cache=(
-                self.cache_config.enable_mamba_fine_grained_prefix_cache
+            enable_mamba_shared_prefix_checkpoint=(
+                self.cache_config.enable_mamba_shared_prefix_checkpoint
             ),
         )
-        # Bind GPU block pool to the KV connector. This must happen after
-        # kv_cache_manager is constructed so block_pool is available.
+        # Bind after construction so connectors can access the cache manager.
         if self.connector is not None:
-            self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool)
+            self.connector.bind_kv_cache_manager(self.kv_cache_manager)
 
         self.use_pp = self.parallel_config.pipeline_parallel_size > 1
         self.use_v2_model_runner = vllm_config.use_v2_model_runner
@@ -369,9 +376,9 @@ def __init__(
         # manager decides whether it can check-point there (per-group eagle bit,
         # no MTP re-prefill tail); splitting for a stop it would refuse costs a
         # forward pass and displaces the block-boundary stop.
-        self.mamba_fine_grained_prefix_cache = (
+        self.mamba_shared_prefix_checkpoint = (
             self.mamba_partial_cache_hit
-            and self.kv_cache_manager.mamba_fine_grained_prefix_cache
+            and self.kv_cache_manager.mamba_shared_prefix_checkpoint
         )
 
         # Counts of non-empty steps scheduled / processed. update_from_output
@@ -386,26 +393,13 @@ def __init__(
         if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:
             self.perf_metrics = ModelMetrics(vllm_config)
 
-        self.enable_return_routed_experts = (
-            vllm_config.model_config.enable_return_routed_experts
+        self.aux_output_connector = (
+            AuxOutputSchedulerConnector()
+            if vllm_config.aux_output_config.enabled
+            else None
         )
         self.return_sampling_mask = vllm_config.model_config.return_sampling_mask
 
-        if self.enable_return_routed_experts:
-            assert self.dcp_world_size == 1 and self.pcp_world_size == 1, (
-                "enable_return_routed_experts does not support context parallelism "
-                "(dcp_world_size > 1 or pcp_world_size > 1)"
-            )
-
-            self.routed_experts_mgr = RoutedExpertsManager(
-                vllm_config=vllm_config,
-                kv_cache_config=kv_cache_config,
-            )
-            # Block-ID snapshot taken at schedule time (before forward),
-            # so update_from_output can read slot data even if a later
-            # schedule() frees the blocks (async scheduling race).
-            self._re_block_ids: dict[str, list[int]] = {}
-
         self._pause_state: PauseState = PauseState.UNPAUSED
 
         # In-flight requests still prefilling (prefill chunks + in-progress
@@ -502,7 +496,7 @@ def _mamba_block_aligned_split(
         # output tokens can still observe a junction there.
         junction_stop = (
             junction
-            if self.mamba_fine_grained_prefix_cache
+            if self.mamba_shared_prefix_checkpoint
             and junction <= request.num_prompt_tokens
             else block_floored
         )
@@ -609,6 +603,24 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
             throttle_prefills and not self.prefill_capacity_bound
         ) and any(not r.is_prefill_chunk for r in self.running)
 
+        # `long_prefill_token_threshold` exists to stop a long prefill from
+        # starving other requests of the token budget. When it is the only
+        # request there is nobody to starve, so let it use the whole budget.
+        num_eligible_reqs = (
+            len(self.running) + len(self.waiting) + len(self.skipped_waiting)
+        )
+        long_prefill_token_threshold = (
+            self.scheduler_config.long_prefill_token_threshold
+            if num_eligible_reqs > 1
+            else 0
+        )
+        if long_prefill_token_threshold > 0 and self.adaptive_long_prefill_threshold:
+            # Floor the cap at a fair share of the input budget so it never
+            # cuts a request below max_num_batched_tokens / num requests.
+            long_prefill_token_threshold = max(
+                long_prefill_token_threshold, input_budget // num_eligible_reqs
+            )
+
         # First, schedule the RUNNING requests.
         req_index = 0
         while req_index < len(self.running) and token_budget > 0:
@@ -660,8 +672,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 + request.num_output_placeholders
                 - request.num_computed_tokens
             )
-            if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
-                num_new_tokens = self.scheduler_config.long_prefill_token_threshold
+            if 0 < long_prefill_token_threshold < num_new_tokens:
+                num_new_tokens = long_prefill_token_threshold
             num_new_tokens = min(
                 num_new_tokens, token_budget, input_budget - draft_slots
             )
@@ -738,6 +750,12 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                         # The request can be scheduled.
                         break
 
+                    if (
+                        self.connector is not None
+                        and self.connector.has_pending_block_frees()
+                    ):
+                        break
+
                     # The request cannot be scheduled.
                     # Preempt the lowest-priority request.
                     if self.policy == SchedulingPolicy.PRIORITY:
@@ -745,13 +763,16 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                             self.running,
                             key=lambda r: (r.priority, r.arrival_time),
                         )
-                        # Record the index of the preemption victim to
-                        # maintain accurate loop state.
+                    else:
+                        preempted_req = self.running[-1]
+
+                    # A deferred free will not help with immediate allocation.
+                    if not self._request_blocks_can_be_freed(preempted_req):
+                        break
+
+                    if self.policy == SchedulingPolicy.PRIORITY:
                         victim_index = self.running.index(preempted_req)
                         del self.running[victim_index]
-                        # Decrement the loop cursor if the removed request
-                        # preceded the current iteration, preventing the
-                        # silent omission of the subsequent request.
                         if victim_index < req_index:
                             req_index -= 1
 
@@ -854,7 +875,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not
                 # in `running` but still hold a model-runner request slot.
                 num_running = len(self.running) + self.num_waiting_for_streaming_input
-                if num_running >= self.max_num_running_reqs:
+                if num_running >= self.max_num_active_reqs:
                     break
 
                 request_queue = self._select_waiting_queue_for_scheduling()
@@ -931,6 +952,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                                 request, block_aligned_local
                             )
                         )
+                        if request.skip_reading_prefix_cache:
+                            ext_tokens, load_kv_async = 0, False
 
                         if ext_tokens is None:
                             # The request cannot be scheduled because
@@ -940,6 +963,17 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                             step_skipped_waiting.prepend_request(request)
                             continue
 
+                        if self.prefix_replay_tokens:
+                            # SWA bounded replay recomputes the hit's last
+                            # window, from the block holding its first
+                            # token; the sliding-window groups retire
+                            # whole blocks below the window of the hit's
+                            # next token. Hits end on a block boundary, or
+                            # a hit ending one token short of one would
+                            # retire the block the replay starts in.
+                            ext_tokens -= ext_tokens % self.block_size
+                            load_kv_async = load_kv_async and ext_tokens > 0
+
                         if partial_tail and ext_tokens > partial_tail:
                             # Remote strictly exceeds the full local hit: drop the
                             # sub-block tail so no CoW is needed, and let the load
@@ -982,6 +1016,17 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                         num_new_local_computed_tokens + num_external_computed_tokens
                     )
                     assert num_computed_tokens <= request.num_tokens
+                    if 0 < num_computed_tokens <= self.prefix_replay_tokens:
+                        # SWA bounded replay: a hit no longer than the replayed
+                        # window would be recomputed in full and save nothing.
+                        new_computed_blocks = (
+                            self.kv_cache_manager.empty_kv_cache_blocks
+                        )
+                        num_new_local_computed_tokens = 0
+                        num_external_computed_tokens = 0
+                        connector_prefix_cache_hits = 0
+                        num_computed_tokens = 0
+                        load_kv_async = False
 
                     # Skip request with pending mm encoding prefetches
                     if self._ec_transfer_pending(request, num_computed_tokens):
@@ -1016,6 +1061,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 new_encoder_compute_budget = encoder_compute_budget
                 pad_spec_decode = False
 
+                # SWA bounded replay: recompute the tail of the hit without
+                # rewriting its cached KV. An async load replays once the KV
+                # has arrived (_update_waiting_for_remote_kv).
+                num_replay_tokens = 0
+
                 if load_kv_async:
                     # KVTransfer: loading remote KV, do not allocate for new work.
                     assert num_external_computed_tokens > 0
@@ -1026,6 +1076,12 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                     break
                 else:
                     request_token_budget = min(token_budget, input_budget - draft_slots)
+                    if did_prefix_cache_lookup:
+                        # Fresh admission: local and/or sync external hit.
+                        num_replay_tokens = self._mark_prefix_replay(
+                            request, num_computed_tokens
+                        )
+                        num_computed_tokens -= num_replay_tokens
                     # Number of tokens to be scheduled.
                     # We use `request.num_tokens` instead of
                     # `request.num_prompt_tokens` to consider the resumed
@@ -1056,9 +1112,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                             num_new_tokens = padded_num_tokens
                             pad_spec_decode = True
 
-                    threshold = self.scheduler_config.long_prefill_token_threshold
-                    if 0 < threshold < num_new_tokens:
-                        num_new_tokens = threshold
+                    if 0 < long_prefill_token_threshold < num_new_tokens:
+                        num_new_tokens = long_prefill_token_threshold
 
                     # chunked prefill has to be enabled explicitly to allow
                     # pooling requests to be chunked
@@ -1145,13 +1200,20 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 if load_kv_async:
                     # An async load holds its blocks for the whole transfer with
                     # no forward progress and isn't preemptible here. Admit it
-                    # only if it fits in (free - other in-flight reservations), to
-                    # avoid deadlock and predictable preemptions.
-                    reserved_blocks = self._inflight_prefill_reserved_blocks()
+                    # only if it fits in (free - other in-flight reservations)
+                    # plus its own spec decode step blocks, to avoid deadlock and
+                    # predictable preemptions.
+                    reserved_blocks = (
+                        self._inflight_prefill_reserved_blocks()
+                        + self._spec_decode_step_blocks()
+                    )
 
+                # Replayed tokens are already counted in the adopted hit; a
+                # chunk that ends inside the replayed range needs no new slots.
+                num_tokens_past_hit = max(num_new_tokens - num_replay_tokens, 0)
                 new_blocks = self.kv_cache_manager.allocate_slots(
                     request,
-                    num_new_tokens,
+                    num_tokens_past_hit,
                     num_new_computed_tokens=num_new_local_computed_tokens,
                     new_computed_blocks=new_computed_blocks,
                     num_lookahead_tokens=effective_lookahead_tokens,
@@ -1429,6 +1491,13 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
             meta = self._build_kv_connector_meta(self.connector, scheduler_output)
             scheduler_output.kv_connector_metadata = meta
 
+        if self.aux_output_connector is not None:
+            scheduler_output.aux_output_connector_metadata = (
+                self.aux_output_connector.build_connector_meta(
+                    scheduler_output, self.requests
+                )
+            )
+
         # Build the connector meta for ECConnector
         if self.ec_connector is not None:
             ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta(
@@ -1483,6 +1552,8 @@ def _preempt_request(
         assert request.status == RequestStatus.RUNNING, (
             "Only running requests can be preempted"
         )
+        if self.aux_output_connector is not None:
+            self.aux_output_connector.request_finished(request)
         self._free_request_blocks(request)
         self.encoder_cache_manager.free(request)
         self._inflight_prefills.discard(request)
@@ -1538,22 +1609,6 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
             if not request.is_prefill_chunk:
                 self._inflight_prefills.discard(request)
 
-        # Snapshot block IDs for routed experts before forward starts.
-        # A concurrent schedule() may preempt requests and free blocks
-        # before update_from_output runs; the snapshot survives that.
-        # Use update() to preserve entries from the previous step that
-        # have not yet been consumed by update_from_output (async
-        # scheduling may call _update_after_schedule again before the
-        # prior update_from_output runs).
-        if self.enable_return_routed_experts:
-            gid = self.routed_experts_mgr.attn_gid
-            self._re_block_ids.update(
-                {
-                    rid: self.kv_cache_manager.get_blocks(rid).get_block_ids()[gid]
-                    for rid in num_scheduled_tokens
-                }
-            )
-
         # Clear the finished and preempted request IDs.
         # NOTE: We shouldn't just clear() here because it will also affect
         # the scheduler output.
@@ -1563,12 +1618,10 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
     def _update_request_as_session(
         self, session: Request, update: StreamingUpdate
     ) -> None:
-        """
-        Updates the waiting session with the next streaming update.
+        """Updates the waiting session with the next streaming update.
 
         Discards the last sampled output token from the prior input chunk.
         """
-
         # Current streaming input behaviour: Keep only computed output tokens
         # (discard final sampled output token).
         num_computed_tokens = session.num_computed_tokens
@@ -1662,6 +1715,29 @@ def _make_cached_request_data(
             num_output_tokens=num_output_tokens,
         )
 
+    def _reject_on_encoder_cache_embed_mismatch(
+        self,
+        request: Request,
+        input_id: int,
+        identifier: str,
+        cached_num_encoder_embeds: int,
+        requested_num_encoder_embeds: int,
+    ) -> bool:
+        if cached_num_encoder_embeds == requested_num_encoder_embeds:
+            return False
+        logger.warning(
+            "Encoder cache entry for multimodal identifier %r holds %d "
+            "embeddings, but request %s input %d expects %d. A client-provided "
+            "multimodal UUID may have been reused for different content.",
+            identifier,
+            cached_num_encoder_embeds,
+            request.request_id,
+            input_id,
+            requested_num_encoder_embeds,
+        )
+        self.encoder_cache_mismatch_reqs.add(request.request_id)
+        return True
+
     def _try_schedule_encoder_inputs(
         self,
         request: Request,
@@ -1670,8 +1746,7 @@ def _try_schedule_encoder_inputs(
         encoder_compute_budget: int,
         shift_computed_tokens: int = 0,
     ) -> tuple[list[int], int, int, list[int]]:
-        """
-        Determine which encoder inputs need to be scheduled in the current step,
+        """Determine which encoder inputs need to be scheduled in the current step,
         and update `num_new_tokens` and encoder token budget accordingly.
 
         An encoder input will be scheduled if:
@@ -1701,7 +1776,7 @@ def _try_schedule_encoder_inputs(
         # NOTE: since scheduler operates on the request level (possibly with
         # multiple encoder inputs per request), we need to create temporary
         # trackers for accounting at the encoder input level.
-        mm_hashes_to_schedule = set()
+        mm_hashes_to_schedule: dict[str, int] = {}
         num_embeds_to_schedule = 0
 
         encoder_window_end = (
@@ -1743,10 +1818,35 @@ def _try_schedule_encoder_inputs(
                 # We are not using the encoder cache for encoder-decoder models,
                 # yet.
                 if item_identifier in mm_hashes_to_schedule:
+                    scheduled_num_encoder_embeds = mm_hashes_to_schedule[
+                        item_identifier
+                    ]
+                    if self._reject_on_encoder_cache_embed_mismatch(
+                        request,
+                        i,
+                        item_identifier,
+                        scheduled_num_encoder_embeds,
+                        num_encoder_embeds,
+                    ):
+                        return [], 0, encoder_compute_budget, []
                     # The same encoder input has already been scheduled in the
                     # current step.
                     continue
 
+                cached_num_encoder_embeds = (
+                    self.encoder_cache_manager.get_cached_num_encoder_embeds(request, i)
+                )
+                if cached_num_encoder_embeds is not None and (
+                    self._reject_on_encoder_cache_embed_mismatch(
+                        request,
+                        i,
+                        item_identifier,
+                        cached_num_encoder_embeds,
+                        num_encoder_embeds,
+                    )
+                ):
+                    return [], 0, encoder_compute_budget, []
+
                 if self.encoder_cache_manager.check_and_update_cache(request, i):
                     # The encoder input is already computed and cached from a
                     # previous step.
@@ -1806,14 +1906,14 @@ def _try_schedule_encoder_inputs(
             if self.ec_connector is not None and self.ec_connector.has_cache_item(
                 item_identifier
             ):
-                mm_hashes_to_schedule.add(item_identifier)
+                mm_hashes_to_schedule[item_identifier] = num_encoder_embeds
                 external_load_encoder_input.append(i)
                 num_embeds_to_schedule += num_encoder_embeds
                 continue
 
             num_embeds_to_schedule += num_encoder_embeds
             encoder_compute_budget -= num_encoder_embeds
-            mm_hashes_to_schedule.add(item_identifier)
+            mm_hashes_to_schedule[item_identifier] = num_encoder_embeds
             encoder_inputs_to_schedule.append(i)
 
         return (
@@ -1892,37 +1992,21 @@ def update_from_output(
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
 
-        failed_kv_load_req_ids = None
+        failed_kv_load_req_ids: set[str] = set()
+        if kv_connector_output and kv_connector_output.failed_recving:
+            failed_kv_load_req_ids.update(
+                self._handle_failed_recving(kv_connector_output.failed_recving)
+            )
         if kv_connector_output and kv_connector_output.invalid_block_ids:
             # These blocks contain externally computed tokens that failed to
             # load. Identify affected requests and adjust their computed token
             # count to trigger recomputation of the invalid blocks.
-            failed_kv_load_req_ids = self._handle_invalid_blocks(
-                kv_connector_output.invalid_block_ids,
-                num_scheduled_tokens,
-            )
-
-        # Persist per-step routed experts into the scheduler-side slot
-        # buffer (CPU->CPU fancy-index assign; ~few MB per step).
-        # MUST precede the per-request routing reads below: stopped
-        # requests may terminate on tokens generated in this very step,
-        # whose routing was just D2H'd into model_runner_output.
-        routing_data = None
-        routing_offsets: dict[str, int] = {}
-        if model_runner_output.routed_experts is not None:
-            re = model_runner_output.routed_experts
-            self.routed_experts_mgr.store_batch(re.routing_data, re.slot_mapping)
-            routing_data = re.routing_data.astype(
-                self.routed_experts_mgr.routed_experts_by_slot.dtype,
-                copy=False,
+            failed_kv_load_req_ids.update(
+                self._handle_invalid_blocks(
+                    kv_connector_output.invalid_block_ids,
+                    num_scheduled_tokens,
+                )
             )
-            # Build offset map using model runner's request order
-            # (input_batch ordering), NOT scheduler dict order.
-            offset = 0
-            for rid in model_runner_output.req_ids:
-                routing_offsets[rid] = offset
-                offset += num_scheduled_tokens[rid]
-
         # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
         # the below loop can be a performance bottleneck. We should do our best
         # to avoid expensive operations inside the loop.
@@ -2015,7 +2099,6 @@ def update_from_output(
             ec_transfer_params = None
             prefill_stats = None
             status_before_stop = request.status
-            num_output_tokens_before = len(request._output_token_ids)
 
             # Check for stop and update request status.
             if new_token_ids:
@@ -2038,77 +2121,27 @@ def update_from_output(
                 request.status = RequestStatus.FINISHED_STOPPED
                 stopped = True
 
-            if new_token_ids and self.structured_output_manager.should_advance(
-                request, new_token_ids=new_token_ids
+            if new_token_ids and not self.structured_output_manager.accept_tokens(
+                request, new_token_ids
             ):
-                struct_output_request = request.structured_output_request
-                assert struct_output_request is not None
-                grammar = struct_output_request.grammar
-                assert isinstance(grammar, StructuredOutputGrammar)
-                # new_token_ids can be a mixed block of reasoning content, then
-                # the reasoning end marker, then the start of the grammar content.
-                # Trim the reasoning content so the grammar only sees grammar content.
-                advance_token_ids = (
-                    self.structured_output_manager.trim_reasoning_for_advance(
-                        request, new_token_ids
-                    )
+                logger.error(
+                    "Unexpected: grammar rejected tokens %s for request %s. "
+                    "Terminating request.",
+                    new_token_ids,
+                    req_id,
                 )
-                if advance_token_ids and not grammar.accept_tokens(
-                    req_id, advance_token_ids
-                ):
-                    logger.error(
-                        "Unexpected: grammar rejected tokens %s for request %s. "
-                        "Terminating request.",
-                        advance_token_ids,
-                        req_id,
-                    )
-                    request.status = RequestStatus.FINISHED_ERROR
-                    request.resumable = False
-                    stopped = True
+                request.status = RequestStatus.FINISHED_ERROR
+                request.resumable = False
+                stopped = True
 
             routed_experts = None
-            if (
-                self.enable_return_routed_experts
-                and routing_data is not None
-                and new_token_ids
-            ):
-                req_offset = routing_offsets[req_id]
-                end = req_offset + num_tokens_scheduled
-                block_ids = self._re_block_ids.pop(req_id, [])
-                if num_output_tokens_before == 0:
-                    # Prefill completed: read full prompt routing from
-                    # slot buffer using the block-ID snapshot taken at
-                    # schedule time (immune to async preemption).
-                    if (
-                        request.sampling_params is not None
-                        and request.sampling_params.routed_experts_prompt_start
-                        is not None
-                    ):
-                        prompt_start = (
-                            request.sampling_params.routed_experts_prompt_start
-                        )
-                        assert prompt_start < request.num_prompt_tokens
-                    else:
-                        prompt_start = 0
-                    routed_experts = self.routed_experts_mgr.get(
-                        block_ids,
-                        request.num_prompt_tokens,
-                        token_start=prompt_start,
-                    )
-                else:
-                    if scheduled_spec_token_ids:
-                        # Spec decode: accepted tokens at the START of
-                        # the scheduled range, rejected at the end.
-                        routed_experts = routing_data[
-                            req_offset : req_offset + len(new_token_ids)
-                        ]
-                    else:
-                        # Normal decode / re-prefill: token(s) at the END.
-                        routed_experts = routing_data[end - len(new_token_ids) : end]
-
             should_emit_output = bool(
                 new_token_ids or pooler_output is not None or stopped
             )
+            if self.aux_output_connector is not None and should_emit_output:
+                routed_experts = self.aux_output_connector.take_output(
+                    request, model_runner_output.aux_output_connector_output
+                )
             if should_emit_output:
                 prefill_stats = request.take_prefill_stats()
                 if prefill_stats is not None:
@@ -2124,6 +2157,17 @@ def update_from_output(
                 finished = self._handle_stopped_request(request)
                 if finished:
                     kv_transfer_params, ec_transfer_params = self._free_request(request)
+                    if (
+                        prefill_stats is not None
+                        and kv_transfer_params is not None
+                        and kv_transfer_params.get("do_remote_prefill")
+                    ):
+                        # P-side cache hits, so D can report them in
+                        # prompt_tokens_details instead of its own (~100%)
+                        # hit rate from the KV transfer.
+                        kv_transfer_params["remote_prefill_cached_tokens"] = (
+                            prefill_stats.num_cached_tokens
+                        )
 
                 if status_before_stop == RequestStatus.RUNNING:
                     stopped_running_reqs.add(request)
@@ -2190,6 +2234,8 @@ def update_from_output(
 
         error_req_ids = set(self.grammar_compile_error_reqs)
         self.grammar_compile_error_reqs.clear()
+        error_req_ids.update(self.encoder_cache_mismatch_reqs)
+        self.encoder_cache_mismatch_reqs.clear()
         if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
             error_req_ids.update(failed_kv_load_req_ids)
         if self.ec_connector is not None:
@@ -2237,6 +2283,23 @@ def update_from_output(
                     else scheduler_kv_connector_stats
                 )
 
+        # Worker-side EC connector stats from the model runner output.
+        ec_connector_stats: ECConnectorStats | None = (
+            ec_connector_output.ec_connector_stats if ec_connector_output else None
+        )
+        if self.ec_connector:
+            # Scheduler-side EC connector stats collected after connector update.
+            scheduler_ec_connector_stats = self.ec_connector.get_ec_connector_stats()
+            if (
+                scheduler_ec_connector_stats is not None
+                and not scheduler_ec_connector_stats.is_empty()
+            ):
+                ec_connector_stats = (
+                    ec_connector_stats.aggregate(scheduler_ec_connector_stats)
+                    if ec_connector_stats is not None
+                    else scheduler_ec_connector_stats
+                )
+
         # collect KV cache events from KV cache manager
         events = self.kv_cache_manager.take_events()
 
@@ -2281,6 +2344,7 @@ def update_from_output(
                 kv_connector_stats,
                 cudagraph_stats,
                 perf_stats,
+                ec_connector_stats=ec_connector_stats,
             )
         ) is not None:
             # Return stats to only one of the front-ends.
@@ -2420,11 +2484,9 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
                 continue
 
             # Add newly generated spec token ids to the request.
-            if self.structured_output_manager.should_advance(request):
-                metadata = request.structured_output_request
-                spec_token_ids = strip_speculative_padding(spec_token_ids)
-                spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids)  # type: ignore[union-attr]
-            request.spec_token_ids = spec_token_ids
+            request.spec_token_ids = self.structured_output_manager.validate_tokens(
+                request, spec_token_ids
+            )
 
     def update_draft_token_ids_in_output(
         self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput
@@ -2450,10 +2512,9 @@ def update_draft_token_ids_in_output(
             # (needed for chunked prefill case for example).
             del spec_token_ids[orig_num_spec_tokens:]
             # Filter out spec tokens which do not adhere to the grammar.
-            if self.structured_output_manager.should_advance(request):
-                metadata = request.structured_output_request
-                spec_token_ids = strip_speculative_padding(spec_token_ids)
-                spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids)  # type: ignore[union-attr]
+            spec_token_ids = self.structured_output_manager.validate_tokens(
+                request, spec_token_ids
+            )
             # Pad to original number of spec tokens.
             num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)
             if num_invalid_tokens:
@@ -2513,6 +2574,7 @@ def finish_requests(
         Returns:
             List of requests that were aborted. Will not include any that were
             already finished.
+
         """
         assert RequestStatus.is_finished(finished_status)
         if isinstance(request_ids, str):
@@ -2568,6 +2630,8 @@ def _free_request(
     ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
         assert request.is_finished()
 
+        if self.aux_output_connector is not None:
+            self.aux_output_connector.request_finished(request)
         self._inflight_prefills.discard(request)
         connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
 
@@ -2605,15 +2669,18 @@ def set_pause_state(self, pause_state: PauseState) -> None:
         logger.info("setting pause state to %s", pause_state.name)
         self._pause_state = pause_state
 
+    def _request_blocks_can_be_freed(self, request: Request) -> bool:
+        # We must defer freeing blocks if an async kv connector may
+        # write to them immediately (not ordered with GPU stream).
+        return not self.defer_block_free or (
+            request.last_sched_seq <= self.processed_step_seq
+        )
+
     def _free_request_blocks(self, request: Request):
         """Free the request's KV blocks, deferring the return to the block
         pool when an in-flight GPU step may still write them.
         """
-        if not self.defer_block_free or (
-            # Last scheduled step already processed: no in-flight write remains
-            # (always the case for a normal finish), so free now.
-            request.last_sched_seq <= self.processed_step_seq
-        ):
+        if self._request_blocks_can_be_freed(request):
             self.kv_cache_manager.free(request)
             return
         blocks = self.kv_cache_manager.pop_blocks_for_free(request)
@@ -2697,6 +2764,16 @@ def reset_prefix_cache(
         Otherwise, this method will only reset the KV prefix cache when there
         is no running requests taking KV cache.
         """
+        if reset_running_requests and self.aux_output_connector is not None:
+            if self._pause_state != PauseState.PAUSED_ALL:
+                raise RuntimeError(
+                    "AuxOutput Connector only supports resetting running requests "
+                    "after pause(mode='keep')."
+                )
+            if any(request.num_in_flight_tokens for request in self.requests.values()):
+                raise RuntimeError(
+                    "AuxOutput Connector cannot reset while model output is in flight."
+                )
         if reset_running_requests:
             # For logging.
             timestamp = time.monotonic()
@@ -2727,6 +2804,9 @@ def reset_prefix_cache(
         if reset_connector:
             reset_successful = self.reset_connector_cache() and reset_successful
 
+        if reset_successful and self.aux_output_connector is not None:
+            self.aux_output_connector.reset()
+
         return reset_successful
 
     def reset_connector_cache(self) -> bool:
@@ -2765,6 +2845,7 @@ def make_stats(
         kv_connector_stats: KVConnectorStats | None = None,
         cudagraph_stats: CUDAGraphStat | None = None,
         perf_stats: PerfStats | None = None,
+        ec_connector_stats: ECConnectorStats | None = None,
     ) -> SchedulerStats | None:
         if not self.log_stats:
             return None
@@ -2783,6 +2864,9 @@ def make_stats(
         connector_stats_payload = (
             kv_connector_stats.to_dict() if kv_connector_stats else None
         )
+        ec_connector_stats_payload = (
+            ec_connector_stats.data if ec_connector_stats else None
+        )
         return SchedulerStats(
             num_running_reqs=len(self.running),
             num_waiting_reqs=len(self.waiting),
@@ -2795,6 +2879,7 @@ def make_stats(
             kv_connector_stats=connector_stats_payload,
             cudagraph_stats=cudagraph_stats,
             perf_stats=perf_stats,
+            ec_connector_stats=ec_connector_stats_payload,
         )
 
     def make_spec_decoding_stats(
@@ -2844,8 +2929,7 @@ def get_kv_event_publisher_config(self) -> KVEventsConfig | None:
     def _connector_finished(
         self, request: Request
     ) -> tuple[bool, dict[str, Any] | None]:
-        """
-        Invoke the KV connector request_finished() method if applicable.
+        """Invoke the KV connector request_finished() method if applicable.
 
         Returns optional kv transfer parameters to be included with the
         request outputs.
@@ -2898,9 +2982,9 @@ def _connector_finished(
         return delay_free or partial_tail_delay, kv_xfer_params
 
     def _request_remaining_blocks(self, request: Request) -> int:
-        """Blocks `request` still needs to allocate to hold its full sequence."""
+        """Blocks `request` still needs to hold its full sequence and be promoted."""
         full_num_tokens = min(request.num_tokens, self.max_model_len)
-        return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate(
+        num_blocks = self.kv_cache_manager.coordinator.get_num_blocks_to_allocate(
             request_id=request.request_id,
             num_tokens=full_num_tokens,
             new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks,
@@ -2910,17 +2994,43 @@ def _request_remaining_blocks(self, request: Request) -> int:
             num_tokens_main_model=full_num_tokens,
             apply_admission_cap=True,
         )
+        return num_blocks + self._spec_decode_step_blocks()
+
+    def _spec_decode_step_blocks(self) -> int:
+        """Number of blocks for the extra KV slots a spec decode step needs.
+
+        When using async kv load, scheduler must reserve enough blocks for
+        full sequence + the spec decode step, otherwise request cannot be
+        able to run after the async_load if we are out of kv blocks.
+        """
+        if not self.num_spec_tokens:
+            return 0
+        return cdiv(
+            1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size
+        )
 
     def _inflight_prefill_reserved_blocks(self) -> int:
         """Num blocks in-flight prefills still need to finish (their reservation)."""
-
         return sum(
             self._request_remaining_blocks(req) for req in self._inflight_prefills
         )
 
-    def _update_waiting_for_remote_kv(self, request: Request) -> None:
+    def _mark_prefix_replay(self, request: Request, num_hit_tokens: int) -> int:
+        """Record where a prefix hit's replay starts and return the number of
+        hit tokens to recompute; the caller rewinds the computed count by it.
+
+        The replayed tokens are the hit's last window (see
+        ``Request.replay_start``): the worker rebuilds their sliding-window KV
+        and leaves their cached KV alone.
         """
-        KV Connector: update request state after async recv is finished.
+        if not self.prefix_replay_tokens:
+            return 0
+        num_replay_tokens = min(self.prefix_replay_tokens, num_hit_tokens)
+        request.replay_start = num_hit_tokens - num_replay_tokens
+        return num_replay_tokens
+
+    def _update_waiting_for_remote_kv(self, request: Request) -> None:
+        """KV Connector: update request state after async recv is finished.
 
         When the kv transfer is ready, we cache the blocks
         and the request state will be moved back to WAITING from
@@ -2954,17 +3064,20 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None:
             # This will cache the blocks iff caching is enabled.
             self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
 
-            # on a full prompt hit, we need to re-compute the last token
-            # in order to be able to sample the next token
-            if request.num_computed_tokens == request.num_tokens:
-                request.num_computed_tokens = request.num_tokens - 1
-
+        # SWA bounded replay recomputes the tail of the hit, which covers the
+        # last token; otherwise a full prompt hit re-computes that token so the
+        # next one can be sampled.
+        num_replay_tokens = self._mark_prefix_replay(
+            request, request.num_computed_tokens
+        )
+        if num_replay_tokens > 0:
+            request.num_computed_tokens -= num_replay_tokens
+        elif request.num_computed_tokens == request.num_tokens:
+            request.num_computed_tokens = request.num_tokens - 1
         self.finished_recving_kv_req_ids.remove(request.request_id)
 
     def _try_promote_blocked_waiting_request(self, request: Request) -> bool:
-        """
-        Try to promote a blocked waiting request back to schedulable states.
-        """
+        """Try to promote a blocked waiting request back to schedulable states."""
         if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
             # finished_recving_kv_req_ids is populated during
             # update_from_output(), based on worker-side connector signals
@@ -2998,8 +3111,7 @@ def _try_promote_blocked_waiting_request(self, request: Request) -> bool:
         )
 
     def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
-        """
-        KV Connector: update the scheduler state based on the output.
+        """KV Connector: update the scheduler state based on the output.
 
         The Worker side connectors add finished_recving and
         finished_sending reqs to the output.
@@ -3007,7 +3119,6 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
         # if finished_recving: add to state so we can
             schedule the request during the next step.
         """
-
         if self.connector is not None:
             self.connector.update_connector_output(kv_connector_output)
 
@@ -3033,8 +3144,7 @@ def _update_requests_with_invalid_blocks(
         num_scheduled_tokens: dict[str, int],
         evict_blocks: bool = True,
     ) -> tuple[set[str], int, set[int]]:
-        """
-        Identify and update requests affected by invalid KV cache blocks.
+        """Identify and update requests affected by invalid KV cache blocks.
 
         This method scans the given requests, detects those with invalid blocks
         and adjusts their `num_computed_tokens` to the longest valid prefix.
@@ -3056,6 +3166,7 @@ def _update_requests_with_invalid_blocks(
                 be recomputed across all affected requests.
                 - blocks_to_evict (set[int]): Block IDs to evict from cache,
                 including invalid blocks and downstream dependent blocks.
+
         """
         affected_req_ids: set[str] = set()
         total_affected_tokens = 0
@@ -3068,7 +3179,8 @@ def _update_requests_with_invalid_blocks(
             is_affected = False
             marked_invalid_block = False
             req_id = request.request_id
-            # TODO (davidb): add support for hybrid memory allocator
+            # One block-id list per KV cache group; single-group layouts are
+            # enforced by _handle_invalid_blocks.
             (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id)
             # We iterate only over blocks that may contain externally computed
             # tokens
@@ -3129,15 +3241,41 @@ def _update_requests_with_invalid_blocks(
 
         return affected_req_ids, total_affected_tokens, blocks_to_evict
 
+    def _handle_failed_recving(self, failed_req_ids: set[str]) -> set[str]:
+        """Fail closed for layouts whose block IDs are not globally unique."""
+        affected_req_ids: set[str] = set()
+        for req_id in failed_req_ids:
+            request = self.requests.get(req_id)
+            if (
+                request is not None
+                and request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+            ):
+                affected_req_ids.add(req_id)
+                if self.recompute_kv_load_failures:
+                    request.num_computed_tokens = 0
+                    self.failed_recving_kv_req_ids.add(req_id)
+        return set() if self.recompute_kv_load_failures else affected_req_ids
+
     def _handle_invalid_blocks(
         self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int]
     ) -> set[str]:
-        """
-        Handle requests affected by invalid KV cache blocks.
+        """Handle requests affected by invalid KV cache blocks.
 
         Returns:
             Set of affected request IDs to skip in update_from_output main loop.
+
         """
+        if len(self.kv_cache_config.kv_cache_groups) > 1:
+            # Block IDs are only unique within a group, so a flat set of
+            # invalid block IDs cannot be mapped back to requests.
+            raise RuntimeError(
+                "A KV connector reported block-level load failures "
+                "(invalid_block_ids) which is not supported for models "
+                "with multiple KV cache groups. "
+                "Connectors must report failed requests via "
+                "KVConnectorTransferResults.failed_recving instead."
+            )
+
         should_fail = not self.recompute_kv_load_failures
 
         # handle async KV loads (not cached yet, evict_blocks=False)
diff --git a/vllm/v1/core/sched/utils.py b/vllm/v1/core/sched/utils.py
index 31fa840e7783..045dacaa9829 100644
--- a/vllm/v1/core/sched/utils.py
+++ b/vllm/v1/core/sched/utils.py
@@ -30,11 +30,14 @@ def check_sequence_repetition(
     params: RepetitionDetectionParams,
 ) -> bool:
     """Check if a sequence of token IDs has a repetition pattern.
+
     Args:
         token_ids: List of token IDs
         params: Repetition detection parameters.
+
     Returns:
         True if a repetition pattern is found, False otherwise.
+
     """
     max_pattern_size = params.max_pattern_size
     min_pattern_size = params.min_pattern_size
@@ -77,6 +80,7 @@ def remove_all(lst: list, items_to_remove: set) -> list:
     Note:
         For single item removal, this modifies the original list in-place
         and returns it. For multiple items, it creates and returns a new list.
+
     """
     if not items_to_remove:
         return lst
diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py
index 15b35d9f7b57..70d56c6543bd 100644
--- a/vllm/v1/core/single_type_kv_cache_manager.py
+++ b/vllm/v1/core/single_type_kv_cache_manager.py
@@ -4,8 +4,9 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Sequence
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
+from vllm.distributed.kv_events import MEDIUM_CPU
 from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.block_pool import BlockPool
@@ -16,6 +17,7 @@
     KVCacheBlock,
     resolve_block_hashes,
 )
+from vllm.v1.hisparse.block_pool import SharedEventQueueBlockPool
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     ChunkedLocalAttentionSpec,
@@ -23,7 +25,10 @@
     CrossAttentionSpec,
     FullAttentionSpec,
     HiddenStateCacheSpec,
+    HiSparseHotSpec,
+    HiSparseResidentSpec,
     KpoolTailSpec,
+    KVCacheGroupRole,
     KVCacheSpec,
     MambaSpec,
     MLAAttentionSpec,
@@ -37,17 +42,32 @@
 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.request import Request
 
+if TYPE_CHECKING:
+    from vllm.v1.hisparse.coordinator import HiSparseCoordinator
+
 logger = init_logger(__name__)
 
 
 class SingleTypeKVCacheManager(ABC):
-    """
-    An abstract base class for a manager that handle the kv cache management
+    """An abstract base class for a manager that handle the kv cache management
     logic of one specific type of attention layer.
     """
 
     supports_fine_grained_hash_lookup: ClassVar[bool] = False
 
+    # Keep this group's longer prefix until external cache lookup completes.
+    retains_longer_hit: bool = False
+
+    @property
+    def has_positionally_stable_blocks(self) -> bool:
+        """Whether positional offload scans can follow this block table.
+
+        True means already-scanned indices will not be reused for a different
+        token range. Managers may still null blocks that left their retention
+        window, as long as the cursor's positional history remains valid.
+        """
+        return True
+
     def __init__(
         self,
         kv_cache_spec: KVCacheSpec,
@@ -60,14 +80,17 @@ def __init__(
         needs_kv_cache_zeroing: bool = False,
         max_admission_blocks_per_request: int | None = None,
     ) -> None:
-        """
-        Initializes the SingleTypeKVCacheManager.
+        """Initializes the SingleTypeKVCacheManager.
+
         Args:
             kv_cache_spec: The kv_cache_spec for this manager.
             block_pool: The block pool.
+            enable_caching: Whether prefix caching is enabled.
             kv_cache_group_id: The id of the kv cache group of this manager.
             scheduler_block_size: The scheduling granularity (LCM of all group
                 block sizes); a multiple of this manager's ``block_size``.
+            dcp_world_size: Decode context parallel world size.
+            pcp_world_size: Prefill context parallel world size.
             needs_kv_cache_zeroing: Whether worker-side KV cache zeroing needs
                 newly allocated block IDs from this manager.
             max_admission_blocks_per_request: Recycling-aware per-request
@@ -76,6 +99,7 @@ def __init__(
                 chunked-local); `None` (the default) means no cap, which is
                 correct for full-attention-style specs that hold every
                 block until the request finishes.
+
         """
         self.scheduler_block_size = scheduler_block_size
         # Hybrid fine-grained lookup may lower this after all participating
@@ -119,9 +143,9 @@ def __init__(
         # aligned segment (SWA). Initialized lazily by the coordinator after
         # determining the attention groups.
         self.use_eagle = False
-        # ``CacheConfig.enable_mamba_fine_grained_prefix_cache``, narrowed and set
+        # ``CacheConfig.enable_mamba_shared_prefix_checkpoint``, narrowed and set
         # by ``KVCacheManager``; only an EAGLE Mamba "align" group ever gets it.
-        self.fine_grained_prefix_cache = False
+        self.shared_prefix_checkpoint = False
         # Partial-hit copy-on-write bookkeeping. Populated only by fine-grained
         # managers (full attention, mamba "align"); harmlessly empty elsewhere.
         self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {}
@@ -163,8 +187,7 @@ def get_num_blocks_to_allocate(
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
     ) -> int:
-        """
-        Get the number of blocks needed to be allocated for the request.
+        """Get the number of blocks needed to be allocated for the request.
 
         Args:
             request_id: The request ID.
@@ -185,8 +208,8 @@ def get_num_blocks_to_allocate(
 
         Returns:
             The number of blocks to allocate.
-        """
 
+        """
         num_required_blocks = cdiv(num_tokens, self.block_size)
         if apply_admission_cap and self._max_admission_blocks_per_request is not None:
             # Recycling-aware specs (SWA, chunked-local) cap the per-request
@@ -248,8 +271,7 @@ def add_local_computed_blocks(
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
     ) -> None:
-        """
-        Add the locally cached (prefix-hit) blocks to the request:
+        """Add the locally cached (prefix-hit) blocks to the request:
         1. Touch the computed blocks (paired with adding them to `req_blocks`)
            so their ref_cnt exactly tracks the referencing requests.
         1.5. (Optional) For sliding window, skipped blocks are padded with nulls.
@@ -261,6 +283,7 @@ def add_local_computed_blocks(
                 prefix cache.
             num_local_computed_tokens: The number of local computed tokens.
             num_external_computed_tokens: The number of external computed tokens.
+
         """
         # The coordinator only calls this for first-time allocations (running
         # requests are short-circuited there), so the request has no blocks yet.
@@ -306,8 +329,7 @@ def allocate_external_computed_blocks(
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
     ) -> None:
-        """
-        Allocate new blocks for external (KV-connector) computed tokens.
+        """Allocate new blocks for external (KV-connector) computed tokens.
 
         Must run only after every group's local blocks have been touched via
         `add_local_computed_blocks`, so this group's `get_new_blocks` cannot
@@ -317,6 +339,7 @@ def allocate_external_computed_blocks(
             request_id: The request ID.
             num_local_computed_tokens: The number of local computed tokens.
             num_external_computed_tokens: The number of external computed tokens.
+
         """
         num_total_computed_tokens = (
             num_local_computed_tokens + num_external_computed_tokens
@@ -343,8 +366,7 @@ def allocate_external_computed_blocks(
     def allocate_new_blocks(
         self, request_id: str, num_tokens: int, num_tokens_main_model: int
     ) -> list[KVCacheBlock]:
-        """
-        Allocate new blocks for the request to give it at least `num_tokens`
+        """Allocate new blocks for the request to give it at least `num_tokens`
         token slots.
 
         Args:
@@ -354,8 +376,10 @@ def allocate_new_blocks(
             num_tokens_main_model: The number of tokens for the main model (aka target
                 model in spec decode). w/o spec decode, it is num_tokens;
                 with spec decode, it is num_tokens - num_lookahead_tokens.
+
         Returns:
             The new allocated blocks.
+
         """
         cow_blocks: list[KVCacheBlock] = []
         if request_id in self._partial_hit_reqs:
@@ -454,8 +478,7 @@ def cache_blocks(
         *,
         replay_boundaries: Sequence[int],
     ) -> None:
-        """
-        Cache the blocks for the request.
+        """Cache the blocks for the request.
 
         Args:
             request: The request.
@@ -467,7 +490,10 @@ def cache_blocks(
                 a tail once per that-sized segment. Only SWA acts on it.
             replay_boundaries: Positions a later request replaying this prompt
                 can resume at, from ``get_replay_boundaries``.
+
         """
+        if not self.kv_cache_spec.prefix_cacheable:
+            return
         num_cached_blocks = self.num_cached_block.get(request.request_id, 0)
         num_full_blocks = num_tokens // self.block_size
 
@@ -531,8 +557,7 @@ def reachable_block_mask(
         return None
 
     def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
-        """
-        Pop the request's bookkeeping and return its blocks without yet
+        """Pop the request's bookkeeping and return its blocks without yet
         returning them to the block pool. The caller is responsible for
         eventually passing the returned blocks to `block_pool.free_blocks`,
         freeing them in reverse order (so that tail blocks are evicted first).
@@ -542,6 +567,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
 
         Returns:
             The request's blocks in allocation order.
+
         """
         # Default to [] in case a request is freed (aborted) before alloc.
         req_blocks = self.req_to_blocks.pop(request_id, [])
@@ -550,19 +576,18 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
         return req_blocks
 
     def free(self, request_id: str) -> None:
-        """
-        Free the blocks for the request.
+        """Free the blocks for the request.
 
         Args:
             request_id: The request ID.
+
         """
         # Free blocks in reverse order so that the tail blocks are freed first.
         self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id)))
 
     @abstractmethod
     def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
-        """
-        Get the number of common prefix blocks for all requests with allocated
+        """Get the number of common prefix blocks for all requests with allocated
         KV cache.
 
         Args:
@@ -571,8 +596,8 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
         Returns:
             The number of common prefix blocks for all requests with allocated
             KV cache.
-        """
 
+        """
         raise NotImplementedError
 
     @classmethod
@@ -589,8 +614,7 @@ def find_longest_cache_hit(
         dcp_world_size: int = 1,
         pcp_world_size: int = 1,
     ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
-        """
-        Get the longest cache hit prefix of the blocks that is not longer than
+        """Get the longest cache hit prefix of the blocks that is not longer than
         `max_length`. The prefix should be a common prefix hit for all the
         kv cache groups in `kv_cache_group_ids`. If no cache hit is found,
         return an empty list.
@@ -626,8 +650,8 @@ def find_longest_cache_hit(
             For example, sliding window manager should return a list like
             ([NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)]) for block size 4
             and sliding window 8 and len(kv_cache_group_ids) = 1.
-        """
 
+        """
         raise NotImplementedError
 
     def _remove_blocks_in_range(
@@ -663,8 +687,8 @@ def remove_skipped_blocks(
         processed_computed_tokens: int,
         num_prompt_tokens: int | None = None,
     ) -> None:
-        """
-        Remove and free the blocks that are no longer needed for attention computation.
+        """Remove and free blocks no longer needed for attention computation.
+
         The removed blocks should be replaced by null_block.
 
         This function depends on `get_num_skipped_tokens`, which need to be implemented
@@ -677,6 +701,7 @@ def remove_skipped_blocks(
             num_prompt_tokens: Optional prompt length for attention types (e.g.
                 R-SWA) that evict a middle gap rather than a head prefix. Ignored
                 by the default implementation.
+
         """
         del num_prompt_tokens
         # Remove the blocks that will be skipped during attention computation.
@@ -697,14 +722,14 @@ def remove_skipped_blocks(
         self._remove_blocks_in_range(request_id, 0, num_skipped_blocks)
 
     def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
-        """
-        Get the number of tokens that will be skipped for attention computation.
+        """Get the number of tokens that will be skipped for attention computation.
 
         Args:
             num_computed_tokens: The number of tokens that have been computed.
 
         Returns:
             The number of tokens that will be skipped for attention computation.
+
         """
         # The default behavior is to not skip any tokens.
         return 0
@@ -1123,8 +1148,7 @@ def reachable_block_mask(
         return mask
 
     def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
-        """
-        Get the number of tokens that will be skipped for attention computation.
+        """Get the number of tokens that will be skipped for attention computation.
 
         For sliding window, this corresponds to the tokens that are prior to
         the current sliding window.
@@ -1153,6 +1177,7 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
 
         Returns:
             The number of tokens that will be skipped for attention computation.
+
         """
         return max(
             0,
@@ -1160,8 +1185,7 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
         )
 
     def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
-        """
-        NOTE(Chen): The prefix blocks are null blocks for sliding window layers.
+        """NOTE(Chen): The prefix blocks are null blocks for sliding window layers.
         So it's not correct to count ref_cnt like FullAttentionManager. Return
         0 here for correctness. Need to support cascade attention + sliding
         window in the future.
@@ -1280,8 +1304,7 @@ def find_longest_cache_hit(
         dcp_world_size: int = 1,
         pcp_world_size: int = 1,
     ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
-        """
-        For chunked local attention, we need to find the longest cache hit
+        """For chunked local attention, we need to find the longest cache hit
         prefix of the blocks that is not longer than `max_length`. The prefix
         should be a common prefix hit for all the kv cache groups in
         `kv_cache_group_ids`. If no cache hit is found, return an empty list.
@@ -1314,6 +1337,7 @@ def find_longest_cache_hit(
 
         Returns:
             A list of cached blocks
+
         """
         assert isinstance(kv_cache_spec, ChunkedLocalAttentionSpec), (
             "ChunkedLocalAttentionManager can only be used for "
@@ -1368,8 +1392,7 @@ def find_longest_cache_hit(
         return computed_blocks, hit_length
 
     def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
-        """
-        Get the number of tokens that will be skipped for attention computation.
+        """Get the number of tokens that will be skipped for attention computation.
 
         For chunked local attention, this corresponds to the tokens that are on
         the left side of the current chunk.
@@ -1407,6 +1430,7 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
 
         Returns:
             The number of tokens that will be skipped for attention computation.
+
         """
         num_skipped_tokens = (
             num_computed_tokens // self.attention_chunk_size
@@ -1414,15 +1438,19 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
         return num_skipped_tokens
 
     def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
-        """
-        cascade attention is not supported by chunked local attention.
-        """
+        """Cascade attention is not supported by chunked local attention."""
         return 0
 
 
 class MambaManager(SingleTypeKVCacheManager):
     supports_fine_grained_hash_lookup: ClassVar[bool] = True
 
+    @property
+    def has_positionally_stable_blocks(self) -> bool:
+        # Align-mode Mamba can null interior states and relocate speculative
+        # blocks in place. Other modes retain positional identity.
+        return self.mamba_cache_mode != "align"
+
     def __init__(
         self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs
     ) -> None:
@@ -1444,6 +1472,7 @@ def __init__(
             # Mapping from request ID to the index of the block
             # allocated in the previous step
             self.last_state_block_idx: dict[str, int] = {}
+            self._num_retired_blocks: dict[str, int] = {}
             # The set of the requests that have been allocated blocks
             self._allocated_block_reqs: set[str] = set()
             # checkpoint position and reserved block index for the current
@@ -1594,6 +1623,27 @@ def reachable_block_mask(
 
         return mask
 
+    def _remove_blocks_in_range(
+        self, request_id: str, first_block: int, last_block: int
+    ) -> None:
+        if self.mamba_cache_mode != "align":
+            return super()._remove_blocks_in_range(request_id, first_block, last_block)
+        blocks = self.req_to_blocks.get(request_id, [])
+        first_block = max(first_block, self._num_retired_blocks.get(request_id, 0))
+        last_block = min(last_block, len(blocks))
+        if first_block >= last_block:
+            return
+        freed: list[KVCacheBlock] = []
+        # Mamba prefill leaves null gaps between states awaiting retirement.
+        for i in range(last_block - 1, first_block - 1, -1):
+            if blocks[i].is_null:
+                continue
+            freed.append(blocks[i])
+            blocks[i] = self._null_block
+        if freed:
+            self.block_pool.free_blocks(freed)
+        self._num_retired_blocks[request_id] = last_block
+
     def remove_skipped_blocks(
         self,
         request_id: str,
@@ -1625,9 +1675,7 @@ def remove_skipped_blocks(
                     blocks[last_state_block_idx] = self._null_block
 
     def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
-        """
-        cascade attention is not supported by mamba
-        """
+        """Cascade attention is not supported by mamba"""
         return 0
 
     def _needs_internal_checkpoint(
@@ -1747,10 +1795,10 @@ def get_num_blocks_to_allocate(
                     self._checkpoints.pop(request_id, None)
             if num_new_blocks > 0:
                 blocks_allocated = request_id in self._allocated_block_reqs
-                if not (checkpoint_block and blocks_allocated):
-                    num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block
-                    if not blocks_allocated:
-                        num_new_blocks += self.num_speculative_blocks
+                physical_block_cap = 1 + int(has_partial_hit) + checkpoint_block
+                if not blocks_allocated or checkpoint_block:
+                    physical_block_cap += self.num_speculative_blocks
+                num_new_blocks = min(num_new_blocks, physical_block_cap)
 
             num_evictable_computed_blocks = self._get_num_evictable_blocks(
                 new_computed_blocks
@@ -1914,6 +1962,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
         if self.mamba_cache_mode == "align":
             self._allocated_block_reqs.discard(request_id)
             self.last_state_block_idx.pop(request_id, None)
+            self._num_retired_blocks.pop(request_id, None)
             self._checkpoints.pop(request_id, None)
             self._producer_partial_tail_reqs.pop(request_id, None)
             # An offer is only guaranteed to hold committed bytes until the end
@@ -1929,8 +1978,7 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
         return super().pop_blocks_for_free(request_id)
 
     def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
-        """
-        Get the number of tokens whose mamba state are not needed anymore. Mamba only
+        """Get the number of tokens whose mamba state are not needed anymore. Mamba only
         need to keep the state of the last computed token, so we return
         num_computed_tokens - 1.
         """
@@ -1953,7 +2001,9 @@ def cache_blocks(
         )
         num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)
         if self.mamba_cache_mode == "align":
-            partial_hash = self._cache_partial_tail_block(request, num_tokens)
+            partial_hash = self._cache_partial_tail_block(
+                request, num_tokens, retention_interval=retention_interval
+            )
             if partial_hash is not None:
                 self.cached_blocks_this_step.add(partial_hash)
         if num_cached_blocks_after > num_cached_blocks_before:
@@ -1988,6 +2038,7 @@ def _cache_partial_tail_block(
         self,
         request: Request,
         num_tokens: int,
+        retention_interval: int | None,
     ) -> BlockHashWithGroupId | None:
         hash_block_size = self.block_pool.hash_block_size
         # Re-key the reserved block at its exported checkpoint boundary.
@@ -1997,6 +2048,17 @@ def _cache_partial_tail_block(
             blocks = self.req_to_blocks[request.request_id]
             assert 0 <= checkpoint_idx < len(blocks)
             checkpoint_block = blocks[checkpoint_idx]
+            if (
+                retention_interval == 0
+                and num_tokens < request.num_prompt_tokens
+                and checkpoint_position != request.shared_prefix_boundary
+            ):
+                # retention_interval == 0 keeps this transient checkpoint
+                # request-local. The slot may carry a hash from this step's
+                # full-block pass; that must go too, since the checkpoint
+                # state is about to overwrite the block.
+                self.block_pool._maybe_evict_cached_block(checkpoint_block)
+                return None
             if checkpoint_block.block_hash_num_tokens == checkpoint_position:
                 return None
             return self.block_pool.cache_partial_block(
@@ -2028,7 +2090,7 @@ def _cache_partial_tail_block(
         # running state block, mutated in place, which equals what its key
         # promises only after that step's forward.
         if num_tokens != latest_prompt_hash_boundary and not (
-            self.fine_grained_prefix_cache
+            self.shared_prefix_checkpoint
             and num_tokens == request.shared_prefix_boundary
             and request.num_computed_tokens < num_tokens <= request.num_prompt_tokens
         ):
@@ -2152,14 +2214,395 @@ def __init__(
         assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0
 
 
+class HiSparseSourceManager(FullAttentionManager):
+    """Host-tier manager with a private pool; publishes hashes once durable.
+
+    Host capacity is best effort: a page that cannot get a host block keeps a
+    null host entry, so its GPU copy is never written back and stays pinned.
+    """
+
+    coordinator: "HiSparseCoordinator | None" = None
+    # The host tier keeps prefixes the device groups have already lost.
+    retains_longer_hit = True
+
+    @property
+    def records_new_block_ids(self) -> bool:
+        return False
+
+    def take_new_block_ids(self) -> list[int]:
+        self.new_block_ids = []
+        return []
+
+    def bind_host_pool(self, num_blocks: int) -> None:
+        """Replace the device pool this group was built with by a host one."""
+        device_pool = self.block_pool
+        self.block_pool = SharedEventQueueBlockPool(
+            num_gpu_blocks=num_blocks,
+            enable_caching=self.enable_caching and self.kv_cache_spec.prefix_cacheable,
+            hash_block_size=device_pool.hash_block_size,
+            enable_kv_cache_events=device_pool.enable_kv_cache_events,
+            metrics_collector=device_pool.metrics_collector,
+            medium=MEDIUM_CPU,
+            event_owner=device_pool,
+        )
+        self._null_block = self.block_pool.null_block
+
+    def take_pending_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]:
+        """Host copies never reach the worker's generic block-copy path."""
+        return []
+
+    def take_host_cow_copies(self) -> list[tuple[KVCacheBlock, KVCacheBlock]]:
+        """Drain host copies for the coordinator to hand to the connector."""
+        copies = self._pending_cow_copies
+        self._pending_cow_copies = []
+        return copies
+
+    def get_num_blocks_to_allocate(
+        self,
+        request_id: str,
+        num_tokens: int,
+        new_computed_blocks: Sequence[KVCacheBlock],
+        total_computed_tokens: int,
+        num_local_computed_tokens: int,
+        num_tokens_main_model: int,
+        apply_admission_cap: bool = False,
+    ) -> int:
+        if (
+            total_computed_tokens > num_local_computed_tokens
+            or self._has_partial_local_hit(
+                new_computed_blocks, num_local_computed_tokens
+            )
+        ):
+            # External loads need real destinations; future GPU-computed pages
+            # remain best effort. Use the same admission sentinel as Mamba.
+            required = super().get_num_blocks_to_allocate(
+                request_id,
+                total_computed_tokens,
+                new_computed_blocks,
+                total_computed_tokens,
+                num_local_computed_tokens,
+                total_computed_tokens,
+            )
+            if required > self.block_pool.get_num_free_blocks():
+                assert self.coordinator is not None
+                assert self.coordinator.gpu_pool is not None
+                return self.coordinator.gpu_pool.num_gpu_blocks + 1
+        return 0
+
+    def allocate_new_blocks(
+        self, request_id: str, num_tokens: int, num_tokens_main_model: int
+    ) -> list[KVCacheBlock]:
+        req_blocks = self.req_to_blocks[request_id]
+        num_required = cdiv(num_tokens, self.block_size)
+        num_free = self.block_pool.get_num_free_blocks()
+        if request_id in self._partial_hit_reqs:
+            num_free -= 1
+        fit_blocks = min(num_required, len(req_blocks) + max(0, num_free))
+        new_blocks = super().allocate_new_blocks(
+            request_id,
+            fit_blocks * self.block_size,
+            min(num_tokens_main_model, fit_blocks * self.block_size),
+        )
+        req_blocks.extend([self._null_block] * (num_required - len(req_blocks)))
+        return new_blocks
+
+    def allocate_external_computed_blocks(
+        self,
+        request_id: str,
+        num_local_computed_tokens: int,
+        num_external_computed_tokens: int,
+    ) -> None:
+        if num_external_computed_tokens <= 0:
+            return
+        # The connector writes these pages; only a successful receive makes
+        # them readable, not advancing the request's computed-token count.
+        assert self.coordinator is not None
+        self.coordinator.record_pending_host_import(
+            request_id, num_local_computed_tokens + num_external_computed_tokens
+        )
+        super().allocate_external_computed_blocks(
+            request_id, num_local_computed_tokens, num_external_computed_tokens
+        )
+
+    def cache_blocks(
+        self,
+        request: Request,
+        num_tokens: int,
+        retention_interval: int | None = None,
+        *,
+        replay_boundaries: Sequence[int],
+    ) -> None:
+        assert self.coordinator is not None
+        self.coordinator.publish_when_ready(
+            request,
+            num_tokens,
+            retention_interval,
+            replay_boundaries=replay_boundaries,
+        )
+
+    def publish_blocks(
+        self,
+        request: Request,
+        num_tokens: int,
+        retention_interval: int | None = None,
+        *,
+        replay_boundaries: Sequence[int],
+    ) -> None:
+        super().cache_blocks(
+            request,
+            num_tokens,
+            retention_interval=retention_interval,
+            replay_boundaries=replay_boundaries,
+        )
+
+    def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
+        assert self.coordinator is not None
+        self.coordinator.free(request_id)
+        return super().pop_blocks_for_free(request_id)
+
+
+class _HiSparseAuxiliaryManager(SingleTypeKVCacheManager):
+    """Base for ephemeral groups whose host source owns prefix caching."""
+
+    coordinator: "HiSparseCoordinator | None" = None
+
+    def __init__(self, kv_cache_spec: KVCacheSpec, **kwargs) -> None:
+        # Never prefix-cached, but the per-step ``cache_blocks`` hook is where
+        # residency work runs, so stay opted in regardless of prefix caching.
+        kwargs["enable_caching"] = True
+        super().__init__(kv_cache_spec, **kwargs)
+
+    def cache_blocks(
+        self,
+        request: Request,
+        num_tokens: int,
+        retention_interval: int | None = None,
+        *,
+        replay_boundaries: Sequence[int],
+    ) -> None:
+        return None
+
+    def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
+        return 0
+
+    @classmethod
+    def find_longest_cache_hit(
+        cls,
+        block_hashes: BlockHashList,
+        max_length: int,
+        kv_cache_group_ids: list[int],
+        block_pool: BlockPool,
+        kv_cache_spec: KVCacheSpec,
+        drop_eagle_block: bool,
+        alignment_tokens: int,
+        dcp_world_size: int = 1,
+        pcp_world_size: int = 1,
+    ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
+        return tuple([] for _ in kv_cache_group_ids), 0
+
+
+class HiSparseHotManager(_HiSparseAuxiliaryManager):
+    """Allocate a hot region only after a request acquires CPU-only history."""
+
+    def __init__(self, kv_cache_spec: HiSparseHotSpec, **kwargs) -> None:
+        super().__init__(kv_cache_spec, **kwargs)
+        self.blocks_per_request = kv_cache_spec.blocks_per_request
+        self.hot_required: set[str] = set()
+
+    def require_hot(self, request_id: str) -> None:
+        self.hot_required.add(request_id)
+
+    def has_hot(self, request_id: str) -> bool:
+        return len(self.req_to_blocks.get(request_id, ())) == self.blocks_per_request
+
+    def get_num_blocks_to_allocate(
+        self,
+        request_id: str,
+        num_tokens: int,
+        new_computed_blocks: Sequence[KVCacheBlock],
+        total_computed_tokens: int,
+        num_local_computed_tokens: int,
+        num_tokens_main_model: int,
+        apply_admission_cap: bool = False,
+    ) -> int:
+        assert not new_computed_blocks
+        # A hot region is needed to read host-backed history: an external
+        # import, a new request resuming a host prefix, or one already asked
+        # to transition. Running requests keep their earlier answer.
+        host_import = total_computed_tokens > num_local_computed_tokens
+        resumes_host_prefix = (
+            num_local_computed_tokens > 0 and request_id not in self.num_cached_block
+        )
+        if host_import or resumes_host_prefix or request_id in self.hot_required:
+            return self.get_num_required_blocks(request_id)
+        return 0
+
+    def get_num_required_blocks(self, request_id: str) -> int:
+        return max(
+            self.blocks_per_request - len(self.req_to_blocks.get(request_id, ())),
+            0,
+        )
+
+    def add_local_computed_blocks(
+        self,
+        request_id: str,
+        new_computed_blocks: Sequence[KVCacheBlock],
+        num_local_computed_tokens: int,
+        num_external_computed_tokens: int,
+    ) -> None:
+        assert not new_computed_blocks
+        self.num_cached_block[request_id] = 0
+        assert self.coordinator is not None
+        if num_local_computed_tokens > 0 or not self.coordinator.resident_managers:
+            self.require_hot(request_id)
+
+    def allocate_external_computed_blocks(
+        self,
+        request_id: str,
+        num_local_computed_tokens: int,
+        num_external_computed_tokens: int,
+    ) -> None:
+        self.require_hot(request_id)
+
+    def allocate_new_blocks(
+        self, request_id: str, num_tokens: int, num_tokens_main_model: int
+    ) -> list[KVCacheBlock]:
+        # Cold admissions can bypass add_local_computed_blocks.
+        self.num_cached_block[request_id] = 0
+        if request_id not in self.hot_required:
+            return []
+        req_blocks = self.req_to_blocks[request_id]
+        num_new_blocks = self.blocks_per_request - len(req_blocks)
+        if num_new_blocks <= 0:
+            return []
+        new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
+        req_blocks.extend(new_blocks)
+        return new_blocks
+
+    def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
+        self.hot_required.discard(request_id)
+        return super().pop_blocks_for_free(request_id)
+
+
+class HiSparseResidentManager(_HiSparseAuxiliaryManager):
+    """Track GPU-resident pages for otherwise host-backed KV."""
+
+    def get_num_blocks_to_allocate(
+        self,
+        request_id: str,
+        num_tokens: int,
+        new_computed_blocks: Sequence[KVCacheBlock],
+        total_computed_tokens: int,
+        num_local_computed_tokens: int,
+        num_tokens_main_model: int,
+        apply_admission_cap: bool = False,
+    ) -> int:
+        del num_tokens_main_model
+        assert not new_computed_blocks
+        if total_computed_tokens > num_local_computed_tokens:
+            # Keep the last imported page writable, including the last-token
+            # replay on a full prompt hit. Only earlier pages stay host-only.
+            imported_pages = cdiv(total_computed_tokens, self.block_size)
+            existing = len(self.req_to_blocks.get(request_id, ()))
+            needs_tail = self.get_resident_page(request_id, imported_pages - 1) is None
+            return int(needs_tail) + max(
+                cdiv(num_tokens, self.block_size) - max(imported_pages, existing), 0
+            )
+        existing = len(self.req_to_blocks.get(request_id, ()))
+        host_pages = cdiv(num_local_computed_tokens, self.block_size)
+        required = cdiv(num_tokens, self.block_size)
+        if apply_admission_cap:
+            assert self._max_admission_blocks_per_request is not None
+            required = min(required, self._max_admission_blocks_per_request)
+        return max(required - max(existing, host_pages), 0)
+
+    def allocate_external_computed_blocks(
+        self,
+        request_id: str,
+        num_local_computed_tokens: int,
+        num_external_computed_tokens: int,
+    ) -> None:
+        """Reserve a writable final page; earlier imported pages stay on host."""
+        assert num_external_computed_tokens > 0
+        num_tokens = num_local_computed_tokens + num_external_computed_tokens
+        blocks = self.req_to_blocks[request_id]
+        tail_page = (num_tokens - 1) // self.block_size
+        if len(blocks) <= tail_page:
+            blocks.extend([self._null_block] * (tail_page + 1 - len(blocks)))
+        if blocks[tail_page].is_null:
+            blocks[tail_page] = self.block_pool.get_new_blocks(1)[0]
+
+    def add_local_computed_blocks(
+        self,
+        request_id: str,
+        new_computed_blocks: Sequence[KVCacheBlock],
+        num_local_computed_tokens: int,
+        num_external_computed_tokens: int,
+    ) -> None:
+        assert not new_computed_blocks
+        req_blocks = self.req_to_blocks[request_id]
+        assert not req_blocks
+        num_host_pages = cdiv(num_local_computed_tokens, self.block_size)
+        req_blocks.extend([self._null_block] * num_host_pages)
+        self.num_cached_block[request_id] = 0
+        assert self.coordinator is not None
+        self.coordinator.commit_computed_blocks(request_id, num_host_pages)
+
+    def cache_blocks(
+        self,
+        request: Request,
+        num_tokens: int,
+        retention_interval: int | None = None,
+        *,
+        replay_boundaries: Sequence[int],
+    ) -> None:
+        assert self.coordinator is not None
+        self.coordinator.plan_prefix_materialization(request.request_id, num_tokens)
+        self.coordinator.update_residency(request.request_id)
+
+    def allocate_new_blocks(
+        self, request_id: str, num_tokens: int, num_tokens_main_model: int
+    ) -> list[KVCacheBlock]:
+        del num_tokens_main_model
+        req_blocks = self.req_to_blocks[request_id]
+        num_new_blocks = cdiv(num_tokens, self.block_size) - len(req_blocks)
+        if num_new_blocks <= 0:
+            return []
+        new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
+        req_blocks.extend(new_blocks)
+        return new_blocks
+
+    def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
+        assert self.coordinator is not None
+        self.coordinator.free(request_id)
+        return super().pop_blocks_for_free(request_id)
+
+    def adopt_resident_page(
+        self, request_id: str, block_idx: int, block: KVCacheBlock
+    ) -> bool:
+        """Point a null prefix page at a pinned GPU copy of its contents."""
+        blocks = self.req_to_blocks.get(request_id)
+        if blocks is None or block_idx >= len(blocks) or not blocks[block_idx].is_null:
+            return False
+        blocks[block_idx] = block
+        return True
+
+    def get_resident_page(self, request_id: str, block_idx: int) -> KVCacheBlock | None:
+        blocks = self.req_to_blocks.get(request_id)
+        if blocks is None or block_idx >= len(blocks):
+            return None
+        block = blocks[block_idx]
+        return None if block.is_null else block
+
+
 def get_manager_for_kv_cache_spec(
     kv_cache_spec: KVCacheSpec,
     max_in_flight_tokens: int,
     max_model_len: int,
+    role: str | None = None,
     **kwargs,
 ) -> SingleTypeKVCacheManager:
-    """
-    Get the appropriate manager for a given KVCacheSpec.
+    """Get the appropriate manager for a given KVCacheSpec.
 
     Uses the KVCacheSpecRegistry to look up the manager class, supporting
     both built-in and custom specs registered via @register_kv_cache_spec
@@ -2170,10 +2613,13 @@ def get_manager_for_kv_cache_spec(
         max_in_flight_tokens: The max tokens scheduled but not yet settled
             (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`
         max_model_len: The maximum context length the model could serve
+        role: Optional KV cache group role used to select a role-specific
+            manager
     Returns:
         An instance of the appropriate SingleTypeKVCacheManager subclass
+
     """
-    manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec)
+    manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec, role)
     assert manager_class is not None, (
         f"No manager registered for KVCacheSpec {type(kv_cache_spec)}"
     )
@@ -2185,7 +2631,7 @@ def get_manager_for_kv_cache_spec(
     # FullAttentionSpec sizing without a separate admission cap.
     if isinstance(
         kv_cache_spec,
-        (SlidingWindowSpec, ChunkedLocalAttentionSpec),
+        (SlidingWindowSpec, ChunkedLocalAttentionSpec, HiSparseResidentSpec),
     ):
         kwargs["max_admission_blocks_per_request"] = (
             kv_cache_spec.max_admission_blocks_per_request(
@@ -2199,11 +2645,24 @@ def get_manager_for_kv_cache_spec(
 
 def register_all_kvcache_specs(vllm_config):
     """Built-in spec registration"""
+    KVCacheSpecRegistry.register_role_manager(
+        KVCacheGroupRole.HISPARSE_SOURCE, HiSparseSourceManager
+    )
     KVCacheSpecRegistry.register(
         FullAttentionSpec,
         FullAttentionManager,
         uniform_type_base_spec=FullAttentionSpec,
     )
+    KVCacheSpecRegistry.register(
+        HiSparseHotSpec,
+        HiSparseHotManager,
+        uniform_type_base_spec=HiSparseHotSpec,
+    )
+    KVCacheSpecRegistry.register(
+        HiSparseResidentSpec,
+        HiSparseResidentManager,
+        uniform_type_base_spec=HiSparseResidentSpec,
+    )
 
     KVCacheSpecRegistry.register(
         SlidingWindowSpec,
diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py
index 6a48b6282d43..0c1fb0b2c4c9 100644
--- a/vllm/v1/cudagraph_dispatcher.py
+++ b/vllm/v1/cudagraph_dispatcher.py
@@ -13,8 +13,7 @@
 
 
 class CudagraphDispatcher:
-    """
-    Runtime cudagraph dispatcher to dispatch keys for multiple set of
+    """Runtime cudagraph dispatcher to dispatch keys for multiple set of
     cudagraphs.
 
     The dispatcher stores two sets of dispatch keys, one for PIECEWISE and one
@@ -109,8 +108,7 @@ def _compute_bs_to_padded_graph_size(self) -> None:
                         )
 
     def _get_lora_cases(self) -> list[int]:
-        """
-        Returns list of has_lora values for CUDA graph capture.
+        """Returns list of has_lora values for CUDA graph capture.
         This is the single source of truth for LoRA capture cases.
         """
         lora_config = self.vllm_config.lora_config
@@ -241,8 +239,7 @@ def dispatch(
         valid_modes: AbstractSet[CUDAGraphMode] | None = None,
         invalid_modes: AbstractSet[CUDAGraphMode] | None = None,
     ) -> tuple[CUDAGraphMode, BatchDescriptor]:
-        """
-        Given conditions(e.g.,batch descriptor and if using piecewise only),
+        """Given conditions(e.g.,batch descriptor and if using piecewise only),
         dispatch to a cudagraph runtime mode and the valid batch descriptor.
         A new batch descriptor is returned as we might dispatch a uniform batch
         to a graph that supports a more general batch (uniform to non-uniform).
@@ -259,6 +256,7 @@ def dispatch(
                 valid_modes to compute allowed modes. (e.g., {FULL} for
                 features like cascade attention not supported by full
                 cudagraphs). None means no modes are excluded.
+
         """
         allowed_modes = valid_modes or CUDAGraphMode.valid_runtime_modes()
 
@@ -324,13 +322,13 @@ def dispatch(
         return CUDAGraphMode.NONE, BatchDescriptor(num_tokens)
 
     def get_capture_descs(self) -> list[tuple[CUDAGraphMode, list[BatchDescriptor]]]:
-        """
-        Returns capture descriptors for cudagraph capturing.
+        """Returns capture descriptors for cudagraph capturing.
 
         Returns:
             List of (runtime_mode, batch_descriptors) tuples, ordered PIECEWISE
             first then FULL. Batch descriptors are sorted largest-first for
             memory efficiency.
+
         """
         if not self.keys_initialized or self.cudagraph_mode == CUDAGraphMode.NONE:
             return []
diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py
index 5ae9ee0cac83..e5bbb96a9cf7 100644
--- a/vllm/v1/engine/__init__.py
+++ b/vllm/v1/engine/__init__.py
@@ -16,6 +16,7 @@
 from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
+from vllm.v1.kv_hints import KvHintsEnvelope
 from vllm.v1.metrics.stats import (
     PrefillStats,
     RequestSpecDecodeMetrics,
@@ -45,8 +46,7 @@ class EEPNotificationType(enum.Enum):
 
 
 class FinishReason(enum.IntEnum):
-    """
-    Reason a request finished - stop, length, abort, error, or repetition.
+    """Reason a request finished - stop, length, abort, error, or repetition.
 
     Int rather than Str for more compact serialization.
 
@@ -102,6 +102,8 @@ class EngineCoreReadyResponse:
     weight_transfer_backend: str | None = None
     enable_sleep_mode: bool = False
     supports_draft_weight_updates: bool = False
+    # Full-attention block size in tokens after initialization, or unavailable.
+    effective_attention_block_size: int | None = None
 
 
 class EngineCoreRequest(
@@ -156,6 +158,7 @@ class EngineCoreRequest(
     abort_immediately: bool = False
 
     session_id: str | None = None
+    kv_hints: KvHintsEnvelope | None = None
 
     @property
     def params(self) -> SamplingParams | PoolingParams:
@@ -282,8 +285,7 @@ def __post_init__(self):
 
 
 class EngineCoreRequestType(enum.Enum):
-    """
-    Request types defined as hex byte strings, so it can be sent over sockets
+    """Request types defined as hex byte strings, so it can be sent over sockets
     without separate encoding step.
     """
 
@@ -308,9 +310,7 @@ class ReconfigureDistributedRequest(msgspec.Struct):
 
 
 class ReconfigureRankType(enum.IntEnum):
-    """
-    Rank type for reconfiguring distributed request.
-    """
+    """Rank type for reconfiguring distributed request."""
 
     KEEP_CURRENT_RANK = -1
     SHUTDOWN_CURRENT_RANK = -2
diff --git a/vllm/v1/engine/admission_control.py b/vllm/v1/engine/admission_control.py
new file mode 100644
index 000000000000..b72d50f3465d
--- /dev/null
+++ b/vllm/v1/engine/admission_control.py
@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any, Protocol, cast
+
+
+class AdmissionCounterBuffer(Protocol):
+    def __getitem__(self, index: int) -> int: ...
+
+    def __setitem__(self, index: int, value: int) -> None: ...
+
+
+class SharedAdmissionStats:
+    """Lock-free admission stats shared by API server processes.
+
+    Each process writes exclusively to its own cache-line-sized slot. Readers
+    aggregate the slots, avoiding both inter-process locks and cache-line
+    contention on the request path.
+    """
+
+    # Space per-process counters 64 bytes apart to avoid false sharing on CPUs
+    # with 64-byte cache lines: one process's writes would otherwise invalidate
+    # the cache line containing another process's counter.
+    CACHE_LINE_BYTES = 64
+    COUNTER_BYTES = 8
+    COUNTERS_PER_SLOT = CACHE_LINE_BYTES // COUNTER_BYTES
+
+    def __init__(
+        self,
+        client_addresses: dict[str, Any],
+        client_count: int = 1,
+        client_index: int = 0,
+    ):
+        self._counters = cast(
+            AdmissionCounterBuffer,
+            client_addresses.pop("mp_admission_counters"),
+        )
+        self._client_count = client_count
+        self._slot_offset = client_index * self.COUNTERS_PER_SLOT
+
+    @classmethod
+    def num_counters(cls, client_count: int) -> int:
+        return client_count * cls.COUNTERS_PER_SLOT
+
+    def set_num_requests(self, count: int) -> None:
+        self._counters[self._slot_offset] = count
+
+    def get_num_requests(self) -> int:
+        return sum(
+            self._counters[client_index * self.COUNTERS_PER_SLOT]
+            for client_index in range(self._client_count)
+        )
diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py
index 91618c409362..7d71761990b0 100644
--- a/vllm/v1/engine/async_llm.py
+++ b/vllm/v1/engine/async_llm.py
@@ -12,6 +12,7 @@
 import vllm.envs as envs
 from vllm import TokensPrompt
 from vllm.config import VllmConfig
+from vllm.config.kv_events import KVEventsConfig
 from vllm.distributed.weight_transfer.base import (
     WeightTransferInitRequest,
     WeightTransferUpdateRequest,
@@ -44,6 +45,7 @@
 from vllm.utils.async_utils import cancel_task_threadsafe
 from vllm.utils.collection_utils import as_list
 from vllm.v1.engine import EngineCoreRequest, PauseMode
+from vllm.v1.engine.admission_control import SharedAdmissionStats
 from vllm.v1.engine.core_client import EngineCoreClient
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 from vllm.v1.engine.input_processor import InputProcessor
@@ -51,6 +53,7 @@
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.executor import Executor
 from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult
+from vllm.v1.kv_hints import KvHintsEnvelope
 from vllm.v1.metrics.loggers import (
     StatLoggerFactory,
     StatLoggerManager,
@@ -93,8 +96,7 @@ def __init__(
         client_index: int = 0,
         profiler: TorchProfilerWrapper | None = None,
     ) -> None:
-        """
-        Create an AsyncLLM.
+        """Create an AsyncLLM.
 
         Args:
             vllm_config: global configuration.
@@ -104,13 +106,21 @@ def __init__(
             mm_registry: Multi-modal registry.
             log_requests: Whether to log requests.
             start_engine_loop: Whether to start the engine loop.
+            client_addresses: ZMQ addresses of an externally launched engine,
+                when not launching one in-process.
+            client_count: Number of API-server clients sharing the engine.
+            client_index: Index of this client within ``client_count``.
+            aggregate_engine_logging: Whether to aggregate per-engine logs
+                into a single stat logger.
             stat_loggers: customized stat loggers for the engine.
                 If not provided, default stat loggers will be used.
                 PLEASE BE AWARE THAT STAT LOGGER IS NOT STABLE
                 IN V1, AND ITS BASE CLASS INTERFACE MIGHT CHANGE.
+            profiler: Torch profiler wrapper used to trace the frontend.
 
         Returns:
             None
+
         """
         # Ensure we can serialize custom transformer configs
         maybe_register_config_serialize_by_value()
@@ -144,12 +154,20 @@ def __init__(
         # Convert EngineInput --> EngineCoreRequest.
         self.input_processor = InputProcessor(self.vllm_config, renderer)
 
+        self.admission_stats = (
+            SharedAdmissionStats(client_addresses, client_count, client_index)
+            if client_addresses is not None
+            and "mp_admission_counters" in client_addresses
+            else None
+        )
+
         # Converts EngineCoreOutputs --> RequestOutput.
         self.output_processor = OutputProcessor(
             renderer.tokenizer,
             log_stats=self.log_stats,
             stream_interval=self.vllm_config.scheduler_config.stream_interval,
             tracing_enabled=tracing_endpoint is not None,
+            admission_stats=self.admission_stats,
         )
 
         # EngineCore (starts the engine in background process).
@@ -191,9 +209,11 @@ def __init__(
             pass
 
         self.profiler = profiler
+        configured_activities = vllm_config.profiler_config.torch_profiler_activities
         if (
             vllm_config.profiler_config.profiler == "torch"
             and not vllm_config.profiler_config.ignore_frontend
+            and (configured_activities is None or "CPU" in configured_activities)
         ):
             profiler_dir = vllm_config.profiler_config.torch_profiler_dir
             logger.info(
@@ -246,7 +266,6 @@ def from_engine_args(
         stat_loggers: list[StatLoggerFactory] | None = None,
     ) -> "AsyncLLM":
         """Create an AsyncLLM from the EngineArgs."""
-
         # Create the engine configs.
         vllm_config = engine_args.create_engine_config(usage_context)
         executor_class = Executor.get_class(vllm_config)
@@ -291,8 +310,9 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None:
         Both limits return HTTP 503 (Service Unavailable) so that load
         balancers and client SDKs retry on a different instance.
 
-        - ``max_num_queued_reqs``: hard cap on the number of unfinished requests
-          (waiting + running).  A request with ``n > 1`` counts as ``n`` slots.
+        - ``max_num_queued_reqs``: approximate global cap on the number of
+          unfinished requests (waiting + running). A request with ``n > 1``
+          counts as ``n`` slots.
         - ``max_num_queued_tokens``: TTFT QoS — cap on the total prompt
           tokens of requests still in prefill.
 
@@ -309,16 +329,21 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None:
         Raises:
             QueueOverflowError: If ``max_num_queued_reqs`` would be exceeded.
             MaxQueuedTokensError: If ``max_num_queued_tokens`` would be exceeded.
+
         """
         max_num_reqs = self.scheduler_config.max_num_queued_reqs
         if max_num_reqs is not None:
-            current = self.get_num_unfinished_requests()
-            if current + n > max_num_reqs:
+            current_requests = (
+                self.admission_stats.get_num_requests()
+                if self.admission_stats is not None
+                else self.get_num_unfinished_requests()
+            )
+            if current_requests + n > max_num_reqs:
                 logger.info(
                     "Request queue full - rejecting request %s "
                     "(current=%d, n=%d, max=%d).",
                     request_id,
-                    current,
+                    current_requests,
                     n,
                     max_num_reqs,
                 )
@@ -362,9 +387,9 @@ async def add_request(
         prompt_text: str | None = None,
         reasoning_ended: bool | None = None,
         reasoning_parser_kwargs: dict[str, Any] | None = None,
+        kv_hints: KvHintsEnvelope | None = None,
     ) -> RequestOutputCollector:
         """Add new request to the AsyncLLM."""
-
         if self.errored:
             raise EngineDeadError()
 
@@ -381,10 +406,6 @@ async def add_request(
                 "prompt logprobs"
             )
 
-        if isinstance(params, SamplingParams) and params.n > 1:
-            # TODO (NickLucche) Batch check admission check for all n requests
-            self.check_admission(params.n, request_id)
-
         if isinstance(prompt, AsyncGenerator):
             if reasoning_ended is not None or reasoning_parser_kwargs is not None:
                 raise NotImplementedError
@@ -401,6 +422,7 @@ async def add_request(
                 priority,
                 data_parallel_rank,
                 session_id,
+                kv_hints,
             )
 
         # Convert Input --> Request.
@@ -434,6 +456,7 @@ async def add_request(
                     priority=priority,
                     data_parallel_rank=data_parallel_rank,
                     session_id=session_id,
+                    kv_hints=kv_hints,
                 )
             else:
                 # Raw prompts require tokenization and possibly multimodal
@@ -450,6 +473,7 @@ async def add_request(
                     priority=priority,
                     data_parallel_rank=data_parallel_rank,
                     session_id=session_id,
+                    kv_hints=kv_hints,
                 )
             prompt_text, _, _ = extract_prompt_components(self.model_config, prompt)
 
@@ -480,14 +504,28 @@ async def add_request(
 
         # Fan out child requests (for n>1).
         parent_request = ParentRequest(request)
+        child_requests: list[EngineCoreRequest] = []
         for idx in range(parent_params.n):
             request_id, child_params = parent_request.get_child_info(idx)
             child_request = request if idx == parent_params.n - 1 else copy(request)
             child_request.request_id = request_id
             child_request.sampling_params = child_params
-            await self._add_request(
-                child_request, prompt_text, parent_request, idx, queue
-            )
+            child_requests.append(child_request)
+
+        self.check_admission(parent_params.n, parent_request.request_id)
+        try:
+            for idx, child_request in enumerate(child_requests):
+                self.output_processor.add_request(
+                    child_request, prompt_text, parent_request, idx, queue
+                )
+
+            for child_request in child_requests:
+                await self.engine_core.add_request_async(child_request)
+                if self.log_requests:
+                    logger.info("Added request %s.", child_request.request_id)
+        except BaseException:
+            await self.abort(parent_request.request_id, internal=True)
+            raise
         return queue
 
     async def _add_request(
@@ -524,6 +562,7 @@ async def _add_streaming_input_request(
         priority: int = 0,
         data_parallel_rank: int | None = None,
         session_id: str | None = None,
+        kv_hints: KvHintsEnvelope | None = None,
     ) -> RequestOutputCollector:
         self._validate_streaming_input_sampling_params(sampling_params)
 
@@ -536,6 +575,7 @@ async def _add_streaming_input_request(
             priority=priority,
             data_parallel_rank=data_parallel_rank,
             session_id=session_id,
+            kv_hints=kv_hints,
         )
 
         if not sampling_params.skip_clone:
@@ -637,11 +677,11 @@ async def generate(
         priority: int = 0,
         data_parallel_rank: int | None = None,
         session_id: str | None = None,
+        kv_hints: KvHintsEnvelope | None = None,
         reasoning_ended: bool | None = None,
         reasoning_parser_kwargs: dict[str, Any] | None = None,
     ) -> AsyncGenerator[RequestOutput, None]:
-        """
-        Main function called by the API server to kick off a request
+        """Main function called by the API server to kick off a request
             * 1) Making an AsyncStream corresponding to the Request.
             * 2) Processing the Input.
             * 3) Adding the Request to the Detokenizer.
@@ -664,8 +704,8 @@ async def generate(
             >>> params = self.renderer.default_cmpl_tok_params
             >>> (engine_input,) = self.renderer.render_cmpl([parsed], params)
             >>> gen = self.generate(engine_input, sampling_params, request_id)
-        """
 
+        """
         q: RequestOutputCollector | None = None
         try:
             q = await self.add_request(
@@ -678,6 +718,7 @@ async def generate(
                 priority=priority,
                 data_parallel_rank=data_parallel_rank,
                 session_id=session_id,
+                kv_hints=kv_hints,
                 prompt_text=prompt_text,
                 reasoning_ended=reasoning_ended,
                 reasoning_parser_kwargs=reasoning_parser_kwargs,
@@ -749,7 +790,6 @@ async def generate(
 
     def _run_output_handler(self):
         """Background loop: pulls from EngineCore and pushes to AsyncStreams."""
-
         if self.output_handler is not None:
             return
 
@@ -835,7 +875,6 @@ async def abort(
         self, request_id: str | Iterable[str], internal: bool = False
     ) -> None:
         """Abort RequestId in OutputProcessor and EngineCore."""
-
         request_ids = (
             (request_id,) if isinstance(request_id, str) else as_list(request_id)
         )
@@ -879,8 +918,7 @@ async def pause_generation(
         wait_for_inflight_requests: bool | None = None,
         clear_cache: bool = True,
     ) -> None:
-        """
-        Pause generation to allow model weight updates.
+        """Pause generation to allow model weight updates.
 
         All mode handling (abort / wait / keep) and cache clearing is done
         in the engine. New generation/encoding requests will not be scheduled
@@ -896,6 +934,7 @@ async def pause_generation(
             wait_for_inflight_requests: DEPRECATED: use mode argument.
             clear_cache: Whether to clear KV cache and prefix cache after
                 draining. Set to ``False`` to preserve cache for faster resume.
+
         """
         if wait_for_inflight_requests:
             warnings.warn(
@@ -936,8 +975,7 @@ async def encode(
         tokenization_kwargs: dict[str, Any] | None = None,
         reasoning_ended: bool | None = None,
     ) -> AsyncGenerator[PoolingRequestOutput, None]:
-        """
-        Main function called by the API server to kick off a request
+        """Main function called by the API server to kick off a request
             * 1) Making an AsyncStream corresponding to the Request.
             * 2) Processing the Input.
             * 3) Adding the Request to the EngineCore (separate process).
@@ -949,7 +987,6 @@ async def encode(
         The caller of generate() iterates the returned AsyncGenerator,
         returning the RequestOutput back to the caller.
         """
-
         q: RequestOutputCollector | None = None
         try:
             q = await self.add_request(
@@ -1064,10 +1101,17 @@ async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         if self.logger_manager is not None:
             self.logger_manager.record_sleep_state(1, level)
 
-    async def wake_up(self, tags: list[str] | None = None) -> None:
-        await self.engine_core.wake_up_async(tags)
+    async def release_kv_cache_memory(self) -> None:
+        await self.renderer.clear_mm_cache_async()
+        await self.engine_core.release_kv_cache_memory_async()
 
         if self.logger_manager is not None:
+            self.logger_manager.record_sleep_state(1, 0)
+
+    async def wake_up(self, tags: list[str] | None = None) -> None:
+        fully_awake = await self.engine_core.wake_up_async(tags)
+
+        if self.logger_manager is not None and fully_awake:
             self.logger_manager.record_sleep_state(0, 0)
 
     async def checkpoint_prepare(self) -> None:
@@ -1102,9 +1146,7 @@ async def collective_rpc(
         args: tuple = (),
         kwargs: dict | None = None,
     ):
-        """
-        Perform a collective RPC call to the given path.
-        """
+        """Perform a collective RPC call to the given path."""
         return await self.engine_core.collective_rpc_async(
             method, timeout, args, kwargs
         )
@@ -1184,7 +1226,7 @@ async def _scale_elastic_ep(
     async def handle_fault(
         self, fault_tolerance_request: FaultToleranceRequest
     ) -> FaultToleranceResult:
-        """send fault tolerance instruction to the engine"""
+        """Send fault tolerance instruction to the engine."""
         return await self.engine_core.handle_fault(fault_tolerance_request)
 
     async def get_status(self):
@@ -1210,11 +1252,11 @@ def dead_error(self) -> BaseException:
     async def init_weight_transfer_engine(
         self, request: WeightTransferInitRequest
     ) -> None:
-        """
-        Initialize weight transfer for RL training.
+        """Initialize weight transfer for RL training.
 
         Args:
             request: Weight transfer initialization request with backend-specific info
+
         """
         await self.collective_rpc(
             "init_weight_transfer_engine", kwargs={"init_info": request.init_info}
@@ -1229,11 +1271,11 @@ async def start_draft_weight_update(self) -> None:
         await self.collective_rpc("start_draft_weight_update")
 
     async def update_weights(self, request: WeightTransferUpdateRequest) -> None:
-        """
-        Batched weight update for RL training.
+        """Batched weight update for RL training.
 
         Args:
             request: Weight update request with backend-specific update info
+
         """
         await self.collective_rpc(
             "update_weights", kwargs={"update_info": request.update_info}
@@ -1252,3 +1294,6 @@ async def update_weight_version(self, new_version: str) -> None:
     async def get_weight_version(self) -> str:
         """Return the latest committed weight version."""
         return await self.engine_core.get_weight_version_async()
+
+    def get_kv_event_sources(self) -> dict[int, KVEventsConfig]:
+        return self.engine_core.get_kv_event_sources()
diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
index ffc9a5de7e68..9b27f50edc03 100644
--- a/vllm/v1/engine/core.py
+++ b/vllm/v1/engine/core.py
@@ -31,8 +31,10 @@
 from vllm.logger import init_logger
 from vllm.logging_utils.dump_input import dump_engine_exception
 from vllm.lora.request import LoRARequest
-from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.cache import MultiModalCacheMissError
+from vllm.multimodal.cache import (
+    MultiModalCacheMissError,
+    engine_receiver_cache_from_config,
+)
 from vllm.tasks import POOLING_TASKS, SupportedTask
 from vllm.tracing import instrument, maybe_init_worker_tracer
 from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
@@ -85,7 +87,11 @@
     EngineCoreSentinel,
     fault_tolerant_wrapper,
 )
-from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    get_kv_cache_spec_kind,
+    is_full_attention_spec,
+)
 from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
@@ -168,6 +174,7 @@ def __init__(
             block_size=scheduler_block_size,
             hash_block_size=hash_block_size,
         )
+        self._initialize_effective_attention_block_size()
         self.use_spec_decode = vllm_config.speculative_config is not None
         self.check_for_draft_tokens = (
             self.use_spec_decode or vllm_config.model_config.is_diffusion
@@ -177,10 +184,7 @@ def __init__(
         if self.scheduler.ec_connector is not None:  # type: ignore
             self.model_executor.init_ec_output_aggregator()
 
-        mm_registry = MULTIMODAL_REGISTRY
-        self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config(
-            vllm_config
-        )
+        self.mm_receiver_cache = engine_receiver_cache_from_config(vllm_config)
 
         # If a KV connector is initialized for scheduler, we want to collect
         # handshake metadata from all workers so the connector in the scheduler
@@ -382,6 +386,21 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig:
             )
         return scheduler_kv_cache_config
 
+    def _initialize_effective_attention_block_size(self) -> None:
+        cache_config = self.vllm_config.cache_config
+        cache_config.effective_attention_block_size = None
+        cache_manager = getattr(self.scheduler, "kv_cache_manager", None)
+        if cache_manager is None:
+            return
+        block_sizes = {
+            manager.block_size
+            for manager in cache_manager.coordinator.single_type_managers
+            if is_full_attention_spec(manager.kv_cache_spec)
+        }
+        cache_config.effective_attention_block_size = (
+            block_sizes.pop() if len(block_sizes) == 1 else None
+        )
+
     def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
         supported_tasks = self.model_executor.supported_tasks
         self._log_pooler_config(supported_tasks)
@@ -441,6 +460,32 @@ def log_field(name: str, field: str) -> str:
             supported_pooling_tasks,
         )
 
+    def get_kv_cache_group_metadata(self) -> list[dict[str, int | str | None]]:
+        """Return serializable KV cache metadata for external event consumers
+        (e.g., Dynamo). This function has no call sites within vLLM but needs to
+        be kept.
+        """
+        kv_cache_config = getattr(self.scheduler, "kv_cache_config", None)
+        if kv_cache_config is None:
+            return []
+
+        kv_cache_manager = cast(Any, self.scheduler).kv_cache_manager
+        managers = kv_cache_manager.coordinator.single_type_managers
+        metadata: list[dict[str, int | str | None]] = []
+        for group_idx, (group, manager) in enumerate(
+            zip(kv_cache_config.kv_cache_groups, managers, strict=True)
+        ):
+            spec = group.kv_cache_spec
+            metadata.append(
+                {
+                    "group_idx": group_idx,
+                    "kind": get_kv_cache_spec_kind(spec).value,
+                    "block_size": manager.block_size,
+                    "sliding_window": getattr(spec, "sliding_window", None),
+                }
+            )
+        return metadata
+
     def add_request(self, request: Request, request_wave: int = 0):
         """Add request to the scheduler.
 
@@ -489,7 +534,6 @@ def add_request(self, request: Request, request_wave: int = 0):
 
     def abort_requests(self, request_ids: list[str]):
         """Abort requests from the scheduler."""
-
         # TODO: The scheduler doesn't really need to know the
         # specific finish reason, TBD whether we propagate that
         # (i.e. client-aborted vs stop criteria met).
@@ -592,7 +636,6 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]:
         Returns tuple of outputs and a flag indicating whether the model
         was executed.
         """
-
         # Check for any requests remaining in the scheduler - unfinished,
         # or finished and not yet removed from the batch.
         if not self.scheduler.has_requests():
@@ -643,7 +686,6 @@ def step_with_batch_queue(
         batch in the job queue is finished.
         3. Update the scheduler from the output.
         """
-
         batch_queue = self.batch_queue
         assert batch_queue is not None
 
@@ -824,10 +866,11 @@ def _reset_caches(
         # reset_connector=True so external connectors clear alongside
         # local caches, matching the pause_generation(clear_cache=True)
         # contract. No-op when no connector is configured.
-        self.reset_prefix_cache(
+        if not self.reset_prefix_cache(
             reset_running_requests=reset_running_requests,
             reset_connector=reset_connector,
-        )
+        ):
+            raise RuntimeError("Failed to reset the KV connector cache.")
         self.reset_mm_cache()
         self.reset_encoder_cache()
 
@@ -887,8 +930,8 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None | Future:
                 - Level 2: Discard all GPU memory.
             mode: Pause mode - how to deal with any existing requests, see
                 documentation of pause_scheduler method.
-        """
 
+        """
         # Pause scheduler before sleeping.
         clear_prefix_cache = level >= 1
         pause_future = self.pause_scheduler(mode=mode, clear_cache=clear_prefix_cache)
@@ -914,11 +957,15 @@ def pause_complete(f: Future):
         pause_future.add_done_callback(pause_complete)
         return future
 
-    def wake_up(self, tags: list[str] | None = None):
+    def wake_up(self, tags: list[str] | None = None) -> bool:
         """Wake up the engine from sleep.
 
         Args:
             tags: Tags to wake up. Use ["scheduling"] for level 0 wake up.
+
+        Returns:
+            Whether all executor memory is resident again (fully awake).
+
         """
         if tags is not None and "scheduling" in tags:
             # Remove "scheduling" from tags if there are other tags to process.
@@ -929,8 +976,30 @@ def wake_up(self, tags: list[str] | None = None):
 
         # Partial wakes intentionally keep the remaining allocations asleep.
         # Resume scheduling only once all executor memory is resident again.
-        if not self.model_executor.is_sleeping:
+        fully_awake = not self.model_executor.is_sleeping
+        if fully_awake:
             self.resume_scheduler()
+        return fully_awake
+
+    def release_kv_cache_memory(self) -> None:
+        """Discard KV cache physical memory. Requires a completed pause
+        and all executor memory to be resident. Kept requests are recomputed
+        after wake-up.
+        """
+        if not (
+            self.is_scheduler_paused()
+            and not self.scheduler.has_requests()
+            and not self.batch_queue
+        ):
+            raise RuntimeError(
+                "release_kv_cache_memory() requires a completed pause first"
+            )
+        if self.model_executor.is_sleeping:
+            raise RuntimeError(
+                "release_kv_cache_memory() requires all executor memory to be resident"
+            )
+        self._reset_caches()
+        self.model_executor.discard(("kv_cache",))
 
     def is_sleeping(self) -> bool:
         """Check if engine is sleeping at any level."""
@@ -1146,8 +1215,7 @@ def _perform_handshakes(
         vllm_config: VllmConfig,
         client_handshake_address: str | None,
     ) -> Generator[EngineZmqAddresses, None, None]:
-        """
-        Perform startup handshakes.
+        """Perform startup handshakes.
 
         For DP=1 or offline mode, this is with the colocated front-end process.
 
@@ -1282,7 +1350,6 @@ def startup_handshake(
     @staticmethod
     def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs):
         """Launch EngineCore busy loop in background process."""
-
         # Ensure we can serialize transformer config after spawning
         maybe_register_config_serialize_by_value()
 
@@ -1428,7 +1495,6 @@ def _maybe_publish_request_counts(self):
 
     def _process_input_queue(self):
         """Exits when an engine step needs to be performed."""
-
         waited = False
         while not self.has_work() and self.is_running():
             # Notify callbacks waiting for engine to become idle.
@@ -1459,7 +1525,6 @@ def _process_input_queue(self):
 
     def _process_engine_step(self) -> bool:
         """Called only when there are unfinished local requests."""
-
         # Step the engine core.
         outputs, model_executed = self.step_fn()
         # Put EngineCoreOutputs into the output queue.
@@ -1533,7 +1598,6 @@ def _handle_client_request(
         self, request_type: EngineCoreRequestType, request: Any
     ) -> None:
         """Dispatch request from client."""
-
         if request_type == EngineCoreRequestType.WAKEUP:
             return
         elif request_type == EngineCoreRequestType.ADD:
@@ -1629,7 +1693,6 @@ def _convert_msgspec_args(method, args):
 
     def _send_engine_dead(self):
         """Send EngineDead status to the EngineCoreClient."""
-
         # Put ENGINE_CORE_DEAD in the queue.
         self.output_queue.put_nowait(EngineCoreProc.ENGINE_CORE_DEAD)
 
@@ -1649,6 +1712,9 @@ def _make_ready_response(self) -> EngineCoreReadyResponse:
             num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0,
             block_size=self.vllm_config.cache_config.block_size,
             mamba_block_size=self.vllm_config.cache_config.mamba_block_size,
+            effective_attention_block_size=(
+                self.vllm_config.cache_config.effective_attention_block_size
+            ),
             dp_stats_address=self.frontend_stats_publish_address,
             dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."),
             vllm_version=VLLM_VERSION,
@@ -1680,6 +1746,8 @@ def _make_ready_response(self) -> EngineCoreReadyResponse:
             enable_sleep_mode=self.vllm_config.model_config.enable_sleep_mode,
             supports_draft_weight_updates=(
                 self.model_executor.supports_draft_weight_updates()
+                if self.vllm_config.weight_transfer_config is not None
+                else False
             ),
         )
 
@@ -1691,7 +1759,6 @@ def process_input_sockets(
         ready_event: threading.Event,
     ):
         """Input socket IO thread."""
-
         # Msgpack serialization decoding with optional tensor IPC receiver.
         add_request_decoder = MsgpackDecoder(
             EngineCoreRequest, oob_tensor_provider=self.tensor_ipc_receiver
@@ -1790,7 +1857,6 @@ def process_output_sockets(
         self, output_paths: list[str], coord_output_path: str | None, engine_index: int
     ):
         """Output socket IO thread."""
-
         # Msgpack serialization encoding.
         encoder = MsgpackEncoder()
         # Send buffers to reuse.
@@ -1935,8 +2001,12 @@ def pause_scheduler(
             raise ValueError(f"Invalid pause mode: {mode}")
 
         def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None:
-            engine._finish_pause(clear_cache)
-            future.set_result(None)
+            try:
+                engine._finish_pause(clear_cache)
+            except Exception as e:
+                future.set_exception(e)
+            else:
+                future.set_result(None)
 
         if mode == "abort":
             aborted_reqs = self.scheduler.finish_requests(
@@ -2180,7 +2250,6 @@ def _should_throttle_prefills(self) -> bool:
     @fault_tolerant_wrapper
     def run_busy_loop(self):
         """Core busy loop of the EngineCore for data parallel case."""
-
         # Loop until process is sent a SIGINT or SIGTERM
         while self._handle_shutdown():
             # 1) Poll the input queue until there is work to do.
@@ -2345,8 +2414,7 @@ def commit_prepared_elastic_ep(self) -> None:
     def _eep_send_engine_core_notification(
         self, notification_type: EEPNotificationType
     ):
-        """
-        Send notifications to EngineCoreClient, which can then forward
+        """Send notifications to EngineCoreClient, which can then forward
         the notifications to other engine core processes. It is used for:
         1) In scale down: removing core engines to notify EngineCoreClient
            so EngineCoreClient can release their ray placement groups;
@@ -2396,9 +2464,7 @@ def _eep_scale_up_before_kv_init(self):
 
 
 class EngineCoreActorMixin:
-    """
-    Ray actor for running EngineCore in a data parallel context
-    """
+    """Ray actor for running EngineCore in a data parallel context."""
 
     def __init__(
         self,
@@ -2452,13 +2518,10 @@ def _set_nixl_side_channel_host():
     def _set_visible_devices(self, vllm_config: VllmConfig, local_dp_rank: int):
         from vllm.platforms import current_platform
 
-        if current_platform.is_xpu():
-            pass
-        else:
-            device_control_env_var = current_platform.device_control_env_var
-            self._set_assigned_physical_gpu_ids(
-                vllm_config, local_dp_rank, device_control_env_var
-            )
+        device_control_env_var = current_platform.device_control_env_var
+        self._set_assigned_physical_gpu_ids(
+            vllm_config, local_dp_rank, device_control_env_var
+        )
 
     def _set_assigned_physical_gpu_ids(
         self,
@@ -2494,16 +2557,14 @@ def _perform_handshakes(
         vllm_config: VllmConfig,
         client_handshake_address: str | None,
     ):
-        """
-        For Ray, we don't need to actually perform handshake.
+        """For Ray, we don't need to actually perform handshake.
         All addresses information is known before the actor creation.
         Therefore, we simply yield these addresses.
         """
         yield self.addresses
 
     def wait_for_init(self):
-        """
-        Wait until the engine core is initialized.
+        """Wait until the engine core is initialized.
 
         This is just an empty method. When ray.get() on this method
         (or any other method of the actor) returns, it is guaranteed
@@ -2512,9 +2573,7 @@ def wait_for_init(self):
         pass
 
     def run(self):
-        """
-        Run the engine core busy loop.
-        """
+        """Run the engine core busy loop."""
         try:
             self.run_busy_loop()  # type: ignore[attr-defined]
         except SystemExit:
diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py
index ffd10ab5572f..5d59bdeaea3a 100644
--- a/vllm/v1/engine/core_client.py
+++ b/vllm/v1/engine/core_client.py
@@ -23,6 +23,7 @@
 
 from vllm import envs
 from vllm.config import VllmConfig
+from vllm.config.kv_events import KVEventsConfig
 from vllm.envs import VLLM_ENGINE_READY_TIMEOUT_S
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
@@ -77,8 +78,7 @@
 
 
 class EngineCoreClient(ABC):
-    """
-    EngineCoreClient: subclasses handle different methods for pushing
+    """EngineCoreClient: subclasses handle different methods for pushing
         and pulling from the EngineCore for asyncio / multiprocessing.
 
     Subclasses:
@@ -87,6 +87,10 @@ class EngineCoreClient(ABC):
     * AsyncMPClient: ZMQ + background proc EngineCore w/ asyncio (for AsyncLLM)
     """
 
+    def get_kv_event_sources(self) -> dict[int, KVEventsConfig]:
+        """KV-event publisher config of each ready engine, keyed by DP rank."""
+        return {}
+
     @staticmethod
     def make_client(
         multiprocess_mode: bool,
@@ -194,7 +198,10 @@ def reset_encoder_cache(self) -> None:
     def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         raise NotImplementedError
 
-    def wake_up(self, tags: list[str] | None = None) -> None:
+    def release_kv_cache_memory(self) -> None:
+        raise NotImplementedError
+
+    def wake_up(self, tags: list[str] | None = None) -> bool:
         raise NotImplementedError
 
     def is_sleeping(self) -> bool:
@@ -286,7 +293,10 @@ async def reset_encoder_cache_async(self) -> None:
     async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:
         raise NotImplementedError
 
-    async def wake_up_async(self, tags: list[str] | None = None) -> None:
+    async def release_kv_cache_memory_async(self) -> None:
+        raise NotImplementedError
+
+    async def wake_up_async(self, tags: list[str] | None = None) -> bool:
         raise NotImplementedError
 
     async def is_sleeping_async(self) -> bool:
@@ -331,8 +341,7 @@ async def get_status(self):
 
 
 class InprocClient(EngineCoreClient):
-    """
-    InprocClient: client for in-process EngineCore. Intended
+    """InprocClient: client for in-process EngineCore. Intended
     for use in LLMEngine for V0-style add_request() and step()
         EngineCore setup in this process (no busy loop).
 
@@ -357,6 +366,11 @@ def __init__(
             log_stats,
             executor_fail_callback=executor_fail_callback,
         )
+        # Release the executor's workers when the client is collected, as
+        # MPClient does; they hold GPU memory for the life of the process
+        # otherwise. The callback must not reference self, or self would stay
+        # reachable and the finalizer would never run.
+        self._finalizer = weakref.finalize(self, self.engine_core.shutdown)
 
     def get_output(self) -> EngineCoreOutputs:
         outputs, model_executed = self.engine_core.step_fn()
@@ -375,7 +389,8 @@ def abort_requests(self, request_ids: list[str]) -> None:
             self.engine_core.abort_requests(request_ids)
 
     def shutdown(self, timeout: float | None = None) -> None:
-        self.engine_core.shutdown()
+        if self._finalizer.detach() is not None:
+            self.engine_core.shutdown()
 
     def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
         self.engine_core.profile(is_start, profile_prefix)
@@ -399,8 +414,11 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         result = self.engine_core.sleep(level, mode)
         assert result is None
 
-    def wake_up(self, tags: list[str] | None = None) -> None:
-        self.engine_core.wake_up(tags)
+    def release_kv_cache_memory(self) -> None:
+        self.engine_core.release_kv_cache_memory()
+
+    def wake_up(self, tags: list[str] | None = None) -> bool:
+        return self.engine_core.wake_up(tags)
 
     def is_sleeping(self) -> bool:
         return self.engine_core.is_sleeping()
@@ -469,7 +487,6 @@ class BackgroundResources:
 
     def __call__(self):
         """Clean up background resources."""
-
         logger.debug_once("[shutdown] MPClient: background resource cleanup start")
         self.engine_dead = True
         if self.engine_manager is not None:
@@ -543,16 +560,15 @@ class ElasticScalingCache:
 
 
 class MPClient(EngineCoreClient):
-    """
-    MPClient: base client for multi-proc EngineCore.
-        EngineCore runs in a background process busy loop, getting
-        new EngineCoreRequests and returning EngineCoreOutputs
+    """MPClient: base client for multi-proc EngineCore.
+    EngineCore runs in a background process busy loop, getting
+    new EngineCoreRequests and returning EngineCoreOutputs
 
-        * pushes EngineCoreRequests via input_socket
-        * pulls EngineCoreOutputs via output_socket
+    * pushes EngineCoreRequests via input_socket
+    * pulls EngineCoreOutputs via output_socket
 
-        * AsyncMPClient subclass for AsyncLLM usage
-        * SyncMPClient subclass for LLM usage
+    * AsyncMPClient subclass for AsyncLLM usage
+    * SyncMPClient subclass for LLM usage
     """
 
     def __init__(
@@ -566,6 +582,8 @@ def __init__(
     ):
         self.vllm_config = vllm_config
         self._renderer: BaseRenderer | None = renderer
+        self._effective_attention_block_sizes: set[int | None] = set()
+        self._kv_event_sources: dict[int, KVEventsConfig] = {}
 
         # ZMQ setup.
         sync_ctx = zmq.Context(io_threads=2)
@@ -803,10 +821,21 @@ def monitor_engine_cores():
     def _apply_ready_response(self, payload: bytes) -> None:
         """Decode an EngineCoreReadyResponse and sync any post-initialization
         config changes (e.g. auto-fitted max_model_len) back to the frontend."""
-        if not payload:
-            return
         vllm_config = self.vllm_config
-        response = msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse)
+        response = (
+            msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse)
+            if payload
+            else None
+        )
+        self._effective_attention_block_sizes.add(
+            response.effective_attention_block_size if response is not None else None
+        )
+        sizes = self._effective_attention_block_sizes
+        vllm_config.cache_config.effective_attention_block_size = (
+            next(iter(sizes)) if len(sizes) == 1 else None
+        )
+        if response is None:
+            return
         vllm_config.model_config.max_model_len = min(
             vllm_config.model_config.max_model_len, response.max_model_len
         )
@@ -843,6 +872,14 @@ def _apply_ready_response(self, payload: bytes) -> None:
             else:
                 assert response.dp_stats_address == self.stats_update_address
 
+        if response.kv_events_config is not None:
+            self._kv_event_sources[response.data_parallel_rank] = (
+                response.kv_events_config
+            )
+
+    def get_kv_event_sources(self) -> dict[int, KVEventsConfig]:
+        return dict(self._kv_event_sources)
+
 
 def _process_utility_output(
     output: UtilityOutput, utility_results: dict[int, AnyFuture]
@@ -1013,8 +1050,11 @@ def pin_lora(self, lora_id: int) -> bool:
     def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
         self.call_utility("sleep", level, mode)
 
-    def wake_up(self, tags: list[str] | None = None) -> None:
-        self.call_utility("wake_up", tags)
+    def release_kv_cache_memory(self) -> None:
+        self.call_utility("release_kv_cache_memory")
+
+    def wake_up(self, tags: list[str] | None = None) -> bool:
+        return self.call_utility("wake_up", tags)
 
     def is_sleeping(self) -> bool:
         return self.call_utility("is_sleeping")
@@ -1257,8 +1297,11 @@ async def reset_encoder_cache_async(self) -> None:
     async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:
         await self.call_utility_async("sleep", level, mode)
 
-    async def wake_up_async(self, tags: list[str] | None = None) -> None:
-        await self.call_utility_async("wake_up", tags)
+    async def release_kv_cache_memory_async(self) -> None:
+        await self.call_utility_async("release_kv_cache_memory")
+
+    async def wake_up_async(self, tags: list[str] | None = None) -> bool:
+        return await self.call_utility_async("wake_up", tags)
 
     async def is_sleeping_async(self) -> bool:
         return await self.call_utility_async("is_sleeping")
@@ -1711,6 +1754,12 @@ async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None:
             "Only ray DP backend supports scaling elastic EP"
         )
         parallel_config = self.vllm_config.parallel_config
+        if new_data_parallel_size > parallel_config.elastic_ep_max_dp_size:
+            raise ValueError(
+                f"Cannot scale to data_parallel_size {new_data_parallel_size}; "
+                "--elastic-ep-max-dp-size is "
+                f"{parallel_config.elastic_ep_max_dp_size}."
+            )
         num_experts = self.vllm_config.model_config.get_num_experts()
         num_physical_experts = (
             num_experts + parallel_config.eplb_config.num_redundant_experts
@@ -1736,8 +1785,7 @@ async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None:
         self._prepared_elastic_ep = new_data_parallel_size, num_redundant_experts
 
     def _eep_wait_for_setup_switch_complete(self) -> asyncio.Future:
-        """
-        Wait for core engines to switch to the new setup.
+        """Wait for core engines to switch to the new setup.
 
         In eep_process_engine_core_notification(), a dummy UtilityOutput with
         EEP_NOTIFICATION_CALL_ID will be set when RECONFIGURE_FINISHED
@@ -1916,6 +1964,8 @@ async def _commit_scale_down_elastic_ep(self, new_data_parallel_size: int) -> No
         self.lb_engines = self.lb_engines[:new_data_parallel_size]
         # Pending coordinator snapshots must use the surviving ranks too.
         self.engine_ranks_managed = self.engine_ranks_managed[:new_data_parallel_size]
+        for rank in range(new_data_parallel_size, cur_data_parallel_size):
+            self._kv_event_sources.pop(rank, None)
         removed_dp_size = cur_data_parallel_size - new_data_parallel_size
         pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size
         pause_futures = [
diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py
index f04fa30a1856..9658c030ef94 100644
--- a/vllm/v1/engine/detokenizer.py
+++ b/vllm/v1/engine/detokenizer.py
@@ -94,8 +94,7 @@ def __init__(self, request: EngineCoreRequest):
         self.output_text = ""
 
     def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None:
-        """
-        Update RequestState for the request_id by:
+        """Update RequestState for the request_id by:
             1) Detokenize the new token ids incrementally.
             2) Evaluate stop criteria.
 
@@ -149,7 +148,6 @@ def decode_next(self, next_token_id: int) -> str:
     def get_next_output_text(self, finished: bool, delta: bool) -> str:
         """If delta is True, only new text since the last call to
         this method is returned"""
-
         # We return the full output text if the sequence is finished.
         buffer_length = 0 if finished else self.stop_buffer_length
         if not delta:
diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py
index 4cc433cbfe38..369c31ca270a 100644
--- a/vllm/v1/engine/input_processor.py
+++ b/vllm/v1/engine/input_processor.py
@@ -2,7 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import time
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
+from functools import partial
 from typing import Any, Literal
 
 import vllm.envs as envs
@@ -29,8 +30,10 @@
 from vllm.tokenizers import TokenizerLike
 from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
 from vllm.utils.async_utils import make_async
+from vllm.utils.diffusion import validate_diffusion_sampling_params
 from vllm.utils.jsontree import json_iter_leaves
 from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.kv_hints import KvHintsEnvelope
 
 logger = init_logger(__name__)
 
@@ -51,12 +54,18 @@ def __init__(
         self.speculative_config = vllm_config.speculative_config
         self.structured_outputs_config = vllm_config.structured_outputs_config
         self.observability_config = vllm_config.observability_config
+        self.diffusion_config = vllm_config.diffusion_config
+        # Load the custom logits processor classes once; the returned callable
+        # runs their validate_params hooks per request at admission.
+        self.validate_logits_processors_params = (
+            self._build_logits_processors_params_validator()
+        )
 
         self.generation_config_fields = model_config.try_get_generation_config()
 
         self.renderer = renderer or renderer_from_config(vllm_config)
 
-        self.supports_mm_inputs = mm_registry.supports_multimodal_inputs(model_config)
+        self.supports_mm_inputs = model_config.supports_multimodal_inputs
         self.mm_encoder_cache_size = 0
         self.skip_prompt_length_check = False
         if self.supports_mm_inputs:
@@ -81,6 +90,28 @@ def tokenizer(self) -> TokenizerLike | None:
     def get_tokenizer(self) -> TokenizerLike:
         return self.renderer.get_tokenizer()
 
+    def _build_logits_processors_params_validator(
+        self,
+    ) -> Callable[[SamplingParams], None]:
+        """Load the custom logits processor classes and return the per-request
+        params validator for the active model runner."""
+        custom_logitsprocs = self.model_config.logits_processors
+        if self.vllm_config.use_v2_model_runner:
+            if self.model_config.runner_type == "pooling" and not custom_logitsprocs:
+                return lambda _: None
+
+            from vllm.v1.worker.gpu.sample.logits_processor import (
+                build_custom_logits_processors_params_validator,
+            )
+
+            return build_custom_logits_processors_params_validator(custom_logitsprocs)
+
+        from vllm.v1.sample.logits_processor import (
+            validate_logits_processors_parameters,
+        )
+
+        return partial(validate_logits_processors_parameters, custom_logitsprocs)
+
     def _validate_params(
         self,
         params: SamplingParams | PoolingParams,
@@ -101,6 +132,22 @@ def _validate_params(
                 self.tokenizer,
             )
 
+            self.validate_logits_processors_params(params)
+
+            if self.model_config.is_diffusion:
+                # Without --diffusion-config the served canvas is unknown here;
+                # the ids and the read-only normalisation are still checked.
+                validate_diffusion_sampling_params(
+                    params,
+                    canvas_length=(
+                        self.diffusion_config.canvas_length
+                        if self.diffusion_config is not None
+                        else None
+                    ),
+                    vocab_size=self.model_config.get_vocab_size(),
+                    async_scheduling=self.vllm_config.scheduler_config.async_scheduling,
+                )
+
             if self.model_config.return_sampling_mask:
                 if params.temperature <= 0:
                     raise ValueError(
@@ -206,8 +253,7 @@ def _get_mm_identifier(
         mm_hash: str,
         lora_request: LoRARequest | None,
     ) -> str:
-        """
-        When enable_tower_connector_lora is True, multi-modal embeddings
+        """When enable_tower_connector_lora is True, multi-modal embeddings
         vary depending on the LoRA request. Therefore, the mm_hash must be
         generated based on the LoRA request to prevent incorrect cache hits.
         """
@@ -292,6 +338,7 @@ def process_inputs(
         data_parallel_rank: int | None = None,
         resumable: bool = False,
         session_id: str | None = None,
+        kv_hints: KvHintsEnvelope | None = None,
     ) -> EngineCoreRequest:
         self._validate_params(params, supported_tasks)
         self._validate_lora(lora_request)
@@ -354,12 +401,28 @@ def process_inputs(
         if isinstance(params, SamplingParams):
             # TODO: can we avoid cloning here in multiproc case?
             sampling_params = params.clone()
+            prompt_len = length_from_prompt_token_ids_or_embeds(
+                prompt_token_ids, prompt_embeds
+            )
+            if not 0 <= sampling_params.routed_experts_prompt_start <= prompt_len:
+                raise VLLMValidationError(
+                    f"routed_experts_prompt_start must be between 0 and "
+                    f"the prompt length ({prompt_len}), inclusive.",
+                    parameter="routed_experts_prompt_start",
+                    value=sampling_params.routed_experts_prompt_start,
+                )
             # If unset max tokens, then generate up to the max_model_len.
             if sampling_params.max_tokens is None:
-                seq_len = length_from_prompt_token_ids_or_embeds(
-                    prompt_token_ids, prompt_embeds
+                sampling_params.max_tokens = (
+                    self.model_config.max_model_len - prompt_len
                 )
-                sampling_params.max_tokens = self.model_config.max_model_len - seq_len
+                # min_tokens is not checked while max_tokens is unset.
+                if sampling_params.min_tokens > sampling_params.max_tokens:
+                    raise VLLMValidationError(
+                        f"min_tokens must be less than or equal to "
+                        f"max_tokens={sampling_params.max_tokens}, got "
+                        f"{sampling_params.min_tokens}."
+                    )
 
             sampling_params.update_from_generation_config(
                 self.generation_config_fields,
@@ -368,12 +431,7 @@ def process_inputs(
             if self.tokenizer is not None:
                 sampling_params.update_from_tokenizer(self.tokenizer)
             if sampling_params.trace_decode_token_ids:
-                self._normalize_trace_replay_params(
-                    sampling_params,
-                    length_from_prompt_token_ids_or_embeds(
-                        prompt_token_ids, prompt_embeds
-                    ),
-                )
+                self._normalize_trace_replay_params(sampling_params, prompt_len)
         else:
             pooling_params = params.clone()
 
@@ -431,6 +489,7 @@ def process_inputs(
             trace_headers=trace_headers,
             resumable=resumable,
             session_id=session_id,
+            kv_hints=kv_hints,
         )
 
     def _validate_prompt_len(
@@ -500,6 +559,15 @@ def _validate_model_input(
         prompt_len = length_from_prompt_token_ids_or_embeds(prompt_ids, prompt_embeds)
         self._validate_prompt_len(prompt_len, prompt_type)
 
+        if prompt_input["type"] == "embeds":
+            is_token_ids = prompt_input.get("is_token_ids")
+            if is_token_ids is not None and len(is_token_ids) != prompt_len:
+                raise VLLMValidationError(
+                    "prompt_is_token_ids must have the same length as prompt_embeds "
+                    f"(expected {prompt_len}, got {len(is_token_ids)}).",
+                    parameter="prompt_is_token_ids",
+                )
+
         if prompt_input["type"] == "multimodal":
             decoder_mm_positions = prompt_input["mm_placeholders"]
             for modality, mm_positions in decoder_mm_positions.items():
diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py
index 32096079e066..c925ba90e8e0 100644
--- a/vllm/v1/engine/llm_engine.py
+++ b/vllm/v1/engine/llm_engine.py
@@ -34,6 +34,7 @@
 from vllm.v1.engine.output_processor import OutputProcessor
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.executor import Executor
+from vllm.v1.kv_hints import KvHintsEnvelope
 from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
 from vllm.v1.metrics.reader import Metric, get_metrics_snapshot
 from vllm.v1.metrics.stats import IterationStats
@@ -171,7 +172,6 @@ def from_engine_args(
         enable_multiprocessing: bool = False,
     ) -> "LLMEngine":
         """Creates an LLM engine from the engine arguments."""
-
         # Create the engine configs.
         vllm_config = engine_args.create_engine_config(usage_context)
         executor_class = Executor.get_class(vllm_config)
@@ -216,7 +216,6 @@ def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
 
     def abort_request(self, request_ids: list[str], internal: bool = False) -> None:
         """Remove request_ids from EngineCore and Detokenizer."""
-
         request_ids = self.output_processor.abort_requests(request_ids, internal)
         self.engine_core.abort_requests(request_ids)
 
@@ -232,6 +231,7 @@ def add_request(
         priority: int = 0,
         session_id: str | None = None,
         prompt_text: str | None = None,
+        kv_hints: KvHintsEnvelope | None = None,
     ) -> str:
         # Validate the request_id type.
         if not isinstance(request_id, str):
@@ -265,6 +265,7 @@ def add_request(
                 trace_headers=trace_headers,
                 priority=priority,
                 session_id=session_id,
+                kv_hints=kv_hints,
             )
             prompt_text, _, _ = extract_prompt_components(self.model_config, prompt)
 
@@ -377,10 +378,17 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort"):
         if self.logger_manager is not None:
             self.logger_manager.record_sleep_state(1, level)
 
-    def wake_up(self, tags: list[str] | None = None):
-        self.engine_core.wake_up(tags)
+    def release_kv_cache_memory(self) -> None:
+        self.renderer.clear_mm_cache()
+        self.engine_core.release_kv_cache_memory()
 
         if self.logger_manager is not None:
+            self.logger_manager.record_sleep_state(1, 0)
+
+    def wake_up(self, tags: list[str] | None = None):
+        fully_awake = self.engine_core.wake_up(tags)
+
+        if self.logger_manager is not None and fully_awake:
             self.logger_manager.record_sleep_state(0, 0)
 
     def is_sleeping(self) -> bool:
diff --git a/vllm/v1/engine/logprobs.py b/vllm/v1/engine/logprobs.py
index 8eef567b0715..ff483e128ee2 100644
--- a/vllm/v1/engine/logprobs.py
+++ b/vllm/v1/engine/logprobs.py
@@ -76,7 +76,6 @@ def _update_sample_logprobs(self, logprobs_lists: LogprobsLists) -> None:
           logprobs_lists: the lists of logprob tokens, logprobs, and ranks.
 
         """
-
         assert self.num_logprobs is not None
         assert self.logprobs is not None
         assert self.cumulative_logprob is not None
@@ -129,7 +128,6 @@ def _update_prompt_logprobs(
                                    tensors.
 
         """
-
         # Prompt logprobs are enabled.
         assert self.num_prompt_logprobs is not None
         assert self.prompt_logprobs is not None
@@ -187,7 +185,7 @@ def _update_prompt_logprobs(
             )
 
     def pop_prompt_logprobs(self) -> PromptLogprobs | None:
-        """Pop and return all request prompt logprobs
+        """Pop and return all request prompt logprobs.
 
         The logprobs processor aggregates prompt chunk logprobs
         over one or more prefill chunks. This method returns
@@ -199,6 +197,7 @@ def pop_prompt_logprobs(self) -> PromptLogprobs | None:
         Returns:
           None if prompt logprobs are disabled for this request.
           List of all prompt logprobs, otherwise.
+
         """
         plp = self.prompt_logprobs
         if plp:
@@ -223,6 +222,7 @@ def _get_sampled_context_ids(
 
         Returns:
             List of sampled token IDs, oldest first, most recent last.
+
         """
         if not logprobs_source:
             return []
@@ -265,6 +265,7 @@ def _correct_decoded_token(
         Returns:
             The corrected decoded string, or empty string if the byte
             sequence is genuinely incomplete at this point.
+
         """
         assert self.tokenizer is not None
 
@@ -325,6 +326,7 @@ def _verify_tokens(
             context_token_ids: Preceding sampled token IDs providing
                 sequential context. If None, extracted from
                 self.logprobs.
+
         """
         if context_token_ids is None:
             context_token_ids = self._get_sampled_context_ids(self.logprobs)
diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py
index 419e3aa8e99d..b4f7338048c7 100644
--- a/vllm/v1/engine/output_processor.py
+++ b/vllm/v1/engine/output_processor.py
@@ -5,7 +5,7 @@
 from collections import defaultdict, deque
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import torch
@@ -41,13 +41,15 @@
 )
 from vllm.v1.outputs import SamplingMaskLists
 
+if TYPE_CHECKING:
+    from vllm.v1.engine.admission_control import SharedAdmissionStats
+
 # shared empty CPU tensor used as a placeholder pooling output
 EMPTY_CPU_TENSOR = torch.empty(0, device="cpu")
 
 
 class RequestOutputCollector:
-    """
-    Collects streamed RequestOutputs per individual request,
+    """Collects streamed RequestOutputs per individual request,
     for hand-off to the consuming asyncio generate task.
 
     When streaming deltas, RequestOutputs are merged if the
@@ -152,6 +154,7 @@ def __init__(
         n: int | None = None,
         temperature: float | None = None,
         stream_input: bool = False,
+        remote_prefill_cached_tokens: int | None = None,
     ):
         self.request_id = request_id
         self.external_req_id = external_req_id
@@ -176,6 +179,7 @@ def __init__(
         self.queue = queue
         self.num_cached_tokens = 0
         self.num_cache_creation_tokens = 0
+        self.remote_prefill_cached_tokens = remote_prefill_cached_tokens
         # Per-sequence spec-decode accumulator; arrives once (on finish) via
         # EngineCoreOutput, then attached to this sequence's CompletionOutput.
         self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None
@@ -227,7 +231,18 @@ def from_new_request(
         log_stats: bool,
         stream_interval: int,
     ) -> "RequestState":
+        remote_prefill_cached_tokens = None
         if sampling_params := request.sampling_params:
+            # In a remote prefill scenario, report cached tokens
+            # as cache hit rate on the remote P worker.
+            if sampling_params.extra_args:
+                kv_transfer_params = sampling_params.extra_args.get(
+                    "kv_transfer_params"
+                )
+                if kv_transfer_params and kv_transfer_params.get("do_remote_prefill"):
+                    cached = kv_transfer_params.get("remote_prefill_cached_tokens")
+                    if isinstance(cached, int):
+                        remote_prefill_cached_tokens = cached
             if not sampling_params.detokenize:
                 tokenizer = None
             output_kind = sampling_params.output_kind
@@ -278,6 +293,7 @@ def from_new_request(
             log_stats=log_stats,
             stream_interval=stream_interval,
             stream_input=request.resumable,
+            remote_prefill_cached_tokens=remote_prefill_cached_tokens,
         )
 
     def make_request_output(
@@ -455,6 +471,7 @@ def __init__(
         log_stats: bool,
         stream_interval: int = 1,
         tracing_enabled: bool = False,
+        admission_stats: "SharedAdmissionStats | None" = None,
     ):
         self.log_stats = log_stats
         self.tokenizer = tokenizer
@@ -464,6 +481,7 @@ def __init__(
         self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list)
         self.lora_states = LoRARequestStates(log_stats)
         self.tracing_enabled = tracing_enabled
+        self.admission_stats = admission_stats
 
     def get_num_unfinished_requests(self):
         return len(self.request_states)
@@ -487,7 +505,6 @@ def has_unfinished_requests(self) -> bool:
 
     def propagate_error(self, e: Exception):
         """Propagate error to all generate() tasks."""
-
         for _, state in self.request_states.items():
             assert state.queue is not None
             state.queue.put(e)
@@ -553,6 +570,7 @@ def abort_requests(self, request_ids: Iterable[str], internal: bool) -> list[str
                     child_reqs = self.abort_requests(child_reqs, internal=True)
                     request_ids_to_abort.extend(child_reqs)
                 self.parent_requests.pop(request_id, None)
+        self._update_admission_stats()
         return request_ids_to_abort
 
     def add_request(
@@ -585,6 +603,7 @@ def add_request(
 
         # Track the external_req_id -> [internal_req_id, ...] mapping
         self.external_req_ids[req_state.external_req_id].append(request_id)
+        self._update_admission_stats()
 
     def _update_streaming_request_state(
         self, req_state: RequestState, request: EngineCoreRequest, prompt: str | None
@@ -625,8 +644,7 @@ def process_outputs(
         engine_core_timestamp: float | None = None,
         iteration_stats: IterationStats | None = None,
     ) -> OutputProcessorOutput:
-        """
-        Process the EngineCoreOutputs:
+        """Process the EngineCoreOutputs:
         1) Compute stats for logging
         2) Detokenize
         3) Create and handle RequestOutput objects:
@@ -646,7 +664,6 @@ def process_outputs(
         If you need to touch every element of the batch, do it from
         within the loop below.
         """
-
         request_outputs: list[RequestOutput | PoolingRequestOutput] = []
         reqs_to_abort: list[str] = []
         for engine_core_output in engine_core_outputs:
@@ -680,6 +697,8 @@ def process_outputs(
                     req_state.num_cache_creation_tokens = (
                         engine_core_output.prefill_stats.num_cache_creation_tokens
                     )
+                if req_state.remote_prefill_cached_tokens is not None:
+                    req_state.num_cached_tokens = req_state.remote_prefill_cached_tokens
                 req_state.is_prefilling = False
 
             if engine_core_output.spec_decode_metrics is not None:
@@ -763,6 +782,11 @@ def _finish_request(self, req_state: RequestState) -> None:
         parent_req = req_state.parent_req
         if parent_req and not parent_req.child_requests:
             self.parent_requests.pop(parent_req.request_id, None)
+        self._update_admission_stats()
+
+    def _update_admission_stats(self) -> None:
+        if self.admission_stats is not None:
+            self.admission_stats.set_num_requests(self.get_num_unfinished_requests())
 
     def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None):
         self.lora_states.update_scheduler_stats(scheduler_stats)
diff --git a/vllm/v1/engine/parallel_sampling.py b/vllm/v1/engine/parallel_sampling.py
index 8eb6fa057d37..9baf785903a1 100644
--- a/vllm/v1/engine/parallel_sampling.py
+++ b/vllm/v1/engine/parallel_sampling.py
@@ -53,7 +53,7 @@ def _get_child_sampling_params(
         self,
         index: int,
     ) -> SamplingParams:
-        """Efficiently obtain child `sampling_params`
+        """Efficiently obtain child `sampling_params`.
 
         If `sampling_params.seed` is not `None` then
         each child request requires a unique clone of
@@ -64,6 +64,7 @@ def _get_child_sampling_params(
 
         Returns:
           Child `sampling_params` instance.
+
         """
         seed = self.sampling_params.seed
         if self.cached_child_sampling_params:
@@ -88,6 +89,7 @@ def get_child_info(self, index: int) -> tuple[str, SamplingParams]:
 
         Returns:
           (request ID, sampling_params) tuple
+
         """
         child_req_id = f"{index}_{self.request_id}"
         self.child_requests.add(child_req_id)
diff --git a/vllm/v1/engine/tensor_ipc.py b/vllm/v1/engine/tensor_ipc.py
index 2a8391b17e6d..2b8fbadd97a5 100644
--- a/vllm/v1/engine/tensor_ipc.py
+++ b/vllm/v1/engine/tensor_ipc.py
@@ -29,8 +29,7 @@
 
 @dataclasses.dataclass
 class TensorIpcData:
-    """
-    Data sent via torch.multiprocessing.Queue for zero-copy IPC.
+    """Data sent via torch.multiprocessing.Queue for zero-copy IPC.
 
     Contains the tensor_id and the actual tensor. The tensor is
     shared in memory (GPU or CPU) for efficient inter-process communication.
@@ -130,7 +129,6 @@ def __call__(
         the queue, buffering them, until the requested tensor is found.
         Works for CUDA and CPU.
         """
-
         # Create lookup key from handle
         sender_id: str = meta["sender_id"]
         message_id: int = meta["message_id"]
diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py
index 9bce2b7687ef..c1002ba1dc5b 100644
--- a/vllm/v1/engine/utils.py
+++ b/vllm/v1/engine/utils.py
@@ -142,8 +142,7 @@ def _node_ip_from_resources(node_resources: dict) -> str | None:
 
 
 class CoreEngineProcManager:
-    """
-    Utility class to handle creation, readiness, and shutdown
+    """Utility class to handle creation, readiness, and shutdown
     of background processes used by the AsyncLLM and LLMEngine.
     """
 
@@ -255,7 +254,6 @@ def shutdown(self, timeout: float | None = None) -> None:
 
     def monitor_engine_liveness(self) -> None:
         """Monitor engine core process liveness."""
-
         sentinel_to_proc = {proc.sentinel: proc for proc in self.processes}
         sentinels = set(sentinel_to_proc.keys())
 
@@ -316,8 +314,7 @@ def set_assigned_physical_gpu_ids_for_dp_rank(
     local_dp_rank: int,
     user_assigned_gpu_ids: list[int] | None = None,
 ) -> None:
-    """
-    Populate assigned_physical_gpu_ids on the config for the given DP rank.
+    """Populate assigned_physical_gpu_ids on the config for the given DP rank.
 
     user_assigned_gpu_ids is the full (un-sharded) --device-ids list, if the
     user provided one; this DP rank's shard is sliced from it. It is passed
@@ -345,8 +342,7 @@ def get_physical_gpu_ids_for_local_dp_rank(
     local_world_size: int | None = None,
     user_assigned_gpu_ids: list[int] | None = None,
 ) -> list[int]:
-    """
-    Returns list of physical GPU IDs for the specified
+    """Returns list of physical GPU IDs for the specified
     data parallel rank.
 
     For example, if world_size=2 and local_dp_rank=1, and there are 4 devices,
@@ -399,8 +395,7 @@ def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None:
 
 
 class CoreEngineActorManager:
-    """
-    Utility class to handle creation, readiness, and shutdown
+    """Utility class to handle creation, readiness, and shutdown
     of core engine Ray actors used by the AsyncLLM and LLMEngine.
 
     Different from CoreEngineProcManager, this class manages
@@ -552,10 +547,7 @@ def __init__(
     def create_dp_placement_groups(
         vllm_config: VllmConfig,
     ) -> tuple[list["PlacementGroup"], list[int]]:
-        """
-        Create placement groups for data parallel.
-        """
-
+        """Create placement groups for data parallel."""
         import ray
         from ray._private.state import available_resources_per_node
 
@@ -769,9 +761,7 @@ def create_dp_placement_groups(
     def add_dp_placement_groups(
         old_vllm_config: VllmConfig, new_data_parallel_size: int
     ) -> tuple[list["PlacementGroup"], list[int]]:
-        """
-        Add placement groups for new data parallel size.
-        """
+        """Add placement groups for new data parallel size."""
         import ray
         from ray._private.state import (
             available_resources_per_node,
@@ -1108,7 +1098,6 @@ def launch_core_engines(
     addresses: EngineZmqAddresses,
 ) -> Iterator[CoreEngineLaunch]:
     """Launch engine and DP coordinator processes as needed."""
-
     parallel_config = vllm_config.parallel_config
     dp_size = parallel_config.data_parallel_size
     local_engine_count = parallel_config.data_parallel_size_local
diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py
index 964f16ed2421..f234e943b4fd 100644
--- a/vllm/v1/executor/abstract.py
+++ b/vllm/v1/executor/abstract.py
@@ -32,11 +32,14 @@
 
 _R = TypeVar("_R")
 
+# Memory-pool tags managed by sleep/wake_up/discard.
+SLEEP_TAGS = frozenset(("weights", "kv_cache"))
+
 FailureCallback = Callable[[], None]
 
 
 class Executor(ABC):
-    """Abstract base class for vLLM executors."
+    """Abstract base class for vLLM executors.
 
     An executor is responsible for executing the model on one device,
     or it can be a distributed executor that can execute the model on multiple devices.
@@ -108,7 +111,6 @@ def __init__(
         self.speculative_config = vllm_config.speculative_config
         self.observability_config = vllm_config.observability_config
         self._init_executor()
-        self.is_sleeping = False
         self.sleeping_tags: set[str] = set()
         self.kv_output_aggregator: KVOutputAggregator | None = None
         self.ec_output_aggregator: ECOutputAggregator | None = None
@@ -139,8 +141,7 @@ def compile_or_warm_up_model(self) -> None:
             )
 
     def register_failure_callback(self, callback: FailureCallback):  # noqa: B027
-        """
-        Register a function to be called if the executor enters a permanent
+        """Register a function to be called if the executor enters a permanent
         failed state.
         """
         pass
@@ -168,8 +169,29 @@ def collective_rpc(
         kwargs: dict | None = None,
         non_block: Literal[False] = False,
     ) -> list[_R]:
-        """
-        Execute an RPC call on all workers.
+        pass
+
+    @overload
+    def collective_rpc(
+        self,
+        method: str | Callable[[WorkerBase], _R],
+        timeout: float | None = None,
+        args: tuple = (),
+        kwargs: dict | None = None,
+        non_block: Literal[True] = True,
+    ) -> Future[list[_R]]:
+        pass
+
+    @abstractmethod
+    def collective_rpc(
+        self,
+        method: str | Callable[[WorkerBase], _R],
+        timeout: float | None = None,
+        args: tuple = (),
+        kwargs: dict | None = None,
+        non_block: bool = False,
+    ) -> list[_R] | Future[list[_R]]:
+        """Execute an RPC call on all workers.
 
         Args:
             method: Name of the worker method to execute, or a callable that
@@ -191,24 +213,8 @@ def collective_rpc(
         Note:
             It is recommended to use this API to only pass control messages,
             and set up data-plane communication to pass data.
-        """
-        pass
 
-    @overload
-    def collective_rpc(
-        self,
-        method: str | Callable[[WorkerBase], _R],
-        timeout: float | None = None,
-        args: tuple = (),
-        kwargs: dict | None = None,
-        non_block: Literal[True] = True,
-    ) -> Future[list[_R]]:
-        pass
-
-    @abstractmethod
-    def collective_rpc(
-        self, method, timeout=None, args=(), kwargs=None, non_block: bool = False
-    ):
+        """
         raise NotImplementedError
 
     def get_kv_connector_handshake_metadata(
@@ -288,7 +294,7 @@ def shutdown(self) -> None:
         self.collective_rpc("shutdown")
 
     def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None:
-        """Init KVOutputAggregator"""
+        """Init KVOutputAggregator."""
         self.kv_output_aggregator = KVOutputAggregator.from_connector(
             connector, self.parallel_config.world_size
         )
@@ -334,15 +340,18 @@ def reset_encoder_cache(self) -> None:
         """Reset the encoder cache in each worker to clear cached encoder outputs."""
         self.collective_rpc("reset_encoder_cache")
 
+    @property
+    def is_sleeping(self) -> bool:
+        return bool(self.sleeping_tags)
+
     def sleep(self, level: int = 1):
-        if self.is_sleeping:
+        if "weights" in self.sleeping_tags:
             logger.warning("Executor is already sleeping.")
             return
         time_before_sleep = time.perf_counter()
         self.collective_rpc("sleep", kwargs=dict(level=level))
         time_after_sleep = time.perf_counter()
-        self.sleeping_tags = {"weights", "kv_cache"}
-        self.is_sleeping = True
+        self.sleeping_tags |= SLEEP_TAGS
         logger.info(
             "It took %.6f seconds to fall asleep.", time_after_sleep - time_before_sleep
         )
@@ -371,8 +380,23 @@ def wake_up(self, tags: list[str] | None = None):
                 self.sleeping_tags.remove(tag)
         else:
             self.sleeping_tags.clear()
-        if not self.sleeping_tags:
-            self.is_sleeping = False
+
+    def discard(self, tags: tuple[str, ...]) -> None:
+        tags_to_discard = set(tags) - self.sleeping_tags
+        if not tags_to_discard:
+            logger.warning("Tags %s are already sleeping.", tags)
+            return
+        time_before_discard = time.perf_counter()
+        try:
+            self.collective_rpc("discard", args=(tuple(tags_to_discard),))
+        finally:
+            self.sleeping_tags |= tags_to_discard
+        time_after_discard = time.perf_counter()
+        logger.info(
+            "It took %.6f seconds to discard tags %s.",
+            time_after_discard - time_before_discard,
+            tags_to_discard,
+        )
 
     def reinitialize_distributed(
         self, reconfig_request: ReconfigureDistributedRequest
@@ -381,9 +405,7 @@ def reinitialize_distributed(
 
     @classmethod
     def supports_async_scheduling(cls) -> bool:
-        """
-        Whether the executor supports async scheduling.
-        """
+        """Whether the executor supports async scheduling."""
         return False
 
 
diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py
index 67a631d3ce20..dbaefe218665 100644
--- a/vllm/v1/executor/multiproc_executor.py
+++ b/vllm/v1/executor/multiproc_executor.py
@@ -500,7 +500,7 @@ def wait_for_termination(procs, timeout):
                 p.kill()
 
     def shutdown(self):
-        """Properly shut down the executor and its workers"""
+        """Properly shut down the executor and its workers."""
         if not getattr(self, "shutting_down", False):
             worker_count = len(getattr(self, "workers", None) or [])
             logger.debug(
@@ -853,7 +853,6 @@ def death_pipe_monitor(queues_to_shutdown: list[MessageQueue]):
     def worker_main(*args, **kwargs):
         """Worker initialization and execution loops.
         This runs a background process"""
-
         # Signal handler used for graceful termination.
         # SystemExit exception is only raised once to allow this and worker
         # processes to terminate without error
@@ -1010,7 +1009,6 @@ def handle_output(self, output: Any):
 
     def async_output_busy_loop(self):
         """Entrypoint for the thread which handles outputs asynchronously."""
-
         # set device to the worker device for the thread.
         # a thread will not inherit the context of the main thread.
         # when calling any cuda runtime functions, it will implicitly
@@ -1027,7 +1025,7 @@ def async_output_busy_loop(self):
             self.enqueue_output(output)
 
     def worker_busy_loop(self):
-        """Main busy loop for Multiprocessing Workers"""
+        """Main busy loop for Multiprocessing Workers."""
         assert self.rpc_broadcast_mq is not None
         while True:
             self._execute_worker_rpc(self.rpc_broadcast_mq.dequeue(indefinite=True))
@@ -1099,7 +1097,6 @@ def set_multiprocessing_worker_envs(local_world_size: int = 1):
     """Set up environment variables that should be used when there are workers
     in a multiprocessing environment. This should be called by the parent
     process before worker processes are created"""
-
     _maybe_force_spawn()
 
     if current_platform.is_cpu() or "OMP_NUM_THREADS" in os.environ:
diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py
index b2c385a862b8..d101885ce2b6 100644
--- a/vllm/v1/executor/ray_executor.py
+++ b/vllm/v1/executor/ray_executor.py
@@ -52,8 +52,7 @@
 
 @dataclass
 class RayWorkerMetaData:
-    """
-    Metadata for a Ray worker.
+    """Metadata for a Ray worker.
     The order of ray worker creation can be random,
     and we need to reset the rank after creating all workers.
     """
@@ -65,7 +64,7 @@ class RayWorkerMetaData:
 
 
 class RayDistributedExecutor(Executor):
-    """Ray-based distributed executor"""
+    """Ray-based distributed executor."""
 
     uses_ray: bool = True
     supports_pp: bool = True
@@ -249,8 +248,7 @@ def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwar
             ip_counts[ip] = ip_counts.get(ip, 0) + 1
 
         def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
-            """
-            Sort the workers based on 3 properties:
+            """Sort the workers based on 3 properties:
             1. If the worker is on the same node as the driver (vllm engine),
                 it should be placed first.
             2. Then, if the worker is on a node with fewer workers, it should
@@ -439,6 +437,7 @@ def sample_tokens(  # type: ignore[override]
 
         Returns:
             The model runner output.
+
         """
         scheduler_output = self.scheduler_output
         if scheduler_output is None:
diff --git a/vllm/v1/executor/ray_executor_v2.py b/vllm/v1/executor/ray_executor_v2.py
index 8f4b84ed0f13..b8a2eec73f77 100644
--- a/vllm/v1/executor/ray_executor_v2.py
+++ b/vllm/v1/executor/ray_executor_v2.py
@@ -72,7 +72,7 @@ class RayWorkerHandle:
     """run() ObjectRef used as a sentinel for health monitoring"""
 
     def run(self):
-        """Start the worker's busy loop"""
+        """Start the worker's busy loop."""
         self.run_ref = self.actor.run.remote()
 
 
@@ -190,8 +190,7 @@ def initialize_worker(
     def _init_message_queues(
         self, input_shm_handle: Handle, vllm_config: VllmConfig
     ) -> None:
-        """
-        Workers on the same node as the executor use shared memory for
+        """Workers on the same node as the executor use shared memory for
         both the broadcast (input) MQ and the response MQ. Workers on
         different nodes use TCP (n_local_reader=0).
         """
diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py
index 71d4cf74717e..a5edb67dee85 100644
--- a/vllm/v1/executor/ray_utils.py
+++ b/vllm/v1/executor/ray_utils.py
@@ -68,8 +68,7 @@ def __init__(self, *args, **kwargs) -> None:
         rpc_rank: int
 
         def adjust_rank(self, rank_mapping: dict[int, int]) -> None:
-            """
-            Adjust the rpc_rank based on the given mapping.
+            """Adjust the rpc_rank based on the given mapping.
             It is only used during the initialization of the executor,
             to adjust the rpc_rank of workers after we create all workers.
             """
@@ -234,14 +233,10 @@ def _copy_if_readonly(arr):
                 token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens
             )
 
-    if output.routed_experts is not None:
-        routing_data, slot_mapping = output.routed_experts
-        routing_data_c = _copy_if_readonly(routing_data)
-        slot_mapping_c = _copy_if_readonly(slot_mapping)
-        if routing_data_c is not routing_data or slot_mapping_c is not slot_mapping:
-            output.routed_experts = type(output.routed_experts)(
-                routing_data_c, slot_mapping_c
-            )
+    aux_output = output.aux_output_connector_output
+    if aux_output is not None:
+        for request_output in aux_output.values():
+            request_output.rows = _copy_if_readonly(request_output.rows)
 
 
 class FutureWrapper(Future):
@@ -363,8 +358,7 @@ def get_bundles_for_indices(
     bundle_indices: list[int],
     world_size: int,
 ) -> list[tuple[int, str, str]]:
-    """
-    Return GPU bundle indices paired with node IDs and node IPs for
+    """Return GPU bundle indices paired with node IDs and node IPs for
     explicit bundle indices specified via VLLM_RAY_BUNDLE_INDICES.
     """
     assert len(bundle_indices) == world_size, (
@@ -391,8 +385,7 @@ def get_bundles_for_indices(
 def get_bundles_sorted_by_node(
     placement_group: "PlacementGroup",
 ) -> list[tuple[int, str, str]]:
-    """
-    Return GPU bundle indices paired with node IDs and node IPs,
+    """Return GPU bundle indices paired with node IDs and node IPs,
     sorted driver-first.
 
     This utility has to be invoked from the driver node.
@@ -533,6 +526,7 @@ def initialize_ray_cluster(
             on the current (driver) node and pin the first PG bundle to it.
             Set to False for executors like RayExecutorV2 where all GPU work
             is delegated to remote Ray actors.
+
     """
     assert_ray_available()
     from vllm.platforms import current_platform
diff --git a/vllm/v1/hisparse/__init__.py b/vllm/v1/hisparse/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/v1/hisparse/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py
new file mode 100644
index 000000000000..716211fbb30a
--- /dev/null
+++ b/vllm/v1/hisparse/binding.py
@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Worker-side wiring of HiSparse caches to attention layers."""
+
+from collections.abc import Mapping
+from copy import copy
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.v1.hisparse.layout import (
+    HISPARSE_HOT_SUFFIX,
+    HISPARSE_RESIDENT_SUFFIX,
+)
+from vllm.v1.hisparse.runtime import (
+    HiSparseCacheHandle,
+    HiSparseHostPool,
+    initialize_hisparse_runtime_buffers,
+    release_pinned_state,
+)
+from vllm.v1.kv_cache_interface import (
+    HiSparseHotSpec,
+    HiSparseResidentSpec,
+    KVCacheConfig,
+    KVCacheLayout,
+    KVCacheSpec,
+    MLAAttentionSpec,
+    UniformTypeKVCacheSpecs,
+    create_kv_cache_views,
+)
+from vllm.v1.worker.utils import allocate_kv_cache, select_common_block_size
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+    from vllm.v1.worker.gpu.block_table import BlockTables
+
+
+def resolve_hisparse_block_size(
+    vllm_config: VllmConfig,
+    kv_cache_spec: dict[str, KVCacheSpec],
+    attn_layers: Mapping[str, "AttentionLayerBase"],
+) -> None:
+    """Resolve a common kernel block size in-place for HiSparse MLA specs."""
+    if vllm_config.attention_config.hisparse_config is None:
+        return
+    mla_specs = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if isinstance(spec, MLAAttentionSpec)
+    }
+    if not mla_specs:
+        return
+    block_sizes = {spec.block_size for spec in mla_specs.values()}
+    if len(block_sizes) != 1:
+        raise ValueError("HiSparse requires one scheduler block size.")
+    backends = [attn_layers[name].get_attn_backend() for name in mla_specs]
+    try:
+        block_size = select_common_block_size(block_sizes.pop(), backends)
+    except ValueError as error:
+        raise ValueError(
+            "HiSparse requires a GPU block size supported by every sparse "
+            f"attention and indexer backend: {error}"
+        ) from error
+    kv_cache_spec.update(
+        (name, spec.copy_with_new_block_size(block_size))
+        for name, spec in mla_specs.items()
+    )
+
+
+def allocate_hisparse_kv_caches(
+    kv_cache_config: KVCacheConfig,
+    device: torch.device,
+    layout: KVCacheLayout,
+    kernel_block_sizes: list[int],
+    host_pool: HiSparseHostPool,
+) -> dict[str, torch.Tensor]:
+    """Allocate the host pool separately from the shared device backing."""
+    device_config = copy(kv_cache_config)
+    device_config.kv_cache_tensors = [
+        tensor
+        for tensor in kv_cache_config.kv_cache_tensors
+        if not tensor.host_resident
+    ]
+    kv_caches = allocate_kv_cache(device_config, device, layout, kernel_block_sizes)
+    host_tensors = [
+        tensor for tensor in kv_cache_config.kv_cache_tensors if tensor.host_resident
+    ]
+    (host_size,) = {tensor.size for tensor in host_tensors}
+    backing = host_pool.allocate(host_size)
+    (host_group_id,) = kv_cache_config.host_group_ids
+    host_spec = kv_cache_config.kv_cache_groups[host_group_id].kv_cache_spec
+    for tensor in host_tensors:
+        spec = (
+            host_spec.kv_cache_specs[tensor.layers[0]]
+            if isinstance(host_spec, UniformTypeKVCacheSpecs)
+            else host_spec
+        )
+        kernel_block_size = kernel_block_sizes[host_group_id]
+        if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None:
+            kernel_block_size = spec.storage_block_size
+        views = create_kv_cache_views(
+            backing,
+            spec,
+            kv_cache_config.num_blocks_of(tensor),
+            layout,
+            tensor,
+            kernel_block_size=kernel_block_size,
+        )
+        kv_caches.update(zip(tensor.layers, views))
+    return kv_caches
+
+
+def init_hisparse_kv_cache(
+    kv_cache_config: KVCacheConfig,
+    device: torch.device,
+    kernel_block_sizes: list[int],
+    vllm_config: VllmConfig,
+    forward_context: dict[str, Any],
+    block_tables: "BlockTables",
+) -> dict[str, torch.Tensor]:
+    """Allocate and bind HiSparse caches within the caller's allocation context."""
+    host_pool = HiSparseHostPool(vllm_config, kv_cache_config)
+    try:
+        kv_caches = allocate_hisparse_kv_caches(
+            kv_cache_config,
+            device,
+            vllm_config.cache_config.get_resolved_kv_cache_layout(),
+            kernel_block_sizes,
+            host_pool,
+        )
+        cache_handles = bind_hisparse_kv_caches(
+            forward_context=forward_context,
+            kv_cache_config=kv_cache_config,
+            kv_caches=kv_caches,
+            block_tables=block_tables,
+            host_pool=host_pool,
+        )
+        initialize_hisparse_runtime_buffers(
+            cache_handles,
+            max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
+            max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
+        )
+        return kv_caches
+    except Exception:
+        if host_pool.shared_region is not None:
+            host_pool.shared_region.cleanup()
+        raise
+
+
+def _get_hisparse_cache(
+    forward_context: dict[str, Any], layer_name: str
+) -> HiSparseCacheHandle:
+    attention_layer = forward_context[layer_name]
+    hisparse_cache = attention_layer.hisparse_cache
+    assert hisparse_cache is not None
+    return hisparse_cache
+
+
+def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None:
+    cache_handles = [
+        cache
+        for layer in forward_context.values()
+        if (cache := getattr(layer, "hisparse_cache", None)) is not None
+    ]
+    runtimes = {
+        id(cache.runtime): cache.runtime
+        for cache in cache_handles
+        if hasattr(cache.runtime, "_host_cache")
+    }
+    if not runtimes:
+        return
+
+    shared_regions = {
+        id(runtime.shared_host_region): runtime.shared_host_region
+        for runtime in runtimes.values()
+        if runtime.shared_host_region is not None
+    }
+    assert len(shared_regions) <= 1
+    shared_region = next(iter(shared_regions.values()), None)
+    registered_pools = (
+        []
+        if shared_region is not None
+        else list(
+            {
+                runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool
+                for runtime in runtimes.values()
+            }.values()
+        )
+    )
+    release_pinned_state(list(runtimes.values()), registered_pools, shared_region)
+    for cache in cache_handles:
+        cache.mirror_staging_cache = None
+        cache.mirror_staging_slots = None
+
+
+def bind_hisparse_kv_caches(
+    *,
+    forward_context: dict[str, Any],
+    kv_cache_config: KVCacheConfig,
+    kv_caches: dict[str, torch.Tensor],
+    block_tables: "BlockTables",
+    host_pool: HiSparseHostPool,
+) -> list[HiSparseCacheHandle]:
+    """Bind existing cache storage and block tables; return the bound handles."""
+    assert host_pool.registered is not None
+    tensor_configs = {
+        name: (
+            tensor_config,
+            tensor_config.offset + layer_index * tensor_config.layer_stride,
+        )
+        for tensor_config in kv_cache_config.kv_cache_tensors
+        for layer_index, name in enumerate(tensor_config.layers)
+    }
+    resident_source_index = 0
+    for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
+        if not isinstance(group.kv_cache_spec, HiSparseResidentSpec):
+            continue
+        for cache_name in group.layer_names:
+            assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX)
+            layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)]
+            tensor_config, byte_offset = tensor_configs[cache_name]
+            assert not tensor_config.host_resident
+            cache_handle = _get_hisparse_cache(forward_context, layer_name)
+            cache_handle.bind_cache(
+                kv_caches[cache_name],
+                byte_offset=byte_offset,
+                block_stride=tensor_config.block_stride,
+                num_blocks=kv_cache_config.num_blocks,
+                block_size=group.kv_cache_spec.block_size,
+                block_table=block_tables.input_block_tables[group_id],
+                slot_mapping=block_tables.slot_mappings[group_id],
+            )
+            assert cache_handle.view is not None
+            kv_caches[cache_name] = cache_handle.view.cache
+            cache_handle.runtime.resident_source_index = resident_source_index
+        resident_source_index += 1
+
+    hot_backing: torch.Tensor | None = None
+    cache_handles: list[HiSparseCacheHandle] = []
+    for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
+        if not isinstance(group.kv_cache_spec, HiSparseHotSpec):
+            continue
+        for cache_name in group.layer_names:
+            assert cache_name.endswith(HISPARSE_HOT_SUFFIX)
+            raw_tensor = kv_caches[cache_name]
+            if hot_backing is None:
+                hot_backing = raw_tensor
+            elif hot_backing.untyped_storage().data_ptr() != (
+                raw_tensor.untyped_storage().data_ptr()
+            ):
+                raise RuntimeError("HiSparse hot tensors must share one GPU backing.")
+            layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)]
+            cache_handle = _get_hisparse_cache(forward_context, layer_name)
+            tensor_config, byte_offset = tensor_configs[cache_name]
+            assert not tensor_config.host_resident
+            cache_handle.runtime.bind_hot_cache(
+                raw_tensor,
+                byte_offset=byte_offset,
+                block_stride=tensor_config.block_stride,
+                num_blocks=kv_cache_config.num_blocks,
+                block_size=group.kv_cache_spec.block_size,
+                block_table=block_tables.input_block_tables[group_id],
+            )
+            resident = cache_handle.view
+            hot = cache_handle.runtime.hot
+            assert resident is not None
+            if (
+                resident.cache.untyped_storage().data_ptr()
+                != hot.cache.untyped_storage().data_ptr()
+                or resident.cache.data_ptr() != hot.cache.data_ptr()
+                or resident.cache.stride() != hot.cache.stride()
+            ):
+                raise RuntimeError("HiSparse resident and hot layouts must match.")
+            source_cache = kv_caches[layer_name]
+            assert source_cache.untyped_storage().data_ptr() == (
+                host_pool.registered.untyped_storage().data_ptr()
+            )
+            cache_handle.runtime.shared_host_region = host_pool.shared_region
+            cache_handle.runtime.bind_source_cache(
+                source_cache,
+                registered_host_pool=host_pool.registered,
+            )
+            cache_handles.append(cache_handle)
+            # The hot slab is raw storage owned by the runtime, not a layer
+            # cache: nothing downstream binds or registers it.
+            del kv_caches[cache_name]
+
+    if hot_backing is None or not cache_handles:
+        raise RuntimeError("HiSparse found no hot-cache handles.")
+    (source_group_id,) = kv_cache_config.host_group_ids
+    for cache_handle in cache_handles:
+        cache_handle.source_block_table = block_tables.input_block_tables[
+            source_group_id
+        ]
+        cache_handle.mirror_slot_mapping = block_tables.slot_mappings[source_group_id]
+    return cache_handles
diff --git a/vllm/v1/hisparse/block_pool.py b/vllm/v1/hisparse/block_pool.py
new file mode 100644
index 000000000000..cbb620d64769
--- /dev/null
+++ b/vllm/v1/hisparse/block_pool.py
@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.distributed.kv_events import KVCacheEvent
+from vllm.v1.core.block_pool import BlockPool
+
+
+class SharedEventQueueBlockPool(BlockPool):
+    """A pool that publishes into another pool's live KV event queue.
+
+    The owner rebinds its queue on every drain, so this reads it through the
+    owner rather than holding a reference.
+    """
+
+    def __init__(self, *args, event_owner: BlockPool, **kwargs) -> None:
+        self._event_owner = event_owner
+        super().__init__(*args, **kwargs)
+
+    @property
+    def kv_event_queue(self) -> list[KVCacheEvent]:
+        return self._event_owner.kv_event_queue
+
+    @kv_event_queue.setter
+    def kv_event_queue(self, events: list[KVCacheEvent]) -> None:
+        # ``BlockPool.__init__`` seeds an empty queue and ``take_events``
+        # swaps in a fresh one; both belong to the owner, which drains it.
+        assert not events
diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py
new file mode 100644
index 000000000000..056e19479a33
--- /dev/null
+++ b/vllm/v1/hisparse/coordinator.py
@@ -0,0 +1,778 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from vllm.utils.math_utils import cdiv
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_utils import (
+    BlockHashWithGroupId,
+    KVCacheBlock,
+    KVCacheBlockCopy,
+)
+from vllm.v1.core.single_type_kv_cache_manager import (
+    HiSparseHotManager,
+    HiSparseResidentManager,
+    HiSparseSourceManager,
+    SingleTypeKVCacheManager,
+)
+from vllm.v1.hisparse.types import (
+    SparseKVOffloadCommand,
+    SparseKVPageTransfer,
+    SparseKVRowMirror,
+)
+from vllm.v1.kv_cache_interface import (
+    HiSparseResidentSpec,
+    KVCacheConfig,
+)
+from vllm.v1.request import Request
+
+if TYPE_CHECKING:
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+# Sealed pages this many positions behind the block-table tail stay pinned so
+# a page written by an in-flight step is never handed out under it.
+_ACTIVE_TAIL_PAGES = 2
+
+
+@dataclass
+class _PendingPublication:
+    request: Request
+    num_computed_tokens: int
+    num_pages: int
+    retention_interval: int | None
+    replay_boundaries: Sequence[int]
+
+
+@dataclass
+class _HiSparseRequestState:
+    """Residency bookkeeping for one request.
+
+    A resident page moves through three states: dirty (only the GPU copy
+    exists), clean (``valid_pages``: the host copy is complete) and unpinned
+    (``unpinned_pages``: clean, and its allocation reference was released so
+    the pool may reuse it). ``pinned_clean`` tracks clean pages still holding
+    their reference, which only happens before the request can read from host.
+    """
+
+    valid_pages: set[int] = field(default_factory=set)
+    ready_prefix_pages: int = 0
+    pending_pages: dict[int, int] = field(default_factory=dict)
+    publication: _PendingPublication | None = None
+    copies_recorded_blocks: int = 0
+    pinned_clean: set[int] = field(default_factory=set)
+    unpinned_pages: set[int] = field(default_factory=set)
+
+
+@dataclass
+class _PendingSpill:
+    transfer_id: int
+    page: tuple[str, int]
+    request_state: _HiSparseRequestState
+    host_block: KVCacheBlock
+    resident_blocks: tuple[KVCacheBlock, ...]
+    restore: bool = False
+    expected_worker_completions: int = 0
+    worker_completions: int = 0
+    enqueue_applied: bool = False
+
+
+class HiSparseCoordinator:
+    """Own HiSparse host allocation, host publication, spills, and GPU residency.
+
+    GPU-resident pages are a write-back cache of the host tier. Once a page's
+    host copy is durable and its owner can read from host (it has, or has
+    asked for, a hot region), the page's allocation reference is released with
+    ``BlockPool.unpin_blocks``: the block stays in the block table and keeps
+    being read, but the pool counts it as free and may hand it out, at which
+    point the owner's page is nulled and its block table republished.
+    """
+
+    def __init__(
+        self,
+        kv_cache_config: KVCacheConfig,
+        managers: tuple[SingleTypeKVCacheManager, ...],
+        max_model_len: int,
+    ) -> None:
+        self.managers = managers
+        self.max_model_len = max_model_len
+        groups = kv_cache_config.kv_cache_groups
+
+        resident_managers: list[HiSparseResidentManager] = []
+        hot_managers: list[HiSparseHotManager] = []
+        resident_group_ids: list[int] = []
+        for group_id, (group, manager) in enumerate(zip(groups, managers)):
+            if isinstance(group.kv_cache_spec, HiSparseResidentSpec):
+                if not isinstance(manager, HiSparseResidentManager):
+                    raise TypeError(
+                        "HiSparse resident specs require resident cache managers."
+                    )
+                resident_managers.append(manager)
+                resident_group_ids.append(group_id)
+                assert not group.host_resident
+            if isinstance(manager, HiSparseHotManager):
+                hot_managers.append(manager)
+                assert not group.host_resident
+            if isinstance(manager, (HiSparseHotManager, HiSparseResidentManager)):
+                manager.coordinator = self
+        self.resident_managers = tuple(resident_managers)
+        self.hot_managers = tuple(hot_managers)
+        self.host_manager: HiSparseSourceManager | None = None
+        self.host_group_id: int | None = None
+        self.max_spill_pages = 0
+        for group_id, manager in enumerate(managers):
+            if isinstance(manager, HiSparseSourceManager):
+                if self.host_manager is not None:
+                    raise ValueError("Only one HiSparse host group is supported.")
+                num_host_blocks = kv_cache_config.hisparse_host_num_blocks
+                if num_host_blocks is None:
+                    raise ValueError("HiSparse host group needs host capacity.")
+                self.host_group_id = group_id
+                self.host_manager = manager
+                manager.coordinator = self
+                manager.bind_host_pool(num_host_blocks)
+        self.has_host_cache = self.host_manager is not None
+        self.gpu_pool: BlockPool | None = None
+        self.transition_watermark = 0
+        if self.resident_managers:
+            assert self.host_manager is not None
+            assert self.host_group_id is not None
+            if self.host_group_id > min(resident_group_ids):
+                raise ValueError("HiSparse host group must precede resident groups.")
+            resident_block_sizes = {
+                manager.block_size for manager in self.resident_managers
+            }
+            if len(resident_block_sizes) != 1:
+                raise ValueError(
+                    "HiSparse resident cache groups must use one block size."
+                )
+            resident_block_size = resident_block_sizes.pop()
+            if self.host_manager.block_size != resident_block_size:
+                raise ValueError("HiSparse host and resident block sizes must match.")
+            self.max_spill_pages = max_model_len // resident_block_size
+            self.gpu_pool = self.resident_managers[0].block_pool
+            # Requests start reading from host once the shared pool runs this
+            # low, so admissions find evictable pages rather than pinned ones.
+            hot_cost = sum(manager.blocks_per_request for manager in hot_managers)
+            self.transition_watermark = max(hot_cost, kv_cache_config.num_blocks // 10)
+
+        # Host copies handed to the worker, by destination block id, until it
+        # reports having run them.
+        self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {}
+        # request -> prefix length an external load is still filling in.
+        self._pending_imports: dict[str, int] = {}
+        self.block_table_updates: set[str] = set()
+        self.spills_to_send: list[SparseKVPageTransfer] = []
+        self.pending_spills: dict[int, _PendingSpill] = {}
+        self.request_states: dict[str, _HiSparseRequestState] = {}
+        self.next_spill_id = 0
+        # Published host block hash -> GPU copies of that page, readable until
+        # the pool reuses them. Lets a host prefix hit come back GPU-resident.
+        self.copies: dict[BlockHashWithGroupId, tuple[KVCacheBlock, ...]] = {}
+        self._copy_by_block: dict[int, BlockHashWithGroupId] = {}
+        # Unpinned GPU block id -> (request, page) pairs still reading it.
+        self._owners: dict[int, set[tuple[str, int]]] = {}
+
+    def _get_request_state(self, request_id: str) -> _HiSparseRequestState:
+        state = self.request_states.get(request_id)
+        if state is None:
+            state = _HiSparseRequestState()
+            self.request_states[request_id] = state
+        return state
+
+    def commit_computed_blocks(self, request_id: str, num_host_pages: int) -> None:
+        """Account for a prefix hit on host pages the request now references."""
+        if not self.resident_managers or self.host_manager is None:
+            return
+        host_blocks = self.host_manager.req_to_blocks.get(request_id, ())
+        num_host_pages = min(num_host_pages, len(host_blocks))
+        if num_host_pages == 0:
+            return
+        state = self._get_request_state(request_id)
+        state.valid_pages.update(range(num_host_pages))
+        state.ready_prefix_pages = max(num_host_pages, state.ready_prefix_pages)
+        self._adopt_copies(request_id, state, host_blocks[:num_host_pages])
+        state.copies_recorded_blocks = max(state.copies_recorded_blocks, num_host_pages)
+
+    def take_host_block_copies(self) -> tuple[KVCacheBlockCopy, ...]:
+        """Drain host copy-on-write work, retaining both endpoints.
+
+        The worker runs the copies in the step that carries them, so the
+        endpoints are released when that step's output comes back.
+        """
+        if self.host_manager is None:
+            return ()
+        pairs = self.host_manager.take_host_cow_copies()
+        for source_block, cow_block in pairs:
+            self._retained_copies[cow_block.block_id] = (source_block, cow_block)
+        return tuple(
+            KVCacheBlockCopy(
+                src_block_id=source_block.block_id,
+                dst_block_id=cow_block.block_id,
+            )
+            for source_block, cow_block in pairs
+        )
+
+    def release_completed_host_copies(self, dst_block_ids: Iterable[int]) -> None:
+        """Release the endpoints of host copies the worker reports having run."""
+        blocks: list[KVCacheBlock] = []
+        for dst_block_id in dst_block_ids:
+            pair = self._retained_copies.pop(dst_block_id, None)
+            if pair is not None:
+                blocks.extend(pair)
+        if not blocks:
+            return
+        assert self.host_manager is not None
+        self.host_manager.block_pool.free_blocks(reversed(blocks))
+
+    def get_host_block_pool(self) -> BlockPool | None:
+        manager = self.host_manager
+        return manager.block_pool if manager is not None else None
+
+    # ------------------------------------------------------------------
+    # GPU copies of published host pages
+    # ------------------------------------------------------------------
+
+    def _adopt_copies(
+        self,
+        request_id: str,
+        state: _HiSparseRequestState,
+        host_blocks: Sequence[KVCacheBlock],
+    ) -> None:
+        """Point null prefix pages at readable GPU copies of their contents.
+
+        Runs after the hit length is fully reconciled, so copy hits and misses
+        can be scattered per page without affecting the prefix hit.
+        """
+        if not self.copies:
+            return
+        for host_idx, host_block in enumerate(host_blocks):
+            if host_block.is_null or host_block.block_hash is None:
+                continue
+            blocks = self.copies.get(host_block.block_hash)
+            if blocks is None:
+                continue
+            for manager, block in zip(self.resident_managers, blocks):
+                if manager.adopt_resident_page(request_id, host_idx, block):
+                    manager.block_pool.touch([block])
+                    state.pinned_clean.add(host_idx)
+
+    def _record_copies(self, request_id: str, num_computed_tokens: int) -> None:
+        """Index the GPU copies of just-published host blocks for later hits."""
+        if not self.resident_managers:
+            return
+        assert self.host_manager is not None
+        host_blocks = self.host_manager.req_to_blocks.get(request_id)
+        if not host_blocks:
+            return
+        state = self._get_request_state(request_id)
+        num_host_blocks = min(
+            num_computed_tokens // self.host_manager.block_size, len(host_blocks)
+        )
+        for host_idx in range(state.copies_recorded_blocks, num_host_blocks):
+            if host_idx in state.pending_pages:
+                continue
+            self._record_copy(request_id, host_idx, host_blocks[host_idx])
+        state.copies_recorded_blocks = max(
+            state.copies_recorded_blocks, num_host_blocks
+        )
+
+    def _record_copy(
+        self, request_id: str, page_idx: int, host_block: KVCacheBlock
+    ) -> None:
+        if host_block.is_null or host_block.block_hash is None:
+            return
+        blocks: list[KVCacheBlock] = []
+        for manager in self.resident_managers:
+            block = manager.get_resident_page(request_id, page_idx)
+            if block is None:
+                return
+            blocks.append(block)
+        self._drop_copy(host_block.block_hash)
+        self.copies[host_block.block_hash] = tuple(blocks)
+        for block in blocks:
+            self._copy_by_block[block.block_id] = host_block.block_hash
+
+    def _drop_copy(self, host_hash: BlockHashWithGroupId) -> None:
+        blocks = self.copies.pop(host_hash, None)
+        if blocks is None:
+            return
+        for block in blocks:
+            self._copy_by_block.pop(block.block_id, None)
+
+    # ------------------------------------------------------------------
+    # Residency: pinning, unpinning and losing pages
+    # ------------------------------------------------------------------
+
+    def _can_read_from_host(self, request_id: str) -> bool:
+        """Whether resident pages may be dropped under the request."""
+        return bool(self.hot_managers) and all(
+            manager.has_hot(request_id) for manager in self.hot_managers
+        )
+
+    def _resident_page_blocks(
+        self, request_id: str, page_idx: int
+    ) -> list[KVCacheBlock] | None:
+        blocks: list[KVCacheBlock] = []
+        for manager in self.resident_managers:
+            req_blocks = manager.req_to_blocks.get(request_id)
+            if (
+                req_blocks is None
+                or page_idx >= len(req_blocks) - _ACTIVE_TAIL_PAGES
+                or req_blocks[page_idx].is_null
+            ):
+                return None
+            blocks.append(req_blocks[page_idx])
+        return blocks
+
+    def _unpin_page(
+        self, request_id: str, state: _HiSparseRequestState, page_idx: int
+    ) -> bool:
+        if page_idx in state.pending_pages or page_idx not in state.valid_pages:
+            return False
+        blocks = self._resident_page_blocks(request_id, page_idx)
+        if blocks is None:
+            return False
+        for manager, block in zip(self.resident_managers, blocks):
+            self._owners.setdefault(block.block_id, set()).add((request_id, page_idx))
+            manager.block_pool.unpin_blocks([block], self._on_block_reused)
+        state.pinned_clean.discard(page_idx)
+        state.unpinned_pages.add(page_idx)
+        return True
+
+    def _unpin_clean_pages(self, request_id: str, state: _HiSparseRequestState) -> None:
+        for page_idx in sorted(state.pinned_clean):
+            self._unpin_page(request_id, state, page_idx)
+
+    def _page_became_clean(
+        self, request_id: str, state: _HiSparseRequestState, page_idx: int
+    ) -> None:
+        state.pinned_clean.add(page_idx)
+        if self._can_read_from_host(request_id):
+            self._unpin_page(request_id, state, page_idx)
+
+    def _on_block_reused(self, block: KVCacheBlock) -> None:
+        host_hash = self._copy_by_block.get(block.block_id)
+        if host_hash is not None:
+            self._drop_copy(host_hash)
+        for request_id, page_idx in self._owners.pop(block.block_id, ()):
+            self._lose_page(request_id, page_idx)
+
+    def _lose_page(self, request_id: str, page_idx: int) -> None:
+        state = self.request_states.get(request_id)
+        if state is None:
+            return
+        for manager in self.resident_managers:
+            blocks = manager.req_to_blocks.get(request_id)
+            if blocks is None or page_idx >= len(blocks) or blocks[page_idx].is_null:
+                continue
+            block = blocks[page_idx]
+            blocks[page_idx] = manager._null_block
+            owners = self._owners.get(block.block_id)
+            if owners is not None:
+                owners.discard((request_id, page_idx))
+        state.unpinned_pages.discard(page_idx)
+        self.block_table_updates.add(request_id)
+        for hot_manager in self.hot_managers:
+            hot_manager.require_hot(request_id)
+
+    def update_residency(self, request_id: str) -> None:
+        """Per-step residency policy for a scheduled request.
+
+        A request that can read from host releases every clean sealed page to
+        the pool. One that cannot keeps its pages pinned until the shared pool
+        runs low, then asks for a hot region. Pages remain pinned until that region
+        is allocated on a subsequent scheduling pass.
+        """
+        if not self.resident_managers:
+            return
+        state = self._get_request_state(request_id)
+        if not self._can_read_from_host(request_id):
+            assert self.gpu_pool is not None
+            if self.gpu_pool.get_num_free_blocks() >= self.transition_watermark:
+                return
+            for manager in self.hot_managers:
+                manager.require_hot(request_id)
+            return
+        self._unpin_clean_pages(request_id, state)
+
+    # ------------------------------------------------------------------
+    # Host publication and spills
+    # ------------------------------------------------------------------
+
+    def plan_prefix_materialization(
+        self, request_id: str, num_computed_tokens: int
+    ) -> None:
+        """Eagerly write back every sealed page that has no host copy yet."""
+        if not self.resident_managers:
+            return
+        assert self.host_manager is not None
+        host_block_size = self.host_manager.block_size
+        num_pages = num_computed_tokens // host_block_size
+        importing_pages = cdiv(
+            self._pending_imports.get(request_id, 0), host_block_size
+        )
+        state = self._get_request_state(request_id)
+        budget = max(self.max_spill_pages - len(self.spills_to_send), 0)
+        for page_idx in range(num_pages):
+            if page_idx < importing_pages:
+                continue
+            if budget == 0:
+                break
+            if page_idx in state.pending_pages:
+                continue
+            if page_idx in state.valid_pages:
+                continue
+            if not all(
+                manager.get_resident_page(request_id, page_idx) is not None
+                for manager in self.resident_managers
+            ):
+                continue
+            if self._plan_page_transfer(request_id, page_idx, after_forward=True):
+                budget -= 1
+
+    def publish_when_ready(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+        retention_interval: int | None,
+        *,
+        replay_boundaries: Sequence[int],
+    ) -> None:
+        """Publish host-source hashes only after their pages are durable."""
+        manager = self.host_manager
+        if manager is None:
+            return
+        num_pages = num_computed_tokens // manager.block_size
+        request_id = request.request_id
+        state = self._get_request_state(request_id)
+        if state.ready_prefix_pages >= num_pages:
+            manager.publish_blocks(
+                request,
+                num_computed_tokens,
+                retention_interval=retention_interval,
+                replay_boundaries=replay_boundaries,
+            )
+            self._record_copies(request_id, num_computed_tokens)
+            state.publication = None
+            return
+        state.publication = _PendingPublication(
+            request=request,
+            num_computed_tokens=num_computed_tokens,
+            num_pages=num_pages,
+            retention_interval=retention_interval,
+            replay_boundaries=replay_boundaries,
+        )
+
+    def _publish_host_blocks_if_ready(self, request_id: str) -> None:
+        state = self.request_states.get(request_id)
+        if state is None:
+            return
+        publication = state.publication
+        if publication is None or state.ready_prefix_pages < publication.num_pages:
+            return
+        assert self.host_manager is not None
+        self.host_manager.publish_blocks(
+            publication.request,
+            publication.num_computed_tokens,
+            retention_interval=publication.retention_interval,
+            replay_boundaries=publication.replay_boundaries,
+        )
+        self._record_copies(request_id, publication.num_computed_tokens)
+        state.publication = None
+
+    def record_pending_host_import(self, request_id: str, num_tokens: int) -> None:
+        """Note a prefix an external load is populating in host pages."""
+        self._pending_imports[request_id] = num_tokens
+
+    def finish_host_import(self, request_id: str, *, failed: bool) -> None:
+        num_tokens = self._pending_imports.pop(request_id, None)
+        if num_tokens is not None and not failed:
+            self.complete_host_import(request_id, num_tokens)
+
+    def complete_host_import(self, request_id: str, num_computed_tokens: int) -> None:
+        """Publish externally populated host pages after connector completion."""
+        if not self.resident_managers:
+            return
+        block_size = self.resident_managers[0].block_size
+        num_pages = num_computed_tokens // block_size
+        state = self._get_request_state(request_id)
+        state.valid_pages = set(range(num_pages))
+        state.ready_prefix_pages = num_pages
+        if num_computed_tokens:
+            self._plan_page_transfer(
+                request_id,
+                (num_computed_tokens - 1) // block_size,
+                after_forward=False,
+                restore=True,
+            )
+        self._publish_host_blocks_if_ready(request_id)
+
+    def _plan_page_transfer(
+        self,
+        request_id: str,
+        page_idx: int,
+        *,
+        after_forward: bool,
+        restore: bool = False,
+    ) -> bool:
+        assert self.host_manager is not None
+        state = self._get_request_state(request_id)
+        if page_idx in state.pending_pages:
+            return False
+        host_blocks = self.host_manager.req_to_blocks.get(request_id)
+        if host_blocks is None or page_idx >= len(host_blocks):
+            return False
+        host_block = host_blocks[page_idx]
+        if host_block.is_null:
+            return False
+
+        blocks: list[KVCacheBlock] = []
+        for manager in self.resident_managers:
+            block = manager.get_resident_page(request_id, page_idx)
+            if block is None:
+                return False
+            blocks.append(block)
+
+        self.host_manager.block_pool.touch([host_block])
+        for manager, block in zip(self.resident_managers, blocks):
+            manager.block_pool.touch([block])
+        spill_id = self.next_spill_id
+        self.next_spill_id += 1
+        plan = SparseKVPageTransfer(
+            transfer_id=spill_id,
+            host_block_id=host_block.block_id,
+            resident_block_ids=tuple(block.block_id for block in blocks),
+            after_forward=after_forward,
+            restore=restore,
+        )
+        self.pending_spills[spill_id] = _PendingSpill(
+            transfer_id=spill_id,
+            page=(request_id, page_idx),
+            request_state=state,
+            host_block=host_block,
+            resident_blocks=tuple(blocks),
+            restore=restore,
+        )
+        state.pending_pages[page_idx] = spill_id
+        self.spills_to_send.append(plan)
+        return True
+
+    # ------------------------------------------------------------------
+    # Worker-facing views
+    # ------------------------------------------------------------------
+
+    def build_row_mirrors(
+        self,
+        requests: Iterable[tuple[str, int, int]],
+    ) -> tuple[SparseKVRowMirror, ...]:
+        """Map scheduled rows to scheduler-owned resident and host blocks."""
+        if not self.resident_managers or self.host_manager is None:
+            return ()
+        block_size = self.resident_managers[0].block_size
+        mirrors: list[SparseKVRowMirror] = []
+        for request_id, num_computed_tokens, num_scheduled_tokens in requests:
+            host_blocks = self.host_manager.req_to_blocks.get(request_id)
+            resident_blocks = [
+                manager.req_to_blocks.get(request_id)
+                for manager in self.resident_managers
+            ]
+            if host_blocks is None or any(blocks is None for blocks in resident_blocks):
+                continue
+            token_position = num_computed_tokens
+            end_position = token_position + num_scheduled_tokens
+            while token_position < end_position:
+                page_idx, row_offset = divmod(token_position, block_size)
+                num_rows = min(block_size - row_offset, end_position - token_position)
+                # A page without a GPU copy has no mirror source, and a page
+                # without a host block has no destination. Skip just that page:
+                # its neighbours in the window are still mirrorable, and a
+                # window only ever moves forward, so dropping them here would
+                # leave their host rows permanently stale.
+                source_starts = []
+                for blocks in resident_blocks:
+                    assert blocks is not None
+                    if page_idx >= len(blocks) or blocks[page_idx].is_null:
+                        break
+                    source_starts.append(
+                        blocks[page_idx].block_id * block_size + row_offset
+                    )
+                if len(source_starts) != len(resident_blocks):
+                    token_position += num_rows
+                    continue
+                host_block_idx = page_idx
+                if host_block_idx >= len(host_blocks):
+                    token_position += num_rows
+                    continue
+                host_block = host_blocks[host_block_idx]
+                if host_block.is_null:
+                    token_position += num_rows
+                    continue
+                mirrors.append(
+                    SparseKVRowMirror(
+                        source_starts=tuple(source_starts),
+                        destination_start=(
+                            host_block.block_id * block_size + row_offset
+                        ),
+                        num_rows=num_rows,
+                    )
+                )
+                token_position += num_rows
+        return tuple(mirrors)
+
+    def all_context_pages_resident(
+        self,
+        requests: Iterable[tuple[str, int, int]],
+    ) -> bool:
+        """Return whether every scheduled request can read only resident KV."""
+        if not self.resident_managers:
+            return False
+        block_size = self.resident_managers[0].block_size
+        for request_id, num_computed_tokens, num_scheduled_tokens in requests:
+            num_pages = cdiv(num_computed_tokens + num_scheduled_tokens, block_size)
+            for page_idx in range(num_pages):
+                if any(
+                    manager.get_resident_page(request_id, page_idx) is None
+                    for manager in self.resident_managers
+                ):
+                    return False
+        return True
+
+    def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]:
+        updates = {
+            request_id: tuple(
+                [block.block_id for block in manager.req_to_blocks.get(request_id, [])]
+                for manager in self.managers
+            )
+            for request_id in self.block_table_updates
+        }
+        self.block_table_updates.clear()
+        return updates
+
+    def build_offload_command(self) -> SparseKVOffloadCommand | None:
+        """Package residency decisions for the worker connector."""
+        if not self.resident_managers:
+            return None
+        command = SparseKVOffloadCommand(page_transfers=self.spills_to_send)
+        self.spills_to_send = []
+        return command
+
+    def has_pending_reclamation(self) -> bool:
+        """Whether an in-flight spill will free GPU blocks on completion."""
+        return any(
+            pending.page[0] in self.request_states
+            and self._can_read_from_host(pending.page[0])
+            for pending in self.pending_spills.values()
+        )
+
+    def has_pending_work(self) -> bool:
+        return bool(self.spills_to_send or self.pending_spills or self._retained_copies)
+
+    def update_spills(
+        self,
+        enqueued_counts: Mapping[int, int],
+        completed_counts: Mapping[int, int],
+    ) -> None:
+        for spill_id, count in enqueued_counts.items():
+            pending = self.pending_spills.get(spill_id)
+            if pending is not None and pending.expected_worker_completions == 0:
+                pending.expected_worker_completions = count
+        for spill_id, count in completed_counts.items():
+            pending = self.pending_spills.get(spill_id)
+            if pending is not None:
+                pending.worker_completions += count
+
+        self._apply_enqueued_spills()
+        self._complete_host_writes()
+
+    def _apply_enqueued_spills(self) -> None:
+        """Drop the GPU pins once every worker has completed the transfer."""
+        for pending in self.pending_spills.values():
+            if (
+                pending.enqueue_applied
+                or not pending.expected_worker_completions
+                or pending.worker_completions < pending.expected_worker_completions
+            ):
+                continue
+            for manager, block in zip(self.resident_managers, pending.resident_blocks):
+                manager.block_pool.free_blocks([block])
+            pending.enqueue_applied = True
+
+    def _complete_host_writes(self) -> None:
+        completed = [
+            pending
+            for pending in self.pending_spills.values()
+            if pending.enqueue_applied
+            and pending.expected_worker_completions
+            and pending.worker_completions >= pending.expected_worker_completions
+        ]
+        completed_request_ids: set[str] = set()
+        for pending in completed:
+            self.pending_spills.pop(pending.transfer_id, None)
+            request_id, page_idx = pending.page
+            pending.request_state.pending_pages.pop(page_idx, None)
+            if self.request_states.get(request_id) is pending.request_state:
+                if not pending.restore:
+                    pending.request_state.valid_pages.add(page_idx)
+                if page_idx in pending.request_state.valid_pages:
+                    self._page_became_clean(request_id, pending.request_state, page_idx)
+                    if pending.restore:
+                        self._record_copy(request_id, page_idx, pending.host_block)
+                completed_request_ids.add(request_id)
+            assert self.host_manager is not None
+            self.host_manager.block_pool.free_blocks([pending.host_block])
+        for request_id in completed_request_ids:
+            state = self.request_states[request_id]
+            while state.ready_prefix_pages in state.valid_pages:
+                state.ready_prefix_pages += 1
+            self._publish_host_blocks_if_ready(request_id)
+
+    def free(self, request_id: str) -> None:
+        """Detach the request; its clean pages stay readable copies in the pool."""
+        self._pending_imports.pop(request_id, None)
+        state = self.request_states.pop(request_id, None)
+        if state is None:
+            return
+        for manager in self.resident_managers:
+            blocks = manager.req_to_blocks.get(request_id)
+            if not blocks:
+                continue
+            for page_idx, block in enumerate(blocks):
+                if block.is_null:
+                    continue
+                if page_idx in state.unpinned_pages:
+                    owners = self._owners.get(block.block_id)
+                    if owners is not None:
+                        owners.discard((request_id, page_idx))
+                elif (
+                    page_idx in state.valid_pages
+                    and page_idx not in state.pending_pages
+                ):
+                    manager.block_pool.unpin_blocks([block], self._on_block_reused)
+                else:
+                    continue
+                blocks[page_idx] = manager._null_block
+
+
+def get_hisparse_coordinator(
+    kv_cache_manager: "KVCacheManager",
+) -> HiSparseCoordinator:
+    """Return the coordinator shared by a KV cache manager's HiSparse groups.
+
+    Built on first call and memoised on the HiSparse managers it binds itself
+    to. Must run before the first request is admitted.
+    """
+    managers = tuple(kv_cache_manager.coordinator.single_type_managers)
+    for manager in managers:
+        coordinator = getattr(manager, "coordinator", None)
+        if coordinator is not None:
+            assert isinstance(coordinator, HiSparseCoordinator)
+            return coordinator
+    coordinator = HiSparseCoordinator(
+        kv_cache_manager.kv_cache_config, managers, kv_cache_manager.max_model_len
+    )
+    if coordinator.host_manager is None:
+        raise ValueError("No HiSparse cache group is configured.")
+    return coordinator
diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py
new file mode 100644
index 000000000000..b3b45d3334a5
--- /dev/null
+++ b/vllm/v1/hisparse/layout.py
@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import cast
+
+from vllm.config import VllmConfig
+from vllm.config.kv_transfer import hisparse_host_pool_gib
+from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
+from vllm.v1.hisparse.runtime import (
+    ResolvedHiSparseConfig,
+    get_hisparse_host_block_stride,
+    use_shared_hisparse_host_pool,
+)
+from vllm.v1.kv_cache_interface import (
+    HiSparseHotSpec,
+    HiSparseResidentSpec,
+    KVCacheConfig,
+    KVCacheGroupRole,
+    KVCacheGroupSpec,
+    KVCacheSpec,
+    KVCacheTensor,
+    MLAAttentionSpec,
+    SparseCacheRole,
+    UniformTypeKVCacheSpecs,
+    compute_layout_strides,
+)
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+logger = init_logger(__name__)
+
+HISPARSE_HOT_SUFFIX = ".hisparse_hot"
+HISPARSE_RESIDENT_SUFFIX = ".hisparse_resident"
+
+
+@dataclass(frozen=True)
+class HiSparseLayout:
+    source_group: KVCacheGroupSpec
+    device_groups: list[KVCacheGroupSpec]
+    host_num_blocks: int
+    host_block_stride: int
+    shared_host_pool: bool
+
+
+def get_hisparse_kv_cache_groups(
+    vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec]
+) -> list[KVCacheGroupSpec] | None:
+    attention_config = getattr(vllm_config, "attention_config", None)
+    if attention_config is None or attention_config.hisparse_config is None:
+        return None
+
+    mla_specs: dict[str, KVCacheSpec] = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if isinstance(spec, MLAAttentionSpec)
+    }
+    other_specs = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if not isinstance(spec, MLAAttentionSpec)
+    }
+    if not mla_specs or not other_specs:
+        return None
+
+    from vllm.v1.core.kv_cache_utils import get_kv_cache_groups
+
+    mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs)
+    assert mla_group_spec is not None
+    mla_group = KVCacheGroupSpec(list(mla_specs), mla_group_spec)
+    return [mla_group, *get_kv_cache_groups(vllm_config, other_specs)]
+
+
+def get_hisparse_host_pool_bytes(vllm_config: VllmConfig) -> int:
+    host_pool_gib = hisparse_host_pool_gib(vllm_config.kv_transfer_config)
+    if host_pool_gib is None:
+        raise ValueError("HiSparse requires HiSparseConnector with host_pool_gib")
+    return int(host_pool_gib * 2**30)
+
+
+def _partition_hisparse_specs(
+    groups: list[KVCacheGroupSpec],
+) -> tuple[dict[str, MLAAttentionSpec], dict[str, MLAAttentionSpec]]:
+    group_spec = groups[0].kv_cache_spec
+    if not isinstance(group_spec, UniformTypeKVCacheSpecs):
+        raise ValueError("HiSparse requires uniform sparse-MLA cache specs.")
+
+    specs = group_spec.kv_cache_specs
+    if any(
+        isinstance(spec, MLAAttentionSpec) and spec.model_version == "deepseek_v4"
+        for spec in specs.values()
+    ):
+        raise ValueError("HiSparse does not support DeepSeek V4.")
+    if not all(isinstance(spec, MLAAttentionSpec) for spec in specs.values()):
+        raise ValueError("HiSparse requires its first cache group to contain MLA only.")
+
+    source_specs = {
+        name: spec
+        for name, spec in specs.items()
+        if isinstance(spec, MLAAttentionSpec)
+        and spec.cache_role is SparseCacheRole.SPARSE
+    }
+    indexer_specs = {
+        name: spec
+        for name, spec in specs.items()
+        if isinstance(spec, MLAAttentionSpec)
+        and spec.cache_role is SparseCacheRole.INDEXER
+    }
+    if not source_specs or not indexer_specs:
+        raise ValueError("HiSparse requires sparse-MLA and indexer cache specs.")
+    return source_specs, indexer_specs
+
+
+def get_hisparse_gpu_memory_usage(
+    vllm_config: VllmConfig,
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> int:
+    _, indexer_specs = _partition_hisparse_specs(kv_cache_groups)
+    return sum(
+        spec.max_memory_usage_bytes(vllm_config) for spec in indexer_specs.values()
+    ) + sum(
+        group.kv_cache_spec.max_memory_usage_bytes(vllm_config)
+        for group in kv_cache_groups[1:]
+    )
+
+
+def create_hisparse_layout(
+    vllm_config: VllmConfig,
+    groups: list[KVCacheGroupSpec],
+    host_budget: int,
+) -> HiSparseLayout:
+    source_specs, indexer_specs = _partition_hisparse_specs(groups)
+    block_sizes = {
+        spec.block_size for spec in (*source_specs.values(), *indexer_specs.values())
+    }
+    if len(block_sizes) != 1:
+        raise ValueError("HiSparse requires one resolved GPU block size.")
+    gpu_block_size = block_sizes.pop()
+
+    config = ResolvedHiSparseConfig.from_vllm_config(
+        vllm_config,
+        vllm_config.model_config.hf_config.index_topk,
+        gpu_block_size,
+    )
+    assert config is not None
+    indexer_group_spec = UniformTypeKVCacheSpecs.from_specs(
+        cast(dict[str, KVCacheSpec], indexer_specs)
+    )
+    assert indexer_group_spec is not None
+    indexer_group = KVCacheGroupSpec(
+        list(indexer_specs),
+        indexer_group_spec,
+        enable_kv_transfer=True,
+        role=KVCacheGroupRole.HISPARSE_INDEXER,
+    )
+
+    indexer_page = sum(spec.page_size_bytes for spec in indexer_specs.values())
+    hot_blocks_per_request = cdiv(config.device_buffer_size, gpu_block_size)
+    hot_units: list[list[tuple[str, MLAAttentionSpec]]] = []
+    for layer_name, layer_spec in source_specs.items():
+        if layer_spec.is_index_group_leader or not hot_units:
+            hot_units.append([])
+        hot_units[-1].append((f"{layer_name}{HISPARSE_HOT_SUFFIX}", layer_spec))
+
+    resident_groups: list[KVCacheGroupSpec] = []
+    hot_groups: list[KVCacheGroupSpec] = []
+
+    def append_hot_group(layers: list[tuple[str, MLAAttentionSpec]]) -> None:
+        page_sizes = {spec.page_size_bytes for _, spec in layers}
+        if len(page_sizes) != 1:
+            raise ValueError(
+                "HiSparse hot-cache groups require one page size, got "
+                f"{sorted(page_sizes)}."
+            )
+        page_size = page_sizes.pop()
+        names = [name for name, _ in layers]
+        resident_groups.append(
+            KVCacheGroupSpec(
+                [
+                    name[: -len(HISPARSE_HOT_SUFFIX)] + HISPARSE_RESIDENT_SUFFIX
+                    for name in names
+                ],
+                HiSparseResidentSpec(
+                    block_size=gpu_block_size,
+                    page_size=page_size,
+                ),
+                enable_kv_transfer=False,
+            )
+        )
+        hot_groups.append(
+            KVCacheGroupSpec(
+                names,
+                HiSparseHotSpec(
+                    block_size=gpu_block_size,
+                    page_size=page_size,
+                    blocks_per_request=hot_blocks_per_request,
+                ),
+                enable_kv_transfer=False,
+            )
+        )
+
+    current: list[tuple[str, MLAAttentionSpec]] = []
+    current_page = 0
+    for unit in hot_units:
+        unit_page = sum(spec.page_size_bytes for _, spec in unit)
+        if current and current_page + unit_page > indexer_page:
+            append_hot_group(current)
+            current = []
+            current_page = 0
+        current.extend(unit)
+        current_page += unit_page
+    if current:
+        append_hot_group(current)
+
+    source_group_spec = UniformTypeKVCacheSpecs.from_specs(
+        cast(dict[str, KVCacheSpec], source_specs)
+    )
+    assert source_group_spec is not None
+    source_group = KVCacheGroupSpec(
+        list(source_specs),
+        source_group_spec,
+        role=KVCacheGroupRole.HISPARSE_SOURCE,
+        host_resident=True,
+        enable_kv_transfer=True,
+    )
+    regular_groups = groups[1:]
+    gpu_groups = [indexer_group, *resident_groups, *hot_groups, *regular_groups]
+
+    shared_host_pool = use_shared_hisparse_host_pool(vllm_config)
+    host_block_stride = get_hisparse_host_block_stride(
+        sum(spec.page_size_bytes for spec in source_specs.values()),
+        use_shared_host_pool=shared_host_pool,
+    )
+    host_num_blocks = host_budget // host_block_stride
+    if host_num_blocks <= 0:
+        raise ValueError("HiSparse has no allocatable host blocks.")
+
+    return HiSparseLayout(
+        source_group=source_group,
+        device_groups=gpu_groups,
+        host_num_blocks=host_num_blocks,
+        host_block_stride=host_block_stride,
+        shared_host_pool=shared_host_pool,
+    )
+
+
+def _build_hisparse_kv_cache_tensors(
+    kv_cache_groups: list[KVCacheGroupSpec],
+    num_blocks: int,
+    size: int,
+    layout: KVCacheLayout,
+    bytes_per_block: int,
+    *,
+    host_resident: bool = False,
+) -> list[KVCacheTensor]:
+    interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
+    tensors: list[KVCacheTensor] = []
+    for group in kv_cache_groups:
+        group_spec = group.kv_cache_spec
+        layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list)
+        if isinstance(group_spec, UniformTypeKVCacheSpecs):
+            for layer_name, spec in group_spec.kv_cache_specs.items():
+                layers_by_spec[spec].append(layer_name)
+        elif group.layer_names:
+            layers_by_spec[group_spec].extend(group.layer_names)
+
+        byte_offset = 0
+        for spec, layer_names in layers_by_spec.items():
+            if isinstance(spec, (HiSparseHotSpec, HiSparseResidentSpec)):
+                if not layout.is_block_outermost:
+                    raise ValueError("HiSparse requires a block-outermost KV layout.")
+                layer_stride = spec.page_size_bytes
+                block_stride = bytes_per_block
+            else:
+                layer_stride, block_stride, _, _, _ = compute_layout_strides(
+                    spec,
+                    num_blocks,
+                    len(layer_names),
+                    layout,
+                    fixed_strides=(None, interleaved_block_stride, None, None, None),
+                )
+            offset = (
+                byte_offset
+                * max(layer_stride, spec.page_size_bytes)
+                // spec.page_size_bytes
+            )
+            tensors.append(
+                KVCacheTensor(
+                    size=size,
+                    layers=layer_names,
+                    layer_stride=layer_stride,
+                    block_stride=block_stride,
+                    offset=offset,
+                    host_resident=host_resident,
+                )
+            )
+            byte_offset += len(layer_names) * spec.page_size_bytes
+    return tensors
+
+
+def get_hisparse_kv_cache_config(
+    vllm_config: VllmConfig,
+    kv_cache_groups: list[KVCacheGroupSpec],
+    available_memory: int,
+    host_budget: int,
+) -> KVCacheConfig:
+    from vllm.v1.core.kv_cache_utils import (
+        _get_kv_cache_bytes_per_block,
+        may_override_num_blocks,
+        validate_kv_cache_layout,
+    )
+
+    hisparse_layout = create_hisparse_layout(vllm_config, kv_cache_groups, host_budget)
+    device_groups = hisparse_layout.device_groups
+    layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
+    validate_kv_cache_layout(layout, device_groups)
+    bytes_per_block = _get_kv_cache_bytes_per_block(device_groups)
+    num_blocks = may_override_num_blocks(
+        vllm_config, available_memory // bytes_per_block
+    )
+    size = bytes_per_block * num_blocks
+    kv_cache_tensors = _build_hisparse_kv_cache_tensors(
+        device_groups, num_blocks, size, layout, bytes_per_block
+    )
+
+    host_groups = [hisparse_layout.source_group]
+    host_layout = KVCacheLayout.LBNHC
+    validate_kv_cache_layout(host_layout, host_groups)
+    host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups)
+    host_size = host_bytes_per_block * hisparse_layout.host_num_blocks
+    kv_cache_tensors[:0] = _build_hisparse_kv_cache_tensors(
+        host_groups,
+        hisparse_layout.host_num_blocks,
+        host_size,
+        host_layout,
+        host_bytes_per_block,
+        host_resident=True,
+    )
+    logger.info_once(
+        "HiSparse HMA: %.1f GiB host source (%d blocks), %.1f GiB shared "
+        "GPU indexer/resident/hot pool (%d blocks).",
+        host_size / 2**30,
+        hisparse_layout.host_num_blocks,
+        size / 2**30,
+        num_blocks,
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=kv_cache_tensors,
+        kv_cache_groups=[*host_groups, *device_groups],
+        hisparse_host_num_blocks=hisparse_layout.host_num_blocks,
+        hisparse_host_block_stride=hisparse_layout.host_block_stride,
+        hisparse_shared_host_pool=hisparse_layout.shared_host_pool,
+        prefix_cache_retention_interval=(
+            vllm_config.cache_config.prefix_cache_retention_interval
+        ),
+    )
diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py
new file mode 100644
index 000000000000..5100379083c3
--- /dev/null
+++ b/vllm/v1/hisparse/runtime.py
@@ -0,0 +1,1291 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVIDIA HiSparse resident, host, and hot-cache data plane."""
+
+from __future__ import annotations
+
+import math
+import mmap
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, TypeAlias
+
+import psutil
+import torch
+
+from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.distributed import get_tp_group
+from vllm.forward_context import get_forward_context, in_piecewise_cudagraph
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.math_utils import round_up
+from vllm.utils.torch_utils import current_stream
+from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
+from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
+
+logger = init_logger(__name__)
+
+HOST_REGISTER_CHUNK_BYTES = 256 * 2**30
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+
+# fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE.
+FP8_DS_MLA_ROW_BYTES = 656
+HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+
+
+def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool:
+    return (
+        attn_metadata is not None
+        and attn_metadata.num_decode_tokens == attn_metadata.num_actual_tokens
+        and attn_metadata.max_query_len == 1
+        and attn_metadata.num_reqs == attn_metadata.num_actual_tokens
+    )
+
+
+def _get_max_decode_query_len(vllm_config: VllmConfig) -> int:
+    speculative_config = getattr(vllm_config, "speculative_config", None)
+    if (
+        speculative_config is not None
+        and speculative_config.num_speculative_tokens is not None
+    ):
+        return 1 + speculative_config.num_speculative_tokens * (
+            2 if speculative_config.parallel_drafting else 1
+        )
+    return 1
+
+
+def _get_max_swap_rows(vllm_config: VllmConfig) -> int:
+    max_query_len = _get_max_decode_query_len(vllm_config)
+    scheduler_config = vllm_config.scheduler_config
+    return min(
+        scheduler_config.max_num_batched_tokens,
+        scheduler_config.max_num_seqs * max_query_len,
+    )
+
+
+@dataclass(frozen=True)
+class PagedCacheView:
+    cache: torch.Tensor
+    attention_cache: torch.Tensor
+    block_size: int
+    attention_block_stride: int
+
+    @classmethod
+    def bind(
+        cls,
+        raw_tensor: torch.Tensor,
+        *,
+        dtype: torch.dtype,
+        row_width: int,
+        byte_offset: int,
+        block_stride: int,
+        num_blocks: int,
+        block_size: int,
+    ) -> PagedCacheView:
+        itemsize = dtype.itemsize
+        assert byte_offset % itemsize == 0 and block_stride % itemsize == 0
+        cache = torch.as_strided(
+            raw_tensor.view(dtype),
+            size=(num_blocks, block_size, row_width),
+            stride=(block_stride // itemsize, row_width, 1),
+            storage_offset=byte_offset // itemsize,
+        )
+        row_bytes = row_width * itemsize
+        if block_stride % (block_size * row_bytes):
+            return cls(cache, cache, block_size, block_size)
+        attention_cache = (
+            raw_tensor[byte_offset:].view(dtype).view(-1, block_size, row_width)
+        )
+        return cls(cache, attention_cache, block_size, block_stride // row_bytes)
+
+
+@dataclass(frozen=True)
+class ResolvedHiSparseConfig:
+    top_k: int
+    device_buffer_size: int
+    eager_host_mirror: bool = True
+
+    @classmethod
+    def from_vllm_config(
+        cls,
+        vllm_config: VllmConfig,
+        model_top_k: int,
+        block_size: int | None = None,
+    ) -> ResolvedHiSparseConfig | None:
+        config = vllm_config.attention_config.hisparse_config
+        if config is None:
+            return None
+
+        max_query_len = _get_max_decode_query_len(vllm_config)
+        min_device_buffer_size = max_query_len * model_top_k
+        # Retain the worst-case union of every speculative step's top-k and
+        # leave one additional top-k of LRU slack.
+        configured_size = config.device_buffer_size
+        if configured_size is None:
+            device_buffer_size = (max_query_len + 1) * model_top_k
+        else:
+            device_buffer_size = configured_size
+
+        if device_buffer_size < min_device_buffer_size:
+            raise ValueError(
+                "HiSparse device_buffer_size must cover every decode query's "
+                "index_topk rows. Got "
+                f"device_buffer_size={device_buffer_size}, "
+                f"max_decode_query_len={max_query_len}, "
+                f"index_topk={model_top_k}; expected at least "
+                f"{min_device_buffer_size}."
+            )
+        max_device_buffer_size = torch.iinfo(torch.int16).max + 1
+        if device_buffer_size > max_device_buffer_size:
+            raise ValueError(
+                "HiSparse device_buffer_size exceeds the int16 slot-index "
+                f"limit: got {device_buffer_size}, maximum is "
+                f"{max_device_buffer_size}."
+            )
+        if configured_size is not None and block_size is not None:
+            padding = -device_buffer_size % block_size
+            if padding:
+                logger.warning(
+                    "HiSparse device_buffer_size=%d is not aligned to the "
+                    "%d-token kernel block size and allocates %d unused "
+                    "padding rows per hot group. Use %d to use all allocated "
+                    "rows.",
+                    device_buffer_size,
+                    block_size,
+                    padding,
+                    device_buffer_size + padding,
+                )
+        return cls(
+            top_k=model_top_k,
+            device_buffer_size=device_buffer_size,
+            eager_host_mirror=config.eager_host_mirror,
+        )
+
+
+def check_hisparse_host_memory(pool_bytes: int) -> None:
+    """Reject a physical host-pool allocation that cannot fit in RAM."""
+    mem = psutil.virtual_memory()
+    if pool_bytes > mem.available * 0.95:
+        raise ValueError(
+            f"HiSparse pinned host pool needs ~{pool_bytes / 2**30:.0f} GiB "
+            f"but only {mem.available / 2**30:.0f} GiB of RAM is available. "
+            "Lower the HiSparseConnector host_pool_gib or leave headroom "
+            "for co-tenants."
+        )
+
+
+def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Allocate and deterministically register an exact-size host KV region."""
+    page = 4096
+    padded_size = round_up(size, page)
+    backing = torch.empty(padded_size + page, dtype=torch.int8, device="cpu")
+    aligned_offset = (-backing.data_ptr()) % page
+    registered = backing[aligned_offset : aligned_offset + padded_size]
+    pin_tensor(registered)
+    return registered[:size], registered
+
+
+class HiSparseHostPool:
+    """Keep the single host backing and its exact CUDA registration range."""
+
+    def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:
+        self.vllm_config = vllm_config
+        self.kv_cache_config = kv_cache_config
+        self.backing: torch.Tensor | None = None
+        self.registered: torch.Tensor | None = None
+        self.shared_region: SharedOffloadRegion | None = None
+
+    def allocate(self, size: int) -> torch.Tensor:
+        assert self.backing is None, "HiSparse host pool is already allocated."
+        config = self.kv_cache_config
+        if config.hisparse_shared_host_pool:
+            tensor_sizes = [
+                size
+                for _, size in sorted(
+                    (tensor.offset + i * tensor.layer_stride, tensor.layer_stride)
+                    for tensor in config.kv_cache_tensors
+                    if tensor.host_resident
+                    for i in range(len(tensor.layers))
+                )
+            ]
+            assert config.hisparse_host_num_blocks is not None
+            assert config.hisparse_host_block_stride is not None
+            _, _, self.shared_region = allocate_hisparse_host_pools(
+                self.vllm_config,
+                tensor_sizes,
+                config.hisparse_host_num_blocks,
+                config.hisparse_host_block_stride,
+                use_shared_host_pool=True,
+            )
+            assert self.shared_region is not None
+            self.registered = self.shared_region.base_tensor
+            self.backing = self.registered[:size]
+        else:
+            check_hisparse_host_memory(size)
+            self.backing, self.registered = allocate_pinned_host_pool(size)
+        return self.backing
+
+
+def use_shared_hisparse_host_pool(vllm_config: VllmConfig) -> bool:
+    """Whether replicated MLA host KV can share one local mmap."""
+    parallel = vllm_config.parallel_config
+    return (
+        current_platform.is_cuda_alike()
+        and parallel.tensor_parallel_size > 1
+        and parallel.pipeline_parallel_size == 1
+        and parallel.prefill_context_parallel_size == 1
+        and parallel.decode_context_parallel_size == 1
+        and parallel.world_size == parallel.tensor_parallel_size
+        and parallel.distributed_executor_backend == "mp"
+        and parallel.nnodes_within_dp == 1
+    )
+
+
+def get_hisparse_host_block_stride(
+    page_size_bytes: int, *, use_shared_host_pool: bool
+) -> int:
+    if use_shared_host_pool:
+        return round_up(page_size_bytes, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT)
+    return page_size_bytes
+
+
+def _hisparse_registration_ranges(
+    tensor_sizes: list[int],
+    num_blocks: int,
+    host_block_stride: int,
+    max_chunk_bytes: int = HOST_REGISTER_CHUNK_BYTES,
+) -> tuple[tuple[int, int], ...]:
+    """Split a shared pool without bisecting any tensor's host block.
+
+    CUDA batch copies reject descriptors spanning adjacent registrations.
+    """
+    total_size = num_blocks * host_block_stride
+    if total_size % mmap.PAGESIZE:
+        raise ValueError("HiSparse host pool must be OS-page aligned.")
+    if max_chunk_bytes < mmap.PAGESIZE:
+        raise ValueError("HiSparse registration chunks must span at least one page.")
+
+    tensors: list[tuple[int, int, int]] = []
+    tensor_start = 0
+    for tensor_size in tensor_sizes:
+        if tensor_size % num_blocks:
+            raise ValueError(
+                f"HiSparse host tensor size {tensor_size} is not divisible by "
+                f"{num_blocks} blocks."
+            )
+        tensor_end = tensor_start + tensor_size
+        tensors.append((tensor_start, tensor_end, tensor_size // num_blocks))
+        tensor_start = tensor_end
+    if tensor_start > total_size:
+        raise ValueError("HiSparse host tensors exceed the shared host pool.")
+
+    def last_safe_boundary(lower: int, upper: int) -> int | None:
+        best = None
+        for start, end, block_bytes in tensors:
+            first_block = max(0, (lower - start) // block_bytes + 1)
+            last_block = min(num_blocks, (min(upper, end) - start) // block_bytes)
+            if first_block > last_block:
+                continue
+            divisor = math.gcd(block_bytes, mmap.PAGESIZE)
+            if start % divisor:
+                continue
+            period = mmap.PAGESIZE // divisor
+            if period == 1:
+                residue = 0
+            else:
+                # Find block boundaries that are also OS-page boundaries.
+                residue = (
+                    (-start // divisor) * pow(block_bytes // divisor, -1, period)
+                ) % period
+            block = last_block - (last_block - residue) % period
+            if block >= first_block:
+                candidate = start + block * block_bytes
+                best = candidate if best is None else max(best, candidate)
+
+        if upper >= tensor_start:
+            candidate = upper // mmap.PAGESIZE * mmap.PAGESIZE
+            if candidate >= tensor_start and candidate > lower:
+                best = candidate if best is None else max(best, candidate)
+        return best
+
+    ranges: list[tuple[int, int]] = []
+    start = 0
+    while total_size - start > max_chunk_bytes:
+        end = last_safe_boundary(start, start + max_chunk_bytes)
+        if end is None:
+            raise ValueError(
+                "A HiSparse host block is too large for the registration chunk limit."
+            )
+        ranges.append((start, end))
+        start = end
+    ranges.append((start, total_size))
+    return tuple(ranges)
+
+
+def allocate_hisparse_host_pools(
+    vllm_config: VllmConfig,
+    tensor_sizes: list[int],
+    num_blocks: int,
+    host_block_stride: int,
+    *,
+    use_shared_host_pool: bool,
+) -> tuple[list[torch.Tensor], list[torch.Tensor], SharedOffloadRegion | None]:
+    """Allocate private host tensors or one mmap shared by local TP ranks."""
+    if not use_shared_host_pool:
+        check_hisparse_host_memory(sum(tensor_sizes))
+        private_pools = [allocate_pinned_host_pool(size) for size in tensor_sizes]
+        return (
+            [pool for pool, _ in private_pools],
+            [registered for _, registered in private_pools],
+            None,
+        )
+
+    region = SharedOffloadRegion(
+        engine_id=(
+            f"hisparse_{vllm_config.instance_id}_"
+            f"dp{vllm_config.parallel_config.data_parallel_index}"
+        ),
+        num_chunks=1,
+        rank=0,
+        kv_bytes_per_chunk=num_blocks * host_block_stride,
+        cpu_page_size=sum(tensor_sizes),
+        barrier=get_tp_group().barrier,
+        creator_memory_check=check_hisparse_host_memory,
+        populate_only_on_creator=True,
+    )
+    try:
+        for start, end in _hisparse_registration_ranges(
+            tensor_sizes, num_blocks, host_block_stride
+        ):
+            tensor = region.base_tensor[start:end]
+            pin_tensor(tensor)
+            region.pinned_addresses.append(tensor.data_ptr())
+            region.is_pinned = True
+        pools = [
+            region.create_next_canonical_view(size).view(-1) for size in tensor_sizes
+        ]
+    except Exception:
+        region.cleanup()
+        raise
+    return pools, [], region
+
+
+def release_pinned_state(
+    runtimes: list[HiSparseRuntime],
+    pinned_host_pools: list[torch.Tensor],
+    shared_host_region: SharedOffloadRegion | None = None,
+) -> None:
+    """Synchronize and release registered host KV pools."""
+    if pinned_host_pools or shared_host_region is not None:
+        try:
+            torch.accelerator.synchronize()
+        except RuntimeError as e:
+            logger.warning(
+                "HiSparse: CUDA context unusable at teardown (%s); leaving "
+                "%d host-pool tensors pinned for kernel exit reclaim.",
+                e,
+                len(pinned_host_pools) + int(shared_host_region is not None),
+            )
+            return
+
+        cudart = torch.cuda.cudart()
+        release_start = time.perf_counter()
+        freed_bytes = 0
+        while pinned_host_pools:
+            tensor = pinned_host_pools[-1]
+            err = cudart.cudaHostUnregister(tensor.data_ptr())
+            if err.value != 0:
+                logger.warning(
+                    "HiSparse: cudaHostUnregister failed (code=%d); leaving "
+                    "%d host-pool tensors pinned for kernel exit reclaim.",
+                    err.value,
+                    len(pinned_host_pools),
+                )
+                cudart.cudaGetLastError()
+                break
+            freed_bytes += tensor.nbytes
+            pinned_host_pools.pop()
+        if freed_bytes:
+            logger.info(
+                "HiSparse: unpinned %.1f GiB of host pool in %.1fs.",
+                freed_bytes / 2**30,
+                time.perf_counter() - release_start,
+            )
+
+    for runtime in runtimes:
+        del runtime._host_cache
+        del runtime.registered_host_pool
+        del runtime.hot_backing
+    if shared_host_region is not None:
+        shared_host_region.cleanup()
+
+
+def hisparse_prefill_staging_remap(
+    block_table: torch.Tensor, block_size: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Renumber blocks without a data-dependent CUDA output allocation."""
+    block_ids = torch.cat(
+        (
+            torch.zeros(1, dtype=torch.int32, device=block_table.device),
+            block_table.flatten().clamp(min=0).to(torch.int32),
+        )
+    )
+    sorted_ids, permutation = torch.sort(block_ids)
+    is_new = torch.cat(
+        (
+            torch.ones(1, dtype=torch.bool, device=block_table.device),
+            sorted_ids[1:] != sorted_ids[:-1],
+        )
+    )
+    compact_sorted = torch.cumsum(is_new, dim=0, dtype=torch.int32) - 1
+    inverse = torch.empty_like(compact_sorted)
+    inverse.scatter_(0, permutation, compact_sorted)
+
+    unique_ids = torch.full_like(sorted_ids, -1)
+    unique_ids.scatter_(0, compact_sorted, sorted_ids)
+    valid = (
+        torch.arange(sorted_ids.numel(), device=block_table.device)
+        <= compact_sorted[-1]
+    )
+    unique_ids = torch.where(valid, unique_ids, -1)
+    new_bt = inverse[1:].view_as(block_table)
+    offsets = torch.arange(block_size, dtype=torch.int32, device=block_table.device)
+    row_ids = (
+        torch.where(
+            valid[:, None],
+            unique_ids[:, None] * block_size + offsets,
+            -1,
+        )
+    ).view(1, -1)
+    return new_bt, row_ids
+
+
+@dataclass
+class HiSparsePrefillStagingPlan:
+    block_table: torch.Tensor
+    row_ids: torch.Tensor
+    dst_rows: torch.Tensor
+    miss_mask: torch.Tensor
+    block_size: int
+    # Host rows with a GPU-resident copy (adopted shadow pages): the flat
+    # resident-cache row to read instead of DMAing from host, -1 for misses.
+    gpu_row_ids: torch.Tensor | None = None
+    gpu_source_key: tuple[int, int] | None = None
+
+    def ensure_gpu_sources(
+        self,
+        resident_block_table: torch.Tensor,
+        resident_block_size: int,
+    ) -> None:
+        """Resolve which staged rows can be served from the resident cache.
+
+        Computed once per plan (the resident block table is shared by every
+        layer in the group); non-null resident pages become miss_mask=0 rows
+        gathered device-to-device by ``gather_prefill_cache``.
+        """
+        source_key = (resident_block_table.data_ptr(), resident_block_size)
+        if self.gpu_source_key == source_key:
+            return
+        block_size = self.block_size
+        if resident_block_size <= 0 or block_size % resident_block_size != 0:
+            return
+        device = self.row_ids.device
+        num_unique = self.row_ids.shape[1] // block_size
+        host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size
+        new_bt = self.block_table.to(torch.int64)
+        num_rows, num_cols = new_bt.shape
+        if num_rows == 0 or resident_block_table.shape[0] < num_rows:
+            return
+        # One representative (row, col) per unique host block: any request
+        # referencing the block holds an equivalent (refcounted) resident view.
+        flat_pos = torch.arange(num_rows * num_cols, device=device)
+        rep = torch.full(
+            (num_unique,), num_rows * num_cols, device=device, dtype=torch.int64
+        )
+        rep.scatter_reduce_(
+            0, new_bt.reshape(-1), flat_pos, reduce="amin", include_self=True
+        )
+        rep_row = (rep // num_cols).clamp(max=resident_block_table.shape[0] - 1)
+        rep_col = rep % num_cols
+        pages_per_block = block_size // resident_block_size
+        res_pos = rep_col[:, None] * pages_per_block + torch.arange(
+            pages_per_block, device=device
+        )
+        res_cols = resident_block_table.shape[1]
+        res_blocks = torch.where(
+            res_pos < res_cols,
+            torch.gather(
+                resident_block_table.to(torch.int64)[rep_row],
+                1,
+                res_pos.clamp(max=max(res_cols - 1, 0)),
+            ),
+            torch.zeros_like(res_pos),
+        )
+        offsets = torch.arange(block_size, device=device)
+        per_off_block = res_blocks[:, offsets // resident_block_size]
+        gpu_rows = (
+            per_off_block * resident_block_size
+            + (offsets % resident_block_size)[None, :]
+        )
+        # Block id 0 is the null block in both id spaces.
+        valid_rows = self.row_ids.view(num_unique, block_size) >= 0
+        hit = (per_off_block > 0) & (host_ids[:, None] > 0) & valid_rows
+        self.gpu_row_ids = torch.where(hit, gpu_rows, -1).reshape(1, -1).to(torch.int32)
+        self.miss_mask = ((self.gpu_row_ids < 0) & valid_rows.view(1, -1)).int()
+        self.gpu_source_key = source_key
+
+
+def build_hisparse_prefill_staging_plan(
+    block_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    staging_block_capacity: int,
+) -> HiSparsePrefillStagingPlan:
+    """Build an asynchronous layer-independent host-cache staging remap."""
+    device = block_table.device
+    used = (seq_lens.to(torch.int64) + block_size - 1) // block_size
+    packed_offsets = torch.arange(staging_block_capacity, device=device)
+    row_ends = torch.cumsum(used, dim=0)
+    rows = torch.searchsorted(row_ends, packed_offsets, right=True)
+    valid = rows < block_table.shape[0]
+    rows = rows.clamp(max=block_table.shape[0] - 1)
+    row_starts = torch.cat((torch.zeros_like(row_ends[:1]), row_ends[:-1]))
+    columns = packed_offsets - row_starts[rows]
+    valid &= columns < block_table.shape[1]
+    flat_positions = rows * block_table.shape[1] + columns.clamp(
+        max=block_table.shape[1] - 1
+    )
+    referenced = torch.where(valid, block_table.flatten()[flat_positions], 0)
+    referenced_bt, row_ids = hisparse_prefill_staging_remap(referenced, block_size)
+    scratch_index = block_table.numel()
+    scatter_positions = torch.where(valid, flat_positions, scratch_index)
+    remapped = torch.where(valid, referenced_bt.flatten(), 0)
+    new_bt_storage = torch.zeros(scratch_index + 1, dtype=torch.int32, device=device)
+    new_bt_storage.scatter_(0, scatter_positions, remapped)
+    new_bt = new_bt_storage[:-1].view_as(block_table)
+    valid_rows = row_ids >= 0
+    dst_rows = torch.arange(row_ids.shape[1], dtype=torch.int32, device=device).view(
+        1, -1
+    )
+    return HiSparsePrefillStagingPlan(
+        block_table=new_bt,
+        row_ids=row_ids,
+        dst_rows=dst_rows,
+        miss_mask=valid_rows.to(torch.int32),
+        block_size=block_size,
+    )
+
+
+def _has_hisparse_ops() -> bool:
+    if not hasattr(torch.ops, "_C_cache_ops"):
+        return False
+    return (
+        hasattr(torch.ops._C_cache_ops, "hisparse_resolve_residency")
+        and hasattr(torch.ops._C_cache_ops, "hisparse_invalidate_written_slots")
+        and hasattr(torch.ops._C_cache_ops, "hisparse_gather_plan")
+        and hasattr(torch.ops._C_cache_ops, "hisparse_gather_compact")
+    )
+
+
+class _SharedTopKState:
+    """CUDA-graph-safe residency result shared by index-sharing layers."""
+
+    def __init__(self, device: torch.device, max_rows: int, top_k: int) -> None:
+        self.device_topk_rows = torch.full(
+            (max_rows, top_k), -1, dtype=torch.int32, device=device
+        )
+        self.physical_topk_indices = torch.empty_like(self.device_topk_rows)
+        self.swap_host_physical_rows = torch.empty(
+            (max_rows, top_k), dtype=torch.int32, device=device
+        )
+        self.swap_device_physical_rows = torch.empty_like(self.swap_host_physical_rows)
+        self.swap_counts = torch.empty(max_rows, dtype=torch.int32, device=device)
+        self.valid_topk_counts = torch.empty(max_rows, dtype=torch.int32, device=device)
+
+
+def _create_shared_topk_state(
+    device: torch.device, max_rows: int, top_k: int
+) -> _SharedTopKState:
+    return _SharedTopKState(device, max_rows, top_k)
+
+
+def _create_copy_stream(device: torch.device) -> torch.Stream:
+    return torch.Stream(device=device)
+
+
+class HiSparseIndexGroup:
+    """GPU index-resolution state shared by one leader and its followers."""
+
+    leader: HiSparseRuntime
+
+    def __init__(
+        self,
+        device: torch.device,
+        max_num_reqs: int,
+        region_stride: int,
+        max_swap_rows: int,
+        top_k: int,
+        copy_stream: torch.Stream | None = None,
+        logical_topk_ready: torch.Event | None = None,
+    ) -> None:
+        self.device_global_indices = torch.full(
+            (max_num_reqs, region_stride),
+            -1,
+            dtype=torch.int32,
+            device=device,
+        )
+        self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device)
+        self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous()
+        self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k)
+        self.copy_stream = (
+            copy_stream if copy_stream is not None else _create_copy_stream(device)
+        )
+        self.logical_topk_ready = logical_topk_ready
+        self.followers: list[HiSparseRuntime] = []
+        self.swap_stats = torch.zeros(2, dtype=torch.uint64, device=device)
+        self.swap_stats_host = torch.empty(
+            2, dtype=torch.uint64, device="cpu", pin_memory=True
+        )
+        self.stats_row_bytes = 0
+
+
+class HiSparseRuntime:
+    """Per-cache host/hot data plane and GPU replacement state."""
+
+    hot: PagedCacheView
+    hot_backing: torch.Tensor
+    hot_block_table: torch.Tensor
+    _host_cache: torch.Tensor
+    registered_host_pool: torch.Tensor
+
+    def __init__(
+        self,
+        config: ResolvedHiSparseConfig,
+        max_num_reqs: int,
+        row_width: int,
+        kv_dtype: torch.dtype,
+        device: torch.device | str,
+        max_swap_rows: int | None = None,
+        index_group: HiSparseIndexGroup | None = None,
+        copy_stream: torch.Stream | None = None,
+        logical_topk_ready: torch.Event | None = None,
+    ) -> None:
+        if not _has_hisparse_ops():
+            raise RuntimeError(
+                "HiSparse requires its compiled _C_cache_ops CUDA kernels "
+                "(host-resident decode has no Python fallback). Rebuild vLLM "
+                "from source so "
+                "csrc/libtorch_stable/hisparse_kernels.cu is included."
+            )
+        self.max_num_reqs = max_num_reqs
+        self.row_width = row_width
+        self.kv_dtype = kv_dtype
+        self.device = torch.device(device)
+        # Logical slots per request. Physical rows come from its ephemeral HMA
+        # block table and need not be contiguous in the shared slab.
+        self.region_stride = config.device_buffer_size
+
+        row_bytes = row_width * kv_dtype.itemsize
+        if row_bytes % 16 != 0:
+            raise ValueError(
+                f"HiSparse requires 16-byte aligned KV rows, got {row_bytes}B."
+            )
+
+        self._layer_ready_event: torch.Event | None = None
+        self._swap_staged = False
+        self._swap_step = 0
+        self.is_group_leader = index_group is None
+        if index_group is None:
+            index_group = HiSparseIndexGroup(
+                self.device,
+                max_num_reqs,
+                self.region_stride,
+                max_swap_rows or max_num_reqs,
+                config.top_k,
+                copy_stream,
+                logical_topk_ready,
+            )
+            index_group.leader = self
+        else:
+            index_group.followers.append(self)
+        self.index_group = index_group
+        index_group.stats_row_bytes += row_bytes
+
+        self.eager_host_mirror = config.eager_host_mirror
+        self.resident_source_index = -1
+        self.request_state_indices: torch.Tensor | None = None
+        self.shared_host_region: SharedOffloadRegion | None = None
+
+    @property
+    def host_cache(self) -> torch.Tensor:
+        return self._host_cache
+
+    def bind_hot_cache(
+        self,
+        raw_tensor: torch.Tensor,
+        *,
+        byte_offset: int,
+        block_stride: int,
+        num_blocks: int,
+        block_size: int,
+        block_table: torch.Tensor,
+    ) -> None:
+        """Bind this runtime's strided view into the shared GPU HMA slab."""
+        self.hot = PagedCacheView.bind(
+            raw_tensor,
+            dtype=self.kv_dtype,
+            row_width=self.row_width,
+            byte_offset=byte_offset,
+            block_stride=block_stride,
+            num_blocks=num_blocks,
+            block_size=block_size,
+        )
+        self.hot_backing = raw_tensor
+        self.hot_block_table = block_table
+
+    def bind_source_cache(
+        self,
+        kv_cache: torch.Tensor,
+        *,
+        registered_host_pool: torch.Tensor | None = None,
+    ) -> None:
+        if kv_cache.dtype != self.kv_dtype or kv_cache.shape[-1] != self.row_width:
+            raise ValueError(
+                "HiSparse runtime bound to a KV cache with mismatched "
+                f"layout: expected ({self.row_width}, {self.kv_dtype}), got "
+                f"({kv_cache.shape[-1]}, {kv_cache.dtype})."
+            )
+        # Host-resident pool: the cache itself is the only full-size store;
+        # every allocated slot is written before the indexer can select it.
+        if kv_cache.device.type != "cpu":
+            raise ValueError(
+                "HiSparse requires a host-resident KV pool; got a KV cache "
+                f"on {kv_cache.device}."
+            )
+        # The pool is pinned via cudaHostRegister (exact-size, deterministic
+        # unpin at shutdown); torch's is_pinned() only recognizes its own
+        # caching-host-allocator memory, so also accept ranges the model
+        # allocator explicitly registered.
+        if not (kv_cache.is_pinned() or registered_host_pool is not None):
+            raise ValueError("HiSparse host-resident KV pool must be pinned memory.")
+
+        self._host_cache = kv_cache.view(-1, kv_cache.shape[-1])
+        self.registered_host_pool = (
+            registered_host_pool if registered_host_pool is not None else kv_cache
+        )
+
+    def gather_prefill_cache(
+        self,
+        kv_cache: torch.Tensor,
+        plan: HiSparsePrefillStagingPlan,
+        resident_cache: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Gather this runtime's host cache using a shared staging plan."""
+        if kv_cache.shape[1] != plan.block_size:
+            raise ValueError(
+                f"HiSparse staging block size {plan.block_size} does not match "
+                f"the KV cache block size {kv_cache.shape[1]}."
+            )
+        row_width = kv_cache.shape[-1]
+
+        staged = torch.empty(
+            (
+                plan.row_ids.shape[1] // plan.block_size,
+                plan.block_size,
+                row_width,
+            ),
+            dtype=kv_cache.dtype,
+            device=plan.block_table.device,
+        )
+        if plan.gpu_row_ids is not None and resident_cache is not None:
+            gpu_rows = plan.gpu_row_ids[0].to(torch.long)
+            src = gpu_rows.clamp_min(0)
+            resident_block_size = resident_cache.shape[1]
+            rows = resident_cache[src // resident_block_size, src % resident_block_size]
+            if rows.dtype != staged.dtype:
+                rows = rows.contiguous().view(staged.dtype)
+            staged.view(-1, row_width).copy_(rows)
+        torch.ops._C_cache_ops.hisparse_gather_plan(
+            kv_cache.view(-1, row_width),
+            staged,
+            plan.row_ids,
+            plan.dst_rows,
+            plan.miss_mask,
+            None,
+            None,
+            0,
+        )
+        return staged
+
+    def reset_hot_state(self) -> None:
+        """Drop all hot-buffer bookkeeping (hits become misses)."""
+        group = self.index_group
+        compute_stream = current_stream()
+        compute_stream.wait_stream(group.copy_stream)
+        group.device_global_indices.fill_(-1)
+        group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots))
+        group.swap_stats.zero_()
+        group.copy_stream.wait_stream(compute_stream)
+
+    def invalidate_sorted_slots(
+        self,
+        sorted_slots: torch.Tensor,
+        state_indices: torch.Tensor,
+    ) -> None:
+        """Drop hot copies using shared, preprocessed invalidation inputs."""
+        device_global_indices = self.index_group.device_global_indices
+        active_indices = device_global_indices.index_select(0, state_indices)
+        positions = torch.searchsorted(sorted_slots, active_indices)
+        positions.clamp_(max=sorted_slots.numel() - 1)
+        active_indices.masked_fill_(
+            sorted_slots[positions] == active_indices,
+            -1,
+        )
+        device_global_indices.index_copy_(0, state_indices, active_indices)
+
+    def invalidate_written_slots(
+        self,
+        written_slots: torch.Tensor,
+        req_id_per_token: torch.Tensor,
+    ) -> None:
+        num_tokens = min(written_slots.numel(), req_id_per_token.numel())
+        if num_tokens == 0:
+            return
+        assert self.request_state_indices is not None
+        torch.ops._C_cache_ops.hisparse_invalidate_written_slots(
+            self.index_group.device_global_indices,
+            self.request_state_indices,
+            req_id_per_token[:num_tokens],
+            written_slots[:num_tokens],
+        )
+
+    def begin_forward(self) -> None:
+        self._swap_step = 0
+
+    def _step_rows(self, num_tokens: int) -> slice:
+        start = self._swap_step * num_tokens
+        stop = start + num_tokens
+        if stop > self.index_group.shared_topk.physical_topk_indices.shape[0]:
+            raise ValueError(
+                "HiSparse swap rows exceed the configured speculative decode "
+                f"capacity: stop={stop}, capacity="
+                f"{self.index_group.shared_topk.physical_topk_indices.shape[0]}."
+            )
+        self._swap_step += 1
+        return slice(start, stop)
+
+    def _resolve_residency(
+        self,
+        *,
+        resident: HiSparseCacheHandle | None = None,
+        req_id_per_token: torch.Tensor,
+        block_table: torch.Tensor,
+        topk_indices: torch.Tensor,
+        block_size: int,
+        return_valid_counts: bool = False,
+        shared_rows: slice,
+        attention_indices_out: torch.Tensor | None = None,
+        valid_counts_out: torch.Tensor | None = None,
+    ) -> None:
+        """Resolve logical top-k positions and compact rows requiring swaps."""
+        num_tokens = topk_indices.shape[0]
+        hot = self.hot
+        host_cache = self.host_cache
+        assert hot.block_size == block_size
+        request_state_indices = self.request_state_indices
+        assert request_state_indices is not None
+        group = self.index_group
+        shared = group.shared_topk
+
+        device_topk_rows = shared.device_topk_rows[shared_rows]
+        valid_topk_counts = None
+        if return_valid_counts:
+            valid_topk_counts = (
+                shared.valid_topk_counts[shared_rows]
+                if valid_counts_out is None
+                else valid_counts_out
+            )
+        physical_topk_indices = (
+            shared.physical_topk_indices[shared_rows]
+            if attention_indices_out is None
+            else attention_indices_out
+        )
+        swap_host_physical_rows = shared.swap_host_physical_rows[shared_rows]
+        swap_device_physical_rows = shared.swap_device_physical_rows[shared_rows]
+        swap_counts = shared.swap_counts[shared_rows]
+
+        # Padded rows are skipped by the kernel (request_state_indices) and must
+        # come out as -1 so the attention kernel masks them.
+        torch.ops._C_cache_ops.hisparse_resolve_residency(
+            host_cache,
+            hot.cache,
+            self.hot_block_table,
+            topk_indices,
+            device_topk_rows,
+            group.device_global_indices,
+            group.lru_slots,
+            request_state_indices,
+            self.region_stride,
+            None,
+            group.swap_stats,
+            physical_topk_indices,
+            hot.attention_block_stride,
+            req_id_per_token[:num_tokens].contiguous(),
+            block_table,
+            block_size,
+            None,
+            valid_topk_counts,
+            swap_host_physical_rows,
+            swap_device_physical_rows,
+            swap_counts,
+            resident.block_table if resident is not None else None,
+            resident.view.block_size
+            if resident is not None and resident.view is not None
+            else 0,
+            0,
+        )
+
+    def _swap_rows(self, shared_rows: slice) -> None:
+        hot = self.hot
+        shared = self.index_group.shared_topk
+        torch.ops._C_cache_ops.hisparse_gather_compact(
+            self.host_cache,
+            hot.cache,
+            shared.swap_host_physical_rows[shared_rows],
+            shared.swap_device_physical_rows[shared_rows],
+            shared.swap_counts[shared_rows],
+        )
+
+    def _resolve_and_stage_group(
+        self,
+        *,
+        resident: HiSparseCacheHandle,
+        req_id_per_token: torch.Tensor,
+        block_table: torch.Tensor,
+        logical_topk_indices: torch.Tensor,
+        block_size: int,
+        return_valid_counts: bool,
+        shared_rows: slice,
+        attention_indices_out: torch.Tensor | None,
+        valid_counts_out: torch.Tensor | None,
+    ) -> None:
+        group = self.index_group
+        compute_stream = current_stream()
+        if group.logical_topk_ready is not None and not in_piecewise_cudagraph():
+            group.copy_stream.wait_event(group.logical_topk_ready)
+        else:
+            group.copy_stream.wait_stream(compute_stream)
+        with group.copy_stream:
+            self._resolve_residency(
+                resident=resident,
+                req_id_per_token=req_id_per_token,
+                block_table=block_table,
+                topk_indices=logical_topk_indices,
+                block_size=block_size,
+                return_valid_counts=return_valid_counts,
+                shared_rows=shared_rows,
+                attention_indices_out=attention_indices_out,
+                valid_counts_out=valid_counts_out,
+            )
+            runtimes = [self]
+            if resident.decode_batch:
+                runtimes.extend(group.followers)
+            for runtime in runtimes:
+                runtime._swap_rows(shared_rows)
+                if runtime._layer_ready_event is None:
+                    runtime._layer_ready_event = torch.Event()
+                runtime._layer_ready_event.record(group.copy_stream)
+                runtime._swap_staged = True
+
+    def _stage_this_layer(self, shared_rows: slice) -> None:
+        group = self.index_group
+        compute_stream = current_stream()
+        group.copy_stream.wait_stream(compute_stream)
+        with group.copy_stream:
+            self._swap_rows(shared_rows)
+            if self._layer_ready_event is None:
+                self._layer_ready_event = torch.Event()
+            self._layer_ready_event.record(group.copy_stream)
+            self._swap_staged = True
+
+    def wait_for_staged_swap(self) -> None:
+        if not self._swap_staged or self._layer_ready_event is None:
+            raise RuntimeError("HiSparse layer swap was not prefetched.")
+        current_stream().wait_event(self._layer_ready_event)
+        self._swap_staged = False
+
+    def swap_in(
+        self,
+        *,
+        resident: HiSparseCacheHandle,
+        req_id_per_token: torch.Tensor,
+        block_table: torch.Tensor,
+        logical_topk_indices: torch.Tensor,
+        block_size: int,
+        return_valid_counts: bool = False,
+        attention_indices_out: torch.Tensor | None = None,
+        valid_counts_out: torch.Tensor | None = None,
+    ) -> HiSparseTopKResult:
+        """Resolve once per group and ensure this layer's rows are available."""
+        num_tokens = logical_topk_indices.shape[0]
+        shared_rows = self._step_rows(num_tokens)
+        group = self.index_group
+        hot = self.hot
+        assert hot.block_size == block_size
+        assert hot.attention_block_stride == group.leader.hot.attention_block_stride
+
+        if self.is_group_leader:
+            self._resolve_and_stage_group(
+                resident=resident,
+                req_id_per_token=req_id_per_token,
+                block_table=block_table,
+                logical_topk_indices=logical_topk_indices,
+                block_size=block_size,
+                return_valid_counts=return_valid_counts,
+                shared_rows=shared_rows,
+                attention_indices_out=attention_indices_out,
+                valid_counts_out=valid_counts_out,
+            )
+
+        if not self._swap_staged:
+            self._stage_this_layer(shared_rows)
+        self.wait_for_staged_swap()
+
+        physical_topk_indices = group.shared_topk.physical_topk_indices[shared_rows]
+        if return_valid_counts:
+            return (
+                physical_topk_indices,
+                group.shared_topk.valid_topk_counts[shared_rows],
+            )
+        return physical_topk_indices
+
+
+class HiSparseCacheHandle:
+    """Attention-facing handle for resident KV and sparse offload state."""
+
+    def __init__(self, runtime: HiSparseRuntime) -> None:
+        self.view: PagedCacheView | None = None
+        self.block_table: torch.Tensor | None = None
+        self.source_block_table: torch.Tensor | None = None
+        self.slot_mapping: torch.Tensor | None = None
+        self.runtime = runtime
+        self.dummy_batch = False
+        self.decode_batch = False
+        self.all_context_pages_resident = True
+        self.num_actual_tokens = 0
+        self.num_decode_tokens = 0
+        self.req_id_per_token: torch.Tensor | None = None
+        self.host_mirror_required = False
+        self.mirror_from_resident = False
+        self.mirror_slot_mapping: torch.Tensor | None = None
+        self.mirror_staging_cache: torch.Tensor | None = None
+        self.mirror_staging_slots: torch.Tensor | None = None
+        self.submit_layer_mirror: Callable[[], None] | None = None
+        self.index_group_caches: list[HiSparseCacheHandle] = [self]
+
+    def prepare_group_for_batch(self, attn_metadata: Any | None) -> None:
+        assert self.runtime.is_group_leader
+        for cache in self.index_group_caches:
+            cache._prepare_for_batch(attn_metadata)
+
+    def _prepare_for_batch(self, attn_metadata: Any | None) -> None:
+        self.dummy_batch = attn_metadata is None
+        self.runtime.begin_forward()
+        self.num_actual_tokens = (
+            attn_metadata.num_actual_tokens if attn_metadata is not None else 0
+        )
+        self.num_decode_tokens = (
+            attn_metadata.num_decode_tokens if attn_metadata is not None else 0
+        )
+        self.req_id_per_token = (
+            attn_metadata.req_id_per_token if attn_metadata is not None else None
+        )
+        self.decode_batch = is_hisparse_decode_batch(attn_metadata)
+        self.host_mirror_required = attn_metadata is not None and (
+            not self.decode_batch or self.runtime.eager_host_mirror
+        )
+
+    def write_target(
+        self, num_input_rows: int, num_slot_rows: int
+    ) -> tuple[torch.Tensor, torch.Tensor, int]:
+        assert self.view is not None and self.slot_mapping is not None
+        num_rows = min(num_input_rows, num_slot_rows, self.num_actual_tokens)
+        if self.decode_batch:
+            num_rows = min(num_rows, self.runtime.max_num_reqs)
+        return self.view.cache, self.slot_mapping[:num_rows], num_rows
+
+    def mirror_write_target(
+        self, num_rows: int
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        if (
+            not self.host_mirror_required
+            or self.decode_batch
+            or self.mirror_from_resident
+        ):
+            return None
+        cache = self.mirror_staging_cache
+        slots = self.mirror_staging_slots
+        if cache is None or slots is None:
+            raise RuntimeError("HiSparse prefill mirror staging is not bound.")
+        if num_rows > slots.numel():
+            raise RuntimeError(
+                "HiSparse prefill mirror exceeds staging capacity: "
+                f"{num_rows} > {slots.numel()}."
+            )
+        return cache, slots[:num_rows]
+
+    def finish_kv_update(self) -> None:
+        if self.dummy_batch:
+            return
+        if (
+            self.submit_layer_mirror is not None
+            and not self.decode_batch
+            and get_forward_context().cudagraph_runtime_mode != CUDAGraphMode.FULL
+        ):
+            self.submit_layer_mirror()
+
+    def bind_cache(
+        self,
+        raw_tensor: torch.Tensor,
+        *,
+        byte_offset: int,
+        block_stride: int,
+        num_blocks: int,
+        block_size: int,
+        block_table: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ) -> None:
+        self.view = PagedCacheView.bind(
+            raw_tensor,
+            dtype=self.runtime.kv_dtype,
+            row_width=self.runtime.row_width,
+            byte_offset=byte_offset,
+            block_stride=block_stride,
+            num_blocks=num_blocks,
+            block_size=block_size,
+        )
+        self.block_table = block_table
+        self.slot_mapping = slot_mapping
+
+    def swap_in(
+        self,
+        req_id_per_token: torch.Tensor,
+        block_table: torch.Tensor,
+        logical_topk_indices: torch.Tensor,
+        *,
+        block_size: int,
+        return_valid_counts: bool = False,
+        attention_indices_out: torch.Tensor | None = None,
+        valid_counts_out: torch.Tensor | None = None,
+    ) -> HiSparseTopKResult:
+        return self.runtime.swap_in(
+            resident=self,
+            req_id_per_token=req_id_per_token[: logical_topk_indices.shape[0]],
+            block_table=block_table,
+            logical_topk_indices=logical_topk_indices,
+            block_size=block_size,
+            return_valid_counts=return_valid_counts,
+            attention_indices_out=attention_indices_out,
+            valid_counts_out=valid_counts_out,
+        )
+
+
+def initialize_hisparse_runtime_buffers(
+    cache_handles: list[HiSparseCacheHandle],
+    *,
+    max_num_reqs: int,
+    max_num_batched_tokens: int,
+) -> None:
+    """Allocate shared request state and per-layer prefill staging buffers."""
+    assert cache_handles
+    resident = cache_handles[0].view
+    assert resident is not None
+    device = resident.cache.device
+    request_state_indices = torch.full(
+        (max_num_reqs,), -1, dtype=torch.int32, device=device
+    )
+    staging_blocks = (
+        max_num_batched_tokens + resident.block_size - 1
+    ) // resident.block_size
+    mirror_staging_caches = torch.empty(
+        (
+            len(cache_handles),
+            staging_blocks,
+            resident.block_size,
+            resident.cache.shape[-1],
+        ),
+        dtype=resident.cache.dtype,
+        device=device,
+    )
+    mirror_staging_slots = torch.arange(
+        max_num_batched_tokens, dtype=torch.int64, device=device
+    )
+    for layer_index, cache_handle in enumerate(cache_handles):
+        cache_handle.runtime.request_state_indices = request_state_indices
+        cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index]
+        cache_handle.mirror_staging_slots = mirror_staging_slots
+
+
+def create_hisparse_cache_handle(
+    vllm_config: VllmConfig,
+    model_top_k: int,
+    *,
+    is_index_group_leader: bool,
+    row_width: int,
+    kv_dtype: torch.dtype,
+    index_group: HiSparseMLAIndexGroup | None = None,
+    device: torch.device | str | None = None,
+) -> HiSparseCacheHandle | None:
+    config = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k)
+    if config is None:
+        return None
+
+    max_num_reqs = vllm_config.scheduler_config.max_num_seqs
+    max_swap_rows = _get_max_swap_rows(vllm_config)
+    if device is None:
+        device = torch.device(
+            current_platform.device_type, torch.accelerator.current_device_index()
+        )
+
+    hisparse_group = (
+        None
+        if is_index_group_leader or index_group is None
+        else index_group.hisparse_group
+    )
+    runtime = HiSparseRuntime(
+        config=config,
+        max_num_reqs=max_num_reqs,
+        max_swap_rows=max_swap_rows,
+        row_width=row_width,
+        kv_dtype=kv_dtype,
+        device=device,
+        index_group=hisparse_group,
+        copy_stream=index_group.side_stream if index_group is not None else None,
+        logical_topk_ready=(
+            index_group.logical_topk_ready
+            if index_group is not None and index_group.has_indexer
+            else None
+        ),
+    )
+    if is_index_group_leader and index_group is not None:
+        index_group.hisparse_group = runtime.index_group
+    logger.info_once(
+        "Enabled experimental HiSparse HMA hot cache: top_k=%d, "
+        "device_buffer_size=%d (%d LRU rows), max_num_seqs=%d.",
+        config.top_k,
+        config.device_buffer_size,
+        config.device_buffer_size,
+        max_num_reqs,
+    )
+    handle = HiSparseCacheHandle(runtime)
+    speculative_config = vllm_config.speculative_config
+    handle.mirror_from_resident = bool(
+        vllm_config.scheduler_config.async_scheduling
+        and speculative_config is not None
+        and speculative_config.uses_draft_kv_cache()
+    )
+    return handle
diff --git a/vllm/v1/hisparse/types.py b/vllm/v1/hisparse/types.py
new file mode 100644
index 000000000000..269793ea8521
--- /dev/null
+++ b/vllm/v1/hisparse/types.py
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class SparseKVPageTransfer:
+    """Copy one logical KV page between cache-manager and worker-owned tiers."""
+
+    transfer_id: int
+    host_block_id: int
+    resident_block_ids: tuple[int, ...]
+    after_forward: bool
+    restore: bool = False
+
+
+@dataclass(frozen=True)
+class SparseKVRowMirror:
+    """Mirror one contiguous resident-row span into the host cache."""
+
+    source_starts: tuple[int, ...]
+    destination_start: int
+    num_rows: int
+
+
+@dataclass
+class SparseKVOffloadCommand:
+    """Opaque scheduler-to-worker command for one model step."""
+
+    page_transfers: list[SparseKVPageTransfer]
diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py
index 1888404a573a..fc8b88921a39 100644
--- a/vllm/v1/kv_cache_interface.py
+++ b/vllm/v1/kv_cache_interface.py
@@ -147,11 +147,14 @@ class KVCacheSpecKind(str, Enum):
     UNKNOWN = "unknown"
 
 
+class SparseCacheRole(str, Enum):
+    SPARSE = "sparse"
+    INDEXER = "indexer"
+
+
 @dataclass(frozen=True)
 class KVCacheSpec:
-    """
-    A base class for specifying the KV cache format of one layer.
-    """
+    """A base class for specifying the KV cache format of one layer."""
 
     # number of tokens in a block
     block_size: int
@@ -161,6 +164,11 @@ def prefix_cacheable(self) -> bool:
         """Whether this spec's group participates in prefix caching."""
         return True
 
+    @property
+    def prefix_replay_tokens(self) -> int:
+        """DeepSeek-V4.1 only: bounded replay. Currently only for DSV41 SWA."""
+        return 0
+
     @property
     def num_heads(self) -> int:
         raise NotImplementedError
@@ -175,11 +183,11 @@ def state_content_size_bytes(self) -> int:
 
     @property
     def page_size_bytes(self) -> int:
-        """
-        The size of a page with `block_size` tokens in bytes.
+        """The size of a page with `block_size` tokens in bytes.
 
         Returns:
             The page size
+
         """
         raise NotImplementedError
 
@@ -192,38 +200,38 @@ def get_num_kernel_states(self, kernel_block_size: int) -> int:
             return kernel_block_size // self.tokens_per_state
         return 1
 
+    @property
+    def block_table_token_alignment(self) -> int | None:
+        return 128
+
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
-        """
-        The maximum possible memory usage of this KV cache in bytes.
+        """The maximum possible memory usage of this KV cache in bytes.
 
         Returns:
             The KV cache size in bytes
+
         """
         raise NotImplementedError
 
     def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
-        """
-        The number of block table entries needed per request, i.e. the row
+        """The number of block table entries needed per request, i.e. the row
         length of the worker-side block table for this cache group.
 
         Args:
             vllm_config: The vllm config.
             max_len: The maximum sequence length to size for, including the
                 encoder length for encoder-decoder models.
+
         """
         return cdiv(max_len, self.block_size)
 
     def copy_with_new_block_size(self, block_size: int) -> Self:
-        """
-        Create a new KVCacheSpec from self but replacing the block size.
-        """
+        """Create a new KVCacheSpec from self but replacing the block size."""
         return replace(self, block_size=block_size)
 
     @classmethod
     def merge(cls, specs: list[Self]) -> Self:
-        """
-        Merge a list of KVCacheSpec objects into a single KVCacheSpec object.
-        """
+        """Merge a list of KVCacheSpec objects into a single KVCacheSpec object."""
         if not all(spec == specs[0] for spec in specs[1:]):
             raise AssertionError(
                 "All layers in the same KV cache group must be the same."
@@ -233,9 +241,7 @@ def merge(cls, specs: list[Self]) -> Self:
     def is_uniform_with_collection(
         self, kv_cache_specs: dict[str, KVCacheSpec]
     ) -> bool:
-        """
-        Whether this KVCacheSpec is uniform with all specs of all layers.
-        """
+        """Whether this KVCacheSpec is uniform with all specs of all layers."""
         uniform_type_base_spec = KVCacheSpecRegistry.get_uniform_type_base_spec(self)
         assert uniform_type_base_spec is not None, (
             f"Unsupported KV cache spec type: {type(self)}. "
@@ -245,6 +251,25 @@ def is_uniform_with_collection(
             isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values()
         )
 
+    @property
+    def has_layer_views(self) -> bool:
+        """Whether generic allocation creates per-layer cache views.
+
+        Specs without a per-layer shape keep the raw backing tensor and lay it
+        out themselves. They also have no attention module of their own, so
+        they take no part in attention-backend or kernel-block selection.
+        """
+        return True
+
+    @property
+    def uses_slot_mapping(self) -> bool:
+        """Whether the worker computes a per-token slot mapping for this spec.
+
+        Specs that address their pages themselves (raw storage, ring buffers)
+        take no slot mapping row.
+        """
+        return True
+
 
 def group_kernel_blocks(cache: torch.Tensor, num_blocks: int) -> torch.Tensor:
     """View a kernel-block-granular layer cache with manager blocks as dim 0.
@@ -383,6 +408,70 @@ def create_kv_cache_views(
     return views
 
 
+@dataclass(frozen=True, kw_only=True)
+class HiSparseHotSpec(KVCacheSpec):
+    """Ephemeral per-request HiSparse hot-cache allocation."""
+
+    page_size: int
+    blocks_per_request: int
+
+    @property
+    def prefix_cacheable(self) -> bool:
+        return False
+
+    @property
+    def page_size_bytes(self) -> int:
+        return self.page_size
+
+    @property
+    def block_table_token_alignment(self) -> int | None:
+        return None
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        return self.blocks_per_request * self.page_size
+
+    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
+        return self.blocks_per_request
+
+    @property
+    def has_layer_views(self) -> bool:
+        return False
+
+    @property
+    def uses_slot_mapping(self) -> bool:
+        return False
+
+
+@dataclass(frozen=True, kw_only=True)
+class HiSparseResidentSpec(KVCacheSpec):
+    """Reclaimable GPU-resident pages for host-backed HiSparse KV."""
+
+    page_size: int
+
+    @property
+    def prefix_cacheable(self) -> bool:
+        return False
+
+    @property
+    def page_size_bytes(self) -> int:
+        return self.page_size
+
+    def max_admission_blocks_per_request(
+        self, max_in_flight_tokens: int, max_model_len: int
+    ) -> int:
+        num_tokens = min(max_in_flight_tokens, max_model_len)
+        return cdiv(num_tokens, self.block_size)
+
+    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
+        return cdiv(vllm_config.model_config.max_model_len, self.block_size) * (
+            self.page_size
+        )
+
+    @property
+    def has_layer_views(self) -> bool:
+        return False
+
+
 @dataclass(frozen=True, kw_only=True)
 class AttentionSpec(KVCacheSpec):
     num_kv_heads: int
@@ -445,8 +534,7 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
 
 @dataclass(frozen=True, kw_only=True)
 class FullAttentionSpec(AttentionSpec):
-    """
-    When hybrid allocator is disabled and the model contains both full
+    """When hybrid allocator is disabled and the model contains both full
     attention layers and sliding window attention layers, sliding
     window attention are regarded as full attention in KV cache manager
     (blocks are allocated for all tokens), while computed as sliding window
@@ -490,8 +578,7 @@ def merge_window_sizes(cls, window_sizes: set[int]) -> int | None:
 
     @classmethod
     def merge(cls, specs: list[Self]) -> Self:
-        """
-        Merge a list of FullAttentionSpec objects into a single
+        """Merge a list of FullAttentionSpec objects into a single
         FullAttentionSpec object.
         """
         assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (
@@ -557,8 +644,16 @@ class MLAAttentionSpec(FullAttentionSpec):
     # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults.
     alignment: int | None = None  # Default to None for no padding.
     model_version: str | None = None
+    cache_role: SparseCacheRole = SparseCacheRole.SPARSE
+    is_index_group_leader: bool = False
     storage_block_size: int | None = None
     """Token width used to view storage when it differs from the kernel block."""
+    block_stride_alignment: int | None = None
+    """Required alignment, in bytes, of the distance between consecutive
+    blocks of this cache. In block-major layouts that distance is the whole
+    block (all layers' pages), so the allocator rounds the block up to it.
+    DeepGEMM's paged sparse MQA-logits kernels address pages as
+    ``base + page * stride`` and need it 512B-aligned."""
     # Group capability enabled when any member flattens a non-causal query block
     # into decode rows. Runtime metadata still selects causal vs. non-causal mode.
     non_causal_multi_token_decode: bool = False
@@ -577,16 +672,22 @@ def merge(cls, specs: list[Self]) -> Self:
         cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
         tokens_per_state_set = set(spec.tokens_per_state for spec in specs)
         model_version_set = set(spec.model_version for spec in specs)
+        cache_role_set = {spec.cache_role for spec in specs}
+        index_group_leader_set = {spec.is_index_group_leader for spec in specs}
         storage_block_size_set = set(spec.storage_block_size for spec in specs)
+        block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
         assert (
             len(cache_dtype_str_set) == 1
             and len(tokens_per_state_set) == 1
             and len(model_version_set) == 1
+            and len(cache_role_set) == 1
+            and len(index_group_leader_set) == 1
             and len(storage_block_size_set) == 1
+            and len(block_stride_alignment_set) == 1
         ), (
             "All attention layers in the same KV cache group must use the same "
-            "quantization method, tokens per state, model version, and storage "
-            "block size."
+            "quantization method, tokens per state, model version, cache role, "
+            "index-sharing role, storage block size and block stride alignment."
         )
         merged_spec = cls(
             block_size=specs[0].block_size,
@@ -600,7 +701,10 @@ def merge(cls, specs: list[Self]) -> Self:
             cache_dtype_str=cache_dtype_str_set.pop(),
             tokens_per_state=tokens_per_state_set.pop(),
             model_version=model_version_set.pop(),
+            cache_role=cache_role_set.pop(),
+            is_index_group_leader=index_group_leader_set.pop(),
             storage_block_size=storage_block_size_set.pop(),
+            block_stride_alignment=block_stride_alignment_set.pop(),
             non_causal_multi_token_decode=any(
                 spec.non_causal_multi_token_decode for spec in specs
             ),
@@ -772,6 +876,10 @@ class CircularBufferSpec(AttentionSpec):
     reads the open group's committed keys from the ring.
     """
 
+    @property
+    def block_table_token_alignment(self) -> int | None:
+        return None
+
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         # The ring occupies one block per request for its whole lifetime.
         del vllm_config
@@ -792,6 +900,10 @@ def is_uniform_with_collection(
     def prefix_cacheable(self) -> bool:
         return False
 
+    @property
+    def uses_slot_mapping(self) -> bool:
+        return False
+
 
 @dataclass(frozen=True, kw_only=True)
 class SlidingWindowMLASpec(SlidingWindowSpec):
@@ -801,6 +913,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
     # DeepseekV4-only: see MLAAttentionSpec.model_version.
     alignment: int | None = None  # Default to None for no padding.
     model_version: str | None = None
+    bounded_replay: bool = False
 
     # MLA stores a single latent vector per state; there is no separate V.
     head_size_v: int = 0
@@ -812,6 +925,14 @@ def __post_init__(self):
         super().__post_init__()
         _apply_alignment_padding(self)
 
+    @property
+    def prefix_cacheable(self) -> bool:
+        return not self.bounded_replay
+
+    @property
+    def prefix_replay_tokens(self) -> int:
+        return self.sliding_window if self.bounded_replay else 0
+
     @classmethod
     def merge(cls, specs: list[Self]) -> Self:
         assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), (
@@ -823,16 +944,18 @@ def merge(cls, specs: list[Self]) -> Self:
         model_version_set = set(spec.model_version for spec in specs)
         sliding_window_set = set(spec.sliding_window for spec in specs)
         extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
+        bounded_replay_set = set(spec.bounded_replay for spec in specs)
         assert (
             len(cache_dtype_str_set) == 1
             and len(tokens_per_state_set) == 1
             and len(model_version_set) == 1
             and len(sliding_window_set) == 1
             and len(extra_retained_set) == 1
+            and len(bounded_replay_set) == 1
         ), (
             "All attention layers in the same KV cache group must use the same "
             "quantization method, tokens per state, model version, sliding "
-            "window size, and retained token count."
+            "window size, retained token count, and replay policy."
         )
         return cls(
             block_size=specs[0].block_size,
@@ -847,6 +970,7 @@ def merge(cls, specs: list[Self]) -> Self:
             cache_dtype_str=cache_dtype_str_set.pop(),
             tokens_per_state=tokens_per_state_set.pop(),
             model_version=model_version_set.pop(),
+            bounded_replay=bounded_replay_set.pop(),
         )
 
     def is_uniform_with_collection(
@@ -855,6 +979,7 @@ def is_uniform_with_collection(
         return all(
             isinstance(spec, SlidingWindowMLASpec)
             and spec.sliding_window == self.sliding_window
+            and spec.bounded_replay == self.bounded_replay
             for spec in kv_cache_specs.values()
         )
 
@@ -880,6 +1005,10 @@ def is_uniform_with_collection(
     def prefix_cacheable(self) -> bool:
         return False
 
+    @property
+    def uses_slot_mapping(self) -> bool:
+        return False
+
 
 @dataclass(frozen=True)
 class MambaSpec(KVCacheSpec):
@@ -919,6 +1048,10 @@ def page_size_bytes(self) -> int:
             return self.page_size_padded
         return page_size
 
+    @property
+    def block_table_token_alignment(self) -> int | None:
+        return None
+
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         if vllm_config.cache_config.mamba_cache_mode == "all":
             max_model_len = vllm_config.model_config.max_model_len
@@ -1003,9 +1136,7 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
 
 @dataclass(frozen=True)
 class CrossAttentionSpec(AttentionSpec):
-    """
-    KV cache spec for cross-attention layers in encoder-decoder models.
-    """
+    """KV cache spec for cross-attention layers in encoder-decoder models."""
 
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         # For cross-attention, we need to cache encoder states
@@ -1020,8 +1151,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
 
     @classmethod
     def merge(cls, specs: list[Self]) -> Self:
-        """
-        Merge a list of FullAttentionSpec objects into a single
+        """Merge a list of FullAttentionSpec objects into a single
         FullAttentionSpec object.
         """
         assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (
@@ -1071,8 +1201,7 @@ def merge(cls, specs: list[Self]) -> Self:
 
 @dataclass(frozen=True)
 class UniformTypeKVCacheSpecs(KVCacheSpec):
-    """
-    A KV cache spec for multiple layers with the same type of attention. Here,
+    """A KV cache spec for multiple layers with the same type of attention. Here,
     same types means always need the same number of token slots. For example,
     sliding window attentions with different window sizes are not the same type
     and should not be merged into one UniformTypeKVCacheSpecs.
@@ -1084,6 +1213,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
     def prefix_cacheable(self) -> bool:
         return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values())
 
+    @property
+    def prefix_replay_tokens(self) -> int:
+        return max(spec.prefix_replay_tokens for spec in self.kv_cache_specs.values())
+
     @property
     def first_spec(self) -> KVCacheSpec:
         """Return the first spec in the group."""
@@ -1093,6 +1226,10 @@ def first_spec(self) -> KVCacheSpec:
     def page_size_bytes(self) -> int:
         return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values())
 
+    @property
+    def block_table_token_alignment(self) -> int | None:
+        return self.first_spec.block_table_token_alignment
+
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         max_num_pages = max(
             cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes)
@@ -1116,8 +1253,7 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
 
     @classmethod
     def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:
-        """
-        Whether all layers have the same type of KV cache spec.
+        """Whether all layers have the same type of KV cache spec.
 
         Uses the registry to determine grouping base classes, so custom specs
         that inherit from FullAttentionSpec are treated as full attention.
@@ -1131,8 +1267,7 @@ def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:
 
     @classmethod
     def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None:
-        """
-        Return a SameTypeKVCacheSpecs object if all layers have the same type
+        """Return a SameTypeKVCacheSpecs object if all layers have the same type
         of KV cache spec. Return None if not.
         """
         if cls.is_uniform_type(kv_cache_specs):
@@ -1154,6 +1289,10 @@ def max_memory_usage_pages(self, vllm_config: VllmConfig) -> int:
             for spec in self.kv_cache_specs.values()
         )
 
+    @property
+    def uses_slot_mapping(self) -> bool:
+        return self.first_spec.uses_slot_mapping
+
 
 def iter_layer_specs(kv_cache_spec: KVCacheSpec) -> Collection[KVCacheSpec]:
     """The per-layer specs a KV cache group spec covers.
@@ -1239,8 +1378,7 @@ def get_kv_cache_spec_sliding_window(kv_cache_spec: KVCacheSpec) -> int | None:
 
 @dataclass
 class KVCacheTensor:
-    """
-    A class for specifying how the workers should initialize the KV cache.
+    """A class for specifying how the workers should initialize the KV cache.
 
     Placement of a set of same-shaped layers in the KV cache allocation.
     Layer ``layers[l]``'s page for block ``b`` starts at
@@ -1259,12 +1397,18 @@ class KVCacheTensor:
     layer_stride: int
     block_stride: int
     offset: int = 0  # byte offset of layers[0]'s block 0
+    host_resident: bool = False
+
+
+class KVCacheGroupRole(str, Enum):
+    DEFAULT = "default"
+    HISPARSE_SOURCE = "hisparse_source"
+    HISPARSE_INDEXER = "hisparse_indexer"
 
 
 @dataclass
 class KVCacheGroupSpec:
-    """
-    Represents a group of model layers that share the same KV cache block table.
+    """Represents a group of model layers that share the same KV cache block table.
     These layers are regarded as one layer in the KV cache manager.
     """
 
@@ -1274,15 +1418,18 @@ class KVCacheGroupSpec:
     kv_cache_spec: KVCacheSpec
     # Whether this group contains EAGLE/MTP draft attention layers.
     is_eagle_group: bool = False
+    # Host groups use the dedicated HiSparse pool; others share the device pool.
+    host_resident: bool = False
     # Whether this group is part of the externally transferable KV state.
+    # Ephemeral accelerator-side replicas are rebuilt from their transferable
+    # source group after a connector load.
     enable_kv_transfer: bool = True
+    role: KVCacheGroupRole = KVCacheGroupRole.DEFAULT
 
 
 @dataclass
 class KVCacheConfig:
-    """
-    The KV cache configuration of a model.
-    """
+    """The KV cache configuration of a model."""
 
     num_blocks: int
     """The number of KV cache blocks"""
@@ -1300,6 +1447,14 @@ class KVCacheConfig:
     """Resolved retention policy for local prefix-cache checkpoints."""
     kv_cache_layout: str | None = None
     """The KV cache layout resolved by the engine core, adopted by all workers."""
+    hisparse_host_num_blocks: int | None = None
+    """Capacity of the dedicated HiSparse host-block manager, when enabled."""
+
+    hisparse_host_block_stride: int | None = None
+    """Physical bytes between consecutive HiSparse host blocks."""
+
+    hisparse_shared_host_pool: bool = False
+    """Whether local TP ranks share one physical HiSparse host pool."""
 
     @cached_property
     def transfer_group_ids(self) -> tuple[int, ...]:
@@ -1310,6 +1465,23 @@ def transfer_group_ids(self) -> tuple[int, ...]:
             if group.enable_kv_transfer
         )
 
+    @cached_property
+    def prefix_cacheable_group_ids(self) -> tuple[int, ...]:
+        """IDs of transferable groups eligible for hash-addressed stores."""
+        return tuple(
+            group_id
+            for group_id in self.transfer_group_ids
+            if self.kv_cache_groups[group_id].kv_cache_spec.prefix_cacheable
+        )
+
+    @cached_property
+    def prefix_cacheable_groups(self) -> tuple[KVCacheGroupSpec, ...]:
+        """Transferable groups eligible for hash-addressed stores."""
+        return tuple(
+            self.kv_cache_groups[group_id]
+            for group_id in self.prefix_cacheable_group_ids
+        )
+
     @cached_property
     def transfer_groups(self) -> tuple[KVCacheGroupSpec, ...]:
         """Cache groups that participate in external KV transfer."""
@@ -1337,6 +1509,21 @@ def select_transfer_block_ids(
             )
         return tuple(block_ids[group_id] for group_id in self.transfer_group_ids)
 
+    @property
+    def host_group_ids(self) -> tuple[int, ...]:
+        return tuple(
+            group_id
+            for group_id, group in enumerate(self.kv_cache_groups)
+            if group.host_resident
+        )
+
+    def num_blocks_of(self, tensor: KVCacheTensor) -> int:
+        """Number of blocks addressable by the pool backing ``tensor``."""
+        if not tensor.host_resident:
+            return self.num_blocks
+        assert self.hisparse_host_num_blocks is not None
+        return self.hisparse_host_num_blocks
+
     @property
     def has_mamba_layers(self) -> bool:
         return any(
@@ -1347,9 +1534,11 @@ def has_mamba_layers(self) -> bool:
 
     @property
     def has_mixed_precision_kv_cache(self) -> bool:
-        """Whether attention groups store their KV cache at more than one precision."""
+        """Whether device attention caches use more than one precision."""
         kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set()
         for group in self.kv_cache_groups:
+            if group.host_resident:
+                continue
             kv_cache_precisions.update(
                 (spec.dtype, spec.kv_quant_mode)
                 for spec in iter_layer_specs(group.kv_cache_spec)
diff --git a/vllm/v1/kv_cache_spec_registry.py b/vllm/v1/kv_cache_spec_registry.py
index 816a6862dae2..3a03fb376ca4 100644
--- a/vllm/v1/kv_cache_spec_registry.py
+++ b/vllm/v1/kv_cache_spec_registry.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Registry for KVCacheSpec types and their associated managers.
+"""Registry for KVCacheSpec types and their associated managers.
 
 This module provides a pluggable architecture for registering custom KVCacheSpec
 subclasses without modifying vLLM core code. Out-of-tree platforms can define
@@ -34,6 +33,7 @@ class KVCacheSpecMetadata:
 
 
 _REGISTRY_KVCACHESPEC_LIST: dict[type["KVCacheSpec"], KVCacheSpecMetadata] = {}
+_REGISTRY_ROLE_MANAGERS: dict[str, type["SingleTypeKVCacheManager"]] = {}
 
 
 class KVCacheSpecRegistry:
@@ -41,9 +41,7 @@ class KVCacheSpecRegistry:
 
     @classmethod
     def _ensure_registered(cls, vllm_config=None) -> None:
-        """
-        Run full KVCacheSpec registration if the registration is not done.
-        """
+        """Run full KVCacheSpec registration if the registration is not done."""
         if _REGISTRY_KVCACHESPEC_LIST:
             return
 
@@ -66,8 +64,7 @@ def register(
         manager_class: type["SingleTypeKVCacheManager"] | None = None,
         uniform_type_base_spec: type["KVCacheSpec"] | None = None,
     ) -> None:
-        """
-        Register a KVCacheSpec class with its manager and base spec.
+        """Register a KVCacheSpec class with its manager and base spec.
 
         Args:
             kvcache_spec_cls: The KVCacheSpec subclass to register
@@ -76,6 +73,7 @@ def register(
                 instead of being grouped to different kvcache group, `kvcache_spec_cls`
                 and `uniform_type_base_spec` will be trated as uniform type.
                 If None, defaults to kvcache_spec_cls itself (for built-in base specs).
+
         """
         assert manager_class is not None, "manager_class is required"
         if uniform_type_base_spec is None:
@@ -102,20 +100,34 @@ def register(
             uniform_type_base_spec=uniform_type_base_spec,
         )
 
+    @classmethod
+    def register_role_manager(
+        cls, role: str, manager_class: type["SingleTypeKVCacheManager"]
+    ) -> None:
+        """Select ``manager_class`` for every group with the given role."""
+        registered = _REGISTRY_ROLE_MANAGERS.get(role)
+        assert registered is None or registered is manager_class, (
+            f"Conflicting manager registration for group role {role}"
+        )
+        _REGISTRY_ROLE_MANAGERS[role] = manager_class
+
     @classmethod
     def get_manager_class(
-        cls, kvcache_spec: "KVCacheSpec"
+        cls, kvcache_spec: "KVCacheSpec", role: str | None = None
     ) -> type["SingleTypeKVCacheManager"] | None:
-        """
-        Get the single type kvcache manager class for a given kvcache spec instance.
+        """Get the single type kvcache manager class for a given kvcache spec instance.
 
         Args:
             kvcache_spec: A KVCacheSpec instance
+            role: The group's role; a role-registered manager takes precedence.
 
         Returns:
             The SingleTypeKVCacheManager class to use for this kvcache_spec
+
         """
         cls._ensure_registered()
+        if role is not None and (manager_class := _REGISTRY_ROLE_MANAGERS.get(role)):
+            return manager_class
         kvcache_spec_cls = type(kvcache_spec)
 
         # Walk up the MRO to find a registered base class
@@ -129,8 +141,7 @@ def get_manager_class(
     def get_uniform_type_base_spec(
         cls, kvcache_spec: "KVCacheSpec"
     ) -> type["KVCacheSpec"] | None:
-        """
-        Get the base kvcache spec class for grouping compatibility checks.
+        """Get the base kvcache spec class for grouping compatibility checks.
         KVCacheSpecs with uniform_type_base_spec will be trated as one group.
 
         Args:
@@ -138,6 +149,7 @@ def get_uniform_type_base_spec(
 
         Returns:
             The base KVCacheSpec class for checking uniform type kvcache specs
+
         """
         cls._ensure_registered()
         kvcache_spec_cls = type(kvcache_spec)
@@ -153,9 +165,7 @@ def get_uniform_type_base_spec(
     def check_kv_cache_spec_registry(
         cls, kv_cache_spec: dict[str, "KVCacheSpec"]
     ) -> None:
-        """
-        Check if the KVCacheSpecs of each layer are registered as expected.
-        """
+        """Check if the KVCacheSpecs of each layer are registered as expected."""
         cls._ensure_registered()
         for layer_name, spec in kv_cache_spec.items():
             # use raise instead of assert to make it effective in production environment
@@ -177,8 +187,7 @@ def register_kv_cache_spec(
     manager_class: type["SingleTypeKVCacheManager"] | None = None,
     uniform_type_base_spec: type["KVCacheSpec"] | None = None,
 ):
-    """
-    Decorator to register a custom KVCacheSpec class.
+    """Decorator to register a custom KVCacheSpec class.
 
     Args:
         manager_class: The SingleTypeKVCacheManager to use for this spec.
@@ -196,6 +205,7 @@ def register_kv_cache_spec(
         @dataclass(frozen=True, kw_only=True)
         class CustomFullAttentionSpec(FullAttentionSpec):
             pass
+
     """
 
     def decorator(kvcache_spec_cls: type["KVCacheSpec"]) -> type["KVCacheSpec"]:
diff --git a/vllm/v1/kv_hints/__init__.py b/vllm/v1/kv_hints/__init__.py
new file mode 100644
index 000000000000..d0807572c7a8
--- /dev/null
+++ b/vllm/v1/kv_hints/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.v1.kv_hints.protocol import KvHintAction, KvHintsEnvelope
+
+__all__ = ["KvHintAction", "KvHintsEnvelope"]
diff --git a/vllm/v1/kv_hints/protocol.py b/vllm/v1/kv_hints/protocol.py
new file mode 100644
index 000000000000..a3f93ab0bf23
--- /dev/null
+++ b/vllm/v1/kv_hints/protocol.py
@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import msgspec
+
+
+class KvHintAction(msgspec.Struct, frozen=True):  # type: ignore[call-arg]
+    """A versioned orchestrator-provided KV hint action."""
+
+    action_id: str
+    action_type: str
+    action_version: str
+    # `payload` must contain JSON-compatible values. Action-specific schemas are
+    # validated by the component implementing the action.
+    payload: dict[str, Any]
+
+
+class KvHintsEnvelope(msgspec.Struct, frozen=True):  # type: ignore[call-arg]
+    """Orchestrator-provided KV hints optionally attached to one inference request.
+    A KV hint envelope is a collection of KV hint action(s)."""
+
+    protocol_version: str
+    message_id: str
+    actions: list[KvHintAction]
diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py
index 93f8e4bda4e9..f8e92abed604 100644
--- a/vllm/v1/kv_offload/base.py
+++ b/vllm/v1/kv_offload/base.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Core abstractions for KV cache offloading in vLLM v1.
-"""
+"""Core abstractions for KV cache offloading in vLLM v1."""
 
 from abc import ABC, abstractmethod
 from collections.abc import Collection, Iterable, Sequence
@@ -18,6 +16,7 @@
         OffloadingConnectorStats,
     )
 
+from vllm.v1.kv_hints import KvHintsEnvelope
 from vllm.v1.kv_offload.config import OffloadingConfig
 
 # `OffloadKey` identifies an offloaded block. It combines a block hash with
@@ -91,11 +90,18 @@ def allows(self, medium: Medium | None, locality: Locality | None) -> bool:
 class ReqContext:
     req_id: str
     kv_transfer_params: dict[str, Any] | None = None
+    kv_hints: KvHintsEnvelope | None = None
     load_tier_filter: TierFilter = TierFilter.ALL
     # Per-request scratch space keyed by value type, so a tier can parse
-    # kv_transfer_params once (in on_new_request) and read the result back
-    # on later calls for the same request.
+    # kv_transfer_params and kv_hints once (in on_new_request) and read the
+    # result back on later calls for the same request.
     _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False)
+    # End-token position for each key in this request. The scheduler records
+    # these positions so managers can recover prefix order even when store
+    # calls arrive out of order (for example, SWA backfills).
+    _offload_key_positions: dict[OffloadKey, int] = field(
+        default_factory=dict, repr=False, init=False
+    )
 
     def set_state(self, val: Any) -> None:
         self._state[type(val)] = val
@@ -103,6 +109,12 @@ def set_state(self, val: Any) -> None:
     def get_state(self, cls: type[_T]) -> _T | None:
         return self._state.get(cls)
 
+    def set_offload_key_position(self, key: OffloadKey, end_token: int) -> None:
+        self._offload_key_positions[key] = end_token
+
+    def get_offload_key_position(self, key: OffloadKey) -> int | None:
+        return self._offload_key_positions.get(key)
+
 
 class LookupResult(Enum):
     """Result of OffloadingManager.lookup()."""
@@ -137,8 +149,7 @@ class ScheduleEndContext(NamedTuple):
 
 
 class LoadStoreSpec:
-    """
-    Metadata that encapsulates information allowing a worker
+    """Metadata that encapsulates information allowing a worker
     to load, and optionally also to store, blocks of KV data.
     """
 
@@ -223,8 +234,7 @@ class OffloadingKVEventsConfig:
 class OffloadingManager(ABC):
     @abstractmethod
     def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
-        """
-        Checks whether a single block is offloaded and ready to be read.
+        """Checks whether a single block is offloaded and ready to be read.
 
         Args:
             key: the key identifying the block to lookup.
@@ -234,6 +244,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
             HIT if the block is offloaded and ready, MISS if not found,
             HIT_PENDING if found but not yet readable, or RETRY if the
             lookup should be retried later.
+
         """
         pass
 
@@ -243,8 +254,7 @@ def prepare_load(
         keys: Collection[OffloadKey],
         req_context: ReqContext,
     ) -> LoadStoreSpec:
-        """
-        Prepare the given blocks to be read.
+        """Prepare the given blocks to be read.
         The given blocks will be protected from eviction until
         complete_load is called.
         It assumes all given blocks are offloaded.
@@ -256,27 +266,28 @@ def prepare_load(
         Returns:
             A LoadStoreSpec that can be used by a worker to locate and load
             the actual offloaded KV data.
+
         """
         pass
 
     def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
-        """
-        Mark the given blocks as recently used.
+        """Mark the given blocks as recently used.
         This could in practice mean moving them to the end of an LRU list.
 
         Args:
             keys: the keys identifying the blocks.
             req_context: per-request context (e.g. kv_transfer_params).
+
         """
         return
 
     def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
-        """
-        Marks previous blocks that were prepared to load as done loading.
+        """Marks previous blocks that were prepared to load as done loading.
 
         Args:
             keys: the keys identifying the blocks.
             req_context: per-request context (e.g. kv_transfer_params).
+
         """
         return
 
@@ -286,8 +297,7 @@ def prepare_store(
         keys: Collection[OffloadKey],
         req_context: ReqContext,
     ) -> PrepareStoreOutput | None:
-        """
-        Prepare the given blocks to be offloaded.
+        """Prepare the given blocks to be offloaded.
         The given blocks will be protected from eviction until
         complete_store is called.
 
@@ -300,6 +310,7 @@ def prepare_store(
             where to store them (LoadStoreSpec), and list of blocks that
             were evicted as a result.
             None is returned if the blocks cannot be stored.
+
         """
         pass
 
@@ -309,8 +320,7 @@ def complete_store(
         req_context: ReqContext,
         success: bool = True,
     ):
-        """
-        Marks blocks which were previously prepared to be stored, as stored.
+        """Marks blocks which were previously prepared to be stored, as stored.
         Following this call, the blocks become loadable.
         If success is False, blocks that were not marked as stored will be
         removed.
@@ -319,25 +329,25 @@ def complete_store(
             keys: the keys identifying the blocks.
             req_context: per-request context (e.g. kv_transfer_params).
             success: whether the blocks were stored successfully.
+
         """
         return
 
     @abstractmethod
     def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
-        """
-        Called when a new request is first seen by the scheduler.
+        """Called when a new request is first seen by the scheduler.
 
         Returns a RequestOffloadingContext indicating how this request's
         blocks should be offloaded.
 
         Args:
             req_context: per-request context.
+
         """
         pass
 
     def on_request_finished(self, req_context: ReqContext) -> None:
-        """
-        Called when a request has finished.
+        """Called when a request has finished.
 
         By the time this is called, the scheduler will issue no more
         submit-side calls for this request, such as prepare_store() and
@@ -352,12 +362,12 @@ def on_request_finished(self, req_context: ReqContext) -> None:
 
         Args:
             req_context: per-request context.
+
         """
         return
 
     def take_events(self) -> Iterable[OffloadingEvent]:
-        """
-        Take the offloading events from the manager.
+        """Take the offloading events from the manager.
 
         A tier manager emits only events for storage state it owns. A
         composing manager may aggregate child event streams, but should not
@@ -365,6 +375,7 @@ def take_events(self) -> Iterable[OffloadingEvent]:
 
         Yields:
             New OffloadingEvents collected since the last call.
+
         """
         return ()
 
@@ -398,8 +409,7 @@ def shutdown(self) -> None:
 
 
 class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
-    """
-    Spec for loading/storing KV blocks from given block numbers.
+    """Spec for loading/storing KV blocks from given block numbers.
 
     Subclass semantics differ: GPULoadStoreSpec.block_ids are GPU block
     indices; CPULoadStoreSpec.block_ids are CPU cache chunk indices.
@@ -413,8 +423,7 @@ def __repr__(self) -> str:
 
 
 class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
-    """
-    Spec for loading/storing a KV block to GPU memory.
+    """Spec for loading/storing a KV block to GPU memory.
 
     If there are multiple KV groups, the blocks are expected to be
     ordered by the group index.
@@ -448,8 +457,7 @@ def __init__(
 
 @dataclass
 class CanonicalKVCacheTensor:
-    """
-    A canonicalized KV cache tensor whose first dimension is num_blocks.
+    """A canonicalized KV cache tensor whose first dimension is num_blocks.
 
     With standardized layouts (RFC #42082) num_blocks is always the leading
     logical dimension.
@@ -505,8 +513,7 @@ def is_writer(self, block_id: int) -> bool:
 
 @dataclass
 class CanonicalKVCacheRef:
-    """
-    Per-layer (or group of layers) reference to a specific (by index)
+    """Per-layer (or group of layers) reference to a specific (by index)
     CanonicalKVCacheTensor and records the un-padded page size used by that layer.
     """
 
@@ -520,8 +527,7 @@ class CanonicalKVCacheRef:
 
 @dataclass
 class CanonicalKVCaches:
-    """
-    Canonicalized block-level representation of the KV caches.
+    """Canonicalized block-level representation of the KV caches.
 
     Composed of:
         - Unique list of KV cache data tensors,
@@ -608,8 +614,7 @@ def __init__(self, config: OffloadingConfig):
 
     @abstractmethod
     def get_manager(self) -> OffloadingManager:
-        """
-        Get an OffloadingManager that will be used
+        """Get an OffloadingManager that will be used
         by the scheduler-side offloading connector to track
         offloaded blocks and manage evictions.
         """
@@ -617,13 +622,13 @@ def get_manager(self) -> OffloadingManager:
 
     @abstractmethod
     def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:
-        """
-        Get an OffloadingWorker that handles async KV transfers for this spec.
+        """Get an OffloadingWorker that handles async KV transfers for this spec.
 
         Args:
             kv_caches: Canonicalized KV caches.
 
         Returns:
             An OffloadingWorker instance for this medium.
+
         """
         pass
diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py
index 637d3bb664dc..74b619dc0f8a 100644
--- a/vllm/v1/kv_offload/config.py
+++ b/vllm/v1/kv_offload/config.py
@@ -14,6 +14,8 @@ class OffloadingGroupConfig:
     tokens_per_block: int
     # Layer names belonging to this group.
     layer_names: tuple[str, ...]
+    # Original KVCacheConfig group index.
+    group_id: int
 
 
 @dataclass(frozen=True)
diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py
index 8db9e805ff1a..797abed5eaa4 100644
--- a/vllm/v1/kv_offload/cpu/gpu_worker.py
+++ b/vllm/v1/kv_offload/cpu/gpu_worker.py
@@ -11,6 +11,7 @@
 import torch
 
 from vllm import _custom_ops as ops
+from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON, triton
@@ -38,10 +39,13 @@
 def _select_swap_blocks_fn(
     layer_refs_per_group: list[list[CanonicalKVCacheRef]],
     gpu_to_cpu: bool,
+    host_memory_is_pinned: bool = True,
 ):
     """Resolve the swap_blocks function for a handler at init time."""
     # GPU->CPU is bandwidth-bound; the dedicated copy engine beats Triton.
-    if gpu_to_cpu:
+    # The Triton kernel dereferences CPU pointers on the GPU, which is only
+    # valid for pinned host memory.
+    if gpu_to_cpu or not host_memory_is_pinned:
         return ops.swap_blocks_batch
     # Fall back to the C++ DMA path on platforms where Triton isn't usable
     # (e.g. ROCm host mappings) or where GPU kernels cannot directly
@@ -80,8 +84,7 @@ def compute_sub_block_ptrs(
     tensor: torch.Tensor,
     skip_count: int = 0,
 ):
-    """
-    Compute byte pointers for sub-blocks of the given block IDs.
+    """Compute byte pointers for sub-blocks of the given block IDs.
 
     Each block in block_ids contains blocks_per_chunk sub-blocks.
     The pointer for sub-block j of block b is:
@@ -98,6 +101,7 @@ def compute_sub_block_ptrs(
         output: pre-allocated pointer array to write pointers into.
         tensor: the source or destination tensor.
         skip_count: sub-blocks to skip in the first block.
+
     """
     assert skip_count < blocks_per_chunk
 
@@ -190,8 +194,12 @@ def _canonical_block_sizes(
     return canonical_bytes_per_block
 
 
+# Bound registration size to avoid driver limits on large host allocations.
+MAX_HOST_REGISTER_CHUNK_BYTES = 64 * 1024**3
+
+
 def pin_mmap_region(region: SharedOffloadRegion) -> None:
-    """Register the entire mmap as CUDA pinned memory via cudaHostRegister."""
+    """Register row-aligned chunks, rolling back on failure."""
     if not current_platform.is_cuda_alike():
         logger.info(
             "Skipping mmap host registration on %s; cudaHostRegister is only "
@@ -201,23 +209,65 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None:
         return
 
     rank = region.rank
+    try:
+        cudart = CudaRTLibrary()
+    except (AssertionError, AttributeError, OSError):
+        logger.warning(
+            "Could not load the CUDA runtime for host registration on rank=%d; "
+            "the offload region stays pageable",
+            rank,
+            exc_info=True,
+        )
+        return
 
     base_ptr = region._base.data_ptr()
-    result = torch.cuda.cudart().cudaHostRegister(base_ptr, region.total_size_bytes, 0)
-    if result.value != 0:
+    total_size = region.total_size_bytes
+    # Chunks end on block-row boundaries, which are page aligned, so neither the
+    # driver's page rounding nor any single block transfer straddles two
+    # registrations.
+    rows_per_chunk = max(MAX_HOST_REGISTER_CHUNK_BYTES // region._row_stride, 1)
+    chunk_size = rows_per_chunk * region._row_stride
+
+    # Register, drain and roll back through the same runtime handle, so a
+    # failed chunk leaves neither a pending error nor a partly pinned region.
+    addresses: list[int] = []
+    for offset in range(0, total_size, chunk_size):
+        address = base_ptr + offset
+        size = min(chunk_size, total_size - offset)
+        result = cudart.cudaHostRegister(address, size)
+        if result == 0:
+            addresses.append(address)
+            continue
+        cudart.cudaGetLastError()
         logger.warning(
-            "cudaHostRegister failed for rank=%d (code=%d) — "
-            "transfers will still work but may be slower (unpinned DMA)",
+            "cudaHostRegister failed for rank=%d at %.2f of %.2f GB (code=%d); "
+            "the offload region stays pageable",
             rank,
+            offset / 1e9,
+            total_size / 1e9,
             result,
         )
-    else:
-        logger.debug(
-            "cudaHostRegister rank=%d %.2f GB",
-            rank,
-            region.total_size_bytes / 1e9,
-        )
-        region.is_pinned = True
+        for registered in reversed(addresses):
+            unregister_result = cudart.cudaHostUnregister(registered)
+            if unregister_result != 0:
+                cudart.cudaGetLastError()
+                logger.warning(
+                    "cudaHostUnregister failed for rank=%d at %#x (code=%d); "
+                    "that chunk stays registered until the process exits",
+                    rank,
+                    registered,
+                    unregister_result,
+                )
+        return
+
+    region.pinned_addresses.extend(addresses)
+    region.is_pinned = True
+    logger.debug(
+        "cudaHostRegister rank=%d %.2f GB in %d chunk(s)",
+        rank,
+        total_size / 1e9,
+        len(addresses),
+    )
 
 
 def _new_descriptor_buffers(
@@ -234,8 +284,7 @@ def _new_descriptor_buffers(
 
 
 class SingleDirectionOffloadingHandler:
-    """
-    Handles transfers for a single direction, either CPU->GPU or GPU->CPU.
+    """Handles transfers for a single direction, either CPU->GPU or GPU->CPU.
     Transfers are guaranteed to be executed in order of their submission.
     Each transfer uses a unique CUDA stream, and its stream will start
     executing only after the streams of previous transfers have finished.
@@ -249,9 +298,9 @@ def __init__(
         layer_refs_per_group: list[list[CanonicalKVCacheRef]],
         gpu_to_cpu: bool,
         canonical_layout: bool = False,
+        host_memory_is_pinned: bool = True,
     ):
-        """
-        Initialize a SingleDirectionOffloadingHandler.
+        """Initialize a SingleDirectionOffloadingHandler.
 
         Args:
             gpu_tensors: list of GPU KV cache tensors.
@@ -259,10 +308,14 @@ def __init__(
             cpu_tensors: list of CPU KV cache tensors.
                 Each of shape (num_cpu_chunks, cpu_page_size_bytes) with dtype int8.
                 Order should match gpu_tensors.
+            blocks_per_chunk: number of blocks transferred per chunk.
             layer_refs_per_group: list of CanonicalKVCacheRef per group.
             gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU.
             canonical_layout: if True, CPU pages use the canonical layout
                 described by the refs' mappings.
+            host_memory_is_pinned: whether the CPU tensors are pinned, so GPU
+                kernels may dereference them directly.
+
         """
         assert len(gpu_tensors) == len(cpu_tensors)
         assert len(gpu_tensors) > 0
@@ -299,7 +352,7 @@ def __init__(
         self.gpu_to_cpu: bool = gpu_to_cpu
         self.layer_refs_per_group = layer_refs_per_group
         self._swap_blocks_batch = _select_swap_blocks_fn(
-            layer_refs_per_group, gpu_to_cpu
+            layer_refs_per_group, gpu_to_cpu, host_memory_is_pinned
         )
 
         # GPU blocks may be smaller
@@ -646,7 +699,7 @@ def transfer_async(
             last_event = last_transfer.end_event
             # assure job will start only after the previous one completes
             stream.wait_event(last_event)
-        # CPU->GPU reads from host pinned memory, which is never written
+        # CPU->GPU reads from host memory, which is never written
         # by a concurrent GPU stream, so CU_MEMCPY_SRC_ACCESS_ORDER_ANY is
         # safe and lets the driver pipeline source reads. GPU->CPU reads
         # from the live GPU KV cache, which the compute stream keeps
@@ -764,6 +817,9 @@ def __init__(
         logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))
         if mmap_region is not None and pin_memory:
             pin_mmap_region(mmap_region)
+        host_memory_is_pinned = pin_memory and (
+            mmap_region is None or mmap_region.is_pinned
+        )
 
         canonical_bytes_per_block = (
             _canonical_block_sizes(kv_caches.group_data_refs, len(kv_caches.tensors))
@@ -813,6 +869,7 @@ def __init__(
             layer_refs_per_group=kv_caches.group_data_refs,
             gpu_to_cpu=True,
             canonical_layout=canonical_layout,
+            host_memory_is_pinned=host_memory_is_pinned,
         )
 
         self._load_handler = SingleDirectionOffloadingHandler(
@@ -822,6 +879,7 @@ def __init__(
             layer_refs_per_group=kv_caches.group_data_refs,
             gpu_to_cpu=False,
             canonical_layout=canonical_layout,
+            host_memory_is_pinned=host_memory_is_pinned,
         )
 
     def submit_store(
diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py
index 67b469b6a394..e42cb47f2530 100644
--- a/vllm/v1/kv_offload/cpu/manager.py
+++ b/vllm/v1/kv_offload/cpu/manager.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections import OrderedDict
 from collections.abc import Collection, Iterable
+from dataclasses import dataclass, field
 
 from typing_extensions import override
 
@@ -18,6 +19,7 @@
     PrepareStoreOutput,
     ReqContext,
     RequestOffloadingContext,
+    get_offload_group_idx,
 )
 from vllm.v1.kv_offload.cpu.common import (
     CPULoadStoreSpec,
@@ -27,9 +29,22 @@
 from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
 
 
+@dataclass(slots=True)
+class _RequestCacheAccess:
+    """Cache keys observed by one request, grouped in prefix order."""
+
+    owner: object
+    cache_generation: int
+    key_groups: dict[int, list[OffloadKey]] = field(default_factory=dict)
+    seen_keys: set[OffloadKey] = field(default_factory=set)
+    inserted_keys: set[OffloadKey] = field(default_factory=set)
+    reused_keys: set[OffloadKey] = field(default_factory=set)
+    store_miss_keys: set[OffloadKey] = field(default_factory=set)
+    finished: bool = False
+
+
 class CPUOffloadingManager(OffloadingManager):
-    """
-    An OffloadingManager with a pluggable CachePolicy, resolved by name via
+    """An OffloadingManager with a pluggable CachePolicy, resolved by name via
     CachePolicyFactory (built in: "lru", "arc"; external policies can either
     register their own or be loaded out-of-tree via cache_policy_module_path).
 
@@ -66,6 +81,7 @@ def __init__(
         self.max_tracker_size: int = max_tracker_size
         self.stores_skipped_in_current_batch: int = 0
         self.allocation_sizes_in_current_batch: list[int] = []
+        self._cache_generation = 0
 
         # Number of chunk references. It is ordered so can evict the LRU entry in O(1).
         self.counts: OrderedDict[OffloadKey, int] | None = (
@@ -124,10 +140,48 @@ def _record_accesses(self, keys: Collection[OffloadKey]) -> None:
                 num_unprotected -= 1
             self.counts[key] = 1
 
+    def _get_request_cache_access(self, req_context: ReqContext) -> _RequestCacheAccess:
+        state = req_context.get_state(_RequestCacheAccess)
+        if (
+            state is None
+            or state.owner is not self
+            or state.cache_generation != self._cache_generation
+        ):
+            state = _RequestCacheAccess(
+                owner=self, cache_generation=self._cache_generation
+            )
+            req_context.set_state(state)
+        return state
+
+    def _record_request_cache_access(
+        self,
+        keys: Iterable[OffloadKey],
+        req_context: ReqContext,
+        inserted_keys: Iterable[OffloadKey] = (),
+        reused_keys: Iterable[OffloadKey] = (),
+    ) -> None:
+        state = self._get_request_cache_access(req_context)
+        for key in keys:
+            if key in state.seen_keys:
+                continue
+            assert not state.finished, (
+                "New cache keys observed after request finalization"
+            )
+            group_idx = get_offload_group_idx(key)
+            state.key_groups.setdefault(group_idx, []).append(key)
+            state.seen_keys.add(key)
+        state.inserted_keys.update(inserted_keys)
+        # Re-reading a chunk inserted by this request is an internal transfer
+        # (for example, a tiering cascade), not a second cache access.
+        state.reused_keys.update(
+            key for key in reused_keys if key not in state.inserted_keys
+        )
+
     # --- OffloadingManager interface ---
 
     @override
     def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
+        self._get_request_cache_access(req_context)
         return RequestOffloadingContext()
 
     @override
@@ -144,6 +198,15 @@ def prepare_load(
         self,
         keys: Collection[OffloadKey],
         req_context: ReqContext,
+    ) -> LoadStoreSpec:
+        return self._prepare_load(keys, req_context, record_access=True)
+
+    def _prepare_load(
+        self,
+        keys: Collection[OffloadKey],
+        req_context: ReqContext,
+        *,
+        record_access: bool,
     ) -> LoadStoreSpec:
         chunks = []
         for key in keys:
@@ -156,6 +219,8 @@ def prepare_load(
                 assert self._num_evictable_cache_chunks >= 0
             chunk.ref_cnt += 1
             chunks.append(chunk)
+        if record_access:
+            self._record_request_cache_access(keys, req_context, reused_keys=keys)
         return self._get_load_store_spec(keys, chunks)
 
     @override
@@ -181,13 +246,52 @@ def prepare_store(
         keys: Collection[OffloadKey],
         req_context: ReqContext,
     ) -> PrepareStoreOutput | None:
+        keys = list(keys)
         if self.counts is not None:
-            num_keys = len(keys)
             self._record_accesses(keys)
-            keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold]
-            self.stores_skipped_in_current_batch += num_keys - len(keys)
-        # filter out chunks that are already stored
-        keys_to_store = [k for k in keys if self._policy.get(k) is None]
+        # Partition the original offer before admission filtering. Ready
+        # resident chunks are request reuses even if their tracker entry was
+        # aged out; the threshold applies only to new store candidates.
+        keys_to_store: list[OffloadKey] = []
+        ready_existing_keys: list[OffloadKey] = []
+        for key in keys:
+            chunk = self._policy.get(key)
+            if chunk is None:
+                keys_to_store.append(key)
+            else:
+                if chunk.is_ready:
+                    ready_existing_keys.append(key)
+
+        if self.counts is not None:
+            num_store_candidates = len(keys_to_store)
+            keys_to_store = [
+                key
+                for key in keys_to_store
+                if self.counts.get(key, 0) >= self.store_threshold
+            ]
+            self.stores_skipped_in_current_batch += num_store_candidates - len(
+                keys_to_store
+            )
+
+        state = self._get_request_cache_access(req_context)
+        new_store_misses = [
+            key for key in keys_to_store if key not in state.store_miss_keys
+        ]
+        if new_store_misses:
+            # ARC learns from B1/B2 before insert() removes the ghost entry.
+            # Deduplication makes this one policy observation per request.
+            self._policy.on_store_miss(new_store_misses, req_context)
+            state.store_miss_keys.update(new_store_misses)
+
+        recorded_keys = set(keys_to_store)
+        recorded_keys.update(ready_existing_keys)
+        # Record once before any allocation-related early return. Store
+        # candidates are classified as insertions only after allocation succeeds.
+        self._record_request_cache_access(
+            (key for key in keys if key in recorded_keys),
+            req_context,
+            reused_keys=ready_existing_keys,
+        )
 
         if not keys_to_store:
             return PrepareStoreOutput(
@@ -239,6 +343,7 @@ def prepare_store(
         for key, chunk in zip(keys_to_store, chunks):
             self._policy.insert(key, chunk)
         self._num_write_pending_chunks += len(keys_to_store)
+        state.inserted_keys.update(keys_to_store)
 
         # build store specs for allocated chunks
         store_spec = self._get_load_store_spec(keys_to_store, chunks)
@@ -284,6 +389,35 @@ def complete_store(
                 )
             )
 
+    @override
+    def on_request_finished(self, req_context: ReqContext) -> None:
+        state = req_context.get_state(_RequestCacheAccess)
+        if (
+            state is None
+            or state.owner is not self
+            or state.cache_generation != self._cache_generation
+            or state.finished
+        ):
+            return
+        state.finished = True
+        key_groups = []
+        for group_idx in sorted(state.key_groups):
+            keys = state.key_groups[group_idx]
+            positions = {
+                key: position
+                for key in keys
+                if (position := req_context.get_offload_key_position(key)) is not None
+            }
+            if len(positions) == len(keys):
+                keys = sorted(keys, key=positions.__getitem__)
+            key_groups.append(tuple(keys))
+        self._policy.on_request_finished(
+            tuple(key_groups),
+            state.inserted_keys - state.reused_keys,
+            state.reused_keys,
+            req_context,
+        )
+
     @override
     def reset_cache(self) -> None:
         # Clear ALL chunks unconditionally. The scheduler's _stale_job_threshold
@@ -292,6 +426,7 @@ def reset_cache(self) -> None:
         # flushes in-flight load job IDs to the workers before any new stores
         # can begin, preventing a cross-direction data race on reused offload chunk IDs.
         self._policy.clear()
+        self._cache_generation += 1
         self._num_evictable_cache_chunks = 0
         self._num_write_pending_chunks = 0
 
diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py
index e413bd603eb4..c3942e90d850 100644
--- a/vllm/v1/kv_offload/cpu/policies/arc.py
+++ b/vllm/v1/kv_offload/cpu/policies/arc.py
@@ -1,17 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections import OrderedDict
-from collections.abc import Iterable, Iterator
+from collections.abc import Iterable, Iterator, Sequence
 
 from typing_extensions import override
 
 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
-from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
+from vllm.v1.kv_offload.cpu.policies.base import (
+    CachePolicy,
+    ChunkStatus,
+    order_request_keys,
+)
 
 
 class ARCCachePolicy(CachePolicy):
-    """
-    ARC (Adaptive Replacement Cache) cache policy.
+    """ARC (Adaptive Replacement Cache) cache policy.
 
     Data Structures:
         T1: Recent cache containing chunks accessed once.
@@ -24,12 +27,11 @@ class ARCCachePolicy(CachePolicy):
            Searches T1 and T2 for chunk hashes and counts consecutive hits
            until a miss or non-ready chunk is encountered.
 
-        2. Cache touch (touch) - Adaptive Learning:
-           For each key (in reverse order):
-           - If in T1: Move to T2 (promotion from recent to frequent).
-           - If in T2: Move to MRU position (end of queue).
-           - If in B1 ghost list: Increase target_t1_size.
-           - If in B2 ghost list: Decrease target_t1_size.
+        2. Request access - Adaptive Learning:
+           - Ready chunks reused by a request move from T1 to T2 once, or
+             move to the MRU end of T2.
+           - B1/B2 misses adjust target_t1_size once per request before
+             insertion removes their ghost entries.
 
         3. Chunk eviction (evict) - Adaptive Replacement:
            Determines eviction source based on adaptive target:
@@ -39,7 +41,7 @@ class ARCCachePolicy(CachePolicy):
 
         4. Chunk insertion (insert):
            New chunks are always inserted into T1 and removed from B1/B2 if
-           present. Chunks may later be promoted to T2 during touch operations.
+           present. A later request reuse may promote them to T2.
 
     Adaptive Behavior:
         The algorithm self-tunes the recency vs. frequency trade-off:
@@ -71,6 +73,26 @@ def remove(self, key: OffloadKey) -> None:
         if self.t1.pop(key, None) is None:
             self.t2.pop(key, None)
 
+    def _adapt_to_ghost_hit(self, key: OffloadKey) -> bool:
+        if key in self.b1:
+            delta = max(1, len(self.b2) / len(self.b1))
+            self.target_t1_size = min(self.target_t1_size + delta, self.cache_capacity)
+            self.b1.move_to_end(key)
+            return True
+        if key in self.b2:
+            delta = max(1, len(self.b1) / len(self.b2))
+            self.target_t1_size = max(self.target_t1_size - delta, 0)
+            self.b2.move_to_end(key)
+            return True
+        return False
+
+    @override
+    def on_store_miss(
+        self, keys: Iterable[OffloadKey], req_context: ReqContext
+    ) -> None:
+        for key in reversed(list(keys)):
+            self._adapt_to_ghost_hit(key)
+
     @override
     def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
         for key in reversed(list(keys)):
@@ -86,19 +108,33 @@ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
             elif key in self.t2:
                 self.t2.move_to_end(key)
 
-            elif key in self.b1:
-                delta = max(1, len(self.b2) / len(self.b1))
-                self.target_t1_size = min(
-                    self.target_t1_size + delta, self.cache_capacity
-                )
-                # move to MRU position (end) to keep it fresh in the ghost list
-                self.b1.move_to_end(key)
-
-            elif key in self.b2:
-                delta = max(1, len(self.b1) / len(self.b2))
-                self.target_t1_size = max(self.target_t1_size - delta, 0)
-                # move to MRU position (end) to keep it fresh in the ghost list
-                self.b2.move_to_end(key)
+            else:
+                self._adapt_to_ghost_hit(key)
+
+    @override
+    def on_request_finished(
+        self,
+        key_groups: Sequence[Sequence[OffloadKey]],
+        insertion_only_keys: set[OffloadKey],
+        reused_keys: set[OffloadKey],
+        req_context: ReqContext,
+    ) -> None:
+        for key in reversed(order_request_keys(key_groups, req_context)):
+            if key in insertion_only_keys:
+                # A store is not a frequency hit. Preserve T1 membership
+                # while still restoring tail-to-head recency.
+                if key in self.t1:
+                    self.t1.move_to_end(key)
+                elif key in self.t2:
+                    self.t2.move_to_end(key)
+                continue
+
+            # Ready chunks reused by this request count as one access.
+            if key in reused_keys and key in self.t1:
+                chunk = self.t1.pop(key)
+                self.t2[key] = chunk
+            elif key in reused_keys and key in self.t2:
+                self.t2.move_to_end(key)
 
     @override
     def clear(self) -> None:
diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py
index 4059dd742f7c..d3ef8b859a9b 100644
--- a/vllm/v1/kv_offload/cpu/policies/base.py
+++ b/vllm/v1/kv_offload/cpu/policies/base.py
@@ -2,14 +2,28 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import ctypes
 from abc import ABC, abstractmethod
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 
 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
 
 
+def order_request_keys(
+    key_groups: Sequence[Sequence[OffloadKey]], req_context: ReqContext
+) -> list[OffloadKey]:
+    """Return one head-to-tail order across all KV cache groups."""
+    keys = [key for group in key_groups for key in group]
+    positions = {
+        key: position
+        for key in keys
+        if (position := req_context.get_offload_key_position(key)) is not None
+    }
+    if len(positions) == len(keys):
+        keys.sort(key=positions.__getitem__)
+    return keys
+
+
 class ChunkStatus(ctypes.Structure):
-    """
-    Offloading status for a single chunk of KV data.
+    """Offloading status for a single chunk of KV data.
     Holds the following information:
 
     ref_cnt - the current number of transfers using this chunk as a source.
@@ -27,15 +41,12 @@ def __init__(self, chunk_id: int):
 
     @property
     def is_ready(self) -> bool:
-        """
-        Returns whether the chunk is ready to be read.
-        """
+        """Returns whether the chunk is ready to be read."""
         return self.ref_cnt >= 0
 
 
 class CachePolicy(ABC):
-    """
-    Encapsulates both chunk organization (data structures) and replacement
+    """Encapsulates both chunk organization (data structures) and replacement
     decisions (which chunk to evict). LRU and ARC differ in both dimensions —
     ARC's ghost lists and target_t1_size live at the intersection of storage
     and eviction, so they cannot be separated cleanly.
@@ -58,20 +69,54 @@ def remove(self, key: OffloadKey) -> None:
 
     @abstractmethod
     def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
-        """
-        Mark chunks as recently used.
+        """Mark chunks as recently used.
 
         Args:
             keys: Chunks to mark as recently used.
             req_context: Per-request context for the request touching these chunks.
+
         """
 
+    def on_store_miss(
+        self, keys: Iterable[OffloadKey], req_context: ReqContext
+    ) -> None:
+        """Observe store misses before their cache entries are inserted.
+
+        The default delegates to ``touch`` for compatibility with external
+        policies. Policies may override this hook when a miss has distinct
+        semantics, such as ARC adapting to a ghost-list hit.
+        """
+        self.touch(keys, req_context)
+
+    def on_request_finished(
+        self,
+        key_groups: Sequence[Sequence[OffloadKey]],
+        insertion_only_keys: set[OffloadKey],
+        reused_keys: set[OffloadKey],
+        req_context: ReqContext,
+    ) -> None:
+        """Apply one request-scoped cache access in prefix order.
+
+        ``key_groups`` contains keys observed for each KV cache group in
+        head-to-tail order. ``insertion_only_keys`` and ``reused_keys``
+        distinguish chunks only created by this request from ready chunks it
+        actually reused, which matters for policies such as ARC where reuse
+        changes frequency but insertion and pending observations do not.
+
+        The default forwards one touch per group, preserving compatibility
+        for experimental out-of-tree policies while moving those touches to
+        request finalization. Policies that distinguish insertion from reuse
+        can override this hook and inspect the access classifications.
+        """
+        del insertion_only_keys, reused_keys
+        for keys in key_groups:
+            self.touch(keys, req_context)
+
     @abstractmethod
     def evict(
         self, n: int, protected: set[OffloadKey]
     ) -> list[tuple[OffloadKey, ChunkStatus]] | None:
-        """
-        Evict exactly n chunks, skipping any in protected.
+        """Evict exactly n chunks, skipping any in protected.
 
         Returns a list of (key, chunk) for the evicted chunks,
         or None if n evictions cannot be satisfied. The operation is atomic:
@@ -83,8 +128,7 @@ def evict(
 
     @abstractmethod
     def clear(self) -> None:
-        """
-        Remove ALL chunks regardless of ref_cnt.
+        """Remove ALL chunks regardless of ref_cnt.
 
         Ghost lists and adaptive state are also reset.
         """
diff --git a/vllm/v1/kv_offload/cpu/policies/factory.py b/vllm/v1/kv_offload/cpu/policies/factory.py
index eedcc6807697..e4f811d49f0e 100644
--- a/vllm/v1/kv_offload/cpu/policies/factory.py
+++ b/vllm/v1/kv_offload/cpu/policies/factory.py
@@ -57,6 +57,7 @@ def get_cache_policy_cls(
 
         Raises ValueError if the cache policy is neither registered nor
         resolvable via `module_path`.
+
         """
         if name in cls._registry:
             return cls._registry[name]()
diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py
index c1dc631d4fe9..6ca9cb32f42a 100644
--- a/vllm/v1/kv_offload/cpu/policies/lru.py
+++ b/vllm/v1/kv_offload/cpu/policies/lru.py
@@ -1,28 +1,73 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections import OrderedDict
-from collections.abc import Iterable
+import heapq
+from collections.abc import Iterable, Sequence
 
 from typing_extensions import override
 
 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
-from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
+from vllm.v1.kv_offload.cpu.policies.base import (
+    CachePolicy,
+    ChunkStatus,
+    order_request_keys,
+)
 
 
 class LRUCachePolicy(CachePolicy):
-    """
-    LRU Caching policy that keeps a dedicated evictable list for fast eviction.
+    """LRU caching policy with logical recency independent of transfer pinning.
+
+    Evictable chunks live in a lazy-invalidating min-heap. A chunk's recency
+    can therefore be updated while it is pinned; when it later becomes
+    evictable, it enters the heap with the order assigned by the request
+    rather than its transfer-completion time.
+
     A use is indicated by,
      - First time the key is added (store).
-     - Load job completion
-     - touch
+     - A request-scoped access.
     """
 
     def __init__(self, cache_capacity: int):
         super().__init__(cache_capacity)
-        # Chunks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU
-        self.evictable_chunks: OrderedDict[OffloadKey, None] = OrderedDict()
         self.chunks: dict[OffloadKey, ChunkStatus] = {}
+        self._ranks: dict[OffloadKey, int] = {}
+        self._evictable: set[OffloadKey] = set()
+        self._heap: list[tuple[int, OffloadKey]] = []
+        self._next_rank = 0
+
+    def _assign_new_rank(self, key: OffloadKey) -> bool:
+        if key not in self.chunks:
+            return False
+        self._next_rank += 1
+        self._ranks[key] = self._next_rank
+        return key in self._evictable
+
+    def _push_evictable(self, key: OffloadKey) -> None:
+        heapq.heappush(
+            self._heap,
+            (self._ranks[key], key),
+        )
+
+    def _is_current(self, entry: tuple[int, OffloadKey]) -> bool:
+        rank, key = entry
+        return key in self._evictable and self._ranks.get(key) == rank
+
+    def _maybe_compact_heap(self) -> None:
+        if len(self._heap) < max(64, 2 * len(self._evictable)):
+            return
+        self._heap = [(self._ranks[key], key) for key in self._evictable]
+        heapq.heapify(self._heap)
+
+    def _update_recency(self, keys: Iterable[OffloadKey]) -> None:
+        updated_evictable = [key for key in keys if self._assign_new_rank(key)]
+        # Rebuilding is linear and substantially cheaper than k heap pushes
+        # for a long prefix. Small updates retain the incremental path.
+        if len(updated_evictable) >= max(64, len(self._evictable) // 4):
+            self._heap = [(self._ranks[key], key) for key in self._evictable]
+            heapq.heapify(self._heap)
+        else:
+            for key in updated_evictable:
+                self._push_evictable(key)
+            self._maybe_compact_heap()
 
     @override
     def get(self, key: OffloadKey) -> ChunkStatus | None:
@@ -31,26 +76,40 @@ def get(self, key: OffloadKey) -> ChunkStatus | None:
     @override
     def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None:
         self.chunks[key] = chunk
+        self._next_rank += 1
+        self._ranks[key] = self._next_rank
         if chunk.ref_cnt == 0:
-            self.evictable_chunks[key] = None
+            self._evictable.add(key)
+            self._push_evictable(key)
 
     @override
     def remove(self, key: OffloadKey) -> None:
         del self.chunks[key]
-        self.evictable_chunks.pop(key, None)
+        self._ranks.pop(key, None)
+        self._evictable.discard(key)
 
     @override
     def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
-        for key in reversed(list(keys)):
-            if key in self.evictable_chunks:
-                self.evictable_chunks.move_to_end(key)
-            # active chunks are untouched as they are non-evictable now. They
-            # will eventually reach the end of evictable_chunks when they finish.
+        self._update_recency(reversed(list(keys)))
+
+    @override
+    def on_request_finished(
+        self,
+        key_groups: Sequence[Sequence[OffloadKey]],
+        insertion_only_keys: set[OffloadKey],
+        reused_keys: set[OffloadKey],
+        req_context: ReqContext,
+    ) -> None:
+        del insertion_only_keys, reused_keys
+        self._update_recency(reversed(order_request_keys(key_groups, req_context)))
 
     @override
     def clear(self) -> None:
-        self.evictable_chunks.clear()
         self.chunks.clear()
+        self._ranks.clear()
+        self._evictable.clear()
+        self._heap.clear()
+        self._next_rank = 0
 
     @override
     def evict(
@@ -59,22 +118,44 @@ def evict(
         if n == 0:
             return []
 
-        candidates: list[tuple[OffloadKey, ChunkStatus]] = []
-        for key, _ in self.evictable_chunks.items():
+        selected: list[tuple[tuple[int, OffloadKey], ChunkStatus]] = []
+        selected_keys: set[OffloadKey] = set()
+        deferred: list[tuple[int, OffloadKey]] = []
+        while self._heap and len(selected) < n:
+            entry = heapq.heappop(self._heap)
+            if not self._is_current(entry):
+                continue
+            key = entry[1]
+            # Re-pinning without a recency change can leave an equivalent
+            # lazy entry behind. Select each cache key at most once.
+            if key in selected_keys:
+                continue
             if key in protected:
+                deferred.append(entry)
                 continue
-
             chunk = self.chunks[key]
             assert chunk.ref_cnt == 0
-            candidates.append((key, chunk))
-            if len(candidates) == n:
-                break
-
-        if len(candidates) < n:
+            selected.append((entry, chunk))
+            selected_keys.add(key)
+
+        if len(selected) < n:
+            for entry, _ in selected:
+                heapq.heappush(self._heap, entry)
+            for entry in deferred:
+                heapq.heappush(self._heap, entry)
             return None
-        for key, _ in candidates:
-            del self.evictable_chunks[key]
+
+        for entry in deferred:
+            heapq.heappush(self._heap, entry)
+
+        candidates: list[tuple[OffloadKey, ChunkStatus]] = []
+        for entry, chunk in selected:
+            key = entry[1]
+            self._evictable.remove(key)
             del self.chunks[key]
+            del self._ranks[key]
+            candidates.append((key, chunk))
+        self._maybe_compact_heap()
         return candidates
 
     @override
@@ -82,9 +163,11 @@ def mark_evictable(self, key: OffloadKey) -> None:
         # chunks can become evictable when,
         # store completes - i.e. ref_cnt -1 -> 0 # not in evictable list
         # all loads complete - i.e ref_cnt 1 -> 0  # not in evictable list
-        self.evictable_chunks[key] = None
+        self._evictable.add(key)
+        self._push_evictable(key)
+        self._maybe_compact_heap()
 
     @override
     def mark_non_evictable(self, key: OffloadKey) -> None:
-        # key must have been in the evictable list.
-        del self.evictable_chunks[key]
+        # key must have been evictable before it was pinned.
+        self._evictable.remove(key)
diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py
index 1a5dde21e62d..af183fbc4f06 100644
--- a/vllm/v1/kv_offload/cpu/shared_offload_region.py
+++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
 import errno
 import mmap
 import os
@@ -27,6 +28,10 @@ def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> N
     while True:
         if os.fstat(fd).st_size >= expected_size:
             return
+        if os.fstat(fd).st_nlink == 0:
+            raise RuntimeError(
+                "Shared offload region creator failed during initialization."
+            )
         if time.monotonic() > deadline:
             raise TimeoutError(
                 f"Timed out waiting for mmap file to reach {expected_size} bytes"
@@ -64,8 +69,7 @@ def _get_populate_write_fn(
 
 
 class SharedOffloadRegion:
-    """
-    Single mmap-backed memory region shared across all workers for a
+    """Single mmap-backed memory region shared across all workers for a
     vLLM instance.  Workers coordinate via the filesystem: the first worker
     to open the file with O_EXCL becomes the creator and calls ftruncate;
     the rest open the existing file and wait until it reaches the expected
@@ -75,6 +79,9 @@ class SharedOffloadRegion:
     given, the path is unlinked once every worker has mapped the file, so
     the kernel reclaims the memory when the last worker exits, no matter
     how it exits; mappings taken before the unlink stay valid.
+
+    Creator-only population pre-faults the entire region before the barrier
+    and requires that barrier to keep joiners from using unpopulated pages.
     """
 
     BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE
@@ -87,7 +94,12 @@ def __init__(
         kv_bytes_per_chunk: int,
         cpu_page_size: int,
         barrier: Callable[[], None] | None = None,
+        *,
+        creator_memory_check: Callable[[int], None] | None = None,
+        populate_only_on_creator: bool = False,
     ) -> None:
+        if populate_only_on_creator and barrier is None:
+            raise ValueError("Creator-only population requires a barrier.")
         self.page_size = mmap.PAGESIZE
         assert kv_bytes_per_chunk % self.page_size == 0
 
@@ -112,25 +124,21 @@ def __init__(
                 # Joiner path — another worker won O_EXCL. Reopen and wait
                 # for the file to reach expected size.
                 self.fd = os.open(self.mmap_path, os.O_RDWR)
-                try:
-                    _wait_for_file_size(self.fd, self.total_size_bytes)
-                except (TimeoutError, OSError):
-                    os.close(self.fd)
-                    raise
+                _wait_for_file_size(self.fd, self.total_size_bytes)
                 logger.info("Opened existing mmap file %s", self.mmap_path)
             else:
                 # Creator path. We won O_EXCL, so we own the file: any
                 # failure here must clean up so concurrent joiners don't
                 # land on a 0-byte stub and spin in _wait_for_file_size
                 # for the full 30 s timeout.
-                try:
-                    check_shm_free_space(self.total_size_bytes)
-                    os.ftruncate(self.fd, self.total_size_bytes)
-                except (RuntimeError, OSError):
-                    os.unlink(self.mmap_path)
-                    os.close(self.fd)
-                    raise
                 self._creator = True
+                if creator_memory_check is not None:
+                    creator_memory_check(self.total_size_bytes)
+                check_shm_free_space(
+                    self.total_size_bytes,
+                    allocation_name="CPU KV offload shared region in /dev/shm",
+                )
+                os.ftruncate(self.fd, self.total_size_bytes)
                 logger.info(
                     "Created mmap file %s (%.2f GB)",
                     self.mmap_path,
@@ -143,10 +151,21 @@ def __init__(
                 flags=mmap.MAP_SHARED,
                 prot=mmap.PROT_READ | mmap.PROT_WRITE,
             )
+
+            if populate_only_on_creator and self._creator:
+                populate_write_fn = _get_populate_write_fn(self.mmap_obj)
+                populate_write_fn(self.mmap_obj, 0, self.total_size_bytes)
         except Exception:
             if self._creator:
-                os.unlink(self.mmap_path)
+                with contextlib.suppress(FileNotFoundError):
+                    os.unlink(self.mmap_path)
                 self._creator = False
+            if hasattr(self, "mmap_obj") and self.mmap_obj is not None:
+                self.mmap_obj.close()
+                self.mmap_obj = None
+            if hasattr(self, "fd") and self.fd is not None:
+                os.close(self.fd)
+                self.fd = None
             # Peers block inside the barrier until the collective times out if
             # we die before reaching it.  Arrive anyway so every worker calls
             # barrier() exactly once and they fail on their own errors instead
@@ -169,16 +188,28 @@ def __init__(
                 barrier()
             except Exception:
                 if self._creator:
-                    os.unlink(self.mmap_path)
+                    with contextlib.suppress(FileNotFoundError):
+                        os.unlink(self.mmap_path)
                     self._creator = False
                 self.mmap_obj.close()
                 os.close(self.fd)
+                self.mmap_obj = None
+                self.fd = None
                 raise
             if self._creator:
                 os.unlink(self.mmap_path)
                 self._creator = False
                 logger.info("Unlinked mmap file %s", self.mmap_path)
 
+        self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8)
+        self._views: list[torch.Tensor] = []
+        self._canonical_offset = 0
+        self.is_pinned: bool = False
+        self.pinned_addresses: list[int] = []
+
+        if populate_only_on_creator:
+            return
+
         populate_write_fn = _get_populate_write_fn(self.mmap_obj)
 
         if rank is not None:
@@ -205,10 +236,11 @@ def __init__(
                 "MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0
             )
 
-        self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8)
-        self._views: list[torch.Tensor] = []
-        self._canonical_offset = 0
-        self.is_pinned: bool = False
+    @property
+    def base_tensor(self) -> torch.Tensor:
+        if self._base is None:
+            raise RuntimeError("Shared offload region has been released.")
+        return self._base
 
     def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor:
         """Allocate a strided int8 view for this worker, one canonical tensor.
@@ -231,6 +263,7 @@ def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor:
 
         Args:
             tensor_page_size: Bytes per chunk for this tensor.
+
         """
         assert self.rank is not None
         new_offset = self._worker_offset + tensor_page_size
@@ -280,6 +313,7 @@ def create_next_canonical_view(self, tensor_page_size: int) -> torch.Tensor:
 
         Args:
             tensor_page_size: Canonical bytes per chunk for this tensor.
+
         """
         new_offset = self._canonical_offset + tensor_page_size
         assert new_offset <= self._row_stride
@@ -311,13 +345,18 @@ def cleanup(self) -> None:
         if self.is_pinned and self._base is not None:
             if current_platform.is_cuda_alike():
                 base_ptr = self._base.data_ptr()
-                result = torch.cuda.cudart().cudaHostUnregister(base_ptr)
-                if result.value != 0:
-                    logger.warning(
-                        "cudaHostUnregister failed for rank=%d (code=%d)",
-                        self.rank,
-                        result,
-                    )
+                addresses = self.pinned_addresses or [base_ptr]
+                for address in reversed(addresses):
+                    result = torch.cuda.cudart().cudaHostUnregister(address)
+                    if result.value != 0:
+                        logger.warning(
+                            "cudaHostUnregister failed for rank=%d, "
+                            "address=%#x (code=%d)",
+                            self.rank,
+                            address,
+                            result.value,
+                        )
+                self.pinned_addresses.clear()
             self.is_pinned = False
         # Release views before _base: each view holds a _base reference and a
         # direct StorageImpl reference.  Freeing views first lets both refcounts
diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py
index 00b36e863a07..fddf56608590 100644
--- a/vllm/v1/kv_offload/cpu/spec.py
+++ b/vllm/v1/kv_offload/cpu/spec.py
@@ -162,7 +162,7 @@ def get_manager(self) -> OffloadingManager:
     def _uses_shared_region(self) -> bool:
         """Whether the worker CPU buffer is the shared mmap region (vs a private
         per-rank tensor); replicated-layout dedup is gated on this being True."""
-        return current_platform.is_cuda_alike()
+        return current_platform.is_cuda_alike() and not current_platform.is_rocm()
 
     def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
         mmap_region: SharedOffloadRegion | None = None
diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py
index 82ebe31e2d33..1ae0d2d1fb07 100644
--- a/vllm/v1/kv_offload/file_mapper.py
+++ b/vllm/v1/kv_offload/file_mapper.py
@@ -19,9 +19,7 @@
 
 
 class FileMapper:
-    """
-    FileMapper maps KV blocks (given by their hash) to file names.
-    """
+    """FileMapper maps KV blocks (given by their hash) to file names."""
 
     def __init__(
         self,
@@ -41,8 +39,7 @@ def __init__(
         replicated_layout: bool = False,
         canonical_format: str | None = None,
     ):
-        """
-        Initialize the file mapper. Each worker constructs its own, but
+        """Initialize the file mapper. Each worker constructs its own, but
         `config.json` is shared across workers since rank lives outside the hash.
         When `parallel_agnostic=True`, tp/pp/pcp/dcp are forced to 1 and rank
         to 0 so multiple parallelism layouts collapse into the same folder.
@@ -136,8 +133,7 @@ def get_config_file_path(self) -> str:
 
     @staticmethod
     def _compute_base_path(root_dir: str, fields: dict) -> str:
-        """
-        Layout: <root_dir>/<safe_model_name>_<sha256-prefix>/.
+        """Layout: <root_dir>/<safe_model_name>_<sha256-prefix>/.
         safe_model_name replaces '/' with '_' so HuggingFace IDs don't nest.
         """
         canonical = json.dumps(fields, sort_keys=True, separators=(",", ":"))
diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py
index f111f233b613..a6879d8c44e5 100644
--- a/vllm/v1/kv_offload/tiering/async_lookup.py
+++ b/vllm/v1/kv_offload/tiering/async_lookup.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-AsyncLookupManager: per-tier async lookup manager for secondary tier
+"""AsyncLookupManager: per-tier async lookup manager for secondary tier
 existence checks.
 
 Each secondary tier that wants non-blocking lookups composes its own
@@ -27,8 +26,10 @@
 flush() is called once per step from the tier's on_schedule_end(), posting
 the entire batch as a single queue item so the background thread sees one
 batch per step.
-drain_results() is called before any lookup() calls in the same step, so
-lookup() is a pure OrderedDict operation.
+Results are drained on the first lookup after each flush, at flush(), and
+after worker shutdown. In-flight lookups with no remaining request references
+are retained until their results are drained, allowing new requests to share
+the same probe.
 """
 
 import queue
@@ -36,6 +37,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Collection, Iterable
 from dataclasses import dataclass, field
+from enum import Enum, auto
 
 from vllm.logger import init_logger
 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
@@ -43,16 +45,26 @@
 logger = init_logger(__name__)
 
 
+class LookupPhase(Enum):
+    """Lifecycle phase of a lookup probe."""
+
+    PENDING = auto()  # Accumulated in _lookup_batch, but not yet submitted.
+    IN_FLIGHT = auto()  # Submitted to the worker, but not yet resolved.
+    RESOLVED = auto()  # A final verdict has been applied to the state.
+
+
 @dataclass(slots=True)
 class LookupState:
     generation: int
-    result: bool | None = None  # True (found), False (not found), None
+    phase: LookupPhase = LookupPhase.PENDING
+    # None while pending/in flight; True if the key exists; False if absent or
+    # explicitly marked missing after a failed load.
+    result: bool | None = None
     request_ids: set[str] = field(default_factory=set)  # requests asking for the lookup
 
 
 class AsyncLookupManager(ABC):
-    """
-    Per-tier async lookup manager for secondary tier existence checks.
+    """Per-tier async lookup manager for secondary tier existence checks.
 
     Each secondary tier that wants non-blocking lookups composes its own
     AsyncLookupManager instance internally. The manager maintains lookup
@@ -112,8 +124,7 @@ def __init__(
     def batch_lookup(
         self, keys: list[OffloadKey], req_context: ReqContext
     ) -> Iterable[bool]:
-        """
-        Check whether a batch of blocks exist in this tier.
+        """Check whether a batch of blocks exist in this tier.
 
         Called from the worker thread — must be synchronous and must not
         touch the primary tier or scheduler state.
@@ -127,13 +138,13 @@ def batch_lookup(
     # ------------------------------------------------------------------
 
     def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
-        """
-        Non-blocking lookup called from the scheduler thread.
+        """Non-blocking lookup called from the scheduler thread.
 
         Returns:
             True  — block is present in this tier.
             False — block is not present in this tier.
             None  — result not yet available; retry next step.
+
         """
         if self._need_to_drain:
             self.drain_results()
@@ -155,24 +166,28 @@ def flush(self) -> None:
         Called once per step from on_schedule_end() after all lookup() calls
         are done. The worker receives the full batch and processes it during
         the model-execution window, maximising time available before the next
-        step's drain_results().  Safe to call with an empty batch (no-op).
+        step's drain_results(). Also drains completed lookups when there
+        are no new keys to submit.
         """
+        self.drain_results()
         self._need_to_drain = True
         batch = self._lookup_batch
         self._lookup_batch = []
-        batch = [
-            (key, req_context, generation)
-            for key, req_context, generation in batch
-            if (state := self._lookup_state.get(key)) is not None
-            and state.generation == generation
-        ]
-        if batch:
-            self._lookup_queue.put(batch)
+        in_flight_batch = []
+        for key, req_context, generation in batch:
+            state = self._lookup_state.get(key)
+            if state is None or state.generation != generation:
+                continue
+            assert state.phase is LookupPhase.PENDING
+            state.phase = LookupPhase.IN_FLIGHT
+            in_flight_batch.append((key, req_context, generation))
+        if in_flight_batch:
+            self._lookup_queue.put(in_flight_batch)
 
     def drain_results(self) -> None:
         """Apply pending worker results to _lookup_state.
 
-        Called from lookup() before checking state.
+        Called from lookup(), flush(), and shutdown() on the scheduler thread.
         """
         while True:
             try:
@@ -183,6 +198,10 @@ def drain_results(self) -> None:
                 state = self._lookup_state.get(key)
                 if state is None or state.generation != generation:
                     continue
+                if not state.request_ids:
+                    del self._lookup_state[key]
+                    continue
+                assert state.phase is LookupPhase.IN_FLIGHT
                 # Each lookup generation is enqueued exactly once. A matching
                 # generation must not receive a second result; stale
                 # generations were discarded above.
@@ -192,6 +211,7 @@ def drain_results(self) -> None:
                     "failed-load livelock"
                 )
                 state.result = result
+                state.phase = LookupPhase.RESOLVED
 
     def mark_miss(self, keys: Collection[OffloadKey]) -> None:
         """Force the cached verdict for ``keys`` to False after a failed load, so
@@ -201,9 +221,10 @@ def mark_miss(self, keys: Collection[OffloadKey]) -> None:
             state = self._lookup_state.get(key)
             if state is not None:
                 state.result = False
+                state.phase = LookupPhase.RESOLVED
 
     def cleanup(self, req_id: str) -> None:
-        """Remove entries no longer needed by any active request.
+        """Release request references, retaining in-flight lookups.
 
         Called from the tier's on_request_finished(). Uses the reverse
         index to visit only keys associated with this request.
@@ -211,13 +232,14 @@ def cleanup(self, req_id: str) -> None:
         for key in self._req_keys.pop(req_id, ()):
             state = self._lookup_state[key]
             state.request_ids.discard(req_id)
-            if not state.request_ids:
+            if not state.request_ids and state.phase is not LookupPhase.IN_FLIGHT:
                 del self._lookup_state[key]
 
     def shutdown(self) -> None:
-        """Stop the worker thread."""
+        """Stop the worker thread and drain completed lookups."""
         self._lookup_queue.put(None)  # unblock _worker from _lookup_queue.get()
         self._thread.join()
+        self.drain_results()
 
     # ------------------------------------------------------------------
     # Internal helpers
diff --git a/vllm/v1/kv_offload/tiering/backpressure.py b/vllm/v1/kv_offload/tiering/backpressure.py
new file mode 100644
index 000000000000..deba37c94b2d
--- /dev/null
+++ b/vllm/v1/kv_offload/tiering/backpressure.py
@@ -0,0 +1,384 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+import time
+from abc import ABC, abstractmethod
+
+
+class BackpressurePolicy(ABC):
+    """Decides what action to take when a detector signals pressure."""
+
+    @abstractmethod
+    def should_store(self, detector: BackpressureDetector) -> bool: ...
+
+    @abstractmethod
+    def on_store_skipped(self, num_blocks: int) -> None: ...
+
+    @abstractmethod
+    def pop_stores_dropped(self) -> tuple[int, int]:
+        """Return and reset (stores_dropped, blocks_dropped)."""
+        ...
+
+    @abstractmethod
+    def reset(self) -> None: ...
+
+
+class DropAccountingPolicy(BackpressurePolicy):
+    """Base policy that tracks dropped store/block counters.
+
+    Subclasses implement ``should_store``; the shared drop accounting
+    (``on_store_skipped``/``pop_stores_dropped``/``reset``) lives here so
+    every policy reports these counters identically. Subclasses that add
+    their own state override ``reset`` and call ``super().reset()``.
+    """
+
+    def __init__(self):
+        self._stores_dropped: int = 0
+        self._blocks_dropped: int = 0
+
+    def on_store_skipped(self, num_blocks: int) -> None:
+        self._stores_dropped += 1
+        self._blocks_dropped += num_blocks
+
+    def pop_stores_dropped(self) -> tuple[int, int]:
+        stores, blocks = self._stores_dropped, self._blocks_dropped
+        self._stores_dropped = 0
+        self._blocks_dropped = 0
+        return stores, blocks
+
+    def reset(self) -> None:
+        self._stores_dropped = 0
+        self._blocks_dropped = 0
+
+
+class DropStorePolicy(DropAccountingPolicy):
+    """Silently drop stores to pressured tiers."""
+
+    def should_store(self, detector) -> bool:
+        return not detector.is_under_pressure()
+
+
+class ThrottledDropPolicy(DropAccountingPolicy):
+    """Drop stores proportionally to pressure severity.
+
+    Instead of dropping all stores when pressure is detected, the drop
+    rate scales with how far the EMA exceeds the high watermark::
+
+        overshoot = (ema - high_water) / high_water
+        drop_rate = min(1.0, overshoot / ramp_factor)
+
+    With the default ``ramp_factor=2.0``, stores ramp from 0% drop at
+    the high watermark to 100% drop at 3x the high watermark.  This
+    keeps secondary tiers populated under moderate pressure, reducing
+    cache misses (and therefore recomputation / ITL regression) while
+    still protecting TTFT under severe congestion.
+    """
+
+    def __init__(self, ramp_factor: float = 2.0):
+        super().__init__()
+        self._call_count: int = 0
+        self._ramp_factor = ramp_factor
+
+    def should_store(self, detector: BackpressureDetector) -> bool:
+        # Proportional throttling needs the EMA and high watermark, which
+        # only the EMA detector exposes.
+        assert isinstance(detector, EMABackpressureDetector), (
+            "ThrottledDropPolicy requires an EMABackpressureDetector"
+        )
+        if not detector.is_under_pressure():
+            return True
+        ema = detector.store_latency_ema
+        high = detector._high
+        if high <= 0:
+            return False
+        overshoot = max(0.0, ema - high) / high
+        drop_rate = min(1.0, overshoot / self._ramp_factor)
+        if drop_rate <= 0:
+            return True
+        if drop_rate >= 1.0:
+            return False
+        self._call_count += 1
+        period = max(2, round(1.0 / drop_rate))
+        return (self._call_count % period) != 0
+
+    def reset(self) -> None:
+        super().reset()
+        self._call_count = 0
+
+
+class BackpressureDetector(ABC):
+    """Observes store completion signals and determines pressure state."""
+
+    def __init__(
+        self,
+        policy: BackpressurePolicy | None = None,
+    ):
+        self._policy = policy or ThrottledDropPolicy()
+
+    @property
+    def policy(self) -> BackpressurePolicy:
+        return self._policy
+
+    @abstractmethod
+    def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None: ...
+
+    @abstractmethod
+    def is_under_pressure(self) -> bool: ...
+
+    @abstractmethod
+    def reset(self) -> None: ...
+
+    def update(self, submit_time: float, num_bytes: int) -> None:
+        """Update pressure state from a completed store job.
+
+        Args:
+            submit_time: ``time.monotonic()`` when the job was submitted.
+            num_bytes: Total bytes written (num_blocks * block_size_bytes).
+
+        """
+        if num_bytes <= 0:
+            return
+        elapsed = time.monotonic() - submit_time
+        self.on_store_completed(elapsed, num_bytes)
+
+    def should_store(self, num_blocks: int) -> bool:
+        """Check policy and record skip if rejected."""
+        if not self._policy.should_store(self):
+            self._policy.on_store_skipped(num_blocks)
+            return False
+        return True
+
+    @classmethod
+    def default_config(
+        cls, tier_type: str, *, locality: str | None = None
+    ) -> dict | None:
+        """Return default constructor kwargs for ``tier_type``, or None.
+
+        Args:
+            tier_type: The tier type string (e.g. ``"fs"``, ``"obj"``).
+            locality: Optional locality hint from the tier config
+                (``"LOCAL"`` or ``"REMOTE"``). When ``"REMOTE"``, tiers
+                that would normally get local-storage watermarks (e.g.
+                ``"fs"``) receive network watermarks instead.
+
+        """
+        return None
+
+    @property
+    def stats(self) -> dict[str, float]:
+        return {}
+
+
+class EMABackpressureDetector(BackpressureDetector):
+    """EMA of store latency normalized by transfer size.
+
+    The EMA tracks seconds per megabyte (s/MiB) so that the metric is
+    comparable regardless of how many blocks are in a job or how large
+    each block is.  Water marks are in the same unit.
+
+    Default water marks are derived from fio benchmarks on the WDC H100
+    cluster.  Two presets are provided:
+
+      LOCAL (NVMe/SSD): NVMe sustains ~5 GB/s writes.
+        high=0.005 s/MiB (~200 MB/s) catches severe congestion;
+        low=0.001 s/MiB (~1 GB/s) requires meaningful recovery.
+
+      NETWORK (CephFS, object store, ``obj``tiers, or any tier
+        with ``"locality": "REMOTE"``): CephFS sustains ~1.5 GB/s.
+        high=0.020 s/MiB (~50 MB/s); low=0.005 s/MiB (~200 MB/s).
+
+    ``obj`` tiers get NETWORK defaults automatically. An
+    ``fs`` tier defaults to LOCAL (appropriate for NVMe/SSD); set
+    ``"locality": "REMOTE"`` in the tier config for network-backed
+    filesystems like CephFS.
+    """
+
+    _MIB = 1 << 20
+
+    # EMA smoothing factor: higher values (→1) react faster to latency
+    # spikes but are noisier; lower values (→0) smooth more but lag.
+    DEFAULT_ALPHA = 0.15
+
+    # Number of store completions to collect before seeding the EMA.
+    # During warmup, pressure is never signalled; the EMA is initialized
+    # to the mean of the warmup samples to avoid cold-start false positives.
+    DEFAULT_WARMUP_COMPLETIONS = 10
+
+    # When under pressure and no completions arrive, the EMA decays
+    # toward zero so the detector can eventually recover.  The half-life
+    # controls how quickly: after this many seconds of silence the EMA
+    # halves.  With a 2 s half-life and high_water=0.005, an EMA frozen
+    # at 0.010 reaches low_water=0.001 in ~6.6 s.
+    DEFAULT_DECAY_HALF_LIFE_S = 2.0
+
+    # After pressure clears, stores are throttled to 50% for this many
+    # seconds to prevent a burst of stores from re-overwhelming the tier.
+    DEFAULT_COOLDOWN_S = 1.0
+
+    # Number of consecutive sub-low-watermark completions before the
+    # detector enters a "healthy" fast-path that bypasses policy checks.
+    _HEALTHY_THRESHOLD = 10
+
+    LOCAL_HIGH_WATER_S = 0.010
+    LOCAL_LOW_WATER_S = 0.002
+
+    NETWORK_HIGH_WATER_S = 0.050
+    NETWORK_LOW_WATER_S = 0.010
+
+    _NETWORK_TIER_TYPES = frozenset({"obj"})
+
+    @classmethod
+    def default_config(
+        cls, tier_type: str, *, locality: str | None = None
+    ) -> dict | None:
+        is_remote = locality is not None and locality.upper() == "REMOTE"
+        if is_remote or tier_type in cls._NETWORK_TIER_TYPES:
+            return {
+                "high_water_s": cls.NETWORK_HIGH_WATER_S,
+                "low_water_s": cls.NETWORK_LOW_WATER_S,
+            }
+        return {
+            "high_water_s": cls.LOCAL_HIGH_WATER_S,
+            "low_water_s": cls.LOCAL_LOW_WATER_S,
+        }
+
+    def __init__(
+        self,
+        high_water_s: float,
+        low_water_s: float,
+        alpha: float = DEFAULT_ALPHA,
+        warmup_completions: int = DEFAULT_WARMUP_COMPLETIONS,
+        decay_half_life_s: float = DEFAULT_DECAY_HALF_LIFE_S,
+        cooldown_s: float = DEFAULT_COOLDOWN_S,
+        policy: BackpressurePolicy | None = None,
+    ):
+        super().__init__(policy=policy)
+        if not (0 < alpha <= 1):
+            raise ValueError(f"alpha must be in (0, 1], got {alpha}")
+        if low_water_s > high_water_s:
+            raise ValueError(
+                f"low_water_s ({low_water_s}) must be <= high_water_s ({high_water_s})"
+            )
+
+        self._alpha = alpha
+        self._high = high_water_s
+        self._low = low_water_s
+        self._warmup_completions = warmup_completions
+        self._decay_half_life_s = decay_half_life_s
+        self._cooldown_s = cooldown_s
+        # Exponential moving average of store latency in seconds.
+        self._ema: float = 0.0
+        # Whether the tier is currently considered under pressure.
+        self._under_pressure: bool = False
+        # Number of store completions observed (for warmup gating).
+        self._completions: int = 0
+        # Samples collected during warmup; used to seed the EMA.
+        self._warmup_samples: list[float] = []
+        # Monotonic timestamp of the last EMA update; used for idle decay.
+        self._last_update: float = 0.0
+        # Consecutive completions with EMA below low watermark.
+        self._healthy_streak: int = 0
+        # Monotonic timestamp when pressure was last cleared (for cooldown).
+        self._pressure_cleared_at: float = 0.0
+        self._cooldown_count: int = 0
+
+    def _apply_idle_decay(self) -> None:
+        """Decay the EMA toward zero based on elapsed idle time.
+
+        Uses the formula ``ema * 0.5^(dt / half_life)`` so the EMA
+        halves every ``decay_half_life_s`` seconds of silence.  This
+        lets the detector recover from pressure when no stores are
+        attempted (and therefore no completions arrive to update the
+        EMA organically).
+        """
+        now = time.monotonic()
+        dt = now - self._last_update
+        if dt <= 0 or self._decay_half_life_s <= 0:
+            return
+        self._ema *= 0.5 ** (dt / self._decay_half_life_s)
+        self._last_update = now
+
+    def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None:
+        now = time.monotonic()
+        s_per_mib = elapsed_s / (num_bytes / self._MIB)
+        self._completions += 1
+        if self._completions <= self._warmup_completions:
+            self._warmup_samples.append(s_per_mib)
+            if self._completions == self._warmup_completions:
+                self._ema = sum(self._warmup_samples) / len(self._warmup_samples)
+                self._last_update = now
+            return
+        self._ema = self._alpha * s_per_mib + (1 - self._alpha) * self._ema
+        self._last_update = now
+        if self._ema > self._high:
+            self._under_pressure = True
+            self._pressure_cleared_at = 0.0
+            self._healthy_streak = 0
+        elif self._ema < self._low:
+            if self._under_pressure:
+                self._pressure_cleared_at = now
+                self._cooldown_count = 0
+            self._under_pressure = False
+            self._healthy_streak += 1
+        else:
+            self._healthy_streak = 0
+
+    def is_under_pressure(self) -> bool:
+        if self._under_pressure and self._completions >= self._warmup_completions:
+            self._apply_idle_decay()
+            if self._ema < self._low:
+                self._under_pressure = False
+                self._pressure_cleared_at = time.monotonic()
+                self._cooldown_count = 0
+        return self._under_pressure
+
+    def is_healthy(self) -> bool:
+        """True when the tier has been consistently fast.
+
+        After ``_HEALTHY_THRESHOLD`` consecutive store completions with
+        EMA below the low watermark (and no active pressure), the
+        detector enters a healthy fast-path: ``should_store()`` returns
+        ``True`` immediately, skipping the policy check entirely.
+        """
+        return (
+            not self._under_pressure and self._healthy_streak >= self._HEALTHY_THRESHOLD
+        )
+
+    def should_store(self, num_blocks: int) -> bool:
+        under_pressure = self.is_under_pressure()
+        if self.is_healthy():
+            return True
+        if (
+            not under_pressure
+            and self._pressure_cleared_at > 0
+            and time.monotonic() - self._pressure_cleared_at < self._cooldown_s
+        ):
+            self._cooldown_count += 1
+            if self._cooldown_count % 2 == 0:
+                self._policy.on_store_skipped(num_blocks)
+                return False
+        return super().should_store(num_blocks)
+
+    def reset(self) -> None:
+        self._ema = 0.0
+        self._under_pressure = False
+        self._completions = 0
+        self._warmup_samples.clear()
+        self._last_update = 0.0
+        self._healthy_streak = 0
+        self._pressure_cleared_at = 0.0
+        self._cooldown_count = 0
+        self._policy.reset()
+
+    @property
+    def store_latency_ema(self) -> float:
+        return self._ema
+
+    @store_latency_ema.setter
+    def store_latency_ema(self, value: float) -> None:
+        self._ema = value
+
+    @property
+    def stats(self) -> dict[str, float]:
+        return {"store_latency_ema": self._ema}
diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py
index 34aa6ddf78ad..a90d33eb1136 100644
--- a/vllm/v1/kv_offload/tiering/base.py
+++ b/vllm/v1/kv_offload/tiering/base.py
@@ -1,12 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Abstract interfaces and data types for the secondary tiering layer.
-"""
+"""Abstract interfaces and data types for the secondary tiering layer."""
 
+import time
 from abc import ABC, abstractmethod
 from collections.abc import Collection, Iterable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
@@ -22,6 +21,9 @@
     RequestOffloadingContext,
     ScheduleEndContext,
 )
+from vllm.v1.kv_offload.tiering.backpressure import (
+    BackpressureDetector,
+)
 
 if TYPE_CHECKING:
     from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
@@ -54,6 +56,11 @@ class TieringOffloadingMetrics:
     )
     ACTIVE_PROMOTION_JOBS = "vllm:kv_offload_tiering_active_promotion_jobs"
     ACTIVE_CASCADE_JOBS = "vllm:kv_offload_tiering_active_cascade_jobs"
+    BACKPRESSURE_STORE_LATENCY_EMA = (
+        "vllm:kv_offload_tiering_backpressure_store_latency_ema"
+    )
+    BACKPRESSURE_STORES_DROPPED = "vllm:kv_offload_tiering_backpressure_stores_dropped"
+    BACKPRESSURE_BLOCKS_DROPPED = "vllm:kv_offload_tiering_backpressure_blocks_dropped"
 
 
 @dataclass
@@ -65,6 +72,7 @@ class TransferJob:
     chunk_ids: np.ndarray
     is_promotion: bool
     req_context: ReqContext
+    submit_time: float = field(default_factory=time.monotonic)
 
 
 @dataclass
@@ -79,6 +87,7 @@ class JobResult:
     # indicated by `success`. Must be a subset of the job's original keys.
     successful_keys: Collection[OffloadKey] | None = None
     transfer_time: float | None = None
+    transfer_bytes: int | None = None
 
 
 class ParentManager(ABC):
@@ -119,8 +128,7 @@ def on_request_finished(self, req_context: ReqContext) -> None: ...
 
 
 class SecondaryTierManager(ABC):
-    """
-    Abstract interface for managing a single non-primary offloading tier.
+    """Abstract interface for managing a single non-primary offloading tier.
 
     Secondary tiers cannot directly access GPU memory. All data transfers
     must go through the CPU (primary) tier:
@@ -139,23 +147,37 @@ def __init__(
         offloading_spec: "OffloadingSpec",
         primary_kv_view: memoryview,
         tier_type: str,
+        backpressure_detector: BackpressureDetector | None = None,
     ) -> None:
-        """
-        Args:
-            offloading_spec: Offloading configuration.
-            primary_kv_view: Memoryview of the primary tier's CPU KV cache.
-            tier_type: Tier type identifier, set by SecondaryTierFactory
-                from the registered tier type.
+        """Args:
+        offloading_spec: Offloading configuration.
+        primary_kv_view: Memoryview of the primary tier's CPU KV cache.
+        tier_type: Tier type identifier, set by SecondaryTierFactory
+            from the registered tier type.
+        backpressure_detector: Optional `BackpressureDetector`.
+
         """
         self._offloading_spec = offloading_spec
         self._primary_kv_view: memoryview = primary_kv_view
+        assert primary_kv_view.strides is not None, (
+            "primary_kv_view.strides cannot be None"
+        )
+        self._block_size_bytes: int = primary_kv_view.strides[0]
         self.tier_type = tier_type
         self.locality: Locality | None = None
+        self._bp_detector = backpressure_detector
+
+    @property
+    def block_size_bytes(self) -> int:
+        return self._block_size_bytes
+
+    @property
+    def bp_detector(self) -> BackpressureDetector | None:
+        return self._bp_detector
 
     @abstractmethod
     def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
-        """
-        Check whether a chunk exists in this secondary tier.
+        """Check whether a chunk exists in this secondary tier.
 
         Args:
             key: Offload key to look up.
@@ -165,13 +187,13 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
             HIT if the chunk is present and ready,
             MISS if not found,
             or RETRY if the chunk is being transferred (retry later).
+
         """
         pass
 
     @abstractmethod
     def submit_store(self, job_metadata: TransferJob) -> None:
-        """
-        Submit an async job to store chunks from the primary tier to this
+        """Submit an async job to store chunks from the primary tier to this
         secondary tier.
 
         This method must be lightweight and non-blocking: allocate metadata
@@ -193,13 +215,13 @@ def submit_store(self, job_metadata: TransferJob) -> None:
         Args:
             job_metadata: Job metadata including job_id, keys, and chunk_ids
                           identifying the primary-tier slots to read from.
+
         """
         pass
 
     @abstractmethod
     def submit_load(self, job_metadata: TransferJob) -> None:
-        """
-        Submit an async job to load chunks from this secondary tier to the
+        """Submit an async job to load chunks from this secondary tier to the
         primary tier.
 
         This method must be lightweight and non-blocking: mark chunks as
@@ -218,13 +240,13 @@ def submit_load(self, job_metadata: TransferJob) -> None:
         Args:
             job_metadata: Job metadata including job_id, keys, and chunk_ids
                           identifying the primary-tier slots to write into.
+
         """
         pass
 
     @abstractmethod
     def get_finished_jobs(self) -> Iterable[JobResult]:
-        """
-        Return all jobs (loads and stores) that completed since the last call.
+        """Return all jobs (loads and stores) that completed since the last call.
 
         The framework uses these results to release resources and finalize
         transfers.
@@ -232,6 +254,7 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
         Returns:
             Iterable of JobResult objects for jobs finished since the
             last call.
+
         """
         pass
 
@@ -248,31 +271,30 @@ def take_events(self) -> Iterable[OffloadingEvent]:
         return ()
 
     def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
-        """
-        Mark chunks as recently used for eviction policy.
+        """Mark chunks as recently used for eviction policy.
 
         Args:
             keys: Offload keys to mark as recently used.
             req_context: Per-request context.
+
         """
         return
 
     @abstractmethod
     def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
-        """
-        Called when a new request is first seen by the scheduler.
+        """Called when a new request is first seen by the scheduler.
 
         Returns a RequestOffloadingContext expressing this tier's preference
         for how chunks should be offloaded for this request.
 
         Args:
             req_context: Per-request context.
+
         """
         pass
 
     def on_request_finished(self, req_context: ReqContext) -> None:
-        """
-        Called when a request has finished.
+        """Called when a request has finished.
 
         By the time this is called, all per-request calls for this request
         (submit_store, submit_load, touch) have already been issued, and none
@@ -283,6 +305,7 @@ def on_request_finished(self, req_context: ReqContext) -> None:
 
         Args:
             req_context: per-request context.
+
         """
         return
 
@@ -300,6 +323,7 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None:
 
         Args:
             context: Per-step context from the scheduler.
+
         """
         return
 
diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py
index 967f7cc49826..e3f7c93d7adb 100644
--- a/vllm/v1/kv_offload/tiering/example/manager.py
+++ b/vllm/v1/kv_offload/tiering/example/manager.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ExampleSecondaryTierManager: A simple in-memory secondary tier.
+"""ExampleSecondaryTierManager: A simple in-memory secondary tier.
 
 This implementation provides a minimal secondary tier that stores chunks
 in memory (using a dictionary) with immediate completion. It serves as a
@@ -32,11 +31,11 @@
 
 if TYPE_CHECKING:
     from vllm.v1.kv_offload.base import OffloadingSpec
+    from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector
 
 
 class ExampleSecondaryTierManager(SecondaryTierManager):
-    """
-    A simple in-memory secondary tier.
+    """A simple in-memory secondary tier.
 
     This implementation:
     - Stores chunks in a dictionary (key -> True)
@@ -51,17 +50,23 @@ def __init__(
         primary_kv_view: memoryview,
         tier_type: str,
         custom_param: int = 0,
+        backpressure_detector: "BackpressureDetector | None" = None,
     ):
-        """
-        Initialize the example secondary tier.
+        """Initialize the example secondary tier.
 
         Args:
+            offloading_spec: The offloading spec this tier belongs to.
+            primary_kv_view: Memoryview over the primary tier's KV buffer.
+            tier_type: Name identifying this tier type.
             custom_param: Dummy parameter demonstrating custom args.
+            backpressure_detector: Optional backpressure detector.
+
         """
         super().__init__(
             offloading_spec=offloading_spec,
             primary_kv_view=primary_kv_view,
             tier_type=tier_type,
+            backpressure_detector=backpressure_detector,
         )
 
         logger.info(
@@ -78,8 +83,7 @@ def __init__(
 
     @override
     def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
-        """
-        Check whether a chunk exists in this secondary tier.
+        """Check whether a chunk exists in this secondary tier.
 
         Args:
             key: Offload key to look up.
@@ -87,17 +91,18 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
 
         Returns:
             HIT if the chunk is present, MISS if not found.
+
         """
         return LookupResult.HIT if key in self.chunks else LookupResult.MISS
 
     @override
     def submit_store(self, job_metadata: TransferJob) -> None:
-        """
-        Submit a job to store chunks from primary tier to this tier.
+        """Submit a job to store chunks from primary tier to this tier.
 
         Args:
             job_metadata: Job metadata including job_id, keys, and
                           spec for reading chunks from the primary tier.
+
         """
         keys = job_metadata.keys
         chunk_ids = job_metadata.chunk_ids
@@ -118,12 +123,12 @@ def submit_store(self, job_metadata: TransferJob) -> None:
 
     @override
     def submit_load(self, job_metadata: TransferJob) -> None:
-        """
-        Submit a job to load chunks from this tier to primary tier.
+        """Submit a job to load chunks from this tier to primary tier.
 
         Args:
             job_metadata: Job metadata including job_id, keys, and
                           spec for writing chunks into the primary tier.
+
         """
         keys = job_metadata.keys
         chunk_ids = job_metadata.chunk_ids
@@ -149,12 +154,12 @@ def submit_load(self, job_metadata: TransferJob) -> None:
 
     @override
     def get_finished_jobs(self) -> Iterable[JobResult]:
-        """
-        Poll for finished jobs.
+        """Poll for finished jobs.
 
         Returns:
             Iterable of JobResult objects for all jobs that have
             finished since the last call.
+
         """
         result = self.completed_jobs
         self.completed_jobs = []
diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py
index 17f8430204d9..6a460724de94 100644
--- a/vllm/v1/kv_offload/tiering/factory.py
+++ b/vllm/v1/kv_offload/tiering/factory.py
@@ -8,11 +8,33 @@
 from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
 
 if TYPE_CHECKING:
+    from vllm.v1.kv_offload.backpressure import BackpressureDetector
     from vllm.v1.kv_offload.base import OffloadingSpec
 
 logger = init_logger(__name__)
 
 
+_DEFAULT_PACKAGE = "vllm.v1.kv_offload.tiering.backpressure"
+
+
+def _import_class(name: str) -> type:
+    """Import a class by name.
+
+    If ``name`` contains a dot it is treated as a fully qualified path
+    (e.g. ``'my_package.MyDetector'``).  Otherwise it is looked up in
+    the default backpressure module
+    (``vllm.v1.kv_offload.tiering.backpressure``), so short names like
+    ``'DropStorePolicy'`` or ``'EMABackpressureDetector'`` work out of
+    the box.
+    """
+    if "." in name:
+        module_path, _, class_name = name.rpartition(".")
+    else:
+        module_path, class_name = _DEFAULT_PACKAGE, name
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)
+
+
 class SecondaryTierFactory:
     """Registry for SecondaryTierManager implementations, resolved by type.
 
@@ -47,10 +69,30 @@ def create_secondary_tier(
         config = tier_config.copy()
         tier_type = config.pop("type")
         config.pop("module_path", None)
+        locality = config.get("locality")
+        bp_config = config.pop("backpressure", None)
+        bp_detector = None
+        if bp_config is not None:
+            bp_config = bp_config.copy()
+            bp_cls_name = bp_config.pop("backpressure_cls", None)
+            if bp_cls_name is None:
+                raise ValueError(
+                    "backpressure config must include 'backpressure_cls' "
+                    "(e.g. 'EMABackpressureDetector')"
+                )
+            policy_cls_name = bp_config.pop("policy_cls", None)
+            policy = _import_class(policy_cls_name)() if policy_cls_name else None
+            detector_cls: type[BackpressureDetector] = _import_class(bp_cls_name)
+            defaults = detector_cls.default_config(tier_type, locality=locality)
+            if defaults:
+                for k, v in defaults.items():
+                    bp_config.setdefault(k, v)
+            bp_detector = detector_cls(policy=policy, **bp_config)
         return tier_cls(
             offloading_spec=offloading_spec,
             primary_kv_view=primary_kv_view,
             tier_type=tier_type,
+            backpressure_detector=bp_detector,
             **config,
         )
 
@@ -71,6 +113,7 @@ def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]:
         Raises:
             ValueError: If 'type' is missing, or is neither registered nor
                 resolvable via module_path.
+
         """
         tier_type = tier_config.get("type")
         if not tier_type:
@@ -123,3 +166,9 @@ def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]:
     "vllm.v1.kv_offload.tiering.obj.manager",
     "ObjectStoreSecondaryTierManager",
 )
+
+SecondaryTierFactory.register_tier(
+    "kvcr",
+    "vllm.v1.kv_offload.tiering.kvcr.manager",
+    "KVCRSecondaryTierManager",
+)
diff --git a/vllm/v1/kv_offload/tiering/fs/io.py b/vllm/v1/kv_offload/tiering/fs/io.py
index 9943c90b5472..cda7e46e5964 100644
--- a/vllm/v1/kv_offload/tiering/fs/io.py
+++ b/vllm/v1/kv_offload/tiering/fs/io.py
@@ -94,9 +94,7 @@ def _store_block(
     block_size: int,
     use_o_direct: bool = True,
 ) -> None:
-    """
-    Store callback: Writes to a temp file then atomically replaces the destination.
-    """
+    """Store callback: write to a temp file then atomically replace the target."""
     # Check if block already exists to avoid redundant writes
     if os.path.exists(dest_path):
         return
@@ -172,8 +170,7 @@ def batch_store_block(
     block_size: int,
     use_o_direct: bool = True,
 ) -> None:
-    """
-    Store a batch of KV blocks from a shared buffer to disk in one call.
+    """Store a batch of KV blocks from a shared buffer to disk in one call.
 
     Each block buffer[offsets[i] : offsets[i]+block_size] is written atomically
     to dest_paths[i] via a temp-file rename.  Raises on first error.
@@ -197,8 +194,7 @@ def batch_load_block(
     block_size: int,
     use_o_direct: bool = True,
 ) -> None:
-    """
-    Load a batch of KV blocks from disk into a shared buffer in one call.
+    """Load a batch of KV blocks from disk into a shared buffer in one call.
 
     Block i is read from source_paths[i] into view[offsets[i] : offsets[i]+block_size].
     Raises on first error (see _load_block for the delete-on-short-read policy).
diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py
index f73f70e77ffb..6522cfd23646 100644
--- a/vllm/v1/kv_offload/tiering/fs/manager.py
+++ b/vllm/v1/kv_offload/tiering/fs/manager.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-FileSystemTierManager: Pure-Python file system secondary tier for KV cache offloading.
+"""FileSystemTierManager: pure-Python filesystem tier for KV cache offloading.
 
 Store path:
     Data is written to a temp file (<dest_path.tmp>) via os.write,
@@ -41,6 +40,7 @@
 )
 from vllm.v1.kv_offload.file_mapper import FileMapper
 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
+from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector
 from vllm.v1.kv_offload.tiering.base import (
     JobId,
     JobResult,
@@ -84,8 +84,7 @@ def batch_lookup(
 
 
 class FileSystemTierManager(SecondaryTierManager):
-    """
-    Pure-Python disk-backed secondary tier.
+    """Pure-Python disk-backed secondary tier.
 
     Read-priority threads service load jobs preferentially; write-priority
     threads service store jobs preferentially.  Both groups can drain either
@@ -117,23 +116,27 @@ def __init__(
         n_write_threads: int = 16,
         enable_kv_events: bool = False,
         locality: str | None = None,
+        backpressure_detector: BackpressureDetector | None = None,
     ):
+        """Args:
+        offloading_spec: Contains normalized offloading configuration and
+            blocks_per_chunk.
+        primary_kv_view: Memoryview of the primary tier's CPU KV cache.
+        tier_type: Tier type identifier, set by SecondaryTierFactory.
+        root_dir: Root directory for block files.
+        n_read_threads: Number of read-priority I/O threads.
+        n_write_threads: Number of write-priority I/O threads.
+        enable_kv_events: Emit BlockStored KV events for blocks
+            successfully stored to this tier. Effective only when KV
+            cache events are enabled globally (kv_events_config).
+        locality: Whether this tier's storage is LOCAL or REMOTE relative
+            to the publishing vLLM instance.
+        backpressure_detector: Optional backpressure detector.
+
         """
-        Args:
-            offloading_spec: Contains normalized offloading configuration and
-                blocks_per_chunk.
-            primary_kv_view: Memoryview of the primary tier's CPU KV cache.
-            tier_type: Tier type identifier, set by SecondaryTierFactory.
-            root_dir: Root directory for block files.
-            n_read_threads: Number of read-priority I/O threads.
-            n_write_threads: Number of write-priority I/O threads.
-            enable_kv_events: Emit BlockStored KV events for blocks
-                successfully stored to this tier. Effective only when KV
-                cache events are enabled globally (kv_events_config).
-            locality: Whether this tier's storage is LOCAL or REMOTE relative
-                to the publishing vLLM instance.
-        """
-        super().__init__(offloading_spec, primary_kv_view, tier_type)
+        super().__init__(
+            offloading_spec, primary_kv_view, tier_type, backpressure_detector
+        )
         self.locality = Locality(locality) if locality is not None else None
 
         self.events: list[OffloadingEvent] | None = None
@@ -152,6 +155,8 @@ def __init__(
         # Keys of in-flight load (promotion) jobs, so a failed load can mark
         # its own cached lookup verdicts False (see get_finished_jobs).
         self._load_job_keys: dict[JobId, list[OffloadKey]] = {}
+        # Block count per in-flight job, used to report transfer_bytes.
+        self._job_block_counts: dict[JobId, int] = {}
         # Per load job: how many blocks loaded before a failure (partial keep).
         # Written by the pool worker inside the load task before it raises (so
         # before task_done publishes the job); read on the scheduler thread in
@@ -227,6 +232,7 @@ def submit_store(self, job_metadata: TransferJob) -> None:
             self._block_size,
             self._use_o_direct,
         )
+        self._job_block_counts[job_metadata.job_id] = len(keys)
         self._pool.enqueue_store(job_metadata.job_id, 1, [task])
 
     @override
@@ -236,6 +242,7 @@ def submit_load(self, job_metadata: TransferJob) -> None:
         # keys as a miss (see get_finished_jobs).
         keys = list(job_metadata.keys)
         self._load_job_keys[job_id] = keys
+        self._job_block_counts[job_id] = len(keys)
         paths = [self.file_mapper.get_file_name(key) for key in keys]
         offsets = [int(cid) * self._block_size for cid in job_metadata.chunk_ids]
 
@@ -274,6 +281,8 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
         as a miss here (scheduler thread)."""
         results = []
         for job_id, success, transfer_time in self._pool.get_finished():
+            block_count = self._job_block_counts.pop(job_id, 0)
+            transfer_bytes = block_count * self._block_size if block_count else None
             if self.events is not None:
                 keys = self._store_job_keys.pop(job_id, None)
                 if success and keys:
@@ -301,6 +310,7 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
                         success=False,
                         successful_keys=tuple(successful) if successful else None,
                         transfer_time=transfer_time,
+                        transfer_bytes=transfer_bytes,
                     )
                 )
                 continue
@@ -309,6 +319,7 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
                     job_id=job_id,
                     success=success,
                     transfer_time=transfer_time,
+                    transfer_bytes=transfer_bytes,
                 )
             )
         return results
@@ -333,8 +344,7 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None:
 
     @override
     def shutdown(self) -> None:
-        """
-        Release resources held by this tier.
+        """Release resources held by this tier.
 
         Shuts down the lookup manager and the thread pool,
         clearing pending tasks and waiting for active threads to complete.
diff --git a/vllm/v1/kv_offload/tiering/fs/thread_pool.py b/vllm/v1/kv_offload/tiering/fs/thread_pool.py
index 9db0629192bd..3c246708cffb 100644
--- a/vllm/v1/kv_offload/tiering/fs/thread_pool.py
+++ b/vllm/v1/kv_offload/tiering/fs/thread_pool.py
@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Thread pool:
-    Two queues (load, store) and two sets of threads:
-      - Load-priority threads: drain the load queue first, then the store queue.
-      - Store-priority threads: drain the store queue first, then the load queue.
-    Load jobs are enqueued to the load queue; store jobs to the store queue.
+"""Thread pool:
+Two queues (load, store) and two sets of threads:
+- Load-priority threads: drain the load queue first, then the store queue.
+- Store-priority threads: drain the store queue first, then the load queue.
+Load jobs are enqueued to the load queue; store jobs to the store queue.
 """
 
 import threading
@@ -20,8 +19,7 @@
 
 
 class JobState:
-    """
-    Thread-safe completion tracker for a set of per-block I/O tasks.
+    """Thread-safe completion tracker for a set of per-block I/O tasks.
 
     Each task calls task_done(success) when it finishes.
     """
@@ -50,7 +48,7 @@ def job_id(self) -> JobId:
     def task_done(
         self, success: bool, transfer_time: float
     ) -> tuple[bool, bool, float]:
-        """Returns if job completed and success flag"""
+        """Returns if job completed and success flag."""
         with self._lock:
             self._completed += 1
             self._transfer_time += transfer_time
@@ -60,8 +58,7 @@ def task_done(
 
 
 class DualQueueThreadPool:
-    """
-    Thread pool with two task queues (load and store) and two thread groups.
+    """Thread pool with two task queues (load and store) and two thread groups.
 
     Load-priority threads drain the load queue first, then fall back to the
     store queue.  Store-priority threads do the reverse.  Both queues share
diff --git a/vllm/v1/kv_offload/tiering/kvcr/__init__.py b/vllm/v1/kv_offload/tiering/kvcr/__init__.py
new file mode 100644
index 000000000000..208f01a7cb5e
--- /dev/null
+++ b/vllm/v1/kv_offload/tiering/kvcr/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py
new file mode 100644
index 000000000000..18da2a6c2764
--- /dev/null
+++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py
@@ -0,0 +1,632 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""vLLM secondary-tier adapter for KVCR."""
+
+import ctypes
+import mmap
+import socket
+import time
+import uuid
+from collections.abc import Collection, Iterable
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import msgspec
+from kvcr import (
+    DURATION_METRIC,
+    KVCR,
+    ROUTER_HINT_CAPABILITIES,
+    STATE_METRIC,
+    TRANSFER_BLOCKS_METRIC,
+    TRANSFER_BYTES_METRIC,
+    KVCRBindings,
+)
+from kvcr.config import (
+    FrameworkDramInput,
+    G3Options,
+    KeyAdapter,
+    KVCRBackendConfigs,
+    KVCRConfig,
+    KVCRGuardConfig,
+    LocalDramOptions,
+    RemoteFWDramOptions,
+)
+from kvcr.control_channels import ZmqPeerControlChannel
+from kvcr.policy import (
+    FIFOPolicy,
+    G3FIFOPolicy,
+    G3LRUPolicy,
+    KVCachePolicy,
+    LRUPolicy,
+)
+from kvcr.types import (
+    BlockKey,
+    CacheTier,
+    InventoryEvent,
+    MemDescriptor,
+    OpHandle,
+    PinRequestId,
+    PinResult,
+    QueryStatus,
+)
+from typing_extensions import override
+
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
+    OffloadingConnectorStats,
+)
+from vllm.logger import init_logger
+from vllm.utils.import_utils import resolve_obj_by_qualname
+from vllm.v1.core.kv_cache_utils import (
+    BlockHash,
+    maybe_convert_block_hash,
+)
+from vllm.v1.kv_offload.base import (
+    LookupResult,
+    Medium,
+    OffloadingCounterMetadata,
+    OffloadingEvent,
+    OffloadingGaugeMetadata,
+    OffloadingHistogramMetadata,
+    OffloadingMetricMetadata,
+    OffloadKey,
+    ReqContext,
+    RequestOffloadingContext,
+    get_offload_block_hash,
+)
+from vllm.v1.kv_offload.tiering.base import (
+    JobResult,
+    ParentManager,
+    SecondaryTierManager,
+    TransferJob,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.kv_offload.base import OffloadingSpec
+
+
+_REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES
+
+logger = init_logger(__name__)
+
+_BUILTIN_POLICIES: dict[str, type[KVCachePolicy]] = {
+    "fifo": FIFOPolicy,
+    "lru": LRUPolicy,
+    "g3_fifo": G3FIFOPolicy,
+    "g3_lru": G3LRUPolicy,
+}
+
+
+def _resolve_policy(name: str | None) -> KVCachePolicy | None:
+    if name is None:
+        return None
+    policy_type = _BUILTIN_POLICIES.get(name)
+    if policy_type is None:
+        if "." not in name:
+            raise ValueError(
+                f"Unknown KVCR policy: {name!r}. "
+                f"Supported: {list(_BUILTIN_POLICIES)}; external policies "
+                "must use a fully qualified module.Class path."
+            )
+        policy_type = resolve_obj_by_qualname(name)
+        if not isinstance(policy_type, type) or not issubclass(
+            policy_type, KVCachePolicy
+        ):
+            raise TypeError(f"{name} is not a KVCachePolicy")
+    return policy_type()
+
+
+def _vllm_metric_name(name: str) -> str:
+    return f"vllm:{name}"
+
+
+def _kvcr_metric_definitions() -> dict[str, OffloadingMetricMetadata]:
+    return {
+        _vllm_metric_name(DURATION_METRIC): OffloadingHistogramMetadata(
+            documentation="KVCR operation and stage duration in seconds.",
+            labelnames=("scope", "result"),
+        ),
+        _vllm_metric_name(TRANSFER_BYTES_METRIC): OffloadingCounterMetadata(
+            documentation="Bytes transferred by KVCR operations.",
+            labelnames=("operation",),
+        ),
+        _vllm_metric_name(TRANSFER_BLOCKS_METRIC): OffloadingCounterMetadata(
+            documentation="KV blocks transferred by successful KVCR operations.",
+            labelnames=("operation",),
+        ),
+        _vllm_metric_name(STATE_METRIC): OffloadingGaugeMetadata(
+            documentation="Current KVCR metadata and operation counts.",
+            labelnames=("resource",),
+        ),
+    }
+
+
+class _FrameworkPinAdapter:
+    """Adapt KVCR framework pin callbacks to primary-tier store jobs."""
+
+    def __init__(self, tier: "KVCRSecondaryTierManager") -> None:
+        self._tier = tier
+        self._next_request_id = 0
+        self._next_pin_request_id = 0
+        self._pending_pins: dict[PinRequestId, tuple[BlockKey, ...]] = {}
+        self._completed_pins: list[tuple[PinRequestId, PinResult]] = []
+        self._pin_jobs: dict[str, int] = {}
+        self._pin_job_results: list[JobResult] = []
+
+    def request_pin(self, keys: Collection[BlockKey]) -> PinRequestId:
+        # The parent callback is valid only from serve_external_requests().
+        request = PinRequestId(self._next_pin_request_id)
+        self._next_pin_request_id += 1
+        self._pending_pins[request] = tuple(keys)
+        return request
+
+    def process_pending(self, parent: ParentManager) -> None:
+        pending = self._pending_pins
+        self._pending_pins = {}
+        self._completed_pins.extend(
+            (request, self._resolve_pin(parent, keys))
+            for request, keys in pending.items()
+        )
+
+    def poll_pin_results(self) -> list[tuple[PinRequestId, PinResult]]:
+        completed = self._completed_pins
+        self._completed_pins = []
+        return completed
+
+    def cancel_pin_request(self, request: PinRequestId) -> None:
+        self._pending_pins.pop(request, None)
+
+    def _resolve_pin(
+        self, parent: ParentManager, keys: Collection[BlockKey]
+    ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None:
+        offload_keys = tuple(OffloadKey(bytes(key)) for key in keys)
+        request_id = f"kvcr-source:{self._next_request_id}"
+        self._next_request_id += 1
+        req_context = ReqContext(req_id=request_id)
+        started = False
+        job: TransferJob | None = None
+        pin_handle: str | None = None
+        try:
+            parent.on_new_request(req_context)
+            started = True
+            hit_keys = tuple(
+                key
+                for key in offload_keys
+                if parent.lookup(key, req_context) is LookupResult.HIT
+            )
+            if not hit_keys:
+                return None
+
+            job = parent.create_store_job(hit_keys, req_context)
+            job_keys = tuple(job.keys)
+            chunk_ids = tuple(int(chunk_id) for chunk_id in job.chunk_ids)
+            if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys):
+                return None
+
+            descriptors: dict[BlockKey, list[MemDescriptor] | None] = {
+                BlockKey(bytes(key)): None for key in offload_keys
+            }
+            descriptors.update(
+                (
+                    BlockKey(bytes(key)),
+                    [self._tier._make_descriptor(chunk_id)],
+                )
+                for key, chunk_id in zip(job_keys, chunk_ids)
+            )
+            pin_handle = f"kvcr-job-{job.job_id}"
+            self._pin_jobs[pin_handle] = job.job_id
+            return pin_handle, descriptors
+        except Exception:
+            logger.warning("KVCR source acquisition failed", exc_info=True)
+            return None
+        finally:
+            if started:
+                try:
+                    parent.on_request_finished(req_context)
+                except Exception:
+                    logger.warning("KVCR source request cleanup failed", exc_info=True)
+            if job is not None and pin_handle is None:
+                self._pin_job_results.append(
+                    JobResult(job_id=job.job_id, success=False)
+                )
+
+    def release_pin(self, pin_handle: str) -> bool:
+        job_id = self._pin_jobs.pop(pin_handle, None)
+        if job_id is not None:
+            self._pin_job_results.append(JobResult(job_id=job_id, success=True))
+        return True
+
+    def has_active_pins(self) -> bool:
+        return bool(self._pin_jobs)
+
+    def take_pin_job_results(self) -> list[JobResult]:
+        results = self._pin_job_results
+        self._pin_job_results = []
+        return results
+
+
+class _VllmKeyAdapter:
+    """Adapt vLLM's local key format to KVCR."""
+
+    def encode(self, framework_key: object) -> BlockKey:
+        if not isinstance(framework_key, bytes):
+            raise TypeError("vLLM offload keys must be bytes")
+        return BlockKey(framework_key)
+
+    def decode(self, key: BlockKey) -> int | bytes:
+        block_hash = BlockHash(get_offload_block_hash(OffloadKey(key)))
+        return maybe_convert_block_hash(block_hash)
+
+
+# Job ID, remaining blocks, aggregate success, and successful load keys.
+_JobState = tuple[int, int, bool, set[OffloadKey] | None]
+
+
+class KVCRSecondaryTierManager(SecondaryTierManager):
+    """Secondary tier wrapper around the KVCR KV P2P API."""
+
+    @classmethod
+    @override
+    def build_metric_definitions(
+        cls, extra_config: dict[str, Any]
+    ) -> dict[str, OffloadingMetricMetadata]:
+        if not extra_config.get("enable_telemetry", False):
+            return {}
+        return _kvcr_metric_definitions()
+
+    def __init__(
+        self,
+        offloading_spec: "OffloadingSpec",
+        primary_kv_view: memoryview,
+        tier_type: str,
+        router_capabilities: Iterable[str] | None = None,
+        control_host: str = "0.0.0.0",
+        control_ports: list[int] | None = None,
+        control_advertise_host: str | None = None,
+        eager_ctrl_connect: bool = True,
+        opportunistic_query: bool = True,
+        enable_telemetry: bool = False,
+        operation_timeout_ms: int = 1000,
+        abandon_timeout_ms: int = 5000,
+        metadata_retry_interval_ms: int = 100,
+        secondary_g2_slots: int = 0,
+        kvcr_service_socket_path: str | None = None,
+        compatibility_digest: str | None = None,
+        policy: str | None = None,
+        g3: dict[str, Any] | None = None,
+        local_dram_backend: str = "UCX",
+        remote_fw_dram_backend: str = "UCX",
+    ) -> None:
+        super().__init__(offloading_spec, primary_kv_view, tier_type)
+        selected_policy = _resolve_policy(policy)
+        if (kvcr_service_socket_path is None) != (compatibility_digest is None):
+            raise ValueError(
+                "kvcr_service_socket_path and compatibility_digest must be "
+                "configured together"
+            )
+        missing = _REQUIRED_ROUTER_CAPABILITIES - set(router_capabilities or ())
+        if missing:
+            raise ValueError(
+                f"KVCR secondary tier requires router capabilities: {sorted(missing)}"
+            )
+        events_enabled = offloading_spec.kv_events_config.enable_kv_cache_events
+        self_describing_events = (
+            offloading_spec.kv_events_config.self_describing_kv_events
+        )
+        if (
+            events_enabled
+            and (
+                secondary_g2_slots > 0
+                or kvcr_service_socket_path is not None
+                or g3 is not None
+            )
+            and not self_describing_events
+        ):
+            raise ValueError(
+                "KVCR local G2/G3 inventory requires self_describing_kv_events"
+            )
+        self._key_adapter: KeyAdapter = _VllmKeyAdapter()
+        advertise_host = control_advertise_host or socket.gethostname()
+        dp_local_rank = offloading_spec.config.parallel.data_parallel_rank_local
+        if dp_local_rank is None:
+            if control_ports is None or len(control_ports) != 1:
+                raise ValueError(
+                    "control_ports must contain exactly one port when "
+                    "data_parallel_rank_local is unset"
+                )
+            dp_local_rank = 0
+        if control_ports is None or not 0 <= dp_local_rank < len(control_ports):
+            raise ValueError(
+                "control_ports must contain a port for "
+                f"local data-parallel rank {dp_local_rank}"
+            )
+        if any(not 1 <= port <= 65535 for port in control_ports):
+            raise ValueError("control_ports must be between 1 and 65535")
+        control = ZmqPeerControlChannel(
+            control_host,
+            control_ports[dp_local_rank],
+            advertise_host,
+        )
+        # Give colocated workers distinct transport listen ports.
+        with socket.socket() as _s:
+            _s.bind(("", 0))
+            _nixl_listen_port = _s.getsockname()[1]
+        self._primary_base_addr = ctypes.addressof(
+            ctypes.c_char.from_buffer(primary_kv_view)
+        )
+        if primary_kv_view.strides is None:
+            raise ValueError("primary KV memoryview must expose strides")
+        self._primary_row_stride = int(primary_kv_view.strides[0])
+        if secondary_g2_slots < 0:
+            raise ValueError("secondary_g2_slots must be non-negative")
+        if (
+            secondary_g2_slots or kvcr_service_socket_path is not None
+        ) and not events_enabled:
+            logger.warning(
+                "KVCR local DRAM is enabled, but KV cache events are disabled; "
+                "local DRAM inventory will not be published"
+            )
+        local_mapping: mmap.mmap | None = None
+        local_dram: LocalDramOptions | None = None
+        if secondary_g2_slots:
+            if kvcr_service_socket_path is not None:
+                logger.warning(
+                    "secondary_g2_slots is ignored when "
+                    "kvcr_service_socket_path is configured"
+                )
+            else:
+                local_mapping = mmap.mmap(
+                    -1, secondary_g2_slots * self._primary_row_stride
+                )
+                local_dram = LocalDramOptions(
+                    pools=[
+                        (
+                            "",
+                            ctypes.addressof(ctypes.c_char.from_buffer(local_mapping)),
+                            len(local_mapping),
+                        )
+                    ],
+                    backend=local_dram_backend,
+                )
+        guard_config = (
+            KVCRGuardConfig(
+                kvcr_service_socket_path=kvcr_service_socket_path,
+                guard_index=dp_local_rank,
+                compatibility_digest=compatibility_digest,
+            )
+            if kvcr_service_socket_path is not None and compatibility_digest is not None
+            else None
+        )
+        g3_config = None
+        if g3 is not None:
+            g3_values = dict(g3)
+            g3_values["paths"] = tuple(Path(path) for path in g3_values["paths"])
+            g3_config = G3Options(**g3_values)
+        nixl_agent_name = f"KVCR-{uuid.uuid4()}"
+        self._framework_pin_adapter = _FrameworkPinAdapter(self)
+        self._inventory_events: list[OffloadingEvent] = []
+
+        try:
+            self._kvcr = KVCR(
+                KVCRConfig(
+                    pool_layouts=[("", self._primary_row_stride)],
+                    enable_telemetry=enable_telemetry,
+                    operation_timeout_ms=operation_timeout_ms,
+                    abandon_timeout_ms=abandon_timeout_ms,
+                    nixl_agent_name=nixl_agent_name,
+                    nixl_listen_port=_nixl_listen_port,
+                ),
+                KVCRBindings(
+                    request_pin=self._framework_pin_adapter.request_pin,
+                    poll_pin_results=self._framework_pin_adapter.poll_pin_results,
+                    release_pin=self._framework_pin_adapter.release_pin,
+                    cancel_pin_request=(self._framework_pin_adapter.cancel_pin_request),
+                    framework_control=control,
+                    key_adapter=self._key_adapter,
+                    inventory_sink=(self._record_inventory if events_enabled else None),
+                    stats_factory=(
+                        OffloadingConnectorStats if enable_telemetry else None
+                    ),
+                    policy=selected_policy,
+                ),
+                KVCRBackendConfigs(
+                    framework_dram=FrameworkDramInput(
+                        self._primary_base_addr, primary_kv_view.nbytes
+                    ),
+                    local_dram=local_dram,
+                    g3=g3_config,
+                    remote_fw_dram=RemoteFWDramOptions(
+                        eager_ctrl_connect=eager_ctrl_connect,
+                        opportunistic_query=opportunistic_query,
+                        metadata_retry_interval_ms=metadata_retry_interval_ms,
+                        backend=remote_fw_dram_backend,
+                    ),
+                ),
+                guard_config,
+            )
+        except BaseException:
+            control.close()
+            if local_mapping is not None:
+                local_mapping.close()
+            raise
+        self._local_dram_mmap = local_mapping
+        self._finished_jobs: list[JobResult] = []
+        self._jobs_by_op: dict[OpHandle, _JobState] = {}
+
+    @override
+    def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
+        block_key = self._key_adapter.encode(key)
+        status, _ = self._kvcr.query((block_key,), req_context.req_id)[0]
+        if status is QueryStatus.FETCHING:
+            return LookupResult.RETRY
+        if status in (QueryStatus.HIT, QueryStatus.FETCHABLE):
+            return LookupResult.HIT
+        return LookupResult.MISS
+
+    @override
+    def submit_load(self, job_metadata: TransferJob) -> None:
+        blocks = {
+            self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))]
+            for key, chunk_id in zip(
+                job_metadata.keys, job_metadata.chunk_ids, strict=True
+            )
+        }
+        if not blocks:
+            self._finished_jobs.append(
+                JobResult(job_id=job_metadata.job_id, success=True)
+            )
+            return
+
+        op_handle = self._kvcr.deliver(
+            blocks, request_id=job_metadata.req_context.req_id
+        )
+        self._jobs_by_op[op_handle] = (
+            job_metadata.job_id,
+            len(blocks),
+            True,
+            set(),
+        )
+
+    @override
+    def submit_store(self, job_metadata: TransferJob) -> None:
+        blocks = {
+            self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))]
+            for key, chunk_id in zip(
+                job_metadata.keys, job_metadata.chunk_ids, strict=True
+            )
+        }
+        if not blocks:
+            self._finished_jobs.append(
+                JobResult(job_id=job_metadata.job_id, success=True)
+            )
+            return
+        op_handle = self._kvcr.deposit(blocks)
+        self._jobs_by_op[op_handle] = (
+            job_metadata.job_id,
+            len(blocks),
+            True,
+            None,
+        )
+
+    @override
+    def get_finished_jobs(self) -> Iterable[JobResult]:
+        results = self._finished_jobs
+        self._finished_jobs = []
+        results.extend(self._poll_finished_jobs())
+        return results
+
+    @override
+    def get_stats(self) -> OffloadingConnectorStats | None:
+        stats = cast(OffloadingConnectorStats | None, self._kvcr.get_stats())
+        if stats is not None:
+            # Add the framework namespace only at the vLLM boundary.
+            for metrics in stats.data.values():
+                for name in tuple(metrics):
+                    metrics[_vllm_metric_name(name)] = metrics.pop(name)
+        return stats
+
+    @override
+    def serve_external_requests(self, parent: ParentManager) -> None:
+        self._framework_pin_adapter.process_pending(parent)
+
+    def _poll_finished_jobs(self) -> list[JobResult]:
+        results: list[JobResult] = []
+        for op_handle, entries in self._kvcr.poll_completed():
+            job_state = self._jobs_by_op.get(op_handle)
+            if job_state is None:
+                continue
+            job_id, remaining, success, successful_keys = job_state
+            remaining -= len(entries)
+            success = success and all(entry.success for entry in entries.values())
+            if successful_keys is not None:
+                successful_keys.update(
+                    OffloadKey(bytes(key))
+                    for key, entry in entries.items()
+                    if entry.success
+                )
+            if remaining > 0:
+                self._jobs_by_op[op_handle] = (
+                    job_id,
+                    remaining,
+                    success,
+                    successful_keys,
+                )
+                continue
+            self._jobs_by_op.pop(op_handle, None)
+            results.append(
+                JobResult(
+                    job_id=job_id,
+                    success=success,
+                    successful_keys=successful_keys if not success else None,
+                )
+            )
+        results.extend(self._framework_pin_adapter.take_pin_job_results())
+        return results
+
+    @override
+    def has_pending_work(self) -> bool:
+        # Keep scheduler polling: control messages may arrive without a load job.
+        return True
+
+    @override
+    def take_events(self) -> Iterable[OffloadingEvent]:
+        events = self._inventory_events
+        self._inventory_events = []
+        return events
+
+    @override
+    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
+        if (hint := req_context.kv_hints) is not None:
+            self._kvcr.submit_hint(
+                request_id=req_context.req_id,
+                hints=msgspec.to_builtins(hint),
+            )
+        return RequestOffloadingContext()
+
+    @override
+    def on_request_finished(self, req_context: ReqContext) -> None:
+        self._kvcr.discard_hint(req_context.req_id)
+
+    @override
+    def drain_jobs(self) -> None:
+        while self._jobs_by_op or self._framework_pin_adapter.has_active_pins():
+            drained = self._poll_finished_jobs()
+            if drained:
+                self._finished_jobs.extend(drained)
+            else:
+                time.sleep(0.001)
+
+    @override
+    def shutdown(self) -> None:
+        self.drain_jobs()
+        # TODO: Keep registered buffers alive if KVCR.close() fails.
+        self._kvcr.close()
+        if self._local_dram_mmap is not None:
+            self._local_dram_mmap.close()
+            self._local_dram_mmap = None
+
+    def _record_inventory(self, event: InventoryEvent) -> None:
+        if event.tier is CacheTier.LOCAL_G2:
+            medium = Medium.CPU
+        elif event.tier is CacheTier.G3:
+            medium = Medium.STORAGE
+        else:
+            raise ValueError(f"unsupported KVCR inventory tier: {event.tier}")
+        self._inventory_events.append(
+            OffloadingEvent(
+                keys=[OffloadKey(bytes(key)) for key in event.keys],
+                medium=medium,
+                removed=event.removed,
+                ownership="kvcr",
+                removal_expected=True,
+            )
+        )
+
+    def _make_descriptor(self, chunk_id: int) -> MemDescriptor:
+        return MemDescriptor(
+            end_point_name=self._kvcr.config.nixl_agent_name,
+            mem_type="DRAM",
+            addr=self._primary_base_addr + chunk_id * self._primary_row_stride,
+            size=self._primary_row_stride,
+            device_Id=0,
+            info="",
+        )
diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py
index 0940aaeae07b..ce518042b656 100644
--- a/vllm/v1/kv_offload/tiering/manager.py
+++ b/vllm/v1/kv_offload/tiering/manager.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-TieringOffloadingManager: Multi-tier KV cache offloading orchestrator.
+"""TieringOffloadingManager: Multi-tier KV cache offloading orchestrator.
 
 This manager coordinates between a CPU primary tier (with direct GPU access)
 and zero or more secondary tiers (Storage, Network, etc.) to provide
@@ -110,13 +109,22 @@ def __init__(
         # read/write is for CPU<->secondary transfers,
         # load/store is for CPU<->GPU transfers.
         # These aliases avoid calling prepare_load inside a store path.
-        self.prepare_read = self.prepare_load
         self.complete_read = self.complete_load
         self.prepare_write = self.prepare_store
         self.complete_write = self.complete_store
 
         self._kv_memoryview = mmap_region.create_kv_memoryview()
 
+    def prepare_read(
+        self, keys: Collection[OffloadKey], req_context: ReqContext
+    ) -> LoadStoreSpec:
+        """Pin chunks for a CPU-to-secondary transfer.
+
+        Cascade reads are implementation details of tiering, not additional
+        request accesses, so they must not alter request-scoped recency.
+        """
+        return self._prepare_load(keys, req_context, record_access=False)
+
     def get_kv_memoryview(self) -> memoryview:
         """Return the memoryview over the primary tier's KV cache buffer.
 
@@ -165,8 +173,7 @@ def on_request_finished(self, req_context: ReqContext) -> None:
 
 
 class TieringOffloadingManager(OffloadingManager):
-    """
-    Orchestrates multi-tier KV cache offloading.
+    """Orchestrates multi-tier KV cache offloading.
 
     This manager coordinates between a CPU primary tier (with direct GPU access)
     and zero or more secondary tiers (Storage, Network, etc.) to provide
@@ -185,13 +192,13 @@ def __init__(
         primary_tier: CPUPrimaryTierOffloadingManager,
         secondary_tiers: list[SecondaryTierManager] | None = None,
     ):
-        """
-        Initialize the TieringOffloadingManager.
+        """Initialize the TieringOffloadingManager.
 
         Args:
             primary_tier: The primary tier manager (CPU-based).
             secondary_tiers: List of secondary tier managers (e.g., Storage,
                             Network). Can be None or empty list.
+
         """
         self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier
         self.secondary_tiers = secondary_tiers or []
@@ -231,6 +238,10 @@ def __init__(
             for tier_idx, tier in enumerate(self.secondary_tiers)
         }
 
+        self._tier_index: dict[SecondaryTierManager, int] = {
+            tier: i for i, tier in enumerate(self.secondary_tiers)
+        }
+
     @property
     def _transfer_jobs(self) -> dict[JobId, JobMetadata]:
         return self._jobs
@@ -250,8 +261,7 @@ def _pop_job(self, job_id: JobId) -> JobMetadata | None:
         return self._jobs.pop(job_id, None)
 
     def _maybe_process_finished_jobs(self):
-        """
-        Poll secondary tiers for completed jobs (at most once per step).
+        """Poll secondary tiers for completed jobs (at most once per step).
 
         Guarded by _processed_jobs_this_step: the first call in an engine step
         does the actual polling; subsequent calls are no-ops. The flag is reset
@@ -297,8 +307,7 @@ def _complete_promotion(
             )
 
     def _process_finished_jobs(self):
-        """
-        Unconditionally poll all secondary tiers for completed jobs.
+        """Unconditionally poll all secondary tiers for completed jobs.
 
         This method:
         1. Calls get_finished_jobs() on each secondary tier
@@ -332,6 +341,43 @@ def _process_finished_jobs(self):
                     self.primary_tier.complete_read(
                         transfer_job.keys, transfer_job.req_context
                     )
+                    if completed_job.success:
+                        self._update_backpressure(tier, job_metadata, completed_job)
+
+    def _should_store_to_tier(
+        self, tier: SecondaryTierManager, num_blocks: int
+    ) -> bool:
+        detector = tier.bp_detector
+        if detector is None:
+            return True
+        return detector.should_store(num_blocks)
+
+    def _update_backpressure(
+        self,
+        tier: SecondaryTierManager,
+        job_metadata: JobMetadata,
+        completed_job: JobResult,
+    ) -> None:
+        detector = tier.bp_detector
+        if detector is None:
+            return
+        was_under_pressure = detector.is_under_pressure()
+        tj = job_metadata.transfer_job
+        num_bytes = (
+            completed_job.transfer_bytes
+            if completed_job.transfer_bytes is not None
+            else len(tj.keys) * tier.block_size_bytes
+        )
+        detector.update(tj.submit_time, num_bytes)
+        if detector.is_under_pressure() != was_under_pressure:
+            tier_idx = self._tier_index[tier]
+            logger.info(
+                "Tier #%d (%s) back-pressure %s (stats=%s)",
+                tier_idx,
+                tier.tier_type,
+                "activated" if detector.is_under_pressure() else "cleared",
+                detector.stats,
+            )
 
     @override
     def lookup(
@@ -341,8 +387,7 @@ def lookup(
         *,
         exclude_tier_idx: int | None = None,
     ) -> LookupResult:
-        """
-        Check whether a single chunk is offloaded and ready.
+        """Check whether a single chunk is offloaded and ready.
 
         Algorithm:
             1. Process any completed async jobs first.
@@ -353,6 +398,7 @@ def lookup(
         Args:
             key: Chunk hash to look up.
             req_context: Per-request context.
+            exclude_tier_idx: Skip this tier index during the lookup.
 
         Returns:
             HIT       — chunk is ready in the primary tier.
@@ -361,6 +407,7 @@ def lookup(
             RETRY     — promotion started or a secondary tier is busy.
             MISS      — chunk not found in any tier, or primary is full
                         and cannot accept a promotion.
+
         """
         # Poll first so a promotion that finished since the last call is
         # already reflected as HIT (not stale HIT_PENDING/MISS) below, and
@@ -423,8 +470,7 @@ def _initiate_promotion(
         key: OffloadKey,
         req_context: ReqContext,
     ) -> bool:
-        """
-        Queue a chunk for promotion from a secondary tier to the primary tier.
+        """Queue a chunk for promotion from a secondary tier to the primary tier.
 
         Allocates space in the primary tier immediately (sets ref_cnt=-1 so
         subsequent lookups within the same step see the slot as in-flight),
@@ -439,6 +485,7 @@ def _initiate_promotion(
 
         Returns:
             True if promotion was initiated, False if primary tier is full.
+
         """
         # Allocate space in primary tier for promoted chunk.
         # Must happen immediately so primary.lookup() returns None (in-flight)
@@ -496,8 +543,7 @@ def _flush_pending_promotions(self) -> None:
     def prepare_load(
         self, keys: Collection[OffloadKey], req_context: ReqContext
     ) -> LoadStoreSpec:
-        """
-        Prepare chunks to be loaded from primary tier to GPU.
+        """Prepare chunks to be loaded from primary tier to GPU.
 
         Callers only pass keys already confirmed HIT by lookup() earlier this
         step.
@@ -511,17 +557,18 @@ def prepare_load(
 
         Returns:
             LoadStoreSpec for reading from primary tier.
+
         """
         return self.primary_tier.prepare_load(keys, req_context)
 
     @override
     def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
-        """
-        Mark chunks as recently used in all tiers.
+        """Mark chunks as recently used in all tiers.
 
         Args:
             keys: Chunks to mark as recently used.
             req_context: Per-request context.
+
         """
         self.primary_tier.touch(keys, req_context)
         for tier in self.secondary_tiers:
@@ -529,8 +576,7 @@ def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
 
     @override
     def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
-        """
-        Mark chunks as done loading from primary tier to GPU.
+        """Mark chunks as done loading from primary tier to GPU.
 
         This decrements ref_cnt on the chunks in the primary tier, allowing
         them to be evicted again.
@@ -538,6 +584,7 @@ def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
         Args:
             keys: Chunks that finished loading.
             req_context: Per-request context.
+
         """
         self.primary_tier.complete_load(keys, req_context)
 
@@ -545,8 +592,7 @@ def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
     def prepare_store(
         self, keys: Collection[OffloadKey], req_context: ReqContext
     ) -> PrepareStoreOutput | None:
-        """
-        Prepare chunks to be stored from GPU to primary tier.
+        """Prepare chunks to be stored from GPU to primary tier.
 
         CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure
         that any completed async transfers have their ref_cnt decremented
@@ -562,6 +608,7 @@ def prepare_store(
         Returns:
             PrepareStoreOutput describing where to store chunks and what was
             evicted, or None if store cannot proceed.
+
         """
         # Step 1: Poll for completed async jobs FIRST
         # _process_finished_jobs() handles two kinds of completions here:
@@ -608,8 +655,7 @@ def _cascade_existing_chunks_to_request_level_tiers(
         req_context: ReqContext,
         request_level_tiers: set[int],
     ) -> None:
-        """
-        For tiers that requested request-level policy, submit_store() for
+        """For tiers that requested request-level policy, submit_store() for
         chunks that are already present in the primary tier.
 
         A key whose primary write is still in flight (HIT_PENDING) cannot be
@@ -638,8 +684,10 @@ def _cascade_existing_chunks_to_request_level_tiers(
             return
 
         for tier_idx in request_level_tiers:
-            job_metadata = self.create_store_job(ready_keys, req_context, tier_idx)
             tier = self.secondary_tiers[tier_idx]
+            if not self._should_store_to_tier(tier, len(ready_keys)):
+                continue
+            job_metadata = self.create_store_job(ready_keys, req_context, tier_idx)
             tier.submit_store(job_metadata)
 
     def _flush_pending_cascades(self) -> None:
@@ -665,8 +713,7 @@ def complete_store(
         req_context: ReqContext,
         success: bool = True,
     ) -> None:
-        """
-        Mark chunks as done storing from GPU to primary tier.
+        """Mark chunks as done storing from GPU to primary tier.
 
         This is where secondary tier cascading happens — after chunks are
         confirmed to be in the primary tier, they are cascaded to ALL
@@ -682,6 +729,7 @@ def complete_store(
             keys: Chunks that finished storing.
             success: Whether the GPU→primary transfer succeeded.
             req_context: Per-request context forwarded to primary.prepare_read().
+
         """
         # Step 1: Complete store in primary tier (makes chunks loadable)
         self.primary_tier.complete_store(keys, req_context, success)
@@ -693,6 +741,8 @@ def complete_store(
             # eviction during the async transfer). One prepare_read() call per
             # secondary tier.
             for tier_idx, tier in enumerate(self.secondary_tiers):
+                if not self._should_store_to_tier(tier, len(keys)):
+                    continue
                 job_metadata = self.create_store_job(keys, req_context, tier_idx)
                 tier.submit_store(job_metadata)
 
@@ -739,8 +789,7 @@ def on_new_request(
         *,
         exclude_tier_idx: int | None = None,
     ) -> RequestOffloadingContext:
-        """
-        Query each secondary tier for its offload policy preference.
+        """Query each secondary tier for its offload policy preference.
 
         Returns REQUEST_LEVEL if ANY secondary tier wants request-level.
         Only stores REQUEST_LEVEL tier decisions for use in prepare_store.
@@ -781,11 +830,10 @@ def _maybe_finalize_request(
         req_id: str,
         exclude_tier_idx: int | None = None,
     ) -> None:
-        """Finalize secondary tiers once no more store cascades can be submitted.
+        """Finalize secondary tiers once no more cascades can be submitted.
 
-        Finalization means forwarding on_request_finished() to secondary tiers.
-        It is delayed until pending GPU->primary stores finish, since their
-        complete_store() callbacks may still submit primary->secondary stores.
+        Their finalization is delayed until pending GPU->primary stores
+        finish, since those callbacks may still submit secondary stores.
         """
         state = self._req_state[req_id]
         if not state.is_finished:
@@ -850,6 +898,7 @@ def take_events(self) -> Iterable[OffloadingEvent]:
 
         Yields:
             New OffloadingEvents collected by each tier since the last call.
+
         """
         yield from self.primary_tier.take_events()
         for tier in self.secondary_tiers:
@@ -900,6 +949,10 @@ def reset_cache(self) -> None:
             del self._req_state[req_id]
         self._processed_jobs_this_step = False
 
+        for tier in self.secondary_tiers:
+            if tier.bp_detector is not None:
+                tier.bp_detector.reset()
+
     @override
     def get_stats(self) -> OffloadingConnectorStats | None:
         stats = self.primary_tier.get_stats()
@@ -907,6 +960,7 @@ def get_stats(self) -> OffloadingConnectorStats | None:
         if stats is not None and stats.is_empty():
             stats = None
 
+        self._metrics.record_backpressure(self.secondary_tiers)
         metrics_stats = self._metrics.take_stats()
         if metrics_stats is not None:
             if stats is None:
diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py
index 55d56999398f..240ea57dd70a 100644
--- a/vllm/v1/kv_offload/tiering/metrics.py
+++ b/vllm/v1/kv_offload/tiering/metrics.py
@@ -12,6 +12,7 @@
 from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext
 from vllm.v1.kv_offload.tiering.base import (
     JobResult,
+    SecondaryTierManager,
     TieringOffloadingMetrics,
     TransferJob,
 )
@@ -130,6 +131,36 @@ def on_promotion_allocation_failure(self) -> None:
             TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES
         )
 
+    def record_backpressure(
+        self,
+        tiers: list[SecondaryTierManager],
+    ) -> None:
+        for i, tier in enumerate(tiers):
+            detector = tier.bp_detector
+            if detector is None:
+                continue
+            label = self.tier_label(i)
+            ema = detector.stats.get("store_latency_ema")
+            if ema is not None:
+                self._stats.set_gauge(
+                    TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA,
+                    ema,
+                    labelvalues=label,
+                )
+            stores_dropped, blocks_dropped = detector.policy.pop_stores_dropped()
+            if stores_dropped > 0:
+                self._stats.increase_counter(
+                    TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED,
+                    stores_dropped,
+                    labelvalues=label,
+                )
+            if blocks_dropped > 0:
+                self._stats.increase_counter(
+                    TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED,
+                    blocks_dropped,
+                    labelvalues=label,
+                )
+
     def take_stats(self) -> OffloadingConnectorStats | None:
         active_transfer_stats = OffloadingConnectorStats()
         self._observe_active_transfer_stats(active_transfer_stats)
diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py
index 50a7e611f674..1886e3b81168 100644
--- a/vllm/v1/kv_offload/tiering/obj/manager.py
+++ b/vllm/v1/kv_offload/tiering/obj/manager.py
@@ -22,6 +22,7 @@
 )
 from vllm.v1.kv_offload.file_mapper import FileMapper
 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
+from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector
 from vllm.v1.kv_offload.tiering.base import (
     JobId,
     JobResult,
@@ -112,22 +113,29 @@ def __init__(
         io_threads: int = 4,
         enable_kv_events: bool = False,
         locality: str | None = None,
+        backpressure_detector: BackpressureDetector | None = None,
     ):
+        """Args:
+        offloading_spec: Offloading configuration.
+        primary_kv_view: Memoryview of the primary tier's CPU KV cache.
+        tier_type: Tier type identifier, set by SecondaryTierFactory.
+        store_config: Object store connection parameters (see ObjStoreConfig).
+        prefix: Key prefix prepended to all object keys.
+        io_threads: Number of NIXL I/O threads.
+        enable_kv_events: Emit BlockStored KV events for blocks
+            successfully stored to this tier. Effective only when KV
+            cache events are enabled globally (kv_events_config).
+        locality: Whether this tier's storage is LOCAL or REMOTE relative
+            to the publishing vLLM instance.
+        backpressure_detector: Optional backpressure detector.
+
         """
-        Args:
-            offloading_spec: Offloading configuration.
-            primary_kv_view: Memoryview of the primary tier's CPU KV cache.
-            tier_type: Tier type identifier, set by SecondaryTierFactory.
-            store_config: Object store connection parameters (see ObjStoreConfig).
-            prefix: Key prefix prepended to all object keys.
-            io_threads: Number of NIXL I/O threads.
-            enable_kv_events: Emit BlockStored KV events for blocks
-                successfully stored to this tier. Effective only when KV
-                cache events are enabled globally (kv_events_config).
-            locality: Whether this tier's storage is LOCAL or REMOTE relative
-                to the publishing vLLM instance.
-        """
-        super().__init__(offloading_spec, primary_kv_view, tier_type)
+        super().__init__(
+            offloading_spec,
+            primary_kv_view,
+            tier_type,
+            backpressure_detector,
+        )
         self.locality = Locality(locality) if locality is not None else None
 
         self.events: list[OffloadingEvent] | None = None
diff --git a/vllm/v1/kv_offload/tiering/p2p/control/base.py b/vllm/v1/kv_offload/tiering/p2p/control/base.py
index 7016cfb162ae..5a7030737bdb 100644
--- a/vllm/v1/kv_offload/tiering/p2p/control/base.py
+++ b/vllm/v1/kv_offload/tiering/p2p/control/base.py
@@ -139,6 +139,7 @@ def connect(self, peer_id: str) -> ControlConnection:
         A connection to peer_id that is registered but no longer alive is
         retired and replaced; a live one is a duplicate and must not be
         replaced.
+
         """
         ...
 
@@ -158,6 +159,7 @@ def poll(self) -> Sequence[ControlConnection]:
             The returned sequence is read-only — callers must not mutate
             it. The caller is responsible for creating sessions for
             these connections.
+
         """
         ...
 
diff --git a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py
index 42dc2c69846b..b28ed76dac97 100644
--- a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py
+++ b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-ZMQ-based transport layer for P2P KV cache sharing.
+"""ZMQ-based transport layer for P2P KV cache sharing.
 
 Provides ZmqConnection (per-peer messaging) and ZmqTransport (connection
 management). Message-content agnostic.
diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py
index ffe609d71c1e..62304b7544a3 100644
--- a/vllm/v1/kv_offload/tiering/p2p/data/base.py
+++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py
@@ -87,6 +87,7 @@ class PollResult(NamedTuple):
     Attributes:
         done: Transfer IDs that completed successfully.
         failed: Transfer IDs that failed (error, timeout, etc.).
+
     """
 
     done: Sequence[int]
@@ -174,6 +175,7 @@ def add_remote_peer(
             base_addr: Base address of the peer's block memory region.
             num_blocks: Number of blocks in the peer's region.
             block_len: Size of each block (must match local block_len).
+
         """
         ...
 
@@ -203,6 +205,7 @@ def write_blocks(
         Returns:
             A unique transfer_id (int) to track this transfer, or
             None if the peer is not registered or submission failed.
+
         """
         ...
 
@@ -224,6 +227,7 @@ def poll(self, peer_id: str | None = None) -> PollResult:
             Completed/failed transfers are removed from the inflight set.
 
         Must be called periodically to drive progress checking.
+
         """
         ...
 
@@ -255,6 +259,7 @@ def cancel(
             For mode="wait", the subset of *transfer_ids* still
             tracked as inflight after the cancel attempt. For
             mode="immediate", always [].
+
         """
         ...
 
diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
index bf3adf4b604b..e8fb66dd6e30 100644
--- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
+++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL.
-"""
+"""NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL."""
 
 from __future__ import annotations
 
diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py
index 2cd490171006..59b9c320695f 100644
--- a/vllm/v1/kv_offload/tiering/p2p/manager.py
+++ b/vllm/v1/kv_offload/tiering/p2p/manager.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-P2PSecondaryTierManager: Secondary tier for P2P KV cache sharing.
+"""P2PSecondaryTierManager: Secondary tier for P2P KV cache sharing.
 
 Owns transports and a single bidirectional P2PSession per remote peer.
 """
@@ -47,10 +46,18 @@
 # Reap unbound store batches that have been parked without a FetchMsg
 # binding them to a session for longer than this. Protects against the
 # prefiller buffering blocks for a decoder that never asks (decoder died,
-# network partition, lost kv_request_id). Must be longer than the per-store
-# deadline so the store-timeout path fires first for individual jobs.
+# network partition, lost kv_request_id). Applies only while a batch is
+# unbound: binding hands it to the session, which times the job out under
+# _STORE_TIMEOUT_S instead.
+# Overridable per tier via the ``unbound_store_timeout_s`` config key.
 _UNBOUND_STORE_TIMEOUT_S = 60.0
 
+# How long a reaped kv_request_id is remembered so a fetch that arrives
+# after the reap can be rejected immediately. Must exceed the consumer's
+# load timeout (session.client._LOAD_TIMEOUT_S): past that the consumer has
+# already given up on its own, so there is nobody left to reject.
+_REAPED_ID_RETENTION_S = 60.0
+
 # Time we wait during shutdown for inflight transfers to drain via
 # cancel(mode="wait") before falling back to mode="immediate". Bounded
 # so a wedged peer can't hang shutdown.
@@ -208,6 +215,7 @@ def __init__(
         port: int | None = None,
         backends: list[str] | None = None,
         num_threads: int = 4,
+        unbound_store_timeout_s: float = _UNBOUND_STORE_TIMEOUT_S,
         **kwargs: Any,
     ) -> None:
         """Initialize the P2P secondary tier manager.
@@ -245,9 +253,48 @@ def __init__(
             num_threads: NIXL agent worker threads for the UCX-only
                 branch. Ignored when ``backends`` contains a non-UCX
                 entry.
+            unbound_store_timeout_s: Seconds a producer holds stored blocks
+                for a consumer that has not fetched them yet, before
+                ``_reap_unbound_stores`` drops them. Raise it for
+                deployments whose prefills outlast the default; the blocks
+                keep primary-tier slots pinned for the whole window.
             **kwargs: Reserved for future tier-specific options.
+
+        Raises:
+            ValueError: If ``unbound_store_timeout_s`` is not a positive
+                number, or anything convertible to one.
+
         """
-        super().__init__(offloading_spec, primary_kv_view, tier_type)
+        backpressure_detector = kwargs.pop("backpressure_detector", None)
+        if backpressure_detector is not None:
+            # The generic (store-latency) detector is unreliable for P2P: in
+            # PD mode dropping a store can leave the decoder waiting until the
+            # load timeout instead of failing fast, and rendezvous time makes
+            # store latency a poor pressure signal. Reject it until a
+            # P2P-specific fail-fast path exists.
+            raise ValueError(
+                "Backpressure is not supported for the P2P secondary tier. "
+                "The generic store-latency detector cannot fail fast in PD "
+                "mode and rendezvous time makes store latency an unreliable "
+                "pressure signal. Remove the 'backpressure' config from the "
+                "p2p tier."
+            )
+        super().__init__(
+            offloading_spec,
+            primary_kv_view,
+            tier_type,
+        )
+        try:
+            timeout_s = float(unbound_store_timeout_s)
+        except (TypeError, ValueError):
+            timeout_s = float("nan")
+        if not timeout_s > 0:
+            raise ValueError(
+                f"unbound_store_timeout_s must be a positive number, got "
+                f"{unbound_store_timeout_s!r}"
+            )
+        self._unbound_store_timeout_s = timeout_s
+
         # Block hashes chain from NONE_HASH (see v1/core/kv_cache_utils.py).
         # Peers whose seeds differ compute different hashes for identical
         # content, so lookups silently miss and no KV crosses the wire. The
@@ -303,8 +350,16 @@ def __init__(
         # kv_request_id → list of batches submit_store'd before any peer
         # asked for that id. Drained into a session by _on_session_fetch
         # when the corresponding FetchMsg arrives, or surfaced as failures
-        # by _reap_unbound_stores after _UNBOUND_STORE_TIMEOUT_S.
+        # by _reap_unbound_stores after _unbound_store_timeout_s.
         self._unbound_stores: dict[str, list[_UnboundStoreBatch]] = {}
+        # kv_request_id → time its parked batches were reaped. Nothing under
+        # one of these ids can be served any more, so a FetchMsg for it is
+        # rejected on the spot by _poll_once instead of parking demand until
+        # the consumer's own load timeout, and submit_store fails rather than
+        # re-parking under it — a prefill that outran the timeout would
+        # otherwise keep pinning fresh primary-tier slots, one timeout at a
+        # time. Entries are pruned after _REAPED_ID_RETENTION_S.
+        self._reaped_stores: dict[str, float] = {}
 
         self._finished_jobs: list[JobResult] = []
         # kv_request_ids that hit a transport/session failure; On load lookup()
@@ -390,7 +445,7 @@ def on_request_finished(self, req_context: ReqContext) -> None:
         kv_transfer_params; if a session has bound the id, finish it. If
         no session has bound the id yet, this is a no-op: parked batches
         in `_unbound_stores` are left in place and cleaned up only by
-        `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`.
+        `_reap_unbound_stores` after `_unbound_store_timeout_s`.
         """
         source = req_context.get_state(P2PSourceInfo)
         dest = req_context.get_state(P2PDestInfo)
@@ -447,6 +502,22 @@ def submit_store(self, job_metadata: TransferJob) -> None:
             self._finished_jobs.append(JobResult(job_id=job_id, success=False))
             return
 
+        # This id's earlier batches were already reaped, so the consumer has
+        # either given up or will be rejected by the fetch path. Parking more
+        # would pin primary-tier slots for another full timeout on blocks
+        # nobody can fetch, so fail the job now — a prefill that outran the
+        # timeout falls back to local recompute on the decoder either way.
+        if kv_request_id in self._reaped_stores:
+            logger.warning(
+                "P2P %s: submit_store for reaped kv_request_id=%s job_id=%d "
+                "— failing without parking",
+                self._local_id,
+                kv_request_id,
+                job_id,
+            )
+            self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+            return
+
         # Fast path: a session has already received FetchMsg for this id,
         # so we can route the batch straight into its ServerRole.
         session = self._kv_to_session.get(kv_request_id)
@@ -707,13 +778,27 @@ def _reap_unbound_stores(self) -> None:
         """Time out submit_store batches that no peer has ever fetched.
 
         Walks `_unbound_stores` for entries whose oldest batch is older
-        than `_UNBOUND_STORE_TIMEOUT_S`. Drops the kv_request_id, surfaces
-        every batched job as failed, and adds the id to `_failed_req_ids`
-        so a late inbound FetchMsg short-circuits to a clean rejection.
+        than `_unbound_store_timeout_s`. Drops the kv_request_id, surfaces
+        every batched job as failed, and records the id in `_reaped_stores`
+        so a late inbound FetchMsg is rejected in one round trip and a late
+        `submit_store` fails instead of re-parking under the reaped id.
+
+        Also expires `_reaped_stores` entries whose consumer can no longer
+        be waiting, which is why this runs before the empty check below.
         """
+        now = time.monotonic()
+        if self._reaped_stores:
+            retention_deadline = now - _REAPED_ID_RETENTION_S
+            for kid in [
+                kid
+                for kid, reaped_at in self._reaped_stores.items()
+                if reaped_at <= retention_deadline
+            ]:
+                del self._reaped_stores[kid]
+
         if not self._unbound_stores:
             return
-        deadline = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S
+        deadline = now - self._unbound_store_timeout_s
         expired: list[str] | None = None
         for kid, batches in self._unbound_stores.items():
             # Batches are appended in arrival order, so the head is oldest.
@@ -725,7 +810,7 @@ def _reap_unbound_stores(self) -> None:
             return
         for kid in expired:
             batches = self._unbound_stores.pop(kid)
-            self._failed_req_ids.add(kid)
+            self._reaped_stores[kid] = now
             for batch in batches:
                 self._finished_jobs.append(
                     JobResult(job_id=batch.job_id, success=False)
@@ -735,7 +820,7 @@ def _reap_unbound_stores(self) -> None:
                 "without a fetch — failing %d job(s)",
                 self._local_id,
                 kid,
-                _UNBOUND_STORE_TIMEOUT_S,
+                self._unbound_store_timeout_s,
                 len(batches),
             )
 
@@ -779,6 +864,24 @@ def _poll_once(self) -> None:
             # inline in dispatch, so the replayed add_stored_blocks calls
             # match that demand and submit transfers immediately.
             for kv_request_id in result.new_fetch_ids:
+                if kv_request_id in self._reaped_stores:
+                    # This id's blocks were already reaped, so no
+                    # submit_store will ever satisfy the demand on_fetch
+                    # just recorded. Finalize the round now: the peer gets
+                    # TransferDoneMsg(success=False) on this tick and falls
+                    # back to local prefill, instead of parking until its
+                    # own load timeout expires. The id is deliberately left
+                    # unbound so any later submit_store keeps taking the
+                    # reject path. Batches re-parked between the reap and the
+                    # retention prune are failed here rather than left to a
+                    # second reap, which would pin their slots for another
+                    # full timeout.
+                    session.finish_request(kv_request_id)
+                    for batch in self._unbound_stores.pop(kv_request_id, ()):
+                        self._finished_jobs.append(
+                            JobResult(job_id=batch.job_id, success=False)
+                        )
+                    continue
                 self._kv_to_session[kv_request_id] = session
                 for batch in self._unbound_stores.pop(kv_request_id, ()):
                     session.add_stored_blocks(
@@ -810,6 +913,7 @@ def shutdown(self) -> None:
                     JobResult(job_id=batch.job_id, success=False)
                 )
         self._unbound_stores.clear()
+        self._reaped_stores.clear()
         self._control.close()
         self._data.close()
 
diff --git a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py
index ec4153e25b39..8369bbffe83d 100644
--- a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py
+++ b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-P2P KV cache sharing protocol constants and documentation.
+"""P2P KV cache sharing protocol constants and documentation.
 
 Protocol Overview
 =================
diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py
index ba2aca6438c6..f43e1a5aab5a 100644
--- a/vllm/v1/kv_offload/tiering/p2p/session/session.py
+++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-P2PSession — bidirectional session combining client + server roles.
+"""P2PSession — bidirectional session combining client + server roles.
 
 A single P2PSession per remote peer handles BOTH directions of the P2P
 protocol on one ControlConnection: it can request blocks from the peer
diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py
index 9922a4095c4e..93c6cbabcb85 100644
--- a/vllm/v1/kv_offload/tiering/spec.py
+++ b/vllm/v1/kv_offload/tiering/spec.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-TieringOffloadingSpec: Spec for multi-tier KV cache offloading.
+"""TieringOffloadingSpec: Spec for multi-tier KV cache offloading.
 
 This spec creates a TieringOffloadingManager with a CPU primary tier
 and configurable secondary tiers (e.g., Storage, Network).
@@ -82,8 +81,7 @@
 
 
 class TieringOffloadingSpec(CPUOffloadingSpec):
-    """
-    Spec for multi-tier KV cache offloading.
+    """Spec for multi-tier KV cache offloading.
 
     Creates a TieringOffloadingManager with:
     - Primary tier: CPU (LRU or ARC eviction policy)
@@ -249,6 +247,34 @@ def build_metric_definitions(
             assert isinstance(tier_config, dict)
             tier_cls = SecondaryTierFactory.get_tier_class(tier_config)
             metrics.update(tier_cls.build_metric_definitions(tier_config))
+
+        metrics[TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA] = (
+            OffloadingGaugeMetadata(
+                documentation=(
+                    "Exponential moving average of store latency "
+                    "for back-pressure detection, in s/MiB."
+                ),
+                labelnames=("tier",),
+            )
+        )
+        metrics[TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED] = (
+            OffloadingCounterMetadata(
+                documentation=(
+                    "Number of store operations dropped due to "
+                    "back-pressure on a secondary tier."
+                ),
+                labelnames=("tier",),
+            )
+        )
+        metrics[TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED] = (
+            OffloadingCounterMetadata(
+                documentation=(
+                    "Number of blocks dropped due to back-pressure on a secondary tier."
+                ),
+                labelnames=("tier",),
+            )
+        )
+
         return metrics
 
     def __init__(self, config: OffloadingConfig):
@@ -261,6 +287,32 @@ def __init__(self, config: OffloadingConfig):
         if not isinstance(self.secondary_tier_configs, list):
             raise ValueError("secondary_tiers must be a list of tier configurations")
 
+        # Backpressure config is merged field-by-field in priority order
+        # (highest first):
+        #   1. Per-tier ``backpressure`` dict in the tier config
+        #   2. Top-level ``backpressure`` in kv_connector_extra_config
+        # Merging per field (rather than per whole dict) means a partial
+        # tier override still inherits missing fields from the top-level
+        # default, so the resolved dict reaching the factory is complete.
+        # Within each tier's resolved dict, tier-type-aware water marks
+        # are filled in last so a bare ``"backpressure": {}`` picks up
+        # sensible thresholds for the storage medium.
+        bp_defaults = self.extra_config.get("backpressure")
+
+        for tier_cfg in self.secondary_tier_configs:
+            tier_override = tier_cfg.get("backpressure")
+            # Overlay from lowest to highest precedence so higher-precedence
+            # fields win while lower-precedence ones fill in the gaps.
+            merged: dict[str, Any] = {}
+            for source in (bp_defaults, tier_override):
+                if source:
+                    merged.update(source)
+            # Only set a resolved dict when at least one source contributed
+            # (or an explicit ``backpressure`` key was present, e.g. ``{}``),
+            # so tiers without any backpressure config stay unconfigured.
+            if merged or "backpressure" in tier_cfg:
+                tier_cfg["backpressure"] = merged
+
         # Scheduler-side mmap (rank=None); kept for cleanup
         self._scheduler_mmap: SharedOffloadRegion | None = None
 
@@ -276,8 +328,7 @@ def __init__(self, config: OffloadingConfig):
 
     @override
     def get_manager(self) -> OffloadingManager:
-        """
-        Get the TieringOffloadingManager.
+        """Get the TieringOffloadingManager.
 
         Creates a TieringOffloadingManager with:
         - Primary tier: CPU (LRU or ARC)
@@ -285,6 +336,7 @@ def get_manager(self) -> OffloadingManager:
 
         Returns:
             TieringOffloadingManager instance
+
         """
         if not self._manager:
             if int(self.extra_config.get("store_threshold", 0)) >= 2:
diff --git a/vllm/v1/metrics/buckets.py b/vllm/v1/metrics/buckets.py
new file mode 100644
index 000000000000..7b9efb8f6dc2
--- /dev/null
+++ b/vllm/v1/metrics/buckets.py
@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Default bucket boundaries for the Prometheus histograms emitted by vLLM.
+
+Every core engine histogram created by the Prometheus stat logger draws its
+bucket boundaries from exactly one of the families defined here, so each
+default list has a single source of truth. Histograms owned by other
+subsystems, such as the KV connector and offloading metrics, keep their own
+boundaries.
+"""
+
+from typing import Literal, get_args
+
+BucketFamilyKey = Literal[
+    "request_latency",
+    "time_to_first_token",
+    "inter_token_latency",
+    "iteration_tokens",
+    "request_params_n",
+    "request_num_preemptions",
+    "request_tokens",
+    "kv_cache_residency",
+]
+"""Canonical keys for the histogram bucket families."""
+
+BUCKET_FAMILY_KEYS: frozenset[str] = frozenset(get_args(BucketFamilyKey))
+"""All bucket family keys, for membership checks."""
+
+REQUEST_LATENCY_BUCKETS: tuple[float, ...] = (
+    0.3,
+    0.5,
+    0.8,
+    1.0,
+    1.5,
+    2.0,
+    2.5,
+    5.0,
+    10.0,
+    15.0,
+    20.0,
+    30.0,
+    40.0,
+    50.0,
+    60.0,
+    120.0,
+    240.0,
+    480.0,
+    960.0,
+    1920.0,
+    7680.0,
+)
+"""Sub-second scheduling delays through multi-hour batch requests; shared by
+the accumulated request phase-timing histograms (e2e, queue, inference,
+prefill, decode)."""
+
+TIME_TO_FIRST_TOKEN_BUCKETS: tuple[float, ...] = (
+    0.001,
+    0.005,
+    0.01,
+    0.02,
+    0.04,
+    0.06,
+    0.08,
+    0.1,
+    0.25,
+    0.5,
+    0.75,
+    1.0,
+    2.5,
+    5.0,
+    7.5,
+    10.0,
+    20.0,
+    40.0,
+    80.0,
+    160.0,
+    640.0,
+    2560.0,
+)
+"""Millisecond-scale prefill for tiny prompts up to ~40-minute worst cases,
+with the densest resolution around interactive (sub-second) latencies."""
+
+INTER_TOKEN_LATENCY_BUCKETS: tuple[float, ...] = (
+    0.01,
+    0.025,
+    0.05,
+    0.075,
+    0.1,
+    0.15,
+    0.2,
+    0.3,
+    0.4,
+    0.5,
+    0.75,
+    1.0,
+    2.5,
+    5.0,
+    7.5,
+    10.0,
+    20.0,
+    40.0,
+    80.0,
+)
+"""Per-decode-step latencies: 10 ms fast decode up to multi-second stalls."""
+
+ITERATION_TOKENS_BUCKETS: tuple[float, ...] = (
+    1,
+    8,
+    16,
+    32,
+    64,
+    128,
+    256,
+    512,
+    1024,
+    2048,
+    4096,
+    8192,
+    16384,
+)
+"""Tokens processed per engine step: powers of two up to a typical
+max-num-batched-tokens budget."""
+
+REQUEST_PARAMS_N_BUCKETS: tuple[float, ...] = (1, 2, 5, 10, 20)
+"""Small integer counts for the ``n`` sampling parameter."""
+
+REQUEST_NUM_PREEMPTIONS_BUCKETS: tuple[float, ...] = (1, 2, 3, 4, 5, 10, 20)
+"""Per-request preemption counts, recorded for every finished request: the
+first bucket therefore holds never-preempted requests together with
+singly-preempted ones, 2 through 5 are separated individually, and the tail
+coarsens to 20. Distinct from `REQUEST_PARAMS_N_BUCKETS` because these are the
+boundaries the metric shipped with, not a harmonized choice."""
+
+KV_CACHE_RESIDENCY_BUCKETS: tuple[float, ...] = (
+    0.001,
+    0.002,
+    0.005,
+    0.01,
+    0.02,
+    0.05,
+    0.1,
+    0.2,
+    0.5,
+    1,
+    2,
+    5,
+    10,
+    20,
+    30,
+    60,
+    120,
+    300,
+    600,
+    1200,
+    1800,
+)
+"""KV cache block residency times: millisecond-scale reuse gaps up to
+30-minute block lifetimes."""
+
+_STATIC_FAMILY_DEFAULTS: dict[str, tuple[float, ...]] = {
+    "request_latency": REQUEST_LATENCY_BUCKETS,
+    "time_to_first_token": TIME_TO_FIRST_TOKEN_BUCKETS,
+    "inter_token_latency": INTER_TOKEN_LATENCY_BUCKETS,
+    "iteration_tokens": ITERATION_TOKENS_BUCKETS,
+    "request_params_n": REQUEST_PARAMS_N_BUCKETS,
+    "request_num_preemptions": REQUEST_NUM_PREEMPTIONS_BUCKETS,
+    "kv_cache_residency": KV_CACHE_RESIDENCY_BUCKETS,
+}
+
+
+def build_buckets(mantissa_lst: list[int], max_value: int) -> list[float]:
+    """Build buckets from mantissas scaled by increasing powers of 10.
+
+    Args:
+        mantissa_lst: Mantissa values multiplied by each power of 10.
+        max_value: Largest bucket value to include.
+
+    Returns:
+        Bucket values in increasing order, capped at `max_value`.
+
+    """
+    exponent = 0
+    buckets: list[float] = []
+    while True:
+        for m in mantissa_lst:
+            value = m * 10**exponent
+            if value <= max_value:
+                buckets.append(value)
+            else:
+                return buckets
+        exponent += 1
+
+
+def build_1_2_5_buckets(max_value: int) -> list[float]:
+    """Build a 1-2-5 series of buckets capped at `max_value`.
+
+    Example:
+        >>> build_1_2_5_buckets(100)
+        [1, 2, 5, 10, 20, 50, 100]
+
+    """
+    return build_buckets([1, 2, 5], max_value)
+
+
+def histogram_buckets(
+    family: BucketFamilyKey,
+    max_model_len: int | None = None,
+) -> list[float]:
+    """Return the default bucket boundaries for a histogram family.
+
+    Args:
+        family: Canonical bucket family key.
+        max_model_len: Cap for the token-count series of the
+            `request_tokens` family; required for that family and ignored
+            otherwise.
+
+    Returns:
+        A fresh list of bucket upper bounds; callers may mutate it freely.
+
+    Raises:
+        ValueError: If `family` is `request_tokens` and `max_model_len`
+            is None.
+
+    """
+    if family == "request_tokens":
+        if max_model_len is None:
+            raise ValueError(
+                "max_model_len is required for the 'request_tokens' bucket family"
+            )
+        return build_1_2_5_buckets(max_model_len)
+    return list(_STATIC_FAMILY_DEFAULTS[family])
diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py
index 4d0b77add7f0..26a4142df9f1 100644
--- a/vllm/v1/metrics/loggers.py
+++ b/vllm/v1/metrics/loggers.py
@@ -11,6 +11,10 @@
 import vllm.envs as envs
 from vllm.compilation.cuda_graph import CUDAGraphLogging
 from vllm.config import SupportsMetricsInfo, VllmConfig
+from vllm.distributed.ec_transfer.ec_connector.metrics import (
+    ECConnectorLogging,
+    ECConnectorProm,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
     KVConnectorLogging,
     KVConnectorProm,
@@ -18,6 +22,7 @@
 from vllm.logger import init_logger
 from vllm.plugins import STAT_LOGGER_PLUGINS_GROUP, load_plugins_by_group
 from vllm.v1.engine import FinishReason
+from vllm.v1.metrics.buckets import histogram_buckets
 from vllm.v1.metrics.perf import PerfMetricsLogging, PerfMetricsProm
 from vllm.v1.metrics.prometheus import unregister_vllm_metrics
 from vllm.v1.metrics.stats import (
@@ -100,6 +105,8 @@ class LoggingStatLogger(StatLoggerBase):
     def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
         self.engine_index = engine_index
         self.vllm_config = vllm_config
+        device_type = vllm_config.device_config.device_type
+        self.kv_cache_device = "GPU" if device_type == "cuda" else device_type.upper()
         self._reset(time.monotonic())
 
         self.last_scheduler_stats = SchedulerStats()
@@ -115,6 +122,8 @@ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
         self.spec_decoding_logging = SpecDecodingLogging(is_diffusion=is_diffusion)
         kv_transfer_config = self.vllm_config.kv_transfer_config
         self.kv_connector_logging = KVConnectorLogging(kv_transfer_config)
+        ec_transfer_config = self.vllm_config.ec_transfer_config
+        self.ec_connector_logging = ECConnectorLogging(ec_transfer_config)
         self.cudagraph_logging = None
         if self.vllm_config.observability_config.cudagraph_metrics:
             self.cudagraph_logging = CUDAGraphLogging(
@@ -183,7 +192,7 @@ def _log_iteration_details(
             "%sIteration(%d): %d context requests, %d context tokens, "
             "%d generation requests, %d generation tokens, "
             "iteration elapsed time: %.2f ms%s, "
-            "GPU KV cache usage: %.1f%%%s",
+            "%s KV cache usage: %.1f%%%s",
             self._log_prefix_for_engine(engine_idx),
             details.iteration_index,
             details.num_ctx_requests,
@@ -192,6 +201,7 @@ def _log_iteration_details(
             details.num_generation_tokens,
             details.elapsed_ms,
             " (dummy)" if details.is_dummy else "",
+            self.kv_cache_device,
             scheduler_stats.kv_cache_usage * 100,
             encoder_msg,
         )
@@ -220,6 +230,8 @@ def record(
                 self.spec_decoding_logging.observe(scheduler_stats.spec_decoding_stats)
             if kv_connector_stats := scheduler_stats.kv_connector_stats:
                 self.kv_connector_logging.observe(kv_connector_stats)
+            if ec_connector_stats := scheduler_stats.ec_connector_stats:
+                self.ec_connector_logging.observe(ec_connector_stats)
             if (
                 self.cudagraph_logging is not None
                 and scheduler_stats.cudagraph_stats is not None
@@ -284,18 +296,16 @@ def log(self):
             log_parts.append("Preemptions: %d")
             log_args.append(self.num_preemptions)
 
-        log_parts.extend(
-            [
-                "GPU KV cache usage: %.1f%%",
-                "Prefix cache hit rate: %.1f%%",
-            ]
-        )
+        log_parts.append("%s KV cache usage: %.1f%%")
         log_args.extend(
             [
+                self.kv_cache_device,
                 self.last_scheduler_stats.kv_cache_usage * 100,
-                self.prefix_caching_metrics.hit_rate * 100,
             ]
         )
+        if not self.prefix_caching_metrics.empty:
+            log_parts.append("Prefix cache hit rate: %.1f%%")
+            log_args.append(self.prefix_caching_metrics.hit_rate * 100)
 
         if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
             log_parts.append("Corrupted: %d reqs")
@@ -314,6 +324,7 @@ def log(self):
 
         self.spec_decoding_logging.log(log_fn=log_fn)
         self.kv_connector_logging.log(log_fn=log_fn)
+        self.ec_connector_logging.log(log_fn=log_fn)
         if self.cudagraph_logging is not None:
             self.cudagraph_logging.log(log_fn=log_fn)
         if self._enable_perf_stats():
@@ -446,6 +457,7 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
     _histogram_cls = Histogram
     _spec_decoding_cls = SpecDecodingProm
     _kv_connector_cls = KVConnectorProm
+    _ec_connector_cls = ECConnectorProm
     _perf_metrics_cls = PerfMetricsProm
 
     def __init__(
@@ -483,6 +495,9 @@ def __init__(
         self.kv_connector_prom = self._kv_connector_cls(
             vllm_config, labelnames, per_engine_labelvalues
         )
+        self.ec_connector_prom = self._ec_connector_cls(
+            vllm_config, labelnames, per_engine_labelvalues
+        )
         self.perf_metrics_prom = self._perf_metrics_cls(
             vllm_config, labelnames, per_engine_labelvalues
         )
@@ -724,13 +739,23 @@ def __init__(
                 for idx in engine_indexes
             }
 
+        request_tokens_buckets = histogram_buckets(
+            "request_tokens", max_model_len=max_model_len
+        )
+        iteration_tokens_buckets = histogram_buckets("iteration_tokens")
+        request_params_n_buckets = histogram_buckets("request_params_n")
+        time_to_first_token_buckets = histogram_buckets("time_to_first_token")
+        inter_token_latency_buckets = histogram_buckets("inter_token_latency")
+        request_latency_buckets = histogram_buckets("request_latency")
+        request_num_preemptions_buckets = histogram_buckets("request_num_preemptions")
+
         #
         # Histograms of counts
         #
         histogram_num_prompt_tokens_request = self._histogram_cls(
             name="vllm:request_prompt_tokens",
             documentation="Number of prefill tokens processed.",
-            buckets=build_1_2_5_buckets(max_model_len),
+            buckets=request_tokens_buckets,
             labelnames=labelnames,
         )
         self.histogram_num_prompt_tokens_request = create_metric_per_engine(
@@ -740,7 +765,7 @@ def __init__(
         histogram_num_generation_tokens_request = self._histogram_cls(
             name="vllm:request_generation_tokens",
             documentation="Number of generation tokens processed.",
-            buckets=build_1_2_5_buckets(max_model_len),
+            buckets=request_tokens_buckets,
             labelnames=labelnames,
         )
         self.histogram_num_generation_tokens_request = create_metric_per_engine(
@@ -753,7 +778,7 @@ def __init__(
         histogram_iteration_tokens = self._histogram_cls(
             name="vllm:iteration_tokens_total",
             documentation="Histogram of number of tokens per engine_step.",
-            buckets=[1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384],
+            buckets=iteration_tokens_buckets,
             labelnames=labelnames,
         )
         self.histogram_iteration_tokens = create_metric_per_engine(
@@ -763,7 +788,7 @@ def __init__(
         histogram_max_num_generation_tokens_request = self._histogram_cls(
             name="vllm:request_max_num_generation_tokens",
             documentation="Histogram of maximum number of requested generation tokens.",
-            buckets=build_1_2_5_buckets(max_model_len),
+            buckets=request_tokens_buckets,
             labelnames=labelnames,
         )
         self.histogram_max_num_generation_tokens_request = create_metric_per_engine(
@@ -773,7 +798,7 @@ def __init__(
         histogram_n_request = self._histogram_cls(
             name="vllm:request_params_n",
             documentation="Histogram of the n request parameter.",
-            buckets=[1, 2, 5, 10, 20],
+            buckets=request_params_n_buckets,
             labelnames=labelnames,
         )
         self.histogram_n_request = create_metric_per_engine(
@@ -783,7 +808,7 @@ def __init__(
         histogram_max_tokens_request = self._histogram_cls(
             name="vllm:request_params_max_tokens",
             documentation="Histogram of the max_tokens request parameter.",
-            buckets=build_1_2_5_buckets(max_model_len),
+            buckets=request_tokens_buckets,
             labelnames=labelnames,
         )
         self.histogram_max_tokens_request = create_metric_per_engine(
@@ -796,30 +821,7 @@ def __init__(
         histogram_time_to_first_token = self._histogram_cls(
             name="vllm:time_to_first_token_seconds",
             documentation="Histogram of time to first token in seconds.",
-            buckets=[
-                0.001,
-                0.005,
-                0.01,
-                0.02,
-                0.04,
-                0.06,
-                0.08,
-                0.1,
-                0.25,
-                0.5,
-                0.75,
-                1.0,
-                2.5,
-                5.0,
-                7.5,
-                10.0,
-                20.0,
-                40.0,
-                80.0,
-                160.0,
-                640.0,
-                2560.0,
-            ],
+            buckets=time_to_first_token_buckets,
             labelnames=labelnames,
         )
         self.histogram_time_to_first_token = create_metric_per_engine(
@@ -829,27 +831,7 @@ def __init__(
         histogram_inter_token_latency = self._histogram_cls(
             name="vllm:inter_token_latency_seconds",
             documentation="Histogram of inter-token latency in seconds.",
-            buckets=[
-                0.01,
-                0.025,
-                0.05,
-                0.075,
-                0.1,
-                0.15,
-                0.2,
-                0.3,
-                0.4,
-                0.5,
-                0.75,
-                1.0,
-                2.5,
-                5.0,
-                7.5,
-                10.0,
-                20.0,
-                40.0,
-                80.0,
-            ],
+            buckets=inter_token_latency_buckets,
             labelnames=labelnames,
         )
         self.histogram_inter_token_latency = create_metric_per_engine(
@@ -859,56 +841,13 @@ def __init__(
         histogram_request_time_per_output_token = self._histogram_cls(
             name="vllm:request_time_per_output_token_seconds",
             documentation="Histogram of time_per_output_token_seconds per request.",
-            buckets=[
-                0.01,
-                0.025,
-                0.05,
-                0.075,
-                0.1,
-                0.15,
-                0.2,
-                0.3,
-                0.4,
-                0.5,
-                0.75,
-                1.0,
-                2.5,
-                5.0,
-                7.5,
-                10.0,
-                20.0,
-                40.0,
-                80.0,
-            ],
+            buckets=inter_token_latency_buckets,
             labelnames=labelnames,
         )
         self.histogram_request_time_per_output_token = create_metric_per_engine(
             histogram_request_time_per_output_token, per_engine_labelvalues
         )
 
-        request_latency_buckets = [
-            0.3,
-            0.5,
-            0.8,
-            1.0,
-            1.5,
-            2.0,
-            2.5,
-            5.0,
-            10.0,
-            15.0,
-            20.0,
-            30.0,
-            40.0,
-            50.0,
-            60.0,
-            120.0,
-            240.0,
-            480.0,
-            960.0,
-            1920.0,
-            7680.0,
-        ]
         histogram_e2e_time_request = self._histogram_cls(
             name="vllm:e2e_request_latency_seconds",
             documentation="Histogram of e2e request latency in seconds.",
@@ -962,7 +901,7 @@ def __init__(
         histogram_request_num_preemptions = self._histogram_cls(
             name="vllm:request_num_preemptions",
             documentation="Histogram of the number of times a request was preempted.",
-            buckets=[1, 2, 3, 4, 5, 10, 20],
+            buckets=request_num_preemptions_buckets,
             labelnames=labelnames,
         )
         self.histogram_request_num_preemptions = create_metric_per_engine(
@@ -975,7 +914,7 @@ def __init__(
                 "Histogram of new KV tokens computed during prefill "
                 "(excluding cached tokens)."
             ),
-            buckets=build_1_2_5_buckets(max_model_len),
+            buckets=request_tokens_buckets,
             labelnames=labelnames,
         )
         self.histogram_prefill_kv_computed_request = create_metric_per_engine(
@@ -986,29 +925,7 @@ def __init__(
         # KV Cache residency metrics
         #
         if self.kv_cache_metrics_enabled:
-            kv_cache_residency_buckets = [
-                0.001,
-                0.002,
-                0.005,
-                0.01,
-                0.02,
-                0.05,
-                0.1,
-                0.2,
-                0.5,
-                1,
-                2,
-                5,
-                10,
-                20,
-                30,
-                60,
-                120,
-                300,
-                600,
-                1200,
-                1800,
-            ]
+            kv_cache_residency_buckets = histogram_buckets("kv_cache_residency")
 
             histogram_kv_block_lifetime = self._histogram_cls(
                 name="vllm:kv_block_lifetime_seconds",
@@ -1157,6 +1074,11 @@ def record(
                     scheduler_stats.kv_connector_stats, engine_idx
                 )
 
+            if scheduler_stats.ec_connector_stats is not None:
+                self.ec_connector_prom.observe(
+                    scheduler_stats.ec_connector_stats, engine_idx
+                )
+
             if scheduler_stats.perf_stats is not None:
                 self.perf_metrics_prom.observe(scheduler_stats.perf_stats, engine_idx)
 
@@ -1294,44 +1216,16 @@ def log_engine_initialized(self):
         self.log_metrics_info("cache_config", self.vllm_config.cache_config)
 
 
-def build_buckets(mantissa_lst: list[int], max_value: int) -> list[int]:
-    """
-    Builds a list of buckets with increasing powers of 10 multiplied by
-    mantissa values until the value exceeds the specified maximum.
-
-    """
-    exponent = 0
-    buckets: list[int] = []
-    while True:
-        for m in mantissa_lst:
-            value = m * 10**exponent
-            if value <= max_value:
-                buckets.append(value)
-            else:
-                return buckets
-        exponent += 1
-
-
-def build_1_2_5_buckets(max_value: int) -> list[int]:
-    """
-    Example:
-    >>> build_1_2_5_buckets(100)
-    [1, 2, 5, 10, 20, 50, 100]
-    """
-    return build_buckets([1, 2, 5], max_value)
-
-
 class StatLoggerManager:
-    """
-    StatLoggerManager:
-        Logging happens at the level of the EngineCore (per scheduler).
-         * DP: >1 EngineCore per AsyncLLM - loggers for each EngineCore.
-         * With Local Logger, just make N copies for N EngineCores.
-         * With Prometheus, we need a single logger with N "labels"
-
-        This class abstracts away this implementation detail from
-        the AsyncLLM, allowing the AsyncLLM to just call .record()
-        and .log() to a simple interface.
+    """StatLoggerManager:
+    Logging happens at the level of the EngineCore (per scheduler).
+     * DP: >1 EngineCore per AsyncLLM - loggers for each EngineCore.
+     * With Local Logger, just make N copies for N EngineCores.
+     * With Prometheus, we need a single logger with N "labels"
+
+    This class abstracts away this implementation detail from
+    the AsyncLLM, allowing the AsyncLLM to just call .record()
+    and .log() to a simple interface.
     """
 
     def __init__(
diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py
index 75524ef12858..1cd8fa79f1d3 100644
--- a/vllm/v1/metrics/perf.py
+++ b/vllm/v1/metrics/perf.py
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""
-Analytic flops/memory estimation module for transformer components,
+"""Analytic flops/memory estimation module for transformer components,
 to help derive MFU (Model Flops Utilization) stats for a running model.
 """
 
@@ -10,7 +9,7 @@
 import time
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 from typing import Any, Protocol
 
 import prometheus_client
@@ -33,8 +32,7 @@
 
 
 class InvalidComponent(Exception):
-    """
-    Custom exception to indicate that a certain ComponentMetric is not
+    """Custom exception to indicate that a certain ComponentMetric is not
     applicable to the given VllmConfig.
     """
 
@@ -99,8 +97,7 @@ class PerfStats:
 
 @dataclass
 class ExecutionContext:
-    """
-    Represents an execution context for a batch of requests.
+    """Represents an execution context for a batch of requests.
 
     This class aggregates statistics across multiple requests in a batch,
     separately tracking prefill and decode phases.
@@ -124,8 +121,12 @@ class ExecutionContext:
     decode_context_len: int = 0  # sum of context_len for decode requests
     decode_token_context_product: int = 0  # sum of (num_tokens * context_len)
 
+    # Per-request records: (num_tokens, context_len, is_prefill)
+    requests: list[tuple[int, int, bool]] = field(default_factory=list)
+
     def add(self, num_tokens: int, context_len: int, is_prefill: bool) -> None:
         """Add a single request's statistics to this batch context."""
+        self.requests.append((num_tokens, context_len, is_prefill))
         if is_prefill:
             self.num_prefill_requests += 1
             self.prefill_num_tokens += num_tokens
@@ -141,9 +142,40 @@ def total_num_tokens(self) -> int:
         """Total number of tokens across all requests in the batch."""
         return self.prefill_num_tokens + self.decode_num_tokens
 
-    def total_token_context_product(self) -> int:
-        """Total sum of (num_tokens * context_len) across all requests."""
-        return self.prefill_token_context_product + self.decode_token_context_product
+    def total_token_context_product(self, sliding_window: int | None = None) -> int:
+        """Total sum of (num_tokens * context_len) across all requests.
+
+        If sliding_window is provided, context length per request is clamped to
+        sliding_window.
+        """
+        if sliding_window is None:
+            return (
+                self.prefill_token_context_product + self.decode_token_context_product
+            )
+        return sum(
+            num_tokens * min(context_len, sliding_window)
+            for num_tokens, context_len, _ in self.requests
+        )
+
+    def prefill_context_len_for_window(self, sliding_window: int | None = None) -> int:
+        """Sum of context_len for prefill requests, clamped to sliding_window if set."""
+        if sliding_window is None:
+            return self.prefill_context_len
+        return sum(
+            min(context_len, sliding_window)
+            for _, context_len, is_prefill in self.requests
+            if is_prefill
+        )
+
+    def decode_context_len_for_window(self, sliding_window: int | None = None) -> int:
+        """Sum of context_len for decode requests, clamped to sliding_window if set."""
+        if sliding_window is None:
+            return self.decode_context_len
+        return sum(
+            min(context_len, sliding_window)
+            for _, context_len, is_prefill in self.requests
+            if not is_prefill
+        )
 
     def num_logits_tokens(self) -> int:
         """Number of tokens that require logits computation (unembedding).
@@ -167,8 +199,7 @@ def from_single_request(
 
 
 class ParsedArgs:
-    """
-    Syntactic sugar so that Parsers can use dot notations
+    """Syntactic sugar so that Parsers can use dot notations
     to access/update the parsed arguments.
 
     e.g.)
@@ -192,16 +223,14 @@ def model_dump(self) -> dict[str, Any]:
 
 class Parser(Protocol):
     def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
-        """
-        Parse the vllm config and update the current ParsedArgs and pass it on.
+        """Parse the vllm config and update the current ParsedArgs and pass it on.
         If the parser isn't applicable to the vllm_config, it will do nothing.
         """
         ...
 
 
 class ParserChain:
-    """
-    Applies chain of parser in a sequential order.
+    """Applies chain of parser in a sequential order.
     Later parsers might overwrite results from previous parsers,
     so parsers should be chained in the appropriate order if they
     are not mutually exclusive.
@@ -214,6 +243,7 @@ def add_parser(self, parser: Parser) -> None:
         self.parsers.append(parser)
 
     def parse(self, vllm_config: VllmConfig) -> ParsedArgs:
+        """Apply all parsers in sequence and return the accumulated ParsedArgs."""
         args = ParsedArgs()
         for parser in self.parsers:
             args = parser.parse(args, vllm_config)
@@ -224,8 +254,7 @@ def parse(self, vllm_config: VllmConfig) -> ParsedArgs:
 
 
 class ComponentMetrics(BaseModel, ABC):
-    """
-    Each concrete ComponentMetrics class is associated with:
+    """Each concrete ComponentMetrics class is associated with:
     - fields that are required for metric derivation
       (fields are specified/validated through pydantic model)
     - parser to parse VllmConfig into fields
@@ -234,13 +263,14 @@ class ComponentMetrics(BaseModel, ABC):
 
     @classmethod
     @abstractmethod
-    def component_type(cls) -> str: ...
+    def component_type(cls) -> str:
+        """Return the unique string key identifying this component in the registry."""
+        ...
 
     @classmethod
     @abstractmethod
     def get_parser(cls) -> ParserChain:
-        """
-        Return a ParserChain that provides values for all required fields.
+        """Return a ParserChain that provides values for all required fields.
         The returned parser chain must populate ParsedArgs with values for every
         field defined on this ComponentMetrics class. Missing fields will cause
         a ValidationError when from_vllm_config() is called.
@@ -254,11 +284,9 @@ def __init_subclass__(cls):
 
     @classmethod
     def from_vllm_config(cls, vllm_config: VllmConfig) -> Self:
-        """
-        Instantiate this class from VllmConfig.
+        """Instantiate this class from VllmConfig.
         Raises ValidationError if parsing fails.
         """
-
         parser = cls.get_parser()
         parsed_args = parser.parse(vllm_config)
         try:
@@ -273,17 +301,23 @@ def registered_metrics(cls) -> Iterable[type["ComponentMetrics"]]:
     @abstractmethod
     def get_num_flops_breakdown(
         self, ctx: ExecutionContext, per_gpu: bool = True
-    ) -> dict[str, int]: ...
+    ) -> dict[str, int]:
+        """Return per-operation FLOPs breakdown for this component."""
+        ...
 
     @abstractmethod
     def get_read_bytes_breakdown(
         self, ctx: ExecutionContext, per_gpu: bool = True
-    ) -> dict[str, int]: ...
+    ) -> dict[str, int]:
+        """Return per-operation read-memory-traffic breakdown for this component."""
+        ...
 
     @abstractmethod
     def get_write_bytes_breakdown(
         self, ctx: ExecutionContext, per_gpu: bool = True
-    ) -> dict[str, int]: ...
+    ) -> dict[str, int]:
+        """Return per-operation write-memory-traffic breakdown for this component."""
+        ...
 
     def get_num_flops(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
         return sum(self.get_num_flops_breakdown(ctx, per_gpu).values())
@@ -299,13 +333,13 @@ def get_write_bytes(self, ctx: ExecutionContext, per_gpu: bool = True) -> int:
 
 
 class BaseConfigParser(Parser):
-    """
-    Parses base model configuration.
+    """Parses base model configuration.
     Provides: vocab_size, hidden_size, num_attention_heads, num_hidden_layers,
     weight_byte_size, activation_byte_size, dp_size, tp_size, pp_size, enable_ep
     """
 
     def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
+        """Parse base model dimensions, dtype byte sizes, and parallelism config."""
         model_config = vllm_config.model_config
 
         args.vocab_size = model_config.get_vocab_size()
@@ -335,9 +369,10 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
 
         args.weight_byte_size = get_dtype_size(torch_dtype)
 
-        # FIXME: handle this better by parsing whether activations use
-        # bf16, fp32, etc...
-        args.activation_byte_size = 2
+        # Activations are produced in the model's compute dtype. Quantization
+        # overrides weight_byte_size below but leaves activations alone, so
+        # this stays keyed on the model dtype.
+        args.activation_byte_size = get_dtype_size(torch_dtype)
 
         args.dp_size = vllm_config.parallel_config.data_parallel_size
         args.tp_size = vllm_config.parallel_config.tensor_parallel_size
@@ -351,12 +386,12 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
 
 
 class BaseAttentionConfigParser(Parser):
-    """
-    Parses attention-specific configuration.
+    """Parses attention-specific configuration.
     Provides: num_key_value_heads, head_dim, cache_byte_size
     """
 
     def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
+        """Parse KV head count, head dimension, and KV-cache dtype byte size."""
         model_config = vllm_config.model_config
 
         args.num_key_value_heads = model_config.get_total_num_kv_heads()
@@ -371,13 +406,60 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
         return args
 
 
-class AttentionQuantizationConfigParser(Parser):
+class SlidingWindowAttentionParser(Parser):
+    """Parses sliding window attention configuration and layer breakdown.
+    Provides: sliding_window, num_swa_layers
     """
-    Parses quantization configuration for attention layers.
+
+    def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
+        """Determine sliding_window size and count of SWA layers.
+
+        Detects per-layer type lists (e.g. Llama-4), Gemma 2 alternating
+        pattern, and uniform SWA models (Mistral, Qwen2.5).
+        """
+        model_config = vllm_config.model_config
+
+        sliding_window: int | None = None
+        if not getattr(model_config, "disable_sliding_window", False):
+            sliding_window = model_config.get_sliding_window()
+            if sliding_window is None and hasattr(vllm_config, "cache_config"):
+                sliding_window = getattr(
+                    vllm_config.cache_config, "sliding_window", None
+                )
+
+        args.sliding_window = sliding_window
+
+        num_layers = args.num_hidden_layers
+        if sliding_window is not None and sliding_window > 0:
+            hf_text_config = getattr(model_config, "hf_text_config", None)
+            layer_types = getattr(hf_text_config, "layer_types", None)
+            if layer_types is not None and isinstance(layer_types, list):
+                num_swa = sum(
+                    1 for lt in layer_types[:num_layers] if "sliding" in lt.lower()
+                )
+                args.num_swa_layers = num_swa
+            elif (
+                getattr(getattr(model_config, "hf_config", None), "model_type", None)
+                == "gemma2"
+            ):
+                # Gemma 2 alternates sliding window: even layers (0, 2, ...) are sliding
+                args.num_swa_layers = (num_layers + 1) // 2
+            else:
+                # Full SWA model (e.g. Mistral, Qwen2.5 with sliding_window)
+                args.num_swa_layers = num_layers
+        else:
+            args.num_swa_layers = 0
+
+        return args
+
+
+class AttentionQuantizationConfigParser(Parser):
+    """Parses quantization configuration for attention layers.
     Overrides: weight_byte_size
     """
 
     def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
+        """Override weight_byte_size based on the active quantization method."""
         cfg = vllm_config.quant_config
 
         if cfg is None:
@@ -395,12 +477,12 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
 
 
 class AttentionDetectionParser(Parser):
-    """
-    Prevents standard AttentionMetrics from being instantiated for MLA models.
+    """Prevents standard AttentionMetrics from being instantiated for MLA models.
     MLA models should use MLAAttentionMetrics instead.
     """
 
     def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
+        """Raise InvalidComponent for MLA models so AttentionMetrics is skipped."""
         if vllm_config.model_config.is_deepseek_mla:
             raise InvalidComponent(
                 "Model uses MLA attention; use MLAAttentionMetrics instead"
@@ -425,53 +507,131 @@ class AttentionMetrics(ComponentMetrics):
     # From BaseConfig Parser, overridden by AttentionQuantizationConfigParser
     weight_byte_size: int | float = Field(..., gt=0)
 
-    # TODO: discern cases where we have mixture of different attention layer types
-    # such as SWA, MLA, etc.
+    # From SlidingWindowAttentionParser
+    sliding_window: int | None = Field(None)
+    num_swa_layers: int = Field(0, ge=0)
 
     @classmethod
     def component_type(cls) -> str:
+        """Return the component registry key for standard (non-MLA) attention."""
         return "attn"
 
     @classmethod
     def get_parser(cls) -> ParserChain:
+        """Return the parser chain for AttentionMetrics.
+
+        Parsers run in order: MLA detection guard, base model config,
+        attention-specific config, sliding-window breakdown, quantization
+        weight-size override.
+        """
         return ParserChain(
             AttentionDetectionParser(),
             BaseConfigParser(),
             BaseAttentionConfigParser(),
+            SlidingWindowAttentionParser(),
             AttentionQuantizationConfigParser(),
         )
 
+    def _layer_counts(self, per_gpu: bool) -> tuple[float, float, float]:
+        """Return ``(L, L_full, L_swa)``, the layer counts to bill for.
+
+        Per-GPU counts divide each global count by ``pp_size`` independently.
+        Deriving the sublayer split after reducing ``L`` instead rounds a
+        minority sublayer away: a 6-layer model with 1 sliding-window layer at
+        ``pp_size=2`` yields ``round(3 * 1 / 6) == 0``, billing all three layers
+        as full attention. Dividing each count separately keeps
+        ``L_full + L_swa == L`` and keeps the per-GPU values summing back to the
+        whole model across stages.
+
+        The counts are fractional because these metrics are built from config
+        alone (see ``from_vllm_config``); no pipeline rank is available, so a
+        per-GPU figure is the average stage rather than a specific one. A model
+        whose layers do not divide evenly across stages has no single integer
+        answer here.
+        """
+        L = float(self.num_hidden_layers)
+        L_swa = float(self.num_swa_layers)
+        L_full = L - L_swa
+
+        if per_gpu:
+            L /= self.pp_size
+            L_swa /= self.pp_size
+            L_full /= self.pp_size
+
+        return L, L_full, L_swa
+
     def get_num_flops_breakdown(
         self, ctx: ExecutionContext, per_gpu: bool = True
     ) -> dict[str, int]:
-        L, D, q, kv, d = (
-            self.num_hidden_layers,
+        """Compute FLOPs breakdown for attention layers.
+
+        Accounts for hybrid models by splitting total layers into full-attention
+        (L_full) and sliding-window (L_swa) sublayers. See ``_layer_counts`` for
+        how those are scaled under pipeline parallelism.
+
+        Args:
+            ctx: Execution context describing the current batch.
+            per_gpu: If True, scale counts down by tensor/pipeline parallelism.
+
+        Returns:
+            Dict with keys: qkv_proj, out_proj, attn_qk, attn_av.
+
+        """
+        D, q, kv, d = (
             self.hidden_size,
             self.num_attention_heads,
             self.num_key_value_heads,
             self.head_dim,
         )
         T = ctx.total_num_tokens()
-        TC = ctx.total_token_context_product()
+        TC_full = ctx.total_token_context_product()
+
+        L, L_full, L_swa = self._layer_counts(per_gpu)
 
         if per_gpu:
-            L //= self.pp_size
             # tensor parallel along heads
             q = max(1, q // self.tp_size)
             kv = max(1, kv // self.tp_size)
 
-        return {
-            "qkv_proj": 2 * T * D * (q + 2 * kv) * d * L,
-            "attn_qk": 2 * q * TC * d * L,
-            "attn_av": 2 * q * TC * d * L,
-            "out_proj": 2 * T * D * q * d * L,
+        flops = {
+            "qkv_proj": int(2 * T * D * (q + 2 * kv) * d * L),
+            "out_proj": int(2 * T * D * q * d * L),
         }
 
+        attn_qk = 0.0
+        attn_av = 0.0
+        if L_full > 0:
+            attn_qk += 2 * q * TC_full * d * L_full
+            attn_av += 2 * q * TC_full * d * L_full
+
+        if L_swa > 0 and self.sliding_window is not None:
+            TC_swa = ctx.total_token_context_product(self.sliding_window)
+            attn_qk += 2 * q * TC_swa * d * L_swa
+            attn_av += 2 * q * TC_swa * d * L_swa
+
+        flops["attn_qk"] = int(attn_qk)
+        flops["attn_av"] = int(attn_av)
+
+        return flops
+
     def get_read_bytes_breakdown(
         self, ctx: ExecutionContext, per_gpu: bool = True
     ) -> dict[str, int]:
-        L, D, q, kv, d = (
-            self.num_hidden_layers,
+        """Compute read-memory-traffic breakdown for attention layers.
+
+        Splits layers into full-attention and sliding-window sublayers using the
+        same counts as get_num_flops_breakdown; see ``_layer_counts``.
+
+        Args:
+            ctx: Execution context describing the current batch.
+            per_gpu: If True, scale counts down by tensor/pipeline parallelism.
+
+        Returns:
+            Dict with keys: qkv_input, qkv_weight, attn_input (conditional),
+            out_input, out_weight.
+
+        """
+        D, q, kv, d = (
             self.hidden_size,
             self.num_attention_heads,
             self.num_key_value_heads,
@@ -479,35 +639,63 @@ def get_read_bytes_breakdown(
         )
         T = ctx.total_num_tokens()
 
+        L, L_full, L_swa = self._layer_counts(per_gpu)
+
         if per_gpu:
-            L //= self.pp_size
             # tensor parallel along heads
             q = max(1, q // self.tp_size)
             kv = max(1, kv // self.tp_size)
 
         read_bytes = {}
 
-        read_bytes["qkv_input"] = T * D * self.activation_byte_size * L
+        read_bytes["qkv_input"] = int(T * D * self.activation_byte_size * L)
         read_bytes["qkv_weight"] = int(D * (q + 2 * kv) * d * self.weight_byte_size * L)
 
-        # Attention input reads differ between prefill and decode
-        # Prefill: read Q, K, V activations (all in activation_byte_size)
-        if ctx.prefill_num_tokens > 0:
-            read_bytes["attn_input"] = (
-                (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv)
-                * d
-                * self.activation_byte_size
-                * L
-            )
+        attn_input = 0.0
 
-        # Decode: read Q activations + read K, V from cache (in cache_byte_size)
-        if ctx.decode_num_tokens > 0:
-            read_bytes["attn_input"] = read_bytes.get("attn_input", 0) + (
-                ctx.decode_num_tokens * q * d * self.activation_byte_size * L
-                + 2 * ctx.decode_context_len * kv * d * self.cache_byte_size * L
-            )
+        # Full attention layers
+        if L_full > 0:
+            if ctx.prefill_num_tokens > 0:
+                attn_input += (
+                    (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv)
+                    * d
+                    * self.activation_byte_size
+                    * L_full
+                )
+            if ctx.decode_num_tokens > 0:
+                attn_input += (
+                    ctx.decode_num_tokens * q * d * self.activation_byte_size * L_full
+                    + 2
+                    * ctx.decode_context_len
+                    * kv
+                    * d
+                    * self.cache_byte_size
+                    * L_full
+                )
 
-        read_bytes["out_input"] = T * q * d * self.activation_byte_size * L
+        # Sliding window attention layers
+        if L_swa > 0:
+            sw = self.sliding_window
+            prefill_sw_len = ctx.prefill_context_len_for_window(sw)
+            decode_sw_len = ctx.decode_context_len_for_window(sw)
+
+            if ctx.prefill_num_tokens > 0:
+                attn_input += (
+                    (ctx.prefill_num_tokens * q + 2 * prefill_sw_len * kv)
+                    * d
+                    * self.activation_byte_size
+                    * L_swa
+                )
+            if ctx.decode_num_tokens > 0:
+                attn_input += (
+                    ctx.decode_num_tokens * q * d * self.activation_byte_size * L_swa
+                    + 2 * decode_sw_len * kv * d * self.cache_byte_size * L_swa
+                )
+
+        if ctx.prefill_num_tokens > 0 or ctx.decode_num_tokens > 0:
+            read_bytes["attn_input"] = int(attn_input)
+
+        read_bytes["out_input"] = int(T * q * d * self.activation_byte_size * L)
         read_bytes["out_weight"] = int(q * d * D * self.weight_byte_size * L)
 
         return read_bytes
@@ -515,9 +703,13 @@ def get_read_bytes_breakdown(
     def get_write_bytes_breakdown(
         self, ctx: ExecutionContext, per_gpu: bool = True
     ) -> dict[str, int]:
-        """Calculate write memory traffic for attention layers."""
-        L, D, q, kv, d = (
-            self.num_hidden_layers,
+        """Calculate write memory traffic for attention layers.
+
+        Writes do not depend on the attention window, so only the total layer
+        count is needed; it comes from ``_layer_counts`` so all three breakdowns
+        bill the same number of layers.
+        """
+        D, q, kv, d = (
             self.hidden_size,
             self.num_attention_heads,
             self.num_key_value_heads,
@@ -525,16 +717,17 @@ def get_write_bytes_breakdown(
         )
         T = ctx.total_num_tokens()
 
+        L, _, _ = self._layer_counts(per_gpu)
+
         if per_gpu:
-            L //= self.pp_size
             # tensor parallel along heads
             q = max(1, q // self.tp_size)
             kv = max(1, kv // self.tp_size)
 
         return {
-            "qkv_output": T * (q + 2 * kv) * d * self.activation_byte_size * L,
-            "kv_cache": 2 * T * kv * d * self.cache_byte_size * L,
-            "out_output": T * D * self.activation_byte_size * L,
+            "qkv_output": int(T * (q + 2 * kv) * d * self.activation_byte_size * L),
+            "kv_cache": int(2 * T * kv * d * self.cache_byte_size * L),
+            "out_output": int(T * D * self.activation_byte_size * L),
         }
 
 
@@ -542,26 +735,26 @@ def get_write_bytes_breakdown(
 
 
 class MLADetectionParser(Parser):
-    """
-    Validates that the model uses MLA attention.
+    """Validates that the model uses MLA attention.
     Raises InvalidComponent if the model does not use MLA,
     so MLAAttentionMetrics is silently skipped for non-MLA models.
     """
 
     def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
+        """Raise InvalidComponent if the model does not use MLA attention."""
         if not vllm_config.model_config.is_deepseek_mla:
             raise InvalidComponent("Model does not use MLA attention")
         return args
 
 
 class MLAConfigParser(Parser):
-    """
-    Parses MLA-specific configuration fields.
+    """Parses MLA-specific configuration fields.
     Provides: kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim,
     v_head_dim, q_lora_rank
     """
 
     def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
+        """Parse MLA-specific compression dimensions and KV-cache dtype byte size."""
         model_config = vllm_config.model_config
         cfg = model_config.hf_text_config
 
@@ -580,8 +773,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
 
 
 class MLAAttentionMetrics(ComponentMetrics):
-    """
-    Performance metrics for Multi-Latent Attention (MLA) layers.
+    """Performance metrics for Multi-Latent Attention (MLA) layers.
 
     MLA uses a compressed latent representation for KV cache:
     - KV cache stores a single compressed vector of size
@@ -616,10 +808,16 @@ class MLAAttentionMetrics(ComponentMetrics):
 
     @classmethod
     def component_type(cls) -> str:
+        """Return the component registry key for MLA attention."""
         return "mla_attn"
 
     @classmethod
     def get_parser(cls) -> ParserChain:
+        """Return the parser chain for MLAAttentionMetrics.
+
+        Parsers run in order: MLA detection guard, base model config,
+        MLA-specific config, quantization weight-size override.
+        """
         return ParserChain(
             MLADetectionParser(),
             BaseConfigParser(),
@@ -812,8 +1010,7 @@ def get_write_bytes_breakdown(
 
 
 class BaseFfnConfigParser(Parser):
-    """
-    Parses FFN and MoE configuration.
+    """Parses FFN and MoE configuration.
     Provides: intermediate_size, num_experts, num_experts_per_tok,
     moe_intermediate_size, num_shared_experts, num_moe_layers
     """
@@ -843,8 +1040,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
 
 
 class FfnParallelParser(Parser):
-    """
-    Parses FFN parallelism configuration.
+    """Parses FFN parallelism configuration.
 
     Provides: ffn_tp_size, ffn_ep_size
     """
@@ -864,8 +1060,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
 
 
 class InterleaveMoeLayerStepParser(Parser):
-    """
-    Parses interleave_moe_layer_step field for models like Llama4.
+    """Parses interleave_moe_layer_step field for models like Llama4.
 
     Overrides: num_moe_layers
     """
@@ -891,8 +1086,7 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
 
 
 class MoeLayerFreqParser(Parser):
-    """
-    Parses moe_layer_freq and first_k_dense_replace fields for models like Deepseek.
+    """Parses moe_layer_freq and first_k_dense_replace fields for models like Deepseek.
 
     Overrides: num_moe_layers
     """
@@ -916,13 +1110,13 @@ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
 
 
 class FfnQuantizationConfigParser(Parser):
-    """
-    Parses quantization configuration for FFN layers.
+    """Parses quantization configuration for FFN layers.
 
     Overrides: weight_byte_size
     """
 
     def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs:
+        """Override weight_byte_size based on the active FFN quantization method."""
         cfg = vllm_config.quant_config
 
         if cfg is None:
@@ -980,10 +1174,16 @@ def validate_moe_fields(self) -> Self:
 
     @classmethod
     def component_type(cls) -> str:
+        """Return the component registry key for FFN layers."""
         return "ffn"
 
     @classmethod
     def get_parser(cls) -> ParserChain:
+        """Return the parser chain for FfnMetrics.
+
+        Parsers run in order: base model config, FFN parallelism, base FFN
+        config, MoE interleave step, MoE layer frequency, quantization override.
+        """
         return ParserChain(
             BaseConfigParser(),
             FfnParallelParser(),
@@ -1207,10 +1407,12 @@ class UnembedMetrics(ComponentMetrics):
 
     @classmethod
     def component_type(cls) -> str:
+        """Return the component registry key for the unembedding (LM head) layer."""
         return "unembed"
 
     @classmethod
     def get_parser(cls) -> ParserChain:
+        """Return the parser chain for UnembedMetrics (base model config only)."""
         return ParserChain(
             BaseConfigParser(),
         )
@@ -1264,11 +1466,9 @@ def get_write_bytes_breakdown(
 
 class ModelMetrics:
     def __init__(self, vllm_config: VllmConfig) -> None:
-        """
-        Parse vllm_config to instantiate metrics for each component.
+        """Parse vllm_config to instantiate metrics for each component.
         is_enabled() will return False if no component metrics could be instantiated.
         """
-
         self.vllm_config = vllm_config
 
         self.metrics: list[ComponentMetrics] = []
@@ -1336,10 +1536,7 @@ def get_write_bytes_breakdown(
     def get_step_perf_stats_per_gpu(
         self, scheduler_output: SchedulerOutput
     ) -> PerfStats:
-        """
-        Calculate perf stats for the current step based on scheduled tokens.
-        """
-
+        """Calculate perf stats for the current step based on scheduled tokens."""
         t0 = time.monotonic()
 
         # Build a single batch context
@@ -1527,7 +1724,7 @@ def log(self, log_fn=logger.info, log_prefix: str = "") -> None:
             )
 
         log_fn(
-            "%sMFU: %.1f TF/s/GPU %.1f GB/s/GPU",
+            "%sMFU: %.1f TF/s/device %.1f GB/s/device",
             log_prefix,
             avg_tflops_per_gpu,
             avg_gbps_per_gpu,
diff --git a/vllm/v1/metrics/prometheus.py b/vllm/v1/metrics/prometheus.py
index c8740276713b..91bb88310d69 100644
--- a/vllm/v1/metrics/prometheus.py
+++ b/vllm/v1/metrics/prometheus.py
@@ -42,6 +42,7 @@ def get_prometheus_registry() -> CollectorRegistry:
 
     Returns:
         Registry: A prometheus registry
+
     """
     if os.getenv("PROMETHEUS_MULTIPROC_DIR") is not None:
         logger.debug("Using multiprocess registry for prometheus metrics")
@@ -70,7 +71,6 @@ def unregister_vllm_metrics():
 
 def shutdown_prometheus():
     """Shutdown prometheus metrics."""
-
     path = _prometheus_multiproc_dir
     if path is None:
         return
diff --git a/vllm/v1/metrics/ray_wrappers.py b/vllm/v1/metrics/ray_wrappers.py
index 7e2100546e82..dedc60ba5425 100644
--- a/vllm/v1/metrics/ray_wrappers.py
+++ b/vllm/v1/metrics/ray_wrappers.py
@@ -68,8 +68,7 @@ def labels(self, *labels, **labelskwargs) -> "RayPrometheusMetric":
 
     @staticmethod
     def _get_sanitized_opentelemetry_name(name: str) -> str:
-        """
-        For compatibility with Ray + OpenTelemetry, the metric name must be
+        """For compatibility with Ray + OpenTelemetry, the metric name must be
         sanitized. In particular, this replaces disallowed character (e.g., ':')
         with '_' in the metric name.
         Allowed characters: a-z, A-Z, 0-9, _
@@ -78,7 +77,6 @@ def _get_sanitized_opentelemetry_name(name: str) -> str:
         Ref: https://github.com/open-telemetry/opentelemetry-cpp/blob/main/sdk/src/metrics/instrument_metadata_validator.cc#L22-L23
         Ref: https://github.com/ray-project/ray/blob/master/src/ray/stats/metric.cc#L107
         """
-
         return re.sub(r"[^a-zA-Z0-9_]", "_", name)
 
 
@@ -169,8 +167,7 @@ def observe(self, value: int | float):
 
 
 class RaySpecDecodingProm(SpecDecodingProm):
-    """
-    RaySpecDecodingProm is used by RayMetrics to log to Ray metrics.
+    """RaySpecDecodingProm is used by RayMetrics to log to Ray metrics.
     Provides the same metrics as SpecDecodingProm but uses Ray's
     util.metrics library.
     """
@@ -179,8 +176,7 @@ class RaySpecDecodingProm(SpecDecodingProm):
 
 
 class RayKVConnectorProm(KVConnectorProm):
-    """
-    RayKVConnectorProm is used by RayMetrics to log Ray
+    """RayKVConnectorProm is used by RayMetrics to log Ray
     metrics. Provides the same metrics as KV connectors but
     uses Ray's util.metrics library.
     """
@@ -191,8 +187,7 @@ class RayKVConnectorProm(KVConnectorProm):
 
 
 class RayPerfMetricsProm(PerfMetricsProm):
-    """
-    RayPerfMetricsProm is used by RayMetrics to log Ray
+    """RayPerfMetricsProm is used by RayMetrics to log Ray
     metrics. Provides the same MFU metrics as PerfMetricsProm
     uses Ray's util.metrics library.
     """
diff --git a/vllm/v1/metrics/reader.py b/vllm/v1/metrics/reader.py
index 48c88e5b61cb..e847e862537a 100644
--- a/vllm/v1/metrics/reader.py
+++ b/vllm/v1/metrics/reader.py
@@ -82,6 +82,7 @@ def get_metrics_snapshot() -> list[Metric]:
         ...         print(f"    count = {metric.count}")
         ...         for bucket_le, value in metrics.buckets.items():
         ...             print(f"    {bucket_le} = {value}")
+
     """
     collected: list[Metric] = []
     for metric in REGISTRY.collect():
diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py
index 24e07e57cc7a..66088c7954aa 100644
--- a/vllm/v1/metrics/stats.py
+++ b/vllm/v1/metrics/stats.py
@@ -34,9 +34,10 @@ class BaseCacheStats:
 
 class CachingMetrics:
     """Metrics for caching with a hit rate of the most recent N requests.
+
     Args:
-        interval: The number of the most recent requests to aggregate.
-            Defaults to 1000.
+        max_recent_requests: The number of the most recent requests to aggregate.
+
     """
 
     def __init__(self, max_recent_requests: int = 1000) -> None:
@@ -62,6 +63,7 @@ def observe(self, stats: BaseCacheStats):
 
         Args:
             stats: The prefix cache stats.
+
         """
         # reset_prefix_cache was invoked before the current update.
         # Reset the metrics before aggregating the current stats.
@@ -113,8 +115,7 @@ def hit_rate(self) -> float:
 
 @dataclass
 class PrefixCacheStats(BaseCacheStats):
-    """
-    Stores prefix cache hit statistics.
+    """Stores prefix cache hit statistics.
     - `reset`: Whether `reset_prefix_cache` was invoked.
     - `queries`: Refers to the number of tokens that were queried.
     """
@@ -144,8 +145,7 @@ def record(self, num_tokens: int, num_hits: int, preempted: bool) -> None:
 
 @dataclass
 class MultiModalCacheStats(BaseCacheStats):
-    """
-    Stores multi-modal cache hit statistics.
+    """Stores multi-modal cache hit statistics.
     - `reset`: Whether `reset_mm_cache` was invoked.
     - `queries`: Refers to the number of multi-modal data items
       that were queried.
@@ -205,6 +205,7 @@ class SchedulerStats:
 
     spec_decoding_stats: SpecDecodingStats | None = None
     kv_connector_stats: dict[str, Any] | None = None
+    ec_connector_stats: dict[str, Any] | None = None
 
     waiting_lora_adapters: dict[str, int] = field(default_factory=dict)
     running_lora_adapters: dict[str, int] = field(default_factory=dict)
diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py
index e94380825f0e..c38eafa4a480 100644
--- a/vllm/v1/outputs.py
+++ b/vllm/v1/outputs.py
@@ -15,13 +15,16 @@
 from vllm.v1.core.sched.output import SchedulerOutput
 
 if TYPE_CHECKING:
+    from vllm.distributed.aux_output_connector.connector import AuxRequestOutput
     from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata
+    from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats
     from vllm.distributed.kv_events import KVConnectorKVEvents
     from vllm.distributed.kv_transfer.kv_connector.v1.base import (
         KVConnectorWorkerMetadata,
     )
     from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 else:
+    AuxRequestOutput = object
     KVConnectorStats = object
     KVConnectorWorkerMetadata = object
     KVConnectorKVEvents = object
@@ -166,7 +169,6 @@ def empty_cpu(
         num_positions: int, num_tokens_per_position: int
     ) -> "LogprobsTensors":
         """Create empty LogprobsTensors on CPU."""
-
         logprob_token_ids = torch.empty(
             (num_positions, num_tokens_per_position),
             dtype=torch.int32,
@@ -188,73 +190,6 @@ def empty_cpu(
         )
 
 
-class RoutedExpertsTensors(NamedTuple):
-    """Device-side snapshot of routed experts data, pending async D2H.
-
-    Produced by :class:`GPUModelRunner` at the end of each async-scheduled
-    step. The copy stream waits on the default stream, then issues
-    non-blocking D2H via :meth:`to_cpu_nonblocking` into a pinned CPU
-    buffer; :class:`AsyncGPUModelRunnerOutput.get_output` synchronizes
-    the copy before the scheduler reads it.
-
-    Sliced to ``total_num_scheduled_tokens`` (step-level, across all
-    requests — NOT per-request). Both ``routing_data`` and
-    ``slot_mapping`` must be private clones when sourced from shared
-    capturer / prepare-input buffers, so the next forward pass /
-    ``_prepare_inputs`` on the default stream does not race with a
-    D2H still pending on the copy stream.
-    """
-
-    # (num_scheduled_tokens, num_layers, num_experts_per_tok)
-    routing_data: torch.Tensor
-    # (num_scheduled_tokens,)
-    slot_mapping: torch.Tensor
-
-    def to_cpu_nonblocking(self) -> "RoutedExpertsTensors":
-        """Issue non-blocking D2H on the current stream.
-
-        NOTE: ``non_blocking=True`` only delivers true overlap when the
-        CPU target is pinned. The current fallback here allocates a
-        new pageable CPU tensor per call, which silently degrades to a
-        synchronous copy; acceptable because the sync happens on the
-        dedicated copy stream, not the default stream.
-        """
-        if self.routing_data.device.type == "cpu":
-            return self
-        return RoutedExpertsTensors(
-            self.routing_data.to("cpu", non_blocking=True),
-            self.slot_mapping.to("cpu", non_blocking=True),
-        )
-
-    def tolists(self) -> "RoutedExpertsLists":
-        """Convert to the numpy-backed form consumed by the scheduler.
-
-        ``.cpu()`` is a no-op when the tensor is already on CPU, so this
-        is cheap for the post-D2H case; for raw device tensors it will
-        synchronously block, which is only reached in tests.
-        """
-        return RoutedExpertsLists(
-            self.routing_data.cpu().numpy(),
-            self.slot_mapping.cpu().numpy(),
-        )
-
-
-class RoutedExpertsLists(NamedTuple):
-    """CPU-side routed experts, the form :meth:`RoutedExpertsManager.store_batch`
-    consumes.
-
-    Batched per scheduler step: the leading dim is the number of tokens
-    scheduled across all requests in this step (``total_num_scheduled_tokens``),
-    not per-request tokens. ``slot_mapping[i]`` tells the scheduler which
-    physical KV-cache slot row ``i`` of ``routing_data`` belongs to.
-    """
-
-    # (num_scheduled_tokens, num_layers, num_experts_per_tok)
-    routing_data: np.ndarray
-    # (num_scheduled_tokens,)
-    slot_mapping: np.ndarray
-
-
 # [num_reqs, <dynamic>]
 # The shape of each element depends on the pooler used
 PoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor] | list[torch.Tensor | None]
@@ -281,6 +216,9 @@ class KVConnectorOutput:
     # IDs of externally computed KV blocks that failed to load.
     # Requests referencing these blocks should be rescheduled to recompute them
     invalid_block_ids: set[int] = field(default_factory=set)
+    # Receive failures keyed by request identity. This remains unambiguous for
+    # hybrid/multi-pool cache layouts where numeric block IDs overlap.
+    failed_recving: set[str] = field(default_factory=set)
     # Configuration describing how many finished sending/receiving
     # notifications should be expected for each request. This allows
     # handshake-based connectors like Nixl to update the KVOutputAggregator.
@@ -295,6 +233,7 @@ def is_empty(self):
             and not self.kv_connector_stats
             and not self.kv_cache_events
             and not self.invalid_block_ids
+            and not self.failed_recving
             and not self.kv_connector_worker_meta
         )
 
@@ -304,12 +243,14 @@ class ECConnectorOutput:
     # [mm_hash]
     finished_sending: set[str] | None = None
     finished_recving: set[str] | None = None
+    ec_connector_stats: "ECConnectorStats | None" = None
     ec_connector_worker_meta: ECConnectorWorkerMetadata | None = None
 
     def is_empty(self):
         return (
             not self.finished_sending
             and not self.finished_recving
+            and not self.ec_connector_stats
             and not self.ec_connector_worker_meta
         )
 
@@ -355,16 +296,7 @@ class ModelRunnerOutput:
     # information related to cudagraph execution
     cudagraph_stats: CUDAGraphStat | None = None
 
-    # Per-step routed experts data captured by the worker.
-    # ``routing_data`` shape: (num_scheduled_tokens, num_layers,
-    #                         num_experts_per_tok); expert IDs as uint8/uint16.
-    # ``slot_mapping`` shape: (num_scheduled_tokens,); physical KV-cache
-    #                         slot for each row of routing_data.
-    # ``num_scheduled_tokens`` is step-level (total across all requests
-    # in this step), not per-request. The scheduler persists this into
-    # its slot buffer via ``slot_buffer[slot_mapping] = routing_data``.
-    # ``None`` when ``enable_return_routed_experts`` is off.
-    routed_experts: RoutedExpertsLists | None = None
+    aux_output_connector_output: dict[str, AuxRequestOutput] | None = None
 
     # ``None`` when ``return_sampling_mask`` is off.
     sampling_masks: SamplingMaskLists | None = None
@@ -433,8 +365,7 @@ class DraftTokenIds:
 def make_empty_encoder_model_runner_output(
     scheduler_output: "SchedulerOutput",
 ) -> ModelRunnerOutput:
-    """
-    Create a ModelRunnerOutput stub that contains the correct
+    """Create a ModelRunnerOutput stub that contains the correct
     per-request bookkeeping but no generated data yet.
     """
     if not scheduler_output.num_scheduled_tokens:
diff --git a/vllm/v1/request.py b/vllm/v1/request.py
index 8b453a09069e..4b1b9ecc2eed 100644
--- a/vllm/v1/request.py
+++ b/vllm/v1/request.py
@@ -20,6 +20,7 @@
     EngineCoreRequest,
     FinishReason,
 )
+from vllm.v1.kv_hints import KvHintsEnvelope
 from vllm.v1.metrics.stats import PrefillStats, RequestSpecDecodeMetrics
 from vllm.v1.structured_output.request import StructuredOutputRequest
 from vllm.v1.utils import ConstantList
@@ -78,6 +79,7 @@ def __init__(
         reasoning_ended: bool | None = None,
         reasoning_parser_kwargs: dict[str, Any] | None = None,
         abort_immediately: bool = False,
+        kv_hints: KvHintsEnvelope | None = None,
     ) -> None:
         self.request_id = request_id
         self.client_index = client_index
@@ -193,6 +195,7 @@ def __init__(
         # trace_headers
         self.trace_headers = trace_headers
         self.session_id = session_id
+        self.kv_hints = kv_hints
 
         # True if this request is scheduled as a non-final prefill chunk.
         self.is_prefill_chunk = False
@@ -201,6 +204,9 @@ def __init__(
         # in the (sparse) prefix cache; 0 means none. Set at admission for
         # hybrid/Mamba models when a shared prefix is detected (Marconi-style).
         self.shared_prefix_boundary = 0
+        # DeepSeek-V4.1 only: SWA bounded replay. The request holds no
+        # sliding-window KV below this position; 0 when nothing replays.
+        self.replay_start = 0
 
         # The number of NaNs in logits. A value greater than 0
         # indicates that the output is corrupted
@@ -257,6 +263,7 @@ def from_engine_core_request(
             block_hasher=block_hasher,
             resumable=request.resumable,
             session_id=request.session_id,
+            kv_hints=request.kv_hints,
             reasoning_ended=request.reasoning_ended,
             reasoning_parser_kwargs=request.reasoning_parser_kwargs,
             abort_immediately=request.abort_immediately,
@@ -348,8 +355,7 @@ def take_prefill_stats(self) -> PrefillStats | None:
         return prefill_stats
 
     def __lt__(self, other: "Request") -> bool:
-        """
-        Compare two requests based on priority, arrival time, and request ID.
+        """Compare two requests based on priority, arrival time, and request ID.
         Used in priority scheduling.
         """
         if self.priority != other.priority:
diff --git a/vllm/v1/sample/logits_processor/__init__.py b/vllm/v1/sample/logits_processor/__init__.py
index f319097495d0..ff9a6428c556 100644
--- a/vllm/v1/sample/logits_processor/__init__.py
+++ b/vllm/v1/sample/logits_processor/__init__.py
@@ -55,8 +55,7 @@
 
 
 def _load_logitsprocs_plugins() -> list[type[LogitsProcessor]]:
-    """Load all installed logit processor plugins"""
-
+    """Load all installed logit processor plugins."""
     from importlib.metadata import entry_points
 
     installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP)
@@ -171,6 +170,7 @@ def _load_custom_logitsprocs(
 
     Returns:
       A list of all loaded logitproc types
+
     """
     from vllm.platforms import current_platform
 
@@ -238,7 +238,7 @@ def validate_logits_processors_parameters(
 
 
 class AdapterLogitsProcessor(LogitsProcessor):
-    """Wrapper for per-request logits processors
+    """Wrapper for per-request logits processors.
 
     To wrap a specific per-request logits processor,
     * Subclass `AdapterLogitsProcessor`
@@ -264,7 +264,6 @@ def __init__(
         these arguments are used, the vLLM logits processor interface requires
         all three arguments to be present.
         """
-
         # Map req index -> logits processor state
         #
         # State representation is a partial[Tensor] comprising a request-level
@@ -301,7 +300,7 @@ def _new_state(
         prompt_ids: list[int] | None,
         output_ids: list[int],
     ) -> partial[torch.Tensor] | None:
-        """Return state representation for new request
+        """Return state representation for new request.
 
         Returns None if logits processor is not applicable to request
 
diff --git a/vllm/v1/sample/logits_processor/builtin.py b/vllm/v1/sample/logits_processor/builtin.py
index 45b5c47711f2..99dd157e78e9 100644
--- a/vllm/v1/sample/logits_processor/builtin.py
+++ b/vllm/v1/sample/logits_processor/builtin.py
@@ -45,7 +45,7 @@ def __init__(
         self.min_p: torch.Tensor = self.min_p_device[:0]
 
     def is_argmax_invariant(self) -> bool:
-        """Min-p never impacts greedy sampling"""
+        """Min-p never impacts greedy sampling."""
         return True
 
     def get_min_p_by_index(self, index: int) -> float:
@@ -377,7 +377,6 @@ def process_dict_updates(
     new_state: Callable[[SamplingParams, list[int] | None, list[int]], T | None],
 ) -> bool:
     """Utility function to update dict state for sparse LogitsProcessors."""
-
     if not batch_update:
         # Nothing to do.
         return False
diff --git a/vllm/v1/sample/logits_processor/interface.py b/vllm/v1/sample/logits_processor/interface.py
index dce0706335e0..355ab6a6f88f 100644
--- a/vllm/v1/sample/logits_processor/interface.py
+++ b/vllm/v1/sample/logits_processor/interface.py
@@ -35,7 +35,7 @@ class MoveDirectionality(Enum):
 
 @dataclass(frozen=True)
 class BatchUpdate:
-    """Persistent batch state change info for logitsprocs"""
+    """Persistent batch state change info for logitsprocs."""
 
     batch_size: int  # Current num reqs in batch
 
@@ -104,5 +104,6 @@ def update_state(
         Args:
             batch_update: Non-None iff there have been changes
                 to the batch makeup.
+
         """
         raise NotImplementedError
diff --git a/vllm/v1/sample/logits_processor/state.py b/vllm/v1/sample/logits_processor/state.py
index 41cbba8dffb3..7988e8c6e921 100644
--- a/vllm/v1/sample/logits_processor/state.py
+++ b/vllm/v1/sample/logits_processor/state.py
@@ -81,6 +81,7 @@ def removed_append(self, index: int) -> None:
 
         Args:
           index: request index
+
         """
         if self._is_removed_sorted:
             raise RuntimeError(
@@ -93,14 +94,14 @@ def has_removed(self) -> bool:
         return bool(self._removed)
 
     def peek_removed(self) -> int | None:
-        """Return lowest removed request index"""
+        """Return lowest removed request index."""
         if self.has_removed():
             self._ensure_removed_sorted()
             return self._removed[-1]
         return None
 
     def pop_removed(self) -> int | None:
-        """Pop lowest removed request index"""
+        """Pop lowest removed request index."""
         if self.has_removed():
             self._ensure_removed_sorted()
             return self._removed.pop()
@@ -125,6 +126,7 @@ def get_and_reset(self, batch_size: int) -> BatchUpdate | None:
 
         Returns:
           Frozen logitsprocs batch update instance; `None` if no updates
+
         """
         # Reset removal-sorting logic
         self._is_removed_sorted = False
diff --git a/vllm/v1/sample/ops/logprobs.py b/vllm/v1/sample/ops/logprobs.py
index 9cdcfa456a80..359a36a81611 100644
--- a/vllm/v1/sample/ops/logprobs.py
+++ b/vllm/v1/sample/ops/logprobs.py
@@ -9,8 +9,7 @@
 
 @torch.compile(backend=current_platform.simple_compile_backend)
 def batched_count_greater_than(x: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
-    """
-    Counts elements in each row of x that are greater than the corresponding
+    """Counts elements in each row of x that are greater than the corresponding
     value in values.  Use torch.compile to generate an optimized kernel for
     this function. otherwise, it will create additional copies of the input
     tensors and cause memory issues.
@@ -21,6 +20,7 @@ def batched_count_greater_than(x: torch.Tensor, values: torch.Tensor) -> torch.T
 
     Returns:
         torch.Tensor: A 1D tensor of shape (batch_size,) with the counts.
+
     """
     torch._check(x.shape[0] >= 1)
     torch._check(x.shape[0] == values.shape[0])
diff --git a/vllm/v1/sample/ops/penalties.py b/vllm/v1/sample/ops/penalties.py
index 7bc6ec7ab893..2e2981bc4c46 100644
--- a/vllm/v1/sample/ops/penalties.py
+++ b/vllm/v1/sample/ops/penalties.py
@@ -15,9 +15,7 @@ def apply_all_penalties(
     repetition_penalties: torch.Tensor,
     output_token_ids: list[list[int]],
 ) -> torch.Tensor:
-    """
-    Applies presence, frequency and repetition penalties to the logits.
-    """
+    """Applies presence, frequency and repetition penalties to the logits."""
     _, vocab_size = logits.shape
     output_tokens_t = _convert_to_tensors(output_token_ids, vocab_size, logits.device)
 
@@ -41,9 +39,7 @@ def apply_all_penalties(
 def _convert_to_tensors(
     output_token_ids: list[list[int]], vocab_size: int, device: torch.device
 ) -> torch.Tensor:
-    """
-    Convert the different list data structures to tensors.
-    """
+    """Convert the different list data structures to tensors."""
     output_tokens_tensor = make_tensor_with_pad(
         output_token_ids,
         # Use the value of vocab_size as a pad since we don't have a
diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py
index 1e8bb03ca7e3..8c26c053b3f3 100644
--- a/vllm/v1/sample/ops/topk_topp_sampler.py
+++ b/vllm/v1/sample/ops/topk_topp_sampler.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-
 import torch
 import torch.nn as nn
 
@@ -10,14 +9,31 @@
 from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
 from vllm.logger import init_logger
 from vllm.platforms import CpuArchEnum, current_platform
+from vllm.platforms.interface import DeviceCapability
 from vllm.triton_utils import HAS_TRITON
 
 if HAS_TRITON:
-    from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
+    from vllm.v1.sample.ops.topk_topp_triton import (
+        _topk_topp,
+        _topp_split_mask,
+        _topp_split_stats,
+        _topp_split_step,
+        apply_top_k_top_p_triton,
+    )
 
 logger = init_logger(__name__)
 
 
+def register_top_k_top_p_warmups() -> None:
+    """Register every native accelerator sampling kernel used at runtime."""
+    if HAS_TRITON and not current_platform.is_cpu():
+        _topk_topp.register_warmup()
+        if current_platform.is_cuda_alike():
+            _topp_split_stats.register_warmup()
+            _topp_split_step.register_warmup()
+            _topp_split_mask.register_warmup()
+
+
 def _aiter_temp_gumbel_sample(
     logits: torch.Tensor, generators: dict[int, torch.Generator]
 ) -> torch.Tensor:
@@ -48,13 +64,47 @@ def _aiter_temp_gumbel_sample(
     return out.to(torch.int64)
 
 
+def _flashinfer_jit_unsupported_reason(capability: DeviceCapability) -> str | None:
+    """Return why FlashInfer JIT codegen cannot target the current GPU, or
+    None if it can.
+
+    FlashInfer swallows arch-detection errors when building its compilation
+    context (e.g. SM 12.x with a CUDA toolkit older than 12.9), leaving an
+    empty target-arch set that makes every JIT spec fail with a misleading
+    "requires sm75 or higher" error at first use — killing the engine during
+    startup profiling (https://github.com/vllm-project/vllm/issues/42393).
+    """
+    try:
+        from flashinfer.jit.core import check_cuda_arch
+    except ImportError:
+        return None
+    try:
+        check_cuda_arch()
+        return None
+    except RuntimeError as e:
+        reason = str(e)
+    try:
+        # Re-derive the real error that FlashInfer swallowed during arch
+        # detection, e.g. "SM 12.x requires CUDA >= 12.9".
+        from flashinfer.compilation_context import CompilationContext
+
+        CompilationContext._normalize_cuda_arch(capability.major, capability.minor)
+    except RuntimeError as e:
+        reason = str(e)
+    except Exception:
+        # FlashInfer internals changed; keep the original error message.
+        pass
+    return reason
+
+
 def flashinfer_sampler_supported() -> bool:
     """Decide whether FlashInfer's top-p/top-k sampler can be used.
 
     Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER``
     is 0, when the platform isn't CUDA, when the GPU's compute capability is
-    unsupported. Raises ``RuntimeError`` if the user explicitly opted in
-    via the env var but FlashInfer is unavailable.
+    unsupported, when the GPU has 16 or fewer SMs, or when FlashInfer cannot
+    JIT-compile for the current GPU/CUDA toolchain. Raises ``RuntimeError`` if
+    the user explicitly opted in via the env var but FlashInfer is unavailable.
 
     Assumes flashinfer is installed, as guaranteed by ``requirements/cuda.txt``;
     otherwise importing the FlashInfer backend below raises ``ImportError``.
@@ -80,6 +130,18 @@ def flashinfer_sampler_supported() -> bool:
         unsupported_reason = (
             f"unsupported compute capability {capability.as_version_str()}"
         )
+    elif (
+        num_sms := current_platform.num_compute_units(
+            torch.accelerator.current_device_index()
+        )
+    ) <= 16:
+        # FlashInfer 0.7+ rejects multi-CTA top-k masking on <=16 SMs because
+        # its cross-CTA software barrier cannot guarantee forward progress.
+        unsupported_reason = (
+            f"top-k masking requires more than 16 SMs; device has {num_sms}"
+        )
+    else:
+        unsupported_reason = _flashinfer_jit_unsupported_reason(capability)
 
     if unsupported_reason is None:
         logger.info_once("Using FlashInfer for top-p & top-k sampling.", scope="global")
@@ -97,9 +159,71 @@ def flashinfer_sampler_supported() -> bool:
     return False
 
 
-class TopKTopPSampler(nn.Module):
+def xpu_sampler_supported() -> bool:
+    """Decide whether the fused XPU top-k/top-p sampler kernel can be used.
+
+    Returns False (with appropriate logging) when the platform isn't XPU, when
+    ``VLLM_XPU_USE_SAMPLER_KERNEL`` is 0.
+
+    Note: callers must additionally ensure no request needs a per-request seed
+    or greedy sampling, since the kernel draws from the device's default
+    generator and always samples randomly.
+    """
+    if not current_platform.is_xpu():
+        return False
+    if not envs.VLLM_XPU_USE_SAMPLER_KERNEL:
+        logger.info_once(
+            "Fused XPU top-p/top-k sampling disabled via VLLM_XPU_USE_SAMPLER_KERNEL=0."
+        )
+        return False
+
+    logger.info_once("Using the fused XPU kernel for top-p & top-k sampling.")
+    return True
+
+
+def xpu_sample(
+    logits: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor | None,
+    logprobs_mode: LogprobsMode = "raw_logprobs",
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Sample from the logits using the fused XPU top-k/top-p kernel.
+
+    Statistically equivalent to `random_sample`, but avoids sorting the vocab
+    and never materializes the probability tensor. The noise comes from the
+    device's default generator, so per-request generators aren't supported.
+
+    Returns the sampled token ids and, for processed logprobs modes, the
+    post-top-k/top-p logits (or logprobs); otherwise None.
     """
-    Module that performs optional top-k and top-p filtering followed by
+    logits = logits.to(dtype=torch.float32).contiguous()
+    sampled = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device)
+    logits_to_return = (
+        torch.empty_like(logits) if logprobs_mode in PROCESSED_LOGPROBS_MODES else None
+    )
+
+    generator = torch.xpu.default_generators[logits.device.index]
+    state = generator.get_state()
+    seed, offset = state.view(torch.int64).tolist()
+    seeds = torch.tensor([seed, offset], dtype=torch.int64, device="cpu")
+    # The XPU kernel expects k as int64 (Long), but the input batch
+    # stores top_k as int32. Cast here to avoid dtype mismatch.
+    if k is not None:
+        k = k.to(torch.int64)
+    torch.ops.vllm.xpu_topk_topp_sampler(
+        sampled, logits_to_return, logits, k, p, logprobs_mode, seeds
+    )
+    # The custom XPU sampler kernel consumes RNG values internally, so advance
+    # the default generator's offset to keep future draws deterministic.
+    # pytorch: offset must be multiple of 4
+    offset = (offset + logits.numel() + 3) // 4 * 4
+    state.view(torch.int64)[1] = offset
+    generator.set_state(state)
+    return sampled, logits_to_return
+
+
+class TopKTopPSampler(nn.Module):
+    """Module that performs optional top-k and top-p filtering followed by
     weighted random sampling of logits.
 
     Implementations may update the logits tensor in-place.
@@ -133,9 +257,14 @@ def __init__(
             else:
                 self.forward = self.forward_cpu
         elif current_platform.is_xpu():
-            if envs.VLLM_XPU_USE_SAMPLER_KERNEL:
+            if xpu_sampler_supported() and not envs.VLLM_BATCH_INVARIANT:
                 self.forward = self.forward_xpu
             else:
+                if envs.VLLM_BATCH_INVARIANT:
+                    logger.info_once(
+                        "VLLM_BATCH_INVARIANT is enabled. Using the "
+                        "PyTorch-native sampler on XPU."
+                    )
                 self.forward = self.forward_native
         elif (
             logprobs_mode not in PROCESSED_LOGPROBS_MODES
@@ -155,6 +284,9 @@ def __init__(
         else:
             self.forward = self.forward_native
 
+        # Every accelerator backend can fall back to native sampling at runtime.
+        register_top_k_top_p_warmups()
+
     def forward_native(
         self,
         logits: torch.Tensor,
@@ -162,8 +294,7 @@ def forward_native(
         k: torch.Tensor | None,
         p: torch.Tensor | None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """
-        PyTorch-native implementation of top-k and top-p sampling.
+        """PyTorch-native implementation of top-k and top-p sampling.
 
         The logits tensor may be updated in-place.
         """
@@ -215,8 +346,7 @@ def forward_cpu(
         k: torch.Tensor | None,
         p: torch.Tensor | None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """
-        PyTorch-native implementation of top-k and top-p sampling for CPU.
+        """PyTorch-native implementation of top-k and top-p sampling for CPU.
 
         The logits tensor may be updated in-place.
         """
@@ -339,37 +469,7 @@ def forward_xpu(
                 "PyTorch-native implementation."
             )
             return self.forward_native(logits, generators, k, p)
-        random_sampled = torch.empty(
-            logits.shape[0], dtype=torch.int64, device=logits.device
-        )
-        logits_to_return = None
-        if self.logprobs_mode in PROCESSED_LOGPROBS_MODES:
-            logits_to_return = torch.empty_like(logits)
-
-        assert len(generators) != logits.shape[0], (
-            "xpu kernel topk_topp_sampler does not support batch-wise generators."
-        )
-        generator = torch.xpu.default_generators[logits.device.index]
-
-        state = generator.get_state()
-        seed, offset = state.view(torch.int64)
-        seeds = torch.tensor(
-            [seed, offset], dtype=torch.int64, device=torch.device("cpu")
-        )
-        # The XPU kernel expects k as int64 (Long), but the input batch
-        # stores top_k as int32. Cast here to avoid dtype mismatch.
-        if k is not None:
-            k = k.to(torch.int64)
-        torch.ops.vllm.xpu_topk_topp_sampler(
-            random_sampled, logits_to_return, logits, k, p, self.logprobs_mode, seeds
-        )
-        # The custom XPU sampler kernel consumes RNG values internally, so advance
-        # the default generator's offset to keep future draws deterministic.
-        # pytorch: offset must be multiple of 4
-        offset = (offset + logits.numel() + 3) // 4 * 4
-        state.view(torch.int64)[1] = offset
-        generator.set_state(state)
-        return random_sampled, logits_to_return
+        return xpu_sample(logits, k, p, self.logprobs_mode)
 
 
 # Note: this is a workaround for
@@ -440,8 +540,7 @@ def apply_top_k_top_p_pytorch(
 
 
 def apply_top_k_only(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
-    """
-    Apply top-k mask to the logits.
+    """Apply top-k mask to the logits.
 
     This implementation doesn't involve sorting the entire vocab.
     Note however that it involves a GPU->CPU sync which can be detrimental for
diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py
index 7c1f080fcdf7..20b636e745a1 100755
--- a/vllm/v1/sample/ops/topk_topp_triton.py
+++ b/vllm/v1/sample/ops/topk_topp_triton.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Combined Top-K and Top-P Triton kernels.
+"""Combined Top-K and Top-P Triton kernels.
 
 Based on the paper "Qrita: High-performance Top-k and Top-p Algorithm for GPUs
 using Pivot-based Truncation and Selection" By Park et al.
@@ -9,8 +8,21 @@
 
 """
 
+from typing import Any
+
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import (
+    WarmupChoices,
+    WarmupIntRange,
+    _when,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.math_utils import next_power_of_2
@@ -863,6 +875,90 @@ def _topk_topp_kernel(
                 tl.store(LOGITS_ROW + offs_n, logits_blk, mask=mask_n)
 
 
+def _max_sampler_batch_size(vllm_config: Any) -> int:
+    return vllm_config.scheduler_config.max_num_seqs * (
+        vllm_config.num_speculative_tokens + 1
+    )
+
+
+def _topk_topp_warmup_inputs(vllm_config: Any) -> dict[str, Any]:
+    vocab_size = vllm_config.model_config.get_vocab_size()
+    split_enabled = current_platform.is_cuda_alike()
+    mode: Any = WarmupChoices((True, False), (True, True), (False, True))
+    logits_stride: Any = WarmupChoices(16, 2)
+    batch_size: Any = WarmupIntRange(1, _max_sampler_batch_size(vllm_config) + 1)
+    topk_enabled = mode[0]
+    topp_enabled = mode[1]
+    _when(topk_enabled or not (split_enabled and batch_size <= _SPLIT_MAX_BATCH))
+    logits = TritonWarmupTensor(
+        torch.float32,
+        shape=(batch_size, vocab_size),
+        strides=(logits_stride, 1),
+    )
+    return dict(
+        logits=logits,
+        buffer=TritonWarmupTensor(torch.float32),
+        percentile_to_std_table=TritonWarmupTensor(torch.float32),
+        normal_cdf_to_sigma_table=TritonWarmupTensor(torch.float32),
+        k=TritonWarmupTensor(torch.int32) if topk_enabled else None,
+        p=TritonWarmupTensor(torch.float32) if topp_enabled else None,
+        mask_value=float("-inf"),
+        num_sm=num_compute_units(),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_topk_topp_kernel,
+    warmup_inputs=_topk_topp_warmup_inputs,
+)
+def _topk_topp(
+    logits: torch.Tensor,
+    buffer: torch.Tensor,
+    percentile_to_std_table: torch.Tensor,
+    normal_cdf_to_sigma_table: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor | None,
+    mask_value: float,
+    num_sm: int,
+) -> DispatchSpec:
+    batch_size, vocab_size = logits.shape
+    topk_enabled = k is not None
+    topp_enabled = p is not None
+    k_ptr = k if k is not None else logits
+    p_ptr = p if p is not None else logits
+    num_programs = min(num_sm, batch_size)
+    # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.
+    # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot
+    # approximation; use smaller tiles for accurate top-p results.
+    if logits.device.type == "cpu":
+        block_size, block_size_trunc, num_warps = 256, 128, 4
+    elif logits.device.type == "xpu":
+        block_size, block_size_trunc, num_warps = 4096, 2048, 4
+    else:
+        block_size, block_size_trunc, num_warps = 8192, 4096, 8
+        # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so
+        # per-tile latency bounds kernel latency, and Triton's default of 4
+        # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is
+        # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not.
+    split_covers_ponly = (
+        current_platform.is_cuda_alike()
+        and topp_enabled
+        and batch_size <= _SPLIT_MAX_BATCH
+    )
+    return (num_programs,), dict(
+        K=k_ptr,
+        P=p_ptr,
+        BATCH_SIZE=batch_size,
+        VOCAB_SIZE=vocab_size,
+        BLOCK_SIZE=block_size,
+        BLOCK_SIZE_TRUNC=block_size_trunc,
+        TOPK_ENABLED=topk_enabled,
+        TOPP_ENABLED=topp_enabled,
+        SPLIT_COVERS_PONLY=split_covers_ponly,
+        num_warps=num_warps,
+    )
+
+
 # ---------------------------------------------------------------------------
 # Split-row pipeline for small batches.
 #
@@ -892,6 +988,13 @@ def _topk_topp_kernel(
 _TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {}
 
 
+def _topp_split_count(batch_size: int, num_sm: int) -> int:
+    splits = 1
+    while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:
+        splits *= 2
+    return splits
+
+
 @triton.jit
 def _topp_sb_stats_kernel(
     LOGITS,
@@ -939,6 +1042,51 @@ def _topp_sb_stats_kernel(
     tl.store(STATS + base + 3, cnt)
 
 
+def _topp_split_stats_warmup_inputs(vllm_config: Any) -> dict[str, Any]:
+    vocab_size = vllm_config.model_config.get_vocab_size()
+    num_sm = num_compute_units()
+    logits_stride: Any = WarmupChoices(16, 2)
+    batch_size: Any = WarmupIntRange(
+        1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1
+    )
+    has_k: Any = WarmupChoices(False, True)
+    logits = TritonWarmupTensor(
+        torch.float32,
+        shape=(batch_size, vocab_size),
+        strides=(logits_stride, 1),
+    )
+    return dict(
+        logits=logits,
+        stats=TritonWarmupTensor(torch.float32),
+        k=TritonWarmupTensor(torch.int32) if has_k else None,
+        p=TritonWarmupTensor(torch.float32),
+        num_sm=num_sm,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_topp_sb_stats_kernel,
+    warmup_inputs=_topp_split_stats_warmup_inputs,
+)
+def _topp_split_stats(
+    logits: torch.Tensor,
+    stats: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor,
+    num_sm: int,
+) -> DispatchSpec:
+    batch_size, vocab_size = logits.shape
+    splits = _topp_split_count(batch_size, num_sm)
+    return (logits.shape[0] * splits,), dict(
+        K=k if k is not None else logits,
+        HAS_K=k is not None,
+        VOCAB_SIZE=vocab_size,
+        S=splits,
+        BLOCK=8192,
+        num_warps=8,
+    )
+
+
 @triton.jit
 def _topp_sb_row_stats(STATS, row, S: tl.constexpr):
     """Combine per-slice partials into row (max, sum_exp, min, n_finite)."""
@@ -1183,6 +1331,58 @@ def _topp_sb_step_kernel(
     tl.store(pbase + fidx * 3 + 2, nmin)
 
 
+def _topp_split_step_warmup_inputs(vllm_config: Any) -> dict[str, Any]:
+    vocab_size = vllm_config.model_config.get_vocab_size()
+    num_sm = num_compute_units()
+    logits_stride: Any = WarmupChoices(16, 2)
+    batch_size: Any = WarmupIntRange(
+        1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1
+    )
+    round: Any = WarmupIntRange(0, _SPLIT_ROUNDS)
+    has_k: Any = WarmupChoices(False, True)
+    logits = TritonWarmupTensor(
+        torch.float32,
+        shape=(batch_size, vocab_size),
+        strides=(logits_stride, 1),
+    )
+    return dict(
+        logits=logits,
+        stats=TritonWarmupTensor(torch.float32),
+        parts=TritonWarmupTensor(torch.float32),
+        k=TritonWarmupTensor(torch.int32) if has_k else None,
+        p=TritonWarmupTensor(torch.float32),
+        round=round,
+        num_sm=num_sm,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_topp_sb_step_kernel,
+    warmup_inputs=_topp_split_step_warmup_inputs,
+)
+def _topp_split_step(
+    logits: torch.Tensor,
+    stats: torch.Tensor,
+    parts: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor,
+    round: int,
+    num_sm: int,
+) -> DispatchSpec:
+    batch_size, vocab_size = logits.shape
+    splits = _topp_split_count(batch_size, num_sm)
+    return (logits.shape[0] * splits,), dict(
+        K=k if k is not None else logits,
+        HAS_K=k is not None,
+        S=splits,
+        F=_SPLIT_FANOUT,
+        NUM_ROUNDS=_SPLIT_ROUNDS,
+        VOCAB_SIZE=vocab_size,
+        BLOCK=2048,
+        num_warps=8,
+    )
+
+
 @triton.jit
 def _topp_sb_mask_kernel(
     LOGITS,
@@ -1313,6 +1513,57 @@ def _topp_sb_mask_kernel(
         tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n)
 
 
+def _topp_split_mask_warmup_inputs(vllm_config: Any) -> dict[str, Any]:
+    vocab_size = vllm_config.model_config.get_vocab_size()
+    num_sm = num_compute_units()
+    logits_stride: Any = WarmupChoices(16, 2)
+    batch_size: Any = WarmupIntRange(
+        1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1
+    )
+    has_k: Any = WarmupChoices(False, True)
+    logits = TritonWarmupTensor(
+        torch.float32,
+        shape=(batch_size, vocab_size),
+        strides=(logits_stride, 1),
+    )
+    return dict(
+        logits=logits,
+        stats=TritonWarmupTensor(torch.float32),
+        parts=TritonWarmupTensor(torch.float32),
+        k=TritonWarmupTensor(torch.int32) if has_k else None,
+        p=TritonWarmupTensor(torch.float32),
+        mask_value=float("-inf"),
+        num_sm=num_sm,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_topp_sb_mask_kernel,
+    warmup_inputs=_topp_split_mask_warmup_inputs,
+)
+def _topp_split_mask(
+    logits: torch.Tensor,
+    stats: torch.Tensor,
+    parts: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor,
+    mask_value: float,
+    num_sm: int,
+) -> DispatchSpec:
+    batch_size, vocab_size = logits.shape
+    splits = _topp_split_count(batch_size, num_sm)
+    return (logits.shape[0] * splits,), dict(
+        K=k if k is not None else logits,
+        HAS_K=k is not None,
+        S=splits,
+        F=_SPLIT_FANOUT,
+        NUM_ROUNDS=_SPLIT_ROUNDS,
+        VOCAB_SIZE=vocab_size,
+        BLOCK=8192,
+        num_warps=8,
+    )
+
+
 def _apply_topp_split(
     logits: torch.Tensor,
     k: torch.Tensor | None,
@@ -1321,11 +1572,6 @@ def _apply_topp_split(
     num_sm: int,
 ) -> None:
     """Split-row top-p pipeline for small batches; masks logits in place."""
-    batch_size, vocab_size = logits.shape
-    splits = 1
-    while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:
-        splits *= 2
-
     ws = _TRITON_SPLIT_CACHE.get(logits.device)
     if ws is None:
         ws = {
@@ -1336,53 +1582,31 @@ def _apply_topp_split(
         }
         _TRITON_SPLIT_CACHE[logits.device] = ws
 
-    k_ptr = k if k is not None else logits  # dummy pointer when HAS_K=False
-    has_k = k is not None
-
-    _topp_sb_stats_kernel[(batch_size * splits,)](
+    _topp_split_stats(
         logits,
-        logits.stride(0),
         ws["stats"],
-        k_ptr,
+        k,
         p,
-        HAS_K=has_k,
-        VOCAB_SIZE=vocab_size,
-        S=splits,
-        BLOCK=8192,
-        num_warps=8,
+        num_sm,
     )
     for round_i in range(_SPLIT_ROUNDS):
-        _topp_sb_step_kernel[(batch_size * splits,)](
+        _topp_split_step(
             logits,
-            logits.stride(0),
             ws["stats"],
             ws["parts"],
-            k_ptr,
+            k,
             p,
             round_i,
-            HAS_K=has_k,
-            S=splits,
-            F=_SPLIT_FANOUT,
-            NUM_ROUNDS=_SPLIT_ROUNDS,
-            VOCAB_SIZE=vocab_size,
-            BLOCK=2048,
-            num_warps=8,
+            num_sm,
         )
-    _topp_sb_mask_kernel[(batch_size * splits,)](
+    _topp_split_mask(
         logits,
-        logits.stride(0),
         ws["stats"],
         ws["parts"],
-        k_ptr,
+        k,
         p,
-        HAS_K=has_k,
-        MASK_VALUE=mask_value,
-        S=splits,
-        F=_SPLIT_FANOUT,
-        NUM_ROUNDS=_SPLIT_ROUNDS,
-        VOCAB_SIZE=vocab_size,
-        BLOCK=8192,
-        num_warps=8,
+        mask_value,
+        num_sm,
     )
 
 
@@ -1392,8 +1616,7 @@ def apply_top_k_top_p_triton(
     p: torch.Tensor | None,
     mask_value: float = float("-inf"),
 ) -> torch.Tensor:
-    """
-    Apply combined top-k and top-p masking using Triton.
+    """Apply combined top-k and top-p masking using Triton.
 
     Top-k is applied first (by logit value), then top-p is applied
     to the remaining k values (by probability).
@@ -1409,6 +1632,7 @@ def apply_top_k_top_p_triton(
 
     Returns:
         The masked logits tensor. It may or may not be modified in-place.
+
     """
     assert logits.ndim == 2
     assert logits.dtype == torch.float32
@@ -1477,39 +1701,15 @@ def apply_top_k_top_p_triton(
     else:
         normal_cdf_to_sigma_table, percentile_to_std_table = tables
 
-    # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.
-    # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot
-    # approximation; use smaller tiles for accurate top-p results.
-    launch_kwargs = {}
-    if logits.device.type == "cpu":
-        block_size, block_size_trunc = 256, 128
-    elif logits.device.type == "xpu":
-        block_size, block_size_trunc = 4096, 2048
-    else:
-        block_size, block_size_trunc = 8192, 4096
-        # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so
-        # per-tile latency bounds kernel latency, and Triton's default of 4
-        # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is
-        # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not.
-        launch_kwargs["num_warps"] = 8
-
-    _topk_topp_kernel[(NUM_PROGRAMS,)](
+    _topk_topp(
         logits,
-        logits.stride(0),
         buffer,
         percentile_to_std_table,
         normal_cdf_to_sigma_table,
-        k_ptr,
-        p_ptr,
-        BATCH_SIZE=batch_size,
-        MASK_VALUE=mask_value,
-        VOCAB_SIZE=vocab_size,
-        BLOCK_SIZE=block_size,
-        BLOCK_SIZE_TRUNC=block_size_trunc,
-        TOPK_ENABLED=topk_enabled,
-        TOPP_ENABLED=topp_enabled,
-        SPLIT_COVERS_PONLY=use_split,
-        **launch_kwargs,
+        k_ptr if topk_enabled else None,
+        p_ptr if topp_enabled else None,
+        mask_value,
+        num_sm,
     )
     if use_split:
         _apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm)
diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py
index a1baf0a80851..d5f01f0fb40e 100644
--- a/vllm/v1/sample/rejection_sampler.py
+++ b/vllm/v1/sample/rejection_sampler.py
@@ -5,14 +5,21 @@
 
 from collections.abc import Sequence
 from dataclasses import replace
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn as nn
 
 from vllm.config.model import PROCESSED_LOGPROBS_MODES
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.jit_warmup import WarmupChoices
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
 from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput
 from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor
 from vllm.v1.sample.metadata import SamplingMetadata
@@ -36,8 +43,7 @@
 
 
 class RejectionSampler(nn.Module):
-    """
-    The implementation strictly follows the algorithm described in
+    """The implementation strictly follows the algorithm described in
         https://arxiv.org/abs/2211.17192.
     However, we want to clarify the terminology used in the implementation:
     accepted tokens: tokens that are accepted based on the relationship
@@ -88,6 +94,14 @@ def __init__(
                 device=device,
             )
         self.synthetic_mode = self.synthetic_conditional_rates is not None
+        _rejection_greedy_sample.register_warmup(synthetic_mode=self.synthetic_mode)
+        _rejection_random_sample.register_warmup(
+            synthetic_mode=self.synthetic_mode,
+        )
+        _expand.register_warmup()
+        _sample_recovered_tokens.register_warmup(
+            use_fp64_gumbel=self.use_fp64_gumbel,
+        )
 
     def forward(
         self,
@@ -98,8 +112,7 @@ def forward(
         logits: torch.Tensor,
         sampling_metadata: SamplingMetadata,
     ) -> SamplerOutput:
-        """
-        Args:
+        """Args:
             metadata:
                 Metadata for spec decoding.
             draft_probs (Optional[torch.Tensor]):
@@ -115,10 +128,12 @@ def forward(
             sampling_metadata (vllm.v1.sample.metadata.SamplingMetadata):
                 Additional metadata needed for sampling, such as temperature,
                 top-k/top-p parameters, or other relevant information.
+
         Returns:
             SamplerOutput:
                 Contains the final output token IDs and their logprobs if
                 requested.
+
         """
         assert metadata.max_spec_len <= MAX_SPEC_LEN
 
@@ -257,6 +272,7 @@ def parse_output(
         logprobs_tensors: LogprobsTensors | None = None,
     ) -> tuple[list[list[int]], LogprobsLists | None]:
         """Parse the output of the rejection sampler.
+
         Args:
             output_token_ids: The sampled token IDs in shape
                 [batch_size, max_spec_len + 1]. The rejected tokens are
@@ -265,8 +281,10 @@ def parse_output(
             vocab_size: The size of the vocabulary.
             discard_req_indices: Optional row indices to discard tokens in.
             logprobs_tensors: Optional logprobs tensors to filter.
+
         Returns:
             A list of lists of token IDs.
+
         """
         output_token_ids_np = output_token_ids.cpu().numpy()
         # Create mask for valid tokens.
@@ -313,9 +331,7 @@ def apply_logits_processors(
             num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu")
             original_indices = torch.arange(num_requests, device="cpu")
             repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens)
-            repeat_indices = repeat_indices_cpu.to(
-                device=logits.device, non_blocking=True
-            )
+            repeat_indices = async_tensor_h2d(repeat_indices_cpu, logits.device)
             logits = self.apply_penalties(
                 logits, sampling_metadata, metadata, repeat_indices, output_token_ids
             )
@@ -453,7 +469,7 @@ def rejection_sample(
     if not sampling_metadata.all_random:
         # Rejection sampling for greedy sampling requests.
         target_argmax = target_logits.argmax(dim=-1)
-        rejection_greedy_sample_kernel[(batch_size,)](
+        _rejection_greedy_sample(
             output_token_ids,
             cu_num_draft_tokens,
             draft_token_ids,
@@ -463,7 +479,6 @@ def rejection_sample(
             max_spec_len,
             uniform_probs,
             synthetic_conditional_rates,
-            SYNTHETIC_MODE=synthetic_mode,
         )
         if sampling_metadata.all_greedy:
             return output_token_ids
@@ -488,7 +503,7 @@ def rejection_sample(
 
     # Rejection sampling for random sampling requests.
     assert uniform_probs is not None
-    rejection_random_sample_kernel[(batch_size,)](
+    _rejection_random_sample(
         output_token_ids,
         cu_num_draft_tokens,
         draft_token_ids,
@@ -499,10 +514,7 @@ def rejection_sample(
         uniform_probs,
         is_greedy,
         max_spec_len,
-        vocab_size,
         synthetic_conditional_rates,
-        NO_DRAFT_PROBS=draft_probs is None,
-        SYNTHETIC_MODE=synthetic_mode,
     )
     return output_token_ids
 
@@ -527,6 +539,7 @@ def apply_sampling_constraints(
     Returns:
         torch.Tensor: Processed logits if non-greedy sampling is used,
         otherwise returns the original logits.
+
     """
     assert logits.ndim == 2
     assert cu_num_draft_tokens.ndim == 1
@@ -588,19 +601,20 @@ def expand_batch_to_tokens(
             Value to be replaced if it is found in x.
         replace_to: int = 0
             Value to replace with when replace_from is found.
+
     Returns:
         expanded_x: [num_tokens] tensor.
+
     """
     batch_size = x.shape[0]
     assert cu_num_tokens.shape[0] == batch_size
     expanded_x = x.new_empty(num_tokens)
-    expand_kernel[(batch_size,)](
+    _expand(
         expanded_x,
         x,
         cu_num_tokens,
         replace_from,
         replace_to,
-        MAX_NUM_TOKENS=MAX_SPEC_LEN,  # To avoid recompilation.
     )
     return expanded_x
 
@@ -611,8 +625,7 @@ def generate_uniform_probs(
     generators: dict[int, torch.Generator],
     device: torch.device,
 ) -> torch.Tensor:
-    """
-    Generates a batch of uniform random samples, with optional seeding
+    """Generates a batch of uniform random samples, with optional seeding
     if available.
 
     This method creates a tensor of shape `(num_tokens, )` filled
@@ -631,10 +644,12 @@ def generate_uniform_probs(
             `torch.Generator` objects.
         device: torch.device
             The device on which to allocate the tensor.
+
     Returns:
         uniform_rand: torch.Tensor
             A tensor of shape `(num_tokens, )` containing uniform
             random values in the range [0, 1).
+
     """
     # NOTE(woosuk): We deliberately use float64 instead of float32 here
     # because when using float32, there's a non-negligible chance that
@@ -694,18 +709,14 @@ def sample_recovered_tokens(
     inv_q = q.reciprocal()
 
     recovered_token_ids = torch.empty_like(draft_token_ids)
-    BLOCK_SIZE = 8192
-    sample_recovered_tokens_kernel[(batch_size, max_spec_len)](
+    _sample_recovered_tokens(
         recovered_token_ids,
         cu_num_draft_tokens,
         draft_token_ids,
         draft_probs,
         target_probs,
         inv_q,
-        vocab_size,
-        BLOCK_SIZE,
-        NO_DRAFT_PROBS=draft_probs is None,
-        USE_FP64_GUMBEL=use_fp64_gumbel,
+        max_spec_len=max_spec_len,
     )
     return recovered_token_ids
 
@@ -769,6 +780,50 @@ def rejection_greedy_sample_kernel(
         )
 
 
+def _rejection_greedy_sample_warmup_inputs(
+    *, synthetic_mode: bool
+) -> dict[str, object]:
+    all_greedy = WarmupChoices(False, True)
+    int32 = TritonWarmupTensor(torch.int32)
+    return dict(
+        output_token_ids=int32,
+        cu_num_draft_tokens=int32,
+        draft_token_ids=int32,
+        target_argmax=TritonWarmupTensor(torch.int64),
+        bonus_token_ids=TritonWarmupTensor(torch.int64),
+        is_greedy=None if all_greedy else TritonWarmupTensor(torch.bool),
+        max_spec_len=5,
+        uniform_probs=(
+            None
+            if all_greedy and not synthetic_mode
+            else TritonWarmupTensor(torch.float64)
+        ),
+        synthetic_conditional_rates=(
+            TritonWarmupTensor(torch.float32) if synthetic_mode else None
+        ),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=rejection_greedy_sample_kernel,
+    warmup_inputs=_rejection_greedy_sample_warmup_inputs,
+)
+def _rejection_greedy_sample(
+    output_token_ids: torch.Tensor,
+    cu_num_draft_tokens: torch.Tensor,
+    draft_token_ids: torch.Tensor,
+    target_argmax: torch.Tensor,
+    bonus_token_ids: torch.Tensor,
+    is_greedy: torch.Tensor | None,
+    max_spec_len: int,
+    uniform_probs: torch.Tensor | None,
+    synthetic_conditional_rates: torch.Tensor | None,
+) -> DispatchSpec:
+    return (output_token_ids.shape[0],), dict(
+        SYNTHETIC_MODE=synthetic_conditional_rates is not None,
+    )
+
+
 # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
 @triton.jit(do_not_specialize=["max_spec_len"])
 def rejection_random_sample_kernel(
@@ -845,6 +900,57 @@ def rejection_random_sample_kernel(
         )
 
 
+def _rejection_random_sample_warmup_inputs(
+    *, synthetic_mode: bool
+) -> dict[str, object]:
+    vocab_size: Any = WarmupChoices(2, 16)
+    no_draft_probs = WarmupChoices(False, True)
+    int32 = TritonWarmupTensor(torch.int32)
+    probabilities = TritonWarmupTensor(
+        torch.float32,
+        shape=(1, vocab_size),
+    )
+    return dict(
+        output_token_ids=int32,
+        cu_num_draft_tokens=int32,
+        draft_token_ids=int32,
+        draft_probs=None if no_draft_probs else probabilities,
+        target_probs=probabilities,
+        bonus_token_ids=TritonWarmupTensor(torch.int64),
+        recovered_token_ids=int32,
+        uniform_probs=TritonWarmupTensor(torch.float64),
+        is_greedy=TritonWarmupTensor(torch.bool),
+        max_spec_len=5,
+        synthetic_conditional_rates=(
+            TritonWarmupTensor(torch.float32) if synthetic_mode else None
+        ),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=rejection_random_sample_kernel,
+    warmup_inputs=_rejection_random_sample_warmup_inputs,
+)
+def _rejection_random_sample(
+    output_token_ids: torch.Tensor,
+    cu_num_draft_tokens: torch.Tensor,
+    draft_token_ids: torch.Tensor,
+    draft_probs: torch.Tensor | None,
+    target_probs: torch.Tensor,
+    bonus_token_ids: torch.Tensor,
+    recovered_token_ids: torch.Tensor,
+    uniform_probs: torch.Tensor,
+    is_greedy: torch.Tensor,
+    max_spec_len: int,
+    synthetic_conditional_rates: torch.Tensor | None,
+) -> DispatchSpec:
+    return (output_token_ids.shape[0],), dict(
+        vocab_size=target_probs.shape[1],
+        NO_DRAFT_PROBS=draft_probs is None,
+        SYNTHETIC_MODE=synthetic_conditional_rates is not None,
+    )
+
+
 # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
 @triton.jit(do_not_specialize=["replace_from", "replace_to"])
 def expand_kernel(
@@ -869,6 +975,31 @@ def expand_kernel(
     tl.store(output_ptr + start_idx + offset, src_val, mask=offset < num_tokens)
 
 
+def _expand_warmup_inputs() -> dict[str, object]:
+    dtype = WarmupChoices(torch.float32, torch.int32, torch.int64)
+    data = TritonWarmupTensor(dtype)
+    return dict(
+        output=data,
+        input=data,
+        cu_num_tokens=TritonWarmupTensor(torch.int32),
+        replace_from=0,
+        replace_to=0,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=expand_kernel, warmup_inputs=_expand_warmup_inputs
+)
+def _expand(
+    output: torch.Tensor,
+    input: torch.Tensor,
+    cu_num_tokens: torch.Tensor,
+    replace_from: int,
+    replace_to: int,
+) -> DispatchSpec:
+    return (input.shape[0],), dict(MAX_NUM_TOKENS=MAX_SPEC_LEN)
+
+
 @triton.jit
 def sample_recovered_tokens_kernel(
     output_token_ids_ptr,  # [num_tokens]
@@ -951,3 +1082,49 @@ def sample_recovered_tokens_kernel(
 
     recovered_id = tl.minimum(recovered_id, vocab_size - 1)
     tl.store(output_token_ids_ptr + token_idx, recovered_id)
+
+
+def _sample_recovered_tokens_warmup_inputs(
+    *, use_fp64_gumbel: bool
+) -> dict[str, object]:
+    vocab_size: Any = WarmupChoices(2, 16)
+    no_draft_probs = WarmupChoices(False, True)
+    int32 = TritonWarmupTensor(torch.int32)
+    probabilities = TritonWarmupTensor(
+        torch.float32,
+        shape=(1, vocab_size),
+    )
+    return dict(
+        output_token_ids=int32,
+        cu_num_draft_tokens=int32,
+        draft_token_ids=int32,
+        draft_probs=None if no_draft_probs else probabilities,
+        target_probs=probabilities,
+        inv_q=TritonWarmupTensor(
+            torch.float64 if use_fp64_gumbel else torch.float32,
+            shape=(1, vocab_size),
+        ),
+        max_spec_len=5,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=sample_recovered_tokens_kernel,
+    warmup_inputs=_sample_recovered_tokens_warmup_inputs,
+)
+def _sample_recovered_tokens(
+    output_token_ids: torch.Tensor,
+    cu_num_draft_tokens: torch.Tensor,
+    draft_token_ids: torch.Tensor,
+    draft_probs: torch.Tensor | None,
+    target_probs: torch.Tensor,
+    inv_q: torch.Tensor,
+    *,
+    max_spec_len: int,
+) -> DispatchSpec:
+    return (cu_num_draft_tokens.shape[0], max_spec_len), dict(
+        vocab_size=target_probs.shape[1],
+        BLOCK_SIZE=8192,
+        NO_DRAFT_PROBS=draft_probs is None,
+        USE_FP64_GUMBEL=inv_q.dtype == torch.float64,
+    )
diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py
index fc1143af831f..1b3eacd33f57 100644
--- a/vllm/v1/sample/sampler.py
+++ b/vllm/v1/sample/sampler.py
@@ -75,8 +75,7 @@ def _aiter_raw_fused_sample(logits: torch.Tensor, sm: SamplingMetadata) -> torch
 
 
 class Sampler(nn.Module):
-    """
-    A layer that samples the next tokens from the model's outputs
+    """A layer that samples the next tokens from the model's outputs
     with the following steps in order:
 
     1. If logprobs are requested:
@@ -231,6 +230,7 @@ def gather_specific_token_logprobs(
         Returns:
             LogprobsTensors with logprobs for the specified tokens, or None
             if no requests have logprob_token_ids.
+
         """
         if not logprob_token_ids:
             return None
@@ -313,7 +313,6 @@ def sample(
         The various logits processing functions called in this method
         may update the logits tensor in-place.
         """
-
         logprobs_mode = logprobs_mode_override or self.logprobs_mode
         assert not (sampling_metadata.all_greedy and sampling_metadata.all_random)
         if sampling_metadata.all_random:
@@ -373,8 +372,7 @@ def gather_logprobs(
         num_logprobs: int,
         token_ids: torch.Tensor,
     ) -> LogprobsTensors:
-        """
-        Gather logprobs for topk and sampled/prompt token.
+        """Gather logprobs for topk and sampled/prompt token.
 
         Args:
           logprobs: (num tokens) x (vocab) tensor
@@ -390,6 +388,7 @@ def gather_logprobs(
           Top-k int indices tensor, (num tokens) x (num_logprobs + 1)
           Top-k float logprobs tensor, (num tokens) x (num_logprobs + 1)
           Sampled token rank tensor, (num tokens)
+
         """
         assert token_ids.dtype == torch.int64
         # Find the topK values.
diff --git a/vllm/v1/serial_utils.py b/vllm/v1/serial_utils.py
index bc4619a7eb3d..882bc44917fd 100644
--- a/vllm/v1/serial_utils.py
+++ b/vllm/v1/serial_utils.py
@@ -57,8 +57,7 @@
 class OOBTensorConsumer(ABC):
     @abstractmethod
     def __call__(self, tensor: torch.Tensor) -> dict | None:
-        """
-        Called with tensors for the current message.
+        """Called with tensors for the current message.
         Returns None to reject the tensor (falls back to regular serialization),
         otherwise a dict with arbitrary placeholder data to be included
         in the serialized message.
@@ -489,8 +488,7 @@ def run_method(
     args: tuple[Any, ...],
     kwargs: dict[str, Any],
 ) -> Any:
-    """
-    Run a method of an object with the given arguments and keyword arguments.
+    """Run a method of an object with the given arguments and keyword arguments.
     If the method is string, it will be converted to a method using getattr.
     If the method is serialized bytes and will be deserialized using
     cloudpickle.
@@ -528,8 +526,7 @@ class PydanticMsgspecMixin:
     def __get_pydantic_core_schema__(
         cls, source_type: Any, handler: GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
-        """
-        Make msgspec.Struct compatible with Pydantic, respecting defaults.
+        """Make msgspec.Struct compatible with Pydantic, respecting defaults.
         Handle JSON=>msgspec.Struct. Used when exposing msgspec.Struct to the
         API as input or in `/docs`. Note this is cached by Pydantic and not
         called on every validation.
diff --git a/vllm/v1/simple_kv_offload/disk_backend.py b/vllm/v1/simple_kv_offload/disk_backend.py
index f39be7d05f56..dfa717cc9ac1 100644
--- a/vllm/v1/simple_kv_offload/disk_backend.py
+++ b/vllm/v1/simple_kv_offload/disk_backend.py
@@ -99,9 +99,12 @@ def init(
             t.stride(0) * t.element_size() for t in gpu_caches.values()
         ]
 
-        assert total_block_bytes % _ALIGNMENT == 0, (
-            f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}"
-        )
+        # The stride alignment is an O_DIRECT kernel-DMA requirement; buffered
+        # I/O through the page cache has no alignment constraint.
+        if not use_page_cache:
+            assert total_block_bytes % _ALIGNMENT == 0, (
+                f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}"
+            )
 
         # Separate buffer pools for store and load threads
         self._store_buffer_caches = {}
diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py
index afb51fb1d32b..96bfb35a6760 100644
--- a/vllm/v1/simple_kv_offload/manager.py
+++ b/vllm/v1/simple_kv_offload/manager.py
@@ -5,6 +5,7 @@
 import contextlib
 from collections.abc import Iterable
 from dataclasses import dataclass, field, replace
+from enum import Enum, auto
 from typing import TYPE_CHECKING, Any
 
 from vllm.config import VllmConfig
@@ -26,6 +27,7 @@
 from vllm.v1.core.kv_cache_utils import (
     BlockHashWithGroupId,
     ExternalBlockHash,
+    dcp_world_size_for_kv_cache_spec,
     generate_block_hash_extra_keys,
     get_block_hash,
     get_group_id,
@@ -49,7 +51,7 @@
 
 if TYPE_CHECKING:
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
-    from vllm.v1.core.kv_cache_utils import KVCacheBlock
+    from vllm.v1.core.kv_cache_utils import BlockHashList, KVCacheBlock
     from vllm.v1.kv_cache_interface import KVCacheConfig
     from vllm.v1.request import Request
 
@@ -99,6 +101,36 @@ class StoreRequestState:
     finished: bool = False
 
 
+class _StoreAdmission(Enum):
+    READY = auto()
+    NULL_BLOCK = auto()
+    NOT_HASHED = auto()
+    IN_FLIGHT = auto()
+    ALREADY_CACHED = auto()
+
+
+@dataclass
+class BoundaryStoreStats:
+    """Cumulative counters for boundary hand-off stores.
+
+    Every branch that declines a hand-off is silent by design: the boundary is
+    simply not offloaded and a later request misses. Without counters the only
+    symptom is a lower cache hit rate with nothing in the logs, so each decline
+    reason is counted and exposed through
+    ``SimpleCPUOffloadConnector.get_boundary_store_stats()``. Reset by
+    ``reset()``; not otherwise cleared.
+    """
+
+    published: int = 0
+    stored: int = 0
+    dropped_cpu_full: int = 0
+    dropped_request_gone: int = 0
+    skipped_already_cached: int = 0
+    skipped_in_flight: int = 0
+    dropped_null_block: int = 0
+    dropped_not_hashed: int = 0
+
+
 class SimpleCPUOffloadScheduler:
     """Scheduler-side manager for CPU offloading."""
 
@@ -133,6 +165,9 @@ def __init__(
             kv_cache_config, offload_capacity
         )
         self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
+        self.prefix_cacheable_group_ids = (
+            self.cpu_kv_cache_config.prefix_cacheable_group_ids
+        )
         self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU
         # Find the full attention kv group for prefix cache matching.
         self.fa_gidx = -1
@@ -141,16 +176,6 @@ def __init__(
                 self.fa_gidx = g_idx
                 break
         assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups)
-        # FA group's own block_size; divides scheduler_block_size (the LCM)
-        # but is NOT assumed to equal it.
-        self.fa_block_size: int = (
-            self.cpu_kv_cache_config.kv_cache_groups[
-                self.fa_gidx
-            ].kv_cache_spec.block_size
-            * self.cp_world_size
-        )
-        assert self.block_size % self.fa_block_size == 0
-
         logger.info(
             "SimpleCPUOffloadScheduler: Allocating %d offload blocks "
             "(%.2f GB, mode=%s, backend=%s)",
@@ -175,7 +200,13 @@ def __init__(
             pcp_world_size=1,
             scheduler_block_size=self.block_size,
             hash_block_size=self.hash_block_size,
+            allow_partial_hash_hits=not lazy_offload,
         )
+        self.group_block_sizes = self.cpu_coordinator.group_block_sizes
+        # FA group's own resolved block_size; divides scheduler_block_size (the
+        # LCM) but is NOT assumed to equal it.
+        self.fa_block_size: int = self.group_block_sizes[self.fa_gidx]
+        assert self.block_size % self.fa_block_size == 0
         self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool
         # GPU block pool reference - bound after scheduler builds kv_cache_manager
         self._gpu_block_pool: BlockPool | None = None
@@ -186,10 +217,13 @@ def __init__(
         # the worker reports completions by event index, not request id.
         self._load_event_to_reqs: dict[int, list[str]] = {}
 
-        # Pending (cpu_hit_blocks, hit_length) tuples from find_longest_cache_hit,
-        # kept pinned via touch() while awaiting update_state_after_alloc().
+        # Pending (cpu_hit_blocks, hit_length, num_computed_tokens) tuples from
+        # find_longest_cache_hit, kept pinned via touch() while awaiting
+        # update_state_after_alloc(). ``num_computed_tokens`` is the local
+        # prefix the hit was resolved against; the load path needs it to place
+        # external blocks and must not re-derive it from block metadata.
         self._pending_cpu_hits: dict[
-            str, tuple[tuple[list[KVCacheBlock], ...], int]
+            str, tuple[tuple[list[KVCacheBlock], ...], int, int]
         ] = {}
 
         # Store metadata
@@ -216,6 +250,7 @@ def __init__(
         # Event counters
         self._load_event_counter: int = 0
         self._store_event_counter: int = 0
+        self.boundary_store_stats = BoundaryStoreStats()
 
         # For TP/PP: track partial store completions across steps.
         # Events must be reported by all world_size workers before considered complete.
@@ -265,9 +300,13 @@ def _estimate_lazy_target_blocks(
         """GPU blocks to keep available (free/offloaded) per step in lazy mode."""
         WATERMARK_RATIO = 1.0  # Reserve larger space to avoid running out of GPU blocks
         target = 0
-        for g in kv_cache_config.kv_cache_groups:
+        for g in kv_cache_config.prefix_cacheable_groups:
             spec = g.kv_cache_spec
-            block_size = spec.block_size * cp_world_size
+            # Only full attention is sharded across DCP ranks; replicated specs
+            # (mamba, sliding window, chunked-local) keep their own block size.
+            block_size = spec.block_size * dcp_world_size_for_kv_cache_spec(
+                spec, cp_world_size
+            )
             if isinstance(spec, MambaSpec):
                 target += 2
             elif isinstance(spec, SlidingWindowSpec):
@@ -285,7 +324,6 @@ def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int | None, bool]:
         """Return (num_new_tokens, is_async) from consecutive CPU cache hits."""
-
         # Pins found CPU blocks so they survive LRU eviction until
         # update_state_after_alloc() consumes them. Any pin from an earlier
         # call on the same request (e.g. retry after a failed allocate_slots)
@@ -296,6 +334,20 @@ def get_num_new_matched_tokens(
         if request.skip_reading_prefix_cache:
             return 0, False
 
+        if num_computed_tokens % self.block_size != 0:
+            # Transfers are whole-block copies, so an external suffix cannot
+            # start in the middle of a destination block. When the GPU-side
+            # coordinator also serves fine-grained hits, the local prefix can
+            # land on a hash boundary that is not a scheduler-block boundary;
+            # such a request keeps its local hit and skips the external one.
+            logger.warning_once(
+                "SimpleCPU external lookup requires scheduler-block-aligned "
+                "local tokens, got %d tokens with scheduler block size %d.",
+                num_computed_tokens,
+                self.block_size,
+            )
+            return 0, False
+
         num_skipped_hashes = num_computed_tokens // self.hash_block_size
         remaining_hashes = request.block_hashes[num_skipped_hashes:]
 
@@ -318,6 +370,7 @@ def get_num_new_matched_tokens(
             self._pending_cpu_hits[request.request_id] = (
                 cpu_hit_blocks,
                 hit_length,
+                num_computed_tokens,
             )
             return hit_length, True
         return 0, False
@@ -371,20 +424,19 @@ def update_state_after_alloc(
             )
             return
 
-        cpu_hit_blocks_full, _ = pending
+        cpu_hit_blocks_full, _, num_computed_tokens = pending
 
-        # ``num_external_tokens`` is LCM-aligned (checked per-group below),
-        # so this counts whole scheduler-aligned chunks of incoming tokens.
-        num_blocks_to_load = num_external_tokens // self.block_size
-        assert num_blocks_to_load > 0
-        num_cached_fa_blocks = sum(
-            blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx]
+        assert num_computed_tokens % self.block_size == 0, (
+            "SimpleCPU external loads must start at a scheduler-block boundary"
+        )
+        # Fine-grained hits are hash-block aligned, not group-block aligned, so
+        # the per-group block count below rounds up instead of dividing exactly.
+        # Hash alignment is still required: it is what makes the accepted
+        # external suffix start on a boundary every group can address.
+        assert num_external_tokens % self.hash_block_size == 0, (
+            f"num_external_tokens={num_external_tokens} is not aligned to "
+            f"hash_block_size={self.hash_block_size}"
         )
-        num_computed_tokens = num_cached_fa_blocks * self.fa_block_size
-
-        # Build transfer pairs across all groups.
-        total_computed_tokens = num_computed_tokens + num_external_tokens
-        kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups
 
         # The scheduler may have accepted fewer blocks than
         # get_num_new_matched_tokens() reported.
@@ -393,14 +445,11 @@ def update_state_after_alloc(
         # the rest will be released along with the temp pin below.
         cpu_hit_blocks: list[list[KVCacheBlock]] = []
         for g in range(num_groups):
-            g_block_size = (
-                kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size
-            )
-            assert num_external_tokens % g_block_size == 0, (
-                f"num_external_tokens={num_external_tokens} not aligned to "
-                f"group {g} block_size={g_block_size}"
-            )
-            n_take_g = num_external_tokens // g_block_size
+            if g not in self.prefix_cacheable_group_ids:
+                cpu_hit_blocks.append([])
+                continue
+            g_block_size = self.group_block_sizes[g]
+            n_take_g = cdiv(num_external_tokens, g_block_size)
             cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])
 
         gpu_block_ids: list[int] = []
@@ -413,14 +462,8 @@ def update_state_after_alloc(
             if n_ext_g == 0:
                 continue
 
-            # Number of blocks in the computed range for this group.
-            g_block_size = (
-                kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size
-            )
-            n_computed_g = cdiv(total_computed_tokens, g_block_size)
-
-            # Back-trace: ext blocks sit at the tail of the computed range.
-            gpu_ext_start = n_computed_g - n_ext_g
+            g_block_size = self.group_block_sizes[g]
+            gpu_ext_start = num_computed_tokens // g_block_size
             group_gpu_ids = block_ids_by_group[g]
 
             for i, cpu_blk in enumerate(cpu_blocks_g):
@@ -598,8 +641,8 @@ def _prepare_eager_store_specs(
         Returns:
             (gpu_block_ids, cpu_block_ids, req_ids, block_meta) for the store
             event. ``block_meta`` is None when kv cache events are disabled.
-        """
 
+        """
         merged_gpu_block_ids: list[int] = []
         merged_cpu_block_ids: list[int] = []
         req_ids: list[str] = []
@@ -614,6 +657,72 @@ def _prepare_eager_store_specs(
         num_groups = len(self.cpu_kv_cache_config.kv_cache_groups)
         # Dedup against blocks already scheduled.
         in_flight = self._in_flight_store_gpu_blocks
+        num_free = cpu_block_pool.get_num_free_blocks()
+
+        block_state = scheduler_output.kv_connector_block_state
+        boundary_offloads = (
+            block_state.boundary_state_offloads if block_state is not None else {}
+        )
+        stats = self.boundary_store_stats
+        preempted_req_ids = scheduler_output.preempted_req_ids or ()
+        for req_id, entries in boundary_offloads.items():
+            # The scheduler drains handoffs without filtering by request
+            # liveness. A request that finished or was preempted in this step
+            # is giving up its blocks, which may already have been reallocated
+            # to another request; reading them would publish unrelated KV under
+            # a valid hash. Drop conservatively, as the mooncake store does.
+            store_state = self._reqs_to_store.get(req_id)
+            if (
+                store_state is None
+                or store_state.finished
+                or req_id in preempted_req_ids
+                or req_id in scheduler_output.finished_req_ids
+            ):
+                stats.published += len(entries)
+                stats.dropped_request_gone += len(entries)
+                continue
+            scheduled_for_req = False
+            # ``boundary_tokens`` is not needed here: the cache key comes from
+            # the handed-off block itself, not from the offered boundary.
+            for _, gpu_block_id, _ in entries:
+                stats.published += 1
+                gpu_block = gpu_block_pool.blocks[gpu_block_id]
+                admission = self._classify_store_candidate(gpu_block)
+                if admission is _StoreAdmission.NULL_BLOCK:
+                    stats.dropped_null_block += 1
+                    continue
+                if admission is _StoreAdmission.NOT_HASHED:
+                    stats.dropped_not_hashed += 1
+                    continue
+                if admission is _StoreAdmission.IN_FLIGHT:
+                    stats.skipped_in_flight += 1
+                    continue
+                if admission is _StoreAdmission.ALREADY_CACHED:
+                    stats.skipped_already_cached += 1
+                    continue
+                if num_free <= 0:
+                    stats.dropped_cpu_full += 1
+                    logger.warning_once(
+                        "SimpleCPU dropped a boundary-state handoff because "
+                        "the CPU block pool is full; this boundary will not be retried."
+                    )
+                    continue
+                cpu_block = cpu_block_pool.get_new_blocks(1)[0]
+                merged_gpu_block_ids.append(gpu_block_id)
+                merged_cpu_block_ids.append(cpu_block.block_id)
+                if merged_block_meta is not None:
+                    # Keep the metadata list index-parallel with the block ids.
+                    # The hand-off block's own hashes are resolved at completion,
+                    # so there is nothing to capture here; emit the event without
+                    # a payload, as the size-mismatch fallback already does.
+                    merged_block_meta.append({})
+                in_flight.add(gpu_block_id)
+                gpu_block_pool.touch([gpu_block])
+                num_free -= 1
+                stats.stored += 1
+                scheduled_for_req = True
+            if scheduled_for_req:
+                req_ids.append(req_id)
 
         for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
             state = self._reqs_to_store.get(req_id)
@@ -641,7 +750,7 @@ def _prepare_eager_store_specs(
                 self._select_eager_blocks_to_store(state, block_ids_by_group)
             )
 
-            # --- Phase 2: Batch allocate CPU blocks ---
+            # Batch allocate the CPU destinations.
             n_to_alloc = len(gpu_block_ids)
             if n_to_alloc > 0:
                 cpu_blocks_alloc = cpu_block_pool.get_new_blocks(n_to_alloc)
@@ -674,8 +783,23 @@ def _prepare_eager_store_specs(
             for g in range(num_groups):
                 state.num_stored_blocks[g] += advanced_per_group[g]
 
+        # A request can contribute both boundary and positional blocks to the
+        # same event. Completion tracking is set-based, so keep one request ID.
+        req_ids = list(dict.fromkeys(req_ids))
         return merged_gpu_block_ids, merged_cpu_block_ids, req_ids, merged_block_meta
 
+    def _cached_gpu_block(
+        self, resolved_hashes: "BlockHashList", block_idx: int, group_id: int
+    ) -> "KVCacheBlock | None":
+        """Return the GPU block still cached under this group's block hash."""
+        if block_idx >= len(resolved_hashes):
+            return None
+        assert self._gpu_block_pool is not None
+        blocks = self._gpu_block_pool.get_cached_block(
+            resolved_hashes[block_idx], [group_id]
+        )
+        return blocks[0] if blocks else None
+
     def _select_eager_blocks_to_store(
         self,
         state: StoreRequestState,
@@ -694,52 +818,68 @@ def _select_eager_blocks_to_store(
         )
         request = state.request
         confirmed_tokens = request.num_computed_tokens - request.num_output_placeholders
-        aligned_tokens = confirmed_tokens // self.block_size * self.block_size
+        # Truncate to the granularity a lookup can actually land on. With
+        # fine-grained hits that is hash_block_size; otherwise hits only land on
+        # the scheduler block (the group LCM). Using the LCM unconditionally
+        # would drop whole blocks that sit between the last LCM boundary and a
+        # reachable fine-grained boundary, so the other groups would hold that
+        # boundary and this one would not, and the joint hybrid lookup would
+        # reconcile to zero.
+        store_alignment = (
+            self.hash_block_size
+            if self.cpu_coordinator.enable_partial_hash_hits
+            else self.block_size
+        )
+        aligned_tokens = confirmed_tokens // store_alignment * store_alignment
         num_free = self.cpu_block_pool.get_num_free_blocks()
 
         for g, group_gpu_ids in enumerate(block_ids_by_group):
+            if g not in self.prefix_cacheable_group_ids:
+                continue
             if len(gpu_block_ids) >= num_free:
                 break
+            group_manager = self.cpu_coordinator.single_type_managers[g]
+            if not group_manager.has_positionally_stable_blocks:
+                continue
             # FIXME (yifan): handle CPU cache eviction, where
             # num_stored_blocks can be stale and omit evicted blocks in
             # the middle of the request.
-            group_size = (
-                self.cpu_kv_cache_config.kv_cache_groups[g].kv_cache_spec.block_size
-                * self.cp_world_size
-            )
+            group_size = self.group_block_sizes[g]
             ready = min(len(group_gpu_ids), aligned_tokens // group_size)
-            resolved_hashes = (
-                resolve_block_hashes(
-                    request.block_hashes, self.hash_block_size, group_size
-                )
-                if block_meta is not None
-                else None
+            resolved_hashes = resolve_block_hashes(
+                request.block_hashes, self.hash_block_size, group_size
             )
             curr_mm_idx = 0
             secondary_mm_idx = 0
             start = state.num_stored_blocks[g]
             for i, gpu_block_id in enumerate(group_gpu_ids[start:ready], start=start):
                 gpu_block = self._gpu_block_pool.blocks[gpu_block_id]
-                if gpu_block.is_null or gpu_block.block_hash is None:
-                    advanced_per_group[g] += 1
-                    continue
+                if gpu_block.is_null:
+                    # Sliding-window groups null pages that left the window
+                    # before the connector sees the block table, but the
+                    # retained prefix-cache tail stays hashed in the GPU free
+                    # queue; store it from there.
+                    recovered = self._cached_gpu_block(resolved_hashes, i, g)
+                    if recovered is None:
+                        advanced_per_group[g] += 1
+                        continue
+                    gpu_block = recovered
+                    gpu_block_id = gpu_block.block_id
                 if (
-                    gpu_block_id in self._in_flight_store_gpu_blocks
-                    or self.cpu_block_pool.cached_block_hash_to_block.get_one_block(
-                        gpu_block.block_hash
-                    )
-                    is not None
+                    self._classify_store_candidate(gpu_block)
+                    is not _StoreAdmission.READY
                 ):
                     advanced_per_group[g] += 1
                     continue
                 if len(gpu_block_ids) >= num_free:
                     break
+                primary_block_hash = gpu_block.block_hash
+                assert primary_block_hash is not None
                 gpu_block_ids.append(gpu_block_id)
                 advanced_per_group[g] += 1
                 if block_meta is not None:
                     token_start = i * group_size
                     token_end = token_start + group_size
-                    assert resolved_hashes is not None
                     parent_hash = (
                         None
                         if i == 0
@@ -750,7 +890,7 @@ def _select_eager_blocks_to_store(
                         request, token_start, token_end, curr_mm_idx
                     )
                     meta_by_hash = {
-                        gpu_block.block_hash: BlockStoreMeta(
+                        primary_block_hash: BlockStoreMeta(
                             token_ids=list(
                                 request.all_token_ids[token_start:token_end]
                             ),
@@ -766,7 +906,7 @@ def _select_eager_blocks_to_store(
                         block_hash = make_block_hash_with_group_id(
                             request.block_hashes[hash_idx], g
                         )
-                        if block_hash == gpu_block.block_hash:
+                        if block_hash is None or block_hash == primary_block_hash:
                             continue
                         hash_start = hash_idx * self.hash_block_size
                         hash_end = hash_start + self.hash_block_size
@@ -796,6 +936,25 @@ def _select_eager_blocks_to_store(
 
         return gpu_block_ids, advanced_per_group, block_meta
 
+    def _classify_store_candidate(self, gpu_block: "KVCacheBlock") -> _StoreAdmission:
+        if gpu_block.is_null:
+            return _StoreAdmission.NULL_BLOCK
+        if gpu_block.block_hash is None:
+            return _StoreAdmission.NOT_HASHED
+        if gpu_block.block_id in self._in_flight_store_gpu_blocks:
+            return _StoreAdmission.IN_FLIGHT
+        if (
+            self.cpu_block_pool.cached_block_hash_to_block.get_one_block(
+                gpu_block.block_hash
+            )
+            is not None
+        ):
+            return _StoreAdmission.ALREADY_CACHED
+        return _StoreAdmission.READY
+
+    def get_boundary_store_stats(self) -> BoundaryStoreStats:
+        return replace(self.boundary_store_stats)
+
     def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
         """Handle async transfer completions from worker.
 
@@ -886,12 +1045,7 @@ def _process_store_completion(
             if self.enable_kv_cache_events:
                 meta_by_hash = block_meta[i] if block_meta is not None else {}
                 primary_group_idx = get_group_id(primary_hash)
-                primary_spec = self.cpu_kv_cache_config.kv_cache_groups[
-                    primary_group_idx
-                ].kv_cache_spec
-                primary_block_size = primary_spec.block_size * (
-                    self.cp_world_size if not isinstance(primary_spec, MambaSpec) else 1
-                )
+                primary_block_size = self.group_block_sizes[primary_group_idx]
                 events_to_emit: list[
                     tuple[BlockHashWithGroupId, BlockStoreMeta | None, int]
                 ] = [
@@ -1028,6 +1182,21 @@ def _queue_finished_eager_store(
         if state is None or gpu_pool is None:
             return
         gpu_ids, _, block_meta = self._select_eager_blocks_to_store(state, block_ids)
+        partial_tail = self._find_fa_partial_tail_source(request, block_ids)
+        if (
+            partial_tail is not None
+            # Decode tokens can fill the boundary block, in which case the
+            # positional scan above already selected it. Appending again would
+            # allocate two CPU blocks for one GPU block.
+            and partial_tail not in gpu_ids
+            and len(gpu_ids) < self.cpu_block_pool.get_num_free_blocks()
+        ):
+            gpu_ids.append(partial_tail)
+            if block_meta is not None:
+                # The block's own hashes are registered at completion, so no
+                # capture is needed here; emit the event without a payload, as
+                # the size-mismatch fallback above already does.
+                block_meta.append({})
         if not gpu_ids:
             return
         cpu_blocks = self.cpu_block_pool.get_new_blocks(len(gpu_ids))
@@ -1037,9 +1206,64 @@ def _queue_finished_eager_store(
         self._in_flight_store_gpu_blocks.update(gpu_ids)
         gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids])
 
-    def _free_pending_cpu_hit(self, pending: tuple) -> None:
+    def _find_fa_partial_tail_source(
+        self, request: "Request", block_ids: tuple[list[int], ...]
+    ) -> int | None:
+        """Locate the full-attention prompt-tail block for a fine-grained hit.
+
+        Mirrors ``FullAttentionManager._cache_partial_tail_block``: only the
+        final prompt hash boundary is eligible, and boundaries that land on a
+        physical block edge are already covered by the positional scan.
+
+        Attention block tables are append-only, so the block is located
+        positionally, as the mooncake store and offloading connector also do;
+        no new core hand-off is required. The boundary key is registered at
+        completion along with the block's other hashes, so nothing has to be
+        captured here.
+        """
+        if not self.cpu_coordinator.enable_partial_hash_hits:
+            return None
+        assert self._gpu_block_pool is not None
+        boundary_tokens = (
+            request.num_prompt_tokens // self.hash_block_size * self.hash_block_size
+        )
+        if boundary_tokens == 0 or boundary_tokens > request.num_computed_tokens:
+            return None
+        if boundary_tokens % self.fa_block_size == 0:
+            return None
+        block_idx = boundary_tokens // self.fa_block_size
+        fa_block_ids = block_ids[self.fa_gidx]
+        if block_idx >= len(fa_block_ids):
+            return None
+        gpu_block_id = fa_block_ids[block_idx]
+        gpu_block = self._gpu_block_pool.blocks[gpu_block_id]
+        if gpu_block.is_null or gpu_block_id in self._in_flight_store_gpu_blocks:
+            return None
+        hash_idx = boundary_tokens // self.hash_block_size - 1
+        if hash_idx >= len(request.block_hashes):
+            return None
+        block_hash = make_block_hash_with_group_id(
+            request.block_hashes[hash_idx], self.fa_gidx
+        )
+        # Only worth copying when the boundary key is actually registered on
+        # the source block; completion replays the block's own hashes.
+        if (
+            gpu_block.block_hash != block_hash
+            and block_hash
+            not in self._gpu_block_pool.cached_block_hashes_by_block.get(
+                gpu_block_id, ()
+            )
+        ):
+            return None
+        if self.cpu_block_pool.cached_block_hash_to_block.get_one_block(block_hash):
+            return None
+        return gpu_block_id
+
+    def _free_pending_cpu_hit(
+        self, pending: tuple[tuple[list["KVCacheBlock"], ...], int, int]
+    ) -> None:
         """Release the temporary CPU block pin taken in get_num_new_matched_tokens()."""
-        cpu_hit_blocks, _ = pending
+        cpu_hit_blocks, _hit_length, _num_computed_tokens = pending
         blocks_to_free = [
             blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null
         ]
@@ -1113,7 +1337,6 @@ def reset(self) -> bool:
         the transfer finished, then release refs without caching abandoned
         store results.
         """
-
         self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks)
         for transfer in self._pending_finished_stores:
             self._release_transfer_refs(transfer)
@@ -1140,6 +1363,7 @@ def reset(self) -> bool:
             if event_idx in self._abandoned_store_event_to_blocks
         }
         self._cursor = None
+        self.boundary_store_stats = BoundaryStoreStats()
         # NOTE: _load_event_counter / _store_event_counter are not
         # reset as they are monotonic and must stay ahead of the workers
         # high-water marks to avoid event index collisions
diff --git a/vllm/v1/simple_kv_offload/metadata.py b/vllm/v1/simple_kv_offload/metadata.py
index 8c8d4511ee60..68ced9b89c73 100644
--- a/vllm/v1/simple_kv_offload/metadata.py
+++ b/vllm/v1/simple_kv_offload/metadata.py
@@ -14,8 +14,7 @@
 
 @dataclass
 class SimpleCPUOffloadMetadata(KVConnectorMetadata):
-    """
-    Metadata passed from scheduler to worker for CPU offload operations.
+    """Metadata passed from scheduler to worker for CPU offload operations.
 
     The worker receives flat block lists keyed by a monotonic event_idx.
     Job->req_id translation is handled by the scheduler-side manager
diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py
index 9efc544f2f5a..9fc16014c259 100644
--- a/vllm/v1/simple_kv_offload/worker.py
+++ b/vllm/v1/simple_kv_offload/worker.py
@@ -67,6 +67,7 @@ def __init__(
 
         # Metadata for the current step
         self._connector_metadata: SimpleCPUOffloadMetadata | None = None
+        self._store_submitted = False
 
         # Compute-done event recorded before each store; reused across steps
         # (get_finished runs once per step, copy queue is FIFO).
@@ -90,6 +91,7 @@ def register_kv_caches(
                 tensor (attention layers) or a list of tensors (Mamba layers
                 in hybrid models). All values are included for offloading
                 by resolving to their underlying raw storage.
+
         """
         if not kv_caches:
             logger.warning("No KV caches to offload.")
@@ -244,12 +246,15 @@ def _init_cpu_mode(
 
     def bind_connector_metadata(self, metadata: SimpleCPUOffloadMetadata) -> None:
         self._connector_metadata = metadata
+        self._store_submitted = False
         if metadata.load_event >= 0:
             self._pending_load_event_indices.add(metadata.load_event)
         if metadata.store_event >= 0:
             self._pending_store_event_indices.add(metadata.store_event)
 
     def clear_connector_metadata(self) -> None:
+        # No-forward steps skip the normal wait_for_save hook.
+        self.wait_for_save()
         self._connector_metadata = None
 
     def start_load_kv(self) -> None:
@@ -278,7 +283,11 @@ def wait_for_save(self) -> None:
         #45704 for the bug and #39306 for the srcAccessOrder rationale.
         """
         metadata = self._connector_metadata
-        if metadata is not None and metadata.store_gpu_blocks:
+        if (
+            metadata is not None
+            and metadata.store_gpu_blocks
+            and not self._store_submitted
+        ):
             backend = self._backend
             assert backend is not None
             if self._store_compute_done is None:
@@ -292,6 +301,7 @@ def wait_for_save(self) -> None:
                 events_list=self._store_events,
                 wait_event=self._store_compute_done,
             )
+            self._store_submitted = True
 
     def get_finished(
         self, finished_req_ids: set[str]
@@ -304,6 +314,7 @@ def get_finished(
             tuple of (finished_sending, finished_recving).
             - finished_sending: always None (stores use worker metadata).
             - finished_recving: req_ids whose loads have completed.
+
         """
         metadata = self._connector_metadata
         finished_recving: set[str] = set()
diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py
index 626dd36dea9b..d7d3e3adb02b 100644
--- a/vllm/v1/spec_decode/dflash.py
+++ b/vllm/v1/spec_decode/dflash.py
@@ -13,7 +13,7 @@
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
 from vllm.v1.spec_decode.utils import (
-    copy_and_expand_dflash_inputs_kernel,
+    _copy_and_expand_dflash_inputs,
     next_power_of_2,
 )
 
@@ -79,6 +79,13 @@ def __init__(
         self.dflash_causal = not dflash_has_any_non_causal(
             self.draft_model_config.hf_config
         )
+        _copy_and_expand_dflash_inputs.register_warmup(
+            block_table_stride=(self.max_model_len + self.block_size - 1)
+            // self.block_size,
+            parallel_drafting_token_id=self.parallel_drafting_token_id,
+            block_size=self.block_size,
+            num_speculative_tokens=self.num_speculative_tokens,
+        )
 
     @override
     def _create_draft_vllm_config(self) -> VllmConfig:
@@ -137,11 +144,9 @@ def set_inputs_first_pass(
         max_ctx_per_req = cad.max_query_len
         max_tokens_per_req = max_ctx_per_req + num_query_per_req
         BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req))
-        num_blocks = (max_tokens_per_req + BLOCK_SIZE - 1) // BLOCK_SIZE
-        grid = (batch_size, num_blocks)
 
         has_num_rejected = num_rejected_tokens_gpu is not None
-        copy_and_expand_dflash_inputs_kernel[grid](
+        _copy_and_expand_dflash_inputs(
             # Inputs
             next_token_ids_ptr=next_token_ids,
             target_positions_ptr=target_positions,
@@ -166,8 +171,9 @@ def set_inputs_first_pass(
             num_query_per_req=num_query_per_req,
             num_speculative_tokens=self.num_speculative_tokens,
             total_input_tokens=num_context,
-            BLOCK_SIZE=BLOCK_SIZE,
-            HAS_NUM_REJECTED=has_num_rejected,
+            max_tokens_per_req=max_tokens_per_req,
+            triton_block_size=BLOCK_SIZE,
+            has_num_rejected=has_num_rejected,
         )
 
         query_slot_mapping = self._slot_mapping_buffer[:num_query_total]
@@ -192,8 +198,6 @@ def set_inputs_first_pass(
                 torch.from_numpy(self.token_arange_np[: batch_size + 1]).clone()
                 * num_query_per_req
             ),
-            _seq_lens_cpu=None,
-            _num_computed_tokens_cpu=None,
             seq_lens_cpu_upper_bound=new_seq_lens_cpu_upper_bound,
             num_reqs=cad.num_reqs,
             num_actual_tokens=num_query_total,
@@ -215,8 +219,7 @@ def dummy_run(
         is_graph_capturing: bool = False,
         slot_mappings: dict[str, torch.Tensor] | None = None,
     ) -> None:
-        """
-        Key differences to default dummy_run:
+        """Key differences to default dummy_run:
         - Only one forward pass due to parallel drafting
         - DFlash uses context states as unpadded metadata, so hidden_states will
         use the unpadded num_tokens instead of num_input_tokens
diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py
index 6325c1bf5daa..18ae2700e937 100644
--- a/vllm/v1/spec_decode/extract_hidden_states.py
+++ b/vllm/v1/spec_decode/extract_hidden_states.py
@@ -106,6 +106,7 @@ def propose(
         The main purpose is to cache hidden states, not to speculate.
 
         Args:
+            num_speculative_tokens: Number of draft tokens to return per request
             sampled_token_ids: Sampled token IDs from the target model
             target_hidden_states: List of hidden state tensors from target model
                                 (one per aux hidden state layer)
@@ -117,6 +118,7 @@ def propose(
             Tuple of:
                 - Draft tokens matching sampled tokens, shape [batch_size, 1]
                 - KV connector output (if KV transfer is active), else None
+
         """
         assert num_speculative_tokens == self.num_speculative_tokens
         assert self.model is not None and isinstance(target_hidden_states, list)
@@ -318,14 +320,12 @@ def prepare_next_token_ids_padded(
         gpu_input_batch: InputBatch,
         discard_request_mask: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Prepare next token IDs for speculative decoding.
+        """Prepare next token IDs for speculative decoding.
 
         Since num_speculative_tokens == 1, sampled_token_ids has shape
         (batch_size, 1). For each request we either use the sampled token
         (if valid and not discarded) or a backup token from the request state.
         """
-
         # Precompute backup token IDs for discarded requests.
         num_reqs = gpu_input_batch.num_reqs
         for i in range(num_reqs):
@@ -359,6 +359,7 @@ def load_model(self, target_model: nn.Module) -> None:
         Args:
             target_model: The target model (passed for compatibility with
                          EagleProposer interface, but not used here)
+
         """
         # Get the target model's attention layers before loading draft model
         target_attn_layer_names = set(
diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py
index 4b078c2aea4a..e4718bb3d7d1 100644
--- a/vllm/v1/spec_decode/llm_base_proposer.py
+++ b/vllm/v1/spec_decode/llm_base_proposer.py
@@ -25,6 +25,7 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.model_loader.utils import get_draft_load_config
 from vllm.model_executor.models import (
     supports_multimodal,
     supports_multimodal_embeddings,
@@ -35,7 +36,6 @@
 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
 from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM
 from vllm.model_executor.models.qwen3_eagle3 import Eagle3Qwen3ForCausalLM
-from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
 from vllm.v1.attention.backend import CommonAttentionMetadata
@@ -52,10 +52,11 @@
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.spec_decode.utils import (
     PADDING_SLOT_ID,
+    _copy_and_expand_eagle_inputs,
+    _eagle_prepare_inputs_padded,
+    _eagle_prepare_next_token_padded,
+    _eagle_step_slot_mapping_metadata,
     compute_new_slot_mapping,
-    copy_and_expand_eagle_inputs_kernel,
-    eagle_prepare_inputs_padded_kernel,
-    eagle_prepare_next_token_padded_kernel,
     eagle_step_update_slot_mapping_and_metadata,
     extend_all_queries_by_N,
     next_power_of_2,
@@ -140,6 +141,23 @@ def __init__(
 
         self.max_batch_size = vllm_config.scheduler_config.max_num_seqs
         self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        _eagle_step_slot_mapping_metadata.register_warmup(
+            max_model_len=self.max_model_len,
+            max_batch_size=self.max_batch_size,
+        )
+        _eagle_prepare_inputs_padded.register_warmup(max_batch_size=self.max_batch_size)
+        _eagle_prepare_next_token_padded.register_warmup(
+            vocab_size=vllm_config.model_config.get_vocab_size(),
+            num_sampled_tokens_per_req=self.num_speculative_tokens + 1,
+            max_batch_size=self.max_batch_size,
+        )
+        if self.needs_extra_input_slots:
+            _copy_and_expand_eagle_inputs.register_warmup(
+                parallel_drafting_token_id=self.parallel_drafting_token_id,
+                num_padding_slots_per_request=self.extra_slots_per_request,
+                shift_input_ids=self.pass_hidden_states_to_model,
+                max_num_tokens=self.max_num_tokens,
+            )
         self.token_arange_np = np.arange(self.max_num_tokens, dtype=np.int32)
 
         # Can be specialized by methods like DFlash to reduce the limit
@@ -147,10 +165,7 @@ def __init__(
         self.max_positions = self.max_num_tokens
 
         # Multi-modal data support
-        self.mm_registry = MULTIMODAL_REGISTRY
-        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
-            vllm_config.model_config
-        )
+        self.supports_mm_inputs = vllm_config.model_config.supports_multimodal_inputs
 
         self.draft_attn_groups: list[AttentionGroup] = []
         self.kv_cache_gid: int = -1
@@ -237,7 +252,6 @@ def __init__(
         self.backup_next_token_ids = CpuGpuBuffer(
             self.max_batch_size,
             dtype=torch.int32,
-            pin_memory=PIN_MEMORY,
             device=device,
             with_numpy=True,
         )
@@ -669,9 +683,6 @@ def propose(
         # (i.e., not the first proposal).
         if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None:
             common_attn_metadata.seq_lens -= num_rejected_tokens_gpu
-            # Invalidate the CPU-side shadows to avoid H<>D sync.
-            common_attn_metadata._seq_lens_cpu = None
-            common_attn_metadata._num_computed_tokens_cpu = None
 
         block_size = self.block_size
         assert block_size > 0, "block_size has not been initialized."
@@ -798,10 +809,6 @@ def _update_positions_dependent_metadata(
             self.max_model_len,
         )
 
-        if common_attn_metadata._seq_lens_cpu is not None:
-            common_attn_metadata._seq_lens_cpu += 1
-        if common_attn_metadata._num_computed_tokens_cpu is not None:
-            common_attn_metadata._num_computed_tokens_cpu += 1
         if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
             common_attn_metadata.seq_lens_cpu_upper_bound += 1
 
@@ -879,13 +886,12 @@ def set_inputs_first_pass(
             )
 
             # Kernel grid: one program per request (row)
-            grid = (batch_size, num_blocks)
             query_start_loc = cad.query_start_loc
             query_end_loc = cad.query_start_loc[1:] - 1
             if num_rejected_tokens_gpu is not None:
                 query_end_loc = query_end_loc - num_rejected_tokens_gpu
 
-            copy_and_expand_eagle_inputs_kernel[grid](
+            _copy_and_expand_eagle_inputs(
                 # (Padded) Inputs from the target model
                 target_token_ids_ptr=target_token_ids,
                 target_positions_ptr=target_positions,
@@ -908,6 +914,8 @@ def set_inputs_first_pass(
                 num_padding_slots_per_request=self.extra_slots_per_request,
                 shift_input_ids=self.pass_hidden_states_to_model,
                 BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,
+                batch_size=batch_size,
+                num_blocks=num_blocks,
             )
             if self.pass_hidden_states_to_model:
                 assert self.parallel_drafting_hidden_state_tensor is not None
@@ -1013,8 +1021,7 @@ def prepare_next_token_ids_cpu(
         gpu_input_batch: InputBatch,
         num_scheduled_tokens: dict[str, int],
     ) -> torch.Tensor:
-        """
-        This function is used to prepare the inputs for speculative decoding.
+        """This function is used to prepare the inputs for speculative decoding.
         It calculates the next token ids for each request based on the sampled
         token ids from the CPU. If a request has no sampled token ids (e.g.,
         during the initial decoding steps), it falls back to using the request
@@ -1046,8 +1053,7 @@ def prepare_next_token_ids_padded(
         gpu_input_batch: InputBatch,
         discard_request_mask: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        This function is used to prepare the inputs for speculative decoding.
+        """This function is used to prepare the inputs for speculative decoding.
         It calculates the next token ids and the number of valid sampled tokens
         for each request, considering the "discarded" requests whose next token
         is not sampled and comes from `request.get_token_id()` instead. This is denoted
@@ -1072,11 +1078,10 @@ def prepare_next_token_ids_padded(
         valid_sampled_tokens_count = next_token_ids.new_empty(batch_size)
 
         # Kernel grid: one program per request (row)
-        grid = (batch_size,)
 
         # Find the next power of 2 for block sizes
         BLOCK_SIZE_TOKENS = next_power_of_2(num_tokens)
-        eagle_prepare_next_token_padded_kernel[grid](
+        _eagle_prepare_next_token_padded(
             sampled_token_ids,
             discard_request_mask,
             backup_tokens_gpu,
@@ -1097,8 +1102,7 @@ def prepare_inputs_padded(
         spec_decode_metadata: SpecDecodeMetadata,
         valid_sampled_tokens_count: torch.Tensor,
     ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:
-        """
-        This function is used to prepare the inputs for speculative decoding
+        """This function is used to prepare the inputs for speculative decoding
         It updates the common_attn_metadata for speculative decoding,
         but does not consider the rejected tokens. Instead, all tokens
         are included as inputs to the speculator, with the rejected tokens
@@ -1115,8 +1119,7 @@ def prepare_inputs_padded(
             (num_reqs,), dtype=torch.int32, device=device
         )
 
-        grid = (num_reqs,)
-        eagle_prepare_inputs_padded_kernel[grid](
+        _eagle_prepare_inputs_padded(
             spec_decode_metadata.cu_num_draft_tokens,
             valid_sampled_tokens_count,
             common_attn_metadata.query_start_loc,
@@ -1134,8 +1137,6 @@ def prepare_inputs_padded(
             query_start_loc=common_attn_metadata.query_start_loc,
             seq_lens=common_attn_metadata.seq_lens,
             query_start_loc_cpu=query_start_loc_cpu,
-            _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,
-            _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
             seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound,
             num_reqs=common_attn_metadata.num_reqs,
             num_actual_tokens=total_num_tokens,
@@ -1159,8 +1160,7 @@ def prepare_inputs(
         sampled_token_ids: list[list[int]],
         num_draft_tokens: list[int],
     ) -> tuple[CommonAttentionMetadata, torch.Tensor]:
-        """
-        This function is used to prepare the inputs for speculative decoding.
+        """This function is used to prepare the inputs for speculative decoding.
         It updates to the common_attn_metadata to account for the rejected
         tokens (and newly sampled tokens). It also returns the token indices
         of the tokens that should be fed to the speculator.
@@ -1245,8 +1245,6 @@ def prepare_inputs(
             query_start_loc=async_tensor_h2d(new_query_start_loc_cpu, device=device),
             seq_lens=async_tensor_h2d(new_seq_lens_cpu, device=device),
             query_start_loc_cpu=new_query_start_loc_cpu,
-            _seq_lens_cpu=new_seq_lens_cpu,
-            _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
             seq_lens_cpu_upper_bound=new_seq_lens_cpu,
             num_reqs=common_attn_metadata.num_reqs,
             num_actual_tokens=total_num_tokens,
@@ -1304,8 +1302,7 @@ def _create_draft_vllm_config(self) -> VllmConfig:
         return base
 
     def _get_model(self) -> nn.Module:
-        """
-        Default method to call get_model(). Can be overridden by subclasses which
+        """Default method to call get_model(). Can be overridden by subclasses which
         need to customize model loading.
         """
         from vllm.compilation.backends import set_model_tag
@@ -1315,7 +1312,7 @@ def _get_model(self) -> nn.Module:
             model = get_model(
                 vllm_config=draft_vllm_config,
                 model_config=self.speculative_config.draft_model_config,
-                load_config=self.speculative_config.draft_load_config,
+                load_config=get_draft_load_config(draft_vllm_config),
             )
         return model
 
@@ -1417,8 +1414,7 @@ def load_model(self, target_model: nn.Module) -> None:
                 self.parallel_drafting_hidden_state_tensor.copy_(flat_mask)
 
     def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None:
-        """
-        Some draft models may not have their own embedding layers, and some may
+        """Some draft models may not have their own embedding layers, and some may
         have a duplicate copy of the target model's embedding layers. In these cases,
         we share the target model's embedding layers with the draft model to save
         memory.
@@ -1510,8 +1506,7 @@ def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None:
             )
 
     def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:
-        """
-        Some draft models may not have their own LM head, and some may have a
+        """Some draft models may not have their own LM head, and some may have a
         duplicate copy of the target model's LM head. In these cases, we share
         the target model's LM head with the draft model to save memory.
         """
@@ -1675,8 +1670,7 @@ def dummy_run(
                 self.model(**kwargs)
 
     def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:
-        """
-        Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary
+        """Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary
         hidden states and directly uses the last layer output just like eagle1.
         They might indicate this by setting "use_aux_hidden_state" to False
         inside the "eagle_config" dict of their hf_config.
@@ -1691,8 +1685,7 @@ def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:
         return use_aux_hidden_state
 
     def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
-        """
-        Validate that all drafting layers belong to the same KVCacheGroup.
+        """Validate that all drafting layers belong to the same KVCacheGroup.
         Need this assumption to ensure all drafting layers can use the
         same AttentionMetadata.
         May extend to multiple AttentionMetadata in the future.
@@ -1718,8 +1711,7 @@ def initialize_attn_backend(
         kv_cache_config: KVCacheConfig,
         kernel_block_sizes: list[int] | None = None,
     ) -> None:
-        """
-        Initialize AttentionGroups for draft layers using kv_cache_config.
+        """Initialize AttentionGroups for draft layers using kv_cache_config.
         Called from the model runner's initialize_metadata_builders.
         """
         all_attn_layers = get_layers_from_vllm_config(
diff --git a/vllm/v1/spec_decode/medusa.py b/vllm/v1/spec_decode/medusa.py
index 7adf7cff5f77..7c926ba8372f 100644
--- a/vllm/v1/spec_decode/medusa.py
+++ b/vllm/v1/spec_decode/medusa.py
@@ -16,9 +16,7 @@
 
 
 class MedusaProposer:
-    """
-    Medusa proposer class for generating token sequences
-    """
+    """Medusa proposer class for generating token sequences."""
 
     def __init__(
         self,
diff --git a/vllm/v1/spec_decode/metrics.py b/vllm/v1/spec_decode/metrics.py
index a3ccfb29e737..826dfbd16a48 100644
--- a/vllm/v1/spec_decode/metrics.py
+++ b/vllm/v1/spec_decode/metrics.py
@@ -187,12 +187,12 @@ class SpecDecodingProm:
 
       1 + (
       rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) /
-      rate(vllm:spec_decode_num_drafts[$interval]))
+      rate(vllm:spec_decode_num_drafts_total[$interval]))
 
     A per-position acceptance rate vector can be computed using
 
-      vllm:spec_decode_num_accepted_tokens_per_pos[$interval] /
-      vllm:spec_decode_num_drafts[$interval]
+      vllm:spec_decode_num_accepted_tokens_per_pos_total[$interval] /
+      vllm:spec_decode_num_drafts_total[$interval]
     """
 
     _counter_cls = prometheus_client.Counter
diff --git a/vllm/v1/spec_decode/ngram_proposer.py b/vllm/v1/spec_decode/ngram_proposer.py
index e0240d0e66b8..5e0d2cfc3c43 100644
--- a/vllm/v1/spec_decode/ngram_proposer.py
+++ b/vllm/v1/spec_decode/ngram_proposer.py
@@ -72,6 +72,8 @@ def batch_propose(
         """Batch version of ngram proposer using numba for acceleration.
 
         Args:
+            num_requests:
+                Number of requests in the batch.
             valid_ngram_requests:
                 Set of indices of requests that need ngram proposals.
             num_tokens_no_spec:
@@ -87,6 +89,7 @@ def batch_propose(
             list[list[int]]:
                 A list where each element is a list of proposed
                 token IDs for the corresponding request.
+
         """
         draft_token_ids: list[list[int]] = []
 
@@ -211,8 +214,7 @@ def _find_longest_matched_ngram_and_propose_tokens(
     max_model_len: int,
     k: int,
 ) -> np.ndarray:
-    """
-    Find the longest n-gram which matches the suffix of the given tokens
+    """Find the longest n-gram which matches the suffix of the given tokens
     whose length is within [min_ngram, max_ngram] (inclusive).
 
     If found, we will extract k right after the matched ngram.
diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py
index 2de5c3be007f..061a8f5ea93b 100644
--- a/vllm/v1/spec_decode/ngram_proposer_gpu.py
+++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-GPU-accelerated N-gram proposer using fully async PyTorch tensor operations.
+"""GPU-accelerated N-gram proposer using fully async PyTorch tensor operations.
 
 This version uses a fully vectorized approach with unfold and argmax for
 finding the first match across all sequences in parallel.
@@ -52,8 +51,7 @@ def _find_first_and_extract_all_n_parallel(
         max_ngram_len: int,
         num_draft_tokens: int,
     ) -> torch.Tensor:
-        """
-        Find suffix n-gram matches and extract following tokens.
+        """Find suffix n-gram matches and extract following tokens.
         Searches for the earliest prior occurrence of the trailing n-gram,
         tries multiple lengths, and picks the longest valid match.
 
@@ -66,6 +64,7 @@ def _find_first_and_extract_all_n_parallel(
 
         Returns:
             Draft token predictions; -1 means invalid/no match.
+
         """
         batch_size = token_ids.shape[0]
         max_seq_len = token_ids.shape[1]
@@ -165,8 +164,7 @@ def forward(
         token_ids_gpu: torch.Tensor,
         combined_mask: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for N-gram proposal using GPU tensor operations.
+        """Forward pass for N-gram proposal using GPU tensor operations.
 
         Args:
             num_tokens_no_spec: Number of tokens for each sequence [batch_size]
@@ -177,8 +175,8 @@ def forward(
             draft_tokens: [batch_size, k] on GPU
             num_valid_draft_tokens: [batch_size] int32 on GPU, count of
                 leading valid (non -1) tokens per request.
-        """
 
+        """
         device = token_ids_gpu.device
 
         # Infer batch size to preserve dynamic shape.
@@ -282,8 +280,7 @@ def _generate_dummy_data(
         pattern_len: int,
         device: str = "cuda",
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Generate random test data with n-gram repetitions.
+        """Generate random test data with n-gram repetitions.
 
         Args:
             batch_size: Number of sequences in the batch
@@ -296,6 +293,7 @@ def _generate_dummy_data(
             num_tokens: [batch_size] tensor
             sampled_flags: [batch_size] bool tensor
             valid_mask: [batch_size] bool tensor
+
         """
         token_ids = torch.zeros(
             batch_size,
@@ -321,8 +319,7 @@ def propose(
         valid_sampled_token_ids_gpu: torch.Tensor,  # [batch_size, num_spec_tokens + 1]
         valid_sampled_tokens_count: torch.Tensor,  # [batch_size]
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Propose draft tokens using GPU-accelerated n-gram matching.
+        """Propose draft tokens using GPU-accelerated n-gram matching.
 
         Scatter sampled tokens into `token_ids_gpu`, compute temporary
         updated lengths, then run the kernel.
@@ -338,6 +335,7 @@ def propose(
             draft_tokens: Proposed draft token IDs [batch_size, k]
             num_valid_draft_tokens: Count of leading valid draft tokens
                 per request [batch_size]
+
         """
         assert num_speculative_tokens == self.k
         assert token_ids_gpu.device == self.device
@@ -398,8 +396,7 @@ def update_token_ids_ngram(
         num_tokens_no_spec: torch.Tensor,
         discard_request_mask: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Prepare speculative decoding inputs on device:
+        """Prepare speculative decoding inputs on device:
         compute next token ids and valid counts, honoring discarded requests
         and rejected tokens, without CPU-GPU sync.
         """
@@ -485,6 +482,7 @@ def update_scheduler_for_invalid_drafts(
         num_valid_draft_tokens_cpu: CPU buffer of valid draft counts.
         scheduler_output: Scheduler metadata to update in-place.
         req_id_to_index: Request-id to batch-index mapping.
+
     """
     req_data = scheduler_output.scheduled_cached_reqs
     num_valid_draft_tokens_event.synchronize()
@@ -652,9 +650,7 @@ def copy_num_valid_draft_tokens(
     num_valid_draft_tokens: torch.Tensor | None,
     batch_size: int,
 ) -> None:
-    """
-    Async D2H copy of per-request valid draft counts.
-    """
+    """Async D2H copy of per-request valid draft counts."""
     if num_valid_draft_tokens is None:
         return
 
diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py
index 47821cc9b2d5..7a73473adc2f 100644
--- a/vllm/v1/spec_decode/step3p5.py
+++ b/vllm/v1/spec_decode/step3p5.py
@@ -155,7 +155,6 @@ def build_per_group_and_layer_attn_metadata(
 
     def _maybe_share_lm_head(self, target_language_model: torch.nn.Module) -> None:
         """Step3.5 MTP uses the lm_head stored in each MTP layer."""
-
         # The base MTP path shares target lm_head into shared_head.head.
         # Step3.5 checkpoints carry per-MTP-layer shared_head weights.
         return
@@ -386,8 +385,6 @@ def propose(
 
         if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None:
             common_attn_metadata.seq_lens -= num_rejected_tokens_gpu
-            common_attn_metadata._seq_lens_cpu = None
-            common_attn_metadata._num_computed_tokens_cpu = None
 
         block_size = self.block_size
         assert block_size > 0, "block_size has not been initialized."
diff --git a/vllm/v1/spec_decode/suffix_decoding.py b/vllm/v1/spec_decode/suffix_decoding.py
index 66137a006316..5d100ce9ed97 100644
--- a/vllm/v1/spec_decode/suffix_decoding.py
+++ b/vllm/v1/spec_decode/suffix_decoding.py
@@ -7,8 +7,7 @@
 
 
 class SuffixDecodingProposer:
-    """
-    Speculative decoding proposer for Suffix Decoding (https://arxiv.org/pdf/2411.04975).
+    """Speculative decoding proposer for Suffix Decoding (https://arxiv.org/pdf/2411.04975).
     This class imports and uses the official implementation from Arctic Inference
     (https://github.com/snowflakedb/ArcticInference).
     """
diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py
index 65b9408a8901..471d51d355e5 100644
--- a/vllm/v1/spec_decode/utils.py
+++ b/vllm/v1/spec_decode/utils.py
@@ -1,7 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import (
+    WarmupChoices,
+    WarmupIntRange,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.utils import (
@@ -39,8 +50,7 @@ def eagle_step_slot_mapping_metadata_kernel(
     PAD_ID: tl.constexpr,
     batch_size,
 ):
-    """
-    Fused kernel for EAGLE autoregressive step: updates positions, slot mapping,
+    """Fused kernel for EAGLE autoregressive step: updates positions, slot mapping,
     and sequence lengths in a single kernel to reduce launch overhead.
 
     Launched with input_batch_size threads. Threads with req_idx >= batch_size
@@ -85,6 +95,49 @@ def eagle_step_slot_mapping_metadata_kernel(
     tl.store(seq_lens_ptr + req_idx, new_seq_len)
 
 
+def _eagle_step_warmup_inputs(*, max_model_len: int, max_batch_size: int):
+    block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2)
+    n_blocks_per_req = triton.cdiv(max_model_len, block_size)
+    return dict(
+        positions=TritonWarmupTensor(torch.int64),
+        block_table=TritonWarmupTensor(
+            torch.int32, shape=(1, n_blocks_per_req), strides=(n_blocks_per_req, 1)
+        ),
+        block_table_stride=n_blocks_per_req,
+        seq_lens=TritonWarmupTensor(torch.int32),
+        out_clamped_positions=TritonWarmupTensor(torch.int64),
+        out_slot_mapping=TritonWarmupTensor(torch.int64),
+        block_size=block_size,
+        max_model_len=max_model_len,
+        n_blocks_per_req=n_blocks_per_req,
+        PAD_ID=PADDING_SLOT_ID,
+        batch_size=WarmupIntRange(1, max_batch_size + 1),
+        input_batch_size=1,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=eagle_step_slot_mapping_metadata_kernel,
+    warmup_inputs=_eagle_step_warmup_inputs,
+)
+def _eagle_step_slot_mapping_metadata(
+    positions: torch.Tensor,
+    block_table: torch.Tensor,
+    block_table_stride: int,
+    seq_lens: torch.Tensor,
+    out_clamped_positions: torch.Tensor,
+    out_slot_mapping: torch.Tensor,
+    *,
+    block_size: int,
+    max_model_len: int,
+    n_blocks_per_req: int,
+    PAD_ID: int,
+    batch_size: int,
+    input_batch_size: int,
+) -> DispatchSpec:
+    return (input_batch_size,), {}
+
+
 def eagle_step_update_slot_mapping_and_metadata(
     positions_1d: torch.Tensor,
     block_table_tensor: torch.Tensor,
@@ -95,8 +148,7 @@ def eagle_step_update_slot_mapping_and_metadata(
     out_slot_mapping: torch.Tensor,
     input_batch_size: int | None = None,
 ) -> None:
-    """
-    Fused update of slot mapping and metadata for one EAGLE autoregressive step.
+    """Fused update of slot mapping and metadata for one EAGLE autoregressive step.
     Updates seq_lens in place. Writes to out_clamped_positions and out_slot_mapping.
 
     When input_batch_size > batch_size, threads beyond batch_size write
@@ -112,13 +164,14 @@ def eagle_step_update_slot_mapping_and_metadata(
         out_slot_mapping: [input_batch_size] output buffer for slot mapping
         input_batch_size: total batch size including cudagraph padding;
             defaults to batch_size (no padding)
+
     """
     batch_size = positions_1d.shape[0]
     if input_batch_size is None:
         input_batch_size = batch_size
 
     n_blocks_per_req = block_table_tensor.shape[1]
-    eagle_step_slot_mapping_metadata_kernel[(input_batch_size,)](
+    _eagle_step_slot_mapping_metadata(
         positions_1d,
         block_table_tensor,
         block_table_tensor.stride(0),
@@ -130,6 +183,7 @@ def eagle_step_update_slot_mapping_and_metadata(
         n_blocks_per_req=n_blocks_per_req,
         PAD_ID=PADDING_SLOT_ID,
         batch_size=batch_size,
+        input_batch_size=input_batch_size,
     )
 
 
@@ -142,8 +196,7 @@ def eagle_prepare_inputs_padded_kernel(
     num_rejected_tokens_gpu_ptr,  # [num_reqs] (output)
     num_reqs,  # tl.int32
 ):
-    """
-    Fused kernel for Eagle prepare_input_padded. This kernel computes the
+    """Fused kernel for Eagle prepare_input_padded. This kernel computes the
     token index to sample for each request, taking into account the number
     of draft tokens and the number of valid sampled tokens (which is one more than
     the number of accepted tokens).
@@ -175,6 +228,33 @@ def eagle_prepare_inputs_padded_kernel(
     tl.store(num_rejected_tokens_gpu_ptr + req_idx, num_rejected_tokens)
 
 
+def _eagle_prepare_inputs_warmup_inputs(*, max_batch_size: int):
+    int32 = TritonWarmupTensor(torch.int32)
+    return dict(
+        cu_num_draft_tokens=int32,
+        valid_sampled_tokens_count=int32,
+        query_start_loc_gpu=int32,
+        token_indices_to_sample=int32,
+        num_rejected_tokens_gpu=int32,
+        num_reqs=WarmupIntRange(1, max_batch_size + 1),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=eagle_prepare_inputs_padded_kernel,
+    warmup_inputs=_eagle_prepare_inputs_warmup_inputs,
+)
+def _eagle_prepare_inputs_padded(
+    cu_num_draft_tokens: torch.Tensor,
+    valid_sampled_tokens_count: torch.Tensor,
+    query_start_loc_gpu: torch.Tensor,
+    token_indices_to_sample: torch.Tensor,
+    num_rejected_tokens_gpu: torch.Tensor,
+    num_reqs: int,
+) -> DispatchSpec:
+    return (num_reqs,), {}
+
+
 @triton.jit
 def eagle_prepare_next_token_padded_kernel(
     sampled_token_ids_ptr,  # [num_reqs, num_sampled_tokens_per_req]
@@ -188,8 +268,7 @@ def eagle_prepare_next_token_padded_kernel(
     stride_sampled_token_ids,  # tl.int32 (stride for dim 0)
     BLOCK_SIZE_TOKENS: tl.constexpr,  # Power-of-2 >= num_sampled_tokens_per_req
 ):
-    """
-    Fused kernel for Eagle prepare_next_token_ids_padded. This kernel computes the
+    """Fused kernel for Eagle prepare_next_token_ids_padded. This kernel computes the
     number of valid (1 + accepted) tokens for each request, and the corresponding
     "next" token id to sample from during speculative decoding. This is the
     "last accepted token" from the sampled tokens, or the backup token if no
@@ -238,6 +317,49 @@ def eagle_prepare_next_token_padded_kernel(
         tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)
 
 
+def _eagle_prepare_next_token_warmup_inputs(
+    *, vocab_size: int, num_sampled_tokens_per_req: int, max_batch_size: int
+):
+    int32 = TritonWarmupTensor(torch.int32)
+    sampled_width: Any = WarmupIntRange(1, num_sampled_tokens_per_req + 1)
+    return dict(
+        sampled_token_ids=TritonWarmupTensor(
+            torch.int32,
+            shape=(1, sampled_width),
+            strides=(sampled_width, 1),
+        ),
+        discard_request_mask=TritonWarmupTensor(torch.bool),
+        backup_next_token_ids=int32,
+        next_token_ids=int32,
+        valid_sampled_tokens_count=int32,
+        vocab_size=vocab_size,
+        num_sampled_tokens_per_req=sampled_width,
+        num_reqs=WarmupIntRange(1, max_batch_size + 1),
+        stride_sampled_token_ids=sampled_width,
+        BLOCK_SIZE_TOKENS=next_power_of_2(sampled_width),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=eagle_prepare_next_token_padded_kernel,
+    warmup_inputs=_eagle_prepare_next_token_warmup_inputs,
+)
+def _eagle_prepare_next_token_padded(
+    sampled_token_ids: torch.Tensor,
+    discard_request_mask: torch.Tensor,
+    backup_next_token_ids: torch.Tensor,
+    next_token_ids: torch.Tensor,
+    valid_sampled_tokens_count: torch.Tensor,
+    vocab_size: int,
+    num_sampled_tokens_per_req: int,
+    num_reqs: int,
+    stride_sampled_token_ids: int,
+    *,
+    BLOCK_SIZE_TOKENS: int,
+) -> DispatchSpec:
+    return (num_reqs,), {}
+
+
 def compute_new_slot_mapping(
     cad: CommonAttentionMetadata,
     new_positions: torch.Tensor,
@@ -276,8 +398,7 @@ def extend_all_queries_by_N(
     arange: torch.Tensor,
     new_slot_mapping: torch.Tensor,
 ) -> CommonAttentionMetadata:
-    """
-    Creates a new CommonAttentionMetadata with all query lengths increased by N.
+    """Creates a new CommonAttentionMetadata with all query lengths increased by N.
     Also all seq lens are increased by N.
     This is useful e.g. in speculative decoding with parallel drafting, where we
     extend each sequence by N tokens and predict all tokens in one pass.
@@ -328,8 +449,7 @@ def copy_and_expand_eagle_inputs_kernel(
     shift_input_ids,  # tl.bool
     BLOCK_SIZE_TOKENS: tl.constexpr,  # Blocks along token dim to handle prefills
 ):
-    """
-    Copy and expand inputs from the target model to the drafting buffers for Eagle
+    """Copy and expand inputs from the target model to the drafting buffers for Eagle
     speculative decoding. This kernel handles padding slots and parallel drafting
     tokens, if enabled.
     """
@@ -454,6 +574,68 @@ def copy_and_expand_eagle_inputs_kernel(
     )
 
 
+def _copy_and_expand_eagle_warmup_inputs(
+    *,
+    parallel_drafting_token_id: int,
+    num_padding_slots_per_request: int,
+    shift_input_ids: bool,
+    max_num_tokens: int,
+):
+    int32 = TritonWarmupTensor(torch.int32)
+    int64 = TritonWarmupTensor(torch.int64)
+    boolean = TritonWarmupTensor(torch.bool)
+    return dict(
+        target_token_ids_ptr=int32,
+        target_positions_ptr=int64,
+        next_token_ids_ptr=int32,
+        out_input_ids_ptr=int32,
+        out_positions_ptr=int64,
+        out_is_rejected_token_mask_ptr=boolean,
+        out_is_masked_token_mask_ptr=boolean,
+        out_new_token_indices_ptr=int32,
+        out_hidden_state_mapping_ptr=int32,
+        query_start_loc_ptr=int32,
+        query_end_loc_ptr=int32,
+        padding_token_id=0,
+        parallel_drafting_token_id=parallel_drafting_token_id,
+        total_input_tokens=WarmupIntRange(1, max_num_tokens + 1),
+        num_padding_slots_per_request=num_padding_slots_per_request,
+        shift_input_ids=shift_input_ids,
+        BLOCK_SIZE_TOKENS=WarmupIntRange(1, 257, advance=lambda value: value * 2),
+        batch_size=1,
+        num_blocks=1,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=copy_and_expand_eagle_inputs_kernel,
+    warmup_inputs=_copy_and_expand_eagle_warmup_inputs,
+)
+def _copy_and_expand_eagle_inputs(
+    target_token_ids_ptr: torch.Tensor,
+    target_positions_ptr: torch.Tensor,
+    next_token_ids_ptr: torch.Tensor,
+    out_input_ids_ptr: torch.Tensor,
+    out_positions_ptr: torch.Tensor,
+    out_is_rejected_token_mask_ptr: torch.Tensor,
+    out_is_masked_token_mask_ptr: torch.Tensor,
+    out_new_token_indices_ptr: torch.Tensor,
+    out_hidden_state_mapping_ptr: torch.Tensor,
+    query_start_loc_ptr: torch.Tensor,
+    query_end_loc_ptr: torch.Tensor,
+    padding_token_id: int,
+    parallel_drafting_token_id: int,
+    total_input_tokens: int,
+    num_padding_slots_per_request: int,
+    shift_input_ids: bool,
+    *,
+    BLOCK_SIZE_TOKENS: int,
+    batch_size: int,
+    num_blocks: int,
+) -> DispatchSpec:
+    return (batch_size, num_blocks), {}
+
+
 @triton.jit
 def copy_and_expand_dflash_inputs_kernel(
     # Inputs
@@ -481,8 +663,7 @@ def copy_and_expand_dflash_inputs_kernel(
     BLOCK_SIZE: tl.constexpr,
     HAS_NUM_REJECTED: tl.constexpr = False,
 ):
-    """
-    Fused kernel for DFlash first-pass input setup.
+    """Fused kernel for DFlash first-pass input setup.
 
     Per request, this kernel:
       1. Copies context positions from target_positions to
@@ -562,6 +743,75 @@ def copy_and_expand_dflash_inputs_kernel(
     )
 
 
+def _copy_and_expand_dflash_warmup_inputs(
+    *,
+    block_table_stride: int,
+    parallel_drafting_token_id: int,
+    block_size: int,
+    num_speculative_tokens: int,
+):
+    int32 = TritonWarmupTensor(torch.int32)
+    int64 = TritonWarmupTensor(torch.int64)
+    triton_block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2)
+    has_num_rejected = WarmupChoices(False, True)
+    return dict(
+        next_token_ids_ptr=int32,
+        target_positions_ptr=int64,
+        out_input_ids_ptr=int32,
+        out_context_positions_ptr=int64,
+        out_query_positions_ptr=int64,
+        out_context_slot_mapping_ptr=int64,
+        out_query_slot_mapping_ptr=int64,
+        out_token_indices_ptr=int32,
+        block_table_ptr=int32,
+        query_start_loc_ptr=int32,
+        num_rejected_tokens_ptr=(int32 if has_num_rejected else 0),
+        block_table_stride=block_table_stride,
+        parallel_drafting_token_id=parallel_drafting_token_id,
+        block_size=block_size,
+        num_query_per_req=1 + num_speculative_tokens,
+        num_speculative_tokens=num_speculative_tokens,
+        total_input_tokens=WarmupChoices(1, 2, 16),
+        max_tokens_per_req=triton_block_size,
+        triton_block_size=triton_block_size,
+        has_num_rejected=has_num_rejected,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=copy_and_expand_dflash_inputs_kernel,
+    warmup_inputs=_copy_and_expand_dflash_warmup_inputs,
+)
+def _copy_and_expand_dflash_inputs(
+    *,
+    next_token_ids_ptr: torch.Tensor,
+    target_positions_ptr: torch.Tensor,
+    out_input_ids_ptr: torch.Tensor,
+    out_context_positions_ptr: torch.Tensor,
+    out_query_positions_ptr: torch.Tensor,
+    out_context_slot_mapping_ptr: torch.Tensor,
+    out_query_slot_mapping_ptr: torch.Tensor,
+    out_token_indices_ptr: torch.Tensor,
+    block_table_ptr: torch.Tensor,
+    query_start_loc_ptr: torch.Tensor,
+    num_rejected_tokens_ptr: torch.Tensor | int,
+    block_table_stride: int,
+    parallel_drafting_token_id: int,
+    block_size: int,
+    num_query_per_req: int,
+    num_speculative_tokens: int,
+    total_input_tokens: int,
+    max_tokens_per_req: int,
+    triton_block_size: int,
+    has_num_rejected: bool,
+) -> DispatchSpec:
+    num_reqs = query_start_loc_ptr.shape[0] - 1
+    num_blocks = triton.cdiv(max_tokens_per_req, triton_block_size)
+    return (num_reqs, num_blocks), dict(
+        BLOCK_SIZE=triton_block_size,
+    )
+
+
 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
 def update_num_computed_tokens_for_batch_change(
     num_computed_tokens: torch.Tensor,
diff --git a/vllm/v1/spec_decode/vocab_mapping.py b/vllm/v1/spec_decode/vocab_mapping.py
index 9a6bbe7cfaf4..dad059fa77d9 100644
--- a/vllm/v1/spec_decode/vocab_mapping.py
+++ b/vllm/v1/spec_decode/vocab_mapping.py
@@ -137,16 +137,12 @@ def __init__(
 
     def map_target_to_draft_ids(self, target_ids):
         draft_ids = self.target_to_draft_ids[target_ids]  # new tensor; no clone needed
-        missing = draft_ids == -1
-        if missing.any():
-            draft_ids[missing] = self.draft_unk_token_id
+        draft_ids[draft_ids == -1] = self.draft_unk_token_id
         return draft_ids.to(target_ids.dtype)
 
     def map_draft_to_target_ids(self, draft_ids):
         target_ids = self.draft_to_target_ids[draft_ids]  # new tensor; no clone needed
-        missing = target_ids == -1
-        if missing.any():
-            target_ids[missing] = self.target_unk_token_id
+        target_ids[target_ids == -1] = self.target_unk_token_id
         return target_ids.to(draft_ids.dtype)
 
     def constrain_draft_logits(self, logits):
diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py
index fd48c99d1cf9..7fd311c873c1 100644
--- a/vllm/v1/structured_output/__init__.py
+++ b/vllm/v1/structured_output/__init__.py
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import itertools
 import multiprocessing
 from collections.abc import Iterable, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
@@ -18,7 +17,7 @@
     StructuredOutputGrammar,
 )
 from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
-from vllm.v1.structured_output.utils import maybe_wrap_mistral_common_tokenizer
+from vllm.v1.structured_output.utils import strip_speculative_padding
 
 if TYPE_CHECKING:
     import numpy as np
@@ -77,8 +76,8 @@ def __init__(self, vllm_config: VllmConfig):
             # of CPUs.
             max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2)
             self.executor = ThreadPoolExecutor(max_workers=max_workers)
-            self.tokenizer = maybe_wrap_mistral_common_tokenizer(
-                cached_tokenizer_from_config(model_config=self.vllm_config.model_config)
+            self.tokenizer = cached_tokenizer_from_config(
+                model_config=self.vllm_config.model_config
             )
             reasoning_parser_plugin = (
                 self.vllm_config.structured_outputs_config.reasoning_parser_plugin
@@ -218,6 +217,100 @@ def _async_submit_fill_bitmask(
     ) -> Future:
         return self.executor_for_fillmask.submit(self._fill_bitmasks, batch)
 
+    def _get_constraint_start(
+        self,
+        request: "Request",
+        spec_tokens: Sequence[int],
+        spec_tokens_committed: bool = False,
+    ) -> int:
+        """Return the index into `spec_tokens` where tokens should start being
+        constrained by the grammar.
+        Assumes `spec_tokens` are stripped of -1 padding tokens.
+        Returns `len(spec_tokens) + 1` if no token should be constrained after
+        accepting all `spec_tokens`.
+
+        `spec_tokens_committed` is True if spec_tokens is already in
+        `request.all_token_ids`.
+        """
+        if self.enable_in_reasoning:
+            return 0
+
+        structured_req = request.structured_output_request
+        assert structured_req is not None
+        if structured_req.reasoning_ended:
+            return 0
+
+        reasoner = self._get_reasoner(request)
+        if reasoner is None:
+            return 0
+
+        if structured_req.reasoning_ended is None:
+            # This should be removed here, but since `openai_gptoss`
+            # is an independent code path, it is kept for now.
+            # After unifying the `openai_gptoss` and non-`openai_gptoss` styles,
+            # it can be removed.
+            if reasoner.is_reasoning_end(request.prompt_token_ids or []):
+                return 0
+            structured_req.reasoning_ended = False
+
+        num_spec_tokens = len(spec_tokens)
+        if num_spec_tokens <= 0:
+            return num_spec_tokens + 1
+
+        # Use `find_reasoning_end_offset` to find constraint start if supported
+        if (
+            isinstance(reasoner, ParserEngineReasoningAdapter)
+            and reasoner.reasoning_end_token_ids
+        ):
+            offset = reasoner.find_reasoning_end_offset(spec_tokens)
+            if offset is not None:
+                return offset + 1
+
+        # Fallback to `find_reasoning_end_offset`
+        # TODO: Build a read-only Sequence view over all_token_ids
+        # instead of copying the entire all_token_ids into input_ids.
+        if spec_tokens_committed:
+            if not spec_tokens or not reasoner.is_reasoning_end_streaming(
+                request.all_token_ids, spec_tokens
+            ):
+                return num_spec_tokens + 1
+
+            input_ids = request.all_token_ids.copy()
+            delta_ids = list(spec_tokens)
+        else:
+            input_ids = request.all_token_ids.copy()
+            input_ids.extend(spec_tokens)
+            if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens):
+                return num_spec_tokens + 1
+            delta_ids = list(spec_tokens)
+
+        for i in range(num_spec_tokens - 1, 0, -1):
+            input_ids.pop()
+            delta_ids.pop()
+            if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids):
+                return i + 1
+        return 1
+
+    def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]:
+        """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`."""
+        if not request.use_structured_output:
+            return spec_tokens
+
+        spec_tokens = strip_speculative_padding(spec_tokens)
+        constraint_start = self._get_constraint_start(request, spec_tokens)
+        if constraint_start >= len(spec_tokens):
+            return spec_tokens
+
+        structured_req = request.structured_output_request
+        if TYPE_CHECKING:
+            assert structured_req is not None
+        grammar = structured_req.grammar
+        if TYPE_CHECKING:
+            assert isinstance(grammar, StructuredOutputGrammar)
+        prefix = spec_tokens[:constraint_start]
+        validated = grammar.validate_tokens(spec_tokens[constraint_start:])
+        return prefix + validated
+
     def grammar_bitmask(
         self,
         requests: dict[str, "Request"],
@@ -265,7 +358,7 @@ def grammar_bitmask(
                 if TYPE_CHECKING:
                     assert isinstance(grammar, StructuredOutputGrammar)
 
-                apply_bitmask = self.should_fill_bitmask(request)
+                apply_bitmask = self._get_constraint_start(request, ()) == 0
                 batch.append((grammar, cumulative_index, apply_bitmask))
                 if len(batch) == self.fill_bitmask_parallel_batch_size:
                     promises.append(self._async_submit_fill_bitmask(batch))
@@ -289,76 +382,47 @@ def grammar_bitmask(
                 grammar = structured_output_request.grammar
                 if TYPE_CHECKING:
                     assert isinstance(grammar, StructuredOutputGrammar)
-                apply_bitmask = self.should_fill_bitmask(request)
-
-                reasoner = None if apply_bitmask else self._get_reasoner(request)
 
+                req_tokens = scheduled_spec_decode_tokens.get(req_id, list())
+                constraint_start = self._get_constraint_start(
+                    request, strip_speculative_padding(req_tokens)
+                )
                 state_advancements = 0
-                post_reasoning_end_in_window = False
-                req_tokens = scheduled_spec_decode_tokens.get(req_id, ())
-                # Locate the reasoning-end marker in the draft window once.
-                # Rejected-draft padding (-1) is trailing, so the window is
-                # cut there and the reasoner never sees a placeholder.
-                window_start = len(request.all_token_ids)
-                reasoning_end_index: int | None = None
-                if reasoner is not None and req_tokens:
-                    window = req_tokens
-                    if -1 in window:
-                        window = window[: window.index(-1)]
-                    reasoning_end_index = self._find_reasoning_end_index(
-                        reasoner,
-                        request.all_token_ids,
-                        window_start,
-                        window,
-                        delta_appended=False,
-                    )
+                seen_padding = False
+                failed = False
                 for i, token in enumerate(req_tokens):
+                    apply_bitmask = (
+                        not failed and not seen_padding and i >= constraint_start
+                    )
                     self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))
-                    advance_grammar = apply_bitmask
                     if token == -1:
-                        apply_bitmask = False
-                        advance_grammar = False
-                    elif not apply_bitmask and window_start + i == reasoning_end_index:
-                        # Reasoning ended mid-window. Constrain the rest
-                        # of the window via bitmask. Skip grammar advance
-                        # through the marker (it is reasoning content);
-                        # try to advance through subsequent drafts so the
-                        # next bitmask row reflects the post-advance state,
-                        # but tolerate rejection since those drafts predate
-                        # the bitmask and are not guaranteed valid.
-                        apply_bitmask = True
-                        advance_grammar = False
-                        post_reasoning_end_in_window = True
-                    if advance_grammar and not grammar.is_terminated():
-                        if post_reasoning_end_in_window:
-                            accepted = bool(grammar.validate_tokens([token]))
-                            if accepted:
-                                accepted = grammar.accept_tokens(req_id, [token])
-                        else:
-                            accepted = grammar.accept_tokens(req_id, [token])
-                        if accepted:
+                        seen_padding = True
+                    elif apply_bitmask:
+                        if not grammar.is_terminated() and grammar.accept_tokens(
+                            req_id, [token]
+                        ):
                             state_advancements += 1
-                        elif not post_reasoning_end_in_window:
-                            raise AssertionError(
-                                (token, req_id, scheduled_spec_decode_tokens)
+                        else:
+                            failed = True
+                            logger.error(
+                                "Unexpected: grammar terminated or rejected draft "
+                                "token %s for request %s during bitmask fill.",
+                                token,
+                                req_id,
                             )
                     cumulative_index += 1
+
                 # Diffusion LLMs don't sample a bonus token after the
                 # scheduled positions, so skip its bitmask in that case.
                 if not (self.vllm_config.model_config.is_diffusion and req_tokens):
-                    # bonus_apply must be True when the bonus-row position
-                    # should be grammar-constrained. Two triggers:
-                    # - should_fill_bitmask(request): reasoning was already
-                    #   over at step start (or no reasoner /
-                    #   enable_in_reasoning).
-                    # - apply_bitmask: reasoning ended mid-window in this
-                    #   call and was flipped True after the marker;
-                    #   should_fill_bitmask still returns False here because
-                    #   reasoning_ended is only persisted later by
-                    #   should_advance.
-                    bonus_apply = self.should_fill_bitmask(request) or apply_bitmask
+                    bonus_apply = (
+                        not failed
+                        and not seen_padding
+                        and constraint_start <= len(req_tokens)
+                    )
                     self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),))
                     cumulative_index += 1
+
                 if state_advancements > 0:
                     grammar.rollback(state_advancements)
 
@@ -371,161 +435,29 @@ def grammar_bitmask(
         # and deserialization when sending this to the GPU workers.
         return bitmask_tensor.numpy()
 
-    def should_fill_bitmask(self, request: "Request") -> bool:
-        # NOTE (Hanchen) if enable_in_reasoning is True, it means that
-        # the model needs to be constrained in reasoning. So we should always
-        # enable the bitmask filling.
-        structured_req = request.structured_output_request
-        if self.enable_in_reasoning or (
-            structured_req is not None and structured_req.reasoning_ended is True
-        ):
+    def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool:
+        """Advance grammar with accepted tokens. Returns False on rejection."""
+        if not request.use_structured_output or not new_token_ids:
             return True
 
-        reasoner = self._get_reasoner(request)
-        if reasoner is not None:
+        structured_req = request.structured_output_request
+        if TYPE_CHECKING:
             assert structured_req is not None
-            if structured_req.reasoning_ended is None:
-                # This should be removed here, but since `openai_gptoss`
-                # is an independent code path, it is kept for now.
-                # After unifying the `openai_gptoss` and non-`openai_gptoss` styles,
-                # it can be removed.
-                structured_req.reasoning_ended = reasoner.is_reasoning_end(
-                    request.prompt_token_ids or []
-                )
-            return structured_req.reasoning_ended
-        return True
-
-    def should_advance(
-        self,
-        request: "Request",
-        new_token_ids: list[int] | None = None,
-    ) -> bool:
-        if not request.use_structured_output:
-            return False
-
-        # To determine whether we can advance the FSM.
-        # Supports thinking usage where we skip the reasoning components.
+        grammar = structured_req.grammar
         if TYPE_CHECKING:
-            assert request.structured_output_request is not None
-            assert request.structured_output_request.grammar is not None
-        structured_req = request.structured_output_request
-        if self.enable_in_reasoning or (
-            structured_req is not None and structured_req.reasoning_ended is True
-        ):
-            return True
-
-        # by default, we should always advance
-        # for cases that don't use thinking mode.
-        reasoner = self._get_reasoner(request)
-        if reasoner is None:
-            return True
+            assert isinstance(grammar, StructuredOutputGrammar)
 
-        assert structured_req is not None
-
-        # Check if reasoning ends in *this* step.
-        # When the caller passes new_token_ids (the tokens that were just
-        # appended this step), use it directly as the delta window. The
-        # placeholder-derived fallback assumes num_output_placeholders ==
-        # len(new_token_ids), which breaks under async scheduling + spec
-        # decode when some drafts are rejected (#43388): the placeholder
-        # count remains > 0 after the step and the computed delta window
-        # starts past the reasoning-end marker.
-        all_token_ids = request.all_token_ids
-        if new_token_ids:
-            # The tokens were already appended this step, so the step window
-            # starts exactly len(new_token_ids) from the end.
-            start = len(all_token_ids) - len(new_token_ids)
-            delta_ids: Sequence[int] = new_token_ids
-        else:
-            # Draft-validation callers arrive here after update_from_output
-            # has already persisted any boundary found this step, so this
-            # window only needs to cover tokens not examined yet.
-            delta_from = request.num_computed_tokens - request.num_output_placeholders
-            start = (
-                delta_from
-                if delta_from >= 0
-                else max(len(all_token_ids) + delta_from, 0)
-            )
-            delta_ids = all_token_ids[start:]
-        end_index = self._find_reasoning_end_index(
-            reasoner, all_token_ids, start, delta_ids, delta_appended=True
+        constraint_start = self._get_constraint_start(
+            request, new_token_ids, spec_tokens_committed=True
         )
-        if end_index is not None:
-            structured_req.reasoning_ended = True
-            # Record the boundary so the scheduler can exclude reasoning tokens.
-            structured_req.reasoning_end_token_index = end_index
+        # Early return only when the constraint hasn't started.
+        # Otherwise, latch structured_req.reasoning_ended.
+        if constraint_start > len(new_token_ids):
             return True
-
-        return False
-
-    @staticmethod
-    def _find_reasoning_end_index(
-        reasoner: "ReasoningParser",
-        all_token_ids: Sequence[int],
-        start: int,
-        delta_ids: Sequence[int],
-        *,
-        delta_appended: bool,
-    ) -> int | None:
-        """Locates the token that ends reasoning within ``delta_ids``.
-
-        ``delta_ids`` occupy positions ``start..``; ``delta_appended`` says
-        whether they are already in ``all_token_ids`` (accepted tokens) or
-        not (speculative drafts).
-
-        Legacy parsers are probed one token at a time. Accepted tokens get a
-        whole-delta check as well, and a marker seen only that way reports
-        the final token. Drafts skip it: some predicates are order-sensitive
-        over a multi-token window and would miss the marker.
-
-        Returns:
-            Absolute index of the last reasoning token, or ``None``.
-        """
-        if (
-            isinstance(reasoner, ParserEngineReasoningAdapter)
-            and reasoner.reasoning_end_token_ids
-        ):
-            offset = reasoner.find_reasoning_end_offset(delta_ids)
-            return None if offset is None else start + offset
-
-        if delta_appended and not reasoner.is_reasoning_end_streaming(
-            all_token_ids, delta_ids
-        ):
-            return None
-        prefix = list(itertools.islice(all_token_ids, start))
-        for offset, token in enumerate(delta_ids):
-            prefix.append(token)
-            if reasoner.is_reasoning_end_streaming(prefix, [token]):
-                return start + offset
-        if delta_appended:
-            return len(all_token_ids) - 1
-        return None
-
-    def trim_reasoning_for_advance(
-        self, request: "Request", new_token_ids: list[int]
-    ) -> list[int]:
-        """Drops reasoning content from tokens about to advance the grammar.
-
-        When reasoning ends mid-step (see should_advance), the step's output
-        still contains reasoning tokens up to and including the end marker.
-        Those are not grammar content: feeding them to accept_tokens makes
-        the grammar reject the marker and kills the request (#44006).
-
-        Returns:
-            The suffix of ``new_token_ids`` that follows the reasoning-end
-            marker. Steps fully after the boundary are returned unchanged.
-        """
-        structured_req = request.structured_output_request
-        if structured_req is None:
-            return new_token_ids
-        end_idx = structured_req.reasoning_end_token_index
-        if end_idx is None:
-            return new_token_ids
-        first_idx = len(request.all_token_ids) - len(new_token_ids)
-        num_reasoning = end_idx + 1 - first_idx
-        if num_reasoning <= 0:
-            return new_token_ids
-        return new_token_ids[num_reasoning:]
+        structured_req.reasoning_ended = True
+        return grammar.accept_tokens(
+            request.request_id, new_token_ids[constraint_start:]
+        )
 
     def clear_backend(self) -> None:
         if self.backend is not None:
diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py
index 16e1f54be6b8..5f43c5a890fc 100644
--- a/vllm/v1/structured_output/backend_guidance.py
+++ b/vllm/v1/structured_output/backend_guidance.py
@@ -8,7 +8,6 @@
 from typing import TYPE_CHECKING, Any
 
 import torch
-from transformers import MistralCommonBackend
 
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
@@ -98,10 +97,6 @@ def __post_init__(self):
 
         if is_mistral_tokenizer(self.tokenizer):
             self.ll_tokenizer = self.tokenizer.llg_tokenizer
-        elif isinstance(self.tokenizer, MistralCommonBackend):
-            from mistral_common.guidance.tokenizer import from_mistral_tokenizer
-
-            self.ll_tokenizer = from_mistral_tokenizer(self.tokenizer.tokenizer)
         else:
             self.ll_tokenizer = llguidance_hf.from_tokenizer(
                 self.tokenizer, max(self.vocab_size, len(self.tokenizer))
@@ -166,7 +161,6 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
         Returns True if the parser was advanced successfully.
         Returns False if the parser failed to advance.
         """
-
         if self.ll_tokenizer.eos_token in tokens:
             if self.ll_matcher.is_stopped() and not self.terminated:
                 self.rollback_lag = 1
diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py
index dd460f96ec91..bb49fb189b72 100644
--- a/vllm/v1/structured_output/backend_outlines.py
+++ b/vllm/v1/structured_output/backend_outlines.py
@@ -283,7 +283,7 @@ def subpattern_consumes(parsed) -> bool:
 
 
 def _check_unsupported(parsed) -> None:
-    """Check for regex features unsupported by regex-automata"""
+    """Check for regex features unsupported by regex-automata."""
     tokens = parsed.data if hasattr(parsed, "data") else parsed
     for ttype, tval in tokens:
         # backreference
@@ -310,8 +310,7 @@ def _check_unsupported(parsed) -> None:
 
 
 def validate_regex_is_buildable(pattern: str) -> None:
-    """
-    Validates that the input regex is not using unsupported features
+    """Validates that the input regex is not using unsupported features
     of the `regex-automata` crate (outlines_core regex engine) and has a
     universal start state.
     definition of universal start state used can be found at:
diff --git a/vllm/v1/structured_output/backend_types.py b/vllm/v1/structured_output/backend_types.py
index d9826c18f534..b3695efd009c 100644
--- a/vllm/v1/structured_output/backend_types.py
+++ b/vllm/v1/structured_output/backend_types.py
@@ -33,8 +33,7 @@ class StructuredOutputGrammar(ABC):
 
     @abstractmethod
     def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
-        """
-        Determines whether the provided tokens are accepted for the
+        """Determines whether the provided tokens are accepted for the
         given request.
 
         Args:
@@ -43,12 +42,12 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
 
         Returns:
             bool: True if the tokens are accepted, False otherwise.
+
         """
 
     @abstractmethod
     def validate_tokens(self, tokens: list[int]) -> list[int]:
-        """
-        Validates the provided tokens against the grammar.
+        """Validates the provided tokens against the grammar.
         Will not advance the FSM.
 
         Args:
@@ -57,42 +56,41 @@ def validate_tokens(self, tokens: list[int]) -> list[int]:
         Returns:
             list[int]: A list of accepted token IDs. Will be a prefix
                 of the input tokens, and empty if none are accepted.
+
         """
 
     @abstractmethod
     def rollback(self, num_tokens: int) -> None:
-        """
-        Rolls back the state of the grammar by a specified number of tokens.
+        """Rolls back the state of the grammar by a specified number of tokens.
         Will also revert counters for the number of processed tokens.
 
         Args:
             num_tokens (int): The number of tokens to roll back.
+
         """
 
     @abstractmethod
     def fill_bitmask(self, bitmask: "torch.Tensor", batch_index: int) -> None:
-        """
-        Fills the bitmask for a specific batch index.
+        """Fills the bitmask for a specific batch index.
 
         Args:
             bitmask (torch.Tensor): The bitmask to fill
             batch_index (int): The index in the bitmask to fill
+
         """
 
     @abstractmethod
     def is_terminated(self) -> bool:
-        """
-        Checks whether the structured output process has terminated.
+        """Checks whether the structured output process has terminated.
 
         Returns:
             bool: True if the process is terminated, False otherwise.
+
         """
 
     @abstractmethod
     def reset(self):
-        """
-        Resets the state of the structured output grammar.
-        """
+        """Resets the state of the structured output grammar."""
 
 
 @dataclass
@@ -110,8 +108,7 @@ def compile_grammar(
         grammar_spec: str,
         stop_token_ids: set[int] | None = None,
     ) -> StructuredOutputGrammar:
-        """
-        Compiles a grammar specification into a structured output grammar.
+        """Compiles a grammar specification into a structured output grammar.
 
         Args:
             request_type (StructuredOutputOptions): The type of structured
@@ -123,20 +120,19 @@ def compile_grammar(
 
         Returns:
             StructuredOutputGrammar: The compiled structured output grammar.
+
         """
 
     @abstractmethod
     def allocate_token_bitmask(self, max_num_seqs: int) -> "torch.Tensor":
-        """
-        Allocates a token bitmask for the specified maximum number of sequences.
+        """Allocates a token bitmask for the specified maximum number of sequences.
 
         Args:
             max_num_seqs (int): The maximum number of sequences for which
                 to allocate the bitmask.
+
         """
 
     @abstractmethod
     def destroy(self):
-        """
-        Backend-specific cleanup.
-        """
+        """Backend-specific cleanup."""
diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py
index 3809ba2d91b0..3c48ff64a0bd 100644
--- a/vllm/v1/structured_output/backend_xgrammar.py
+++ b/vllm/v1/structured_output/backend_xgrammar.py
@@ -236,6 +236,24 @@ def reset(self):
 }
 
 
+def _has_pattern_and_length_bounds(schema: dict[str, Any]) -> bool:
+    return ("pattern" in schema or "format" in schema) and (
+        "minLength" in schema or "maxLength" in schema
+    )
+
+
+# FIXME(arpera): The approach used here needs to be redesigned because of
+# existing bugs: https://github.com/vllm-project/vllm/issues/57550
+def _schema_types(schema: dict[str, Any]) -> set[str]:
+    """Normalize a scalar or list-valued JSON Schema type."""
+    schema_type = schema.get("type")
+    if isinstance(schema_type, str):
+        return {schema_type}
+    if isinstance(schema_type, list):
+        return {item for item in schema_type if isinstance(item, str)}
+    return set()
+
+
 def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool:
     """Check if JSON schema contains features unsupported by xgrammar."""
 
@@ -243,12 +261,14 @@ def check_object(obj: dict[str, Any]) -> bool:
         if not isinstance(obj, dict):
             return False
 
+        schema_types = _schema_types(obj)
+
         # Check for numeric ranges
-        if obj.get("type") in ("integer", "number") and ("multipleOf" in obj):
+        if (schema_types & {"integer", "number"}) and ("multipleOf" in obj):
             return True
 
         # Check for array unsupported keywords
-        if obj.get("type") == "array" and any(
+        if "array" in schema_types and any(
             key in obj
             for key in ("uniqueItems", "contains", "minContains", "maxContains")
         ):
@@ -256,7 +276,7 @@ def check_object(obj: dict[str, Any]) -> bool:
 
         # Unsupported keywords for strings
         if (
-            obj.get("type") == "string"
+            "string" in schema_types
             and "format" in obj
             and obj["format"] not in STRING_SUPPORTED_FORMATS
         ):
@@ -269,16 +289,39 @@ def check_object(obj: dict[str, Any]) -> bool:
         # the compiled EBNF: pattern/format grammars come out byte-identical
         # with and without the length keywords, while maxLength alone lowers
         # to {0, N} correctly.
+        if "string" in schema_types and _has_pattern_and_length_bounds(obj):
+            return True
+
+        # propertyNames validates names, so it is a string schema even when it
+        # omits "type", which is the form that escapes the check above.
         if (
-            obj.get("type") == "string"
-            and ("pattern" in obj or "format" in obj)
-            and ("minLength" in obj or "maxLength" in obj)
+            "object" in schema_types
+            and isinstance(obj.get("propertyNames"), dict)
+            and _has_pattern_and_length_bounds(obj["propertyNames"])
         ):
             return True
 
-        # Unsupported keywords for objects
-        if obj.get("type") == "object" and any(
-            key in obj for key in ("patternProperties", "propertyNames")
+        # FIXME: propertyNames conflicts with properties/patternProperties/
+        # additionalProperties/unevaluatedProperties under xgrammar.
+        # https://github.com/mlc-ai/xgrammar/issues/826
+        if (
+            "object" in schema_types
+            and "propertyNames" in obj
+            and (
+                "properties" in obj
+                or "patternProperties" in obj
+                or isinstance(obj.get("additionalProperties"), dict)
+                or obj.get("unevaluatedProperties", True) is not True
+            )
+        ):
+            return True
+
+        # FIXME: multiple patternProperties, or patternProperties alongside
+        # properties, conflict under xgrammar.
+        if (
+            "object" in schema_types
+            and isinstance(obj.get("patternProperties"), dict)
+            and ("properties" in obj or len(obj["patternProperties"]) > 1)
         ):
             return True
 
diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py
index b0f4f3c3a8fd..cd044f54669f 100644
--- a/vllm/v1/structured_output/request.py
+++ b/vllm/v1/structured_output/request.py
@@ -4,7 +4,6 @@
 import functools
 import json
 from concurrent.futures import Future
-from concurrent.futures._base import TimeoutError
 from typing import TYPE_CHECKING, Any, cast
 
 from vllm.sampling_params import SamplingParams, StructuredOutputsParams
@@ -25,12 +24,6 @@ class StructuredOutputRequest:
         Future[StructuredOutputGrammar] | StructuredOutputGrammar | Exception | None
     ) = None
     reasoning_ended: bool | None = None
-    # Absolute index into the request's all_token_ids of the last reasoning
-    # token (the reasoning-end marker). Tokens at or before this index are
-    # reasoning content and must never be fed to the grammar. Only set when
-    # reasoning ends in a step whose tokens the scheduler advances immediately
-    # (structural tags + speculative decoding, see #42452).
-    reasoning_end_token_index: int | None = None
     reasoning_parser_kwargs: dict[str, Any] | None = None
     # Cached per request; do not share reasoning parsers across requests because
     # their behavior can depend on reasoning_parser_kwargs.
@@ -49,11 +42,10 @@ def from_sampling_params(
 
     def _check_grammar_completion(self) -> bool:
         if isinstance(self._grammar, Future):
-            try:
-                # We will check whether the future is ready within 100 us
-                self._grammar = self._grammar.result(timeout=0.0001)
-            except TimeoutError:
+            if not self._grammar.done():
                 return False
+            try:
+                self._grammar = self._grammar.result()
             except Exception as e:
                 self._grammar = e
         return True
diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py
index aa5d59c9c5b4..4bb438d6f380 100644
--- a/vllm/v1/structured_output/utils.py
+++ b/vllm/v1/structured_output/utils.py
@@ -14,7 +14,6 @@
 import regex as re
 import torch
 from cachetools import LRUCache
-from transformers import MistralCommonBackend
 
 import vllm.envs as envs
 from vllm.logger import init_logger
@@ -75,6 +74,7 @@ def compile_regex_with_timeout(fn: Callable[[str], _T], pattern: str) -> _T:
 
     Raises:
         ValueError: If compilation exceeds the configured timeout.
+
     """
     timeout = envs.VLLM_REGEX_COMPILATION_TIMEOUT_S
     if timeout <= 0:
@@ -104,13 +104,14 @@ def apply_grammar_bitmask(
     input_batch: InputBatch,
     logits: torch.Tensor,
 ) -> None:
-    """
-    Apply grammar bitmask to output logits of the model with xgrammar function.
+    """Apply grammar bitmask to output logits of the model with xgrammar function.
 
     Args:
         scheduler_output (SchedulerOutput): The result of engine scheduling.
+        grammar_output (GrammarOutput): The grammar bitmask to apply.
         input_batch (InputBatch): The input of model runner.
         logits (torch.Tensor): The output logits of model forward.
+
     """
     # Serialization of np.ndarray is much more efficient than a tensor,
     # so we receive it in that format.
@@ -191,8 +192,7 @@ def apply_grammar_bitmask(
 
 
 class OutlinesVocabulary:
-    """
-    Wrapper class for `outlines_core.Vocabulary`,
+    """Wrapper class for `outlines_core.Vocabulary`,
     which allows us to store a hash with the vocabulary
     """
 
@@ -207,7 +207,7 @@ def __init__(self, vocabulary: oc.Vocabulary) -> None:
 
 
 def get_outlines_cache_path() -> str:
-    """Get the context object that contains previously-computed return values"""
+    """Get the context object that contains previously-computed return values."""
     outlines_cache_dir = os.getenv("OUTLINES_CACHE_DIR")
     xdg_cache_home = os.getenv("XDG_CACHE_HOME")
     home_dir = os.path.expanduser("~")
@@ -295,8 +295,7 @@ def clear(self):
 
 
 def get_outlines_cache():
-    """Get the Cache instance to be used for index caching"""
-
+    """Get the Cache instance to be used for index caching."""
     cache_dir = get_outlines_cache_path()
     if envs.VLLM_V1_USE_OUTLINES_CACHE:
         logger.warning(
@@ -320,23 +319,12 @@ def get_outlines_cache():
 re_replacement_seq = re.compile(r"^.{0,6}�+.{0,6}$")
 
 
-def maybe_wrap_mistral_common_tokenizer(tokenizer: TokenizerLike) -> TokenizerLike:
-    """The grammar backends cannot consume a `MistralCommonBackend` directly."""
-    if not isinstance(tokenizer, MistralCommonBackend):
-        return tokenizer
-
-    # Deferred: `vllm.tokenizers.mistral` pulls in a large dependency tree, and
-    # this module is imported by `vllm.v1.worker.gpu_model_runner`.
-    from vllm.tokenizers.mistral import MistralTokenizer
-
-    return MistralTokenizer(tokenizer)
-
-
 def _reduced_vocabulary(tokenizer: TokenizerLike) -> dict[bytes, list[int]]:
     """Create a map from vocabulary tokens to lists of equivalent token ids.
 
     Returns:
         A Dict of token string -> equivalent token ids
+
     """
     eos_token_id = tokenizer.eos_token_id
 
@@ -416,8 +404,7 @@ def get_outlines_vocabulary(tokenizer: TokenizerLike) -> oc.Vocabulary:
 
 
 def grammar_is_likely_lark(grammar_str: str) -> bool:
-    """
-    Check if grammar appears to use Lark syntax.
+    """Check if grammar appears to use Lark syntax.
 
     Args:
         grammar_str: Input grammar string
@@ -430,6 +417,7 @@ def grammar_is_likely_lark(grammar_str: str) -> bool:
         True
         >>> grammar_is_likely_lark("rule ::= 'abc'")
         False
+
     """
     if not grammar_str or not isinstance(grammar_str, str):
         return False
@@ -448,8 +436,7 @@ def grammar_is_likely_lark(grammar_str: str) -> bool:
 
 
 def convert_lark_to_ebnf(grammar_str: str) -> str:
-    """
-    Convert a Lark grammar string to EBNF format.
+    """Convert a Lark grammar string to EBNF format.
 
     EBNF reference:
     https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md
@@ -466,6 +453,7 @@ def convert_lark_to_ebnf(grammar_str: str) -> str:
         >>> print(convert_lark_to_ebnf("rule: 'hello'"))
         root ::= rule
         rule ::= "hello"
+
     """
     if not isinstance(grammar_str, str):
         raise ValueError(f"Grammar must be a string, got {type(grammar_str)}")
@@ -579,9 +567,18 @@ def extract_references(text: str) -> set[str]:
 
 def choice_as_grammar(choice: list[str]) -> str:
     def escape_ebnf_string(s: str) -> str:
-        """Escape special characters in a EBNF string."""
-        # Escape double quotes and backslashes
-        return re.sub(r'(["\\])', r"\\\1", s)
+        """Escape EBNF literals, including raw LF, CR, and NUL terminators."""
+        escapes = {"\\": r"\\", '"': r"\"", "\n": r"\n", "\r": r"\r", "\t": r"\t"}
+
+        def escape_char(ch: str) -> str:
+            if ch in escapes:
+                return escapes[ch]
+            # Escape remaining C0 controls (U+0000-U+001F) and DEL (U+007F).
+            if ord(ch) < 0x20 or ord(ch) == 0x7F:
+                return f"\\u{ord(ch):04x}"
+            return ch
+
+        return "".join(escape_char(ch) for ch in s)
 
     escaped_choices = (escape_ebnf_string(c) for c in choice)
     grammar = "root ::= " + " | ".join(f'"{c}"' for c in escaped_choices)
diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py
index b0411dc9e8b1..a89ccbb9ed7c 100644
--- a/vllm/v1/utils.py
+++ b/vllm/v1/utils.py
@@ -137,9 +137,10 @@ def __init__(
             self.np = self.cpu.numpy()
 
     def copy_to_gpu(self, n: int | None = None) -> torch.Tensor:
-        if n is None:
-            return self.gpu.copy_(self.cpu, non_blocking=True)
-        return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)
+        cpu, gpu = self.cpu, self.gpu
+        if n is not None:
+            cpu, gpu = cpu[:n], gpu[:n]
+        return gpu.copy_(cpu.pin_memory() if PIN_MEMORY else cpu, non_blocking=True)
 
     def copy_to_cpu(self, n: int | None = None) -> torch.Tensor:
         """NOTE: Because this method is non-blocking, explicit synchronization
@@ -200,6 +201,7 @@ def __init__(
             output_addresses: Output addresses for each API server
             stats_update_address: Optional stats update address
             tensor_queue: Optional tensor IPC queue for sharing MM tensors
+
         """
         self.listen_address = listen_address
         self.sock = sock
@@ -209,6 +211,15 @@ def __init__(
         self.processes: list[BaseProcess] = []
         self._address_pipes: list[connection.Connection] = []
 
+        admission_counters = None
+        if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None:
+            from vllm.v1.engine.admission_control import SharedAdmissionStats
+
+            admission_counters = spawn_context.RawArray(
+                "q",
+                SharedAdmissionStats.num_counters(num_servers),
+            )
+
         for i, in_addr, out_addr in zip(
             range(num_servers), input_addresses, output_addresses
         ):
@@ -218,6 +229,8 @@ def __init__(
                 "client_count": num_servers,
                 "client_index": i,
             }
+            if admission_counters is not None:
+                client_config["mp_admission_counters"] = admission_counters
             if stats_update_address is not None:
                 client_config["stats_update_address"] = stats_update_address
             if tensor_queue is not None:
@@ -507,7 +520,6 @@ def run_api_server_worker_proc(
     listen_address, sock, args, client_config=None, **uvicorn_kwargs
 ) -> None:
     """Entrypoint for individual API server worker processes."""
-
     from vllm.entrypoints.launchers.api_server.entry import run_server_worker
 
     client_config = client_config or {}
@@ -538,8 +550,8 @@ def wait_for_completion_or_failure(
             If CoreEngineProcManager, it manages local engines;
             if CoreEngineActorManager, it manages all engines.
         coordinator: The coordinator for data parallel.
-    """
 
+    """
     try:
         logger.info("Waiting for API servers to complete ...")
         # Create a mapping of sentinels to their corresponding processes
@@ -601,6 +613,7 @@ def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None:
     Args:
         procs: List of processes to shutdown
         timeout: Maximum time in seconds to wait for graceful shutdown
+
     """
     if timeout is None:
         # Keep a small grace period for best-effort cleanup paths that do not
@@ -655,8 +668,7 @@ def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None:
 def copy_slice(
     from_tensor: torch.Tensor, to_tensor: torch.Tensor, length: int
 ) -> torch.Tensor:
-    """
-    Copy the first length elements of a tensor into another tensor in a
+    """Copy the first length elements of a tensor into another tensor in a
     non-blocking manner.
 
     Used to copy pinned CPU tensor data to pre-allocated GPU tensors.
@@ -670,7 +682,6 @@ def report_usage_stats(
     vllm_config, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT
 ) -> None:
     """Report usage statistics if enabled."""
-
     if not is_usage_stats_enabled():
         return
 
@@ -791,6 +802,7 @@ def tensor_data(tensor: torch.Tensor) -> memoryview:
 
     Returns:
         A memoryview of the tensor data as uint8.
+
     """
     return tensor.flatten().cpu().contiguous().view(torch.uint8).numpy().data
 
@@ -814,8 +826,7 @@ def __repr__(self) -> str:
 
 
 def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDetails:
-    """
-    Compute the number of context/generation requests and tokens
+    """Compute the number of context/generation requests and tokens
     for the current iteration's scheduler output. A requests is regarded
     as a context request if its output tokens are still 0, an extended chunk
     of chunked prefill falls into this category.
@@ -826,6 +837,7 @@ def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDet
     Returns:
         An IterationDetails object containing the number of
         context/generation requests and tokens.
+
     """
     num_context_requests = 0
     num_context_tokens = 0
diff --git a/vllm/v1/watermarking/__init__.py b/vllm/v1/watermarking/__init__.py
index bb53981461ce..23e9fd20fda6 100644
--- a/vllm/v1/watermarking/__init__.py
+++ b/vllm/v1/watermarking/__init__.py
@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from vllm.config.watermarking import derive_watermark_key
 from vllm.v1.watermarking.detector import WatermarkDetection, WatermarkDetector
 from vllm.v1.watermarking.factory import create_watermarker
 from vllm.v1.watermarking.gumbel import (
+    DualKeyGumbelWatermarkDetector,
+    DualKeyGumbelWatermarker,
     GumbelWatermarkDetector,
     GumbelWatermarker,
 )
@@ -12,12 +15,19 @@
     WatermarkPRF,
     create_prf,
 )
-from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
+from vllm.v1.watermarking.watermarker import (
+    SupportsSpeculativeDecoding,
+    Watermarker,
+    WatermarkSample,
+)
 
 __all__ = [
+    "DualKeyGumbelWatermarkDetector",
+    "DualKeyGumbelWatermarker",
     "GumbelWatermarkDetector",
     "GumbelWatermarker",
     "PhiloxPRF",
+    "SupportsSpeculativeDecoding",
     "WatermarkDetection",
     "WatermarkDetector",
     "WatermarkPRF",
@@ -25,4 +35,5 @@
     "WatermarkSample",
     "create_watermarker",
     "create_prf",
+    "derive_watermark_key",
 ]
diff --git a/vllm/v1/watermarking/factory.py b/vllm/v1/watermarking/factory.py
index d9111e7a9e8c..2bc1919aef01 100644
--- a/vllm/v1/watermarking/factory.py
+++ b/vllm/v1/watermarking/factory.py
@@ -1,24 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from collections.abc import Callable
-
 from vllm.config.watermarking import WatermarkConfig
-from vllm.v1.watermarking.gumbel import GumbelWatermarker
+from vllm.v1.watermarking.gumbel import (
+    DualKeyGumbelWatermarker,
+    GumbelWatermarker,
+)
 from vllm.v1.watermarking.watermarker import Watermarker
 
 
-def _create_gumbel(config: WatermarkConfig) -> Watermarker:
-    return GumbelWatermarker(config.key, config.context_width, config.prf)
-
-
-_WATERMARKERS: dict[str, Callable[[WatermarkConfig], Watermarker]] = {
-    "gumbel": _create_gumbel,
-}
-
-
 def create_watermarker(config: WatermarkConfig) -> Watermarker:
-    watermarker_factory = _WATERMARKERS.get(config.algorithm)
-    if watermarker_factory is None:
-        raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}")
-    return watermarker_factory(config)
+    if config.algorithm == "gumbel":
+        return GumbelWatermarker(config.key, config.context_width, config.prf)
+    if config.algorithm == "dual_key_gumbel":
+        return DualKeyGumbelWatermarker(
+            config.key,
+            config.context_width,
+            config.prf,
+            config.alpha,
+        )
+    raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}")
diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py
index 7e943cc04035..a4e6e2dd5bb0 100644
--- a/vllm/v1/watermarking/gpu_sampler.py
+++ b/vllm/v1/watermarking/gpu_sampler.py
@@ -1,16 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from typing import Literal
+
 import numpy as np
 import torch
 
 from vllm.logger import init_logger
 from vllm.sampling_params import SamplingParams
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
-from vllm.v1.watermarking.watermarker import Watermarker
+from vllm.v1.watermarking.watermarker import RandomSampler, Watermarker
 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.sample.watermark import repeated_context_mask
 
 logger = init_logger(__name__)
 
@@ -20,20 +22,22 @@ def __init__(
         self,
         watermarker: Watermarker,
         *args,
+        deduplicate_contexts: Literal["none", "single_turn", "all"] = "single_turn",
+        deduplicate_contexts_max_history: int | None = 8192,
         **kwargs,
     ) -> None:
         super().__init__(*args, **kwargs)
         self.watermarker = watermarker
+        self.deduplicate_contexts = deduplicate_contexts
+        self.deduplicate_contexts_max_history = deduplicate_contexts_max_history
         self.watermarking = UvaBackedTensor(
             self.sampling_states.max_num_reqs, dtype=torch.bool
         )
         self.watermarking.np.fill(True)
         self.watermarking.copy_to_uva()
 
-    def add_request(
-        self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
-    ) -> None:
-        super().add_request(req_idx, prompt_len, sampling_params)
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
+        super().add_request(req_idx, sampling_params)
         self.watermarking.np[req_idx] = sampling_params.watermarking
         if sampling_params.watermarking and sampling_params.temperature == 0:
             logger.warning_once(
@@ -54,7 +58,7 @@ def _sample_random(
         pos: torch.Tensor,
         top_k: torch.Tensor | None,
         top_p: torch.Tensor | None,
-        use_flashinfer: bool,
+        use_fused_sampler: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         enabled = self.watermarking.np[idx_mapping_np] & (
             self.sampling_states.temperature.np[idx_mapping_np] != 0
@@ -67,42 +71,43 @@ def _sample_random(
                 pos,
                 top_k,
                 top_p,
-                use_flashinfer,
+                use_fused_sampler,
             )
 
         processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
         contexts = self._get_contexts(expanded_idx_mapping)
+        repeated_contexts = None
+        if self.deduplicate_contexts != "none":
+            repeated_contexts = self._get_repeated_contexts(
+                expanded_idx_mapping, contexts
+            )
 
-        def random_sample(sample_logits: torch.Tensor) -> torch.Tensor:
-            return gumbel_sample(
-                sample_logits,
-                expanded_idx_mapping,
-                self.sampling_states.temperature.gpu,
-                self.sampling_states.seeds.gpu,
-                pos,
-                apply_temperature=False,
-                is_drafting=False,
-                use_fp64=self.use_fp64_gumbel,
+        temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping]
+        needs_mixed_sampling = repeated_contexts is not None or not np.all(enabled)
+        skip_mask = None
+        if needs_mixed_sampling:
+            watermarking = self.watermarking.gpu[expanded_idx_mapping] & (
+                temperatures != 0
             )
+            if repeated_contexts is not None:
+                watermarking &= ~repeated_contexts
+            skip_mask = ~watermarking
 
+        random_sampler = RandomSampler(
+            expanded_idx_mapping=expanded_idx_mapping,
+            temperatures=self.sampling_states.temperature.gpu,
+            seeds=self.sampling_states.seeds.gpu,
+            positions=pos,
+            use_fp64=self.use_fp64_gumbel,
+        )
         output = self.watermarker.sample(
             processed_logits,
             contexts,
-            random_sample,
+            random_sampler,
+            skip_mask=skip_mask,
         )
-        temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping]
-        watermarking = self.watermarking.gpu[expanded_idx_mapping] & (temperatures != 0)
-        if not np.all(enabled):
-            unwatermarked = random_sample(processed_logits)
-            sampled = torch.where(watermarking, output.token_ids, unwatermarked)
-            output_logits = output.logits
-            if output.logits is not processed_logits:
-                output_logits = torch.where(
-                    watermarking.unsqueeze(-1), output.logits, processed_logits
-                )
-        else:
-            sampled = output.token_ids
-            output_logits = output.logits
+        sampled = output.token_ids
+        output_logits = output.logits
         sampled = torch.where(
             temperatures == 0,
             processed_logits.argmax(dim=-1),
@@ -110,22 +115,78 @@ def random_sample(sample_logits: torch.Tensor) -> torch.Tensor:
         )
         return sampled, output_logits
 
-    def _get_contexts(self, expanded_idx_mapping: torch.Tensor) -> torch.Tensor:
+    def _get_repeated_contexts(
+        self,
+        expanded_idx_mapping: torch.Tensor,
+        contexts: torch.Tensor,
+    ) -> torch.Tensor:
+        return repeated_context_mask(
+            self.req_states.all_token_ids.gpu,
+            expanded_idx_mapping,
+            self.req_states.prompt_len.gpu,
+            self.req_states.total_len.gpu,
+            contexts,
+            self.deduplicate_contexts_max_history,
+            include_prompt=self.deduplicate_contexts == "all",
+            skip_partial_context=self.deduplicate_contexts == "all",
+        )
+
+    def _get_contexts(
+        self,
+        expanded_idx_mapping: torch.Tensor,
+        expanded_local_pos: torch.Tensor | None = None,
+        draft_sampled: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Build contexts for flattened verification rows.
+
+        With committed tokens ``[10, 11]``, drafts ``[20, 21]``, and width 2,
+        local positions ``[0, 1, 2]`` use contexts ``[10, 11]``, ``[11, 20]``,
+        and ``[20, 21]``. The first verification row precedes the first draft.
+        """
         context_width = self.watermarker.context_width
         req_indices = expanded_idx_mapping.to(torch.int64)
         valid_reqs = req_indices >= 0
         safe_req_indices = req_indices.clamp_min(0)
         total_lens = self.req_states.total_len.gpu[safe_req_indices].to(torch.int64)
         prompt_lens = self.req_states.prompt_len.gpu[safe_req_indices].to(torch.int64)
+        if expanded_local_pos is None:
+            expanded_local_pos = torch.zeros_like(req_indices)
+        else:
+            expanded_local_pos = expanded_local_pos.to(torch.int64)
         offsets = torch.arange(
             -context_width, 0, dtype=torch.int64, device=req_indices.device
         )
-        positions = total_lens.unsqueeze(-1) + offsets
-        valid_positions = valid_reqs.unsqueeze(-1) & (
-            positions >= prompt_lens.unsqueeze(-1)
+        positions = (
+            total_lens.unsqueeze(-1) + expanded_local_pos.unsqueeze(-1) + offsets
         )
-        positions = positions.clamp_min(0)
-        contexts = self.req_states.all_token_ids.gpu[
-            safe_req_indices.unsqueeze(-1), positions
+        committed = positions < total_lens.unsqueeze(-1)
+        valid_committed = (
+            valid_reqs.unsqueeze(-1)
+            & committed
+            & (positions >= prompt_lens.unsqueeze(-1))
+        )
+        committed_contexts = self.req_states.all_token_ids.gpu[
+            safe_req_indices.unsqueeze(-1),
+            positions.clamp(0, self.req_states.all_token_ids.gpu.shape[1] - 1),
         ]
-        return torch.where(valid_positions, contexts, -1)
+        contexts = torch.where(valid_committed, committed_contexts, -1)
+        if draft_sampled is None:
+            return contexts
+
+        row_indices = torch.arange(
+            len(req_indices), dtype=torch.int64, device=req_indices.device
+        )
+        draft_offsets = positions - total_lens.unsqueeze(-1)
+        draft_indices = (
+            row_indices.unsqueeze(-1)
+            - expanded_local_pos.unsqueeze(-1)
+            + draft_offsets
+            + 1
+        )
+        valid_drafts = (
+            valid_reqs.unsqueeze(-1)
+            & ~committed
+            & (draft_offsets < expanded_local_pos.unsqueeze(-1))
+        )
+        draft_contexts = draft_sampled[draft_indices.clamp(0, len(draft_sampled) - 1)]
+        return torch.where(valid_drafts, draft_contexts, contexts)
diff --git a/vllm/v1/watermarking/gumbel.py b/vllm/v1/watermarking/gumbel.py
index 856d697aeaf9..cacb620b5e4b 100644
--- a/vllm/v1/watermarking/gumbel.py
+++ b/vllm/v1/watermarking/gumbel.py
@@ -8,13 +8,14 @@
 
 import torch
 
-from vllm.config.watermarking import WatermarkPRFName
+from vllm.config.watermarking import WatermarkPRFName, derive_watermark_key
 from vllm.v1.watermarking.detector import (
     WatermarkDetector,
 )
 from vllm.v1.watermarking.prfs import PhiloxPRF, WatermarkPRF, create_prf
 from vllm.v1.watermarking.watermarker import (
     RandomSampler,
+    SupportsSpeculativeDecoding,
     Watermarker,
     WatermarkSample,
 )
@@ -54,11 +55,10 @@ def __init__(
     def context_width(self) -> int:
         return self._context_width
 
-    def sample(
+    def _sample_watermarked(
         self,
         logits: torch.Tensor,
         contexts: torch.Tensor,
-        _random_sample: RandomSampler,
     ) -> WatermarkSample:
         if type(self.prf) is PhiloxPRF and logits.device.type == "cuda":
             from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
@@ -72,6 +72,110 @@ def sample(
         noise = -torch.log(-torch.log(uniforms))
         return WatermarkSample(torch.argmax(logits + noise, dim=-1), logits)
 
+    def _try_sample_mixed(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
+        skip_mask: torch.Tensor,
+        random_sampler: RandomSampler,
+    ) -> WatermarkSample | None:
+        if type(self.prf) is not PhiloxPRF or logits.device.type != "cuda":
+            return None
+
+        from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
+
+        token_ids = philox_gumbel_sample(
+            logits,
+            contexts,
+            self.prf.key,
+            skip_mask=skip_mask,
+            expanded_idx_mapping=random_sampler.expanded_idx_mapping,
+            temperatures=random_sampler.temperatures,
+            seeds=random_sampler.seeds,
+            positions=random_sampler.positions,
+            use_fp64=random_sampler.use_fp64,
+        )
+        return WatermarkSample(token_ids, logits)
+
+
+class DualKeyGumbelWatermarker(Watermarker, SupportsSpeculativeDecoding):
+    def __init__(
+        self,
+        key: int,
+        context_width: int = 4,
+        prf: WatermarkPRFName = "philox",
+        alpha: float = 0.1,
+    ) -> None:
+        if not 0 <= alpha <= 1:
+            raise ValueError("alpha must be between 0 and 1")
+        self.alpha = alpha
+        self.draft_watermarker = GumbelWatermarker(
+            derive_watermark_key(key, b"key_a"), context_width, prf
+        )
+        self.target_watermarker = GumbelWatermarker(
+            derive_watermark_key(key, b"key_b"), context_width, prf
+        )
+        self._routing_logits_cache: dict[torch.device, torch.Tensor] = {}
+
+    def _routing_logits(self, device: torch.device) -> torch.Tensor:
+        # Building this tensor on every CUDA call stages a pageable H2D copy and
+        # synchronizes the stream.
+        cached = self._routing_logits_cache.get(device)
+        if cached is None:
+            cached = torch.tensor(
+                [1 - self.alpha, self.alpha],
+                dtype=torch.float32,
+                device=device,
+            ).log()
+            self._routing_logits_cache[device] = cached
+        return cached
+
+    @property
+    def context_width(self) -> int:
+        return self.draft_watermarker.context_width
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
+        random_sampler: RandomSampler | None = None,
+        skip_mask: torch.Tensor | None = None,
+    ) -> WatermarkSample:
+        if self.alpha == 0:
+            return self.draft_watermarker.sample(
+                logits, contexts, random_sampler, skip_mask
+            )
+        if self.alpha == 1:
+            return self.target_watermarker.sample(
+                logits, contexts, random_sampler, skip_mask
+            )
+        if random_sampler is None:
+            raise ValueError("dual-key Gumbel routing requires a random sampler")
+
+        key_a_sample = self.draft_watermarker.sample(
+            logits, contexts, random_sampler, skip_mask
+        )
+        key_b_sample = self.target_watermarker.sample(
+            logits, contexts, random_sampler, skip_mask
+        )
+        routing_logits = self._routing_logits(logits.device).expand(logits.shape[0], -1)
+        use_key_a = random_sampler(routing_logits) == 0
+        return WatermarkSample(
+            torch.where(
+                use_key_a,
+                key_a_sample.token_ids,
+                key_b_sample.token_ids,
+            ),
+            logits,
+        )
+
+    def _sample_watermarked(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
+    ) -> WatermarkSample:
+        raise ValueError("dual-key Gumbel routing requires a random sampler")
+
 
 class GumbelWatermarkDetector(WatermarkDetector):
     def __init__(
@@ -100,6 +204,61 @@ def _aggregate_scores(self, token_scores: torch.Tensor) -> float:
         return token_scores.sum().item()
 
 
+class DualKeyGumbelWatermarkDetector(GumbelWatermarkDetector):
+    """Detect a dual-key watermark with weighted early fusion.
+
+    ``alpha`` is the weight assigned to key B; key A receives ``1 - alpha``.
+    """
+
+    def __init__(
+        self,
+        key: int,
+        context_width: int = 4,
+        p_value_threshold: float = 0.01,
+        prf: WatermarkPRFName = "philox",
+        deduplicate_contexts: bool = True,
+        alpha: float = 0.2,
+    ) -> None:
+        if not 0 <= alpha <= 1:
+            raise ValueError("alpha must be between 0 and 1")
+        super().__init__(
+            derive_watermark_key(key, b"key_a"),
+            context_width,
+            p_value_threshold,
+            prf,
+            deduplicate_contexts,
+        )
+        self.key_b_prf = create_prf(prf, derive_watermark_key(key, b"key_b"))
+        self.alpha = alpha
+
+    def _score_tokens(
+        self, contexts: torch.Tensor, targets: torch.Tensor
+    ) -> torch.Tensor:
+        scores = torch.stack(
+            [
+                -torch.log1p(
+                    -prf.uniform(contexts, targets.unsqueeze(-1))
+                    .squeeze(-1)
+                    .to(torch.float64)
+                )
+                for prf in (self.prf, self.key_b_prf)
+            ],
+            dim=-1,
+        )
+        weights = scores.new_tensor([1 - self.alpha, self.alpha])
+        return scores @ weights
+
+    def _get_p_value(self, score: float, num_scored_tokens: int) -> float:
+        variance = (1 - self.alpha) ** 2 + self.alpha**2
+        if self.alpha in (0, 0.5, 1):
+            return _gamma_survival_integer_shape(
+                score / variance, round(num_scored_tokens / variance)
+            )
+        shape = torch.tensor(num_scored_tokens / variance, dtype=torch.float64)
+        scaled_score = torch.tensor(score / variance, dtype=torch.float64)
+        return torch.special.gammaincc(shape, scaled_score).item()
+
+
 def _validate_context_width(context_width: int) -> None:
     if context_width < 1:
         raise ValueError("context_width must be positive")
diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py
new file mode 100644
index 000000000000..c3675a3ab953
--- /dev/null
+++ b/vllm/v1/watermarking/spec_decode.py
@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.config.watermarking import WatermarkConfig
+from vllm.v1.watermarking.factory import create_watermarker
+from vllm.v1.watermarking.gumbel import GumbelWatermarker
+from vllm.v1.watermarking.prfs import PhiloxPRF
+from vllm.v1.watermarking.watermarker import (
+    SupportsSpeculativeDecoding,
+    Watermarker,
+)
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample
+
+
+class DraftWatermarker:
+    def __init__(
+        self,
+        watermarker: Watermarker,
+        max_num_reqs: int,
+        device: torch.device,
+    ) -> None:
+        self.watermarker = watermarker
+        self.contexts = torch.zeros(
+            max_num_reqs,
+            watermarker.context_width,
+            dtype=torch.int64,
+            device=device,
+        )
+        self.enabled = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)
+
+    def prepare(self, contexts: torch.Tensor, enabled: torch.Tensor) -> None:
+        num_reqs = contexts.shape[0]
+        self.contexts[:num_reqs].copy_(contexts)
+        self.enabled[:num_reqs].copy_(enabled)
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        ordinary_sampled: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        temperature: torch.Tensor,
+    ) -> torch.Tensor:
+        request_temperatures = temperature[idx_mapping]
+        processed_logits = logits / torch.where(
+            request_temperatures == 0, 1, request_temperatures
+        ).unsqueeze(-1)
+        watermarked = self.watermarker.sample(
+            processed_logits,
+            self.contexts[: logits.shape[0]],
+        ).token_ids
+        enabled = self.enabled[: logits.shape[0]] & (request_temperatures != 0)
+        sampled = torch.where(enabled, watermarked, ordinary_sampled)
+        contexts = self.contexts[: logits.shape[0]]
+        contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1))
+        return sampled
+
+
+def create_speculative_target_watermarker(watermarker: Watermarker) -> Watermarker:
+    if isinstance(watermarker, SupportsSpeculativeDecoding):
+        return watermarker.target_watermarker
+    return watermarker
+
+
+def create_speculative_draft_watermarker(
+    watermarker: Watermarker,
+    max_num_reqs: int,
+    device: torch.device,
+    allow_target_only: bool,
+) -> DraftWatermarker | None:
+    if isinstance(watermarker, SupportsSpeculativeDecoding):
+        return DraftWatermarker(watermarker.draft_watermarker, max_num_reqs, device)
+    if allow_target_only:
+        return None
+    raise ValueError(
+        f"The {type(watermarker).__name__} watermarking algorithm does not support "
+        "speculative decoding. Set allow_target_only_watermarking=true to "
+        "leave draft tokens unwatermarked."
+    )
+
+
+def _philox_key(watermarker: Watermarker) -> int:
+    if not (
+        isinstance(watermarker, GumbelWatermarker)
+        and type(watermarker.prf) is PhiloxPRF
+    ):
+        raise NotImplementedError(
+            "In-kernel watermarked recovery supports only the Philox PRF "
+            "GumbelWatermarker"
+        )
+    return watermarker.prf.key
+
+
+def _resolve_watermark_key(watermarker: Watermarker) -> int:
+    """Return the Philox recovery key."""
+    if isinstance(watermarker, SupportsSpeculativeDecoding):
+        raise ValueError(
+            f"{type(watermarker).__name__} keys the target role separately from "
+            "the draft. Pass create_speculative_target_watermarker(watermarker) "
+            "so the recovery draw carries the target's key."
+        )
+    return _philox_key(watermarker)
+
+
+def speculative_target_watermark_key(
+    watermark_config: WatermarkConfig | None,
+) -> int | None:
+    """Resolve the resample kernel's Philox key from configuration."""
+    if watermark_config is None:
+        return None
+
+    return _resolve_watermark_key(
+        create_speculative_target_watermarker(create_watermarker(watermark_config))
+    )
+
+
+def watermarked_rejection_sample(
+    target_logits: torch.Tensor,
+    draft_logits: torch.Tensor | None,
+    draft_sampled: torch.Tensor,
+    cu_num_logits: torch.Tensor,
+    pos: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    expanded_local_pos: torch.Tensor,
+    temperature: torch.Tensor,
+    seed: torch.Tensor,
+    num_speculative_steps: int,
+    contexts: torch.Tensor,
+    watermarking: torch.Tensor,
+    watermarker: Watermarker,
+    use_fp64: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Use the target key for rejection recovery and bonus tokens."""
+    assert contexts.shape == (target_logits.shape[0], watermarker.context_width)
+    return rejection_sample(
+        target_logits,
+        draft_logits,
+        draft_sampled,
+        cu_num_logits,
+        pos,
+        idx_mapping,
+        expanded_idx_mapping,
+        expanded_local_pos,
+        temperature,
+        seed,
+        num_speculative_steps,
+        use_fp64=use_fp64,
+        contexts=contexts,
+        watermarking=watermarking,
+        watermark_key=_resolve_watermark_key(watermarker),
+    )
diff --git a/vllm/v1/watermarking/watermarker.py b/vllm/v1/watermarking/watermarker.py
index f00659393efd..12570f4ae5ca 100644
--- a/vllm/v1/watermarking/watermarker.py
+++ b/vllm/v1/watermarking/watermarker.py
@@ -2,13 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TypeAlias
+from typing import Protocol, runtime_checkable
 
 import torch
 
-RandomSampler: TypeAlias = Callable[[torch.Tensor], torch.Tensor]
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 
 
 @dataclass(frozen=True)
@@ -17,17 +16,81 @@ class WatermarkSample:
     logits: torch.Tensor
 
 
+@dataclass(frozen=True)
+class RandomSampler:
+    expanded_idx_mapping: torch.Tensor
+    temperatures: torch.Tensor
+    seeds: torch.Tensor
+    positions: torch.Tensor
+    use_fp64: bool = False
+
+    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
+        return gumbel_sample(
+            logits,
+            self.expanded_idx_mapping,
+            self.temperatures,
+            self.seeds,
+            self.positions,
+            apply_temperature=False,
+            is_drafting=False,
+            use_fp64=self.use_fp64,
+        )
+
+
 class Watermarker(ABC):
     @property
     @abstractmethod
     def context_width(self) -> int:
         raise NotImplementedError
 
-    @abstractmethod
     def sample(
         self,
         logits: torch.Tensor,
         contexts: torch.Tensor,
-        random_sample: RandomSampler,
+        random_sampler: RandomSampler | None = None,
+        skip_mask: torch.Tensor | None = None,
+    ) -> WatermarkSample:
+        if skip_mask is not None:
+            if random_sampler is None:
+                raise ValueError("skip_mask requires a random sampler")
+            mixed = self._try_sample_mixed(logits, contexts, skip_mask, random_sampler)
+            if mixed is not None:
+                return mixed
+
+        watermarked = self._sample_watermarked(logits, contexts)
+        if skip_mask is None:
+            return watermarked
+
+        assert random_sampler is not None
+        token_ids = torch.where(
+            skip_mask, random_sampler(logits), watermarked.token_ids
+        )
+        output_logits = watermarked.logits
+        if output_logits is not logits:
+            output_logits = torch.where(skip_mask.unsqueeze(-1), logits, output_logits)
+        return WatermarkSample(token_ids, output_logits)
+
+    @abstractmethod
+    def _sample_watermarked(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
     ) -> WatermarkSample:
         raise NotImplementedError
+
+    def _try_sample_mixed(
+        self,
+        logits: torch.Tensor,
+        contexts: torch.Tensor,
+        skip_mask: torch.Tensor,
+        random_sampler: RandomSampler,
+    ) -> WatermarkSample | None:
+        return None
+
+
+@runtime_checkable
+class SupportsSpeculativeDecoding(Protocol):
+    """A watermarker compatible with standard speculative rejection sampling."""
+
+    draft_watermarker: Watermarker
+    target_watermarker: Watermarker
diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py
index e8e6f808534f..c63908dad0d0 100644
--- a/vllm/v1/worker/block_table.py
+++ b/vllm/v1/worker/block_table.py
@@ -11,11 +11,11 @@
 
 from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
     TritonWarmupTensor,
     VllmTritonJitKernel,
-    kernel_launcher,
     triton_scalar_specialization_rep,
 )
 from vllm.triton_utils import tl, triton
@@ -67,20 +67,20 @@ def __init__(
         cp_kv_cache_interleave_size: int,
         slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT,
     ):
-        """
-        Args:
-            block_size: Block size used for KV cache memory allocation
-            max_num_reqs: Maximum number of concurrent requests supported.
-            max_num_blocks_per_req: Maximum number of blocks per request.
-            max_num_batched_tokens: Maximum number of tokens in a batch.
-            pin_memory: Whether to pin memory for faster GPU transfers.
-            device: Target device for the block table.
-            kernel_block_size: The block_size of underlying attention kernel.
-                Will be the same as `block_size` if `block_size` is supported
-                by the attention kernel.
-            slot_mapping_mode: How this cache group maps scheduled tokens to
-                cache slots. Mamba-like state caches do not use token slot
-                mappings and should use SlotMappingMode.NONE.
+        """Args:
+        block_size: Block size used for KV cache memory allocation
+        max_num_reqs: Maximum number of concurrent requests supported.
+        max_num_blocks_per_req: Maximum number of blocks per request.
+        max_num_batched_tokens: Maximum number of tokens in a batch.
+        pin_memory: Whether to pin memory for faster GPU transfers.
+        device: Target device for the block table.
+        kernel_block_size: The block_size of underlying attention kernel.
+            Will be the same as `block_size` if `block_size` is supported
+            by the attention kernel.
+        slot_mapping_mode: How this cache group maps scheduled tokens to
+            cache slots. Mamba-like state caches do not use token slot
+            mappings and should use SlotMappingMode.NONE.
+
         """
         self.max_num_reqs = max_num_reqs
         self.max_num_batched_tokens = max_num_batched_tokens
@@ -254,6 +254,7 @@ def map_to_kernel_blocks(
             # kv_manager_block_id 0 → kernel block id [0, 1]
             # kv_manager_block_id 1 → kernel block id [2, 3]
             # kv_manager_block_id 2 → kernel block id [4, 5]
+
         """
         if blocks_per_kv_block == 1:
             return kv_manager_block_ids
diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py
index 7e208da39327..79651759aad4 100644
--- a/vllm/v1/worker/cp_utils.py
+++ b/vllm/v1/worker/cp_utils.py
@@ -18,16 +18,20 @@
 logger = init_logger(__name__)
 
 
-def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
+def check_attention_cp_compatibility(
+    vllm_config: VllmConfig,
+    target_layer_names: set[str] | None = None,
+) -> None:
     pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
     dcp_size = vllm_config.parallel_config.decode_context_parallel_size
     interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size
     if pcp_size * dcp_size > 1:
         layer_type = cast(type[Any], AttentionLayerBase)
         layers = get_layers_from_vllm_config(vllm_config, layer_type)
-        for layer in layers.values():
+        for layer_name, layer in layers.items():
+            check_pcp = target_layer_names is None or layer_name in target_layer_names
             get_attn_backend = getattr(layer, "get_attn_backend", None)
-            if pcp_size > 1 and get_attn_backend is not None:
+            if pcp_size > 1 and check_pcp and get_attn_backend is not None:
                 backend = get_attn_backend()
                 assert backend.supports_pcp(), (
                     "PCP requires attention backend support, "
diff --git a/vllm/v1/worker/cpu/buffer_utils.py b/vllm/v1/worker/cpu/buffer_utils.py
index 8ee257a28f40..cdc76ab8b47f 100644
--- a/vllm/v1/worker/cpu/buffer_utils.py
+++ b/vllm/v1/worker/cpu/buffer_utils.py
@@ -4,13 +4,11 @@
 
 import torch
 
-from vllm.utils.platform_utils import is_uva_available
-
 
 class UvaBuffer:
     def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
-        if not is_uva_available():
-            raise RuntimeError("UVA is not available")
         self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
         self.np = self.cpu.numpy()
-        self.uva = self.cpu
+
+    def uva(self, n: int | None = None):
+        return self.cpu[:n] if n is not None else self.cpu
diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py
index f92c6c3bf28d..b615f4b5b5c8 100644
--- a/vllm/v1/worker/cpu/shm.py
+++ b/vllm/v1/worker/cpu/shm.py
@@ -76,14 +76,17 @@ def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]:
 
 def async_tensor_h2d(
     data: list | np.ndarray | torch.Tensor,
-    device: str | torch.device,
+    device: str | torch.device | None = None,
     dtype: torch.dtype | None = None,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     if isinstance(data, np.ndarray):
         data = torch.from_numpy(data)
-    if isinstance(data, torch.Tensor):
+    if not isinstance(data, torch.Tensor):
+        data = torch.tensor(data, dtype=dtype, device="cpu")
+    elif out is None:
         return data.to(dtype=dtype)
-    return torch.tensor(data, dtype=dtype, device="cpu")
+    return data if out is None else out.copy_(data)
 
 
 torch_utils.async_tensor_h2d = async_tensor_h2d
diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py
index 67ddf83baec5..04ca2b9194e4 100644
--- a/vllm/v1/worker/cpu_model_runner.py
+++ b/vllm/v1/worker/cpu_model_runner.py
@@ -80,16 +80,15 @@ def _postprocess_triton(self) -> None:
 
         # Speculative decoding fallbacks
         import vllm.v1.sample.rejection_sampler
-        import vllm.v1.spec_decode.llm_base_proposer
         import vllm.v1.spec_decode.utils as spec_decode_utils
 
-        vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_inputs_padded_kernel = (
+        spec_decode_utils._eagle_prepare_inputs_padded.kernel = (
             cpu_tl.eagle_prepare_inputs_padded_kernel
         )
-        vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_next_token_padded_kernel = (
+        spec_decode_utils._eagle_prepare_next_token_padded.kernel = (
             cpu_tl.eagle_prepare_next_token_padded_kernel
         )
-        vllm.v1.spec_decode.llm_base_proposer.copy_and_expand_eagle_inputs_kernel = (
+        spec_decode_utils._copy_and_expand_eagle_inputs.kernel = (
             cpu_tl.copy_and_expand_eagle_inputs_kernel
         )
         spec_decode_utils.copy_and_expand_dflash_inputs_kernel = (
diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py
index 2ee2164dc31e..e36b15300b25 100644
--- a/vllm/v1/worker/cpu_worker.py
+++ b/vllm/v1/worker/cpu_worker.py
@@ -7,7 +7,6 @@
 import math
 import os
 import sys
-from typing import Any
 
 import psutil
 import torch
@@ -16,7 +15,6 @@
 from vllm.config import CompilationMode, VllmConfig
 from vllm.logger import init_logger
 from vllm.platforms import CpuArchEnum, current_platform
-from vllm.profiler.wrapper import TorchProfilerWrapper
 from vllm.utils.cpu_resource_utils import (
     get_allowed_cpu_list,
     get_memory_node_info,
@@ -106,18 +104,6 @@ def __init__(
 
         self.parallel_config.disable_custom_all_reduce = True
 
-        # Torch profiler. Enabled and configured through profiler_config.
-        self.profiler: Any | None = None
-        profiler_config = vllm_config.profiler_config
-        if profiler_config.profiler == "torch":
-            worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
-            self.profiler = TorchProfilerWrapper(
-                profiler_config,
-                worker_name=worker_name,
-                local_rank=self.local_rank,
-                activities=["CPU"],
-            )
-
     def init_device(self):
         self.device = torch.device("cpu")
 
@@ -198,6 +184,9 @@ def wake_up(self, tags: list[str] | None = None) -> None:
         logger.warning("sleep mode is not supported on CPU, ignore it.")
         pass
 
+    def discard(self, tags: tuple[str, ...]) -> None:
+        raise RuntimeError("Selective discard is not supported on CPU.")
+
     def _should_warm_up_model(self) -> bool:
         # VLLM_CPU_CI_ENV always skips warmup to save CI time.
         if envs.VLLM_CPU_CI_ENV:
@@ -284,11 +273,3 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
             language_model=self.compilation_config.compilation_time,
             encoder=self.compilation_config.encoder_compilation_time,
         )
-
-    def profile(self, is_start: bool = True, profile_prefix: str | None = None):
-        if self.profiler is None:
-            raise RuntimeError("Profiler is not enabled.")
-        if is_start:
-            self.profiler.start()
-        else:
-            self.profiler.stop()
diff --git a/vllm/v1/worker/dp_utils.py b/vllm/v1/worker/dp_utils.py
index 07bda704833e..b2f2fe1506a9 100644
--- a/vllm/v1/worker/dp_utils.py
+++ b/vllm/v1/worker/dp_utils.py
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
+from contextvars import ContextVar
 
 import torch
 import torch.distributed as dist
@@ -18,6 +19,22 @@
 
 logger = init_logger(__name__)
 
+_skip_dp_coordination = ContextVar("skip_dp_coordination", default=False)
+
+
+@contextmanager
+def skip_dp_coordination():
+    """Run without coordinating DP metadata with other DP ranks."""
+    token = _skip_dp_coordination.set(True)
+    try:
+        yield
+    finally:
+        _skip_dp_coordination.reset(token)
+
+
+def should_skip_dp_coordination() -> bool:
+    return _skip_dp_coordination.get()
+
 
 def _get_device_and_group(parallel_config: ParallelConfig):
     # Use the actual device assigned to the DP group, not just the device type
@@ -95,8 +112,7 @@ def _post_process_dp_padding(tensor: torch.Tensor, should_dp_pad: bool) -> torch
 
 
 def _post_process_cudagraph_mode(tensor: torch.Tensor) -> int:
-    """
-    Synchronize cudagraph_mode across DP ranks by taking the minimum.
+    """Synchronize cudagraph_mode across DP ranks by taking the minimum.
     If any rank has NONE (0), all ranks use NONE.
     This ensures all ranks send consistent values (all padded or all unpadded).
     """
@@ -110,8 +126,7 @@ def _synchronize_dp_ranks(
     cudagraph_mode: int,
     parallel_config: ParallelConfig,
 ) -> tuple[bool, torch.Tensor | None, int]:
-    """
-    1. Decides if each DP rank is going to microbatch. Either all ranks
+    """1. Decides if each DP rank is going to microbatch. Either all ranks
     run with microbatching or none of them do.
 
     2. Determines the total number of tokens that each rank will run.
@@ -180,8 +195,7 @@ def coordinate_batch_across_dp(
     uniform_decode: bool | None = None,
     cudagraph_mode: int = 0,
 ) -> tuple[bool, torch.Tensor | None, int]:
-    """
-    Coordinates amongst all DP ranks to determine if and how the full batch
+    """Coordinates amongst all DP ranks to determine if and how the full batch
     should be split into microbatches.
 
     Args:
@@ -223,6 +237,19 @@ def coordinate_batch_across_dp(
     if num_tokens_padded is None:
         num_tokens_padded = num_tokens_unpadded
 
+    if should_skip_dp_coordination():
+        should_ubatch = should_attempt_ubatching and not is_last_ubatch_empty(
+            num_tokens_unpadded,
+            num_tokens_padded,
+            parallel_config.num_ubatches,
+        )
+        num_tokens_after_padding = torch.full(
+            (parallel_config.data_parallel_size,),
+            num_tokens_padded,
+            dtype=torch.int32,
+        )
+        return should_ubatch, num_tokens_after_padding, cudagraph_mode
+
     (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) = (
         _synchronize_dp_ranks(
             num_tokens_unpadded,
diff --git a/vllm/v1/worker/ec_connector_model_runner_mixin.py b/vllm/v1/worker/ec_connector_model_runner_mixin.py
index be676666e84d..8a32f996601a 100644
--- a/vllm/v1/worker/ec_connector_model_runner_mixin.py
+++ b/vllm/v1/worker/ec_connector_model_runner_mixin.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Define EC connector functionality mixin for model runners.
-"""
+"""Define EC connector functionality mixin for model runners."""
 
 from collections.abc import Generator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py
index a60e0d1892f0..3b587ad71f80 100644
--- a/vllm/v1/worker/encoder_cudagraph.py
+++ b/vllm/v1/worker/encoder_cudagraph.py
@@ -351,6 +351,7 @@ def _find_smallest_fitting_budget_given_tokens(
 
         Returns:
             Token budget if found, None if no fitting budget.
+
         """
         budgets = budgets if budgets is not None else self.token_budgets
         for budget in budgets:
@@ -411,8 +412,11 @@ def _run_budget_graph(
             mm_kwargs: Multimodal inputs for the batch.
             token_budget: Token budget to use.
             path: Configured encoder path.
+            axis_keys: Resolved capture-axis keys for this batch.
+
         Returns:
             Encoder outputs, or None if graph not captured.
+
         """
         graph_set = self._get_graph_set(path)
         num_items = len(self._get_item_specs(mm_kwargs))
@@ -562,6 +566,7 @@ def _dp_shard(
             images_per_rank: Number of items per rank.
             max_output_tokens_per_rank: Max output tokens across all ranks
                 (for padding during all_gather).
+
         """
         tp_size = get_tensor_model_parallel_world_size()
         current_rank = get_tensor_model_parallel_rank()
@@ -687,6 +692,7 @@ def execute(
 
         Returns:
             List of encoder outputs (one per item).
+
         """
         if self.use_dp:
             per_item_out_tokens = self._get_per_item_out_tokens(mm_kwargs)
diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py
index 0ab0ab3f3cd8..02592980a856 100644
--- a/vllm/v1/worker/gpu/async_utils.py
+++ b/vllm/v1/worker/gpu/async_utils.py
@@ -15,12 +15,12 @@
     LogprobsTensors,
     ModelRunnerOutput,
     PoolerOutput,
-    RoutedExpertsTensors,
 )
 from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
 from vllm.v1.worker.utils import raise_if_nan_logits
 
 if TYPE_CHECKING:
+    from vllm.distributed.aux_output_connector.worker import PendingAuxOutput
     from vllm.v1.worker.gpu.input_batch import InputBatch
 
 
@@ -120,8 +120,8 @@ def __init__(
         num_sampled_tokens: torch.Tensor,
         main_stream: torch.cuda.Stream,
         copy_stream: torch.cuda.Stream,
-        check_ep_fault: bool = False,
-        routed_experts: RoutedExpertsTensors | None = None,
+        check_ep_fault: bool,
+        pending_aux_output: "PendingAuxOutput | None",
     ):
         # NOTE(woosuk): We must retain references to the GPU tensors,
         # as the copy operations are performed on a different CUDA stream than
@@ -129,7 +129,7 @@ def __init__(
         self.model_runner_output = model_runner_output
         self.sampler_output = sampler_output
         self.num_sampled_tokens = num_sampled_tokens
-        self.routed_experts = routed_experts
+        self.pending_aux_output = pending_aux_output
         # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
         self.copy_event = torch.cuda.Event(blocking=True)
         self._has_fault: torch.Tensor | None = None
@@ -152,13 +152,15 @@ def __init__(
                 self.sampling_mask_tensors = (
                     sampler_output.sampling_mask_tensors.to_cpu_nonblocking()
                 )
-            self.routed_experts_cpu: RoutedExpertsTensors | None = None
-            if routed_experts is not None:
-                self.routed_experts_cpu = routed_experts.to_cpu_nonblocking()
             self.prompt_logprobs_dict = {
                 k: v.to_cpu_nonblocking() if v is not None else None
                 for k, v in self.model_runner_output.prompt_logprobs_dict.items()
             }
+            if self.pending_aux_output is not None:
+                self.pending_aux_output.enqueue_cpu_copy(
+                    num_sampled=self.num_sampled_tokens_np,
+                    num_rejected=async_copy_to_np(sampler_output.num_rejected),
+                )
             if check_ep_fault:
                 has_fault = get_ep_all2all_manager().query_fault()
                 self._has_fault = has_fault.to("cpu", non_blocking=True)
@@ -192,8 +194,10 @@ def get_output(self) -> ModelRunnerOutput:
         if self.logprobs_tensors is not None:
             self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
         self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
-        if self.routed_experts_cpu is not None:
-            self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists()
+        if self.pending_aux_output is not None:
+            self.model_runner_output.aux_output_connector_output = (
+                self.pending_aux_output.process_output()
+            )
 
         if self._has_fault is not None and self._has_fault.item():
             mask = get_ep_all2all_manager().query_active_mask()
diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py
index 05949c25ac21..e8b4da8f3629 100644
--- a/vllm/v1/worker/gpu/attn_utils.py
+++ b/vllm/v1/worker/gpu/attn_utils.py
@@ -3,11 +3,16 @@
 from collections.abc import Mapping, Sequence
 from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
+import numpy as np
 import torch
 
-from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.config import (
+    VllmConfig,
+    get_layers_from_vllm_config,
+)
+from vllm.config.compilation import CUDAGraphMode
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.multimodal.inputs import MultiModalFeatureSpec
@@ -15,6 +20,11 @@
     AttentionCGSupport,
     CommonAttentionMetadata,
 )
+from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend
+from vllm.v1.hisparse.binding import (
+    init_hisparse_kv_cache,
+    resolve_hisparse_block_size,
+)
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     KVCacheConfig,
@@ -27,10 +37,17 @@
     AttentionGroup,
     add_kv_sharing_layers_to_kv_cache_groups,
     allocate_kv_cache,
-    bind_kv_cache,
+    bind_kv_cache_to_layers,
     prepare_kernel_block_sizes,
 )
 
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.block_table import BlockTables
+    from vllm.v1.worker.gpu.cudagraph_utils import (
+        BatchExecutionDescriptor,
+        CudaGraphManager,
+    )
+
 
 @dataclass(frozen=True)
 class AttentionCGSupportInfo:
@@ -50,6 +67,67 @@ def narrow(
         return self
 
 
+@dataclass(frozen=True)
+class FastPrefillBatchMetadata:
+    """Per-step inputs for the KV-sharing fast prefill path."""
+
+    logits_indices_padded: torch.Tensor
+    num_logits_indices: int
+    max_logits_per_req: int
+
+
+class FastPrefillHelper:
+    """Decides per step whether to arm the KV-sharing fast prefill path, and
+    stages the logits indices it runs on.
+    """
+
+    def __init__(self, cudagraph_manager: "CudaGraphManager", max_num_tokens: int):
+        self.max_num_tokens = max_num_tokens
+        self.cudagraph_manager = cudagraph_manager
+        self.logits_indices_buf = torch.zeros(
+            max_num_tokens, dtype=torch.int32, device=cudagraph_manager.device
+        )
+
+    def prepare(
+        self,
+        logits_indices: torch.Tensor,
+        num_reqs: int,
+        cu_num_logits_np: np.ndarray,
+        has_prefill: bool,
+        batch_desc: "BatchExecutionDescriptor",
+        num_active_loras: int = 0,
+    ) -> FastPrefillBatchMetadata | None:
+        if (
+            not has_prefill
+            or batch_desc.cg_mode == CUDAGraphMode.FULL
+            or batch_desc.num_ubatches > 1
+        ):
+            return None
+        buf = self.logits_indices_buf
+        num_logits = logits_indices.shape[0]
+        assert num_logits > 0
+        buf[:num_logits].copy_(logits_indices)
+        # There might be leftover indices in buf[num_logits:] from previous
+        # iterations. Broadcast the scalar GPU-side to keep them valid.
+        buf[num_logits:] = logits_indices[-1]
+        # Pad so the model's KV-sharing layers run their reduced (logits-only)
+        # batch at a captured piecewise cudagraph size. Pre-capture, or with
+        # cudagraphs off, dispatch returns the unpadded count.
+        desc = self.cudagraph_manager.dispatch(
+            num_reqs=num_reqs,
+            num_tokens=num_logits,
+            uniform_token_count=None,
+            num_active_loras=num_active_loras,
+        )
+        num_logits_padded = min(desc.num_tokens, self.max_num_tokens)
+        return FastPrefillBatchMetadata(
+            logits_indices_padded=buf[:num_logits_padded],
+            num_logits_indices=num_logits,
+            # Largest per-request logits count, known on the host.
+            max_logits_per_req=int(np.diff(cu_num_logits_np).max()),
+        )
+
+
 def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
     kv_cache_spec: dict[str, KVCacheSpec] = {}
     layer_type = cast(type[Any], AttentionLayerBase)
@@ -63,6 +141,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
             if isinstance(spec, AttentionSpec):
                 spec = attn_module.get_attn_backend().customize_spec(spec)
             kv_cache_spec[layer_name] = spec
+    resolve_hisparse_block_size(vllm_config, kv_cache_spec, attn_layers)
     return kv_cache_spec
 
 
@@ -75,11 +154,42 @@ def get_shared_kv_cache_layers(vllm_config: VllmConfig):
     }
 
 
+def get_kv_sharing_fast_prefill_eligible_layers(
+    vllm_config: VllmConfig, draft_layer_names: set[str] | None = None
+) -> set[str]:
+    """Trailing run of KV-sharing layers, eligible for fast prefill.
+
+    In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other similar
+    KV sharing setups, only the layers that generate KV caches are involved in
+    the prefill phase, enabling prefill to early exit. Layers are registered in
+    execution order, so the eligible layers are the contiguous suffix of
+    KV-sharing layers.
+
+    Speculator draft layers register after the target model's layers (and may
+    themselves share KV), so they are excluded from the walk.
+    """
+    if not vllm_config.cache_config.kv_sharing_fast_prefill:
+        return set()
+    shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
+    if not shared_kv_cache_layers:
+        return set()
+    eligible_layers: set[str] = set()
+    attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
+    for layer_name in reversed(attn_layers):
+        if draft_layer_names is not None and layer_name in draft_layer_names:
+            continue
+        if layer_name not in shared_kv_cache_layers:
+            break
+        eligible_layers.add(layer_name)
+    return eligible_layers
+
+
 def init_attn_backend(
     kv_cache_config: KVCacheConfig,
     vllm_config: VllmConfig,
     device: torch.device,
     active_layer_names: set[str] | None = None,
+    draft_layer_names: set[str] | None = None,
 ) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]:
     # Phase 1: discover attention groups for each kv cache group.
     attn_groups: list[list[AttentionGroup]] = []
@@ -89,12 +199,18 @@ def init_attn_backend(
     add_kv_sharing_layers_to_kv_cache_groups(
         get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups
     )
+    fast_prefill_eligible_layers = get_kv_sharing_fast_prefill_eligible_layers(
+        vllm_config, draft_layer_names
+    )
 
     # Phase 1: discover attention groups for each kv cache group.
     for kv_cache_group_id, kv_cache_group_spec in enumerate(
         kv_cache_config.kv_cache_groups
     ):
         layer_names = kv_cache_group_spec.layer_names
+        if not kv_cache_group_spec.kv_cache_spec.has_layer_views:
+            attn_groups.append([])
+            continue
         if active_layer_names is not None:
             layer_names = list(active_layer_names.intersection(layer_names))
 
@@ -106,6 +222,10 @@ def init_attn_backend(
 
         for layer_name in layer_names:
             attn_backend = attn_layers[layer_name].get_attn_backend()
+            if layer_name in fast_prefill_eligible_layers:
+                attn_backend = create_fast_prefill_custom_backend(
+                    "FastPrefill", attn_backend
+                )
 
             layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec
             if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -209,22 +329,34 @@ def get_query_lens_mismatch_unsupported_backend(
 
 
 def init_kv_cache(
-    runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
     forward_context: dict[str, Any],
     kv_cache_config: KVCacheConfig,
     device: torch.device,
     kernel_block_sizes: list[int],
     vllm_config: VllmConfig,
     kv_cache_allocation_context: AbstractContextManager | None = None,
+    *,
+    block_tables: "BlockTables | None" = None,
 ) -> dict[str, Any]:
     allocation_context = kv_cache_allocation_context or nullcontext()
     with allocation_context:
-        kv_caches = allocate_kv_cache(
-            kv_cache_config,
-            device,
-            vllm_config.cache_config.get_resolved_kv_cache_layout(),
-            kernel_block_sizes,
-        )
+        if vllm_config.attention_config.hisparse_config is not None:
+            assert block_tables is not None
+            kv_caches = init_hisparse_kv_cache(
+                kv_cache_config,
+                device,
+                kernel_block_sizes,
+                vllm_config,
+                forward_context,
+                block_tables,
+            )
+        else:
+            kv_caches = allocate_kv_cache(
+                kv_cache_config,
+                device,
+                vllm_config.cache_config.get_resolved_kv_cache_layout(),
+                kernel_block_sizes,
+            )
     for layer_name, target in get_shared_kv_cache_layers(vllm_config).items():
         kv_caches[layer_name] = kv_caches[target]
     # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per
@@ -235,10 +367,12 @@ def init_kv_cache(
         in ("longcat_flash", "longcat_flash_ngram")
         else 1
     )
-    bind_kv_cache(
-        kv_caches,
+    bindable_caches = {
+        name: cache for name, cache in kv_caches.items() if name in forward_context
+    }
+    bind_kv_cache_to_layers(
+        bindable_caches,
         forward_context,
-        runner_kv_caches,
         num_attn_module,
         kv_cache_groups=kv_cache_config.kv_cache_groups,
     )
@@ -270,6 +404,7 @@ def build_attn_metadata(
     kv_cache_config: KVCacheConfig,
     seq_lens_cpu_upper_bound: torch.Tensor | None = None,
     dcp_local_seq_lens: torch.Tensor | None = None,
+    dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None,
     positions: torch.Tensor | None = None,
     is_prefilling: torch.Tensor | None = None,
     mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,
@@ -278,16 +413,25 @@ def build_attn_metadata(
     causal: bool | torch.Tensor | Mapping[int, bool] = True,
     rswa_prefix_lens: torch.Tensor | None = None,
     ubatch_idx: int = 0,
+    fast_prefill: FastPrefillBatchMetadata | None = None,
+    req_idx: np.ndarray | None = None,
 ) -> dict[str, Any]:
     seq_lens = seq_lens[:num_reqs]
     if dcp_local_seq_lens is not None:
         dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]
+    if dcp_local_seq_lens_cpu_upper_bound is not None:
+        dcp_local_seq_lens_cpu_upper_bound = dcp_local_seq_lens_cpu_upper_bound[
+            :num_reqs
+        ]
     if seq_lens_cpu_upper_bound is not None:
         seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs]
 
     attn_metadata: dict[str, Any] = {}
+    token_to_req_indices: torch.Tensor | None = None
     num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
     for i in range(num_kv_cache_groups):
+        if not attn_groups[i]:
+            continue
         block_table = block_tables[i]
         slot_mapping = slot_mappings[i]
         # Per-group causal for hybrid drafters (mixed SWA/full attention).
@@ -305,6 +449,12 @@ def build_attn_metadata(
         group_is_prefilling = common_attn_metadata_extra_kwargs.pop(
             "is_prefilling", is_prefilling
         )
+        if fast_prefill is not None:
+            common_attn_metadata_extra_kwargs.update(
+                logits_indices_padded=fast_prefill.logits_indices_padded,
+                num_logits_indices=fast_prefill.num_logits_indices,
+                max_logits_per_req=fast_prefill.max_logits_per_req,
+            )
         common_attn_metadata = CommonAttentionMetadata(
             query_start_loc=query_start_loc_gpu,
             query_start_loc_cpu=query_start_loc_cpu,
@@ -318,10 +468,13 @@ def build_attn_metadata(
             slot_mapping=slot_mapping,
             causal=group_causal,
             dcp_local_seq_lens=dcp_local_seq_lens,
+            dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound,
             positions=positions,
             is_prefilling=group_is_prefilling,
             mm_req_doc_ranges=mm_req_doc_ranges,
             rswa_prefix_lens=rswa_prefix_lens,
+            req_idx=req_idx,
+            _token_to_req_indices_cache=token_to_req_indices,
             **common_attn_metadata_extra_kwargs,
         )
 
@@ -347,6 +500,7 @@ def build_attn_metadata(
                 )
             for layer_name in attn_group.layer_names:
                 attn_metadata[layer_name] = metadata
+        token_to_req_indices = common_attn_metadata._token_to_req_indices_cache
     return attn_metadata
 
 
diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py
index c01895e63545..096d9b99aa66 100644
--- a/vllm/v1/worker/gpu/block_table.py
+++ b/vllm/v1/worker/gpu/block_table.py
@@ -59,8 +59,7 @@ def __init__(
             self.block_tables.append(block_table)
 
         self.num_blocks = UvaBackedTensor(
-            (self.num_kv_cache_groups, self.max_num_reqs),
-            dtype=torch.int32,
+            (self.num_kv_cache_groups, self.max_num_reqs), dtype=torch.int32
         )
         self.fused_writer: FusedStagedWriter | None = None
         if self.num_kv_cache_groups > 1:
@@ -315,6 +314,9 @@ def _compute_slot_mappings_kernel(
     mapping_enabled = tl.load(slot_mapping_enabled + group_id)
 
     req_state_idx = tl.load(idx_mapping + batch_idx)
+    # idx_mapping == -1 marks a dummy (or CUDA-graph padding) request that owns
+    # no blocks: never read its block-table row and emit PAD for its tokens.
+    is_real_req = req_state_idx >= 0
     start_idx = tl.load(query_start_loc + batch_idx)
     end_idx = tl.load(query_start_loc + batch_idx + 1)
     for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE):
@@ -342,12 +344,12 @@ def _compute_slot_mappings_kernel(
         block_offsets = local_positions % kernel_block_size
         block_numbers = tl.load(
             block_table_ptr + req_state_idx * block_table_stride + block_indices,
-            mask=is_local,
+            mask=is_local & is_real_req,
             other=0,
         )
         slot_ids = block_numbers * kernel_block_size + block_offsets
         if CP_SIZE != 1:
             slot_ids = tl.where(is_local, slot_ids, PAD_ID)
 
-        slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID)
+        slot_ids = tl.where(mapping_enabled & is_real_req, slot_ids, PAD_ID)
         tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py
index 2c1dac2d0710..0731072e5968 100644
--- a/vllm/v1/worker/gpu/buffer_utils.py
+++ b/vllm/v1/worker/gpu/buffer_utils.py
@@ -6,6 +6,7 @@
 import numpy as np
 import torch
 
+from vllm.logger import init_logger
 from vllm.triton_utils import tl, triton
 from vllm.utils.platform_utils import is_uva_available
 from vllm.utils.torch_utils import (
@@ -13,6 +14,8 @@
     get_accelerator_view_from_cpu_tensor,
 )
 
+logger = init_logger(__name__)
+
 # Default round-robin depth for the UVA buffer pools. Must be >= the number of
 # concurrent in-flight steps (engine batch_queue_size).
 _DEFAULT_MAX_CONCURRENCY = 2
@@ -23,31 +26,36 @@ def set_default_max_concurrency(n: int) -> None:
     _DEFAULT_MAX_CONCURRENCY = max(2, n)
 
 
-def async_copy_to_gpu(
-    x: torch.Tensor | np.ndarray,
-    out: torch.Tensor | None = None,
-    device: torch.device | None = None,
-) -> torch.Tensor:
-    if isinstance(x, np.ndarray):
-        x = torch.from_numpy(x)
-    assert x.is_cpu
-
-    if out is None:
-        assert device is not None
-        out = torch.empty_like(x, device=device)
-
-    # pin_memory() is no-op if the memory is already pinned.
-    pinned = x.pin_memory()
-    return out.copy_(pinned, non_blocking=True)
-
-
 class UvaBuffer:
     def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
         if not is_uva_available():
             raise RuntimeError("UVA is not available")
         self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True)
         self.np = self.cpu.numpy()
-        self.uva = get_accelerator_view_from_cpu_tensor(self.cpu)
+        self._uva = get_accelerator_view_from_cpu_tensor(self.cpu)
+
+    def uva(self, n: int | None = None) -> torch.Tensor:
+        return self._uva[:n] if n is not None else self._uva
+
+
+class NonUvaBuffer:
+    """Explicit-copy fallback for platforms without pinned memory."""
+
+    def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
+        from vllm.platforms import current_platform
+
+        logger.warning_once(
+            "Pinned memory is not available on this platform; falling back "
+            "to device memory for UVA buffers."
+        )
+        self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
+        self.np = self.cpu.numpy()
+        self._uva = torch.zeros(size, dtype=dtype, device=current_platform.device_type)
+
+    def uva(self, n: int | None = None) -> torch.Tensor:
+        if n is None:
+            return self._uva.copy_(self.cpu)
+        return self._uva[:n].copy_(self.cpu[:n])
 
 
 class UvaBufferPool:
@@ -69,7 +77,8 @@ def __init__(
         self.max_concurrency = max_concurrency
 
         # UVA buffers for concurrency
-        self._uva_bufs = [UvaBuffer(size, dtype) for _ in range(max_concurrency)]
+        self._buffer_cls = UvaBuffer if is_uva_available() else NonUvaBuffer
+        self._uva_bufs = [self._buffer_cls(size, dtype) for _ in range(max_concurrency)]
         # Current buffer index
         self._curr = 0
 
@@ -80,12 +89,12 @@ def copy_to_uva(self, x: torch.Tensor | np.ndarray | list) -> torch.Tensor:
         n = len(x)
         if n > buf.cpu.shape[0]:
             capacity = 1 << (n - 1).bit_length()
-            buf = UvaBuffer((capacity, *buf.cpu.shape[1:]), self.dtype)
+            buf = self._buffer_cls((capacity, *buf.cpu.shape[1:]), self.dtype)
             self._uva_bufs[self._curr] = buf
         # CPU-to-CPU copy
         dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np
         dst[:n] = x
-        return buf.uva[:n]
+        return buf.uva(n)
 
     def copy_to_gpu(
         self,
@@ -141,6 +150,9 @@ def __init__(
         self.device = device
         self.max_concurrency = max_concurrency
 
+        # Without UVA, large-but-cold tensors live on the GPU directly.
+        uva_instead_of_gpu = uva_instead_of_gpu and is_uva_available()
+
         if not uva_instead_of_gpu:
             # Create a GPU tensor (default)
             self.gpu = torch.zeros(size, dtype=dtype, device=device)
@@ -148,7 +160,7 @@ def __init__(
             # For a large but not-frequently-accessed tensor, we can use UVA instead of
             # GPU to save GPU memory
             self._uva_buf = UvaBuffer(size, dtype)
-            self.gpu = self._uva_buf.uva
+            self.gpu = self._uva_buf.uva()
 
         self._staged_write_indices: list[int] = []
         self._staged_write_starts: list[int] = []
diff --git a/vllm/v1/worker/gpu/cp_utils.py b/vllm/v1/worker/gpu/cp_utils.py
index 470271a95bd7..1c906f4214a6 100644
--- a/vllm/v1/worker/gpu/cp_utils.py
+++ b/vllm/v1/worker/gpu/cp_utils.py
@@ -5,7 +5,7 @@
 from vllm.triton_utils import tl, triton
 
 
-def maybe_prepare_dcp_local_seq_lens(
+def prepare_dcp_local_seq_lens(
     dcp_local_seq_lens: torch.Tensor,
     seq_lens: torch.Tensor,
     num_reqs: int,
@@ -16,8 +16,7 @@ def maybe_prepare_dcp_local_seq_lens(
     num_reqs_padded: int | None = None,
 ) -> torch.Tensor | None:
     """Populate caller-owned storage and return its padded view, or None without DCP."""
-    if dcp_size == 1:
-        return None
+    assert dcp_size > 1
 
     max_num_reqs = dcp_local_seq_lens.shape[0]
     BLOCK_SIZE = 128
diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py
index 287e64ce300c..8f8658077c13 100644
--- a/vllm/v1/worker/gpu/cudagraph_utils.py
+++ b/vllm/v1/worker/gpu/cudagraph_utils.py
@@ -4,7 +4,8 @@
 import itertools
 from collections import defaultdict
 from collections.abc import Callable, Iterable
-from dataclasses import dataclass
+from contextlib import ExitStack
+from dataclasses import dataclass, replace
 from itertools import groupby, product
 from typing import TYPE_CHECKING, Any, NamedTuple, Protocol
 
@@ -34,18 +35,21 @@
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import round_up
 from vllm.utils.torch_utils import current_stream
+from vllm.v1.hisparse.binding import release_hisparse_profiling_cache
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.interface import ModelState
+from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches
 from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches
 
 if TYPE_CHECKING:
     from vllm.v1.worker.gpu.model_runner import GPUModelRunner
     from vllm.v1.worker.gpu.pcp_manager import PCPManager
+    from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner
 
 logger = init_logger(__name__)
 
@@ -143,6 +147,7 @@ def __init__(
         decode_query_len: int,
         lora_capture_cases: list[int] | None = None,
         varlen_decode: bool = False,
+        ubatch_runner: "UBatchRunner | None" = None,
     ):
         self.vllm_config = vllm_config
         self.device = device
@@ -152,6 +157,8 @@ def __init__(
         self.cudagraph_mode = cudagraph_mode
         self.decode_query_len = decode_query_len
         self.varlen_decode = varlen_decode
+        # DBO supports FULL CUDA graphs only.
+        self.ubatch_runner = ubatch_runner
 
         self.dp_size = vllm_config.parallel_config.data_parallel_size
         self.tp_size = vllm_config.parallel_config.tensor_parallel_size
@@ -211,6 +218,33 @@ def _resolve_effective_loras(self, num_active_loras: int) -> int:
         # Counts above the largest captured case clamp to it.
         return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case)
 
+    def _maybe_ubatch_twin(
+        self, desc: BatchExecutionDescriptor
+    ) -> BatchExecutionDescriptor | None:
+        """Return a microbatched capture candidate when eligible.
+
+        Uniform query lengths preserve the captured request split. Use the DP
+        dispatch thresholds so all ranks generate the same candidates.
+        """
+        if self.ubatch_runner is None or desc.cg_mode != CUDAGraphMode.FULL:
+            return None
+        if desc.num_reqs is None:
+            return None
+        uniform_token_count, remainder = divmod(desc.num_tokens, desc.num_reqs)
+        if remainder or desc.uniform_token_count not in (None, uniform_token_count):
+            return None
+        parallel_config = self.vllm_config.parallel_config
+        num_ubatches = get_num_ubatches(parallel_config)
+        if desc.num_tokens < num_ubatches:
+            return None
+        if not check_ubatch_thresholds(
+            parallel_config, desc.num_tokens, uniform_decode=True
+        ):
+            return None
+        return replace(
+            desc, num_ubatches=num_ubatches, uniform_token_count=uniform_token_count
+        )
+
     def _init_candidates(self) -> None:
         """Build priority-ordered candidate lists for each token count."""
         capture_sizes = self.compilation_config.cudagraph_capture_sizes
@@ -299,6 +333,12 @@ def _init_candidates(self) -> None:
                     if desc not in descs_by_mode[decode_mode]:
                         descs_by_mode[decode_mode].append(desc)
 
+                    ubatch_desc = self._maybe_ubatch_twin(desc)
+                    if ubatch_desc is not None and (
+                        ubatch_desc not in descs_by_mode[decode_mode]
+                    ):
+                        descs_by_mode[decode_mode].append(ubatch_desc)
+
             # recoverSSM cannot capture a dummy query wider than its workspace.
             if mixed_mode and (
                 not self.vllm_config.cache_config.use_kda_recoverssm
@@ -320,6 +360,10 @@ def _init_candidates(self) -> None:
                 )
                 descs_by_mode[mixed_mode].append(desc)
 
+                ubatch_desc = self._maybe_ubatch_twin(desc)
+                if ubatch_desc is not None:
+                    descs_by_mode[mixed_mode].append(ubatch_desc)
+
         for mode, descs in descs_by_mode.items():
             descs.sort(key=lambda d: d.num_tokens, reverse=True)
             self._capture_descs[mode] = descs
@@ -343,6 +387,13 @@ def _init_candidates(self) -> None:
     def needs_capture(self) -> bool:
         return len(self._capture_descs) > 0
 
+    def _capture_stream(self, desc: BatchExecutionDescriptor) -> torch.cuda.Stream:
+        """Capture on the stream used by the microbatch threads."""
+        if desc.num_ubatches > 1:
+            assert self.ubatch_runner is not None
+            return self.ubatch_runner.capture_stream
+        return current_stream()
+
     @torch.inference_mode()
     def capture(
         self,
@@ -357,8 +408,13 @@ def capture(
                 it is invoked once with warmup=True and again with warmup=False
                 because attention backends may mutate or lazily initialize
                 metadata during warmup.
+            progress_bar_desc: Description shown on the capture progress bar.
+
         """
-        with graph_capture(device=self.device):
+        with graph_capture(device=self.device), ExitStack() as stack:
+            if self.ubatch_runner is not None:
+                # Join parked threads on failure to avoid blocking later captures.
+                stack.callback(self.ubatch_runner.abort_pending_run)
             # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger
             # activations so FULL activations should fit in already allocated
             # buffers in the graph pool.
@@ -386,7 +442,9 @@ def capture(
 
                     # Capture
                     logger.debug(
-                        "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc
+                        "CG Capture: mode=%s, batch_desc=%s",
+                        desc.cg_mode.name,
+                        desc,
                     )
                     if (
                         desc.cg_mode == CUDAGraphMode.PIECEWISE
@@ -414,7 +472,7 @@ def capture(
                             torch.accelerator.synchronize()
                             free_before = torch.accelerator.get_memory_info()[0]
                         with torch.cuda.graph(
-                            graph, self.pool, stream=current_stream()
+                            graph, self.pool, stream=self._capture_stream(desc)
                         ):
                             forward_fn(CUDAGraphMode.NONE)
                             # Join offloader's copy stream after forward to avoid
@@ -428,6 +486,7 @@ def capture(
                             self._capture_mem_samples.append(free_before - free_after)
                         self.graphs[desc] = graph
                         compilation_counter.num_cudagraph_captured += 1
+
         self._graphs_captured = True
 
     def captured_token_counts(self) -> list[int]:
@@ -446,7 +505,6 @@ def dispatch(
         num_ubatches: int = 1,
     ) -> BatchExecutionDescriptor:
         """Find matching cudagraph descriptor from priority-ordered candidates."""
-
         effective_loras = self._resolve_effective_loras(num_active_loras)
         key = (num_tokens, effective_loras)
         if self._graphs_captured and num_tokens > 0 and key in self._candidates:
@@ -509,6 +567,7 @@ def __init__(
         decode_query_len: int,
         lora_capture_cases: list[int] | None = None,
         varlen_decode: bool = False,
+        ubatch_runner: "UBatchRunner | None" = None,
     ):
         super().__init__(
             vllm_config,
@@ -517,6 +576,7 @@ def __init__(
             decode_query_len,
             lora_capture_cases=lora_capture_cases,
             varlen_decode=varlen_decode,
+            ubatch_runner=ubatch_runner,
         )
         self.hidden_states: torch.Tensor | None = None
         self.aux_hidden_states: list[torch.Tensor] = []
@@ -553,9 +613,37 @@ def capture(
                 "Set VLLM_USE_BREAKABLE_CUDAGRAPH=1 or cudagraph_mode=NONE/FULL."
             )
 
+        def store_capture_output(num_tokens: int, model_output: Any) -> None:
+            """Copy outputs to persistent buffers, allocating on first use."""
+            if self.is_last_pp_rank:
+                # Last PP rank (common case).
+                if self.use_aux_hidden_state_outputs:
+                    hidden_states, aux_hidden_states = model_output
+                else:
+                    hidden_states = model_output
+                    aux_hidden_states = []
+                if self.hidden_states is None:
+                    self.hidden_states = torch.empty_like(hidden_states)
+                self.hidden_states[:num_tokens] = hidden_states
+                if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:
+                    self.aux_hidden_states = [
+                        torch.empty_like(x) for x in aux_hidden_states
+                    ]
+                for i, aux in enumerate(aux_hidden_states):
+                    self.aux_hidden_states[i][:num_tokens] = aux
+            else:
+                # Non-last PP rank.
+                assert isinstance(model_output, IntermediateTensors)
+                intermediate_tensors = model_output
+                if self.intermediate_tensors is None:
+                    self.intermediate_tensors = IntermediateTensors.empty_like(
+                        intermediate_tensors
+                    )
+                for k, v in intermediate_tensors.tensors.items():
+                    self.intermediate_tensors[k][:num_tokens] = v
+
         def create_forward_fn(
-            desc: BatchExecutionDescriptor,
-            warmup: bool,
+            desc: BatchExecutionDescriptor, warmup: bool
         ) -> Callable[[CUDAGraphMode], None]:
             num_tokens = desc.num_tokens
             num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
@@ -573,6 +661,7 @@ def create_forward_fn(
             model_inputs = {
                 "input_ids": input_buffers.input_ids[:num_tokens],
                 "positions": input_buffers.positions[:num_tokens],
+                "intermediate_tensors": None,
                 **model_state.prepare_dummy_inputs(num_reqs, num_tokens),
             }
             if not self.is_first_pp_rank:
@@ -582,6 +671,30 @@ def create_forward_fn(
                 assert intermediate_tensors is not None
                 model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens]
 
+            if desc.num_ubatches > 1:
+                # Prepare and park threads before capture; finish runs inside it.
+                assert self.ubatch_runner is not None
+                ubatch_state = self.ubatch_runner.prepare(
+                    InputBatch.make_dummy(num_reqs, num_tokens, input_buffers),
+                    block_tables.get_dummy_block_tables(num_reqs),
+                    block_tables.get_dummy_slot_mappings(num_tokens),
+                    cg_mode=CUDAGraphMode.FULL,
+                    for_capture=True,
+                )
+                # Capture with dummy rows marked as padding.
+                input_buffers.is_padding.fill_(True)
+                finish = self.ubatch_runner.begin_capturable_run(
+                    model, model_inputs, ubatch_state, for_capture=True
+                )
+
+                def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None:
+                    assert cg_mode != CUDAGraphMode.PIECEWISE, (
+                        "DBO does not support PIECEWISE cudagraphs"
+                    )
+                    store_capture_output(num_tokens, finish())
+
+                return ubatch_forward_fn
+
             attn_metadata, slot_mappings = prepare_inputs_to_capture(
                 num_reqs,
                 num_tokens,
@@ -591,7 +704,7 @@ def create_forward_fn(
                 attn_groups,
                 kv_cache_config,
                 full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
-                max_query_len=desc.max_query_len,
+                max_query_len=desc.max_query_len or desc.uniform_token_count,
                 pcp_manager=pcp_manager,
             )
 
@@ -628,32 +741,7 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None:
                     # model outputs. No need to keep track of the hidden states.
                     return None
 
-                if self.is_last_pp_rank:
-                    # Last PP rank (common case).
-                    if self.use_aux_hidden_state_outputs:
-                        hidden_states, aux_hidden_states = model_output
-                    else:
-                        hidden_states = model_output
-                        aux_hidden_states = []
-                    if self.hidden_states is None:
-                        self.hidden_states = torch.empty_like(hidden_states)
-                    self.hidden_states[:num_tokens] = hidden_states
-                    if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:
-                        self.aux_hidden_states = [
-                            torch.empty_like(x) for x in aux_hidden_states
-                        ]
-                    for i, aux in enumerate(aux_hidden_states):
-                        self.aux_hidden_states[i][:num_tokens] = aux
-                else:
-                    # Non-last PP rank.
-                    assert isinstance(model_output, IntermediateTensors)
-                    intermediate_tensors = model_output
-                    if self.intermediate_tensors is None:
-                        self.intermediate_tensors = IntermediateTensors.empty_like(
-                            intermediate_tensors
-                        )
-                    for k, v in intermediate_tensors.tensors.items():
-                        self.intermediate_tensors[k][:num_tokens] = v
+                store_capture_output(num_tokens, model_output)
 
             return forward_fn
 
@@ -687,27 +775,33 @@ def prepare_inputs_to_capture(
     max_query_len: int | None = None,
     pcp_manager: "PCPManager | None" = None,
 ) -> AttentionState:
+    if full_cudagraph and max_query_len is None:
+        # Mixed graphs can replay a single prefill spanning the entire batch,
+        # even when the dummy batch distributes one token to each request.
+        max_query_len = num_tokens
     input_batch = InputBatch.make_dummy(
         num_reqs, num_tokens, input_buffers, max_query_len=max_query_len
     )
-    input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
-    slot_mapping_provider: BlockTables | PCPManager = block_tables
     if pcp_manager is not None:
-        slot_mapping_provider = pcp_manager
-    slot_mappings = slot_mapping_provider.get_dummy_slot_mappings(num_tokens)
+        input_batch = pcp_manager.prepare_inputs_to_capture(input_batch)
+
+    block_table_provider = pcp_manager or block_tables
+    input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs)
+    slot_mappings = block_table_provider.get_dummy_slot_mappings(num_tokens)
     slot_mappings_by_layer = build_slot_mappings_by_layer(
         slot_mappings, kv_cache_config
     )
 
-    input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
-        input_buffers.dcp_local_seq_lens,
-        input_batch.seq_lens,
-        input_batch.num_reqs,
-        block_tables.cp_size,
-        block_tables.cp_rank,
-        block_tables.cp_interleave,
-        num_reqs_padded=input_batch.num_reqs_after_padding,
-    )
+    if block_tables.cp_size > 1:
+        input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
+            input_buffers.dcp_local_seq_lens,
+            input_batch.seq_lens,
+            input_batch.num_reqs,
+            block_tables.cp_size,
+            block_tables.cp_rank,
+            block_tables.cp_interleave,
+            num_reqs_padded=input_batch.num_reqs_after_padding,
+        )
 
     # NOTE(woosuk): Attention metadata is required not just by standard attention
     # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,
@@ -931,6 +1025,7 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
     if (model := getattr(runner, "model", None)) is not None:
         layers = itertools.chain(layers, model.modules())
     clear_layer_kv_caches(layers)
+    release_hisparse_profiling_cache(runner.compilation_config.static_forward_context)
     runner.cache_config.num_gpu_blocks = None
     runner.maybe_remove_all_loras(runner.lora_config)
     gc.collect()
diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py
index dddf630cc335..8eabc3adbf1b 100644
--- a/vllm/v1/worker/gpu/dp_utils.py
+++ b/vllm/v1/worker/gpu/dp_utils.py
@@ -10,6 +10,7 @@
 from vllm.config import ParallelConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.distributed.parallel_state import get_dp_group
+from vllm.v1.worker.dp_utils import should_skip_dp_coordination
 from vllm.v1.worker.gpu.cudagraph_utils import (
     BatchExecutionDescriptor,
     CudaGraphManager,
@@ -54,8 +55,7 @@ def sync_cudagraph_and_dp_padding(
     allow_ubatching: bool = False,
     uniform_decode: bool = False,
 ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]:
-    """
-    Coordinates the batch descriptor and DP padding across all ranks.
+    """Coordinates the batch descriptor and DP padding across all ranks.
 
     `parallel_config` is only needed to decide whether to microbatch, so callers
     that never do (`allow_ubatching=False`) can leave it out.
@@ -71,7 +71,10 @@ def sync_cudagraph_and_dp_padding(
     tensor[3][dp_rank] = max_query_len or -1  # (-1 means None)
     tensor[4][dp_rank] = int(allow_ubatching)
     tensor[5][dp_rank] = num_reqs
-    dist.all_reduce(tensor, group=group)
+    if should_skip_dp_coordination():
+        tensor[:] = tensor[:, dp_rank, None].clone()
+    else:
+        dist.all_reduce(tensor, group=group)
 
     num_tokens_across_dp = tensor[0]
     cg_mode_across_dp = tensor[1]
@@ -108,23 +111,45 @@ def sync_cudagraph_and_dp_padding(
             int(num_tokens_across_dp.min()),
             uniform_decode=uniform_decode_across_dp,
         ):
-            # Microbatching is all-or-nothing: the expert all-to-all is
-            # collective, so every rank splits and pads to the same token count.
-            # A rank too small to fill a microbatch does no work in it, like a
-            # dummy run. Microbatched steps run eager; nothing is captured yet.
+            # Expert all-to-all requires every rank to split and pad equally.
+            # Empty microbatches run as dummy batches.
             ubatch_num_tokens = int(num_tokens_across_dp.max())
-            return BatchExecutionDescriptor(
-                cg_mode=CUDAGraphMode.NONE,
-                num_tokens=ubatch_num_tokens,
-                num_reqs=num_reqs,
-                num_ubatches=get_num_ubatches(parallel_config),
-            ), DPSyncState(
+            num_ubatches = get_num_ubatches(parallel_config)
+            ubatch_desc = None
+            if cudagraph_manager is not None:
+                # Match a FULL microbatched CUDA graph, falling back to an eager
+                # descriptor if no graph matches.
+                ubatch_desc = cudagraph_manager.dispatch(
+                    num_reqs,
+                    ubatch_num_tokens,
+                    synced_uniform_token_count,
+                    num_active_loras=num_active_loras,
+                    num_ubatches=num_ubatches,
+                )
+                if 2 * int(num_tokens_across_dp.min()) < ubatch_desc.num_tokens:
+                    # If one rank has an empty second microbatch, run without
+                    # CUDA graphs.
+                    ubatch_desc = None
+            if ubatch_desc is None:
+                ubatch_desc = BatchExecutionDescriptor(
+                    cg_mode=CUDAGraphMode.NONE,
+                    num_tokens=ubatch_num_tokens,
+                    num_reqs=num_reqs,
+                    num_ubatches=num_ubatches,
+                )
+            # Refresh the token count to include CUDA graph padding.
+            ubatch_num_tokens = ubatch_desc.num_tokens
+            num_reqs = int(num_reqs_across_dp.max())
+            if ubatch_desc.cg_mode == CUDAGraphMode.FULL:
+                assert ubatch_desc.num_reqs is not None
+                num_reqs = ubatch_desc.num_reqs
+            return ubatch_desc, DPSyncState(
                 num_tokens_across_dp=torch.full_like(
                     num_tokens_across_dp, ubatch_num_tokens
                 ),
                 uniform_token_count=synced_uniform_token_count,
-                eager=True,
-                num_reqs=int(num_reqs_across_dp.max()),
+                eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE,
+                num_reqs=num_reqs,
             )
 
     synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item()))
@@ -222,6 +247,12 @@ def dispatch_cg_and_sync_dp(
         need_eager: Force `CUDAGraphMode.NONE` instead of dispatching.
         num_active_loras: Active LoRA count for this rank. Does not need
             cross-rank agreement; it never changes a bucket's token count.
+        parallel_config: Only needed to decide whether to microbatch, so
+            callers that never do (`allow_ubatching=False`) can leave it out.
+        allow_ubatching: Whether this rank's batch may be split into
+            microbatches. Agreed across ranks before it takes effect.
+        uniform_decode: Whether this rank's batch is a uniform decode, used
+            to pick the microbatching split.
         dp_sync: Agreement from a prior dispatch over this same batch, to reuse.
             Must come from a batch with this same padded `num_tokens` and the
             same `uniform_token_count`; `num_reqs` may differ, as neither
@@ -231,6 +262,7 @@ def dispatch_cg_and_sync_dp(
     Returns:
         (batch_desc, sync), where `sync` is this batch's agreement for a later
         dispatch to reuse. It is None when `dp_size` is 1 or no rank has work.
+
     """
     reuse_eager = dp_sync is not None and dp_sync.eager
 
diff --git a/vllm/v1/worker/gpu/ec_connector.py b/vllm/v1/worker/gpu/ec_connector.py
index 1285d9fd14db..72f09f74e742 100644
--- a/vllm/v1/worker/gpu/ec_connector.py
+++ b/vllm/v1/worker/gpu/ec_connector.py
@@ -29,10 +29,7 @@ def maybe_get_output(
     ) -> Generator[ECConnectorOutput | None, None, None]:
         yield None
 
-    def no_forward(
-        self,
-        scheduler_output: "SchedulerOutput",
-    ) -> ModelRunnerOutput:
+    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
         return EMPTY_MODEL_RUNNER_OUTPUT
 
 
@@ -80,13 +77,11 @@ def maybe_get_output(
             output.finished_sending, output.finished_recving = (
                 ec_connector.get_finished(scheduler_output.finished_req_ids)
             )
+            output.ec_connector_stats = ec_connector.get_ec_connector_stats()
             output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta()
             ec_connector.clear_connector_metadata()
 
-    def no_forward(
-        self,
-        scheduler_output: "SchedulerOutput",
-    ) -> ModelRunnerOutput:
+    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
         # EC send/recv even if no work to do.
         with self.maybe_get_output(scheduler_output) as ec_connector_output:
             pass
@@ -98,8 +93,7 @@ def no_forward(
 
 
 def get_ec_connector(
-    vllm_config: VllmConfig,
-    encoder_cache: "EncoderCache | None",
+    vllm_config: VllmConfig, encoder_cache: "EncoderCache | None"
 ) -> ECConnector:
     if (
         not has_ec_transfer()
diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py
index 24c0f0c4e1c1..d1bbd3338e07 100644
--- a/vllm/v1/worker/gpu/eplb_utils.py
+++ b/vllm/v1/worker/gpu/eplb_utils.py
@@ -77,19 +77,13 @@ def maybe_register_speculator(
         assert speculative_config is not None
         assert speculative_config.draft_model_config is not None
         assert self.state is not None
-        self.state.add_model(
-            draft_moe_model,
-            speculative_config.draft_model_config,
-        )
+        self.state.add_model(draft_moe_model, speculative_config.draft_model_config)
         speculator.set_eplb_state(self.state)
         self._has_registered_models = True
         return True
 
     def maybe_register_model(
-        self,
-        model: nn.Module,
-        model_config: Any,
-        load_dummy_weights: bool,
+        self, model: nn.Module, model_config: Any, load_dummy_weights: bool
     ) -> bool:
         if not self.parallel_config.enable_eplb or load_dummy_weights:
             return False
@@ -110,11 +104,7 @@ def maybe_start_async_loop(self, eplb_models_added: bool) -> None:
         if eplb_models_added and self.state is not None and self.state.is_async:
             self.state.start_async_loop()
 
-    def step(
-        self,
-        is_dummy: bool = False,
-        is_profile: bool = False,
-    ) -> None:
+    def step(self, is_dummy: bool = False, is_profile: bool = False) -> None:
         if (
             not self.parallel_config.enable_eplb
             or self.suppressed
@@ -123,11 +113,8 @@ def step(
         ):
             return
 
-        self.state.step(
-            is_dummy,
-            is_profile,
-            log_stats=self.parallel_config.eplb_config.log_balancedness,
-        )
+        log_stats = self.parallel_config.eplb_config.log_balancedness
+        self.state.step(is_dummy, is_profile, log_stats=log_stats)
 
     def prepare_forward(
         self,
diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py
index 2a38c3d71061..75bbef5d8d04 100644
--- a/vllm/v1/worker/gpu/input_batch.py
+++ b/vllm/v1/worker/gpu/input_batch.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -11,16 +12,12 @@
 from vllm.utils.math_utils import cdiv
 
 if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.attn_utils import FastPrefillBatchMetadata
     from vllm.v1.worker.gpu.block_table import BlockTables
 
 
 class InputBuffers:
-    def __init__(
-        self,
-        max_num_reqs: int,
-        max_num_tokens: int,
-        device: torch.device,
-    ):
+    def __init__(self, max_num_reqs: int, max_num_tokens: int, device: torch.device):
         self.max_num_reqs = max_num_reqs
         self.max_num_tokens = max_num_tokens
         self.device = device
@@ -109,6 +106,13 @@ class InputBatch:
     # stays valid for every replay the graph serves.
     max_query_len: int | None = None
 
+    # Arms the KV-sharing fast prefill path for this step. Absent for dummy
+    # (cudagraph capture) batches, which run the KV-sharing layers in full.
+    fast_prefill: "FastPrefillBatchMetadata | None" = None
+
+    # [num_reqs] set only under PCP+DCP (see CommonAttentionMetadata).
+    dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None
+
     @classmethod
     def make_dummy(
         cls,
@@ -116,13 +120,14 @@ def make_dummy(
         num_tokens: int,
         input_buffers: InputBuffers,
         max_query_len: int | None = None,
+        is_padding: bool = True,
     ) -> "InputBatch":
         assert 0 < num_reqs <= num_tokens
         device = input_buffers.device
 
         req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
         idx_mapping_np = np.arange(num_reqs, dtype=np.intp)
-        idx_mapping = torch.arange(num_reqs, dtype=torch.int64, device=device)
+        idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
         expanded_idx_mapping = idx_mapping
         expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)
 
@@ -159,7 +164,7 @@ def make_dummy(
         input_ids = input_buffers.input_ids[:num_tokens].zero_()
         positions = input_buffers.positions[:num_tokens].zero_()
 
-        input_buffers.is_padding[:num_tokens].fill_(True)
+        input_buffers.is_padding[:num_tokens].fill_(is_padding)
         is_padding = input_buffers.is_padding[:num_tokens]
 
         logits_indices = query_start_loc[1:] - 1
@@ -209,9 +214,12 @@ def set_dummy_context(
     context_len: int,
     num_kv_blocks: int,
     max_model_len: int,
+    input_block_tables: Sequence[torch.Tensor] | None = None,
 ) -> None:
     """Give each dummy request context_len of context, used when profiling step cost."""
-    if not block_tables.input_block_tables:
+    if input_block_tables is None:
+        input_block_tables = block_tables.input_block_tables
+    if not input_block_tables:
         # Attention-free models have no KV context to fabricate.
         return
     num_reqs = input_batch.num_reqs
@@ -233,7 +241,7 @@ def set_dummy_context(
 
     seq_len = context_len + query_len
     for block_table, block_size, bpk in zip(
-        block_tables.input_block_tables,
+        input_block_tables,
         block_tables.kernel_block_sizes,
         block_tables.blocks_per_kv_block,
     ):
@@ -461,11 +469,7 @@ def combine_sampled_and_draft_tokens(
     num_reqs = idx_mapping.shape[0]
     num_speculative_steps = draft_tokens.shape[-1]
 
-    logits_indices = torch.empty(
-        num_logits,
-        dtype=torch.int64,
-        device=input_ids.device,
-    )
+    logits_indices = torch.empty(num_logits, dtype=torch.int64, device=input_ids.device)
     _combine_sampled_and_draft_tokens_kernel[(num_reqs,)](
         input_ids,
         idx_mapping,
@@ -698,9 +702,7 @@ def expand_idx_mapping(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     num_reqs = idx_mapping.shape[0]
     expanded_idx_mapping = idx_mapping.new_empty(total_num_logits)
-    expanded_local_pos = torch.empty(
-        total_num_logits, dtype=torch.int32, device=idx_mapping.device
-    )
+    expanded_local_pos = idx_mapping.new_empty(total_num_logits)
     _expand_idx_mapping_kernel[(num_reqs,)](
         idx_mapping,
         expanded_idx_mapping,
diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py
index eaef0c291661..88b7b0cf9d29 100644
--- a/vllm/v1/worker/gpu/kv_connector.py
+++ b/vllm/v1/worker/gpu/kv_connector.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 
@@ -29,7 +29,10 @@
 class KVConnector:
     """KVConnector interface used by GPUModelRunner."""
 
-    def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
+    def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None:
+        pass
+
+    def finish_forward(self) -> None:
         pass
 
     def post_forward(
@@ -43,6 +46,9 @@ def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
     def set_disabled(self, disabled: bool) -> None:
         pass
 
+    def reset_capture_state(self) -> None:
+        pass
+
 
 class ActiveKVConnector(KVConnector):
     def __init__(
@@ -54,10 +60,10 @@ def __init__(
         self.kv_connector.register_kv_caches(kv_caches_dict)
         self.kv_connector.set_host_xfer_buffer_ops(copy_kv_blocks)
 
-        self._pending_load_start = False
+        self._pending_load_kwargs: dict[str, Any] | None = None
         self._disabled = False
 
-    def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
+    def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None:
         if self._disabled:
             return
 
@@ -65,23 +71,31 @@ def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
         assert kv_connector_metadata is not None
         self.kv_connector.handle_preemptions(kv_connector_metadata)
         self.kv_connector.bind_connector_metadata(kv_connector_metadata)
+        self._pending_load_kwargs = kwargs
 
         if scheduler_output.has_sync_kv_loads:
             # Sync loads need to run before this step's forward.
             self._start_load_kv()
-        else:
-            # Start any async loads in post-forward instead, keeping
-            # their host-side submission cost off the critical path.
-            self._pending_load_start = True
+        # Otherwise start the async load after forward to keep submission
+        # off the critical path.
 
     def _start_load_kv(self) -> None:
-        self._pending_load_start = False
+        load_kwargs = self._pending_load_kwargs
+        assert load_kwargs is not None
+        self._pending_load_kwargs = None
         # TODO: sort out KV Connectors' use of forward_context
         if is_forward_context_available():
-            self.kv_connector.start_load_kv(get_forward_context())
+            self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs)
         else:
             with set_forward_context(None, self.vllm_config):
-                self.kv_connector.start_load_kv(get_forward_context())
+                self.kv_connector.start_load_kv(get_forward_context(), **load_kwargs)
+
+    def finish_forward(self) -> None:
+        if not self._disabled:
+            self.kv_connector.finish_forward()
+
+    def reset_capture_state(self) -> None:
+        self.kv_connector.reset_capture_state()
 
     def post_forward(
         self, finished_req_ids: set[str], wait_for_save: bool = True
@@ -89,15 +103,16 @@ def post_forward(
         if self._disabled:
             return None
 
-        if self._pending_load_start:
+        if self._pending_load_kwargs is not None:
             self._start_load_kv()
 
         output = KVConnectorOutput()
         if wait_for_save:
             self.kv_connector.wait_for_save()
-        output.finished_sending, output.finished_recving = (
-            self.kv_connector.get_finished(finished_req_ids)
-        )
+        transfer_results = self.kv_connector.get_transfer_results(finished_req_ids)
+        output.finished_sending = transfer_results.finished_sending or None
+        output.finished_recving = transfer_results.finished_recving or None
+        output.failed_recving = transfer_results.failed_recving
         output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors()
         output.kv_connector_stats = self.kv_connector.get_kv_connector_stats()
         output.kv_cache_events = self.kv_connector.get_kv_connector_kv_cache_events()
@@ -112,6 +127,7 @@ def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
             return EMPTY_MODEL_RUNNER_OUTPUT
 
         self.pre_forward(scheduler_output)
+        self.finish_forward()
         finished_req_ids = scheduler_output.finished_req_ids
         kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False)
         return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
diff --git a/vllm/v1/worker/gpu/lora_utils.py b/vllm/v1/worker/gpu/lora_utils.py
index 6a00b35e4bce..eb81a83662a1 100644
--- a/vllm/v1/worker/gpu/lora_utils.py
+++ b/vllm/v1/worker/gpu/lora_utils.py
@@ -21,8 +21,7 @@ def get_lora_capture_cases(
     lora_config: "LoRAConfig | None",
     compilation_config: "CompilationConfig",
 ) -> list[int]:
-    """
-    Return num_active_loras values for cudagraph capture.
+    """Return num_active_loras values for cudagraph capture.
 
     When cudagraph_specialize_lora=True: powers of 2 up to max_loras, plus
     max_loras+1. When False: [0, max_loras+1]. When LoRA disabled: [0].
diff --git a/vllm/v1/worker/gpu/mm/encoder_cache.py b/vllm/v1/worker/gpu/mm/encoder_cache.py
index 21d9218182ca..b9aba7b31a58 100644
--- a/vllm/v1/worker/gpu/mm/encoder_cache.py
+++ b/vllm/v1/worker/gpu/mm/encoder_cache.py
@@ -24,8 +24,7 @@ def remove_request(self, req_id: str) -> None:
         self.mm_features.pop(req_id, None)
 
     def reset_mm_cache(self) -> None:
-        """
-        Clear the multi-modal cache that was used during profiling,
+        """Clear the multi-modal cache that was used during profiling,
         but no longer needed during inference.
         """
         # NOTE: v2 encoder cache profiling skips the multi-modal cache
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index ec5ce6350377..5c617533f182 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-NOTE: Coding style guide for this file:
+"""NOTE: Coding style guide for this file:
 This model runner is shared by all models: text and multimodal, generative
 and embedding, public and private. As a result, this file must only contain
 code that is common to every model. Model-specific behavior belongs in the
@@ -33,17 +32,14 @@
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
-from vllm.distributed.parallel_state import (
-    get_dcp_group,
-    get_pp_group,
+from vllm.distributed.aux_output_connector.worker import (
+    AuxOutputWorkerConnector,
+    get_aux_output_connector,
 )
+from vllm.distributed.parallel_state import get_dcp_group, get_pp_group
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    RoutedExpertsCapturer,
-    bind_routed_experts_capturer,
-)
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     initialize_mamba_ssu_backend,
 )
@@ -58,16 +54,15 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.encoder_budget import (
     MultiModalBudget,
-    get_dummy_encoder_profile_inputs,
 )
+from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
 from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.kv_cache_interface import (
-    CircularBufferSpec,
     KVCacheConfig,
     MambaSpec,
     UniformTypeKVCacheSpecs,
@@ -76,10 +71,13 @@
     DraftTokenIds,
     ECConnectorOutput,
     ModelRunnerOutput,
-    RoutedExpertsTensors,
 )
 from vllm.v1.watermarking import create_watermarker
 from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
+from vllm.v1.watermarking.spec_decode import (
+    create_speculative_target_watermarker,
+    speculative_target_watermark_key,
+)
 from vllm.v1.worker.block_table import get_block_table_width
 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
 from vllm.v1.worker.gpu import pcp_manager as pcp
@@ -89,20 +87,19 @@
     StepTimingCollector,
 )
 from vllm.v1.worker.gpu.attn_utils import (
+    FastPrefillHelper,
     build_slot_mappings_by_layer,
     get_kv_cache_spec,
     init_attn_backend,
     init_kv_cache,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.buffer_utils import (
-    async_copy_to_gpu,
-    set_default_max_concurrency,
-)
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import (
     BatchExecutionDescriptor,
     ModelCudaGraphManager,
+    has_compiled_submodule,
     make_cudagraph_stats,
 )
 from vllm.v1.worker.gpu.cudagraph_utils import (
@@ -143,6 +140,7 @@
     all_to_all_logits,
     gather_sampler_output,
 )
+from vllm.v1.worker.gpu.sample.logits_processor import build_custom_logits_processors
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
 from vllm.v1.worker.gpu.sample.sampler import Sampler
@@ -151,6 +149,7 @@
 from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
     AdaptiveVerificationManager,
     maybe_create_adaptive_verification_manager,
+    resolve_adaptive_cudagraph_mode,
 )
 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
     set_eagle3_aux_hidden_state_layers,
@@ -256,9 +255,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
 
         # Multimodal
         self.mm_registry = MULTIMODAL_REGISTRY
-        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
-            self.model_config
-        )
+        self.supports_mm_inputs = self.model_config.supports_multimodal_inputs
         self.uses_inputs_embeds = (
             self.supports_mm_inputs or self.model_config.enable_prompt_embeds
         )
@@ -293,15 +290,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.is_pooling_model = self.model_config.runner_type == "pooling"
         self.pooling_runner: PoolingRunner | None = None
 
-        # Multi-module MTP feeds its modules the next num_speculative_steps prefill
-        # tokens during chunked prefill. Other speculators only read the immediate
-        # next one.
-        num_prefill_lookahead = (
-            self.num_speculative_steps
-            if self.speculative_config is not None
-            and self.speculative_config.use_multi_module_mtp()
-            else 1
-        )
+        num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens)
 
         self.step_timing = StepTimingCollector()
 
@@ -321,6 +310,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
             max_num_tokens=self.max_num_tokens,
             device=self.device,
         )
+        self.fast_prefill: FastPrefillHelper | None = None
         if self.use_pp:
             self.pp_handler = PPHandler(
                 max_num_reqs=self.max_num_reqs,
@@ -353,7 +343,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
 
         # Expert parallelism load balancer.
         self.eplb = EPLBController(self.parallel_config, self.device)
-        self.routed_experts_capturer: RoutedExpertsCapturer | None = None
+        # The AuxOutput Connector owns R3 capture, copying, and storage.
+        self.aux_output_connector: AuxOutputWorkerConnector | None = None
 
         set_offloader(create_offloader(self.vllm_config.offload_config))
 
@@ -361,15 +352,6 @@ def update_max_model_len(self, max_model_len: int) -> None:
         self.max_model_len = max_model_len
         self.req_states.max_model_len = max_model_len
 
-    def init_routed_experts_capturer(self) -> None:
-        """Initialize target-model capture on every participating worker."""
-        self.routed_experts_capturer = RoutedExpertsCapturer(
-            max_num_batched_tokens=self.max_num_tokens,
-            vllm_config=self.vllm_config,
-            kv_cache_config=self.kv_cache_config,
-        )
-        bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
-
     def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
         tasks: list[SupportedTask] = []
         if self.model_config.runner_type == "generate":
@@ -453,8 +435,23 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
                 )
 
         # Initialize samplers. Model states may override via custom_sampler().
+        config_processors = self.model_config.logits_processors or ()
+        custom_logits_processors = build_custom_logits_processors(
+            self.vllm_config, self.req_states, self.is_pooling_model, config_processors
+        )
         if self.is_last_pp_rank and not self.is_pooling_model:
+            # V2 bypasses TopKTopPSampler, which registers native warmups.
+            # ROCm only: on CUDA these warmups add ~2 min to every engine start.
+            if current_platform.is_rocm():
+                from vllm.v1.sample.ops.topk_topp_sampler import (
+                    register_top_k_top_p_warmups,
+                )
+
+                with self.jit_warmup_registry.activate():
+                    register_top_k_top_p_warmups()
+
             sampler_kwargs: dict[str, Any] = {
+                "vllm_config": self.vllm_config,
                 "max_num_reqs": self.max_num_reqs,
                 "vocab_size": self.vocab_size,
                 "device": self.device,
@@ -463,14 +460,24 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
                 "num_speculative_tokens": self.decode_query_len,
                 "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
                 "enable_trace_replay": self.model_config.enable_trace_replay,
-                "reasoning_config": self.vllm_config.reasoning_config,
                 "return_sampling_mask": self.model_config.return_sampling_mask,
+                "custom_logits_processors": custom_logits_processors,
             }
             if self.vllm_config.watermark_config is None:
                 self.sampler = Sampler(**sampler_kwargs)
             else:
-                watermarker = create_watermarker(self.vllm_config.watermark_config)
-                self.sampler = GPUWatermarkSampler(watermarker, **sampler_kwargs)
+                wm_config = self.vllm_config.watermark_config
+                watermarker = create_watermarker(wm_config)
+                if self.speculative_config is not None:
+                    watermarker = create_speculative_target_watermarker(watermarker)
+                self.sampler = GPUWatermarkSampler(
+                    watermarker,
+                    deduplicate_contexts=wm_config.deduplicate_contexts,
+                    deduplicate_contexts_max_history=(
+                        wm_config.deduplicate_contexts_max_history
+                    ),
+                    **sampler_kwargs,
+                )
             custom = self.model_state.custom_sampler(self.sampler)
 
             if custom:
@@ -481,10 +488,12 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
                     self.sampler,
                     self.speculative_config,
                     self.device,
+                    watermark_key=speculative_target_watermark_key(
+                        self.vllm_config.watermark_config
+                    ),
                 )
             self.prompt_logprobs_worker = PromptLogprobsWorker(
-                self.max_num_reqs,
-                logprobs_mode=self.model_config.logprobs_mode,
+                self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode
             )
             self.structured_outputs_worker = StructuredOutputsWorker(
                 max_num_logits=self.max_num_reqs * self.decode_query_len,
@@ -497,9 +506,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
         if self.is_pooling_model and self.is_last_pp_rank:
             self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
         eplb_models_added |= self.eplb.maybe_register_model(
-            self.model,
-            self.model_config,
-            load_dummy_weights,
+            self.model, self.model_config, load_dummy_weights
         )
         self.eplb.maybe_start_async_loop(eplb_models_added)
 
@@ -586,7 +593,7 @@ def initialize_kv_cache(
             layer_spec = (
                 spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
             )
-            slot_mapping_enabled.append(not isinstance(layer_spec, CircularBufferSpec))
+            slot_mapping_enabled.append(layer_spec.uses_slot_mapping)
             # Let each cache type account for CP. Attention KV is DCP-sharded,
             # while Mamba/GDN recurrent state is replicated across DCP ranks.
             max_num_blocks = spec.max_num_blocks_per_req(
@@ -594,12 +601,11 @@ def initialize_kv_cache(
             )
             # Preserve each cache type's alignment requirements after applying
             # its topology-aware block-table width.
-            if isinstance(layer_spec, (MambaSpec, CircularBufferSpec)):
-                max_num_blocks = get_block_table_width(
-                    max_num_blocks, spec.block_size, token_alignment=None
-                )
-            else:
-                max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size)
+            max_num_blocks = get_block_table_width(
+                max_num_blocks,
+                spec.block_size,
+                token_alignment=spec.block_table_token_alignment,
+            )
             max_num_blocks_per_group.append(max_num_blocks)
 
         target_attn_layer_names = None
@@ -610,11 +616,21 @@ def initialize_kv_cache(
                 for group in self.kv_cache_config.kv_cache_groups
                 for layer_name in group.layer_names
             } - self.speculator.draft_attn_layer_names
-        self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
-            self.kv_cache_config,
-            self.vllm_config,
-            self.device,
-        )
+        draft_attn_layer_names = None
+        if isinstance(self.speculator, DraftModelSpeculator):
+            draft_attn_layer_names = self.speculator.draft_attn_layer_names
+        # Metadata builders select attention kernels that need JIT warmup.
+        with self.jit_warmup_registry.activate():
+            (
+                self.attn_groups,
+                attn_cg_support,
+                self.kernel_block_sizes,
+            ) = init_attn_backend(
+                self.kv_cache_config,
+                self.vllm_config,
+                self.device,
+                draft_layer_names=draft_attn_layer_names,
+            )
         additional_attn_cg_support = self.model_state.get_additional_cg_support()
         attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
         # The speculator clears the flag at load time when the checkpoint has
@@ -650,7 +666,6 @@ def initialize_kv_cache(
             self.vllm_config,
             self.device,
             self.supports_mm_inputs,
-            self.req_states,
             self.block_tables,
             cls=self.pcp_manager_cls,
         )
@@ -662,13 +677,21 @@ def initialize_kv_cache(
             self.kv_cache_config,
             self.max_num_reqs,
         )
+        if self.speculator is not None:
+            self.speculator.pcp_manager = self.pcp_manager
         initialize_mamba_ssu_backend(
             self.vllm_config.mamba_config,
             self.kv_cache_config,
-            use_replayssm=self.vllm_config.cache_config.use_replayssm,
+            use_replayssm=self.cache_config.use_replayssm,
+        )
+        piecewise_capture_available = bool(
+            envs.VLLM_USE_BREAKABLE_CUDAGRAPH or has_compiled_submodule(self.model)
         )
         if self.adaptive_verification is not None:
-            self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
+            self.compilation_config.cudagraph_mode = resolve_adaptive_cudagraph_mode(
+                self.compilation_config.cudagraph_mode,
+                piecewise_capture_available=piecewise_capture_available,
+            )
         cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
             attn_cg_support.min_cg_support,
             attn_cg_support.min_cg_attn_backend,
@@ -678,6 +701,7 @@ def initialize_kv_cache(
             kv_cache_config=self.kv_cache_config,
             max_num_reqs=self.max_num_reqs,
             is_profiling=is_profiling,
+            piecewise_capture_available=piecewise_capture_available,
         )
         self.cudagraph_manager = ModelCudaGraphManager(
             self.vllm_config,
@@ -686,8 +710,13 @@ def initialize_kv_cache(
             decode_query_len=self.decode_query_len,
             lora_capture_cases=self.lora_capture_cases,
             varlen_decode=self.adaptive_verification is not None,
+            ubatch_runner=self.ubatch_runner,
         )
-        check_attention_cp_compatibility(self.vllm_config)
+        if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None:
+            self.fast_prefill = FastPrefillHelper(
+                self.cudagraph_manager, self.max_num_tokens
+            )
+        check_attention_cp_compatibility(self.vllm_config, target_attn_layer_names)
         if isinstance(self.speculator, DraftModelSpeculator):
             # HACK(woosuk)
             self.speculator.set_attn(
@@ -702,23 +731,31 @@ def initialize_kv_cache(
             # to its own attention support.
             self.speculator.init_cudagraph_manager(cudagraph_mode)
 
-        self.kv_caches: list[torch.Tensor] = []
         # Capture warmup providers that depend on allocated KV-cache strides.
         with self.jit_warmup_registry.activate():
             kv_caches_dict = init_kv_cache(
-                self.kv_caches,
                 self.compilation_config.static_forward_context,
                 self.kv_cache_config,
                 self.device,
                 self.kernel_block_sizes,
                 self.vllm_config,
                 kv_cache_allocation_context=kv_cache_allocation_context,
+                block_tables=self.block_tables,
             )
+        self.kv_caches = [
+            cache for cache in kv_caches_dict.values() if cache.device == self.device
+        ]
         if is_profiling:
             self.kv_connector = NO_OP_KV_CONNECTOR
         else:
             self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
 
+            # AuxOutput connector requires resolved kv_cache_config.
+            if self.vllm_config.aux_output_config.enabled:
+                self.aux_output_connector = get_aux_output_connector(
+                    self.model, self.vllm_config, kv_cache_config
+                )
+
     def _init_kv_zero_meta(self) -> None:
         """Build KV-block zeroing metadata; invoked from gpu_worker."""
         self.kv_block_zeroer = KVBlockZeroer(
@@ -758,6 +795,11 @@ def _dummy_run(
             num_tokens = max(num_tokens, self.decode_query_len)
             num_reqs = num_tokens // self.decode_query_len
             assert num_tokens % self.decode_query_len == 0
+        elif self.speculator is not None:
+            num_reqs = min(
+                num_reqs,
+                self.max_num_tokens // self.speculator.num_query_per_req,
+            )
         # Distribute the remainder evenly so no dummy request exceeds
         # ceil(num_tokens / num_reqs) <= max_model_len tokens.
         num_tokens_per_request = [
@@ -823,14 +865,10 @@ def _dummy_run(
             self.step_timing.drafter_start()
             mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
             if self.speculator.supports_mm_inputs:
-                mm_inputs = (
-                    [],
-                    torch.zeros(
-                        input_batch.num_tokens,
-                        dtype=torch.bool,
-                        device="cpu",
-                    ),
+                all_false = torch.zeros(
+                    input_batch.num_tokens, dtype=torch.bool, device="cpu"
                 )
+                mm_inputs = [], all_false
 
             # Let the target override the hidden state fed to the drafter
             # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
@@ -840,6 +878,11 @@ def _dummy_run(
             if hasattr(self.model, "get_mtp_target_hidden_states"):
                 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
                 spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]  # type: ignore[union-attr]
+            if isinstance(self.sampler, GPUWatermarkSampler):
+                self.speculator.prepare_watermarking(
+                    self.sampler._get_contexts(input_batch.idx_mapping),
+                    self.sampler.watermarking.gpu[input_batch.idx_mapping],
+                )
             with use_workspace_lane(self._draft_workspace_lane):
                 self.speculator.propose(
                     input_batch=input_batch,
@@ -894,14 +937,9 @@ def profile_run(self) -> None:
             mm_config = self.model_config.multimodal_config
             if mm_config is not None and not mm_config.skip_mm_profiling:
                 mm_budget = MultiModalBudget(
-                    self.vllm_config,
-                    self.mm_registry,
-                    enable_cache=False,
-                )
-                dummy_mm_inputs = get_dummy_encoder_profile_inputs(
-                    self.mm_registry,
-                    mm_budget,
+                    self.vllm_config, self.mm_registry, enable_cache=False
                 )
+                dummy_mm_inputs = mm_budget.get_dummy_encoder_profile_inputs()
                 self.model_state.encoder_runner.profile_encoder_cache(
                     dummy_mm_inputs, mm_budget
                 )
@@ -938,6 +976,16 @@ def profile_cudagraph_memory(self) -> int:
         """Estimate the GPU memory required to capture CUDA graphs."""
         return _profile_cudagraph_memory(self)
 
+    def needs_cudagraph_capture(self) -> bool:
+        """Whether capture_model() has any CUDA graphs left to capture."""
+        return (
+            self.model_state.supports_mm_inputs
+            and self.model_state.encoder_runner.has_cudagraph()
+        ) or (
+            self.cudagraph_manager is not None
+            and self.cudagraph_manager.needs_capture()
+        )
+
     @torch.inference_mode()
     def capture_model(self, *, profile_only: bool = False) -> int:
         assert self.cudagraph_manager is not None
@@ -994,6 +1042,7 @@ def capture_model(self, *, profile_only: bool = False) -> int:
                             ):
                                 self._dummy_run(**batch)
                         self.adaptive_verification.set_initial_cost_curves(timings)
+                    self.kv_connector.reset_capture_state()
 
             end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
 
@@ -1085,7 +1134,6 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None:
 
             if self.pooling_runner is not None:
                 assert new_req_data.pooling_params is not None
-                assert new_req_data.prompt_token_ids is not None
                 self.pooling_runner.add_request(
                     req_id,
                     req_index,
@@ -1104,9 +1152,7 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None:
 
             if self.is_last_pp_rank and new_req_data.sampling_params is not None:
                 assert self.sampler is not None
-                self.sampler.add_request(
-                    req_index, prompt_len, new_req_data.sampling_params
-                )
+                self.sampler.add_request(req_index, new_req_data.sampling_params)
                 assert self.prompt_logprobs_worker is not None
                 self.prompt_logprobs_worker.add_request(
                     req_id, req_index, new_req_data.sampling_params
@@ -1121,13 +1167,17 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None:
     def update_requests(self, scheduler_output: SchedulerOutput) -> None:
         # Add new blocks and update num_computed_tokens for the existing requests.
         reqs = scheduler_output.scheduled_cached_reqs
+        table_updates = scheduler_output.block_table_updates or {}
+        for req_id, block_ids in table_updates.items():
+            if (idx := self.req_states.req_id_to_index.get(req_id)) is not None:
+                self.block_tables.append_block_ids(idx, block_ids, overwrite=True)
         num_computed_tokens_np = self.req_states.num_computed_tokens_np
         for req_id, num_computed_tokens, req_new_block_ids in zip(
             reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids
         ):
             req_index = self.req_states.req_id_to_index[req_id]
             num_computed_tokens_np[req_index] = num_computed_tokens
-            if req_new_block_ids is not None:
+            if req_new_block_ids is not None and req_id not in table_updates:
                 self.block_tables.append_block_ids(
                     req_index, req_new_block_ids, overwrite=False
                 )
@@ -1207,32 +1257,42 @@ def gather_batch_req_state(
             num_reqs, num_toks, max_query_len, batch_state.has_prefill
         )
 
+    def _prepare_padding_mask(
+        self, num_tokens: int, num_tokens_after_padding: int
+    ) -> torch.Tensor:
+        is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
+        is_padding[:num_tokens].fill_(False)
+        is_padding[num_tokens:].fill_(True)
+        return is_padding
+
     def prepare_inputs(
         self,
         scheduler_output: SchedulerOutput,
         batch_req_state: "BatchReqState",
         batch_desc: BatchExecutionDescriptor,
+        num_active_loras: int,
     ) -> InputBatch:
         num_tokens = batch_req_state.num_tokens
         num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
         assert num_tokens > 0
+        is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
         if envs.VLLM_MOE_SKIP_PADDING:
-            # Mark trailing cudagraph-padding rows so kernels can skip work for
-            # them when supported.
-            is_padding = self.input_buffers.is_padding
-            is_padding[:num_tokens].fill_(False)
-            is_padding[num_tokens:num_tokens_after_padding].fill_(True)
+            is_padding = self._prepare_padding_mask(
+                num_tokens, num_tokens_after_padding
+            )
 
         req_ids = batch_req_state.req_ids
         num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
         idx_mapping_np = batch_req_state.idx_mapping_np
-        idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
+        idx_mapping = async_tensor_h2d(
+            idx_mapping_np, device=self.device, dtype=torch.int32
+        )
         num_reqs = len(req_ids)
 
         # Get the number of draft tokens for each request.
         draft_tokens = scheduler_output.scheduled_spec_decode_tokens
         num_draft_tokens_per_req = None
-        if not draft_tokens:
+        if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1:
             # No draft token scheduled (common case).
             total_num_draft_tokens = 0
             total_num_logits = num_reqs
@@ -1261,11 +1321,11 @@ def prepare_inputs(
             cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
             cu_num_logits_np[0] = 0
             np.cumsum(num_logits, out=cu_num_logits_np[1:])
-            cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
+            cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device)
 
-        adaptive_verification = (
-            self.adaptive_verification if num_draft_tokens_per_req is not None else None
-        )
+        # The general branch also serves a no-draft batch with k != 1, and
+        # compact_batch's budget is only primed when drafts are scheduled.
+        adaptive_verification = self.adaptive_verification if draft_tokens else None
         num_scheduled_tokens_upper_bound = num_scheduled_tokens_np
         if adaptive_verification is not None:
             # num_scheduled_tokens represents the draft budget evenly distributed across
@@ -1289,13 +1349,13 @@ def prepare_inputs(
         # Some attention backends like FA3 require query_start_loc to be non-decreasing.
         query_start_loc_np[num_reqs + 1 :] = num_tokens
         query_start_loc = self.input_buffers.query_start_loc
-        async_copy_to_gpu(query_start_loc_np, out=query_start_loc)
+        async_tensor_h2d(query_start_loc_np, out=query_start_loc)
         if adaptive_verification is not None:
             cu_num_logits, query_start_loc, total_num_draft_tokens = (
                 adaptive_verification.reallocate_drafts(req_ids, idx_mapping)
             )
             total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
-        if draft_tokens:
+        if num_draft_tokens_per_req is not None:
             expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
                 idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len
             )
@@ -1339,6 +1399,17 @@ def prepare_inputs(
             self.model_state.num_new_sampled_tokens_per_step,
         )
 
+        fast_prefill = None
+        if self.fast_prefill is not None:
+            fast_prefill = self.fast_prefill.prepare(
+                logits_indices,
+                num_reqs,
+                cu_num_logits_np,
+                batch_req_state.has_prefill,
+                batch_desc,
+                num_active_loras=num_active_loras,
+            )
+
         # CPU upper bound on seq_lens; padded entries left at zero.
         num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
         seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
@@ -1379,23 +1450,21 @@ def prepare_inputs(
             has_prefill=batch_req_state.has_prefill,
             input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
             positions=self.input_buffers.positions[:num_tokens_after_padding],
-            is_padding=self.input_buffers.is_padding[:num_tokens_after_padding],
+            is_padding=is_padding,
             logits_indices=logits_indices,
             cu_num_logits=cu_num_logits,
             cu_num_logits_np=cu_num_logits_np,
             has_structured_output_reqs=scheduler_output.has_structured_output_requests,
             prompt_lens=prompt_lens,
+            fast_prefill=fast_prefill,
             max_query_len=(
                 int(num_scheduled_tokens_upper_bound.max())
                 if adaptive_verification is not None
                 else None
             ),
         )
-        input_batch = pcp.maybe_partition_pcp_batch(
-            self.pcp_manager,
-            input_batch,
-            padded_num_tokens=batch_desc.num_tokens,
-        )
+        if self.pcp_manager is not None:
+            input_batch = self.pcp_manager.partition_batch(input_batch, batch_desc)
         return input_batch
 
     def prepare_attn(
@@ -1406,8 +1475,7 @@ def prepare_attn(
 
         # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
         block_tables = self.block_tables.gather_block_tables(
-            input_batch.idx_mapping,
-            num_reqs_padded=input_batch.num_reqs_after_padding,
+            input_batch.idx_mapping, num_reqs_padded=input_batch.num_reqs_after_padding
         )
         # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
         # Kernel pads beyond num_tokens with PAD_SLOT_ID.
@@ -1422,18 +1490,16 @@ def prepare_attn(
     def prepare_dummy_attn(
         self, input_batch: InputBatch, valid_state_slots: bool = False
     ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
-        block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
+        block_table_provider = self.pcp_manager or self.block_tables
+        block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs)
         if valid_state_slots:
             state_slots = torch.arange(
-                1,
-                input_batch.num_reqs + 1,
-                dtype=torch.int32,
-                device=self.device,
+                1, input_batch.num_reqs + 1, dtype=torch.int32, device=self.device
             )
             for block_table in block_tables:
                 block_table[:, 0].copy_(state_slots)
-        slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(
-            self.pcp_manager, self.block_tables, input_batch.num_tokens
+        slot_mappings = block_table_provider.get_dummy_slot_mappings(
+            input_batch.num_tokens
         )
         return block_tables, slot_mappings
 
@@ -1462,7 +1528,9 @@ def sample(
             sample_hidden_states = hidden_states[input_batch.logits_indices]
             logits = self.model.compute_logits(sample_hidden_states)
 
-        if grammar_output is not None:
+        # A diffusion prefill has no logit rows even when a bitmask row
+        # arrived for it.
+        if grammar_output is not None and logits.shape[0] > 0:
             # Apply grammar bitmask to the logits in-place.
             assert self.structured_outputs_worker is not None
             self.structured_outputs_worker.apply_grammar_bitmask(
@@ -1524,11 +1592,10 @@ def postprocess_sampled(
         query_start_loc: torch.Tensor | None = None,
     ) -> None:
         # Update the number of computed tokens.
+        output_bin_counts = None
         if self.is_last_pp_rank:
             assert self.sampler is not None
             output_bin_counts = self.sampler.penalties_state.output_bin_counts
-        else:
-            output_bin_counts = None
         post_update(
             idx_mapping,
             self.req_states.num_computed_tokens.gpu,
@@ -1574,6 +1641,11 @@ def execute_model(
             self.add_requests(scheduler_output)
             self.update_requests(scheduler_output)
             self.block_tables.apply_staged_writes()
+            if self.aux_output_connector is not None:
+                # Register this step before the GPU forward.
+                self.aux_output_connector.begin_step(
+                    scheduler_output.aux_output_connector_metadata
+                )
             if scheduler_output.total_num_scheduled_tokens == 0:
                 # No need to run the model.
                 empty_output = self.kv_connector.no_forward(scheduler_output)
@@ -1640,7 +1712,7 @@ def execute_model(
                 cudagraph_stats = make_cudagraph_stats(batch_desc, num_toks)
             assert batch_req_state is not None
             input_batch = self.prepare_inputs(
-                scheduler_output, batch_req_state, batch_desc
+                scheduler_output, batch_req_state, batch_desc, num_active_loras
             )
             block_tables, slot_mappings = self.prepare_attn(input_batch)
             # Mamba "align" pre-copy: migrate recurrent state across block
@@ -1670,8 +1742,19 @@ def execute_model(
                 batch_desc.num_tokens,
                 self.input_buffers,
                 max_query_len=batch_desc.max_query_len,
+                # Profiling and warmup must route the dummy tokens to experts
+                # so MoE memory is measured and MoE kernels are exercised.
+                is_padding=not is_profile,
             )
-            if not skip_attn_for_dummy_run:
+            if self.pcp_manager is not None:
+                input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)
+            if skip_attn_for_dummy_run:
+                assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
+                    "Attention metadata must be prepared for dummy runs when using "
+                    "FULL cudagraph mode."
+                )
+                block_tables, slot_mappings = None, None
+            else:
                 block_tables, slot_mappings = self.prepare_dummy_attn(
                     input_batch, valid_dummy_state_slots
                 )
@@ -1682,24 +1765,15 @@ def execute_model(
                         context_len,
                         self.kv_cache_config.num_blocks,
                         self.max_model_len,
+                        input_block_tables=block_tables,
                     )
-            else:
-                assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
-                    "Attention metadata must be prepared for dummy runs when using "
-                    "FULL cudagraph mode."
-                )
-                block_tables = None
-                slot_mappings = None
 
-        attn_metadata = None
-        slot_mappings_by_layer = None
-        if not (dummy_run and skip_attn_for_dummy_run):
-            input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
-                (
-                    self.pcp_manager.input_buffers
-                    if self.pcp_manager is not None
-                    else self.input_buffers
-                ).dcp_local_seq_lens,
+        if self.dcp_size > 1 and not (dummy_run and skip_attn_for_dummy_run):
+            input_buffers = self.input_buffers
+            if self.pcp_manager is not None:
+                input_buffers = self.pcp_manager.input_buffers
+            input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
+                input_buffers.dcp_local_seq_lens,
                 input_batch.seq_lens,
                 input_batch.num_reqs,
                 self.dcp_size,
@@ -1707,12 +1781,19 @@ def execute_model(
                 self.cp_interleave,
                 num_reqs_padded=input_batch.num_reqs_after_padding,
             )
+
+        attn_metadata = None
+        slot_mappings_by_layer = None
         ubatch_state: UBatchState | None = None
         if batch_desc.num_ubatches > 1:
             assert self.ubatch_runner is not None
             assert block_tables is not None and slot_mappings is not None
             ubatch_state = self.ubatch_runner.prepare(
-                input_batch, block_tables, slot_mappings
+                input_batch,
+                block_tables,
+                slot_mappings,
+                cg_mode=batch_desc.cg_mode,
+                for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
             )
         elif not (dummy_run and skip_attn_for_dummy_run):
             assert slot_mappings is not None
@@ -1810,13 +1891,22 @@ def execute_model(
         )
         self.step_timing.forward_start()
 
+        connector_kwargs = dict(
+            scheduler_output=scheduler_output,
+            request_state_indices=input_batch.idx_mapping,
+            request_ids=input_batch.req_ids,
+            num_tokens=input_batch.num_tokens,
+        )
+
         # Run model.
         if batch_desc.cg_mode == CUDAGraphMode.FULL:
             # Use explicit cudagraph replay for FULL mode.
             # NOTE(woosuk): Here, we don't need to pass the input tensors,
             # because they are already copied to the CUDA graph input buffers.
             assert self.cudagraph_manager is not None
-            self.kv_connector.pre_forward(scheduler_output)
+            self.kv_connector.pre_forward(
+                **connector_kwargs, attn_metadata=attn_metadata
+            )
             model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
         else:
             # For piecewise and eager mode, just call model().
@@ -1840,7 +1930,7 @@ def execute_model(
                 skip_compiled=skip_compiled,
                 is_padding=input_batch.is_padding,
             ):
-                self.kv_connector.pre_forward(scheduler_output)
+                self.kv_connector.pre_forward(**connector_kwargs)
                 if ubatch_state is not None:
                     assert self.ubatch_runner is not None
                     model_output = self.ubatch_runner.run(
@@ -1858,6 +1948,8 @@ def execute_model(
                     # Eager (NONE): call the raw model directly.
                     model_output = self.model(**model_inputs)
 
+        self.kv_connector.finish_forward()
+
         if self.is_last_pp_rank:
             if self.use_aux_hidden_state_outputs:
                 assert isinstance(model_output, tuple)
@@ -1873,11 +1965,6 @@ def execute_model(
             aux_hidden_states = None
             output_intermediate_tensors = model_output
 
-        routed_experts = None
-        if not dummy_run and (capturer := self.routed_experts_capturer) is not None:
-            assert slot_mappings is not None
-            routed_experts = capturer.get_routed_experts(slot_mappings, num_toks)
-
         finished_req_ids = scheduler_output.finished_req_ids
         self.execute_model_state = ExecuteModelState(
             input_batch=input_batch,
@@ -1888,7 +1975,6 @@ def execute_model(
             dp_sync=dp_sync,
             finished_req_ids=finished_req_ids,
             ec_connector_output=ec_connector_output,
-            routed_experts=routed_experts,
             cudagraph_stats=cudagraph_stats,
         )
 
@@ -1918,7 +2004,6 @@ def sample_tokens(
         dp_sync = self.execute_model_state.dp_sync
         finished_req_ids = self.execute_model_state.finished_req_ids
         ec_connector_output = self.execute_model_state.ec_connector_output
-        routed_experts = self.execute_model_state.routed_experts
         cudagraph_stats = self.execute_model_state.cudagraph_stats
         self.execute_model_state = None
 
@@ -1943,9 +2028,12 @@ def sample_tokens(
             return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
 
         # Last rank: sample tokens
-        hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(
-            self.pcp_manager, hidden_states, input_batch
-        )
+        assert hidden_states is not None
+        draft_hidden_states = hidden_states
+        if self.pcp_manager is not None:
+            hidden_states, aux_hidden_states, input_batch = (
+                self.pcp_manager.restore_for_sampling(hidden_states, aux_hidden_states)
+            )
 
         sampler_output, num_sampled, num_rejected = self.sample(
             hidden_states, input_batch, grammar_output
@@ -1980,6 +2068,10 @@ def sample_tokens(
             prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
             cudagraph_stats=cudagraph_stats,
         )
+        pending_aux_output = None
+        if self.aux_output_connector is not None:
+            pending_aux_output = self.aux_output_connector.prepare_output(input_batch)
+
         # Start async output copy here so that it can overlap with speculator proposal.
         async_output = AsyncOutput(
             model_runner_output=model_runner_output,
@@ -1988,11 +2080,17 @@ def sample_tokens(
             main_stream=self.main_stream,
             copy_stream=self.output_copy_stream,
             check_ep_fault=self.check_ep_fault,
-            routed_experts=routed_experts,
+            pending_aux_output=pending_aux_output,
         )
 
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
-        if self.speculator is not None and self.speculator.supports_mm_inputs:
+        # The encoder runner exists only on the first PP rank, so later ranks
+        # have no cached embeddings to gather.
+        if (
+            self.speculator is not None
+            and self.speculator.supports_mm_inputs
+            and self.model_state.supports_mm_inputs
+        ):
             # Get cached multimodal embeddings for draft forward.
             # NOTE: This is done here because postprocess updates
             # num_computed_prefill_tokens.
@@ -2018,14 +2116,23 @@ def sample_tokens(
 
         if self.speculator is not None:
             assert self.sampler is not None
+            if isinstance(self.speculator, DraftModelSpeculator):
+                self.speculator.observe_verification(
+                    input_batch.idx_mapping, num_sampled, num_rejected
+                )
             # Let the target override the hidden state fed to the drafter
             # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
             # target returns a persistent buffer sized at max_num_batched_tokens;
             # slice to the active token count that propose() expects.
-            spec_hidden_states = hidden_states
+            spec_hidden_states = draft_hidden_states
             if hasattr(self.model, "get_mtp_target_hidden_states"):
                 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
-                spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]  # type: ignore[union-attr]
+                spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
+            if isinstance(self.sampler, GPUWatermarkSampler):
+                self.speculator.prepare_watermarking(
+                    self.sampler._get_contexts(input_batch.idx_mapping),
+                    self.sampler.watermarking.gpu[input_batch.idx_mapping],
+                )
             with use_workspace_lane(self._draft_workspace_lane):
                 draft_tokens = self.speculator.propose(
                     input_batch,
@@ -2126,7 +2233,12 @@ def shutdown(self) -> None:
         """Release GPU tensors (model weights, KV caches, workspace) so that
         memory is reclaimable when running in the same process."""
         torch.accelerator.synchronize()
+        if self.aux_output_connector is not None:
+            self.aux_output_connector.close()
+        set_offloader(None)
         self.cudagraph_manager = None
+        self.fast_prefill = None
+        self.pooling_runner = None
         if hasattr(self, "kv_caches"):
             self.kv_caches.clear()
         if hasattr(self, "attn_groups"):
@@ -2171,13 +2283,10 @@ def eep_eplb_suppressed(self, suppressed: bool) -> None:
         self.eplb.suppressed = suppressed
 
     def setup_eplb_from_mapping(
-        self,
-        expanded_physical_to_logical: torch.Tensor,
+        self, expanded_physical_to_logical: torch.Tensor
     ) -> None:
         self.eplb.setup_from_mapping(
-            self.model,
-            self.model_config,
-            expanded_physical_to_logical,
+            self.model, self.model_config, expanded_physical_to_logical
         )
 
     ########### EPLB methods end ###########
@@ -2197,7 +2306,6 @@ class ExecuteModelState(NamedTuple):
     dp_sync: DPSyncState | None
     finished_req_ids: set[str]
     ec_connector_output: ECConnectorOutput | None
-    routed_experts: RoutedExpertsTensors | None
     cudagraph_stats: CUDAGraphStat | None
 
 
diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py
index 7d42745d1064..35ca9ccba102 100644
--- a/vllm/v1/worker/gpu/model_states/default.py
+++ b/vllm/v1/worker/gpu/model_states/default.py
@@ -16,7 +16,10 @@
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.mm.rope import get_rope_state
-from vllm.v1.worker.gpu.model_states.interface import ModelState
+from vllm.v1.worker.gpu.model_states.interface import (
+    ModelSpecificAttnMetadata,
+    ModelState,
+)
 from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner
 from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState
 from vllm.v1.worker.gpu.states import RequestState
@@ -173,6 +176,7 @@ def prepare_attn(
         kv_cache_config: KVCacheConfig,
         for_capture: bool = False,
         ubatch_idx: int = 0,
+        model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
     ) -> dict[str, Any]:
         if cudagraph_mode == CUDAGraphMode.FULL:
             # Use padded sizes - padding is handled by model_runner.prepare_attn.
@@ -220,11 +224,15 @@ def prepare_attn(
             kv_cache_config=kv_cache_config,
             seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
             dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
+            dcp_local_seq_lens_cpu_upper_bound=input_batch.dcp_local_seq_lens_cpu_upper_bound,
             positions=input_batch.positions,
             is_prefilling=torch.from_numpy(input_batch.is_prefilling_np),
             mm_req_doc_ranges=req_doc_ranges,
             for_cudagraph_capture=for_capture,
             rswa_prefix_lens=input_batch.prompt_lens,
+            req_idx=input_batch.idx_mapping_np,
             ubatch_idx=ubatch_idx,
+            fast_prefill=input_batch.fast_prefill,
+            model_specific_attn_metadata=model_specific_attn_metadata,
         )
         return attn_metadata
diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py
index a61384ed7f5a..bebb9cd41be2 100644
--- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py
+++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py
@@ -27,9 +27,7 @@ class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata):
     encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]]
 
     def get_extra_common_attn_kwargs(
-        self,
-        kv_cache_group_id: int,
-        num_reqs: int,
+        self, kv_cache_group_id: int, num_reqs: int
     ) -> dict[str, Any]:
         encoder_seq_lens = self.encoder_seq_lens.get(kv_cache_group_id)
         if encoder_seq_lens is None:
diff --git a/vllm/v1/worker/gpu/model_states/encoder_only.py b/vllm/v1/worker/gpu/model_states/encoder_only.py
index 32a70e7331ee..1b23723fdce8 100644
--- a/vllm/v1/worker/gpu/model_states/encoder_only.py
+++ b/vllm/v1/worker/gpu/model_states/encoder_only.py
@@ -19,6 +19,7 @@
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
 from vllm.v1.worker.gpu.states import RequestState
 from vllm.v1.worker.utils import AttentionGroup
 
@@ -168,6 +169,7 @@ def prepare_attn(
         kv_cache_config: KVCacheConfig,
         for_capture: bool = False,
         ubatch_idx: int = 0,
+        model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
     ) -> dict[str, Any]:
         assert ubatch_idx == 0, "DBO is not supported"
         attn_metadata = super().prepare_attn(
@@ -178,6 +180,7 @@ def prepare_attn(
             attn_groups,
             kv_cache_config,
             for_capture,
+            model_specific_attn_metadata=model_specific_attn_metadata,
         )
         attn_metadata.update(
             self._build_encoder_attn_metadata(input_batch, cudagraph_mode, for_capture)
@@ -185,10 +188,7 @@ def prepare_attn(
         return attn_metadata
 
     def _build_encoder_attn_metadata(
-        self,
-        input_batch: InputBatch,
-        cudagraph_mode: CUDAGraphMode,
-        for_capture: bool,
+        self, input_batch: InputBatch, cudagraph_mode: CUDAGraphMode, for_capture: bool
     ) -> dict[str, Any]:
         if cudagraph_mode == CUDAGraphMode.FULL:
             num_reqs = input_batch.num_reqs_after_padding
diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py
index a5c8d6f4cfd7..ee3e519a6e74 100644
--- a/vllm/v1/worker/gpu/model_states/interface.py
+++ b/vllm/v1/worker/gpu/model_states/interface.py
@@ -28,16 +28,12 @@ class ModelSpecificAttnMetadata:
     """Base class for model-specific attention metadata."""
 
     def get_extra_common_attn_kwargs(
-        self,
-        kv_cache_group_id: int,
-        num_reqs: int,
+        self, kv_cache_group_id: int, num_reqs: int
     ) -> dict[str, Any]:
         return {}
 
     def get_extra_attn_kwargs(
-        self,
-        attn_metadata_builder: Any,
-        num_reqs: int,
+        self, attn_metadata_builder: Any, num_reqs: int
     ) -> dict[str, Any]:
         return {}
 
@@ -59,7 +55,6 @@ def __init__(
         self.model = model
         self.device = device
 
-        self.max_model_len = self.model_config.max_model_len
         self.max_num_reqs = self.scheduler_config.max_num_seqs
         self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
         self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
@@ -99,6 +94,11 @@ def __init__(
                 ),
             )
 
+    @property
+    def max_model_len(self) -> int:
+        # Auto-fit can reduce the limit after model state initialization.
+        return self.model_config.max_model_len
+
     def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
         from vllm.model_executor.models.interfaces import (
             supports_realtime,
diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
index 0ac3f2f63b96..df527bd891b7 100644
--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
@@ -43,16 +43,12 @@ class MambaHybridAttnMetadata(ModelSpecificAttnMetadata):
     num_decode_draft_tokens_cpu: torch.Tensor | None = None
 
     def get_extra_common_attn_kwargs(
-        self,
-        kv_cache_group_id: int,
-        num_reqs: int,
+        self, kv_cache_group_id: int, num_reqs: int
     ) -> dict[str, Any]:
         return {"is_prefilling": self.is_prefilling[:num_reqs]}
 
     def get_extra_attn_kwargs(
-        self,
-        attn_metadata_builder: Any,
-        num_reqs: int,
+        self, attn_metadata_builder: Any, num_reqs: int
     ) -> dict[str, Any]:
         if not isinstance(
             attn_metadata_builder,
@@ -239,8 +235,10 @@ def prepare_attn(
         kv_cache_config: KVCacheConfig,
         for_capture: bool = False,
         ubatch_idx: int = 0,
+        model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
     ) -> dict[str, Any]:
         assert ubatch_idx == 0, "DBO is not supported"
+        assert model_specific_attn_metadata is None
         if cudagraph_mode == CUDAGraphMode.FULL:
             num_reqs = input_batch.num_reqs_after_padding
             num_tokens = input_batch.num_tokens_after_padding
@@ -248,7 +246,11 @@ def prepare_attn(
             num_reqs = input_batch.num_reqs
             num_tokens = input_batch.num_tokens
         query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
-        max_query_len = input_batch.num_scheduled_tokens.max().item()
+        # Prefer the promised bound: a capture dummy's measured max is its even
+        # split, not the length the graph must replay.
+        max_query_len = input_batch.max_query_len
+        if max_query_len is None:
+            max_query_len = input_batch.num_scheduled_tokens.max().item()
         seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound
         if for_capture:
             # Capture with worst-case max_seq_len so the graph is valid at any replay.
@@ -276,10 +278,11 @@ def prepare_attn(
             num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32)
             num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req
             if num_draft_tokens_per_req is not None:
-                # A row is a spec-decode row only when its whole prompt is already
-                # computed, i.e. exactly one non-draft (decode) token is scheduled.
-                is_decode = (
-                    input_batch.num_scheduled_tokens == num_draft_tokens_per_req + 1
+                # Test request state, not num_scheduled_tokens == draft_count+1:
+                # adaptive rewrites num_scheduled_tokens to an even split, so that
+                # equality rarely holds and would demote every verify row to decode.
+                is_decode = (~input_batch.is_prefilling_np) & (
+                    input_batch.num_scheduled_tokens > 0
                 )
                 spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode
                 num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(
@@ -331,9 +334,7 @@ def prepare_attn(
         )
         if self.recoverssm is not None:
             self.recoverssm.record_step(
-                attn_metadata,
-                attn_groups,
-                for_capture=for_capture,
+                attn_metadata, attn_groups, for_capture=for_capture
             )
         return attn_metadata
 
diff --git a/vllm/v1/worker/gpu/model_states/prompt_embeds.py b/vllm/v1/worker/gpu/model_states/prompt_embeds.py
index 251558f9c1ba..b9c82bb84958 100644
--- a/vllm/v1/worker/gpu/model_states/prompt_embeds.py
+++ b/vllm/v1/worker/gpu/model_states/prompt_embeds.py
@@ -37,8 +37,8 @@ def __init__(
 
         # Indexed by req_state index. Stale entries after removal are
         # harmless: add_request rewrites all fields for every index it claims.
-        self.embeds_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
-        self.mask_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
+        self.embeds_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.uint64)
+        self.mask_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.uint64)
         self.embeds_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
 
     def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
@@ -100,8 +100,8 @@ def apply(
 def _apply_prompt_embeds_kernel(
     inputs_embeds_ptr,
     inputs_embeds_stride,
-    embeds_ptrs_ptr,  # int64 [max_num_reqs], device pointers (0-len = unused)
-    mask_ptrs_ptr,  # int64 [max_num_reqs], 0 = no is-token-ids mask
+    embeds_ptrs_ptr,  # uint64 [max_num_reqs], device pointers (0-len = unused)
+    mask_ptrs_ptr,  # uint64 [max_num_reqs], 0 = no is-token-ids mask
     embeds_lens_ptr,  # int32 [max_num_reqs]
     idx_mapping_ptr,
     query_start_loc_ptr,
diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py
index abec427b3793..aff8f04cf300 100644
--- a/vllm/v1/worker/gpu/pcp_manager.py
+++ b/vllm/v1/worker/gpu/pcp_manager.py
@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterator
 from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
@@ -9,16 +10,16 @@
 from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group
 from vllm.logger import init_logger
-from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
 from vllm.v1.worker.gpu.input_batch import (
     InputBatch,
     InputBuffers,
-    combine_sampled_and_draft_tokens,
-    prepare_pos_seq_lens,
 )
-from vllm.v1.worker.gpu.states import RequestState
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 
 logger = init_logger(__name__)
 
@@ -48,7 +49,6 @@ def __init__(
         pcp_world_size: int,
         pcp_rank: int,
         device: torch.device,
-        req_states: RequestState | None = None,
         max_num_reqs: int | None = None,
         max_num_tokens: int | None = None,
         block_tables: BlockTables | None = None,
@@ -64,7 +64,9 @@ def __init__(
         self.cp_interleave = cp_interleave
 
         self._global_batch: InputBatch | None = None
-        self._req_states = req_states
+        self._local_batch: InputBatch | None = None
+        self._local_gather_idx: torch.Tensor | None = None
+        self.draft_prefill_batch: InputBatch | None = None
         self._block_tables = block_tables
         self._hidden_restore_idx: torch.Tensor | None = None
         self._padded_gather_idx: torch.Tensor | None = None
@@ -77,11 +79,6 @@ def __init__(
             if max_num_local_reqs is not None and max_num_tokens is not None
             else None
         )
-        self._local_req_idx = (
-            torch.arange(max_num_local_reqs, dtype=torch.int32, device=device)
-            if max_num_local_reqs is not None
-            else None
-        )
         self._local_block_tables: tuple[torch.Tensor, ...] | None
         self._local_block_table_ptrs: torch.Tensor | None
         if block_tables is not None and max_num_local_reqs is not None:
@@ -144,43 +141,70 @@ def validate_config(
             raise NotImplementedError("MRV2 PCP does not support MM inputs yet.")
         if vllm_config.lora_config is not None:
             raise NotImplementedError("MRV2 PCP does not support LoRA yet.")
-        if vllm_config.speculative_config is not None:
-            raise NotImplementedError(
-                "MRV2 PCP does not support speculative decoding yet."
-            )
+        speculative_config = vllm_config.speculative_config
+        if speculative_config is not None:
+            if speculative_config.use_dspark():
+                dcp_size = parallel_config.decode_context_parallel_size
+                if dcp_size not in (1, pcp_size):
+                    raise NotImplementedError(
+                        "MRV2 PCP DSpark requires DCP=1 or DCP=PCP; got "
+                        f"DCP={dcp_size}, PCP={pcp_size}."
+                    )
+            elif (
+                speculative_config.method != "mtp"
+                or speculative_config.use_multi_module_mtp()
+            ):
+                raise NotImplementedError(
+                    "MRV2 PCP only supports DSpark or single-module MTP "
+                    "speculative decoding."
+                )
+        cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
         is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk")
+        if parallel_config.decode_context_parallel_size > 1 and not is_sparse_mla:
+            # Dense MLA prefill sizes its DCP KV gather from each rank's own
+            # chunk rows, so the ranks' collectives diverge (#53573).
+            raise NotImplementedError("MRV2 PCP + DCP supports sparse MLA models only.")
         if (
             is_sparse_mla
-            and vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
+            and parallel_config.decode_context_parallel_size == 1
+            and cudagraph_mode != CUDAGraphMode.NONE
         ):
             raise NotImplementedError(
                 "MRV2 sparse MLA PCP does not support CUDA graphs yet. "
                 "Set -cc.cudagraph_mode=NONE."
             )
-        if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs():
-            raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.")
+        if (
+            cudagraph_mode.has_full_cudagraphs()
+            and not cudagraph_mode.separate_routine()
+        ):
+            raise NotImplementedError(
+                "MRV2 PCP supports full CUDA graphs for decode-only routines. "
+                "Use FULL_DECODE_ONLY, FULL_AND_PIECEWISE, PIECEWISE, or NONE."
+            )
+        if (
+            parallel_config.decode_context_parallel_size > 1
+            and parallel_config.dcp_comm_backend != "ag_rs"
+        ):
+            raise NotImplementedError(
+                "MRV2 PCP + DCP requires dcp_comm_backend='ag_rs'; got "
+                f"'{parallel_config.dcp_comm_backend}'."
+            )
 
     @staticmethod
     def _reorder_segments(
-        segments: list[RankSegment],
-        num_computed_tokens: np.ndarray,
-        is_prefilling: np.ndarray,
-        query_start_loc_np: np.ndarray,
+        segments: list[RankSegment], is_prefilling: np.ndarray
     ) -> list[RankSegment]:
-        """Move pure prefills last to match the batch ordering expected by
-        attention backends like MLA and sparse MLA.
-        """
+        """Order this rank's rows decodes-first, then prefills, canonically."""
 
-        def is_pure_prefill(segment: RankSegment) -> bool:
+        def sort_key(segment: RankSegment) -> tuple[bool, int, int]:
             req_idx = segment.global_batch_req_idx
-            start_pos = (
-                num_computed_tokens[req_idx]
-                + segment.global_batch_slice.start
-                - query_start_loc_np[req_idx]
+            return (
+                bool(is_prefilling[req_idx]),
+                req_idx,
+                segment.global_batch_slice.start,
             )
-            return is_prefilling[req_idx] and start_pos == 0
 
-        segments.sort(key=is_pure_prefill)
+        segments.sort(key=sort_key)
         rank_offset = 0
         for index, segment in enumerate(segments):
             segments[index] = replace(
@@ -192,6 +216,19 @@ def is_pure_prefill(segment: RankSegment) -> bool:
             rank_offset += segment.num_tokens
         return segments
 
+    def replicated_requests(
+        self, num_scheduled_tokens: np.ndarray, is_prefilling: np.ndarray
+    ) -> np.ndarray:
+        """Per global request, whether every PCP rank gets the whole query."""
+        num_chunks = 2 * self.pcp_world_size
+        query_lens = np.asarray(num_scheduled_tokens, dtype=np.int64)
+        replicated = ~np.asarray(is_prefilling, dtype=np.bool_)
+        if self.dcp_world_size > 1:
+            chunk_sizes = (query_lens + num_chunks - 1) // num_chunks
+            drops_a_chunk = (num_chunks - 1) * chunk_sizes >= query_lens
+            replicated |= drops_a_chunk
+        return replicated
+
     def _iter_rank_chunks(
         self,
         rank: int,
@@ -207,17 +244,20 @@ def _iter_rank_chunks(
             rank 1:      1                   6
             rank 2:          2           5
             rank 3:              3   4
+
+        Decodes, and prefills too short to fill all eight, are replicated instead.
         """
         num_chunks = 2 * self.pcp_world_size
+        replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling)
         for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens):
             query_len = int(num_tokens)
             if query_len == 0:
                 continue
             chunk_indices: tuple[int, ...]
-            if bool(is_prefilling[global_batch_req_idx]):
+            if not replicated[global_batch_req_idx]:
                 chunk_size = (query_len + num_chunks - 1) // num_chunks
                 chunk_indices = (rank, num_chunks - 1 - rank)
-            else:  # decodes are replicated
+            else:  # decodes, and short prefills under DCP, are replicated
                 chunk_size = query_len
                 chunk_indices = (0,)
 
@@ -232,7 +272,6 @@ def _get_rank_segments(
         self,
         rank: int,
         num_scheduled_tokens: np.ndarray,
-        num_computed_tokens: np.ndarray,
         is_prefilling: np.ndarray,
         query_start_loc_np: np.ndarray,
     ) -> list[RankSegment]:
@@ -251,12 +290,7 @@ def _get_rank_segments(
                 )
             )
             rank_offset += chunk_len
-        return self._reorder_segments(
-            rank_segments,
-            num_computed_tokens,
-            is_prefilling,
-            query_start_loc_np,
-        )
+        return self._reorder_segments(rank_segments, is_prefilling)
 
     def _build_batch_layout(
         self,
@@ -266,13 +300,13 @@ def _build_batch_layout(
         query_start_loc_np: np.ndarray,
         padded_num_tokens: int | None = None,
     ) -> tuple[list[list[RankSegment]], list[int]]:
+        replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling)
         segments_by_rank = []
         per_rank_num_tokens = []
         for rank in range(self.pcp_world_size):
             segments = self._get_rank_segments(
                 rank,
                 num_scheduled_tokens,
-                num_computed_tokens,
                 is_prefilling,
                 query_start_loc_np,
             )
@@ -312,7 +346,7 @@ def _build_batch_layout(
                     dtype=np.int64,
                 )
                 # Cache insertion pairs one slot entry with each rank's local decode.
-                if not bool(is_prefilling[segment.global_batch_req_idx]) and rank != 0:
+                if replicated[segment.global_batch_req_idx] and rank != 0:
                     continue
                 gathered_kv_write_mask[padded_gathered_slice] = True
                 hidden_restore_idx[segment.global_batch_slice] = np.arange(
@@ -321,13 +355,13 @@ def _build_batch_layout(
                     dtype=np.int64,
                 )
 
-        self._hidden_restore_idx = async_copy_to_gpu(
+        self._hidden_restore_idx = async_tensor_h2d(
             hidden_restore_idx, device=self.device
         )
-        self._padded_gather_idx = async_copy_to_gpu(
+        self._padded_gather_idx = async_tensor_h2d(
             padded_gather_idx, device=self.device
         )
-        self._gathered_kv_write_mask = async_copy_to_gpu(
+        self._gathered_kv_write_mask = async_tensor_h2d(
             gathered_kv_write_mask, device=self.device
         )
         return segments_by_rank, per_rank_num_tokens
@@ -348,22 +382,32 @@ def get_num_tokens_for_dispatch(
             for rank in range(self.pcp_world_size)
         )
 
+    @staticmethod
+    def _resolve_num_reqs_after_padding(
+        input_batch: InputBatch,
+        padded_num_reqs: int | None,
+        num_local_reqs: int,
+    ) -> int:
+        if padded_num_reqs is None:
+            return num_local_reqs
+        if input_batch.has_prefill:
+            raise RuntimeError("PCP FULL graphs require a decode-only batch.")
+        assert padded_num_reqs >= num_local_reqs, (
+            "PCP graph request capacity must cover the rank-local batch: "
+            f"{padded_num_reqs} < {num_local_reqs}."
+        )
+        return padded_num_reqs
+
     @property
     def input_buffers(self) -> InputBuffers:
         assert self._input_buffers is not None
         return self._input_buffers
 
     def partition_batch(
-        self,
-        input_batch: InputBatch,
-        padded_num_tokens: int | None = None,
+        self, input_batch: InputBatch, batch_desc: "BatchExecutionDescriptor"
     ) -> InputBatch:
-        assert self._req_states is not None
         assert self._input_buffers is not None
-        req_states = self._req_states
         input_buffers = self._input_buffers
-        if input_batch.num_draft_tokens > 0:
-            raise NotImplementedError("MRV2 PCP does not support spec decode yet.")
 
         global_batch = input_batch
         self._global_batch = global_batch
@@ -372,6 +416,13 @@ def partition_batch(
         num_computed_tokens = global_batch.num_computed_tokens_np
         is_prefilling = global_batch.is_prefilling_np
 
+        padded_num_tokens = None
+        padded_num_reqs = None
+        if batch_desc.cg_mode != CUDAGraphMode.NONE:
+            padded_num_tokens = batch_desc.num_tokens
+            if batch_desc.cg_mode == CUDAGraphMode.FULL:
+                padded_num_reqs = batch_desc.num_reqs
+
         segments_by_rank, per_rank_num_tokens = self._build_batch_layout(
             num_scheduled_tokens,
             num_computed_tokens,
@@ -391,10 +442,13 @@ def partition_batch(
             ]
 
         num_local_reqs = len(local_segments)
-        if num_local_reqs > input_buffers.max_num_reqs:
+        num_reqs_after_padding = self._resolve_num_reqs_after_padding(
+            global_batch, padded_num_reqs, num_local_reqs
+        )
+        if num_reqs_after_padding > input_buffers.max_num_reqs:
             raise RuntimeError(
-                "PCP local request count exceeds the MRV2 input buffer size: "
-                f"{num_local_reqs} > {input_buffers.max_num_reqs}."
+                "PCP padded local request count exceeds the MRV2 input buffer size: "
+                f"{num_reqs_after_padding} > {input_buffers.max_num_reqs}."
             )
 
         local_to_global_batch_req_idx_np = np.fromiter(
@@ -458,12 +512,21 @@ def partition_batch(
         local_gather_idx = self._padded_gather_idx[
             rank_token_start : rank_token_start + num_local_tokens_padded
         ]
+        self._local_gather_idx = local_gather_idx
         torch.index_select(
             global_batch.input_ids,
             0,
             local_gather_idx,
             out=input_buffers.input_ids[:num_local_tokens_padded],
         )
+        # Keep the GPU request-state cursor materialized by prepare_inputs().
+        # The CPU cursor can lag after speculative rejection.
+        torch.index_select(
+            global_batch.positions,
+            0,
+            local_gather_idx,
+            out=input_buffers.positions[:num_local_tokens_padded],
+        )
 
         local_query_start_loc_np = np.empty(
             input_buffers.max_num_reqs + 1, dtype=np.int32
@@ -472,23 +535,26 @@ def partition_batch(
         local_query_start_loc_out = local_query_start_loc_np[1 : num_local_reqs + 1]
         np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out)
         local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens
-        async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc)
-        local_query_start_loc = input_buffers.query_start_loc[: num_local_reqs + 1]
+        async_tensor_h2d(local_query_start_loc_np, out=input_buffers.query_start_loc)
+        local_query_start_loc = input_buffers.query_start_loc[
+            : num_reqs_after_padding + 1
+        ]
 
-        local_to_global_req_idx = async_copy_to_gpu(
+        local_to_global_req_idx = async_tensor_h2d(
             local_to_global_req_idx_np, device=self.device
         )
-        local_start_pos = async_copy_to_gpu(local_start_pos_np, device=self.device)
-
-        assert self._local_req_idx is not None
-        prepare_pos_seq_lens(
-            self._local_req_idx[:num_local_reqs],
-            local_query_start_loc,
-            local_start_pos,
-            input_buffers.positions,
-            input_buffers.seq_lens[:num_local_reqs],
-        )
-        seq_lens = input_buffers.seq_lens[:num_local_reqs]
+        seq_lens = input_buffers.seq_lens[:num_reqs_after_padding]
+        real_seq_lens = seq_lens[:num_local_reqs]
+        if num_local_tokens > 0:
+            local_end_positions = torch.index_select(
+                input_buffers.positions,
+                0,
+                local_query_start_loc[1 : num_local_reqs + 1] - 1,
+            )
+            real_seq_lens.copy_(local_end_positions + 1)
+        else:
+            real_seq_lens.zero_()
+        seq_lens[num_local_reqs:].zero_()
         is_padding = input_buffers.is_padding[:num_local_tokens_padded]
         is_padding[:num_local_tokens].fill_(False)
         is_padding[num_local_tokens:].fill_(True)
@@ -511,18 +577,9 @@ def partition_batch(
             cu_num_logits = torch.zeros(
                 num_local_reqs + 1, device=self.device, dtype=torch.int32
             )
-        logits_indices = combine_sampled_and_draft_tokens(
-            input_buffers.input_ids,
-            local_to_global_req_idx,
-            req_states.last_sampled_tokens,
-            local_query_start_loc,
-            seq_lens,
-            req_states.prefill_len.gpu,
-            req_states.draft_tokens,
-            cu_num_logits,
-            total_num_logits,
-            1,
-        )
+        # Local logits are never sampled. The complete hidden-state tensor is
+        # restored first and sampled with the untouched global InputBatch.
+        logits_indices = local_query_start_loc[1:] - 1
 
         local_prefill_len_np = global_batch.prefill_len_np[
             local_to_global_batch_req_idx_np
@@ -530,17 +587,40 @@ def partition_batch(
         local_num_computed_prefill_tokens_np = np.minimum(
             local_start_pos_np, local_prefill_len_np
         )
-        local_is_prefilling_np = (
+        real_local_is_prefilling_np = (
             local_num_computed_prefill_tokens_np < local_prefill_len_np
         )
-        seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32)
-        seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens
+        local_is_prefilling_np = np.zeros(num_reqs_after_padding, dtype=np.bool_)
+        local_is_prefilling_np[:num_local_reqs] = real_local_is_prefilling_np
+        seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_after_padding, dtype=np.int32)
+        seq_lens_cpu_upper_bound_np[:num_local_reqs] = (
+            local_start_pos_np + local_num_scheduled_tokens
+        )
+        dcp_local_seq_lens_cpu_upper_bound = None
+        if self.dcp_world_size > 1:
+            # The largest DCP shard of each row's whole request, identical on
+            # every PCP rank: the sparse backends pad their KV gather to it.
+            request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[
+                local_to_global_batch_req_idx_np
+            ]
+            real_dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens(
+                torch.from_numpy(request_seq_lens.astype(np.int32)),
+                self.dcp_world_size,
+                0,
+                self.cp_interleave,
+            )
+            dcp_local_seq_lens_cpu_upper_bound = torch.zeros(
+                num_reqs_after_padding, dtype=torch.int32
+            )
+            dcp_local_seq_lens_cpu_upper_bound[:num_local_reqs].copy_(
+                real_dcp_local_seq_lens_cpu_upper_bound
+            )
 
-        return replace(
+        self._local_batch = replace(
             input_batch,
             req_ids=local_req_ids,
             num_reqs=num_local_reqs,
-            num_reqs_after_padding=num_local_reqs,
+            num_reqs_after_padding=num_reqs_after_padding,
             idx_mapping=local_to_global_req_idx,
             idx_mapping_np=local_to_global_req_idx_np,
             expanded_idx_mapping=local_to_global_req_idx,
@@ -553,10 +633,11 @@ def partition_batch(
             num_draft_tokens=0,
             num_draft_tokens_per_req=None,
             query_start_loc=local_query_start_loc,
-            query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1],
+            query_start_loc_np=local_query_start_loc_np[: num_reqs_after_padding + 1],
             seq_lens=seq_lens,
             seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np),
             dcp_local_seq_lens=None,
+            dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound,
             num_computed_tokens_np=local_start_pos_np,
             prefill_len_np=local_prefill_len_np,
             num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,
@@ -570,6 +651,32 @@ def partition_batch(
             cu_num_logits_np=cu_num_logits_np,
             prompt_lens=None,
         )
+        return self._local_batch
+
+    def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch:
+        """Stage a capture or dummy batch in persistent PCP input buffers."""
+        input_buffers = self.input_buffers
+        num_reqs = input_batch.num_reqs_after_padding
+        num_tokens = input_batch.num_tokens_after_padding
+        input_batch = replace(
+            input_batch,
+            input_ids=input_buffers.input_ids[:num_tokens].copy_(input_batch.input_ids),
+            positions=input_buffers.positions[:num_tokens].copy_(input_batch.positions),
+            is_padding=input_buffers.is_padding[:num_tokens].copy_(
+                input_batch.is_padding
+            ),
+            query_start_loc=input_buffers.query_start_loc[: num_reqs + 1].copy_(
+                input_batch.query_start_loc
+            ),
+            seq_lens=input_buffers.seq_lens[:num_reqs].copy_(input_batch.seq_lens),
+        )
+        return input_batch
+
+    def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
+        assert self._local_block_tables is not None
+        return tuple(
+            block_table[:num_reqs].zero_() for block_table in self._local_block_tables
+        )
 
     def prepare_attn(
         self, input_batch: InputBatch
@@ -606,8 +713,7 @@ def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
         return self._gathered_kv_slot_mappings[:, : num_tokens * self.pcp_world_size]
 
     def _convert_to_gathered_slot_mappings(
-        self,
-        global_batch_slot_mappings: torch.Tensor,
+        self, global_batch_slot_mappings: torch.Tensor
     ) -> torch.Tensor:
         assert self._padded_gather_idx is not None
         assert self._gathered_kv_write_mask is not None
@@ -640,53 +746,61 @@ def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
         gathered = get_pcp_group().all_gather(hidden_states, dim=0)
         return gathered[self._hidden_restore_idx]
 
-    def restore_for_sampling(
-        self,
-        hidden_states: torch.Tensor,
-    ) -> tuple[torch.Tensor, InputBatch]:
-        assert self._global_batch is not None
-        return self.restore_hidden_states(hidden_states), self._global_batch
-
-
-def maybe_partition_pcp_batch(
-    manager: PCPManager | None,
-    input_batch: InputBatch,
-    padded_num_tokens: int | None = None,
-) -> InputBatch:
-    if manager is None:
-        return input_batch
-    return manager.partition_batch(
-        input_batch,
-        padded_num_tokens=padded_num_tokens,
-    )
-
+    def get_draft_input_buffers(
+        self, input_buffers: InputBuffers
+    ) -> InputBatch | InputBuffers:
+        return self.draft_prefill_batch or input_buffers
 
-def maybe_get_pcp_dummy_slot_mappings(
-    manager: PCPManager | None,
-    block_tables: BlockTables,
-    num_tokens: int,
-) -> torch.Tensor:
-    if manager is None:
-        return block_tables.get_dummy_slot_mappings(num_tokens)
-    return manager.get_dummy_slot_mappings(num_tokens)
+    def prepare_draft_prefill(
+        self, input_batch: InputBatch, input_ids: torch.Tensor
+    ) -> None:
+        self.draft_prefill_batch = None
+        if input_batch is not self._global_batch or self._local_batch is None:
+            return
+        local_batch = self._local_batch
+        assert self._local_gather_idx is not None
+        num_local_tokens = self._local_gather_idx.shape[0]
+        torch.index_select(
+            input_ids,
+            0,
+            self._local_gather_idx,
+            out=local_batch.input_ids[:num_local_tokens],
+        )
+        self.draft_prefill_batch = local_batch
 
+    def restore_draft_prefill(
+        self,
+        last_hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.draft_prefill_batch is None:
+            return last_hidden_states, hidden_states
+        local_last_hidden_states = last_hidden_states
+        last_hidden_states = self.restore_hidden_states(local_last_hidden_states)
+        hidden_states = (
+            last_hidden_states
+            if local_last_hidden_states is hidden_states
+            else self.restore_hidden_states(hidden_states)
+        )
+        self.draft_prefill_batch = None
+        return last_hidden_states, hidden_states
 
-def maybe_restore_pcp_for_sampling(
-    manager: PCPManager | None,
-    hidden_states: torch.Tensor | None,
-    input_batch: InputBatch,
-) -> tuple[torch.Tensor, InputBatch]:
-    assert hidden_states is not None
-    if manager is None:
-        return hidden_states, input_batch
-    return manager.restore_for_sampling(hidden_states)
+    def restore_for_sampling(
+        self, hidden_states: torch.Tensor, aux_hidden_states: list[torch.Tensor] | None
+    ) -> tuple[torch.Tensor, list[torch.Tensor] | None, InputBatch]:
+        assert self._global_batch is not None
+        hidden_states = self.restore_hidden_states(hidden_states)
+        if aux_hidden_states is not None:
+            aux_hidden_states = [
+                self.restore_hidden_states(states) for states in aux_hidden_states
+            ]
+        return hidden_states, aux_hidden_states, self._global_batch
 
 
 def maybe_build_pcp_manager(
     vllm_config: VllmConfig,
     device: torch.device,
     supports_mm_inputs: bool,
-    req_states: RequestState,
     block_tables: BlockTables,
     cls: type[PCPManager] = PCPManager,
 ) -> PCPManager | None:
@@ -705,7 +819,6 @@ def maybe_build_pcp_manager(
         pcp_world_size=pcp_size,
         pcp_rank=pcp_rank,
         device=device,
-        req_states=req_states,
         max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
         max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
         block_tables=block_tables,
diff --git a/vllm/v1/worker/gpu/pool/pooling_runner.py b/vllm/v1/worker/gpu/pool/pooling_runner.py
index 839edc8b8cab..ce947648514d 100644
--- a/vllm/v1/worker/gpu/pool/pooling_runner.py
+++ b/vllm/v1/worker/gpu/pool/pooling_runner.py
@@ -64,7 +64,7 @@ def add_request(
         req_id: str,
         req_index: int,
         pooling_params: PoolingParams,
-        prompt_token_ids: list[int],
+        prompt_token_ids: list[int] | None,
     ) -> None:
         task = pooling_params.task
         if task not in self.supported_tasks:
@@ -77,6 +77,7 @@ def add_request(
         self.pooling_params[req_index] = pooling_params
         self.pooling_states[req_index] = PoolingStates()
         if pooling_params.requires_token_ids:
+            assert prompt_token_ids is not None
             self.prompt_token_ids[req_index] = torch.tensor(
                 prompt_token_ids, dtype=torch.int64
             )
diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py
index 8c65a23ef7ec..c2aaa1fba25a 100644
--- a/vllm/v1/worker/gpu/pp_utils.py
+++ b/vllm/v1/worker/gpu/pp_utils.py
@@ -11,7 +11,7 @@
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.worker.gpu.input_batch import InputBatch
 
 
@@ -39,7 +39,6 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
     that produce a sampled token this step, and therefore must have that token
     (and the draft block proposed from it) propagated to the earlier PP stages.
     Returns None if no request in the batch produces a sample."""
-
     old_computed = input_batch.num_computed_tokens_np
     prefill_len = input_batch.prefill_len_np
     # Exclude non-final prefill chunks (they don't produce a sample).
@@ -138,7 +137,7 @@ def get_prev_sampled_outputs(
                 return None
             # Filter excluded request indices.
             idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np)
-            idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
+            idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device)
 
         self.main_stream.wait_event(slot.event)
         if slot.draft_tokens is not None and draft_tokens_to_update is not None:
@@ -148,7 +147,7 @@ def get_prev_sampled_outputs(
                 keep = ~exclude_mask
                 keep_t = torch.as_tensor(keep, device=self.device)
                 draft_tokens = draft_tokens[keep_t]
-                draft_idx_mapping = async_copy_to_gpu(
+                draft_idx_mapping = async_tensor_h2d(
                     slot.idx_mapping_np[keep], device=self.device
                 )
             draft_tokens_to_update[draft_idx_mapping] = draft_tokens
diff --git a/vllm/v1/worker/gpu/sample/bad_words.py b/vllm/v1/worker/gpu/sample/bad_words.py
index 3fb9d148ecac..974b24359889 100644
--- a/vllm/v1/worker/gpu/sample/bad_words.py
+++ b/vllm/v1/worker/gpu/sample/bad_words.py
@@ -1,43 +1,47 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import numpy as np
+from typing import TYPE_CHECKING
+
 import torch
 
 import vllm.envs as envs
 from vllm.sampling_params import SamplingParams
 from vllm.triton_utils import tl, triton
 from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
-from vllm.v1.worker.gpu.states import RequestState
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
+    LogitsContext,
+    LogitsProcessor,
+    LogitsProcRequestState,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
 
 
-class BadWordsState:
-    def __init__(self, req_states: RequestState):
+class BadWordsState(LogitsProcessor):
+    def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
         self.req_states = req_states
-        self.max_num_reqs = req_states.max_num_reqs
-        self.device = req_states.device
+        max_num_reqs = req_states.max_num_reqs
+        device = req_states.device
 
         max_total_tokens = envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS
         max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS
         # flattened bad word tokens: [max_num_reqs, VLLM_MAX_BAD_WORDS_TOTAL_TOKENS]
         self.bad_word_token_ids = StagedWriteTensor(
-            (self.max_num_reqs, max_total_tokens),
-            dtype=torch.int32,
-            device=self.device,
+            (max_num_reqs, max_total_tokens), dtype=torch.int32, device=device
         )
         # cumulative offsets of bad words: [max_num_reqs, VLLM_MAX_NUM_BAD_WORDS + 1]
         self.bad_word_offsets = StagedWriteTensor(
-            (self.max_num_reqs, max_num_bad_words + 1),
-            dtype=torch.int32,
-            device=self.device,
+            (max_num_reqs, max_num_bad_words + 1), dtype=torch.int32, device=device
         )
         # number of bad words per request
-        self.num_bad_words = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
+        self.num_bad_words = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
 
-    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
         bad_words_token_ids = sampling_params.bad_words_token_ids
         if not bad_words_token_ids:
             self.num_bad_words.np[req_idx] = 0
-            return
+            return False
 
         num_bad_words = len(bad_words_token_ids)
         max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS
@@ -65,38 +69,33 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
         self.bad_word_token_ids.stage_write(req_idx, 0, flattened_tokens)
         self.bad_word_offsets.stage_write(req_idx, 0, offsets)
         self.num_bad_words.np[req_idx] = num_bad_words
+        return True
 
     def apply_staged_writes(self) -> None:
         self.num_bad_words.copy_to_uva()
         self.bad_word_token_ids.apply_write()
         self.bad_word_offsets.apply_write()
 
-    def apply_bad_words(
-        self,
-        logits: torch.Tensor,
-        expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
-        input_ids: torch.Tensor,
-        expanded_local_pos: torch.Tensor,
-    ) -> None:
-        max_num_bad_words = int(self.num_bad_words.np[idx_mapping_np].max())
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        max_num_bad_words = int(self.num_bad_words.np[ctx.idx_mapping_np].max())
         if max_num_bad_words == 0:
             # No request uses bad words. Skip the kernel launch.
-            return
+            return logits
 
         apply_bad_words(
             logits,
-            expanded_idx_mapping,
+            ctx.expanded_idx_mapping,
             self.bad_word_token_ids.gpu,
             self.bad_word_offsets.gpu,
             self.num_bad_words.gpu,
             self.req_states.all_token_ids.gpu,
             self.req_states.prompt_len.gpu,
             self.req_states.total_len.gpu,
-            input_ids,
-            expanded_local_pos,
+            ctx.input_ids,
+            ctx.expanded_local_pos,
             max_num_bad_words,
         )
+        return logits
 
 
 @triton.jit
diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py
index 0f484f1a1710..3d41c49720f2 100644
--- a/vllm/v1/worker/gpu/sample/batch_shard.py
+++ b/vllm/v1/worker/gpu/sample/batch_shard.py
@@ -200,6 +200,9 @@ def shard_sampler_inputs(
             num_logits_per_req_np[local_req_indices_np], out=local_cu_num_logits_np[1:]
         )
         max_num_logits_per_req = int(num_logits_per_req_np.max()) if num_reqs else 1
+        local_seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[
+            torch.from_numpy(local_req_indices_np)
+        ]
 
         # Shard the input batch GPU tensors.
         sorted_logits_indices = torch.empty(
@@ -269,6 +272,7 @@ def shard_sampler_inputs(
             expanded_idx_mapping=local_expanded_idx_mapping,
             expanded_local_pos=local_expanded_local_pos,
             seq_lens=local_seq_lens,
+            seq_lens_cpu_upper_bound=local_seq_lens_cpu_upper_bound,
             logits_indices=local_logits_indices,
             cu_num_logits=local_cu_num_logits,
             cu_num_logits_np=local_cu_num_logits_np,
@@ -613,8 +617,6 @@ def gather_sampler_output(
         device=device,
     )
     if local_output is not None:
-        assert local_output.num_sampled is not None
-        assert local_output.num_rejected is not None
         assert not gather_num_nans or local_output.num_nans is not None
         num_src_cols = min(
             local_output.sampled_token_ids.shape[1], max_num_logits_per_req
diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py
index 4bf1f47280cf..e54a45870914 100644
--- a/vllm/v1/worker/gpu/sample/gumbel.py
+++ b/vllm/v1/worker/gpu/sample/gumbel.py
@@ -2,22 +2,20 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import torch
 
-from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton
+from vllm.triton_utils import tl, triton
 
-# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp
-# zero draws before the flipped Gumbel transform below.
+# Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the
+# retained Philox helper for rejection sampling.
 #
 # Triton requires globals accessed from `@triton.jit` functions to be wrapped
-# in `tl.constexpr(...)`. We can only do that when Triton is actually
-# available — on the CPU worker path `tl` is a placeholder whose `constexpr`
-# attribute is `None`, and `tl.constexpr(...)` would crash at import time.
-_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10
+# in `tl.constexpr(...)`.
+_TL_RAND_MIN = tl.constexpr(4.6566127342e-10)
 
 # Offset salt keeping the draft's Gumbel noise disjoint from the target's.
 # Verification is a probability-ratio test, not a Gumbel coupling, so a proposal
 # and the residual it is resampled from must not share a noise vector.
 # Positions are int64 and never approach 2**30, so the streams cannot collide.
-_DRAFT_NOISE_SALT = tl.constexpr(1 << 30) if HAS_TRITON else (1 << 30)
+_DRAFT_NOISE_SALT = tl.constexpr(1 << 30)
 
 
 @triton.jit
@@ -64,21 +62,6 @@ def apply_temperature(
     )
 
 
-@triton.jit
-def tl_rand64(seed, offset, includes_zero: tl.constexpr):
-    lo, hi, _, _ = tl.randint4x(seed, offset)
-    lo = lo.to(tl.uint32, bitcast=True).to(tl.uint64)
-    hi = hi.to(tl.uint32, bitcast=True).to(tl.uint64)
-    r = (hi << 32) | lo
-
-    # 1 / 2**64
-    scale = 5.421010862427522170037e-20
-    u = r.to(tl.float64) * scale
-    if not includes_zero:
-        u = tl.maximum(u, 2.2250738585072014e-308)  # float64 tiny
-    return u
-
-
 @triton.jit
 def tl_rand32(seed, offset, includes_zero: tl.constexpr):
     u = tl.rand(seed, offset)
@@ -87,6 +70,97 @@ def tl_rand32(seed, offset, includes_zero: tl.constexpr):
     return u
 
 
+@triton.jit
+def _murmur3_rotl32(value, shift: tl.constexpr):
+    return (value << shift) | (value >> (32 - shift))
+
+
+@triton.jit
+def _murmur3_mix(h, key):
+    key *= 0xCC9E2D51
+    key = _murmur3_rotl32(key, 15)
+    key *= 0x1B873593
+    h ^= key
+    h = _murmur3_rotl32(h, 13)
+    return h * 5 + 0xE6546B64
+
+
+@triton.jit
+def _murmur3_fmix32(h):
+    h ^= h >> 16
+    h *= 0x85EBCA6B
+    h ^= h >> 13
+    h *= 0xC2B2AE35
+    return h ^ (h >> 16)
+
+
+@triton.jit
+def murmur3_hash32(seed, pos, offset, domain: tl.constexpr = 0):
+    seed = seed.to(tl.int64)
+    pos = pos.to(tl.int64)
+    offset = offset.to(tl.uint32)
+    # Keep the request-wide prefix scalar until the token offset is mixed in.
+    h = (seed ^ seed).to(tl.uint32)
+    h ^= domain
+    h = _murmur3_mix(h, (seed & 0xFFFFFFFF).to(tl.uint32))
+    h = _murmur3_mix(h, ((seed >> 32) & 0xFFFFFFFF).to(tl.uint32))
+    h = _murmur3_mix(h, (pos & 0xFFFFFFFF).to(tl.uint32))
+    h = _murmur3_mix(h, offset)
+    return _murmur3_fmix32(h ^ 16)
+
+
+@triton.jit
+def murmur3_uniform32(seed, pos, offset):
+    random32 = murmur3_hash32(seed, pos, offset)
+    # Split the uint32 before converting to fp32 so backends without a native
+    # uint32-to-float conversion can still use all 32 source bits. Both 16-bit
+    # halves convert exactly; their sum is the correctly rounded fp32 value of
+    # (random32 + 0.5) * 2**-32. In particular, the u -> 0 winning tail keeps
+    # the full 32-bit source resolution instead of being truncated to 24 bits.
+    hi16 = (random32 >> 16).to(tl.int32)
+    lo16 = (random32 & 0xFFFF).to(tl.int32)
+    return (
+        hi16.to(tl.float32) * 1.52587890625e-05
+        + (lo16.to(tl.float32) + 0.5) * 2.3283064365386963e-10
+    )
+
+
+@triton.jit
+def _uniform64_from_random53(random53):
+    uniform = (random53.to(tl.float64) + 0.5) * 1.1102230246251565e-16
+    # The largest midpoint rounds to 1.0 in fp64; keep the uniform open without
+    # relying on a near-one literal that Triton may materialize in fp32.
+    return tl.where(uniform == 1.0, uniform - 1.1102230246251565e-16, uniform)
+
+
+@triton.jit
+def murmur3_uniform64(seed, pos, offset):
+    lo = murmur3_hash32(seed, pos, offset).to(tl.uint64)
+    hi = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.uint64)
+    random53 = ((hi << 32) | lo) >> 11
+    return _uniform64_from_random53(random53)
+
+
+@triton.jit
+def _log1p_neg_stable(value):
+    # Preserve precision for the positive Gumbel tail without relying on a
+    # backend-specific libdevice log1p. The degree-8 series has absolute error
+    # below 6e-7 on [0, 0.25]; subtraction is well-conditioned elsewhere for
+    # the part of the distribution that can win the argmax.
+    polynomial = 1.0 / 8.0
+    polynomial = 1.0 / 7.0 + value * polynomial
+    polynomial = 1.0 / 6.0 + value * polynomial
+    polynomial = 1.0 / 5.0 + value * polynomial
+    polynomial = 1.0 / 4.0 + value * polynomial
+    polynomial = 1.0 / 3.0 + value * polynomial
+    polynomial = 1.0 / 2.0 + value * polynomial
+    polynomial = 1.0 + value * polynomial
+    series = -value * polynomial
+
+    direct = tl.log(tl.maximum(1.0 - value, 5.960464477539063e-08))
+    return tl.where(value < 0.25, series, direct)
+
+
 @triton.jit
 def gumbel_noised_argmax(
     logits,
@@ -117,14 +191,15 @@ def gumbel_noised_argmax(
     if temp != 0.0:
         if IS_DRAFTING:
             pos = pos + _DRAFT_NOISE_SALT
-        gumbel_seed = tl.randint(seed, pos)
         if USE_FP64:
-            u = tl_rand64(gumbel_seed, keys, includes_zero=False)
+            u = murmur3_uniform64(seed, pos, keys)
             gumbel_noise = -tl.log(-tl.log(u))
         else:
-            u = tl_rand32(gumbel_seed, keys, includes_zero=False)
-            # log1p keeps the winning tail at u -> 0, where fp32 resolves it.
-            gumbel_noise = -tl.log(-tldevice.log1p(-u))
+            u = murmur3_uniform32(seed, pos, keys)
+            # Draw the large-noise tail (which decides the argmax winner) from
+            # u -> 0, where fp32 has fine resolution. Avoid backend-specific
+            # log1p while preserving precision in the winning tail.
+            gumbel_noise = -tl.log(-_log1p_neg_stable(u))
         logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
 
     return tl.max(logits, axis=0, return_indices=True)
diff --git a/vllm/v1/worker/gpu/sample/logit_bias.py b/vllm/v1/worker/gpu/sample/logit_bias.py
index 846edbe1e74d..53e20802b69f 100644
--- a/vllm/v1/worker/gpu/sample/logit_bias.py
+++ b/vllm/v1/worker/gpu/sample/logit_bias.py
@@ -1,60 +1,61 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING
+
 import numpy as np
 import torch
 
 from vllm.sampling_params import SamplingParams
 from vllm.triton_utils import tl, triton
 from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
+    LogitsContext,
+    LogitsProcessor,
+    LogitsProcRequestState,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
 
 MAX_NUM_ALLOWED_TOKEN_IDS = 1024
 MAX_NUM_LOGIT_BIAS_TOKENS = 1024
 MAX_NUM_STOP_TOKEN_IDS = 128
 
 
-class LogitBiasState:
-    def __init__(self, max_num_reqs: int, device: torch.device):
-        self.max_num_reqs = max_num_reqs
+class LogitBiasState(LogitsProcessor):
+    def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
+        self.prompt_len_np = req_states.prompt_len.np
+        max_num_reqs = req_states.max_num_reqs
+        device = req_states.device
 
         # Allowed token IDs.
-        self.num_allowed_token_ids = UvaBackedTensor(
-            self.max_num_reqs, dtype=torch.int32
-        )
+        self.num_allowed_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
         self.allowed_token_ids = StagedWriteTensor(
-            (self.max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS),
-            dtype=torch.int32,
-            device=device,
+            (max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS), dtype=torch.int32, device=device
         )
         # Logit bias.
-        self.num_logit_bias = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
+        self.num_logit_bias = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
         self.logit_bias_token_ids = StagedWriteTensor(
-            (self.max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
-            dtype=torch.int32,
-            device=device,
+            (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS), dtype=torch.int32, device=device
         )
         self.logit_bias = StagedWriteTensor(
-            (self.max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
+            (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
             dtype=torch.float32,
             device=device,
         )
         # Min tokens.
-        self.min_lens = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
-        self.num_stop_token_ids = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
-        self.restore_when_all_masked = UvaBackedTensor(
-            self.max_num_reqs, dtype=torch.int32
-        )
+        self.min_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
+        self.num_stop_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
+        self.restore_when_all_masked = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
         self.stop_token_ids = StagedWriteTensor(
-            (self.max_num_reqs, MAX_NUM_STOP_TOKEN_IDS),
-            dtype=torch.int32,
-            device=device,
+            (max_num_reqs, MAX_NUM_STOP_TOKEN_IDS), dtype=torch.int32, device=device
         )
 
         # Using any of the above.
         self.use_logit_bias = np.zeros(max_num_reqs, dtype=bool)
 
-    def add_request(
-        self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
-    ) -> None:
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        prompt_len = int(self.prompt_len_np[req_idx])
         # Using any logit bias.
         use_logit_bias = False
 
@@ -112,6 +113,7 @@ def add_request(
             self.restore_when_all_masked.np[req_idx] = 0
 
         self.use_logit_bias[req_idx] = use_logit_bias
+        return use_logit_bias
 
     def apply_staged_writes(self) -> None:
         self.num_allowed_token_ids.copy_to_uva()
@@ -126,25 +128,19 @@ def apply_staged_writes(self) -> None:
         self.restore_when_all_masked.copy_to_uva()
         self.stop_token_ids.apply_write()
 
-    def apply_logit_bias(
-        self,
-        logits: torch.Tensor,
-        expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
-        pos: torch.Tensor,
-    ) -> None:
-        if not np.any(self.use_logit_bias[idx_mapping_np]):
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        if not np.any(self.use_logit_bias[ctx.idx_mapping_np]):
             # No request uses logit bias. Skip the kernel launch.
-            return
+            return logits
 
         enable_stop_token_restore = bool(
-            np.any(self.restore_when_all_masked.np[idx_mapping_np])
+            np.any(self.restore_when_all_masked.np[ctx.idx_mapping_np])
         )
 
         apply_logit_bias(
             logits,
-            expanded_idx_mapping,
-            pos,
+            ctx.expanded_idx_mapping,
+            ctx.pos,
             self.num_allowed_token_ids.gpu,
             self.allowed_token_ids.gpu,
             self.num_logit_bias.gpu,
@@ -156,6 +152,7 @@ def apply_logit_bias(
             self.stop_token_ids.gpu,
             enable_stop_token_restore,
         )
+        return logits
 
 
 @triton.jit
diff --git a/vllm/v1/worker/gpu/sample/logits_processor/__init__.py b/vllm/v1/worker/gpu/sample/logits_processor/__init__.py
new file mode 100644
index 000000000000..f76eed35ceb3
--- /dev/null
+++ b/vllm/v1/worker/gpu/sample/logits_processor/__init__.py
@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Custom logits processors for the V2 model runner.
+
+Kept import-light: the frontend process imports this package to validate
+per-request params, so nothing here may pull in model-runner side modules
+(torch, triton, worker state) at import time.
+"""
+
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
+    LogitsContext,
+    LogitsProcessor,
+    LogitsProcRequestState,
+)
+from vllm.v1.worker.gpu.sample.logits_processor.loader import (
+    build_custom_logits_processors,
+    build_custom_logits_processors_params_validator,
+)
+
+__all__ = [
+    "LogitsProcRequestState",
+    "LogitsContext",
+    "LogitsProcessor",
+    "build_custom_logits_processors",
+    "build_custom_logits_processors_params_validator",
+]
diff --git a/vllm/v1/worker/gpu/sample/logits_processor/interface.py b/vllm/v1/worker/gpu/sample/logits_processor/interface.py
new file mode 100644
index 000000000000..9934492ec227
--- /dev/null
+++ b/vllm/v1/worker/gpu/sample/logits_processor/interface.py
@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""V2 logits processor interface.
+
+Import-light by design: custom processor classes are also loaded in the
+frontend process (to validate per-request params), so this module must not
+pull in model-runner side modules (torch, triton, worker state) at import
+time. Heavy imports live under TYPE_CHECKING; annotations are deferred.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+    import torch
+
+    from vllm.config import VllmConfig
+    from vllm.sampling_params import SamplingParams
+    from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
+    from vllm.v1.worker.gpu.states import RequestState
+
+
+@dataclass(frozen=True)
+class LogitsProcRequestState:
+    """State associated with active requests, shared with logits processors.
+
+    Wraps the model runner's per-slot buffers, which are mutated in place,
+    so reads always see current values. Processors must treat every field
+    as read-only.
+    """
+
+    device: torch.device
+    max_num_reqs: int
+    vocab_size: int
+
+    # [max_num_reqs, max_model_len] committed token ids per request slot.
+    all_token_ids: StagedWriteTensor
+    # [max_num_reqs] tokens in the user-provided prompt.
+    prompt_len: UvaBackedTensor
+    # [max_num_reqs] tokens fed at the latest (re)fill: the prompt plus any
+    # partial output on resumption after preemption.
+    prefill_len: UvaBackedTensor
+    # [max_num_reqs] prompt_len + output_len; grows as the request progresses.
+    total_len: StagedWriteTensor
+
+    @classmethod
+    def from_request_state(cls, req_states: RequestState) -> LogitsProcRequestState:
+        return cls(
+            device=req_states.device,
+            max_num_reqs=req_states.max_num_reqs,
+            vocab_size=req_states.vocab_size,
+            all_token_ids=req_states.all_token_ids,
+            prompt_len=req_states.prompt_len,
+            prefill_len=req_states.prefill_len,
+            total_len=req_states.total_len,
+        )
+
+
+@dataclass(frozen=True)
+class LogitsContext:
+    """The current step's batch layout, passed to every ``apply()`` call.
+
+    A row is a logits row, not a request: rows are reordered every step, and
+    under speculative decoding a request owns one row per draft token.
+    Committed tokens live in ``req_states.all_token_ids`` (valid up to
+    ``total_len``); this step's draft tokens are only in ``input_ids``.
+    """
+
+    # [num_logits_rows] row -> persistent request slot.
+    expanded_idx_mapping: torch.Tensor
+    # [num_reqs] batch position -> persistent request slot.
+    idx_mapping: torch.Tensor
+    # [num_reqs] batch position -> persistent request slot, on the host, for
+    # skipping work without a device sync.
+    idx_mapping_np: np.ndarray
+    # [num_logits_rows] row -> its offset among the rows of its own request.
+    expanded_local_pos: torch.Tensor
+    # [num_logits_rows] token fed to the model at each row's input position.
+    input_ids: torch.Tensor
+    # [num_logits_rows] position of each row within its sequence.
+    pos: torch.Tensor
+    # [num_reqs] batch position -> upper bound of tokens visible to the token
+    # being sampled this step, on the host. Exact when spec decoding isn't in use.
+    # Exact per-row lengths on device are `pos + 1`.
+    seq_lens_upper_bound_np: np.ndarray
+
+
+class LogitsProcessor(ABC):
+    """Custom logits processor for Model Runner V2.
+
+    Per-request state is keyed by the request slot index; slots are recycled
+    through a free list, so per-slot state must be fully (re)initialized in
+    ``add_request()``.
+
+    ``apply()`` runs after the built-in bias, penalty, bad-words and grammar
+    stages and before temperature, min_p and top-k/top-p, so it sees unscaled
+    logits and must not re-inflate grammar-masked tokens. Thinking-budget
+    forcing runs after ``apply()`` and wins over its edits for requests whose
+    budget is exhausted.
+
+    State that is constant for a request belongs in ``__init__()`` or
+    ``add_request()``; ``apply()`` receives only what changes per step.
+    """
+
+    def __init__(  # noqa: B027
+        self, vllm_config: VllmConfig, req_states: LogitsProcRequestState
+    ):
+        """Capture what stays constant for the processor's lifetime.
+
+        ``req_states`` exposes the on-device token history and batch constants a
+        processor may read. Treat it as read-only.
+        """
+
+    @classmethod  # noqa: B027
+    def validate_params(cls, sampling_params: SamplingParams) -> None:
+        """Raise ``ValueError`` for invalid per-request arguments.
+
+        Runs at request admission, so invalid arguments fail the request
+        with an error instead of reaching the sampler.
+        """
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        """Initialize per-slot state for a request entering the batch.
+
+        The slot may hold a previous occupant's state; overwrite or
+        neutralize all of it here.
+
+        Returns whether this processor modifies logits for the request.
+        """
+        return True
+
+    def apply_staged_writes(self) -> None:  # noqa: B027
+        """Flush any host-side writes staged by ``add_request()`` to the device.
+
+        Called once per step before the forward pass, after the model runner
+        has flushed ``req_states``, so a processor that stages writes here can
+        read the request's tokens back on device.
+        """
+
+    @abstractmethod
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        """Modify logits in place or return a new tensor. In-place modification
+        is preferred for efficiency.
+
+        ``apply()`` is called once for the whole batch, including rows of
+        requests this processor declined in ``add_request()``, so filter rows
+        via ``ctx.expanded_idx_mapping``.
+
+        Args:
+            logits: [num_logits_rows, vocab_size] float32 tensor.
+            ctx: this step's batch layout.
+
+        """
+        raise NotImplementedError
diff --git a/vllm/v1/worker/gpu/sample/logits_processor/loader.py b/vllm/v1/worker/gpu/sample/logits_processor/loader.py
new file mode 100644
index 000000000000..1382abd5373f
--- /dev/null
+++ b/vllm/v1/worker/gpu/sample/logits_processor/loader.py
@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Loading of custom logits processor classes for the V2 model runner.
+
+This module is deliberately import-light: the frontend process uses it to
+validate per-request params, so importing it must not pull in model-runner
+side modules (torch, triton, worker state). Imports needed to instantiate
+processors are deferred to ``build_custom_logits_processors``.
+"""
+
+import importlib
+from collections.abc import Callable, Sequence
+from functools import lru_cache
+from importlib.metadata import entry_points
+from typing import TYPE_CHECKING
+
+from vllm.exceptions import VLLMValidationError
+from vllm.logger import init_logger
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
+    LogitsProcessor,
+    LogitsProcRequestState,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.sampling_params import SamplingParams
+    from vllm.v1.worker.gpu.states import RequestState
+
+logger = init_logger(__name__)
+
+# Same entry-points group as
+# vllm.v1.sample.logits_processor.LOGITSPROCS_GROUP; duplicated here so this
+# module stays importable without the V1 sampler machinery.
+LOGITSPROCS_GROUP = "vllm.logits_processors"
+
+
+def _load_v2_logitsprocs_plugins() -> list[type[LogitsProcessor]]:
+    """Load installed logits processor plugins.
+
+    Plugins registered under ``LOGITSPROCS_GROUP`` ship V1-interface
+    processors today; those crash under the V2 sampler, so reject them at
+    load time.
+    """
+    from vllm.utils.torch_utils import guard_cuda_initialization
+
+    installed_logitsprocs_plugins = entry_points(group=LOGITSPROCS_GROUP)
+    if len(installed_logitsprocs_plugins) == 0:
+        logger.debug("No logitsprocs plugins installed (group %s).", LOGITSPROCS_GROUP)
+        return []
+
+    logger.debug("Loading installed logitsprocs plugins (group %s):", LOGITSPROCS_GROUP)
+    classes: list[type[LogitsProcessor]] = []
+    for entrypoint in installed_logitsprocs_plugins:
+        try:
+            logger.debug(
+                "- Loading logitproc plugin entrypoint=%s target=%s",
+                entrypoint.name,
+                entrypoint.value,
+            )
+            with guard_cuda_initialization():
+                cls = entrypoint.load()
+        except Exception as e:
+            logger.error("Failed to load LogitsProcessor plugin %s: %s", entrypoint, e)
+            raise RuntimeError(
+                f"Failed to load LogitsProcessor plugin {entrypoint}"
+            ) from e
+        if not (isinstance(cls, type) and issubclass(cls, LogitsProcessor)):
+            raise ValueError(
+                f"LogitsProcessor plugin {entrypoint.name} is not a subclass of "
+                "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor; "
+                "V1-interface plugins are not supported by the V2 model runner."
+            )
+        classes.append(cls)
+    return classes
+
+
+def _load_v2_logitsprocs_by_fqcns(
+    logits_processors: Sequence[str | type],
+) -> list[type[LogitsProcessor]]:
+    """Resolve a mixed list of processor types and FQCN strings into types.
+
+    FQCN syntax is <module>:<type> i.e. x.y.z:CustomLogitProc.
+    """
+    from vllm.utils.torch_utils import guard_cuda_initialization
+
+    classes: list[type[LogitsProcessor]] = []
+    for ldx, logitproc in enumerate(logits_processors):
+        if isinstance(logitproc, type):
+            logger.debug(" - Already-loaded logit processor: %s", logitproc.__name__)
+            if not issubclass(logitproc, LogitsProcessor):
+                raise ValueError(
+                    f"{logitproc.__name__} is not a subclass of "
+                    "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor"
+                )
+            classes.append(logitproc)
+            continue
+
+        logger.debug("- Loading logits processor %s", logitproc)
+        parts = logitproc.split(":")
+        if len(parts) != 2:
+            raise ValueError(
+                f"Invalid logits processor FQCN {logitproc!r}. "
+                "Expected format: '<module>:<type>'"
+            )
+        module_path, qualname = parts
+
+        try:
+            with guard_cuda_initialization():
+                module = importlib.import_module(module_path)
+        except Exception as e:
+            logger.error(
+                "Failed to load %sth LogitsProcessor plugin %s: %s",
+                ldx,
+                logitproc,
+                e,
+            )
+            raise RuntimeError(
+                f"Failed to load {ldx}th LogitsProcessor plugin {logitproc}"
+            ) from e
+
+        obj = module
+        for attr in qualname.split("."):
+            obj = getattr(obj, attr)
+        if not isinstance(obj, type):
+            raise ValueError("Loaded logit processor must be a type.")
+        if not issubclass(obj, LogitsProcessor):
+            raise ValueError(
+                f"{obj.__name__} is not a subclass of "
+                "vllm.v1.worker.gpu.sample.logits_processor.LogitsProcessor"
+            )
+        classes.append(obj)
+
+    return classes
+
+
+def _load_v2_logitsprocs(
+    logits_processors: Sequence[str | type] | None,
+) -> list[type[LogitsProcessor]]:
+    """Load all custom logits processors for the V2 model runner.
+
+    Combines installed plugins (validated as V2 processors) with the
+    user-specified list of processor types and FQCN strings.
+    """
+    return _load_v2_logitsprocs_plugins() + _load_v2_logitsprocs_by_fqcns(
+        logits_processors or ()
+    )
+
+
+@lru_cache
+def _cached_load_v2_logitsprocs(
+    custom_logitsprocs: tuple[str | type, ...],
+) -> list[type[LogitsProcessor]]:
+    return _load_v2_logitsprocs(list(custom_logitsprocs))
+
+
+def build_custom_logits_processors(
+    vllm_config: "VllmConfig",
+    req_states: "RequestState",
+    is_pooling_model: bool,
+    custom_logitsprocs: Sequence[str | type] = (),
+) -> list[LogitsProcessor]:
+    """Load and instantiate custom logits processors, entrypoint plugins first.
+
+    Raises:
+        ValueError: if a pooling model specifies custom processors, or a
+            loaded class does not implement the V2 interface.
+        RuntimeError: if an FQCN fails to import.
+
+    """
+    from vllm.v1.sample.logits_processor import STR_POOLING_REJECTS_LOGITSPROCS
+
+    if is_pooling_model:
+        if custom_logitsprocs:
+            raise ValueError(STR_POOLING_REJECTS_LOGITSPROCS)
+        logger.debug(
+            "Skipping logits processor loading because pooling models"
+            " do not support logits processors."
+        )
+        return []
+
+    custom_logitsprocs_classes = _load_v2_logitsprocs(custom_logitsprocs)
+    lp_req_state = LogitsProcRequestState.from_request_state(req_states)
+    return [ctor(vllm_config, lp_req_state) for ctor in custom_logitsprocs_classes]
+
+
+def build_custom_logits_processors_params_validator(
+    custom_logitsprocs: Sequence[str | type] | None,
+) -> "Callable[[SamplingParams], None]":
+    """Load custom processor classes once and return a params validator.
+
+    Called from the frontend at startup. The returned callable runs each
+    processor's ``validate_params`` at request admission.
+
+    Raises:
+        ValueError: if a loaded class does not implement the V2 interface.
+        RuntimeError: if an FQCN fails to import.
+
+    """
+    classes = _cached_load_v2_logitsprocs(tuple(custom_logitsprocs or ()))
+
+    if not classes:
+        return lambda _: None
+
+    def validate_params(sampling_params: "SamplingParams") -> None:
+        for cls in classes:
+            try:
+                cls.validate_params(sampling_params)
+            except ValueError as e:
+                raise VLLMValidationError(str(e)) from e
+
+    return validate_params
diff --git a/vllm/v1/worker/gpu/sample/output.py b/vllm/v1/worker/gpu/sample/output.py
index ff3675db130f..aedff5ce1873 100644
--- a/vllm/v1/worker/gpu/sample/output.py
+++ b/vllm/v1/worker/gpu/sample/output.py
@@ -17,8 +17,8 @@ class SamplerOutput:
     sampled_token_ids: torch.Tensor
     logprobs_tensors: LogprobsTensors | None
     num_nans: torch.Tensor | None
-    num_sampled: torch.Tensor | None
-    num_rejected: torch.Tensor | None = None
+    num_sampled: torch.Tensor
+    num_rejected: torch.Tensor
     sampling_mask_tensors: SamplingMaskTensors | None = None
 
 
diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py
index 25cb2f211d94..eb751e8be355 100644
--- a/vllm/v1/worker/gpu/sample/penalties.py
+++ b/vllm/v1/worker/gpu/sample/penalties.py
@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING
+
 import numpy as np
 import torch
 
@@ -8,16 +10,22 @@
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
-from vllm.v1.worker.gpu.states import RequestState
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
+    LogitsContext,
+    LogitsProcessor,
+    LogitsProcRequestState,
+)
 
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
 
-class PenaltiesState:
-    def __init__(self, req_states: RequestState):
-        self.req_states = req_states
 
-        max_num_reqs = req_states.max_num_reqs
-        self.vocab_size = req_states.vocab_size
+class PenaltiesState(LogitsProcessor):
+    def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
+        self.req_states = req_states
         self.device = req_states.device
+        max_num_reqs = req_states.max_num_reqs
+        vocab_size = req_states.vocab_size
 
         self.repetition_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
         self.frequency_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
@@ -30,20 +38,17 @@ def __init__(self, req_states: RequestState):
 
         # Statistics for penalties.
         self.prompt_bin_mask = torch.zeros(
-            max_num_reqs,
-            cdiv(self.vocab_size, 32),
-            dtype=torch.int32,
-            device=self.device,
+            max_num_reqs, cdiv(vocab_size, 32), dtype=torch.int32, device=self.device
         )
         # TODO(woosuk): This tensor is rarely used but can be very large, taking up
         # GBs of GPU memory. Optimize the memory usage.
         self.output_bin_counts = torch.zeros(
-            max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
+            max_num_reqs, vocab_size, dtype=torch.int32, device=self.device
         )
 
         self._new_penalties_reqs: list[int] = []
 
-    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
         self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty
         self.frequency_penalty.np[req_idx] = sampling_params.frequency_penalty
         self.presence_penalty.np[req_idx] = sampling_params.presence_penalty
@@ -52,13 +57,12 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
         self.use_penalty[req_idx] = do_penalty
         if do_penalty:
             self._new_penalties_reqs.append(req_idx)
+        return do_penalty
 
     def apply_staged_writes(self) -> None:
         if self._new_penalties_reqs:
             idx_mapping = async_tensor_h2d(
-                self._new_penalties_reqs,
-                dtype=torch.int32,
-                device=self.device,
+                self._new_penalties_reqs, dtype=torch.int32, device=self.device
             )
 
             prefill_lens = self.req_states.prefill_len.np[self._new_penalties_reqs]
@@ -78,29 +82,23 @@ def apply_staged_writes(self) -> None:
         self.frequency_penalty.copy_to_uva()
         self.presence_penalty.copy_to_uva()
 
-    def apply_penalties(
-        self,
-        logits: torch.Tensor,
-        expanded_idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
-        input_ids: torch.Tensor,
-        expanded_local_pos: torch.Tensor,
-    ) -> None:
-        if not np.any(self.use_penalty[idx_mapping_np]):
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
+        if not np.any(self.use_penalty[ctx.idx_mapping_np]):
             # No request uses penalties. Skip the kernel launch.
-            return
+            return logits
 
         apply_penalties(
             logits,
-            expanded_idx_mapping,
-            input_ids,
-            expanded_local_pos,
+            ctx.expanded_idx_mapping,
+            ctx.input_ids,
+            ctx.expanded_local_pos,
             self.repetition_penalty.gpu,
             self.frequency_penalty.gpu,
             self.presence_penalty.gpu,
             self.prompt_bin_mask,
             self.output_bin_counts,
         )
+        return logits
 
 
 @triton.jit
diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py
index b069604c2964..802adb6f642f 100644
--- a/vllm/v1/worker/gpu/sample/sampler.py
+++ b/vllm/v1/worker/gpu/sample/sampler.py
@@ -1,23 +1,32 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Sequence
+
 import numpy as np
 import torch
 
 import vllm.envs as envs
+from vllm.config import VllmConfig
 from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
-from vllm.config.reasoning import ReasoningConfig
 from vllm.sampling_params import SamplingParams
 from vllm.v1.sample.ops.topk_topp_sampler import (
     apply_top_k_top_p,
     flashinfer_sample,
     flashinfer_sampler_supported,
+    xpu_sample,
+    xpu_sampler_supported,
 )
 from vllm.v1.worker.gpu.input_batch import InputBatch, get_num_sampled_and_rejected
 from vllm.v1.worker.gpu.metrics.logits import get_num_nans
 from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
+from vllm.v1.worker.gpu.sample.logits_processor.interface import (
+    LogitsContext,
+    LogitsProcessor,
+    LogitsProcRequestState,
+)
 from vllm.v1.worker.gpu.sample.logprob import (
     LogprobTokenIdsState,
     compute_topk_scores,
@@ -33,6 +42,7 @@
 class Sampler:
     def __init__(
         self,
+        vllm_config: VllmConfig,
         max_num_reqs: int,
         vocab_size: int,
         device: torch.device,
@@ -41,8 +51,8 @@ def __init__(
         num_speculative_tokens: int = 1,
         use_fp64_gumbel: bool = False,
         enable_trace_replay: bool = False,
-        reasoning_config: ReasoningConfig | None = None,
         return_sampling_mask: bool = False,
+        custom_logits_processors: Sequence[LogitsProcessor] = (),
     ):
         self.logprobs_mode = logprobs_mode
         self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS  # False by default.
@@ -50,11 +60,25 @@ def __init__(
 
         self.req_states = req_states
         self.sampling_states = SamplingStates(max_num_reqs, vocab_size)
-        self.penalties_state = PenaltiesState(req_states)
-        self.logit_bias_state = LogitBiasState(max_num_reqs, device)
-        self.bad_words_state = BadWordsState(req_states)
+
+        lp_req_state = LogitsProcRequestState.from_request_state(req_states)
+        self.penalties_state = PenaltiesState(vllm_config, lp_req_state)
+        logit_bias_state = LogitBiasState(vllm_config, lp_req_state)
+        bad_words_state = BadWordsState(vllm_config, lp_req_state)
+
+        # List order is pipeline order: bias adds, penalties scale, so the
+        # two do not commute.
+        self.logits_processors: list[LogitsProcessor] = [
+            logit_bias_state,
+            self.penalties_state,
+            bad_words_state,
+            *custom_logits_processors,
+        ]
+
         self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
-        self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config)
+        self.thinking_budget_state = ThinkingBudgetState(
+            req_states, vllm_config.reasoning_config
+        )
         self.trace_replay_state = (
             TraceReplayState(req_states) if enable_trace_replay else None
         )
@@ -64,42 +88,30 @@ def __init__(
         self.use_flashinfer = (
             not return_sampling_mask and flashinfer_sampler_supported()
         )
+        # The XPU kernel draws fp32 exponential noise, so it can't honor fp64.
+        self.use_xpu_sampler = (
+            not return_sampling_mask and not use_fp64_gumbel and xpu_sampler_supported()
+        )
+
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
+        needs_processing = self.sampling_states.add_request(req_idx, sampling_params)
+        needs_processing |= self.thinking_budget_state.add_request(
+            req_idx, sampling_params
+        )
+        for processor in self.logits_processors:
+            needs_processing |= processor.add_request(req_idx, sampling_params)
+        self.needs_logits_processing[req_idx] = needs_processing
 
-    def add_request(
-        self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
-    ) -> None:
-        self.sampling_states.add_request(req_idx, sampling_params)
-        self.penalties_state.add_request(req_idx, sampling_params)
-        self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params)
-        self.bad_words_state.add_request(req_idx, sampling_params)
         self.logprob_token_ids_state.add_request(req_idx, sampling_params)
-        self.thinking_budget_state.add_request(req_idx, sampling_params)
         if self.trace_replay_state is not None:
             self.trace_replay_state.add_request(req_idx, sampling_params)
 
-        states = self.sampling_states
-        temperature = states.temperature.np[req_idx]
-        self.needs_logits_processing[req_idx] = (
-            self.logit_bias_state.use_logit_bias[req_idx]
-            or self.penalties_state.use_penalty[req_idx]
-            or self.bad_words_state.num_bad_words.np[req_idx] > 0
-            or (
-                self.thinking_budget_state.enabled
-                and self.thinking_budget_state.use_thinking_budget[req_idx]
-            )
-            or (temperature != 0.0 and temperature != 1.0)
-            or states.min_p.np[req_idx] != 0.0
-            or states.top_k.np[req_idx] != states.vocab_size
-            or states.top_p.np[req_idx] != 1.0
-        )
-
     def apply_staged_writes(self) -> None:
         self.sampling_states.apply_staged_writes()
-        self.penalties_state.apply_staged_writes()
-        self.logit_bias_state.apply_staged_writes()
-        self.bad_words_state.apply_staged_writes()
-        self.logprob_token_ids_state.apply_staged_writes()
+        for processor in self.logits_processors:
+            processor.apply_staged_writes()
         self.thinking_budget_state.apply_staged_writes()
+        self.logprob_token_ids_state.apply_staged_writes()
         if self.trace_replay_state is not None:
             self.trace_replay_state.apply_staged_writes()
 
@@ -131,6 +143,7 @@ def __call__(
         expanded_local_pos = input_batch.expanded_local_pos
         pos = input_batch.positions[input_batch.logits_indices]
         input_ids = input_batch.input_ids[input_batch.logits_indices]
+        seq_lens_upper_bound_np = input_batch.seq_lens_cpu_upper_bound.numpy()
 
         # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
         # that num_nans is computed before applying penalties and temperature.
@@ -146,6 +159,7 @@ def __call__(
             pos,
             input_ids,
             expanded_local_pos,
+            seq_lens_upper_bound_np,
             return_logprobs=logprobs_dims is not None,
         )
 
@@ -215,6 +229,7 @@ def apply_sampling_params(
         pos: torch.Tensor,
         input_ids: torch.Tensor,
         expanded_local_pos: torch.Tensor,
+        seq_lens_upper_bound_np: np.ndarray,
         skip_top_k_top_p: bool = False,
     ) -> torch.Tensor:
         if not np.any(self.needs_logits_processing[idx_mapping_np]):
@@ -223,39 +238,23 @@ def apply_sampling_params(
         # Copy logits to a new FP32 tensor.
         logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
 
-        # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
-        self.logit_bias_state.apply_logit_bias(
-            logits, expanded_idx_mapping, idx_mapping_np, pos
-        )
-
-        # Apply penalties in place.
-        self.penalties_state.apply_penalties(
-            logits,
-            expanded_idx_mapping,
-            idx_mapping_np,
-            input_ids,
-            expanded_local_pos,
+        ctx = LogitsContext(
+            expanded_idx_mapping=expanded_idx_mapping,
+            idx_mapping=idx_mapping,
+            idx_mapping_np=idx_mapping_np,
+            expanded_local_pos=expanded_local_pos,
+            input_ids=input_ids,
+            pos=pos,
+            seq_lens_upper_bound_np=seq_lens_upper_bound_np,
         )
 
-        # Apply bad words masking in place.
-        self.bad_words_state.apply_bad_words(
-            logits,
-            expanded_idx_mapping,
-            idx_mapping_np,
-            input_ids,
-            expanded_local_pos,
-        )
+        # Apply logits processors (native + any custom).
+        for processor in self.logits_processors:
+            logits = processor.apply(logits, ctx)
 
-        # Force the reasoning end marker once a request's thinking budget is
-        # reached; applied before temperature so the forced token is always kept.
-        self.thinking_budget_state.apply(
-            logits,
-            expanded_idx_mapping,
-            idx_mapping,
-            idx_mapping_np,
-            input_ids,
-            expanded_local_pos,
-        )
+        # Forcing runs last so no stage can overwrite the forced end marker
+        # or weaken it by scaling.
+        self.thinking_budget_state.apply(logits, ctx)
 
         # Apply temperature in place.
         self.sampling_states.apply_temperature(
@@ -282,6 +281,7 @@ def sample(
         pos: torch.Tensor,
         input_ids: torch.Tensor,
         expanded_local_pos: torch.Tensor,
+        seq_lens_upper_bound_np: np.ndarray,
         return_logprobs: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         processed_logits = self.apply_sampling_params(
@@ -292,20 +292,24 @@ def sample(
             pos,
             input_ids,
             expanded_local_pos,
+            seq_lens_upper_bound_np,
             skip_top_k_top_p=True,
         )
         top_k, top_p = self.sampling_states.get_top_k_top_p(
             expanded_idx_mapping, idx_mapping_np
         )
-        use_flashinfer = self.use_flashinfer and not (
-            # Don't use FI sampler if no requests use top_k/top_p, if there are
-            # any greedy requests or per-request seeds, or if post-processed
-            # logprobs need to be returned for any requests.
+        # Don't use a fused sampler if no requests use top_k/top_p, if there are
+        # any greedy requests or per-request seeds, or if post-processed
+        # logprobs need to be returned for any requests.
+        fused_sampler_eligible = not (
             (top_k is None and top_p is None)
             or (return_logprobs and self.logprobs_mode in PROCESSED_LOGPROBS_MODES)
             or self.sampling_states.any_greedy(idx_mapping_np)
             or self.sampling_states.any_explicit_seed(idx_mapping_np)
         )
+        use_fused_sampler = (
+            self.use_flashinfer or self.use_xpu_sampler
+        ) and fused_sampler_eligible
 
         return self._sample_random(
             processed_logits,
@@ -314,7 +318,7 @@ def sample(
             pos,
             top_k,
             top_p,
-            use_flashinfer,
+            use_fused_sampler,
         )
 
     def _sample_random(
@@ -325,10 +329,15 @@ def _sample_random(
         pos: torch.Tensor,
         top_k: torch.Tensor | None,
         top_p: torch.Tensor | None,
-        use_flashinfer: bool,
+        use_fused_sampler: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        if use_flashinfer:
-            sampled = flashinfer_sample(processed_logits, top_k, top_p).to(torch.int64)
+        if use_fused_sampler:
+            if self.use_flashinfer:
+                sampled = flashinfer_sample(processed_logits, top_k, top_p).to(
+                    torch.int64
+                )
+            else:  # Use XPU sampler
+                sampled, _ = xpu_sample(processed_logits, top_k, top_p)
         else:
             processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
             sampled = gumbel_sample(
diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py
index fe4dee6a6b10..47cf5de856e8 100644
--- a/vllm/v1/worker/gpu/sample/states.py
+++ b/vllm/v1/worker/gpu/sample/states.py
@@ -38,14 +38,18 @@ def __init__(self, max_num_reqs: int, vocab_size: int):
         # -1 means no logprobs are requested.
         self.num_logprobs.fill(NO_LOGPROBS)
 
-    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
-        self.temperature.np[req_idx] = sampling_params.temperature
-        self.top_p.np[req_idx] = sampling_params.top_p
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
+        temperature = sampling_params.temperature
         top_k = sampling_params.top_k
+        top_p = sampling_params.top_p
+        min_p = sampling_params.min_p
+
+        self.temperature.np[req_idx] = temperature
+        self.top_p.np[req_idx] = top_p
         if top_k <= 0 or top_k > self.vocab_size:
             top_k = self.vocab_size
         self.top_k.np[req_idx] = top_k
-        self.min_p.np[req_idx] = sampling_params.min_p
+        self.min_p.np[req_idx] = min_p
 
         seed = sampling_params.seed
         self.seeds_set[req_idx] = seed is not None
@@ -60,6 +64,13 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
             num_logprobs = self.vocab_size
         self.num_logprobs[req_idx] = num_logprobs
 
+        return temperature != 0.0 and (
+            temperature != 1.0
+            or min_p != 0.0
+            or top_k != self.vocab_size
+            or top_p != 1.0
+        )
+
     def apply_staged_writes(self) -> None:
         self.temperature.copy_to_uva()
         self.top_p.copy_to_uva()
diff --git a/vllm/v1/worker/gpu/sample/thinking_budget.py b/vllm/v1/worker/gpu/sample/thinking_budget.py
index eca792986027..48c960cbeb68 100644
--- a/vllm/v1/worker/gpu/sample/thinking_budget.py
+++ b/vllm/v1/worker/gpu/sample/thinking_budget.py
@@ -9,6 +9,7 @@
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
+from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext
 from vllm.v1.worker.gpu.states import RequestState
 
 if TYPE_CHECKING:
@@ -78,11 +79,12 @@ def __init__(
             end_ids, dtype=torch.int32, device=self.device
         )
 
-    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
+    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
         if not self.enabled:
-            return
+            return False
         budget = sampling_params.thinking_token_budget
-        self.use_thinking_budget[req_idx] = budget is not None
+        use_thinking_budget = budget is not None
+        self.use_thinking_budget[req_idx] = use_thinking_budget
         if budget is None:
             budget = -1
         else:
@@ -91,6 +93,7 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
         if self.thinking_token_budget.np[req_idx] != budget:
             self.thinking_token_budget.np[req_idx] = budget
             self._budget_dirty = True
+        return use_thinking_budget
 
     def apply_staged_writes(self) -> None:
         if not self.enabled:
@@ -107,27 +110,20 @@ def apply_staged_writes(self) -> None:
             self.thinking_token_budget.copy_to_uva()
             self._budget_dirty = False
 
-    def apply(
-        self,
-        logits: torch.Tensor,
-        expanded_idx_mapping: torch.Tensor,
-        idx_mapping: torch.Tensor,
-        idx_mapping_np: np.ndarray,
-        input_ids: torch.Tensor,
-        expanded_local_pos: torch.Tensor,
-    ) -> None:
+    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> None:
+        idx_mapping_np = ctx.idx_mapping_np
         if not self.enabled or not np.any(self.use_thinking_budget[idx_mapping_np]):
             return
 
         apply_thinking_budget(
             logits,
-            idx_mapping,
-            expanded_idx_mapping,
+            ctx.idx_mapping,
+            ctx.expanded_idx_mapping,
             self.thinking_token_budget.gpu,
             self.req_states.all_token_ids.gpu,
             self.req_states.total_len.gpu,
-            input_ids,
-            expanded_local_pos,
+            ctx.input_ids,
+            ctx.expanded_local_pos,
             self.cached_last_start,
             self.cached_last_end,
             self.cached_scan_pos,
diff --git a/vllm/v1/worker/gpu/sample/watermark.py b/vllm/v1/worker/gpu/sample/watermark.py
index a0105cc1b7f0..b70d8ae271db 100644
--- a/vllm/v1/worker/gpu/sample/watermark.py
+++ b/vllm/v1/worker/gpu/sample/watermark.py
@@ -25,6 +25,7 @@
 from vllm.v1.watermarking.prfs.philox import (
     _UINT32_MASK as _UINT32_MASK_VALUE,
 )
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax
 
 if HAS_TRITON:
     from triton.language import math as tl_math
@@ -42,6 +43,163 @@
 _TOKEN_DOMAIN = tl.constexpr(_TOKEN_DOMAIN_VALUE) if HAS_TRITON else _TOKEN_DOMAIN_VALUE
 
 
+@triton.jit
+def _repeated_context_mask_kernel(
+    output_ptr,
+    all_token_ids_ptr,
+    all_token_ids_stride,
+    req_indices_ptr,
+    prompt_lens_ptr,
+    total_lens_ptr,
+    contexts_ptr,
+    context_stride,
+    CONTEXT_WIDTH: tl.constexpr,
+    MAX_HISTORY: tl.constexpr,
+    INCLUDE_PROMPT: tl.constexpr,
+    SKIP_PARTIAL_CONTEXT: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    req_idx = tl.load(req_indices_ptr + row).to(tl.int64)
+    if req_idx < 0:
+        tl.store(output_ptr + row, 0)
+        return
+    prompt_len = tl.load(prompt_lens_ptr + req_idx)
+    total_len = tl.load(total_lens_ptr + req_idx)
+    if SKIP_PARTIAL_CONTEXT and total_len - prompt_len < CONTEXT_WIDTH:
+        tl.store(output_ptr + row, 1)
+        return
+    sequence_start = 0 if INCLUDE_PROMPT else prompt_len
+    history_len = total_len - sequence_start
+
+    offsets = tl.arange(0, BLOCK)
+    repeated = tl.full((), 0, tl.int32)
+    scan_start = 0
+    if MAX_HISTORY > 0:
+        scan_start = tl.maximum(history_len - MAX_HISTORY, 0)
+    # Aligned loop bounds let the per-lane offsets fold into immediates.
+    aligned_start = (scan_start // BLOCK) * BLOCK
+    for block_start in tl.range(aligned_start, history_len, BLOCK):
+        previous_pos = block_start + offsets
+        in_window = (previous_pos >= scan_start) & (previous_pos < history_len)
+        matches = in_window
+        for offset in range(CONTEXT_WIDTH):
+            context_token = tl.load(contexts_ptr + row * context_stride + offset)
+            historical_pos = previous_pos + offset - CONTEXT_WIDTH
+            historical_token = tl.load(
+                all_token_ids_ptr
+                + req_idx * all_token_ids_stride
+                + sequence_start
+                + tl.maximum(historical_pos, 0),
+                mask=in_window & (historical_pos >= 0),
+                other=-1,
+            )
+            matches &= historical_token == context_token
+        repeated |= tl.max(matches.to(tl.int32), axis=0)
+    tl.store(output_ptr + row, repeated)
+
+
+def _repeated_context_mask_cpu(
+    all_token_ids: torch.Tensor,
+    req_indices: torch.Tensor,
+    prompt_lens: torch.Tensor,
+    total_lens: torch.Tensor,
+    contexts: torch.Tensor,
+    max_history: int | None = None,
+    include_prompt: bool = False,
+    skip_partial_context: bool = False,
+) -> torch.Tensor:
+    """Reference implementation; the parity tests check the Triton kernel against it."""
+    repeated = torch.zeros(len(req_indices), dtype=torch.bool)
+    for row, req_idx_tensor in enumerate(req_indices):
+        req_idx = int(req_idx_tensor)
+        if req_idx < 0:
+            continue
+        prompt_len = int(prompt_lens[req_idx])
+        total_len = int(total_lens[req_idx])
+        if skip_partial_context and total_len - prompt_len < contexts.shape[-1]:
+            repeated[row] = True
+            continue
+        sequence_start = 0 if include_prompt else prompt_len
+        history_tokens = all_token_ids[req_idx, sequence_start:total_len].tolist()
+        prefix = [-1] * contexts.shape[-1]
+        current_context = tuple(contexts[row].tolist())
+        history_start = (
+            max(0, len(history_tokens) - max_history) if max_history is not None else 0
+        )
+        for history_pos, token_id in enumerate(history_tokens):
+            if history_pos >= history_start and (
+                tuple(prefix[-contexts.shape[-1] :]) == current_context
+            ):
+                repeated[row] = True
+                break
+            prefix.append(token_id)
+    return repeated
+
+
+def repeated_context_mask(
+    all_token_ids: torch.Tensor,
+    req_indices: torch.Tensor,
+    prompt_lens: torch.Tensor,
+    total_lens: torch.Tensor,
+    contexts: torch.Tensor,
+    max_history: int | None = None,
+    include_prompt: bool = False,
+    skip_partial_context: bool = False,
+) -> torch.Tensor:
+    """Return, per row, whether the row's context already occurred in its history.
+
+    Args:
+        all_token_ids: `[max_num_reqs, max_model_len]` token ids of every request.
+        req_indices: Request slot per sampled row; -1 marks a padding row, which
+            is reported as not repeated.
+        prompt_lens: Prompt length per request slot.
+        total_lens: Prompt plus generated length per request slot.
+        contexts: `[num_rows, context_width]` context per row, padded with -1
+            before the start of the scanned history.
+        max_history: Number of most recent history positions searched, or
+            ``None`` for all of them. The window compared at each position
+            reaches `context_width` tokens further back.
+        include_prompt: Search the prompt as well as the generated tokens.
+        skip_partial_context: Mark contexts containing start padding so they use
+            ordinary sampling.
+
+    """
+    if max_history is not None and max_history < 1:
+        raise ValueError("max_history must be positive or None")
+    if all_token_ids.device.type == "cpu":
+        return _repeated_context_mask_cpu(
+            all_token_ids,
+            req_indices,
+            prompt_lens,
+            total_lens,
+            contexts,
+            max_history,
+            include_prompt,
+            skip_partial_context,
+        )
+
+    if contexts.stride(-1) != 1:
+        contexts = contexts.contiguous()
+    repeated = torch.empty(len(req_indices), dtype=torch.bool, device=contexts.device)
+    _repeated_context_mask_kernel[(len(req_indices),)](
+        repeated,
+        all_token_ids,
+        all_token_ids.stride(0),
+        req_indices,
+        prompt_lens,
+        total_lens,
+        contexts,
+        contexts.stride(0),
+        CONTEXT_WIDTH=contexts.shape[-1],
+        MAX_HISTORY=0 if max_history is None else max_history,
+        INCLUDE_PROMPT=include_prompt,
+        SKIP_PARTIAL_CONTEXT=skip_partial_context,
+        BLOCK=512,
+    )
+    return repeated
+
+
 @triton.jit
 def _mulhilo32(multiplier: tl.constexpr, value):
     high = tl_math.umulhi(multiplier, value)
@@ -74,13 +232,101 @@ def _uint32_to_uniform(value):
 
 
 @triton.jit
-def _gumbel_value(logits_ptr, output, mask):
-    logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32)
+def _philox_context_state(
+    contexts_row_ptr,
+    key_0,
+    key_1,
+    CONTEXT_WIDTH: tl.constexpr,
+):
+    state_0 = tl.full((), _CONTEXT_DOMAIN, tl.uint32)
+    state_1 = tl.full((), CONTEXT_WIDTH, tl.uint32)
+    state_2 = tl.full((), 0, tl.uint32)
+    state_3 = tl.full((), 0, tl.uint32)
+    for offset in range(0, CONTEXT_WIDTH, 4):
+        context_0 = tl.load(contexts_row_ptr + offset).to(tl.uint32)
+        if offset + 1 < CONTEXT_WIDTH:
+            context_1 = tl.load(contexts_row_ptr + offset + 1).to(tl.uint32)
+        else:
+            context_1 = tl.full((), _UINT32_MASK - 1, tl.uint32)
+        if offset + 2 < CONTEXT_WIDTH:
+            context_2 = tl.load(contexts_row_ptr + offset + 2).to(tl.uint32)
+        else:
+            context_2 = tl.full((), _UINT32_MASK - 2, tl.uint32)
+        if offset + 3 < CONTEXT_WIDTH:
+            context_3 = tl.load(contexts_row_ptr + offset + 3).to(tl.uint32)
+        else:
+            context_3 = tl.full((), _UINT32_MASK - 3, tl.uint32)
+        state_0, state_1, state_2, state_3 = _philox4x32_10(
+            (state_0 ^ context_0) & _UINT32_MASK,
+            (state_1 ^ context_1) & _UINT32_MASK,
+            (state_2 ^ context_2) & _UINT32_MASK,
+            (state_3 ^ context_3) & _UINT32_MASK,
+            (key_0 ^ offset) & _UINT32_MASK,
+            key_1,
+        )
+    return state_0, state_1, state_2, state_3
+
+
+@triton.jit
+def _philox_candidate_words(groups, state_0, state_1, state_2, state_3, key_0, key_1):
+    candidate_words = groups.to(tl.uint32)
+    vector_zero = candidate_words * 0
+    return _philox4x32_10(
+        candidate_words & _UINT32_MASK,
+        state_0 + vector_zero,
+        state_1 + vector_zero,
+        state_2 + vector_zero,
+        ((key_0 ^ state_3) & _UINT32_MASK) + vector_zero,
+        ((key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK) + vector_zero,
+    )
+
+
+@triton.jit
+def _philox_gumbel_from_logits(logits, word):
     logits = tl.where(logits != logits, float("-inf"), logits)
-    uniform = _uint32_to_uniform(output)
+    uniform = _uint32_to_uniform(word)
     return logits - tl.log(-tl.log(uniform))
 
 
+@triton.jit
+def _gumbel_value(logits_ptr, output, mask):
+    logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32)
+    return _philox_gumbel_from_logits(logits, output)
+
+
+@triton.jit
+def philox_gumbel_block_argmax(
+    logits,
+    mask,
+    block_idx,
+    contexts_row_ptr,
+    key_0,
+    key_1,
+    CONTEXT_WIDTH: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Apply the detector-compatible PRF to one vocabulary block."""
+    tl.static_assert(BLOCK_SIZE % 4 == 0)
+    state_0, state_1, state_2, state_3 = _philox_context_state(
+        contexts_row_ptr, key_0, key_1, CONTEXT_WIDTH
+    )
+    groups = block_idx * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4)
+    output_0, output_1, output_2, output_3 = _philox_candidate_words(
+        groups, state_0, state_1, state_2, state_3, key_0, key_1
+    )
+    # Lay the per-group words out in token order: 4g+0, 4g+1, 4g+2, 4g+3.
+    words = tl.interleave(
+        tl.interleave(output_0, output_2),
+        tl.interleave(output_1, output_3),
+    )
+    values = tl.where(
+        mask,
+        _philox_gumbel_from_logits(logits.to(tl.float32), words),
+        float("-inf"),
+    )
+    return tl.max(values, axis=0, return_indices=True)
+
+
 @triton.jit
 def _select_max(best_value, best_token, candidate_value, candidate):
     take_candidate = candidate_value > best_value
@@ -100,60 +346,63 @@ def _philox_gumbel_kernel(
     logits_stride,
     contexts_ptr,
     context_stride,
+    skip_mask_ptr,
+    expanded_idx_mapping_ptr,
+    seeds_ptr,
+    pos_ptr,
+    temp_ptr,
     key_0_value,
     key_1_value,
     vocab_size,
     CONTEXT_WIDTH: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
+    USE_FP64: tl.constexpr,
 ):
     row = tl.program_id(0).to(tl.int64)
     block_index = tl.program_id(1)
     groups = block_index * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4)
 
-    key_0 = key_0_value.to(tl.uint32)
-    key_1 = key_1_value.to(tl.uint32)
-    state_0 = tl.full((), _CONTEXT_DOMAIN, tl.uint32)
-    state_1 = tl.full((), CONTEXT_WIDTH, tl.uint32)
-    state_2 = tl.full((), 0, tl.uint32)
-    state_3 = tl.full((), 0, tl.uint32)
-    for offset in range(0, CONTEXT_WIDTH, 4):
-        context_0 = tl.load(contexts_ptr + row * context_stride + offset).to(tl.uint32)
-        if offset + 1 < CONTEXT_WIDTH:
-            context_1 = tl.load(contexts_ptr + row * context_stride + offset + 1).to(
-                tl.uint32
+    if skip_mask_ptr is not None:
+        req_state_idx = tl.load(expanded_idx_mapping_ptr + row).to(tl.int64)
+        valid_req = req_state_idx >= 0
+        if tl.load(skip_mask_ptr + row):
+            temp = tl.load(temp_ptr + req_state_idx, mask=valid_req, other=0.0).to(
+                tl.float32
             )
-        else:
-            context_1 = tl.full((), _UINT32_MASK - 1, tl.uint32)
-        if offset + 2 < CONTEXT_WIDTH:
-            context_2 = tl.load(contexts_ptr + row * context_stride + offset + 2).to(
-                tl.uint32
+            seed = tl.load(seeds_ptr + req_state_idx, mask=valid_req, other=0)
+            pos = tl.load(pos_ptr + row)
+            candidate = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = candidate < vocab_size
+            logits_row = logits_ptr + row * logits_stride
+            logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to(
+                tl.float32
             )
-        else:
-            context_2 = tl.full((), _UINT32_MASK - 2, tl.uint32)
-        if offset + 3 < CONTEXT_WIDTH:
-            context_3 = tl.load(contexts_ptr + row * context_stride + offset + 3).to(
-                tl.uint32
+            value, index = gumbel_noised_argmax(
+                logits,
+                candidate,
+                mask,
+                seed,
+                pos,
+                temp,
+                IS_DRAFTING=False,
+                USE_FP64=USE_FP64,
+                APPLY_TEMPERATURE=False,
             )
-        else:
-            context_3 = tl.full((), _UINT32_MASK - 3, tl.uint32)
-        state_0, state_1, state_2, state_3 = _philox4x32_10(
-            (state_0 ^ context_0) & _UINT32_MASK,
-            (state_1 ^ context_1) & _UINT32_MASK,
-            (state_2 ^ context_2) & _UINT32_MASK,
-            (state_3 ^ context_3) & _UINT32_MASK,
-            (key_0 ^ offset) & _UINT32_MASK,
-            key_1,
-        )
+            token_id = block_index * BLOCK_SIZE + index
+            tl.store(
+                local_argmax_ptr + row * local_argmax_stride + block_index,
+                token_id,
+            )
+            tl.store(local_max_ptr + row * local_max_stride + block_index, value)
+            return
 
-    candidate_words = groups.to(tl.uint32)
-    vector_zero = candidate_words * 0
-    output_0, output_1, output_2, output_3 = _philox4x32_10(
-        candidate_words & _UINT32_MASK,
-        state_0 + vector_zero,
-        state_1 + vector_zero,
-        state_2 + vector_zero,
-        ((key_0 ^ state_3) & _UINT32_MASK) + vector_zero,
-        ((key_1 ^ _TOKEN_DOMAIN) & _UINT32_MASK) + vector_zero,
+    key_0 = key_0_value.to(tl.uint32)
+    key_1 = key_1_value.to(tl.uint32)
+    state_0, state_1, state_2, state_3 = _philox_context_state(
+        contexts_ptr + row * context_stride, key_0, key_1, CONTEXT_WIDTH
+    )
+    output_0, output_1, output_2, output_3 = _philox_candidate_words(
+        groups, state_0, state_1, state_2, state_3, key_0, key_1
     )
     candidate_0 = groups * 4
     candidate_1 = candidate_0 + 1
@@ -184,17 +433,42 @@ def _philox_gumbel_kernel(
 
 
 def philox_gumbel_sample(
-    logits: torch.Tensor, contexts: torch.Tensor, key: int
+    logits: torch.Tensor,
+    contexts: torch.Tensor,
+    key: int,
+    *,
+    skip_mask: torch.Tensor | None = None,
+    expanded_idx_mapping: torch.Tensor | None = None,
+    temperatures: torch.Tensor | None = None,
+    seeds: torch.Tensor | None = None,
+    positions: torch.Tensor | None = None,
+    use_fp64: bool = False,
 ) -> torch.Tensor:
+    sampling_state = (expanded_idx_mapping, temperatures, seeds, positions)
+    if skip_mask is None and any(value is not None for value in sampling_state):
+        raise ValueError("sampling state requires skip_mask")
+    if skip_mask is not None and any(value is None for value in sampling_state):
+        raise ValueError("skip_mask requires complete sampling state")
+
     if logits.stride(-1) != 1:
         logits = logits.contiguous()
     if contexts.stride(-1) != 1:
         contexts = contexts.contiguous()
+    if skip_mask is not None:
+        skip_mask = skip_mask.contiguous()
+    if expanded_idx_mapping is not None:
+        expanded_idx_mapping = expanded_idx_mapping.contiguous()
+    if positions is not None:
+        positions = positions.contiguous()
     num_tokens, vocab_size = logits.shape
     block_size = 1024
     num_blocks = triton.cdiv(vocab_size, block_size)
     local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)
-    local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float32)
+    local_max = logits.new_empty(
+        num_tokens,
+        num_blocks,
+        dtype=torch.float64 if use_fp64 else torch.float32,
+    )
     _philox_gumbel_kernel[(num_tokens, num_blocks)](
         local_argmax,
         local_argmax.stride(0),
@@ -204,11 +478,17 @@ def philox_gumbel_sample(
         logits.stride(0),
         contexts,
         contexts.stride(0),
+        skip_mask,
+        expanded_idx_mapping,
+        seeds,
+        positions,
+        temperatures,
         key & _UINT32_MASK_VALUE,
         key >> 32,
         vocab_size,
         CONTEXT_WIDTH=contexts.shape[-1],
         BLOCK_SIZE=block_size,
+        USE_FP64=use_fp64,
     )
     max_block_index = local_max.argmax(dim=-1, keepdim=True)
     return local_argmax.gather(dim=-1, index=max_block_index).view(-1)
diff --git a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py
new file mode 100644
index 000000000000..46065d0546e7
--- /dev/null
+++ b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py
@@ -0,0 +1,501 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
+    _compute_max_and_sumexp,
+)
+
+# The feature is a log-odds, log(q / (1 - q)), which a row carrying all of its
+# mass on one token sends to infinity. Clamp it symmetrically.
+_MAX_LOG_ODDS = 40.0
+
+# Cold-start coefficients: the median slope and intercept fitted online across
+# DeepSeek-V4-Flash, Muse-Glimmer-30B, MiMo-V2.5-Pro, Inkling and Kimi-K2.5.
+# Log-loss is flat near the optimum, so one starting point serves all of them
+# and a real fit replaces it within a refit or two.
+_INIT_SLOPE = 0.5
+_INIT_BIAS = -0.2
+
+
+@triton.jit
+def _accumulate_kernel(
+    info_ptr,
+    grad_ptr,
+    idx_mapping_ptr,
+    num_sampled_ptr,
+    num_rejected_ptr,
+    features_ptr,
+    features_stride,
+    pred_ptr,
+    pred_stride,
+    counts_ptr,
+    num_reqs,
+    BLOCK_R: tl.constexpr,
+):
+    step = tl.program_id(0).to(tl.int64)
+    req_block = tl.arange(0, BLOCK_R)
+    req_mask = req_block < num_reqs
+    req_state_block = tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0).to(
+        tl.int64
+    )
+
+    num_sampled = tl.load(num_sampled_ptr + req_block, mask=req_mask, other=0).to(
+        tl.int64
+    )
+    # Subtract bonus token.
+    num_accepted = tl.maximum(num_sampled - 1, 0)
+    num_rejected = tl.load(num_rejected_ptr + req_block, mask=req_mask, other=0).to(
+        tl.int64
+    )
+    num_admitted = num_accepted + num_rejected
+
+    feature = tl.load(
+        features_ptr + req_state_block * features_stride + step,
+        mask=req_mask,
+        other=0.0,
+    ).to(tl.float32)
+    pred = tl.load(
+        pred_ptr + req_state_block * pred_stride + step,
+        mask=req_mask,
+        other=0.0,
+    ).to(tl.float32)
+
+    # The accepted tokens and the first rejected draft token (if one exists) are
+    # "observed" and given labels.
+    observed = req_mask & (step <= num_accepted) & (step < num_admitted)
+    label = tl.where(step < num_accepted, 1.0, 0.0)
+    w = tl.where(observed, pred * (1.0 - pred), 0.0)
+    resid = tl.where(observed, label - pred, 0.0)
+
+    # This block's contribution to the normal equations for design row
+    # x = [feature, eₖ]: the XᵀWX entries aₖ, bₖ, cₖ and the Xᵀr entries g₀ₖ, g₁ₖ.
+    # aₖ and g₀ₖ are partials of scalars shared by every position, which
+    # _refit_kernel sums.
+    a_k = tl.sum(w * feature * feature, axis=0)
+    b_k = tl.sum(w * feature, axis=0)
+    c_k = tl.sum(w, axis=0)
+    g0_k = tl.sum(resid * feature, axis=0)
+    g1_k = tl.sum(resid, axis=0)
+    count = tl.sum(tl.where(observed, 1.0, 0.0), axis=0)
+    tl.store(info_ptr + step * 3 + 0, tl.load(info_ptr + step * 3 + 0) + a_k)
+    tl.store(info_ptr + step * 3 + 1, tl.load(info_ptr + step * 3 + 1) + b_k)
+    tl.store(info_ptr + step * 3 + 2, tl.load(info_ptr + step * 3 + 2) + c_k)
+    tl.store(grad_ptr + step * 2 + 0, tl.load(grad_ptr + step * 2 + 0) + g0_k)
+    tl.store(grad_ptr + step * 2 + 1, tl.load(grad_ptr + step * 2 + 1) + g1_k)
+    tl.store(counts_ptr + step, tl.load(counts_ptr + step) + count)
+
+
+@triton.jit
+def _refit_kernel(
+    slope_ptr,
+    intercepts_ptr,
+    info_ptr,
+    info_row_stride,
+    grad_ptr,
+    grad_row_stride,
+    counts_ptr,
+    NUM_SPECULATIVE_STEPS: tl.constexpr,
+    L2: tl.constexpr,
+    DAMPING: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    k = tl.arange(0, BLOCK)
+    mask = k < NUM_SPECULATIVE_STEPS
+
+    # Load the stats accumulated over this round to solve (XᵀWX + λI) Δθ = Xᵀr.
+    a_k = tl.load(info_ptr + k * info_row_stride + 0, mask=mask, other=0.0)
+    b_k = tl.load(info_ptr + k * info_row_stride + 1, mask=mask, other=0.0)
+    c_k = tl.load(info_ptr + k * info_row_stride + 2, mask=mask, other=0.0) + L2
+    a = tl.sum(a_k, axis=0) + L2
+    g0_k = tl.load(grad_ptr + k * grad_row_stride + 0, mask=mask, other=0.0)
+    g1_k = tl.load(grad_ptr + k * grad_row_stride + 1, mask=mask, other=0.0)
+    g0 = tl.sum(g0_k, axis=0)
+    n = tl.load(counts_ptr + k, mask=mask, other=0.0)
+
+    # Load the current shared slope and per-position intercepts.
+    w = tl.load(slope_ptr)
+    bias = tl.load(intercepts_ptr + k, mask=mask, other=0.0)
+
+    # Solve for the shared slope delta, Δω, in the following arrowhead system:
+    #   ⎡ a+λ   b₀     b₁    …  bₙ₋₁    ⎤ ⎡  Δω   ⎤   ⎡  g₀   ⎤
+    #   ⎢ b₀    c₀+λ                    ⎥ ⎢  Δβ₀  ⎥   ⎢  g₁₀  ⎥
+    #   ⎢ b₁           c₁+λ             ⎥ ⎢  Δβ₁  ⎥ = ⎢  g₁₁  ⎥
+    #   ⎢ ⋮                  ⋱          ⎥ ⎢  ⋮    ⎥   ⎢  ⋮    ⎥
+    #   ⎣ bₙ₋₁                  cₙ₋₁+λ   ⎦ ⎣ Δβₙ₋₁ ⎦   ⎣ g₁ₙ₋₁ ⎦
+    #               XᵀWX + λI                Δθ        Xᵀ(Y-P)
+    shrink = b_k * b_k / c_k
+    coupling = b_k * g1_k / c_k
+    step_w = (g0 - tl.sum(coupling, axis=0)) / (a - tl.sum(shrink, axis=0))
+
+    # The slope learns from every position, so it is damped by the round's
+    # total sample count.
+    total_n = tl.sum(n, axis=0)
+    step_w *= total_n / (total_n + DAMPING)
+    # Mask out NaNs and steps that would drive the slope negative.
+    step_w = tl.where((step_w == step_w) & (w + step_w >= 0.0), step_w, 0.0)
+    # Update the shared slope.
+    new_w = w + step_w
+
+    # Solve for the per-position intercept deltas, Δβₖ, in the same system.
+    step_bias = (g1_k - b_k * step_w) / c_k
+    # Damp by the position's sample count.
+    step_bias *= n / (n + DAMPING)
+    # Mask out NaNs.
+    step_bias = tl.where(step_bias == step_bias, step_bias, 0.0)
+    # Update the per-position intercepts.
+    new_bias = bias + step_bias
+
+    tl.store(slope_ptr, new_w)
+    tl.store(intercepts_ptr + k, new_bias, mask=mask)
+
+    # Start the round clean by clearing the per-round statistics.
+    tl.store(info_ptr + k * info_row_stride + 0, 0.0, mask=mask)
+    tl.store(info_ptr + k * info_row_stride + 1, 0.0, mask=mask)
+    tl.store(info_ptr + k * info_row_stride + 2, 0.0, mask=mask)
+    tl.store(grad_ptr + k * grad_row_stride + 0, 0.0, mask=mask)
+    tl.store(grad_ptr + k * grad_row_stride + 1, 0.0, mask=mask)
+    tl.store(counts_ptr + k, 0.0, mask=mask)
+
+
+@triton.jit
+def _local_max_sumexp_kernel(
+    local_max_ptr,
+    local_max_stride,
+    local_sumexp_ptr,
+    local_sumexp_stride,
+    logits_ptr,
+    logits_stride,
+    idx_mapping_ptr,
+    idx_mapping_stride,
+    temperature_ptr,
+    num_tokens,
+    vocab_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    token_idx = tl.program_id(0).to(tl.int64)
+    if token_idx >= num_tokens:
+        return
+
+    block_idx = tl.program_id(1)
+    req_state_idx = tl.load(idx_mapping_ptr + token_idx * idx_mapping_stride).to(
+        tl.int64
+    )
+    if req_state_idx < 0:
+        # Cudagraph-padded requests carry -1. Write an empty partial so the
+        # buffer holds nothing uninitialized; _predict_kernel skips these rows.
+        tl.store(
+            local_max_ptr + token_idx * local_max_stride + block_idx, float("-inf")
+        )
+        tl.store(local_sumexp_ptr + token_idx * local_sumexp_stride + block_idx, 0.0)
+        return
+
+    # Draft logits are pre-temperature, and acceptance is decided on the scaled
+    # distribution the drafter actually sampled from, so scale first. A request
+    # with temp == 0 drafts greedily, so leave its logits alone rather than dividing
+    # by zero, matching what the sampler does.
+    temp = tl.load(temperature_ptr + req_state_idx).to(tl.float32)
+    temp = tl.where(temp > 0.0, temp, 1.0)
+
+    offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < vocab_size
+    logits = (
+        tl.load(
+            logits_ptr + token_idx * logits_stride + offsets,
+            mask=mask,
+            other=float("-inf"),
+        ).to(tl.float32)
+        / temp
+    )
+    block_max, block_sumexp = _compute_max_and_sumexp(logits)
+    tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, block_max)
+    tl.store(
+        local_sumexp_ptr + token_idx * local_sumexp_stride + block_idx, block_sumexp
+    )
+
+
+@triton.jit
+def _predict_kernel(
+    features_ptr,
+    features_stride,
+    pred_ptr,
+    pred_stride,
+    conf_ptr,
+    conf_stride,
+    slope_ptr,
+    intercepts_ptr,
+    local_max_ptr,
+    local_max_stride,
+    local_sumexp_ptr,
+    local_sumexp_stride,
+    idx_mapping_ptr,
+    idx_mapping_stride,
+    step_ptr,
+    num_tokens,
+    vocab_num_blocks,
+    per_token_step: tl.constexpr,
+    NUM_SPECULATIVE_STEPS: tl.constexpr,
+    MAX_LOG_ODDS: tl.constexpr,
+    PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
+):
+    token_idx = tl.program_id(0).to(tl.int64)
+    if token_idx >= num_tokens:
+        return
+
+    req_state_idx = tl.load(idx_mapping_ptr + token_idx * idx_mapping_stride).to(
+        tl.int64
+    )
+    if req_state_idx < 0:
+        # Cudagraph-padded requests carry -1. Skip them so that they don't
+        # scatter garbage over a live request's features.
+        return
+
+    # A per-row step means the speculator sampled the whole block in one pass, so
+    # each request owns NUM_SPECULATIVE_STEPS consecutive rows; a step shared by
+    # every row means one row per request. Batch position, unlike req_state_idx,
+    # is only meaningful for the step that computed it, and is what adaptive
+    # verification indexes by.
+    if per_token_step:
+        step = tl.load(step_ptr + token_idx).to(tl.int64)
+        batch_idx = token_idx // NUM_SPECULATIVE_STEPS
+    else:
+        step = tl.load(step_ptr).to(tl.int64)
+        batch_idx = token_idx
+
+    blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
+    blocks_mask = blocks < vocab_num_blocks
+    maxes = tl.load(
+        local_max_ptr + token_idx * local_max_stride + blocks,
+        mask=blocks_mask,
+        other=float("-inf"),
+    )
+    sumexps = tl.load(
+        local_sumexp_ptr + token_idx * local_sumexp_stride + blocks,
+        mask=blocks_mask,
+        other=0.0,
+    )
+
+    # The feature, logit(q(m*)) for the row's max draft logit, m*, is:
+    #   q(x) = e^(x - m*) / Σ_z e^(z - m*)
+    #   logit(q(m*)) = log(q(m*) / (1 - q(m*)))
+    #   	= m* - log(Σ_{z≠m*} e^(z - m*)).
+    # It scores the draft distribution, not the token drawn from it, so it is
+    # a function of the prefix alone. Trimming a draft therefore never
+    # conditions on the token being trimmed, and rejection sampling stays
+    # lossless. Scoring the drawn token instead would select on the proposal
+    # and bias the emitted distribution toward the drafter's confident modes.
+
+    # Rescale each block's partial to the global max and sum. The most likely
+    # token contributes exactly e^(m - m) = 1, so removing it leaves the sum
+    # over every other token.
+    global_max = tl.max(maxes, axis=0)
+    total_sumexp = tl.sum(
+        tl.where(blocks_mask, sumexps * tl.exp(maxes - global_max), 0.0), axis=0
+    )
+    complement_sumexp = tl.maximum(total_sumexp - 1.0, 0.0)
+
+    # Shifted by m*, the numerator is e^(m* - m*) = 1, so the feature is just the
+    # negated complement log-sum-exponential. Clamp to the min/max log-odds.
+    feature = -tl.log(complement_sumexp)
+    feature = tl.minimum(tl.maximum(feature, -MAX_LOG_ODDS), MAX_LOG_ODDS)
+    tl.store(features_ptr + req_state_idx * features_stride + step, feature)
+
+    # Predict the acceptance probability from the feature and current coefficients.
+    weight = tl.load(slope_ptr)
+    bias = tl.load(intercepts_ptr + step)
+    prob = tl.sigmoid(weight * feature + bias)
+    tl.store(pred_ptr + req_state_idx * pred_stride + step, prob)
+    tl.store(conf_ptr + batch_idx * conf_stride + step, prob)
+
+
+class OnlineAcceptanceEstimator:
+    """Predicts per-position acceptance, and calibrates itself while serving.
+
+    Lifecycle per step, driven by ``DraftModelSpeculator``:
+
+    1. ``step`` folds the previous step's drafts, now graded by the target,
+       into the IRLS accumulators, and periodically solves for new coefficients,
+       independently but identically on every rank.
+    2. ``predict`` runs inside the captured draft graph, turning this step's
+       draft logits into acceptance probabilities for adaptive verification.
+
+    Trimming starts immediately: survival is a running product of sigmoids, so it
+    decreases with draft position whatever the coefficients are, and an unfitted
+    estimator degrades to uniform-depth truncation rather than to anything harmful.
+    """
+
+    # Refit every this many steps, accumulating samples in between.
+    REFIT_INTERVAL = 100
+    # Newton steps are damped by n / (n + DAMPING_OBSERVATIONS), so a round
+    # carrying this many observations moves a parameter half of a full step. It
+    # replaces a hard minimum-sample gate: a data-poor position keeps learning,
+    # just slowly, instead of freezing until it crosses a threshold. Lower values
+    # track a drifting workload faster, higher ones are steadier on thin data.
+    DAMPING_OBSERVATIONS = 50.0
+    # Ridge on the Newton solve.
+    L2 = 1e-3
+
+    def __init__(
+        self,
+        max_num_reqs: int,
+        num_speculative_steps: int,
+        device: torch.device,
+    ):
+        self.num_speculative_steps = num_speculative_steps
+        self.device = device
+        self._steps_since_refit = 0
+        self._refits = 0
+
+        # The learned coefficients used for predicting acceptance.
+        self.slope = torch.full((1,), _INIT_SLOPE, dtype=torch.float32, device=device)
+        self.intercepts = torch.full(
+            (num_speculative_steps,), _INIT_BIAS, dtype=torch.float32, device=device
+        )
+        # Holds logit(max q), which is used as the feature for the logistic.
+        # Stored in stable slots keyed by persistent request-state index.
+        self.features = torch.zeros(
+            max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device
+        )
+        # Holds the acceptance predictions made at draft time, computed as:
+        #   p = σ(ω * f + βₖ)
+        # where f is feature (logit(max q)) at draft position k, ω is the learned
+        # shared slope, and βₖ is the learned per-position intercept. Stored in
+        # stable slots keyed by the persistent request-state index.
+        self.predictions = torch.zeros(
+            max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device
+        )
+
+        # Per-round Newton-IRLS statistics, cleared after each refit. For observation
+        # i the design row is xᵢ = [fᵢ, eₖ], the IRLS weight is wᵢ = pᵢ(1-pᵢ), and the
+        # residual is yᵢ - pᵢ, where yᵢ is the label (1 for accepted, 0 for rejected).
+        # XᵀWX is then the negative Hessian of the Bernoulli log-likelihood loss and
+        # Xᵀ(Y-P) is its gradient, so the Newton step for θ = [ω, β₀ … βₙ₋₁] is the
+        # the solution of (XᵀWX + λI) Δθ = Xᵀ(Y-P).
+        #
+        # Columns [aₖ, bₖ, cₖ] are each position's contribution to the XᵀWX arrowhead:
+        #   ⎡ a    b₀   b₁ … bₙ₋₁ ⎤
+        #   ⎢ b₀   c₀             ⎥
+        #   ⎢ b₁        c₁        ⎥
+        #   ⎢ ⋮            ⋱      ⎥
+        #   ⎣ bₙ₋₁           cₙ₋₁ ⎦
+        # bₖ and cₖ are per-position, but a = Σₖ aₖ is one scalar shared by all of
+        # them.
+        self.info = torch.zeros(
+            num_speculative_steps, 3, dtype=torch.float32, device=device
+        )
+        # Columns [g₀ₖ, g₁ₖ] are each position's contribution to the Xᵀ(Y-P) vector:
+        #   ⎡ g₀    ⎤
+        #   ⎢ g₁₀   ⎥
+        #   ⎢ g₁₁   ⎥
+        #   ⎢  ⋮    ⎥
+        #   ⎣ g₁ₙ₋₁ ⎦
+        # As above, g₁ₖ is per-position and g₀ = Σₖ g₀ₖ is shared.
+        self.grad = torch.zeros(
+            num_speculative_steps, 2, dtype=torch.float32, device=device
+        )
+        # Number of graded drafts per-position in the current round, zeroed after
+        # each refit. _refit_kernel uses this to damp coefficient updates by:
+        # n / (n + DAMPING_OBSERVATIONS)
+        self.counts = torch.zeros(
+            num_speculative_steps, dtype=torch.float32, device=device
+        )
+
+    def step(
+        self,
+        idx_mapping: torch.Tensor,
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+    ) -> None:
+        # Accumulate the previous step's graded drafts into the IRLS statistics.
+        num_reqs = idx_mapping.shape[0]
+        _accumulate_kernel[(self.num_speculative_steps,)](
+            self.info,
+            self.grad,
+            idx_mapping,
+            num_sampled,
+            num_rejected,
+            self.features,
+            self.features.stride(0),
+            self.predictions,
+            self.predictions.stride(0),
+            self.counts,
+            num_reqs,
+            BLOCK_R=triton.next_power_of_2(max(num_reqs, 1)),
+        )
+
+        self._steps_since_refit += 1
+        if self._steps_since_refit < self.REFIT_INTERVAL:
+            return
+        self._steps_since_refit = 0
+
+        # Fit the coefficients to the accumulated statistics gathered over the
+        # course of the last REFIT_INTERVAL steps.
+        _refit_kernel[(1,)](
+            self.slope,
+            self.intercepts,
+            self.info,
+            self.info.stride(0),
+            self.grad,
+            self.grad.stride(0),
+            self.counts,
+            NUM_SPECULATIVE_STEPS=self.num_speculative_steps,
+            L2=self.L2,
+            DAMPING=self.DAMPING_OBSERVATIONS,
+            BLOCK=triton.next_power_of_2(self.num_speculative_steps),
+        )
+        self._refits += 1
+
+    def predict(
+        self,
+        logits: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        draft_step: torch.Tensor,
+        confidence_probs: torch.Tensor,
+        temperature: torch.Tensor,
+    ) -> None:
+        num_tokens, vocab_size = logits.shape
+        VOCAB_BLOCK_SIZE = 4096
+        num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE)
+        local_max = torch.empty(
+            num_tokens, num_blocks, dtype=torch.float32, device=self.device
+        )
+        local_sumexp = torch.empty_like(local_max)
+        _local_max_sumexp_kernel[(num_tokens, num_blocks)](
+            local_max,
+            local_max.stride(0),
+            local_sumexp,
+            local_sumexp.stride(0),
+            logits,
+            logits.stride(0),
+            idx_mapping,
+            idx_mapping.stride(0),
+            temperature,
+            num_tokens,
+            vocab_size,
+            BLOCK_SIZE=VOCAB_BLOCK_SIZE,
+        )
+        _predict_kernel[(num_tokens,)](
+            self.features,
+            self.features.stride(0),
+            self.predictions,
+            self.predictions.stride(0),
+            confidence_probs,
+            confidence_probs.stride(0),
+            self.slope,
+            self.intercepts,
+            local_max,
+            local_max.stride(0),
+            local_sumexp,
+            local_sumexp.stride(0),
+            idx_mapping,
+            idx_mapping.stride(0),
+            draft_step,
+            num_tokens,
+            num_blocks,
+            per_token_step=draft_step.dim() > 0,
+            NUM_SPECULATIVE_STEPS=self.num_speculative_steps,
+            MAX_LOG_ODDS=_MAX_LOG_ODDS,
+            PADDED_VOCAB_NUM_BLOCKS=triton.next_power_of_2(num_blocks),
+        )
diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
index d3b3743d67fd..f47b4b7cbfe2 100644
--- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
+++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
@@ -10,9 +10,11 @@
 import torch
 
 import vllm.envs as envs
+from vllm.config.compilation import CUDAGraphMode
 from vllm.distributed.parallel_state import get_tp_group
 from vllm.logger import init_logger
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.utils import CpuGpuBuffer
 from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream
@@ -20,7 +22,6 @@
     get_attn_cg_support,
     get_query_lens_mismatch_unsupported_backend,
 )
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
 
 logger = init_logger(__name__)
 _PROFILE_REPLAYS = 5
@@ -33,6 +34,16 @@
     from vllm.v1.worker.utils import AttentionGroup
 
 
+def resolve_adaptive_cudagraph_mode(
+    mode: CUDAGraphMode, *, piecewise_capture_available: bool
+) -> CUDAGraphMode:
+    """Select a separate decode route for variable-length verification."""
+    assert mode.has_full_cudagraphs()
+    if mode == CUDAGraphMode.FULL_DECODE_ONLY or not piecewise_capture_available:
+        return CUDAGraphMode.FULL_DECODE_ONLY
+    return CUDAGraphMode.FULL_AND_PIECEWISE
+
+
 def _assign_draft_token_budget(
     confidence_probs: torch.Tensor,
     idx_mapping: torch.Tensor,
@@ -77,9 +88,14 @@ def build_cost_tables_from_curves(
     """Build cost tables: graph-padded below the capture limit, smooth above.
 
     Args:
+        draft_curve: (size, cost) samples for the draft model.
+        verify_curve: (size, cost) samples for the verify model.
+        max_num_reqs: Largest request count to build a table for.
+        max_batch_tokens: Largest token count to build a table for.
         cudagraph_limit: Largest cudagraph-captured size. At or below it,
             execution pads up to the next captured size, so cost is a step
             function. Above it there is no padding, so cost is continuous.
+
     """
 
     def build_table(limit: int, curve: list[tuple[int, float]]) -> np.ndarray:
@@ -239,9 +255,7 @@ def set_cost_curves(
         logger.debug("DSpark cost tables: %s", self.cost_tables)
 
     def record_confidences(
-        self,
-        confidence_probs: torch.Tensor,
-        input_batch: "InputBatch",
+        self, confidence_probs: torch.Tensor, input_batch: "InputBatch"
     ) -> None:
         """Publish this step's raw confidences for the ranking kernel and start
         copying them to the CPU, where a later step's budget reads them."""
@@ -289,9 +303,9 @@ def get_num_tokens(
         )
         num_non_draft_tokens = scheduled_tokens - scheduled_drafts
         slots = np.fromiter(
-            (self.req_states.req_id_to_index[req_id] for req_id in req_ids),
+            map(self.req_states.req_id_to_index.__getitem__, req_ids),
             dtype=np.int32,
-            count=len(req_ids),
+            count=num_reqs,
         )
         stale_confidences = self._stale_confidences[self._stale_idx].np[slots]
         survival_probability = np.cumprod(stale_confidences.astype(np.float64), axis=1)
@@ -402,7 +416,7 @@ def reallocate_drafts(
         if draft_budget == 0:
             capacities.zero_()
         else:
-            async_copy_to_gpu(scheduled_drafts, out=capacities)
+            async_tensor_h2d(scheduled_drafts, out=capacities)
             if draft_budget < int(scheduled_drafts.sum()):
                 _assign_draft_token_budget_compiled(
                     self._confidence_probs,
@@ -413,10 +427,7 @@ def reallocate_drafts(
                 )
 
         num_non_draft_tokens_gpu = self._num_non_draft_tokens[:num_reqs]
-        async_copy_to_gpu(
-            num_non_draft_tokens,
-            out=num_non_draft_tokens_gpu,
-        )
+        async_tensor_h2d(num_non_draft_tokens, out=num_non_draft_tokens_gpu)
         self._cu_num_logits[:1].zero_()
         torch.cumsum(
             capacities + self.num_bonus_tokens,
@@ -456,8 +467,7 @@ def maybe_create_adaptive_verification_manager(
     # The selector rejects unsupported backends, but models that
     # hard-wire theirs (e.g. DeepSeek-V4) never go through it.
     backend = get_query_lens_mismatch_unsupported_backend(
-        attn_groups,
-        checked_layer_names=target_layer_names,
+        attn_groups, checked_layer_names=target_layer_names
     )
     if backend is not None:
         raise ValueError(
diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py
index ef3b6e2ed535..fe3bdef269ea 100644
--- a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py
@@ -57,6 +57,7 @@ def create_forward_fn(
                 attn_groups,
                 kv_cache_config,
                 full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
+                max_query_len=desc.max_query_len or desc.uniform_token_count,
             )
 
             return lambda cg_mode: forward_fn(
diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
index bba4e2b4bd31..ca0662076b31 100644
--- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
@@ -53,10 +53,7 @@ def load_model(self, target_model: nn.Module) -> None:
             return
 
         self.inputs_embeds = torch.zeros(
-            self.max_num_tokens,
-            self.hidden_size,
-            dtype=self.dtype,
-            device=self.device,
+            self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=self.device
         )
 
     # Lifecycle hooks for model-specific optimizations. Subclasses override
@@ -73,8 +70,7 @@ def on_multi_step_decode_end(self, num_reqs: int) -> None: ...
 
     @property
     def advance_draft_positions(self) -> bool:
-        """
-        Whether to increment positions and seq_lens between draft steps.
+        """Whether to increment positions and seq_lens between draft steps.
 
         True for Eagle/standard MTP (each step produces new KV).
         False for Gemma4 MTP (Q-only, shares target KV, constant positions).
@@ -141,10 +137,7 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
 
         # Initialize cudagraph manager for draft decodes (draft positions > 0).
         self.decode_cudagraph_manager = SpeculatorCudaGraphManager(
-            self.vllm_config,
-            self.device,
-            cudagraph_mode,
-            decode_query_len=1,
+            self.vllm_config, self.device, cudagraph_mode, decode_query_len=1
         )
 
     def capture(self) -> None:
@@ -250,13 +243,12 @@ def propose(
             )
         else:
             hidden_states = last_hidden_states
-        self.hidden_states[:num_tokens_padded].copy_(hidden_states)
-
         self._copy_request_inputs(
             num_reqs,
             input_batch.idx_mapping,
             temperature,
             seeds,
+            dummy_run=dummy_run,
         )
 
         # Get the input ids and last token indices for the speculator.
@@ -272,6 +264,16 @@ def propose(
             self.max_num_reqs,
         )
 
+        if self.pcp_manager is not None:
+            self.pcp_manager.prepare_draft_prefill(
+                input_batch, self.input_buffers.input_ids[:num_tokens_padded]
+            )
+            prefill = self.pcp_manager.draft_prefill_batch
+            if prefill is not None:
+                num_tokens = prefill.num_tokens
+                num_tokens_padded = prefill.num_tokens_after_padding
+        self.hidden_states[:num_tokens_padded].copy_(hidden_states)
+
         # When all requests are decoding (no true prefills), each has
         # num_speculative_steps + 1 tokens, enabling FULL graph replay.
         uniform_token_count = get_uniform_decode_token_count(
@@ -324,6 +326,11 @@ def propose(
             # Early exit.
             return self.draft_tokens[:num_reqs, :1]
 
+        if self.pcp_manager is not None and not dummy_run:
+            self.block_tables.gather_block_tables(
+                input_batch.idx_mapping, num_reqs_padded=num_reqs
+            )
+
         # Prepare the inputs for the decode steps.
         prepare_decode_inputs(
             self.draft_tokens[:num_reqs, 0],
@@ -387,6 +394,11 @@ def _run_model(
         cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_buffers: InputBatch | InputBuffers = self.input_buffers
+        is_padding = None
+        if self.pcp_manager is not None:
+            input_buffers = self.pcp_manager.get_draft_input_buffers(self.input_buffers)
+            is_padding = input_buffers.is_padding[:num_tokens]
         batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
         with set_forward_context(
             attn_metadata,
@@ -396,6 +408,7 @@ def _run_model(
             num_tokens_across_dp=num_tokens_across_dp,
             slot_mapping=slot_mappings,
             batch_descriptor=batch_descriptor,
+            is_padding=is_padding,
         ):
             inputs_embeds = None
             if self.supports_mm_inputs:
@@ -406,15 +419,15 @@ def _run_model(
                     is_mm_embed.shape[0] if is_mm_embed is not None else num_tokens
                 )
                 self.inputs_embeds[:num_input_tokens] = self.model.embed_input_ids(
-                    self.input_buffers.input_ids[:num_input_tokens],
+                    input_buffers.input_ids[:num_input_tokens],
                     multimodal_embeddings=mm_embeds,
                     is_multimodal=is_mm_embed,
                 )
                 inputs_embeds = self.inputs_embeds[:num_tokens]
 
             model_inputs = dict(
-                input_ids=self.input_buffers.input_ids[:num_tokens],
-                positions=self.input_buffers.positions[:num_tokens],
+                input_ids=input_buffers.input_ids[:num_tokens],
+                positions=input_buffers.positions[:num_tokens],
                 hidden_states=self.hidden_states[:num_tokens],
                 inputs_embeds=inputs_embeds,
             )
@@ -463,6 +476,10 @@ def _prefill(
             cudagraph_runtime_mode=cudagraph_runtime_mode,
             mm_inputs=mm_inputs,
         )
+        if self.pcp_manager is not None:
+            last_hidden_states, hidden_states = self.pcp_manager.restore_draft_prefill(
+                last_hidden_states, hidden_states
+            )
 
         sample_hidden_states = last_hidden_states[last_token_indices]
         self.draft_tokens[:num_reqs, 0] = self.sample_draft(
@@ -508,11 +525,10 @@ def _multi_step_decode(
                 slot_mappings_by_layer = build_slot_mappings_by_layer(
                     slot_mappings, self.kv_cache_config
                 )
-                attn_metadata = self._build_draft_attn_metadata(
+                attn_metadata = self._build_uniform_attn_metadata(
                     num_reqs=num_reqs,
-                    num_reqs_padded=batch_desc.num_reqs or num_reqs,
-                    # One query per request; exclude DP-only model padding.
-                    num_tokens_padded=batch_desc.num_reqs or num_reqs,
+                    batch_desc=batch_desc,
+                    num_query_per_req=1,
                     seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
                     step=step,
                 )
@@ -557,11 +573,10 @@ def _fused_multi_step_decode(
                 slot_mappings_by_layer = build_slot_mappings_by_layer(
                     slot_mappings, self.kv_cache_config
                 )
-            attn_metadata = self._build_draft_attn_metadata(
+            attn_metadata = self._build_uniform_attn_metadata(
                 num_reqs=num_reqs,
-                num_reqs_padded=batch_desc.num_reqs or num_reqs,
-                # One query per request; exclude DP-only model padding.
-                num_tokens_padded=batch_desc.num_reqs or num_reqs,
+                batch_desc=batch_desc,
+                num_query_per_req=1,
                 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
                 step=1,
             )
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
index 725714fcdbc0..e8f8be7b64aa 100644
--- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
+++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
@@ -11,7 +11,7 @@
     build_slot_mappings_by_layer,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import (
     AttentionState,
     BatchExecutionDescriptor,
@@ -42,15 +42,16 @@ def _prepare_dflash_inputs_to_capture(
     attn_metadata = None
     if not skip_attn:
         query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
-        input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
-            input_buffers.dcp_local_seq_lens,
-            input_batch.seq_lens,
-            input_batch.num_reqs,
-            block_tables.cp_size,
-            block_tables.cp_rank,
-            block_tables.cp_interleave,
-            num_reqs_padded=input_batch.num_reqs_after_padding,
-        )
+        if block_tables.cp_size > 1:
+            input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
+                input_buffers.dcp_local_seq_lens,
+                input_batch.seq_lens,
+                input_batch.num_reqs,
+                block_tables.cp_size,
+                block_tables.cp_rank,
+                block_tables.cp_interleave,
+                num_reqs_padded=input_batch.num_reqs_after_padding,
+            )
         attn_metadata = build_attn_metadata(
             attn_groups=attn_groups,
             num_reqs=num_reqs,
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
index a3e8df79243a..a5710829e0ff 100644
--- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
@@ -1,10 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import copy
-from collections.abc import Mapping
 from typing import Any
 
-import numpy as np
 import torch
 import torch.nn as nn
 
@@ -38,6 +36,13 @@ class DFlashSpeculator(DraftModelSpeculator):
     _speculator_name = "DFlash"  # For logging, so we can share methods with subclasses
 
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        parallel_config = vllm_config.parallel_config
+        if parallel_config.prefill_context_parallel_size > 1:
+            vllm_config = copy.copy(vllm_config)
+            vllm_config.parallel_config = replace(
+                parallel_config,
+                prefill_context_parallel_size=1,
+            )
         super().__init__(vllm_config, device)
 
         self.hidden_states = torch.zeros(
@@ -292,33 +297,6 @@ def _generate_draft(
             num_reqs, self.num_speculative_steps
         )
 
-    def _build_draft_attn_metadata(
-        self,
-        num_reqs: int,
-        num_reqs_padded: int,
-        num_tokens_padded: int,
-        seq_lens_cpu_upper_bound: torch.Tensor,
-        step: int,
-        num_query_per_req: int | None = None,
-        causal: bool | Mapping[int, bool] = False,
-        query_start_loc_np: np.ndarray | None = None,
-        dcp_local_seq_lens: torch.Tensor | None = None,
-    ) -> dict[str, Any] | None:
-        if not self.draft_attn_layer_names:
-            return None
-        assert num_query_per_req is None  # Omitted for DFlash, read from self instead
-        return super()._build_draft_attn_metadata(
-            num_reqs,
-            num_reqs_padded,
-            num_tokens_padded,
-            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
-            step=step,
-            num_query_per_req=self.num_query_per_req,
-            causal=causal,
-            query_start_loc_np=query_start_loc_np,
-            dcp_local_seq_lens=dcp_local_seq_lens,
-        )
-
     @torch.inference_mode()
     def propose(
         self,
@@ -388,6 +366,11 @@ def propose(
             )
             return self.draft_tokens[:num_reqs]
 
+        if self.pcp_manager is not None and not dummy_run:
+            self.block_tables.gather_block_tables(
+                input_batch.idx_mapping, num_reqs_padded=num_reqs
+            )
+
         # The query slot mapping is written into the shared BlockTables slot_mappings.
         # That buffer's address is what the captured CUDA graph reads from at replay.
         assert self.draft_kv_cache_group_id >= 0
@@ -459,7 +442,6 @@ def propose(
             need_eager=is_profile,
             dp_sync=batch_sync,
         )
-        num_reqs_padded = batch_desc.num_reqs or num_reqs
         num_tokens_padded = batch_desc.num_tokens
         num_tokens_across_dp = (
             batch_sync.num_tokens_across_dp if batch_sync is not None else None
@@ -467,10 +449,10 @@ def propose(
 
         # Rebuild the draft attention metadata even when replaying the FULL
         # graph so that any attention metadata builder state is updated.
-        draft_attn_metadata = self._build_draft_attn_metadata(
+        draft_attn_metadata = self._build_uniform_attn_metadata(
             num_reqs=num_reqs,
-            num_reqs_padded=num_reqs_padded,
-            num_tokens_padded=num_tokens_padded,
+            batch_desc=batch_desc,
+            num_query_per_req=self.num_query_per_req,
             seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
             step=self.num_query_per_req,
             causal=self._group_causal,
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py
index 6226dea1c2db..2e5826362346 100644
--- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py
@@ -4,11 +4,13 @@
 
 from vllm.config import VllmConfig, replace
 from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.model_loader.utils import get_draft_load_config
 from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
     _should_share,
     get_target_lm_head,
     maybe_share_target_embed,
 )
+from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
 
 
 def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
@@ -37,6 +39,7 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
             if speculative_config.kv_cache_dtype is not None
             else vllm_config.cache_config
         ),
+        load_config=get_pp_safe_draft_load_config(get_draft_load_config(vllm_config)),
     )
     with set_model_tag("dflash_head"):
         dflash_model = get_model(
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py
index 1cd19acdab23..7db52844979e 100644
--- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py
@@ -213,5 +213,11 @@ def _generate_draft(
             anchor_token_ids,
         )
         self._sample_path(candidate_ids, scores, num_reqs)
+        if self.enable_adaptive_verification:
+            self._maybe_predict_acceptance(
+                self._selector_scores[:num_reqs].flatten(0, 1),
+                self.sample_idx_mapping[:num_sample],
+                self.sample_col[:num_sample],
+            )
         if self.draft_logits is not None:
             self._cache_draft_logits(candidate_ids, num_sample)
diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
index e9eca550ca9a..b86f58a45063 100644
--- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
@@ -79,12 +79,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
             self.draft_model_config.hf_config, "dspark_draft_topk", None
         )
 
-        self.draft_token_confidence_probs = torch.empty_like(
-            self.draft_tokens, dtype=torch.float32
-        )
-        self.enable_adaptive_verification = (
-            self.speculative_config.enable_adaptive_verification
-        )
+        self.use_confidence_head: bool = False
 
     def load_draft_model(
         self,
@@ -108,13 +103,14 @@ def load_draft_model(
                 dtype=self.draft_logits.dtype,
                 device=self.device,
             )
-        if self.enable_adaptive_verification and model.model.confidence_head is None:
-            raise ValueError(
-                "Adaptive verification needs a DSpark checkpoint with a confidence "
-                "head, and this one has none. Pass "
-                "enable_adaptive_verification=false in the speculative config to verify"
-                " a fixed number of drafts instead."
-            )
+        self.use_confidence_head = (
+            self.enable_adaptive_verification
+            and model.model.confidence_head is not None
+        )
+        if self.use_confidence_head:
+            # The acceptance estimator is not needed when a trained confidence head
+            # is available.
+            self.use_acceptance_estimator = False
         return model
 
     def _sample_logits(
@@ -124,8 +120,10 @@ def _sample_logits(
         sample_pos: torch.Tensor,
         step: int,
     ) -> torch.Tensor:
+        self._maybe_predict_acceptance(logits, idx_map, self._step_cols[step])
         if self.draft_logits is None:
-            return self.model.map_draft_to_target(logits.argmax(dim=-1))
+            draft_ids = logits.argmax(dim=-1)
+            return self.model.map_draft_to_target(draft_ids)
 
         # Probabilistic sampling and rejection operate in target-vocabulary
         # space. A reduced draft vocabulary is scattered into its target rows.
@@ -137,7 +135,7 @@ def _sample_logits(
 
         # sample_pos is the predicted token's position P. Sampling keys a draw
         # by the position before the sampled token, P-1.
-        return gumbel_sample(
+        sampled = gumbel_sample(
             logits,
             idx_map,
             self.temperature,
@@ -149,6 +147,11 @@ def _sample_logits(
             logits_cache_col=self._step_cols[step],
             use_fp64=self.use_fp64_gumbel,
         )
+        if self.draft_watermarker is not None:
+            sampled = self.draft_watermarker.sample(
+                logits, sampled, idx_map, self.temperature
+            )
+        return sampled
 
     def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
         if self._draft_topk is not None:
@@ -176,7 +179,7 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
         for i in range(n_spec):
             # Sequential stage: Markov bias from the previously sampled token.
             markov_embed = self.model.markov_embed(prev)
-            if self.enable_adaptive_verification:
+            if self.use_confidence_head:
                 confidence_markov_embeds.append(markov_embed)
             bias = self.model.markov_bias(markov_embed)
             logits_i = base_logits[:, i] + bias
@@ -186,7 +189,7 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
             self.draft_tokens[:num_reqs, i] = draft_sampled_i
             prev = draft_sampled_i
 
-        if self.enable_adaptive_verification:
+        if self.use_confidence_head:
             confidence = self.model.compute_confidence(
                 sample_hidden,
                 torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1),
@@ -221,7 +224,7 @@ def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> N
 
         for i in range(n_spec):
             markov_embed = self.model.markov_embed(prev)
-            if self.enable_adaptive_verification:
+            if self.use_confidence_head:
                 confidence_markov_embeds.append(markov_embed)
             logits_i = self.model.apply_markov_bias_gathered(
                 markov_embed,
@@ -235,7 +238,7 @@ def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> N
             self.draft_tokens[:num_reqs, i] = draft_sampled_i
             prev = draft_sampled_i
 
-        if self.enable_adaptive_verification:
+        if self.use_confidence_head:
             confidence = self.model.compute_confidence(
                 sample_hidden,
                 torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1),
diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py
index f252a35a51f8..489d0da554a2 100644
--- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py
@@ -3,9 +3,11 @@
 
 import torch.nn as nn
 
-from vllm.config import ModelConfig, VllmConfig, replace
+from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace
 from vllm.logger import init_logger
+from vllm.model_executor.model_loader.utils import get_draft_load_config
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
 
 logger = init_logger(__name__)
 
@@ -17,9 +19,9 @@ def _resolve_dspark_attention_backend(
 ) -> AttentionBackendEnum | None:
     if draft_backend is not None:
         return draft_backend
-    # DeepSeek-V4 draft layers share the target's KV-cache layout. Other
+    # DeepSeek-V4(.1) draft layers share the target's KV-cache layout. Other
     # DSpark architectures may use a different attention kind.
-    if draft_model_config.hf_config.model_type == "deepseek_v4":
+    if draft_model_config.hf_config.model_type in ("deepseek_v4", "deepseek_v41"):
         if target_backend is not None:
             logger.info_once(
                 "Using the target model's %s attention backend for the "
@@ -30,6 +32,29 @@ def _resolve_dspark_attention_backend(
     return None
 
 
+def _get_dspark_parallel_config(
+    parallel_config: ParallelConfig,
+    tensor_parallel_size: int,
+) -> ParallelConfig:
+    if parallel_config.enable_eplb:
+        logger.warning_once(
+            "EPLB is disabled for the DSpark draft model. EPLB remains enabled "
+            "for the target model."
+        )
+
+    return replace(
+        parallel_config,
+        pipeline_parallel_size=1,
+        tensor_parallel_size=tensor_parallel_size,
+        enable_eplb=False,
+        eplb_config=replace(
+            parallel_config.eplb_config,
+            num_redundant_experts=0,
+        ),
+        enable_elastic_ep=False,
+    )
+
+
 def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
     speculative_config = vllm_config.speculative_config
     assert speculative_config is not None
@@ -53,12 +78,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
 
     draft_vllm_config = replace(
         vllm_config,
-        parallel_config=replace(
+        parallel_config=_get_dspark_parallel_config(
             vllm_config.parallel_config,
-            pipeline_parallel_size=1,
-            tensor_parallel_size=(
-                speculative_config.draft_parallel_config.tensor_parallel_size
-            ),
+            speculative_config.draft_parallel_config.tensor_parallel_size,
         ),
         attention_config=replace(
             vllm_config.attention_config,
@@ -73,6 +95,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
             if speculative_config.kv_cache_dtype is not None
             else vllm_config.cache_config
         ),
+        load_config=get_pp_safe_draft_load_config(get_draft_load_config(vllm_config)),
     )
     # VllmConfig post-init restores the target's quant config because the target
     # config is retained for DSpark's target-layer metadata, so we must override it.
diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
index 0dca4588c332..6b0c5fe51d7a 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
@@ -112,7 +112,14 @@ def get_eagle3_aux_layers_from_config(
     if not layer_ids:
         dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None)
         if dspark_layer_ids:
-            layer_ids = [i + 1 for i in dspark_layer_ids]
+            if getattr(hf_config, "model_type", None) == "deepseek_v41":
+                # v4.1 reads the attention *inputs* of its target layers, and
+                # the target model captures the entry stream of layer L when
+                # idx+1 == L, so the ids are used as-is. (v4's ids are in
+                # capture-after semantics and keep the +1.)
+                layer_ids = list(dspark_layer_ids)
+            else:
+                layer_ids = [i + 1 for i in dspark_layer_ids]
     if not layer_ids:
         # Dense DSpark (e.g. Qwen3) also uses different aux layer semantics.
         target_layer_ids = getattr(hf_config, "target_layer_ids", None)
diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
index a85f21a3cd66..06bed9391e03 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py
@@ -7,12 +7,13 @@
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.lora.layers.base import BaseLayerWithLoRA
 from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.model_loader.utils import get_draft_load_config
 from vllm.model_executor.models.utils import PPMissingLayer
+from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
 
 
 def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
     """Share when the draft has no own copy, or its copy matches the target."""
-
     if not getattr(eagle, flag, False) or draft is None:
         return True
     if target is None:
@@ -75,34 +76,12 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
     speculative_config = vllm_config.speculative_config
     assert speculative_config is not None
     draft_model_config = speculative_config.draft_model_config
-    if speculative_config.moe_backend is not None:
-        # Otherwise the draft inherits the target's --moe-backend, which
-        # fails when the draft is unquantized and that backend is not.
-        vllm_config = replace(
-            vllm_config,
-            kernel_config=replace(
-                vllm_config.kernel_config,
-                moe_backend=speculative_config.moe_backend,
-            ),
-        )
-    if speculative_config.kv_cache_dtype is not None:
-        vllm_config = replace(
-            vllm_config,
-            cache_config=replace(
-                vllm_config.cache_config,
-                cache_dtype=speculative_config.kv_cache_dtype,
-            ),
-        )
-    if speculative_config.attention_backend is not None:
-        # Before get_model(): the backend is read off the constructed layers.
-        # Only when set, so the draft keeps a KV cache layout the target shares.
-        vllm_config = replace(
-            vllm_config,
-            attention_config=replace(
-                vllm_config.attention_config,
-                backend=speculative_config.attention_backend,
-            ),
-        )
+    vllm_config = speculative_config.apply_draft_overrides(vllm_config)
+    draft_load_config = get_pp_safe_draft_load_config(
+        get_draft_load_config(vllm_config)
+    )
+    if draft_load_config is not vllm_config.load_config:
+        vllm_config = replace(vllm_config, load_config=draft_load_config)
     with set_model_tag("eagle_head"):
         eagle_model = get_model(
             vllm_config=vllm_config, model_config=draft_model_config
diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py
index dfa2c680109d..6f76e5055a82 100644
--- a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py
@@ -16,6 +16,7 @@
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.model_loader.utils import get_draft_load_config
 from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
     AutoRegressiveSpeculator,
 )
@@ -23,6 +24,27 @@
 logger = init_logger(__name__)
 
 
+def _copy_target_kv_scales(attn: nn.Module, target_attn: nn.Module) -> None:
+    """Copy target KV scales while preserving their tensor representation.
+
+    Default attention scales are scalar buffers, while some quantization
+    methods replace them with length-one or per-head parameters. Preserve the
+    draft layer's registration type so the shared KV cache is interpreted with
+    the target's values and shapes without aliasing target parameters.
+    """
+    for scale_name in ("_k_scale", "_v_scale"):
+        draft_scale = getattr(attn, scale_name)
+        target_scale = getattr(target_attn, scale_name)
+        scale = target_scale.detach().clone()
+        if isinstance(draft_scale, nn.Parameter):
+            scale = nn.Parameter(scale, requires_grad=False)
+        setattr(attn, scale_name, scale)
+    for scale_name in ("_k_scale_float", "_v_scale_float"):
+        setattr(attn, scale_name, getattr(target_attn, scale_name))
+    for scale_name in ("_k_scale_cpu", "_v_scale_cpu"):
+        getattr(attn, scale_name).copy_(getattr(target_attn, scale_name))
+
+
 class Gemma4Speculator(AutoRegressiveSpeculator):
     @property
     def advance_draft_positions(self) -> bool:
@@ -40,7 +62,7 @@ def load_draft_model(
             draft_model = get_model(
                 vllm_config=draft_vllm_config,
                 model_config=self.speculative_config.draft_model_config,
-                load_config=self.speculative_config.draft_load_config,
+                load_config=get_draft_load_config(draft_vllm_config),
             )
         self._setup_gemma4_kv_sharing(draft_model, target_attn_layer_names)
         self._share_embeddings(draft_model, target_model)
@@ -76,7 +98,7 @@ def _setup_gemma4_kv_sharing(
         model: nn.Module,
         target_attn_layer_names: set[str],
     ) -> None:
-        """Wire draft layers to share KV with the target model.
+        """Wire draft layers to share KV and KV scales with the target model.
 
         Each draft decoder layer is mapped to the last non-KV-shared
         target layer of the same attention type (sliding or full).
@@ -91,7 +113,10 @@ def _setup_gemma4_kv_sharing(
             return
 
         target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)
-        num_non_shared = len(target_layer_types) - target_num_kv_shared
+        target_num_layers = getattr(
+            target_text_config, "num_hidden_layers", len(target_layer_types)
+        )
+        num_non_shared = target_num_layers - target_num_kv_shared
         type_to_target_indices: dict[str, list[int]] = defaultdict(list)
         for idx, lt in enumerate(target_layer_types[:num_non_shared]):
             type_to_target_indices[lt].append(idx)
@@ -127,6 +152,17 @@ def _setup_gemma4_kv_sharing(
             target_idx = candidates[-1]
             target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
             attn.kv_sharing_target_layer_name = target_layer_name
+
+            # KV-cache sharing aliases the cache tensor during allocation, but
+            # the quantization scales live on the Attention modules themselves.
+            # The BF16 draft model has no quantization config, so its K/V scales
+            # otherwise remain at the default 1.0 while it reads the target's
+            # calibrated FP8 cache.
+            target_attn = self.vllm_config.compilation_config.static_forward_context[
+                target_layer_name
+            ]
+            _copy_target_kv_scales(attn, target_attn)
+
             logger.info(
                 "Gemma4 MTP: draft layer %d (%s) -> %s",
                 draft_idx,
diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py
index 4889e9a750f8..d94702612d9b 100644
--- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py
@@ -28,7 +28,8 @@ def load_draft_model(
         # toggles skip_topk so step 0 computes MTP's own indices and
         # steps 1+ reuse them.
         self.share_mtp_topk_indices = (
-            getattr(draft_hf_config, "index_share_for_mtp_iteration", False)
+            self.vllm_config.parallel_config.prefill_context_parallel_size == 1
+            and getattr(draft_hf_config, "index_share_for_mtp_iteration", False)
             and hasattr(draft_model.model, "set_skip_topk")
             and hasattr(draft_model.model, "compact_topk_indices")
         )
diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
index 90572c03ccd5..fc0924e3a3f4 100644
--- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
@@ -166,6 +166,7 @@ def propose(
             input_batch.idx_mapping,
             temperature,
             seeds,
+            dummy_run=dummy_run,
         )
 
         num_tokens = input_batch.num_tokens
@@ -226,10 +227,9 @@ def propose(
             slot_mappings = build_slot_mappings_by_layer(
                 slot_mappings_tensor, self.kv_cache_config
             )
-            draft_attn_metadata = self._build_draft_attn_metadata(
+            draft_attn_metadata = self._build_attn_metadata(
                 num_reqs=num_reqs,
-                num_reqs_padded=batch_desc.num_reqs or num_reqs,
-                num_tokens_padded=batch_desc.num_tokens,
+                batch_desc=batch_desc,
                 query_start_loc_np=input_batch.query_start_loc_np,
                 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
                 step=0,
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
index a9027e0550f8..743d66010d00 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable, Iterator
+from typing import Any
 
 import numpy as np
 import torch
@@ -10,6 +11,7 @@
 from vllm.triton_utils import tl, triton
 from vllm.v1.outputs import LogprobsTensors
 from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
 from vllm.v1.worker.gpu.input_batch import (
     InputBatch,
     get_num_sampled_and_rejected,
@@ -78,8 +80,13 @@ def __init__(
         sampler: Sampler,
         spec_config: SpeculativeConfig,
         device: torch.device,
+        *,
+        watermark_key: int | None = None,
     ):
         self.sampler = sampler
+        self.watermark_key = watermark_key
+        if watermark_key is not None:
+            assert isinstance(sampler, GPUWatermarkSampler)
         self.num_speculative_steps = spec_config.num_speculative_tokens
         self.enable_adaptive_verification = spec_config.enable_adaptive_verification
         rejection_sample_method = spec_config.rejection_sample_method
@@ -127,8 +134,7 @@ def _get_logprobs_tensors(
         if expanded_logits:
             if self.enable_adaptive_verification:
                 # Adaptive verification keeps the true per-request boundaries
-                # on device only; cu_num_logits_np holds the pre-compacted
-                # layout.
+                # on device only; cu_num_logits_np holds the pre-compacted layout.
                 cu_num_generated_tokens = cu_num_logits.clone()
             else:
                 cu_num_generated_tokens = cu_num_logits_np.tolist()
@@ -141,6 +147,23 @@ def _get_logprobs_tensors(
             in ("raw_logits", "processed_logits"),
         )
 
+    def _watermarking_kwargs(
+        self,
+        draft_sampled: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        expanded_local_pos: torch.Tensor,
+    ) -> dict[str, Any]:
+        if self.watermark_key is None:
+            return {}
+        assert isinstance(self.sampler, GPUWatermarkSampler)
+        return {
+            "contexts": self.sampler._get_contexts(
+                expanded_idx_mapping, expanded_local_pos, draft_sampled
+            ),
+            "watermarking": self.sampler.watermarking.gpu,
+            "watermark_key": self.watermark_key,
+        }
+
     def _verify(
         self,
         logits: torch.Tensor,
@@ -152,6 +175,7 @@ def _verify(
         idx_mapping_np: np.ndarray,
         expanded_idx_mapping: torch.Tensor,
         expanded_local_pos: torch.Tensor,
+        seq_lens_upper_bound_np: np.ndarray,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         processed_logits = self.sampler.apply_sampling_params(
             logits,
@@ -161,6 +185,7 @@ def _verify(
             pos,
             draft_sampled,
             expanded_local_pos,
+            seq_lens_upper_bound_np,
         )
         sampled, num_sampled = rejection_sample(
             processed_logits,
@@ -177,6 +202,9 @@ def _verify(
             self.synthetic_conditional_rates,
             use_fp64=self.sampler.use_fp64_gumbel,
             use_block_verification=self.use_block_verification,
+            **self._watermarking_kwargs(
+                draft_sampled, expanded_idx_mapping, expanded_local_pos
+            ),
         )
         return processed_logits, sampled, num_sampled
 
@@ -224,6 +252,7 @@ def _verify_in_chunks(
                 input_batch.idx_mapping_np[start:end],
                 input_batch.expanded_idx_mapping[lo:hi],
                 input_batch.expanded_local_pos[lo:hi],
+                input_batch.seq_lens_cpu_upper_bound.numpy()[start:end],
             )
             chunk_logprobs = self._get_logprobs_tensors(
                 sampled,
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
index 29b20afea69a..cf804e6a6bd4 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
@@ -4,6 +4,7 @@
 
 from vllm.triton_utils import tl, tldevice, triton
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32
+from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax
 
 
 @triton.jit
@@ -693,6 +694,39 @@ def _rejection_kernel(
 
 
 @triton.jit
+def _seeded_resample_argmax(
+    residual_logits,
+    block,
+    mask,
+    resample_token_idx,
+    expanded_idx_mapping_ptr,
+    temp_ptr,
+    seed_ptr,
+    pos_ptr,
+    vocab_size,
+    USE_FP64: tl.constexpr,
+):
+    return gumbel_block_argmax(
+        residual_logits,
+        block,
+        mask,
+        resample_token_idx,
+        expanded_idx_mapping_ptr,
+        temp_ptr,
+        seed_ptr,
+        pos_ptr,
+        None,  # logits_cache_ptr
+        0,  # logits_cache_stride_0
+        0,  # logits_cache_stride_1
+        None,  # logits_cache_col_ptr
+        vocab_size,
+        IS_DRAFTING=False,
+        APPLY_TEMPERATURE=False,
+        USE_FP64=USE_FP64,
+    )
+
+
+@triton.jit(do_not_specialize=["watermark_key_0", "watermark_key_1"])
 def _resample_kernel(
     # [num_reqs, num_blocks]
     resampled_local_argmax_ptr,
@@ -727,11 +761,20 @@ def _resample_kernel(
     pos_ptr,
     # [num_logits]
     cumulative_log_p_ptr,
+    # [num_logits, CONTEXT_WIDTH]
+    contexts_ptr,
+    contexts_stride,
+    # [max_num_reqs], uint8 view of a bool tensor
+    watermarking_ptr,
+    watermark_key_0,
+    watermark_key_1,
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
     USE_FP64: tl.constexpr,
     USE_BLOCK_VERIFICATION: tl.constexpr,
+    CONTEXT_WIDTH: tl.constexpr,
+    WATERMARK: tl.constexpr,
 ):
     req_idx = tl.program_id(0)
     resample_idx = tl.load(rejected_step_ptr + req_idx)
@@ -823,24 +866,55 @@ def _resample_kernel(
         ).to(tl.float32)
 
     # Resample the rejected/bonus token.
-    value, idx = gumbel_block_argmax(
-        residual_logits,
-        block,
-        mask,
-        resample_token_idx,
-        expanded_idx_mapping_ptr,
-        temp_ptr,
-        seed_ptr,
-        pos_ptr,
-        None,  # logits_cache_ptr
-        0,  # logits_cache_stride_0
-        0,  # logits_cache_stride_1
-        None,  # logits_cache_col_ptr
-        vocab_size,
-        IS_DRAFTING=False,
-        APPLY_TEMPERATURE=False,
-        USE_FP64=USE_FP64,
-    )
+    if WATERMARK:
+        # Padded and greedy rows retain the stock draw.
+        is_watermarked = (
+            tl.load(
+                watermarking_ptr + req_state_idx,
+                mask=req_state_idx >= 0,
+                other=0,
+            )
+            != 0
+        )
+        if is_watermarked & (temp != 0.0):
+            watermark_value, idx = philox_gumbel_block_argmax(
+                residual_logits,
+                mask,
+                block_idx,
+                contexts_ptr + resample_token_idx * contexts_stride,
+                watermark_key_0.to(tl.uint32),
+                watermark_key_1.to(tl.uint32),
+                CONTEXT_WIDTH,
+                BLOCK_SIZE,
+            )
+            # Detector compatibility fixes the keyed draw at fp32.
+            value = watermark_value.to(tl.float64) if USE_FP64 else watermark_value
+        else:
+            value, idx = _seeded_resample_argmax(
+                residual_logits,
+                block,
+                mask,
+                resample_token_idx,
+                expanded_idx_mapping_ptr,
+                temp_ptr,
+                seed_ptr,
+                pos_ptr,
+                vocab_size,
+                USE_FP64=USE_FP64,
+            )
+    else:
+        value, idx = _seeded_resample_argmax(
+            residual_logits,
+            block,
+            mask,
+            resample_token_idx,
+            expanded_idx_mapping_ptr,
+            temp_ptr,
+            seed_ptr,
+            pos_ptr,
+            vocab_size,
+            USE_FP64=USE_FP64,
+        )
     token_id = block_idx * BLOCK_SIZE + idx
     tl.store(
         resampled_local_argmax_ptr
@@ -947,6 +1021,9 @@ def rejection_sample(
     synthetic_conditional_rates: torch.Tensor | None = None,
     use_fp64: bool = False,
     use_block_verification: bool = False,
+    contexts: torch.Tensor | None = None,
+    watermarking: torch.Tensor | None = None,
+    watermark_key: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert target_logits.ndim == 2 and target_logits.stride(-1) == 1
     assert draft_logits is None or (
@@ -954,6 +1031,30 @@ def rejection_sample(
     )
     num_reqs = cu_num_logits.shape[0] - 1
     num_logits, vocab_size = target_logits.shape
+
+    watermark = contexts is not None
+    assert watermark == (watermarking is not None) == (watermark_key is not None), (
+        "contexts, watermarking and watermark_key must be set together."
+    )
+    contexts_stride = 0
+    context_width = 1
+    watermarking_bytes: torch.Tensor | None = None
+    watermark_key_0 = 0
+    watermark_key_1 = 0
+    if contexts is not None:
+        assert watermarking is not None and watermark_key is not None
+        assert contexts.ndim == 2 and contexts.shape[0] == num_logits
+        # Context words are hashed as uint32, so int32 (the request-state token
+        # dtype) and int64 both land on the same PRF stream, -1 included.
+        assert contexts.dtype in (torch.int32, torch.int64)
+        if contexts.stride(-1) != 1:
+            contexts = contexts.contiguous()
+        assert watermarking.ndim == 1 and watermarking.dtype == torch.bool
+        contexts_stride = contexts.stride(0)
+        context_width = contexts.shape[-1]
+        watermarking_bytes = watermarking.view(torch.uint8)
+        watermark_key_0 = watermark_key & 0xFFFFFFFF
+        watermark_key_1 = watermark_key >> 32
     draft_logits_stride_0 = 0
     draft_logits_stride_1 = 0
     if has_draft_logits := draft_logits is not None:
@@ -1165,11 +1266,18 @@ def rejection_sample(
         seed,
         pos,
         cumulative_log_p,
+        contexts,
+        contexts_stride,
+        watermarking_bytes,
+        watermark_key_0,
+        watermark_key_1,
         vocab_size,
         BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
         HAS_DRAFT_LOGITS=has_draft_logits,
         USE_FP64=use_fp64,
         USE_BLOCK_VERIFICATION=use_block_verification,
+        CONTEXT_WIDTH=context_width,
+        WATERMARK=watermark,
     )
 
     # Insert the resampled tokens into the output sampled.
diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py
index 3c510e88850c..d1d6d888aabe 100644
--- a/vllm/v1/worker/gpu/spec_decode/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/speculator.py
@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
 from dataclasses import replace
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
@@ -15,20 +15,31 @@
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.models import supports_multimodal_embeddings
-from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.watermarking import create_watermarker
+from vllm.v1.watermarking.spec_decode import (
+    DraftWatermarker,
+    create_speculative_draft_watermarker,
+)
 from vllm.v1.worker.gpu.attn_utils import (
     build_attn_metadata,
     init_attn_backend,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.dp_utils import DPSyncState
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
+    OnlineAcceptanceEstimator,
+)
 from vllm.v1.worker.utils import AttentionGroup
 
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.pcp_manager import PCPManager
+
 logger = init_logger(__name__)
 
 
@@ -46,6 +57,8 @@ def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
 
 
 class BaseSpeculator(ABC):
+    num_query_per_req: int = 1
+
     @abstractmethod
     def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
         pass
@@ -146,9 +159,17 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
             dtype=torch.int64,
             device=device,
         )
-        self.arange = torch.arange(
-            self.max_num_reqs + 1, dtype=torch.int32, device="cpu"
+        self.arange_np = np.arange(self.max_num_reqs + 1, dtype=np.int32)
+        self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool)
+
+        self.draft_token_confidence_probs = torch.empty_like(
+            self.draft_tokens, dtype=torch.float32
+        )
+        self.enable_adaptive_verification = (
+            self.speculative_config.enable_adaptive_verification
         )
+        self.use_acceptance_estimator = self.enable_adaptive_verification
+        self.acceptance_estimator: OnlineAcceptanceEstimator | None = None
 
         self.draft_logits: torch.Tensor | None = None
         if self.speculative_config.draft_sample_method == "probabilistic":
@@ -165,13 +186,23 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
                 device=device,
             )
 
+        self.draft_watermarker: DraftWatermarker | None = None
+        watermark_config = getattr(vllm_config, "watermark_config", None)
+        if watermark_config is not None:
+            watermarker = create_watermarker(watermark_config)
+            self.draft_watermarker = create_speculative_draft_watermarker(
+                watermarker,
+                self.max_num_reqs,
+                device,
+                watermark_config.allow_target_only_watermarking,
+            )
+
         self.supports_mm_inputs = False
+        self.pcp_manager: PCPManager | None = None
 
     @abstractmethod
     def load_draft_model(
-        self,
-        target_model: nn.Module,
-        target_attn_layer_names: set[str],
+        self, target_model: nn.Module, target_attn_layer_names: set[str]
     ) -> nn.Module:
         pass
 
@@ -194,9 +225,7 @@ def load_model(self, target_model: nn.Module) -> None:
         )
         self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names
 
-        target_supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
-            self.vllm_config.model_config
-        )
+        target_supports_mm = self.vllm_config.model_config.supports_multimodal_inputs
         draft_supports_mm = supports_multimodal_embeddings(self.model)
         self.supports_mm_inputs = target_supports_mm and draft_supports_mm
         if target_supports_mm and not draft_supports_mm:
@@ -207,6 +236,20 @@ def load_model(self, target_model: nn.Module) -> None:
                 type(self.model).__name__,
             )
 
+        if self.use_acceptance_estimator:
+            if self.use_local_argmax_reduction:
+                raise ValueError(
+                    "Adaptive verification without a confidence head estimates "
+                    "per-position acceptance from the draft logits, which "
+                    "use_local_argmax_reduction never materializes. Disable one "
+                    "of them."
+                )
+            self.acceptance_estimator = OnlineAcceptanceEstimator(
+                self.max_num_reqs,
+                self.num_speculative_steps,
+                self.device,
+            )
+
     def set_eplb_state(self, eplb_state: EplbState) -> None:
         """Inject EPLB state after construction."""
         self.eplb_state = eplb_state
@@ -249,42 +292,37 @@ def set_attn(
         self.target_input_buffers = target_input_buffers
         self.target_attn_groups = target_attn_groups
 
-    def _build_draft_attn_metadata(
+    def _build_attn_metadata(
         self,
         num_reqs: int,
-        num_reqs_padded: int,
-        num_tokens_padded: int,
+        batch_desc: BatchExecutionDescriptor,
+        query_start_loc_np: np.ndarray,
         seq_lens_cpu_upper_bound: torch.Tensor,
         step: int,
-        num_query_per_req: int = 1,
         causal: bool | Mapping[int, bool] = True,
-        query_start_loc_np: np.ndarray | None = None,
         dcp_local_seq_lens: torch.Tensor | None = None,
     ) -> dict[str, Any] | None:
-        if query_start_loc_np is not None:
-            # Non-uniform query layout (e.g. multi-module MTP's mixed
-            # prefill/decode queries); num_query_per_req is ignored.
-            query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32)
-            query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy(
-                query_start_loc_np[: num_reqs + 1]
-            )
-            query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs]
-            max_query_len = int(
-                (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max()
-            )
-        else:
-            # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.
-            # Clamp keeps the series non-decreasing past num_reqs, which some
-            # attention backends require.
-            query_start_loc_cpu = (
-                torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs)
-                * num_query_per_req
-            )
-            max_query_len = num_query_per_req
+        num_reqs_padded = batch_desc.num_reqs or num_reqs
+        # A FULL graph replays a captured shape whose padded requests each hold
+        # a full query width, so attention must see the padded token count.
+        # PIECEWISE/eager needs the actual token count, because batch_desc may
+        # carry graph or DP padding that no request owns, which would desync it
+        # from query_start_loc.
+        num_tokens = (
+            batch_desc.num_tokens
+            if batch_desc.cg_mode == CUDAGraphMode.FULL
+            else int(query_start_loc_np[-1])
+        )
+        query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32)
+        query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy(
+            query_start_loc_np[: num_reqs + 1]
+        )
+        query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs]
+        max_query_len = int((query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max())
         block_tables = [
             x[:num_reqs_padded] for x in self.block_tables.input_block_tables
         ]
-        slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded]
+        slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]
         draft_seq_lens_cpu_upper_bound = torch.zeros(
             num_reqs_padded, dtype=torch.int32, device="cpu"
         )
@@ -297,7 +335,7 @@ def _build_draft_attn_metadata(
         if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
             # Draft steps advance and rewind their own global sequence lengths,
             # so the target model's DCP-local lengths may already be stale.
-            dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
+            dcp_local_seq_lens = prepare_dcp_local_seq_lens(
                 self.input_buffers.dcp_local_seq_lens,
                 self.input_buffers.seq_lens,
                 num_reqs,
@@ -308,7 +346,7 @@ def _build_draft_attn_metadata(
         attn_metadata = build_attn_metadata(
             attn_groups=self.attn_groups,
             num_reqs=num_reqs_padded,
-            num_tokens=num_tokens_padded,
+            num_tokens=num_tokens,
             query_start_loc_gpu=self.input_buffers.query_start_loc[
                 : num_reqs_padded + 1
             ],
@@ -326,6 +364,7 @@ def _build_draft_attn_metadata(
             kv_cache_config=self.kv_cache_config,
             causal=causal,
             seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound,
+            is_prefilling=self.draft_is_prefilling[:num_reqs_padded],
         )
         return attn_metadata
 
@@ -355,12 +394,6 @@ def _validate_local_argmax_reduction(self) -> None:
             "(communication: O(2*tp_size) vs O(vocab_size))."
         )
 
-    def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self.use_local_argmax_reduction:
-            return self.model.get_top_tokens(hidden_states)
-        logits = self.model.compute_logits(hidden_states)
-        return logits.argmax(dim=-1)
-
     def sample_draft(
         self,
         hidden_states: torch.Tensor,
@@ -373,7 +406,7 @@ def sample_draft(
     ) -> torch.Tensor:
         if draft_logits is not None:
             logits = self.model.compute_logits(hidden_states)
-            return gumbel_sample(
+            sampled = gumbel_sample(
                 logits,
                 idx_mapping,
                 temperature,
@@ -385,7 +418,50 @@ def sample_draft(
                 logits_cache_col=draft_step,
                 use_fp64=self.use_fp64_gumbel,
             )
-        return self._greedy_sample_draft(hidden_states)
+            if self.draft_watermarker is not None:
+                sampled = self.draft_watermarker.sample(
+                    logits, sampled, idx_mapping, temperature
+                )
+        elif self.use_local_argmax_reduction:
+            return self.model.get_top_tokens(hidden_states)
+        else:
+            logits = self.model.compute_logits(hidden_states)
+            sampled = logits.argmax(dim=-1)
+        self._maybe_predict_acceptance(logits, idx_mapping, draft_step)
+        return sampled
+
+    def _maybe_predict_acceptance(
+        self, logits: torch.Tensor, idx_mapping: torch.Tensor, draft_step: torch.Tensor
+    ) -> None:
+        if self.acceptance_estimator is not None:
+            self.acceptance_estimator.predict(
+                logits,
+                idx_mapping,
+                draft_step,
+                self.draft_token_confidence_probs,
+                self.temperature,
+            )
+
+    def observe_verification(
+        self,
+        idx_mapping: torch.Tensor,
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+    ) -> None:
+        """Fold the target's verdict on the last drafts into the estimator.
+
+        Must run before the next `propose`, which overwrites the per-slot
+        features the verdict grades.
+        """
+        if self.acceptance_estimator is not None:
+            self.acceptance_estimator.step(idx_mapping, num_sampled, num_rejected)
+
+    def prepare_watermarking(
+        self, contexts: torch.Tensor, watermarking: torch.Tensor
+    ) -> None:
+        if self.draft_watermarker is None:
+            return
+        self.draft_watermarker.prepare(contexts, watermarking)
 
     def _copy_request_inputs(
         self,
@@ -396,6 +472,7 @@ def _copy_request_inputs(
         temperature: torch.Tensor,
         # [max_num_reqs]
         seeds: torch.Tensor,
+        dummy_run: bool = False,
     ) -> None:
         # Copy temperature, seeds, and idx mapping to the pre-allocated buffers.
         # NOTE(woosuk): For draft sampling, we only consider the temperature
@@ -405,10 +482,15 @@ def _copy_request_inputs(
         # affect the output distribution after rejection sampling.
         self.temperature.copy_(temperature)
         self.seeds.copy_(seeds)
-        self.idx_mapping[:num_reqs].copy_(idx_mapping)
-        # idx_mapping for CG padded requests points to -1, which is ignored
-        # during sampling to prevent writing stale values to draft logits.
-        self.idx_mapping[num_reqs:].fill_(-1)
+        # idx_mapping == -1 marks a row the drafter must not act on: sampling
+        # skips it and compute_slot_mappings emits PAD. CUDA-graph padded rows
+        # always get it; a dummy batch gets it for every row, since its arange
+        # idx_mapping names request-state slots with stale block tables.
+        if dummy_run:
+            self.idx_mapping.fill_(-1)
+        else:
+            self.idx_mapping[:num_reqs].copy_(idx_mapping)
+            self.idx_mapping[num_reqs:].fill_(-1)
 
     def _build_uniform_batch_dp_sync(
         self,
@@ -428,3 +510,24 @@ def _build_uniform_batch_dp_sync(
             uniform_token_count=num_query_per_req,
             eager=False,
         ), num_batch_tokens
+
+    def _build_uniform_attn_metadata(
+        self,
+        batch_desc: BatchExecutionDescriptor,
+        num_reqs: int,
+        num_query_per_req: int,
+        seq_lens_cpu_upper_bound: torch.Tensor,
+        step: int,
+        causal: bool | Mapping[int, bool] = True,
+        dcp_local_seq_lens: torch.Tensor | None = None,
+    ) -> dict[str, Any] | None:
+        query_start_loc_np = self.arange_np[: num_reqs + 1] * num_query_per_req
+        return self._build_attn_metadata(
+            num_reqs=num_reqs,
+            batch_desc=batch_desc,
+            query_start_loc_np=query_start_loc_np,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            step=step,
+            causal=causal,
+            dcp_local_seq_lens=dcp_local_seq_lens,
+        )
diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py
index e25672e953bd..f73acb2bdaed 100644
--- a/vllm/v1/worker/gpu/spec_decode/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/utils.py
@@ -3,10 +3,28 @@
 import numpy as np
 import torch
 
+from vllm.config import LoadConfig, replace
+from vllm.distributed import get_pp_group
+from vllm.logger import init_logger
 from vllm.v1.outputs import DraftTokenIds
 from vllm.v1.worker.gpu.async_utils import async_copy_to_np
 from vllm.v1.worker.gpu.input_batch import InputBatch
 
+logger = init_logger(__name__)
+
+
+def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig:
+    """Avoid collectives that include PP ranks without a draft model."""
+    if get_pp_group().world_size > 1 and load_config.load_format == "fastsafetensors":
+        logger.warning_once(
+            "fastsafetensors cannot load a draft model instantiated on only "
+            "one pipeline stage; falling back to the standard safetensors "
+            "loader for the draft model. The target model still uses "
+            "fastsafetensors."
+        )
+        return replace(load_config, load_format="auto")
+    return load_config
+
 
 class DraftTokensHandler:
     def __init__(self, device: torch.device | None = None):
diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py
index 025c8300fcd2..db27bf013f73 100644
--- a/vllm/v1/worker/gpu/states.py
+++ b/vllm/v1/worker/gpu/states.py
@@ -78,10 +78,7 @@ def __init__(
         )
 
         self.next_prefill_tokens = torch.zeros(
-            num_prefill_lookahead,
-            self.max_num_reqs,
-            dtype=torch.int32,
-            device=device,
+            num_prefill_lookahead, self.max_num_reqs, dtype=torch.int32, device=device
         )
 
     @property
diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py
index 8437b0aed9a4..057d6f7154c7 100644
--- a/vllm/v1/worker/gpu/structured_outputs.py
+++ b/vllm/v1/worker/gpu/structured_outputs.py
@@ -5,8 +5,7 @@
 
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
-from vllm.utils.torch_utils import PIN_MEMORY
-from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
 from vllm.v1.worker.gpu.input_batch import InputBatch
 
 
@@ -68,7 +67,7 @@ def apply_grammar_bitmask(
 
         # Asynchronously copy the bitmask to GPU.
         with torch.cuda.stream(self.copy_stream):
-            bitmask = async_copy_to_gpu(
+            bitmask = async_tensor_h2d(
                 grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]]
             )
 
diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py
index 012ce6c471a5..d6f24fcc695c 100644
--- a/vllm/v1/worker/gpu/ubatch_utils.py
+++ b/vllm/v1/worker/gpu/ubatch_utils.py
@@ -3,6 +3,8 @@
 """Microbatching (DBO) helpers for the V2 GPU model runner."""
 
 import threading
+from collections.abc import Callable
+from contextlib import ExitStack
 from dataclasses import replace
 from typing import Any, NamedTuple
 
@@ -22,7 +24,7 @@
 from vllm.utils.torch_utils import current_stream
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.ubatch_utils import (
@@ -148,7 +150,7 @@ def _slice_input_batch(
     dcp_local_seq_lens = None
     if dcp_size > 1:
         assert dcp_local_seq_lens_buf is not None
-        dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
+        dcp_local_seq_lens = prepare_dcp_local_seq_lens(
             dcp_local_seq_lens_buf,
             seq_lens,
             num_reqs,
@@ -187,6 +189,7 @@ def _slice_input_batch(
             if input_batch.prompt_lens is None
             else input_batch.prompt_lens[req_start:req_stop]
         ),
+        fast_prefill=None,
     )
 
 
@@ -316,8 +319,11 @@ def __init__(
             for _ in range(self.num_ubatches)
         ]
         self.comm_stream = torch.cuda.Stream(device=device)
+        # Threads and graph capture must use the same compute stream.
+        self.capture_stream = torch.cuda.Stream(device=device)
         # The microbatch threads plus the thread that starts them.
         self.ready_barrier = threading.Barrier(self.num_ubatches + 1)
+        self._pending_finish: Callable[[], Any] | None = None
         self.sm_control = create_sm_control_context(self.parallel_config)
 
     def prepare(
@@ -325,11 +331,16 @@ def prepare(
         input_batch: InputBatch,
         block_tables: tuple[torch.Tensor, ...],
         slot_mappings: torch.Tensor,
+        cg_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+        for_capture: bool = False,
     ) -> UBatchState:
         """Split the batch into the microbatches the step will run on.
 
         Attention metadata is built per microbatch and carried in the forward
         contexts the threads install, not in the caller's.
+
+        FULL uses padded sizes; `for_capture` refreshes metadata from dummy
+        block tables.
         """
         ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches)
 
@@ -354,12 +365,13 @@ def prepare(
             attn_metadata.append(
                 self.model_state.prepare_attn(
                     ubatch,
-                    CUDAGraphMode.NONE,
+                    cg_mode,
                     ubatch_block_tables,
                     ubatch_slot_mappings,
                     self.attn_groups,
                     self.kv_cache_config,
                     ubatch_idx=i,
+                    for_capture=for_capture,
                 )
             )
             slot_mappings_by_layer.append(
@@ -414,14 +426,34 @@ def run(
         model_inputs: dict[str, Any],
         ubatch_state: UBatchState,
     ) -> Any:
+        finish = self.begin_capturable_run(model, model_inputs, ubatch_state)
+        return finish()
+
+    def begin_capturable_run(
+        self,
+        model: Any,
+        model_inputs: dict[str, Any],
+        ubatch_state: UBatchState,
+        for_capture: bool = False,
+    ) -> Callable[[], Any]:
+        """Start microbatch threads and return a callback to run and join them.
+
+        For capture, start threads outside the graph and call the callback inside.
+        `for_capture` selects the capture stream and initializes cuBLAS workspaces
+        before capture.
+
+        Call the callback exactly once on this thread, or call `abort_pending_run`
+        on failure to release parked threads.
+        """
         ubatch_slices = ubatch_state.slices
         assert len(ubatch_slices) == len(ubatch_state.forward_contexts)
         assert len(ubatch_slices) == self.num_ubatches
 
+        compute_stream = self.capture_stream if for_capture else current_stream()
         ubatch_contexts = make_ubatch_contexts(
             num_micro_batches=self.num_ubatches,
             comm_stream=self.comm_stream,
-            compute_stream=current_stream(),
+            compute_stream=compute_stream,
             forward_contexts=ubatch_state.forward_contexts,
             ready_barrier=self.ready_barrier,
         )
@@ -432,6 +464,13 @@ def run(
         @torch.inference_mode()
         def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None:
             try:
+                # New threads must select the worker's device.
+                torch.accelerator.set_device_index(self.device.index)
+                if for_capture:
+                    # Allocate cuBLAS workspaces before capture.
+                    for stream in (compute_stream, self.comm_stream):
+                        with torch.cuda.stream(stream):
+                            torch.cuda.current_blas_handle()
                 with ubatch_context:
                     outputs[ubatch_context.id] = model(**inputs)
             except BaseException as e:  # noqa: BLE001
@@ -444,33 +483,56 @@ def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None:
                 # is deliberately out of scope here.
                 errors[ubatch_context.id] = e
 
-        # The threads manage the forward context themselves; clear it here so
-        # it is restored correctly once they are done.
-        with override_forward_context(None), self.sm_control:
-            threads = []
-            for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices):
-                thread = threading.Thread(
-                    target=run_ubatch,
-                    args=(
-                        ubatch_context,
-                        slice_model_inputs(model_inputs, ubatch_slice.token_slice),
-                    ),
-                )
-                threads.append(thread)
-                thread.start()
+        # Threads own the forward context; keep it cleared and SMs reserved
+        # until finish joins them.
+        stack = ExitStack()
+        stack.enter_context(override_forward_context(None))
+        stack.enter_context(self.sm_control)
+
+        threads = []
+        for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices):
+            thread = threading.Thread(
+                target=run_ubatch,
+                args=(
+                    ubatch_context,
+                    slice_model_inputs(model_inputs, ubatch_slice.token_slice),
+                ),
+            )
+            threads.append(thread)
+            thread.start()
 
-            # Wait for every thread to reach its context, then start the first.
-            self.ready_barrier.wait()
-            ubatch_contexts[0].cpu_wait_event.set()
-            for thread in threads:
-                thread.join()
+        # Park all threads before the caller starts capture.
+        self.ready_barrier.wait()
 
-        if errors:
-            failed = min(errors)
-            raise RuntimeError(
-                f"Microbatch {failed} of {self.num_ubatches} failed"
-            ) from errors[failed]
-        return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)])
+        def finish() -> Any:
+            try:
+                ubatch_contexts[0].cpu_wait_event.set()
+                for thread in threads:
+                    thread.join()
+            finally:
+                self._pending_finish = None
+                stack.close()
+
+            if errors:
+                failed = min(errors)
+                raise RuntimeError(
+                    f"Microbatch {failed} of {self.num_ubatches} failed"
+                ) from errors[failed]
+            return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)])
+
+        self._pending_finish = finish
+        return finish
+
+    def abort_pending_run(self) -> None:
+        """Release and join parked threads after a failed capture."""
+        if self._pending_finish is None:
+            return
+        try:
+            self._pending_finish()
+        except Exception:
+            logger.exception("Failed to join microbatch threads after an abort")
+        finally:
+            self._pending_finish = None
 
 
 def maybe_build_ubatch_runner(
@@ -490,10 +552,7 @@ def maybe_build_ubatch_runner(
     if not parallel_config.use_ubatching or parallel_config.data_parallel_size <= 1:
         return None
 
-    logger.info_once(
-        "Dual batch overlap is enabled. Microbatched steps run without "
-        "CUDA graphs on the V2 model runner."
-    )
+    logger.info_once("Dual batch overlap is enabled.")
     return UBatchRunner(
         vllm_config, device, model_state, attn_groups, kv_cache_config, max_num_reqs
     )
diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py
index 7072210cb79e..56f7d438fc3f 100644
--- a/vllm/v1/worker/gpu/warmup.py
+++ b/vllm/v1/worker/gpu/warmup.py
@@ -224,10 +224,20 @@ def warmup_kernels(
     # fixed draft counts here, then restore the manager for capture and serving.
     adaptive_verification = model_runner.adaptive_verification
     model_runner.adaptive_verification = None
+    rejection_sampler = model_runner.rejection_sampler
+    adaptive_sampling = (
+        rejection_sampler is not None and rejection_sampler.enable_adaptive_verification
+    )
+    if adaptive_sampling:
+        assert rejection_sampler is not None
+        rejection_sampler.enable_adaptive_verification = False
     try:
         _warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens)
     finally:
         model_runner.adaptive_verification = adaptive_verification
+        if adaptive_sampling:
+            assert rejection_sampler is not None
+            rejection_sampler.enable_adaptive_verification = True
 
 
 def _warmup_kernels(
diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py
index 01b6b5b0cc90..79db440cff5b 100644
--- a/vllm/v1/worker/gpu_input_batch.py
+++ b/vllm/v1/worker/gpu_input_batch.py
@@ -323,7 +323,6 @@ def _register_add_request(self, request: "CachedRequestState") -> int:
         """Track add-request operations for logits processors.
         Not applicable to pooling models.
         """
-
         # Fill the next empty index if there is one.
         if (new_req_index := self.batch_update_builder.pop_removed()) is None:
             # Append to end otherwise.
@@ -531,8 +530,8 @@ def remove_request(self, req_id: str) -> int | None:
 
         Returns:
           Removed request index, or `None` if `req_id` not recognized
-        """
 
+        """
         req_index = self.req_id_to_index.pop(req_id, None)
         if req_index is None:
             return None
@@ -710,6 +709,7 @@ def condense(self) -> None:
         Returns:
           swaps: list of (from,to) swap tuples for moved requests
           empty_req_indices: indices not filled by condensation
+
         """
         num_reqs = self.num_reqs
 
@@ -835,7 +835,6 @@ def condense(self) -> None:
 
     def refresh_metadata(self):
         """Apply any batch updates to sampling metadata."""
-
         if self.is_pooling_model:
             batch_changed = self.batch_update_builder.reset()
             if batch_changed:
@@ -1001,18 +1000,19 @@ def _make_prompt_token_ids_cpu_tensor(self) -> torch.Tensor:
     def make_lora_inputs(
         self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray
     ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
-        """
-        Given the num_scheduled_tokens for each request in the batch, return
+        """Given the num_scheduled_tokens for each request in the batch, return
         datastructures used to activate the current LoRAs.
+
         Returns:
             1. prompt_lora_mapping: A tuple of size np.sum(num_sampled_tokens)
-               where, prompt_lora_mapping[i] is the LoRA id to use for the ith
-               sampled token.
+                where, prompt_lora_mapping[i] is the LoRA id to use for the ith
+                sampled token.
             2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens)
-               where, token_lora_mapping[i] is the LoRA id to use for ith token.
+                where, token_lora_mapping[i] is the LoRA id to use for ith
+                token.
             3. lora_requests: Set of relevant LoRA requests.
-        """
 
+        """
         req_lora_mapping = self.request_lora_mapping[: self.num_reqs]
         prompt_lora_mapping = tuple(req_lora_mapping.repeat(num_sampled_tokens))
         token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens))
@@ -1028,8 +1028,7 @@ def set_async_sampled_token_ids(
         sampled_token_ids_cpu: torch.Tensor,
         async_copy_ready_event: torch.Event,
     ) -> None:
-        """
-        In async scheduling case, store ref to sampled_token_ids_cpu
+        """In async scheduling case, store ref to sampled_token_ids_cpu
         tensor and corresponding copy-ready event. Used to repair
         output_token_ids prior to sampling, if needed by logits processors.
         """
@@ -1041,8 +1040,7 @@ def set_async_sampled_token_ids(
             self.async_copy_ready_event = None
 
     def update_async_output_token_ids(self) -> None:
-        """
-        In async scheduling case, update output_token_ids in sampling metadata
+        """In async scheduling case, update output_token_ids in sampling metadata
         from prior steps sampled token ids once they've finished copying to CPU.
         This is called right before they are needed by the logits processors.
         """
@@ -1088,8 +1086,7 @@ def update_async_output_token_ids(self) -> None:
             # ^ Implicitly resizes to (first_placeholder + num_to_replace)
 
     def update_async_spec_token_ids(self, draft_token_ids: list[list[int]]) -> None:
-        """
-        In async scheduling case, update spec_token_ids in sampling metadata with
+        """In async scheduling case, update spec_token_ids in sampling metadata with
         real draft token ids from prior step. This is called right before they are
         needed by the rejection sampler for penalty/bad_words computation.
         """
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index d6b8ed025c82..00e33b47c609 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -46,7 +46,6 @@
     copy_kv_blocks,
 )
 from vllm.distributed.parallel_state import (
-    GraphCaptureContext,
     get_dcp_group,
     get_pp_group,
     get_tp_group,
@@ -62,10 +61,6 @@
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    RoutedExpertsCapturer,
-    bind_routed_experts_capturer,
-)
 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncsByType
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     initialize_mamba_ssu_backend,
@@ -128,7 +123,6 @@
 from vllm.utils.torch_utils import (
     PIN_MEMORY,
     async_tensor_h2d,
-    current_stream,
     kv_cache_dtype_str_to_dtype,
 )
 from vllm.v1.attention.backend import (
@@ -178,8 +172,6 @@
     LogprobsTensors,
     ModelRunnerOutput,
     PoolerOutput,
-    RoutedExpertsLists,
-    RoutedExpertsTensors,
     SamplerOutput,
     make_empty_encoder_model_runner_output,
 )
@@ -295,7 +287,6 @@ def __init__(
         invalid_req_indices: list[int],
         async_output_copy_stream: torch.cuda.Stream,
         vocab_size: int,
-        routed_experts: RoutedExpertsTensors | None = None,
         check_ep_fault: bool = False,
         num_nans: torch.Tensor | None = None,
     ):
@@ -311,7 +302,6 @@ def __init__(
         self._sampled_token_ids = sampled_token_ids
         self.vocab_size = vocab_size
         self._logprobs_tensors = logprobs_tensors
-        self._routed_experts = routed_experts
         self._num_nans = num_nans
         self._has_fault: torch.Tensor | None = None
 
@@ -327,11 +317,6 @@ def __init__(
                 if self._logprobs_tensors
                 else None
             )
-            self._routed_experts_cpu = (
-                self._routed_experts.to_cpu_nonblocking()
-                if self._routed_experts is not None
-                else None
-            )
             self._num_nans_cpu = (
                 self._num_nans.to("cpu", non_blocking=True)
                 if self._num_nans is not None
@@ -372,10 +357,6 @@ def get_output(self) -> ModelRunnerOutput:
         output.sampled_token_ids = valid_sampled_token_ids
         output.logprobs = logprobs_lists
 
-        if self._routed_experts_cpu is not None:
-            output.routed_experts = self._routed_experts_cpu.tolists()
-        del self._routed_experts
-
         if self._num_nans_cpu is not None:
             output.num_nans_in_logits = nans_to_dict(
                 self._num_nans_cpu.tolist(), output.req_id_to_index
@@ -498,6 +479,7 @@ class ExecuteModelState(NamedTuple):
 class GPUModelRunner(
     LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin
 ):
+    @JitWarmupRegistry.capture
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -514,7 +496,6 @@ def __init__(
         self.scheduler_config = vllm_config.scheduler_config
         self.speculative_config = vllm_config.speculative_config
         self.observability_config = vllm_config.observability_config
-        self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
 
         model_config = self.model_config
         cache_config = self.cache_config
@@ -539,9 +520,6 @@ def __init__(
         # These will be overridden in load_model()
         self.is_multimodal_pruning_enabled = False
         self.requires_sequential_video_encoding = False
-        # Set to True after init_routed_experts_capturer() completes.
-        # Prevents routed experts code from running during profiling/dummy run.
-        self.routed_experts_initialized = False
         self.max_model_len = model_config.max_model_len
 
         # Always set to false after the first forward pass
@@ -572,9 +550,7 @@ def __init__(
         self.mm_registry = MULTIMODAL_REGISTRY
         self.uses_mrope = model_config.uses_mrope
         self.mrope_num_dims = model_config.mrope_num_dims
-        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
-            model_config
-        )
+        self.supports_mm_inputs = model_config.supports_multimodal_inputs
 
         if self.model_config.is_encoder_decoder:
             # Maximum length of the encoder input, only for encoder-decoder
@@ -750,36 +726,35 @@ def __init__(
             self.parallel_config.cp_kv_cache_interleave_size
         )
         # Capture warmup providers registered by the initial placeholder InputBatch
-        with self.jit_warmup_registry.activate():
-            self.input_batch = InputBatch(
-                max_num_reqs=self.max_num_reqs,
-                # We need to use the encoder length for encoder-decoder
-                # because of KV cache for cross-attention.
-                max_model_len=max(self.max_model_len, self.max_encoder_len),
-                max_num_batched_tokens=self.max_num_tokens,
-                device=self.device,
-                vocab_size=self.model_config.get_vocab_size(),
-                block_sizes=[placeholder_block_size],
-                kernel_block_sizes=[placeholder_block_size],
-                max_num_blocks_per_req=[placeholder_max_num_blocks],
-                num_spec_tokens=self.num_spec_tokens,
-                logitsprocs=build_logitsprocs(
-                    self.vllm_config,
-                    self.device,
-                    PIN_MEMORY,
-                    self.is_pooling_model,
-                    custom_logitsprocs,
-                ),
-                # We currently don't know whether a particular custom logits processor
-                # uses output token ids so we set this conservatively. Thinking-budget
-                # tracking is requested dynamically when a budgeted request is in the
-                # batch.
-                logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
-                is_pooling_model=self.is_pooling_model,
-                cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
-                reasoning_config=self.vllm_config.reasoning_config,
-                use_replayssm=self.cache_config.use_replayssm,
-            )
+        self.input_batch = InputBatch(
+            max_num_reqs=self.max_num_reqs,
+            # We need to use the encoder length for encoder-decoder
+            # because of KV cache for cross-attention.
+            max_model_len=max(self.max_model_len, self.max_encoder_len),
+            max_num_batched_tokens=self.max_num_tokens,
+            device=self.device,
+            vocab_size=self.model_config.get_vocab_size(),
+            block_sizes=[placeholder_block_size],
+            kernel_block_sizes=[placeholder_block_size],
+            max_num_blocks_per_req=[placeholder_max_num_blocks],
+            num_spec_tokens=self.num_spec_tokens,
+            logitsprocs=build_logitsprocs(
+                self.vllm_config,
+                self.device,
+                PIN_MEMORY,
+                self.is_pooling_model,
+                custom_logitsprocs,
+            ),
+            # We currently don't know whether a particular custom logits processor
+            # uses output token ids so we set this conservatively. Thinking-budget
+            # tracking is requested dynamically when a budgeted request is in the
+            # batch.
+            logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
+            is_pooling_model=self.is_pooling_model,
+            cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
+            reasoning_config=self.vllm_config.reasoning_config,
+            use_replayssm=self.cache_config.use_replayssm,
+        )
 
         # Separate cuda stream for overlapping transfer of sampled token ids from
         # GPU to CPU when async scheduling is enabled.
@@ -813,9 +788,15 @@ def __init__(
 
         # Persistent buffers for CUDA graphs.
         self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
+        # [max_num_reqs, depth] prompt token ids preceding each request's chunk
+        # start, for models exposing `token_lookback_depth`. Sized after load.
+        self.lookback_token_ids: CpuGpuBuffer | None = None
         self.positions = torch.zeros(
             self.max_num_tokens, dtype=torch.int64, device=self.device
         )
+        self.is_padding = torch.zeros(
+            self.max_num_tokens, dtype=torch.bool, device=self.device
+        )
         self.query_start_loc = self._make_buffer(
             self.max_num_reqs + 1, dtype=torch.int32
         )
@@ -1001,8 +982,7 @@ def update_max_model_len(self, max_model_len: int) -> None:
                 self.effective_drafter_max_model_len = self.max_model_len
 
     def reset_mm_cache(self) -> None:
-        """
-        Clear the multi-modal cache that was used during profiling,
+        """Clear the multi-modal cache that was used during profiling,
         but no longer needed during inference.
         """
         if self.mm_budget:
@@ -1065,9 +1045,36 @@ def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers:
             )
         return self._mamba_bufs
 
-    def _init_model_kwargs(self):
+    def _prepare_lookback_token_ids(self, num_reqs: int) -> torch.Tensor:
+        """Gather, per request, the `depth` prompt token ids preceding its
+        first scheduled token (column j is position start - 1 - j); -1 where
+        the position is before the prompt or already past it. Generated
+        positions are left to the model: under async scheduling the CPU token
+        table holds placeholders for them."""
+        buf = self.lookback_token_ids
+        assert buf is not None
+        buf.np.fill(-1)
+        if num_reqs > 0:
+            depth = buf.np.shape[1]
+            starts = self.input_batch.num_computed_tokens_cpu[:num_reqs, None]
+            pos = starts - np.arange(1, depth + 1)
+            num_prompt = self.input_batch.num_prompt_tokens[:num_reqs, None]
+            valid = (pos >= 0) & (pos < num_prompt)
+            rows = np.arange(num_reqs)[:, None]
+            ids = self.input_batch.token_ids_cpu[rows, np.clip(pos, 0, None)]
+            buf.np[:num_reqs] = np.where(valid, ids, -1)
+        return buf.copy_to_gpu()
+
+    def _init_model_kwargs(self, num_reqs: int | None = None):
         model_kwargs = dict[str, Any]()
 
+        if self.lookback_token_ids is not None:
+            if num_reqs is None:
+                num_reqs = self.input_batch.num_reqs
+            model_kwargs["lookback_token_ids"] = self._prepare_lookback_token_ids(
+                num_reqs
+            )
+
         if not self.is_pooling_model:
             return model_kwargs
 
@@ -1107,14 +1114,14 @@ def _init_model_kwargs(self):
         return model_kwargs
 
     def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
-        """
-        Update the order of requests in the batch based on the attention
+        """Update the order of requests in the batch based on the attention
         backend's needs. For example, some attention backends (namely MLA) may
         want to separate requests based on if the attention computation will be
         compute-bound or memory-bound.
 
         Args:
             scheduler_output: The scheduler output.
+
         """
         # Attention free models have zero kv_cache_groups, however models
         # like Mamba are also attention free but use the kv_cache for
@@ -1153,7 +1160,6 @@ def _zero_block_ids(self, block_ids: list[int]) -> None:
     # Note: used for model runner override.
     def _init_device_properties(self) -> None:
         """Initialize attributes from torch.cuda.get_device_properties"""
-
         self.num_sms = num_compute_units(self.device.index)
 
     # Note: used for model runner override.
@@ -1774,7 +1780,6 @@ def _prepare_input_ids(
         Uses self.prev_positions[:num_reqs] which maps current pos -> prev pos
         (-1 for new requests).
         """
-
         # Sync the async PP broadcast before reading sampled tokens.
         if self._pp_recv_work is not None:
             self._pp_recv_work.wait()
@@ -1834,8 +1839,12 @@ def _prepare_input_ids(
         total_without_spec = total_num_scheduled_tokens - total_num_spec_tokens
         if self.enable_prompt_embeds:
             # The multimodal embed path reads is_token_ids.gpu; its .cpu copy is
-            # refreshed every step but the async fast paths below only scatter
-            # input_ids.gpu, so refresh is_token_ids.gpu here too.
+            # refreshed every step but the async fast path below replaces draft
+            # tokens directly on the GPU. Mark those positions as token IDs before
+            # refreshing the GPU mask so their embeddings are rebuilt. Sampled
+            # tokens are already marked by bookkeeping.
+            if spec_flattened_indices:
+                self.is_token_ids.np[spec_flattened_indices] = True
             self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
         if num_common_tokens < total_without_spec:
             # If not all requests are decodes from the last iteration,
@@ -1948,9 +1957,9 @@ def _prepare_inputs(
         SpecDecodeMetadata | None,
         int,
     ]:
-        """
-        Returns:
-            tuple[logits_indices, spec_decode_metadata, max_num_sampled_tokens]
+        """Returns:
+        tuple[logits_indices, spec_decode_metadata, max_num_sampled_tokens]
+
         """
         total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         assert total_num_scheduled_tokens > 0
@@ -2281,9 +2290,9 @@ def _build_attention_metadata(
         cascade_attn_prefix_lens: list[list[int]] | None = None,
         slot_mappings: dict[int, torch.Tensor] | None = None,
     ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]:
-        """
-        Returns:
-            tuple[attn_metadata, spec_decode_common_attn_metadata]
+        """Returns:
+        tuple[attn_metadata, spec_decode_common_attn_metadata]
+
         """
         # Attention metadata is not needed for attention free models
         if len(self.kv_cache_config.kv_cache_groups) == 0:
@@ -2329,26 +2338,13 @@ def _get_block_table(kv_cache_gid: int):
         block_table_gid_0 = _get_block_table(0)
         slot_mapping_gid_0 = slot_mappings[0]
 
-        if self.routed_experts_initialized:
-            # Copy this step's attention slot_mapping into our private
-            # device buffer. The shared ``slot_mappings[attn_gid]`` is
-            # owned by the attention block table and will be overwritten
-            # by the next ``_prepare_inputs``; we need a stable snapshot
-            # because the async D2H may still be in flight on the copy
-            # stream when the next step runs.
-            slot_mapping_attn = slot_mappings[self.routed_experts_capturer.attn_gid]
-            self.routed_experts_slot_mapping_device[:num_tokens].copy_(
-                slot_mapping_attn[:num_tokens]
-            )
-
         num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[
             :num_reqs_padded
         ]
         num_prompt_tokens_cpu = self.input_batch.num_prompt_tokens_cpu_tensor[
             :num_reqs_padded
         ]
-        seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs_padded]
-        seq_lens_cpu_upper_bound = seq_lens_cpu
+        seq_lens_cpu_upper_bound = self.optimistic_seq_lens_cpu[:num_reqs_padded]
 
         # is_prefilling: True if request is still in prefill phase.
         # Used by mamba backends to distinguish actual decodes from
@@ -2358,11 +2354,6 @@ def _get_block_table(kv_cache_gid: int):
         # misclassify padding as prefill in CUDA graph mode.
         is_prefilling[num_reqs:] = False
 
-        if self.use_async_spec_decode:
-            # GPU tensors are authoritative in async mode.
-            seq_lens_cpu = None
-            num_computed_tokens_cpu = None
-
         # Compute mm_prefix bidirectional ranges before building
         # attention metadata so builders handle them during build().
         # By default, ranges exceeding sliding_window are skipped to prevent
@@ -2436,8 +2427,6 @@ def _get_block_table(kv_cache_gid: int):
             query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],
             query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],
             seq_lens=self.seq_lens[:num_reqs_padded],
-            _seq_lens_cpu=seq_lens_cpu,
-            _num_computed_tokens_cpu=num_computed_tokens_cpu,
             seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
             replayssm_decode_base_cpu=replayssm_decode_base_cpu,
             num_reqs=num_reqs_padded,
@@ -2464,7 +2453,7 @@ def _get_block_table(kv_cache_gid: int):
             self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded)
 
             cm_base.dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded]
-            cm_base.dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[
+            cm_base.dcp_local_seq_lens_cpu_upper_bound = self.dcp_local_seq_lens.cpu[
                 :num_reqs_padded
             ]
 
@@ -2630,14 +2619,13 @@ def _compute_cascade_attn_prefix_lens(
         num_computed_tokens: np.ndarray,
         num_common_prefix_blocks: list[int],
     ) -> list[list[int]] | None:
-        """
-        Returns:
-            Optional[cascade_attn_prefix_lens]
-                cascade_attn_prefix_lens is 2D:
-                ``[kv_cache_group_id][attn_group_idx]``,
-                None if we should not use cascade attention
-        """
+        """Returns:
+        Optional[cascade_attn_prefix_lens]
+            cascade_attn_prefix_lens is 2D:
+            ``[kv_cache_group_id][attn_group_idx]``,
+            None if we should not use cascade attention
 
+        """
         use_cascade_attn = False
         num_kv_cache_groups = len(self.kv_cache_config.kv_cache_groups)
         cascade_attn_prefix_lens: list[list[int]] = [
@@ -2682,12 +2670,16 @@ def _compute_cascade_attn_prefix_len(
 
         Args:
             num_scheduled_tokens: Number of tokens scheduled per request.
+            num_computed_tokens: Number of tokens already computed per request.
             num_common_prefix_blocks: Number of shared KV cache blocks.
+            kv_cache_spec: The KV cache spec of the group being considered.
+            attn_metadata_builder: Builder used to decide whether cascade
+                attention can be used.
 
         Returns:
             int: Length of common prefix in tokens.
-        """
 
+        """
         common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size
         if common_prefix_len == 0:
             # Common case.
@@ -2926,6 +2918,7 @@ def _batch_mm_inputs_from_scheduler(
             - mm_hashes: List of multimodal hashes for each item
             - mm_kwargs: List of multimodal kwargs for each item
             - mm_lora_refs: List of (req_id, placeholder_range) for each item
+
         """
         scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
         if not scheduled_encoder_inputs:
@@ -3383,9 +3376,7 @@ def sync_and_gather_intermediate_tensors(
         )
 
     def eplb_step(self, is_dummy: bool = False, is_profile: bool = False) -> None:
-        """
-        Step for the EPLB (Expert Parallelism Load Balancing) state.
-        """
+        """Step for the EPLB (Expert Parallelism Load Balancing) state."""
         if not self.parallel_config.enable_eplb or self.eep_eplb_suppressed:
             return
 
@@ -3478,6 +3469,14 @@ def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int:
             return round_up(num_scheduled_tokens, tp_size)
         return num_scheduled_tokens
 
+    def _prepare_padding_mask(
+        self, num_tokens_unpadded: int, num_tokens_padded: int
+    ) -> torch.Tensor:
+        padding_mask = self.is_padding[:num_tokens_padded]
+        padding_mask[:num_tokens_unpadded].fill_(False)
+        padding_mask[num_tokens_unpadded:].fill_(True)
+        return padding_mask
+
     def _prepare_mm_inputs(
         self, num_tokens: int
     ) -> tuple[torch.Tensor | None, torch.Tensor]:
@@ -3712,20 +3711,6 @@ def _bookkeeping_sync(
         invalid_req_indices = []
         logprobs_lists = None
         if not self.use_async_scheduling:
-            # Sync scheduling: issue routed experts D2H into the pinned
-            # CPU buffer BEFORE ``_to_list`` below. ``_to_list`` does
-            # ``event.synchronize()`` on the async copy stream which
-            # waits for every D2H queued on the default stream since
-            # the last sync, so this enqueue is naturally covered
-            # without requiring its own synchronize.
-            if self.routed_experts_initialized:
-                buf = self.routed_experts_capturer.get_device_buffer()
-                total = scheduler_output.total_num_scheduled_tokens
-                self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)
-                self.routed_experts_slot_mapping_cpu[:total].copy_(
-                    self.routed_experts_slot_mapping_device[:total],
-                    non_blocking=True,
-                )
             with gpu_sync_allowed():
                 # Get the valid generated tokens.
                 max_gen_len = sampled_token_ids.shape[-1]
@@ -3854,6 +3839,7 @@ def _model_forward(
 
         Returns:
             Model output tensor
+
         """
         return self.model(
             input_ids=input_ids,
@@ -3871,8 +3857,7 @@ def _is_uniform_decode(
         num_reqs: int,
         force_uniform_decode: bool | None = None,
     ) -> bool:
-        """
-        Checks if it's a decode batch with same amount scheduled tokens
+        """Checks if it's a decode batch with same amount scheduled tokens
         across all requests.
         """
         return (
@@ -4044,11 +4029,9 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None):
         )
 
     def _register_layerwise_nvtx_hooks(self) -> None:
-        """
-        Register layerwise NVTX hooks if --enable-layerwise-nvtx-tracing is enabled
+        """Register layerwise NVTX hooks if --enable-layerwise-nvtx-tracing is enabled
         to trace detailed information of each layer or module in the model.
         """
-
         if (
             self.vllm_config.observability_config.enable_layerwise_nvtx_tracing
             and not self.layerwise_nvtx_hooks_registered
@@ -4089,8 +4072,7 @@ def _get_slot_mappings(
         dict[int, torch.Tensor] | None,
         dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None,
     ]:
-        """
-        Build slot mappings in both formats needed by the system.
+        """Build slot mappings in both formats needed by the system.
 
         Args:
             num_tokens_padded: Total number of tokens (padded)
@@ -4102,6 +4084,7 @@ def _get_slot_mappings(
             A tuple of:
             - slot_mappings_by_gid: dict[int, torch.Tensor] for attention metadata
             - slot_mappings_by_layer: dict[str, torch.Tensor] or list for ForwardContext
+
         """
         if not (
             hasattr(self, "kv_cache_config")
@@ -4421,6 +4404,7 @@ def execute_model(
                 num_tokens_unpadded,
                 ubatch_slices_padded,
             )
+        is_padding = self._prepare_padding_mask(num_tokens_unpadded, num_tokens_padded)
         with (
             set_forward_context(
                 attn_metadata,
@@ -4432,6 +4416,7 @@ def execute_model(
                 ubatch_slices=ubatch_slices_padded,
                 slot_mapping=slot_mappings,
                 skip_compiled=has_encoder_input,
+                is_padding=is_padding,
             ),
             record_function_or_nullcontext("gpu_model_runner: forward"),
             self.maybe_get_kv_connector_output(
@@ -4769,39 +4754,14 @@ def propose_draft_token_ids(sampled_token_ids):
                 else None,
                 num_nans_in_logits=num_nans_in_logits,
                 cudagraph_stats=cudagraph_stats,
-                routed_experts=None,
             )
 
         if not self.use_async_scheduling:
-            if self.routed_experts_initialized:
-                # Sync path: D2H was issued in ``_bookkeeping_sync`` and
-                # synchronized by ``_to_list``'s event.synchronize(), so
-                # the pinned buffers are ready to be wrapped as numpy.
-                total = scheduler_output.total_num_scheduled_tokens
-                output.routed_experts = RoutedExpertsLists(
-                    routing_data=self.routed_experts_cpu[:total].numpy(),
-                    slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),
-                )
             return output
 
         with record_function_or_nullcontext(
             "gpu_model_runner: AsyncGPUModelRunnerOutput"
         ):
-            # Async path: produce a device-side snapshot that the async
-            # copy stream can D2H later. Both tensors must be private
-            # clones because:
-            #   - ``routing_data`` source is the shared capturer buffer,
-            #     which the next forward overwrites on the default stream.
-            #   - ``slot_mapping`` source is our own
-            #     ``routed_experts_slot_mapping_device``, which the
-            #     next ``_prepare_inputs`` overwrites on the default
-            #     stream while the D2H is still pending on the copy
-            #     stream.
-            # Without clones, the copy stream would read torn data.
-            routed_experts_snapshot = self.get_routed_experts(
-                scheduler_output.total_num_scheduled_tokens
-            )
-
             async_output = AsyncGPUModelRunnerOutput(
                 model_runner_output=output,
                 sampled_token_ids=sampler_output.sampled_token_ids,
@@ -4809,7 +4769,6 @@ def propose_draft_token_ids(sampled_token_ids):
                 invalid_req_indices=invalid_req_indices,
                 async_output_copy_stream=self._get_or_create_async_output_copy_stream(),
                 vocab_size=self.input_batch.vocab_size,
-                routed_experts=routed_experts_snapshot,
                 check_ep_fault=self.check_ep_fault,
                 num_nans=num_nans_device,
             )
@@ -5289,9 +5248,9 @@ def update_config(self, overrides: dict[str, Any]) -> None:
 
     @instrument(span_name="Loading (GPU)")
     def load_model(self, load_dummy_weights: bool = False) -> None:
-        """
-        Args:
-            load_dummy_weights: load dummy weights instead of real weights.
+        """Args:
+        load_dummy_weights: load dummy weights instead of real weights.
+
         """
         logger.info_once(
             "Starting to load model %s...",
@@ -5310,10 +5269,15 @@ def load_model(self, load_dummy_weights: bool = False) -> None:
                     self.load_config.load_format = "dummy"
                 model_loader = get_model_loader(self.load_config)
                 # Capture warmup providers selected while constructing the model.
-                with self.jit_warmup_registry.activate():
+                with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
                     self.model = model_loader.load_model(
                         vllm_config=self.vllm_config, model_config=self.model_config
                     )
+                lookback_depth = getattr(self.model, "token_lookback_depth", 0)
+                if lookback_depth > 0:
+                    self.lookback_token_ids = self._make_buffer(
+                        self.max_num_reqs, lookback_depth, dtype=torch.int32
+                    )
                 if self.lora_config:
                     self.model = self.load_lora_model(
                         self.model, self.vllm_config, self.device
@@ -5434,13 +5398,22 @@ def load_model(self, load_dummy_weights: bool = False) -> None:
             and cudagraph_mode != CUDAGraphMode.NONE
             and not self.parallel_config.use_ubatching
         ):
-            self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config)
+            # Scoped to PIECEWISE dispatch; FULL cudagraphs (below) are
+            # unaffected. PIECEWISE dispatch can also arise after wrapping
+            # (drafters under a FULL target mode, or a later FULL ->
+            # FULL_AND_PIECEWISE upgrade in _check_and_update_cudagraph_mode).
+            self.model = BreakableCUDAGraphWrapper(
+                self.model, self.vllm_config, runtime_mode=CUDAGraphMode.PIECEWISE
+            )
             drafter = getattr(self, "drafter", None)
             if drafter is not None and hasattr(drafter, "model"):
                 drafter.model = BreakableCUDAGraphWrapper(
-                    drafter.model, self.vllm_config
+                    drafter.model,
+                    self.vllm_config,
+                    runtime_mode=CUDAGraphMode.PIECEWISE,
                 )
-        elif (
+
+        if (
             cudagraph_mode.has_full_cudagraphs()
             and not self.parallel_config.use_ubatching
         ):
@@ -5490,6 +5463,7 @@ def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None:
         Returns:
             Tuple of layer indices if found in draft model config,
             None otherwise.
+
         """
         if not (self.speculative_config and self.speculative_config.draft_model_config):
             return None
@@ -5521,8 +5495,7 @@ def reload_weights(
         weights_path: str | None = None,
         is_checkpoint_format: bool = True,
     ) -> None:
-        """
-        Reload weights from a weights iterator or from disk
+        """Reload weights from a weights iterator or from disk
 
         Args:
             weights_iterator: weights to load into model
@@ -5530,6 +5503,7 @@ def reload_weights(
                 provided. Use path of original model if neither is provided.
             is_checkpoint_format: set to False if weights have already been
                 processed into kernel format (repacking, renaming, etc.)
+
         """
         # TODO(@kylesayrs): generalize to all runners and loaders
         # argument validation
@@ -5739,13 +5713,11 @@ def maybe_randomize_inputs(
         inputs_embeds: torch.Tensor | None,
         randomize_inputs: bool = False,
     ):
-        """
-        Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set.
+        """Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set.
         This is to help balance expert-selection
          - during profile_run
          - during DP rank dummy run
         """
-
         dp_size = self.vllm_config.parallel_config.data_parallel_size
         randomize_inputs = randomize_inputs or (
             envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1
@@ -5791,21 +5763,15 @@ def _get_mm_dummy_batch(
         assert self.mm_budget is not None
 
         # Don't use `max_items_per_batch` here to avoid redundant computation
-        dummy_mm_inputs = self.mm_registry.get_dummy_mm_inputs(
-            self.model_config,
-            mm_counts={modality: 1},
-            cache=self.mm_budget.cache,
+        dummy_mm_inputs = self.mm_budget.get_dummy_encoder_profile_inputs(
+            modality,
+            max_items_per_batch,
         )
-        dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0]
-
-        # We use the cache so that the item is saved to the cache,
-        # but not read from the cache
-        assert dummy_mm_item is not None, "Item should not already be cached"
 
         return next(
             mm_kwargs_batch
             for _, _, mm_kwargs_batch in group_and_batch_mm_kwargs(
-                [(modality, dummy_mm_item)] * max_items_per_batch,
+                dummy_mm_inputs,
                 device=self.device,
                 pin_memory=PIN_MEMORY,
             )
@@ -5828,8 +5794,7 @@ def _dummy_run(
         profile_seq_lens: int | None = None,
         randomize_inputs: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Run a dummy forward pass to warm up/profile run or capture the
+        """Run a dummy forward pass to warm up/profile run or capture the
         CUDA graph for the model.
 
         Args:
@@ -5854,6 +5819,12 @@ def _dummy_run(
             profile_seq_lens: If provided, use this value for seq_lens instead
                 of max_query_len. Used to profile attention workspace that
                 scales with context length.
+            allow_microbatching: Whether the dummy batch may be split into
+                microbatches.
+            is_graph_capturing: Whether this run is a cudagraph capture.
+            randomize_inputs: Whether to fill the dummy inputs with random
+                values rather than zeros.
+
         """
         mm_config = self.vllm_config.model_config.multimodal_config
         if mm_config and mm_config.mm_encoder_only:
@@ -6076,7 +6047,7 @@ def _dummy_run(
         ):
             # Make sure padding doesn't exceed max_num_tokens
             assert num_tokens_padded <= self.max_num_tokens
-            model_kwargs = self._init_model_kwargs()
+            model_kwargs = self._init_model_kwargs(num_reqs=0)
             if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
                 input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)
 
@@ -6087,7 +6058,7 @@ def _dummy_run(
             elif self.enable_prompt_embeds:
                 input_ids = None
                 inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
-                model_kwargs = self._init_model_kwargs()
+                model_kwargs = self._init_model_kwargs(num_reqs=0)
             else:
                 input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
@@ -6113,6 +6084,11 @@ def _dummy_run(
                     num_tokens_padded, None, False
                 )
 
+            num_tokens_unpadded = num_tokens_padded if is_profile else 0
+            is_padding = self._prepare_padding_mask(
+                num_tokens_unpadded, num_tokens_padded
+            )
+
             if ubatch_slices_padded is not None:
                 # Adjust values to reflect a single ubatch.
                 # TODO(sage,lucas): this is cruft that should be addressed in
@@ -6134,6 +6110,7 @@ def _dummy_run(
                     batch_descriptor=batch_desc,
                     ubatch_slices=ubatch_slices_padded,
                     slot_mapping=slot_mappings,
+                    is_padding=is_padding,
                 ),
             ):
                 outputs = self.model(
@@ -6676,31 +6653,11 @@ def profile_cudagraph_memory(self) -> int:
         per_graph_estimate = {}
         encoder_memory_estimate = 0
 
-        # On ROCm, capture these throwaway profiling graphs on vLLM's dedicated
-        # compute stream instead of the fresh side stream graph_capture()
-        # allocates by default. torch's allocator pools free blocks per stream,
-        # so a side-stream forward strands a persistent aiter scratch buffer in
-        # a separate pool, shifting the physical placement of the real KV cache
-        # allocated afterward and slowing bandwidth-bound decode ~20%. The
-        # graphs are discarded, so a side stream is unnecessary here.
-        # Use current_stream(), not torch.cuda.current_stream(): before vLLM
-        # initializes its dedicated stream, torch returns the per-thread default
-        # stream (cuda_stream=0), which cannot be used for cudagraph capture.
-        # cap_ctx=None keeps the side-stream path on CUDA.
-        cap_ctx = (
-            GraphCaptureContext(current_stream())
-            if current_platform.is_rocm()
-            else None
-        )
-
         # Cleanup-only guard: CUDA graph capture errors should still propagate
         # because encoder graph capture is opt-in.
         try:
             set_cudagraph_capturing_enabled(True)
-            with (
-                self._freeze_gc(),
-                graph_capture(device=self.device, graph_capture_context=cap_ctx),
-            ):
+            with self._freeze_gc(), graph_capture(device=self.device):
                 torch.accelerator.synchronize()
                 torch.accelerator.empty_cache()
 
@@ -7003,9 +6960,7 @@ def initialize_attn_backend(
         kv_cache_config: KVCacheConfig,
         is_profiling: bool = False,
     ) -> None:
-        """
-        Initialize the attention backends and attention metadata builders.
-        """
+        """Initialize the attention backends and attention metadata builders."""
         assert len(self.attn_groups) == 0, "Attention backends are already initialized"
 
         class AttentionGroupKey(NamedTuple):
@@ -7108,9 +7063,7 @@ def create_attn_groups(
     def initialize_metadata_builders(
         self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
     ) -> None:
-        """
-        Create the metadata builders for all KV cache groups and attn groups.
-        """
+        """Create the metadata builders for all KV cache groups and attn groups."""
         for kv_cache_group_id in range(len(kv_cache_config.kv_cache_groups)):
             for attn_group in self.attn_groups[kv_cache_group_id]:
                 attn_group.create_metadata_builders(
@@ -7145,8 +7098,7 @@ def _check_and_update_cudagraph_mode(
         kv_cache_groups: list[KVCacheGroupSpec],
         is_profiling: bool = False,
     ) -> None:
-        """
-        Resolve the cudagraph_mode when there are multiple attention
+        """Resolve the cudagraph_mode when there are multiple attention
         groups with potential conflicting CUDA graph support.
         Then initialize the cudagraph_dispatcher based on the resolved
         cudagraph_mode.
@@ -7199,8 +7151,7 @@ def _check_and_update_cudagraph_mode(
             self.drafter.initialize_cudagraph_keys(cudagraph_mode)
 
     def calculate_reorder_batch_threshold(self) -> None:
-        """
-        Choose the minimum reorder batch threshold from all attention groups.
+        """Choose the minimum reorder batch threshold from all attention groups.
         Backends should be able to support lower threshold then what they request
         just may have a performance penalty due to that backend treating decodes
         as prefills.
@@ -7221,8 +7172,7 @@ def calculate_reorder_batch_threshold(self) -> None:
     def may_reinitialize_input_batch(
         self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
     ) -> None:
-        """
-        Re-initialize the input batch if the block sizes are different from
+        """Re-initialize the input batch if the block sizes are different from
         what it was originally created with. This happens when the final
         block size (determined after model loading) differs from the
         placeholder used during __init__, or when there are multiple
@@ -7231,6 +7181,7 @@ def may_reinitialize_input_batch(
         Args:
             kv_cache_config: The KV cache configuration.
             kernel_block_sizes: The kernel block sizes for each KV cache group.
+
         """
         block_sizes = []
         max_num_blocks = []
@@ -7268,7 +7219,7 @@ def may_reinitialize_input_batch(
                 self.parallel_config.cp_kv_cache_interleave_size
             )
             # Capture warmup providers registered after final KV-cache geometry is known
-            with self.jit_warmup_registry.activate():
+            with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
                 self.input_batch = InputBatch(
                     max_num_reqs=self.max_num_reqs,
                     max_model_len=max_model_len,
@@ -7312,18 +7263,19 @@ def initialize_kv_cache_tensors(
         kernel_block_sizes: list[int],
         kv_cache_allocation_context: AbstractContextManager | None = None,
     ) -> dict[str, torch.Tensor]:
-        """
-        Initialize the memory buffer for KV cache.
+        """Initialize the memory buffer for KV cache.
 
         Args:
             kv_cache_config: The KV cache config
             kernel_block_sizes: The kernel block sizes for each KV cache group.
+            kv_cache_allocation_context: Context manager the allocation runs
+                inside, e.g. to place the KV cache in a specific memory pool.
 
         Returns:
             Dict[str, torch.Tensor]: A map between layer names to their
             corresponding memory buffer for KV cache.
-        """
 
+        """
         allocation_context = kv_cache_allocation_context or nullcontext()
         with allocation_context:
             kv_caches = allocate_kv_cache(
@@ -7353,8 +7305,7 @@ def initialize_kv_cache_tensors(
     def maybe_add_kv_sharing_layers_to_kv_cache_groups(
         self, kv_cache_config: KVCacheConfig
     ) -> None:
-        """
-        Add layers that re-use KV cache to KV cache group of its target layer.
+        """Add layers that re-use KV cache to KV cache group of its target layer.
         Mapping of KV cache tensors happens in `initialize_kv_cache_tensors()`
         """
         if not self.shared_kv_cache_layers:
@@ -7384,11 +7335,15 @@ def initialize_kv_cache(
         is_profiling: bool = False,
         kv_cache_allocation_context: AbstractContextManager | None = None,
     ) -> None:
-        """
-        Initialize KV cache based on `kv_cache_config`.
+        """Initialize KV cache based on `kv_cache_config`.
+
         Args:
             kv_cache_config: Configuration for the KV cache, including the KV
-            cache size of each layer
+                cache size of each layer
+            is_profiling: Whether this call is part of a profiling run.
+            kv_cache_allocation_context: Context manager the allocation runs
+                inside, e.g. to place the KV cache in a specific memory pool.
+
         """
         kv_cache_config = deepcopy(kv_cache_config)
         self.kv_cache_config = kv_cache_config
@@ -7418,7 +7373,7 @@ def initialize_kv_cache(
         # Reinitialize need to after initialize_attn_backend
         self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
         # Capture warmup providers that depend on allocated KV-cache strides.
-        with self.jit_warmup_registry.activate():
+        with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
             kv_caches = self.initialize_kv_cache_tensors(
                 kv_cache_config,
                 kernel_block_sizes,
@@ -7439,65 +7394,8 @@ def initialize_kv_cache(
             kv_transfer_group.register_kv_caches(kv_caches)
             kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
 
-    def get_routed_experts(
-        self,
-        num_tokens: int,
-    ) -> RoutedExpertsTensors | None:
-        if not self.routed_experts_initialized:
-            return None
-
-        device_buffer = self.routed_experts_capturer.get_device_buffer()
-        return RoutedExpertsTensors(
-            routing_data=device_buffer[:num_tokens].clone(),
-            slot_mapping=self.routed_experts_slot_mapping_device[:num_tokens].clone(),
-        )
-
-    def init_routed_experts_capturer(self):
-        logger.info(
-            "Initializing routed experts capturer, enable_return_routed_experts: %s",
-            self.model_config.enable_return_routed_experts,
-        )
-        self.routed_experts_capturer = RoutedExpertsCapturer(
-            max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
-            vllm_config=self.vllm_config,
-            kv_cache_config=self.kv_cache_config,
-        )
-        bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
-
-        # Pinned CPU buffer for non-blocking D2H of ``routing_data`` on
-        # the sync scheduling path. Shape / dtype mirror the device
-        # capturer exactly so ``copy_`` is a straight memcpy.
-        self.routed_experts_cpu = torch.empty(
-            self.routed_experts_capturer.device_buffer.shape,
-            dtype=self.routed_experts_capturer.device_buffer.dtype,
-            device="cpu",
-            pin_memory=PIN_MEMORY,
-        )
-        # ``slot_mapping`` dtype is fixed to int64 by
-        # ``block_table.slot_mapping``; we mirror that here.
-        max_tokens = self.scheduler_config.max_num_batched_tokens
-        self.routed_experts_slot_mapping_cpu = torch.empty(
-            (max_tokens,),
-            dtype=torch.int64,
-            device="cpu",
-            pin_memory=PIN_MEMORY,
-        )
-        # Private device buffer so the shared ``block_table.slot_mapping``
-        # can be overwritten by the next ``_prepare_inputs`` while the
-        # D2H is still pending on the copy stream. Written in
-        # ``_prepare_inputs``, read in ``_bookkeeping_sync`` (sync path)
-        # or cloned into a snapshot (async path).
-        self.routed_experts_slot_mapping_device = torch.empty(
-            (max_tokens,),
-            dtype=torch.int64,
-            device=self.device,
-        )
-        self.routed_experts_initialized = True
-
     def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
-        """
-        Add encoder-only layers to the KV cache config.
-        """
+        """Add encoder-only layers to the KV cache config."""
         block_size = self.vllm_config.cache_config.block_size
         encoder_only_attn_specs: dict[AttentionSpec, list[str]] = defaultdict(list)
         attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
@@ -7521,12 +7419,13 @@ def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
             )
 
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
-        """
-        Generates the KVCacheSpec by parsing the kv cache format from each
+        """Generates the KVCacheSpec by parsing the kv cache format from each
         Attention module in the static forward context.
+
         Returns:
             KVCacheSpec: A dictionary mapping layer names to their KV cache
             format. Layers that do not need KV cache are not included.
+
         """
         if has_ec_transfer() and not get_ec_transfer().is_consumer:
             return {}
@@ -7570,11 +7469,11 @@ def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
         return pinned.tolist()
 
     def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
-        """
-        Get encoder timing stats for all requests and clear the registry.
+        """Get encoder timing stats for all requests and clear the registry.
 
         Returns:
             Dictionary mapping request_id to stats dict.
+
         """
         with self._encoder_timing_lock:
             stats = {
@@ -7592,14 +7491,14 @@ def timed_encoder_operation(
         current_item_idx: int,
         num_items: int,
     ):
-        """
-        Context manager to time encoder forward operations.
+        """Context manager to time encoder forward operations.
 
         Args:
             should_time: Whether timing is enabled
             group_lora_refs: Full list of (request_id, pos_info) tuples
             current_item_idx: Starting index for this group
             num_items: Number of items in this group
+
         """
         if not should_time:
             yield
diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py
index 0735b39f81c2..f2cab7898be5 100644
--- a/vllm/v1/worker/gpu_ubatch_wrapper.py
+++ b/vllm/v1/worker/gpu_ubatch_wrapper.py
@@ -119,8 +119,7 @@ def unwrap(self) -> Callable:
         return self.runnable
 
     def _capture_ubatches(self, ubatch_metadata, model) -> torch.Tensor:
-        """
-        Capture a cudagraph for a microbatched run.
+        """Capture a cudagraph for a microbatched run.
 
         The logic here is somewhat complicated because we need to make sure that
         each of the ubatch threads initialize the cuda context before we start
@@ -262,6 +261,7 @@ def _make_ubatch_metadata(
         dp_metadata,
         batch_descriptor,
         cudagraph_runtime_mode,
+        is_padding,
     ) -> list[UbatchMetadata]:
         # Create one forward context per ubatch
         forward_contexts = []
@@ -277,6 +277,11 @@ def _make_ubatch_metadata(
                     batch_descriptor=batch_descriptor,
                     cudagraph_runtime_mode=cudagraph_runtime_mode,
                     slot_mapping=slot_mapping[i] if has_slot_mapping else None,
+                    is_padding=(
+                        is_padding[ubatch_slice.token_slice]
+                        if is_padding is not None
+                        else None
+                    ),
                 )
             )
 
@@ -374,6 +379,7 @@ def __call__(self, *args, **kwargs):
 
         attn_metadata = forward_context.attn_metadata
         slot_mapping = forward_context.slot_mapping
+        is_padding = forward_context.is_padding
         num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices)
         input_ids = kwargs["input_ids"]
         positions = kwargs["positions"]
@@ -415,6 +421,7 @@ def __call__(self, *args, **kwargs):
                 dp_metadata=ubatch_dp_metadata,
                 batch_descriptor=batch_descriptor,
                 cudagraph_runtime_mode=CUDAGraphMode.NONE,
+                is_padding=is_padding,
             )
             with self.sm_control:
                 return self._capture_ubatches(ubatch_metadata, self.runnable)
@@ -441,6 +448,7 @@ def __call__(self, *args, **kwargs):
                 dp_metadata=ubatch_dp_metadata,
                 batch_descriptor=batch_descriptor,
                 cudagraph_runtime_mode=CUDAGraphMode.NONE,
+                is_padding=is_padding,
             )
             with self.sm_control:
                 return self._run_ubatches(ubatch_metadata, self.runnable)
diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py
index 53bc0550a524..7e4831447c04 100644
--- a/vllm/v1/worker/gpu_worker.py
+++ b/vllm/v1/worker/gpu_worker.py
@@ -8,8 +8,9 @@
 from collections.abc import Callable
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from datetime import timedelta
+from fnmatch import filter as fnmatch_filter
 from types import NoneType
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import regex as re
@@ -45,6 +46,7 @@
     Handle,
     checkpoint_prepare_distributed_state,
     checkpoint_restore_distributed_state,
+    get_pcp_group,
     get_pp_group,
     get_tp_group,
     resume_device_comms,
@@ -60,9 +62,9 @@
 from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory
 from vllm.platforms import current_platform
 from vllm.profiler.wrapper import (
-    CudaProfilerWrapper,
-    ProtonProfilerWrapper,
-    TorchProfilerWrapper,
+    create_graph_capture_profiler,
+    create_worker_profiler,
+    validate_worker_profiler_config,
 )
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
@@ -96,6 +98,7 @@
 from vllm.v1.worker.workspace import init_workspace_manager
 
 from ...model_executor.model_loader import TensorizerLoader
+from .gpu.cudagraph_utils import has_compiled_submodule
 from .gpu.warmup import warmup_kernels
 from .utils import request_memory
 
@@ -141,6 +144,7 @@ def maybe_rocm_profiling_fallback(profile_result: MemoryProfilingResult) -> int
 if TYPE_CHECKING:
     from vllm.device_allocator.sleep_mode_backend import SleepModeBackend
     from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 
@@ -205,6 +209,7 @@ def __init__(
             self.worker_sentinel = WorkerSentinel(worker=self)
         # Buffers saved before sleep
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
+        self._sleep_saved_parameters: dict[str, torch.Tensor] = {}
         self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {}
 
         # Weight transfer engine is created in `load_model` once the model
@@ -218,6 +223,7 @@ def __init__(
         # so we have all the information needed for proper trace naming.
         self.profiler: Any | None = None
         self.profiler_config = vllm_config.profiler_config
+        validate_worker_profiler_config(self.profiler_config)
 
         self.use_v2_model_runner = vllm_config.use_v2_model_runner
 
@@ -239,6 +245,30 @@ def sleep_mode_backend(self) -> "SleepModeBackend":
             )
         return self._sleep_mode_backend
 
+    def _save_sleep_parameters(self, model: nn.Module) -> None:
+        patterns = self.model_config.sleep_preserve_parameter_names
+        if not patterns:
+            self._sleep_saved_parameters = {}
+            return
+        parameters = dict(model.named_parameters())
+        names: set[str] = set()
+        for pattern in patterns:
+            matches = fnmatch_filter(parameters, pattern)
+            if not matches:
+                raise ValueError(f"No parameter matches sleep retention: {pattern}")
+            names.update(matches)
+        self._sleep_saved_parameters = {
+            name: param.detach().to("cpu")
+            for name, param in parameters.items()
+            if name in names and not param.is_cpu
+        }
+
+    @torch.no_grad()
+    def _restore_sleep_parameters(self, model: nn.Module) -> None:
+        for name, value in self._sleep_saved_parameters.items():
+            model.get_parameter(name).copy_(value)
+        self._sleep_saved_parameters.clear()
+
     def sleep(self, level: int = 1) -> None:
         torch.accelerator.synchronize()
         free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]
@@ -246,6 +276,7 @@ def sleep(self, level: int = 1) -> None:
         # Save the buffers before level 2 sleep
         if level == 2:
             model = self.model_runner.model
+            self._save_sleep_parameters(model)
             self._sleep_saved_buffers = {
                 name: buffer.cpu().clone() for name, buffer in model.named_buffers()
             }
@@ -283,6 +314,8 @@ def wake_up(self, tags: list[str] | None = None) -> None:
 
         # Restore the buffers after level 2 sleep
         wake_weights = tags is None or "weights" in tags
+        if wake_weights and self._sleep_saved_parameters:
+            self._restore_sleep_parameters(self.model_runner.model)
         if wake_weights and len(self._sleep_saved_buffers):
             model = self.model_runner.model
             for name, buffer in model.named_buffers():
@@ -300,6 +333,9 @@ def wake_up(self, tags: list[str] | None = None) -> None:
 
         self.synchronize_device()
 
+    def discard(self, tags: tuple[str, ...]) -> None:
+        self.sleep_mode_backend.discard(tags)
+
     def checkpoint_prepare(self) -> None:
         checkpoint_prepare_distributed_state()
 
@@ -370,13 +406,14 @@ def init_device(self):
                 if dp_local_rank is None:
                     dp_local_rank = self.parallel_config.data_parallel_index
 
-                tp_pp_world_size = (
+                tp_pcp_pp_world_size = (
                     self.parallel_config.pipeline_parallel_size
+                    * self.parallel_config.prefill_context_parallel_size
                     * self.parallel_config.tensor_parallel_size
                 )
 
-                # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
-                self.local_rank += dp_local_rank * tp_pp_world_size
+                # DP_LOCAL_RANK * TP_PCP_PP_WORLD_SIZE + TP_LOCAL_RANK
+                self.local_rank += dp_local_rank * tp_pcp_pp_world_size
 
             # Publish the logical-to-physical mapping for topology queries
             # such as NIC affinity and P2P checks.
@@ -514,6 +551,11 @@ def load_model(self, *, load_dummy_weights: bool = False) -> None:
                 self.model_runner.get_model(),
             )
 
+        # Preserve parallel weight loading, then use serving's thread count
+        # for profiling and compilation so Dynamo's global-state guards remain
+        # valid when requests arrive.
+        set_torch_threads_for_runtime()
+
     def update_config(self, overrides: dict[str, Any]) -> None:
         self.model_runner.update_config(overrides)
 
@@ -533,6 +575,7 @@ def determine_available_memory(self) -> int:
         Tip:
             You may limit the usage of GPU memory
             by adjusting the `gpu_memory_utilization` parameter.
+
         """
         maybe_apply_startup_plan(self)
 
@@ -573,12 +616,10 @@ def determine_available_memory(self) -> int:
         # torch.accelerator.get_memory_info (reliable on ROCm, as used by
         # the AMD-CI mem tests), and graph_pool_handle resolves to the same
         # torch.cuda handle the live capture path already uses on ROCm.
-        # XPU stays excluded (see #39977).
         cudagraph_memory_estimate = 0
         if (
-            current_platform.is_cuda_alike()
-            and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
-        ):
+            current_platform.is_cuda_alike() or current_platform.is_xpu()
+        ) and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
             cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()
 
         # Respect the opt-in flag as originally designed.
@@ -691,7 +732,6 @@ def get_kv_connector_handshake_metadata(
 
         Returned dict is keyed by `(pp_rank, tp_rank)`.
         """
-
         if not has_kv_transfer_group():
             return None
 
@@ -703,6 +743,14 @@ def get_kv_connector_handshake_metadata(
 
         pp_rank = get_pp_group().rank_in_group
         tp_rank = get_tp_group().rank_in_group
+        parallel_config = self.vllm_config.parallel_config
+        if (
+            parallel_config.prefill_context_parallel_size > 1
+            and parallel_config.decode_context_parallel_size > 1
+        ):
+            tp_rank += (
+                get_pcp_group().rank_in_group * parallel_config.tensor_parallel_size
+            )
         return {(pp_rank, tp_rank): metadata}
 
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
@@ -723,7 +771,6 @@ def update_max_model_len(self, max_model_len: int) -> None:
     @instrument(span_name="Allocate KV cache")
     def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
         """Allocate GPU KV cache with the specified kv_cache_config."""
-
         # Update local config with adjusted num blocks after profiling,
         # so that it's available to the warmup stage.
         self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks
@@ -747,9 +794,6 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
             ),
         )
 
-        if self.model_config.enable_return_routed_experts:
-            self.model_runner.init_routed_experts_capturer()
-
         # Build KV-zero metadata outside the CuMem pool so the bookkeeping
         # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch
         # allocator and are not discarded during sleep/wake cycles.
@@ -762,7 +806,10 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
     def compile_or_warm_up_model(self) -> CompilationTimes:
         warmup_sizes: list[int] = []
 
-        if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE:
+        if (
+            self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE
+            and has_compiled_submodule(self.model_runner.get_model())
+        ):
             # warm up sizes that are not in cudagraph capture sizes,
             # but users still want to compile for better performance,
             # e.g. for the max-num-batched token size in chunked prefill.
@@ -795,13 +842,17 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
         # cuda graph capture.
         kernel_warmup(self)
 
-        if self.use_v2_model_runner:
-            # A workspace resize after capture frees what the graphs point at.
-            warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)
+        if self.use_v2_model_runner:  # noqa: SIM102
+            if self.vllm_config.kernel_config.enable_jit_warmup:
+                # A workspace resize after capture frees what the graphs point at.
+                warmup_kernels(
+                    self.model_runner, self.execute_model, self.sample_tokens
+                )
 
         cuda_graph_memory_bytes = 0
         if not self.model_config.enforce_eager:
-            cuda_graph_memory_bytes = self.model_runner.capture_model()
+            with self._get_cudagraph_capture_context():
+                cuda_graph_memory_bytes = self.model_runner.capture_model()
 
         # Compare actual vs estimated CUDA graph memory (if we did profiling)
         if (
@@ -931,15 +982,26 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
         # gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
         enable_gpu_sync_check()
 
-        # Startup is done; steady-state serving gets no benefit from torch
-        # intra-op parallelism.
-        set_torch_threads_for_runtime()
-
         return CompilationTimes(
             language_model=self.compilation_config.compilation_time,
             encoder=self.compilation_config.encoder_compilation_time,
         )
 
+    def _get_cudagraph_capture_context(self) -> AbstractContextManager[None]:
+        """Let the configured profiler observe CUDA graph capture."""
+        if not self.use_v2_model_runner:
+            return nullcontext()
+        if self.profiler is None:
+            model_runner = cast("GPUModelRunnerV2", self.model_runner)
+            if not model_runner.needs_cudagraph_capture():
+                return nullcontext()
+            self.profiler = create_graph_capture_profiler(
+                self.profiler_config, global_rank=self.rank
+            )
+            if self.profiler is None:
+                return nullcontext()
+        return self.profiler.capture_cuda_graphs()
+
     def reset_mm_cache(self) -> None:
         self.model_runner.reset_mm_cache()
 
@@ -1005,7 +1067,7 @@ def annotate_profile(self, scheduler_output):
             return nullcontext()
 
         self.profiler.step()
-        if not self.profiler.is_running:
+        if not self.profiler.should_annotate:
             return nullcontext()
 
         iteration_details = compute_iteration_details(scheduler_output)
@@ -1211,6 +1273,8 @@ def execute_model(
         )
         self._pp_send_work = handles[1:]
 
+        if self.use_v2_model_runner and self.model_runner.is_pooling_model:
+            return self.model_runner.pool()  # type: ignore
         return None
 
     def take_draft_token_ids(self) -> DraftTokenIds | None:
@@ -1227,7 +1291,6 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None):
             )
 
         if is_start:
-            profiler_type = self.profiler_config.profiler
             # Generate the trace name by combining prefix with comprehensive rank suffix
             from vllm.distributed.utils import get_worker_rank_suffix
 
@@ -1239,33 +1302,16 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None):
             else:
                 trace_name = rank_suffix
 
+            if self.profiler_config.profiler == "proton" and self.profiler is not None:
+                self.profiler.set_output_name(trace_name)
+
             # Create the profiler wrapper only on the first start call
             if self.profiler is None:
-                if profiler_type == "torch":
-                    self.profiler = TorchProfilerWrapper(
-                        self.profiler_config,
-                        worker_name=trace_name,
-                        local_rank=self.local_rank,
-                        activities=["CPU", "CUDA"],
-                    )
-                    logger.debug(
-                        "Starting torch profiler with trace name: %s", trace_name
-                    )
-                elif profiler_type == "cuda":
-                    self.profiler = CudaProfilerWrapper(self.profiler_config)
-                    logger.debug("Starting CUDA profiler")
-                elif profiler_type == "proton":
-                    self.profiler = ProtonProfilerWrapper(
-                        self.profiler_config, worker_name=trace_name
-                    )
-                    logger.debug(
-                        "Starting Proton profiler with trace name: %s", trace_name
-                    )
-                else:
-                    # Config validation should prevent this code being reached
-                    raise ValueError(
-                        f"Invalid profiler value of {self.profiler_config.profiler}"
-                    )
+                self.profiler = create_worker_profiler(
+                    self.profiler_config,
+                    worker_name=trace_name,
+                    local_rank=self.local_rank,
+                )
 
             self.profiler.start()
         else:
@@ -1275,7 +1321,9 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None):
             try:
                 self.profiler.stop()
             finally:
-                if self.profiler_config.profiler == "proton":
+                if self.profiler_config.profiler == "proton" and not (
+                    self.profiler.has_cuda_graph_session
+                ):
                     # Proton output names are fixed when the wrapper is constructed.
                     # Recreate it so the next profile_prefix is honored.
                     self.profiler = None
@@ -1330,12 +1378,12 @@ def _check_weight_transfer_engine(self) -> None:
             )
 
     def init_weight_transfer_engine(self, init_info: dict) -> None:
-        """
-        Initialize weight transfer mechanism.
+        """Initialize weight transfer mechanism.
         For NCCL backend, this creates a process group with the trainer.
 
         Args:
             init_info: Dictionary containing backend-specific initialization info
+
         """
         self._check_weight_transfer_engine()
         assert self.weight_transfer_engine is not None
@@ -1344,8 +1392,7 @@ def init_weight_transfer_engine(self, init_info: dict) -> None:
         self.weight_transfer_engine.init_transfer_engine(typed_init_info)
 
     def start_weight_update(self) -> None:
-        """
-        Start a new weight update session.
+        """Start a new weight update session.
 
         Delegates engine-specific preparation (e.g. layerwise reload setup) to
         the configured weight transfer engine. The worker only tracks that a
@@ -1355,8 +1402,7 @@ def start_weight_update(self) -> None:
             self._start_weight_update()
 
     def start_draft_weight_update(self) -> None:
-        """
-        Like start_weight_update, but retargets the engine at the speculative
+        """Like start_weight_update, but retargets the engine at the speculative
         draft model for this session.
         """
         with set_current_vllm_config(self.vllm_config):
@@ -1389,8 +1435,7 @@ def _start_weight_update(self, is_draft: bool = False) -> None:
         self._weight_update_is_draft = is_draft
 
     def update_weights(self, update_info: dict | list[dict]) -> None:
-        """
-        Receive one weight update chunk from the trainer.
+        """Receive one weight update chunk from the trainer.
 
         start_weight_update must be called before update_weights and
         finish_weight_update must be called after all chunks have been sent.
@@ -1400,6 +1445,7 @@ def update_weights(self, update_info: dict | list[dict]) -> None:
         Args:
             update_info: Backend-specific update info, or a list indexed by
                 global worker rank across data parallel replicas.
+
         """
         self._check_weight_transfer_engine()
         assert self.weight_transfer_engine is not None
@@ -1414,7 +1460,7 @@ def update_weights(self, update_info: dict | list[dict]) -> None:
                 if isinstance(update_info, list):
                     parallel_config = self.vllm_config.parallel_config
                     local_update_info = update_info[
-                        parallel_config.data_parallel_rank * parallel_config.world_size
+                        parallel_config.data_parallel_index * parallel_config.world_size
                         + self.rank
                     ]
                 else:
diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py
index 5d6235bf3359..26ffc159f3c0 100644
--- a/vllm/v1/worker/kv_connector_model_runner_mixin.py
+++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Define KV connector functionality mixin for model runners.
-"""
+"""Define KV connector functionality mixin for model runners."""
 
 from collections.abc import Generator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
@@ -94,9 +92,12 @@ def _get_kv_connector_output(
             if wait_for_save and not defer_finalize:
                 kv_connector.wait_for_save()
 
-            output.finished_sending, output.finished_recving = (
-                kv_connector.get_finished(scheduler_output.finished_req_ids)
+            transfer_results = kv_connector.get_transfer_results(
+                scheduler_output.finished_req_ids
             )
+            output.finished_sending = transfer_results.finished_sending
+            output.finished_recving = transfer_results.finished_recving
+            output.failed_recving = transfer_results.failed_recving
             output.invalid_block_ids = kv_connector.get_block_ids_with_load_errors()
 
             output.kv_connector_stats = kv_connector.get_kv_connector_stats()
diff --git a/vllm/v1/worker/lora_model_runner_mixin.py b/vllm/v1/worker/lora_model_runner_mixin.py
index 27e850b66a67..9d94e62edb84 100644
--- a/vllm/v1/worker/lora_model_runner_mixin.py
+++ b/vllm/v1/worker/lora_model_runner_mixin.py
@@ -1,12 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Define LoRA functionality mixin for model runners.
-"""
+"""Define LoRA functionality mixin for model runners."""
 
 from collections.abc import Callable
 from contextlib import contextmanager
-from typing import TypeAlias
 
 import numpy as np
 import torch
@@ -19,10 +16,7 @@
 from vllm.lora.request import LoRARequest
 from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
 from vllm.model_executor.models import supports_lora
-from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch
-from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch
-
-InputBatch: TypeAlias = TPUInputBatch | GPUInputBatch
+from vllm.v1.worker.gpu_input_batch import InputBatch
 
 logger = init_logger(__name__)
 
@@ -154,16 +148,17 @@ def maybe_select_dummy_loras(
         num_sampled_tokens: np.ndarray | None = None,
         num_active_loras: int = 0,
     ):
-        """
-        Context manager to select dummy LoRAs for capture/warmup.
+        """Context manager to select dummy LoRAs for capture/warmup.
 
         Args:
             lora_config: LoRA configuration, or None if LoRA is disabled.
             num_scheduled_tokens: Array of scheduled token counts per request.
+            mapping_type: Which LoRA mapping to build (language or encoder).
             num_sampled_tokens: Array of sampled token counts per request.
             num_active_loras: Number of distinct active LoRAs to use.
                 - 0: No LoRA active (set up zero mappings).
                 - >0: Use exactly this many distinct LoRAs.
+
         """
         if num_sampled_tokens is None:
             num_sampled_tokens = np.ones_like(num_scheduled_tokens, dtype=np.int32)
@@ -259,8 +254,7 @@ def maybe_dummy_run_with_lora(
         num_active_loras: int = 0,
         mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
     ):
-        """
-        Context manager for dummy runs with LoRA.
+        """Context manager for dummy runs with LoRA.
 
         Args:
             lora_config: LoRA configuration.
@@ -269,6 +263,8 @@ def maybe_dummy_run_with_lora(
             remove_lora: Whether to remove LoRAs after the context exits.
             num_active_loras: Number of distinct active LoRAs to use.
                 LoRA is activated when num_active_loras > 0.
+            mapping_type: Which LoRA mapping to build (language or encoder).
+
         """
         with (
             self.maybe_setup_dummy_loras(lora_config, remove_lora),
diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py
index 724c714057f7..c2e0120e4613 100644
--- a/vllm/v1/worker/mamba_utils.py
+++ b/vllm/v1/worker/mamba_utils.py
@@ -26,8 +26,7 @@
     UniformTypeKVCacheSpecs,
 )
 from vllm.v1.utils import CpuGpuBuffer
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
-from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
 logger = init_logger(__name__)
 
@@ -416,8 +415,7 @@ def postprocess_mamba_fused_kernel(
     # the existing 2D-grid contract.
     TEMPORAL_TILES: tl.constexpr = 1,
 ):
-    """
-    Fused GPU kernel for postprocess_mamba that computes decisions AND performs
+    """Fused GPU kernel for postprocess_mamba that computes decisions AND performs
     mamba state copies without any CPU-GPU synchronization.
 
     Grid: (num_reqs, num_states [, TEMPORAL_TILES])
@@ -770,8 +768,7 @@ def create(
 
 @dataclasses.dataclass
 class MambaSpecDecodeGPUContext:
-    """
-    Context for GPU-side Mamba state copy operations during the
+    """Context for GPU-side Mamba state copy operations during the
     fused postprocess path.
 
     Only used when speculative decoding is enabled on a hybrid model
@@ -928,8 +925,7 @@ def initialize_from_forward_context(
         mamba_state_copy_funcs: MambaStateCopyFuncsByType,
         block_tables: list[torch.Tensor],
     ) -> None:
-        """
-        Extract and cache memory layout metadata from Mamba state tensors.
+        """Extract and cache memory layout metadata from Mamba state tensors.
 
         This method populates the pre-allocated metadata tensors with information
         needed by `postprocess_mamba_fused_kernel` to perform state copies entirely
@@ -961,6 +957,7 @@ def initialize_from_forward_context(
             block_tables: per-mamba-group persistent block-table tensors, in
                 the same order as `mamba_group_ids`. Their `data_ptr()` /
                 `stride(0)` are captured once for the kernel to index into.
+
         """
         if self.is_initialized:
             return
@@ -1097,7 +1094,7 @@ def compute_aligned_state_indices(
         seq_lens: torch.Tensor,
         num_reqs: int,
     ) -> torch.Tensor:
-        """compute every Mamba group's aligned physical state IDs in one launch."""
+        """Compute every Mamba group's aligned physical state IDs in one launch."""
         assert self.is_initialized
         assert seq_lens.is_cuda
         assert 0 <= num_reqs <= seq_lens.shape[0]
@@ -1138,8 +1135,7 @@ def run_fused_postprocess(
         num_computed_tokens_gpu: torch.Tensor,
         num_draft_tokens_gpu: torch.Tensor,
     ) -> None:
-        """
-        Run the fused postprocess_mamba kernel on GPU.
+        """Run the fused postprocess_mamba kernel on GPU.
 
         This computes decisions and performs mamba state copies entirely on GPU,
         eliminating the CPU-GPU sync that was previously needed.
@@ -1151,6 +1147,7 @@ def run_fused_postprocess(
             num_scheduled_tokens_gpu: [num_reqs] scheduled token counts
             num_computed_tokens_gpu: [num_reqs] computed token counts
             num_draft_tokens_gpu: [num_reqs] draft token counts
+
         """
         if num_reqs == 0 or not self.is_initialized:
             return
@@ -1206,6 +1203,7 @@ def run_fused_precopy(
             token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1).
             idx_mapping: optional [num_reqs] batch_idx -> req_state_idx.
                 None means V1 batch order already equals request state order.
+
         """
         if num_reqs == 0 or not self.is_initialized:
             return
@@ -1436,15 +1434,14 @@ def preprocess_mamba(
     kv_cache_config: KVCacheConfig,
     cache_config: CacheConfig,
     mamba_state_idx: dict[str, int],
-    input_batch: GPUInputBatch,
+    input_batch: InputBatch,
     requests: dict[str, CachedRequestState],
     forward_context: dict[str, Any],
     mamba_state_copy_funcs: MambaStateCopyFuncsByType,
     copy_bufs: MambaCopyBuffers,
     align_ctx: MambaSpecDecodeGPUContext | None = None,
 ):
-    """
-    Copy the mamba state of previous step to the last
+    """Copy the mamba state of previous step to the last
     (1 + num_speculative_blocks) block.
     """
     fused = _resolve_fused_precopy(align_ctx)
@@ -1541,7 +1538,7 @@ def preprocess_mamba(
 def postprocess_mamba_all(
     scheduler_output: SchedulerOutput,
     kv_cache_config: KVCacheConfig,
-    input_batch: GPUInputBatch,
+    input_batch: InputBatch,
     requests: dict[str, CachedRequestState],
     mamba_state_idx: dict[str, int],
     num_spec_tokens: int,
@@ -1572,7 +1569,7 @@ def postprocess_mamba_all(
 
 def preprocess_mamba_all_specdec(
     scheduler_output: SchedulerOutput,
-    input_batch: GPUInputBatch,
+    input_batch: InputBatch,
     mamba_state_idx: dict[str, int],
     num_reqs: int,
     prev_last_scheduled_idx_buf: CpuGpuBuffer,
@@ -1591,7 +1588,7 @@ def postprocess_mamba_align_gpu(
     num_reqs: int,
     num_accepted_tokens_gpu: torch.Tensor,
     num_accepted_tokens_cpu_tensor: torch.Tensor,
-    input_batch: GPUInputBatch,
+    input_batch: InputBatch,
     kv_cache_config: KVCacheConfig,
     forward_context: dict[str, Any],
     mamba_state_copy_funcs: MambaStateCopyFuncsByType,
diff --git a/vllm/v1/worker/mm_encoder_model_runner.py b/vllm/v1/worker/mm_encoder_model_runner.py
index 8208a6ce3df9..f0d19fa12ef2 100644
--- a/vllm/v1/worker/mm_encoder_model_runner.py
+++ b/vllm/v1/worker/mm_encoder_model_runner.py
@@ -4,9 +4,10 @@
 
 An encoder-only instance -- `--mm-encoder-only`, or the producer side of
 encoder-cache disaggregation -- encodes the multi-modal items and publishes the
-embeddings. It runs no language model: no KV cache, no sampler, no CUDA graphs.
+embeddings. It runs no language model: no KV cache, sampler, or decoder graphs.
 """
 
+import time
 from collections.abc import Iterator
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any
@@ -15,7 +16,9 @@
 
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
+from vllm.logger import init_logger
 from vllm.sequence import IntermediateTensors
+from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
 from vllm.v1.kv_cache_interface import KVCacheSpec
 from vllm.v1.outputs import (
     ModelRunnerOutput,
@@ -24,10 +27,13 @@
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+from vllm.v1.worker.workspace import lock_workspace
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
 
+logger = init_logger(__name__)
+
 
 class MMEncoderModelRunner(GPUModelRunner):
     """Encoder-only variant of the V2 GPU model runner."""
@@ -67,8 +73,34 @@ def input_tensor_semaphore(self) -> Iterator[None]:
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
         return {}
 
+    @torch.inference_mode()
     def capture_model(self, *, profile_only: bool = False) -> int:
-        return 0
+        encoder_runner = self.model_state.encoder_runner
+        if not encoder_runner.has_cudagraph():
+            return 0
+
+        start_time = time.perf_counter()
+        with freeze_gc_for_cudagraph_capture():
+            torch.accelerator.empty_cache()
+            start_free_memory = torch.accelerator.get_memory_info()[0]
+            encoder_runner.capture()
+            graph_memory = start_free_memory - torch.accelerator.get_memory_info()[0]
+
+        if not profile_only:
+            lock_workspace()
+
+        manager = encoder_runner.cudagraph_manager
+        assert manager is not None
+        logger.info(
+            "Encoder graph capturing finished in %.0f secs, took %.2f GiB (%d graphs)",
+            time.perf_counter() - start_time,
+            graph_memory / (1 << 30),
+            manager.get_num_graphs_to_capture(),
+        )
+        return graph_memory
+
+    def needs_cudagraph_capture(self) -> bool:
+        return False
 
     def _dummy_run(
         self, *args: Any, **kwargs: Any
@@ -111,13 +143,15 @@ def execute_model(
 
             batch_req_state, _ = self.gather_batch_req_state(scheduler_output, False)
             assert batch_req_state is not None
-            # No CUDA graph, and no DP peer to agree a padded shape with.
+            # No decoder graph or DP padding; encoder graphs run independently.
             batch_desc = BatchExecutionDescriptor(
                 cg_mode=CUDAGraphMode.NONE,
                 num_tokens=batch_req_state.num_tokens,
                 num_reqs=None,
             )
-            self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
+            self.prepare_inputs(
+                scheduler_output, batch_req_state, batch_desc, num_active_loras=0
+            )
 
         scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
         if self.lora_config is not None:
diff --git a/vllm/v1/worker/tpu_input_batch.py b/vllm/v1/worker/tpu_input_batch.py
deleted file mode 100644
index c16156a5e3ef..000000000000
--- a/vllm/v1/worker/tpu_input_batch.py
+++ /dev/null
@@ -1,559 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Datastructures defining a TPU input batch
-
-from typing import cast
-
-import numpy as np
-import torch
-
-from vllm.lora.request import LoRARequest
-from vllm.sampling_params import SamplingType
-from vllm.utils import length_from_prompt_token_ids_or_embeds
-from vllm.utils.collection_utils import swap_dict_values
-from vllm.v1.outputs import LogprobsTensors
-from vllm.v1.worker.block_table import MultiGroupBlockTable
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
-
-_SAMPLING_EPS = 1e-5
-
-
-class InputBatch:
-    def __init__(
-        self,
-        max_num_reqs: int,
-        max_model_len: int,
-        max_num_batched_tokens: int,
-        device: torch.device,
-        pin_memory: bool,
-        vocab_size: int,
-        block_sizes: list[int],  # The block_size of each kv cache group
-        kernel_block_sizes: list[int],
-        max_num_blocks_per_req: list[int],
-    ):
-        self.max_num_reqs = max_num_reqs
-        self.max_model_len = max_model_len
-        self.max_num_batched_tokens = max_num_batched_tokens
-        self.device = device
-        self.pin_memory = pin_memory
-        self.vocab_size = vocab_size
-
-        self._req_ids: list[str | None] = []
-        self.req_id_to_index: dict[str, int] = {}
-
-        # TODO(woosuk): This buffer could be too large if max_model_len is big.
-        # Find a way to reduce the CPU memory usage.
-        # This buffer is not directly transferred to the GPU, so it does not
-        # need to be pinned.
-        self.token_ids_cpu_tensor = torch.zeros(
-            (max_num_reqs, max_model_len),
-            device="cpu",
-            dtype=torch.int32,
-            pin_memory=False,
-        )
-        self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
-        self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
-        self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
-        self.num_computed_tokens_cpu_tensor = torch.zeros(
-            (max_num_reqs,),
-            device="cpu",
-            dtype=torch.int32,
-            pin_memory=pin_memory,
-        )
-        self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy()
-
-        # Block table.
-        self.block_table = MultiGroupBlockTable(
-            max_num_reqs=max_num_reqs,
-            max_num_batched_tokens=max_num_batched_tokens,
-            pin_memory=pin_memory,
-            device=device,
-            block_sizes=block_sizes,
-            kernel_block_sizes=kernel_block_sizes,
-            max_num_blocks=max_num_blocks_per_req,
-        )
-
-        # Sampling-related.
-        self.temperature = torch.empty(
-            (max_num_reqs,), dtype=torch.float32, device=device
-        )
-        self.temperature_cpu_tensor = torch.empty(
-            (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory
-        )
-        self.temperature_cpu = self.temperature_cpu_tensor.numpy()
-        self.greedy_reqs: set[str] = set()
-        self.random_reqs: set[str] = set()
-
-        self.top_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device)
-        self.top_p_cpu_tensor = torch.empty(
-            (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory
-        )
-        self.top_p_cpu = self.top_p_cpu_tensor.numpy()
-        self.top_p_reqs: set[str] = set()
-
-        self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device)
-        self.top_k_cpu_tensor = torch.empty(
-            (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory
-        )
-        self.top_k_cpu = self.top_k_cpu_tensor.numpy()
-        self.top_k_reqs: set[str] = set()
-
-        self.min_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device)
-        self.min_p_cpu_tensor = torch.empty(
-            (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory
-        )
-        self.min_p_cpu = self.min_p_cpu_tensor.numpy()
-        self.min_p_reqs: set[str] = set()
-
-        # Frequency penalty related data structures
-        self.frequency_penalties = torch.empty(
-            (max_num_reqs,), dtype=torch.float, device=device
-        )
-        self.frequency_penalties_cpu_tensor = torch.empty(
-            (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
-        )
-        self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy()
-        self.frequency_penalties_reqs: set[str] = set()
-
-        # Presence penalty related data structures
-        self.presence_penalties = torch.empty(
-            (max_num_reqs,), dtype=torch.float, device=device
-        )
-        self.presence_penalties_cpu_tensor = torch.empty(
-            (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
-        )
-        self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy()
-        self.presence_penalties_reqs: set[str] = set()
-
-        # Repetition penalty related data structures
-        self.repetition_penalties = torch.empty(
-            (max_num_reqs,), dtype=torch.float, device=device
-        )
-        self.repetition_penalties_cpu_tensor = torch.empty(
-            (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
-        )
-        self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy()
-        self.repetition_penalties_reqs: set[str] = set()
-
-        # req_index -> (min_tokens, stop_token_ids)
-        self.min_tokens: dict[int, tuple[int, set[int]]] = {}
-
-        # lora related
-        self.request_lora_mapping = np.zeros((self.max_num_reqs,), dtype=np.int64)
-        self.lora_id_to_request_ids: dict[int, set[str]] = {}
-        self.lora_id_to_lora_request: dict[int, LoRARequest] = {}
-
-        # req_index -> generator
-        # NOTE(woosuk): The indices of the requests that do not have their own
-        # generator should not be included in the dictionary.
-        self.generators: dict[int, torch.Generator] = {}
-
-        self.num_logprobs: dict[str, int] = {}
-
-        # To accumulate prompt logprobs tensor chunks across prefill steps.
-        self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
-
-        self.logit_bias: list[dict[int, float] | None] = [None] * max_num_reqs
-        self.has_allowed_token_ids: set[str] = set()
-        # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
-        # the value is False. Since we use masked_fill_ to set -inf.
-        self.allowed_token_ids_mask: torch.Tensor | None = None
-        self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None
-
-        # req_index -> bad_words_token_ids
-        self.bad_words_token_ids: dict[int, list[list[int]]] = {}
-
-        self.req_output_token_ids: list[list[int] | None] = []
-
-    @property
-    def req_ids(self) -> list[str]:
-        # None elements should only be present transiently
-        # while performing state updates to the batch.
-        return cast(list[str], self._req_ids)
-
-    def add_request(
-        self,
-        request: "CachedRequestState",
-        req_index: int | None = None,
-    ) -> None:
-        if req_index is None:
-            req_index = self.num_reqs
-        assert req_index < self.max_num_reqs
-
-        req_id = request.req_id
-        if req_index == len(self._req_ids):
-            self._req_ids.append(req_id)
-            self.req_output_token_ids.append(request.output_token_ids)
-        else:
-            self._req_ids[req_index] = req_id
-            self.req_output_token_ids[req_index] = request.output_token_ids
-
-        self.req_id_to_index[req_id] = req_index
-
-        # Copy the prompt token ids and output token ids.
-        num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
-            request.prompt_token_ids, request.prompt_embeds
-        )
-        # TODO: copy prompt_embeds
-        self.num_prompt_tokens[req_index] = num_prompt_tokens
-        self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids
-        start_idx = num_prompt_tokens
-        end_idx = start_idx + len(request.output_token_ids)
-        self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids
-        # Number of tokens without spec decode tokens.
-        self.num_tokens_no_spec[req_index] = request.num_tokens
-
-        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
-        self.block_table.add_row(request.block_ids, req_index)
-
-        sampling_params = request.sampling_params
-        assert sampling_params is not None, "pooling requests not supported yet"
-        if sampling_params.sampling_type == SamplingType.GREEDY:
-            # Should avoid division by zero later when apply_temperature.
-            self.temperature_cpu[req_index] = 0.0
-            self.greedy_reqs.add(req_id)
-        else:
-            self.temperature_cpu[req_index] = sampling_params.temperature
-            self.random_reqs.add(req_id)
-
-        self.top_p_cpu[req_index] = sampling_params.top_p
-        if sampling_params.top_p < 1:
-            self.top_p_reqs.add(req_id)
-        top_k = sampling_params.top_k
-        if 0 < top_k < self.vocab_size:
-            self.top_k_reqs.add(req_id)
-        else:
-            top_k = self.vocab_size
-        self.top_k_cpu[req_index] = top_k
-        self.min_p_cpu[req_index] = sampling_params.min_p
-        self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty
-        if sampling_params.min_p > _SAMPLING_EPS:
-            self.min_p_reqs.add(req_id)
-        if sampling_params.frequency_penalty != 0.0:
-            self.frequency_penalties_reqs.add(req_id)
-        self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty
-        if sampling_params.presence_penalty != 0.0:
-            self.presence_penalties_reqs.add(req_id)
-        self.repetition_penalties_cpu[req_index] = sampling_params.repetition_penalty
-        if sampling_params.repetition_penalty != 1.0:
-            self.repetition_penalties_reqs.add(req_id)
-        if sampling_params.min_tokens:
-            self.min_tokens[req_index] = (
-                sampling_params.min_tokens,
-                sampling_params.all_stop_token_ids,
-            )
-
-        # NOTE(woosuk): self.generators should not include the requests that
-        # do not have their own generator.
-        if request.generator is not None:
-            self.generators[req_index] = request.generator
-
-        if sampling_params.logprobs is not None:
-            self.num_logprobs[req_id] = sampling_params.logprobs
-        if sampling_params.logit_bias is not None:
-            self.logit_bias[req_index] = sampling_params.logit_bias
-
-        if sampling_params.allowed_token_ids:
-            self.has_allowed_token_ids.add(req_id)
-            if self.allowed_token_ids_mask_cpu_tensor is None:
-                # Lazy allocation for this tensor, which can be large.
-                # False means we don't fill with -inf.
-                self.allowed_token_ids_mask = torch.zeros(
-                    self.max_num_reqs,
-                    self.vocab_size,
-                    dtype=torch.bool,
-                    device=self.device,
-                )
-                self.allowed_token_ids_mask_cpu_tensor = torch.zeros(
-                    self.max_num_reqs, self.vocab_size, dtype=torch.bool, device="cpu"
-                )
-            self.allowed_token_ids_mask_cpu_tensor[req_index] = True
-            # False means we don't fill with -inf.
-            self.allowed_token_ids_mask_cpu_tensor[req_index][
-                sampling_params.allowed_token_ids
-            ] = False
-
-        if sampling_params.bad_words_token_ids:
-            self.bad_words_token_ids[req_index] = sampling_params.bad_words_token_ids
-
-        # Add request lora ID
-        if request.lora_request:
-            lora_id = request.lora_request.lora_int_id
-            if lora_id not in self.lora_id_to_request_ids:
-                self.lora_id_to_request_ids[lora_id] = set()
-
-            self.request_lora_mapping[req_index] = lora_id
-            self.lora_id_to_request_ids[lora_id].add(request.req_id)
-            self.lora_id_to_lora_request[lora_id] = request.lora_request
-        else:
-            # No LoRA
-            self.request_lora_mapping[req_index] = 0
-
-    def remove_request(self, req_id: str) -> int | None:
-        """This method must always be followed by a call to condense()."""
-
-        req_index = self.req_id_to_index.pop(req_id, None)
-        if req_index is None:
-            return None
-        self._req_ids[req_index] = None
-        self.req_output_token_ids[req_index] = None
-
-        self.greedy_reqs.discard(req_id)
-        self.random_reqs.discard(req_id)
-        self.top_p_reqs.discard(req_id)
-        self.top_k_reqs.discard(req_id)
-        self.min_p_reqs.discard(req_id)
-        self.min_tokens.pop(req_index, None)
-        self.frequency_penalties_reqs.discard(req_id)
-        self.presence_penalties_reqs.discard(req_id)
-        self.repetition_penalties_reqs.discard(req_id)
-        self.generators.pop(req_index, None)
-        self.num_logprobs.pop(req_id, None)
-        self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
-
-        # LoRA
-        lora_id = self.request_lora_mapping[req_index]
-        if lora_id != 0:
-            self.lora_id_to_request_ids[lora_id].discard(req_id)
-            if len(self.lora_id_to_request_ids[lora_id]) == 0:
-                self.lora_id_to_request_ids.pop(lora_id)
-                self.lora_id_to_lora_request.pop(lora_id)
-            self.request_lora_mapping[req_index] = 0
-
-        self.logit_bias[req_index] = None
-        self.has_allowed_token_ids.discard(req_id)
-        if self.allowed_token_ids_mask_cpu_tensor is not None:
-            # False means we don't fill with -inf.
-            self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
-        self.bad_words_token_ids.pop(req_index, None)
-        return req_index
-
-    def swap_states(self, i1: int, i2: int) -> None:
-        old_id_i1 = self._req_ids[i1]
-        old_id_i2 = self._req_ids[i2]
-        self._req_ids[i1], self._req_ids[i2] = self._req_ids[i2], self._req_ids[i1]  # noqa
-        self.req_output_token_ids[i1], self.req_output_token_ids[i2] = (
-            self.req_output_token_ids[i2],
-            self.req_output_token_ids[i1],
-        )
-        assert old_id_i1 is not None and old_id_i2 is not None
-        self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] = (
-            self.req_id_to_index[old_id_i2],
-            self.req_id_to_index[old_id_i1],
-        )
-        self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = (
-            self.num_tokens_no_spec[i2],
-            self.num_tokens_no_spec[i1],
-        )
-        self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] = (
-            self.num_prompt_tokens[i2],
-            self.num_prompt_tokens[i1],
-        )
-        self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = (
-            self.num_computed_tokens_cpu[i2],
-            self.num_computed_tokens_cpu[i1],
-        )
-        self.temperature_cpu[i1], self.temperature_cpu[i2] = (
-            self.temperature_cpu[i2],
-            self.temperature_cpu[i1],
-        )
-        self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1]
-        self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1]
-        self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = (
-            self.frequency_penalties_cpu[i2],
-            self.frequency_penalties_cpu[i1],
-        )
-        self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] = (
-            self.presence_penalties_cpu[i2],
-            self.presence_penalties_cpu[i1],
-        )
-        self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = (
-            self.repetition_penalties_cpu[i2],
-            self.repetition_penalties_cpu[i1],
-        )
-        self.min_p_cpu[i1], self.min_p_cpu[i2] = self.min_p_cpu[i2], self.min_p_cpu[i1]
-
-        # NOTE: the following is unsafe
-        # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
-        #     self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]
-        # instead, we need to temporarily copy the data for one of the indices
-        n = max(self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2])
-        tmp = self.token_ids_cpu[i1, :n].copy()
-        self.token_ids_cpu[i1, :n] = self.token_ids_cpu[i2, :n]
-        self.token_ids_cpu[i2, :n] = tmp
-
-        swap_dict_values(self.generators, i1, i2)
-        swap_dict_values(self.min_tokens, i1, i2)
-        swap_dict_values(self.bad_words_token_ids, i1, i2)
-
-        self.request_lora_mapping[i1], self.request_lora_mapping[i2] = (
-            self.request_lora_mapping[i2],
-            self.request_lora_mapping[i1],
-        )
-        self.logit_bias[i1], self.logit_bias[i2] = (
-            self.logit_bias[i2],
-            self.logit_bias[i1],
-        )
-
-        if self.allowed_token_ids_mask_cpu_tensor is not None:
-            (
-                self.allowed_token_ids_mask_cpu_tensor[i1],
-                self.allowed_token_ids_mask_cpu_tensor[i2],
-            ) = (
-                self.allowed_token_ids_mask_cpu_tensor[i2],
-                self.allowed_token_ids_mask_cpu_tensor[i1],
-            )
-        self.block_table.swap_row(i1, i2)
-
-    def condense(self, empty_req_indices: list[int]) -> None:
-        """Move non-empty requests down into lower, empty indices.
-
-        Args:
-          empty_req_indices: empty batch indices, sorted descending.
-        """
-        num_reqs = self.num_reqs
-        if num_reqs == 0:
-            # The batched states are empty.
-            self._req_ids.clear()
-            self.req_output_token_ids.clear()
-            return
-
-        # NOTE(woosuk): This function assumes that the empty_req_indices
-        # is sorted in descending order.
-        last_req_index = num_reqs + len(empty_req_indices) - 1
-        while empty_req_indices:
-            # Find the largest non-empty index.
-            while last_req_index in empty_req_indices:
-                last_req_index -= 1
-
-            # Find the smallest empty index.
-            empty_index = empty_req_indices.pop()
-            if empty_index >= last_req_index:
-                break
-
-            # Swap the states.
-            req_id = self._req_ids[last_req_index]
-            output_token_ids = self.req_output_token_ids[last_req_index]
-            assert req_id is not None
-            self._req_ids[empty_index] = req_id
-            self._req_ids[last_req_index] = None
-            self.req_output_token_ids[empty_index] = output_token_ids
-            self.req_output_token_ids[last_req_index] = None
-            self.req_id_to_index[req_id] = empty_index
-
-            num_tokens = self.num_tokens_no_spec[last_req_index]
-            self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
-                last_req_index, :num_tokens
-            ]
-            self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
-                last_req_index
-            ]
-            self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index]
-            self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[
-                last_req_index
-            ]
-            self.block_table.move_row(last_req_index, empty_index)
-            self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index]
-            self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index]
-            self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index]
-            self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[
-                last_req_index
-            ]
-            self.presence_penalties_cpu[empty_index] = self.presence_penalties_cpu[
-                last_req_index
-            ]
-            self.repetition_penalties_cpu[empty_index] = self.repetition_penalties_cpu[
-                last_req_index
-            ]
-            self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index]
-            generator = self.generators.pop(last_req_index, None)
-            if generator is not None:
-                self.generators[empty_index] = generator
-
-            min_token = self.min_tokens.pop(last_req_index, None)
-            if min_token is not None:
-                self.min_tokens[empty_index] = min_token
-
-            self.request_lora_mapping[empty_index] = self.request_lora_mapping[
-                last_req_index
-            ]
-
-            self.logit_bias[empty_index] = self.logit_bias[last_req_index]
-
-            if self.allowed_token_ids_mask_cpu_tensor is not None:
-                self.allowed_token_ids_mask_cpu_tensor[empty_index] = (
-                    self.allowed_token_ids_mask_cpu_tensor[last_req_index]
-                )
-
-            bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None)
-            if bad_words_token_ids is not None:
-                self.bad_words_token_ids[empty_index] = bad_words_token_ids
-            # Decrement last_req_index since it is now empty.
-            last_req_index -= 1
-
-        # Trim lists to the batch size.
-        del self._req_ids[self.num_reqs :]
-        del self.req_output_token_ids[self.num_reqs :]
-
-    def make_lora_inputs(
-        self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray
-    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
-        """
-        Given the num_scheduled_tokens for each request in the batch, return
-        datastructures used to activate the current LoRAs.
-        Returns:
-            1. prompt_lora_mapping: A tuple of size self.num_reqs where,
-               prompt_lora_mapping[i] is the LoRA id to use for the ith prompt.
-            2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens)
-               where, token_lora_mapping[i] is the LoRA id to use for ith token.
-            3. lora_requests: Set of relevant LoRA requests.
-        """
-
-        req_lora_mapping = self.request_lora_mapping[: self.num_reqs]
-        prompt_lora_mapping = tuple(req_lora_mapping)
-        token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens))
-        active_lora_requests: set[LoRARequest] = set(
-            self.lora_id_to_lora_request.values()
-        )
-
-        return prompt_lora_mapping, token_lora_mapping, active_lora_requests
-
-    @property
-    def num_reqs(self) -> int:
-        return len(self.req_id_to_index)
-
-    @property
-    def all_greedy(self) -> bool:
-        return len(self.random_reqs) == 0
-
-    @property
-    def all_random(self) -> bool:
-        return len(self.greedy_reqs) == 0
-
-    @property
-    def no_top_p(self) -> bool:
-        return len(self.top_p_reqs) == 0
-
-    @property
-    def no_top_k(self) -> bool:
-        return len(self.top_k_reqs) == 0
-
-    @property
-    def no_min_p(self) -> bool:
-        return len(self.min_p_reqs) == 0
-
-    @property
-    def no_penalties(self) -> bool:
-        return (
-            len(self.presence_penalties_reqs) == 0
-            and len(self.frequency_penalties_reqs) == 0
-            and len(self.repetition_penalties_reqs) == 0
-        )
-
-    @property
-    def max_num_logprobs(self) -> int | None:
-        return max(self.num_logprobs.values()) if self.num_logprobs else None
-
-    @property
-    def no_allowed_token_ids(self) -> bool:
-        return len(self.has_allowed_token_ids) == 0
diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py
index f9612b1e7f82..329c261a0a38 100644
--- a/vllm/v1/worker/ubatch_utils.py
+++ b/vllm/v1/worker/ubatch_utils.py
@@ -43,8 +43,7 @@ def __init__(
         set_comm_sms: Callable[[int], None],
         set_compute_sms: Callable[[int], None],
     ):
-        """
-        Context manager for controlling SM (Streaming Multiprocessor)
+        """Context manager for controlling SM (Streaming Multiprocessor)
         allocation. Upon entering the context, it sets the number of SMs
         allocated for communication and computation to comm_sms and
         total_sms - comm_sms respectively. Upon exiting, it restores the
@@ -57,8 +56,8 @@ def __init__(
                 A function that sets the number of SMs for communication.
             set_compute_sms (Callable[[int], None]):
                 A function that sets the number of SMs for computation.
-        """
 
+        """
         assert current_platform.is_cuda() or current_platform.is_rocm(), (
             "SM/CU control is supported on CUDA and ROCm platforms"
         )
@@ -221,8 +220,7 @@ def slice_query_start_locs(
     query_start_loc: torch.Tensor,
     request_slice: slice,
 ) -> torch.Tensor:
-    """
-    Creates a new query_start_loc that corresponds to the requests in
+    """Creates a new query_start_loc that corresponds to the requests in
     request_slice.
 
     Note: This function creates a new tensor to hold the new query_start_locs.
@@ -237,11 +235,9 @@ def slice_query_start_locs(
 def _make_metadata_with_slice(
     ubatch_slice: UBatchSlice, attn_metadata: CommonAttentionMetadata
 ) -> CommonAttentionMetadata:
-    """
-    This function creates a new CommonAttentionMetadata that corresponds to
+    """This function creates a new CommonAttentionMetadata that corresponds to
     the requests included in ubatch_slice
     """
-
     assert not ubatch_slice.is_empty(), f"Ubatch slice {ubatch_slice} is empty"
 
     request_slice = ubatch_slice.request_slice
@@ -280,23 +276,11 @@ def _make_metadata_with_slice(
         query_start_loc[1:] -= tokens_skipped
         query_start_loc_cpu[1:] -= tokens_skipped
     seq_lens = attn_metadata.seq_lens[request_slice]
-    # Read raw fields to avoid triggering the deprecated D2H-syncing properties.
-    seq_lens_cpu = (
-        attn_metadata._seq_lens_cpu[request_slice]
-        if attn_metadata._seq_lens_cpu is not None
-        else None
-    )
     seq_lens_cpu_upper_bound = (
         attn_metadata.seq_lens_cpu_upper_bound[request_slice]
         if attn_metadata.seq_lens_cpu_upper_bound is not None
         else None
     )
-    num_computed_tokens_cpu = (
-        attn_metadata._num_computed_tokens_cpu[request_slice]
-        if attn_metadata._num_computed_tokens_cpu is not None
-        else None
-    )
-
     if splits_last_request:
         # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate
         # the tokens skipped because query_start_loc_cpu might have been modified
@@ -309,9 +293,6 @@ def _make_metadata_with_slice(
         #  (not cudagraph compatible)
         seq_lens = seq_lens.clone()
         seq_lens[-1] -= tokens_skipped
-        if seq_lens_cpu is not None:
-            seq_lens_cpu = seq_lens_cpu.clone()
-            seq_lens_cpu[-1] -= tokens_skipped
         if seq_lens_cpu_upper_bound is not None:
             seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound.clone()
             seq_lens_cpu_upper_bound[-1] -= tokens_skipped
@@ -346,8 +327,6 @@ def _make_metadata_with_slice(
         block_table_tensor=block_table_tensor,
         slot_mapping=slot_mapping,
         seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
-        _seq_lens_cpu=seq_lens_cpu,
-        _num_computed_tokens_cpu=num_computed_tokens_cpu,
     )
 
 
@@ -355,8 +334,7 @@ def split_attn_metadata(
     ubatch_slices: list[UBatchSlice],
     common_attn_metadata: CommonAttentionMetadata,
 ) -> list[CommonAttentionMetadata]:
-    """
-    Creates a new CommonAttentionMetadata instance that corresponds to the
+    """Creates a new CommonAttentionMetadata instance that corresponds to the
     requests for each UBatchSlice in ubatch_slices.
 
     Note: This function does not modify common_attn_metadata
diff --git a/vllm/v1/worker/ubatching.py b/vllm/v1/worker/ubatching.py
index 6d06ba2135f5..a773ee715d14 100644
--- a/vllm/v1/worker/ubatching.py
+++ b/vllm/v1/worker/ubatching.py
@@ -18,9 +18,7 @@
 
 
 class UBatchContext:
-    """
-    Context manager for micro-batching synchronization using threading events.
-    """
+    """Context manager for micro-batching synchronization using threading events."""
 
     def __init__(
         self,
diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py
index 067a2b0d2267..6c2a05a6260d 100644
--- a/vllm/v1/worker/utils.py
+++ b/vllm/v1/worker/utils.py
@@ -166,6 +166,8 @@ def __init__(
                 kv = static_forward_context[layer_name].kv_cache
                 if not isinstance(kv, torch.Tensor):
                     continue
+                if kv.device.type != self.device.type:
+                    continue
                 dp = kv.data_ptr()
 
                 assert kv.shape[0] % num_blocks == 0, (
@@ -329,8 +331,7 @@ def select_common_block_size(
     kv_manager_block_size: int,
     backends: list[type[AttentionBackend]],
 ) -> int:
-    """
-    Select a block size that is supported by all backends and is a factor of
+    """Select a block size that is supported by all backends and is a factor of
     kv_manager_block_size.
 
     If kv_manager_block_size is supported by all backends, return it directly.
@@ -345,6 +346,7 @@ def select_common_block_size(
 
     Raises:
         ValueError: If no valid block size found.
+
     """
 
     def block_size_is_supported(
@@ -366,32 +368,22 @@ def block_size_is_supported(
                 return False
         return True
 
-    # Case 1: if the block_size of kv cache manager is supported by all backends,
-    # return it directly.
     if block_size_is_supported(backends, kv_manager_block_size):
         return kv_manager_block_size
 
-    # Case 2: otherwise, the block_size must be an `int`-format supported size of
-    # at least one backend. Iterate over all `int`-format supported sizes in
-    # descending order and return the first one that is supported by all backends.
-    # Simple proof:
-    # If the supported size b is in MultipleOf(x_i) format for all attention
-    # backends i, and b a factor of kv_manager_block_size, then
-    # kv_manager_block_size also satisfies MultipleOf(x_i) for all i. We will
-    # return kv_manager_block_size in case 1.
-    all_int_supported_sizes = set(
-        supported_size
+    # MultipleOf constraints also accept the manager size if they accept a divisor.
+    # Any remaining candidate must therefore be an explicit size from a backend.
+    candidates = {
+        size
         for backend in backends
-        for supported_size in backend.get_supported_kernel_block_sizes()
-        if isinstance(supported_size, int)
-    )
+        for size in backend.get_supported_kernel_block_sizes()
+        if isinstance(size, int) and kv_manager_block_size % size == 0
+    }
 
-    for supported_size in sorted(all_int_supported_sizes, reverse=True):
-        if kv_manager_block_size % supported_size != 0:
-            continue
-        if block_size_is_supported(backends, supported_size):
-            return supported_size
-    raise ValueError(f"No common block size for {kv_manager_block_size}. ")
+    for size in sorted(candidates, reverse=True):
+        if block_size_is_supported(backends, size):
+            return size
+    raise ValueError(f"No common block size for {kv_manager_block_size}.")
 
 
 def allocate_kv_cache(
@@ -440,7 +432,11 @@ def allocate_kv_cache(
         if isinstance(spec, UniformTypeKVCacheSpecs):
             spec = spec.kv_cache_specs[layer_name]
 
-        num_blocks = kv_cache_config.num_blocks
+        if not spec.has_layer_views:
+            kv_caches.update((name, buf) for name in tensor.layers)
+            continue
+
+        num_blocks = kv_cache_config.num_blocks_of(tensor)
         kernel_block_size = None
         if kernel_block_sizes is not None and group_id < len(kernel_block_sizes):
             kernel_block_size = kernel_block_sizes[group_id]
@@ -462,8 +458,7 @@ def allocate_kv_cache(
 def prepare_kernel_block_sizes(
     kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]]
 ) -> list[int]:
-    """
-    Generate kernel_block_sizes that matches each block_size.
+    """Generate kernel_block_sizes that matches each block_size.
 
     For attention backends that support virtual block splitting,
     use the supported block sizes from the backend.
@@ -475,6 +470,7 @@ def prepare_kernel_block_sizes(
 
     Returns:
         List of kernel block sizes for each cache group.
+
     """
     kernel_block_sizes = []
     for kv_cache_gid, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
@@ -485,7 +481,9 @@ def prepare_kernel_block_sizes(
             kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values()))
         if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
             continue
-        if isinstance(kv_cache_spec, AttentionSpec):
+        if not kv_cache_spec.has_layer_views:
+            kernel_block_sizes.append(kv_cache_spec.block_size)
+        elif isinstance(kv_cache_spec, AttentionSpec):
             # This is an attention backend that supports virtual block splitting.
             kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size
             group_backends = [g.backend for g in attn_groups[kv_cache_gid]]
@@ -507,8 +505,7 @@ def sanity_check_mm_encoder_outputs(
     mm_embeddings: MultiModalEmbeddings,
     expected_num_items: int,
 ) -> None:
-    """
-    Perform sanity checks for the result of
+    """Perform sanity checks for the result of
     [`vllm.model_executor.models.SupportsMultiModal.embed_multimodal`][].
     """
     assert isinstance(mm_embeddings, (list, tuple, torch.Tensor)), (
@@ -534,8 +531,7 @@ def sanity_check_mm_encoder_outputs(
 
 
 def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> int:
-    """
-    Calculate the amount of memory required by vLLM, then validate
+    """Calculate the amount of memory required by vLLM, then validate
     that the current amount of free memory is sufficient for that.
     """
     requested_memory = math.ceil(
@@ -561,8 +557,7 @@ def add_kv_sharing_layers_to_kv_cache_groups(
     kv_cache_groups: list[KVCacheGroupSpec],
     runner_only_attn_layers: set[str] | None = None,
 ) -> None:
-    """
-    Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches`
+    """Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches`
     for layers that do not allocate its own KV cache, based on the mapping in
     `shared_kv_cache_layers`. Adds these layers to the corresponding KV cache
     group, which is needed to ensure that attention metadata is assigned later.
@@ -573,6 +568,9 @@ def add_kv_sharing_layers_to_kv_cache_groups(
             means this layer will perform attention using the keys and values
             from the KV cache of `shared_kv_cache_layers[layer_name]`.
         kv_cache_groups: The KV cache groups of the model.
+        runner_only_attn_layers: Attention layers handled by the runner only,
+            which are excluded from the KV cache groups.
+
     """
     if not shared_kv_cache_layers:
         return
@@ -597,8 +595,7 @@ def bind_kv_cache(
     num_attn_module: int = 1,
     kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
 ) -> None:
-    """
-    Bind the allocated KV cache to both ModelRunner and forward context so
+    """Bind the allocated KV cache to both ModelRunner and forward context so
     that the KV cache can be used in the forward pass.
 
     This function:
@@ -609,9 +606,13 @@ def bind_kv_cache(
 
     Args:
         kv_caches: The allocated kv_caches with layer names as keys.
+        num_attn_module: Number of attention modules per layer entry.
         forward_context: The global forward context containing all Attention
             layers with layer names as keys.
         runner_kv_caches: The kv_cache declared by ModelRunner.
+        kv_cache_groups: The KV cache groups of the model, used to resolve
+            layers that share a KV cache.
+
     """
     # Bind kv_caches to ModelRunner
     assert len(runner_kv_caches) == 0
@@ -621,7 +622,6 @@ def bind_kv_cache(
     for layer_name in kv_caches:
         index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name)
 
-    ordered_layer_names: list[str] = []
     for layer_index in sorted(index2name.keys()):
         layer_names = index2name[layer_index]
         if len(layer_names) > 1:
@@ -635,8 +635,19 @@ def bind_kv_cache(
             current_platform.check_runner_kv_caches_multi_layer()
         for layer_name in layer_names:
             runner_kv_caches.append(kv_caches[layer_name])
-            ordered_layer_names.append(layer_name)
 
+    bind_kv_cache_to_layers(
+        kv_caches, forward_context, num_attn_module, kv_cache_groups
+    )
+
+
+def bind_kv_cache_to_layers(
+    kv_caches: dict[str, torch.Tensor],
+    forward_context: dict[str, Attention],
+    num_attn_module: int = 1,
+    kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
+) -> None:
+    """Bind layer caches and share ReplaySSM trackers in model-layer order."""
     # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks
     # its raw allocation into the per-layer view(s) it needs (e.g. Mamba
     # splits conv/ssm), so the kv_caches dict can hold a single tensor per
@@ -644,6 +655,9 @@ def bind_kv_cache(
     for layer_name, kv_cache in kv_caches.items():
         forward_context[layer_name].bind_kv_cache(kv_cache)
 
+    ordered_layer_names = sorted(
+        kv_caches, key=lambda name: extract_layer_index(name, num_attn_module)
+    )
     share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups)
 
 
diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py
index cf5b0e7ce334..b0839bc4cae2 100644
--- a/vllm/v1/worker/worker_base.py
+++ b/vllm/v1/worker/worker_base.py
@@ -12,7 +12,7 @@
 from vllm.distributed.kv_transfer.kv_connector.utils import get_current_attn_backends
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
-from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.cache import worker_receiver_cache_from_config
 from vllm.tracing import instrument
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.system_utils import update_environment_variables
@@ -55,8 +55,7 @@ def __init__(
         distributed_init_method: str,
         is_driver_worker: bool = False,
     ) -> None:
-        """
-        Initialize common worker components.
+        """Initialize common worker components.
 
         Args:
             vllm_config: Complete vLLM configuration
@@ -65,6 +64,7 @@ def __init__(
             distributed_init_method: Distributed initialization method
             is_driver_worker: Whether this worker handles driver
                 responsibilities
+
         """
         self.vllm_config = vllm_config
         self.model_config = vllm_config.model_config
@@ -118,6 +118,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
 
         Returns:
             Compilation times (language_model, encoder) in seconds.
+
         """
         raise NotImplementedError
 
@@ -209,8 +210,7 @@ def shutdown(self) -> None:
 
 
 class WorkerWrapperBase:
-    """
-    This class represents one process in an executor/engine. It is responsible
+    """This class represents one process in an executor/engine. It is responsible
     for lazily initializing the worker and handling the worker's lifecycle.
     We first instantiate the WorkerWrapper, which remembers the worker module
     and class name. Then, when we call `update_environment_variables`, and the
@@ -222,8 +222,7 @@ def __init__(
         rpc_rank: int = 0,
         global_rank: int | None = None,
     ) -> None:
-        """
-        Initialize the worker wrapper with the given vllm_config and rpc_rank.
+        """Initialize the worker wrapper with the given vllm_config and rpc_rank.
         Note: rpc_rank is the rank of the worker in the executor. In most cases,
         it is also the rank of the worker in the distributed group. However,
         when multiple executors work together, they can be different.
@@ -252,8 +251,7 @@ def update_environment_variables(
 
     @instrument(span_name="Worker init")
     def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None:
-        """
-        Here we inject some common logic before initializing the worker.
+        """Here we inject some common logic before initializing the worker.
         Arguments are passed to the worker class constructor.
         """
         kwargs = all_kwargs[self.rpc_rank]
@@ -317,26 +315,10 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None:
             )
 
         shared_worker_lock = kwargs.pop("shared_worker_lock", None)
-        if shared_worker_lock is None:
-            msg = (
-                "Missing `shared_worker_lock` argument from executor. "
-                "This argument is needed for mm_processor_cache_type='shm'."
-            )
-
-            mm_config = vllm_config.model_config.multimodal_config
-            if mm_config and mm_config.mm_processor_cache_type == "shm":
-                raise ValueError(msg)
-            else:
-                logger.warning_once(msg)
-
-            self.mm_receiver_cache = None
-        else:
-            self.mm_receiver_cache = (
-                MULTIMODAL_REGISTRY.worker_receiver_cache_from_config(
-                    vllm_config,
-                    shared_worker_lock,
-                )
-            )
+        self.mm_receiver_cache = worker_receiver_cache_from_config(
+            vllm_config,
+            shared_worker_lock,
+        )
 
         with set_current_vllm_config(self.vllm_config):
             # To make vLLM config available during worker initialization
diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py
index 39b3b10349c3..a86110f2f864 100644
--- a/vllm/v1/worker/workspace.py
+++ b/vllm/v1/worker/workspace.py
@@ -3,11 +3,12 @@
 
 import inspect
 import os
-from collections.abc import Iterator
+from collections.abc import Callable, Hashable, Iterator
 from contextlib import contextmanager
 from contextvars import ContextVar
 from itertools import accumulate
 from math import prod
+from typing import TypeVar, cast
 
 import torch
 
@@ -17,6 +18,7 @@
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 logger = init_logger(__name__)
+T = TypeVar("T")
 
 
 def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int:
@@ -47,8 +49,8 @@ def use_workspace_lane(lane: int) -> Iterator[None]:
 class WorkspaceManager:
     """Manager for workspace allocation.
 
-    Manages one workspace buffer per active ``(ubatch, lane)`` slot.
-    Can be locked to prevent further growth during execution.
+    Manages shared scratch and keyed persistent resources per ``(ubatch, lane)``.
+    Can be locked to prevent growth or new resources during execution.
     """
 
     def __init__(
@@ -66,6 +68,9 @@ def __init__(
         self._current_workspaces: list[torch.Tensor | None] = [None] * (
             self._num_ubatches * self._num_lanes
         )
+        self._persistent_resources: list[dict[Hashable, object]] = [
+            {} for _ in self._current_workspaces
+        ]
         self._locked: bool = False
 
     @staticmethod
@@ -76,10 +81,10 @@ def _workspace_size_bytes(workspace: torch.Tensor | None) -> int:
         return workspace.numel() * workspace.element_size()
 
     def lock(self) -> None:
-        """Lock the workspace to prevent further growth.
+        """Lock the workspace to prevent growth or new persistent resources.
 
-        After locking, any attempt to allocate a larger workspace will raise
-        an assertion error. This ensures workspace size is fixed during execution.
+        Larger scratch requests and new persistent resources raise an assertion
+        error. Existing resources retain their storage throughout execution.
         """
         self._locked = True
         if envs.VLLM_DEBUG_WORKSPACE:
@@ -113,6 +118,63 @@ def is_locked(self) -> bool:
         """Check if workspace is locked."""
         return self._locked
 
+    def _get_workspace_id(self) -> int:
+        lane = _workspace_lane.get()
+        if lane >= self._num_lanes:
+            raise RuntimeError(
+                f"Workspace lane {lane} is not configured; manager has "
+                f"{self._num_lanes} lane(s)."
+            )
+        return dbo_current_ubatch_id() * self._num_lanes + lane
+
+    def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T:
+        """Return the resource cached for ``key`` in the current ubatch and lane.
+
+        The first request calls ``factory`` and stores its result for the lifetime
+        of this manager. Later requests return the same object without calling
+        the factory again. Each ubatch/lane pair has a separate cache. Once the
+        manager is locked, requesting a key missing from that cache raises an
+        AssertionError.
+
+        Keys must identify the resource type and configuration. Callers must not
+        resize or replace tensor storage held by a cached resource. Uses of the
+        same resource must not overlap; include the CUDA stream in the key when
+        streams can execute concurrently.
+
+        With torch.compile, call this inside a runtime custom op so resource
+        selection happens during execution.
+        """
+        resources = self._persistent_resources[self._get_workspace_id()]
+        if key not in resources:
+            if self._locked:
+                raise AssertionError(
+                    f"Workspace is locked but {key!r} was not allocated during "
+                    "warmup for the current ubatch and lane."
+                )
+            resources[key] = factory()
+        return cast(T, resources[key])
+
+    def get_persistent(
+        self,
+        key: Hashable,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        *,
+        zero_init: bool = False,
+    ) -> torch.Tensor:
+        """Get a fixed-shape tensor, optionally zeroed on first allocation only."""
+        factory = torch.zeros if zero_init else torch.empty
+        workspace = self.get_persistent_resource(
+            (torch.Tensor, key),
+            lambda: factory(shape, dtype=dtype, device=self._device),
+        )
+        if workspace.shape != shape or workspace.dtype != dtype:
+            raise ValueError(
+                f"Persistent workspace {key!r} has shape={workspace.shape}, "
+                f"dtype={workspace.dtype}; requested shape={shape}, dtype={dtype}."
+            )
+        return workspace
+
     def get_simultaneous(
         self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype]
     ) -> list[torch.Tensor]:
@@ -123,6 +185,7 @@ def get_simultaneous(
 
         Returns:
             List of tensor views into the workspace buffer, one per shape/dtype pair.
+
         """
         actual_bytes = [_compute_bytes(s, d) for s, d in shapes_and_dtypes]
         aligned_bytes = [round_up(actual, 256) for actual in actual_bytes]
@@ -148,15 +211,9 @@ def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor:
 
         Returns:
             The current workspace tensor.
+
         """
-        ubatch_id = dbo_current_ubatch_id()
-        lane = _workspace_lane.get()
-        if lane >= self._num_lanes:
-            raise RuntimeError(
-                f"Workspace lane {lane} is not configured; manager has "
-                f"{self._num_lanes} lane(s)."
-            )
-        workspace_id = ubatch_id * self._num_lanes + lane
+        workspace_id = self._get_workspace_id()
         current_workspace = self._current_workspaces[workspace_id]
         current_size = self._workspace_size_bytes(current_workspace)
 
@@ -216,8 +273,8 @@ def get_caller_info() -> str:
                     get_caller_info(),
                     current_size / _MB,
                     required_bytes / _MB,
-                    ubatch_id,
-                    lane,
+                    workspace_id // self._num_lanes,
+                    workspace_id % self._num_lanes,
                 )
 
         return current_workspace
@@ -228,6 +285,7 @@ def is_workspace_manager_initialized() -> bool:
 
     Returns:
         True if workspace manager is initialized, False otherwise.
+
     """
     return _manager is not None
 
@@ -237,6 +295,7 @@ def current_workspace_manager() -> "WorkspaceManager":
 
     Raises:
         AssertionError: If workspace manager has not been initialized.
+
     """
     assert _manager is not None, (
         "WorkspaceManager not initialized. Call init_workspace_manager() "
@@ -259,6 +318,7 @@ def init_workspace_manager(
         device: The device to allocate workspace on.
         num_ubatches: Number of workspace ubatch slots. Defaults to 1.
         num_lanes: Number of independent execution lanes per ubatch. Defaults to 1.
+
     """
     global _manager
     if _manager is not None:
@@ -272,12 +332,11 @@ def init_workspace_manager(
 
 
 def lock_workspace() -> None:
-    """Lock the workspace to prevent further growth.
+    """Lock the workspace to prevent growth or new persistent resources.
 
-    After calling this function, any attempt to allocate a workspace larger
-    than the current size will raise an AssertionError. This ensures that
-    workspace size is fixed during execution and prevents unexpected memory
-    allocations in the hot path.
+    Larger scratch requests and new persistent resources raise an AssertionError.
+    All required resources must be initialized in each execution slot during
+    warmup, so their storage remains fixed during execution.
 
     Example:
         # During initialization
@@ -289,6 +348,7 @@ def lock_workspace() -> None:
         lock_workspace()
 
         # Now all get_workspace calls must fit in pre-allocated size
+
     """
     current_workspace_manager().lock()
 
diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py
index 68d6bde3a049..02cc2c9e87d5 100644
--- a/vllm/v1/worker/xpu_worker.py
+++ b/vllm/v1/worker/xpu_worker.py
@@ -8,7 +8,6 @@
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from vllm.profiler.wrapper import TorchProfilerWrapper
 from vllm.utils.mem_utils import MemorySnapshot, format_gib
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.utils import report_usage_stats
@@ -115,27 +114,55 @@ def init_device(self):
 
             self.local_rank += dp_local_rank * replica_world_size
 
-            assert self.local_rank < visible_device_count, (
-                f"DP adjusted local rank {self.local_rank} is out of bounds. "
-            )
-            assert parallel_config.local_world_size <= visible_device_count, (
-                f"local_world_size ({parallel_config.local_world_size}) must "
-                f"be less than or equal to the number of visible devices "
-                f"({visible_device_count})."
-            )
-
         device = self.device_config.device
         if (
             isinstance(device, torch.device)
             and device.type == "xpu"
             and current_platform.is_xpu()
         ):
-            self.device = torch.device(f"xpu:{self.local_rank}")
+            assigned_physical_gpu_ids = self.parallel_config.assigned_physical_gpu_ids
+            if assigned_physical_gpu_ids is not None:
+                from vllm.platforms.interface import set_assigned_physical_gpu_ids
+
+                set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
+                assert self.local_rank < len(assigned_physical_gpu_ids), (
+                    f"local_rank {self.local_rank} is out of bounds for "
+                    f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"
+                )
+                # NOTE: local_world_size is derived from parallel_config.nnodes,
+                # which is only set for the "mp" multi-node backend. With the
+                # "ray"/"external_launcher" backends nnodes stays 1, so
+                # local_world_size collapses to the full world_size and this
+                # check wrongly fires on cross-node deployments.
+                # assigned_physical_gpu_ids is already per-node and the
+                # local_rank bound above fully validates the mapping for
+                # these backends, so skip the check for them.
+                if parallel_config.distributed_executor_backend not in (
+                    "ray",
+                    "external_launcher",
+                ):
+                    assert parallel_config.local_world_size <= len(
+                        assigned_physical_gpu_ids
+                    ), (
+                        f"local_world_size ({parallel_config.local_world_size})"
+                        " exceeds assigned_physical_gpu_ids count "
+                        f"({len(assigned_physical_gpu_ids)})"
+                    )
+            else:
+                assert self.local_rank < torch.accelerator.device_count(), (
+                    f"DP adjusted local rank {self.local_rank} is out of "
+                    f"bounds for {torch.accelerator.device_count()} devices."
+                )
+
+            visible_device_index = (
+                current_platform.logical_device_id_to_visible_device_id(self.local_rank)
+            )
+            self.device = torch.device(f"xpu:{visible_device_index}")
             torch.accelerator.set_device_index(self.device)
             current_platform.check_if_supports_dtype(self.model_config.dtype)
             torch.accelerator.empty_cache()
             self.init_gpu_memory = torch.xpu.get_device_properties(
-                self.local_rank
+                visible_device_index
             ).total_memory
         else:
             raise RuntimeError(f"Unsupported device type: {self.device_config.device}")
@@ -197,33 +224,6 @@ def init_device(self):
             # If usage stat is enabled, collect relevant info.
             report_usage_stats(self.vllm_config)
 
-    def profile(self, is_start: bool = True, profile_prefix: str | None = None):
-        if self.profiler_config is None or self.profiler_config.profiler is None:
-            raise RuntimeError(
-                "Profiling is not enabled. Please set --profiler-config to enable "
-                "profiling. Example: "
-                "'--profiler-config.profiler=torch --profiler-config.torch_profiler_dir"
-                "=YOUR_DIR_PATH_TO_DUMP_TRACE'"
-            )
-
-        if is_start and self.profiler is None:
-            from vllm.distributed.utils import get_worker_rank_suffix
-
-            rank_suffix = get_worker_rank_suffix(global_rank=self.rank)
-            trace_name = (
-                f"{profile_prefix}_{rank_suffix}" if profile_prefix else rank_suffix
-            )
-
-            self.profiler = TorchProfilerWrapper(
-                self.profiler_config,
-                worker_name=trace_name,
-                local_rank=self.local_rank,
-                activities=["CPU", "XPU"],
-            )
-            logger.debug("Starting torch profiler with trace name: %s", trace_name)
-
-        super().profile(is_start=is_start, profile_prefix=profile_prefix)
-
     def shutdown(self) -> None:
         logger.info(
             "XPUWorker shutdown: cleaning up (rank=%d, local_rank=%d)",
diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py
index 81768e12092e..3c4c0d8948be 100644
--- a/vllm/vllm_flash_attn/flash_attn_interface.py
+++ b/vllm/vllm_flash_attn/flash_attn_interface.py
@@ -261,11 +261,41 @@ def flash_attn_varlen_func(
         return_attn_probs: bool. Whether to return the attention probabilities. This option is for
            testing only. The returned probabilities are not guaranteed to be correct
            (they might not have the right scaling).
+        seqused_k: (batch_size,), dtype torch.int32. Actual number of key tokens per
+           sequence. Mutually exclusive with cu_seqlens_k, and required when
+           block_table is given.
+        q_v: optional value tensor fused with q. FA3 only.
+        block_table: (batch_size, max_blocks_per_seq), dtype torch.int32. Paged KV
+           cache block table. Requires seqused_k.
+        return_softmax_lse: bool. Whether to also return softmax_lse.
+        out: optional tensor to write the output into.
+        scheduler_metadata: precomputed scheduling metadata. Not supported by FA2.
+        q_descale: descale factor for FP8 q. Not supported by FA2, and ignored by
+           FA4 when the inputs are not FP8.
+        k_descale: descale factor for FP8 k. Same constraints as q_descale.
+        v_descale: descale factor for FP8 v. Same constraints as q_descale.
+        num_splits: int. Number of splits for the split-KV kernel; 0 lets the kernel
+           choose. FA2 only supports 0 or 1.
+        output_scale: scale for fused FP8 output quantization. FA4 only.
+        fa_version: int. FlashAttention major version to dispatch to.
+        s_aux: learnable per-head attention sink logits. Not supported by FA2.
+        cp_world_size: int. Context parallel world size.
+        cp_rank: int. Rank of this shard within the context parallel group.
+        cp_tot_seqused_k: (batch_size,), dtype torch.int32. Key lengths across the
+           whole context parallel group.
+        mask_mod: optional callable applying a custom mask. Not supported by FA2.
+        block_sparse_tensors: block-sparse index tensors. FA4 only, and must
+           materialize full_block_cnt and full_block_idx.
+        aux_tensors: auxiliary tensors consumed by mask_mod. FA4 only.
+        aux_tensor_leading_dims: leading dimensions of each entry in aux_tensors.
+        dynamic_causal: optional per-sequence causal offsets. FA4 only.
+
     Return:
         out: (total, nheads, headdim).
         softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The
             logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
             normalization factor).
+
     """
     assert cu_seqlens_k is not None or seqused_k is not None, (
         "cu_seqlens_k or seqused_k must be provided"
@@ -404,6 +434,14 @@ def flash_attn_varlen_func(
 
         from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd
 
+        # SM90 FA4 fp8-KV path: fp8 e4m3 paged K/V dequantized and the K/V descale folded
+        # in-kernel; accepts bf16/fp16 Q and writes O in its native dtype (no Q cast, no
+        # output copy). Only the (batch, num_kv_heads) f32 K/V descales are forwarded.
+        fa4_fp8_kv_dequant = (
+            k.dtype == torch.float8_e4m3fn
+            and torch.cuda.get_device_capability()[0] == 9
+        )
+
         out, softmax_lse, _, _ = _flash_attn_fwd(
             q,
             k,
@@ -428,10 +466,11 @@ def flash_attn_varlen_func(
             block_sparse_tensors=block_sparse_tensors,
             aux_tensors=aux_tensors,
             aux_tensor_leading_dims=aux_tensor_leading_dims,
-            q_descale=q_descale,
+            q_descale=None if fa4_fp8_kv_dequant else q_descale,
             k_descale=k_descale,
             v_descale=v_descale,
             output_scale=output_scale,
+            fp8_kv_dequant=fa4_fp8_kv_dequant,
         )
     else:
         raise ValueError(f"Unsupported FA version: {fa_version}")
@@ -445,6 +484,7 @@ def compile_flash_attn_varlen_func_from_specs(
     v_shape: tuple[int, ...],
     q_dtype: torch.dtype,
     v_stride: tuple[int, ...] | None = None,
+    k_stride: tuple[int, ...] | None = None,
     cu_seqlens_q_shape: tuple[int, ...] | None = None,
     cu_seqlens_k_shape: tuple[int, ...] | None = None,
     max_seqlen_q: int | None = None,
@@ -457,6 +497,9 @@ def compile_flash_attn_varlen_func_from_specs(
     return_softmax_lse=False,
     num_splits: int = 0,
     fa_version: int = DEFAULT_FA_VERSION,
+    seqused_k_shape: tuple[int, ...] | None = None,
+    page_table_shape: tuple[int, ...] | None = None,
+    softcap: float | None = None,
 ) -> None:
     if fa_version != 4:
         raise ValueError(
@@ -467,7 +510,8 @@ def compile_flash_attn_varlen_func_from_specs(
     del deterministic
 
     from vllm.vllm_flash_attn.cute.interface import (
-        compile_flash_attn_varlen_func_from_specs as _fa4_compile_flash_attn_varlen_func_from_specs,
+        _flash_attn_fwd,
+        _make_compile_only_tensor_spec,
     )
 
     real_window_size: tuple[int, int]
@@ -480,21 +524,50 @@ def compile_flash_attn_varlen_func_from_specs(
     if softmax_scale is None:
         softmax_scale = q_shape[-1] ** (-0.5)
 
-    return _fa4_compile_flash_attn_varlen_func_from_specs(
-        q_shape=q_shape,
-        k_shape=k_shape,
-        v_shape=v_shape,
-        q_dtype=q_dtype,
-        v_stride=v_stride,
-        cu_seqlens_q_shape=cu_seqlens_q_shape,
-        cu_seqlens_k_shape=cu_seqlens_k_shape,
+    q = _make_compile_only_tensor_spec(q_shape, q_dtype)
+    k = _make_compile_only_tensor_spec(k_shape, q_dtype, stride=k_stride)
+    v = _make_compile_only_tensor_spec(v_shape, q_dtype, stride=v_stride)
+    cu_seqlens_q = _make_compile_only_tensor_spec(cu_seqlens_q_shape, torch.int32, 4)
+    out = _make_compile_only_tensor_spec(
+        (*q_shape[:-1], v_shape[-1]),
+        q_dtype,
+    )
+    lse = None
+    if return_softmax_lse:
+        assert q is not None
+        if cu_seqlens_q_shape is None:
+            lse_shape = (*q.shape[:-3], q.shape[-2], q.shape[-3])
+            lse_stride = None
+        else:
+            lse_shape = (q.shape[-2], q.shape[0])
+            lse_stride = (None, 1)
+        lse = _make_compile_only_tensor_spec(
+            lse_shape,
+            torch.float32,
+            4,
+            stride=lse_stride,
+        )
+
+    _flash_attn_fwd(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=_make_compile_only_tensor_spec(cu_seqlens_k_shape, torch.int32, 4),
+        seqused_k=_make_compile_only_tensor_spec(seqused_k_shape, torch.int32, 4),
+        page_table=_make_compile_only_tensor_spec(page_table_shape, torch.int32, 4),
         max_seqlen_q=max_seqlen_q,
         max_seqlen_k=max_seqlen_k,
         softmax_scale=softmax_scale,
         causal=causal,
-        window_size=real_window_size,
+        softcap=softcap,
+        window_size_left=real_window_size[0],
+        window_size_right=real_window_size[1],
         num_splits=num_splits,
         return_lse=return_softmax_lse,
+        out=out,
+        lse=lse,
+        compile_only=True,
     )
 
 
@@ -543,11 +616,16 @@ def sparse_attn_func(
         return_attn_probs: bool. Whether to return the attention probabilities. This option is for
            testing only. The returned probabilities are not guaranteed to be correct
            (they might not have the right scaling).
+        softcap: float. Anything > 0 activates softcapping attention.
+        return_softmax_lse: bool. Whether to also return softmax_lse.
+        out: optional tensor to write the output into.
+
     Return:
         out: (batch_size, seqlen, nheads, headdim).
         softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
             logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
             normalization factor).
+
     """
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** (-0.5)
@@ -629,11 +707,15 @@ def sparse_attn_varlen_func(
         return_attn_probs: bool. Whether to return the attention probabilities. This option is for
            testing only. The returned probabilities are not guaranteed to be correct
            (they might not have the right scaling).
+        return_softmax_lse: bool. Whether to also return softmax_lse.
+        out: optional tensor to write the output into.
+
     Return:
         out: (total, nheads, headdim).
         softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The
             logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
             normalization factor).
+
     """
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** (-0.5)